Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
755a2d9d96 | ||
|
|
c174c95ab0 | ||
|
|
e375b12680 | ||
|
|
52c9ce2dac | ||
|
|
9a1270e6cc | ||
|
|
8e9b497387 | ||
|
|
39e3fd3d12 | ||
|
|
d388583e0b | ||
|
|
7b876324c2 | ||
|
|
b8bd6fbbdb | ||
|
|
37e5954acf | ||
|
|
22ffb85063 | ||
|
|
b33d536f63 | ||
|
|
0fe76857ac | ||
|
|
e8d22d4fa8 | ||
|
|
77e8ba0a7b | ||
|
|
c6edb34d04 | ||
|
|
fb4204ebf6 | ||
|
|
e24f86ae2e | ||
|
|
678af19240 | ||
|
|
5632b78b3d | ||
|
|
345504b68f | ||
|
|
d20a6b5b39 | ||
|
|
9ebfbf9fde | ||
|
|
2e067bbc24 | ||
|
|
19980bb37c | ||
|
|
c8efdbd671 | ||
|
|
95a07f7732 | ||
|
|
47edee1a79 | ||
|
|
05c6c9a0b6 | ||
|
|
d14ea37842 | ||
|
|
1a7fb55b5f | ||
|
|
3a32cd6780 | ||
|
|
f327cd1c89 | ||
|
|
d4cbdff9cd | ||
|
|
8384dead56 | ||
|
|
3aa84c4742 | ||
|
|
fd9a0ee27f | ||
|
|
e3b0afbb97 | ||
|
|
ead87aa521 | ||
|
|
d8cc7a0409 | ||
|
|
a664e40716 | ||
|
|
ee15f39ea7 | ||
|
|
9337bec9ad | ||
|
|
9d1392dcfc | ||
|
|
e18cb89d5a | ||
|
|
554233208f | ||
|
|
e1cdd0c23b | ||
|
|
e82c0a5e1d | ||
|
|
8e1bbc9ab5 | ||
|
|
2903f528d2 | ||
|
|
e6cf26efb7 | ||
|
|
7a34298743 | ||
|
|
b60a661b3f | ||
|
|
339a6da28c | ||
|
|
25b451a552 | ||
|
|
0cdec8210f | ||
|
|
3f504b45d2 | ||
|
|
c2578f7131 | ||
|
|
5f029ee4d5 | ||
|
|
b001151e0d | ||
|
|
ad7c966b1d | ||
|
|
52b53367ef | ||
|
|
c1ba86ab5b | ||
|
|
de4e47237a |
No files matched your search
@@ -11,15 +11,6 @@ set -o pipefail
|
||||
export PYTHONPATH=".."
|
||||
|
||||
# Print ROCm version
|
||||
echo "--- Confirming Clean Initial State"
|
||||
while true; do
|
||||
sleep 3
|
||||
if grep -q clean /opt/amdgpu/etc/gpu_state; then
|
||||
echo "GPUs state is \"clean\""
|
||||
break
|
||||
fi
|
||||
done
|
||||
|
||||
echo "--- ROCm info"
|
||||
rocminfo
|
||||
|
||||
@@ -51,18 +42,6 @@ cleanup_docker() {
|
||||
# Call the cleanup docker function
|
||||
cleanup_docker
|
||||
|
||||
echo "--- Resetting GPUs"
|
||||
|
||||
echo "reset" > /opt/amdgpu/etc/gpu_state
|
||||
|
||||
while true; do
|
||||
sleep 3
|
||||
if grep -q clean /opt/amdgpu/etc/gpu_state; then
|
||||
echo "GPUs state is \"clean\""
|
||||
break
|
||||
fi
|
||||
done
|
||||
|
||||
echo "--- Pulling container"
|
||||
## Temporary change to use AMD Docker Hub to store the vllm-omni image
|
||||
# to bypass the rate limit issue with ECR Public Gallery.
|
||||
|
||||
@@ -275,6 +275,46 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Audio Streaming Input Test with H100"
|
||||
timeout_in_minutes: 30
|
||||
depends_on: upload-merge-pipeline
|
||||
commands:
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY="1"
|
||||
- pytest -s -v tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py -m "advanced_model" --run-level "advanced_model"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 2
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Diffusion Image Edit Test with H100 (1 GPU)"
|
||||
timeout_in_minutes: 20
|
||||
depends_on: upload-merge-pipeline
|
||||
|
||||
+107
-12
@@ -13,7 +13,7 @@ steps:
|
||||
- label: ":full_moon: Omni · Function Test"
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and H100 and omni" --run-level "advanced_model"
|
||||
- pytest -s -v tests/e2e/ -m "full_model and H100 and omni" --run-level "full_model" --ignore=tests/e2e/accuracy
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
@@ -52,7 +52,7 @@ steps:
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
- pytest -s -v tests/examples/ -m "advanced_model and omni and L4" --run-level "advanced_model"
|
||||
- pytest -s -v tests/examples/ -m "full_model and omni and L4" --run-level "full_model"
|
||||
agents:
|
||||
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
|
||||
plugins:
|
||||
@@ -70,7 +70,52 @@ steps:
|
||||
- label: ":full_moon: Omni · Doc Test with H100"
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- pytest -s -v tests/examples/ -m "advanced_model and omni and H100" --run-level "advanced_model"
|
||||
- pytest -s -v tests/examples/ -m "full_model and omni and H100" --run-level "full_model"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 2
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: Omni · Accuracy Test"
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- export SEED_TTS_WER_EVAL=1
|
||||
- export SEED_TTS_EVAL_DEVICE=cuda:1
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/e2e/accuracy/qwen3_omni/test_qwen3_omni.py -m "full_model" --run-level full_model
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/e2e/accuracy/qwen3_omni/results/qwen_omni_acc/*.json"
|
||||
exit $$EXIT
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
@@ -160,7 +205,7 @@ steps:
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
- pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and L4 and omni" --run-level "advanced_model"
|
||||
- pytest -s -v tests/e2e/ -m "full_model and L4 and omni" --run-level "full_model" --ignore=tests/e2e/accuracy
|
||||
agents:
|
||||
queue: "gpu_1_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
|
||||
plugins:
|
||||
@@ -233,7 +278,7 @@ steps:
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100"
|
||||
timeout_in_minutes: 120
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_*_expansion.py -k "not test_wan22_expansion and not test_wan_2_1_vace_expansion and not hunyuan" -m "advanced_model and diffusion and H100" --run-level "advanced_model"
|
||||
- pytest -sv tests/e2e/ -k "not test_wan and not hunyuan" -m "full_model and diffusion and H100" --run-level "full_model" --ignore=tests/e2e/accuracy
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
@@ -271,7 +316,7 @@ steps:
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with L4"
|
||||
timeout_in_minutes: 60
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_*_expansion.py -k "not test_wan22_expansion and not test_wan_2_1_vace_expansion and not hunyuan" -m "advanced_model and diffusion and L4" --run-level "advanced_model"
|
||||
- pytest -sv tests/e2e/ -k "not test_wan and not hunyuan" -m "full_model and diffusion and L4" --run-level "full_model" --ignore=tests/e2e/accuracy
|
||||
agents:
|
||||
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
|
||||
plugins:
|
||||
@@ -290,7 +335,7 @@ steps:
|
||||
timeout_in_minutes: 60
|
||||
commands:
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY="1"
|
||||
- pytest -s -v tests/examples/*/test_text_to_image.py -m "advanced_model and example and H100" --run-level "advanced_model"
|
||||
- pytest -s -v tests/examples/*/test_text_to_image.py -m "full_model and example and H100" --run-level "full_model"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
@@ -328,7 +373,7 @@ steps:
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · GEBench Accuracy Test"
|
||||
timeout_in_minutes: 60
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/accuracy/test_gebench_h100_smoke.py --run-level advanced_model --gebench-model Qwen/Qwen-Image-2512 --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gebench-port 8093 --accuracy-workers 1
|
||||
- pytest -s -v tests/e2e/accuracy/test_gebench_h100_smoke.py --run-level full_model --gebench-model Qwen/Qwen-Image-2512 --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gebench-port 8093 --accuracy-workers 1
|
||||
- buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gebench_qwen-image-2512/summary*.json"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
@@ -367,7 +412,7 @@ steps:
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · GEdit-Bench Accuracy Test"
|
||||
timeout_in_minutes: 60
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/accuracy/test_gedit_bench_h100_smoke.py --run-level advanced_model --gedit-model Qwen/Qwen-Image-Edit --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gedit-port 8093 --gedit-samples-per-group 20 --accuracy-workers 1
|
||||
- pytest -s -v tests/e2e/accuracy/test_gedit_bench_h100_smoke.py --run-level full_model --gedit-model Qwen/Qwen-Image-Edit --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gedit-port 8093 --gedit-samples-per-group 20 --accuracy-workers 1
|
||||
- buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gedit_scores_qwen-image-edit/qwen-image-edit_all_all_vie_score_*.csv"
|
||||
- buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gedit_scores_qwen-image-edit/qwen-image-edit_all_all_summary_*.json"
|
||||
agents:
|
||||
@@ -409,7 +454,7 @@ steps:
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Accuracy Test"
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/accuracy/test_qwen_image*.py --run-level advanced_model
|
||||
- pytest -s -v tests/e2e/accuracy/test_qwen_image*.py --run-level full_model
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
@@ -514,7 +559,7 @@ steps:
|
||||
- label: ":full_moon: Diffusion X2V · Function Test"
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_hunyuan_video_15_expansion.py -m "advanced_model" --run-level "advanced_model"
|
||||
- pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_hunyuan_video_15_expansion.py --run-level "full_model"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
@@ -552,7 +597,55 @@ steps:
|
||||
- label: ":full_moon: Diffusion X2V · Accuracy Test"
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py -m advanced_model --run-level advanced_model
|
||||
- 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/test_ltx2_3_video_similarity.py -m advanced_model --run-level advanced_model
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 2
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: Diffusion X2V · Perf Test"
|
||||
key: nightly-diffusion-x2v-performance
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json
|
||||
EXIT1=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
|
||||
exit $$EXIT1
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
@@ -636,6 +729,7 @@ steps:
|
||||
- nightly-omni-performance
|
||||
- nightly-tts-performance
|
||||
- nightly-diffusion-x2iat-performance
|
||||
- nightly-diffusion-x2v-performance
|
||||
- nightly-testcase-statistics
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
commands:
|
||||
@@ -645,6 +739,7 @@ steps:
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-tts-performance
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance
|
||||
- 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
|
||||
- python tools/nightly/generate_nightly_perf_html.py
|
||||
|
||||
@@ -120,23 +120,6 @@ steps:
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Audio Generation Model Test"
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
- timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_expansion.py -m "advanced_model and diffusion and L4" --run-level advanced_model
|
||||
agents:
|
||||
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Diffusion Cache Backend Test"
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
|
||||
@@ -282,10 +282,11 @@ For Omni or custom models, create:
|
||||
|
||||
Required updates:
|
||||
1. `docs/user_guide/diffusion/parallelism_acceleration.md` — parallelism support table
|
||||
2. `docs/user_guide/diffusion/teacache.md` — if TeaCache supported
|
||||
3. `docs/user_guide/diffusion/cache_dit_acceleration.md` — if Cache-DiT supported
|
||||
4. `examples/offline_inference/xxx/README.md` — offline example docs
|
||||
5. `examples/online_serve/xxx/README.md` — online serve docs
|
||||
2. `docs/user_guide/diffusion/cpu_offload_diffusion.md` — if CPU offload supported (add to supported models table)
|
||||
3. `docs/user_guide/diffusion/teacache.md` — if TeaCache supported
|
||||
4. `docs/user_guide/diffusion/cache_dit_acceleration.md` — if Cache-DiT supported
|
||||
5. `examples/offline_inference/xxx/README.md` — offline example docs
|
||||
6. `examples/online_serve/xxx/README.md` — online serve docs
|
||||
|
||||
### Step 8: Add E2E Tests (Recommended)
|
||||
|
||||
@@ -512,6 +513,37 @@ After adding parallelism support, update:
|
||||
1. `docs/user_guide/diffusion/parallelism_acceleration.md` — add your model to the support table
|
||||
2. Record which parallelism methods are supported (USP, Ring, CFG, TP, HSDP, VAE-Patch)
|
||||
|
||||
### Step 11: Add CPU Offload Support
|
||||
|
||||
Implement `SupportsModuleOffload` on your pipeline class to enable
|
||||
`--enable-cpu-offload` and `--enable-layerwise-offload`. The protocol
|
||||
declares which submodules the offloader should manage:
|
||||
|
||||
```python
|
||||
from typing import ClassVar
|
||||
from vllm_omni.diffusion.models.interface import SupportsModuleOffload
|
||||
|
||||
class YourPipeline(nn.Module, SupportsModuleOffload):
|
||||
_dit_modules: ClassVar[list[str]] = ["transformer"]
|
||||
_encoder_modules: ClassVar[list[str]] = ["text_encoder"]
|
||||
_vae_modules: ClassVar[list[str]] = ["vae"]
|
||||
_resident_modules: ClassVar[list[str]] = [] # optional
|
||||
```
|
||||
|
||||
- `_dit_modules`: denoising submodules (kept on GPU during diffusion loop)
|
||||
- `_encoder_modules`: encoder/vision submodules (offloaded to CPU during diffusion loop)
|
||||
- `_vae_modules`: VAE(s) (handled by both sequential and layerwise backends)
|
||||
- `_resident_modules`: additional modules to pin on GPU during layerwise
|
||||
offloading (e.g. embedders, connectors). Only used by the layerwise
|
||||
backend. Optional — defaults to `[]`.
|
||||
|
||||
All attribute names support dotted paths for nested submodules
|
||||
(e.g. `"pipe.transformer"`, `"bagel.time_embedder"`).
|
||||
|
||||
Pipelines without `SupportsModuleOffload` fall back to scanning
|
||||
well-known attribute names (`transformer`, `text_encoder`, `vae`,
|
||||
etc.), which fails for non-standard names.
|
||||
|
||||
---
|
||||
|
||||
## Iterative Development Tips
|
||||
|
||||
@@ -8,10 +8,90 @@ description: "Integrate a new text-to-speech model into vLLM-Omni from HuggingFa
|
||||
## Overview
|
||||
|
||||
```
|
||||
HF Reference -> Stage Separation -> Online Serving -> Async Chunk -> CUDA Graph
|
||||
(Phase 1) (Phase 2) (Phase 3) (Phase 4) (Phase 5)
|
||||
HF Reference -> Stage Separation -> Online Serving -> Async Chunk -> CUDA Graph -> Pre-commit/DCO
|
||||
(Phase 1) (Phase 2) (Phase 3) (Phase 4) (Phase 5) (Phase 6)
|
||||
```
|
||||
|
||||
Three architecture patterns are supported:
|
||||
|
||||
- **Two-stage pipeline** (e.g. Qwen3-TTS, Fish Speech, CosyVoice3): AR
|
||||
code-predictor → audio decoder, connected via async_chunk for low-latency
|
||||
streaming. Use this for maximum performance.
|
||||
- **Single-stage AR via generator** (e.g. MOSS-TTS-Nano): entire model runs
|
||||
inside one AR worker, streaming audio chunks through a per-request
|
||||
`inference_stream()` generator. Use this when the upstream model bundles AR
|
||||
+ codec inseparably. See [references/single-stage-ar.md](references/single-stage-ar.md).
|
||||
- **Single-stage, vLLM-native base LM + side computation** (e.g. VoxCPM2):
|
||||
the base language model runs under vLLM's PagedAttention as a normal AR
|
||||
model; diffusion / VAE / side computations run outside vLLM and are
|
||||
attached via the runner post-processing hook. This is a distinct pattern
|
||||
from the generator approach above — do not confuse the two.
|
||||
|
||||
The single-stage variants skip Phase 4 (async_chunk) but Phase 5 (CUDA graph)
|
||||
is still encouraged for the inner AR loop.
|
||||
|
||||
## Cross-Cutting Invariants
|
||||
|
||||
These rules apply to every TTS model regardless of architecture (AR vs AR+diffusion, single-stage vs two-stage, codec-based vs VAE-based). They surface repeatedly across PRs — check them at the end of every phase.
|
||||
|
||||
### I1. Streaming output contract
|
||||
|
||||
Pick exactly one per-step semantics for `forward()` and document it in the docstring:
|
||||
|
||||
- **Delta**: yield only new audio samples produced this step. Preferred — linear cost, low memory.
|
||||
- **Cumulative**: re-decode from step 0 every call. O(N²); only acceptable if the codec has no streaming decode path.
|
||||
|
||||
If you choose **delta**, verify the full emit→consolidate→consume chain:
|
||||
|
||||
1. `forward()` returns `{"model_outputs": <new_chunk_only>, ...}`
|
||||
2. `_consolidate_multimodal_tensors()` in `vllm_omni/engine/output_processor.py` concatenates the audio key into one tensor at finish. If it skips the key (`continue`), offline consumers receive only the final chunk. See `output_processor.py` for the concrete list of handled modality keys.
|
||||
3. Streaming consumers (SSE, Gradio) receive per-step deltas; offline consumers (`engine.generate()`) receive a single concatenated tensor.
|
||||
|
||||
Cumulative-vs-delta mismatch is the most common silent bug — offline RTF benchmarks pass, but users hear replays or truncation.
|
||||
|
||||
### I2. Multimodal output consumer hygiene
|
||||
|
||||
`outputs[0].outputs[0].multimodal_output[<key>]` can be any of `Tensor`, `list[Tensor]` (pre-consolidation snapshot), `np.ndarray`, or scalar. When writing tests, examples, and benchmarks:
|
||||
|
||||
- **Never** use `dict.get("a") or dict.get("b")` on tensor values — Python evaluates the tensor's boolean, raising `RuntimeError: Boolean value of Tensor with more than one value is ambiguous`. Use explicit `if x is None` chains.
|
||||
- Always defensively handle the list form: `if isinstance(x, list): x = torch.cat([t.reshape(-1) for t in x], dim=0)`.
|
||||
- Assert `shape` / `dtype` / `duration` explicitly; do not rely on truthiness for presence checks.
|
||||
|
||||
### I3. Hot-loop GPU discipline
|
||||
|
||||
Inside any per-step model loop (AR decode, diffusion solver, CFM Euler, vocoder block loop):
|
||||
|
||||
- No `tensor.item()`, `.cpu()`, or `.tolist()` — each triggers a GPU→CPU sync; at 10 steps × 60 frames × 4 ops that is 2400 syncs per request.
|
||||
- Prefer `dst.copy_(src)` over `dst.fill_(src.item())` when writing a scalar tensor into a buffer.
|
||||
- Prefer `torch.compile(Model.forward, fullgraph=False)` on the whole forward over per-submodule compile — fewer dispatch boundaries, larger fusion regions. Measure before choosing granularity.
|
||||
- No Python-side control flow that depends on tensor values; use `torch.where` / masking instead.
|
||||
|
||||
Profile first, optimize second. See the profiling docs / project memory for the trace-analysis workflow.
|
||||
|
||||
### I4. Validation pyramid
|
||||
|
||||
Offline RTF alone is necessary but not sufficient. Every new TTS model must pass all three:
|
||||
|
||||
| Layer | Catches | Tool |
|
||||
|-------|---------|------|
|
||||
| Offline RTF / duration check | Throughput regressions, missing audio, wrong sample rate | `end2end.py`, pytest e2e |
|
||||
| Browser streaming playback | Delta/cumulative bugs, chunk boundary glitches, TTFP regressions | Gradio demo over `/v1/audio/speech?stream=true` |
|
||||
| Concurrent requests | Per-request state leaks, codec window round-robin gaps | `max_num_seqs>1` smoke test with 4+ parallel prompts |
|
||||
|
||||
Declaring a model "done" without all three has shipped regressions more than once.
|
||||
|
||||
### I5. Per-request state is owned by the request, not the model
|
||||
|
||||
If the model caches *anything* across `forward()` calls (streaming generators, codec buffers, sliding-window pads, CUDA graph state), key it by request ID:
|
||||
|
||||
```python
|
||||
self._state: dict[str, YourState] = {} # request_key → state
|
||||
# fetch: request_key = str(info.get("_omni_req_id", "0"))
|
||||
# free on finish: del self._state[request_key]
|
||||
```
|
||||
|
||||
A shared buffer silently corrupts audio across concurrent requests — the symptom is crosstalk or truncation only under load.
|
||||
|
||||
## Phase 1: HuggingFace Reference
|
||||
|
||||
**Goal**: Understand the reference implementation and verify it produces correct audio.
|
||||
@@ -85,6 +165,48 @@ When audio output is wrong, check in this order:
|
||||
6. **Tensor layout**: Codebook-major vs frame-major ordering
|
||||
7. **dtype**: Float32 for codec decoders (autocast can corrupt audio)
|
||||
|
||||
### Streaming Correctness Rules (single-stage and two-stage)
|
||||
|
||||
These bugs appear in almost every new TTS PR. Check all before the first push. See also the cross-cutting invariants I1 (output contract) and I5 (per-request state) above — the rules below are the Phase 2-specific instances of those invariants:
|
||||
|
||||
- **Accumulate codes across AR steps** — each `forward()` appends new codes; do not reset between steps or audio will be truncated (fish speech: `fix: accumulate audio_codes across steps`)
|
||||
- **Emit delta audio, not full waveform** — in streaming mode yield only the new chunk per step, not the re-decoded full waveform from step 0 (fish speech: `fix: emit delta audio not full waveform`)
|
||||
- **All return paths must emit `model_outputs`** — if any early-return branch skips setting `model_outputs`, the serving layer silently drops that step's audio (fish speech: `fix: ensure ALL return paths emit model_outputs`)
|
||||
- **Per-request state isolation** — for batched concurrent requests, key all state by request ID; a shared buffer corrupts audio across requests (fish speech: `fix: per-request vocode + delta emission`)
|
||||
- **Codec tensor device** — move codec codes to the codec decoder's device before calling decode; mismatches cause silent CPU fallback or crashes (fish speech: `fix: use model device for CUDA stream`)
|
||||
- **AR stage `max_num_seqs`** — set to at least 4 in production configs; for single-stage models this is the only stage. For two-stage models, Stage 0 (AR) needs `max_num_seqs ≥ 4` to pipeline concurrent requests; Stage 1 (codec decoder) typically uses `max_num_seqs: 1` intentionally. Default of 1 everywhere causes audio gaps under concurrency because the codec window round-robins across requests (RFC #2568)
|
||||
|
||||
### Optional Dependency Handling
|
||||
|
||||
Patch optional dependencies (`torchaudio` / `torchcodec` / `soundfile`) at
|
||||
the top of `load_weights()`, not at module import. Failures to do so cause
|
||||
cryptic errors only on environments missing the optional package — after
|
||||
the model is already deployed. See
|
||||
[references/optional-deps.md](references/optional-deps.md) for the full
|
||||
pattern, signature constraints, and MOSS-TTS-Nano reference.
|
||||
|
||||
### Single-Stage AR Pattern (alternative to two-stage)
|
||||
|
||||
When the upstream model cannot be cleanly split into an AR stage and a
|
||||
separate decoder, run the full pipeline inside a single AR worker and
|
||||
stream audio through a per-request `inference_stream()` generator keyed by
|
||||
`_omni_req_id`. Stage config must set `worker_type: ar`,
|
||||
`engine_output_type: audio`, `final_output: true`, `is_comprehension: true`,
|
||||
and `async_chunk: false` at the top level. Only extract params from
|
||||
`additional_information` that you actually forward, or pre-commit fails
|
||||
`ruff F841`.
|
||||
|
||||
Full walkthrough with the complete `forward()` / `_create_stream_gen()`
|
||||
skeleton and stage-config fields:
|
||||
[references/single-stage-ar.md](references/single-stage-ar.md). For an
|
||||
in-tree reference, look for any single-stage AR model under
|
||||
`vllm_omni/model_executor/models/` — e.g. the MOSS-TTS-Nano integration when
|
||||
it lands.
|
||||
|
||||
**VoxCPM2 is a different pattern** and should not reuse this skeleton — it
|
||||
runs the base LM under vLLM PagedAttention with external side-computation.
|
||||
See `plan/voxcpm2_native_ar_design.md`.
|
||||
|
||||
### Deliverables
|
||||
|
||||
- Model files in `vllm_omni/model_executor/models/<model_name>/`
|
||||
@@ -98,10 +220,66 @@ When audio output is wrong, check in this order:
|
||||
|
||||
### Steps
|
||||
|
||||
1. **Register in `serving_speech.py`**:
|
||||
- Add model stage name to `_TTS_MODEL_STAGES` set
|
||||
- Add model detection flag (e.g., `_is_fish_speech`)
|
||||
- Implement prompt builder method (e.g., `_build_fish_speech_prompt()`)
|
||||
1. **Register in `serving_speech.py`** — add all 5 points in a **single commit**;
|
||||
partial integration causes hard-to-debug failures. This file is modified by every
|
||||
model PR and is the most common source of rebase conflicts — see conflict note below.
|
||||
|
||||
**Point 1** — stage constant (near the top, alongside the other `_*_TTS_MODEL_STAGES` sets):
|
||||
```python
|
||||
_YOUR_MODEL_TTS_MODEL_STAGES = {"your_stage_key"}
|
||||
```
|
||||
|
||||
**Point 2** — union into `_TTS_MODEL_STAGES`:
|
||||
```python
|
||||
_TTS_MODEL_STAGES: set[str] = (
|
||||
...
|
||||
| _YOUR_MODEL_TTS_MODEL_STAGES
|
||||
)
|
||||
```
|
||||
|
||||
**Point 3** — model type detection in `_detect_tts_model_type()`:
|
||||
```python
|
||||
if model_stage in _YOUR_MODEL_TTS_MODEL_STAGES:
|
||||
return "your_model"
|
||||
```
|
||||
|
||||
**Point 4** — validation dispatch in `_validate_tts_request()`:
|
||||
```python
|
||||
if self._tts_model_type == "your_model":
|
||||
return self._validate_your_model_request(request)
|
||||
```
|
||||
|
||||
**Point 5** — validation + parameter-builder methods:
|
||||
```python
|
||||
def _validate_your_model_request(self, request) -> str | None:
|
||||
if not request.input or not request.input.strip():
|
||||
return "Input text cannot be empty"
|
||||
return None
|
||||
|
||||
def _build_your_model_params(self, request) -> dict:
|
||||
params = {"text": [request.input]}
|
||||
if request.voice is not None:
|
||||
params["voice"] = [request.voice]
|
||||
return params
|
||||
```
|
||||
Wire `_build_your_model_params` into `_create_tts_request()` alongside the other
|
||||
model-specific param builders.
|
||||
|
||||
> **Two dispatch patterns coexist**: Fish Speech uses a `self._is_fish_speech` boolean
|
||||
> instance attribute checked before `elif self._is_tts`, while all newer models
|
||||
> (CosyVoice3, MOSS-TTS-Nano) use the `_tts_model_type` string returned by
|
||||
> `_detect_tts_model_type()`. For new models, always use the `_tts_model_type` string
|
||||
> pattern — do not add new `_is_*` flags.
|
||||
|
||||
> **Unused variable rule**: only extract fields in `_build_your_model_params` that
|
||||
> are actually forwarded to the model. Unused extractions fail `ruff F841`.
|
||||
> For voice-cloning fields (`ref_audio` → `prompt_audio_path`, `ref_text` →
|
||||
> `prompt_text`), add them to the param builder and verify they reach the model call.
|
||||
|
||||
**Rebase conflict note**: when rebasing onto `main` after another model was merged,
|
||||
`serving_speech.py` will conflict. Resolution: always keep *both* the upstream
|
||||
model's additions and your own — never discard either side.
|
||||
|
||||
2. **Handle model-specific parameters**:
|
||||
- Voice cloning: `ref_audio` encoding and prompt injection
|
||||
- `max_new_tokens` override in sampling params
|
||||
@@ -129,11 +307,29 @@ def build_voice_clone_prompt(ref_audio_path: str, text: str, codec) -> list:
|
||||
|
||||
### Deliverables
|
||||
|
||||
- Updated `serving_speech.py` with model-specific prompt builder
|
||||
- Updated `serving_speech.py` with all 5 integration points (single commit)
|
||||
- Client scripts and server launcher
|
||||
- Gradio demo with streaming and voice cloning UI
|
||||
- E2E online serving test (`tests/e2e/online_serving/test_<model>.py`)
|
||||
- Buildkite CI entry in `.buildkite/test-merge.yml`
|
||||
- Documentation (offline + online serving docs)
|
||||
|
||||
### E2E test pitfalls to avoid
|
||||
|
||||
- **One `OmniServerParams` set per file.** `omni_server` is module-scoped; a second
|
||||
id in the same file forces mid-module teardown/restart and exposes startup
|
||||
races (`APIConnectionError` on the first request post-restart). Split variants
|
||||
into separate files instead.
|
||||
- **No external URL fetches from the server.** CI and some dev hosts can't
|
||||
reach `raw.githubusercontent.com` over TLS. Inline ref audio as
|
||||
`data:audio/wav;base64,...`; the serving layer accepts both URL and data URL.
|
||||
- **Use the harness readiness gate.** The fixture waits for HTTP 200 on
|
||||
`/health`; don't add `time.sleep` in tests. If warmup is incomplete, make
|
||||
`/health` return non-200 until you're actually ready.
|
||||
- **Mark with `@pytest.mark.core_model` + `hardware_test(res={"cuda": "H100"})`**
|
||||
so the test lands in `test-ready.yml` (triggered by the `ready` label) rather
|
||||
than only nightly.
|
||||
|
||||
## Phase 4: Async Chunk (Streaming)
|
||||
|
||||
**Goal**: Enable inter-stage streaming so audio chunks are produced while AR generation continues.
|
||||
@@ -199,35 +395,10 @@ Stage 0 (AR) Stage 1 (Decoder)
|
||||
- Fixed batch size (fall back to eager for other sizes)
|
||||
- No dynamic control flow inside the graph
|
||||
|
||||
### Example: Code Predictor CUDA Graph (Qwen3-TTS)
|
||||
|
||||
```python
|
||||
import torch
|
||||
|
||||
class CodePredictorGraph:
|
||||
"""Captures the 16-step code predictor AR loop as a single CUDA graph."""
|
||||
|
||||
def setup_graph(self, device: torch.device, kv_heads: int = 4, head_dim: int = 64):
|
||||
self.num_steps = 16
|
||||
self.kv_cache = torch.zeros(1, kv_heads, self.num_steps, head_dim, device=device)
|
||||
self.positions = torch.arange(self.num_steps, device=device)
|
||||
self.causal_mask = torch.tril(torch.ones(self.num_steps, self.num_steps, device=device))
|
||||
self.input_buf = torch.zeros(1, 1, kv_heads * head_dim, device=device)
|
||||
self.output_buf = torch.zeros(1, self.num_steps, device=device, dtype=torch.long)
|
||||
# Warm up, then: self.graph = torch.cuda.CUDAGraph(); self.graph.capture(...)
|
||||
|
||||
def run_graph(self, initial_input: torch.Tensor) -> torch.Tensor:
|
||||
self.input_buf.copy_(initial_input)
|
||||
self.graph.replay()
|
||||
return self.output_buf.clone()
|
||||
```
|
||||
|
||||
### Performance Expectations
|
||||
|
||||
Based on Qwen3-TTS code predictor experience:
|
||||
- **3-5x speedup** for the graphed component
|
||||
- Only effective for fixed batch sizes (typically batch_size=1)
|
||||
- Falls back to eager mode for unsupported configurations
|
||||
See [references/cuda-graph-example.md](references/cuda-graph-example.md) for
|
||||
a worked skeleton (Qwen3-TTS code predictor, 16-step AR loop), performance
|
||||
expectations (3–5× on the graphed component for fixed batch_size=1), and the
|
||||
graph-safety constraints you must honor inside the captured region.
|
||||
|
||||
### Deliverables
|
||||
|
||||
@@ -235,10 +406,33 @@ Based on Qwen3-TTS code predictor experience:
|
||||
- Benchmark script comparing eager vs graph performance
|
||||
- Documentation of constraints and fallback behavior
|
||||
|
||||
## Phase 6: Pre-commit and DCO
|
||||
|
||||
**Goal**: Every commit passes `pre-commit` lint and carries a DCO
|
||||
`Signed-off-by` line that matches the author email.
|
||||
|
||||
- Install hooks once: `pre-commit install`.
|
||||
- Run `pre-commit run --files <changed-files>` before every push; accept any
|
||||
auto-fixes, stage, re-commit.
|
||||
- Sign every commit with `git commit -s`. DCO checks that author email and
|
||||
`Signed-off-by` email match — `git config user.email` must match your
|
||||
GitHub account email.
|
||||
|
||||
Common pre-commit failures, recovery commands for missing sign-off, and the
|
||||
full `pre-commit run` invocation for a TTS model:
|
||||
[references/precommit-dco.md](references/precommit-dco.md).
|
||||
|
||||
## Integration Checklist
|
||||
|
||||
Use this checklist when integrating a new TTS model:
|
||||
|
||||
### Cross-Cutting Invariants (verify at end of every phase)
|
||||
- [ ] I1: `forward()` docstring states cumulative vs delta; consolidation path audited end-to-end
|
||||
- [ ] I2: Tests / examples / benchmarks never use `dict.get(a) or dict.get(b)` on tensor values; list form handled
|
||||
- [ ] I3: No `.item()` / `.cpu()` / Python branch on tensor values inside per-step loops
|
||||
- [ ] I4: Offline RTF, browser streaming playback, and concurrent-request smoke test all pass
|
||||
- [ ] I5: Any cross-step cache keyed by `_omni_req_id`; entries freed when the request finishes
|
||||
|
||||
### Phase 1: HF Reference
|
||||
- [ ] Reference model runs and produces correct audio
|
||||
- [ ] Architecture documented (stages, codebooks, tokens, sample rate)
|
||||
@@ -248,17 +442,27 @@ Use this checklist when integrating a new TTS model:
|
||||
- [ ] Model registered in `registry.py`
|
||||
- [ ] Config classes created with `model_type` registration
|
||||
- [ ] Stage 0 (AR) implemented and generates correct tokens
|
||||
- [ ] Stage 1 (Decoder) produces correct audio from tokens
|
||||
- [ ] Stage 1 (Decoder) produces correct audio from tokens — dtype float32 for codec decoder
|
||||
- [ ] Stage 1 `max_num_seqs` ≥ 4 in production config (default 1 causes gaps under concurrency)
|
||||
- [ ] Optional dependency fallbacks handled at `load_weights()` time (torchaudio/soundfile/etc.)
|
||||
- [ ] Streaming: codec codes accumulated across AR steps (not reset per step)
|
||||
- [ ] Streaming: delta audio emitted per chunk, not full re-decoded waveform
|
||||
- [ ] Streaming: all `forward()` return paths emit `model_outputs`
|
||||
- [ ] Streaming: per-request state keyed by request ID (not shared across requests)
|
||||
- [ ] Streaming: codec tensors moved to codec decoder device before decode
|
||||
- [ ] Stage config YAML created
|
||||
- [ ] `end2end.py` produces audio matching reference quality
|
||||
- [ ] README.md written
|
||||
|
||||
### Phase 3: Online Serving
|
||||
- [ ] Model added to `serving_speech.py`
|
||||
- [ ] All 5 `serving_speech.py` integration points added in one commit
|
||||
- [ ] Only extract params in `_build_*_params` that are forwarded to the model call (ruff F841)
|
||||
- [ ] Prompt builder handles text input correctly
|
||||
- [ ] Voice cloning works (if supported)
|
||||
- [ ] All response formats work (wav, mp3, flac, pcm)
|
||||
- [ ] Client scripts and server launcher created
|
||||
- [ ] E2E online serving test written (`tests/e2e/online_serving/test_<model>.py`)
|
||||
- [ ] Buildkite CI entry added to `.buildkite/test-merge.yml`
|
||||
- [ ] Gradio demo working
|
||||
- [ ] Documentation added (offline + online docs, nav, supported models)
|
||||
|
||||
@@ -276,9 +480,25 @@ Use this checklist when integrating a new TTS model:
|
||||
- [ ] Benchmark shows meaningful speedup
|
||||
- [ ] Fallback to eager works for unsupported configs
|
||||
|
||||
### Phase 6: Pre-commit and DCO
|
||||
- [ ] `pre-commit run --files <changed>` passes before every push
|
||||
- [ ] Every commit has `Signed-off-by` matching the author email (`git commit -s`)
|
||||
- [ ] `git config user.email` matches the email registered on your GitHub account
|
||||
- [ ] Details and failure-recovery commands: [references/precommit-dco.md](references/precommit-dco.md)
|
||||
|
||||
## References
|
||||
|
||||
- [TTS audio skill](../vllm-omni-audio-tts/SKILL.md) -- supported models and usage
|
||||
- [Fish Speech integration](../vllm-omni-audio-tts/references/fish-speech.md) -- complete example of Phases 1-3
|
||||
- [Qwen3-TTS reference](../vllm-omni-audio-tts/references/qwen-tts.md) -- complete example of all 5 phases
|
||||
In-skill references (details split out of the main body):
|
||||
|
||||
- [references/single-stage-ar.md](references/single-stage-ar.md) — full `forward()` / generator skeleton for the MOSS-TTS-Nano-style pattern
|
||||
- [references/optional-deps.md](references/optional-deps.md) — torchaudio / torchcodec fallback pattern
|
||||
- [references/cuda-graph-example.md](references/cuda-graph-example.md) — Qwen3-TTS code-predictor CUDA graph skeleton
|
||||
- [references/precommit-dco.md](references/precommit-dco.md) — full pre-commit invocation, failure table, DCO recovery
|
||||
|
||||
Project docs and adjacent skills:
|
||||
|
||||
- [TTS audio skill](../vllm-omni-audio-tts/SKILL.md) — supported models and usage
|
||||
- [Fish Speech integration](../vllm-omni-audio-tts/references/fish-speech.md) — complete example of Phases 1–3
|
||||
- [Qwen3-TTS reference](../vllm-omni-audio-tts/references/qwen-tts.md) — complete example of all 5 phases
|
||||
- [Adding a TTS model (developer guide)](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/model/adding_tts_model.md)
|
||||
- `plan/voxcpm2_native_ar_design.md` — VoxCPM2's vLLM-native AR + side-computation pattern (distinct from the generator-based single-stage described above)
|
||||
@@ -0,0 +1,42 @@
|
||||
# CUDA Graph Example: Qwen3-TTS Code Predictor
|
||||
|
||||
Reference sketch for capturing the 16-step code-predictor AR loop as a single
|
||||
CUDA graph. Adapt the shapes, number of steps, and KV-head layout to your
|
||||
model.
|
||||
|
||||
```python
|
||||
import torch
|
||||
|
||||
class CodePredictorGraph:
|
||||
"""Captures the 16-step code predictor AR loop as a single CUDA graph."""
|
||||
|
||||
def setup_graph(self, device: torch.device, kv_heads: int = 4, head_dim: int = 64):
|
||||
self.num_steps = 16
|
||||
self.kv_cache = torch.zeros(1, kv_heads, self.num_steps, head_dim, device=device)
|
||||
self.positions = torch.arange(self.num_steps, device=device)
|
||||
self.causal_mask = torch.tril(torch.ones(self.num_steps, self.num_steps, device=device))
|
||||
self.input_buf = torch.zeros(1, 1, kv_heads * head_dim, device=device)
|
||||
self.output_buf = torch.zeros(1, self.num_steps, device=device, dtype=torch.long)
|
||||
# Warm up, then: self.graph = torch.cuda.CUDAGraph(); self.graph.capture(...)
|
||||
|
||||
def run_graph(self, initial_input: torch.Tensor) -> torch.Tensor:
|
||||
self.input_buf.copy_(initial_input)
|
||||
self.graph.replay()
|
||||
return self.output_buf.clone()
|
||||
```
|
||||
|
||||
## Performance expectations (Qwen3-TTS code predictor)
|
||||
|
||||
- **3–5× speedup** on the graphed component.
|
||||
- Effective only for fixed batch sizes (typically `batch_size=1`).
|
||||
- Fall back to eager for any shape/config that wasn't captured — do not try
|
||||
to recapture per request.
|
||||
|
||||
## Graph-safety constraints
|
||||
|
||||
- `torch.argmax` instead of `torch.multinomial`.
|
||||
- Fixed batch size.
|
||||
- No Python control flow that branches on tensor values inside the captured
|
||||
region (use `torch.where` / masks).
|
||||
- No `.item()`, `.cpu()`, `.tolist()` — each would break the capture or
|
||||
cause a GPU→CPU sync during replay.
|
||||
@@ -0,0 +1,47 @@
|
||||
# Optional Dependency Handling
|
||||
|
||||
Models that rely on `torchaudio`, `torchcodec`, `soundfile`, or other optional
|
||||
packages must handle the missing-package case at import time, not at call
|
||||
time. Failing to do this causes cryptic errors only on environments without
|
||||
the optional package — after the model is already deployed.
|
||||
|
||||
## Pattern (used in MOSS-TTS-Nano)
|
||||
|
||||
```python
|
||||
def _patch_torchaudio_load() -> None:
|
||||
"""Fallback torchaudio.load/save to soundfile if torchcodec is unavailable."""
|
||||
try:
|
||||
import torchcodec # noqa: F401
|
||||
return # torchcodec present, torchaudio works as-is
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
import soundfile as sf
|
||||
import torchaudio
|
||||
|
||||
def _sf_load(path, **kwargs):
|
||||
data, sr = sf.read(str(path), dtype="float32", always_2d=True)
|
||||
return torch.from_numpy(data).T, sr
|
||||
|
||||
torchaudio.load = _sf_load
|
||||
# patch .save similarly if needed
|
||||
```
|
||||
|
||||
## Rules
|
||||
|
||||
- Mirror the full signature of the replaced function. `torchaudio.load`
|
||||
accepts `frame_offset`, `num_frames`, `normalize`, `channels_first`,
|
||||
`format` — missing any of them causes `TypeError` from calling code.
|
||||
- Catch `except Exception`, not just `ImportError`. `import torchaudio`
|
||||
itself can fail with non-`ImportError` errors on broken installs. Log the
|
||||
exception type and message (`logger.warning("torchaudio probe failed: %s: %s",
|
||||
type(exc).__name__, exc)`) before falling back, so unrelated errors are not
|
||||
silently swallowed.
|
||||
- Call the patch function at the top of `load_weights()` before loading any
|
||||
audio assets. Do not call it at module import time.
|
||||
|
||||
## Reference implementation
|
||||
|
||||
Any in-tree model that patches `torchaudio.load` in its `load_weights()` —
|
||||
e.g. MOSS-TTS-Nano's `modeling_moss_tts_nano.py` once that integration
|
||||
lands.
|
||||
@@ -0,0 +1,54 @@
|
||||
# Pre-commit and DCO
|
||||
|
||||
Every commit must pass `pre-commit` lint and carry a `Signed-off-by` line
|
||||
that matches the commit author email.
|
||||
|
||||
## Pre-commit
|
||||
|
||||
Install hooks once:
|
||||
|
||||
```bash
|
||||
pre-commit install
|
||||
```
|
||||
|
||||
Run before every push on the files you changed:
|
||||
|
||||
```bash
|
||||
pre-commit run --files \
|
||||
vllm_omni/model_executor/models/<model_name>/*.py \
|
||||
vllm_omni/entrypoints/openai/serving_speech.py \
|
||||
vllm_omni/model_executor/models/registry.py \
|
||||
tests/e2e/offline_inference/test_<model_name>.py \
|
||||
tests/e2e/online_serving/test_<model_name>.py
|
||||
```
|
||||
|
||||
When pre-commit **modifies files** (ruff format auto-fix), it exits non-zero
|
||||
but the changes are correct — stage the modified files and re-commit.
|
||||
|
||||
| Failure | Root cause | Fix |
|
||||
|---------|-----------|-----|
|
||||
| `ruff F841` | Variable extracted but never forwarded to model call | Remove the extraction or wire it through |
|
||||
| `ruff E402` | Import added below function definitions | Move to top-level import block |
|
||||
| `ruff format` | Line length, spacing, quote style | Accept auto-fix, stage, re-commit |
|
||||
|
||||
## DCO sign-off
|
||||
|
||||
Every commit must carry `Signed-off-by: Your Name <your@email.com>`. Use
|
||||
`-s`:
|
||||
|
||||
```bash
|
||||
git commit -s -m "feat(<model>): add <Model> TTS support"
|
||||
```
|
||||
|
||||
Or set it permanently: `git config format.signOff true`.
|
||||
|
||||
The DCO check verifies that the commit author email matches the
|
||||
`Signed-off-by` line. Confirm `git config user.email` matches your GitHub
|
||||
account email before committing.
|
||||
|
||||
Fix a missing or mismatched sign-off on the latest commit:
|
||||
|
||||
```bash
|
||||
git commit --amend -s --no-edit
|
||||
git push origin <branch> --force-with-lease
|
||||
```
|
||||
@@ -0,0 +1,108 @@
|
||||
# Single-Stage AR Pattern
|
||||
|
||||
When the upstream model cannot be cleanly split into an AR stage and a separate
|
||||
decoder (e.g. MOSS-TTS-Nano, or any model that bundles AR + codec via an
|
||||
`inference_stream()` generator), run the whole pipeline inside a single AR
|
||||
worker that yields audio chunks per request.
|
||||
|
||||
This is distinct from VoxCPM2's pattern, which also runs in a single stage but
|
||||
uses vLLM's native PagedAttention on the base language model with diffusion /
|
||||
VAE side-computation outside vLLM — see
|
||||
`plan/voxcpm2_native_ar_design.md` for that variant.
|
||||
|
||||
## Implementation
|
||||
|
||||
1. **Single model file** — load both AR LM and codec inside
|
||||
`modeling_<model>.py`.
|
||||
2. **Load weights in `load_weights()`**, not `__init__()` — vLLM initializes
|
||||
distributed state before any CUDA allocations.
|
||||
3. **Stream via a per-request generator** stored in `self._stream_gens`:
|
||||
|
||||
```python
|
||||
class YourModelForCausalLM(nn.Module):
|
||||
def __init__(self, *, vllm_config, prefix=""):
|
||||
super().__init__()
|
||||
self._lm = None # populated in load_weights()
|
||||
self._stream_gens: dict = {} # request_key → generator
|
||||
|
||||
def load_weights(self, weights):
|
||||
# Load self._lm here, after vLLM distributed init
|
||||
...
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids,
|
||||
positions,
|
||||
intermediate_tensors=None,
|
||||
inputs_embeds=None,
|
||||
runtime_additional_information: list[dict] | None = None, # one dict per request
|
||||
**kwargs,
|
||||
) -> OmniOutput:
|
||||
infos = runtime_additional_information or [{}]
|
||||
# Skip dummy/profiling calls
|
||||
if not runtime_additional_information or all(i.get("_is_dummy") for i in infos):
|
||||
self._ar_emit_stop_token = True
|
||||
return OmniOutput(...) # return empty outputs
|
||||
|
||||
outputs, last_flags = [], []
|
||||
for info in infos:
|
||||
request_key = str(info.get("_omni_req_id", "0")) # per-request ID from vLLM
|
||||
if request_key not in self._stream_gens:
|
||||
self._stream_gens[request_key] = self._create_stream_gen(info)
|
||||
try:
|
||||
chunk, is_last = next(self._stream_gens[request_key])
|
||||
except StopIteration:
|
||||
chunk, is_last = torch.zeros(0), True
|
||||
if is_last:
|
||||
del self._stream_gens[request_key]
|
||||
outputs.append(chunk)
|
||||
last_flags.append(is_last)
|
||||
|
||||
self._ar_emit_stop_token = all(last_flags)
|
||||
return OmniOutput(multimodal_outputs={"model_outputs": outputs, ...})
|
||||
|
||||
def _create_stream_gen(self, info: dict):
|
||||
"""Yield (waveform_tensor, is_last) tuples from inference_stream()."""
|
||||
for event in self._lm.inference_stream(...):
|
||||
if event["type"] == "audio":
|
||||
yield event["waveform"], False
|
||||
elif event["type"] == "result":
|
||||
# Fallback: some models emit a single "result" event instead of
|
||||
# incremental "audio" events — handle both paths
|
||||
yield event.get("waveform", torch.zeros(0)), True
|
||||
return
|
||||
yield torch.zeros(0), True
|
||||
|
||||
def compute_logits(self, hidden_states, sampling_metadata):
|
||||
# Emit EOS only after the last chunk so the AR scheduler ends the request
|
||||
...
|
||||
```
|
||||
|
||||
## Key points
|
||||
|
||||
- `runtime_additional_information` is the correct parameter name (not
|
||||
`**kwargs`) — it carries one dict per request in the batch.
|
||||
- The request ID is `info.get("_omni_req_id")` — set by vLLM, not by user code.
|
||||
- Handle both `"audio"` (incremental) and `"result"` (final combined) event
|
||||
types from upstream models.
|
||||
|
||||
## Stage config
|
||||
|
||||
Single stage with `worker_type: ar`, `engine_output_type: audio`,
|
||||
`final_output: true`, `is_comprehension: true`, and `async_chunk: false` at
|
||||
the top level. Omitting any of these causes silent misclassification in the
|
||||
serving layer.
|
||||
|
||||
## Lint discipline
|
||||
|
||||
Only extract variables from `additional_information` that you actually
|
||||
forward to the model call — unused extractions trip `ruff F841` in
|
||||
pre-commit.
|
||||
|
||||
## Reference implementation
|
||||
|
||||
Look for any single-stage AR model under
|
||||
`vllm_omni/model_executor/models/` — e.g. `moss_tts_nano/` when its
|
||||
integration lands. If none is in tree yet, follow the skeleton above and
|
||||
cross-check against the `is_comprehension: true` / `async_chunk: false`
|
||||
dispatch in `vllm_omni/entrypoints/openai/serving_speech.py`.
|
||||
@@ -1,629 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
PR Reviewer using GLM API for vllm-omni project.
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, TypedDict
|
||||
|
||||
import requests
|
||||
|
||||
|
||||
# Type definitions for API responses
|
||||
class PRDetails(TypedDict):
|
||||
"""Type definition for GitHub PR details response."""
|
||||
|
||||
title: str
|
||||
body: str
|
||||
number: int
|
||||
state: str
|
||||
user: dict[str, Any]
|
||||
|
||||
|
||||
class GLMMessage(TypedDict):
|
||||
"""Type definition for GLM API message."""
|
||||
|
||||
role: str
|
||||
content: str
|
||||
|
||||
|
||||
class GLMChoice(TypedDict):
|
||||
"""Type definition for GLM API choice."""
|
||||
|
||||
message: GLMMessage
|
||||
finish_reason: str
|
||||
|
||||
|
||||
class GLMResponse(TypedDict):
|
||||
"""Type definition for GLM API response."""
|
||||
|
||||
choices: list[GLMChoice]
|
||||
usage: dict[str, int] | None
|
||||
|
||||
|
||||
class GitHubComment(TypedDict):
|
||||
"""Type definition for GitHub comment."""
|
||||
|
||||
id: int
|
||||
body: str
|
||||
created_at: str
|
||||
user: dict[str, Any]
|
||||
|
||||
|
||||
# Configuration
|
||||
TRIGGER_PHRASE: str = "@vllm-omni-reviewer"
|
||||
DEFAULT_GLM_API_URL: str = "https://open.bigmodel.cn/api/paas/v4/chat/completions" # noqa: E501
|
||||
DEFAULT_GLM_MODEL: str = "glm-5"
|
||||
DEFAULT_COOLDOWN_MINUTES: int = 5
|
||||
DEFAULT_MAX_RETRIES: int = 3
|
||||
DEFAULT_RETRY_DELAY: float = 1.0
|
||||
MAX_DIFF_SIZE: int = 100_000 # Maximum diff size in characters
|
||||
|
||||
|
||||
@dataclass
|
||||
class Config:
|
||||
"""Configuration for the PR reviewer."""
|
||||
|
||||
glm_api_url: str
|
||||
glm_model: str
|
||||
cooldown_minutes: int
|
||||
max_retries: int
|
||||
retry_delay: float
|
||||
max_diff_size: int
|
||||
|
||||
|
||||
# Setup logging
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="[PR Reviewer] %(asctime)s - %(levelname)s - %(message)s",
|
||||
datefmt="%Y-%m-%d %H:%M:%S",
|
||||
)
|
||||
logger: logging.Logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def get_config() -> Config:
|
||||
"""Load configuration from environment variables with defaults."""
|
||||
return Config(
|
||||
glm_api_url=os.getenv("GLM_API_URL", DEFAULT_GLM_API_URL),
|
||||
glm_model=os.getenv("GLM_MODEL", DEFAULT_GLM_MODEL),
|
||||
cooldown_minutes=int(
|
||||
os.getenv(
|
||||
"PR_REVIEWER_COOLDOWN_MINUTES",
|
||||
str(DEFAULT_COOLDOWN_MINUTES),
|
||||
)
|
||||
),
|
||||
max_retries=int(
|
||||
os.getenv(
|
||||
"PR_REVIEWER_MAX_RETRIES",
|
||||
str(DEFAULT_MAX_RETRIES),
|
||||
)
|
||||
),
|
||||
retry_delay=float(os.getenv("PR_REVIEWER_RETRY_DELAY", str(DEFAULT_RETRY_DELAY))),
|
||||
max_diff_size=int(os.getenv("PR_REVIEWER_MAX_DIFF_SIZE", str(MAX_DIFF_SIZE))), # noqa: E501
|
||||
)
|
||||
|
||||
|
||||
def get_env_var(name: str) -> str:
|
||||
"""
|
||||
Get an environment variable or raise an error.
|
||||
|
||||
Args:
|
||||
name: Name of the environment variable.
|
||||
|
||||
Returns:
|
||||
The value of the environment variable.
|
||||
|
||||
Raises:
|
||||
SystemExit: If the environment variable is not set.
|
||||
"""
|
||||
value = os.environ.get(name)
|
||||
if not value:
|
||||
logger.error(f"Environment variable {name} is not set")
|
||||
sys.exit(1)
|
||||
return value
|
||||
|
||||
|
||||
def check_trigger(comment_body: str) -> bool:
|
||||
"""
|
||||
Check if the comment contains the trigger phrase.
|
||||
|
||||
Args:
|
||||
comment_body: The body of the comment to check.
|
||||
|
||||
Returns:
|
||||
True if the trigger phrase is found, False otherwise.
|
||||
"""
|
||||
return TRIGGER_PHRASE in comment_body
|
||||
|
||||
|
||||
def fetch_pr_diff(
|
||||
repo_name: str,
|
||||
pr_number: int,
|
||||
token: str,
|
||||
max_size: int = MAX_DIFF_SIZE,
|
||||
) -> str | None:
|
||||
"""
|
||||
Fetch the diff for a pull request.
|
||||
|
||||
Args:
|
||||
repo_name: The repository name in format "owner/repo".
|
||||
pr_number: The pull request number.
|
||||
token: GitHub authentication token.
|
||||
max_size: Maximum diff size in characters.
|
||||
|
||||
Returns:
|
||||
The diff content as a string, or None if fetching failed.
|
||||
Returns empty string if diff is larger than max_size.
|
||||
"""
|
||||
url: str = f"https://api.github.com/repos/{repo_name}/pulls/{pr_number}"
|
||||
headers: dict[str, str] = {
|
||||
"Authorization": f"Bearer {token}",
|
||||
"Accept": "application/vnd.github.v3.diff",
|
||||
}
|
||||
|
||||
logger.info(f"Fetching PR diff from {url}")
|
||||
response = requests.get(url, headers=headers, timeout=30)
|
||||
|
||||
if response.status_code == 200:
|
||||
diff: str = response.text
|
||||
if len(diff) > max_size:
|
||||
logger.warning(
|
||||
f"Diff size ({len(diff)} bytes) exceeds maximum "
|
||||
f"({max_size} bytes), truncating to first "
|
||||
f"{max_size} bytes"
|
||||
)
|
||||
return diff[:max_size] + "\n\n... [Diff truncated due to size] ..."
|
||||
logger.info(f"Successfully fetched diff ({len(diff)} bytes)")
|
||||
return diff
|
||||
else:
|
||||
logger.error(f"Failed to fetch PR diff: {response.status_code}")
|
||||
logger.error(f"Response: {response.text}")
|
||||
return None
|
||||
|
||||
|
||||
def fetch_pr_details(
|
||||
repo_name: str,
|
||||
pr_number: int,
|
||||
token: str,
|
||||
) -> PRDetails | None:
|
||||
"""
|
||||
Fetch PR details including title and description.
|
||||
|
||||
Args:
|
||||
repo_name: The repository name in format "owner/repo".
|
||||
pr_number: The pull request number.
|
||||
token: GitHub authentication token.
|
||||
|
||||
Returns:
|
||||
A dictionary containing PR details, or None if fetching failed.
|
||||
"""
|
||||
url: str = f"https://api.github.com/repos/{repo_name}/pulls/{pr_number}"
|
||||
headers: dict[str, str] = {
|
||||
"Authorization": f"Bearer {token}",
|
||||
"Accept": "application/vnd.github.v3+json",
|
||||
}
|
||||
|
||||
logger.info(f"Fetching PR details from {url}")
|
||||
response = requests.get(url, headers=headers, timeout=30)
|
||||
|
||||
if response.status_code == 200:
|
||||
return response.json()
|
||||
else:
|
||||
logger.error(f"Failed to fetch PR details: {response.status_code}")
|
||||
return None
|
||||
|
||||
|
||||
def build_review_prompt(pr_title: str, pr_description: str, diff: str) -> str:
|
||||
"""
|
||||
Build the prompt for the GLM-4.7 API.
|
||||
|
||||
Args:
|
||||
pr_title: The title of the pull request.
|
||||
pr_description: The description/body of the pull request.
|
||||
diff: The diff content of the pull request.
|
||||
|
||||
Returns:
|
||||
The formatted prompt string for the API.
|
||||
"""
|
||||
return f"""You are an expert code reviewer for the VLLM-Omni project. \
|
||||
Please review the following pull request:
|
||||
|
||||
## Pull Request Details
|
||||
**Title:** {pr_title}
|
||||
|
||||
**Description:**
|
||||
{pr_description if pr_description else "No description provided."}
|
||||
|
||||
## Code Changes (Diff)
|
||||
{diff}
|
||||
|
||||
## Review Guidelines
|
||||
|
||||
Please provide a comprehensive code review with the following sections:
|
||||
|
||||
### 1. Overview
|
||||
- Brief summary of the changes
|
||||
- Overall assessment (positive, neutral, or concerns)
|
||||
|
||||
### 2. Code Quality
|
||||
- Code style and consistency
|
||||
- Potential bugs or edge cases
|
||||
- Performance considerations
|
||||
- Error handling
|
||||
|
||||
### 3. Architecture & Design
|
||||
- Integration with existing codebase
|
||||
- Design patterns and best practices
|
||||
- Potential improvements
|
||||
|
||||
### 4. Security & Safety
|
||||
- Security concerns (if any)
|
||||
- Resource management
|
||||
- Input validation
|
||||
|
||||
### 5. Testing & Documentation
|
||||
- Test coverage considerations
|
||||
- Documentation completeness
|
||||
- Examples and usage clarity
|
||||
|
||||
### 6. Specific Suggestions
|
||||
- Line-by-line specific feedback (use `file:line` format)
|
||||
- Concrete actionable suggestions
|
||||
- Code examples for improvements (if applicable)
|
||||
|
||||
### 7. Approval Status
|
||||
- **LGTM** (Looks Good To Me) if the PR is ready to merge
|
||||
- **LGTM with suggestions** if the PR is good but has minor suggestions
|
||||
- **Changes requested** if significant changes are needed
|
||||
|
||||
## Important Notes
|
||||
- Be constructive and helpful
|
||||
- Focus on objective technical feedback
|
||||
- Acknowledge good practices when you see them
|
||||
- Prioritize critical issues over nitpicks
|
||||
- If the diff is empty or minimal, acknowledge this and provide
|
||||
any relevant context-specific guidance
|
||||
|
||||
Please format your response in Markdown with clear section headers.
|
||||
"""
|
||||
|
||||
|
||||
def validate_glm_response(data: dict[str, Any]) -> str | None:
|
||||
"""
|
||||
Validate and extract content from GLM API response.
|
||||
|
||||
Args:
|
||||
data: The response data from GLM API.
|
||||
|
||||
Returns:
|
||||
The review content string if valid, None otherwise.
|
||||
"""
|
||||
# Check if choices exists and is a non-empty list
|
||||
if "choices" not in data:
|
||||
logger.error("GLM API response missing 'choices' field")
|
||||
logger.error(f"Response structure: {json.dumps(data, indent=2)}")
|
||||
return None
|
||||
|
||||
choices = data["choices"]
|
||||
if not isinstance(choices, list):
|
||||
logger.error(f"GLM API 'choices' is not a list: {type(choices)}")
|
||||
return None
|
||||
|
||||
if len(choices) == 0:
|
||||
logger.error("GLM API 'choices' is an empty list")
|
||||
return None
|
||||
|
||||
# Check if first choice has message
|
||||
try:
|
||||
first_choice = choices[0]
|
||||
if not isinstance(first_choice, dict):
|
||||
logger.error(f"GLM API choice is not a dict: {type(first_choice)}")
|
||||
return None
|
||||
|
||||
if "message" not in first_choice:
|
||||
logger.error("GLM API choice missing 'message' field")
|
||||
logger.error(f"Choice structure: {json.dumps(first_choice, indent=2)}") # noqa: E501
|
||||
return None
|
||||
|
||||
message = first_choice["message"]
|
||||
if not isinstance(message, dict):
|
||||
logger.error(f"GLM API message is not a dict: {type(message)}")
|
||||
return None
|
||||
|
||||
if "content" not in message:
|
||||
logger.error("GLM API message missing 'content' field")
|
||||
logger.error(f"Message structure: {json.dumps(message, indent=2)}")
|
||||
return None
|
||||
|
||||
content = message["content"]
|
||||
if not isinstance(content, str):
|
||||
logger.error(f"GLM API content is not a string: {type(content)}")
|
||||
return None
|
||||
|
||||
return content
|
||||
|
||||
except (KeyError, IndexError, TypeError) as e:
|
||||
logger.error(f"Failed to parse GLM API response: {e}")
|
||||
logger.error(f"Response: {json.dumps(data, indent=2)}")
|
||||
return None
|
||||
|
||||
|
||||
def call_glm_api(prompt: str, api_key: str, config: Config) -> str | None:
|
||||
"""
|
||||
Call the GLM-4.7 API to get code review with retry logic.
|
||||
|
||||
Args:
|
||||
prompt: The prompt to send to the API.
|
||||
api_key: The GLM API key.
|
||||
config: Configuration object.
|
||||
|
||||
Returns:
|
||||
The review content as a string, or None if all retries failed.
|
||||
"""
|
||||
headers: dict[str, str] = {
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
payload: dict[str, Any] = {
|
||||
"model": config.glm_model,
|
||||
"messages": [{"role": "user", "content": prompt}],
|
||||
"temperature": 0.3,
|
||||
"max_tokens": 32000,
|
||||
"top_p": 0.9,
|
||||
}
|
||||
|
||||
last_error: str | None = None
|
||||
|
||||
for attempt in range(config.max_retries):
|
||||
try:
|
||||
logger.info(f"Calling GLM API ({config.glm_model}) - Attempt {attempt + 1}/{config.max_retries}")
|
||||
response = requests.post(
|
||||
config.glm_api_url,
|
||||
headers=headers,
|
||||
json=payload,
|
||||
timeout=120,
|
||||
)
|
||||
|
||||
if response.status_code == 200:
|
||||
data = response.json()
|
||||
review = validate_glm_response(data)
|
||||
if review:
|
||||
logger.info(f"Successfully received review ({len(review)} chars)") # noqa: E501
|
||||
return review
|
||||
else:
|
||||
last_error = "Failed to validate API response structure"
|
||||
logger.error(last_error)
|
||||
else:
|
||||
last_error = f"GLM API request failed: {response.status_code} - {response.text}"
|
||||
logger.error(last_error)
|
||||
|
||||
except requests.exceptions.Timeout:
|
||||
last_error = f"GLM API request timed out (attempt {attempt + 1})"
|
||||
logger.error(last_error)
|
||||
except requests.exceptions.RequestException as e:
|
||||
last_error = f"GLM API request exception: {e}"
|
||||
logger.error(last_error)
|
||||
except json.JSONDecodeError as e:
|
||||
last_error = f"Failed to decode GLM API response as JSON: {e}"
|
||||
logger.error(last_error)
|
||||
|
||||
# Exponential backoff before retry
|
||||
if attempt < config.max_retries - 1:
|
||||
wait_time: float = config.retry_delay * (2**attempt)
|
||||
logger.info(f"Waiting {wait_time}s before retry...") # noqa: E501
|
||||
time.sleep(wait_time)
|
||||
|
||||
logger.error(
|
||||
f"All {config.max_retries} attempts failed. Last error: {last_error}" # noqa: E501
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def check_cooldown( # noqa: E501
|
||||
repo_name: str,
|
||||
pr_number: int,
|
||||
token: str,
|
||||
cooldown_minutes: int,
|
||||
) -> bool:
|
||||
"""
|
||||
Check if the PR is within the cooldown period.
|
||||
|
||||
Args:
|
||||
repo_name: The repository name in format "owner/repo".
|
||||
pr_number: The pull request number.
|
||||
token: GitHub authentication token.
|
||||
cooldown_minutes: Cooldown period in minutes.
|
||||
|
||||
Returns:
|
||||
True if within cooldown period (should skip), False otherwise.
|
||||
"""
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
url: str = (
|
||||
f"https://api.github.com/repos/{repo_name}/issues/"
|
||||
f"{pr_number}/comments" # noqa: E501
|
||||
)
|
||||
headers: dict[str, str] = {
|
||||
"Authorization": f"Bearer {token}",
|
||||
"Accept": "application/vnd.github.v3+json",
|
||||
}
|
||||
|
||||
logger.info(f"Checking cooldown period ({cooldown_minutes} minutes)")
|
||||
response = requests.get(url, headers=headers, timeout=30)
|
||||
|
||||
if response.status_code != 200:
|
||||
logger.warning(f"Failed to check cooldown: {response.status_code}, proceeding with review")
|
||||
return False
|
||||
|
||||
comments: list[dict[str, Any]] = response.json()
|
||||
cutoff_time: datetime = datetime.utcnow() - timedelta(minutes=cooldown_minutes) # noqa: E501
|
||||
|
||||
for comment in reversed(comments):
|
||||
# Check if this is a bot comment
|
||||
body: str = comment.get("body", "")
|
||||
if "VLLM-Omni PR Review" in body or "PR Reviewer Bot" in body:
|
||||
created_at_str: str = comment.get("created_at", "")
|
||||
try:
|
||||
# Parse GitHub timestamp format
|
||||
created_at = datetime.fromisoformat(created_at_str.replace("Z", "+00:00"))
|
||||
created_at = created_at.replace(tzinfo=None)
|
||||
if created_at > cutoff_time:
|
||||
logger.info(f"PR is within cooldown period (last review: {created_at_str})")
|
||||
return True
|
||||
except ValueError:
|
||||
logger.warning(f"Failed to parse comment timestamp: {created_at_str}") # noqa: E501
|
||||
continue
|
||||
|
||||
logger.info("PR is outside cooldown period, proceeding with review")
|
||||
return False
|
||||
|
||||
|
||||
def post_review_comment( # noqa: E501
|
||||
repo_name: str,
|
||||
pr_number: int,
|
||||
token: str,
|
||||
review: str,
|
||||
) -> bool:
|
||||
"""
|
||||
Post the review as a comment on the PR.
|
||||
|
||||
Args:
|
||||
repo_name: The repository name in format "owner/repo".
|
||||
pr_number: The pull request number.
|
||||
token: GitHub authentication token.
|
||||
review: The review content to post.
|
||||
|
||||
Returns:
|
||||
True if posting succeeded, False otherwise.
|
||||
"""
|
||||
url: str = (
|
||||
f"https://api.github.com/repos/{repo_name}/issues/"
|
||||
f"{pr_number}/comments" # noqa: E501
|
||||
)
|
||||
headers: dict[str, str] = {
|
||||
"Authorization": f"Bearer {token}",
|
||||
"Accept": "application/vnd.github.v3+json",
|
||||
}
|
||||
|
||||
# Format the review comment
|
||||
comment_body: str = f"""## 🤖 VLLM-Omni PR Review
|
||||
|
||||
{review}
|
||||
|
||||
---
|
||||
*This review was generated automatically by the VLLM-Omni PR Reviewer Bot
|
||||
using {os.getenv("GLM_MODEL", DEFAULT_GLM_MODEL)}.*
|
||||
"""
|
||||
|
||||
payload: dict[str, str] = {"body": comment_body}
|
||||
|
||||
logger.info(f"Posting review comment to PR #{pr_number}")
|
||||
response = requests.post(url, headers=headers, json=payload, timeout=30)
|
||||
|
||||
if response.status_code == 201:
|
||||
logger.info("Successfully posted review comment")
|
||||
return True
|
||||
else:
|
||||
logger.error(f"Failed to post comment: {response.status_code}")
|
||||
logger.error(f"Response: {response.text}")
|
||||
return False
|
||||
|
||||
|
||||
def main() -> int:
|
||||
"""
|
||||
Main entry point for the PR reviewer bot.
|
||||
|
||||
Returns:
|
||||
0 on success, 1 on error.
|
||||
"""
|
||||
logger.info("VLLM-Omni PR Reviewer Bot starting...")
|
||||
|
||||
# Load configuration
|
||||
config: Config = get_config()
|
||||
logger.info(
|
||||
f"Configuration: model={config.glm_model}, "
|
||||
f"cooldown={config.cooldown_minutes}min, "
|
||||
f"max_retries={config.max_retries}"
|
||||
)
|
||||
|
||||
# Get environment variables
|
||||
token: str = get_env_var("GITHUB_TOKEN")
|
||||
api_key: str = get_env_var("GLM_API_KEY")
|
||||
repo_name: str = get_env_var("REPO_NAME")
|
||||
pr_number_str: str = get_env_var("PR_NUMBER")
|
||||
comment_body: str = get_env_var("COMMENT_BODY")
|
||||
|
||||
try:
|
||||
pr_number: int = int(pr_number_str)
|
||||
except ValueError:
|
||||
logger.error(f"Invalid PR number: {pr_number_str}")
|
||||
return 1
|
||||
|
||||
logger.info(f"Repository: {repo_name}")
|
||||
logger.info(f"PR Number: {pr_number}")
|
||||
|
||||
# Check if the comment contains the trigger phrase
|
||||
if not check_trigger(comment_body):
|
||||
logger.info(
|
||||
f"Comment does not contain trigger phrase '{TRIGGER_PHRASE}', exiting" # noqa: E501
|
||||
)
|
||||
return 0
|
||||
|
||||
logger.info("Trigger phrase detected! Starting review process...")
|
||||
|
||||
# Check cooldown period
|
||||
if check_cooldown(repo_name, pr_number, token, config.cooldown_minutes):
|
||||
logger.info("Skipping review due to cooldown period")
|
||||
return 0
|
||||
|
||||
# Fetch PR details
|
||||
logger.info("Step 1/4: Fetching PR details...")
|
||||
pr_details: PRDetails | None = fetch_pr_details(repo_name, pr_number, token) # noqa: E501
|
||||
if not pr_details:
|
||||
logger.error("Failed to fetch PR details")
|
||||
return 1
|
||||
|
||||
pr_title: str = pr_details.get("title", "Unknown")
|
||||
pr_description: str = pr_details.get("body", "")
|
||||
|
||||
logger.info(f"PR Title: {pr_title}")
|
||||
|
||||
# Fetch PR diff
|
||||
logger.info("Step 2/4: Fetching PR diff...")
|
||||
diff: str | None = fetch_pr_diff(repo_name, pr_number, token, config.max_diff_size)
|
||||
if diff is None:
|
||||
logger.error("Failed to fetch PR diff")
|
||||
return 1
|
||||
|
||||
if not diff:
|
||||
logger.warning("Warning: Empty diff - this might be a draft PR or no code changes")
|
||||
|
||||
# Build prompt
|
||||
logger.info("Step 3/4: Building review prompt...")
|
||||
prompt: str = build_review_prompt(pr_title, pr_description, diff)
|
||||
|
||||
# Call GLM API
|
||||
logger.info("Step 4/4: Calling GLM API...")
|
||||
review: str | None = call_glm_api(prompt, api_key, config)
|
||||
if not review:
|
||||
logger.error("Failed to get review from GLM API")
|
||||
return 1
|
||||
|
||||
# Post review comment
|
||||
logger.info("Posting review comment...")
|
||||
if not post_review_comment(repo_name, pr_number, token, review):
|
||||
logger.error("Failed to post review comment")
|
||||
return 1
|
||||
|
||||
logger.info("PR review completed successfully!")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -1,62 +0,0 @@
|
||||
name: VLLM-Omni PR Reviewer
|
||||
|
||||
on:
|
||||
issue_comment:
|
||||
types: [created]
|
||||
|
||||
permissions:
|
||||
contents: read
|
||||
pull-requests: write
|
||||
issues: write
|
||||
|
||||
jobs:
|
||||
pr-reviewer:
|
||||
name: Review Pull Request
|
||||
runs-on: ubuntu-latest
|
||||
timeout-minutes: 10
|
||||
# Only run when the comment is from a collaborator/owner/member
|
||||
if: |
|
||||
github.event_name == 'issue_comment' &&
|
||||
github.event.issue.pull_request != null &&
|
||||
contains(github.event.comment.body, '@vllm-omni-reviewer') &&
|
||||
(github.event.comment.author_association == 'MEMBER' ||
|
||||
github.event.comment.author_association == 'COLLABORATOR' ||
|
||||
github.event.comment.author_association == 'OWNER')
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v6.0.2
|
||||
with:
|
||||
fetch-depth: 0
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v6.2.0
|
||||
with:
|
||||
python-version: '3.11'
|
||||
cache: 'pip'
|
||||
|
||||
- name: Install dependencies
|
||||
run: |
|
||||
python -m pip install --upgrade pip==24.0
|
||||
pip install requests==2.31.0 pyyaml==6.0.1
|
||||
|
||||
- name: Run PR Reviewer
|
||||
id: reviewer
|
||||
env:
|
||||
GLM_API_KEY: ${{ secrets.GLM_API_KEY }}
|
||||
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
PR_NUMBER: ${{ github.event.issue.number || github.event.pull_request.number }}
|
||||
COMMENT_BODY: ${{ github.event.comment.body }}
|
||||
REPO_NAME: ${{ github.repository }}
|
||||
PR_REVIEWER_COOLDOWN_MINUTES: 5
|
||||
PR_REVIEWER_MAX_RETRIES: 3
|
||||
run: |
|
||||
python .github/scripts/pr_reviewer.py 2>&1 | tee "/tmp/pr_review_${PR_NUMBER}.log"
|
||||
|
||||
- name: Upload review logs
|
||||
if: always()
|
||||
uses: actions/upload-artifact@v7.0.0
|
||||
with:
|
||||
name: pr-review-logs-${{ github.event.issue.number || github.event.pull_request.number }}
|
||||
path: /tmp/pr_review_*.log
|
||||
retention-days: 7
|
||||
if-no-files-found: ignore
|
||||
@@ -264,5 +264,6 @@ tmp_test
|
||||
vllm_omni/_version.py
|
||||
# output files
|
||||
*.wav
|
||||
.worktrees/
|
||||
# CI overlay yamls materialized from tests/utils.py:_CI_OVERLAYS at test time
|
||||
tests/.ci_generated/
|
||||
@@ -23,5 +23,5 @@ Test guidance:
|
||||
- Local static/self-checks live in `tests/benchmarks/test_accuracy_bench_utils.py`.
|
||||
- End-to-end generation/evaluation should be validated in a remote GPU
|
||||
environment. In the current repo marker system there is `L4` but no `L5`
|
||||
marker, so benchmark smoke tests should be wired as `advanced_model +
|
||||
benchmark + L4` when GPU capacity is available.
|
||||
marker, so benchmark smoke tests should be wired as `full_model +
|
||||
benchmark + L4` for nightly when GPU capacity is available.
|
||||
@@ -99,5 +99,5 @@ Notes:
|
||||
- This flow requires the optional Hugging Face `datasets` package.
|
||||
- `generate` writes `generation_manifest.json` with local output coverage.
|
||||
- The current repo marker set exposes `L4` but not `L5`, so if you promote an
|
||||
end-to-end smoke test into CI, use the existing `advanced_model`, `benchmark`,
|
||||
and `L4` markers or introduce a new repo-wide marker explicitly first.
|
||||
end-to-end smoke test into CI, use the `full_model`, `benchmark`,
|
||||
and `L4` markers for nightly (or `advanced_model` for merge) or introduce a new repo-wide marker explicitly first.
|
||||
@@ -0,0 +1,20 @@
|
||||
vd001|||The quick brown fox jumps over the lazy dog.|A warm, friendly female voice with a slight American Midwest accent, speaking at a moderate pace with natural inflection.
|
||||
vd002|||Welcome to the future of text-to-speech synthesis.|A deep, authoritative male news anchor voice, clear and professional with a measured cadence.
|
||||
vd003|||The sunset painted the sky in brilliant shades of orange and pink.|A gentle elderly female voice, soft and wise, with a slight Southern American accent.
|
||||
vd004|||Scientists have discovered a new species of deep-sea creature.|A young male voice with an Australian accent, curious and enthusiastic.
|
||||
vd005|||Breaking news: a major climate summit opens today in Geneva.|A crisp female newsreader voice, neutral accent, confident and precise.
|
||||
vd006|||In the beginning, there was darkness and silence across the universe.|A rich, dramatic bass male narrator voice, slow and deeply resonant.
|
||||
vd007|||Come closer, I have something important to tell you.|A soft, intimate female voice, slightly whispery, warm and gentle.
|
||||
vd008|||And they're off! The horses race toward the first turn at incredible speed.|An energetic male sports commentator, fast-paced and excited.
|
||||
vd009|||Once upon a time, in a land far away, lived a very clever fox.|A light, playful voice with childlike enthusiasm, bright and clear.
|
||||
vd010|||The ancient manuscript reveals secrets hidden for a thousand years.|A wise, measured elderly male voice, slow and deliberate, British English accent.
|
||||
vd011|||Good evening, ladies and gentlemen, and welcome to our show.|A sophisticated female voice with a slight French accent speaking English, elegant and refined.
|
||||
vd012|||System initialized. Running diagnostics. All systems nominal.|A clear, precise robotic-sounding voice, neutral and monotone with slight synthetic quality.
|
||||
vd013|||I hear what you are saying, and it is completely understandable to feel that way.|A warm, empathetic female therapist voice, calm and reassuring, unhurried pace.
|
||||
vd014|||Attention all units: proceed to grid reference seven-seven-alpha.|A firm, authoritative military male voice, clipped and commanding.
|
||||
vd015|||Oh my goodness, you have to try this amazing new recipe I just found!|An enthusiastic, bubbly female voice, high energy and friendly.
|
||||
vd016|||Dude, the waves were totally amazing out there today. Super happy about it!|A relaxed male voice with a California accent, casual and laid-back.
|
||||
vd017|||The quarterly results exceed expectations across all major metrics.|A sharp, businesslike female voice, confident and efficient, fast-paced delivery.
|
||||
vd018|||Chapter one. The morning sun filtered gently through the forest canopy.|A smooth, rich male audiobook narrator voice, expressive and engaging.
|
||||
vd019|||To be or not to be, that is the question.|A theatrical female voice, dramatic and expressive, stage projection quality.
|
||||
vd020|||And that is all for tonight. Stay well out there, everyone.|A warm, velvety male late-night radio DJ voice, smooth and intimate.
|
||||
@@ -0,0 +1,20 @@
|
||||
smoke001|||The quick brown fox jumps over the lazy dog near the riverbank at sunset.
|
||||
smoke002|||Welcome to the future of text-to-speech synthesis in production systems.
|
||||
smoke003|||Yesterday the team finished rolling out the new authentication flow.
|
||||
smoke004|||She walked carefully across the wet cobblestones, careful not to slip.
|
||||
smoke005|||The conference call is scheduled for nine in the morning, Pacific time.
|
||||
smoke006|||Please remember to save your work before closing the editor.
|
||||
smoke007|||Two plus two equals four, but five hundred and forty three digits is long.
|
||||
smoke008|||I would like a coffee with oat milk and a chocolate croissant please.
|
||||
smoke009|||The library closes at eight on weekdays and six on Saturdays.
|
||||
smoke010|||During the Renaissance, art and science flourished in European cities.
|
||||
smoke011|||He whispered the secret word so quietly that no one else could hear.
|
||||
smoke012|||Our flight departs from gate twenty three at eleven fifteen.
|
||||
smoke013|||The storm knocked out power for six hours, but the backup generator kicked in.
|
||||
smoke014|||Reading a good book on a rainy afternoon is one of life's great pleasures.
|
||||
smoke015|||When the kettle whistled, she poured the hot water over the fresh tea leaves.
|
||||
smoke016|||The algorithm runs in linear time, which is a big improvement over the previous approach.
|
||||
smoke017|||In the distance, the mountains were shrouded in thick morning fog.
|
||||
smoke018|||Our company reported record revenue for the fourth quarter of the fiscal year.
|
||||
smoke019|||She explained the new policy in detail during the staff meeting this morning.
|
||||
smoke020|||The children laughed and played in the garden until the sun began to set.
|
||||
@@ -122,6 +122,18 @@ async def async_request_chat_completions(
|
||||
output.peak_memory_mb = first_item.get("peak_memory_mb", 0.0)
|
||||
except (IndexError, TypeError, AttributeError):
|
||||
pass
|
||||
|
||||
if (not output.stage_durations or output.peak_memory_mb == 0.0) and isinstance(
|
||||
resp_json.get("metrics"), dict
|
||||
):
|
||||
m = resp_json["metrics"]
|
||||
if not output.stage_durations and isinstance(m.get("stage_durations"), dict):
|
||||
output.stage_durations = m.get("stage_durations") or {}
|
||||
if output.peak_memory_mb == 0.0 and m.get("peak_memory_mb") is not None:
|
||||
try:
|
||||
output.peak_memory_mb = float(m.get("peak_memory_mb") or 0.0)
|
||||
except (TypeError, ValueError):
|
||||
pass
|
||||
else:
|
||||
output.error = f"HTTP {response.status}: {await response.text()}"
|
||||
output.success = False
|
||||
|
||||
@@ -12,15 +12,15 @@ Supports multiple backends:
|
||||
- v1/videos: Use /v1/videos endpoint
|
||||
|
||||
Usage:
|
||||
# Video (vllm-omni backend)
|
||||
# Video (v1/videos backend)
|
||||
t2v:
|
||||
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
--backend vllm-omni --dataset vbench --task t2v --num-prompts 10 \
|
||||
--backend v1/videos --dataset vbench --task t2v --num-prompts 10 \
|
||||
--height 480 --width 640 --fps 16 --num-frames 80
|
||||
|
||||
i2v:
|
||||
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
--backend vllm-omni --dataset vbench --task i2v --num-prompts 10
|
||||
--backend v1/videos --dataset vbench --task i2v --num-prompts 10
|
||||
|
||||
|
||||
# Image (vllm-omni backend)
|
||||
@@ -49,7 +49,7 @@ Usage:
|
||||
--backend openai --dataset vbench --task t2i --num-prompts 10 \
|
||||
--height 1024 --width 1024 --port 3000
|
||||
|
||||
# Video (v1/vedeos)
|
||||
# Video (v1/videos)
|
||||
t2v:
|
||||
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
--backend v1/videos --dataset random --task t2v --num-prompts 1 \
|
||||
|
||||
@@ -1,86 +0,0 @@
|
||||
# Benchmarks Guide
|
||||
|
||||
This README explains how to (1) prepare benchmark datasets and (2) run the provided Qwen3-Omni benchmarks.
|
||||
|
||||
## 1) Prepare the dataset (SeedTTS top100)
|
||||
|
||||
```bash
|
||||
cd benchmarks/build_dataset
|
||||
pip install gdown
|
||||
|
||||
# Download SeedTTS test set from Google Drive
|
||||
gdown 1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP
|
||||
|
||||
# Extract
|
||||
tar -xf seedtts_testset.tar
|
||||
|
||||
# Copy metadata and extract top-100 prompts
|
||||
cp seedtts_testset/en/meta.lst meta.lst
|
||||
python extract_prompts.py -i meta.lst -o top100.txt -n 100
|
||||
|
||||
# (Optional) clean up to save space
|
||||
rm -rf seedtts_testset seedtts_testset.tar meta.lst
|
||||
```
|
||||
|
||||
Artifacts:
|
||||
- `benchmarks/build_dataset/top100.txt` — 100 text prompts (one per line).
|
||||
|
||||
## 2) Run benchmarks
|
||||
|
||||
All commands assume repo root (`vllm-omni`).
|
||||
|
||||
### A. Transformers benchmark (offline, HF Transformers)
|
||||
|
||||
```
|
||||
bash benchmarks/qwen3-omni/transformers/eval_qwen3_moe_omni_transformers.sh
|
||||
```
|
||||
|
||||
What it does:
|
||||
- Runs `qwen3_omni_moe_transformers.py` over `top100.txt` with `--num-prompts 100`.
|
||||
- Outputs to `benchmarks/qwen3-omni/transformers/benchmark_results/`:
|
||||
- `perf_stats.json` — aggregated & per-prompt TPS/latency (thinker/talker/code2wav/overall).
|
||||
- `results.json` — per-prompt outputs and audio paths.
|
||||
- `audio/` — ~100 generated `.wav` files.
|
||||
|
||||
Key checks:
|
||||
- `overall_tps` and `*_tps_avg` should be non-zero and reasonably stable.
|
||||
- Investigate any 0/NaN or unusually low TPS / long-tail latency.
|
||||
|
||||
### B. vLLM Omni end-to-end benchmark (pipeline)
|
||||
|
||||
```
|
||||
bash benchmarks/qwen3-omni/vllm_omni/eval_qwen3_moe_omni.sh
|
||||
```
|
||||
|
||||
What it does:
|
||||
- Runs `examples/offline_inference/qwen3_omni/end2end.py` with `--log-stats`.
|
||||
- Uses `benchmarks/build_dataset/top100.txt` and writes to:
|
||||
- Logs: `benchmarks/qwen3-omni/vllm_omni/logs/`
|
||||
- `omni_pipeline_text.orchestrator.stats.jsonl` — per-stage latency stats.
|
||||
- `omni_pipeline_text.overall.stats.jsonl` — end-to-end latency/TPS.
|
||||
- `omni_pipeline_text.stage{0,1,2}.log` — per-stage detailed logs/errors.
|
||||
- Outputs: `benchmarks/qwen3-omni/vllm_omni/outputs/` — ~100 text and `.wav` files.
|
||||
|
||||
Key checks:
|
||||
- Overall stats: end-to-end latency/TPS should be reasonable.
|
||||
- Orchestrator stats: per-stage latency should be stable; investigate long tails.
|
||||
- Stage logs: ensure no errors and no unusually slow stages.
|
||||
|
||||
|
||||
## Performance snapshot
|
||||
|
||||
The chart below summarizes our measured Qwen3-Omni MoE end-to-end benchmark, comparing vLLM-Omni against HF Transformers. It shows the overall throughput advantage for vLLM-Omni. These are actual experiment results—please refer to this performance when evaluating or reproducing the benchmark.
|
||||
|
||||

|
||||
|
||||
## Directory layout
|
||||
- `benchmarks/build_dataset/` — dataset prep utilities (e.g., SeedTTS top100).
|
||||
- `benchmarks/<model>/vllm_omni/` — vLLM-Omni pipeline benchmarks, logs, outputs.
|
||||
- Add new tasks under `benchmarks/<model>/...` with the same pattern: `transformers/`, `vllm_omni/`, task-specific README, and (optionally) dataset prep notes.
|
||||
- `benchmarks/<model>/vllm-omni-vs-hf.png` — current performance snapshot (overall throughput comparison).
|
||||
- `benchmarks/<model>/transformers/` — HF Transformers benchmarks (offline reference).
|
||||
|
||||
## Troubleshooting
|
||||
- Make sure GPU/driver/FlashAttention2 requirements are met for the chosen model.
|
||||
- If downloads fail, confirm network access to Google Drive (`gdown`) and Hugging Face.
|
||||
- If audio files are missing, check for errors in stage logs or model generation.***
|
||||
@@ -1,26 +0,0 @@
|
||||
#!/bin/bash
|
||||
# Qwen3-Omni Transformers Benchmark Evaluation Script
|
||||
# This script must be run from the vllm-omni root directory
|
||||
|
||||
# Get the directory where this script is located
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
|
||||
# Navigate to vllm-omni root directory (4 levels up from script location)
|
||||
VLLM_OMNI_ROOT="$(cd "$SCRIPT_DIR/../../.." && pwd)"
|
||||
cd "$VLLM_OMNI_ROOT" || { echo "Error: Failed to navigate to vllm-omni directory"; exit 1; }
|
||||
|
||||
echo "Working directory: $(pwd)"
|
||||
# Verify we're in the correct directory and run benchmark
|
||||
if [[ ! -f "benchmarks/qwen3-omni/transformers/qwen3_omni_moe_transformers.py" ]]; then
|
||||
echo "Error: Not in vllm-omni root directory. Please run from vllm-omni folder."
|
||||
else
|
||||
cd benchmarks/qwen3-omni/transformers
|
||||
|
||||
python qwen3_omni_moe_transformers.py --prompts-file ../../build_dataset/top100.txt --num-prompts 100
|
||||
|
||||
echo "Logs and outputs are saved to $(pwd)/benchmark_results:"
|
||||
echo " - perf_stats.json Aggregated/per-prompt TPS and latency (thinker/talker/code2wav/overall)"
|
||||
echo " - results.json Per-prompt outputs and audio paths"
|
||||
echo " - audio/ Generated wav files, there should be 100 wav file generated"
|
||||
echo "Key checks: overall_tps and *_tps_avg should be non-zero and stable; investigate 0/NaN or unusually low TPS/long-tail latency."
|
||||
fi
|
||||
@@ -1,265 +0,0 @@
|
||||
import time
|
||||
|
||||
import torch
|
||||
from transformers import Qwen3OmniMoeForConditionalGeneration
|
||||
|
||||
|
||||
class Qwen3OmniMoeForConditionalGenerationWithLogging(Qwen3OmniMoeForConditionalGeneration):
|
||||
@torch.no_grad()
|
||||
def generate(
|
||||
self,
|
||||
input_ids: torch.Tensor | None = None,
|
||||
speaker: str = "Ethan",
|
||||
use_audio_in_video: bool = False,
|
||||
return_audio: bool | None = None,
|
||||
thinker_max_new_tokens: int = 1024,
|
||||
thinker_eos_token_id: int = 151645,
|
||||
talker_max_new_tokens: int = 4096,
|
||||
talker_do_sample: bool = True,
|
||||
talker_top_k: int = 50,
|
||||
talker_top_p: float = 1.0,
|
||||
talker_temperature: float = 0.9,
|
||||
talker_repetition_penalty: float = 1.05,
|
||||
**kwargs,
|
||||
):
|
||||
total_t0 = time.time()
|
||||
perf_stats = {
|
||||
"thinker_tokens": 0,
|
||||
"thinker_time_s": 0.0,
|
||||
"thinker_tps": 0.0,
|
||||
"talker_tokens": 0,
|
||||
"talker_time_s": 0.0,
|
||||
"talker_tps": 0.0,
|
||||
"code2wav_tokens": 0,
|
||||
"code2wav_time_s": 0.0,
|
||||
"code2wav_tps": 0.0,
|
||||
"total_tokens": 0,
|
||||
"total_time_s": 0.0,
|
||||
"total_tps": 0.0,
|
||||
}
|
||||
if return_audio and not self.has_talker:
|
||||
raise ValueError(
|
||||
"Cannot use talker when talker module not initialized. "
|
||||
"Use `enable_talker` method or set enable_talker in config "
|
||||
"to enable talker."
|
||||
)
|
||||
if return_audio is None:
|
||||
return_audio = self.has_talker
|
||||
|
||||
shared_kwargs = {"use_audio_in_video": use_audio_in_video}
|
||||
thinker_kwargs = {
|
||||
"max_new_tokens": thinker_max_new_tokens,
|
||||
"eos_token_id": thinker_eos_token_id,
|
||||
}
|
||||
|
||||
talker_kwargs = {}
|
||||
token2wav_kwargs = {}
|
||||
if return_audio:
|
||||
speaker_id = self.config.talker_config.speaker_id.get(speaker.lower())
|
||||
if speaker_id is None:
|
||||
raise NotImplementedError(f"Speaker {speaker} not implemented")
|
||||
if input_ids.shape[0] != 1:
|
||||
raise NotImplementedError("Qwen3-Omni currently does not support batched inference with audio output")
|
||||
talker_suppressed_tokens = [
|
||||
i
|
||||
for i in range(
|
||||
self.config.talker_config.text_config.vocab_size - 1024,
|
||||
self.config.talker_config.text_config.vocab_size,
|
||||
)
|
||||
if i != self.config.talker_config.codec_eos_token_id
|
||||
] # Suppress additional special tokens, should not be predicted
|
||||
talker_kwargs = {
|
||||
"max_new_tokens": talker_max_new_tokens,
|
||||
"do_sample": talker_do_sample,
|
||||
"top_k": talker_top_k,
|
||||
"top_p": talker_top_p,
|
||||
"temperature": talker_temperature,
|
||||
"eos_token_id": self.config.talker_config.codec_eos_token_id,
|
||||
"repetition_penalty": talker_repetition_penalty,
|
||||
"suppress_tokens": talker_suppressed_tokens,
|
||||
"output_hidden_states": True,
|
||||
"return_dict_in_generate": True,
|
||||
}
|
||||
token2wav_kwargs = {}
|
||||
|
||||
for key, value in kwargs.items():
|
||||
if key.startswith("thinker_"):
|
||||
thinker_kwargs[key[len("thinker_") :]] = value
|
||||
elif key.startswith("talker_"):
|
||||
talker_kwargs[key[len("talker_") :]] = value
|
||||
elif key.startswith("token2wav_"):
|
||||
token2wav_kwargs[key[len("token2wav_") :]] = value
|
||||
# Process special input values
|
||||
elif key == "feature_attention_mask":
|
||||
thinker_kwargs[key] = value
|
||||
talker_kwargs["audio_feature_lengths"] = torch.sum(value, dim=1)
|
||||
elif key in ("input_features", "attention_mask"):
|
||||
thinker_kwargs[key] = value
|
||||
# Put other key to shared kwargs
|
||||
else:
|
||||
shared_kwargs[key] = value
|
||||
|
||||
# Merge kwargs
|
||||
for key, value in shared_kwargs.items():
|
||||
if key not in thinker_kwargs:
|
||||
thinker_kwargs[key] = value
|
||||
if key not in talker_kwargs and key in ["image_grid_thw", "video_grid_thw", "video_second_per_grid"]:
|
||||
talker_kwargs[key] = value
|
||||
if key not in token2wav_kwargs:
|
||||
token2wav_kwargs[key] = value
|
||||
|
||||
# 1. Generate from thinker module
|
||||
generate_audio = return_audio and self.has_talker
|
||||
if generate_audio:
|
||||
thinker_kwargs["output_hidden_states"] = True
|
||||
thinker_kwargs["return_dict_in_generate"] = True
|
||||
|
||||
t0 = time.time()
|
||||
thinker_result = self.thinker.generate(input_ids=input_ids, **thinker_kwargs)
|
||||
t1 = time.time()
|
||||
perf_stats["thinker_time_s"] = max(0.0, t1 - t0)
|
||||
try:
|
||||
prompt_len = int(input_ids.shape[1]) if input_ids is not None else 0
|
||||
total_len = int(thinker_result.sequences.shape[-1])
|
||||
thinker_out_len = max(0, total_len - prompt_len)
|
||||
except Exception:
|
||||
thinker_out_len = 0
|
||||
perf_stats["thinker_tokens"] = thinker_out_len
|
||||
perf_stats["thinker_tps"] = (
|
||||
(thinker_out_len / perf_stats["thinker_time_s"]) if perf_stats["thinker_time_s"] > 0 else 0.0
|
||||
)
|
||||
|
||||
if not generate_audio:
|
||||
perf_stats["total_tokens"] = perf_stats["thinker_tokens"]
|
||||
perf_stats["total_time_s"] = time.time() - total_t0
|
||||
perf_stats["total_tps"] = (
|
||||
(perf_stats["total_tokens"] / perf_stats["total_time_s"]) if perf_stats["total_time_s"] > 0 else 0.0
|
||||
)
|
||||
# attach stats to self
|
||||
setattr(self, "_perf_stats_last", perf_stats)
|
||||
if not hasattr(self, "_perf_stats_history"):
|
||||
setattr(self, "_perf_stats_history", [])
|
||||
self._perf_stats_history.append(perf_stats)
|
||||
return thinker_result, None
|
||||
|
||||
# 2. Prepare talker input
|
||||
thinker_embed = torch.cat([hidden_states[0] for hidden_states in thinker_result.hidden_states], dim=1).to(
|
||||
self.talker.device
|
||||
) # [1 t d]
|
||||
thinker_hidden = torch.cat(
|
||||
[
|
||||
hidden_states[self.config.talker_config.accept_hidden_layer]
|
||||
for hidden_states in thinker_result.hidden_states
|
||||
],
|
||||
dim=1,
|
||||
).to(self.talker.device) # [1 t d]
|
||||
|
||||
im_start_indexes = torch.cat(
|
||||
(
|
||||
torch.nonzero(input_ids[0] == self.config.im_start_token_id).squeeze(),
|
||||
torch.tensor([thinker_result.sequences.shape[-1]], device=input_ids.device, dtype=input_ids.dtype),
|
||||
),
|
||||
dim=-1,
|
||||
).to(self.talker.device) # Shape [n_starts + 1]; Take batch 0 since batched inference is not supported here.
|
||||
multimodal_mask = (
|
||||
(thinker_result.sequences == self.config.thinker_config.audio_token_id) |
|
||||
(thinker_result.sequences == self.config.thinker_config.image_token_id) |
|
||||
(thinker_result.sequences == self.config.thinker_config.video_token_id)
|
||||
).to(self.talker.device) # [1 t] # fmt: skip
|
||||
|
||||
talker_special_tokens = torch.tensor(
|
||||
[[self.config.tts_bos_token_id, self.config.tts_eos_token_id, self.config.tts_pad_token_id]],
|
||||
device=self.thinker.device,
|
||||
dtype=input_ids.dtype,
|
||||
)
|
||||
tts_bos_embed, tts_eos_embed, tts_pad_embed = (
|
||||
self.talker.text_projection(self.thinker.get_input_embeddings()(talker_special_tokens))
|
||||
.to(self.talker.device)
|
||||
.chunk(3, dim=1)
|
||||
) # 3 * [1 1 d]
|
||||
|
||||
talker_input_embeds = [] # [1 t d]
|
||||
talker_input_ids = []
|
||||
# For every chatml parts
|
||||
for i in range(len(im_start_indexes) - 1):
|
||||
im_start_index = im_start_indexes[i]
|
||||
segment_end_index = im_start_indexes[i + 1]
|
||||
role_token = input_ids[0][im_start_index + 1]
|
||||
# Talker should ignore thinker system prompt
|
||||
if role_token == self.config.system_token_id:
|
||||
continue
|
||||
# Talker takes word embeddings for tokens and hidden state from `accept_hidden_layer` for multimodal inputs
|
||||
elif role_token == self.config.user_token_id:
|
||||
talker_user_part = self._get_talker_user_parts(
|
||||
im_start_index, segment_end_index, multimodal_mask, thinker_hidden, thinker_embed
|
||||
)
|
||||
talker_input_embeds.append(talker_user_part)
|
||||
talker_input_ids.append(thinker_result.sequences[:, im_start_index:segment_end_index])
|
||||
# Take assistant output (for now)
|
||||
elif role_token == self.config.assistant_token_id and i == len(im_start_indexes) - 2:
|
||||
talker_assistant_embeds, talker_assistant_ids, trailing_text_hidden = self._get_talker_assistant_parts(
|
||||
im_start_index,
|
||||
segment_end_index,
|
||||
speaker_id,
|
||||
thinker_embed,
|
||||
tts_pad_embed,
|
||||
tts_bos_embed,
|
||||
tts_eos_embed,
|
||||
)
|
||||
talker_input_embeds.append(talker_assistant_embeds)
|
||||
talker_input_ids.append(talker_assistant_ids)
|
||||
# History assistant output (ignore for now)
|
||||
elif role_token == self.config.assistant_token_id and i != len(im_start_indexes) - 2:
|
||||
continue
|
||||
else:
|
||||
raise AssertionError("Expect role id after <|im_start|> (assistant, user, system)")
|
||||
talker_input_embed = torch.cat([embed.to(self.talker.device) for embed in talker_input_embeds], dim=1)
|
||||
talker_input_id = torch.cat([embed.to(self.talker.device) for embed in talker_input_ids], dim=1)
|
||||
t2 = time.time()
|
||||
talker_result = self.talker.generate(
|
||||
inputs_embeds=talker_input_embed,
|
||||
trailing_text_hidden=trailing_text_hidden,
|
||||
tts_pad_embed=tts_pad_embed,
|
||||
talker_input_ids=talker_input_id, # Not use input_ids to prevent repetition penalty out of bound
|
||||
**talker_kwargs,
|
||||
)
|
||||
t3 = time.time()
|
||||
perf_stats["talker_time_s"] = max(0.0, t3 - t2)
|
||||
talker_codes = (
|
||||
torch.stack([hid[-1] for hid in talker_result.hidden_states if hid[-1] is not None], dim=1)
|
||||
.transpose(1, 2)
|
||||
.to(self.code2wav.device)
|
||||
)
|
||||
try:
|
||||
# codes shape: (B, num_quantizers, T). We log T as token length.
|
||||
perf_stats["talker_tokens"] = int(talker_codes.shape[-1])
|
||||
except Exception:
|
||||
perf_stats["talker_tokens"] = 0
|
||||
perf_stats["talker_tps"] = (
|
||||
(perf_stats["talker_tokens"] / perf_stats["talker_time_s"]) if perf_stats["talker_time_s"] > 0 else 0.0
|
||||
)
|
||||
t4 = time.time()
|
||||
talker_wavs = self.code2wav.chunked_decode(talker_codes, chunk_size=300, left_context_size=25).float()
|
||||
t5 = time.time()
|
||||
perf_stats["code2wav_time_s"] = max(0.0, t5 - t4)
|
||||
perf_stats["code2wav_tokens"] = perf_stats["talker_tokens"] # same T, not times 16
|
||||
perf_stats["code2wav_tps"] = (
|
||||
(perf_stats["code2wav_tokens"] / perf_stats["code2wav_time_s"])
|
||||
if perf_stats["code2wav_time_s"] > 0
|
||||
else 0.0
|
||||
)
|
||||
perf_stats["total_tokens"] = perf_stats["thinker_tokens"] + perf_stats["talker_tokens"]
|
||||
perf_stats["total_time_s"] = time.time() - total_t0
|
||||
perf_stats["total_tps"] = (
|
||||
(perf_stats["total_tokens"] / perf_stats["total_time_s"]) if perf_stats["total_time_s"] > 0 else 0.0
|
||||
)
|
||||
setattr(self, "_perf_stats_last", perf_stats)
|
||||
if not hasattr(self, "_perf_stats_history"):
|
||||
setattr(self, "_perf_stats_history", [])
|
||||
self._perf_stats_history.append(perf_stats)
|
||||
return thinker_result, talker_wavs.float()
|
||||
|
||||
|
||||
__all__ = [
|
||||
"Qwen3OmniMoeForConditionalGenerationWithLogging",
|
||||
]
|
||||
@@ -1,275 +0,0 @@
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
|
||||
import soundfile as sf
|
||||
from qwen3_omni_moe_model import Qwen3OmniMoeForConditionalGenerationWithLogging
|
||||
from qwen_omni_utils import process_mm_info
|
||||
from tqdm import tqdm
|
||||
from transformers import Qwen3OmniMoeProcessor
|
||||
|
||||
MODEL_PATH = "Qwen/Qwen3-Omni-30B-A3B-Instruct"
|
||||
# MODEL_PATH = "Qwen/Qwen3-Omni-30B-A3B-Thinking"
|
||||
|
||||
|
||||
def load_prompts(prompts_file: str) -> list[str]:
|
||||
"""Load prompts from a text file, one prompt per line."""
|
||||
prompts = []
|
||||
with open(prompts_file, encoding="utf-8") as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
if line:
|
||||
prompts.append(line)
|
||||
return prompts
|
||||
|
||||
|
||||
def run_benchmark(
|
||||
model,
|
||||
processor,
|
||||
prompts: list[str],
|
||||
output_dir: str = "benchmark_results",
|
||||
speaker: str = "Ethan",
|
||||
use_audio_in_video: bool = True,
|
||||
):
|
||||
"""
|
||||
Run benchmark on a list of prompts and collect performance stats.
|
||||
|
||||
Args:
|
||||
model: The Qwen3OmniMoe model
|
||||
processor: The Qwen3OmniMoe processor
|
||||
prompts: List of text prompts to process
|
||||
output_dir: Directory to save results
|
||||
speaker: Speaker voice for audio output
|
||||
use_audio_in_video: Whether to use audio in video
|
||||
|
||||
Returns:
|
||||
tuple: (aggregated_stats, results, audio_outputs)
|
||||
- aggregated_stats: dict with aggregated performance statistics
|
||||
- results: list of dicts with per-prompt results
|
||||
- audio_outputs: list of audio tensors/arrays (or None if no audio)
|
||||
"""
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
audio_dir = os.path.join(output_dir, "audio")
|
||||
os.makedirs(audio_dir, exist_ok=True)
|
||||
|
||||
all_stats = []
|
||||
results = []
|
||||
audio_outputs = []
|
||||
|
||||
for idx, prompt in enumerate(tqdm(prompts, desc="Processing prompts")):
|
||||
conversation = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "text": prompt}],
|
||||
},
|
||||
]
|
||||
|
||||
# Preparation for inference
|
||||
text = processor.apply_chat_template(conversation, add_generation_prompt=True, tokenize=False)
|
||||
audios, images, videos = process_mm_info(conversation, use_audio_in_video=use_audio_in_video)
|
||||
inputs = processor(
|
||||
text=text,
|
||||
audio=audios,
|
||||
images=images,
|
||||
videos=videos,
|
||||
return_tensors="pt",
|
||||
padding=True,
|
||||
use_audio_in_video=use_audio_in_video,
|
||||
)
|
||||
inputs = inputs.to(model.device).to(model.dtype)
|
||||
|
||||
# Inference: Generation of the output text and audio
|
||||
text_ids, audio = model.generate(
|
||||
**inputs, speaker=speaker, thinker_return_dict_in_generate=True, use_audio_in_video=use_audio_in_video
|
||||
)
|
||||
|
||||
# Decode output text
|
||||
output_text = processor.batch_decode(
|
||||
text_ids.sequences[:, inputs["input_ids"].shape[1] :],
|
||||
skip_special_tokens=True,
|
||||
clean_up_tokenization_spaces=False,
|
||||
)[0]
|
||||
|
||||
# Collect performance stats
|
||||
perf_stats = None
|
||||
if hasattr(model, "_perf_stats_last"):
|
||||
perf_stats = model._perf_stats_last.copy()
|
||||
perf_stats["prompt_idx"] = idx
|
||||
perf_stats["prompt"] = prompt
|
||||
all_stats.append(perf_stats)
|
||||
|
||||
# Save audio and collect audio output
|
||||
audio_path = None
|
||||
audio_data = None
|
||||
if audio is not None:
|
||||
audio_data = audio.reshape(-1).detach().cpu().numpy()
|
||||
audio_path = os.path.join(audio_dir, f"output_{idx:04d}.wav")
|
||||
sf.write(
|
||||
audio_path,
|
||||
audio_data,
|
||||
samplerate=24000,
|
||||
)
|
||||
audio_outputs.append(audio_data)
|
||||
else:
|
||||
audio_outputs.append(None)
|
||||
|
||||
# Save result
|
||||
result = {
|
||||
"idx": idx,
|
||||
"prompt": prompt,
|
||||
"output": output_text,
|
||||
"audio_path": audio_path,
|
||||
"perf_stats": perf_stats,
|
||||
}
|
||||
results.append(result)
|
||||
|
||||
# Aggregate statistics
|
||||
aggregated_stats = aggregate_stats(all_stats)
|
||||
|
||||
# Save all results
|
||||
results_path = os.path.join(output_dir, "results.json")
|
||||
with open(results_path, "w", encoding="utf-8") as f:
|
||||
json.dump(results, f, ensure_ascii=False, indent=2)
|
||||
|
||||
# Save aggregated stats
|
||||
stats_path = os.path.join(output_dir, "perf_stats.json")
|
||||
with open(stats_path, "w", encoding="utf-8") as f:
|
||||
json.dump({"aggregated": aggregated_stats, "per_prompt": all_stats}, f, ensure_ascii=False, indent=2)
|
||||
|
||||
# Count saved audio files
|
||||
num_audio_saved = sum(1 for a in audio_outputs if a is not None)
|
||||
print(f"\nSaved {num_audio_saved} audio files to {audio_dir}/")
|
||||
|
||||
return aggregated_stats, results, audio_outputs
|
||||
|
||||
|
||||
def aggregate_stats(all_stats: list[dict]) -> dict:
|
||||
"""Aggregate performance statistics from multiple runs."""
|
||||
if not all_stats:
|
||||
return {}
|
||||
|
||||
keys = [
|
||||
"thinker_tokens",
|
||||
"thinker_time_s",
|
||||
"thinker_tps",
|
||||
"talker_tokens",
|
||||
"talker_time_s",
|
||||
"talker_tps",
|
||||
"code2wav_tokens",
|
||||
"code2wav_time_s",
|
||||
"code2wav_tps",
|
||||
"total_tokens",
|
||||
"total_time_s",
|
||||
"total_tps",
|
||||
]
|
||||
|
||||
aggregated = {
|
||||
"num_samples": len(all_stats),
|
||||
}
|
||||
|
||||
for key in keys:
|
||||
values = [s.get(key, 0) for s in all_stats if key in s]
|
||||
if values:
|
||||
aggregated[f"{key}_sum"] = sum(values)
|
||||
aggregated[f"{key}_avg"] = sum(values) / len(values)
|
||||
aggregated[f"{key}_min"] = min(values)
|
||||
aggregated[f"{key}_max"] = max(values)
|
||||
|
||||
# Calculate overall throughput
|
||||
total_tokens = aggregated.get("total_tokens_sum", 0)
|
||||
total_time = aggregated.get("total_time_s_sum", 0)
|
||||
if total_time > 0:
|
||||
aggregated["overall_tps"] = total_tokens / total_time
|
||||
|
||||
return aggregated
|
||||
|
||||
|
||||
def print_stats(stats: dict):
|
||||
"""Print performance statistics in a formatted way."""
|
||||
print("\n" + "=" * 60)
|
||||
print("Performance Statistics Summary")
|
||||
print("=" * 60)
|
||||
|
||||
print(f"\nNumber of samples: {stats.get('num_samples', 0)}")
|
||||
|
||||
print("\n--- Thinker ---")
|
||||
print(f" Total tokens: {stats.get('thinker_tokens_sum', 0):.0f}")
|
||||
print(f" Total time: {stats.get('thinker_time_s_sum', 0):.2f}s")
|
||||
print(f" Avg TPS: {stats.get('thinker_tps_avg', 0):.2f}")
|
||||
print(f" Min TPS: {stats.get('thinker_tps_min', 0):.2f}")
|
||||
print(f" Max TPS: {stats.get('thinker_tps_max', 0):.2f}")
|
||||
|
||||
print("\n--- Talker ---")
|
||||
print(f" Total tokens: {stats.get('talker_tokens_sum', 0):.0f}")
|
||||
print(f" Total time: {stats.get('talker_time_s_sum', 0):.2f}s")
|
||||
print(f" Avg TPS: {stats.get('talker_tps_avg', 0):.2f}")
|
||||
print(f" Min TPS: {stats.get('talker_tps_min', 0):.2f}")
|
||||
print(f" Max TPS: {stats.get('talker_tps_max', 0):.2f}")
|
||||
|
||||
print("\n--- Code2Wav ---")
|
||||
print(f" Total tokens: {stats.get('code2wav_tokens_sum', 0):.0f}")
|
||||
print(f" Total time: {stats.get('code2wav_time_s_sum', 0):.2f}s")
|
||||
print(f" Avg TPS: {stats.get('code2wav_tps_avg', 0):.2f}")
|
||||
print(f" Min TPS: {stats.get('code2wav_tps_min', 0):.2f}")
|
||||
print(f" Max TPS: {stats.get('code2wav_tps_max', 0):.2f}")
|
||||
|
||||
print("\n--- Overall ---")
|
||||
print(f" Total tokens: {stats.get('total_tokens_sum', 0):.0f}")
|
||||
print(f" Total time: {stats.get('total_time_s_sum', 0):.2f}s")
|
||||
print(f" Overall TPS: {stats.get('overall_tps', 0):.2f}")
|
||||
print(f" Avg TPS: {stats.get('total_tps_avg', 0):.2f}")
|
||||
print(f" Min TPS: {stats.get('total_tps_min', 0):.2f}")
|
||||
print(f" Max TPS: {stats.get('total_tps_max', 0):.2f}")
|
||||
|
||||
print("=" * 60 + "\n")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Qwen3-Omni Benchmark Script")
|
||||
parser.add_argument(
|
||||
"--prompts-file",
|
||||
type=str,
|
||||
default="benchmark/build_dataset/top100.txt",
|
||||
help="Path to the prompts file (one prompt per line)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-dir", type=str, default="benchmark_results", help="Directory to save benchmark results"
|
||||
)
|
||||
parser.add_argument("--model-path", type=str, default=MODEL_PATH, help="Path to the model")
|
||||
parser.add_argument("--speaker", type=str, default="Ethan", help="Speaker voice for audio output")
|
||||
parser.add_argument("--num-prompts", type=int, default=None, help="Number of prompts to process (default: all)")
|
||||
args = parser.parse_args()
|
||||
|
||||
# Load model and processor
|
||||
print(f"Loading model from {args.model_path}...")
|
||||
model = Qwen3OmniMoeForConditionalGenerationWithLogging.from_pretrained(
|
||||
args.model_path,
|
||||
dtype="auto",
|
||||
device_map="auto",
|
||||
attn_implementation="flash_attention_2",
|
||||
)
|
||||
processor = Qwen3OmniMoeProcessor.from_pretrained(args.model_path)
|
||||
|
||||
# Benchmark mode
|
||||
print(f"Loading prompts from {args.prompts_file}...")
|
||||
prompts = load_prompts(args.prompts_file)
|
||||
|
||||
if args.num_prompts:
|
||||
prompts = prompts[: args.num_prompts]
|
||||
|
||||
print(f"Running benchmark on {len(prompts)} prompts...")
|
||||
|
||||
aggregated_stats, results, audio_outputs = run_benchmark(
|
||||
model=model,
|
||||
processor=processor,
|
||||
prompts=prompts,
|
||||
output_dir=args.output_dir,
|
||||
speaker=args.speaker,
|
||||
)
|
||||
|
||||
print_stats(aggregated_stats)
|
||||
print(f"\nResults saved to {args.output_dir}/")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 115 KiB |
@@ -1,37 +0,0 @@
|
||||
#!/bin/bash
|
||||
# Qwen3-Omni Benchmark Evaluation Script
|
||||
# This script must be run from the vllm-omni root directory
|
||||
|
||||
# Get the directory where this script is located
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
|
||||
# Navigate to vllm-omni root directory (4 levels up from script location)
|
||||
VLLM_OMNI_ROOT="$(cd "$SCRIPT_DIR/../../.." && pwd)"
|
||||
cd "$VLLM_OMNI_ROOT" || { echo "Error: Failed to navigate to vllm-omni directory"; exit 1; }
|
||||
|
||||
echo "Working directory: $(pwd)"
|
||||
|
||||
# Verify we're in the correct directory and run benchmark
|
||||
if [[ ! -d "benchmarks/qwen3-omni/vllm_omni" ]]; then
|
||||
echo "Error: Not in vllm-omni root directory. Please run from vllm-omni folder."
|
||||
else
|
||||
log_dir=benchmarks/qwen3-omni/vllm_omni/logs
|
||||
outputs_dir=benchmarks/qwen3-omni/vllm_omni/outputs
|
||||
end2end_script_path=examples/offline_inference/qwen3_omni/end2end.py
|
||||
build_dataset_path=benchmarks/build_dataset/top100.txt
|
||||
|
||||
python $end2end_script_path --output-wav $outputs_dir \
|
||||
--query-type text \
|
||||
--txt-prompts $build_dataset_path \
|
||||
--log-stats \
|
||||
--log-dir $log_dir
|
||||
echo "Logs and outputs are saved in ${log_dir} and ${outputs_dir} respectively:"
|
||||
echo " - omni_pipeline_text run dir/base name"
|
||||
echo " - omni_pipeline_text.orchestrator.stats.jsonl orchestrator-stage latency stats"
|
||||
echo " - omni_pipeline_text.overall.stats.jsonl overall latency/TPS stats"
|
||||
echo " - omni_pipeline_text.stage0.log per-stage detailed logs"
|
||||
echo " - omni_pipeline_text.stage1.log"
|
||||
echo " - omni_pipeline_text.stage2.log"
|
||||
echo "Key checks: overall.stats.jsonl for end-to-end latency/TPS; orchestrator.stats.jsonl for stable per-stage latency; stage*.log for errors or long tails."
|
||||
echo " - outputs/ Generated txt and wav files, there should be 100 text and wav files generated respectively"
|
||||
fi
|
||||
@@ -1,107 +0,0 @@
|
||||
# Qwen3-TTS Benchmark
|
||||
|
||||
Benchmarks for Qwen3-TTS text-to-speech models, comparing vLLM-Omni streaming serving against HuggingFace Transformers offline inference.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
```bash
|
||||
pip install matplotlib aiohttp soundfile numpy tqdm
|
||||
pip install qwen_tts # for HF baseline
|
||||
```
|
||||
|
||||
## Quick Start
|
||||
|
||||
Run the full benchmark (vllm-omni + HF baseline) with a single command:
|
||||
|
||||
```bash
|
||||
cd benchmarks/qwen3-tts
|
||||
bash run_benchmark.sh
|
||||
```
|
||||
|
||||
Results (JSON + PNG plots) are saved to `results/`.
|
||||
|
||||
### Common options
|
||||
|
||||
```bash
|
||||
# Only vllm-omni (skip HF baseline)
|
||||
bash run_benchmark.sh --async-only
|
||||
|
||||
# Only HF baseline
|
||||
bash run_benchmark.sh --hf-only
|
||||
|
||||
# Use a different model (e.g. 1.7B)
|
||||
MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice bash run_benchmark.sh --async-only
|
||||
|
||||
# Use a Voice Clone model
|
||||
MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-Base TASK_TYPE=Base bash run_benchmark.sh --async-only
|
||||
|
||||
# Use batch size 16 for higher throughput
|
||||
BATCH_SIZE=16 bash run_benchmark.sh --async-only
|
||||
|
||||
# Custom GPU, prompt count, concurrency levels
|
||||
GPU_DEVICE=1 NUM_PROMPTS=20 CONCURRENCY="1 4" bash run_benchmark.sh
|
||||
```
|
||||
|
||||
## Manual Steps
|
||||
|
||||
### 1) Start the vLLM-Omni server
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=0 python -m vllm_omni.entrypoints.cli.main serve \
|
||||
"Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice" \
|
||||
--omni --host 127.0.0.1 --port 8000 \
|
||||
--deploy-config vllm_omni/deploy/qwen3_tts.yaml \
|
||||
--stage-overrides '{"0":{"max_num_seqs":1,"gpu_memory_utilization":0.3,"max_num_batched_tokens":512},"1":{"max_num_seqs":1,"gpu_memory_utilization":0.3,"max_num_batched_tokens":8192}}' \
|
||||
--trust-remote-code
|
||||
```
|
||||
|
||||
### 2) Run online serving benchmark
|
||||
|
||||
```bash
|
||||
python benchmarks/qwen3-tts/vllm_omni/bench_tts_serve.py \
|
||||
--port 8000 \
|
||||
--num-prompts 50 \
|
||||
--max-concurrency 1 4 10 \
|
||||
--config-name "async_chunk" \
|
||||
--result-dir results/
|
||||
```
|
||||
|
||||
### 3) Run HuggingFace baseline
|
||||
|
||||
```bash
|
||||
python benchmarks/qwen3-tts/transformers/bench_tts_hf.py \
|
||||
--model "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice" \
|
||||
--num-prompts 50 \
|
||||
--gpu-device 0 \
|
||||
--result-dir results/
|
||||
```
|
||||
|
||||
### 4) Generate comparison plots
|
||||
|
||||
```bash
|
||||
python benchmarks/qwen3-tts/plot_results.py \
|
||||
--results results/bench_async_chunk_*.json results/bench_hf_transformers_*.json \
|
||||
--labels "vllm-omni" "hf_transformers" \
|
||||
--output results/comparison.png
|
||||
```
|
||||
|
||||
## Batch-size presets
|
||||
|
||||
The bench script loads the bundled production deploy (`vllm_omni/deploy/qwen3_tts.yaml`) and layers per-stage budgets on top via `--stage-overrides`, driven by the `BATCH_SIZE` env var. Each batch size picks compatible per-stage `max_num_seqs`, `max_num_batched_tokens`, and `gpu_memory_utilization` defaults:
|
||||
|
||||
| `BATCH_SIZE` | Description |
|
||||
|:--:|-------------|
|
||||
| `1` (default) | Single-request processing (lowest latency) |
|
||||
| `4` | Moderate-throughput concurrent processing |
|
||||
| `16` | High-throughput concurrent processing |
|
||||
|
||||
The 2-stage pipeline (Talker -> Code2Wav) runs with `async_chunk` streaming enabled via the prod deploy; the `SharedMemoryConnector` streams codec frames (25-frame chunks with 25-frame context overlap) between stages.
|
||||
|
||||
The model is specified via the CLI `--model` flag (or `MODEL` env var), so the same bench script works for both the 0.6B and 1.7B model variants.
|
||||
|
||||
## Metrics
|
||||
|
||||
- **TTFP (Time to First Audio Packet)**: Time from request to first audio chunk (streaming latency)
|
||||
- **E2E (End-to-End Latency)**: Total time from request to complete audio response
|
||||
- **RTF (Real-Time Factor)**: E2E latency / audio duration. RTF < 1.0 means faster-than-real-time synthesis
|
||||
- **Throughput**: Total audio seconds generated per wall-clock second
|
||||
@@ -1,254 +0,0 @@
|
||||
"""Plot Qwen3-TTS benchmark results.
|
||||
|
||||
Generates comparison bar charts similar to the async_chunk design doc:
|
||||
- TTFP (Time-to-First-Packet) across concurrency levels
|
||||
- E2E latency across concurrency levels
|
||||
- RTF (Real-Time Factor) across concurrency levels
|
||||
|
||||
Usage:
|
||||
# Compare two configs (async_chunk vs no_async_chunk):
|
||||
python plot_results.py \
|
||||
--results results/bench_async_chunk_*.json results/bench_no_async_chunk_*.json \
|
||||
--labels "async_chunk" "no_async_chunk" \
|
||||
--output results/qwen3_tts_benchmark.png
|
||||
|
||||
# Single config:
|
||||
python plot_results.py \
|
||||
--results results/bench_async_chunk_*.json \
|
||||
--labels "async_chunk" \
|
||||
--output results/qwen3_tts_benchmark.png
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
|
||||
|
||||
def load_results(result_files: list[str]) -> list[list[dict]]:
|
||||
"""Load benchmark results from JSON files."""
|
||||
all_results = []
|
||||
for f in result_files:
|
||||
with open(f) as fh:
|
||||
data = json.load(fh)
|
||||
all_results.append(data)
|
||||
return all_results
|
||||
|
||||
|
||||
def plot_comparison(
|
||||
all_results: list[list[dict]],
|
||||
labels: list[str],
|
||||
output_path: str,
|
||||
title_prefix: str = "Qwen3-TTS",
|
||||
):
|
||||
"""Generate comparison bar charts."""
|
||||
n_configs = len(all_results)
|
||||
|
||||
# Collect concurrency levels present in ALL configs (skip missing data)
|
||||
all_concurrencies = [set(r["concurrency"] for r in results) for results in all_results]
|
||||
concurrencies = sorted(set.union(*all_concurrencies))
|
||||
|
||||
# Build data arrays, using None for missing concurrency levels
|
||||
ttfp_data = {label: [] for label in labels}
|
||||
e2e_data = {label: [] for label in labels}
|
||||
rtf_data = {label: [] for label in labels}
|
||||
throughput_data = {label: [] for label in labels}
|
||||
|
||||
for results, label in zip(all_results, labels):
|
||||
conc_map = {r["concurrency"]: r for r in results}
|
||||
for c in concurrencies:
|
||||
r = conc_map.get(c)
|
||||
ttfp_data[label].append(r["mean_ttfp_ms"] if r else None)
|
||||
e2e_data[label].append(r["mean_e2e_ms"] if r else None)
|
||||
rtf_data[label].append(r["mean_rtf"] if r else None)
|
||||
throughput_data[label].append(r["audio_throughput"] if r else None)
|
||||
|
||||
fig, axes = plt.subplots(2, 2, figsize=(14, 10))
|
||||
fig.suptitle(f"{title_prefix} Performance Benchmark", fontsize=16, fontweight="bold")
|
||||
|
||||
x = np.arange(len(concurrencies))
|
||||
width = 0.35 if n_configs == 2 else 0.5
|
||||
if n_configs > 1:
|
||||
offsets = np.linspace(-width / 2 * (n_configs - 1), width / 2 * (n_configs - 1), n_configs)
|
||||
else:
|
||||
offsets = [0]
|
||||
|
||||
colors = ["#2196F3", "#FF5722", "#4CAF50", "#FFC107"]
|
||||
|
||||
def plot_metric(ax, data_dict, ylabel, title, fmt=".1f"):
|
||||
bars = []
|
||||
for i, (label, values) in enumerate(data_dict.items()):
|
||||
# Replace None with 0 for plotting, but track which are missing
|
||||
plot_values = [v if v is not None else 0 for v in values]
|
||||
color = colors[i % len(colors)]
|
||||
bar = ax.bar(x + offsets[i], plot_values, width, label=label, color=color, alpha=0.85)
|
||||
bars.append(bar)
|
||||
# Add value labels on bars (skip None/missing data)
|
||||
max_val = max((v for v in values if v is not None), default=1)
|
||||
for rect, val in zip(bar, values):
|
||||
if val is not None and val > 0:
|
||||
ax.text(
|
||||
rect.get_x() + rect.get_width() / 2,
|
||||
rect.get_height() + max_val * 0.02,
|
||||
f"{val:{fmt}}",
|
||||
ha="center",
|
||||
va="bottom",
|
||||
fontsize=9,
|
||||
fontweight="bold",
|
||||
)
|
||||
ax.set_xlabel("Concurrency", fontsize=12)
|
||||
ax.set_ylabel(ylabel, fontsize=12)
|
||||
ax.set_title(title, fontsize=13, fontweight="bold")
|
||||
ax.set_xticks(x)
|
||||
ax.set_xticklabels([str(c) for c in concurrencies])
|
||||
ax.legend(fontsize=10)
|
||||
ax.grid(axis="y", alpha=0.3)
|
||||
ax.set_axisbelow(True)
|
||||
|
||||
plot_metric(axes[0, 0], ttfp_data, "TTFP (ms)", "Time to First Audio Packet (TTFP)")
|
||||
plot_metric(axes[0, 1], e2e_data, "E2E Latency (ms)", "End-to-End Latency (E2E)")
|
||||
plot_metric(axes[1, 0], rtf_data, "RTF", "Real-Time Factor (RTF)", fmt=".3f")
|
||||
plot_metric(axes[1, 1], throughput_data, "Audio-sec / Wall-sec", "Audio Throughput", fmt=".2f")
|
||||
|
||||
plt.tight_layout()
|
||||
plt.savefig(output_path, dpi=150, bbox_inches="tight")
|
||||
print(f"Plot saved to {output_path}")
|
||||
plt.close()
|
||||
|
||||
|
||||
def plot_single_summary(results: list[dict], label: str, output_path: str):
|
||||
"""Generate a single-config summary with percentile breakdown."""
|
||||
concurrencies = [r["concurrency"] for r in results]
|
||||
|
||||
fig, axes = plt.subplots(1, 3, figsize=(16, 5))
|
||||
fig.suptitle(f"Qwen3-TTS Benchmark - {label}", fontsize=15, fontweight="bold")
|
||||
|
||||
# TTFP breakdown
|
||||
ax = axes[0]
|
||||
means = [r["mean_ttfp_ms"] for r in results]
|
||||
medians = [r["median_ttfp_ms"] for r in results]
|
||||
p90s = [r["p90_ttfp_ms"] for r in results]
|
||||
p99s = [r["p99_ttfp_ms"] for r in results]
|
||||
x = np.arange(len(concurrencies))
|
||||
w = 0.2
|
||||
ax.bar(x - 1.5 * w, means, w, label="mean", color="#2196F3")
|
||||
ax.bar(x - 0.5 * w, medians, w, label="median", color="#4CAF50")
|
||||
ax.bar(x + 0.5 * w, p90s, w, label="p90", color="#FF9800")
|
||||
ax.bar(x + 1.5 * w, p99s, w, label="p99", color="#F44336")
|
||||
ax.set_xticks(x)
|
||||
ax.set_xticklabels([str(c) for c in concurrencies])
|
||||
ax.set_xlabel("Concurrency")
|
||||
ax.set_ylabel("TTFP (ms)")
|
||||
ax.set_title("Time to First Audio Packet")
|
||||
ax.legend(fontsize=9)
|
||||
ax.grid(axis="y", alpha=0.3)
|
||||
|
||||
# E2E breakdown
|
||||
ax = axes[1]
|
||||
means = [r["mean_e2e_ms"] for r in results]
|
||||
medians = [r["median_e2e_ms"] for r in results]
|
||||
p90s = [r["p90_e2e_ms"] for r in results]
|
||||
p99s = [r["p99_e2e_ms"] for r in results]
|
||||
ax.bar(x - 1.5 * w, means, w, label="mean", color="#2196F3")
|
||||
ax.bar(x - 0.5 * w, medians, w, label="median", color="#4CAF50")
|
||||
ax.bar(x + 0.5 * w, p90s, w, label="p90", color="#FF9800")
|
||||
ax.bar(x + 1.5 * w, p99s, w, label="p99", color="#F44336")
|
||||
ax.set_xticks(x)
|
||||
ax.set_xticklabels([str(c) for c in concurrencies])
|
||||
ax.set_xlabel("Concurrency")
|
||||
ax.set_ylabel("E2E Latency (ms)")
|
||||
ax.set_title("End-to-End Latency")
|
||||
ax.legend(fontsize=9)
|
||||
ax.grid(axis="y", alpha=0.3)
|
||||
|
||||
# RTF
|
||||
ax = axes[2]
|
||||
means = [r["mean_rtf"] for r in results]
|
||||
medians = [r["median_rtf"] for r in results]
|
||||
ax.bar(x - 0.15, means, 0.3, label="mean", color="#2196F3")
|
||||
ax.bar(x + 0.15, medians, 0.3, label="median", color="#4CAF50")
|
||||
ax.set_xticks(x)
|
||||
ax.set_xticklabels([str(c) for c in concurrencies])
|
||||
ax.set_xlabel("Concurrency")
|
||||
ax.set_ylabel("RTF")
|
||||
ax.set_title("Real-Time Factor")
|
||||
ax.legend(fontsize=9)
|
||||
ax.grid(axis="y", alpha=0.3)
|
||||
|
||||
plt.tight_layout()
|
||||
plt.savefig(output_path, dpi=150, bbox_inches="tight")
|
||||
print(f"Plot saved to {output_path}")
|
||||
plt.close()
|
||||
|
||||
|
||||
def print_comparison_table(all_results: list[list[dict]], labels: list[str]):
|
||||
"""Print a markdown-formatted comparison table."""
|
||||
concurrencies = sorted(set(r["concurrency"] for r in all_results[0]))
|
||||
|
||||
print("\n## Benchmark Results\n")
|
||||
header = "| Metric | Concurrency |"
|
||||
sep = "| --- | --- |"
|
||||
for label in labels:
|
||||
header += f" {label} |"
|
||||
sep += " --- |"
|
||||
print(header)
|
||||
print(sep)
|
||||
|
||||
for metric, key, fmt in [
|
||||
("TTFP (ms)", "mean_ttfp_ms", ".1f"),
|
||||
("E2E (ms)", "mean_e2e_ms", ".1f"),
|
||||
("RTF", "mean_rtf", ".3f"),
|
||||
("Throughput (audio-s/s)", "audio_throughput", ".2f"),
|
||||
]:
|
||||
for c in concurrencies:
|
||||
row = f"| {metric} | {c} |"
|
||||
for results in all_results:
|
||||
conc_map = {r["concurrency"]: r for r in results}
|
||||
val = conc_map.get(c, {}).get(key, 0)
|
||||
row += f" {val:{fmt}} |"
|
||||
print(row)
|
||||
|
||||
# Improvement calculation (only if 2 configs)
|
||||
if len(all_results) == 2:
|
||||
print(f"\n## Improvement ({labels[0]} vs {labels[1]})\n")
|
||||
print("| Metric | Concurrency | Improvement |")
|
||||
print("| --- | --- | --- |")
|
||||
for metric, key in [("TTFP", "mean_ttfp_ms"), ("E2E", "mean_e2e_ms"), ("RTF", "mean_rtf")]:
|
||||
for c in concurrencies:
|
||||
m0 = {r["concurrency"]: r for r in all_results[0]}
|
||||
m1 = {r["concurrency"]: r for r in all_results[1]}
|
||||
v0 = m0.get(c, {}).get(key, 0)
|
||||
v1 = m1.get(c, {}).get(key, 0)
|
||||
if v1 > 0:
|
||||
pct = (v1 - v0) / v1 * 100
|
||||
print(f"| {metric} | {c} | {pct:+.1f}% |")
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Plot Qwen3-TTS benchmark results")
|
||||
parser.add_argument(
|
||||
"--results", type=str, nargs="+", required=True, help="Path(s) to result JSON files (one per config)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--labels", type=str, nargs="+", required=True, help="Labels for each config (must match --results count)"
|
||||
)
|
||||
parser.add_argument("--output", type=str, default="results/qwen3_tts_benchmark.png", help="Output image path")
|
||||
parser.add_argument("--title", type=str, default="Qwen3-TTS", help="Title prefix for the plot")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
assert len(args.results) == len(args.labels), "--results and --labels must have the same count"
|
||||
|
||||
all_results = load_results(args.results)
|
||||
print_comparison_table(all_results, args.labels)
|
||||
|
||||
Path(args.output).parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
if len(all_results) == 1:
|
||||
plot_single_summary(all_results[0], args.labels[0], args.output)
|
||||
else:
|
||||
plot_comparison(all_results, args.labels, args.output, title_prefix=args.title)
|
||||
@@ -1,3 +0,0 @@
|
||||
# Benchmark results are machine-specific - do not commit
|
||||
*
|
||||
!.gitignore
|
||||
@@ -1,285 +0,0 @@
|
||||
#!/bin/bash
|
||||
# Qwen3-TTS Benchmark Runner
|
||||
#
|
||||
# Compares vllm-omni streaming serving vs HuggingFace transformers offline inference.
|
||||
# Produces JSON results and comparison plots.
|
||||
#
|
||||
# Usage:
|
||||
# # Full comparison (vllm-omni + HF):
|
||||
# bash run_benchmark.sh
|
||||
#
|
||||
# # Only vllm-omni async_chunk config:
|
||||
# bash run_benchmark.sh --async-only
|
||||
#
|
||||
# # Only HuggingFace baseline:
|
||||
# bash run_benchmark.sh --hf-only
|
||||
#
|
||||
# # vllm-omni only (skip HF):
|
||||
# bash run_benchmark.sh --skip-hf
|
||||
#
|
||||
# # Custom settings:
|
||||
# GPU_DEVICE=1 NUM_PROMPTS=20 CONCURRENCY="1 4" bash run_benchmark.sh
|
||||
#
|
||||
# # Use 1.7B model:
|
||||
# MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice bash run_benchmark.sh --async-only
|
||||
#
|
||||
# # Use Voice Clone model
|
||||
# MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-Base TASK_TYPE=Base bash run_benchmark.sh --async-only
|
||||
#
|
||||
# # Use batch_size=4:
|
||||
# BATCH_SIZE=4 bash run_benchmark.sh --async-only
|
||||
#
|
||||
# Environment variables:
|
||||
# GPU_DEVICE - GPU index to use (default: 0)
|
||||
# NUM_PROMPTS - Number of prompts per concurrency level (default: 50)
|
||||
# CONCURRENCY - Space-separated concurrency levels (default: "1 4 10")
|
||||
# MODEL - Model name (default: Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice)
|
||||
# PORT - Server port (default: 8000)
|
||||
# BATCH_SIZE - Per-stage ``max_num_seqs`` for both talker and code2wav (default: 1)
|
||||
# GPU_MEM_TALKER - gpu_memory_utilization for talker stage (default: 0.3 at bs=1, else 0.2)
|
||||
# GPU_MEM_CODE2WAV - gpu_memory_utilization for code2wav stage (default: 0.3 at bs=1, else 0.2)
|
||||
# TASK_TYPE - Task type: CustomVoice, VoiceDesign, Base (default: CustomVoice)
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
PROJECT_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)"
|
||||
|
||||
# Defaults
|
||||
GPU_DEVICE="${GPU_DEVICE:-0}"
|
||||
NUM_PROMPTS="${NUM_PROMPTS:-50}"
|
||||
CONCURRENCY="${CONCURRENCY:-1 4 10}"
|
||||
MODEL="${MODEL:-Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice}"
|
||||
PORT="${PORT:-8000}"
|
||||
BATCH_SIZE="${BATCH_SIZE:-1}"
|
||||
DEFAULT_MEM=$([ "${BATCH_SIZE}" = "1" ] && echo "0.3" || echo "0.2")
|
||||
GPU_MEM_TALKER="${GPU_MEM_TALKER:-${DEFAULT_MEM}}"
|
||||
GPU_MEM_CODE2WAV="${GPU_MEM_CODE2WAV:-${DEFAULT_MEM}}"
|
||||
NUM_WARMUPS="${NUM_WARMUPS:-3}"
|
||||
DEPLOY_CONFIG="vllm_omni/deploy/qwen3_tts.yaml"
|
||||
RESULT_DIR="${SCRIPT_DIR}/results"
|
||||
TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
|
||||
TASK_TYPE="${TASK_TYPE:-CustomVoice}"
|
||||
|
||||
# Build --stage-overrides JSON from BATCH_SIZE + GPU_MEM_*.
|
||||
STAGE_OVERRIDES=$(
|
||||
BATCH_SIZE="${BATCH_SIZE}" \
|
||||
GPU_MEM_TALKER="${GPU_MEM_TALKER}" \
|
||||
GPU_MEM_CODE2WAV="${GPU_MEM_CODE2WAV}" \
|
||||
python - <<'PYEOF'
|
||||
import json, os
|
||||
bs = int(os.environ["BATCH_SIZE"])
|
||||
mem_t = float(os.environ["GPU_MEM_TALKER"])
|
||||
mem_c = float(os.environ["GPU_MEM_CODE2WAV"])
|
||||
# Prefill budget grows with batch size on both stages.
|
||||
talker_batched = 512 if bs <= 4 else 4096
|
||||
code2wav_batched = 8192 if bs <= 4 else 32768
|
||||
print(json.dumps({
|
||||
"0": {"max_num_seqs": bs, "gpu_memory_utilization": mem_t, "max_num_batched_tokens": talker_batched},
|
||||
"1": {"max_num_seqs": bs, "gpu_memory_utilization": mem_c, "max_num_batched_tokens": code2wav_batched},
|
||||
}))
|
||||
PYEOF
|
||||
)
|
||||
|
||||
# Parse args
|
||||
RUN_ASYNC=true
|
||||
RUN_HF=true
|
||||
for arg in "$@"; do
|
||||
case "$arg" in
|
||||
--async-only) RUN_HF=false ;;
|
||||
--hf-only) RUN_ASYNC=false ;;
|
||||
--skip-hf) RUN_HF=false ;;
|
||||
esac
|
||||
done
|
||||
|
||||
mkdir -p "${RESULT_DIR}"
|
||||
|
||||
echo "============================================================"
|
||||
echo " Qwen3-TTS Benchmark"
|
||||
echo "============================================================"
|
||||
echo " GPU: ${GPU_DEVICE}"
|
||||
echo " Model: ${MODEL}"
|
||||
echo " Prompts: ${NUM_PROMPTS}"
|
||||
echo " Concurrency: ${CONCURRENCY}"
|
||||
echo " Port: ${PORT}"
|
||||
echo " Deploy config: ${DEPLOY_CONFIG}"
|
||||
echo " Batch size: ${BATCH_SIZE}"
|
||||
echo " GPU mem T/C: ${GPU_MEM_TALKER} / ${GPU_MEM_CODE2WAV}"
|
||||
echo " Results: ${RESULT_DIR}"
|
||||
echo " Task type: ${TASK_TYPE}"
|
||||
echo "============================================================"
|
||||
|
||||
# Start server and wait for it to be ready
|
||||
start_server() {
|
||||
local config_name="$1"
|
||||
local log_file="${RESULT_DIR}/server_${config_name}_${TIMESTAMP}.log"
|
||||
|
||||
echo ""
|
||||
echo "Starting server with config: ${config_name}"
|
||||
echo " Deploy config: ${DEPLOY_CONFIG}"
|
||||
echo " Stage overrides: ${STAGE_OVERRIDES}"
|
||||
echo " Log file: ${log_file}"
|
||||
|
||||
VLLM_WORKER_MULTIPROC_METHOD=spawn \
|
||||
CUDA_VISIBLE_DEVICES="${GPU_DEVICE}" \
|
||||
python -m vllm_omni.entrypoints.cli.main serve "${MODEL}" \
|
||||
--omni \
|
||||
--host 127.0.0.1 \
|
||||
--port "${PORT}" \
|
||||
--deploy-config "${DEPLOY_CONFIG}" \
|
||||
--stage-overrides "${STAGE_OVERRIDES}" \
|
||||
--stage-init-timeout 120 \
|
||||
--trust-remote-code \
|
||||
--disable-log-stats \
|
||||
> "${log_file}" 2>&1 &
|
||||
|
||||
SERVER_PID=$!
|
||||
echo " Server PID: ${SERVER_PID}"
|
||||
|
||||
# Wait for server to be ready
|
||||
echo " Waiting for server to be ready..."
|
||||
local max_wait=300
|
||||
local waited=0
|
||||
while [ ${waited} -lt ${max_wait} ]; do
|
||||
if curl -sf "http://127.0.0.1:${PORT}/v1/models" > /dev/null 2>&1; then
|
||||
echo " Server is ready! (waited ${waited}s)"
|
||||
return 0
|
||||
fi
|
||||
# Check if process is still alive
|
||||
if ! kill -0 ${SERVER_PID} 2>/dev/null; then
|
||||
echo " ERROR: Server process died. Check log: ${log_file}"
|
||||
tail -20 "${log_file}"
|
||||
return 1
|
||||
fi
|
||||
sleep 2
|
||||
waited=$((waited + 2))
|
||||
done
|
||||
|
||||
echo " ERROR: Server did not start within ${max_wait}s. Check log: ${log_file}"
|
||||
kill ${SERVER_PID} 2>/dev/null || true
|
||||
return 1
|
||||
}
|
||||
|
||||
# Stop the server
|
||||
stop_server() {
|
||||
if [ -n "${SERVER_PID:-}" ]; then
|
||||
echo " Stopping server (PID: ${SERVER_PID})..."
|
||||
kill ${SERVER_PID} 2>/dev/null || true
|
||||
wait ${SERVER_PID} 2>/dev/null || true
|
||||
# Kill any remaining child processes on the port
|
||||
local pids
|
||||
pids=$(lsof -ti:${PORT} 2>/dev/null || true)
|
||||
if [ -n "${pids}" ]; then
|
||||
echo " Cleaning up remaining processes on port ${PORT}..."
|
||||
echo "${pids}" | xargs kill -9 2>/dev/null || true
|
||||
fi
|
||||
echo " Server stopped."
|
||||
SERVER_PID=""
|
||||
fi
|
||||
}
|
||||
|
||||
# Cleanup on exit
|
||||
trap 'stop_server' EXIT
|
||||
|
||||
# Run benchmark for a given config
|
||||
run_bench() {
|
||||
local config_name="$1"
|
||||
|
||||
echo ""
|
||||
echo "============================================================"
|
||||
echo " Benchmarking: ${config_name}"
|
||||
echo "============================================================"
|
||||
|
||||
start_server "${config_name}"
|
||||
|
||||
# Convert concurrency string to args
|
||||
local conc_args=""
|
||||
for c in ${CONCURRENCY}; do
|
||||
conc_args="${conc_args} ${c}"
|
||||
done
|
||||
|
||||
cd "${PROJECT_ROOT}"
|
||||
python "${SCRIPT_DIR}/vllm_omni/bench_tts_serve.py" \
|
||||
--host 127.0.0.1 \
|
||||
--port "${PORT}" \
|
||||
--num-prompts "${NUM_PROMPTS}" \
|
||||
--max-concurrency ${conc_args} \
|
||||
--num-warmups "${NUM_WARMUPS}" \
|
||||
--config-name "${config_name}" \
|
||||
--result-dir "${RESULT_DIR}" \
|
||||
--task-type "${TASK_TYPE}"
|
||||
|
||||
stop_server
|
||||
|
||||
# Allow GPU memory to settle
|
||||
sleep 5
|
||||
}
|
||||
|
||||
# Run vllm-omni benchmark
|
||||
if [ "${RUN_ASYNC}" = true ]; then
|
||||
run_bench "async_chunk"
|
||||
fi
|
||||
|
||||
# Run HuggingFace baseline benchmark
|
||||
if [ "${RUN_HF}" = true ]; then
|
||||
echo ""
|
||||
echo "============================================================"
|
||||
echo " Benchmarking: HuggingFace transformers (offline)"
|
||||
echo "============================================================"
|
||||
|
||||
cd "${PROJECT_ROOT}"
|
||||
python "${SCRIPT_DIR}/transformers/bench_tts_hf.py" \
|
||||
--model "${MODEL}" \
|
||||
--num-prompts "${NUM_PROMPTS}" \
|
||||
--num-warmups "${NUM_WARMUPS}" \
|
||||
--gpu-device "${GPU_DEVICE}" \
|
||||
--config-name "hf_transformers" \
|
||||
--result-dir "${RESULT_DIR}" \
|
||||
--task-type "${TASK_TYPE}"
|
||||
|
||||
# Allow GPU memory to settle
|
||||
sleep 5
|
||||
fi
|
||||
|
||||
# Plot results
|
||||
echo ""
|
||||
echo "============================================================"
|
||||
echo " Generating plots..."
|
||||
echo "============================================================"
|
||||
|
||||
RESULT_FILES=""
|
||||
LABELS=""
|
||||
|
||||
if [ "${RUN_ASYNC}" = true ]; then
|
||||
ASYNC_FILE=$(ls -t "${RESULT_DIR}"/bench_async_chunk_*.json 2>/dev/null | head -1)
|
||||
if [ -n "${ASYNC_FILE}" ]; then
|
||||
RESULT_FILES="${ASYNC_FILE}"
|
||||
LABELS="async_chunk"
|
||||
fi
|
||||
fi
|
||||
|
||||
if [ "${RUN_HF}" = true ]; then
|
||||
HF_FILE=$(ls -t "${RESULT_DIR}"/bench_hf_transformers_*.json 2>/dev/null | head -1)
|
||||
if [ -n "${HF_FILE}" ]; then
|
||||
if [ -n "${RESULT_FILES}" ]; then
|
||||
RESULT_FILES="${RESULT_FILES} ${HF_FILE}"
|
||||
LABELS="${LABELS} hf_transformers"
|
||||
else
|
||||
RESULT_FILES="${HF_FILE}"
|
||||
LABELS="hf_transformers"
|
||||
fi
|
||||
fi
|
||||
fi
|
||||
|
||||
if [ -n "${RESULT_FILES}" ]; then
|
||||
python "${SCRIPT_DIR}/plot_results.py" \
|
||||
--results ${RESULT_FILES} \
|
||||
--labels ${LABELS} \
|
||||
--output "${RESULT_DIR}/qwen3_tts_benchmark_${TIMESTAMP}.png"
|
||||
fi
|
||||
|
||||
echo ""
|
||||
echo "============================================================"
|
||||
echo " Benchmark complete!"
|
||||
echo " Results: ${RESULT_DIR}"
|
||||
echo "============================================================"
|
||||
@@ -1,301 +0,0 @@
|
||||
"""Benchmark Qwen3-TTS using HuggingFace transformers (qwen_tts library).
|
||||
|
||||
Measures E2E latency, RTF, and audio duration for offline (non-serving) inference.
|
||||
Results are saved in the same JSON format as bench_tts_serve.py for unified plotting.
|
||||
|
||||
Usage:
|
||||
python bench_tts_hf.py \
|
||||
--model Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice \
|
||||
--num-prompts 50 \
|
||||
--num-warmups 3 \
|
||||
--gpu-device 0 \
|
||||
--result-dir results/
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import time
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
import torch
|
||||
|
||||
PROMPTS = [
|
||||
"Hello, welcome to the voice synthesis benchmark test.",
|
||||
"She said she would be here by noon, but nobody showed up.",
|
||||
"The quick brown fox jumps over the lazy dog near the riverbank.",
|
||||
"I can't believe how beautiful the sunset looks from up here on the mountain.",
|
||||
"Please remember to bring your identification documents to the appointment tomorrow morning.",
|
||||
"Have you ever wondered what it would be like to travel through time and visit ancient civilizations?",
|
||||
"The restaurant on the corner serves the best pasta I have ever tasted in my entire life.",
|
||||
"After the meeting, we should discuss the quarterly results and plan for the next phase.",
|
||||
"Learning a new language takes patience, practice, and a genuine curiosity about other cultures.",
|
||||
"The train leaves at half past seven, so we need to arrive at the station before then.",
|
||||
"Could you please turn down the music a little bit, I'm trying to concentrate on my work.",
|
||||
"It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.",
|
||||
]
|
||||
|
||||
REF_AUDIO = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-TTS-Repo/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."
|
||||
INSTRUCT = "Speak in an incredulous tone, but with a hint of panic beginning to creep into your voice."
|
||||
|
||||
|
||||
@dataclass
|
||||
class BenchmarkResult:
|
||||
config_name: str = ""
|
||||
concurrency: int = 1 # always 1 for offline
|
||||
num_prompts: int = 0
|
||||
completed: int = 0
|
||||
failed: int = 0
|
||||
duration_s: float = 0.0
|
||||
# TTFP stats - not applicable for HF offline, set to E2E for compatibility
|
||||
mean_ttfp_ms: float = 0.0
|
||||
median_ttfp_ms: float = 0.0
|
||||
std_ttfp_ms: float = 0.0
|
||||
p90_ttfp_ms: float = 0.0
|
||||
p95_ttfp_ms: float = 0.0
|
||||
p99_ttfp_ms: float = 0.0
|
||||
# E2E stats (ms)
|
||||
mean_e2e_ms: float = 0.0
|
||||
median_e2e_ms: float = 0.0
|
||||
std_e2e_ms: float = 0.0
|
||||
p90_e2e_ms: float = 0.0
|
||||
p95_e2e_ms: float = 0.0
|
||||
p99_e2e_ms: float = 0.0
|
||||
# RTF stats
|
||||
mean_rtf: float = 0.0
|
||||
median_rtf: float = 0.0
|
||||
std_rtf: float = 0.0
|
||||
p99_rtf: float = 0.0
|
||||
# Audio stats
|
||||
mean_audio_duration_s: float = 0.0
|
||||
total_audio_duration_s: float = 0.0
|
||||
audio_throughput: float = 0.0
|
||||
request_throughput: float = 0.0
|
||||
# Per-request details
|
||||
per_request: list = field(default_factory=list)
|
||||
|
||||
|
||||
def generate_audio(model, prompt: str, args):
|
||||
if args.task_type == "Base":
|
||||
return model.generate_voice_clone(
|
||||
text=prompt,
|
||||
language=args.language,
|
||||
ref_audio=REF_AUDIO,
|
||||
ref_text=REF_TEXT,
|
||||
)
|
||||
|
||||
if args.task_type == "VoiceDesign":
|
||||
return model.generate_voice_design(
|
||||
text=prompt,
|
||||
language=args.language,
|
||||
instruct=INSTRUCT,
|
||||
)
|
||||
|
||||
return model.generate_custom_voice(
|
||||
text=prompt,
|
||||
language=args.language,
|
||||
speaker=args.voice,
|
||||
)
|
||||
|
||||
|
||||
def run_benchmark(args):
|
||||
from qwen_tts import Qwen3TTSModel
|
||||
|
||||
device = f"cuda:{args.gpu_device}"
|
||||
print(f"Loading model: {args.model} on {device}")
|
||||
model = Qwen3TTSModel.from_pretrained(
|
||||
args.model,
|
||||
device_map=device,
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
print("Model loaded.")
|
||||
|
||||
# Build prompt list
|
||||
prompts = [PROMPTS[i % len(PROMPTS)] for i in range(args.num_prompts)]
|
||||
|
||||
# Warmup
|
||||
if args.num_warmups > 0:
|
||||
print(f"Warming up with {args.num_warmups} requests...")
|
||||
for i in range(args.num_warmups):
|
||||
p = PROMPTS[i % len(PROMPTS)]
|
||||
wavs, sr = generate_audio(model, p, args)
|
||||
# Sync GPU
|
||||
torch.cuda.synchronize(device)
|
||||
print("Warmup done.")
|
||||
|
||||
# Benchmark
|
||||
print(f"Running {args.num_prompts} requests sequentially...")
|
||||
e2e_times = []
|
||||
rtfs = []
|
||||
audio_durations = []
|
||||
per_request = []
|
||||
failed = 0
|
||||
|
||||
audio_dir = None
|
||||
if args.save_audio:
|
||||
audio_dir = Path(args.result_dir) / "audio_hf"
|
||||
audio_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
total_start = time.perf_counter()
|
||||
|
||||
for i, prompt in enumerate(prompts):
|
||||
try:
|
||||
torch.cuda.synchronize(device)
|
||||
st = time.perf_counter()
|
||||
|
||||
wavs, sr = generate_audio(model, prompt, args)
|
||||
|
||||
torch.cuda.synchronize(device)
|
||||
elapsed = time.perf_counter() - st
|
||||
|
||||
# Compute audio duration
|
||||
audio_samples = wavs[0]
|
||||
if isinstance(audio_samples, torch.Tensor):
|
||||
audio_samples = audio_samples.cpu().numpy()
|
||||
audio_dur = len(audio_samples) / sr
|
||||
|
||||
rtf = elapsed / audio_dur if audio_dur > 0 else 0.0
|
||||
|
||||
e2e_times.append(elapsed)
|
||||
rtfs.append(rtf)
|
||||
audio_durations.append(audio_dur)
|
||||
per_request.append(
|
||||
{
|
||||
"e2e_ms": elapsed * 1000,
|
||||
"ttfp_ms": elapsed * 1000, # no streaming, TTFP = E2E
|
||||
"rtf": rtf,
|
||||
"audio_duration_s": audio_dur,
|
||||
"prompt": prompt,
|
||||
}
|
||||
)
|
||||
|
||||
if audio_dir:
|
||||
sf.write(str(audio_dir / f"output_{i:04d}.wav"), audio_samples, sr)
|
||||
|
||||
if (i + 1) % 10 == 0 or i == 0:
|
||||
print(
|
||||
f" [{i + 1}/{args.num_prompts}] e2e={elapsed * 1000:.0f}ms rtf={rtf:.3f} audio={audio_dur:.2f}s"
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
print(f" [{i + 1}/{args.num_prompts}] FAILED: {e}")
|
||||
failed += 1
|
||||
|
||||
total_duration = time.perf_counter() - total_start
|
||||
completed = len(e2e_times)
|
||||
|
||||
# Compute stats
|
||||
result = BenchmarkResult(
|
||||
config_name=args.config_name,
|
||||
concurrency=1,
|
||||
num_prompts=args.num_prompts,
|
||||
completed=completed,
|
||||
failed=failed,
|
||||
duration_s=total_duration,
|
||||
)
|
||||
|
||||
if e2e_times:
|
||||
e2e_ms = [t * 1000 for t in e2e_times]
|
||||
|
||||
result.mean_e2e_ms = float(np.mean(e2e_ms))
|
||||
result.median_e2e_ms = float(np.median(e2e_ms))
|
||||
result.std_e2e_ms = float(np.std(e2e_ms))
|
||||
result.p90_e2e_ms = float(np.percentile(e2e_ms, 90))
|
||||
result.p95_e2e_ms = float(np.percentile(e2e_ms, 95))
|
||||
result.p99_e2e_ms = float(np.percentile(e2e_ms, 99))
|
||||
|
||||
# For HF offline, TTFP = E2E (no streaming)
|
||||
result.mean_ttfp_ms = result.mean_e2e_ms
|
||||
result.median_ttfp_ms = result.median_e2e_ms
|
||||
result.std_ttfp_ms = result.std_e2e_ms
|
||||
result.p90_ttfp_ms = result.p90_e2e_ms
|
||||
result.p95_ttfp_ms = result.p95_e2e_ms
|
||||
result.p99_ttfp_ms = result.p99_e2e_ms
|
||||
|
||||
result.mean_rtf = float(np.mean(rtfs))
|
||||
result.median_rtf = float(np.median(rtfs))
|
||||
result.std_rtf = float(np.std(rtfs))
|
||||
result.p99_rtf = float(np.percentile(rtfs, 99))
|
||||
|
||||
result.mean_audio_duration_s = float(np.mean(audio_durations))
|
||||
result.total_audio_duration_s = float(np.sum(audio_durations))
|
||||
result.audio_throughput = result.total_audio_duration_s / total_duration
|
||||
result.request_throughput = completed / total_duration
|
||||
result.per_request = per_request
|
||||
|
||||
# Print summary in standardized performance template
|
||||
W = 50
|
||||
print("")
|
||||
print(f"{'=' * W}")
|
||||
print(f"{'Serving Benchmark Result':^{W}}")
|
||||
print(f"{'=' * W}")
|
||||
print(f"{'Successful requests:':<40}{completed:<10}")
|
||||
print(f"{'Failed requests:':<40}{failed:<10}")
|
||||
print(f"{'Maximum request concurrency:':<40}{1:<10}")
|
||||
print(f"{'Benchmark duration (s):':<40}{total_duration:<10.2f}")
|
||||
print(f"{'Request throughput (req/s):':<40}{result.request_throughput:<10.2f}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'End-to-end Latency':^{W}}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Mean E2EL (ms):':<40}{result.mean_e2e_ms:<10.2f}")
|
||||
print(f"{'Median E2EL (ms):':<40}{result.median_e2e_ms:<10.2f}")
|
||||
print(f"{'P99 E2EL (ms):':<40}{result.p99_e2e_ms:<10.2f}")
|
||||
print(f"{'=' * W}")
|
||||
print(f"{'Audio Result':^{W}}")
|
||||
print(f"{'=' * W}")
|
||||
print(f"{'Total audio duration generated (s):':<40}{result.total_audio_duration_s:<10.2f}")
|
||||
print(f"{'Audio throughput (audio duration/s):':<40}{result.audio_throughput:<10.2f}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Time to First Packet':^{W}}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Mean AUDIO_TTFP (ms):':<40}{result.mean_ttfp_ms:<10.2f}")
|
||||
print(f"{'Median AUDIO_TTFP (ms):':<40}{result.median_ttfp_ms:<10.2f}")
|
||||
print(f"{'P99 AUDIO_TTFP (ms):':<40}{result.p99_ttfp_ms:<10.2f}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Real Time Factor':^{W}}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Mean AUDIO_RTF:':<40}{result.mean_rtf:<10.3f}")
|
||||
print(f"{'Median AUDIO_RTF:':<40}{result.median_rtf:<10.3f}")
|
||||
print(f"{'P99 AUDIO_RTF:':<40}{result.p99_rtf:<10.3f}")
|
||||
print(f"{'=' * W}")
|
||||
print("")
|
||||
|
||||
# Save results (as a list with single concurrency=1 entry, matching serve format)
|
||||
result_dir = Path(args.result_dir)
|
||||
result_dir.mkdir(parents=True, exist_ok=True)
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
result_file = result_dir / f"bench_{args.config_name}_{timestamp}.json"
|
||||
|
||||
with open(result_file, "w") as f:
|
||||
json.dump([asdict(result)], f, indent=2)
|
||||
print(f"Results saved to {result_file}")
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Qwen3-TTS HuggingFace Benchmark")
|
||||
parser.add_argument(
|
||||
"--model", type=str, default="Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice", help="HuggingFace model name or path"
|
||||
)
|
||||
parser.add_argument("--num-prompts", type=int, default=50)
|
||||
parser.add_argument("--num-warmups", type=int, default=3)
|
||||
parser.add_argument("--gpu-device", type=int, default=0)
|
||||
parser.add_argument("--voice", type=str, default="Vivian")
|
||||
parser.add_argument("--language", type=str, default="English")
|
||||
parser.add_argument("--task-type", type=str, default="CustomVoice", choices=["CustomVoice", "VoiceDesign", "Base"])
|
||||
parser.add_argument(
|
||||
"--config-name", type=str, default="hf_transformers", help="Label for this config (used in filenames)"
|
||||
)
|
||||
parser.add_argument("--result-dir", type=str, default="results")
|
||||
parser.add_argument("--save-audio", action="store_true", help="Save generated audio files")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
run_benchmark(args)
|
||||
@@ -1,301 +0,0 @@
|
||||
"""Benchmark comparing async_chunk on vs off for Qwen3-TTS.
|
||||
|
||||
Measures TTFP (Time-to-First-Packet), E2E latency, and RTF across
|
||||
concurrency levels for both async_chunk modes. Saves results as JSON.
|
||||
|
||||
Usage:
|
||||
# Run against a server already serving with a given config:
|
||||
python bench_async_chunk.py \
|
||||
--host 127.0.0.1 --port 8000 \
|
||||
--config-name async_chunk_on \
|
||||
--num-prompts 50 \
|
||||
--max-concurrency 1 10 \
|
||||
--result-dir results/
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import time
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
import aiohttp
|
||||
import numpy as np
|
||||
from tqdm.asyncio import tqdm
|
||||
|
||||
PROMPTS = [
|
||||
"Hello, welcome to the voice synthesis benchmark test.",
|
||||
"She said she would be here by noon, but nobody showed up.",
|
||||
"The quick brown fox jumps over the lazy dog near the riverbank.",
|
||||
"I can't believe how beautiful the sunset looks from up here on the mountain.",
|
||||
"Please remember to bring your identification documents to the appointment tomorrow morning.",
|
||||
"Have you ever wondered what it would be like to travel through time and visit ancient civilizations?",
|
||||
"The restaurant on the corner serves the best pasta I have ever tasted in my entire life.",
|
||||
"After the meeting, we should discuss the quarterly results and plan for the next phase.",
|
||||
"Learning a new language takes patience, practice, and a genuine curiosity about other cultures.",
|
||||
"The train leaves at half past seven, so we need to arrive at the station before then.",
|
||||
"Could you please turn down the music a little bit, I'm trying to concentrate on my work.",
|
||||
"It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.",
|
||||
]
|
||||
|
||||
|
||||
@dataclass
|
||||
class RequestResult:
|
||||
success: bool = False
|
||||
ttfp: float = 0.0
|
||||
e2e: float = 0.0
|
||||
audio_bytes: int = 0
|
||||
audio_duration: float = 0.0
|
||||
rtf: float = 0.0
|
||||
prompt: str = ""
|
||||
error: str = ""
|
||||
|
||||
|
||||
@dataclass
|
||||
class BenchmarkResult:
|
||||
config_name: str = ""
|
||||
concurrency: int = 0
|
||||
num_prompts: int = 0
|
||||
completed: int = 0
|
||||
failed: int = 0
|
||||
duration_s: float = 0.0
|
||||
mean_ttfp_ms: float = 0.0
|
||||
median_ttfp_ms: float = 0.0
|
||||
std_ttfp_ms: float = 0.0
|
||||
p90_ttfp_ms: float = 0.0
|
||||
p95_ttfp_ms: float = 0.0
|
||||
p99_ttfp_ms: float = 0.0
|
||||
mean_e2e_ms: float = 0.0
|
||||
median_e2e_ms: float = 0.0
|
||||
std_e2e_ms: float = 0.0
|
||||
p90_e2e_ms: float = 0.0
|
||||
p95_e2e_ms: float = 0.0
|
||||
p99_e2e_ms: float = 0.0
|
||||
mean_rtf: float = 0.0
|
||||
median_rtf: float = 0.0
|
||||
std_rtf: float = 0.0
|
||||
mean_audio_duration_s: float = 0.0
|
||||
total_audio_duration_s: float = 0.0
|
||||
audio_throughput: float = 0.0
|
||||
request_throughput: float = 0.0
|
||||
per_request: list = field(default_factory=list)
|
||||
|
||||
|
||||
def pcm_bytes_to_duration(num_bytes: int, sample_rate: int = 24000, sample_width: int = 2) -> float:
|
||||
return num_bytes / sample_width / sample_rate
|
||||
|
||||
|
||||
async def send_tts_request(
|
||||
session: aiohttp.ClientSession,
|
||||
api_url: str,
|
||||
prompt: str,
|
||||
voice: str = "vivian",
|
||||
language: str = "English",
|
||||
stream: bool = True,
|
||||
pbar: tqdm | None = None,
|
||||
) -> RequestResult:
|
||||
payload = {
|
||||
"input": prompt,
|
||||
"voice": voice,
|
||||
"language": language,
|
||||
"stream": stream,
|
||||
"response_format": "pcm",
|
||||
}
|
||||
|
||||
result = RequestResult(prompt=prompt)
|
||||
st = time.perf_counter()
|
||||
|
||||
try:
|
||||
async with session.post(api_url, json=payload) as response:
|
||||
if response.status != 200:
|
||||
result.error = f"HTTP {response.status}: {await response.text()}"
|
||||
return result
|
||||
|
||||
first_chunk = True
|
||||
total_bytes = 0
|
||||
|
||||
async for chunk in response.content.iter_any():
|
||||
if first_chunk and len(chunk) > 0:
|
||||
result.ttfp = time.perf_counter() - st
|
||||
first_chunk = False
|
||||
total_bytes += len(chunk)
|
||||
|
||||
result.e2e = time.perf_counter() - st
|
||||
result.audio_bytes = total_bytes
|
||||
result.audio_duration = pcm_bytes_to_duration(total_bytes)
|
||||
if result.audio_duration > 0:
|
||||
result.rtf = result.e2e / result.audio_duration
|
||||
result.success = True
|
||||
|
||||
except Exception as e:
|
||||
result.error = str(e)
|
||||
result.e2e = time.perf_counter() - st
|
||||
|
||||
if pbar:
|
||||
pbar.update(1)
|
||||
return result
|
||||
|
||||
|
||||
async def run_benchmark(
|
||||
host: str,
|
||||
port: int,
|
||||
num_prompts: int,
|
||||
max_concurrency: int,
|
||||
num_warmups: int = 3,
|
||||
voice: str = "vivian",
|
||||
language: str = "English",
|
||||
stream: bool = True,
|
||||
) -> BenchmarkResult:
|
||||
api_url = f"http://{host}:{port}/v1/audio/speech"
|
||||
|
||||
connector = aiohttp.TCPConnector(limit=max_concurrency, limit_per_host=max_concurrency, keepalive_timeout=60)
|
||||
session = aiohttp.ClientSession(connector=connector, timeout=aiohttp.ClientTimeout(total=600))
|
||||
|
||||
if num_warmups > 0:
|
||||
print(f" Warming up with {num_warmups} requests...")
|
||||
warmup_tasks = [
|
||||
send_tts_request(session, api_url, PROMPTS[i % len(PROMPTS)], voice, language, stream)
|
||||
for i in range(num_warmups)
|
||||
]
|
||||
await asyncio.gather(*warmup_tasks)
|
||||
print(" Warmup done.")
|
||||
|
||||
request_prompts = [PROMPTS[i % len(PROMPTS)] for i in range(num_prompts)]
|
||||
|
||||
print(f" Running {num_prompts} requests with concurrency={max_concurrency}...")
|
||||
semaphore = asyncio.Semaphore(max_concurrency)
|
||||
pbar = tqdm(total=num_prompts, desc=f" concurrency={max_concurrency}")
|
||||
|
||||
async def limited_request(prompt):
|
||||
async with semaphore:
|
||||
return await send_tts_request(session, api_url, prompt, voice, language, stream, pbar)
|
||||
|
||||
start_time = time.perf_counter()
|
||||
tasks = [asyncio.create_task(limited_request(p)) for p in request_prompts]
|
||||
results: list[RequestResult] = await asyncio.gather(*tasks)
|
||||
duration = time.perf_counter() - start_time
|
||||
pbar.close()
|
||||
|
||||
await session.close()
|
||||
|
||||
successful = [r for r in results if r.success]
|
||||
failed = [r for r in results if not r.success]
|
||||
|
||||
bench = BenchmarkResult(
|
||||
concurrency=max_concurrency,
|
||||
num_prompts=num_prompts,
|
||||
completed=len(successful),
|
||||
failed=len(failed),
|
||||
duration_s=duration,
|
||||
)
|
||||
|
||||
if successful:
|
||||
ttfps = [r.ttfp * 1000 for r in successful]
|
||||
e2es = [r.e2e * 1000 for r in successful]
|
||||
rtfs = [r.rtf for r in successful]
|
||||
audio_durs = [r.audio_duration for r in successful]
|
||||
|
||||
bench.mean_ttfp_ms = float(np.mean(ttfps))
|
||||
bench.median_ttfp_ms = float(np.median(ttfps))
|
||||
bench.std_ttfp_ms = float(np.std(ttfps))
|
||||
bench.p90_ttfp_ms = float(np.percentile(ttfps, 90))
|
||||
bench.p95_ttfp_ms = float(np.percentile(ttfps, 95))
|
||||
bench.p99_ttfp_ms = float(np.percentile(ttfps, 99))
|
||||
|
||||
bench.mean_e2e_ms = float(np.mean(e2es))
|
||||
bench.median_e2e_ms = float(np.median(e2es))
|
||||
bench.std_e2e_ms = float(np.std(e2es))
|
||||
bench.p90_e2e_ms = float(np.percentile(e2es, 90))
|
||||
bench.p95_e2e_ms = float(np.percentile(e2es, 95))
|
||||
bench.p99_e2e_ms = float(np.percentile(e2es, 99))
|
||||
|
||||
bench.mean_rtf = float(np.mean(rtfs))
|
||||
bench.median_rtf = float(np.median(rtfs))
|
||||
bench.std_rtf = float(np.std(rtfs))
|
||||
|
||||
bench.mean_audio_duration_s = float(np.mean(audio_durs))
|
||||
bench.total_audio_duration_s = float(np.sum(audio_durs))
|
||||
bench.audio_throughput = bench.total_audio_duration_s / duration
|
||||
bench.request_throughput = len(successful) / duration
|
||||
|
||||
bench.per_request = [
|
||||
{
|
||||
"ttfp_ms": r.ttfp * 1000,
|
||||
"e2e_ms": r.e2e * 1000,
|
||||
"rtf": r.rtf,
|
||||
"audio_duration_s": r.audio_duration,
|
||||
"prompt": r.prompt,
|
||||
}
|
||||
for r in successful
|
||||
]
|
||||
|
||||
print(f"\n{'=' * 60}")
|
||||
print(f" Concurrency: {max_concurrency} | Completed: {bench.completed} | Failed: {bench.failed}")
|
||||
print(f" Duration: {duration:.2f}s | Throughput: {bench.request_throughput:.2f} req/s")
|
||||
print(
|
||||
f" TTFP (ms): mean={bench.mean_ttfp_ms:.1f} median={bench.median_ttfp_ms:.1f}"
|
||||
f" p90={bench.p90_ttfp_ms:.1f} p99={bench.p99_ttfp_ms:.1f}"
|
||||
)
|
||||
print(
|
||||
f" E2E (ms): mean={bench.mean_e2e_ms:.1f} median={bench.median_e2e_ms:.1f}"
|
||||
f" p90={bench.p90_e2e_ms:.1f} p99={bench.p99_e2e_ms:.1f}"
|
||||
)
|
||||
print(f" RTF: mean={bench.mean_rtf:.3f} median={bench.median_rtf:.3f}")
|
||||
print(f" Throughput: {bench.audio_throughput:.2f} audio-sec/wall-sec")
|
||||
print(f"{'=' * 60}\n")
|
||||
|
||||
if failed:
|
||||
for r in failed[:3]:
|
||||
print(f" [ERROR] {r.error[:200]}")
|
||||
|
||||
return bench
|
||||
|
||||
|
||||
async def main(args):
|
||||
all_results = []
|
||||
|
||||
for concurrency in args.max_concurrency:
|
||||
result = await run_benchmark(
|
||||
host=args.host,
|
||||
port=args.port,
|
||||
num_prompts=args.num_prompts,
|
||||
max_concurrency=concurrency,
|
||||
num_warmups=args.num_warmups,
|
||||
voice=args.voice,
|
||||
language=args.language,
|
||||
stream=args.stream,
|
||||
)
|
||||
result.config_name = args.config_name
|
||||
all_results.append(asdict(result))
|
||||
|
||||
result_dir = Path(args.result_dir)
|
||||
result_dir.mkdir(parents=True, exist_ok=True)
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
result_file = result_dir / f"bench_{args.config_name}_{timestamp}.json"
|
||||
|
||||
with open(result_file, "w") as f:
|
||||
json.dump(all_results, f, indent=2)
|
||||
print(f"Results saved to {result_file}")
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Qwen3-TTS async_chunk benchmark client")
|
||||
parser.add_argument("--host", type=str, default="127.0.0.1")
|
||||
parser.add_argument("--port", type=int, default=8000)
|
||||
parser.add_argument("--num-prompts", type=int, default=50)
|
||||
parser.add_argument("--max-concurrency", type=int, nargs="+", default=[1, 10])
|
||||
parser.add_argument("--num-warmups", type=int, default=3)
|
||||
parser.add_argument("--voice", type=str, default="vivian")
|
||||
parser.add_argument("--language", type=str, default="English")
|
||||
parser.add_argument("--stream", action="store_true", default=True)
|
||||
parser.add_argument("--no-stream", dest="stream", action="store_false")
|
||||
parser.add_argument("--config-name", type=str, default="async_chunk_on")
|
||||
parser.add_argument("--result-dir", type=str, default="results")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
asyncio.run(main(args))
|
||||
@@ -1,371 +0,0 @@
|
||||
"""Benchmark client for Qwen3-TTS via /v1/audio/speech endpoint.
|
||||
|
||||
Measures TTFP (Time-to-First-Packet), E2E latency, and RTF (Real-Time Factor)
|
||||
across configurable concurrency levels. Saves results as JSON for plotting.
|
||||
|
||||
Usage:
|
||||
python bench_tts_serve.py \
|
||||
--host 127.0.0.1 --port 8000 \
|
||||
--num-prompts 50 \
|
||||
--max-concurrency 1 4 10 \
|
||||
--result-dir results/
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import time
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
import aiohttp
|
||||
import numpy as np
|
||||
from tqdm.asyncio import tqdm
|
||||
|
||||
PROMPTS = [
|
||||
"Hello, welcome to the voice synthesis benchmark test.",
|
||||
"She said she would be here by noon, but nobody showed up.",
|
||||
"The quick brown fox jumps over the lazy dog near the riverbank.",
|
||||
"I can't believe how beautiful the sunset looks from up here on the mountain.",
|
||||
"Please remember to bring your identification documents to the appointment tomorrow morning.",
|
||||
"Have you ever wondered what it would be like to travel through time and visit ancient civilizations?",
|
||||
"The restaurant on the corner serves the best pasta I have ever tasted in my entire life.",
|
||||
"After the meeting, we should discuss the quarterly results and plan for the next phase.",
|
||||
"Learning a new language takes patience, practice, and a genuine curiosity about other cultures.",
|
||||
"The train leaves at half past seven, so we need to arrive at the station before then.",
|
||||
"Could you please turn down the music a little bit, I'm trying to concentrate on my work.",
|
||||
"It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.",
|
||||
]
|
||||
REF_AUDIO = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-TTS-Repo/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."
|
||||
INSTRUCT = "Speak in an incredulous tone, but with a hint of panic beginning to creep into your voice."
|
||||
|
||||
|
||||
@dataclass
|
||||
class RequestResult:
|
||||
success: bool = False
|
||||
ttfp: float = 0.0 # Time to first audio packet (seconds)
|
||||
e2e: float = 0.0 # End-to-end latency (seconds)
|
||||
audio_bytes: int = 0 # Total audio bytes received
|
||||
audio_duration: float = 0.0 # Audio duration in seconds (estimated from PCM)
|
||||
rtf: float = 0.0 # Real-time factor = e2e / audio_duration
|
||||
prompt: str = ""
|
||||
error: str = ""
|
||||
|
||||
|
||||
@dataclass
|
||||
class BenchmarkResult:
|
||||
config_name: str = ""
|
||||
concurrency: int = 0
|
||||
num_prompts: int = 0
|
||||
completed: int = 0
|
||||
failed: int = 0
|
||||
duration_s: float = 0.0
|
||||
# TTFP stats (ms)
|
||||
mean_ttfp_ms: float = 0.0
|
||||
median_ttfp_ms: float = 0.0
|
||||
std_ttfp_ms: float = 0.0
|
||||
p90_ttfp_ms: float = 0.0
|
||||
p95_ttfp_ms: float = 0.0
|
||||
p99_ttfp_ms: float = 0.0
|
||||
# E2E stats (ms)
|
||||
mean_e2e_ms: float = 0.0
|
||||
median_e2e_ms: float = 0.0
|
||||
std_e2e_ms: float = 0.0
|
||||
p90_e2e_ms: float = 0.0
|
||||
p95_e2e_ms: float = 0.0
|
||||
p99_e2e_ms: float = 0.0
|
||||
# RTF stats
|
||||
mean_rtf: float = 0.0
|
||||
median_rtf: float = 0.0
|
||||
std_rtf: float = 0.0
|
||||
p99_rtf: float = 0.0
|
||||
# Audio stats
|
||||
mean_audio_duration_s: float = 0.0
|
||||
total_audio_duration_s: float = 0.0
|
||||
audio_throughput: float = 0.0 # audio_duration / wall_time
|
||||
request_throughput: float = 0.0 # requests / second
|
||||
# Per-request details
|
||||
per_request: list = field(default_factory=list)
|
||||
|
||||
|
||||
def pcm_bytes_to_duration(num_bytes: int, sample_rate: int = 24000, sample_width: int = 2) -> float:
|
||||
"""Convert raw PCM byte count to duration in seconds."""
|
||||
num_samples = num_bytes / sample_width
|
||||
return num_samples / sample_rate
|
||||
|
||||
|
||||
def create_payload(
|
||||
prompt: str, task_type: str = "CustomVoice", voice: str = "vivian", language: str = "English"
|
||||
) -> dict:
|
||||
payload = {
|
||||
"input": prompt,
|
||||
"language": language,
|
||||
"stream": True,
|
||||
"response_format": "pcm",
|
||||
"task_type": task_type,
|
||||
}
|
||||
|
||||
if task_type == "Base":
|
||||
payload["ref_audio"] = REF_AUDIO
|
||||
payload["ref_text"] = REF_TEXT
|
||||
elif task_type == "CustomVoice":
|
||||
payload["voice"] = voice
|
||||
elif task_type == "VoiceDesign":
|
||||
payload["instructions"] = INSTRUCT
|
||||
|
||||
return payload
|
||||
|
||||
|
||||
async def send_tts_request(
|
||||
session: aiohttp.ClientSession,
|
||||
api_url: str,
|
||||
prompt: str,
|
||||
task_type: str = "CustomVoice",
|
||||
voice: str = "vivian",
|
||||
language: str = "English",
|
||||
pbar: tqdm | None = None,
|
||||
) -> RequestResult:
|
||||
"""Send a streaming TTS request and measure latency metrics."""
|
||||
payload = create_payload(prompt, task_type, voice, language)
|
||||
|
||||
result = RequestResult(prompt=prompt)
|
||||
st = time.perf_counter()
|
||||
|
||||
try:
|
||||
async with session.post(api_url, json=payload) as response:
|
||||
if response.status != 200:
|
||||
result.error = f"HTTP {response.status}: {await response.text()}"
|
||||
result.success = False
|
||||
return result
|
||||
|
||||
first_chunk = True
|
||||
total_bytes = 0
|
||||
|
||||
async for chunk in response.content.iter_any():
|
||||
if first_chunk and len(chunk) > 0:
|
||||
result.ttfp = time.perf_counter() - st
|
||||
first_chunk = False
|
||||
total_bytes += len(chunk)
|
||||
|
||||
result.e2e = time.perf_counter() - st
|
||||
result.audio_bytes = total_bytes
|
||||
result.audio_duration = pcm_bytes_to_duration(total_bytes)
|
||||
|
||||
if result.audio_duration > 0:
|
||||
result.rtf = result.e2e / result.audio_duration
|
||||
result.success = True
|
||||
|
||||
except Exception as e:
|
||||
result.error = str(e)
|
||||
result.success = False
|
||||
result.e2e = time.perf_counter() - st
|
||||
|
||||
if pbar:
|
||||
pbar.update(1)
|
||||
return result
|
||||
|
||||
|
||||
async def run_benchmark(
|
||||
host: str,
|
||||
port: int,
|
||||
num_prompts: int,
|
||||
max_concurrency: int,
|
||||
num_warmups: int = 3,
|
||||
task_type: str = "CustomVoice",
|
||||
voice: str = "vivian",
|
||||
language: str = "English",
|
||||
) -> BenchmarkResult:
|
||||
"""Run benchmark at a given concurrency level."""
|
||||
api_url = f"http://{host}:{port}/v1/audio/speech"
|
||||
|
||||
connector = aiohttp.TCPConnector(
|
||||
limit=max_concurrency,
|
||||
limit_per_host=max_concurrency,
|
||||
keepalive_timeout=60,
|
||||
)
|
||||
session = aiohttp.ClientSession(
|
||||
connector=connector,
|
||||
timeout=aiohttp.ClientTimeout(total=600),
|
||||
)
|
||||
|
||||
# Warmup
|
||||
if num_warmups > 0:
|
||||
print(f" Warming up with {num_warmups} requests...")
|
||||
warmup_tasks = []
|
||||
for i in range(num_warmups):
|
||||
prompt = PROMPTS[i % len(PROMPTS)]
|
||||
warmup_tasks.append(send_tts_request(session, api_url, prompt, task_type, voice, language))
|
||||
await asyncio.gather(*warmup_tasks)
|
||||
print(" Warmup done.")
|
||||
|
||||
# Build request list
|
||||
request_prompts = [PROMPTS[i % len(PROMPTS)] for i in range(num_prompts)]
|
||||
|
||||
# Run benchmark
|
||||
print(f" Running {num_prompts} requests with concurrency={max_concurrency}...")
|
||||
semaphore = asyncio.Semaphore(max_concurrency)
|
||||
pbar = tqdm(total=num_prompts, desc=f" concurrency={max_concurrency}")
|
||||
|
||||
async def limited_request(prompt):
|
||||
async with semaphore:
|
||||
return await send_tts_request(session, api_url, prompt, task_type, voice, language, pbar)
|
||||
|
||||
start_time = time.perf_counter()
|
||||
tasks = [asyncio.create_task(limited_request(p)) for p in request_prompts]
|
||||
results: list[RequestResult] = await asyncio.gather(*tasks)
|
||||
duration = time.perf_counter() - start_time
|
||||
pbar.close()
|
||||
|
||||
await session.close()
|
||||
|
||||
# Compute stats
|
||||
successful = [r for r in results if r.success]
|
||||
failed = [r for r in results if not r.success]
|
||||
|
||||
bench = BenchmarkResult(
|
||||
concurrency=max_concurrency,
|
||||
num_prompts=num_prompts,
|
||||
completed=len(successful),
|
||||
failed=len(failed),
|
||||
duration_s=duration,
|
||||
)
|
||||
|
||||
if successful:
|
||||
ttfps = [r.ttfp * 1000 for r in successful] # convert to ms
|
||||
e2es = [r.e2e * 1000 for r in successful]
|
||||
rtfs = [r.rtf for r in successful]
|
||||
audio_durs = [r.audio_duration for r in successful]
|
||||
|
||||
bench.mean_ttfp_ms = float(np.mean(ttfps))
|
||||
bench.median_ttfp_ms = float(np.median(ttfps))
|
||||
bench.std_ttfp_ms = float(np.std(ttfps))
|
||||
bench.p90_ttfp_ms = float(np.percentile(ttfps, 90))
|
||||
bench.p95_ttfp_ms = float(np.percentile(ttfps, 95))
|
||||
bench.p99_ttfp_ms = float(np.percentile(ttfps, 99))
|
||||
|
||||
bench.mean_e2e_ms = float(np.mean(e2es))
|
||||
bench.median_e2e_ms = float(np.median(e2es))
|
||||
bench.std_e2e_ms = float(np.std(e2es))
|
||||
bench.p90_e2e_ms = float(np.percentile(e2es, 90))
|
||||
bench.p95_e2e_ms = float(np.percentile(e2es, 95))
|
||||
bench.p99_e2e_ms = float(np.percentile(e2es, 99))
|
||||
|
||||
bench.mean_rtf = float(np.mean(rtfs))
|
||||
bench.median_rtf = float(np.median(rtfs))
|
||||
bench.std_rtf = float(np.std(rtfs))
|
||||
bench.p99_rtf = float(np.percentile(rtfs, 99))
|
||||
|
||||
bench.mean_audio_duration_s = float(np.mean(audio_durs))
|
||||
bench.total_audio_duration_s = float(np.sum(audio_durs))
|
||||
bench.audio_throughput = bench.total_audio_duration_s / duration
|
||||
bench.request_throughput = len(successful) / duration
|
||||
|
||||
bench.per_request = [
|
||||
{
|
||||
"ttfp_ms": r.ttfp * 1000,
|
||||
"e2e_ms": r.e2e * 1000,
|
||||
"rtf": r.rtf,
|
||||
"audio_duration_s": r.audio_duration,
|
||||
"prompt": r.prompt,
|
||||
}
|
||||
for r in successful
|
||||
]
|
||||
|
||||
# Print summary in standardized performance template
|
||||
W = 50
|
||||
print("")
|
||||
print(f"{'=' * W}")
|
||||
print(f"{'Serving Benchmark Result':^{W}}")
|
||||
print(f"{'=' * W}")
|
||||
print(f"{'Successful requests:':<40}{bench.completed:<10}")
|
||||
print(f"{'Failed requests:':<40}{bench.failed:<10}")
|
||||
print(f"{'Maximum request concurrency:':<40}{max_concurrency:<10}")
|
||||
print(f"{'Benchmark duration (s):':<40}{duration:<10.2f}")
|
||||
print(f"{'Request throughput (req/s):':<40}{bench.request_throughput:<10.2f}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'End-to-end Latency':^{W}}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Mean E2EL (ms):':<40}{bench.mean_e2e_ms:<10.2f}")
|
||||
print(f"{'Median E2EL (ms):':<40}{bench.median_e2e_ms:<10.2f}")
|
||||
print(f"{'P99 E2EL (ms):':<40}{bench.p99_e2e_ms:<10.2f}")
|
||||
print(f"{'=' * W}")
|
||||
print(f"{'Audio Result':^{W}}")
|
||||
print(f"{'=' * W}")
|
||||
print(f"{'Total audio duration generated (s):':<40}{bench.total_audio_duration_s:<10.2f}")
|
||||
print(f"{'Audio throughput (audio duration/s):':<40}{bench.audio_throughput:<10.2f}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Time to First Packet':^{W}}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Mean AUDIO_TTFP (ms):':<40}{bench.mean_ttfp_ms:<10.2f}")
|
||||
print(f"{'Median AUDIO_TTFP (ms):':<40}{bench.median_ttfp_ms:<10.2f}")
|
||||
print(f"{'P99 AUDIO_TTFP (ms):':<40}{bench.p99_ttfp_ms:<10.2f}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Real Time Factor':^{W}}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Mean AUDIO_RTF:':<40}{bench.mean_rtf:<10.3f}")
|
||||
print(f"{'Median AUDIO_RTF:':<40}{bench.median_rtf:<10.3f}")
|
||||
print(f"{'P99 AUDIO_RTF:':<40}{bench.p99_rtf:<10.3f}")
|
||||
print(f"{'=' * W}")
|
||||
print("")
|
||||
|
||||
if failed:
|
||||
for r in failed[:3]:
|
||||
print(f" [ERROR] {r.error[:200]}")
|
||||
|
||||
return bench
|
||||
|
||||
|
||||
async def main(args):
|
||||
all_results = []
|
||||
|
||||
for concurrency in args.max_concurrency:
|
||||
result = await run_benchmark(
|
||||
host=args.host,
|
||||
port=args.port,
|
||||
num_prompts=args.num_prompts,
|
||||
max_concurrency=concurrency,
|
||||
num_warmups=args.num_warmups,
|
||||
task_type=args.task_type,
|
||||
voice=args.voice,
|
||||
language=args.language,
|
||||
)
|
||||
result.config_name = args.config_name
|
||||
all_results.append(asdict(result))
|
||||
|
||||
# Save results
|
||||
result_dir = Path(args.result_dir)
|
||||
result_dir.mkdir(parents=True, exist_ok=True)
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
result_file = result_dir / f"bench_{args.config_name}_{timestamp}.json"
|
||||
|
||||
with open(result_file, "w") as f:
|
||||
json.dump(all_results, f, indent=2)
|
||||
print(f"Results saved to {result_file}")
|
||||
|
||||
return all_results
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Qwen3-TTS Benchmark Client")
|
||||
parser.add_argument("--host", type=str, default="127.0.0.1")
|
||||
parser.add_argument("--port", type=int, default=8000)
|
||||
parser.add_argument("--num-prompts", type=int, default=50, help="Number of prompts per concurrency level")
|
||||
parser.add_argument( # noqa: E501
|
||||
"--max-concurrency", type=int, nargs="+", default=[1, 4, 10], help="Concurrency levels to test"
|
||||
)
|
||||
parser.add_argument("--num-warmups", type=int, default=3)
|
||||
parser.add_argument("--task-type", type=str, default="CustomVoice", choices=["CustomVoice", "VoiceDesign", "Base"])
|
||||
parser.add_argument("--voice", type=str, default="vivian")
|
||||
parser.add_argument("--language", type=str, default="English")
|
||||
parser.add_argument(
|
||||
"--config-name", type=str, default="async_chunk", help="Label for this config (used in filenames)"
|
||||
)
|
||||
parser.add_argument("--result-dir", type=str, default="results")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
asyncio.run(main(args))
|
||||
@@ -1,249 +0,0 @@
|
||||
"""Plot TTFP comparison: async_chunk off vs on.
|
||||
|
||||
Generates a bar chart with improvement arrows, matching the Qwen3-Omni
|
||||
async_chunk benchmark figure style.
|
||||
|
||||
Usage:
|
||||
python plot_async_chunk.py \
|
||||
--off results/bench_async_chunk_off_*.json \
|
||||
--on results/bench_async_chunk_on_*.json \
|
||||
--output results/qwen3_tts_async_chunk_ttfp.png
|
||||
|
||||
# Also supports E2E and RTF metrics:
|
||||
python plot_async_chunk.py \
|
||||
--off results/bench_async_chunk_off_*.json \
|
||||
--on results/bench_async_chunk_on_*.json \
|
||||
--metric e2e \
|
||||
--output results/qwen3_tts_async_chunk_e2e.png
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
|
||||
METRIC_CONFIG = {
|
||||
"ttfp": {
|
||||
"key": "mean_ttfp_ms",
|
||||
"ylabel": "TTFP (s)",
|
||||
"title": "TTFP (Time to First Audio Packet) - Qwen3-TTS, by concurrency",
|
||||
"to_seconds": True,
|
||||
},
|
||||
"e2e": {
|
||||
"key": "mean_e2e_ms",
|
||||
"ylabel": "E2E (s)",
|
||||
"title": "E2E Latency - Qwen3-TTS, by concurrency",
|
||||
"to_seconds": True,
|
||||
},
|
||||
"rtf": {
|
||||
"key": "mean_rtf",
|
||||
"ylabel": "RTF",
|
||||
"title": "Real-Time Factor - Qwen3-TTS, by concurrency",
|
||||
"to_seconds": False,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def load_results(path: str) -> list[dict]:
|
||||
with open(path) as f:
|
||||
return json.load(f)
|
||||
|
||||
|
||||
def plot_ttfp_comparison(
|
||||
off_results: list[dict],
|
||||
on_results: list[dict],
|
||||
metric: str,
|
||||
output_path: str,
|
||||
title_override: str | None = None,
|
||||
):
|
||||
cfg = METRIC_CONFIG[metric]
|
||||
key = cfg["key"]
|
||||
to_seconds = cfg["to_seconds"]
|
||||
|
||||
off_map = {r["concurrency"]: r for r in off_results}
|
||||
on_map = {r["concurrency"]: r for r in on_results}
|
||||
concurrencies = sorted(set(off_map.keys()) & set(on_map.keys()))
|
||||
|
||||
off_vals = []
|
||||
on_vals = []
|
||||
for c in concurrencies:
|
||||
v_off = off_map[c][key]
|
||||
v_on = on_map[c][key]
|
||||
if to_seconds:
|
||||
v_off /= 1000.0
|
||||
v_on /= 1000.0
|
||||
off_vals.append(v_off)
|
||||
on_vals.append(v_on)
|
||||
|
||||
fig, ax = plt.subplots(figsize=(8, 6))
|
||||
|
||||
x = np.arange(len(concurrencies))
|
||||
width = 0.3
|
||||
|
||||
ax.bar(x - width / 2, off_vals, width, label="async_chunk off", color="#87CEEB", edgecolor="none")
|
||||
ax.bar(x + width / 2, on_vals, width, label="async_chunk on", color="#FFF8DC", edgecolor="#DDD8B8")
|
||||
|
||||
# Draw improvement arrows and labels
|
||||
for i in range(len(concurrencies)):
|
||||
v_off = off_vals[i]
|
||||
v_on = on_vals[i]
|
||||
if v_on > 0:
|
||||
improvement = v_off / v_on
|
||||
else:
|
||||
improvement = float("inf")
|
||||
|
||||
# Arrow from top of off-bar to top of on-bar
|
||||
arrow_start_x = x[i] - width / 2
|
||||
arrow_start_y = v_off * 0.95
|
||||
arrow_end_x = x[i] + width / 2
|
||||
arrow_end_y = v_on * 1.05
|
||||
|
||||
ax.annotate(
|
||||
"",
|
||||
xy=(arrow_end_x, arrow_end_y),
|
||||
xytext=(arrow_start_x, arrow_start_y),
|
||||
arrowprops=dict(arrowstyle="->", color="red", lw=1.5),
|
||||
)
|
||||
|
||||
# Improvement label
|
||||
label_x = (arrow_start_x + arrow_end_x) / 2
|
||||
label_y = arrow_start_y + (v_off - v_on) * 0.15
|
||||
ax.text(
|
||||
label_x,
|
||||
label_y,
|
||||
f"{improvement:.1f}x improvement",
|
||||
ha="center",
|
||||
va="bottom",
|
||||
fontsize=10,
|
||||
color="red",
|
||||
fontweight="bold",
|
||||
)
|
||||
|
||||
title = title_override or cfg["title"]
|
||||
ax.set_title(title, fontsize=13, fontweight="bold")
|
||||
ax.set_ylabel(cfg["ylabel"], fontsize=12)
|
||||
ax.set_xlabel("Max concurrency", fontsize=12)
|
||||
ax.set_xticks(x)
|
||||
ax.set_xticklabels([str(c) for c in concurrencies])
|
||||
ax.set_yscale("log")
|
||||
ax.legend(loc="upper left", fontsize=11)
|
||||
ax.grid(axis="y", alpha=0.3, linestyle="--")
|
||||
ax.set_axisbelow(True)
|
||||
|
||||
plt.tight_layout()
|
||||
Path(output_path).parent.mkdir(parents=True, exist_ok=True)
|
||||
plt.savefig(output_path, dpi=150, bbox_inches="tight")
|
||||
print(f"Plot saved to {output_path}")
|
||||
plt.close()
|
||||
|
||||
|
||||
def plot_all_metrics(off_results: list[dict], on_results: list[dict], output_path: str):
|
||||
"""Generate a 1x3 subplot with TTFP, E2E, and RTF comparisons."""
|
||||
off_map = {r["concurrency"]: r for r in off_results}
|
||||
on_map = {r["concurrency"]: r for r in on_results}
|
||||
concurrencies = sorted(set(off_map.keys()) & set(on_map.keys()))
|
||||
|
||||
fig, axes = plt.subplots(1, 3, figsize=(18, 6))
|
||||
fig.suptitle("Qwen3-TTS: async_chunk on vs off", fontsize=15, fontweight="bold")
|
||||
|
||||
for ax, metric in zip(axes, ["ttfp", "e2e", "rtf"]):
|
||||
cfg = METRIC_CONFIG[metric]
|
||||
key = cfg["key"]
|
||||
to_seconds = cfg["to_seconds"]
|
||||
|
||||
off_vals = []
|
||||
on_vals = []
|
||||
for c in concurrencies:
|
||||
v_off = off_map[c][key]
|
||||
v_on = on_map[c][key]
|
||||
if to_seconds:
|
||||
v_off /= 1000.0
|
||||
v_on /= 1000.0
|
||||
off_vals.append(v_off)
|
||||
on_vals.append(v_on)
|
||||
|
||||
x = np.arange(len(concurrencies))
|
||||
width = 0.3
|
||||
ax.bar(x - width / 2, off_vals, width, label="async_chunk off", color="#87CEEB")
|
||||
ax.bar(x + width / 2, on_vals, width, label="async_chunk on", color="#FFF8DC", edgecolor="#DDD8B8")
|
||||
|
||||
for i in range(len(concurrencies)):
|
||||
if on_vals[i] > 0:
|
||||
improvement = off_vals[i] / on_vals[i]
|
||||
ax.annotate(
|
||||
"",
|
||||
xy=(x[i] + width / 2, on_vals[i] * 1.05),
|
||||
xytext=(x[i] - width / 2, off_vals[i] * 0.95),
|
||||
arrowprops=dict(arrowstyle="->", color="red", lw=1.5),
|
||||
)
|
||||
label_y = off_vals[i] * 0.85
|
||||
ax.text(x[i], label_y, f"{improvement:.1f}x", ha="center", fontsize=10, color="red", fontweight="bold")
|
||||
|
||||
ax.set_title(cfg["title"].split(" - ")[0], fontsize=12, fontweight="bold")
|
||||
ax.set_ylabel(cfg["ylabel"], fontsize=11)
|
||||
ax.set_xlabel("Max concurrency", fontsize=11)
|
||||
ax.set_xticks(x)
|
||||
ax.set_xticklabels([str(c) for c in concurrencies])
|
||||
if metric != "rtf":
|
||||
ax.set_yscale("log")
|
||||
ax.legend(fontsize=9)
|
||||
ax.grid(axis="y", alpha=0.3, linestyle="--")
|
||||
ax.set_axisbelow(True)
|
||||
|
||||
plt.tight_layout()
|
||||
Path(output_path).parent.mkdir(parents=True, exist_ok=True)
|
||||
plt.savefig(output_path, dpi=150, bbox_inches="tight")
|
||||
print(f"Plot saved to {output_path}")
|
||||
plt.close()
|
||||
|
||||
|
||||
def print_table(off_results: list[dict], on_results: list[dict]):
|
||||
off_map = {r["concurrency"]: r for r in off_results}
|
||||
on_map = {r["concurrency"]: r for r in on_results}
|
||||
concurrencies = sorted(set(off_map.keys()) & set(on_map.keys()))
|
||||
|
||||
print("\n## Benchmark Results: async_chunk off vs on\n")
|
||||
print("| Metric | Concurrency | async_chunk off | async_chunk on | Improvement |")
|
||||
print("| --- | --- | --- | --- | --- |")
|
||||
|
||||
for name, key, fmt in [
|
||||
("TTFP (ms)", "mean_ttfp_ms", ".1f"),
|
||||
("E2E (ms)", "mean_e2e_ms", ".1f"),
|
||||
("RTF", "mean_rtf", ".3f"),
|
||||
("Throughput", "audio_throughput", ".2f"),
|
||||
]:
|
||||
for c in concurrencies:
|
||||
v_off = off_map[c].get(key, 0)
|
||||
v_on = on_map[c].get(key, 0)
|
||||
if v_on > 0 and key != "audio_throughput":
|
||||
ratio = f"{v_off / v_on:.1f}x"
|
||||
elif v_off > 0 and key == "audio_throughput":
|
||||
ratio = f"{v_on / v_off:.1f}x"
|
||||
else:
|
||||
ratio = "N/A"
|
||||
print(f"| {name} | {c} | {v_off:{fmt}} | {v_on:{fmt}} | {ratio} |")
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Plot async_chunk comparison for Qwen3-TTS")
|
||||
parser.add_argument("--off", type=str, required=True, help="JSON results for async_chunk off")
|
||||
parser.add_argument("--on", type=str, required=True, help="JSON results for async_chunk on")
|
||||
parser.add_argument("--metric", type=str, default="ttfp", choices=["ttfp", "e2e", "rtf", "all"])
|
||||
parser.add_argument("--output", type=str, default="results/qwen3_tts_async_chunk.png")
|
||||
parser.add_argument("--title", type=str, default=None, help="Custom title override")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
off_results = load_results(args.off)
|
||||
on_results = load_results(args.on)
|
||||
|
||||
print_table(off_results, on_results)
|
||||
|
||||
if args.metric == "all":
|
||||
plot_all_metrics(off_results, on_results, args.output)
|
||||
else:
|
||||
plot_ttfp_comparison(off_results, on_results, args.metric, args.output, args.title)
|
||||
@@ -1,169 +0,0 @@
|
||||
#!/bin/bash
|
||||
# Qwen3-TTS async_chunk on vs off Benchmark
|
||||
#
|
||||
# Starts two servers (async_chunk on and off), benchmarks both,
|
||||
# and generates comparison plots.
|
||||
#
|
||||
# Usage:
|
||||
# bash run_async_chunk_benchmark.sh
|
||||
#
|
||||
# Environment variables:
|
||||
# GPU_DEVICE - GPU index (default: 0)
|
||||
# MODEL - Model path (default: Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice)
|
||||
# NUM_PROMPTS - Prompts per concurrency level (default: 50)
|
||||
# CONCURRENCY - Space-separated concurrency levels (default: "1 10")
|
||||
# PORT_ON - Port for async_chunk on server (default: 8000)
|
||||
# PORT_OFF - Port for async_chunk off server (default: 8001)
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
PROJECT_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)"
|
||||
cd "$PROJECT_ROOT"
|
||||
|
||||
GPU_DEVICE="${GPU_DEVICE:-0}"
|
||||
MODEL="${MODEL:-Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice}"
|
||||
NUM_PROMPTS="${NUM_PROMPTS:-50}"
|
||||
CONCURRENCY="${CONCURRENCY:-1 10}"
|
||||
NUM_WARMUPS="${NUM_WARMUPS:-3}"
|
||||
PORT_ON="${PORT_ON:-8000}"
|
||||
PORT_OFF="${PORT_OFF:-8001}"
|
||||
RESULT_DIR="${SCRIPT_DIR}/results"
|
||||
TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
|
||||
|
||||
# The bundled ``vllm_omni/deploy/qwen3_tts.yaml`` is auto-loaded by the model
|
||||
# registry; no ``--deploy-config`` flag needed on the default (ON) path.
|
||||
# async_chunk OFF is selected by the ``--no-async-chunk`` CLI flag —
|
||||
# the single ``qwen3_tts`` pipeline dispatches to the end-to-end codec
|
||||
# processor when ``deploy.async_chunk`` is false.
|
||||
|
||||
mkdir -p "${RESULT_DIR}"
|
||||
|
||||
echo "============================================================"
|
||||
echo " Qwen3-TTS async_chunk Benchmark"
|
||||
echo "============================================================"
|
||||
echo " GPU: ${GPU_DEVICE}"
|
||||
echo " Model: ${MODEL}"
|
||||
echo " Prompts: ${NUM_PROMPTS}"
|
||||
echo " Concurrency: ${CONCURRENCY}"
|
||||
echo " Port (on/off): ${PORT_ON} / ${PORT_OFF}"
|
||||
echo " Results: ${RESULT_DIR}"
|
||||
echo "============================================================"
|
||||
|
||||
cleanup() {
|
||||
echo "Cleaning up servers..."
|
||||
kill "$PID_ON" 2>/dev/null || true
|
||||
kill "$PID_OFF" 2>/dev/null || true
|
||||
wait "$PID_ON" 2>/dev/null || true
|
||||
wait "$PID_OFF" 2>/dev/null || true
|
||||
}
|
||||
trap cleanup EXIT
|
||||
|
||||
wait_for_server() {
|
||||
local port=$1
|
||||
local name=$2
|
||||
local max_wait=300
|
||||
local elapsed=0
|
||||
echo "Waiting for ${name} server on port ${port}..."
|
||||
while ! curl -s "http://localhost:${port}/health" >/dev/null 2>&1; do
|
||||
sleep 5
|
||||
elapsed=$((elapsed + 5))
|
||||
if [ $elapsed -ge $max_wait ]; then
|
||||
echo "ERROR: ${name} server failed to start within ${max_wait}s"
|
||||
exit 1
|
||||
fi
|
||||
done
|
||||
echo "${name} server ready (${elapsed}s)"
|
||||
}
|
||||
|
||||
# ---- Phase 1: Start async_chunk ON server ----
|
||||
echo ""
|
||||
echo "[Phase 1] Starting async_chunk ON server on port ${PORT_ON}..."
|
||||
CUDA_VISIBLE_DEVICES=${GPU_DEVICE} vllm-omni serve "${MODEL}" \
|
||||
--host 0.0.0.0 --port "${PORT_ON}" \
|
||||
--trust-remote-code --enforce-eager --omni \
|
||||
> "${RESULT_DIR}/server_on_${TIMESTAMP}.log" 2>&1 &
|
||||
PID_ON=$!
|
||||
|
||||
wait_for_server "${PORT_ON}" "async_chunk_on"
|
||||
|
||||
echo "[Phase 1] Benchmarking async_chunk ON..."
|
||||
# shellcheck disable=SC2086
|
||||
python "${SCRIPT_DIR}/bench_async_chunk.py" \
|
||||
--host 127.0.0.1 --port "${PORT_ON}" \
|
||||
--config-name "async_chunk_on" \
|
||||
--num-prompts "${NUM_PROMPTS}" \
|
||||
--max-concurrency ${CONCURRENCY} \
|
||||
--num-warmups "${NUM_WARMUPS}" \
|
||||
--result-dir "${RESULT_DIR}"
|
||||
|
||||
echo "[Phase 1] Stopping async_chunk ON server..."
|
||||
kill "$PID_ON" 2>/dev/null || true
|
||||
wait "$PID_ON" 2>/dev/null || true
|
||||
sleep 5
|
||||
|
||||
# ---- Phase 2: Start async_chunk OFF server ----
|
||||
echo ""
|
||||
echo "[Phase 2] Starting async_chunk OFF server on port ${PORT_OFF}..."
|
||||
CUDA_VISIBLE_DEVICES=${GPU_DEVICE} vllm-omni serve "${MODEL}" \
|
||||
--no-async-chunk \
|
||||
--host 0.0.0.0 --port "${PORT_OFF}" \
|
||||
--trust-remote-code --enforce-eager --omni \
|
||||
> "${RESULT_DIR}/server_off_${TIMESTAMP}.log" 2>&1 &
|
||||
PID_OFF=$!
|
||||
|
||||
wait_for_server "${PORT_OFF}" "async_chunk_off"
|
||||
|
||||
echo "[Phase 2] Benchmarking async_chunk OFF (non-streaming)..."
|
||||
# shellcheck disable=SC2086
|
||||
python "${SCRIPT_DIR}/bench_async_chunk.py" \
|
||||
--host 127.0.0.1 --port "${PORT_OFF}" \
|
||||
--config-name "async_chunk_off" \
|
||||
--num-prompts "${NUM_PROMPTS}" \
|
||||
--max-concurrency ${CONCURRENCY} \
|
||||
--num-warmups "${NUM_WARMUPS}" \
|
||||
--no-stream \
|
||||
--result-dir "${RESULT_DIR}"
|
||||
|
||||
echo "[Phase 2] Stopping async_chunk OFF server..."
|
||||
kill "$PID_OFF" 2>/dev/null || true
|
||||
wait "$PID_OFF" 2>/dev/null || true
|
||||
|
||||
# ---- Phase 3: Plot results ----
|
||||
echo ""
|
||||
echo "[Phase 3] Generating plots..."
|
||||
|
||||
# Find the latest result files
|
||||
RESULT_ON=$(ls -t "${RESULT_DIR}"/bench_async_chunk_on_*.json 2>/dev/null | head -1)
|
||||
RESULT_OFF=$(ls -t "${RESULT_DIR}"/bench_async_chunk_off_*.json 2>/dev/null | head -1)
|
||||
|
||||
if [ -z "$RESULT_ON" ] || [ -z "$RESULT_OFF" ]; then
|
||||
echo "ERROR: Could not find result files. Check logs in ${RESULT_DIR}/"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo " ON results: ${RESULT_ON}"
|
||||
echo " OFF results: ${RESULT_OFF}"
|
||||
|
||||
# TTFP comparison (main figure)
|
||||
python "${SCRIPT_DIR}/plot_async_chunk.py" \
|
||||
--off "${RESULT_OFF}" \
|
||||
--on "${RESULT_ON}" \
|
||||
--metric ttfp \
|
||||
--output "${RESULT_DIR}/qwen3_tts_async_chunk_ttfp.png"
|
||||
|
||||
# All metrics comparison
|
||||
python "${SCRIPT_DIR}/plot_async_chunk.py" \
|
||||
--off "${RESULT_OFF}" \
|
||||
--on "${RESULT_ON}" \
|
||||
--metric all \
|
||||
--output "${RESULT_DIR}/qwen3_tts_async_chunk_all.png"
|
||||
|
||||
echo ""
|
||||
echo "============================================================"
|
||||
echo " Benchmark complete!"
|
||||
echo " Results: ${RESULT_DIR}/"
|
||||
echo " Plots:"
|
||||
echo " - ${RESULT_DIR}/qwen3_tts_async_chunk_ttfp.png"
|
||||
echo " - ${RESULT_DIR}/qwen3_tts_async_chunk_all.png"
|
||||
echo "============================================================"
|
||||
@@ -0,0 +1,227 @@
|
||||
# TTS Universal Benchmark
|
||||
|
||||
A model-agnostic serving benchmark for TTS models in vllm-omni. One CLI
|
||||
(`bench_tts.py`) + one YAML registry (`model_configs.yaml`) drive perf and
|
||||
quality runs for every registered checkpoint: **Qwen3-TTS** (Base / CustomVoice)
|
||||
and **VoxCPM2** today, more to come.
|
||||
|
||||
The same three task types — `voice_clone`, `default_voice`, `voice_design` —
|
||||
are wired into both the manual CLI and the DFX nightly CI matrix
|
||||
(`tests/dfx/perf/tests/test_tts.json`).
|
||||
|
||||
## Quick start
|
||||
|
||||
### 1. Start the server
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base --omni --port 8000
|
||||
```
|
||||
|
||||
The server auto-loads its Deploy YAML from `vllm_omni/deploy/qwen3_tts.yaml`
|
||||
(Pipeline + Deploy schema introduced in #2383). No `--stage-configs-path` or
|
||||
`--deploy-config` flag is needed for any registered model.
|
||||
|
||||
### 2. Run the benchmark (`vllm bench serve --omni`)
|
||||
|
||||
The primary, directly-controllable path. Copy-paste one of these and tweak
|
||||
any bench flag (sampling params, endpoint, extra body, warmups, etc.):
|
||||
|
||||
#### voice_clone (Qwen3-TTS-Base, seed-tts dataset)
|
||||
|
||||
```bash
|
||||
vllm bench serve --omni \
|
||||
--host 127.0.0.1 --port 8000 \
|
||||
--model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
|
||||
--backend openai-audio-speech \
|
||||
--endpoint /v1/audio/speech \
|
||||
--dataset-name seed-tts \
|
||||
--dataset-path /path/to/seed-tts-eval \
|
||||
--seed-tts-locale en \
|
||||
--num-prompts 20 --num-warmups 2 \
|
||||
--extra-body '{"task_type":"Base"}' \
|
||||
--max-concurrency 1 --request-rate inf \
|
||||
--percentile-metrics ttft,e2el,audio_rtf,audio_ttfp,audio_duration \
|
||||
--save-result --result-dir ./results
|
||||
```
|
||||
|
||||
#### default_voice (Qwen3-TTS-CustomVoice, bundled seed_tts_smoke)
|
||||
|
||||
```bash
|
||||
vllm bench serve --omni \
|
||||
--host 127.0.0.1 --port 8000 \
|
||||
--model Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
|
||||
--backend openai-audio-speech \
|
||||
--endpoint /v1/audio/speech \
|
||||
--dataset-name seed-tts-text \
|
||||
--dataset-path benchmarks/build_dataset/seed_tts_smoke \
|
||||
--seed-tts-locale en \
|
||||
--num-prompts 20 --num-warmups 2 \
|
||||
--extra-body '{"voice":"Vivian","language":"English","task_type":"CustomVoice"}' \
|
||||
--max-concurrency 1 --request-rate inf \
|
||||
--percentile-metrics ttft,e2el,audio_rtf,audio_ttfp,audio_duration \
|
||||
--save-result --result-dir ./results
|
||||
```
|
||||
|
||||
#### voice_design (Qwen3-TTS-CustomVoice, bundled seed_tts_design)
|
||||
|
||||
```bash
|
||||
vllm bench serve --omni \
|
||||
--host 127.0.0.1 --port 8000 \
|
||||
--model Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
|
||||
--backend openai-audio-speech \
|
||||
--endpoint /v1/audio/speech \
|
||||
--dataset-name seed-tts-design \
|
||||
--dataset-path benchmarks/build_dataset/seed_tts_design \
|
||||
--seed-tts-locale en \
|
||||
--num-prompts 20 --num-warmups 2 \
|
||||
--extra-body '{"task_type":"VoiceDesign","language":"English"}' \
|
||||
--max-concurrency 1 --request-rate inf \
|
||||
--percentile-metrics ttft,e2el,audio_rtf,audio_ttfp,audio_duration \
|
||||
--save-result --result-dir ./results
|
||||
```
|
||||
|
||||
#### Add WER / SIM / UTMOS to any of the above
|
||||
|
||||
Append `--seed-tts-wer-eval` (and optionally `SEED_TTS_EVAL_DEVICE=cuda:0`
|
||||
in the env, per PR #2558). This triggers the seed-tts-eval protocol:
|
||||
Whisper-large-v3 ASR → WER, WavLM embeddings → SIM, balacoon/utmos → UTMOS.
|
||||
|
||||
### 3. Convenience wrapper (`bench_tts.py`)
|
||||
|
||||
If you're running the **canonical** configuration for a registered model,
|
||||
`bench_tts.py` loads the right defaults from `model_configs.yaml` and
|
||||
emits the exact `vllm bench serve --omni` command above — useful for
|
||||
concurrency sweeps and multi-task runs:
|
||||
|
||||
```bash
|
||||
# Smallest smoke — 5 prompts, concurrency=1
|
||||
python benchmarks/tts/bench_tts.py \
|
||||
--model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
|
||||
--task voice_clone \
|
||||
--dataset-path /path/to/seed-tts-eval \
|
||||
--concurrency 1 --num-prompts 5 \
|
||||
--output-dir ./results
|
||||
|
||||
# Full concurrency sweep
|
||||
python benchmarks/tts/bench_tts.py \
|
||||
--model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
|
||||
--task voice_clone \
|
||||
--dataset-path /path/to/seed-tts-eval \
|
||||
--concurrency 1 2 4 8 16 32 \
|
||||
--num-prompts 20 \
|
||||
--output-dir ./results
|
||||
|
||||
# With WER / SIM / UTMOS quality eval (adds ASR + embedding compute)
|
||||
python benchmarks/tts/bench_tts.py \
|
||||
--model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
|
||||
--task voice_clone \
|
||||
--dataset-path /path/to/seed-tts-eval \
|
||||
--wer-eval \
|
||||
--concurrency 4 --num-prompts 200 \
|
||||
--output-dir ./results
|
||||
```
|
||||
|
||||
### 4. Plot a sweep
|
||||
|
||||
```bash
|
||||
python benchmarks/tts/plot_results.py \
|
||||
--results ./results/*.json \
|
||||
--output ./results/curve.png
|
||||
```
|
||||
|
||||
Outputs TTFP / RTF / throughput curves (and a markdown table) for every
|
||||
`(task, concurrency)` combination in the result set.
|
||||
|
||||
## Task types
|
||||
|
||||
| Task | Dataset | Request body | Checkpoints that support it |
|
||||
|-----------------|-------------------|-----------------------------------------------------|------------------------------------------|
|
||||
| `voice_clone` | `seed-tts` | `ref_audio` + `ref_text` + `task_type=Base` | `Qwen3-TTS-*-Base`, `VoxCPM2` |
|
||||
| `default_voice` | `seed-tts-text` | `voice=Vivian` + `task_type=CustomVoice` | `Qwen3-TTS-*-CustomVoice` |
|
||||
| `voice_design` | `seed-tts-design` | `instructions=<natural-language description>` + `task_type=VoiceDesign` | `Qwen3-TTS-*-CustomVoice` |
|
||||
|
||||
**`-CustomVoice` checkpoints do NOT ship `speaker_encoder` weights**, so
|
||||
voice_clone requests raise `ValueError` at model runtime. Use `-Base` for
|
||||
voice_clone.
|
||||
|
||||
## Adding a new TTS model
|
||||
|
||||
Drop an entry into `model_configs.yaml` — no Python changes required:
|
||||
|
||||
```yaml
|
||||
models:
|
||||
<org>/<model-id>:
|
||||
supported_tasks: [voice_clone] # or default_voice / voice_design
|
||||
backend: openai-audio-speech # vllm bench serve backend
|
||||
endpoint: /v1/audio/speech # OpenAI-compatible endpoint
|
||||
task_extra_body: # merged into every request's body
|
||||
voice_clone:
|
||||
task_type: Base
|
||||
```
|
||||
|
||||
Then add the model's Deploy YAML under `vllm_omni/deploy/<model>.yaml`
|
||||
(Pipeline + Deploy schema) and it's immediately benchable.
|
||||
|
||||
## Datasets
|
||||
|
||||
| Dataset | Bundled? | Format | Source |
|
||||
|--------------------|----------|-------------------|----------------------------------------------------------------|
|
||||
| `seed-tts-design` | ✅ | 5-field meta.lst | `benchmarks/build_dataset/seed_tts_design/en/meta.lst` (20 prompts) |
|
||||
| `seed_tts_smoke` | ✅ | 4-field meta.lst | `benchmarks/build_dataset/seed_tts_smoke/en/meta.lst` (20 text-only) |
|
||||
| `seed-tts` | ❌ | 4-field meta.lst + WAVs | Google-Drive: [BytedanceSpeech/seed-tts-eval][seedtts] (~1.2 GB) |
|
||||
| `seed-tts-text` | ❌ | 4-field meta.lst | Same archive as `seed-tts` (wav column unused) |
|
||||
|
||||
[seedtts]: https://github.com/BytedanceSpeech/seed-tts-eval
|
||||
|
||||
For manual voice_clone / default_voice runs against the full corpus, follow
|
||||
`benchmarks/build_dataset/download_process_data_seedtts.md` and point
|
||||
`--dataset-path` at the extracted `seedtts_testset` directory.
|
||||
|
||||
## DFX nightly CI
|
||||
|
||||
`tests/dfx/perf/tests/test_tts.json` wires three perf regimes plus quality:
|
||||
|
||||
| eval_phase | concurrency | purpose | Baseline metrics |
|
||||
|---------------|-------------|---------------------------------------------------------|-----------------------------------------|
|
||||
| `latency` | 1 | Single-request TTFP / RTF SLO | `median_audio_ttfp_ms`, `median_audio_rtf` |
|
||||
| `throughput` | 8 | Codec-batching cliff sentinel (PDF #272 concurrency≥8) | `median_audio_ttfp_ms`, `median_audio_rtf` |
|
||||
| `quality` | 4 | WER / SIM / UTMOS regression (disabled in CI by default)| `mean_audio_rtf` |
|
||||
|
||||
Why `median_*` for latency/throughput and `mean_*` for quality: latency
|
||||
distributions have cold-start tails that drag the mean; quality aggregates
|
||||
over 200 prompts so single-request outliers don't matter.
|
||||
|
||||
Quality entries are `enabled: false` in CI because seed-tts-eval is not
|
||||
staged in the Buildkite container (matches the precedent in
|
||||
PR #2558 — quality runs are manual / release-validation, not nightly).
|
||||
|
||||
## Concurrency cliff regression sentinel
|
||||
|
||||
Observed on H20-3e, Qwen3-TTS-1.7B (measured pre-merge on this branch):
|
||||
|
||||
| Task | Model | c=1 | c=4 | **c=8** | c=16 | c=32 |
|
||||
|---------------|---------------|--------|--------|------------|--------|--------|
|
||||
| voice_clone | 1.7B-Base | RTF 0.15 / TTFP 165ms | 0.28 / 412ms | **0.49 / 1701ms** | 0.72 / 3355ms | 0.77 / 3772ms |
|
||||
| voice_design | 1.7B-CustomVoice | RTF 0.08 / TTFP 53ms | 0.11 / 154ms | **0.21 / 872ms** | 0.33 / 1801ms | 0.38 / 1989ms |
|
||||
|
||||
Both models show a **4–6× TTFP jump from c=4 to c=8** while audio throughput
|
||||
saturates around c=4–8 — the codec-bs=1 bottleneck documented in
|
||||
vllm-project/vllm-omni#272. The `throughput` CI regime at c=8 is the
|
||||
sentinel for regressions in this area.
|
||||
|
||||
## File layout
|
||||
|
||||
```
|
||||
benchmarks/tts/
|
||||
├── README.md (this file)
|
||||
├── bench_tts.py CLI — serve-mode benchmark driver
|
||||
├── bench_voxcpm_offline.py CLI — offline VoxCPM benchmark (sync + streaming)
|
||||
├── plot_results.py Generate per-task / per-concurrency curves
|
||||
└── model_configs.yaml Model registry (supported tasks + extra body)
|
||||
```
|
||||
|
||||
## Related
|
||||
|
||||
- Upstream seed-tts-eval integration: vllm-project/vllm-omni#2558
|
||||
- Pipeline + Deploy schema: vllm-project/vllm-omni#2383
|
||||
- Concurrency cliff RFC: vllm-project/vllm-omni#272
|
||||
@@ -0,0 +1,308 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Universal TTS benchmark CLI for vllm-omni.
|
||||
|
||||
Runs ``vllm bench serve --omni`` with model-aware defaults loaded from
|
||||
``model_configs.yaml``. Supports Qwen3-TTS, VoxCPM2, and any future TTS
|
||||
model registered in the config file -- no code changes needed to add models.
|
||||
|
||||
Usage::
|
||||
|
||||
python benchmarks/tts/bench_tts.py \\
|
||||
--model Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \\
|
||||
--task voice_clone \\
|
||||
--locale en \\
|
||||
--concurrency 1 4 \\
|
||||
--num-prompts 20 \\
|
||||
--dataset-path /path/to/seed-tts-eval \\
|
||||
--host localhost --port 8000
|
||||
|
||||
See ``--help`` for full option list.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import yaml
|
||||
|
||||
|
||||
def _vllm_omni_bin() -> str:
|
||||
"""Return the vllm-omni (or vllm) binary co-located with the current Python."""
|
||||
bin_dir = Path(sys.executable).parent
|
||||
for candidate in ("vllm-omni", "vllm"):
|
||||
p = bin_dir / candidate
|
||||
if p.is_file():
|
||||
return str(p)
|
||||
return "vllm-omni" # fall back and let the shell resolve it
|
||||
|
||||
|
||||
_REPO_ROOT = Path(__file__).resolve().parent.parent.parent
|
||||
_SCRIPT_DIR = Path(__file__).resolve().parent
|
||||
_DEFAULT_MODEL_CONFIGS = _SCRIPT_DIR / "model_configs.yaml"
|
||||
|
||||
# Maps task name to the dataset_name used with vllm bench serve
|
||||
_TASK_TO_DATASET: dict[str, str] = {
|
||||
"voice_clone": "seed-tts",
|
||||
"default_voice": "seed-tts-text",
|
||||
"voice_design": "seed-tts-design",
|
||||
}
|
||||
|
||||
# Default design dataset path (bundled with the repo)
|
||||
_DEFAULT_DESIGN_DATASET_PATH = str(_REPO_ROOT / "benchmarks" / "build_dataset" / "seed_tts_design")
|
||||
|
||||
|
||||
def load_model_configs(path: Path) -> dict[str, Any]:
|
||||
"""Load model registry from YAML file."""
|
||||
with open(path, encoding="utf-8") as f:
|
||||
data = yaml.safe_load(f)
|
||||
return data.get("models", {})
|
||||
|
||||
|
||||
def build_bench_args(
|
||||
*,
|
||||
host: str,
|
||||
port: int,
|
||||
model: str,
|
||||
task: str,
|
||||
model_cfg: dict[str, Any],
|
||||
locale: str,
|
||||
num_prompts: int,
|
||||
concurrency: int | None,
|
||||
dataset_path: str | None,
|
||||
wer_eval: bool,
|
||||
output_dir: str | None,
|
||||
result_filename: str | None,
|
||||
extra_cli_args: list[str],
|
||||
) -> list[str]:
|
||||
"""Build the ``vllm bench serve --omni`` command for one (task, concurrency) run."""
|
||||
dataset_name = _TASK_TO_DATASET[task]
|
||||
backend: str = model_cfg["backend"]
|
||||
endpoint: str = model_cfg["endpoint"]
|
||||
task_extra_body: dict[str, Any] = (model_cfg.get("task_extra_body") or {}).get(task) or {}
|
||||
|
||||
# Resolve dataset path
|
||||
if dataset_path:
|
||||
resolved_dataset_path = dataset_path
|
||||
elif task == "voice_design":
|
||||
resolved_dataset_path = _DEFAULT_DESIGN_DATASET_PATH
|
||||
else:
|
||||
resolved_dataset_path = None
|
||||
|
||||
cmd = [
|
||||
_vllm_omni_bin(),
|
||||
"bench",
|
||||
"serve",
|
||||
"--omni",
|
||||
"--host",
|
||||
host,
|
||||
"--port",
|
||||
str(port),
|
||||
"--model",
|
||||
model,
|
||||
"--backend",
|
||||
backend,
|
||||
"--endpoint",
|
||||
endpoint,
|
||||
"--dataset-name",
|
||||
dataset_name,
|
||||
"--num-prompts",
|
||||
str(num_prompts),
|
||||
"--num-warmups",
|
||||
"2",
|
||||
"--percentile-metrics",
|
||||
"ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
]
|
||||
|
||||
if resolved_dataset_path:
|
||||
cmd += ["--dataset-path", resolved_dataset_path]
|
||||
|
||||
if locale:
|
||||
cmd += ["--seed-tts-locale", locale]
|
||||
|
||||
if task_extra_body:
|
||||
cmd += ["--extra-body", json.dumps(task_extra_body, separators=(",", ":"))]
|
||||
|
||||
if concurrency is not None:
|
||||
cmd += ["--max-concurrency", str(concurrency), "--request-rate", "inf"]
|
||||
|
||||
if wer_eval:
|
||||
cmd.append("--seed-tts-wer-eval")
|
||||
|
||||
if output_dir or result_filename:
|
||||
out_dir = output_dir or "."
|
||||
os.makedirs(out_dir, exist_ok=True)
|
||||
cmd += ["--save-result", "--result-dir", out_dir]
|
||||
if result_filename:
|
||||
cmd += ["--result-filename", result_filename]
|
||||
|
||||
cmd += extra_cli_args
|
||||
return cmd
|
||||
|
||||
|
||||
def run_one_benchmark(cmd: list[str]) -> dict[str, Any] | None:
|
||||
"""Run a single benchmark subprocess and return parsed JSON result if available."""
|
||||
print(f"\n{'=' * 60}")
|
||||
print("Running:", " ".join(cmd))
|
||||
print("=" * 60)
|
||||
result = subprocess.run(cmd, check=False)
|
||||
if result.returncode != 0:
|
||||
print(f"[bench_tts] WARNING: benchmark exited with code {result.returncode}")
|
||||
return None
|
||||
# If --save-result was used, find the result file
|
||||
try:
|
||||
result_dir_idx = cmd.index("--result-dir")
|
||||
result_dir = Path(cmd[result_dir_idx + 1])
|
||||
if "--result-filename" in cmd:
|
||||
fname_idx = cmd.index("--result-filename")
|
||||
result_file = result_dir / cmd[fname_idx + 1]
|
||||
else:
|
||||
# find most recently modified json
|
||||
jsons = sorted(result_dir.glob("result_*.json"), key=lambda p: p.stat().st_mtime)
|
||||
result_file = jsons[-1] if jsons else None
|
||||
if result_file and result_file.is_file():
|
||||
return json.loads(result_file.read_text(encoding="utf-8"))
|
||||
except (ValueError, IndexError, OSError):
|
||||
pass
|
||||
return None
|
||||
|
||||
|
||||
def print_summary_table(results: list[dict[str, Any]]) -> None:
|
||||
"""Print a unified metrics table across all (task, concurrency) runs."""
|
||||
if not results:
|
||||
return
|
||||
header = (
|
||||
f"{'Task':<16} {'Concurrency':>11} {'RTF mean':>10} "
|
||||
f"{'TTFP (ms)':>10} {'Throughput':>12} {'WER':>7} {'SIM':>7} {'UTMOS':>7}"
|
||||
)
|
||||
print(f"\n{'=' * len(header)}")
|
||||
print("BENCHMARK SUMMARY")
|
||||
print("=" * len(header))
|
||||
print(header)
|
||||
print("-" * len(header))
|
||||
for r in results:
|
||||
task = r.get("_task", "?")
|
||||
conc = r.get("_concurrency", "?")
|
||||
rtf = r.get("mean_audio_rtf", float("nan"))
|
||||
ttfp = r.get("mean_audio_ttfp_ms", float("nan"))
|
||||
throughput = r.get("audio_throughput", float("nan"))
|
||||
wer = r.get("seed_tts_mean_wer", float("nan"))
|
||||
sim = r.get("seed_tts_mean_sim", float("nan"))
|
||||
utmos = r.get("seed_tts_mean_utmos", float("nan"))
|
||||
|
||||
def fmt(v: float, digits: int = 3) -> str:
|
||||
return f"{v:.{digits}f}" if not math.isnan(v) else " n/a"
|
||||
|
||||
print(
|
||||
f"{task:<16} {str(conc):>11} {fmt(rtf):>10} {fmt(ttfp, 0):>10} "
|
||||
f"{fmt(throughput):>12} {fmt(wer):>7} {fmt(sim):>7} {fmt(utmos):>7}"
|
||||
)
|
||||
print("=" * len(header))
|
||||
|
||||
|
||||
def main() -> None:
|
||||
"""Entry point for the universal TTS benchmark CLI."""
|
||||
parser = argparse.ArgumentParser(
|
||||
description=__doc__,
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model", required=True, help="HuggingFace model ID (e.g. Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice)"
|
||||
)
|
||||
parser.add_argument("--task", default="all", help="Task type: voice_clone | default_voice | voice_design | all")
|
||||
parser.add_argument("--locale", default="en", choices=["en", "zh"])
|
||||
parser.add_argument("--concurrency", type=int, nargs="+", default=[1, 4], metavar="N")
|
||||
parser.add_argument(
|
||||
"--num-prompts",
|
||||
type=int,
|
||||
nargs="+",
|
||||
default=[20],
|
||||
metavar="N",
|
||||
help="Number of prompts per run. If one value, applied to all concurrency levels.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dataset-path", default=None, help="Root of seed-tts-eval dataset (required for voice_clone/default_voice)"
|
||||
)
|
||||
parser.add_argument("--wer-eval", action="store_true", help="Enable WER/SIM/UTMOS quality eval")
|
||||
parser.add_argument("--output-dir", default=None, help="Directory to save result JSON files")
|
||||
parser.add_argument("--host", default="localhost")
|
||||
parser.add_argument("--port", type=int, default=8000)
|
||||
parser.add_argument("--model-configs", default=str(_DEFAULT_MODEL_CONFIGS), help="Path to model_configs.yaml")
|
||||
parser.add_argument("extra", nargs=argparse.REMAINDER, help="Extra args passed directly to vllm bench serve")
|
||||
args = parser.parse_args()
|
||||
|
||||
model_configs = load_model_configs(Path(args.model_configs))
|
||||
if args.model not in model_configs:
|
||||
known = "\n ".join(model_configs.keys())
|
||||
print(f"[bench_tts] ERROR: model '{args.model}' not in model_configs.yaml.\nKnown models:\n {known}")
|
||||
sys.exit(1)
|
||||
|
||||
model_cfg = model_configs[args.model]
|
||||
supported_tasks: list[str] = model_cfg.get("supported_tasks", [])
|
||||
|
||||
tasks_to_run: list[str]
|
||||
if args.task == "all":
|
||||
tasks_to_run = supported_tasks
|
||||
elif args.task in supported_tasks:
|
||||
tasks_to_run = [args.task]
|
||||
else:
|
||||
print(
|
||||
f"[bench_tts] ERROR: task '{args.task}' not supported by {args.model}.\nSupported tasks: {supported_tasks}"
|
||||
)
|
||||
sys.exit(1)
|
||||
|
||||
# Align num_prompts list with concurrency list
|
||||
num_prompts_list: list[int] = args.num_prompts
|
||||
if len(num_prompts_list) == 1:
|
||||
num_prompts_list = num_prompts_list * len(args.concurrency)
|
||||
elif len(num_prompts_list) != len(args.concurrency):
|
||||
print(
|
||||
f"[bench_tts] ERROR: --num-prompts ({len(num_prompts_list)} values) must be "
|
||||
f"length 1 or match --concurrency ({len(args.concurrency)} values)."
|
||||
)
|
||||
sys.exit(1)
|
||||
|
||||
all_results: list[dict[str, Any]] = []
|
||||
|
||||
for task in tasks_to_run:
|
||||
for concurrency, num_prompts in zip(args.concurrency, num_prompts_list):
|
||||
ts = datetime.now().strftime("%Y%m%d-%H%M%S")
|
||||
result_filename = f"bench_tts_{args.model.replace('/', '_')}_{task}_c{concurrency}_{ts}.json"
|
||||
cmd = build_bench_args(
|
||||
host=args.host,
|
||||
port=args.port,
|
||||
model=args.model,
|
||||
task=task,
|
||||
model_cfg=model_cfg,
|
||||
locale=args.locale,
|
||||
num_prompts=num_prompts,
|
||||
concurrency=concurrency,
|
||||
dataset_path=args.dataset_path,
|
||||
wer_eval=args.wer_eval,
|
||||
output_dir=args.output_dir,
|
||||
result_filename=result_filename,
|
||||
extra_cli_args=args.extra or [],
|
||||
)
|
||||
result = run_one_benchmark(cmd)
|
||||
if result is not None:
|
||||
result["_task"] = task
|
||||
result["_concurrency"] = concurrency
|
||||
all_results.append(result)
|
||||
# Persist the metadata so plot_results.py can pick it up.
|
||||
if args.output_dir and result_filename:
|
||||
result_path = Path(args.output_dir) / result_filename
|
||||
if result_path.is_file():
|
||||
result_path.write_text(json.dumps(result, indent=2), encoding="utf-8")
|
||||
|
||||
print_summary_table(all_results)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+36
-4
@@ -6,6 +6,27 @@ Supports both:
|
||||
- text-only synthesis
|
||||
- voice cloning
|
||||
- text/clone batch inputs from txt or jsonl
|
||||
|
||||
Usage::
|
||||
|
||||
# Sync (default voice)
|
||||
python benchmarks/tts/bench_voxcpm_offline.py \\
|
||||
--model /path/to/VoxCPM \\
|
||||
--text "Hello world" \\
|
||||
--output-dir results/audio/
|
||||
|
||||
# Streaming (async_chunk)
|
||||
python benchmarks/tts/bench_voxcpm_offline.py \\
|
||||
--model /path/to/VoxCPM \\
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml \\
|
||||
--txt-prompts prompts.txt \\
|
||||
--output-dir results/audio/
|
||||
|
||||
# Voice cloning batch via JSONL
|
||||
python benchmarks/tts/bench_voxcpm_offline.py \\
|
||||
--model /path/to/VoxCPM \\
|
||||
--jsonl-prompts prompts.jsonl \\
|
||||
--output-dir results/audio/
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -26,7 +47,21 @@ from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
from vllm_omni import AsyncOmni, Omni
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[3]
|
||||
|
||||
def _find_repo_root(start: Path) -> Path:
|
||||
"""Walk up from ``start`` until a repo marker is found.
|
||||
|
||||
Falls back to ``parents[2]`` for backwards compatibility if no marker hits
|
||||
(which can only happen in unusual checkouts — the tree should always have
|
||||
pyproject.toml + vllm_omni/ at the top level).
|
||||
"""
|
||||
for candidate in [start, *start.parents]:
|
||||
if (candidate / "pyproject.toml").is_file() and (candidate / "vllm_omni").is_dir():
|
||||
return candidate
|
||||
return start.parents[2]
|
||||
|
||||
|
||||
REPO_ROOT = _find_repo_root(Path(__file__).resolve())
|
||||
DEFAULT_STAGE_ASYNC = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "voxcpm_async_chunk.yaml"
|
||||
DEFAULT_STAGE_SYNC = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "voxcpm.yaml"
|
||||
|
||||
@@ -473,9 +508,6 @@ def parse_args():
|
||||
|
||||
def _is_streaming_stage_config(stage_configs_path: str) -> bool:
|
||||
cfg_name = Path(stage_configs_path).name.lower()
|
||||
# Keep routing purely config-path based:
|
||||
# - voxcpm.yaml => sync
|
||||
# - voxcpm_async_chunk.yaml => streaming
|
||||
return "async_chunk" in cfg_name
|
||||
|
||||
|
||||
@@ -0,0 +1,39 @@
|
||||
# Universal TTS benchmark model registry.
|
||||
# Maps HuggingFace model ID → supported tasks + per-task extra body fields.
|
||||
# To add a new TTS model: add an entry here. No code changes required.
|
||||
#
|
||||
# The server auto-loads its Deploy YAML from vllm_omni/deploy/<model>.yaml via
|
||||
# the Pipeline + Deploy schema introduced in #2383, so no stage_config path
|
||||
# is tracked here.
|
||||
|
||||
models:
|
||||
# -CustomVoice checkpoints lack speaker_encoder weights, so voice_clone is
|
||||
# NOT supported (an attempt raises ValueError from _extract_speaker_embedding
|
||||
# at model runtime). Use -Base for voice_clone.
|
||||
Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice:
|
||||
supported_tasks: [default_voice, voice_design]
|
||||
backend: openai-audio-speech
|
||||
endpoint: /v1/audio/speech
|
||||
task_extra_body:
|
||||
default_voice:
|
||||
voice: Vivian
|
||||
language: English
|
||||
task_type: CustomVoice
|
||||
voice_design:
|
||||
task_type: VoiceDesign
|
||||
language: English
|
||||
|
||||
Qwen/Qwen3-TTS-12Hz-1.7B-Base:
|
||||
supported_tasks: [voice_clone]
|
||||
backend: openai-audio-speech
|
||||
endpoint: /v1/audio/speech
|
||||
task_extra_body:
|
||||
voice_clone:
|
||||
task_type: Base
|
||||
|
||||
openbmb/VoxCPM2:
|
||||
supported_tasks: [voice_clone]
|
||||
backend: openai-audio-speech
|
||||
endpoint: /v1/audio/speech
|
||||
task_extra_body:
|
||||
voice_clone: {}
|
||||
@@ -0,0 +1,324 @@
|
||||
"""Plot universal TTS benchmark results.
|
||||
|
||||
Reads JSON files saved by ``bench_tts.py`` (via ``vllm bench serve --omni``)
|
||||
and generates comparison bar charts grouped by task type.
|
||||
|
||||
Metrics plotted:
|
||||
- AUDIO_TTFP (mean audio time-to-first-packet, ms)
|
||||
- E2EL (mean end-to-end latency, ms)
|
||||
- Audio RTF (mean real-time factor)
|
||||
- Audio throughput (audio-seconds / wall-second)
|
||||
|
||||
Quality metrics (WER / SIM / UTMOS) are printed in a table when present.
|
||||
|
||||
Usage::
|
||||
|
||||
# Single run — one JSON per task, all in results/
|
||||
python benchmarks/tts/plot_results.py \\
|
||||
--results results/bench_tts_*.json \\
|
||||
--output results/tts_benchmark.png
|
||||
|
||||
# Compare two runs (e.g. async_chunk on vs off)
|
||||
python benchmarks/tts/plot_results.py \\
|
||||
--results run_a/bench_tts_*.json \\
|
||||
--results run_b/bench_tts_*.json \\
|
||||
--labels "async_chunk_on" "async_chunk_off" \\
|
||||
--output results/comparison.png
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import math
|
||||
from pathlib import Path
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# JSON loading
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def load_run(paths: list[str]) -> list[dict]:
|
||||
"""Load and merge all JSON files for one run into a flat list of records.
|
||||
|
||||
Each record is expected to have at least ``_concurrency`` (int) and
|
||||
``_task`` (str) keys injected by ``bench_tts.py``. Records that come
|
||||
from a file that contains a list are flattened.
|
||||
"""
|
||||
records: list[dict] = []
|
||||
for p in paths:
|
||||
raw = json.loads(Path(p).read_text(encoding="utf-8"))
|
||||
if isinstance(raw, list):
|
||||
records.extend(raw)
|
||||
elif isinstance(raw, dict):
|
||||
records.append(raw)
|
||||
return records
|
||||
|
||||
|
||||
def _get(record: dict, key: str) -> float:
|
||||
v = record.get(key, float("nan"))
|
||||
if v is None or (isinstance(v, float) and math.isnan(v)):
|
||||
return float("nan")
|
||||
try:
|
||||
return float(v)
|
||||
except (TypeError, ValueError):
|
||||
return float("nan")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Plotting helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _bar_group(
|
||||
ax: plt.Axes,
|
||||
x: np.ndarray,
|
||||
data_per_label: dict[str, list[float]],
|
||||
width: float,
|
||||
colors: list[str],
|
||||
ylabel: str,
|
||||
title: str,
|
||||
concurrency_labels: list[str],
|
||||
fmt: str = ".1f",
|
||||
) -> None:
|
||||
n = len(data_per_label)
|
||||
offsets = np.linspace(-(n - 1) * width / 2, (n - 1) * width / 2, n) if n > 1 else [0.0]
|
||||
|
||||
for i, (label, values) in enumerate(data_per_label.items()):
|
||||
plot_vals = [0.0 if math.isnan(v) else v for v in values]
|
||||
bar = ax.bar(x + offsets[i], plot_vals, width, label=label, color=colors[i % len(colors)], alpha=0.85)
|
||||
max_val = max((v for v in values if not math.isnan(v)), default=1.0)
|
||||
for rect, val in zip(bar, values):
|
||||
if not math.isnan(val) and val > 0:
|
||||
ax.text(
|
||||
rect.get_x() + rect.get_width() / 2,
|
||||
rect.get_height() + max_val * 0.02,
|
||||
f"{val:{fmt}}",
|
||||
ha="center",
|
||||
va="bottom",
|
||||
fontsize=8,
|
||||
fontweight="bold",
|
||||
)
|
||||
|
||||
ax.set_xlabel("Concurrency", fontsize=11)
|
||||
ax.set_ylabel(ylabel, fontsize=11)
|
||||
ax.set_title(title, fontsize=12, fontweight="bold")
|
||||
ax.set_xticks(x)
|
||||
ax.set_xticklabels(concurrency_labels)
|
||||
ax.legend(fontsize=9)
|
||||
ax.grid(axis="y", alpha=0.3)
|
||||
ax.set_axisbelow(True)
|
||||
|
||||
|
||||
COLORS = ["#2196F3", "#FF5722", "#4CAF50", "#FFC107", "#9C27B0"]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Comparison plot (multiple labels / runs)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def plot_comparison(
|
||||
all_runs: list[list[dict]],
|
||||
labels: list[str],
|
||||
output_path: str,
|
||||
task_filter: str | None = None,
|
||||
title_prefix: str = "TTS",
|
||||
) -> None:
|
||||
"""One 2×2 subplot per task found in the data."""
|
||||
# Determine tasks to plot
|
||||
tasks: list[str] = []
|
||||
for run in all_runs:
|
||||
for r in run:
|
||||
t = r.get("_task", "unknown")
|
||||
if t not in tasks:
|
||||
tasks.append(t)
|
||||
if task_filter:
|
||||
tasks = [t for t in tasks if t == task_filter]
|
||||
|
||||
n_tasks = len(tasks)
|
||||
if n_tasks == 0:
|
||||
print("[plot_results] No tasks found in data.")
|
||||
return
|
||||
|
||||
fig, axes_grid = plt.subplots(n_tasks, 4, figsize=(18, 4.5 * n_tasks))
|
||||
fig.suptitle(f"{title_prefix} Benchmark", fontsize=15, fontweight="bold")
|
||||
|
||||
# Ensure axes_grid is always 2D
|
||||
if n_tasks == 1:
|
||||
axes_grid = [axes_grid]
|
||||
|
||||
for row_idx, task in enumerate(tasks):
|
||||
# Collect concurrencies across all runs for this task
|
||||
all_concs: set[int] = set()
|
||||
for run in all_runs:
|
||||
for r in run:
|
||||
if r.get("_task") == task:
|
||||
c = r.get("_concurrency")
|
||||
if c is not None:
|
||||
all_concs.add(int(c))
|
||||
concurrencies = sorted(all_concs)
|
||||
x = np.arange(len(concurrencies))
|
||||
conc_labels = [str(c) for c in concurrencies]
|
||||
|
||||
def _series(run: list[dict], metric_key: str) -> list[float]:
|
||||
conc_map = {int(r["_concurrency"]): r for r in run if r.get("_task") == task and "_concurrency" in r}
|
||||
return [_get(conc_map.get(c, {}), metric_key) for c in concurrencies]
|
||||
|
||||
metrics = [
|
||||
("mean_audio_ttfp_ms", "TTFP (ms)", "Time-to-First-Packet", ".0f"),
|
||||
("mean_e2el_ms", "E2E Latency (ms)", "End-to-End Latency", ".0f"),
|
||||
("mean_audio_rtf", "RTF", "Real-Time Factor (RTF)", ".3f"),
|
||||
("audio_throughput", "audio-s / wall-s", "Audio Throughput", ".2f"),
|
||||
]
|
||||
|
||||
axes_row = axes_grid[row_idx]
|
||||
for col_idx, (key, ylabel, subtitle, fmt) in enumerate(metrics):
|
||||
data_per_label = {lbl: _series(run, key) for lbl, run in zip(labels, all_runs)}
|
||||
_bar_group(
|
||||
axes_row[col_idx],
|
||||
x,
|
||||
data_per_label,
|
||||
width=0.3 if len(labels) > 1 else 0.5,
|
||||
colors=COLORS,
|
||||
ylabel=ylabel,
|
||||
title=f"{task} — {subtitle}",
|
||||
concurrency_labels=conc_labels,
|
||||
fmt=fmt,
|
||||
)
|
||||
|
||||
plt.tight_layout()
|
||||
Path(output_path).parent.mkdir(parents=True, exist_ok=True)
|
||||
plt.savefig(output_path, dpi=150, bbox_inches="tight")
|
||||
print(f"Plot saved to {output_path}")
|
||||
plt.close()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Markdown comparison table
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def print_comparison_table(all_runs: list[list[dict]], labels: list[str]) -> None:
|
||||
tasks: list[str] = []
|
||||
for run in all_runs:
|
||||
for r in run:
|
||||
t = r.get("_task", "unknown")
|
||||
if t not in tasks:
|
||||
tasks.append(t)
|
||||
|
||||
perf_metrics = [
|
||||
("TTFP (ms)", "mean_audio_ttfp_ms", ".1f"),
|
||||
("E2E (ms)", "mean_e2el_ms", ".1f"),
|
||||
("RTF", "mean_audio_rtf", ".3f"),
|
||||
("Throughput (a-s/s)", "audio_throughput", ".2f"),
|
||||
]
|
||||
quality_metrics = [
|
||||
("WER (%)", "seed_tts_mean_wer", ".1f"),
|
||||
("SIM", "seed_tts_mean_sim", ".3f"),
|
||||
("UTMOS", "seed_tts_mean_utmos", ".2f"),
|
||||
]
|
||||
|
||||
for task in tasks:
|
||||
all_concs: set[int] = set()
|
||||
for run in all_runs:
|
||||
for r in run:
|
||||
if r.get("_task") == task:
|
||||
c = r.get("_concurrency")
|
||||
if c is not None:
|
||||
all_concs.add(int(c))
|
||||
concurrencies = sorted(all_concs)
|
||||
|
||||
print(f"\n## {task}\n")
|
||||
col_header = "| Metric | Concurrency |" + "".join(f" {lbl} |" for lbl in labels)
|
||||
sep = "| --- | --- |" + " --- |" * len(labels)
|
||||
print(col_header)
|
||||
print(sep)
|
||||
|
||||
for metric, key, fmt in perf_metrics + quality_metrics:
|
||||
for c in concurrencies:
|
||||
row = f"| {metric} | {c} |"
|
||||
for run in all_runs:
|
||||
conc_map = {
|
||||
int(r["_concurrency"]): r for r in run if r.get("_task") == task and "_concurrency" in r
|
||||
}
|
||||
val = _get(conc_map.get(c, {}), key)
|
||||
row += f" {val:{fmt}} |" if not math.isnan(val) else " n/a |"
|
||||
print(row)
|
||||
|
||||
# Improvement column (2-run comparison only)
|
||||
if len(all_runs) == 2:
|
||||
print(f"\n### Improvement ({labels[0]} vs {labels[1]})\n")
|
||||
print("| Metric | Concurrency | Change |")
|
||||
print("| --- | --- | --- |")
|
||||
for metric, key, _ in perf_metrics:
|
||||
for c in concurrencies:
|
||||
conc_map0 = {
|
||||
int(r["_concurrency"]): r for r in all_runs[0] if r.get("_task") == task and "_concurrency" in r
|
||||
}
|
||||
conc_map1 = {
|
||||
int(r["_concurrency"]): r for r in all_runs[1] if r.get("_task") == task and "_concurrency" in r
|
||||
}
|
||||
v0 = _get(conc_map0.get(c, {}), key)
|
||||
v1 = _get(conc_map1.get(c, {}), key)
|
||||
if not math.isnan(v0) and not math.isnan(v1) and v1 > 0:
|
||||
pct = (v1 - v0) / v1 * 100
|
||||
print(f"| {metric} | {c} | {pct:+.1f}% |")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CLI
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description=__doc__,
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--results",
|
||||
type=str,
|
||||
nargs="+",
|
||||
action="append",
|
||||
required=True,
|
||||
metavar="FILE",
|
||||
help="JSON result file(s) for one run. Repeat --results for multiple runs to compare.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--labels",
|
||||
type=str,
|
||||
nargs="+",
|
||||
default=None,
|
||||
help="Label for each --results group (must match the number of --results groups).",
|
||||
)
|
||||
parser.add_argument("--output", type=str, default="results/tts_benchmark.png", help="Output image path.")
|
||||
parser.add_argument("--title", type=str, default="TTS", help="Title prefix for the plot.")
|
||||
parser.add_argument("--task", type=str, default=None, help="Filter to a single task (e.g. voice_clone).")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
|
||||
# args.results is a list-of-lists due to action="append"
|
||||
all_runs: list[list[dict]] = [load_run(group) for group in args.results]
|
||||
n_runs = len(all_runs)
|
||||
|
||||
labels: list[str]
|
||||
if args.labels:
|
||||
if len(args.labels) != n_runs:
|
||||
raise SystemExit(f"--labels count ({len(args.labels)}) must match --results groups ({n_runs})")
|
||||
labels = args.labels
|
||||
else:
|
||||
labels = [f"run{i + 1}" for i in range(n_runs)]
|
||||
|
||||
print_comparison_table(all_runs, labels)
|
||||
plot_comparison(all_runs, labels, args.output, task_filter=args.task, title_prefix=args.title)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,119 +0,0 @@
|
||||
# VoxCPM Benchmark
|
||||
|
||||
This directory contains both:
|
||||
|
||||
- online serving benchmark through the OpenAI-compatible `/v1/audio/speech` API
|
||||
- offline benchmark for `Omni` / `AsyncOmni`
|
||||
- full offline smoke-matrix orchestration
|
||||
|
||||
Both benchmark paths report:
|
||||
|
||||
- TTFP: time to first PCM packet
|
||||
- E2E latency
|
||||
- RTF: real-time factor (`e2e / audio_duration`)
|
||||
|
||||
## Offline Benchmark
|
||||
|
||||
Single offline benchmark run:
|
||||
|
||||
```bash
|
||||
python benchmarks/voxcpm/vllm_omni/bench_tts_offline.py \
|
||||
--model /path/to/voxcpm-model \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm.yaml \
|
||||
--text "This is a split-stage VoxCPM synthesis example running on vLLM Omni." \
|
||||
--warmup-runs 1 \
|
||||
--output-dir benchmarks/voxcpm/results/offline_single
|
||||
```
|
||||
|
||||
Streaming offline benchmark:
|
||||
|
||||
```bash
|
||||
python benchmarks/voxcpm/vllm_omni/bench_tts_offline.py \
|
||||
--model /path/to/voxcpm-model \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml \
|
||||
--text "This is a split-stage VoxCPM streaming example running on vLLM Omni." \
|
||||
--warmup-runs 1 \
|
||||
--output-dir benchmarks/voxcpm/results/offline_streaming
|
||||
```
|
||||
|
||||
Full fixed offline matrix, equivalent to the old `examples/offline_inference/voxcpm/test.py`:
|
||||
|
||||
```bash
|
||||
python benchmarks/voxcpm/vllm_omni/run_offline_matrix.py \
|
||||
--model /path/to/voxcpm-model \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--ref-text "The exact transcript spoken in reference.wav." \
|
||||
--output-root benchmarks/voxcpm/results/offline_matrix
|
||||
```
|
||||
|
||||
The full matrix covers both routes:
|
||||
|
||||
- streaming: `voxcpm_async_chunk.yaml`
|
||||
- sync: `voxcpm.yaml`
|
||||
|
||||
And these six scenarios under each route:
|
||||
|
||||
- warmup + single TTS
|
||||
- warmup + single voice cloning
|
||||
- warmup + batch TTS
|
||||
- warmup + batch voice cloning
|
||||
- cold single TTS
|
||||
- cold single voice cloning
|
||||
|
||||
`bench_tts_offline.py` itself no longer writes `summary.json` / `results.json`; it prints TTFP / RTF inline and saves generated WAV files only. The matrix runner keeps only per-case `run.log`.
|
||||
|
||||
## Start the Server
|
||||
|
||||
Async-chunk:
|
||||
|
||||
```bash
|
||||
vllm serve /path/to/voxcpm-model \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml \
|
||||
--trust-remote-code \
|
||||
--enforce-eager \
|
||||
--omni \
|
||||
--port 8091
|
||||
```
|
||||
|
||||
Non-streaming:
|
||||
|
||||
```bash
|
||||
vllm serve /path/to/voxcpm-model \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm.yaml \
|
||||
--trust-remote-code \
|
||||
--enforce-eager \
|
||||
--omni \
|
||||
--port 8091
|
||||
```
|
||||
|
||||
## Run the Benchmark
|
||||
|
||||
```bash
|
||||
python benchmarks/voxcpm/vllm_omni/bench_tts_serve.py \
|
||||
--host 127.0.0.1 \
|
||||
--port 8091 \
|
||||
--num-prompts 20 \
|
||||
--max-concurrency 1 \
|
||||
--result-dir /tmp/voxcpm_bench
|
||||
```
|
||||
|
||||
Voice cloning benchmark:
|
||||
|
||||
```bash
|
||||
python benchmarks/voxcpm/vllm_omni/bench_tts_serve.py \
|
||||
--host 127.0.0.1 \
|
||||
--port 8091 \
|
||||
--num-prompts 10 \
|
||||
--max-concurrency 1 \
|
||||
--ref-audio https://example.com/reference.wav \
|
||||
--ref-text "The exact transcript spoken in the reference audio." \
|
||||
--result-dir /tmp/voxcpm_clone_bench
|
||||
```
|
||||
|
||||
## Notes
|
||||
|
||||
- The benchmark uses `stream=true` and `response_format=pcm` so TTFP is measured from the first audio packet.
|
||||
- `RTF < 1.0` means the server generates audio faster than real time.
|
||||
- For `voxcpm_async_chunk.yaml`, keep concurrency at `1`. This matches native VoxCPM streaming more closely.
|
||||
- Do not benchmark concurrent online streaming on `voxcpm_async_chunk.yaml`; use `voxcpm.yaml` for multi-request throughput runs.
|
||||
- For the offline matrix mode, `--ref-audio` and `--ref-text` are required because clone cases are part of the fixed coverage set.
|
||||
@@ -1,283 +0,0 @@
|
||||
"""Benchmark VoxCPM via /v1/audio/speech.
|
||||
|
||||
Reports TTFP (time to first packet), E2E latency, and RTF (real-time factor).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import time
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
import aiohttp
|
||||
import numpy as np
|
||||
from tqdm.asyncio import tqdm
|
||||
|
||||
DEFAULT_MODEL = "OpenBMB/VoxCPM1.5"
|
||||
DEFAULT_SAMPLE_RATE = 24000
|
||||
PROMPTS = [
|
||||
"Hello, welcome to the VoxCPM speech benchmark.",
|
||||
"This is a short benchmark prompt for online text-to-speech generation.",
|
||||
"The quick brown fox jumps over the lazy dog near the riverbank.",
|
||||
"Please remember to bring your identification documents tomorrow morning.",
|
||||
"Learning a new language takes patience, practice, and curiosity.",
|
||||
"This benchmark reports TTFP and RTF for the VoxCPM online serving path.",
|
||||
]
|
||||
|
||||
|
||||
@dataclass
|
||||
class RequestResult:
|
||||
success: bool = False
|
||||
ttfp: float = 0.0
|
||||
e2e: float = 0.0
|
||||
audio_bytes: int = 0
|
||||
audio_duration: float = 0.0
|
||||
rtf: float = 0.0
|
||||
prompt: str = ""
|
||||
error: str = ""
|
||||
|
||||
|
||||
@dataclass
|
||||
class BenchmarkResult:
|
||||
concurrency: int = 0
|
||||
num_prompts: int = 0
|
||||
completed: int = 0
|
||||
failed: int = 0
|
||||
duration_s: float = 0.0
|
||||
mean_ttfp_ms: float = 0.0
|
||||
median_ttfp_ms: float = 0.0
|
||||
p95_ttfp_ms: float = 0.0
|
||||
mean_e2e_ms: float = 0.0
|
||||
median_e2e_ms: float = 0.0
|
||||
p95_e2e_ms: float = 0.0
|
||||
mean_rtf: float = 0.0
|
||||
median_rtf: float = 0.0
|
||||
p95_rtf: float = 0.0
|
||||
total_audio_duration_s: float = 0.0
|
||||
request_throughput: float = 0.0
|
||||
per_request: list[dict[str, float | str]] = field(default_factory=list)
|
||||
|
||||
|
||||
def pcm_bytes_to_duration(num_bytes: int, sample_rate: int = DEFAULT_SAMPLE_RATE, sample_width: int = 2) -> float:
|
||||
num_samples = num_bytes / sample_width
|
||||
return num_samples / sample_rate
|
||||
|
||||
|
||||
async def send_tts_request(
|
||||
session: aiohttp.ClientSession,
|
||||
api_url: str,
|
||||
*,
|
||||
model: str,
|
||||
prompt: str,
|
||||
ref_audio: str | None,
|
||||
ref_text: str | None,
|
||||
pbar: tqdm | None = None,
|
||||
) -> RequestResult:
|
||||
payload: dict[str, object] = {
|
||||
"model": model,
|
||||
"input": prompt,
|
||||
"stream": True,
|
||||
"response_format": "pcm",
|
||||
}
|
||||
if ref_audio is not None:
|
||||
payload["ref_audio"] = ref_audio
|
||||
if ref_text is not None:
|
||||
payload["ref_text"] = ref_text
|
||||
|
||||
result = RequestResult(prompt=prompt)
|
||||
started_at = time.perf_counter()
|
||||
|
||||
try:
|
||||
async with session.post(api_url, json=payload) as response:
|
||||
if response.status != 200:
|
||||
result.error = f"HTTP {response.status}: {await response.text()}"
|
||||
return result
|
||||
|
||||
first_chunk = True
|
||||
total_bytes = 0
|
||||
async for chunk in response.content.iter_any():
|
||||
if not chunk:
|
||||
continue
|
||||
if first_chunk:
|
||||
result.ttfp = time.perf_counter() - started_at
|
||||
first_chunk = False
|
||||
total_bytes += len(chunk)
|
||||
|
||||
result.e2e = time.perf_counter() - started_at
|
||||
result.audio_bytes = total_bytes
|
||||
result.audio_duration = pcm_bytes_to_duration(total_bytes)
|
||||
if result.audio_duration > 0:
|
||||
result.rtf = result.e2e / result.audio_duration
|
||||
result.success = True
|
||||
except Exception as e:
|
||||
result.error = str(e)
|
||||
result.e2e = time.perf_counter() - started_at
|
||||
|
||||
if pbar is not None:
|
||||
pbar.update(1)
|
||||
return result
|
||||
|
||||
|
||||
async def run_benchmark(
|
||||
*,
|
||||
host: str,
|
||||
port: int,
|
||||
model: str,
|
||||
num_prompts: int,
|
||||
max_concurrency: int,
|
||||
num_warmups: int,
|
||||
ref_audio: str | None,
|
||||
ref_text: str | None,
|
||||
) -> BenchmarkResult:
|
||||
api_url = f"http://{host}:{port}/v1/audio/speech"
|
||||
connector = aiohttp.TCPConnector(limit=max_concurrency, limit_per_host=max_concurrency, keepalive_timeout=60)
|
||||
timeout = aiohttp.ClientTimeout(total=600)
|
||||
|
||||
async with aiohttp.ClientSession(connector=connector, timeout=timeout) as session:
|
||||
if num_warmups > 0:
|
||||
print(f" Warming up with {num_warmups} requests...")
|
||||
warmup_tasks = [
|
||||
send_tts_request(
|
||||
session,
|
||||
api_url,
|
||||
model=model,
|
||||
prompt=PROMPTS[i % len(PROMPTS)],
|
||||
ref_audio=ref_audio,
|
||||
ref_text=ref_text,
|
||||
)
|
||||
for i in range(num_warmups)
|
||||
]
|
||||
await asyncio.gather(*warmup_tasks)
|
||||
print(" Warmup done.")
|
||||
|
||||
request_prompts = [PROMPTS[i % len(PROMPTS)] for i in range(num_prompts)]
|
||||
semaphore = asyncio.Semaphore(max_concurrency)
|
||||
pbar = tqdm(total=num_prompts, desc=f" concurrency={max_concurrency}")
|
||||
|
||||
async def limited_request(prompt: str) -> RequestResult:
|
||||
async with semaphore:
|
||||
return await send_tts_request(
|
||||
session,
|
||||
api_url,
|
||||
model=model,
|
||||
prompt=prompt,
|
||||
ref_audio=ref_audio,
|
||||
ref_text=ref_text,
|
||||
pbar=pbar,
|
||||
)
|
||||
|
||||
started_at = time.perf_counter()
|
||||
results = await asyncio.gather(*[asyncio.create_task(limited_request(prompt)) for prompt in request_prompts])
|
||||
duration = time.perf_counter() - started_at
|
||||
pbar.close()
|
||||
|
||||
succeeded = [result for result in results if result.success]
|
||||
bench = BenchmarkResult(
|
||||
concurrency=max_concurrency,
|
||||
num_prompts=num_prompts,
|
||||
completed=len(succeeded),
|
||||
failed=len(results) - len(succeeded),
|
||||
duration_s=duration,
|
||||
)
|
||||
|
||||
if not succeeded:
|
||||
return bench
|
||||
|
||||
ttfps = np.array([result.ttfp * 1000 for result in succeeded], dtype=np.float64)
|
||||
e2es = np.array([result.e2e * 1000 for result in succeeded], dtype=np.float64)
|
||||
rtfs = np.array([result.rtf for result in succeeded], dtype=np.float64)
|
||||
audio_durations = np.array([result.audio_duration for result in succeeded], dtype=np.float64)
|
||||
|
||||
bench.mean_ttfp_ms = float(np.mean(ttfps))
|
||||
bench.median_ttfp_ms = float(np.median(ttfps))
|
||||
bench.p95_ttfp_ms = float(np.percentile(ttfps, 95))
|
||||
bench.mean_e2e_ms = float(np.mean(e2es))
|
||||
bench.median_e2e_ms = float(np.median(e2es))
|
||||
bench.p95_e2e_ms = float(np.percentile(e2es, 95))
|
||||
bench.mean_rtf = float(np.mean(rtfs))
|
||||
bench.median_rtf = float(np.median(rtfs))
|
||||
bench.p95_rtf = float(np.percentile(rtfs, 95))
|
||||
bench.total_audio_duration_s = float(np.sum(audio_durations))
|
||||
bench.request_throughput = len(succeeded) / duration if duration > 0 else 0.0
|
||||
bench.per_request = [
|
||||
{
|
||||
"prompt": result.prompt,
|
||||
"ttfp_ms": result.ttfp * 1000,
|
||||
"e2e_ms": result.e2e * 1000,
|
||||
"rtf": result.rtf,
|
||||
"audio_duration_s": result.audio_duration,
|
||||
}
|
||||
for result in succeeded
|
||||
]
|
||||
|
||||
return bench
|
||||
|
||||
|
||||
def print_summary(result: BenchmarkResult) -> None:
|
||||
width = 54
|
||||
print("")
|
||||
print("=" * width)
|
||||
print(f"{'VoxCPM Serving Benchmark':^{width}}")
|
||||
print("=" * width)
|
||||
print(f"concurrency : {result.concurrency}")
|
||||
print(f"requests : {result.completed}/{result.num_prompts} succeeded")
|
||||
print(f"wall time (s) : {result.duration_s:.3f}")
|
||||
print(f"mean TTFP (ms) : {result.mean_ttfp_ms:.2f}")
|
||||
print(f"p95 TTFP (ms) : {result.p95_ttfp_ms:.2f}")
|
||||
print(f"mean E2E (ms) : {result.mean_e2e_ms:.2f}")
|
||||
print(f"p95 E2E (ms) : {result.p95_e2e_ms:.2f}")
|
||||
print(f"mean RTF : {result.mean_rtf:.3f}")
|
||||
print(f"p95 RTF : {result.p95_rtf:.3f}")
|
||||
print(f"request throughput : {result.request_throughput:.2f} req/s")
|
||||
print("=" * width)
|
||||
|
||||
|
||||
async def main_async(args) -> None:
|
||||
result_dir = Path(args.result_dir)
|
||||
result_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
all_results: list[BenchmarkResult] = []
|
||||
for concurrency in args.max_concurrency:
|
||||
result = await run_benchmark(
|
||||
host=args.host,
|
||||
port=args.port,
|
||||
model=args.model,
|
||||
num_prompts=args.num_prompts,
|
||||
max_concurrency=concurrency,
|
||||
num_warmups=args.num_warmups,
|
||||
ref_audio=args.ref_audio,
|
||||
ref_text=args.ref_text,
|
||||
)
|
||||
print_summary(result)
|
||||
all_results.append(result)
|
||||
|
||||
payload = {
|
||||
"model": args.model,
|
||||
"created_at": datetime.utcnow().isoformat() + "Z",
|
||||
"results": [asdict(result) for result in all_results],
|
||||
}
|
||||
result_path = result_dir / "bench_tts_serve.json"
|
||||
result_path.write_text(json.dumps(payload, indent=2), encoding="utf-8")
|
||||
print(f"Saved results to: {result_path}")
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Benchmark VoxCPM via /v1/audio/speech")
|
||||
parser.add_argument("--host", default="127.0.0.1", help="Server host")
|
||||
parser.add_argument("--port", type=int, default=8091, help="Server port")
|
||||
parser.add_argument("--model", default=DEFAULT_MODEL, help="Model name or path")
|
||||
parser.add_argument("--num-prompts", type=int, default=20, help="Number of prompts to send")
|
||||
parser.add_argument("--max-concurrency", type=int, nargs="+", default=[1], help="Concurrency levels to benchmark")
|
||||
parser.add_argument("--num-warmups", type=int, default=3, help="Warmup request count")
|
||||
parser.add_argument("--ref-audio", default=None, help="Reference audio URL or data URL for voice cloning")
|
||||
parser.add_argument("--ref-text", default=None, help="Reference audio transcript for voice cloning")
|
||||
parser.add_argument("--result-dir", default="results", help="Directory to save benchmark JSON")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main_async(parse_args()))
|
||||
@@ -1,303 +0,0 @@
|
||||
"""Run the full offline VoxCPM smoke matrix.
|
||||
|
||||
This script keeps the old `test.py` coverage, but delegates each case to
|
||||
`bench_tts_offline.py` so the benchmark runner itself stays focused on a
|
||||
single execution path.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import shlex
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[3]
|
||||
BENCH_SCRIPT = Path(__file__).with_name("bench_tts_offline.py")
|
||||
DEFAULT_STAGE_ASYNC = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "voxcpm_async_chunk.yaml"
|
||||
DEFAULT_STAGE_SYNC = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "voxcpm.yaml"
|
||||
DEFAULT_OUTPUT_ROOT = BENCH_SCRIPT.parents[1] / "results" / "offline_matrix"
|
||||
|
||||
SINGLE_TTS_TEXT = "This is a single text-to-speech smoke test for VoxCPM on vLLM Omni."
|
||||
SINGLE_CLONE_TEXT = "This sentence is synthesized with the cloned voice for validation."
|
||||
BATCH_TTS_TEXTS = [
|
||||
"The first batch text-to-speech sample validates sequential batch execution.",
|
||||
"The second batch text-to-speech sample checks another prompt in the same file.",
|
||||
"The third batch text-to-speech sample completes the sequential batch path.",
|
||||
]
|
||||
BATCH_CLONE_TEXTS = [
|
||||
"The first cloned sample validates sequential batch voice cloning.",
|
||||
"The second cloned sample checks the same reference voice on another prompt.",
|
||||
"The third cloned sample finishes the shared-reference clone batch path.",
|
||||
]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ModeSpec:
|
||||
name: str
|
||||
stage_config: Path
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class CaseSpec:
|
||||
name: str
|
||||
warmup_runs: int
|
||||
prompt_kind: str
|
||||
voice_clone: bool
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class CaseResult:
|
||||
mode: str
|
||||
case: str
|
||||
returncode: int
|
||||
elapsed_s: float
|
||||
output_dir: Path
|
||||
log_path: Path
|
||||
|
||||
@property
|
||||
def ok(self) -> bool:
|
||||
return self.returncode == 0
|
||||
|
||||
|
||||
MODE_SPECS = [
|
||||
ModeSpec(name="streaming", stage_config=DEFAULT_STAGE_ASYNC),
|
||||
ModeSpec(name="sync", stage_config=DEFAULT_STAGE_SYNC),
|
||||
]
|
||||
|
||||
CASE_SPECS = [
|
||||
CaseSpec(name="warmup_single_tts", warmup_runs=1, prompt_kind="single", voice_clone=False),
|
||||
CaseSpec(name="warmup_single_clone", warmup_runs=1, prompt_kind="single", voice_clone=True),
|
||||
CaseSpec(name="warmup_batch_tts", warmup_runs=1, prompt_kind="batch", voice_clone=False),
|
||||
CaseSpec(name="warmup_batch_clone", warmup_runs=1, prompt_kind="batch", voice_clone=True),
|
||||
CaseSpec(name="cold_single_tts", warmup_runs=0, prompt_kind="single", voice_clone=False),
|
||||
CaseSpec(name="cold_single_clone", warmup_runs=0, prompt_kind="single", voice_clone=True),
|
||||
]
|
||||
|
||||
|
||||
def _write_lines(path: Path, lines: list[str]) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text("\n".join(lines) + "\n", encoding="utf-8")
|
||||
|
||||
|
||||
def _prepare_batch_inputs(output_root: Path) -> tuple[Path, Path]:
|
||||
input_dir = output_root / "inputs"
|
||||
batch_tts_path = input_dir / "batch_tts_prompts.txt"
|
||||
batch_clone_path = input_dir / "batch_clone_prompts.txt"
|
||||
_write_lines(batch_tts_path, BATCH_TTS_TEXTS)
|
||||
_write_lines(batch_clone_path, BATCH_CLONE_TEXTS)
|
||||
return batch_tts_path, batch_clone_path
|
||||
|
||||
|
||||
def _base_command(args, mode: ModeSpec, output_dir: Path) -> list[str]:
|
||||
cmd = [
|
||||
args.python,
|
||||
str(BENCH_SCRIPT),
|
||||
"--model",
|
||||
args.model,
|
||||
"--stage-configs-path",
|
||||
str(mode.stage_config),
|
||||
"--output-dir",
|
||||
str(output_dir),
|
||||
"--num-runs",
|
||||
str(args.num_runs),
|
||||
"--stage-init-timeout",
|
||||
str(args.stage_init_timeout),
|
||||
]
|
||||
cmd.append("--log-stats" if args.log_stats else "--no-log-stats")
|
||||
cmd.extend(["--cfg-value", str(args.cfg_value)])
|
||||
cmd.extend(["--inference-timesteps", str(args.inference_timesteps)])
|
||||
cmd.extend(["--min-len", str(args.min_len)])
|
||||
cmd.extend(["--max-new-tokens", str(args.max_new_tokens)])
|
||||
if args.streaming_prefix_len is not None:
|
||||
cmd.extend(["--streaming-prefix-len", str(args.streaming_prefix_len)])
|
||||
if args.enable_profiler:
|
||||
profiler_dir = Path(args.profiler_dir) if args.profiler_dir is not None else (output_dir / "profiler")
|
||||
cmd.append("--enable-profiler")
|
||||
cmd.extend(["--profiler-dir", str(profiler_dir)])
|
||||
cmd.extend(["--profiler-wait-seconds", str(args.profiler_wait_seconds)])
|
||||
if args.profiler_stages is not None:
|
||||
cmd.append("--profiler-stages")
|
||||
cmd.extend(str(stage_id) for stage_id in args.profiler_stages)
|
||||
return cmd
|
||||
|
||||
|
||||
def _build_case_command(
|
||||
args,
|
||||
mode: ModeSpec,
|
||||
case: CaseSpec,
|
||||
*,
|
||||
batch_tts_path: Path,
|
||||
batch_clone_path: Path,
|
||||
output_dir: Path,
|
||||
) -> list[str]:
|
||||
cmd = _base_command(args, mode, output_dir)
|
||||
cmd.extend(["--warmup-runs", str(case.warmup_runs)])
|
||||
if case.prompt_kind == "single":
|
||||
cmd.extend(["--text", SINGLE_CLONE_TEXT if case.voice_clone else SINGLE_TTS_TEXT])
|
||||
else:
|
||||
cmd.extend(["--txt-prompts", str(batch_clone_path if case.voice_clone else batch_tts_path)])
|
||||
if case.voice_clone:
|
||||
cmd.extend(["--ref-audio", args.ref_audio, "--ref-text", args.ref_text])
|
||||
return cmd
|
||||
|
||||
|
||||
def _run_case(
|
||||
args,
|
||||
mode: ModeSpec,
|
||||
case: CaseSpec,
|
||||
*,
|
||||
batch_tts_path: Path,
|
||||
batch_clone_path: Path,
|
||||
output_root: Path,
|
||||
) -> CaseResult:
|
||||
case_output_dir = output_root / mode.name / case.name
|
||||
case_output_dir.mkdir(parents=True, exist_ok=True)
|
||||
case_log_path = case_output_dir / "run.log"
|
||||
cmd = _build_case_command(
|
||||
args,
|
||||
mode,
|
||||
case,
|
||||
batch_tts_path=batch_tts_path,
|
||||
batch_clone_path=batch_clone_path,
|
||||
output_dir=case_output_dir,
|
||||
)
|
||||
|
||||
print()
|
||||
print("=" * 80)
|
||||
print(f"[{mode.name}] {case.name}")
|
||||
print(f"Output directory: {case_output_dir}")
|
||||
print(shlex.join(cmd))
|
||||
|
||||
start = time.perf_counter()
|
||||
with case_log_path.open("w", encoding="utf-8") as log_fp:
|
||||
process = subprocess.Popen(
|
||||
cmd,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.STDOUT,
|
||||
text=True,
|
||||
bufsize=1,
|
||||
)
|
||||
assert process.stdout is not None
|
||||
for line in process.stdout:
|
||||
print(line, end="")
|
||||
log_fp.write(line)
|
||||
process.wait()
|
||||
|
||||
elapsed_s = time.perf_counter() - start
|
||||
status = "PASS" if (process.returncode or 0) == 0 else f"FAIL({process.returncode})"
|
||||
print(f"[{mode.name}] {case.name} -> {status} ({elapsed_s:.2f}s)")
|
||||
return CaseResult(
|
||||
mode=mode.name,
|
||||
case=case.name,
|
||||
returncode=int(process.returncode or 0),
|
||||
elapsed_s=elapsed_s,
|
||||
output_dir=case_output_dir,
|
||||
log_path=case_log_path,
|
||||
)
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = FlexibleArgumentParser(description="Run the full offline VoxCPM smoke matrix.")
|
||||
parser.add_argument("--model", type=str, required=True, help="Local VoxCPM model directory.")
|
||||
parser.add_argument("--ref-audio", type=str, required=True, help="Reference audio path for clone cases.")
|
||||
parser.add_argument("--ref-text", type=str, required=True, help="Exact transcript spoken in --ref-audio.")
|
||||
parser.add_argument("--output-root", type=str, default=str(DEFAULT_OUTPUT_ROOT), help="Root directory for outputs.")
|
||||
parser.add_argument("--python", type=str, default=sys.executable, help="Python executable used to launch cases.")
|
||||
parser.add_argument("--stage-init-timeout", type=int, default=600, help="Stage initialization timeout in seconds.")
|
||||
parser.add_argument("--log-stats", dest="log_stats", action="store_true", help="Enable vLLM Omni stats logging.")
|
||||
parser.add_argument(
|
||||
"--no-log-stats",
|
||||
dest="log_stats",
|
||||
action="store_false",
|
||||
help="Disable vLLM Omni stats logging.",
|
||||
)
|
||||
parser.set_defaults(log_stats=True)
|
||||
parser.add_argument("--num-runs", type=int, default=1, help="Number of measured runs per case.")
|
||||
parser.add_argument("--cfg-value", type=float, default=2.0, help="Classifier-free guidance value for VoxCPM.")
|
||||
parser.add_argument("--inference-timesteps", type=int, default=10, help="Number of inference timesteps.")
|
||||
parser.add_argument("--min-len", type=int, default=2, help="Minimum generated token length.")
|
||||
parser.add_argument("--max-new-tokens", type=int, default=4096, help="Maximum generated token length.")
|
||||
parser.add_argument(
|
||||
"--streaming-prefix-len",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Optional VoxCPM streaming window passed to streaming cases.",
|
||||
)
|
||||
parser.add_argument("--enable-profiler", action="store_true", help="Enable torch profiler for each case.")
|
||||
parser.add_argument(
|
||||
"--profiler-dir",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Profiler output root. Defaults to <case-output-dir>/profiler.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--profiler-stages",
|
||||
type=int,
|
||||
nargs="*",
|
||||
default=None,
|
||||
help="Optional stage ids to profile. Defaults to all configured stages.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--profiler-wait-seconds",
|
||||
type=float,
|
||||
default=30.0,
|
||||
help="Seconds to wait after stopping profiler for traces to flush.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
if args.num_runs < 1:
|
||||
parser.error("--num-runs must be >= 1")
|
||||
return args
|
||||
|
||||
|
||||
def main(args) -> int:
|
||||
output_root = Path(args.output_root)
|
||||
output_root.mkdir(parents=True, exist_ok=True)
|
||||
batch_tts_path, batch_clone_path = _prepare_batch_inputs(output_root)
|
||||
|
||||
print(f"Model: {args.model}")
|
||||
print(f"Reference audio: {args.ref_audio}")
|
||||
print(f"Reference text: {args.ref_text}")
|
||||
print(f"Python: {args.python}")
|
||||
print(f"Output root: {output_root}")
|
||||
print(f"Cases: {len(MODE_SPECS) * len(CASE_SPECS)}")
|
||||
|
||||
results: list[CaseResult] = []
|
||||
for mode in MODE_SPECS:
|
||||
for case in CASE_SPECS:
|
||||
results.append(
|
||||
_run_case(
|
||||
args,
|
||||
mode,
|
||||
case,
|
||||
batch_tts_path=batch_tts_path,
|
||||
batch_clone_path=batch_clone_path,
|
||||
output_root=output_root,
|
||||
)
|
||||
)
|
||||
|
||||
failed = [result for result in results if not result.ok]
|
||||
print()
|
||||
print("=" * 80)
|
||||
print("Summary:")
|
||||
for result in results:
|
||||
status = "PASS" if result.ok else f"FAIL({result.returncode})"
|
||||
print(f"- [{result.mode}] {result.case}: {status} ({result.elapsed_s:.2f}s)")
|
||||
print(f" output_dir={result.output_dir}")
|
||||
print(f" log={result.log_path}")
|
||||
|
||||
print(f"Passed: {len(results) - len(failed)}/{len(results)}")
|
||||
if failed:
|
||||
print("Failed cases:")
|
||||
for result in failed:
|
||||
print(f"- [{result.mode}] {result.case}: see {result.log_path}")
|
||||
return 1
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main(parse_args()))
|
||||
+3
-1
@@ -10,6 +10,7 @@ nav:
|
||||
- Image Generation: serving/image_generation_api.md
|
||||
- Image Edit: serving/image_edit_api.md
|
||||
- Text to Speech: serving/speech_api.md
|
||||
- Streaming Video Input: serving/video_stream_api.md
|
||||
- Examples:
|
||||
- examples/README.md
|
||||
- Offline Inference:
|
||||
@@ -34,6 +35,7 @@ nav:
|
||||
- Online Serving:
|
||||
- BAGEL-7B-MoT: user_guide/examples/online_serving/bagel.md
|
||||
- vLLM-Omni Helm Chart: user_guide/examples/online_serving/chart-helm.md
|
||||
- Diffusers Backend Adapter Example: user_guide/examples/online_serving/diffusers_pipeline_adapter.md
|
||||
- Fish Speech S2 Pro: user_guide/examples/online_serving/fish_speech.md
|
||||
- GLM-Image Online Serving: user_guide/examples/online_serving/glm_image.md
|
||||
- Image-To-Image: user_guide/examples/online_serving/image_to_image.md
|
||||
@@ -64,6 +66,7 @@ nav:
|
||||
- FP8: user_guide/diffusion/quantization/fp8.md
|
||||
- Int8: user_guide/diffusion/quantization/int8.md
|
||||
- GGUF: user_guide/diffusion/quantization/gguf.md
|
||||
- Attention Backends: user_guide/diffusion/attention_backends.md
|
||||
- Frame Interpolation: user_guide/diffusion/frame_interpolation.md
|
||||
- Parallelism:
|
||||
- Overview: user_guide/diffusion/parallelism/overview.md
|
||||
@@ -81,7 +84,6 @@ nav:
|
||||
- Developer Guide:
|
||||
- General:
|
||||
- contributing/README.md
|
||||
- pr_reviewer.md
|
||||
- glob: contributing/*
|
||||
flatten_single_child_sections: true
|
||||
- Model Implementation:
|
||||
|
||||
@@ -1,5 +1,59 @@
|
||||
# vllm-omni serve
|
||||
|
||||
## Stage-based CLI quickstart
|
||||
|
||||
The stage-based CLI is designed for deployments that require launching each pipeline stage in an isolated process
|
||||
(e.g., across separate operating system processes, distinct GPUs, or distributed hosts).
|
||||
|
||||
- For **migrated models** that utilize the bundled deployment YAML configurations located in
|
||||
`vllm_omni/deploy/`, the `--deploy-config` flag is only required to override the default configuration. By default, executing `vllm serve MODEL --omni ...`
|
||||
automatically loads the bundled deployment configuration.
|
||||
- For **legacy models** utilizing configuration files located in
|
||||
`vllm_omni/model_executor/stage_configs/`, the `--stage-configs-path` parameter remains mandatory.
|
||||
|
||||
Example: Initializing Stage 0 (Orchestrator and API Server):
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
--port 8091 \
|
||||
--stage-id 0 \
|
||||
--omni-master-address 127.0.0.1 \
|
||||
--omni-master-port 26000
|
||||
```
|
||||
|
||||
Example: Initializing a Headless Worker Stage (Stage 1):
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
--stage-id 1 \
|
||||
--headless \
|
||||
--omni-master-address 127.0.0.1 \
|
||||
--omni-master-port 26000
|
||||
```
|
||||
|
||||
When utilizing a custom deployment YAML based on the new schema, append `--deploy-config /path/to/override.yaml` to each command execution. Conversely, for legacy models, substitute this parameter with `--stage-configs-path /path/to/stage_configs.yaml`.
|
||||
|
||||
In the standard execution paradigm, the `--stage-overrides` argument is utilized to apply stage-specific configurations from a single CLI command.
|
||||
However, under the **stage-based CLI** paradigm, where each process strictly encapsulates a single stage, it is recommended to specify tuning parameters directly via discrete command-line flags for the respective stage, rather than constructing a composite `--stage-overrides` JSON string.
|
||||
|
||||
For example, as an alternative to the following composite configuration:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
|
||||
--stage-overrides '{"1": {"gpu_memory_utilization": 0.5}}'
|
||||
```
|
||||
|
||||
the stage-based CLI permits the direct initialization of Stage 1 with explicit parameters:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
--stage-id 1 \
|
||||
--headless \
|
||||
--gpu-memory-utilization 0.5 \
|
||||
--omni-master-address 127.0.0.1 \
|
||||
--omni-master-port 26000
|
||||
```
|
||||
|
||||
## JSON CLI Arguments
|
||||
|
||||
--8<-- "docs/cli/json_tip.inc.md"
|
||||
|
||||
@@ -19,7 +19,7 @@ The new deploy schema lives under `vllm_omni/deploy/` and is paired with a froze
|
||||
| `platforms` | dict | optional | `null` | Keyed by `npu` / `rocm` / `xpu`, each contains a `stages:` list with per-platform overrides applied on top of the CUDA defaults. |
|
||||
| `pipeline` | str | optional | `null` | Override the auto-detected pipeline registry key (used for structural variants like `qwen2_5_omni_thinker_only`). |
|
||||
| `trust_remote_code` | bool | optional | `true` | **Pipeline-wide.** Trust HF remote code on model load; applies to every stage. |
|
||||
| `distributed_executor_backend` | str | optional | `"mp"` | **Pipeline-wide.** Executor backend (`"mp"` or `"ray"`). |
|
||||
| `distributed_executor_backend` | str \| null | optional | `null` | **Pipeline-wide.** Distributed executor backend forwarded to vLLM (`"mp"`, `"ray"`, `"external_launcher"`). If omitted, vLLM auto-selects backend from runtime topology. |
|
||||
| `dtype` | str \| null | optional | `null` | **Pipeline-wide.** Model dtype for every stage. |
|
||||
| `quantization` | str \| null | optional | `null` | **Pipeline-wide.** Quantization method for every stage. |
|
||||
| `enable_prefix_caching` | bool | optional | `false` | **Pipeline-wide.** Prefix cache toggle applied to every stage. |
|
||||
@@ -27,6 +27,9 @@ The new deploy schema lives under `vllm_omni/deploy/` and is paired with a froze
|
||||
| `data_parallel_size` | int | optional | `1` | **Pipeline-wide.** DP degree for every stage. |
|
||||
| `pipeline_parallel_size` | int | optional | `1` | **Pipeline-wide.** PP degree for every stage. |
|
||||
|
||||
Note: for diffusion path, `distributed_executor_backend` currently defaults to
|
||||
`mp`, and `ray` / `external_launcher` are not fully supported yet.
|
||||
|
||||
### Stage fields
|
||||
|
||||
Each entry under `stages:` accepts any `StageDeployConfig` field directly (no nested `engine_args:`). Only fields whose value legitimately varies across stages live here; pipeline-wide settings (trust_remote_code, distributed_executor_backend, dtype, quantization, prefix/chunked prefill, DP/PP sizes) are declared at the top level and applied to every stage. Unknown keys fall through to `engine_extras:` and are forwarded to the engine.
|
||||
@@ -88,6 +91,55 @@ stages:
|
||||
| `--async-chunk` / `--no-async-chunk` | Flip the deploy YAML's `async_chunk:` bool. Unset (default) leaves the YAML value in force. |
|
||||
| `--stage-configs-path` | **Deprecated.** Accepts legacy `stage_args` yamls and (auto-detected) new deploy yamls; emits a deprecation warning. Migrate to `--deploy-config`. To be removed in a follow-up PR. |
|
||||
|
||||
### Stage-Based CLI Paradigm
|
||||
|
||||
The stage-based CLI paradigm facilitates the execution of discrete pipeline stages within isolated processes:
|
||||
|
||||
- **Stage 0** typically encapsulates the orchestrator and the primary API server. Invocation requires `--stage-id 0`,
|
||||
`--omni-master-address`, `--omni-master-port`, and standard port declarations (e.g., `--port`).
|
||||
- **Worker Stages** operate without a distinct API server (i.e., using `--headless`), are assigned sequential `--stage-id` identifiers, and must reference the corresponding
|
||||
`--omni-master-address` and `--omni-master-port` parameters to successfully register with Stage 0.
|
||||
|
||||
For migrated architectures, the system automatically resolves and loads the bundled deployment YAML. Consequently, the primary execution path
|
||||
does **not** necessitate the explicit definition of `--deploy-config`:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
--port 8091 \
|
||||
--stage-id 0 \
|
||||
--omni-master-address 127.0.0.1 \
|
||||
--omni-master-port 26000
|
||||
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
--stage-id 1 \
|
||||
--headless \
|
||||
--omni-master-address 127.0.0.1 \
|
||||
--omni-master-port 26000
|
||||
```
|
||||
|
||||
When instantiating a custom deployment YAML conforming to the updated schema, append the `--deploy-config /path/to/override.yaml` directive
|
||||
to all node invocations. For legacy architectures (e.g., BAGEL) configured via deprecated `stage_args:` schemas, continue to specify the relevant configuration via `--stage-configs-path /path/to/config.yaml`.
|
||||
|
||||
In the context of standard initialization architectures, utilizing the `--stage-overrides` parameter operates as the optimal methodology
|
||||
for delineating stage-specific tuning from the CLI interface:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
|
||||
--stage-overrides '{"1": {"gpu_memory_utilization": 0.5}}'
|
||||
```
|
||||
|
||||
Conversely, in the context of the **stage-based CLI** paradigm, given that each execution process exclusively instantiates a single pipeline stage, configuration override attributes
|
||||
can be defined uniformly via explicit CLI flags on the corresponding instantiation command, rendering composite `--stage-overrides` JSON strings unnecessary:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
--stage-id 1 \
|
||||
--headless \
|
||||
--gpu-memory-utilization 0.5 \
|
||||
--omni-master-address 127.0.0.1 \
|
||||
--omni-master-port 26000
|
||||
```
|
||||
|
||||
### Precedence
|
||||
|
||||
From highest to lowest:
|
||||
@@ -133,6 +185,17 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
|
||||
--stage-overrides '{"0": {"max_num_seqs": 8}}'
|
||||
```
|
||||
|
||||
Within the stage-based CLI paradigm, equivalent configuration parameters can inherently be passed directly
|
||||
as command-line arguments to the designated single-stage process instantiation:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
--stage-id 0 \
|
||||
--max-num-seqs 8 \
|
||||
--omni-master-address 127.0.0.1 \
|
||||
--omni-master-port 26000
|
||||
```
|
||||
|
||||
Effective config per stage after the merge:
|
||||
|
||||
| Stage | Field | Final value | Source |
|
||||
@@ -153,9 +216,14 @@ Therefore, as a core part of vLLM-Omni, the stage configs for a model have sever
|
||||
- Input and output dependencies for each stage.
|
||||
- Default input parameters.
|
||||
|
||||
If users want to modify some part of it. The custom stage_configs file can be input as input argument in both online and offline. Just like examples below:
|
||||
To override specific parameters, explicitly inject the customized configuration schema
|
||||
in both online and offline instantiation flows. Prioritize the `--deploy-config` flag
|
||||
when loading the new-schema deploy YAML schemas, reserving the `--stage-configs-path` parameter
|
||||
exclusively to maintain compatibility with legacy `stage_args` YAML constructs.
|
||||
|
||||
For offline (Assume necessary dependencies have ben imported):
|
||||
Examples:
|
||||
|
||||
For offline (Assume necessary dependencies have been imported):
|
||||
```python
|
||||
model_name = "Qwen/Qwen2.5-Omni-7B"
|
||||
omni = Omni(model=model_name, stage_configs_path="/path/to/custom_stage_configs.yaml")
|
||||
@@ -163,7 +231,13 @@ omni = Omni(model=model_name, stage_configs_path="/path/to/custom_stage_configs.
|
||||
|
||||
For online serving:
|
||||
```bash
|
||||
vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
|
||||
vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --deploy-config /path/to/deploy_config.yaml
|
||||
```
|
||||
|
||||
Legacy online serving:
|
||||
|
||||
```bash
|
||||
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
|
||||
```
|
||||
!!! important
|
||||
We are actively iterating on the definition of stage configs, and we welcome all feedbacks from both community users and developers to help us shape the development!
|
||||
|
||||
@@ -105,7 +105,8 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
|
||||
<td> Depends on reality </td>
|
||||
<td>
|
||||
<strong>Stability:</strong><br>
|
||||
/tests/dfx/stability/tests/test.json<br>
|
||||
/tests/dfx/stability/tests/test_qwen3_omni.json<br>
|
||||
/tests/dfx/stability/tests/test_wan22.json<br>
|
||||
<strong>Reliability:</strong><br>
|
||||
tests/e2e/reliability/test_{model_name}.py
|
||||
</td>
|
||||
@@ -418,13 +419,13 @@ L3 level testing executes after code is merged into the main branch. Its core pu
|
||||
|
||||
**Explanation**:
|
||||
|
||||
@pytest.mark.advanced_model: Marks the test as L3 or L4 level, indicating that this test case performs deep validation, using real models for performance, integration, and accuracy testing. This forms a "basic-advanced" correspondence with the core_model mark at the L2 level.
|
||||
@pytest.mark.advanced_model: Marks the test as L3 merge level, indicating deep validation with real models. @pytest.mark.full_model: Marks L4 nightly-only suites (e.g. `test_*_expansion.py`, doc examples).
|
||||
|
||||
@pytest.mark.core_model: Marks the test as L1 or L2 level, indicating that this test case validates the basic functionality of the core model. It uses mock weights and only checks if the relevant interface functions correctly.
|
||||
|
||||
@pytest.mark.parametrize: A parameterization decorator that allows abstracting test data into parameters, enabling reuse of the same test logic across different data configurations. indirect=True indicates that parameters will be passed to the fixture for processing.
|
||||
|
||||
**Notes**: If you believe the test case only needs to execute basic run logic at the PR-level CI, you can mark it only with @pytest.mark.core_model. If you believe it only needs to execute deep validation run logic at the merge or nightly level, you can mark it only with @pytest.mark.advanced_model. If you believe the test case needs to accommodate both basic run and deep validation test logic, you should mark it with both @pytest.mark.core_model and @pytest.mark.advanced_model.
|
||||
**Notes**: If you believe the test case only needs to execute basic run logic at the PR-level CI, you can mark it only with @pytest.mark.core_model. If you believe it only needs to execute deep validation at merge (L3), use @pytest.mark.advanced_model. For L4 nightly-only expansion and doc-example tests, use @pytest.mark.full_model with `--run-level full_model`. If the test case needs both basic run and deep validation, mark with @pytest.mark.core_model and the appropriate L3/L4 marker (`advanced_model` and/or `full_model`).
|
||||
|
||||
**2.4.2 Test Function Definition and Documentation**
|
||||
|
||||
@@ -516,9 +517,11 @@ L3 level testing executes after code is merged into the main branch. Its core pu
|
||||
|
||||
**Single Request**: The comment clearly states this is a single-request completion test. For concurrent testing, it can be extended to multiple requests using request_num = n.
|
||||
|
||||
**Implicit Validation**: The `send_omni_request` and `send_diffusion_request` methods internally includes validation logic dynamically selected based on the --run-level parameter: core_model performs basic validation, while advanced_model performs deep validation.
|
||||
**Implicit Validation**: The `send_omni_request` and `send_diffusion_request` methods internally includes validation logic dynamically selected based on the --run-level parameter: core_model performs basic validation, while advanced_model and full_model perform deep validation.
|
||||
|
||||
- ***Run Command***: `pytest -s -v /tests/e2e/online_serving/test_{model_name}.py -m advanced_model --run-level=advanced_model`
|
||||
- ***Run Command (L3 merge)***: `pytest -s -v /tests/e2e/online_serving/test_{model_name}.py -m advanced_model --run-level=advanced_model`
|
||||
|
||||
- ***Run Command (L4 nightly expansion)***: `pytest -s -v /tests/e2e/online_serving/test_{model_name}_expansion.py -m full_model --run-level=full_model`
|
||||
|
||||
## Chapter 3: L4 Level Testing - Full Functionality, Performance, and Documentation Testing
|
||||
|
||||
@@ -572,12 +575,12 @@ L5 level testing focuses on the performance of model services under ***long-runn
|
||||
|
||||
### 4.2 Testing Content and Scope
|
||||
|
||||
- ***Long-term Stability (Stability) Testing***: Uses the `tests/dfx/stability/tests/test.json` configuration to run the service under moderate load for an extended period (e.g., over 12 hours), monitoring whether metrics like memory/VRAM usage, response time, and throughput degrade over time, and whether the service process remains stable.
|
||||
- ***Long-term Stability (Stability) Testing***: Uses JSON under `tests/dfx/stability/tests/` (for example `test_qwen3_omni.json` and `test_wan22.json`) to run the service under moderate load for an extended period (e.g., over 12 hours), monitoring whether metrics like memory/VRAM usage, response time, and throughput degrade over time, and whether the service process remains stable.
|
||||
- ***Reliability Testing***: Uses `tests/e2e/reliability/test_{model_name}.py` to actively simulate various fault and abnormal scenarios, such as: dependent service interruption, abnormal input data, network flicker, hardware resource preemption, etc., to verify the system's fault tolerance, self-healing, and graceful degradation capabilities.
|
||||
|
||||
### 4.3 Test Directory and Execution Files
|
||||
|
||||
- ***Stability Test Configuration***: `tests/dfx/stability/tests/test.json`
|
||||
- ***Stability Test Configuration***: `tests/dfx/stability/tests/test_qwen3_omni.json`, `tests/dfx/stability/tests/test_wan22.json` (one JSON per model / runner family)
|
||||
- ***Reliability Test Suite***: `tests/e2e/reliability/test_{model_name}.py`
|
||||
|
||||
### 4.4 Execution Method and Example
|
||||
@@ -588,7 +591,7 @@ L5 level testing focuses on the performance of model services under ***long-runn
|
||||
<details>
|
||||
<summary> Test Examples</summary>
|
||||
|
||||
When you want to add L5-level stability test cases, you can refer to the following format for case addition in `tests/dfx/stability/tests/test.json`:
|
||||
When you want to add L5-level stability test cases, add or extend the appropriate JSON file under `tests/dfx/stability/tests/` (for example `test_qwen3_omni.json` for Omni bench traffic, or `test_wan22.json` for diffusion `/v1/videos` workloads). The following illustrates the Qwen3-Omni shape:
|
||||
|
||||
```json
|
||||
{
|
||||
@@ -659,7 +662,7 @@ All other optional parameters follow the same rules as the in Chapter 3.4.
|
||||
|
||||
</details>
|
||||
|
||||
- - ***Stability***: `pytest -s -v tests/dfx/stability/scripts/test_{model_name}.py`
|
||||
- - ***Stability***: `pytest -s -v tests/dfx/stability/scripts/test_stability_qwen3_omni.py` or `pytest -s -v tests/dfx/stability/scripts/test_stability_wan22.py` (or add `test_stability_<model>.py` alongside a matching JSON config)
|
||||
- ***Reliability***: `pytest -s -v tests/e2e/reliability/test_{model_name}.py`
|
||||
|
||||
## Summary
|
||||
|
||||
@@ -37,7 +37,7 @@ Currently all the features are available in online serving mode. Hence, only nee
|
||||
**Code Style**
|
||||
|
||||
- Validation: test that the multimodal output files of your model have the correct shapes. `OpenAIClientHandler.send_diffusion_request` should have taken care of this.
|
||||
- Test marks: always add `advanced_model` and `diffusion`. Add GPU-related marks if needed. Ref: [Markers for Tests](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/tests_markers/).
|
||||
- Test marks: always add `full_model` and `diffusion` for L4 nightly `test_*_expansion.py` cases. Add GPU-related marks if needed. Ref: [Markers for Tests](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/tests_markers/).
|
||||
- To maximize code reuse, you may refer to
|
||||
- `tests/conftest.py` for `omni_server` (running server in subprocess) and `openai_client` fixtures (sending requests and validating output), `generate_synthetic_image` and `assert_XXX_valid` helper.
|
||||
- `tests/helpers/mark.py` for `@hardware_test(...)` and `hardware_marks`.
|
||||
|
||||
@@ -42,31 +42,63 @@ Our test scripts use the pytest framework. First, please use `git clone https://
|
||||
```
|
||||
The latest test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-ready.yml).
|
||||
|
||||
=== "L3 level & L4 level"
|
||||
=== "L3 level"
|
||||
|
||||
```bash
|
||||
pytest -s -v -m "advanced_model" --run-level=advanced_model
|
||||
```
|
||||
If you only want to run L3 test case, you can use:
|
||||
```bash
|
||||
pytest -s -v e2e/ --ignore-glob='*expansion.py' -m "advanced_model" --run-level=advanced_model
|
||||
```
|
||||
If you only want to run a specific test case, you can use:
|
||||
```bash
|
||||
pytest -s -v test_xxxx.py --run-level=advanced_model
|
||||
```
|
||||
If you only want to run specific test cases on a particular platform, you can use:
|
||||
```bash
|
||||
pytest -s -v -m "core_model and distributed_cuda and L4" --run-level=core_model
|
||||
pytest -s -v -m "advanced_model and distributed_cuda and L4" --run-level=advanced_model
|
||||
```
|
||||
The latest L3 test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-merge.yml).
|
||||
|
||||
|
||||
=== "L4 level"
|
||||
|
||||
```bash
|
||||
cd tests
|
||||
pytest -s -v -m "full_model" --run-level=full_model
|
||||
```
|
||||
If you only want to run a specific test case, you can use:
|
||||
```bash
|
||||
pytest -s -v test_xxxx.py --run-level=full_model
|
||||
```
|
||||
If you only want to run specific test cases on a particular platform, you can use:
|
||||
```bash
|
||||
pytest -s -v -m "full_model and distributed_cuda and L4" --run-level=full_model
|
||||
```
|
||||
Note: To run performance tests (defaults to ``test_qwen_omni.json``; use ``--test-config-file tests/dfx/perf/tests/test_tts.json`` for TTS):
|
||||
```bash
|
||||
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py
|
||||
```
|
||||
The latest L4 (nightly) test commands use the `full_model` marker and `--run-level full_model` (see [test-nightly.yml](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-nightly.yml) and [test-nightly-diffusion.yml](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-nightly-diffusion.yml)). Example:
|
||||
|
||||
The latest L3 test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-merge.yml).
|
||||
```bash
|
||||
cd tests
|
||||
pytest -s -v -m "full_model and omni and H100" --run-level=full_model
|
||||
```
|
||||
|
||||
The latest L4 test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-nightly.yml).
|
||||
=== "L5 level"
|
||||
|
||||
L5 includes stability and reliability testing. Typical commands:
|
||||
|
||||
```bash
|
||||
cd tests
|
||||
|
||||
# Stability: Qwen3-Omni
|
||||
pytest -s -v dfx/stability/scripts/test_stability_qwen3_omni.py
|
||||
|
||||
# Stability: Wan2.2 (v1/videos diffusion benchmark loop)
|
||||
pytest -s -v dfx/stability/scripts/test_stability_wan22.py
|
||||
|
||||
```
|
||||
|
||||
The latest L5 commands for CI can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-ready.yml).
|
||||
|
||||
You can find more information about markers in the documentation: [marker doc](./tests_markers.md)
|
||||
|
||||
|
||||
@@ -8,7 +8,8 @@ Defined in `pyproject.toml`:
|
||||
| Marker | Description |
|
||||
| ------------------ | --------------------------------------------------------- |
|
||||
| `core_model` | L1&L2 tests (run in each PR) |
|
||||
| `advanced_model` | L3&L4 level tests (run in each merge or nightly) |
|
||||
| `advanced_model` | L3 tests (run on each merge to main) |
|
||||
| `full_model` | L4 tests (run nightly) |
|
||||
| `diffusion` | Diffusion model tests |
|
||||
| `omni` | Omni model tests |
|
||||
| `cache` | Cache backend tests |
|
||||
|
||||
@@ -813,16 +813,16 @@ class WanTransformer3DModel(nn.Module):
|
||||
|
||||
---
|
||||
|
||||
### Diffusion Timing (Performance Profiling)
|
||||
### Diffusion Pipeline Profiler (Performance Profiling)
|
||||
When adapting a new diffusion model, it is often useful to analyze the latency of key components such as text encoding, diffusion denoising, and VAE decoding.
|
||||
vLLM-Omni provides a timing utility via `DiffusionPipelineProfilerMixin` to help developers quickly identify performance bottlenecks.
|
||||
|
||||
!!! info
|
||||
`DiffusionPipelineProfilerMixin` is different from using `torch.profiler` for diffusion models, as introduced in this [tutorial](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/profiling.md#3-profiling-diffusion-models). `DiffusionPipelineProfilerMixin` only prints the timing information of multiple functions (such as `vae.decode`), while `torch.profiler` saves detailed GPU/CPU computation time, call/execution steps.
|
||||
`DiffusionPipelineProfilerMixin` is different from using `torch.profiler` for diffusion models, as introduced in this [tutorial](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/profiling.md). `DiffusionPipelineProfilerMixin` only prints the timing information of multiple functions (such as `vae.decode`), while `torch.profiler` saves detailed GPU/CPU computation time, call/execution steps.
|
||||
|
||||
This tool automatically measures the execution time of selected pipeline modules and prints the results in the logs.
|
||||
|
||||
**Enabling Diffusion Timing**
|
||||
**Enabling Diffusion Pipeline Profiler**
|
||||
|
||||
|
||||
Enable timing by setting:
|
||||
@@ -843,7 +843,7 @@ If not specified, the default targets are used:
|
||||
**Adding DiffusionPipelineProfilerMixin to a Pipeline**
|
||||
To enable timing support in your pipeline, inherit from DiffusionPipelineProfilerMixin.
|
||||
```python
|
||||
from vllm_omni.diffusion.utils.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
|
||||
from vllm_omni.diffusion.profiler import DiffusionPipelineProfilerMixin
|
||||
|
||||
class YourModelPipeline(nn.Module, DiffusionPipelineProfilerMixin):
|
||||
# Optional: Specify custom timing targets
|
||||
@@ -862,7 +862,9 @@ class YourModelPipeline(nn.Module, DiffusionPipelineProfilerMixin):
|
||||
...
|
||||
|
||||
# initialize timing profiler
|
||||
self.setup_diffusion_pipeline_profiler(enable_diffusion_pipeline_profiler)
|
||||
self.setup_diffusion_pipeline_profiler(
|
||||
enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler
|
||||
)
|
||||
```
|
||||
The mixin dynamically wraps selected methods and records their execution time during inference.
|
||||
|
||||
@@ -906,9 +908,9 @@ tokenizer.forward
|
||||
|
||||
When enabled, timing logs appear like this:
|
||||
```
|
||||
[DiffusionTiming] text_encoder.forward took 0.018s
|
||||
[DiffusionTiming] diffuse took 2.412s
|
||||
[DiffusionTiming] vae.decode took 0.063s
|
||||
[DiffusionPipelineProfiler] text_encoder.forward took 0.018s
|
||||
[DiffusionPipelineProfiler] diffuse took 2.412s
|
||||
[DiffusionPipelineProfiler] vae.decode took 0.063s
|
||||
```
|
||||
These measurements help identify bottlenecks during model adaptation and optimization
|
||||
|
||||
|
||||
@@ -1,20 +1,93 @@
|
||||
# Adding a TTS Model
|
||||
|
||||
This guide walks through adding a new TTS model to vLLM-Omni, using **Qwen3-TTS**
|
||||
as a reference. Qwen3-TTS demonstrates the standard two-stage TTS pipeline and the
|
||||
key optimizations all TTS models in this repo should follow.
|
||||
This guide walks through adding a new TTS model to vLLM-Omni. Two patterns are
|
||||
supported:
|
||||
|
||||
- **Two-stage pipeline** (e.g. Qwen3-TTS, Fish Speech): an AR code-predictor stage
|
||||
feeds an audio decoder stage via the `async_chunk` framework. This is the standard
|
||||
pattern for maximum streaming performance.
|
||||
- **Single-stage AR model** (e.g. MOSS-TTS-Nano): the model runs entirely inside one
|
||||
AR worker and streams audio chunks directly from its own `inference_stream()` generator.
|
||||
|
||||
Qwen3-TTS is used as the reference for the two-stage pattern. For the single-stage
|
||||
pattern, refer to MOSS-TTS-Nano.
|
||||
|
||||
## Table of Contents
|
||||
|
||||
1. [Overview](#overview)
|
||||
2. [Directory Structure](#directory-structure)
|
||||
3. [Step-by-Step Implementation](#step-by-step-implementation)
|
||||
4. [Key Components](#key-components)
|
||||
5. [Model Registration](#model-registration)
|
||||
6. [Stage Configuration](#stage-configuration)
|
||||
7. [Stage Input Processors](#stage-input-processors)
|
||||
8. [Testing](#testing)
|
||||
9. [Summary](#summary)
|
||||
2. [Cross-Cutting Invariants](#cross-cutting-invariants)
|
||||
3. [Directory Structure](#directory-structure)
|
||||
4. [Step-by-Step Implementation](#step-by-step-implementation)
|
||||
5. [Key Components](#key-components)
|
||||
6. [Model Registration](#model-registration)
|
||||
7. [Stage Configuration](#stage-configuration)
|
||||
8. [Stage Input Processors](#stage-input-processors)
|
||||
9. [Online Serving Integration](#online-serving-integration)
|
||||
10. [Single-Stage Models](#single-stage-models)
|
||||
11. [Testing](#testing)
|
||||
12. [Pre-commit and DCO](#pre-commit-and-dco)
|
||||
13. [Summary](#summary)
|
||||
|
||||
## Cross-Cutting Invariants
|
||||
|
||||
These rules apply to every TTS model regardless of architecture (AR vs AR+diffusion,
|
||||
single-stage vs two-stage, codec-based vs VAE-based). Each has surfaced as a silent
|
||||
bug in a shipped PR — check them at the end of every phase, not just at the start.
|
||||
|
||||
**I1. Streaming output contract.** Pick one per-step semantics for `forward()` and
|
||||
document it in the docstring:
|
||||
|
||||
- *Delta*: yield only new audio samples produced this step. Preferred — linear cost.
|
||||
- *Cumulative*: re-decode from step 0 every call. O(N²); only acceptable when the
|
||||
codec exposes no streaming decode.
|
||||
|
||||
If you choose delta, audit the full chain: `forward()` returns the new chunk →
|
||||
`_consolidate_multimodal_tensors()` in `vllm_omni/engine/output_processor.py`
|
||||
concatenates the audio key into a single tensor at finish → streaming consumers
|
||||
receive per-step chunks, offline consumers receive the concatenated tensor. A
|
||||
mismatch (consolidator skips the key with `continue`, or consumers expect a list
|
||||
but receive a tensor) is invisible in offline RTF benchmarks — users hear replays
|
||||
or truncation only under live playback.
|
||||
|
||||
**I2. Multimodal output consumer hygiene.** `outputs[0].outputs[0].multimodal_output[key]`
|
||||
can be `Tensor`, `list[Tensor]` (pre-consolidation snapshot), `np.ndarray`, or
|
||||
scalar. In every test, example, and benchmark:
|
||||
|
||||
- Never write `dict.get("a") or dict.get("b")` on tensor values — Python evaluates
|
||||
the tensor's truthiness and raises `Boolean value of Tensor with more than one
|
||||
value is ambiguous`. Use explicit `if x is None` chains.
|
||||
- Defensively handle the list form:
|
||||
`if isinstance(x, list): x = torch.cat([t.reshape(-1) for t in x], dim=0)`.
|
||||
- Assert `shape` / `dtype` / `duration` explicitly — do not rely on truthiness for
|
||||
presence checks.
|
||||
|
||||
**I3. Hot-loop GPU discipline.** Inside any per-step model loop (AR decode,
|
||||
diffusion solver, CFM Euler step, per-frame vocoder):
|
||||
|
||||
- No `tensor.item()`, `.cpu()`, or `.tolist()` — each triggers a GPU→CPU sync; a
|
||||
10-step × 60-frame × 4-op loop creates 2400 syncs per request.
|
||||
- Prefer `dst.copy_(src)` over `dst.fill_(src.item())` for scalar-into-buffer writes.
|
||||
- Whole-model `torch.compile(Model.forward, fullgraph=False)` usually outperforms
|
||||
per-submodule compile — fewer dispatch boundaries, larger fusion regions. Measure
|
||||
before choosing granularity.
|
||||
- No Python control flow that depends on tensor values; use `torch.where` or masking.
|
||||
|
||||
Profile before optimizing.
|
||||
|
||||
**I4. Validation pyramid.** Offline RTF alone is necessary but not sufficient. A
|
||||
new TTS model must pass all three levels:
|
||||
|
||||
| Layer | Catches | Tool |
|
||||
|-------|---------|------|
|
||||
| Offline RTF / duration | Throughput regressions, missing audio, wrong sample rate | `end2end.py`, pytest e2e |
|
||||
| Browser streaming playback | Delta-vs-cumulative bugs, chunk boundary glitches, TTFP regressions | Gradio demo over `/v1/audio/speech?stream=true` |
|
||||
| Concurrent requests | Per-request state leaks, codec window round-robin gaps | `max_num_seqs>1` smoke with 4+ parallel prompts |
|
||||
|
||||
**I5. Per-request state belongs to the request.** If the model caches anything
|
||||
across `forward()` calls (streaming generators, codec buffers, sliding-window pads,
|
||||
CUDA graph state), key it by `info.get("_omni_req_id")` and free the entry on
|
||||
request finish. A shared buffer silently corrupts audio across concurrent requests —
|
||||
the symptom is crosstalk or truncation under load, nothing in single-request tests.
|
||||
|
||||
## Overview
|
||||
|
||||
@@ -556,6 +629,302 @@ Recommended test cases for a new TTS model:
|
||||
|
||||
Reference test: `tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py`
|
||||
|
||||
### E2E Online Serving Tests (`tests/e2e/online_serving/test_<your_model>.py`)
|
||||
|
||||
The `omni_server` fixture in `tests/conftest.py` is **module-scoped**. Each distinct
|
||||
`OmniServerParams` id in the same test file forces the fixture to tear the server
|
||||
down and spawn a new one mid-module. A few rules that save real CI debugging time:
|
||||
|
||||
- **Prefer a single `OmniServerParams` set per file.** If you need to exercise two
|
||||
deploy variants (e.g. `model.yaml` and `model_async_chunk.yaml`), either use one
|
||||
variant and exercise streaming via request args, or split into two test files so
|
||||
each file does exactly one server lifecycle. Mid-module teardown/restart is the
|
||||
fragile path and surfaces startup races first.
|
||||
- **Never depend on server-side fetching of external URLs** for reference audio or
|
||||
other fixture data. CI runners (and China-hosted dev boxes) routinely fail on
|
||||
SSL/DNS for public URLs. Inline the payload as a `data:audio/wav;base64,...`
|
||||
ref_audio value — the serving layer accepts both forms.
|
||||
- **Don't roll your own readiness probe.** The harness already waits for HTTP 200
|
||||
on `/health` before releasing the server to the test. If your model needs extra
|
||||
warmup signals, expose them through `/health` rather than adding `time.sleep(...)`
|
||||
inside the test. (Bare TCP `connect_ex` probes were insufficient; see
|
||||
`tests/conftest.py` `OmniServer.wait_for_ready`.)
|
||||
- **Use `core_model` marker + H100 hardware_test** to match the `test-ready.yml`
|
||||
pipeline so your test is picked up by the `ready` label, not only nightly.
|
||||
|
||||
## Online Serving Integration
|
||||
|
||||
To expose your model through the `/v1/audio/speech` OpenAI-compatible endpoint, add
|
||||
**all five** of the following integration points to
|
||||
`vllm_omni/entrypoints/openai/serving_speech.py` in a **single commit**. Adding them
|
||||
piecemeal causes partial-integration failures that are hard to debug.
|
||||
|
||||
### 1. Stage constant
|
||||
|
||||
Near the top of the file, alongside the other `_*_TTS_MODEL_STAGES` constants:
|
||||
|
||||
```python
|
||||
_YOUR_MODEL_TTS_MODEL_STAGES = {"your_model_stage_key"}
|
||||
```
|
||||
|
||||
### 2. Union into `_TTS_MODEL_STAGES`
|
||||
|
||||
Add to the `_TTS_MODEL_STAGES` set union:
|
||||
|
||||
```python
|
||||
_TTS_MODEL_STAGES: set[str] = (
|
||||
...
|
||||
| _YOUR_MODEL_TTS_MODEL_STAGES
|
||||
)
|
||||
```
|
||||
|
||||
### 3. Model type detection
|
||||
|
||||
In `_detect_tts_model_type()`, add before the final `return None`:
|
||||
|
||||
```python
|
||||
if model_stage in _YOUR_MODEL_TTS_MODEL_STAGES:
|
||||
return "your_model"
|
||||
```
|
||||
|
||||
### 4. Request validation dispatch
|
||||
|
||||
In `_validate_tts_request()`, add before the fallback `return`:
|
||||
|
||||
```python
|
||||
if self._tts_model_type == "your_model":
|
||||
return self._validate_your_model_request(request)
|
||||
```
|
||||
|
||||
### 5. Validation and parameter-builder methods
|
||||
|
||||
Add two new methods:
|
||||
|
||||
```python
|
||||
def _validate_your_model_request(
|
||||
self, request: OpenAICreateSpeechRequest
|
||||
) -> str | None:
|
||||
"""Validate YourModel request. Returns an error string or None."""
|
||||
if not request.input or not request.input.strip():
|
||||
return "Input text cannot be empty"
|
||||
return None
|
||||
|
||||
def _build_your_model_params(
|
||||
self, request: OpenAICreateSpeechRequest
|
||||
) -> dict[str, Any]:
|
||||
"""Build additional_information dict for YourModel."""
|
||||
params: dict[str, Any] = {"text": [request.input]}
|
||||
if request.voice is not None:
|
||||
params["voice"] = [request.voice]
|
||||
# Add any other model-specific fields here
|
||||
return params
|
||||
```
|
||||
|
||||
Then wire `_build_your_model_params` into the request-dispatch block in
|
||||
`_create_tts_request()` (search for the equivalent `_build_*_params` call for an
|
||||
existing model to find the right location). If the model supports voice cloning
|
||||
(`ref_audio` → `prompt_audio_path`, `ref_text` → `prompt_text`), add those mappings
|
||||
here too — follow any existing `_build_<model>_params` in `serving_speech.py` (e.g.
|
||||
`_build_moss_tts_params` for the voice-cloning variant) for the pattern.
|
||||
|
||||
> **Two dispatch patterns coexist:** Fish Speech uses a `self._is_fish_speech` boolean
|
||||
> checked *before* `elif self._is_tts`. All newer models use the `_tts_model_type`
|
||||
> string pattern shown above. For new models, always use the string pattern — do not
|
||||
> add new `_is_*` boolean flags.
|
||||
|
||||
> **Note on unused variables:** Only extract parameters in `_build_your_model_params`
|
||||
> that you actually pass to the model's generate / `inference_stream` call. Extracting
|
||||
> a variable without forwarding it will trigger a `ruff F841` pre-commit failure.
|
||||
|
||||
### Merge conflicts
|
||||
|
||||
`serving_speech.py` is modified by every new model PR and is the most common source of
|
||||
rebase conflicts. When rebasing onto `main` and a conflict appears here, the resolution
|
||||
is always to **keep both** the upstream model's additions and your own — never discard
|
||||
either side. After resolving:
|
||||
|
||||
```bash
|
||||
git add vllm_omni/entrypoints/openai/serving_speech.py
|
||||
git rebase --continue
|
||||
```
|
||||
|
||||
## Single-Stage Models
|
||||
|
||||
Some TTS models (e.g. MOSS-TTS-Nano) do not use a two-stage pipeline. Instead the
|
||||
entire AR LM and audio decoder run inside a single AR worker, streaming audio chunks
|
||||
directly from the model's own generator.
|
||||
|
||||
### Directory structure
|
||||
|
||||
```
|
||||
vllm_omni/model_executor/models/your_model_name/
|
||||
__init__.py
|
||||
modeling_your_model_name.py # unified class: load_weights + forward + streaming
|
||||
|
||||
vllm_omni/model_executor/stage_configs/your_model_name.yaml
|
||||
```
|
||||
|
||||
No stage input processor is needed.
|
||||
|
||||
### Stage config
|
||||
|
||||
Use a single stage with `worker_type: ar`. The `is_comprehension: true` field and the
|
||||
top-level `async_chunk: false` are required — omitting them causes silent
|
||||
misclassification in the serving layer. Set `max_num_seqs` to at least 4 for
|
||||
concurrent production use.
|
||||
|
||||
```yaml
|
||||
# stage_configs/your_model_name.yaml
|
||||
async_chunk: false
|
||||
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
stage_type: llm
|
||||
is_comprehension: true # required for serving_speech.py dispatch
|
||||
runtime:
|
||||
devices: "0"
|
||||
engine_args:
|
||||
model_stage: your_model_stage_key
|
||||
model_arch: YourModelForCausalLM
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
engine_output_type: audio
|
||||
max_num_seqs: 4 # min 4 for concurrent requests; default 1 causes gaps
|
||||
final_output: true
|
||||
final_output_type: audio
|
||||
```
|
||||
|
||||
### Generator-based streaming pattern
|
||||
|
||||
This is the MOSS-TTS-Nano pattern, distinct from VoxCPM2's vLLM-native AR pattern
|
||||
(see `plan/voxcpm2_native_ar_design.md` for that variant). Load model weights in
|
||||
`load_weights()` (not `__init__`) so vLLM finishes distributed initialisation before
|
||||
any CUDA allocations. Stream via a per-request generator stored in an instance dict:
|
||||
|
||||
```python
|
||||
class YourModelForCausalLM(nn.Module):
|
||||
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
||||
super().__init__()
|
||||
self._lm: nn.Module | None = None # populated in load_weights()
|
||||
self._stream_gens: dict[str, Any] = {} # request_key → generator
|
||||
|
||||
def load_weights(self, weights):
|
||||
# Load self._lm here, after vLLM distributed init
|
||||
...
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_ids,
|
||||
positions,
|
||||
intermediate_tensors=None,
|
||||
inputs_embeds=None,
|
||||
runtime_additional_information: list[dict] | None = None, # one dict per request
|
||||
**kwargs,
|
||||
) -> OmniOutput:
|
||||
infos = runtime_additional_information or [{}]
|
||||
# Return empty output during dummy/profiling calls
|
||||
if not runtime_additional_information or all(i.get("_is_dummy") for i in infos):
|
||||
self._ar_emit_stop_token = True
|
||||
return OmniOutput(...)
|
||||
|
||||
outputs, last_flags = [], []
|
||||
for info in infos:
|
||||
request_key = str(info.get("_omni_req_id", "0")) # set by vLLM, not user code
|
||||
if request_key not in self._stream_gens:
|
||||
self._stream_gens[request_key] = self._create_stream_gen(info)
|
||||
try:
|
||||
chunk, is_last = next(self._stream_gens[request_key])
|
||||
except StopIteration:
|
||||
chunk, is_last = torch.zeros(0), True
|
||||
if is_last:
|
||||
del self._stream_gens[request_key]
|
||||
outputs.append(chunk)
|
||||
last_flags.append(is_last)
|
||||
|
||||
self._ar_emit_stop_token = all(last_flags)
|
||||
return OmniOutput(multimodal_outputs={"model_outputs": outputs, "is_last": last_flags})
|
||||
|
||||
def _create_stream_gen(self, info: dict):
|
||||
"""Yield (waveform_tensor, is_last) from the model's inference_stream().
|
||||
|
||||
Handle both incremental ("audio" events) and batch ("result" event) models:
|
||||
some upstream implementations emit one "result" event with the full waveform
|
||||
instead of incremental "audio" events. Both paths must be covered.
|
||||
"""
|
||||
for event in self._lm.inference_stream(...):
|
||||
if event["type"] == "audio":
|
||||
yield event["waveform"], False
|
||||
elif event["type"] == "result":
|
||||
# Fallback for models that don't emit incremental audio events
|
||||
yield event.get("waveform", torch.zeros(0)), True
|
||||
return
|
||||
yield torch.zeros(0), True
|
||||
|
||||
def compute_logits(self, hidden_states, sampling_metadata):
|
||||
# Emit EOS only when the last chunk has been yielded so the AR
|
||||
# scheduler ends the request at the right time.
|
||||
...
|
||||
```
|
||||
|
||||
For an in-tree reference, look for any single-stage AR model under
|
||||
`vllm_omni/model_executor/models/` (for example
|
||||
`moss_tts_nano/modeling_moss_tts_nano.py` once its integration has landed).
|
||||
|
||||
## Pre-commit and DCO
|
||||
|
||||
All contributions must pass the pre-commit checks and the Developer Certificate of
|
||||
Origin (DCO) sign-off before merging.
|
||||
|
||||
### Running pre-commit
|
||||
|
||||
Install the hooks once with `pre-commit install`. Then run before committing:
|
||||
|
||||
```bash
|
||||
pre-commit run --files \
|
||||
vllm_omni/model_executor/models/your_model_name/*.py \
|
||||
vllm_omni/entrypoints/openai/serving_speech.py \
|
||||
vllm_omni/model_executor/models/registry.py \
|
||||
tests/e2e/offline_inference/test_your_model_name.py \
|
||||
tests/e2e/online_serving/test_your_model_name.py
|
||||
```
|
||||
|
||||
When pre-commit **modifies files**, it exits with a non-zero code but the reformatting
|
||||
is correct. Stage the modified files and commit again — do not revert the changes.
|
||||
|
||||
Common failures and fixes:
|
||||
|
||||
| Check | Cause | Fix |
|
||||
|-------|-------|-----|
|
||||
| `ruff F841` | Local variable assigned but never used | Remove the extraction or forward it to the model call |
|
||||
| `ruff E402` | Module-level import not at top of file | Move import to the top-level import block |
|
||||
| `ruff format` | Line length, spacing, or quote style | Accept the auto-fix, stage, and re-commit |
|
||||
|
||||
### DCO sign-off
|
||||
|
||||
Every commit must carry a `Signed-off-by` trailer. Use the `-s` flag when committing:
|
||||
|
||||
```bash
|
||||
git commit -s -m "feat(your-model): add YourModel TTS support"
|
||||
```
|
||||
|
||||
Or configure git to add it automatically:
|
||||
|
||||
```bash
|
||||
git config format.signOff true
|
||||
```
|
||||
|
||||
To fix a missing sign-off on the most recent commit:
|
||||
|
||||
```bash
|
||||
git commit --amend -s --no-edit
|
||||
git push origin your-branch --force-with-lease
|
||||
```
|
||||
|
||||
> The DCO check verifies that the commit author email matches the `Signed-off-by` email.
|
||||
> Make sure `git config user.email` is set to the address associated with your GitHub
|
||||
> account before committing.
|
||||
|
||||
## Adding a Model Recipe
|
||||
|
||||
After implementing and testing your model, add a model recipe to the
|
||||
@@ -567,15 +936,19 @@ for the expected format.
|
||||
|
||||
Adding a TTS model to vLLM-Omni involves:
|
||||
|
||||
1. **Create model directory** with AR stage, decoder stage, and unified class
|
||||
1. **Create model directory** with AR stage, decoder stage, and unified class (two-stage)
|
||||
or a single unified class with generator-based streaming (single-stage)
|
||||
2. **AR stage** - use vLLM's native decoder layers with fused QKV; do not wrap HF directly
|
||||
3. **Decoder stage** - thin wrapper around your audio decoder; implement `chunked_decode_streaming()`
|
||||
4. **Unified class** - dispatches on `model_stage`; same structure as `Qwen3TTSModelForGeneration`
|
||||
5. **Register** all stage classes in `registry.py`
|
||||
6. **YAML configs** - provide both batch and `async_chunk` variants
|
||||
7. **Stage input processor** - buffer Stage 0 outputs and forward in chunks of 25
|
||||
8. **Tests** - cover single request, batching, and async_chunk streaming
|
||||
9. **Model recipe** - add to [vllm-project/recipes](https://github.com/vllm-project/recipes)
|
||||
6. **YAML configs** - provide both batch and `async_chunk` variants (two-stage), or a single-stage AR config
|
||||
7. **Stage input processor** - buffer Stage 0 outputs and forward in chunks of 25 (two-stage only)
|
||||
8. **Online serving** - add all 5 integration points to `serving_speech.py` in one commit
|
||||
9. **Tests** - cover single request, batching, and streaming
|
||||
10. **Pre-commit + DCO** - run `pre-commit` before pushing; sign every commit with `git commit -s`
|
||||
11. **Model recipe** - add to [vllm-project/recipes](https://github.com/vllm-project/recipes)
|
||||
12. **Invariants** - re-check I1–I5 (streaming contract, consumer hygiene, hot-loop discipline, validation pyramid, per-request state) at the end of every phase
|
||||
|
||||
### Qwen3-TTS Reference Files
|
||||
|
||||
|
||||
+147
-75
@@ -1,15 +1,17 @@
|
||||
# Profiling vLLM-Omni
|
||||
# Profiling Diffusion Models
|
||||
|
||||
> **Warning:** Profiling is for development and debugging only. It adds significant overhead and should not be enabled in production.
|
||||
|
||||
vLLM-Omni supports two profiler backends through `profiler_config`:
|
||||
Diffusion profiling supports two backends through `profiler_config`:
|
||||
|
||||
- `torch`: detailed CPU/CUDA traces written to `torch_profiler_dir`
|
||||
- `torch`: detailed CPU/CUDA traces, operator tables, and optional memory snapshots
|
||||
- `cuda`: low-overhead CUDA range control for NVIDIA Nsight Systems (`nsys`)
|
||||
|
||||
## 1. Configure Profiling
|
||||
## 1. Configure `profiler_config`
|
||||
|
||||
Use the same `profiler_config` shape everywhere:
|
||||
Use `profiler_config` to enable profiling for a diffusion model. For diffusion usage, pass it directly to `Omni(...)` or `vllm serve`.
|
||||
|
||||
Minimal torch-profiler config:
|
||||
|
||||
```yaml
|
||||
profiler_config:
|
||||
@@ -21,54 +23,114 @@ Supported fields:
|
||||
|
||||
| Field | Description |
|
||||
|---|---|
|
||||
| `profiler` | Profiler backend. Supported values: `torch`, `cuda`. |
|
||||
| `torch_profiler_dir` | Output directory for torch traces. Required when `profiler: torch`. |
|
||||
| `profiler` | Profiler backend. Supported values: `torch`, `cuda`. Use `torch` for `trace.json`, Excel operator tables, and optional memory snapshots. Use `cuda` for Nsight Systems only. |
|
||||
| `torch_profiler_dir` | Output directory for torch-profiler artifacts. Required when `profiler: torch`. |
|
||||
| `torch_profiler_use_gzip` | Compress `trace_rank*.json` into `trace_rank*.json.gz`. |
|
||||
| `torch_profiler_record_shapes` | Record input shapes and add a `by_shape` sheet to `ops_rank*.xlsx`. |
|
||||
| `torch_profiler_with_stack` | Record call stacks, add a `by_stack` sheet to `ops_rank*.xlsx`, and export `stacks_cpu_rank*.txt` and `stacks_cuda_rank*.txt`. |
|
||||
| `torch_profiler_with_memory` | Enable memory profiling and attempt to dump `memory_snapshot_rank*.pickle`. The pickle is only generated when the current backend supports memory history and snapshot APIs. |
|
||||
| `torch_profiler_with_flops` | Enable FLOPs collection in `torch.profiler`. This does not add a separate output file. |
|
||||
| `torch_profiler_dump_cuda_time_total` | Export an additional text summary `profiler_out_<rank>.txt` sorted by `self_cuda_time_total`. |
|
||||
| `delay_iterations` | Number of worker iterations to skip before profiling starts. |
|
||||
| `max_iterations` | Maximum number of worker iterations to capture before auto-stop. |
|
||||
| `wait_iterations` | Torch-profiler wait iterations before warmup. |
|
||||
| `warmup_iterations` | Torch-profiler warmup iterations. |
|
||||
| `active_iterations` | Torch-profiler active iterations. |
|
||||
| `wait_iterations` | Torch-profiler wait iterations before warmup. |
|
||||
|
||||
For multi-stage omni pipelines, put `profiler_config` under the target stage's `engine_args`.
|
||||
### Minimal configurations by output
|
||||
|
||||
```yaml
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
stage_type: llm
|
||||
engine_args:
|
||||
profiler_config:
|
||||
profiler: torch
|
||||
torch_profiler_dir: ./perf
|
||||
```
|
||||
|
||||
For single-stage diffusion usage, pass `profiler_config` directly to `Omni(...)` or `vllm serve`.
|
||||
|
||||
## 2. Profiling Omni Pipelines
|
||||
|
||||
It is usually best to profile only the stages you need.
|
||||
Only collect trace output:
|
||||
|
||||
```python
|
||||
# Profile all stages.
|
||||
omni.start_profile()
|
||||
|
||||
# Profile selected stages only.
|
||||
omni.start_profile(stages=[0, 2])
|
||||
...
|
||||
omni.stop_profile(stages=[0, 2])
|
||||
profiler_config = {
|
||||
"profiler": "torch",
|
||||
"torch_profiler_dir": "./perf",
|
||||
}
|
||||
```
|
||||
|
||||
Always stop the same stage set that you started. If only some stages have `profiler_config`, pass an explicit `stages=[...]` list instead of relying on the default "all stages" behavior.
|
||||
Outputs:
|
||||
|
||||
Examples:
|
||||
- `trace_rank*.json`
|
||||
- `ops_rank*.xlsx` with a `summary` sheet
|
||||
|
||||
1. [Qwen2.5-Omni end2end](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen2_5_omni/end2end.py)
|
||||
2. [Qwen3-Omni end2end](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen3_omni/end2end.py)
|
||||
Collect compressed trace output:
|
||||
|
||||
## 3. Profiling Single-Stage Diffusion
|
||||
```python
|
||||
profiler_config = {
|
||||
"profiler": "torch",
|
||||
"torch_profiler_dir": "./perf",
|
||||
"torch_profiler_use_gzip": True,
|
||||
}
|
||||
```
|
||||
|
||||
Single-stage diffusion models use the same `start_profile()` / `stop_profile()` controls, but you must provide `profiler_config` explicitly.
|
||||
Outputs:
|
||||
|
||||
### PyTorch profiler
|
||||
- `trace_rank*.json.gz`
|
||||
- `ops_rank*.xlsx` with a `summary` sheet
|
||||
|
||||
Collect trace and full operator tables:
|
||||
|
||||
```python
|
||||
profiler_config = {
|
||||
"profiler": "torch",
|
||||
"torch_profiler_dir": "./perf",
|
||||
"torch_profiler_record_shapes": True,
|
||||
"torch_profiler_with_stack": True,
|
||||
}
|
||||
```
|
||||
|
||||
Outputs:
|
||||
|
||||
- `trace_rank*.json`
|
||||
- `ops_rank*.xlsx` with `summary`, `by_shape`, and `by_stack`
|
||||
- `stacks_cpu_rank*.txt`
|
||||
- `stacks_cuda_rank*.txt`
|
||||
|
||||
Collect trace, operator tables, and memory snapshots:
|
||||
|
||||
```python
|
||||
profiler_config = {
|
||||
"profiler": "torch",
|
||||
"torch_profiler_dir": "./perf",
|
||||
"torch_profiler_record_shapes": True,
|
||||
"torch_profiler_with_stack": True,
|
||||
"torch_profiler_with_memory": True,
|
||||
}
|
||||
```
|
||||
|
||||
Outputs:
|
||||
|
||||
- `trace_rank*.json`
|
||||
- `ops_rank*.xlsx` with `summary`, `by_shape`, and `by_stack`
|
||||
- `stacks_cpu_rank*.txt`
|
||||
- `stacks_cuda_rank*.txt`
|
||||
- `memory_snapshot_rank*.pickle` when supported by the current backend
|
||||
|
||||
### Full torch-profiler configuration
|
||||
|
||||
If you want to enable the commonly used torch-profiler options together:
|
||||
|
||||
```python
|
||||
profiler_config = {
|
||||
"profiler": "torch",
|
||||
"torch_profiler_dir": "./perf",
|
||||
"torch_profiler_use_gzip": False,
|
||||
"torch_profiler_record_shapes": True,
|
||||
"torch_profiler_with_stack": True,
|
||||
"torch_profiler_with_memory": True,
|
||||
"torch_profiler_with_flops": False,
|
||||
"torch_profiler_dump_cuda_time_total": False,
|
||||
"delay_iterations": 0,
|
||||
"max_iterations": 0,
|
||||
"wait_iterations": 0,
|
||||
"warmup_iterations": 0,
|
||||
"active_iterations": 0,
|
||||
}
|
||||
```
|
||||
|
||||
## 2. Profiling Diffusion with PyTorch Profiler
|
||||
|
||||
Single-stage diffusion models use `start_profile()` / `stop_profile()` controls. The profiler only writes artifacts after profiling has been started and then stopped.
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
@@ -86,7 +148,29 @@ omni.start_profile()
|
||||
omni.stop_profile()
|
||||
```
|
||||
|
||||
### Nsight Systems (`nsys`)
|
||||
For diffusion offline example scripts under `examples/offline_inference/`, pass `--profiler-config` as a JSON object. The script enables profiling when this argument is set and wraps generation with `start_profile()` / `stop_profile()`.
|
||||
|
||||
Example:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/image_to_video/image_to_video.py \
|
||||
--model Wan-AI/Wan2.2-I2V-A14B-Diffusers \
|
||||
--image input.jpg \
|
||||
--prompt "A cat playing with yarn" \
|
||||
--profiler-config '{
|
||||
"profiler": "torch",
|
||||
"torch_profiler_dir": "./perf",
|
||||
"torch_profiler_record_shapes": true,
|
||||
"torch_profiler_with_stack": true
|
||||
}'
|
||||
```
|
||||
|
||||
Examples:
|
||||
|
||||
1. [Image edit example](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/image_to_image/image_edit.py)
|
||||
2. [Image to video example](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_video)
|
||||
|
||||
## 3. Profiling Diffusion with Nsight Systems (`nsys`)
|
||||
|
||||
For Nsight Systems, use `profiler: cuda` and wrap the process with `nsys profile`.
|
||||
|
||||
@@ -103,35 +187,20 @@ nsys profile \
|
||||
The Python process being profiled must create the diffusion engine with:
|
||||
|
||||
```python
|
||||
profiler_config={"profiler": "cuda"}
|
||||
profiler_config = {"profiler": "cuda"}
|
||||
```
|
||||
|
||||
Then call `start_profile()` before the requests you want to capture and `stop_profile()` after them. The diffusion worker processes open and close the CUDA capture range themselves, so `nsys` sees the actual GPU work instead of only the parent process.
|
||||
|
||||
Examples:
|
||||
|
||||
1. [Image edit example](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/image_to_image/image_edit.py)
|
||||
2. [Image to video example](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_video)
|
||||
|
||||
## 4. Profiling Online Serving
|
||||
|
||||
When any stage has `profiler_config.profiler` set, the server exposes:
|
||||
When `profiler_config.profiler` is set for a diffusion model, the server exposes:
|
||||
|
||||
- `POST /start_profile`
|
||||
- `POST /stop_profile`
|
||||
|
||||
### Start the server
|
||||
|
||||
Multi-stage omni serving:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen2.5-Omni-7B \
|
||||
--omni \
|
||||
--port 8091
|
||||
```
|
||||
|
||||
(The default deploy config at `vllm_omni/deploy/qwen2_5_omni.yaml` is loaded automatically. Pass `--deploy-config /path/to/custom.yaml` to override.)
|
||||
|
||||
Single-stage diffusion serving with torch profiler:
|
||||
|
||||
```bash
|
||||
@@ -144,21 +213,17 @@ vllm serve Wan-AI/Wan2.2-I2V-A14B-Diffusers \
|
||||
"torch_profiler_with_stack": true,
|
||||
"torch_profiler_with_flops": false,
|
||||
"torch_profiler_use_gzip": true,
|
||||
"torch_profiler_dump_cuda_time_total": true,
|
||||
"torch_profiler_dump_cuda_time_total": false,
|
||||
"torch_profiler_record_shapes": true,
|
||||
"torch_profiler_with_memory": true,
|
||||
"delay_iterations": 0,
|
||||
"max_iterations": 0
|
||||
"max_iterations": 0,
|
||||
"wait_iterations": 0,
|
||||
"warmup_iterations": 0,
|
||||
"active_iterations": 0
|
||||
}'
|
||||
```
|
||||
|
||||
Useful optional fields:
|
||||
|
||||
- `torch_profiler_with_stack`: export `by_stack` operator views and stack text files.
|
||||
- `torch_profiler_record_shapes`: export `by_shape` operator views.
|
||||
- `torch_profiler_with_memory`: dump `memory_snapshot_rank*.pickle` when the backend supports memory history.
|
||||
- `torch_profiler_use_gzip`: write the trace as `trace_rank*.json.gz`.
|
||||
|
||||
Single-stage diffusion serving with Nsight Systems:
|
||||
|
||||
```bash
|
||||
@@ -176,30 +241,37 @@ nsys profile \
|
||||
|
||||
### Control capture
|
||||
|
||||
Example profiling flow for an online Qwen-Image request:
|
||||
|
||||
```bash
|
||||
# Start profiling on all profiled stages.
|
||||
# Start profiling.
|
||||
curl -X POST http://localhost:8091/start_profile
|
||||
|
||||
# Start profiling on selected stages.
|
||||
curl -X POST http://localhost:8091/start_profile \
|
||||
# Send a Qwen-Image generation request while profiling is active.
|
||||
curl http://localhost:8091/v1/images/generations \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"stages": [0]}'
|
||||
-d '{
|
||||
"model": "Qwen/Qwen-Image",
|
||||
"prompt": "A red vintage bicycle parked beside a quiet canal at sunset"
|
||||
}'
|
||||
|
||||
# Stop profiling.
|
||||
# Stop profiling and flush profiler artifacts.
|
||||
curl -X POST http://localhost:8091/stop_profile
|
||||
```
|
||||
|
||||
For mixed-stage pipelines, use explicit `stages` and pass the same stage list to both endpoints.
|
||||
## 5. Diffusion Pipeline Profiler
|
||||
|
||||
## 5. Analyze Results
|
||||
For lightweight per-stage pipeline timing such as `vae.decode` or `diffuse`, see [Diffusion Pipeline Profiler](model/adding_diffusion_model.md#diffusion-pipeline-profiler-performance-profiling). That utility logs stage durations only and does not generate torch-profiler artifacts such as `trace.json`, Excel tables, or memory snapshots.
|
||||
|
||||
Torch profiler output:
|
||||
## 6. Analyze Results
|
||||
|
||||
Torch-profiler output:
|
||||
|
||||
- Chrome/Perfetto trace: `trace_rank*.json` or `trace_rank*.json.gz`
|
||||
- Excel workbook: `ops_rank*.xlsx` with `summary`, and optional `by_shape` / `by_stack` sheets
|
||||
- Stack exports: `stacks_cpu_rank*.txt` and `stacks_cuda_rank*.txt` when stack capture is enabled
|
||||
- Memory snapshot: `memory_snapshot_rank*.pickle` when memory capture is enabled and supported by the backend
|
||||
- Optional aggregated CUDA-time tables under the same session directory
|
||||
- Optional CUDA-time text summary: `profiler_out_<rank>.txt` when `torch_profiler_dump_cuda_time_total` is enabled
|
||||
|
||||
CUDA profiler / Nsight Systems output:
|
||||
|
||||
@@ -211,4 +283,4 @@ Recommended viewers:
|
||||
- `nsys stats <report>.nsys-rep` for CLI summaries
|
||||
- Nsight Systems GUI for CUDA kernel timelines
|
||||
|
||||
vLLM-Omni reuses the vLLM profiling infrastructure where possible. For the upstream reference, see the [vLLM profiling guide](https://docs.vllm.ai/en/stable/contributing/profiling/).
|
||||
For upstream background on the underlying vLLM profiling infrastructure, see the [vLLM profiling guide](https://docs.vllm.ai/en/stable/contributing/profiling/).
|
||||
@@ -121,6 +121,7 @@ def extract_omni_serve_subparser_init():
|
||||
"FlexibleArgumentParser": _FlexibleArgumentParser,
|
||||
"make_arg_parser": lambda parser: parser, # no-op for doc
|
||||
"_ensure_vllm_platform": lambda: None, # no-op for doc
|
||||
"nullify_stage_engine_defaults": lambda parser: None, # no-op for doc
|
||||
"VLLM_SUBCMD_PARSER_EPILOG": "",
|
||||
"logger": logger,
|
||||
"DummySubparsers": DummySubparsers,
|
||||
|
||||
@@ -36,6 +36,8 @@ th {
|
||||
| `LTX2ImageToVideoPipeline` | LTX-2-I2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
|
||||
| `LTX2TwoStagesPipeline` | LTX-2-T2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
|
||||
| `LTX2ImageToVideoTwoStagesPipeline` | LTX-2-I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
|
||||
| `LTX23Pipeline` | LTX-2.3-T2V | `dg845/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
|
||||
| `LTX23ImageToVideoPipeline` | LTX-2.3-I2V | `dg845/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
|
||||
| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | |
|
||||
| `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | |
|
||||
| `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
@@ -54,7 +56,8 @@ th {
|
||||
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `MiMoAudioForConditionalGeneration` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` | ✅︎ | ✅︎ | | |
|
||||
| `MiMoAudioModel` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` | ✅︎ | ✅︎ | | |
|
||||
| `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | |
|
||||
| `Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | ✅︎ | ✅︎ | | |
|
||||
| `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | |
|
||||
| `DreamIDOmniPipeline` | DreamID-Omni | `XuGuo699/DreamID-Omni` | ✅︎ | ✅︎ | | |
|
||||
|
||||
@@ -1,255 +0,0 @@
|
||||
# VLLM-Omni PR Reviewer
|
||||
|
||||
## Overview
|
||||
|
||||
The VLLM-Omni PR Reviewer is an automated code review bot powered by GLM-4.7 AI model. It helps maintain code quality by providing intelligent feedback on pull requests.
|
||||
|
||||
## Features
|
||||
|
||||
- **Intelligent Code Analysis**: Leverages GLM-4.7 for understanding code context and providing meaningful feedback
|
||||
- **Comprehensive Reviews**: Covers code quality, architecture, security, testing, and documentation
|
||||
- **Structured Output**: Provides well-formatted reviews with clear sections and actionable suggestions
|
||||
- **Rate Limiting**: Built-in cooldown mechanism to prevent excessive API usage
|
||||
- **Retry Logic**: Automatic retries with exponential backoff for transient API failures
|
||||
- **Defensive Parsing**: Robust validation of API responses to handle malformed data
|
||||
- **Cost Control**: Only repository members/collaborators/owners can trigger reviews
|
||||
|
||||
## How to Use
|
||||
|
||||
### Triggering a Review
|
||||
|
||||
To trigger an automated PR review, mention the bot in a PR comment:
|
||||
|
||||
```
|
||||
@vllm-omni-reviewer please review
|
||||
```
|
||||
|
||||
Or include in your PR description:
|
||||
|
||||
```
|
||||
@vllm-omni-reviewer
|
||||
```
|
||||
|
||||
The bot will automatically review your changes and post a detailed comment.
|
||||
|
||||
## What Gets Reviewed
|
||||
|
||||
- **vLLM Architecture Compatibility**: Ensures changes align with vLLM's design patterns
|
||||
- **Multi-modal Integration**: Reviews audio, vision, and text processing implementations
|
||||
- **Performance Implications**: Analyzes impact on inference latency and throughput
|
||||
- **Code Quality**: Checks Python best practices, type hints, and documentation
|
||||
- **Security Considerations**: Identifies potential security vulnerabilities
|
||||
- **Testing Coverage**: Recommends additional test cases when needed
|
||||
|
||||
## Review Output
|
||||
|
||||
The bot posts a structured review comment with:
|
||||
|
||||
- **Overview**: Brief summary of the PR's purpose
|
||||
- **Critical Issues (Must Fix)**: Blocking issues that need to be addressed
|
||||
- **Important Issues (Should Fix)**: Significant concerns that should be resolved
|
||||
- **Minor Issues & Suggestions**: Small improvements and optional suggestions
|
||||
- **Positive Aspects**: Highlights well-implemented features
|
||||
- **Performance Considerations**: Analysis of performance impact
|
||||
- **Testing Recommendations**: Suggestions for additional tests
|
||||
- **Overall Assessment**: Final recommendation (Approve/Request Changes/Needs Major Work)
|
||||
|
||||
## Rate Limiting and Cooldown
|
||||
|
||||
The bot includes a cooldown mechanism to prevent excessive API usage:
|
||||
|
||||
- **Default cooldown**: 5 minutes between reviews per PR
|
||||
- **Configurable**: Can be adjusted via `PR_REVIEWER_COOLDOWN_MINUTES` environment variable
|
||||
- **Smart detection**: Checks for previous bot comments before starting a review
|
||||
|
||||
If you trigger a review within the cooldown period, the bot will log a message and skip the review.
|
||||
|
||||
## Architecture
|
||||
|
||||
```
|
||||
┌─────────────────┐
|
||||
│ PR Comment │
|
||||
│ @vllm-omni- │
|
||||
│ reviewer │
|
||||
└────────┬────────┘
|
||||
│
|
||||
▼
|
||||
┌─────────────────────────────────┐
|
||||
│ GitHub Actions Workflow │
|
||||
│ (.github/workflows/ │
|
||||
│ pr-reviewer.yml) │
|
||||
│ │
|
||||
│ - Python 3.11 │
|
||||
│ - requests==2.31.0 │
|
||||
│ - pyyaml==6.0.1 │
|
||||
└────────┬────────────────────────┘
|
||||
│
|
||||
▼
|
||||
┌─────────────────────────────────┐
|
||||
│ PR Reviewer Script │
|
||||
│ (.github/scripts/ │
|
||||
│ pr_reviewer.py) │
|
||||
│ │
|
||||
│ 1. Check cooldown │
|
||||
│ 2. Fetch PR details & diff │
|
||||
│ 3. Build review prompt │
|
||||
│ 4. Call GLM-4.7 API │
|
||||
│ (with retry logic) │
|
||||
│ 5. Validate response │
|
||||
│ 6. Post review comment │
|
||||
└────────┬────────────────────────┘
|
||||
│
|
||||
▼
|
||||
┌─────────────────────────────────┐
|
||||
│ GLM-4.7 API │
|
||||
│ (open.bigmodel.cn) │
|
||||
└─────────────────────────────────┘
|
||||
```
|
||||
|
||||
1. **GitHub Actions Workflow** (`.github/workflows/pr-reviewer.yml`): Triggers on @mention
|
||||
2. **Python Script** (`.github/scripts/pr_reviewer.py`): Fetches PR data and calls GLM-4.7 API
|
||||
3. **GLM-4.7 API**: Provides intelligent code analysis
|
||||
|
||||
## Testing
|
||||
|
||||
### Testing the PR Reviewer Bot
|
||||
|
||||
To test the PR reviewer bot before deploying to production:
|
||||
|
||||
1. **Create a test PR** - Make a small, safe change (e.g., documentation update)
|
||||
2. **Open the PR** - Create a pull request with a descriptive title
|
||||
3. **Trigger the review** - Comment `@vllm-omni-reviewer` on the PR
|
||||
4. **Monitor results** - Check the Actions tab for workflow execution logs
|
||||
|
||||
### Running Unit Tests
|
||||
|
||||
The bot includes comprehensive unit tests that can be run locally:
|
||||
|
||||
```bash
|
||||
# Run all tests
|
||||
pytest .github/tests/test_pr_reviewer.py -v
|
||||
|
||||
# Run specific test
|
||||
pytest .github/tests/test_pr_reviewer.py::TestCheckTrigger -v
|
||||
|
||||
# Run with coverage
|
||||
pytest .github/tests/test_pr_reviewer.py --cov=.github/scripts/pr_reviewer.py --cov-report=term-missing
|
||||
```
|
||||
|
||||
### What to Look For
|
||||
|
||||
When testing, verify that:
|
||||
- [ ] The workflow triggers on the `@vllm-omni-reviewer` comment
|
||||
- [ ] The cooldown mechanism works correctly
|
||||
- [ ] The GLM API call completes without errors (with retry if needed)
|
||||
- [ ] A review comment is posted to the PR
|
||||
- [ ] The review content is meaningful and well-structured
|
||||
- [ ] The cost is within the expected range (0.50-5 CNY)
|
||||
|
||||
### Safe Test Changes
|
||||
|
||||
For testing, consider making these types of safe changes:
|
||||
- Documentation updates (like adding this Testing section)
|
||||
- Comment improvements
|
||||
- README enhancements
|
||||
- Non-functional file additions
|
||||
|
||||
### Example Test PR
|
||||
|
||||
A good test PR might:
|
||||
- Update a documentation file
|
||||
- Add explanatory comments
|
||||
- Improve code formatting
|
||||
- Fix a minor typo
|
||||
|
||||
These changes are safe to merge if the test is successful and won't affect functionality.
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### Bot Doesn't Respond
|
||||
|
||||
1. **Check permissions** - Verify you have Owner/Member/Collaborator access
|
||||
2. **Check Actions tab** - Look for workflow execution and view logs
|
||||
3. **Check cooldown** - If another review was posted recently, wait for the cooldown period
|
||||
4. **Check API key** - Ensure `GLM_API_KEY` is configured in repository secrets
|
||||
|
||||
### API Errors
|
||||
|
||||
If the GLM API call fails:
|
||||
- Check the Actions tab for detailed error logs
|
||||
- Verify the `GLM_API_KEY` secret is correctly configured
|
||||
- Ensure sufficient API quota is available
|
||||
- The bot will automatically retry up to 3 times with exponential backoff
|
||||
|
||||
### Review Seems Truncated
|
||||
|
||||
If the review appears incomplete:
|
||||
- Large diffs may be truncated at 100,000 characters
|
||||
- Check the logs for truncation warnings
|
||||
- Consider breaking large PRs into smaller chunks
|
||||
|
||||
## Configuration
|
||||
|
||||
### Required Secrets
|
||||
|
||||
The following secret must be configured in the repository settings:
|
||||
|
||||
- `GLM_API_KEY` - Your GLM (BigModel) API key for accessing the GLM-4.7 API
|
||||
|
||||
To add the secret:
|
||||
1. Go to repository Settings → Secrets and variables → Actions
|
||||
2. Click "New repository secret"
|
||||
3. Name: `GLM_API_KEY`
|
||||
4. Value: Your GLM API key
|
||||
|
||||
### Optional Configuration
|
||||
|
||||
The following optional environment variables can be set in the workflow file:
|
||||
|
||||
| Variable | Default | Description |
|
||||
|----------|---------|-------------|
|
||||
| `GLM_API_URL` | `https://open.bigmodel.cn/api/paas/v4/chat/completions` | GLM API endpoint |
|
||||
| `GLM_MODEL` | `glm-4.7` | Model to use for reviews |
|
||||
| `PR_REVIEWER_COOLDOWN_MINUTES` | `5` | Cooldown period between reviews |
|
||||
| `PR_REVIEWER_MAX_RETRIES` | `3` | Maximum API retry attempts |
|
||||
| `PR_REVIEWER_RETRY_DELAY` | `1.0` | Base delay for retry backoff (seconds) |
|
||||
| `PR_REVIEWER_MAX_DIFF_SIZE` | `100000` | Maximum diff size before truncation |
|
||||
|
||||
### Workflow Customization
|
||||
|
||||
The workflow can be customized in `.github/workflows/pr-reviewer.yml`:
|
||||
- Change Python version (default: 3.11)
|
||||
- Adjust timeout value (default: 10 minutes)
|
||||
- Modify trigger conditions
|
||||
- Add additional dependencies
|
||||
|
||||
## Code Quality
|
||||
|
||||
The PR reviewer script follows vllm-omni coding standards:
|
||||
|
||||
- **Type hints**: All functions have complete type hints following mypy strict mode
|
||||
- **Logging**: Uses Python's logging module for structured logging
|
||||
- **Testing**: Comprehensive unit tests with pytest
|
||||
- **Pre-commit**: Script is checked by pre-commit hooks (flake8)
|
||||
|
||||
## Cost Estimate
|
||||
|
||||
| Component | Cost |
|
||||
|-----------|------|
|
||||
| GitHub Actions (public repo) | Free |
|
||||
| GLM API (glm-4.7) | ~0.50-5 CNY per PR (varies by size) |
|
||||
| Total (20 PRs/month) | ~10-100 CNY/month (~$2-15 USD) |
|
||||
|
||||
## Contributing
|
||||
|
||||
To improve the PR reviewer bot:
|
||||
|
||||
1. Edit `.github/scripts/pr_reviewer.py` for logic changes
|
||||
2. Edit `.github/workflows/pr-reviewer.yml` for workflow changes
|
||||
3. Add tests to `.github/tests/test_pr_reviewer.py`
|
||||
4. Run `pre-commit run --files .github/scripts/pr_reviewer.py` to check code quality
|
||||
5. Test thoroughly with a test PR before deploying to production
|
||||
|
||||
## License
|
||||
|
||||
This bot is part of the VLLM-Omni project and follows the same license terms.
|
||||
@@ -22,21 +22,10 @@ vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
|
||||
--enforce-eager
|
||||
|
||||
# Fish Speech S2 Pro
|
||||
vllm-omni serve fishaudio/s2-pro \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/fish_speech_s2_pro.yaml \
|
||||
--omni \
|
||||
--port 8091 \
|
||||
--trust-remote-code \
|
||||
--enforce-eager \
|
||||
--gpu-memory-utilization 0.9
|
||||
vllm serve fishaudio/s2-pro --omni --port 8091
|
||||
|
||||
# Voxtral TTS
|
||||
vllm serve mistralai/Voxtral-4B-TTS-2603 \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
|
||||
--omni \
|
||||
--port 8091 \
|
||||
--trust-remote-code \
|
||||
--enforce-eager
|
||||
vllm serve mistralai/Voxtral-4B-TTS-2603 --omni --port 8091
|
||||
```
|
||||
|
||||
### Generate Speech
|
||||
|
||||
@@ -0,0 +1,93 @@
|
||||
# Streaming Video Input API
|
||||
|
||||
vLLM-Omni provides a WebSocket API for streaming video frames and optional audio chunks into Qwen3-Omni, then asking questions over the buffered session context.
|
||||
|
||||
Each server instance runs a single model specified at startup with `vllm serve <model> --omni`.
|
||||
|
||||
## Quick Start
|
||||
|
||||
### Start the Server
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct \
|
||||
--deploy-config vllm_omni/deploy/qwen3_omni.yaml \
|
||||
--omni \
|
||||
--port 8000 \
|
||||
--trust-remote-code
|
||||
```
|
||||
|
||||
### Run the Example Client
|
||||
|
||||
```bash
|
||||
python examples/online_serving/qwen3_omni/streaming_video_client.py \
|
||||
--url ws://localhost:8000/v1/video/chat/stream \
|
||||
--video /path/to/video.mp4 \
|
||||
--query "Describe what is happening in the video."
|
||||
```
|
||||
|
||||
## API Reference
|
||||
|
||||
### Endpoint
|
||||
|
||||
```text
|
||||
WebSocket /v1/video/chat/stream
|
||||
```
|
||||
|
||||
### Protocol
|
||||
|
||||
| Direction | Type | Required fields | Description |
|
||||
|-----------|------|-----------------|-------------|
|
||||
| Client -> Server | `session.config` | none | First message. Configures output modalities, frame sampling, EVS, and prompts. |
|
||||
| Client -> Server | `video.frame` | `data` | Base64 JPEG/PNG frame. |
|
||||
| Client -> Server | `audio.chunk` | `data` | Base64 PCM16 16 kHz mono audio bytes. |
|
||||
| Client -> Server | `video.query` | `text` | Ask a question over the buffered frames and audio. |
|
||||
| Client -> Server | `video.done` | none | End the WebSocket session. |
|
||||
| Server -> Client | `response.start` | none | Query generation started. |
|
||||
| Server -> Client | `response.text.delta` | `delta` | Incremental text output. |
|
||||
| Server -> Client | `response.text.done` | `text` | Final text output for the query. |
|
||||
| Server -> Client | `response.audio.delta` | `data`, `format` | Incremental generated audio, base64 WAV. |
|
||||
| Server -> Client | `response.audio.done` | none | Audio output finished. |
|
||||
| Server -> Client | `session.done` | none | Session closed. |
|
||||
| Server -> Client | `error` | `message` | Recoverable protocol or generation error. |
|
||||
|
||||
### `session.config` Fields
|
||||
|
||||
| Field | Type | Default | Description |
|
||||
|-------|------|---------|-------------|
|
||||
| `model` | string or null | null | Optional model name. Usually omitted because the server hosts one model. |
|
||||
| `modalities` | list[string] | `["text", "audio"]` | Output modalities. Use `["text"]`, `["audio"]`, or both. |
|
||||
| `num_frames` | integer, 1-128 | `4` | Number of buffered frames sampled for each query. |
|
||||
| `max_frames` | integer, 1-256 | `50` | Maximum retained frame buffer size. Oldest frames are evicted first. |
|
||||
| `system_prompt` | string or null | null | Optional custom system prompt. |
|
||||
| `use_audio_in_video` | bool | `true` | Include streamed audio chunks in multimodal video understanding when audio is present. |
|
||||
| `sampling_params_list` | list or null | null | Optional per-stage sampling parameter overrides. |
|
||||
| `enable_frame_filter` | bool | `true` | Enable EVS near-duplicate frame filtering. |
|
||||
| `frame_filter_threshold` | float, 0.0-1.0 | `0.95` | EVS similarity threshold. Higher keeps more frames; lower drops more near-duplicates. |
|
||||
|
||||
### Legacy Aliases
|
||||
|
||||
The server accepts these legacy field names and rewrites them before validation. New clients should send the canonical names above.
|
||||
|
||||
| Legacy field | Canonical field |
|
||||
|--------------|-----------------|
|
||||
| `num_sample_frames` | `num_frames` |
|
||||
| `evs_enabled` | `enable_frame_filter` |
|
||||
| `evs_threshold` | `frame_filter_threshold` |
|
||||
|
||||
### Environment Variables
|
||||
|
||||
| Variable | Values | Default | Description |
|
||||
|----------|--------|---------|-------------|
|
||||
| `VLLM_VIDEO_ASYNC_CHUNK` | `on`, `off` | `on` | Wire-level streaming switch. `off` buffers server-side deltas and emits coalesced outputs at the end of a query. |
|
||||
| `VLLM_VIDEO_AUDIO_DELTA_MODE` | `fast`, `slow` | `fast` | Audio delta extraction strategy. `fast` emits only newly produced chunks; `slow` recomputes from accumulated audio and exists for A/B verification. |
|
||||
|
||||
## EVS Semantics
|
||||
|
||||
EVS compares downsampled frames and drops near-duplicate frames before they enter the session frame buffer. `frame_filter_threshold` controls retention: higher values are more permissive and keep more frames; lower values are more aggressive and drop more similar frames.
|
||||
|
||||
## Known Limitations
|
||||
|
||||
- Session KV reuse and incremental prefill are not implemented in this PR. Each `video.query` rebuilds the model prompt from the retained frame and audio buffers.
|
||||
- Back-to-back short replies can still expose an engine-layer scheduler race. The PR notes an observed workaround of at least 200 ms idle between turns when clients repeatedly see idle timeouts.
|
||||
- If the audio buffer exceeds the server limit, the server emits `Audio buffer overflow` and clears the currently buffered audio for the session.
|
||||
- The API is intended for Qwen3-Omni streaming video understanding; other models may not support the same multimodal processor arguments.
|
||||
@@ -0,0 +1,120 @@
|
||||
# Diffusion Attention Backends
|
||||
|
||||
This document describes the diffusion attention backends available in vLLM-Omni, how to select them, and how to use SageAttention.
|
||||
|
||||
## Overview
|
||||
|
||||
Diffusion attention backend selection is controlled by the `DIFFUSION_ATTENTION_BACKEND` environment variable and resolved in `vllm_omni.diffusion.attention.selector`.
|
||||
|
||||
This backend is used by diffusion attention layers such as the DiT attention in video and image generation models.
|
||||
|
||||
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 `DIFFUSION_ATTENTION_BACKEND` is unset, vLLM-Omni asks the current platform to choose the default backend. On CUDA, that normally means `FLASH_ATTN` when available, otherwise `TORCH_SDPA`.
|
||||
|
||||
## Backend Options
|
||||
|
||||
| Value | Notes |
|
||||
|---|---|
|
||||
| `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. |
|
||||
|
||||
## Selection Priority
|
||||
|
||||
Diffusion attention backend selection follows this order:
|
||||
|
||||
1. `DIFFUSION_ATTENTION_BACKEND`
|
||||
2. Platform default
|
||||
|
||||
Example:
|
||||
|
||||
```bash
|
||||
export DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN
|
||||
```
|
||||
|
||||
## SageAttention Installation
|
||||
|
||||
vLLM-Omni expects SageAttention to be installed into the same Python environment as vLLM-Omni.
|
||||
|
||||
Build from source:
|
||||
|
||||
```bash
|
||||
git clone https://github.com/thu-ml/SageAttention.git
|
||||
cd SageAttention
|
||||
|
||||
export EXT_PARALLEL=4 NVCC_APPEND_FLAGS="--threads 8" MAX_JOBS=32
|
||||
pip install . --no-build-isolation
|
||||
```
|
||||
|
||||
Quick check:
|
||||
|
||||
```bash
|
||||
python -c "import sageattention; print(sageattention.__file__)"
|
||||
```
|
||||
|
||||
## Usage
|
||||
|
||||
### Enable SageAttention
|
||||
|
||||
Example: HunyuanVideo-1.5 text-to-video
|
||||
|
||||
```bash
|
||||
DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN 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 30 --seed 42 --guidance-scale 6.0 \
|
||||
--tensor-parallel-size 2 \
|
||||
--output ../tmp/hv15_modelopt_sage.mp4
|
||||
```
|
||||
|
||||
Example: Wan2.2 TI2V 5B
|
||||
|
||||
```bash
|
||||
DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
--model Wan-AI/Wan2.2-TI2V-5B-Diffusers \
|
||||
--prompt "A dog running across a field of golden wheat." \
|
||||
--height 704 --width 1280 --num-frames 49 \
|
||||
--num-inference-steps 30 --seed 42 --guidance-scale 5.0 \
|
||||
--tensor-parallel-size 2 \
|
||||
--output outputs/wan22_sage.mp4
|
||||
```
|
||||
|
||||
### Compare Against FlashAttention
|
||||
|
||||
Unset the backend override, or explicitly use `FLASH_ATTN`:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
--model Wan-AI/Wan2.2-TI2V-5B-Diffusers \
|
||||
--prompt "A dog running across a field of golden wheat." \
|
||||
--height 704 --width 1280 --num-frames 49 \
|
||||
--num-inference-steps 30 --seed 42 --guidance-scale 5.0 \
|
||||
--tensor-parallel-size 2 \
|
||||
--output outputs/wan22_fa3.mp4
|
||||
```
|
||||
|
||||
## Validation Guidance
|
||||
|
||||
Do not assume that a faster attention backend is numerically interchangeable with `FLASH_ATTN`.
|
||||
|
||||
Always compare:
|
||||
|
||||
- End-to-end runtime
|
||||
- DiT / diffusion stage runtime
|
||||
- Output quality against a known-good baseline
|
||||
|
||||
At minimum, keep the same:
|
||||
|
||||
- model
|
||||
- prompt
|
||||
- seed
|
||||
- resolution
|
||||
- frame count
|
||||
- inference steps
|
||||
- parallel config
|
||||
@@ -283,3 +283,10 @@ Using Cache-DiT acceleration:
|
||||
|
||||
1. ✅ **Enable Cache-DiT** - Set `cache_backend="cache_dit"` to get 1.5x-3x speedup with optimized defaults
|
||||
2. ✅ **(Optional) Customize** - Adjust `cache_config` parameters for specific speed/quality trade-offs
|
||||
|
||||
---
|
||||
|
||||
## Additional Resources
|
||||
|
||||
- [Cache-DiT documentation](https://cache-dit.readthedocs.io/en/latest/)
|
||||
- [Cache-DiT API reference](https://cache-dit.readthedocs.io/en/latest/user_guide/CACHE_API/)
|
||||
@@ -36,6 +36,45 @@ m = Omni(model="Wan-AI/Wan2.2-T2V-A14B-Diffusers", enable_cpu_offload=True)
|
||||
vllm-omni serve diffusion Wan-AI/Wan2.2-T2V-A14B-Diffusers --enable-cpu-offload
|
||||
```
|
||||
|
||||
### To Support a Model
|
||||
|
||||
Implement the `SupportsModuleOffload` protocol to declare which
|
||||
submodules participate in offloading:
|
||||
|
||||
```python
|
||||
from typing import ClassVar
|
||||
from vllm_omni.diffusion.models.interface import SupportsModuleOffload
|
||||
|
||||
class MyPipeline(nn.Module, SupportsModuleOffload):
|
||||
_dit_modules: ClassVar[list[str]] = ["transformer"]
|
||||
_encoder_modules: ClassVar[list[str]] = ["text_encoder", "vision_model"]
|
||||
_vae_modules: ClassVar[list[str]] = ["vae"]
|
||||
_resident_modules: ClassVar[list[str]] = [] # optional
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.transformer = ... # DiT — stays on GPU during denoising
|
||||
self.text_encoder = ... # Encoder — offloaded to CPU during denoising
|
||||
self.vision_model = ... # Encoder — offloaded to CPU during denoising
|
||||
self.vae = ... # VAE — always on GPU
|
||||
```
|
||||
|
||||
- `_dit_modules`: attribute names of denoising submodules (kept on GPU
|
||||
during the diffusion loop).
|
||||
- `_encoder_modules`: attribute names of encoder/vision submodules
|
||||
(offloaded to CPU during the diffusion loop).
|
||||
- `_vae_modules`: attribute names of VAE(s) (always kept on GPU, not
|
||||
part of the mutual exclusion hooks).
|
||||
- `_resident_modules`: attribute names of small submodules that must
|
||||
stay on GPU during layerwise offloading (e.g. embedders, connectors).
|
||||
Optional — defaults to `[]`.
|
||||
|
||||
All attribute names support dotted paths for nested submodules
|
||||
(e.g. `"pipe.transformer"`, `"bagel.time_embedder"`).
|
||||
|
||||
Both DiT and encoder lists are needed because the offload hooks use
|
||||
mutual exclusion: when one group runs, the other moves to CPU.
|
||||
|
||||
### Limitations
|
||||
- Cold start latency increases
|
||||
- Adds overhead from CPU-GPU transfers between encoder and denoising phases
|
||||
@@ -116,11 +155,19 @@ class Flux2Transformer2DModel(nn.Module):
|
||||
|
||||
**Module Discovery**
|
||||
|
||||
The offloader automatically discovers pipeline components:
|
||||
The offloader discovers pipeline components in two ways:
|
||||
|
||||
- **DiT modules**: `transformer`, `transformer_2`, `dit`
|
||||
- **Encoders**: `text_encoder`, `text_encoder_2`, `text_encoder_3`, `image_encoder`
|
||||
- **VAE**: `vae`
|
||||
1. **Protocol-based** (preferred): If the pipeline implements
|
||||
`SupportsModuleOffload`, its `_dit_modules`, `_encoder_modules`,
|
||||
`_vae_modules`, and `_resident_modules` class variables are used
|
||||
directly. All attribute names support dotted paths (e.g.
|
||||
`"pipe.transformer"`, `"bagel.time_embedder"`) for nested submodules.
|
||||
|
||||
2. **Fallback attribute scan**: Otherwise, the offloader scans for
|
||||
well-known attribute names:
|
||||
- **DiT modules**: `transformer`, `transformer_2`, `dit`, `sr_dit`, `language_model`, `transformer_blocks`, `model`
|
||||
- **Encoders**: `text_encoder`, `text_encoder_2`, `text_encoder_3`, `image_encoder`
|
||||
- **VAE**: `vae`, `audio_vae`
|
||||
|
||||
**Hook System**
|
||||
|
||||
@@ -148,6 +195,7 @@ Factory function `get_offload_backend()` selects the appropriate backend based o
|
||||
| StableDiffusion3Pipeline | `stabilityai/stable-diffusion-3.5-medium` | `SD3Transformer2DModel` | - | ✓ | `"transformer_blocks"` |
|
||||
| Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` |
|
||||
| Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` |
|
||||
| BagelPipeline | `ByteDance-Seed/BAGEL-7B-MoT` | `Qwen2MoTModel` | - | ✓ | `"layers"`, `"customized modules"` |
|
||||
|
||||
**Notes:**
|
||||
- Model-Level Offloading is expected to be supported by all common diffusion models (DiT and encoders) naturally
|
||||
|
||||
@@ -72,6 +72,8 @@ At load time:
|
||||
| Model | HF Checkpoint | Scheme | Group Size | Backend |
|
||||
|-------|--------------|--------|------------|---------|
|
||||
| FLUX.1-dev | `vllm-project-org/FLUX.1-dev-AutoRound-w4a16` | W4A16 | 128 | GPTQ-Marlin |
|
||||
| Qwen2.5-Omni-7B | `Intel/Qwen2.5-Omni-7B-int4-AutoRound` | W4A16 | 128 | GPTQ-Marlin |
|
||||
| Qwen3-Omni-30B-A3B-Instruct | `Intel/Qwen3-Omni-30B-A3B-Instruct-int4-AutoRound` | W4A16 | 128 | GPTQ-Marlin |
|
||||
|
||||
## Creating a Quantized Checkpoint
|
||||
|
||||
|
||||
@@ -65,6 +65,7 @@ The available `ignored_layers` names depend on the model architecture (e.g., `to
|
||||
| Flux | `black-forest-labs/FLUX.1-dev` | All layers | None |
|
||||
| HunyuanImage-3 | `tencent/HunyuanImage3` | All layers | None |
|
||||
| HunyuanVideo-1.5 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `720p_t2v`, `480p_i2v` | All layers | None |
|
||||
| GLM-Image | `zai-org/GLM-Image` | All layers | None |
|
||||
| Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | All layers | None |
|
||||
|
||||
## Combining with Other Features
|
||||
|
||||
@@ -0,0 +1,56 @@
|
||||
# msModelSlim Quantization
|
||||
|
||||
## Overview
|
||||
|
||||
[msModelSlim](https://github.com/Ascend/msmodelslim) is an Ascend-friendly compression tool focused on acceleration, using compression techniques, and built for Ascend hardware. It includes a series of inference optimization technologies such as quantization and compression, aiming to accelerate large language dense models, MoE models, multimodal understanding models, multimodal generation models, etc.
|
||||
|
||||
Once you have a quantized model which is generated by **msModelSlim**, you can use vLLM Omni for inference by specifying the --quantization ascend parameter to enable quantization features.
|
||||
|
||||
### Supported Schemes
|
||||
|
||||
| Scheme | Bits | Status |
|
||||
|--------|------|--------|
|
||||
| W8A8 | 8 | ✅ Supported |
|
||||
| W4A4 | 4 | Planned |
|
||||
|
||||
W8A8 is the first supported scheme. Additional schemes will be added in future releases.
|
||||
|
||||
## Model Quantization
|
||||
|
||||
The following example shows how to generate W8A8 quantized weights for the [Wan2_2 model](https://gitcode.com/Ascend/msmodelslim/blob/master/example/multimodal_sd/Wan2_2/README.md).
|
||||
|
||||
**Quantization Script:**
|
||||
|
||||
```bash
|
||||
msmodelslim quant \
|
||||
--model_path /path/to/wan2_2_t2v_float_weights \
|
||||
--save_path /path/to/wan2_2_t2v_quantized_weights \
|
||||
--device npu \
|
||||
--model_type Wan2_2 \
|
||||
--config_path /lab_practice/wan2_2/wan2_2_w8a8f8_mxfp_t2v.yaml \
|
||||
--trust_remote_code True
|
||||
```
|
||||
|
||||
After quantization completes, the output directory will contain the quantized model files.
|
||||
|
||||
For more examples, refer to the [official examples](https://gitcode.com/Ascend/msit/tree/master/msmodelslim/example).
|
||||
|
||||
## Configuration
|
||||
|
||||
1. **CLI**: pass `--quantization ascend`.
|
||||
|
||||
```bash
|
||||
# Offline inference
|
||||
python text_to_image.py --model <your-model> --quantization ascend
|
||||
|
||||
# Online serving
|
||||
vllm serve <your-model> --omni --quantization ascend
|
||||
```
|
||||
|
||||
## Supported Models
|
||||
|
||||
| Model | HF Models | Recommendation | `ignored_layers` |
|
||||
|-------|-----------|---------------|------------------|
|
||||
| HunyuanImage-3.0 | - | All layers | None |
|
||||
|
||||
Currently, quantized HunyuanImage-3.0 weights have not been uploaded to public model platforms such as Hugging Face. You can use a [HunyuanImage-3.0-adapted msModelSlim version](https://gitcode.com/betta18/msmodelslim/tree/hyimage3_mxfp8) to generate the quantized weights manually. We will upload the quantized weights as soon as possible.
|
||||
@@ -107,7 +107,7 @@ The following tables show which models support each feature:
|
||||
|
||||
| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
|
||||
|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:|
|
||||
| **Bagel** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **Bagel** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **FLUX.1-dev** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
|
||||
| **FLUX.1-schnell** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
|
||||
| **FLUX.2-klein** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
|
||||
@@ -141,6 +141,7 @@ The following tables show which models support each feature:
|
||||
| **Wan2.2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (encode/decode) | ❌ | ❌ |
|
||||
| **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **LTX-2.3** | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **Helios** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ |
|
||||
| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
@@ -148,13 +149,13 @@ The following tables show which models support each feature:
|
||||
**Frame Interpolation Support**
|
||||
|
||||
- **Supported**: Wan2.2 text-to-video, image-to-video, and TI2V pipelines
|
||||
- **Not supported**: Wan2.1-VACE, LTX-2, Helios, HunyuanVideo-1.5, DreamID-Omni
|
||||
- **Not supported**: Wan2.1-VACE, LTX-2, LTX-2.3, Helios, HunyuanVideo-1.5, DreamID-Omni
|
||||
|
||||
### AudioGen
|
||||
|
||||
| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
|
||||
|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:|
|
||||
| **Stable-Audio-Open** | ✅ | ❌ | ❓ | ❓ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ |
|
||||
| **Stable-Audio-Open** | ✅ | ❌ | ❓ | ❓ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ |
|
||||
|
||||
|
||||
## Feature Compatibility
|
||||
|
||||
@@ -2,46 +2,61 @@
|
||||
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/bagel>.
|
||||
|
||||
|
||||
## Set up
|
||||
## Setup
|
||||
|
||||
Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
|
||||
|
||||
## Run examples
|
||||
## Architecture
|
||||
|
||||
**Note**: These examples work with the default configuration on an **NVIDIA A100 (80GB)**. We also tested on dual **NVIDIA RTX 5000 Ada (32GB each)**. For dual-GPU setups, please modify the stage configuration to distribute the model across devices.
|
||||
BAGEL-7B-MoT is a Mixture-of-Transformers (MoT) model supporting both image generation and understanding. It offers two deployment topologies:
|
||||
|
||||
Get into the bagel folder
|
||||
| Topology | Stages | Description |
|
||||
| :------- | :----- | :---------- |
|
||||
| **Two-stage** (default) | Stage 0 (Thinker, AR) + Stage 1 (DiT, Diffusion) | Thinker handles text/understanding via vLLM AR engine; DiT handles image generation. KV cache is transferred between stages. |
|
||||
| **Single-stage** | Stage 0 (DiT, Diffusion) only | The DiT stage contains a full LLM, ViT, VAE, and tokenizer internally. All modalities are handled within a single diffusion process. |
|
||||
|
||||
Both topologies support all four modalities: `text2img`, `img2img`, `img2text`, `text2text`.
|
||||
|
||||
## Quick Start
|
||||
|
||||
```bash
|
||||
cd examples/offline_inference/bagel
|
||||
|
||||
# Default two-stage mode (auto-detected)
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2img \
|
||||
--prompts "A cute cat"
|
||||
|
||||
# Single-stage mode
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2img \
|
||||
--prompts "A cute cat" \
|
||||
--deploy-config vllm_omni/deploy/bagel_single_stage.yaml
|
||||
```
|
||||
|
||||
### Modality Control
|
||||
> **Note**: These examples work with the default configuration on an **NVIDIA A100 (80GB)**. For dual-GPU setups, modify the deploy YAML to distribute stages across devices.
|
||||
|
||||
BAGEL-7B-MoT supports multiple modality modes. You can control the mode using the `--modality` argument:
|
||||
## Modality Control
|
||||
|
||||
#### Text to Image (text2img)
|
||||
Control the mode using the `--modality` argument:
|
||||
|
||||
- **Pipeline**: Text → Thinker → DiT → VAE Decode → Image
|
||||
- **Stages Used**: Stage 0 (Thinker) + Stage 1 (DiT)
|
||||
- **KV Transfer**: Thinker sends KV cache to DiT for conditioned generation
|
||||
| Modality | Input | Output | Description |
|
||||
| :------- | :---- | :----- | :---------- |
|
||||
| `text2img` | Text | Image | Generate images from text prompts |
|
||||
| `img2img` | Image + Text | Image | Transform images using text guidance |
|
||||
| `img2text` | Image + Text | Text | Generate text descriptions from images |
|
||||
| `text2text` | Text | Text | Pure text generation (language model mode) |
|
||||
|
||||
Generate images from text prompts:
|
||||
### Text to Image (text2img)
|
||||
|
||||
```bash
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2img \
|
||||
--prompts "A cute cat"
|
||||
--prompts "A cute cat" \
|
||||
--steps 50
|
||||
```
|
||||
|
||||
#### Image to Image (img2img)
|
||||
|
||||
- **Pipeline**: Image → VAE Encode → DiT → VAE Decode → New Image
|
||||
- **Stages Used**: Stage 1 (DiT) only
|
||||
- **Special**: Bypasses the Thinker stage, direct image-to-image transformation
|
||||
|
||||
Transform images based on text prompts:
|
||||
### Image to Image (img2img)
|
||||
|
||||
```bash
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
@@ -50,13 +65,7 @@ python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--prompts "Let the woman wear a blue dress"
|
||||
```
|
||||
|
||||
#### Image to Text (img2text)
|
||||
|
||||
- **Pipeline**: Image → ViT + VAE Encode → Thinker → Text Output
|
||||
- **Stages Used**: Stage 0 (Thinker) only
|
||||
- **Special**: Uses both VAE latent encoding AND ViT semantic encoding for comprehensive image understanding
|
||||
|
||||
Generate text descriptions from images:
|
||||
### Image to Text (img2text)
|
||||
|
||||
```bash
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
@@ -65,196 +74,210 @@ python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--prompts "Describe this image in detail"
|
||||
```
|
||||
|
||||
#### Text to Text (text2text)
|
||||
|
||||
- **Pipeline**: Text → Thinker → Text Output
|
||||
- **Stages Used**: Stage 0 (Thinker) only
|
||||
- **Special**: No visual components involved, operates as pure language model
|
||||
|
||||
Pure text generation:
|
||||
### Text to Text (text2text)
|
||||
|
||||
```bash
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2text \
|
||||
--prompts "What is the capital of France?"
|
||||
|
||||
# You can load prompts from a text file (one prompt per line):
|
||||
# Load prompts from a text file (one prompt per line):
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2text \
|
||||
--txt-prompts /path/to/prompts.txt
|
||||
```
|
||||
|
||||
### Inference Steps
|
||||
## Think Mode
|
||||
|
||||
Control the number of inference steps for image generation:
|
||||
Think mode enables the model to generate `<think>...</think>` planning/reasoning tokens before producing the final output. This improves generation quality for complex prompts.
|
||||
|
||||
- **Two-stage**: The Thinker (AR) stage decodes think tokens, then transfers the augmented KV cache to the DiT stage for image generation.
|
||||
- **Single-stage**: The DiT's internal LLM generates think tokens in-place before proceeding to denoise.
|
||||
|
||||
```bash
|
||||
# You can adjust steps to 100 to improve image quality
|
||||
# Think + text2img: plan before generating
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2img \
|
||||
--steps 50 \
|
||||
--prompts "A cute cat"
|
||||
```
|
||||
--prompts "A futuristic city with flying cars" \
|
||||
--think \
|
||||
--max-think-tokens 1000
|
||||
|
||||
### Key arguments
|
||||
|
||||
BAGEL-7B-MoT supports **multiple modality modes** for different use cases.
|
||||
|
||||
The default yaml configuration deploys Thinker and DiT on the same GPU. You can use the default configuration file: [`bagel.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/model_executor/stage_configs/bagel.yaml)
|
||||
|
||||
#### 📌 Command Line Arguments (end2end.py)
|
||||
|
||||
| Argument | Type | Default | Description |
|
||||
| :--------------------- | :----- | :---------------------------- | :----------------------------------------------------------- |
|
||||
| `--model` | string | `ByteDance-Seed/BAGEL-7B-MoT` | Model path or name |
|
||||
| `--modality` | choice | `text2img` | Modality mode: `text2img`, `img2img`, `img2text`, `text2text` |
|
||||
| `--prompts` | list | `None` | Input text prompts directly |
|
||||
| `--txt-prompts` | string | `None` | Path to txt file with one prompt per line |
|
||||
| `--image-path` | string | `None` | Input image path (for `img2img`/`img2text`) |
|
||||
| `--steps` | int | `50` | Number of inference steps |
|
||||
| `--stage-configs-path` | string | `None` | Custom stage config file path |
|
||||
| `--worker-backend` | choice | `process` | Worker backend: `process` or `ray` |
|
||||
| `--ray-address` | string | `None` | Ray cluster address |
|
||||
| `--enable-stats` | flag | `False` | Enable statistics logging |
|
||||
| `--init-sleep-seconds` | int | `20` | Initialization sleep time |
|
||||
| `--batch-timeout` | int | `5` | Batch timeout |
|
||||
| `--init-timeout` | int | `300` | Initialization timeout |
|
||||
|
||||
------
|
||||
|
||||
#### ⚙️ Stage Configuration Parameters (bagel.yaml)
|
||||
|
||||
**Stage 0 - Thinker (LLM Stage)**
|
||||
|
||||
| Parameter | Value | Description |
|
||||
| :------------------------------- | :------------------------------ | :----------------------- |
|
||||
| `stage_type` | `llm` | Stage type |
|
||||
| `devices` | `"0"` | GPU device ID |
|
||||
| `max_num_seqs` | `1` | Maximum batch size |
|
||||
| `model_stage` | `thinker` | Model stage identifier |
|
||||
| `model_arch` | `BagelForConditionalGeneration` | Model architecture |
|
||||
| `gpu_memory_utilization` | `0.4` | GPU memory utilization |
|
||||
| `tensor_parallel_size` | `1` | Tensor parallel size |
|
||||
| `max_num_batched_tokens` | `32768` | Maximum batched tokens |
|
||||
| `omni_kv_config.need_send_cache` | `true` | Whether to send KV cache |
|
||||
|
||||
------
|
||||
|
||||
**Stage 1 - DiT (Diffusion Stage)**
|
||||
|
||||
| Parameter | Value | Description |
|
||||
| :------------------------------- | :---------- | :-------------------------- |
|
||||
| `stage_type` | `diffusion` | Stage type |
|
||||
| `devices` | `"0"` | GPU device ID |
|
||||
| `max_num_seqs` | `1` | Maximum batch size |
|
||||
| `model_stage` | `dit` | Model stage identifier |
|
||||
| `gpu_memory_utilization` | `0.4` | GPU memory utilization |
|
||||
| `omni_kv_config.need_recv_cache` | `true` | Whether to receive KV cache |
|
||||
| `engine_input_source` | `[0]` | Input source from Stage 0 |
|
||||
|
||||
------
|
||||
|
||||
#### Tensor Parallelism (TP)
|
||||
|
||||
For larger models or multi-GPU environments, you can enable Tensor Parallelism (TP) by modifying the stage configuration (e.g., [`bagel.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/model_executor/stage_configs/bagel.yaml)).
|
||||
|
||||
1. **Set `tensor_parallel_size`**: Increase this value (e.g., to `2` or `4`).
|
||||
2. **Set `devices`**: Specify the comma-separated GPU IDs to be used for the stage (e.g., `"0,1"`).
|
||||
|
||||
Example configuration for TP=2 on GPUs 0 and 1:
|
||||
```yaml
|
||||
engine_args:
|
||||
tensor_parallel_size: 2
|
||||
...
|
||||
runtime:
|
||||
devices: "0,1"
|
||||
```
|
||||
|
||||
------
|
||||
|
||||
#### 🔗 Runtime Configuration
|
||||
|
||||
| Parameter | Value | Description |
|
||||
| :-------------------- | :------ | :------------------------------- |
|
||||
| `shm_threshold_bytes` | `65536` | Shared memory threshold (64KB) |
|
||||
|
||||
## Using Mooncake Connector
|
||||
|
||||
[Mooncake](https://github.com/kvcache-ai/Mooncake) is a high-performance distributed KV cache transfer engine that enables efficient cross-node data movement via TCP or RDMA, making it ideal for multi-node disaggregated inference.
|
||||
|
||||
By default, BAGEL uses `SharedMemoryConnector` for inter-stage communication. You can switch to the Mooncake connector for better performance on multi-GPU setups and to enable multi-node deployment.
|
||||
|
||||
### Prerequisites
|
||||
|
||||
Install the Mooncake transfer engine:
|
||||
|
||||
```bash
|
||||
# For CUDA-enabled systems (recommended)
|
||||
pip install mooncake-transfer-engine
|
||||
|
||||
# For non-CUDA systems
|
||||
pip install mooncake-transfer-engine-non-cuda
|
||||
```
|
||||
|
||||
### Step 1: Start the Mooncake Master
|
||||
|
||||
On the **primary node**, start the Mooncake master service (run in a separate terminal or background with `&`):
|
||||
|
||||
```bash
|
||||
# Optional: enable disk-backed storage by creating a directory and passing --root_fs_dir.
|
||||
# Without it, Mooncake runs in memory-only mode, which is sufficient for KV cache transfer.
|
||||
mkdir -p ./mc_storage
|
||||
|
||||
mooncake_master \
|
||||
--rpc_port=50051 \
|
||||
--enable_http_metadata_server=true \
|
||||
--http_metadata_server_host=0.0.0.0 \
|
||||
--http_metadata_server_port=8080 \
|
||||
--metrics_port=9003 \
|
||||
--root_fs_dir=./mc_storage/ \
|
||||
--cluster_id=mc-local-1 &
|
||||
```
|
||||
|
||||
### Step 2: Run Offline Inference with Mooncake
|
||||
|
||||
Use the provided Mooncake stage config [`bagel_multiconnector.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml). Before launching, update the `metadata_server` and `master` addresses in the YAML to match your Mooncake master node's IP (use `127.0.0.1` for single-node testing).
|
||||
|
||||
```bash
|
||||
cd examples/offline_inference/bagel
|
||||
|
||||
# Text to Image with Mooncake
|
||||
# Think + img2img: reason about the edit
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2img \
|
||||
--prompts "A cute cat" \
|
||||
--stage-configs-path ../../../vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
|
||||
--modality img2img \
|
||||
--image-path /path/to/image.jpg \
|
||||
--prompts "Make it look like a watercolor painting" \
|
||||
--think
|
||||
|
||||
# Image to Text with Mooncake
|
||||
# Think + img2text: reason before describing
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality img2text \
|
||||
--image-path /path/to/image.jpg \
|
||||
--prompts "Describe this image" \
|
||||
--stage-configs-path ../../../vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
|
||||
--prompts "What is happening in this image?" \
|
||||
--think
|
||||
|
||||
# Text to Text with Mooncake
|
||||
# Think + text2text: chain-of-thought reasoning
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2text \
|
||||
--prompts "What is the capital of France?" \
|
||||
--stage-configs-path ../../../vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
|
||||
--prompts "Solve: 23 * 47" \
|
||||
--think
|
||||
```
|
||||
|
||||
For more details on the Mooncake connector and multi-node setup, see the [Mooncake Store Connector documentation](https://github.com/vllm-project/vllm-omni/tree/main/docs/design/feature/omni_connectors/mooncake_store_connector.md).
|
||||
Think mode parameters:
|
||||
|
||||
------
|
||||
| Argument | Default | Description |
|
||||
| :------- | :------ | :---------- |
|
||||
| `--think` | `False` | Enable thinking mode |
|
||||
| `--max-think-tokens` | `1000` | Maximum tokens for think generation |
|
||||
| `--do-sample` | `False` | Enable sampling (vs. greedy) for text generation |
|
||||
| `--text-temperature` | `0.3` | Temperature for text generation sampling |
|
||||
|
||||
## Classifier-Free Guidance (CFG)
|
||||
|
||||
CFG controls the trade-off between prompt fidelity and diversity. These parameters apply to image generation modalities (`text2img`, `img2img`).
|
||||
|
||||
```bash
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2img \
|
||||
--prompts "A photorealistic portrait" \
|
||||
--cfg-text-scale 6.0 \
|
||||
--cfg-img-scale 2.0 \
|
||||
--negative-prompt "blurry, low quality, distorted" \
|
||||
--cfg-interval 0.4 1.0 \
|
||||
--cfg-renorm-type global \
|
||||
--cfg-renorm-min 0.0
|
||||
```
|
||||
|
||||
| Argument | Default | Description |
|
||||
| :------- | :------ | :---------- |
|
||||
| `--cfg-text-scale` | `4.0` | Text CFG scale (higher = more prompt-adherent) |
|
||||
| `--cfg-img-scale` | `1.5` | Image CFG scale (for img2img) |
|
||||
| `--negative-prompt` | `None` | Negative prompt for CFG conditioning |
|
||||
| `--cfg-interval` | pipeline default | CFG active interval `[start, end]` as fractions of total timesteps |
|
||||
| `--cfg-renorm-type` | `None` | Renormalization type: `global`, `text_channel`, `channel` |
|
||||
| `--cfg-renorm-min` | `None` | Minimum renormalization value |
|
||||
| `--cfg-parallel-size` | `1` | CFG parallel size: `1` = batched (single GPU), `2` = 2-branch parallel, `3` = full 3-GPU parallel |
|
||||
|
||||
## Deployment Topologies
|
||||
|
||||
### Two-Stage (Default)
|
||||
|
||||
The default topology auto-detected from the model. No extra flags needed.
|
||||
|
||||
```bash
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2img \
|
||||
--prompts "A cute cat"
|
||||
```
|
||||
|
||||
The pipeline is defined in [`bagel.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/deploy/bagel.yaml). Stage 0 (Thinker) and Stage 1 (DiT) share GPU 0 by default. For dual-GPU setups, customize the deploy YAML and set `devices: "1"` for stage 1.
|
||||
|
||||
### Single-Stage
|
||||
|
||||
Pass the single-stage deploy config via `--deploy-config`:
|
||||
|
||||
```bash
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2img \
|
||||
--prompts "A cute cat" \
|
||||
--deploy-config vllm_omni/deploy/bagel_single_stage.yaml
|
||||
```
|
||||
|
||||
See [`bagel_single_stage.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/deploy/bagel_single_stage.yaml) for configuration details. The `pipeline: bagel_single_stage` field selects the single-stage topology from the pipeline registry.
|
||||
|
||||
### Tensor Parallelism (TP)
|
||||
|
||||
For larger models or multi-GPU environments, customize the deploy YAML (see [`bagel.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/deploy/bagel.yaml)) and set per-stage `tensor_parallel_size` and `devices`:
|
||||
|
||||
```yaml
|
||||
# Example: TP=2 on GPUs 0,1 for the Thinker stage
|
||||
stages:
|
||||
- stage_id: 0
|
||||
tensor_parallel_size: 2
|
||||
devices: "0,1"
|
||||
```
|
||||
|
||||
Then pass the custom deploy YAML:
|
||||
|
||||
```bash
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2img \
|
||||
--prompts "A cute cat" \
|
||||
--deploy-config /path/to/custom_bagel.yaml
|
||||
```
|
||||
|
||||
### FP8 Quantization
|
||||
|
||||
```bash
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2img \
|
||||
--prompts "A cute cat" \
|
||||
--quantization fp8
|
||||
```
|
||||
|
||||
## Command Line Reference
|
||||
|
||||
### Core Arguments
|
||||
|
||||
| Argument | Type | Default | Description |
|
||||
| :------- | :--- | :------ | :---------- |
|
||||
| `--model` | string | `ByteDance-Seed/BAGEL-7B-MoT` | Model path or HuggingFace name |
|
||||
| `--modality` | choice | `text2img` | `text2img`, `img2img`, `img2text`, `text2text` |
|
||||
| `--prompts` | list | `None` | Input text prompts |
|
||||
| `--txt-prompts` | string | `None` | Path to text file with one prompt per line |
|
||||
| `--image-path` | string | `None` | Input image path (required for `img2img`/`img2text`) |
|
||||
| `--output` | string | `.` | Output directory for saved images |
|
||||
| `--steps` | int | `50` | Number of diffusion inference steps |
|
||||
| `--seed` | int | `None` | Random seed for reproducibility |
|
||||
|
||||
### Think Mode Arguments
|
||||
|
||||
| Argument | Type | Default | Description |
|
||||
| :------- | :--- | :------ | :---------- |
|
||||
| `--think` | flag | `False` | Enable `<think>...</think>` planning/reasoning |
|
||||
| `--max-think-tokens` | int | `1000` | Maximum tokens for think generation |
|
||||
| `--do-sample` | flag | `False` | Use sampling instead of greedy decoding |
|
||||
| `--text-temperature` | float | `0.3` | Sampling temperature for text generation |
|
||||
|
||||
### CFG Arguments
|
||||
|
||||
| Argument | Type | Default | Description |
|
||||
| :------- | :--- | :------ | :---------- |
|
||||
| `--cfg-text-scale` | float | `4.0` | Text CFG guidance scale |
|
||||
| `--cfg-img-scale` | float | `1.5` | Image CFG guidance scale |
|
||||
| `--negative-prompt` | string | `None` | Negative prompt for CFG |
|
||||
| `--cfg-parallel-size` | int | `1` | CFG parallel GPU count (1, 2, or 3) |
|
||||
| `--cfg-interval` | float[2] | pipeline default | CFG active window `[start, end]` |
|
||||
| `--cfg-renorm-type` | string | `None` | `global`, `text_channel`, or `channel` |
|
||||
| `--cfg-renorm-min` | float | `None` | Minimum renormalization value |
|
||||
|
||||
### Engine Arguments
|
||||
|
||||
| Argument | Type | Default | Description |
|
||||
| :------- | :--- | :------ | :---------- |
|
||||
| `--deploy-config` | string | `None` | Path to deploy YAML (auto-detected if omitted) |
|
||||
| `--stage-configs-path` | string | `None` | [Deprecated] Legacy path to `stage_args` YAML; prefer `--deploy-config` |
|
||||
| `--worker-backend` | choice | `process` | `process` or `ray` |
|
||||
| `--ray-address` | string | `None` | Ray cluster address |
|
||||
| `--quantization` | string | `None` | Quantization method (e.g. `fp8`) |
|
||||
| `--log-stats` | flag | `False` | Enable statistics logging |
|
||||
| `--init-timeout` | int | `300` | Initialization timeout (seconds) |
|
||||
| `--batch-timeout` | int | `5` | Batch timeout (seconds) |
|
||||
| `--enable-diffusion-pipeline-profiler` | flag | `False` | Profile diffusion stage durations |
|
||||
|
||||
## FAQ
|
||||
|
||||
- If you don’t know how much VRAM is needed for the model or encounter the OOM error, you can try to decrease the max_model_len.
|
||||
- If you encounter OOM errors, try decreasing `max_model_len` or `gpu_memory_utilization` in the deploy YAML.
|
||||
|
||||
| Stage | VRAM |
|
||||
| :------------------ | :--------------------------- |
|
||||
| Stage-0 (Thinker) | **15.04 GiB** **+ KV Cache** |
|
||||
| Stage-1 (DiT) | **26.50 GiB** |
|
||||
| Total | **~42 GiB + KV Cache** |
|
||||
**Two-stage VRAM usage:**
|
||||
|
||||
| Stage | VRAM |
|
||||
| :---- | :--- |
|
||||
| Stage 0 (Thinker) | **15.04 GiB + KV Cache** |
|
||||
| Stage 1 (DiT) | **26.50 GiB** |
|
||||
| Total | **~42 GiB + KV Cache** |
|
||||
|
||||
**Single-stage VRAM usage:** The DiT loads the full model (~42 GiB) in one process.
|
||||
|
||||
## Example materials
|
||||
|
||||
|
||||
@@ -61,10 +61,17 @@ Key components live in `vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py
|
||||
- Stage 0 uses `CosyVoice3LM` and outputs speech tokens + conditioning features.
|
||||
- Stage 1 runs the flow model (DiT-based CFM) and HiFiGAN to synthesize waveform.
|
||||
|
||||
Stage wiring is configured in `vllm_omni/model_executor/stage_configs/cosyvoice3.yaml`.
|
||||
Pipeline topology lives in `vllm_omni/model_executor/models/cosyvoice3/pipeline.py`;
|
||||
runtime tunables (batch size, memory limits, sampling) live in
|
||||
`vllm_omni/deploy/cosyvoice3.yaml`. The deploy config auto-loads by
|
||||
HF `model_type` and defaults to `async_chunk: true` (shared-memory
|
||||
streaming). Pass `--no-async-chunk` on `vllm serve` to switch to the
|
||||
legacy sync path where stage 1 runs `text2flow` over the full
|
||||
speech-token sequence.
|
||||
|
||||
- Stage 0 emits latent speech tokens.
|
||||
- Stage 1 consumes them via `custom_process_input_func` and outputs audio.
|
||||
- Stage 1 consumes them via `sync_process_input_func` (sync mode) or the
|
||||
shared-memory connector (async-chunk mode) and outputs audio.
|
||||
|
||||
## Example materials
|
||||
|
||||
|
||||
@@ -38,7 +38,6 @@ Run a single sample for basic TTS:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type tts_sft
|
||||
```
|
||||
@@ -47,7 +46,6 @@ Run batch samples for basic TTS:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type tts_sft \
|
||||
--num-prompts {batch_size}
|
||||
@@ -65,7 +63,6 @@ Generate speech from text input:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type tts_sft \
|
||||
--text "The weather is so nice today."
|
||||
@@ -77,7 +74,6 @@ Generate speech with explicit voice style instructions:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type tts_sft_with_instruct \
|
||||
--text "The weather is so nice today." \
|
||||
@@ -90,7 +86,6 @@ Generate speech using an audio reference for voice cloning:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type tts_sft_with_audio \
|
||||
--text "The weather is so nice today." \
|
||||
@@ -103,7 +98,6 @@ Generate speech from text containing natural voice descriptions:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type tts_sft_with_natural_instruction \
|
||||
--text "In a panting young male voice, he said: I can't run anymore, wait for me!"
|
||||
@@ -115,7 +109,6 @@ Transcribe audio to text:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type audio_trancribing_sft \
|
||||
--audio-path "./spoken_dialogue_assistant_turn_1.wav"
|
||||
@@ -127,7 +120,6 @@ Understand and analyze audio content with text queries:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type audio_understanding_sft \
|
||||
--text "Summarize the audio." \
|
||||
@@ -140,7 +132,6 @@ Audio understanding with reasoning chain:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type audio_understanding_sft_with_thinking \
|
||||
--text "Summarize the audio." \
|
||||
@@ -153,7 +144,6 @@ Multi-turn dialogue with audio input and output:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type spoken_dialogue_sft_multiturn \
|
||||
--audio-path "./prompt_speech_zh_m.wav"
|
||||
@@ -167,7 +157,6 @@ Multi-turn dialogue converting speech to text:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type speech2text_dialogue_sft_multiturn
|
||||
```
|
||||
@@ -180,7 +169,6 @@ Multi-turn text-only dialogue:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type text_dialogue_sft_multiturn
|
||||
```
|
||||
|
||||
@@ -29,6 +29,22 @@ python text_to_audio.py \
|
||||
--output stable_audio_output.wav
|
||||
```
|
||||
|
||||
To reduce per-GPU memory for multi-GPU inference, launch with HSDP:
|
||||
|
||||
```bash
|
||||
python text_to_audio.py \
|
||||
--model stabilityai/stable-audio-open-1.0 \
|
||||
--prompt "The sound of a hammer hitting a wooden surface" \
|
||||
--negative-prompt "Low quality" \
|
||||
--seed 42 \
|
||||
--guidance-scale 7.0 \
|
||||
--audio-length 10.0 \
|
||||
--num-inference-steps 100 \
|
||||
--use-hsdp \
|
||||
--hsdp-shard-size 2 \
|
||||
--output stable_audio_output.wav
|
||||
```
|
||||
|
||||
Key arguments:
|
||||
|
||||
- `--prompt`: text description (string).
|
||||
@@ -37,6 +53,9 @@ Key arguments:
|
||||
- `--guidance-scale`: classifier-free guidance scale.
|
||||
- `--audio-length`: audio duration in seconds.
|
||||
- `--num-inference-steps`: diffusion sampling steps.(more steps = higher quality, slower).
|
||||
- `--use-hsdp`: enable HSDP weight sharding for the Stable Audio DiT.
|
||||
- `--hsdp-shard-size`: number of GPUs used for HSDP sharding.
|
||||
- `--hsdp-replicate-size`: number of HSDP replica groups.
|
||||
- `--output`: path to save the generated WAV file.
|
||||
|
||||
## Example materials
|
||||
|
||||
@@ -5,6 +5,8 @@ Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inf
|
||||
|
||||
A unified script for text-to-video generation. Supports multiple models with model-aware defaults.
|
||||
|
||||
For backend selection and SageAttention usage, see the [Diffusion Attention Backends](../../diffusion/attention_backends.md) guide.
|
||||
|
||||
## Supported Models
|
||||
|
||||
| Model | Default Resolution | Default Frames | Default Steps | Guidance | VRAM (BF16) |
|
||||
|
||||
@@ -13,28 +13,24 @@ When `mistral_common` has `SpeechRequest` support, prompt token IDs are built vi
|
||||
```bash
|
||||
# Basic single-prompt with cheerful_female voice preset
|
||||
python3 examples/offline_inference/voxtral_tts/end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
|
||||
--write-audio --voice cheerful_female \
|
||||
--model mistralai/Voxtral-4B-TTS-2603 \
|
||||
--text "That eerie silence after the first storm was just the calm before another round of chaos, wasn't it?"
|
||||
|
||||
# 32 replicate prompts with cheerful_female voice preset
|
||||
python3 examples/offline_inference/voxtral_tts/end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
|
||||
--num-prompts 32 --write-audio --voice cheerful_female \
|
||||
--model mistralai/Voxtral-4B-TTS-2603 \
|
||||
--text "That eerie silence after the first storm was just the calm before another round of chaos, wasn't it?"
|
||||
|
||||
# Streaming with neutral_female voice preset
|
||||
python3 examples/offline_inference/voxtral_tts/end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
|
||||
--streaming --write-audio --voice neutral_female \
|
||||
--model mistralai/Voxtral-4B-TTS-2603 \
|
||||
--text "That eerie silence after the first storm was just the calm before another round of chaos, wasn't it?"
|
||||
|
||||
# 32 prompts, 8 concurrent requests per wave, streaming with neutral_female voice
|
||||
python3 examples/offline_inference/voxtral_tts/end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
|
||||
--num-prompts 32 --concurrency 8 --streaming --write-audio --voice neutral_female \
|
||||
--model mistralai/Voxtral-4B-TTS-2603 \
|
||||
--text "That eerie silence after the first storm was just the calm before another round of chaos, wasn't it?"
|
||||
@@ -42,7 +38,6 @@ python3 examples/offline_inference/voxtral_tts/end2end.py \
|
||||
# Short debug prompt with reference audio
|
||||
# Note: Reference audio capability is not yet released.
|
||||
python3 examples/offline_inference/voxtral_tts/end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
|
||||
--write-audio \
|
||||
--model mistralai/Voxtral-4B-TTS-2603 \
|
||||
--text "This is a test message." \
|
||||
@@ -57,7 +52,7 @@ python3 examples/offline_inference/voxtral_tts/end2end.py \
|
||||
| `--text TEXT` | Text to synthesize (default: `"This is a test message."`) |
|
||||
| `--audio-path PATH` | Path to reference audio file for voice cloning |
|
||||
| `--output-dir DIR` | Directory to write output WAV files (default: `output_audio`) |
|
||||
| `--stage-configs-path PATH` | Path to stage configs YAML (currently it must be set for VoxtralTTS) |
|
||||
| `--deploy-config PATH` | Override the deploy config path. If unset, auto-loads `vllm_omni/deploy/voxtral_tts.yaml` from the HF `model_type`. |
|
||||
| `--num-prompts N` | Number of replicate prompts to run for measuring performance (default: 1) |
|
||||
| `--streaming` | Use streaming generation via `AsyncOmni` (default: blocking `Omni`) |
|
||||
| `--concurrency N` | Max concurrent requests per wave (must be used with `--streaming`, must evenly divide `--num-prompts`) |
|
||||
|
||||
@@ -2,147 +2,112 @@
|
||||
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/bagel>.
|
||||
|
||||
|
||||
## 🛠️ Installation
|
||||
## Installation
|
||||
|
||||
Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/README.md)
|
||||
|
||||
## Run examples (BAGEL-7B-MoT)
|
||||
## Architecture
|
||||
|
||||
**Note**: These examples work with the default configuration on an **NVIDIA A100 (80GB)**. We also tested on dual **NVIDIA RTX 5000 Ada (32GB each)**. For dual-GPU setups, please modify the stage configuration to distribute the model across devices.
|
||||
BAGEL-7B-MoT is a Mixture-of-Transformers (MoT) model supporting both image generation and understanding. It offers two deployment topologies:
|
||||
|
||||
### Launch the Server
|
||||
| Topology | Stages | Description |
|
||||
| :------- | :----- | :---------- |
|
||||
| **Two-stage** (default) | Stage 0 (Thinker, AR) + Stage 1 (DiT, Diffusion) | Thinker handles text/understanding via vLLM AR engine; DiT handles image generation. KV cache is transferred between stages. |
|
||||
| **Single-stage** | Stage 0 (DiT, Diffusion) only | The DiT stage contains a full LLM, ViT, VAE, and tokenizer internally. All modalities are handled within a single diffusion process. |
|
||||
|
||||
Both topologies support all four modalities: `text2img`, `img2img`, `img2text`, `text2text`.
|
||||
|
||||
> **Note**: These examples work with the default configuration on an **NVIDIA A100 (80GB)**. We also tested on dual **NVIDIA RTX 5000 Ada (32GB each)**. For dual-GPU setups, modify the deploy YAML to distribute stages across devices.
|
||||
|
||||
## Launch the Server
|
||||
|
||||
### Two-Stage (Default)
|
||||
|
||||
The default pipeline is auto-detected from the model. No extra flags needed:
|
||||
|
||||
```bash
|
||||
# Use default configuration
|
||||
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091
|
||||
```
|
||||
|
||||
Or use the convenience script:
|
||||
|
||||
```bash
|
||||
cd /workspace/vllm-omni/examples/online_serving/bagel
|
||||
cd examples/online_serving/bagel
|
||||
bash run_server.sh
|
||||
|
||||
# Launch a single stage per terminal
|
||||
bash run_server_stage_cli.sh --stage 0
|
||||
bash run_server_stage_cli.sh --stage 1
|
||||
```
|
||||
|
||||
```bash
|
||||
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
|
||||
```
|
||||
|
||||
#### 🚀 Tensor Parallelism (TP)
|
||||
|
||||
For larger models or multi-GPU environments, you can enable Tensor Parallelism (TP) for the server.
|
||||
|
||||
1. **Modify Stage Config**: Create or modify a stage configuration yaml (e.g., [`bagel.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/model_executor/stage_configs/bagel.yaml)). Set `tensor_parallel_size` to `2` (or more) and update `devices` to include multiple GPU IDs (e.g., `"0,1"`).
|
||||
|
||||
```yaml
|
||||
engine_args:
|
||||
tensor_parallel_size: 2
|
||||
...
|
||||
runtime:
|
||||
devices: "0,1"
|
||||
```
|
||||
|
||||
2. **Launch Server**:
|
||||
```bash
|
||||
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --stage-configs-path /path/to/your/custom_bagel.yaml
|
||||
```
|
||||
|
||||
#### Using Mooncake Connector
|
||||
|
||||
By default, BAGEL uses `SharedMemoryConnector` for inter-stage communication. You can use the [Mooncake](https://github.com/kvcache-ai/Mooncake) connector to transfer KV cache between stages, which also enables multi-node deployment.
|
||||
|
||||
**1. Install Mooncake**
|
||||
|
||||
```bash
|
||||
# For CUDA-enabled systems (recommended)
|
||||
pip install mooncake-transfer-engine
|
||||
|
||||
# For non-CUDA systems
|
||||
pip install mooncake-transfer-engine-non-cuda
|
||||
```
|
||||
|
||||
**2. Start Mooncake Master** on the primary node:
|
||||
|
||||
```bash
|
||||
# Optional: enable disk-backed storage by creating a directory and passing --root_fs_dir.
|
||||
# Without it, Mooncake runs in memory-only mode, which is sufficient for KV cache transfer.
|
||||
mkdir -p ./mc_storage
|
||||
|
||||
mooncake_master \
|
||||
--rpc_port=50051 \
|
||||
--enable_http_metadata_server=true \
|
||||
--http_metadata_server_host=0.0.0.0 \
|
||||
--http_metadata_server_port=8080 \
|
||||
--metrics_port=9003 \
|
||||
--root_fs_dir=./mc_storage/ \
|
||||
--cluster_id=mc-local-1 &
|
||||
```
|
||||
|
||||
**3. Launch the server** with the Mooncake stage config:
|
||||
To use a custom deploy YAML (note: `--stage-configs-path` is deprecated in favor of `--deploy-config`):
|
||||
|
||||
```bash
|
||||
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
|
||||
--deploy-config /path/to/deploy_config.yaml
|
||||
```
|
||||
|
||||
> **Note**: Before launching, edit [`bagel_multiconnector.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml) and replace the `metadata_server` and `master` addresses with your Mooncake master node's actual IP. For single-node testing, `127.0.0.1` works.
|
||||
See [`bagel.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/deploy/bagel.yaml) for the default two-stage deploy configuration.
|
||||
|
||||
The client-side usage is identical to the default setup -- the Mooncake connector is transparent to the API. See the requests section below.
|
||||
### Single-Stage
|
||||
|
||||
For more details on the Mooncake connector configuration, see the [Mooncake Store Connector documentation](https://github.com/vllm-project/vllm-omni/tree/main/docs/design/feature/omni_connectors/mooncake_store_connector.md).
|
||||
|
||||
#### Multi-Node Deployment
|
||||
|
||||
You can deploy each stage on a **separate node** for better resource utilization. In this example, the orchestrator (Stage 0 / Thinker) and Stage 1 (DiT) run on different machines, connected via Mooncake.
|
||||
|
||||
Replace `<ORCHESTRATOR_IP>` below with the actual IP address of your orchestrator node (e.g., `10.244.227.244`).
|
||||
|
||||
> [!WARNING]
|
||||
> **Before launching**, edit [`bagel_multiconnector.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml) and replace the `metadata_server` and `master` addresses with your Mooncake master node's actual IP. Mismatched addresses will cause silent connection failures.
|
||||
|
||||
**1. Start Mooncake Master** (on the orchestrator node):
|
||||
The DiT stage contains a full LLM, ViT, VAE, and tokenizer, so it can handle all modalities (text2img, img2img, img2text, text2text, think) without a separate Thinker stage:
|
||||
|
||||
```bash
|
||||
mooncake_master \
|
||||
--rpc_port=50051 \
|
||||
--enable_http_metadata_server=true \
|
||||
--http_metadata_server_host=<ORCHESTRATOR_IP> \
|
||||
--http_metadata_server_port=8080 \
|
||||
--metrics_port=9003
|
||||
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \
|
||||
--deploy-config vllm_omni/deploy/bagel_single_stage.yaml
|
||||
```
|
||||
|
||||
**2. Launch Stage 0 (Thinker / Orchestrator)** on the orchestrator node:
|
||||
See [`bagel_single_stage.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/deploy/bagel_single_stage.yaml) for configuration. The `pipeline: bagel_single_stage` field selects the single-stage topology from the pipeline registry.
|
||||
|
||||
### Tensor Parallelism (TP)
|
||||
|
||||
For larger models or multi-GPU environments, enable TP via CLI:
|
||||
|
||||
```bash
|
||||
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --tensor-parallel-size 2
|
||||
```
|
||||
|
||||
Or set `tensor_parallel_size` per stage in a custom deploy YAML.
|
||||
|
||||
### Multi-Node Deployment
|
||||
|
||||
Deploy each stage on a **separate node** for better resource utilization. Replace `<ORCHESTRATOR_IP>` with the actual IP address of your orchestrator node.
|
||||
|
||||
**1. Launch Stage 0 (Thinker / Orchestrator)** on the orchestrator node:
|
||||
|
||||
```bash
|
||||
# API server port for client requests: 8000
|
||||
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni \
|
||||
--port 8000 \ # API server port for client requests
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml \
|
||||
--port 8000 \
|
||||
--stage-id 0 \
|
||||
-oma <ORCHESTRATOR_IP> \
|
||||
-omp 8091
|
||||
--omni-master-address <ORCHESTRATOR_IP> \
|
||||
--omni-master-port 8091
|
||||
```
|
||||
|
||||
**3. Launch Stage 1 (DiT)** on the remote node in headless mode:
|
||||
**2. Launch Stage 1 (DiT)** on the remote node in headless mode:
|
||||
|
||||
```bash
|
||||
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml \
|
||||
--stage-id 1 \
|
||||
--headless \
|
||||
-oma <ORCHESTRATOR_IP> \
|
||||
-omp 8091
|
||||
--omni-master-address <ORCHESTRATOR_IP> \
|
||||
--omni-master-port 8091
|
||||
```
|
||||
|
||||
**Mooncake Master arguments:**
|
||||
Or use the convenience script:
|
||||
|
||||
| Argument | Description |
|
||||
| :------- | :---------- |
|
||||
| `--rpc_port` | Mooncake RPC port for control-plane coordination between stages |
|
||||
| `--enable_http_metadata_server` | Enable the HTTP metadata server for service discovery |
|
||||
| `--http_metadata_server_host` | IP address to bind the metadata server (use the orchestrator node's IP) |
|
||||
| `--http_metadata_server_port` | Port for the HTTP metadata server |
|
||||
| `--metrics_port` | Port for Prometheus-compatible metrics endpoint |
|
||||
```bash
|
||||
# Terminal 1: Stage 0
|
||||
bash run_server_stage_cli.sh --stage 0
|
||||
|
||||
# Terminal 2: Stage 1
|
||||
bash run_server_stage_cli.sh --stage 1
|
||||
|
||||
# With extra args
|
||||
bash run_server_stage_cli.sh --stage 0 -- --tensor-parallel-size 2
|
||||
bash run_server_stage_cli.sh --stage 1 -- --gpu-memory-utilization 0.9
|
||||
```
|
||||
|
||||
**vllm serve arguments:**
|
||||
|
||||
@@ -150,85 +115,31 @@ vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni \
|
||||
| :------- | :---------- |
|
||||
| `--stage-id` | Which stage this process runs (0 = Thinker, 1 = DiT) |
|
||||
| `--headless` | Run without the API server (worker-only mode) |
|
||||
| `-oma` | Orchestrator master address |
|
||||
| `-omp` | Orchestrator master port for Stage 1 to connect to Stage 0 for task coordination |
|
||||
| `-oma` / `--omni-master-address` | Orchestrator master address |
|
||||
| `-omp` / `--omni-master-port` | Orchestrator master port |
|
||||
|
||||
> [!IMPORTANT]
|
||||
> **Startup Order**: Stage 0 (orchestrator) must be launched **before** Stage 1 (headless).
|
||||
> Stage 0 will appear to hang on startup until Stage 1 (worker) connects — this is expected behavior.
|
||||
|
||||
**Network Requirements**
|
||||
### Inter-Stage Connectors
|
||||
|
||||
All nodes must have network connectivity to each other. Ensure the following ports are open **between all participating nodes**:
|
||||
When deploying stages across nodes, configure the connector type in the deploy YAML:
|
||||
|
||||
| Port | Protocol | Service | Direction |
|
||||
| :--- | :------- | :------ | :-------- |
|
||||
| 50051 | TCP | Mooncake Master RPC | Worker → Orchestrator |
|
||||
| 8080 | TCP | Mooncake HTTP Metadata Server | Worker → Orchestrator |
|
||||
| 8091 | TCP | Orchestrator Master (`-omp`) | Worker → Orchestrator |
|
||||
| 8000 | TCP | API Server (`--port`) | Client → Orchestrator |
|
||||
| 9003 | TCP | Metrics (optional) | Monitoring → Orchestrator |
|
||||
- **SharedMemoryConnector** (default): Used for single-node deployments. No explicit configuration needed.
|
||||
- **MooncakeTransferEngineConnector**: For multi-node setups with RDMA hardware. Defined in [`bagel.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/deploy/bagel.yaml) under `connectors.rdma_connector`.
|
||||
|
||||
> **Tip**: If nodes are behind a firewall or in different VPCs/security groups, make sure the above ports are allowed in ingress/egress rules. All nodes should be reachable via their IP addresses (no NAT). Using nodes on the same subnet or VPC is recommended to minimize latency for Mooncake KV cache transfers.
|
||||
To use Mooncake, create a custom deploy YAML that binds `output_connectors` / `input_connectors` on each stage to the `rdma_connector` defined in the `connectors` section.
|
||||
|
||||
### Send Multi-modal Request
|
||||
|
||||
Get into the bagel folder:
|
||||
## Send Requests
|
||||
|
||||
```bash
|
||||
cd examples/online_serving/bagel
|
||||
```
|
||||
|
||||
Send request via Python
|
||||
|
||||
```bash
|
||||
python openai_chat_client.py --prompt "A cute cat" --modality text2img
|
||||
```
|
||||
|
||||
The Python client supports the following command-line arguments:
|
||||
|
||||
- `--prompt` (or `-p`): Text prompt for generation (default: `A cute cat`)
|
||||
- `--output` (or `-o`): Output file path for image results (default: `bagel_output.png`)
|
||||
- `--server` (or `-s`): Server URL (default: `http://localhost:8091`)
|
||||
- `--image-url` (or `-i`): Input image URL or local file path (for img2img/img2text modes)
|
||||
- `--modality` (or `-m`): Task modality (default: `text2img`). Options: `text2img`, `img2img`, `img2text`, `text2text`
|
||||
- `--height`: Image height in pixels (default: 512)
|
||||
- `--width`: Image width in pixels (default: 512)
|
||||
- `--steps`: Number of inference steps (default: 25)
|
||||
- `--seed`: Random seed (default: 42)
|
||||
- `--negative`: Negative prompt for image generation
|
||||
|
||||
Example with custom parameters:
|
||||
|
||||
```bash
|
||||
python openai_chat_client.py \
|
||||
--prompt "A futuristic city" \
|
||||
--modality text2img \
|
||||
--height 768 \
|
||||
--width 768 \
|
||||
--steps 50 \
|
||||
--seed 42 \
|
||||
--negative "blurry, low quality"
|
||||
```
|
||||
|
||||
## Modality Control
|
||||
|
||||
BAGEL-7B-MoT supports **multiple modality modes** for different use cases.
|
||||
|
||||
The default yaml configuration deploys Thinker and DiT on the same GPU. You can use the default configuration file: [`bagel.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/model_executor/stage_configs/bagel.yaml)
|
||||
|
||||
| Modality | Input | Output | Description |
|
||||
| ----------- | ------------ | ------ | -------------------------------------- |
|
||||
| `text2img` | Text | Image | Generate images from text prompts |
|
||||
| `img2img` | Image + Text | Image | Transform images using text guidance |
|
||||
| `img2text` | Image + Text | Text | Generate text descriptions from images |
|
||||
| `text2text` | Text | Text | Pure text generation |
|
||||
|
||||
### Text to Image (text2img)
|
||||
|
||||
Generate images from text prompts:
|
||||
|
||||
**Using Python client**
|
||||
**Python client:**
|
||||
|
||||
```bash
|
||||
python openai_chat_client.py \
|
||||
@@ -238,7 +149,7 @@ python openai_chat_client.py \
|
||||
--steps 50
|
||||
```
|
||||
|
||||
**Using curl**
|
||||
**curl:**
|
||||
|
||||
```bash
|
||||
curl http://localhost:8091/v1/chat/completions \
|
||||
@@ -253,12 +164,9 @@ curl http://localhost:8091/v1/chat/completions \
|
||||
}'
|
||||
```
|
||||
|
||||
|
||||
### Image to Image (img2img)
|
||||
|
||||
Transform images based on text prompts:
|
||||
|
||||
**Using Python client**
|
||||
**Python client:**
|
||||
|
||||
```bash
|
||||
python openai_chat_client.py \
|
||||
@@ -268,7 +176,7 @@ python openai_chat_client.py \
|
||||
--output transformed.png
|
||||
```
|
||||
|
||||
**Using curl**
|
||||
**curl:**
|
||||
|
||||
```bash
|
||||
IMAGE_BASE64=$(base64 -w 0 cat.jpg)
|
||||
@@ -293,14 +201,11 @@ EOF
|
||||
curl http://localhost:8091/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d @payload.json
|
||||
|
||||
```
|
||||
|
||||
### Image to Text (img2text)
|
||||
|
||||
Generate text descriptions from images:
|
||||
|
||||
**Using Python client**
|
||||
**Python client:**
|
||||
|
||||
```bash
|
||||
python openai_chat_client.py \
|
||||
@@ -309,7 +214,7 @@ python openai_chat_client.py \
|
||||
--image-url /path/to/image.jpg
|
||||
```
|
||||
|
||||
**Using curl**
|
||||
**curl:**
|
||||
|
||||
```bash
|
||||
IMAGE_BASE64=$(base64 -w 0 cat.jpg)
|
||||
@@ -334,9 +239,7 @@ curl http://localhost:8091/v1/chat/completions \
|
||||
|
||||
### Text to Text (text2text)
|
||||
|
||||
Pure text generation:
|
||||
|
||||
**Using Python client**
|
||||
**Python client:**
|
||||
|
||||
```bash
|
||||
python openai_chat_client.py \
|
||||
@@ -344,26 +247,81 @@ python openai_chat_client.py \
|
||||
--modality text2text
|
||||
```
|
||||
|
||||
**Using curl**
|
||||
**curl:**
|
||||
|
||||
```bash
|
||||
curl http://localhost:8091/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"messages": [{"role": "user", "content": [{"type": "text", "text": "<|im_start|>user\nWhat is the capital of France?<|im_end|>\n<|im_start|>assistant\n"}]}]
|
||||
"messages": [{"role": "user", "content": [{"type": "text", "text": "<|im_start|>user\nWhat is the capital of France?<|im_end|>\n<|im_start|>assistant\n"}]}],
|
||||
"modalities": ["text"]
|
||||
}'
|
||||
```
|
||||
|
||||
### Python Client Arguments
|
||||
|
||||
| Argument | Default | Description |
|
||||
| :------- | :------ | :---------- |
|
||||
| `--prompt` / `-p` | `A cute cat` | Text prompt |
|
||||
| `--output` / `-o` | `bagel_output.png` | Output file path |
|
||||
| `--server` / `-s` | `http://localhost:8091` | Server URL |
|
||||
| `--image-url` / `-i` | `None` | Input image URL or local path (img2img/img2text) |
|
||||
| `--modality` / `-m` | `text2img` | `text2img`, `img2img`, `img2text`, `text2text` |
|
||||
| `--height` | `512` | Image height in pixels |
|
||||
| `--width` | `512` | Image width in pixels |
|
||||
| `--steps` | `25` | Number of inference steps |
|
||||
| `--seed` | `42` | Random seed |
|
||||
| `--negative` | `None` | Negative prompt for CFG |
|
||||
|
||||
Example with custom parameters:
|
||||
|
||||
```bash
|
||||
python openai_chat_client.py \
|
||||
--prompt "A futuristic city" \
|
||||
--modality text2img \
|
||||
--height 768 \
|
||||
--width 768 \
|
||||
--steps 50 \
|
||||
--seed 42 \
|
||||
--negative "blurry, low quality"
|
||||
```
|
||||
|
||||
## Configuration Reference
|
||||
|
||||
### Deploy YAML Files
|
||||
|
||||
| File | Description |
|
||||
| :--- | :---------- |
|
||||
| [`bagel.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/deploy/bagel.yaml) | Two-stage default (Thinker + DiT on GPU 0) |
|
||||
| [`bagel_single_stage.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/deploy/bagel_single_stage.yaml) | Single-stage (DiT only) |
|
||||
|
||||
### Key Deploy YAML Fields
|
||||
|
||||
| Field | Scope | Description |
|
||||
| :---- | :---- | :---------- |
|
||||
| `pipeline` | top-level | Override auto-detected pipeline (e.g. `bagel_single_stage`) |
|
||||
| `stages[].stage_id` | per-stage | Stage identifier (0, 1, ...) |
|
||||
| `stages[].devices` | per-stage | GPU device IDs (e.g. `"0"`, `"0,1"`) |
|
||||
| `stages[].max_num_seqs` | per-stage | Maximum concurrent sequences |
|
||||
| `stages[].gpu_memory_utilization` | per-stage | Fraction of GPU memory to use |
|
||||
| `stages[].enforce_eager` | per-stage | Disable CUDA graphs |
|
||||
| `stages[].tensor_parallel_size` | per-stage | TP degree for this stage |
|
||||
| `connectors` | top-level | Define available connector instances (SHM, Mooncake) |
|
||||
| `platforms` | top-level | Platform-specific overrides (e.g. `xpu`) |
|
||||
|
||||
## FAQ
|
||||
|
||||
- If you don’t know how much VRAM is needed for the model or encounter the OOM error, you can try to decrease the max_model_len.
|
||||
- If you encounter OOM errors, try decreasing `max_model_len` or `gpu_memory_utilization` in the deploy YAML.
|
||||
|
||||
| Stage | VRAM |
|
||||
| :------------------ | :--------------------------- |
|
||||
| Stage-0 (Thinker) | **15.04 GiB** **+ KV Cache** |
|
||||
| Stage-1 (DiT) | **26.50 GiB** |
|
||||
| Total | **~42 GiB + KV Cache** |
|
||||
**Two-stage VRAM usage:**
|
||||
|
||||
| Stage | VRAM |
|
||||
| :---- | :--- |
|
||||
| Stage 0 (Thinker) | **15.04 GiB + KV Cache** |
|
||||
| Stage 1 (DiT) | **26.50 GiB** |
|
||||
| Total | **~42 GiB + KV Cache** |
|
||||
|
||||
**Single-stage VRAM usage:** The DiT loads the full model (~42 GiB) in one process.
|
||||
|
||||
## Example materials
|
||||
|
||||
|
||||
@@ -0,0 +1,93 @@
|
||||
# Diffusers Backend Adapter Example
|
||||
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/diffusers_pipeline_adapter>.
|
||||
|
||||
|
||||
This example demonstrates how to serve any 🤗 Diffusers pipeline through vLLM-Omni
|
||||
using the `diffusers` load format.
|
||||
|
||||
## Supported Models
|
||||
|
||||
Any model loadable via `DiffusionPipeline.from_pretrained()` should be supported, including text-to-image, image-to-image, text-to-video, image-to-video, and text-to-audio.
|
||||
|
||||
## Limitations
|
||||
|
||||
The diffusers backend is a black-box adapter. The following features are NOT yet supported.
|
||||
It is not guaranteed whether they will be supported in the future.
|
||||
|
||||
- CFG parallel execution
|
||||
- Sequence parallel execution
|
||||
- TeaCache / Cache-DiT acceleration
|
||||
- Step-wise execution (continuous batching)
|
||||
|
||||
For these features, it is recommended to use natively supported pipelines instead.
|
||||
|
||||
## Usage
|
||||
|
||||
### Option 1: CLI arguments
|
||||
|
||||
```bash
|
||||
vllm serve "stable-diffusion-v1-5/stable-diffusion-v1-5" \
|
||||
--omni \
|
||||
--diffusion-load-format diffusers \
|
||||
--diffusers-load-kwargs '{"use_safetensors": true}' \
|
||||
--diffusers-call-kwargs '{"num_inference_steps": 30, "guidance_scale": 7.5}'
|
||||
```
|
||||
|
||||
`--diffusers-load-kwargs` and `--diffusers-call-kwargs` are only valid together with `--diffusion-load-format diffusers`.
|
||||
|
||||
### Option 2: Stage config YAML
|
||||
|
||||
```bash
|
||||
vllm serve stable-diffusion-v1-5/stable-diffusion-v1-5 --stage-configs-path examples/online_serving/diffusers_pipeline_adapter/stage_config.yaml --omni
|
||||
```
|
||||
|
||||
The particular fields of interest are `model`, `diffusion_load_format`, `diffusers_load_kwargs`, and `diffusers_call_kwargs` under `engine_args`. They are the same as the CLI arguments.
|
||||
|
||||
## Send a Request
|
||||
|
||||
```bash
|
||||
curl http://localhost:8000/v1/images/generations \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "stable-diffusion-v1-5/stable-diffusion-v1-5",
|
||||
"prompt": "a photo of an astronaut riding a horse on mars",
|
||||
"n": 1,
|
||||
"size": "512x512"
|
||||
}'
|
||||
```
|
||||
|
||||
Or refer to other documentation pages on how to request a particular input/output modality, such as `examples/online_serving/text_to_image/openai_chat_client.py`.
|
||||
|
||||
## Configuration Reference
|
||||
|
||||
For the diffusers adapter, set options under **`engine_args`**:
|
||||
|
||||
### `diffusion_load_format: "diffusers"`
|
||||
|
||||
This field selects the Hugging Face diffusers adapter path (see `DiffusersPipelineLoader`).
|
||||
|
||||
### `diffusers_load_kwargs`
|
||||
|
||||
Passed to `DiffusionPipeline.from_pretrained()`.
|
||||
|
||||
This is suitable for model-specific configurations not available through the vLLM-Omni interface (such as `Omni.__init__()`, `vllm serve` CLI arguments, and stage config YAML fields outside `diffusers_load_kwargs`).
|
||||
|
||||
When a parameter is available in the vLLM-Omni interface, it will be adapted here.
|
||||
But if that parameter is simultaneously set in both the vLLM-Omni interface and `diffusers_load_kwargs`, the **latter** will take precedence.
|
||||
|
||||
### `diffusers_call_kwargs`
|
||||
|
||||
Passed to `pipeline.__call__()`.
|
||||
|
||||
This is suitable for sampling parameters not available through the vLLM-Omni interface (such as `Omni.generate()` and online serving payloads).
|
||||
|
||||
When a parameter is available in the vLLM-Omni interface, it will be adapted here.
|
||||
But if that parameter is simultaneously set in both the vLLM-Omni interface and `diffusers_call_kwargs`, the **former** will take precedence (because it is set at request time).
|
||||
|
||||
## Example materials
|
||||
|
||||
??? abstract "stage_config.yaml"
|
||||
``````yaml
|
||||
--8<-- "examples/online_serving/diffusers_pipeline_adapter/stage_config.yaml"
|
||||
``````
|
||||
@@ -41,15 +41,11 @@ Features:
|
||||
## Launch the Server
|
||||
|
||||
```bash
|
||||
vllm-omni serve fishaudio/s2-pro \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/fish_speech_s2_pro.yaml \
|
||||
--omni \
|
||||
--port 8091 \
|
||||
--trust-remote-code \
|
||||
--enforce-eager \
|
||||
--gpu-memory-utilization 0.9
|
||||
vllm serve fishaudio/s2-pro --omni --port 8091
|
||||
```
|
||||
|
||||
The deploy config is auto-loaded from `vllm_omni/deploy/fish_qwen3_omni.yaml`.
|
||||
|
||||
Or use the convenience script:
|
||||
|
||||
```bash
|
||||
|
||||
@@ -104,9 +104,9 @@ guides.
|
||||
|
||||
When using `/v1/chat/completions`, pass these inside `extra_body` in the curl
|
||||
JSON, or via the `extra_body` keyword argument in the OpenAI Python SDK (see the
|
||||
[Diffusion Chat API guide](../../../../serving/diffusion_chat_api.md)).
|
||||
When using the dedicated [`/v1/images/generations`](../../../../serving/image_generation_api.md)
|
||||
or [`/v1/images/edits`](../../../../serving/image_edit_api.md) endpoints, pass
|
||||
[Diffusion Chat API guide](../../../serving/diffusion_chat_api.md)).
|
||||
When using the dedicated [`/v1/images/generations`](../../../serving/image_generation_api.md)
|
||||
or [`/v1/images/edits`](../../../serving/image_edit_api.md) endpoints, pass
|
||||
the supported generation controls as top-level fields directly. For image
|
||||
dimensions and count, use `size` and `n` rather than `height` or `width`.
|
||||
|
||||
|
||||
@@ -13,10 +13,10 @@ Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/
|
||||
```bash
|
||||
export MIMO_AUDIO_TOKENIZER_PATH="XiaomiMiMo/MiMo-Audio-Tokenizer"
|
||||
|
||||
vllm-omni serve XiaomiMiMo/MiMo-Audio-7B-Instruct --omni \
|
||||
--served-model-name "MiMo-Audio-7B-Instruct" \
|
||||
--port 18091 --stage-configs-path ./vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--chat-template ./examples/online_serving/mimo_audio/chat_template.jinja
|
||||
vllm serve XiaomiMiMo/MiMo-Audio-7B-Instruct --omni \
|
||||
--served-model-name "MiMo-Audio-7B-Instruct" \
|
||||
--port 18091 \
|
||||
--chat-template ./examples/online_serving/mimo_audio/chat_template.jinja
|
||||
```
|
||||
> ⚠️ **Important**
|
||||
> **MiMo-Audio is not compatible with the default chat template.**
|
||||
|
||||
@@ -15,15 +15,72 @@ Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
|
||||
```
|
||||
|
||||
If you want to open async chunking for qwen3-omni, launch the server with command below
|
||||
The default deployment configuration situated at `vllm_omni/deploy/qwen3_omni_moe.yaml` is resolved and loaded
|
||||
automatically via the model registry, obviating the necessity for the `--deploy-config` flag in standard deployment topologies.
|
||||
Asynchronous chunk streaming is **enabled by default** within the bundled configuration.
|
||||
|
||||
To explicitly utilize a custom deployment YAML, specify the configuration path:
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --deploy-config /vllm_omni/deploy/qwen3_omni_moe.yaml
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
|
||||
--deploy-config /path/to/deploy_config_file
|
||||
```
|
||||
|
||||
If you have custom stage configs file, launch the server with command below
|
||||
### Launch individual stages (stage-based CLI)
|
||||
|
||||
Adopt the stage-based CLI architecture to independently instantiate execution processes per functional stage.
|
||||
|
||||
**1. Stage 0 (Thinker + API server)**
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --deploy-config /path/to/deploy_config_file
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
--port 8091 \
|
||||
--stage-id 0 \
|
||||
--omni-master-address 127.0.0.1 \
|
||||
--omni-master-port 26000
|
||||
```
|
||||
|
||||
**2. Stage 1 (Talker)**
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
--stage-id 1 \
|
||||
--headless \
|
||||
--omni-master-address 127.0.0.1 \
|
||||
--omni-master-port 26000
|
||||
```
|
||||
|
||||
**3. Stage 2 (Code2Wav)**
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
--stage-id 2 \
|
||||
--headless \
|
||||
--omni-master-address 127.0.0.1 \
|
||||
--omni-master-port 26000
|
||||
```
|
||||
|
||||
Add `--deploy-config /path/to/deploy_config_file` to every command if you want
|
||||
to override the bundled deploy YAML.
|
||||
|
||||
For the regular one-process launch, stage-specific CLI tuning is usually done
|
||||
with `--stage-overrides`, for example:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
|
||||
--stage-overrides '{"1": {"gpu_memory_utilization": 0.5}}'
|
||||
```
|
||||
|
||||
For the stage-based CLI, you usually do **not** need `--stage-overrides` for
|
||||
that kind of change. Since each command launches one stage, just pass the knob
|
||||
directly on that stage command:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
--stage-id 1 \
|
||||
--headless \
|
||||
--gpu-memory-utilization 0.5 \
|
||||
--omni-master-address 127.0.0.1 \
|
||||
--omni-master-port 26000
|
||||
```
|
||||
|
||||
### Send Multi-modal Request
|
||||
|
||||
@@ -1,44 +1,60 @@
|
||||
# BAGEL-7B-MoT
|
||||
|
||||
## Set up
|
||||
## Setup
|
||||
|
||||
Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
|
||||
|
||||
## Run examples
|
||||
## Architecture
|
||||
|
||||
**Note**: These examples work with the default configuration on an **NVIDIA A100 (80GB)**. We also tested on dual **NVIDIA RTX 5000 Ada (32GB each)**. For dual-GPU setups, please modify the stage configuration to distribute the model across devices.
|
||||
BAGEL-7B-MoT is a Mixture-of-Transformers (MoT) model supporting both image generation and understanding. It offers two deployment topologies:
|
||||
|
||||
Get into the bagel folder
|
||||
| Topology | Stages | Description |
|
||||
| :------- | :----- | :---------- |
|
||||
| **Two-stage** (default) | Stage 0 (Thinker, AR) + Stage 1 (DiT, Diffusion) | Thinker handles text/understanding via vLLM AR engine; DiT handles image generation. KV cache is transferred between stages. |
|
||||
| **Single-stage** | Stage 0 (DiT, Diffusion) only | The DiT stage contains a full LLM, ViT, VAE, and tokenizer internally. All modalities are handled within a single diffusion process. |
|
||||
|
||||
Both topologies support all four modalities: `text2img`, `img2img`, `img2text`, `text2text`.
|
||||
|
||||
## Quick Start
|
||||
|
||||
```bash
|
||||
cd examples/offline_inference/bagel
|
||||
|
||||
# Default two-stage mode (auto-detected)
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2img \
|
||||
--prompts "A cute cat"
|
||||
|
||||
# Single-stage mode
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2img \
|
||||
--prompts "A cute cat" \
|
||||
--deploy-config vllm_omni/deploy/bagel_single_stage.yaml
|
||||
```
|
||||
|
||||
### Modality Control
|
||||
> **Note**: These examples work with the default configuration on an **NVIDIA A100 (80GB)**. For dual-GPU setups, modify the deploy YAML to distribute stages across devices.
|
||||
|
||||
BAGEL-7B-MoT supports multiple modality modes. You can control the mode using the `--modality` argument:
|
||||
## Modality Control
|
||||
|
||||
#### Text to Image (text2img)
|
||||
Control the mode using the `--modality` argument:
|
||||
|
||||
- **Pipeline**: Text → Thinker → DiT → VAE Decode → Image
|
||||
- **Stages Used**: Stage 0 (Thinker) + Stage 1 (DiT)
|
||||
- **KV Transfer**: Thinker sends KV cache to DiT for conditioned generation
|
||||
| Modality | Input | Output | Description |
|
||||
| :------- | :---- | :----- | :---------- |
|
||||
| `text2img` | Text | Image | Generate images from text prompts |
|
||||
| `img2img` | Image + Text | Image | Transform images using text guidance |
|
||||
| `img2text` | Image + Text | Text | Generate text descriptions from images |
|
||||
| `text2text` | Text | Text | Pure text generation (language model mode) |
|
||||
|
||||
Generate images from text prompts:
|
||||
### Text to Image (text2img)
|
||||
|
||||
```bash
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2img \
|
||||
--prompts "A cute cat"
|
||||
--prompts "A cute cat" \
|
||||
--steps 50
|
||||
```
|
||||
|
||||
#### Image to Image (img2img)
|
||||
|
||||
- **Pipeline**: Image → VAE Encode → DiT → VAE Decode → New Image
|
||||
- **Stages Used**: Stage 1 (DiT) only
|
||||
- **Special**: Bypasses the Thinker stage, direct image-to-image transformation
|
||||
|
||||
Transform images based on text prompts:
|
||||
### Image to Image (img2img)
|
||||
|
||||
```bash
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
@@ -47,13 +63,7 @@ python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--prompts "Let the woman wear a blue dress"
|
||||
```
|
||||
|
||||
#### Image to Text (img2text)
|
||||
|
||||
- **Pipeline**: Image → ViT + VAE Encode → Thinker → Text Output
|
||||
- **Stages Used**: Stage 0 (Thinker) only
|
||||
- **Special**: Uses both VAE latent encoding AND ViT semantic encoding for comprehensive image understanding
|
||||
|
||||
Generate text descriptions from images:
|
||||
### Image to Text (img2text)
|
||||
|
||||
```bash
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
@@ -62,193 +72,206 @@ python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--prompts "Describe this image in detail"
|
||||
```
|
||||
|
||||
#### Text to Text (text2text)
|
||||
|
||||
- **Pipeline**: Text → Thinker → Text Output
|
||||
- **Stages Used**: Stage 0 (Thinker) only
|
||||
- **Special**: No visual components involved, operates as pure language model
|
||||
|
||||
Pure text generation:
|
||||
### Text to Text (text2text)
|
||||
|
||||
```bash
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2text \
|
||||
--prompts "What is the capital of France?"
|
||||
|
||||
# You can load prompts from a text file (one prompt per line):
|
||||
# Load prompts from a text file (one prompt per line):
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2text \
|
||||
--txt-prompts /path/to/prompts.txt
|
||||
```
|
||||
|
||||
### Inference Steps
|
||||
## Think Mode
|
||||
|
||||
Control the number of inference steps for image generation:
|
||||
Think mode enables the model to generate `<think>...</think>` planning/reasoning tokens before producing the final output. This improves generation quality for complex prompts.
|
||||
|
||||
- **Two-stage**: The Thinker (AR) stage decodes think tokens, then transfers the augmented KV cache to the DiT stage for image generation.
|
||||
- **Single-stage**: The DiT's internal LLM generates think tokens in-place before proceeding to denoise.
|
||||
|
||||
```bash
|
||||
# You can adjust steps to 100 to improve image quality
|
||||
# Think + text2img: plan before generating
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2img \
|
||||
--steps 50 \
|
||||
--prompts "A cute cat"
|
||||
```
|
||||
--prompts "A futuristic city with flying cars" \
|
||||
--think \
|
||||
--max-think-tokens 1000
|
||||
|
||||
### Key arguments
|
||||
|
||||
BAGEL-7B-MoT supports **multiple modality modes** for different use cases.
|
||||
|
||||
The default yaml configuration deploys Thinker and DiT on the same GPU. You can use the default configuration file: [`bagel.yaml`](../../../vllm_omni/model_executor/stage_configs/bagel.yaml)
|
||||
|
||||
#### 📌 Command Line Arguments (end2end.py)
|
||||
|
||||
| Argument | Type | Default | Description |
|
||||
| :--------------------- | :----- | :---------------------------- | :----------------------------------------------------------- |
|
||||
| `--model` | string | `ByteDance-Seed/BAGEL-7B-MoT` | Model path or name |
|
||||
| `--modality` | choice | `text2img` | Modality mode: `text2img`, `img2img`, `img2text`, `text2text` |
|
||||
| `--prompts` | list | `None` | Input text prompts directly |
|
||||
| `--txt-prompts` | string | `None` | Path to txt file with one prompt per line |
|
||||
| `--image-path` | string | `None` | Input image path (for `img2img`/`img2text`) |
|
||||
| `--steps` | int | `50` | Number of inference steps |
|
||||
| `--stage-configs-path` | string | `None` | Custom stage config file path |
|
||||
| `--worker-backend` | choice | `process` | Worker backend: `process` or `ray` |
|
||||
| `--ray-address` | string | `None` | Ray cluster address |
|
||||
| `--enable-stats` | flag | `False` | Enable statistics logging |
|
||||
| `--init-sleep-seconds` | int | `20` | Initialization sleep time |
|
||||
| `--batch-timeout` | int | `5` | Batch timeout |
|
||||
| `--init-timeout` | int | `300` | Initialization timeout |
|
||||
|
||||
------
|
||||
|
||||
#### ⚙️ Stage Configuration Parameters (bagel.yaml)
|
||||
|
||||
**Stage 0 - Thinker (LLM Stage)**
|
||||
|
||||
| Parameter | Value | Description |
|
||||
| :------------------------------- | :------------------------------ | :----------------------- |
|
||||
| `stage_type` | `llm` | Stage type |
|
||||
| `devices` | `"0"` | GPU device ID |
|
||||
| `max_num_seqs` | `1` | Maximum batch size |
|
||||
| `model_stage` | `thinker` | Model stage identifier |
|
||||
| `model_arch` | `BagelForConditionalGeneration` | Model architecture |
|
||||
| `gpu_memory_utilization` | `0.4` | GPU memory utilization |
|
||||
| `tensor_parallel_size` | `1` | Tensor parallel size |
|
||||
| `max_num_batched_tokens` | `32768` | Maximum batched tokens |
|
||||
| `omni_kv_config.need_send_cache` | `true` | Whether to send KV cache |
|
||||
|
||||
------
|
||||
|
||||
**Stage 1 - DiT (Diffusion Stage)**
|
||||
|
||||
| Parameter | Value | Description |
|
||||
| :------------------------------- | :---------- | :-------------------------- |
|
||||
| `stage_type` | `diffusion` | Stage type |
|
||||
| `devices` | `"0"` | GPU device ID |
|
||||
| `max_num_seqs` | `1` | Maximum batch size |
|
||||
| `model_stage` | `dit` | Model stage identifier |
|
||||
| `gpu_memory_utilization` | `0.4` | GPU memory utilization |
|
||||
| `omni_kv_config.need_recv_cache` | `true` | Whether to receive KV cache |
|
||||
| `engine_input_source` | `[0]` | Input source from Stage 0 |
|
||||
|
||||
------
|
||||
|
||||
#### Tensor Parallelism (TP)
|
||||
|
||||
For larger models or multi-GPU environments, you can enable Tensor Parallelism (TP) by modifying the stage configuration (e.g., [`bagel.yaml`](../../../vllm_omni/model_executor/stage_configs/bagel.yaml)).
|
||||
|
||||
1. **Set `tensor_parallel_size`**: Increase this value (e.g., to `2` or `4`).
|
||||
2. **Set `devices`**: Specify the comma-separated GPU IDs to be used for the stage (e.g., `"0,1"`).
|
||||
|
||||
Example configuration for TP=2 on GPUs 0 and 1:
|
||||
```yaml
|
||||
engine_args:
|
||||
tensor_parallel_size: 2
|
||||
...
|
||||
runtime:
|
||||
devices: "0,1"
|
||||
```
|
||||
|
||||
------
|
||||
|
||||
#### 🔗 Runtime Configuration
|
||||
|
||||
| Parameter | Value | Description |
|
||||
| :-------------------- | :------ | :------------------------------- |
|
||||
| `shm_threshold_bytes` | `65536` | Shared memory threshold (64KB) |
|
||||
|
||||
## Using Mooncake Connector
|
||||
|
||||
[Mooncake](https://github.com/kvcache-ai/Mooncake) is a high-performance distributed KV cache transfer engine that enables efficient cross-node data movement via TCP or RDMA, making it ideal for multi-node disaggregated inference.
|
||||
|
||||
By default, BAGEL uses `SharedMemoryConnector` for inter-stage communication. You can switch to the Mooncake connector for better performance on multi-GPU setups and to enable multi-node deployment.
|
||||
|
||||
### Prerequisites
|
||||
|
||||
Install the Mooncake transfer engine:
|
||||
|
||||
```bash
|
||||
# For CUDA-enabled systems (recommended)
|
||||
pip install mooncake-transfer-engine
|
||||
|
||||
# For non-CUDA systems
|
||||
pip install mooncake-transfer-engine-non-cuda
|
||||
```
|
||||
|
||||
### Step 1: Start the Mooncake Master
|
||||
|
||||
On the **primary node**, start the Mooncake master service (run in a separate terminal or background with `&`):
|
||||
|
||||
```bash
|
||||
# Optional: enable disk-backed storage by creating a directory and passing --root_fs_dir.
|
||||
# Without it, Mooncake runs in memory-only mode, which is sufficient for KV cache transfer.
|
||||
mkdir -p ./mc_storage
|
||||
|
||||
mooncake_master \
|
||||
--rpc_port=50051 \
|
||||
--enable_http_metadata_server=true \
|
||||
--http_metadata_server_host=0.0.0.0 \
|
||||
--http_metadata_server_port=8080 \
|
||||
--metrics_port=9003 \
|
||||
--root_fs_dir=./mc_storage/ \
|
||||
--cluster_id=mc-local-1 &
|
||||
```
|
||||
|
||||
### Step 2: Run Offline Inference with Mooncake
|
||||
|
||||
Use the provided Mooncake stage config [`bagel_multiconnector.yaml`](../../../vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml). Before launching, update the `metadata_server` and `master` addresses in the YAML to match your Mooncake master node's IP (use `127.0.0.1` for single-node testing).
|
||||
|
||||
```bash
|
||||
cd examples/offline_inference/bagel
|
||||
|
||||
# Text to Image with Mooncake
|
||||
# Think + img2img: reason about the edit
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2img \
|
||||
--prompts "A cute cat" \
|
||||
--stage-configs-path ../../../vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
|
||||
--modality img2img \
|
||||
--image-path /path/to/image.jpg \
|
||||
--prompts "Make it look like a watercolor painting" \
|
||||
--think
|
||||
|
||||
# Image to Text with Mooncake
|
||||
# Think + img2text: reason before describing
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality img2text \
|
||||
--image-path /path/to/image.jpg \
|
||||
--prompts "Describe this image" \
|
||||
--stage-configs-path ../../../vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
|
||||
--prompts "What is happening in this image?" \
|
||||
--think
|
||||
|
||||
# Text to Text with Mooncake
|
||||
# Think + text2text: chain-of-thought reasoning
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2text \
|
||||
--prompts "What is the capital of France?" \
|
||||
--stage-configs-path ../../../vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
|
||||
--prompts "Solve: 23 * 47" \
|
||||
--think
|
||||
```
|
||||
|
||||
For more details on the Mooncake connector and multi-node setup, see the [Mooncake Store Connector documentation](../../../docs/design/feature/omni_connectors/mooncake_store_connector.md).
|
||||
Think mode parameters:
|
||||
|
||||
------
|
||||
| Argument | Default | Description |
|
||||
| :------- | :------ | :---------- |
|
||||
| `--think` | `False` | Enable thinking mode |
|
||||
| `--max-think-tokens` | `1000` | Maximum tokens for think generation |
|
||||
| `--do-sample` | `False` | Enable sampling (vs. greedy) for text generation |
|
||||
| `--text-temperature` | `0.3` | Temperature for text generation sampling |
|
||||
|
||||
## Classifier-Free Guidance (CFG)
|
||||
|
||||
CFG controls the trade-off between prompt fidelity and diversity. These parameters apply to image generation modalities (`text2img`, `img2img`).
|
||||
|
||||
```bash
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2img \
|
||||
--prompts "A photorealistic portrait" \
|
||||
--cfg-text-scale 6.0 \
|
||||
--cfg-img-scale 2.0 \
|
||||
--negative-prompt "blurry, low quality, distorted" \
|
||||
--cfg-interval 0.4 1.0 \
|
||||
--cfg-renorm-type global \
|
||||
--cfg-renorm-min 0.0
|
||||
```
|
||||
|
||||
| Argument | Default | Description |
|
||||
| :------- | :------ | :---------- |
|
||||
| `--cfg-text-scale` | `4.0` | Text CFG scale (higher = more prompt-adherent) |
|
||||
| `--cfg-img-scale` | `1.5` | Image CFG scale (for img2img) |
|
||||
| `--negative-prompt` | `None` | Negative prompt for CFG conditioning |
|
||||
| `--cfg-interval` | pipeline default | CFG active interval `[start, end]` as fractions of total timesteps |
|
||||
| `--cfg-renorm-type` | `None` | Renormalization type: `global`, `text_channel`, `channel` |
|
||||
| `--cfg-renorm-min` | `None` | Minimum renormalization value |
|
||||
| `--cfg-parallel-size` | `1` | CFG parallel size: `1` = batched (single GPU), `2` = 2-branch parallel, `3` = full 3-GPU parallel |
|
||||
|
||||
## Deployment Topologies
|
||||
|
||||
### Two-Stage (Default)
|
||||
|
||||
The default topology auto-detected from the model. No extra flags needed.
|
||||
|
||||
```bash
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2img \
|
||||
--prompts "A cute cat"
|
||||
```
|
||||
|
||||
The pipeline is defined in [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml). Stage 0 (Thinker) and Stage 1 (DiT) share GPU 0 by default. For dual-GPU setups, customize the deploy YAML and set `devices: "1"` for stage 1.
|
||||
|
||||
### Single-Stage
|
||||
|
||||
Pass the single-stage deploy config via `--deploy-config`:
|
||||
|
||||
```bash
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2img \
|
||||
--prompts "A cute cat" \
|
||||
--deploy-config vllm_omni/deploy/bagel_single_stage.yaml
|
||||
```
|
||||
|
||||
See [`bagel_single_stage.yaml`](../../../vllm_omni/deploy/bagel_single_stage.yaml) for configuration details. The `pipeline: bagel_single_stage` field selects the single-stage topology from the pipeline registry.
|
||||
|
||||
### Tensor Parallelism (TP)
|
||||
|
||||
For larger models or multi-GPU environments, customize the deploy YAML (see [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml)) and set per-stage `tensor_parallel_size` and `devices`:
|
||||
|
||||
```yaml
|
||||
# Example: TP=2 on GPUs 0,1 for the Thinker stage
|
||||
stages:
|
||||
- stage_id: 0
|
||||
tensor_parallel_size: 2
|
||||
devices: "0,1"
|
||||
```
|
||||
|
||||
Then pass the custom deploy YAML:
|
||||
|
||||
```bash
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2img \
|
||||
--prompts "A cute cat" \
|
||||
--deploy-config /path/to/custom_bagel.yaml
|
||||
```
|
||||
|
||||
### FP8 Quantization
|
||||
|
||||
```bash
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2img \
|
||||
--prompts "A cute cat" \
|
||||
--quantization fp8
|
||||
```
|
||||
|
||||
## Command Line Reference
|
||||
|
||||
### Core Arguments
|
||||
|
||||
| Argument | Type | Default | Description |
|
||||
| :------- | :--- | :------ | :---------- |
|
||||
| `--model` | string | `ByteDance-Seed/BAGEL-7B-MoT` | Model path or HuggingFace name |
|
||||
| `--modality` | choice | `text2img` | `text2img`, `img2img`, `img2text`, `text2text` |
|
||||
| `--prompts` | list | `None` | Input text prompts |
|
||||
| `--txt-prompts` | string | `None` | Path to text file with one prompt per line |
|
||||
| `--image-path` | string | `None` | Input image path (required for `img2img`/`img2text`) |
|
||||
| `--output` | string | `.` | Output directory for saved images |
|
||||
| `--steps` | int | `50` | Number of diffusion inference steps |
|
||||
| `--seed` | int | `None` | Random seed for reproducibility |
|
||||
|
||||
### Think Mode Arguments
|
||||
|
||||
| Argument | Type | Default | Description |
|
||||
| :------- | :--- | :------ | :---------- |
|
||||
| `--think` | flag | `False` | Enable `<think>...</think>` planning/reasoning |
|
||||
| `--max-think-tokens` | int | `1000` | Maximum tokens for think generation |
|
||||
| `--do-sample` | flag | `False` | Use sampling instead of greedy decoding |
|
||||
| `--text-temperature` | float | `0.3` | Sampling temperature for text generation |
|
||||
|
||||
### CFG Arguments
|
||||
|
||||
| Argument | Type | Default | Description |
|
||||
| :------- | :--- | :------ | :---------- |
|
||||
| `--cfg-text-scale` | float | `4.0` | Text CFG guidance scale |
|
||||
| `--cfg-img-scale` | float | `1.5` | Image CFG guidance scale |
|
||||
| `--negative-prompt` | string | `None` | Negative prompt for CFG |
|
||||
| `--cfg-parallel-size` | int | `1` | CFG parallel GPU count (1, 2, or 3) |
|
||||
| `--cfg-interval` | float[2] | pipeline default | CFG active window `[start, end]` |
|
||||
| `--cfg-renorm-type` | string | `None` | `global`, `text_channel`, or `channel` |
|
||||
| `--cfg-renorm-min` | float | `None` | Minimum renormalization value |
|
||||
|
||||
### Engine Arguments
|
||||
|
||||
| Argument | Type | Default | Description |
|
||||
| :------- | :--- | :------ | :---------- |
|
||||
| `--deploy-config` | string | `None` | Path to deploy YAML (auto-detected if omitted) |
|
||||
| `--worker-backend` | choice | `process` | `process` or `ray` |
|
||||
| `--ray-address` | string | `None` | Ray cluster address |
|
||||
| `--quantization` | string | `None` | Quantization method (e.g. `fp8`) |
|
||||
| `--log-stats` | flag | `False` | Enable statistics logging |
|
||||
| `--init-timeout` | int | `300` | Initialization timeout (seconds) |
|
||||
| `--batch-timeout` | int | `5` | Batch timeout (seconds) |
|
||||
| `--enable-diffusion-pipeline-profiler` | flag | `False` | Profile diffusion stage durations |
|
||||
|
||||
## FAQ
|
||||
|
||||
- If you don’t know how much VRAM is needed for the model or encounter the OOM error, you can try to decrease the max_model_len.
|
||||
- If you encounter OOM errors, try decreasing `max_model_len` or `gpu_memory_utilization` in the deploy YAML.
|
||||
|
||||
| Stage | VRAM |
|
||||
| :------------------ | :--------------------------- |
|
||||
| Stage-0 (Thinker) | **15.04 GiB** **+ KV Cache** |
|
||||
| Stage-1 (DiT) | **26.50 GiB** |
|
||||
| Total | **~42 GiB + KV Cache** |
|
||||
**Two-stage VRAM usage:**
|
||||
|
||||
| Stage | VRAM |
|
||||
| :---- | :--- |
|
||||
| Stage 0 (Thinker) | **15.04 GiB + KV Cache** |
|
||||
| Stage 1 (DiT) | **26.50 GiB** |
|
||||
| Total | **~42 GiB + KV Cache** |
|
||||
|
||||
**Single-stage VRAM usage:** The DiT loads the full model (~42 GiB) in one process.
|
||||
@@ -53,7 +53,12 @@ def parse_args():
|
||||
parser.add_argument("--shm-threshold-bytes", type=int, default=65536)
|
||||
parser.add_argument("--worker-backend", type=str, default="process", choices=["process", "ray"])
|
||||
parser.add_argument("--ray-address", type=str, default=None)
|
||||
parser.add_argument("--stage-configs-path", type=str, default=None)
|
||||
parser.add_argument(
|
||||
"--deploy-config",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to deploy YAML. If unset, auto-loads vllm_omni/deploy/bagel.yaml based on the HF model_type.",
|
||||
)
|
||||
parser.add_argument("--steps", type=int, default=50, help="Number of inference steps.")
|
||||
|
||||
parser.add_argument("--cfg-text-scale", type=float, default=4.0, help="Text CFG scale (default: 4.0)")
|
||||
@@ -116,6 +121,9 @@ def parse_args():
|
||||
help="Temperature for text generation sampling (default: 0.3).",
|
||||
)
|
||||
|
||||
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||
|
||||
nullify_stage_engine_defaults(parser)
|
||||
args = parser.parse_args()
|
||||
return args
|
||||
|
||||
@@ -146,14 +154,12 @@ def main():
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
omni_kwargs = {}
|
||||
stage_configs_path = args.stage_configs_path
|
||||
is_single_stage = stage_configs_path and "single_stage" in stage_configs_path
|
||||
if args.think and stage_configs_path is None:
|
||||
stage_configs_path = "vllm_omni/model_executor/stage_configs/bagel_think.yaml"
|
||||
print(f"[Info] Think mode enabled, using stage config: {stage_configs_path}")
|
||||
if stage_configs_path:
|
||||
omni_kwargs["stage_configs_path"] = stage_configs_path
|
||||
is_single_stage = "single_stage" in stage_configs_path
|
||||
deploy_config = args.deploy_config
|
||||
if args.think and deploy_config is None:
|
||||
deploy_config = "vllm_omni/deploy/bagel_think.yaml"
|
||||
print(f"[Info] Think mode enabled, using deploy config: {deploy_config}")
|
||||
if deploy_config:
|
||||
omni_kwargs["deploy_config"] = deploy_config
|
||||
|
||||
omni_kwargs.update(
|
||||
{
|
||||
@@ -170,7 +176,7 @@ def main():
|
||||
if args.quantization:
|
||||
omni_kwargs["quantization_config"] = args.quantization
|
||||
|
||||
omni = Omni(model=model_name, **omni_kwargs)
|
||||
omni = Omni.from_cli_args(args, model=model_name, **omni_kwargs)
|
||||
|
||||
formatted_prompts = []
|
||||
for p in prompts:
|
||||
@@ -215,9 +221,11 @@ def main():
|
||||
formatted_prompts.append(prompt_dict)
|
||||
|
||||
params_list = omni.default_sampling_params_list
|
||||
# Bagel exposes 1 sampling param set for single-stage (DiT-only) and
|
||||
# 2 for two-stage (Thinker + DiT). This heuristic may need updating
|
||||
# if future pipelines break that 1:1 mapping.
|
||||
is_single_stage = len(params_list) == 1
|
||||
|
||||
# For single-stage DiT, think/text params go into the diffusion sampling params extra_args.
|
||||
# For 2-stage, diffusion params are at index 1.
|
||||
diffusion_params_idx = 0 if is_single_stage else (1 if len(params_list) > 1 else 0)
|
||||
diffusion_params = params_list[diffusion_params_idx]
|
||||
|
||||
|
||||
@@ -58,7 +58,14 @@ Key components live in `vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py
|
||||
- Stage 0 uses `CosyVoice3LM` and outputs speech tokens + conditioning features.
|
||||
- Stage 1 runs the flow model (DiT-based CFM) and HiFiGAN to synthesize waveform.
|
||||
|
||||
Stage wiring is configured in `vllm_omni/model_executor/stage_configs/cosyvoice3.yaml`.
|
||||
Pipeline topology lives in `vllm_omni/model_executor/models/cosyvoice3/pipeline.py`;
|
||||
runtime tunables (batch size, memory limits, sampling) live in
|
||||
`vllm_omni/deploy/cosyvoice3.yaml`. The deploy config auto-loads by
|
||||
HF `model_type` and defaults to `async_chunk: true` (shared-memory
|
||||
streaming). Pass `--no-async-chunk` on `vllm serve` to switch to the
|
||||
legacy sync path where stage 1 runs `text2flow` over the full
|
||||
speech-token sequence.
|
||||
|
||||
- Stage 0 emits latent speech tokens.
|
||||
- Stage 1 consumes them via `custom_process_input_func` and outputs audio.
|
||||
- Stage 1 consumes them via `sync_process_input_func` (sync mode) or the
|
||||
shared-memory connector (async-chunk mode) and outputs audio.
|
||||
@@ -24,7 +24,13 @@ def run_e2e():
|
||||
required=True,
|
||||
help="Path to CosyVoice3 model directory (e.g., pretrained_models/Fun-CosyVoice3-0.5B/).",
|
||||
)
|
||||
parser.add_argument("--stage-config", type=str, default="vllm_omni/model_executor/stage_configs/cosyvoice3.yaml")
|
||||
parser.add_argument(
|
||||
"--deploy-config",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Override the deploy config path. If unset, auto-loads "
|
||||
"vllm_omni/deploy/cosyvoice3.yaml based on the HF model_type.",
|
||||
)
|
||||
parser.add_argument("--prompt", type=str, default="Hello, this is a test of the CosyVoice system capability.")
|
||||
parser.add_argument(
|
||||
"--prompt-text",
|
||||
@@ -43,19 +49,14 @@ def run_e2e():
|
||||
if not os.path.exists(args.tokenizer):
|
||||
raise FileNotFoundError(f"{args.tokenizer} does not exist!")
|
||||
|
||||
# Ensure stage config exists
|
||||
if not os.path.exists(args.stage_config):
|
||||
raise FileNotFoundError(f"{args.stage_config} does not exist!")
|
||||
if args.deploy_config is not None and not os.path.exists(args.deploy_config):
|
||||
raise FileNotFoundError(f"{args.deploy_config} does not exist!")
|
||||
|
||||
print(f"Initializing cosyvoice E2E with model={args.model}")
|
||||
|
||||
# Initialize Omni
|
||||
# This spins up the engine(s) based on the stage config
|
||||
# We pass trust_remote_code=True same as Qwen examples
|
||||
omni = Omni(
|
||||
model=args.model,
|
||||
stage_configs_path=args.stage_config,
|
||||
trust_remote_code=True,
|
||||
deploy_config=args.deploy_config,
|
||||
tokenizer=args.tokenizer,
|
||||
log_stats=True,
|
||||
)
|
||||
|
||||
@@ -44,9 +44,11 @@ Layered RGBA output:
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from PIL import Image
|
||||
@@ -58,6 +60,16 @@ from vllm_omni.outputs import OmniRequestOutput
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
|
||||
def parse_profiler_config(value: str) -> dict[str, Any]:
|
||||
try:
|
||||
config = json.loads(value)
|
||||
except json.JSONDecodeError as e:
|
||||
raise argparse.ArgumentTypeError(f"--profiler-config must be valid JSON: {e}") from e
|
||||
if not isinstance(config, dict):
|
||||
raise argparse.ArgumentTypeError("--profiler-config must be a JSON object")
|
||||
return config
|
||||
|
||||
|
||||
# ===========================
|
||||
# Argument Parser
|
||||
# ===========================
|
||||
@@ -99,7 +111,16 @@ def parse_args() -> argparse.Namespace:
|
||||
parser.add_argument("--vae-use-slicing", action="store_true")
|
||||
parser.add_argument("--vae-use-tiling", action="store_true")
|
||||
parser.add_argument("--enable-cpu-offload", action="store_true")
|
||||
parser.add_argument(
|
||||
"--profiler-config",
|
||||
type=parse_profiler_config,
|
||||
default=None,
|
||||
help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
|
||||
)
|
||||
|
||||
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||
|
||||
nullify_stage_engine_defaults(parser)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
@@ -158,12 +179,13 @@ async def main():
|
||||
enable_cpu_offload=args.enable_cpu_offload,
|
||||
diffusion_load_format="dummy",
|
||||
custom_pipeline_args={"pipeline_class": "custom_pipeline.CustomPipeline"},
|
||||
profiler_config=args.profiler_config,
|
||||
)
|
||||
|
||||
print(">>> Pipeline loaded successfully")
|
||||
|
||||
# ---- Profiling + Info ----
|
||||
profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
|
||||
profiler_enabled = args.profiler_config is not None
|
||||
print(f"\n{'=' * 60}")
|
||||
print("Generation Configuration")
|
||||
print(f"Model: {args.model}")
|
||||
|
||||
@@ -970,6 +970,9 @@ def parse_args(repo_root: Path) -> argparse.Namespace:
|
||||
parser.add_argument("--vq-model-audio-local-files-only", action=argparse.BooleanOptionalAction, default=None)
|
||||
|
||||
parser.add_argument("--disable-hf-xet", action=argparse.BooleanOptionalAction, default=True)
|
||||
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||
|
||||
nullify_stage_engine_defaults(parser)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
|
||||
@@ -1,145 +0,0 @@
|
||||
# GLM-Image Multistage End-to-End Inference
|
||||
|
||||
This example demonstrates how to run GLM-Image with the vLLM-Omni multistage architecture.
|
||||
|
||||
## Architecture
|
||||
|
||||
GLM-Image uses a 2-stage pipeline:
|
||||
|
||||
```
|
||||
┌─────────────────────────────────────────────────────────────┐
|
||||
│ GLM-Image Pipeline │
|
||||
├─────────────────────────────────────────────────────────────┤
|
||||
│ │
|
||||
│ Stage 0 (AR Model) Stage 1 (Diffusion) │
|
||||
│ ┌─────────────────┐ ┌─────────────────────┐ │
|
||||
│ │ vLLM-optimized │ │ GlmImagePipeline │ │
|
||||
│ │ GlmImageFor │ prior │ ┌───────────────┐ │ │
|
||||
│ │ Conditional │──tokens───►│ │ DiT Denoiser │ │ │
|
||||
│ │ Generation │ │ └───────────────┘ │ │
|
||||
│ │ (9B AR model) │ │ │ │ │
|
||||
│ └─────────────────┘ │ ▼ │ │
|
||||
│ ▲ │ ┌───────────────┐ │ │
|
||||
│ │ │ │ VAE Decode │──┼──► Image
|
||||
│ Text/Image │ └───────────────┘ │ │
|
||||
│ Input └─────────────────────┘ │
|
||||
│ │
|
||||
└─────────────────────────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
## Features
|
||||
|
||||
- **vLLM-optimized AR**: Uses PagedAttention and tensor parallelism for faster prior token generation
|
||||
- **Flexible deployment**: AR and Diffusion stages can run on different GPUs
|
||||
- **Text-to-Image**: Generate images from text descriptions
|
||||
- **Image-to-Image**: Edit existing images with text prompts
|
||||
|
||||
## Usage
|
||||
|
||||
### Text-to-Image
|
||||
|
||||
```bash
|
||||
python end2end.py \
|
||||
--config-path ../../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
|
||||
--prompt "A beautiful sunset over the ocean with sailing boats" \
|
||||
--height 1024 \
|
||||
--width 1024 \
|
||||
--output output_t2i.png
|
||||
```
|
||||
|
||||
### Image-to-Image (Image Editing)
|
||||
|
||||
```bash
|
||||
python end2end.py \
|
||||
--config-path ../../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
|
||||
--prompt "Transform this scene into a winter wonderland" \
|
||||
--image input.png \
|
||||
--output output_i2i.png
|
||||
```
|
||||
|
||||
### With Custom Parameters
|
||||
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model-path /path/to/glm-image \
|
||||
--config-path ../../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
|
||||
--prompt "A photorealistic cat sitting on a window sill" \
|
||||
--height 1024 \
|
||||
--width 1024 \
|
||||
--num-inference-steps 50 \
|
||||
--guidance-scale 1.5 \
|
||||
--seed 42 \
|
||||
--output output.png
|
||||
```
|
||||
|
||||
## Shell Scripts
|
||||
|
||||
### Run Text-to-Image
|
||||
|
||||
```bash
|
||||
./run_t2i.sh
|
||||
```
|
||||
|
||||
### Run Image-to-Image
|
||||
|
||||
```bash
|
||||
./run_i2i.sh --image /path/to/input.png
|
||||
```
|
||||
|
||||
## Stage Configuration
|
||||
|
||||
The stage config (`glm_image.yaml`) defines:
|
||||
|
||||
- **Stage 0 (AR)**: Uses `GPUARWorker` with vLLM engine
|
||||
|
||||
- Model: `GlmImageForConditionalGeneration`
|
||||
- Output: `token_ids` (prior tokens)
|
||||
|
||||
- **Stage 1 (Diffusion)**: Uses diffusion engine
|
||||
- Model: `GlmImagePipeline`
|
||||
- Output: Generated image
|
||||
|
||||
See `vllm_omni/model_executor/stage_configs/glm_image.yaml` for full configuration.
|
||||
|
||||
## Comparison with Single-Stage
|
||||
|
||||
| Aspect | Single-Stage (transformers) | Multistage (vLLM) |
|
||||
| ----------- | --------------------------- | ------------------- |
|
||||
| AR Model | transformers native | vLLM PagedAttention |
|
||||
| Memory | Higher (no KV cache opt) | Lower (optimized) |
|
||||
| Throughput | Lower | Higher |
|
||||
| Flexibility | Single GPU | Multi-GPU support |
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### OOM Error
|
||||
|
||||
Try reducing memory usage:
|
||||
|
||||
```bash
|
||||
# In glm_image.yaml, adjust:
|
||||
gpu_memory_utilization: 0.5 # Reduce from 0.6
|
||||
```
|
||||
|
||||
### Slow Initialization
|
||||
|
||||
The first run loads model weights. Subsequent runs are faster:
|
||||
|
||||
```bash
|
||||
--stage-init-timeout 900 # Increase timeout for slow storage
|
||||
```
|
||||
|
||||
### `Transformers does not recognize this architecture` Error
|
||||
|
||||
Your have to upgrade `transformers` package to `5.3.0` or above:
|
||||
|
||||
```
|
||||
pip install --upgrade transformers
|
||||
```
|
||||
|
||||
## Requirements
|
||||
|
||||
- vLLM-Omni with GLM-Image support
|
||||
- CUDA-capable GPU (recommended: H100/A100 with 80GB)
|
||||
- GLM-Image model weights
|
||||
- `transformers` v5.3.0 or above
|
||||
@@ -1,525 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
End-to-end offline inference example for GLM-Image with multistage architecture.
|
||||
|
||||
This script tests the multistage pipeline where:
|
||||
- Stage 0 (AR): vLLM-optimized GlmImageForConditionalGeneration generates prior_token_ids
|
||||
- Stage 1 (Diffusion): GlmImagePipeline performs DiT denoising + VAE decode
|
||||
|
||||
Usage (text-to-image):
|
||||
python end2end.py \
|
||||
--model-path /path/to/glm-image \
|
||||
--config-path /path/to/glm_image.yaml \
|
||||
--prompt "A beautiful sunset over the ocean" \
|
||||
--output output_t2i.png
|
||||
|
||||
Usage (image-to-image / image edit):
|
||||
python end2end.py \
|
||||
--model-path /path/to/glm-image \
|
||||
--config-path /path/to/glm_image.yaml \
|
||||
--prompt "Make it look like winter" \
|
||||
--image input.png \
|
||||
--output output_i2i.png
|
||||
|
||||
Usage (with custom parameters):
|
||||
python end2end.py \
|
||||
--model-path /path/to/glm-image \
|
||||
--config-path /path/to/glm_image.yaml \
|
||||
--prompt "A cat sitting on a window sill" \
|
||||
--height 1024 \
|
||||
--width 1024 \
|
||||
--num-inference-steps 50 \
|
||||
--guidance-scale 1.5 \
|
||||
--seed 42
|
||||
|
||||
For more options, run:
|
||||
python end2end.py --help
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
from PIL import Image
|
||||
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
# Default stage config path (relative to vllm_omni package)
|
||||
DEFAULT_CONFIG_PATH = "vllm_omni/model_executor/stage_configs/glm_image.yaml"
|
||||
|
||||
SEED = 42
|
||||
|
||||
# GLM-Image special tokens
|
||||
GLM_IMAGE_EOS_TOKEN_ID = 16385 # eos_token_id from generation_config.json
|
||||
GLM_IMAGE_VISION_VOCAB_SIZE = 16512 # top_k should be vision_vocab_size
|
||||
|
||||
|
||||
def compute_max_tokens(height: int, width: int, factor: int = 32, is_i2i: bool = False) -> int:
|
||||
"""
|
||||
Compute max_new_tokens for GLM-Image AR generation.
|
||||
|
||||
GLM-Image generation differs by mode:
|
||||
- text-to-image (t2i): small preview + large target + EOS
|
||||
- image-to-image (i2i): large target + EOS
|
||||
|
||||
Args:
|
||||
height: Target image height in pixels
|
||||
width: Target image width in pixels
|
||||
factor: Downsampling factor (32 for GLM-Image AR output)
|
||||
is_i2i: Whether the request is image-to-image mode
|
||||
|
||||
Returns:
|
||||
Total number of tokens to generate for the specified mode
|
||||
"""
|
||||
# Large image tokens (target resolution)
|
||||
token_h = height // factor
|
||||
token_w = width // factor
|
||||
large_tokens = token_h * token_w
|
||||
|
||||
# Small preview tokens (half resolution in each dimension)
|
||||
import math
|
||||
|
||||
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
|
||||
|
||||
if is_i2i:
|
||||
return large_tokens + 1
|
||||
return small_tokens + large_tokens + 1
|
||||
|
||||
|
||||
def load_image(image_path: str) -> Image.Image:
|
||||
"""Load an image from file path."""
|
||||
if not os.path.exists(image_path):
|
||||
raise FileNotFoundError(f"Image file not found: {image_path}")
|
||||
return Image.open(image_path).convert("RGB")
|
||||
|
||||
|
||||
def save_image(image: Image.Image, output_path: str) -> None:
|
||||
"""Save an image to file path."""
|
||||
output_dir = os.path.dirname(output_path)
|
||||
if output_dir:
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
image.save(output_path)
|
||||
print(f"Image saved to: {output_path}")
|
||||
|
||||
|
||||
def build_prompt_for_t2i(
|
||||
prompt: str,
|
||||
height: int = 1024,
|
||||
width: int = 1024,
|
||||
) -> dict:
|
||||
"""
|
||||
Build prompt dict for text-to-image generation.
|
||||
|
||||
Args:
|
||||
prompt: Text description for image generation
|
||||
height: Target image height
|
||||
width: Target image width
|
||||
|
||||
Returns:
|
||||
Dict containing prompt and generation parameters
|
||||
"""
|
||||
return {
|
||||
"prompt": prompt,
|
||||
"height": height,
|
||||
"width": width,
|
||||
# Pass target dimensions to AR processor for proper grid token generation
|
||||
"mm_processor_kwargs": {
|
||||
"target_h": height,
|
||||
"target_w": width,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def build_prompt_for_i2i(
|
||||
prompt: str,
|
||||
image: Image.Image,
|
||||
height: int | None = None,
|
||||
width: int | None = None,
|
||||
) -> dict:
|
||||
"""
|
||||
Build prompt dict for image-to-image generation.
|
||||
|
||||
Args:
|
||||
prompt: Text description for image editing
|
||||
image: Source image for editing
|
||||
height: Target image height (default: use source image size)
|
||||
width: Target image width (default: use source image size)
|
||||
|
||||
Returns:
|
||||
Dict containing prompt, image, and generation parameters
|
||||
"""
|
||||
# Use source image dimensions if not specified
|
||||
if height is None:
|
||||
height = image.height
|
||||
if width is None:
|
||||
width = image.width
|
||||
|
||||
return {
|
||||
"prompt": prompt,
|
||||
"multi_modal_data": {
|
||||
"image": image,
|
||||
},
|
||||
"height": height,
|
||||
"width": width,
|
||||
# Pass target dimensions to AR processor for proper grid token generation
|
||||
"mm_processor_kwargs": {
|
||||
"target_h": height,
|
||||
"target_w": width,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def main(args: argparse.Namespace) -> None:
|
||||
"""Main entry point for GLM-Image end-to-end inference."""
|
||||
print("=" * 60)
|
||||
print("GLM-Image Multistage End-to-End Inference")
|
||||
print("=" * 60)
|
||||
|
||||
# Validate arguments
|
||||
if not args.prompt:
|
||||
raise ValueError("--prompt is required")
|
||||
|
||||
# Determine config path
|
||||
config_path = args.config_path
|
||||
if config_path is None:
|
||||
# Try to find default config
|
||||
if os.path.exists(DEFAULT_CONFIG_PATH):
|
||||
config_path = DEFAULT_CONFIG_PATH
|
||||
else:
|
||||
# Try relative to script location
|
||||
script_dir = Path(__file__).parent.parent.parent.parent
|
||||
config_path = script_dir / "vllm_omni/model_executor/stage_configs/glm_image.yaml"
|
||||
if not config_path.exists():
|
||||
raise FileNotFoundError(
|
||||
f"Stage config not found. Please specify --config-path. Tried: {DEFAULT_CONFIG_PATH}"
|
||||
)
|
||||
config_path = str(config_path)
|
||||
|
||||
print(f"Model path: {args.model_path}")
|
||||
print(f"Config path: {config_path}")
|
||||
print(f"Prompt: {args.prompt}")
|
||||
|
||||
# Load source image for image-to-image mode
|
||||
source_image = None
|
||||
if args.image:
|
||||
print(f"Source image: {args.image}")
|
||||
source_image = load_image(args.image)
|
||||
print(f" Image size: {source_image.width}x{source_image.height}")
|
||||
|
||||
# Build prompt based on mode
|
||||
if source_image is not None:
|
||||
# Image-to-image mode
|
||||
prompt_dict = build_prompt_for_i2i(
|
||||
prompt=args.prompt,
|
||||
image=source_image,
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
)
|
||||
mode = "image-to-image"
|
||||
else:
|
||||
# Text-to-image mode
|
||||
prompt_dict = build_prompt_for_t2i(
|
||||
prompt=args.prompt,
|
||||
height=args.height or 1024,
|
||||
width=args.width or 1024,
|
||||
)
|
||||
mode = "text-to-image"
|
||||
|
||||
print(f"Mode: {mode}")
|
||||
print(f"Target size: {prompt_dict.get('height', 1024)}x{prompt_dict.get('width', 1024)}")
|
||||
|
||||
# Add generation parameters to prompt
|
||||
prompt_dict["seed"] = args.seed
|
||||
prompt_dict["num_inference_steps"] = args.num_inference_steps
|
||||
prompt_dict["guidance_scale"] = args.guidance_scale
|
||||
|
||||
if args.negative_prompt:
|
||||
prompt_dict["negative_prompt"] = args.negative_prompt
|
||||
|
||||
# Build cache-dit config if requested
|
||||
cache_config = None
|
||||
if args.cache_backend == "cache_dit":
|
||||
cache_config = {
|
||||
"Fn_compute_blocks": 1,
|
||||
"Bn_compute_blocks": 0,
|
||||
"max_warmup_steps": 4,
|
||||
"residual_diff_threshold": 0.24,
|
||||
"max_continuous_cached_steps": 3,
|
||||
"enable_taylorseer": False,
|
||||
"taylorseer_order": 1,
|
||||
"scm_steps_mask_policy": None,
|
||||
"scm_steps_policy": "dynamic",
|
||||
}
|
||||
|
||||
# Initialize Omni with multistage config
|
||||
print("\nInitializing Omni with multistage pipeline...")
|
||||
print(f" Cache backend: {args.cache_backend or 'None (no acceleration)'}")
|
||||
start_time = time.time()
|
||||
|
||||
omni = Omni(
|
||||
model=args.model_path,
|
||||
stage_configs_path=config_path,
|
||||
log_stats=args.enable_stats,
|
||||
stage_init_timeout=args.stage_init_timeout,
|
||||
cache_backend=args.cache_backend,
|
||||
cache_config=cache_config,
|
||||
enable_cache_dit_summary=getattr(args, "enable_cache_dit_summary", False),
|
||||
enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
|
||||
)
|
||||
|
||||
init_time = time.time() - start_time
|
||||
print(f"Initialization completed in {init_time:.2f}s")
|
||||
|
||||
# Prepare prompts (support batch generation)
|
||||
prompts = [prompt_dict for _ in range(args.num_prompts)]
|
||||
|
||||
# No explicit sampling_params for diffusion - parameters are in prompt_dict
|
||||
# For multistage, the AR stage may need sampling params
|
||||
from vllm import SamplingParams
|
||||
|
||||
# Compute max_tokens dynamically based on target image size
|
||||
target_height = prompt_dict.get("height", 1024)
|
||||
target_width = prompt_dict.get("width", 1024)
|
||||
is_i2i = source_image is not None
|
||||
calculated_max_tokens = compute_max_tokens(target_height, target_width, is_i2i=is_i2i)
|
||||
|
||||
# Use calculated value unless user explicitly specified a different value
|
||||
# Default args.max_tokens is 16384 (very large), so prefer calculated value
|
||||
effective_max_tokens = calculated_max_tokens if args.max_tokens == 16384 else args.max_tokens
|
||||
|
||||
if args.verbose:
|
||||
print(
|
||||
f"AR max_tokens: {effective_max_tokens} "
|
||||
f"(calculated: {calculated_max_tokens}, arg: {args.max_tokens}, mode: {'i2i' if is_i2i else 't2i'})"
|
||||
)
|
||||
|
||||
# IMPORTANT: GLM-Image AR model requires these exact sampling parameters
|
||||
# from generation_config.json for proper image token generation.
|
||||
# - temperature=0.9, top_p=0.75, top_k=16512 (vision_vocab_size)
|
||||
# - stop_token_ids=[16385] (eos_token_id) is CRITICAL to stop generation
|
||||
ar_sampling_params = SamplingParams(
|
||||
temperature=0.9, # From generation_config.json
|
||||
top_p=0.75, # From generation_config.json
|
||||
top_k=GLM_IMAGE_VISION_VOCAB_SIZE, # 16512, vision vocabulary size
|
||||
max_tokens=effective_max_tokens,
|
||||
stop_token_ids=[GLM_IMAGE_EOS_TOKEN_ID], # 16385, CRITICAL for stopping
|
||||
seed=args.seed,
|
||||
detokenize=False,
|
||||
# Keep target size available in runner/model for deterministic M-RoPE
|
||||
# decode grids in t2i (no mm_features available in this path).
|
||||
extra_args={
|
||||
"target_h": int(target_height),
|
||||
"target_w": int(target_width),
|
||||
},
|
||||
)
|
||||
|
||||
# For diffusion stage, sampling_params contains diffusion-specific parameters
|
||||
# These are passed as kwargs to the diffusion engine
|
||||
diffusion_sampling_params = OmniDiffusionSamplingParams(
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
guidance_scale=args.guidance_scale,
|
||||
height=prompt_dict.get("height", 1024),
|
||||
width=prompt_dict.get("width", 1024),
|
||||
seed=args.seed,
|
||||
)
|
||||
|
||||
# For multistage, we need sampling_params for each stage
|
||||
# Stage 0 (AR): SamplingParams for vLLM
|
||||
# Stage 1 (Diffusion): dict with diffusion kwargs
|
||||
sampling_params_list = [ar_sampling_params, diffusion_sampling_params]
|
||||
|
||||
# Run generation
|
||||
print(f"\nGenerating {args.num_prompts} image(s)...")
|
||||
gen_start_time = time.time()
|
||||
|
||||
output_dir = os.path.dirname(args.output) if args.output else "outputs"
|
||||
if output_dir:
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
output_count = 0
|
||||
for stage_outputs in omni.generate(prompts, sampling_params_list, py_generator=True):
|
||||
output = stage_outputs.request_output
|
||||
if stage_outputs.final_output_type == "image":
|
||||
request_id = output.request_id
|
||||
|
||||
# Get generated images
|
||||
images = output.images if hasattr(output, "images") else []
|
||||
if not images and hasattr(output, "multimodal_output"):
|
||||
images = output.multimodal_output.get("images", [])
|
||||
|
||||
# Save each generated image
|
||||
for idx, img in enumerate(images):
|
||||
if args.num_prompts == 1 and len(images) == 1:
|
||||
output_path = args.output
|
||||
else:
|
||||
base, ext = os.path.splitext(args.output)
|
||||
output_path = f"{base}_{request_id}_{idx}{ext}"
|
||||
|
||||
if isinstance(img, Image.Image):
|
||||
save_image(img, output_path)
|
||||
else:
|
||||
print(f"Warning: Unexpected image type for request {request_id}: {type(img)}")
|
||||
|
||||
output_count += 1
|
||||
|
||||
elif stage_outputs.final_output_type == "text":
|
||||
# AR stage output (intermediate, for debugging)
|
||||
if args.verbose:
|
||||
print(f"AR output for request {output.request_id}:")
|
||||
print(f" Token count: {len(output.outputs[0].token_ids)}")
|
||||
|
||||
gen_time = time.time() - gen_start_time
|
||||
print(f"\nGeneration completed in {gen_time:.2f}s")
|
||||
print(f"Generated {output_count} image(s)")
|
||||
|
||||
# Cleanup
|
||||
omni.close()
|
||||
print("\nDone!")
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
"""Parse command line arguments."""
|
||||
parser = argparse.ArgumentParser(
|
||||
description="GLM-Image Multistage End-to-End Inference",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog=__doc__,
|
||||
)
|
||||
|
||||
# Required arguments
|
||||
parser.add_argument(
|
||||
"--model-path",
|
||||
type=str,
|
||||
default="zai-org/GLM-Image",
|
||||
help="Path to GLM-Image model directory or HuggingFace model ID",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompt",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Text prompt for image generation",
|
||||
)
|
||||
|
||||
# Optional arguments
|
||||
parser.add_argument(
|
||||
"--config-path",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to stage config YAML file (default: auto-detect)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--image",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to source image for image-to-image mode",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
type=str,
|
||||
default="output_glm_image.png",
|
||||
help="Output image path (default: output_glm_image.png)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--negative-prompt",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Negative prompt for classifier-free guidance",
|
||||
)
|
||||
|
||||
# Generation parameters
|
||||
parser.add_argument(
|
||||
"--height",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Output image height (default: 1024 for t2i, source size for i2i)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--width",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Output image width (default: 1024 for t2i, source size for i2i)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-inference-steps",
|
||||
type=int,
|
||||
default=50,
|
||||
help="Number of diffusion denoising steps (default: 50)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--guidance-scale",
|
||||
type=float,
|
||||
default=1.5,
|
||||
help="Classifier-free guidance scale (default: 1.5)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--seed",
|
||||
type=int,
|
||||
default=SEED,
|
||||
help=f"Random seed for reproducibility (default: {SEED})",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-tokens",
|
||||
type=int,
|
||||
default=16384,
|
||||
help="Maximum tokens for AR generation (default: 16384)",
|
||||
)
|
||||
|
||||
# Batch processing
|
||||
parser.add_argument(
|
||||
"--num-prompts",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of images to generate (default: 1)",
|
||||
)
|
||||
|
||||
# Cache acceleration
|
||||
parser.add_argument(
|
||||
"--cache-backend",
|
||||
type=str,
|
||||
default=None,
|
||||
choices=["cache_dit"],
|
||||
help="Cache backend for DiT acceleration. Default: None (no cache).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enable-cache-dit-summary",
|
||||
action="store_true",
|
||||
help="Enable cache-dit summary logging after diffusion forward passes.",
|
||||
)
|
||||
|
||||
# Runtime options
|
||||
parser.add_argument(
|
||||
"--enable-stats",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Enable statistics logging",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stage-init-timeout",
|
||||
type=int,
|
||||
default=600,
|
||||
help="Timeout for stage initialization in seconds (default: 600)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--verbose",
|
||||
"-v",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Enable verbose output",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enable-diffusion-pipeline-profiler",
|
||||
action="store_true",
|
||||
help="Enable diffusion pipeline profiler to display stage durations.",
|
||||
)
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
main(args)
|
||||
@@ -1,93 +0,0 @@
|
||||
#!/bin/bash
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# Run GLM-Image image-to-image (editing) with multistage pipeline
|
||||
|
||||
set -e
|
||||
|
||||
# Default values
|
||||
MODEL_PATH="${MODEL_PATH:-/path/to/glm-image}"
|
||||
CONFIG_PATH="${CONFIG_PATH:-vllm_omni/model_executor/stage_configs/glm_image.yaml}"
|
||||
PROMPT="${PROMPT:-Transform this image into an oil painting style}"
|
||||
INPUT_IMAGE=""
|
||||
OUTPUT="${OUTPUT:-output_i2i.png}"
|
||||
NUM_STEPS="${NUM_STEPS:-50}"
|
||||
GUIDANCE="${GUIDANCE:-1.5}"
|
||||
SEED="${SEED:-42}"
|
||||
|
||||
# Parse command line arguments
|
||||
while [[ $# -gt 0 ]]; do
|
||||
case $1 in
|
||||
--model-path)
|
||||
MODEL_PATH="$2"
|
||||
shift 2
|
||||
;;
|
||||
--config-path)
|
||||
CONFIG_PATH="$2"
|
||||
shift 2
|
||||
;;
|
||||
--prompt)
|
||||
PROMPT="$2"
|
||||
shift 2
|
||||
;;
|
||||
--image)
|
||||
INPUT_IMAGE="$2"
|
||||
shift 2
|
||||
;;
|
||||
--output)
|
||||
OUTPUT="$2"
|
||||
shift 2
|
||||
;;
|
||||
--num-steps)
|
||||
NUM_STEPS="$2"
|
||||
shift 2
|
||||
;;
|
||||
--guidance)
|
||||
GUIDANCE="$2"
|
||||
shift 2
|
||||
;;
|
||||
--seed)
|
||||
SEED="$2"
|
||||
shift 2
|
||||
;;
|
||||
*)
|
||||
echo "Unknown option: $1"
|
||||
exit 1
|
||||
;;
|
||||
esac
|
||||
done
|
||||
|
||||
# Check if input image is provided
|
||||
if [ -z "${INPUT_IMAGE}" ]; then
|
||||
echo "Error: --image is required for image-to-image mode"
|
||||
echo "Usage: ./run_i2i.sh --image /path/to/input.png [--prompt \"edit instruction\"]"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
if [ ! -f "${INPUT_IMAGE}" ]; then
|
||||
echo "Error: Input image not found: ${INPUT_IMAGE}"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "=============================================="
|
||||
echo "GLM-Image Image-to-Image Generation"
|
||||
echo "=============================================="
|
||||
echo "Model: ${MODEL_PATH}"
|
||||
echo "Config: ${CONFIG_PATH}"
|
||||
echo "Input: ${INPUT_IMAGE}"
|
||||
echo "Prompt: ${PROMPT}"
|
||||
echo "Output: ${OUTPUT}"
|
||||
echo "Steps: ${NUM_STEPS}"
|
||||
echo "Guidance: ${GUIDANCE}"
|
||||
echo "Seed: ${SEED}"
|
||||
echo "=============================================="
|
||||
|
||||
python end2end.py \
|
||||
--model-path "${MODEL_PATH}" \
|
||||
--config-path "${CONFIG_PATH}" \
|
||||
--prompt "${PROMPT}" \
|
||||
--image "${INPUT_IMAGE}" \
|
||||
--output "${OUTPUT}" \
|
||||
--num-inference-steps "${NUM_STEPS}" \
|
||||
--guidance-scale "${GUIDANCE}" \
|
||||
--seed "${SEED}" \
|
||||
--verbose
|
||||
@@ -1,87 +0,0 @@
|
||||
#!/bin/bash
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# Run GLM-Image text-to-image generation with multistage pipeline
|
||||
|
||||
set -e
|
||||
|
||||
# Default values
|
||||
MODEL_PATH="${MODEL_PATH:-/path/to/glm-image}"
|
||||
CONFIG_PATH="${CONFIG_PATH:-vllm_omni/model_executor/stage_configs/glm_image.yaml}"
|
||||
PROMPT="${PROMPT:-A beautiful mountain landscape with snow-capped peaks and a clear blue lake}"
|
||||
OUTPUT="${OUTPUT:-output_t2i.png}"
|
||||
HEIGHT="${HEIGHT:-1024}"
|
||||
WIDTH="${WIDTH:-1024}"
|
||||
NUM_STEPS="${NUM_STEPS:-50}"
|
||||
GUIDANCE="${GUIDANCE:-1.5}"
|
||||
SEED="${SEED:-42}"
|
||||
|
||||
# Parse command line arguments
|
||||
while [[ $# -gt 0 ]]; do
|
||||
case $1 in
|
||||
--model-path)
|
||||
MODEL_PATH="$2"
|
||||
shift 2
|
||||
;;
|
||||
--config-path)
|
||||
CONFIG_PATH="$2"
|
||||
shift 2
|
||||
;;
|
||||
--prompt)
|
||||
PROMPT="$2"
|
||||
shift 2
|
||||
;;
|
||||
--output)
|
||||
OUTPUT="$2"
|
||||
shift 2
|
||||
;;
|
||||
--height)
|
||||
HEIGHT="$2"
|
||||
shift 2
|
||||
;;
|
||||
--width)
|
||||
WIDTH="$2"
|
||||
shift 2
|
||||
;;
|
||||
--num-steps)
|
||||
NUM_STEPS="$2"
|
||||
shift 2
|
||||
;;
|
||||
--guidance)
|
||||
GUIDANCE="$2"
|
||||
shift 2
|
||||
;;
|
||||
--seed)
|
||||
SEED="$2"
|
||||
shift 2
|
||||
;;
|
||||
*)
|
||||
echo "Unknown option: $1"
|
||||
exit 1
|
||||
;;
|
||||
esac
|
||||
done
|
||||
|
||||
echo "=============================================="
|
||||
echo "GLM-Image Text-to-Image Generation"
|
||||
echo "=============================================="
|
||||
echo "Model: ${MODEL_PATH}"
|
||||
echo "Config: ${CONFIG_PATH}"
|
||||
echo "Prompt: ${PROMPT}"
|
||||
echo "Output: ${OUTPUT}"
|
||||
echo "Size: ${WIDTH}x${HEIGHT}"
|
||||
echo "Steps: ${NUM_STEPS}"
|
||||
echo "Guidance: ${GUIDANCE}"
|
||||
echo "Seed: ${SEED}"
|
||||
echo "=============================================="
|
||||
|
||||
python end2end.py \
|
||||
--model-path "${MODEL_PATH}" \
|
||||
--config-path "${CONFIG_PATH}" \
|
||||
--prompt "${PROMPT}" \
|
||||
--output "${OUTPUT}" \
|
||||
--height "${HEIGHT}" \
|
||||
--width "${WIDTH}" \
|
||||
--num-inference-steps "${NUM_STEPS}" \
|
||||
--guidance-scale "${GUIDANCE}" \
|
||||
--seed "${SEED}" \
|
||||
--verbose
|
||||
@@ -196,6 +196,9 @@ def parse_args() -> argparse.Namespace:
|
||||
parser.add_argument("--cfg-parallel-size", type=int, default=1, choices=[1, 2], help="CFG parallel size.")
|
||||
parser.add_argument("--tensor-parallel-size", type=int, default=1, help="Tensor parallelism size.")
|
||||
|
||||
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||
|
||||
nullify_stage_engine_defaults(parser)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
|
||||
@@ -135,6 +135,9 @@ def parse_args():
|
||||
parser.add_argument("--init-timeout", type=int, default=300, help="Initialization timeout in seconds.")
|
||||
parser.add_argument("--enforce-eager", action="store_true", help="Disable torch.compile.")
|
||||
|
||||
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||
|
||||
nullify_stage_engine_defaults(parser)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
|
||||
@@ -87,9 +87,11 @@ For more options, run:
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from PIL import Image
|
||||
@@ -101,6 +103,16 @@ from vllm_omni.outputs import OmniRequestOutput
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
|
||||
def parse_profiler_config(value: str) -> dict[str, Any]:
|
||||
try:
|
||||
config = json.loads(value)
|
||||
except json.JSONDecodeError as e:
|
||||
raise argparse.ArgumentTypeError(f"--profiler-config must be valid JSON: {e}") from e
|
||||
if not isinstance(config, dict):
|
||||
raise argparse.ArgumentTypeError("--profiler-config must be a JSON object")
|
||||
return config
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Edit an image with Qwen-Image-Edit.")
|
||||
parser.add_argument(
|
||||
@@ -196,6 +208,13 @@ def parse_args() -> argparse.Namespace:
|
||||
default=1,
|
||||
help="Number of GPUs used for ulysses sequence parallelism.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ulysses-mode",
|
||||
type=str,
|
||||
default="strict",
|
||||
choices=["strict", "advanced_uaa"],
|
||||
help="Ulysses sequence-parallel mode: 'strict' (divisibility required) or 'advanced_uaa' (UAA).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ring-degree",
|
||||
type=int,
|
||||
@@ -325,11 +344,43 @@ def parse_args() -> argparse.Namespace:
|
||||
action="store_true",
|
||||
help="Enable layerwise (blockwise) offloading on DiT modules.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--vae-patch-parallel-size",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of GPUs used for VAE patch/tile parallelism (decode).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--use-hsdp",
|
||||
action="store_true",
|
||||
help="Enable HSDP (Hybrid Sharded Data Parallel) for diffusion models.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--hsdp-shard-size",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of GPUs to shard weights across for HSDP.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--hsdp-replicate-size",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of HSDP replica groups.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enable-diffusion-pipeline-profiler",
|
||||
action="store_true",
|
||||
help="Enable diffusion pipeline profiler to display stage durations.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--profiler-config",
|
||||
type=parse_profiler_config,
|
||||
default=None,
|
||||
help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
|
||||
)
|
||||
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||
|
||||
nullify_stage_engine_defaults(parser)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
@@ -395,11 +446,11 @@ def main():
|
||||
enforce_eager=args.enforce_eager,
|
||||
enable_cpu_offload=args.enable_cpu_offload,
|
||||
enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
|
||||
profiler_config=args.profiler_config,
|
||||
)
|
||||
print("Pipeline loaded")
|
||||
|
||||
# Check if profiling is requested via environment variable
|
||||
profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
|
||||
profiler_enabled = args.profiler_config is not None
|
||||
|
||||
# Time profiling for generation
|
||||
print(f"\n{'=' * 60}")
|
||||
|
||||
@@ -33,9 +33,10 @@ Usage:
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import json
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import PIL.Image
|
||||
@@ -48,6 +49,16 @@ from vllm_omni.outputs import OmniRequestOutput
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
|
||||
def parse_profiler_config(value: str) -> dict[str, Any]:
|
||||
try:
|
||||
config = json.loads(value)
|
||||
except json.JSONDecodeError as e:
|
||||
raise argparse.ArgumentTypeError(f"--profiler-config must be valid JSON: {e}") from e
|
||||
if not isinstance(config, dict):
|
||||
raise argparse.ArgumentTypeError("--profiler-config must be a JSON object")
|
||||
return config
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Generate a video from an image (Wan2.2, LTX2, HunyuanVideo-1.5).")
|
||||
parser.add_argument(
|
||||
@@ -194,6 +205,15 @@ def parse_args() -> argparse.Namespace:
|
||||
action="store_true",
|
||||
help="Enable diffusion pipeline profiler to display stage durations.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--profiler-config",
|
||||
type=parse_profiler_config,
|
||||
default=None,
|
||||
help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
|
||||
)
|
||||
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||
|
||||
nullify_stage_engine_defaults(parser)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
@@ -274,8 +294,7 @@ def main():
|
||||
"rel_l1_thresh": 0.2,
|
||||
}
|
||||
|
||||
# Check if profiling is requested via environment variable
|
||||
profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
|
||||
profiler_enabled = args.profiler_config is not None
|
||||
parallel_config = DiffusionParallelConfig(
|
||||
ulysses_degree=args.ulysses_degree,
|
||||
ring_degree=args.ring_degree,
|
||||
@@ -300,6 +319,7 @@ def main():
|
||||
cache_backend=args.cache_backend,
|
||||
cache_config=cache_config,
|
||||
enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
|
||||
profiler_config=args.profiler_config,
|
||||
)
|
||||
|
||||
if profiler_enabled:
|
||||
|
||||
@@ -24,6 +24,9 @@ def parse_args():
|
||||
parser.add_argument("--width", type=int, default=448, help="Video width.")
|
||||
parser.add_argument("--num-inference-steps", type=int, default=8, help="Number of denoising steps.")
|
||||
parser.add_argument("--seed", type=int, default=52, help="Random seed for generation.")
|
||||
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||
|
||||
nullify_stage_engine_defaults(parser)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
|
||||
@@ -19,6 +19,7 @@ from vllm import SamplingParams
|
||||
from vllm.multimodal.image import convert_image_mode
|
||||
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||
|
||||
DEFAULT_SYSTEM = "You are a helpful assistant."
|
||||
DEFAULT_QUESTION = "Please summarize the content of this image."
|
||||
@@ -48,6 +49,7 @@ def parse_args() -> argparse.Namespace:
|
||||
action="store_true",
|
||||
help="Enable diffusion pipeline profiler to display stage durations.",
|
||||
)
|
||||
nullify_stage_engine_defaults(parser)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
|
||||
@@ -29,6 +29,7 @@ from PIL import Image
|
||||
from vllm.sampling_params import SamplingParams
|
||||
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -117,6 +118,7 @@ def parse_args() -> argparse.Namespace:
|
||||
)
|
||||
p.add_argument("--out", type=str, default="output.png", help="Path to save the generated image.")
|
||||
p.add_argument("--trust-remote-code", action="store_true", help="Trust remote code when loading the model.")
|
||||
nullify_stage_engine_defaults(p)
|
||||
args = p.parse_args()
|
||||
if not args.prompt:
|
||||
args.prompt = ["A stylish woman with sunglasses riding a motorcycle in NYC."]
|
||||
|
||||
@@ -39,7 +39,6 @@ Run a single sample for basic TTS:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type tts_sft
|
||||
```
|
||||
@@ -48,7 +47,6 @@ Run batch samples for basic TTS:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type tts_sft \
|
||||
--num-prompts {batch_size}
|
||||
@@ -66,7 +64,6 @@ Generate speech from text input:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type tts_sft \
|
||||
--text "The weather is so nice today."
|
||||
@@ -78,7 +75,6 @@ Generate speech with explicit voice style instructions:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type tts_sft_with_instruct \
|
||||
--text "The weather is so nice today." \
|
||||
@@ -91,7 +87,6 @@ Generate speech using an audio reference for voice cloning:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type tts_sft_with_audio \
|
||||
--text "The weather is so nice today." \
|
||||
@@ -104,7 +99,6 @@ Generate speech from text containing natural voice descriptions:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type tts_sft_with_natural_instruction \
|
||||
--text "In a panting young male voice, he said: I can't run anymore, wait for me!"
|
||||
@@ -116,7 +110,6 @@ Transcribe audio to text:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type audio_trancribing_sft \
|
||||
--audio-path "./spoken_dialogue_assistant_turn_1.wav"
|
||||
@@ -128,7 +121,6 @@ Understand and analyze audio content with text queries:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type audio_understanding_sft \
|
||||
--text "Summarize the audio." \
|
||||
@@ -141,7 +133,6 @@ Audio understanding with reasoning chain:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type audio_understanding_sft_with_thinking \
|
||||
--text "Summarize the audio." \
|
||||
@@ -154,7 +145,6 @@ Multi-turn dialogue with audio input and output:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type spoken_dialogue_sft_multiturn \
|
||||
--audio-path "./prompt_speech_zh_m.wav"
|
||||
@@ -168,7 +158,6 @@ Multi-turn dialogue converting speech to text:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type speech2text_dialogue_sft_multiturn
|
||||
```
|
||||
@@ -181,7 +170,6 @@ Multi-turn text-only dialogue:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type text_dialogue_sft_multiturn
|
||||
```
|
||||
|
||||
@@ -182,7 +182,7 @@ def main(args):
|
||||
|
||||
omni = Omni(
|
||||
model=model_name,
|
||||
stage_configs_path=args.stage_configs_path,
|
||||
deploy_config=args.deploy_config,
|
||||
log_stats=args.enable_stats,
|
||||
log_file=("omni_pipeline.log" if args.enable_stats else None),
|
||||
init_sleep_seconds=args.init_sleep_seconds,
|
||||
@@ -309,7 +309,10 @@ def main(args):
|
||||
lines.append("Prompt:\n")
|
||||
lines.append(str(prompt_text) + "\n")
|
||||
lines.append("vllm_text_output:\n")
|
||||
lines.append(str(text_output).strip() + "\n")
|
||||
output_text = str(text_output)
|
||||
if "<chinese>" in output_text or "<english>" in output_text:
|
||||
output_text = output_text.replace("<chinese>", "").replace("<english>", "").strip()
|
||||
lines.append(output_text + "\n")
|
||||
try:
|
||||
with open(out_txt, "w", encoding="utf-8") as f:
|
||||
print("lines", lines)
|
||||
@@ -351,7 +354,7 @@ def parse_args():
|
||||
"--text",
|
||||
"-t",
|
||||
type=str,
|
||||
default="The weather is so nice today.",
|
||||
default="",
|
||||
help="input text",
|
||||
)
|
||||
parser.add_argument(
|
||||
@@ -428,10 +431,11 @@ def parse_args():
|
||||
help="Sampling rate for audio.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stage-configs-path",
|
||||
"--deploy-config",
|
||||
type=str,
|
||||
default="../../../model_executor/stage_configs/mimo_audio.yaml",
|
||||
help="Path to a stage configs file.",
|
||||
default=None,
|
||||
help="Override the deploy config path. If unset, auto-loads "
|
||||
"vllm_omni/deploy/mimo_audio.yaml based on the HF model_type.",
|
||||
)
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
@@ -1,75 +1,91 @@
|
||||
# Ming-flash-omni 2.0
|
||||
|
||||
[Ming-flash-omni-2.0](https://github.com/inclusionAI/Ming) is an omni-modal model supporting text, image, video, and audio understanding, with outputs in text, image, and audio. For now, Ming-flash-omni-2.0 in vLLM-Omni is supported with thinker stage (multi-modal understanding).
|
||||
[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 two deployment modes:
|
||||
|
||||
| Mode | Stage config | Output |
|
||||
|------|-------------|--------|
|
||||
| Thinker only (multimodal understanding) | `ming_flash_omni_thinker.yaml` (default `--omni`) | Text |
|
||||
| Thinker + Talker (omni-speech) | `ming_flash_omni.yaml` | Text + Audio |
|
||||
|
||||
For standalone TTS (talker only), see [`examples/offline_inference/ming_flash_omni_tts/`](../ming_flash_omni_tts/).
|
||||
|
||||
## Setup
|
||||
|
||||
Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
|
||||
|
||||
The default `--omni` flag runs thinker only. For omni-speech, pass the two-stage config explicitly:
|
||||
|
||||
```bash
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/ming_flash_omni.yaml
|
||||
```
|
||||
|
||||
## Run examples
|
||||
|
||||
### Text-only
|
||||
The end-to-end script defaults to built-in assets; pass `--image-path`,
|
||||
`--audio-path`, or `--video-path` to override.
|
||||
|
||||
```bash
|
||||
# Text-only
|
||||
python examples/offline_inference/ming_flash_omni/end2end.py --query-type text
|
||||
|
||||
# Image / audio / video / mixed understanding
|
||||
python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_image
|
||||
python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_audio
|
||||
python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_video --num-frames 16
|
||||
python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_mixed_modalities \
|
||||
--image-path /path/to/image.jpg --audio-path /path/to/audio.wav
|
||||
```
|
||||
|
||||
#### Reasoning (Thinking Mode)
|
||||
|
||||
Reasoning (Thinking) mode is enabled via applying "detailed thinking on" when building the system prompt template (in `apply_chat_template`).
|
||||
|
||||
In the end2end example, a default problem for thinking mode is provided, as referred to the example usage of Ming's cookbook;
|
||||
To utilize it, you have to download the example figure from https://github.com/inclusionAI/Ming/blob/3954fcb880ff5e61ff128bcf7f1ec344d46a6fe3/figures/cases/3_0.png
|
||||
Reasoning ("detailed thinking on") is applied by the script when
|
||||
`--query-type reasoning` is set. The default prompt matches Ming's cookbook
|
||||
and expects the reference figure from the upstream repo — see
|
||||
`get_reasoning_query` in `end2end.py`.
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/ming_flash_omni/end2end.py -q reasoning --image-path ./3_0.png
|
||||
```
|
||||
|
||||
### Image understanding
|
||||
```bash
|
||||
python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_image
|
||||
### Omni-speech (thinker + talker)
|
||||
|
||||
# With a local image
|
||||
python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_image --image-path /path/to/image.jpg
|
||||
To enable spoken output, use the two-stage config and request `audio` (or `text,audio`) modalities.
|
||||
The thinker processes your multimodal input, generates text, then the talker synthesises the response as speech.
|
||||
|
||||
**Audio-only output** (speech response, no text):
|
||||
```bash
|
||||
python examples/offline_inference/ming_flash_omni/end2end.py \
|
||||
--query-type text \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/ming_flash_omni.yaml \
|
||||
--modalities audio \
|
||||
--output-dir output_ming_omni_speech
|
||||
```
|
||||
|
||||
### Audio understanding
|
||||
**Both text and audio output**:
|
||||
```bash
|
||||
python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_audio
|
||||
|
||||
# With a local audio file
|
||||
python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_audio --audio-path /path/to/audio.wav
|
||||
python examples/offline_inference/ming_flash_omni/end2end.py \
|
||||
--query-type use_audio \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/ming_flash_omni.yaml \
|
||||
--modalities text,audio \
|
||||
--output-dir output_ming_omni_speech
|
||||
```
|
||||
|
||||
### Video understanding
|
||||
```bash
|
||||
python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_video
|
||||
Generated `.wav` files are saved to `--output-dir` (default `output_ming`), one per request.
|
||||
|
||||
# With a local video and custom frame count
|
||||
python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_video --video-path /path/to/video.mp4 --num-frames 16
|
||||
```
|
||||
|
||||
### Mixed modalities (image + audio)
|
||||
```bash
|
||||
python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_mixed_modalities \
|
||||
--image-path /path/to/image.jpg \
|
||||
--audio-path /path/to/audio.wav
|
||||
```
|
||||
|
||||
If media file paths are not provided, the script uses built-in default assets.
|
||||
The stage config allocates thinker on GPUs 0–3 and talker on GPU 3 by default. Adjust `devices` in the YAML to match your hardware.
|
||||
|
||||
### Modality control
|
||||
To control output modalities (e.g. text-only output):
|
||||
```bash
|
||||
python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_audio --modalities text
|
||||
```
|
||||
|
||||
*For now, only text output is supported*
|
||||
| `--modalities` | Thinker output | Talker | Saved files |
|
||||
|---------------|----------------|--------|-------------|
|
||||
| `text` (default) | Text | Not run | `<id>.txt` |
|
||||
| `audio` | Text (internal) | Runs | `<id>.wav` |
|
||||
| `text,audio` | Text | Runs | `<id>.txt` + `<id>.wav` |
|
||||
|
||||
### Custom stage config
|
||||
```bash
|
||||
python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_image \
|
||||
--stage-configs-path /path/to/your_config.yaml
|
||||
```
|
||||
Pass `--stage-configs-path /path/to/your_config.yaml` to any of the commands
|
||||
above to override the stage config.
|
||||
|
||||
## Online serving
|
||||
|
||||
|
||||
@@ -6,8 +6,8 @@ import os
|
||||
import time
|
||||
from typing import NamedTuple
|
||||
|
||||
import librosa
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
import vllm
|
||||
from PIL import Image
|
||||
from transformers import AutoProcessor
|
||||
@@ -16,6 +16,7 @@ from vllm.assets.audio import AudioAsset
|
||||
from vllm.assets.image import ImageAsset
|
||||
from vllm.assets.video import VideoAsset, video_to_ndarrays
|
||||
from vllm.multimodal.image import convert_image_mode
|
||||
from vllm.multimodal.media.audio import load_audio
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
import vllm_omni
|
||||
@@ -91,7 +92,7 @@ def get_audio_query(
|
||||
if audio_path:
|
||||
if not os.path.exists(audio_path):
|
||||
raise FileNotFoundError(f"Audio file not found: {audio_path}")
|
||||
audio_signal, sr = librosa.load(audio_path, sr=sampling_rate)
|
||||
audio_signal, sr = load_audio(audio_path, sr=sampling_rate)
|
||||
audio_data = (audio_signal.astype(np.float32), sr)
|
||||
else:
|
||||
audio_data = AudioAsset("mary_had_lamb").audio_and_sample_rate
|
||||
@@ -172,7 +173,7 @@ def get_mixed_modalities_query(
|
||||
if audio_path:
|
||||
if not os.path.exists(audio_path):
|
||||
raise FileNotFoundError(f"Audio file not found: {audio_path}")
|
||||
sig, sr = librosa.load(audio_path, sr=sampling_rate)
|
||||
sig, sr = load_audio(audio_path, sr=sampling_rate)
|
||||
audio_data = (sig.astype(np.float32), sr)
|
||||
else:
|
||||
audio_data = AudioAsset("mary_had_lamb").audio_and_sample_rate
|
||||
@@ -319,7 +320,16 @@ def main(args):
|
||||
seed=SEED,
|
||||
detokenize=True,
|
||||
)
|
||||
sampling_params_list = [thinker_sampling_params]
|
||||
# Talker (ming_tts) uses a custom generation loop (CFM + AudioVAE);
|
||||
# vLLM sampling is a no-op here — max_tokens=1 just satisfies the scheduler.
|
||||
talker_sampling_params = SamplingParams(
|
||||
temperature=0.0,
|
||||
max_tokens=1,
|
||||
)
|
||||
all_sampling_params = [thinker_sampling_params, talker_sampling_params]
|
||||
# Match sampling params to the number of configured stages
|
||||
# (thinker-only yaml → 1, thinker+talker yaml → 2).
|
||||
sampling_params_list = all_sampling_params[: omni.num_stages]
|
||||
|
||||
prompts = [query_result.inputs for _ in range(args.num_prompts)]
|
||||
|
||||
@@ -362,7 +372,19 @@ def main(args):
|
||||
print(f"Failed to write output file {out_txt}: {e}")
|
||||
|
||||
elif stage_outputs.final_output_type == "audio":
|
||||
raise NotImplementedError("Add audio example after talker supported.")
|
||||
request_id = output.request_id
|
||||
mm = output.outputs[0].multimodal_output
|
||||
if mm and "audio" in mm:
|
||||
audio = mm["audio"]
|
||||
sr_raw = mm.get("sr", 44100)
|
||||
sample_rate = int(sr_raw.item() if hasattr(sr_raw, "item") else sr_raw)
|
||||
audio_numpy = audio.float().squeeze().cpu().numpy()
|
||||
output_wav = os.path.join(output_dir, f"{request_id}.wav")
|
||||
sf.write(output_wav, audio_numpy, samplerate=sample_rate, format="WAV")
|
||||
print(
|
||||
f"Request ID: {request_id}, audio saved to {output_wav} "
|
||||
f"({len(audio_numpy) / sample_rate:.2f}s, {sample_rate}Hz)"
|
||||
)
|
||||
|
||||
processed_count += 1
|
||||
if profiler_enabled and processed_count >= total_requests:
|
||||
|
||||
@@ -0,0 +1,47 @@
|
||||
# Ming-flash-omni Standalone TTS (Offline)
|
||||
|
||||
This example runs **Ming-flash-omni-2.0 talker-only** offline inference with:
|
||||
|
||||
- `model`: `Jonathan1909/Ming-flash-omni-2.0`
|
||||
- `stage config`: `vllm_omni/model_executor/stage_configs/ming_flash_omni_tts.yaml`
|
||||
|
||||
It follows the Ming cookbook parameter style:
|
||||
|
||||
- `prompt`: `"Please generate speech based on the following description.\n"`
|
||||
- `max_decode_steps`: `200`
|
||||
- `cfg`: `2.0`
|
||||
- `sigma`: `0.25`
|
||||
- `temperature`: `0.0`
|
||||
|
||||
## Quick Start
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/ming_flash_omni_tts/end2end.py --case style
|
||||
```
|
||||
|
||||
## Cases
|
||||
|
||||
```bash
|
||||
# Style
|
||||
python examples/offline_inference/ming_flash_omni_tts/end2end.py --case style
|
||||
|
||||
# IP
|
||||
python examples/offline_inference/ming_flash_omni_tts/end2end.py --case ip
|
||||
|
||||
# Basic (speed/pitch/volume control)
|
||||
python examples/offline_inference/ming_flash_omni_tts/end2end.py --case basic
|
||||
```
|
||||
|
||||
## Useful Arguments
|
||||
|
||||
- `--text`: override default text in the selected case
|
||||
- `--output`: custom output wav path
|
||||
- `--model`: local model path or HF repo id
|
||||
- `--stage-configs-path`: custom talker stage config path
|
||||
- `--log-stats`: enable runtime stats logs
|
||||
|
||||
## Notes
|
||||
|
||||
- This directory is for **standalone talker deployment (TTS)**.
|
||||
- For Ming thinker multimodal understanding examples, see:
|
||||
`examples/offline_inference/ming_flash_omni/`.
|
||||
@@ -0,0 +1,128 @@
|
||||
"""Offline e2e example for Ming-flash-omni-2.0 standalone talker (TTS)."""
|
||||
|
||||
import os
|
||||
from typing import Any
|
||||
|
||||
import soundfile as sf
|
||||
import torch
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniTokensPrompt
|
||||
from vllm_omni.model_executor.models.ming_flash_omni.prompt_utils import (
|
||||
DEFAULT_PROMPT,
|
||||
create_instruction,
|
||||
)
|
||||
|
||||
MODEL_NAME = "Jonathan1909/Ming-flash-omni-2.0"
|
||||
DEFAULT_STAGE_CONFIG = "vllm_omni/model_executor/stage_configs/ming_flash_omni_tts.yaml"
|
||||
|
||||
|
||||
def get_messages(case: str, text_override: str | None) -> dict[str, Any]:
|
||||
if case == "style":
|
||||
text = text_override or "我会一直在这里陪着你,直到你慢慢、慢慢地沉入那个最温柔的梦里……好吗?"
|
||||
instruction = create_instruction(
|
||||
{
|
||||
"风格": "这是一种ASMR耳语,属于一种旨在引发特殊感官体验的创意风格。这个女性使用轻柔的普通话进行耳语,声音气音成分重。音量极低,紧贴麦克风,语速极慢,旨在制造触发听者颅内快感的声学刺激。",
|
||||
}
|
||||
)
|
||||
return {
|
||||
"prompt": DEFAULT_PROMPT,
|
||||
"text": text,
|
||||
"instruction": instruction,
|
||||
"use_zero_spk_emb": True,
|
||||
}
|
||||
if case == "ip":
|
||||
text = text_override or "这款产品的名字,叫变态坑爹牛肉丸。"
|
||||
return {
|
||||
"prompt": DEFAULT_PROMPT,
|
||||
"text": text,
|
||||
"instruction": create_instruction({"IP": "灵小甄"}),
|
||||
"use_zero_spk_emb": True,
|
||||
}
|
||||
if case == "basic":
|
||||
text = text_override or "我们当迎着阳光辛勤耕作,去摘取,去制作,去品尝,去馈赠。"
|
||||
return {
|
||||
"prompt": DEFAULT_PROMPT,
|
||||
"text": text,
|
||||
"instruction": create_instruction({"语速": "快速", "基频": "中", "音量": "中"}),
|
||||
"use_zero_spk_emb": True,
|
||||
}
|
||||
raise ValueError(f"Unknown case: {case}")
|
||||
|
||||
|
||||
def save_audio(mm: dict[str, Any], output_path: str) -> None:
|
||||
if not mm or "audio" not in mm:
|
||||
raise RuntimeError("No audio found in model output")
|
||||
audio = mm["audio"]
|
||||
sr_raw = mm.get("sr", 44100)
|
||||
if isinstance(sr_raw, torch.Tensor):
|
||||
sample_rate = int(sr_raw.item())
|
||||
else:
|
||||
sample_rate = int(sr_raw)
|
||||
waveform = audio.squeeze().float().cpu().numpy()
|
||||
sf.write(output_path, waveform, sample_rate)
|
||||
print(f"Saved {output_path} ({len(waveform) / sample_rate:.2f}s, {sample_rate}Hz)")
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = FlexibleArgumentParser(description="Ming-flash-omni standalone talker offline e2e example")
|
||||
parser.add_argument("--model", type=str, default=MODEL_NAME, help="Model name or local path.")
|
||||
parser.add_argument(
|
||||
"--stage-configs-path",
|
||||
type=str,
|
||||
default=DEFAULT_STAGE_CONFIG,
|
||||
help="Path to stage configs yaml for standalone talker deployment.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--case",
|
||||
type=str,
|
||||
default="style",
|
||||
choices=["style", "ip", "basic"],
|
||||
help="Example case.",
|
||||
)
|
||||
parser.add_argument("--text", type=str, default=None, help="Override default text for the selected case.")
|
||||
parser.add_argument("--output", type=str, default=None, help="Output wav path.")
|
||||
parser.add_argument("--log-stats", action="store_true", default=False, help="Enable stats logging.")
|
||||
parser.add_argument("--init-timeout", type=int, default=600, help="Engine init timeout in seconds.")
|
||||
parser.add_argument("--stage-init-timeout", type=int, default=300, help="Single stage init timeout in seconds.")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
|
||||
omni = Omni(
|
||||
model=args.model,
|
||||
stage_configs_path=args.stage_configs_path,
|
||||
log_stats=args.log_stats,
|
||||
init_timeout=args.init_timeout,
|
||||
stage_init_timeout=args.stage_init_timeout,
|
||||
)
|
||||
|
||||
messages = get_messages(args.case, args.text)
|
||||
decode_args = {
|
||||
# Standalone TTS deployment
|
||||
"ming_task": "instruct",
|
||||
"max_decode_steps": 200,
|
||||
"cfg": 2.0,
|
||||
"sigma": 0.25,
|
||||
"temperature": 0.0,
|
||||
}
|
||||
req = OmniTokensPrompt(
|
||||
prompt_token_ids=[0],
|
||||
additional_information={**messages, **decode_args},
|
||||
)
|
||||
|
||||
outputs = omni.generate(req)
|
||||
mm = outputs[0].outputs[0].multimodal_output
|
||||
|
||||
output_path = args.output or f"output_{args.case}.wav"
|
||||
save_audio(mm, output_path)
|
||||
omni.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -89,7 +89,6 @@ def run_e2e():
|
||||
omni = Omni(
|
||||
model=args.model,
|
||||
stage_configs_path=args.stage_config,
|
||||
trust_remote_code=True,
|
||||
log_stats=True,
|
||||
)
|
||||
|
||||
|
||||
Loaded 100 of 427 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user