Compare commits
119
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
b32c697b05 | ||
|
|
4a24a517ab | ||
|
|
5d4b16e6e3 | ||
|
|
3c85ca5536 | ||
|
|
c483a23deb | ||
|
|
2856ff7aea | ||
|
|
56ca5d9a8f | ||
|
|
81ab2f98da | ||
|
|
b25ea13cb0 | ||
|
|
687a44e5c8 | ||
|
|
19f8f42822 | ||
|
|
0a8204cb81 | ||
|
|
b8bd75837e | ||
|
|
576afb6f53 | ||
|
|
1e8dc84150 | ||
|
|
28558cc374 | ||
|
|
282e0b6642 | ||
|
|
1e5f288a91 | ||
|
|
e969d2e99f | ||
|
|
6f784cbc50 | ||
|
|
369a47d5a1 | ||
|
|
b076006c35 | ||
|
|
9f81a0a4b0 | ||
|
|
f36d891ed1 | ||
|
|
354511b805 | ||
|
|
005621ba7f | ||
|
|
4e41b7bc23 | ||
|
|
353ac8a402 | ||
|
|
e49fbd8a2d | ||
|
|
2b85474c23 | ||
|
|
5cf3f7947b | ||
|
|
44cde33eb5 | ||
|
|
a77c56725d | ||
|
|
fde38ca54f | ||
|
|
b64ab05d09 | ||
|
|
a0918ce583 | ||
|
|
f19891e4f1 | ||
|
|
d0b531520d | ||
|
|
1a93818e54 | ||
|
|
bb239fa949 | ||
|
|
703e31fc47 | ||
|
|
e1b30061b6 | ||
|
|
c007d40b14 | ||
|
|
c708aecb18 | ||
|
|
c306511512 | ||
|
|
5fc0bfe01e | ||
|
|
9a53703676 | ||
|
|
bb69cbc9f1 | ||
|
|
33586d845c | ||
|
|
3a26799507 | ||
|
|
6bb18af119 | ||
|
|
21a3035c10 | ||
|
|
e50a066548 | ||
|
|
4c06bb09ed | ||
|
|
5dabbb2f18 | ||
|
|
7491b67446 | ||
|
|
d49356e4ca | ||
|
|
6a837849dd | ||
|
|
5e82d7f311 | ||
|
|
368afea2f2 | ||
|
|
e403acda3f | ||
|
|
04bfbbae61 | ||
|
|
6509303a04 | ||
|
|
68a484589d | ||
|
|
5916a78551 | ||
|
|
38d440109b | ||
|
|
cc38bf195c | ||
|
|
01f500a5bf | ||
|
|
9de3806624 | ||
|
|
704b6bb6e6 | ||
|
|
2cd64cdccb | ||
|
|
01ebc0cdf5 | ||
|
|
77e1e5c01c | ||
|
|
e637726f94 | ||
|
|
bede8e02d8 | ||
|
|
da60663f78 | ||
|
|
0309b3bd19 | ||
|
|
aa0587c407 | ||
|
|
16310c8ef3 | ||
|
|
15a1143695 | ||
|
|
7fc528f7a9 | ||
|
|
fa945b2626 | ||
|
|
ac66282746 | ||
|
|
bfb730ed50 | ||
|
|
359e474db4 | ||
|
|
b1ff695029 | ||
|
|
b9789091e9 | ||
|
|
7f2b1e9386 | ||
|
|
dd0fa02547 | ||
|
|
08cf5a6083 | ||
|
|
4c2bea782e | ||
|
|
b50fe1d8b1 | ||
|
|
954d0642a4 | ||
|
|
c6c7fd9b6d | ||
|
|
49c9b89fac | ||
|
|
7a8b428bdb | ||
|
|
9de131e6b0 | ||
|
|
335e2eb9e4 | ||
|
|
1069a144a8 | ||
|
|
569899e8d9 | ||
|
|
c93359bb35 | ||
|
|
0470e96771 | ||
|
|
539fa57006 | ||
|
|
0016006974 | ||
|
|
239abcb480 | ||
|
|
d34093b360 | ||
|
|
5f53dccd15 | ||
|
|
00c11c8b87 | ||
|
|
a150194b11 | ||
|
|
52f17bf292 | ||
|
|
afa2b09684 | ||
|
|
62c86c9d62 | ||
|
|
8e2e7c74fb | ||
|
|
6bcf2c8894 | ||
|
|
9da06690ea | ||
|
|
e03ab82dd2 | ||
|
|
6e35bfbbae | ||
|
|
6b52db9e2b | ||
|
|
2390350533 |
No files matched your search
@@ -10,7 +10,7 @@ steps:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
# Buildkite will automatically replace this with the actual commit hash
|
||||
VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
|
||||
VLLM_VERSION: "v0.19.0"
|
||||
VLLM_VERSION: "v0.20.0"
|
||||
priority: 100
|
||||
timeout_in_minutes: 60
|
||||
soft_fail: true
|
||||
soft_fail: false
|
||||
+17
-4
@@ -1,8 +1,9 @@
|
||||
# Document 1: Buildkite loads only this block on first parse. The next step resolves docs-only skip-ci
|
||||
# from git diff, then uploads document 2. When docs-only skip applies, image-build still runs if nightly-test
|
||||
# / main NIGHTLY so upload-nightly is not skipped together with test-ready/test-merge.
|
||||
#
|
||||
# Document 2: appended after `---`; same file, read by upload_pipeline_with_skip_ci.sh (not evaluated as a second pipeline by Buildkite).
|
||||
# from git diff, then uploads document 2. When docs-only skip applies, image-build still runs for the nightly
|
||||
# exception in upload_pipeline_with_skip_ci.sh (nightly-test label or main NIGHTLY=1); otherwise upload-nightly
|
||||
# can be skipped together with test-ready/test-merge. Weekly uploads also depend on image-build.
|
||||
|
||||
# Document 2: appended after `---`; same file, read by upload_pipeline_with_skip_ci.sh (not evaluated as a second pipeline by Buildkite). Child uploads include test-ready.yml, test-merge.yml, test-nightly.yml, and test-weekly.yml.
|
||||
steps:
|
||||
- label: ":github: Resolve skip-ci & upload pipeline"
|
||||
key: upload-ci-pipeline
|
||||
@@ -61,3 +62,15 @@ steps:
|
||||
- buildkite-agent pipeline upload .buildkite/test-nightly.yml
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
# L5 Test — main+WEEKLY=1 (scheduled), or PR with weekly-test label
|
||||
- label: "Upload Weekly Pipeline"
|
||||
depends_on: image-build
|
||||
key: upload-weekly-pipeline
|
||||
if: >-
|
||||
(build.branch == "main" && build.env("WEEKLY") == "1") ||
|
||||
(build.branch != "main" && build.pull_request.labels includes "weekly-test")
|
||||
commands:
|
||||
- buildkite-agent pipeline upload .buildkite/test-weekly.yml
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
@@ -141,5 +141,5 @@ docker run \
|
||||
"${image_name}" \
|
||||
bash -c '
|
||||
set -e
|
||||
VLLM_USE_MODELSCOPE=True pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
|
||||
VLLM_USE_MODELSCOPE=True pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni_expansion.py
|
||||
'
|
||||
@@ -7,7 +7,9 @@ steps:
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_ROCM_USE_AITER=0
|
||||
- "timeout 20m pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
|
||||
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
|
||||
# that is not supported on CPU
|
||||
- "timeout 20m pytest -v -s -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
|
||||
|
||||
- label: "Diffusion Model Test"
|
||||
agent_pool: mi325_2
|
||||
@@ -103,16 +105,17 @@ steps:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 20m pytest -s -v tests/engine/test_async_omni_engine_abort.py
|
||||
|
||||
# This is from test-nightly.yml.
|
||||
# This model seems to be important, so I am trying to run it here.
|
||||
# Since this is from test-nightly.yml,
|
||||
# I keep this non-blocking.
|
||||
- label: "Omni Model Test Qwen2-5-Omni"
|
||||
agent_pool: mi325_2
|
||||
agent_pool: mi325_4
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 20m pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
|
||||
- timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen2_5_omni.py -m "advanced_model" --run-level "advanced_model"
|
||||
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
- pytest -s -v tests/e2e/*/test_qwen2_5_omni_expansion.py --run-level "core_model"
|
||||
|
||||
- label: "Omni Model Test Qwen3-Omni"
|
||||
agent_pool: mi325_2
|
||||
@@ -120,10 +123,9 @@ steps:
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
- timeout 30m pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py tests/e2e/online_serving/test_mimo_audio.py -m "advanced_model" --run-level "advanced_model"
|
||||
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py tests/e2e/online_serving/test_mimo_audio.py -m "advanced_model" --run-level "advanced_model"
|
||||
|
||||
- label: "Qwen3-TTS CustomVoice E2E Test"
|
||||
agent_pool: mi325_1
|
||||
|
||||
@@ -7,7 +7,9 @@ steps:
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_ROCM_USE_AITER=0
|
||||
- "timeout 20m pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
|
||||
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
|
||||
# that is not supported on CPU
|
||||
- "timeout 20m pytest -vvvv -s -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
|
||||
|
||||
- label: "Voxtral TTS CUDA Unit Test"
|
||||
agent_pool: mi325_1
|
||||
@@ -71,6 +73,17 @@ steps:
|
||||
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# - 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
|
||||
|
||||
- label: "AudioX Online Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 20m pytest -s -v tests/e2e/online_serving/test_audiox_online.py
|
||||
|
||||
- label: "Diffusion Cache Backend Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
@@ -112,19 +125,6 @@ steps:
|
||||
pytest -s -v tests/engine/test_async_omni_engine_abort.py
|
||||
'
|
||||
|
||||
|
||||
# NOTE: This test is not running any thing. It is skipped and deselected.
|
||||
# Currently it is = 1 skipped, 1 deselected, 17 warnings in 0.03s ======
|
||||
# - label: "Omni Model Test Qwen2-5-Omni"
|
||||
# agent_pool: mi325_2
|
||||
# depends_on: amd-build
|
||||
# mirror_hardwares: [amdproduction]
|
||||
# grade: Blocking
|
||||
# commands:
|
||||
# - export VLLM_LOGGING_LEVEL=DEBUG
|
||||
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# - timeout 20m pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py -m "core_model" --run-level "core_model"
|
||||
|
||||
# - label: "Omni Model Test Qwen3-Omni"
|
||||
# agent_pool: mi325_2
|
||||
# depends_on: amd-build
|
||||
|
||||
@@ -88,7 +88,7 @@ steps:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
|
||||
- pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni_expansion.py
|
||||
- pytest -s -v tests/engine/test_async_omni_engine_abort.py
|
||||
|
||||
- label: "Omni Model Test Qwen3-Omni"
|
||||
|
||||
+23
-20
@@ -172,25 +172,6 @@ steps:
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Omni Model Test"
|
||||
timeout_in_minutes: 30
|
||||
depends_on: upload-merge-pipeline
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py tests/e2e/online_serving/test_qwen2_5_omni.py -m "advanced_model" --run-level "advanced_model"
|
||||
agents:
|
||||
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Qwen3-TTS CustomVoice E2E Test"
|
||||
depends_on: upload-merge-pipeline
|
||||
commands:
|
||||
@@ -235,6 +216,28 @@ steps:
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "MOSS-TTS-Nano E2E Test"
|
||||
depends_on: upload-merge-pipeline
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
pytest -s -v tests/e2e/online_serving/test_moss_tts_nano.py tests/e2e/offline_inference/test_moss_tts_nano.py -m "advanced_model" --run-level "advanced_model"
|
||||
'
|
||||
agents:
|
||||
queue: "gpu_1_queue"
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Omni Model Test with H100"
|
||||
timeout_in_minutes: 30
|
||||
depends_on: upload-merge-pipeline
|
||||
@@ -363,7 +366,7 @@ steps:
|
||||
set -e
|
||||
export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
export VLLM_IMAGE_FETCH_TIMEOUT=60
|
||||
pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "advanced_model" --run-level "advanced_model" -k "shared_memory"
|
||||
pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "advanced_model" --run-level "advanced_model"
|
||||
pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py tests/e2e/online_serving/test_bagel_online.py -m "advanced_model" --run-level "advanced_model"
|
||||
'
|
||||
agents:
|
||||
|
||||
+198
-5
@@ -10,7 +10,7 @@ steps:
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "omni-test"
|
||||
steps:
|
||||
- label: ":full_moon: Omni · Function Test"
|
||||
- label: ":full_moon: Omni · Function Test with H100"
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/ -m "full_model and H100 and omni" --run-level "full_model" --ignore=tests/e2e/accuracy
|
||||
@@ -48,6 +48,25 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: Omni · Function Test with L4"
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
- pytest -s -v tests/e2e/*/test_qwen2_5_omni_expansion.py --run-level "core_model"
|
||||
agents:
|
||||
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: ":full_moon: Omni · Doc Test with L4"
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
@@ -196,6 +215,44 @@ steps:
|
||||
type: DirectoryOrCreate
|
||||
|
||||
|
||||
- label: ":full_moon: Omni · Multi-Replica Startup Test with 4x H100"
|
||||
timeout_in_minutes: 45
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py -m "core_model" --run-level "core_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: 4
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- group: ":card_index_dividers: TTS Model Test"
|
||||
key: nightly-tts-test-group
|
||||
depends_on: upload-nightly-pipeline
|
||||
@@ -205,7 +262,7 @@ steps:
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
- pytest -s -v tests/e2e/ -m "full_model and L4 and omni" --run-level "full_model" --ignore=tests/e2e/accuracy
|
||||
- pytest -s -v tests/e2e/ -m "full_model and L4 and omni" -k "not test_qwen2_5" --run-level "full_model" --ignore=tests/e2e/accuracy
|
||||
agents:
|
||||
queue: "gpu_1_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
|
||||
plugins:
|
||||
@@ -278,7 +335,10 @@ steps:
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100"
|
||||
timeout_in_minutes: 120
|
||||
commands:
|
||||
- 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
|
||||
- pytest -sv tests/e2e/ -k "not test_wan and not test_bagel_expansion and not hunyuan" -m "full_model and diffusion and H100" --run-level "full_model" --ignore=tests/e2e/accuracy
|
||||
# revoke after bagel optimization finished
|
||||
# Keep Bagel expansion and multi-replica tests in their dedicated H100 jobs below.
|
||||
# - pytest -sv tests/e2e/ -k "not test_wan and not test_bagel_expansion and not hunyuan" -m "full_model and diffusion and H100" --run-level "full_model" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_bagel_multi_replicas.py
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
@@ -313,10 +373,87 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Bagel Function Test with H100"
|
||||
timeout_in_minutes: 120
|
||||
commands:
|
||||
- pytest -sv tests/e2e/online_serving/test_bagel_expansion.py -m "full_model and diffusion 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: 3
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
# revoke after bagel optimization finished
|
||||
# - label: ":full_moon: Diffusion X2I(&A&T) · BAGEL Multi-Replica with H100"
|
||||
# timeout_in_minutes: 120
|
||||
# commands:
|
||||
# - pytest -s -v tests/e2e/online_serving/test_bagel_multi_replicas.py -m "full_model and diffusion 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: 4
|
||||
# volumeMounts:
|
||||
# - name: devshm
|
||||
# mountPath: /dev/shm
|
||||
# - name: hf-cache
|
||||
# mountPath: /root/.cache/huggingface
|
||||
# env:
|
||||
# - name: HF_HOME
|
||||
# value: /root/.cache/huggingface
|
||||
# - name: HF_TOKEN
|
||||
# valueFrom:
|
||||
# secretKeyRef:
|
||||
# name: hf-token-secret
|
||||
# key: token
|
||||
# nodeSelector:
|
||||
# node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
# volumes:
|
||||
# - name: devshm
|
||||
# emptyDir:
|
||||
# medium: Memory
|
||||
# - name: hf-cache
|
||||
# hostPath:
|
||||
# path: /mnt/hf-cache
|
||||
# type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with L4"
|
||||
timeout_in_minutes: 60
|
||||
commands:
|
||||
- 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
|
||||
- pytest -sv tests/e2e/ -k "not test_wan and not test_bagel_expansion 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:
|
||||
@@ -455,6 +592,7 @@ steps:
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/accuracy/test_qwen_image*.py --run-level full_model
|
||||
- pytest -s -v tests/e2e/accuracy/test_diffusers_backend_similarity.py -k '2i_matches_diffusers' -m full_model --run-level full_model
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
@@ -493,7 +631,7 @@ steps:
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test"
|
||||
key: nightly-diffusion-x2iat-performance
|
||||
timeout_in_minutes: 180
|
||||
timeout_in_minutes: 360
|
||||
commands:
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
|
||||
@@ -547,6 +685,60 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · HunyuanImage3 · DiT Perf Test"
|
||||
key: nightly-hunyuan-image3-performance
|
||||
soft_fail: true
|
||||
timeout_in_minutes: 120
|
||||
if: build.env("RUN_HUNYUAN_IMAGE3_PERF") == "1"
|
||||
commands:
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
|
||||
- export CACHE_DIT_VERSION=1.3.0
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_hunyuan_image_tp4_fp8.json
|
||||
EXIT1=$$?
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_hunyuan_image_tp2_fp8_sp2.json
|
||||
EXIT2=$$?
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_hunyuan_image_tp2_fp8_cfgp2.json
|
||||
EXIT3=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
|
||||
exit $$((EXIT1 | EXIT2 | EXIT3))
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 4
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
# Diffusion x2v only (Wan, HunyuanVideo, …). x2i/x2a/x2t live in the X2I group above, not here.
|
||||
- group: ":card_index_dividers: Diffusion X2V Model Test"
|
||||
key: nightly-diffusion-x2v-group
|
||||
@@ -599,6 +791,7 @@ steps:
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py -m full_model --run-level full_model
|
||||
- pytest -s -v tests/e2e/accuracy/test_ltx2_3_video_similarity.py -m advanced_model --run-level advanced_model
|
||||
- pytest -s -v tests/e2e/accuracy/test_diffusers_backend_similarity.py -k '2v_matches_diffusers' -m full_model --run-level full_model
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
|
||||
@@ -120,6 +120,23 @@ steps:
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "AudioX Online Test"
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
- timeout 20m pytest -s -v tests/e2e/online_serving/test_audiox_online.py -m "core_model and diffusion" --run-level core_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:
|
||||
|
||||
@@ -0,0 +1,85 @@
|
||||
env:
|
||||
VLLM_WORKER_MULTIPROC_METHOD: spawn
|
||||
HF_HUB_DOWNLOAD_TIMEOUT: 300
|
||||
HF_HUB_ETAG_TIMEOUT: 60
|
||||
|
||||
steps:
|
||||
- label: "Reliability Test - qwen3-omni"
|
||||
timeout_in_minutes: 180
|
||||
depends_on: upload-weekly-pipeline
|
||||
if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test"
|
||||
commands:
|
||||
- pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m "slow"
|
||||
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: "Reliability Test - wan22"
|
||||
timeout_in_minutes: 180
|
||||
depends_on: upload-weekly-pipeline
|
||||
if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test"
|
||||
commands:
|
||||
- pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m "slow"
|
||||
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
|
||||
@@ -544,6 +544,15 @@ Pipelines without `SupportsModuleOffload` fall back to scanning
|
||||
well-known attribute names (`transformer`, `text_encoder`, `vae`,
|
||||
etc.), which fails for non-standard names.
|
||||
|
||||
### Step 12: Performance Profiling
|
||||
|
||||
After verifying correctness and implementing parallelism/caching, profile the model's performance to identify bottlenecks and ensure optimal execution.
|
||||
|
||||
See the [Profiling Single-Stage Diffusion](../../../docs/contributing/profiling.md#3-profiling-single-stage-diffusion) guide for detailed instructions on:
|
||||
1. Using the PyTorch profiler (`profiler: "torch"`) to capture detailed CPU/CUDA traces.
|
||||
2. Using Nsight Systems (`nsys`) with `profiler: "cuda"` for low-overhead CUDA traces.
|
||||
3. Controlling profiling via `omni.start_profile()` and `omni.stop_profile()`.
|
||||
|
||||
---
|
||||
|
||||
## Iterative Development Tips
|
||||
|
||||
@@ -211,8 +211,8 @@ See `plan/voxcpm2_native_ar_design.md`.
|
||||
|
||||
- Model files in `vllm_omni/model_executor/models/<model_name>/`
|
||||
- Stage config YAML
|
||||
- Working `end2end.py` with correct audio output
|
||||
- README.md in the example directory
|
||||
- Working `end2end.py` at `examples/offline_inference/text_to_speech/<model>/end2end.py`
|
||||
- New section in `examples/offline_inference/text_to_speech/README.md` (table row + per-model section). Do **not** create a top-level `examples/offline_inference/<model>/` dir or a per-model `README.md` inside `text_to_speech/<model>/` — the hub README is the documented surface and the mkdocs `generate_examples` hook only descends one level into `examples/<category>/`.
|
||||
|
||||
## Phase 3: Online Serving
|
||||
|
||||
@@ -308,11 +308,11 @@ def build_voice_clone_prompt(ref_audio_path: str, text: str, codec) -> list:
|
||||
### Deliverables
|
||||
|
||||
- Updated `serving_speech.py` with all 5 integration points (single commit)
|
||||
- Client scripts and server launcher
|
||||
- Gradio demo with streaming and voice cloning UI
|
||||
- Client scripts and server launcher under `examples/online_serving/text_to_speech/<model>/`
|
||||
- Gradio demo with streaming and voice cloning UI in the same dir
|
||||
- E2E online serving test (`tests/e2e/online_serving/test_<model>.py`)
|
||||
- Buildkite CI entry in `.buildkite/test-merge.yml`
|
||||
- Documentation (offline + online serving docs)
|
||||
- New section in `examples/online_serving/text_to_speech/README.md` (table row + per-model section). Do **not** create a top-level `examples/online_serving/<model>/` dir or a per-model `README.md` inside `text_to_speech/<model>/`.
|
||||
|
||||
### E2E test pitfalls to avoid
|
||||
|
||||
|
||||
@@ -27,6 +27,7 @@
|
||||
*.mov binary
|
||||
*.mp4 binary
|
||||
*.mp3 binary
|
||||
*.wav binary
|
||||
*.flv binary
|
||||
*.fla binary
|
||||
*.swf binary
|
||||
|
||||
@@ -0,0 +1,223 @@
|
||||
pull_request_rules:
|
||||
- name: label-documentation
|
||||
description: Automatically apply documentation label
|
||||
conditions:
|
||||
- label != stale
|
||||
- or:
|
||||
- files~=^[^/]+\.md$
|
||||
- files~=^docs/
|
||||
- files~=^examples/
|
||||
actions:
|
||||
label:
|
||||
add:
|
||||
- documentation
|
||||
comment:
|
||||
message: "Documentation preview: https://vllm--{{number}}.org.readthedocs.build/en/{{number}}/"
|
||||
|
||||
- name: comment-pre-commit-failure
|
||||
description: Comment on PR when pre-commit check fails
|
||||
conditions:
|
||||
- status-failure=pre-commit
|
||||
- -closed
|
||||
- -draft
|
||||
actions:
|
||||
comment:
|
||||
message: |
|
||||
Hi @{{author}}, the pre-commit checks have failed. Please run:
|
||||
|
||||
```bash
|
||||
uv pip install pre-commit
|
||||
pre-commit install
|
||||
pre-commit run --all-files
|
||||
```
|
||||
|
||||
Then, commit the changes and push to your branch.
|
||||
|
||||
For future commits, `pre-commit` will run automatically on changed files before each commit.
|
||||
|
||||
> [!TIP]
|
||||
> <details>
|
||||
> <summary>Is <code>mypy</code> or <code>markdownlint</code> failing?</summary>
|
||||
> <br/>
|
||||
> <code>mypy</code> and <code>markdownlint</code> are run differently in CI. If the failure is related to either of these checks, please use the following commands to run them locally:
|
||||
>
|
||||
> ```bash
|
||||
> # For mypy (substitute "3.10" with the failing version if needed)
|
||||
> pre-commit run --hook-stage manual mypy-3.10
|
||||
> # For markdownlint
|
||||
> pre-commit run --hook-stage manual markdownlint
|
||||
> ```
|
||||
> </details>
|
||||
|
||||
- name: comment-dco-failure
|
||||
description: Comment on PR when DCO check fails
|
||||
conditions:
|
||||
- status-failure=dco
|
||||
- -closed
|
||||
- -draft
|
||||
actions:
|
||||
comment:
|
||||
message: |
|
||||
Hi @{{author}}, the DCO check has failed. Please click on DCO in the Checks section for instructions on how to resolve this.
|
||||
|
||||
- name: label-ci-build
|
||||
description: Automatically apply ci/build label
|
||||
conditions:
|
||||
- label != stale
|
||||
- or:
|
||||
- files~=^\.github/
|
||||
- files~=\.buildkite/
|
||||
- files~=^cmake/
|
||||
- files=CMakeLists.txt
|
||||
- files~=^docker/Dockerfile
|
||||
- files~=^requirements.*\.txt
|
||||
- files=setup.py
|
||||
actions:
|
||||
label:
|
||||
add:
|
||||
- ci/build
|
||||
|
||||
- name: label-frontend
|
||||
description: Automatically apply frontend label
|
||||
conditions:
|
||||
- label != stale
|
||||
- files~=^vllm_omni/entrypoints/
|
||||
actions:
|
||||
label:
|
||||
add:
|
||||
- frontend
|
||||
|
||||
- name: label-diffusion
|
||||
description: Automatically apply diffusion label
|
||||
conditions:
|
||||
- label != stale
|
||||
- files~=^vllm_omni/diffusion/
|
||||
actions:
|
||||
label:
|
||||
add:
|
||||
- diffusion
|
||||
|
||||
- name: label-new-model
|
||||
description: Automatically apply new-model label
|
||||
conditions:
|
||||
- label != stale
|
||||
- or:
|
||||
- files=vllm_omni/model_executor/models/registry.py
|
||||
- files=vllm_omni/diffusion/registry.py
|
||||
actions:
|
||||
label:
|
||||
add:
|
||||
- new-model
|
||||
|
||||
- name: label-npu
|
||||
description: Automatically apply npu label
|
||||
conditions:
|
||||
- label != stale
|
||||
- or:
|
||||
- files~=^vllm_omni/platforms/npu/
|
||||
- title~=(?i)NPU
|
||||
- title~=(?i)Ascend
|
||||
actions:
|
||||
label:
|
||||
add:
|
||||
- npu
|
||||
assign:
|
||||
users:
|
||||
- "gcanlin"
|
||||
|
||||
- name: label-rocm
|
||||
description: Automatically apply rocm label
|
||||
conditions:
|
||||
- label != stale
|
||||
- or:
|
||||
- files~=^vllm_omni/platforms/rocm/
|
||||
- title~=(?i)AMD
|
||||
- title~=(?i)ROCm
|
||||
actions:
|
||||
label:
|
||||
add:
|
||||
- rocm
|
||||
assign:
|
||||
users:
|
||||
- "tjtanaa"
|
||||
|
||||
- name: label-xpu
|
||||
description: Automatically apply xpu label
|
||||
conditions:
|
||||
- label != stale
|
||||
- or:
|
||||
- files~=^vllm_omni/platforms/xpu/
|
||||
- title~=(?i)Intel
|
||||
- title~=(?i)XPU
|
||||
actions:
|
||||
label:
|
||||
add:
|
||||
- xpu
|
||||
assign:
|
||||
users:
|
||||
- "yma11"
|
||||
|
||||
- name: label-modelrunner
|
||||
description: Automatically apply modelrunner label
|
||||
conditions:
|
||||
- label != stale
|
||||
- or:
|
||||
- files~=^vllm_omni/worker/
|
||||
- title~=(?i)ModelRunner
|
||||
- title~=(?i)worker
|
||||
actions:
|
||||
label:
|
||||
add:
|
||||
- ModelRunner
|
||||
assign:
|
||||
users:
|
||||
- "gcanlin" # To check whether need to sync to NPU
|
||||
|
||||
- name: auto-rebase if approved, ready, and 30 commits behind main
|
||||
conditions:
|
||||
- base = main
|
||||
- label=ready
|
||||
- "#approved-reviews-by >= 1"
|
||||
- "#commits-behind >= 30"
|
||||
- -closed
|
||||
- -draft
|
||||
- -conflict
|
||||
actions:
|
||||
rebase: {}
|
||||
|
||||
- name: ping author on conflicts and add 'needs-rebase' label
|
||||
conditions:
|
||||
- label != stale
|
||||
- conflict
|
||||
- -closed
|
||||
actions:
|
||||
label:
|
||||
add:
|
||||
- needs-rebase
|
||||
comment:
|
||||
message: |
|
||||
This pull request has merge conflicts that must be resolved before it can be
|
||||
merged. Please rebase the PR, @{{author}}.
|
||||
|
||||
https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork
|
||||
|
||||
- name: remove 'needs-rebase' label when conflict is resolved
|
||||
conditions:
|
||||
- -conflict
|
||||
- -closed
|
||||
actions:
|
||||
label:
|
||||
remove:
|
||||
- needs-rebase
|
||||
|
||||
- name: label-bug
|
||||
description: Automatically apply bug label
|
||||
conditions:
|
||||
- label != stale
|
||||
- or:
|
||||
- title~=(?i)\bbug\b
|
||||
- title~=(?i)\bbugfix\b
|
||||
actions:
|
||||
label:
|
||||
add:
|
||||
- bug
|
||||
@@ -264,6 +264,8 @@ tmp_test
|
||||
vllm_omni/_version.py
|
||||
# output files
|
||||
*.wav
|
||||
# Vendored test reference audio (small, intentionally checked in)
|
||||
!tests/assets/**/*.wav
|
||||
.worktrees/
|
||||
# CI overlay yamls materialized from tests/utils.py:_CI_OVERLAYS at test time
|
||||
tests/.ci_generated/
|
||||
@@ -32,6 +32,8 @@ python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
|
||||
- The benchmark talks to `http://<host>:<port>/v1/chat/completions`.
|
||||
- If you run the server on another host or port, pass `--base-url` accordingly.
|
||||
- To compare diffusion multi-replica serving, restart the server with different
|
||||
`runtime.num_replicas` values and keep the benchmark parameters unchanged.
|
||||
|
||||
## 2. Supported Datasets
|
||||
|
||||
@@ -93,6 +95,35 @@ You can point to your own trace using `--dataset-path`.
|
||||
|
||||
## 3. Benchmark Parameters
|
||||
|
||||
### Multi-replica sweep
|
||||
|
||||
For replica scaling, use the same model, prompt count, image/video size, and
|
||||
step count for every run. Change only the serving stage config:
|
||||
|
||||
```yaml
|
||||
runtime:
|
||||
devices: "0,1,2,3"
|
||||
num_replicas: 4
|
||||
```
|
||||
|
||||
Then run a fixed benchmark and write one JSON file per replica count:
|
||||
|
||||
```bash
|
||||
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
--base-url http://localhost:8099 \
|
||||
--model Qwen/Qwen-Image \
|
||||
--task t2i \
|
||||
--dataset random \
|
||||
--num-prompts 32 \
|
||||
--max-concurrency 8 \
|
||||
--request-rate inf \
|
||||
--num-inference-steps 10 \
|
||||
--width 512 --height 512 \
|
||||
--output-file diffusion_replica_4.json
|
||||
```
|
||||
|
||||
Repeat for `num_replicas=1`, `2`, `3`, and `4`, using matching `runtime.devices`.
|
||||
|
||||
### Basic flags
|
||||
|
||||
- `--base-url`: Server address (the script calls `.../v1/chat/completions`).
|
||||
@@ -130,9 +161,21 @@ Warmup flags:
|
||||
|
||||
- `--warmup-requests`: Number of warmup requests.
|
||||
- `--warmup-num-inference-steps`: Steps used during warmup.
|
||||
- `--warmup-concurrency`: Maximum concurrent warmup requests. Use this to warm
|
||||
the same batch shape as the measured run instead of warming only batch=`1`.
|
||||
- For `--task t2v`: warmup requests are forced to use `num_frames=1` to make warmup faster and less noisy.
|
||||
|
||||
Traffic / concurrency flags:
|
||||
|
||||
- `--request-rate`: Target request rate (requests/second). If set to `inf`, the script sends all requests immediately.
|
||||
- `--max-concurrency`: Max number of in-flight requests (default: `1`). This can hard-cap the achieved QPS: if it is too small, requests will queue behind the semaphore, and both achieved throughput and observed SLO attainment can be skewed.
|
||||
|
||||
### Batched warmup note
|
||||
|
||||
For batched serving runs, warm the same in-flight shape you plan to measure.
|
||||
For example, a run with `--max-concurrency 8` should usually also use
|
||||
`--warmup-requests 8 --warmup-concurrency 8`; otherwise the first measured
|
||||
batch may still pay compile or CUDA-graph capture cost.
|
||||
|
||||
For a Qwen-Image continuous-batching replay example, see
|
||||
[`performance_dashboard/qwen_image_serving_performance.md`](./performance_dashboard/qwen_image_serving_performance.md).
|
||||
@@ -751,6 +751,52 @@ async def iter_requests(
|
||||
yield req
|
||||
|
||||
|
||||
def _make_warmup_request(
|
||||
requests_list: list[RequestFuncInput],
|
||||
index: int,
|
||||
args,
|
||||
) -> RequestFuncInput:
|
||||
warm_req = requests_list[index % len(requests_list)]
|
||||
if args.warmup_num_inference_steps is not None:
|
||||
warm_req = replace(
|
||||
warm_req,
|
||||
num_inference_steps=args.warmup_num_inference_steps,
|
||||
)
|
||||
if args.task == "t2v":
|
||||
warm_req = replace(warm_req, num_frames=1)
|
||||
return warm_req
|
||||
|
||||
|
||||
async def _run_warmups(
|
||||
requests_list: list[RequestFuncInput],
|
||||
args,
|
||||
session: aiohttp.ClientSession,
|
||||
request_func,
|
||||
) -> list[tuple[RequestFuncInput, RequestFuncOutput]]:
|
||||
if not args.warmup_requests or not requests_list:
|
||||
return []
|
||||
|
||||
warmup_requests = [_make_warmup_request(requests_list, i, args) for i in range(args.warmup_requests)]
|
||||
warmup_concurrency = min(int(args.warmup_concurrency), len(warmup_requests))
|
||||
warmup_semaphore = asyncio.Semaphore(warmup_concurrency)
|
||||
|
||||
print(
|
||||
f"Running {len(warmup_requests)} warmup request(s) "
|
||||
f"with num_inference_steps={args.warmup_num_inference_steps} "
|
||||
f"and warmup_concurrency={warmup_concurrency}..."
|
||||
)
|
||||
|
||||
async def limited_warmup_request_func(
|
||||
req: RequestFuncInput,
|
||||
) -> RequestFuncOutput:
|
||||
async with warmup_semaphore:
|
||||
return await request_func(req, session, None)
|
||||
|
||||
warmup_tasks = [asyncio.create_task(limited_warmup_request_func(req)) for req in warmup_requests]
|
||||
warmup_outputs = await asyncio.gather(*warmup_tasks)
|
||||
return list(zip(warmup_requests, warmup_outputs))
|
||||
|
||||
|
||||
def calculate_metrics(
|
||||
outputs: list[RequestFuncOutput],
|
||||
total_duration: float,
|
||||
@@ -900,23 +946,12 @@ async def benchmark(args):
|
||||
pbar = tqdm(total=len(requests_list), disable=args.disable_tqdm)
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
warmup_pairs: list[tuple[RequestFuncInput, RequestFuncOutput]] = []
|
||||
if args.warmup_requests and requests_list:
|
||||
print(
|
||||
f"Running {args.warmup_requests} warmup request(s) \
|
||||
with num_inference_steps={args.warmup_num_inference_steps}..."
|
||||
)
|
||||
for i in range(args.warmup_requests):
|
||||
warm_req = requests_list[i % len(requests_list)]
|
||||
if args.warmup_num_inference_steps is not None:
|
||||
warm_req = replace(
|
||||
warm_req,
|
||||
num_inference_steps=args.warmup_num_inference_steps,
|
||||
)
|
||||
if args.task == "t2v":
|
||||
warm_req = replace(warm_req, num_frames=1)
|
||||
warm_out = await limited_request_func(warm_req, session, None)
|
||||
warmup_pairs.append((warm_req, warm_out))
|
||||
warmup_pairs = await _run_warmups(
|
||||
requests_list=requests_list,
|
||||
args=args,
|
||||
session=session,
|
||||
request_func=request_func,
|
||||
)
|
||||
|
||||
if args.slo:
|
||||
# Prefer trace-provided per-request slo_ms. Only populate when missing.
|
||||
@@ -1072,6 +1107,13 @@ if __name__ == "__main__":
|
||||
default=1,
|
||||
help="num_inference_steps used for warmup requests.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--warmup-concurrency",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Maximum number of warmup requests to run concurrently. "
|
||||
"Set this to match the real batch shape when warming up torch.compile or CUDA graphs.",
|
||||
)
|
||||
parser.add_argument("--width", type=int, default=None, help="Image/Video width.")
|
||||
parser.add_argument("--height", type=int, default=None, help="Image/Video height.")
|
||||
parser.add_argument("--num-frames", type=int, default=None, help="Number of frames (for video).")
|
||||
@@ -1135,5 +1177,4 @@ if __name__ == "__main__":
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
asyncio.run(benchmark(args))
|
||||
@@ -33,6 +33,12 @@ vllm serve Qwen/Qwen-Image --omni \
|
||||
--port 8091
|
||||
```
|
||||
|
||||
To replay step-wise continuous batching, compare
|
||||
`vllm serve Qwen/Qwen-Image --omni --port 8089 --step-execution --max-num-seqs 1`
|
||||
against the same command with `--max-num-seqs 8`. `--step-execution` is the
|
||||
feature gate, and increasing `--max-num-seqs` above `1` lets the scheduler
|
||||
keep more compatible requests active at once.
|
||||
|
||||
## 3.2 Key Parameters
|
||||
|
||||
| Parameter | Description |
|
||||
@@ -70,6 +76,7 @@ python benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
| `--task` | Task type (e.g., `t2i`) |
|
||||
| `--num-prompts` | Total number of requests |
|
||||
| `--max-concurrency` | Client-side concurrency |
|
||||
| `--warmup-concurrency` | Warmup concurrency for batch-shape prewarming |
|
||||
| `--random-request-config`| JSON string defining random request |
|
||||
|
||||
---
|
||||
@@ -133,6 +140,12 @@ python benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
]'
|
||||
```
|
||||
|
||||
For a continuous-batching A/B replay, run the same benchmark command twice
|
||||
against the two serving configs above and keep traffic plus warmup settings
|
||||
identical across both runs. If the measured run uses `--max-concurrency 8`,
|
||||
warm with `--warmup-requests 8 --warmup-concurrency 8` so the first measured
|
||||
batch does not include compile or CUDA-graph capture overhead.
|
||||
|
||||
---
|
||||
|
||||
# 6. Performance Metrics
|
||||
@@ -163,6 +176,7 @@ To ensure consistent and comparable benchmark results:
|
||||
* Record parallel configuration
|
||||
* Record benchmark parameters (resolution, concurrency, number of prompts)
|
||||
* Ensure no background workload on GPUs during testing
|
||||
* When comparing continuous batching, keep warmup settings the same across runs
|
||||
|
||||
---
|
||||
|
||||
|
||||
+2
-2
@@ -3,11 +3,11 @@
|
||||
Measures TTFP improvement from DAC-code caching when using uploaded voices.
|
||||
|
||||
Setup:
|
||||
1. Start vllm-omni with Fish Speech S2 Pro (use our feat branch)
|
||||
1. Start vllm-omni with Fish Speech S2 Pro
|
||||
2. Provide a reference audio file for voice cloning
|
||||
|
||||
Usage:
|
||||
python bench_voice_cache.py \
|
||||
python bench_speaker_cache.py \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--ref-text "Transcript of the reference audio." \
|
||||
--num-prompts 20 \
|
||||
@@ -0,0 +1,157 @@
|
||||
# GLM-Image Benchmarks
|
||||
|
||||
Benchmark GLM-Image T2I (text-to-image) and I2I (image-to-image) performance across three backends: HuggingFace baseline, vLLM-Omni offline, and vLLM-Omni online serving.
|
||||
|
||||
## Benchmarks
|
||||
|
||||
| Benchmark | Script | Description |
|
||||
|-----------|--------|-------------|
|
||||
| HuggingFace Baseline | `huggingface/inference.py` | Single-GPU transformers + diffusers pipeline |
|
||||
| vLLM-Omni Offline | `vllm-omni/inference.py` | Offline inference with continuous batching |
|
||||
| vLLM-Omni Online | `benchmark_glm_image.py` | Online serving via `/v1/chat/completions` |
|
||||
|
||||
## HuggingFace Baseline
|
||||
|
||||
Single-request sequential inference using the reference HuggingFace pipeline.
|
||||
|
||||
```bash
|
||||
# T2I
|
||||
CUDA_VISIBLE_DEVICES=0 python benchmarks/glm_image/huggingface/inference.py \
|
||||
--model-path /path/to/GLM-Image --mode t2i --num-prompts 10
|
||||
|
||||
# I2I
|
||||
CUDA_VISIBLE_DEVICES=0 python benchmarks/glm_image/huggingface/inference.py \
|
||||
--model-path /path/to/GLM-Image --mode i2i --num-prompts 10
|
||||
```
|
||||
|
||||
### Options
|
||||
|
||||
| Flag | Default | Description |
|
||||
|------|---------|-------------|
|
||||
| `--model-path` | `zai-org/GLM-Image` | Model path |
|
||||
| `--mode` | `t2i` | `t2i` or `i2i` |
|
||||
| `--dataset-path` | `prompt/prompt.json` | Path to prompt.json |
|
||||
| `--num-prompts` | `10` | Number of images to generate |
|
||||
| `--width` / `--height` | `1024` | Output image size |
|
||||
| `--num-inference-steps` | `50` | Diffusion denoising steps |
|
||||
| `--output-dir` | `benchmarks/glm_image/huggingface/outputs` | Output directory |
|
||||
| `--output-file` | - | JSON file for metrics |
|
||||
|
||||
## vLLM-Omni Offline
|
||||
|
||||
Multi-GPU offline inference with pipeline parallelism and continuous batching.
|
||||
|
||||
```bash
|
||||
# T2I
|
||||
CUDA_VISIBLE_DEVICES=0,1 python benchmarks/glm_image/vllm-omni/inference.py \
|
||||
--model-path /path/to/GLM-Image --mode t2i --num-prompts 10
|
||||
|
||||
# I2I
|
||||
CUDA_VISIBLE_DEVICES=0,1 python benchmarks/glm_image/vllm-omni/inference.py \
|
||||
--model-path /path/to/GLM-Image --mode i2i --num-prompts 10
|
||||
```
|
||||
|
||||
### Options
|
||||
|
||||
| Flag | Default | Description |
|
||||
|------|---------|-------------|
|
||||
| `--model-path` | `zai-org/GLM-Image` | Model path |
|
||||
| `--deploy-config` | - | Deploy config YAML |
|
||||
| `--mode` | `t2i` | `t2i` or `i2i` |
|
||||
| `--dataset-path` | `prompt/prompt.json` | Path to prompt.json |
|
||||
| `--num-prompts` | `10` | Number of images to generate |
|
||||
| `--width` / `--height` | `1024` | Output image size |
|
||||
| `--num-inference-steps` | `50` | Diffusion denoising steps |
|
||||
| `--output-dir` | `benchmarks/glm_image/vllm-omni/outputs` | Output directory |
|
||||
| `--output-file` | - | JSON file for metrics |
|
||||
| `--stage-init-timeout` | `600` | Stage initialization timeout (s) |
|
||||
|
||||
### Latency Computation
|
||||
|
||||
In offline mode all requests are submitted simultaneously and processed with continuous batching. The per-request latency is computed by summing the actual per-stage times (with `stage_0_gen_ms` diffed against the previous request to remove accumulated queue/scheduling wait).
|
||||
|
||||
## vLLM-Omni Online Serving
|
||||
|
||||
### Start the server
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=0,1 vllm serve /path/to/GLM-Image \
|
||||
--omni --port 8091 --host 0.0.0.0 \
|
||||
--served-model-name glm-image
|
||||
```
|
||||
|
||||
### Run the benchmark
|
||||
|
||||
```bash
|
||||
# T2I
|
||||
python benchmarks/glm_image/benchmark_glm_image.py \
|
||||
--mode t2i --num-prompts 10 --model glm-image
|
||||
|
||||
# I2I
|
||||
python benchmarks/glm_image/benchmark_glm_image.py \
|
||||
--mode i2i --num-prompts 10 --model glm-image
|
||||
|
||||
# Custom dataset
|
||||
python benchmarks/glm_image/benchmark_glm_image.py \
|
||||
--mode i2i --dataset custom \
|
||||
--dataset-path prompts.json --num-prompts 5
|
||||
```
|
||||
|
||||
### Options
|
||||
|
||||
| Flag | Default | Description |
|
||||
|------|---------|-------------|
|
||||
| `--mode` | `t2i` | `t2i` or `i2i` |
|
||||
| `--dataset` | `prompt` | `prompt`, `random`, or `custom` |
|
||||
| `--dataset-path` | - | JSON file path (required for `custom`) |
|
||||
| `--num-prompts` | `10` | Number of benchmark requests |
|
||||
| `--max-concurrency` | `1` | Max concurrent requests |
|
||||
| `--request-rate` | `inf` | Requests per second (Poisson arrival) |
|
||||
| `--warmup-requests` | `1` | Warmup requests before measurement |
|
||||
| `--width` / `--height` | `1024` | Output image size |
|
||||
| `--num-inference-steps` | `50` | Diffusion denoising steps |
|
||||
| `--seed` | - | Random seed |
|
||||
| `--model` | `default` | Model name (must match `--served-model-name`) |
|
||||
| `--host` | `localhost` | Server host |
|
||||
| `--port` | `8091` | Server port |
|
||||
| `--output-file` | - | JSON output file for metrics |
|
||||
| `--num-input-images` | `1` | Number of input images for random I2I |
|
||||
|
||||
## Dataset
|
||||
|
||||
The default dataset is hosted on [HuggingFace](https://huggingface.co/datasets/JaredforReal/glm-image-bench) (`prompt.json`). It is automatically downloaded and cached to `prompt/prompt.json` on first run. No manual setup needed.
|
||||
|
||||
Each entry contains:
|
||||
|
||||
- `t2i_prompt`: Text prompt for text-to-image generation
|
||||
- `i2i_prompt`: Text prompt for image-to-image editing
|
||||
- `image_url`: Source image URL for I2I (downloaded and cached on first use)
|
||||
|
||||
Custom datasets use the same JSON format and can be provided via `--dataset-path`.
|
||||
|
||||
## Pipeline Timings
|
||||
|
||||
All three benchmarks report per-stage pipeline timings (in milliseconds):
|
||||
|
||||
| Key | Description |
|
||||
|-----|-------------|
|
||||
| `preprocess_ms` | Input preprocessing (tokenization, multimodal encoding) |
|
||||
| `stage_0_gen_ms` | AR (autoregressive) model generation time |
|
||||
| `ar2diffusion_ms` | AR output to diffusion input conversion |
|
||||
| `stage_1_gen_ms` | Diffusion model denoising time |
|
||||
| `queue_wait_ms` | Queue wait time before processing |
|
||||
|
||||
The stages are ordered by execution: `preprocess → stage_0 (AR) → ar2diffusion → stage_1 (Diffusion)`.
|
||||
|
||||
## Sample Results
|
||||
|
||||
Tested on 2x GPU with 10 prompts, 1024x1024, 50 denoising steps:
|
||||
|
||||
| Backend | Mode | Latency Mean (s) | Throughput (img/s) |
|
||||
|---------|------|-------------------|--------------------|
|
||||
| HuggingFace | T2I | 72.6 | 0.014 |
|
||||
| HuggingFace | I2I | 70.9 | 0.014 |
|
||||
| vLLM-Omni Offline | T2I | 35.0 | 0.044 |
|
||||
| vLLM-Omni Offline | I2I | 31.0 | 0.053 |
|
||||
| vLLM-Omni Online | T2I | 38.8 | 0.026 |
|
||||
| vLLM-Omni Online | I2I | 34.7 | 0.029 |
|
||||
Whitespace-only changes.
@@ -0,0 +1,464 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Online serving benchmark for GLM-Image (T2I and I2I modes).
|
||||
|
||||
Sends requests to the /v1/chat/completions endpoint and reports end-to-end
|
||||
latency, throughput, and per-stage durations (when the server is started with
|
||||
--enable-diffusion-pipeline-profiler and/or --enable-ar-profiler).
|
||||
|
||||
Supports three dataset types:
|
||||
- prompt: Use prompt.json (default). T2I uses t2i_prompt, I2I uses i2i_prompt
|
||||
and sends source images from image_url.
|
||||
- random: Generate synthetic prompts (and random images for I2I).
|
||||
- custom: Load from a user-specified JSON file.
|
||||
|
||||
Usage:
|
||||
# T2I with prompt.json (default)
|
||||
python benchmarks/glm_image/benchmark_glm_image.py \
|
||||
--mode t2i --num-prompts 10
|
||||
|
||||
# I2I with prompt.json (downloads source images automatically)
|
||||
python benchmarks/glm_image/benchmark_glm_image.py \
|
||||
--mode i2i --num-prompts 10
|
||||
|
||||
# Random dataset
|
||||
python benchmarks/glm_image/benchmark_glm_image.py \
|
||||
--mode t2i --dataset random --num-prompts 20
|
||||
|
||||
# Custom dataset
|
||||
python benchmarks/glm_image/benchmark_glm_image.py \
|
||||
--mode i2i --dataset custom \
|
||||
--dataset-path my_prompts.json --num-prompts 5
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import aiohttp
|
||||
import numpy as np
|
||||
import requests as sync_requests
|
||||
from PIL import Image
|
||||
from tqdm.asyncio import tqdm
|
||||
|
||||
# Import backends from the diffusion benchmark (add parent dirs to path)
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "diffusion"))
|
||||
from backends import RequestFuncOutput
|
||||
|
||||
BENCHMARK_DIR = Path(__file__).resolve().parent
|
||||
DEFAULT_PROMPT_JSON = BENCHMARK_DIR / "prompt" / "prompt.json"
|
||||
IMAGE_CACHE_DIR = BENCHMARK_DIR / "prompt" / "images"
|
||||
|
||||
DATASET_REPO = "JaredforReal/glm-image-bench"
|
||||
DATASET_FILE = "prompt.json"
|
||||
|
||||
|
||||
def _ensure_prompt_json(dataset_path: str | None) -> str:
|
||||
"""Return path to prompt.json, downloading from HuggingFace if needed."""
|
||||
if dataset_path:
|
||||
return dataset_path
|
||||
local = DEFAULT_PROMPT_JSON
|
||||
if local.exists():
|
||||
return str(local)
|
||||
print(f"Downloading {DATASET_FILE} from {DATASET_REPO} ...")
|
||||
try:
|
||||
from huggingface_hub import hf_hub_download
|
||||
|
||||
downloaded = hf_hub_download(
|
||||
repo_id=DATASET_REPO,
|
||||
filename=DATASET_FILE,
|
||||
repo_type="dataset",
|
||||
)
|
||||
local.parent.mkdir(parents=True, exist_ok=True)
|
||||
import shutil
|
||||
|
||||
shutil.copy2(downloaded, local)
|
||||
print(f"Saved to {local}")
|
||||
except ImportError:
|
||||
url = f"https://huggingface.co/datasets/{DATASET_REPO}/resolve/main/{DATASET_FILE}"
|
||||
import urllib.request
|
||||
|
||||
local.parent.mkdir(parents=True, exist_ok=True)
|
||||
urllib.request.urlretrieve(url, local)
|
||||
print(f"Saved to {local}")
|
||||
return str(local)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class GLMImageRequest:
|
||||
prompt: str
|
||||
image_path: str | None = None # Only for I2I mode
|
||||
|
||||
|
||||
def download_image(url: str) -> str:
|
||||
"""Download an image to cache and return the local path."""
|
||||
IMAGE_CACHE_DIR.mkdir(parents=True, exist_ok=True)
|
||||
fname = url.rsplit("/", 1)[-1]
|
||||
local_path = IMAGE_CACHE_DIR / fname
|
||||
if local_path.exists():
|
||||
return str(local_path)
|
||||
resp = sync_requests.get(url, timeout=30)
|
||||
resp.raise_for_status()
|
||||
local_path.write_bytes(resp.content)
|
||||
return str(local_path)
|
||||
|
||||
|
||||
def encode_image_as_data_url(path: str) -> str:
|
||||
"""Encode a local image file as a base64 data URL."""
|
||||
with open(path, "rb") as f:
|
||||
encoded = base64.b64encode(f.read()).decode("utf-8")
|
||||
ext = Path(path).suffix.lower()
|
||||
mime = {"png": "image/png", ".jpg": "image/jpeg", ".jpeg": "image/jpeg"}.get(ext, "image/png")
|
||||
return f"data:{mime};base64,{encoded}"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Datasets
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class PromptDataset:
|
||||
"""Load from prompt.json. T2I uses t2i_prompt, I2I uses i2i_prompt + image_url."""
|
||||
|
||||
def __init__(self, args: argparse.Namespace):
|
||||
path = _ensure_prompt_json(args.dataset_path)
|
||||
with open(path, encoding="utf-8") as f:
|
||||
raw = json.load(f)
|
||||
|
||||
prompt_key = "t2i_prompt" if args.mode == "t2i" else "i2i_prompt"
|
||||
self.items: list[GLMImageRequest] = []
|
||||
|
||||
for entry in raw:
|
||||
prompt = entry.get(prompt_key, "").strip()
|
||||
if not prompt:
|
||||
continue
|
||||
image_path = None
|
||||
if args.mode == "i2i":
|
||||
url = entry.get("image_url", "")
|
||||
if url:
|
||||
image_path = download_image(url)
|
||||
self.items.append(GLMImageRequest(prompt=prompt, image_path=image_path))
|
||||
|
||||
if args.num_prompts and len(self.items) > args.num_prompts:
|
||||
self.items = self.items[: args.num_prompts]
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.items)
|
||||
|
||||
def __getitem__(self, idx: int) -> GLMImageRequest:
|
||||
return self.items[idx]
|
||||
|
||||
def get_requests(self) -> list[GLMImageRequest]:
|
||||
return list(self.items)
|
||||
|
||||
|
||||
class RandomDataset:
|
||||
"""Generate synthetic prompts (and optional random images for I2I)."""
|
||||
|
||||
def __init__(self, args: argparse.Namespace):
|
||||
self.args = args
|
||||
self.num_prompts = args.num_prompts
|
||||
self._random_image_paths: list[str] | None = None
|
||||
if args.mode == "i2i":
|
||||
self._random_image_paths = self._generate_random_images()
|
||||
|
||||
def _generate_random_images(self) -> list[str]:
|
||||
paths: list[str] = []
|
||||
for i in range(self.args.num_input_images):
|
||||
img = Image.new("RGB", (512, 512), (128 + i * 30 % 128, 64, 192))
|
||||
path = os.path.join(tempfile.gettempdir(), f"glm_image_bench_input_{i}.png")
|
||||
img.save(path)
|
||||
paths.append(path)
|
||||
return paths
|
||||
|
||||
def __len__(self) -> int:
|
||||
return self.num_prompts
|
||||
|
||||
def __getitem__(self, idx: int) -> GLMImageRequest:
|
||||
image_path = None
|
||||
if self._random_image_paths is not None:
|
||||
image_path = self._random_image_paths[idx % len(self._random_image_paths)]
|
||||
return GLMImageRequest(
|
||||
prompt=f"A beautiful scene with vivid colors and intricate details, prompt {idx}",
|
||||
image_path=image_path,
|
||||
)
|
||||
|
||||
def get_requests(self) -> list[GLMImageRequest]:
|
||||
return [self[i] for i in range(len(self))]
|
||||
|
||||
|
||||
class CustomDataset:
|
||||
"""Load from a user-specified JSON file.
|
||||
|
||||
Expected format:
|
||||
[
|
||||
{"prompt": "A cat sitting on a windowsill"},
|
||||
{"prompt": "Make it look like winter", "image_path": "/path/to/img.png"}
|
||||
]
|
||||
"""
|
||||
|
||||
def __init__(self, args: argparse.Namespace):
|
||||
if not args.dataset_path:
|
||||
raise ValueError("--dataset-path is required for custom dataset")
|
||||
with open(args.dataset_path, encoding="utf-8") as f:
|
||||
raw = json.load(f)
|
||||
self.items: list[GLMImageRequest] = []
|
||||
for item in raw:
|
||||
self.items.append(
|
||||
GLMImageRequest(
|
||||
prompt=item.get("prompt", ""),
|
||||
image_path=item.get("image_path"),
|
||||
)
|
||||
)
|
||||
if args.num_prompts and len(self.items) > args.num_prompts:
|
||||
self.items = self.items[: args.num_prompts]
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.items)
|
||||
|
||||
def __getitem__(self, idx: int) -> GLMImageRequest:
|
||||
return self.items[idx]
|
||||
|
||||
def get_requests(self) -> list[GLMImageRequest]:
|
||||
return list(self.items)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Async request for GLM-Image (chat completions with image support)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def async_glm_image_request(
|
||||
req: GLMImageRequest,
|
||||
api_url: str,
|
||||
model: str,
|
||||
session: aiohttp.ClientSession,
|
||||
pbar: Any,
|
||||
args: argparse.Namespace,
|
||||
) -> RequestFuncOutput:
|
||||
"""Send a single T2I or I2I request via chat completions endpoint."""
|
||||
output = RequestFuncOutput()
|
||||
output.start_time = time.perf_counter()
|
||||
|
||||
# Build messages
|
||||
if req.image_path and args.mode == "i2i":
|
||||
data_url = encode_image_as_data_url(req.image_path)
|
||||
content = [
|
||||
{"type": "text", "text": req.prompt},
|
||||
{"type": "image_url", "image_url": {"url": data_url}},
|
||||
]
|
||||
else:
|
||||
content = req.prompt
|
||||
|
||||
messages = [{"role": "user", "content": content}]
|
||||
|
||||
extra_body: dict[str, Any] = {}
|
||||
if args.height:
|
||||
extra_body["height"] = args.height
|
||||
if args.width:
|
||||
extra_body["width"] = args.width
|
||||
if args.num_inference_steps:
|
||||
extra_body["num_inference_steps"] = args.num_inference_steps
|
||||
if args.seed is not None:
|
||||
extra_body["seed"] = args.seed
|
||||
|
||||
payload: dict[str, Any] = {
|
||||
"model": model,
|
||||
"messages": messages,
|
||||
}
|
||||
if extra_body:
|
||||
payload["extra_body"] = extra_body
|
||||
|
||||
try:
|
||||
async with session.post(api_url, json=payload) as response:
|
||||
if response.status == 200:
|
||||
resp_json = await response.json()
|
||||
output.response_body = resp_json
|
||||
output.success = True
|
||||
try:
|
||||
choices = resp_json.get("choices", [])
|
||||
if choices and isinstance(choices, list):
|
||||
msg = choices[0].get("message", {})
|
||||
if isinstance(msg, dict):
|
||||
resp_content = msg.get("content", [])
|
||||
if resp_content and isinstance(resp_content, list) and len(resp_content) > 0:
|
||||
first_item = resp_content[0]
|
||||
if isinstance(first_item, dict):
|
||||
output.stage_durations = first_item.get("stage_durations") or {}
|
||||
output.peak_memory_mb = first_item.get("peak_memory_mb", 0.0)
|
||||
except (IndexError, TypeError, AttributeError):
|
||||
pass
|
||||
else:
|
||||
output.error = f"HTTP {response.status}: {await response.text()}"
|
||||
output.success = False
|
||||
except Exception as e:
|
||||
output.error = str(e)
|
||||
output.success = False
|
||||
|
||||
output.latency = time.perf_counter() - output.start_time
|
||||
if pbar:
|
||||
pbar.update(1)
|
||||
return output
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Benchmark
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def iter_requests(n: int, request_rate: float) -> Any:
|
||||
import random as _random
|
||||
|
||||
for i in range(n):
|
||||
if request_rate != float("inf") and i > 0:
|
||||
await asyncio.sleep(_random.expovariate(request_rate))
|
||||
yield i
|
||||
|
||||
|
||||
def calculate_metrics(outputs: list[RequestFuncOutput], total_duration: float) -> dict[str, Any]:
|
||||
success = [o for o in outputs if o.success]
|
||||
errors = [o for o in outputs if not o.success]
|
||||
latencies = [o.latency for o in success]
|
||||
peak_mems = [o.peak_memory_mb for o in success if o.peak_memory_mb > 0]
|
||||
|
||||
stage_duration_lists: dict[str, list[float]] = {}
|
||||
for o in success:
|
||||
for stage, dur in (o.stage_durations or {}).items():
|
||||
stage_duration_lists.setdefault(stage, []).append(dur)
|
||||
|
||||
return {
|
||||
"duration": total_duration,
|
||||
"completed_requests": len(success),
|
||||
"failed_requests": len(errors),
|
||||
"throughput_qps": len(success) / total_duration if total_duration > 0 else 0,
|
||||
"latency_mean": float(np.mean(latencies)) if latencies else 0,
|
||||
"latency_median": float(np.median(latencies)) if latencies else 0,
|
||||
"latency_p99": float(np.percentile(latencies, 99)) if latencies else 0,
|
||||
"latency_p95": float(np.percentile(latencies, 95)) if latencies else 0,
|
||||
"peak_memory_mb_max": max(peak_mems) if peak_mems else 0,
|
||||
"stage_durations_mean": {s: float(np.mean(v)) for s, v in stage_duration_lists.items()},
|
||||
"stage_durations_p50": {s: float(np.percentile(v, 50)) for s, v in stage_duration_lists.items()},
|
||||
}
|
||||
|
||||
|
||||
async def benchmark(args: argparse.Namespace) -> None:
|
||||
api_url = f"http://{args.host}:{args.port}/v1/chat/completions"
|
||||
|
||||
# Load dataset
|
||||
if args.dataset == "prompt":
|
||||
dataset = PromptDataset(args)
|
||||
elif args.dataset == "random":
|
||||
dataset = RandomDataset(args)
|
||||
elif args.dataset == "custom":
|
||||
dataset = CustomDataset(args)
|
||||
else:
|
||||
raise ValueError(f"Unknown dataset: {args.dataset}")
|
||||
|
||||
glm_requests = dataset.get_requests()
|
||||
print(f"Prepared {len(glm_requests)} requests (mode={args.mode}, dataset={args.dataset})")
|
||||
|
||||
semaphore = asyncio.Semaphore(args.max_concurrency) if args.max_concurrency else None
|
||||
|
||||
async def limited_request(idx: int, req: GLMImageRequest, session: aiohttp.ClientSession, pbar: Any):
|
||||
if semaphore:
|
||||
async with semaphore:
|
||||
return await async_glm_image_request(req, api_url, args.model, session, pbar, args)
|
||||
return await async_glm_image_request(req, api_url, args.model, session, pbar, args)
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
# Warmup
|
||||
if args.warmup_requests and glm_requests:
|
||||
print(f"Running {args.warmup_requests} warmup request(s)...")
|
||||
for i in range(args.warmup_requests):
|
||||
await limited_request(i, glm_requests[i % len(glm_requests)], session, None)
|
||||
|
||||
# Main benchmark
|
||||
pbar = tqdm(total=len(glm_requests), disable=args.disable_tqdm)
|
||||
start_time = time.perf_counter()
|
||||
tasks = []
|
||||
async for idx in iter_requests(len(glm_requests), args.request_rate):
|
||||
tasks.append(asyncio.create_task(limited_request(idx, glm_requests[idx], session, pbar)))
|
||||
outputs = await asyncio.gather(*tasks)
|
||||
total_duration = time.perf_counter() - start_time
|
||||
pbar.close()
|
||||
|
||||
# Metrics
|
||||
metrics = calculate_metrics(outputs, total_duration)
|
||||
metrics["mode"] = args.mode
|
||||
metrics["model"] = args.model
|
||||
metrics["dataset"] = args.dataset
|
||||
|
||||
print(f"\n{' GLM-Image Online Benchmark Result ':=^60}")
|
||||
print(f"{'Mode:':<40} {args.mode}")
|
||||
print(f"{'Model:':<40} {args.model}")
|
||||
print(f"{'Dataset:':<40} {args.dataset}")
|
||||
print("-" * 50)
|
||||
print(f"{'Benchmark duration (s):':<40} {metrics['duration']:.2f}")
|
||||
print(f"{'Request rate:':<40} {args.request_rate}")
|
||||
print(f"{'Max concurrency:':<40} {args.max_concurrency}")
|
||||
print(f"{'Successful requests:':<40} {metrics['completed_requests']}/{len(glm_requests)}")
|
||||
print("-" * 50)
|
||||
print(f"{'Throughput (req/s):':<40} {metrics['throughput_qps']:.2f}")
|
||||
print(f"{'Latency Mean (s):':<40} {metrics['latency_mean']:.4f}")
|
||||
print(f"{'Latency Median (s):':<40} {metrics['latency_median']:.4f}")
|
||||
print(f"{'Latency P95 (s):':<40} {metrics['latency_p95']:.4f}")
|
||||
print(f"{'Latency P99 (s):':<40} {metrics['latency_p99']:.4f}")
|
||||
|
||||
if metrics["peak_memory_mb_max"] > 0:
|
||||
print("-" * 50)
|
||||
print(f"{'Peak Memory Max (MB):':<40} {metrics['peak_memory_mb_max']:.2f}")
|
||||
|
||||
if metrics["stage_durations_mean"]:
|
||||
print("-" * 50)
|
||||
print("Stage Durations Mean:")
|
||||
for stage, val in sorted(metrics["stage_durations_mean"].items()):
|
||||
unit = "ms" if stage.endswith("_ms") else "s"
|
||||
print(f" {stage + ':':<38} {val:.4f} ({unit})")
|
||||
|
||||
print("=" * 60)
|
||||
|
||||
if args.output_file:
|
||||
with open(args.output_file, "w") as f:
|
||||
json.dump(metrics, f, indent=2)
|
||||
print(f"Metrics saved to {args.output_file}")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Benchmark GLM-Image T2I/I2I online serving.")
|
||||
parser.add_argument("--mode", type=str, default="t2i", choices=["t2i", "i2i"])
|
||||
parser.add_argument("--dataset", type=str, default="prompt", choices=["prompt", "random", "custom"])
|
||||
parser.add_argument("--dataset-path", type=str, default=None)
|
||||
parser.add_argument("--num-prompts", type=int, default=10)
|
||||
parser.add_argument("--max-concurrency", type=int, default=1)
|
||||
parser.add_argument("--request-rate", type=float, default=float("inf"))
|
||||
parser.add_argument("--warmup-requests", type=int, default=1)
|
||||
parser.add_argument("--width", type=int, default=1024)
|
||||
parser.add_argument("--height", type=int, default=1024)
|
||||
parser.add_argument("--num-inference-steps", type=int, default=50)
|
||||
parser.add_argument("--seed", type=int, default=None)
|
||||
parser.add_argument("--model", type=str, default="default")
|
||||
parser.add_argument("--host", type=str, default="localhost")
|
||||
parser.add_argument("--port", type=int, default=8091)
|
||||
parser.add_argument("--output-file", type=str, default=None)
|
||||
parser.add_argument("--disable-tqdm", action="store_true")
|
||||
parser.add_argument("--num-input-images", type=int, default=1, help="For random I2I dataset.")
|
||||
args = parser.parse_args()
|
||||
asyncio.run(benchmark(args))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,291 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
HuggingFace (transformers + diffusers) baseline benchmark for GLM-Image.
|
||||
|
||||
Supports T2I and I2I modes with the prompt.json dataset.
|
||||
Downloads source images for I2I from image_url on first run and caches locally.
|
||||
|
||||
Usage:
|
||||
# T2I mode (text-to-image, no source images needed)
|
||||
python benchmarks/glm_image/huggingface/inference.py \
|
||||
--model-path zai-org/GLM-Image \
|
||||
--mode t2i --num-prompts 10
|
||||
|
||||
# I2I mode (image-to-image, downloads source images)
|
||||
python benchmarks/glm_image/huggingface/inference.py \
|
||||
--model-path zai-org/GLM-Image \
|
||||
--mode i2i --num-prompts 10
|
||||
|
||||
# With custom prompt.json
|
||||
python benchmarks/glm_image/huggingface/inference.py \
|
||||
--model-path zai-org/GLM-Image \
|
||||
--mode i2i --dataset-path prompts.json --num-prompts 5
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import requests
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
BENCHMARK_DIR = Path(__file__).resolve().parent.parent
|
||||
DEFAULT_PROMPT_JSON = BENCHMARK_DIR / "prompt" / "prompt.json"
|
||||
IMAGE_CACHE_DIR = BENCHMARK_DIR / "prompt" / "images"
|
||||
|
||||
DATASET_REPO = "JaredforReal/glm-image-bench"
|
||||
DATASET_FILE = "prompt.json"
|
||||
|
||||
|
||||
def _ensure_prompt_json(dataset_path: str | None) -> str:
|
||||
"""Return path to prompt.json, downloading from HuggingFace if needed."""
|
||||
if dataset_path:
|
||||
return dataset_path
|
||||
local = DEFAULT_PROMPT_JSON
|
||||
if local.exists():
|
||||
return str(local)
|
||||
print(f"Downloading {DATASET_FILE} from {DATASET_REPO} ...")
|
||||
try:
|
||||
from huggingface_hub import hf_hub_download
|
||||
|
||||
downloaded = hf_hub_download(
|
||||
repo_id=DATASET_REPO,
|
||||
filename=DATASET_FILE,
|
||||
repo_type="dataset",
|
||||
)
|
||||
local.parent.mkdir(parents=True, exist_ok=True)
|
||||
import shutil
|
||||
|
||||
shutil.copy2(downloaded, local)
|
||||
print(f"Saved to {local}")
|
||||
except ImportError:
|
||||
url = f"https://huggingface.co/datasets/{DATASET_REPO}/resolve/main/{DATASET_FILE}"
|
||||
import urllib.request
|
||||
|
||||
local.parent.mkdir(parents=True, exist_ok=True)
|
||||
urllib.request.urlretrieve(url, local)
|
||||
print(f"Saved to {local}")
|
||||
return str(local)
|
||||
|
||||
|
||||
HEIGHT = 1024
|
||||
WIDTH = 1024
|
||||
SEED = 42
|
||||
NUM_INFERENCE_STEPS = 50
|
||||
GUIDANCE_SCALE = 1.5
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Dataset
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def load_dataset(
|
||||
dataset_path: str | None,
|
||||
mode: str,
|
||||
num_prompts: int,
|
||||
) -> list[dict]:
|
||||
"""Load prompts from prompt.json and prepare per-request data."""
|
||||
path = _ensure_prompt_json(dataset_path)
|
||||
with open(path, encoding="utf-8") as f:
|
||||
raw = json.load(f)
|
||||
|
||||
items = []
|
||||
for entry in raw:
|
||||
if mode == "t2i":
|
||||
prompt_key = "t2i_prompt"
|
||||
else:
|
||||
prompt_key = "i2i_prompt"
|
||||
|
||||
prompt_text = entry.get(prompt_key, "").strip()
|
||||
if not prompt_text:
|
||||
continue
|
||||
|
||||
item = {"prompt": prompt_text}
|
||||
if mode == "i2i":
|
||||
item["image_url"] = entry.get("image_url", "")
|
||||
items.append(item)
|
||||
|
||||
if num_prompts and len(items) > num_prompts:
|
||||
items = items[:num_prompts]
|
||||
return items
|
||||
|
||||
|
||||
def download_image(url: str, cache_dir: Path) -> str:
|
||||
"""Download an image to cache_dir and return the local path."""
|
||||
cache_dir.mkdir(parents=True, exist_ok=True)
|
||||
fname = url.rsplit("/", 1)[-1]
|
||||
local_path = cache_dir / fname
|
||||
if local_path.exists():
|
||||
return str(local_path)
|
||||
print(f" Downloading {url} ...")
|
||||
resp = requests.get(url, timeout=30)
|
||||
resp.raise_for_status()
|
||||
local_path.write_bytes(resp.content)
|
||||
return str(local_path)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Benchmark
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def benchmark(args: argparse.Namespace) -> None:
|
||||
from diffusers.pipelines.glm_image import GlmImagePipeline
|
||||
|
||||
print("=" * 60)
|
||||
print("GLM-Image HuggingFace Baseline Benchmark")
|
||||
print(f"Mode: {args.mode} | Model: {args.model_path}")
|
||||
print(f"Size: {args.height}x{args.width} | Steps: {args.num_inference_steps}")
|
||||
print("=" * 60)
|
||||
|
||||
# Load dataset
|
||||
items = load_dataset(args.dataset_path, args.mode, args.num_prompts)
|
||||
if not items:
|
||||
print("No prompts loaded. Exiting.")
|
||||
return
|
||||
print(f"Loaded {len(items)} prompts for {args.mode} mode")
|
||||
|
||||
# Download I2I source images
|
||||
if args.mode == "i2i":
|
||||
print("Preparing source images...")
|
||||
for item in items:
|
||||
url = item.get("image_url", "")
|
||||
if url:
|
||||
item["image_path"] = download_image(url, IMAGE_CACHE_DIR)
|
||||
else:
|
||||
item["image_path"] = None
|
||||
|
||||
# Load pipeline
|
||||
print(f"\nLoading pipeline from {args.model_path} ...")
|
||||
t0 = time.perf_counter()
|
||||
pipe = GlmImagePipeline.from_pretrained(
|
||||
args.model_path,
|
||||
torch_dtype=torch.bfloat16,
|
||||
device_map="cuda",
|
||||
)
|
||||
init_time = time.perf_counter() - t0
|
||||
print(f"Pipeline loaded in {init_time:.2f}s")
|
||||
|
||||
# Create output dir
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
|
||||
# Run benchmark
|
||||
generator = torch.Generator(device="cuda").manual_seed(args.seed)
|
||||
latencies = []
|
||||
success = 0
|
||||
failed = 0
|
||||
|
||||
print(f"\nRunning {len(items)} requests sequentially...")
|
||||
print("-" * 60)
|
||||
|
||||
for i, item in enumerate(items):
|
||||
prompt = item["prompt"]
|
||||
gen_kwargs: dict = {
|
||||
"prompt": prompt,
|
||||
"height": args.height,
|
||||
"width": args.width,
|
||||
"num_inference_steps": args.num_inference_steps,
|
||||
"guidance_scale": args.guidance_scale,
|
||||
"generator": generator,
|
||||
}
|
||||
|
||||
if args.mode == "i2i":
|
||||
img_path = item.get("image_path")
|
||||
if img_path and os.path.exists(img_path):
|
||||
gen_kwargs["image"] = [Image.open(img_path).convert("RGB")]
|
||||
else:
|
||||
print(f" [{i + 1}] SKIP: no source image")
|
||||
failed += 1
|
||||
continue
|
||||
|
||||
t_start = time.perf_counter()
|
||||
try:
|
||||
result = pipe(**gen_kwargs)
|
||||
image = result.images[0]
|
||||
elapsed = time.perf_counter() - t_start
|
||||
latencies.append(elapsed)
|
||||
success += 1
|
||||
|
||||
out_path = os.path.join(args.output_dir, f"{i:04d}.png")
|
||||
image.save(out_path)
|
||||
print(f" [{i + 1}/{len(items)}] {elapsed:.3f}s -> {out_path}")
|
||||
except Exception as e:
|
||||
elapsed = time.perf_counter() - t_start
|
||||
failed += 1
|
||||
print(f" [{i + 1}/{len(items)}] FAILED ({elapsed:.3f}s): {e}")
|
||||
|
||||
# Report
|
||||
total_gen_time = sum(latencies) if latencies else 0
|
||||
print("\n" + "=" * 60)
|
||||
print("HuggingFace Baseline Results")
|
||||
print("=" * 60)
|
||||
print(f"{'Mode:':<40} {args.mode}")
|
||||
print(f"{'Model:':<40} {args.model_path}")
|
||||
print(f"{'Image size:':<40} {args.height}x{args.width}")
|
||||
print(f"{'Num inference steps:':<40} {args.num_inference_steps}")
|
||||
print("-" * 50)
|
||||
print(f"{'Pipeline init time (s):':<40} {init_time:.2f}")
|
||||
print(f"{'Successful:':<40} {success}/{len(items)}")
|
||||
print(f"{'Failed:':<40} {failed}")
|
||||
print("-" * 50)
|
||||
if latencies:
|
||||
arr = np.array(latencies)
|
||||
print(f"{'Total generation time (s):':<40} {total_gen_time:.2f}")
|
||||
print(f"{'Throughput (img/s):':<40} {success / total_gen_time:.4f}")
|
||||
print(f"{'Latency Mean (s):':<40} {arr.mean():.4f}")
|
||||
print(f"{'Latency Median (s):':<40} {np.median(arr):.4f}")
|
||||
print(f"{'Latency P95 (s):':<40} {np.percentile(arr, 95):.4f}")
|
||||
print(f"{'Latency P99 (s):':<40} {np.percentile(arr, 99):.4f}")
|
||||
|
||||
print(f"\n{'Output dir:':<40} {args.output_dir}")
|
||||
print("=" * 60)
|
||||
|
||||
# Save metrics JSON
|
||||
metrics = {
|
||||
"backend": "huggingface",
|
||||
"mode": args.mode,
|
||||
"model": args.model_path,
|
||||
"height": args.height,
|
||||
"width": args.width,
|
||||
"num_inference_steps": args.num_inference_steps,
|
||||
"init_time_s": init_time,
|
||||
"completed_requests": success,
|
||||
"failed_requests": failed,
|
||||
"total_gen_time_s": total_gen_time,
|
||||
"throughput_qps": success / total_gen_time if total_gen_time > 0 else 0,
|
||||
"latency_mean": float(np.mean(latencies)) if latencies else 0,
|
||||
"latency_median": float(np.median(latencies)) if latencies else 0,
|
||||
"latency_p95": float(np.percentile(latencies, 95)) if latencies else 0,
|
||||
"latency_p99": float(np.percentile(latencies, 99)) if latencies else 0,
|
||||
}
|
||||
if args.output_file:
|
||||
with open(args.output_file, "w") as f:
|
||||
json.dump(metrics, f, indent=2)
|
||||
print(f"Metrics saved to {args.output_file}")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="GLM-Image HuggingFace baseline benchmark")
|
||||
parser.add_argument("--model-path", type=str, default="zai-org/GLM-Image")
|
||||
parser.add_argument("--mode", type=str, default="t2i", choices=["t2i", "i2i"])
|
||||
parser.add_argument("--dataset-path", type=str, default=None, help="Path to prompt.json")
|
||||
parser.add_argument("--num-prompts", type=int, default=10)
|
||||
parser.add_argument("--height", type=int, default=HEIGHT)
|
||||
parser.add_argument("--width", type=int, default=WIDTH)
|
||||
parser.add_argument("--num-inference-steps", type=int, default=NUM_INFERENCE_STEPS)
|
||||
parser.add_argument("--guidance-scale", type=float, default=GUIDANCE_SCALE)
|
||||
parser.add_argument("--seed", type=int, default=SEED)
|
||||
parser.add_argument("--output-dir", type=str, default="benchmarks/glm_image/huggingface/outputs")
|
||||
parser.add_argument("--output-file", type=str, default=None, help="JSON file for metrics")
|
||||
args = parser.parse_args()
|
||||
benchmark(args)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,505 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
vLLM-Omni offline benchmark for GLM-Image.
|
||||
|
||||
Supports T2I and I2I modes with the prompt.json dataset.
|
||||
Downloads source images for I2I from image_url on first run and caches locally.
|
||||
|
||||
Usage:
|
||||
# T2I mode
|
||||
python benchmarks/glm_image/vllm-omni/inference.py \
|
||||
--model-path zai-org/GLM-Image \
|
||||
--mode t2i --num-prompts 10
|
||||
|
||||
# I2I mode (downloads source images)
|
||||
python benchmarks/glm_image/vllm-omni/inference.py \
|
||||
--model-path zai-org/GLM-Image \
|
||||
--mode i2i --num-prompts 10
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import requests
|
||||
from PIL import Image
|
||||
from vllm import SamplingParams
|
||||
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
BENCHMARK_DIR = Path(__file__).resolve().parent.parent
|
||||
DEFAULT_PROMPT_JSON = BENCHMARK_DIR / "prompt" / "prompt.json"
|
||||
IMAGE_CACHE_DIR = BENCHMARK_DIR / "prompt" / "images"
|
||||
DEFAULT_DEPLOY_CONFIG = "vllm_omni/deploy/glm_image.yaml"
|
||||
|
||||
DATASET_REPO = "JaredforReal/glm-image-bench"
|
||||
DATASET_FILE = "prompt.json"
|
||||
|
||||
|
||||
def _ensure_prompt_json(dataset_path: str | None) -> str:
|
||||
"""Return path to prompt.json, downloading from HuggingFace if needed."""
|
||||
if dataset_path:
|
||||
return dataset_path
|
||||
local = DEFAULT_PROMPT_JSON
|
||||
if local.exists():
|
||||
return str(local)
|
||||
print(f"Downloading {DATASET_FILE} from {DATASET_REPO} ...")
|
||||
try:
|
||||
from huggingface_hub import hf_hub_download
|
||||
|
||||
downloaded = hf_hub_download(
|
||||
repo_id=DATASET_REPO,
|
||||
filename=DATASET_FILE,
|
||||
repo_type="dataset",
|
||||
)
|
||||
local.parent.mkdir(parents=True, exist_ok=True)
|
||||
import shutil
|
||||
|
||||
shutil.copy2(downloaded, local)
|
||||
print(f"Saved to {local}")
|
||||
except ImportError:
|
||||
url = f"https://huggingface.co/datasets/{DATASET_REPO}/resolve/main/{DATASET_FILE}"
|
||||
import urllib.request
|
||||
|
||||
local.parent.mkdir(parents=True, exist_ok=True)
|
||||
urllib.request.urlretrieve(url, local)
|
||||
print(f"Saved to {local}")
|
||||
return str(local)
|
||||
|
||||
|
||||
SEED = 42
|
||||
HEIGHT = 1024
|
||||
WIDTH = 1024
|
||||
NUM_INFERENCE_STEPS = 50
|
||||
GUIDANCE_SCALE = 1.5
|
||||
|
||||
GLM_IMAGE_EOS_TOKEN_ID = 16385
|
||||
GLM_IMAGE_VISION_VOCAB_SIZE = 16512
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Dataset
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def load_dataset(
|
||||
dataset_path: str | None,
|
||||
mode: str,
|
||||
num_prompts: int,
|
||||
) -> list[dict]:
|
||||
path = _ensure_prompt_json(dataset_path)
|
||||
with open(path, encoding="utf-8") as f:
|
||||
raw = json.load(f)
|
||||
|
||||
items = []
|
||||
for entry in raw:
|
||||
prompt_key = "t2i_prompt" if mode == "t2i" else "i2i_prompt"
|
||||
prompt_text = entry.get(prompt_key, "").strip()
|
||||
if not prompt_text:
|
||||
continue
|
||||
|
||||
item = {"prompt": prompt_text}
|
||||
if mode == "i2i":
|
||||
item["image_url"] = entry.get("image_url", "")
|
||||
items.append(item)
|
||||
|
||||
if num_prompts and len(items) > num_prompts:
|
||||
items = items[:num_prompts]
|
||||
return items
|
||||
|
||||
|
||||
def download_image(url: str, cache_dir: Path) -> str:
|
||||
cache_dir.mkdir(parents=True, exist_ok=True)
|
||||
fname = url.rsplit("/", 1)[-1]
|
||||
local_path = cache_dir / fname
|
||||
if local_path.exists():
|
||||
return str(local_path)
|
||||
print(f" Downloading {url} ...")
|
||||
resp = requests.get(url, timeout=30)
|
||||
resp.raise_for_status()
|
||||
local_path.write_bytes(resp.content)
|
||||
return str(local_path)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def compute_max_tokens(height: int, width: int, is_i2i: bool = False) -> int:
|
||||
factor = 32
|
||||
token_h = height // factor
|
||||
token_w = width // factor
|
||||
large_tokens = token_h * token_w
|
||||
|
||||
# Small preview tokens (half resolution in each dimension)
|
||||
|
||||
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
|
||||
|
||||
# Mode-dependent totals:
|
||||
# - t2i: small + large + EOS
|
||||
# - i2i: large + EOS
|
||||
if is_i2i:
|
||||
return large_tokens + 1
|
||||
return small_tokens + large_tokens + 1
|
||||
|
||||
|
||||
def build_prompt_t2i(prompt: str, height: int, width: int, **gen_kw) -> dict:
|
||||
return {
|
||||
"prompt": prompt,
|
||||
"height": height,
|
||||
"width": width,
|
||||
"mm_processor_kwargs": {"target_h": height, "target_w": width},
|
||||
**gen_kw,
|
||||
}
|
||||
|
||||
|
||||
def build_prompt_i2i(prompt: str, image_path: str, height: int, width: int, **gen_kw) -> dict:
|
||||
return {
|
||||
"prompt": prompt,
|
||||
"height": height,
|
||||
"width": width,
|
||||
"mm_processor_kwargs": {"target_h": height, "target_w": width},
|
||||
"multi_modal_data": {"image": Image.open(image_path).convert("RGB")},
|
||||
**gen_kw,
|
||||
}
|
||||
|
||||
|
||||
def resolve_deploy_config(args: argparse.Namespace) -> str:
|
||||
if args.deploy_config:
|
||||
return args.deploy_config
|
||||
if os.path.exists(DEFAULT_DEPLOY_CONFIG):
|
||||
return DEFAULT_DEPLOY_CONFIG
|
||||
fallback = Path(__file__).resolve().parents[3] / DEFAULT_DEPLOY_CONFIG
|
||||
if fallback.exists():
|
||||
return str(fallback)
|
||||
raise FileNotFoundError("Deploy config not found. Specify --deploy-config.")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Benchmark
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def benchmark(args: argparse.Namespace) -> None:
|
||||
is_i2i = args.mode == "i2i"
|
||||
|
||||
print("=" * 60)
|
||||
print("GLM-Image vLLM-Omni Benchmark")
|
||||
print(f"Mode: {args.mode} | Model: {args.model_path}")
|
||||
print(f"Size: {args.height}x{args.width} | Steps: {args.num_inference_steps}")
|
||||
print("=" * 60)
|
||||
|
||||
# Load dataset
|
||||
items = load_dataset(args.dataset_path, args.mode, args.num_prompts)
|
||||
if not items:
|
||||
print("No prompts loaded. Exiting.")
|
||||
return
|
||||
print(f"Loaded {len(items)} prompts for {args.mode} mode")
|
||||
|
||||
# Download I2I source images
|
||||
if is_i2i:
|
||||
print("Preparing source images...")
|
||||
for item in items:
|
||||
url = item.get("image_url", "")
|
||||
if url:
|
||||
item["image_path"] = download_image(url, IMAGE_CACHE_DIR)
|
||||
else:
|
||||
item["image_path"] = None
|
||||
|
||||
# Init Omni
|
||||
deploy_config = resolve_deploy_config(args)
|
||||
print(f"\nInitializing vLLM-Omni (deploy config: {deploy_config}) ...")
|
||||
t0 = time.perf_counter()
|
||||
|
||||
omni = Omni(
|
||||
model=args.model_path,
|
||||
deploy_config=deploy_config,
|
||||
log_stats=args.log_stats,
|
||||
stage_init_timeout=args.stage_init_timeout,
|
||||
enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
|
||||
enable_ar_profiler=args.enable_ar_profiler,
|
||||
)
|
||||
|
||||
init_time = time.perf_counter() - t0
|
||||
print(f"Initialized in {init_time:.2f}s")
|
||||
|
||||
# Sampling params
|
||||
max_tokens = compute_max_tokens(args.height, args.width, is_i2i=is_i2i)
|
||||
ar_params = SamplingParams(
|
||||
temperature=0.9,
|
||||
top_p=0.75,
|
||||
top_k=GLM_IMAGE_VISION_VOCAB_SIZE,
|
||||
max_tokens=max_tokens,
|
||||
stop_token_ids=[GLM_IMAGE_EOS_TOKEN_ID],
|
||||
seed=args.seed,
|
||||
detokenize=False,
|
||||
extra_args={"target_h": args.height, "target_w": args.width},
|
||||
)
|
||||
diff_params = OmniDiffusionSamplingParams(
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
guidance_scale=args.guidance_scale,
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
seed=args.seed,
|
||||
)
|
||||
sampling_params_list = [ar_params, diff_params]
|
||||
|
||||
# Build all prompts
|
||||
gen_kw = {
|
||||
"seed": args.seed,
|
||||
"num_inference_steps": args.num_inference_steps,
|
||||
"guidance_scale": args.guidance_scale,
|
||||
}
|
||||
all_prompts = []
|
||||
for item in items:
|
||||
if is_i2i:
|
||||
img_path = item.get("image_path")
|
||||
if not img_path or not os.path.exists(img_path):
|
||||
continue
|
||||
all_prompts.append(build_prompt_i2i(item["prompt"], img_path, args.height, args.width, **gen_kw))
|
||||
else:
|
||||
all_prompts.append(build_prompt_t2i(item["prompt"], args.height, args.width, **gen_kw))
|
||||
|
||||
valid = len(all_prompts)
|
||||
print(f"Valid prompts: {valid}")
|
||||
|
||||
# Create output dir
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
|
||||
# Warmup: run 1 request to prime caches, CUDA graphs, etc.
|
||||
if all_prompts:
|
||||
print("Running warmup request...")
|
||||
try:
|
||||
warmup_prompt = [all_prompts[0]]
|
||||
omni.generate(warmup_prompt, sampling_params_list, py_generator=False)
|
||||
print("Warmup done.\n")
|
||||
except Exception as e:
|
||||
print(f"Warmup failed (continuing): {e}")
|
||||
|
||||
# Run
|
||||
print(f"\nRunning {valid} requests...")
|
||||
print("-" * 60)
|
||||
|
||||
latencies = []
|
||||
all_stage_durations: list[dict[str, float]] = []
|
||||
success = 0
|
||||
failed = 0
|
||||
wall_start = time.perf_counter()
|
||||
|
||||
try:
|
||||
output_idx = 0
|
||||
for stage_outputs in omni.generate(all_prompts, sampling_params_list, py_generator=True):
|
||||
if stage_outputs.final_output_type == "image":
|
||||
request_output = stage_outputs.request_output
|
||||
request_id = getattr(request_output, "request_id", "")
|
||||
|
||||
images = getattr(request_output, "images", [])
|
||||
if not images and hasattr(request_output, "multimodal_output"):
|
||||
mm = request_output.multimodal_output
|
||||
if isinstance(mm, dict):
|
||||
images = mm.get("images", [])
|
||||
|
||||
elapsed = time.perf_counter() - wall_start
|
||||
if images:
|
||||
for img in images:
|
||||
if isinstance(img, Image.Image):
|
||||
out_path = os.path.join(args.output_dir, f"{output_idx:04d}.png")
|
||||
img.save(out_path)
|
||||
success += 1
|
||||
latencies.append(elapsed)
|
||||
stage_durations = getattr(stage_outputs, "stage_durations", {})
|
||||
if stage_durations:
|
||||
all_stage_durations.append(stage_durations)
|
||||
# Show wall-clock elapsed and pipeline breakdown if available
|
||||
preprocess_str = ""
|
||||
if "preprocess_ms" in stage_durations:
|
||||
preprocess_str = f" preprocess={stage_durations['preprocess_ms'] / 1000.0:.2f}s"
|
||||
print(f" [{success}/{valid}] id={request_id[:8]} {elapsed:.2f}s{preprocess_str}")
|
||||
output_idx += 1
|
||||
else:
|
||||
failed += 1
|
||||
except Exception as e:
|
||||
print(f"Error: {e}")
|
||||
failed = valid - success
|
||||
|
||||
total_gen_time = time.perf_counter() - wall_start
|
||||
|
||||
# Diff stage_0_gen_ms with previous request to remove accumulated wait time.
|
||||
# stage_0_gen_ms is measured from submit_ts (same for all requests submitted
|
||||
# at once), so it accumulates queue/scheduling overhead across requests.
|
||||
# Other stages and pipeline timings are per-request already.
|
||||
_TIMING_ORDER = [
|
||||
"preprocess_ms",
|
||||
"stage_0_gen_ms",
|
||||
"ar2diffusion_ms",
|
||||
"stage_1_gen_ms",
|
||||
"queue_wait_ms",
|
||||
]
|
||||
|
||||
per_request_actual: list[dict[str, float]] = []
|
||||
prev_stage_0_ms = 0.0
|
||||
for sd in all_stage_durations:
|
||||
actual = dict(sd)
|
||||
s0 = sd.get("stage_0_gen_ms", 0.0)
|
||||
actual["stage_0_gen_ms"] = s0 - prev_stage_0_ms
|
||||
prev_stage_0_ms = s0
|
||||
per_request_actual.append(actual)
|
||||
|
||||
per_request_e2e_ms: list[float] = []
|
||||
for actual in per_request_actual:
|
||||
e2e_ms = sum(v for k, v in actual.items() if k in _TIMING_ORDER)
|
||||
if e2e_ms > 0:
|
||||
per_request_e2e_ms.append(e2e_ms)
|
||||
|
||||
# Report
|
||||
print("\n" + "=" * 60)
|
||||
print("vLLM-Omni Benchmark Results")
|
||||
print("=" * 60)
|
||||
print(f"{'Mode:':<40} {args.mode}")
|
||||
print(f"{'Model:':<40} {args.model_path}")
|
||||
print(f"{'Image size:':<40} {args.height}x{args.width}")
|
||||
print(f"{'Num inference steps:':<40} {args.num_inference_steps}")
|
||||
print("-" * 50)
|
||||
print(f"{'Init time (s):':<40} {init_time:.2f}")
|
||||
print(f"{'Successful:':<40} {success}/{valid}")
|
||||
print(f"{'Failed:':<40} {failed}")
|
||||
print("-" * 50)
|
||||
|
||||
if per_request_e2e_ms:
|
||||
per_request_s = np.array(per_request_e2e_ms) / 1000.0
|
||||
print(f"{'Total generation time (s):':<40} {total_gen_time:.2f}")
|
||||
print(f"{'Throughput (img/s):':<40} {success / total_gen_time:.4f}")
|
||||
print(f"{'Latency Mean (s):':<40} {per_request_s.mean():.4f}")
|
||||
print(f"{'Latency Median (s):':<40} {np.median(per_request_s):.4f}")
|
||||
print(f"{'Latency P95 (s):':<40} {np.percentile(per_request_s, 95):.4f}")
|
||||
print(f"{'Latency P99 (s):':<40} {np.percentile(per_request_s, 99):.4f}")
|
||||
print(f"{'Latency Min (s):':<40} {per_request_s.min():.4f}")
|
||||
print(f"{'Latency Max (s):':<40} {per_request_s.max():.4f}")
|
||||
elif latencies:
|
||||
per_request = np.diff([0.0] + list(latencies))
|
||||
print(f"{'Total generation time (s):':<40} {total_gen_time:.2f}")
|
||||
print(f"{'Throughput (img/s):':<40} {success / total_gen_time:.4f}")
|
||||
print(f"{'Latency Mean (s) [wall-clock]:':<40} {per_request.mean():.4f}")
|
||||
print(f"{'Latency Median (s) [wall-clock]:':<40} {np.median(per_request):.4f}")
|
||||
print(f"{'Latency P95 (s) [wall-clock]:':<40} {np.percentile(per_request, 95):.4f}")
|
||||
print(f"{'Latency P99 (s) [wall-clock]:':<40} {np.percentile(per_request, 99):.4f}")
|
||||
print(f"{'Latency Min (s) [wall-clock]:':<40} {per_request.min():.4f}")
|
||||
print(f"{'Latency Max (s) [wall-clock]:':<40} {per_request.max():.4f}")
|
||||
|
||||
if per_request_actual:
|
||||
print("-" * 50)
|
||||
print("Pipeline Timings Mean:")
|
||||
for key in _TIMING_ORDER:
|
||||
vals = [d.get(key, 0.0) for d in per_request_actual]
|
||||
if any(v != 0 for v in vals):
|
||||
unit = "ms" if key.endswith("_ms") else "s"
|
||||
print(f" {key + ':':<38} {np.mean(vals):.4f} ({unit})")
|
||||
# Show any extra keys not in the ordered list
|
||||
ordered_set = set(_TIMING_ORDER)
|
||||
extra_keys = sorted(k for k in per_request_actual[0].keys() if k not in ordered_set)
|
||||
for key in extra_keys:
|
||||
vals = [d.get(key, 0.0) for d in per_request_actual]
|
||||
if any(v != 0 for v in vals):
|
||||
unit = "ms" if key.endswith("_ms") else "s"
|
||||
print(f" {key + ':':<38} {np.mean(vals):.4f} ({unit})")
|
||||
|
||||
print(f"\n{'Output dir:':<40} {args.output_dir}")
|
||||
print("=" * 60)
|
||||
|
||||
# Metrics JSON
|
||||
metrics = {
|
||||
"backend": "vllm-omni",
|
||||
"mode": args.mode,
|
||||
"model": args.model_path,
|
||||
"height": args.height,
|
||||
"width": args.width,
|
||||
"num_inference_steps": args.num_inference_steps,
|
||||
"init_time_s": init_time,
|
||||
"completed_requests": success,
|
||||
"failed_requests": failed,
|
||||
"total_gen_time_s": total_gen_time,
|
||||
"throughput_qps": success / total_gen_time if total_gen_time > 0 else 0,
|
||||
}
|
||||
if per_request_e2e_ms:
|
||||
per_request_s = np.array(per_request_e2e_ms) / 1000.0
|
||||
metrics["latency_mean"] = float(per_request_s.mean())
|
||||
metrics["latency_median"] = float(np.median(per_request_s))
|
||||
metrics["latency_p95"] = float(np.percentile(per_request_s, 95))
|
||||
metrics["latency_p99"] = float(np.percentile(per_request_s, 99))
|
||||
elif latencies:
|
||||
per_request = np.diff([0.0] + list(latencies))
|
||||
metrics["latency_mean"] = float(per_request.mean())
|
||||
metrics["latency_median"] = float(np.median(per_request))
|
||||
metrics["latency_p95"] = float(np.percentile(per_request, 95))
|
||||
metrics["latency_p99"] = float(np.percentile(per_request, 99))
|
||||
else:
|
||||
metrics["latency_mean"] = 0
|
||||
metrics["latency_median"] = 0
|
||||
metrics["latency_p95"] = 0
|
||||
metrics["latency_p99"] = 0
|
||||
if per_request_actual:
|
||||
all_keys = list(_TIMING_ORDER) + sorted(k for k in per_request_actual[0].keys() if k not in set(_TIMING_ORDER))
|
||||
stage_metrics = {}
|
||||
for key in all_keys:
|
||||
vals = [d.get(key, 0.0) for d in per_request_actual]
|
||||
stage_metrics[key] = {
|
||||
"mean": float(np.mean(vals)),
|
||||
"median": float(np.median(vals)),
|
||||
"p95": float(np.percentile(vals, 95)),
|
||||
}
|
||||
metrics["stage_durations"] = stage_metrics
|
||||
if args.output_file:
|
||||
with open(args.output_file, "w") as f:
|
||||
json.dump(metrics, f, indent=2)
|
||||
print(f"Metrics saved to {args.output_file}")
|
||||
|
||||
omni.close()
|
||||
print("Done!")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="GLM-Image vLLM-Omni offline benchmark")
|
||||
parser.add_argument("--model-path", type=str, default="zai-org/GLM-Image")
|
||||
parser.add_argument("--deploy-config", type=str, default=None, help="Deploy config YAML")
|
||||
parser.add_argument("--mode", type=str, default="t2i", choices=["t2i", "i2i"])
|
||||
parser.add_argument("--dataset-path", type=str, default=None, help="Path to prompt.json")
|
||||
parser.add_argument("--num-prompts", type=int, default=10)
|
||||
parser.add_argument("--height", type=int, default=HEIGHT)
|
||||
parser.add_argument("--width", type=int, default=WIDTH)
|
||||
parser.add_argument("--num-inference-steps", type=int, default=NUM_INFERENCE_STEPS)
|
||||
parser.add_argument("--guidance-scale", type=float, default=GUIDANCE_SCALE)
|
||||
parser.add_argument("--seed", type=int, default=SEED)
|
||||
parser.add_argument("--output-dir", type=str, default="benchmarks/glm_image/vllm-omni/outputs")
|
||||
parser.add_argument("--output-file", type=str, default=None, help="JSON file for metrics")
|
||||
parser.add_argument("--stage-init-timeout", type=int, default=600)
|
||||
parser.add_argument(
|
||||
"--enable-diffusion-pipeline-profiler",
|
||||
action="store_true",
|
||||
help="Enable diffusion pipeline profiler for stage-level timing",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enable-ar-profiler",
|
||||
action="store_true",
|
||||
help="Enable AR stage profiler to include AR timing in stage_durations",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--log-stats",
|
||||
action="store_true",
|
||||
help="Enable detailed per-request pipeline stats logging",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
benchmark(args)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,5 +1,5 @@
|
||||
ARG VLLM_BASE_IMAGE=vllm/vllm-openai
|
||||
ARG VLLM_BASE_TAG=v0.19.0
|
||||
ARG VLLM_BASE_TAG=v0.20.0
|
||||
FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
|
||||
ARG APP_DIR=/workspace/vllm-omni
|
||||
WORKDIR ${APP_DIR}
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
ARG BASE_IMAGE=vllm/vllm-openai:v0.19.0
|
||||
ARG BASE_IMAGE=vllm/vllm-openai:v0.20.0
|
||||
FROM ${BASE_IMAGE}
|
||||
|
||||
ARG COMMON_WORKDIR=/app
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.19.0
|
||||
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.20.0
|
||||
FROM ${BASE_IMAGE} AS base
|
||||
|
||||
# Declare a variable to know if we want to use the nightly build or the stable build.
|
||||
|
||||
+29
-22
@@ -14,8 +14,9 @@ RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRO
|
||||
RUN apt clean && apt-get update -y && \
|
||||
apt-get install -y --no-install-recommends --fix-missing \
|
||||
curl \
|
||||
espeak-ng \
|
||||
ffmpeg \
|
||||
git \
|
||||
libsndfile1 \
|
||||
libsm6 \
|
||||
libxext6 \
|
||||
libgl1 \
|
||||
@@ -28,8 +29,20 @@ RUN apt clean && apt-get update -y && \
|
||||
python3.12-dev \
|
||||
python3-pip
|
||||
|
||||
RUN apt update && apt upgrade -y && \
|
||||
apt install -y intel-oneapi-compiler-dpcpp-cpp-2025.3
|
||||
# Add oneAPI repo, pin oneAPI to 2025.3, then install pinned packages in one layer.
|
||||
RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | gpg --dearmor | tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null && \
|
||||
echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list && \
|
||||
printf '%s\n' \
|
||||
'Package: intel-oneapi-* intel-deep-learning-essentials* intel-pti*' \
|
||||
'Pin: version 2025.3*' \
|
||||
'Pin-Priority: 1001' \
|
||||
> /etc/apt/preferences.d/oneapi-2025.3.pref && \
|
||||
apt-get update -y && \
|
||||
apt-get install -y --no-install-recommends \
|
||||
intel-oneapi-compiler-dpcpp-cpp-2025.3 \
|
||||
intel-oneapi-mkl-devel-2025.3 \
|
||||
intel-oneapi-dnnl-devel-2025.3 && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Install UMD with fixed version
|
||||
RUN mkdir neo && \
|
||||
@@ -52,9 +65,9 @@ RUN curl -LsSf https://astral.sh/uv/install.sh | sh
|
||||
RUN uv venv --python ${PYTHON_VERSION} --seed ${VIRTUAL_ENV}
|
||||
ENV PATH="$VIRTUAL_ENV/bin:$PATH"
|
||||
|
||||
# This oneccl contains the BMG support which is not the case for default version of oneapi 2025.2.
|
||||
ARG ONECCL_INSTALLER="intel-oneccl-2021.15.7.8_offline.sh"
|
||||
RUN wget "https://github.com/uxlfoundation/oneCCL/releases/download/2021.15.7/${ONECCL_INSTALLER}" && \
|
||||
# This oneccl contains the BMG support which is not the case for default version of oneapi 2025.3.
|
||||
ARG ONECCL_INSTALLER="intel-oneccl-2021.15.9.14_offline.sh"
|
||||
RUN wget "https://github.com/uxlfoundation/oneCCL/releases/download/2021.15.9/${ONECCL_INSTALLER}" && \
|
||||
bash "${ONECCL_INSTALLER}" -a --silent --eula accept && \
|
||||
rm "${ONECCL_INSTALLER}" && \
|
||||
echo "source /opt/intel/oneapi/setvars.sh --force" >> /root/.bashrc && \
|
||||
@@ -66,6 +79,7 @@ SHELL ["bash", "-c"]
|
||||
CMD ["bash", "-c", "source /root/.bashrc && exec bash"]
|
||||
|
||||
WORKDIR /workspace/
|
||||
|
||||
ENV UV_HTTP_TIMEOUT=500
|
||||
|
||||
# Configure package index for XPU
|
||||
@@ -74,24 +88,18 @@ ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
|
||||
ENV UV_INDEX_STRATEGY="unsafe-best-match"
|
||||
ENV UV_LINK_MODE="copy"
|
||||
|
||||
ARG VLLM_VERSION=v0.19.0
|
||||
ARG VLLM_VERSION=v0.20.0
|
||||
RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm
|
||||
WORKDIR /workspace/vllm
|
||||
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install --upgrade pip && \
|
||||
uv pip install -r requirements/xpu.txt
|
||||
|
||||
# used for suffix method speculative decoding
|
||||
# build deps for proto + nanobind-based extensions to set up the build environment
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install grpcio-tools protobuf nanobind
|
||||
# arctic-inference is built from source which needs torch-xpu properly installed first
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install -r requirements/xpu.txt && \
|
||||
uv pip install setuptools_scm &&\
|
||||
uv pip install grpcio-tools protobuf nanobind && \
|
||||
source /opt/intel/oneapi/setvars.sh --force && \
|
||||
source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force && \
|
||||
export CMAKE_PREFIX_PATH="$(python -c 'import site; print(site.getsitepackages()[0])'):${CMAKE_PREFIX_PATH}" && \
|
||||
uv pip install --no-build-isolation arctic-inference==0.1.1
|
||||
export CMAKE_PREFIX_PATH="$(python3 -c 'import site; print(site.getsitepackages()[0])'):${CMAKE_PREFIX_PATH}"
|
||||
|
||||
ENV LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/local/lib/"
|
||||
|
||||
@@ -105,6 +113,9 @@ CMD ["/bin/bash"]
|
||||
|
||||
FROM vllm-base AS vllm-openai
|
||||
|
||||
# install nixl from source code
|
||||
RUN uv pip install nixl
|
||||
|
||||
# install additional dependencies for openai api server
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip install accelerate hf_transfer pytest pytest_asyncio lm_eval[api] modelscope
|
||||
@@ -112,10 +123,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
# install development dependencies (for testing)
|
||||
RUN uv pip install -e tests/vllm_test_utils
|
||||
|
||||
# install nixl from source code
|
||||
ENV NIXL_VERSION=0.7.0
|
||||
RUN python /workspace/vllm/tools/install_nixl_from_source_ubuntu.py
|
||||
|
||||
# ensure vllm is properly installed
|
||||
RUN python -c "import vllm, inspect; print(vllm.__file__)"
|
||||
RUN uv pip show vllm
|
||||
@@ -135,7 +142,7 @@ RUN uv pip install --no-cache-dir ".[dev]" --no-build-isolation
|
||||
# FIX triton
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip uninstall triton triton-xpu && \
|
||||
uv pip install triton-xpu==3.6.0
|
||||
uv pip install triton-xpu==3.7.0
|
||||
|
||||
# remove torch bundled oneccl to avoid conflicts
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
|
||||
+14
-13
@@ -10,40 +10,38 @@ nav:
|
||||
- Image Generation: serving/image_generation_api.md
|
||||
- Image Edit: serving/image_edit_api.md
|
||||
- Text to Speech: serving/speech_api.md
|
||||
- Audio Generation: serving/audio_generate_api.md
|
||||
- Streaming Video Input: serving/video_stream_api.md
|
||||
- Examples:
|
||||
- examples/README.md
|
||||
- Offline Inference:
|
||||
- BAGEL-7B-MoT: user_guide/examples/offline_inference/bagel.md
|
||||
- CosyVoice3: user_guide/examples/offline_inference/cosyvoice3.md
|
||||
- Fish Speech S2 Pro: user_guide/examples/offline_inference/fish_speech.md
|
||||
- GLM-Image Multistage End-to-End Inference: user_guide/examples/offline_inference/glm_image.md
|
||||
- Helios Video Generation: user_guide/examples/offline_inference/helios.md
|
||||
- HunyuanImage-3.0 Image-to-Text Inference: user_guide/examples/offline_inference/hunyuan_image3.md
|
||||
- Image-To-Image: user_guide/examples/offline_inference/image_to_image.md
|
||||
- Image-To-Video: user_guide/examples/offline_inference/image_to_video.md
|
||||
- InternVLA-A1: user_guide/examples/offline_inference/internvla_a1.md
|
||||
- MammothModa2-Preview: user_guide/examples/offline_inference/mammothmodal2_preview.md
|
||||
- MiMo-Audio Offline Inference: user_guide/examples/offline_inference/mimo_audio.md
|
||||
- Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md
|
||||
- Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md
|
||||
- Qwen3-TTS: user_guide/examples/offline_inference/qwen3_tts.md
|
||||
- Text-To-Audio: user_guide/examples/offline_inference/text_to_audio.md
|
||||
- Text-To-Image: user_guide/examples/offline_inference/text_to_image.md
|
||||
- Text-To-Speech: user_guide/examples/offline_inference/text_to_speech.md
|
||||
- Text-To-Video: user_guide/examples/offline_inference/text_to_video.md
|
||||
- Voxtral TTS Offline Inference: user_guide/examples/offline_inference/voxtral_tts.md
|
||||
- X-To-Video-Audio: user_guide/examples/offline_inference/x_to_video_audio.md
|
||||
- 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
|
||||
- Diffusers Backend Adapter: user_guide/examples/online_serving/diffusers_pipeline_adapter.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
|
||||
- Image-To-Video: user_guide/examples/online_serving/image_to_video.md
|
||||
- Online serving Example of vLLM-Omni for MiMo-Audio: user_guide/examples/online_serving/mimo_audio.md
|
||||
- Qwen2.5-Omni: user_guide/examples/online_serving/qwen2_5_omni.md
|
||||
- Qwen3-Omni: user_guide/examples/online_serving/qwen3_omni.md
|
||||
- Qwen3-TTS: user_guide/examples/online_serving/qwen3_tts.md
|
||||
- Text-To-Speech: user_guide/examples/online_serving/text_to_speech.md
|
||||
- Text-To-Image: user_guide/examples/online_serving/text_to_image.md
|
||||
- Text-To-Video: user_guide/examples/online_serving/text_to_video.md
|
||||
- General:
|
||||
@@ -61,11 +59,6 @@ nav:
|
||||
- Cache Acceleration:
|
||||
- TeaCache: user_guide/diffusion/cache_acceleration/teacache.md
|
||||
- Cache-DiT: user_guide/diffusion/cache_acceleration/cache_dit.md
|
||||
- Quantization:
|
||||
- Overview: user_guide/diffusion/quantization/overview.md
|
||||
- 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:
|
||||
@@ -80,6 +73,14 @@ nav:
|
||||
- LoRA: user_guide/diffusion/lora.md
|
||||
- Custom Pipeline: features/custom_pipeline.md
|
||||
- Step Execution: user_guide/diffusion/step_execution.md
|
||||
- Quantization:
|
||||
- Overview: user_guide/quantization/overview.md
|
||||
- Online Quantization: user_guide/quantization/online.md
|
||||
- FP8 W8A8: user_guide/quantization/fp8.md
|
||||
- Int8 W8A8: user_guide/quantization/int8.md
|
||||
- GGUF: user_guide/quantization/gguf.md
|
||||
- AutoRound: user_guide/quantization/autoround.md
|
||||
- msModelSlim: user_guide/quantization/msmodelslim.md
|
||||
- ComfyUI: features/comfyui.md
|
||||
- Developer Guide:
|
||||
- General:
|
||||
@@ -91,7 +92,6 @@ nav:
|
||||
- contributing/model/adding_omni_model.md
|
||||
- contributing/model/adding_tts_model.md
|
||||
- contributing/model/adding_diffusion_model.md
|
||||
- contributing/model/adding_quantization_model.md
|
||||
- CI: contributing/ci
|
||||
- Design Documents:
|
||||
- design/index.md
|
||||
@@ -112,6 +112,7 @@ nav:
|
||||
- design/feature/async_chunk.md
|
||||
- design/feature/vae_parallel.md
|
||||
- design/feature/diffusion_step_execution.md
|
||||
- design/feature/diffusion_continuous_batching.md
|
||||
- Module Design:
|
||||
- design/module/ar_module.md
|
||||
- design/module/dit_module.md
|
||||
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 96 KiB After Width: | Height: | Size: 182 KiB |
+5
-3
@@ -12,9 +12,11 @@ The stage-based CLI is designed for deployments that require launching each pipe
|
||||
`vllm_omni/model_executor/stage_configs/`, the `--stage-configs-path` parameter remains mandatory.
|
||||
|
||||
Example: Initializing Stage 0 (Orchestrator and API Server):
|
||||
The commands below show a common device mapping where Stage 0 uses GPU 0 and
|
||||
worker stages use GPU 1 via `CUDA_VISIBLE_DEVICES`.
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
--port 8091 \
|
||||
--stage-id 0 \
|
||||
--omni-master-address 127.0.0.1 \
|
||||
@@ -24,7 +26,7 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
Example: Initializing a Headless Worker Stage (Stage 1):
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
--stage-id 1 \
|
||||
--headless \
|
||||
--omni-master-address 127.0.0.1 \
|
||||
@@ -46,7 +48,7 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
|
||||
the stage-based CLI permits the direct initialization of Stage 1 with explicit parameters:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
--stage-id 1 \
|
||||
--headless \
|
||||
--gpu-memory-utilization 0.5 \
|
||||
|
||||
@@ -102,15 +102,17 @@ The stage-based CLI paradigm facilitates the execution of discrete pipeline stag
|
||||
|
||||
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`:
|
||||
the example below uses `CUDA_VISIBLE_DEVICES=0` for Stage 0 and
|
||||
`CUDA_VISIBLE_DEVICES=1` for Stage 1.
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
CUDA_VISIBLE_DEVICES=0 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 \
|
||||
CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
--stage-id 1 \
|
||||
--headless \
|
||||
--omni-master-address 127.0.0.1 \
|
||||
@@ -132,7 +134,7 @@ Conversely, in the context of the **stage-based CLI** paradigm, given that each
|
||||
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 \
|
||||
CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
--stage-id 1 \
|
||||
--headless \
|
||||
--gpu-memory-utilization 0.5 \
|
||||
@@ -189,7 +191,7 @@ Within the stage-based CLI paradigm, equivalent configuration parameters can inh
|
||||
as command-line arguments to the designated single-stage process instantiation:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
--stage-id 0 \
|
||||
--max-num-seqs 8 \
|
||||
--omni-master-address 127.0.0.1 \
|
||||
|
||||
@@ -108,7 +108,8 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
|
||||
/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
|
||||
tests/dfx/reliability/test_reliability_{model_key}.py<br>
|
||||
(e.g. <code>test_reliability_qwen3_omni.py</code>, <code>test_reliability_wan22.py</code>)
|
||||
</td>
|
||||
<td>
|
||||
<a href="#chapter-4-l5-level-testing-stability-and-reliability-testing">Chapter 4</a><br>
|
||||
@@ -576,12 +577,17 @@ 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 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.
|
||||
- ***Reliability Testing***: Uses pytest suites under `tests/dfx/reliability/` to inject controlled faults against a **live** `vllm_omni serve` instance (same **`omni_server` / `omni_server_function`** fixture style as E2E). Current suites emphasize **GPU memory pressure** (CUDA sidecar “memory hog”), **worker / runtime process kill** (`SIGKILL` on `VLLM::Worker` for Qwen3-Omni or `multiprocessing.spawn` for Wan2.2 video workers), **large multimodal chat** or **`/v1/videos`** jobs under OOM, **`/health` → 503** and **fast-fail / non-hanging concurrent** requests after kill, and **OpenAI-style 5xx error contracts** (e.g. text vs text+audio under OOM). **Post-fault recovery** checks exist where enabled (some cases may be `skip` while issues are tracked). See the Reliability `<details>` block in Section 4.4 for file-level responsibilities and CI markers (`slow`, `hardware_test`, POSIX-only kill).
|
||||
|
||||
### 4.3 Test Directory and Execution Files
|
||||
|
||||
- ***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`
|
||||
- ***Reliability Test Suite*** (`tests/dfx/reliability/`):
|
||||
- `test_reliability_qwen3_omni.py` — Qwen3-Omni chat / multimodal reliability (GPU OOM, process kill, recovery, error contract under `--async-chunk` vs default).
|
||||
- `test_reliability_wan22.py` — Wan2.2 T2V video API reliability (`/v1/videos` under OOM and process kill, recovery).
|
||||
- `helpers.py` — Shared primitives used by current suites: raw HTTP probes for `/v1/chat/completions` and `/health`, OpenAI-style error parsing, GPU OOM sidecar (`inject_gpu_oom` / `stop_gpu_oom_hogs`), and `pgrep`-based process-kill injector construction (`make_process_kill_fault_injector`).
|
||||
- `conftest.py` — `fault_injector` and `omni_server_after_fault` / `omni_server_after_fault_function` fixtures to run a callable **after** the server is ready.
|
||||
- `README.md` — Short local run commands for this directory.
|
||||
|
||||
### 4.4 Execution Method and Example
|
||||
|
||||
@@ -662,8 +668,57 @@ All other optional parameters follow the same rules as the in Chapter 3.4.
|
||||
|
||||
</details>
|
||||
|
||||
<details>
|
||||
<summary> Reliability test suite (<code>tests/dfx/reliability</code>)</summary>
|
||||
|
||||
#### Purpose and relationship to stability
|
||||
|
||||
Reliability tests are **short fault-injection** integration runs (L5 **(b)** in `tests/dfx/reliability/README.md`). They complement **stability** JSON-driven long runs: instead of hours of steady traffic, they **perturb** the server (GPU OOM sidecar, fatal signals on selected processes) and check **failure mode** and **latency bounds** (e.g. chat or `/v1/videos` must not hang under concurrent fault-time load).
|
||||
|
||||
#### Directory layout
|
||||
|
||||
| Path | Responsibility |
|
||||
| ---- | -------------- |
|
||||
| `helpers.py` | Shared helpers used by current reliability suites: raw `POST`/`GET` probes (`/v1/chat/completions`, `/health`), OpenAI error parsing (`extract_openai_error_contract_from_bytes`), GPU OOM sidecar lifecycle (`inject_gpu_oom`, `stop_gpu_oom_hogs`), and process-kill injector builder (`make_process_kill_fault_injector`). |
|
||||
| `conftest.py` | Pytest fixtures: indirect `fault_injector`, `omni_server_after_fault` / `omni_server_after_fault_function` (run injector after server is ready, then yield server). |
|
||||
| `test_reliability_qwen3_omni.py` | Qwen3-Omni: OOM vs **text vs text+audio** error contract, large multimodal chat under OOM, concurrent pressure, **SIGKILL** on `VLLM::Worker`, `/health` → 503 + fast-fail + concurrent chat; optional OOM recovery scenario (may be skipped while tracked in issues). |
|
||||
| `test_reliability_wan22.py` | Wan2.2 T2V: large `/v1/videos` under OOM, **SIGKILL** on `multiprocessing.spawn` chain, health / fast-fail / concurrent video requests; optional recovery test (may be skipped). |
|
||||
| `README.md` | Minimal run / collect examples. |
|
||||
|
||||
#### Parametrization and markers
|
||||
|
||||
- Each test module defines a **`RELIABILITY_SCENARIOS`** list (`test_name`, `server_params`: model, `stage_config_name` or diffusion `server_args`, etc.). **`create_reliability_omni_server_params()`** in `tests/dfx/conftest.py` resolves stage paths (including XPU substitutions where applicable) and builds **`OmniServerParams`** lists consumed by **`@pytest.mark.parametrize(..., indirect=True)`** on `omni_server` or `omni_server_function`.
|
||||
- Cases are tagged **`@pytest.mark.slow`** for weekly / selective CI. GPU-heavy suites use **`@hardware_test(res={"cuda": "H100"}, num_cards=...)`** (Qwen3-Omni paths often require **2** cards; Wan2.2 video paths **1** card).
|
||||
- **Process-kill** tests use **`@pytest.mark.skipif(os.name == "nt", ...)`** because injection uses POSIX **`pgrep` / `kill`**.
|
||||
|
||||
#### CI trigger
|
||||
|
||||
Weekly Buildkite (`.buildkite/test-weekly.yml`) runs, for example:
|
||||
|
||||
```bash
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m "slow"
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m "slow"
|
||||
```
|
||||
|
||||
#### Local commands
|
||||
|
||||
```bash
|
||||
pytest --collect-only tests/dfx/reliability
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m slow
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m slow
|
||||
```
|
||||
|
||||
#### Adding a new model suite
|
||||
|
||||
1. Add `test_reliability_<model>.py` under `tests/dfx/reliability/`.
|
||||
2. Define **`RELIABILITY_SCENARIOS`** and pass them through **`create_reliability_omni_server_params()`** with the correct deploy or e2e stage-config directory (same pattern as existing files).
|
||||
3. Reuse **`helpers`** for OOM / kill / raw HTTP; prefer **`assert_fault_exception()`** and **`resolve_oom_device_spec()`** from `tests/dfx/conftest.py` for consistent device selection vs stage YAML.
|
||||
4. Register **`slow`** (and **`hardware_test`** if needed); extend **`.buildkite/test-weekly.yml`** when the suite should run in weekly L5.
|
||||
|
||||
</details>
|
||||
|
||||
- - ***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`
|
||||
- ***Reliability***: `pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m slow` and/or `pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m slow` (add `test_reliability_<suite>.py` for new models)
|
||||
|
||||
## Summary
|
||||
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
**Scope**
|
||||
|
||||
For diffusion models, the L4 functionality test covers all or common *diffusion features* that are supported by this model, including several [parallelism acceleration methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism_acceleration/), [CPU offloading](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cpu_offload_diffusion/), [TeaCache](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/teacache/) and [Cache-DiT](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cache_dit_acceleration/) cache backends, [quantization methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/quantization/overview/).
|
||||
For diffusion models, the L4 functionality test covers all or common *diffusion features* that are supported by this model, including several [parallelism acceleration methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism_acceleration/), [CPU offloading](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cpu_offload_diffusion/), [TeaCache](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/teacache/) and [Cache-DiT](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cache_dit_acceleration/) cache backends, [quantization methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/quantization/overview/).
|
||||
|
||||
**Test Case Design**
|
||||
|
||||
|
||||
@@ -737,6 +737,16 @@ step-scoped phases. The reference implementation is
|
||||
Do not enable `step_execution=True` until those four methods are implemented
|
||||
and validated against the request-level path.
|
||||
|
||||
If you want the pipeline to work with the experimental batched step-wise path
|
||||
(`max_num_seqs > 1`), also see:
|
||||
[Continuous Batching for Step-Wise Diffusion](../../design/feature/diffusion_continuous_batching.md).
|
||||
|
||||
If you expose this in example scripts or recipes, keep it opt-in. Surface
|
||||
runtime features like `step_execution` as optional flags instead of silently
|
||||
turning them on. For Qwen-Image-style serving examples, document
|
||||
`--step-execution` as the feature gate and `--max-num-seqs N` as the
|
||||
companion batching knob.
|
||||
|
||||
### Cache Acceleration
|
||||
|
||||
#### TeaCache
|
||||
|
||||
@@ -1,478 +0,0 @@
|
||||
# Support FP8 Quantization
|
||||
|
||||
This section describes how to add FP8 quantization support to a diffusion model. We use the Qwen-Image pipeline as the reference implementation.
|
||||
|
||||
---
|
||||
|
||||
## Table of Contents
|
||||
|
||||
- [Overview](#overview)
|
||||
- [Step-by-Step Implementation](#step-by-step-implementation)
|
||||
- [Testing](#testing)
|
||||
- [Troubleshooting](#troubleshooting)
|
||||
- [Reference Implementations](#reference-implementations)
|
||||
- [Summary](#summary)
|
||||
|
||||
---
|
||||
|
||||
## Overview
|
||||
|
||||
### What is FP8 Quantization?
|
||||
|
||||
FP8 quantization reduces model memory footprint by storing weights in 8-bit floating point (`float8_e4m3fn`) instead of 16-bit (`bfloat16`/`float16`). This saves ~50% of weight memory.
|
||||
|
||||
A diffusion pipeline typically has three quantizable components:
|
||||
|
||||
| Component | Layer Types | Quantization Approach |
|
||||
|-----------|-----------|----------------------|
|
||||
| **DiT (transformer)** | `nn.Linear` | Replace with vLLM quantized linear layers |
|
||||
| **Text encoder** | `nn.Linear` | FP8 weight storage via hooks |
|
||||
| **VAE** | `nn.Conv2d`, `nn.Conv3d` | FP8 weight storage via hooks |
|
||||
|
||||
### Architecture
|
||||
|
||||
The quantization system has two mechanisms:
|
||||
|
||||
**1. vLLM Quantized Linear Layers (for DiT)**
|
||||
|
||||
vLLM provides drop-in replacements for `nn.Linear` that perform FP8 W8A8 computation on supported hardware. These are configured via `quant_config` passed to layer constructors.
|
||||
|
||||
Key classes:
|
||||
|
||||
| Class | Location | Purpose |
|
||||
|-------|----------|---------|
|
||||
| `build_quant_config` | `vllm_omni/quantization/__init__.py` | Build quantization config from string/dict/per-component spec |
|
||||
| `ComponentQuantizationConfig` | `vllm_omni/quantization/component_config.py` | Per-component routing by layer prefix |
|
||||
|
||||
**2. FP8 Weight Storage via Hooks (for text encoder and VAE)**
|
||||
|
||||
For models loaded via `from_pretrained()` (outside vLLM's weight pipeline), we use forward hooks to store weights in FP8 and dequantize on-the-fly:
|
||||
|
||||
```
|
||||
At rest: weight.data = fp8_tensor (half memory)
|
||||
Pre-hook: weight.data = fp8_weight.to(bf16) * scale
|
||||
Forward: Normal computation in BF16
|
||||
Post-hook: weight.data = fp8_weight (back to FP8)
|
||||
```
|
||||
|
||||
Key function:
|
||||
|
||||
| Function | Location | Purpose |
|
||||
|----------|----------|---------|
|
||||
| `apply_fp8_weight_storage` | `vllm_omni/diffusion/models/utils.py` | Apply FP8 hooks to Linear/Conv layers (introduced in [#1412](https://github.com/vllm-project/vllm-omni/pull/1412)) |
|
||||
|
||||
---
|
||||
|
||||
## Step-by-Step Implementation
|
||||
|
||||
### Step 1: Add `quant_config` to the Transformer
|
||||
|
||||
Pass `quant_config` to vLLM parallel linear layers in the transformer. This enables FP8 W8A8 for DiT linear layers.
|
||||
|
||||
#### 1.1: Accept `quant_config` in Constructor
|
||||
|
||||
```python
|
||||
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
|
||||
|
||||
class YourTransformer2DModel(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
od_config: OmniDiffusionConfig,
|
||||
quant_config: QuantizationConfig | None = None,
|
||||
# ... other params
|
||||
):
|
||||
self.quant_config = quant_config
|
||||
# Pass to sub-modules
|
||||
self.blocks = nn.ModuleList([
|
||||
YourTransformerBlock(quant_config=quant_config, ...)
|
||||
for _ in range(num_layers)
|
||||
])
|
||||
```
|
||||
|
||||
#### 1.2: Pass `quant_config` to vLLM Linear Layers
|
||||
|
||||
Replace `nn.Linear` with vLLM parallel linear layers that accept `quant_config`:
|
||||
|
||||
```python
|
||||
from vllm.model_executor.layers.linear import (
|
||||
ColumnParallelLinear,
|
||||
QKVParallelLinear,
|
||||
RowParallelLinear,
|
||||
)
|
||||
|
||||
class YourAttentionBlock(nn.Module):
|
||||
def __init__(self, hidden_size, num_heads, quant_config=None, prefix=""):
|
||||
# QKV projection
|
||||
self.to_qkv = QKVParallelLinear(
|
||||
hidden_size=hidden_size,
|
||||
head_size=hidden_size // num_heads,
|
||||
total_num_heads=num_heads,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.to_qkv",
|
||||
)
|
||||
# Output projection
|
||||
self.to_out = RowParallelLinear(
|
||||
input_size=hidden_size,
|
||||
output_size=hidden_size,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.to_out",
|
||||
)
|
||||
|
||||
class YourFeedForward(nn.Module):
|
||||
def __init__(self, hidden_size, intermediate_size, quant_config=None, prefix=""):
|
||||
self.gate_up_proj = ColumnParallelLinear(
|
||||
input_size=hidden_size,
|
||||
output_size=intermediate_size * 2,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.gate_up_proj",
|
||||
)
|
||||
self.down_proj = RowParallelLinear(
|
||||
input_size=intermediate_size,
|
||||
output_size=hidden_size,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.down_proj",
|
||||
)
|
||||
```
|
||||
|
||||
!!! important
|
||||
Always pass the `prefix` parameter. It is used by `ignored_layers` matching to skip sensitive layers.
|
||||
|
||||
#### 1.3: Identify Sensitive Layers
|
||||
|
||||
Some layers degrade output quality when quantized. Common sensitive layers:
|
||||
|
||||
| Layer | Why Sensitive | Typical Models | Quant Methods Affected |
|
||||
|-------|-------------|---------------|----------------------|
|
||||
| `img_mlp` | Processes denoising latents with shifting dynamic range | Qwen-Image | FP8, Int8 |
|
||||
| `feed_forward` | FFN layers in DiT blocks, large dynamic range | Z-Image | Int8 |
|
||||
| `proj_out` | Final output projection, small quantization errors amplified | Various | FP8, Int8 |
|
||||
| `lm_head` | Output vocabulary projection, precision-critical | Omni (thinker) | NVFP4 |
|
||||
| `mlp.gate` | MoE router gates, routing decisions are precision-critical | MoE models | NVFP4 |
|
||||
|
||||
**Quantization quality benchmarks** (LPIPS: lower = better, <0.01 imperceptible, >0.1 noticeable):
|
||||
|
||||
| Model | Method | All Layers | With `ignored_layers` | Recommendation |
|
||||
|-------|--------|-----------|----------------------|----------------|
|
||||
| Qwen-Image-2512 | Int8 | LPIPS 0.0197 | 0.0027 (skip `img_mlp`) | Skip `img_mlp` |
|
||||
| Z-Image-Turbo | Int8 | LPIPS 0.1597 | 0.0290 (skip `feed_forward`) | Skip `feed_forward` |
|
||||
| Z-Image-Turbo | FP8 | LPIPS ~0.005 | — | All layers OK |
|
||||
|
||||
To identify sensitive layers for a new model:
|
||||
|
||||
1. Run inference with all layers quantized
|
||||
2. Compare output quality with BF16 baseline using LPIPS (see `benchmarks/diffusion/quantization_quality.py`)
|
||||
3. Selectively disable quantization per layer (via `ignored_layers`) until quality matches
|
||||
4. Document the recommended `ignored_layers` in the model's supported models table
|
||||
|
||||
### Step 2: Wire Up Quantization in the Pipeline
|
||||
|
||||
#### 2.1: Create `quant_config` and Pass to Transformer
|
||||
|
||||
In the pipeline's `__init__`, extract the vLLM quant config and pass it to the transformer:
|
||||
|
||||
```python
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
|
||||
class YourPipeline(nn.Module):
|
||||
def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
|
||||
super().__init__()
|
||||
# ... load scheduler, text_encoder, vae ...
|
||||
|
||||
# Build quant config from OmniDiffusionConfig
|
||||
quant_config = build_quant_config(od_config.quantization_config)
|
||||
|
||||
# Pass to transformer
|
||||
self.transformer = YourTransformer2DModel(
|
||||
od_config=od_config,
|
||||
quant_config=quant_config,
|
||||
# ... other kwargs
|
||||
)
|
||||
```
|
||||
|
||||
#### 2.2: Apply FP8 Weight Storage to Text Encoder and VAE
|
||||
|
||||
For models loaded via `from_pretrained()`, apply hook-based FP8 weight storage.
|
||||
|
||||
!!! note
|
||||
The `apply_fp8_weight_storage` function is provided by `vllm_omni/diffusion/models/utils.py`, introduced in [PR #1412](https://github.com/vllm-project/vllm-omni/pull/1412).
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.models.utils import apply_fp8_weight_storage
|
||||
|
||||
class YourPipeline(nn.Module):
|
||||
def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
|
||||
super().__init__()
|
||||
|
||||
self.text_encoder = SomeTextEncoder.from_pretrained(...)
|
||||
self.vae = SomeVAE.from_pretrained(...).to(self.device)
|
||||
|
||||
# Apply FP8 weight storage to text encoder and VAE
|
||||
if (
|
||||
od_config.quantization_config is not None
|
||||
and getattr(od_config.quantization_config, "quant_method", None) == "fp8"
|
||||
):
|
||||
apply_fp8_weight_storage(self.vae)
|
||||
apply_fp8_weight_storage(self.text_encoder)
|
||||
|
||||
# ... rest of init
|
||||
```
|
||||
|
||||
`apply_fp8_weight_storage` walks all `nn.Linear`, `nn.Conv2d`, and `nn.Conv3d` modules in the model and applies FP8 hooks to each.
|
||||
|
||||
#### 2.3: Fix `load_weights` to Include Pre-loaded Components
|
||||
|
||||
Since text encoders and VAEs are loaded via `from_pretrained()` (not through vLLM's weight pipeline), their parameters won't appear in the `loaded_weights` set. Mark them explicitly:
|
||||
|
||||
```python
|
||||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||||
loader = AutoWeightsLoader(self)
|
||||
loaded_weights = loader.load_weights(weights)
|
||||
# Mark pre-loaded component weights as loaded
|
||||
loaded_weights |= {f"vae.{name}" for name, _ in self.vae.named_parameters()}
|
||||
loaded_weights |= {f"text_encoder.{name}" for name, _ in self.text_encoder.named_parameters()}
|
||||
return loaded_weights
|
||||
```
|
||||
|
||||
### Step 3: Add to Supported Models Documentation
|
||||
|
||||
Update the FP8 supported models table in `docs/user_guide/diffusion/quantization/fp8.md`:
|
||||
|
||||
```markdown
|
||||
| Your-Model | `org/Your-Model` | ✅ | ✅ | ✅ | `sensitive_layer` |
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Testing
|
||||
|
||||
### Quick Verification
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
# Test with FP8 enabled
|
||||
omni = Omni(model="your-model", quantization="fp8")
|
||||
|
||||
outputs = omni.generate(
|
||||
"A cat sitting on a windowsill",
|
||||
OmniDiffusionSamplingParams(num_inference_steps=50),
|
||||
)
|
||||
```
|
||||
|
||||
### Quality Gate Test (LPIPS)
|
||||
|
||||
We provide a pytest-based quality gate at `tests/diffusion/quantization/test_quantization_quality.py`.
|
||||
It generates outputs with both BF16 and your quantized method using the same seed, computes
|
||||
[LPIPS](https://github.com/richzhang/PerceptualSimilarity) perceptual distance, and **fails if it
|
||||
exceeds a threshold**. This is the recommended way to validate that a quantization method does not
|
||||
unacceptably degrade output quality.
|
||||
|
||||
**Step 1: Add your test config** — edit `QUALITY_CONFIGS` in the test file:
|
||||
|
||||
```python
|
||||
QualityTestConfig(
|
||||
id="int8_z_image", # pytest ID
|
||||
model="Tongyi-MAI/Z-Image-Turbo",
|
||||
quantization="int8", # your quantization method
|
||||
task="t2i", # "t2i" or "t2v"
|
||||
prompt="a cup of coffee on a wooden table, morning light",
|
||||
max_lpips=0.15, # fail threshold (see guide below)
|
||||
num_inference_steps=20,
|
||||
),
|
||||
```
|
||||
|
||||
**Step 2: Run the test:**
|
||||
|
||||
```bash
|
||||
# Install dependency
|
||||
pip install lpips
|
||||
|
||||
# Run all quality tests
|
||||
pytest tests/diffusion/quantization/test_quantization_quality.py -v -m ""
|
||||
|
||||
# Run only your method
|
||||
pytest tests/diffusion/quantization/test_quantization_quality.py -v -m "" -k "int8"
|
||||
|
||||
# Run a specific model
|
||||
pytest tests/diffusion/quantization/test_quantization_quality.py -v -m "" -k "z_image"
|
||||
```
|
||||
|
||||
**Step 3: Read the output:**
|
||||
|
||||
```
|
||||
Quantization Quality: int8_z_image
|
||||
Model: Tongyi-MAI/Z-Image-Turbo
|
||||
Method: int8
|
||||
LPIPS: 0.0290 (threshold: 0.15)
|
||||
BF16 memory: 19.15 GiB
|
||||
Quant memory: 14.80 GiB (23% reduction)
|
||||
Result: PASS
|
||||
```
|
||||
|
||||
**LPIPS threshold guide:**
|
||||
|
||||
| LPIPS Range | Meaning | Recommendation |
|
||||
|-------------|---------|----------------|
|
||||
| < 0.01 | Imperceptible | Ideal — no visible difference |
|
||||
| 0.01 – 0.05 | Minor differences | Good — acceptable for most use cases |
|
||||
| 0.05 – 0.15 | Noticeable on close inspection | Acceptable with `ignored_layers` tuning |
|
||||
| 0.15 – 0.30 | Clearly visible | Investigate — likely needs `ignored_layers` |
|
||||
| > 0.30 | Significant degradation | Quantization method may not suit this model |
|
||||
|
||||
!!! tip
|
||||
If your test fails, don't raise the threshold — instead, identify sensitive layers and add
|
||||
them to `ignored_layers` (see next section). Then re-run the test with the lower LPIPS.
|
||||
|
||||
### Quality Benchmark (Detailed)
|
||||
|
||||
For a full multi-prompt benchmark with Markdown output (suitable for PR descriptions), use the
|
||||
standalone benchmark script:
|
||||
|
||||
```bash
|
||||
python benchmarks/diffusion/quantization_quality.py \
|
||||
--model Tongyi-MAI/Z-Image-Turbo \
|
||||
--quantization fp8 \
|
||||
--prompts \
|
||||
"a cup of coffee on the table" \
|
||||
"an aerial view of a coral reef" \
|
||||
--height 1024 --width 1024 \
|
||||
--num-inference-steps 50 --seed 42
|
||||
```
|
||||
|
||||
This generates BF16 and quantized outputs, computes per-prompt LPIPS, and prints a Markdown
|
||||
summary table ready to paste into your PR.
|
||||
|
||||
### Identifying Sensitive Layers with `ignored_layers`
|
||||
|
||||
Not all layers tolerate quantization equally. **`ignored_layers`** lets you keep specific layers
|
||||
in BF16 while quantizing everything else. This is critical for achieving good quality with
|
||||
aggressive quantization methods (Int8, NVFP4).
|
||||
|
||||
**Why some layers are sensitive:**
|
||||
|
||||
| Layer Type | Why Sensitive | Typical Impact |
|
||||
|------------|--------------|----------------|
|
||||
| `img_mlp` | Processes denoising latents with shifting dynamic range across timesteps. No built-in normalization to absorb quantization error. | Color shifts, blurriness |
|
||||
| `feed_forward` | FFN layers with large dynamic range in DiT blocks. | Artifacts, loss of detail |
|
||||
| `proj_out` | Final output projection — small errors get amplified since there are no subsequent layers to correct them. | Overall quality loss |
|
||||
| `lm_head` | Vocabulary projection in LLMs — precision-critical for token selection. | Garbage text output |
|
||||
| `mlp.gate` | MoE router gates — routing decisions are binary and precision-critical. | Wrong expert selection, quality collapse |
|
||||
|
||||
**How to find sensitive layers for your model:**
|
||||
|
||||
1. Run the quality gate test with all layers quantized:
|
||||
```bash
|
||||
pytest tests/diffusion/quantization/test_quantization_quality.py -v -m "" -k "your_model"
|
||||
```
|
||||
|
||||
2. If LPIPS exceeds your threshold, try skipping common sensitive layers one at a time:
|
||||
```python
|
||||
QualityTestConfig(
|
||||
id="fp8_your_model_skip_mlp",
|
||||
model="your-org/Your-Model",
|
||||
quantization="fp8",
|
||||
task="t2i",
|
||||
prompt="a cup of coffee on a wooden table",
|
||||
max_lpips=0.05,
|
||||
),
|
||||
```
|
||||
|
||||
And in the test, modify the Omni call:
|
||||
```python
|
||||
omni_qt = Omni(
|
||||
model=config.model,
|
||||
quantization_config={
|
||||
"method": config.quantization,
|
||||
"ignored_layers": ["img_mlp"],
|
||||
},
|
||||
)
|
||||
```
|
||||
|
||||
3. Compare LPIPS with and without `ignored_layers` to isolate the problematic layers.
|
||||
|
||||
4. Document the recommended `ignored_layers` in the supported models table.
|
||||
|
||||
**Real-world examples:**
|
||||
|
||||
| Model | Method | All Layers (LPIPS) | With `ignored_layers` | Which Layers to Skip |
|
||||
|-------|--------|-------------------|----------------------|---------------------|
|
||||
| Qwen-Image-2512 | Int8 | 0.0197 | 0.0027 | `img_mlp` |
|
||||
| Z-Image-Turbo | Int8 | 0.1597 | 0.0290 | `feed_forward` |
|
||||
| Z-Image-Turbo | FP8 | ~0.005 | — | None needed |
|
||||
|
||||
### What to Verify (Checklist)
|
||||
|
||||
1. **Model loads without errors** — Check logs for quantization initialization messages
|
||||
2. **Quality gate passes** — LPIPS within threshold vs BF16 baseline
|
||||
3. **Memory reduction** — Expect ~30-50% savings on quantized components
|
||||
4. **Sensitive layers identified** — Document which layers (if any) need `ignored_layers`
|
||||
5. **`ignored_layers` pattern works** — Verify that the layer name prefix matches your transformer's module names
|
||||
|
||||
---
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### Issue: Output Quality Degradation
|
||||
|
||||
**Symptoms:** Generated images have visible artifacts, color shifts, or blurriness compared to BF16 baseline.
|
||||
|
||||
**Causes & Solutions:**
|
||||
|
||||
- **Sensitive layers being quantized:**
|
||||
|
||||
**Solution:** Identify and skip sensitive layers:
|
||||
```python
|
||||
omni = Omni(
|
||||
model="your-model",
|
||||
quantization_config={
|
||||
"method": "fp8",
|
||||
"ignored_layers": ["img_mlp"],
|
||||
},
|
||||
)
|
||||
```
|
||||
|
||||
- **All layers sensitive (rare):**
|
||||
|
||||
**Solution:** The model may not be suitable for FP8 quantization. Consider weight-only quantization or other methods.
|
||||
|
||||
### Issue: Missing Weights Warning
|
||||
|
||||
**Symptoms:** Warnings about unloaded weights for `vae.*` or `text_encoder.*` parameters.
|
||||
|
||||
**Solution:** Ensure `load_weights` marks pre-loaded component weights. See [Step 2.3](#23-fix-load_weights-to-include-pre-loaded-components).
|
||||
|
||||
### Issue: CUDA Out of Memory
|
||||
|
||||
**Symptoms:** OOM errors during model loading despite FP8 being enabled.
|
||||
|
||||
**Causes & Solutions:**
|
||||
|
||||
- **Text encoder/VAE loaded in BF16 before FP8 conversion:**
|
||||
|
||||
FP8 weight storage is applied after `from_pretrained()`, so peak load memory is still BF16. This is a one-time cost during initialization.
|
||||
|
||||
---
|
||||
|
||||
## Reference Implementations
|
||||
|
||||
| Model | Pipeline File | Transformer File | Notes |
|
||||
|-------|-------------|-----------------|-------|
|
||||
| **Z-Image** | `models/z_image/pipeline_z_image.py` | `models/z_image/z_image_transformer.py` | DiT + text encoder FP8 |
|
||||
| **Qwen-Image** | `models/qwen_image/pipeline_qwen_image.py` | `models/qwen_image/qwen_image_transformer.py` | DiT + text encoder + VAE FP8 |
|
||||
|
||||
All files are under `vllm_omni/diffusion/`.
|
||||
|
||||
---
|
||||
|
||||
## Summary
|
||||
|
||||
Adding FP8 quantization support to a new model:
|
||||
|
||||
1. ✅ **Add `quant_config` to transformer** — Accept and pass to vLLM parallel linear layers with proper `prefix`
|
||||
2. ✅ **Identify sensitive layers** — Profile which layers need `ignored_layers` for quality
|
||||
3. ✅ **Apply FP8 weight storage** — Call `apply_fp8_weight_storage()` on text encoder and VAE in pipeline `__init__`
|
||||
4. ✅ **Fix `load_weights`** — Mark pre-loaded component weights as loaded
|
||||
5. ✅ **Update docs** — Add model to the supported models table in `docs/user_guide/diffusion/quantization/fp8.md`
|
||||
6. ✅ **Test** — Verify quality, memory savings, and sensitive layer behavior
|
||||
@@ -186,6 +186,17 @@ vllm_omni/model_executor/stage_configs/
|
||||
your_model_name_async_chunk.yaml # Streaming mode config
|
||||
```
|
||||
|
||||
### Example placement
|
||||
|
||||
TTS examples live in the consolidated text-to-speech hub, **not** in their
|
||||
own top-level directory. Place per-model scripts under
|
||||
`examples/offline_inference/text_to_speech/<your_model>/` and
|
||||
`examples/online_serving/text_to_speech/<your_model>/`, and add a section
|
||||
to the hub `README.md` files (table row + per-model section) instead of a
|
||||
new per-model `README.md`. The mkdocs `generate_examples` hook treats the
|
||||
`text_to_speech/` parent as a single example, so per-model READMEs inside
|
||||
it would not be picked up — the hub README is the documented surface.
|
||||
|
||||
**Qwen3-TTS reference files:**
|
||||
|
||||
| File | Purpose |
|
||||
|
||||
@@ -37,6 +37,8 @@ Supported fields:
|
||||
| `warmup_iterations` | Torch-profiler warmup iterations. |
|
||||
| `active_iterations` | Torch-profiler active iterations. |
|
||||
|
||||
For detailed explanations of the fields, please refer to upstream vLLM implementation [vllm/config/profiler.py](https://github.com/vllm-project/vllm/blob/v0.20.1/vllm/config/profiler.py)
|
||||
|
||||
### Minimal configurations by output
|
||||
|
||||
Only collect trace output:
|
||||
@@ -124,7 +126,7 @@ profiler_config = {
|
||||
"max_iterations": 0,
|
||||
"wait_iterations": 0,
|
||||
"warmup_iterations": 0,
|
||||
"active_iterations": 0,
|
||||
"active_iterations": 1,
|
||||
}
|
||||
```
|
||||
|
||||
@@ -220,7 +222,7 @@ vllm serve Wan-AI/Wan2.2-I2V-A14B-Diffusers \
|
||||
"max_iterations": 0,
|
||||
"wait_iterations": 0,
|
||||
"warmup_iterations": 0,
|
||||
"active_iterations": 0
|
||||
"active_iterations": 1
|
||||
}'
|
||||
```
|
||||
|
||||
|
||||
@@ -0,0 +1,156 @@
|
||||
# Continuous Batching for Step-Wise Diffusion
|
||||
|
||||
!!! warning "Experimental Feature"
|
||||
This feature is experimental. It currently applies only to native
|
||||
diffusion pipelines running with `step_execution=True`.
|
||||
|
||||
This document describes the batching extension built on top of
|
||||
[Diffusion Step Execution](diffusion_step_execution.md). The base
|
||||
step-execution contract is unchanged. The batching work is mainly in the
|
||||
scheduler and runner layers.
|
||||
|
||||
## Why It Helps
|
||||
|
||||
Step-wise execution breaks a long denoise loop into scheduler-visible units.
|
||||
That gives the runtime a place to admit other compatible requests between
|
||||
steps instead of waiting for an entire request to finish.
|
||||
|
||||
This matters most in low-MFU or bursty serving scenarios:
|
||||
|
||||
- one request's denoise step may not fully saturate the GPU
|
||||
- several compatible requests can share the same denoise forward pass
|
||||
- throughput and device utilization can improve without changing request-local
|
||||
scheduler state
|
||||
|
||||
This is **not** a guaranteed single-request latency win. The main benefit is
|
||||
usually higher utilization and better throughput when the workload contains
|
||||
multiple in-flight compatible requests.
|
||||
|
||||
## Overview
|
||||
|
||||
With continuous batching enabled:
|
||||
|
||||
- the scheduler may keep multiple compatible requests active at the same time
|
||||
- the runner packs request-local step state into one `InputBatch`
|
||||
- `denoise_step()` runs on that batch
|
||||
- `step_scheduler()` and `post_decode()` still run per request
|
||||
|
||||
The current implementation is conservative:
|
||||
|
||||
- only compatible requests are batched together
|
||||
- request-mode diffusion still runs with `max_num_seqs=1`
|
||||
- per-request progress and completion remain independent
|
||||
|
||||
## Enablement
|
||||
|
||||
Use `--step-execution` as the feature gate, then increase `--max-num-seqs`
|
||||
above `1` if you want batching:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen-Image --omni \
|
||||
--port 8091 \
|
||||
--step-execution \
|
||||
--max-num-seqs 8
|
||||
```
|
||||
|
||||
`--max-num-seqs 1` keeps the step-wise path without enabling batching.
|
||||
|
||||
For a reproducible replay flow using the bundled serving benchmark, see the
|
||||
Qwen-Image replay commands in
|
||||
[`benchmarks/diffusion/README.md`](gh-file:benchmarks/diffusion/README.md)
|
||||
and
|
||||
[`benchmarks/diffusion/performance_dashboard/qwen_image_serving_performance.md`](gh-file:benchmarks/diffusion/performance_dashboard/qwen_image_serving_performance.md).
|
||||
|
||||
## Scheduler
|
||||
|
||||
The scheduler derives its batch capacity from `max_num_seqs` through
|
||||
`max_num_running_reqs`.
|
||||
|
||||
Batch admission is gated by
|
||||
[`SamplingParamsKey`](gh-file:vllm_omni/diffusion/sched/interface.py),
|
||||
which is built from shape-sensitive and CFG-sensitive sampling fields. This is
|
||||
the core correctness rule for batching: requests are only co-batched when they
|
||||
share the same denoise tensor contract.
|
||||
|
||||
There are two important details:
|
||||
|
||||
- `num_inference_steps` is not part of the key, so requests with different
|
||||
total step counts can still share a batch
|
||||
- requests also do not need to be at the same current denoise progress; active
|
||||
requests can continue batching even when their current step indices diverge
|
||||
- admission is still FIFO, so an incompatible request at the head of the
|
||||
waiting queue blocks later compatible requests
|
||||
|
||||
Today that compatibility rule is still shape-sensitive. `height`, `width`,
|
||||
`num_frames`, and CFG-related fields remain part of the key, so different
|
||||
resolutions or incompatible guidance settings do **not** co-batch yet.
|
||||
|
||||
The current batching unit is one `OmniDiffusionRequest`. Requests with
|
||||
multiple prompts do not participate in batching today.
|
||||
|
||||
## Runner
|
||||
|
||||
The runner keeps persistent per-request execution state in
|
||||
[`DiffusionRequestState`](gh-file:vllm_omni/diffusion/worker/utils.py),
|
||||
while the scheduler owns a separate lightweight request state for queueing and
|
||||
lifecycle tracking.
|
||||
|
||||
For each step, the runner builds an
|
||||
[`InputBatch`](gh-file:vllm_omni/diffusion/worker/input_batch.py) from the
|
||||
active request states:
|
||||
|
||||
- prompt embeddings and masks are normalized and padded
|
||||
- dynamic tensors such as `latents` and `timesteps` are gathered each step
|
||||
- buffers are reused when batch composition stays the same
|
||||
|
||||
The step-wise batched path is:
|
||||
|
||||
1. Run `prepare_encode()` for newly admitted requests.
|
||||
2. Build or refresh `InputBatch`.
|
||||
3. Run one batched `denoise_step(input_batch)`.
|
||||
4. Slice the batched `noise_pred` back per request.
|
||||
5. Run per-request `step_scheduler()`.
|
||||
6. Run `post_decode()` only for requests that finished denoising.
|
||||
7. Scatter updated latents back into persistent request state with
|
||||
[`scatter_latents()`](gh-file:vllm_omni/diffusion/worker/input_batch.py).
|
||||
|
||||
This keeps the shared work limited to the denoise forward pass while preserving
|
||||
request-local scheduler state and outputs.
|
||||
|
||||
## Engine
|
||||
|
||||
[`DiffusionEngine`](gh-file:vllm_omni/diffusion/diffusion_engine.py) provides
|
||||
the background loop and async add-request path needed for multiple requests to
|
||||
accumulate in the scheduler.
|
||||
|
||||
This is supporting infrastructure, not the main design point. The batching
|
||||
behavior is defined by scheduler-side compatibility gating and runner-side
|
||||
batch packing.
|
||||
|
||||
## Current Limitations
|
||||
|
||||
- Experimental feature; use `max_num_seqs=1` for the older conservative path.
|
||||
- Only native pipelines that already support `step_execution=True`.
|
||||
- Request-mode diffusion still clamps `max_num_seqs` back to `1`.
|
||||
- Only homogeneous batches keyed by `SamplingParamsKey` are supported.
|
||||
- Multi-prompt requests are not batched.
|
||||
- `cache_backend`, KV transfer, and other request-mode extras are not wired
|
||||
into the batched step-wise path yet.
|
||||
- Future work can relax the current same-shape restriction with richer
|
||||
heterogeneous batching policies such as bucketing or padded execution for
|
||||
different resolutions.
|
||||
|
||||
## Related Files
|
||||
|
||||
- Scheduler base:
|
||||
[`vllm_omni/diffusion/sched/base_scheduler.py`](gh-file:vllm_omni/diffusion/sched/base_scheduler.py)
|
||||
- Scheduler interface:
|
||||
[`vllm_omni/diffusion/sched/interface.py`](gh-file:vllm_omni/diffusion/sched/interface.py)
|
||||
- Step scheduler:
|
||||
[`vllm_omni/diffusion/sched/step_scheduler.py`](gh-file:vllm_omni/diffusion/sched/step_scheduler.py)
|
||||
- Runner:
|
||||
[`vllm_omni/diffusion/worker/diffusion_model_runner.py`](gh-file:vllm_omni/diffusion/worker/diffusion_model_runner.py)
|
||||
- Input batch:
|
||||
[`vllm_omni/diffusion/worker/input_batch.py`](gh-file:vllm_omni/diffusion/worker/input_batch.py)
|
||||
- Tests:
|
||||
[`tests/diffusion/test_diffusion_scheduler.py`](gh-file:tests/diffusion/test_diffusion_scheduler.py)
|
||||
@@ -1,4 +1,4 @@
|
||||
# Diffusion Step Execution
|
||||
# Adding Step Execution Support for Diffusion Pipelines
|
||||
|
||||
This guide documents vLLM-Omni's stepwise diffusion contract for model authors
|
||||
and contributors implementing `step_execution=True` support for a diffusion
|
||||
@@ -7,12 +7,20 @@ pipeline.
|
||||
For end-user enablement, supported models, and current limitations, see
|
||||
[Step Execution](../../user_guide/diffusion/step_execution.md).
|
||||
|
||||
This document describes the base step-execution contract only. For the
|
||||
experimental batching policy layered on top of the step-wise path, see
|
||||
[Continuous Batching for Step-Wise Diffusion](diffusion_continuous_batching.md).
|
||||
|
||||
## Current Support Scope
|
||||
|
||||
`step_execution` is **not** a generic diffusion toggle. It only works for
|
||||
pipelines that implement the segmented stateful contract in
|
||||
[`vllm_omni/diffusion/models/interface.py`](gh-file:vllm_omni/diffusion/models/interface.py).
|
||||
|
||||
This page is intentionally author-facing. Treat runtime enablement
|
||||
(`step_execution=True` in Python or `--step-execution` in serving) as an
|
||||
opt-in user knob layered on top of the implementation contract below.
|
||||
|
||||
Current in-tree support:
|
||||
|
||||
| Pipeline | Example models | Step execution |
|
||||
@@ -22,7 +30,9 @@ Current in-tree support:
|
||||
|
||||
Current engine/runtime limitations:
|
||||
|
||||
- `StepScheduler` only schedules `batch_size=1`.
|
||||
- Continuous batching with `max_num_seqs > 1` is experimental and documented in
|
||||
[Continuous Batching for Step-Wise Diffusion](diffusion_continuous_batching.md).
|
||||
Keep `max_num_seqs=1` if you want the older conservative behavior.
|
||||
- `cache_backend` is not supported in step mode.
|
||||
- Request-mode extras such as KV transfer are not wired into step mode yet.
|
||||
- Unsupported pipelines now fail early during model loading instead of failing on the first request.
|
||||
|
||||
@@ -10,6 +10,8 @@ This section contains design documents and architecture specifications for vLLM-
|
||||
|
||||
- [Disaggregated Inference](feature/disaggregated_inference.md)
|
||||
- [Ray-based Execution](feature/ray_based_execution.md)
|
||||
- [Adding Step Execution Support for Diffusion Pipelines](feature/diffusion_step_execution.md)
|
||||
- [Continuous Batching for Step-Wise Diffusion](feature/diffusion_continuous_batching.md)
|
||||
|
||||
## Module Design Documents
|
||||
|
||||
|
||||
@@ -201,7 +201,7 @@ class _BaseScheduler(SchedulerInterface):
|
||||
|
||||
- **Shared cleanup logic**: Request-id registration, finish handling, and state removal are centralized instead of duplicated in each policy.
|
||||
|
||||
- **Current constraint**: `max_num_running_reqs` remains `1` because the current engine path is still synchronous request-mode execution.
|
||||
- **Current constraint boundary**: `_BaseScheduler` derives `max_num_running_reqs` from `max_num_seqs`, but request-mode diffusion is still clamped back to `1` by the engine. The step-wise path can keep this above `1` for compatible-request batching.
|
||||
|
||||
#### 2.4 Current `RequestScheduler` Policy
|
||||
|
||||
@@ -217,7 +217,7 @@ class RequestScheduler(_BaseScheduler):
|
||||
|
||||
- **FIFO request scheduling**: Waiting requests are promoted in queue order.
|
||||
|
||||
- **Single-request admission**: The current policy only admits one active request at a time.
|
||||
- **Single-request admission**: `RequestScheduler` still admits one active request at a time because request-mode execution completes a whole request per dispatch.
|
||||
|
||||
- **Executor result feedback**: `update_from_output()` converts executor output into `FINISHED_COMPLETED` or `FINISHED_ERROR` and returns finished scheduler ids.
|
||||
|
||||
@@ -237,7 +237,7 @@ while True:
|
||||
|
||||
- **No scheduler-owned IPC**: Scheduler no longer talks to workers directly.
|
||||
|
||||
- **Conservative concurrency**: The current request-mode implementation still allows only one active request at a time.
|
||||
- **Split concurrency model**: Request-mode diffusion remains single-active-request, while the step-wise path can keep multiple compatible requests running and advance them independently between denoise steps.
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -39,13 +39,13 @@ uv pip install 'vllm-omni[demo]'
|
||||
# --8<-- [start:build-wheel-from-source]
|
||||
|
||||
#### Installation of vLLM
|
||||
If you do not need to modify source code of vLLM, you can directly install the stable 0.19.0 release version of the library
|
||||
If you do not need to modify source code of vLLM, you can directly install the stable 0.20.0 release version of the library
|
||||
|
||||
```bash
|
||||
uv pip install vllm==0.19.0 --torch-backend=auto
|
||||
uv pip install vllm==0.20.0 --torch-backend=auto
|
||||
```
|
||||
|
||||
The 0.19.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
|
||||
The 0.20.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
|
||||
|
||||
#### Installation of vLLM-Omni
|
||||
Since vllm-omni is rapidly evolving, it's recommended to install it from source
|
||||
@@ -66,12 +66,12 @@ If you want to check, modify or debug with source code of vLLM, install the libr
|
||||
```bash
|
||||
git clone https://github.com/vllm-project/vllm.git
|
||||
cd vllm
|
||||
git checkout v0.19.0
|
||||
git checkout v0.20.0
|
||||
```
|
||||
Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file.
|
||||
```bash
|
||||
#For CUDA 13.0
|
||||
export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.19.0/vllm-0.19.0+cu130-cp38-abi3-manylinux_2_35_x86_64.whl
|
||||
#For CUDA 13.0 (the default for v0.20.0; the wheel filename has no `+cu130` suffix)
|
||||
export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.20.0/vllm-0.20.0-cp38-abi3-manylinux_2_35_x86_64.whl
|
||||
```
|
||||
Install vllm with command below (If you have no existing PyTorch).
|
||||
```bash
|
||||
@@ -80,7 +80,7 @@ uv pip install --editable .
|
||||
Install vllm with command below (If you already have PyTorch).
|
||||
```bash
|
||||
python use_existing_torch.py
|
||||
uv pip install -r requirements/build.txt
|
||||
uv pip install -r requirements/build/cuda.txt
|
||||
uv pip install --no-build-isolation --editable .
|
||||
```
|
||||
</details>
|
||||
|
||||
@@ -13,7 +13,7 @@ vLLM-Omni current recommends the steps in under setup through Docker Images.
|
||||
|
||||
vLLM-Omni is built based on vLLM. Please install it with command below.
|
||||
```bash
|
||||
uv pip install vllm==0.19.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.19.0/rocm700
|
||||
uv pip install vllm==0.20.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.20.0/rocm721
|
||||
```
|
||||
|
||||
#### Installation of vLLM-Omni
|
||||
@@ -34,13 +34,13 @@ uv pip install onnxruntime-rocm
|
||||
# --8<-- [start:build-wheel-from-source]
|
||||
|
||||
#### Installation of vLLM
|
||||
If you do not need to modify source code of vLLM, you can directly install the stable 0.19.0 release version of the library
|
||||
If you do not need to modify source code of vLLM, you can directly install the stable 0.20.0 release version of the library
|
||||
|
||||
```bash
|
||||
uv pip install vllm==0.19.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.19.0/rocm700
|
||||
uv pip install vllm==0.20.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.20.0/rocm721
|
||||
```
|
||||
|
||||
The pre-built 0.19.0 vLLM wheel targets ROCm 7.0. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
|
||||
The pre-built 0.20.0 vLLM wheel targets ROCm 7.2.1. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
|
||||
|
||||
#### Installation of vLLM-Omni
|
||||
Since vllm-omni is rapidly evolving, it's recommended to install it from source
|
||||
@@ -58,7 +58,7 @@ If you want to check, modify or debug with source code of vLLM, install the libr
|
||||
```bash
|
||||
git clone https://github.com/vllm-project/vllm.git
|
||||
cd vllm
|
||||
git checkout v0.19.0
|
||||
git checkout v0.20.0
|
||||
python3 -m pip install -r requirements/rocm.txt
|
||||
python3 setup.py develop
|
||||
```
|
||||
@@ -130,7 +130,7 @@ docker run --rm \
|
||||
-v ~/.cache/huggingface:/root/.cache/huggingface \
|
||||
--env "HF_TOKEN=$HF_TOKEN" \
|
||||
-p 8091:8091 \
|
||||
vllm/vllm-omni-rocm:v0.19.0 \
|
||||
vllm/vllm-omni-rocm:v0.20.0 \
|
||||
--model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
|
||||
```
|
||||
|
||||
@@ -149,7 +149,7 @@ docker run --rm -it \
|
||||
-v ~/.cache/huggingface:/root/.cache/huggingface \
|
||||
--env "HF_TOKEN=$HF_TOKEN" \
|
||||
--entrypoint bash \
|
||||
vllm/vllm-omni-rocm:v0.19.0
|
||||
vllm/vllm-omni-rocm:v0.20.0
|
||||
```
|
||||
|
||||
# --8<-- [end:pre-built-images]
|
||||
@@ -19,10 +19,10 @@ uv venv --python 3.12 --seed
|
||||
source .venv/bin/activate
|
||||
|
||||
# On CUDA
|
||||
uv pip install vllm==0.19.0 --torch-backend=auto
|
||||
uv pip install vllm==0.20.0 --torch-backend=auto
|
||||
|
||||
# On ROCm
|
||||
uv pip install vllm==0.19.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.19.0/rocm700
|
||||
uv pip install vllm==0.20.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.20.0/rocm721
|
||||
|
||||
git clone https://github.com/vllm-project/vllm-omni.git
|
||||
cd vllm-omni
|
||||
@@ -31,6 +31,12 @@ uv pip install -e .
|
||||
|
||||
For additional installation methods — please see the [installation guide](installation/README.md).
|
||||
|
||||
|
||||
!!! note
|
||||
It is important to install the same major & minor version of vLLM and vLLM Omni, otherwise things may not work as expected. If the versions are misaligned, you will see a warning when you import vLLM Omni.
|
||||
|
||||
If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.20.0` and vLLM Omni `0.20.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.
|
||||
|
||||
## Offline Inference
|
||||
|
||||
Text-to-image generation quickstart with vLLM-Omni:
|
||||
|
||||
@@ -20,6 +20,7 @@ th {
|
||||
| `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `InternVLAA1Pipeline` | InternVLA-A1 | `InternRobotics/InternVLA-A1-3B` | ✅︎ | ✅︎ | | |
|
||||
| `HunyuanImage3ForCausalMM` | HunyuanImage3.0 (DiT-only) | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
@@ -31,6 +32,7 @@ th {
|
||||
| `ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `WanSpeechToVideoPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `LTX2Pipeline` | LTX-2-T2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
|
||||
| `LTX2ImageToVideoPipeline` | LTX-2-I2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
|
||||
@@ -52,6 +54,7 @@ th {
|
||||
| `FluxPipeline` | FLUX.1-schnell | `black-forest-labs/FLUX.1-schnell` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `AudioXPipeline` | AudioX | `zhangj1an/AudioX` | ✅︎ | ✅︎ | | |
|
||||
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
@@ -65,5 +68,5 @@ th {
|
||||
| `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | |
|
||||
| `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | |
|
||||
|`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | |
|
||||
|
||||
| `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
✅︎ indicates the model is supported on that backend. Empty cells mean not listed as supported on that backend.
|
||||
@@ -0,0 +1,338 @@
|
||||
# Audio Generate API
|
||||
|
||||
vLLM-Omni provides an API for text-to-audio generation using diffusion-based models such as Stable Audio.
|
||||
|
||||
Unlike the [Speech API](speech_api.md) which targets text-to-speech synthesis, the Audio Generate API is designed for general-purpose audio generation from text descriptions (sound effects, music, ambient soundscapes, etc.).
|
||||
|
||||
Each server instance runs a single model (specified at startup via `vllm-omni serve <model> --omni`).
|
||||
|
||||
## Quick Start
|
||||
|
||||
### Start the Server
|
||||
|
||||
```bash
|
||||
vllm-omni serve stabilityai/stable-audio-open-1.0 \
|
||||
--host 0.0.0.0 \
|
||||
--port 8091 \
|
||||
--gpu-memory-utilization 0.9 \
|
||||
--trust-remote-code \
|
||||
--enforce-eager \
|
||||
--omni
|
||||
```
|
||||
|
||||
### Generate Audio
|
||||
|
||||
**Using curl:**
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/generate \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "The sound of a cat purring",
|
||||
"audio_length": 10.0
|
||||
}' --output cat.wav
|
||||
```
|
||||
|
||||
**Using Python:**
|
||||
|
||||
```python
|
||||
import httpx
|
||||
|
||||
response = httpx.post(
|
||||
"http://localhost:8091/v1/audio/generate",
|
||||
json={
|
||||
"input": "The sound of a cat purring",
|
||||
"audio_length": 10.0,
|
||||
},
|
||||
timeout=300.0,
|
||||
)
|
||||
|
||||
with open("cat.wav", "wb") as f:
|
||||
f.write(response.content)
|
||||
```
|
||||
|
||||
## API Reference
|
||||
|
||||
### Endpoint
|
||||
|
||||
```
|
||||
POST /v1/audio/generate
|
||||
Content-Type: application/json
|
||||
```
|
||||
|
||||
### Request Parameters
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| `input` | string | **required** | Text prompt describing the audio to generate |
|
||||
| `model` | string | server's model | Model to use (optional, should match server if specified) |
|
||||
| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
|
||||
| `speed` | float | 1.0 | Playback speed (0.25 - 4.0) |
|
||||
|
||||
#### Diffusion Parameters
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| `audio_length` | float | null | Audio duration in seconds (default value is the max ~47s for `stable-audio-open-1.0`) |
|
||||
| `audio_start` | float | 0.0 | Audio start time in seconds |
|
||||
| `negative_prompt` | string | null | Text describing what to avoid in generation |
|
||||
| `guidance_scale` | float | model default | Classifier-free guidance scale (higher = more adherence to prompt) |
|
||||
| `num_inference_steps` | int | model default | Number of denoising steps (higher = better quality, slower) |
|
||||
| `seed` | int | null | Random seed for reproducible generation |
|
||||
|
||||
### Response Format
|
||||
|
||||
Returns binary audio data with the appropriate `Content-Type` header:
|
||||
|
||||
| `response_format` | Content-Type |
|
||||
|--------------------|--------------|
|
||||
| `wav` | `audio/wav` |
|
||||
| `mp3` | `audio/mpeg` |
|
||||
| `flac` | `audio/flac` |
|
||||
| `pcm` | `audio/pcm` |
|
||||
| `aac` | `audio/aac` |
|
||||
| `opus` | `audio/opus` |
|
||||
|
||||
## Examples
|
||||
|
||||
### Basic Generation
|
||||
|
||||
Generate audio with only a text prompt (model defaults for all other parameters):
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/generate \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "The sound of ocean waves crashing on a beach"
|
||||
}' --output ocean.wav
|
||||
```
|
||||
|
||||
### Custom Duration
|
||||
|
||||
Specify an explicit audio length in seconds:
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/generate \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "A dog barking",
|
||||
"audio_length": 5.0
|
||||
}' --output dog_5s.wav
|
||||
```
|
||||
|
||||
### High Quality with Negative Prompt
|
||||
|
||||
Use a negative prompt to steer generation away from undesired characteristics, and increase inference steps for higher quality:
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/generate \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "A piano playing a gentle melody",
|
||||
"audio_length": 10.0,
|
||||
"negative_prompt": "Low quality, distorted, noisy",
|
||||
"guidance_scale": 8.0,
|
||||
"num_inference_steps": 150
|
||||
}' --output piano_hq.wav
|
||||
```
|
||||
|
||||
### Reproducible Generation
|
||||
|
||||
Set a `seed` to get deterministic results across runs:
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/generate \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "Thunder and rain sounds",
|
||||
"audio_length": 15.0,
|
||||
"seed": 42
|
||||
}' --output thunder.wav
|
||||
```
|
||||
|
||||
### Full Control
|
||||
|
||||
Combine all parameters for precise control over generation:
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/generate \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "Thunder and rain sounds",
|
||||
"audio_length": 15.0,
|
||||
"negative_prompt": "Low quality",
|
||||
"guidance_scale": 7.0,
|
||||
"num_inference_steps": 100,
|
||||
"seed": 42
|
||||
}' --output thunder_rain.wav
|
||||
```
|
||||
|
||||
### Quick Generation (Fewer Steps)
|
||||
|
||||
For faster generation with slightly lower quality:
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/generate \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "Birds chirping in a forest",
|
||||
"audio_length": 8.0,
|
||||
"num_inference_steps": 50
|
||||
}' --output birds_quick.wav
|
||||
```
|
||||
|
||||
### Python Client
|
||||
|
||||
```python
|
||||
import httpx
|
||||
|
||||
response = httpx.post(
|
||||
"http://localhost:8091/v1/audio/generate",
|
||||
json={
|
||||
"input": "Thunder and rain",
|
||||
"audio_length": 15.0,
|
||||
"negative_prompt": "Low quality",
|
||||
"guidance_scale": 7.0,
|
||||
"num_inference_steps": 100,
|
||||
"seed": 42,
|
||||
"response_format": "wav",
|
||||
},
|
||||
timeout=300.0,
|
||||
)
|
||||
|
||||
with open("thunder.wav", "wb") as f:
|
||||
f.write(response.content)
|
||||
```
|
||||
|
||||
## Parameter Tuning Guide
|
||||
|
||||
### `guidance_scale`
|
||||
|
||||
Controls how closely the generated audio follows the text prompt.
|
||||
|
||||
| Range | Behaviour |
|
||||
|-------|-----------|
|
||||
| 3 - 5 | More creative / varied output |
|
||||
| 7 (default) | Balanced adherence |
|
||||
| 10+ | Strict adherence to the prompt |
|
||||
|
||||
### `num_inference_steps`
|
||||
|
||||
Controls the number of denoising steps in the diffusion process.
|
||||
|
||||
| Steps | Quality | Speed | Use Case |
|
||||
|-------|---------|-------|----------|
|
||||
| 50 | Good | Fast | Quick previews |
|
||||
| 100 | Very Good | Medium | General purpose |
|
||||
| 150+ | Excellent | Slow | Final / critical audio |
|
||||
|
||||
### `audio_length`
|
||||
|
||||
Duration of the generated audio clip. For `stable-audio-open-1.0`, the maximum is approximately 47 seconds. If omitted, the model uses its own default length.
|
||||
|
||||
### `negative_prompt`
|
||||
|
||||
Describes characteristics to avoid. Common negative prompts include:
|
||||
|
||||
- `"Low quality, distorted, noisy"`
|
||||
- `"Silence, static"`
|
||||
- `"Music"` (when generating sound effects only)
|
||||
|
||||
## Supported Models
|
||||
|
||||
| Model | Description |
|
||||
|-------|-------------|
|
||||
| `stabilityai/stable-audio-open-1.0` | Open-source audio generation model, up to ~47 seconds, 44.1 kHz stereo |
|
||||
|
||||
## Error Responses
|
||||
|
||||
### 400 Bad Request
|
||||
|
||||
Invalid or missing parameters:
|
||||
|
||||
```json
|
||||
{
|
||||
"error": {
|
||||
"message": "Audio generation model did not produce audio output.",
|
||||
"type": "BadRequestError",
|
||||
"param": null,
|
||||
"code": 400
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### 404 Not Found
|
||||
|
||||
Model mismatch:
|
||||
|
||||
```json
|
||||
{
|
||||
"error": {
|
||||
"message": "The model `xxx` does not exist.",
|
||||
"type": "NotFoundError",
|
||||
"param": "model",
|
||||
"code": 404
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
### 422 Unprocessable Entity
|
||||
|
||||
Pydantic validation failure (e.g. invalid `response_format`, `speed` out of range):
|
||||
|
||||
```json
|
||||
{
|
||||
"detail": [
|
||||
{
|
||||
"type": "literal_error",
|
||||
"msg": "Input should be 'wav', 'pcm', 'flac', 'mp3', 'aac' or 'opus'",
|
||||
...
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### "Audio generation model did not produce audio output"
|
||||
|
||||
The model finished but returned no audio data. Verify the server started successfully and the model loaded without errors.
|
||||
|
||||
### Server Not Responding
|
||||
|
||||
```bash
|
||||
# Check if the server is healthy
|
||||
curl http://localhost:8091/health
|
||||
```
|
||||
|
||||
### Audio Quality Issues
|
||||
|
||||
- Increase `num_inference_steps` (e.g. 150).
|
||||
- Add a negative prompt: `"Low quality, distorted, noisy"`.
|
||||
- Increase `guidance_scale` for stronger prompt adherence.
|
||||
|
||||
### Generation Timeout
|
||||
|
||||
- Reduce `num_inference_steps`.
|
||||
- Reduce `audio_length`.
|
||||
- Check GPU memory with `nvidia-smi`.
|
||||
|
||||
### Out of Memory
|
||||
|
||||
- Lower `--gpu-memory-utilization` (e.g. 0.8).
|
||||
- Reduce `audio_length`.
|
||||
|
||||
## Development
|
||||
|
||||
Enable debug logging:
|
||||
|
||||
```bash
|
||||
vllm-omni serve stabilityai/stable-audio-open-1.0 \
|
||||
--host 0.0.0.0 \
|
||||
--port 8091 \
|
||||
--gpu-memory-utilization 0.9 \
|
||||
--trust-remote-code \
|
||||
--enforce-eager \
|
||||
--omni \
|
||||
--uvicorn-log-level debug
|
||||
```
|
||||
@@ -358,6 +358,31 @@ curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
}' --output cloned.wav
|
||||
```
|
||||
|
||||
### Voice Storage & Caching
|
||||
|
||||
Uploaded voices are persisted to disk as a single `.safetensors` file per voice
|
||||
(audio samples + metadata — name, consent, ref_text, sample_rate, created_at —
|
||||
in the file header). On server restart the directory is scanned and all
|
||||
previously uploaded voices are restored automatically, so uploads survive
|
||||
process restarts.
|
||||
|
||||
Uploading an existing name overwrites the previous entry (a warning is logged).
|
||||
|
||||
Feature extraction artifacts (ref_code, speaker_embedding, DAC codes, etc.)
|
||||
are cached in-process with a shared LRU so repeated requests with the same
|
||||
`voice=...` skip the extraction pipeline. The cache is a true singleton across
|
||||
all TTS model types; deleting a voice invalidates every model-type slot at
|
||||
once.
|
||||
|
||||
**Configuration (environment variables):**
|
||||
|
||||
| Variable | Default | Description |
|
||||
|----------|---------|-------------|
|
||||
| `SPEAKER_SAMPLES_DIR` | `~/.cache/vllm-omni/speakers` | Directory for persisted uploaded speakers (`.safetensors` files). |
|
||||
| `SPEAKER_MAX_UPLOADED` | `1000` | Maximum number of uploaded speakers kept on disk. Upload requests past the cap return 400. |
|
||||
|
||||
The in-memory LRU has a fixed 512 MiB byte budget.
|
||||
|
||||
## Batch Speech Generation
|
||||
|
||||
The batch endpoint synthesizes multiple texts in a single request, returning all results as JSON with base64-encoded audio.
|
||||
@@ -535,7 +560,7 @@ The bundled `qwen3_tts.yaml` uses `max_num_seqs: 1` (single request) on both sta
|
||||
|-------|-------------|
|
||||
| `fishaudio/s2-pro` | 4B dual-AR TTS with DAC codec (44.1 kHz). Supports text-to-speech and voice cloning. |
|
||||
|
||||
Fish Speech uses `ref_audio` and `ref_text` for voice cloning (no `task_type` needed). The `voice` field should be set to `"default"`. See the [Fish Speech online serving example](../user_guide/examples/online_serving/fish_speech.md) for details.
|
||||
Fish Speech uses `ref_audio` and `ref_text` for voice cloning (no `task_type` needed). The `voice` field should be set to `"default"`. See the [Fish Speech section of the online TTS hub](../user_guide/examples/online_serving/text_to_speech.md#fish-speech-s2-pro) for details.
|
||||
|
||||
### Voxtral TTS
|
||||
|
||||
@@ -543,6 +568,30 @@ Fish Speech uses `ref_audio` and `ref_text` for voice cloning (no `task_type` ne
|
||||
|-------|-------------|
|
||||
| `mistralai/Voxtral-4B-TTS-2603` | 3B AR + FlowMatching TTS. Supports text-to-speech with preset voices. |
|
||||
|
||||
### CosyVoice3
|
||||
|
||||
| Model | Description |
|
||||
|-------|-------------|
|
||||
| `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | Voice cloning from `ref_audio` + `ref_text`. No built-in voice presets — upload a voice or pass `ref_audio`/`ref_text` per request. |
|
||||
|
||||
### OmniVoice
|
||||
|
||||
| Model | Description |
|
||||
|-------|-------------|
|
||||
| `k2-fsa/OmniVoice` | Pure-diffusion TTS. Supports voice cloning via `ref_audio` (with optional `ref_text`); no built-in voice presets. |
|
||||
|
||||
### VoxCPM2
|
||||
|
||||
| Model | Description |
|
||||
|-------|-------------|
|
||||
| `openbmb/VoxCPM2` | TTS + voice cloning with built-in speaker presets and uploaded-voice support. Accepts `voice` (preset or uploaded) or `ref_audio` + optional `ref_text`. |
|
||||
|
||||
### MOSS-TTS-Nano
|
||||
|
||||
| Model | Description |
|
||||
|-------|-------------|
|
||||
| `OpenMOSS-Team/MOSS-TTS-Nano` | Voice cloning only. Requires `ref_audio` (or an uploaded `voice`); no built-in voice presets. `ref_text` is accepted but ignored — upstream's `voice_clone` mode does not consume a transcript. |
|
||||
|
||||
## Error Responses
|
||||
|
||||
### 400 Bad Request
|
||||
|
||||
@@ -17,7 +17,7 @@ Both strategies use pinned memory for faster CPU-GPU transfers. The strategies a
|
||||
Model-level offloading implements mutual exclusion between DiT transformer and encoder modules using pre forward hooks:
|
||||
|
||||
- **When encoders run**: DiT transformer is offloaded to CPU
|
||||
- **When DiT runs**: Encoders are offloaded to CPU
|
||||
- **When DiT runs**: Encoders are offloaded to CPU, if more than one dit models, only one loaded on GPU, others get offloaded to CPU.
|
||||
- **VAE**: Stays resident on GPU
|
||||
|
||||
Before each module's forward pass, the hook automatically moves it to GPU while offloading the other module group to CPU. Transfers use pinned memory for speed.
|
||||
|
||||
@@ -0,0 +1,80 @@
|
||||
# Diffusion Multi-Replica Serving
|
||||
|
||||
Diffusion stages can be replicated in online serving by setting `runtime.num_replicas` on a `stage_type: diffusion` stage. Each replica owns an independent diffusion engine process and receives requests through round-robin routing with request affinity.
|
||||
|
||||
## When to Use
|
||||
|
||||
Use multi-replica serving when one diffusion process does not saturate the available GPUs, or when you want higher request throughput for many independent image/video generation requests. It improves concurrency throughput; it does not make a single request faster.
|
||||
|
||||
## Configuration
|
||||
|
||||
A replica consumes the devices required by that diffusion stage. For single-device diffusion pipelines, set `runtime.devices` to one device per replica:
|
||||
|
||||
```yaml
|
||||
stages:
|
||||
- stage_id: 0
|
||||
stage_type: diffusion
|
||||
runtime:
|
||||
devices: "0,1,2,3"
|
||||
num_replicas: 4
|
||||
engine_args:
|
||||
model: Qwen/Qwen-Image
|
||||
parallel_config:
|
||||
world_size: 1
|
||||
```
|
||||
|
||||
For multi-device diffusion pipelines, `parallel_config.world_size` is treated as the number of devices per replica. For example, four 2-GPU replicas need eight devices:
|
||||
|
||||
```yaml
|
||||
stages:
|
||||
- stage_id: 0
|
||||
stage_type: diffusion
|
||||
runtime:
|
||||
devices: "0,1,2,3,4,5,6,7"
|
||||
num_replicas: 4
|
||||
engine_args:
|
||||
model: Wan-AI/Wan2.2-T2V-A14B-Diffusers
|
||||
parallel_config:
|
||||
world_size: 2
|
||||
```
|
||||
|
||||
The runtime splits the device list into contiguous chunks per replica. With the second example, replica 0 uses `0,1`, replica 1 uses `2,3`, replica 2 uses `4,5`, and replica 3 uses `6,7`.
|
||||
|
||||
## Serving
|
||||
|
||||
Start the server with the stage config:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen-Image \
|
||||
--omni \
|
||||
--port 8099 \
|
||||
--stage-configs-path /path/to/diffusion_multi_replica.yaml
|
||||
```
|
||||
|
||||
Then send normal OpenAI-compatible diffusion chat requests. The orchestrator routes new requests across replicas and keeps updates/abort calls on the same replica for each request.
|
||||
|
||||
## Benchmarking
|
||||
|
||||
Benchmark each replica count by changing `runtime.num_replicas` and `runtime.devices`, restarting the server, then running:
|
||||
|
||||
```bash
|
||||
python benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
--base-url http://localhost:8099 \
|
||||
--model Qwen/Qwen-Image \
|
||||
--task t2i \
|
||||
--dataset random \
|
||||
--num-prompts 8 \
|
||||
--request-rate inf \
|
||||
--num-inference-steps 2 \
|
||||
--width 256 --height 256 \
|
||||
--output-file diffusion_replica_4.json
|
||||
```
|
||||
|
||||
For a quick scaling sweep, test `num_replicas` values `1`, `2`, `3`, and `4` with the same prompt count and generation parameters. Compare throughput and end-to-end latency from the JSON outputs.
|
||||
|
||||
## Notes and Limitations
|
||||
|
||||
- Multi-replica serving is currently supported for diffusion stages. Single-stage mode still rejects non-diffusion replicas.
|
||||
- `runtime.devices` must provide exactly `num_replicas * devices_per_replica` devices.
|
||||
- Request-level latency may stay similar or slightly increase because each request still runs on one replica.
|
||||
- Throughput improves only when enough concurrent requests are issued and each replica has enough GPU memory.
|
||||
@@ -1,93 +0,0 @@
|
||||
# AutoRound Quantization
|
||||
|
||||
## Overview
|
||||
|
||||
[AutoRound](https://github.com/intel/auto-round) is an advanced quantization toolkit designed for Large Language Models (LLMs), Vision-Language Models (VLMs), and diffusion models. It achieves high accuracy at ultra-low bit widths (2–4 bits) with minimal tuning by leveraging sign-gradient descent, while providing broad hardware compatibility with multi-datatype support.
|
||||
|
||||
The quantization config is auto-detected from the checkpoint's `config.json` (`quantization_config.quant_method = "auto-round"`). No extra CLI flags are needed.
|
||||
|
||||
### Supported Schemes
|
||||
|
||||
| Scheme | Bits | Status |
|
||||
|--------|------|--------|
|
||||
| W4A16 | 4 | ✅ Supported |
|
||||
| W8A16 | 8 | Planned |
|
||||
|
||||
W4A16 is the first supported scheme. Additional schemes will be added in future releases.
|
||||
|
||||
## Configuration
|
||||
|
||||
1. **Python API**: point `model` at a pre-quantized checkpoint. The quantization is detected automatically.
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
omni = Omni(model="vllm-project-org/FLUX.1-dev-AutoRound-w4a16")
|
||||
|
||||
outputs = omni.generate(
|
||||
"A cat sitting on a windowsill",
|
||||
OmniDiffusionSamplingParams(num_inference_steps=28),
|
||||
)
|
||||
outputs[0].save_images("output.png")
|
||||
```
|
||||
|
||||
2. **CLI**: pass the quantized model path directly.
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_image/text_to_image.py \
|
||||
--model vllm-project-org/FLUX.1-dev-AutoRound-w4a16 \
|
||||
--prompt "A cat sitting on a windowsill" \
|
||||
--num-inference-steps 28 \
|
||||
--output outputs/flux_w4a16.png
|
||||
```
|
||||
|
||||
No `--quantization` flag is needed — the quantization method is read from the checkpoint.
|
||||
|
||||
## How It Works
|
||||
|
||||
The checkpoint's `config.json` contains:
|
||||
|
||||
```json
|
||||
{
|
||||
"quantization_config": {
|
||||
"quant_method": "auto-round",
|
||||
"bits": 4,
|
||||
"group_size": 128,
|
||||
"sym": true,
|
||||
"packing_format": "auto_round:auto_gptq",
|
||||
"block_name_to_quantize": "transformer_blocks,single_transformer_blocks"
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
At load time:
|
||||
|
||||
1. `TransformerConfig.from_dict()` parses the `quantization_config` section and builds a vLLM `INCConfig` via `build_quant_config("auto-round", ...)`.
|
||||
2. `OmniDiffusionConfig.set_tf_model_config()` propagates the detected config to the engine.
|
||||
3. The appropriate compute kernel (e.g. GPTQ-Marlin for W4A16) is selected automatically based on the checkpoint's bit-width and packing format.
|
||||
|
||||
## Supported Models
|
||||
|
||||
| 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
|
||||
|
||||
Use [AutoRound](https://github.com/intel/auto-round) to quantize a BF16 model. The `--scheme` flag selects the quantization scheme:
|
||||
|
||||
```bash
|
||||
# W4A16 (4-bit weight, 16-bit activation)
|
||||
auto-round \
|
||||
--model black-forest-labs/FLUX.1-dev \
|
||||
--scheme W4A16 \
|
||||
--batch_size 1 \
|
||||
--disable_opt_rtn \
|
||||
--dataset coco2014 \
|
||||
--iters 0
|
||||
```
|
||||
|
||||
The output directory can be used directly as the `model` argument. See the [AutoRound documentation](https://github.com/intel/auto-round) for all available schemes and options.
|
||||
@@ -1,82 +0,0 @@
|
||||
# FP8 Quantization
|
||||
|
||||
## Overview
|
||||
|
||||
FP8 quantization converts BF16/FP16 weights to FP8 at model load time. No calibration or pre-quantized checkpoint needed.
|
||||
|
||||
Depending on the model, either all layers can be quantized, or some sensitive layers should stay in BF16. See the [per-model table](#supported-models) for which case applies.
|
||||
|
||||
Common sensitive layers in DiT-based diffusion models include **image-stream MLPs** (`img_mlp`). These are particularly vulnerable to FP8 precision loss because they process denoising latents whose dynamic range shifts significantly across timesteps, and unlike attention projections (which benefit from QK-Norm stabilization), MLPs have no built-in normalization to absorb quantization error. In deep architectures (e.g., 60+ residual blocks), small per-layer errors compound and degrade output quality. Other layers such as **attention projections** (`to_qkv`, `to_out`) and **text-stream MLPs** (`txt_mlp`) are generally more robust due to normalization or more stable input statistics.
|
||||
|
||||
## Configuration
|
||||
|
||||
1. **Python API**: set `quantization="fp8"`. To skip sensitive layers, use `quantization_config` with `ignored_layers`.
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
# All layers quantized
|
||||
omni = Omni(model="<your-model>", quantization="fp8")
|
||||
|
||||
# Skip sensitive layers
|
||||
omni = Omni(
|
||||
model="<your-model>",
|
||||
quantization_config={
|
||||
"method": "fp8",
|
||||
"ignored_layers": ["<layer-name>"],
|
||||
},
|
||||
)
|
||||
|
||||
outputs = omni.generate(
|
||||
"A cat sitting on a windowsill",
|
||||
OmniDiffusionSamplingParams(num_inference_steps=50),
|
||||
)
|
||||
```
|
||||
|
||||
2. **CLI**: pass `--quantization fp8` and optionally `--ignored-layers`.
|
||||
|
||||
```bash
|
||||
# All layers
|
||||
python text_to_image.py --model <your-model> --quantization fp8
|
||||
|
||||
# Skip sensitive layers
|
||||
python text_to_image.py --model <your-model> --quantization fp8 --ignored-layers "img_mlp"
|
||||
|
||||
# Online serving
|
||||
vllm serve <your-model> --omni --quantization fp8
|
||||
```
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| `method` | str | — | Quantization method (`"fp8"`) |
|
||||
| `ignored_layers` | list[str] | `[]` | Layer name patterns to keep in BF16 |
|
||||
| `activation_scheme` | str | `"dynamic"` | `"dynamic"` (no calibration) or `"static"` |
|
||||
| `weight_block_size` | list[int] \| None | `None` | Block size for block-wise weight quantization |
|
||||
|
||||
The available `ignored_layers` names depend on the model architecture (e.g., `to_qkv`, `to_out`, `img_mlp`, `txt_mlp`). Consult the transformer source for your target model.
|
||||
|
||||
## Supported Models
|
||||
|
||||
| Model | HF Models | Recommendation | `ignored_layers` |
|
||||
|-------|-----------|---------------|------------------|
|
||||
| Z-Image | `Tongyi-MAI/Z-Image-Turbo` | All layers | None |
|
||||
| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Skip sensitive layers | `img_mlp` |
|
||||
| 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
|
||||
|
||||
FP8 quantization can be combined with cache acceleration:
|
||||
|
||||
```python
|
||||
omni = Omni(
|
||||
model="<your-model>",
|
||||
quantization="fp8",
|
||||
cache_backend="tea_cache",
|
||||
cache_config={"rel_l1_thresh": 0.2},
|
||||
)
|
||||
```
|
||||
@@ -1,227 +0,0 @@
|
||||
# GGUF Quantization
|
||||
|
||||
## Goals
|
||||
1. Reuse vLLM quantization configs and weight loaders as much as possible.
|
||||
2. Add native GGUF support to diffusion transformers without changing model definitions.
|
||||
3. Keep user-facing knobs minimal and consistent across offline and online flows.
|
||||
|
||||
## Scope
|
||||
1. Models: Z-Image, Flux2-klein, and Qwen-Image family.
|
||||
2. Components: diffusion transformer weights, loader paths, and quantization configs.
|
||||
3. Modes: native GGUF (transformer-only weights).
|
||||
|
||||
## Architecture Overview
|
||||
1. `OmniDiffusionConfig` accepts `quantization` or `quantization_config`.
|
||||
2. Diffusion quantization wrapper (`DiffusionGGUFConfig`) produces vLLM `QuantizationConfig` objects for linear layers.
|
||||
3. `DiffusersPipelineLoader` branches on quantization method and loads either HF weights or GGUF weights for the transformer.
|
||||
4. GGUF transformer loading is routed through model-specific adapters (e.g., Flux2Klein).
|
||||
5. vLLM GGUF path uses `GGUFConfig` and `GGUFLinearMethod` for matmul.
|
||||
|
||||
## Call Chain (Offline)
|
||||
```
|
||||
CLI (examples/offline_inference/text_to_image/text_to_image.py)
|
||||
|
|
||||
v
|
||||
Omni (vllm_omni/entrypoints/__init__.py)
|
||||
|
|
||||
v
|
||||
AsyncOmniEngine
|
||||
|
|
||||
v
|
||||
Orchestrator
|
||||
|
|
||||
v
|
||||
StageDiffusionClient
|
||||
|
|
||||
v
|
||||
DiffusionWorker
|
||||
|
|
||||
v
|
||||
DiffusionModelRunner
|
||||
|
|
||||
v
|
||||
DiffusersPipelineLoader
|
||||
|
|
||||
v
|
||||
Pipeline.forward (Flux2/Qwen/Z-Image)
|
||||
|
|
||||
v
|
||||
DiffusionEngine
|
||||
|
|
||||
v
|
||||
OmniRequestOutput
|
||||
|
|
||||
v
|
||||
Client (saved PNG)
|
||||
```
|
||||
|
||||
## Call Chain (Online)
|
||||
```
|
||||
Client
|
||||
|
|
||||
| POST /v1/images/generations
|
||||
v
|
||||
APIServer (vllm_omni/entrypoints/openai/api_server.py)
|
||||
|
|
||||
v
|
||||
_generate_with_async_omni
|
||||
|
|
||||
v
|
||||
AsyncOmni
|
||||
|
|
||||
v
|
||||
AsyncOmniEngine
|
||||
|
|
||||
v
|
||||
Orchestrator
|
||||
|
|
||||
v
|
||||
StageDiffusionClient
|
||||
|
|
||||
v
|
||||
OmniRequestOutput
|
||||
|
|
||||
v
|
||||
encode_image_base64
|
||||
|
|
||||
v
|
||||
ImageGenerationResponse
|
||||
|
|
||||
v
|
||||
Client
|
||||
```
|
||||
|
||||
## Call Chain (GGUF Operator Path)
|
||||
```
|
||||
Pipeline.forward (Flux2/Qwen/Z-Image)
|
||||
|
|
||||
v
|
||||
Transformer blocks
|
||||
|
|
||||
v
|
||||
QKVParallelLinear / ColumnParallelLinear / RowParallelLinear
|
||||
|
|
||||
v
|
||||
LinearBase.forward
|
||||
|
|
||||
v
|
||||
QuantMethod.apply (GGUFLinearMethod.apply)
|
||||
|
|
||||
v
|
||||
fused_mul_mat_gguf
|
||||
|
|
||||
v
|
||||
_fused_mul_mat_gguf (custom op)
|
||||
|
|
||||
v
|
||||
ops.ggml_dequantize
|
||||
|
|
||||
v
|
||||
x @ weight.T
|
||||
```
|
||||
|
||||
## GGUF Weight Loading Path (Transformer-Only)
|
||||
1. `DiffusersPipelineLoader.load_model` detects `quantization_config.method == "gguf"`.
|
||||
2. `gguf_model` is resolved as one of: local file, `repo/file.gguf`, or `repo:quant_type`.
|
||||
3. GGUF weights are routed through adapters in `vllm_omni/diffusion/model_loader/gguf_adapters/`.
|
||||
4. Name mapping is applied per-architecture (Z-Image, Flux2Klein, Qwen-Image).
|
||||
5. GGUF weights are loaded into transformer modules, remaining non-transformer weights come from the HF checkpoint.
|
||||
|
||||
## GGUF Adapter Design
|
||||
1. `GGUFAdapter` is an abstract base class for model-specific adapters.
|
||||
2. `Flux2KleinGGUFAdapter` implements Flux2-Klein remapping + qkv split + adaLN swap.
|
||||
3. `ZImageGGUFAdapter` implements Z-Image qkv + ffn shard handling and linear qweight routing.
|
||||
4. `QwenImageGGUFAdapter` implements Qwen-Image fused attention projection remapping.
|
||||
5. `get_gguf_adapter(...)` strictly selects by model class/config; unsupported models raise an error (no fallback adapter).
|
||||
|
||||
Adapter paths:
|
||||
- Base: `vllm_omni/diffusion/model_loader/gguf_adapters/base.py`
|
||||
- Z-Image: `vllm_omni/diffusion/model_loader/gguf_adapters/z_image.py`
|
||||
- Flux2-Klein: `vllm_omni/diffusion/model_loader/gguf_adapters/flux2_klein.py`
|
||||
- Qwen-Image: `vllm_omni/diffusion/model_loader/gguf_adapters/qwen_image.py`
|
||||
|
||||
## User Usage (Offline)
|
||||
|
||||
### Baseline BF16
|
||||
```bash
|
||||
python examples/offline_inference/text_to_image/text_to_image.py \
|
||||
--model /workspace/models/black-forest-labs/FLUX.2-klein-4B \
|
||||
--prompt "a photo of a forest with mist swirling around the tree trunks. The word 'FLUX.2' is painted over it in big, red brush strokes with visible texture" \
|
||||
--height 768 \
|
||||
--width 1360 \
|
||||
--seed 42 \
|
||||
--cfg-scale 4.0 \
|
||||
--num-images-per-prompt 1 \
|
||||
--num-inference-steps 4 \
|
||||
--output outputs/flux2_klein_4b.png
|
||||
```
|
||||
|
||||
### Native GGUF (Transformer Only)
|
||||
```bash
|
||||
python examples/offline_inference/text_to_image/text_to_image.py \
|
||||
--model /workspace/models/black-forest-labs/FLUX.2-klein-4B \
|
||||
--gguf-model "/workspace/models/unsloth/FLUX.2-klein-4B-GGUF/flux-2-klein-4b-Q8_0.gguf" \
|
||||
--quantization gguf \
|
||||
--prompt "a photo of a forest with mist swirling around the tree trunks. The word 'FLUX.2' is painted over it in big, red brush strokes with visible texture" \
|
||||
--height 768 \
|
||||
--width 1360 \
|
||||
--seed 42 \
|
||||
--cfg-scale 4.0 \
|
||||
--num-images-per-prompt 1 \
|
||||
--num-inference-steps 4 \
|
||||
--output outputs/flux2_klein_4b_gguf.png
|
||||
```
|
||||
|
||||
Notes for GGUF:
|
||||
1. Many GGUF repos do not ship `model_index.json` and configs. Use the base repo for `--model` and only pass the GGUF file via `--gguf-model`.
|
||||
2. `gguf_model` supports local path, `repo/file.gguf`, or `repo:quant_type`.
|
||||
|
||||
## User Usage (Online)
|
||||
|
||||
### Start Server (Native GGUF via CLI)
|
||||
```bash
|
||||
vllm serve /workspace/models/black-forest-labs/FLUX.2-klein-4B \
|
||||
--omni \
|
||||
--port 8000 \
|
||||
--quantization-config '{"method":"gguf","gguf_model":"/workspace/models/unsloth/FLUX.2-klein-4B-GGUF/flux-2-klein-4b-Q8_0.gguf"}'
|
||||
```
|
||||
|
||||
### Online Request (Images API)
|
||||
```bash
|
||||
curl -X POST http://localhost:8000/v1/images/generations \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"prompt": "a dragon laying over the spine of the Green Mountains of Vermont",
|
||||
"size": "1024x1024",
|
||||
"seed": 42,
|
||||
"num_inference_steps": 4
|
||||
}'
|
||||
```
|
||||
|
||||
## Qwen-Image Example
|
||||
|
||||
### Offline
|
||||
```bash
|
||||
python examples/offline_inference/text_to_image/text_to_image.py \
|
||||
--model Qwen/Qwen-Image \
|
||||
--gguf-model QuantStack/Qwen-Image-GGUF/Qwen_Image-Q4_K_M.gguf \
|
||||
--quantization gguf \
|
||||
--prompt "a red paper kite hanging from a pine tree in a winter courtyard" \
|
||||
--height 1024 \
|
||||
--width 1024 \
|
||||
--seed 42 \
|
||||
--num_inference_steps 20 \
|
||||
--output outputs/qwen_image_q4km.png
|
||||
```
|
||||
|
||||
### Online
|
||||
```bash
|
||||
vllm serve Qwen/Qwen-Image \
|
||||
--omni \
|
||||
--port 8000 \
|
||||
--quantization-config '{"method":"gguf","gguf_model":"QuantStack/Qwen-Image-GGUF/Qwen_Image-Q4_K_M.gguf"}'
|
||||
```
|
||||
|
||||
Notes:
|
||||
1. The GGUF repo provides transformer weights only; the base model repo still provides tokenizer, text encoder, scheduler, and VAE.
|
||||
2. The same adapter path covers `QwenImagePipeline`, `QwenImageEditPipeline`, `QwenImageEditPlusPipeline`, and `QwenImageLayeredPipeline`.
|
||||
@@ -1,75 +0,0 @@
|
||||
# Int8 Quantization
|
||||
|
||||
## Overview
|
||||
|
||||
Int8 quantization converts BF16/FP16 weights to Int8 at model load time. No calibration or pre-quantized checkpoint needed.
|
||||
|
||||
Depending on the model, either all layers can be quantized, or some sensitive layers should stay in BF16/FP16. See the [per-model table](#supported-models) for which case applies.
|
||||
|
||||
## Configuration
|
||||
|
||||
1. **Python API**: set `quantization="int8"`. To skip sensitive layers, use `quantization_config` with `ignored_layers`.
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
# All layers quantized
|
||||
omni = Omni(model="<your-model>", quantization="int8")
|
||||
|
||||
# Skip sensitive layers
|
||||
omni = Omni(
|
||||
model="<your-model>",
|
||||
quantization_config={
|
||||
"method": "int8",
|
||||
"ignored_layers": ["<layer-name>"],
|
||||
},
|
||||
)
|
||||
|
||||
outputs = omni.generate(
|
||||
"A cat sitting on a windowsill",
|
||||
OmniDiffusionSamplingParams(num_inference_steps=50),
|
||||
)
|
||||
```
|
||||
|
||||
2. **CLI**: pass `--quantization int8` and optionally `--ignored-layers`.
|
||||
|
||||
```bash
|
||||
# All layers
|
||||
python text_to_image.py --model <your-model> --quantization int8
|
||||
|
||||
# Skip sensitive layers
|
||||
python text_to_image.py --model <your-model> --quantization int8 --ignored-layers "img_mlp"
|
||||
|
||||
# Online serving
|
||||
vllm serve <your-model> --omni --quantization int8
|
||||
```
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| `method` | str | — | Quantization method (`"int8"`) |
|
||||
| `ignored_layers` | list[str] | `[]` | Layer name patterns to keep in BF16/FP16 |
|
||||
| `activation_scheme` | str | `"dynamic"` | `"dynamic"` (no calibration) |
|
||||
|
||||
|
||||
The available `ignored_layers` names depend on the model architecture (e.g., `to_qkv`, `to_out`, `img_mlp`, `txt_mlp`). Consult the transformer source for your target model.
|
||||
|
||||
## Supported Models
|
||||
|
||||
| Model | HF Models | Recommendation | `ignored_layers` |
|
||||
|-------|-----------|---------------|------------------|
|
||||
| Z-Image | `Tongyi-MAI/Z-Image-Turbo` | All layers | None |
|
||||
| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | All layers | None |
|
||||
|
||||
## Combining with Other Features
|
||||
|
||||
Int8 quantization can be combined with cache acceleration:
|
||||
|
||||
```python
|
||||
omni = Omni(
|
||||
model="<your-model>",
|
||||
quantization="int8",
|
||||
cache_backend="tea_cache",
|
||||
cache_config={"rel_l1_thresh": 0.2},
|
||||
)
|
||||
```
|
||||
@@ -1,56 +0,0 @@
|
||||
# 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.
|
||||
@@ -1,153 +0,0 @@
|
||||
# Quantization
|
||||
|
||||
vLLM-Omni provides a unified quantization framework that supports both diffusion models (DiT, text encoders, VAEs) and multi-stage omni models (e.g., Qwen3-Omni thinker). Quantization reduces memory usage and accelerates inference.
|
||||
|
||||
## Supported Methods
|
||||
|
||||
### Diffusion Models
|
||||
|
||||
| Method | Guide | Description | Tested Models | Min GPU |
|
||||
|--------|-------|-------------|---------------|---------|
|
||||
| FP8 | [FP8](fp8.md) | FP8 W8A8, dynamic or static | Z-Image, Qwen-Image, Flux, Bagel | SM 89 (Ada) |
|
||||
| Int8 | [Int8](int8.md) | Int8 W8A8 | Z-Image, Qwen-Image | SM 89 (Ada) / Ascend NPU |
|
||||
| GGUF | [GGUF](gguf.md) | GGUF format, dequant+GEMM for N-D tensors | Z-Image, Flux | SM 60 |
|
||||
| AutoRound | [AutoRound](autoround.md) | W4A16 (pre-quantized) | Flux | SM 80 (Ampere) |
|
||||
|
||||
### Multi-stage Omni Models (Pre-quantized Checkpoints)
|
||||
|
||||
| Method | Description | Tested Models | Min GPU |
|
||||
|--------|-------------|---------------|---------|
|
||||
| ModelOpt FP8 | Pre-quantized FP8 via NVIDIA ModelOpt | Qwen3-Omni (thinker) | SM 89 (Ada/Hopper) |
|
||||
| ModelOpt NVFP4 | Pre-quantized NVFP4 via NVIDIA ModelOpt | Qwen3-Omni (experimental, quality issues) | SM 100 (Blackwell) |
|
||||
|
||||
!!! note
|
||||
AWQ, GPTQ, and BitsAndBytes are supported by vLLM's upstream quantization registry but have **not been tested** with vLLM-Omni pipelines. They may work via `build_quant_config()` but are not validated.
|
||||
|
||||
### Pre-quantized LLM Checkpoints (Multi-stage Models)
|
||||
|
||||
For multi-stage models like Qwen3-Omni, the unified quantization framework auto-detects
|
||||
pre-quantized checkpoints via `quantization_config` in the HF config. Quantization is
|
||||
automatically scoped to the thinker's `language_model` — audio encoder, vision encoder,
|
||||
talker, and code2wav remain in BF16.
|
||||
|
||||
| Format | `quant_algo` | Hardware | Example | Status |
|
||||
|--------|-------------|----------|---------|--------|
|
||||
| ModelOpt FP8 | `FP8` | Ada/Hopper (SM 89+) | `asdazd/Qwen3-Omni-30B-A3B-Instruct_modelopt_FP8` | Tested — 47% memory reduction, comparable throughput |
|
||||
| ModelOpt NVFP4 | `NVFP4` | Blackwell (SM 100+) | `shunyang90/Qwen3-Omni-30B-A3B-Instruct-NVFP4` | Experimental — loads and runs but output quality is unacceptable |
|
||||
|
||||
**Tested FP8 results (Qwen3-Omni, 1×H200):**
|
||||
|
||||
| Config | Model Memory (GiB) | Mem Reduction | Decode (tok/s) |
|
||||
|--------|-------------------|---------------|----------------|
|
||||
| BF16 baseline | 59.26 | — | 41.6 |
|
||||
| FP8 (ModelOpt) | 31.41 | 47% | 39.9 |
|
||||
|
||||
FP8 enables the full 3-stage pipeline (thinker + talker + code2wav) on a single 64GB GPU,
|
||||
which is impossible with BF16 (thinker alone requires 59.26 GiB).
|
||||
|
||||
## Quantization Scope
|
||||
|
||||
### Diffusion Models
|
||||
|
||||
When `--quantization fp8` is enabled for diffusion models:
|
||||
|
||||
| Component | What Gets Quantized | Mechanism |
|
||||
|-----------|-------------------|-----------|
|
||||
| **DiT (transformer)** | `nn.Linear` layers | vLLM W8A8 FP8 compute (Ada/Hopper) or weight-only (older GPUs) |
|
||||
| **Text encoder** | `nn.Linear` layers | FP8 weight storage, BF16 compute |
|
||||
| **VAE** | `nn.Conv2d`, `nn.Conv3d` layers | FP8 weight storage, BF16 compute |
|
||||
|
||||
### Multi-stage Omni Models
|
||||
|
||||
For pre-quantized omni model checkpoints, quantization is scoped per component:
|
||||
|
||||
| Component | Quantized? | Notes |
|
||||
|-----------|-----------|-------|
|
||||
| **Thinker (language_model)** | Yes | FP8 or NVFP4 via ModelOpt |
|
||||
| **Audio encoder** | No | Stays BF16 |
|
||||
| **Vision encoder** | No | Stays BF16 |
|
||||
| **Talker** | No | Stays BF16 |
|
||||
| **Code2Wav** | No | Stays BF16 |
|
||||
|
||||
!!! note
|
||||
Not all models support all components. See the [FP8 supported models table](fp8.md#supported-models) for per-model details.
|
||||
|
||||
## Device Compatibility for FP8
|
||||
|
||||
| Device | Example Hardware | FP8 | NVFP4 |
|
||||
|--------|-----------------|-----|-------|
|
||||
| Blackwell GPU (SM 100+) | RTX 5090, B100, B200 | Yes | Yes (native FP4 HW) |
|
||||
| Ada/Hopper GPU (SM 89+) | RTX 4090, H100, H200 | Yes (W8A8 native) | No |
|
||||
| Turing/Ampere GPU (SM 75-86) | RTX 3090, A100 | Yes (weight-only Marlin) | No |
|
||||
| Ascend NPU | Atlas 800T A2 (910B) | Not yet supported | No |
|
||||
|
||||
Kernel selection is automatic on CUDA GPUs.
|
||||
|
||||
## Device Compatibility for Int8
|
||||
|
||||
| Device Type | Generation | Example | Int8 Mode |
|
||||
|-------------|---------------|-------------------|----------|
|
||||
| NVIDIA GPU | Ada/Hopper (SM 89+) | RTX 4090, H100, H200 | Full W8A8 with native hardware |
|
||||
| Ascend NPU | Atlas A2/Atlas A3 | Atlas 800T A2/Atlas 900 A3 | Full W8A8 with native hardware |
|
||||
|
||||
## Python API
|
||||
|
||||
The `build_quant_config()` factory accepts multiple input formats:
|
||||
|
||||
```python
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
|
||||
# String
|
||||
config = build_quant_config("fp8")
|
||||
|
||||
# Dict with parameters
|
||||
config = build_quant_config({"method": "fp8", "activation_scheme": "static"})
|
||||
|
||||
# AutoRound / INC (auto-detected from checkpoint, or explicit)
|
||||
config = build_quant_config("auto-round", bits=4, group_size=128)
|
||||
|
||||
# Per-component dict
|
||||
config = build_quant_config({
|
||||
"transformer": {"method": "fp8"},
|
||||
"vae": None,
|
||||
})
|
||||
|
||||
# None / "none"
|
||||
config = build_quant_config(None) # returns None
|
||||
```
|
||||
|
||||
## Architecture
|
||||
|
||||
```
|
||||
build_quant_config(spec)
|
||||
├── str → _build_single(method) → vLLM registry / _OVERRIDES
|
||||
├── dict with "method" → _build_single(method, **kwargs)
|
||||
├── per-component dict → ComponentQuantizationConfig
|
||||
│ routes get_quant_method() by longest-prefix match
|
||||
├── QuantizationConfig → passthrough
|
||||
└── None / "none" → None
|
||||
```
|
||||
|
||||
## Migration Guide
|
||||
|
||||
### Before (v0.14.x)
|
||||
|
||||
```python
|
||||
# Old diffusion-specific API (removed)
|
||||
from vllm_omni.diffusion.quantization import get_diffusion_quant_config
|
||||
config = get_diffusion_quant_config("fp8", activation_scheme="static")
|
||||
```
|
||||
|
||||
### After (v0.16.0+)
|
||||
|
||||
```python
|
||||
# Unified API — delegates to vLLM's registry
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
config = build_quant_config({"method": "fp8", "activation_scheme": "static"})
|
||||
|
||||
# Per-component (new)
|
||||
config = build_quant_config({
|
||||
"transformer": {"method": "fp8"},
|
||||
"vae": None,
|
||||
})
|
||||
```
|
||||
@@ -8,6 +8,8 @@ implement the stepwise contract support it today.
|
||||
|
||||
## Quick Start
|
||||
|
||||
### Python API
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
@@ -25,6 +27,18 @@ outputs = omni.generate(
|
||||
)
|
||||
```
|
||||
|
||||
### Serving
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen-Image --omni \
|
||||
--port 8091 \
|
||||
--step-execution \
|
||||
--max-num-seqs 8
|
||||
```
|
||||
|
||||
For serving, `--step-execution` enables the step-wise runtime. Continuous
|
||||
batching only becomes relevant when `--max-num-seqs > 1`.
|
||||
|
||||
## Supported Pipelines
|
||||
|
||||
| Pipeline | Example models | Step execution |
|
||||
@@ -32,9 +46,17 @@ outputs = omni.generate(
|
||||
| `QwenImagePipeline` | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Yes |
|
||||
| All other diffusion pipelines | `QwenImageEditPipeline`, `QwenImageEditPlusPipeline`, `QwenImageLayeredPipeline`, GLM-Image, Wan, Flux, etc. | No |
|
||||
|
||||
!!! warning "Experimental continuous batching"
|
||||
When `--step-execution` is enabled and `max_num_seqs > 1` is configured,
|
||||
the step-wise path can batch
|
||||
compatible requests together. This is experimental. Requests with
|
||||
incompatible sampling parameters are intentionally kept in separate batches,
|
||||
and `max_num_seqs=1` remains the conservative default.
|
||||
|
||||
## Current Limitations
|
||||
|
||||
- `step_execution` currently supports `batch_size=1` only.
|
||||
- Continuous batching under `step_execution` is experimental and only batches
|
||||
compatible requests.
|
||||
- `cache_backend` is not supported together with step execution.
|
||||
- Unsupported pipelines fail early during model loading.
|
||||
- Request-mode extras such as KV transfer are not wired into step mode yet.
|
||||
@@ -45,6 +67,11 @@ Use step execution only when you specifically need the pipeline to run through
|
||||
its stepwise request state machine. For normal diffusion inference, leave it
|
||||
disabled unless your workflow depends on this mode.
|
||||
|
||||
For Qwen-Image online serving, the usual progression is:
|
||||
|
||||
- start with `--step-execution --max-num-seqs 1` if you only need the step-wise path
|
||||
- increase `--max-num-seqs` after that if you want the experimental compatible-request batching behavior
|
||||
|
||||
If you are looking for general diffusion speedups, see
|
||||
[Diffusion Features Overview](../diffusion_features.md).
|
||||
|
||||
@@ -59,3 +86,7 @@ pipeline does not support step execution.
|
||||
If you want to add step execution support to a new diffusion pipeline, see the
|
||||
implementation guide:
|
||||
[Diffusion Step Execution Design](../../design/feature/diffusion_step_execution.md).
|
||||
|
||||
If you also want that pipeline to participate in the experimental batched
|
||||
step-wise path, see:
|
||||
[Continuous Batching for Step-Wise Diffusion](../../design/feature/diffusion_continuous_batching.md).
|
||||
@@ -15,7 +15,7 @@ vLLM-Omni supports various advanced features for diffusion models:
|
||||
- Acceleration: **cache methods**, **parallelism methods**, **startup optimizations**
|
||||
- Memory optimization: **cpu offloading**, **quantization**
|
||||
- Extensions: **LoRA inference**, **frame interpolation**
|
||||
- Execution modes: **step execution**
|
||||
- Execution modes: **step execution**, **multi-replica serving**
|
||||
|
||||
## Supported Features
|
||||
|
||||
@@ -43,6 +43,7 @@ Parallelism methods distribute computation across GPUs without quality loss (mat
|
||||
| **[Tensor Parallelism](diffusion/parallelism/tensor_parallel.md)** | Shards model weights across devices | Large models that don't fit in single GPU, with 2+ GPUs |
|
||||
| **[HSDP](diffusion/parallelism/hsdp.md)** | Weight sharding via FSDP2, redistributed on-demand at runtime | Very large models (14B+) on limited VRAM, combinable with SP |
|
||||
| **[Expert Parallelism](diffusion/parallelism/expert_parallel.md)** | Shards MoE expert MLP blocks across devices | MoE diffusion models (e.g., HunyuanImage3.0) |
|
||||
| **[Multi-Replica Serving](diffusion/multi_replica_serving.md)** | Runs multiple independent diffusion stage replicas and routes requests across them | Online serving throughput with many concurrent requests |
|
||||
|
||||
#### Startup Optimization
|
||||
|
||||
@@ -59,7 +60,7 @@ Memory optimization methods help reduce GPU memory usage, enabling inference on
|
||||
| Method | Description | Best For |
|
||||
|--------|-------------|----------|
|
||||
| **[CPU Offload](diffusion/cpu_offload_diffusion.md)** | Offloads model components to CPU memory | Limited VRAM, large models on consumer GPUs |
|
||||
| **[Quantization](diffusion/quantization/overview.md)** | Reduces DiT layers from BF16 to FP8/INT8/etc. | Limited VRAM, minimal accuracy loss |
|
||||
| **[Quantization](quantization/overview.md)** | Reduces transformer stages from BF16 to FP8/INT8/etc. | Limited VRAM, minimal accuracy loss |
|
||||
| **[VAE Patch Parallelism](diffusion/parallelism/vae_patch_parallel.md)** | Distributes VAE decode tiling across GPUs | High-resolution generation with reduced VAE memory peak |
|
||||
|
||||
### Extensions
|
||||
@@ -86,9 +87,9 @@ Execution modes control how the diffusion pipeline processes denoise steps.
|
||||
|
||||
| Method | Configuration | Description | Best For |
|
||||
|--------|--------------|-------------|----------|
|
||||
| **[FP8](diffusion/quantization/fp8.md)** | `quantization="fp8"` | FP8 W8A8 on Ada/Hopper, weight-only on older GPUs | Memory reduction, inference speedup |
|
||||
| **[INT8](diffusion/quantization/int8.md)** | `quantization="int8"` | INT8 weight-only, no calibration or pre-quantized checkpoint needed | Memory reduction, broad GPU compatibility |
|
||||
| **[GGUF](diffusion/quantization/gguf.md)** | `quantization="gguf"` | Native GGUF transformer-only weights (Q4, Q8, etc.) | Memory reduction on consumer GPUs |
|
||||
| **[FP8](quantization/fp8.md)** | `quantization="fp8"` | FP8 W8A8 on validated transformer stages | Memory reduction, inference speedup |
|
||||
| **[INT8](quantization/int8.md)** | `quantization="int8"` | INT8 W8A8 on validated transformer stages | Memory reduction, broad GPU compatibility |
|
||||
| **[GGUF](quantization/gguf.md)** | `quantization="gguf"` | Native GGUF transformer-only weights (Q4, Q8, etc.) | Memory reduction on consumer GPUs |
|
||||
|
||||
## Supported Models
|
||||
|
||||
@@ -101,14 +102,14 @@ The following tables show which models support each feature:
|
||||
> Notes:
|
||||
|
||||
> 1. CPU Offload has two methods: Module-wise (default for models with DiT + text encoder) and Layerwise. The tables below show **Layerwise support** only.
|
||||
> 2. The **💾Quantization** column is collapsed for readability. See [Quantization Overview](diffusion/quantization/overview.md) for per-method (FP8, GGUF, …) and per-model support details.
|
||||
> 2. The **💾Quantization** column is collapsed for readability. See [Quantization Overview](quantization/overview.md) for per-method and per-model support details.
|
||||
|
||||
### ImageGen
|
||||
|
||||
| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
|
||||
|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:|
|
||||
| **Bagel** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **FLUX.1-dev** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
|
||||
| **FLUX.1-dev** | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
|
||||
| **FLUX.1-schnell** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
|
||||
| **FLUX.2-klein** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
|
||||
| **FLUX.1-Kontext-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
|
||||
@@ -129,6 +130,7 @@ The following tables show which models support each feature:
|
||||
| **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ✅ | ❌ | ✅ (decode) | ✅ | ❌ |
|
||||
| **ERNIE-Image** | ❌ | ✅ | ✅ | ❓ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||
|
||||
> Notes:
|
||||
> 1. Nextstep_1(T2I) does not support cache acceleration methods such as TeaCache or Cache-DiT.
|
||||
@@ -144,7 +146,7 @@ The following tables show which models support each feature:
|
||||
| **LTX-2.3** | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **Helios** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ |
|
||||
| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||
|
||||
**Frame Interpolation Support**
|
||||
|
||||
@@ -155,7 +157,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 |
|
||||
|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:|
|
||||
| **Stable-Audio-Open** | ✅ | ❌ | ❓ | ❓ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ |
|
||||
| **Stable-Audio-Open** | ✅ | ❌ | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ |
|
||||
|
||||
|
||||
## Feature Compatibility
|
||||
@@ -261,7 +263,7 @@ Measured on NVIDIA H800:
|
||||
|
||||
- **[CPU Offload Guide](diffusion/cpu_offload_diffusion.md)** - Offload model components to CPU, reduce GPU memory usage
|
||||
- **[VAE Patch Parallelism Guide](diffusion/parallelism/vae_patch_parallel.md)** - Distribute VAE decode tiling across GPUs for high-resolution images
|
||||
- **[Quantization Overview](diffusion/quantization/overview.md)** - Overview of quantization methods for diffusion models
|
||||
- **[Quantization Overview](quantization/overview.md)** - Overview of quantization methods for diffusion, multi-stage omni/TTS, and multi-stage diffusion models
|
||||
|
||||
**Extensions:**
|
||||
|
||||
|
||||
@@ -1,81 +0,0 @@
|
||||
# CosyVoice3
|
||||
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/cosyvoice3>.
|
||||
|
||||
|
||||
## Setup
|
||||
|
||||
Install dependencies:
|
||||
```
|
||||
uv pip install -e .
|
||||
```
|
||||
|
||||
> **Note:** This includes required libraries such as `soundfile`,
|
||||
> `onnxruntime`, `x-transformers`, and `einops` via
|
||||
> `requirements/common.txt` and platform-specific requirements files.
|
||||
|
||||
Download the model snapshot:
|
||||
```
|
||||
from huggingface_hub import snapshot_download
|
||||
snapshot_download('FunAudioLLM/Fun-CosyVoice3-0.5B-2512', local_dir='pretrained_models/Fun-CosyVoice3-0.5B')
|
||||
```
|
||||
|
||||
Add `config.json` in `pretrained_models/Fun-CosyVoice3-0.5B/`:
|
||||
```json
|
||||
{
|
||||
"model_type": "cosyvoice3",
|
||||
"architectures": [
|
||||
"CosyVoice3Model"
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
> **Why `config.json` is required:**
|
||||
> `AutoConfig.register("cosyvoice3", CosyVoice3Config)` only registers a class mapping.
|
||||
> The loader still needs `model_type: "cosyvoice3"` from `config.json` to select that class.
|
||||
> If no `config.json` is present, model type cannot be inferred automatically.
|
||||
> If your downloaded checkpoint already includes a valid `config.json` with
|
||||
> `model_type: "cosyvoice3"`, this manual step can be skipped.
|
||||
|
||||
Run the offline verification script:
|
||||
```
|
||||
python examples/offline_inference/cosyvoice3/verify_e2e_cosyvoice.py \
|
||||
--model pretrained_models/Fun-CosyVoice3-0.5B \
|
||||
--tokenizer pretrained_models/Fun-CosyVoice3-0.5B/CosyVoice-BlankEN
|
||||
```
|
||||
|
||||
## Implementation Overview
|
||||
|
||||
CosyVoice3 runs as a **2-stage Omni pipeline**.
|
||||
|
||||
- **Stage 0** (`talker`) converts text + prompt audio to speech tokens.
|
||||
- **Stage 1** (`code2wav`) converts speech tokens + prompt features to waveform via flow matching and HiFiGAN.
|
||||
|
||||
Key components live in `vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py`.
|
||||
|
||||
- `CosyVoice3MultiModalProcessor` builds the multimodal inputs:
|
||||
- Tokenizes `prompt` and `prompt_text`.
|
||||
- Extracts speech tokens and mel features from the prompt audio.
|
||||
- Extracts a speaker embedding.
|
||||
- `CosyVoice3Model` implements both stages:
|
||||
- Stage 0 uses `CosyVoice3LM` and outputs speech tokens + conditioning features.
|
||||
- Stage 1 runs the flow model (DiT-based CFM) and HiFiGAN to synthesize waveform.
|
||||
|
||||
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 `sync_process_input_func` (sync mode) or the
|
||||
shared-memory connector (async-chunk mode) and outputs audio.
|
||||
|
||||
## Example materials
|
||||
|
||||
??? abstract "verify_e2e_cosyvoice.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/cosyvoice3/verify_e2e_cosyvoice.py"
|
||||
``````
|
||||
@@ -1,85 +0,0 @@
|
||||
# Fish Speech S2 Pro
|
||||
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/fish_speech>.
|
||||
|
||||
|
||||
This directory contains an offline demo for running Fish Speech S2 Pro (`fishaudio/s2-pro`) with vLLM Omni. It supports text-only synthesis and voice cloning with reference audio.
|
||||
|
||||
## Model Overview
|
||||
|
||||
[Fish Speech S2 Pro](https://huggingface.co/fishaudio/s2-pro) is a 4B dual-AR text-to-speech model by FishAudio. It uses:
|
||||
|
||||
- **Slow AR**: Generates semantic tokens autoregressively (Qwen3-based backbone)
|
||||
- **Fast AR**: Predicts residual codebook tokens from semantic tokens
|
||||
- **DAC Decoder**: Converts codec codes to 44.1 kHz audio waveform
|
||||
|
||||
The model produces high-quality speech with voice cloning capabilities.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
Install the `fish-speech` package for the DAC codec:
|
||||
|
||||
```bash
|
||||
pip install fish-speech
|
||||
```
|
||||
|
||||
## 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.
|
||||
|
||||
## Quick Start
|
||||
|
||||
### Text-Only Synthesis
|
||||
|
||||
```bash
|
||||
python end2end.py --text "Hello, this is a test of the Fish Speech text to speech system."
|
||||
```
|
||||
|
||||
Generated audio files are saved to `output_audio/` by default.
|
||||
|
||||
### Voice Cloning
|
||||
|
||||
Provide a reference audio file and its transcript to clone a voice:
|
||||
|
||||
```bash
|
||||
python end2end.py \
|
||||
--text "Hello, this is a cloned voice." \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--ref-text "Transcript of the reference audio."
|
||||
```
|
||||
|
||||
## Streaming Mode
|
||||
|
||||
Add `--streaming` to stream audio chunks progressively via `AsyncOmni` (requires `async_chunk: true` in the stage config):
|
||||
|
||||
```bash
|
||||
python end2end.py --text "Hello, this is a test." --streaming
|
||||
```
|
||||
|
||||
Each DAC decoder chunk is logged as it arrives. The final WAV file is written once generation completes.
|
||||
|
||||
## Configuration
|
||||
|
||||
| Argument | Default | Description |
|
||||
|----------|---------|-------------|
|
||||
| `--text` | `"Hello, this is a test..."` | Text to synthesize |
|
||||
| `--ref-audio` | None | Path to reference audio for voice cloning |
|
||||
| `--ref-text` | None | Transcript of the reference audio |
|
||||
| `--model` | `fishaudio/s2-pro` | HuggingFace model path |
|
||||
| `--stage-configs-path` | `fish_speech_s2_pro.yaml` | Path to stage configs YAML |
|
||||
| `--output-dir` | `output_audio` | Output directory for WAV files |
|
||||
| `--streaming` | False | Enable streaming via AsyncOmni |
|
||||
| `--stage-init-timeout` | 600 | Stage initialization timeout (seconds) |
|
||||
|
||||
## Notes
|
||||
|
||||
- Output audio is 44.1 kHz mono WAV.
|
||||
- The DAC codec weights (`codec.pth`) are loaded lazily from the model directory.
|
||||
- Voice cloning encodes the reference audio using the DAC codec to extract semantic codes, then prepends them as a system message.
|
||||
|
||||
## Example materials
|
||||
|
||||
??? abstract "end2end.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/fish_speech/end2end.py"
|
||||
``````
|
||||
@@ -1,154 +1,87 @@
|
||||
# GLM-Image Multistage End-to-End Inference
|
||||
# GLM-Image Offline Inference
|
||||
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/glm_image>.
|
||||
|
||||
|
||||
This example demonstrates how to run GLM-Image with the vLLM-Omni multistage architecture.
|
||||
GLM-Image is a 2-stage image generation model (AR + Diffusion) supported by vLLM-Omni's
|
||||
declarative config system. The pipeline topology and stage structure are declared in
|
||||
`vllm_omni/model_executor/models/glm_image/pipeline.py`; deployment knobs live in
|
||||
`vllm_omni/deploy/glm_image.yaml`.
|
||||
|
||||
## 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 └─────────────────────┘ │
|
||||
│ │
|
||||
└─────────────────────────────────────────────────────────────┘
|
||||
Stage 0 (AR Model) Stage 1 (Diffusion)
|
||||
┌───────────────────┐ ┌─────────────────────┐
|
||||
│ vLLM-optimized │ prior │ GlmImagePipeline │
|
||||
│ GlmImageFor │──tokens──►│ ┌───────────────┐ │
|
||||
│ Conditional │ │ │ DiT Denoiser │ │
|
||||
│ Generation │ │ └───────┬───────┘ │
|
||||
│ (9B AR model) │ │ ▼ │
|
||||
└───────────────────┘ │ ┌───────────────┐ │
|
||||
▲ │ │ VAE Decode │──┼──► Image
|
||||
│ │ └───────────────┘ │
|
||||
Text / Image └─────────────────────┘
|
||||
Input
|
||||
```
|
||||
|
||||
## Features
|
||||
## Text-to-Image
|
||||
|
||||
- **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
|
||||
```python
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
## 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
|
||||
if __name__ == "__main__":
|
||||
omni = Omni(model="zai-org/GLM-Image")
|
||||
outputs = omni.generate(
|
||||
"A photorealistic mountain landscape at sunset",
|
||||
sampling_params={
|
||||
"height": 1024,
|
||||
"width": 1024,
|
||||
"num_inference_steps": 50,
|
||||
"guidance_scale": 1.5,
|
||||
"seed": 42,
|
||||
},
|
||||
)
|
||||
outputs[0].request_output.images[0].save("output.png")
|
||||
```
|
||||
|
||||
### Image-to-Image (Image Editing)
|
||||
## 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
|
||||
```python
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
if __name__ == "__main__":
|
||||
omni = Omni(model="zai-org/GLM-Image")
|
||||
outputs = omni.generate(
|
||||
{
|
||||
"prompt": "Convert this image to watercolor style",
|
||||
"multi_modal_data": {
|
||||
"image": "input.png",
|
||||
},
|
||||
},
|
||||
sampling_params={
|
||||
"height": 1024,
|
||||
"width": 1024,
|
||||
"num_inference_steps": 50,
|
||||
"guidance_scale": 1.5,
|
||||
"seed": 42,
|
||||
},
|
||||
)
|
||||
outputs[0].request_output.images[0].save("output.png")
|
||||
```
|
||||
|
||||
### With Custom Parameters
|
||||
## Generation 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
|
||||
```
|
||||
| Parameter | Type | Default | Description |
|
||||
| --------------------- | ----- | ------- | ----------------------------------- |
|
||||
| `height` | int | 1024 | Image height in pixels |
|
||||
| `width` | int | 1024 | Image width in pixels |
|
||||
| `num_inference_steps` | int | 50 | Number of diffusion denoising steps |
|
||||
| `guidance_scale` | float | 1.5 | Classifier-free guidance scale |
|
||||
| `seed` | int | None | Optional random seed |
|
||||
| `negative_prompt` | str | None | Negative prompt |
|
||||
|
||||
## Shell Scripts
|
||||
## VRAM Requirements
|
||||
|
||||
### 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
|
||||
```
|
||||
|
||||
## Requirements
|
||||
|
||||
- vLLM-Omni with GLM-Image support
|
||||
- CUDA-capable GPU (recommended: H100/A100 with 80GB)
|
||||
- GLM-Image model weights
|
||||
|
||||
## Example materials
|
||||
|
||||
??? abstract "end2end.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/glm_image/end2end.py"
|
||||
``````
|
||||
??? abstract "run_i2i.sh"
|
||||
``````sh
|
||||
--8<-- "examples/offline_inference/glm_image/run_i2i.sh"
|
||||
``````
|
||||
??? abstract "run_t2i.sh"
|
||||
``````sh
|
||||
--8<-- "examples/offline_inference/glm_image/run_t2i.sh"
|
||||
``````
|
||||
| Stage | VRAM |
|
||||
| :---------------- | :--------------------- |
|
||||
| Stage-0 (AR) | **~18 GiB + KV Cache** |
|
||||
| Stage-1 (DiT+VAE) | **~20 GiB** |
|
||||
| Total | **~38 GiB + KV Cache** |
|
||||
@@ -0,0 +1,205 @@
|
||||
# InternVLA-A1
|
||||
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/internvla_a1>.
|
||||
|
||||
|
||||
## Setup
|
||||
This example is adapted from: https://github.com/InternRobotics/InternVLA-A1/blob/master/tests/policies/internvla_a1_3b/open_loop_genie1_real.ipynb
|
||||
|
||||
This example runs the single-path vLLM offline inference workflow for InternVLA-A1 open-loop action prediction.
|
||||
|
||||
Before running the script, export the required local paths:
|
||||
|
||||
```bash
|
||||
# hf Jia-Zeng/InternVLA-A1-3B-FineTuned-Place_Markpen
|
||||
export INTERNVLA_A1_MODEL_DIR=/path/to/InternVLA-A1-3B-ft-pen
|
||||
# hf download InternRobotics/InternData-A1 real_lerobotv30/genie1/Genie1-Place_Markpen.tar.gz --repo-type dataset --local-dir /path/to/Genie1-Place_Markpen
|
||||
export INTERNVLA_A1_DATASET_DIR=/path/to/Genie1-Place_Markpen
|
||||
export INTERNVLA_A1_PROCESSOR_DIR=/path/to/Qwen3-VL-2B-Instruct
|
||||
# hf tenstep/Cosmos-Tokenizer-CI8x8-SafeTensors
|
||||
export INTERNVLA_A1_COSMOS_DIR=/path/to/Cosmos-Tokenizer-CI8x8-SafeTensor
|
||||
```
|
||||
|
||||
The shell entrypoint also accepts these variables as defaults, so you can keep the command line itself short.
|
||||
|
||||
`INTERNVLA_A1_COSMOS_DIR` is expected to contain:
|
||||
|
||||
- `encoder.safetensors`
|
||||
- `decoder.safetensors`
|
||||
|
||||
Reference Hugging Face repo: `tenstep/Cosmos-Tokenizer-CI8x8-SafeTensors`
|
||||
|
||||
## Run examples
|
||||
|
||||
Get into the example folder:
|
||||
|
||||
```bash
|
||||
cd examples/offline_inference/internvla_a1
|
||||
```
|
||||
|
||||
### Run one sample
|
||||
|
||||
This runs a single sample through the maintained vLLM registry path.
|
||||
|
||||
```bash
|
||||
bash run.sh --num-samples 1 --num-episodes 0
|
||||
```
|
||||
|
||||
### Run open-loop evaluation against GT
|
||||
|
||||
This runs one episode, writes evaluation metrics, and saves prediction-vs-GT plots when `matplotlib` is available.
|
||||
|
||||
```bash
|
||||
bash run.sh --num-episodes 1
|
||||
```
|
||||
|
||||
Outputs are written under:
|
||||
|
||||
```bash
|
||||
outputs/internvla_a1/vllm_infer/
|
||||
```
|
||||
|
||||
Typical files:
|
||||
|
||||
- `summary.json`: top-level run summary
|
||||
- `registry/log.json`: per-episode GT comparison metrics
|
||||
- `registry/plots/*.jpg`: prediction-vs-GT figures
|
||||
|
||||
### Eager float32 validation
|
||||
|
||||
Use eager attention and float32 when you want the most stable numerical comparison baseline.
|
||||
|
||||
```bash
|
||||
bash run.sh \
|
||||
--num-episodes 1 \
|
||||
--attn-implementation eager \
|
||||
--dtype float32
|
||||
```
|
||||
|
||||
### Optional runtime switches
|
||||
|
||||
- `--dtype {bfloat16,float32}`: choose inference dtype
|
||||
- `--attn-implementation {eager,sdpa}`: switch attention backend
|
||||
- `--enable-regional-compile`: enable regional `torch.compile`
|
||||
- `--enable-warmup`: run pipeline warmup in initialization
|
||||
- `--skip-plots`: skip plot generation even if `matplotlib` is installed
|
||||
|
||||
### Collect results and performance logs
|
||||
|
||||
Use the helper below when you want to collect one-sample output, forward latency, GT comparison, plots, timing logs, and optional pytest output into a single result directory:
|
||||
|
||||
```bash
|
||||
bash collect_results.sh
|
||||
```
|
||||
|
||||
The script writes a timestamped directory under:
|
||||
|
||||
```bash
|
||||
outputs/internvla_a1/collected_results/
|
||||
```
|
||||
|
||||
Typical result files:
|
||||
|
||||
- `env_summary.txt`: environment and path summary
|
||||
- `sample_run.log`: one-sample inference log
|
||||
- `forward_benchmark/forward_latency.json`: isolated `pipeline.forward` latency summary
|
||||
- `eval_run.log`: open-loop GT evaluation log
|
||||
- `eval_outputs/summary.json`: top-level output summary
|
||||
- `eval_outputs/registry/log.json`: GT comparison metrics
|
||||
- `eval_outputs/registry/plots/*.jpg`: prediction-vs-GT figures
|
||||
- `*_time.txt`: timing output
|
||||
- `*_gpu.csv`: sampled GPU usage
|
||||
|
||||
### Benchmark forward latency
|
||||
|
||||
Use the dedicated mode below when you want the isolated `pipeline.forward` latency instead of end-to-end script time:
|
||||
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model-dir "$INTERNVLA_A1_MODEL_DIR" \
|
||||
--dataset-dir "$INTERNVLA_A1_DATASET_DIR" \
|
||||
--benchmark-forward \
|
||||
--dtype bfloat16 \
|
||||
--attn-implementation eager \
|
||||
--warmup-iters 3 \
|
||||
--benchmark-iters 10 \
|
||||
--output-dir outputs/internvla_a1/forward_benchmark
|
||||
```
|
||||
|
||||
The output JSON is written to `outputs/internvla_a1/forward_benchmark/forward_latency.json` and contains:
|
||||
|
||||
- `cold_start_ms`
|
||||
- `warmup_summary`
|
||||
- `benchmark_summary`
|
||||
- `benchmark_samples_ms`
|
||||
|
||||
### Reference results
|
||||
|
||||
Reference run collected on `1x NVIDIA H200`, `bfloat16`, `eager`:
|
||||
|
||||
- one-sample end-to-end run: `38s`
|
||||
- one-episode GT evaluation run: `45s`
|
||||
- `average_mse = 1.7173260857816786e-05`
|
||||
- `average_mae = 0.0011860118247568607`
|
||||
- `average_mse_joint = 7.42028441891307e-06`
|
||||
- `average_mae_joint = 0.0010777723509818316`
|
||||
- `average_mse_gripper = 8.544408774469048e-05`
|
||||
- `average_mae_gripper = 0.0019436875591054559`
|
||||
|
||||
The reference plot filename is:
|
||||
|
||||
- `eval_outputs/registry/plots/vllm_registry_open_loop_ep0.jpg`
|
||||
|
||||
## Measure latency and VRAM
|
||||
|
||||
You can collect end-to-end runtime and peak host memory with:
|
||||
|
||||
```bash
|
||||
/usr/bin/time -v bash run.sh --num-episodes 1
|
||||
```
|
||||
|
||||
For GPU memory snapshots, use:
|
||||
|
||||
```bash
|
||||
nvidia-smi --query-gpu=name,memory.total,memory.used --format=csv
|
||||
```
|
||||
|
||||
Run `nvidia-smi` before and during inference to record the peak device memory for your PR report.
|
||||
|
||||
## FAQ
|
||||
|
||||
If `matplotlib` is missing, evaluation still runs and only plot generation is skipped. To enable plots:
|
||||
|
||||
```bash
|
||||
pip install matplotlib
|
||||
```
|
||||
|
||||
## Example materials
|
||||
|
||||
- `end2end.py`: main offline inference and GT evaluation entrypoint
|
||||
- `run.sh`: shell wrapper with local path env vars
|
||||
- `collect_results.sh`: helper to gather result summaries and performance logs into one directory
|
||||
- `internvla_a1_common.py`: dataset, evaluation, and plotting helpers
|
||||
|
||||
## Embedded source listings
|
||||
|
||||
??? abstract "collect_results.sh"
|
||||
``````sh
|
||||
--8<-- "examples/offline_inference/internvla_a1/collect_results.sh"
|
||||
``````
|
||||
??? abstract "end2end.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/internvla_a1/end2end.py"
|
||||
``````
|
||||
??? abstract "internvla_a1_common.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/internvla_a1/internvla_a1_common.py"
|
||||
``````
|
||||
??? abstract "run.sh"
|
||||
``````sh
|
||||
--8<-- "examples/offline_inference/internvla_a1/run.sh"
|
||||
``````
|
||||
??? abstract "standalone/outputs/qwena1/standalone_infer/log.json"
|
||||
``````json
|
||||
--8<-- "examples/offline_inference/internvla_a1/standalone/outputs/qwena1/standalone_infer/log.json"
|
||||
``````
|
||||
@@ -1,172 +0,0 @@
|
||||
# Qwen3-TTS
|
||||
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/qwen3_tts>.
|
||||
|
||||
|
||||
This directory contains an offline demo for running Qwen3 TTS models with vLLM Omni. It builds task-specific inputs and generates WAV files locally.
|
||||
|
||||
## Model Overview
|
||||
|
||||
Qwen3 TTS provides multiple task variants for speech generation:
|
||||
|
||||
- **CustomVoice**: Generate speech with a known speaker identity (speaker ID) and optional instruction.
|
||||
- **VoiceDesign**: Generate speech from text plus a descriptive instruction that designs a new voice.
|
||||
- **Base**: Voice cloning using a reference audio + reference transcript, with optional mode selection. The `ref_audio` field accepts a local file path, HTTP URL, or base64 data URL.
|
||||
|
||||
## 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.
|
||||
|
||||
### ROCm Dependencies
|
||||
|
||||
You will need to install the dependency `onnxruntime-rocm`.
|
||||
|
||||
```
|
||||
pip uninstall onnxruntime # should be removed before we can install onnxruntime-rocm
|
||||
pip install onnxruntime-rocm
|
||||
```
|
||||
|
||||
## Quick Start
|
||||
|
||||
Run a single sample for a task:
|
||||
|
||||
```
|
||||
python end2end.py --query-type CustomVoice
|
||||
```
|
||||
|
||||
Generated audio files are saved to `output_audio/` by default.
|
||||
|
||||
## Task Usage
|
||||
|
||||
### CustomVoice
|
||||
|
||||
Single sample:
|
||||
|
||||
```
|
||||
python end2end.py --query-type CustomVoice
|
||||
```
|
||||
|
||||
Batch sample (multiple prompts in one run):
|
||||
|
||||
```
|
||||
python end2end.py --query-type CustomVoice --use-batch-sample
|
||||
```
|
||||
|
||||
### VoiceDesign
|
||||
|
||||
Single sample:
|
||||
|
||||
```
|
||||
python end2end.py --query-type VoiceDesign
|
||||
```
|
||||
|
||||
Batch sample:
|
||||
|
||||
```
|
||||
python end2end.py --query-type VoiceDesign --use-batch-sample
|
||||
```
|
||||
|
||||
### Base (Voice Clone)
|
||||
|
||||
Single sample:
|
||||
|
||||
```
|
||||
python end2end.py --query-type Base
|
||||
```
|
||||
|
||||
Batch sample:
|
||||
|
||||
```
|
||||
python end2end.py --query-type Base --use-batch-sample
|
||||
```
|
||||
|
||||
Mode selection for Base:
|
||||
|
||||
- `--mode-tag icl` (default): standard mode
|
||||
- `--mode-tag xvec_only`: enable `x_vector_only_mode` in the request
|
||||
|
||||
Examples:
|
||||
|
||||
```
|
||||
python end2end.py --query-type Base --mode-tag icl
|
||||
```
|
||||
|
||||
## Voice and Language Control
|
||||
|
||||
### Supported Voices (CustomVoice)
|
||||
|
||||
Predefined speaker voices are set via the `speaker` (or `voice_type`) field in `additional_information`. Available speakers depend on the loaded checkpoint; check `talker_config.spk_id` in the model config for the full list. Common voices include `vivian`, `ryan`, `aiden`, `ethan`, `serena` (case-insensitive).
|
||||
|
||||
Pass the speaker name in your request:
|
||||
|
||||
```python
|
||||
additional_information = {
|
||||
"text": ["你好,我是通义千问"],
|
||||
"task_type": ["CustomVoice"],
|
||||
"speaker": ["Vivian"],
|
||||
"language": ["Chinese"],
|
||||
}
|
||||
```
|
||||
|
||||
### Supported Languages
|
||||
|
||||
The `language` field controls the codec-level language tag. Use `"Auto"` (default) for automatic detection.
|
||||
|
||||
Supported values: `Auto`, `Chinese`, `English`, `Japanese`, `Korean`, `German`, `French`, `Russian`, `Portuguese`, `Spanish`, `Italian`.
|
||||
|
||||
```python
|
||||
additional_information = {
|
||||
"text": ["Hello, nice to meet you."],
|
||||
"task_type": ["CustomVoice"],
|
||||
"speaker": ["Aiden"],
|
||||
"language": ["English"],
|
||||
}
|
||||
```
|
||||
|
||||
### VoiceDesign and Base
|
||||
|
||||
- **VoiceDesign**: Use `instruct` for natural-language voice description; no `speaker` needed.
|
||||
- **Base**: Use `ref_audio` and `ref_text` for voice cloning; `language` is optional.
|
||||
|
||||
## Streaming Mode
|
||||
|
||||
Add `--streaming` to stream audio chunks progressively via `AsyncOmni` (requires `async_chunk: true` in the stage config):
|
||||
|
||||
```bash
|
||||
python end2end.py --query-type CustomVoice --streaming --output-dir /tmp/out_stream
|
||||
```
|
||||
|
||||
Each Code2Wav chunk is logged as it arrives (default 25 frames; configurable via `codec_chunk_frames`
|
||||
in the stage config). The initial chunk size is dynamically selected based on server load for reduced
|
||||
TTFA, and can be overridden per-request via the `initial_codec_chunk_frames` API field. The final WAV file is written once generation
|
||||
completes. This demonstrates that audio data is available progressively rather than only at the end.
|
||||
|
||||
> **Note:** Streaming uses `AsyncOmni` internally. The non-streaming path (`Omni`) is unchanged.
|
||||
|
||||
## Batched Decoding
|
||||
|
||||
The Code2Wav stage (stage 1) supports batched decoding, where multiple requests are decoded in a single forward pass through the SpeechTokenizer. To use it, set `max_num_seqs > 1` on both stages via `--stage-overrides` and pass multiple prompts via `--txt-prompts` with a matching `--batch-size`.
|
||||
|
||||
```
|
||||
python end2end.py --query-type CustomVoice \
|
||||
--txt-prompts benchmark_prompts.txt \
|
||||
--batch-size 4 \
|
||||
--stage-overrides '{"0":{"max_num_seqs":4,"gpu_memory_utilization":0.2},"1":{"max_num_seqs":4,"gpu_memory_utilization":0.2}}'
|
||||
```
|
||||
|
||||
**Important:** `--batch-size` must match a CUDA graph capture size (1, 2, 4, 8, 16...) because the Talker's code predictor KV cache is sized to `max_num_seqs`, and CUDA graphs pad the batch to the next capture size. Both stages need `max_num_seqs >= batch_size` in the stage config for batching to take effect. If only stage 1 has a higher `max_num_seqs`, it won't help — stage 1 can only batch chunks from requests that are in-flight simultaneously, which requires stage 0 to also process multiple requests concurrently.
|
||||
|
||||
## Notes
|
||||
|
||||
- The script uses the model paths embedded in `end2end.py`. Update them if your local cache path differs.
|
||||
- Use `--output-dir` to change the output folder.
|
||||
|
||||
## Example materials
|
||||
|
||||
??? abstract "benchmark_prompts.txt"
|
||||
``````txt
|
||||
--8<-- "examples/offline_inference/qwen3_tts/benchmark_prompts.txt"
|
||||
``````
|
||||
??? abstract "end2end.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/qwen3_tts/end2end.py"
|
||||
``````
|
||||
@@ -57,6 +57,8 @@ Key arguments:
|
||||
- `--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.
|
||||
- `--enable-cpu-offload`: enabling model-wise offloading to save gpu memory
|
||||
- `--enable-layerwise-offload`: enabling layerwise offloading to save gpu memory
|
||||
|
||||
## Example materials
|
||||
|
||||
|
||||
@@ -0,0 +1,395 @@
|
||||
# Text-To-Speech
|
||||
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_speech>.
|
||||
|
||||
|
||||
vLLM-Omni supports several autoregressive TTS models. They share a common
|
||||
CLI shape (`--text`, `--ref-audio`, `--ref-text`, `--output-dir`) and live
|
||||
together in this hub. Each model has its own subdirectory containing a
|
||||
single `end2end.py` script; this README is the single doc entry point.
|
||||
|
||||
For online serving, see `examples/online_serving/<model>/`. For the full
|
||||
list of supported architectures across all modalities, see
|
||||
[Supported Models](https://github.com/vllm-project/vllm-omni/tree/main/docs/models/supported_models.md).
|
||||
|
||||
## Supported Models
|
||||
|
||||
| Model | HuggingFace repo | Stages | Voice cloning | Streaming | Special modes | Sample rate |
|
||||
|---|---|---|---|---|---|---|
|
||||
| VoxCPM2 | `openbmb/VoxCPM2` | single (native AR) | ✓ | — | continuation (`--ref-audio` + `--ref-text`) | 48 kHz |
|
||||
| VoxCPM | local model dir | split | ✓ | ✓ (`voxcpm_async_chunk.yaml`) | — | 24 kHz |
|
||||
| CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ (`async_chunk: true` default) | — | 22.05 kHz |
|
||||
| Fish Speech S2 Pro | `fishaudio/s2-pro` | dual-AR | ✓ | ✓ (`--streaming`) | — | 44.1 kHz |
|
||||
| OmniVoice | `k2-fsa/OmniVoice` | 2 (gen + dec) | ✓ | — | voice design (`--instruct`), language (`--lang`) | 24 kHz |
|
||||
| Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | 2 (talker + code2wav) | ✓ (Base) | ✓ | 3 task variants (`--query-type`) | 24 kHz |
|
||||
| Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | varies | ✓ | ✓ | voice presets (`--voice`) | 24 kHz |
|
||||
|
||||
## Common Quick Start
|
||||
|
||||
Most models share this invocation shape:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/<model>/end2end.py \
|
||||
--text "Hello, this is a test." \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--ref-text "Transcript of the reference audio."
|
||||
```
|
||||
|
||||
`--ref-audio` and `--ref-text` are optional (text-only synthesis works without them) and must be provided together for voice cloning. The exotic scripts — Qwen3-TTS, Voxtral TTS, CosyVoice3 — accept additional model-specific flags documented in their per-model section below.
|
||||
|
||||
---
|
||||
|
||||
## VoxCPM2
|
||||
|
||||
Single-stage native AR TTS at 48 kHz. Pipeline: `feat_encoder → MiniCPM4 → FSQ → residual_lm → LocDiT → AudioVAE`.
|
||||
|
||||
### Prerequisites
|
||||
```bash
|
||||
pip install voxcpm
|
||||
# or, for a local source checkout:
|
||||
export VLLM_OMNI_VOXCPM_CODE_PATH=/path/to/voxcpm
|
||||
```
|
||||
|
||||
### Quick start
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/voxcpm2/end2end.py \
|
||||
--model openbmb/VoxCPM2 \
|
||||
--text "Hello, this is a VoxCPM2 demo."
|
||||
```
|
||||
|
||||
### Voice cloning
|
||||
Pass a reference audio for isolated cloning, or both `--ref-audio` + `--ref-text` for prompt continuation:
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/voxcpm2/end2end.py \
|
||||
--text "Hello, this is a voice clone demo." \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--ref-text "Transcript of the reference audio."
|
||||
```
|
||||
|
||||
### Notes
|
||||
- Output: 48 kHz mono WAV.
|
||||
- Stage config: `vllm_omni/model_executor/stage_configs/voxcpm2.yaml` (default).
|
||||
|
||||
---
|
||||
|
||||
## VoxCPM
|
||||
|
||||
Split-stage TTS. The hub example covers single TTS, single voice cloning, and streaming. Use `benchmarks/voxcpm/` for warmup, batch JSONL prompts, profiler injection, and offline TTFP / RTF measurement.
|
||||
|
||||
### Prerequisites
|
||||
```bash
|
||||
pip install voxcpm soundfile
|
||||
# or use a local source tree:
|
||||
export VLLM_OMNI_VOXCPM_CODE_PATH=/path/to/VoxCPM/src
|
||||
```
|
||||
|
||||
If the native VoxCPM `config.json` does not contain HF metadata such as `model_type`, prepare a persistent HF-compatible config directory and point the stage configs to it via `VLLM_OMNI_VOXCPM_HF_CONFIG_PATH`:
|
||||
|
||||
```bash
|
||||
export VOXCPM_MODEL=/path/to/voxcpm-model
|
||||
export VLLM_OMNI_VOXCPM_HF_CONFIG_PATH=/tmp/voxcpm_hf_config
|
||||
mkdir -p "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH"
|
||||
cp "$VOXCPM_MODEL/config.json" "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH/config.json"
|
||||
cp "$VOXCPM_MODEL/generation_config.json" "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH/generation_config.json" 2>/dev/null || true
|
||||
python3 -c 'import json, os; p=os.path.join(os.environ["VLLM_OMNI_VOXCPM_HF_CONFIG_PATH"], "config.json"); cfg=json.load(open(p, "r", encoding="utf-8")); cfg["model_type"]="voxcpm"; cfg.setdefault("architectures", ["VoxCPMForConditionalGeneration"]); json.dump(cfg, open(p, "w", encoding="utf-8"), indent=2, ensure_ascii=False)'
|
||||
```
|
||||
|
||||
### Quick start
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/voxcpm/end2end.py \
|
||||
--model "$VOXCPM_MODEL" \
|
||||
--text "This is a split-stage VoxCPM synthesis example running on vLLM Omni."
|
||||
```
|
||||
|
||||
### Voice cloning
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/voxcpm/end2end.py \
|
||||
--model "$VOXCPM_MODEL" \
|
||||
--text "This sentence is synthesized with a cloned voice." \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--ref-text "The exact transcript spoken in reference.wav."
|
||||
```
|
||||
|
||||
### Streaming
|
||||
Pass the async-chunk stage config:
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/voxcpm/end2end.py \
|
||||
--model "$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."
|
||||
```
|
||||
|
||||
### Notes
|
||||
- `voxcpm.yaml` is the default non-streaming stage config; `voxcpm_async_chunk.yaml` enables streaming.
|
||||
- Streaming is currently single-request oriented.
|
||||
- `--ref-text` must be the real transcript of `--ref-audio`; mismatched text degrades quality.
|
||||
- For online serving, see [`examples/online_serving/voxcpm`](https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/voxcpm/README.md). For benchmark reporting, see [`benchmarks/voxcpm`](https://github.com/vllm-project/vllm-omni/tree/main/benchmarks/voxcpm/README.md).
|
||||
|
||||
---
|
||||
|
||||
## CosyVoice3
|
||||
|
||||
2-stage TTS pipeline (`talker` + `code2wav`) at 22.05 kHz.
|
||||
|
||||
### Prerequisites
|
||||
```bash
|
||||
uv pip install -e .
|
||||
# Includes soundfile, onnxruntime, x-transformers, einops via requirements.
|
||||
```
|
||||
|
||||
Download the model snapshot:
|
||||
```python
|
||||
from huggingface_hub import snapshot_download
|
||||
snapshot_download('FunAudioLLM/Fun-CosyVoice3-0.5B-2512',
|
||||
local_dir='pretrained_models/Fun-CosyVoice3-0.5B')
|
||||
```
|
||||
|
||||
If your downloaded checkpoint lacks `config.json`, add it:
|
||||
```json
|
||||
{
|
||||
"model_type": "cosyvoice3",
|
||||
"architectures": ["CosyVoice3Model"]
|
||||
}
|
||||
```
|
||||
This is required because `AutoConfig.register("cosyvoice3", CosyVoice3Config)` only registers the class mapping; the loader still reads `model_type` from `config.json` to select the class.
|
||||
|
||||
### Quick start
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/cosyvoice3/end2end.py \
|
||||
--model pretrained_models/Fun-CosyVoice3-0.5B \
|
||||
--tokenizer pretrained_models/Fun-CosyVoice3-0.5B/CosyVoice-BlankEN
|
||||
```
|
||||
|
||||
### Voice cloning
|
||||
Pass a reference audio. Note that CosyVoice3's `--prompt-text` is a system-style prompt for the GPT stage, not a reference transcript:
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/cosyvoice3/end2end.py \
|
||||
--model pretrained_models/Fun-CosyVoice3-0.5B \
|
||||
--tokenizer pretrained_models/Fun-CosyVoice3-0.5B/CosyVoice-BlankEN \
|
||||
--ref-audio prompt.wav \
|
||||
--prompt-text "You are a helpful assistant.<|endofprompt|>Testing my voices. Why should I not?"
|
||||
```
|
||||
|
||||
### Notes
|
||||
- Stage 0 (`talker`) emits speech tokens; stage 1 (`code2wav`) runs flow matching + HiFiGAN to synthesize waveform.
|
||||
- Deploy config auto-loads from `vllm_omni/deploy/cosyvoice3.yaml` based on HF `model_type`. Pass `--deploy-config <path>` to override.
|
||||
- `async_chunk: true` is the default; pass `--no-async-chunk` to switch to the legacy synchronous path.
|
||||
|
||||
---
|
||||
|
||||
## Fish Speech S2 Pro
|
||||
|
||||
4B dual-AR text-to-speech model from FishAudio with the DAC codec at 44.1 kHz.
|
||||
|
||||
### Prerequisites
|
||||
```bash
|
||||
pip install fish-speech
|
||||
```
|
||||
|
||||
### Quick start
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/fish_speech/end2end.py \
|
||||
--text "Hello, this is a test of the Fish Speech text to speech system."
|
||||
```
|
||||
|
||||
### Voice cloning
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/fish_speech/end2end.py \
|
||||
--text "Hello, this is a cloned voice." \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--ref-text "Transcript of the reference audio."
|
||||
```
|
||||
|
||||
### Streaming
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/fish_speech/end2end.py \
|
||||
--text "Hello, this is a streaming test." \
|
||||
--streaming
|
||||
```
|
||||
Streaming requires `async_chunk: true` in the stage config.
|
||||
|
||||
### Notes
|
||||
- Output: 44.1 kHz mono WAV.
|
||||
- DAC codec weights (`codec.pth`) are loaded lazily from the model directory.
|
||||
|
||||
---
|
||||
|
||||
## OmniVoice
|
||||
|
||||
Zero-shot multilingual TTS supporting 600+ languages, with three modes (auto / clone / design).
|
||||
|
||||
### Prerequisites
|
||||
```bash
|
||||
huggingface-cli download k2-fsa/OmniVoice
|
||||
```
|
||||
Voice cloning requires `transformers>=5.3.0`. Auto and design modes work with `transformers>=4.57.0`.
|
||||
|
||||
### Quick start (auto voice)
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/omnivoice/end2end.py \
|
||||
--model k2-fsa/OmniVoice \
|
||||
--text "Hello, this is a test."
|
||||
```
|
||||
|
||||
### Voice cloning
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/omnivoice/end2end.py \
|
||||
--model k2-fsa/OmniVoice \
|
||||
--text "Hello, this is a test." \
|
||||
--ref-audio ref.wav \
|
||||
--ref-text "This is the reference transcription."
|
||||
```
|
||||
|
||||
### Voice design
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/omnivoice/end2end.py \
|
||||
--model k2-fsa/OmniVoice \
|
||||
--text "Hello, this is a test." \
|
||||
--instruct "female, low pitch, british accent"
|
||||
```
|
||||
|
||||
### Language hint
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/omnivoice/end2end.py \
|
||||
--model k2-fsa/OmniVoice \
|
||||
--text "你好,这是一个测试。" \
|
||||
--lang zh
|
||||
```
|
||||
|
||||
### Notes
|
||||
- Stage 0 (Generator): Qwen3-0.6B with 32-step iterative unmasking.
|
||||
- Stage 1 (Decoder): HiggsAudioV2 RVQ + DAC at 24 kHz.
|
||||
|
||||
---
|
||||
|
||||
## Qwen3-TTS
|
||||
|
||||
3-task-variant TTS with 24 kHz output. Has its own argparse surface (this script does not follow the common `--text` / `--ref-audio` shape).
|
||||
|
||||
### Prerequisites
|
||||
For ROCm builds, replace `onnxruntime` with `onnxruntime-rocm`:
|
||||
```bash
|
||||
pip uninstall onnxruntime
|
||||
pip install onnxruntime-rocm
|
||||
```
|
||||
|
||||
### Task variants
|
||||
- `CustomVoice`: predefined speaker (speaker ID) with optional style instruction.
|
||||
- `VoiceDesign`: text + descriptive instruction designs a new voice.
|
||||
- `Base`: voice cloning from reference audio + transcript.
|
||||
|
||||
```bash
|
||||
# Single sample
|
||||
python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py --query-type CustomVoice
|
||||
python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py --query-type VoiceDesign
|
||||
python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py --query-type Base
|
||||
|
||||
# Base variant has an additional mode flag:
|
||||
python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py --query-type Base --mode-tag icl # default
|
||||
python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py --query-type Base --mode-tag xvec_only # x_vector_only_mode
|
||||
|
||||
# Batch (multiple prompts in one run)
|
||||
python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py --query-type CustomVoice --use-batch-sample
|
||||
```
|
||||
|
||||
### Streaming
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \
|
||||
--query-type CustomVoice \
|
||||
--streaming \
|
||||
--output-dir /tmp/out_stream
|
||||
```
|
||||
Streaming requires `async_chunk: true` in the stage config.
|
||||
|
||||
### Batched decoding
|
||||
The Code2Wav stage supports batched decoding through the SpeechTokenizer. Configure both stages with `max_num_seqs > 1` via `--stage-overrides` and pass multiple prompts via `--txt-prompts`:
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \
|
||||
--query-type CustomVoice \
|
||||
--txt-prompts examples/offline_inference/text_to_speech/qwen3_tts/benchmark_prompts.txt \
|
||||
--batch-size 4 \
|
||||
--stage-overrides '{"0":{"max_num_seqs":4,"gpu_memory_utilization":0.2},"1":{"max_num_seqs":4,"gpu_memory_utilization":0.2}}'
|
||||
```
|
||||
`--batch-size` must match a CUDA-graph capture size (1, 2, 4, 8, 16…).
|
||||
|
||||
### Notes
|
||||
- Run `--help` for the full argument surface.
|
||||
- See `qwen3_tts/end2end.py` for the prompt-length-estimation logic the Talker uses.
|
||||
|
||||
---
|
||||
|
||||
## Voxtral TTS
|
||||
|
||||
Voxtral-4B-TTS (Mistral). Has its own argparse surface; uses voice presets and the `mistral_common` `SpeechRequest` protocol.
|
||||
|
||||
### Prerequisites
|
||||
Latest `mistral_common` with `SpeechRequest` support:
|
||||
```bash
|
||||
pip install -e /path/to/mistral-common # or upgrade from PyPI when available
|
||||
```
|
||||
|
||||
### Quick start (voice preset)
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/voxtral_tts/end2end.py \
|
||||
--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?"
|
||||
```
|
||||
|
||||
### Voice cloning (capability gated upstream)
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/voxtral_tts/end2end.py \
|
||||
--write-audio \
|
||||
--model mistralai/Voxtral-4B-TTS-2603 \
|
||||
--text "This is a test message." \
|
||||
--ref-audio path/to/reference_audio.wav
|
||||
```
|
||||
|
||||
### Streaming + concurrency
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/voxtral_tts/end2end.py \
|
||||
--num-prompts 32 --concurrency 8 --streaming --write-audio --voice neutral_female \
|
||||
--model mistralai/Voxtral-4B-TTS-2603 \
|
||||
--text "..."
|
||||
```
|
||||
Available voice presets are listed on the HF model card (`mistralai/Voxtral-4B-TTS-2603`).
|
||||
|
||||
### Notes
|
||||
- `--num-prompts N` replicates the prompt for performance measurement.
|
||||
- `--concurrency M` requires `--streaming` and must evenly divide `--num-prompts`.
|
||||
- Run `--help` for the full argument surface.
|
||||
|
||||
## Example materials
|
||||
|
||||
??? abstract "cosyvoice3/end2end.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/text_to_speech/cosyvoice3/end2end.py"
|
||||
``````
|
||||
??? abstract "fish_speech/end2end.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/text_to_speech/fish_speech/end2end.py"
|
||||
``````
|
||||
??? abstract "omnivoice/end2end.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/text_to_speech/omnivoice/end2end.py"
|
||||
``````
|
||||
??? abstract "qwen3_tts/benchmark_prompts.txt"
|
||||
``````txt
|
||||
--8<-- "examples/offline_inference/text_to_speech/qwen3_tts/benchmark_prompts.txt"
|
||||
``````
|
||||
??? abstract "qwen3_tts/end2end.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/text_to_speech/qwen3_tts/end2end.py"
|
||||
``````
|
||||
??? abstract "voxcpm/end2end.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/text_to_speech/voxcpm/end2end.py"
|
||||
``````
|
||||
??? abstract "voxcpm2/end2end.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/text_to_speech/voxcpm2/end2end.py"
|
||||
``````
|
||||
??? abstract "voxtral_tts/end2end.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/text_to_speech/voxtral_tts/end2end.py"
|
||||
``````
|
||||
@@ -1,69 +0,0 @@
|
||||
# Voxtral TTS Offline Inference
|
||||
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/voxtral_tts>.
|
||||
|
||||
|
||||
`end2end.py` runs Voxtral TTS end-to-end offline inference using vLLM. It supports both blocking (`Omni`) and streaming (`AsyncOmni`) generation, batched prompts with configurable concurrency, and voice selection via preset name or reference audio file.
|
||||
|
||||
When `mistral_common` has `SpeechRequest` support, prompt token IDs are built via `encode_speech_request`. Otherwise, the script falls back to manual token construction.
|
||||
|
||||
## Usage Examples
|
||||
|
||||
|
||||
```bash
|
||||
# Basic single-prompt with cheerful_female voice preset
|
||||
python3 examples/offline_inference/voxtral_tts/end2end.py \
|
||||
--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 \
|
||||
--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 \
|
||||
--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 \
|
||||
--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?"
|
||||
|
||||
# Short debug prompt with reference audio
|
||||
# Note: Reference audio capability is not yet released.
|
||||
python3 examples/offline_inference/voxtral_tts/end2end.py \
|
||||
--write-audio \
|
||||
--model mistralai/Voxtral-4B-TTS-2603 \
|
||||
--text "This is a test message." \
|
||||
--audio-path path/to/reference_audio.wav
|
||||
```
|
||||
|
||||
## Arguments
|
||||
|
||||
| Argument | Description |
|
||||
|---|---|
|
||||
| `--model PATH` | HuggingFace repo ID or local directory path (default: `mistralai/Voxtral-4B-TTS-2603`) |
|
||||
| `--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`) |
|
||||
| `--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`) |
|
||||
| `--voice NAME` | Voice preset to use instead of reference audio file (e.g., casual_female, casual_male, cheerful_female, neutral_female, neutral_male) |
|
||||
| `--write-audio` | Write generated audio to WAV files |
|
||||
| `--profiling-mode` | Enable profiling mode (reduces max tokens to 50) |
|
||||
| `--log-stats` | Enable detailed statistics logging |
|
||||
|
||||
## Example materials
|
||||
|
||||
??? abstract "end2end.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/voxtral_tts/end2end.py"
|
||||
``````
|
||||
@@ -47,6 +47,24 @@ python x_to_video_audio.py \
|
||||
```
|
||||
In the current test scenario (2 images + 2 audio inputs), the VRAM requirement is 72GB, regardless of whether cfg-parallel is enabled or disabled.
|
||||
The VRAM usage can be reduced by enabling CPU offload via --enable-cpu-offload.
|
||||
For multi-GPU memory reduction on the fused DreamID-Omni transformer, you can also enable HSDP:
|
||||
|
||||
```python
|
||||
python x_to_video_audio.py \
|
||||
--model /path/to/dreamid_omni \
|
||||
--prompt "Two people walking together and singing happily" \
|
||||
--image-path ./example0.png ./example1.png \
|
||||
--audio-path ./example0.wav ./example1.wav \
|
||||
--video-negative-prompt "jitter, bad hands, blur, distortion" \
|
||||
--audio-negative-prompt "robotic, muffled, echo, distorted" \
|
||||
--cfg-parallel-size 2 \
|
||||
--num-inference-steps 45 \
|
||||
--height 704 \
|
||||
--width 1280 \
|
||||
--use-hsdp \
|
||||
--hsdp-shard-size 2 \
|
||||
--output out_dreamid_omni_twoip.mp4
|
||||
```
|
||||
|
||||
|
||||
You could take reference images/audios from the test cases in the official repo: https://github.com/Guoxu1233/DreamID-Omni
|
||||
@@ -77,6 +95,9 @@ Key arguments:
|
||||
- `--image-path`: path to the input image list.
|
||||
- `--audio-path`: path to the input audio list, indicate the timbre of the output video.
|
||||
- `--cfg-parallel-size`: number of parallel cfg parallel (defaults 1).
|
||||
- `--use-hsdp`: enable HSDP weight sharding for DreamID-Omni fused blocks.
|
||||
- `--hsdp-shard-size`: number of GPUs used for HSDP sharding.
|
||||
- `--hsdp-replicate-size`: number of HSDP replica groups.
|
||||
- `--num-inference-steps`: number of denoising steps (defaults 45).
|
||||
- `--video-negative-prompt`: negative prompt for video generation.
|
||||
- `--audio-negative-prompt`: negative prompt for audio generation.
|
||||
|
||||
@@ -1,18 +1,14 @@
|
||||
# Diffusers Backend Adapter Example
|
||||
# Diffusers Backend Adapter
|
||||
|
||||
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.
|
||||
vLLM-Omni supports running diffusion models with the diffusers backend, directly serving any 🤗 Diffusers pipeline online without implementing them natively.
|
||||
|
||||
## Limitations
|
||||
|
||||
The diffusers backend is a black-box adapter. The following features are NOT yet supported.
|
||||
The diffusers backend is a black-box adapter. Its primary focus is to serve diffusion models online.
|
||||
Currently, the following features are NOT yet supported.
|
||||
It is not guaranteed whether they will be supported in the future.
|
||||
|
||||
- CFG parallel execution
|
||||
@@ -22,72 +18,87 @@ It is not guaranteed whether they will be supported in the future.
|
||||
|
||||
For these features, it is recommended to use natively supported pipelines instead.
|
||||
|
||||
## Usage
|
||||
## Model Support
|
||||
|
||||
### Option 1: CLI arguments
|
||||
Any model loadable via `DiffusionPipeline.from_pretrained()` should run, including text-to-image, image-to-image, text-to-video, image-to-video, and text-to-audio.
|
||||
|
||||
However, as we strive to ensure output similarity between vLLM-Omni's diffuser backend and plain diffusers library, the following models are particularly verified:
|
||||
|
||||
- Qwen/Qwen-Image
|
||||
- Tongyi-MAI/Z-Image-Turbo
|
||||
- Wan2.2-I2V-A14B-Diffusers
|
||||
|
||||
If you find that a model not listed above also produces different outputs from running diffusers model directly.
|
||||
Please consider file an issue.
|
||||
|
||||
## Usage
|
||||
|
||||
```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}'
|
||||
--diffusion-load-format diffusers
|
||||
```
|
||||
|
||||
`--diffusers-load-kwargs` and `--diffusers-call-kwargs` are only valid together with `--diffusion-load-format diffusers`.
|
||||
Users turn on the diffusers backend primarily through `--diffusion-load-format diffusers` argument.
|
||||
There are two more optional arguments, `--diffusers-load-kwargs` and `--diffusers-call-kwargs`,
|
||||
which 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`.
|
||||
After launching the model, users send a request as usual. 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`**:
|
||||
### `--diffusers-load-kwargs`
|
||||
|
||||
### `diffusion_load_format: "diffusers"`
|
||||
Passed as-is to `DiffusionPipeline.from_pretrained()`.
|
||||
|
||||
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`).
|
||||
This is suitable for model-specific configurations not available through the vLLM-Omni interface.
|
||||
For example: `--diffusers-load-kwargs '{"use_safetensors": true}'`.
|
||||
|
||||
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`
|
||||
### `--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).
|
||||
This is suitable for sampling parameters not available through the vLLM-Omni interface (such as 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
|
||||
### Attention Backends
|
||||
|
||||
??? abstract "stage_config.yaml"
|
||||
``````yaml
|
||||
--8<-- "examples/online_serving/diffusers_pipeline_adapter/stage_config.yaml"
|
||||
``````
|
||||
The diffusers backend converts
|
||||
[vLLM-Omni standard of attention backend setting](../../../docs/user_guide/diffusion/attention_backends.md)
|
||||
to [diffusers standard](https://huggingface.co/docs/diffusers/optimization/attention_backends#available-backends).
|
||||
|
||||
Specifically for `FLASH_ATTN`, it will first attempt to use FlashAttention-3 and then FlashAttention-2.
|
||||
|
||||
For each attempted version of `FLASH_ATTN` and `SAGE_ATTN`, it will first try to load the attention backend from HuggingFace `kernels` library, then without.
|
||||
|
||||
For unsuccessful attention selection or `TORCH_SDPA`, it will use the PyTorch's default attention backend.
|
||||
|
||||
The loaded attention backend and the failed attempts (if any) are logged to console.
|
||||
|
||||
### Model Specific Settings
|
||||
|
||||
The model loading and inference strictly follows the diffusers library, and they may be different from vLLM-Omni's native interface for some specific models.
|
||||
Users are encouraged to double-check the model pipeline's interface in [diffusers' official documentation](https://huggingface.co/docs/diffusers/api/pipelines/overview).
|
||||
Some particular examples are below.
|
||||
|
||||
#### Wan Series
|
||||
|
||||
The Wan series video generation models takes `boundary_ratio` and `flow_shift` during model initialization ([ref](https://huggingface.co/docs/diffusers/api/pipelines/wan)), not during inference.
|
||||
|
||||
Since our `OmniDiffusionConfig` contains these two values ([source](https://github.com/vllm-project/vllm-omni/blob/main/vllm_omni/diffusion/data.py)), we can directly pass `--boundary-ratio` and `--flow-shift` arguments to `vllm serve` command.
|
||||
|
||||
```bash
|
||||
vllm serve "Wan2.2-T2V-A14B-Diffusers" \
|
||||
--omni \
|
||||
--boundary-ratio 0.875 \
|
||||
--flow-shift 3 \
|
||||
--diffusion-load-format diffusers
|
||||
```
|
||||
|
||||
These extra CLI args will be attempted to pass as-is to the `OmniDiffusionConfig` dataclass and being accessible during model loading time.
|
||||
Special routines inside the pipeline adapter ensures that they are set properly.
|
||||
@@ -1,199 +0,0 @@
|
||||
# Fish Speech S2 Pro
|
||||
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/fish_speech>.
|
||||
|
||||
|
||||
## Model
|
||||
|
||||
| Model | Description |
|
||||
|-------|-------------|
|
||||
| `fishaudio/s2-pro` | Fish Speech S2 Pro -- 4B dual-AR TTS model with DAC codec (44.1 kHz) |
|
||||
|
||||
## Gradio Demo
|
||||
|
||||
!!! note "Gradio is an optional dependency"
|
||||
The Gradio demo requires the `[demo]` extras. Install them first:
|
||||
|
||||
```bash
|
||||
pip install 'vllm-omni[demo]'
|
||||
```
|
||||
|
||||
Or, if installing from source: `pip install -e '.[demo]'`
|
||||
|
||||
An interactive Gradio demo is available with text-to-speech synthesis, voice cloning, and streaming support.
|
||||
|
||||
```bash
|
||||
# Option 1: Launch server + Gradio together
|
||||
./run_gradio_demo.sh
|
||||
|
||||
# Option 2: If server is already running
|
||||
python gradio_demo.py --api-base http://localhost:8091
|
||||
```
|
||||
|
||||
Then open http://localhost:7860 in your browser.
|
||||
|
||||
Features:
|
||||
|
||||
- Text-to-speech synthesis
|
||||
- Voice cloning from uploaded audio or URL
|
||||
- Streaming mode (progressive PCM playback)
|
||||
|
||||
## Launch the Server
|
||||
|
||||
```bash
|
||||
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
|
||||
./run_server.sh
|
||||
```
|
||||
|
||||
## Send TTS Request
|
||||
|
||||
### Using curl
|
||||
|
||||
```bash
|
||||
# Basic TTS
|
||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "Hello, how are you?",
|
||||
"voice": "default",
|
||||
"response_format": "wav"
|
||||
}' --output output.wav
|
||||
```
|
||||
|
||||
### Voice Cloning
|
||||
|
||||
Provide a reference audio (URL or base64 data URL) and its transcript:
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "Hello, this is a cloned voice.",
|
||||
"voice": "default",
|
||||
"ref_audio": "https://example.com/reference.wav",
|
||||
"ref_text": "Transcript of the reference audio."
|
||||
}' --output cloned.wav
|
||||
```
|
||||
|
||||
### Using Python
|
||||
|
||||
```python
|
||||
import httpx
|
||||
|
||||
# Basic TTS
|
||||
response = httpx.post(
|
||||
"http://localhost:8091/v1/audio/speech",
|
||||
json={
|
||||
"input": "Hello, how are you?",
|
||||
"voice": "default",
|
||||
"response_format": "wav",
|
||||
},
|
||||
timeout=300.0,
|
||||
)
|
||||
|
||||
with open("output.wav", "wb") as f:
|
||||
f.write(response.content)
|
||||
```
|
||||
|
||||
### Using the CLI Client
|
||||
|
||||
```bash
|
||||
cd examples/online_serving/fish_speech
|
||||
|
||||
# Basic TTS
|
||||
python speech_client.py --text "Hello, how are you?"
|
||||
|
||||
# Voice cloning
|
||||
python speech_client.py \
|
||||
--text "Hello, this is a cloned voice." \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--ref-text "Transcript of the reference audio."
|
||||
|
||||
# Streaming PCM output
|
||||
python speech_client.py --text "Hello world" --stream --output output.pcm
|
||||
```
|
||||
|
||||
The CLI client supports:
|
||||
|
||||
- `--api-base`: API base URL (default: `http://localhost:8091`)
|
||||
- `--model` (or `-m`): Model name (default: `fishaudio/s2-pro`)
|
||||
- `--text`: Text to synthesize (required)
|
||||
- `--ref-audio`: Reference audio for voice cloning (local path or URL)
|
||||
- `--ref-text`: Transcript of the reference audio
|
||||
- `--stream`: Enable streaming (PCM output)
|
||||
- `--response-format`: Audio format: wav, mp3, flac, pcm, aac, opus (default: wav)
|
||||
- `--output` (or `-o`): Output file path
|
||||
|
||||
## Streaming
|
||||
|
||||
Set `stream=true` with `response_format="pcm"` to receive raw PCM audio chunks as they are decoded:
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "Hello, how are you?",
|
||||
"voice": "default",
|
||||
"stream": true,
|
||||
"response_format": "pcm"
|
||||
}' --no-buffer | play -t raw -r 44100 -e signed -b 16 -c 1 -
|
||||
```
|
||||
|
||||
**Note:** Fish Speech outputs at 44.1 kHz (unlike Qwen3-TTS which outputs at 24 kHz).
|
||||
|
||||
## API Parameters
|
||||
|
||||
Fish Speech uses the same `/v1/audio/speech` endpoint as Qwen3-TTS. See the [Speech API reference](https://docs.vllm.ai/projects/vllm-omni/en/latest/serving/speech_api/) for full parameter documentation.
|
||||
|
||||
Key parameters for Fish Speech:
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| `input` | string | **required** | Text to synthesize |
|
||||
| `voice` | string | "default" | Voice name (use "default" for Fish Speech) |
|
||||
| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
|
||||
| `ref_audio` | string | null | Reference audio URL or base64 data URL for voice cloning |
|
||||
| `ref_text` | string | null | Transcript of reference audio (required for voice cloning) |
|
||||
| `max_new_tokens` | int | 4096 | Maximum tokens to generate |
|
||||
| `stream` | bool | false | Stream raw PCM chunks |
|
||||
|
||||
## Prerequisites
|
||||
|
||||
Install the `fish-speech` package for the DAC codec:
|
||||
|
||||
```bash
|
||||
pip install fish-speech
|
||||
```
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
1. **No audio output**: Make sure the `fish-speech` package is installed for the DAC decoder
|
||||
2. **Connection refused**: Ensure the server is running on the correct port
|
||||
3. **Flashinfer version mismatch**: Set `FLASHINFER_DISABLE_VERSION_CHECK=1` if you see version warnings
|
||||
4. **Out of memory**: Reduce `--gpu-memory-utilization` or use a GPU with more VRAM
|
||||
|
||||
## Example materials
|
||||
|
||||
??? abstract "gradio_demo.py"
|
||||
``````py
|
||||
--8<-- "examples/online_serving/fish_speech/gradio_demo.py"
|
||||
``````
|
||||
??? abstract "run_gradio_demo.sh"
|
||||
``````sh
|
||||
--8<-- "examples/online_serving/fish_speech/run_gradio_demo.sh"
|
||||
``````
|
||||
??? abstract "run_server.sh"
|
||||
``````sh
|
||||
--8<-- "examples/online_serving/fish_speech/run_server.sh"
|
||||
``````
|
||||
??? abstract "speech_client.py"
|
||||
``````py
|
||||
--8<-- "examples/online_serving/fish_speech/speech_client.py"
|
||||
``````
|
||||
@@ -1,99 +1,96 @@
|
||||
# GLM-Image Online Serving
|
||||
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/glm_image>.
|
||||
GLM-Image is a 2-stage image generation model (AR + Diffusion) supported by vLLM-Omni's
|
||||
declarative config system. The pipeline topology and stage structure are declared in
|
||||
`vllm_omni/model_executor/models/glm_image/pipeline.py`; deployment knobs (GPU placement,
|
||||
memory, sampling params) live in `vllm_omni/deploy/glm_image.yaml`.
|
||||
|
||||
|
||||
This example demonstrates how to deploy GLM-Image for online image generation using vLLM-Omni.
|
||||
|
||||
## 🛠️ Installation
|
||||
|
||||
Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/README.md)
|
||||
|
||||
## Run examples (GLM-Image)
|
||||
|
||||
**Note**: These examples work with the default configuration on **2× NVIDIA A100 (80GB)** or equivalent. Stage 0 (AR) and Stage 1 (Diffusion) each use one GPU by default. For single-GPU setups, modify the stage configuration to share the same device.
|
||||
|
||||
### Launch the Server
|
||||
## Start Server
|
||||
|
||||
```bash
|
||||
# Use default configuration
|
||||
vllm serve zai-org/GLM-Image --omni --port 8091
|
||||
```
|
||||
|
||||
Or use the convenience script:
|
||||
The config system auto-detects the pipeline from the model's `model_index.json` — no
|
||||
manual `--stage-configs-path` or `--deploy-config` needed.
|
||||
|
||||
By default, stage 0 (AR) runs on GPU 0 and stage 1 (Diffusion) on GPU 1. To colocate
|
||||
both stages on a single GPU, override per stage:
|
||||
|
||||
```bash
|
||||
cd examples/online_serving/glm_image
|
||||
bash run_server.sh
|
||||
vllm serve zai-org/GLM-Image --omni --port 8091 \
|
||||
--stage-0-devices 0 --stage-1-devices 0
|
||||
```
|
||||
|
||||
If you have a custom stage configs file:
|
||||
|
||||
```bash
|
||||
vllm serve zai-org/GLM-Image --omni --port 8091 --stage-configs-path /path/to/glm_image.yaml
|
||||
```
|
||||
|
||||
### Send Requests
|
||||
|
||||
Get into the glm_image folder:
|
||||
|
||||
```bash
|
||||
cd examples/online_serving/glm_image
|
||||
```
|
||||
|
||||
Send request via Python:
|
||||
|
||||
```bash
|
||||
python openai_chat_client.py --prompt "A cute cat sitting on a window sill"
|
||||
```
|
||||
|
||||
The Python client supports the following command-line arguments:
|
||||
|
||||
- `--prompt` (or `-p`): Text prompt for generation (default: `A beautiful sunset over the ocean with sailing boats`)
|
||||
- `--output` (or `-o`): Output file path (default: `glm_image_output.png`)
|
||||
- `--server` (or `-s`): Server URL (default: `http://localhost:8091`)
|
||||
- `--image` (or `-i`): Input image path (for image-to-image editing)
|
||||
- `--height`: Image height in pixels (default: 1024)
|
||||
- `--width`: Image width in pixels (default: 1024)
|
||||
- `--steps`: Number of inference steps (default: 50)
|
||||
- `--guidance-scale`: Classifier-free guidance scale (default: 1.5)
|
||||
- `--seed`: Random seed (default: 42)
|
||||
- `--negative`: Negative prompt
|
||||
|
||||
## Modality Control
|
||||
|
||||
GLM-Image supports **text-to-image** and **image-to-image** modes.
|
||||
|
||||
The default yaml configuration deploys AR on GPU 0 and DiT on GPU 1. You can use the default configuration file: [`glm_image.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/model_executor/stage_configs/glm_image.yaml)
|
||||
|
||||
| Mode | Input | Output | Description |
|
||||
| -------------- | ------------ | ------ | ---------------------------------- |
|
||||
| Text-to-Image | Text | Image | Generate images from text prompts |
|
||||
| Image-to-Image | Image + Text | Image | Edit images with text instructions |
|
||||
## API Calls
|
||||
|
||||
### Text-to-Image
|
||||
|
||||
```bash
|
||||
python openai_chat_client.py \
|
||||
--prompt "A photorealistic mountain landscape at sunset" \
|
||||
--height 1024 \
|
||||
--width 1024 \
|
||||
--output landscape.png
|
||||
|
||||
# Or use the curl script:
|
||||
bash run_curl_text_to_image.sh "A futuristic city skyline at night"
|
||||
curl -s http://localhost:8091/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"messages": [
|
||||
{"role": "user", "content": "A photorealistic mountain landscape at sunset"}
|
||||
],
|
||||
"extra_body": {
|
||||
"height": 1024,
|
||||
"width": 1024,
|
||||
"num_inference_steps": 50,
|
||||
"guidance_scale": 1.5,
|
||||
"seed": 42
|
||||
}
|
||||
}' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > output.png
|
||||
```
|
||||
|
||||
### Image-to-Image (Image Editing)
|
||||
|
||||
```bash
|
||||
python openai_chat_client.py \
|
||||
--prompt "Convert this image to watercolor style" \
|
||||
--image input.png \
|
||||
--output watercolor.png
|
||||
curl -s http://localhost:8091/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Convert this image to watercolor style"},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,$(base64 -w0 input.png)}"}
|
||||
]
|
||||
}
|
||||
],
|
||||
"extra_body": {
|
||||
"height": 1024,
|
||||
"width": 1024,
|
||||
"num_inference_steps": 50,
|
||||
"guidance_scale": 1.5
|
||||
}
|
||||
}' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > output.png
|
||||
```
|
||||
|
||||
# Or use the curl script:
|
||||
bash run_curl_image_edit.sh input.png "Convert to watercolor style"
|
||||
### Using the OpenAI Python SDK
|
||||
|
||||
```python
|
||||
from openai import OpenAI
|
||||
import base64
|
||||
|
||||
client = OpenAI(base_url="http://localhost:8091/v1", api_key="none")
|
||||
|
||||
response = client.chat.completions.create(
|
||||
model="zai-org/GLM-Image",
|
||||
messages=[{"role": "user", "content": "A beautiful sunset over the ocean"}],
|
||||
extra_body={
|
||||
"height": 1024,
|
||||
"width": 1024,
|
||||
"num_inference_steps": 50,
|
||||
"guidance_scale": 1.5,
|
||||
"seed": 42,
|
||||
},
|
||||
)
|
||||
|
||||
img_url = response.choices[0].message.content[0].image_url.url
|
||||
_, b64_data = img_url.split(",", 1)
|
||||
with open("output.png", "wb") as f:
|
||||
f.write(base64.b64decode(b64_data))
|
||||
```
|
||||
|
||||
For general-purpose request methods (curl, OpenAI SDK, Python `requests`), see
|
||||
@@ -116,7 +113,7 @@ dimensions and count, use `size` and `n` rather than `height` or `width`.
|
||||
| `width` | int | 1024 | Image width in pixels |
|
||||
| `num_inference_steps` | int | 50 | Number of diffusion denoising steps |
|
||||
| `guidance_scale` | float | 1.5 | Classifier-free guidance scale |
|
||||
| `seed` | int | None | Optional random seed; `/v1/images/*` generates one server-side if omitted |
|
||||
| `seed` | int | None | Optional random seed |
|
||||
| `negative_prompt` | str | None | Negative prompt |
|
||||
|
||||
## Response Format
|
||||
@@ -150,13 +147,12 @@ dimensions and count, use `size` and `n` rather than `height` or `width`.
|
||||
## Extract Image
|
||||
|
||||
```bash
|
||||
# From a saved JSON response
|
||||
cat response.json | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > output.png
|
||||
```
|
||||
|
||||
## Architecture
|
||||
|
||||
GLM-Image uses a 2-stage multistage pipeline:
|
||||
GLM-Image uses a 2-stage pipeline:
|
||||
|
||||
```
|
||||
Stage 0 (AR Model) Stage 1 (Diffusion)
|
||||
@@ -181,41 +177,13 @@ Stage 0 (AR Model) Stage 1 (Diffusion)
|
||||
| Stage-1 (DiT+VAE) | **~20 GiB** |
|
||||
| Total | **~38 GiB + KV Cache** |
|
||||
|
||||
## File Description
|
||||
|
||||
| File | Description |
|
||||
| --------------------------- | ------------------------------------- |
|
||||
| `run_server.sh` | Server startup script |
|
||||
| `run_curl_text_to_image.sh` | Text-to-image curl example |
|
||||
| `run_curl_image_edit.sh` | Image-to-image (editing) curl example |
|
||||
| `openai_chat_client.py` | Python client (t2i + i2i) |
|
||||
|
||||
## FAQ
|
||||
|
||||
- If you encounter OOM errors, adjust `gpu_memory_utilization` in the stage config:
|
||||
- If you encounter OOM errors, adjust `gpu_memory_utilization` in the deploy config:
|
||||
|
||||
```yaml
|
||||
# In glm_image.yaml, reduce from default 0.6:
|
||||
# In vllm_omni/deploy/glm_image.yaml, reduce from default 0.6:
|
||||
gpu_memory_utilization: 0.5
|
||||
```
|
||||
|
||||
- The first request may be slow due to model warmup. Subsequent requests will be faster.
|
||||
|
||||
## Example materials
|
||||
|
||||
??? abstract "openai_chat_client.py"
|
||||
``````py
|
||||
--8<-- "examples/online_serving/glm_image/openai_chat_client.py"
|
||||
``````
|
||||
??? abstract "run_curl_image_edit.sh"
|
||||
``````sh
|
||||
--8<-- "examples/online_serving/glm_image/run_curl_image_edit.sh"
|
||||
``````
|
||||
??? abstract "run_curl_text_to_image.sh"
|
||||
``````sh
|
||||
--8<-- "examples/online_serving/glm_image/run_curl_text_to_image.sh"
|
||||
``````
|
||||
??? abstract "run_server.sh"
|
||||
``````sh
|
||||
--8<-- "examples/online_serving/glm_image/run_server.sh"
|
||||
``````
|
||||
@@ -28,11 +28,13 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
|
||||
### Launch individual stages (stage-based CLI)
|
||||
|
||||
Adopt the stage-based CLI architecture to independently instantiate execution processes per functional stage.
|
||||
The example below pins Stage 0 to GPU 0 and Stage 1/2 to GPU 1 via
|
||||
`CUDA_VISIBLE_DEVICES`.
|
||||
|
||||
**1. Stage 0 (Thinker + API server)**
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
CUDA_VISIBLE_DEVICES=0 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
--port 8091 \
|
||||
--stage-id 0 \
|
||||
--omni-master-address 127.0.0.1 \
|
||||
@@ -42,7 +44,7 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
**2. Stage 1 (Talker)**
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
--stage-id 1 \
|
||||
--headless \
|
||||
--omni-master-address 127.0.0.1 \
|
||||
@@ -52,7 +54,7 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
**3. Stage 2 (Code2Wav)**
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
--stage-id 2 \
|
||||
--headless \
|
||||
--omni-master-address 127.0.0.1 \
|
||||
@@ -75,7 +77,7 @@ 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 \
|
||||
CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
--stage-id 1 \
|
||||
--headless \
|
||||
--gpu-memory-utilization 0.5 \
|
||||
|
||||
@@ -1,451 +0,0 @@
|
||||
# Qwen3-TTS
|
||||
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/qwen3_tts>.
|
||||
|
||||
|
||||
## 🛠️ Installation
|
||||
|
||||
Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/README.md)
|
||||
|
||||
## Supported Models
|
||||
|
||||
| Model | Task Type | Description |
|
||||
| -------------------------------------- | ----------- | ----------------------------------------------------- |
|
||||
| `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | CustomVoice | Predefined speaker voices with optional style control |
|
||||
| `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | VoiceDesign | Natural language voice style description |
|
||||
| `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | Base | Voice cloning from reference audio |
|
||||
| `Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice` | CustomVoice | Smaller/faster variant |
|
||||
| `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | Base | Smaller/faster variant for voice cloning |
|
||||
|
||||
## Gradio Demo
|
||||
|
||||
!!! note "Gradio is an optional dependency"
|
||||
The Gradio demo requires the `[demo]` extras. Install them first:
|
||||
|
||||
```bash
|
||||
pip install 'vllm-omni[demo]'
|
||||
```
|
||||
|
||||
Or, if installing from source: `pip install -e '.[demo]'`
|
||||
|
||||
Two interactive Gradio demos are available, both supporting all 3 task types:
|
||||
|
||||
| Demo | File | Transport | Streaming Quality |
|
||||
| -------- | ------------------------ | ------------ | -------------------------------------------------- |
|
||||
| Standard | `gradio_demo.py` | HTTP chunked | May have small gaps between chunks |
|
||||
| FastRTC | `gradio_fastrtc_demo.py` | WebRTC | Gapless streaming (requires `pip install fastrtc`) |
|
||||
|
||||
```bash
|
||||
# Option 1: Launch server + Standard Gradio together
|
||||
./run_gradio_demo.sh # CustomVoice (default)
|
||||
./run_gradio_demo.sh --task-type VoiceDesign # VoiceDesign
|
||||
./run_gradio_demo.sh --task-type Base # Voice cloning
|
||||
|
||||
# Option 2: If server is already running
|
||||
python gradio_demo.py --api-base http://localhost:8000
|
||||
|
||||
# Option 3: FastRTC demo (gapless streaming)
|
||||
pip install fastrtc
|
||||
python gradio_fastrtc_demo.py --api-base http://localhost:8000
|
||||
```
|
||||
|
||||
Then open http://localhost:7860 in your browser.
|
||||
|
||||
## Run examples (Qwen3-TTS)
|
||||
|
||||
### Launch the Server
|
||||
|
||||
```bash
|
||||
# CustomVoice model (predefined speakers)
|
||||
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
|
||||
--deploy-config vllm_omni/deploy/qwen3_tts.yaml \
|
||||
--omni \
|
||||
--port 8091 \
|
||||
--trust-remote-code \
|
||||
--enforce-eager
|
||||
|
||||
# VoiceDesign model
|
||||
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign \
|
||||
--deploy-config vllm_omni/deploy/qwen3_tts.yaml \
|
||||
--omni \
|
||||
--port 8091 \
|
||||
--trust-remote-code \
|
||||
--enforce-eager
|
||||
|
||||
# Base model (voice cloning)
|
||||
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base \
|
||||
--deploy-config vllm_omni/deploy/qwen3_tts.yaml \
|
||||
--omni \
|
||||
--port 8091 \
|
||||
--trust-remote-code \
|
||||
--enforce-eager
|
||||
```
|
||||
|
||||
If you have custom stage configs file, launch the server with command below
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
|
||||
--stage-configs-path /path/to/stage_configs_file \
|
||||
--omni \
|
||||
--port 8091 \
|
||||
--trust-remote-code \
|
||||
--enforce-eager
|
||||
```
|
||||
|
||||
Alternatively, use the convenience script:
|
||||
```bash
|
||||
./run_server.sh # Default: CustomVoice model
|
||||
./run_server.sh CustomVoice # CustomVoice model
|
||||
./run_server.sh VoiceDesign # VoiceDesign model
|
||||
./run_server.sh Base # Base (voice clone) model
|
||||
```
|
||||
|
||||
### Send TTS Request
|
||||
|
||||
Get into the example folder
|
||||
```bash
|
||||
cd examples/online_serving/qwen3_tts
|
||||
```
|
||||
|
||||
#### Send request via python
|
||||
|
||||
```bash
|
||||
# CustomVoice: Use predefined speaker
|
||||
python openai_speech_client.py \
|
||||
--text "你好,我是通义千问" \
|
||||
--speaker vivian \
|
||||
--language Chinese
|
||||
|
||||
# CustomVoice with style instruction
|
||||
python openai_speech_client.py \
|
||||
--text "今天天气真好" \
|
||||
--speaker ryan \
|
||||
--instructions "用开心的语气说"
|
||||
|
||||
# VoiceDesign: Describe the voice style
|
||||
python openai_speech_client.py \
|
||||
--model Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign \
|
||||
--task-type VoiceDesign \
|
||||
--text "哥哥,你回来啦" \
|
||||
--instructions "体现撒娇稚嫩的萝莉女声,音调偏高"
|
||||
|
||||
# Base: Voice cloning
|
||||
python openai_speech_client.py \
|
||||
--model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
|
||||
--task-type Base \
|
||||
--text "Hello, this is a cloned voice" \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--ref-text "Original transcript of the reference audio"
|
||||
```
|
||||
|
||||
The Python client supports the following command-line arguments:
|
||||
|
||||
- `--api-base`: API base URL (default: `http://localhost:8091`)
|
||||
- `--model` (or `-m`): Model name/path (default: `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice`)
|
||||
- `--task-type` (or `-t`): TTS task type. Options: `CustomVoice`, `VoiceDesign`, `Base`
|
||||
- `--text`: Text to synthesize (required)
|
||||
- `--speaker`: Speaker name (default: `vivian`). Options: `vivian`, `ryan`, `aiden`, etc.
|
||||
- `--language`: Language. Options: `Auto`, `Chinese`, `English`, `Japanese`, `Korean`, `German`, `French`, `Russian`, `Portuguese`, `Spanish`, `Italian`
|
||||
- `--instructions`: Voice style/emotion instructions
|
||||
- `--ref-audio`: Reference audio file path or URL for voice cloning (Base task). Local paths are automatically base64-encoded by the client before sending to the server.
|
||||
- `--ref-text`: Reference audio transcript for voice cloning (Base task).
|
||||
- `--response-format`: Audio output format (default: `wav`). Options: `wav`, `mp3`, `flac`, `pcm`, `aac`, `opus`
|
||||
- `--output` (or `-o`): Output audio file path (default: `tts_output.wav`)
|
||||
|
||||
#### Send request via curl
|
||||
|
||||
```bash
|
||||
# Simple TTS request
|
||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "Hello, how are you?",
|
||||
"voice": "vivian",
|
||||
"language": "English"
|
||||
}' --output output.wav
|
||||
|
||||
# With style instruction
|
||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "I am so excited!",
|
||||
"voice": "vivian",
|
||||
"instructions": "Speak with great enthusiasm"
|
||||
}' --output excited.wav
|
||||
|
||||
# List available voices in CustomVoice models
|
||||
curl http://localhost:8091/v1/audio/voices
|
||||
```
|
||||
|
||||
### Using OpenAI SDK
|
||||
|
||||
```python
|
||||
from openai import OpenAI
|
||||
|
||||
client = OpenAI(base_url="http://localhost:8091/v1", api_key="none")
|
||||
|
||||
response = client.audio.speech.create(
|
||||
model="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
|
||||
voice="vivian",
|
||||
input="Hello, how are you?",
|
||||
)
|
||||
|
||||
response.stream_to_file("output.wav")
|
||||
```
|
||||
|
||||
### Using Python httpx
|
||||
|
||||
```python
|
||||
import httpx
|
||||
|
||||
response = httpx.post(
|
||||
"http://localhost:8091/v1/audio/speech",
|
||||
json={
|
||||
"input": "Hello, how are you?",
|
||||
"voice": "vivian",
|
||||
"language": "English",
|
||||
},
|
||||
timeout=300.0,
|
||||
)
|
||||
|
||||
with open("output.wav", "wb") as f:
|
||||
f.write(response.content)
|
||||
```
|
||||
|
||||
## API Reference
|
||||
|
||||
### Voices Endpoint
|
||||
|
||||
#### GET /v1/audio/voices
|
||||
|
||||
List all available voices/speakers from the loaded model, including both built-in model voices and uploaded custom voices.
|
||||
|
||||
**Response Example:**
|
||||
```json
|
||||
{
|
||||
"voices": ["vivian", "ryan", "custom_voice_1"],
|
||||
"uploaded_voices": [
|
||||
{
|
||||
"name": "custom_voice_1",
|
||||
"consent": "user_consent_id",
|
||||
"created_at": 1738660000,
|
||||
"file_size": 1024000,
|
||||
"mime_type": "audio/wav",
|
||||
"ref_text": "The exact transcript of the audio sample.",
|
||||
"speaker_description": "warm narrator"
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
Fields `ref_text` and `speaker_description` are omitted per-entry when not provided at upload time.
|
||||
|
||||
#### POST /v1/audio/voices
|
||||
|
||||
Upload a new voice sample for voice cloning in Base task TTS requests.
|
||||
|
||||
**Form Parameters:**
|
||||
- `audio_sample` (required): Audio file (max 10MB, supported formats: wav, mp3, flac, ogg, aac, webm, mp4)
|
||||
- `consent` (required): Consent recording ID
|
||||
- `name` (required): Name for the new voice
|
||||
- `ref_text` (optional): Transcript of the audio. Enables in-context voice cloning (higher quality).
|
||||
- `speaker_description` (optional): Free-form description of the voice (e.g. "warm narrator", "energetic presenter"). Stored as metadata.
|
||||
|
||||
**Response Example:**
|
||||
```json
|
||||
{
|
||||
"success": true,
|
||||
"voice": {
|
||||
"name": "custom_voice_1",
|
||||
"consent": "user_consent_id",
|
||||
"created_at": 1738660000,
|
||||
"mime_type": "audio/wav",
|
||||
"file_size": 1024000,
|
||||
"ref_text": "The exact transcript of the audio sample.",
|
||||
"speaker_description": "warm narrator"
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
Fields `ref_text` and `speaker_description` are omitted when not provided at upload time.
|
||||
|
||||
**Usage Example:**
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/voices \
|
||||
-F "audio_sample=@/path/to/voice_sample.wav" \
|
||||
-F "consent=user_consent_id" \
|
||||
-F "name=custom_voice_1" \
|
||||
-F "ref_text=The exact transcript of the audio sample." \
|
||||
-F "speaker_description=warm narrator"
|
||||
```
|
||||
|
||||
### Endpoint
|
||||
|
||||
```
|
||||
POST /v1/audio/speech
|
||||
Content-Type: application/json
|
||||
```
|
||||
|
||||
This endpoint follows the [OpenAI Audio Speech API](https://platform.openai.com/docs/api-reference/audio/createSpeech) format with additional Qwen3-TTS parameters.
|
||||
|
||||
### Request Body
|
||||
|
||||
```json
|
||||
{
|
||||
"input": "Text to synthesize",
|
||||
"voice": "vivian",
|
||||
"response_format": "wav",
|
||||
"task_type": "CustomVoice",
|
||||
"language": "Auto",
|
||||
"instructions": "Optional style instructions",
|
||||
"ref_audio": "HTTP URL, base64 data URL, or file:// URI for voice cloning",
|
||||
"ref_text": "Reference audio transcript",
|
||||
"x_vector_only_mode": false,
|
||||
"max_new_tokens": 2048
|
||||
}
|
||||
```
|
||||
|
||||
> **Note:** The `model` field is optional when serving a single model, as the server already knows which model is loaded.
|
||||
|
||||
### Response
|
||||
|
||||
Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/wav`).
|
||||
|
||||
### Voice and language (summary)
|
||||
|
||||
- **Speaker**: Use the `voice` request field to select the speaker (e.g., `vivian`, `ryan`, `aiden`). List available speakers with `GET /v1/audio/voices`.
|
||||
- **Language**: Use the `language` field for the codec language tag (`Auto`, `Chinese`, `English`, etc.). Default is `Auto` for automatic detection.
|
||||
- **CustomVoice**: Requires a valid `voice` from the model’s speaker set. **VoiceDesign**: Use `instructions` to describe the voice. **Base**: Use `ref_audio` and `ref_text` for voice cloning.
|
||||
|
||||
## Parameters
|
||||
|
||||
### OpenAI Standard Parameters
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
| ----------------- | ------ | -------------- | ----------------------------------------------------------- |
|
||||
| `input` | string | **required** | Text to synthesize |
|
||||
| `model` | string | server's model | Model to use (optional, should match server if specified) |
|
||||
| `voice` | string | "vivian" | Speaker name (e.g., vivian, ryan, aiden) |
|
||||
| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
|
||||
| `speed` | float | 1.0 | Playback speed (0.25-4.0, not supported with `stream=true`) |
|
||||
|
||||
### vLLM-Omni Extension Parameters
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
| ---------------------------- | ------ | ------------- | -------------------------------------------------------------------------------------------------------------------- |
|
||||
| `task_type` | string | "CustomVoice" | Task: CustomVoice, VoiceDesign, or Base |
|
||||
| `language` | string | "Auto" | Language (see supported languages below) |
|
||||
| `instructions` | string | "" | Voice style/emotion instructions |
|
||||
| `max_new_tokens` | int | 2048 | Maximum tokens to generate |
|
||||
| `initial_codec_chunk_frames` | int | null | Per-request initial chunk size override for TTFA tuning. When null, IC is computed dynamically based on server load. |
|
||||
| `stream` | bool | false | Stream raw PCM chunks as they are decoded (requires `response_format="pcm"`) |
|
||||
|
||||
**Supported languages:** Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian
|
||||
|
||||
### Voice Clone Parameters (Base task)
|
||||
|
||||
| Parameter | Type | Required | Description |
|
||||
| -------------------- | ------ | -------- | ----------------------------------------------------------------------------------------------- |
|
||||
| `ref_audio` | string | **Yes** | Reference audio (HTTP URL, base64 data URL, or `file://` URI with `--allowed-local-media-path`) |
|
||||
| `ref_text` | string | No | Transcript of reference audio (for ICL mode) |
|
||||
| `x_vector_only_mode` | bool | No | Use speaker embedding only (no ICL) |
|
||||
|
||||
## Streaming
|
||||
|
||||
Set `stream=true` with `response_format="pcm"` to receive raw PCM audio chunks as they are decoded
|
||||
(one chunk per Code2Wav window, default 25 frames; configurable in the stage config):
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "Hello, how are you?",
|
||||
"voice": "vivian",
|
||||
"language": "English",
|
||||
"stream": true,
|
||||
"response_format": "pcm"
|
||||
}' --no-buffer | play -t raw -r 24000 -e signed -b 16 -c 1 -
|
||||
```
|
||||
|
||||
**Constraints:**
|
||||
- `stream=true` requires `response_format="pcm"` (raw 16-bit signed PCM, 24 kHz mono).
|
||||
- `speed` adjustment is not supported when streaming.
|
||||
- Requires the server stage config to have `async_chunk: true` (default in `qwen3_tts.yaml`).
|
||||
|
||||
## Streaming Text Input (WebSocket)
|
||||
|
||||
The `/v1/audio/speech/stream` WebSocket endpoint accepts text incrementally, buffers and splits it at sentence boundaries, and generates audio per sentence.
|
||||
|
||||
When `stream_audio=true`, each sentence is emitted as `audio.start`, one or more binary PCM frames, and `audio.done`.
|
||||
|
||||
### Quick Start
|
||||
|
||||
```bash
|
||||
python streaming_speech_client.py \
|
||||
--text "Hello world. How are you? I am fine."
|
||||
|
||||
python streaming_speech_client.py \
|
||||
--text "Hello world. How are you? I am fine." \
|
||||
--simulate-stt --stt-delay 0.1
|
||||
```
|
||||
|
||||
### WebSocket Protocol
|
||||
|
||||
Client -> Server:
|
||||
|
||||
```jsonc
|
||||
{"type": "session.config", "voice": "Vivian", "task_type": "CustomVoice", "language": "Auto", "split_granularity": "sentence", "stream_audio": true, "response_format": "pcm"}
|
||||
{"type": "input.text", "text": "Hello, how are you? "}
|
||||
{"type": "input.done"}
|
||||
```
|
||||
|
||||
Server -> Client:
|
||||
|
||||
```jsonc
|
||||
{"type": "audio.start", "sentence_index": 0, "sentence_text": "Hello, how are you?", "format": "pcm", "sample_rate": 24000}
|
||||
// binary PCM frame(s)
|
||||
{"type": "audio.done", "sentence_index": 0, "total_bytes": 96000, "error": false}
|
||||
{"type": "session.done", "total_sentences": 1}
|
||||
```
|
||||
|
||||
## Limitations
|
||||
|
||||
- **Single request**: Batch processing is not yet optimized for online serving.
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
1. **TTS model did not produce audio output**: Ensure you're using the correct model variant for your task type (CustomVoice task → CustomVoice model, etc.)
|
||||
2. **Connection refused**: Make sure the server is running on the correct port
|
||||
3. **Out of memory**: Use smaller model variant (`Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice`) or reduce `--gpu-memory-utilization`
|
||||
4. **Unsupported speaker**: Use `/v1/audio/voices` to list available voices for the loaded model
|
||||
5. **Voice clone fails**: Ensure you're using the Base model variant for voice cloning
|
||||
|
||||
## Example materials
|
||||
|
||||
??? abstract "gradio_demo.py"
|
||||
``````py
|
||||
--8<-- "examples/online_serving/qwen3_tts/gradio_demo.py"
|
||||
``````
|
||||
??? abstract "gradio_fastrtc_demo.py"
|
||||
``````py
|
||||
--8<-- "examples/online_serving/qwen3_tts/gradio_fastrtc_demo.py"
|
||||
``````
|
||||
??? abstract "openai_speech_client.py"
|
||||
``````py
|
||||
--8<-- "examples/online_serving/qwen3_tts/openai_speech_client.py"
|
||||
``````
|
||||
??? abstract "run_gradio_demo.sh"
|
||||
``````sh
|
||||
--8<-- "examples/online_serving/qwen3_tts/run_gradio_demo.sh"
|
||||
``````
|
||||
??? abstract "run_server.sh"
|
||||
``````sh
|
||||
--8<-- "examples/online_serving/qwen3_tts/run_server.sh"
|
||||
``````
|
||||
??? abstract "streaming_speech_client.py"
|
||||
``````py
|
||||
--8<-- "examples/online_serving/qwen3_tts/streaming_speech_client.py"
|
||||
``````
|
||||
??? abstract "tts_common.py"
|
||||
``````py
|
||||
--8<-- "examples/online_serving/qwen3_tts/tts_common.py"
|
||||
``````
|
||||
@@ -0,0 +1,193 @@
|
||||
# Text-To-Audio
|
||||
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/stable_audio>.
|
||||
|
||||
This example demonstrates how to deploy Stable Audio models for online text-to-audio generation using vLLM-Omni.
|
||||
|
||||
## Supported Models
|
||||
|
||||
| Model | Description |
|
||||
|-------|-------------|
|
||||
| `stabilityai/stable-audio-open-1.0` | Open-source audio generation, up to ~47 seconds, 44.1 kHz stereo |
|
||||
|
||||
## Start Server
|
||||
|
||||
### Basic Start
|
||||
|
||||
```bash
|
||||
vllm-omni serve stabilityai/stable-audio-open-1.0 \
|
||||
--host 0.0.0.0 \
|
||||
--port 8091 \
|
||||
--gpu-memory-utilization 0.9 \
|
||||
--trust-remote-code \
|
||||
--enforce-eager \
|
||||
--omni
|
||||
```
|
||||
|
||||
## API Calls
|
||||
|
||||
### Method 1: Using curl
|
||||
|
||||
```bash
|
||||
# Run all curl examples
|
||||
bash curl_examples.sh
|
||||
|
||||
# Or execute directly
|
||||
curl -X POST http://localhost:8091/v1/audio/generate \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "The sound of a cat purring",
|
||||
"audio_length": 10.0
|
||||
}' --output cat.wav
|
||||
```
|
||||
|
||||
### Method 2: Using Python Client
|
||||
|
||||
```bash
|
||||
cd examples/online_serving/stable_audio
|
||||
|
||||
# Simple generation
|
||||
python stable_audio_client.py \
|
||||
--text "The sound of a cat purring"
|
||||
|
||||
# With custom duration
|
||||
python stable_audio_client.py \
|
||||
--text "A dog barking" \
|
||||
--audio_length 5.0
|
||||
|
||||
# With all parameters
|
||||
python stable_audio_client.py \
|
||||
--text "Thunder and rain" \
|
||||
--audio_length 15.0 \
|
||||
--negative_prompt "Low quality" \
|
||||
--guidance_scale 7.0 \
|
||||
--num_inference_steps 100 \
|
||||
--seed 42 \
|
||||
--output thunder.wav
|
||||
```
|
||||
|
||||
The Python client supports the following command-line arguments:
|
||||
|
||||
- `--api_url`: API endpoint URL (default: `http://localhost:8091/v1/audio/generate`)
|
||||
- `--text`: Text prompt for audio generation (default: `"The sound of a cat purring"`)
|
||||
- `--audio_length`: Audio length in seconds (default: `10.0`, max ~47s for `stable-audio-open-1.0`)
|
||||
- `--audio_start`: Audio start time in seconds (default: `0.0`)
|
||||
- `--negative_prompt`: Negative prompt for classifier-free guidance (default: `"Low quality"`)
|
||||
- `--guidance_scale`: Guidance scale for diffusion (default: `7.0`)
|
||||
- `--num_inference_steps`: Number of inference steps (default: `100`)
|
||||
- `--seed`: Random seed for reproducibility (default: `None`)
|
||||
- `--response_format`: Audio output format (default: `wav`). Options: `wav`, `mp3`, `flac`, `pcm`
|
||||
- `--output`: Output file path (default: `stable_audio_output.wav`)
|
||||
|
||||
### Method 3: Using Python httpx
|
||||
|
||||
```python
|
||||
import httpx
|
||||
|
||||
response = httpx.post(
|
||||
"http://localhost:8091/v1/audio/generate",
|
||||
json={
|
||||
"input": "The sound of ocean waves crashing on a beach",
|
||||
"audio_length": 10.0,
|
||||
"negative_prompt": "Low quality, distorted",
|
||||
"guidance_scale": 7.0,
|
||||
"num_inference_steps": 100,
|
||||
},
|
||||
timeout=300.0,
|
||||
)
|
||||
|
||||
with open("ocean.wav", "wb") as f:
|
||||
f.write(response.content)
|
||||
```
|
||||
|
||||
## Request Format
|
||||
|
||||
### Simple Generation
|
||||
|
||||
```json
|
||||
{
|
||||
"input": "The sound of ocean waves"
|
||||
}
|
||||
```
|
||||
|
||||
### Generation with Parameters
|
||||
|
||||
```json
|
||||
{
|
||||
"input": "A piano playing a gentle melody",
|
||||
"audio_length": 10.0,
|
||||
"negative_prompt": "Low quality, distorted, noisy",
|
||||
"guidance_scale": 8.0,
|
||||
"num_inference_steps": 150,
|
||||
"seed": 42,
|
||||
"response_format": "wav"
|
||||
}
|
||||
```
|
||||
|
||||
## API Reference
|
||||
|
||||
### Endpoint
|
||||
|
||||
```
|
||||
POST /v1/audio/generate
|
||||
Content-Type: application/json
|
||||
```
|
||||
|
||||
### Generation Parameters
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| `input` | string | **required** | Text prompt describing the audio to generate |
|
||||
| `model` | string | server's model | Model to use (optional, should match server if specified) |
|
||||
| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
|
||||
| `speed` | float | 1.0 | Playback speed (0.25 - 4.0) |
|
||||
| `audio_length` | float | null | Audio duration in seconds (max ~47s for `stable-audio-open-1.0`) |
|
||||
| `audio_start` | float | 0.0 | Audio start time in seconds |
|
||||
| `negative_prompt` | string | null | Text describing what to avoid in generation |
|
||||
| `guidance_scale` | float | model default | Classifier-free guidance scale (higher = more adherence to prompt) |
|
||||
| `num_inference_steps` | int | model default | Number of denoising steps (higher = better quality, slower) |
|
||||
| `seed` | int | null | Random seed for reproducible generation |
|
||||
|
||||
### Response Format
|
||||
|
||||
Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/wav`).
|
||||
|
||||
## Tuning Tips
|
||||
|
||||
1. **Audio Length**: Keep under 47 seconds for `stable-audio-open-1.0`.
|
||||
2. **Quality vs Speed**:
|
||||
- 50 steps: Fast, decent quality (quick previews)
|
||||
- 100 steps: Good balance (general purpose)
|
||||
- 150+ steps: High quality, slower (final / critical audio)
|
||||
3. **Guidance Scale**:
|
||||
- Lower (3 - 5): More creative / varied output
|
||||
- Default (7): Good balance
|
||||
- Higher (10+): Strict adherence to the prompt
|
||||
4. **Negative Prompts**: Use to avoid unwanted characteristics such as `"Low quality"`, `"distorted"`, `"noisy"`.
|
||||
5. **Seeds**: Set a fixed seed to get deterministic, reproducible results.
|
||||
|
||||
## File Description
|
||||
|
||||
| File | Description |
|
||||
|------|-------------|
|
||||
| `curl_examples.sh` | Curl examples covering common use cases |
|
||||
| `stable_audio_client.py` | Python client with full CLI argument support |
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
1. **Audio generation model did not produce audio output**: Verify the server started successfully and the model loaded without errors.
|
||||
2. **Connection refused**: Make sure the server is running on the correct port.
|
||||
3. **Generation timeout**: Reduce `num_inference_steps` or `audio_length`, and check GPU memory with `nvidia-smi`.
|
||||
4. **Out of memory**: Lower `--gpu-memory-utilization` or reduce `audio_length`.
|
||||
5. **Audio quality issues**: Increase `num_inference_steps`, add a negative prompt, or raise `guidance_scale`.
|
||||
|
||||
## Example materials
|
||||
|
||||
??? abstract "stable_audio_client.py"
|
||||
``````py
|
||||
--8<-- "examples/online_serving/stable_audio/stable_audio_client.py"
|
||||
``````
|
||||
??? abstract "curl_examples.sh"
|
||||
``````sh
|
||||
--8<-- "examples/online_serving/stable_audio/curl_examples.sh"
|
||||
``````
|
||||
@@ -0,0 +1,489 @@
|
||||
# Text-To-Speech (Online Serving)
|
||||
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/text_to_speech>.
|
||||
|
||||
|
||||
vLLM-Omni exposes TTS models through the OpenAI-compatible
|
||||
[`POST /v1/audio/speech`](https://github.com/vllm-project/vllm-omni/tree/main/docs/serving/speech_api.md) endpoint,
|
||||
launched with `vllm serve <model> --omni`. Each TTS model has its own
|
||||
subdirectory containing client snippets, gradio demos, and helper
|
||||
scripts; this README is the single doc entry point for all of them.
|
||||
|
||||
For offline inference, see [`examples/offline_inference/text_to_speech`](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_speech/README.md).
|
||||
For the full list of supported architectures across all modalities, see
|
||||
[Supported Models](https://github.com/vllm-project/vllm-omni/tree/main/docs/models/supported_models.md).
|
||||
|
||||
## Supported Models
|
||||
|
||||
| Model | HuggingFace repo | Voice cloning | Streaming | Voice presets / upload | Gradio demo |
|
||||
|---|---|---|---|---|---|
|
||||
| Fish Speech S2 Pro | `fishaudio/s2-pro` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | ✓ |
|
||||
| OmniVoice | `k2-fsa/OmniVoice` | (offline only) | — | — | — |
|
||||
| Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | ✓ (Base) | ✓ (PCM + WebSocket) | ✓ (presets + `/v1/audio/voices` upload) | ✓ (standard + FastRTC) |
|
||||
| VoxCPM | local model dir | ✓ | ✓ (PCM stream) | — | — |
|
||||
| VoxCPM2 | `openbmb/VoxCPM2` | ✓ | ✓ (AudioWorklet via gradio) | — | ✓ |
|
||||
| Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✓ (gated upstream) | ✓ | ✓ (presets) | ✓ |
|
||||
|
||||
CosyVoice3 is intentionally absent: no online example exists for it yet. See its [offline section](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_speech/README.md#cosyvoice3) instead.
|
||||
|
||||
## Common Quick Start
|
||||
|
||||
Launch the server (defaults shown — adjust `--port`, `--gpu-memory-utilization`, etc. as needed):
|
||||
|
||||
```bash
|
||||
vllm serve <hf-repo-or-local-path> --omni --port 8091
|
||||
```
|
||||
|
||||
Send a TTS request via curl:
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "Hello, how are you?",
|
||||
"voice": "default",
|
||||
"response_format": "wav"
|
||||
}' --output output.wav
|
||||
```
|
||||
|
||||
Or via Python httpx:
|
||||
|
||||
```python
|
||||
import httpx
|
||||
|
||||
response = httpx.post(
|
||||
"http://localhost:8091/v1/audio/speech",
|
||||
json={
|
||||
"input": "Hello, how are you?",
|
||||
"voice": "default",
|
||||
"response_format": "wav",
|
||||
},
|
||||
timeout=300.0,
|
||||
)
|
||||
open("output.wav", "wb").write(response.content)
|
||||
```
|
||||
|
||||
Or via the OpenAI SDK:
|
||||
|
||||
```python
|
||||
from openai import OpenAI
|
||||
|
||||
client = OpenAI(base_url="http://localhost:8091/v1", api_key="none")
|
||||
response = client.audio.speech.create(
|
||||
model="<hf-repo>",
|
||||
voice="default",
|
||||
input="Hello, how are you?",
|
||||
)
|
||||
response.stream_to_file("output.wav")
|
||||
```
|
||||
|
||||
Streaming PCM output (where supported) — set `stream=true` with `response_format="pcm"`:
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "Hello, how are you?",
|
||||
"voice": "default",
|
||||
"stream": true,
|
||||
"response_format": "pcm"
|
||||
}' --no-buffer | play -t raw -r 24000 -e signed -b 16 -c 1 -
|
||||
```
|
||||
|
||||
Adjust the player's sample rate to match the model (44.1 kHz for Fish Speech, 48 kHz for VoxCPM2, 24 kHz for the others).
|
||||
|
||||
For full request-shape documentation (all parameters, response formats, error codes), see the [Speech API reference](https://github.com/vllm-project/vllm-omni/tree/main/docs/serving/speech_api.md).
|
||||
|
||||
---
|
||||
|
||||
## Fish Speech S2 Pro
|
||||
|
||||
4B dual-AR TTS at 44.1 kHz. Server uses the DAC codec.
|
||||
|
||||
### Prerequisites
|
||||
```bash
|
||||
pip install fish-speech
|
||||
```
|
||||
|
||||
### Launch
|
||||
```bash
|
||||
vllm serve fishaudio/s2-pro --omni --port 8091
|
||||
# or:
|
||||
./fish_speech/run_server.sh
|
||||
```
|
||||
The deploy config auto-loads from `vllm_omni/deploy/fish_qwen3_omni.yaml` (the HF `model_type` on the fishaudio checkpoint is `fish_qwen3_omni`).
|
||||
|
||||
### Voice cloning
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "Hello, this is a cloned voice.",
|
||||
"voice": "default",
|
||||
"ref_audio": "https://example.com/reference.wav",
|
||||
"ref_text": "Transcript of the reference audio."
|
||||
}' --output cloned.wav
|
||||
```
|
||||
|
||||
### CLI client
|
||||
```bash
|
||||
cd examples/online_serving/text_to_speech/fish_speech
|
||||
python speech_client.py --text "Hello, how are you?"
|
||||
python speech_client.py --text "Hello world" --stream --output output.pcm
|
||||
```
|
||||
|
||||
### Gradio demo
|
||||
```bash
|
||||
./fish_speech/run_gradio_demo.sh # launches server + Gradio
|
||||
python fish_speech/gradio_demo.py --api-base http://localhost:8091 # if server already running
|
||||
```
|
||||
|
||||
### Notes
|
||||
- Output: 44.1 kHz mono.
|
||||
- Streaming PCM player command must use `-r 44100`.
|
||||
|
||||
---
|
||||
|
||||
## OmniVoice
|
||||
|
||||
Zero-shot multilingual TTS (600+ languages). Online serving currently exposes **auto voice** only; voice cloning and voice design are available offline.
|
||||
|
||||
### Prerequisites
|
||||
```bash
|
||||
huggingface-cli download k2-fsa/OmniVoice
|
||||
```
|
||||
Voice cloning (offline) needs `transformers>=5.3.0`; auto voice works with `transformers>=4.57.0`.
|
||||
|
||||
### Launch
|
||||
```bash
|
||||
vllm serve k2-fsa/OmniVoice --omni --port 8091 --trust-remote-code
|
||||
# or:
|
||||
./omnivoice/run_server.sh
|
||||
```
|
||||
|
||||
### CLI client
|
||||
```bash
|
||||
cd examples/online_serving/text_to_speech/omnivoice
|
||||
python speech_client.py --text "Hello, how are you?"
|
||||
python speech_client.py --text "Bonjour, comment allez-vous?" --language French
|
||||
```
|
||||
|
||||
The client supports `--api-base`, `--model`, `--text`, `--response-format`, `--language`, `--output`.
|
||||
|
||||
### Notes
|
||||
- Voice cloning and voice design require offline inference; see the [offline OmniVoice section](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_speech/README.md#omnivoice).
|
||||
|
||||
---
|
||||
|
||||
## Qwen3-TTS
|
||||
|
||||
Three model variants exposed via separate checkpoints:
|
||||
|
||||
| Variant | HF repo | Use |
|
||||
|---|---|---|
|
||||
| CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | Predefined speakers (`vivian`, `ryan`, …) with optional style instructions |
|
||||
| VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | Natural-language voice style description |
|
||||
| Base | `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | Voice cloning from a reference audio |
|
||||
|
||||
Each variant ships smaller `0.6B` companions where available.
|
||||
|
||||
### Launch
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni --port 8091
|
||||
# or:
|
||||
./qwen3_tts/run_server.sh # default: CustomVoice
|
||||
./qwen3_tts/run_server.sh VoiceDesign
|
||||
./qwen3_tts/run_server.sh Base
|
||||
```
|
||||
|
||||
### Choosing an executor backend (uniproc vs mp)
|
||||
Stage configs ship with the chunked-streaming default. To use the uniproc executor (lower IPC overhead for the Base cloning task), pass `--stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts_uniproc.yaml`. See [#2603](https://github.com/vllm-project/vllm-omni/issues/2603) and [#2604](https://github.com/vllm-project/vllm-omni/pull/2604) for the full investigation.
|
||||
|
||||
To opt out of chunked streaming, pass `--no-async-chunk` instead — the pipeline auto-dispatches to the end-to-end codec processor.
|
||||
|
||||
### Sending requests
|
||||
```bash
|
||||
# CustomVoice with a predefined speaker
|
||||
python qwen3_tts/openai_speech_client.py \
|
||||
--model Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
|
||||
--text "今天天气真好" \
|
||||
--voice ryan \
|
||||
--instructions "用开心的语气说"
|
||||
|
||||
# VoiceDesign with a style description
|
||||
python qwen3_tts/openai_speech_client.py \
|
||||
--model Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign \
|
||||
--task-type VoiceDesign \
|
||||
--text "哥哥,你回来啦" \
|
||||
--instructions "体现撒娇稚嫩的萝莉女声,音调偏高"
|
||||
|
||||
# Base voice cloning
|
||||
python qwen3_tts/openai_speech_client.py \
|
||||
--model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
|
||||
--task-type Base \
|
||||
--text "Hello, this is a cloned voice" \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--ref-text "Original transcript of the reference audio"
|
||||
```
|
||||
|
||||
### Voices endpoint
|
||||
List available voices, or upload a custom one for Base cloning:
|
||||
```bash
|
||||
# List
|
||||
curl http://localhost:8091/v1/audio/voices
|
||||
|
||||
# Upload
|
||||
curl -X POST http://localhost:8091/v1/audio/voices \
|
||||
-F "audio_sample=@/path/to/voice_sample.wav" \
|
||||
-F "consent=user_consent_id" \
|
||||
-F "name=custom_voice_1" \
|
||||
-F "ref_text=The exact transcript of the audio sample." \
|
||||
-F "speaker_description=warm narrator"
|
||||
```
|
||||
Uploaded voices are then usable as `voice="custom_voice_1"` on subsequent requests.
|
||||
|
||||
### Streaming PCM
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "Hello, how are you?",
|
||||
"voice": "vivian",
|
||||
"language": "English",
|
||||
"stream": true,
|
||||
"response_format": "pcm"
|
||||
}' --no-buffer | play -t raw -r 24000 -e signed -b 16 -c 1 -
|
||||
```
|
||||
Streaming requires `response_format="pcm"` and `async_chunk: true` on the stage config (default in `qwen3_tts.yaml`). `speed` is not supported when streaming.
|
||||
|
||||
### Streaming WebSocket
|
||||
The `/v1/audio/speech/stream` endpoint accepts text incrementally, splits it at sentence boundaries, and emits one PCM stream per sentence:
|
||||
```bash
|
||||
python qwen3_tts/streaming_speech_client.py --text "Hello world. How are you? I am fine."
|
||||
python qwen3_tts/streaming_speech_client.py --text "..." --simulate-stt --stt-delay 0.1
|
||||
```
|
||||
|
||||
### Gradio demos
|
||||
```bash
|
||||
./qwen3_tts/run_gradio_demo.sh # CustomVoice (default)
|
||||
./qwen3_tts/run_gradio_demo.sh --task-type VoiceDesign
|
||||
./qwen3_tts/run_gradio_demo.sh --task-type Base
|
||||
|
||||
# FastRTC variant (gapless WebRTC streaming):
|
||||
pip install fastrtc
|
||||
python qwen3_tts/gradio_fastrtc_demo.py --api-base http://localhost:8000
|
||||
```
|
||||
|
||||
### Speaker embedding interpolation
|
||||
`qwen3_tts/speaker_embedding_interpolation.py` blends two predefined speakers' embeddings to produce intermediate voices. See the script for usage.
|
||||
|
||||
### Batch client
|
||||
`qwen3_tts/batch_speech_client.py` issues many concurrent requests for throughput measurement.
|
||||
|
||||
### Notes
|
||||
- Base voice cloning has uniproc-vs-mp tradeoffs depending on per-request reference audio cost; see the executor-backend section above.
|
||||
- `vllm_omni/deploy/qwen3_tts.yaml` is the default deploy config (loaded by HF `model_type`); per-stage runtime overrides are available via `--stage-N-<field> <value>`.
|
||||
|
||||
---
|
||||
|
||||
## VoxCPM
|
||||
|
||||
Split-stage TTS at 24 kHz.
|
||||
|
||||
### Prerequisites
|
||||
```bash
|
||||
pip install voxcpm
|
||||
# or use a local source tree:
|
||||
export VLLM_OMNI_VOXCPM_CODE_PATH=/path/to/VoxCPM/src
|
||||
```
|
||||
|
||||
If the native VoxCPM `config.json` lacks HF `model_type`, set up an HF-compatible config dir:
|
||||
```bash
|
||||
export VOXCPM_MODEL=/path/to/voxcpm-model
|
||||
export VLLM_OMNI_VOXCPM_HF_CONFIG_PATH=/tmp/voxcpm_hf_config
|
||||
mkdir -p "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH"
|
||||
cp "$VOXCPM_MODEL/config.json" "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH/config.json"
|
||||
cp "$VOXCPM_MODEL/generation_config.json" "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH/generation_config.json" 2>/dev/null || true
|
||||
python3 -c 'import json, os; p=os.path.join(os.environ["VLLM_OMNI_VOXCPM_HF_CONFIG_PATH"], "config.json"); cfg=json.load(open(p, "r", encoding="utf-8")); cfg["model_type"]="voxcpm"; cfg.setdefault("architectures", ["VoxCPMForConditionalGeneration"]); json.dump(cfg, open(p, "w", encoding="utf-8"), indent=2, ensure_ascii=False)'
|
||||
```
|
||||
|
||||
### Launch
|
||||
```bash
|
||||
export VOXCPM_MODEL=/path/to/voxcpm-model
|
||||
./voxcpm/run_server.sh # async-chunk streaming (default)
|
||||
./voxcpm/run_server.sh sync # non-streaming
|
||||
```
|
||||
Or directly:
|
||||
```bash
|
||||
vllm serve "$VOXCPM_MODEL" \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml \
|
||||
--trust-remote-code --enforce-eager --omni --port 8091
|
||||
```
|
||||
|
||||
### Sending requests
|
||||
```bash
|
||||
# Basic TTS
|
||||
python voxcpm/openai_speech_client.py \
|
||||
--model "$VOXCPM_MODEL" \
|
||||
--text "This is a VoxCPM online text-to-speech example."
|
||||
|
||||
# Voice cloning
|
||||
python voxcpm/openai_speech_client.py \
|
||||
--model "$VOXCPM_MODEL" \
|
||||
--text "This sentence is synthesized with a cloned voice." \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--ref-text "The exact transcript spoken in reference.wav."
|
||||
|
||||
# Streaming PCM
|
||||
python voxcpm/openai_speech_client.py \
|
||||
--model "$VOXCPM_MODEL" \
|
||||
--text "This is a streaming VoxCPM request." \
|
||||
--stream --output voxcpm_stream.pcm
|
||||
```
|
||||
|
||||
### Notes
|
||||
- `voxcpm.yaml` for one-shot decode; `voxcpm_async_chunk.yaml` for single-request streaming. Do not use the async-chunk config for concurrent requests or `/v1/audio/speech/batch`.
|
||||
- Generic TTS fields not supported by VoxCPM: `voice`, `instructions`, `language`, `speaker_embedding`, `x_vector_only_mode`.
|
||||
- For benchmark measurement, see [`benchmarks/voxcpm`](https://github.com/vllm-project/vllm-omni/tree/main/benchmarks/voxcpm/README.md).
|
||||
|
||||
---
|
||||
|
||||
## VoxCPM2
|
||||
|
||||
Single-stage native AR TTS at 48 kHz.
|
||||
|
||||
### Launch
|
||||
```bash
|
||||
vllm serve openbmb/VoxCPM2 --omni --host 0.0.0.0 --port 8000
|
||||
```
|
||||
Deploy config auto-loads from `vllm_omni/deploy/voxcpm2.yaml`. Pass `--deploy-config <path>` to override or `--stage-N-<field> <value>` for per-stage runtime tweaks.
|
||||
|
||||
### Sending requests
|
||||
```bash
|
||||
# Zero-shot synthesis
|
||||
python voxcpm2/openai_speech_client.py --text "Hello, this is VoxCPM2."
|
||||
|
||||
# Voice cloning
|
||||
python voxcpm2/openai_speech_client.py \
|
||||
--text "This should sound like the reference speaker." \
|
||||
--ref-audio /path/to/reference.wav
|
||||
```
|
||||
The `ref_audio` field accepts local file paths (auto-base64), HTTP URLs, or `data:audio/wav;base64,...` data URIs.
|
||||
|
||||
### Gradio demo (gapless streaming via AudioWorklet)
|
||||
```bash
|
||||
python voxcpm2/gradio_demo.py
|
||||
```
|
||||
Uses an AudioWorklet-based player adapted from the Qwen3-TTS demo for gap-free playback. Audio is streamed from the OpenAI Speech endpoint with `stream=true`.
|
||||
|
||||
---
|
||||
|
||||
## Voxtral TTS
|
||||
|
||||
Voxtral-4B-TTS (Mistral). Uses the `mistral_common` `SpeechRequest` protocol; voice presets are model-specific.
|
||||
|
||||
### Prerequisites
|
||||
Latest `mistral_common` with `SpeechRequest` support:
|
||||
```bash
|
||||
pip install -e /path/to/mistral-common # or upgrade from PyPI when available
|
||||
```
|
||||
|
||||
### Launch
|
||||
```bash
|
||||
vllm serve mistralai/Voxtral-4B-TTS-2603 --omni --port 8091
|
||||
```
|
||||
Deploy config auto-loads from `vllm_omni/deploy/voxtral_tts.yaml`.
|
||||
|
||||
### Gradio demo
|
||||
```bash
|
||||
python voxtral_tts/gradio_demo.py
|
||||
```
|
||||
The demo handles voice-preset selection and reference-audio upload. `voxtral_tts/text_preprocess.py` provides the text-normalization helpers used by the demo (also available for other clients).
|
||||
|
||||
### Notes
|
||||
- Voice presets are listed on the HF model card (`mistralai/Voxtral-4B-TTS-2603`).
|
||||
- Voice cloning is gated upstream and may require a recent `mistral_common`.
|
||||
- A standalone CLI client is not yet shipped; the gradio demo is the canonical reference for now.
|
||||
|
||||
## Example materials
|
||||
|
||||
??? abstract "fish_speech/gradio_demo.py"
|
||||
``````py
|
||||
--8<-- "examples/online_serving/text_to_speech/fish_speech/gradio_demo.py"
|
||||
``````
|
||||
??? abstract "fish_speech/run_gradio_demo.sh"
|
||||
``````sh
|
||||
--8<-- "examples/online_serving/text_to_speech/fish_speech/run_gradio_demo.sh"
|
||||
``````
|
||||
??? abstract "fish_speech/run_server.sh"
|
||||
``````sh
|
||||
--8<-- "examples/online_serving/text_to_speech/fish_speech/run_server.sh"
|
||||
``````
|
||||
??? abstract "fish_speech/speech_client.py"
|
||||
``````py
|
||||
--8<-- "examples/online_serving/text_to_speech/fish_speech/speech_client.py"
|
||||
``````
|
||||
??? abstract "omnivoice/run_server.sh"
|
||||
``````sh
|
||||
--8<-- "examples/online_serving/text_to_speech/omnivoice/run_server.sh"
|
||||
``````
|
||||
??? abstract "omnivoice/speech_client.py"
|
||||
``````py
|
||||
--8<-- "examples/online_serving/text_to_speech/omnivoice/speech_client.py"
|
||||
``````
|
||||
??? abstract "qwen3_tts/batch_speech_client.py"
|
||||
``````py
|
||||
--8<-- "examples/online_serving/text_to_speech/qwen3_tts/batch_speech_client.py"
|
||||
``````
|
||||
??? abstract "qwen3_tts/gradio_demo.py"
|
||||
``````py
|
||||
--8<-- "examples/online_serving/text_to_speech/qwen3_tts/gradio_demo.py"
|
||||
``````
|
||||
??? abstract "qwen3_tts/openai_speech_client.py"
|
||||
``````py
|
||||
--8<-- "examples/online_serving/text_to_speech/qwen3_tts/openai_speech_client.py"
|
||||
``````
|
||||
??? abstract "qwen3_tts/run_gradio_demo.sh"
|
||||
``````sh
|
||||
--8<-- "examples/online_serving/text_to_speech/qwen3_tts/run_gradio_demo.sh"
|
||||
``````
|
||||
??? abstract "qwen3_tts/run_server.sh"
|
||||
``````sh
|
||||
--8<-- "examples/online_serving/text_to_speech/qwen3_tts/run_server.sh"
|
||||
``````
|
||||
??? abstract "qwen3_tts/speaker_embedding_interpolation.py"
|
||||
``````py
|
||||
--8<-- "examples/online_serving/text_to_speech/qwen3_tts/speaker_embedding_interpolation.py"
|
||||
``````
|
||||
??? abstract "qwen3_tts/streaming_speech_client.py"
|
||||
``````py
|
||||
--8<-- "examples/online_serving/text_to_speech/qwen3_tts/streaming_speech_client.py"
|
||||
``````
|
||||
??? abstract "qwen3_tts/tts_common.py"
|
||||
``````py
|
||||
--8<-- "examples/online_serving/text_to_speech/qwen3_tts/tts_common.py"
|
||||
``````
|
||||
??? abstract "voxcpm/openai_speech_client.py"
|
||||
``````py
|
||||
--8<-- "examples/online_serving/text_to_speech/voxcpm/openai_speech_client.py"
|
||||
``````
|
||||
??? abstract "voxcpm/run_server.sh"
|
||||
``````sh
|
||||
--8<-- "examples/online_serving/text_to_speech/voxcpm/run_server.sh"
|
||||
``````
|
||||
??? abstract "voxcpm2/gradio_demo.py"
|
||||
``````py
|
||||
--8<-- "examples/online_serving/text_to_speech/voxcpm2/gradio_demo.py"
|
||||
``````
|
||||
??? abstract "voxcpm2/openai_speech_client.py"
|
||||
``````py
|
||||
--8<-- "examples/online_serving/text_to_speech/voxcpm2/openai_speech_client.py"
|
||||
``````
|
||||
??? abstract "voxtral_tts/gradio_demo.py"
|
||||
``````py
|
||||
--8<-- "examples/online_serving/text_to_speech/voxtral_tts/gradio_demo.py"
|
||||
``````
|
||||
??? abstract "voxtral_tts/text_preprocess.py"
|
||||
``````py
|
||||
--8<-- "examples/online_serving/text_to_speech/voxtral_tts/text_preprocess.py"
|
||||
``````
|
||||
@@ -0,0 +1,127 @@
|
||||
# AutoRound Quantization
|
||||
|
||||
## Overview
|
||||
|
||||
[AutoRound](https://github.com/intel/auto-round) produces pre-quantized
|
||||
checkpoints for LLMs, VLMs, and diffusion models. vLLM-Omni reads the
|
||||
checkpoint's `config.json` and auto-detects
|
||||
`quantization_config.quant_method = "auto-round"`.
|
||||
|
||||
AutoRound is static quantization: no `--quantization` flag is needed at
|
||||
inference time when the checkpoint already contains the quantization config.
|
||||
|
||||
## Hardware Support
|
||||
|
||||
| Device | Support |
|
||||
|--------|---------|
|
||||
| NVIDIA Blackwell GPU (SM 100+) | ✅ |
|
||||
| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ |
|
||||
| NVIDIA Ampere GPU (SM 80+) | ✅ |
|
||||
| AMD ROCm | ⭕ |
|
||||
| Intel XPU | ✅ |
|
||||
| Ascend NPU | ❌ |
|
||||
|
||||
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
|
||||
guide. AutoRound is Intel-supported.
|
||||
|
||||
## Model Type Support
|
||||
|
||||
### Diffusion Model (Qwen-Image, Wan2.2)
|
||||
|
||||
| Model | Checkpoint | Scope | Scheme | Backend |
|
||||
|-------|------------|-------|--------|---------|
|
||||
| FLUX.1-dev | `vllm-project-org/FLUX.1-dev-AutoRound-w4a16` | Diffusion transformer | W4A16 | GPTQ-Marlin or Intel-supported AutoRound backend |
|
||||
| Qwen-Image | Not listed | Diffusion transformer | W4A16 | Not validated |
|
||||
| Wan2.2 | Not listed | Diffusion transformer | W4A16 | Not validated |
|
||||
|
||||
### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
|
||||
|
||||
| Model | Checkpoint | Scope | Scheme | Backend |
|
||||
|-------|------------|-------|--------|---------|
|
||||
| Qwen2.5-Omni-7B | `Intel/Qwen2.5-Omni-7B-int4-AutoRound` | Language-model stage | W4A16 | AutoRound |
|
||||
| Qwen3-Omni-30B-A3B-Instruct | `Intel/Qwen3-Omni-30B-A3B-Instruct-int4-AutoRound` | Thinker language-model stage | W4A16 | AutoRound |
|
||||
| Qwen3-TTS | Not listed | TTS language-model stage | W4A16 | Not validated |
|
||||
|
||||
AutoRound support is checkpoint-driven. A model is supported when its
|
||||
checkpoint uses a compatible INC/AutoRound config and the target stage maps to
|
||||
vLLM-Omni's runtime module names.
|
||||
|
||||
### Multi-Stage Diffusion Model (BAGEL, GLM-Image)
|
||||
|
||||
| Model | Scope | Status | Notes |
|
||||
|-------|-------|--------|-------|
|
||||
| BAGEL | Checkpoint-defined diffusion or transformer stage | Not validated | Requires a compatible AutoRound checkpoint |
|
||||
| GLM-Image | Checkpoint-defined diffusion or transformer stage | Not validated | Requires a compatible AutoRound checkpoint |
|
||||
|
||||
## Configuration
|
||||
|
||||
Python API:
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
omni = Omni(model="vllm-project-org/FLUX.1-dev-AutoRound-w4a16")
|
||||
|
||||
outputs = omni.generate(
|
||||
"A cat sitting on a windowsill",
|
||||
OmniDiffusionSamplingParams(num_inference_steps=28),
|
||||
)
|
||||
outputs[0].save_images("output.png")
|
||||
```
|
||||
|
||||
CLI:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_image/text_to_image.py \
|
||||
--model vllm-project-org/FLUX.1-dev-AutoRound-w4a16 \
|
||||
--prompt "A cat sitting on a windowsill" \
|
||||
--num-inference-steps 28 \
|
||||
--output outputs/flux_w4a16.png
|
||||
```
|
||||
|
||||
## Parameters
|
||||
|
||||
| Field | Type | Description |
|
||||
|-------|------|-------------|
|
||||
| `quant_method` | str | Must be `"auto-round"` |
|
||||
| `bits` | int | Quantized weight bit width, usually `4` |
|
||||
| `group_size` | int | Quantization group size |
|
||||
| `packing_format` | str | AutoRound packing format, for example `auto_round:auto_gptq` |
|
||||
| `block_name_to_quantize` | str | Checkpoint block names that should map to runtime module names |
|
||||
|
||||
The checkpoint should contain a config like:
|
||||
|
||||
```json
|
||||
{
|
||||
"quantization_config": {
|
||||
"quant_method": "auto-round",
|
||||
"bits": 4,
|
||||
"group_size": 128,
|
||||
"sym": true,
|
||||
"packing_format": "auto_round:auto_gptq",
|
||||
"block_name_to_quantize": "transformer_blocks,single_transformer_blocks"
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
## Validation and Notes
|
||||
|
||||
At load time, vLLM-Omni builds an `OmniINCConfig`, remaps checkpoint block names
|
||||
to runtime module names, and selects the matching vLLM compute backend.
|
||||
|
||||
Example checkpoint creation:
|
||||
|
||||
```bash
|
||||
auto-round \
|
||||
--model black-forest-labs/FLUX.1-dev \
|
||||
--scheme W4A16 \
|
||||
--batch_size 1 \
|
||||
--disable_opt_rtn \
|
||||
--dataset coco2014 \
|
||||
--iters 0
|
||||
```
|
||||
|
||||
Use the generated output directory directly as the `model` argument. See the
|
||||
[AutoRound documentation](https://github.com/intel/auto-round) for all
|
||||
available schemes and options.
|
||||
@@ -0,0 +1,121 @@
|
||||
# FP8 Quantization
|
||||
|
||||
## Overview
|
||||
|
||||
FP8 quantization converts BF16/FP16 weights to FP8 at model load time, or loads
|
||||
a checkpoint whose target stage already declares an FP8 quantization config.
|
||||
Online activation scaling is the default and does not require calibration.
|
||||
Static activation scaling is supported when calibrated scale information is
|
||||
available.
|
||||
|
||||
Some architectures can quantize all linear layers. Others have
|
||||
quality-sensitive layers that should stay in BF16 through `ignored_layers`.
|
||||
Image-stream MLPs (`img_mlp`) are a common sensitive target because denoising
|
||||
latent ranges shift across timesteps and small per-layer errors can compound
|
||||
in deep DiT blocks.
|
||||
|
||||
## Hardware Support
|
||||
|
||||
| Device | Support |
|
||||
|--------|---------|
|
||||
| NVIDIA Blackwell GPU (SM 100+) | ✅ |
|
||||
| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ |
|
||||
| NVIDIA Ampere GPU (SM 80+) | ✅ |
|
||||
| AMD ROCm | ⭕ |
|
||||
| Intel XPU | ⭕ |
|
||||
| Ascend NPU | ❌ |
|
||||
|
||||
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
|
||||
guide. FP8 on Ampere may use a weight-only path where available.
|
||||
|
||||
## Model Type Support
|
||||
|
||||
### Diffusion Model (Qwen-Image, Wan2.2)
|
||||
|
||||
| Model | HF models | Online | Pre-calibrated | Recommendation | `ignored_layers` | Text-Encoder quantization |
|
||||
|-------|-----------|:-------:|:------:|----------------|------------------|------------------|
|
||||
| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Yes | Yes | Skip sensitive image-stream MLPs when quality regresses | `img_mlp` | |
|
||||
| Wan2.2 | Wan2.2 diffusion pipelines | Not validated | Not validated | Validate against BF16 before documenting as supported | TBD | |
|
||||
| Z-Image | `Tongyi-MAI/Z-Image-Turbo` | Yes | Yes | All layers | None | ✅︎ |
|
||||
| FLUX.1 | `black-forest-labs/FLUX.1-dev`, `black-forest-labs/FLUX.1-schnell` | Yes | Yes | All layers | None | |
|
||||
| FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B` | Yes | Yes | All layers | None | |
|
||||
| HunyuanImage-3.0 | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | Yes | Yes | All layers; use the Hunyuan stage config for multi-stage runs | None | |
|
||||
| HunyuanVideo-1.5 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `720p_t2v`, `480p_i2v` | Yes | Yes | All layers | None | |
|
||||
|
||||
### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
|
||||
|
||||
| Model | Scope | Format | Status |
|
||||
|-------|-------|--------|--------|
|
||||
| Qwen3-Omni | Thinker language-model stage | ModelOpt `quant_algo=FP8` | Tested for thinker memory reduction |
|
||||
| Qwen3-TTS | TTS language-model stage | Checkpoint config | Not validated |
|
||||
|
||||
Audio encoder, vision encoder, talker, and code2wav stay in BF16 unless a
|
||||
model-specific guide says otherwise.
|
||||
|
||||
### Multi-Stage Diffusion Model (BAGEL, GLM-Image)
|
||||
|
||||
| Model | Scope | Status | Notes |
|
||||
|-------|-------|--------|-------|
|
||||
| BAGEL | Stage-specific transformer or DiT module | Not validated | Route FP8 to the intended stage before enabling |
|
||||
| GLM-Image | Stage-specific transformer or DiT module | Not validated | Validate quality against BF16 baseline |
|
||||
|
||||
## Configuration
|
||||
|
||||
Python API:
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
omni = Omni(model="<your-model>", quantization="fp8")
|
||||
|
||||
omni_with_skips = Omni(
|
||||
model="<your-model>",
|
||||
quantization_config={
|
||||
"method": "fp8",
|
||||
"ignored_layers": ["img_mlp"],
|
||||
},
|
||||
)
|
||||
|
||||
outputs = omni.generate(
|
||||
"A cat sitting on a windowsill",
|
||||
OmniDiffusionSamplingParams(num_inference_steps=50),
|
||||
)
|
||||
```
|
||||
|
||||
CLI:
|
||||
|
||||
```bash
|
||||
python text_to_image.py --model <your-model> --quantization fp8
|
||||
python text_to_image.py --model <your-model> --quantization fp8 --ignored-layers "img_mlp"
|
||||
vllm serve <your-model> --omni --quantization fp8
|
||||
```
|
||||
|
||||
## Parameters
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| `method` | str | - | Quantization method (`"fp8"`) |
|
||||
| `ignored_layers` | list[str] | `[]` | Layer name patterns to keep in BF16 |
|
||||
| `activation_scheme` | str | `"dynamic"` | `"dynamic"` selects online activation scaling, or `"static"` when scales are available |
|
||||
| `weight_block_size` | list[int] \| None | `None` | Block size for block-wise weight quantization |
|
||||
|
||||
The available `ignored_layers` names depend on the model architecture, for
|
||||
example `to_qkv`, `to_out`, `img_mlp`, or `txt_mlp`.
|
||||
|
||||
## Validation and Notes
|
||||
|
||||
FP8 quantization can be combined with cache acceleration:
|
||||
|
||||
```python
|
||||
omni = Omni(
|
||||
model="<your-model>",
|
||||
quantization="fp8",
|
||||
cache_backend="tea_cache",
|
||||
cache_config={"rel_l1_thresh": 0.2},
|
||||
)
|
||||
```
|
||||
|
||||
Compare generated outputs with a BF16 baseline before adding a new model to the
|
||||
supported table. GLM-Image and Helios are not listed as FP8-supported diffusion
|
||||
models until they have method-specific validation.
|
||||
@@ -0,0 +1,107 @@
|
||||
# GGUF Quantization
|
||||
|
||||
## Overview
|
||||
|
||||
GGUF support loads pre-quantized diffusion transformer weights while keeping
|
||||
the rest of the pipeline on the base Hugging Face checkpoint. Use the base
|
||||
model for tokenizer, text encoder, scheduler, and VAE, then pass the GGUF file
|
||||
for the transformer.
|
||||
|
||||
GGUF is static quantization: the quantized weights are produced before serving.
|
||||
|
||||
## Hardware Support
|
||||
|
||||
| Device | Support |
|
||||
|--------|---------|
|
||||
| NVIDIA Blackwell GPU (SM 100+) | ✅ |
|
||||
| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ |
|
||||
| NVIDIA Ampere GPU (SM 80+) | ✅ |
|
||||
| AMD ROCm | ⭕ |
|
||||
| Intel XPU | ⭕ |
|
||||
| Ascend NPU | ❌ |
|
||||
|
||||
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
|
||||
guide.
|
||||
|
||||
## Model Type Support
|
||||
|
||||
### Diffusion Model (Qwen-Image, Wan2.2)
|
||||
|
||||
| Model | HF base model | GGUF input | Scope | Adapter |
|
||||
|-------|---------------|------------|-------|---------|
|
||||
| Qwen-Image family | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512`, edit and layered Qwen-Image pipelines | Local `.gguf`, `repo/file.gguf`, or `repo:quant_type` | Transformer only | `QwenImageGGUFAdapter` |
|
||||
| Wan2.2 | Wan2.2 diffusion pipelines | Not validated | Transformer only | No validated adapter listed |
|
||||
| Z-Image | `Tongyi-MAI/Z-Image-Turbo` | Local `.gguf`, `repo/file.gguf`, or `repo:quant_type` | Transformer only | `ZImageGGUFAdapter` |
|
||||
| FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B` | Local `.gguf`, `repo/file.gguf`, or `repo:quant_type` | Transformer only | `Flux2KleinGGUFAdapter` |
|
||||
|
||||
Generic FLUX.1 GGUF checkpoints are not listed here; the implemented adapter is
|
||||
for the FLUX.2-klein path.
|
||||
|
||||
### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
|
||||
|
||||
| Model | Scope | Status | Notes |
|
||||
|-------|-------|--------|-------|
|
||||
| Qwen3-Omni | Thinker language-model stage | Not validated | GGUF is not documented for omni/TTS AR stages |
|
||||
| Qwen3-TTS | TTS language-model stage | Not validated | GGUF is not documented for TTS stages |
|
||||
|
||||
### Multi-Stage Diffusion Model (BAGEL, GLM-Image)
|
||||
|
||||
| Model | Scope | Status | Notes |
|
||||
|-------|-------|--------|-------|
|
||||
| BAGEL | Stage-specific transformer weights | Not validated | Requires a model-specific GGUF adapter |
|
||||
| GLM-Image | Stage-specific transformer weights | Not validated | Requires a model-specific GGUF adapter |
|
||||
|
||||
## Configuration
|
||||
|
||||
Offline:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_image/text_to_image.py \
|
||||
--model Qwen/Qwen-Image \
|
||||
--gguf-model QuantStack/Qwen-Image-GGUF/Qwen_Image-Q4_K_M.gguf \
|
||||
--quantization gguf \
|
||||
--prompt "a red paper kite hanging from a pine tree in a winter courtyard" \
|
||||
--height 1024 \
|
||||
--width 1024 \
|
||||
--seed 42 \
|
||||
--num_inference_steps 20 \
|
||||
--output outputs/qwen_image_q4km.png
|
||||
```
|
||||
|
||||
Online:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen-Image \
|
||||
--omni \
|
||||
--port 8000 \
|
||||
--quantization-config '{"method":"gguf","gguf_model":"QuantStack/Qwen-Image-GGUF/Qwen_Image-Q4_K_M.gguf"}'
|
||||
```
|
||||
|
||||
## Parameters
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| `method` | str | - | Quantization method (`"gguf"`) |
|
||||
| `gguf_model` | str | - | Local GGUF file, explicit Hugging Face file, or `repo:quant_type` selector |
|
||||
|
||||
`gguf_model` accepts:
|
||||
|
||||
| Form | Example |
|
||||
|------|---------|
|
||||
| Local file | `/models/z-image-Q4_K_M.gguf` |
|
||||
| Explicit HF file | `QuantStack/Qwen-Image-GGUF/Qwen_Image-Q4_K_M.gguf` |
|
||||
| HF repo plus quant type | `owner/repo:Q4_K_M` |
|
||||
|
||||
## Validation and Notes
|
||||
|
||||
1. `OmniDiffusionConfig` receives `{"method": "gguf", "gguf_model": ...}`.
|
||||
2. `DiffusersPipelineLoader` resolves the GGUF file.
|
||||
3. A model-specific adapter remaps GGUF tensor names to vLLM-Omni transformer
|
||||
names.
|
||||
4. Only transformer weights are loaded from GGUF. Missing non-transformer
|
||||
weights are loaded from the base model repository.
|
||||
5. vLLM's GGUF linear method performs dequantization and GEMM at runtime.
|
||||
|
||||
Unsupported models fail fast with a clear "No GGUF adapter matched" error
|
||||
instead of falling back to a generic mapper. Many GGUF repositories do not
|
||||
include `model_index.json`; always pass the normal base model through `--model`.
|
||||
@@ -0,0 +1,107 @@
|
||||
# Int8 Quantization
|
||||
|
||||
## Overview
|
||||
|
||||
Int8 quantization supports W8A8 diffusion transformer inference on CUDA and
|
||||
Ascend NPU. It can quantize BF16/FP16 weights at load time, or load serialized
|
||||
Int8 checkpoints that already contain quantized weights and scales.
|
||||
|
||||
Only online activation scaling is currently supported.
|
||||
|
||||
## Hardware Support
|
||||
|
||||
| Device | Support |
|
||||
|--------|---------|
|
||||
| NVIDIA Blackwell GPU (SM 100+) | ✅ |
|
||||
| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ |
|
||||
| NVIDIA Ampere GPU (SM 80+) | ✅ |
|
||||
| AMD ROCm | ⭕ |
|
||||
| Intel XPU | ⭕ |
|
||||
| Ascend NPU | ✅ |
|
||||
|
||||
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
|
||||
guide.
|
||||
|
||||
## Model Type Support
|
||||
|
||||
### Diffusion Model (Qwen-Image, Wan2.2)
|
||||
|
||||
| Model | HF models | CUDA | Ascend NPU | Mode | Recommendation |
|
||||
|-------|-----------|:----:|:----------:|------|----------------|
|
||||
| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Yes | Yes | Online W8A8 | All layers |
|
||||
| Wan2.2 | Wan2.2 diffusion pipelines | Not validated | Not validated | Online W8A8 | Validate before enabling in docs |
|
||||
| Z-Image | `Tongyi-MAI/Z-Image-Turbo` | Yes | Yes | Online W8A8 | All layers |
|
||||
|
||||
Other diffusion models may work if their transformer uses supported linear
|
||||
layers, but they are not validated in this guide.
|
||||
|
||||
### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
|
||||
|
||||
| Model | Scope | Status | Notes |
|
||||
|-------|-------|--------|-------|
|
||||
| Qwen3-Omni | Thinker language-model stage | Not validated | Prefer checkpoint-supported ModelOpt FP8 or AutoRound paths |
|
||||
| Qwen3-TTS | TTS language-model stage | Not validated | No Int8 TTS stage support is documented |
|
||||
|
||||
### Multi-Stage Diffusion Model (BAGEL, GLM-Image)
|
||||
|
||||
| Model | Scope | Status | Notes |
|
||||
|-------|-------|--------|-------|
|
||||
| BAGEL | Stage-specific transformer or DiT module | Not validated | Requires explicit stage routing |
|
||||
| GLM-Image | Stage-specific transformer or DiT module | Not validated | Requires quality comparison with BF16 |
|
||||
|
||||
## Configuration
|
||||
|
||||
Python API:
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
omni = Omni(model="<your-model>", quantization="int8")
|
||||
|
||||
omni_with_skips = Omni(
|
||||
model="<your-model>",
|
||||
quantization_config={
|
||||
"method": "int8",
|
||||
"ignored_layers": ["<layer-name>"],
|
||||
},
|
||||
)
|
||||
|
||||
outputs = omni.generate(
|
||||
"A cat sitting on a windowsill",
|
||||
OmniDiffusionSamplingParams(num_inference_steps=50),
|
||||
)
|
||||
```
|
||||
|
||||
CLI:
|
||||
|
||||
```bash
|
||||
python text_to_image.py --model <your-model> --quantization int8
|
||||
python text_to_image.py --model <your-model> --quantization int8 --ignored-layers "img_mlp"
|
||||
vllm serve <your-model> --omni --quantization int8
|
||||
```
|
||||
|
||||
## Parameters
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| `method` | str | - | Quantization method (`"int8"`) |
|
||||
| `activation_scheme` | str | `"dynamic"` | `"dynamic"` selects online activation scaling; static is not supported |
|
||||
| `ignored_layers` | list[str] | `[]` | Layer name patterns to keep in BF16/FP16 |
|
||||
| `is_checkpoint_int8_serialized` | bool | `False` | Set by checkpoint config when loading serialized Int8 weights |
|
||||
|
||||
## Validation and Notes
|
||||
|
||||
Int8 quantization can be combined with cache acceleration:
|
||||
|
||||
```python
|
||||
omni = Omni(
|
||||
model="<your-model>",
|
||||
quantization="int8",
|
||||
cache_backend="tea_cache",
|
||||
cache_config={"rel_l1_thresh": 0.2},
|
||||
)
|
||||
```
|
||||
|
||||
Only add a new model to the supported table after comparing the Int8 output
|
||||
against a BF16 baseline and documenting any required `ignored_layers`.
|
||||
@@ -0,0 +1,95 @@
|
||||
# msModelSlim Quantization
|
||||
|
||||
## Overview
|
||||
|
||||
[msModelSlim](https://github.com/Ascend/msmodelslim) is an Ascend compression
|
||||
toolkit for producing pre-quantized model checkpoints. In vLLM-Omni, these
|
||||
checkpoints run through the Ascend/NPU path with `--quantization ascend`.
|
||||
|
||||
msModelSlim is static quantization: quantized weights are generated offline
|
||||
before vLLM-Omni inference starts.
|
||||
|
||||
## Hardware Support
|
||||
|
||||
| Device | Support |
|
||||
|--------|---------|
|
||||
| NVIDIA Blackwell GPU (SM 100+) | ❌ |
|
||||
| NVIDIA Ada/Hopper GPU (SM 89+) | ❌ |
|
||||
| NVIDIA Ampere GPU (SM 80+) | ❌ |
|
||||
| AMD ROCm | ❌ |
|
||||
| Intel XPU | ❌ |
|
||||
| Ascend NPU | ✅ |
|
||||
|
||||
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
|
||||
guide.
|
||||
|
||||
## Model Type Support
|
||||
|
||||
### Diffusion Model (Qwen-Image, Wan2.2)
|
||||
|
||||
| Model | Base model | Scope | Hardware | Notes |
|
||||
|-------|------------|-------|----------|-------|
|
||||
| Wan2.2 | Wan2.2 diffusion weights | DiT or diffusion stage | Ascend NPU | Upstream msModelSlim provides a Wan2.2 quantization recipe; vLLM-Omni inference validation is not listed |
|
||||
| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | DiT or diffusion stage | Ascend NPU | Not validated in this guide |
|
||||
| HunyuanImage-3.0 | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | DiT or diffusion stage | Ascend A2/A3 NPU | Generate quantized weights with the HunyuanImage-3.0 msModelSlim adaptation |
|
||||
|
||||
Public Hugging Face quantized weights are not available yet. Use the
|
||||
[HunyuanImage-3.0 msModelSlim adaptation](https://gitcode.com/betta18/msmodelslim/tree/hyimage3_mxfp8)
|
||||
to generate the checkpoint manually.
|
||||
|
||||
### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
|
||||
|
||||
| Model | Scope | Status | Notes |
|
||||
|-------|-------|--------|-------|
|
||||
| Qwen3-Omni | Thinker or language-model stage | Not validated | No msModelSlim omni checkpoint path is documented |
|
||||
| Qwen3-TTS | TTS language-model stage | Not validated | No msModelSlim TTS checkpoint path is documented |
|
||||
|
||||
### Multi-Stage Diffusion Model (BAGEL, GLM-Image)
|
||||
|
||||
| Model | Scope | Status | Notes |
|
||||
|-------|-------|--------|-------|
|
||||
| BAGEL | Stage-specific diffusion or transformer weights | Not validated | Requires a model-specific Ascend adaptation |
|
||||
| GLM-Image | Stage-specific diffusion or transformer weights | Not validated | Requires a model-specific Ascend adaptation |
|
||||
|
||||
## Configuration
|
||||
|
||||
Offline inference:
|
||||
|
||||
```bash
|
||||
python text_to_image.py --model <quantized-model-path> --quantization ascend
|
||||
```
|
||||
|
||||
Online serving:
|
||||
|
||||
```bash
|
||||
vllm serve <quantized-model-path> --omni --quantization ascend
|
||||
```
|
||||
|
||||
## Parameters
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| `--quantization` | str | - | Use `ascend` for msModelSlim-produced checkpoints |
|
||||
| `model` | str | - | Path to the quantized checkpoint generated by Ascend tooling |
|
||||
|
||||
Example msModelSlim command for a Wan2.2 W8A8 checkpoint:
|
||||
|
||||
```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 /path/to/wan2_2_w8a8f8_mxfp_t2v.yaml \
|
||||
--trust_remote_code True
|
||||
```
|
||||
|
||||
For HunyuanImage-3.0, use the Hunyuan-specific adaptation linked above.
|
||||
|
||||
## Validation and Notes
|
||||
|
||||
1. Run with the Ascend/NPU installation and environment.
|
||||
2. The `ascend` quantization method expects weights produced by the Ascend
|
||||
tooling; it is not a load-time CUDA quantizer.
|
||||
3. Keep the quantized checkpoint aligned with the same model architecture and
|
||||
stage config used for BF16 inference.
|
||||
@@ -0,0 +1,93 @@
|
||||
# Online Quantization
|
||||
|
||||
## Overview
|
||||
|
||||
Online quantization means vLLM-Omni computes quantized weights and scales while
|
||||
loading the model. Use it when you want memory savings without preparing a
|
||||
separate quantized checkpoint.
|
||||
|
||||
This mode is different from pre-quantized checkpoint formats such as GGUF,
|
||||
AutoRound, msModelSlim, or serialized Int8 checkpoints. Those formats are
|
||||
prepared before serving and are documented in their method-specific guides.
|
||||
|
||||
## Hardware Support
|
||||
|
||||
| Device | FP8 W8A8 | Int8 W8A8 |
|
||||
|--------|----------|-----------|
|
||||
| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ |
|
||||
| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ |
|
||||
| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ |
|
||||
| AMD ROCm | ⭕ | ⭕ |
|
||||
| Intel XPU | ⭕ | ⭕ |
|
||||
| Ascend NPU | ❌ | ✅ |
|
||||
|
||||
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
|
||||
guide. FP8 on Ampere may use a weight-only path where available.
|
||||
|
||||
## Model Type Support
|
||||
|
||||
### Diffusion Model (Qwen-Image, Wan2.2)
|
||||
|
||||
| Method | Guide | Example models | Status |
|
||||
|--------|-------|----------------|--------|
|
||||
| FP8 W8A8 | [FP8](fp8.md) | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image family and other DiT models |
|
||||
| Int8 W8A8 | [Int8](int8.md) | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image and Z-Image |
|
||||
|
||||
### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
|
||||
|
||||
Online quantization is not currently validated for the omni/TTS stages. For
|
||||
Qwen3-Omni and related models, prefer checkpoint-declared ModelOpt or
|
||||
AutoRound paths when available.
|
||||
|
||||
### Multi-Stage Diffusion Model (BAGEL, GLM-Image)
|
||||
|
||||
Online quantization must be routed to the intended stage. BAGEL and GLM-Image
|
||||
need model-specific validation before they are listed as supported targets.
|
||||
|
||||
## Configuration
|
||||
|
||||
Python API:
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
|
||||
omni_fp8 = Omni(model="<your-model>", quantization="fp8")
|
||||
omni_int8 = Omni(model="<your-model>", quantization="int8")
|
||||
```
|
||||
|
||||
CLI:
|
||||
|
||||
```bash
|
||||
vllm serve <your-model> --omni --quantization fp8
|
||||
vllm serve <your-model> --omni --quantization int8
|
||||
```
|
||||
|
||||
Per-component routing:
|
||||
|
||||
```python
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
|
||||
config = build_quant_config({
|
||||
"transformer": {"method": "fp8"},
|
||||
"vae": None,
|
||||
})
|
||||
```
|
||||
|
||||
## Parameters
|
||||
|
||||
| Parameter | Methods | Description |
|
||||
|-----------|---------|-------------|
|
||||
| `method` | FP8, Int8 | Quantization method, usually `"fp8"` or `"int8"` |
|
||||
| `ignored_layers` | FP8, Int8 | Layer name patterns to keep in BF16/FP16 |
|
||||
| `activation_scheme` | FP8, Int8 | The runtime value `"dynamic"` selects online activation scaling |
|
||||
| `weight_block_size` | FP8 | Optional block-wise FP8 weight quantization size |
|
||||
|
||||
## Validation and Notes
|
||||
|
||||
1. Compare the online-quantized output against a BF16 baseline with the same
|
||||
seed and generation parameters.
|
||||
2. Use `ignored_layers` for quality-sensitive MLPs or output projections.
|
||||
3. Document any required skipped layers in the method page before marking a new
|
||||
model as supported.
|
||||
4. If a model already ships quantized weights, use the matching pre-quantized
|
||||
method guide instead of online quantization.
|
||||
@@ -0,0 +1,134 @@
|
||||
# Quantization
|
||||
|
||||
vLLM-Omni exposes quantization through the unified `quantization_config`
|
||||
path. The same configuration entrypoint is used across diffusion-only models,
|
||||
multi-stage omni/TTS models, and multi-stage diffusion models, but each model
|
||||
type has a different quantization scope.
|
||||
|
||||
## Quantization Modes
|
||||
|
||||
| Mode | Guide | Description | Methods |
|
||||
|------|-------|-------------|---------|
|
||||
| Online quantization | [Online Quantization](online.md) | vLLM-Omni computes quantized weights and scales while loading the model. | FP8 W8A8, Int8 W8A8 |
|
||||
| Pre-quantized checkpoints | Method-specific guides | The checkpoint or an offline quantizer provides quantized weights and scales before serving. | GGUF, AutoRound, msModelSlim, serialized Int8 |
|
||||
|
||||
## Hardware Support
|
||||
|
||||
| Device | FP8 W8A8 | Int8 W8A8 | GGUF | AutoRound | msModelSlim |
|
||||
|--------|----------|-----------|------|-----------|-------------|
|
||||
| NVIDIA Blackwell GPU (SM 100+) | ✅ | ✅ | ✅ | ✅ | ❌ |
|
||||
| NVIDIA Ada/Hopper GPU (SM 89+) | ✅ | ✅ | ✅ | ✅ | ❌ |
|
||||
| NVIDIA Ampere GPU (SM 80+) | ✅ | ✅ | ✅ | ✅ | ❌ |
|
||||
| AMD ROCm | ⭕ | ⭕ | ⭕ | ⭕ | ❌ |
|
||||
| Intel XPU | ⭕ | ⭕ | ⭕ | ✅ | ❌ |
|
||||
| Ascend NPU | ❌ | ✅ | ❌ | ❌ | ✅ |
|
||||
|
||||
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
|
||||
guide. FP8 on Ampere may use a weight-only path where available.
|
||||
|
||||
## Model Type Support
|
||||
|
||||
### Diffusion Model (Qwen-Image, Wan2.2)
|
||||
|
||||
These models run a diffusion transformer as the primary inference module. The
|
||||
default quantization target is the transformer; tokenizer, scheduler, text
|
||||
encoder, and VAE stay on the base checkpoint unless a method guide says
|
||||
otherwise.
|
||||
|
||||
| Method | Guide | Mode | Example models | Status |
|
||||
|--------|-------|------|----------------|--------|
|
||||
| FP8 W8A8 | [FP8](fp8.md) | Online W8A8 or checkpoint FP8 | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image family and other DiT models |
|
||||
| Int8 W8A8 | [Int8](int8.md) | Online or serialized W8A8 | Qwen-Image; Wan2.2 is not validated | Validated for Qwen-Image and Z-Image |
|
||||
| GGUF | [GGUF](gguf.md) | Pre-quantized transformer weights | Qwen-Image | Validated where a model-specific GGUF adapter exists |
|
||||
| AutoRound | [AutoRound](autoround.md) | Pre-quantized W4A16 checkpoints | FLUX.1-dev; Qwen-Image/Wan2.2 not validated | Checkpoint-driven |
|
||||
| msModelSlim | [msModelSlim](msmodelslim.md) | Pre-quantized Ascend checkpoints | Wan2.2 recipe; HunyuanImage-3.0 inference target | Ascend/NPU path |
|
||||
|
||||
### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
|
||||
|
||||
These models combine an AR language model with audio, vision, talker, or TTS
|
||||
stages. Quantization is scoped to the AR language-model stage when the
|
||||
checkpoint contains a supported `quantization_config`; the non-AR stages stay
|
||||
in BF16 unless the model guide explicitly adds support.
|
||||
|
||||
| Method | Guide | Scope | Example models | Status |
|
||||
|--------|-------|-------|----------------|--------|
|
||||
| FP8 | [FP8](fp8.md) | Thinker or language-model checkpoint config | Qwen3-Omni thinker | ModelOpt checkpoint path |
|
||||
| Int8 | [Int8](int8.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
|
||||
| GGUF | [GGUF](gguf.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
|
||||
| AutoRound | [AutoRound](autoround.md) | Thinker or language-model checkpoint config | Qwen2.5-Omni, Qwen3-Omni | Supported through AutoRound checkpoints |
|
||||
| msModelSlim | [msModelSlim](msmodelslim.md) | Not currently validated for omni/TTS stages | Qwen3-Omni, Qwen3-TTS | Not validated |
|
||||
|
||||
### Multi-Stage Diffusion Model (BAGEL, GLM-Image)
|
||||
|
||||
These models split generation across multiple stages. Quantization must be
|
||||
attached to the intended stage rather than applied globally.
|
||||
|
||||
| Method | Guide | Scope | Example models | Status |
|
||||
|--------|-------|-------|----------------|--------|
|
||||
| FP8 | [FP8](fp8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Requires model-specific validation |
|
||||
| Int8 | [Int8](int8.md) | Stage-specific DiT or transformer module | BAGEL, GLM-Image | Requires model-specific validation |
|
||||
| GGUF | [GGUF](gguf.md) | Stage-specific transformer weights | BAGEL, GLM-Image | No validated adapter listed |
|
||||
| AutoRound | [AutoRound](autoround.md) | Checkpoint-defined stage | BAGEL, GLM-Image | No validated checkpoint listed |
|
||||
| msModelSlim | [msModelSlim](msmodelslim.md) | Ascend-generated stage weights | GLM-Image | Requires model-specific adaptation |
|
||||
|
||||
!!! note
|
||||
"Online quantization" means vLLM-Omni computes the quantization data while
|
||||
loading the model. "Pre-quantized" means the checkpoint or external
|
||||
quantizer provides the required quantized weights and scales.
|
||||
|
||||
## Quantization Scope
|
||||
|
||||
### Diffusion Model (Qwen-Image, Wan2.2)
|
||||
|
||||
The default target is the diffusion transformer. Component routing is available
|
||||
through `build_quant_config()`:
|
||||
|
||||
```python
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
|
||||
config = build_quant_config({
|
||||
"transformer": {"method": "fp8"},
|
||||
"vae": None,
|
||||
})
|
||||
```
|
||||
|
||||
| Component | Default quantized? | Notes |
|
||||
|-----------|--------------------|-------|
|
||||
| Diffusion transformer | Yes | Primary target for FP8, Int8, GGUF, AutoRound, and msModelSlim |
|
||||
| Text encoder | No | Keep BF16 unless a method-specific guide documents support |
|
||||
| VAE | No | Keep BF16; storage-only paths are method-specific |
|
||||
| Scheduler/tokenizer | No | Loaded from the base model repository |
|
||||
|
||||
### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
|
||||
|
||||
| Component | Default quantized? | Notes |
|
||||
|-----------|--------------------|-------|
|
||||
| Thinker or AR language model | Yes, when checkpoint config is supported | ModelOpt FP8/NVFP4 or AutoRound checkpoint config |
|
||||
| Audio encoder | No | BF16 |
|
||||
| Vision encoder | No | BF16 |
|
||||
| Talker or TTS stage | No | BF16 unless model-specific support is documented |
|
||||
| Code2Wav | No | BF16 |
|
||||
|
||||
### Multi-Stage Diffusion Model (BAGEL, GLM-Image)
|
||||
|
||||
| Component | Default quantized? | Notes |
|
||||
|-----------|--------------------|-------|
|
||||
| Selected diffusion or transformer stage | Method-specific | Must be routed to the intended stage |
|
||||
| Other generation stages | No | Keep BF16 unless separately validated |
|
||||
| VAE, tokenizer, scheduler | No | Loaded from the base checkpoint |
|
||||
|
||||
## Python API
|
||||
|
||||
`build_quant_config()` accepts strings, dictionaries, per-component
|
||||
dictionaries, existing `QuantizationConfig` objects, or `None`.
|
||||
|
||||
```python
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
|
||||
build_quant_config("fp8")
|
||||
build_quant_config({"method": "fp8", "activation_scheme": "static"})
|
||||
build_quant_config("auto-round", bits=4, group_size=128)
|
||||
build_quant_config({"method": "gguf", "gguf_model": "/path/to/model.gguf"})
|
||||
build_quant_config({"transformer": {"method": "fp8"}, "vae": None})
|
||||
build_quant_config(None)
|
||||
```
|
||||
@@ -0,0 +1,40 @@
|
||||
# AudioX offline inference
|
||||
|
||||
Generate audio with the [AudioX](https://zeyuet.github.io/AudioX/) MMDiT diffusion
|
||||
pipeline (`AudioXPipeline`). Six tasks: `t2a`, `t2m`, `v2a`, `v2m`, `tv2a`, `tv2m`.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
Download a vLLM-Omni weight bundle (component-sharded safetensors):
|
||||
|
||||
```bash
|
||||
huggingface-cli download zhangj1an/AudioX --local-dir ./audiox_weights
|
||||
```
|
||||
|
||||
The Hugging Face id `zhangj1an/AudioX` also works directly without prefetching.
|
||||
|
||||
## Usage
|
||||
|
||||
```bash
|
||||
# Text-to-audio only (default uses zhangj1an/AudioX from the Hub):
|
||||
python end2end.py --tasks t2a
|
||||
|
||||
# All six tasks against a local bundle and a sample video for v2*/tv2*:
|
||||
python end2end.py \
|
||||
--model ./audiox_weights \
|
||||
--video https://zeyuet.github.io/AudioX/static/samples/V2M/1XeBotOFqHA.mp4
|
||||
|
||||
# Subset of tasks, custom seed and steps:
|
||||
python end2end.py --tasks t2a tv2a --num-inference-steps 100 --seed 0
|
||||
```
|
||||
|
||||
## Arguments
|
||||
|
||||
- `--model`: HF id or local bundle path (default: `zhangj1an/AudioX`).
|
||||
- `--tasks`: any subset of `t2a t2m v2a v2m tv2a tv2m` (default: all).
|
||||
- `--video`: video file/URL — required for `v2*` and `tv2*`.
|
||||
- `--reference-audio`: optional audio prompt (audio-conditioned generation).
|
||||
- `--num-inference-steps`, `--guidance-scale`, `--seed`, `--seconds-total`,
|
||||
`--sample-rate`, `--output-dir`: generation knobs.
|
||||
|
||||
Outputs land in `<output-dir>/<task>.wav` as 16-bit stereo WAV.
|
||||
@@ -0,0 +1,110 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""End-to-end AudioX offline example covering the 6 t2*/v2*/tv2* tasks.
|
||||
|
||||
Provide a directory with the **vLLM-Omni AudioX safetensors bundle** (e.g. from
|
||||
``zhangj1an/AudioX`` on Hugging Face)::
|
||||
|
||||
huggingface-cli download zhangj1an/AudioX --local-dir ./audiox_weights
|
||||
python end2end.py --model ./audiox_weights
|
||||
python end2end.py --model ./audiox_weights --tasks t2a tv2a
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import soundfile
|
||||
import torch
|
||||
import torchaudio.functional as TF
|
||||
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
ROOT = Path(__file__).resolve().parent
|
||||
|
||||
SAMPLE_PROMPTS: dict[str, str] = {
|
||||
"t2a": "Fireworks burst twice, followed by a period of silence before a clock begins ticking.",
|
||||
"t2m": "Uplifting ukulele tune for a travel vlog",
|
||||
"v2a": "",
|
||||
"v2m": "",
|
||||
"tv2a": "drum beating sound and human talking",
|
||||
"tv2m": "uplifting music matching the scene",
|
||||
}
|
||||
|
||||
ALL_TASKS = ("t2a", "t2m", "v2a", "v2m", "tv2a", "tv2m")
|
||||
VIDEO_TASKS = frozenset({"v2a", "v2m", "tv2a", "tv2m"})
|
||||
TEXT_TASKS = frozenset({"t2a", "t2m", "tv2a", "tv2m"})
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
p = argparse.ArgumentParser(description="AudioX offline end-to-end (6 t2*/v2*/tv2* tasks).")
|
||||
p.add_argument("--model", default="zhangj1an/AudioX", help="HF id or local AudioX bundle path.")
|
||||
p.add_argument("--tasks", nargs="+", default=list(ALL_TASKS), choices=ALL_TASKS)
|
||||
p.add_argument("--video", default="", help="Video path / URL (required for v2*/tv2*).")
|
||||
p.add_argument("--reference-audio", default="", help="Optional audio prompt for audio-conditioned generation.")
|
||||
p.add_argument("--output-dir", default=str(ROOT / "audiox_task_outputs"))
|
||||
p.add_argument("--num-inference-steps", type=int, default=250)
|
||||
p.add_argument("--seconds-total", type=float, default=10.0)
|
||||
p.add_argument("--guidance-scale", type=float, default=6.0)
|
||||
p.add_argument("--seed", type=int, default=42)
|
||||
p.add_argument("--sample-rate", type=int, default=48000, help="Output WAV rate (resampled if != model rate).")
|
||||
return p.parse_args()
|
||||
|
||||
|
||||
def save_wav(audio: torch.Tensor, path: Path, sample_rate: int) -> None:
|
||||
"""Write 16-bit PCM WAV. ``audio`` is ``[channels, samples]`` float in [-1, 1]."""
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
soundfile.write(str(path), audio.clamp(-1.0, 1.0).cpu().T.numpy(), sample_rate, subtype="PCM_16")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
|
||||
omni = Omni(model=args.model, model_class_name="AudioXPipeline")
|
||||
|
||||
for task in args.tasks:
|
||||
if task in VIDEO_TASKS and not args.video:
|
||||
raise SystemExit(f"task={task!r} requires --video")
|
||||
prompt = SAMPLE_PROMPTS[task] if task in TEXT_TASKS else ""
|
||||
extra: dict = {"audiox_task": task, "seconds_start": 0.0, "seconds_total": float(args.seconds_total)}
|
||||
if task in VIDEO_TASKS:
|
||||
extra["video_path"] = args.video
|
||||
if args.reference_audio:
|
||||
extra["audio_path"] = args.reference_audio
|
||||
|
||||
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
|
||||
t0 = time.perf_counter()
|
||||
outputs = omni.generate(
|
||||
prompt,
|
||||
OmniDiffusionSamplingParams(
|
||||
generator=generator,
|
||||
guidance_scale=args.guidance_scale,
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
seed=args.seed,
|
||||
extra_args=extra,
|
||||
),
|
||||
)
|
||||
audio = outputs[0].request_output.multimodal_output.get("audio")
|
||||
if audio is None:
|
||||
raise RuntimeError(f"No audio produced for task {task!r}")
|
||||
audio = torch.as_tensor(audio).detach().cpu().float()
|
||||
if audio.ndim == 3:
|
||||
audio = audio[0]
|
||||
|
||||
model_sr = int(outputs[0].request_output.multimodal_output.get("audio_sample_rate") or 44100)
|
||||
if model_sr != args.sample_rate:
|
||||
audio = TF.resample(audio, model_sr, args.sample_rate)
|
||||
|
||||
out_path = Path(args.output_dir) / f"{task}.wav"
|
||||
save_wav(audio, out_path, args.sample_rate)
|
||||
print(f"[{task}] saved {out_path} ({time.perf_counter() - t0:.2f}s)")
|
||||
|
||||
omni.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -153,7 +153,7 @@ def main():
|
||||
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
omni_kwargs = {}
|
||||
omni_kwargs = vars(args).copy()
|
||||
deploy_config = args.deploy_config
|
||||
if args.think and deploy_config is None:
|
||||
deploy_config = "vllm_omni/deploy/bagel_think.yaml"
|
||||
@@ -161,22 +161,12 @@ def main():
|
||||
if deploy_config:
|
||||
omni_kwargs["deploy_config"] = deploy_config
|
||||
|
||||
omni_kwargs.update(
|
||||
{
|
||||
"log_stats": args.log_stats,
|
||||
"init_sleep_seconds": args.init_sleep_seconds,
|
||||
"batch_timeout": args.batch_timeout,
|
||||
"init_timeout": args.init_timeout,
|
||||
"shm_threshold_bytes": args.shm_threshold_bytes,
|
||||
"worker_backend": args.worker_backend,
|
||||
"ray_address": args.ray_address,
|
||||
"enable_diffusion_pipeline_profiler": args.enable_diffusion_pipeline_profiler,
|
||||
}
|
||||
)
|
||||
if args.quantization:
|
||||
omni_kwargs["quantization_config"] = args.quantization
|
||||
|
||||
omni = Omni.from_cli_args(args, model=model_name, **omni_kwargs)
|
||||
# Override CLI --model with the derived model_name.
|
||||
omni_kwargs["model"] = model_name
|
||||
omni = Omni(**omni_kwargs)
|
||||
|
||||
formatted_prompts = []
|
||||
for p in prompts:
|
||||
|
||||
@@ -1,71 +0,0 @@
|
||||
# CosyVoice3
|
||||
|
||||
## Setup
|
||||
|
||||
Install dependencies:
|
||||
```
|
||||
uv pip install -e .
|
||||
```
|
||||
|
||||
> **Note:** This includes required libraries such as `soundfile`,
|
||||
> `onnxruntime`, `x-transformers`, and `einops` via
|
||||
> `requirements/common.txt` and platform-specific requirements files.
|
||||
|
||||
Download the model snapshot:
|
||||
```
|
||||
from huggingface_hub import snapshot_download
|
||||
snapshot_download('FunAudioLLM/Fun-CosyVoice3-0.5B-2512', local_dir='pretrained_models/Fun-CosyVoice3-0.5B')
|
||||
```
|
||||
|
||||
Add `config.json` in `pretrained_models/Fun-CosyVoice3-0.5B/`:
|
||||
```json
|
||||
{
|
||||
"model_type": "cosyvoice3",
|
||||
"architectures": [
|
||||
"CosyVoice3Model"
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
> **Why `config.json` is required:**
|
||||
> `AutoConfig.register("cosyvoice3", CosyVoice3Config)` only registers a class mapping.
|
||||
> The loader still needs `model_type: "cosyvoice3"` from `config.json` to select that class.
|
||||
> If no `config.json` is present, model type cannot be inferred automatically.
|
||||
> If your downloaded checkpoint already includes a valid `config.json` with
|
||||
> `model_type: "cosyvoice3"`, this manual step can be skipped.
|
||||
|
||||
Run the offline verification script:
|
||||
```
|
||||
python examples/offline_inference/cosyvoice3/verify_e2e_cosyvoice.py \
|
||||
--model pretrained_models/Fun-CosyVoice3-0.5B \
|
||||
--tokenizer pretrained_models/Fun-CosyVoice3-0.5B/CosyVoice-BlankEN
|
||||
```
|
||||
|
||||
## Implementation Overview
|
||||
|
||||
CosyVoice3 runs as a **2-stage Omni pipeline**.
|
||||
|
||||
- **Stage 0** (`talker`) converts text + prompt audio to speech tokens.
|
||||
- **Stage 1** (`code2wav`) converts speech tokens + prompt features to waveform via flow matching and HiFiGAN.
|
||||
|
||||
Key components live in `vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py`.
|
||||
|
||||
- `CosyVoice3MultiModalProcessor` builds the multimodal inputs:
|
||||
- Tokenizes `prompt` and `prompt_text`.
|
||||
- Extracts speech tokens and mel features from the prompt audio.
|
||||
- Extracts a speaker embedding.
|
||||
- `CosyVoice3Model` implements both stages:
|
||||
- Stage 0 uses `CosyVoice3LM` and outputs speech tokens + conditioning features.
|
||||
- Stage 1 runs the flow model (DiT-based CFM) and HiFiGAN to synthesize waveform.
|
||||
|
||||
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 `sync_process_input_func` (sync mode) or the
|
||||
shared-memory connector (async-chunk mode) and outputs audio.
|
||||
@@ -118,9 +118,6 @@ def parse_args() -> argparse.Namespace:
|
||||
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()
|
||||
|
||||
|
||||
|
||||
@@ -18,6 +18,8 @@ import numpy as np
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||
|
||||
TASK_CHOICES = ("t2t", "t2i", "t2s", "i2i", "i2t", "s2t", "v2t")
|
||||
|
||||
TASK_DEFAULT_RUNTIME = {
|
||||
@@ -970,7 +972,6 @@ 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,75 +0,0 @@
|
||||
# Fish Speech S2 Pro
|
||||
|
||||
This directory contains an offline demo for running Fish Speech S2 Pro (`fishaudio/s2-pro`) with vLLM Omni. It supports text-only synthesis and voice cloning with reference audio.
|
||||
|
||||
## Model Overview
|
||||
|
||||
[Fish Speech S2 Pro](https://huggingface.co/fishaudio/s2-pro) is a 4B dual-AR text-to-speech model by FishAudio. It uses:
|
||||
|
||||
- **Slow AR**: Generates semantic tokens autoregressively (Qwen3-based backbone)
|
||||
- **Fast AR**: Predicts residual codebook tokens from semantic tokens
|
||||
- **DAC Decoder**: Converts codec codes to 44.1 kHz audio waveform
|
||||
|
||||
The model produces high-quality speech with voice cloning capabilities.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
Install the `fish-speech` package for the DAC codec:
|
||||
|
||||
```bash
|
||||
pip install fish-speech
|
||||
```
|
||||
|
||||
## 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.
|
||||
|
||||
## Quick Start
|
||||
|
||||
### Text-Only Synthesis
|
||||
|
||||
```bash
|
||||
python end2end.py --text "Hello, this is a test of the Fish Speech text to speech system."
|
||||
```
|
||||
|
||||
Generated audio files are saved to `output_audio/` by default.
|
||||
|
||||
### Voice Cloning
|
||||
|
||||
Provide a reference audio file and its transcript to clone a voice:
|
||||
|
||||
```bash
|
||||
python end2end.py \
|
||||
--text "Hello, this is a cloned voice." \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--ref-text "Transcript of the reference audio."
|
||||
```
|
||||
|
||||
## Streaming Mode
|
||||
|
||||
Add `--streaming` to stream audio chunks progressively via `AsyncOmni` (requires `async_chunk: true` in the stage config):
|
||||
|
||||
```bash
|
||||
python end2end.py --text "Hello, this is a test." --streaming
|
||||
```
|
||||
|
||||
Each DAC decoder chunk is logged as it arrives. The final WAV file is written once generation completes.
|
||||
|
||||
## Configuration
|
||||
|
||||
| Argument | Default | Description |
|
||||
|----------|---------|-------------|
|
||||
| `--text` | `"Hello, this is a test..."` | Text to synthesize |
|
||||
| `--ref-audio` | None | Path to reference audio for voice cloning |
|
||||
| `--ref-text` | None | Transcript of the reference audio |
|
||||
| `--model` | `fishaudio/s2-pro` | HuggingFace model path |
|
||||
| `--stage-configs-path` | `fish_speech_s2_pro.yaml` | Path to stage configs YAML |
|
||||
| `--output-dir` | `output_audio` | Output directory for WAV files |
|
||||
| `--streaming` | False | Enable streaming via AsyncOmni |
|
||||
| `--stage-init-timeout` | 600 | Stage initialization timeout (seconds) |
|
||||
|
||||
## Notes
|
||||
|
||||
- Output audio is 44.1 kHz mono WAV.
|
||||
- The DAC codec weights (`codec.pth`) are loaded lazily from the model directory.
|
||||
- Voice cloning encodes the reference audio using the DAC codec to extract semantic codes, then prepends them as a system message.
|
||||
@@ -196,9 +196,6 @@ 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()
|
||||
|
||||
|
||||
|
||||
@@ -31,6 +31,14 @@ python end2end.py --model tencent/HunyuanImage-3.0-Instruct \
|
||||
--prompts "A cute cat sitting on a windowsill watching the sunset"
|
||||
```
|
||||
|
||||
**With VAE tiling (required on A100 GPUs):**
|
||||
```bash
|
||||
python end2end.py --model tencent/HunyuanImage-3.0-Instruct \
|
||||
--modality text2img \
|
||||
--prompts "A cute cat sitting on a windowsill watching the sunset" \
|
||||
--vae-use-tiling
|
||||
```
|
||||
|
||||
#### Image to Image (img2img)
|
||||
|
||||
- **Pipeline**: Image + Text → AR (CoT + recaption + latent) → DiT → Edited Image
|
||||
@@ -103,6 +111,7 @@ python end2end.py --modality text2img \
|
||||
| `--stage-configs-path` | string | auto | Custom stage config YAML path |
|
||||
| `--enforce-eager` | flag | `False` | Disable torch.compile |
|
||||
| `--init-timeout` | int | `300` | Initialization timeout (seconds) |
|
||||
| `--vae-use-tiling` | flag | `False` | Enable VAE tiling for memory optimization (required to avoid OOM on A100) |
|
||||
|
||||
------
|
||||
|
||||
@@ -135,23 +144,25 @@ python end2end.py --model tencent/HunyuanImage-3.0-Instruct \
|
||||
|
||||
## Prompt Format
|
||||
|
||||
HunyuanImage-3.0 uses a pretrain template format:
|
||||
HunyuanImage-3.0-Instruct uses an instruct chat template:
|
||||
|
||||
```
|
||||
<|startoftext|>{system_prompt}{<img>}{trigger_tag}{user_prompt}
|
||||
<|startoftext|>{system_prompt}\n\nUser: {<img>?}{user_prompt}\n\nAssistant: {trigger_tag?}
|
||||
```
|
||||
|
||||
- `<img>`: Placeholder for each input image (auto-inserted by `prompt_utils.py`)
|
||||
- Trigger tags: `<think>` (CoT), `<recaption>` (recaptioning)
|
||||
- `<img>`: Placeholder for each input image (single token; expanded by the multimodal pipeline)
|
||||
- Trigger tags: `<think>` (CoT), `<recaption>` (recaptioning) — placed AFTER `Assistant: `
|
||||
- System prompt: Auto-selected based on task
|
||||
- `t2i_vanilla` is the only task that uses the bare pretrain template (no chat structure)
|
||||
|
||||
The `prompt_utils.build_prompt()` handles this formatting automatically.
|
||||
The shared `vllm_omni.diffusion.models.hunyuan_image3.prompt_utils.build_prompt_tokens()`
|
||||
helper handles segment-by-segment tokenization (matches HF `apply_chat_template` byte-for-byte).
|
||||
|
||||
------
|
||||
|
||||
## FAQ
|
||||
|
||||
- **OOM errors**: Decrease `gpu_memory_utilization` in the YAML stage config, or use a smaller `max_num_batched_tokens`.
|
||||
- **OOM errors**: Decrease `gpu_memory_utilization` in the YAML stage config, use a smaller `max_num_batched_tokens`, or enable VAE tiling with `--vae-use-tiling` (required on A100 GPUs).
|
||||
- **Custom image sizes**: Use `--height` and `--width` flags (multiples of 16 recommended).
|
||||
|
||||
| Stage | VRAM (approx) |
|
||||
|
||||
@@ -16,13 +16,14 @@ Usage:
|
||||
import argparse
|
||||
import os
|
||||
|
||||
from vllm_omni.diffusion.models.hunyuan_image3.system_prompt import (
|
||||
get_system_prompt,
|
||||
from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import (
|
||||
build_prompt_tokens,
|
||||
)
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniPromptType
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
# task → (sys_type, bot_task, trigger_tag)
|
||||
# task -> (sys_type, bot_task, trigger_tag)
|
||||
_TASK_PRESETS: dict[str, tuple[str, str | None, str | None]] = {
|
||||
"t2t": ("en_unified", None, None),
|
||||
"i2t": ("en_unified", None, None),
|
||||
@@ -42,36 +43,6 @@ _MODALITY_TASK_MAP = {
|
||||
}
|
||||
|
||||
|
||||
def build_prompt(
|
||||
user_prompt: str,
|
||||
task: str = "it2i_think",
|
||||
sys_type: str | None = None,
|
||||
custom_system_prompt: str | None = None,
|
||||
) -> str:
|
||||
"""Build a HunyuanImage-3.0 prompt using pretrain template format."""
|
||||
if task not in _TASK_PRESETS:
|
||||
raise ValueError(f"Unknown task {task!r}. Choose from: {sorted(_TASK_PRESETS)}")
|
||||
|
||||
preset_sys_type, preset_bot_task, trigger_tag = _TASK_PRESETS[task]
|
||||
effective_sys_type = sys_type or preset_sys_type
|
||||
|
||||
system_prompt = get_system_prompt(effective_sys_type, preset_bot_task, custom_system_prompt)
|
||||
sys_text = system_prompt.strip() if system_prompt else ""
|
||||
|
||||
has_image_input = task.startswith("i2t") or task.startswith("it2i")
|
||||
|
||||
parts = ["<|startoftext|>"]
|
||||
if sys_text:
|
||||
parts.append(sys_text)
|
||||
if has_image_input:
|
||||
parts.append("<img>")
|
||||
if trigger_tag:
|
||||
parts.append(trigger_tag)
|
||||
parts.append(user_prompt)
|
||||
|
||||
return "".join(parts)
|
||||
|
||||
|
||||
# Modality → default stage config
|
||||
_MODALITY_DEFAULT_CONFIG = {
|
||||
"text2img": "hunyuan_image3_t2i.yaml",
|
||||
@@ -114,6 +85,11 @@ def parse_args():
|
||||
parser.add_argument("--seed", type=int, default=42, help="Random seed.")
|
||||
parser.add_argument("--height", type=int, default=1024, help="Output image height.")
|
||||
parser.add_argument("--width", type=int, default=1024, help="Output image width.")
|
||||
parser.add_argument(
|
||||
"--vae-use-tiling",
|
||||
action="store_true",
|
||||
help="Enable VAE tiling for memory optimization.",
|
||||
)
|
||||
|
||||
# Prompt configuration
|
||||
parser.add_argument(
|
||||
@@ -134,6 +110,32 @@ def parse_args():
|
||||
parser.add_argument("--log-stats", action="store_true", default=False)
|
||||
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.")
|
||||
parser.add_argument(
|
||||
"--quantization",
|
||||
type=str,
|
||||
default=None,
|
||||
choices=["fp8", "int8", "gguf"],
|
||||
help="Quantization method for the transformer. "
|
||||
"Options: 'fp8' (FP8 W8A8 on Ada/Hopper, weight-only on older GPUs), 'int8' (Int8 W8A8), 'gguf' (GGUF quantized weights). "
|
||||
"Default: None (no quantization, uses BF16).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--gguf-model",
|
||||
type=str,
|
||||
default=None,
|
||||
help=("GGUF file path or HF reference for transformer weights. Required when --quantization gguf is set."),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ignored-layers",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Comma-separated list of layer name patterns to skip quantization. "
|
||||
"Only used when --quantization is set. "
|
||||
"Available layers: to_qkv, to_out, add_kv_proj, to_add_out, img_mlp, txt_mlp, proj_out. "
|
||||
"Example: --ignored-layers 'add_kv_proj,to_add_out'",
|
||||
)
|
||||
|
||||
current_omni_platform.pre_register_and_update(parser)
|
||||
|
||||
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||
|
||||
@@ -150,14 +152,34 @@ def main():
|
||||
|
||||
# Determine stage config
|
||||
stage_configs_path = args.stage_configs_path or _MODALITY_DEFAULT_CONFIG[args.modality]
|
||||
# Build quantization kwargs: use quantization_config dict when
|
||||
# ignored_layers is specified so the list flows through OmniDiffusionConfig
|
||||
quant_kwargs: dict[str, Any] = {}
|
||||
ignored_layers = [s.strip() for s in args.ignored_layers.split(",") if s.strip()] if args.ignored_layers else None
|
||||
if args.quantization == "gguf":
|
||||
if not args.gguf_model:
|
||||
raise ValueError("--gguf-model is required when --quantization gguf is set.")
|
||||
quant_kwargs["quantization_config"] = {
|
||||
"method": "gguf",
|
||||
"gguf_model": args.gguf_model,
|
||||
}
|
||||
elif args.quantization and ignored_layers:
|
||||
quant_kwargs["quantization_config"] = {
|
||||
"method": args.quantization,
|
||||
"ignored_layers": ignored_layers,
|
||||
}
|
||||
elif args.quantization:
|
||||
quant_kwargs["quantization"] = args.quantization
|
||||
|
||||
# Build Omni
|
||||
omni_kwargs = {
|
||||
"model": args.model,
|
||||
"vae_use_tiling": args.vae_use_tiling,
|
||||
"stage_configs_path": stage_configs_path,
|
||||
"log_stats": args.log_stats,
|
||||
"init_timeout": args.init_timeout,
|
||||
"enforce_eager": args.enforce_eager,
|
||||
**quant_kwargs,
|
||||
}
|
||||
if args.modality in ("text2img", "img2img"):
|
||||
omni_kwargs["mode"] = "text-to-image"
|
||||
@@ -179,12 +201,28 @@ def main():
|
||||
|
||||
input_image = Image.open(args.image_path).convert("RGB")
|
||||
|
||||
# Load tokenizer for segment-wise prompt tokenization (matches HF
|
||||
# apply_chat_template byte-for-byte; see build_prompt_tokens docstring).
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
|
||||
|
||||
# Format prompts
|
||||
formatted_prompts: list[OmniPromptType] = []
|
||||
for p in prompts:
|
||||
formatted_text = build_prompt(p, task=task, sys_type=args.sys_type)
|
||||
token_ids = build_prompt_tokens(p, tokenizer, task=task, sys_type=args.sys_type)
|
||||
preset_sys_type, _, _ = _TASK_PRESETS[task]
|
||||
effective_sys_type = args.sys_type or preset_sys_type
|
||||
|
||||
prompt_dict: dict = {"prompt": formatted_text}
|
||||
# `prompt_token_ids` drives the AR stage (matches HF byte-for-byte).
|
||||
# `prompt` and `use_system_prompt` are forwarded by ar2diffusion to
|
||||
# the DiT stage so the diffusion pipeline can rebuild the same
|
||||
# system prefix when constructing its model inputs.
|
||||
prompt_dict: dict = {
|
||||
"prompt_token_ids": token_ids,
|
||||
"prompt": p,
|
||||
"use_system_prompt": effective_sys_type,
|
||||
}
|
||||
|
||||
if args.modality == "text2img":
|
||||
prompt_dict["modalities"] = ["image"]
|
||||
|
||||
@@ -378,9 +378,6 @@ def parse_args() -> argparse.Namespace:
|
||||
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()
|
||||
|
||||
|
||||
|
||||
@@ -211,9 +211,6 @@ def parse_args() -> argparse.Namespace:
|
||||
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()
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,30 @@
|
||||
# InternVLA-A1
|
||||
|
||||
Full usage and result-reporting guidance lives in [docs/user_guide/examples/offline_inference/internvla_a1.md](../../../docs/user_guide/examples/offline_inference/internvla_a1.md).
|
||||
|
||||
Quick start:
|
||||
|
||||
```bash
|
||||
export INTERNVLA_A1_MODEL_DIR=/path/to/InternVLA-A1-3B-ft-pen
|
||||
export INTERNVLA_A1_DATASET_DIR=/path/to/Genie1-Place_Markpen
|
||||
export INTERNVLA_A1_PROCESSOR_DIR=/path/to/Qwen3-VL-2B-Instruct
|
||||
# hf tenstep/Cosmos-Tokenizer-CI8x8-SafeTensors
|
||||
export INTERNVLA_A1_COSMOS_DIR=/path/to/Cosmos-Tokenizer-CI8x8-SafeTensor
|
||||
|
||||
bash run.sh --num-samples 1 --num-episodes 0
|
||||
bash run.sh --num-episodes 1
|
||||
bash collect_results.sh
|
||||
```
|
||||
|
||||
Expected files under `INTERNVLA_A1_COSMOS_DIR`:
|
||||
|
||||
- `encoder.safetensors`
|
||||
- `decoder.safetensors`
|
||||
|
||||
Reference Hugging Face repo: `tenstep/Cosmos-Tokenizer-CI8x8-SafeTensors`
|
||||
|
||||
Key entrypoints:
|
||||
|
||||
- `run.sh`: wrapper for offline inference and GT evaluation
|
||||
- `collect_results.sh`: collect sample output, latency, metrics, plots, and logs
|
||||
- `end2end.py`: underlying Python entrypoint
|
||||
@@ -0,0 +1,220 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
REPO_ROOT="$(cd "$ROOT_DIR/../../.." && pwd)"
|
||||
TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
|
||||
RESULT_ROOT="${INTERNVLA_A1_RESULT_DIR:-$REPO_ROOT/outputs/internvla_a1/collected_results/$TIMESTAMP}"
|
||||
|
||||
: "${INTERNVLA_A1_MODEL_DIR:?Please export INTERNVLA_A1_MODEL_DIR=/path/to/InternVLA-A1-3B-ft-pen}"
|
||||
: "${INTERNVLA_A1_DATASET_DIR:?Please export INTERNVLA_A1_DATASET_DIR=/path/to/Genie1-Place_Markpen}"
|
||||
: "${INTERNVLA_A1_PROCESSOR_DIR:?Please export INTERNVLA_A1_PROCESSOR_DIR=/path/to/Qwen3-VL-2B-Instruct}"
|
||||
: "${INTERNVLA_A1_COSMOS_DIR:?Please export INTERNVLA_A1_COSMOS_DIR=/path/to/Cosmos-Tokenizer-CI8x8-SafeTensor}"
|
||||
|
||||
EVAL_OUTPUT_DIR="$RESULT_ROOT/eval_outputs"
|
||||
GPU_MONITOR_INTERVAL="${GPU_MONITOR_INTERVAL:-1}"
|
||||
|
||||
mkdir -p "$RESULT_ROOT"
|
||||
|
||||
export INTERNVLA_A1_MODEL_DIR
|
||||
export INTERNVLA_A1_DATASET_DIR
|
||||
export INTERNVLA_A1_PROCESSOR_DIR
|
||||
export INTERNVLA_A1_COSMOS_DIR
|
||||
|
||||
write_env_summary() {
|
||||
cat >"$RESULT_ROOT/env_summary.txt" <<EOF
|
||||
timestamp=$TIMESTAMP
|
||||
repo_root=$REPO_ROOT
|
||||
result_root=$RESULT_ROOT
|
||||
python=$(command -v python || true)
|
||||
python_version=$(python --version 2>&1 || true)
|
||||
model_dir=$INTERNVLA_A1_MODEL_DIR
|
||||
dataset_dir=$INTERNVLA_A1_DATASET_DIR
|
||||
processor_dir=$INTERNVLA_A1_PROCESSOR_DIR
|
||||
cosmos_dir=$INTERNVLA_A1_COSMOS_DIR
|
||||
pwd=$(pwd)
|
||||
EOF
|
||||
}
|
||||
|
||||
capture_gpu_snapshot() {
|
||||
local output_file="$1"
|
||||
if command -v nvidia-smi >/dev/null 2>&1; then
|
||||
nvidia-smi --query-gpu=index,name,memory.total,memory.used,utilization.gpu --format=csv >"$output_file" 2>&1 || true
|
||||
else
|
||||
echo "nvidia-smi not found" >"$output_file"
|
||||
fi
|
||||
}
|
||||
|
||||
start_gpu_monitor() {
|
||||
local output_file="$1"
|
||||
if ! command -v nvidia-smi >/dev/null 2>&1; then
|
||||
echo "nvidia-smi not found" >"$output_file"
|
||||
return 1
|
||||
fi
|
||||
|
||||
(
|
||||
echo "timestamp,index,name,memory.total [MiB],memory.used [MiB],utilization.gpu [%]"
|
||||
while true; do
|
||||
local now
|
||||
now="$(date '+%Y-%m-%d %H:%M:%S')"
|
||||
nvidia-smi --query-gpu=index,name,memory.total,memory.used,utilization.gpu --format=csv,noheader,nounits \
|
||||
| awk -F',' -v ts="$now" '{gsub(/^[ \t]+|[ \t]+$/, "", $0); print ts "," $0}'
|
||||
sleep "$GPU_MONITOR_INTERVAL"
|
||||
done
|
||||
) >"$output_file" 2>/dev/null &
|
||||
echo $!
|
||||
}
|
||||
|
||||
stop_gpu_monitor() {
|
||||
local monitor_pid="${1:-}"
|
||||
if [[ -n "$monitor_pid" ]] && kill -0 "$monitor_pid" 2>/dev/null; then
|
||||
kill "$monitor_pid" 2>/dev/null || true
|
||||
wait "$monitor_pid" 2>/dev/null || true
|
||||
fi
|
||||
}
|
||||
|
||||
copy_eval_outputs() {
|
||||
local source_dir="$1"
|
||||
local target_dir="$2"
|
||||
local source_real=""
|
||||
local target_real=""
|
||||
source_real="$(realpath "$source_dir" 2>/dev/null || echo "$source_dir")"
|
||||
target_real="$(realpath "$target_dir" 2>/dev/null || echo "$target_dir")"
|
||||
if [[ "$source_real" == "$target_real" ]]; then
|
||||
return 0
|
||||
fi
|
||||
mkdir -p "$target_dir"
|
||||
if [[ -d "$source_dir" ]]; then
|
||||
cp -r "$source_dir"/. "$target_dir"/
|
||||
fi
|
||||
}
|
||||
|
||||
run_with_artifacts() {
|
||||
local name="$1"
|
||||
shift
|
||||
|
||||
local log_file="$RESULT_ROOT/${name}.log"
|
||||
local time_file="$RESULT_ROOT/${name}_time.txt"
|
||||
local gpu_file="$RESULT_ROOT/${name}_gpu.csv"
|
||||
local status_file="$RESULT_ROOT/${name}_status.txt"
|
||||
local monitor_pid=""
|
||||
|
||||
monitor_pid="$(start_gpu_monitor "$gpu_file" || true)"
|
||||
set +e
|
||||
local exit_code=0
|
||||
local start_ts=0
|
||||
local end_ts=0
|
||||
start_ts="$(date +%s)"
|
||||
if [[ -x /usr/bin/time ]]; then
|
||||
/usr/bin/time -v -o "$time_file" "$@" >"$log_file" 2>&1
|
||||
exit_code=$?
|
||||
if [[ ! -s "$time_file" ]]; then
|
||||
end_ts="$(date +%s)"
|
||||
{
|
||||
echo "timing_mode=usr_bin_time_empty_fallback"
|
||||
echo "start_ts=$start_ts"
|
||||
echo "end_ts=$end_ts"
|
||||
echo "elapsed_seconds=$((end_ts - start_ts))"
|
||||
} >"$time_file"
|
||||
fi
|
||||
else
|
||||
"$@" >"$log_file" 2>&1
|
||||
exit_code=$?
|
||||
end_ts="$(date +%s)"
|
||||
{
|
||||
echo "timing_mode=shell_date"
|
||||
echo "start_ts=$start_ts"
|
||||
echo "end_ts=$end_ts"
|
||||
echo "elapsed_seconds=$((end_ts - start_ts))"
|
||||
} >"$time_file"
|
||||
fi
|
||||
set -e
|
||||
stop_gpu_monitor "$monitor_pid"
|
||||
|
||||
echo "exit_code=$exit_code" >"$status_file"
|
||||
if [[ $exit_code -ne 0 ]]; then
|
||||
echo "[error] command failed for $name, see $log_file"
|
||||
return $exit_code
|
||||
fi
|
||||
}
|
||||
|
||||
write_manifest() {
|
||||
cat >"$RESULT_ROOT/README.txt" <<EOF
|
||||
InternVLA-A1 collected results
|
||||
|
||||
Key files:
|
||||
- env_summary.txt: environment and path summary
|
||||
- sample_run.log / sample_run_time.txt: one-sample functional run
|
||||
- forward_benchmark.log / forward_benchmark_time.txt: pure pipeline.forward latency benchmark
|
||||
- eval_run.log / eval_run_time.txt: GT evaluation run
|
||||
- pytest_e2e.log / pytest_e2e_time.txt: offline e2e pytest result
|
||||
- gpu_info_before.csv / gpu_info_after.csv: point-in-time GPU snapshots
|
||||
- *_gpu.csv: sampled GPU usage during each run
|
||||
- eval_outputs/: copied output directory from the GT evaluation run
|
||||
|
||||
Important outputs:
|
||||
- forward_benchmark/forward_latency.json
|
||||
- eval_outputs/summary.json
|
||||
- eval_outputs/registry/log.json
|
||||
- eval_outputs/registry/plots/
|
||||
EOF
|
||||
}
|
||||
|
||||
write_skip_artifact() {
|
||||
local name="$1"
|
||||
local reason="$2"
|
||||
echo "$reason" >"$RESULT_ROOT/${name}.log"
|
||||
echo "timing_mode=skipped" >"$RESULT_ROOT/${name}_time.txt"
|
||||
echo "exit_code=0" >"$RESULT_ROOT/${name}_status.txt"
|
||||
echo "skipped_reason=$reason" >>"$RESULT_ROOT/${name}_status.txt"
|
||||
if [[ ! -f "$RESULT_ROOT/${name}_gpu.csv" ]]; then
|
||||
echo "skipped,$reason" >"$RESULT_ROOT/${name}_gpu.csv"
|
||||
fi
|
||||
}
|
||||
|
||||
write_env_summary
|
||||
write_manifest
|
||||
capture_gpu_snapshot "$RESULT_ROOT/gpu_info_before.csv"
|
||||
|
||||
run_with_artifacts \
|
||||
"sample_run" \
|
||||
bash "$ROOT_DIR/run.sh" \
|
||||
--output-dir "$RESULT_ROOT/sample_outputs" \
|
||||
--num-samples 1 \
|
||||
--num-episodes 0
|
||||
|
||||
run_with_artifacts \
|
||||
"forward_benchmark" \
|
||||
python "$ROOT_DIR/end2end.py" \
|
||||
--model-dir "$INTERNVLA_A1_MODEL_DIR" \
|
||||
--dataset-dir "$INTERNVLA_A1_DATASET_DIR" \
|
||||
--benchmark-forward \
|
||||
--dtype bfloat16 \
|
||||
--attn-implementation eager \
|
||||
--warmup-iters 3 \
|
||||
--benchmark-iters 10 \
|
||||
--output-dir "$RESULT_ROOT/forward_benchmark"
|
||||
|
||||
run_with_artifacts \
|
||||
"eval_run" \
|
||||
bash "$ROOT_DIR/run.sh" \
|
||||
--output-dir "$EVAL_OUTPUT_DIR" \
|
||||
--num-episodes 1
|
||||
|
||||
copy_eval_outputs "$EVAL_OUTPUT_DIR" "$RESULT_ROOT/eval_outputs"
|
||||
|
||||
if python - <<'PY' >/dev/null 2>&1
|
||||
import importlib.util
|
||||
raise SystemExit(0 if importlib.util.find_spec("pytest") else 1)
|
||||
PY
|
||||
then
|
||||
run_with_artifacts \
|
||||
"pytest_e2e" \
|
||||
python -m pytest -sv tests/e2e/offline_inference/test_internvla_a1.py -m advanced_model
|
||||
else
|
||||
write_skip_artifact "pytest_e2e" "pytest is not installed in the current python environment"
|
||||
fi
|
||||
|
||||
capture_gpu_snapshot "$RESULT_ROOT/gpu_info_after.csv"
|
||||
|
||||
echo "Results written to: $RESULT_ROOT"
|
||||
+327
@@ -0,0 +1,327 @@
|
||||
#!/usr/bin/env python3
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
import statistics
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
|
||||
_WORKSPACE_ROOT = Path(__file__).resolve().parent
|
||||
_REPO_ROOT = _WORKSPACE_ROOT.parents[2]
|
||||
if str(_REPO_ROOT) not in sys.path:
|
||||
# Keep the shared helper import clean by doing script-only path bootstrap here.
|
||||
sys.path.insert(0, str(_REPO_ROOT))
|
||||
|
||||
from internvla_a1_common import ( # noqa: E402
|
||||
A2DOpenLoopDataset,
|
||||
collate_open_loop_samples,
|
||||
make_shared_noise,
|
||||
run_open_loop_evaluation,
|
||||
select_indices,
|
||||
tensor_dtype,
|
||||
tensor_sha256,
|
||||
)
|
||||
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig # noqa: E402
|
||||
from vllm_omni.diffusion.models.internvla_a1 import ( # noqa: E402
|
||||
InternVLAA1Config,
|
||||
InternVLAA1TrainMetadata,
|
||||
)
|
||||
from vllm_omni.diffusion.models.internvla_a1.config import OBS_STATE # noqa: E402
|
||||
from vllm_omni.diffusion.registry import initialize_model # noqa: E402
|
||||
from vllm_omni.diffusion.request import OmniDiffusionRequest # noqa: E402
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams # noqa: E402
|
||||
|
||||
|
||||
def _required_path_arg(env_name: str, cli_value: str | None) -> str:
|
||||
value = cli_value or os.getenv(env_name)
|
||||
if not value:
|
||||
raise ValueError(f"Missing required path: set --{env_name.lower().replace('_', '-')} or {env_name}.")
|
||||
return value
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Run vLLM InternVLA-A1 inference on a few samples.")
|
||||
parser.add_argument("--model-dir")
|
||||
parser.add_argument("--dataset-dir")
|
||||
parser.add_argument("--num-samples", type=int, default=1)
|
||||
parser.add_argument("--num-episodes", type=int, default=0)
|
||||
parser.add_argument("--seed", type=int, default=1234)
|
||||
parser.add_argument("--dtype", choices=["bfloat16", "float32"], default="bfloat16")
|
||||
parser.add_argument("--device", default="cuda")
|
||||
parser.add_argument("--compile-model", action="store_true")
|
||||
parser.add_argument("--attn-implementation", default="eager")
|
||||
parser.add_argument("--enable-regional-compile", action="store_true")
|
||||
parser.add_argument("--enable-warmup", action="store_true")
|
||||
parser.add_argument("--strict-load", action="store_true")
|
||||
parser.add_argument("--output-dir", default="outputs/internvla_a1/vllm_infer")
|
||||
parser.add_argument("--skip-plots", action="store_true")
|
||||
parser.add_argument("--benchmark-forward", action="store_true")
|
||||
parser.add_argument("--warmup-iters", type=int, default=3)
|
||||
parser.add_argument("--benchmark-iters", type=int, default=10)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def build_od_config(args: argparse.Namespace, processor_model_name: str) -> OmniDiffusionConfig:
|
||||
return OmniDiffusionConfig(
|
||||
model=str(Path(args.model_dir).resolve()),
|
||||
model_class_name="InternVLAA1Pipeline",
|
||||
dtype=tensor_dtype(args.dtype),
|
||||
custom_pipeline_args={
|
||||
"device": args.device,
|
||||
"dtype": args.dtype,
|
||||
"compile_model": args.compile_model,
|
||||
"attn_implementation": args.attn_implementation,
|
||||
"enable_regional_compile": args.enable_regional_compile,
|
||||
"enable_warmup": args.enable_warmup,
|
||||
"strict_load": args.strict_load,
|
||||
"processor_model_name": processor_model_name,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def build_dataset(args: argparse.Namespace) -> tuple[A2DOpenLoopDataset, InternVLAA1Config, InternVLAA1TrainMetadata]:
|
||||
model_dir = Path(args.model_dir)
|
||||
config = InternVLAA1Config.from_pretrained(model_dir)
|
||||
config.device = args.device
|
||||
config.dtype = args.dtype
|
||||
config.compile_model = args.compile_model
|
||||
config.attn_implementation = args.attn_implementation
|
||||
config.enable_regional_compile = args.enable_regional_compile
|
||||
|
||||
train_meta = InternVLAA1TrainMetadata.from_pretrained(model_dir)
|
||||
with open(model_dir / "stats.json", encoding="utf-8") as f:
|
||||
train_stats = json.load(f)["a2d"]
|
||||
|
||||
dataset = A2DOpenLoopDataset(
|
||||
args.dataset_dir,
|
||||
config=config,
|
||||
train_stats=train_stats,
|
||||
)
|
||||
return dataset, config, train_meta
|
||||
|
||||
|
||||
def run_one_path(
|
||||
pipeline,
|
||||
dataset: A2DOpenLoopDataset,
|
||||
config: InternVLAA1Config,
|
||||
args: argparse.Namespace,
|
||||
indices: list[int],
|
||||
) -> list[dict[str, object]]:
|
||||
results: list[dict[str, object]] = []
|
||||
for index in indices:
|
||||
sample = dataset.get_sample(index)
|
||||
batch_inputs, _ = collate_open_loop_samples([sample], device=args.device, dtype=tensor_dtype(args.dtype))
|
||||
noise = make_shared_noise(
|
||||
args.seed,
|
||||
index,
|
||||
(
|
||||
batch_inputs[OBS_STATE].shape[0],
|
||||
config.chunk_size,
|
||||
config.max_action_dim,
|
||||
),
|
||||
args.device,
|
||||
)
|
||||
pred = run_pipeline_forward(pipeline, batch_inputs, noise)
|
||||
pred = pred[:, :, : dataset.physical_action_dim].to(torch.float32).cpu()
|
||||
results.append(
|
||||
{
|
||||
"path": "registry",
|
||||
"index": index,
|
||||
"episode_index": sample.episode_index,
|
||||
"task": sample.task,
|
||||
"seed": args.seed,
|
||||
"shape": list(pred.shape),
|
||||
"mean": float(pred.mean().item()),
|
||||
"std": float(pred.std().item()),
|
||||
"action_sha256": tensor_sha256(pred),
|
||||
"first_action_prefix": pred[0, 0, :8].tolist(),
|
||||
}
|
||||
)
|
||||
return results
|
||||
|
||||
|
||||
def run_pipeline_forward(
|
||||
pipeline,
|
||||
batch_inputs: dict[str, torch.Tensor],
|
||||
noise: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
output = pipeline.forward(
|
||||
OmniDiffusionRequest(
|
||||
prompts=[""],
|
||||
sampling_params=OmniDiffusionSamplingParams(
|
||||
extra_args={
|
||||
"batch_inputs": batch_inputs,
|
||||
"noise": noise,
|
||||
"decode_image": False,
|
||||
}
|
||||
),
|
||||
)
|
||||
)
|
||||
if output.error:
|
||||
raise RuntimeError(output.error)
|
||||
if output.output is None:
|
||||
raise RuntimeError("InternVLAA1Pipeline.forward returned no output tensor.")
|
||||
return output.output
|
||||
|
||||
|
||||
def _synchronize(device: str) -> None:
|
||||
if device.startswith("cuda") and torch.cuda.is_available():
|
||||
torch.cuda.synchronize()
|
||||
|
||||
|
||||
def _latency_summary(values_ms: list[float]) -> dict[str, float]:
|
||||
sorted_values = sorted(values_ms)
|
||||
p50_index = min(len(sorted_values) - 1, round(0.50 * (len(sorted_values) - 1)))
|
||||
p90_index = min(len(sorted_values) - 1, round(0.90 * (len(sorted_values) - 1)))
|
||||
return {
|
||||
"mean_ms": float(statistics.mean(sorted_values)),
|
||||
"stdev_ms": float(statistics.pstdev(sorted_values)) if len(sorted_values) > 1 else 0.0,
|
||||
"min_ms": float(sorted_values[0]),
|
||||
"max_ms": float(sorted_values[-1]),
|
||||
"p50_ms": float(sorted_values[p50_index]),
|
||||
"p90_ms": float(sorted_values[p90_index]),
|
||||
}
|
||||
|
||||
|
||||
def benchmark_forward(
|
||||
pipeline,
|
||||
dataset: A2DOpenLoopDataset,
|
||||
config: InternVLAA1Config,
|
||||
args: argparse.Namespace,
|
||||
index: int,
|
||||
output_dir: Path,
|
||||
) -> dict[str, object]:
|
||||
sample = dataset.get_sample(index)
|
||||
batch_inputs, _ = collate_open_loop_samples([sample], device=args.device, dtype=tensor_dtype(args.dtype))
|
||||
noise = make_shared_noise(
|
||||
args.seed,
|
||||
index,
|
||||
(
|
||||
batch_inputs[OBS_STATE].shape[0],
|
||||
config.chunk_size,
|
||||
config.max_action_dim,
|
||||
),
|
||||
args.device,
|
||||
)
|
||||
|
||||
_synchronize(args.device)
|
||||
cold_start_begin = time.perf_counter()
|
||||
pred = run_pipeline_forward(pipeline, batch_inputs, noise)
|
||||
_synchronize(args.device)
|
||||
cold_start_ms = (time.perf_counter() - cold_start_begin) * 1000.0
|
||||
|
||||
warmup_ms: list[float] = []
|
||||
for _ in range(args.warmup_iters):
|
||||
_synchronize(args.device)
|
||||
begin = time.perf_counter()
|
||||
_ = run_pipeline_forward(pipeline, batch_inputs, noise)
|
||||
_synchronize(args.device)
|
||||
warmup_ms.append((time.perf_counter() - begin) * 1000.0)
|
||||
|
||||
benchmark_ms: list[float] = []
|
||||
for _ in range(args.benchmark_iters):
|
||||
_synchronize(args.device)
|
||||
begin = time.perf_counter()
|
||||
_ = run_pipeline_forward(pipeline, batch_inputs, noise)
|
||||
_synchronize(args.device)
|
||||
benchmark_ms.append((time.perf_counter() - begin) * 1000.0)
|
||||
|
||||
pred = pred[:, :, : dataset.physical_action_dim].to(torch.float32).cpu()
|
||||
summary = {
|
||||
"mode": "forward_latency",
|
||||
"model_dir": str(Path(args.model_dir).resolve()),
|
||||
"dataset_dir": str(Path(args.dataset_dir).resolve()),
|
||||
"sample_index": index,
|
||||
"episode_index": sample.episode_index,
|
||||
"task": sample.task,
|
||||
"device": args.device,
|
||||
"dtype": args.dtype,
|
||||
"attn_implementation": args.attn_implementation,
|
||||
"enable_regional_compile": args.enable_regional_compile,
|
||||
"warmup_iters": args.warmup_iters,
|
||||
"benchmark_iters": args.benchmark_iters,
|
||||
"output_shape": list(pred.shape),
|
||||
"cold_start_ms": cold_start_ms,
|
||||
"warmup_summary": _latency_summary(warmup_ms) if warmup_ms else {},
|
||||
"benchmark_summary": _latency_summary(benchmark_ms) if benchmark_ms else {},
|
||||
"benchmark_samples_ms": benchmark_ms,
|
||||
}
|
||||
with open(output_dir / "forward_latency.json", "w", encoding="utf-8") as f:
|
||||
json.dump(summary, f, ensure_ascii=False, indent=2)
|
||||
print(json.dumps(summary, ensure_ascii=False, indent=2))
|
||||
return summary
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
args.model_dir = _required_path_arg("INTERNVLA_A1_MODEL_DIR", args.model_dir)
|
||||
args.dataset_dir = _required_path_arg("INTERNVLA_A1_DATASET_DIR", args.dataset_dir)
|
||||
output_dir = Path(args.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
dataset, config, train_meta = build_dataset(args)
|
||||
indices = select_indices(dataset, args.num_samples)
|
||||
od_config = build_od_config(args, train_meta.processor_model_name)
|
||||
|
||||
eval_summaries: dict[str, object] = {}
|
||||
pipeline = initialize_model(od_config)
|
||||
if args.benchmark_forward:
|
||||
benchmark_forward(
|
||||
pipeline=pipeline,
|
||||
dataset=dataset,
|
||||
config=config,
|
||||
args=args,
|
||||
index=indices[0],
|
||||
output_dir=output_dir,
|
||||
)
|
||||
return
|
||||
results = run_one_path(
|
||||
pipeline=pipeline,
|
||||
dataset=dataset,
|
||||
config=config,
|
||||
args=args,
|
||||
indices=indices,
|
||||
)
|
||||
if args.num_episodes > 0:
|
||||
eval_summaries["registry"] = run_open_loop_evaluation(
|
||||
mode="vllm_registry",
|
||||
policy=pipeline,
|
||||
config=config,
|
||||
dataset=dataset,
|
||||
train_meta=train_meta,
|
||||
collate_samples=collate_open_loop_samples,
|
||||
run_sample_actions=lambda policy, batch_inputs, noise: run_pipeline_forward(policy, batch_inputs, noise),
|
||||
num_episodes=args.num_episodes,
|
||||
seed=args.seed,
|
||||
device=args.device,
|
||||
dtype=tensor_dtype(args.dtype),
|
||||
output_dir=output_dir / "registry",
|
||||
skip_plots=args.skip_plots,
|
||||
)
|
||||
|
||||
summary = {
|
||||
"mode": "registry",
|
||||
"model_dir": str(Path(args.model_dir).resolve()),
|
||||
"dataset_dir": str(Path(args.dataset_dir).resolve()),
|
||||
"device": args.device,
|
||||
"dtype": args.dtype,
|
||||
"attn_implementation": args.attn_implementation,
|
||||
"enable_regional_compile": args.enable_regional_compile,
|
||||
"seed": args.seed,
|
||||
"indices": indices,
|
||||
"results": results,
|
||||
"output_dir": str(output_dir.resolve()),
|
||||
"eval_summaries": eval_summaries,
|
||||
}
|
||||
with open(output_dir / "summary.json", "w", encoding="utf-8") as f:
|
||||
json.dump(summary, f, ensure_ascii=False, indent=2)
|
||||
print(json.dumps(summary, ensure_ascii=False, indent=2))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,471 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import math
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import pyarrow.parquet as pq
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import torchvision
|
||||
|
||||
from vllm_omni.diffusion.models.internvla_a1.config import ( # noqa: E402
|
||||
OBS_IMAGES,
|
||||
OBS_STATE,
|
||||
OBS_TASK,
|
||||
InternVLAA1Config,
|
||||
)
|
||||
|
||||
|
||||
def tensor_dtype(name: str) -> torch.dtype:
|
||||
return torch.bfloat16 if name == "bfloat16" else torch.float32
|
||||
|
||||
|
||||
def select_indices(dataset, num_samples: int) -> list[int]:
|
||||
indices: list[int] = []
|
||||
for _, episode_indices in dataset.episode_start_indices(max_episodes=num_samples):
|
||||
if episode_indices:
|
||||
indices.append(episode_indices[0])
|
||||
return indices[:num_samples]
|
||||
|
||||
|
||||
def make_shared_noise(seed: int, sample_index: int, shape: tuple[int, ...], device: str) -> torch.Tensor:
|
||||
generator = torch.Generator(device="cpu")
|
||||
generator.manual_seed(seed + sample_index)
|
||||
noise = torch.randn(shape, generator=generator, dtype=torch.float32)
|
||||
return noise.to(device=device)
|
||||
|
||||
|
||||
def tensor_sha256(tensor: torch.Tensor) -> str:
|
||||
array = tensor.detach().contiguous().cpu().numpy()
|
||||
return hashlib.sha256(array.tobytes()).hexdigest()
|
||||
|
||||
|
||||
def _clamp_index(index: int, start: int, end: int) -> int:
|
||||
return max(start, min(end - 1, index))
|
||||
|
||||
|
||||
def _load_parquet_rows(path: Path) -> list[dict[str, Any]]:
|
||||
return pq.read_table(path).to_pylist()
|
||||
|
||||
|
||||
def _load_json(path: Path) -> dict[str, Any]:
|
||||
with open(path, encoding="utf-8") as f:
|
||||
return json.load(f)
|
||||
|
||||
|
||||
def _stack_stats(stats: dict[str, Any], keys: list[str]) -> dict[str, torch.Tensor]:
|
||||
result = {}
|
||||
for stat_name in ("mean", "std"):
|
||||
values = []
|
||||
for key in keys:
|
||||
values.extend(stats[key][stat_name])
|
||||
result[stat_name] = torch.tensor(values, dtype=torch.float32)
|
||||
return result
|
||||
|
||||
|
||||
def normalize_vector(values: torch.Tensor, stats: dict[str, torch.Tensor]) -> torch.Tensor:
|
||||
denom = torch.where(stats["std"] == 0, torch.ones_like(stats["std"]), stats["std"])
|
||||
return (values - stats["mean"]) / denom
|
||||
|
||||
|
||||
def unnormalize_vector(values: torch.Tensor, stats: dict[str, torch.Tensor]) -> torch.Tensor:
|
||||
return values * stats["std"] + stats["mean"]
|
||||
|
||||
|
||||
class TorchvisionVideoReaderCache:
|
||||
def __init__(self, backend: str = "pyav") -> None:
|
||||
self.backend = backend
|
||||
self._readers: dict[str, Any] = {}
|
||||
torchvision.set_video_backend(backend)
|
||||
|
||||
def get(self, path: str) -> Any:
|
||||
reader = self._readers.get(path)
|
||||
if reader is None:
|
||||
reader = torchvision.io.VideoReader(path, "video")
|
||||
self._readers[path] = reader
|
||||
return reader
|
||||
|
||||
def decode_frames(self, path: str, timestamps: list[float], tolerance_s: float = 1e-4) -> torch.Tensor:
|
||||
reader = self.get(path)
|
||||
first_ts = min(timestamps)
|
||||
last_ts = max(timestamps)
|
||||
reader.seek(first_ts, keyframes_only=self.backend == "pyav")
|
||||
|
||||
loaded_frames: list[torch.Tensor] = []
|
||||
loaded_ts: list[float] = []
|
||||
for frame in reader:
|
||||
current_ts = float(frame["pts"])
|
||||
loaded_frames.append(frame["data"])
|
||||
loaded_ts.append(current_ts)
|
||||
if current_ts >= last_ts:
|
||||
break
|
||||
|
||||
query_ts = torch.tensor(timestamps, dtype=torch.float32)
|
||||
loaded_ts_tensor = torch.tensor(loaded_ts, dtype=torch.float32)
|
||||
distances = torch.cdist(query_ts[:, None], loaded_ts_tensor[:, None], p=1)
|
||||
min_dist, argmin = distances.min(dim=1)
|
||||
if not torch.all(min_dist < tolerance_s):
|
||||
raise RuntimeError(
|
||||
f"Video timestamps are outside tolerance: query={query_ts.tolist()} "
|
||||
f"loaded={loaded_ts_tensor.tolist()} path={path}"
|
||||
)
|
||||
return torch.stack([loaded_frames[i] for i in argmin]).float() / 255.0
|
||||
|
||||
|
||||
@dataclass
|
||||
class A2DOpenLoopSample:
|
||||
index: int
|
||||
episode_index: int
|
||||
task: str
|
||||
state_raw: torch.Tensor
|
||||
action_raw: torch.Tensor
|
||||
inputs: dict[str, torch.Tensor]
|
||||
|
||||
|
||||
class A2DOpenLoopDataset:
|
||||
image_keys = [
|
||||
"observation.images.head",
|
||||
"observation.images.hand_left",
|
||||
"observation.images.hand_right",
|
||||
]
|
||||
state_keys = [
|
||||
"observation.states.joint.position",
|
||||
"observation.states.effector.position",
|
||||
]
|
||||
action_keys = [
|
||||
"actions.joint.position",
|
||||
"actions.effector.position",
|
||||
]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dataset_root: str | Path,
|
||||
*,
|
||||
config: InternVLAA1Config,
|
||||
train_stats: dict[str, Any],
|
||||
image_offsets: tuple[int, int] = (-15, 0),
|
||||
tolerance_s: float = 1e-4,
|
||||
) -> None:
|
||||
self.root = Path(dataset_root)
|
||||
self.config = config
|
||||
self.info = _load_json(self.root / "meta" / "info.json")
|
||||
self.dataset_stats = _load_json(self.root / "meta" / "stats.json")
|
||||
self.data_rows = _load_parquet_rows(self.root / "data" / "chunk-000" / "file-000.parquet")
|
||||
self.episode_rows = _load_parquet_rows(self.root / "meta" / "episodes" / "chunk-000" / "file-000.parquet")
|
||||
self.task_rows = _load_parquet_rows(self.root / "meta" / "tasks.parquet")
|
||||
|
||||
self.state_stats = _stack_stats(train_stats, self.state_keys)
|
||||
self.action_stats = _stack_stats(train_stats, self.action_keys)
|
||||
self.image_offsets = image_offsets
|
||||
self.tolerance_s = tolerance_s
|
||||
self.video_reader = TorchvisionVideoReaderCache(backend="pyav")
|
||||
|
||||
@property
|
||||
def num_episodes(self) -> int:
|
||||
return len(self.episode_rows)
|
||||
|
||||
@property
|
||||
def physical_action_dim(self) -> int:
|
||||
return 16
|
||||
|
||||
def episode_start_indices(self, max_episodes: int | None = None) -> list[tuple[int, list[int]]]:
|
||||
rows = self.episode_rows if max_episodes is None else self.episode_rows[:max_episodes]
|
||||
result = []
|
||||
for ep in rows:
|
||||
start = int(ep["dataset_from_index"])
|
||||
end = int(ep["dataset_to_index"])
|
||||
result.append((int(ep["episode_index"]), list(range(start, end, self.config.chunk_size))))
|
||||
return result
|
||||
|
||||
def _task_text(self, task_index: int) -> str:
|
||||
return self.task_rows[task_index]["__index_level_0__"]
|
||||
|
||||
def _episode_for_index(self, row: dict[str, Any]) -> dict[str, Any]:
|
||||
return self.episode_rows[int(row["episode_index"])]
|
||||
|
||||
def _state_vector(self, row: dict[str, Any]) -> torch.Tensor:
|
||||
return torch.tensor(row[self.state_keys[0]] + row[self.state_keys[1]], dtype=torch.float32)
|
||||
|
||||
def _action_vector(self, row: dict[str, Any]) -> torch.Tensor:
|
||||
return torch.tensor(row[self.action_keys[0]] + row[self.action_keys[1]], dtype=torch.float32)
|
||||
|
||||
def _query_rows(self, idx: int, deltas: list[int]) -> list[dict[str, Any]]:
|
||||
row = self.data_rows[idx]
|
||||
episode = self._episode_for_index(row)
|
||||
start = int(episode["dataset_from_index"])
|
||||
end = int(episode["dataset_to_index"])
|
||||
return [self.data_rows[_clamp_index(idx + delta, start, end)] for delta in deltas]
|
||||
|
||||
def _decode_camera_history(
|
||||
self, episode: dict[str, Any], camera_key: str, rows: list[dict[str, Any]]
|
||||
) -> torch.Tensor:
|
||||
timestamps = [float(r["timestamp"]) for r in rows]
|
||||
shifted = [float(episode[f"videos/{camera_key}/from_timestamp"]) + ts for ts in timestamps]
|
||||
chunk_idx = int(episode[f"videos/{camera_key}/chunk_index"])
|
||||
file_idx = int(episode[f"videos/{camera_key}/file_index"])
|
||||
path = self.root / self.info["video_path"].format(
|
||||
video_key=camera_key,
|
||||
chunk_index=chunk_idx,
|
||||
file_index=file_idx,
|
||||
)
|
||||
frames = self.video_reader.decode_frames(str(path), shifted, tolerance_s=self.tolerance_s)
|
||||
return frames
|
||||
|
||||
def get_sample(self, idx: int) -> A2DOpenLoopSample:
|
||||
row = self.data_rows[idx]
|
||||
episode = self._episode_for_index(row)
|
||||
image_rows = self._query_rows(idx, list((-15, 0)))
|
||||
action_rows = self._query_rows(idx, list(range(self.config.chunk_size)))
|
||||
camera_images = [self._decode_camera_history(episode, camera_key, image_rows) for camera_key in self.image_keys]
|
||||
state_raw = self._state_vector(row)
|
||||
state_norm = normalize_vector(state_raw, self.state_stats)
|
||||
action_raw = torch.stack([self._action_vector(action_row) for action_row in action_rows], dim=0)
|
||||
task = self._task_text(int(row["task_index"]))
|
||||
inputs = {
|
||||
OBS_STATE: state_norm,
|
||||
OBS_TASK: task,
|
||||
f"{OBS_IMAGES}.image0": camera_images[0],
|
||||
f"{OBS_IMAGES}.image1": camera_images[1],
|
||||
f"{OBS_IMAGES}.image2": camera_images[2],
|
||||
f"{OBS_IMAGES}.image0_mask": torch.tensor(True),
|
||||
f"{OBS_IMAGES}.image1_mask": torch.tensor(True),
|
||||
f"{OBS_IMAGES}.image2_mask": torch.tensor(True),
|
||||
}
|
||||
return A2DOpenLoopSample(
|
||||
index=idx,
|
||||
episode_index=int(row["episode_index"]),
|
||||
task=task,
|
||||
state_raw=state_raw,
|
||||
action_raw=action_raw,
|
||||
inputs=inputs,
|
||||
)
|
||||
|
||||
|
||||
def collate_open_loop_samples(
|
||||
samples: list[A2DOpenLoopSample],
|
||||
*,
|
||||
device: str,
|
||||
dtype: torch.dtype,
|
||||
) -> tuple[dict[str, torch.Tensor], dict[str, Any]]:
|
||||
first = samples[0]
|
||||
batch_inputs: dict[str, torch.Tensor] = {}
|
||||
for key in first.inputs:
|
||||
values = [sample.inputs[key] for sample in samples]
|
||||
if isinstance(values[0], torch.Tensor):
|
||||
tensor = torch.stack(values, dim=0)
|
||||
if tensor.dtype in (torch.int64, torch.bool):
|
||||
batch_inputs[key] = tensor.to(device=device)
|
||||
else:
|
||||
batch_inputs[key] = tensor.to(device=device, dtype=dtype)
|
||||
else:
|
||||
batch_inputs[key] = values
|
||||
|
||||
metadata = {
|
||||
"indices": [sample.index for sample in samples],
|
||||
"episode_indices": [sample.episode_index for sample in samples],
|
||||
"tasks": [sample.task for sample in samples],
|
||||
"state_raw": torch.stack([sample.state_raw for sample in samples], dim=0),
|
||||
"action_raw": torch.stack([sample.action_raw for sample in samples], dim=0),
|
||||
}
|
||||
return batch_inputs, metadata
|
||||
|
||||
|
||||
def plot_prediction_series(
|
||||
*,
|
||||
series: dict[str, np.ndarray],
|
||||
output_path: Path,
|
||||
title: str,
|
||||
) -> None:
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
first_series = next(iter(series.values()))
|
||||
num_dims = int(first_series.shape[1])
|
||||
num_cols = 2
|
||||
num_rows = math.ceil(num_dims / num_cols)
|
||||
fig, axs = plt.subplots(num_rows, num_cols, figsize=(16, max(6, 3 * num_rows)))
|
||||
axs = np.array(axs).reshape(-1)
|
||||
x_values = np.arange(first_series.shape[0])
|
||||
styles = [
|
||||
("Ground Truth", "blue", "-", 1.5),
|
||||
("VLLM", "green", "-.", 1.4),
|
||||
("Predicted", "red", "--", 1.4),
|
||||
]
|
||||
style_map = {name: (color, linestyle, linewidth) for name, color, linestyle, linewidth in styles}
|
||||
|
||||
for dim in range(num_dims):
|
||||
ax = axs[dim]
|
||||
for name, values in series.items():
|
||||
color, linestyle, linewidth = style_map.get(name, (None, "-", 1.2))
|
||||
ax.plot(
|
||||
x_values,
|
||||
values[:, dim],
|
||||
label=name,
|
||||
color=color,
|
||||
linestyle=linestyle,
|
||||
linewidth=linewidth,
|
||||
)
|
||||
ax.set_title(f"Dimension {dim + 1}")
|
||||
ax.set_xlabel("Time Step / Sample Index")
|
||||
ax.set_ylabel(f"Value Dim {dim + 1}")
|
||||
ax.legend(loc="upper right")
|
||||
ax.grid(True, linestyle="--", alpha=0.7)
|
||||
|
||||
for dim in range(num_dims, len(axs)):
|
||||
axs[dim].axis("off")
|
||||
|
||||
fig.suptitle(title, fontsize=16)
|
||||
plt.tight_layout(rect=(0, 0, 1, 0.97))
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
plt.savefig(output_path)
|
||||
plt.close(fig)
|
||||
|
||||
|
||||
def summarize_prediction_metrics(
|
||||
*,
|
||||
gt_tensor: torch.Tensor,
|
||||
pred_tensor: torch.Tensor,
|
||||
joint_dims: int = 14,
|
||||
) -> dict[str, float]:
|
||||
return {
|
||||
"mse": float(F.mse_loss(gt_tensor, pred_tensor, reduction="mean").item()),
|
||||
"mae": float(F.l1_loss(gt_tensor, pred_tensor, reduction="mean").item()),
|
||||
"mse_joint": float(F.mse_loss(gt_tensor[:, :joint_dims], pred_tensor[:, :joint_dims], reduction="mean").item()),
|
||||
"mae_joint": float(F.l1_loss(gt_tensor[:, :joint_dims], pred_tensor[:, :joint_dims], reduction="mean").item()),
|
||||
"mse_gripper": float(
|
||||
F.mse_loss(gt_tensor[:, joint_dims:], pred_tensor[:, joint_dims:], reduction="mean").item()
|
||||
),
|
||||
"mae_gripper": float(
|
||||
F.l1_loss(gt_tensor[:, joint_dims:], pred_tensor[:, joint_dims:], reduction="mean").item()
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def summarize_metric_list(values: list[float]) -> float | None:
|
||||
return float(np.mean(values)) if values else None
|
||||
|
||||
|
||||
def run_open_loop_evaluation(
|
||||
*,
|
||||
mode: str,
|
||||
policy,
|
||||
config,
|
||||
dataset,
|
||||
train_meta,
|
||||
collate_samples: Callable[..., tuple[dict[str, torch.Tensor], dict[str, Any]]],
|
||||
run_sample_actions: Callable[[Any, dict[str, torch.Tensor], torch.Tensor], torch.Tensor],
|
||||
num_episodes: int,
|
||||
seed: int,
|
||||
device: str,
|
||||
dtype: torch.dtype,
|
||||
output_dir: str | Path,
|
||||
skip_plots: bool,
|
||||
) -> dict[str, Any]:
|
||||
output_dir = Path(output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
plots_dir = output_dir / "plots"
|
||||
plots_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
mse_values: list[float] = []
|
||||
mae_values: list[float] = []
|
||||
mse_joint_values: list[float] = []
|
||||
mae_joint_values: list[float] = []
|
||||
mse_gripper_values: list[float] = []
|
||||
mae_gripper_values: list[float] = []
|
||||
per_episode: list[dict[str, Any]] = []
|
||||
|
||||
plotting_available = True
|
||||
if not skip_plots:
|
||||
try:
|
||||
import matplotlib.pyplot as _plt # noqa: F401
|
||||
except ModuleNotFoundError:
|
||||
plotting_available = False
|
||||
print("[warning] matplotlib is not installed; plots will be skipped.")
|
||||
|
||||
episode_specs = dataset.episode_start_indices(max_episodes=num_episodes)
|
||||
for episode_index, indices in episode_specs:
|
||||
print(f"[{mode}] episode: {episode_index}")
|
||||
pred_chunks = []
|
||||
gt_chunks = []
|
||||
visited_indices = []
|
||||
task = None
|
||||
|
||||
for index in indices:
|
||||
sample = dataset.get_sample(index)
|
||||
task = sample.task
|
||||
batch_inputs, meta = collate_samples([sample], device=device, dtype=dtype)
|
||||
noise = make_shared_noise(
|
||||
seed,
|
||||
index,
|
||||
(batch_inputs[OBS_STATE].shape[0], config.chunk_size, config.max_action_dim),
|
||||
device,
|
||||
)
|
||||
with torch.no_grad():
|
||||
pred = run_sample_actions(policy, batch_inputs, noise)
|
||||
|
||||
pred = pred[:, :, : dataset.physical_action_dim].to(torch.float32).cpu()
|
||||
pred_phys = unnormalize_vector(pred, dataset.action_stats)
|
||||
if train_meta.action_mode == "delta":
|
||||
pred_phys[:, :, :14] += meta["state_raw"][:, None, :14]
|
||||
|
||||
gt_phys = meta["action_raw"][:, :, : dataset.physical_action_dim].to(torch.float32).cpu()
|
||||
pred_chunks.append(pred_phys[0])
|
||||
gt_chunks.append(gt_phys[0])
|
||||
visited_indices.append(index)
|
||||
|
||||
pred_tensor = torch.cat(pred_chunks, dim=0)
|
||||
gt_tensor = torch.cat(gt_chunks, dim=0)
|
||||
metrics = summarize_prediction_metrics(gt_tensor=gt_tensor, pred_tensor=pred_tensor)
|
||||
|
||||
mse_values.append(metrics["mse"])
|
||||
mae_values.append(metrics["mae"])
|
||||
mse_joint_values.append(metrics["mse_joint"])
|
||||
mae_joint_values.append(metrics["mae_joint"])
|
||||
mse_gripper_values.append(metrics["mse_gripper"])
|
||||
mae_gripper_values.append(metrics["mae_gripper"])
|
||||
|
||||
if not skip_plots and plotting_available:
|
||||
plot_prediction_series(
|
||||
series={
|
||||
"Ground Truth": gt_tensor.numpy(),
|
||||
"VLLM": pred_tensor.numpy(),
|
||||
},
|
||||
output_path=plots_dir / f"{mode}_open_loop_ep{episode_index}.jpg",
|
||||
title=f"{mode} Ground Truth vs VLLM",
|
||||
)
|
||||
|
||||
episode_log = {
|
||||
"episode_id": int(episode_index),
|
||||
"task": task,
|
||||
"visited_indices": visited_indices,
|
||||
"num_pred_steps": int(pred_tensor.shape[0]),
|
||||
**metrics,
|
||||
}
|
||||
per_episode.append(episode_log)
|
||||
print(json.dumps({"mode": mode, **episode_log}, ensure_ascii=False))
|
||||
|
||||
summary = {
|
||||
"mode": mode,
|
||||
"num_episodes": len(per_episode),
|
||||
"mse": mse_values,
|
||||
"mae": mae_values,
|
||||
"average_mse": summarize_metric_list(mse_values),
|
||||
"average_mae": summarize_metric_list(mae_values),
|
||||
"mse_joint": mse_joint_values,
|
||||
"mae_joint": mae_joint_values,
|
||||
"average_mse_joint": summarize_metric_list(mse_joint_values),
|
||||
"average_mae_joint": summarize_metric_list(mae_joint_values),
|
||||
"mse_gripper": mse_gripper_values,
|
||||
"mae_gripper": mae_gripper_values,
|
||||
"average_mse_gripper": summarize_metric_list(mse_gripper_values),
|
||||
"average_mae_gripper": summarize_metric_list(mae_gripper_values),
|
||||
"episodes": per_episode,
|
||||
}
|
||||
with open(output_dir / "log.json", "w", encoding="utf-8") as f:
|
||||
json.dump(summary, f, indent=2, ensure_ascii=False)
|
||||
return summary
|
||||
+23
@@ -0,0 +1,23 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
REPO_ROOT="$(cd "$ROOT_DIR/../../.." && pwd)"
|
||||
: "${INTERNVLA_A1_MODEL_DIR:?Please export INTERNVLA_A1_MODEL_DIR=/path/to/InternVLA-A1-3B-ft-pen}"
|
||||
: "${INTERNVLA_A1_DATASET_DIR:?Please export INTERNVLA_A1_DATASET_DIR=/path/to/Genie1-Place_Markpen}"
|
||||
: "${INTERNVLA_A1_PROCESSOR_DIR:?Please export INTERNVLA_A1_PROCESSOR_DIR=/path/to/Qwen3-VL-2B-Instruct}"
|
||||
: "${INTERNVLA_A1_COSMOS_DIR:?Please export INTERNVLA_A1_COSMOS_DIR=/path/to/Cosmos-Tokenizer-CI8x8-SafeTensor}"
|
||||
INTERNVLA_A1_OUTPUT_DIR="${INTERNVLA_A1_OUTPUT_DIR:-$REPO_ROOT/outputs/internvla_a1/vllm_infer}"
|
||||
|
||||
export INTERNVLA_A1_MODEL_DIR
|
||||
export INTERNVLA_A1_DATASET_DIR
|
||||
export INTERNVLA_A1_PROCESSOR_DIR
|
||||
export INTERNVLA_A1_COSMOS_DIR
|
||||
|
||||
python "$ROOT_DIR/end2end.py" \
|
||||
--model-dir "$INTERNVLA_A1_MODEL_DIR" \
|
||||
--dataset-dir "$INTERNVLA_A1_DATASET_DIR" \
|
||||
--output-dir "$INTERNVLA_A1_OUTPUT_DIR" \
|
||||
--num-episodes "${INTERNVLA_A1_NUM_EPISODES:-1}" \
|
||||
--attn-implementation eager \
|
||||
"$@"
|
||||
@@ -24,9 +24,6 @@ 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()
|
||||
|
||||
|
||||
|
||||
@@ -23,6 +23,7 @@ from message_convert import (
|
||||
from vllm import SamplingParams
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniTokensPrompt
|
||||
|
||||
@@ -438,6 +439,7 @@ def parse_args():
|
||||
"vllm_omni/deploy/mimo_audio.yaml based on the HF model_type.",
|
||||
)
|
||||
|
||||
nullify_stage_engine_defaults(parser)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
|
||||
Loaded 100 of 552 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user