Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
 hsliuustc0106
|
202961e56c |
No files matched your search
@@ -52,7 +52,6 @@ time timeout -k 30 30m docker run \
|
||||
pip install tblib==3.1.0
|
||||
cd /workspace/vllm-omni
|
||||
pytest -v -s -m "core_model and xpu and B60"
|
||||
pytest -v -s tests/diffusion/quantization/test_mxfp8_config.py
|
||||
export VLLM_XPU_USE_SAMPLER_KERNEL=0 # NOTE: Remove this after vLLM v0.21.1 is merged. Fixes Qwen2-5 omni-expansion tests.
|
||||
pytest -v -s -m "advanced_model and xpu and B60"
|
||||
'
|
||||
@@ -389,31 +389,6 @@ steps:
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Omni · Higgs-Audio V2 Test"
|
||||
timeout_in_minutes: 30
|
||||
commands:
|
||||
- |
|
||||
timeout 30m bash -c "
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
export VLLM_USE_DEEP_GEMM=0
|
||||
export VLLM_MOE_USE_DEEP_GEMM=0
|
||||
pytest -s -v tests/e2e/online_serving/test_higgs_audio_v2.py tests/e2e/offline_inference/test_higgs_audio_v2.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
"
|
||||
agents:
|
||||
queue: "gpu_1_queue"
|
||||
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: "Diffusion · Qwen Image Edit Test"
|
||||
timeout_in_minutes: 20
|
||||
commands:
|
||||
|
||||
+5
-105
@@ -305,51 +305,6 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: Omni · Perf Test · Multi-Replica"
|
||||
key: nightly-omni-performance-multi-replicas
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- export BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen3_omni_multi_replicas.json
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
|
||||
exit $$EXIT
|
||||
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: 3
|
||||
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 · Multi-Replica Startup Test with 4x H100"
|
||||
timeout_in_minutes: 45
|
||||
@@ -570,7 +525,6 @@ steps:
|
||||
tests/e2e/online_serving/test_longcat_image_expansion.py
|
||||
tests/e2e/online_serving/test_longcat_image_edit_expansion.py
|
||||
tests/e2e/online_serving/test_flux_2_dev_expansion.py
|
||||
tests/e2e/offline_inference/test_glm_image_autoround_w4a16_expansion.py
|
||||
-m "full_model and diffusion and H100 and distributed_cuda"
|
||||
--run-level "full_model"
|
||||
agents:
|
||||
@@ -843,8 +797,8 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Edit-2511"
|
||||
key: nightly-diffusion-x2iat-performance-qwen-image-edit-2511
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Edit-2509"
|
||||
key: nightly-diffusion-x2iat-performance-qwen-image-edit-2509
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
@@ -852,7 +806,7 @@ steps:
|
||||
- export CACHE_DIT_VERSION=1.3.0
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2509_vllm_omni.json
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
|
||||
@@ -939,54 +893,6 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · BAGEL"
|
||||
key: nightly-diffusion-x2iat-performance-bagel
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
|
||||
- export CACHE_DIT_VERSION=1.3.0
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_bagel_vllm_omni.json
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
|
||||
exit $$EXIT
|
||||
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 X2I(&A&T) · HunyuanImage3 · DiT Perf Test"
|
||||
key: nightly-hunyuan-image3-performance
|
||||
soft_fail: true
|
||||
@@ -1092,8 +998,6 @@ steps:
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py -m full_model --run-level full_model
|
||||
- pytest -s -v tests/e2e/accuracy/hunyuanvideo15_t2v/test_hunyuanvideo15_t2v_video_similarity.py -m full_model --run-level full_model
|
||||
- pytest -s -v tests/e2e/accuracy/hunyuanvideo15_i2v/test_hunyuanvideo15_i2v_video_similarity.py -m full_model --run-level full_model
|
||||
- pytest -s -v tests/e2e/accuracy/test_ltx2_3_video_similarity.py -m full_model --run-level full_model
|
||||
- pytest -s -v tests/e2e/accuracy/test_diffusers_backend_similarity.py -k '2v_matches_diffusers' -m full_model --run-level full_model
|
||||
agents:
|
||||
@@ -1284,13 +1188,11 @@ steps:
|
||||
- nightly-omni-performance-no-async-chunk
|
||||
- nightly-omni-performance-async-chunk
|
||||
- nightly-omni-performance-vllm-text
|
||||
- nightly-omni-performance-multi-replicas
|
||||
- nightly-tts-performance
|
||||
- nightly-diffusion-x2iat-performance-qwen-image
|
||||
- nightly-diffusion-x2iat-performance-qwen-image-edit
|
||||
- nightly-diffusion-x2iat-performance-qwen-image-edit-2511
|
||||
- nightly-diffusion-x2iat-performance-qwen-image-edit-2509
|
||||
- nightly-diffusion-x2iat-performance-qwen-image-layered
|
||||
- nightly-diffusion-x2iat-performance-bagel
|
||||
- nightly-diffusion-x2v-performance
|
||||
- nightly-testcase-statistics
|
||||
if: build.env("NIGHTLY") == "1" && build.env("EMAIL_DISTRIBUTION") == "1"
|
||||
@@ -1302,12 +1204,10 @@ steps:
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-no-async-chunk
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-async-chunk
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-vllm-text
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-multi-replicas
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-edit
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-edit-2511
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-edit-2509
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-qwen-image-layered
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance-bagel
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2v-performance
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.html" . --step nightly-testcase-statistics
|
||||
- python tools/nightly/generate_nightly_perf_excel.py
|
||||
|
||||
@@ -21,9 +21,6 @@ Before proposing changes, ask for the optimization scene if it is not already kn
|
||||
- Current enabled strategies: USP/SP, CFG parallel, HSDP/FSDP, VAE patch parallel, torch.compile, profiler options.
|
||||
- Optimization target: latency, throughput, memory, cost, or quality-preserving speed.
|
||||
- Precision/quality tolerance: bf16/fp8/quantization/sparsity/approximate attention allowed or not.
|
||||
- Quality validation method: exact-output comparison when feasible, or metrics
|
||||
such as SSIM, PSNR, LPIPS, cosine similarity, MAE, MSE, temporal flicker
|
||||
checks, baseline self-run variance, and reviewable output artifacts.
|
||||
|
||||
## Workflow
|
||||
|
||||
@@ -82,19 +79,7 @@ Before proposing changes, ask for the optimization scene if it is not already kn
|
||||
- Record latency, stage timings, memory, output quality, and logs.
|
||||
- Report comparison tables in ms. Include at least end-to-end client time, server `inference_time`, server stage generation time if available, `vae.encode`, `diffuse`, `vae.decode`, and peak memory.
|
||||
|
||||
6. **Enforce the quality and precision gate for every optimization.**
|
||||
- Do not mark an optimization ready until quality is checked on the same
|
||||
workload: model, prompt/media, seed, shape, steps, scheduler, dtype,
|
||||
backend, parallelism, and output encoding.
|
||||
- For math-preserving changes, require exact/near-exact agreement or show
|
||||
that differences are within baseline self-run variance.
|
||||
- For precision, quantization, approximate attention, sparse/custom kernels,
|
||||
or backend changes, include reviewable artifacts and metrics such as SSIM,
|
||||
PSNR, LPIPS, cosine similarity, MAE, MSE, and temporal flicker checks.
|
||||
- If quality tolerance is not stated, default to quality-preserving behavior.
|
||||
Failed or inconclusive quality validation blocks a ready-to-merge claim.
|
||||
|
||||
7. **Collect diagnostic trace only after narrowing hypotheses.**
|
||||
6. **Collect diagnostic trace only after narrowing hypotheses.**
|
||||
- Use torch profiler for a small number of requests.
|
||||
- Run two separate diagnostic traces instead of mixing concerns:
|
||||
- **Operator/shape trace**: enable `torch_profiler_record_shapes=True` and keep stack collection disabled. Use this to rank CUDA kernels, NCCL collectives, attention/MLP/norm/RoPE work, and shape-specific hot operators.
|
||||
@@ -106,16 +91,15 @@ Before proposing changes, ask for the optimization scene if it is not already kn
|
||||
- Diagnostic reports must be written to disk and preserve: server command, profiler config, warmup/request/polling commands, trace artifact paths, rank analyzed, analyzer output or summary, and the decision about whether additional ranks are necessary.
|
||||
- Analyze both rank-level balance and device-level free bubbles when additional ranks are opened.
|
||||
|
||||
8. **Analyze host, communication, and operators.**
|
||||
7. **Analyze host, communication, and operators.**
|
||||
- Find GPU idle/free intervals and map each large gap to the enclosing CPU/Python code.
|
||||
- Separate real GPU idle from profiler overhead such as CUPTI `Command Buffer Full`.
|
||||
- Compare NCCL kernel time to user annotations; annotations can overcount nested intervals.
|
||||
- Rank operator work by total CUDA time and by repeated small-kernel launch count.
|
||||
|
||||
9. **Produce an optimization plan.**
|
||||
8. **Produce an optimization plan.**
|
||||
- Classify candidates as P0/P1/P2.
|
||||
- For each candidate, state necessity, expected benefit, implementation path, validation plan, and quality risk.
|
||||
- Include a concrete quality/precision validation gate for each candidate.
|
||||
- Do not implement high-risk operator rewrites before proving the operator is a bottleneck for the target shapes.
|
||||
- End the plan with a user-facing candidate selection table. The assistant
|
||||
should not automatically choose a risky optimization just because it is
|
||||
@@ -140,13 +124,7 @@ Before proposing changes, ask for the optimization scene if it is not already kn
|
||||
- **P1:** meaningful code changes with contained risk. Examples: cross-attention KV caching, VAE gather/broadcast reduction, AdaLayerNorm/RMSNorm/RoPE fusion after trace evidence.
|
||||
- **P2:** high implementation or quality risk. Examples: FA to LA replacement, custom Triton/CUDA fused kernels, FP8/quantization, sparsity/Rainfusion-style acceleration.
|
||||
|
||||
Every implemented optimization needs A/B validation and a passing
|
||||
quality/precision gate, including math-preserving changes such as padding trim,
|
||||
layout cleanup, cache reuse, or host/runtime cleanup. Use objective metrics
|
||||
when available, keep reviewable artifacts, and compare against baseline
|
||||
self-run variance for precision, quantization, approximate kernels, or backend
|
||||
changes. If quality validation fails or is inconclusive, do not present the
|
||||
optimization as ready to merge.
|
||||
Every implemented optimization needs A/B validation and quality regression. A/B means same workload and hardware before/after. Quality regression means checking generated image/video stability, artifacts, temporal flicker, and seed behavior when precision or approximate kernels change.
|
||||
|
||||
## Optimization Layers
|
||||
|
||||
|
||||
@@ -10,11 +10,6 @@ build:
|
||||
jobs:
|
||||
post_checkout:
|
||||
- git fetch --unshallow || true
|
||||
pre_install:
|
||||
# Pre-install CPU torch so the subsequent `pip install .[docs]` does not
|
||||
# pull ~2 GB of CUDA 13 libs (cuBLAS/cuDNN/NCCL/...) that exceed RTD's
|
||||
# 1 GB memory cap and time out the mkdocs build.
|
||||
- python -m pip install --index-url https://download.pytorch.org/whl/cpu torch
|
||||
|
||||
mkdocs:
|
||||
configuration: mkdocs.yml
|
||||
|
||||
+35
-51
@@ -1,60 +1,44 @@
|
||||
# Benchmarks
|
||||
# Benchmarks Overview and Architecture
|
||||
|
||||
This directory contains benchmark suites for evaluating different model families and infrastructure components in vLLM-Omni. Each subfolder targets a different benchmark family with its own scripts, configs, and metrics. See the per-subfolder READMEs for detailed usage.
|
||||
This document explains the benchmark architecture across all benchmark assets in this repo. It describes what we measure, and where to find or plug in new scenarios. Per-task details remain in subfolder READMEs (e.g., `benchmarks/<model>/README.md`).
|
||||
|
||||
## Benchmark families
|
||||
## Scope and goals
|
||||
- Establish repeatable latency/throughput measurements for multimodal LLM pipelines.
|
||||
- Provide both HF Transformers (offline) and vLLM-Omni (multi-stage/pipeline) baselines.
|
||||
- Make it easy to plug in new datasets and models with minimal changes to the runner scripts.
|
||||
|
||||
### [TTS](tts/README.md) — Text-to-Speech
|
||||
## Dataset and inputs
|
||||
- Default example: SeedTTS top-100 prompts (`benchmarks/build_dataset/top100.txt`) via `benchmarks/build_dataset/`.
|
||||
- Extensible: drop in new prompt files or modality-aligned payloads; keep the expected format for the consuming scripts (e.g., one prompt per line).
|
||||
- If you add a new dataset, document it under `benchmarks/<model>/README.md` and point scripts to your data path.
|
||||
|
||||
Model-agnostic serving benchmarks for TTS models, including Qwen3-TTS and VoxCPM2.
|
||||
## Directory layout
|
||||
- `benchmarks/build_dataset/` — dataset prep utilities (e.g., SeedTTS top100).
|
||||
- `benchmarks/<model>/vllm_omni/` — vLLM-Omni pipeline benchmarks, logs, outputs.
|
||||
- `benchmarks/accuracy/` — accuracy benchmark integrations that adapt external
|
||||
benchmark suites to vLLM-Omni serving and evaluation flows.
|
||||
- Add new tasks under `benchmarks/<model>/...` with the same pattern: `transformers/`, `vllm_omni/`, task-specific README, and (optionally) dataset prep notes.
|
||||
|
||||
- **Layout**: `tts/bench_tts.py` (serving benchmark driver), `tts/model_configs.yaml` (model registry), `tts/plot_results.py` (result plotting)
|
||||
- **Dataset**: Seed-TTS full or text-only datasets, plus bundled smoke/design prompts under `build_dataset/`
|
||||
- **Key metrics**: TTFP (time to first audio packet), E2E latency, RTF (real-time factor), throughput (audio seconds / wall-clock second)
|
||||
## Reference workflows
|
||||
- **HF Transformers (offline, single process)**
|
||||
Script (example): `benchmarks/<model>/transformers/eval_qwen3_moe_omni_transformers.sh`
|
||||
Outputs: `benchmark_results/perf_stats.json`, `benchmark_results/results.json`, `benchmark_results/audio/` (if audio is produced).
|
||||
|
||||
### [Diffusion](diffusion/README.md) — Image and Video Generation
|
||||
- **vLLM-Omni end-to-end pipeline**
|
||||
Script (example): `benchmarks/<model>/vllm_omni/eval_qwen3_moe_omni.sh`
|
||||
Outputs: `vllm_omni/logs/*.stats.jsonl` (per-stage/overall latency & TPS), `vllm_omni/logs/stage*.log`, `vllm_omni/outputs/` (text/audio artifacts).
|
||||
|
||||
Online-serving benchmark for diffusion image/video models, sending requests to the configured vLLM serving endpoint (`/v1/chat/completions`, `/v1/images/generations`, or `/v1/videos`, depending on backend/task).
|
||||
- **Adding a new task/model**
|
||||
1) Create `benchmarks/<model>/transformers/` and/or `benchmarks/<model>/vllm_omni/` with scripts referencing your model and dataset.
|
||||
2) Add a task README describing dataset, configs, and expected outputs.
|
||||
3) Keep the output/log structure similar for easy comparison (perf_stats/results/audio or text outputs; stats.jsonl/logs for pipeline).
|
||||
|
||||
- **Tasks**: text-to-image, text-to-video, image-to-image, image-to-video, text+image-to-image, text+image-to-video
|
||||
- **Datasets**: `vbench`, `trace`, `random`
|
||||
- **Key metrics**: request throughput, latency percentiles, optional SLO attainment
|
||||
## Metrics to watch
|
||||
- **Throughput**: `overall_tps`, `*_tps_avg` per stage.
|
||||
- **Latency distribution**: look for long tails in `*.stats.jsonl`.
|
||||
- **Quality/completeness**: missing outputs or errors in stage logs indicate pipeline failures or misconfigurations.
|
||||
|
||||
### [GLM-Image](glm_image/README.md) — Text-to-Image and Image-to-Image
|
||||
|
||||
Benchmarks for GLM-Image performance across HuggingFace baseline, vLLM-Omni offline inference, and vLLM-Omni online serving.
|
||||
|
||||
- **Layout**: `glm_image/huggingface/` (HF baseline), `glm_image/vllm-omni/` (offline inference), `glm_image/benchmark_glm_image.py` (online serving)
|
||||
- **Tasks**: text-to-image and image-to-image
|
||||
- **Key metrics**: request/image throughput, latency percentiles, optional per-stage pipeline timings
|
||||
|
||||
### [Distributed](distributed/omni_connectors/README.md) — RDMA Connector Testing
|
||||
|
||||
RDMA environment setup and transfer tests for `MooncakeTransferEngineConnector`, including pytest-based single-node checks and manual cross-node benchmarks.
|
||||
|
||||
- **Transfer modes**: `copy`, `zerocopy`, `gpu` (GPUDirect)
|
||||
- **Supports**: single-node pytest suites and manual multi-node/cross-node transfer testing
|
||||
|
||||
### [Accuracy](accuracy/README.md) — Image Generation and Editing Quality
|
||||
|
||||
Accuracy benchmarks for image generation/editing models, adapting external suites to vLLM-Omni serving and local judge-evaluation flows.
|
||||
|
||||
- **Layout**: `accuracy/text_to_image/` (GEBench), `accuracy/image_to_image/` (GEdit-Bench)
|
||||
- **Method**: generation and judge scoring both run through local `vllm-omni serve` endpoints
|
||||
|
||||
### Common serving metrics framework
|
||||
|
||||
`vllm_omni/benchmarks/` extends `vllm bench serve --omni` with Omni-specific datasets, backends, and multimodal metrics. Key metrics include:
|
||||
|
||||
- **Text output**: TTFT (time to first token), TPOT (time per output token), ITL (inter-token latency)
|
||||
- **Audio output**: TTFP (time to first audio packet), E2E latency, RTF (real-time factor)
|
||||
- **Throughput**: request throughput, output token throughput, total token throughput, audio throughput
|
||||
|
||||
See `vllm_omni/benchmarks/serve.py` for the `vllm bench serve --omni` runner wrapper and `vllm_omni/benchmarks/metrics/` for Omni metric definitions.
|
||||
|
||||
## Adding a new benchmark
|
||||
|
||||
1. Create a subfolder under `benchmarks/<name>/` with scripts, configs if needed, and a `README.md`.
|
||||
2. If comparing against another runtime, use clear backend subfolders where applicable, such as `huggingface/` and `vllm-omni/`, or follow the shared TTS serving benchmark pattern in `tts/`.
|
||||
3. Add reusable dataset or prompt-building utilities to `build_dataset/` if applicable.
|
||||
4. Update this README with a link to the new benchmark family.
|
||||
## Troubleshooting
|
||||
- Verify GPU/driver/FlashAttention2 requirements for your chosen model/config.
|
||||
- Ensure network access for dataset/model downloads (Google Drive, Hugging Face, etc.).
|
||||
- If outputs are missing or slow, inspect per-stage logs and `*.stats.jsonl` for errors, stragglers, or contention.
|
||||
@@ -1,328 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Standalone diagnostic for the diffusion attention backends added in #3079.
|
||||
|
||||
Exercises the same synthetic attention shape through each backend and each
|
||||
torch SDPA sub-kernel, with and without an attn_mask, so we can pinpoint why
|
||||
CUDNN_ATTN underperforms on some SKUs. Two typical causes:
|
||||
|
||||
(a) cuDNN lacks a tuned kernel for this (SM, head_dim, seq) combination
|
||||
and silently falls back to MATH.
|
||||
(b) attn_mask is non-None and the cuDNN SDPA dispatch rejects it,
|
||||
walking CUDNN -> FLASH -> MATH.
|
||||
|
||||
Run:
|
||||
python benchmarks/diffusion/bench_attention_backends.py --preset hv15
|
||||
python benchmarks/diffusion/bench_attention_backends.py --preset wan22
|
||||
python benchmarks/diffusion/bench_attention_backends.py \
|
||||
--batch 1 --heads 24 --seq 14336 --head-dim 128
|
||||
|
||||
Optional extras:
|
||||
pip install --pre flash-attn-4 # FA4 is currently pre-release only
|
||||
pip install -U flashinfer # latest FlashInfer (0.6.9)
|
||||
|
||||
The table at the end is the data we want on the PR — surface the row where a
|
||||
backend is >1.5x the SDPA baseline, that's the one to gate off the auto-route.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import time
|
||||
|
||||
import torch
|
||||
from torch.nn.attention import SDPBackend, sdpa_kernel
|
||||
|
||||
# Presets mirror the DiT hot-path attention call (single-stream block) at the
|
||||
# shapes we run in the PR's validation configs.
|
||||
_PRESETS = {
|
||||
# HunyuanVideo-1.5 480p/33f: (30 * 52 * 9) latent + ~256 text tokens.
|
||||
"hv15": {"batch": 1, "heads": 24, "seq": 14336, "head_dim": 128},
|
||||
# Wan 2.2 480p/33f rough estimate; adjust when we measure real shapes.
|
||||
"wan22": {"batch": 1, "heads": 40, "seq": 16384, "head_dim": 128},
|
||||
# Smaller image-gen shape for quick smoke tests.
|
||||
"flux": {"batch": 1, "heads": 24, "seq": 4096, "head_dim": 128},
|
||||
}
|
||||
|
||||
_SDPA_BACKENDS = [
|
||||
("CUDNN_ATTENTION", [SDPBackend.CUDNN_ATTENTION]),
|
||||
("FLASH_ATTENTION", [SDPBackend.FLASH_ATTENTION]),
|
||||
("EFFICIENT_ATTENTION", [SDPBackend.EFFICIENT_ATTENTION]),
|
||||
("MATH", [SDPBackend.MATH]),
|
||||
# The PR's CUDNN_ATTN impl uses this priority chain; keep it here to show
|
||||
# what torch actually picks when multiple are allowed.
|
||||
("CUDNN_ATTN_CHAIN", [SDPBackend.CUDNN_ATTENTION, SDPBackend.FLASH_ATTENTION, SDPBackend.MATH]),
|
||||
]
|
||||
|
||||
|
||||
def _env_report() -> None:
|
||||
if not torch.cuda.is_available():
|
||||
raise SystemExit("CUDA not available; this diagnostic needs a GPU.")
|
||||
cc = torch.cuda.get_device_capability()
|
||||
print("=" * 72)
|
||||
print(f"GPU : {torch.cuda.get_device_name()}")
|
||||
print(f"SM capability : sm_{cc[0]}{cc[1]}")
|
||||
print(f"torch : {torch.__version__}")
|
||||
print(f"cuDNN version : {torch.backends.cudnn.version()}")
|
||||
print("flashinfer : ", end="")
|
||||
try:
|
||||
import flashinfer
|
||||
|
||||
print(getattr(flashinfer, "__version__", "present"))
|
||||
except Exception as e:
|
||||
print(f"not installed ({type(e).__name__})")
|
||||
print("=" * 72)
|
||||
|
||||
|
||||
def _make_qkv(batch: int, heads: int, seq: int, head_dim: int, device: str, dtype: torch.dtype):
|
||||
# Layout: (B, S, H, D) — what the PR's backends accept as input.
|
||||
q = torch.randn(batch, seq, heads, head_dim, device=device, dtype=dtype)
|
||||
k = torch.randn(batch, seq, heads, head_dim, device=device, dtype=dtype)
|
||||
v = torch.randn(batch, seq, heads, head_dim, device=device, dtype=dtype)
|
||||
return q, k, v
|
||||
|
||||
|
||||
def _make_mask(batch: int, seq: int, device: str, dtype: torch.dtype, pad_tokens: int = 128) -> torch.Tensor:
|
||||
# Mirrors a text-encoder padding mask: last `pad_tokens` positions masked.
|
||||
mask = torch.zeros(batch, 1, seq, seq, device=device, dtype=dtype)
|
||||
if pad_tokens > 0:
|
||||
mask[..., -pad_tokens:] = float("-inf")
|
||||
return mask
|
||||
|
||||
|
||||
def _time_call(fn, *args, warmup: int = 3, iters: int = 10) -> tuple[float, str]:
|
||||
"""Return (median ms, error string). Skips on any exception so that a
|
||||
backend that rejects our args (wrong dtype, missing JIT module, unsupported
|
||||
kwarg value) doesn't abort the whole sweep."""
|
||||
try:
|
||||
for _ in range(warmup):
|
||||
fn(*args)
|
||||
torch.accelerator.synchronize()
|
||||
times = []
|
||||
for _ in range(iters):
|
||||
torch.accelerator.synchronize()
|
||||
t0 = time.perf_counter()
|
||||
fn(*args)
|
||||
torch.accelerator.synchronize()
|
||||
times.append((time.perf_counter() - t0) * 1000.0)
|
||||
times.sort()
|
||||
return times[len(times) // 2], ""
|
||||
except Exception as e: # noqa: BLE001 — probe script, keep the whole table going
|
||||
msg = str(e).split("\n", 1)[0][:60]
|
||||
return float("nan"), f"{type(e).__name__}: {msg}" if msg else type(e).__name__
|
||||
|
||||
|
||||
def _run_sdpa_variants(q, k, v, attn_mask, scale: float) -> list[tuple[str, float, str]]:
|
||||
# q/k/v come in (B, S, H, D); F.sdpa wants (B, H, S, D).
|
||||
q_t = q.permute(0, 2, 1, 3)
|
||||
k_t = k.permute(0, 2, 1, 3)
|
||||
v_t = v.permute(0, 2, 1, 3)
|
||||
rows: list[tuple[str, float, str]] = []
|
||||
for name, backends in _SDPA_BACKENDS:
|
||||
|
||||
def _call(mask=attn_mask):
|
||||
with sdpa_kernel(backends):
|
||||
return torch.nn.functional.scaled_dot_product_attention(
|
||||
q_t, k_t, v_t, attn_mask=mask, dropout_p=0.0, is_causal=False, scale=scale
|
||||
)
|
||||
|
||||
ms, err = _time_call(_call)
|
||||
rows.append((name, ms, err))
|
||||
return rows
|
||||
|
||||
|
||||
def _run_flashinfer(q, k, v, scale: float, backend: str | None = None, attn_mask=None) -> tuple[float, str]:
|
||||
"""Call FlashInfer's dense single-prefill.
|
||||
|
||||
``backend`` hints at cutlass/fa3/trtllm if the installed version exposes
|
||||
that kwarg. ``attn_mask`` is the same additive 4D mask we pass to SDPA —
|
||||
we convert it to the 2D boolean ``custom_mask`` form FlashInfer accepts.
|
||||
Per FlashInfer docs, ``custom_mask`` only applies when ``causal=False``.
|
||||
"""
|
||||
try:
|
||||
import inspect
|
||||
|
||||
from flashinfer.prefill import single_prefill_with_kv_cache
|
||||
except Exception as e:
|
||||
return float("nan"), f"import-{type(e).__name__}"
|
||||
|
||||
kwargs: dict = {"sm_scale": scale, "causal": False, "return_lse": False}
|
||||
if backend is not None:
|
||||
sig = inspect.signature(single_prefill_with_kv_cache)
|
||||
if "backend" not in sig.parameters:
|
||||
return float("nan"), "no-backend-kwarg"
|
||||
kwargs["backend"] = backend
|
||||
|
||||
if attn_mask is not None:
|
||||
# Collapse (B, 1, S, S) additive float mask to (S, S) boolean.
|
||||
# FlashInfer expects True = keep, False = masked out.
|
||||
mask_2d = attn_mask[0, 0]
|
||||
kwargs["custom_mask"] = mask_2d != float("-inf")
|
||||
|
||||
def _call():
|
||||
out = single_prefill_with_kv_cache(q[0], k[0], v[0], **kwargs)
|
||||
return out.unsqueeze(0)
|
||||
|
||||
return _time_call(_call)
|
||||
|
||||
|
||||
def _run_fa4(q, k, v, scale: float) -> tuple[float, str]:
|
||||
"""Call FlashAttention-4 directly (``pip install flash-attn-4``). FA4
|
||||
ships a Blackwell-native kernel via CuTe-DSL; on sm_120 it should beat
|
||||
cuDNN by ~20%. API lives under ``flash_attn.cute``; the older
|
||||
``flash_attn.flash_attn_func`` path is FA2/FA3 only."""
|
||||
try:
|
||||
from flash_attn.cute import flash_attn_func
|
||||
except Exception as e:
|
||||
return float("nan"), f"import-{type(e).__name__}"
|
||||
|
||||
def _call():
|
||||
# FA4 accepts (B, S, H, D) directly — same layout the PR's backends use.
|
||||
return flash_attn_func(q, k, v, softmax_scale=scale, causal=False)
|
||||
|
||||
return _time_call(_call)
|
||||
|
||||
|
||||
def _run_flashinfer_cudnn_batch(q, k, v, scale: float) -> tuple[float, str]:
|
||||
"""Call FlashInfer's direct cuDNN wrapper. Bypasses PyTorch SDPA dispatch
|
||||
(which has a few hundred ns of overhead per call) and talks to cuDNN FMHA
|
||||
straight. Useful as a ceiling for 'pure cuDNN, no SDPA' on Blackwell."""
|
||||
try:
|
||||
from flashinfer.prefill import cudnn_batch_prefill_with_kv_cache
|
||||
except Exception as e:
|
||||
return float("nan"), f"import-{type(e).__name__}"
|
||||
|
||||
b, s, h, d = q.shape
|
||||
qo_indptr = torch.tensor([0, s], dtype=torch.int32, device=q.device)
|
||||
kv_indptr = torch.tensor([0, s], dtype=torch.int32, device=q.device)
|
||||
|
||||
def _call():
|
||||
return cudnn_batch_prefill_with_kv_cache(
|
||||
q.reshape(b * s, h, d),
|
||||
k.reshape(b * s, h, d),
|
||||
v.reshape(b * s, h, d),
|
||||
qo_indptr=qo_indptr,
|
||||
kv_indptr=kv_indptr,
|
||||
max_qo_len=s,
|
||||
max_kv_len=s,
|
||||
sm_scale=scale,
|
||||
causal=False,
|
||||
)
|
||||
|
||||
return _time_call(_call)
|
||||
|
||||
|
||||
def _print_table(title: str, rows: list[tuple[str, float, str]], baseline_name: str | None = None) -> None:
|
||||
print(f"\n{title}")
|
||||
print("-" * 88)
|
||||
header = f"{'backend':<24} {'median (ms)':>14} {'vs baseline':>12} status"
|
||||
print(header)
|
||||
print("-" * 88)
|
||||
baseline_ms = None
|
||||
if baseline_name is not None:
|
||||
for name, ms, _ in rows:
|
||||
if name == baseline_name and ms == ms: # not NaN
|
||||
baseline_ms = ms
|
||||
break
|
||||
for name, ms, err in rows:
|
||||
ms_str = f"{ms:>14.3f}" if ms == ms else f"{'n/a':>14}"
|
||||
if baseline_ms is not None and ms == ms:
|
||||
ratio = baseline_ms / ms
|
||||
ratio_str = f"{ratio:>11.2f}x"
|
||||
else:
|
||||
ratio_str = f"{'—':>12}"
|
||||
status = "ok" if not err else f"FAILED ({err})"
|
||||
print(f"{name:<24} {ms_str} {ratio_str} {status}")
|
||||
|
||||
|
||||
def _pick_winner(rows: list[tuple[str, float, str]]) -> tuple[str, float] | None:
|
||||
"""Return (backend, ms) of the fastest non-failing row, or None."""
|
||||
ok_rows = [(n, m) for n, m, err in rows if not err and m == m]
|
||||
if not ok_rows:
|
||||
return None
|
||||
return min(ok_rows, key=lambda x: x[1])
|
||||
|
||||
|
||||
def _bench_one_shape(shape: dict, dtype: torch.dtype, device: str) -> tuple[list, list]:
|
||||
torch.manual_seed(0)
|
||||
q, k, v = _make_qkv(**shape, device=device, dtype=dtype)
|
||||
scale = 1.0 / (shape["head_dim"] ** 0.5)
|
||||
|
||||
rows_nomask = _run_sdpa_variants(q, k, v, attn_mask=None, scale=scale)
|
||||
rows_nomask.append(("FLASHINFER (default)", *_run_flashinfer(q, k, v, scale)))
|
||||
# `trtllm-gen` has no sm_120 cubins (NVIDIA/TensorRT-LLM#11799) — skipped.
|
||||
for fi_backend in ("fa2", "fa3", "cutlass", "auto"):
|
||||
rows_nomask.append((f"FLASHINFER ({fi_backend})", *_run_flashinfer(q, k, v, scale, backend=fi_backend)))
|
||||
rows_nomask.append(("FLASHINFER (cudnn-batch)", *_run_flashinfer_cudnn_batch(q, k, v, scale)))
|
||||
rows_nomask.append(("FA4 (direct)", *_run_fa4(q, k, v, scale)))
|
||||
|
||||
mask = _make_mask(shape["batch"], shape["seq"], device=device, dtype=dtype, pad_tokens=256)
|
||||
rows_mask = _run_sdpa_variants(q, k, v, attn_mask=mask, scale=scale)
|
||||
rows_mask.append(("FLASHINFER (dense)", *_run_flashinfer(q, k, v, scale, attn_mask=mask)))
|
||||
rows_mask.append(("FA4 (direct)", float("nan"), "mask-not-supported"))
|
||||
|
||||
return rows_nomask, rows_mask
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
parser.add_argument("--preset", choices=list(_PRESETS.keys()), default="hv15")
|
||||
parser.add_argument("--sweep", action="store_true", help="Run all presets and print a ranking")
|
||||
parser.add_argument("--batch", type=int, default=None)
|
||||
parser.add_argument("--heads", type=int, default=None)
|
||||
parser.add_argument("--seq", type=int, default=None)
|
||||
parser.add_argument("--head-dim", type=int, default=None)
|
||||
parser.add_argument("--dtype", choices=["bf16", "fp16"], default="bf16")
|
||||
parser.add_argument("--device", default="cuda:0")
|
||||
args = parser.parse_args()
|
||||
|
||||
dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16
|
||||
_env_report()
|
||||
print(f"dtype : {args.dtype}")
|
||||
|
||||
presets = list(_PRESETS.keys()) if args.sweep else [args.preset]
|
||||
summary: dict[str, tuple[str, float] | None] = {}
|
||||
|
||||
for preset_name in presets:
|
||||
shape = dict(_PRESETS[preset_name])
|
||||
if not args.sweep:
|
||||
for k in ("batch", "heads", "seq", "head_dim"):
|
||||
v = getattr(args, k if k != "head_dim" else "head_dim")
|
||||
if v is not None:
|
||||
shape[k] = v
|
||||
|
||||
print("\n" + "=" * 88)
|
||||
print(
|
||||
f"Preset: {preset_name} | batch={shape['batch']} heads={shape['heads']} "
|
||||
f"seq={shape['seq']} head_dim={shape['head_dim']}"
|
||||
)
|
||||
print("=" * 88)
|
||||
|
||||
rows_nomask, rows_mask = _bench_one_shape(shape, dtype, args.device)
|
||||
_print_table(f"[{preset_name}] No attention mask", rows_nomask, baseline_name="CUDNN_ATTENTION")
|
||||
_print_table(
|
||||
f"[{preset_name}] With attention mask (pad 256 tokens)", rows_mask, baseline_name="CUDNN_ATTENTION"
|
||||
)
|
||||
|
||||
summary[preset_name] = _pick_winner(rows_nomask)
|
||||
|
||||
if args.sweep or len(presets) > 1:
|
||||
print("\n" + "=" * 88)
|
||||
print("Winners per preset (no-mask path)")
|
||||
print("=" * 88)
|
||||
for preset_name, winner in summary.items():
|
||||
if winner is None:
|
||||
print(f" {preset_name:<10} — no successful backend")
|
||||
else:
|
||||
name, ms = winner
|
||||
print(f" {preset_name:<10} {name:<24} {ms:>8.3f} ms")
|
||||
|
||||
print("\nNotes:")
|
||||
print(" * Ratios are relative to CUDNN_ATTENTION. >1.0x means faster than cuDNN.")
|
||||
print(" * Mask-path winner inherits CUDNN_ATTN's fallback in the PR's backends.")
|
||||
print(" * `trtllm-gen` and FA4 4.0.0b10 are known-broken on sm_120 as of Apr 2026.")
|
||||
print(" * For e2e timings, run text_to_video.py / text_to_image.py while varying DIFFUSION_ATTENTION_BACKEND.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -47,13 +47,6 @@ Usage:
|
||||
--endpoint /v1/images/edits --dataset random --task ti2i --num-prompts 10 \
|
||||
--bot-task think
|
||||
|
||||
python benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
--endpoint /v1/images/edits \
|
||||
--dataset custom \
|
||||
--dataset-path custom_requests.jsonl \
|
||||
--task ti2i \
|
||||
--bot-task think
|
||||
|
||||
i2i (chat-based models such as Qwen-Image-Edit):
|
||||
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
--endpoint /v1/chat/completions --dataset vbench --task i2i --num-prompts 10
|
||||
@@ -573,135 +566,6 @@ class TraceDataset(BaseDataset):
|
||||
return [self[i] for i in range(len(self))]
|
||||
|
||||
|
||||
class CustomDataset(BaseDataset):
|
||||
"""
|
||||
Custom dataset that loads requests from a JSONL file.
|
||||
|
||||
Each line in the JSONL file should be a JSON object with the following fields:
|
||||
- prompt (required): The text prompt for the request
|
||||
- width (optional): Image/video width
|
||||
- height (optional): Image height
|
||||
- num_inference_steps (optional): Number of diffusion steps
|
||||
- seed (optional): Random seed
|
||||
- image_paths (optional): List of input image paths for i2i/i2v/ti2i tasks
|
||||
- image_urls (optional): List of input image URLs (alternative to image_paths)
|
||||
|
||||
Example JSONL for ti2i:
|
||||
{"prompt": "Add sunset lighting", "width": 1024, "height": 1024, "image_urls": ["https://example.com/image.jpg"]}
|
||||
"""
|
||||
|
||||
def __init__(self, args, api_url: str, model: str):
|
||||
super().__init__(args, api_url, model)
|
||||
self.dataset_path = args.dataset_path
|
||||
if not self.dataset_path:
|
||||
raise ValueError("--dataset-path must be provided when using 'custom' dataset")
|
||||
self.load_data()
|
||||
|
||||
def load_data(self) -> None:
|
||||
"""Load data from JSONL file."""
|
||||
import pandas as pd
|
||||
|
||||
if not self.dataset_path.endswith(".jsonl"):
|
||||
raise ValueError("Custom dataset must be a JSONL file")
|
||||
|
||||
try:
|
||||
df = pd.read_json(path_or_buf=self.dataset_path, lines=True)
|
||||
except Exception as e:
|
||||
raise ValueError(f"Failed to load JSONL file {self.dataset_path}: {e}")
|
||||
|
||||
if "prompt" not in df.columns:
|
||||
raise ValueError("JSONL file must contain a 'prompt' column")
|
||||
|
||||
self.data = df.to_dict("records")
|
||||
print(f"Loaded {len(self.data)} requests from {self.dataset_path}")
|
||||
|
||||
def _resolve_image_paths(self, item: dict) -> list[str] | None:
|
||||
# Handle local file paths
|
||||
if "image_paths" in item and item["image_paths"]:
|
||||
paths = item["image_paths"]
|
||||
if isinstance(paths, str):
|
||||
paths = [paths]
|
||||
|
||||
# Verify all paths exist
|
||||
valid_paths = []
|
||||
for path in paths:
|
||||
if os.path.exists(path):
|
||||
valid_paths.append(path)
|
||||
else:
|
||||
raise ValueError(f"Image file not found: {path}")
|
||||
|
||||
return valid_paths if valid_paths else None
|
||||
|
||||
# Handle URLs - download to temp files
|
||||
if "image_urls" in item and item["image_urls"]:
|
||||
urls = item["image_urls"]
|
||||
if isinstance(urls, str):
|
||||
urls = [urls]
|
||||
|
||||
downloaded_paths = []
|
||||
for url in urls:
|
||||
try:
|
||||
response = requests.get(url, timeout=30)
|
||||
response.raise_for_status()
|
||||
|
||||
# Create temp file with appropriate extension
|
||||
suffix = os.path.splitext(url)[1] or ".png"
|
||||
with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
|
||||
tmp.write(response.content)
|
||||
downloaded_paths.append(tmp.name)
|
||||
except Exception as e:
|
||||
raise ValueError(f"Failed to download image from {url}: {e}")
|
||||
|
||||
return downloaded_paths if downloaded_paths else None
|
||||
|
||||
return None
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.data)
|
||||
|
||||
def __getitem__(self, idx: int) -> RequestFuncInput:
|
||||
item = self.data[idx]
|
||||
|
||||
# Extract fields with fallback to args defaults
|
||||
prompt = item["prompt"]
|
||||
width = item.get("width", self.args.width)
|
||||
height = item.get("height", self.args.height)
|
||||
num_inference_steps = item.get("num_inference_steps", self.args.num_inference_steps)
|
||||
seed = item.get("seed", self.args.seed)
|
||||
reserved_keys = {"prompt", "width", "height", "num_inference_steps", "seed", "image_paths", "image_urls"}
|
||||
extra_body = {k: v for k, v in item.items() if k not in reserved_keys}
|
||||
|
||||
# Handle image paths/URLs
|
||||
image_paths = self._resolve_image_paths(item)
|
||||
|
||||
return RequestFuncInput(
|
||||
prompt=prompt,
|
||||
api_url=self.api_url,
|
||||
model=self.model,
|
||||
seed=seed,
|
||||
image_paths=image_paths,
|
||||
extra_body=extra_body,
|
||||
width=width,
|
||||
height=height,
|
||||
num_inference_steps=num_inference_steps,
|
||||
)
|
||||
|
||||
def get_requests(self) -> list[RequestFuncInput]:
|
||||
"""Get all requests, cycling through data if num_prompts > dataset size."""
|
||||
num_requests = getattr(self.args, "num_prompts", len(self.data))
|
||||
|
||||
if num_requests <= len(self.data):
|
||||
return [self[i] for i in range(num_requests)]
|
||||
|
||||
# Cycle through the dataset to reach num_prompts
|
||||
requests = []
|
||||
for i in range(num_requests):
|
||||
idx = i % len(self.data)
|
||||
requests.append(self[idx])
|
||||
|
||||
return requests
|
||||
|
||||
|
||||
class RandomDataset(BaseDataset):
|
||||
def __init__(self, args, api_url: str, model: str, enable_negative_prompt: bool = False):
|
||||
super().__init__(args, api_url, model)
|
||||
@@ -1086,8 +950,6 @@ async def benchmark(args):
|
||||
dataset = TraceDataset(args, api_url, args.model)
|
||||
elif args.dataset == "random":
|
||||
dataset = RandomDataset(args, api_url, args.model, args.enable_negative_prompt)
|
||||
elif args.dataset == "custom":
|
||||
dataset = CustomDataset(args, api_url, args.model)
|
||||
else:
|
||||
raise ValueError(f"Unknown dataset: {args.dataset}")
|
||||
|
||||
@@ -1235,7 +1097,7 @@ if __name__ == "__main__":
|
||||
"--dataset",
|
||||
type=str,
|
||||
default="vbench",
|
||||
choices=["vbench", "trace", "random", "custom"],
|
||||
choices=["vbench", "trace", "random"],
|
||||
help="Dataset to use.",
|
||||
)
|
||||
parser.add_argument(
|
||||
|
||||
@@ -169,13 +169,15 @@ class DiffusionRequestStatus(enum.IntEnum):
|
||||
|
||||
@dataclass
|
||||
class DiffusionRequestState:
|
||||
request_id: str
|
||||
sched_req_id: str
|
||||
req: OmniDiffusionRequest
|
||||
status: DiffusionRequestStatus = DiffusionRequestStatus.WAITING
|
||||
```
|
||||
|
||||
**Design Features**:
|
||||
|
||||
- **Scheduler-owned ID**: Each `OmniDiffusionRequest` is tracked by an internal `sched_req_id`, separated from public `request_id` values.
|
||||
|
||||
- **Explicit lifecycle**: Requests move through waiting, running, optional preemption, and terminal states.
|
||||
|
||||
- **Centralized error handling**: Completion, abort, and error states are all normalized in the scheduler layer.
|
||||
@@ -186,6 +188,7 @@ class DiffusionRequestState:
|
||||
class _BaseScheduler(SchedulerInterface):
|
||||
def __init__(self) -> None:
|
||||
self._request_states = {}
|
||||
self._request_id_to_sched_req_id = {}
|
||||
self._waiting = deque()
|
||||
self._running = []
|
||||
self._finished_req_ids = set()
|
||||
@@ -196,8 +199,7 @@ class _BaseScheduler(SchedulerInterface):
|
||||
|
||||
- **Common state storage**: Shared request maps and waiting/running sets live in the base class.
|
||||
|
||||
- **Shared cleanup logic**: Duplicate-request checks, finish handling, and state removal are centralized instead of
|
||||
duplicated in each policy.
|
||||
- **Shared cleanup logic**: Request-id registration, finish handling, and state removal are centralized instead of duplicated in each policy.
|
||||
|
||||
- **Current constraint boundary**: `_BaseScheduler` derives `max_num_running_reqs` from `max_num_seqs`, but request-mode diffusion is still clamped back to `1` by the engine. The step-wise path can keep this above `1` for compatible-request batching.
|
||||
|
||||
@@ -217,12 +219,12 @@ class RequestScheduler(_BaseScheduler):
|
||||
|
||||
- **Single-request admission**: `RequestScheduler` still admits one active request at a time because request-mode execution completes a whole request per dispatch.
|
||||
|
||||
- **Executor result feedback**: `update_from_output()` converts executor output into `FINISHED_COMPLETED` or `FINISHED_ERROR` and returns finished request ids.
|
||||
- **Executor result feedback**: `update_from_output()` converts executor output into `FINISHED_COMPLETED` or `FINISHED_ERROR` and returns finished scheduler ids.
|
||||
|
||||
#### 2.5 Engine-Driven Execution Loop
|
||||
|
||||
```python
|
||||
request_id = scheduler.add_request(request)
|
||||
sched_req_id = scheduler.add_request(request)
|
||||
while True:
|
||||
sched_output = scheduler.schedule()
|
||||
output = executor.add_req(req)
|
||||
|
||||
@@ -16,7 +16,6 @@ ROOT_DIR_RELATIVE = "../../../../.."
|
||||
EXAMPLE_DIR = ROOT_DIR / "examples"
|
||||
EXAMPLE_DOC_DIR = ROOT_DIR / "docs/user_guide/examples"
|
||||
NAV_FILE = ROOT_DIR / "docs/.nav.yml"
|
||||
MAX_INLINE_MATERIAL_SIZE = 8 * 1024
|
||||
|
||||
|
||||
def fix_case(text: str) -> str:
|
||||
@@ -192,21 +191,18 @@ class Example:
|
||||
|
||||
return re.sub(link_pattern, replace_link, content)
|
||||
|
||||
def github_url(self, path: Path) -> str:
|
||||
url = "https://github.com/vllm-project/vllm-omni/"
|
||||
url += "tree/main" if path.is_dir() else "blob/main"
|
||||
return f"{url}/{path.relative_to(ROOT_DIR).as_posix()}"
|
||||
|
||||
def generate(self) -> str:
|
||||
content = f"# {self.title}\n\n"
|
||||
content += f"Source <{self.github_url(self.path)}>.\n\n"
|
||||
url = "https://github.com/vllm-project/vllm-omni/"
|
||||
url += "tree/main" if self.path.is_dir() else "blob/main"
|
||||
content += f"Source <{url}/{self.path.relative_to(ROOT_DIR)}>.\n\n"
|
||||
|
||||
# Use long code fence to avoid issues with
|
||||
# included files containing code fences too
|
||||
code_fence = "``````"
|
||||
|
||||
if self.is_code:
|
||||
main_file_rel = self.main_file.relative_to(ROOT_DIR).as_posix()
|
||||
main_file_rel = self.main_file.relative_to(ROOT_DIR)
|
||||
content += f'{code_fence}{self.main_file.suffix[1:]}\n--8<-- "{main_file_rel}"\n{code_fence}\n'
|
||||
else:
|
||||
with open(self.main_file, encoding="utf-8") as f:
|
||||
@@ -220,15 +216,10 @@ class Example:
|
||||
|
||||
content += "## Example materials\n\n"
|
||||
for file in sorted(self.other_files):
|
||||
content += f'??? abstract "{file.relative_to(self.path).as_posix()}"\n'
|
||||
if file.stat().st_size > MAX_INLINE_MATERIAL_SIZE:
|
||||
content += (
|
||||
f" Large file omitted from the rendered docs. View it on GitHub: <{self.github_url(file)}>.\n\n"
|
||||
)
|
||||
continue
|
||||
content += f'??? abstract "{file.relative_to(self.path)}"\n'
|
||||
if file.suffix != ".md":
|
||||
content += f" {code_fence}{file.suffix[1:]}\n"
|
||||
content += f' --8<-- "{file.relative_to(ROOT_DIR).as_posix()}"\n'
|
||||
content += f' --8<-- "{file.relative_to(ROOT_DIR)}"\n'
|
||||
if file.suffix != ".md":
|
||||
content += f" {code_fence}\n"
|
||||
|
||||
@@ -300,7 +291,7 @@ def update_nav_file(examples: list[Example]):
|
||||
for example in category_examples:
|
||||
doc_path = EXAMPLE_DOC_DIR / example.category / f"{example.path.stem}.md"
|
||||
rel_path = doc_path.relative_to(ROOT_DIR / "docs")
|
||||
category_items.append({example.title: rel_path.as_posix()})
|
||||
category_items.append({example.title: str(rel_path)})
|
||||
|
||||
if category_items:
|
||||
# Format category name (e.g., "offline_inference" -> "Offline Inference")
|
||||
|
||||
@@ -1,29 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Reduce browser search-index memory for generated API reference pages."""
|
||||
|
||||
from material.plugins.search.plugin import SearchIndex
|
||||
|
||||
_add_entry_from_context = SearchIndex.add_entry_from_context
|
||||
|
||||
|
||||
def add_entry_from_context(self: SearchIndex, page) -> None:
|
||||
search = page.meta.get("search") or {}
|
||||
if search.get("exclude"):
|
||||
return
|
||||
|
||||
if str(page.url).startswith("api/"):
|
||||
entry = {
|
||||
"location": page.url,
|
||||
"title": str(page.meta.get("title", page.title)),
|
||||
"text": "",
|
||||
}
|
||||
if "boost" in search:
|
||||
entry["boost"] = search["boost"]
|
||||
self.entries.append(entry)
|
||||
return
|
||||
|
||||
_add_entry_from_context(self, page)
|
||||
|
||||
|
||||
SearchIndex.add_entry_from_context = add_entry_from_context
|
||||
@@ -1,44 +1,20 @@
|
||||
// Enables MathJax rendering only on pages that contain math.
|
||||
// Enables MathJax rendering
|
||||
window.MathJax = {
|
||||
tex: {
|
||||
inlineMath: [["\\(", "\\)"]],
|
||||
displayMath: [["\\[", "\\]"]],
|
||||
processEscapes: true,
|
||||
processEnvironments: true,
|
||||
},
|
||||
options: {
|
||||
ignoreHtmlClass: ".*|",
|
||||
processHtmlClass: "arithmatex",
|
||||
},
|
||||
};
|
||||
tex: {
|
||||
inlineMath: [["\\(", "\\)"]],
|
||||
displayMath: [["\\[", "\\]"]],
|
||||
processEscapes: true,
|
||||
processEnvironments: true
|
||||
},
|
||||
options: {
|
||||
ignoreHtmlClass: ".*|",
|
||||
processHtmlClass: "arithmatex"
|
||||
}
|
||||
};
|
||||
|
||||
let mathJaxLoading;
|
||||
|
||||
function loadMathJax() {
|
||||
if (window.MathJax.typesetPromise) {
|
||||
return Promise.resolve();
|
||||
}
|
||||
if (!mathJaxLoading) {
|
||||
mathJaxLoading = new Promise((resolve, reject) => {
|
||||
const script = document.createElement("script");
|
||||
script.src = "https://unpkg.com/mathjax@3.2.2/es5/tex-mml-chtml.js";
|
||||
script.async = true;
|
||||
script.onload = resolve;
|
||||
script.onerror = reject;
|
||||
document.head.appendChild(script);
|
||||
});
|
||||
}
|
||||
return mathJaxLoading;
|
||||
}
|
||||
|
||||
document$.subscribe(() => {
|
||||
if (!document.querySelector(".arithmatex")) {
|
||||
return;
|
||||
}
|
||||
loadMathJax().then(() => {
|
||||
MathJax.startup.output.clearCache();
|
||||
MathJax.typesetClear();
|
||||
MathJax.texReset();
|
||||
MathJax.typesetPromise();
|
||||
});
|
||||
});
|
||||
document$.subscribe(() => {
|
||||
MathJax.startup.output.clearCache()
|
||||
MathJax.typesetClear()
|
||||
MathJax.texReset()
|
||||
MathJax.typesetPromise()
|
||||
})
|
||||
@@ -1,2 +0,0 @@
|
||||
// Opt out of bfcache so large generated API pages are released after navigation.
|
||||
window.addEventListener("unload", () => {});
|
||||
@@ -1,42 +1,21 @@
|
||||
// Load Mermaid only on pages that contain diagrams.
|
||||
let mermaidLoading;
|
||||
// Initialize Mermaid for diagram rendering
|
||||
mermaid.initialize({
|
||||
startOnLoad: false,
|
||||
theme: 'default',
|
||||
securityLevel: 'loose',
|
||||
flowchart: {
|
||||
useMaxWidth: true,
|
||||
htmlLabels: true
|
||||
}
|
||||
});
|
||||
|
||||
function loadMermaid() {
|
||||
if (window.mermaid) {
|
||||
return Promise.resolve();
|
||||
// Render Mermaid diagrams when page content is ready
|
||||
document$.subscribe(() => {
|
||||
const mermaidElements = document.querySelectorAll('.mermaid');
|
||||
if (mermaidElements.length > 0) {
|
||||
mermaid.run({
|
||||
querySelector: '.mermaid',
|
||||
nodes: mermaidElements
|
||||
});
|
||||
}
|
||||
if (!mermaidLoading) {
|
||||
mermaidLoading = new Promise((resolve, reject) => {
|
||||
const script = document.createElement("script");
|
||||
script.src = "https://unpkg.com/mermaid@10/dist/mermaid.min.js";
|
||||
script.async = true;
|
||||
script.onload = () => {
|
||||
mermaid.initialize({
|
||||
startOnLoad: false,
|
||||
theme: "default",
|
||||
securityLevel: "loose",
|
||||
flowchart: {
|
||||
useMaxWidth: true,
|
||||
htmlLabels: true,
|
||||
},
|
||||
});
|
||||
resolve();
|
||||
};
|
||||
script.onerror = reject;
|
||||
document.head.appendChild(script);
|
||||
});
|
||||
}
|
||||
return mermaidLoading;
|
||||
}
|
||||
|
||||
document$.subscribe(() => {
|
||||
if (!document.querySelector(".mermaid")) {
|
||||
return;
|
||||
}
|
||||
loadMermaid().then(() => {
|
||||
mermaid.run({
|
||||
querySelector: ".mermaid",
|
||||
nodes: document.querySelectorAll(".mermaid"),
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -19,7 +19,6 @@ th {
|
||||
|--------------|--------|-------------------|------------|---------|-----|-----------|
|
||||
| `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `MingFlashOmniForConditionalGeneration` + `MingImagePipeline` | Ming-flash-omni-2.0 (omni-speech + imagegen<sup>1</sup>) | `Jonathan1909/Ming-flash-omni-2.0` | ✅︎ | | | |
|
||||
| `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `InternVLAA1Pipeline` | InternVLA-A1 | `InternRobotics/InternVLA-A1-3B` | ✅︎ | ✅︎ | | |
|
||||
| `HunyuanImage3ForCausalMM` | HunyuanImage3.0 (DiT-only) | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
|
||||
@@ -13,7 +13,7 @@ For example...
|
||||
```bash
|
||||
# Qwen-Image
|
||||
vllm serve Qwen/Qwen-Image-Edit-2511 --omni --port 8000
|
||||
```
|
||||
|
||||
|
||||
### Generate Images
|
||||
|
||||
@@ -53,7 +53,7 @@ result = client.images.edit(
|
||||
output_format='jpeg',
|
||||
# url格式
|
||||
extra_body={
|
||||
"url": [input_image_url,input_image_url],
|
||||
"url": [input_image_url1,input_image_url],
|
||||
"num_inference_steps": 50,
|
||||
"guidance_scale": 1,
|
||||
"seed": 777,
|
||||
@@ -92,8 +92,7 @@ Content-Type: multipart/form-data
|
||||
| `user` | string | null | User identifier for tracking |
|
||||
| `output_format` | string | "png" | The format in which the generated images are returned. Must be one of "png", "jpg", "jpeg", "webp". |
|
||||
| `output_compression` | integer | 100 | The compression level (0-100%) for the generated images. |
|
||||
| `background` | string or null | "auto" | Allows to set transparency for the background of the generated image(s). |
|
||||
| `stream` | boolean | false | Return Server-Sent Events instead of a single JSON response. Currently supported only by multi-stage image edit pipelines that produce AR text before the final image, such as HunyuanImage3 IT2I AR+DiT. |
|
||||
| `background` | string or null | "auto" | Allows to set transparency for the background of the generated image(s).
|
||||
|
||||
#### vllm-omni Extension Parameters
|
||||
|
||||
@@ -110,8 +109,6 @@ Content-Type: multipart/form-data
|
||||
|
||||
### Response Format
|
||||
|
||||
When `stream=false` or omitted, the endpoint returns the standard image edit response:
|
||||
|
||||
```json
|
||||
{
|
||||
"created": 1701234567,
|
||||
@@ -123,56 +120,10 @@ When `stream=false` or omitted, the endpoint returns the standard image edit res
|
||||
}
|
||||
],
|
||||
"output_format": null,
|
||||
"size": null
|
||||
"size": null,
|
||||
}
|
||||
```
|
||||
|
||||
### Streaming Response Format
|
||||
|
||||
Set `stream=true` when you want to receive the HunyuanImage3 IT2I AR
|
||||
recaption text before the final edited image is ready. Streaming is only
|
||||
available for multi-stage image edit pipelines where stage 0 produces AR text
|
||||
and a later diffusion stage produces the image. Single-stage diffusion image
|
||||
edit pipelines reject `stream=true` with HTTP 400.
|
||||
|
||||
The response uses Server-Sent Events with `Content-Type: text/event-stream`.
|
||||
Each event is sent as one `data:` line. The stream order is:
|
||||
|
||||
1. One or more AR text delta chunks.
|
||||
2. One final image chunk.
|
||||
3. `data: [DONE]`.
|
||||
|
||||
AR text chunks expose only the generated text delta and the AR completion
|
||||
`index`. They do not expose token ids, ratio tokens, KV cache data, or internal
|
||||
prompt token ids.
|
||||
|
||||
Example AR delta event:
|
||||
|
||||
```text
|
||||
data: {"object":"image.edit.chunk","type":"ar_delta","delta":"A close-up product photo...","index":0,"created":1701234567,"model":"tencent/HunyuanImage-3.0-Instruct"}
|
||||
```
|
||||
|
||||
Example final image event:
|
||||
|
||||
```text
|
||||
data: {"object":"image.edit.chunk","type":"image","data":[{"b64_json":"<base64-encoded PNG>","url":null,"revised_prompt":null}],"output_format":"png","size":"1024x1024","created":1701234567,"model":"tencent/HunyuanImage-3.0-Instruct"}
|
||||
```
|
||||
|
||||
Terminal event:
|
||||
|
||||
```text
|
||||
data: [DONE]
|
||||
```
|
||||
|
||||
If the engine fails after the SSE response has started, the stream emits an
|
||||
error chunk followed by `[DONE]`:
|
||||
|
||||
```text
|
||||
data: {"object":"error","created":1701234567,"model":"tencent/HunyuanImage-3.0-Instruct","error":{"message":"<error message>","type":"server_error","code":500}}
|
||||
|
||||
data: [DONE]
|
||||
```
|
||||
|
||||
## Examples
|
||||
|
||||
### Multiple Images input
|
||||
@@ -183,28 +134,12 @@ curl -s -D >(grep -i x-request-id >&2) \
|
||||
-X POST "http://localhost:8000/v1/images/edits" \
|
||||
-F "model=xxx" \
|
||||
-F "image=@xx.png" \
|
||||
-F "image=@xx.png" \
|
||||
-F "image=@xx.png"
|
||||
-F "prompt='this bear is wearing sportwear. holding a basketball, and bending one leg.'" \
|
||||
-F "size=1024x1024" \
|
||||
-F "output_format=png"
|
||||
```
|
||||
|
||||
### Streaming HunyuanImage3 IT2I AR Text
|
||||
|
||||
```bash
|
||||
curl -N -X POST "http://localhost:8000/v1/images/edits" \
|
||||
-F "model=tencent/HunyuanImage-3.0-Instruct" \
|
||||
-F "image=@./input.png" \
|
||||
-F "prompt=Turn the product into a clean studio advertisement." \
|
||||
-F "size=1024x1024" \
|
||||
-F "output_format=png" \
|
||||
-F "stream=true"
|
||||
```
|
||||
|
||||
Use the `ar_delta` events for progressive display of the AR-generated
|
||||
recaption. Decode the `b64_json` field from the final `image` event to get the
|
||||
edited image.
|
||||
|
||||
|
||||
## Parameter Handling
|
||||
|
||||
|
||||
@@ -1,24 +1,28 @@
|
||||
# Diffusion Attention Backends
|
||||
|
||||
This document describes the diffusion attention backends available in vLLM-Omni, how to select them globally and per-role, the per-platform defaults, and how to use SageAttention.
|
||||
This document describes the diffusion attention backends available in vLLM-Omni, how to select them globally and per-role, and how to use SageAttention.
|
||||
|
||||
## Overview
|
||||
|
||||
Diffusion attention backend selection is resolved in `vllm_omni.diffusion.attention.selector`. It looks up the backend from a structured `AttentionConfig` carried on `OmniDiffusionConfig` and falls back to the platform default when nothing is configured.
|
||||
|
||||
This backend is used by diffusion attention layers such as the DiT attention in video and image generation models. It does **not** affect autoregressive (LLM) attention paths — those go through vLLM's own attention backend selector.
|
||||
This backend is used by diffusion attention layers such as the DiT attention in video and image generation models.
|
||||
|
||||
The full set of backends and their platform defaults is in the **Backend Options** and **Platform Defaults** sections below. If no attention backend is configured, vLLM-Omni asks the current platform to choose the default.
|
||||
On CUDA, the practical choices today are:
|
||||
|
||||
- `FLASH_ATTN`: FlashAttention backend. This is the default on supported CUDA systems when FlashAttention is installed.
|
||||
- `TORCH_SDPA`: PyTorch `scaled_dot_product_attention`.
|
||||
- `SAGE_ATTN`: SageAttention backend, if `sageattention` is installed.
|
||||
|
||||
If no attention backend is configured, vLLM-Omni asks the current platform to choose the default. On CUDA, that normally means `FLASH_ATTN` when available, otherwise `TORCH_SDPA`.
|
||||
|
||||
## Backend Options
|
||||
|
||||
| Value | Notes |
|
||||
|---|---|
|
||||
| `FLASH_ATTN` | Wraps FlashAttention 2. Default on Hopper / Ada / Ampere when `flash-attn` is installed. |
|
||||
| `CUDNN_ATTN` | Pins `sdpa_kernel([CUDNN_ATTENTION])`. Default on Blackwell (sm_10x / sm_12x) with cuDNN ≥ 9.5. Wins on mask-heavy DiTs (HunyuanVideo-1.5: 2× e2e vs SDPA). |
|
||||
| `FLASHINFER_ATTN` | Calls FlashInfer's dense `single_prefill_with_kv_cache` directly with `custom_mask` for non-causal masked attention. Used as Blackwell fallback when cuDNN is unavailable. Requires `flashinfer`. |
|
||||
| `TORCH_SDPA` | PyTorch `scaled_dot_product_attention` with the default backend dispatcher. Most conservative; always available. |
|
||||
| `SAGE_ATTN` | SageAttention 2.2 — INT8-quantized attention with FP16 accumulation. Lossy but typically visually indistinguishable on diffusion outputs. Requires `sageattention`. |
|
||||
| `FLASH_ATTN` | Default on CUDA when FlashAttention is available. Good default for most diffusion workloads. |
|
||||
| `TORCH_SDPA` | Most conservative fallback. Useful for debugging or compatibility. |
|
||||
| `SAGE_ATTN` | Requires `sageattention`. Can improve performance on some workloads, but output quality must be validated model-by-model. |
|
||||
|
||||
## Configuration
|
||||
|
||||
@@ -91,80 +95,6 @@ config = OmniDiffusionConfig(
|
||||
|
||||
A plain dict is also accepted and normalized to `AttentionConfig`.
|
||||
|
||||
## Platform Defaults
|
||||
|
||||
### Blackwell (sm_100 / sm_103 / sm_120 / sm_121)
|
||||
|
||||
Auto-route preference, in order:
|
||||
|
||||
1. `CUDNN_ATTN` — when cuDNN ≥ 9.5 is available (ships in PyTorch 2.5+ wheels)
|
||||
2. `FLASHINFER_ATTN` — when `flashinfer` is installed but cuDNN < 9.5
|
||||
3. `FLASH_ATTN` — when `flash-attn` is installed with the Blackwell CUTE kernel
|
||||
4. `TORCH_SDPA` — last resort
|
||||
|
||||
The startup log line `Defaulting to diffusion attention backend CUDNN_ATTN (Blackwell sm_120, cuDNN 91002)` confirms the route.
|
||||
|
||||
**Why CUDNN_ATTN by default**: on mask-heavy diffusion models (HunyuanVideo-1.5, Qwen-Image), cuDNN's pinned FMHA kernel sidesteps a PyTorch SDPA dispatch quirk where the unpinned dispatcher picks `EFFICIENT_ATTENTION` (~25 ms) for masked calls instead of cuDNN (~11 ms). The pin gives 2× e2e on HV-1.5 with no regression on lighter models.
|
||||
|
||||
### Hopper (sm_90) / Ada (sm_89) / Ampere (sm_80–sm_86)
|
||||
|
||||
Auto-route preference:
|
||||
|
||||
1. `FLASH_ATTN` — when `flash-attn` is installed
|
||||
2. `TORCH_SDPA` — fallback
|
||||
|
||||
`CUDNN_ATTN` and `FLASHINFER_ATTN` are still selectable via env var on these GPUs but are not in the auto-route — FlashAttention 2 is the well-tuned path on pre-Blackwell hardware.
|
||||
|
||||
## End-to-End Benchmark (BF16, sm_120 RTX Pro 6000 Blackwell)
|
||||
|
||||
Same prompt and seed across runs. `Total generation time` from `text_to_video.py` / `text_to_image.py`.
|
||||
|
||||
| Model | Shape | TORCH_SDPA | CUDNN_ATTN | FLASHINFER_ATTN |
|
||||
|---|---|---|---|---|
|
||||
| HunyuanVideo-1.5 (T2V) | 480p / 33f / 50 steps | 147.05 s | **73.02 s** | 127.84 s |
|
||||
| Wan 2.2 14B (T2V) | 480p / 33f / 40 steps | 117.75 s | 117.17 s | **115.07 s** |
|
||||
| Qwen-Image (T2I) | 1024² / 50 steps | 17.41 s | **15.14 s** | 16.02 s |
|
||||
| FLUX.2-dev (T2I) | 1024² / 50 steps, TP=2 | 53.62 s | **53.30 s** | 54.94 s |
|
||||
|
||||
Pattern: mask-heavy DiTs (HV-1.5, Qwen-Image) favor `CUDNN_ATTN`; lighter-mask DiTs and TP-saturated configs (Wan 2.2, FLUX.2 TP=2) tie within noise.
|
||||
|
||||
## Known Limitations
|
||||
|
||||
### LTX-2.0: `CUDNN_ATTN` crashes under torch.compile
|
||||
|
||||
LTX-2's audio attention has a symbolic head_dim under torch.compile tracing. cuDNN's SDPA backend selector rejects symbolic dims and Dynamo aborts compilation. Tracked in [#3121](https://github.com/vllm-project/vllm-omni/issues/3121).
|
||||
|
||||
**Workaround**: explicitly select `FLASHINFER_ATTN` or `TORCH_SDPA` for LTX-2.0:
|
||||
|
||||
```bash
|
||||
DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
--model Lightricks/LTX-2 ...
|
||||
```
|
||||
|
||||
### FA4 not yet integrated
|
||||
|
||||
FlashAttention-4 (released March 2026) targets Blackwell natively and reportedly beats cuDNN by ~20% on B200. As of this writing the `flash-attn-4 4.0.0b10` wheel crashes with `AttributeError: 'NoneType' object has no attribute '_trait'` during JIT on sm_120. Not yet wired into vLLM-Omni; revisit when stable lands.
|
||||
|
||||
## Choosing a Backend Manually
|
||||
|
||||
### When to override the default
|
||||
|
||||
- **Quality validation**: compare a new backend against `TORCH_SDPA` as the reference, since SDPA's default dispatcher is the most extensively tested.
|
||||
- **Lossy speedup hunting**: try `SAGE_ATTN` (INT8 quantized) on diffusion outputs — typically indistinguishable visually but always validate.
|
||||
- **Workaround for known issues**: see Known Limitations above.
|
||||
|
||||
### Verifying which backend is in use
|
||||
|
||||
The startup log prints one of:
|
||||
|
||||
```
|
||||
Using diffusion attention backend 'CUDNN_ATTN' # explicit override
|
||||
Defaulting to diffusion attention backend CUDNN_ATTN ... # auto-route
|
||||
Defaulting to diffusion attention backend SDPA # nothing else available
|
||||
```
|
||||
|
||||
If you don't see one of these, the model didn't reach diffusion stage init — check earlier logs for failures.
|
||||
|
||||
## SageAttention Installation
|
||||
|
||||
vLLM-Omni expects SageAttention to be installed into the same Python environment as vLLM-Omni.
|
||||
@@ -185,33 +115,11 @@ Quick check:
|
||||
python -c "import sageattention; print(sageattention.__file__)"
|
||||
```
|
||||
|
||||
## Usage Examples
|
||||
## Usage
|
||||
|
||||
### Default (auto-route)
|
||||
### Enable SageAttention
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
--model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \
|
||||
--prompt "A dog running across a field of golden wheat." \
|
||||
--height 480 --width 832 --num-frames 33 \
|
||||
--num-inference-steps 50 --seed 42 --guidance-scale 6.0 \
|
||||
--output hv15.mp4
|
||||
```
|
||||
|
||||
On Blackwell this picks `CUDNN_ATTN` automatically. Check the log for the `Defaulting to ...` line.
|
||||
|
||||
### Explicit backend selection
|
||||
|
||||
```bash
|
||||
DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
--model Lightricks/LTX-2 \
|
||||
--prompt "A dog running across a field of golden wheat." \
|
||||
--height 480 --width 832 --num-frames 33 \
|
||||
--num-inference-steps 40 --seed 42 --guidance-scale 4.0 \
|
||||
--output ltx2.mp4
|
||||
```
|
||||
|
||||
### SageAttention (lossy)
|
||||
Example: HunyuanVideo-1.5 text-to-video
|
||||
|
||||
```bash
|
||||
DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
@@ -220,7 +128,7 @@ DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN python examples/offline_inference/text_to_
|
||||
--height 480 --width 832 --num-frames 33 \
|
||||
--num-inference-steps 30 --seed 42 --guidance-scale 6.0 \
|
||||
--tensor-parallel-size 2 \
|
||||
--output hv15_sage.mp4
|
||||
--output ../tmp/hv15_modelopt_sage.mp4
|
||||
```
|
||||
|
||||
Example: Wan2.2 TI2V 5B
|
||||
@@ -265,13 +173,13 @@ python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
|
||||
## Validation Guidance
|
||||
|
||||
Don't assume a faster attention backend is numerically interchangeable with `TORCH_SDPA`.
|
||||
Do not assume that a faster attention backend is numerically interchangeable with `FLASH_ATTN`.
|
||||
|
||||
Always compare:
|
||||
|
||||
- End-to-end runtime
|
||||
- Diffusion-stage runtime (`add_req_and_wait` line in DiffusionEngine.step breakdown)
|
||||
- Output quality against a known-good baseline (CLIP similarity, frame-level diff, or visual review)
|
||||
- DiT / diffusion stage runtime
|
||||
- Output quality against a known-good baseline
|
||||
|
||||
At minimum, keep the same:
|
||||
|
||||
@@ -279,19 +187,6 @@ At minimum, keep the same:
|
||||
- prompt
|
||||
- seed
|
||||
- resolution
|
||||
- frame count / step count
|
||||
- parallel config (TP / CFG-parallel / Ulysses degrees)
|
||||
|
||||
## Reproducing the Benchmark Table
|
||||
|
||||
The end-to-end numbers above were collected by running `text_to_video.py` /
|
||||
`text_to_image.py` with the same prompt and seed while varying
|
||||
`DIFFUSION_ATTENTION_BACKEND`. For a quick kernel-level comparison of the
|
||||
backends without loading a model:
|
||||
|
||||
```bash
|
||||
python benchmarks/diffusion/bench_attention_backends.py --preset hv15
|
||||
```
|
||||
|
||||
It runs all three BF16 backends on representative DiT attention shapes and
|
||||
prints a ranking table at the end.
|
||||
- frame count
|
||||
- inference steps
|
||||
- parallel config
|
||||
@@ -113,7 +113,7 @@ The following tables show which models support each feature:
|
||||
| **FLUX.1-schnell** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ |
|
||||
| **FLUX.2-klein** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ |
|
||||
| **FLUX.1-Kontext-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **FLUX.2-dev** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **FLUX.2-dev** | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **GLM-Image** | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **Hidream-I1-Full** | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **HunyuanImage3** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ |
|
||||
@@ -145,7 +145,7 @@ The following tables show which models support each feature:
|
||||
| **Wan2.2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (encode/decode) | ❌ | ❌ |
|
||||
| **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **LTX-2.3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **LTX-2.3** | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ |
|
||||
| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||
|
||||
@@ -23,7 +23,6 @@ list of supported architectures across all modalities, see
|
||||
| OmniVoice | `k2-fsa/OmniVoice` | 2 (gen + dec) | ✓ | — | voice design (`--instruct`), language (`--lang`) | 24 kHz |
|
||||
| Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | 2 (talker + code2wav) | ✓ (Base) | ✓ | 3 task variants (`--query-type`) | 24 kHz |
|
||||
| Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | varies | ✓ | ✓ | voice presets (`--voice`) | 24 kHz |
|
||||
| higgs-audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` (+ `k2-fsa/OmniVoice/audio_tokenizer/` codec) | 2 (talker + code2wav, DualFFN) | ✓ (`--ref-audio` + `--ref-text`) | — | — | 24 kHz |
|
||||
|
||||
## Common Quick Start
|
||||
|
||||
@@ -334,46 +333,6 @@ Available voice presets are listed on the HF model card (`mistralai/Voxtral-4B-T
|
||||
- `--concurrency M` requires `--streaming` and must evenly divide `--num-prompts`.
|
||||
- Run `--help` for the full argument surface.
|
||||
|
||||
---
|
||||
|
||||
## higgs-audio v2
|
||||
|
||||
2-stage TTS at 24 kHz: a vLLM-native Llama-3.2-3B talker with a DualFFN audio expert (Stage 0) feeding a HiggsAudio codec decoder (Stage 1) over the shared-memory connector. Stage 1 builds on the HiggsAudio decoder kernel at `vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_decoder.py`.
|
||||
|
||||
### Prerequisites
|
||||
|
||||
Voice clone uses HF's `HiggsAudioV2TokenizerModel`, instantiated from `k2-fsa/OmniVoice/audio_tokenizer/` — the boson-ai standalone tokenizer Hub repo's `model.safetensors` is actually the 3B talker LM, so we point HF at k2's repackaged codec weights instead. Only the `audio_tokenizer/` subdir (~806 MB) is downloaded.
|
||||
|
||||
```bash
|
||||
pip install -U "transformers>=5.3.0"
|
||||
```
|
||||
|
||||
### Quick start (plain TTS)
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py \
|
||||
--texts "Hello world." "The quick brown fox jumps over the lazy dog." \
|
||||
--output-dir results/wavs
|
||||
```
|
||||
|
||||
### Voice cloning
|
||||
|
||||
Pass both `--ref-audio` and `--ref-text` together:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py \
|
||||
--texts "Hello, this is a cloned voice." \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--ref-text "Exact transcript spoken in reference.wav." \
|
||||
--output-dir results/wavs
|
||||
```
|
||||
|
||||
### Notes
|
||||
|
||||
- Deploy config auto-loads from `vllm_omni/deploy/higgs_audio_v2.yaml`.
|
||||
- `--ref-text` must be the real transcript of `--ref-audio`; mismatched text degrades cloned-voice quality.
|
||||
- Out of scope (rejected with 4xx by the request validator): multi-speaker `[SPEAKERn]` dialogue, `profile:` text-only speaker descriptions, the `ref_audio_in_system_message` system-block variant, chunked long-form generation, and per-request `voice` / `instructions` / `task_type` / `language` / `speed != 1.0` / `x_vector_only_mode` / `speaker_embedding`.
|
||||
|
||||
## Example materials
|
||||
|
||||
??? abstract "cosyvoice3/end2end.py"
|
||||
@@ -384,9 +343,6 @@ python examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py \
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/text_to_speech/fish_speech/end2end.py"
|
||||
``````
|
||||
??? abstract "higgs_audio_v2/end2end.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py"
|
||||
??? abstract "glm_tts/end2end.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/text_to_speech/glm_tts/end2end.py"
|
||||
|
||||
@@ -19,7 +19,6 @@ For the full list of supported architectures across all modalities, see
|
||||
|---|---|---|---|---|---|
|
||||
| CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | — |
|
||||
| Fish Speech S2 Pro | `fishaudio/s2-pro` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | ✓ |
|
||||
| higgs-audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` | ✓ (`ref_audio`+`ref_text`) | ✓ (codec_streaming) | — | — |
|
||||
| GLM-TTS | `zai-org/GLM-TTS` | ✓ (`ref_audio`+`ref_text`, required) | ✓ (PCM stream) | — | ✓ |
|
||||
| OmniVoice | `k2-fsa/OmniVoice` | (offline only) | — | — | — |
|
||||
| Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | ✓ (Base) | ✓ (PCM + WebSocket) | ✓ (presets + `/v1/audio/voices` upload) | ✓ (standard + FastRTC) |
|
||||
@@ -194,45 +193,6 @@ python fish_speech/gradio_demo.py --api-base http://localhost:8091 # if server
|
||||
|
||||
---
|
||||
|
||||
## higgs-audio v2
|
||||
|
||||
2-stage TTS at 24 kHz: a vLLM-native Llama-3.2-3B talker with a DualFFN audio expert (Stage 0) feeding a HiggsAudio codec decoder (Stage 1) that streams chunks back to the client.
|
||||
|
||||
### Prerequisites
|
||||
|
||||
Voice clone uses HF's `HiggsAudioV2TokenizerModel` loaded from `k2-fsa/OmniVoice/audio_tokenizer/` (~806 MB subdir; the boson-ai standalone tokenizer Hub repo's `model.safetensors` is the 3B talker LM, not the codec):
|
||||
|
||||
```bash
|
||||
pip install -U "transformers>=5.3.0"
|
||||
```
|
||||
|
||||
### Launch
|
||||
```bash
|
||||
GPUS=6,7 PORT=8094 bash examples/online_serving/text_to_speech/higgs_audio_v2/run_server.sh
|
||||
```
|
||||
Deploy config auto-loads from `vllm_omni/deploy/higgs_audio_v2.yaml`.
|
||||
|
||||
### Sending requests
|
||||
```bash
|
||||
# Plain TTS
|
||||
python higgs_audio_v2/batch_speech_client.py \
|
||||
--base-url http://localhost:8094 \
|
||||
--output-dir /tmp/higgs_out \
|
||||
--prompts "Hello world." "The quick brown fox jumps over the lazy dog."
|
||||
|
||||
# Voice cloning — pass a reference clip and its transcript together
|
||||
python higgs_audio_v2/batch_speech_client.py \
|
||||
--base-url http://localhost:8094 \
|
||||
--output-dir /tmp/higgs_clone \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--ref-text "Exact transcript spoken in reference.wav." \
|
||||
--prompts "Hello, this is a cloned voice."
|
||||
```
|
||||
|
||||
### Notes
|
||||
- Output: 24 kHz mono.
|
||||
- `--ref-text` must be the real transcript of `--ref-audio`; mismatched text degrades cloned-voice quality.
|
||||
- Out of scope (rejected with explicit 4xx): multi-speaker `[SPEAKERn]` tags inside `input`, `profile:` text-only speaker descriptions, the `ref_audio_in_system_message` system-block variant, chunked long-form generation, and per-request `voice` / `instructions` / `task_type` / `language` / `speed != 1.0` / `x_vector_only_mode` / `speaker_embedding`.
|
||||
## GLM-TTS
|
||||
|
||||
2-stage TTS (AR + DiT flow-matching) at 24 kHz. Every request requires `ref_audio` + `ref_text`.
|
||||
@@ -499,13 +459,6 @@ The demo handles voice-preset selection and reference-audio upload. `voxtral_tts
|
||||
``````py
|
||||
--8<-- "examples/online_serving/text_to_speech/fish_speech/speech_client.py"
|
||||
``````
|
||||
??? abstract "higgs_audio_v2/batch_speech_client.py"
|
||||
``````py
|
||||
--8<-- "examples/online_serving/text_to_speech/higgs_audio_v2/batch_speech_client.py"
|
||||
``````
|
||||
??? abstract "higgs_audio_v2/run_server.sh"
|
||||
``````sh
|
||||
--8<-- "examples/online_serving/text_to_speech/higgs_audio_v2/run_server.sh"
|
||||
??? abstract "glm_tts/gradio_demo.py"
|
||||
``````py
|
||||
--8<-- "examples/online_serving/text_to_speech/glm_tts/gradio_demo.py"
|
||||
|
||||
@@ -52,8 +52,8 @@ vLLM-Omni's runtime module names.
|
||||
|
||||
| Model | Scope | Status | Notes |
|
||||
|-------|-------|--------|-------|
|
||||
| GLM-Image | Diffusion transformer | ✅ | `Intel/GLM-Image-int4-AutoRound` |
|
||||
| BAGEL | Checkpoint-defined diffusion or transformer stage | Not validated | Requires a compatible AutoRound checkpoint |
|
||||
| GLM-Image | Checkpoint-defined diffusion or transformer stage | Not validated | Requires a compatible AutoRound checkpoint |
|
||||
|
||||
## Configuration
|
||||
|
||||
|
||||
@@ -7,28 +7,26 @@ using the OCP MX format: groups of 32 K-dimension elements share a single
|
||||
`float8_e8m0fnu` exponent scale. This gives better accuracy than channel-wise
|
||||
FP8 while keeping the same 8-bit weight footprint.
|
||||
|
||||
This method supports three modes:
|
||||
This method supports two modes:
|
||||
|
||||
| Mode | Description |
|
||||
|------|-------------|
|
||||
| **Online** | BF16 weights are quantized to MXFP8 at load time — no pre-processing needed |
|
||||
| **Offline (Native)** | msModelSlim-exported MXFP8 weights converted to diffusers format via `merge_mxfp8_checkpoint.py` — weights and scales are loaded directly from the preprocessed checkpoint |
|
||||
| **Offline (AutoRound)** | AutoRound MXFP8 checkpoints with `data_type="mx_fp"` — auto-detected from `config.json` |
|
||||
| **Offline** | msModelSlim-exported MXFP8 weights converted to diffusers format via `merge_mxfp8_checkpoint.py` — weights and scales are loaded directly from the preprocessed checkpoint |
|
||||
|
||||
## Hardware Support
|
||||
|
||||
| Device | Online | Offline (Native) | Offline (AutoRound) |
|
||||
|--------|--------|------------------|---------------------|
|
||||
| NVIDIA Blackwell GPU (SM 100+) | ⭕ | ⭕ | ⭕ |
|
||||
| NVIDIA Ada/Hopper GPU (SM 89+) | ⭕ | ⭕ | ⭕ |
|
||||
| NVIDIA Ampere GPU (SM 80+) | ⭕ | ⭕ | ⭕ |
|
||||
| AMD ROCm | ⭕ | ⭕ | ⭕ |
|
||||
| Intel XPU | ✅ | ❌ | ✅ |
|
||||
| Ascend NPU (Atlas 950 A5) | ✅ | ✅ | ⭕ |
|
||||
| Device | Support |
|
||||
|--------|---------|
|
||||
| NVIDIA Blackwell GPU (SM 100+) | ⭕ |
|
||||
| NVIDIA Ada/Hopper GPU (SM 89+) | ⭕ |
|
||||
| NVIDIA Ampere GPU (SM 80+) | ⭕ |
|
||||
| AMD ROCm | ⭕ |
|
||||
| Intel XPU | ⭕ |
|
||||
| Ascend NPU (Atlas 950 A5) | ✅ |
|
||||
|
||||
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this guide.
|
||||
|
||||
**Note**: Intel XPU only supports AutoRound MXFP8 for offline mode. Use AutoRound quantized checkpoints or online mode for XPU.
|
||||
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
|
||||
guide.
|
||||
|
||||
## Model Type Support
|
||||
|
||||
@@ -84,9 +82,9 @@ python text_to_video.py --model <your-model> --quantization mxfp8
|
||||
vllm serve <your-model> --omni --quantization mxfp8
|
||||
```
|
||||
|
||||
### Offline Mode (Native)
|
||||
### Offline Mode
|
||||
|
||||
Native offline mode loads a pre-quantized checkpoint from msModelSlim. A preprocessing
|
||||
Offline mode loads a pre-quantized checkpoint from msModelSlim. A preprocessing
|
||||
step converts the raw quantized output to the diffusers format expected by
|
||||
vLLM-Omni and injects the quantization config into `transformer/config.json` so
|
||||
that vLLM-Omni auto-detects the offline path without a `--quantization` flag.
|
||||
@@ -161,41 +159,10 @@ omni = Omni(model="/path/to/Wan2.2-TI2V-5B-MXFP8")
|
||||
```
|
||||
|
||||
!!! note
|
||||
No `--quantization` flag is needed for native offline mode. The preprocessing
|
||||
No `--quantization` flag is needed for offline mode. The preprocessing
|
||||
script injects `quantization_config` into each `transformer/config.json`,
|
||||
which vLLM-Omni reads automatically to activate the offline MXFP8 method.
|
||||
|
||||
### Offline Mode (AutoRound)
|
||||
|
||||
AutoRound MXFP8 checkpoints declare `quant_method="auto-round"` with `data_type="mx_fp"` in their `config.json`. These are automatically detected and use the `IncMxfp8OfflineLinearMethod` backend.
|
||||
|
||||
To use an AutoRound MXFP8 checkpoint:
|
||||
|
||||
```bash
|
||||
python text_to_video.py --model <autoround-mxfp8-model>
|
||||
|
||||
# Online serving
|
||||
vllm serve <autoround-mxfp8-model> --omni
|
||||
```
|
||||
|
||||
Python API:
|
||||
|
||||
```python
|
||||
omni = Omni(model="<autoround-mxfp8-model>")
|
||||
```
|
||||
|
||||
!!! note
|
||||
AutoRound MXFP8 checkpoints are auto-detected from `config.json` and do not require a `--quantization` flag. The config must include:
|
||||
```json
|
||||
{
|
||||
"quantization_config": {
|
||||
"quant_method": "auto-round",
|
||||
"data_type": "mx_fp",
|
||||
...
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
## Parameters
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|
||||
@@ -21,7 +21,7 @@ by msModelSlim and the merge tools.
|
||||
| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ⭕ | ⭕ |
|
||||
| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ⭕ | ⭕ |
|
||||
| AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ |
|
||||
| Intel XPU | ⭕ | ⭕ | ✅ | ⭕ |
|
||||
| Intel XPU | ⭕ | ⭕ | ⭕ | ⭕ |
|
||||
| Ascend NPU | ❌ | ✅ | ✅ | ✅ |
|
||||
|
||||
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
|
||||
@@ -36,7 +36,7 @@ MXFP4 are documented for the Ascend NPU path.
|
||||
|--------|-------|----------------|--------|
|
||||
| FP8 W8A8 | [FP8](fp8.md) | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image family and other DiT models |
|
||||
| Int8 W8A8 | [Int8](int8.md) | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image and Z-Image |
|
||||
| MXFP8 W8A8 | [MXFP8](mxfp8.md) | Wan2.2-T2V-A14B, Wan2.2-I2V-A14B, Wan2.2-TI2V-5B | Validated on Ascend NPU and Intel XPU |
|
||||
| MXFP8 W8A8 | [MXFP8](mxfp8.md) | Wan2.2-T2V-A14B, Wan2.2-I2V-A14B, Wan2.2-TI2V-5B | Ascend NPU only |
|
||||
| MXFP4 W4A4 | [MXFP4](mxfp4.md) | Wan2.2-T2V-A14B, Wan2.2-I2V-A14B | Ascend NPU only; TI2V-5B is not supported |
|
||||
|
||||
### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
|
||||
|
||||
@@ -17,7 +17,7 @@ from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniPromptType
|
||||
|
||||
_REPO_ROOT = Path(__file__).resolve().parents[3]
|
||||
_DEFAULT_DEPLOY_CONFIG = str(_REPO_ROOT / "vllm_omni" / "deploy" / "hunyuan_image_3_moe.yaml")
|
||||
_DEFAULT_DEPLOY_CONFIG = str(_REPO_ROOT / "vllm_omni" / "deploy" / "hunyuan_image3.yaml")
|
||||
_DEFAULT_AR_DEPLOY_CONFIG = str(_REPO_ROOT / "vllm_omni" / "deploy" / "hunyuan_image3_ar.yaml")
|
||||
|
||||
_MODALITY_TASK_MAP: dict[str, tuple[str, str | None]] = {
|
||||
|
||||
@@ -127,12 +127,7 @@ def run_one_path(
|
||||
),
|
||||
args.device,
|
||||
)
|
||||
pred = run_pipeline_forward(
|
||||
pipeline,
|
||||
batch_inputs,
|
||||
noise,
|
||||
request_id=f"internvla-a1-sample-{index}",
|
||||
)
|
||||
pred = run_pipeline_forward(pipeline, batch_inputs, noise)
|
||||
pred = pred[:, :, : dataset.physical_action_dim].to(torch.float32).cpu()
|
||||
results.append(
|
||||
{
|
||||
@@ -155,7 +150,6 @@ def run_pipeline_forward(
|
||||
pipeline,
|
||||
batch_inputs: dict[str, torch.Tensor],
|
||||
noise: torch.Tensor,
|
||||
request_id: str,
|
||||
) -> torch.Tensor:
|
||||
output = pipeline.forward(
|
||||
OmniDiffusionRequest(
|
||||
@@ -167,7 +161,6 @@ def run_pipeline_forward(
|
||||
"decode_image": False,
|
||||
}
|
||||
),
|
||||
request_id=request_id,
|
||||
)
|
||||
)
|
||||
if output.error:
|
||||
@@ -219,38 +212,23 @@ def benchmark_forward(
|
||||
|
||||
_synchronize(args.device)
|
||||
cold_start_begin = time.perf_counter()
|
||||
pred = run_pipeline_forward(
|
||||
pipeline,
|
||||
batch_inputs,
|
||||
noise,
|
||||
request_id=f"internvla-a1-benchmark-{index}-cold",
|
||||
)
|
||||
pred = run_pipeline_forward(pipeline, batch_inputs, noise)
|
||||
_synchronize(args.device)
|
||||
cold_start_ms = (time.perf_counter() - cold_start_begin) * 1000.0
|
||||
|
||||
warmup_ms: list[float] = []
|
||||
for iter_idx in range(args.warmup_iters):
|
||||
for _ in range(args.warmup_iters):
|
||||
_synchronize(args.device)
|
||||
begin = time.perf_counter()
|
||||
_ = run_pipeline_forward(
|
||||
pipeline,
|
||||
batch_inputs,
|
||||
noise,
|
||||
request_id=f"internvla-a1-benchmark-{index}-warmup-{iter_idx}",
|
||||
)
|
||||
_ = run_pipeline_forward(pipeline, batch_inputs, noise)
|
||||
_synchronize(args.device)
|
||||
warmup_ms.append((time.perf_counter() - begin) * 1000.0)
|
||||
|
||||
benchmark_ms: list[float] = []
|
||||
for iter_idx in range(args.benchmark_iters):
|
||||
for _ in range(args.benchmark_iters):
|
||||
_synchronize(args.device)
|
||||
begin = time.perf_counter()
|
||||
_ = run_pipeline_forward(
|
||||
pipeline,
|
||||
batch_inputs,
|
||||
noise,
|
||||
request_id=f"internvla-a1-benchmark-{index}-iter-{iter_idx}",
|
||||
)
|
||||
_ = run_pipeline_forward(pipeline, batch_inputs, noise)
|
||||
_synchronize(args.device)
|
||||
benchmark_ms.append((time.perf_counter() - begin) * 1000.0)
|
||||
|
||||
@@ -317,12 +295,7 @@ def main() -> None:
|
||||
dataset=dataset,
|
||||
train_meta=train_meta,
|
||||
collate_samples=collate_open_loop_samples,
|
||||
run_sample_actions=lambda policy, batch_inputs, noise: run_pipeline_forward(
|
||||
policy,
|
||||
batch_inputs,
|
||||
noise,
|
||||
request_id="internvla-a1-open-loop",
|
||||
),
|
||||
run_sample_actions=lambda policy, batch_inputs, noise: run_pipeline_forward(policy, batch_inputs, noise),
|
||||
num_episodes=args.num_episodes,
|
||||
seed=args.seed,
|
||||
device=args.device,
|
||||
|
||||
@@ -2,13 +2,12 @@
|
||||
|
||||
[Ming-flash-omni-2.0](https://github.com/inclusionAI/Ming) is an omni-modal model supporting text, image, video, and audio understanding, with text and speech outputs.
|
||||
|
||||
vLLM-Omni supports three deployment modes:
|
||||
vLLM-Omni supports two deployment modes:
|
||||
|
||||
| Mode | Deploy config | Output |
|
||||
|------|--------------|--------|
|
||||
| Thinker + Talker (omni-speech, default) | `vllm_omni/deploy/ming_flash_omni.yaml` | Text + Audio |
|
||||
| Thinker only (multimodal understanding) | `vllm_omni/deploy/ming_flash_omni_thinker_only.yaml` | Text |
|
||||
| Thinker + Imagegen (text-to-image / img2img) | `vllm_omni/model_executor/stage_configs/ming_flash_omni_dual.yaml` | Image (online-serving only at the moment) |
|
||||
|
||||
For standalone TTS (talker only), see the [Ming-flash-omni-TTS section in the Text-To-Speech hub](../text_to_speech/README.md#ming-flash-omni-tts).
|
||||
|
||||
@@ -88,17 +87,6 @@ The default deploy YAML allocates thinker on GPUs 0–3 and talker on GPU 3 for
|
||||
|
||||
Pass `--deploy-config /path/to/your_deploy.yaml` to any of the commands above to override the bundled deploy config.
|
||||
|
||||
### Image generation (text-to-image / img2img)
|
||||
|
||||
The diffusion-stage image-generation path is currently only wired through
|
||||
the online OpenAI-compatible chat endpoint (`/v1/chat/completions` with
|
||||
`"modalities": ["image"]`). For payloads, optional knobs
|
||||
(`steps`/`cfg`/`seed`/`byte5_text`/`negative_prompt`), and the img2img
|
||||
reference-image flow, see the
|
||||
[image-gen section in the online-serving README](../../online_serving/ming_flash_omni/README.md#image-generation-text-to-image).
|
||||
`end2end.py` does not yet exercise the imagegen stage; that is tracked as
|
||||
follow-up work.
|
||||
|
||||
## Online serving
|
||||
|
||||
For online serving via the OpenAI-compatible API, see [examples/online_serving/ming_flash_omni/README.md](../../online_serving/ming_flash_omni/README.md).
|
||||
@@ -1,38 +0,0 @@
|
||||
# higgs-audio v2 — offline example
|
||||
|
||||
Drives Stage 0 (DualFFN talker) + Stage 1 (HiggsAudio codec) for `bosonai/higgs-audio-v2-generation-3B-base` end-to-end through the vLLM-Omni engine and writes a 24 kHz mono WAV per prompt.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
Voice clone needs `transformers>=5.3.0` — vllm-omni loads the audio codec via HF's `HiggsAudioV2TokenizerModel`, instantiated from the `k2-fsa/OmniVoice/audio_tokenizer/` subdirectory (only that ~806 MB subdir is downloaded). The boson-ai standalone tokenizer repo's `model.safetensors` is actually a copy of the 3B talker LM, so HF can't load it directly; the k2 bundle ships the same codec weights repackaged with HF-compatible key naming.
|
||||
|
||||
```bash
|
||||
pip install -U "transformers>=5.3.0"
|
||||
```
|
||||
|
||||
## Quick start
|
||||
|
||||
Plain TTS:
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py \
|
||||
--texts "Hello world." "The quick brown fox jumps over the lazy dog." \
|
||||
--output-dir results/wavs
|
||||
```
|
||||
|
||||
## Voice cloning
|
||||
|
||||
Pass both `--ref-audio` and `--ref-text` together:
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py \
|
||||
--texts "Hello, this is a cloned voice." \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--ref-text "Exact transcript spoken in reference.wav." \
|
||||
--output-dir results/wavs
|
||||
```
|
||||
|
||||
## Notes
|
||||
|
||||
- Output: 24 kHz mono WAV.
|
||||
- Deploy config: `vllm_omni/deploy/higgs_audio_v2.yaml` (auto-loaded by HF `model_type`).
|
||||
- `--ref-text` must be the real transcript of `--ref-audio`; mismatched text degrades cloned-voice quality.
|
||||
- For online serving, see [`examples/online_serving/text_to_speech/higgs_audio_v2/`](../../../online_serving/text_to_speech/higgs_audio_v2/).
|
||||
@@ -1,205 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Offline higgs-audio v2 inference example.
|
||||
|
||||
Runs Stage 0 (DualFFN talker) + Stage 1 (HiggsAudio codec) end-to-end
|
||||
through the vLLM-Omni engine without going through the HTTP server, and
|
||||
saves a 24 kHz mono WAV per prompt.
|
||||
|
||||
Example:
|
||||
|
||||
python examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py \\
|
||||
--texts "Hello world." \\
|
||||
"The quick brown fox jumps over the lazy dog." \\
|
||||
--output-dir results/wavs
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
# DeepGEMM FP8 kernels require an optional backend that may not be installed.
|
||||
# Disable the warmup before importing vLLM so engine startup falls back to the
|
||||
# regular gemm path. Users with deep_gemm installed can override these.
|
||||
os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0")
|
||||
os.environ.setdefault("VLLM_MOE_USE_DEEP_GEMM", "0")
|
||||
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
import torch
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||
|
||||
SAMPLE_RATE = 24_000
|
||||
DEFAULT_TEXTS = (
|
||||
"Hello world.",
|
||||
"The quick brown fox jumps over the lazy dog.",
|
||||
)
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = FlexibleArgumentParser(description="Offline higgs-audio v2 inference")
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
default="bosonai/higgs-audio-v2-generation-3B-base",
|
||||
help="Stage-0 talker model id or path.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--texts",
|
||||
type=str,
|
||||
nargs="+",
|
||||
default=list(DEFAULT_TEXTS),
|
||||
help="One or more plain-text prompts to synthesize.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-dir",
|
||||
type=str,
|
||||
default="results/wavs",
|
||||
help="Directory to write per-prompt WAV files.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--deploy-config",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Override the deploy config path. Auto-loads "
|
||||
"vllm_omni/deploy/higgs_audio_v2.yaml from the HF model_type by default.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-new-tokens",
|
||||
type=int,
|
||||
default=500,
|
||||
help="Cap on Stage-0 codec frames per request.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ref-audio",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Reference clip for voice clone (path to a WAV file). Paired with --ref-text.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ref-text",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Transcript of the reference clip. Required when --ref-audio is set.",
|
||||
)
|
||||
nullify_stage_engine_defaults(parser)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def _slugify(text: str) -> str:
|
||||
import re
|
||||
|
||||
slug = re.sub(r"\s+", "_", text.strip().lower())
|
||||
slug = re.sub(r"[^a-z0-9_]+", "", slug)
|
||||
return slug[:48] or "prompt"
|
||||
|
||||
|
||||
def _extract_pcm(multimodal_output: dict) -> torch.Tensor:
|
||||
"""Pull the final concatenated PCM tensor out of a request's multimodal_output."""
|
||||
audio = multimodal_output.get("model_outputs")
|
||||
if audio is None:
|
||||
audio = multimodal_output.get("audio")
|
||||
if audio is None:
|
||||
raise ValueError(f"no audio key in multimodal_output: {list(multimodal_output.keys())}")
|
||||
if isinstance(audio, list):
|
||||
valid = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio if a is not None]
|
||||
if not valid:
|
||||
raise ValueError("audio list is empty")
|
||||
return torch.cat(valid, dim=0) if len(valid) > 1 else valid[0]
|
||||
return torch.as_tensor(audio).float().cpu().reshape(-1)
|
||||
|
||||
|
||||
def _pcm_to_int16(pcm: torch.Tensor) -> np.ndarray:
|
||||
arr = pcm.numpy()
|
||||
if arr.dtype.kind == "f":
|
||||
arr = np.clip(arr, -1.0, 1.0)
|
||||
arr = (arr * 32767.0).astype(np.int16)
|
||||
else:
|
||||
arr = arr.astype(np.int16)
|
||||
return arr
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
if (args.ref_audio is None) != (args.ref_text is None):
|
||||
raise SystemExit("--ref-audio and --ref-text must be supplied together")
|
||||
output_dir = Path(args.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
engine = Omni(model=args.model, deploy_config=args.deploy_config)
|
||||
|
||||
# Build prompt_token_ids using the same path serving_speech.py takes online.
|
||||
from transformers import AutoProcessor
|
||||
|
||||
from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_v2_tokenizer import (
|
||||
build_plain_text_prompt,
|
||||
build_voice_clone_prompt,
|
||||
input_ids_to_python_list,
|
||||
)
|
||||
|
||||
processor = AutoProcessor.from_pretrained(args.model, trust_remote_code=True)
|
||||
|
||||
# Voice-clone path: load the reference clip once. The HF processor will
|
||||
# encode it via the bundled HiggsAudioV2TokenizerModel each time we build
|
||||
# a prompt below. This is cheap on CPU for a few-second clip.
|
||||
ref_wav: np.ndarray | None = None
|
||||
ref_sr: int | None = None
|
||||
if args.ref_audio is not None:
|
||||
ref_wav, ref_sr = sf.read(args.ref_audio, always_2d=False)
|
||||
if ref_wav.ndim == 2:
|
||||
ref_wav = ref_wav.mean(axis=1)
|
||||
|
||||
print(f"Model : {args.model}")
|
||||
print(f"Prompts : {len(args.texts)}")
|
||||
print(f"Output dir : {output_dir}")
|
||||
print(f"Voice clone : {'yes' if ref_wav is not None else 'no'}")
|
||||
|
||||
# Run one prompt at a time. The Stage-0 talker's per-slot audio state is
|
||||
# request-scoped; submitting multiple prompts in the same engine.generate()
|
||||
# call would batch them in the AR runner and exercise a code path that is
|
||||
# not validated for this model yet.
|
||||
total_elapsed = 0.0
|
||||
total_dur = 0.0
|
||||
for text in args.texts:
|
||||
if ref_wav is not None:
|
||||
out = build_voice_clone_prompt(processor, text, ref_wav, int(ref_sr), args.ref_text)
|
||||
prompt = {
|
||||
"prompt_token_ids": out["prompt_token_ids"],
|
||||
# Bare tensors (NOT list-wrapped): the msgspec serializer in
|
||||
# vllm_omni.data_entry_keys routes torch.Tensor → tensor_data
|
||||
# and list[Tensor] → list_data (which silently strips tensors).
|
||||
"additional_information": {
|
||||
"audio_input_ids": out["audio_input_ids"],
|
||||
"audio_input_ids_mask": out["audio_input_ids_mask"],
|
||||
},
|
||||
}
|
||||
else:
|
||||
inputs = build_plain_text_prompt(processor, text)
|
||||
prompt = {"prompt_token_ids": input_ids_to_python_list(inputs)}
|
||||
t_start = time.perf_counter()
|
||||
outputs = engine.generate([prompt])
|
||||
elapsed = time.perf_counter() - t_start
|
||||
total_elapsed += elapsed
|
||||
|
||||
mm = outputs[0].outputs[0].multimodal_output
|
||||
pcm = _extract_pcm(mm)
|
||||
slug = _slugify(text)
|
||||
out_path = output_dir / f"{slug}.wav"
|
||||
sf.write(str(out_path), _pcm_to_int16(pcm), SAMPLE_RATE, format="WAV", subtype="PCM_16")
|
||||
dur = pcm.numel() / SAMPLE_RATE
|
||||
total_dur += dur
|
||||
print(f" {slug:<50} dur={dur:5.2f}s -> {out_path}")
|
||||
|
||||
rtf = total_elapsed / total_dur if total_dur > 0 else float("inf")
|
||||
print(f"Total infer : {total_elapsed:.2f}s total audio: {total_dur:.2f}s RTF: {rtf:.3f}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
|
||||
main()
|
||||
@@ -1,105 +1,24 @@
|
||||
# VACE Video Generation
|
||||
|
||||
Generate videos from text prompts, images, or video conditions using vLLM-Omni's VACE diffusion pipeline.
|
||||
[VACE](https://github.com/ali-vilab/VACE) (Video All-in-one Creation Engine) supports multiple video tasks through a single model.
|
||||
|
||||
- `vace_video_generation.py`: command-line script for multi-mode video generation with advanced options.
|
||||
| Model | Architecture | Model Weights (bf16) | HuggingFace |
|
||||
|-------|-------------|----------------------|-------------|
|
||||
| Wan2.1-VACE (1.3B) | Wan2.1 | ~10 GB | [Wan-AI/Wan2.1-VACE-1.3B-diffusers](https://huggingface.co/Wan-AI/Wan2.1-VACE-1.3B-diffusers) |
|
||||
| Wan2.1-VACE (14B) | Wan2.1 | ~38 GB | [Wan-AI/Wan2.1-VACE-14B-diffusers](https://huggingface.co/Wan-AI/Wan2.1-VACE-14B-diffusers) |
|
||||
|
||||
## Table of Contents
|
||||
|
||||
- [Overview](#overview)
|
||||
- [Quick Start](#quick-start)
|
||||
- [Key Arguments](#key-arguments)
|
||||
- [More CLI Examples](#more-cli-examples)
|
||||
|
||||
## Overview
|
||||
|
||||
[VACE](https://github.com/ali-vilab/VACE) (Video All-in-one Creation Engine) supports multiple video tasks through a single unified model, including text-to-video, image-to-video, first-last-frame interpolation, inpainting, and reference image-guided generation.
|
||||
|
||||
### Supported Models
|
||||
|
||||
| Model | Architecture | Peak VRAM (GiB) * | Model Weights (GiB) | HuggingFace |
|
||||
|-------|-------------|-------------------|---------------------|-------------|
|
||||
| Wan2.1-VACE (1.3B) | Wan2.1 | TBD | ~10 | [Wan-AI/Wan2.1-VACE-1.3B-diffusers](https://huggingface.co/Wan-AI/Wan2.1-VACE-1.3B-diffusers) |
|
||||
| Wan2.1-VACE (14B) | Wan2.1 | TBD | ~38 | [Wan-AI/Wan2.1-VACE-14B-diffusers](https://huggingface.co/Wan-AI/Wan2.1-VACE-14B-diffusers) |
|
||||
|
||||
!!! info
|
||||
*Peak VRAM: based on basic single-card usage, 480x832 resolution, 81 frames, 30 inference steps, without any acceleration/optimization features.
|
||||
|
||||
Default model: `Wan-AI/Wan2.1-VACE-14B-diffusers`
|
||||
|
||||
## Quick Start
|
||||
|
||||
### Python API
|
||||
|
||||
Text-to-video generation:
|
||||
|
||||
```python
|
||||
import torch
|
||||
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
if __name__ == "__main__":
|
||||
omni = Omni(model="Wan-AI/Wan2.1-VACE-1.3B-diffusers")
|
||||
prompt = "A sleek robot stands in a vast warehouse filled with boxes"
|
||||
outputs = omni.generate(
|
||||
prompt,
|
||||
OmniDiffusionSamplingParams(
|
||||
height=480,
|
||||
width=832,
|
||||
num_frames=81,
|
||||
num_inference_steps=30,
|
||||
guidance_scale=5.0,
|
||||
),
|
||||
)
|
||||
video = outputs[0].images
|
||||
|
||||
from diffusers.utils import export_to_video
|
||||
export_to_video(list(video[0]), "t2v_output.mp4", fps=16)
|
||||
omni.close()
|
||||
```
|
||||
|
||||
### Local CLI Usage
|
||||
## Text-to-Video (T2V)
|
||||
|
||||
```bash
|
||||
python vace_video_generation.py \
|
||||
--mode t2v \
|
||||
--prompt "A sleek robot stands in a vast warehouse filled with boxes" \
|
||||
--height 480 --width 832 --num-frames 81 \
|
||||
--num-inference-steps 30 --guidance-scale 5.0 --flow-shift 5.0 \
|
||||
--output t2v_output.mp4
|
||||
```
|
||||
|
||||
## Key Arguments
|
||||
|
||||
| Argument | Type | Default | Description |
|
||||
| -------- | ---- | ------- | ----------- |
|
||||
| `--mode` | str | `"t2v"` | VACE task mode: `t2v`, `i2v`, `v2lf`, `flf2v`, `inpaint`, `r2v` |
|
||||
| `--model` | str | `"Wan-AI/Wan2.1-VACE-14B-diffusers"` | Model ID or local path |
|
||||
| `--prompt` | str | `"A cat walking in a garden"` | Text description of desired video |
|
||||
| `--negative-prompt` | str | `""` | Negative prompt for classifier-free guidance |
|
||||
| `--image` | str | `None` | Input image path (for I2V, R2V, FLF2V, inpaint modes) |
|
||||
| `--last-image` | str | `None` | Last frame image path (for FLF2V mode) |
|
||||
| `--height` | int | `480` | Output video height in pixels (should be a multiple of 16) |
|
||||
| `--width` | int | `832` | Output video width in pixels (should be a multiple of 16) |
|
||||
| `--num-frames` | int | `81` | Number of video frames to generate |
|
||||
| `--num-inference-steps` | int | `30` | Number of denoising steps (more steps = higher quality, slower) |
|
||||
| `--guidance-scale` | float | `5.0` | Classifier-free guidance scale |
|
||||
| `--flow-shift` | float | `5.0` | Scheduler flow shift parameter |
|
||||
| `--seed` | int | `42` | Random seed for deterministic sampling |
|
||||
| `--fps` | int | `16` | Frames per second for the saved MP4 |
|
||||
| `--output` | str | `"vace_output.mp4"` | Path to save the generated video |
|
||||
| `--vae-use-tiling` | flag | on | Enable VAE tiling for memory optimization |
|
||||
| `--cfg-parallel-size` | int | `1` | Set to `2` to enable CFG Parallel |
|
||||
| `--ulysses-degree` | int | `1` | Ulysses sequence parallel degree for multi-GPU inference |
|
||||
| `--ring-degree` | int | `1` | Ring sequence parallel degree for hybrid Ulysses + Ring inference |
|
||||
| `--tensor-parallel-size` | int | — | Tensor parallel size |
|
||||
| `--enforce-eager` | flag | off | Disable torch.compile |
|
||||
|
||||
> If you encounter OOM errors, try `--vae-use-tiling` or multi-GPU parallelism options (`--ulysses-degree`, `--cfg-parallel-size`).
|
||||
|
||||
## More CLI Examples
|
||||
|
||||
### Image-to-Video (I2V)
|
||||
## Image-to-Video (I2V)
|
||||
|
||||
First frame is kept, remaining frames are generated:
|
||||
|
||||
@@ -112,7 +31,7 @@ python vace_video_generation.py \
|
||||
--output i2v_output.mp4
|
||||
```
|
||||
|
||||
### First-Last-Frame Interpolation (FLF2V)
|
||||
## First-Last-Frame Interpolation (FLF2V)
|
||||
|
||||
```bash
|
||||
python vace_video_generation.py \
|
||||
@@ -123,7 +42,7 @@ python vace_video_generation.py \
|
||||
--output flf2v_output.mp4
|
||||
```
|
||||
|
||||
### Inpainting
|
||||
## Inpainting
|
||||
|
||||
Center vertical stripe is masked and regenerated:
|
||||
|
||||
@@ -136,7 +55,7 @@ python vace_video_generation.py \
|
||||
--output inpaint_output.mp4
|
||||
```
|
||||
|
||||
### Reference Image-guided (R2V)
|
||||
## Reference Image-guided (R2V)
|
||||
|
||||
```bash
|
||||
python vace_video_generation.py \
|
||||
@@ -146,3 +65,24 @@ python vace_video_generation.py \
|
||||
--height 480 --width 832 --num-frames 81 \
|
||||
--output r2v_output.mp4
|
||||
```
|
||||
|
||||
## Key Arguments
|
||||
|
||||
- `--mode`: VACE task mode (`t2v`, `i2v`, `flf2v`, `inpaint`, `r2v`).
|
||||
- `--model`: Model ID (default: `Wan-AI/Wan2.1-VACE-1.3B-diffusers`).
|
||||
- `--image`: Input image for I2V, inpainting, and R2V modes.
|
||||
- `--last-image`: Last frame image for FLF2V mode.
|
||||
- `--prompt`: Text description of desired video.
|
||||
- `--height/--width`: Output resolution (default 480x832). Dimensions should be multiples of 16.
|
||||
- `--num-frames`: Number of frames (default 81).
|
||||
- `--guidance-scale`: CFG scale (default 5.0).
|
||||
- `--flow-shift`: Scheduler flow shift (default 5.0).
|
||||
- `--num-inference-steps`: Number of denoising steps (default 30).
|
||||
- `--fps`: Frames per second for the saved MP4 (default 16).
|
||||
- `--output`: Path to save the generated video.
|
||||
- `--vae-use-tiling`: Enable VAE tiling for memory optimization.
|
||||
- `--ulysses-degree`: Ulysses sequence parallelism degree for multi-GPU.
|
||||
- `--cfg-parallel-size`: CFG parallel size for multi-GPU.
|
||||
- `--tensor-parallel-size`: Tensor parallel size.
|
||||
|
||||
> If you encounter OOM errors, try `--vae-use-tiling` or multi-GPU parallelism options.
|
||||
@@ -102,5 +102,3 @@ Key arguments:
|
||||
- `--audio-negative-prompt`: negative prompt for audio generation.
|
||||
- `--enable-cpu-offload`: enable CPU offload (defaults False).
|
||||
- `--cache-backend`: enable `cache_dit` for acceleration.
|
||||
- `--quantization`: online (dynamic) quantization method — `fp8` or `int8`.
|
||||
(VAEs, the T5 text encoder, norms and modulation stay bf16.)
|
||||
@@ -4,7 +4,6 @@
|
||||
import argparse
|
||||
import re
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
@@ -42,12 +41,6 @@ def parse_args() -> argparse.Namespace:
|
||||
choices=[1, 2, 3, 4],
|
||||
help="Number of GPUs used for classifier free guidance parallel size (max 4 branches).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--tensor-parallel-size",
|
||||
type=int,
|
||||
default=1,
|
||||
help=("Number of GPUs used for tensor parallelism (TP) inside the DiT."),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--video-negative-prompt",
|
||||
default="jitter, bad hands, blur, distortion",
|
||||
@@ -59,13 +52,6 @@ def parse_args() -> argparse.Namespace:
|
||||
help="Negative prompt for audio.",
|
||||
)
|
||||
parser.add_argument("--output", default="dreamid_output.mp4", help="Output video path.")
|
||||
parser.add_argument(
|
||||
"--quantization",
|
||||
type=str,
|
||||
default=None,
|
||||
choices=["fp8", "int8"],
|
||||
help=("Online (dynamic) quantization method for the DreamID-Omni transformer. "),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enable-cpu-offload",
|
||||
action="store_true",
|
||||
@@ -114,27 +100,6 @@ def load_image_and_audio(image_paths, audio_paths):
|
||||
return image, audio
|
||||
|
||||
|
||||
def _extract_peak_memory_mb(result: Any) -> float:
|
||||
"""Pull worker-reported peak VRAM (MiB) generation result.
|
||||
|
||||
Mirrors vllm_omni/entrypoints/openai/serving_video.py:_extract_peak_memory_mb.
|
||||
"""
|
||||
if isinstance(result, list):
|
||||
result = result[0] if result else None
|
||||
if result is None:
|
||||
return 0.0
|
||||
val = getattr(result, "peak_memory_mb", 0.0)
|
||||
if not val:
|
||||
inner = getattr(result, "request_output", None)
|
||||
if isinstance(inner, list):
|
||||
inner = inner[0] if inner else None
|
||||
val = getattr(inner, "peak_memory_mb", 0.0)
|
||||
try:
|
||||
return float(val or 0.0)
|
||||
except (TypeError, ValueError):
|
||||
return 0.0
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
if args.prompt is None and args.prompt_file is None:
|
||||
@@ -177,7 +142,6 @@ def main() -> None:
|
||||
|
||||
parallel_config = DiffusionParallelConfig(
|
||||
cfg_parallel_size=args.cfg_parallel_size,
|
||||
tensor_parallel_size=args.tensor_parallel_size,
|
||||
use_hsdp=args.use_hsdp,
|
||||
hsdp_shard_size=args.hsdp_shard_size,
|
||||
hsdp_replicate_size=args.hsdp_replicate_size,
|
||||
@@ -198,7 +162,7 @@ def main() -> None:
|
||||
"scm_steps_policy": "dynamic",
|
||||
}
|
||||
|
||||
omni_kwargs = dict(
|
||||
omni = Omni(
|
||||
model=args.model,
|
||||
parallel_config=parallel_config,
|
||||
model_type=args.model_type,
|
||||
@@ -207,9 +171,6 @@ def main() -> None:
|
||||
cache_backend=args.cache_backend,
|
||||
cache_config=cache_config,
|
||||
)
|
||||
if args.quantization is not None:
|
||||
omni_kwargs["quantization"] = args.quantization
|
||||
omni = Omni(**omni_kwargs)
|
||||
start = time.perf_counter()
|
||||
outputs = omni.generate(prompt, sampling_params)
|
||||
elapsed = time.perf_counter() - start
|
||||
@@ -247,9 +208,6 @@ def main() -> None:
|
||||
f.write(video_bytes)
|
||||
print(f"Saved generated video to {output_path}")
|
||||
print(f"Total time: {elapsed:.2f}s")
|
||||
peak_mb = _extract_peak_memory_mb(outputs)
|
||||
if peak_mb:
|
||||
print(f"Worker peak GPU memory (reserved): {peak_mb:.2f} MiB ({peak_mb / 1024:.2f} GiB)")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -10,7 +10,6 @@ Please refer to [README.md](../../../README.md)
|
||||
|------|---------------|--------|
|
||||
| Thinker + Talker (omni-speech, default) | `vllm serve ... --omni` | Text + Audio |
|
||||
| Thinker only (multimodal understanding) | `vllm serve ... --omni --deploy-config vllm_omni/deploy/ming_flash_omni_thinker_only.yaml` | Text |
|
||||
| Thinker + Imagegen (text-to-image / img2img) | `vllm serve ... --omni --stage-configs-path vllm_omni/model_executor/stage_configs/ming_flash_omni_dual.yaml` | Image |
|
||||
|
||||
For standalone TTS (talker only), see the [Ming-flash-omni-TTS section in the Text-To-Speech hub](../text_to_speech/README.md#ming-flash-omni-tts).
|
||||
|
||||
@@ -57,98 +56,8 @@ bash run_curl_multimodal_generation.sh use_image
|
||||
bash run_curl_multimodal_generation.sh use_audio
|
||||
bash run_curl_multimodal_generation.sh use_video
|
||||
bash run_curl_multimodal_generation.sh use_mixed_modalities
|
||||
bash run_curl_multimodal_generation.sh use_image_gen
|
||||
```
|
||||
|
||||
## Image generation (text-to-image)
|
||||
|
||||
Ming-flash-omni-2.0 also exposes an image-generation (diffusion) stage.
|
||||
Launch with the dual-stage YAML, which adds an image-gen stage behind the
|
||||
thinker:
|
||||
|
||||
```bash
|
||||
vllm serve Jonathan1909/Ming-flash-omni-2.0 --omni \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/ming_flash_omni_dual.yaml \
|
||||
--stage-init-timeout 1800 \
|
||||
--port 8188
|
||||
```
|
||||
|
||||
Then request image output by passing `"modalities": ["image"]`:
|
||||
|
||||
```bash
|
||||
curl http://127.0.0.1:8188/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "Jonathan1909/Ming-flash-omni-2.0",
|
||||
"messages": [{"role": "user", "content": "Draw a beautiful girl with short black hair and red dress."}],
|
||||
"modalities": ["image"]
|
||||
}' -o /tmp/ming_response.json
|
||||
|
||||
python - <<'PY'
|
||||
import base64, json
|
||||
r = json.load(open('/tmp/ming_response.json'))
|
||||
url = r['choices'][0]['message']['content'][0]['image_url']['url']
|
||||
open('/tmp/ming_gen.png', 'wb').write(base64.b64decode(url.split(',', 1)[1]))
|
||||
print('saved /tmp/ming_gen.png')
|
||||
PY
|
||||
```
|
||||
|
||||
### Optional knobs
|
||||
|
||||
Pass image-gen overrides via the diffusion-stage `sampling_params_list[1].extra_args.image_gen`:
|
||||
|
||||
| Key | Default | Description |
|
||||
| --- | --- | --- |
|
||||
| `height` / `width` | from config (1024) | Output resolution (multiples of `vae_scale_factor * 2`, currently 16). |
|
||||
| `steps` | 30 | Number of FlowMatchEuler denoise steps. |
|
||||
| `cfg` | 2.0 | Classifier-free guidance scale. |
|
||||
| `seed` | 42 | Per-request RNG seed (deterministic when ByT5 is also seed-stable). |
|
||||
| `byte5_text` | (auto) | Override the glyph text for ByT5 enhancement; raw strings are auto-wrapped to Ming's `Text "...". ` format. |
|
||||
| `negative_prompt` | empty | Real CFG negative conditioning (must be set on **stage-0 thinker** `sampling_params` so `expand_cfg_prompts` spawns the companion). |
|
||||
|
||||
Example with all the knobs:
|
||||
|
||||
```bash
|
||||
curl http://127.0.0.1:8188/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "Jonathan1909/Ming-flash-omni-2.0",
|
||||
"modalities": ["image"],
|
||||
"sampling_params_list": [
|
||||
{"temperature":0.4,"top_p":0.9,"top_k":1,"max_tokens":1,"seed":42,
|
||||
"extra_args":{"image_gen":{"negative_prompt":"ugly, blurry, distorted"}}},
|
||||
{"seed":42,"extra_args":{"image_gen":{
|
||||
"steps":6,"cfg":1.5,"height":512,"width":512,"seed":123,
|
||||
"byte5_text":["理解与生成统一"]}}}
|
||||
],
|
||||
"messages": [{"role": "user", "content": "Draw a poster."}]
|
||||
}' -o /tmp/ming_response.json
|
||||
```
|
||||
|
||||
### img2img (reference image)
|
||||
|
||||
Add an `image_url` content part to the user message; the parser routes it
|
||||
into the diffusion stage as `extra[reference_image]`:
|
||||
|
||||
```jsonc
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Change the background to a sandy beach at sunset."},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,<base64>"}}
|
||||
]
|
||||
}]
|
||||
```
|
||||
|
||||
### GPU layout
|
||||
|
||||
The shipped `ming_flash_omni_dual.yaml` allocates the thinker on GPUs 0–3
|
||||
(TP=4) and the diffusion stage on GPU 4 (TP=1). Copy the YAML and edit
|
||||
`runtime.devices` per stage to relocate; with fewer GPUs available, drop
|
||||
the thinker TP to 2 and run the diffusion stage on a free card. Image-gen
|
||||
warmup takes roughly an extra 30–60 s on top of the thinker — set
|
||||
`--stage-init-timeout 1800` if the default 300 s is too tight.
|
||||
|
||||
## Modality control
|
||||
|
||||
| `modalities` | Server config | Output |
|
||||
@@ -156,7 +65,6 @@ warmup takes roughly an extra 30–60 s on top of the thinker — set
|
||||
| `["text"]` or omitted | Thinker only | Text |
|
||||
| `["audio"]` | Thinker + Talker | Audio (speech) |
|
||||
| `["text", "audio"]` | Thinker + Talker | Text + Audio |
|
||||
| `["image"]` | Thinker + Imagegen (dual YAML) | Image (PNG, base64 in `choices[0].message.content`) |
|
||||
|
||||
For ready-to-copy curl examples (text / audio / multimodal input, SSE
|
||||
streaming, reasoning mode), see the recipe at
|
||||
|
||||
@@ -7,15 +7,14 @@ PORT="${PORT:-8091}"
|
||||
QUERY_TYPE="${1:-text}"
|
||||
|
||||
# Validate query type
|
||||
if [[ ! "$QUERY_TYPE" =~ ^(text|use_audio|use_image|use_video|use_mixed_modalities|use_image_gen)$ ]]; then
|
||||
if [[ ! "$QUERY_TYPE" =~ ^(text|use_audio|use_image|use_video|use_mixed_modalities)$ ]]; then
|
||||
echo "Error: Invalid query type '$QUERY_TYPE'"
|
||||
echo "Usage: $0 [text|use_audio|use_image|use_video|use_mixed_modalities|use_image_gen]"
|
||||
echo "Usage: $0 [text|use_audio|use_image|use_video|use_mixed_modalities]"
|
||||
echo " text: Text-only query"
|
||||
echo " use_audio: Audio + Text query"
|
||||
echo " use_image: Image + Text query"
|
||||
echo " use_video: Video + Text query"
|
||||
echo " use_mixed_modalities: Audio + Image + Video + Text query"
|
||||
echo " use_image_gen: Text-to-image (diffusion stage) — saves the generated PNG locally"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
@@ -101,43 +100,6 @@ esac
|
||||
echo "Running query type: $QUERY_TYPE"
|
||||
echo ""
|
||||
|
||||
if [[ "$QUERY_TYPE" == "use_image_gen" ]]; then
|
||||
# Image-gen branch: dual-stage diffusion endpoint, modalities=["image"].
|
||||
# The shipped ming_flash_omni_dual.yaml runs the thinker on cards 0-3
|
||||
# (TP=4) and the diffusion stage on card 4 (TP=1). See the image-gen
|
||||
# section in README.md for knobs and img2img.
|
||||
out_path="${OUT_PATH:-/tmp/ming_imagegen.png}"
|
||||
request_body=$(cat <<'EOF'
|
||||
{
|
||||
"model": "Jonathan1909/Ming-flash-omni-2.0",
|
||||
"modalities": ["image"],
|
||||
"messages": [
|
||||
{"role": "user", "content": "Draw a beautiful girl with short black hair and red dress."}
|
||||
]
|
||||
}
|
||||
EOF
|
||||
)
|
||||
output=$(curl -sS --retry 3 --retry-delay 3 --retry-connrefused \
|
||||
-X POST "http://localhost:${PORT}/v1/chat/completions" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d "$request_body")
|
||||
# Server returns base64-encoded PNG either as a data URL in
|
||||
# ``choices[0].message.content[0].image_url.url`` or as a raw string.
|
||||
b64=$(echo "$output" | jq -r '.choices[0].message.content
|
||||
| if type == "array"
|
||||
then (map(select(.image_url? != null)) | .[0].image_url.url // "")
|
||||
elif type == "string" then . else "" end')
|
||||
if [[ -z "$b64" ]]; then
|
||||
echo "Error: no image returned. Raw response:"
|
||||
echo "$output" | jq '.'
|
||||
exit 1
|
||||
fi
|
||||
if [[ "$b64" == data:image* ]]; then b64="${b64#*,}"; fi
|
||||
printf '%s' "$b64" | base64 --decode > "$out_path"
|
||||
echo "Saved generated image: $out_path ($(stat -c%s "$out_path" 2>/dev/null || stat -f%z "$out_path") bytes)"
|
||||
exit 0
|
||||
fi
|
||||
|
||||
request_body=$(cat <<EOF
|
||||
{
|
||||
"model": "Jonathan1909/Ming-flash-omni-2.0",
|
||||
|
||||
@@ -143,26 +143,6 @@ bash examples/online_serving/text_to_speech/glm_tts/run_gradio_demo.sh
|
||||
pip install fish-speech
|
||||
```
|
||||
|
||||
### Kvcache attention fast path
|
||||
|
||||
Fish Speech S2 Pro uses a Triton decode-only kvcache attention fast path by
|
||||
default on CUDA builds. Set `VLLM_OMNI_FISH_KVCACHE_ATTN=0` to disable it, or
|
||||
`VLLM_OMNI_FISH_KVCACHE_ATTN=required` to fail fast if the fast path cannot be
|
||||
installed.
|
||||
|
||||
```bash
|
||||
# Verify fast path availability.
|
||||
python - <<'PY'
|
||||
from vllm_omni.attention import fish_kvcache_attn
|
||||
|
||||
print(fish_kvcache_attn.is_available())
|
||||
print(fish_kvcache_attn.load_error())
|
||||
PY
|
||||
|
||||
# Optional: disable the runtime fast path.
|
||||
export VLLM_OMNI_FISH_KVCACHE_ATTN=0
|
||||
```
|
||||
|
||||
### Launch
|
||||
```bash
|
||||
vllm serve fishaudio/s2-pro --omni --port 8091
|
||||
|
||||
@@ -1,63 +0,0 @@
|
||||
# higgs-audio v2 online example
|
||||
|
||||
This directory contains the online-serving entry points for boson-ai's higgs-audio v2 as integrated by vllm-omni: a 2-stage TTS pipeline (Llama-3.2-3B talker with DualFFN audio expert + HiggsAudio codec decoder) emitting 24 kHz mono speech.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
Voice clone uses HF's `HiggsAudioV2TokenizerModel` loaded from `k2-fsa/OmniVoice/audio_tokenizer/` (the boson-ai standalone tokenizer Hub repo's `model.safetensors` is the 3B talker LM, not the codec). Only that ~806 MB subdir is downloaded.
|
||||
|
||||
```bash
|
||||
pip install -U "transformers>=5.3.0"
|
||||
```
|
||||
|
||||
## Files
|
||||
|
||||
- `run_server.sh` — launch the vllm-omni server with the bundled `vllm_omni/deploy/higgs_audio_v2.yaml` deploy config.
|
||||
- `batch_speech_client.py` — send a list of prompts to `/v1/audio/speech` and save the returned WAV / PCM bytes to a directory; optionally passes `--ref-audio` + `--ref-text` for shallow voice clone.
|
||||
|
||||
## Launching the server
|
||||
|
||||
```bash
|
||||
GPUS=6,7 PORT=8094 bash examples/online_serving/text_to_speech/higgs_audio_v2/run_server.sh
|
||||
```
|
||||
|
||||
Environment overrides:
|
||||
|
||||
- `MODEL` — HF id of the talker (default `bosonai/higgs-audio-v2-generation-3B-base`).
|
||||
- `PORT` — server port (default `8094`).
|
||||
- `GPUS` — `CUDA_VISIBLE_DEVICES` value (default `6,7`).
|
||||
- `GPU_UTIL` — `--gpu-memory-utilization` (default `0.4`).
|
||||
|
||||
The script also exports `VLLM_USE_DEEP_GEMM=0` / `VLLM_MOE_USE_DEEP_GEMM=0` so the example works on images without the optional `deep_gemm` backend.
|
||||
|
||||
The deploy YAML ships with `async_chunk: false` and `codec_streaming: true`, i.e. Stage 0 finishes its codec frames before Stage 1 starts decoding, and Stage 1 streams WAV/PCM bytes to the client chunk-by-chunk.
|
||||
|
||||
## Driving the server
|
||||
|
||||
Plain TTS:
|
||||
|
||||
```bash
|
||||
python examples/online_serving/text_to_speech/higgs_audio_v2/batch_speech_client.py \
|
||||
--base-url http://localhost:8094 \
|
||||
--model bosonai/higgs-audio-v2-generation-3B-base \
|
||||
--output-dir /tmp/higgs_audio_v2_batch \
|
||||
--prompts "Hello world." \
|
||||
"The quick brown fox jumps over the lazy dog."
|
||||
```
|
||||
|
||||
Voice clone — pass a reference clip and its transcript (both required together):
|
||||
|
||||
```bash
|
||||
python examples/online_serving/text_to_speech/higgs_audio_v2/batch_speech_client.py \
|
||||
--base-url http://localhost:8094 \
|
||||
--model bosonai/higgs-audio-v2-generation-3B-base \
|
||||
--output-dir /tmp/higgs_audio_v2_clone \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--ref-text "Exact transcript spoken in reference.wav." \
|
||||
--prompts "Hello, this is a cloned voice."
|
||||
```
|
||||
|
||||
## Notes
|
||||
|
||||
- `--ref-text` must be the real transcript of `--ref-audio`; mismatched text degrades cloned-voice quality.
|
||||
- Out of scope (rejected with explicit 4xx by the request validator): multi-speaker `[SPEAKERn]` tags inside `input`, `profile:` text-only speaker descriptions, the `ref_audio_in_system_message` system-block variant, chunked long-form generation, and per-request `voice` / `instructions` / `task_type` / `language` / `speed != 1.0` / `x_vector_only_mode` / `speaker_embedding`.
|
||||
@@ -1,122 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Batch client for the higgs-audio v2 online server.
|
||||
|
||||
Sends a fixed list of prompts to ``/v1/audio/speech`` and saves the returned
|
||||
WAV files (or raw PCM bytes when ``--format pcm``) into ``--output-dir``.
|
||||
|
||||
Usage (plain text -> speech):
|
||||
|
||||
python examples/online_serving/text_to_speech/higgs_audio_v2/batch_speech_client.py \
|
||||
--base-url http://localhost:8094 \
|
||||
--output-dir /tmp/higgs_audio_v2_batch \
|
||||
--prompts "Hello world." "The quick brown fox jumps over the lazy dog."
|
||||
|
||||
Usage (shallow voice clone — pass a reference clip + its transcript):
|
||||
|
||||
python examples/online_serving/text_to_speech/higgs_audio_v2/batch_speech_client.py \
|
||||
--base-url http://localhost:8094 \
|
||||
--output-dir /tmp/higgs_audio_v2_clone \
|
||||
--ref-audio path/to/reference.wav \
|
||||
--ref-text "the transcript of the reference clip" \
|
||||
--prompts "Hello world."
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
DEFAULT_PROMPTS = (
|
||||
"Hello world.",
|
||||
"The quick brown fox jumps over the lazy dog.",
|
||||
"It was the night before my birthday.",
|
||||
"Innovation distinguishes between a leader and a follower.",
|
||||
)
|
||||
|
||||
|
||||
def _slug(text: str) -> str:
|
||||
import re
|
||||
|
||||
s = re.sub(r"\s+", "_", text.strip().lower())
|
||||
return re.sub(r"[^a-z0-9_]+", "", s)[:32] or "prompt"
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
parser.add_argument("--base-url", default="http://localhost:8094")
|
||||
parser.add_argument("--model", default="higgs_audio_v2")
|
||||
parser.add_argument("--prompts", nargs="+", default=list(DEFAULT_PROMPTS))
|
||||
parser.add_argument("--output-dir", type=Path, default=Path("/tmp/higgs_audio_v2_batch"))
|
||||
parser.add_argument("--format", choices=("wav", "pcm"), default="wav")
|
||||
parser.add_argument("--max-new-tokens", type=int, default=300)
|
||||
parser.add_argument("--seed", type=int, default=42)
|
||||
parser.add_argument("--timeout-s", type=float, default=120.0)
|
||||
parser.add_argument(
|
||||
"--ref-audio",
|
||||
type=Path,
|
||||
default=None,
|
||||
help="Reference clip for voice clone (path to a WAV file). Must be paired with --ref-text.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ref-text",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Transcript of the reference clip. Required when --ref-audio is set.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
if (args.ref_audio is None) != (args.ref_text is None):
|
||||
print("--ref-audio and --ref-text must be supplied together", file=sys.stderr)
|
||||
return 2
|
||||
|
||||
ref_audio_data_url: str | None = None
|
||||
if args.ref_audio is not None:
|
||||
if not args.ref_audio.exists():
|
||||
print(f"ref-audio file not found: {args.ref_audio}", file=sys.stderr)
|
||||
return 2
|
||||
mime = "audio/wav" if args.ref_audio.suffix.lower() == ".wav" else "audio/mpeg"
|
||||
ref_b64 = base64.b64encode(args.ref_audio.read_bytes()).decode("ascii")
|
||||
ref_audio_data_url = f"data:{mime};base64,{ref_b64}"
|
||||
|
||||
try:
|
||||
import httpx
|
||||
except ImportError:
|
||||
print(
|
||||
"this client needs `httpx`. Install with `pip install httpx`.",
|
||||
file=sys.stderr,
|
||||
)
|
||||
return 2
|
||||
|
||||
args.output_dir.mkdir(parents=True, exist_ok=True)
|
||||
url = args.base_url.rstrip("/") + "/v1/audio/speech"
|
||||
failures = 0
|
||||
with httpx.Client(timeout=args.timeout_s) as client:
|
||||
for prompt in args.prompts:
|
||||
payload = {
|
||||
"model": args.model,
|
||||
"input": prompt,
|
||||
"response_format": args.format,
|
||||
"max_new_tokens": args.max_new_tokens,
|
||||
"seed": args.seed,
|
||||
}
|
||||
if ref_audio_data_url is not None:
|
||||
payload["ref_audio"] = ref_audio_data_url
|
||||
payload["ref_text"] = args.ref_text
|
||||
resp = client.post(url, json=payload)
|
||||
if resp.status_code != 200:
|
||||
print(f"[FAIL] {prompt!r} -> {resp.status_code}: {resp.text[:200]}", file=sys.stderr)
|
||||
failures += 1
|
||||
continue
|
||||
suffix = ".wav" if args.format == "wav" else ".pcm"
|
||||
out = args.output_dir / f"{_slug(prompt)}{suffix}"
|
||||
out.write_bytes(resp.content)
|
||||
print(f"[ ok ] {prompt!r} -> {out} ({len(resp.content)} bytes)")
|
||||
|
||||
return 1 if failures else 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -1,37 +0,0 @@
|
||||
#!/bin/bash
|
||||
# Launch vLLM-Omni server for higgs-audio v2.
|
||||
#
|
||||
# v1 scope: plain text -> 24 kHz speech only. Voice cloning, multi-speaker,
|
||||
# ChatML rich content, and language overrides are rejected by the validator
|
||||
# with explicit 4xx (see vllm_omni/entrypoints/openai/serving_speech.py).
|
||||
#
|
||||
# Usage:
|
||||
# ./run_server.sh # default port 8094, GPUs 6 and 7
|
||||
# PORT=8095 GPUS=6,7 ./run_server.sh
|
||||
# MODEL=bosonai/higgs-audio-v2-generation-3B-base ./run_server.sh
|
||||
|
||||
set -e
|
||||
|
||||
MODEL="${MODEL:-bosonai/higgs-audio-v2-generation-3B-base}"
|
||||
PORT="${PORT:-8094}"
|
||||
GPUS="${GPUS:-6,7}"
|
||||
GPU_UTIL="${GPU_UTIL:-0.4}"
|
||||
|
||||
echo "Starting higgs-audio v2 server"
|
||||
echo " MODEL=$MODEL"
|
||||
echo " PORT=$PORT"
|
||||
echo " CUDA_VISIBLE_DEVICES=$GPUS"
|
||||
|
||||
# DeepGEMM FP8 kernels are optional and trip warmup on builds without
|
||||
# the deep_gemm backend; disable them so the example works out of the box.
|
||||
# Users with deep_gemm installed can re-enable via the same env vars.
|
||||
CUDA_VISIBLE_DEVICES="$GPUS" \
|
||||
VLLM_USE_DEEP_GEMM=0 \
|
||||
VLLM_MOE_USE_DEEP_GEMM=0 \
|
||||
vllm-omni serve "$MODEL" \
|
||||
--deploy-config vllm_omni/deploy/higgs_audio_v2.yaml \
|
||||
--host 0.0.0.0 \
|
||||
--port "$PORT" \
|
||||
--gpu-memory-utilization "$GPU_UTIL" \
|
||||
--trust-remote-code \
|
||||
--omni
|
||||
+4
-4
@@ -38,12 +38,13 @@ theme:
|
||||
features:
|
||||
- content.action.edit
|
||||
- content.code.copy
|
||||
- navigation.instant
|
||||
- navigation.instant.progress
|
||||
- navigation.tracking
|
||||
- navigation.tabs
|
||||
- navigation.tabs.sticky
|
||||
- navigation.sections
|
||||
- navigation.indexes
|
||||
- navigation.prune
|
||||
- navigation.top
|
||||
- search.suggest
|
||||
- search.highlight
|
||||
@@ -55,7 +56,6 @@ theme:
|
||||
custom_dir: docs/mkdocs/overrides
|
||||
|
||||
hooks:
|
||||
- docs/mkdocs/hooks/search_index.py
|
||||
- docs/mkdocs/hooks/generate_api_readme.py
|
||||
- docs/mkdocs/hooks/url_schemes.py
|
||||
- docs/mkdocs/hooks/generate_examples.py
|
||||
@@ -117,7 +117,6 @@ plugins:
|
||||
separate_signature: true
|
||||
show_overloads: true
|
||||
signature_crossrefs: true
|
||||
show_source: false
|
||||
inventories:
|
||||
- https://docs.python.org/3/objects.inv
|
||||
- https://typing-extensions.readthedocs.io/en/latest/objects.inv
|
||||
@@ -168,8 +167,9 @@ extra_css:
|
||||
- mkdocs/stylesheets/extra.css
|
||||
|
||||
extra_javascript:
|
||||
- mkdocs/javascript/memory.js
|
||||
- mkdocs/javascript/mathjax.js
|
||||
- https://unpkg.com/mathjax@3.2.2/es5/tex-mml-chtml.js
|
||||
- https://unpkg.com/mermaid@10/dist/mermaid.min.js
|
||||
- mkdocs/javascript/mermaid.js
|
||||
- mkdocs/javascript/edit_and_feedback.js
|
||||
- mkdocs/javascript/slack_and_forum.js
|
||||
@@ -249,7 +249,6 @@ extend-ignore-identifiers-re = [
|
||||
".*[Oo]no_[Aa]nna.*",
|
||||
".*cann.*",
|
||||
".*[Nn][Dd]3.*",
|
||||
"byt5.*",
|
||||
]
|
||||
extend-ignore-re = [
|
||||
"\\.\\.\\.n,d->\\.\\.\\.nd",
|
||||
@@ -262,4 +261,3 @@ Vai = "Vai"
|
||||
tockens = "tockens"
|
||||
CANN = "CANN"
|
||||
ND = "ND"
|
||||
byt5 = "byt5"
|
||||
@@ -86,6 +86,7 @@ curl -X POST http://localhost:8000/v1/videos \
|
||||
- Duration: Controlled by `num_frames` (frames = duration_seconds * 24 + 1)
|
||||
- Known limitations:
|
||||
- No image-to-video support yet (LTX23ImageToVideoPipeline is a placeholder)
|
||||
- No CFG-parallel support (single-GPU only)
|
||||
- Requires `diffusers >= 0.38.0` (not yet on PyPI, install from git)
|
||||
|
||||
## Hardware Support
|
||||
|
||||
@@ -50,26 +50,6 @@ fi
|
||||
pip install fish-speech
|
||||
```
|
||||
|
||||
#### Kvcache attention fast path
|
||||
|
||||
Fish Speech S2 Pro uses a Triton decode-only kvcache attention fast path by
|
||||
default on CUDA builds. Set `VLLM_OMNI_FISH_KVCACHE_ATTN=0` to disable it, or
|
||||
`VLLM_OMNI_FISH_KVCACHE_ATTN=required` to fail fast if the fast path cannot be
|
||||
installed.
|
||||
|
||||
```bash
|
||||
# Verify fast path availability.
|
||||
python - <<'PY'
|
||||
from vllm_omni.attention import fish_kvcache_attn
|
||||
|
||||
print(fish_kvcache_attn.is_available())
|
||||
print(fish_kvcache_attn.load_error())
|
||||
PY
|
||||
|
||||
# Optional: disable the runtime fast path.
|
||||
export VLLM_OMNI_FISH_KVCACHE_ATTN=0
|
||||
```
|
||||
|
||||
#### Command
|
||||
|
||||
```bash
|
||||
|
||||
@@ -62,12 +62,6 @@ def detect_target_device() -> str:
|
||||
else:
|
||||
print(f"Warning: Invalid VLLM_OMNI_TARGET_DEVICE '{target_device}', falling back to auto-detection")
|
||||
|
||||
# Priority 1.5: ReadTheDocs builds use CPU requirements (no GPU, 1 GB RAM cap).
|
||||
# Avoids pulling ~2 GB of CUDA libs that push the mkdocs build into swap.
|
||||
if os.environ.get("READTHEDOCS"):
|
||||
print("Detected READTHEDOCS=1; using cpu requirements for docs build")
|
||||
return "cpu"
|
||||
|
||||
# Priority 2: Torch backend detection
|
||||
# This is a code path for when user is using
|
||||
# --no-build-isolation flag
|
||||
|
||||
File diff suppressed because it is too large.
Load diff
@@ -28,7 +28,7 @@ class FakeAdapter:
|
||||
req = running.pop()
|
||||
self.waiting_for_chunk_running_requests.append(req)
|
||||
|
||||
def restore_queues(self, waiting, running, scheduler_requests=None):
|
||||
def restore_queues(self, waiting, running):
|
||||
"""Put requests back."""
|
||||
self.restore_called = True
|
||||
running.extend(self.waiting_for_chunk_running_requests)
|
||||
|
||||
@@ -1,104 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit coverage for _process_pending_input_timeouts.
|
||||
|
||||
Verifies that the mixin correctly *delegates* timed-out requests to the
|
||||
base scheduler's ``finish_requests`` API with ``RequestStatus.FINISHED_ERROR``.
|
||||
The end-to-end effect (queue removal + status set + per-request cleanup +
|
||||
client-facing FINISHED_ERROR emission) is the responsibility of upstream
|
||||
vLLM's ``finish_requests`` implementation and is covered by upstream tests;
|
||||
this file only asserts the wiring from the mixin to that API.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
|
||||
|
||||
|
||||
class _FakeCoordinator:
|
||||
def __init__(self, timed_out_ids):
|
||||
self._timed_out_ids = set(timed_out_ids)
|
||||
self.calls = []
|
||||
|
||||
def collect_timed_out_request_ids(self, timeout_s):
|
||||
self.calls.append(timeout_s)
|
||||
return set(self._timed_out_ids)
|
||||
|
||||
|
||||
class _FakeScheduler(OmniSchedulerMixin):
|
||||
def __init__(self, requests, coordinator):
|
||||
self.requests = requests
|
||||
self.input_coordinator = coordinator
|
||||
self.finish_calls = []
|
||||
|
||||
def finish_requests(self, req_ids, status):
|
||||
self.finish_calls.append((set(req_ids), status))
|
||||
|
||||
|
||||
def test_process_pending_input_timeouts_delegates_to_finish_requests():
|
||||
"""Timed-out request present in self.requests is forwarded to finish_requests."""
|
||||
req_id = "stuck-req"
|
||||
requests = {req_id: SimpleNamespace(request_id=req_id)}
|
||||
coord = _FakeCoordinator(timed_out_ids={req_id})
|
||||
scheduler = _FakeScheduler(requests, coord)
|
||||
|
||||
scheduler._process_pending_input_timeouts()
|
||||
|
||||
assert len(coord.calls) == 1, "coordinator should be polled once"
|
||||
assert coord.calls[0] > 0, "timeout must be positive when enabled"
|
||||
|
||||
assert len(scheduler.finish_calls) == 1
|
||||
finished_ids, status = scheduler.finish_calls[0]
|
||||
assert finished_ids == {req_id}
|
||||
# RequestStatus is the upstream enum; the mixin imports it as
|
||||
# RequestStatus.FINISHED_ERROR. Check by name to avoid hard import here.
|
||||
assert getattr(status, "name", str(status)).endswith("FINISHED_ERROR")
|
||||
|
||||
|
||||
def test_process_pending_input_timeouts_skips_already_freed_request():
|
||||
"""Timed-out id no longer in self.requests must not be forwarded."""
|
||||
coord = _FakeCoordinator(timed_out_ids={"already-freed"})
|
||||
scheduler = _FakeScheduler(requests={}, coordinator=coord)
|
||||
|
||||
scheduler._process_pending_input_timeouts()
|
||||
|
||||
assert coord.calls == [coord.calls[0]] and coord.calls[0] > 0
|
||||
assert scheduler.finish_calls == []
|
||||
|
||||
|
||||
def test_process_pending_input_timeouts_noop_without_coordinator():
|
||||
"""No coordinator => no finish_requests call, no crash."""
|
||||
|
||||
class _NoCoord(OmniSchedulerMixin):
|
||||
def __init__(self):
|
||||
self.requests = {}
|
||||
self.input_coordinator = None
|
||||
self.finish_calls = []
|
||||
|
||||
def finish_requests(self, req_ids, status):
|
||||
self.finish_calls.append((set(req_ids), status))
|
||||
|
||||
scheduler = _NoCoord()
|
||||
scheduler._process_pending_input_timeouts()
|
||||
assert scheduler.finish_calls == []
|
||||
|
||||
|
||||
def test_process_pending_input_timeouts_disabled_when_timeout_zero(monkeypatch):
|
||||
"""Setting DEFAULT_INPUT_WAIT_TIMEOUT_S <= 0 disables the safety net."""
|
||||
from vllm_omni.core.sched import omni_scheduler_mixin
|
||||
|
||||
monkeypatch.setattr(omni_scheduler_mixin, "DEFAULT_INPUT_WAIT_TIMEOUT_S", 0.0)
|
||||
|
||||
coord = _FakeCoordinator(timed_out_ids={"r1"})
|
||||
scheduler = _FakeScheduler(requests={"r1": SimpleNamespace(request_id="r1")}, coordinator=coord)
|
||||
scheduler._process_pending_input_timeouts()
|
||||
assert coord.calls == [], "coordinator must not be polled when timeout is disabled"
|
||||
assert scheduler.finish_calls == []
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
@@ -16,7 +16,7 @@ import torch
|
||||
import vllm_omni.core.sched.omni_scheduling_coordinator as coord_mod
|
||||
from vllm_omni.core.sched.omni_scheduling_coordinator import (
|
||||
OmniSchedulingCoordinator,
|
||||
uses_full_payload_input_coordinator,
|
||||
uses_qwen3_omni_full_payload_input_coordinator,
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
@@ -92,96 +92,47 @@ class MockQueue:
|
||||
|
||||
|
||||
class TestFullPayloadCoordinatorSelection(unittest.TestCase):
|
||||
"""Tests for the (model_arch, model_stage) whitelist gate.
|
||||
|
||||
The init_omni_connectors arch allowlist is keyed by ``model_arch`` and
|
||||
is a superset of the stages registered here -- consumer-wait stages
|
||||
must be registered explicitly in ``_FULL_PAYLOAD_INPUT_STAGES``, while
|
||||
the init allowlist covers both producer- and consumer-side runners.
|
||||
These tests pin which ``(arch, stage)`` pairs the gate fires for today.
|
||||
"""
|
||||
|
||||
# Expected whitelist (model_arch, model_stage). Hardcoded to avoid the
|
||||
# tautology of importing _FULL_PAYLOAD_INPUT_STAGES and asserting it
|
||||
# against itself; any drift between this matrix and the whitelist will
|
||||
# fail loudly here.
|
||||
EXPECTED_FULL_PAYLOAD_INPUT_STAGES: frozenset[tuple[str, str]] = frozenset(
|
||||
{
|
||||
("Qwen3OmniMoeForConditionalGeneration", "talker"),
|
||||
("Qwen3OmniMoeForConditionalGeneration", "code2wav"),
|
||||
("Qwen2_5OmniForConditionalGeneration", "talker"),
|
||||
("Qwen2_5OmniForConditionalGeneration", "code2wav"),
|
||||
("CovoAudioForConditionalGeneration", "code2wav"),
|
||||
("MiMoAudioModel", "code2wav"),
|
||||
("Qwen3TTSCode2Wav", "code2wav"),
|
||||
("CosyVoice3Model", "cosyvoice3_code2wav"),
|
||||
("DyninOmniForConditionalGeneration", "token2image"),
|
||||
("DyninOmniForConditionalGeneration", "token2audio"),
|
||||
}
|
||||
)
|
||||
|
||||
def test_whitelist_matches_expected_matrix(self):
|
||||
"""_FULL_PAYLOAD_INPUT_STAGES must equal the hardcoded expected matrix.
|
||||
|
||||
Catches both accidental additions (which would silently enable the
|
||||
consumer-wait gate for a new arch) and accidental removals (which
|
||||
would silently disable an enabled arch).
|
||||
"""
|
||||
from vllm_omni.core.sched.omni_scheduling_coordinator import _FULL_PAYLOAD_INPUT_STAGES
|
||||
|
||||
self.assertEqual(
|
||||
frozenset(_FULL_PAYLOAD_INPUT_STAGES),
|
||||
self.EXPECTED_FULL_PAYLOAD_INPUT_STAGES,
|
||||
msg="_FULL_PAYLOAD_INPUT_STAGES drifted from the expected matrix; "
|
||||
"update EXPECTED_FULL_PAYLOAD_INPUT_STAGES if intentional.",
|
||||
)
|
||||
|
||||
def test_all_whitelisted_arch_stage_pairs_fire_gate(self):
|
||||
"""Every (arch, stage) pair in the expected matrix must fire
|
||||
the gate when stage_id > 0 and async_chunk=False.
|
||||
"""
|
||||
for arch, stage in self.EXPECTED_FULL_PAYLOAD_INPUT_STAGES:
|
||||
def test_qwen3_omni_talker_and_code2wav_use_full_payload_input_coordinator(self):
|
||||
for model_stage in ("talker", "code2wav"):
|
||||
model_config = SimpleNamespace(
|
||||
stage_id=1,
|
||||
async_chunk=False,
|
||||
model_arch=arch,
|
||||
model_stage=stage,
|
||||
)
|
||||
self.assertTrue(
|
||||
uses_full_payload_input_coordinator(model_config),
|
||||
msg=f"expected gate to fire for {arch}/{stage}",
|
||||
model_arch="Qwen3OmniMoeForConditionalGeneration",
|
||||
model_stage=model_stage,
|
||||
)
|
||||
|
||||
def test_other_arch_or_stage_or_mode_does_not_fire(self):
|
||||
self.assertTrue(uses_qwen3_omni_full_payload_input_coordinator(model_config))
|
||||
|
||||
def test_async_chunk_and_non_qwen3_omni_do_not_use_full_payload_input_coordinator(self):
|
||||
cases = [
|
||||
SimpleNamespace(
|
||||
stage_id=1, async_chunk=True, model_arch="Qwen3OmniMoeForConditionalGeneration", model_stage="talker"
|
||||
),
|
||||
SimpleNamespace(
|
||||
stage_id=0, async_chunk=False, model_arch="Qwen3OmniMoeForConditionalGeneration", model_stage="thinker"
|
||||
stage_id=1,
|
||||
async_chunk=True,
|
||||
model_arch="Qwen3OmniMoeForConditionalGeneration",
|
||||
model_stage="talker",
|
||||
),
|
||||
SimpleNamespace(
|
||||
stage_id=1,
|
||||
async_chunk=False,
|
||||
model_arch="Qwen3TTSForConditionalGeneration",
|
||||
model_stage="code2wav",
|
||||
),
|
||||
SimpleNamespace(
|
||||
stage_id=1,
|
||||
async_chunk=False,
|
||||
model_arch="Qwen2_5OmniForConditionalGeneration",
|
||||
model_stage="talker",
|
||||
),
|
||||
SimpleNamespace(
|
||||
stage_id=0,
|
||||
async_chunk=False,
|
||||
model_arch="Qwen3OmniMoeForConditionalGeneration",
|
||||
model_stage="some_future_stage",
|
||||
),
|
||||
SimpleNamespace(
|
||||
stage_id=1, async_chunk=False, model_arch="Qwen3TTSForConditionalGeneration", model_stage="code2wav"
|
||||
),
|
||||
SimpleNamespace(
|
||||
stage_id=1, async_chunk=False, model_arch="MingFlashOmniForConditionalGeneration", model_stage="talker"
|
||||
),
|
||||
SimpleNamespace(stage_id=1, async_chunk=False, model_arch=None, model_stage="talker"),
|
||||
SimpleNamespace(
|
||||
stage_id=1, async_chunk=False, model_arch="Qwen3OmniMoeForConditionalGeneration", model_stage=None
|
||||
model_stage="thinker",
|
||||
),
|
||||
]
|
||||
|
||||
for model_config in cases:
|
||||
self.assertFalse(
|
||||
uses_full_payload_input_coordinator(model_config),
|
||||
msg=f"expected gate OFF for {model_config}",
|
||||
)
|
||||
self.assertFalse(uses_qwen3_omni_full_payload_input_coordinator(model_config))
|
||||
|
||||
|
||||
class TestChunkCoordinatorStateTransition(unittest.TestCase):
|
||||
|
||||
@@ -13,7 +13,7 @@ DEFAULT_SHAPE = torch.Size([NUM_BLOCKS, BLOCK_SIZE, HIDDEN_SIZE])
|
||||
|
||||
|
||||
class MockInputBatch:
|
||||
def __init__(self, num_computed_tokens_cpu, block_table=None):
|
||||
def __init__(self, num_computed_tokens_cpu):
|
||||
self.req_ids = ["req1", "req2"]
|
||||
self.req_id_to_index = {req_id: i for i, req_id in enumerate(self.req_ids)}
|
||||
self.num_computed_tokens_cpu = num_computed_tokens_cpu
|
||||
@@ -21,33 +21,11 @@ class MockInputBatch:
|
||||
# Block table is only mocked for validation of length;
|
||||
# we don't actually need to add valid values here since
|
||||
# we patch the table when testing.
|
||||
if block_table is None:
|
||||
class _DummyBlockTable:
|
||||
pass
|
||||
|
||||
class _DummyBlockTable:
|
||||
pass
|
||||
|
||||
self.block_table = _DummyBlockTable()
|
||||
self.block_table.block_tables = [None]
|
||||
else:
|
||||
|
||||
class _TensorWrapper:
|
||||
def __init__(self, tensor):
|
||||
self.cpu = tensor
|
||||
|
||||
class _BlockGroup:
|
||||
def __init__(self, tensor):
|
||||
self.block_table = _TensorWrapper(tensor)
|
||||
|
||||
class _BlockTable:
|
||||
def __init__(self, tensor):
|
||||
self._group = _BlockGroup(tensor)
|
||||
self.block_tables = [self._group.block_table]
|
||||
|
||||
def __getitem__(self, idx):
|
||||
assert idx == 0
|
||||
return self._group
|
||||
|
||||
self.block_table = _BlockTable(block_table)
|
||||
self.block_table = _DummyBlockTable()
|
||||
self.block_table.block_tables = [None]
|
||||
|
||||
|
||||
def get_omni_pcache_with_mm_tensors(feat_dims, seq_len) -> OmniTensorPrefixCache:
|
||||
@@ -193,94 +171,6 @@ def test_update_with_multimodal_outputs(feat_dims):
|
||||
assert torch.all(slot_states == new_output)
|
||||
|
||||
|
||||
def test_deferred_multimodal_cache_write_commits_on_completion():
|
||||
cache = get_omni_pcache()
|
||||
mm_key = "codes.audio"
|
||||
block_table = torch.tensor(
|
||||
[
|
||||
[2, 3, 4, 5],
|
||||
[6, 7, 8, 9],
|
||||
],
|
||||
dtype=torch.long,
|
||||
)
|
||||
input_batch = MockInputBatch(
|
||||
num_computed_tokens_cpu=torch.tensor([4, 0], dtype=torch.long),
|
||||
block_table=block_table,
|
||||
)
|
||||
|
||||
first_chunk = torch.arange(8, dtype=torch.long).reshape(4, 2)
|
||||
cache.update_omni_tensor_prefix_cache(
|
||||
hidden_states=None,
|
||||
multimodal_outputs={mm_key: first_chunk},
|
||||
num_tokens_unpadded=4,
|
||||
slot_mapping=torch.arange(8, 12),
|
||||
skip_mm_cache_keys={mm_key},
|
||||
)
|
||||
|
||||
rows = cache.mm_outputs_cache[mm_key].view(-1, 2)
|
||||
assert torch.all(rows[8:12] == 0)
|
||||
|
||||
cache.stage_deferred_mm_outputs(
|
||||
query_start_loc=torch.tensor([0, 4], dtype=torch.long),
|
||||
input_batch=input_batch,
|
||||
multimodal_outputs={mm_key: first_chunk},
|
||||
num_scheduled_tokens={"req1": 4, "req2": 0},
|
||||
deferred_mm_cache_keys={mm_key},
|
||||
)
|
||||
|
||||
second_chunk = torch.arange(8, 16, dtype=torch.long).reshape(4, 2)
|
||||
input_batch.num_computed_tokens_cpu = torch.tensor([8, 0], dtype=torch.long)
|
||||
cache.stage_deferred_mm_outputs(
|
||||
query_start_loc=torch.tensor([0, 4], dtype=torch.long),
|
||||
input_batch=input_batch,
|
||||
multimodal_outputs={mm_key: second_chunk},
|
||||
num_scheduled_tokens={"req1": 4, "req2": 0},
|
||||
deferred_mm_cache_keys={mm_key},
|
||||
)
|
||||
|
||||
cache.commit_deferred_mm_outputs({"req1"}, input_batch)
|
||||
|
||||
expected = torch.cat([first_chunk, second_chunk], dim=0)
|
||||
rows = cache.mm_outputs_cache[mm_key].view(-1, 2)
|
||||
assert torch.equal(rows[8:16], expected)
|
||||
|
||||
|
||||
def test_deferred_multimodal_cache_can_be_merged_on_full_block_hit():
|
||||
cache = get_omni_pcache()
|
||||
mm_key = "codes.audio"
|
||||
block_table = torch.tensor(
|
||||
[
|
||||
[2, 3, 4, 5],
|
||||
[6, 7, 8, 9],
|
||||
],
|
||||
dtype=torch.long,
|
||||
)
|
||||
input_batch = MockInputBatch(
|
||||
num_computed_tokens_cpu=torch.tensor([8, 0], dtype=torch.long),
|
||||
block_table=block_table,
|
||||
)
|
||||
cached_codes = torch.arange(16, dtype=torch.long).reshape(8, 2)
|
||||
cache.stage_deferred_mm_outputs(
|
||||
query_start_loc=torch.tensor([0, 8], dtype=torch.long),
|
||||
input_batch=input_batch,
|
||||
multimodal_outputs={mm_key: cached_codes},
|
||||
num_scheduled_tokens={"req1": 8, "req2": 0},
|
||||
deferred_mm_cache_keys={mm_key},
|
||||
)
|
||||
cache.commit_deferred_mm_outputs({"req1"}, input_batch)
|
||||
|
||||
new_codes = torch.tensor([[99, 100]], dtype=torch.long)
|
||||
cache.add_prefix_cached_new_req_id("req1")
|
||||
merged = cache.get_merged_multimodal_states(
|
||||
query_start_loc=torch.tensor([0, 1], dtype=torch.long),
|
||||
input_batch=input_batch,
|
||||
multimodal_outputs={mm_key: new_codes},
|
||||
num_scheduled_tokens={"req1": 1, "req2": 0},
|
||||
)
|
||||
|
||||
assert torch.equal(merged[mm_key]["req1"], torch.cat([cached_codes, new_codes], dim=0))
|
||||
|
||||
|
||||
### Tests for Merging
|
||||
def fake_get_cached_block_ids(self, req_idx, *args, **kwargs):
|
||||
"""Fake block table lookup.
|
||||
|
||||
@@ -41,72 +41,6 @@ pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
os.environ.setdefault("DIFFUSION_ATTENTION_BACKEND", "FLASH_ATTN")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Inline field processing
|
||||
# ---------------------------------------------------------------------------
|
||||
def _process_inline_fields(obj: Any, parent_key: str = "") -> None:
|
||||
"""Recursively process '*-inline' fields into temp files."""
|
||||
|
||||
if isinstance(obj, list):
|
||||
for item in obj:
|
||||
_process_inline_fields(item, parent_key)
|
||||
return
|
||||
|
||||
if not isinstance(obj, dict):
|
||||
return
|
||||
|
||||
import atexit
|
||||
|
||||
import yaml
|
||||
|
||||
for key in list(obj.keys()):
|
||||
value = obj[key]
|
||||
|
||||
if not key.endswith("-inline"):
|
||||
_process_inline_fields(value, key)
|
||||
continue
|
||||
|
||||
base_key = key[:-7]
|
||||
full_key = f"{parent_key}.{key}" if parent_key else key
|
||||
|
||||
try:
|
||||
if not isinstance(value, dict):
|
||||
raise ValueError("must be a dict")
|
||||
|
||||
file_type = value.get("type")
|
||||
content = value.get("content")
|
||||
|
||||
if file_type not in {"yaml", "jsonl"}:
|
||||
raise ValueError(f"invalid type: {file_type}")
|
||||
|
||||
fd, path = tempfile.mkstemp(
|
||||
suffix=f".{file_type}",
|
||||
prefix=f"{base_key}_",
|
||||
)
|
||||
|
||||
atexit.register(Path(path).unlink, missing_ok=True)
|
||||
|
||||
with os.fdopen(fd, "w", encoding="utf-8") as f:
|
||||
if file_type == "jsonl":
|
||||
items = content if isinstance(content, list) else [content]
|
||||
f.writelines(json.dumps(x, ensure_ascii=False) + "\n" for x in items)
|
||||
else:
|
||||
yaml.dump(
|
||||
content,
|
||||
f,
|
||||
allow_unicode=True,
|
||||
sort_keys=False,
|
||||
indent=2,
|
||||
)
|
||||
|
||||
obj[base_key] = path
|
||||
del obj[key]
|
||||
|
||||
except Exception as e:
|
||||
print(f"Warning: failed processing '{full_key}': {e}")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Paths
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -168,14 +102,11 @@ def _resolve_refs(configs: list[dict[str, Any]], config_dir: Path) -> list[dict[
|
||||
|
||||
|
||||
def load_configs(config_path: str) -> list[dict[str, Any]]:
|
||||
"""Load benchmark configs from JSON file and process inline fields."""
|
||||
try:
|
||||
abs_path = Path(config_path).resolve()
|
||||
with open(abs_path, encoding="utf-8") as f:
|
||||
configs = json.load(f)
|
||||
configs = _resolve_refs(configs, abs_path.parent)
|
||||
_process_inline_fields(configs)
|
||||
return configs
|
||||
return _resolve_refs(configs, abs_path.parent)
|
||||
except json.JSONDecodeError as e:
|
||||
raise ValueError(f"JSON parsing error: {str(e)}")
|
||||
except FileNotFoundError:
|
||||
@@ -723,47 +654,9 @@ def run_benchmark(
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _resolve_baseline_value(
|
||||
baseline_raw: Any,
|
||||
*,
|
||||
sweep_index: int | None,
|
||||
max_concurrency: Any = None,
|
||||
request_rate: Any = None,
|
||||
) -> Any:
|
||||
"""Pick the baseline threshold for this sweep step."""
|
||||
if baseline_raw is None:
|
||||
return 100000
|
||||
if isinstance(baseline_raw, dict):
|
||||
if max_concurrency is not None:
|
||||
for key in (max_concurrency, str(max_concurrency)):
|
||||
if key in baseline_raw:
|
||||
return baseline_raw[key]
|
||||
if request_rate is not None:
|
||||
for key in (request_rate, str(request_rate)):
|
||||
if key in baseline_raw:
|
||||
return baseline_raw[key]
|
||||
raise KeyError(
|
||||
f"baseline dict has no key for max_concurrency={max_concurrency!r} "
|
||||
f"or request_rate={request_rate!r}; keys={list(baseline_raw.keys())!r}"
|
||||
)
|
||||
if isinstance(baseline_raw, (list, tuple)):
|
||||
if sweep_index is None:
|
||||
raise ValueError("sweep_index is required when baseline is a list or tuple")
|
||||
return baseline_raw[sweep_index]
|
||||
return baseline_raw
|
||||
|
||||
|
||||
def assert_result(
|
||||
result: dict[str, Any],
|
||||
params: dict[str, Any],
|
||||
num_prompts: int,
|
||||
*,
|
||||
sweep_index: int | None = None,
|
||||
max_concurrency: Any = None,
|
||||
request_rate: Any = None,
|
||||
assert_baseline: bool = True,
|
||||
) -> None:
|
||||
"""Assert that benchmark metrics satisfy the configured baselines."""
|
||||
def assert_result(result: dict[str, Any], params: dict[str, Any], *, assert_baseline: bool) -> None:
|
||||
"""Assert request completion; optionally assert metrics against ``baseline`` in *params*."""
|
||||
num_prompts = params.get("num-prompts", 10)
|
||||
completed = result.get("completed_requests", result.get("completed", 0))
|
||||
assert completed == num_prompts, f"Expected {num_prompts} completed requests, got {completed}"
|
||||
|
||||
@@ -773,15 +666,9 @@ def assert_result(
|
||||
print("Skipping performance assertions.")
|
||||
return
|
||||
|
||||
for metric, baseline_raw in params.get("baseline", {}).items():
|
||||
for metric, threshold in params.get("baseline", {}).items():
|
||||
current = result.get(metric)
|
||||
assert current is not None, f"Metric '{metric}' not found in result: {list(result.keys())}"
|
||||
threshold = _resolve_baseline_value(
|
||||
baseline_raw,
|
||||
sweep_index=sweep_index,
|
||||
max_concurrency=max_concurrency,
|
||||
request_rate=request_rate,
|
||||
)
|
||||
if "throughput" in metric:
|
||||
assert current >= threshold, f"{metric}: {current:.4f} < baseline {threshold}"
|
||||
else:
|
||||
@@ -805,111 +692,6 @@ def _resolve_benchmark_endpoint(server_cfg: dict[str, Any], params: dict[str, An
|
||||
return _default_benchmark_endpoint_for_task(cast(str, params.get("task", "t2i")))
|
||||
|
||||
|
||||
def _to_list(value: Any) -> list[Any]:
|
||||
if value is None:
|
||||
return []
|
||||
return [value] if not isinstance(value, (list, tuple)) else list(value)
|
||||
|
||||
|
||||
def _build_run_params(
|
||||
params: dict[str, Any],
|
||||
*,
|
||||
num_prompts: int,
|
||||
sweep_index: int | None = None,
|
||||
request_rate: Any | None = None,
|
||||
max_concurrency: Any | None = None,
|
||||
) -> dict[str, Any]:
|
||||
run_params = {
|
||||
key: value for key, value in params.items() if key not in {"request-rate", "max-concurrency", "num-prompts"}
|
||||
}
|
||||
run_params["num-prompts"] = num_prompts
|
||||
if request_rate is not None:
|
||||
run_params["request-rate"] = request_rate
|
||||
if max_concurrency is not None:
|
||||
run_params["max-concurrency"] = max_concurrency
|
||||
if "baseline" in params:
|
||||
run_params["baseline"] = {
|
||||
metric: _resolve_baseline_value(
|
||||
baseline_raw,
|
||||
sweep_index=sweep_index,
|
||||
max_concurrency=max_concurrency,
|
||||
request_rate=request_rate,
|
||||
)
|
||||
for metric, baseline_raw in params["baseline"].items()
|
||||
}
|
||||
return run_params
|
||||
|
||||
|
||||
def _iter_sweep_runs(params: dict[str, Any]) -> list[dict[str, Any]]:
|
||||
request_rate_list = _to_list(params.get("request-rate"))
|
||||
num_prompt_list = _to_list(params.get("num-prompts", 10))
|
||||
max_concurrency_list = _to_list(params.get("max-concurrency"))
|
||||
|
||||
max_len = max(len(request_rate_list), len(max_concurrency_list))
|
||||
if len(num_prompt_list) == 1 and max_len > 1:
|
||||
num_prompt_list = num_prompt_list * max_len
|
||||
elif max_len == 1 and len(num_prompt_list) > 1:
|
||||
if len(request_rate_list) == 1:
|
||||
request_rate_list = request_rate_list * len(num_prompt_list)
|
||||
if len(max_concurrency_list) == 1:
|
||||
max_concurrency_list = max_concurrency_list * len(num_prompt_list)
|
||||
max_len = max(len(request_rate_list), len(max_concurrency_list))
|
||||
elif len(num_prompt_list) != max_len and max_len > 0:
|
||||
raise ValueError("The number of prompts does not match the request-rate or max-concurrency")
|
||||
|
||||
sweep_runs: list[dict[str, Any]] = []
|
||||
|
||||
for i, (request_rate, num_prompts) in enumerate(zip(request_rate_list, num_prompt_list)):
|
||||
sweep_runs.append(
|
||||
{
|
||||
"params": _build_run_params(
|
||||
params,
|
||||
request_rate=request_rate,
|
||||
num_prompts=num_prompts,
|
||||
sweep_index=i,
|
||||
),
|
||||
"num_prompts": num_prompts,
|
||||
"sweep_index": i,
|
||||
"request_rate": request_rate,
|
||||
"max_concurrency": None,
|
||||
}
|
||||
)
|
||||
|
||||
for i, (max_concurrency, num_prompts) in enumerate(zip(max_concurrency_list, num_prompt_list)):
|
||||
sweep_runs.append(
|
||||
{
|
||||
"params": _build_run_params(
|
||||
params,
|
||||
max_concurrency=max_concurrency,
|
||||
num_prompts=num_prompts,
|
||||
sweep_index=i,
|
||||
request_rate="inf",
|
||||
),
|
||||
"num_prompts": num_prompts,
|
||||
"sweep_index": i,
|
||||
"request_rate": None,
|
||||
"max_concurrency": max_concurrency,
|
||||
}
|
||||
)
|
||||
|
||||
if not sweep_runs:
|
||||
default_num_prompts = num_prompt_list[0]
|
||||
sweep_runs.append(
|
||||
{
|
||||
"params": _build_run_params(
|
||||
params,
|
||||
num_prompts=default_num_prompts,
|
||||
),
|
||||
"num_prompts": default_num_prompts,
|
||||
"sweep_index": None,
|
||||
"request_rate": None,
|
||||
"max_concurrency": None,
|
||||
}
|
||||
)
|
||||
|
||||
return sweep_runs
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test entry point
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -933,47 +715,36 @@ def test_diffusion_performance_benchmark(diffusion_server, benchmark_params, req
|
||||
test_name = benchmark_params["test_name"]
|
||||
params = benchmark_params["params"]
|
||||
server_cfg = getattr(diffusion_server, "server_cfg", {})
|
||||
sweep_runs = _iter_sweep_runs(params)
|
||||
endpoint = _resolve_benchmark_endpoint(server_cfg, params)
|
||||
|
||||
for sweep_run in sweep_runs:
|
||||
endpoint = _resolve_benchmark_endpoint(server_cfg, sweep_run["params"])
|
||||
result = run_benchmark(
|
||||
host=diffusion_server.host,
|
||||
port=diffusion_server.port,
|
||||
model=diffusion_server.model,
|
||||
params=sweep_run["params"],
|
||||
test_name=test_name,
|
||||
endpoint=endpoint,
|
||||
server_cfg=server_cfg,
|
||||
source_file=cast(str, CONFIG_FILE_PATH),
|
||||
)
|
||||
result = run_benchmark(
|
||||
host=diffusion_server.host,
|
||||
port=diffusion_server.port,
|
||||
model=diffusion_server.model,
|
||||
params=params,
|
||||
test_name=test_name,
|
||||
endpoint=endpoint,
|
||||
server_cfg=server_cfg,
|
||||
source_file=cast(str, CONFIG_FILE_PATH),
|
||||
)
|
||||
|
||||
print(f"\n{'=' * 60}")
|
||||
print(f"Results for {test_name} (server={diffusion_server.server_type}, endpoint={endpoint}):")
|
||||
for key in (
|
||||
"throughput_qps",
|
||||
"latency_mean",
|
||||
"latency_median",
|
||||
"latency_p50",
|
||||
"latency_p99",
|
||||
"peak_memory_mb_max",
|
||||
"peak_memory_mb_mean",
|
||||
"peak_memory_mb_median",
|
||||
):
|
||||
if key in result:
|
||||
print(f" {key}: {result[key]:.4f}")
|
||||
print(f"\n{'=' * 60}")
|
||||
print(f"Results for {test_name} (server={diffusion_server.server_type}, endpoint={endpoint}):")
|
||||
for key in (
|
||||
"throughput_qps",
|
||||
"latency_mean",
|
||||
"latency_median",
|
||||
"latency_p50",
|
||||
"latency_p99",
|
||||
"peak_memory_mb_max",
|
||||
"peak_memory_mb_mean",
|
||||
"peak_memory_mb_median",
|
||||
):
|
||||
if key in result:
|
||||
print(f" {key}: {result[key]:.4f}")
|
||||
|
||||
print(f"\n Aggregated results: {AGGREGATED_RESULT_FILE}")
|
||||
print("=" * 60)
|
||||
print(f"\n Aggregated results: {AGGREGATED_RESULT_FILE}")
|
||||
print("=" * 60)
|
||||
|
||||
assert_baseline = request.config.getoption("--assert-baseline", default=False)
|
||||
|
||||
assert_result(
|
||||
result,
|
||||
params,
|
||||
sweep_run["num_prompts"],
|
||||
sweep_index=sweep_run["sweep_index"],
|
||||
max_concurrency=sweep_run["max_concurrency"],
|
||||
request_rate=sweep_run["request_rate"],
|
||||
assert_baseline=assert_baseline,
|
||||
)
|
||||
assert_baseline = request.config.getoption("--assert-baseline", default=False)
|
||||
assert_result(result, params, assert_baseline=assert_baseline)
|
||||
@@ -1,152 +0,0 @@
|
||||
[
|
||||
{
|
||||
"test_name": "test_bagel_single_device_single_stage_t2i",
|
||||
"description": "Single-stage BAGEL (TP=1, CFG=1), t2i Task",
|
||||
"server_type": "vllm-omni",
|
||||
"server_params": {
|
||||
"model": "ByteDance-Seed/BAGEL-7B-MoT",
|
||||
"serve_args": {
|
||||
"enable-diffusion-pipeline-profiler": true,
|
||||
"deploy-config": "../vllm_omni/deploy/bagel_single_stage.yaml",
|
||||
"no-async-chunk": true,
|
||||
"tensor-parallel-size": 1,
|
||||
"cfg-parallel-size": 1,
|
||||
"gpu-memory-utilization": 0.90,
|
||||
"max-model-len": 8192,
|
||||
"max-num-batched-tokens": 9216
|
||||
}
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"name": "512x512_t2i_steps20",
|
||||
"dataset": "random",
|
||||
"task": "t2i",
|
||||
"width": 512,
|
||||
"height": 512,
|
||||
"num-inference-steps": 20,
|
||||
"num-prompts": 20,
|
||||
"max-concurrency": 1,
|
||||
"enable-negative-prompt": true,
|
||||
"baseline": {
|
||||
"throughput_qps": 0.3678,
|
||||
"latency_mean": 2.691,
|
||||
"peak_memory_mb_max": 33083.6,
|
||||
"peak_memory_mb_mean": 33082.5
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"test_name": "test_bagel_single_device_single_stage_i2i",
|
||||
"description": "Single-stage BAGEL (TP=1, CFG=1), i2i Task",
|
||||
"server_type": "vllm-omni",
|
||||
"server_params": {
|
||||
"model": "ByteDance-Seed/BAGEL-7B-MoT",
|
||||
"serve_args": {
|
||||
"enable-diffusion-pipeline-profiler": true,
|
||||
"deploy-config": "../vllm_omni/deploy/bagel_single_stage.yaml",
|
||||
"no-async-chunk": true,
|
||||
"tensor-parallel-size": 1,
|
||||
"cfg-parallel-size": 1,
|
||||
"gpu-memory-utilization": 0.90,
|
||||
"max-model-len": 8192,
|
||||
"max-num-batched-tokens": 9216
|
||||
}
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"name": "512x512_i2i_steps20",
|
||||
"dataset": "random",
|
||||
"task": "i2i",
|
||||
"width": 512,
|
||||
"height": 512,
|
||||
"num-inference-steps": 20,
|
||||
"num-prompts": 20,
|
||||
"max-concurrency": 1,
|
||||
"enable-negative-prompt": true,
|
||||
"num-input-images": 1,
|
||||
"baseline": {
|
||||
"throughput_qps": 0.3124,
|
||||
"latency_mean": 3.1686,
|
||||
"peak_memory_mb_max": 33448.8,
|
||||
"peak_memory_mb_mean": 33447.7
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"test_name": "test_bagel_single_device_multi_stage_t2i",
|
||||
"description": "Two-stage BAGEL (bagel.yaml), TP=1, CFG=1, t2i Task",
|
||||
"server_type": "vllm-omni",
|
||||
"server_params": {
|
||||
"model": "ByteDance-Seed/BAGEL-7B-MoT",
|
||||
"serve_args": {
|
||||
"enable-diffusion-pipeline-profiler": true,
|
||||
"deploy-config": "../vllm_omni/deploy/bagel.yaml",
|
||||
"no-async-chunk": true,
|
||||
"tensor-parallel-size": 1,
|
||||
"cfg-parallel-size": 1,
|
||||
"gpu-memory-utilization": 0.45,
|
||||
"max-model-len": 8192,
|
||||
"max-num-batched-tokens": 9216
|
||||
}
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"name": "512x512_t2i_steps20",
|
||||
"dataset": "random",
|
||||
"task": "t2i",
|
||||
"width": 512,
|
||||
"height": 512,
|
||||
"num-inference-steps": 20,
|
||||
"num-prompts": 20,
|
||||
"max-concurrency": 1,
|
||||
"enable-negative-prompt": true,
|
||||
"baseline": {
|
||||
"throughput_qps": 0.3746,
|
||||
"latency_mean": 2.6425,
|
||||
"peak_memory_mb_max": 33077.0,
|
||||
"peak_memory_mb_mean": 33077.0
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"test_name": "test_bagel_single_device_multi_stage_i2i",
|
||||
"description": "Two-stage BAGEL (bagel.yaml), TP=1, CFG=1, i2i Task",
|
||||
"server_type": "vllm-omni",
|
||||
"server_params": {
|
||||
"model": "ByteDance-Seed/BAGEL-7B-MoT",
|
||||
"serve_args": {
|
||||
"enable-diffusion-pipeline-profiler": true,
|
||||
"deploy-config": "../vllm_omni/deploy/bagel.yaml",
|
||||
"no-async-chunk": true,
|
||||
"tensor-parallel-size": 1,
|
||||
"cfg-parallel-size": 1,
|
||||
"gpu-memory-utilization": 0.45,
|
||||
"max-model-len": 8192,
|
||||
"max-num-batched-tokens": 9216
|
||||
}
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"name": "512x512_i2i_steps20",
|
||||
"dataset": "random",
|
||||
"task": "i2i",
|
||||
"width": 512,
|
||||
"height": 512,
|
||||
"num-inference-steps": 20,
|
||||
"num-prompts": 20,
|
||||
"max-concurrency": 1,
|
||||
"enable-negative-prompt": true,
|
||||
"num-input-images": 1,
|
||||
"baseline": {
|
||||
"throughput_qps": 0.1349,
|
||||
"latency_mean": 7.339,
|
||||
"peak_memory_mb_max": 33321.2,
|
||||
"peak_memory_mb_mean": 33321.2
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
@@ -1,131 +0,0 @@
|
||||
[
|
||||
{
|
||||
"test_name": "test_hunyuan_image3_ti2i_tp4_tp4_inline",
|
||||
"description": "AR (TP=4) + DiT (TP=4) - using inline dataset and deploy-config",
|
||||
"server_type": "vllm-omni",
|
||||
"benchmark_endpoint": "/v1/images/edits",
|
||||
"server_params": {
|
||||
"model": "tencent/HunyuanImage-3.0-Instruct",
|
||||
"serve_args": {
|
||||
"deploy-config-inline": {
|
||||
"type": "yaml",
|
||||
"content": {
|
||||
"pipeline": "hunyuan_image_3_moe",
|
||||
"async_chunk": false,
|
||||
"trust_remote_code": true,
|
||||
"connectors": {
|
||||
"rdma_connector": {
|
||||
"name": "MooncakeTransferEngineConnector",
|
||||
"extra": {
|
||||
"host": "auto",
|
||||
"zmq_port": 50051,
|
||||
"protocol": "rdma",
|
||||
"device_name": "",
|
||||
"memory_pool_size": 4294967296,
|
||||
"memory_pool_device": "cpu"
|
||||
}
|
||||
}
|
||||
},
|
||||
"stages": [
|
||||
{
|
||||
"stage_id": 0,
|
||||
"is_comprehension": true,
|
||||
"final_output": true,
|
||||
"final_output_type": "text",
|
||||
"max_num_seqs": 4,
|
||||
"gpu_memory_utilization": 0.95,
|
||||
"enforce_eager": false,
|
||||
"enable_prefix_caching": false,
|
||||
"max_num_batched_tokens": 32768,
|
||||
"devices": "0,1,2,3",
|
||||
"tensor_parallel_size": 4,
|
||||
"hf_overrides": {
|
||||
"rope_parameters": {
|
||||
"mrope_section": [0, 32, 32],
|
||||
"rope_type": "default"
|
||||
}
|
||||
},
|
||||
"omni_kv_config": {
|
||||
"need_send_cache": true
|
||||
},
|
||||
"output_connectors": {
|
||||
"to_stage_1": "rdma_connector"
|
||||
},
|
||||
"default_sampling_params": {
|
||||
"temperature": 0.0,
|
||||
"top_p": 1,
|
||||
"top_k": -1,
|
||||
"max_tokens": 8192,
|
||||
"detokenize": true,
|
||||
"skip_special_tokens": false,
|
||||
"stop_token_ids": [128025]
|
||||
}
|
||||
},
|
||||
{
|
||||
"stage_id": 1,
|
||||
"max_num_seqs": 1,
|
||||
"enforce_eager": true,
|
||||
"devices": "4,5,6,7",
|
||||
"distributed_executor_backend": "mp",
|
||||
"omni_kv_config": {
|
||||
"need_recv_cache": true
|
||||
},
|
||||
"parallel_config": {
|
||||
"tensor_parallel_size": 4,
|
||||
"enable_expert_parallel": true
|
||||
},
|
||||
"input_connectors": {
|
||||
"from_stage_0": "rdma_connector"
|
||||
},
|
||||
"default_sampling_params": {
|
||||
"num_inference_steps": 50,
|
||||
"guidance_scale": 0
|
||||
}
|
||||
}
|
||||
],
|
||||
"edges": [
|
||||
{
|
||||
"from": 0,
|
||||
"to": 1,
|
||||
"window_size": -1,
|
||||
"max_inflight": 1
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
"enable-diffusion-pipeline-profiler": true
|
||||
}
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"name": "inline_dataset_5_requests",
|
||||
"dataset": "custom",
|
||||
"dataset-path-inline": {
|
||||
"type": "jsonl",
|
||||
"content": [
|
||||
{
|
||||
"prompt": "新年宠物海报，Q版圆润的可爱标题“新年快乐汪”，副标题“HAPPY NEW YEAR”。 鱼眼镜头，背景是房间门口，近景，上传的主体歪头笑，围着红色围巾，戴着红色毛线帽，高清，绒毛细节，面部特写。 宝丽莱相纸，超现实主义，写实主义，胶片摄影，打印颗粒感肌理。肌理，超写实，复古感。",
|
||||
"image_urls": ["https://raw.githubusercontent.com/Tencent-Hunyuan/HunyuanImage-3.0/main/assets/demo_instruct_imgs/input_0_0.png"],
|
||||
"num_inference_steps": 8,
|
||||
"seed": 42,
|
||||
"guidance_scale": 1.0
|
||||
}
|
||||
]
|
||||
},
|
||||
"task": "ti2i",
|
||||
"bot-task": "think_recaption",
|
||||
"max-concurrency": 4,
|
||||
"num-prompts": 16,
|
||||
"skip-performance-assertion": true,
|
||||
"baseline": {
|
||||
"throughput_qps": 0.1,
|
||||
"latency_p99": 44.46,
|
||||
"latency_mean": 35.57,
|
||||
"latency_std_dev": 0.5,
|
||||
"peak_memory_mb_max": 70000,
|
||||
"peak_memory_mb_mean": 70000
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
@@ -1,115 +0,0 @@
|
||||
[
|
||||
{
|
||||
"test_name": "test_qwen3_omni_3gpu_replica2",
|
||||
"server_params": {
|
||||
"model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
|
||||
"stage_overrides": {
|
||||
"1": {"num_replicas": 2, "devices": "1,2"},
|
||||
"2": {"num_replicas": 2, "devices": "1,2"}
|
||||
},
|
||||
"extra_cli_args": [
|
||||
"--async-chunk"
|
||||
]
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"dataset_name": "random",
|
||||
"backend": "openai-chat-omni",
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"num_prompts": [128, 128, 128, 128],
|
||||
"max_concurrency": [8, 16, 24, 32],
|
||||
"random_input_len": 2500,
|
||||
"random_output_len": 900,
|
||||
"ignore_eos": true,
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"mean_ttft_ms": [1000, 2000, 3000, 4000],
|
||||
"mean_tpot_ms": [200, 200, 200, 200],
|
||||
"mean_audio_ttfp_ms": [10000, 15000, 20000, 25000],
|
||||
"mean_audio_rtf": [1.5, 1.5, 1.5, 1.5]
|
||||
}
|
||||
},
|
||||
{
|
||||
"dataset_name": "random-mm",
|
||||
"backend": "openai-chat-omni",
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"num_prompts": [10],
|
||||
"request_rate": [0.1],
|
||||
"random_input_len": 100,
|
||||
"random_output_len": 100,
|
||||
"random_range_ratio": 0.0,
|
||||
"ignore_eos": true,
|
||||
"random_mm_base_items_per_request": 1,
|
||||
"random_mm_num_mm_items_range_ratio": 0.5,
|
||||
"random_mm_limit_mm_per_prompt": {
|
||||
"audio": 1
|
||||
},
|
||||
"random_mm_bucket_config": {
|
||||
"(0, 60, 3)": 1.0
|
||||
},
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"mean_ttft_ms": [2000],
|
||||
"mean_audio_ttfp_ms": [2000],
|
||||
"mean_audio_rtf": [0.25]
|
||||
}
|
||||
},
|
||||
{
|
||||
"dataset_name": "random-mm",
|
||||
"backend": "openai-chat-omni",
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"num_prompts": [40],
|
||||
"request_rate": [0.5],
|
||||
"random_input_len": 100,
|
||||
"random_output_len": 100,
|
||||
"random_range_ratio": 0.0,
|
||||
"ignore_eos": true,
|
||||
"random_mm_base_items_per_request": 2,
|
||||
"random_mm_num_mm_items_range_ratio": 0.5,
|
||||
"random_mm_limit_mm_per_prompt": {
|
||||
"image": 1,
|
||||
"video": 1
|
||||
},
|
||||
"random_mm_bucket_config": {
|
||||
"(256, 256, 1)": 0.5,
|
||||
"(720, 1280, 2)": 0.5
|
||||
},
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"mean_ttft_ms": [6000],
|
||||
"mean_audio_ttfp_ms": [6000],
|
||||
"mean_audio_rtf": [0.7]
|
||||
}
|
||||
},
|
||||
{
|
||||
"dataset_name": "random-mm",
|
||||
"backend": "openai-chat-omni",
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"num_prompts": [100],
|
||||
"request_rate": [1.0],
|
||||
"random_input_len": 100,
|
||||
"random_output_len": 100,
|
||||
"random_range_ratio": 0.0,
|
||||
"ignore_eos": true,
|
||||
"random_mm_base_items_per_request": 3,
|
||||
"random_mm_num_mm_items_range_ratio": 0.5,
|
||||
"random_mm_limit_mm_per_prompt": {
|
||||
"image": 1,
|
||||
"video": 1,
|
||||
"audio": 1
|
||||
},
|
||||
"random_mm_bucket_config": {
|
||||
"(256, 256, 1)": 0.34,
|
||||
"(720, 1280, 2)": 0.33,
|
||||
"(0, 60, 3)": 0.33
|
||||
},
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"mean_ttft_ms": [12000],
|
||||
"mean_audio_ttfp_ms": [12000],
|
||||
"mean_audio_rtf": [1.0]
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
+24
-6
@@ -1,10 +1,10 @@
|
||||
[
|
||||
{
|
||||
"test_name": "test_qwen_image_edit_2511_single_device",
|
||||
"test_name": "test_qwen_image_edit_2509_single_device",
|
||||
"description": "Single-device baseline (two input images)",
|
||||
"server_type": "vllm-omni",
|
||||
"server_params": {
|
||||
"model": "Qwen/Qwen-Image-Edit-2511",
|
||||
"model": "Qwen/Qwen-Image-Edit-2509",
|
||||
"serve_args": {
|
||||
"enable-diffusion-pipeline-profiler": true
|
||||
}
|
||||
@@ -49,11 +49,11 @@
|
||||
]
|
||||
},
|
||||
{
|
||||
"test_name": "test_qwen_image_edit_2511_ulysses2_cfg2_vae_patch4",
|
||||
"test_name": "test_qwen_image_edit_2509_ulysses2_cfg2_vae_patch4",
|
||||
"description": "Ulysses SP=2 + CFG=2 + VAE patch parallel=4",
|
||||
"server_type": "vllm-omni",
|
||||
"server_params": {
|
||||
"model": "Qwen/Qwen-Image-Edit-2511",
|
||||
"model": "Qwen/Qwen-Image-Edit-2509",
|
||||
"serve_args": {
|
||||
"ulysses-degree": 2,
|
||||
"cfg-parallel-size": 2,
|
||||
@@ -63,6 +63,24 @@
|
||||
}
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"name": "512x512_steps20_i2i_2img",
|
||||
"dataset": "random",
|
||||
"task": "i2i",
|
||||
"width": 512,
|
||||
"height": 512,
|
||||
"num-inference-steps": 20,
|
||||
"num-prompts": 10,
|
||||
"max-concurrency": 1,
|
||||
"num-input-images": 2,
|
||||
"enable-negative-prompt": true,
|
||||
"baseline": {
|
||||
"throughput_qps": 0.15,
|
||||
"latency_mean": 10,
|
||||
"peak_memory_mb_max": 55000,
|
||||
"peak_memory_mb_mean": 55000
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "1536x1536_steps35_i2i_2img",
|
||||
"dataset": "random",
|
||||
@@ -84,11 +102,11 @@
|
||||
]
|
||||
},
|
||||
{
|
||||
"test_name": "test_qwen_image_edit_2511_ulysses2_cfg2_cache_dit",
|
||||
"test_name": "test_qwen_image_edit_2509_ulysses2_cfg2_cache_dit",
|
||||
"description": "Ulysses SP=2 + CFG=2 + CacheDiT",
|
||||
"server_type": "vllm-omni",
|
||||
"server_params": {
|
||||
"model": "Qwen/Qwen-Image-Edit-2511",
|
||||
"model": "Qwen/Qwen-Image-Edit-2509",
|
||||
"serve_args": {
|
||||
"ulysses-degree": 2,
|
||||
"cfg-parallel-size": 2,
|
||||
@@ -61,6 +61,23 @@
|
||||
}
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"name": "512x512_steps20_i2i",
|
||||
"dataset": "random",
|
||||
"task": "i2i",
|
||||
"width": 512,
|
||||
"height": 512,
|
||||
"num-inference-steps": 20,
|
||||
"num-prompts": 10,
|
||||
"max-concurrency": 1,
|
||||
"enable-negative-prompt": true,
|
||||
"baseline": {
|
||||
"throughput_qps": 0.21,
|
||||
"latency_mean": 10,
|
||||
"peak_memory_mb_max": 56000,
|
||||
"peak_memory_mb_mean": 56000
|
||||
}
|
||||
},
|
||||
{
|
||||
"name": "1536x1536_steps35_i2i",
|
||||
"dataset": "random",
|
||||
|
||||
@@ -181,39 +181,5 @@
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"test_name": "test_qwen_image_single_device_step_execution_high_concurrency",
|
||||
"description": "Single-device Qwen-Image step-level execution benchmark at 512x512, 20 steps, concurrency sweep",
|
||||
"server_type": "vllm-omni",
|
||||
"server_params": {
|
||||
"model": "Qwen/Qwen-Image",
|
||||
"serve_args": {
|
||||
"enable-diffusion-pipeline-profiler": true,
|
||||
"step-execution": true,
|
||||
"max-num-seqs": 8
|
||||
}
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"name": "512x512_steps20_high_concurrency",
|
||||
"dataset": "random",
|
||||
"task": "t2i",
|
||||
"width": 512,
|
||||
"height": 512,
|
||||
"num-inference-steps": 20,
|
||||
"num-prompts": [20, 40, 80, 160],
|
||||
"max-concurrency": [1, 2, 4, 8],
|
||||
"warmup-requests": 8,
|
||||
"warmup-concurrency": 8,
|
||||
"warmup-num-inference-steps": 20,
|
||||
"enable-negative-prompt": true,
|
||||
"baseline": {
|
||||
"throughput_qps": [0.30, 0.50, 0.60, 0.70],
|
||||
"latency_mean": [3.50, 4.00, 7.00, 12.00],
|
||||
"peak_memory_mb_mean": [67000, 67000, 67000, 67000]
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
@@ -60,7 +60,6 @@ def _make_request(**overrides) -> OmniDiffusionRequest:
|
||||
output_type="pil",
|
||||
generator_device="cpu",
|
||||
),
|
||||
"request_id": "diffusers-backend",
|
||||
}
|
||||
defaults.update(overrides)
|
||||
return OmniDiffusionRequest(**defaults)
|
||||
|
||||
@@ -1,169 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for Bagel.forward_cache_update_vae packed-sequence assembly."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from PIL import Image
|
||||
from transformers import PretrainedConfig
|
||||
from vllm.transformers_utils.configs.bagel import BagelConfig
|
||||
|
||||
from vllm_omni.diffusion.models.bagel.bagel_transformer import (
|
||||
Bagel,
|
||||
BaseNavitOutputWithPast,
|
||||
NaiveCache,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
HIDDEN_SIZE = 32
|
||||
NUM_LAYERS = 2
|
||||
|
||||
|
||||
def _make_bagel_config() -> BagelConfig:
|
||||
llm_config = PretrainedConfig()
|
||||
llm_config.hidden_size = HIDDEN_SIZE
|
||||
llm_config.layer_module = "Qwen2MoTDecoderLayer"
|
||||
llm_config.num_hidden_layers = NUM_LAYERS
|
||||
llm_config.num_attention_heads = 4
|
||||
|
||||
vae_config = PretrainedConfig()
|
||||
vae_config.downsample = 8
|
||||
vae_config.z_channels = 4
|
||||
|
||||
return BagelConfig(
|
||||
llm_config=llm_config,
|
||||
vae_config=vae_config,
|
||||
visual_gen=True,
|
||||
visual_und=False,
|
||||
latent_patch_size=2,
|
||||
max_latent_size=32,
|
||||
)
|
||||
|
||||
|
||||
def _make_tracking_language_model() -> MagicMock:
|
||||
"""Return a mock LM that records forward() kwargs."""
|
||||
calls: list[dict] = []
|
||||
hidden_size = HIDDEN_SIZE
|
||||
|
||||
def forward(*_args, **kwargs):
|
||||
calls.append(dict(kwargs))
|
||||
if kwargs.get("return_embeddings_only"):
|
||||
text_ids = kwargs["packed_text_ids"]
|
||||
emb = torch.randn(text_ids.shape[0], hidden_size)
|
||||
return BaseNavitOutputWithPast(packed_query_sequence=emb)
|
||||
|
||||
packed_query_sequence = kwargs.get("packed_query_sequence")
|
||||
if packed_query_sequence is None:
|
||||
text_ids = kwargs["packed_text_ids"]
|
||||
packed_query_sequence = torch.randn(text_ids.shape[0], hidden_size)
|
||||
return BaseNavitOutputWithPast(
|
||||
packed_query_sequence=packed_query_sequence,
|
||||
past_key_values=kwargs.get("past_key_values"),
|
||||
)
|
||||
|
||||
lm = MagicMock()
|
||||
lm.forward = forward
|
||||
lm._calls = calls
|
||||
return lm
|
||||
|
||||
|
||||
def _vae_transform(image: Image.Image) -> torch.Tensor:
|
||||
arr = torch.from_numpy(np.array(image.convert("RGB"), dtype=np.float32))
|
||||
arr = arr / 127.5 - 1.0
|
||||
return arr.permute(2, 0, 1)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def bagel_and_vae_input():
|
||||
language_model = _make_tracking_language_model()
|
||||
bagel = Bagel(language_model=language_model, vit_model=None, config=_make_bagel_config())
|
||||
|
||||
image = Image.new("RGB", (64, 64), color=(128, 64, 32))
|
||||
new_token_ids = {"start_of_image": 100, "end_of_image": 101}
|
||||
|
||||
gen_input, _newlens, _new_rope = bagel.prepare_vae_images(
|
||||
curr_kvlens=[0],
|
||||
curr_rope=[0],
|
||||
images=[image],
|
||||
transforms=_vae_transform,
|
||||
new_token_ids=new_token_ids,
|
||||
)
|
||||
|
||||
vae_model = MagicMock()
|
||||
patch_shapes = gen_input["patchified_vae_latent_shapes"]
|
||||
p = bagel.latent_patch_size
|
||||
c = bagel.latent_channel
|
||||
|
||||
def encode(padded_images: torch.Tensor) -> torch.Tensor:
|
||||
# Match AutoEncoder.encode layout: [B, z_channels, h*p, w*p] per image.
|
||||
batch = padded_images.shape[0]
|
||||
assert batch == len(patch_shapes)
|
||||
latents = [torch.randn(c, h * p, w * p) for h, w in patch_shapes]
|
||||
return torch.stack(latents, dim=0)
|
||||
|
||||
vae_model.encode = encode
|
||||
|
||||
expected_seq_len = int(gen_input["packed_seqlens"].sum().item())
|
||||
num_text_tokens = int(gen_input["packed_text_ids"].numel())
|
||||
|
||||
return SimpleNamespace(
|
||||
bagel=bagel,
|
||||
language_model=language_model,
|
||||
vae_model=vae_model,
|
||||
gen_input=gen_input,
|
||||
past_key_values=NaiveCache(NUM_LAYERS),
|
||||
expected_seq_len=expected_seq_len,
|
||||
num_text_tokens=num_text_tokens,
|
||||
)
|
||||
|
||||
|
||||
class TestForwardCacheUpdateVae:
|
||||
"""Regression tests for single-stage img2img VAE cache prefill."""
|
||||
|
||||
def test_forwards_full_packed_query_sequence(self, bagel_and_vae_input):
|
||||
"""Update forward must use packed_query_sequence spanning text + VAE tokens."""
|
||||
ctx = bagel_and_vae_input
|
||||
assert ctx.expected_seq_len > ctx.num_text_tokens, "test setup should include many VAE tokens"
|
||||
|
||||
ctx.bagel.forward_cache_update_vae(
|
||||
ctx.vae_model,
|
||||
ctx.past_key_values,
|
||||
**ctx.gen_input,
|
||||
)
|
||||
|
||||
update_calls = [c for c in ctx.language_model._calls if not c.get("return_embeddings_only")]
|
||||
assert len(update_calls) == 1
|
||||
|
||||
call = update_calls[0]
|
||||
assert "packed_query_sequence" in call
|
||||
assert "packed_text_ids" not in call
|
||||
|
||||
packed = call["packed_query_sequence"]
|
||||
assert packed.shape[0] == ctx.expected_seq_len
|
||||
assert call.get("mode") == "gen"
|
||||
|
||||
def test_packed_sequence_covers_text_and_vae_index_ranges(self, bagel_and_vae_input):
|
||||
"""Text and VAE index ranges must lie within packed_query_sequence."""
|
||||
ctx = bagel_and_vae_input
|
||||
ctx.bagel.forward_cache_update_vae(
|
||||
ctx.vae_model,
|
||||
ctx.past_key_values,
|
||||
**ctx.gen_input,
|
||||
)
|
||||
|
||||
update_call = next(c for c in ctx.language_model._calls if not c.get("return_embeddings_only"))
|
||||
packed = update_call["packed_query_sequence"]
|
||||
|
||||
text_idx = ctx.gen_input["packed_text_indexes"]
|
||||
vae_idx = ctx.gen_input["packed_vae_token_indexes"]
|
||||
assert int(text_idx.max()) < packed.shape[0]
|
||||
assert int(vae_idx.max()) < packed.shape[0]
|
||||
assert int(text_idx.min()) >= 0
|
||||
assert int(vae_idx.min()) >= 0
|
||||
@@ -48,7 +48,6 @@ def _make_request(prompts=None, **sp_kwargs) -> OmniDiffusionRequest:
|
||||
return OmniDiffusionRequest(
|
||||
prompts=prompts or [{"prompt": "a cat dancing"}],
|
||||
sampling_params=sp,
|
||||
request_id="dmd2-sanitize",
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -48,7 +48,7 @@ def _make_pipeline(cls):
|
||||
|
||||
def _make_request(**sp_kwargs) -> OmniDiffusionRequest:
|
||||
sp = OmniDiffusionSamplingParams(**sp_kwargs)
|
||||
return OmniDiffusionRequest(prompts=[{"prompt": "a cat"}], sampling_params=sp, request_id="dmd2-scheduler")
|
||||
return OmniDiffusionRequest(prompts=[{"prompt": "a cat"}], sampling_params=sp)
|
||||
|
||||
|
||||
@pytest.fixture(
|
||||
|
||||
@@ -73,7 +73,7 @@ def test_should_apply_pe_respects_sampling_extra_args():
|
||||
|
||||
|
||||
def test_should_apply_pe_disables_dummy_warmup_request():
|
||||
req = SimpleNamespace(is_dummy_run=lambda: True, sampling_params=SimpleNamespace(extra_args={}))
|
||||
req = SimpleNamespace(request_ids=["dummy_req_id"], sampling_params=SimpleNamespace(extra_args={}))
|
||||
|
||||
assert ErnieImagePipeline._should_apply_pe(req) is False
|
||||
|
||||
|
||||
@@ -1,741 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for GLM-Image quantization support (W4A16/AutoRound).
|
||||
|
||||
These tests verify that the GLM-Image DiT transformer correctly accepts and uses
|
||||
quantization configs for W4A16/AutoRound quantization support.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from pytest_mock import MockerFixture
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
from vllm_omni.diffusion.models.glm_image.glm_image_transformer import (
|
||||
ColumnParallelGELU,
|
||||
ColumnParallelSiLU,
|
||||
GlmImageAdaLayerNormContinuous,
|
||||
GlmImageAdaLayerNormZero,
|
||||
GlmImageAttention,
|
||||
GlmImageFeedForward,
|
||||
GlmImageImageProjector,
|
||||
GlmImagePrepare,
|
||||
GlmImageRotaryPosEmbed,
|
||||
GlmImageTransformer2DModel,
|
||||
GlmImageTransformerBlock,
|
||||
_positive_divisors,
|
||||
validate_glm_image_tp_constraints,
|
||||
)
|
||||
from vllm_omni.model_executor.models.glm_image.pipeline import GLM_IMAGE_PIPELINE
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
@pytest.fixture(scope="function", autouse=True)
|
||||
def _init_distributed():
|
||||
"""Initialize the minimal distributed environment required by
|
||||
ReplicatedLinear (tensor-parallel group must exist)."""
|
||||
from vllm.distributed.parallel_state import (
|
||||
cleanup_dist_env_and_memory,
|
||||
init_distributed_environment,
|
||||
initialize_model_parallel,
|
||||
)
|
||||
|
||||
os.environ.setdefault("MASTER_ADDR", "localhost")
|
||||
os.environ.setdefault("MASTER_PORT", "29502")
|
||||
init_distributed_environment(
|
||||
world_size=1,
|
||||
rank=0,
|
||||
local_rank=0,
|
||||
distributed_init_method="env://",
|
||||
)
|
||||
initialize_model_parallel()
|
||||
yield
|
||||
cleanup_dist_env_and_memory()
|
||||
|
||||
|
||||
@pytest.fixture(scope="function", autouse=True)
|
||||
def _force_default_gemm(monkeypatch):
|
||||
"""Force CPU-compatible GEMM dispatch for tests using CPU tensors.
|
||||
|
||||
vLLM's dispatch_unquantized_gemm() selects the backend by platform
|
||||
(e.g. rocm_unquantized_gemm on AMD machines), not by tensor device.
|
||||
CPU test tensors crash with NotImplementedError on ROCm. Monkeypatch
|
||||
the dispatcher to always return the default (torch.nn.functional.linear)
|
||||
implementation which works on any device."""
|
||||
from vllm.model_executor.layers.utils import default_unquantized_gemm
|
||||
|
||||
monkeypatch.setattr(
|
||||
"vllm.model_executor.layers.linear.dispatch_unquantized_gemm",
|
||||
lambda: default_unquantized_gemm,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="function", autouse=True)
|
||||
def setup_mocks(mocker: MockerFixture):
|
||||
"""Set up common mocks for all tests."""
|
||||
mocker.patch(
|
||||
"vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size",
|
||||
return_value=1,
|
||||
)
|
||||
mock_get_tp_group = mocker.patch("vllm.distributed.parallel_state.get_tp_group")
|
||||
mock_tp_group = mocker.MagicMock()
|
||||
mock_tp_group.world_size = 1
|
||||
mock_get_tp_group.return_value = mock_tp_group
|
||||
yield
|
||||
|
||||
|
||||
class TestPositiveDivisors:
|
||||
"""Test _positive_divisors helper function."""
|
||||
|
||||
def test_divisors_of_1(self):
|
||||
assert _positive_divisors(1) == {1}
|
||||
|
||||
def test_divisors_of_12(self):
|
||||
assert _positive_divisors(12) == {1, 2, 3, 4, 6, 12}
|
||||
|
||||
def test_divisors_of_prime(self):
|
||||
assert _positive_divisors(7) == {1, 7}
|
||||
|
||||
def test_divisors_of_0_returns_empty(self):
|
||||
assert _positive_divisors(0) == set()
|
||||
|
||||
def test_divisors_of_negative_returns_empty(self):
|
||||
assert _positive_divisors(-5) == set()
|
||||
|
||||
|
||||
class TestValidateGlmImageTpConstraints:
|
||||
"""Test TP constraint validation for GLM-Image."""
|
||||
|
||||
def test_valid_tp_size_1(self):
|
||||
"""TP=1 should always be valid."""
|
||||
result = validate_glm_image_tp_constraints(
|
||||
dim=2560,
|
||||
num_heads=64,
|
||||
ffn_hidden_dim=10240,
|
||||
tensor_parallel_size=1,
|
||||
)
|
||||
assert 1 in result
|
||||
|
||||
def test_valid_tp_size_2_for_divisible_dim(self):
|
||||
"""TP=2 is valid when all dims are divisible by 2."""
|
||||
result = validate_glm_image_tp_constraints(
|
||||
dim=2560,
|
||||
num_heads=64,
|
||||
ffn_hidden_dim=10240,
|
||||
tensor_parallel_size=2,
|
||||
)
|
||||
assert 1 in result
|
||||
assert 2 in result
|
||||
|
||||
def test_valid_tp_size_4_for_divisible_dim(self):
|
||||
"""TP=4 is valid when all dims are divisible by 4."""
|
||||
result = validate_glm_image_tp_constraints(
|
||||
dim=2560,
|
||||
num_heads=64,
|
||||
ffn_hidden_dim=10240,
|
||||
tensor_parallel_size=4,
|
||||
)
|
||||
assert 1 in result
|
||||
assert 2 in result
|
||||
assert 4 in result
|
||||
|
||||
def test_invalid_tp_size_3_for_divisible_dim(self):
|
||||
"""TP=3 is invalid when dim is not divisible by 3."""
|
||||
with pytest.raises(ValueError, match="dim % tensor_parallel_size == 0"):
|
||||
validate_glm_image_tp_constraints(
|
||||
dim=2560, # 2560 % 3 != 0
|
||||
num_heads=64,
|
||||
ffn_hidden_dim=10240,
|
||||
tensor_parallel_size=3,
|
||||
)
|
||||
|
||||
def test_invalid_tp_size_zero(self):
|
||||
"""TP=0 should raise error."""
|
||||
with pytest.raises(ValueError, match="tensor_parallel_size must be > 0"):
|
||||
validate_glm_image_tp_constraints(
|
||||
dim=2560,
|
||||
num_heads=64,
|
||||
ffn_hidden_dim=10240,
|
||||
tensor_parallel_size=0,
|
||||
)
|
||||
|
||||
def test_invalid_tp_size_negative(self):
|
||||
"""Negative TP size should raise error."""
|
||||
with pytest.raises(ValueError, match="tensor_parallel_size must be > 0"):
|
||||
validate_glm_image_tp_constraints(
|
||||
dim=2560,
|
||||
num_heads=64,
|
||||
ffn_hidden_dim=10240,
|
||||
tensor_parallel_size=-1,
|
||||
)
|
||||
|
||||
|
||||
class TestGlmImageAdaLayerNormZeroQuantization:
|
||||
"""Test GlmImageAdaLayerNormZero with quantization config."""
|
||||
|
||||
def test_accepts_quant_config_parameter(self, mocker: MockerFixture):
|
||||
"""Verify the class accepts quant_config parameter."""
|
||||
mock_quant_config = mocker.MagicMock()
|
||||
layer = GlmImageAdaLayerNormZero(
|
||||
embedding_dim=512,
|
||||
dim=2560,
|
||||
quant_config=mock_quant_config,
|
||||
prefix="test.norm1",
|
||||
)
|
||||
assert layer.linear.quant_config is mock_quant_config
|
||||
|
||||
def test_accepts_none_quant_config(self):
|
||||
"""Verify quant_config=None is accepted."""
|
||||
layer = GlmImageAdaLayerNormZero(
|
||||
embedding_dim=512,
|
||||
dim=2560,
|
||||
quant_config=None,
|
||||
)
|
||||
assert layer.linear.quant_config is None
|
||||
|
||||
def test_forward_handles_tuple_return_from_linear(self):
|
||||
"""Verify forward handles tuple returns from ReplicatedLinear."""
|
||||
layer = GlmImageAdaLayerNormZero(
|
||||
embedding_dim=512,
|
||||
dim=2560,
|
||||
quant_config=None,
|
||||
)
|
||||
|
||||
batch_size = 2
|
||||
seq_len = 10
|
||||
hidden_states = torch.randn(batch_size, seq_len, 2560)
|
||||
encoder_hidden_states = torch.randn(batch_size, seq_len, 2560)
|
||||
temb = torch.randn(batch_size, 512)
|
||||
|
||||
# This should work regardless of whether linear returns tuple or tensor
|
||||
result = layer(hidden_states, encoder_hidden_states, temb)
|
||||
assert len(result) == 10 # Should return 10 chunks
|
||||
|
||||
|
||||
class TestGlmImageAdaLayerNormContinuousQuantization:
|
||||
"""Test GlmImageAdaLayerNormContinuous with quantization config."""
|
||||
|
||||
def test_accepts_quant_config_parameter(self, mocker: MockerFixture):
|
||||
"""Verify the class accepts quant_config parameter."""
|
||||
mock_quant_config = mocker.MagicMock()
|
||||
layer = GlmImageAdaLayerNormContinuous(
|
||||
embedding_dim=2560,
|
||||
conditioning_embedding_dim=512,
|
||||
quant_config=mock_quant_config,
|
||||
prefix="test.norm_out",
|
||||
)
|
||||
assert layer.linear.quant_config is mock_quant_config
|
||||
|
||||
def test_accepts_none_quant_config(self):
|
||||
"""Verify quant_config=None is accepted."""
|
||||
layer = GlmImageAdaLayerNormContinuous(
|
||||
embedding_dim=2560,
|
||||
conditioning_embedding_dim=512,
|
||||
quant_config=None,
|
||||
)
|
||||
assert layer.linear.quant_config is None
|
||||
|
||||
def test_forward_handles_tuple_return_from_linear(self):
|
||||
"""Verify forward handles tuple returns from ReplicatedLinear."""
|
||||
layer = GlmImageAdaLayerNormContinuous(
|
||||
embedding_dim=2560,
|
||||
conditioning_embedding_dim=512,
|
||||
quant_config=None,
|
||||
)
|
||||
|
||||
batch_size = 2
|
||||
seq_len = 10
|
||||
x = torch.randn(batch_size, seq_len, 2560)
|
||||
conditioning_embedding = torch.randn(batch_size, 512)
|
||||
|
||||
result = layer(x, conditioning_embedding)
|
||||
assert result.shape == (batch_size, seq_len, 2560)
|
||||
|
||||
|
||||
class TestGlmImageAttentionQuantization:
|
||||
"""Test GlmImageAttention with quantization config."""
|
||||
|
||||
def test_accepts_quant_config_parameter(self, mocker: MockerFixture):
|
||||
"""Verify GlmImageAttention accepts quant_config parameter."""
|
||||
mock_quant_config = mocker.MagicMock()
|
||||
attn = GlmImageAttention(
|
||||
dim=2560,
|
||||
num_heads=64,
|
||||
head_dim=40,
|
||||
quant_config=mock_quant_config,
|
||||
prefix="test.attn1",
|
||||
)
|
||||
assert attn.to_qkv.quant_config is mock_quant_config
|
||||
|
||||
def test_accepts_none_quant_config(self):
|
||||
"""Verify quant_config=None is accepted."""
|
||||
attn = GlmImageAttention(
|
||||
dim=2560,
|
||||
num_heads=64,
|
||||
head_dim=40,
|
||||
quant_config=None,
|
||||
)
|
||||
assert attn.to_qkv.quant_config is None
|
||||
|
||||
|
||||
class TestColumnParallelModulesQuantization:
|
||||
"""Test ColumnParallelGELU and ColumnParallelSiLU with quantization config."""
|
||||
|
||||
def test_column_parallel_gelu_accepts_quant_config(self, mocker: MockerFixture):
|
||||
"""Verify ColumnParallelGELU accepts quant_config."""
|
||||
mock_quant_config = mocker.MagicMock()
|
||||
layer = ColumnParallelGELU(
|
||||
dim_in=2560,
|
||||
dim_out=10240,
|
||||
quant_config=mock_quant_config,
|
||||
prefix="test.gelu",
|
||||
)
|
||||
assert layer.proj.quant_config is mock_quant_config
|
||||
|
||||
def test_column_parallel_silu_accepts_quant_config(self, mocker: MockerFixture):
|
||||
"""Verify ColumnParallelSiLU accepts quant_config."""
|
||||
mock_quant_config = mocker.MagicMock()
|
||||
layer = ColumnParallelSiLU(
|
||||
dim_in=2560,
|
||||
dim_out=10240,
|
||||
quant_config=mock_quant_config,
|
||||
prefix="test.silu",
|
||||
)
|
||||
assert layer.proj.quant_config is mock_quant_config
|
||||
|
||||
|
||||
class TestGlmImageFeedForwardQuantization:
|
||||
"""Test GlmImageFeedForward with quantization config."""
|
||||
|
||||
def test_accepts_quant_config_parameter(self, mocker: MockerFixture):
|
||||
"""Verify GlmImageFeedForward accepts quant_config parameter."""
|
||||
mock_quant_config = mocker.MagicMock()
|
||||
ff = GlmImageFeedForward(
|
||||
dim=2560,
|
||||
dim_out=2560,
|
||||
inner_dim=10240,
|
||||
activation_fn="gelu-approximate",
|
||||
quant_config=mock_quant_config,
|
||||
prefix="test.ff",
|
||||
)
|
||||
# Check that the first layer (ColumnParallelGELU) has quant_config
|
||||
gelu_layer = ff.net[0]
|
||||
assert gelu_layer.proj.quant_config is mock_quant_config
|
||||
|
||||
def test_accepts_none_quant_config(self):
|
||||
"""Verify quant_config=None is accepted."""
|
||||
ff = GlmImageFeedForward(
|
||||
dim=2560,
|
||||
dim_out=2560,
|
||||
inner_dim=10240,
|
||||
activation_fn="gelu-approximate",
|
||||
quant_config=None,
|
||||
)
|
||||
gelu_layer = ff.net[0]
|
||||
assert gelu_layer.proj.quant_config is None
|
||||
|
||||
def test_linear_silu_activation(self):
|
||||
"""Test linear-silu activation function initialization works."""
|
||||
# This test verifies the module can be instantiated with linear-silu activation.
|
||||
# Full forward testing requires proper TP group setup, so we just verify construction.
|
||||
ff = GlmImageFeedForward(
|
||||
dim=2560,
|
||||
dim_out=2560,
|
||||
inner_dim=10240,
|
||||
activation_fn="linear-silu",
|
||||
quant_config=None,
|
||||
)
|
||||
# Verify the FFN has the correct structure
|
||||
assert len(ff.net) == 3 # ColumnParallelSiLU, Identity, RowParallelLinear
|
||||
assert isinstance(ff.net[0], ColumnParallelSiLU)
|
||||
|
||||
|
||||
class TestGlmImageTransformerBlockQuantization:
|
||||
"""Test GlmImageTransformerBlock with quantization config."""
|
||||
|
||||
def test_accepts_quant_config_parameter(self, mocker: MockerFixture):
|
||||
"""Verify GlmImageTransformerBlock accepts quant_config parameter."""
|
||||
mock_quant_config = mocker.MagicMock()
|
||||
parallel_config = DiffusionParallelConfig(
|
||||
tensor_parallel_size=1,
|
||||
sequence_parallel_size=1,
|
||||
)
|
||||
block = GlmImageTransformerBlock(
|
||||
dim=2560,
|
||||
num_attention_heads=64,
|
||||
attention_head_dim=40,
|
||||
time_embed_dim=512,
|
||||
parallel_config=parallel_config,
|
||||
quant_config=mock_quant_config,
|
||||
prefix="test.block",
|
||||
)
|
||||
# Check that inner modules have quant_config
|
||||
assert block.norm1.linear.quant_config is mock_quant_config
|
||||
assert block.attn1.to_qkv.quant_config is mock_quant_config
|
||||
assert block.ff.net[0].proj.quant_config is mock_quant_config
|
||||
|
||||
def test_accepts_none_quant_config(self):
|
||||
"""Verify quant_config=None is accepted."""
|
||||
block = GlmImageTransformerBlock(
|
||||
dim=2560,
|
||||
num_attention_heads=64,
|
||||
attention_head_dim=40,
|
||||
time_embed_dim=512,
|
||||
quant_config=None,
|
||||
)
|
||||
assert block.norm1.linear.quant_config is None
|
||||
|
||||
|
||||
class TestGlmImagePrepareModule:
|
||||
"""Test GlmImagePrepare module."""
|
||||
|
||||
def test_prepare_module_exists(self):
|
||||
"""Verify GlmImagePrepare module exists and works."""
|
||||
projector = GlmImageImageProjector(in_channels=16, hidden_size=2560, patch_size=2)
|
||||
rope = GlmImageRotaryPosEmbed(dim=40, patch_size=2)
|
||||
prepare = GlmImagePrepare(
|
||||
image_projector=projector,
|
||||
rope=rope,
|
||||
patch_size=2,
|
||||
)
|
||||
|
||||
# Test forward pass
|
||||
hidden_states = torch.randn(1, 16, 64, 64) # [B, C, H, W]
|
||||
result = prepare(hidden_states)
|
||||
|
||||
assert len(result) == 5
|
||||
hidden_out, rope_cos, rope_sin, height, width = result
|
||||
assert hidden_out.shape[0] == 1 # batch size
|
||||
assert hidden_out.shape[1] == 1024 # seq_len (32 * 32)
|
||||
assert hidden_out.shape[2] == 2560 # hidden_dim
|
||||
|
||||
|
||||
class TestGlmImagePipelineConfig:
|
||||
"""Test GLM_IMAGE_PIPELINE configuration."""
|
||||
|
||||
def test_glm_image_pipeline_config_exists(self):
|
||||
"""Verify GLM_IMAGE_PIPELINE config exists."""
|
||||
assert GLM_IMAGE_PIPELINE is not None
|
||||
|
||||
def test_pipeline_has_correct_model_type(self):
|
||||
"""Verify pipeline has correct model_type."""
|
||||
assert GLM_IMAGE_PIPELINE.model_type == "glm_image"
|
||||
|
||||
def test_pipeline_has_correct_model_arch(self):
|
||||
"""Verify pipeline has correct model_arch."""
|
||||
assert GLM_IMAGE_PIPELINE.model_arch == "GlmImageForConditionalGeneration"
|
||||
|
||||
def test_pipeline_has_two_stages(self):
|
||||
"""Verify pipeline has two stages (AR + DiT)."""
|
||||
assert len(GLM_IMAGE_PIPELINE.stages) == 2
|
||||
|
||||
def test_stage_0_is_ar_llm(self):
|
||||
"""Verify stage 0 is the AR LLM stage."""
|
||||
stage_0 = GLM_IMAGE_PIPELINE.stages[0]
|
||||
assert stage_0.stage_id == 0
|
||||
assert stage_0.model_stage == "ar"
|
||||
assert stage_0.execution_type.value == "llm_ar"
|
||||
assert stage_0.owns_tokenizer is True
|
||||
|
||||
def test_stage_1_is_diffusion(self):
|
||||
"""Verify stage 1 is the DiT diffusion stage."""
|
||||
stage_1 = GLM_IMAGE_PIPELINE.stages[1]
|
||||
assert stage_1.stage_id == 1
|
||||
assert stage_1.model_stage == "dit"
|
||||
assert stage_1.execution_type.value == "diffusion"
|
||||
assert stage_1.final_output is True
|
||||
assert stage_1.final_output_type == "image"
|
||||
assert stage_1.input_sources == (0,) # Takes input from stage 0
|
||||
|
||||
def test_pipeline_has_diffusers_class_name(self):
|
||||
"""Verify pipeline has diffusers_class_name."""
|
||||
assert GLM_IMAGE_PIPELINE.diffusers_class_name == "GlmImagePipeline"
|
||||
|
||||
|
||||
class TestGlmImageImageProjector:
|
||||
"""Test GlmImageImageProjector module."""
|
||||
|
||||
def test_projector_output_shape(self):
|
||||
"""Verify projector produces correct output shape."""
|
||||
projector = GlmImageImageProjector(
|
||||
in_channels=16,
|
||||
hidden_size=2560,
|
||||
patch_size=2,
|
||||
)
|
||||
|
||||
# Input: [B, C, H, W] = [1, 16, 64, 64]
|
||||
hidden_states = torch.randn(1, 16, 64, 64)
|
||||
output = projector(hidden_states)
|
||||
|
||||
# After patchify: [B, H/2*W/2, D] = [1, 32*32, 2560] = [1, 1024, 2560]
|
||||
assert output.shape == (1, 1024, 2560)
|
||||
|
||||
def test_projector_with_different_patch_sizes(self):
|
||||
"""Verify projector works with different patch sizes."""
|
||||
for patch_size in [2, 4]:
|
||||
projector = GlmImageImageProjector(
|
||||
in_channels=16,
|
||||
hidden_size=2560,
|
||||
patch_size=patch_size,
|
||||
)
|
||||
h, w = 64, 64
|
||||
hidden_states = torch.randn(1, 16, h, w)
|
||||
output = projector(hidden_states)
|
||||
expected_seq_len = (h // patch_size) * (w // patch_size)
|
||||
assert output.shape[1] == expected_seq_len
|
||||
|
||||
|
||||
class TestGlmImageRotaryPosEmbed:
|
||||
"""Test GlmImageRotaryPosEmbed module."""
|
||||
|
||||
def test_rope_output_shape(self):
|
||||
"""Verify RoPE produces correct output shape."""
|
||||
rope = GlmImageRotaryPosEmbed(dim=40, patch_size=2, theta=10000.0)
|
||||
hidden_states = torch.randn(1, 16, 64, 64) # [B, C, H, W]
|
||||
|
||||
cos, sin = rope(hidden_states)
|
||||
|
||||
# After patchify: height=32, width=32
|
||||
# Output: [32*32, dim] = [1024, 40]
|
||||
assert cos.shape[0] == 1024
|
||||
assert cos.shape[1] == 40
|
||||
assert sin.shape == cos.shape
|
||||
|
||||
def test_rope_value_range(self):
|
||||
"""Verify RoPE produces values in valid range."""
|
||||
rope = GlmImageRotaryPosEmbed(dim=40, patch_size=2, theta=10000.0)
|
||||
hidden_states = torch.randn(2, 16, 32, 32)
|
||||
|
||||
cos, sin = rope(hidden_states)
|
||||
|
||||
# cos and sin should be in [-1, 1]
|
||||
assert cos.min() >= -1.0 and cos.max() <= 1.0
|
||||
assert sin.min() >= -1.0 and sin.max() <= 1.0
|
||||
|
||||
def test_rope_consistency(self):
|
||||
"""Verify RoPE is consistent for same input."""
|
||||
rope = GlmImageRotaryPosEmbed(dim=40, patch_size=2, theta=10000.0)
|
||||
hidden_states = torch.randn(1, 16, 32, 32)
|
||||
|
||||
cos1, sin1 = rope(hidden_states)
|
||||
cos2, sin2 = rope(hidden_states)
|
||||
|
||||
# Same input should produce same output
|
||||
assert torch.allclose(cos1, cos2)
|
||||
assert torch.allclose(sin1, sin2)
|
||||
|
||||
|
||||
class TestGlmImageTransformer2DModelQuantization:
|
||||
"""Test GlmImageTransformer2DModel with quantization config."""
|
||||
|
||||
def test_accepts_quant_config_parameter(self, mocker: MockerFixture):
|
||||
"""Verify the model accepts quant_config parameter."""
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
|
||||
mock_quant_config = mocker.MagicMock()
|
||||
parallel_config = DiffusionParallelConfig(
|
||||
tensor_parallel_size=1,
|
||||
sequence_parallel_size=1,
|
||||
)
|
||||
|
||||
# Create a minimal mock od_config
|
||||
mock_tf_config = mocker.MagicMock()
|
||||
mock_tf_config.patch_size = 2
|
||||
mock_tf_config.in_channels = 16
|
||||
mock_tf_config.out_channels = 16
|
||||
mock_tf_config.num_attention_heads = 64
|
||||
mock_tf_config.attention_head_dim = 40
|
||||
mock_tf_config.time_embed_dim = 512
|
||||
mock_tf_config.condition_dim = 256
|
||||
mock_tf_config.prior_vq_quantizer_codebook_size = 16384
|
||||
mock_tf_config.text_embed_dim = 1024
|
||||
mock_tf_config.num_layers = 2 # Small number for testing
|
||||
|
||||
mock_od_config = mocker.MagicMock(spec=OmniDiffusionConfig)
|
||||
mock_od_config.tf_model_config = mock_tf_config
|
||||
mock_od_config.parallel_config = parallel_config
|
||||
|
||||
model = GlmImageTransformer2DModel(
|
||||
od_config=mock_od_config,
|
||||
quant_config=mock_quant_config,
|
||||
)
|
||||
|
||||
# Check that quantization config was passed to transformer blocks
|
||||
for block in model.transformer_blocks:
|
||||
assert block.norm1.linear.quant_config is mock_quant_config
|
||||
assert block.attn1.to_qkv.quant_config is mock_quant_config
|
||||
assert block.ff.net[0].proj.quant_config is mock_quant_config
|
||||
|
||||
def test_accepts_none_quant_config(self, mocker: MockerFixture):
|
||||
"""Verify quant_config=None is accepted."""
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
|
||||
parallel_config = DiffusionParallelConfig(
|
||||
tensor_parallel_size=1,
|
||||
sequence_parallel_size=1,
|
||||
)
|
||||
|
||||
mock_tf_config = mocker.MagicMock()
|
||||
mock_tf_config.patch_size = 2
|
||||
mock_tf_config.in_channels = 16
|
||||
mock_tf_config.out_channels = 16
|
||||
mock_tf_config.num_attention_heads = 64
|
||||
mock_tf_config.attention_head_dim = 40
|
||||
mock_tf_config.time_embed_dim = 512
|
||||
mock_tf_config.condition_dim = 256
|
||||
mock_tf_config.prior_vq_quantizer_codebook_size = 16384
|
||||
mock_tf_config.text_embed_dim = 1024
|
||||
mock_tf_config.num_layers = 2
|
||||
|
||||
mock_od_config = mocker.MagicMock(spec=OmniDiffusionConfig)
|
||||
mock_od_config.tf_model_config = mock_tf_config
|
||||
mock_od_config.parallel_config = parallel_config
|
||||
|
||||
model = GlmImageTransformer2DModel(
|
||||
od_config=mock_od_config,
|
||||
quant_config=None,
|
||||
)
|
||||
|
||||
# Check that quantization config is None
|
||||
for block in model.transformer_blocks:
|
||||
assert block.norm1.linear.quant_config is None
|
||||
assert block.attn1.to_qkv.quant_config is None
|
||||
|
||||
def test_norm_out_has_no_quantization(self, mocker: MockerFixture):
|
||||
"""Verify norm_out (output layer) does NOT use quantization to preserve precision."""
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
|
||||
mock_quant_config = mocker.MagicMock()
|
||||
parallel_config = DiffusionParallelConfig(
|
||||
tensor_parallel_size=1,
|
||||
sequence_parallel_size=1,
|
||||
)
|
||||
|
||||
mock_tf_config = mocker.MagicMock()
|
||||
mock_tf_config.patch_size = 2
|
||||
mock_tf_config.in_channels = 16
|
||||
mock_tf_config.out_channels = 16
|
||||
mock_tf_config.num_attention_heads = 64
|
||||
mock_tf_config.attention_head_dim = 40
|
||||
mock_tf_config.time_embed_dim = 512
|
||||
mock_tf_config.condition_dim = 256
|
||||
mock_tf_config.prior_vq_quantizer_codebook_size = 16384
|
||||
mock_tf_config.text_embed_dim = 1024
|
||||
mock_tf_config.num_layers = 2
|
||||
|
||||
mock_od_config = mocker.MagicMock(spec=OmniDiffusionConfig)
|
||||
mock_od_config.tf_model_config = mock_tf_config
|
||||
mock_od_config.parallel_config = parallel_config
|
||||
|
||||
model = GlmImageTransformer2DModel(
|
||||
od_config=mock_od_config,
|
||||
quant_config=mock_quant_config,
|
||||
)
|
||||
|
||||
# norm_out.linear should NOT have quant_config to preserve output precision
|
||||
assert model.norm_out.linear.quant_config is None
|
||||
|
||||
def test_model_has_sp_plan(self, mocker: MockerFixture):
|
||||
"""Verify model has _sp_plan defined for sequence parallelism."""
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
|
||||
parallel_config = DiffusionParallelConfig(
|
||||
tensor_parallel_size=1,
|
||||
sequence_parallel_size=1,
|
||||
)
|
||||
|
||||
mock_tf_config = mocker.MagicMock()
|
||||
mock_tf_config.patch_size = 2
|
||||
mock_tf_config.in_channels = 16
|
||||
mock_tf_config.out_channels = 16
|
||||
mock_tf_config.num_attention_heads = 64
|
||||
mock_tf_config.attention_head_dim = 40
|
||||
mock_tf_config.time_embed_dim = 512
|
||||
mock_tf_config.condition_dim = 256
|
||||
mock_tf_config.prior_vq_quantizer_codebook_size = 16384
|
||||
mock_tf_config.text_embed_dim = 1024
|
||||
mock_tf_config.num_layers = 2
|
||||
|
||||
mock_od_config = mocker.MagicMock(spec=OmniDiffusionConfig)
|
||||
mock_od_config.tf_model_config = mock_tf_config
|
||||
mock_od_config.parallel_config = parallel_config
|
||||
|
||||
model = GlmImageTransformer2DModel(
|
||||
od_config=mock_od_config,
|
||||
quant_config=None,
|
||||
)
|
||||
|
||||
# Verify _sp_plan exists
|
||||
assert hasattr(model, "_sp_plan")
|
||||
assert "prepare" in model._sp_plan
|
||||
assert "proj_out" in model._sp_plan
|
||||
|
||||
def test_model_has_hsdp_shard_conditions(self, mocker: MockerFixture):
|
||||
"""Verify model has _hsdp_shard_conditions for HSDP parallelism."""
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
|
||||
parallel_config = DiffusionParallelConfig(
|
||||
tensor_parallel_size=1,
|
||||
sequence_parallel_size=1,
|
||||
)
|
||||
|
||||
mock_tf_config = mocker.MagicMock()
|
||||
mock_tf_config.patch_size = 2
|
||||
mock_tf_config.in_channels = 16
|
||||
mock_tf_config.out_channels = 16
|
||||
mock_tf_config.num_attention_heads = 64
|
||||
mock_tf_config.attention_head_dim = 40
|
||||
mock_tf_config.time_embed_dim = 512
|
||||
mock_tf_config.condition_dim = 256
|
||||
mock_tf_config.prior_vq_quantizer_codebook_size = 16384
|
||||
mock_tf_config.text_embed_dim = 1024
|
||||
mock_tf_config.num_layers = 2
|
||||
|
||||
mock_od_config = mocker.MagicMock(spec=OmniDiffusionConfig)
|
||||
mock_od_config.tf_model_config = mock_tf_config
|
||||
mock_od_config.parallel_config = parallel_config
|
||||
|
||||
model = GlmImageTransformer2DModel(
|
||||
od_config=mock_od_config,
|
||||
quant_config=None,
|
||||
)
|
||||
|
||||
assert hasattr(model, "_hsdp_shard_conditions")
|
||||
assert len(model._hsdp_shard_conditions) > 0
|
||||
|
||||
def test_model_creates_kv_cache(self, mocker: MockerFixture):
|
||||
"""Verify model can create KV cache for image editing."""
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
|
||||
parallel_config = DiffusionParallelConfig(
|
||||
tensor_parallel_size=1,
|
||||
sequence_parallel_size=1,
|
||||
)
|
||||
|
||||
mock_tf_config = mocker.MagicMock()
|
||||
mock_tf_config.patch_size = 2
|
||||
mock_tf_config.in_channels = 16
|
||||
mock_tf_config.out_channels = 16
|
||||
mock_tf_config.num_attention_heads = 64
|
||||
mock_tf_config.attention_head_dim = 40
|
||||
mock_tf_config.time_embed_dim = 512
|
||||
mock_tf_config.condition_dim = 256
|
||||
mock_tf_config.prior_vq_quantizer_codebook_size = 16384
|
||||
mock_tf_config.text_embed_dim = 1024
|
||||
mock_tf_config.num_layers = 2
|
||||
|
||||
mock_od_config = mocker.MagicMock(spec=OmniDiffusionConfig)
|
||||
mock_od_config.tf_model_config = mock_tf_config
|
||||
mock_od_config.parallel_config = parallel_config
|
||||
|
||||
model = GlmImageTransformer2DModel(
|
||||
od_config=mock_od_config,
|
||||
quant_config=None,
|
||||
)
|
||||
|
||||
kv_cache = model.create_kv_cache()
|
||||
assert kv_cache is not None
|
||||
assert len(kv_cache) == 2 # num_layers
|
||||
@@ -23,7 +23,6 @@ NUM_HEADS = 4
|
||||
NUM_KV_HEADS = 2
|
||||
HEAD_DIM = 16
|
||||
IMAGE_TOKEN_LEN = 8
|
||||
SUFFIX_TOKEN_LEN = 3
|
||||
SCALING = 1.0 / math.sqrt(HEAD_DIM)
|
||||
|
||||
|
||||
@@ -80,10 +79,6 @@ def _make_known_kv(num_tokens, base=0.0):
|
||||
return k, v
|
||||
|
||||
|
||||
def _gen_timestep_index(bs: int, current_start: int) -> torch.Tensor:
|
||||
return torch.full((bs, 1), current_start, dtype=torch.long)
|
||||
|
||||
|
||||
def _call_mgr(
|
||||
mgr,
|
||||
bs,
|
||||
@@ -95,8 +90,6 @@ def _call_mgr(
|
||||
uncond_cfg_prefill=False,
|
||||
num_image_tokens=IMAGE_TOKEN_LEN,
|
||||
shard_image_size=None,
|
||||
gen_timestep_scatter_index=None,
|
||||
position_ids=None,
|
||||
):
|
||||
query = torch.randn(bs * q_len, NUM_HEADS, HEAD_DIM)
|
||||
attn_mask = torch.zeros(bs, 1, seq_len, seq_len)
|
||||
@@ -111,8 +104,6 @@ def _call_mgr(
|
||||
uncond_cfg_prefill=uncond_cfg_prefill,
|
||||
num_image_tokens=num_image_tokens,
|
||||
shard_image_size=shard_image_size,
|
||||
gen_timestep_scatter_index=gen_timestep_scatter_index,
|
||||
position_ids=position_ids,
|
||||
)
|
||||
|
||||
|
||||
@@ -126,65 +117,54 @@ def test_no_ar_kv(bs):
|
||||
"""
|
||||
No AR KV injected. Tests the basic first_step cache → update reuse path.
|
||||
|
||||
Sequence layout per batch on first_step (q_len=14, IMAGE_TOKEN_LEN=8):
|
||||
[prompt(3) | current timestep/image(8) | suffix(3)]
|
||||
gen_timestep_scatter_index = 3
|
||||
cached_prompt_len = 3
|
||||
Sequence layout per batch on first_step (q_len=12, IMAGE_TOKEN_LEN=8):
|
||||
[prompt(3) | image(8) | eoi(1)]
|
||||
image_size = IMAGE_TOKEN_LEN + 1 = 9
|
||||
cached_prompt_len = seq_len - image_size = 12 - 9 = 3
|
||||
|
||||
After first_step:
|
||||
image_kv_cache_map stores prompt tokens [0:3] for each batch.
|
||||
image_kv_cache_map stores prompt tokens [0:3] for each batch (total 3*bs tokens flat).
|
||||
|
||||
Update step (q_len=IMAGE_TOKEN_LEN=8, seq_len = cached_prompt(3) + current(8) = 11):
|
||||
_reuse_prompt_kv produces [cached_prompt(3) | new_current(8)] per batch.
|
||||
Update step (q_len=IMAGE_TOKEN_LEN=8, seq_len = cached_prompt(3) + eoi(1) + image(8) = 12):
|
||||
_reuse_prompt_kv produces [cached_prompt(3) | new_image(8) | zero_eoi(1)] per batch.
|
||||
"""
|
||||
mgr = _make_cache_mgr()
|
||||
assert mgr.image_kv_cache_map is None
|
||||
|
||||
# --- first_step ---
|
||||
prompt_len = 3
|
||||
q_len = prompt_len + IMAGE_TOKEN_LEN + SUFFIX_TOKEN_LEN
|
||||
q_len = 12
|
||||
k_flat, v_flat = _make_known_kv(bs * q_len, base=1.0)
|
||||
|
||||
_call_mgr(
|
||||
mgr,
|
||||
bs,
|
||||
q_len=q_len,
|
||||
seq_len=q_len,
|
||||
key_flat=k_flat,
|
||||
value_flat=v_flat,
|
||||
first_step=True,
|
||||
gen_timestep_scatter_index=_gen_timestep_index(bs, prompt_len),
|
||||
)
|
||||
_call_mgr(mgr, bs, q_len=q_len, seq_len=q_len, key_flat=k_flat, value_flat=v_flat, first_step=True)
|
||||
|
||||
cached_key, cached_value = mgr.image_kv_cache_map
|
||||
# 3 prompt tokens cached per batch
|
||||
assert cached_key.shape == (bs, prompt_len, NUM_KV_HEADS, HEAD_DIM)
|
||||
assert torch.equal(mgr.image_kv_cache_lens, torch.full((bs,), prompt_len))
|
||||
assert cached_key.shape[0] == 3 * bs
|
||||
for b in range(bs):
|
||||
flat_offset = b * q_len
|
||||
assert torch.allclose(cached_key[b, :prompt_len], k_flat[flat_offset : flat_offset + prompt_len])
|
||||
assert torch.allclose(cached_value[b, :prompt_len], v_flat[flat_offset : flat_offset + prompt_len])
|
||||
cache_offset = b * 3
|
||||
assert torch.allclose(cached_key[cache_offset : cache_offset + 3], k_flat[flat_offset : flat_offset + 3])
|
||||
assert torch.allclose(cached_value[cache_offset : cache_offset + 3], v_flat[flat_offset : flat_offset + 3])
|
||||
|
||||
# --- update step ---
|
||||
img_q_len = IMAGE_TOKEN_LEN
|
||||
update_seq_len = prompt_len + img_q_len # cached_prompt(3) + current(8) = 11
|
||||
update_seq_len = 3 + 1 + img_q_len # cached_prompt(3) + eoi(1) + image(8) = 12
|
||||
new_img_k, new_img_v = _make_known_kv(bs * img_q_len, base=50.0)
|
||||
|
||||
key_input = new_img_k.reshape(bs, img_q_len, NUM_KV_HEADS, HEAD_DIM)
|
||||
val_input = new_img_v.reshape(bs, img_q_len, NUM_KV_HEADS, HEAD_DIM)
|
||||
position_ids = torch.arange(prompt_len, prompt_len + img_q_len).repeat(bs, 1)
|
||||
result_k, result_v = mgr._reuse_prompt_kv(key_input, val_input, update_seq_len, bs=bs, position_ids=position_ids)
|
||||
result_k, result_v = mgr._reuse_prompt_kv(key_input, val_input, update_seq_len, bs=bs)
|
||||
|
||||
assert result_k.shape == (bs, update_seq_len, NUM_KV_HEADS, HEAD_DIM)
|
||||
for b in range(bs):
|
||||
cache_offset = b * 3
|
||||
img_offset = b * img_q_len
|
||||
# Cached prompt preserved
|
||||
assert torch.allclose(result_k[b, :prompt_len], cached_key[b, :prompt_len])
|
||||
assert torch.allclose(result_k[b, :3], cached_key[cache_offset : cache_offset + 3])
|
||||
# New image tokens
|
||||
assert torch.allclose(
|
||||
result_k[b, prompt_len : prompt_len + img_q_len],
|
||||
new_img_k[img_offset : img_offset + img_q_len],
|
||||
)
|
||||
assert torch.allclose(result_k[b, 3 : 3 + img_q_len], new_img_k[img_offset : img_offset + img_q_len])
|
||||
# Zero eoi at the end
|
||||
assert torch.allclose(result_k[b, -1], torch.zeros(NUM_KV_HEADS, HEAD_DIM))
|
||||
|
||||
|
||||
# ============================================================
|
||||
@@ -198,17 +178,16 @@ def test_ar_kv_no_cfg(sp_size):
|
||||
AR KV injected, bs=1, no CFG. Tests AR prefix prepend + cache + reuse.
|
||||
|
||||
sp_size=1:
|
||||
first_step input: q_len=14, ar_len=5, seq_len=19
|
||||
Sequence layout: [ar(5) | prompt(3) | current(8) | suffix(3)]
|
||||
gen_timestep_scatter_index = 3
|
||||
cached_prompt_len = ar(5) + prompt(3) = 8
|
||||
first_step input: q_len=12, ar_len=5, seq_len=17
|
||||
Sequence layout: [ar(5) | prompt(3) | image(8) | eoi(1)]
|
||||
image_size = IMAGE_TOKEN_LEN + 1 = 9
|
||||
cached_prompt_len = seq_len - image_size = 17 - 9 = 8 (= ar(5) + prompt(3))
|
||||
|
||||
Update step: q_len=8, seq_len = cached_prompt(8) + current(8) = 16
|
||||
Result: [cached(8) | new_current(8)]
|
||||
Update step: q_len=8, seq_len = cached_prompt(8) + eoi(1) + image(8) = 17
|
||||
Result: [cached(8) | new_image(8) | zero_eoi(1)]
|
||||
|
||||
sp_size=2:
|
||||
shard_image_size = 4 (passed externally, simulating SP sharding)
|
||||
gen_timestep_scatter_index = 8
|
||||
cached_prompt_len = seq_len - shard_image_size = 17 - 4 = 13 (= ar(5) + prompt(8))
|
||||
Note: with SP, more of the sequence is treated as "prompt" because
|
||||
image tokens are sharded across ranks.
|
||||
@@ -225,23 +204,15 @@ def test_ar_kv_no_cfg(sp_size):
|
||||
mgr._injected_ar_kv = [(ar_k.clone(), ar_v.clone())]
|
||||
|
||||
# --- first_step: call _cache_prompt_kv directly ---
|
||||
bs = 1
|
||||
q_len = 12 if sp_size > 1 else 14
|
||||
seq_len = q_len + ar_len
|
||||
bs, q_len = 1, 12
|
||||
seq_len = q_len + ar_len # 17
|
||||
k_raw, v_raw = _make_known_kv(q_len, base=1.0)
|
||||
key_4d = k_raw.reshape(1, q_len, NUM_KV_HEADS, HEAD_DIM)
|
||||
val_4d = v_raw.reshape(1, q_len, NUM_KV_HEADS, HEAD_DIM)
|
||||
|
||||
shard_image_size = 4 if sp_size > 1 else None
|
||||
current_start = 8 if sp_size > 1 else 3
|
||||
mgr.image_kv_cache_map = None
|
||||
mgr._cache_prompt_kv(
|
||||
key_4d,
|
||||
val_4d,
|
||||
seq_len,
|
||||
shard_image_size,
|
||||
gen_timestep_scatter_index=_gen_timestep_index(bs, current_start),
|
||||
)
|
||||
mgr._cache_prompt_kv(key_4d, val_4d, seq_len, shard_image_size)
|
||||
|
||||
cached_key, cached_value = mgr.image_kv_cache_map
|
||||
if sp_size == 1:
|
||||
@@ -251,32 +222,25 @@ def test_ar_kv_no_cfg(sp_size):
|
||||
# cached = seq_len - shard_image_size = 17 - 4 = 13
|
||||
expected_cached_len = seq_len - shard_image_size
|
||||
|
||||
assert cached_key.shape == (bs, expected_cached_len, NUM_KV_HEADS, HEAD_DIM)
|
||||
assert cached_key.shape[0] == expected_cached_len
|
||||
# AR KV always at the front
|
||||
assert torch.allclose(cached_key[0, :ar_len], ar_k)
|
||||
assert torch.allclose(cached_value[0, :ar_len], ar_v)
|
||||
assert torch.allclose(cached_key[:ar_len], ar_k)
|
||||
assert torch.allclose(cached_value[:ar_len], ar_v)
|
||||
# Prompt tokens follow AR
|
||||
prompt_cached = expected_cached_len - ar_len
|
||||
assert torch.allclose(cached_key[0, ar_len : ar_len + prompt_cached], k_raw[:prompt_cached])
|
||||
assert torch.allclose(cached_key[ar_len : ar_len + prompt_cached], k_raw[:prompt_cached])
|
||||
# AR KV consumed
|
||||
assert mgr._injected_ar_kv is None
|
||||
|
||||
# --- update step: call _reuse_prompt_kv directly ---
|
||||
if sp_size == 1:
|
||||
img_q_len = IMAGE_TOKEN_LEN
|
||||
update_seq_len = expected_cached_len + img_q_len # 8 + 8 = 16
|
||||
update_seq_len = expected_cached_len + 1 + img_q_len # 8 + 1 + 8 = 17
|
||||
new_img_k, new_img_v = _make_known_kv(img_q_len, base=50.0)
|
||||
|
||||
key_input = new_img_k.reshape(1, img_q_len, NUM_KV_HEADS, HEAD_DIM)
|
||||
val_input = new_img_v.reshape(1, img_q_len, NUM_KV_HEADS, HEAD_DIM)
|
||||
position_ids = torch.arange(expected_cached_len, expected_cached_len + img_q_len).reshape(1, img_q_len)
|
||||
result_k, result_v = mgr._reuse_prompt_kv(
|
||||
key_input,
|
||||
val_input,
|
||||
update_seq_len,
|
||||
bs=1,
|
||||
position_ids=position_ids,
|
||||
)
|
||||
result_k, result_v = mgr._reuse_prompt_kv(key_input, val_input, update_seq_len, bs=1)
|
||||
|
||||
assert result_k.shape == (1, update_seq_len, NUM_KV_HEADS, HEAD_DIM)
|
||||
# AR + prompt preserved
|
||||
@@ -284,6 +248,8 @@ def test_ar_kv_no_cfg(sp_size):
|
||||
assert torch.allclose(result_k[0, ar_len : ar_len + prompt_cached], k_raw[:prompt_cached])
|
||||
# New image tokens
|
||||
assert torch.allclose(result_k[0, expected_cached_len : expected_cached_len + img_q_len], new_img_k)
|
||||
# Zero eoi
|
||||
assert torch.allclose(result_k[0, -1], torch.zeros(NUM_KV_HEADS, HEAD_DIM))
|
||||
else:
|
||||
# SP path: _reuse_prompt_kv returns only cached prompt (no image concat)
|
||||
img_q_len = shard_image_size # 4
|
||||
@@ -326,12 +292,13 @@ def test_ar_kv_with_cfg(cfg_parallel, bs):
|
||||
→ _injected_ar_kv becomes [(pos_ar(10), pos_av(10)), (neg_k(10), neg_v(10))]
|
||||
|
||||
first_step (bs=2, q_len=12, seq_len=22):
|
||||
Batch 0 (pos): [pos_ar(10) | prompt(3) | current(8) | suffix(3)]
|
||||
Batch 1 (neg): [neg_ar(10) | prompt(3) | current(8) | suffix(3)]
|
||||
cached_prompt_len per batch = 10 + 3 = 13
|
||||
Batch 0 (pos): [pos_ar(10) | prompt(3) | image(8) | eoi(1)]
|
||||
Batch 1 (neg): [neg_ar(10) | prompt(3) | image(8) | eoi(1)]
|
||||
cached_prompt_len per batch = 22 - 9 = 13
|
||||
Total cached = 13 * 2 = 26
|
||||
|
||||
Update (bs=2, seq_len = 13 + 8 = 21):
|
||||
Result per batch: [cached(13) | new_current(8)]
|
||||
Update (bs=2, seq_len = 13 + 1 + 8 = 22):
|
||||
Result per batch: [cached(13) | new_image(8) | zero_eoi(1)]
|
||||
|
||||
CFG Parallel (cfg_parallel=True, bs=1):
|
||||
This rank handles only the negative branch.
|
||||
@@ -341,11 +308,11 @@ def test_ar_kv_with_cfg(cfg_parallel, bs):
|
||||
→ _injected_ar_kv = [(neg_k(10), neg_v(10))]
|
||||
|
||||
first_step (bs=1, q_len=12, seq_len=22):
|
||||
[neg_ar(10) | prompt(3) | current(8) | suffix(3)]
|
||||
cached_prompt_len = 13
|
||||
[neg_ar(10) | prompt(3) | image(8) | eoi(1)]
|
||||
cached_prompt_len = 22 - 9 = 13
|
||||
|
||||
Update (bs=1, seq_len = 13 + 8 = 21):
|
||||
Result: [cached(13) | new_current(8)]
|
||||
Update (bs=1, seq_len = 13 + 1 + 8 = 22):
|
||||
Result: [cached(13) | new_image(8) | zero_eoi(1)]
|
||||
"""
|
||||
positive_reuse_len = 10
|
||||
negative_reuse_len = 6
|
||||
@@ -384,62 +351,46 @@ def test_ar_kv_with_cfg(cfg_parallel, bs):
|
||||
mgr._injected_ar_kv = [mgr._injected_ar_kv[1]]
|
||||
|
||||
# --- first_step ---
|
||||
prompt_len = 3
|
||||
q_len = prompt_len + IMAGE_TOKEN_LEN + SUFFIX_TOKEN_LEN
|
||||
q_len = 12
|
||||
seq_len = q_len + positive_reuse_len # 22
|
||||
k_flat, v_flat = _make_known_kv(bs * q_len, base=1.0)
|
||||
|
||||
_call_mgr(
|
||||
mgr,
|
||||
bs,
|
||||
q_len=q_len,
|
||||
seq_len=seq_len,
|
||||
key_flat=k_flat,
|
||||
value_flat=v_flat,
|
||||
first_step=True,
|
||||
gen_timestep_scatter_index=_gen_timestep_index(bs, prompt_len),
|
||||
)
|
||||
_call_mgr(mgr, bs, q_len=q_len, seq_len=seq_len, key_flat=k_flat, value_flat=v_flat, first_step=True)
|
||||
|
||||
cached_key, cached_value = mgr.image_kv_cache_map
|
||||
cached_prompt_len_per_batch = positive_reuse_len + prompt_len
|
||||
assert cached_key.shape == (bs, cached_prompt_len_per_batch, NUM_KV_HEADS, HEAD_DIM)
|
||||
cached_prompt_len_per_batch = seq_len - (IMAGE_TOKEN_LEN + 1) # 22 - 9 = 13
|
||||
assert cached_key.shape[0] == cached_prompt_len_per_batch * bs
|
||||
assert mgr._injected_ar_kv is None
|
||||
|
||||
if not cfg_parallel:
|
||||
# bs=2: batch 0 = pos, batch 1 = neg
|
||||
# Batch 0: pos_ar(10) + prompt(3)
|
||||
assert torch.allclose(cached_key[0, :positive_reuse_len], pos_ar_k)
|
||||
assert torch.allclose(cached_key[0, positive_reuse_len:13], k_flat[:3])
|
||||
assert torch.allclose(cached_key[:positive_reuse_len], pos_ar_k)
|
||||
assert torch.allclose(cached_key[positive_reuse_len:13], k_flat[:3])
|
||||
# Batch 1: neg_ar(10) + prompt(3)
|
||||
assert torch.allclose(cached_key[1, :positive_reuse_len], neg_ar_k)
|
||||
assert torch.allclose(cached_key[1, positive_reuse_len:13], k_flat[q_len : q_len + 3])
|
||||
assert torch.allclose(cached_key[13:23], neg_ar_k)
|
||||
assert torch.allclose(cached_key[23:26], k_flat[q_len : q_len + 3])
|
||||
else:
|
||||
# bs=1: only neg branch
|
||||
assert torch.allclose(cached_key[0, :positive_reuse_len], neg_ar_k)
|
||||
assert torch.allclose(cached_key[0, positive_reuse_len:13], k_flat[:3])
|
||||
assert torch.allclose(cached_key[:positive_reuse_len], neg_ar_k)
|
||||
assert torch.allclose(cached_key[positive_reuse_len:13], k_flat[:3])
|
||||
|
||||
# --- update step ---
|
||||
img_q_len = IMAGE_TOKEN_LEN
|
||||
update_seq_len = cached_prompt_len_per_batch + img_q_len # 13 + 8 = 21
|
||||
update_seq_len = cached_prompt_len_per_batch + 1 + img_q_len # 13 + 1 + 8 = 22
|
||||
new_img_k, new_img_v = _make_known_kv(bs * img_q_len, base=50.0)
|
||||
|
||||
key_input = new_img_k.reshape(bs, img_q_len, NUM_KV_HEADS, HEAD_DIM)
|
||||
val_input = new_img_v.reshape(bs, img_q_len, NUM_KV_HEADS, HEAD_DIM)
|
||||
position_ids = torch.arange(cached_prompt_len_per_batch, cached_prompt_len_per_batch + img_q_len).repeat(bs, 1)
|
||||
result_k, result_v = mgr._reuse_prompt_kv(
|
||||
key_input,
|
||||
val_input,
|
||||
update_seq_len,
|
||||
bs=bs,
|
||||
position_ids=position_ids,
|
||||
)
|
||||
result_k, result_v = mgr._reuse_prompt_kv(key_input, val_input, update_seq_len, bs=bs)
|
||||
|
||||
assert result_k.shape == (bs, update_seq_len, NUM_KV_HEADS, HEAD_DIM)
|
||||
for b in range(bs):
|
||||
cache_offset = b * cached_prompt_len_per_batch
|
||||
# Cached prompt preserved
|
||||
assert torch.allclose(
|
||||
result_k[b, :cached_prompt_len_per_batch],
|
||||
cached_key[b, :cached_prompt_len_per_batch],
|
||||
cached_key[cache_offset : cache_offset + cached_prompt_len_per_batch],
|
||||
)
|
||||
# New image tokens
|
||||
img_offset = b * img_q_len
|
||||
@@ -447,6 +398,8 @@ def test_ar_kv_with_cfg(cfg_parallel, bs):
|
||||
result_k[b, cached_prompt_len_per_batch : cached_prompt_len_per_batch + img_q_len],
|
||||
new_img_k[img_offset : img_offset + img_q_len],
|
||||
)
|
||||
# Zero eoi
|
||||
assert torch.allclose(result_k[b, -1], torch.zeros(NUM_KV_HEADS, HEAD_DIM))
|
||||
|
||||
|
||||
# ============================================================
|
||||
@@ -459,7 +412,7 @@ def test_cross_request_isolation():
|
||||
Verify leftover image_kv_cache_map from a previous request is NOT treated as AR KV.
|
||||
|
||||
Setup: mgr has stale cache from a prior request (9 tokens, base=999).
|
||||
New request: first_step with q_len=14, no AR KV.
|
||||
New request: first_step with q_len=12, no AR KV.
|
||||
|
||||
Expected: stale cache is overwritten. New cache = prompt tokens from current request.
|
||||
The stale values (999.x) must NOT appear in the new cache.
|
||||
@@ -468,33 +421,20 @@ def test_cross_request_isolation():
|
||||
|
||||
# Simulate leftover from previous request
|
||||
leftover_k, leftover_v = _make_known_kv(9, base=999.0)
|
||||
mgr.image_kv_cache_map = (
|
||||
leftover_k.reshape(1, 9, NUM_KV_HEADS, HEAD_DIM),
|
||||
leftover_v.reshape(1, 9, NUM_KV_HEADS, HEAD_DIM),
|
||||
)
|
||||
mgr.image_kv_cache_map = (leftover_k, leftover_v)
|
||||
assert mgr._injected_ar_kv is None
|
||||
|
||||
# New request first_step
|
||||
bs = 1
|
||||
prompt_len = 3
|
||||
seq_len = prompt_len + IMAGE_TOKEN_LEN + SUFFIX_TOKEN_LEN
|
||||
bs, seq_len = 1, 12
|
||||
k_flat, v_flat = _make_known_kv(bs * seq_len, base=1.0)
|
||||
|
||||
_call_mgr(
|
||||
mgr,
|
||||
bs,
|
||||
q_len=seq_len,
|
||||
seq_len=seq_len,
|
||||
key_flat=k_flat,
|
||||
value_flat=v_flat,
|
||||
first_step=True,
|
||||
gen_timestep_scatter_index=_gen_timestep_index(bs, prompt_len),
|
||||
)
|
||||
_call_mgr(mgr, bs, q_len=seq_len, seq_len=seq_len, key_flat=k_flat, value_flat=v_flat, first_step=True)
|
||||
|
||||
cached_key, cached_value = mgr.image_kv_cache_map
|
||||
assert cached_key.shape == (bs, prompt_len, NUM_KV_HEADS, HEAD_DIM)
|
||||
# cached_prompt_len = 12 - 9 = 3
|
||||
assert cached_key.shape[0] == 3
|
||||
# Must be from current request, not stale
|
||||
assert torch.allclose(cached_key[0, :prompt_len], k_flat[:prompt_len])
|
||||
assert torch.allclose(cached_value[0, :prompt_len], v_flat[:prompt_len])
|
||||
assert torch.allclose(cached_key[:3], k_flat[:3])
|
||||
assert torch.allclose(cached_value[:3], v_flat[:3])
|
||||
# Stale values must not be present
|
||||
assert not torch.any(cached_key >= 999.0)
|
||||
@@ -15,29 +15,12 @@ These tests verify:
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _make_ltx23_pipeline(sequence_parallel_size: int = 1):
|
||||
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
|
||||
|
||||
pipeline = object.__new__(LTX23Pipeline)
|
||||
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))
|
||||
pipeline.audio_vae = SimpleNamespace(
|
||||
latents_mean=torch.tensor(0.0),
|
||||
latents_std=torch.tensor(1.0),
|
||||
)
|
||||
return pipeline
|
||||
|
||||
|
||||
class TestPipelineIndependence:
|
||||
"""Verify LTX23Pipeline is fully independent from LTX2Pipeline."""
|
||||
|
||||
@@ -65,313 +48,6 @@ class TestPipelineIndependence:
|
||||
|
||||
assert issubclass(LTX23Pipeline, ProgressBarMixin)
|
||||
|
||||
def test_ltx23_pipeline_has_cfg_parallel_mixin(self):
|
||||
"""LTX23Pipeline must use the shared CFG parallel implementation."""
|
||||
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
|
||||
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
|
||||
|
||||
assert issubclass(LTX23Pipeline, CFGParallelMixin)
|
||||
|
||||
|
||||
class TestCFGParallelHelpers:
|
||||
"""Test LTX-2.3 CFG helper math without loading model weights."""
|
||||
|
||||
def test_combine_cfg_noise_matches_x0_space_formula(self):
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
|
||||
|
||||
pipe = object.__new__(LTX23Pipeline)
|
||||
video_sample = torch.tensor([[[1.0, -2.0]]])
|
||||
audio_sample = torch.tensor([[[0.5, 3.0]]])
|
||||
video_pos = torch.tensor([[[0.2, -0.3]]])
|
||||
video_neg = torch.tensor([[[-0.4, 0.1]]])
|
||||
audio_pos = torch.tensor([[[0.7, -0.2]]])
|
||||
audio_neg = torch.tensor([[[0.1, 0.4]]])
|
||||
video_sigma = torch.tensor(0.25)
|
||||
audio_sigma = torch.tensor(0.5)
|
||||
scale = 4.0
|
||||
|
||||
video_combined, audio_combined = pipe.combine_cfg_noise(
|
||||
(video_pos, audio_pos),
|
||||
(video_neg, audio_neg),
|
||||
scale,
|
||||
video_latents=video_sample,
|
||||
audio_latents=audio_sample,
|
||||
video_sigma=video_sigma,
|
||||
audio_sigma=audio_sigma,
|
||||
)
|
||||
|
||||
x0_video_cond = video_sample - video_pos * video_sigma
|
||||
x0_video_uncond = video_sample - video_neg * video_sigma
|
||||
x0_video_guided = x0_video_cond + (scale - 1) * (x0_video_cond - x0_video_uncond)
|
||||
expected_video = (video_sample - x0_video_guided) / video_sigma
|
||||
|
||||
x0_audio_cond = audio_sample - audio_pos * audio_sigma
|
||||
x0_audio_uncond = audio_sample - audio_neg * audio_sigma
|
||||
x0_audio_guided = x0_audio_cond + (scale - 1) * (x0_audio_cond - x0_audio_uncond)
|
||||
expected_audio = (audio_sample - x0_audio_guided) / audio_sigma
|
||||
assert torch.allclose(video_combined, expected_video)
|
||||
assert torch.allclose(audio_combined, expected_audio)
|
||||
|
||||
def test_two_rank_cfg_parallel_smoke_uses_rank_local_branch_and_x0_formula(self, monkeypatch):
|
||||
from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2_3 as ltx23
|
||||
|
||||
pipe = object.__new__(ltx23.LTX23Pipeline)
|
||||
video_sample = torch.tensor([[[1.0, -2.0]]])
|
||||
audio_sample = torch.tensor([[[0.5, 3.0, -1.0]]])
|
||||
video_pos = torch.tensor([[[0.2, -0.3]]])
|
||||
video_neg = torch.tensor([[[-0.4, 0.1]]])
|
||||
audio_pos = torch.tensor([[[0.7, -0.2, 0.3]]])
|
||||
audio_neg = torch.tensor([[[0.1, 0.4, -0.5]]])
|
||||
video_sigma = torch.tensor(0.25)
|
||||
audio_sigma = torch.tensor(0.5)
|
||||
scale = 4.0
|
||||
|
||||
class FakeCfgGroup:
|
||||
def all_gather(self, tensor, separate_tensors=True):
|
||||
assert separate_tensors
|
||||
if tensor.shape == video_pos.shape:
|
||||
return [video_pos, video_neg]
|
||||
return [audio_pos, audio_neg]
|
||||
|
||||
monkeypatch.setattr(ltx23, "get_classifier_free_guidance_world_size", lambda: 2)
|
||||
monkeypatch.setattr(ltx23, "get_cfg_group", lambda: FakeCfgGroup())
|
||||
|
||||
expected_video = ltx23.LTX23Pipeline._combine_x0_space_cfg(
|
||||
video_sample,
|
||||
video_pos,
|
||||
video_neg,
|
||||
video_sigma,
|
||||
scale,
|
||||
)
|
||||
expected_audio = ltx23.LTX23Pipeline._combine_x0_space_cfg(
|
||||
audio_sample,
|
||||
audio_pos,
|
||||
audio_neg,
|
||||
audio_sigma,
|
||||
scale,
|
||||
)
|
||||
|
||||
for rank, expected_branch in ((0, "positive"), (1, "negative")):
|
||||
calls = []
|
||||
monkeypatch.setattr(ltx23, "get_classifier_free_guidance_rank", lambda rank=rank: rank)
|
||||
|
||||
def fake_predict_noise(**kwargs):
|
||||
calls.append(kwargs["branch"])
|
||||
if kwargs["branch"] == "positive":
|
||||
return video_pos, audio_pos
|
||||
return video_neg, audio_neg
|
||||
|
||||
object.__setattr__(pipe, "predict_noise", fake_predict_noise)
|
||||
video_combined, audio_combined = pipe.predict_noise_with_parallel_cfg(
|
||||
true_cfg_scale=scale,
|
||||
positive_kwargs={"branch": "positive"},
|
||||
negative_kwargs={"branch": "negative"},
|
||||
cfg_normalize=False,
|
||||
video_latents=video_sample,
|
||||
audio_latents=audio_sample,
|
||||
video_sigma=video_sigma,
|
||||
audio_sigma=audio_sigma,
|
||||
)
|
||||
|
||||
assert calls == [expected_branch]
|
||||
torch.testing.assert_close(video_combined, expected_video)
|
||||
torch.testing.assert_close(audio_combined, expected_audio)
|
||||
|
||||
assert "_cfg_video_latents" not in pipe.__dict__
|
||||
assert "_cfg_audio_latents" not in pipe.__dict__
|
||||
|
||||
|
||||
class TestCFGParallelForwardPath:
|
||||
"""Test the LTX-2.3 CFG-parallel denoising path without loading model weights."""
|
||||
|
||||
@pytest.mark.parametrize(("cfg_rank", "expected_prompt_value"), [(0, 1.0), (1, 0.0)])
|
||||
def test_forward_cfg_parallel_steps_video_and_audio_scheduler(
|
||||
self,
|
||||
monkeypatch,
|
||||
cfg_rank,
|
||||
expected_prompt_value,
|
||||
):
|
||||
from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2_3 as ltx23
|
||||
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
pipe = object.__new__(ltx23.LTX23Pipeline)
|
||||
torch.nn.Module.__init__(pipe)
|
||||
pipe.device = torch.device("cpu")
|
||||
pipe.tokenizer_max_length = 1
|
||||
pipe.vae_spatial_compression_ratio = 32
|
||||
pipe.vae_temporal_compression_ratio = 1
|
||||
pipe.transformer_spatial_patch_size = 1
|
||||
pipe.transformer_temporal_patch_size = 1
|
||||
pipe.audio_sampling_rate = 1
|
||||
pipe.audio_hop_length = 1
|
||||
pipe.audio_vae_temporal_compression_ratio = 1
|
||||
pipe.audio_vae_mel_compression_ratio = 1
|
||||
pipe.od_config = SimpleNamespace(parallel_config=SimpleNamespace(sequence_parallel_size=1))
|
||||
pipe.tokenizer = SimpleNamespace(padding_side="left")
|
||||
pipe.vae = SimpleNamespace(
|
||||
latents_mean=torch.zeros(2),
|
||||
latents_std=torch.ones(2),
|
||||
config=SimpleNamespace(scaling_factor=1.0),
|
||||
)
|
||||
pipe.audio_vae = SimpleNamespace(
|
||||
latents_mean=torch.zeros(2),
|
||||
latents_std=torch.ones(2),
|
||||
config=SimpleNamespace(mel_bins=2, latent_channels=1),
|
||||
)
|
||||
|
||||
video_pos = torch.tensor([[[0.2, -0.3]]])
|
||||
video_neg = torch.tensor([[[-0.4, 0.1]]])
|
||||
audio_pos = torch.tensor([[[0.7, -0.2]]])
|
||||
audio_neg = torch.tensor([[[0.1, 0.4]]])
|
||||
|
||||
class FakeCfgGroup:
|
||||
def all_gather(self, tensor, separate_tensors=True):
|
||||
assert separate_tensors
|
||||
if torch.equal(tensor, video_pos) or torch.equal(tensor, video_neg):
|
||||
return [video_pos, video_neg]
|
||||
if torch.equal(tensor, audio_pos) or torch.equal(tensor, audio_neg):
|
||||
return [audio_pos, audio_neg]
|
||||
raise AssertionError(f"Unexpected gathered tensor: {tensor}")
|
||||
|
||||
monkeypatch.setattr(ltx23, "get_classifier_free_guidance_world_size", lambda: 2)
|
||||
monkeypatch.setattr(ltx23, "get_classifier_free_guidance_rank", lambda: cfg_rank)
|
||||
monkeypatch.setattr(ltx23, "get_cfg_group", lambda: FakeCfgGroup())
|
||||
|
||||
def fake_retrieve_timesteps(scheduler, num_inference_steps, device, timesteps, sigmas=None, mu=None):
|
||||
scheduler.sigmas = torch.tensor([0.25, 0.25], device=device)
|
||||
return torch.tensor([1.0, 0.5], device=device), 2
|
||||
|
||||
monkeypatch.setattr(ltx23, "retrieve_timesteps", fake_retrieve_timesteps)
|
||||
|
||||
class FakeScheduler:
|
||||
def __init__(self, name="video", calls=None):
|
||||
self.name = name
|
||||
self.calls = [] if calls is None else calls
|
||||
self.config = {
|
||||
"max_image_seq_len": 4096,
|
||||
"base_image_seq_len": 1024,
|
||||
"base_shift": 0.95,
|
||||
"max_shift": 2.05,
|
||||
}
|
||||
self.sigmas = torch.tensor([0.25, 0.25])
|
||||
|
||||
def __deepcopy__(self, memo):
|
||||
return FakeScheduler("audio", self.calls)
|
||||
|
||||
def step(self, noise_pred, t, latents, return_dict=False, generator=None):
|
||||
self.calls.append((self.name, noise_pred.clone(), t.clone(), latents.clone()))
|
||||
return (latents - noise_pred,)
|
||||
|
||||
class FakeConnectors:
|
||||
def to(self, device):
|
||||
return self
|
||||
|
||||
def __call__(self, prompt_embeds, prompt_attention_mask, padding_side):
|
||||
assert padding_side == "left"
|
||||
assert prompt_embeds.shape[0] == 2
|
||||
return prompt_embeds, prompt_embeds, prompt_attention_mask
|
||||
|
||||
class FakeRope:
|
||||
def prepare_video_coords(self, batch_size, num_frames, height, width, device, fps):
|
||||
return torch.zeros(batch_size, num_frames * height * width, 3, device=device)
|
||||
|
||||
def prepare_audio_coords(self, batch_size, num_frames, device):
|
||||
return torch.zeros(batch_size, num_frames, 1, device=device)
|
||||
|
||||
class FakeTransformer:
|
||||
def __init__(self):
|
||||
self.config = SimpleNamespace(in_channels=2)
|
||||
self.rope = FakeRope()
|
||||
self.audio_rope = FakeRope()
|
||||
self.calls = []
|
||||
|
||||
def __call__(self, **kwargs):
|
||||
self.calls.append(kwargs)
|
||||
expected_prompt = torch.full((1, 1, 1), expected_prompt_value)
|
||||
torch.testing.assert_close(kwargs["encoder_hidden_states"], expected_prompt)
|
||||
torch.testing.assert_close(kwargs["audio_encoder_hidden_states"], expected_prompt)
|
||||
assert kwargs["hidden_states"].shape == (1, 1, 2)
|
||||
assert kwargs["audio_hidden_states"].shape == (1, 1, 2)
|
||||
if cfg_rank == 0:
|
||||
return video_pos, audio_pos
|
||||
return video_neg, audio_neg
|
||||
|
||||
class DummyProgress:
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
def update(self):
|
||||
pass
|
||||
|
||||
pipe.scheduler = FakeScheduler()
|
||||
pipe.connectors = FakeConnectors()
|
||||
pipe.transformer = FakeTransformer()
|
||||
object.__setattr__(pipe, "progress_bar", lambda total: DummyProgress())
|
||||
|
||||
def fake_encode_prompt(**kwargs):
|
||||
return (
|
||||
torch.ones(1, 1, 1),
|
||||
torch.ones(1, 1, dtype=torch.bool),
|
||||
torch.zeros(1, 1, 1),
|
||||
torch.ones(1, 1, dtype=torch.bool),
|
||||
)
|
||||
|
||||
object.__setattr__(pipe, "encode_prompt", fake_encode_prompt)
|
||||
|
||||
video_latents = torch.tensor([[[1.0, -2.0]]])
|
||||
audio_latents = torch.tensor([[[0.5, 3.0]]])
|
||||
req = OmniDiffusionRequest(
|
||||
prompts=[{"prompt": "prompt", "negative_prompt": "negative"}],
|
||||
sampling_params=OmniDiffusionSamplingParams(
|
||||
height=32,
|
||||
width=32,
|
||||
num_frames=1,
|
||||
frame_rate=1.0,
|
||||
num_inference_steps=2,
|
||||
guidance_scale=4.0,
|
||||
latents=video_latents,
|
||||
audio_latents=audio_latents,
|
||||
output_type="latent",
|
||||
),
|
||||
request_id="ltx23-cfg-parallel-forward-test",
|
||||
)
|
||||
|
||||
output = pipe.forward(req)
|
||||
|
||||
expected_video_noise = ltx23.LTX23Pipeline._combine_x0_space_cfg(
|
||||
video_latents,
|
||||
video_pos,
|
||||
video_neg,
|
||||
pipe.scheduler.sigmas[0],
|
||||
4.0,
|
||||
)
|
||||
expected_audio_noise = ltx23.LTX23Pipeline._combine_x0_space_cfg(
|
||||
audio_latents,
|
||||
audio_pos,
|
||||
audio_neg,
|
||||
pipe.scheduler.sigmas[0],
|
||||
4.0,
|
||||
)
|
||||
scheduler_call_names = [call[0] for call in pipe.scheduler.calls]
|
||||
assert scheduler_call_names == ["video", "audio", "video", "audio"]
|
||||
assert len(pipe.transformer.calls) == 2
|
||||
torch.testing.assert_close(pipe.scheduler.calls[0][1], expected_video_noise)
|
||||
torch.testing.assert_close(pipe.scheduler.calls[1][1], expected_audio_noise)
|
||||
torch.testing.assert_close(pipe.scheduler.calls[2][1], expected_video_noise)
|
||||
torch.testing.assert_close(pipe.scheduler.calls[3][1], expected_audio_noise)
|
||||
torch.testing.assert_close(pipe.scheduler.calls[2][3], video_latents - expected_video_noise)
|
||||
torch.testing.assert_close(pipe.scheduler.calls[3][3], audio_latents - expected_audio_noise)
|
||||
|
||||
video_out, audio_out = output.output
|
||||
torch.testing.assert_close(video_out, (video_latents - 2 * expected_video_noise).reshape(1, 2, 1, 1, 1))
|
||||
torch.testing.assert_close(audio_out, (audio_latents - 2 * expected_audio_noise).reshape(1, 1, 1, 2))
|
||||
|
||||
|
||||
class TestRegistryIntegration:
|
||||
"""Verify all LTX-2.3 pipeline variants are registered."""
|
||||
@@ -552,85 +228,3 @@ class TestInitExports:
|
||||
for name in expected_classes:
|
||||
assert hasattr(ltx2, name), f"{name} not exported from ltx2 package"
|
||||
assert name in ltx2.__all__, f"{name} not in ltx2.__all__"
|
||||
|
||||
|
||||
class TestAudioLatentSPPadding:
|
||||
def test_prepare_audio_latents_pads_generated_dummy_length_for_sp(self):
|
||||
pipeline = _make_ltx23_pipeline(sequence_parallel_size=2)
|
||||
|
||||
latents, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents(
|
||||
batch_size=1,
|
||||
num_channels_latents=8,
|
||||
num_mel_bins=64,
|
||||
audio_latent_length=1,
|
||||
dtype=torch.float32,
|
||||
device=torch.device("cpu"),
|
||||
)
|
||||
|
||||
assert original_num_frames == 1
|
||||
assert padded_num_frames == 2
|
||||
assert latents.shape == (1, 2, 128)
|
||||
|
||||
def test_prepare_audio_latents_pads_provided_packed_sequence_dim_for_sp(self):
|
||||
pipeline = _make_ltx23_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_prepare_audio_latents_accepts_already_padded_4d_latents_for_sp(self):
|
||||
pipeline = _make_ltx23_pipeline(sequence_parallel_size=4)
|
||||
latents = torch.arange(96, dtype=torch.float32).view(1, 2, 12, 4)
|
||||
|
||||
audio_latent_length = pipeline._resolve_audio_latent_length(10, latents)
|
||||
padded, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents(
|
||||
batch_size=1,
|
||||
num_channels_latents=2,
|
||||
num_mel_bins=16,
|
||||
audio_latent_length=audio_latent_length,
|
||||
dtype=torch.float32,
|
||||
device=torch.device("cpu"),
|
||||
latents=latents,
|
||||
)
|
||||
|
||||
assert audio_latent_length == 10
|
||||
assert original_num_frames == 10
|
||||
assert padded_num_frames == 12
|
||||
assert padded.shape == (1, 12, 8)
|
||||
torch.testing.assert_close(padded, pipeline._pack_audio_latents(latents))
|
||||
|
||||
def test_resolve_audio_latent_length_preserves_legacy_4d_shape_inference(self):
|
||||
pipeline = _make_ltx23_pipeline(sequence_parallel_size=4)
|
||||
latents = torch.zeros(1, 2, 13, 4)
|
||||
|
||||
audio_latent_length = pipeline._resolve_audio_latent_length(10, latents)
|
||||
|
||||
assert audio_latent_length == 13
|
||||
|
||||
def test_prepare_audio_latents_rejects_incompatible_provided_length(self):
|
||||
pipeline = _make_ltx23_pipeline(sequence_parallel_size=4)
|
||||
latents = torch.zeros(1, 11, 4)
|
||||
|
||||
with pytest.raises(ValueError, match="incompatible audio frame count"):
|
||||
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,
|
||||
)
|
||||
@@ -167,7 +167,6 @@ def test_basic_generation(ovis_pipeline):
|
||||
# Setup request
|
||||
req = OmniDiffusionRequest(
|
||||
prompts=["A photo of a cat"],
|
||||
request_id="ovis-basic",
|
||||
sampling_params=OmniDiffusionSamplingParams(
|
||||
height=256,
|
||||
width=256,
|
||||
@@ -207,7 +206,6 @@ def test_guidance_scale(ovis_pipeline, monkeypatch: pytest.MonkeyPatch):
|
||||
"negative_prompt": "bad quality",
|
||||
}
|
||||
],
|
||||
request_id="ovis-guidance",
|
||||
sampling_params=OmniDiffusionSamplingParams(
|
||||
height=256,
|
||||
width=256,
|
||||
@@ -227,7 +225,6 @@ def test_resolution_check(ovis_pipeline):
|
||||
# Pass odd resolution
|
||||
req = OmniDiffusionRequest(
|
||||
prompts=["test"],
|
||||
request_id="ovis-resolution",
|
||||
sampling_params=OmniDiffusionSamplingParams(
|
||||
height=250, # Not divisible by 16 (8*2)
|
||||
width=250,
|
||||
|
||||
@@ -115,64 +115,37 @@ def test_mxfp8_quant_config_structure():
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.skipif(not current_omni_platform.is_npu(), reason="Native MXFP8 offline only supported on NPU")
|
||||
def test_get_quant_method_offline_npu(mocker: MockerFixture):
|
||||
"""Offline (serialized) path must return NPUMxfp8LinearMethod on NPU."""
|
||||
from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
|
||||
def test_get_quant_method_npu_offline(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch):
|
||||
"""Offline (serialized) path on NPU must return NPUMxfp8LinearMethod."""
|
||||
from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config, NPUMxfp8LinearMethod
|
||||
|
||||
config = DiffusionMXFP8Config(is_checkpoint_mxfp8_serialized=True)
|
||||
layer = mocker.Mock(spec=LinearBase)
|
||||
monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True)
|
||||
|
||||
method = config.get_quant_method(layer, "blocks.0.attn1.to_q")
|
||||
assert type(method).__name__ == "NPUMxfp8LinearMethod"
|
||||
assert isinstance(method, NPUMxfp8LinearMethod)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not current_omni_platform.is_xpu(), reason="XPU platform not available")
|
||||
def test_get_quant_method_offline_xpu_raises(mocker: MockerFixture):
|
||||
"""XPU offline mode must raise NotImplementedError (use AutoRound MXFP8 instead)."""
|
||||
from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
|
||||
|
||||
config = DiffusionMXFP8Config(is_checkpoint_mxfp8_serialized=True)
|
||||
layer = mocker.Mock(spec=LinearBase)
|
||||
|
||||
with pytest.raises(NotImplementedError, match="Native MXFP8 offline mode is not supported on XPU"):
|
||||
config.get_quant_method(layer, "blocks.0.attn1.to_q")
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not (current_omni_platform.is_npu() or current_omni_platform.is_xpu()), reason="MXFP8 only supported on NPU and XPU"
|
||||
)
|
||||
def test_get_quant_method_online(mocker: MockerFixture):
|
||||
"""Online (BF16 checkpoint) path must return platform-specific method on current platform."""
|
||||
from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
|
||||
|
||||
# Mock the vLLM online method to avoid config dependency for XPU
|
||||
if current_omni_platform.is_xpu():
|
||||
mock_inner = mocker.Mock()
|
||||
mocker.patch(
|
||||
"vllm_omni.quantization.mxfp8_config.VllmMxfp8OnlineLinearMethod.__init__",
|
||||
lambda self: setattr(self, "_inner", mock_inner),
|
||||
)
|
||||
def test_get_quant_method_npu_online(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch):
|
||||
"""Online (BF16 checkpoint) path on NPU must return NPUMxfp8OnlineLinearMethod."""
|
||||
from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config, NPUMxfp8OnlineLinearMethod
|
||||
|
||||
config = DiffusionMXFP8Config(is_checkpoint_mxfp8_serialized=False)
|
||||
layer = mocker.Mock(spec=LinearBase)
|
||||
monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True)
|
||||
|
||||
method = config.get_quant_method(layer, "blocks.0.attn1.to_q")
|
||||
|
||||
if current_omni_platform.is_npu():
|
||||
assert type(method).__name__ == "NPUMxfp8OnlineLinearMethod"
|
||||
elif current_omni_platform.is_xpu():
|
||||
assert type(method).__name__ == "VllmMxfp8OnlineLinearMethod"
|
||||
assert isinstance(method, NPUMxfp8OnlineLinearMethod)
|
||||
|
||||
|
||||
def test_get_quant_method_unsupported_platform_raises(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch):
|
||||
"""Unsupported platform (CUDA, ROCm) must raise NotImplementedError."""
|
||||
def test_get_quant_method_non_npu_raises(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch):
|
||||
"""Non-NPU platform must raise NotImplementedError (CUDA not yet supported)."""
|
||||
from vllm_omni.quantization.mxfp8_config import DiffusionMXFP8Config
|
||||
|
||||
config = DiffusionMXFP8Config()
|
||||
layer = mocker.Mock(spec=LinearBase)
|
||||
monkeypatch.setattr(current_omni_platform, "is_npu", lambda: False)
|
||||
monkeypatch.setattr(current_omni_platform, "is_xpu", lambda: False)
|
||||
monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: False)
|
||||
|
||||
with pytest.raises(NotImplementedError):
|
||||
@@ -186,8 +159,6 @@ def test_get_quant_method_ignored_layer(mocker: MockerFixture, monkeypatch: pyte
|
||||
config = DiffusionMXFP8Config(ignored_layers=["proj_out"])
|
||||
layer = mocker.Mock(spec=LinearBase)
|
||||
monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True)
|
||||
monkeypatch.setattr(current_omni_platform, "is_xpu", lambda: False)
|
||||
monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: False)
|
||||
|
||||
method = config.get_quant_method(layer, "proj_out")
|
||||
assert isinstance(method, UnquantizedLinearMethod)
|
||||
|
||||
@@ -71,7 +71,7 @@ class QualityTestConfig:
|
||||
prompt: str # generation prompt
|
||||
max_lpips: float # fail threshold — higher = more lenient
|
||||
model: str | None = None # HF model name
|
||||
quantization: str | dict[str, object] | None = None # quantization method/config, e.g. "fp8"
|
||||
quantization: str | None = None # quantization method, e.g. "fp8"
|
||||
baseline_model: str | None = None # explicit BF16/local baseline path
|
||||
quantized_model: str | None = None # explicit quantized/local model path
|
||||
height: int = 1024
|
||||
@@ -118,57 +118,9 @@ QUALITY_CONFIGS = [
|
||||
QualityTestConfig(
|
||||
id="fp8_z_image",
|
||||
model="Tongyi-MAI/Z-Image-Turbo",
|
||||
quantization={
|
||||
"method": "fp8",
|
||||
"ignored_layers": [
|
||||
"img_mlp",
|
||||
*[
|
||||
f"layers.{layer_id}.{suffix}"
|
||||
for layer_id in range(15, 30)
|
||||
for suffix in (
|
||||
"attention.to_qkv",
|
||||
"attention.to_out.0",
|
||||
"feed_forward.w13",
|
||||
"feed_forward.w2",
|
||||
)
|
||||
],
|
||||
*[
|
||||
f"model.layers.{layer_id}.{suffix}"
|
||||
for layer_id in range(28, 36)
|
||||
for suffix in (
|
||||
"self_attn.q_proj",
|
||||
"self_attn.k_proj",
|
||||
"self_attn.v_proj",
|
||||
"self_attn.o_proj",
|
||||
"mlp.gate_proj",
|
||||
"mlp.up_proj",
|
||||
"mlp.down_proj",
|
||||
)
|
||||
],
|
||||
],
|
||||
},
|
||||
quantization="fp8",
|
||||
task="t2i",
|
||||
prompt=(
|
||||
"A breathtaking twilight scene atop a floating archipelago of crystalline islands suspended in an "
|
||||
"endless nebula-drenched sky, where waterfalls of liquid starlight cascade from the edges of each "
|
||||
"island into the cosmic abyss below. The largest island hosts an ancient, overgrown observatory "
|
||||
"crafted from pearlescent white stone and living silverwood trees whose bioluminescent leaves pulse "
|
||||
"with soft cerulean and violet light. At the observatory's center stands a solitary figure—a young "
|
||||
"astronomer in flowing robes woven from woven moonlight and deep-space fabric, their face illuminated "
|
||||
"by the glow of a holographic star chart hovering above an altar of polished obsidian. The sky swirls "
|
||||
"with the birth of a new galaxy: ribbons of magenta and gold gas clouds twist around nascent stars, "
|
||||
"while distant supernovae bloom like cosmic flowers in the far reaches of the void. Below, the abyss "
|
||||
"reveals faint echoes of forgotten civilizations—ghostly silhouettes of submerged cities shimmering "
|
||||
"in layers of atmospheric haze. Cinematic volumetric lighting cuts through the scene as twin moons—one "
|
||||
"copper-hued, one opalescent—rise on opposite horizons, casting long, intersecting shadows across "
|
||||
"moss-covered ruins and crystalline flora that refract light into prismatic halos. Hyper-detailed, "
|
||||
"photorealistic rendering with the atmospheric depth of Roger Deakins' cinematography, the "
|
||||
"architectural grandeur of Zaha Hadid, and the cosmic wonder of James Jean's illustrations. Shot on "
|
||||
"a mythical 150mm lens with shallow depth of field, 8K resolution, Unreal Engine 5 realism, subsurface "
|
||||
"scattering on organic elements, and ray-traced reflections dancing across every water droplet in the "
|
||||
"starlight waterfalls. Ethereal, melancholic, and transcendent mood—like a moment of quiet revelation "
|
||||
"at the edge of existence."
|
||||
),
|
||||
prompt="a cup of coffee on a wooden table, morning light",
|
||||
max_lpips=0.15,
|
||||
num_inference_steps=20,
|
||||
),
|
||||
@@ -439,6 +391,12 @@ _OUTPUT_DIR = Path(os.environ["VLLM_OMNI_QUALITY_OUTPUT_DIR"]) if "VLLM_OMNI_QUA
|
||||
|
||||
def _quality_param(c: QualityTestConfig):
|
||||
marks = list(_marks)
|
||||
if c.id == "fp8_z_image":
|
||||
marks.append(
|
||||
pytest.mark.skip(
|
||||
reason="Z-Image FP8 quality gate temporarily disabled: https://github.com/vllm-project/vllm-omni/issues/3531"
|
||||
)
|
||||
)
|
||||
if c.id == "fp8_qwen_image":
|
||||
marks.append(
|
||||
pytest.mark.skip(reason="Qwen-Image FP8 quality gate temporarily disabled (see CI / issue tracker).")
|
||||
|
||||
@@ -22,15 +22,15 @@ class DiffusionSchedulerOutput:
|
||||
num_waiting_reqs: int = 0
|
||||
|
||||
@property
|
||||
def scheduled_request_ids(self):
|
||||
def scheduled_req_ids(self):
|
||||
ids = [req.request_id for req in self.scheduled_new_reqs]
|
||||
if self.scheduled_cached_reqs:
|
||||
ids.extend(self.scheduled_cached_reqs.request_ids)
|
||||
ids.extend(self.scheduled_cached_reqs.sched_req_ids)
|
||||
return ids
|
||||
|
||||
@property
|
||||
def is_empty(self):
|
||||
return len(self.scheduled_request_ids) == 0
|
||||
return len(self.scheduled_req_ids) == 0
|
||||
|
||||
|
||||
class MockScheduler:
|
||||
@@ -71,36 +71,31 @@ async def test_async_add_req_and_wait_for_response():
|
||||
engine.scheduler = MockScheduler()
|
||||
engine._out_queue = {}
|
||||
engine.abort_queue: queue.Queue[str] = queue.Queue()
|
||||
engine._rpc_queue = queue.Queue()
|
||||
engine._rpc_lock = threading.RLock()
|
||||
engine._cv = threading.Condition(engine._rpc_lock)
|
||||
engine._init_lock = asyncio.Lock()
|
||||
engine._closed = False
|
||||
engine._loop_started = False
|
||||
engine.main_loop = None
|
||||
engine._rpc_lock = threading.Lock()
|
||||
|
||||
engine._finalize_finished_request = lambda rid, out, err: out.result
|
||||
|
||||
def mock_execute_batch(sched_output):
|
||||
request_ids = sched_output.scheduled_request_ids
|
||||
req_ids = sched_output.scheduled_req_ids
|
||||
|
||||
time.sleep(1)
|
||||
|
||||
class MockRunnerOutput:
|
||||
def __init__(self, ids):
|
||||
self.request_id = ids
|
||||
self.req_id = ids
|
||||
self.step_index = 0
|
||||
self.finished = True
|
||||
self._results = {rid: SimpleNamespace(result_data=f"data_{rid}") for rid in ids}
|
||||
|
||||
def get_request_output(self, rid):
|
||||
def get_req_output(self, rid):
|
||||
return SimpleNamespace(result=self._results[rid], step_index=0, finished=True)
|
||||
|
||||
return MockRunnerOutput(request_ids)
|
||||
return MockRunnerOutput(req_ids)
|
||||
|
||||
engine.execute_fn = mock_execute_batch
|
||||
|
||||
await engine._check_and_start_background_loop()
|
||||
engine.stop_event = threading.Event()
|
||||
engine.start_background_loop()
|
||||
|
||||
async def run_task(rid):
|
||||
req = SimpleNamespace(request_id=rid)
|
||||
@@ -110,13 +105,10 @@ async def test_async_add_req_and_wait_for_response():
|
||||
|
||||
task_ids = [f"req_{i}" for i in range(5)]
|
||||
tasks = [run_task(rid) for rid in task_ids]
|
||||
try:
|
||||
results = await asyncio.gather(*tasks)
|
||||
finally:
|
||||
with engine._cv:
|
||||
engine.stop_event.set()
|
||||
engine._cv.notify_all()
|
||||
engine.worker_thread.join(timeout=5)
|
||||
results = await asyncio.gather(*tasks)
|
||||
|
||||
engine.stop_event.set()
|
||||
engine.worker_thread.join(timeout=5)
|
||||
assert len(results) == 5
|
||||
for rid, res, elapsed in results:
|
||||
assert rid in res.result_data
|
||||
|
||||
@@ -17,10 +17,15 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
|
||||
def _make_request(request_id: str) -> OmniDiffusionRequest:
|
||||
def _make_request(
|
||||
request_id: str,
|
||||
*,
|
||||
request_ids: list[str] | None = None,
|
||||
) -> OmniDiffusionRequest:
|
||||
return OmniDiffusionRequest(
|
||||
prompts=[f"prompt_{request_id}"],
|
||||
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
|
||||
request_ids=request_ids or [request_id],
|
||||
request_id=request_id,
|
||||
)
|
||||
|
||||
@@ -74,15 +79,15 @@ def test_handle_finished_requests_ignores_already_drained_waiter() -> None:
|
||||
engine._finalize_finished_request.assert_not_called()
|
||||
|
||||
|
||||
def test_abort_request_id_aborts_scheduler_request() -> None:
|
||||
def test_abort_parent_request_id_aborts_batched_scheduler_request() -> None:
|
||||
engine = _make_engine()
|
||||
request = _make_request("batch-parent")
|
||||
request_id = engine.scheduler.add_request(request)
|
||||
request = _make_request("batch-parent", request_ids=["batch-parent-0", "batch-parent-1"])
|
||||
sched_req_id = engine.scheduler.add_request(request)
|
||||
|
||||
engine.abort("batch-parent")
|
||||
engine._process_aborts_queue()
|
||||
|
||||
state = engine.scheduler.get_request_state(request_id)
|
||||
state = engine.scheduler.get_request_state(sched_req_id)
|
||||
assert state is not None
|
||||
assert state.status == DiffusionRequestStatus.FINISHED_ABORTED
|
||||
|
||||
|
||||
@@ -89,7 +89,7 @@ class _ConcurrencyTrackingExecutor:
|
||||
if method in {"execute_model", "execute_stepwise", "generate"}:
|
||||
# args[0] is the request-like object for execute_model.
|
||||
req = args[0] if args else None
|
||||
tag = req.request_id if req is not None and hasattr(req, "request_id") else "unknown"
|
||||
tag = req.request_ids[0] if req is not None and hasattr(req, "request_ids") else "unknown"
|
||||
return DiffusionOutput(error=f"result_for_{tag}")
|
||||
tag = args[0] if args else method
|
||||
return DiffusionOutput(error=f"rpc_result_for_{tag}")
|
||||
@@ -108,7 +108,7 @@ class _ConcurrencyTrackingExecutor:
|
||||
exec_all_ranks=True,
|
||||
)
|
||||
return RunnerOutput(
|
||||
request_id=new_req.request_id,
|
||||
req_id=new_req.sched_req_id,
|
||||
step_index=None,
|
||||
finished=True,
|
||||
result=result,
|
||||
@@ -120,7 +120,7 @@ class _ConcurrencyTrackingExecutor:
|
||||
|
||||
def _make_request(tag: str):
|
||||
return SimpleNamespace(
|
||||
request_id=tag,
|
||||
request_ids=[tag],
|
||||
prompts=[f"prompt_{tag}"],
|
||||
sampling_params=SimpleNamespace(num_inference_steps=1, **_SAMPLING_KEY_DEFAULTS),
|
||||
)
|
||||
@@ -148,7 +148,6 @@ def _make_engine_with_loop(
|
||||
engine._rpc_lock = threading.RLock()
|
||||
engine._cv = threading.Condition(engine._rpc_lock)
|
||||
engine._out_queue = {}
|
||||
engine._closed = False
|
||||
engine.abort_queue = queue.Queue()
|
||||
engine._rpc_queue = queue.Queue()
|
||||
|
||||
|
||||
@@ -5,7 +5,7 @@ import random
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.diffusion.request import DUMMY_DIFFUSION_REQUEST_ID, OmniDiffusionRequest
|
||||
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
@@ -15,53 +15,9 @@ def _make_request() -> OmniDiffusionRequest:
|
||||
return OmniDiffusionRequest(
|
||||
prompts=[{"prompt": "a cup of coffee on a table"}],
|
||||
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
|
||||
request_id="request-test",
|
||||
)
|
||||
|
||||
|
||||
def test_request_id_is_required():
|
||||
with pytest.raises(TypeError, match="request_id"):
|
||||
OmniDiffusionRequest(
|
||||
prompts=[{"prompt": "a cup of coffee on a table"}],
|
||||
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
|
||||
)
|
||||
|
||||
|
||||
def test_request_ids_identity_list_is_removed():
|
||||
req = _make_request()
|
||||
|
||||
assert req.request_id == "request-test"
|
||||
assert not hasattr(req, "request_ids")
|
||||
|
||||
|
||||
def test_request_id_must_be_non_empty():
|
||||
with pytest.raises(ValueError, match="request_id must be a non-empty string"):
|
||||
OmniDiffusionRequest(
|
||||
prompts=[{"prompt": "a cup of coffee on a table"}],
|
||||
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
|
||||
request_id="",
|
||||
)
|
||||
|
||||
|
||||
def test_dummy_run_request_is_identified_by_reserved_request_id():
|
||||
req = OmniDiffusionRequest(
|
||||
prompts=[{"prompt": "dummy run"}],
|
||||
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
|
||||
request_id=DUMMY_DIFFUSION_REQUEST_ID,
|
||||
)
|
||||
|
||||
assert req.is_dummy_run()
|
||||
assert OmniDiffusionRequest.is_dummy_run_request_id(DUMMY_DIFFUSION_REQUEST_ID)
|
||||
|
||||
|
||||
def test_non_dummy_request_is_not_identified_as_dummy_run():
|
||||
req = _make_request()
|
||||
|
||||
assert not req.is_dummy_run()
|
||||
assert not OmniDiffusionRequest.is_dummy_run_request_id(req.request_id)
|
||||
assert not OmniDiffusionRequest.is_dummy_run_request_id(None)
|
||||
|
||||
|
||||
def test_tp_seed_same_across_ranks_and_varies_across_requests():
|
||||
random.seed(0)
|
||||
n_requests = 5
|
||||
|
||||
@@ -31,13 +31,13 @@ def _make_request(req_id: str) -> OmniDiffusionRequest:
|
||||
return OmniDiffusionRequest(
|
||||
prompts=[f"prompt_{req_id}"],
|
||||
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
|
||||
request_id=req_id,
|
||||
request_ids=[req_id],
|
||||
)
|
||||
|
||||
|
||||
def _make_request_output(req_id: str, *, error: str | None = None, finished: bool = True):
|
||||
return RunnerOutput(
|
||||
request_id=req_id,
|
||||
req_id=req_id,
|
||||
step_index=None,
|
||||
finished=finished,
|
||||
result=DiffusionOutput(output=None, error=error),
|
||||
@@ -52,7 +52,7 @@ def _make_step_output(
|
||||
error: str | None = None,
|
||||
):
|
||||
return RunnerOutput(
|
||||
request_id=req_id,
|
||||
req_id=req_id,
|
||||
step_index=step_index,
|
||||
finished=finished,
|
||||
result=DiffusionOutput(output=None, error=error) if error is not None else None,
|
||||
@@ -73,16 +73,16 @@ def _make_step_request(
|
||||
num_inference_steps=num_inference_steps,
|
||||
step_index=step_index,
|
||||
),
|
||||
request_id=req_id,
|
||||
request_ids=[req_id],
|
||||
)
|
||||
|
||||
|
||||
def _new_ids(sched_output) -> list[str]:
|
||||
return [req.request_id for req in sched_output.scheduled_new_reqs]
|
||||
return [req.sched_req_id for req in sched_output.scheduled_new_reqs]
|
||||
|
||||
|
||||
def _cached_ids(sched_output) -> list[str]:
|
||||
return list(sched_output.scheduled_cached_reqs.request_ids)
|
||||
return list(sched_output.scheduled_cached_reqs.sched_req_ids)
|
||||
|
||||
|
||||
class _StubScheduler(SchedulerInterface):
|
||||
@@ -90,7 +90,7 @@ class _StubScheduler(SchedulerInterface):
|
||||
self._request = request
|
||||
self._output = output
|
||||
self.initialized_with = None
|
||||
self._request_id = request.request_id
|
||||
self._sched_req_id = request.request_ids[0]
|
||||
self._state = None
|
||||
self._scheduled = False
|
||||
self.max_num_running_reqs = 1
|
||||
@@ -100,22 +100,22 @@ class _StubScheduler(SchedulerInterface):
|
||||
|
||||
def add_request(self, request: OmniDiffusionRequest) -> str:
|
||||
assert request is self._request
|
||||
self._state = SimpleNamespace(request_id=self._request_id, req=request)
|
||||
return self._request_id
|
||||
self._state = SimpleNamespace(sched_req_id=self._sched_req_id, req=request)
|
||||
return self._sched_req_id
|
||||
|
||||
def schedule(self):
|
||||
if self._scheduled or self._state is None:
|
||||
return SimpleNamespace(
|
||||
scheduled_new_reqs=[],
|
||||
scheduled_cached_reqs=CachedRequestData.make_empty(),
|
||||
scheduled_request_ids=[],
|
||||
scheduled_req_ids=[],
|
||||
is_empty=True,
|
||||
)
|
||||
self._scheduled = True
|
||||
return SimpleNamespace(
|
||||
scheduled_new_reqs=[NewRequestData.from_state(self._state)],
|
||||
scheduled_cached_reqs=CachedRequestData.make_empty(),
|
||||
scheduled_request_ids=[self._state.request_id],
|
||||
scheduled_req_ids=[self._state.sched_req_id],
|
||||
is_empty=False,
|
||||
)
|
||||
|
||||
@@ -123,25 +123,30 @@ class _StubScheduler(SchedulerInterface):
|
||||
del sched_output
|
||||
assert output is self._output
|
||||
self._state.status = DiffusionRequestStatus.FINISHED_COMPLETED
|
||||
return {self._request_id}
|
||||
return {self._sched_req_id}
|
||||
|
||||
def has_requests(self) -> bool:
|
||||
return not self._scheduled
|
||||
|
||||
def get_request_state(self, request_id: str):
|
||||
del request_id
|
||||
def get_request_state(self, sched_req_id: str):
|
||||
del sched_req_id
|
||||
return self._state
|
||||
|
||||
def pop_request_state(self, request_id: str):
|
||||
del request_id
|
||||
def get_sched_req_id(self, request_id: str) -> str | None:
|
||||
if request_id in self._request.request_ids:
|
||||
return self._sched_req_id
|
||||
return None
|
||||
|
||||
def pop_request_state(self, sched_req_id: str):
|
||||
del sched_req_id
|
||||
return self._state
|
||||
|
||||
def preempt_request(self, request_id: str) -> bool:
|
||||
del request_id
|
||||
def preempt_request(self, sched_req_id: str) -> bool:
|
||||
del sched_req_id
|
||||
return False
|
||||
|
||||
def finish_requests(self, request_ids, status) -> None:
|
||||
del request_ids, status
|
||||
def finish_requests(self, sched_req_ids, status) -> None:
|
||||
del sched_req_ids, status
|
||||
return None
|
||||
|
||||
def close(self) -> None:
|
||||
@@ -166,7 +171,7 @@ class TestGetSamplingParamsKey:
|
||||
return OmniDiffusionRequest(
|
||||
prompts=["prompt"],
|
||||
sampling_params=sp,
|
||||
request_id=f"req-{lora_int_id}-{lora_scale}",
|
||||
request_ids=[f"req-{lora_int_id}-{lora_scale}"],
|
||||
)
|
||||
|
||||
def test_distinguishes_lora_id(self) -> None:
|
||||
@@ -286,7 +291,7 @@ class TestRequestScheduler:
|
||||
OmniDiffusionRequest(
|
||||
prompts=["prompt_b"],
|
||||
sampling_params=OmniDiffusionSamplingParams(width=768),
|
||||
request_id="b",
|
||||
request_ids=["b"],
|
||||
)
|
||||
)
|
||||
scheduler.add_request(_make_request("c"))
|
||||
@@ -328,7 +333,7 @@ class TestRequestScheduler:
|
||||
assert state_a.status == DiffusionRequestStatus.FINISHED_ABORTED
|
||||
|
||||
assert self.scheduler.has_requests() is False
|
||||
assert self.scheduler.schedule().scheduled_request_ids == []
|
||||
assert self.scheduler.schedule().scheduled_req_ids == []
|
||||
|
||||
def test_has_requests_state_transition(self) -> None:
|
||||
assert self.scheduler.has_requests() is False
|
||||
@@ -343,28 +348,60 @@ class TestRequestScheduler:
|
||||
assert self.scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_COMPLETED
|
||||
assert self.scheduler.has_requests() is False
|
||||
|
||||
def test_request_id_is_scheduler_key(self) -> None:
|
||||
def test_request_id_mapping_lifecycle(self) -> None:
|
||||
request = OmniDiffusionRequest(
|
||||
prompts=["prompt_map_a", "prompt_map_b"],
|
||||
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
|
||||
request_ids=["map-a", "map-b"],
|
||||
request_id="map-parent",
|
||||
)
|
||||
|
||||
request_id = self.scheduler.add_request(request)
|
||||
sched_req_id = self.scheduler.add_request(request)
|
||||
|
||||
assert request_id == "map-parent"
|
||||
state = self.scheduler.get_request_state("map-parent")
|
||||
assert state.request_id == "map-parent"
|
||||
assert self.scheduler.get_sched_req_id("map-a") == sched_req_id
|
||||
assert self.scheduler.get_sched_req_id("map-b") == sched_req_id
|
||||
assert self.scheduler.get_sched_req_id("map-parent") == sched_req_id
|
||||
|
||||
self.scheduler.pop_request_state("map-parent")
|
||||
self.scheduler.pop_request_state(sched_req_id)
|
||||
|
||||
assert self.scheduler.get_request_state("map-parent") is None
|
||||
assert self.scheduler.get_sched_req_id("map-a") is None
|
||||
assert self.scheduler.get_sched_req_id("map-b") is None
|
||||
assert self.scheduler.get_sched_req_id("map-parent") is None
|
||||
|
||||
def test_duplicate_request_id_is_rejected(self) -> None:
|
||||
self.scheduler.add_request(_make_request("dup"))
|
||||
def test_parent_request_id_registration_failure_rolls_back_child_ids(self) -> None:
|
||||
self.scheduler.add_request(
|
||||
OmniDiffusionRequest(
|
||||
prompts=["prompt_existing"],
|
||||
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
|
||||
request_ids=["existing-child"],
|
||||
request_id="duplicate-parent",
|
||||
)
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="request_id 'dup' is already active"):
|
||||
self.scheduler.add_request(_make_request("dup"))
|
||||
colliding_request = OmniDiffusionRequest(
|
||||
prompts=["prompt_unique"],
|
||||
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
|
||||
request_ids=["unique-child"],
|
||||
request_id="duplicate-parent",
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="duplicate-parent"):
|
||||
self.scheduler.add_request(colliding_request)
|
||||
|
||||
assert self.scheduler.get_request_state("unique-child") is None
|
||||
assert self.scheduler.get_sched_req_id("unique-child") is None
|
||||
|
||||
valid_request = OmniDiffusionRequest(
|
||||
prompts=["prompt_unique"],
|
||||
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
|
||||
request_ids=["unique-child"],
|
||||
request_id="unique-parent",
|
||||
)
|
||||
sched_req_id = self.scheduler.add_request(valid_request)
|
||||
|
||||
assert sched_req_id == "unique-child"
|
||||
assert self.scheduler.get_sched_req_id("unique-child") == sched_req_id
|
||||
assert self.scheduler.get_sched_req_id("unique-parent") == sched_req_id
|
||||
|
||||
|
||||
class TestDiffusionEngine:
|
||||
@@ -625,7 +662,7 @@ class TestStepScheduler:
|
||||
finished = self.scheduler.update_from_output(
|
||||
sched_output,
|
||||
RunnerOutput(
|
||||
request_id=req_id,
|
||||
req_id=req_id,
|
||||
step_index=None,
|
||||
finished=True,
|
||||
result=None,
|
||||
@@ -897,7 +934,7 @@ class TestStepScheduler:
|
||||
|
||||
for _ in range(expected_steps - 1):
|
||||
sched_output = self.scheduler.schedule()
|
||||
assert sched_output.scheduled_request_ids == [req_id]
|
||||
assert sched_output.scheduled_req_ids == [req_id]
|
||||
next_step = request.sampling_params.step_index + 1
|
||||
assert (
|
||||
self.scheduler.update_from_output(
|
||||
@@ -908,7 +945,7 @@ class TestStepScheduler:
|
||||
)
|
||||
|
||||
final_output = self.scheduler.schedule()
|
||||
assert final_output.scheduled_request_ids == [req_id]
|
||||
assert final_output.scheduled_req_ids == [req_id]
|
||||
assert self.scheduler.update_from_output(
|
||||
final_output,
|
||||
_make_step_output(req_id, step_index=expected_steps, finished=True),
|
||||
|
||||
@@ -206,7 +206,7 @@ class _DistributedStepPipeline(CFGParallelMixin):
|
||||
def _make_step_request(num_inference_steps: int = 2):
|
||||
return SimpleNamespace(
|
||||
prompts=["a prompt"],
|
||||
request_id="req-1",
|
||||
request_ids=["req-1"],
|
||||
sampling_params=SimpleNamespace(
|
||||
generator=None,
|
||||
seed=None,
|
||||
@@ -227,7 +227,7 @@ def _make_engine_request(req_id: str = "req-1", num_inference_steps: int = 2) ->
|
||||
return OmniDiffusionRequest(
|
||||
prompts=[f"prompt-{req_id}"],
|
||||
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=num_inference_steps),
|
||||
request_id=req_id,
|
||||
request_ids=[req_id],
|
||||
)
|
||||
|
||||
|
||||
@@ -274,10 +274,10 @@ def _make_distributed_runner(mode: str, device: torch.device):
|
||||
return runner
|
||||
|
||||
|
||||
def _make_scheduler_output(req, request_id="req-1", step_id=0, finished_req_ids=None):
|
||||
def _make_scheduler_output(req, sched_req_id="req-1", step_id=0, finished_req_ids=None):
|
||||
return DiffusionSchedulerOutput(
|
||||
step_id=step_id,
|
||||
scheduled_new_reqs=[NewRequestData(request_id=request_id, req=req)],
|
||||
scheduled_new_reqs=[NewRequestData(sched_req_id=sched_req_id, req=req)],
|
||||
scheduled_cached_reqs=CachedRequestData.make_empty(),
|
||||
finished_req_ids=set() if finished_req_ids is None else set(finished_req_ids),
|
||||
num_running_reqs=1,
|
||||
@@ -287,7 +287,7 @@ def _make_scheduler_output(req, request_id="req-1", step_id=0, finished_req_ids=
|
||||
|
||||
def _make_batch_scheduler_output(reqs, *, step_id=0, finished_req_ids=None):
|
||||
"""Scheduler output for a homogeneous batch (one NewRequestData per req)."""
|
||||
new_reqs = [NewRequestData(request_id=r.request_id, req=r) for r in reqs]
|
||||
new_reqs = [NewRequestData(sched_req_id=r.request_ids[0], req=r) for r in reqs]
|
||||
return DiffusionSchedulerOutput(
|
||||
step_id=step_id,
|
||||
scheduled_new_reqs=new_reqs,
|
||||
@@ -298,11 +298,11 @@ def _make_batch_scheduler_output(reqs, *, step_id=0, finished_req_ids=None):
|
||||
)
|
||||
|
||||
|
||||
def _make_cached_scheduler_output(request_id="req-1", step_id=1, finished_req_ids=None):
|
||||
def _make_cached_scheduler_output(sched_req_id="req-1", step_id=1, finished_req_ids=None):
|
||||
return DiffusionSchedulerOutput(
|
||||
step_id=step_id,
|
||||
scheduled_new_reqs=[],
|
||||
scheduled_cached_reqs=CachedRequestData(request_ids=[request_id]),
|
||||
scheduled_cached_reqs=CachedRequestData(sched_req_ids=[sched_req_id]),
|
||||
finished_req_ids=set() if finished_req_ids is None else set(finished_req_ids),
|
||||
num_running_reqs=1,
|
||||
num_waiting_reqs=0,
|
||||
@@ -359,7 +359,7 @@ def _distributed_step_worker(local_rank: int, world_size: int, mode: str, master
|
||||
runner,
|
||||
_make_scheduler_output(_make_step_request(num_inference_steps=1), step_id=0),
|
||||
)
|
||||
output = result.get_request_output("req-1")
|
||||
output = result.get_req_output("req-1")
|
||||
|
||||
assert output.finished is True
|
||||
assert output.result is not None
|
||||
@@ -384,16 +384,16 @@ class TestRunner:
|
||||
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
|
||||
|
||||
result = DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req, step_id=0))
|
||||
first = result.get_request_output("req-1")
|
||||
assert first.request_id == "req-1"
|
||||
first = result.get_req_output("req-1")
|
||||
assert first.req_id == "req-1"
|
||||
assert first.step_index == 1
|
||||
assert first.finished is False
|
||||
assert first.result is None
|
||||
assert "req-1" in runner.state_cache
|
||||
|
||||
result = DiffusionModelRunner.execute_stepwise(runner, _make_cached_scheduler_output(step_id=1))
|
||||
second = result.get_request_output("req-1")
|
||||
assert second.request_id == "req-1"
|
||||
second = result.get_req_output("req-1")
|
||||
assert second.req_id == "req-1"
|
||||
assert second.step_index == 2
|
||||
assert second.finished is True
|
||||
assert second.result is not None
|
||||
@@ -410,13 +410,13 @@ class TestRunner:
|
||||
runner = _make_runner()
|
||||
req_1 = _make_step_request()
|
||||
req_2 = _make_step_request()
|
||||
req_2.request_id = "req-2"
|
||||
req_2.request_ids = ["req-2"]
|
||||
|
||||
scheduler_output = DiffusionSchedulerOutput(
|
||||
step_id=0,
|
||||
scheduled_new_reqs=[
|
||||
NewRequestData(request_id="req-1", req=req_1),
|
||||
NewRequestData(request_id="req-2", req=req_2),
|
||||
NewRequestData(sched_req_id="req-1", req=req_1),
|
||||
NewRequestData(sched_req_id="req-2", req=req_2),
|
||||
],
|
||||
scheduled_cached_reqs=CachedRequestData.make_empty(),
|
||||
finished_req_ids=set(),
|
||||
@@ -431,7 +431,7 @@ class TestRunner:
|
||||
runner = _make_runner()
|
||||
|
||||
with pytest.raises(ValueError, match="Missing cached state"):
|
||||
DiffusionModelRunner.execute_stepwise(runner, _make_cached_scheduler_output(request_id="req-missing"))
|
||||
DiffusionModelRunner.execute_stepwise(runner, _make_cached_scheduler_output(sched_req_id="req-missing"))
|
||||
|
||||
def test_interrupt_marks_request_finished_and_clears_state(self, monkeypatch):
|
||||
runner = _make_runner()
|
||||
@@ -440,8 +440,8 @@ class TestRunner:
|
||||
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
|
||||
|
||||
result = DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req, step_id=0))
|
||||
output = result.get_request_output("req-1")
|
||||
assert output.request_id == "req-1"
|
||||
output = result.get_req_output("req-1")
|
||||
assert output.req_id == "req-1"
|
||||
assert output.step_index == 0
|
||||
assert output.finished is True
|
||||
assert output.result is not None
|
||||
@@ -515,7 +515,7 @@ def _make_step_worker(lora_manager=None, *, expected_output=None):
|
||||
worker = object.__new__(DiffusionWorker)
|
||||
worker.lora_manager = lora_manager
|
||||
worker._step_lora_state = {}
|
||||
output = expected_output if expected_output is not None else RunnerOutput(request_id="req-1")
|
||||
output = expected_output if expected_output is not None else RunnerOutput(req_id="req-1")
|
||||
worker.model_runner = SimpleNamespace(execute_stepwise=lambda arg: output)
|
||||
return worker
|
||||
|
||||
@@ -525,9 +525,9 @@ class TestWorker:
|
||||
"""DiffusionWorker.execute_stepwise"""
|
||||
|
||||
def test_delegates_to_model_runner(self):
|
||||
expected = RunnerOutput(request_id="req-1", step_index=1, finished=False, result=None)
|
||||
expected = RunnerOutput(req_id="req-1", step_index=1, finished=False, result=None)
|
||||
worker = _make_step_worker(expected_output=expected)
|
||||
scheduler_output = _make_scheduler_output(_make_engine_request("req-1"), request_id="req-1")
|
||||
scheduler_output = _make_scheduler_output(_make_engine_request("req-1"), sched_req_id="req-1")
|
||||
|
||||
output = DiffusionWorker.execute_stepwise(worker, scheduler_output)
|
||||
|
||||
@@ -536,7 +536,7 @@ class TestWorker:
|
||||
def test_deactivates_lora_when_request_has_no_adapter(self):
|
||||
manager = _RecordingLoRAManager()
|
||||
worker = _make_step_worker(lora_manager=manager)
|
||||
scheduler_output = _make_scheduler_output(_make_engine_request("req-1"), request_id="req-1")
|
||||
scheduler_output = _make_scheduler_output(_make_engine_request("req-1"), sched_req_id="req-1")
|
||||
|
||||
DiffusionWorker.execute_stepwise(worker, scheduler_output)
|
||||
|
||||
@@ -552,7 +552,7 @@ class TestWorker:
|
||||
|
||||
manager = _RecordingLoRAManager()
|
||||
worker = _make_step_worker(lora_manager=manager)
|
||||
scheduler_output = _make_scheduler_output(request, request_id="req-1")
|
||||
scheduler_output = _make_scheduler_output(request, sched_req_id="req-1")
|
||||
|
||||
DiffusionWorker.execute_stepwise(worker, scheduler_output)
|
||||
|
||||
@@ -568,8 +568,8 @@ class TestWorker:
|
||||
|
||||
manager = _RecordingLoRAManager()
|
||||
worker = _make_step_worker(lora_manager=manager)
|
||||
first = _make_scheduler_output(request, request_id="req-1")
|
||||
second = _make_cached_scheduler_output(request_id="req-1", step_id=1)
|
||||
first = _make_scheduler_output(request, sched_req_id="req-1")
|
||||
second = _make_cached_scheduler_output(sched_req_id="req-1", step_id=1)
|
||||
|
||||
DiffusionWorker.execute_stepwise(worker, first)
|
||||
DiffusionWorker.execute_stepwise(worker, second)
|
||||
@@ -610,10 +610,10 @@ class TestWorker:
|
||||
next_request.sampling_params.lora_request = lora_request
|
||||
|
||||
worker = _make_step_worker(lora_manager=_RecordingLoRAManager())
|
||||
first = _make_scheduler_output(finishing, request_id="req-1")
|
||||
first = _make_scheduler_output(finishing, sched_req_id="req-1")
|
||||
next_batch = _make_scheduler_output(
|
||||
next_request,
|
||||
request_id="req-2",
|
||||
sched_req_id="req-2",
|
||||
step_id=1,
|
||||
finished_req_ids={"req-1"},
|
||||
)
|
||||
@@ -633,11 +633,11 @@ class TestExecutor:
|
||||
def test_execute_step_passes_through_runner_output(self, mocker: MockerFixture):
|
||||
executor = object.__new__(MultiprocDiffusionExecutor)
|
||||
executor._ensure_open = lambda: None
|
||||
expected = RunnerOutput(request_id="req-step", step_index=1, finished=False, result=None)
|
||||
expected = RunnerOutput(req_id="req-step", step_index=1, finished=False, result=None)
|
||||
executor.collective_rpc = mocker.Mock(return_value=expected)
|
||||
|
||||
request = _make_engine_request("req-step", num_inference_steps=2)
|
||||
scheduler_output = _make_scheduler_output(request, request_id="req-step")
|
||||
scheduler_output = _make_scheduler_output(request, sched_req_id="req-step")
|
||||
|
||||
output = MultiprocDiffusionExecutor.execute_step(executor, scheduler_output)
|
||||
|
||||
@@ -653,7 +653,7 @@ class TestEngine:
|
||||
[
|
||||
(
|
||||
lambda _: RunnerOutput(
|
||||
request_id="req-error",
|
||||
req_id="req-error",
|
||||
step_index=1,
|
||||
finished=True,
|
||||
result=DiffusionOutput(error="boom"),
|
||||
@@ -688,7 +688,7 @@ class TestEngine:
|
||||
call_count["n"] += 1
|
||||
finished = call_count["n"] == 2
|
||||
return RunnerOutput(
|
||||
request_id="req-step",
|
||||
req_id="req-step",
|
||||
step_index=call_count["n"],
|
||||
finished=finished,
|
||||
result=(DiffusionOutput(output=torch.tensor([2.0])) if finished else None),
|
||||
@@ -714,7 +714,7 @@ class TestEngine:
|
||||
step["n"] += 1
|
||||
engine.abort("req-stop")
|
||||
return RunnerOutput(
|
||||
request_id="req-stop",
|
||||
req_id="req-stop",
|
||||
step_index=1,
|
||||
finished=False,
|
||||
result=None,
|
||||
@@ -741,7 +741,7 @@ class TestEngine:
|
||||
assert sched_output == _make_cached_scheduler_output("req-mid", step_id=1)
|
||||
engine.abort("req-mid")
|
||||
return RunnerOutput(
|
||||
request_id="req-mid",
|
||||
req_id="req-mid",
|
||||
step_index=step["n"],
|
||||
finished=False,
|
||||
result=None,
|
||||
@@ -760,7 +760,7 @@ class TestEngine:
|
||||
engine = _make_engine(
|
||||
scheduler,
|
||||
execute_fn=lambda _: RunnerOutput(
|
||||
request_id="req-missing",
|
||||
req_id="req-missing",
|
||||
step_index=1,
|
||||
finished=True,
|
||||
result=None,
|
||||
@@ -777,7 +777,7 @@ class TestEngine:
|
||||
class TestIPC:
|
||||
def test_pack_unpack_runner_output_shm(self):
|
||||
tensor = torch.zeros(300_000, dtype=torch.float32)
|
||||
output = RunnerOutput(request_id="req-1", finished=True, result=DiffusionOutput(output=tensor))
|
||||
output = RunnerOutput(req_id="req-1", finished=True, result=DiffusionOutput(output=tensor))
|
||||
|
||||
packed = pack_diffusion_output_shm(output)
|
||||
assert isinstance(packed.result.output, dict)
|
||||
|
||||
@@ -36,7 +36,7 @@ def _tagged_output(tag: str) -> DiffusionOutput:
|
||||
def _mock_request(tag: str):
|
||||
"""Return a lightweight request object identifiable by *tag*."""
|
||||
return SimpleNamespace(
|
||||
request_id=tag,
|
||||
request_ids=[tag],
|
||||
prompts=[f"prompt_{tag}"],
|
||||
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
|
||||
)
|
||||
@@ -97,8 +97,8 @@ def _start_worker(req_q, res_q, count=2):
|
||||
req = req_q.get(timeout=10)
|
||||
method = req.get("method", "")
|
||||
args = req.get("args", ())
|
||||
if method in {"generate", "execute_model"} and args and hasattr(args[0], "request_id"):
|
||||
tag = f"result_for_{args[0].request_id}"
|
||||
if method in {"generate", "execute_model"} and args and hasattr(args[0], "request_ids"):
|
||||
tag = f"result_for_{args[0].request_ids[0]}"
|
||||
elif args:
|
||||
tag = f"result_for_{args[0]}"
|
||||
else:
|
||||
|
||||
@@ -68,7 +68,7 @@ def test_process_batch_request_preserves_parent_request_id_and_kv_sender_info():
|
||||
|
||||
request = captured["request"]
|
||||
assert request.request_id == "req-parent"
|
||||
assert not hasattr(request, "request_ids")
|
||||
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"]
|
||||
|
||||
@@ -33,20 +33,6 @@ def test_tensor_serialization():
|
||||
assert torch.equal(tensor, deserialized)
|
||||
|
||||
|
||||
def test_noncontiguous_tensor_serialization():
|
||||
"""Test torch.Tensor serialization for strided views."""
|
||||
import torch
|
||||
|
||||
tensor = torch.arange(40, dtype=torch.long).reshape(4, 10).transpose(0, 1)
|
||||
assert not tensor.is_contiguous()
|
||||
assert tensor.stride(-1) != 1
|
||||
|
||||
serialized = OmniSerializer.serialize({"codes": tensor})
|
||||
deserialized = OmniSerializer.deserialize(serialized)
|
||||
|
||||
assert torch.equal(tensor, deserialized["codes"])
|
||||
|
||||
|
||||
def test_ndarray_serialization():
|
||||
"""Test numpy.ndarray serialization."""
|
||||
import numpy as np
|
||||
|
||||
@@ -37,7 +37,6 @@ def _req(req_id: str, status: RequestStatus, external_req_id: str | None = None)
|
||||
status=status,
|
||||
prompt_token_ids=[],
|
||||
num_computed_tokens=0,
|
||||
num_output_placeholders=0,
|
||||
additional_information=None,
|
||||
is_finished=lambda: status == RequestStatus.FINISHED_STOPPED,
|
||||
)
|
||||
@@ -45,16 +44,9 @@ def _req(req_id: str, status: RequestStatus, external_req_id: str | None = None)
|
||||
|
||||
@pytest.fixture
|
||||
def build_adapter(monkeypatch, mocker: MockerFixture):
|
||||
def _build(
|
||||
*,
|
||||
stage_id: int = 1,
|
||||
model_mode: str = "ar",
|
||||
max_num_seqs: int = 2,
|
||||
connector_extra: dict | None = None,
|
||||
):
|
||||
def _build(*, stage_id: int = 1, model_mode: str = "ar", max_num_seqs: int = 2):
|
||||
connector = mocker.MagicMock()
|
||||
connector.stage_id = stage_id
|
||||
connector.config = {"extra": connector_extra or {}}
|
||||
connector.get.return_value = None
|
||||
connector.put.return_value = (True, 1, {})
|
||||
|
||||
@@ -135,61 +127,6 @@ def test_load_poll(build_adapter):
|
||||
assert "req-1" not in adapter._pending_load_reqs
|
||||
|
||||
|
||||
def test_load_poll_generation_tensor_codes_use_placeholder_prompt(build_adapter):
|
||||
adapter, connector = build_adapter(stage_id=1, model_mode="generation")
|
||||
request = _req("req-tensor", RequestStatus.WAITING, external_req_id="external-tensor")
|
||||
|
||||
codes = torch.tensor([[1, 2], [3, 4]], dtype=torch.long)
|
||||
payload: OmniPayload = {
|
||||
"codes": {"audio": codes},
|
||||
"meta": {
|
||||
"left_context_size": 1,
|
||||
"finished": torch.tensor(False, dtype=torch.bool),
|
||||
},
|
||||
}
|
||||
connector.get.return_value = (payload, 16)
|
||||
|
||||
adapter._poll_single_request(request)
|
||||
|
||||
assert request.prompt_token_ids == [0]
|
||||
assert request.num_computed_tokens == 0
|
||||
assert torch.equal(request.additional_information["codes"]["audio"], codes)
|
||||
assert request.additional_information["meta"]["left_context_size"] == 1
|
||||
assert "finished" not in request.additional_information["meta"]
|
||||
assert "req-tensor" in adapter._finished_load_reqs
|
||||
|
||||
|
||||
def test_load_poll_generation_empty_nonterminal_chunk_keeps_polling(build_adapter):
|
||||
adapter, connector = build_adapter(stage_id=1, model_mode="generation")
|
||||
request = _req("req-empty-tensor", RequestStatus.WAITING, external_req_id="external-empty")
|
||||
|
||||
empty_payload: OmniPayload = {
|
||||
"codes": {"audio": torch.empty((4, 0), dtype=torch.long)},
|
||||
"meta": {
|
||||
"left_context_size": 0,
|
||||
"finished": torch.tensor(False, dtype=torch.bool),
|
||||
},
|
||||
}
|
||||
ready_payload: OmniPayload = {
|
||||
"codes": {"audio": torch.tensor([[1, 2]], dtype=torch.long)},
|
||||
"meta": {
|
||||
"left_context_size": 0,
|
||||
"finished": torch.tensor(False, dtype=torch.bool),
|
||||
},
|
||||
}
|
||||
connector.get.side_effect = [(empty_payload, 16), (ready_payload, 16)]
|
||||
|
||||
assert adapter._poll_single_request(request) is False
|
||||
assert request.request_id not in adapter._finished_load_reqs
|
||||
assert request.request_id not in adapter.requests_with_ready_chunks
|
||||
assert adapter.get_req_chunk[request.request_id] == 1
|
||||
|
||||
assert adapter._poll_single_request(request) is True
|
||||
assert request.request_id in adapter._finished_load_reqs
|
||||
assert torch.equal(request.additional_information["codes"]["audio"], ready_payload["codes"]["audio"])
|
||||
assert adapter.get_req_chunk[request.request_id] == 2
|
||||
|
||||
|
||||
def test_save_async(build_adapter):
|
||||
adapter, _ = build_adapter(stage_id=1)
|
||||
request = _req("req-1", RequestStatus.WAITING, external_req_id="external-1")
|
||||
@@ -203,18 +140,6 @@ def test_save_async(build_adapter):
|
||||
assert task["is_finished"] is False
|
||||
|
||||
|
||||
def test_save_async_uses_confirmed_tokens_for_async_scheduler_watermark(build_adapter):
|
||||
adapter, _ = build_adapter(stage_id=1)
|
||||
request = _req("req-async", RequestStatus.WAITING, external_req_id="external-async")
|
||||
request.num_computed_tokens = 10
|
||||
request.num_output_placeholders = 2
|
||||
|
||||
adapter.save_async(pooling_output=None, request=request)
|
||||
|
||||
assert adapter.requests_num_chunks_sent["external-async"] == 8
|
||||
assert len(adapter._pending_save_reqs) == 1
|
||||
|
||||
|
||||
def test_send_single_request_struct_without_meta_does_not_crash(build_adapter, monkeypatch):
|
||||
"""Producer may return a struct with ``meta=None`` (e.g. payload that
|
||||
carries only ``embed`` or ``codes``). The sender's ``meta is not None``
|
||||
@@ -533,52 +458,6 @@ def test_finish_requests_restores_status(build_adapter):
|
||||
assert req_id not in adapter.requests_origin_status
|
||||
|
||||
|
||||
def test_finish_requests_removes_zombies_from_chunk_waiting_deques(build_adapter):
|
||||
adapter, _ = build_adapter(stage_id=1)
|
||||
zombie = _req("req-zombie", RequestStatus.WAITING_FOR_CHUNK)
|
||||
other = _req("req-live", RequestStatus.WAITING_FOR_CHUNK)
|
||||
adapter.waiting_for_chunk_waiting_requests = deque([zombie, other])
|
||||
adapter.waiting_for_chunk_running_requests = deque([other, zombie])
|
||||
adapter.requests_with_ready_chunks.add("req-zombie")
|
||||
adapter.finished_requests.add("req-zombie")
|
||||
requests_map = {
|
||||
"req-zombie": zombie,
|
||||
"req-live": other,
|
||||
}
|
||||
|
||||
adapter.finish_requests(
|
||||
["req-zombie"],
|
||||
RequestStatus.FINISHED_ABORTED,
|
||||
requests_map,
|
||||
)
|
||||
|
||||
assert [req.request_id for req in adapter.waiting_for_chunk_waiting_requests] == ["req-live"]
|
||||
assert [req.request_id for req in adapter.waiting_for_chunk_running_requests] == ["req-live"]
|
||||
assert "req-zombie" not in adapter.requests_with_ready_chunks
|
||||
assert "req-zombie" not in adapter.finished_requests
|
||||
|
||||
|
||||
def test_restore_queues_skips_requests_missing_from_scheduler_requests(build_adapter):
|
||||
adapter, _ = build_adapter(stage_id=1)
|
||||
zombie = _req("req-zombie", RequestStatus.WAITING_FOR_CHUNK)
|
||||
live = _req("req-live", RequestStatus.WAITING_FOR_CHUNK)
|
||||
waiting_queue = DummyWaitingQueue()
|
||||
running_queue = []
|
||||
adapter.waiting_for_chunk_waiting_requests = deque([zombie, live])
|
||||
adapter.waiting_for_chunk_running_requests = deque([zombie, live])
|
||||
|
||||
adapter.restore_queues(
|
||||
waiting_queue,
|
||||
running_queue,
|
||||
scheduler_requests={"req-live": live},
|
||||
)
|
||||
|
||||
assert [req.request_id for req in waiting_queue] == ["req-live"]
|
||||
assert [req.request_id for req in running_queue] == ["req-live"]
|
||||
assert not adapter.waiting_for_chunk_waiting_requests
|
||||
assert not adapter.waiting_for_chunk_running_requests
|
||||
|
||||
|
||||
# ---------------------------------------------------------------
|
||||
# Scheduler trigger tests
|
||||
# ---------------------------------------------------------------
|
||||
|
||||
@@ -384,7 +384,7 @@ def test_manager_reception(kv_config, mock_connector, common_constants):
|
||||
req = OmniDiffusionRequest(
|
||||
prompts=["test_recv"],
|
||||
sampling_params=OmniDiffusionSamplingParams(),
|
||||
request_id=req_id,
|
||||
request_ids=[req_id],
|
||||
)
|
||||
# req.need_kv_receive = True # Implicitly handled by receive_kv_cache check? No, manager doesn't check it, runner does.
|
||||
# But receive_kv_cache in manager checks request_id. Which we need to fix in manager next.
|
||||
@@ -428,6 +428,7 @@ def test_manager_reception_prefers_parent_request_id_for_batched_request(kv_conf
|
||||
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,
|
||||
)
|
||||
|
||||
@@ -453,6 +454,7 @@ def test_receive_multi_kv_cache_uses_parent_request_id_for_cfg_collection(kv_con
|
||||
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"}
|
||||
@@ -514,7 +516,7 @@ def test_integration_flow(common_constants):
|
||||
req = OmniDiffusionRequest(
|
||||
prompts=["test_integ"],
|
||||
sampling_params=OmniDiffusionSamplingParams(),
|
||||
request_id=req_id,
|
||||
request_ids=[req_id],
|
||||
)
|
||||
|
||||
# Receive
|
||||
|
||||
@@ -62,19 +62,6 @@ def pytest_addoption(parser):
|
||||
default=1200,
|
||||
help="Online serving timeout in seconds for Wan2.2 I2V accuracy tests.",
|
||||
)
|
||||
group.addoption(
|
||||
"--hunyuanvideo15-i2v-image-source",
|
||||
action="store",
|
||||
default=None,
|
||||
help="Image source for HunyuanVideo-1.5 I2V accuracy tests. Can be local path or remote URL",
|
||||
)
|
||||
group.addoption(
|
||||
"--hunyuanvideo15-online-timeout-seconds",
|
||||
action="store",
|
||||
type=int,
|
||||
default=3600,
|
||||
help="Online serving timeout in seconds for HunyuanVideo-1.5 accuracy tests.",
|
||||
)
|
||||
|
||||
|
||||
def _hf_cache_root() -> Path:
|
||||
@@ -181,17 +168,6 @@ def wan22_i2v_online_timeout_seconds(request: pytest.FixtureRequest) -> int:
|
||||
return int(request.config.getoption("wan22_i2v_online_timeout_seconds"))
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def hunyuanvideo15_i2v_image_source(request: pytest.FixtureRequest) -> str | None:
|
||||
value = request.config.getoption("hunyuanvideo15_i2v_image_source")
|
||||
return str(value) if value else None
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def hunyuanvideo15_online_timeout_seconds(request: pytest.FixtureRequest) -> int:
|
||||
return int(request.config.getoption("hunyuanvideo15_online_timeout_seconds"))
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def gebench_samples_per_type(request: pytest.FixtureRequest) -> int:
|
||||
return int(request.config.getoption("gebench_samples_per_type"))
|
||||
|
||||
@@ -1,17 +1,7 @@
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import subprocess
|
||||
import time
|
||||
from base64 import b64decode, b64encode
|
||||
from fractions import Fraction
|
||||
from hashlib import sha1
|
||||
from io import BytesIO
|
||||
from mimetypes import guess_type
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
@@ -73,274 +63,6 @@ def model_output_dir(parent_dir: Path, model: str) -> Path:
|
||||
return path
|
||||
|
||||
|
||||
_SSIM_RE = re.compile(r"All:(?P<score>[0-9.]+)")
|
||||
_PSNR_RE = re.compile(r"average:(?P<score>[0-9.]+)")
|
||||
|
||||
|
||||
def parse_video_metadata(payload: dict[str, Any]) -> dict[str, int | float]:
|
||||
streams = payload.get("streams")
|
||||
if not isinstance(streams, list) or not streams:
|
||||
raise ValueError(f"ffprobe payload did not include video streams: {payload}")
|
||||
|
||||
stream = streams[0]
|
||||
width = int(stream["width"])
|
||||
height = int(stream["height"])
|
||||
fps = float(Fraction(stream["avg_frame_rate"]))
|
||||
frame_count_value = stream.get("nb_read_frames") or stream.get("nb_frames")
|
||||
if frame_count_value is None:
|
||||
raise ValueError(f"ffprobe payload did not include frame count: {payload}")
|
||||
frame_count = int(frame_count_value)
|
||||
return {
|
||||
"width": width,
|
||||
"height": height,
|
||||
"fps": fps,
|
||||
"frame_count": frame_count,
|
||||
}
|
||||
|
||||
|
||||
def parse_ssim_score(output: str) -> float:
|
||||
match = _SSIM_RE.search(output)
|
||||
if match is None:
|
||||
raise ValueError(f"Could not parse SSIM score from ffmpeg output:\n{output}")
|
||||
return float(match.group("score"))
|
||||
|
||||
|
||||
def parse_psnr_score(output: str) -> float:
|
||||
match = _PSNR_RE.search(output)
|
||||
if match is None:
|
||||
raise ValueError(f"Could not parse PSNR score from ffmpeg output:\n{output}")
|
||||
return float(match.group("score"))
|
||||
|
||||
|
||||
def probe_binary(binary: str, *, test_name: str = "video similarity e2e test") -> str:
|
||||
resolved = shutil.which(binary)
|
||||
if resolved is None:
|
||||
pytest.skip(f"{binary} is required for {test_name}.")
|
||||
return resolved
|
||||
|
||||
|
||||
def probe_video(path: Path) -> dict[str, int | float]:
|
||||
probe_binary("ffprobe")
|
||||
result = subprocess.run(
|
||||
[
|
||||
"ffprobe",
|
||||
"-v",
|
||||
"error",
|
||||
"-select_streams",
|
||||
"v:0",
|
||||
"-count_frames",
|
||||
"-show_entries",
|
||||
"stream=width,height,avg_frame_rate,nb_read_frames",
|
||||
"-of",
|
||||
"json",
|
||||
str(path),
|
||||
],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=True,
|
||||
)
|
||||
return parse_video_metadata(json.loads(result.stdout))
|
||||
|
||||
|
||||
def run_ffmpeg_similarity(filter_name: str, first: Path, second: Path) -> str:
|
||||
probe_binary("ffmpeg")
|
||||
result = subprocess.run(
|
||||
[
|
||||
"ffmpeg",
|
||||
"-hide_banner",
|
||||
"-i",
|
||||
str(first),
|
||||
"-i",
|
||||
str(second),
|
||||
"-lavfi",
|
||||
f"[0:v][1:v]{filter_name}",
|
||||
"-f",
|
||||
"null",
|
||||
"-",
|
||||
],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=True,
|
||||
)
|
||||
return result.stderr
|
||||
|
||||
|
||||
def online_timeout_seconds(configured: int | None, *, default: int = 1200) -> int:
|
||||
return int(configured) if configured is not None else default
|
||||
|
||||
|
||||
def is_remote_image_source(source: str) -> bool:
|
||||
return source.startswith("http://") or source.startswith("https://")
|
||||
|
||||
|
||||
def validate_image_source(source: str) -> None:
|
||||
if is_remote_image_source(source):
|
||||
response = requests.get(source, timeout=60)
|
||||
response.raise_for_status()
|
||||
return
|
||||
|
||||
image_path = Path(source)
|
||||
if not image_path.exists():
|
||||
raise FileNotFoundError(f"Local image source does not exist: {image_path}")
|
||||
|
||||
|
||||
def build_online_image_reference(source: str) -> str:
|
||||
if source.startswith("data:image"):
|
||||
return source
|
||||
|
||||
if is_remote_image_source(source):
|
||||
response = requests.get(source, timeout=60)
|
||||
response.raise_for_status()
|
||||
mime_type = response.headers.get("Content-Type", "").split(";", 1)[0].strip()
|
||||
if not mime_type.startswith("image/"):
|
||||
mime_type = guess_type(urlparse(source).path)[0] or "image/png"
|
||||
encoded = b64encode(response.content).decode("ascii")
|
||||
return f"data:{mime_type};base64,{encoded}"
|
||||
|
||||
image_path = Path(source)
|
||||
mime_type = guess_type(image_path.name)[0] or "image/png"
|
||||
encoded = b64encode(image_path.read_bytes()).decode("ascii")
|
||||
return f"data:{mime_type};base64,{encoded}"
|
||||
|
||||
|
||||
def materialize_image_source(source: str, output_dir: Path) -> str:
|
||||
if not is_remote_image_source(source):
|
||||
return source
|
||||
|
||||
from diffusers.utils import load_image
|
||||
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
source_name = Path(urlparse(source).path).name or "input.png"
|
||||
safe_name = re.sub(r"[^A-Za-z0-9_.-]+", "_", source_name)
|
||||
if "." not in safe_name:
|
||||
safe_name = f"{safe_name}.png"
|
||||
image_path = output_dir / safe_name
|
||||
if not image_path.exists():
|
||||
image = load_image(source).convert("RGB")
|
||||
image.save(image_path)
|
||||
return str(image_path)
|
||||
|
||||
|
||||
def video_artifact_dir(result_root: Path, source: str) -> Path:
|
||||
if is_remote_image_source(source):
|
||||
source_name = Path(urlparse(source).path).stem or "remote"
|
||||
else:
|
||||
source_name = Path(source).stem or "local"
|
||||
|
||||
safe_name = re.sub(r"[^A-Za-z0-9_.-]+", "_", source_name)
|
||||
digest = sha1(source.encode("utf-8")).hexdigest()[:8]
|
||||
return result_root / f"{safe_name}-{digest}"
|
||||
|
||||
|
||||
def send_video_request_with_timeout(
|
||||
openai_client,
|
||||
request_config: dict[str, Any],
|
||||
*,
|
||||
timeout_seconds: int,
|
||||
) -> bytes:
|
||||
form_data = request_config.get("form_data")
|
||||
if not isinstance(form_data, dict):
|
||||
raise ValueError("Video request_config must contain 'form_data'")
|
||||
|
||||
if not form_data.get("prompt"):
|
||||
raise ValueError("Video request_config['form_data'] must contain 'prompt'")
|
||||
|
||||
normalized_form_data = {key: str(value) for key, value in form_data.items() if value is not None}
|
||||
|
||||
files: dict[str, tuple[str, BytesIO, str]] = {}
|
||||
image_reference = request_config.get("image_reference")
|
||||
if image_reference:
|
||||
if image_reference.startswith("data:image"):
|
||||
header, encoded = image_reference.split(",", 1)
|
||||
content_type = header.split(";")[0].removeprefix("data:")
|
||||
extension = content_type.split("/")[-1]
|
||||
file_data = b64decode(encoded)
|
||||
files["input_reference"] = (
|
||||
f"reference.{extension}",
|
||||
BytesIO(file_data),
|
||||
content_type,
|
||||
)
|
||||
else:
|
||||
normalized_form_data["image_reference"] = json.dumps({"image_url": image_reference})
|
||||
|
||||
start_time = time.perf_counter()
|
||||
create_url = openai_client._build_url("/v1/videos")
|
||||
response = requests.post(
|
||||
create_url,
|
||||
data=normalized_form_data,
|
||||
files=files,
|
||||
headers={"Accept": "application/json"},
|
||||
timeout=60,
|
||||
)
|
||||
try:
|
||||
response.raise_for_status()
|
||||
except requests.HTTPError as exc:
|
||||
raise requests.HTTPError(f"{exc}; response body: {response.text}", response=response) from exc
|
||||
job_data = response.json()
|
||||
video_id = job_data["id"]
|
||||
openai_client._wait_until_video_completed(video_id, timeout_seconds=timeout_seconds)
|
||||
video_content = openai_client._download_video_content(video_id)
|
||||
print(f"online_video_e2e_latency_s={time.perf_counter() - start_time:.3f}")
|
||||
return video_content
|
||||
|
||||
|
||||
def assert_video_metadata(
|
||||
metadata: dict[str, int | float],
|
||||
*,
|
||||
width: int,
|
||||
height: int,
|
||||
fps: int | float,
|
||||
frame_count: int,
|
||||
) -> None:
|
||||
assert metadata["width"] == width
|
||||
assert metadata["height"] == height
|
||||
assert metadata["fps"] == float(fps)
|
||||
assert metadata["frame_count"] == frame_count
|
||||
|
||||
|
||||
def assert_video_similarity_metrics(
|
||||
*,
|
||||
label: str,
|
||||
online_path: Path,
|
||||
offline_path: Path,
|
||||
ssim_threshold: float,
|
||||
psnr_threshold: float,
|
||||
) -> None:
|
||||
ssim_output = run_ffmpeg_similarity("ssim", online_path, offline_path)
|
||||
psnr_output = run_ffmpeg_similarity("psnr", online_path, offline_path)
|
||||
ssim_score = parse_ssim_score(ssim_output)
|
||||
psnr_score = parse_psnr_score(psnr_output)
|
||||
|
||||
print(f"{label} similarity metrics:")
|
||||
print(
|
||||
" SSIM:"
|
||||
f" value={ssim_score:.6f},"
|
||||
f" threshold>={ssim_threshold:.6f},"
|
||||
" range=[0, 1],"
|
||||
" higher_is_better=True,"
|
||||
" interpretation=structural_similarity"
|
||||
)
|
||||
print(
|
||||
" PSNR:"
|
||||
f" value={psnr_score:.6f} dB,"
|
||||
f" threshold>={psnr_threshold:.6f} dB,"
|
||||
" range=[0, +inf),"
|
||||
" higher_is_better=True,"
|
||||
" interpretation=pixel_error_in_decibels"
|
||||
)
|
||||
print(f"online_video={online_path}")
|
||||
print(f"offline_video={offline_path}")
|
||||
|
||||
assert ssim_score >= ssim_threshold, (
|
||||
f"SSIM below threshold: got {ssim_score:.6f}, expected >= {ssim_threshold:.6f}. "
|
||||
f"online={online_path} offline={offline_path}"
|
||||
)
|
||||
assert psnr_score >= psnr_threshold, (
|
||||
f"PSNR below threshold: got {psnr_score:.6f}, expected >= {psnr_threshold:.6f}. "
|
||||
f"online={online_path} offline={offline_path}"
|
||||
)
|
||||
|
||||
|
||||
def assert_similarity(
|
||||
*,
|
||||
model_name: str,
|
||||
|
||||
@@ -1 +0,0 @@
|
||||
"""HunyuanVideo-1.5 I2V accuracy test cases."""
|
||||
@@ -1,262 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.e2e.accuracy.helpers import (
|
||||
assert_video_metadata,
|
||||
assert_video_similarity_metrics,
|
||||
build_online_image_reference,
|
||||
materialize_image_source,
|
||||
probe_binary,
|
||||
probe_video,
|
||||
send_video_request_with_timeout,
|
||||
validate_image_source,
|
||||
video_artifact_dir,
|
||||
)
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.runtime import OmniServerParams
|
||||
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
|
||||
|
||||
MODEL_NAME = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v"
|
||||
IMAGE_URL = "https://huggingface.co/datasets/YiYiXu/testing-images/resolve/main/wan_i2v_input.JPG"
|
||||
PROMPT = (
|
||||
"Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. "
|
||||
"The fluffy-furred feline gazes directly at the camera with a relaxed expression. "
|
||||
"Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, "
|
||||
"and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, "
|
||||
"as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's "
|
||||
"intricate details and the refreshing atmosphere of the seaside."
|
||||
)
|
||||
WIDTH = 832
|
||||
HEIGHT = 480
|
||||
FPS = 24
|
||||
NUM_FRAMES = 33
|
||||
NUM_INFERENCE_STEPS = 50
|
||||
GUIDANCE_SCALE = 6.0
|
||||
FLOW_SHIFT = 5.0
|
||||
SEED = 42
|
||||
SSIM_THRESHOLD = 0.94
|
||||
PSNR_THRESHOLD = 28.0
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[4]
|
||||
WORKSPACE_ROOT = REPO_ROOT.parent
|
||||
RUNNER_PATH = REPO_ROOT / "examples" / "offline_inference" / "image_to_video" / "image_to_video.py"
|
||||
RESULT_ROOT = Path(__file__).parent / "result"
|
||||
VIDEO_TIMEOUT_SECONDS = 60 * 60
|
||||
SERVER_CASES = [
|
||||
pytest.param(
|
||||
OmniServerParams(
|
||||
model=MODEL_NAME,
|
||||
server_args=[
|
||||
"--flow-shift",
|
||||
str(FLOW_SHIFT),
|
||||
"--enforce-eager",
|
||||
],
|
||||
env_dict={"VLLM_OMNI_STORAGE_PATH": str(RESULT_ROOT / "storage")},
|
||||
use_omni=True,
|
||||
),
|
||||
id="hunyuanvideo15_i2v_default",
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
def _runner_env() -> dict[str, str]:
|
||||
env = os.environ.copy()
|
||||
pythonpath_parts = [
|
||||
str(REPO_ROOT),
|
||||
str(WORKSPACE_ROOT / "diffusers" / "src"),
|
||||
]
|
||||
existing_pythonpath = env.get("PYTHONPATH")
|
||||
if existing_pythonpath:
|
||||
pythonpath_parts.append(existing_pythonpath)
|
||||
env["PYTHONPATH"] = os.pathsep.join(pythonpath_parts)
|
||||
env.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
|
||||
return env
|
||||
|
||||
|
||||
def _resolve_image_source(configured: str | None) -> str:
|
||||
if configured:
|
||||
candidate = Path(configured)
|
||||
if candidate.exists():
|
||||
return str(candidate.resolve())
|
||||
return configured
|
||||
return IMAGE_URL
|
||||
|
||||
|
||||
def _artifact_paths(image_source: str) -> tuple[Path, Path]:
|
||||
artifact_dir = video_artifact_dir(RESULT_ROOT, image_source)
|
||||
artifact_dir.mkdir(parents=True, exist_ok=True)
|
||||
return artifact_dir / "online.mp4", artifact_dir / "offline.mp4"
|
||||
|
||||
|
||||
def _offline_image_source(image_source: str) -> str:
|
||||
artifact_dir = video_artifact_dir(RESULT_ROOT, image_source)
|
||||
return materialize_image_source(image_source, artifact_dir)
|
||||
|
||||
|
||||
def _build_offline_command(*, image_source: str, output_path: Path) -> list[str]:
|
||||
return [
|
||||
sys.executable,
|
||||
str(RUNNER_PATH),
|
||||
"--model",
|
||||
MODEL_NAME,
|
||||
"--image",
|
||||
image_source,
|
||||
"--prompt",
|
||||
PROMPT,
|
||||
"--height",
|
||||
str(HEIGHT),
|
||||
"--width",
|
||||
str(WIDTH),
|
||||
"--num-frames",
|
||||
str(NUM_FRAMES),
|
||||
"--num-inference-steps",
|
||||
str(NUM_INFERENCE_STEPS),
|
||||
"--guidance-scale",
|
||||
str(GUIDANCE_SCALE),
|
||||
"--flow-shift",
|
||||
str(FLOW_SHIFT),
|
||||
"--fps",
|
||||
str(FPS),
|
||||
"--seed",
|
||||
str(SEED),
|
||||
"--enforce-eager",
|
||||
"--output",
|
||||
str(output_path),
|
||||
]
|
||||
|
||||
|
||||
def _generate_offline_video(*, image_source: str) -> Path:
|
||||
_, offline_path = _artifact_paths(image_source)
|
||||
offline_image_source = _offline_image_source(image_source)
|
||||
subprocess.run(
|
||||
_build_offline_command(image_source=offline_image_source, output_path=offline_path),
|
||||
cwd=REPO_ROOT,
|
||||
env=_runner_env(),
|
||||
check=True,
|
||||
timeout=VIDEO_TIMEOUT_SECONDS,
|
||||
)
|
||||
return offline_path
|
||||
|
||||
|
||||
def _generate_online_video(
|
||||
*,
|
||||
omni_server,
|
||||
openai_client,
|
||||
image_source: str,
|
||||
timeout_seconds: int,
|
||||
) -> Path:
|
||||
online_path, _ = _artifact_paths(image_source)
|
||||
online_image_source = _offline_image_source(image_source)
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"form_data": {
|
||||
"prompt": PROMPT,
|
||||
"height": HEIGHT,
|
||||
"width": WIDTH,
|
||||
"fps": FPS,
|
||||
"num_frames": NUM_FRAMES,
|
||||
"guidance_scale": GUIDANCE_SCALE,
|
||||
"flow_shift": FLOW_SHIFT,
|
||||
"num_inference_steps": NUM_INFERENCE_STEPS,
|
||||
"seed": SEED,
|
||||
},
|
||||
"image_reference": build_online_image_reference(online_image_source),
|
||||
}
|
||||
online_video_bytes = send_video_request_with_timeout(
|
||||
openai_client,
|
||||
request_config,
|
||||
timeout_seconds=timeout_seconds,
|
||||
)
|
||||
online_path.write_bytes(online_video_bytes)
|
||||
return online_path
|
||||
|
||||
|
||||
@pytest.mark.benchmark
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
def test_hunyuanvideo15_i2v_diffusers_offline_generates_video(
|
||||
hunyuanvideo15_i2v_image_source: str | None,
|
||||
) -> None:
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip("HunyuanVideo-1.5 I2V offline accuracy test requires CUDA.")
|
||||
|
||||
probe_binary("ffprobe")
|
||||
if not RUNNER_PATH.exists():
|
||||
raise AssertionError(f"Offline runner does not exist: {RUNNER_PATH}")
|
||||
|
||||
image_source = _resolve_image_source(hunyuanvideo15_i2v_image_source)
|
||||
validate_image_source(image_source)
|
||||
offline_path = _generate_offline_video(image_source=image_source)
|
||||
assert offline_path.exists(), f"Expected offline video artifact at {offline_path}"
|
||||
offline_metadata = probe_video(offline_path)
|
||||
assert_video_metadata(offline_metadata, width=WIDTH, height=HEIGHT, fps=FPS, frame_count=NUM_FRAMES)
|
||||
|
||||
|
||||
@pytest.mark.benchmark
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
@pytest.mark.parametrize("omni_server", SERVER_CASES, indirect=True)
|
||||
def test_hunyuanvideo15_i2v_online_serving_generates_video(
|
||||
omni_server,
|
||||
openai_client,
|
||||
hunyuanvideo15_i2v_image_source: str | None,
|
||||
hunyuanvideo15_online_timeout_seconds: int,
|
||||
) -> None:
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip("HunyuanVideo-1.5 I2V online accuracy test requires CUDA.")
|
||||
|
||||
probe_binary("ffprobe")
|
||||
image_source = _resolve_image_source(hunyuanvideo15_i2v_image_source)
|
||||
validate_image_source(image_source)
|
||||
online_path = _generate_online_video(
|
||||
omni_server=omni_server,
|
||||
openai_client=openai_client,
|
||||
image_source=image_source,
|
||||
timeout_seconds=hunyuanvideo15_online_timeout_seconds,
|
||||
)
|
||||
assert online_path.exists(), f"Expected online video artifact at {online_path}"
|
||||
online_metadata = probe_video(online_path)
|
||||
assert_video_metadata(online_metadata, width=WIDTH, height=HEIGHT, fps=FPS, frame_count=NUM_FRAMES)
|
||||
|
||||
|
||||
@pytest.mark.benchmark
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
def test_hunyuanvideo15_i2v_serving_matches_offline_video_similarity(
|
||||
hunyuanvideo15_i2v_image_source: str | None,
|
||||
) -> None:
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip("HunyuanVideo-1.5 I2V video similarity test requires CUDA.")
|
||||
|
||||
probe_binary("ffmpeg")
|
||||
probe_binary("ffprobe")
|
||||
image_source = _resolve_image_source(hunyuanvideo15_i2v_image_source)
|
||||
validate_image_source(image_source)
|
||||
online_path, offline_path = _artifact_paths(image_source)
|
||||
if not online_path.exists():
|
||||
pytest.skip(f"Missing online artifact from prerequisite test: {online_path}")
|
||||
if not offline_path.exists():
|
||||
pytest.skip(f"Missing offline artifact from prerequisite test: {offline_path}")
|
||||
|
||||
online_metadata = probe_video(online_path)
|
||||
offline_metadata = probe_video(offline_path)
|
||||
assert online_metadata == offline_metadata, (
|
||||
f"Video metadata mismatch:\n"
|
||||
f"online={online_metadata}\n"
|
||||
f"offline={offline_metadata}\n"
|
||||
f"online_path={online_path}\n"
|
||||
f"offline_path={offline_path}"
|
||||
)
|
||||
assert_video_metadata(online_metadata, width=WIDTH, height=HEIGHT, fps=FPS, frame_count=NUM_FRAMES)
|
||||
assert_video_similarity_metrics(
|
||||
label="hunyuanvideo15_i2v",
|
||||
online_path=online_path,
|
||||
offline_path=offline_path,
|
||||
ssim_threshold=SSIM_THRESHOLD,
|
||||
psnr_threshold=PSNR_THRESHOLD,
|
||||
)
|
||||
@@ -1 +0,0 @@
|
||||
"""HunyuanVideo-1.5 T2V accuracy test cases."""
|
||||
@@ -1,211 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.e2e.accuracy.helpers import (
|
||||
assert_video_metadata,
|
||||
assert_video_similarity_metrics,
|
||||
probe_binary,
|
||||
probe_video,
|
||||
send_video_request_with_timeout,
|
||||
)
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.runtime import OmniServerParams
|
||||
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
|
||||
|
||||
MODEL_NAME = "hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v"
|
||||
PROMPT = "A cat walks on the grass, realistic style."
|
||||
WIDTH = 832
|
||||
HEIGHT = 480
|
||||
FPS = 24
|
||||
NUM_FRAMES = 33
|
||||
NUM_INFERENCE_STEPS = 50
|
||||
GUIDANCE_SCALE = 6.0
|
||||
FLOW_SHIFT = 5.0
|
||||
SEED = 42
|
||||
SSIM_THRESHOLD = 0.94
|
||||
PSNR_THRESHOLD = 28.0
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[4]
|
||||
WORKSPACE_ROOT = REPO_ROOT.parent
|
||||
RUNNER_PATH = REPO_ROOT / "examples" / "offline_inference" / "text_to_video" / "text_to_video.py"
|
||||
RESULT_ROOT = Path(__file__).parent / "result"
|
||||
ARTIFACT_DIR = RESULT_ROOT / "cat_grass"
|
||||
VIDEO_TIMEOUT_SECONDS = 60 * 60
|
||||
SERVER_CASES = [
|
||||
pytest.param(
|
||||
OmniServerParams(
|
||||
model=MODEL_NAME,
|
||||
server_args=[
|
||||
"--flow-shift",
|
||||
str(FLOW_SHIFT),
|
||||
],
|
||||
env_dict={"VLLM_OMNI_STORAGE_PATH": str(RESULT_ROOT / "storage")},
|
||||
use_omni=True,
|
||||
),
|
||||
id="hunyuanvideo15_t2v_default",
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
def _runner_env() -> dict[str, str]:
|
||||
env = os.environ.copy()
|
||||
pythonpath_parts = [
|
||||
str(REPO_ROOT),
|
||||
str(WORKSPACE_ROOT / "diffusers" / "src"),
|
||||
]
|
||||
existing_pythonpath = env.get("PYTHONPATH")
|
||||
if existing_pythonpath:
|
||||
pythonpath_parts.append(existing_pythonpath)
|
||||
env["PYTHONPATH"] = os.pathsep.join(pythonpath_parts)
|
||||
env.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
|
||||
return env
|
||||
|
||||
|
||||
def _artifact_paths() -> tuple[Path, Path]:
|
||||
ARTIFACT_DIR.mkdir(parents=True, exist_ok=True)
|
||||
return ARTIFACT_DIR / "online.mp4", ARTIFACT_DIR / "offline.mp4"
|
||||
|
||||
|
||||
def _build_offline_command(*, output_path: Path) -> list[str]:
|
||||
return [
|
||||
sys.executable,
|
||||
str(RUNNER_PATH),
|
||||
"--model",
|
||||
MODEL_NAME,
|
||||
"--prompt",
|
||||
PROMPT,
|
||||
"--height",
|
||||
str(HEIGHT),
|
||||
"--width",
|
||||
str(WIDTH),
|
||||
"--num-frames",
|
||||
str(NUM_FRAMES),
|
||||
"--num-inference-steps",
|
||||
str(NUM_INFERENCE_STEPS),
|
||||
"--guidance-scale",
|
||||
str(GUIDANCE_SCALE),
|
||||
"--flow-shift",
|
||||
str(FLOW_SHIFT),
|
||||
"--fps",
|
||||
str(FPS),
|
||||
"--seed",
|
||||
str(SEED),
|
||||
"--output",
|
||||
str(output_path),
|
||||
]
|
||||
|
||||
|
||||
def _generate_offline_video() -> Path:
|
||||
_, offline_path = _artifact_paths()
|
||||
subprocess.run(
|
||||
_build_offline_command(output_path=offline_path),
|
||||
cwd=REPO_ROOT,
|
||||
env=_runner_env(),
|
||||
check=True,
|
||||
timeout=VIDEO_TIMEOUT_SECONDS,
|
||||
)
|
||||
return offline_path
|
||||
|
||||
|
||||
def _generate_online_video(*, omni_server, openai_client, timeout_seconds: int) -> Path:
|
||||
online_path, _ = _artifact_paths()
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"form_data": {
|
||||
"prompt": PROMPT,
|
||||
"height": HEIGHT,
|
||||
"width": WIDTH,
|
||||
"fps": FPS,
|
||||
"num_frames": NUM_FRAMES,
|
||||
"guidance_scale": GUIDANCE_SCALE,
|
||||
"flow_shift": FLOW_SHIFT,
|
||||
"num_inference_steps": NUM_INFERENCE_STEPS,
|
||||
"seed": SEED,
|
||||
},
|
||||
}
|
||||
online_video_bytes = send_video_request_with_timeout(
|
||||
openai_client,
|
||||
request_config,
|
||||
timeout_seconds=timeout_seconds,
|
||||
)
|
||||
online_path.write_bytes(online_video_bytes)
|
||||
return online_path
|
||||
|
||||
|
||||
@pytest.mark.benchmark
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
def test_hunyuanvideo15_t2v_diffusers_offline_generates_video() -> None:
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip("HunyuanVideo-1.5 T2V offline accuracy test requires CUDA.")
|
||||
|
||||
probe_binary("ffprobe")
|
||||
if not RUNNER_PATH.exists():
|
||||
raise AssertionError(f"Offline runner does not exist: {RUNNER_PATH}")
|
||||
|
||||
offline_path = _generate_offline_video()
|
||||
assert offline_path.exists(), f"Expected offline video artifact at {offline_path}"
|
||||
offline_metadata = probe_video(offline_path)
|
||||
assert_video_metadata(offline_metadata, width=WIDTH, height=HEIGHT, fps=FPS, frame_count=NUM_FRAMES)
|
||||
|
||||
|
||||
@pytest.mark.benchmark
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
@pytest.mark.parametrize("omni_server", SERVER_CASES, indirect=True)
|
||||
def test_hunyuanvideo15_t2v_online_serving_generates_video(
|
||||
omni_server,
|
||||
openai_client,
|
||||
hunyuanvideo15_online_timeout_seconds: int,
|
||||
) -> None:
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip("HunyuanVideo-1.5 T2V online accuracy test requires CUDA.")
|
||||
|
||||
probe_binary("ffprobe")
|
||||
online_path = _generate_online_video(
|
||||
omni_server=omni_server,
|
||||
openai_client=openai_client,
|
||||
timeout_seconds=hunyuanvideo15_online_timeout_seconds,
|
||||
)
|
||||
assert online_path.exists(), f"Expected online video artifact at {online_path}"
|
||||
online_metadata = probe_video(online_path)
|
||||
assert_video_metadata(online_metadata, width=WIDTH, height=HEIGHT, fps=FPS, frame_count=NUM_FRAMES)
|
||||
|
||||
|
||||
@pytest.mark.benchmark
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
def test_hunyuanvideo15_t2v_serving_matches_offline_video_similarity() -> None:
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip("HunyuanVideo-1.5 T2V video similarity test requires CUDA.")
|
||||
|
||||
probe_binary("ffmpeg")
|
||||
probe_binary("ffprobe")
|
||||
online_path, offline_path = _artifact_paths()
|
||||
if not online_path.exists():
|
||||
pytest.skip(f"Missing online artifact from prerequisite test: {online_path}")
|
||||
if not offline_path.exists():
|
||||
pytest.skip(f"Missing offline artifact from prerequisite test: {offline_path}")
|
||||
|
||||
online_metadata = probe_video(online_path)
|
||||
offline_metadata = probe_video(offline_path)
|
||||
assert online_metadata == offline_metadata, (
|
||||
f"Video metadata mismatch:\n"
|
||||
f"online={online_metadata}\n"
|
||||
f"offline={offline_metadata}\n"
|
||||
f"online_path={online_path}\n"
|
||||
f"offline_path={offline_path}"
|
||||
)
|
||||
assert_video_metadata(online_metadata, width=WIDTH, height=HEIGHT, fps=FPS, frame_count=NUM_FRAMES)
|
||||
assert_video_similarity_metrics(
|
||||
label="hunyuanvideo15_t2v",
|
||||
online_path=online_path,
|
||||
offline_path=offline_path,
|
||||
ssim_threshold=SSIM_THRESHOLD,
|
||||
psnr_threshold=PSNR_THRESHOLD,
|
||||
)
|
||||
@@ -30,14 +30,10 @@ from tests.e2e.accuracy.helpers import (
|
||||
env_to_apply_ftfy_mock_in_subproc,
|
||||
model_output_dir,
|
||||
)
|
||||
from tests.e2e.accuracy.helpers import (
|
||||
parse_psnr_score as _parse_psnr_score,
|
||||
)
|
||||
from tests.e2e.accuracy.helpers import (
|
||||
parse_ssim_score as _parse_ssim_score,
|
||||
)
|
||||
from tests.e2e.accuracy.helpers import (
|
||||
run_ffmpeg_similarity as _run_ffmpeg_similarity,
|
||||
from tests.e2e.accuracy.wan22_i2v.test_wan22_i2v_video_similarity import (
|
||||
_parse_psnr_score,
|
||||
_parse_ssim_score,
|
||||
_run_ffmpeg_similarity,
|
||||
)
|
||||
from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup
|
||||
from tests.helpers.mark import hardware_test
|
||||
|
||||
@@ -77,9 +77,9 @@ THRESHOLDS = {
|
||||
"text_prefix_match": 10, # First 10 characters must match exactly
|
||||
"cot_semantic_sim": 0.9, # Full CoT semantic similarity
|
||||
# Image comparison
|
||||
"clip_score": 90, # CLIP image semantic similarity
|
||||
"ssim": 0.26, # Structural similarity
|
||||
"psnr": 12.5, # Peak signal-to-noise ratio (dB)
|
||||
"clip_score": 85, # CLIP image semantic similarity
|
||||
"ssim": 0.20, # Structural similarity
|
||||
"psnr": 11.0, # Peak signal-to-noise ratio (dB)
|
||||
}
|
||||
|
||||
QUANT_PROMPT = "A brown and white dog is running on the grass."
|
||||
|
||||
@@ -20,7 +20,7 @@ pytestmark = [pytest.mark.full_model, pytest.mark.diffusion]
|
||||
|
||||
|
||||
SINGLE_MODEL = "Qwen/Qwen-Image-Edit"
|
||||
MULTIPLE_MODEL = "Qwen/Qwen-Image-Edit-2511"
|
||||
MULTIPLE_MODEL = "Qwen/Qwen-Image-Edit-2509"
|
||||
WIDTH = 512
|
||||
HEIGHT = 512
|
||||
NUM_INFERENCE_STEPS = 20
|
||||
@@ -253,6 +253,37 @@ def _diffusers_output_single_image(accuracy_artifact_root: Path, qwen_bear_image
|
||||
)
|
||||
|
||||
|
||||
def _vllm_omni_output_multiple_image(
|
||||
accuracy_artifact_root: Path,
|
||||
qwen_bear_image: Image.Image,
|
||||
rabbit_image: Image.Image,
|
||||
) -> Image.Image:
|
||||
output_dir = model_output_dir(accuracy_artifact_root, MULTIPLE_MODEL)
|
||||
output_path = output_dir / "vllm_omni_multiple.png"
|
||||
with OmniServer(model=MULTIPLE_MODEL, serve_args=SERVER_ARGS) as server:
|
||||
output = _run_vllm_omni_image_edit(
|
||||
omni_server=server,
|
||||
prompt=PROMPT_MULTIPLE_IMAGE,
|
||||
input_images=[qwen_bear_image, rabbit_image],
|
||||
output_path=output_path,
|
||||
)
|
||||
return output
|
||||
|
||||
|
||||
def _diffusers_output_multiple_image(
|
||||
accuracy_artifact_root: Path, qwen_bear_image: Image.Image, rabbit_image: Image.Image
|
||||
) -> Image.Image:
|
||||
output_dir = model_output_dir(accuracy_artifact_root, MULTIPLE_MODEL)
|
||||
output_path = output_dir / "diffusers_multiple.png"
|
||||
return _run_diffusers_image_edit(
|
||||
model=MULTIPLE_MODEL,
|
||||
pipeline_class=QwenImageEditPlusPipeline,
|
||||
prompt=PROMPT_MULTIPLE_IMAGE,
|
||||
input_images=[qwen_bear_image, rabbit_image],
|
||||
output_path=output_path,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.benchmark
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
def test_qwen_image_edit_single_matches_diffusers(
|
||||
@@ -278,6 +309,37 @@ def test_qwen_image_edit_single_matches_diffusers(
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.benchmark
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
@pytest.mark.skip(
|
||||
reason="Skipping as the second image seems to be ignored by the API. Will come back to this later after #2772 is merged."
|
||||
)
|
||||
def test_qwen_image_edit_multiple_matches_diffusers(
|
||||
accuracy_artifact_root: Path,
|
||||
qwen_bear_image: Image.Image,
|
||||
rabbit_image: Image.Image,
|
||||
) -> None:
|
||||
vllm_image = _vllm_omni_output_multiple_image(
|
||||
accuracy_artifact_root=accuracy_artifact_root,
|
||||
qwen_bear_image=qwen_bear_image,
|
||||
rabbit_image=rabbit_image,
|
||||
)
|
||||
diffusers_image = _diffusers_output_multiple_image(
|
||||
accuracy_artifact_root=accuracy_artifact_root,
|
||||
qwen_bear_image=qwen_bear_image,
|
||||
rabbit_image=rabbit_image,
|
||||
)
|
||||
assert_similarity(
|
||||
model_name=MULTIPLE_MODEL,
|
||||
vllm_image=vllm_image,
|
||||
diffusers_image=diffusers_image,
|
||||
width=WIDTH,
|
||||
height=HEIGHT,
|
||||
ssim_threshold=SSIM_THRESHOLD,
|
||||
psnr_threshold=PSNR_THRESHOLD,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.benchmark
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
def test_qwen_image_edit_2511_matches_diffusers_pixelwise(accuracy_artifact_root: Path) -> None:
|
||||
|
||||
@@ -1,10 +1,20 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
from base64 import b64decode
|
||||
import time
|
||||
from base64 import b64decode, b64encode
|
||||
from fractions import Fraction
|
||||
from hashlib import sha1
|
||||
from io import BytesIO
|
||||
from mimetypes import guess_type
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import pytest
|
||||
import requests
|
||||
@@ -12,38 +22,6 @@ import torch
|
||||
from diffusers import UniPCMultistepScheduler
|
||||
from PIL import Image
|
||||
|
||||
from tests.e2e.accuracy.helpers import (
|
||||
assert_video_metadata,
|
||||
assert_video_similarity_metrics,
|
||||
video_artifact_dir,
|
||||
)
|
||||
from tests.e2e.accuracy.helpers import (
|
||||
build_online_image_reference as _build_online_image_reference,
|
||||
)
|
||||
from tests.e2e.accuracy.helpers import (
|
||||
online_timeout_seconds as _online_timeout_seconds,
|
||||
)
|
||||
from tests.e2e.accuracy.helpers import (
|
||||
parse_psnr_score as _parse_psnr_score,
|
||||
)
|
||||
from tests.e2e.accuracy.helpers import (
|
||||
parse_ssim_score as _parse_ssim_score,
|
||||
)
|
||||
from tests.e2e.accuracy.helpers import (
|
||||
parse_video_metadata as _parse_video_metadata,
|
||||
)
|
||||
from tests.e2e.accuracy.helpers import (
|
||||
probe_binary as _probe_binary,
|
||||
)
|
||||
from tests.e2e.accuracy.helpers import (
|
||||
probe_video as _probe_video,
|
||||
)
|
||||
from tests.e2e.accuracy.helpers import (
|
||||
send_video_request_with_timeout as _send_video_request_with_timeout,
|
||||
)
|
||||
from tests.e2e.accuracy.helpers import (
|
||||
validate_image_source as _validate_image_source,
|
||||
)
|
||||
from tests.e2e.accuracy.wan22_i2v.run_wan22_i2v_diffusers_cp import (
|
||||
_configure_scheduler,
|
||||
_ensure_wan_ftfy_fallback,
|
||||
@@ -240,6 +218,8 @@ WORKSPACE_ROOT = REPO_ROOT.parent
|
||||
RUNNER_PATH = Path(__file__).with_name("run_wan22_i2v_diffusers_cp.py")
|
||||
RESULT_ROOT = Path(__file__).parent / "result"
|
||||
VIDEO_TIMEOUT_SECONDS = 60 * 60
|
||||
_SSIM_RE = re.compile(r"All:(?P<score>[0-9.]+)")
|
||||
_PSNR_RE = re.compile(r"average:(?P<score>[0-9.]+)")
|
||||
SERVER_CASES = [
|
||||
pytest.param(
|
||||
OmniServerParams(
|
||||
@@ -258,6 +238,48 @@ SERVER_CASES = [
|
||||
]
|
||||
|
||||
|
||||
def _parse_video_metadata(payload: dict[str, Any]) -> dict[str, int | float]:
|
||||
streams = payload.get("streams")
|
||||
if not isinstance(streams, list) or not streams:
|
||||
raise ValueError(f"ffprobe payload did not include video streams: {payload}")
|
||||
|
||||
stream = streams[0]
|
||||
width = int(stream["width"])
|
||||
height = int(stream["height"])
|
||||
fps = float(Fraction(stream["avg_frame_rate"]))
|
||||
frame_count_value = stream.get("nb_read_frames") or stream.get("nb_frames")
|
||||
if frame_count_value is None:
|
||||
raise ValueError(f"ffprobe payload did not include frame count: {payload}")
|
||||
frame_count = int(frame_count_value)
|
||||
return {
|
||||
"width": width,
|
||||
"height": height,
|
||||
"fps": fps,
|
||||
"frame_count": frame_count,
|
||||
}
|
||||
|
||||
|
||||
def _parse_ssim_score(output: str) -> float:
|
||||
match = _SSIM_RE.search(output)
|
||||
if match is None:
|
||||
raise ValueError(f"Could not parse SSIM score from ffmpeg output:\n{output}")
|
||||
return float(match.group("score"))
|
||||
|
||||
|
||||
def _parse_psnr_score(output: str) -> float:
|
||||
match = _PSNR_RE.search(output)
|
||||
if match is None:
|
||||
raise ValueError(f"Could not parse PSNR score from ffmpeg output:\n{output}")
|
||||
return float(match.group("score"))
|
||||
|
||||
|
||||
def _probe_binary(binary: str) -> str:
|
||||
resolved = shutil.which(binary)
|
||||
if resolved is None:
|
||||
pytest.skip(f"{binary} is required for Wan2.2 video similarity e2e test.")
|
||||
return resolved
|
||||
|
||||
|
||||
def _build_diffusers_command(
|
||||
*,
|
||||
runner_path: Path,
|
||||
@@ -321,8 +343,90 @@ def _resolve_image_source(configured: str | None) -> str:
|
||||
return configured
|
||||
|
||||
|
||||
def _is_remote_image_source(source: str) -> bool:
|
||||
return source.startswith("http://") or source.startswith("https://")
|
||||
|
||||
|
||||
def _validate_image_source(source: str) -> None:
|
||||
if _is_remote_image_source(source):
|
||||
response = requests.get(source, timeout=60)
|
||||
response.raise_for_status()
|
||||
return
|
||||
|
||||
image_path = Path(source)
|
||||
if not image_path.exists():
|
||||
raise FileNotFoundError(f"Local image source does not exist: {image_path}")
|
||||
|
||||
|
||||
def _build_online_image_reference(source: str) -> str:
|
||||
if _is_remote_image_source(source) or source.startswith("data:image"):
|
||||
return source
|
||||
|
||||
image_path = Path(source)
|
||||
mime_type = guess_type(image_path.name)[0] or "image/png"
|
||||
encoded = b64encode(image_path.read_bytes()).decode("ascii")
|
||||
return f"data:{mime_type};base64,{encoded}"
|
||||
|
||||
|
||||
def _probe_video(path: Path) -> dict[str, int | float]:
|
||||
_probe_binary("ffprobe")
|
||||
result = subprocess.run(
|
||||
[
|
||||
"ffprobe",
|
||||
"-v",
|
||||
"error",
|
||||
"-select_streams",
|
||||
"v:0",
|
||||
"-count_frames",
|
||||
"-show_entries",
|
||||
"stream=width,height,avg_frame_rate,nb_read_frames",
|
||||
"-of",
|
||||
"json",
|
||||
str(path),
|
||||
],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=True,
|
||||
)
|
||||
return _parse_video_metadata(json.loads(result.stdout))
|
||||
|
||||
|
||||
def _run_ffmpeg_similarity(filter_name: str, first: Path, second: Path) -> str:
|
||||
_probe_binary("ffmpeg")
|
||||
result = subprocess.run(
|
||||
[
|
||||
"ffmpeg",
|
||||
"-hide_banner",
|
||||
"-i",
|
||||
str(first),
|
||||
"-i",
|
||||
str(second),
|
||||
"-lavfi",
|
||||
f"[0:v][1:v]{filter_name}",
|
||||
"-f",
|
||||
"null",
|
||||
"-",
|
||||
],
|
||||
capture_output=True,
|
||||
text=True,
|
||||
check=True,
|
||||
)
|
||||
return result.stderr
|
||||
|
||||
|
||||
def _online_timeout_seconds(configured: int | None) -> int:
|
||||
return int(configured) if configured is not None else 1200
|
||||
|
||||
|
||||
def _artifact_dir(image_source: str) -> Path:
|
||||
return video_artifact_dir(RESULT_ROOT, image_source)
|
||||
if _is_remote_image_source(image_source):
|
||||
source_name = Path(urlparse(image_source).path).stem or "remote"
|
||||
else:
|
||||
source_name = Path(image_source).stem or "local"
|
||||
|
||||
safe_name = re.sub(r"[^A-Za-z0-9_.-]+", "_", source_name)
|
||||
digest = sha1(image_source.encode("utf-8")).hexdigest()[:8]
|
||||
return RESULT_ROOT / f"{safe_name}-{digest}"
|
||||
|
||||
|
||||
def _artifact_paths(image_source: str) -> tuple[Path, Path, Path]:
|
||||
@@ -335,6 +439,55 @@ def _artifact_paths(image_source: str) -> tuple[Path, Path, Path]:
|
||||
)
|
||||
|
||||
|
||||
def _send_video_request_with_timeout(
|
||||
openai_client,
|
||||
request_config: dict[str, Any],
|
||||
*,
|
||||
timeout_seconds: int,
|
||||
) -> bytes:
|
||||
form_data = request_config.get("form_data")
|
||||
if not isinstance(form_data, dict):
|
||||
raise ValueError("Video request_config must contain 'form_data'")
|
||||
|
||||
if not form_data.get("prompt"):
|
||||
raise ValueError("Video request_config['form_data'] must contain 'prompt'")
|
||||
|
||||
normalized_form_data = {key: str(value) for key, value in form_data.items() if value is not None}
|
||||
|
||||
files: dict[str, tuple[str, BytesIO, str]] = {}
|
||||
image_reference = request_config.get("image_reference")
|
||||
if image_reference:
|
||||
if image_reference.startswith("data:image"):
|
||||
header, encoded = image_reference.split(",", 1)
|
||||
content_type = header.split(";")[0].removeprefix("data:")
|
||||
extension = content_type.split("/")[-1]
|
||||
file_data = b64decode(encoded)
|
||||
files["input_reference"] = (
|
||||
f"reference.{extension}",
|
||||
BytesIO(file_data),
|
||||
content_type,
|
||||
)
|
||||
else:
|
||||
normalized_form_data["image_reference"] = json.dumps({"image_url": image_reference})
|
||||
|
||||
start_time = time.perf_counter()
|
||||
create_url = openai_client._build_url("/v1/videos")
|
||||
response = requests.post(
|
||||
create_url,
|
||||
data=normalized_form_data,
|
||||
files=files,
|
||||
headers={"Accept": "application/json"},
|
||||
timeout=60,
|
||||
)
|
||||
response.raise_for_status()
|
||||
job_data = response.json()
|
||||
video_id = job_data["id"]
|
||||
openai_client._wait_until_video_completed(video_id, timeout_seconds=timeout_seconds)
|
||||
video_content = openai_client._download_video_content(video_id)
|
||||
print(f"online_video_e2e_latency_s={time.perf_counter() - start_time:.3f}")
|
||||
return video_content
|
||||
|
||||
|
||||
def _generate_online_video(
|
||||
*,
|
||||
omni_server,
|
||||
@@ -404,7 +557,10 @@ def test_wan22_i2v_diffusers_offline_generates_video(
|
||||
assert offline_path.exists(), f"Expected offline video artifact at {offline_path}"
|
||||
assert offline_metadata_path.exists(), f"Expected offline metadata artifact at {offline_metadata_path}"
|
||||
offline_metadata = _probe_video(offline_path)
|
||||
assert_video_metadata(offline_metadata, width=WIDTH, height=HEIGHT, fps=FPS, frame_count=NUM_FRAMES)
|
||||
assert offline_metadata["width"] == WIDTH
|
||||
assert offline_metadata["height"] == HEIGHT
|
||||
assert offline_metadata["fps"] == float(FPS)
|
||||
assert offline_metadata["frame_count"] == NUM_FRAMES
|
||||
|
||||
|
||||
@pytest.mark.benchmark
|
||||
@@ -430,7 +586,10 @@ def test_wan22_i2v_online_serving_generates_video(
|
||||
)
|
||||
assert online_path.exists(), f"Expected online video artifact at {online_path}"
|
||||
online_metadata = _probe_video(online_path)
|
||||
assert_video_metadata(online_metadata, width=WIDTH, height=HEIGHT, fps=FPS, frame_count=NUM_FRAMES)
|
||||
assert online_metadata["width"] == WIDTH
|
||||
assert online_metadata["height"] == HEIGHT
|
||||
assert online_metadata["fps"] == float(FPS)
|
||||
assert online_metadata["frame_count"] == NUM_FRAMES
|
||||
|
||||
|
||||
@pytest.mark.benchmark
|
||||
@@ -468,12 +627,42 @@ def test_wan22_i2v_serving_matches_diffusers_video_similarity(
|
||||
f"online_path={online_path}\n"
|
||||
f"offline_path={offline_path}"
|
||||
)
|
||||
assert_video_metadata(online_metadata, width=WIDTH, height=HEIGHT, fps=FPS, frame_count=NUM_FRAMES)
|
||||
assert_video_similarity_metrics(
|
||||
label="wan22_i2v",
|
||||
online_path=online_path,
|
||||
offline_path=offline_path,
|
||||
ssim_threshold=SSIM_THRESHOLD,
|
||||
psnr_threshold=PSNR_THRESHOLD,
|
||||
assert online_metadata["width"] == WIDTH
|
||||
assert online_metadata["height"] == HEIGHT
|
||||
assert online_metadata["fps"] == float(FPS)
|
||||
assert online_metadata["frame_count"] == NUM_FRAMES
|
||||
|
||||
ssim_output = _run_ffmpeg_similarity("ssim", online_path, offline_path)
|
||||
psnr_output = _run_ffmpeg_similarity("psnr", online_path, offline_path)
|
||||
ssim_score = _parse_ssim_score(ssim_output)
|
||||
psnr_score = _parse_psnr_score(psnr_output)
|
||||
|
||||
print("wan22_i2v_similarity metrics:")
|
||||
print(
|
||||
" SSIM:"
|
||||
f" value={ssim_score:.6f},"
|
||||
f" threshold>={SSIM_THRESHOLD:.6f},"
|
||||
" range=[0, 1],"
|
||||
" higher_is_better=True,"
|
||||
" interpretation=structural_similarity"
|
||||
)
|
||||
print(
|
||||
" PSNR:"
|
||||
f" value={psnr_score:.6f} dB,"
|
||||
f" threshold>={PSNR_THRESHOLD:.6f} dB,"
|
||||
" range=[0, +inf),"
|
||||
" higher_is_better=True,"
|
||||
" interpretation=pixel_error_in_decibels"
|
||||
)
|
||||
print(f"online_video={online_path}")
|
||||
print(f"offline_video={offline_path}")
|
||||
print(f"offline_metadata={offline_metadata_path}")
|
||||
|
||||
assert ssim_score >= SSIM_THRESHOLD, (
|
||||
f"SSIM below threshold: got {ssim_score:.6f}, expected >= {SSIM_THRESHOLD:.6f}. "
|
||||
f"online={online_path} offline={offline_path}"
|
||||
)
|
||||
assert psnr_score >= PSNR_THRESHOLD, (
|
||||
f"PSNR below threshold: got {psnr_score:.6f}, expected >= {PSNR_THRESHOLD:.6f}. "
|
||||
f"online={online_path} offline={offline_path}"
|
||||
)
|
||||
@@ -1,238 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""E2E tests for GLM-Image AutoRound W4A16 quantized inference.
|
||||
|
||||
These tests cover text-to-image and image-to-image generation with
|
||||
the W4A16 quantized GLM-Image model.
|
||||
|
||||
Requirements:
|
||||
- 2 CUDA GPUs (H100 or equivalent)
|
||||
- The quantized model checkpoint (Intel/GLM-Image-int4-AutoRound)
|
||||
"""
|
||||
|
||||
import gc
|
||||
import math
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
from PIL import Image
|
||||
from vllm import SamplingParams
|
||||
|
||||
from tests.helpers.env import DeviceMemoryMonitor
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.media import generate_synthetic_image
|
||||
from tests.helpers.runtime import OmniRunnerHandler
|
||||
from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
QUANTIZED_MODEL = os.environ.get("GLM_IMAGE_AUTOROUND_MODEL", "Intel/GLM-Image-int4-AutoRound")
|
||||
|
||||
# Small resolution to keep GPU memory & time manageable
|
||||
HEIGHT = 256
|
||||
WIDTH = 256
|
||||
NUM_STEPS = 2 # minimal for smoke-test
|
||||
|
||||
# GLM-Image AR generation config (from generation_config.json)
|
||||
GLM_IMAGE_EOS_TOKEN_ID = 16385
|
||||
GLM_IMAGE_VISION_VOCAB_SIZE = 16512
|
||||
|
||||
_CI_DEPLOY = get_deploy_config_path("glm_image.yaml")
|
||||
|
||||
|
||||
def _get_stage_config():
|
||||
"""Build a CI-friendly stage config with eager mode for testing."""
|
||||
return modify_stage_config(
|
||||
_CI_DEPLOY,
|
||||
updates={
|
||||
"stages": {
|
||||
0: {"enforce_eager": True},
|
||||
1: {"enforce_eager": True},
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
stage_config = _get_stage_config()
|
||||
|
||||
# (model, stage_config_path) for ``omni_runner`` indirect parametrize
|
||||
_OMNI_RUNNER_PARAM = (QUANTIZED_MODEL, stage_config)
|
||||
|
||||
|
||||
def compute_max_tokens(height: int, width: int, factor: int = 32) -> int:
|
||||
"""Compute max_new_tokens for GLM-Image AR text-to-image generation."""
|
||||
token_h = height // factor
|
||||
token_w = width // factor
|
||||
large_tokens = token_h * token_w
|
||||
|
||||
ratio = token_h / token_w if token_w > 0 else 1.0
|
||||
small_token_h = max(1, int(math.sqrt(ratio) * (factor // 2)))
|
||||
small_token_w = max(1, int(math.sqrt(1 / ratio) * (factor // 2)))
|
||||
small_tokens = small_token_h * small_token_w
|
||||
|
||||
return small_tokens + large_tokens + 1
|
||||
|
||||
|
||||
def _ar_sampling_params(max_tokens: int, height: int, width: int, seed: int = 42) -> SamplingParams:
|
||||
"""Build AR stage SamplingParams for GLM-Image."""
|
||||
return SamplingParams(
|
||||
temperature=0.9,
|
||||
top_p=0.75,
|
||||
top_k=GLM_IMAGE_VISION_VOCAB_SIZE,
|
||||
max_tokens=max_tokens,
|
||||
stop_token_ids=[GLM_IMAGE_EOS_TOKEN_ID],
|
||||
seed=seed,
|
||||
detokenize=False,
|
||||
extra_args={
|
||||
"target_h": height,
|
||||
"target_w": width,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _diffusion_sampling_params(
|
||||
height: int = HEIGHT,
|
||||
width: int = WIDTH,
|
||||
num_steps: int = NUM_STEPS,
|
||||
seed: int = 42,
|
||||
) -> OmniDiffusionSamplingParams:
|
||||
"""Build Diffusion stage OmniDiffusionSamplingParams."""
|
||||
return OmniDiffusionSamplingParams(
|
||||
height=height,
|
||||
width=width,
|
||||
num_inference_steps=num_steps,
|
||||
guidance_scale=0.0,
|
||||
seed=seed,
|
||||
)
|
||||
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.full_model,
|
||||
pytest.mark.diffusion,
|
||||
]
|
||||
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Test: text-to-image generation produces a valid image (quantized)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True)
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=2)
|
||||
def test_glm_image_autoround_w4a16_generates_image(omni_runner_handler: OmniRunnerHandler):
|
||||
"""Load the W4A16 quantized GLM-Image model and verify it produces a valid image."""
|
||||
gc.collect()
|
||||
current_omni_platform.empty_cache()
|
||||
device_index = current_omni_platform.current_device()
|
||||
current_omni_platform.reset_peak_memory_stats()
|
||||
monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02)
|
||||
monitor.start()
|
||||
|
||||
prompt_dict = {
|
||||
"prompt": "A photo of a cat sitting on a laptop keyboard",
|
||||
"height": HEIGHT,
|
||||
"width": WIDTH,
|
||||
"mm_processor_kwargs": {
|
||||
"target_h": HEIGHT,
|
||||
"target_w": WIDTH,
|
||||
},
|
||||
}
|
||||
ar_params = _ar_sampling_params(
|
||||
max_tokens=compute_max_tokens(HEIGHT, WIDTH),
|
||||
height=HEIGHT,
|
||||
width=WIDTH,
|
||||
seed=42,
|
||||
)
|
||||
diffusion_params = _diffusion_sampling_params(
|
||||
height=HEIGHT,
|
||||
width=WIDTH,
|
||||
num_steps=NUM_STEPS,
|
||||
seed=42,
|
||||
)
|
||||
|
||||
outputs = omni_runner_handler.runner.generate(
|
||||
[prompt_dict],
|
||||
[ar_params, diffusion_params],
|
||||
)
|
||||
|
||||
monitor.stop()
|
||||
|
||||
first_output = outputs[0]
|
||||
assert first_output.final_output_type == "image"
|
||||
req_out = first_output.request_output
|
||||
assert isinstance(req_out, OmniRequestOutput) and hasattr(req_out, "images")
|
||||
images = req_out.images
|
||||
|
||||
assert len(images) >= 1, "Expected at least one generated image"
|
||||
img = images[0]
|
||||
assert isinstance(img, Image.Image)
|
||||
assert img.width == WIDTH, f"Expected width {WIDTH}, got {img.width}"
|
||||
assert img.height == HEIGHT, f"Expected height {HEIGHT}, got {img.height}"
|
||||
|
||||
# Sanity: image should not be blank (all-zero)
|
||||
arr = np.array(img)
|
||||
assert arr.std() > 1.0, "Generated image appears blank (std ≈ 0)"
|
||||
|
||||
gc.collect()
|
||||
current_omni_platform.empty_cache()
|
||||
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Test: image-to-image generation (quantized)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True)
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=2)
|
||||
def test_glm_image_autoround_w4a16_image_to_image(omni_runner_handler: OmniRunnerHandler):
|
||||
"""Load the W4A16 quantized GLM-Image and verify image-to-image generation works."""
|
||||
ref_image_arr = generate_synthetic_image(WIDTH, HEIGHT)["np_array"]
|
||||
|
||||
gc.collect()
|
||||
current_omni_platform.empty_cache()
|
||||
current_omni_platform.reset_peak_memory_stats()
|
||||
|
||||
prompt_dict = {
|
||||
"prompt": "Make it look like winter",
|
||||
"multi_modal_data": {"image": ref_image_arr},
|
||||
"height": HEIGHT,
|
||||
"width": WIDTH,
|
||||
"mm_processor_kwargs": {
|
||||
"target_h": HEIGHT,
|
||||
"target_w": WIDTH,
|
||||
},
|
||||
}
|
||||
ar_params = _ar_sampling_params(
|
||||
max_tokens=compute_max_tokens(HEIGHT, WIDTH),
|
||||
height=HEIGHT,
|
||||
width=WIDTH,
|
||||
seed=42,
|
||||
)
|
||||
diffusion_params = _diffusion_sampling_params(
|
||||
height=HEIGHT,
|
||||
width=WIDTH,
|
||||
num_steps=NUM_STEPS,
|
||||
seed=42,
|
||||
)
|
||||
|
||||
outputs = omni_runner_handler.runner.generate(
|
||||
[prompt_dict],
|
||||
[ar_params, diffusion_params],
|
||||
)
|
||||
|
||||
first_output = outputs[0]
|
||||
assert first_output.final_output_type == "image"
|
||||
req_out = first_output.request_output
|
||||
assert isinstance(req_out, OmniRequestOutput) and hasattr(req_out, "images")
|
||||
images = req_out.images
|
||||
|
||||
assert len(images) >= 1, "Expected at least one generated image"
|
||||
img = images[0]
|
||||
assert isinstance(img, Image.Image)
|
||||
assert img.width == WIDTH
|
||||
assert img.height == HEIGHT
|
||||
|
||||
gc.collect()
|
||||
current_omni_platform.empty_cache()
|
||||
@@ -1,231 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
End-to-end offline inference tests for higgs-audio v2.
|
||||
|
||||
Mirrors examples/offline_inference/text_to_speech/higgs_audio_v2/end2end.py:
|
||||
build Stage-0 prompt_token_ids with the model's own processor + plain-text
|
||||
chat template, drive the engine via ``Omni.generate``, and inspect the
|
||||
Stage-1 PCM payload directly. The OmniRunnerHandler.send_audio_speech_request
|
||||
helper applies a Qwen-style ChatML wrap that is not valid for higgs, so this
|
||||
test bypasses it (same approach as test_voxtral_tts.py).
|
||||
|
||||
v1 scope is plain text -> 24 kHz speech; voice cloning, multi-speaker prompts,
|
||||
language overrides, and task_type selection are out of scope.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
# Match the example script: keep DeepGEMM FP8 off so engine init does not
|
||||
# require the optional ``deep_gemm`` backend.
|
||||
os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
|
||||
os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0")
|
||||
os.environ.setdefault("VLLM_MOE_USE_DEEP_GEMM", "0")
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from vllm import SamplingParams
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.media import get_asset_path
|
||||
from tests.helpers.runtime import OmniRunner
|
||||
from tests.helpers.stage_config import get_deploy_config_path
|
||||
|
||||
MODEL = "bosonai/higgs-audio-v2-generation-3B-base"
|
||||
STAGE_CONFIG = get_deploy_config_path("higgs_audio_v2.yaml")
|
||||
SAMPLE_RATE = 24_000
|
||||
# A short sentence at 24 kHz must produce at least ~0.5 s of audio.
|
||||
MIN_AUDIO_SAMPLES = 12_000
|
||||
TEST_TEXT = "Hello world."
|
||||
|
||||
# Reuse the qwen3_tts ref clip + transcript for voice clone tests (see
|
||||
# tests/e2e/online_serving/test_qwen3_tts_base.py:27-28 for the source).
|
||||
_REF_AUDIO_ASSET = "qwen3_tts/clone_2.wav"
|
||||
_REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
|
||||
|
||||
# (model, stage_config_path, extra_omni_kwargs) — same shape as test_voxtral_tts.py.
|
||||
# Stage configs already pin enforce_eager=true per stage, so no Omni-level override.
|
||||
_OMNI_RUNNER_PARAM = (MODEL, STAGE_CONFIG, {"trust_remote_code": True})
|
||||
|
||||
|
||||
def _compose_request(model_name: str, text: str) -> dict:
|
||||
"""Build a higgs prompt dict via the upstream processor + plain-text template.
|
||||
|
||||
The serving layer's ``_build_higgs_audio_v2_params`` does the same thing
|
||||
(see vllm_omni/entrypoints/openai/serving_speech.py:1520) — keeping the
|
||||
offline and online prompt-construction paths identical is what guarantees
|
||||
input-token parity with the HF reference.
|
||||
"""
|
||||
from transformers import AutoProcessor
|
||||
|
||||
from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_v2_tokenizer import (
|
||||
build_plain_text_prompt,
|
||||
input_ids_to_python_list,
|
||||
)
|
||||
|
||||
processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True)
|
||||
inputs = build_plain_text_prompt(processor, text)
|
||||
return {"prompt_token_ids": input_ids_to_python_list(inputs)}
|
||||
|
||||
|
||||
def _extract_pcm(outputs) -> torch.Tensor:
|
||||
"""Pull the concatenated PCM tensor out of the Stage-1 multimodal_output."""
|
||||
for stage_output in outputs:
|
||||
mm = getattr(stage_output, "multimodal_output", None)
|
||||
if not mm:
|
||||
continue
|
||||
audio = mm.get("model_outputs") if "model_outputs" in mm else mm.get("audio")
|
||||
if audio is None:
|
||||
continue
|
||||
if isinstance(audio, list):
|
||||
chunks = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio if a is not None]
|
||||
if not chunks:
|
||||
continue
|
||||
return torch.cat(chunks, dim=0) if len(chunks) > 1 else chunks[0]
|
||||
return torch.as_tensor(audio).float().cpu().reshape(-1)
|
||||
raise AssertionError(f"no audio payload in any stage output (got {len(outputs)} stages)")
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
@pytest.mark.parametrize("omni_runner_function", [_OMNI_RUNNER_PARAM], indirect=True)
|
||||
def test_higgs_audio_v2_offline_plain_text(omni_runner_function: OmniRunner) -> None:
|
||||
"""
|
||||
Offline plain-text -> 24 kHz PCM via the bundled higgs_audio_v2 deploy.
|
||||
|
||||
Deploy Setting: vllm_omni/deploy/higgs_audio_v2.yaml (codec_streaming=true,
|
||||
async_chunk=false, enforce_eager=true per stage).
|
||||
Input Modal: text
|
||||
Output Modal: audio (PCM tensor, 24 kHz mono float)
|
||||
Input Setting: single prompt, default sampling
|
||||
"""
|
||||
omni = omni_runner_function.omni
|
||||
inputs = _compose_request(MODEL, TEST_TEXT)
|
||||
|
||||
# max_tokens caps Stage-0 codec frames; Stage 1 max_tokens lifts the
|
||||
# per-codebook prefill cap. Mirror the deploy yaml's defaults.
|
||||
stage_sampling = [
|
||||
SamplingParams(temperature=1.0, top_p=0.95, top_k=50, max_tokens=1024, seed=42),
|
||||
SamplingParams(temperature=0.0, top_p=1.0, top_k=-1, max_tokens=65536, seed=42),
|
||||
]
|
||||
|
||||
outputs = list(omni.generate([inputs], stage_sampling))
|
||||
assert outputs, "no outputs returned from Omni.generate"
|
||||
|
||||
pcm = _extract_pcm(outputs)
|
||||
arr = pcm.numpy()
|
||||
|
||||
assert arr.size > MIN_AUDIO_SAMPLES, (
|
||||
f"audio too short: {arr.size} samples, expected > {MIN_AUDIO_SAMPLES} "
|
||||
f"({MIN_AUDIO_SAMPLES / SAMPLE_RATE:.2f} s @ {SAMPLE_RATE} Hz)"
|
||||
)
|
||||
# Threshold 5e-4 (not 1e-3) keeps us above true silence while tolerating
|
||||
# the ``load_format=dummy`` Stage-1 codec — random RVQ/DAC weights collapse
|
||||
# the output to a near-constant peak around 7.4e-4. Real-weight runs reach
|
||||
# 0.1+ peak so this still catches genuine silence.
|
||||
assert np.max(np.abs(arr)) > 5e-4, "audio output is silent (peak abs amplitude < 5e-4)"
|
||||
assert np.isfinite(arr).all(), "audio contains non-finite samples (nan/inf)"
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
@pytest.mark.parametrize("omni_runner_function", [_OMNI_RUNNER_PARAM], indirect=True)
|
||||
def test_higgs_audio_v2_offline_batch_two_prompts(omni_runner_function: OmniRunner) -> None:
|
||||
"""
|
||||
Submit two prompts in one engine call and verify both yield non-trivial PCM.
|
||||
|
||||
Guards the Stage-0 talker's per-slot audio state under a batched call —
|
||||
the prior example ran prompts one at a time citing 'request-scoped audio
|
||||
state'; this test pins the batched-2 behavior the engine path actually
|
||||
needs to support for /v1/audio/speech under concurrent traffic.
|
||||
"""
|
||||
omni = omni_runner_function.omni
|
||||
prompts = [
|
||||
_compose_request(MODEL, "Hello world."),
|
||||
_compose_request(MODEL, "The quick brown fox jumps over the lazy dog."),
|
||||
]
|
||||
stage_sampling = [
|
||||
SamplingParams(temperature=1.0, top_p=0.95, top_k=50, max_tokens=1024, seed=42),
|
||||
SamplingParams(temperature=0.0, top_p=1.0, top_k=-1, max_tokens=65536, seed=42),
|
||||
]
|
||||
|
||||
outputs = list(omni.generate(prompts, stage_sampling))
|
||||
assert outputs, "no outputs returned from Omni.generate"
|
||||
|
||||
pcm = _extract_pcm(outputs)
|
||||
# Whether the engine returns one concatenated PCM or per-request tensors is
|
||||
# implementation-defined; for either shape, total samples must clear the
|
||||
# floor for at least one of the two prompts.
|
||||
assert pcm.numel() > MIN_AUDIO_SAMPLES, (
|
||||
f"batched audio too short: {pcm.numel()} samples, expected > {MIN_AUDIO_SAMPLES}"
|
||||
)
|
||||
# See the silence-threshold note in test_higgs_audio_v2_offline_plain_text.
|
||||
assert np.max(np.abs(pcm.numpy())) > 5e-4, "batched audio output is silent"
|
||||
|
||||
|
||||
def _compose_voice_clone_request(model_name: str, text: str, ref_text: str, ref_wav_path) -> dict:
|
||||
"""Build the voice-clone prompt + additional_information payload.
|
||||
|
||||
Calls the same ``build_voice_clone_prompt`` the serving layer uses online —
|
||||
the HF processor encodes the reference clip on the spot via the bundled
|
||||
HiggsAudioV2TokenizerModel. The talker's
|
||||
``_maybe_apply_ref_audio_substitution`` consumes the tensors at prefill.
|
||||
"""
|
||||
import soundfile as sf
|
||||
from transformers import AutoProcessor
|
||||
|
||||
from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_v2_tokenizer import (
|
||||
build_voice_clone_prompt,
|
||||
)
|
||||
|
||||
wav, sr = sf.read(str(ref_wav_path), always_2d=False)
|
||||
if wav.ndim == 2:
|
||||
wav = wav.mean(axis=1) # downmix any stereo asset to mono
|
||||
processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=True)
|
||||
out = build_voice_clone_prompt(processor, text, wav, int(sr), ref_text)
|
||||
return {
|
||||
"prompt_token_ids": out["prompt_token_ids"],
|
||||
# Pass tensors bare (see serving_speech._build_higgs_audio_v2_params
|
||||
# — list-wrapped tensors get dropped by the msgspec serializer).
|
||||
"additional_information": {
|
||||
"audio_input_ids": out["audio_input_ids"],
|
||||
"audio_input_ids_mask": out["audio_input_ids_mask"],
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
@pytest.mark.parametrize("omni_runner_function", [_OMNI_RUNNER_PARAM], indirect=True)
|
||||
def test_higgs_audio_v2_offline_voice_clone(omni_runner_function: OmniRunner) -> None:
|
||||
"""
|
||||
Shallow voice clone via ref_audio + ref_text, offline.
|
||||
|
||||
Deploy Setting: vllm_omni/deploy/higgs_audio_v2.yaml.
|
||||
Input Modal: text + reference audio (qwen3_tts/clone_2.wav vendored asset)
|
||||
Output Modal: audio (PCM tensor, 24 kHz mono float)
|
||||
Input Setting: single prompt
|
||||
"""
|
||||
omni = omni_runner_function.omni
|
||||
ref_path = get_asset_path(_REF_AUDIO_ASSET)
|
||||
inputs = _compose_voice_clone_request(MODEL, TEST_TEXT, _REF_TEXT, ref_path)
|
||||
|
||||
stage_sampling = [
|
||||
SamplingParams(temperature=1.0, top_p=0.95, top_k=50, max_tokens=1024, seed=42),
|
||||
SamplingParams(temperature=0.0, top_p=1.0, top_k=-1, max_tokens=65536, seed=42),
|
||||
]
|
||||
|
||||
outputs = list(omni.generate([inputs], stage_sampling))
|
||||
assert outputs, "no outputs returned from Omni.generate for voice clone"
|
||||
|
||||
pcm = _extract_pcm(outputs)
|
||||
arr = pcm.numpy()
|
||||
assert arr.size > MIN_AUDIO_SAMPLES, f"cloned audio too short: {arr.size} samples, expected > {MIN_AUDIO_SAMPLES}"
|
||||
assert np.max(np.abs(arr)) > 5e-4, "voice-clone output is silent"
|
||||
assert np.isfinite(arr).all(), "voice-clone audio contains non-finite samples"
|
||||
@@ -3,8 +3,6 @@ End-to-end diffusion coverage for FLUX.2-dev in online serving mode.
|
||||
|
||||
Coverage:
|
||||
- CPU offload
|
||||
- Ulysses sequence parallelism
|
||||
- Ring sequence parallelism
|
||||
|
||||
This test verifies that FLUX.2-dev can be launched with CPU offload enabled,
|
||||
accepts text-to-image requests through the OpenAI-compatible API, and returns
|
||||
@@ -30,7 +28,7 @@ PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
|
||||
|
||||
|
||||
def _get_flux_2_dev_feature_cases(model: str):
|
||||
"""Return FLUX.2-dev diffusion feature cases for CPU offload and SP."""
|
||||
"""Return FLUX.2-dev diffusion feature cases for CPU offload."""
|
||||
|
||||
return [
|
||||
pytest.param(
|
||||
@@ -55,30 +53,6 @@ def _get_flux_2_dev_feature_cases(model: str):
|
||||
id="parallel_cfg_2",
|
||||
marks=PARALLEL_FEATURE_MARKS,
|
||||
),
|
||||
pytest.param(
|
||||
OmniServerParams(
|
||||
model=model,
|
||||
server_args=[
|
||||
"--enable-cpu-offload",
|
||||
"--ulysses-degree",
|
||||
"2",
|
||||
],
|
||||
),
|
||||
id="ulysses_2",
|
||||
marks=PARALLEL_FEATURE_MARKS,
|
||||
),
|
||||
pytest.param(
|
||||
OmniServerParams(
|
||||
model=model,
|
||||
server_args=[
|
||||
"--enable-cpu-offload",
|
||||
"--ring-degree",
|
||||
"2",
|
||||
],
|
||||
),
|
||||
id="ring_2",
|
||||
marks=PARALLEL_FEATURE_MARKS,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
|
||||
@@ -1,299 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
End-to-end online tests for higgs-audio v2 against /v1/audio/speech.
|
||||
|
||||
v1 scope is plain text -> 24 kHz speech plus shallow voice clone via
|
||||
ref_audio + ref_text (inline) or voice=<name> (after POST /v1/audio/voices).
|
||||
The model-aware request validator
|
||||
(vllm_omni/entrypoints/openai/serving_speech.py::_validate_higgs_audio_v2_request)
|
||||
rejects multi-speaker tags, language overrides, task_type, and bare
|
||||
voice=<name> for names that do not match an uploaded speaker — this suite
|
||||
exercises both the happy path (plain text in, audio bytes out) and the
|
||||
validator rejections, so a regression that loosens the schema will fail
|
||||
loudly.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
|
||||
# Match run_server.sh: DeepGEMM FP8 kernels are optional and trip warmup on
|
||||
# images without the deep_gemm backend, so disable them by default.
|
||||
os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0")
|
||||
os.environ.setdefault("VLLM_MOE_USE_DEEP_GEMM", "0")
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.media import load_test_audio_data_url
|
||||
from tests.helpers.runtime import OmniServerParams
|
||||
from tests.helpers.stage_config import get_deploy_config_path
|
||||
|
||||
MODEL = "bosonai/higgs-audio-v2-generation-3B-base"
|
||||
STAGE_CONFIG = get_deploy_config_path("higgs_audio_v2.yaml")
|
||||
SERVER_ARGS = ["--trust-remote-code", "--disable-log-stats"]
|
||||
# DeepGEMM warmup is optional; mirror run_server.sh and switch it off in the
|
||||
# server subprocess env too (parent-process os.environ above only affects this
|
||||
# test driver; the engine subprocesses inherit through env_dict).
|
||||
SERVER_ENV = {"VLLM_USE_DEEP_GEMM": "0", "VLLM_MOE_USE_DEEP_GEMM": "0"}
|
||||
|
||||
TEST_PARAMS = [
|
||||
pytest.param(
|
||||
OmniServerParams(
|
||||
model=MODEL,
|
||||
stage_config_path=STAGE_CONFIG,
|
||||
server_args=SERVER_ARGS,
|
||||
env_dict=SERVER_ENV,
|
||||
),
|
||||
id="higgs_audio_v2_plain_text",
|
||||
)
|
||||
]
|
||||
|
||||
DEFAULT_SPEECH_TIMEOUT_S = 180.0
|
||||
# Floor for ~0.5 s of 24 kHz mono PCM_16: 24000 * 0.5 * 2 bytes ≈ 24 KiB.
|
||||
# A WAV header adds 44 bytes; pick a conservative floor that catches truncated
|
||||
# / silence-only outputs without flagging short legitimate clips.
|
||||
_MIN_AUDIO_BYTES = 20_000
|
||||
|
||||
# Reuse the qwen3_tts vendored reference clip (clean ~5 s 24 kHz mono human
|
||||
# speech) + its transcript. See tests/e2e/online_serving/test_qwen3_tts_base.py
|
||||
# for the asset rationale — keeping a single shared reference clip across TTS
|
||||
# tests avoids duplicating WAVs in the repo.
|
||||
_REF_AUDIO_URL = load_test_audio_data_url("qwen3_tts/clone_2.wav")
|
||||
_REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
|
||||
|
||||
|
||||
@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
|
||||
class TestHiggsAudioV2OnlineHappyPath:
|
||||
"""Plain-text -> audio happy paths against the live HTTP server."""
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
def test_plain_text_wav(self, omni_server, openai_client) -> None:
|
||||
"""Single non-streaming WAV request — covers the canonical TTS happy path."""
|
||||
openai_client.send_audio_speech_request(
|
||||
{
|
||||
"model": omni_server.model,
|
||||
"input": "Hello world.",
|
||||
"stream": False,
|
||||
"response_format": "wav",
|
||||
"timeout": DEFAULT_SPEECH_TIMEOUT_S,
|
||||
"min_audio_bytes": _MIN_AUDIO_BYTES,
|
||||
}
|
||||
)
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
def test_plain_text_pcm_streaming(self, omni_server, openai_client) -> None:
|
||||
"""Streaming PCM via the shared-memory connector's codec_streaming path.
|
||||
|
||||
higgs_audio_v2.yaml pins ``codec_streaming: true`` + ``async_chunk: false``,
|
||||
so the only streaming surface exposed to clients is the WAV/PCM bytes
|
||||
served chunk-by-chunk from Stage 1 — exercise it directly.
|
||||
"""
|
||||
openai_client.send_audio_speech_request(
|
||||
{
|
||||
"model": omni_server.model,
|
||||
"input": "Streaming the quick brown fox over the lazy dog.",
|
||||
"stream": True,
|
||||
"response_format": "pcm",
|
||||
"timeout": DEFAULT_SPEECH_TIMEOUT_S,
|
||||
"min_audio_bytes": _MIN_AUDIO_BYTES,
|
||||
}
|
||||
)
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
def test_plain_text_with_max_new_tokens(self, omni_server, openai_client) -> None:
|
||||
"""max_new_tokens is one of the few extra fields the higgs validator accepts."""
|
||||
openai_client.send_audio_speech_request(
|
||||
{
|
||||
"model": omni_server.model,
|
||||
"input": "Innovation distinguishes between a leader and a follower.",
|
||||
"stream": False,
|
||||
"response_format": "wav",
|
||||
"timeout": DEFAULT_SPEECH_TIMEOUT_S,
|
||||
"max_new_tokens": 500,
|
||||
"min_audio_bytes": _MIN_AUDIO_BYTES,
|
||||
}
|
||||
)
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
def test_concurrent_plain_text(self, omni_server, openai_client) -> None:
|
||||
"""Three concurrent non-streaming requests — guards the per-slot audio state."""
|
||||
openai_client.send_audio_speech_request(
|
||||
{
|
||||
"model": omni_server.model,
|
||||
"input": "It was the night before my birthday.",
|
||||
"stream": False,
|
||||
"response_format": "wav",
|
||||
"timeout": DEFAULT_SPEECH_TIMEOUT_S,
|
||||
"min_audio_bytes": _MIN_AUDIO_BYTES,
|
||||
},
|
||||
request_num=3,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Validator rejections served over HTTP. Each case targets one out-of-scope
|
||||
# field; the validator returns 4xx via the OpenAI error path. We keep the
|
||||
# error-message substring loose so a phrasing tweak in the validator does not
|
||||
# break CI, but tight enough to catch a regression that silently accepts the
|
||||
# field.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
|
||||
class TestHiggsAudioV2OnlineValidatorRejections:
|
||||
"""Out-of-scope fields must come back as 4xx with a higgs-named message."""
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
def test_rejects_ref_audio_without_ref_text(self, omni_server, openai_client) -> None:
|
||||
"""Voice clone needs the transcript too — half-supplied is a 4xx."""
|
||||
openai_client.send_audio_speech_request(
|
||||
{
|
||||
"model": omni_server.model,
|
||||
"input": "Hello world.",
|
||||
"ref_audio": _REF_AUDIO_URL,
|
||||
"response_format": "wav",
|
||||
"timeout": DEFAULT_SPEECH_TIMEOUT_S,
|
||||
"status_code": (400, 422),
|
||||
"err_message": "ref_text",
|
||||
}
|
||||
)
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
def test_rejects_ref_text_without_ref_audio(self, omni_server, openai_client) -> None:
|
||||
"""Symmetric guard: ref_text alone is not enough — must come with ref_audio."""
|
||||
openai_client.send_audio_speech_request(
|
||||
{
|
||||
"model": omni_server.model,
|
||||
"input": "Hello world.",
|
||||
"ref_text": "some transcript",
|
||||
"response_format": "wav",
|
||||
"timeout": DEFAULT_SPEECH_TIMEOUT_S,
|
||||
"status_code": (400, 422),
|
||||
"err_message": "ref_audio",
|
||||
}
|
||||
)
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
def test_rejects_task_type(self, omni_server, openai_client) -> None:
|
||||
openai_client.send_audio_speech_request(
|
||||
{
|
||||
"model": omni_server.model,
|
||||
"input": "Hello world.",
|
||||
"task_type": "Base",
|
||||
"response_format": "wav",
|
||||
"timeout": DEFAULT_SPEECH_TIMEOUT_S,
|
||||
"status_code": (400, 422),
|
||||
"err_message": "task_type",
|
||||
}
|
||||
)
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
def test_rejects_language_override(self, omni_server, openai_client) -> None:
|
||||
openai_client.send_audio_speech_request(
|
||||
{
|
||||
"model": omni_server.model,
|
||||
"input": "Hello world.",
|
||||
"language": "Chinese",
|
||||
"response_format": "wav",
|
||||
"timeout": DEFAULT_SPEECH_TIMEOUT_S,
|
||||
"status_code": (400, 422),
|
||||
"err_message": "language",
|
||||
}
|
||||
)
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
def test_rejects_multi_speaker_tag_in_text(self, omni_server, openai_client) -> None:
|
||||
openai_client.send_audio_speech_request(
|
||||
{
|
||||
"model": omni_server.model,
|
||||
"input": "[SPEAKER0] hi",
|
||||
"response_format": "wav",
|
||||
"timeout": DEFAULT_SPEECH_TIMEOUT_S,
|
||||
"status_code": (400, 422),
|
||||
"err_message": "multi-speaker",
|
||||
}
|
||||
)
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
def test_rejects_empty_input(self, omni_server, openai_client) -> None:
|
||||
openai_client.send_audio_speech_request(
|
||||
{
|
||||
"model": omni_server.model,
|
||||
"input": " ",
|
||||
"response_format": "wav",
|
||||
"timeout": DEFAULT_SPEECH_TIMEOUT_S,
|
||||
"status_code": (400, 422),
|
||||
"err_message": "empty",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
|
||||
class TestHiggsAudioV2OnlineVoiceClone:
|
||||
"""Shallow voice clone: ref_audio + ref_text -> speech in the cloned voice.
|
||||
|
||||
The HF processor that the serving layer calls (lazy-loaded at the first
|
||||
request) ships with the bundled HiggsAudioV2TokenizerModel and encodes
|
||||
the reference clip in-process. The talker substitutes the encoded codes
|
||||
at the prompt-side audio placeholders via
|
||||
:meth:`HiggsAudioV2TalkerForConditionalGeneration._maybe_apply_ref_audio_substitution`.
|
||||
"""
|
||||
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
def test_voice_clone_basic(self, omni_server, openai_client) -> None:
|
||||
"""Single non-streaming WAV request driven by the qwen3_tts ref clip."""
|
||||
openai_client.send_audio_speech_request(
|
||||
{
|
||||
"model": omni_server.model,
|
||||
"input": "Hello world.",
|
||||
"ref_audio": _REF_AUDIO_URL,
|
||||
"ref_text": _REF_TEXT,
|
||||
"stream": False,
|
||||
"response_format": "wav",
|
||||
"timeout": DEFAULT_SPEECH_TIMEOUT_S,
|
||||
"min_audio_bytes": _MIN_AUDIO_BYTES,
|
||||
}
|
||||
)
|
||||
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
def test_voice_clone_pcm_streaming(self, omni_server, openai_client) -> None:
|
||||
"""Voice clone over the streaming PCM path."""
|
||||
openai_client.send_audio_speech_request(
|
||||
{
|
||||
"model": omni_server.model,
|
||||
"input": "Innovation distinguishes a leader from a follower.",
|
||||
"ref_audio": _REF_AUDIO_URL,
|
||||
"ref_text": _REF_TEXT,
|
||||
"stream": True,
|
||||
"response_format": "pcm",
|
||||
"timeout": DEFAULT_SPEECH_TIMEOUT_S,
|
||||
"min_audio_bytes": _MIN_AUDIO_BYTES,
|
||||
}
|
||||
)
|
||||
@@ -539,24 +539,3 @@ def test_language_001(omni_server, openai_client) -> None:
|
||||
}
|
||||
|
||||
openai_client.send_omni_request(request_config)
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
||||
def test_text_to_audio_long_output_001(omni_server, openai_client) -> None:
|
||||
"""
|
||||
Input Modal: text only (long-form generation prompt).
|
||||
Output Modal: text, audio (default ``modalities``);
|
||||
Input Setting: stream=True
|
||||
Datasets: single request
|
||||
"""
|
||||
messages = dummy_messages_from_mix_data(
|
||||
system_prompt=get_system_prompt(),
|
||||
content_text="Tell a 300-word story.",
|
||||
)
|
||||
|
||||
request_config = {"model": omni_server.model, "messages": messages, "stream": True}
|
||||
responses = openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
|
||||
text = responses[0].text_content if responses else ""
|
||||
word_count = len(text.split())
|
||||
assert word_count >= 200, f"Expected at least 200 words in long output, got {word_count}"
|
||||
@@ -5,8 +5,8 @@
|
||||
Online serving tests for the Qwen-Image-Edit family (image-to-image via chat completions).
|
||||
|
||||
- ``test_single_image_to_image_001``: ``Qwen/Qwen-Image-Edit`` — one reference image, fixed 512×512.
|
||||
- ``test_multi_images_to_image_001``: ``Qwen/Qwen-Image-Edit-2511`` — two reference images, fixed 512×512.
|
||||
- ``test_different_sizes_001``: ``Qwen/Qwen-Image-Edit-2511`` only, ``advanced_model`` — mixed input
|
||||
- ``test_multi_images_to_image_001``: ``Qwen/Qwen-Image-Edit-2509`` — two reference images, fixed 512×512.
|
||||
- ``test_different_sizes_001``: ``Qwen/Qwen-Image-Edit-2509`` only, ``advanced_model`` — mixed input
|
||||
resolutions; ``extra_body`` uses per-output ``width``/``height`` lists; the test client sends one
|
||||
scalar-size request per list index in parallel and merges images (see ``OpenAIClientHandler.send_diffusion_request``).
|
||||
|
||||
@@ -53,6 +53,7 @@ def _get_diffusion_feature_cases(model: str):
|
||||
|
||||
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.diffusion
|
||||
@pytest.mark.parametrize(
|
||||
"omni_server",
|
||||
@@ -82,16 +83,15 @@ def test_single_image_to_image_001(omni_server: OmniServer, openai_client: OpenA
|
||||
openai_client.send_diffusion_request(request_config)
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.diffusion
|
||||
@pytest.mark.parametrize(
|
||||
"omni_server",
|
||||
_get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2511"),
|
||||
_get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2509"),
|
||||
indirect=True,
|
||||
)
|
||||
def test_multi_images_to_image_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):
|
||||
"""Two-reference edit smoke for ``Qwen/Qwen-Image-Edit-2511``."""
|
||||
"""Two-reference edit smoke for ``Qwen/Qwen-Image-Edit-2509``."""
|
||||
|
||||
image_data_url_list = [f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}" for _ in range(2)]
|
||||
|
||||
@@ -118,7 +118,7 @@ def test_multi_images_to_image_001(omni_server: OmniServer, openai_client: OpenA
|
||||
@pytest.mark.diffusion
|
||||
@pytest.mark.parametrize(
|
||||
"omni_server",
|
||||
_get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2511"),
|
||||
_get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2509"),
|
||||
indirect=True,
|
||||
)
|
||||
def test_different_sizes_001(omni_server: OmniServer, openai_client: OpenAIClientHandler):
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
Comprehensive tests of diffusion features that are available in online serving mode
|
||||
and are supported by the following models:
|
||||
- Qwen-Image-Edit: single image input
|
||||
- Qwen-Image-Edit-2511: single image input and two image inputs
|
||||
- Qwen-Image-Edit-2509: single image input and two image inputs
|
||||
"""
|
||||
|
||||
import pytest
|
||||
@@ -14,7 +14,7 @@ from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHand
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
|
||||
|
||||
EDIT_PROMPT = "Transform this modern, geometrist image into a Vincent van Gogh style impressionist painting."
|
||||
SINGLE_EDIT_PROMPT_2511 = "Restyle this image into a Vincent van Gogh style impressionist painting."
|
||||
SINGLE_EDIT_PROMPT_2509 = "Restyle this image into a Vincent van Gogh style impressionist painting."
|
||||
MULTI_EDIT_PROMPT = (
|
||||
"Transform the first image into a Dadaism collage art. "
|
||||
"Transform the second image into a Vincent van Gogh style painting. "
|
||||
@@ -34,7 +34,7 @@ def _get_diffusion_feature_cases(model: str):
|
||||
model=model,
|
||||
server_args=[
|
||||
"--cache-backend",
|
||||
"cache_dit",
|
||||
"tea_cache", # [TODO] may consider changing to cache_dit after #1779 is resolved. Currently cache_dit and layerwise offload cannot work together.
|
||||
"--enable-layerwise-offload",
|
||||
],
|
||||
),
|
||||
@@ -59,7 +59,7 @@ def _get_diffusion_feature_cases(model: str):
|
||||
model=model,
|
||||
server_args=[
|
||||
"--cache-backend",
|
||||
"tea_cache",
|
||||
"cache_dit",
|
||||
"--ring",
|
||||
"2",
|
||||
],
|
||||
@@ -141,11 +141,11 @@ def test_qwen_image_edit(omni_server: OmniServer, openai_client: OpenAIClientHan
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"omni_server",
|
||||
_get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2511"),
|
||||
_get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2509"),
|
||||
indirect=True,
|
||||
)
|
||||
def test_qwen_image_edit_2511_single_image(omni_server: OmniServer, openai_client: OpenAIClientHandler):
|
||||
"""Test Qwen-Image-Edit-2511 with a single image input.
|
||||
def test_qwen_image_edit_2509_single_image(omni_server: OmniServer, openai_client: OpenAIClientHandler):
|
||||
"""Test Qwen-Image-Edit-2509 with a single image input.
|
||||
|
||||
Regression: with tea_cache enabled and zero_cond_t=True, the TeaCache
|
||||
postprocess closure used the doubled temb (shape 2*batch) without halving
|
||||
@@ -156,7 +156,7 @@ def test_qwen_image_edit_2511_single_image(omni_server: OmniServer, openai_clien
|
||||
"""
|
||||
image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
|
||||
|
||||
messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=SINGLE_EDIT_PROMPT_2511)
|
||||
messages = dummy_messages_from_mix_data(image_data_url=image_data_url, content_text=SINGLE_EDIT_PROMPT_2509)
|
||||
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
@@ -176,11 +176,11 @@ def test_qwen_image_edit_2511_single_image(omni_server: OmniServer, openai_clien
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"omni_server",
|
||||
_get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2511"),
|
||||
_get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2509"),
|
||||
indirect=True,
|
||||
)
|
||||
def test_qwen_image_edit_2511_two_images(omni_server: OmniServer, openai_client: OpenAIClientHandler):
|
||||
"""Test Qwen-Image-Edit-2511 with two image inputs."""
|
||||
def test_qwen_image_edit_2509_two_images(omni_server: OmniServer, openai_client: OpenAIClientHandler):
|
||||
"""Test Qwen-Image-Edit-2509 with two image inputs."""
|
||||
image_data_url_1 = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
|
||||
image_data_url_2 = f"data:image/jpeg;base64,{generate_synthetic_image(512, 512)['base64']}"
|
||||
|
||||
|
||||
@@ -1,10 +1,8 @@
|
||||
import concurrent.futures
|
||||
import importlib
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
import types
|
||||
from contextlib import contextmanager
|
||||
|
||||
import pytest
|
||||
|
||||
@@ -191,12 +189,10 @@ def test_initialize_diffusion_replica_restores_device_visibility_after_local_ini
|
||||
old_env = os.environ.get(env_var)
|
||||
os.environ[env_var] = "0,1"
|
||||
|
||||
@contextmanager
|
||||
def _fake_stage_runtime_setup(_stage_id, _runtime_cfg):
|
||||
def _fake_setup_stage_devices(_stage_id, _runtime_cfg):
|
||||
current_omni_platform.set_device_control_env_var("1")
|
||||
yield
|
||||
|
||||
monkeypatch.setattr(engine_mod, "stage_runtime_setup", _fake_stage_runtime_setup)
|
||||
monkeypatch.setattr(engine_mod, "setup_stage_devices", _fake_setup_stage_devices)
|
||||
monkeypatch.setattr(engine_mod, "inject_kv_stage_info", lambda *_: None)
|
||||
monkeypatch.setattr(engine_mod, "initialize_diffusion_stage", lambda *_, **__: types.SimpleNamespace())
|
||||
|
||||
@@ -225,11 +221,7 @@ def test_initialize_diffusion_replica_passes_stage_init_timeout_and_inline_flag(
|
||||
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
@contextmanager
|
||||
def _noop_stage_runtime_setup(*_):
|
||||
yield
|
||||
|
||||
monkeypatch.setattr(engine_mod, "stage_runtime_setup", _noop_stage_runtime_setup)
|
||||
monkeypatch.setattr(engine_mod, "setup_stage_devices", lambda *_: None)
|
||||
monkeypatch.setattr(engine_mod, "inject_kv_stage_info", lambda *_: None)
|
||||
|
||||
def _capture_initialize_diffusion_stage(
|
||||
@@ -380,44 +372,6 @@ def test_build_logical_stage_init_plans_applies_replica_device_splits(monkeypatc
|
||||
assert all(replica.num_replicas == 3 for replica in stage_plans[1].replicas)
|
||||
|
||||
|
||||
def test_initialize_stage_replicas_reuses_long_lived_executor(monkeypatch):
|
||||
engine = object.__new__(AsyncOmniEngine)
|
||||
engine._stage_init_executor = None
|
||||
|
||||
cfg = types.SimpleNamespace(model_config=types.SimpleNamespace(max_model_len=64))
|
||||
stage_plans = [_make_llm_plan(0, configured_stage_id=0, vllm_config=cfg, num_replicas=1)]
|
||||
|
||||
created_executors = []
|
||||
|
||||
class FakeExecutor:
|
||||
def __init__(self, *args, **kwargs):
|
||||
self.shutdown_calls = []
|
||||
created_executors.append(self)
|
||||
|
||||
def submit(self, fn, *args, **kwargs):
|
||||
future = concurrent.futures.Future()
|
||||
future.set_result(fn(*args, **kwargs))
|
||||
return future
|
||||
|
||||
def shutdown(self, wait=False, cancel_futures=False):
|
||||
self.shutdown_calls.append((wait, cancel_futures))
|
||||
|
||||
monkeypatch.setattr(concurrent.futures, "ThreadPoolExecutor", FakeExecutor)
|
||||
monkeypatch.setattr(
|
||||
engine,
|
||||
"_initialize_replica",
|
||||
lambda plan, _stage_init_timeout, _stage_launch_lock: f"client-{plan.replica_id}",
|
||||
)
|
||||
|
||||
first = engine._initialize_stage_replicas(stage_plans, stage_init_timeout=1)
|
||||
second = engine._initialize_stage_replicas(stage_plans, stage_init_timeout=1)
|
||||
|
||||
assert first == {0: ["client-0"]}
|
||||
assert second == {0: ["client-0"]}
|
||||
assert len(created_executors) == 1
|
||||
assert engine._stage_init_executor is created_executors[0]
|
||||
|
||||
|
||||
def test_initialize_stage_replicas_collects_results_by_stage_and_replica_id(monkeypatch):
|
||||
engine = object.__new__(AsyncOmniEngine)
|
||||
|
||||
@@ -449,42 +403,6 @@ def test_initialize_stage_replicas_collects_results_by_stage_and_replica_id(monk
|
||||
}
|
||||
|
||||
|
||||
def test_shutdown_closes_long_lived_stage_init_executor():
|
||||
engine = object.__new__(AsyncOmniEngine)
|
||||
engine._shutdown_called = False
|
||||
engine._weak_finalizer = None
|
||||
engine._omni_master_server = None
|
||||
engine._coordinator_runtime = None
|
||||
engine.orchestrator_thread = types.SimpleNamespace(join=lambda timeout: None, is_alive=lambda: False)
|
||||
|
||||
class FakeQueue:
|
||||
def __init__(self):
|
||||
self.closed = False
|
||||
self.sync_q = types.SimpleNamespace(put_nowait=lambda _msg: None)
|
||||
|
||||
def close(self):
|
||||
self.closed = True
|
||||
|
||||
class FakeExecutor:
|
||||
def __init__(self):
|
||||
self.shutdown_calls = []
|
||||
|
||||
def shutdown(self, wait=False, cancel_futures=False):
|
||||
self.shutdown_calls.append((wait, cancel_futures))
|
||||
|
||||
engine.request_queue = FakeQueue()
|
||||
engine.output_queue = FakeQueue()
|
||||
engine.rpc_output_queue = FakeQueue()
|
||||
engine._stage_init_executor = FakeExecutor()
|
||||
|
||||
engine.shutdown()
|
||||
|
||||
assert engine._stage_init_executor is None
|
||||
assert engine.request_queue.closed is True
|
||||
assert engine.output_queue.closed is True
|
||||
assert engine.rpc_output_queue.closed is True
|
||||
|
||||
|
||||
def test_initialize_stages_cleans_up_successful_replicas_after_partial_multi_replica_failure(monkeypatch):
|
||||
import vllm_omni.engine.async_omni_engine as engine_mod
|
||||
|
||||
@@ -609,11 +527,7 @@ def test_initialize_llm_replica_passes_stage_init_timeout_to_complete_stage_hand
|
||||
prev_device_env = os.environ.get(device_env_var)
|
||||
os.environ[device_env_var] = "0"
|
||||
|
||||
@contextmanager
|
||||
def _noop_stage_runtime_setup(*_):
|
||||
yield
|
||||
|
||||
monkeypatch.setattr(engine_mod, "stage_runtime_setup", _noop_stage_runtime_setup)
|
||||
monkeypatch.setattr(engine_mod, "setup_stage_devices", lambda *_: None)
|
||||
monkeypatch.setattr(engine_mod, "build_engine_args_dict", lambda *_, **__: {})
|
||||
monkeypatch.setattr(engine_mod, "acquire_device_locks", lambda *_: [])
|
||||
monkeypatch.setattr(engine_mod, "spawn_stage_core", lambda **_: (fake_addresses, fake_proc, "ipc://handshake"))
|
||||
|
||||
@@ -4,7 +4,7 @@ from __future__ import annotations
|
||||
|
||||
import os
|
||||
import threading
|
||||
from contextlib import contextmanager, nullcontext
|
||||
from contextlib import contextmanager
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
|
||||
@@ -835,7 +835,7 @@ class TestSingleStageReplicaInitialization:
|
||||
prev_device_env = os.environ.get(device_env_var)
|
||||
os.environ[device_env_var] = "0"
|
||||
|
||||
mocker.patch.object(engine_mod, "stage_runtime_setup", return_value=nullcontext())
|
||||
mocker.patch.object(engine_mod, "setup_stage_devices")
|
||||
mocker.patch.object(engine_mod, "build_engine_args_dict", return_value={})
|
||||
mocker.patch.object(engine_mod, "acquire_device_locks", return_value=[])
|
||||
mocker.patch.object(engine_mod, "release_device_locks")
|
||||
@@ -909,7 +909,7 @@ class TestSingleStageReplicaInitialization:
|
||||
prev_device_env = os.environ.get(device_env_var)
|
||||
os.environ[device_env_var] = "0"
|
||||
|
||||
mocker.patch.object(engine_mod, "stage_runtime_setup", return_value=nullcontext())
|
||||
mocker.patch.object(engine_mod, "setup_stage_devices")
|
||||
mocker.patch.object(engine_mod, "inject_kv_stage_info")
|
||||
mocker.patch.object(engine_mod, "build_diffusion_config", return_value="diffusion-config")
|
||||
mock_register = mocker.patch.object(
|
||||
@@ -988,7 +988,7 @@ class TestSingleStageReplicaInitialization:
|
||||
prev_device_env = os.environ.get(device_env_var)
|
||||
os.environ[device_env_var] = "0"
|
||||
|
||||
mocker.patch.object(engine_mod, "stage_runtime_setup", return_value=nullcontext())
|
||||
mocker.patch.object(engine_mod, "setup_stage_devices")
|
||||
mocker.patch.object(engine_mod, "inject_kv_stage_info")
|
||||
mocker.patch.object(engine_mod, "build_diffusion_config", return_value="diffusion-config")
|
||||
mocker.patch.object(
|
||||
|
||||
@@ -9,7 +9,6 @@ OpenAI-compatible async text-to-image generation API endpoints in api_server.py.
|
||||
|
||||
import base64
|
||||
import io
|
||||
import json
|
||||
from argparse import Namespace
|
||||
from types import SimpleNamespace
|
||||
|
||||
@@ -20,7 +19,6 @@ from PIL import Image
|
||||
from pytest_mock import MockerFixture
|
||||
from vllm import SamplingParams
|
||||
from vllm.entrypoints.openai.models.protocol import BaseModelPath
|
||||
from vllm.sampling_params import RequestOutputKind
|
||||
|
||||
from vllm_omni.entrypoints.async_omni import AsyncOmni
|
||||
from vllm_omni.entrypoints.openai.api_server import _DiffusionServingModels, router
|
||||
@@ -122,27 +120,6 @@ class MockGenerationResult:
|
||||
self.peak_memory_mb = 0.0
|
||||
|
||||
|
||||
class MockStageResult:
|
||||
"""Mock multi-stage output for streaming image edit tests."""
|
||||
|
||||
def __init__(self, *, stage_id, final_output_type, text="", texts=None, images=None):
|
||||
self.stage_id = stage_id
|
||||
self.final_output_type = final_output_type
|
||||
self.images = images or []
|
||||
if texts is not None:
|
||||
outputs = [SimpleNamespace(text=item, index=index) for index, item in enumerate(texts)]
|
||||
elif text:
|
||||
outputs = [SimpleNamespace(text=text, index=0)]
|
||||
else:
|
||||
outputs = []
|
||||
self.request_output = SimpleNamespace(
|
||||
outputs=outputs,
|
||||
images=self.images,
|
||||
)
|
||||
self.stage_durations = {}
|
||||
self.peak_memory_mb = 0.0
|
||||
|
||||
|
||||
class FakeAsyncOmni:
|
||||
"""Fake AsyncOmni that yields a single diffusion output."""
|
||||
|
||||
@@ -414,72 +391,6 @@ def async_omni_stage_configs_only_client():
|
||||
return TestClient(app)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def streaming_image_edit_client():
|
||||
"""Create a multi-stage client whose engine yields AR text before image output."""
|
||||
from fastapi import FastAPI
|
||||
|
||||
from vllm_omni.entrypoints.async_omni import AsyncOmni
|
||||
from vllm_omni.entrypoints.openai.api_server import router
|
||||
from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
|
||||
|
||||
class FakeAsyncOmniClass(AsyncOmni):
|
||||
def __init__(self):
|
||||
stage_configs = [
|
||||
SimpleNamespace(stage_type="llm", is_comprehension=True),
|
||||
SimpleNamespace(stage_type="diffusion", is_comprehension=False),
|
||||
]
|
||||
default_sampling_params_list = [
|
||||
SamplingParams(temperature=0.1),
|
||||
OmniDiffusionSamplingParams(),
|
||||
]
|
||||
self.engine = SimpleNamespace(
|
||||
stage_configs=stage_configs,
|
||||
default_sampling_params_list=default_sampling_params_list,
|
||||
)
|
||||
self.default_sampling_params_list = default_sampling_params_list
|
||||
self.captured_sampling_params_list = None
|
||||
self.captured_prompt = None
|
||||
self.od_config = SimpleNamespace(supports_multimodal_inputs=True)
|
||||
|
||||
async def generate(self, prompt, request_id, sampling_params=None, sampling_params_list=None):
|
||||
self.captured_prompt = prompt
|
||||
self.captured_sampling_params_list = sampling_params_list or [sampling_params]
|
||||
assert self.captured_sampling_params_list[0].output_kind == RequestOutputKind.DELTA
|
||||
yield MockStageResult(stage_id=0, final_output_type="text", text="recap")
|
||||
yield MockStageResult(stage_id=0, final_output_type="text", text=" done")
|
||||
yield MockStageResult(
|
||||
stage_id=1,
|
||||
final_output_type="image",
|
||||
images=[Image.new("RGB", (32, 24), color="purple")],
|
||||
)
|
||||
|
||||
def __class_getitem__(cls, item):
|
||||
return cls
|
||||
|
||||
def get_diffusion_od_config(self):
|
||||
return self.od_config
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(router)
|
||||
|
||||
engine = FakeAsyncOmniClass()
|
||||
chat_handler = object.__new__(OmniOpenAIServingChat)
|
||||
chat_handler.engine_client = engine
|
||||
chat_handler._diffusion_engine = None
|
||||
app.state.openai_serving_chat = chat_handler
|
||||
app.state.engine_client = engine
|
||||
app.state.stage_configs = [
|
||||
SimpleNamespace(stage_type="llm"),
|
||||
SimpleNamespace(stage_type="diffusion"),
|
||||
]
|
||||
app.state.args = Namespace(
|
||||
default_sampling_params='{"1": {"num_inference_steps":4, "guidance_scale":7.5, "generator_device":"cpu"}}',
|
||||
max_generated_image_size=1024 * 1792,
|
||||
)
|
||||
return TestClient(app)
|
||||
|
||||
|
||||
def test_health_endpoint(test_client):
|
||||
"""Test health check endpoint for diffusion mode"""
|
||||
response = test_client.get("/health")
|
||||
@@ -769,154 +680,6 @@ def test_image_edits_async_omni_stage_configs_only(async_omni_stage_configs_only
|
||||
assert len(captured) == 2
|
||||
|
||||
|
||||
def _parse_sse_payloads(body: str):
|
||||
payloads = []
|
||||
for line in body.splitlines():
|
||||
if not line.startswith("data: "):
|
||||
continue
|
||||
data = line[len("data: ") :]
|
||||
payloads.append(data if data == "[DONE]" else json.loads(data))
|
||||
return payloads
|
||||
|
||||
|
||||
def test_image_edits_streaming_returns_ar_delta_then_image(streaming_image_edit_client):
|
||||
img_bytes = make_test_image_bytes((16, 16))
|
||||
response = streaming_image_edit_client.post(
|
||||
"/v1/images/edits",
|
||||
files=[("image", img_bytes)],
|
||||
data={
|
||||
"prompt": "edit me",
|
||||
"size": "auto",
|
||||
"stream": "true",
|
||||
"output_format": "png",
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.headers["content-type"].startswith("text/event-stream")
|
||||
|
||||
payloads = _parse_sse_payloads(response.text)
|
||||
assert [p["type"] if isinstance(p, dict) and "type" in p else p for p in payloads] == [
|
||||
"ar_delta",
|
||||
"ar_delta",
|
||||
"image",
|
||||
"[DONE]",
|
||||
]
|
||||
assert payloads[0]["delta"] == "recap"
|
||||
assert payloads[0]["index"] == 0
|
||||
assert payloads[1]["delta"] == " done"
|
||||
assert payloads[2]["output_format"] == "png"
|
||||
assert payloads[2]["size"] == "16x16"
|
||||
|
||||
image_payload = payloads[2]["data"][0]
|
||||
img = Image.open(io.BytesIO(base64.b64decode(image_payload["b64_json"])))
|
||||
assert img.size == (32, 24)
|
||||
|
||||
|
||||
def test_image_edits_streaming_ar_delta_chunks_include_index(streaming_image_edit_client):
|
||||
async def generate_multi_output_delta(prompt, request_id, sampling_params=None, sampling_params_list=None):
|
||||
yield MockStageResult(stage_id=0, final_output_type="text", texts=["first", "second"])
|
||||
yield MockStageResult(
|
||||
stage_id=1,
|
||||
final_output_type="image",
|
||||
images=[Image.new("RGB", (32, 24), color="purple")],
|
||||
)
|
||||
|
||||
streaming_image_edit_client.app.state.engine_client.generate = generate_multi_output_delta
|
||||
img_bytes = make_test_image_bytes((16, 16))
|
||||
response = streaming_image_edit_client.post(
|
||||
"/v1/images/edits",
|
||||
files=[("image", img_bytes)],
|
||||
data={
|
||||
"prompt": "edit me",
|
||||
"stream": "true",
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
payloads = _parse_sse_payloads(response.text)
|
||||
assert [(payloads[0]["index"], payloads[0]["delta"]), (payloads[1]["index"], payloads[1]["delta"])] == [
|
||||
(0, "first"),
|
||||
(1, "second"),
|
||||
]
|
||||
|
||||
|
||||
def test_image_edits_streaming_errors_without_final_image(streaming_image_edit_client):
|
||||
async def generate_without_image(prompt, request_id, sampling_params=None, sampling_params_list=None):
|
||||
yield MockStageResult(stage_id=0, final_output_type="text", text="recap")
|
||||
|
||||
streaming_image_edit_client.app.state.engine_client.generate = generate_without_image
|
||||
img_bytes = make_test_image_bytes((16, 16))
|
||||
response = streaming_image_edit_client.post(
|
||||
"/v1/images/edits",
|
||||
files=[("image", img_bytes)],
|
||||
data={
|
||||
"prompt": "edit me",
|
||||
"stream": "true",
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
payloads = _parse_sse_payloads(response.text)
|
||||
assert payloads[0]["type"] == "ar_delta"
|
||||
assert payloads[1]["object"] == "error"
|
||||
assert "without a final image" in payloads[1]["error"]["message"]
|
||||
assert payloads[2] == "[DONE]"
|
||||
|
||||
|
||||
def test_image_edits_streaming_errors_on_empty_final_image(streaming_image_edit_client):
|
||||
async def generate_empty_image(prompt, request_id, sampling_params=None, sampling_params_list=None):
|
||||
yield MockStageResult(stage_id=0, final_output_type="text", text="recap")
|
||||
yield MockStageResult(stage_id=1, final_output_type="image", images=[])
|
||||
|
||||
streaming_image_edit_client.app.state.engine_client.generate = generate_empty_image
|
||||
img_bytes = make_test_image_bytes((16, 16))
|
||||
response = streaming_image_edit_client.post(
|
||||
"/v1/images/edits",
|
||||
files=[("image", img_bytes)],
|
||||
data={
|
||||
"prompt": "edit me",
|
||||
"stream": "true",
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
payloads = _parse_sse_payloads(response.text)
|
||||
assert payloads[0]["type"] == "ar_delta"
|
||||
assert payloads[1]["object"] == "error"
|
||||
assert "empty final image" in payloads[1]["error"]["message"]
|
||||
assert payloads[2] == "[DONE]"
|
||||
|
||||
|
||||
def test_image_edits_streaming_rejects_single_stage(test_client):
|
||||
img_bytes = make_test_image_bytes((16, 16))
|
||||
response = test_client.post(
|
||||
"/v1/images/edits",
|
||||
files=[("image", img_bytes)],
|
||||
data={
|
||||
"prompt": "edit me",
|
||||
"stream": "true",
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 400
|
||||
assert "multi-stage" in response.json()["detail"]
|
||||
|
||||
|
||||
def test_image_edits_streaming_rejects_single_stage_before_loading_url(test_client):
|
||||
response = test_client.post(
|
||||
"/v1/images/edits",
|
||||
data={
|
||||
"prompt": "edit me",
|
||||
"url": "https://example.invalid/not-fetched.png",
|
||||
"stream": "true",
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 400
|
||||
assert "multi-stage" in response.json()["detail"]
|
||||
|
||||
|
||||
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(
|
||||
|
||||
@@ -1,289 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import dataclasses
|
||||
|
||||
import pytest
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class FakeAck:
|
||||
stage_id: int
|
||||
result: str = "ok"
|
||||
|
||||
|
||||
def _make_app(engine_client):
|
||||
from vllm_omni.entrypoints.openai.api_server import router
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(router)
|
||||
app.state.engine_client = engine_client
|
||||
app.state.sleeping_stages = set()
|
||||
return app
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sleep_capable_engine(mocker):
|
||||
engine = mocker.MagicMock()
|
||||
engine.sleep = mocker.AsyncMock(return_value=[FakeAck(stage_id=0), FakeAck(stage_id=1)])
|
||||
engine.wake_up = mocker.AsyncMock(return_value=[FakeAck(stage_id=0), FakeAck(stage_id=1)])
|
||||
return engine
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sleep_incapable_engine(mocker):
|
||||
return mocker.MagicMock(spec=[]) # no 'sleep' or 'wake_up' attributes
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# /v1/omni/sleep
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_sleep_success(sleep_capable_engine):
|
||||
app = _make_app(sleep_capable_engine)
|
||||
client = TestClient(app)
|
||||
|
||||
response = client.post("/v1/omni/sleep", json={"stage_ids": [0, 1], "level": 2})
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["status"] == "SUCCESS"
|
||||
assert [ack["stage_id"] for ack in data["acks"]] == [0, 1]
|
||||
assert app.state.sleeping_stages == {0, 1}
|
||||
sleep_capable_engine.sleep.assert_awaited_once_with(stage_ids=[0, 1], level=2)
|
||||
|
||||
|
||||
def test_sleep_default_level(sleep_capable_engine):
|
||||
app = _make_app(sleep_capable_engine)
|
||||
client = TestClient(app)
|
||||
|
||||
response = client.post("/v1/omni/sleep", json={"stage_ids": [0]})
|
||||
|
||||
assert response.status_code == 200
|
||||
sleep_capable_engine.sleep.assert_awaited_once_with(stage_ids=[0], level=2)
|
||||
|
||||
|
||||
def test_sleep_empty_stage_ids(sleep_capable_engine, mocker):
|
||||
"""Empty stage_ids: engine is still called and returns SUCCESS with no acks."""
|
||||
sleep_capable_engine.sleep = mocker.AsyncMock(return_value=[])
|
||||
app = _make_app(sleep_capable_engine)
|
||||
client = TestClient(app)
|
||||
|
||||
response = client.post("/v1/omni/sleep", json={"stage_ids": [], "level": 2})
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["status"] == "SUCCESS"
|
||||
assert data["acks"] == []
|
||||
assert app.state.sleeping_stages == set()
|
||||
sleep_capable_engine.sleep.assert_awaited_once_with(stage_ids=[], level=2)
|
||||
|
||||
|
||||
def test_sleep_updates_sleeping_set(sleep_capable_engine):
|
||||
app = _make_app(sleep_capable_engine)
|
||||
app.state.sleeping_stages = {5} # pre-existing entry
|
||||
client = TestClient(app)
|
||||
|
||||
client.post("/v1/omni/sleep", json={"stage_ids": [0, 1], "level": 2})
|
||||
|
||||
assert app.state.sleeping_stages == {5, 0, 1}
|
||||
|
||||
|
||||
def test_sleep_engine_not_support(sleep_incapable_engine):
|
||||
app = _make_app(sleep_incapable_engine)
|
||||
client = TestClient(app)
|
||||
|
||||
response = client.post("/v1/omni/sleep", json={"stage_ids": [0], "level": 1})
|
||||
|
||||
assert response.status_code == 501
|
||||
assert "sleep" in response.json()["detail"]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# /v1/omni/wakeup
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_wakeup_success(sleep_capable_engine):
|
||||
app = _make_app(sleep_capable_engine)
|
||||
app.state.sleeping_stages = {0, 1}
|
||||
client = TestClient(app)
|
||||
|
||||
response = client.post("/v1/omni/wakeup", json={"stage_ids": [0, 1]})
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["status"] == "SUCCESS"
|
||||
assert [ack["stage_id"] for ack in data["acks"]] == [0, 1]
|
||||
assert app.state.sleeping_stages == set()
|
||||
sleep_capable_engine.wake_up.assert_awaited_once_with(stage_ids=[0, 1])
|
||||
|
||||
|
||||
def test_wakeup_skipped_when_not_sleeping(sleep_capable_engine):
|
||||
app = _make_app(sleep_capable_engine)
|
||||
app.state.sleeping_stages = set() # nothing sleeping
|
||||
client = TestClient(app)
|
||||
|
||||
response = client.post("/v1/omni/wakeup", json={"stage_ids": [0, 1]})
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["status"] == "SKIPPED"
|
||||
sleep_capable_engine.wake_up.assert_not_awaited()
|
||||
|
||||
|
||||
def test_wakeup_partial_sleeping(sleep_capable_engine):
|
||||
"""Only stage 0 is sleeping; stage 1 is not. Should still proceed (partial match)."""
|
||||
app = _make_app(sleep_capable_engine)
|
||||
app.state.sleeping_stages = {0}
|
||||
client = TestClient(app)
|
||||
|
||||
response = client.post("/v1/omni/wakeup", json={"stage_ids": [0, 1]})
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["status"] == "SUCCESS"
|
||||
assert 0 not in app.state.sleeping_stages
|
||||
|
||||
|
||||
def test_wakeup_empty_stage_ids(sleep_capable_engine):
|
||||
"""Empty stage_ids: any() on empty iterable is False → SKIPPED, engine not called."""
|
||||
app = _make_app(sleep_capable_engine)
|
||||
app.state.sleeping_stages = {0, 1}
|
||||
client = TestClient(app)
|
||||
|
||||
response = client.post("/v1/omni/wakeup", json={"stage_ids": []})
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json()["status"] == "SKIPPED"
|
||||
sleep_capable_engine.wake_up.assert_not_awaited()
|
||||
|
||||
|
||||
def test_wakeup_engine_not_support(sleep_incapable_engine):
|
||||
app = _make_app(sleep_incapable_engine)
|
||||
app.state.sleeping_stages = {0}
|
||||
client = TestClient(app)
|
||||
|
||||
response = client.post("/v1/omni/wakeup", json={"stage_ids": [0]})
|
||||
|
||||
assert response.status_code == 501
|
||||
assert "wake_up" in response.json()["detail"]
|
||||
|
||||
|
||||
def test_wakeup_removes_only_requested_stages(sleep_capable_engine):
|
||||
app = _make_app(sleep_capable_engine)
|
||||
app.state.sleeping_stages = {0, 1, 2}
|
||||
client = TestClient(app)
|
||||
|
||||
client.post("/v1/omni/wakeup", json={"stage_ids": [0]})
|
||||
|
||||
assert app.state.sleeping_stages == {1, 2}
|
||||
|
||||
|
||||
def test_sleep_then_wakeup_roundtrip(sleep_capable_engine):
|
||||
app = _make_app(sleep_capable_engine)
|
||||
client = TestClient(app)
|
||||
|
||||
sleep_resp = client.post("/v1/omni/sleep", json={"stage_ids": [0, 1], "level": 2})
|
||||
assert sleep_resp.status_code == 200
|
||||
assert app.state.sleeping_stages == {0, 1}
|
||||
|
||||
wakeup_resp = client.post("/v1/omni/wakeup", json={"stage_ids": [0, 1]})
|
||||
assert wakeup_resp.status_code == 200
|
||||
assert app.state.sleeping_stages == set()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# pure-diffusion branch: omni_init_app_state initialises sleeping_stages
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def pure_diffusion_engine(mocker):
|
||||
engine = mocker.MagicMock()
|
||||
engine.stage_configs = [{"stage_type": "diffusion"}]
|
||||
engine.sleep = mocker.AsyncMock(return_value=[FakeAck(stage_id=0)])
|
||||
engine.wake_up = mocker.AsyncMock(return_value=[FakeAck(stage_id=0)])
|
||||
# Remove attributes that would make _get_vllm_config return a config
|
||||
del engine.get_vllm_config
|
||||
del engine.vllm_config
|
||||
return engine
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def pure_diffusion_app(pure_diffusion_engine, mocker):
|
||||
"""App whose state was initialised via omni_init_app_state (pure diffusion path)."""
|
||||
import asyncio
|
||||
from argparse import Namespace
|
||||
|
||||
from vllm_omni.entrypoints.openai.api_server import omni_init_app_state, router
|
||||
|
||||
args = Namespace(
|
||||
served_model_name=None,
|
||||
model="fake-diffusion-model",
|
||||
enable_log_requests=False,
|
||||
disable_log_stats=True,
|
||||
max_log_len=0,
|
||||
enable_server_load_tracking=False,
|
||||
)
|
||||
|
||||
app = FastAPI()
|
||||
app.include_router(router)
|
||||
|
||||
mocker.patch(
|
||||
"vllm_omni.entrypoints.openai.api_server.OmniOpenAIServingChat"
|
||||
).for_diffusion.return_value = mocker.MagicMock()
|
||||
mocker.patch(
|
||||
"vllm_omni.entrypoints.openai.api_server.OmniOpenAIServingVideo"
|
||||
).for_diffusion.return_value = mocker.MagicMock()
|
||||
mocker.patch(
|
||||
"vllm_omni.entrypoints.openai.api_server.OmniOpenAIServingSpeech"
|
||||
).for_diffusion.return_value = mocker.MagicMock()
|
||||
mocker.patch("vllm_omni.entrypoints.openai.api_server._DiffusionServingModels")
|
||||
|
||||
loop = asyncio.new_event_loop()
|
||||
try:
|
||||
loop.run_until_complete(omni_init_app_state(pure_diffusion_engine, app.state, args))
|
||||
finally:
|
||||
loop.close()
|
||||
|
||||
return app
|
||||
|
||||
|
||||
def test_pure_diffusion_sleep_success(pure_diffusion_app, pure_diffusion_engine):
|
||||
client = TestClient(pure_diffusion_app)
|
||||
|
||||
response = client.post("/v1/omni/sleep", json={"stage_ids": [0], "level": 2})
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json()["status"] == "SUCCESS"
|
||||
assert pure_diffusion_app.state.sleeping_stages == {0}
|
||||
pure_diffusion_engine.sleep.assert_awaited_once_with(stage_ids=[0], level=2)
|
||||
|
||||
|
||||
def test_pure_diffusion_wakeup_skipped(pure_diffusion_app, pure_diffusion_engine):
|
||||
# sleeping_stages is empty after init — nothing is sleeping
|
||||
client = TestClient(pure_diffusion_app)
|
||||
|
||||
response = client.post("/v1/omni/wakeup", json={"stage_ids": [0]})
|
||||
|
||||
assert response.status_code == 200
|
||||
assert response.json()["status"] == "SKIPPED"
|
||||
pure_diffusion_engine.wake_up.assert_not_awaited()
|
||||
|
||||
|
||||
def test_pure_diffusion_sleep_then_wakeup(pure_diffusion_app, pure_diffusion_engine):
|
||||
client = TestClient(pure_diffusion_app)
|
||||
|
||||
sleep_resp = client.post("/v1/omni/sleep", json={"stage_ids": [0], "level": 2})
|
||||
assert sleep_resp.status_code == 200
|
||||
assert pure_diffusion_app.state.sleeping_stages == {0}
|
||||
|
||||
wakeup_resp = client.post("/v1/omni/wakeup", json={"stage_ids": [0]})
|
||||
assert wakeup_resp.status_code == 200
|
||||
assert pure_diffusion_app.state.sleeping_stages == set()
|
||||
@@ -184,7 +184,7 @@ def test_direct_omni_with_nullified_parser_only_nulls_untyped_override_fields(
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--gpu-memory-utilization", type=float, default=0.9)
|
||||
parser.add_argument("--batch-timeout", type=int, default=10)
|
||||
parser.add_argument("--hsdp-shard-size", type=int, default=-1)
|
||||
nullify_stage_engine_defaults(parser)
|
||||
args = parser.parse_args([])
|
||||
args.model = "fake-model"
|
||||
@@ -192,7 +192,7 @@ def test_direct_omni_with_nullified_parser_only_nulls_untyped_override_fields(
|
||||
Omni(**vars(args))
|
||||
|
||||
assert captured["gpu_memory_utilization"] is None
|
||||
assert captured["batch_timeout"] == 10
|
||||
assert captured["hsdp_shard_size"] == -1
|
||||
assert "_cli_explicit_keys" not in captured
|
||||
|
||||
|
||||
@@ -233,7 +233,7 @@ def test_deprecated_from_cli_args_preserves_legacy_parser_nulling(
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--gpu-memory-utilization", type=float, default=0.9)
|
||||
parser.add_argument("--batch-timeout", type=int, default=10)
|
||||
parser.add_argument("--hsdp-shard-size", type=int, default=-1)
|
||||
args = parser.parse_args([])
|
||||
args.model = "fake-model"
|
||||
|
||||
@@ -241,7 +241,7 @@ def test_deprecated_from_cli_args_preserves_legacy_parser_nulling(
|
||||
Omni.from_cli_args(args, parser=parser)
|
||||
|
||||
assert captured["gpu_memory_utilization"] is None
|
||||
assert captured["batch_timeout"] == 10
|
||||
assert captured["hsdp_shard_size"] == -1
|
||||
|
||||
|
||||
def _make_base():
|
||||
|
||||
@@ -1,35 +0,0 @@
|
||||
"""
|
||||
Offline inference tests: vace video generation.
|
||||
See examples/offline_inference/vace/vace_video_generation.md
|
||||
"""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.examples.helpers import EXAMPLES, ExampleRunner, ReadmeSnippet
|
||||
from tests.helpers.assertions import assert_video_valid
|
||||
from tests.helpers.mark import hardware_marks
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.usefixtures("clean_gpu_memory_between_tests"),
|
||||
pytest.mark.full_model,
|
||||
pytest.mark.example,
|
||||
*hardware_marks(res={"cuda": "H100"}),
|
||||
]
|
||||
|
||||
VACE_SCRIPT = EXAMPLES / "offline_inference" / "vace" / "vace_video_generation.py"
|
||||
README_PATH = VACE_SCRIPT.with_name("vace_video_generation.md")
|
||||
EXAMPLE_OUTPUT_SUBFOLDER = "example_offline_vace"
|
||||
|
||||
README_SNIPPETS = ReadmeSnippet.extract_readme_snippets(README_PATH)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("snippet", README_SNIPPETS, ids=lambda snippet: snippet.test_id)
|
||||
def test_vace_video_generation(snippet: ReadmeSnippet, example_runner: ExampleRunner):
|
||||
should_skip, reason = snippet.skip
|
||||
if should_skip:
|
||||
pytest.skip(reason)
|
||||
result = example_runner.run(snippet, output_subfolder=Path(EXAMPLE_OUTPUT_SUBFOLDER))
|
||||
for asset in result.assets:
|
||||
assert_video_valid(asset)
|
||||
@@ -477,7 +477,6 @@ def assert_omni_response(response: Any, request_config: dict[str, Any], run_leve
|
||||
# Verify similarity (Whisper transcript vs streamed/detokenized text)
|
||||
if "audio" in modalities:
|
||||
audio_ref_text = request_config.get("audio_ref_text")
|
||||
similarity_threshold = request_config.get("similarity_threshold", 0.8)
|
||||
if "text" in modalities:
|
||||
transcript = (response.audio_content or "").strip()
|
||||
text_output = (response.text_content or "").strip()
|
||||
@@ -485,14 +484,14 @@ def assert_omni_response(response: Any, request_config: dict[str, Any], run_leve
|
||||
transcript.lower(),
|
||||
text_output.lower(),
|
||||
)
|
||||
assert similarity > 0.9, "The audio content is not same as the text"
|
||||
print(f"similarity is: {similarity}")
|
||||
assert similarity > similarity_threshold, "The audio content is not same as the text"
|
||||
if audio_ref_text:
|
||||
audio_similarity = cosine_similarity_text(
|
||||
response.audio_content.lower(),
|
||||
str(audio_ref_text).lower(),
|
||||
)
|
||||
assert audio_similarity > similarity_threshold, (
|
||||
assert audio_similarity > 0.9, (
|
||||
f"The audio content does not match reference text: similarity={audio_similarity:.3f}"
|
||||
)
|
||||
|
||||
|
||||
Loaded 100 of 261 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user