Compare commits
122
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
f55ea28005 | ||
|
|
a353594e2e | ||
|
|
d9912eb6e9 | ||
|
|
d0edd3e4e3 | ||
|
|
bd8b9b6142 | ||
|
|
3fe46360f0 | ||
|
|
637abb759f | ||
|
|
c0d90b64c6 | ||
|
|
c139ee4aed | ||
|
|
c51adeacde | ||
|
|
0db2b3c8f2 | ||
|
|
94e6a96fb8 | ||
|
|
49c1abea17 | ||
|
|
61e0ec517d | ||
|
|
9f667ee151 | ||
|
|
6f87a9382f | ||
|
|
4be54ccec4 | ||
|
|
055963b3fd | ||
|
|
e408fd97aa | ||
|
|
f7107a384b | ||
|
|
a67859574d | ||
|
|
590ab5b3e9 | ||
|
|
66d5478dad | ||
|
|
e3d178cd66 | ||
|
|
d96d63cb7f | ||
|
|
574ec99ef4 | ||
|
|
3a5483d339 | ||
|
|
9857ed2839 | ||
|
|
c9e2d2c394 | ||
|
|
15df8ec97f | ||
|
|
2dde219c00 | ||
|
|
c392ce21e9 | ||
|
|
ce916f4233 | ||
|
|
029c85393a | ||
|
|
8305477750 | ||
|
|
fb5d8ffd4e | ||
|
|
d15c91d6d2 | ||
|
|
16b99c3444 | ||
|
|
807f3133d7 | ||
|
|
04a6e5a992 | ||
|
|
9f1d61b0f7 | ||
|
|
cf66ee86a4 | ||
|
|
b569cb56b6 | ||
|
|
29936a3edd | ||
|
|
d2b9f9ff0f | ||
|
|
18d2d082aa | ||
|
|
fc32da7624 | ||
|
|
5466ed419b | ||
|
|
c09b039f37 | ||
|
|
0adbac33a5 | ||
|
|
91a818629c | ||
|
|
ed8cfdfb1b | ||
|
|
c2888a8a36 | ||
|
|
6d439ebad8 | ||
|
|
61092acd82 | ||
|
|
8f8ba9f8d6 | ||
|
|
4febac1efa | ||
|
|
0849c5efd8 | ||
|
|
f3d078e9f2 | ||
|
|
e776082095 | ||
|
|
8a59923fc7 | ||
|
|
482e878ed1 | ||
|
|
8d75700321 | ||
|
|
648a0b81a1 | ||
|
|
44770f9f20 | ||
|
|
69652a2983 | ||
|
|
e7e05a55cf | ||
|
|
7f112041da | ||
|
|
7bcd43f456 | ||
|
|
47ba887b93 | ||
|
|
c7c1c3b04a | ||
|
|
6b93459826 | ||
|
|
2a630313ce | ||
|
|
8619321b90 | ||
|
|
f8777d85bc | ||
|
|
c39225724d | ||
|
|
5c14757ef0 | ||
|
|
fb9f20a921 | ||
|
|
7217557950 | ||
|
|
93f1878229 | ||
|
|
50bd1961bf | ||
|
|
4cfa9bf843 | ||
|
|
9ead0d890d | ||
|
|
b2b061a9b3 | ||
|
|
3cdf28828e | ||
|
|
d16be833e1 | ||
|
|
b239bf0b7d | ||
|
|
ad5e4e5710 | ||
|
|
d71c9ae60f | ||
|
|
4ac25287fd | ||
|
|
8b102bdac1 | ||
|
|
c6b7597126 | ||
|
|
b0efb6ede0 | ||
|
|
9bf8f8ce79 | ||
|
|
ec6ee1e336 | ||
|
|
5aef6b9c71 | ||
|
|
1f5eca5686 | ||
|
|
df93b3a965 | ||
|
|
9eeac248f2 | ||
|
|
77d773ad47 | ||
|
|
d6b55b9196 | ||
|
|
61cd532cff | ||
|
|
5a58a1aac0 | ||
|
|
c14ab819e3 | ||
|
|
65d76a4127 | ||
|
|
560023d5d1 | ||
|
|
8a58394379 | ||
|
|
602f831068 | ||
|
|
b147553d67 | ||
|
|
0833470d41 | ||
|
|
cb6d012d91 | ||
|
|
f1ed949fcf | ||
|
|
a8f1394699 | ||
|
|
d8b35035ec | ||
|
|
af9368d109 | ||
|
|
b4a96b0db4 | ||
|
|
a5574a2fb0 | ||
|
|
28aee51a59 | ||
|
|
e7171b7549 | ||
|
|
a2a23a1499 | ||
|
|
1bd1d9d618 | ||
|
|
15f9d60dd4 |
No files matched your search
@@ -13,7 +13,7 @@ steps:
|
||||
- label: "Upload Ready Pipeline"
|
||||
depends_on: image-build
|
||||
key: upload-ready-pipeline
|
||||
if: build.branch != "main"
|
||||
if: build.branch != "main" && build.pull_request.labels includes "ready"
|
||||
commands:
|
||||
- buildkite-agent pipeline upload .buildkite/test-ready.yml
|
||||
agents:
|
||||
@@ -23,17 +23,17 @@ steps:
|
||||
- label: "Upload Merge Pipeline"
|
||||
depends_on: image-build
|
||||
key: upload-merge-pipeline
|
||||
if: build.branch == "main"
|
||||
if: build.branch == "main" && build.env("NIGHTLY") != "1"
|
||||
commands:
|
||||
- buildkite-agent pipeline upload .buildkite/test-merge.yml
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
# L4 Test
|
||||
# L4 Test — main+NIGHTLY=1 (scheduled), or PR with label nightly-test (e.g. add label then Rebuild)
|
||||
- label: "Upload Nightly Pipeline"
|
||||
depends_on: image-build
|
||||
key: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
if: '(build.branch == "main" && build.env("NIGHTLY") == "1") || (build.branch != "main" && build.pull_request.labels includes "nightly-test")'
|
||||
commands:
|
||||
- buildkite-agent pipeline upload .buildkite/test-nightly.yml
|
||||
agents:
|
||||
|
||||
+43
-39
@@ -112,6 +112,7 @@ steps:
|
||||
depends_on: upload-merge-pipeline
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py
|
||||
- pytest -s -v tests/diffusion/distributed/test_ulysses_uaa_perf.py
|
||||
agents:
|
||||
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
|
||||
plugins:
|
||||
@@ -187,12 +188,13 @@ steps:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Omni Model Test"
|
||||
timeout_in_minutes: 15
|
||||
timeout_in_minutes: 30
|
||||
depends_on: upload-merge-pipeline
|
||||
commands:
|
||||
- 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/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:
|
||||
@@ -214,6 +216,7 @@ steps:
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY="1"
|
||||
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
|
||||
- pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "advanced_model" --run-level "advanced_model"
|
||||
- pytest -s -v tests/e2e/online_serving/test_mimo_audio.py -m "advanced_model" --run-level "advanced_model"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
@@ -336,41 +339,42 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
# - label: "Voxtral TTS E2E Test"
|
||||
# timeout_in_minutes: 20
|
||||
# depends_on: upload-merge-pipeline
|
||||
# commands:
|
||||
# - |
|
||||
# timeout 20m bash -c '
|
||||
# export VLLM_LOGGING_LEVEL=DEBUG
|
||||
# export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
|
||||
# '
|
||||
# agents:
|
||||
# queue: "mithril-h100-pool"
|
||||
# plugins:
|
||||
# - kubernetes:
|
||||
# podSpec:
|
||||
# containers:
|
||||
# - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
# resources:
|
||||
# limits:
|
||||
# nvidia.com/gpu: 1
|
||||
# volumeMounts:
|
||||
# - name: devshm
|
||||
# mountPath: /dev/shm
|
||||
# - name: hf-cache
|
||||
# mountPath: /root/.cache/huggingface
|
||||
# env:
|
||||
# - name: HF_HOME
|
||||
# value: /root/.cache/huggingface
|
||||
# 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: "Voxtral-TTS E2E Test"
|
||||
timeout_in_minutes: 20
|
||||
depends_on: upload-merge-pipeline
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
|
||||
pytest -s -v tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
|
||||
'
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 1
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
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
|
||||
+268
-24
@@ -2,7 +2,7 @@ steps:
|
||||
- label: ":full_moon: Omni Model Test with H100"
|
||||
timeout_in_minutes: 90
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- |
|
||||
@@ -11,10 +11,11 @@ steps:
|
||||
EXIT1=$$?
|
||||
pytest -s -v tests/examples/online_serving/test_qwen3_omni.py -m "advanced_model" --run-level "advanced_model"
|
||||
EXIT2=$$?
|
||||
# TODO: uncomment when mistralai/tts-model is public
|
||||
# pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
|
||||
# EXIT3=$$?
|
||||
exit $$((EXIT1 | EXIT2))
|
||||
pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
|
||||
EXIT3=$$?
|
||||
pytest -s -v tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
|
||||
EXIT4=$$?
|
||||
exit $$((EXIT1 | EXIT2 | EXIT3 | EXIT4))
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
@@ -52,7 +53,7 @@ steps:
|
||||
- label: ":full_moon: Omni Model Test"
|
||||
timeout_in_minutes: 60
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/examples/online_serving/test_qwen2_5_omni.py -m "advanced_model" --run-level "advanced_model"
|
||||
@@ -71,12 +72,117 @@ steps:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: ":full_moon: Diffusion Model Test with H100"
|
||||
timeout_in_minutes: 60
|
||||
timeout_in_minutes: 120
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and diffusion and H100" --run-level "advanced_model"
|
||||
- pytest -s -v tests/e2e/online_serving/test_*_expansion.py -k "not test_wan22_expansion" -m "advanced_model and diffusion and H100" --run-level "advanced_model"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 2
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: Diffusion Model Wan22 completed Test with H100"
|
||||
timeout_in_minutes: 90
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -m "advanced_model" --run-level "advanced_model"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 2
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: Diffusion Model Test with L4"
|
||||
timeout_in_minutes: 60
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and diffusion and L4" --run-level "advanced_model"
|
||||
agents:
|
||||
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
|
||||
- label: ":full_moon: Documentation Example Code Test with H100"
|
||||
timeout_in_minutes: 60
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY="1"
|
||||
- pytest -s -v tests/examples/online_serving/test_text_to_image.py tests/examples/offline_inference/test_text_to_image.py -m "advanced_model and example and H100" --run-level "advanced_model"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
@@ -114,7 +220,7 @@ steps:
|
||||
- label: ":full_moon: Qwen3-TTS Non-Async-Chunk E2E Test"
|
||||
timeout_in_minutes: 30
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
|
||||
commands:
|
||||
- |
|
||||
huggingface-cli download Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice
|
||||
@@ -135,17 +241,17 @@ steps:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: ":full_moon: Omni Model Perf Test & Test Case Statistics"
|
||||
key: nightly-performance
|
||||
key: nightly-omni-performance
|
||||
timeout_in_minutes: 180
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export BENCHMARK_DIR=tests
|
||||
- pytest -s -v tests/perf/scripts/run_benchmark.py
|
||||
- buildkite-agent artifact upload "tests/*.json"
|
||||
- python tools/nightly/buildkite_testcase_statistics.py -o tests/buildkite_testcase_statistics.html
|
||||
- buildkite-agent artifact upload "tests/*.html"
|
||||
- export BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- pytest -s -v tests/dfx/perf/scripts/run_benchmark.py
|
||||
- buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
|
||||
- python tools/nightly/buildkite_testcase_statistics.py -o tests/dfx/perf/results/buildkite_testcase_statistics.html
|
||||
- buildkite-agent artifact upload "tests/dfx/perf/results/*.html"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
@@ -161,6 +267,93 @@ steps:
|
||||
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: GEBench Accuracy Test with H100"
|
||||
key: nightly-gebench-accuracy
|
||||
timeout_in_minutes: 60
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/e2e/accuracy/test_gebench_h100_smoke.py --run-level advanced_model --gebench-model Qwen/Qwen-Image-2512 --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gebench-port 8093 --accuracy-workers 1
|
||||
- buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gebench_qwen-image-2512/summary*.json"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 1
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: GEdit-Bench Accuracy Test with H100"
|
||||
key: nightly-gedit-bench-accuracy
|
||||
timeout_in_minutes: 60
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/e2e/accuracy/test_gedit_bench_h100_smoke.py --run-level advanced_model --gedit-model Qwen/Qwen-Image-Edit --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gedit-port 8093 --gedit-samples-per-group 20 --accuracy-workers 1
|
||||
- buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gedit_scores_qwen-image-edit/qwen-image-edit_all_all_vie_score_*.csv"
|
||||
- buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gedit_scores_qwen-image-edit/qwen-image-edit_all_all_summary_*.json"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 1
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
@@ -182,19 +375,70 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: Qwen-Image Diffusion Perf Test with H100"
|
||||
key: nightly-qwen-image-performance
|
||||
timeout_in_minutes: 300
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export CACHE_DIT_VERSION=1.3.0
|
||||
- pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
|
||||
- buildkite-agent artifact upload "tests/dfx/perf/results/benchmark_results_*.json"
|
||||
- buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 4
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
# No need to run this step for PRs with label nightly-test
|
||||
- label: ":email: Nightly Collection & Email"
|
||||
key: nightly-perf-distribution
|
||||
depends_on:
|
||||
- nightly-performance
|
||||
- nightly-omni-performance
|
||||
- nightly-qwen-image-performance
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
commands:
|
||||
- pip install openpyxl
|
||||
- export DEFAULT_INPUT_DIR=tests
|
||||
- export DEFAULT_OUTPUT_DIR=tests
|
||||
- buildkite-agent artifact download "tests/*.json" . --step nightly-performance
|
||||
- buildkite-agent artifact download "tests/*.html" . --step nightly-performance
|
||||
- export DEFAULT_INPUT_DIR=tests/dfx/perf/results
|
||||
- export DEFAULT_OUTPUT_DIR=tests/dfx/perf/results
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-qwen-image-performance
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.html" . --step nightly-omni-performance
|
||||
- python tools/nightly/generate_nightly_perf_excel.py
|
||||
- python tools/nightly/send_nightly_email.py --report-file "tests/*.xlsx, tests/*.html"
|
||||
- buildkite-agent artifact upload "tests/*.xlsx"
|
||||
- python tools/nightly/generate_nightly_perf_html.py
|
||||
- python tools/nightly/send_nightly_email.py --report-file "tests/dfx/perf/results/*.xlsx, tests/dfx/perf/results/*.html"
|
||||
- buildkite-agent artifact upload "tests/dfx/perf/results/*.xlsx"
|
||||
- buildkite-agent artifact upload "tests/dfx/perf/results/*.html"
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
+116
-39
@@ -50,6 +50,39 @@ steps:
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Diffusion Batching Test"
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
- timeout 20m pytest -s -v tests/e2e/offline_inference/test_qwen_image_diffusion_batching.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: "Custom Pipeline Test"
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
- timeout 20m pytest -s -v tests/e2e/offline_inference/custom_pipeline/ -m "core_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
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Diffusion Model CPU offloading Test"
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
@@ -168,7 +201,7 @@ steps:
|
||||
timeout 17m bash -c '
|
||||
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/online_serving/test_qwen2_5_omni.py -m "core_model" --run-level "core_model"
|
||||
'
|
||||
agents:
|
||||
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
|
||||
@@ -222,6 +255,49 @@ steps:
|
||||
# path: /mnt/hf-cache
|
||||
# type: DirectoryOrCreate
|
||||
|
||||
- label: "MiMo-Audio E2E Test with H100"
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
- |
|
||||
timeout 30m bash -c '
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -s -v tests/e2e/online_serving/test_mimo_audio.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: 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: "Qwen3-TTS E2E Test"
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
@@ -246,44 +322,45 @@ steps:
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
# TODO: uncomment when mistralai/tts-model is public
|
||||
# - label: "Voxtral-TTS E2E Test"
|
||||
# depends_on: upload-ready-pipeline
|
||||
# commands:
|
||||
# - |
|
||||
# timeout 20m bash -c '
|
||||
# export VLLM_LOGGING_LEVEL=DEBUG
|
||||
# export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
|
||||
# '
|
||||
# agents:
|
||||
# queue: "mithril-h100-pool"
|
||||
# plugins:
|
||||
# - kubernetes:
|
||||
# podSpec:
|
||||
# containers:
|
||||
# - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
# resources:
|
||||
# limits:
|
||||
# nvidia.com/gpu: 1
|
||||
# volumeMounts:
|
||||
# - name: devshm
|
||||
# mountPath: /dev/shm
|
||||
# - name: hf-cache
|
||||
# mountPath: /root/.cache/huggingface
|
||||
# env:
|
||||
# - name: HF_HOME
|
||||
# value: /root/.cache/huggingface
|
||||
# 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: "Voxtral-TTS E2E Test"
|
||||
timeout_in_minutes: 20
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
|
||||
pytest -s -v tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
|
||||
'
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 1
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
# - label: "Diffusion Image Edit Test with H100 (1 GPU)"
|
||||
# depends_on: upload-ready-pipeline
|
||||
|
||||
@@ -39,10 +39,10 @@ jobs:
|
||||
python-version: ["3.11", "3.12"]
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/checkout@v6
|
||||
|
||||
- name: Set up Python ${{ matrix.python-version }}
|
||||
uses: actions/setup-python@v5
|
||||
uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version: ${{ matrix.python-version }}
|
||||
|
||||
@@ -57,7 +57,7 @@ jobs:
|
||||
bash scripts/build_wheel.sh --python python
|
||||
|
||||
- name: Upload Wheel Artifact
|
||||
uses: actions/upload-artifact@v4
|
||||
uses: actions/upload-artifact@v7
|
||||
with:
|
||||
name: vllm-omni-wheel-py${{ matrix.python-version }}
|
||||
path: dist/*.whl
|
||||
@@ -24,12 +24,12 @@ jobs:
|
||||
github.event.comment.author_association == 'OWNER')
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v4.2.2
|
||||
uses: actions/checkout@v6.0.2
|
||||
with:
|
||||
fetch-depth: 0
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v5.3.0
|
||||
uses: actions/setup-python@v6.2.0
|
||||
with:
|
||||
python-version: '3.11'
|
||||
cache: 'pip'
|
||||
@@ -54,7 +54,7 @@ jobs:
|
||||
|
||||
- name: Upload review logs
|
||||
if: always()
|
||||
uses: actions/upload-artifact@v4.4.3
|
||||
uses: actions/upload-artifact@v7.0.0
|
||||
with:
|
||||
name: pr-review-logs-${{ github.event.issue.number || github.event.pull_request.number }}
|
||||
path: /tmp/pr_review_*.log
|
||||
|
||||
@@ -18,8 +18,8 @@ jobs:
|
||||
pre-commit:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/setup-python@v5
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/setup-python@v6
|
||||
with:
|
||||
python-version: "3.12"
|
||||
cache: 'pip'
|
||||
|
||||
@@ -15,6 +15,8 @@ This document explains the benchmark architecture across all benchmark assets in
|
||||
## Directory layout
|
||||
- `benchmarks/build_dataset/` — dataset prep utilities (e.g., SeedTTS top100).
|
||||
- `benchmarks/<model>/vllm_omni/` — vLLM-Omni pipeline benchmarks, logs, outputs.
|
||||
- `benchmarks/accuracy/` — accuracy benchmark integrations that adapt external
|
||||
benchmark suites to vLLM-Omni serving and evaluation flows.
|
||||
- Add new tasks under `benchmarks/<model>/...` with the same pattern: `transformers/`, `vllm_omni/`, task-specific README, and (optionally) dataset prep notes.
|
||||
|
||||
## Reference workflows
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
"""Benchmark helpers and runnable benchmark entrypoints."""
|
||||
@@ -0,0 +1,27 @@
|
||||
# Accuracy Benchmarks
|
||||
|
||||
This directory hosts accuracy benchmark integrations that run entirely through a
|
||||
local `vllm-omni serve` deployment.
|
||||
|
||||
Current integrations:
|
||||
|
||||
- `text_to_image/`: GEBench generation + local judge scoring flow.
|
||||
- `image_to_image/`: GEdit-Bench generation + local VIEScore-style scoring flow.
|
||||
|
||||
Design notes:
|
||||
|
||||
- Generation is executed through the OpenAI-compatible endpoints exposed by
|
||||
`vllm-omni serve`.
|
||||
- Evaluation is also executed through a local OpenAI-compatible judge model
|
||||
served by `vllm-omni`.
|
||||
- Both generation and judge requests accept either `http://host:port` or
|
||||
`http://host:port/v1`.
|
||||
- Output directory layout intentionally stays close to the upstream repos.
|
||||
|
||||
Test guidance:
|
||||
|
||||
- Local static/self-checks live in `tests/benchmarks/test_accuracy_bench_utils.py`.
|
||||
- End-to-end generation/evaluation should be validated in a remote GPU
|
||||
environment. In the current repo marker system there is `L4` but no `L5`
|
||||
marker, so benchmark smoke tests should be wired as `advanced_model +
|
||||
benchmark + L4` when GPU capacity is available.
|
||||
@@ -0,0 +1 @@
|
||||
"""Accuracy benchmark integrations for vLLM-Omni."""
|
||||
@@ -0,0 +1,199 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import io
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import requests
|
||||
from PIL import Image
|
||||
|
||||
|
||||
def ensure_dir(path: Path) -> Path:
|
||||
path.mkdir(parents=True, exist_ok=True)
|
||||
return path
|
||||
|
||||
|
||||
def load_json(path: Path) -> dict[str, Any]:
|
||||
with path.open("r", encoding="utf-8") as handle:
|
||||
return json.load(handle)
|
||||
|
||||
|
||||
def write_json(path: Path, payload: dict[str, Any]) -> None:
|
||||
ensure_dir(path.parent)
|
||||
with path.open("w", encoding="utf-8") as handle:
|
||||
json.dump(payload, handle, indent=2, ensure_ascii=False)
|
||||
|
||||
|
||||
def save_image(path: Path, image: Image.Image) -> None:
|
||||
ensure_dir(path.parent)
|
||||
image.save(path)
|
||||
|
||||
|
||||
def find_first_image(folder: Path, stem: str | None = None) -> Path | None:
|
||||
patterns = [f"{stem}.*"] if stem else ["*.png", "*.jpg", "*.jpeg", "*.webp"]
|
||||
for pattern in patterns:
|
||||
for candidate in sorted(folder.glob(pattern)):
|
||||
if candidate.suffix.lower() in {".png", ".jpg", ".jpeg", ".webp"}:
|
||||
return candidate
|
||||
return None
|
||||
|
||||
|
||||
def extract_json_object(raw_text: str) -> dict[str, Any]:
|
||||
raw_text = raw_text.strip()
|
||||
delimiter = "||V^=^V||"
|
||||
if raw_text.count(delimiter) >= 2:
|
||||
start = raw_text.find(delimiter) + len(delimiter)
|
||||
end = raw_text.rfind(delimiter)
|
||||
raw_text = raw_text[start:end].strip()
|
||||
|
||||
start = raw_text.find("{")
|
||||
end = raw_text.rfind("}")
|
||||
if start == -1 or end == -1 or end < start:
|
||||
raise ValueError(f"Could not find JSON object in: {raw_text[:200]}")
|
||||
return json.loads(raw_text[start : end + 1])
|
||||
|
||||
|
||||
def build_openai_url(base_url: str, api_path: str) -> str:
|
||||
base = base_url.rstrip("/")
|
||||
normalized_path = api_path if api_path.startswith("/") else f"/{api_path}"
|
||||
if base.endswith(normalized_path):
|
||||
return base
|
||||
if base.endswith("/v1"):
|
||||
return f"{base}{normalized_path}"
|
||||
return f"{base}/v1{normalized_path}"
|
||||
|
||||
|
||||
def pil_to_base64(image: Image.Image, image_format: str = "PNG") -> str:
|
||||
buffer = io.BytesIO()
|
||||
image.save(buffer, format=image_format)
|
||||
return base64.b64encode(buffer.getvalue()).decode("utf-8")
|
||||
|
||||
|
||||
def pil_to_data_url(image: Image.Image, image_format: str = "PNG") -> str:
|
||||
return f"data:image/{image_format.lower()};base64,{pil_to_base64(image, image_format=image_format)}"
|
||||
|
||||
|
||||
def decode_base64_image(encoded: str) -> Image.Image:
|
||||
image = Image.open(io.BytesIO(base64.b64decode(encoded)))
|
||||
image.load()
|
||||
return image.convert("RGB")
|
||||
|
||||
|
||||
def pil_to_png_bytes(image: Image.Image) -> bytes:
|
||||
buffer = io.BytesIO()
|
||||
image.save(buffer, format="PNG")
|
||||
return buffer.getvalue()
|
||||
|
||||
|
||||
class VllmOmniImageClient:
|
||||
"""Thin OpenAI-compatible image client for vLLM-Omni serving."""
|
||||
|
||||
def __init__(self, base_url: str, api_key: str = "EMPTY", timeout: int = 600):
|
||||
self.base_url = base_url.rstrip("/")
|
||||
self.api_key = api_key
|
||||
self.timeout = timeout
|
||||
|
||||
@property
|
||||
def _headers(self) -> dict[str, str]:
|
||||
return {
|
||||
"Authorization": f"Bearer {self.api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
def generate_text_to_image(
|
||||
self,
|
||||
*,
|
||||
model: str,
|
||||
prompt: str,
|
||||
width: int,
|
||||
height: int,
|
||||
num_inference_steps: int = 20,
|
||||
guidance_scale: float | None = None,
|
||||
seed: int | None = None,
|
||||
output_compression: int | None = None,
|
||||
) -> Image.Image:
|
||||
payload: dict[str, Any] = {
|
||||
"model": model,
|
||||
"prompt": prompt,
|
||||
"n": 1,
|
||||
"size": f"{width}x{height}",
|
||||
"response_format": "b64_json",
|
||||
"num_inference_steps": num_inference_steps,
|
||||
}
|
||||
if guidance_scale is not None:
|
||||
payload["guidance_scale"] = guidance_scale
|
||||
if seed is not None:
|
||||
payload["seed"] = seed
|
||||
if output_compression is not None:
|
||||
payload["output_compression"] = output_compression
|
||||
|
||||
response = requests.post(
|
||||
build_openai_url(self.base_url, "/images/generations"),
|
||||
json=payload,
|
||||
headers=self._headers,
|
||||
timeout=self.timeout,
|
||||
)
|
||||
response.raise_for_status()
|
||||
return decode_base64_image(response.json()["data"][0]["b64_json"])
|
||||
|
||||
def generate_image_edit(
|
||||
self,
|
||||
*,
|
||||
model: str,
|
||||
prompt: str,
|
||||
images: Image.Image | list[Image.Image],
|
||||
width: int,
|
||||
height: int,
|
||||
num_inference_steps: int = 20,
|
||||
guidance_scale: float | None = None,
|
||||
seed: int | None = None,
|
||||
negative_prompt: str | None = None,
|
||||
output_compression: int | None = None,
|
||||
) -> Image.Image:
|
||||
if not isinstance(images, list):
|
||||
images = [images]
|
||||
data: dict[str, Any] = {
|
||||
"model": model,
|
||||
"prompt": prompt,
|
||||
"n": 1,
|
||||
"size": f"{width}x{height}",
|
||||
"response_format": "b64_json",
|
||||
"num_inference_steps": str(num_inference_steps),
|
||||
}
|
||||
if guidance_scale is not None:
|
||||
data["guidance_scale"] = str(guidance_scale)
|
||||
if seed is not None:
|
||||
data["seed"] = str(seed)
|
||||
if negative_prompt:
|
||||
data["negative_prompt"] = negative_prompt
|
||||
if output_compression is not None:
|
||||
data["output_compression"] = str(output_compression)
|
||||
|
||||
files = [
|
||||
(
|
||||
"image[]" if len(images) > 1 else "image",
|
||||
(f"image_{index}.png", pil_to_png_bytes(image), "image/png"),
|
||||
)
|
||||
for index, image in enumerate(images)
|
||||
]
|
||||
|
||||
edit_paths = ["/images/edits", "/images/edit"]
|
||||
last_response: requests.Response | None = None
|
||||
for api_path in edit_paths:
|
||||
response = requests.post(
|
||||
build_openai_url(self.base_url, api_path),
|
||||
data=data,
|
||||
files=files,
|
||||
headers={"Authorization": f"Bearer {self.api_key}"},
|
||||
timeout=self.timeout,
|
||||
)
|
||||
last_response = response
|
||||
if response.status_code != 404:
|
||||
response.raise_for_status()
|
||||
return decode_base64_image(response.json()["data"][0]["b64_json"])
|
||||
|
||||
assert last_response is not None
|
||||
last_response.raise_for_status()
|
||||
raise ValueError("No image payload returned from image edit endpoint")
|
||||
@@ -0,0 +1,103 @@
|
||||
# GEdit-Bench on vLLM-Omni
|
||||
|
||||
This integration adapts the upstream `stepfun-ai/Step1X-Edit/GEdit-Bench`
|
||||
evaluation flow into `vllm-omni/benchmarks/accuracy/image_to_image`.
|
||||
|
||||
Upstream mapping:
|
||||
|
||||
- `run_gedit_score.py` -> `run_gedit_bench.py evaluate`
|
||||
- `calculate_statistics.py` -> `run_gedit_bench.py summarize`
|
||||
- upstream output layout under `results/<method>/fullset/...` is preserved
|
||||
|
||||
What changed:
|
||||
|
||||
- The upstream repo mainly ships evaluation scripts. This integration adds a
|
||||
generation runner that uses the local `vllm-omni` OpenAI-compatible
|
||||
`/v1/images/edits` endpoint to produce benchmark outputs in the expected
|
||||
directory structure.
|
||||
- The evaluator keeps the same VIEScore-style decomposition:
|
||||
- `semantics_score`
|
||||
- `quality_score`
|
||||
- `overall_score = sqrt(semantics_score * quality_score)`
|
||||
- Judge calls are routed to a local OpenAI-compatible model served by
|
||||
`vllm-omni`, not a remote provider.
|
||||
|
||||
Dataset:
|
||||
|
||||
- Default `--dataset-ref` is `stepfun-ai/GEdit-Bench`
|
||||
- You can also pass a local dataset directory previously saved with
|
||||
Hugging Face `datasets`
|
||||
|
||||
Example usage:
|
||||
|
||||
```bash
|
||||
python benchmarks/accuracy/image_to_image/run_gedit_bench.py generate \
|
||||
--output-root benchmarks/accuracy/image_to_image/results \
|
||||
--base-url http://127.0.0.1:8000 \
|
||||
--model Qwen/Qwen-Image-Edit \
|
||||
--model-name qwen_image_edit \
|
||||
--dataset-ref stepfun-ai/GEdit-Bench \
|
||||
--task-type all \
|
||||
--instruction-language en
|
||||
```
|
||||
|
||||
```bash
|
||||
python benchmarks/accuracy/image_to_image/run_gedit_bench.py evaluate \
|
||||
--output-root benchmarks/accuracy/image_to_image/results \
|
||||
--model-name qwen_image_edit \
|
||||
--save-dir benchmarks/accuracy/image_to_image/scores \
|
||||
--dataset-ref stepfun-ai/GEdit-Bench \
|
||||
--judge-base-url http://127.0.0.1:8000 \
|
||||
--judge-model Qwen/Qwen2.5-VL-7B-Instruct \
|
||||
--judge-api-key EMPTY
|
||||
```
|
||||
|
||||
```bash
|
||||
python benchmarks/accuracy/image_to_image/run_gedit_bench.py summarize \
|
||||
--csv-path benchmarks/accuracy/image_to_image/scores/qwen_image_edit_all_all_vie_score.csv \
|
||||
--language en
|
||||
```
|
||||
|
||||
Example summary output:
|
||||
|
||||
```json
|
||||
{
|
||||
"language": "all",
|
||||
"languages": {
|
||||
"en": {
|
||||
"overall": {"count": 110, "Q_SC": 6.58, "Q_PQ": 5.89, "Q_O": 5.86},
|
||||
"intersection": {"count": 78, "Q_SC": 6.50, "Q_PQ": 5.66, "Q_O": 5.65}
|
||||
},
|
||||
"cn": {
|
||||
"overall": {"count": 110, "Q_SC": 6.90, "Q_PQ": 5.78, "Q_O": 6.11},
|
||||
"intersection": {"count": 63, "Q_SC": 7.22, "Q_PQ": 5.59, "Q_O": 6.28}
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
Example generated images to inspect:
|
||||
|
||||
- `benchmarks/accuracy/image_to_image/results/qwen_image_edit/fullset/background_change/en/<sample>.png`
|
||||
- `benchmarks/accuracy/image_to_image/results/qwen_image_edit/fullset/text_change/en/<sample>.png`
|
||||
- `benchmarks/accuracy/image_to_image/results/qwen_image_edit/fullset/subject-replace/cn/<sample>.png`
|
||||
|
||||
Example score artifacts to inspect together with the images:
|
||||
|
||||
- `benchmarks/accuracy/image_to_image/scores/qwen_image_edit_all_all_vie_score.csv`
|
||||
- `benchmarks/accuracy/image_to_image/scores/qwen_image_edit_all_all_summary.json`
|
||||
|
||||
What to expect:
|
||||
|
||||
- `Q_SC` measures instruction following and content preservation.
|
||||
- `Q_PQ` measures image naturalness and artifact quality.
|
||||
- `Q_O` is the combined overall score; higher is better.
|
||||
- `overall.count` is the number of evaluated samples for that language, while `intersection.count` is the subset with `intersection_exist == True`.
|
||||
|
||||
Notes:
|
||||
|
||||
- This flow requires the optional Hugging Face `datasets` package.
|
||||
- `generate` writes `generation_manifest.json` with local output coverage.
|
||||
- The current repo marker set exposes `L4` but not `L5`, so if you promote an
|
||||
end-to-end smoke test into CI, use the existing `advanced_model`, `benchmark`,
|
||||
and `L4` markers or introduce a new repo-wide marker explicitly first.
|
||||
@@ -0,0 +1 @@
|
||||
"""GEdit-Bench integration for vLLM-Omni."""
|
||||
@@ -0,0 +1,718 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import csv
|
||||
import json
|
||||
import logging
|
||||
import math
|
||||
import statistics
|
||||
from collections import defaultdict
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from datetime import datetime, timezone
|
||||
from io import BytesIO
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import requests
|
||||
from PIL import Image
|
||||
from tqdm.auto import tqdm
|
||||
|
||||
from benchmarks.accuracy.common import (
|
||||
VllmOmniImageClient,
|
||||
build_openai_url,
|
||||
extract_json_object,
|
||||
write_json,
|
||||
)
|
||||
|
||||
GROUPS = [
|
||||
"background_change",
|
||||
"color_alter",
|
||||
"material_alter",
|
||||
"motion_change",
|
||||
"ps_human",
|
||||
"style_change",
|
||||
"subject-add",
|
||||
"subject-remove",
|
||||
"subject-replace",
|
||||
"text_change",
|
||||
"tone_transfer",
|
||||
]
|
||||
DEFAULT_SAMPLES_PER_GROUP = 10
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def infer_model_name(model: str) -> str:
|
||||
normalized = model.rstrip("/\\")
|
||||
name = Path(normalized).name
|
||||
return name or normalized
|
||||
|
||||
|
||||
def resolve_model_name(*, model_name: str | None, model: str | None = None, output_root: Path | None = None) -> str:
|
||||
if model_name:
|
||||
return model_name
|
||||
if model:
|
||||
return infer_model_name(model)
|
||||
if output_root is not None:
|
||||
candidates = sorted(path.name for path in output_root.iterdir() if path.is_dir())
|
||||
if len(candidates) == 1:
|
||||
return candidates[0]
|
||||
if not candidates:
|
||||
raise ValueError(f"Could not infer model-name from empty output root: {output_root}")
|
||||
raise ValueError(
|
||||
f"Could not infer model-name from output root {output_root}; multiple candidates found: {candidates}"
|
||||
)
|
||||
raise ValueError("model-name is required when it cannot be inferred from model or output-root")
|
||||
|
||||
|
||||
def parse_score_payload(raw_text: str) -> dict[str, Any]:
|
||||
try:
|
||||
parsed = extract_json_object(raw_text)
|
||||
except Exception:
|
||||
stripped = raw_text.strip()
|
||||
if stripped.startswith("[") and stripped.endswith("]"):
|
||||
parsed = {"score": json.loads(stripped), "reasoning": ""}
|
||||
else:
|
||||
raise
|
||||
score = parsed.get("score", [])
|
||||
if not isinstance(score, list):
|
||||
score = [score]
|
||||
parsed["score"] = [int(value) for value in score]
|
||||
return parsed
|
||||
|
||||
|
||||
def _utc_timestamp() -> str:
|
||||
return datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
|
||||
|
||||
|
||||
def summarize_generated_records(records: list[dict[str, Any]]) -> dict[str, Any]:
|
||||
by_task: dict[str, list[dict[str, Any]]] = defaultdict(list)
|
||||
by_language: dict[str, list[dict[str, Any]]] = defaultdict(list)
|
||||
for record in records:
|
||||
by_task[record["task_type"]].append(record)
|
||||
by_language[record["instruction_language"]].append(record)
|
||||
|
||||
return {
|
||||
"count": len(records),
|
||||
"by_task": {
|
||||
task_type: {
|
||||
"count": len(rows),
|
||||
"samples": sorted(row["key"] for row in rows),
|
||||
}
|
||||
for task_type, rows in sorted(by_task.items())
|
||||
},
|
||||
"by_language": {
|
||||
language: {
|
||||
"count": len(rows),
|
||||
"samples": sorted(row["key"] for row in rows),
|
||||
}
|
||||
for language, rows in sorted(by_language.items())
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def select_balanced_gedit_rows(
|
||||
rows: list[dict[str, Any]],
|
||||
*,
|
||||
task_type: str = "all",
|
||||
instruction_language: str = "all",
|
||||
samples_per_group: int | None,
|
||||
) -> list[dict[str, Any]]:
|
||||
filtered_rows = []
|
||||
for row in rows:
|
||||
if task_type != "all" and row["task_type"] != task_type:
|
||||
continue
|
||||
if instruction_language != "all" and row["instruction_language"] != instruction_language:
|
||||
continue
|
||||
filtered_rows.append(row)
|
||||
|
||||
if samples_per_group is None:
|
||||
return filtered_rows
|
||||
|
||||
def _select_rows_for_group(group_rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||
if instruction_language != "all":
|
||||
return group_rows[:samples_per_group]
|
||||
|
||||
rows_by_language: dict[str, list[dict[str, Any]]] = defaultdict(list)
|
||||
for row in group_rows:
|
||||
rows_by_language[str(row.get("instruction_language", "")).strip()].append(row)
|
||||
|
||||
ordered_languages = ["en", "cn"]
|
||||
per_language_quota = samples_per_group // max(len(ordered_languages), 1)
|
||||
remainder = samples_per_group % max(len(ordered_languages), 1)
|
||||
|
||||
selected_rows: list[dict[str, Any]] = []
|
||||
leftovers: list[dict[str, Any]] = []
|
||||
for index, language in enumerate(ordered_languages):
|
||||
language_rows = rows_by_language.get(language, [])
|
||||
quota = per_language_quota + (1 if index < remainder else 0)
|
||||
selected_rows.extend(language_rows[:quota])
|
||||
leftovers.extend(language_rows[quota:])
|
||||
|
||||
if len(selected_rows) < samples_per_group:
|
||||
selected_rows.extend(leftovers[: samples_per_group - len(selected_rows)])
|
||||
|
||||
return selected_rows
|
||||
|
||||
if task_type != "all":
|
||||
return _select_rows_for_group(filtered_rows)
|
||||
|
||||
grouped_rows: dict[str, list[dict[str, Any]]] = defaultdict(list)
|
||||
for row in filtered_rows:
|
||||
grouped_rows[row["task_type"]].append(row)
|
||||
|
||||
selected: list[dict[str, Any]] = []
|
||||
for group in GROUPS:
|
||||
selected.extend(_select_rows_for_group(grouped_rows.get(group, [])))
|
||||
return selected
|
||||
|
||||
|
||||
def summarize_gedit_rows(rows: list[dict[str, Any]], language: str = "all") -> dict[str, Any]:
|
||||
return summarize_gedit_rows_with_backbone(rows, language=language)
|
||||
|
||||
|
||||
def _mean_or_none(values: list[float]) -> float | None:
|
||||
return statistics.fmean(values) if values else None
|
||||
|
||||
|
||||
def _to_q_metrics(section: dict[str, float | None]) -> dict[str, float | None]:
|
||||
return {
|
||||
"count": section.get("count"),
|
||||
"Q_SC": section.get("avg_semantics"),
|
||||
"Q_PQ": section.get("avg_quality"),
|
||||
"Q_O": section.get("avg_overall"),
|
||||
}
|
||||
|
||||
|
||||
def _summarize_gedit_rows_single_language(
|
||||
rows: list[dict[str, Any]],
|
||||
*,
|
||||
language: str,
|
||||
) -> dict[str, Any]:
|
||||
filtered_rows = [row for row in rows if str(row.get("instruction_language", "")).strip() == language]
|
||||
|
||||
def _to_bool(value: Any) -> bool:
|
||||
if isinstance(value, bool):
|
||||
return value
|
||||
return str(value).strip().lower() in {"1", "true", "yes"}
|
||||
|
||||
per_group: dict[str, dict[str, float | None]] = {}
|
||||
intersection_group: dict[str, dict[str, float | None]] = {}
|
||||
|
||||
for group in GROUPS:
|
||||
group_rows = [row for row in filtered_rows if row.get("task_type") == group]
|
||||
semantics = [float(row["semantics_score"]) for row in group_rows]
|
||||
quality = [float(row["quality_score"]) for row in group_rows]
|
||||
overall = [math.sqrt(float(row["semantics_score"]) * float(row["quality_score"])) for row in group_rows]
|
||||
|
||||
per_group[group] = {
|
||||
"count": len(group_rows),
|
||||
"avg_semantics": _mean_or_none(semantics),
|
||||
"avg_quality": _mean_or_none(quality),
|
||||
"avg_overall": _mean_or_none(overall),
|
||||
}
|
||||
|
||||
intersection_rows = [row for row in group_rows if _to_bool(row.get("intersection_exist", False))]
|
||||
intersection_semantics = [float(row["semantics_score"]) for row in intersection_rows]
|
||||
intersection_quality = [float(row["quality_score"]) for row in intersection_rows]
|
||||
intersection_overall = [
|
||||
math.sqrt(float(row["semantics_score"]) * float(row["quality_score"])) for row in intersection_rows
|
||||
]
|
||||
intersection_group[group] = {
|
||||
"count": len(intersection_rows),
|
||||
"avg_semantics": _mean_or_none(intersection_semantics),
|
||||
"avg_quality": _mean_or_none(intersection_quality),
|
||||
"avg_overall": _mean_or_none(intersection_overall),
|
||||
}
|
||||
|
||||
overall_section = {
|
||||
"count": len(filtered_rows),
|
||||
"avg_semantics": _mean_or_none(
|
||||
[score["avg_semantics"] for score in per_group.values() if score["avg_semantics"] is not None]
|
||||
),
|
||||
"avg_quality": _mean_or_none(
|
||||
[score["avg_quality"] for score in per_group.values() if score["avg_quality"] is not None]
|
||||
),
|
||||
"avg_overall": _mean_or_none(
|
||||
[score["avg_overall"] for score in per_group.values() if score["avg_overall"] is not None]
|
||||
),
|
||||
}
|
||||
intersection_section = {
|
||||
"count": sum(score["count"] for score in intersection_group.values()),
|
||||
"avg_semantics": _mean_or_none(
|
||||
[score["avg_semantics"] for score in intersection_group.values() if score["avg_semantics"] is not None]
|
||||
),
|
||||
"avg_quality": _mean_or_none(
|
||||
[score["avg_quality"] for score in intersection_group.values() if score["avg_quality"] is not None]
|
||||
),
|
||||
"avg_overall": _mean_or_none(
|
||||
[score["avg_overall"] for score in intersection_group.values() if score["avg_overall"] is not None]
|
||||
),
|
||||
}
|
||||
|
||||
return {
|
||||
"language": language,
|
||||
"by_group": {group: _to_q_metrics(section) for group, section in per_group.items()},
|
||||
"overall": _to_q_metrics(overall_section),
|
||||
"intersection": _to_q_metrics(intersection_section),
|
||||
}
|
||||
|
||||
|
||||
def summarize_gedit_rows_with_backbone(
|
||||
rows: list[dict[str, Any]],
|
||||
*,
|
||||
language: str = "all",
|
||||
) -> dict[str, Any]:
|
||||
if language == "all":
|
||||
return {
|
||||
"language": "all",
|
||||
"languages": {
|
||||
single_language: _summarize_gedit_rows_single_language(
|
||||
rows,
|
||||
language=single_language,
|
||||
)
|
||||
for single_language in ["en", "cn"]
|
||||
},
|
||||
}
|
||||
|
||||
return _summarize_gedit_rows_single_language(rows, language=language)
|
||||
|
||||
|
||||
def _require_datasets():
|
||||
try:
|
||||
from datasets import load_dataset, load_from_disk
|
||||
except ImportError as exc:
|
||||
raise ImportError("GEdit-Bench requires the optional `datasets` package.") from exc
|
||||
return load_dataset, load_from_disk
|
||||
|
||||
|
||||
def _load_gedit_dataset(dataset_ref: str):
|
||||
load_dataset, load_from_disk = _require_datasets()
|
||||
dataset_path = Path(dataset_ref)
|
||||
if dataset_path.exists():
|
||||
if (dataset_path / "state.json").exists() and (dataset_path / "dataset_info.json").exists():
|
||||
return load_from_disk(str(dataset_path))
|
||||
return load_dataset(str(dataset_path))
|
||||
return load_dataset(dataset_ref)
|
||||
|
||||
|
||||
def _resolve_gedit_split(dataset_obj: Any) -> Any:
|
||||
if isinstance(dataset_obj, dict):
|
||||
if "train" in dataset_obj:
|
||||
return dataset_obj["train"]
|
||||
return dataset_obj
|
||||
try:
|
||||
return dataset_obj["train"]
|
||||
except Exception:
|
||||
return dataset_obj
|
||||
|
||||
|
||||
def _to_pil_image(value: Any) -> Image.Image:
|
||||
if isinstance(value, Image.Image):
|
||||
return value.convert("RGB")
|
||||
if isinstance(value, dict) and "bytes" in value:
|
||||
image = Image.open(BytesIO(value["bytes"]))
|
||||
image.load()
|
||||
return image.convert("RGB")
|
||||
raise TypeError(f"Unsupported image payload type: {type(value)!r}")
|
||||
|
||||
|
||||
class LocalVIEScorer:
|
||||
def __init__(self, *, base_url: str, api_key: str, model: str, timeout: int = 600):
|
||||
self.base_url = base_url.rstrip("/")
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
self.timeout = timeout
|
||||
self.sc_prompt = (
|
||||
"You are evaluating image editing quality.\n"
|
||||
"Two images are provided: the source image and the edited image.\n"
|
||||
'Return JSON only in the format {"score": [edit_success, content_preservation], '
|
||||
'"reasoning": "..."}.\n'
|
||||
"Each score must be an integer from 0 to 10.\n"
|
||||
"Editing instruction: <instruction>"
|
||||
)
|
||||
self.pq_prompt = (
|
||||
"You are evaluating image quality.\n"
|
||||
'Return JSON only in the format {"score": [naturalness, artifact_free], "reasoning": "..."}.\n'
|
||||
"Each score must be an integer from 0 to 10."
|
||||
)
|
||||
|
||||
def _request(self, prompt: str, images: list[Image.Image]) -> dict[str, Any]:
|
||||
from benchmarks.accuracy.common import pil_to_data_url
|
||||
|
||||
content: list[dict[str, Any]] = [{"type": "text", "text": prompt}]
|
||||
for image in images:
|
||||
content.append({"type": "image_url", "image_url": {"url": pil_to_data_url(image)}})
|
||||
|
||||
response = requests.post(
|
||||
build_openai_url(self.base_url, "/chat/completions"),
|
||||
json={
|
||||
"model": self.model,
|
||||
"messages": [{"role": "user", "content": content}],
|
||||
"temperature": 0,
|
||||
},
|
||||
headers={
|
||||
"Authorization": f"Bearer {self.api_key}",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
timeout=self.timeout,
|
||||
)
|
||||
response.raise_for_status()
|
||||
message_content = response.json()["choices"][0]["message"]["content"]
|
||||
if isinstance(message_content, list):
|
||||
text = "\n".join(part.get("text", "") for part in message_content if part.get("type") == "text")
|
||||
else:
|
||||
text = str(message_content)
|
||||
return parse_score_payload(text)
|
||||
|
||||
def evaluate(self, source_image: Image.Image, edited_image: Image.Image, instruction: str) -> dict[str, float]:
|
||||
sc_payload = self._request(self.sc_prompt.replace("<instruction>", instruction), [source_image, edited_image])
|
||||
pq_payload = self._request(self.pq_prompt, [edited_image])
|
||||
semantics = float(min(sc_payload["score"])) if sc_payload["score"] else 0.0
|
||||
quality = float(min(pq_payload["score"])) if pq_payload["score"] else 0.0
|
||||
overall = math.sqrt(semantics * quality)
|
||||
return {
|
||||
"semantics_score": semantics,
|
||||
"quality_score": quality,
|
||||
"overall_score": overall,
|
||||
}
|
||||
|
||||
|
||||
class GEditBenchRunner:
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
dataset_ref: str,
|
||||
output_root: Path,
|
||||
base_url: str,
|
||||
model: str,
|
||||
api_key: str = "EMPTY",
|
||||
width: int = 512,
|
||||
height: int = 512,
|
||||
num_inference_steps: int = 20,
|
||||
guidance_scale: float | None = None,
|
||||
seed: int | None = 42,
|
||||
):
|
||||
self.dataset_ref = dataset_ref
|
||||
self.output_root = output_root
|
||||
self.model = model
|
||||
self.width = width
|
||||
self.height = height
|
||||
self.num_inference_steps = num_inference_steps
|
||||
self.guidance_scale = guidance_scale
|
||||
self.seed = seed
|
||||
self.client = VllmOmniImageClient(base_url=base_url, api_key=api_key)
|
||||
|
||||
def generate(
|
||||
self,
|
||||
*,
|
||||
model_name: str,
|
||||
task_type: str = "all",
|
||||
instruction_language: str = "all",
|
||||
workers: int = 1,
|
||||
max_samples: int | None = None,
|
||||
samples_per_group: int | None = None,
|
||||
) -> list[dict[str, Any]]:
|
||||
dataset = _resolve_gedit_split(_load_gedit_dataset(self.dataset_ref))
|
||||
rows = select_balanced_gedit_rows(
|
||||
list(dataset),
|
||||
task_type=task_type,
|
||||
instruction_language=instruction_language,
|
||||
samples_per_group=samples_per_group,
|
||||
)
|
||||
if max_samples is not None:
|
||||
rows = rows[:max_samples]
|
||||
|
||||
outputs: list[dict[str, Any]] = []
|
||||
total = len(rows)
|
||||
if workers <= 1:
|
||||
with tqdm(total=total, desc="GEdit generate", unit="sample") as progress:
|
||||
for item in rows:
|
||||
result = self._safe_generate_one(model_name, item)
|
||||
if result:
|
||||
outputs.append(result)
|
||||
progress.update(1)
|
||||
return outputs
|
||||
|
||||
with tqdm(total=total, desc="GEdit generate", unit="sample") as progress:
|
||||
with ThreadPoolExecutor(max_workers=workers) as executor:
|
||||
futures = [executor.submit(self._safe_generate_one, model_name, item) for item in rows]
|
||||
for future in as_completed(futures):
|
||||
result = future.result()
|
||||
if result:
|
||||
outputs.append(result)
|
||||
progress.update(1)
|
||||
return outputs
|
||||
|
||||
def _safe_generate_one(self, model_name: str, item: dict[str, Any]) -> dict[str, Any] | None:
|
||||
try:
|
||||
return self._generate_one(model_name, item)
|
||||
except Exception:
|
||||
logger.exception("Failed to generate GEdit-Bench sample %s", item.get("key", "<unknown>"))
|
||||
return None
|
||||
|
||||
def _generate_one(self, model_name: str, item: dict[str, Any]) -> dict[str, Any] | None:
|
||||
output_path = (
|
||||
self.output_root
|
||||
/ model_name
|
||||
/ "fullset"
|
||||
/ item["task_type"]
|
||||
/ item["instruction_language"]
|
||||
/ f"{item['key']}.png"
|
||||
)
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
if output_path.exists():
|
||||
return {
|
||||
"task_type": item["task_type"],
|
||||
"instruction_language": item["instruction_language"],
|
||||
"key": item["key"],
|
||||
"output_path": str(output_path),
|
||||
}
|
||||
|
||||
source_image = _to_pil_image(item["input_image_raw"])
|
||||
edited_image = self.client.generate_image_edit(
|
||||
model=self.model,
|
||||
prompt=item["instruction"],
|
||||
images=source_image,
|
||||
width=self.width,
|
||||
height=self.height,
|
||||
num_inference_steps=self.num_inference_steps,
|
||||
guidance_scale=self.guidance_scale,
|
||||
seed=self.seed,
|
||||
)
|
||||
edited_image.save(output_path)
|
||||
return {
|
||||
"task_type": item["task_type"],
|
||||
"instruction_language": item["instruction_language"],
|
||||
"key": item["key"],
|
||||
"output_path": str(output_path),
|
||||
}
|
||||
|
||||
|
||||
class GEditBenchEvaluator:
|
||||
def __init__(self, *, dataset_ref: str, output_root: Path, scorer: LocalVIEScorer):
|
||||
self.dataset_ref = dataset_ref
|
||||
self.output_root = output_root
|
||||
self.scorer = scorer
|
||||
|
||||
def evaluate(
|
||||
self,
|
||||
*,
|
||||
model_name: str,
|
||||
save_dir: Path,
|
||||
task_type: str = "all",
|
||||
instruction_language: str = "all",
|
||||
workers: int = 1,
|
||||
max_samples: int | None = None,
|
||||
samples_per_group: int | None = None,
|
||||
) -> dict[str, Any]:
|
||||
dataset = _resolve_gedit_split(_load_gedit_dataset(self.dataset_ref))
|
||||
rows = select_balanced_gedit_rows(
|
||||
list(dataset),
|
||||
task_type=task_type,
|
||||
instruction_language=instruction_language,
|
||||
samples_per_group=samples_per_group,
|
||||
)
|
||||
if max_samples is not None:
|
||||
rows = rows[:max_samples]
|
||||
|
||||
results: list[dict[str, Any]] = []
|
||||
total = len(rows)
|
||||
if workers <= 1:
|
||||
with tqdm(total=total, desc="GEdit evaluate", unit="sample") as progress:
|
||||
for item in rows:
|
||||
result = self._safe_evaluate_one(model_name, item)
|
||||
if result:
|
||||
results.append(result)
|
||||
progress.update(1)
|
||||
else:
|
||||
with tqdm(total=total, desc="GEdit evaluate", unit="sample") as progress:
|
||||
with ThreadPoolExecutor(max_workers=workers) as executor:
|
||||
futures = [executor.submit(self._safe_evaluate_one, model_name, item) for item in rows]
|
||||
for future in as_completed(futures):
|
||||
result = future.result()
|
||||
if result:
|
||||
results.append(result)
|
||||
progress.update(1)
|
||||
|
||||
save_dir.mkdir(parents=True, exist_ok=True)
|
||||
base_name = f"{model_name}_{task_type}_{instruction_language}"
|
||||
timestamp = _utc_timestamp()
|
||||
csv_path = save_dir / f"{base_name}_vie_score.csv"
|
||||
timestamped_csv_path = save_dir / f"{base_name}_vie_score_{timestamp}.csv"
|
||||
|
||||
def _write_csv(path: Path) -> None:
|
||||
with path.open("w", encoding="utf-8", newline="") as handle:
|
||||
writer = csv.DictWriter(
|
||||
handle,
|
||||
fieldnames=[
|
||||
"key",
|
||||
"task_type",
|
||||
"edited_image",
|
||||
"instruction",
|
||||
"semantics_score",
|
||||
"quality_score",
|
||||
"overall_score",
|
||||
"intersection_exist",
|
||||
"instruction_language",
|
||||
],
|
||||
)
|
||||
writer.writeheader()
|
||||
for row in results:
|
||||
writer.writerow(row)
|
||||
|
||||
_write_csv(csv_path)
|
||||
_write_csv(timestamped_csv_path)
|
||||
|
||||
summary = summarize_gedit_rows_with_backbone(
|
||||
results,
|
||||
language=instruction_language,
|
||||
)
|
||||
summary_path = save_dir / f"{base_name}_summary.json"
|
||||
timestamped_summary_path = save_dir / f"{base_name}_summary_{timestamp}.json"
|
||||
write_json(summary_path, summary)
|
||||
write_json(timestamped_summary_path, summary)
|
||||
return {
|
||||
"results": results,
|
||||
"summary": summary,
|
||||
"csv_path": str(csv_path),
|
||||
"summary_path": str(summary_path),
|
||||
"timestamped_csv_path": str(timestamped_csv_path),
|
||||
"timestamped_summary_path": str(timestamped_summary_path),
|
||||
}
|
||||
|
||||
def _safe_evaluate_one(self, model_name: str, item: dict[str, Any]) -> dict[str, Any] | None:
|
||||
try:
|
||||
return self._evaluate_one(model_name, item)
|
||||
except Exception:
|
||||
logger.exception("Failed to evaluate GEdit-Bench sample %s", item.get("key", "<unknown>"))
|
||||
return None
|
||||
|
||||
def _evaluate_one(self, model_name: str, item: dict[str, Any]) -> dict[str, Any] | None:
|
||||
edited_image_path = (
|
||||
self.output_root
|
||||
/ model_name
|
||||
/ "fullset"
|
||||
/ item["task_type"]
|
||||
/ item["instruction_language"]
|
||||
/ f"{item['key']}.png"
|
||||
)
|
||||
if not edited_image_path.exists():
|
||||
return None
|
||||
|
||||
source_image = _to_pil_image(item["input_image_raw"])
|
||||
edited_image = Image.open(edited_image_path).convert("RGB")
|
||||
scores = self.scorer.evaluate(source_image, edited_image, item["instruction"])
|
||||
return {
|
||||
"key": item["key"],
|
||||
"task_type": item["task_type"],
|
||||
"edited_image": str(edited_image_path),
|
||||
"instruction": item["instruction"],
|
||||
"semantics_score": scores["semantics_score"],
|
||||
"quality_score": scores["quality_score"],
|
||||
"overall_score": scores["overall_score"],
|
||||
"intersection_exist": item.get("Intersection_exist", False),
|
||||
"instruction_language": item["instruction_language"],
|
||||
}
|
||||
|
||||
|
||||
def build_parser() -> argparse.ArgumentParser:
|
||||
parser = argparse.ArgumentParser(description="Run the GEdit-Bench integration against a local vLLM-Omni server.")
|
||||
subparsers = parser.add_subparsers(dest="command", required=True)
|
||||
|
||||
generate = subparsers.add_parser("generate")
|
||||
generate.add_argument("--dataset-ref", type=str, default="stepfun-ai/GEdit-Bench")
|
||||
generate.add_argument("--output-root", type=Path, required=True)
|
||||
generate.add_argument("--base-url", type=str, required=True)
|
||||
generate.add_argument("--model", type=str, required=True)
|
||||
generate.add_argument("--model-name", type=str, default=None)
|
||||
generate.add_argument("--api-key", type=str, default="EMPTY")
|
||||
generate.add_argument("--task-type", choices=["all", *GROUPS], default="all")
|
||||
generate.add_argument("--instruction-language", choices=["all", "en", "cn"], default="all")
|
||||
generate.add_argument("--width", type=int, default=512)
|
||||
generate.add_argument("--height", type=int, default=512)
|
||||
generate.add_argument("--num-inference-steps", type=int, default=20)
|
||||
generate.add_argument("--guidance-scale", type=float, default=None)
|
||||
generate.add_argument("--seed", type=int, default=42)
|
||||
generate.add_argument("--workers", type=int, default=1)
|
||||
generate.add_argument("--max-samples", type=int, default=None)
|
||||
generate.add_argument("--samples-per-group", type=int, default=None)
|
||||
|
||||
evaluate = subparsers.add_parser("evaluate")
|
||||
evaluate.add_argument("--dataset-ref", type=str, default="stepfun-ai/GEdit-Bench")
|
||||
evaluate.add_argument("--output-root", type=Path, required=True)
|
||||
evaluate.add_argument("--model-name", type=str, default=None)
|
||||
evaluate.add_argument("--save-dir", type=Path, required=True)
|
||||
evaluate.add_argument("--task-type", choices=["all", *GROUPS], default="all")
|
||||
evaluate.add_argument("--instruction-language", choices=["all", "en", "cn"], default="all")
|
||||
evaluate.add_argument("--judge-base-url", type=str, required=True)
|
||||
evaluate.add_argument("--judge-model", type=str, required=True)
|
||||
evaluate.add_argument("--judge-api-key", type=str, default="EMPTY")
|
||||
evaluate.add_argument("--workers", type=int, default=1)
|
||||
evaluate.add_argument("--max-samples", type=int, default=None)
|
||||
evaluate.add_argument("--samples-per-group", type=int, default=None)
|
||||
|
||||
summarize = subparsers.add_parser("summarize")
|
||||
summarize.add_argument("--csv-path", type=Path, required=True)
|
||||
summarize.add_argument("--language", choices=["all", "en", "cn"], default="all")
|
||||
|
||||
return parser
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
parser = build_parser()
|
||||
args = parser.parse_args(argv)
|
||||
|
||||
if args.command == "generate":
|
||||
model_name = resolve_model_name(model_name=args.model_name, model=args.model)
|
||||
runner = GEditBenchRunner(
|
||||
dataset_ref=args.dataset_ref,
|
||||
output_root=args.output_root,
|
||||
base_url=args.base_url,
|
||||
model=args.model,
|
||||
api_key=args.api_key,
|
||||
width=args.width,
|
||||
height=args.height,
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
guidance_scale=args.guidance_scale,
|
||||
seed=args.seed,
|
||||
)
|
||||
records = runner.generate(
|
||||
model_name=model_name,
|
||||
task_type=args.task_type,
|
||||
instruction_language=args.instruction_language,
|
||||
workers=args.workers,
|
||||
max_samples=args.max_samples,
|
||||
samples_per_group=args.samples_per_group,
|
||||
)
|
||||
payload = {"records": records, "summary": summarize_generated_records(records)}
|
||||
write_json(args.output_root / model_name / "generation_manifest.json", payload)
|
||||
return 0
|
||||
|
||||
if args.command == "evaluate":
|
||||
model_name = resolve_model_name(model_name=args.model_name, output_root=args.output_root)
|
||||
scorer = LocalVIEScorer(
|
||||
base_url=args.judge_base_url,
|
||||
api_key=args.judge_api_key,
|
||||
model=args.judge_model,
|
||||
)
|
||||
evaluator = GEditBenchEvaluator(dataset_ref=args.dataset_ref, output_root=args.output_root, scorer=scorer)
|
||||
evaluator.evaluate(
|
||||
model_name=model_name,
|
||||
save_dir=args.save_dir,
|
||||
task_type=args.task_type,
|
||||
instruction_language=args.instruction_language,
|
||||
workers=args.workers,
|
||||
max_samples=args.max_samples,
|
||||
samples_per_group=args.samples_per_group,
|
||||
)
|
||||
return 0
|
||||
|
||||
if args.command == "summarize":
|
||||
with args.csv_path.open("r", encoding="utf-8", newline="") as handle:
|
||||
rows = list(csv.DictReader(handle))
|
||||
summary = summarize_gedit_rows_with_backbone(rows, language=args.language)
|
||||
print(json.dumps(summary, indent=2, ensure_ascii=False))
|
||||
return 0
|
||||
|
||||
parser.error(f"Unknown command: {args.command}")
|
||||
return 1
|
||||
@@ -0,0 +1,13 @@
|
||||
# ruff: noqa: E402, I001
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[3]
|
||||
if str(REPO_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
from benchmarks.accuracy.image_to_image.gedit_bench import main
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -0,0 +1,103 @@
|
||||
# GEBench on vLLM-Omni
|
||||
|
||||
This integration adapts the upstream `stepfun-ai/GEBench` scripts into
|
||||
`vllm-omni/benchmarks/accuracy/text_to_image`.
|
||||
|
||||
Upstream mapping:
|
||||
|
||||
- `scripts/generate.py` -> `run_gebench.py generate`
|
||||
- `scripts/evaluate.py` -> `run_gebench.py evaluate`
|
||||
- upstream prompt / judge logic -> `gbench.py`
|
||||
|
||||
What changed:
|
||||
|
||||
- Generation calls the local OpenAI-compatible `vllm-omni` endpoints:
|
||||
- `/v1/images/generations` for text-only frame generation
|
||||
- `/v1/images/edits` for image-conditioned GUI transition generation
|
||||
- Evaluation still keeps the GEBench scoring dimensions:
|
||||
- `goal`
|
||||
- `logic`
|
||||
- `cons`
|
||||
- `ui`
|
||||
- `qual`
|
||||
- Judge calls are also routed to a local OpenAI-compatible model served by
|
||||
`vllm-omni` instead of a remote service.
|
||||
|
||||
Expected dataset layout:
|
||||
|
||||
- Clone the benchmark dataset from Hugging Face into a local directory:
|
||||
|
||||
```bash
|
||||
git clone https://huggingface.co/datasets/stepfun-ai/GEBench /path/to/GEBench
|
||||
```
|
||||
|
||||
Example usage:
|
||||
|
||||
```bash
|
||||
python benchmarks/accuracy/text_to_image/run_gebench.py generate \
|
||||
--dataset-root /path/to/GEBench \
|
||||
--output-root benchmarks/accuracy/text_to_image/outputs \
|
||||
--base-url http://127.0.0.1:8000 \
|
||||
--model Tongyi-MAI/Z-Image-Turbo \
|
||||
--data-type type3
|
||||
```
|
||||
|
||||
```bash
|
||||
python benchmarks/accuracy/text_to_image/run_gebench.py evaluate \
|
||||
--dataset-root /path/to/GEBench \
|
||||
--output-root benchmarks/accuracy/text_to_image/outputs \
|
||||
--data-type type3 \
|
||||
--judge-base-url http://127.0.0.1:8000 \
|
||||
--judge-model Qwen/Qwen2.5-VL-7B-Instruct \
|
||||
--judge-api-key EMPTY
|
||||
```
|
||||
|
||||
```bash
|
||||
python benchmarks/accuracy/text_to_image/run_gebench.py summarize \
|
||||
--output-root benchmarks/accuracy/text_to_image/outputs
|
||||
```
|
||||
|
||||
Example summary output:
|
||||
|
||||
```json
|
||||
{
|
||||
"generation": {
|
||||
"count": 20,
|
||||
"by_type": {
|
||||
"type3": {"count": 10},
|
||||
"type4": {"count": 10}
|
||||
}
|
||||
},
|
||||
"evaluation": {
|
||||
"count": 20,
|
||||
"overall_mean": 0.52,
|
||||
"by_type": {
|
||||
"type3": {"count": 10, "overall_mean": 0.50, "overall_mean_100": 50.0},
|
||||
"type4": {"count": 10, "overall_mean": 0.54, "overall_mean_100": 54.0}
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
Example generated images to inspect:
|
||||
|
||||
- `benchmarks/accuracy/text_to_image/outputs/03_trajectory_text_fictionalapp/<lang_device>/<sample>/frame0.png`
|
||||
- `benchmarks/accuracy/text_to_image/outputs/03_trajectory_text_fictionalapp/<lang_device>/<sample>/frame5.png`
|
||||
- `benchmarks/accuracy/text_to_image/outputs/04_trajectory_text_realapp/<lang_device>/<sample>/frame0.png`
|
||||
- `benchmarks/accuracy/text_to_image/outputs/04_trajectory_text_realapp/<lang_device>/<sample>/frame5.png`
|
||||
|
||||
What to expect:
|
||||
|
||||
- `overall_mean` is normalized to `0.0 ~ 1.0`; higher is better.
|
||||
- `frame0.png` is the initial GUI frame; `frame5.png` is the final trajectory frame most often used for quick inspection.
|
||||
- For full debugging, inspect the whole `frame0.png` ... `frame5.png` sequence for one sample directory.
|
||||
|
||||
Notes:
|
||||
|
||||
- GEBench upstream leaves type3/type4 generation unfinished. This integration
|
||||
fills that gap with a trajectory runner that generates `frame0.png` followed
|
||||
by `frame1.png` ... `frame5.png`.
|
||||
- Type1/2/5 require an image-edit capable model exposed through
|
||||
`vllm-omni serve`.
|
||||
- `summarize` will report both generated coverage and any existing evaluation
|
||||
summary files.
|
||||
@@ -0,0 +1 @@
|
||||
"""GEBench integration for vLLM-Omni."""
|
||||
@@ -0,0 +1,927 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import statistics
|
||||
from collections import defaultdict
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import requests
|
||||
from PIL import Image
|
||||
|
||||
from benchmarks.accuracy.common import (
|
||||
VllmOmniImageClient,
|
||||
build_openai_url,
|
||||
ensure_dir,
|
||||
extract_json_object,
|
||||
find_first_image,
|
||||
load_json,
|
||||
pil_to_data_url,
|
||||
save_image,
|
||||
write_json,
|
||||
)
|
||||
|
||||
TYPE_TO_FOLDER = {
|
||||
"type1": "01_single_step",
|
||||
"type2": "02_multi_step",
|
||||
"type3": "03_trajectory_text_fictionalapp",
|
||||
"type4": "04_trajectory_text_realapp",
|
||||
"type5": "05_grounding_data",
|
||||
}
|
||||
SCORE_KEYS = ("goal", "logic", "cons", "ui", "qual")
|
||||
DEFAULT_SAMPLES_PER_TYPE = 10
|
||||
|
||||
|
||||
def _utc_timestamp() -> str:
|
||||
return datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
|
||||
|
||||
|
||||
def _write_json_with_timestamp(path: Path, payload: dict[str, Any]) -> Path:
|
||||
write_json(path, payload)
|
||||
timestamped_path = path.with_name(f"{path.stem}_{_utc_timestamp()}{path.suffix}")
|
||||
write_json(timestamped_path, payload)
|
||||
return timestamped_path
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class GEBenchSampleSpec:
|
||||
sample_path: Path
|
||||
metadata: dict[str, Any]
|
||||
sample_name: str
|
||||
lang_device: str
|
||||
|
||||
|
||||
def summarize_generated_records(records: list[dict[str, Any]]) -> dict[str, Any]:
|
||||
by_type: dict[str, list[dict[str, Any]]] = defaultdict(list)
|
||||
for record in records:
|
||||
by_type[record["data_type"]].append(record)
|
||||
|
||||
return {
|
||||
"count": len(records),
|
||||
"by_type": {
|
||||
data_type: {
|
||||
"count": len(rows),
|
||||
}
|
||||
for data_type, rows in sorted(by_type.items())
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def summarize_gebench_results(results: list[dict[str, Any]]) -> dict[str, Any]:
|
||||
by_type: dict[str, list[dict[str, Any]]] = defaultdict(list)
|
||||
for result in results:
|
||||
by_type[result["data_type"]].append(result)
|
||||
|
||||
summary: dict[str, Any] = {
|
||||
"count": len(results),
|
||||
"overall_mean": statistics.fmean(r["overall"] for r in results) if results else 0.0,
|
||||
"by_type": {},
|
||||
}
|
||||
for data_type, rows in by_type.items():
|
||||
score_means: dict[str, float] = {}
|
||||
all_score_keys = {key for row in rows for key in row.get("scores", {}).keys()}
|
||||
for score_key in all_score_keys:
|
||||
values = [row["scores"][score_key] for row in rows if score_key in row.get("scores", {})]
|
||||
score_means[score_key] = statistics.fmean(values) if values else 0.0
|
||||
overall_mean = statistics.fmean(row["overall"] for row in rows)
|
||||
summary["by_type"][data_type] = {
|
||||
"count": len(rows),
|
||||
"overall_mean": overall_mean,
|
||||
"overall_mean_100": overall_mean * 100.0,
|
||||
"score_means": score_means,
|
||||
}
|
||||
return summary
|
||||
|
||||
|
||||
def select_balanced_gebench_samples(
|
||||
sample_paths_by_type: dict[str, list[Any]],
|
||||
*,
|
||||
samples_per_type: int | None,
|
||||
) -> dict[str, list[Any]]:
|
||||
if samples_per_type is None:
|
||||
return {data_type: list(paths) for data_type, paths in sample_paths_by_type.items()}
|
||||
return {data_type: list(paths)[:samples_per_type] for data_type, paths in sample_paths_by_type.items()}
|
||||
|
||||
|
||||
def collect_gebench_generation_summary(output_root: Path) -> dict[str, Any]:
|
||||
records: list[dict[str, Any]] = []
|
||||
for data_type, folder_name in TYPE_TO_FOLDER.items():
|
||||
type_root = output_root / folder_name
|
||||
if not type_root.exists():
|
||||
continue
|
||||
for lang_dir in sorted(path for path in type_root.iterdir() if path.is_dir()):
|
||||
for sample_dir in sorted(path for path in lang_dir.iterdir() if path.is_dir()):
|
||||
expected = sample_dir / "frame5.png" if data_type in {"type2", "type3", "type4"} else None
|
||||
if expected is None:
|
||||
expected = find_first_image(sample_dir)
|
||||
elif not expected.exists():
|
||||
expected = None
|
||||
if expected is None:
|
||||
continue
|
||||
records.append(
|
||||
{
|
||||
"data_type": data_type,
|
||||
"sample_name": f"{lang_dir.name}/{sample_dir.name}",
|
||||
"output_path": str(expected),
|
||||
}
|
||||
)
|
||||
return summarize_generated_records(records)
|
||||
|
||||
|
||||
def _normalize_score_key(key: str) -> str:
|
||||
mapping = {
|
||||
"goal": "goal",
|
||||
"logic": "logic",
|
||||
"cons": "cons",
|
||||
"consistency": "cons",
|
||||
"ui": "ui",
|
||||
"qual": "qual",
|
||||
"quality": "qual",
|
||||
}
|
||||
return mapping.get(key.lower(), key.lower())
|
||||
|
||||
|
||||
def _normalize_scores(raw_scores: dict[str, Any]) -> dict[str, int]:
|
||||
scores: dict[str, int] = {}
|
||||
for key, value in raw_scores.items():
|
||||
normalized = _normalize_score_key(key)
|
||||
if normalized not in SCORE_KEYS:
|
||||
continue
|
||||
scalar = value.get("s", 0) if isinstance(value, dict) else value
|
||||
try:
|
||||
scores[normalized] = int(scalar)
|
||||
except (TypeError, ValueError):
|
||||
scores[normalized] = 0
|
||||
for key in SCORE_KEYS:
|
||||
scores.setdefault(key, 0)
|
||||
return scores
|
||||
|
||||
|
||||
def _compute_overall(scores: dict[str, int]) -> float:
|
||||
return sum(scores.values()) / (len(SCORE_KEYS) * 5.0)
|
||||
|
||||
|
||||
def _iter_sample_paths(dataset_root: Path, data_type: str) -> list[Path]:
|
||||
data_dir = dataset_root / TYPE_TO_FOLDER[data_type]
|
||||
if not data_dir.exists():
|
||||
raise FileNotFoundError(f"Dataset path not found: {data_dir}")
|
||||
|
||||
samples: list[Path] = []
|
||||
for lang_dir in sorted(path for path in data_dir.iterdir() if path.is_dir()):
|
||||
for child in sorted(lang_dir.iterdir()):
|
||||
if child.is_dir():
|
||||
samples.append(child)
|
||||
elif child.suffix.lower() == ".json":
|
||||
samples.append(child)
|
||||
return samples
|
||||
|
||||
|
||||
def _sample_name_from_metadata(metadata: dict[str, Any], sample_path: Path, item_index: int | None = None) -> str:
|
||||
sample_id = metadata.get("id") or metadata.get("sample_id") or metadata.get("name")
|
||||
if sample_id:
|
||||
return str(sample_id)
|
||||
if item_index is not None:
|
||||
return f"{sample_path.stem}_{item_index:04d}"
|
||||
return sample_path.stem if sample_path.is_file() else sample_path.name
|
||||
|
||||
|
||||
def _expand_sample_path(sample_path: Path) -> list[GEBenchSampleSpec]:
|
||||
if sample_path.is_dir():
|
||||
metadata = _load_metadata(sample_path)
|
||||
return [
|
||||
GEBenchSampleSpec(
|
||||
sample_path=sample_path,
|
||||
metadata=metadata,
|
||||
sample_name=_sample_name(sample_path),
|
||||
lang_device=_lang_device(sample_path, metadata),
|
||||
)
|
||||
]
|
||||
|
||||
payload = load_json(sample_path)
|
||||
if isinstance(payload, dict):
|
||||
return [
|
||||
GEBenchSampleSpec(
|
||||
sample_path=sample_path,
|
||||
metadata=payload,
|
||||
sample_name=_sample_name_from_metadata(payload, sample_path),
|
||||
lang_device=_lang_device(sample_path, payload),
|
||||
)
|
||||
]
|
||||
|
||||
if isinstance(payload, list):
|
||||
specs: list[GEBenchSampleSpec] = []
|
||||
for item_index, item in enumerate(payload):
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
specs.append(
|
||||
GEBenchSampleSpec(
|
||||
sample_path=sample_path,
|
||||
metadata=item,
|
||||
sample_name=_sample_name_from_metadata(item, sample_path, item_index),
|
||||
lang_device=_lang_device(sample_path, item),
|
||||
)
|
||||
)
|
||||
return specs
|
||||
|
||||
raise TypeError(f"Unsupported metadata payload type for sample {sample_path}: {type(payload)!r}")
|
||||
|
||||
|
||||
def _iter_sample_specs(dataset_root: Path, data_type: str) -> list[GEBenchSampleSpec]:
|
||||
sample_specs: list[GEBenchSampleSpec] = []
|
||||
for sample_path in _iter_sample_paths(dataset_root, data_type):
|
||||
sample_specs.extend(_expand_sample_path(sample_path))
|
||||
return sample_specs
|
||||
|
||||
|
||||
def _load_metadata(sample_path: Path) -> dict[str, Any]:
|
||||
if sample_path.is_file():
|
||||
return load_json(sample_path)
|
||||
for candidate in ("meta_data.json", "metadata.json"):
|
||||
meta_path = sample_path / candidate
|
||||
if meta_path.exists():
|
||||
return load_json(meta_path)
|
||||
raise FileNotFoundError(f"Metadata not found for sample: {sample_path}")
|
||||
|
||||
|
||||
def _sample_name(sample_path: Path) -> str:
|
||||
return sample_path.stem if sample_path.is_file() else sample_path.name
|
||||
|
||||
|
||||
def _lang_device(sample_path: Path, metadata: dict[str, Any]) -> str:
|
||||
return str(metadata.get("lang_device") or sample_path.parent.name)
|
||||
|
||||
|
||||
def _resolve_referenced_image(
|
||||
*,
|
||||
metadata: dict[str, Any],
|
||||
sample_path: Path,
|
||||
dataset_root: Path,
|
||||
data_type: str,
|
||||
) -> Image.Image | None:
|
||||
for key in ("image", "input_image", "initial_image", "reference_image"):
|
||||
image_ref = metadata.get(key)
|
||||
if not image_ref:
|
||||
continue
|
||||
candidate = dataset_root / TYPE_TO_FOLDER[data_type] / str(image_ref)
|
||||
if candidate.exists():
|
||||
image = Image.open(candidate)
|
||||
image.load()
|
||||
return image.convert("RGB")
|
||||
if sample_path.is_dir():
|
||||
local_image = find_first_image(sample_path)
|
||||
if local_image:
|
||||
image = Image.open(local_image)
|
||||
image.load()
|
||||
return image.convert("RGB")
|
||||
return None
|
||||
|
||||
|
||||
def _trajectory_steps(metadata: dict[str, Any]) -> list[dict[str, Any]]:
|
||||
for key in ("trajectory", "steps", "frames"):
|
||||
value = metadata.get(key)
|
||||
if isinstance(value, list):
|
||||
return [step for step in value if isinstance(step, dict)]
|
||||
extracted: list[dict[str, Any]] = []
|
||||
for index in range(1, 6):
|
||||
value = metadata.get(f"step{index}") or metadata.get(str(index))
|
||||
if isinstance(value, dict):
|
||||
extracted.append(value)
|
||||
return extracted
|
||||
|
||||
|
||||
def _text_or_default(value: Any, default: str = "") -> str:
|
||||
return str(value).strip() if value is not None else default
|
||||
|
||||
|
||||
def _type1_prompt(metadata: dict[str, Any]) -> str:
|
||||
caption = _text_or_default(metadata.get("caption") or metadata.get("instruction"), "Transform the reference GUI.")
|
||||
return (
|
||||
"Using the reference GUI screenshot, generate the next GUI state after the requested interaction.\n\n"
|
||||
f"Requested change:\n{caption}\n\n"
|
||||
"Requirements:\n"
|
||||
"- Preserve layout, visual identity, and unrelated regions.\n"
|
||||
"- Only apply the requested state change.\n"
|
||||
"- Keep all text and controls readable.\n"
|
||||
)
|
||||
|
||||
|
||||
def _type2_prompt(goal: str, step_num: int) -> str:
|
||||
return (
|
||||
"Generate the next GUI state for a multi-step task.\n\n"
|
||||
f"Overall goal: {goal}\n"
|
||||
f"Current progress step: {step_num}/5\n\n"
|
||||
"Requirements:\n"
|
||||
"- The change should be incremental and plausible.\n"
|
||||
"- Preserve layout and visual identity.\n"
|
||||
"- Make text/buttons readable.\n"
|
||||
)
|
||||
|
||||
|
||||
def _type34_initial_prompt(metadata: dict[str, Any], first_step: dict[str, Any]) -> str:
|
||||
app_name = _text_or_default(metadata.get("app_name"), "App")
|
||||
final_goal = _text_or_default(metadata.get("final_goal") or metadata.get("instruction"), "Complete the task.")
|
||||
visual_description = _text_or_default(
|
||||
metadata.get("visual_description") or first_step.get("visual_description") or first_step.get("description"),
|
||||
"A clean product-quality app home screen.",
|
||||
)
|
||||
return (
|
||||
"Generate the first GUI frame for a task trajectory.\n\n"
|
||||
f"App name: {app_name}\n"
|
||||
f"Final goal: {final_goal}\n"
|
||||
f"Visual description:\n{visual_description}\n\n"
|
||||
"Requirements:\n"
|
||||
"- Generate a production-looking UI screenshot only.\n"
|
||||
"- Keep the layout coherent and readable.\n"
|
||||
)
|
||||
|
||||
|
||||
def _type34_next_prompt(step_num: int, step_info: dict[str, Any]) -> str:
|
||||
action = _text_or_default(step_info.get("action") or step_info.get("instruction"), "Continue the task.")
|
||||
visual_description = _text_or_default(
|
||||
step_info.get("visual_description") or step_info.get("description"),
|
||||
"Reflect the expected next GUI state.",
|
||||
)
|
||||
return (
|
||||
"Using the previous frame as reference, generate the next GUI frame.\n\n"
|
||||
f"Step {step_num} action: {action}\n"
|
||||
f"Expected visual state:\n{visual_description}\n\n"
|
||||
"Requirements:\n"
|
||||
"- Only change UI regions affected by this action.\n"
|
||||
"- Preserve persistent bars, layout, and style.\n"
|
||||
"- Keep text and icons readable.\n"
|
||||
)
|
||||
|
||||
|
||||
def _type5_prompt(metadata: dict[str, Any]) -> str:
|
||||
grounding = metadata.get("grounding") or {}
|
||||
explanation = _text_or_default(
|
||||
metadata.get("grounding_explanation") or grounding.get("effect") or grounding.get("description"),
|
||||
"Predict the immediate GUI reaction to the indicated target.",
|
||||
)
|
||||
return (
|
||||
"Using the reference GUI screenshot, predict the immediate GUI state after the grounded interaction.\n\n"
|
||||
f"Expected effect: {explanation}\n"
|
||||
f"Grounding metadata: {json.dumps(grounding, ensure_ascii=False)}\n\n"
|
||||
"Requirements:\n"
|
||||
"- Apply only the interaction-triggered change.\n"
|
||||
"- Preserve unrelated regions.\n"
|
||||
"- Keep the UI realistic and readable.\n"
|
||||
)
|
||||
|
||||
|
||||
def _make_storyboard_image(
|
||||
frames: list[Image.Image],
|
||||
*,
|
||||
columns: int = 3,
|
||||
background_color: tuple[int, int, int] = (255, 255, 255),
|
||||
) -> Image.Image:
|
||||
if not frames:
|
||||
raise ValueError("Expected at least one frame to build a storyboard image.")
|
||||
|
||||
normalized = [frame.convert("RGB") for frame in frames]
|
||||
frame_width = max(frame.width for frame in normalized)
|
||||
frame_height = max(frame.height for frame in normalized)
|
||||
rows = (len(normalized) + columns - 1) // columns
|
||||
storyboard = Image.new("RGB", (frame_width * columns, frame_height * rows), color=background_color)
|
||||
|
||||
for index, frame in enumerate(normalized):
|
||||
x_offset = (index % columns) * frame_width
|
||||
y_offset = (index // columns) * frame_height
|
||||
if frame.size != (frame_width, frame_height):
|
||||
frame = frame.resize((frame_width, frame_height))
|
||||
storyboard.paste(frame, (x_offset, y_offset))
|
||||
return storyboard
|
||||
|
||||
|
||||
def _trajectory_judge_payload(frames: list[Image.Image]) -> tuple[str, list[Image.Image]]:
|
||||
storyboard = _make_storyboard_image(frames, columns=3)
|
||||
prompt_suffix = (
|
||||
"The attached image is a storyboard containing six frames arranged left-to-right, "
|
||||
"top-to-bottom as frame0, frame1, frame2, frame3, frame4, frame5."
|
||||
)
|
||||
return prompt_suffix, [storyboard]
|
||||
|
||||
|
||||
class LocalJudgeClient:
|
||||
def __init__(self, base_url: str, api_key: str, model: str, timeout: int = 600):
|
||||
self.base_url = base_url.rstrip("/")
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
self.timeout = timeout
|
||||
|
||||
def _build_scoring_prompt(self, task_prompt: str) -> str:
|
||||
return (
|
||||
"You are an expert evaluator for GUI image editing and GUI trajectory generation.\n"
|
||||
"Evaluate whether the generated image(s) satisfy the task.\n\n"
|
||||
"Score these five dimensions from 0 to 5:\n"
|
||||
"- goal: whether the user goal is completed correctly\n"
|
||||
"- logic: whether the transition/state change is logically correct\n"
|
||||
"- cons: whether unrelated regions remain consistent\n"
|
||||
"- ui: whether the UI layout/components remain realistic and coherent\n"
|
||||
"- qual: whether the images are visually clear and artifact-free\n\n"
|
||||
"Return JSON only. Do not add any prose outside JSON.\n"
|
||||
"Use exactly this schema:\n"
|
||||
"{\n"
|
||||
' "goal": 0,\n'
|
||||
' "logic": 0,\n'
|
||||
' "cons": 0,\n'
|
||||
' "ui": 0,\n'
|
||||
' "qual": 0,\n'
|
||||
' "reasoning": "short explanation"\n'
|
||||
"}\n\n"
|
||||
"Scoring task:\n"
|
||||
f"{task_prompt}"
|
||||
)
|
||||
|
||||
def _request_text(self, prompt: str, images: list[Image.Image]) -> str:
|
||||
content: list[dict[str, Any]] = [{"type": "text", "text": prompt}]
|
||||
for image in images:
|
||||
content.append({"type": "image_url", "image_url": {"url": pil_to_data_url(image)}})
|
||||
|
||||
response = requests.post(
|
||||
build_openai_url(self.base_url, "/chat/completions"),
|
||||
json={
|
||||
"model": self.model,
|
||||
"messages": [{"role": "user", "content": content}],
|
||||
"temperature": 0,
|
||||
},
|
||||
headers={
|
||||
"Authorization": f"Bearer {self.api_key}",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
timeout=self.timeout,
|
||||
)
|
||||
response.raise_for_status()
|
||||
message_content = response.json()["choices"][0]["message"]["content"]
|
||||
if isinstance(message_content, list):
|
||||
return "\n".join(part.get("text", "") for part in message_content if part.get("type") == "text")
|
||||
return str(message_content)
|
||||
|
||||
def evaluate(self, *, prompt: str, images: list[Image.Image]) -> dict[str, Any]:
|
||||
primary_prompt = self._build_scoring_prompt(prompt)
|
||||
text = self._request_text(primary_prompt, images)
|
||||
try:
|
||||
return extract_json_object(text)
|
||||
except ValueError:
|
||||
retry_prompt = (
|
||||
self._build_scoring_prompt(prompt) + "\n\nYour previous response was not valid JSON. "
|
||||
"Return only the JSON object with integer scores."
|
||||
)
|
||||
retry_text = self._request_text(retry_prompt, images)
|
||||
try:
|
||||
return extract_json_object(retry_text)
|
||||
except ValueError:
|
||||
return {
|
||||
"goal": 0,
|
||||
"logic": 0,
|
||||
"cons": 0,
|
||||
"ui": 0,
|
||||
"qual": 0,
|
||||
"reasoning": retry_text.strip() or text.strip() or "Judge response was not valid JSON.",
|
||||
}
|
||||
|
||||
|
||||
class GEBenchRunner:
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
dataset_root: Path,
|
||||
output_root: Path,
|
||||
base_url: str,
|
||||
model: str,
|
||||
api_key: str = "EMPTY",
|
||||
width: int = 768,
|
||||
height: int = 576,
|
||||
num_inference_steps: int = 8,
|
||||
output_compression: int | None = 98,
|
||||
guidance_scale: float | None = None,
|
||||
seed: int | None = 42,
|
||||
):
|
||||
self.dataset_root = dataset_root
|
||||
self.output_root = output_root
|
||||
self.model = model
|
||||
self.width = width
|
||||
self.height = height
|
||||
self.num_inference_steps = num_inference_steps
|
||||
self.output_compression = output_compression
|
||||
self.guidance_scale = guidance_scale
|
||||
self.seed = seed
|
||||
self.client = VllmOmniImageClient(base_url=base_url, api_key=api_key)
|
||||
|
||||
def generate(
|
||||
self,
|
||||
*,
|
||||
data_type: str,
|
||||
workers: int = 1,
|
||||
max_samples: int | None = None,
|
||||
samples_per_type: int | None = None,
|
||||
) -> list[dict[str, Any]]:
|
||||
sample_specs = select_balanced_gebench_samples(
|
||||
{data_type: _iter_sample_specs(self.dataset_root, data_type)},
|
||||
samples_per_type=samples_per_type,
|
||||
)[data_type]
|
||||
if max_samples is not None:
|
||||
sample_specs = sample_specs[:max_samples]
|
||||
|
||||
results: list[dict[str, Any]] = []
|
||||
if workers <= 1:
|
||||
for sample_spec in sample_specs:
|
||||
result = self._generate_one(data_type, sample_spec)
|
||||
if result:
|
||||
results.append(result)
|
||||
return results
|
||||
|
||||
with ThreadPoolExecutor(max_workers=workers) as executor:
|
||||
futures = [executor.submit(self._generate_one, data_type, sample_spec) for sample_spec in sample_specs]
|
||||
for future in as_completed(futures):
|
||||
result = future.result()
|
||||
if result:
|
||||
results.append(result)
|
||||
return results
|
||||
|
||||
def _generate_one(self, data_type: str, sample_spec: GEBenchSampleSpec) -> dict[str, Any] | None:
|
||||
sample_path = sample_spec.sample_path
|
||||
metadata = sample_spec.metadata
|
||||
lang_device = sample_spec.lang_device
|
||||
sample_name = sample_spec.sample_name
|
||||
output_dir = ensure_dir(self.output_root / TYPE_TO_FOLDER[data_type] / lang_device / sample_name)
|
||||
|
||||
if data_type == "type1":
|
||||
output_path = output_dir / "generated.png"
|
||||
if output_path.exists():
|
||||
return {
|
||||
"data_type": data_type,
|
||||
"sample_name": f"{lang_device}/{sample_name}",
|
||||
"output_path": str(output_path),
|
||||
}
|
||||
source = _resolve_referenced_image(
|
||||
metadata=metadata, sample_path=sample_path, dataset_root=self.dataset_root, data_type=data_type
|
||||
)
|
||||
if source is None:
|
||||
return None
|
||||
generated = self.client.generate_image_edit(
|
||||
model=self.model,
|
||||
prompt=_type1_prompt(metadata),
|
||||
images=source,
|
||||
width=self.width,
|
||||
height=self.height,
|
||||
num_inference_steps=self.num_inference_steps,
|
||||
output_compression=self.output_compression,
|
||||
guidance_scale=self.guidance_scale,
|
||||
seed=self.seed,
|
||||
)
|
||||
save_image(output_path, generated)
|
||||
return {
|
||||
"data_type": data_type,
|
||||
"sample_name": f"{lang_device}/{sample_name}",
|
||||
"output_path": str(output_path),
|
||||
}
|
||||
|
||||
if data_type == "type2":
|
||||
goal = _text_or_default(metadata.get("question") or metadata.get("caption"), "Complete the task.")
|
||||
source = _resolve_referenced_image(
|
||||
metadata=metadata, sample_path=sample_path, dataset_root=self.dataset_root, data_type=data_type
|
||||
)
|
||||
if source is None:
|
||||
return None
|
||||
frame0_path = output_dir / "frame0.png"
|
||||
if not frame0_path.exists():
|
||||
save_image(frame0_path, source)
|
||||
previous = source
|
||||
for step_num in range(1, 6):
|
||||
frame_path = output_dir / f"frame{step_num}.png"
|
||||
if frame_path.exists():
|
||||
previous = Image.open(frame_path).convert("RGB")
|
||||
continue
|
||||
generated = self.client.generate_image_edit(
|
||||
model=self.model,
|
||||
prompt=_type2_prompt(goal, step_num),
|
||||
images=previous,
|
||||
width=self.width,
|
||||
height=self.height,
|
||||
num_inference_steps=self.num_inference_steps,
|
||||
output_compression=self.output_compression,
|
||||
guidance_scale=self.guidance_scale,
|
||||
seed=self.seed,
|
||||
)
|
||||
save_image(frame_path, generated)
|
||||
previous = generated
|
||||
output_path = output_dir / "frame5.png"
|
||||
return {
|
||||
"data_type": data_type,
|
||||
"sample_name": f"{lang_device}/{sample_name}",
|
||||
"output_path": str(output_path),
|
||||
}
|
||||
|
||||
if data_type in {"type3", "type4"}:
|
||||
steps = _trajectory_steps(metadata)
|
||||
frame0_path = output_dir / "frame0.png"
|
||||
if frame0_path.exists():
|
||||
previous = Image.open(frame0_path).convert("RGB")
|
||||
else:
|
||||
previous = self.client.generate_text_to_image(
|
||||
model=self.model,
|
||||
prompt=_type34_initial_prompt(metadata, steps[0] if steps else {}),
|
||||
width=self.width,
|
||||
height=self.height,
|
||||
num_inference_steps=self.num_inference_steps,
|
||||
output_compression=self.output_compression,
|
||||
guidance_scale=self.guidance_scale,
|
||||
seed=self.seed,
|
||||
)
|
||||
save_image(frame0_path, previous)
|
||||
|
||||
for step_num in range(1, 6):
|
||||
frame_path = output_dir / f"frame{step_num}.png"
|
||||
if frame_path.exists():
|
||||
previous = Image.open(frame_path).convert("RGB")
|
||||
continue
|
||||
step_info = steps[step_num - 1] if step_num - 1 < len(steps) else {}
|
||||
generated = self.client.generate_image_edit(
|
||||
model=self.model,
|
||||
prompt=_type34_next_prompt(step_num, step_info),
|
||||
images=previous,
|
||||
width=self.width,
|
||||
height=self.height,
|
||||
num_inference_steps=self.num_inference_steps,
|
||||
output_compression=self.output_compression,
|
||||
guidance_scale=self.guidance_scale,
|
||||
seed=self.seed,
|
||||
)
|
||||
save_image(frame_path, generated)
|
||||
previous = generated
|
||||
output_path = output_dir / "frame5.png"
|
||||
return {
|
||||
"data_type": data_type,
|
||||
"sample_name": f"{lang_device}/{sample_name}",
|
||||
"output_path": str(output_path),
|
||||
}
|
||||
|
||||
if data_type == "type5":
|
||||
output_path = output_dir / "generated.png"
|
||||
if output_path.exists():
|
||||
return {
|
||||
"data_type": data_type,
|
||||
"sample_name": f"{lang_device}/{sample_name}",
|
||||
"output_path": str(output_path),
|
||||
}
|
||||
source = _resolve_referenced_image(
|
||||
metadata=metadata, sample_path=sample_path, dataset_root=self.dataset_root, data_type=data_type
|
||||
)
|
||||
if source is None:
|
||||
return None
|
||||
generated = self.client.generate_image_edit(
|
||||
model=self.model,
|
||||
prompt=_type5_prompt(metadata),
|
||||
images=source,
|
||||
width=self.width,
|
||||
height=self.height,
|
||||
num_inference_steps=self.num_inference_steps,
|
||||
output_compression=self.output_compression,
|
||||
guidance_scale=self.guidance_scale,
|
||||
seed=self.seed,
|
||||
)
|
||||
save_image(output_path, generated)
|
||||
return {
|
||||
"data_type": data_type,
|
||||
"sample_name": f"{lang_device}/{sample_name}",
|
||||
"output_path": str(output_path),
|
||||
}
|
||||
|
||||
raise ValueError(f"Unsupported data type: {data_type}")
|
||||
|
||||
|
||||
class GEBenchEvaluator:
|
||||
def __init__(self, *, dataset_root: Path, output_root: Path, judge: LocalJudgeClient):
|
||||
self.dataset_root = dataset_root
|
||||
self.output_root = output_root
|
||||
self.judge = judge
|
||||
|
||||
def evaluate(
|
||||
self,
|
||||
*,
|
||||
data_type: str,
|
||||
workers: int = 1,
|
||||
max_samples: int | None = None,
|
||||
samples_per_type: int | None = None,
|
||||
) -> dict[str, Any]:
|
||||
output_type_dir = self.output_root / TYPE_TO_FOLDER[data_type]
|
||||
sample_specs_by_name = {
|
||||
(spec.lang_device, spec.sample_name): spec for spec in _iter_sample_specs(self.dataset_root, data_type)
|
||||
}
|
||||
if not output_type_dir.exists():
|
||||
payload = {"data_type": data_type, "results": [], "summary": summarize_gebench_results([])}
|
||||
write_json(self.output_root / "evaluations" / f"{data_type}.json", payload)
|
||||
return payload
|
||||
sample_dirs = [
|
||||
sample_dir
|
||||
for lang_dir in sorted(path for path in output_type_dir.iterdir() if path.is_dir())
|
||||
for sample_dir in sorted(path for path in lang_dir.iterdir() if path.is_dir())
|
||||
if (lang_dir.name, sample_dir.name) in sample_specs_by_name
|
||||
]
|
||||
sample_dirs = select_balanced_gebench_samples(
|
||||
{data_type: sample_dirs},
|
||||
samples_per_type=samples_per_type,
|
||||
)[data_type]
|
||||
if max_samples is not None:
|
||||
sample_dirs = sample_dirs[:max_samples]
|
||||
results: list[dict[str, Any]] = []
|
||||
if workers <= 1:
|
||||
for sample_dir in sample_dirs:
|
||||
result = self._evaluate_one(
|
||||
data_type,
|
||||
sample_dir,
|
||||
sample_specs_by_name[(sample_dir.parent.name, sample_dir.name)],
|
||||
)
|
||||
if result:
|
||||
results.append(result)
|
||||
else:
|
||||
with ThreadPoolExecutor(max_workers=workers) as executor:
|
||||
futures = [
|
||||
executor.submit(
|
||||
self._evaluate_one,
|
||||
data_type,
|
||||
sample_dir,
|
||||
sample_specs_by_name[(sample_dir.parent.name, sample_dir.name)],
|
||||
)
|
||||
for sample_dir in sample_dirs
|
||||
]
|
||||
for future in as_completed(futures):
|
||||
result = future.result()
|
||||
if result:
|
||||
results.append(result)
|
||||
|
||||
payload = {"data_type": data_type, "results": results, "summary": summarize_gebench_results(results)}
|
||||
write_json(self.output_root / "evaluations" / f"{data_type}.json", payload)
|
||||
return payload
|
||||
|
||||
def _evaluate_one(self, data_type: str, sample_dir: Path, sample_spec: GEBenchSampleSpec) -> dict[str, Any] | None:
|
||||
lang_device = sample_dir.parent.name
|
||||
sample_name = sample_dir.name
|
||||
dataset_sample = sample_spec.sample_path
|
||||
metadata = sample_spec.metadata
|
||||
|
||||
if data_type == "type1":
|
||||
source = _resolve_referenced_image(
|
||||
metadata=metadata, sample_path=dataset_sample, dataset_root=self.dataset_root, data_type=data_type
|
||||
)
|
||||
generated_path = find_first_image(sample_dir)
|
||||
if source is None or generated_path is None:
|
||||
return None
|
||||
generated = Image.open(generated_path).convert("RGB")
|
||||
raw_scores = self.judge.evaluate(prompt=_type1_prompt(metadata), images=[source, generated])
|
||||
elif data_type == "type2":
|
||||
frames = [Image.open(sample_dir / f"frame{i}.png").convert("RGB") for i in range(6)]
|
||||
goal = _text_or_default(metadata.get("question") or metadata.get("caption"), "Complete the task.")
|
||||
prompt_suffix, judge_images = _trajectory_judge_payload(frames)
|
||||
raw_scores = self.judge.evaluate(
|
||||
prompt=f"Evaluate a six-frame GUI trajectory.\nTask: {goal}\n{prompt_suffix}",
|
||||
images=judge_images,
|
||||
)
|
||||
elif data_type in {"type3", "type4"}:
|
||||
frames = [Image.open(sample_dir / f"frame{i}.png").convert("RGB") for i in range(6)]
|
||||
instruction = _text_or_default(metadata.get("instruction") or metadata.get("caption"), "Complete the task.")
|
||||
prompt_suffix, judge_images = _trajectory_judge_payload(frames)
|
||||
raw_scores = self.judge.evaluate(
|
||||
prompt=f"Evaluate a six-frame GUI trajectory.\nInstruction: {instruction}\n{prompt_suffix}",
|
||||
images=judge_images,
|
||||
)
|
||||
elif data_type == "type5":
|
||||
source = _resolve_referenced_image(
|
||||
metadata=metadata, sample_path=dataset_sample, dataset_root=self.dataset_root, data_type=data_type
|
||||
)
|
||||
generated_path = find_first_image(sample_dir)
|
||||
if source is None or generated_path is None:
|
||||
return None
|
||||
generated = Image.open(generated_path).convert("RGB")
|
||||
raw_scores = self.judge.evaluate(prompt=_type5_prompt(metadata), images=[source, generated])
|
||||
else:
|
||||
raise ValueError(f"Unsupported data type: {data_type}")
|
||||
|
||||
scores = _normalize_scores(raw_scores)
|
||||
return {
|
||||
"sample_name": f"{lang_device}/{sample_name}",
|
||||
"data_type": data_type,
|
||||
"scores": scores,
|
||||
"overall": _compute_overall(scores),
|
||||
"raw_scores": raw_scores,
|
||||
}
|
||||
|
||||
|
||||
def _data_types_arg(value: str) -> list[str]:
|
||||
return list(TYPE_TO_FOLDER.keys()) if value == "all" else [value]
|
||||
|
||||
|
||||
def build_parser() -> argparse.ArgumentParser:
|
||||
parser = argparse.ArgumentParser(description="Run local GEBench generation and scoring against vLLM-Omni.")
|
||||
subparsers = parser.add_subparsers(dest="command", required=True)
|
||||
|
||||
generate = subparsers.add_parser("generate")
|
||||
generate.add_argument("--dataset-root", type=Path, required=True)
|
||||
generate.add_argument("--output-root", type=Path, required=True)
|
||||
generate.add_argument("--base-url", type=str, required=True)
|
||||
generate.add_argument("--model", type=str, required=True)
|
||||
generate.add_argument("--data-type", choices=["all", *TYPE_TO_FOLDER.keys()], default="all")
|
||||
generate.add_argument("--api-key", type=str, default="EMPTY")
|
||||
generate.add_argument("--width", type=int, default=768)
|
||||
generate.add_argument("--height", type=int, default=576)
|
||||
generate.add_argument("--num-inference-steps", type=int, default=8)
|
||||
generate.add_argument("--output-compression", type=int, default=98)
|
||||
generate.add_argument("--guidance-scale", type=float, default=None)
|
||||
generate.add_argument("--seed", type=int, default=42)
|
||||
generate.add_argument("--workers", type=int, default=1)
|
||||
generate.add_argument("--max-samples", type=int, default=None)
|
||||
generate.add_argument("--samples-per-type", type=int, default=None)
|
||||
|
||||
evaluate = subparsers.add_parser("evaluate")
|
||||
evaluate.add_argument("--dataset-root", type=Path, required=True)
|
||||
evaluate.add_argument("--output-root", type=Path, required=True)
|
||||
evaluate.add_argument("--data-type", choices=["all", *TYPE_TO_FOLDER.keys()], default="all")
|
||||
evaluate.add_argument("--judge-base-url", type=str, required=True)
|
||||
evaluate.add_argument("--judge-model", type=str, required=True)
|
||||
evaluate.add_argument("--judge-api-key", type=str, default="EMPTY")
|
||||
evaluate.add_argument("--workers", type=int, default=1)
|
||||
evaluate.add_argument("--max-samples", type=int, default=None)
|
||||
evaluate.add_argument("--samples-per-type", type=int, default=None)
|
||||
|
||||
summarize = subparsers.add_parser("summarize")
|
||||
summarize.add_argument("--output-root", type=Path, required=True)
|
||||
|
||||
return parser
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
parser = build_parser()
|
||||
args = parser.parse_args(argv)
|
||||
|
||||
if args.command == "generate":
|
||||
runner = GEBenchRunner(
|
||||
dataset_root=args.dataset_root,
|
||||
output_root=args.output_root,
|
||||
base_url=args.base_url,
|
||||
model=args.model,
|
||||
api_key=args.api_key,
|
||||
width=args.width,
|
||||
height=args.height,
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
output_compression=args.output_compression,
|
||||
guidance_scale=args.guidance_scale,
|
||||
seed=args.seed,
|
||||
)
|
||||
records: list[dict[str, Any]] = []
|
||||
for data_type in _data_types_arg(args.data_type):
|
||||
records.extend(
|
||||
runner.generate(
|
||||
data_type=data_type,
|
||||
workers=args.workers,
|
||||
max_samples=args.max_samples,
|
||||
samples_per_type=args.samples_per_type,
|
||||
)
|
||||
)
|
||||
payload = {"records": records, "summary": summarize_generated_records(records)}
|
||||
write_json(args.output_root / "generation_manifest.json", payload)
|
||||
return 0
|
||||
|
||||
if args.command == "evaluate":
|
||||
judge = LocalJudgeClient(
|
||||
base_url=args.judge_base_url,
|
||||
api_key=args.judge_api_key,
|
||||
model=args.judge_model,
|
||||
)
|
||||
evaluator = GEBenchEvaluator(dataset_root=args.dataset_root, output_root=args.output_root, judge=judge)
|
||||
combined_results: list[dict[str, Any]] = []
|
||||
for data_type in _data_types_arg(args.data_type):
|
||||
payload = evaluator.evaluate(
|
||||
data_type=data_type,
|
||||
workers=args.workers,
|
||||
max_samples=args.max_samples,
|
||||
samples_per_type=args.samples_per_type,
|
||||
)
|
||||
combined_results.extend(payload["results"])
|
||||
_write_json_with_timestamp(
|
||||
args.output_root / "evaluations" / "summary.json",
|
||||
{"summary": summarize_gebench_results(combined_results)},
|
||||
)
|
||||
return 0
|
||||
|
||||
if args.command == "summarize":
|
||||
generation_summary = collect_gebench_generation_summary(args.output_root)
|
||||
evaluation_dir = args.output_root / "evaluations"
|
||||
result_records: list[dict[str, Any]] = []
|
||||
if evaluation_dir.exists():
|
||||
for file_path in sorted(evaluation_dir.glob("type*.json")):
|
||||
payload = load_json(file_path)
|
||||
result_records.extend(payload.get("results", []))
|
||||
payload: dict[str, Any] = {"generation": generation_summary}
|
||||
if result_records:
|
||||
payload["evaluation"] = summarize_gebench_results(result_records)
|
||||
_write_json_with_timestamp(args.output_root / "summary.json", payload)
|
||||
print(json.dumps(payload, indent=2, ensure_ascii=False))
|
||||
return 0
|
||||
|
||||
parser.error(f"Unknown command: {args.command}")
|
||||
return 1
|
||||
@@ -0,0 +1,13 @@
|
||||
# ruff: noqa: E402, I001
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[3]
|
||||
if str(REPO_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
from benchmarks.accuracy.text_to_image.gbench import main
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -109,8 +109,6 @@ async def async_request_chat_completions(
|
||||
resp_json = await response.json()
|
||||
output.response_body = resp_json
|
||||
output.success = True
|
||||
if "peak_memory_mb" in resp_json:
|
||||
output.peak_memory_mb = resp_json["peak_memory_mb"]
|
||||
try:
|
||||
choices = resp_json.get("choices", [])
|
||||
if choices and isinstance(choices, list):
|
||||
@@ -120,7 +118,8 @@ async def async_request_chat_completions(
|
||||
if content and isinstance(content, list) and len(content) > 0:
|
||||
first_item = content[0]
|
||||
if isinstance(first_item, dict):
|
||||
output.stage_durations = first_item.get("stage_durations")
|
||||
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:
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
# adapted from sglang and fastvideo
|
||||
# adapted from fastvideo
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
@@ -289,6 +289,9 @@ class VBenchDataset(BaseDataset):
|
||||
|
||||
def _resize_data(self, data: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||
"""Resize data to match num_prompts."""
|
||||
if not data:
|
||||
raise ValueError("No benchmark data available. Install Pillow or provide --dataset-path.")
|
||||
|
||||
if not self.args.num_prompts:
|
||||
return data
|
||||
|
||||
@@ -512,6 +515,14 @@ class TraceDataset(BaseDataset):
|
||||
timestamp = self._coerce_float(row.get("timestamp"))
|
||||
slo_ms = self._coerce_float(row.get("slo_ms"))
|
||||
image_paths = row.get("image_paths")
|
||||
if not image_paths:
|
||||
single = row.get("image_path")
|
||||
image_paths = [single] if single else None
|
||||
|
||||
if not image_paths and self.args.task in ["i2v", "i2i", "ti2v", "ti2i"]:
|
||||
raise ValueError(
|
||||
f"Task {self.args.task} requires image input, but no image_path or image_paths found in trace row."
|
||||
)
|
||||
|
||||
override_w = self.args.width
|
||||
override_h = self.args.height
|
||||
@@ -714,19 +725,19 @@ async def iter_requests(
|
||||
requests_list: list[RequestFuncInput],
|
||||
request_rate: float,
|
||||
) -> AsyncGenerator[RequestFuncInput, None]:
|
||||
"""Yield requests using a fixed interval if request_rate is set.
|
||||
"""Yield requests using a Poisson process if request_rate is set.
|
||||
|
||||
- If request_rate is inf, all requests are yielded immediately (no sleep).
|
||||
- Otherwise, requests are emitted at a fixed cadence of 1 / request_rate seconds.
|
||||
- Otherwise, inter-arrival times follow an exponential distribution.
|
||||
"""
|
||||
|
||||
if request_rate != float("inf"):
|
||||
if request_rate <= 0:
|
||||
raise ValueError(f"request_rate must be positive or inf, got {request_rate}.")
|
||||
interval_s = 1.0 / float(request_rate)
|
||||
|
||||
for i, req in enumerate(requests_list):
|
||||
if request_rate != float("inf") and i > 0:
|
||||
interval_s = random.expovariate(request_rate)
|
||||
await asyncio.sleep(interval_s)
|
||||
yield req
|
||||
|
||||
@@ -745,6 +756,15 @@ def calculate_metrics(
|
||||
latencies = [o.latency for o in success_outputs]
|
||||
peak_memories = [o.peak_memory_mb for o in success_outputs if o.peak_memory_mb > 0]
|
||||
|
||||
# Aggregate per-stage durations across all successful requests that reported them.
|
||||
stage_duration_lists: dict[str, list[float]] = {}
|
||||
for o in success_outputs:
|
||||
for stage, duration in (o.stage_durations or {}).items():
|
||||
stage_duration_lists.setdefault(stage, []).append(duration)
|
||||
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()}
|
||||
stage_durations_p99 = {s: float(np.percentile(v, 99)) for s, v in stage_duration_lists.items()}
|
||||
|
||||
metrics = {
|
||||
"duration": total_duration,
|
||||
"completed_requests": num_success,
|
||||
@@ -758,6 +778,9 @@ def calculate_metrics(
|
||||
"peak_memory_mb_max": max(peak_memories) if peak_memories else 0,
|
||||
"peak_memory_mb_mean": np.mean(peak_memories) if peak_memories else 0,
|
||||
"peak_memory_mb_median": np.median(peak_memories) if peak_memories else 0,
|
||||
"stage_durations_mean": stage_durations_mean,
|
||||
"stage_durations_p50": stage_durations_p50,
|
||||
"stage_durations_p99": stage_durations_p99,
|
||||
}
|
||||
|
||||
if slo_enabled:
|
||||
@@ -881,6 +904,8 @@ async def benchmark(args):
|
||||
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))
|
||||
|
||||
@@ -953,6 +978,12 @@ async def benchmark(args):
|
||||
print("{:<40} {:<15.2f}".format("Peak Memory Mean (MB):", metrics["peak_memory_mb_mean"]))
|
||||
print("{:<40} {:<15.2f}".format("Peak Memory Median (MB):", metrics["peak_memory_mb_median"]))
|
||||
|
||||
if metrics["stage_durations_mean"]:
|
||||
print(f"{'-' * 50}")
|
||||
print("Stage Durations Mean (s):")
|
||||
for stage, val in metrics["stage_durations_mean"].items():
|
||||
print("{:<40} {:<15.4f}".format(f" {stage}:", val))
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
|
||||
if args.output_file:
|
||||
|
||||
@@ -0,0 +1,460 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
Benchmark quantization quality loss for diffusion models (image & video).
|
||||
|
||||
Generates outputs with BF16 (baseline) and a quantized config using the same
|
||||
seed, then computes LPIPS perceptual distance between them. Results are printed
|
||||
as a Markdown table ready to paste into a PR description.
|
||||
|
||||
Requirements:
|
||||
pip install lpips Pillow numpy
|
||||
|
||||
Image example (text-to-image):
|
||||
python benchmarks/diffusion/quantization_quality.py \
|
||||
--model Tongyi-MAI/Z-Image-Turbo \
|
||||
--task t2i \
|
||||
--quantization fp8 \
|
||||
--prompts \
|
||||
"an aerial view of a coral reef with crystal clear turquoise water" \
|
||||
"a campfire in a dark forest with sparks rising into a starry sky" \
|
||||
"a gourmet dessert plate with chocolate mousse and gold leaf" \
|
||||
--height 1024 --width 1024 \
|
||||
--num-inference-steps 50 --seed 42
|
||||
|
||||
Video example (text-to-video):
|
||||
python benchmarks/diffusion/quantization_quality.py \
|
||||
--model Wan-AI/Wan2.2-T2V-A14B-Diffusers \
|
||||
--task t2v \
|
||||
--quantization fp8 \
|
||||
--prompts \
|
||||
"A serene lakeside sunrise with mist over the water" \
|
||||
"A cat walking across a wooden bridge in autumn" \
|
||||
--height 720 --width 1280 \
|
||||
--num-frames 81 --num-inference-steps 40 --seed 42
|
||||
|
||||
Multiple quantization methods:
|
||||
python benchmarks/diffusion/quantization_quality.py \
|
||||
--model Tongyi-MAI/Z-Image-Turbo \
|
||||
--task t2i \
|
||||
--quantization fp8 int8 bitsandbytes \
|
||||
--prompts "a cup of coffee on the table" \
|
||||
--height 1024 --width 1024 \
|
||||
--num-inference-steps 50 --seed 42
|
||||
|
||||
Output directory structure (--output-dir, default: ./quant_bench_output):
|
||||
quant_bench_output/
|
||||
baseline/ # BF16 outputs
|
||||
<method>/ # Quantized outputs per method
|
||||
results.md # Markdown table
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import gc
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
|
||||
def compute_lpips_images(
|
||||
baseline_images: list,
|
||||
quantized_images: list,
|
||||
net: str = "alex",
|
||||
) -> list[float]:
|
||||
"""Compute LPIPS between paired lists of PIL images."""
|
||||
import lpips
|
||||
from torchvision import transforms
|
||||
|
||||
loss_fn = lpips.LPIPS(net=net).eval()
|
||||
if torch.cuda.is_available():
|
||||
loss_fn = loss_fn.cuda()
|
||||
|
||||
transform = transforms.Compose(
|
||||
[
|
||||
transforms.Resize((256, 256)),
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
|
||||
]
|
||||
)
|
||||
|
||||
scores = []
|
||||
for img_bl, img_qt in zip(baseline_images, quantized_images):
|
||||
t_bl = transform(img_bl.convert("RGB")).unsqueeze(0)
|
||||
t_qt = transform(img_qt.convert("RGB")).unsqueeze(0)
|
||||
if torch.cuda.is_available():
|
||||
t_bl, t_qt = t_bl.cuda(), t_qt.cuda()
|
||||
with torch.no_grad():
|
||||
score = loss_fn(t_bl, t_qt).item()
|
||||
scores.append(score)
|
||||
return scores
|
||||
|
||||
|
||||
def compute_lpips_video(
|
||||
baseline_frames: np.ndarray,
|
||||
quantized_frames: np.ndarray,
|
||||
net: str = "alex",
|
||||
) -> float:
|
||||
"""Compute mean per-frame LPIPS for a video pair.
|
||||
|
||||
Args:
|
||||
baseline_frames: (F, H, W, C) float array in [0, 1].
|
||||
quantized_frames: same shape.
|
||||
|
||||
Returns:
|
||||
Mean LPIPS across all frames.
|
||||
"""
|
||||
import lpips
|
||||
|
||||
loss_fn = lpips.LPIPS(net=net).eval()
|
||||
if torch.cuda.is_available():
|
||||
loss_fn = loss_fn.cuda()
|
||||
|
||||
num_frames = min(len(baseline_frames), len(quantized_frames))
|
||||
scores = []
|
||||
for i in range(num_frames):
|
||||
# Convert (H, W, C) float [0,1] -> (1, C, H, W) float [-1, 1]
|
||||
f_bl = torch.from_numpy(baseline_frames[i]).permute(2, 0, 1).unsqueeze(0).float() * 2 - 1
|
||||
f_qt = torch.from_numpy(quantized_frames[i]).permute(2, 0, 1).unsqueeze(0).float() * 2 - 1
|
||||
if torch.cuda.is_available():
|
||||
f_bl, f_qt = f_bl.cuda(), f_qt.cuda()
|
||||
with torch.no_grad():
|
||||
score = loss_fn(f_bl, f_qt).item()
|
||||
scores.append(score)
|
||||
return float(np.mean(scores))
|
||||
|
||||
|
||||
def _build_omni_kwargs(args, quantization=None):
|
||||
"""Build kwargs dict for Omni() constructor."""
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
|
||||
parallel_config = DiffusionParallelConfig(
|
||||
ulysses_degree=args.ulysses_degree,
|
||||
ring_degree=args.ring_degree,
|
||||
tensor_parallel_size=args.tensor_parallel_size,
|
||||
)
|
||||
kwargs = {
|
||||
"model": args.model,
|
||||
"parallel_config": parallel_config,
|
||||
"enforce_eager": args.enforce_eager,
|
||||
}
|
||||
if quantization:
|
||||
kwargs["quantization_config"] = quantization
|
||||
return kwargs
|
||||
|
||||
|
||||
def _generate_image(omni, args, prompt, seed):
|
||||
"""Generate a single image and return (PIL.Image, time_seconds, memory_gib)."""
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(seed)
|
||||
torch.cuda.reset_peak_memory_stats()
|
||||
start = time.perf_counter()
|
||||
outputs = omni.generate(
|
||||
{"prompt": prompt},
|
||||
OmniDiffusionSamplingParams(
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
generator=generator,
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
),
|
||||
)
|
||||
elapsed = time.perf_counter() - start
|
||||
peak_mem = torch.cuda.max_memory_allocated() / (1024**3)
|
||||
|
||||
first = outputs[0]
|
||||
req_out = first.request_output[0] if hasattr(first, "request_output") else first
|
||||
img = req_out.images[0]
|
||||
return img, elapsed, peak_mem
|
||||
|
||||
|
||||
def _generate_video(omni, args, prompt, seed):
|
||||
"""Generate a video and return (np.ndarray [F,H,W,C], time_seconds, memory_gib)."""
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(seed)
|
||||
torch.cuda.reset_peak_memory_stats()
|
||||
start = time.perf_counter()
|
||||
outputs = omni.generate(
|
||||
{"prompt": prompt, "negative_prompt": ""},
|
||||
OmniDiffusionSamplingParams(
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
generator=generator,
|
||||
guidance_scale=args.guidance_scale,
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
num_frames=args.num_frames,
|
||||
),
|
||||
)
|
||||
elapsed = time.perf_counter() - start
|
||||
peak_mem = torch.cuda.max_memory_allocated() / (1024**3)
|
||||
|
||||
first = outputs[0]
|
||||
if hasattr(first, "request_output") and isinstance(first.request_output, list):
|
||||
inner = first.request_output[0]
|
||||
if isinstance(inner, OmniRequestOutput) and hasattr(inner, "images"):
|
||||
frames = inner.images[0] if inner.images else None
|
||||
else:
|
||||
frames = inner
|
||||
elif hasattr(first, "images") and first.images:
|
||||
frames = first.images
|
||||
else:
|
||||
raise ValueError("Could not extract video frames from output.")
|
||||
|
||||
if isinstance(frames, torch.Tensor):
|
||||
video = frames.detach().cpu()
|
||||
if video.dim() == 5:
|
||||
video = video[0].permute(1, 2, 3, 0) if video.shape[1] in (3, 4) else video[0]
|
||||
elif video.dim() == 4 and video.shape[0] in (3, 4):
|
||||
video = video.permute(1, 2, 3, 0)
|
||||
if video.is_floating_point():
|
||||
video = video.clamp(-1, 1) * 0.5 + 0.5
|
||||
frames_array = video.float().numpy()
|
||||
else:
|
||||
frames_array = np.asarray(frames)
|
||||
if frames_array.ndim == 5:
|
||||
frames_array = frames_array[0]
|
||||
|
||||
return frames_array, elapsed, peak_mem
|
||||
|
||||
|
||||
def _unload_omni(omni):
|
||||
"""Delete Omni instance and free GPU memory."""
|
||||
del omni
|
||||
gc.collect()
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
|
||||
def run_benchmark(args):
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
output_dir = Path(args.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
is_video = args.task == "t2v"
|
||||
prompts = args.prompts
|
||||
seed = args.seed
|
||||
|
||||
# Determine configs to benchmark
|
||||
configs = [] # list of (label, quantization_method)
|
||||
for method in args.quantization:
|
||||
configs.append((method, method))
|
||||
|
||||
# --- Baseline run ---
|
||||
print("\n" + "=" * 60)
|
||||
print("Running BF16 baseline...")
|
||||
print("=" * 60)
|
||||
bl_kwargs = _build_omni_kwargs(args, quantization=None)
|
||||
omni_bl = Omni(**bl_kwargs)
|
||||
|
||||
baseline_outputs = {} # prompt -> (output, time, mem)
|
||||
for prompt in prompts:
|
||||
print(f" Generating: {prompt[:60]}...")
|
||||
if is_video:
|
||||
out, t, mem = _generate_video(omni_bl, args, prompt, seed)
|
||||
else:
|
||||
out, t, mem = _generate_image(omni_bl, args, prompt, seed)
|
||||
baseline_outputs[prompt] = (out, t, mem)
|
||||
|
||||
bl_avg_time = np.mean([v[1] for v in baseline_outputs.values()])
|
||||
bl_mem = baseline_outputs[prompts[0]][2] # use first prompt's memory
|
||||
_unload_omni(omni_bl)
|
||||
|
||||
# Save baseline outputs
|
||||
bl_dir = output_dir / "baseline"
|
||||
bl_dir.mkdir(parents=True, exist_ok=True)
|
||||
for i, prompt in enumerate(prompts):
|
||||
out = baseline_outputs[prompt][0]
|
||||
if is_video:
|
||||
try:
|
||||
from diffusers.utils import export_to_video
|
||||
|
||||
frames_list = list(out) if isinstance(out, np.ndarray) and out.ndim == 4 else out
|
||||
export_to_video(frames_list, str(bl_dir / f"prompt_{i}.mp4"), fps=args.fps)
|
||||
except ImportError:
|
||||
np.save(bl_dir / f"prompt_{i}.npy", out)
|
||||
else:
|
||||
out.save(bl_dir / f"prompt_{i}.png")
|
||||
|
||||
# --- Quantized runs ---
|
||||
all_results = [] # list of dicts
|
||||
|
||||
for config_label, quant_method in configs:
|
||||
print(f"\n{'=' * 60}")
|
||||
print(f"Running: {config_label}...")
|
||||
print("=" * 60)
|
||||
|
||||
qt_kwargs = _build_omni_kwargs(args, quantization=quant_method)
|
||||
omni_qt = Omni(**qt_kwargs)
|
||||
|
||||
qt_outputs = {}
|
||||
for prompt in prompts:
|
||||
print(f" Generating: {prompt[:60]}...")
|
||||
if is_video:
|
||||
out, t, mem = _generate_video(omni_qt, args, prompt, seed)
|
||||
else:
|
||||
out, t, mem = _generate_image(omni_qt, args, prompt, seed)
|
||||
qt_outputs[prompt] = (out, t, mem)
|
||||
|
||||
qt_avg_time = np.mean([v[1] for v in qt_outputs.values()])
|
||||
qt_mem = qt_outputs[prompts[0]][2]
|
||||
_unload_omni(omni_qt)
|
||||
|
||||
# Save quantized outputs
|
||||
qt_dir = output_dir / config_label.replace(" ", "_")
|
||||
qt_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Compute LPIPS per prompt
|
||||
per_prompt = []
|
||||
for i, prompt in enumerate(prompts):
|
||||
bl_out = baseline_outputs[prompt][0]
|
||||
qt_out = qt_outputs[prompt][0]
|
||||
if is_video:
|
||||
lpips_score = compute_lpips_video(bl_out, qt_out, net=args.lpips_net)
|
||||
try:
|
||||
from diffusers.utils import export_to_video
|
||||
|
||||
frames_list = list(qt_out) if isinstance(qt_out, np.ndarray) and qt_out.ndim == 4 else qt_out
|
||||
export_to_video(frames_list, str(qt_dir / f"prompt_{i}.mp4"), fps=args.fps)
|
||||
except ImportError:
|
||||
np.save(qt_dir / f"prompt_{i}.npy", qt_out)
|
||||
else:
|
||||
lpips_score = compute_lpips_images([bl_out], [qt_out], net=args.lpips_net)[0]
|
||||
qt_out.save(qt_dir / f"prompt_{i}.png")
|
||||
per_prompt.append({"prompt": prompt, "lpips": lpips_score})
|
||||
|
||||
mean_lpips = np.mean([p["lpips"] for p in per_prompt])
|
||||
speedup = bl_avg_time / qt_avg_time if qt_avg_time > 0 else float("inf")
|
||||
mem_reduction = (bl_mem - qt_mem) / bl_mem * 100
|
||||
|
||||
all_results.append(
|
||||
{
|
||||
"config": config_label,
|
||||
"avg_time": qt_avg_time,
|
||||
"speedup": speedup,
|
||||
"memory_gib": qt_mem,
|
||||
"mem_reduction_pct": mem_reduction,
|
||||
"mean_lpips": mean_lpips,
|
||||
"per_prompt": per_prompt,
|
||||
}
|
||||
)
|
||||
|
||||
# --- Print results ---
|
||||
print("\n\n")
|
||||
print("=" * 80)
|
||||
print("RESULTS")
|
||||
print("=" * 80)
|
||||
|
||||
# Summary table
|
||||
lines = []
|
||||
lines.append(f"## Quantization Quality Benchmark — {args.model.split('/')[-1]}")
|
||||
lines.append(
|
||||
f"Setup: {args.height}x{args.width}, {args.num_inference_steps} steps, "
|
||||
f"seed={args.seed}, LPIPS ({args.lpips_net})"
|
||||
)
|
||||
if is_video:
|
||||
lines.append(f"Video: {args.num_frames} frames")
|
||||
lines.append("")
|
||||
lines.append("### Summary")
|
||||
lines.append("")
|
||||
lines.append("| Config | Avg Time | Speedup | Memory (GiB) | Mem Reduction | Mean LPIPS |")
|
||||
lines.append("|--------|----------|---------|--------------|---------------|------------|")
|
||||
lines.append(f"| BF16 baseline | {bl_avg_time:.2f}s | 1.00x | {bl_mem:.2f} | — | (ref) |")
|
||||
for r in all_results:
|
||||
lines.append(
|
||||
f"| {r['config']} | {r['avg_time']:.2f}s | {r['speedup']:.2f}x "
|
||||
f"| {r['memory_gib']:.2f} | {r['mem_reduction_pct']:.0f}% "
|
||||
f"| {r['mean_lpips']:.4f} |"
|
||||
)
|
||||
lines.append("")
|
||||
lines.append("> LPIPS < 0.01 = imperceptible, > 0.1 = clearly noticeable.")
|
||||
lines.append("")
|
||||
|
||||
# Per-prompt table
|
||||
if len(prompts) > 1:
|
||||
lines.append("### Per-Prompt LPIPS")
|
||||
lines.append("")
|
||||
header = "| Prompt |"
|
||||
sep = "|--------|"
|
||||
for r in all_results:
|
||||
header += f" {r['config']} |"
|
||||
sep += "--------|"
|
||||
lines.append(header)
|
||||
lines.append(sep)
|
||||
for i, prompt in enumerate(prompts):
|
||||
short = prompt[:50] + "..." if len(prompt) > 50 else prompt
|
||||
row = f"| {short} |"
|
||||
for r in all_results:
|
||||
row += f" {r['per_prompt'][i]['lpips']:.4f} |"
|
||||
lines.append(row)
|
||||
lines.append("")
|
||||
|
||||
md = "\n".join(lines)
|
||||
print(md)
|
||||
|
||||
# Save markdown
|
||||
results_path = output_dir / "results.md"
|
||||
results_path.write_text(md, encoding="utf-8")
|
||||
print(f"\nResults saved to {results_path}")
|
||||
print(f"Baseline outputs in {bl_dir}")
|
||||
for r in all_results:
|
||||
qt_dir = output_dir / r["config"].replace(" ", "_")
|
||||
print(f"Quantized outputs in {qt_dir}")
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Benchmark quantization quality loss for diffusion models.",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
)
|
||||
parser.add_argument("--model", required=True, help="Model name or local path.")
|
||||
parser.add_argument(
|
||||
"--task",
|
||||
default="t2i",
|
||||
choices=["t2i", "t2v"],
|
||||
help="Task type: t2i (text-to-image) or t2v (text-to-video).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--quantization",
|
||||
nargs="+",
|
||||
required=True,
|
||||
help="One or more quantization methods to benchmark (e.g. fp8 int8 bitsandbytes).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompts",
|
||||
nargs="+",
|
||||
default=["a cup of coffee on the table"],
|
||||
help="One or more prompts to generate.",
|
||||
)
|
||||
parser.add_argument("--seed", type=int, default=42)
|
||||
parser.add_argument("--height", type=int, default=1024)
|
||||
parser.add_argument("--width", type=int, default=1024)
|
||||
parser.add_argument("--num-inference-steps", type=int, default=50)
|
||||
parser.add_argument("--num-frames", type=int, default=81, help="Number of video frames (t2v only).")
|
||||
parser.add_argument("--fps", type=int, default=24, help="Video FPS for saving (t2v only).")
|
||||
parser.add_argument("--guidance-scale", type=float, default=4.0, help="CFG scale (used for video).")
|
||||
parser.add_argument("--output-dir", type=str, default="./quant_bench_output", help="Directory to save outputs.")
|
||||
parser.add_argument(
|
||||
"--lpips-net",
|
||||
type=str,
|
||||
default="alex",
|
||||
choices=["alex", "vgg", "squeeze"],
|
||||
help="LPIPS backbone network.",
|
||||
)
|
||||
parser.add_argument("--ulysses-degree", type=int, default=1)
|
||||
parser.add_argument("--ring-degree", type=int, default=1)
|
||||
parser.add_argument("--tensor-parallel-size", type=int, default=1)
|
||||
parser.add_argument("--enforce-eager", action="store_true")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
run_benchmark(args)
|
||||
+14
-1
@@ -2,9 +2,22 @@ ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
|
||||
ARG VLLM_ASCEND_TAG=v0.17.0rc1
|
||||
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
|
||||
|
||||
WORKDIR /vllm-workspace/vllm
|
||||
RUN git fetch origin --tags && git checkout v0.18.0
|
||||
|
||||
WORKDIR /vllm-workspace/vllm-ascend
|
||||
RUN git fetch origin releases/v0.18.0 && git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
|
||||
# Install vllm-ascend
|
||||
# Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
|
||||
RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
|
||||
source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
|
||||
source /usr/local/Ascend/nnal/atb/set_env.sh && \
|
||||
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
|
||||
python3 -m pip install -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
|
||||
python3 -m pip cache purge
|
||||
|
||||
ARG APP_DIR=/vllm-workspace/vllm-omni
|
||||
WORKDIR ${APP_DIR}
|
||||
|
||||
COPY . .
|
||||
|
||||
RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
|
||||
|
||||
@@ -2,9 +2,22 @@ ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
|
||||
ARG VLLM_ASCEND_TAG=v0.17.0rc1-a3
|
||||
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
|
||||
|
||||
WORKDIR /vllm-workspace/vllm
|
||||
RUN git fetch origin --tags && git checkout v0.18.0
|
||||
|
||||
WORKDIR /vllm-workspace/vllm-ascend
|
||||
RUN git fetch origin releases/v0.18.0 && git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
|
||||
# Install vllm-ascend
|
||||
# Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
|
||||
RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
|
||||
source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
|
||||
source /usr/local/Ascend/nnal/atb/set_env.sh && \
|
||||
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
|
||||
python3 -m pip install -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
|
||||
python3 -m pip cache purge
|
||||
|
||||
ARG APP_DIR=/vllm-workspace/vllm-omni
|
||||
WORKDIR ${APP_DIR}
|
||||
|
||||
COPY . .
|
||||
|
||||
RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
|
||||
|
||||
@@ -6,6 +6,7 @@ nav:
|
||||
- getting_started/installation/*
|
||||
- Serving:
|
||||
- OpenAI-Compatible API:
|
||||
- Diffusion Chat API: serving/diffusion_chat_api.md
|
||||
- Image Generation: serving/image_generation_api.md
|
||||
- Image Edit: serving/image_edit_api.md
|
||||
- Text to Speech: serving/speech_api.md
|
||||
@@ -28,6 +29,7 @@ nav:
|
||||
- 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-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
|
||||
@@ -78,6 +80,7 @@ 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
|
||||
@@ -92,11 +95,13 @@ nav:
|
||||
- design/feature/cache_dit.md
|
||||
- design/feature/teacache.md
|
||||
- design/feature/async_chunk_design.md
|
||||
- design/feature/vae_parallel.md
|
||||
- design/feature/diffusion_step_execution.md
|
||||
- Module Design:
|
||||
- design/module/ar_module.md
|
||||
- design/module/dit_module.md
|
||||
- design/module/entrypoint_module.md
|
||||
- design/module/async_omni_architecture.md
|
||||
- Docs Guide: contributing/DOCS_GUIDE.md
|
||||
- API Reference:
|
||||
- api/README.md
|
||||
|
||||
@@ -6,6 +6,7 @@ Main entry points for vLLM-Omni inference and serving.
|
||||
|
||||
- [vllm_omni.entrypoints.async_omni.AsyncOmni][]
|
||||
- [vllm_omni.entrypoints.async_omni_diffusion.AsyncOmniDiffusion][]
|
||||
- [vllm_omni.entrypoints.cfg_companion_tracker.CfgCompanionTracker][]
|
||||
- [vllm_omni.entrypoints.cli.benchmark.base.OmniBenchmarkSubcommandBase][]
|
||||
- [vllm_omni.entrypoints.cli.benchmark.main.OmniBenchmarkSubcommand][]
|
||||
- [vllm_omni.entrypoints.cli.benchmark.serve.OmniBenchmarkServingSubcommand][]
|
||||
@@ -13,11 +14,13 @@ Main entry points for vLLM-Omni inference and serving.
|
||||
- [vllm_omni.entrypoints.client_request_state.ClientRequestState][]
|
||||
- [vllm_omni.entrypoints.omni.Omni][]
|
||||
- [vllm_omni.entrypoints.omni_base.OmniBase][]
|
||||
- [vllm_omni.entrypoints.pd_utils.PDDisaggregationMixin][]
|
||||
|
||||
## Inputs
|
||||
|
||||
Input data structures for multi-modal inputs.
|
||||
|
||||
- [vllm_omni.inputs.data.OmniCustomPrompt][]
|
||||
- [vllm_omni.inputs.data.OmniDiffusionSamplingParams][]
|
||||
- [vllm_omni.inputs.data.OmniEmbedsPrompt][]
|
||||
- [vllm_omni.inputs.data.OmniTextPrompt][]
|
||||
@@ -43,8 +46,18 @@ Engine classes for offline and online inference.
|
||||
- [vllm_omni.engine.OmniEngineCoreRequest][]
|
||||
- [vllm_omni.engine.PromptEmbedsPayload][]
|
||||
- [vllm_omni.engine.arg_utils.OmniEngineArgs][]
|
||||
- [vllm_omni.engine.async_omni_engine.AsyncOmniEngine][]
|
||||
- [vllm_omni.engine.mm_outputs.MultimodalCompletionOutput][]
|
||||
- [vllm_omni.engine.mm_outputs.MultimodalPayload][]
|
||||
- [vllm_omni.engine.orchestrator.Orchestrator][]
|
||||
- [vllm_omni.engine.orchestrator.OrchestratorRequestState][]
|
||||
- [vllm_omni.engine.output_modality.OutputModality][]
|
||||
- [vllm_omni.engine.output_modality.TensorAccumulationStrategy][]
|
||||
- [vllm_omni.engine.output_processor.MultimodalOutputProcessor][]
|
||||
- [vllm_omni.engine.output_processor.OmniRequestState][]
|
||||
- [vllm_omni.engine.stage_engine_core_client.StageEngineCoreClient][]
|
||||
- [vllm_omni.engine.stage_init_utils.StageMetadata][]
|
||||
- [vllm_omni.engine.stage_init_utils.StartedLlmStage][]
|
||||
|
||||
## Core
|
||||
|
||||
@@ -56,6 +69,19 @@ Core scheduling and caching components.
|
||||
- [vllm_omni.core.sched.output.OmniCachedRequestData][]
|
||||
- [vllm_omni.core.sched.output.OmniNewRequestData][]
|
||||
- [vllm_omni.core.sched.output.OmniSchedulerOutput][]
|
||||
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.cfm.BASECFM][]
|
||||
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.cfm.CausalConditionalCFM][]
|
||||
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.cfm.CausalMaskedDiffWithDiT][]
|
||||
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.cfm.ConditionalCFM][]
|
||||
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan.CausalConv1d][]
|
||||
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan.CausalConv1dUpsample][]
|
||||
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan.CausalConvRNNF0Predictor][]
|
||||
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan.CausalHiFTGenerator][]
|
||||
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan.HiFTGenerator][]
|
||||
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan.SineGen][]
|
||||
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan.SineGen2][]
|
||||
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan.Snake][]
|
||||
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan.SourceModuleHnNSF][]
|
||||
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.DistributedGroupResidualVectorQuantization][]
|
||||
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.DistributedResidualVectorQuantization][]
|
||||
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.EuclideanCodebook][]
|
||||
@@ -67,9 +93,17 @@ Core scheduling and caching components.
|
||||
Configuration classes.
|
||||
|
||||
- [vllm_omni.config.model.OmniModelConfig][]
|
||||
- [vllm_omni.config.stage_config.ModelPipeline][]
|
||||
- [vllm_omni.config.stage_config.StageConfig][]
|
||||
- [vllm_omni.config.stage_config.StageConfigFactory][]
|
||||
- [vllm_omni.config.stage_config.StageType][]
|
||||
- [vllm_omni.diffusion.cache.teacache.config.TeaCacheConfig][]
|
||||
- [vllm_omni.distributed.omni_connectors.utils.config.ConnectorSpec][]
|
||||
- [vllm_omni.distributed.omni_connectors.utils.config.OmniTransferConfig][]
|
||||
- [vllm_omni.model_executor.models.cosyvoice3.config.CosyVoice3Config][]
|
||||
- [vllm_omni.model_executor.models.fish_speech.configuration_fish_speech.FishSpeechConfig][]
|
||||
- [vllm_omni.model_executor.models.fish_speech.configuration_fish_speech.FishSpeechFastARConfig][]
|
||||
- [vllm_omni.model_executor.models.fish_speech.configuration_fish_speech.FishSpeechSlowARConfig][]
|
||||
- [vllm_omni.model_executor.models.mimo_audio.config_mimo_audio.MiMoAudioConfig][]
|
||||
- [vllm_omni.model_executor.models.mimo_audio.config_mimo_audio.MiMoAudioTokenizerConfig][]
|
||||
- [vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts.Qwen3TTSConfig][]
|
||||
@@ -83,6 +117,10 @@ Configuration classes.
|
||||
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1DecoderConfig][]
|
||||
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1DecoderDiTConfig][]
|
||||
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1EncoderConfig][]
|
||||
- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Config][]
|
||||
- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLConfig][]
|
||||
- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLTextConfig][]
|
||||
- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLVisionConfig][]
|
||||
|
||||
## Workers
|
||||
|
||||
@@ -93,6 +131,8 @@ Worker classes and model runners for distributed inference.
|
||||
- [vllm_omni.diffusion.worker.diffusion_worker.DiffusionWorker][]
|
||||
- [vllm_omni.diffusion.worker.diffusion_worker.WorkerProc][]
|
||||
- [vllm_omni.diffusion.worker.diffusion_worker.WorkerWrapperBase][]
|
||||
- [vllm_omni.diffusion.worker.utils.DiffusionRequestState][]
|
||||
- [vllm_omni.diffusion.worker.utils.RunnerOutput][]
|
||||
- [vllm_omni.platforms.npu.worker.npu_ar_model_runner.ExecuteModelState][]
|
||||
- [vllm_omni.platforms.npu.worker.npu_ar_model_runner.NPUARModelRunner][]
|
||||
- [vllm_omni.platforms.npu.worker.npu_ar_worker.NPUARWorker][]
|
||||
@@ -110,5 +150,6 @@ Worker classes and model runners for distributed inference.
|
||||
- [vllm_omni.worker.gpu_generation_model_runner.GPUGenerationModelRunner][]
|
||||
- [vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker][]
|
||||
- [vllm_omni.worker.gpu_memory_utils.parse_cuda_visible_devices][]
|
||||
- [vllm_omni.worker.gpu_model_runner.CUDAGraphWrapper][]
|
||||
- [vllm_omni.worker.gpu_model_runner.OmniGPUModelRunner][]
|
||||
- [vllm_omni.worker.mixins.OmniWorkerMixin][]
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 97 KiB After Width: | Height: | Size: 99 KiB |
@@ -0,0 +1,6 @@
|
||||
nav:
|
||||
- CI_5levels.md
|
||||
- failures.md
|
||||
- test_guide.md
|
||||
- test_markers.md
|
||||
- test_style.md
|
||||
@@ -86,7 +86,7 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
|
||||
/tests/e2e/online_serving/test_{model_name}_expansion.py<br>
|
||||
/tests/e2e/offline_inference/test_{model_name}_expansion.py<br>
|
||||
<strong>Performance:</strong><br>
|
||||
/tests/e2e/perf/nightly.json<br>
|
||||
/tests/dfx/perf/tests/test.json<br>
|
||||
<strong>Doc Test:</strong><br>
|
||||
tests/example/online_serving/test_{model_name}.py<br>
|
||||
tests/example/offline_inference/test_{model_name}.py
|
||||
@@ -104,7 +104,7 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
|
||||
<td> Depends on reality </td>
|
||||
<td>
|
||||
<strong>Stability:</strong><br>
|
||||
tests/e2e/stability/weekly.json<br>
|
||||
/tests/dfx/stability/tests/test.json<br>
|
||||
<strong>Reliability:</strong><br>
|
||||
tests/e2e/reliability/test_{model_name}.py
|
||||
</td>
|
||||
@@ -225,9 +225,11 @@ vllm_omni/ tests/
|
||||
│
|
||||
└── e2e (tests side only) → ├── e2e/
|
||||
├── online_serving/ ✅ non-empty
|
||||
│ ├── test_qwen2_5_omni.py
|
||||
│ ├── test_async_omni.py
|
||||
│ ├── test_qwen3_omni.py
|
||||
│ ├── test_qwen3_omni_expansion.py
|
||||
│ ├── test_mimo_audio.py
|
||||
│ ├── test_image_gen_edit.py
|
||||
│ ├── test_images_generations_lora.py
|
||||
│ └── stage_configs/
|
||||
@@ -529,13 +531,13 @@ L4 level testing is a comprehensive quality audit before a version release. It e
|
||||
### 3.2 Testing Content and Scope
|
||||
|
||||
- ***Full Functionality Testing***: Executes all test cases defined in `test_{model_name}_expansion.py`, covering all implemented features, positive flows, boundary conditions, and exception handling.
|
||||
- ***Performance Testing***: Uses the `/tests/e2e/perf/nightly.json` configuration file to drive performance testing tools for stress, load, and endurance tests, collecting metrics like throughput, response time, and resource utilization.
|
||||
- ***Performance Testing***: Uses the `tests/dfx/perf/tests/test.json` configuration file to drive performance testing tools for stress, load, and endurance tests, collecting metrics like throughput, response time, and resource utilization.
|
||||
- ***Documentation Testing***: Verifies whether the example code provided to users is runnable and its results match the description.
|
||||
|
||||
### 3.3 Test Directory and Execution Files
|
||||
|
||||
- ***Functional Testing***: Same directories as L3.
|
||||
- ***Performance Test Configuration***: `/tests/e2e/perf/nightly.json`
|
||||
- ***Performance Test Configuration***: `tests/dfx/perf/tests/test.json`
|
||||
- ***Documentation Example Tests***:
|
||||
- - `tests/example/online_serving/test_{model_name}.py`
|
||||
- `tests/example/offline_inference/test_{model_name}.py`
|
||||
@@ -545,97 +547,104 @@ L4 level testing is a comprehensive quality audit before a version release. It e
|
||||
- ***Trigger Timing***: **`Nightly`**, automatically executed every night.
|
||||
- ***Execution Environment***: ***GPU*** server clusters to meet the resource demands of performance testing.
|
||||
- ***Script Example***:
|
||||
<details>
|
||||
<summary> Test Examples</summary>
|
||||
When you want to add L4-level performance test cases, you can refer to the following format for case addition in tests/perf/tests/test.json:
|
||||
|
||||
```JSON
|
||||
{
|
||||
"test_name": "test_qwen3_omni",
|
||||
"server_params": {
|
||||
"model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
|
||||
"stage_config_name": "qwen3_omni.yaml"
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"dataset_name": "random",
|
||||
"num_prompts": [10, 20],
|
||||
"request_rate": [0.5, 1],
|
||||
"random_input_len": 2500,
|
||||
"random_output_len": 900,
|
||||
"ignore_eos": true,
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"mean_ttft_ms": 100000,
|
||||
"mean_audio_ttfp_ms": 100000,
|
||||
"mean_audio_rtf": 100000
|
||||
???+ example "Test Examples"
|
||||
|
||||
When adding L4-level ***documentation example Tests***, please pay attention to the following guides.
|
||||
|
||||
--8<-- "docs/contributing/ci/test_examples/doc_example_tests.inc.md"
|
||||
|
||||
When you want to add L4-level ***performance test*** cases, you can refer to the following format for case addition in tests/dfx/perf/tests/test.json:
|
||||
|
||||
```JSON
|
||||
{
|
||||
"test_name": "test_qwen3_omni",
|
||||
"server_params": {
|
||||
"model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
|
||||
"stage_config_name": "qwen3_omni.yaml"
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"dataset_name": "random",
|
||||
"num_prompts": [10, 20],
|
||||
"request_rate": [0.5, 1],
|
||||
"random_input_len": 2500,
|
||||
"random_output_len": 900,
|
||||
"ignore_eos": true,
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"mean_ttft_ms": 100000,
|
||||
"mean_audio_ttfp_ms": 100000,
|
||||
"mean_audio_rtf": 100000
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
**Parameter Explanation**
|
||||
|
||||
*Overview*
|
||||
|
||||
| Field | Required | Description |
|
||||
| ---------------- | -------- | --------------------------------------------------------------- |
|
||||
| test_name | Yes | Unique identifier for the test case |
|
||||
| server_params | Yes | Server-side configuration parameters |
|
||||
| benchmark_params | Yes | Benchmark running parameters (supports multiple configurations) |
|
||||
|
||||
**server_params Configuration**
|
||||
|
||||
*Basic Parameters*
|
||||
|
||||
| Parameter | Required | Example | Description |
|
||||
| ----------------- | -------- | ---------------------------------- | ----------------------------- |
|
||||
| model | Yes | "Qwen/Qwen3-Omni-30B-A3B-Instruct" | Model name or path |
|
||||
| stage_config_name | Yes | "qwen3_omni.yaml" | Stage configuration file name |
|
||||
|
||||
*Dynamic Configuration (update/delete)*
|
||||
|
||||
Supports incremental modifications based on the basic configuration:
|
||||
|
||||
| Operation | Description |
|
||||
| --------- | ------------------------------------ |
|
||||
| update | Update or add configuration items |
|
||||
| delete | Delete specified configuration items |
|
||||
|
||||
***Example***:
|
||||
|
||||
```
|
||||
"update": {
|
||||
"async_chunk": true, // Enable asynchronous chunk processing
|
||||
"stage_args": {
|
||||
"0": {
|
||||
"engine_args.custom_process_next_stage_input_func": "vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_async_chunk"
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
#### Parameter Explanation
|
||||
|
||||
***Overview***
|
||||
|
||||
| Field | Required | Description |
|
||||
| ---------------- | -------- | --------------------------------------------------------------- |
|
||||
| test_name | Yes | Unique identifier for the test case |
|
||||
| server_params | Yes | Server-side configuration parameters |
|
||||
| benchmark_params | Yes | Benchmark running parameters (supports multiple configurations) |
|
||||
|
||||
#### server_params Configuration
|
||||
|
||||
##### Basic Parameters
|
||||
|
||||
| Parameter | Required | Example | Description |
|
||||
| ----------------- | -------- | ---------------------------------- | ----------------------------- |
|
||||
| model | Yes | "Qwen/Qwen3-Omni-30B-A3B-Instruct" | Model name or path |
|
||||
| stage_config_name | Yes | "qwen3_omni.yaml" | Stage configuration file name |
|
||||
|
||||
##### Dynamic Configuration (update/delete)
|
||||
|
||||
Supports incremental modifications based on the basic configuration:
|
||||
|
||||
| Operation | Description |
|
||||
| --------- | ------------------------------------ |
|
||||
| update | Update or add configuration items |
|
||||
| delete | Delete specified configuration items |
|
||||
|
||||
***Example***:
|
||||
```
|
||||
"update": {
|
||||
"async_chunk": true, // Enable asynchronous chunk processing
|
||||
"stage_args": {
|
||||
"0": {
|
||||
"engine_args.custom_process_next_stage_input_func": "vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_async_chunk"
|
||||
},
|
||||
"delete": {
|
||||
"stage_args": {
|
||||
"2": ["custom_process_input_func"] // Delete this configuration for stage 2
|
||||
}
|
||||
}
|
||||
},
|
||||
"delete": {
|
||||
"stage_args": {
|
||||
"2": ["custom_process_input_func"] // Delete this configuration for stage 2
|
||||
}
|
||||
}
|
||||
```
|
||||
#### benchmark_params Configuration
|
||||
```
|
||||
|
||||
You can add any benchmark running parameters you need here. For all optional parameters, refer to the [benchmark documentation](https://github.com/vllm-project/vllm-omni/blob/main/docs/cli/bench/serve.md). General modifications are as follows:
|
||||
**benchmark_params Configuration**
|
||||
|
||||
1. Change the ---xxx-xx-xx running parameters to xxx_xx_xx format and fill them as keys in the JSON file.
|
||||
2. For boolean variables in the running parameters, modify them to forms such as ignore_eos: true/false and fill them into the JSON file.
|
||||
3. Add the baseline parameter to specify the required validation values, ensuring the validation metric names match those in the result.json generated by the benchmark.
|
||||
4. The qps and concurrency modes are mutually exclusive. For detailed explanations, see the table below:
|
||||
You can add any benchmark running parameters you need here. For all optional parameters, refer to the [benchmark documentation](https://github.com/vllm-project/vllm-omni/blob/main/docs/cli/bench/serve.md). General modifications are as follows:
|
||||
|
||||
| Parameter | Type | Required | Example/Values | Description |
|
||||
| --------------- | ----------- | -------- | --------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
|
||||
| num_prompts | int / array | Yes | 10,[10, 20, 30] | Number of requests. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of qps or max_concurrency, e.g., [10,10,10]. If an array is used, its length must match the number of qps or max_concurrency. |
|
||||
| request_rate | int / array | No | 1, [1, 2, 3] | Queries per second. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of num_prompts, e.g., [1,1,1]. If an array is used, its length must match the number of num_prompts. |
|
||||
| max_concurrency | int / array | No | 1, [1, 2, 3] | Queries per second. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of num_prompts, e.g., [1,1,1]. If an array is used, its length must match the number of num_prompts. |
|
||||
</details>
|
||||
1. Change the ---xxx-xx-xx running parameters to xxx_xx_xx format and fill them as keys in the JSON file.
|
||||
2. For boolean variables in the running parameters, modify them to forms such as ignore_eos: true/false and fill them into the JSON file.
|
||||
3. Add the baseline parameter to specify the required validation values, ensuring the validation metric names match those in the result.json generated by the benchmark.
|
||||
4. The qps and concurrency modes are mutually exclusive. For detailed explanations, see the table below:
|
||||
|
||||
- - ***Run Command***: (Specific commands would depend on the performance testing tool and configuration defined in `nightly.json`).
|
||||
| Parameter | Type | Required | Example/Values | Description |
|
||||
| --------------- | ----------- | -------- | --------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
|
||||
| num_prompts | int / array | Yes | 10,[10, 20, 30] | Number of requests. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of qps or max_concurrency, e.g., [10,10,10]. If an array is used, its length must match the number of qps or max_concurrency. |
|
||||
| request_rate | int / array | No | 1, [1, 2, 3] | Queries per second. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of num_prompts, e.g., [1,1,1]. If an array is used, its length must match the number of num_prompts. |
|
||||
| max_concurrency | int / array | No | 1, [1, 2, 3] | Queries per second. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of num_prompts, e.g., [1,1,1]. If an array is used, its length must match the number of num_prompts. |
|
||||
</details>
|
||||
|
||||
- - ***Run Command***: (Specific commands would depend on the performance testing tool and configuration defined in `nightly.json`).
|
||||
|
||||
## Chapter 4: L5 Level Testing - Stability and Reliability Testing
|
||||
|
||||
@@ -647,12 +656,12 @@ L5 level testing focuses on the performance of model services under ***long-runn
|
||||
|
||||
### 4.2 Testing Content and Scope
|
||||
|
||||
- ***Long-term Stability (Stability) Testing***: Uses the `tests/e2e/stability/weekly.json` configuration to run the service under moderate load for an extended period (e.g., over 12 hours), monitoring whether metrics like memory/VRAM usage, response time, and throughput degrade over time, and whether the service process remains stable.
|
||||
- ***Long-term Stability (Stability) Testing***: Uses the `tests/dfx/stability/tests/test.json` configuration to run the service under moderate load for an extended period (e.g., over 12 hours), monitoring whether metrics like memory/VRAM usage, response time, and throughput degrade over time, and whether the service process remains stable.
|
||||
- ***Reliability Testing***: Uses `tests/e2e/reliability/test_{model_name}.py` to actively simulate various fault and abnormal scenarios, such as: dependent service interruption, abnormal input data, network flicker, hardware resource preemption, etc., to verify the system's fault tolerance, self-healing, and graceful degradation capabilities.
|
||||
|
||||
### 4.3 Test Directory and Execution Files
|
||||
|
||||
- ***Stability Test Configuration***: `tests/e2e/stability/weekly.json`
|
||||
- ***Stability Test Configuration***: `tests/dfx/stability/tests/test.json`
|
||||
- ***Reliability Test Suite***: `tests/e2e/reliability/test_{model_name}.py`
|
||||
|
||||
### 4.4 Execution Method and Example
|
||||
@@ -662,12 +671,79 @@ L5 level testing focuses on the performance of model services under ***long-runn
|
||||
- ***Script Example***:
|
||||
<details>
|
||||
<summary> Test Examples</summary>
|
||||
```python
|
||||
# WIP
|
||||
|
||||
When you want to add L5-level stability test cases, you can refer to the following format for case addition in `tests/dfx/stability/tests/test.json`:
|
||||
|
||||
```json
|
||||
{
|
||||
"test_name": "test_qwen3_omni_stability",
|
||||
"server_params": {
|
||||
"model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
|
||||
"stage_config_name": "qwen3_omni.yaml"
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"dataset_name": "random",
|
||||
"backend": "openai-chat-omni",
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"duration_sec": 43200,
|
||||
"request_rate": 0.5,
|
||||
"num_prompts_per_batch": 20,
|
||||
"random_input_len": 2500,
|
||||
"random_output_len": 900,
|
||||
"ignore_eos": true,
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration"
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
#### Parameter Explanation
|
||||
|
||||
***Overview***
|
||||
|
||||
| Field | Required | Description |
|
||||
| ---------------- | -------- | --------------------------------------------------------------------------- |
|
||||
| test_name | Yes | Unique identifier for the stability test case |
|
||||
| server_params | Yes | Server-side configuration parameters (model, stage configuration, etc.) |
|
||||
| benchmark_params | Yes | Stability benchmark running parameters (supports multiple configurations) |
|
||||
|
||||
#### server_params Configuration
|
||||
|
||||
##### Basic Parameters
|
||||
|
||||
| Parameter | Required | Example | Description |
|
||||
| ----------------- | -------- | ---------------------------------- | ----------------------------------- |
|
||||
| model | Yes | "Qwen/Qwen3-Omni-30B-A3B-Instruct" | Model name or path |
|
||||
| stage_config_name | Yes | "qwen3_omni.yaml" | Stage configuration file name |
|
||||
|
||||
##### Dynamic Configuration (update/delete)
|
||||
|
||||
Supports incremental modifications based on the basic configuration:
|
||||
|
||||
| Operation | Description |
|
||||
| --------- | ------------------------------------ |
|
||||
| update | Update or add configuration items |
|
||||
| delete | Delete specified configuration items |
|
||||
|
||||
***Example***:
|
||||
You can refer to Test Examples in Chapter 3.4
|
||||
|
||||
#### benchmark_params Configuration
|
||||
|
||||
For stability testing, the key parameters are:
|
||||
|
||||
- **duration_sec**: Total duration (in seconds) during which benchmark traffic is sent. The stability benchmark will keep sending batches until this duration is reached.
|
||||
- **request_rate** / **max_concurrency**: Exactly one of them must be specified. They control how the traffic is generated for each batch:
|
||||
- `request_rate`: Number of requests per second. The benchmark will send `num_prompts_per_batch` requests at the given rate.
|
||||
- `max_concurrency`: Maximum number of concurrent requests. When this is used, `request_rate` is set to `inf` internally.
|
||||
- **num_prompts_per_batch**: Number of prompts sent in each batch. Multiple batches will be executed sequentially within `duration_sec`.
|
||||
|
||||
All other optional parameters follow the same rules as the in Chapter 3.4.
|
||||
|
||||
</details>
|
||||
|
||||
- - ***Stability***: (Execution would be driven by the configuration in `weekly.json`).
|
||||
- - ***Stability***: `pytest -s -v tests/dfx/stability/scripts/test_{model_name}.py`
|
||||
- ***Reliability***: `pytest -s -v tests/e2e/reliability/test_{model_name}.py`
|
||||
|
||||
## Summary
|
||||
|
||||
@@ -0,0 +1,49 @@
|
||||
**Preferred Test Strategy**
|
||||
|
||||
Use one of the following patterns depending on page type:
|
||||
|
||||
- **Dynamic code-block extraction (preferred for offline docs)**
|
||||
- Extract Python/Bash code blocks from markdown AST analyzer, then execute them directly in tests.
|
||||
- Benefit: test logic stays automatically aligned with docs.
|
||||
- Basic idea: Use `ReadmeSnippet.extract_readme_snippets` to extract a list of code blocks as a global variable in file,
|
||||
use this list as `pytest.mark.parametrize` parameters, and pass each snippet item to `example_runner.run` inside the parametrized test.
|
||||
Additionally pass an `output_subfolder` argument for the 2nd-level output folder explained in **Output Directory Structure** below.
|
||||
If any extra environment variable is need for a test (e.g., the example script reads it), `example_runner.run` also accepts a 3rd `env` parameter.
|
||||
- See [tests/examples/offline_inference/test_text_to_image.py](https://github.com/vllm-project/vllm-omni/blob/main/tests/examples/offline_inference/test_text_to_image.py) for reference implementation.
|
||||
|
||||
- **Explicit copied scripts (used by online docs for now until further update)**
|
||||
- For online serving pages, it is acceptable to copy code from docs into dedicated test functions, because only client-side, request-sending scripts are tested.
|
||||
- Benefit: dynamic extraction is overly complex: need to tell server-launch and client-request scripts.
|
||||
- Requirement: copied test code must be kept in sync with doc updates.
|
||||
|
||||
**Test Case Naming Convention**
|
||||
|
||||
- Dynamic code extraction (auto-generated internally):
|
||||
- `test_{single_function_name_matching_file_name}[h2_heading_00X]`
|
||||
- Example: `test_text_to_image[basic_usage_001]`
|
||||
- Explicit copied scripts:
|
||||
- `test_{h2_heading_00X}[{dummy_param_id_for_omni_server}]`
|
||||
- Example: `test_api_calls_001[omni_server0]`
|
||||
|
||||
**Runtime Configuration**
|
||||
|
||||
In the example code tests, do **not** reduce `num_inference_steps` just to speed up the tests unless there is a strong CI reliability reason to do otherwise.
|
||||
|
||||
**Skipping Rules**
|
||||
|
||||
You may skip examples falling in the following categories using `pytest.mark.skip` or `pytest.skip`:
|
||||
|
||||
- Gradio UI scripts
|
||||
- Scenarios that significantly overlap with existing tests and add little new coverage.
|
||||
|
||||
**Output Directory Structure**
|
||||
|
||||
Use a three-layer output structure to store output artifacts:
|
||||
|
||||
1. Root output directory
|
||||
- Auto-detected from `OUTPUT_DIR` env var or auto-generated under `/tmp`.
|
||||
2. Doc-page directory
|
||||
- Define and use a clear page-level folder name in each `test_*.py` yourself (abbreviations are acceptable, e.g., `example_offline_t2i`).
|
||||
3. Test-case directory
|
||||
- Must match the case identifier (e.g., `basic_usage_001`).
|
||||
- Auto-generated for dynamic extracted tests.
|
||||
@@ -130,9 +130,11 @@ vllm_omni/ tests/
|
||||
│
|
||||
└── e2e (tests side only) → ├── e2e/
|
||||
├── online_serving/ ✅ non-empty
|
||||
│ ├── test_qwen2_5_omni.py
|
||||
│ ├── test_async_omni.py
|
||||
│ ├── test_qwen3_omni.py
|
||||
│ ├── test_qwen3_omni_expansion.py
|
||||
│ ├── test_mimo_audio.py
|
||||
│ ├── test_image_gen_edit.py
|
||||
│ ├── test_images_generations_lora.py
|
||||
│ └── stage_configs/
|
||||
@@ -157,6 +159,12 @@ vllm_omni/ tests/
|
||||
├── qwen3_omni_ci.yaml
|
||||
├── bagel_*.yaml
|
||||
└── npu/, rocm/, etc.
|
||||
examples/ tests
|
||||
│ └── examples
|
||||
├── online_serving/ → ├── online_serving/
|
||||
│ └── {doc_page_title}/README.md │ └── test_{doc_page_title}.py ⬜
|
||||
└── offline_inference/ → └── offline_inference/
|
||||
└── {doc_page_title}/README.md └── test_{doc_page_title}.py ⬜
|
||||
```
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,478 @@
|
||||
# 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
|
||||
+106
-38
@@ -4,48 +4,54 @@
|
||||
|
||||
vLLM-Omni uses the PyTorch Profiler to analyze performance across both **multi-stage omni-modality models** and **diffusion models**.
|
||||
|
||||
### 1. Set the Output Directory
|
||||
Before running any script, set this environment variable. The system detects this and automatically saves traces here.
|
||||
### 1. Configure Profiling in the Stage YAML
|
||||
|
||||
```bash
|
||||
export VLLM_TORCH_PROFILER_DIR=./profiles
|
||||
Enable profiling by adding `profiler_config` under `engine_args` for the stage(s) you want to profile in your stage config YAML:
|
||||
|
||||
```yaml
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
stage_type: llm
|
||||
engine_args:
|
||||
# ... other engine args ...
|
||||
profiler_config:
|
||||
profiler: torch
|
||||
torch_profiler_dir: ./perf
|
||||
```
|
||||
|
||||
| Field | Description |
|
||||
|---|---|
|
||||
| `profiler` | Profiler backend to use. Currently supports `torch`. |
|
||||
| `torch_profiler_dir` | Directory where trace files are saved. Created automatically if it doesn't exist. |
|
||||
|
||||
> **Tip:** Only enable `profiler_config` on stages you actually need to profile. Stages without it will not start a profiler, keeping overhead minimal.
|
||||
|
||||
### 2. Profiling Omni-Modality Models
|
||||
|
||||
It is best to limit profiling to one iteration to keep trace files manageable.
|
||||
|
||||
```bash
|
||||
export VLLM_PROFILER_MAX_ITERS=1
|
||||
```
|
||||
|
||||
**Selective Stage Profiling**
|
||||
The profiler is default to function across all stages. But It is highly recommended to profile specific stages by passing the stages list, preventing from producing too large trace files:
|
||||
|
||||
It is highly recommended to profile specific stages to prevent producing overly large trace files:
|
||||
|
||||
```python
|
||||
# Profile all stages
|
||||
omni.start_profile()
|
||||
omni_llm.start_profile()
|
||||
|
||||
# Only profile Stage 1
|
||||
omni.start_profile(stages=[1])
|
||||
omni_llm.start_profile(stages=[1])
|
||||
|
||||
# Stage 0 (Thinker) and Stage 2 (Audio Decoder) for qwen omni
|
||||
omni_llm.start_profile(stages=[0, 2])
|
||||
```
|
||||
|
||||
```python
|
||||
# Stage 0 (Thinker) and Stage 2 (Audio Decoder) for qwen omni
|
||||
omni.start_profile(stages=[0, 2])
|
||||
```
|
||||
> **Important:** Always pass the same `stages` list to both `start_profile()` and `stop_profile()`. If you omit `stages` from `stop_profile()`, it defaults to stopping all stages — including ones that were never started — which will produce errors.
|
||||
|
||||
**Python Usage**: Wrap your generation logic with `start_profile()` and `stop_profile()`.
|
||||
|
||||
```python
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
profiler_stages = [0] # Only profile the stages you need
|
||||
|
||||
omni = Omni(model="Qwen/Qwen3-Omni-30B-A3B-Instruct")
|
||||
|
||||
profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
|
||||
|
||||
# 1. Start profiling if enabled
|
||||
if profiler_enabled:
|
||||
omni.start_profile(stages=[0])
|
||||
# 1. Start profiling
|
||||
omni.start_profile(stages=profiler_stages)
|
||||
|
||||
# Initialize generator
|
||||
omni_generator = omni.generate(prompts, sampling_params_list, py_generator=args.py_generator)
|
||||
@@ -59,42 +65,57 @@ for stage_outputs in omni_generator:
|
||||
# ... [Output processing logic for text/audio would go here] ...
|
||||
|
||||
# Update count to track when to stop profiling
|
||||
processed_count += 1
|
||||
processed_count += len(stage_outputs.request_output)
|
||||
|
||||
# 2. Check if all requests are done to stop the profiler safely
|
||||
if profiler_enabled and processed_count >= total_requests:
|
||||
print(f"[Info] Processed {processed_count}/{total_requests}. Stopping profiler inside active loop...")
|
||||
|
||||
# Stop the profiler while workers are still active
|
||||
omni.stop_profile()
|
||||
# Pass the same stages list used in start_profile()
|
||||
omni_llm.stop_profile(stages=profiler_stages)
|
||||
|
||||
# Wait for traces to flush to disk
|
||||
print("[Info] Waiting 30s for workers to write trace files to disk...")
|
||||
time.sleep(30)
|
||||
print("[Info] Trace export wait time finished.")
|
||||
|
||||
omni.close()
|
||||
omni_llm.close()
|
||||
```
|
||||
|
||||
|
||||
**CLI Usage** (using `end2end.py`):
|
||||
```bash
|
||||
# Profile only Stage 0 (Thinker)
|
||||
python end2end.py --output-wav output_audio \
|
||||
--query-type text --enable-profiler --profiler-stages 0
|
||||
|
||||
# Profile Stage 0 and Stage 2
|
||||
python end2end.py --output-wav output_audio \
|
||||
--query-type text --enable-profiler --profiler-stages 0 2
|
||||
|
||||
# Profile all stages (omit --profiler-stages)
|
||||
python end2end.py --output-wav output_audio \
|
||||
--query-type text --enable-profiler
|
||||
```
|
||||
|
||||
**Examples**:
|
||||
|
||||
1. **Qwen2.5-Omni**: [https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen2_5_omni/end2end.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen2_5_omni/end2end.py)
|
||||
|
||||
2. **Qwen3-Omni**: [https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen3_omni/end2end.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen3_omni/end2end.py)
|
||||
|
||||
|
||||
### 3. Profiling diffusion models
|
||||
|
||||
Diffusion profiling is End-to-End, capturing encoding, denoising loops, and decoding.
|
||||
Diffusion profiling is End-to-End, capturing encoding, denoising loops, and decoding. Standalone diffusion scripts use `--profiler-dir` to enable profiling.
|
||||
|
||||
**CLI Usage:**
|
||||
```python
|
||||
|
||||
```bash
|
||||
python image_to_video.py \
|
||||
--model Wan-AI/Wan2.2-I2V-A14B-Diffusers \
|
||||
--image qwen-bear.png \
|
||||
--prompt "A cat playing with yarn, smooth motion" \
|
||||
--profiler-dir \
|
||||
\
|
||||
# Minimize Spatial Dimensions (Optional but helpful):
|
||||
# Drastically reduces memory usage so the profiler doesn't
|
||||
@@ -124,25 +145,72 @@ python image_to_video.py \
|
||||
--flow-shift 12.0 \
|
||||
--fps 16 \
|
||||
--output i2v_output.mp4
|
||||
|
||||
```
|
||||
|
||||
> **Note:** For diffusion stages within a multi-stage omni pipeline, use `profiler_config` in the stage YAML instead (see Section 1).
|
||||
|
||||
**Examples**:
|
||||
|
||||
1. **Qwen image edit**: [https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/image_to_image/image_edit.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/image_to_image/image_edit.py)
|
||||
|
||||
2. **Wan-AI/Wan2.2-I2V-A14B-Diffusers**: [https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_video](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_video)
|
||||
|
||||
### 4. Analyzing Omni Traces
|
||||
### 4. Profiling Online Serving
|
||||
|
||||
Output files are saved to your configured ```VLLM_TORCH_PROFILER_DIR```.
|
||||
When `profiler_config` is set in the stage YAML, the server automatically exposes `/start_profile` and `/stop_profile` HTTP endpoints.
|
||||
|
||||
**1. Start the server** with a stage YAML that has `profiler_config` enabled:
|
||||
```bash
|
||||
vllm serve Qwen/Qwen2.5-Omni-7B \
|
||||
--omni \
|
||||
--stage-configs-path qwen2_5_omni.yaml \
|
||||
--port 8091
|
||||
```
|
||||
|
||||
Or for one stage diffusion models:
|
||||
|
||||
```bash
|
||||
vllm serve Wan-AI/Wan2.2-I2V-A14B-Diffusers --omni --port 8091 --profiler-config '{"profiler": "torch", "torch_profiler_dir": "./vllm_profile"}'
|
||||
```
|
||||
|
||||
**2. Start profiling** by sending a POST request:
|
||||
```bash
|
||||
# Profile all stages that have profiler_config set
|
||||
curl -X POST http://localhost:8091/start_profile
|
||||
|
||||
# Profile specific stages only
|
||||
curl -X POST http://localhost:8091/start_profile \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"stages": [0]}'
|
||||
```
|
||||
|
||||
**3. Send your inference requests** as normal while the profiler is running.
|
||||
|
||||
**4. Stop profiling** and collect traces:
|
||||
```bash
|
||||
# Stop all stages
|
||||
curl -X POST http://localhost:8091/stop_profile
|
||||
|
||||
# Stop specific stages (must match the stages you started)
|
||||
curl -X POST http://localhost:8091/stop_profile \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"stages": [0]}'
|
||||
```
|
||||
|
||||
Trace files are written to the `torch_profiler_dir` specified in your stage YAML.
|
||||
|
||||
> **Important:** Always stop the same stages you started. Stopping a stage that was never started will produce errors.
|
||||
|
||||
### 5. Analyzing Traces
|
||||
|
||||
Output files are saved to the `torch_profiler_dir` specified in your stage YAML config.
|
||||
|
||||
**Output**
|
||||
**Chrome Trace** (```.json.gz```): Visual timeline of kernels and stages. Open in Perfetto UI.
|
||||
**Chrome Trace** (`.json.gz`): Visual timeline of kernels and stages. Open in Perfetto UI.
|
||||
|
||||
**Viewing Tools:**
|
||||
|
||||
- [Perfetto](https://ui.perfetto.dev/)(recommended)
|
||||
- ```chrome://tracing```(Chrome only)
|
||||
- [Perfetto](https://ui.perfetto.dev/) (recommended)
|
||||
- `chrome://tracing` (Chrome only)
|
||||
|
||||
**Note**: vLLM-Omni reuses the PyTorch Profiler infrastructure from vLLM. See the official vLLM profiler documentation: [vLLM Profiling Guide](https://docs.vllm.ai/en/stable/contributing/profiling/)
|
||||
@@ -113,13 +113,13 @@ def enable_cache_for_dit(pipeline: Any, cache_config: Any) -> Callable[[int], No
|
||||
|
||||
Some models require custom handling:
|
||||
|
||||
- **Dual-transformer:** Models with separate high-noise and low-noise transformers (e.g., Wan2.2)
|
||||
- **Single or dual-transformer:** Models that may use one or two transformers (e.g., Wan2.2)
|
||||
- **Multi-block-list:** Models with multiple block lists in one transformer (e.g., LongCatImage with `transformer_blocks` + `single_transformer_blocks`)
|
||||
- **Special forward patterns:** Models with non-standard block execution patterns
|
||||
|
||||
### Example 1: Dual-Transformer Model (Wan2.2)
|
||||
### Example 1: Single or Dual-Transformer Model (Wan2.2)
|
||||
|
||||
Wan2.2 uses two transformers: one for high-noise steps and one for low-noise steps.
|
||||
Wan2.2 can use either a single transformer or two transformers (one for high-noise steps and one for low-noise steps). The implementation automatically detects the mode based on the presence of `transformer_2`.
|
||||
|
||||
**Key difference:** Use `BlockAdapter` to wrap multiple transformers with separate configurations.
|
||||
|
||||
@@ -267,7 +267,7 @@ Complete examples in the codebase:
|
||||
| Model | Path | Pattern | Notes |
|
||||
|-------|------|---------|-------|
|
||||
| **Standard DiT** | [`cache_dit_backend.py::enable_cache_for_dit`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_dit) | Default enabler | Single transformer, automatic |
|
||||
| **Wan2.2** | [`cache_dit_backend.py::enable_cache_for_wan22`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_wan22) | Dual-transformer | Separate high/low noise transformers |
|
||||
| **Wan2.2** | [`cache_dit_backend.py::enable_cache_for_wan22`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_wan22) | Single or dual-transformer | Auto-detects mode based on transformer_2 presence |
|
||||
| **LongCat** | [`cache_dit_backend.py::enable_cache_for_longcat_image`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_longcat_image) | Multi-block-list | Two block lists in one transformer |
|
||||
| **BAGEL** | [`cache_dit_backend.py::enable_cache_for_bagel`](https://docs.vllm.ai/projects/vllm-omni/en/latest/api/vllm_omni/diffusion/cache/cache_dit_backend/#vllm_omni.diffusion.cache.cache_dit_backend.enable_cache_for_bagel) | Omni model | Complex architecture |
|
||||
|
||||
|
||||
@@ -7,6 +7,7 @@ This section describes how to add Sequence Parallel (SP) to a diffusion transfor
|
||||
## Table of Contents
|
||||
|
||||
- [Overview](#overview)
|
||||
- [UAA Mode (Experimental)](#uaa-mode-experimental)
|
||||
- [Approach 1: Non-Intrusive `_sp_plan` (Recommended)](#approach-1-non-intrusive-_sp_plan-recommended)
|
||||
- [Approach 2: Intrusive Modification (For Complex Cases)](#approach-2-intrusive-modification-for-complex-cases)
|
||||
- [Testing](#testing)
|
||||
@@ -46,6 +47,41 @@ from vllm_omni.diffusion.distributed.sp_sharding import sp_shard, sp_gather
|
||||
|
||||
---
|
||||
|
||||
## UAA Mode (Experimental)
|
||||
|
||||
`ulysses_mode="advanced_uaa"` enables the experimental UAA ("Ulysses Anything Attention") feature, which lets Ulysses attention handle arbitrary sequence lengths and arbitrary attention head counts. The same idea is also supported by [Cache-DiT](https://cache-dit.readthedocs.io/en/latest/user_guide/CONTEXT_PARALLEL/#uaa-ulysses-anything-attention).
|
||||
|
||||
Use it when plain Ulysses-SP would otherwise fail because:
|
||||
|
||||
- the local sequence shards are not evenly divisible after split hooks, or
|
||||
- the attention head count is not divisible by `ulysses_degree`.
|
||||
|
||||
### Design Summary
|
||||
|
||||
1. **Strict mode stays unchanged.**
|
||||
`ulysses_mode="strict"` keeps the original fast path and still requires divisible sequence/head shapes.
|
||||
|
||||
2. **UAA uses variable all-to-all split sizes for sequence shards.**
|
||||
Before the Ulysses Q/K/V exchange, each rank all-gathers its local sequence length and uses those lengths as `all_to_all_single(..., output_split_sizes=seq_lens)`. This lets Ulysses gather the full sequence even when each rank started with a different local shard length.
|
||||
|
||||
3. **UAA pads heads only inside the Ulysses exchange.**
|
||||
If `head_cnt % ulysses_degree != 0`, UAA pads the head dimension up to the next multiple of `ulysses_degree`, performs the forward/reverse all-to-all, then slices the temporary head padding away after the reverse exchange. The same rule is applied to joint attention tensors.
|
||||
|
||||
4. **Hybrid Ulysses + Ring is still shape-constrained.**
|
||||
Ring attention expects every rank in a ring group to exchange exactly the same post-Ulysses sequence shape. UAA therefore validates those shapes before entering the ring path and raises a clear error if ring peers disagree on `S_global`.
|
||||
|
||||
5. **Tiny scalar gathers stay out of TorchDynamo tracing.**
|
||||
`_all_gather_int()` is marked with `@torch.compiler.disable` so the scalar `item()` conversions used by UAA metadata collection do not get traced into `torch.compile`.
|
||||
|
||||
### UAA vs `auto_pad`
|
||||
|
||||
- `auto_pad=True` pads sequence tokens in `_sp_plan` and requires attention backends that support `attention_mask`.
|
||||
- `advanced_uaa` does not depend on mask-based token padding inside Ulysses attention. It is therefore a better fit for non-divisible head counts and uneven Ulysses shard sizes.
|
||||
- `auto_pad=True` remains incompatible with Ring attention because the ring backend does not consume `attention_mask`.
|
||||
- `advanced_uaa` is still experimental and hybrid mode remains limited by Ring's equal-shape requirement.
|
||||
|
||||
---
|
||||
|
||||
## Approach 1: Non-Intrusive `_sp_plan` (Recommended)
|
||||
|
||||
The `_sp_plan` mechanism allows SP **without modifying `forward()` logic**. The framework automatically registers hooks to shard inputs and gather outputs at module boundaries.
|
||||
@@ -364,9 +400,14 @@ _sp_plan = {
|
||||
|
||||
- **Sequence Length not divisible by sp_size:**
|
||||
|
||||
**Problem:** `SequenceParallelInput(auto_pad=False)` - auto_pad should be True to enable automatic sequence padding.
|
||||
**Problem:** strict Ulysses sequence parallel requires divisible shapes. If the shard length is uneven, or if the model head count is not divisible by `ulysses_degree`, the strict path will raise an error.
|
||||
|
||||
**Solution:** In `SequenceParallelInput`, set `auto_pad=True` and add attention mask support.
|
||||
**Solutions:**
|
||||
|
||||
1. Use `ulysses_mode="advanced_uaa"` for Ulysses-SP when you want the experimental uneven-shape path without relying on attention-mask padding.
|
||||
2. If the model already uses `_sp_plan` token padding and the attention backend supports `attention_mask`, set `auto_pad=True` and add attention-mask plumbing.
|
||||
|
||||
> **Experimental Feature:** `ulysses_mode="advanced_uaa"` is experimental. It is intended to relax Ulysses divisibility constraints, but hybrid Ulysses + Ring still requires equal post-Ulysses sequence lengths inside each ring group.
|
||||
|
||||
> **Experimental Feature:** `auto_pad=True` is an experimental feature and may be changed in the future. We plan to improve this solution to involve minimal changes to model files. More details are [here](https://github.com/vllm-project/vllm-omni/issues/1324).
|
||||
|
||||
|
||||
@@ -127,80 +127,117 @@ def step(self, requests: list[OmniDiffusionRequest]):
|
||||
|
||||
## 2. Scheduler
|
||||
|
||||
**Location**: `vllm_omni/diffusion/scheduler.py`
|
||||
**Location**: `vllm_omni/diffusion/sched/`
|
||||
|
||||
### Architecture
|
||||
|
||||
The `Scheduler` is implemented as a **Singleton** pattern to ensure a single coordination point across the system, i.e., only one scheduler instance exists for coordination.
|
||||
The scheduler is a **request-state scheduler**. It owns request lifecycle management and scheduling decisions, while execution stays in `DiffusionEngine` and the executor.
|
||||
|
||||
### Key Components
|
||||
|
||||
#### 2.1 Message Queue System
|
||||
#### 2.1 Scheduler Interface
|
||||
|
||||
```python
|
||||
class Scheduler:
|
||||
def initialize(self, od_config: OmniDiffusionConfig):
|
||||
# Broadcast queue: scheduler -> all workers
|
||||
self.mq = MessageQueue(
|
||||
n_reader=od_config.num_gpus,
|
||||
n_local_reader=od_config.num_gpus,
|
||||
local_reader_ranks=list(range(od_config.num_gpus)),
|
||||
)
|
||||
|
||||
# Result queue: rank 0 worker -> scheduler
|
||||
self.result_mq = None # Initialized later
|
||||
class SchedulerInterface(ABC):
|
||||
def add_request(self, request: OmniDiffusionRequest) -> str: ...
|
||||
def schedule(self) -> DiffusionSchedulerOutput: ...
|
||||
def update_from_output(
|
||||
self,
|
||||
sched_output: DiffusionSchedulerOutput,
|
||||
output: DiffusionOutput,
|
||||
) -> set[str]: ...
|
||||
```
|
||||
|
||||
**Communication Pattern**:
|
||||
**Responsibilities**:
|
||||
|
||||
- **Broadcast Queue**: One-to-many communication (scheduler → all workers)
|
||||
- **Lifecycle contract**: Defines how the engine adds requests, triggers one scheduling cycle, and feeds executor results back.
|
||||
|
||||
- **Result Queue**: One-to-one communication (rank 0 → scheduler)
|
||||
- **Stable boundary**: `DiffusionSchedulerOutput` is the only scheduling result consumed by `DiffusionEngine`.
|
||||
|
||||
- **Shared Memory**: Uses `MessageQueue` (ZMQ-based) for efficient IPC
|
||||
- **Pluggability**: Different scheduler policies can reuse the same engine integration path.
|
||||
|
||||
#### 2.2 Request Distribution
|
||||
#### 2.2 Request State Model
|
||||
|
||||
```python
|
||||
def add_req(self, requests: list[OmniDiffusionRequest]) -> DiffusionOutput:
|
||||
# Broadcast request to all workers
|
||||
self.mq.enqueue(requests)
|
||||
class DiffusionRequestStatus(enum.IntEnum):
|
||||
WAITING = ...
|
||||
RUNNING = ...
|
||||
PREEMPTED = ...
|
||||
FINISHED_COMPLETED = ...
|
||||
FINISHED_ABORTED = ...
|
||||
FINISHED_ERROR = ...
|
||||
|
||||
# Wait for result from Rank 0
|
||||
output = self.result_mq.dequeue()
|
||||
return output
|
||||
@dataclass
|
||||
class DiffusionRequestState:
|
||||
sched_req_id: str
|
||||
req: OmniDiffusionRequest
|
||||
status: DiffusionRequestStatus = DiffusionRequestStatus.WAITING
|
||||
```
|
||||
|
||||
**Design Features**:
|
||||
|
||||
- **Broadcast Model**: All workers receive the same request (for tensor parallelism)
|
||||
- **Scheduler-owned ID**: Each `OmniDiffusionRequest` is tracked by an internal `sched_req_id`, separated from public `request_id` values.
|
||||
|
||||
- **Single Response**: Only rank 0 sends results back (avoids duplicate outputs)
|
||||
- **Explicit lifecycle**: Requests move through waiting, running, optional preemption, and terminal states.
|
||||
|
||||
- **Synchronous**: Blocks until result is received (can be made async)
|
||||
- **Centralized error handling**: Completion, abort, and error states are all normalized in the scheduler layer.
|
||||
|
||||
#### 2.3 Singleton Pattern
|
||||
#### 2.3 Shared Bookkeeping in `_BaseScheduler`
|
||||
|
||||
```python
|
||||
class Scheduler:
|
||||
_instance = None
|
||||
|
||||
def __new__(cls, *args, **kwargs):
|
||||
if not cls._instance:
|
||||
cls._instance = super().__new__(cls)
|
||||
return cls._instance
|
||||
|
||||
# Global singleton instance
|
||||
scheduler = Scheduler()
|
||||
class _BaseScheduler(SchedulerInterface):
|
||||
def __init__(self) -> None:
|
||||
self._request_states = {}
|
||||
self._request_id_to_sched_req_id = {}
|
||||
self._waiting = deque()
|
||||
self._running = []
|
||||
self._finished_req_ids = set()
|
||||
self._max_batch_size = 1
|
||||
```
|
||||
|
||||
**Benefits**:
|
||||
**Design Features**:
|
||||
|
||||
- **Single Point of Control**: Ensures consistent state
|
||||
- **Common state storage**: Shared request maps and waiting/running sets live in the base class.
|
||||
|
||||
- **Easy Access**: Global `scheduler` instance accessible everywhere
|
||||
- **Shared cleanup logic**: Request-id registration, finish handling, and state removal are centralized instead of duplicated in each policy.
|
||||
|
||||
- **Resource Management**: Centralized queue management
|
||||
- **Current constraint**: `_max_batch_size` remains `1` because the current engine path is still synchronous request-mode execution.
|
||||
|
||||
#### 2.4 Current `RequestScheduler` Policy
|
||||
|
||||
```python
|
||||
class RequestScheduler(_BaseScheduler):
|
||||
def schedule(self) -> DiffusionSchedulerOutput:
|
||||
# 1. keep existing RUNNING requests in the scheduling result
|
||||
# 2. pull WAITING requests while capacity remains
|
||||
# 3. move newly admitted requests into RUNNING
|
||||
```
|
||||
|
||||
**Behavior**:
|
||||
|
||||
- **FIFO request scheduling**: Waiting requests are promoted in queue order.
|
||||
|
||||
- **Single-request admission**: The current policy only admits one active request at a time.
|
||||
|
||||
- **Executor result feedback**: `update_from_output()` converts executor output into `FINISHED_COMPLETED` or `FINISHED_ERROR` and returns finished scheduler ids.
|
||||
|
||||
#### 2.5 Engine-Driven Execution Loop
|
||||
|
||||
```python
|
||||
sched_req_id = scheduler.add_request(request)
|
||||
while True:
|
||||
sched_output = scheduler.schedule()
|
||||
output = executor.add_req(req)
|
||||
finished_req_ids = scheduler.update_from_output(sched_output, output)
|
||||
```
|
||||
|
||||
**Design Decisions**:
|
||||
|
||||
- **Separation of concerns**: Scheduler manages state and policy; executor handles runtime execution.
|
||||
|
||||
- **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.
|
||||
|
||||
---
|
||||
|
||||
@@ -880,8 +917,9 @@ def initialize_model_parallel(
|
||||
└─> Model-specific transformations
|
||||
|
||||
3. Scheduling
|
||||
└─> scheduler.add_req(requests)
|
||||
└─> Broadcast via MessageQueue to all workers
|
||||
└─> scheduler.add_request(request)
|
||||
└─> scheduler.schedule()
|
||||
└─> DiffusionEngine submits scheduled request to executor.add_req(req)
|
||||
|
||||
4. Worker Execution
|
||||
└─> WorkerProc.worker_busy_loop()
|
||||
@@ -895,8 +933,9 @@ def initialize_model_parallel(
|
||||
└─> vae.decode()
|
||||
|
||||
5. Result Collection
|
||||
└─> Rank 0 sends DiffusionOutput via result queue
|
||||
└─> Scheduler receives and returns
|
||||
└─> Executor returns DiffusionOutput
|
||||
└─> scheduler.update_from_output(...)
|
||||
└─> DiffusionEngine pops finished request state
|
||||
|
||||
6. Post-processing
|
||||
└─> post_process_func(output)
|
||||
|
||||
@@ -10,10 +10,10 @@ The recommended way to use vLLM-Omni on NPU is through the vllm-ascend pre-built
|
||||
```bash
|
||||
# Update the vllm-ascend image
|
||||
# Atlas A2:
|
||||
# export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0rc1
|
||||
# export IMAGE=quay.io/ascend/vllm-ascend:v0.17.0rc1
|
||||
# Atlas A3:
|
||||
# export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0rc1-a3
|
||||
export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0rc1
|
||||
# export IMAGE=quay.io/ascend/vllm-ascend:v0.17.0rc1-a3
|
||||
export IMAGE=quay.io/ascend/vllm-ascend:v0.17.0rc1
|
||||
docker run --rm \
|
||||
--name vllm-omni-npu \
|
||||
--shm-size=1g \
|
||||
@@ -34,20 +34,19 @@ docker run --rm \
|
||||
-it $IMAGE bash
|
||||
|
||||
cd /vllm-workspace/vllm
|
||||
git pull origin main
|
||||
git fetch origin --tags
|
||||
git checkout v0.16.0
|
||||
git checkout v0.18.0
|
||||
|
||||
# Because vllm-ascend will release v0.16.0rc1 after vllm-omni 0.16.0,
|
||||
# Because vllm-ascend will release v0.18.0rc1 after vllm-omni 0.16.0,
|
||||
# we have to pin vllm-ascend at the current commit.
|
||||
cd /vllm-workspace/vllm-ascend
|
||||
git pull origin main
|
||||
git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
|
||||
git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
|
||||
pip install -v -e .
|
||||
|
||||
# Inside the container, install vLLM-Omni from source
|
||||
cd /vllm-workspace
|
||||
git clone -b v0.16.0 https://github.com/vllm-project/vllm-omni.git
|
||||
git clone -b v0.18.0 https://github.com/vllm-project/vllm-omni.git
|
||||
cd vllm-omni
|
||||
pip install -v -e . --no-build-isolation
|
||||
# or VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e .
|
||||
@@ -70,21 +69,19 @@ You can also build vLLM-Omni from the latest main branch if you want to use the
|
||||
```bash
|
||||
# Pin vLLM version to 0.18.0
|
||||
cd /vllm-workspace/vllm
|
||||
git pull origin main
|
||||
git fetch origin --tags
|
||||
git checkout v0.18.0
|
||||
VLLM_TARGET_DEVICE=empty pip install -v -e .
|
||||
|
||||
# Because vllm-ascend has not yet entered continuous development and has not been officially released, we need to pin it to a specific commit. Please note that this commit may change over time.
|
||||
cd /vllm-workspace/vllm-ascend
|
||||
git pull origin main
|
||||
git fetch origin --tags
|
||||
git checkout 1e05c4908f31737bc4eef865a9f351d030a77c9d
|
||||
git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
|
||||
pip install -v -e .
|
||||
|
||||
# Install vLLM-Omni from the latest main branch
|
||||
cd /vllm-workspace/vllm-omni
|
||||
git clone https://github.com/vllm-project/vllm-omni.git
|
||||
cd /vllm-workspace/vllm-omni
|
||||
pip install -v -e . --no-build-isolation
|
||||
# or VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e .
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
@@ -107,10 +104,10 @@ Here's an example deployment command that has been verified on 4 x NPUs:
|
||||
|
||||
```bash
|
||||
# Atlas A2:
|
||||
# export IMAGE=quay.io/ascend/vllm-omni:v0.16.0
|
||||
# export IMAGE=quay.io/ascend/vllm-omni:v0.18.0
|
||||
# Atlas A3:
|
||||
# export IMAGE=quay.io/ascend/vllm-omni:v0.16.0-a3
|
||||
export IMAGE=quay.io/ascend/vllm-omni:v0.16.0
|
||||
# export IMAGE=quay.io/ascend/vllm-omni:v0.18.0-a3
|
||||
export IMAGE=quay.io/ascend/vllm-omni:v0.18.0
|
||||
docker run --rm \
|
||||
--name vllm-omni-npu \
|
||||
--shm-size=1g \
|
||||
|
||||
@@ -139,6 +139,8 @@ def scan_package(package_name: str = "vllm_omni") -> dict[str, list[str]]:
|
||||
# Skip excluded modules (avoid importing vllm during docs build)
|
||||
excluded_prefixes = [
|
||||
"vllm_omni.diffusion.models.qwen_image",
|
||||
"vllm_omni.diffusion.quantization",
|
||||
"vllm_omni.quantization",
|
||||
"vllm_omni.entrypoints.async_diffusion",
|
||||
"vllm_omni.entrypoints.openai",
|
||||
"vllm_omni.model_executor.models.voxtral_tts.configuration_voxtral_tts",
|
||||
|
||||
@@ -6,7 +6,7 @@ vLLM-Omni supports unified multimodal comprehension and generation models across
|
||||
|
||||
If vLLM-Omni natively supports a model, its implementation can be found in <gh-file:vllm_omni/model_executor/models> and <gh-file:vllm_omni/diffusion/models>.
|
||||
|
||||
## List of Supported Models for Nvidia GPU / AMD GPU
|
||||
## List of Supported Models
|
||||
|
||||
<style>
|
||||
th {
|
||||
@@ -15,74 +15,46 @@ th {
|
||||
}
|
||||
</style>
|
||||
|
||||
| Architecture | Models | Example HF Models |
|
||||
|--------------|--------|-------------------|
|
||||
| `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` |
|
||||
| `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` |
|
||||
| `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` |
|
||||
| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2509 | `Qwen/Qwen-Image-Edit-2509` |
|
||||
| `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` |
|
||||
| `GlmImagePipeline` | GLM-Image | `zai-org/GLM-Image` |
|
||||
|`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` |
|
||||
| `LTX2Pipeline` | LTX-2-T2V | `Lightricks/LTX-2` |
|
||||
| `LTX2ImageToVideoPipeline` | LTX-2-I2V | `Lightricks/LTX-2` |
|
||||
| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` |
|
||||
| `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` |
|
||||
|`LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` |
|
||||
|`LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` |
|
||||
|`StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` |
|
||||
|`CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` |
|
||||
|`MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` |
|
||||
|`Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` |
|
||||
|`FluxPipeline` | FLUX.1-dev | `black-forest-labs/FLUX.1-dev` |
|
||||
|`OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` |
|
||||
|`StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` |
|
||||
|`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` |
|
||||
|`GlmImageForConditionalGeneration` | GLM-Image | `zai-org/GLM-Image` |
|
||||
|`NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` |
|
||||
|`MiMoAudioForConditionalGeneration` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` |
|
||||
|`Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` |
|
||||
|`FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` |
|
||||
|`DreamIDOmniPipeline`| DreamID-Omni | `XuGuo699/DreamID-Omni` |
|
||||
| `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` |
|
||||
| `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` |
|
||||
|`VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` |
|
||||
| Architecture | Models | Example HF Models | NVIDIA GPU | AMD GPU | Ascend NPU | Intel GPU |
|
||||
|--------------|--------|-------------------|------------|---------|-----|-----------|
|
||||
| `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` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2509 | `Qwen/Qwen-Image-Edit-2509` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2511 | `Qwen/Qwen-Image-Edit-2511` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `GlmImagePipeline` | GLM-Image | `zai-org/GLM-Image` | ✅︎ | ✅︎ | | |
|
||||
| `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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `LTX2Pipeline` | LTX-2-T2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
|
||||
| `LTX2ImageToVideoPipeline` | LTX-2-I2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
|
||||
| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | |
|
||||
| `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | |
|
||||
| `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | |
|
||||
| `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `FluxKontextPipeline` | FLUX.1-Kontext-dev | `black-forest-labs/FLUX.1-Kontext-dev` | ✅︎ | ✅︎ | | |
|
||||
| `FluxPipeline` | FLUX.1-dev | `black-forest-labs/FLUX.1-dev` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `MiMoAudioForConditionalGeneration` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` | ✅︎ | ✅︎ | | |
|
||||
| `Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | ✅︎ | ✅︎ | | |
|
||||
| `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | |
|
||||
| `DreamIDOmniPipeline` | DreamID-Omni | `XuGuo699/DreamID-Omni` | ✅︎ | ✅︎ | | |
|
||||
| `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | |
|
||||
| `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | |
|
||||
| `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | |
|
||||
|
||||
|
||||
## List of Supported Models for NPU
|
||||
|
||||
<style>
|
||||
th {
|
||||
white-space: nowrap;
|
||||
min-width: 0 !important;
|
||||
}
|
||||
</style>
|
||||
|
||||
| Architecture | Models | Example HF Models |
|
||||
|--------------|--------|-------------------|
|
||||
| `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` |
|
||||
| `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-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` |
|
||||
| `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` |
|
||||
| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2509 | `Qwen/Qwen-Image-Edit-2509` |
|
||||
| `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` |
|
||||
| `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2511 | `Qwen/Qwen-Image-Edit-2511` |
|
||||
|`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` |
|
||||
|`LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` |
|
||||
|`LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` |
|
||||
|`Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` |
|
||||
|`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` |
|
||||
✅︎ indicates the model is supported on that backend. Empty cells mean not listed as supported on that backend.
|
||||
@@ -0,0 +1,78 @@
|
||||
# Diffusion Chat Completions API
|
||||
|
||||
vLLM-Omni supports generating and editing images via the `/v1/chat/completions`
|
||||
endpoint using diffusion models. This page explains how to pass generation
|
||||
parameters (such as `num_inference_steps`, `height`, `width`) to diffusion
|
||||
models through this endpoint.
|
||||
|
||||
!!! tip
|
||||
For dedicated endpoints that accept generation parameters as top-level
|
||||
fields, see [Image Generation API](image_generation_api.md) and
|
||||
[Image Edit API](image_edit_api.md).
|
||||
|
||||
## Passing Generation Parameters
|
||||
|
||||
The `/v1/chat/completions` endpoint follows the OpenAI Chat API schema, which
|
||||
does not natively include diffusion-specific fields like `num_inference_steps`
|
||||
or `height`. How you pass these extra fields depends on your client.
|
||||
|
||||
### curl / Python `requests`
|
||||
|
||||
Wrap generation parameters inside an `"extra_body"` key in the JSON body:
|
||||
|
||||
```bash
|
||||
curl -s http://localhost:8091/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"messages": [
|
||||
{"role": "user", "content": "A beautiful landscape painting"}
|
||||
],
|
||||
"extra_body": {
|
||||
"num_inference_steps": 50,
|
||||
"seed": 42
|
||||
}
|
||||
}'
|
||||
```
|
||||
|
||||
### OpenAI Python SDK
|
||||
|
||||
Use the `extra_body` **keyword argument**. The SDK automatically merges these
|
||||
fields into the top-level request body:
|
||||
|
||||
```python
|
||||
response = client.chat.completions.create(
|
||||
model="Qwen/Qwen-Image",
|
||||
messages=[{"role": "user", "content": "A beautiful landscape painting"}],
|
||||
extra_body={
|
||||
"num_inference_steps": 50,
|
||||
"seed": 42,
|
||||
},
|
||||
)
|
||||
```
|
||||
|
||||
!!! note "SDK `extra_body` vs. JSON `extra_body`"
|
||||
These two `extra_body` usages look similar but work differently under the
|
||||
hood. The SDK flattens the dict into the top-level request JSON, while the
|
||||
curl/requests approach sends it as a nested `"extra_body"` key. Both are
|
||||
handled correctly by the server.
|
||||
|
||||
!!! note "About the `ignored fields` warning"
|
||||
You may see a log message like:
|
||||
|
||||
```
|
||||
WARNING: The following fields were present in the request but ignored: {'height', 'width', ...}
|
||||
```
|
||||
|
||||
This is **harmless**. It is emitted by vLLM's request validation layer
|
||||
because these fields are not part of the standard OpenAI
|
||||
`ChatCompletionRequest` schema. The fields are still stored internally
|
||||
and correctly forwarded to the diffusion pipeline.
|
||||
|
||||
## Model-Specific Examples
|
||||
|
||||
For complete examples with full request/response details, see the model-specific
|
||||
guides:
|
||||
|
||||
- [Text-to-Image (Qwen-Image)](../user_guide/examples/online_serving/text_to_image.md)
|
||||
- [Image-to-Image (Qwen-Image-Edit, Qwen-Image-Layered)](../user_guide/examples/online_serving/image_to_image.md)
|
||||
- [GLM-Image](../user_guide/examples/online_serving/glm_image.md)
|
||||
+174
-4
@@ -160,6 +160,7 @@ Upload a new voice sample for voice cloning in Base task TTS requests.
|
||||
| `audio_sample` | file | Yes | Audio file (max 10MB, supported formats: wav, mp3, flac, ogg, aac, webm, mp4) |
|
||||
| `consent` | string | Yes | Consent recording ID |
|
||||
| `name` | string | Yes | Name for the new voice |
|
||||
| `ref_text` | string | No | Transcript of the audio. When provided, enables in-context voice cloning (higher quality). Without it, only the speaker embedding is extracted. |
|
||||
|
||||
**Response Example:**
|
||||
|
||||
@@ -182,7 +183,8 @@ Upload a new voice sample for voice cloning in Base task TTS requests.
|
||||
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 "name=custom_voice_1" \
|
||||
-F "ref_text=The exact transcript of the audio sample."
|
||||
```
|
||||
|
||||
## Streaming Text Input (WebSocket)
|
||||
@@ -317,7 +319,9 @@ curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
}' --output cloned.wav
|
||||
```
|
||||
|
||||
upload voice
|
||||
### Upload Voice
|
||||
|
||||
Upload voice (speaker embedding only):
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/voices \
|
||||
-F "audio_sample=@/path/to/voice_sample.wav" \
|
||||
@@ -325,17 +329,183 @@ curl -X POST http://localhost:8091/v1/audio/voices \
|
||||
-F "name=custom_voice_1"
|
||||
```
|
||||
|
||||
use upload voice
|
||||
Upload voice with transcript (in-context cloning, higher quality):
|
||||
```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_2" \
|
||||
-F "ref_text=The exact transcript of the audio sample."
|
||||
```
|
||||
|
||||
### Use Uploaded Voice
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "Hello, this is a cloned voice",
|
||||
"task_type": "Base",
|
||||
"voice": "custom_voice_1"
|
||||
}' --output cloned.wav
|
||||
```
|
||||
|
||||
## Batch Speech Generation
|
||||
|
||||
The batch endpoint synthesizes multiple texts in a single request, returning all results as JSON with base64-encoded audio.
|
||||
|
||||
### Endpoint
|
||||
|
||||
```
|
||||
POST /v1/audio/speech/batch
|
||||
Content-Type: application/json
|
||||
```
|
||||
|
||||
### Request Parameters
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| `items` | array | **required** | List of items to synthesize (1–32) |
|
||||
| `model` | string | server's model | Model to use |
|
||||
| `voice` | string | null | Default voice for all items |
|
||||
| `response_format` | string | "wav" | Default audio format for all items |
|
||||
| `speed` | float | 1.0 | Default playback speed (0.25–4.0) |
|
||||
| `task_type` | string | null | Default TTS task type |
|
||||
| `language` | string | null | Default language |
|
||||
| `instructions` | string | null | Default voice style instructions |
|
||||
| `ref_audio` | string | null | Default reference audio (Base task) |
|
||||
| `ref_text` | string | null | Default reference transcript (Base task) |
|
||||
| `max_new_tokens` | integer | null | Default max tokens |
|
||||
|
||||
Each item in the `items` array requires only `input` (the text). All other fields are optional and override the batch-level defaults when set:
|
||||
|
||||
| Field | Type | Description |
|
||||
|-------|------|-------------|
|
||||
| `input` | string | **required** — text to synthesize |
|
||||
| `voice` | string | Override voice for this item |
|
||||
| `response_format` | string | Override format for this item |
|
||||
| `speed` | float | Override speed for this item |
|
||||
| `task_type` | string | Override task type |
|
||||
| `language` | string | Override language |
|
||||
| `instructions` | string | Override instructions |
|
||||
| `ref_audio` | string | Override reference audio |
|
||||
| `ref_text` | string | Override reference transcript |
|
||||
| `max_new_tokens` | integer | Override max tokens |
|
||||
|
||||
### Response Format
|
||||
|
||||
```json
|
||||
{
|
||||
"id": "speech-batch-abc123",
|
||||
"results": [
|
||||
{
|
||||
"index": 0,
|
||||
"status": "success",
|
||||
"audio_data": "<base64-encoded audio>",
|
||||
"media_type": "audio/wav"
|
||||
},
|
||||
{
|
||||
"index": 1,
|
||||
"status": "error",
|
||||
"error": "Input text cannot be empty"
|
||||
}
|
||||
],
|
||||
"total": 2,
|
||||
"succeeded": 1,
|
||||
"failed": 1
|
||||
}
|
||||
```
|
||||
|
||||
### Examples
|
||||
|
||||
**Basic batch with shared defaults:**
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/speech/batch \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"items": [
|
||||
{"input": "Hello, how are you?"},
|
||||
{"input": "Goodbye, see you later!"}
|
||||
],
|
||||
"voice": "vivian",
|
||||
"language": "English"
|
||||
}'
|
||||
```
|
||||
|
||||
**Per-item overrides (different voices and formats):**
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/speech/batch \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"items": [
|
||||
{"input": "Hello!", "voice": "vivian", "response_format": "mp3"},
|
||||
{"input": "你好!", "voice": "ryan", "language": "Chinese"}
|
||||
],
|
||||
"response_format": "wav"
|
||||
}'
|
||||
```
|
||||
|
||||
**Voice cloning with shared reference audio (Base task):**
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/speech/batch \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"items": [
|
||||
{"input": "First sentence in the cloned voice."},
|
||||
{"input": "Second sentence in the cloned voice."}
|
||||
],
|
||||
"task_type": "Base",
|
||||
"ref_audio": "https://example.com/reference.wav",
|
||||
"ref_text": "Transcript of the reference audio"
|
||||
}'
|
||||
```
|
||||
|
||||
Setting `ref_audio` at the batch level applies it to all items, avoiding the need to repeat it per item.
|
||||
|
||||
**Decoding the response in Python:**
|
||||
|
||||
```python
|
||||
import base64
|
||||
import httpx
|
||||
|
||||
response = httpx.post(
|
||||
"http://localhost:8091/v1/audio/speech/batch",
|
||||
json={
|
||||
"items": [
|
||||
{"input": "First sentence."},
|
||||
{"input": "Second sentence."},
|
||||
],
|
||||
"voice": "vivian",
|
||||
},
|
||||
timeout=300.0,
|
||||
)
|
||||
|
||||
for result in response.json()["results"]:
|
||||
if result["status"] == "success":
|
||||
audio_bytes = base64.b64decode(result["audio_data"])
|
||||
with open(f"output_{result['index']}.wav", "wb") as f:
|
||||
f.write(audio_bytes)
|
||||
```
|
||||
|
||||
### Configuration
|
||||
|
||||
| Parameter | Source | Default | Description |
|
||||
|-----------|--------|---------|-------------|
|
||||
| `tts_batch_max_items` | engine kwarg | 32 | Maximum number of items per batch request |
|
||||
|
||||
All items are fanned out to `generate()` concurrently. The engine's stage worker automatically batches them up to the configured `max_batch_size` and queues the rest — no client-side throttling needed.
|
||||
|
||||
For best throughput, use a batch-optimized stage config with `max_batch_size > 1`:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts_batch.yaml \
|
||||
--omni --port 8091 --trust-remote-code --enforce-eager
|
||||
```
|
||||
|
||||
The default `qwen3_tts.yaml` uses `max_batch_size: 1` (single request). The `qwen3_tts_batch.yaml` config sets `max_batch_size: 4` for ~4x throughput.
|
||||
|
||||
## Supported Models
|
||||
|
||||
### Qwen3-TTS
|
||||
|
||||
@@ -39,7 +39,10 @@ The following table shows which models are currently supported by parallelism me
|
||||
| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ❌ | ✅ | ✅ | ❌ | N/A | ✅ |
|
||||
| **FLUX.2-dev** | `black-forest-labs/FLUX.2-dev` | ❌ | ❌ | ❌ | ✅ | ❌ | N/A | ✅ |
|
||||
| **HunyuanImage3.0** | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ | ❌ |
|
||||
| **Bagel** | `ByteDance-Seed/BAGEL-7B-MoT` | ✅ | ✅ | ✅ | ✅ | ❌ | N/A | ❌ |
|
||||
| **DreamID-Omni** | `XuGuo699/DreamID-Omni` | ❌ | ❌ | ✅ | ❌ | ❌ | N/A | ❌ |
|
||||
| **FLUX.1-Kontext-dev** | `black-forest-labs/FLUX.1-Kontext-dev` | ❌ | ❌ | ❌ | ✅ | ❌ | N/A | ✅ |
|
||||
| **OmniGen2** | `OmniGen2/OmniGen2` | ❌ | ❌ | ❌ | ✅ | ❌ | N/A | ❌ |
|
||||
|
||||
!!! note "TP Limitations for Diffusion Models"
|
||||
We currently implement Tensor Parallelism (TP) only for the DiT (Diffusion Transformer) blocks. This is because the `text_encoder` component in vLLM-Omni uses the original Transformers implementation, which does not yet support TP.
|
||||
@@ -131,6 +134,11 @@ VAE patch parallelism automatically selects between two internal decode methods
|
||||
|
||||
#### Ulysses-SP
|
||||
|
||||
!!! note "Experimental UAA mode"
|
||||
`ulysses_mode="advanced_uaa"` is an experimental extension to Ulysses-SP. It lets Ulysses attention handle arbitrary sequence lengths and arbitrary attention head counts without relying on `attention_mask`-based token padding.
|
||||
|
||||
In hybrid Ulysses + Ring mode, Ring still requires every rank in the same ring group to observe the same post-Ulysses sequence length. If that condition is not met, vLLM-Omni raises a validation error instead of entering the ring kernel with inconsistent shapes.
|
||||
|
||||
##### Offline Inference
|
||||
|
||||
An example of offline inference script using [Ulysses-SP](https://arxiv.org/pdf/2309.14509) is shown below:
|
||||
@@ -153,6 +161,18 @@ outputs = omni.generate(
|
||||
|
||||
See `examples/offline_inference/text_to_image/text_to_image.py` for a complete working example.
|
||||
|
||||
To enable the experimental UAA mode explicitly, use a model/configuration that actually requires it. For example, `Tongyi-MAI/Z-Image-Turbo` has 30 attention heads, so `ulysses_degree=4` requires UAA because 30 is not divisible by 4:
|
||||
|
||||
```python
|
||||
omni = Omni(
|
||||
model="Tongyi-MAI/Z-Image-Turbo",
|
||||
parallel_config=DiffusionParallelConfig(
|
||||
ulysses_degree=4,
|
||||
ulysses_mode="advanced_uaa",
|
||||
),
|
||||
)
|
||||
```
|
||||
|
||||
##### Online Serving
|
||||
|
||||
You can enable Ulysses-SP in online serving for diffusion models via `--usp`:
|
||||
@@ -160,6 +180,9 @@ You can enable Ulysses-SP in online serving for diffusion models via `--usp`:
|
||||
```bash
|
||||
# Text-to-image (requires >= 2 GPUs)
|
||||
vllm serve Qwen/Qwen-Image --omni --port 8091 --usp 2
|
||||
|
||||
# Experimental UAA mode for a model with 30 attention heads
|
||||
vllm serve Tongyi-MAI/Z-Image-Turbo --omni --port 8091 --usp 4 --ulysses-mode advanced_uaa
|
||||
```
|
||||
|
||||
##### Benchmarks
|
||||
@@ -240,6 +263,9 @@ To measure the parallelism methods, we run benchmarks with **Qwen/Qwen-Image** m
|
||||
|
||||
You can combine both Ulysses-SP and Ring-Attention for larger scale parallelism. The total sequence parallel size equals `ulysses_degree × ring_degree`.
|
||||
|
||||
!!! note "Experimental UAA in hybrid mode"
|
||||
`ulysses_mode="advanced_uaa"` can also be used with hybrid Ulysses + Ring, but this does not remove Ring's shape requirement. Every rank in the same ring group must still have the same post-Ulysses sequence length.
|
||||
|
||||
##### Offline Inference
|
||||
|
||||
```python
|
||||
@@ -250,7 +276,10 @@ from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
# Hybrid: 2 Ulysses × 2 Ring = 4 GPUs total
|
||||
omni = Omni(
|
||||
model="Qwen/Qwen-Image",
|
||||
parallel_config=DiffusionParallelConfig(ulysses_degree=2, ring_degree=2)
|
||||
parallel_config=DiffusionParallelConfig(
|
||||
ulysses_degree=2,
|
||||
ring_degree=2,
|
||||
)
|
||||
)
|
||||
|
||||
outputs = omni.generate(
|
||||
|
||||
@@ -63,6 +63,7 @@ The available `ignored_layers` names depend on the model architecture (e.g., `to
|
||||
| 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 |
|
||||
|
||||
## Combining with Other Features
|
||||
|
||||
@@ -6,13 +6,13 @@
|
||||
3. Keep user-facing knobs minimal and consistent across offline and online flows.
|
||||
|
||||
## Scope
|
||||
1. Models: Z-Image, and Flux2-klein.
|
||||
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.
|
||||
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.
|
||||
@@ -124,19 +124,21 @@ x @ weight.T
|
||||
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).
|
||||
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. `get_gguf_adapter(...)` strictly selects by model class/config; unsupported models raise an error (no fallback adapter).
|
||||
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)
|
||||
|
||||
@@ -195,3 +197,31 @@ curl -X POST http://localhost:8000/v1/images/generations \
|
||||
"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,22 +1,86 @@
|
||||
# Quantization for Diffusion Transformers
|
||||
# Quantization
|
||||
|
||||
vLLM-Omni supports quantization of DiT linear layers to reduce memory usage and accelerate inference.
|
||||
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
|
||||
|
||||
| Method | Guide |
|
||||
|--------|-------|
|
||||
| FP8 | [FP8](fp8.md) |
|
||||
| Int8 | [Int8](int8.md) |
|
||||
| GGUF | [GGUF](gguf.md) |
|
||||
### 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 |
|
||||
|
||||
### 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
|
||||
|
||||
| GPU Generation | Example GPUs | FP8 Mode |
|
||||
|---------------|-------------------|----------|
|
||||
| Ada/Hopper (SM 89+) | RTX 4090, H100, H200 | Full W8A8 with native hardware |
|
||||
| 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.
|
||||
Kernel selection is automatic on CUDA GPUs.
|
||||
|
||||
## Device Compatibility for Int8
|
||||
|
||||
@@ -24,3 +88,62 @@ Kernel selection is automatic.
|
||||
|-------------|---------------|-------------------|----------|
|
||||
| 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"})
|
||||
|
||||
# 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,
|
||||
})
|
||||
```
|
||||
@@ -66,11 +66,13 @@ The following table shows which models are currently supported by each accelerat
|
||||
| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
|
||||
| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ (TP=2 only) | ✅ |
|
||||
| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ |
|
||||
| **Bagel** | `ByteDance-Seed/BAGEL-7B-MoT` | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **Bagel** | `ByteDance-Seed/BAGEL-7B-MoT` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
|
||||
| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ |
|
||||
| **NextStep-1.1** | `stepfun-ai/NextStep-1.1` | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ |
|
||||
| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ |
|
||||
| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ |
|
||||
| **FLUX.2-dev** | `black-forest-labs/FLUX.2-dev` | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ |
|
||||
| **FLUX.1-Kontext-dev** | `black-forest-labs/FLUX.1-Kontext-dev` | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ |
|
||||
| **GLM-Image** | `zai-org/GLM-Image` | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ |
|
||||
|
||||
### VideoGen
|
||||
|
||||
@@ -91,6 +93,12 @@ The following table shows which models are currently supported by each accelerat
|
||||
| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ |
|
||||
|
||||
|
||||
### AudioGen
|
||||
|
||||
| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention | CFG-Parallel |
|
||||
|------------------------|------------------------------------------|:--------:|:---------:|:----------:|:--------------:|:------------:|
|
||||
| **Stable-Audio-Open** | `stabilityai/stable-audio-open-1.0` | ✅ | | | | |
|
||||
|
||||
## Performance Benchmarks
|
||||
|
||||
The following benchmarks were measured on **Qwen/Qwen-Image** and **Qwen/Qwen-Image-Edit** models generating 1024x1024 images with 50 inference steps:
|
||||
|
||||
@@ -79,19 +79,7 @@ Each DAC decoder chunk is logged as it arrives. The final WAV file is written on
|
||||
|
||||
## Example materials
|
||||
|
||||
??? abstract "diag_codebook.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/fish_speech/diag_codebook.py"
|
||||
``````
|
||||
??? abstract "diag_reference_compare.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/fish_speech/diag_reference_compare.py"
|
||||
``````
|
||||
??? abstract "end2end.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/fish_speech/end2end.py"
|
||||
``````
|
||||
??? abstract "standalone_test.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/fish_speech/standalone_test.py"
|
||||
``````
|
||||
@@ -43,8 +43,7 @@ GLM-Image uses a 2-stage pipeline:
|
||||
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model-path /path/to/glm-image \
|
||||
--config-path ../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
|
||||
--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 \
|
||||
@@ -55,8 +54,7 @@ python end2end.py \
|
||||
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model-path /path/to/glm-image \
|
||||
--config-path ../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
|
||||
--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
|
||||
@@ -67,7 +65,7 @@ python end2end.py \
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model-path /path/to/glm-image \
|
||||
--config-path ../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
|
||||
--config-path ../../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
|
||||
--prompt "A photorealistic cat sitting on a window sill" \
|
||||
--height 1024 \
|
||||
--width 1024 \
|
||||
|
||||
@@ -0,0 +1,203 @@
|
||||
# Helios Video Generation
|
||||
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/helios>.
|
||||
|
||||
|
||||
Helios is a text-to-video (T2V), image-to-video (I2V), and video-to-video (V2V) diffusion model. This example demonstrates end-to-end video generation using vLLM-Omni with three model variants:
|
||||
|
||||
| Variant | Description | Key Features |
|
||||
|---------|-------------|--------------|
|
||||
| **Helios-Base** | Base model, Stage 1 only | Single-stage denoising, `guidance_scale=5.0` |
|
||||
| **Helios-Mid** | Mid model, Stage 2 pyramid | Multi-stage pyramid denoising, CFG-Zero* support |
|
||||
| **Helios-Distilled** | Distilled model, Stage 2+3 | Few-step inference with DMD, `guidance_scale=1.0` |
|
||||
|
||||
## Setup
|
||||
|
||||
Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
|
||||
|
||||
## Run Examples
|
||||
|
||||
Get into the example folder:
|
||||
```bash
|
||||
cd examples/offline_inference/helios
|
||||
```
|
||||
|
||||
### Text-to-Video (T2V)
|
||||
|
||||
**Helios-Base** (Stage 1 only):
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Base \
|
||||
--sample-type t2v \
|
||||
--prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \
|
||||
--guidance-scale 5.0 \
|
||||
--output helios_t2v_base.mp4
|
||||
```
|
||||
|
||||
**Helios-Mid** (Stage 2 + CFG-Zero*):
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Mid \
|
||||
--sample-type t2v \
|
||||
--prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \
|
||||
--guidance-scale 5.0 \
|
||||
--is-enable-stage2 \
|
||||
--pyramid-num-inference-steps-list 20 20 20 \
|
||||
--use-cfg-zero-star --use-zero-init --zero-steps 1 \
|
||||
--output helios_t2v_mid.mp4
|
||||
```
|
||||
|
||||
**Helios-Distilled** (Stage 2 pyramid + DMD):
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Distilled \
|
||||
--sample-type t2v \
|
||||
--prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \
|
||||
--num-frames 240 \
|
||||
--guidance-scale 1.0 \
|
||||
--is-enable-stage2 \
|
||||
--pyramid-num-inference-steps-list 2 2 2 \
|
||||
--is-amplify-first-chunk \
|
||||
--output helios_t2v_distilled.mp4
|
||||
```
|
||||
|
||||
### Image-to-Video (I2V)
|
||||
|
||||
**Helios-Base**:
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Base \
|
||||
--sample-type i2v \
|
||||
--image-path /path/to/image.jpg \
|
||||
--prompt "A towering emerald wave surges forward, its crest curling with raw power and energy." \
|
||||
--guidance-scale 5.0 \
|
||||
--output helios_i2v_base.mp4
|
||||
```
|
||||
|
||||
**Helios-Mid** (Stage 2 + CFG-Zero*):
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Mid \
|
||||
--sample-type i2v \
|
||||
--image-path /path/to/image.jpg \
|
||||
--prompt "A towering emerald wave surges forward, its crest curling with raw power and energy." \
|
||||
--guidance-scale 5.0 \
|
||||
--is-enable-stage2 \
|
||||
--pyramid-num-inference-steps-list 20 20 20 \
|
||||
--use-cfg-zero-star --use-zero-init --zero-steps 1 \
|
||||
--output helios_i2v_mid.mp4
|
||||
```
|
||||
|
||||
**Helios-Distilled**:
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Distilled \
|
||||
--sample-type i2v \
|
||||
--image-path /path/to/image.jpg \
|
||||
--prompt "A towering emerald wave surges forward, its crest curling with raw power and energy." \
|
||||
--num-frames 240 \
|
||||
--guidance-scale 1.0 \
|
||||
--is-enable-stage2 \
|
||||
--pyramid-num-inference-steps-list 2 2 2 \
|
||||
--is-amplify-first-chunk \
|
||||
--output helios_i2v_distilled.mp4
|
||||
```
|
||||
|
||||
### Video-to-Video (V2V)
|
||||
|
||||
**Helios-Base**:
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Base \
|
||||
--sample-type v2v \
|
||||
--video-path /path/to/video.mp4 \
|
||||
--prompt "A bright yellow Lamborghini speeds along a curving mountain road." \
|
||||
--guidance-scale 5.0 \
|
||||
--output helios_v2v_base.mp4
|
||||
```
|
||||
|
||||
**Helios-Mid** (Stage 2 + CFG-Zero*):
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Mid \
|
||||
--sample-type v2v \
|
||||
--video-path /path/to/video.mp4 \
|
||||
--prompt "A bright yellow Lamborghini speeds along a curving mountain road." \
|
||||
--guidance-scale 5.0 \
|
||||
--is-enable-stage2 \
|
||||
--pyramid-num-inference-steps-list 20 20 20 \
|
||||
--use-cfg-zero-star --use-zero-init --zero-steps 1 \
|
||||
--output helios_v2v_mid.mp4
|
||||
```
|
||||
|
||||
**Helios-Distilled**:
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Distilled \
|
||||
--sample-type v2v \
|
||||
--video-path /path/to/video.mp4 \
|
||||
--prompt "A bright yellow Lamborghini speeds along a curving mountain road." \
|
||||
--num-frames 240 \
|
||||
--guidance-scale 1.0 \
|
||||
--is-enable-stage2 \
|
||||
--pyramid-num-inference-steps-list 2 2 2 \
|
||||
--is-amplify-first-chunk \
|
||||
--output helios_v2v_distilled.mp4
|
||||
```
|
||||
|
||||
## Common Parameters
|
||||
|
||||
| Parameter | Default | Description |
|
||||
|-----------|---------|-------------|
|
||||
| `--model` | `BestWishYsh/Helios-Base` | Model ID or local path |
|
||||
| `--sample-type` | `t2v` | Generation mode: `t2v`, `i2v`, or `v2v` |
|
||||
| `--prompt` | — | Text prompt describing the video |
|
||||
| `--negative-prompt` | *(see source)* | Negative prompt for CFG (includes anti-static terms) |
|
||||
| `--image-path` | — | Input image (required for `i2v`) |
|
||||
| `--video-path` | — | Input video (required for `v2v`) |
|
||||
| `--height` | `384` | Video height in pixels |
|
||||
| `--width` | `640` | Video width in pixels |
|
||||
| `--num-frames` | `99` | Number of output frames |
|
||||
| `--num-inference-steps` | `50` | Denoising steps (Stage 1 only) |
|
||||
| `--guidance-scale` | `5.0` | CFG scale (`1.0` for Distilled) |
|
||||
| `--seed` | `42` | Random seed |
|
||||
| `--fps` | `16` | Output video frame rate |
|
||||
| `--output` | `helios_output.mp4` | Output file path |
|
||||
|
||||
### Stage 2 / Pyramid Parameters
|
||||
|
||||
| Parameter | Default | Description |
|
||||
|-----------|---------|-------------|
|
||||
| `--is-enable-stage2` | off | Enable pyramid multi-stage denoising |
|
||||
| `--pyramid-num-stages` | `3` | Number of pyramid stages |
|
||||
| `--pyramid-num-inference-steps-list` | `10 10 10` | Steps per pyramid stage |
|
||||
| `--is-amplify-first-chunk` | off | DMD amplification (Distilled) |
|
||||
|
||||
### CFG-Zero* Parameters
|
||||
|
||||
| Parameter | Default | Description |
|
||||
|-----------|---------|-------------|
|
||||
| `--use-cfg-zero-star` | off | Enable CFG-Zero* guidance (Mid) |
|
||||
| `--use-zero-init` | off | Zero init for first steps |
|
||||
| `--zero-steps` | `1` | Number of zero-init steps |
|
||||
|
||||
### Memory & Parallelism
|
||||
|
||||
| Parameter | Default | Description |
|
||||
|-----------|---------|-------------|
|
||||
| `--vae-use-slicing` | off | Enable VAE slicing |
|
||||
| `--vae-use-tiling` | off | Enable VAE tiling |
|
||||
| `--enforce-eager` | off | Disable torch.compile |
|
||||
| `--enable-cpu-offload` | off | CPU offloading |
|
||||
| `--enable-layerwise-offload` | off | Layerwise offloading |
|
||||
| `--ulysses-degree` | `1` | Ulysses SP degree |
|
||||
| `--ring-degree` | `1` | Ring SP degree |
|
||||
| `--cfg-parallel-size` | `1` | CFG parallel size (1 or 2) |
|
||||
| `--tensor-parallel-size` | `1` | Tensor parallelism size |
|
||||
|
||||
## Example materials
|
||||
|
||||
??? abstract "end2end.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/helios/end2end.py"
|
||||
``````
|
||||
@@ -0,0 +1,35 @@
|
||||
# HunyuanImage-3.0 Image-to-Text Inference
|
||||
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/hunyuan_image3>.
|
||||
|
||||
|
||||
This example demonstrates how to run HunyuanImage-3.0 Image-to-Text with the vLLM-Omni.
|
||||
|
||||
## Local CLI Usage
|
||||
|
||||
Download the example image:
|
||||
|
||||
```bash
|
||||
wget https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg
|
||||
```
|
||||
|
||||
Run example:
|
||||
|
||||
```bash
|
||||
python image_to_text.py \
|
||||
--image cherry_blossom.jpg \
|
||||
--prompt "<|startoftext|>You are an assistant that understands images and outputs text.<img>Describe the content of the picture."
|
||||
```
|
||||
|
||||
Key arguments:
|
||||
|
||||
- `--model`: Model used. Default is: tencent/HunyuanImage-3.0-Instruct (Optional).
|
||||
- `--image`: Path to input image (required).
|
||||
- `--prompt`: Text description used to guide image understanding (required).
|
||||
|
||||
## Example materials
|
||||
|
||||
??? abstract "image_to_text.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/hunyuan_image3/image_to_text.py"
|
||||
``````
|
||||
@@ -68,6 +68,7 @@ Key arguments:
|
||||
- `--vae-use-slicing`: Enable VAE slicing for memory optimization.
|
||||
- `--vae-use-tiling`: Enable VAE tiling for memory optimization.
|
||||
- `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](https://github.com/vllm-project/vllm-omni/tree/main/docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
|
||||
- `--tensor-parallel-size`: tensor parallel size (effective for models that support TP, e.g. LTX2).
|
||||
- `--enable-cpu-offload`: enable CPU offloading for diffusion models.
|
||||
- `--use-hsdp`: Enable Hybrid Sharded Data Parallel to shard model weights across GPUs.
|
||||
- `--hsdp-shard-size`: Number of GPUs to shard model weights across within each replica group. -1 (default) auto-calculates as world_size / replicate_size.
|
||||
|
||||
@@ -0,0 +1,46 @@
|
||||
# MammothModa2-Preview
|
||||
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/mammothmodal2_preview>.
|
||||
|
||||
|
||||
## Run examples (MammothModa2-Preview)
|
||||
|
||||
Download model
|
||||
```bash
|
||||
hf bytedance-research/MammothModa2-Preview --local-dir ./MammothModa2-Preview
|
||||
```
|
||||
|
||||
### Text-to-Image (T2I)
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py \
|
||||
--model ./MammothModa2-Preview \
|
||||
--stage-config ./vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml \
|
||||
--prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \
|
||||
--height 1024 \
|
||||
--width 1024 \
|
||||
--num-inference-steps 50 \
|
||||
--text-guidance-scale 4.0 \
|
||||
--out output.png
|
||||
```
|
||||
|
||||
### Image Summary
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py \
|
||||
--model ./MammothModa2-Preview \
|
||||
--stage-config ./vllm_omni/model_executor/stage_configs/mammoth_moda2_ar.yaml \
|
||||
--question "Summarize this image." \
|
||||
--image ./image.png
|
||||
```
|
||||
|
||||
## Example materials
|
||||
|
||||
??? abstract "run_mammothmoda2_image_summarize.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py"
|
||||
``````
|
||||
??? abstract "run_mammothmoda2_t2i.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py"
|
||||
``````
|
||||
@@ -0,0 +1,254 @@
|
||||
# MiMo-Audio Offline Inference
|
||||
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/mimo_audio>.
|
||||
|
||||
|
||||
This directory contains an offline demo for running MiMo-Audio models with vLLM Omni. It builds task-specific inputs and generates WAV files or text outputs locally.
|
||||
|
||||
## Model Overview
|
||||
|
||||
MiMo-Audio provides multiple task variants for audio understanding and generation:
|
||||
|
||||
- **tts_sft**: Basic text-to-speech generation from text input.
|
||||
- **tts_sft_with_instruct**: TTS generation with explicit voice style instructions.
|
||||
- **tts_sft_with_audio**: TTS generation with audio reference for voice cloning.
|
||||
- **tts_sft_with_natural_instruction**: TTS generation from natural language descriptions embedded in text.
|
||||
- **audio_trancribing_sft**: Transcribe audio to text (speech-to-text). (note: the upstream task name uses the spelling 'trancribing', don't fix it)
|
||||
- **audio_understanding_sft**: Understand and analyze audio content with text queries.
|
||||
- **audio_understanding_sft_with_thinking**: Audio understanding with reasoning chain.
|
||||
- **spoken_dialogue_sft_multiturn**: Multi-turn spoken dialogue with audio input/output.
|
||||
- **speech2text_dialogue_sft_multiturn**: Multi-turn dialogue converting speech to text.
|
||||
- **text_dialogue_sft_multiturn**: Multi-turn text-only dialogue.
|
||||
|
||||
## 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.
|
||||
|
||||
### Environment Variables
|
||||
|
||||
The `MIMO_AUDIO_TOKENIZER_PATH` environment variable is mandatory due to the specialized architecture:
|
||||
|
||||
```bash
|
||||
export MIMO_AUDIO_TOKENIZER_PATH="XiaomiMiMo/MiMo-Audio-Tokenizer"
|
||||
```
|
||||
|
||||
## Quick Start
|
||||
|
||||
Run a single sample for basic TTS:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type tts_sft
|
||||
```
|
||||
|
||||
Run batch samples for basic TTS:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type tts_sft \
|
||||
--num-prompts {batch_size}
|
||||
```
|
||||
|
||||
When enabling multi-batch processing, if the total number of tokens passed to the next stage exceeds the `max_model_len` value in the `mimo_audio.yaml` configuration file, you must also synchronously update the `max_position_embeddings` value in `MiMo-Audio-7B-Instruct/config.json` to match the modified value.
|
||||
|
||||
Generated audio files are saved to `output_audio/` by default. `--num-prompts` also can be used to all tasks below.
|
||||
|
||||
## Task Usage
|
||||
|
||||
### tts_sft (Basic Text-to-Speech)
|
||||
|
||||
Generate speech from text input:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type tts_sft \
|
||||
--text "The weather is so nice today."
|
||||
```
|
||||
|
||||
### tts_sft_with_instruct (TTS with Voice Instructions)
|
||||
|
||||
Generate speech with explicit voice style instructions:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type tts_sft_with_instruct \
|
||||
--text "The weather is so nice today." \
|
||||
--instruct "Speak happily in a child's voice"
|
||||
```
|
||||
|
||||
### tts_sft_with_audio (TTS with Audio Reference)
|
||||
|
||||
Generate speech using an audio reference for voice cloning:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type tts_sft_with_audio \
|
||||
--text "The weather is so nice today." \
|
||||
--audio-path "./spoken_dialogue_assistant_turn_1.wav"
|
||||
```
|
||||
|
||||
### tts_sft_with_natural_instruction (Natural Language TTS)
|
||||
|
||||
Generate speech from text containing natural voice descriptions:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type tts_sft_with_natural_instruction \
|
||||
--text "In a panting young male voice, he said: I can't run anymore, wait for me!"
|
||||
```
|
||||
|
||||
### audio_trancribing_sft (Speech-to-Text)
|
||||
|
||||
Transcribe audio to text:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type audio_trancribing_sft \
|
||||
--audio-path "./spoken_dialogue_assistant_turn_1.wav"
|
||||
```
|
||||
|
||||
### audio_understanding_sft (Audio Understanding)
|
||||
|
||||
Understand and analyze audio content with text queries:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type audio_understanding_sft \
|
||||
--text "Summarize the audio." \
|
||||
--audio-path "./spoken_dialogue_assistant_turn_1.wav"
|
||||
```
|
||||
|
||||
### audio_understanding_sft_with_thinking (Audio Understanding with Reasoning)
|
||||
|
||||
Audio understanding with reasoning chain:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type audio_understanding_sft_with_thinking \
|
||||
--text "Summarize the audio." \
|
||||
--audio-path "./spoken_dialogue_assistant_turn_1.wav"
|
||||
```
|
||||
|
||||
### spoken_dialogue_sft_multiturn (Multi-turn Spoken Dialogue)
|
||||
|
||||
Multi-turn dialogue with audio input and output:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type spoken_dialogue_sft_multiturn \
|
||||
--audio-path "./prompt_speech_zh_m.wav"
|
||||
```
|
||||
|
||||
Note: This task uses hardcoded audio files in the script. The audio files used in examples are available at: https://github.com/XiaomiMiMo/MiMo-Audio/tree/main/examples
|
||||
|
||||
### speech2text_dialogue_sft_multiturn (Speech-to-Text Dialogue)
|
||||
|
||||
Multi-turn dialogue converting speech to text:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type speech2text_dialogue_sft_multiturn
|
||||
```
|
||||
|
||||
Note: This task uses hardcoded audio files and message lists in the script.
|
||||
|
||||
### text_dialogue_sft_multiturn (Text Dialogue)
|
||||
|
||||
Multi-turn text-only dialogue:
|
||||
|
||||
```bash
|
||||
python3 -u end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--model-name XiaomiMiMo/MiMo-Audio-7B-Instruct \
|
||||
--query-type text_dialogue_sft_multiturn
|
||||
```
|
||||
|
||||
Note: This task uses hardcoded message lists in the script.
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### Audio dependencies (soundfile, librosa)
|
||||
|
||||
This example depends on **soundfile** (read/write WAV) and **librosa** (load audio including MP3). Install the project requirements first:
|
||||
|
||||
```bash
|
||||
pip install -r requirements/common.txt
|
||||
# or at least: pip install soundfile>=0.13.1 librosa>=0.11.0
|
||||
```
|
||||
|
||||
- **`soundfile` / libsndfile not found**
|
||||
`soundfile` uses the C library **libsndfile**. On Linux, install the system package before pip:
|
||||
- Debian/Ubuntu: `sudo apt-get install libsndfile1`
|
||||
- For development builds: `sudo apt-get install libsndfile1-dev`
|
||||
- Then: `pip install soundfile`
|
||||
|
||||
- **`librosa` fails to load MP3 or reports "No backend available"**
|
||||
Loading MP3 (e.g. in `spoken_dialogue_sft_multiturn` with `.mp3` files) uses **ffmpeg** as the backend. Install ffmpeg:
|
||||
- Debian/Ubuntu: `sudo apt-get install ffmpeg`
|
||||
- macOS: `brew install ffmpeg`
|
||||
|
||||
- **`ImportError: No module named 'soundfile'` or `ModuleNotFoundError: ... librosa`**
|
||||
Ensure you are in the same Python environment where vLLM Omni and the example dependencies are installed, and that `requirements/common.txt` (or the packages above) are installed.
|
||||
|
||||
### Tokenizer path
|
||||
|
||||
- **`MIMO_AUDIO_TOKENIZER_PATH` not set or model fails to find tokenizer**
|
||||
Export the tokenizer path before running:
|
||||
```bash
|
||||
export MIMO_AUDIO_TOKENIZER_PATH="XiaomiMiMo/MiMo-Audio-Tokenizer"
|
||||
```
|
||||
See [Environment Variables](#environment-variables) in Setup.
|
||||
|
||||
### Other
|
||||
|
||||
- If the model or stage config fails to load, check [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) for memory and GPU settings.
|
||||
- For errors when reading/writing WAV (e.g. unsupported format), ensure input files are standard WAV/MP3 and that `soundfile` is linked to a working libsndfile (see above).
|
||||
|
||||
## Notes
|
||||
|
||||
- The script uses default model paths and audio files embedded in `end2end.py`. Update them if your local cache path differs.
|
||||
- Use `--output-dir` to change the output folder (default: `./output_audio`).
|
||||
- Use `--num-prompts` to generate multiple prompts in one run (default: 1).
|
||||
- Audio files used in multi-turn dialogue examples are available at: https://github.com/XiaomiMiMo/MiMo-Audio/tree/main/examples
|
||||
- The script supports various configuration options for initialization timeouts, batch timeouts, and shared memory thresholds. See `--help` for details.
|
||||
|
||||
## Example materials
|
||||
|
||||
??? abstract "end2end.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/mimo_audio/end2end.py"
|
||||
``````
|
||||
??? abstract "message_base64_wav.json"
|
||||
``````json
|
||||
--8<-- "examples/offline_inference/mimo_audio/message_base64_wav.json"
|
||||
``````
|
||||
??? abstract "message_convert.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/mimo_audio/message_convert.py"
|
||||
``````
|
||||
??? abstract "process_speechdata.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/mimo_audio/process_speechdata.py"
|
||||
``````
|
||||
@@ -67,6 +67,51 @@ If media file paths are not provided, the script will use default assets. Suppor
|
||||
- `multi_audios`: Multiple audio inputs
|
||||
- `mixed_modalities`: Combination of video, image, and audio inputs
|
||||
|
||||
### Async-chunk (offline)
|
||||
|
||||
For true stage-level concurrency -- where downstream stages (Talker, Code2Wav)
|
||||
start **before** the upstream stage (Thinker) finishes -- use the async_chunk
|
||||
example. This requires:
|
||||
|
||||
1. A stage config YAML with ``async_chunk: true`` (e.g.
|
||||
``qwen3_omni_moe_async_chunk.yaml``).
|
||||
2. Hardware that matches the config (e.g. 2x H100 for the default 3-stage
|
||||
config).
|
||||
|
||||
The async_chunk example uses ``AsyncOmni`` instead of the synchronous ``Omni``
|
||||
class, which enables the async orchestrator to receive stage-0 intermediate
|
||||
outputs and trigger downstream stages early. Chunk data flows directly between
|
||||
stage workers via the in-worker ``OmniChunkTransferAdapter`` / connector,
|
||||
**not** through the orchestrator.
|
||||
|
||||
#### Single prompt
|
||||
```bash
|
||||
cd examples/offline_inference/qwen3_omni
|
||||
bash run_single_prompt_async_chunk.sh
|
||||
```
|
||||
|
||||
#### Multiple prompts with concurrency control
|
||||
```bash
|
||||
bash run_multiple_prompts_async_chunk.sh --max-in-flight 4
|
||||
```
|
||||
|
||||
#### Text-only output (skip audio generation)
|
||||
```bash
|
||||
python end2end_async_chunk.py --query-type text --modalities text
|
||||
```
|
||||
|
||||
#### Custom stage config
|
||||
```bash
|
||||
python end2end_async_chunk.py \
|
||||
--query-type use_audio \
|
||||
--stage-configs-path /path/to/your_async_chunk.yaml
|
||||
```
|
||||
|
||||
> **Note**: The synchronous ``end2end.py`` (using ``Omni``) is still the
|
||||
> recommended entry point for non-async-chunk workflows. Only use the
|
||||
> async_chunk example when you need the stage-level concurrency semantics
|
||||
> described in PR #962 / #1151.
|
||||
|
||||
### FAQ
|
||||
|
||||
If you encounter error about backend of librosa, try to install ffmpeg with command below.
|
||||
@@ -81,14 +126,26 @@ sudo apt install ffmpeg
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/qwen3_omni/end2end.py"
|
||||
``````
|
||||
??? abstract "end2end_async_chunk.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/qwen3_omni/end2end_async_chunk.py"
|
||||
``````
|
||||
??? abstract "run_multiple_prompts.sh"
|
||||
``````sh
|
||||
--8<-- "examples/offline_inference/qwen3_omni/run_multiple_prompts.sh"
|
||||
``````
|
||||
??? abstract "run_multiple_prompts_async_chunk.sh"
|
||||
``````sh
|
||||
--8<-- "examples/offline_inference/qwen3_omni/run_multiple_prompts_async_chunk.sh"
|
||||
``````
|
||||
??? abstract "run_single_prompt.sh"
|
||||
``````sh
|
||||
--8<-- "examples/offline_inference/qwen3_omni/run_single_prompt.sh"
|
||||
``````
|
||||
??? abstract "run_single_prompt_async_chunk.sh"
|
||||
``````sh
|
||||
--8<-- "examples/offline_inference/qwen3_omni/run_single_prompt_async_chunk.sh"
|
||||
``````
|
||||
??? abstract "run_single_prompt_tp.sh"
|
||||
``````sh
|
||||
--8<-- "examples/offline_inference/qwen3_omni/run_single_prompt_tp.sh"
|
||||
|
||||
@@ -90,6 +90,43 @@ 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):
|
||||
|
||||
@@ -3,14 +3,50 @@
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_image>.
|
||||
|
||||
|
||||
This folder provides several entrypoints for experimenting with `Qwen/Qwen-Image` `Qwen/Qwen-Image-2512` `Tongyi-MAI/Z-Image-Turbo` `stepfun-ai/NextStep-1.1` using vLLM-Omni, note that NextStep-1.1 has different architecture so we treat it differently regarding running arguments and pipeline.
|
||||
Generate images from text prompts using vLLM-Omni's diffusion pipeline entrypoints.
|
||||
|
||||
- `text_to_image.py`: command-line script for single image generation with advanced options.
|
||||
- `web_demo.py`: lightweight Gradio UI for interactive prompt/seed/CFG exploration.
|
||||
- `gradio_demo.py`: lightweight Gradio UI for interactive prompt/seed/CFG exploration.
|
||||
|
||||
Note that when you pass in multiple independent prompts, they will be processed sequentially. Batching requests is currently not supported.
|
||||
## Table of Contents
|
||||
|
||||
## Basic Usage
|
||||
- [Overview](#overview)
|
||||
- [Quick Start](#quick-start)
|
||||
- [Key Arguments](#key-arguments)
|
||||
- [More CLI Examples](#more-cli-examples)
|
||||
- [Web UI Demo](#web-ui-demo)
|
||||
|
||||
## Overview
|
||||
|
||||
This folder provides several entrypoints for experimenting with text-to-image diffusion models using vLLM-Omni. Note that `NextStep-1.1` has a different architecture, so it is treated differently regarding running arguments and pipeline.
|
||||
|
||||
### Supported Models
|
||||
|
||||
| Model | Image Shape | Peak VRAM (GiB) * | Model Weights (GiB) |
|
||||
| ----- | ----------- | ----------- | ----------------- |
|
||||
| `Qwen/Qwen-Image` | 1024 x 1024 | 60.0 | 53.7 |
|
||||
| `Qwen/Qwen-Image-2512` |1024 x 1024 | 60.0 | 53.7 |
|
||||
| `Tongyi-MAI/Z-Image-Turbo` | 1024 x 1024 | 24.8 | 19.2 |
|
||||
| `stepfun-ai/NextStep-1.1` | 512 x 512 | 71.8 | 28.1 |
|
||||
| `meituan-longcat/LongCat-Image` | 1024 x 1024 | 71.2 | 27.3 |
|
||||
| `AIDC-AI/Ovis-Image-7B` | 1024 x 1024 | 71.8 | 17.1 |
|
||||
| `OmniGen2/OmniGen2` | 1024 x 1024 | 20.1 | 14.7 |
|
||||
| `stabilityai/stable-diffusion-3.5-medium` | 1024 x 1024 | 20.1 | 15.6 |
|
||||
| `black-forest-labs/FLUX.1-dev` | 1024 x 1024 | 77.6 | 31.4 |
|
||||
| `black-forest-labs/FLUX.2-klein-4B` | 1024 x 1024 | 72.7 | 14.9 |
|
||||
| `black-forest-labs/FLUX.2-klein-9B` | 1024 x 1024 | 37.1 | 32.3 |
|
||||
| `black-forest-labs/FLUX.2-dev` | 1024 x 1024 | 65.7 | >80 (CPU offload required) |
|
||||
|
||||
!!! info
|
||||
*Peak VRAM: based on basic single-card usage, batch size =1, without any acceleration/optimization features. FLUX.2-dev requires `--enable-cpu-offload` on a single 80 GiB GPU.
|
||||
|
||||
Default model: `Qwen/Qwen-Image`
|
||||
|
||||
## Quick Start
|
||||
|
||||
### Python API
|
||||
|
||||
Single-prompt generation:
|
||||
|
||||
```python
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
@@ -23,78 +59,74 @@ if __name__ == "__main__":
|
||||
images[0].save("coffee.png")
|
||||
```
|
||||
|
||||
Or put more than one prompt in a request.
|
||||
### Local CLI Usage
|
||||
|
||||
```python
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
if __name__ == "__main__":
|
||||
omni = Omni(model="Qwen/Qwen-Image")
|
||||
prompts = [
|
||||
"a cup of coffee on a table",
|
||||
"a toy dinosaur on a sandy beach",
|
||||
"a fox waking up in bed and yawning",
|
||||
]
|
||||
outputs = omni.generate(prompts)
|
||||
for i, output in enumerate(outputs):
|
||||
image = output.request_output.images[0].save(f"{i}.jpg")
|
||||
```bash
|
||||
python text_to_image.py \
|
||||
--model Qwen/Qwen-Image \
|
||||
--prompt "a cup of coffee on the table" \
|
||||
--output coffee.png
|
||||
```
|
||||
|
||||
!!! info
|
||||
## Key Arguments
|
||||
|
||||
However, it is not currently recommended to do so
|
||||
because not all models support batch inference,
|
||||
and batch requesting mostly does not provide significant performance improvement (despite the impression that it does).
|
||||
This feature is primarily for the sake of interface compatibility with vLLM and to allow for future improvements.
|
||||
**Common arguments:**
|
||||
|
||||
!!! info
|
||||
| Argument | Type | Default | Description |
|
||||
| -------- | ---- | ------- | ----------- |
|
||||
| `--prompt` | str | `"a cup of coffee on the table"` | Text description for image generation |
|
||||
| `--seed` | int | `142` | Integer seed for deterministic sampling |
|
||||
| `--negative-prompt` | str | `None` | Negative prompt for classifier-free conditional guidance |
|
||||
| `--cfg-scale` | float | `4.0` | True CFG scale (model-specific guidance strength) |
|
||||
| `--guidance-scale` | float | `1.0` | Classifier-free guidance scale |
|
||||
| `--num-images-per-prompt` | int | `1` | Number of images per prompt (saved as `output`, `output_1`, ...) |
|
||||
| `--num-inference-steps` | int | `50` | Diffusion sampling steps (more steps = higher quality, slower) |
|
||||
| `--height` | int | `1024` | Output image height in pixels |
|
||||
| `--width` | int | `1024` | Output image width in pixels |
|
||||
| `--output` | str | `"qwen_image_output.png"` | Path to save the generated image |
|
||||
| `--vae-use-slicing` | flag | off | Enable VAE slicing for memory optimization |
|
||||
| `--vae-use-tiling` | flag | off | Enable VAE tiling for memory optimization |
|
||||
| `--cfg-parallel-size` | int | `1` | Set to `2` to enable CFG Parallel |
|
||||
| `--enable-cpu-offload` | flag | off | Enable CPU offloading for diffusion models |
|
||||
| `--lora-path` | str | — | Path to PEFT LoRA adapter folder |
|
||||
| `--lora-scale` | float | `1.0` | Scale factor for LoRA weights |
|
||||
|
||||
For diffusion pipelines, the stage config field `stage_args.[].engine_args.max_num_seqs` is 1 by default, and the input
|
||||
list is sliced into single-item requests before feeding into the diffusion pipeline. For models that do internally support
|
||||
batched inputs, you can [modify this configuration](https://github.com/vllm-project/vllm-omni/tree/main/configuration/stage_configs.md) to let the model accept a longer batch of prompts.
|
||||
**NextStep-1.1 specific arguments:**
|
||||
|
||||
Apart from string prompt, vLLM-Omni also supports dictionary prompts in the same style as vLLM.
|
||||
This is useful for models that support negative prompts.
|
||||
| Argument | Type | Default | Description |
|
||||
| -------- | ---- | ------- | ----------- |
|
||||
| `--guidance-scale-2` | float | `1.0` | Secondary guidance scale (e.g. image-level CFG) |
|
||||
| `--timesteps-shift` | float | `1.0` | Timesteps shift parameter for sampling |
|
||||
| `--cfg-schedule` | str | `"constant"` | CFG schedule type: `"constant"` or `"linear"` |
|
||||
| `--use-norm` | flag | off | Apply layer normalization to sampled tokens |
|
||||
|
||||
```python
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
> If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.
|
||||
|
||||
if __name__ == "__main__":
|
||||
omni = Omni(model="Qwen/Qwen-Image")
|
||||
outputs = omni.generate([
|
||||
{
|
||||
"prompt": "a cup of coffee on a table",
|
||||
"negative_prompt": "low resolution"
|
||||
},
|
||||
{
|
||||
"prompt": "a toy dinosaur on a sandy beach",
|
||||
"negative_prompt": "cinematic, realistic"
|
||||
}
|
||||
])
|
||||
for i, output in enumerate(outputs):
|
||||
image = output.request_output.images[0].save(f"{i}.jpg")
|
||||
```
|
||||
> Qwen-Image currently publishes best-effort presets at `1328x1328`, `1664x928`, `928x1664`, `1472x1140`, `1140x1472`, `1584x1056`, and `1056x1584`. Adjust `--height/--width` accordingly for the most reliable outcomes.
|
||||
|
||||
## Local CLI Usage
|
||||
## More CLI Examples
|
||||
|
||||
### Qwen/Tongyi Models
|
||||
### Tongyi Models
|
||||
|
||||
```bash
|
||||
python text_to_image.py \
|
||||
--model Tongyi-MAI/Z-Image-Turbo \
|
||||
--prompt "a cup of coffee on the table" \
|
||||
--seed 42 \
|
||||
--cfg-scale 4.0 \
|
||||
--guidance-scale 0.0 \
|
||||
--num-images-per-prompt 1 \
|
||||
--num-inference-steps 50 \
|
||||
--num-inference-steps 9 \
|
||||
--height 1024 \
|
||||
--width 1024 \
|
||||
--output outputs/coffee.png
|
||||
```
|
||||
|
||||
`Tongyi-MAI/Z-Image-Turbo` is a distilled version of Z-Image. Distilled diffusion models usually require less number of inference steps (4~9), and Classifier-Free Guidance (CFG) is usually NOT applied. Similar distilled models are `black-forest-labs/FLUX.2-klein-4B` and `black-forest-labs/FLUX.2-klein-9B`.
|
||||
|
||||
### NextStep Models
|
||||
|
||||
NextStep-1.1 can have extra arguments
|
||||
NextStep-1.1 supports extra arguments for dual-level CFG control:
|
||||
|
||||
```bash
|
||||
python text_to_image.py \
|
||||
--model stepfun-ai/NextStep-1.1 \
|
||||
@@ -109,38 +141,101 @@ python text_to_image.py \
|
||||
--seed 42
|
||||
```
|
||||
|
||||
### Key Arguments
|
||||
### FLUX.2-dev Models
|
||||
|
||||
**Common arguments:**
|
||||
To run FLUX.2-dev on a single GPU, `--enable-cpu-offload` is required because the model weights exceed 80 GiB:
|
||||
|
||||
- `--prompt`: text description (string).
|
||||
- `--seed`: integer seed for deterministic sampling.
|
||||
- `--cfg-scale`: true CFG scale (model-specific guidance strength).
|
||||
- `--num-images-per-prompt`: number of images to generate per prompt (saves as `output`, `output_1`, ...).
|
||||
- `--num-inference-steps`: diffusion sampling steps (more steps = higher quality, slower).
|
||||
- `--height/--width`: output resolution (defaults 1024x1024).
|
||||
- `--output`: path to save the generated PNG.
|
||||
- `--vae-use-slicing`: enable VAE slicing for memory optimization.
|
||||
- `--vae-use-tiling`: enable VAE tiling for memory optimization.
|
||||
- `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](https://github.com/vllm-project/vllm-omni/tree/main/docs/user_guide/diffusion_acceleration.md#using-cfg-parallel).
|
||||
- `--enable-cpu-offload`: enable CPU offloading for diffusion models.
|
||||
- `--guidance-scale`: classifier-free guidance scale.
|
||||
```bash
|
||||
python examples/offline_inference/text_to_image/text_to_image.py \
|
||||
--model black-forest-labs/FLUX.2-dev \
|
||||
--prompt "a lovely bunny holding a sign that says 'vllm-omni'" \
|
||||
--seed 42 \
|
||||
--tensor-parallel-size 1 \
|
||||
--num-images-per-prompt 1 \
|
||||
--num-inference-steps 50 \
|
||||
--guidance-scale 4.0 \
|
||||
--height 1024 \
|
||||
--width 1024 \
|
||||
--enable-cpu-offload \
|
||||
--output flux2-dev.png
|
||||
```
|
||||
|
||||
**NextStep-1.1 specific:**
|
||||
- `--guidance-scale-2`: secondary guidance scale, e.g. image-level CFG (default: 1.0).
|
||||
- `--timesteps-shift`: timesteps shift parameter for sampling (default: 1.0).
|
||||
- `--cfg-schedule`: CFG schedule type, "constant" or "linear" (default: "constant").
|
||||
- `--use-norm`: apply layer normalization to sampled tokens.
|
||||
### Batch Requests (Multiple Prompts)
|
||||
|
||||
> ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.
|
||||
You can pass multiple prompts in a single `generate` call.
|
||||
|
||||
> ℹ️ Qwen-Image currently publishes best-effort presets at `1328x1328`, `1664x928`, `928x1664`, `1472x1140`, `1140x1472`, `1584x1056`, and `1056x1584`. Adjust `--height/--width` accordingly for the most reliable outcomes.
|
||||
```python
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
## LoRA
|
||||
if __name__ == "__main__":
|
||||
omni = Omni(model="Qwen/Qwen-Image")
|
||||
prompts = [
|
||||
"a cup of coffee on a table",
|
||||
"a toy dinosaur on a sandy beach",
|
||||
"a fox waking up in bed and yawning",
|
||||
]
|
||||
outputs = omni.generate(prompts)
|
||||
for i, output in enumerate(outputs):
|
||||
output.request_output.images[0].save(f"{i}.jpg")
|
||||
```
|
||||
|
||||
This example supports Peft-compatible LoRA (Low-Rank Adaptation) adapters for diffusion models. Pass `--lora-path` to use a LoRA adapter and optionally `--lora-scale` (default 1.0); omit it to use the base model only.
|
||||
!!! info
|
||||
|
||||
### Basic usage with LoRA
|
||||
Not all models support batch inference, and batch requesting mostly does not provide significant
|
||||
performance improvement. This feature is primarily for interface compatibility with vLLM and to
|
||||
allow for future improvements.
|
||||
|
||||
!!! info
|
||||
|
||||
For diffusion pipelines, the stage config field `stage_args.[].runtime.max_batch_size` is 1 by
|
||||
default, and the input list is sliced into single-item requests before feeding into the diffusion
|
||||
pipeline. For models that do internally support batched inputs, you can
|
||||
[modify this configuration](https://github.com/vllm-project/vllm-omni/tree/main/configuration/stage_configs.md) to let the model accept a
|
||||
longer batch of prompts.
|
||||
|
||||
### Negative Prompts
|
||||
|
||||
vLLM-Omni supports dictionary prompts for models that accept negative prompts:
|
||||
|
||||
```python
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
if __name__ == "__main__":
|
||||
omni = Omni(model="Qwen/Qwen-Image")
|
||||
outputs = omni.generate([
|
||||
{
|
||||
"prompt": "a cup of coffee on a table",
|
||||
"negative_prompt": "low resolution"
|
||||
},
|
||||
{
|
||||
"prompt": "a toy dinosaur on a sandy beach",
|
||||
"negative_prompt": "cinematic, realistic"
|
||||
}
|
||||
])
|
||||
for i, output in enumerate(outputs):
|
||||
output.request_output.images[0].save(f"{i}.jpg")
|
||||
```
|
||||
|
||||
You can also pass a negative prompt via the CLI argument `--negative-prompt`:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_image/text_to_image.py \
|
||||
--model Qwen/Qwen-Image \
|
||||
--prompt "a cup of coffee on a table" \
|
||||
--negative-prompt "low resolution, blurry" \
|
||||
--output coffee.png
|
||||
```
|
||||
|
||||
### Advanced Features
|
||||
|
||||
#### CFG Parallel
|
||||
|
||||
Set `--cfg-parallel-size 2` to enable CFG Parallel for faster inference on multi-GPU setups.
|
||||
See more examples in the [diffusion acceleration user guide](https://github.com/vllm-project/vllm-omni/tree/main/docs/user_guide/diffusion_acceleration.md#using-cfg-parallel).
|
||||
|
||||
#### LoRA
|
||||
|
||||
This example supports PEFT-compatible LoRA (Low-Rank Adaptation) adapters for diffusion models. Pass `--lora-path` to use a LoRA adapter and optionally `--lora-scale` (default `1.0`); omit it to use the base model only.
|
||||
|
||||
```bash
|
||||
python text_to_image.py \
|
||||
@@ -151,14 +246,7 @@ python text_to_image.py \
|
||||
--output output.png
|
||||
```
|
||||
|
||||
### LoRA parameters
|
||||
|
||||
- `--lora-path`: Path to LoRA adapter folder (PEFT format). Loaded at initialization and used for generation.
|
||||
- `--lora-scale`: Scale factor for LoRA weights (default: 1.0). Higher values increase the influence of the LoRA adapter.
|
||||
|
||||
### LoRA adapter format
|
||||
|
||||
LoRA adapters must be in PEFT (Parameter-Efficient Fine-Tuning) format. A typical LoRA adapter directory structure:
|
||||
LoRA adapters must be in PEFT format. A typical adapter directory structure:
|
||||
|
||||
```
|
||||
lora_adapter/
|
||||
@@ -168,13 +256,13 @@ lora_adapter/
|
||||
|
||||
## Web UI Demo
|
||||
|
||||
Launch the gradio demo:
|
||||
Launch the Gradio demo:
|
||||
|
||||
```bash
|
||||
python gradio_demo.py --port 7862
|
||||
```
|
||||
|
||||
Then open `http://localhost:7862/` on your local browser to interact with the web UI.
|
||||
Then open `http://localhost:7862/` in your local browser to interact with the web UI.
|
||||
|
||||
## Example materials
|
||||
|
||||
|
||||
@@ -3,13 +3,20 @@
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_video>.
|
||||
|
||||
|
||||
The Wan T2V pipeline generates short videos from text prompts. Supported models include:
|
||||
A unified script for text-to-video generation. Supports multiple models with model-aware defaults.
|
||||
|
||||
- `Wan-AI/Wan2.2-T2V-A14B-Diffusers` (14B, requires ~48GB VRAM)
|
||||
- `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` (1.3B, requires ~16GB VRAM, use `--flow-shift 3.0 --boundary-ratio 0.0`)
|
||||
## Supported Models
|
||||
|
||||
| Model | Default Resolution | Default Frames | Default Steps | Guidance | VRAM (BF16) |
|
||||
|---|---|---|---|---|---|
|
||||
| `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 720x1280 | 81 | 40 | 4.0 | ~60 GiB |
|
||||
| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v` | 480x832 | 121 | 50 | 6.0 | 1×A100 80GB |
|
||||
| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required |
|
||||
|
||||
## Local CLI Usage
|
||||
|
||||
### Wan2.2 (default)
|
||||
|
||||
```bash
|
||||
python text_to_video.py \
|
||||
--prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \
|
||||
@@ -25,22 +32,119 @@ python text_to_video.py \
|
||||
--output t2v_out.mp4
|
||||
```
|
||||
|
||||
Key arguments:
|
||||
LTX2 example:
|
||||
|
||||
```bash
|
||||
python text_to_video.py \
|
||||
--model "Lightricks/LTX-2" \
|
||||
--prompt "A cinematic close-up of ocean waves at golden hour." \
|
||||
--negative-prompt "worst quality, inconsistent motion, blurry, jittery, distorted" \
|
||||
--height 512 \
|
||||
--width 768 \
|
||||
--num-frames 121 \
|
||||
--num-inference-steps 40 \
|
||||
--guidance-scale 4.0 \
|
||||
--frame-rate 24 \
|
||||
--output ltx2_out.mp4
|
||||
```
|
||||
|
||||
### HunyuanVideo-1.5 (480p)
|
||||
|
||||
```bash
|
||||
python text_to_video.py \
|
||||
--model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \
|
||||
--prompt "A cat walks through a sunlit garden, flowers swaying gently in the breeze." \
|
||||
--height 480 \
|
||||
--width 832 \
|
||||
--num-frames 121 \
|
||||
--guidance-scale 6.0 \
|
||||
--flow-shift 5.0 \
|
||||
--num-inference-steps 50 \
|
||||
--fps 24 \
|
||||
--output hunyuan_video_15_output.mp4
|
||||
```
|
||||
|
||||
### HunyuanVideo-1.5 (720p)
|
||||
|
||||
```bash
|
||||
python text_to_video.py \
|
||||
--model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v \
|
||||
--prompt "A serene lakeside sunrise with mist over the water." \
|
||||
--height 720 \
|
||||
--width 1280 \
|
||||
--num-frames 121 \
|
||||
--guidance-scale 6.0 \
|
||||
--flow-shift 9.0 \
|
||||
--num-inference-steps 50 \
|
||||
--fps 24 \
|
||||
--output hunyuan_720p.mp4
|
||||
```
|
||||
|
||||
### HunyuanVideo-1.5 with FP8 Quantization
|
||||
|
||||
```bash
|
||||
python text_to_video.py \
|
||||
--model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \
|
||||
--prompt "A dog running across a field of golden wheat." \
|
||||
--quantization fp8 \
|
||||
--height 480 --width 832 --num-frames 121 \
|
||||
--guidance-scale 6.0 --flow-shift 5.0 \
|
||||
--output hunyuan_fp8.mp4
|
||||
```
|
||||
|
||||
Quick test (smaller resolution, fewer frames):
|
||||
|
||||
```bash
|
||||
python text_to_video.py \
|
||||
--model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \
|
||||
--prompt "A serene lakeside sunrise with mist over the water." \
|
||||
--height 320 --width 576 --num-frames 17 --num-inference-steps 30 \
|
||||
--flow-shift 5.0 \
|
||||
--output quick_test.mp4
|
||||
```
|
||||
|
||||
## Key Arguments
|
||||
|
||||
### Common
|
||||
|
||||
- `--model`: Diffusers model ID or local path.
|
||||
- `--prompt`: text description (string).
|
||||
- `--height/--width`: output resolution (defaults 480x832, i.e. 480P). Dimensions should align with Wan VAE downsampling (multiples of 8).
|
||||
- `--num-frames`: Number of frames (Wan default is 81).
|
||||
- `--guidance-scale` and `--guidance-scale-high`: CFG scale (applied to low/high).
|
||||
- `--negative-prompt`: optional list of artifacts to suppress (the PR demo used a long Chinese string).
|
||||
- `--boundary-ratio`: Boundary split ratio for low/high DiT. Default `0.875` uses both transformers for best quality. Set to `1.0` to load only the low-noise transformer (saves noticeable memory with good quality, recommended if memory is limited). Set to `0.0` loads only the high-noise transformer (not recommended, lower quality).
|
||||
- `--fps`: frames per second for the saved MP4 (requires `diffusers` export_to_video).
|
||||
- `--height/--width`: output resolution. Default depends on model.
|
||||
- `--num-frames`: number of frames. Default depends on model.
|
||||
- `--guidance-scale`: CFG scale. Default depends on model.
|
||||
- `--num-inference-steps`: sampling steps. Default depends on model.
|
||||
- `--fps`: frames per second for the saved MP4.
|
||||
- `--output`: path to save the generated video.
|
||||
- `--vae-use-slicing`: enable VAE slicing for memory optimization.
|
||||
- `--vae-use-tiling`: enable VAE tiling for memory optimization.
|
||||
- `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](https://github.com/vllm-project/vllm-omni/tree/main/docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
|
||||
- `--tensor-parallel-size`: tensor parallel size (effective for models that support TP, e.g. LTX2).
|
||||
- `--enable-cpu-offload`: enable CPU offloading for diffusion models.
|
||||
- `--enable-layerwise-offload`: enable layerwise offloading on DiT modules.
|
||||
- `--frame-rate`: generation FPS for pipelines that require it (e.g., LTX2).
|
||||
- `--audio-sample-rate`: audio sample rate for embedded audio (when the pipeline returns audio).
|
||||
- `--quantization`: quantization method (`fp8` for FP8, `gguf` for GGUF).
|
||||
- `--flow-shift`: scheduler flow_shift parameter.
|
||||
|
||||
> ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.
|
||||
### Wan2.2-specific
|
||||
|
||||
- `--negative-prompt`: artifacts to suppress.
|
||||
- `--guidance-scale-high`: separate CFG scale for high-noise stage.
|
||||
- `--boundary-ratio`: boundary split for low/high DiT (default 0.875).
|
||||
- `--flow-shift`: scheduler flow_shift (5.0 for 720p, 12.0 for 480p).
|
||||
- `--cache-backend`: `cache_dit` for acceleration.
|
||||
|
||||
### HunyuanVideo-1.5 Optimal Configs
|
||||
|
||||
| Variant | flow_shift | guidance_scale | steps |
|
||||
|---------|-----------|----------------|-------|
|
||||
| 480p T2V | 5.0 | 6.0 | 50 |
|
||||
| 720p T2V | 9.0 | 6.0 | 50 |
|
||||
| 480p I2V | 5.0 | 6.0 | 50 |
|
||||
| 720p I2V | 7.0 | 6.0 | 50 |
|
||||
| CFG-distilled | (same) | 1.0 | 50 |
|
||||
|
||||
> If you encounter OOM errors, try `--vae-use-slicing`, `--vae-use-tiling`, `--enable-cpu-offload`, or `--quantization fp8`.
|
||||
|
||||
## Example materials
|
||||
|
||||
|
||||
@@ -0,0 +1,74 @@
|
||||
# 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 \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
|
||||
--write-audio --voice cheerful_female \
|
||||
--model mistralai/Voxtral-4B-TTS-2603 \
|
||||
--text "That eerie silence after the first storm was just the calm before another round of chaos, wasn't it?"
|
||||
|
||||
# 32 replicate prompts with cheerful_female voice preset
|
||||
python3 examples/offline_inference/voxtral_tts/end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
|
||||
--num-prompts 32 --write-audio --voice cheerful_female \
|
||||
--model mistralai/Voxtral-4B-TTS-2603 \
|
||||
--text "That eerie silence after the first storm was just the calm before another round of chaos, wasn't it?"
|
||||
|
||||
# Streaming with neutral_female voice preset
|
||||
python3 examples/offline_inference/voxtral_tts/end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
|
||||
--streaming --write-audio --voice neutral_female \
|
||||
--model mistralai/Voxtral-4B-TTS-2603 \
|
||||
--text "That eerie silence after the first storm was just the calm before another round of chaos, wasn't it?"
|
||||
|
||||
# 32 prompts, 8 concurrent requests per wave, streaming with neutral_female voice
|
||||
python3 examples/offline_inference/voxtral_tts/end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
|
||||
--num-prompts 32 --concurrency 8 --streaming --write-audio --voice neutral_female \
|
||||
--model mistralai/Voxtral-4B-TTS-2603 \
|
||||
--text "That eerie silence after the first storm was just the calm before another round of chaos, wasn't it?"
|
||||
|
||||
# Short debug prompt with reference audio
|
||||
# Note: Reference audio capability is not yet released.
|
||||
python3 examples/offline_inference/voxtral_tts/end2end.py \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
|
||||
--write-audio \
|
||||
--model mistralai/Voxtral-4B-TTS-2603 \
|
||||
--text "This is a test message." \
|
||||
--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`) |
|
||||
| `--stage-configs-path PATH` | Path to stage configs YAML (currently it must be set for VoxtralTTS) |
|
||||
| `--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"
|
||||
``````
|
||||
@@ -1,5 +1,8 @@
|
||||
# X-To-Video-Audio
|
||||
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/x_to_video_audio>.
|
||||
|
||||
|
||||
The `DreamID-Omni` pipeline generates short videos from text, image and video.
|
||||
|
||||
## Local CLI Usage
|
||||
@@ -56,3 +59,14 @@ Key arguments:
|
||||
- `--video-negative-prompt`: negative prompt for video generation.
|
||||
- `--audio-negative-prompt`: negative prompt for audio generation.
|
||||
- `--enable-cpu-offload`: enable CPU offload (defaults False).
|
||||
|
||||
## Example materials
|
||||
|
||||
??? abstract "download_dreamid_omni.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/x_to_video_audio/download_dreamid_omni.py"
|
||||
``````
|
||||
??? abstract "x_to_video_audio.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/x_to_video_audio/x_to_video_audio.py"
|
||||
``````
|
||||
@@ -0,0 +1,226 @@
|
||||
# vLLM-Omni Helm Chart
|
||||
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/chart-helm>.
|
||||
|
||||
|
||||
Helm chart for deploying [vLLM-Omni](https://github.com/vllm-project/vllm-omni) on Kubernetes. vLLM-Omni extends vLLM with omni-modality model serving, supporting text-to-image, multimodal chat, text-to-speech, and more.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
- Kubernetes 1.24+
|
||||
- Helm 3.x
|
||||
- NVIDIA GPU nodes with [NVIDIA Device Plugin](https://github.com/NVIDIA/k8s-device-plugin)
|
||||
|
||||
## Quick Start
|
||||
|
||||
```bash
|
||||
helm install my-release ./chart-helm \
|
||||
--set model=Tongyi-MAI/Z-Image-Turbo
|
||||
```
|
||||
|
||||
## Configuration
|
||||
|
||||
### Model Selection
|
||||
|
||||
Set the `model` value to any supported HuggingFace model ID:
|
||||
|
||||
| Model | Type | GPUs | Notes |
|
||||
|-------|------|------|-------|
|
||||
| `Tongyi-MAI/Z-Image-Turbo` | text-to-image | 1 | Small, fast (default) |
|
||||
| `stabilityai/stable-diffusion-3.5-medium` | text-to-image | 1 | ~6GB VRAM |
|
||||
| `Qwen/Qwen-Image` | text-to-image | 1 | Large, ~40GB+ VRAM |
|
||||
| `Qwen/Qwen2.5-Omni-7B` | multimodal | 2 | Text + audio + image + video |
|
||||
| `Qwen/Qwen3-Omni-7B-Chat` | multimodal | 2 | Latest omni model |
|
||||
| `Qwen/Qwen3-TTS` | text-to-speech | 1 | TTS |
|
||||
|
||||
### HuggingFace Token
|
||||
|
||||
For gated models that require authentication:
|
||||
|
||||
```bash
|
||||
helm install my-release ./chart-helm \
|
||||
--set model=Qwen/Qwen2.5-Omni-7B \
|
||||
--set hfToken=hf_xxxxx \
|
||||
--set resources.requests."nvidia\.com/gpu"=2 \
|
||||
--set resources.limits."nvidia\.com/gpu"=2
|
||||
```
|
||||
|
||||
### Omni-Specific Flags
|
||||
|
||||
Enable VAE memory optimizations for diffusion models:
|
||||
|
||||
```bash
|
||||
helm install my-release ./chart-helm \
|
||||
--set model=Qwen/Qwen-Image \
|
||||
--set omniArgs.vaeUseSlicing=true \
|
||||
--set omniArgs.vaeUseTiling=true
|
||||
```
|
||||
|
||||
Enable CPU offloading:
|
||||
|
||||
```bash
|
||||
helm install my-release ./chart-helm \
|
||||
--set model=Qwen/Qwen-Image \
|
||||
--set omniArgs.enableCpuOffload=true
|
||||
```
|
||||
|
||||
Pass additional raw CLI flags:
|
||||
|
||||
```bash
|
||||
helm install my-release ./chart-helm \
|
||||
--set model=Qwen/Qwen-Image \
|
||||
--set omniArgs.extraArgs[0]="--enable-layerwise-offload"
|
||||
```
|
||||
|
||||
### Model Cache
|
||||
|
||||
By default, a PersistentVolumeClaim is created for the HuggingFace model cache to avoid re-downloading models on pod restarts:
|
||||
|
||||
```yaml
|
||||
modelCache:
|
||||
enabled: true
|
||||
storageSize: "50Gi"
|
||||
storageClassName: ""
|
||||
```
|
||||
|
||||
To use an ephemeral volume instead:
|
||||
|
||||
```bash
|
||||
helm install my-release ./chart-helm \
|
||||
--set modelCache.enabled=false
|
||||
```
|
||||
|
||||
### Custom Command Override
|
||||
|
||||
To fully override the container command:
|
||||
|
||||
```bash
|
||||
helm install my-release ./chart-helm \
|
||||
--set image.command[0]=vllm \
|
||||
--set image.command[1]=serve \
|
||||
--set image.command[2]=my-model \
|
||||
--set image.command[3]=--omni \
|
||||
--set image.command[4]=--host \
|
||||
--set image.command[5]=0.0.0.0
|
||||
```
|
||||
|
||||
## API Endpoints
|
||||
|
||||
Once deployed, vLLM-Omni exposes the following OpenAI-compatible endpoints:
|
||||
|
||||
| Endpoint | Method | Description |
|
||||
|----------|--------|-------------|
|
||||
| `/health` | GET | Health check |
|
||||
| `/v1/models` | GET | List available models |
|
||||
| `/v1/chat/completions` | POST | Chat completions (text/multimodal) |
|
||||
| `/v1/images/generations` | POST | Image generation |
|
||||
| `/v1/images/edits` | POST | Image editing |
|
||||
| `/v1/audio/speech` | POST | Text-to-speech |
|
||||
|
||||
## Files
|
||||
|
||||
| File | Description |
|
||||
|------|-------------|
|
||||
| `Chart.yaml` | Chart metadata (name, version, maintainers) |
|
||||
| `values.yaml` | Default configuration values |
|
||||
| `values.schema.json` | JSON schema for validating values |
|
||||
| `templates/_helpers.tpl` | Helper templates for common configurations |
|
||||
| `templates/deployment.yaml` | Kubernetes Deployment |
|
||||
| `templates/service.yaml` | Kubernetes Service (ClusterIP) |
|
||||
| `templates/secrets.yaml` | Secrets (generic + HuggingFace token) |
|
||||
| `templates/pvc.yaml` | PersistentVolumeClaim for model cache |
|
||||
| `templates/configmap.yaml` | Optional ConfigMap |
|
||||
| `templates/hpa.yaml` | HorizontalPodAutoscaler |
|
||||
| `templates/poddisruptionbudget.yaml` | PodDisruptionBudget |
|
||||
| `templates/custom-objects.yaml` | Custom Kubernetes objects |
|
||||
|
||||
## Running Tests
|
||||
|
||||
This chart includes unit tests using [helm-unittest](https://github.com/helm-unittest/helm-unittest). Install the plugin and run tests:
|
||||
|
||||
```bash
|
||||
# Install plugin
|
||||
helm plugin install https://github.com/helm-unittest/helm-unittest
|
||||
|
||||
# Run tests
|
||||
helm unittest .
|
||||
```
|
||||
|
||||
## Example materials
|
||||
|
||||
??? abstract ".helmignore"
|
||||
``````
|
||||
--8<-- "examples/online_serving/chart-helm/.helmignore"
|
||||
``````
|
||||
??? abstract "Chart.yaml"
|
||||
``````yaml
|
||||
--8<-- "examples/online_serving/chart-helm/Chart.yaml"
|
||||
``````
|
||||
??? abstract "ct.yaml"
|
||||
``````yaml
|
||||
--8<-- "examples/online_serving/chart-helm/ct.yaml"
|
||||
``````
|
||||
??? abstract "lintconf.yaml"
|
||||
``````yaml
|
||||
--8<-- "examples/online_serving/chart-helm/lintconf.yaml"
|
||||
``````
|
||||
??? abstract "templates/_helpers.tpl"
|
||||
``````tpl
|
||||
--8<-- "examples/online_serving/chart-helm/templates/_helpers.tpl"
|
||||
``````
|
||||
??? abstract "templates/configmap.yaml"
|
||||
``````yaml
|
||||
--8<-- "examples/online_serving/chart-helm/templates/configmap.yaml"
|
||||
``````
|
||||
??? abstract "templates/custom-objects.yaml"
|
||||
``````yaml
|
||||
--8<-- "examples/online_serving/chart-helm/templates/custom-objects.yaml"
|
||||
``````
|
||||
??? abstract "templates/deployment.yaml"
|
||||
``````yaml
|
||||
--8<-- "examples/online_serving/chart-helm/templates/deployment.yaml"
|
||||
``````
|
||||
??? abstract "templates/hpa.yaml"
|
||||
``````yaml
|
||||
--8<-- "examples/online_serving/chart-helm/templates/hpa.yaml"
|
||||
``````
|
||||
??? abstract "templates/ingress.yaml"
|
||||
``````yaml
|
||||
--8<-- "examples/online_serving/chart-helm/templates/ingress.yaml"
|
||||
``````
|
||||
??? abstract "templates/poddisruptionbudget.yaml"
|
||||
``````yaml
|
||||
--8<-- "examples/online_serving/chart-helm/templates/poddisruptionbudget.yaml"
|
||||
``````
|
||||
??? abstract "templates/pvc.yaml"
|
||||
``````yaml
|
||||
--8<-- "examples/online_serving/chart-helm/templates/pvc.yaml"
|
||||
``````
|
||||
??? abstract "templates/secrets.yaml"
|
||||
``````yaml
|
||||
--8<-- "examples/online_serving/chart-helm/templates/secrets.yaml"
|
||||
``````
|
||||
??? abstract "templates/service.yaml"
|
||||
``````yaml
|
||||
--8<-- "examples/online_serving/chart-helm/templates/service.yaml"
|
||||
``````
|
||||
??? abstract "tests/deployment_test.yaml"
|
||||
``````yaml
|
||||
--8<-- "examples/online_serving/chart-helm/tests/deployment_test.yaml"
|
||||
``````
|
||||
??? abstract "tests/ingress_test.yaml"
|
||||
``````yaml
|
||||
--8<-- "examples/online_serving/chart-helm/tests/ingress_test.yaml"
|
||||
``````
|
||||
??? abstract "tests/pvc_test.yaml"
|
||||
``````yaml
|
||||
--8<-- "examples/online_serving/chart-helm/tests/pvc_test.yaml"
|
||||
``````
|
||||
??? abstract "tests/secrets_test.yaml"
|
||||
``````yaml
|
||||
--8<-- "examples/online_serving/chart-helm/tests/secrets_test.yaml"
|
||||
``````
|
||||
??? abstract "values.yaml"
|
||||
``````yaml
|
||||
--8<-- "examples/online_serving/chart-helm/values.yaml"
|
||||
``````
|
||||
@@ -73,90 +73,42 @@ The default yaml configuration deploys AR on GPU 0 and DiT on GPU 1. You can use
|
||||
|
||||
### Text-to-Image
|
||||
|
||||
Generate images from text prompts:
|
||||
|
||||
**Using Python client**
|
||||
|
||||
```bash
|
||||
python openai_chat_client.py \
|
||||
--prompt "A photorealistic mountain landscape at sunset" \
|
||||
--height 1024 \
|
||||
--width 1024 \
|
||||
--output landscape.png
|
||||
```
|
||||
|
||||
**Using curl**
|
||||
|
||||
```bash
|
||||
curl -s http://localhost:8091/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"messages": [
|
||||
{"role": "user", "content": "A beautiful sunset over the ocean with sailing boats"}
|
||||
],
|
||||
"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
|
||||
```
|
||||
|
||||
Or use the script:
|
||||
|
||||
```bash
|
||||
# Or use the curl script:
|
||||
bash run_curl_text_to_image.sh "A futuristic city skyline at night"
|
||||
```
|
||||
|
||||
### Image-to-Image (Image Editing)
|
||||
|
||||
Edit images with text instructions:
|
||||
|
||||
**Using Python client**
|
||||
|
||||
```bash
|
||||
python openai_chat_client.py \
|
||||
--prompt "Convert this image to watercolor style" \
|
||||
--image input.png \
|
||||
--output watercolor.png
|
||||
```
|
||||
|
||||
**Using curl**
|
||||
|
||||
```bash
|
||||
IMG_B64=$(base64 < input.png | tr -d '\n')
|
||||
|
||||
curl -s http://localhost:8091/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d @- <<EOF | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > output.png
|
||||
{
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Convert this image to watercolor style"},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,'$IMG_B64'"}}
|
||||
]
|
||||
}],
|
||||
"extra_body": {
|
||||
"height": 1024,
|
||||
"width": 1024,
|
||||
"num_inference_steps": 50,
|
||||
"guidance_scale": 1.5,
|
||||
"seed": 42
|
||||
}
|
||||
}
|
||||
EOF
|
||||
```
|
||||
|
||||
Or use the script:
|
||||
|
||||
```bash
|
||||
# Or use the curl script:
|
||||
bash run_curl_image_edit.sh input.png "Convert to watercolor style"
|
||||
```
|
||||
|
||||
## Generation Parameters (extra_body)
|
||||
For general-purpose request methods (curl, OpenAI SDK, Python `requests`), see
|
||||
the [Text-to-Image](text_to_image.md) and [Image-to-Image](image_to_image.md)
|
||||
guides.
|
||||
|
||||
## Generation Parameters
|
||||
|
||||
When using `/v1/chat/completions`, pass these inside `extra_body` in the curl
|
||||
JSON, or via the `extra_body` keyword argument in the OpenAI Python SDK (see the
|
||||
[Diffusion Chat API guide](../../../../serving/diffusion_chat_api.md)).
|
||||
When using the dedicated [`/v1/images/generations`](../../../../serving/image_generation_api.md)
|
||||
or [`/v1/images/edits`](../../../../serving/image_edit_api.md) endpoints, pass
|
||||
the supported generation controls as top-level fields directly. For image
|
||||
dimensions and count, use `size` and `n` rather than `height` or `width`.
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
| --------------------- | ----- | ------- | ----------------------------------- |
|
||||
@@ -164,7 +116,7 @@ bash run_curl_image_edit.sh input.png "Convert to watercolor style"
|
||||
| `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 | 42 | Random seed for reproducibility |
|
||||
| `seed` | int | None | Optional random seed; `/v1/images/*` generates one server-side if omitted |
|
||||
| `negative_prompt` | str | None | Negative prompt |
|
||||
|
||||
## Response Format
|
||||
|
||||
@@ -69,10 +69,49 @@ cat <<EOF > request.json
|
||||
}
|
||||
EOF
|
||||
|
||||
curl -s http://localhost:8092/v1/chat/completions -H "Content-Type: application/json" -d @request.json | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2 | base64 -d > output.png
|
||||
curl -s http://localhost:8092/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d @request.json \
|
||||
| jq -r '.choices[0].message.content[0].image_url.url' \
|
||||
| cut -d',' -f2 | base64 -d > output.png
|
||||
```
|
||||
|
||||
### Method 2: Using Python Client
|
||||
### Method 2: Using OpenAI Python SDK
|
||||
|
||||
```python
|
||||
import base64
|
||||
from openai import OpenAI
|
||||
|
||||
client = OpenAI(base_url="http://localhost:8092/v1", api_key="none")
|
||||
|
||||
with open("input.png", "rb") as f:
|
||||
img_b64 = base64.b64encode(f.read()).decode()
|
||||
|
||||
response = client.chat.completions.create(
|
||||
model="Qwen/Qwen-Image-Edit",
|
||||
messages=[{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Convert to watercolor style"},
|
||||
{"type": "image_url", "image_url": {
|
||||
"url": f"data:image/png;base64,{img_b64}"
|
||||
}},
|
||||
],
|
||||
}],
|
||||
extra_body={
|
||||
"num_inference_steps": 50,
|
||||
"guidance_scale": 1,
|
||||
"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))
|
||||
```
|
||||
|
||||
### Method 3: Using Python Client Script
|
||||
|
||||
```bash
|
||||
python openai_chat_client.py --input input.png --prompt "Convert to oil painting style" --output output.png
|
||||
@@ -81,7 +120,7 @@ python openai_chat_client.py --input input.png --prompt "Convert to oil painting
|
||||
python openai_chat_client.py --input input1.png input2.png --prompt "Combine these images into a single scene" --output output.png
|
||||
```
|
||||
|
||||
### Method 3: Using Gradio Demo
|
||||
### Method 4: Using Gradio Demo
|
||||
|
||||
```bash
|
||||
python gradio_demo.py
|
||||
@@ -124,7 +163,7 @@ python gradio_demo.py
|
||||
|
||||
### Image Editing with Parameters
|
||||
|
||||
Use `extra_body` to pass generation parameters:
|
||||
Wrap generation parameters inside `extra_body` in the request JSON:
|
||||
|
||||
```json
|
||||
{
|
||||
@@ -147,6 +186,149 @@ Use `extra_body` to pass generation parameters:
|
||||
}
|
||||
```
|
||||
|
||||
!!! tip "Using the OpenAI SDK"
|
||||
When using the OpenAI Python SDK, pass these parameters via the `extra_body`
|
||||
keyword argument. The SDK merges them into the top-level request body automatically:
|
||||
|
||||
```python
|
||||
client.chat.completions.create(
|
||||
model="Qwen/Qwen-Image-Edit",
|
||||
messages=[...],
|
||||
extra_body={"num_inference_steps": 50, "guidance_scale": 7.5, "seed": 42},
|
||||
)
|
||||
```
|
||||
|
||||
For details on how generation parameters are handled across different clients, see the
|
||||
[Diffusion Chat API guide](../../../../serving/diffusion_chat_api.md).
|
||||
|
||||
### Layered Image Generation (Qwen-Image-Layered)
|
||||
|
||||
Qwen-Image-Layered generates multiple decomposed layers from a reference image and a text prompt.
|
||||
Start the server with:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen-Image-Layered --omni --port 8093
|
||||
```
|
||||
|
||||
=== "curl"
|
||||
|
||||
```bash
|
||||
IMG_B64=$(base64 -w0 input.png)
|
||||
|
||||
curl -sS http://localhost:8093/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d "$(jq -n --arg img "$IMG_B64" '{
|
||||
messages: [{
|
||||
role: "user",
|
||||
content: [
|
||||
{type: "image_url", image_url: {url: ("data:image/png;base64," + $img)}},
|
||||
{type: "text", text: "a rabbit"}
|
||||
]
|
||||
}],
|
||||
extra_body: {
|
||||
num_inference_steps: 50,
|
||||
cfg_scale: 4.0,
|
||||
seed: 0,
|
||||
layers: 4,
|
||||
resolution: 640
|
||||
}
|
||||
}')" \
|
||||
| jq -r '.choices[0].message.content[] | .image_url.url | split(",")[1]' \
|
||||
| while IFS= read -r b64; do
|
||||
((i++)); echo "$b64" | base64 -d > "layer_${i}.png"
|
||||
done
|
||||
```
|
||||
|
||||
=== "OpenAI SDK"
|
||||
|
||||
```python
|
||||
import base64
|
||||
from openai import OpenAI
|
||||
|
||||
client = OpenAI(base_url="http://localhost:8093/v1", api_key="none")
|
||||
|
||||
with open("input.png", "rb") as f:
|
||||
img_b64 = base64.b64encode(f.read()).decode()
|
||||
|
||||
response = client.chat.completions.create(
|
||||
model="Qwen/Qwen-Image-Layered",
|
||||
messages=[{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image_url", "image_url": {
|
||||
"url": f"data:image/png;base64,{img_b64}"
|
||||
}},
|
||||
{"type": "text", "text": "a rabbit"},
|
||||
],
|
||||
}],
|
||||
extra_body={
|
||||
"num_inference_steps": 50,
|
||||
"cfg_scale": 4.0,
|
||||
"seed": 0,
|
||||
"layers": 4,
|
||||
"resolution": 640,
|
||||
},
|
||||
)
|
||||
|
||||
for i, item in enumerate(response.choices[0].message.content):
|
||||
_, b64_data = item.image_url.url.split(",", 1)
|
||||
with open(f"layer_{i}.png", "wb") as f:
|
||||
f.write(base64.b64decode(b64_data))
|
||||
```
|
||||
|
||||
=== "Python requests"
|
||||
|
||||
```python
|
||||
import base64
|
||||
import requests
|
||||
|
||||
with open("input.png", "rb") as f:
|
||||
img_b64 = base64.b64encode(f.read()).decode()
|
||||
|
||||
payload = {
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image_url", "image_url": {
|
||||
"url": f"data:image/png;base64,{img_b64}"
|
||||
}},
|
||||
{"type": "text", "text": "a rabbit"},
|
||||
],
|
||||
}],
|
||||
"extra_body": {
|
||||
"num_inference_steps": 50,
|
||||
"cfg_scale": 4.0,
|
||||
"seed": 0,
|
||||
"layers": 4,
|
||||
"resolution": 640,
|
||||
},
|
||||
}
|
||||
|
||||
resp = requests.post(
|
||||
"http://localhost:8093/v1/chat/completions",
|
||||
json=payload,
|
||||
timeout=600,
|
||||
)
|
||||
data = resp.json()
|
||||
|
||||
for i, item in enumerate(data["choices"][0]["message"]["content"]):
|
||||
_, b64_data = item["image_url"]["url"].split(",", 1)
|
||||
with open(f"layer_{i}.png", "wb") as f:
|
||||
f.write(base64.b64decode(b64_data))
|
||||
```
|
||||
|
||||
The response contains multiple images in `choices[0].message.content` — one per generated layer.
|
||||
|
||||
#### Qwen-Image-Layered Parameters
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| `layers` | int | 4 | Number of layers to decompose |
|
||||
| `resolution` | int | 640 | Resolution for dimension calculation (640 or 1024) |
|
||||
| `cfg_scale` | float | 4.0 | Classifier-free guidance scale (alias for `true_cfg_scale`) |
|
||||
| `num_inference_steps` | int | 50 | Number of denoising steps |
|
||||
| `seed` | int | None | Random seed for reproducibility |
|
||||
|
||||
### Multi-Image Editing (Qwen-Image-Edit-2509)
|
||||
|
||||
Provide multiple images in `content` (order matters):
|
||||
@@ -166,7 +348,15 @@ Provide multiple images in `content` (order matters):
|
||||
}
|
||||
```
|
||||
|
||||
## Generation Parameters (extra_body)
|
||||
## Generation Parameters
|
||||
|
||||
When using `/v1/chat/completions`, pass these inside `extra_body` in the curl
|
||||
JSON, or via the `extra_body` keyword argument in the OpenAI Python SDK (see the
|
||||
[Diffusion Chat API guide](../../../../serving/diffusion_chat_api.md)).
|
||||
When using the dedicated [`/v1/images/edits`](../../../../serving/image_edit_api.md)
|
||||
endpoint, pass the supported generation controls as top-level form fields
|
||||
directly. For image dimensions and count, use `size` and `n` rather than
|
||||
`height`, `width`, or `num_outputs_per_prompt`.
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
| ------------------------ | ----- | ------- | ------------------------------------- |
|
||||
@@ -174,10 +364,12 @@ Provide multiple images in `content` (order matters):
|
||||
| `width` | int | None | Output image width in pixels |
|
||||
| `size` | str | None | Output image size (e.g., "1024x1024") |
|
||||
| `num_inference_steps` | int | 50 | Number of denoising steps |
|
||||
| `guidance_scale` | float | 7.5 | CFG guidance scale |
|
||||
| `guidance_scale` | float | 1.0 | CFG guidance scale |
|
||||
| `seed` | int | None | Random seed (reproducible) |
|
||||
| `negative_prompt` | str | None | Negative prompt |
|
||||
| `num_outputs_per_prompt` | int | 1 | Number of images to generate |
|
||||
| `layers` | int | 4 | Number of layers (Qwen-Image-Layered) |
|
||||
| `resolution` | int | 640 | Resolution, 640 or 1024 (Qwen-Image-Layered) |
|
||||
|
||||
## Response Format
|
||||
|
||||
|
||||
@@ -144,7 +144,7 @@ curl -X POST http://localhost:8091/v1/videos \
|
||||
|
||||
## Create Response Format
|
||||
|
||||
`POST /v1/videos` returns a job record.
|
||||
`POST /v1/videos` returns a job record, not inline base64 video data.
|
||||
|
||||
```json
|
||||
{
|
||||
@@ -201,6 +201,10 @@ done
|
||||
|
||||
## Example materials
|
||||
|
||||
??? abstract "run_curl_hunyuan_video_15.sh"
|
||||
``````sh
|
||||
--8<-- "examples/online_serving/image_to_video/run_curl_hunyuan_video_15.sh"
|
||||
``````
|
||||
??? abstract "run_curl_image_to_video.sh"
|
||||
``````sh
|
||||
--8<-- "examples/online_serving/image_to_video/run_curl_image_to_video.sh"
|
||||
@@ -209,3 +213,7 @@ done
|
||||
``````sh
|
||||
--8<-- "examples/online_serving/image_to_video/run_server.sh"
|
||||
``````
|
||||
??? abstract "run_server_hunyuan_video_15.sh"
|
||||
``````sh
|
||||
--8<-- "examples/online_serving/image_to_video/run_server_hunyuan_video_15.sh"
|
||||
``````
|
||||
@@ -0,0 +1,72 @@
|
||||
# Online serving Example of vLLM-Omni for MiMo-Audio
|
||||
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/mimo_audio>.
|
||||
|
||||
|
||||
## 🛠️ Installation
|
||||
|
||||
Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/README.md)
|
||||
|
||||
## Run examples (MiMo-Audio)
|
||||
|
||||
### Launch the Server
|
||||
```bash
|
||||
export MIMO_AUDIO_TOKENIZER_PATH="XiaomiMiMo/MiMo-Audio-Tokenizer"
|
||||
|
||||
vllm-omni serve XiaomiMiMo/MiMo-Audio-7B-Instruct --omni \
|
||||
--served-model-name "MiMo-Audio-7B-Instruct" \
|
||||
--port 18091 --stage-configs-path ./vllm_omni/model_executor/stage_configs/mimo_audio.yaml \
|
||||
--chat-template ./examples/online_serving/mimo_audio/chat_template.jinja
|
||||
```
|
||||
> ⚠️ **Important**
|
||||
> **MiMo-Audio is not compatible with the default chat template.**
|
||||
> The provided `chat_template.jinja` implements MiMo-specific role, audio token, and instruction formatting and **must be used for all inference**.
|
||||
|
||||
|
||||
### Send Multi-modal Request
|
||||
|
||||
Get into the example folder
|
||||
```bash
|
||||
cd examples/online_serving/mimo_audio
|
||||
```
|
||||
|
||||
#### Send request via python
|
||||
|
||||
```bash
|
||||
# Audio dialogue task
|
||||
python openai_chat_completion_client_for_multimodal_generation.py \
|
||||
--query-type multi_audios \
|
||||
--message-json ../../offline_inference/mimo_audio/message_base64_wav.json
|
||||
```
|
||||
|
||||
The Python client supports the following command-line arguments:
|
||||
|
||||
- `--query-type` (or `-q`): Query type (default: `multi_audios`)
|
||||
- Options: `multi_audios`, `text`
|
||||
- `--message-json` (or `-m`): Path to `base64` multi rounds audio messages json file
|
||||
- Do not pass any value for "text" query type
|
||||
- Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs, only for "Are these two audio clips the same?" task
|
||||
- Example: `---message-json ./examples/offline_inference/mimo_audio/message_base64_wav.json`
|
||||
- `--prompt` (or `-p`): Custom text prompt/question, only for query type is "text"(TTS task)
|
||||
- Attention! Do not pass any value for "multi_audios" query type
|
||||
- Example: `--prompt "What are the main activities shown in this video?"`
|
||||
|
||||
|
||||
For example, to use multi rounds audios with local files:
|
||||
|
||||
```bash
|
||||
python openai_chat_completion_client_for_multimodal_generation.py \
|
||||
--query-type multi_audios \
|
||||
--message-json ../../offline_inference/mimo_audio/message_base64_wav.json
|
||||
```
|
||||
|
||||
## Example materials
|
||||
|
||||
??? abstract "chat_template.jinja"
|
||||
``````jinja
|
||||
--8<-- "examples/online_serving/mimo_audio/chat_template.jinja"
|
||||
``````
|
||||
??? abstract "openai_chat_completion_client_for_multimodal_generation.py"
|
||||
``````py
|
||||
--8<-- "examples/online_serving/mimo_audio/openai_chat_completion_client_for_multimodal_generation.py"
|
||||
``````
|
||||
@@ -223,10 +223,6 @@ sudo apt install ffmpeg
|
||||
``````py
|
||||
--8<-- "examples/online_serving/qwen2_5_omni/gradio_demo.py"
|
||||
``````
|
||||
??? abstract "openai_chat_completion_client_for_multimodal_generation.py"
|
||||
``````py
|
||||
--8<-- "examples/online_serving/qwen2_5_omni/openai_chat_completion_client_for_multimodal_generation.py"
|
||||
``````
|
||||
??? abstract "run_curl_multimodal_generation.sh"
|
||||
``````sh
|
||||
--8<-- "examples/online_serving/qwen2_5_omni/run_curl_multimodal_generation.sh"
|
||||
|
||||
@@ -21,10 +21,10 @@ Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/
|
||||
|
||||
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`) |
|
||||
| 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
|
||||
@@ -103,13 +103,13 @@ cd examples/online_serving/qwen3_tts
|
||||
# CustomVoice: Use predefined speaker
|
||||
python openai_speech_client.py \
|
||||
--text "你好,我是通义千问" \
|
||||
--voice vivian \
|
||||
--speaker vivian \
|
||||
--language Chinese
|
||||
|
||||
# CustomVoice with style instruction
|
||||
python openai_speech_client.py \
|
||||
--text "今天天气真好" \
|
||||
--voice ryan \
|
||||
--speaker ryan \
|
||||
--instructions "用开心的语气说"
|
||||
|
||||
# VoiceDesign: Describe the voice style
|
||||
@@ -134,11 +134,11 @@ The Python client supports the following command-line arguments:
|
||||
- `--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)
|
||||
- `--voice`: Speaker/voice name (default: `vivian`). Options: `vivian`, `ryan`, `aiden`, etc.
|
||||
- `--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)
|
||||
- `--ref-text`: Reference audio transcript for voice cloning (Base task)
|
||||
- `--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`)
|
||||
|
||||
@@ -150,7 +150,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "Hello, how are you?",
|
||||
"voice": "vivian",
|
||||
"speaker": "vivian",
|
||||
"language": "English"
|
||||
}' --output output.wav
|
||||
|
||||
@@ -159,7 +159,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "I am so excited!",
|
||||
"voice": "vivian",
|
||||
"speaker": "vivian",
|
||||
"instructions": "Speak with great enthusiasm"
|
||||
}' --output excited.wav
|
||||
|
||||
@@ -176,7 +176,7 @@ 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",
|
||||
speaker="vivian",
|
||||
input="Hello, how are you?",
|
||||
)
|
||||
|
||||
@@ -192,7 +192,7 @@ response = httpx.post(
|
||||
"http://localhost:8091/v1/audio/speech",
|
||||
json={
|
||||
"input": "Hello, how are you?",
|
||||
"voice": "vivian",
|
||||
"speaker": "vivian",
|
||||
"language": "English",
|
||||
},
|
||||
timeout=300.0,
|
||||
@@ -242,6 +242,7 @@ Upload a new voice sample for voice cloning in Base task TTS requests.
|
||||
- `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).
|
||||
|
||||
**Response Example:**
|
||||
```json
|
||||
@@ -262,7 +263,8 @@ Upload a new voice sample for voice cloning in Base task TTS requests.
|
||||
curl -X POST http://localhost:8000/v1/audio/voices \
|
||||
-F "audio_sample=@/path/to/voice_sample.wav" \
|
||||
-F "consent=user_consent_id" \
|
||||
-F "name=custom_voice_1"
|
||||
-F "name=custom_voice_1" \
|
||||
-F "ref_text=The exact transcript of the audio sample."
|
||||
```
|
||||
|
||||
### Endpoint
|
||||
@@ -279,12 +281,12 @@ This endpoint follows the [OpenAI Audio Speech API](https://platform.openai.com/
|
||||
```json
|
||||
{
|
||||
"input": "Text to synthesize",
|
||||
"voice": "vivian",
|
||||
"speaker": "vivian",
|
||||
"response_format": "wav",
|
||||
"task_type": "CustomVoice",
|
||||
"language": "Auto",
|
||||
"instructions": "Optional style instructions",
|
||||
"ref_audio": "URL or base64 for voice cloning",
|
||||
"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
|
||||
@@ -297,6 +299,12 @@ This endpoint follows the [OpenAI Audio Speech API](https://platform.openai.com/
|
||||
|
||||
Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/wav`).
|
||||
|
||||
### Voice and language (summary)
|
||||
|
||||
- **Speaker**: Use the `speaker` 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
|
||||
@@ -305,30 +313,30 @@ Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/w
|
||||
| ----------------- | ------ | -------------- | ----------------------------------------------------------- |
|
||||
| `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) |
|
||||
| `speaker` | 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"`) |
|
||||
| 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 (URL or base64 data URL) |
|
||||
| `ref_text` | string | No | Transcript of reference audio (for ICL mode) |
|
||||
| `x_vector_only_mode` | bool | No | Use speaker embedding only (no ICL) |
|
||||
| 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
|
||||
|
||||
@@ -340,7 +348,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "Hello, how are you?",
|
||||
"voice": "vivian",
|
||||
"speaker": "vivian",
|
||||
"language": "English",
|
||||
"stream": true,
|
||||
"response_format": "pcm"
|
||||
@@ -352,6 +360,42 @@ curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
- `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.
|
||||
@@ -386,6 +430,10 @@ curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
``````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"
|
||||
|
||||
@@ -44,7 +44,7 @@ vllm serve Qwen/Qwen-Image --omni --port 8091 --ring 2
|
||||
vllm serve Qwen/Qwen-Image --omni --port 8091 --usp 2 --ring 2
|
||||
```
|
||||
|
||||
For more details on parallelism acceleration, see the [Parallelism Acceleration Guide](../../diffusion/parallelism_acceleration.md).
|
||||
For more details on parallelism acceleration, see the [Parallelism Acceleration Guide](https://github.com/vllm-project/vllm-omni/tree/main/examples/diffusion/parallelism_acceleration.md).
|
||||
|
||||
## API Calls
|
||||
|
||||
@@ -71,13 +71,39 @@ curl -s http://localhost:8091/v1/chat/completions \
|
||||
}' | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > output.png
|
||||
```
|
||||
|
||||
### Method 2: Using Python Client
|
||||
### Method 2: Using 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="Qwen/Qwen-Image",
|
||||
messages=[{"role": "user", "content": "A beautiful landscape painting"}],
|
||||
extra_body={
|
||||
"height": 1024,
|
||||
"width": 1024,
|
||||
"num_inference_steps": 50,
|
||||
"true_cfg_scale": 4.0,
|
||||
"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))
|
||||
```
|
||||
|
||||
### Method 3: Using Python Client Script
|
||||
|
||||
```bash
|
||||
python openai_chat_client.py --prompt "A beautiful landscape painting" --output output.png
|
||||
```
|
||||
|
||||
### Method 3: Using Gradio Demo
|
||||
### Method 4: Using Gradio Demo
|
||||
|
||||
```bash
|
||||
python gradio_demo.py
|
||||
@@ -120,12 +146,12 @@ curl -X POST http://localhost:8091/v1/images/generations \
|
||||
|
||||
### LoRA Parameters
|
||||
|
||||
| Parameter | Type | Description |
|
||||
| ------------- | ------ | ---------------------------------------------------------------- |
|
||||
| `name` | str | LoRA adapter name (optional, defaults to path stem) |
|
||||
| `local_path` | str | Server-local path to LoRA adapter folder (PEFT format, required) |
|
||||
| `scale` | float | LoRA scale factor (default: 1.0) |
|
||||
| `int_id` | int | LoRA integer ID for caching (optional, derived from path if not provided) |
|
||||
| Parameter | Type | Description |
|
||||
| ------------ | ----- | ------------------------------------------------------------------------- |
|
||||
| `name` | str | LoRA adapter name (optional, defaults to path stem) |
|
||||
| `local_path` | str | Server-local path to LoRA adapter folder (PEFT format, required) |
|
||||
| `scale` | float | LoRA scale factor (default: 1.0) |
|
||||
| `int_id` | int | LoRA integer ID for caching (optional, derived from path if not provided) |
|
||||
|
||||
### LoRA Adapter Format
|
||||
|
||||
@@ -151,7 +177,7 @@ lora_adapter/
|
||||
|
||||
### Generation with Parameters
|
||||
|
||||
Use `extra_body` to pass generation parameters:
|
||||
Wrap generation parameters inside `extra_body` in the request JSON:
|
||||
|
||||
```json
|
||||
{
|
||||
@@ -168,6 +194,21 @@ Use `extra_body` to pass generation parameters:
|
||||
}
|
||||
```
|
||||
|
||||
!!! tip "Using the OpenAI SDK"
|
||||
When using the OpenAI Python SDK, pass these parameters via the `extra_body`
|
||||
keyword argument. The SDK merges them into the top-level request body automatically:
|
||||
|
||||
```python
|
||||
client.chat.completions.create(
|
||||
model="Qwen/Qwen-Image",
|
||||
messages=[...],
|
||||
extra_body={"height": 1024, "width": 1024, "num_inference_steps": 50},
|
||||
)
|
||||
```
|
||||
|
||||
For details on how generation parameters are handled across different clients, see the
|
||||
[Diffusion Chat API guide](../../../../serving/diffusion_chat_api.md).
|
||||
|
||||
### Multimodal Input (Text + Structured Content)
|
||||
|
||||
```json
|
||||
@@ -183,7 +224,15 @@ Use `extra_body` to pass generation parameters:
|
||||
}
|
||||
```
|
||||
|
||||
## Generation Parameters (extra_body)
|
||||
## Generation Parameters
|
||||
|
||||
When using `/v1/chat/completions`, pass these inside `extra_body` in the curl
|
||||
JSON, or via the `extra_body` keyword argument in the OpenAI Python SDK (see the
|
||||
[Diffusion Chat API guide](../../../../serving/diffusion_chat_api.md)).
|
||||
When using the dedicated [`/v1/images/generations`](../../../../serving/image_generation_api.md)
|
||||
endpoint, pass the supported generation controls as top-level JSON fields
|
||||
directly. For image dimensions and count, use `size` and `n` rather than
|
||||
`height`, `width`, or `num_outputs_per_prompt`.
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
| ------------------------ | ----- | ------- | ------------------------------ |
|
||||
@@ -195,7 +244,6 @@ Use `extra_body` to pass generation parameters:
|
||||
| `seed` | int | None | Random seed (reproducible) |
|
||||
| `negative_prompt` | str | None | Negative prompt |
|
||||
| `num_outputs_per_prompt` | int | 1 | Number of images to generate |
|
||||
| `--cfg-parallel-size`. | int | 1 | Number of GPUs for CFG parallelism |
|
||||
|
||||
## Response Format
|
||||
|
||||
|
||||
@@ -46,6 +46,7 @@ The main endpoints are:
|
||||
## Storage
|
||||
|
||||
Generated video files are stored on local disk by the async video API.
|
||||
Local file storage behavior can be controlled via the following environment variables:
|
||||
|
||||
- `VLLM_OMNI_STORAGE_PATH`: directory used for generated files (default: `/tmp/storage`)
|
||||
- `VLLM_OMNI_STORAGE_MAX_CONCURRENCY`: max concurrent save/delete operations (default: `4`)
|
||||
@@ -147,7 +148,7 @@ curl -X POST http://localhost:8091/v1/videos \
|
||||
|
||||
## Create Response Format
|
||||
|
||||
`POST /v1/videos` returns a job record.
|
||||
`POST /v1/videos` returns a job record, not inline base64 video data.
|
||||
|
||||
```json
|
||||
{
|
||||
@@ -204,6 +205,10 @@ done
|
||||
|
||||
## Example materials
|
||||
|
||||
??? abstract "run_curl_hunyuan_video_15.sh"
|
||||
``````sh
|
||||
--8<-- "examples/online_serving/text_to_video/run_curl_hunyuan_video_15.sh"
|
||||
``````
|
||||
??? abstract "run_curl_text_to_video.sh"
|
||||
``````sh
|
||||
--8<-- "examples/online_serving/text_to_video/run_curl_text_to_video.sh"
|
||||
@@ -212,3 +217,7 @@ done
|
||||
``````sh
|
||||
--8<-- "examples/online_serving/text_to_video/run_server.sh"
|
||||
``````
|
||||
??? abstract "run_server_hunyuan_video_15.sh"
|
||||
``````sh
|
||||
--8<-- "examples/online_serving/text_to_video/run_server_hunyuan_video_15.sh"
|
||||
``````
|
||||
@@ -70,6 +70,12 @@ def parse_args():
|
||||
action="store_true",
|
||||
help="Enable diffusion pipeline profiler to display stage durations.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--quantization",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Quantization method (e.g. 'fp8').",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
return args
|
||||
@@ -119,6 +125,8 @@ def main():
|
||||
"enable_diffusion_pipeline_profiler": args.enable_diffusion_pipeline_profiler,
|
||||
}
|
||||
)
|
||||
if args.quantization:
|
||||
omni_kwargs["quantization_config"] = args.quantization
|
||||
|
||||
omni = Omni(model=model_name, **omni_kwargs)
|
||||
|
||||
|
||||
@@ -16,9 +16,12 @@ Usage:
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import math
|
||||
import os
|
||||
import tempfile
|
||||
import time
|
||||
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
import torch
|
||||
|
||||
@@ -27,6 +30,12 @@ os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
from vllm_omni import AsyncOmni, Omni
|
||||
from vllm_omni.model_executor.models.fish_speech.dac_utils import DAC_HOP_LENGTH, DAC_SAMPLE_RATE
|
||||
from vllm_omni.model_executor.models.fish_speech.prompt_utils import (
|
||||
build_fish_text_only_prompt_ids,
|
||||
estimate_fish_voice_clone_prompt_len_from_normalized,
|
||||
normalize_fish_voice_clone_texts,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -51,46 +60,64 @@ def build_prompt(
|
||||
) -> dict:
|
||||
"""Build a prompt for Fish Speech S2 Pro.
|
||||
|
||||
Uses Qwen3 chat template format with <|voice|> modality marker:
|
||||
<|im_start|>user\n<|speaker:0|>text<|im_end|>\n<|im_start|>assistant\n<|voice|>
|
||||
|
||||
The model then generates semantic tokens autoregressively.
|
||||
Uses the same text-only / structured-clone protocol as serving.
|
||||
"""
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
|
||||
if ref_audio_path is None and ref_text is None:
|
||||
prompt_ids, normalized_text = build_fish_text_only_prompt_ids(tokenizer, text)
|
||||
additional_information = {
|
||||
"text": [normalized_text],
|
||||
}
|
||||
return {
|
||||
"prompt_token_ids": prompt_ids,
|
||||
"additional_information": additional_information,
|
||||
}
|
||||
|
||||
# Build user message with speaker tag prefix.
|
||||
user_text = f"<|speaker:0|>{text}"
|
||||
if not ref_audio_path or not ref_text:
|
||||
raise ValueError("Fish Speech voice cloning requires both --ref-audio and --ref-text")
|
||||
|
||||
# Use chat template for user turn + assistant generation prompt.
|
||||
messages = [{"role": "user", "content": user_text}]
|
||||
prompt_ids = tokenizer.apply_chat_template(
|
||||
messages,
|
||||
tokenize=True,
|
||||
add_generation_prompt=True,
|
||||
normalized_text, normalized_ref_text = normalize_fish_voice_clone_texts(text, ref_text)
|
||||
ref_audio_wav, ref_audio_sr = sf.read(ref_audio_path, dtype="float32", always_2d=False)
|
||||
semantic_len = _estimate_fish_ref_code_len(ref_audio_wav, ref_audio_sr)
|
||||
ph_len = estimate_fish_voice_clone_prompt_len_from_normalized(
|
||||
tokenizer,
|
||||
normalized_text,
|
||||
normalized_ref_text,
|
||||
semantic_len,
|
||||
)
|
||||
|
||||
# Append <|voice|> token (151673) to signal voice generation mode.
|
||||
voice_token_id = tokenizer.encode("<|voice|>", add_special_tokens=False)
|
||||
prompt_ids = prompt_ids + voice_token_id
|
||||
# The model-side structured clone prefill consumes a temporary .npy file and
|
||||
# removes it after loading. Abnormal termination can still leave the file
|
||||
# behind, which is acceptable for this offline example.
|
||||
with tempfile.NamedTemporaryFile(prefix="fish_ref_", suffix=".npy", delete=False) as f:
|
||||
np.save(f, np.asarray(ref_audio_wav, dtype=np.float32))
|
||||
ref_audio_npy_path = f.name
|
||||
|
||||
additional_information = {
|
||||
"text": [text],
|
||||
"max_new_tokens": [4096],
|
||||
"text": normalized_text,
|
||||
"ref_text": normalized_ref_text,
|
||||
"ref_audio_path": ref_audio_npy_path,
|
||||
"ref_audio_sr": int(ref_audio_sr),
|
||||
"fish_structured_voice_clone": True,
|
||||
}
|
||||
|
||||
if ref_audio_path:
|
||||
additional_information["ref_audio"] = [ref_audio_path]
|
||||
if ref_text:
|
||||
additional_information["ref_text"] = [ref_text]
|
||||
|
||||
return {
|
||||
"prompt_token_ids": prompt_ids,
|
||||
"prompt_token_ids": [1] * ph_len,
|
||||
"additional_information": additional_information,
|
||||
}
|
||||
|
||||
|
||||
def _estimate_fish_ref_code_len(wav: np.ndarray, sample_rate: int) -> int:
|
||||
"""Estimate Fish semantic token length from local reference audio."""
|
||||
n_samples = int(np.asarray(wav).shape[0]) if np.asarray(wav).ndim > 0 else 0
|
||||
if sample_rate <= 0 or n_samples <= 0:
|
||||
raise ValueError("Reference audio must contain samples and a positive sample rate")
|
||||
resampled_len = max(1, math.ceil(n_samples * DAC_SAMPLE_RATE / int(sample_rate)))
|
||||
return max(1, math.ceil(resampled_len / DAC_HOP_LENGTH))
|
||||
|
||||
|
||||
def _save_wav(output_dir: str, request_id: str, mm: dict) -> None:
|
||||
"""Concatenate audio chunks and write to a wav file."""
|
||||
audio_data = mm["audio"]
|
||||
|
||||
@@ -238,8 +238,24 @@ def main(args: argparse.Namespace) -> None:
|
||||
if args.negative_prompt:
|
||||
prompt_dict["negative_prompt"] = args.negative_prompt
|
||||
|
||||
# Build cache-dit config if requested
|
||||
cache_config = None
|
||||
if args.cache_backend == "cache_dit":
|
||||
cache_config = {
|
||||
"Fn_compute_blocks": 1,
|
||||
"Bn_compute_blocks": 0,
|
||||
"max_warmup_steps": 4,
|
||||
"residual_diff_threshold": 0.24,
|
||||
"max_continuous_cached_steps": 3,
|
||||
"enable_taylorseer": False,
|
||||
"taylorseer_order": 1,
|
||||
"scm_steps_mask_policy": None,
|
||||
"scm_steps_policy": "dynamic",
|
||||
}
|
||||
|
||||
# Initialize Omni with multistage config
|
||||
print("\nInitializing Omni with multistage pipeline...")
|
||||
print(f" Cache backend: {args.cache_backend or 'None (no acceleration)'}")
|
||||
start_time = time.time()
|
||||
|
||||
omni = Omni(
|
||||
@@ -247,6 +263,9 @@ def main(args: argparse.Namespace) -> None:
|
||||
stage_configs_path=config_path,
|
||||
log_stats=args.enable_stats,
|
||||
stage_init_timeout=args.stage_init_timeout,
|
||||
cache_backend=args.cache_backend,
|
||||
cache_config=cache_config,
|
||||
enable_cache_dit_summary=getattr(args, "enable_cache_dit_summary", False),
|
||||
enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
|
||||
)
|
||||
|
||||
@@ -444,6 +463,20 @@ def parse_args() -> argparse.Namespace:
|
||||
help="Number of images to generate (default: 1)",
|
||||
)
|
||||
|
||||
# Cache acceleration
|
||||
parser.add_argument(
|
||||
"--cache-backend",
|
||||
type=str,
|
||||
default=None,
|
||||
choices=["cache_dit"],
|
||||
help="Cache backend for DiT acceleration. Default: None (no cache).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enable-cache-dit-summary",
|
||||
action="store_true",
|
||||
help="Enable cache-dit summary logging after diffusion forward passes.",
|
||||
)
|
||||
|
||||
# Runtime options
|
||||
parser.add_argument(
|
||||
"--enable-stats",
|
||||
|
||||
@@ -60,10 +60,13 @@ def main(args: argparse.Namespace) -> None:
|
||||
omni = Omni(
|
||||
model=args.model,
|
||||
enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
|
||||
mode="image-to-text",
|
||||
)
|
||||
|
||||
prompt = "<|startoftext|>You are an assistant that understands images and outputs text.<img>" + args.prompt
|
||||
|
||||
prompt_dict = {
|
||||
"prompt": args.prompt,
|
||||
"prompt": prompt,
|
||||
"modalities": ["text"],
|
||||
}
|
||||
|
||||
|
||||
@@ -10,7 +10,7 @@ MiMo-Audio provides multiple task variants for audio understanding and generatio
|
||||
- **tts_sft_with_instruct**: TTS generation with explicit voice style instructions.
|
||||
- **tts_sft_with_audio**: TTS generation with audio reference for voice cloning.
|
||||
- **tts_sft_with_natural_instruction**: TTS generation from natural language descriptions embedded in text.
|
||||
- **audio_trancribing_sft**: Transcribe audio to text (speech-to-text).
|
||||
- **audio_trancribing_sft**: Transcribe audio to text (speech-to-text). (note: the upstream task name uses the spelling 'trancribing', don't fix it)
|
||||
- **audio_understanding_sft**: Understand and analyze audio content with text queries.
|
||||
- **audio_understanding_sft_with_thinking**: Audio understanding with reasoning chain.
|
||||
- **spoken_dialogue_sft_multiturn**: Multi-turn spoken dialogue with audio input/output.
|
||||
@@ -29,6 +29,10 @@ The `MIMO_AUDIO_TOKENIZER_PATH` environment variable is mandatory due to the spe
|
||||
export MIMO_AUDIO_TOKENIZER_PATH="XiaomiMiMo/MiMo-Audio-Tokenizer"
|
||||
```
|
||||
|
||||
### Flash Attention (audio generation)
|
||||
|
||||
For **audio generation** (e.g. TTS variants, multi-turn spoken dialogue with audio output), install the **`flash-attn`** package with a build that matches your **CUDA** and **PyTorch** versions. On GPU, omitting **`flash-attn`** can cause **generated audio to be noise-only or otherwise unusable**. See the [FlashAttention](https://github.com/Dao-AILab/flash-attention) project for installation options and prebuilt wheels.
|
||||
|
||||
## Quick Start
|
||||
|
||||
Run a single sample for basic TTS:
|
||||
|
||||
@@ -189,7 +189,6 @@ def main(args):
|
||||
batch_timeout=args.batch_timeout,
|
||||
init_timeout=args.init_timeout,
|
||||
shm_threshold_bytes=args.shm_threshold_bytes,
|
||||
enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
|
||||
)
|
||||
|
||||
thinker_sampling_params = SamplingParams(
|
||||
@@ -434,11 +433,6 @@ def parse_args():
|
||||
default="../../../model_executor/stage_configs/mimo_audio.yaml",
|
||||
help="Path to a stage configs file.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enable-diffusion-pipeline-profiler",
|
||||
action="store_true",
|
||||
help="Enable diffusion pipeline profiler to display stage durations.",
|
||||
)
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
@@ -327,7 +327,6 @@ def main(args):
|
||||
batch_timeout=args.batch_timeout,
|
||||
init_timeout=args.init_timeout,
|
||||
shm_threshold_bytes=args.shm_threshold_bytes,
|
||||
enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
|
||||
)
|
||||
thinker_sampling_params = SamplingParams(
|
||||
temperature=0.0, # Deterministic - no randomness
|
||||
@@ -541,11 +540,6 @@ def parse_args():
|
||||
default=False,
|
||||
help="Use py_generator mode. The returned type of Omni.generate() is a Python Generator object.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enable-diffusion-pipeline-profiler",
|
||||
action="store_true",
|
||||
help="Enable diffusion pipeline profiler to display stage durations.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
|
||||
@@ -263,7 +263,7 @@ query_map = {
|
||||
|
||||
|
||||
def main(args):
|
||||
model_name = "Qwen/Qwen3-Omni-30B-A3B-Instruct"
|
||||
model_name = args.model
|
||||
print("=" * 20, "\n", f"vllm version: {vllm.__version__}", "\n", "=" * 20)
|
||||
|
||||
# Get paths from args
|
||||
@@ -296,10 +296,11 @@ def main(args):
|
||||
|
||||
omni = Omni(
|
||||
model=model_name,
|
||||
dtype=args.dtype,
|
||||
stage_configs_path=args.stage_configs_path,
|
||||
log_stats=args.log_stats,
|
||||
stage_init_timeout=args.stage_init_timeout,
|
||||
enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
|
||||
init_timeout=args.init_timeout,
|
||||
)
|
||||
|
||||
thinker_sampling_params = SamplingParams(
|
||||
@@ -333,11 +334,14 @@ def main(args):
|
||||
repetition_penalty=1.1,
|
||||
)
|
||||
|
||||
sampling_params_list = [
|
||||
all_sampling_params = [
|
||||
thinker_sampling_params,
|
||||
talker_sampling_params, # code predictor is integrated into talker for Qwen3 Omni
|
||||
code2wav_sampling_params,
|
||||
]
|
||||
# Match sampling params to the number of configured stages
|
||||
num_stages = omni.num_stages
|
||||
sampling_params_list = all_sampling_params[:num_stages]
|
||||
|
||||
if args.txt_prompts is None:
|
||||
prompts = [query_result.inputs for _ in range(args.num_prompts)]
|
||||
@@ -353,9 +357,9 @@ def main(args):
|
||||
for i, prompt in enumerate(prompts):
|
||||
prompt["modalities"] = output_modalities
|
||||
|
||||
profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
|
||||
profiler_enabled = args.enable_profiler
|
||||
if profiler_enabled:
|
||||
omni.start_profile(stages=[0])
|
||||
omni.start_profile(stages=args.profiler_stages)
|
||||
omni_generator = omni.generate(prompts, sampling_params_list, py_generator=args.py_generator)
|
||||
# Determine output directory: prefer --output-dir; fallback to --output-wav
|
||||
output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav
|
||||
@@ -405,7 +409,7 @@ def main(args):
|
||||
if profiler_enabled and processed_count >= total_requests:
|
||||
print(f"[Info] Processed {processed_count}/{total_requests}. Stopping profiler inside active loop...")
|
||||
# Stop the profiler while workers are still alive
|
||||
omni.stop_profile()
|
||||
omni.stop_profile(stages=args.profiler_stages)
|
||||
|
||||
print("[Info] Waiting 30s for workers to write trace files to disk...")
|
||||
time.sleep(30)
|
||||
@@ -415,6 +419,12 @@ def main(args):
|
||||
|
||||
def parse_args():
|
||||
parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
default="Qwen/Qwen3-Omni-30B-A3B-Instruct",
|
||||
help="Model name or path.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--query-type",
|
||||
"-q",
|
||||
@@ -528,9 +538,23 @@ def parse_args():
|
||||
help="Use py_generator mode. The returned type of Omni.generate() is a Python Generator object.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enable-diffusion-pipeline-profiler",
|
||||
"--enable-profiler",
|
||||
action="store_true",
|
||||
help="Enable diffusion pipeline profiler to display stage durations.",
|
||||
default=False,
|
||||
help="Enables profiling when set.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--profiler-stages",
|
||||
type=int,
|
||||
nargs="*",
|
||||
default=None,
|
||||
help="List of stage IDs to profile. If not set, profiles all stages.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dtype",
|
||||
type=str,
|
||||
default="auto",
|
||||
help="Model dtype (auto, half, float16, bfloat16, float, float32).",
|
||||
)
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
@@ -382,7 +382,6 @@ async def run_all(args):
|
||||
stage_configs_path=args.stage_configs_path,
|
||||
log_stats=args.log_stats,
|
||||
stage_init_timeout=args.stage_init_timeout,
|
||||
enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
|
||||
)
|
||||
|
||||
# Use default sampling params from stage config (they are pre-configured
|
||||
@@ -585,11 +584,6 @@ def parse_args():
|
||||
default=16000,
|
||||
help="Sampling rate for audio loading.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enable-diffusion-pipeline-profiler",
|
||||
action="store_true",
|
||||
help="Enable diffusion pipeline profiler to display stage durations.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
|
||||
@@ -8,7 +8,7 @@ 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.
|
||||
- **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.
|
||||
|
||||
@@ -371,7 +371,6 @@ def main(args):
|
||||
stage_configs_path=args.stage_configs_path,
|
||||
log_stats=args.log_stats,
|
||||
stage_init_timeout=args.stage_init_timeout,
|
||||
enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
|
||||
)
|
||||
|
||||
batch_size = args.batch_size
|
||||
@@ -393,7 +392,6 @@ async def main_streaming(args):
|
||||
stage_configs_path=args.stage_configs_path,
|
||||
log_stats=args.log_stats,
|
||||
stage_init_timeout=args.stage_init_timeout,
|
||||
enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
|
||||
)
|
||||
|
||||
for i, prompt in enumerate(inputs):
|
||||
@@ -535,11 +533,6 @@ def parse_args():
|
||||
default=1,
|
||||
help="Number of prompts per batch (default: 1, sequential).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enable-diffusion-pipeline-profiler",
|
||||
action="store_true",
|
||||
help="Enable diffusion pipeline profiler to display stage durations.",
|
||||
)
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
@@ -1,13 +1,49 @@
|
||||
# Text-To-Image
|
||||
|
||||
This folder provides several entrypoints for experimenting with `Qwen/Qwen-Image` `Qwen/Qwen-Image-2512` `Tongyi-MAI/Z-Image-Turbo` `stepfun-ai/NextStep-1.1` using vLLM-Omni, note that NextStep-1.1 has different architecture so we treat it differently regarding running arguments and pipeline.
|
||||
Generate images from text prompts using vLLM-Omni's diffusion pipeline entrypoints.
|
||||
|
||||
- `text_to_image.py`: command-line script for single image generation with advanced options.
|
||||
- `web_demo.py`: lightweight Gradio UI for interactive prompt/seed/CFG exploration.
|
||||
- `gradio_demo.py`: lightweight Gradio UI for interactive prompt/seed/CFG exploration.
|
||||
|
||||
Note that when you pass in multiple independent prompts, they will be processed sequentially. Batching requests is currently not supported.
|
||||
## Table of Contents
|
||||
|
||||
## Basic Usage
|
||||
- [Overview](#overview)
|
||||
- [Quick Start](#quick-start)
|
||||
- [Key Arguments](#key-arguments)
|
||||
- [More CLI Examples](#more-cli-examples)
|
||||
- [Web UI Demo](#web-ui-demo)
|
||||
|
||||
## Overview
|
||||
|
||||
This folder provides several entrypoints for experimenting with text-to-image diffusion models using vLLM-Omni. Note that `NextStep-1.1` has a different architecture, so it is treated differently regarding running arguments and pipeline.
|
||||
|
||||
### Supported Models
|
||||
|
||||
| Model | Image Shape | Peak VRAM (GiB) * | Model Weights (GiB) |
|
||||
| ----- | ----------- | ----------- | ----------------- |
|
||||
| `Qwen/Qwen-Image` | 1024 x 1024 | 60.0 | 53.7 |
|
||||
| `Qwen/Qwen-Image-2512` |1024 x 1024 | 60.0 | 53.7 |
|
||||
| `Tongyi-MAI/Z-Image-Turbo` | 1024 x 1024 | 24.8 | 19.2 |
|
||||
| `stepfun-ai/NextStep-1.1` | 512 x 512 | 71.8 | 28.1 |
|
||||
| `meituan-longcat/LongCat-Image` | 1024 x 1024 | 71.2 | 27.3 |
|
||||
| `AIDC-AI/Ovis-Image-7B` | 1024 x 1024 | 71.8 | 17.1 |
|
||||
| `OmniGen2/OmniGen2` | 1024 x 1024 | 20.1 | 14.7 |
|
||||
| `stabilityai/stable-diffusion-3.5-medium` | 1024 x 1024 | 20.1 | 15.6 |
|
||||
| `black-forest-labs/FLUX.1-dev` | 1024 x 1024 | 77.6 | 31.4 |
|
||||
| `black-forest-labs/FLUX.2-klein-4B` | 1024 x 1024 | 72.7 | 14.9 |
|
||||
| `black-forest-labs/FLUX.2-klein-9B` | 1024 x 1024 | 37.1 | 32.3 |
|
||||
| `black-forest-labs/FLUX.2-dev` | 1024 x 1024 | 65.7 | >80 (CPU offload required) |
|
||||
|
||||
!!! info
|
||||
*Peak VRAM: based on basic single-card usage, batch size =1, without any acceleration/optimization features. FLUX.2-dev requires `--enable-cpu-offload` on a single 80 GiB GPU.
|
||||
|
||||
Default model: `Qwen/Qwen-Image`
|
||||
|
||||
## Quick Start
|
||||
|
||||
### Python API
|
||||
|
||||
Single-prompt generation:
|
||||
|
||||
```python
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
@@ -20,78 +56,90 @@ if __name__ == "__main__":
|
||||
images[0].save("coffee.png")
|
||||
```
|
||||
|
||||
Or put more than one prompt in a request.
|
||||
### Local CLI Usage
|
||||
|
||||
```python
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
if __name__ == "__main__":
|
||||
omni = Omni(model="Qwen/Qwen-Image")
|
||||
prompts = [
|
||||
"a cup of coffee on a table",
|
||||
"a toy dinosaur on a sandy beach",
|
||||
"a fox waking up in bed and yawning",
|
||||
]
|
||||
outputs = omni.generate(prompts)
|
||||
for i, output in enumerate(outputs):
|
||||
image = output.request_output.images[0].save(f"{i}.jpg")
|
||||
```bash
|
||||
python text_to_image.py \
|
||||
--model Qwen/Qwen-Image \
|
||||
--prompt "a cup of coffee on the table" \
|
||||
--output coffee.png
|
||||
```
|
||||
|
||||
!!! info
|
||||
## Key Arguments
|
||||
|
||||
However, it is not currently recommended to do so
|
||||
because not all models support batch inference,
|
||||
and batch requesting mostly does not provide significant performance improvement (despite the impression that it does).
|
||||
This feature is primarily for the sake of interface compatibility with vLLM and to allow for future improvements.
|
||||
**Common arguments:**
|
||||
|
||||
!!! info
|
||||
| Argument | Type | Default | Description |
|
||||
| -------- | ---- | ------- | ----------- |
|
||||
| `--prompt` | str | `"a cup of coffee on the table"` | Text description for image generation |
|
||||
| `--seed` | int | `142` | Integer seed for deterministic sampling |
|
||||
| `--negative-prompt` | str | `None` | Negative prompt for classifier-free conditional guidance |
|
||||
| `--cfg-scale` | float | `4.0` | True CFG scale (model-specific guidance strength) |
|
||||
| `--guidance-scale` | float | `1.0` | Classifier-free guidance scale |
|
||||
| `--num-images-per-prompt` | int | `1` | Number of images per prompt (saved as `output`, `output_1`, ...) |
|
||||
| `--num-inference-steps` | int | `50` | Diffusion sampling steps (more steps = higher quality, slower) |
|
||||
| `--height` | int | `1024` | Output image height in pixels |
|
||||
| `--width` | int | `1024` | Output image width in pixels |
|
||||
| `--output` | str | `"qwen_image_output.png"` | Path to save the generated image |
|
||||
| `--vae-use-slicing` | flag | off | Enable VAE slicing for memory optimization |
|
||||
| `--vae-use-tiling` | flag | off | Enable VAE tiling for memory optimization |
|
||||
| `--cfg-parallel-size` | int | `1` | Set to `2` to enable CFG Parallel |
|
||||
| `--ulysses-degree` | int | `1` | Ulysses sequence parallel degree for multi-GPU inference |
|
||||
| `--ring-degree` | int | `1` | Ring sequence parallel degree for hybrid Ulysses + Ring inference |
|
||||
| `--ulysses-mode` | str | `"strict"` | Ulysses SP mode: `"strict"` or `"advanced_uaa"` |
|
||||
| `--enable-cpu-offload` | flag | off | Enable CPU offloading for diffusion models |
|
||||
| `--lora-path` | str | — | Path to PEFT LoRA adapter folder |
|
||||
| `--lora-scale` | float | `1.0` | Scale factor for LoRA weights |
|
||||
|
||||
For diffusion pipelines, the stage config field `stage_args.[].engine_args.max_num_seqs` is 1 by default, and the input
|
||||
list is sliced into single-item requests before feeding into the diffusion pipeline. For models that do internally support
|
||||
batched inputs, you can [modify this configuration](../../../configuration/stage_configs.md) to let the model accept a longer batch of prompts.
|
||||
**NextStep-1.1 specific arguments:**
|
||||
|
||||
Apart from string prompt, vLLM-Omni also supports dictionary prompts in the same style as vLLM.
|
||||
This is useful for models that support negative prompts.
|
||||
| Argument | Type | Default | Description |
|
||||
| -------- | ---- | ------- | ----------- |
|
||||
| `--guidance-scale-2` | float | `1.0` | Secondary guidance scale (e.g. image-level CFG) |
|
||||
| `--timesteps-shift` | float | `1.0` | Timesteps shift parameter for sampling |
|
||||
| `--cfg-schedule` | str | `"constant"` | CFG schedule type: `"constant"` or `"linear"` |
|
||||
| `--use-norm` | flag | off | Apply layer normalization to sampled tokens |
|
||||
|
||||
```python
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
> If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.
|
||||
|
||||
if __name__ == "__main__":
|
||||
omni = Omni(model="Qwen/Qwen-Image")
|
||||
outputs = omni.generate([
|
||||
{
|
||||
"prompt": "a cup of coffee on a table",
|
||||
"negative_prompt": "low resolution"
|
||||
},
|
||||
{
|
||||
"prompt": "a toy dinosaur on a sandy beach",
|
||||
"negative_prompt": "cinematic, realistic"
|
||||
}
|
||||
])
|
||||
for i, output in enumerate(outputs):
|
||||
image = output.request_output.images[0].save(f"{i}.jpg")
|
||||
```
|
||||
> Qwen-Image currently publishes best-effort presets at `1328x1328`, `1664x928`, `928x1664`, `1472x1140`, `1140x1472`, `1584x1056`, and `1056x1584`. Adjust `--height/--width` accordingly for the most reliable outcomes.
|
||||
|
||||
## Local CLI Usage
|
||||
## More CLI Examples
|
||||
|
||||
### Qwen/Tongyi Models
|
||||
### Tongyi Models
|
||||
|
||||
```bash
|
||||
python text_to_image.py \
|
||||
--model Tongyi-MAI/Z-Image-Turbo \
|
||||
--prompt "a cup of coffee on the table" \
|
||||
--seed 42 \
|
||||
--cfg-scale 4.0 \
|
||||
--guidance-scale 0.0 \
|
||||
--num-images-per-prompt 1 \
|
||||
--num-inference-steps 50 \
|
||||
--num-inference-steps 9 \
|
||||
--height 1024 \
|
||||
--width 1024 \
|
||||
--output outputs/coffee.png
|
||||
```
|
||||
|
||||
`Tongyi-MAI/Z-Image-Turbo` is a distilled version of Z-Image. Distilled diffusion models usually require less number of inference steps (4~9), and Classifier-Free Guidance (CFG) is usually NOT applied. Similar distilled models are `black-forest-labs/FLUX.2-klein-4B` and `black-forest-labs/FLUX.2-klein-9B`.
|
||||
|
||||
Advanced UAA example (requires 2 GPUs):
|
||||
|
||||
```bash
|
||||
python text_to_image.py \
|
||||
--model Tongyi-MAI/Z-Image-Turbo \
|
||||
--prompt "a cup of coffee on the table" \
|
||||
--ulysses-degree 2 \
|
||||
--ulysses-mode advanced_uaa \
|
||||
--height 1024 \
|
||||
--width 1024 \
|
||||
--output outputs/coffee_hybrid.png
|
||||
```
|
||||
|
||||
### NextStep Models
|
||||
|
||||
NextStep-1.1 can have extra arguments
|
||||
NextStep-1.1 supports extra arguments for dual-level CFG control:
|
||||
|
||||
```bash
|
||||
python text_to_image.py \
|
||||
--model stepfun-ai/NextStep-1.1 \
|
||||
@@ -106,8 +154,10 @@ python text_to_image.py \
|
||||
--seed 42
|
||||
```
|
||||
|
||||
### Flux.2-dev Models
|
||||
To start Flux.2-dev with a single GPU, cpu-offload must be enabled because the total size of its weights exceeds the 80GB memory capacity of the GPU.
|
||||
### FLUX.2-dev Models
|
||||
|
||||
To run FLUX.2-dev on a single GPU, `--enable-cpu-offload` is required because the model weights exceed 80 GiB:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_image/text_to_image.py \
|
||||
--model black-forest-labs/FLUX.2-dev \
|
||||
@@ -123,38 +173,82 @@ python examples/offline_inference/text_to_image/text_to_image.py \
|
||||
--output flux2-dev.png
|
||||
```
|
||||
|
||||
### Key Arguments
|
||||
### Batch Requests (Multiple Prompts)
|
||||
|
||||
**Common arguments:**
|
||||
You can pass multiple prompts in a single `generate` call.
|
||||
|
||||
- `--prompt`: text description (string).
|
||||
- `--seed`: integer seed for deterministic sampling.
|
||||
- `--cfg-scale`: true CFG scale (model-specific guidance strength).
|
||||
- `--num-images-per-prompt`: number of images to generate per prompt (saves as `output`, `output_1`, ...).
|
||||
- `--num-inference-steps`: diffusion sampling steps (more steps = higher quality, slower).
|
||||
- `--height/--width`: output resolution (defaults 1024x1024).
|
||||
- `--output`: path to save the generated PNG.
|
||||
- `--vae-use-slicing`: enable VAE slicing for memory optimization.
|
||||
- `--vae-use-tiling`: enable VAE tiling for memory optimization.
|
||||
- `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion_acceleration.md#using-cfg-parallel).
|
||||
- `--enable-cpu-offload`: enable CPU offloading for diffusion models.
|
||||
- `--guidance-scale`: classifier-free guidance scale.
|
||||
```python
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
**NextStep-1.1 specific:**
|
||||
- `--guidance-scale-2`: secondary guidance scale, e.g. image-level CFG (default: 1.0).
|
||||
- `--timesteps-shift`: timesteps shift parameter for sampling (default: 1.0).
|
||||
- `--cfg-schedule`: CFG schedule type, "constant" or "linear" (default: "constant").
|
||||
- `--use-norm`: apply layer normalization to sampled tokens.
|
||||
if __name__ == "__main__":
|
||||
omni = Omni(model="Qwen/Qwen-Image")
|
||||
prompts = [
|
||||
"a cup of coffee on a table",
|
||||
"a toy dinosaur on a sandy beach",
|
||||
"a fox waking up in bed and yawning",
|
||||
]
|
||||
outputs = omni.generate(prompts)
|
||||
for i, output in enumerate(outputs):
|
||||
output.request_output.images[0].save(f"{i}.jpg")
|
||||
```
|
||||
|
||||
> ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.
|
||||
!!! info
|
||||
|
||||
> ℹ️ Qwen-Image currently publishes best-effort presets at `1328x1328`, `1664x928`, `928x1664`, `1472x1140`, `1140x1472`, `1584x1056`, and `1056x1584`. Adjust `--height/--width` accordingly for the most reliable outcomes.
|
||||
Not all models support batch inference, and batch requesting mostly does not provide significant
|
||||
performance improvement. This feature is primarily for interface compatibility with vLLM and to
|
||||
allow for future improvements.
|
||||
|
||||
## LoRA
|
||||
!!! info
|
||||
|
||||
This example supports Peft-compatible LoRA (Low-Rank Adaptation) adapters for diffusion models. Pass `--lora-path` to use a LoRA adapter and optionally `--lora-scale` (default 1.0); omit it to use the base model only.
|
||||
For diffusion pipelines, the stage config field `stage_args.[].runtime.max_batch_size` is 1 by
|
||||
default, and the input list is sliced into single-item requests before feeding into the diffusion
|
||||
pipeline. For models that do internally support batched inputs, you can
|
||||
[modify this configuration](../../../configuration/stage_configs.md) to let the model accept a
|
||||
longer batch of prompts.
|
||||
|
||||
### Basic usage with LoRA
|
||||
### Negative Prompts
|
||||
|
||||
vLLM-Omni supports dictionary prompts for models that accept negative prompts:
|
||||
|
||||
```python
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
if __name__ == "__main__":
|
||||
omni = Omni(model="Qwen/Qwen-Image")
|
||||
outputs = omni.generate([
|
||||
{
|
||||
"prompt": "a cup of coffee on a table",
|
||||
"negative_prompt": "low resolution"
|
||||
},
|
||||
{
|
||||
"prompt": "a toy dinosaur on a sandy beach",
|
||||
"negative_prompt": "cinematic, realistic"
|
||||
}
|
||||
])
|
||||
for i, output in enumerate(outputs):
|
||||
output.request_output.images[0].save(f"{i}.jpg")
|
||||
```
|
||||
|
||||
You can also pass a negative prompt via the CLI argument `--negative-prompt`:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_image/text_to_image.py \
|
||||
--model Qwen/Qwen-Image \
|
||||
--prompt "a cup of coffee on a table" \
|
||||
--negative-prompt "low resolution, blurry" \
|
||||
--output coffee.png
|
||||
```
|
||||
|
||||
### Advanced Features
|
||||
|
||||
#### CFG Parallel
|
||||
|
||||
Set `--cfg-parallel-size 2` to enable CFG Parallel for faster inference on multi-GPU setups.
|
||||
See more examples in the [diffusion acceleration user guide](../../../docs/user_guide/diffusion_acceleration.md#using-cfg-parallel).
|
||||
|
||||
#### LoRA
|
||||
|
||||
This example supports PEFT-compatible LoRA (Low-Rank Adaptation) adapters for diffusion models. Pass `--lora-path` to use a LoRA adapter and optionally `--lora-scale` (default `1.0`); omit it to use the base model only.
|
||||
|
||||
```bash
|
||||
python text_to_image.py \
|
||||
@@ -165,14 +259,7 @@ python text_to_image.py \
|
||||
--output output.png
|
||||
```
|
||||
|
||||
### LoRA parameters
|
||||
|
||||
- `--lora-path`: Path to LoRA adapter folder (PEFT format). Loaded at initialization and used for generation.
|
||||
- `--lora-scale`: Scale factor for LoRA weights (default: 1.0). Higher values increase the influence of the LoRA adapter.
|
||||
|
||||
### LoRA adapter format
|
||||
|
||||
LoRA adapters must be in PEFT (Parameter-Efficient Fine-Tuning) format. A typical LoRA adapter directory structure:
|
||||
LoRA adapters must be in PEFT format. A typical adapter directory structure:
|
||||
|
||||
```
|
||||
lora_adapter/
|
||||
@@ -182,10 +269,10 @@ lora_adapter/
|
||||
|
||||
## Web UI Demo
|
||||
|
||||
Launch the gradio demo:
|
||||
Launch the Gradio demo:
|
||||
|
||||
```bash
|
||||
python gradio_demo.py --port 7862
|
||||
```
|
||||
|
||||
Then open `http://localhost:7862/` on your local browser to interact with the web UI.
|
||||
Then open `http://localhost:7862/` in your local browser to interact with the web UI.
|
||||
@@ -36,7 +36,17 @@ def parse_args() -> argparse.Namespace:
|
||||
"--model",
|
||||
default="Qwen/Qwen-Image",
|
||||
help="Diffusion model name or local path. Supported models: "
|
||||
"Qwen/Qwen-Image, Tongyi-MAI/Z-Image-Turbo, Qwen/Qwen-Image-2512, stepfun-ai/NextStep-1.1",
|
||||
"Qwen/Qwen-Image, Tongyi-MAI/Z-Image-Turbo, Qwen/Qwen-Image-2512, stepfun-ai/NextStep-1.1, "
|
||||
"black-forest-labs/FLUX.1-dev, black-forest-labs/FLUX.2-klein-9B, "
|
||||
"black-forest-labs/FLUX.2-dev, tencent/HunyuanImage-3.0-Instruct, "
|
||||
"meituan-longcat/LongCat-Image, OvisAI/Ovis-Image, "
|
||||
"stabilityai/stable-diffusion-3.5-medium, Tongyi-MAI/Z-Image-Turbo and etc.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stage-configs-path",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to a YAML file containing stage configurations for Omni.",
|
||||
)
|
||||
parser.add_argument("--prompt", default="a cup of coffee on the table", help="Text prompt for image generation.")
|
||||
parser.add_argument(
|
||||
@@ -99,6 +109,13 @@ def parse_args() -> argparse.Namespace:
|
||||
default=1,
|
||||
help="Number of GPUs used for ulysses sequence parallelism.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ulysses-mode",
|
||||
type=str,
|
||||
default="strict",
|
||||
choices=["strict", "advanced_uaa"],
|
||||
help="Ulysses sequence-parallel mode: 'strict' (divisibility required) or 'advanced_uaa' (UAA).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ring-degree",
|
||||
type=int,
|
||||
@@ -220,6 +237,11 @@ def parse_args() -> argparse.Namespace:
|
||||
action="store_true",
|
||||
help="Enable diffusion pipeline profiler to display stage durations.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--log-stats",
|
||||
action="store_true",
|
||||
help="Enable logging of diffusion pipeline stats.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
@@ -262,6 +284,7 @@ def main():
|
||||
parallel_config = DiffusionParallelConfig(
|
||||
ulysses_degree=args.ulysses_degree,
|
||||
ring_degree=args.ring_degree,
|
||||
ulysses_mode=args.ulysses_mode,
|
||||
cfg_parallel_size=args.cfg_parallel_size,
|
||||
tensor_parallel_size=args.tensor_parallel_size,
|
||||
vae_patch_parallel_size=args.vae_patch_parallel_size,
|
||||
@@ -307,10 +330,14 @@ def main():
|
||||
"parallel_config": parallel_config,
|
||||
"enforce_eager": args.enforce_eager,
|
||||
"enable_cpu_offload": args.enable_cpu_offload,
|
||||
"mode": "text-to-image",
|
||||
"log_stats": args.log_stats,
|
||||
"enable_diffusion_pipeline_profiler": args.enable_diffusion_pipeline_profiler,
|
||||
**lora_args,
|
||||
**quant_kwargs,
|
||||
}
|
||||
if args.stage_configs_path:
|
||||
omni_kwargs["stage_configs_path"] = args.stage_configs_path
|
||||
if use_nextstep:
|
||||
# NextStep-1.1 requires explicit pipeline class
|
||||
omni_kwargs["model_class_name"] = "NextStep11Pipeline"
|
||||
@@ -331,13 +358,17 @@ def main():
|
||||
print(f" Ignored layers: {ignored_layers}")
|
||||
print(
|
||||
f" Parallel configuration: tensor_parallel_size={args.tensor_parallel_size}, "
|
||||
f"ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}, "
|
||||
f"vae_patch_parallel_size={args.vae_patch_parallel_size}, enable_expert_parallel={args.enable_expert_parallel}."
|
||||
f"ulysses_degree={args.ulysses_degree}, ulysses_mode={args.ulysses_mode}, "
|
||||
f"ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}, "
|
||||
f"vae_patch_parallel_size={args.vae_patch_parallel_size}, "
|
||||
f"enable_expert_parallel={args.enable_expert_parallel}."
|
||||
)
|
||||
print(f" CPU offload: {args.enable_cpu_offload}")
|
||||
print(f" Image size: {args.width}x{args.height}")
|
||||
if args.lora_path:
|
||||
print(f" LoRA: scale={args.lora_scale}")
|
||||
if args.stage_configs_path:
|
||||
print(f" stage-configs-path: {args.stage_configs_path}")
|
||||
print(f"{'=' * 60}\n")
|
||||
|
||||
# Build LoRA request when --lora-path is set
|
||||
|
||||
@@ -58,7 +58,7 @@ python3 examples/offline_inference/voxtral_tts/end2end.py \
|
||||
| `--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) |
|
||||
| `--voice NAME` | Voice preset to use instead of reference audio file. Check Huggingface `mistralai/Voxtral-4B-TTS-2603` to get the list of available voices |
|
||||
| `--write-audio` | Write generated audio to WAV files |
|
||||
| `--profiling-mode` | Enable profiling mode (reduces max tokens to 50) |
|
||||
| `--log-stats` | Enable detailed statistics logging |
|
||||
@@ -69,7 +69,9 @@ async def run_streaming(inputs, sampling_params_list, model_name, args, output_d
|
||||
ttfa = None
|
||||
accumulated_sample = 0
|
||||
|
||||
async for stage_output in async_omni.generate(single_input, request_id, sampling_params_list):
|
||||
async for stage_output in async_omni.generate(
|
||||
single_input, request_id=request_id, sampling_params_list=sampling_params_list
|
||||
):
|
||||
mm_output = stage_output.multimodal_output
|
||||
finished = stage_output.finished
|
||||
if not mm_output or "audio" not in mm_output:
|
||||
|
||||
@@ -70,90 +70,41 @@ The default yaml configuration deploys AR on GPU 0 and DiT on GPU 1. You can use
|
||||
|
||||
### Text-to-Image
|
||||
|
||||
Generate images from text prompts:
|
||||
|
||||
**Using Python client**
|
||||
|
||||
```bash
|
||||
python openai_chat_client.py \
|
||||
--prompt "A photorealistic mountain landscape at sunset" \
|
||||
--height 1024 \
|
||||
--width 1024 \
|
||||
--output landscape.png
|
||||
```
|
||||
|
||||
**Using curl**
|
||||
|
||||
```bash
|
||||
curl -s http://localhost:8091/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"messages": [
|
||||
{"role": "user", "content": "A beautiful sunset over the ocean with sailing boats"}
|
||||
],
|
||||
"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
|
||||
```
|
||||
|
||||
Or use the script:
|
||||
|
||||
```bash
|
||||
# Or use the curl script:
|
||||
bash run_curl_text_to_image.sh "A futuristic city skyline at night"
|
||||
```
|
||||
|
||||
### Image-to-Image (Image Editing)
|
||||
|
||||
Edit images with text instructions:
|
||||
|
||||
**Using Python client**
|
||||
|
||||
```bash
|
||||
python openai_chat_client.py \
|
||||
--prompt "Convert this image to watercolor style" \
|
||||
--image input.png \
|
||||
--output watercolor.png
|
||||
```
|
||||
|
||||
**Using curl**
|
||||
|
||||
```bash
|
||||
IMG_B64=$(base64 < input.png | tr -d '\n')
|
||||
|
||||
curl -s http://localhost:8091/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d @- <<EOF | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > output.png
|
||||
{
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Convert this image to watercolor style"},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/png;base64,'$IMG_B64'"}}
|
||||
]
|
||||
}],
|
||||
"extra_body": {
|
||||
"height": 1024,
|
||||
"width": 1024,
|
||||
"num_inference_steps": 50,
|
||||
"guidance_scale": 1.5,
|
||||
"seed": 42
|
||||
}
|
||||
}
|
||||
EOF
|
||||
```
|
||||
|
||||
Or use the script:
|
||||
|
||||
```bash
|
||||
# Or use the curl script:
|
||||
bash run_curl_image_edit.sh input.png "Convert to watercolor style"
|
||||
```
|
||||
|
||||
## Generation Parameters (extra_body)
|
||||
For general-purpose request methods (curl, OpenAI SDK, Python `requests`), see
|
||||
the [Text-to-Image](../text_to_image/README.md) and
|
||||
[Image-to-Image](../image_to_image/README.md) guides.
|
||||
|
||||
## Generation Parameters
|
||||
|
||||
When using `/v1/chat/completions`, pass these inside `extra_body` in the curl
|
||||
JSON, or via the `extra_body` keyword argument in the OpenAI Python SDK.
|
||||
When using the dedicated `/v1/images/generations` or `/v1/images/edits`
|
||||
endpoints, pass the supported generation controls as top-level fields directly.
|
||||
For image dimensions and count, use `size` and `n` rather than `height` or
|
||||
`width`.
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
| --------------------- | ----- | ------- | ----------------------------------- |
|
||||
@@ -161,7 +112,7 @@ bash run_curl_image_edit.sh input.png "Convert to watercolor style"
|
||||
| `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 | 42 | Random seed for reproducibility |
|
||||
| `seed` | int | None | Optional random seed; `/v1/images/*` generates one server-side if omitted |
|
||||
| `negative_prompt` | str | None | Negative prompt |
|
||||
|
||||
## Response Format
|
||||
|
||||
@@ -0,0 +1,175 @@
|
||||
# Helios Video Generation via Online API
|
||||
|
||||
This example demonstrates how to use the `/v1/videos` API with Helios models using the generic `extra_params` field.
|
||||
|
||||
## Overview
|
||||
|
||||
The `/v1/videos` API now supports model-specific parameters through the `extra_params` field, which accepts a JSON object containing any model-specific configuration. This allows supporting new models like Helios without modifying the API.
|
||||
|
||||
## Helios Model Variants
|
||||
|
||||
- **Helios-Base**: Basic T2V/I2V/V2V generation (Stage 1 only)
|
||||
- **Helios-Mid**: Advanced generation with CFG-Zero* (Stage 2)
|
||||
- **Helios-Distilled**: Fast generation with DMD (Stage 2)
|
||||
|
||||
## API Usage Examples
|
||||
|
||||
### 1. T2V (Text-to-Video) - Helios-Base
|
||||
|
||||
Basic text-to-video generation:
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8000/v1/videos \
|
||||
-F "prompt=A serene lakeside sunrise with mist over the water." \
|
||||
-F "model=BestWishYsh/Helios-Base" \
|
||||
-F "width=640" \
|
||||
-F "height=384" \
|
||||
-F "num_frames=99" \
|
||||
-F "num_inference_steps=50" \
|
||||
-F "guidance_scale=5.0" \
|
||||
-F "seed=42"
|
||||
```
|
||||
|
||||
### 2. I2V (Image-to-Video) - Helios-Base
|
||||
|
||||
Generate video from an input image:
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8000/v1/videos \
|
||||
-F "prompt=The lake water gently ripples as morning mist rises." \
|
||||
-F "model=BestWishYsh/Helios-Base" \
|
||||
-F "input_reference=@/path/to/image.jpg" \
|
||||
-F "width=640" \
|
||||
-F "height=384" \
|
||||
-F "num_frames=99" \
|
||||
-F "guidance_scale=5.0"
|
||||
```
|
||||
|
||||
### 3. Helios-Mid with Stage 2 + CFG-Zero*
|
||||
|
||||
Advanced generation with pyramid multi-stage denoising:
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8000/v1/videos \
|
||||
-F "prompt=A serene lakeside sunrise with mist over the water." \
|
||||
-F "model=BestWishYsh/Helios-Mid" \
|
||||
-F "width=640" \
|
||||
-F "height=384" \
|
||||
-F "guidance_scale=5.0" \
|
||||
-F 'extra_params={
|
||||
"is_enable_stage2": true,
|
||||
"pyramid_num_stages": 3,
|
||||
"pyramid_num_inference_steps_list": [20, 20, 20],
|
||||
"use_cfg_zero_star": true,
|
||||
"use_zero_init": true,
|
||||
"zero_steps": 1
|
||||
}'
|
||||
```
|
||||
|
||||
### 4. Helios-Distilled with DMD
|
||||
|
||||
Fast generation with Distribution Matching Distillation:
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8000/v1/videos \
|
||||
-F "prompt=A serene lakeside sunrise with mist over the water." \
|
||||
-F "model=BestWishYsh/Helios-Distilled" \
|
||||
-F "width=640" \
|
||||
-F "height=384" \
|
||||
-F "guidance_scale=1.0" \
|
||||
-F 'extra_params={
|
||||
"is_enable_stage2": true,
|
||||
"pyramid_num_stages": 3,
|
||||
"pyramid_num_inference_steps_list": [2, 2, 2],
|
||||
"is_amplify_first_chunk": true
|
||||
}'
|
||||
```
|
||||
|
||||
## Model-Specific Parameters
|
||||
|
||||
The `extra_params` field accepts a JSON object with model-specific parameters. For Helios models, supported parameters include:
|
||||
|
||||
### Stage 2 Parameters
|
||||
- `is_enable_stage2` (bool): Enable pyramid multi-stage denoising
|
||||
- `pyramid_num_stages` (int): Number of pyramid stages (default: 3)
|
||||
- `pyramid_num_inference_steps_list` (array): Steps per stage, e.g., `[20, 20, 20]`
|
||||
|
||||
### CFG Zero Star Parameters (Helios-Mid)
|
||||
- `use_cfg_zero_star` (bool): Enable CFG Zero Star guidance
|
||||
- `use_zero_init` (bool): Use zero initialization for first steps
|
||||
- `zero_steps` (int): Number of initial zero prediction steps (default: 1)
|
||||
|
||||
### DMD Parameters (Helios-Distilled)
|
||||
- `is_amplify_first_chunk` (bool): Enable DMD amplification for first chunk
|
||||
|
||||
### Video Input (V2V mode)
|
||||
For video-to-video generation, upload a video file via `input_reference`:
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8000/v1/videos \
|
||||
-F "prompt=Transform the video into a watercolor painting style." \
|
||||
-F "model=BestWishYsh/Helios-Base" \
|
||||
-F "input_reference=@/path/to/video.mp4" \
|
||||
-F "guidance_scale=5.0"
|
||||
```
|
||||
|
||||
## Python Example
|
||||
|
||||
```python
|
||||
import requests
|
||||
import json
|
||||
|
||||
url = "http://localhost:8000/v1/videos"
|
||||
|
||||
# Helios-Mid with Stage 2
|
||||
data = {
|
||||
"prompt": "A serene lakeside sunrise with mist over the water.",
|
||||
"model": "BestWishYsh/Helios-Mid",
|
||||
"width": 640,
|
||||
"height": 384,
|
||||
"guidance_scale": 5.0,
|
||||
"extra_params": json.dumps({
|
||||
"is_enable_stage2": True,
|
||||
"pyramid_num_stages": 3,
|
||||
"pyramid_num_inference_steps_list": [20, 20, 20],
|
||||
"use_cfg_zero_star": True,
|
||||
"use_zero_init": True,
|
||||
"zero_steps": 1
|
||||
})
|
||||
}
|
||||
|
||||
response = requests.post(url, data=data)
|
||||
video_job = response.json()
|
||||
print(f"Video job created: {video_job['id']}")
|
||||
|
||||
# Poll for completion
|
||||
import time
|
||||
while True:
|
||||
status_response = requests.get(f"{url}/{video_job['id']}")
|
||||
status = status_response.json()
|
||||
|
||||
if status['status'] == 'completed':
|
||||
print(f"Video generated: {status['file_name']}")
|
||||
break
|
||||
elif status['status'] == 'failed':
|
||||
print(f"Generation failed: {status.get('error')}")
|
||||
break
|
||||
|
||||
print(f"Progress: {status['progress']}%")
|
||||
time.sleep(2)
|
||||
```
|
||||
|
||||
## Benefits of This Approach
|
||||
|
||||
1. **No API Changes**: New models can be supported without modifying the API
|
||||
2. **Backward Compatible**: Existing clients continue to work
|
||||
3. **Flexible**: Any model-specific parameter can be passed
|
||||
4. **Type-Safe**: Parameters are validated at the model level
|
||||
5. **Future-Proof**: Supports models that don't exist yet
|
||||
|
||||
## Notes
|
||||
|
||||
- The `extra_params` field must be a valid JSON object
|
||||
- Parameters are passed directly to the model's `extra_args`
|
||||
- Invalid parameters will be caught by the model implementation
|
||||
- Tensor inputs (images/videos) should use `input_reference`, not `extra_params`
|
||||
@@ -0,0 +1,212 @@
|
||||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
Helios video generation client example using the /v1/videos API.
|
||||
|
||||
This example demonstrates how to use the generic extra_params field
|
||||
to pass Helios-specific parameters without modifying the API.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import requests
|
||||
|
||||
|
||||
def create_video_job(
|
||||
api_url: str,
|
||||
prompt: str,
|
||||
model: str,
|
||||
width: int = 640,
|
||||
height: int = 384,
|
||||
num_frames: int = 99,
|
||||
guidance_scale: float = 5.0,
|
||||
seed: int = 42,
|
||||
extra_params: dict | None = None,
|
||||
input_image: Path | None = None,
|
||||
) -> dict:
|
||||
"""Create a video generation job."""
|
||||
data = {
|
||||
"prompt": prompt,
|
||||
"model": model,
|
||||
"width": width,
|
||||
"height": height,
|
||||
"num_frames": num_frames,
|
||||
"guidance_scale": guidance_scale,
|
||||
"seed": seed,
|
||||
}
|
||||
|
||||
files = {}
|
||||
if input_image:
|
||||
files["input_reference"] = open(input_image, "rb")
|
||||
|
||||
if extra_params:
|
||||
data["extra_params"] = json.dumps(extra_params)
|
||||
|
||||
response = requests.post(f"{api_url}/v1/videos", data=data, files=files)
|
||||
response.raise_for_status()
|
||||
|
||||
if files:
|
||||
files["input_reference"].close()
|
||||
|
||||
return response.json()
|
||||
|
||||
|
||||
def poll_video_status(api_url: str, video_id: str, poll_interval: int = 2) -> dict:
|
||||
"""Poll video generation status until completion."""
|
||||
print(f"Polling video job: {video_id}")
|
||||
|
||||
while True:
|
||||
response = requests.get(f"{api_url}/v1/videos/{video_id}")
|
||||
response.raise_for_status()
|
||||
status_data = response.json()
|
||||
|
||||
status = status_data["status"]
|
||||
progress = status_data.get("progress", 0)
|
||||
|
||||
print(f"Status: {status}, Progress: {progress}%")
|
||||
|
||||
if status == "completed":
|
||||
print("Video generation completed!")
|
||||
return status_data
|
||||
elif status == "failed":
|
||||
error = status_data.get("error", {})
|
||||
raise RuntimeError(f"Video generation failed: {error}")
|
||||
|
||||
time.sleep(poll_interval)
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Helios video generation client")
|
||||
parser.add_argument(
|
||||
"--api-url",
|
||||
default="http://localhost:8000",
|
||||
help="API server URL",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
default="BestWishYsh/Helios-Base",
|
||||
help="Model name (Helios-Base, Helios-Mid, Helios-Distilled)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompt",
|
||||
default="A serene lakeside sunrise with mist over the water.",
|
||||
help="Text prompt",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--width",
|
||||
type=int,
|
||||
default=640,
|
||||
help="Video width",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--height",
|
||||
type=int,
|
||||
default=384,
|
||||
help="Video height",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-frames",
|
||||
type=int,
|
||||
default=99,
|
||||
help="Number of frames",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--guidance-scale",
|
||||
type=float,
|
||||
default=5.0,
|
||||
help="Guidance scale",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--seed",
|
||||
type=int,
|
||||
default=42,
|
||||
help="Random seed",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--input-image",
|
||||
type=Path,
|
||||
help="Input image for I2V mode",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--preset",
|
||||
choices=["base", "mid-stage2", "distilled"],
|
||||
default="base",
|
||||
help="Helios preset configuration",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# Define preset configurations
|
||||
presets = {
|
||||
"base": None, # No extra params for base model
|
||||
"mid-stage2": {
|
||||
"is_enable_stage2": True,
|
||||
"pyramid_num_stages": 3,
|
||||
"pyramid_num_inference_steps_list": [20, 20, 20],
|
||||
"use_cfg_zero_star": True,
|
||||
"use_zero_init": True,
|
||||
"zero_steps": 1,
|
||||
},
|
||||
"distilled": {
|
||||
"is_enable_stage2": True,
|
||||
"pyramid_num_stages": 3,
|
||||
"pyramid_num_inference_steps_list": [2, 2, 2],
|
||||
"is_amplify_first_chunk": True,
|
||||
},
|
||||
}
|
||||
|
||||
extra_params = presets[args.preset]
|
||||
|
||||
print("=" * 50)
|
||||
print("Helios Video Generation")
|
||||
print("=" * 50)
|
||||
print(f"API URL: {args.api_url}")
|
||||
print(f"Model: {args.model}")
|
||||
print(f"Preset: {args.preset}")
|
||||
print(f"Prompt: {args.prompt}")
|
||||
print(f"Size: {args.width}x{args.height}")
|
||||
print(f"Frames: {args.num_frames}")
|
||||
print(f"Guidance Scale: {args.guidance_scale}")
|
||||
if extra_params:
|
||||
print(f"Extra Params: {json.dumps(extra_params, indent=2)}")
|
||||
print()
|
||||
|
||||
# Create video job
|
||||
print("Creating video generation job...")
|
||||
job = create_video_job(
|
||||
api_url=args.api_url,
|
||||
prompt=args.prompt,
|
||||
model=args.model,
|
||||
width=args.width,
|
||||
height=args.height,
|
||||
num_frames=args.num_frames,
|
||||
guidance_scale=args.guidance_scale,
|
||||
seed=args.seed,
|
||||
extra_params=extra_params,
|
||||
input_image=args.input_image,
|
||||
)
|
||||
|
||||
video_id = job["id"]
|
||||
print(f"Video job created: {video_id}")
|
||||
print()
|
||||
|
||||
# Poll for completion
|
||||
result = poll_video_status(args.api_url, video_id)
|
||||
|
||||
print()
|
||||
print("=" * 50)
|
||||
print("Generation Complete")
|
||||
print("=" * 50)
|
||||
print(f"Video ID: {result['id']}")
|
||||
print(f"File: {result.get('file_name', 'N/A')}")
|
||||
print(f"Inference Time: {result.get('inference_time_s', 'N/A')}s")
|
||||
print()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,79 @@
|
||||
#!/bin/bash
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
# Helios-Distilled with DMD Example
|
||||
# Fast generation with Distribution Matching Distillation
|
||||
|
||||
API_URL="${API_URL:-http://localhost:8000}"
|
||||
MODEL="${MODEL:-BestWishYsh/Helios-Distilled}"
|
||||
PROMPT="${PROMPT:-A serene lakeside sunrise with mist over the water.}"
|
||||
|
||||
echo "==================================="
|
||||
echo "Helios-Distilled with DMD"
|
||||
echo "==================================="
|
||||
echo "API URL: $API_URL"
|
||||
echo "Model: $MODEL"
|
||||
echo "Prompt: $PROMPT"
|
||||
echo ""
|
||||
|
||||
# Model-specific parameters for Helios-Distilled
|
||||
extra_params='{
|
||||
"is_enable_stage2": true,
|
||||
"pyramid_num_stages": 3,
|
||||
"pyramid_num_inference_steps_list": [2, 2, 2],
|
||||
"is_amplify_first_chunk": true
|
||||
}'
|
||||
|
||||
echo "Model extra params: $extra_params"
|
||||
echo ""
|
||||
|
||||
# Create video generation job
|
||||
echo "Creating video generation job..."
|
||||
RESPONSE=$(curl -s -X POST "$API_URL/v1/videos" \
|
||||
-F "prompt=$PROMPT" \
|
||||
-F "model=$MODEL" \
|
||||
-F "width=640" \
|
||||
-F "height=384" \
|
||||
-F "guidance_scale=1.0" \
|
||||
-F "seed=42" \
|
||||
-F "extra_params=$extra_params")
|
||||
|
||||
echo "Response: $RESPONSE"
|
||||
echo ""
|
||||
|
||||
# Extract video ID
|
||||
VIDEO_ID=$(echo "$RESPONSE" | grep -o '"id":"[^"]*"' | cut -d'"' -f4)
|
||||
|
||||
if [ -z "$VIDEO_ID" ]; then
|
||||
echo "Error: Failed to create video job"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "Video job created: $VIDEO_ID"
|
||||
echo ""
|
||||
|
||||
# Poll for completion
|
||||
echo "Polling for completion..."
|
||||
while true; do
|
||||
STATUS_RESPONSE=$(curl -s "$API_URL/v1/videos/$VIDEO_ID")
|
||||
STATUS=$(echo "$STATUS_RESPONSE" | grep -o '"status":"[^"]*"' | cut -d'"' -f4)
|
||||
PROGRESS=$(echo "$STATUS_RESPONSE" | grep -o '"progress":[0-9]*' | cut -d':' -f2)
|
||||
|
||||
echo "Status: $STATUS, Progress: $PROGRESS%"
|
||||
|
||||
if [ "$STATUS" = "completed" ]; then
|
||||
echo ""
|
||||
echo "Video generation completed!"
|
||||
echo "Full response:"
|
||||
echo "$STATUS_RESPONSE" | jq '.'
|
||||
break
|
||||
elif [ "$STATUS" = "failed" ]; then
|
||||
echo ""
|
||||
echo "Video generation failed!"
|
||||
echo "$STATUS_RESPONSE" | jq '.'
|
||||
exit 1
|
||||
fi
|
||||
|
||||
sleep 2
|
||||
done
|
||||
@@ -0,0 +1,81 @@
|
||||
#!/bin/bash
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
# Helios-Mid with Stage 2 + CFG-Zero* Example
|
||||
# This demonstrates advanced generation with pyramid multi-stage denoising
|
||||
|
||||
API_URL="${API_URL:-http://localhost:8000}"
|
||||
MODEL="${MODEL:-BestWishYsh/Helios-Mid}"
|
||||
PROMPT="${PROMPT:-A serene lakeside sunrise with mist over the water.}"
|
||||
|
||||
echo "==================================="
|
||||
echo "Helios-Mid Stage 2 + CFG-Zero*"
|
||||
echo "==================================="
|
||||
echo "API URL: $API_URL"
|
||||
echo "Model: $MODEL"
|
||||
echo "Prompt: $PROMPT"
|
||||
echo ""
|
||||
|
||||
# Model-specific parameters for Helios-Mid
|
||||
extra_params='{
|
||||
"is_enable_stage2": true,
|
||||
"pyramid_num_stages": 3,
|
||||
"pyramid_num_inference_steps_list": [20, 20, 20],
|
||||
"use_cfg_zero_star": true,
|
||||
"use_zero_init": true,
|
||||
"zero_steps": 1
|
||||
}'
|
||||
|
||||
echo "Model extra params: $extra_params"
|
||||
echo ""
|
||||
|
||||
# Create video generation job
|
||||
echo "Creating video generation job..."
|
||||
RESPONSE=$(curl -s -X POST "$API_URL/v1/videos" \
|
||||
-F "prompt=$PROMPT" \
|
||||
-F "model=$MODEL" \
|
||||
-F "width=640" \
|
||||
-F "height=384" \
|
||||
-F "guidance_scale=5.0" \
|
||||
-F "seed=42" \
|
||||
-F "extra_params=$extra_params")
|
||||
|
||||
echo "Response: $RESPONSE"
|
||||
echo ""
|
||||
|
||||
# Extract video ID
|
||||
VIDEO_ID=$(echo "$RESPONSE" | grep -o '"id":"[^"]*"' | cut -d'"' -f4)
|
||||
|
||||
if [ -z "$VIDEO_ID" ]; then
|
||||
echo "Error: Failed to create video job"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "Video job created: $VIDEO_ID"
|
||||
echo ""
|
||||
|
||||
# Poll for completion
|
||||
echo "Polling for completion..."
|
||||
while true; do
|
||||
STATUS_RESPONSE=$(curl -s "$API_URL/v1/videos/$VIDEO_ID")
|
||||
STATUS=$(echo "$STATUS_RESPONSE" | grep -o '"status":"[^"]*"' | cut -d'"' -f4)
|
||||
PROGRESS=$(echo "$STATUS_RESPONSE" | grep -o '"progress":[0-9]*' | cut -d':' -f2)
|
||||
|
||||
echo "Status: $STATUS, Progress: $PROGRESS%"
|
||||
|
||||
if [ "$STATUS" = "completed" ]; then
|
||||
echo ""
|
||||
echo "Video generation completed!"
|
||||
echo "Full response:"
|
||||
echo "$STATUS_RESPONSE" | jq '.'
|
||||
break
|
||||
elif [ "$STATUS" = "failed" ]; then
|
||||
echo ""
|
||||
echo "Video generation failed!"
|
||||
echo "$STATUS_RESPONSE" | jq '.'
|
||||
exit 1
|
||||
fi
|
||||
|
||||
sleep 2
|
||||
done
|
||||
@@ -0,0 +1,69 @@
|
||||
#!/bin/bash
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
# Helios Text-to-Video Generation Example
|
||||
# This script demonstrates how to use the /v1/videos API with Helios models
|
||||
|
||||
API_URL="${API_URL:-http://localhost:8000}"
|
||||
MODEL="${MODEL:-BestWishYsh/Helios-Base}"
|
||||
PROMPT="${PROMPT:-A serene lakeside sunrise with mist over the water.}"
|
||||
|
||||
echo "==================================="
|
||||
echo "Helios T2V Generation"
|
||||
echo "==================================="
|
||||
echo "API URL: $API_URL"
|
||||
echo "Model: $MODEL"
|
||||
echo "Prompt: $PROMPT"
|
||||
echo ""
|
||||
|
||||
# Create video generation job
|
||||
echo "Creating video generation job..."
|
||||
RESPONSE=$(curl -s -X POST "$API_URL/v1/videos" \
|
||||
-F "prompt=$PROMPT" \
|
||||
-F "model=$MODEL" \
|
||||
-F "width=640" \
|
||||
-F "height=384" \
|
||||
-F "num_frames=99" \
|
||||
-F "num_inference_steps=50" \
|
||||
-F "guidance_scale=5.0" \
|
||||
-F "seed=42")
|
||||
|
||||
echo "Response: $RESPONSE"
|
||||
echo ""
|
||||
|
||||
# Extract video ID
|
||||
VIDEO_ID=$(echo "$RESPONSE" | grep -o '"id":"[^"]*"' | cut -d'"' -f4)
|
||||
|
||||
if [ -z "$VIDEO_ID" ]; then
|
||||
echo "Error: Failed to create video job"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "Video job created: $VIDEO_ID"
|
||||
echo ""
|
||||
|
||||
# Poll for completion
|
||||
echo "Polling for completion..."
|
||||
while true; do
|
||||
STATUS_RESPONSE=$(curl -s "$API_URL/v1/videos/$VIDEO_ID")
|
||||
STATUS=$(echo "$STATUS_RESPONSE" | grep -o '"status":"[^"]*"' | cut -d'"' -f4)
|
||||
PROGRESS=$(echo "$STATUS_RESPONSE" | grep -o '"progress":[0-9]*' | cut -d':' -f2)
|
||||
|
||||
echo "Status: $STATUS, Progress: $PROGRESS%"
|
||||
|
||||
if [ "$STATUS" = "completed" ]; then
|
||||
echo ""
|
||||
echo "Video generation completed!"
|
||||
echo "Full response:"
|
||||
echo "$STATUS_RESPONSE" | jq '.'
|
||||
break
|
||||
elif [ "$STATUS" = "failed" ]; then
|
||||
echo ""
|
||||
echo "Video generation failed!"
|
||||
echo "$STATUS_RESPONSE" | jq '.'
|
||||
exit 1
|
||||
fi
|
||||
|
||||
sleep 2
|
||||
done
|
||||
@@ -69,7 +69,48 @@ EOF
|
||||
curl -s http://localhost:8092/v1/chat/completions -H "Content-Type: application/json" -d @request.json | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2 | base64 -d > output.png
|
||||
```
|
||||
|
||||
### Method 2: Using Python Client
|
||||
### Method 2: Using OpenAI Python SDK
|
||||
|
||||
```python
|
||||
import base64
|
||||
from openai import OpenAI
|
||||
|
||||
client = OpenAI(base_url="http://localhost:8092/v1", api_key="none")
|
||||
|
||||
with open("input.png", "rb") as f:
|
||||
img_b64 = base64.b64encode(f.read()).decode()
|
||||
|
||||
response = client.chat.completions.create(
|
||||
model="Qwen/Qwen-Image-Edit",
|
||||
messages=[{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Convert to watercolor style"},
|
||||
{"type": "image_url", "image_url": {
|
||||
"url": f"data:image/png;base64,{img_b64}"
|
||||
}},
|
||||
],
|
||||
}],
|
||||
extra_body={
|
||||
"num_inference_steps": 50,
|
||||
"guidance_scale": 1,
|
||||
"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))
|
||||
```
|
||||
|
||||
!!! note
|
||||
The OpenAI SDK's `extra_body` keyword argument merges parameters into the
|
||||
top-level request body automatically. When using curl or Python `requests`,
|
||||
wrap generation parameters inside a literal `"extra_body"` key in the JSON
|
||||
instead (as shown in the curl example above).
|
||||
|
||||
### Method 3: Using Python Client Script
|
||||
|
||||
```bash
|
||||
python openai_chat_client.py --input input.png --prompt "Convert to oil painting style" --output output.png
|
||||
@@ -78,7 +119,7 @@ python openai_chat_client.py --input input.png --prompt "Convert to oil painting
|
||||
python openai_chat_client.py --input input1.png input2.png --prompt "Combine these images into a single scene" --output output.png
|
||||
```
|
||||
|
||||
### Method 3: Using Gradio Demo
|
||||
### Method 4: Using Gradio Demo
|
||||
|
||||
```bash
|
||||
python gradio_demo.py
|
||||
@@ -144,6 +185,97 @@ Use `extra_body` to pass generation parameters:
|
||||
}
|
||||
```
|
||||
|
||||
### Layered Image Generation (Qwen-Image-Layered)
|
||||
|
||||
Qwen-Image-Layered generates multiple decomposed layers from a reference image and a text prompt.
|
||||
Start the server with:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen-Image-Layered --omni --port 8093
|
||||
```
|
||||
|
||||
**Using curl**
|
||||
|
||||
```bash
|
||||
IMG_B64=$(base64 -w0 input.png)
|
||||
|
||||
curl -sS http://localhost:8093/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d "$(jq -n --arg img "$IMG_B64" '{
|
||||
messages: [{
|
||||
role: "user",
|
||||
content: [
|
||||
{type: "image_url", image_url: {url: ("data:image/png;base64," + $img)}},
|
||||
{type: "text", text: "a rabbit"}
|
||||
]
|
||||
}],
|
||||
extra_body: {
|
||||
num_inference_steps: 50,
|
||||
cfg_scale: 4.0,
|
||||
seed: 0,
|
||||
layers: 4,
|
||||
resolution: 640
|
||||
}
|
||||
}')" \
|
||||
| jq -r '.choices[0].message.content[] | .image_url.url | split(",")[1]' \
|
||||
| while IFS= read -r b64; do
|
||||
((i++)); echo "$b64" | base64 -d > "layer_${i}.png"
|
||||
done
|
||||
```
|
||||
|
||||
**Using Python**
|
||||
|
||||
```python
|
||||
import base64
|
||||
import requests
|
||||
|
||||
with open("input.png", "rb") as f:
|
||||
img_b64 = base64.b64encode(f.read()).decode()
|
||||
|
||||
payload = {
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image_url", "image_url": {
|
||||
"url": f"data:image/png;base64,{img_b64}"
|
||||
}},
|
||||
{"type": "text", "text": "a rabbit"},
|
||||
],
|
||||
}],
|
||||
"extra_body": {
|
||||
"num_inference_steps": 50,
|
||||
"cfg_scale": 4.0,
|
||||
"seed": 0,
|
||||
"layers": 4,
|
||||
"resolution": 640,
|
||||
},
|
||||
}
|
||||
|
||||
resp = requests.post(
|
||||
"http://localhost:8093/v1/chat/completions",
|
||||
json=payload,
|
||||
timeout=600,
|
||||
)
|
||||
data = resp.json()
|
||||
|
||||
for i, item in enumerate(data["choices"][0]["message"]["content"]):
|
||||
_, b64_data = item["image_url"]["url"].split(",", 1)
|
||||
with open(f"layer_{i}.png", "wb") as f:
|
||||
f.write(base64.b64decode(b64_data))
|
||||
```
|
||||
|
||||
The response contains multiple images in `choices[0].message.content` — one per generated layer.
|
||||
|
||||
#### Qwen-Image-Layered Parameters
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| `layers` | int | 4 | Number of layers to decompose |
|
||||
| `resolution` | int | 640 | Resolution for dimension calculation (640 or 1024) |
|
||||
| `cfg_scale` | float | 4.0 | Classifier-free guidance scale (alias for `true_cfg_scale`) |
|
||||
| `num_inference_steps` | int | 50 | Number of denoising steps |
|
||||
| `seed` | int | None | Random seed for reproducibility |
|
||||
|
||||
### Multi-Image Editing (Qwen-Image-Edit-2509)
|
||||
|
||||
Provide multiple images in `content` (order matters):
|
||||
@@ -163,7 +295,14 @@ Provide multiple images in `content` (order matters):
|
||||
}
|
||||
```
|
||||
|
||||
## Generation Parameters (extra_body)
|
||||
## Generation Parameters
|
||||
|
||||
When using `/v1/chat/completions`, pass these inside `extra_body` in the curl
|
||||
JSON, or via the `extra_body` keyword argument in the OpenAI Python SDK.
|
||||
When using the dedicated `/v1/images/edits` endpoint, pass the supported
|
||||
generation controls as top-level form fields directly. For image dimensions and
|
||||
count, use `size` and `n` rather than `height`, `width`, or
|
||||
`num_outputs_per_prompt`.
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
| ------------------------ | ----- | ------- | ------------------------------------- |
|
||||
@@ -171,10 +310,12 @@ Provide multiple images in `content` (order matters):
|
||||
| `width` | int | None | Output image width in pixels |
|
||||
| `size` | str | None | Output image size (e.g., "1024x1024") |
|
||||
| `num_inference_steps` | int | 50 | Number of denoising steps |
|
||||
| `guidance_scale` | float | 7.5 | CFG guidance scale |
|
||||
| `guidance_scale` | float | 1.0 | CFG guidance scale |
|
||||
| `seed` | int | None | Random seed (reproducible) |
|
||||
| `negative_prompt` | str | None | Negative prompt |
|
||||
| `num_outputs_per_prompt` | int | 1 | Number of images to generate |
|
||||
| `layers` | int | 4 | Number of layers (Qwen-Image-Layered) |
|
||||
| `resolution` | int | 640 | Resolution, 640 or 1024 (Qwen-Image-Layered) |
|
||||
|
||||
## Response Format
|
||||
|
||||
|
||||
@@ -4,6 +4,9 @@
|
||||
|
||||
Please refer to [README.md](../../../README.md)
|
||||
|
||||
> ⚠️ **Important (audio generation)**
|
||||
> For **audio generation** (TTS, responses that include synthesized audio, etc.), install **`flash-attn`** for your **CUDA** and **PyTorch** stack. Without it on GPU, **output audio may be noise-only or unusable**. See the [FlashAttention](https://github.com/Dao-AILab/flash-attention) repository for compatible builds.
|
||||
|
||||
## Run examples (MiMo-Audio)
|
||||
|
||||
### Launch the Server
|
||||
|
||||
+97
-68
@@ -16,6 +16,14 @@ class QueryResult(NamedTuple):
|
||||
limit_mm_per_prompt: dict[str, int]
|
||||
|
||||
|
||||
def make_audio_output_filename(request_id: str | None, index: int) -> str:
|
||||
"""Build a stable output filename using request ID when available."""
|
||||
if not request_id:
|
||||
request_id = f"unknown_{index}"
|
||||
safe_request_id = "".join(ch if (ch.isalnum() or ch in ("-", "_")) else "_" for ch in request_id)
|
||||
return f"audio_{safe_request_id}_{index}.wav"
|
||||
|
||||
|
||||
def encode_base64_content_from_url(content_url: str) -> str:
|
||||
"""Encode a content retrieved from a remote url to base64 format."""
|
||||
|
||||
@@ -165,6 +173,34 @@ def get_system_prompt():
|
||||
}
|
||||
|
||||
|
||||
def _parse_csv_arg(value: str | None) -> list[str]:
|
||||
if not value:
|
||||
return []
|
||||
return [item.strip() for item in value.split(",") if item.strip()]
|
||||
|
||||
|
||||
def _build_prompt_for_query_type(
|
||||
query_type: str,
|
||||
custom_prompt: str | None,
|
||||
video_path: str | None,
|
||||
image_path: str | None,
|
||||
audio_path: str | None,
|
||||
):
|
||||
query_func = query_map[query_type]
|
||||
if query_type == "use_video":
|
||||
return query_func(video_path=video_path, custom_prompt=custom_prompt)
|
||||
if query_type == "use_image":
|
||||
return query_func(image_path=image_path, custom_prompt=custom_prompt)
|
||||
if query_type == "use_audio":
|
||||
return query_func(audio_path=audio_path, custom_prompt=custom_prompt)
|
||||
if query_type == "text":
|
||||
return query_func(custom_prompt=custom_prompt)
|
||||
if query_type == "use_audio_in_video":
|
||||
return query_func(video_path=video_path, custom_prompt=custom_prompt)
|
||||
# use_mixed_modalities / use_multi_audios
|
||||
return query_func(custom_prompt=custom_prompt)
|
||||
|
||||
|
||||
def get_text_query(custom_prompt: str | None = None):
|
||||
question = (
|
||||
custom_prompt or "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words."
|
||||
@@ -334,44 +370,6 @@ query_map = {
|
||||
|
||||
def run_multimodal_generation(args, client: OpenAI) -> None:
|
||||
model_name = args.model
|
||||
thinker_sampling_params = {
|
||||
"temperature": 0.4, # Deterministic
|
||||
"top_p": 0.9,
|
||||
"top_k": 1,
|
||||
"max_tokens": 16384,
|
||||
"repetition_penalty": 1.05,
|
||||
"stop_token_ids": [151645], # Qwen EOS token <|im_end|>
|
||||
"seed": SEED,
|
||||
}
|
||||
|
||||
# Sampling parameters for Talker stage (codec generation)
|
||||
# Stop at codec EOS token
|
||||
talker_sampling_params = {
|
||||
"temperature": 0.9,
|
||||
"top_k": 50,
|
||||
"max_tokens": 4096,
|
||||
"seed": SEED,
|
||||
"detokenize": False,
|
||||
"repetition_penalty": 1.05,
|
||||
"stop_token_ids": [2150], # TALKER_CODEC_EOS_TOKEN_ID
|
||||
}
|
||||
|
||||
# # Sampling parameters for Code2Wav stage (audio generation)
|
||||
code2wav_sampling_params = {
|
||||
"temperature": 0.0,
|
||||
"top_p": 1.0,
|
||||
"top_k": -1,
|
||||
"max_tokens": 4096 * 16,
|
||||
"seed": SEED,
|
||||
"detokenize": True,
|
||||
"repetition_penalty": 1.1,
|
||||
}
|
||||
|
||||
sampling_params_list = [
|
||||
thinker_sampling_params,
|
||||
talker_sampling_params,
|
||||
code2wav_sampling_params,
|
||||
]
|
||||
|
||||
# Get paths and custom prompt from args
|
||||
video_path = getattr(args, "video_path", None)
|
||||
@@ -379,31 +377,6 @@ def run_multimodal_generation(args, client: OpenAI) -> None:
|
||||
audio_path = getattr(args, "audio_path", None)
|
||||
custom_prompt = getattr(args, "prompt", None)
|
||||
|
||||
# Get the query function and call it with appropriate parameters
|
||||
query_func = query_map[args.query_type]
|
||||
if args.query_type == "use_video":
|
||||
prompt = query_func(video_path=video_path, custom_prompt=custom_prompt)
|
||||
elif args.query_type == "use_image":
|
||||
prompt = query_func(image_path=image_path, custom_prompt=custom_prompt)
|
||||
elif args.query_type == "use_audio":
|
||||
prompt = query_func(audio_path=audio_path, custom_prompt=custom_prompt)
|
||||
elif args.query_type == "text":
|
||||
prompt = query_func(custom_prompt=custom_prompt)
|
||||
elif args.query_type == "use_audio_in_video":
|
||||
prompt = query_func(
|
||||
video_path=video_path,
|
||||
custom_prompt=custom_prompt,
|
||||
)
|
||||
else:
|
||||
prompt = query_func()
|
||||
|
||||
extra_body = {
|
||||
"sampling_params_list": sampling_params_list # Optional, it has a default setting in stage_configs of the corresponding model.
|
||||
}
|
||||
|
||||
if args.query_type == "use_audio_in_video":
|
||||
extra_body["mm_processor_kwargs"] = {"use_audio_in_video": True}
|
||||
|
||||
if args.modalities is not None:
|
||||
output_modalities = args.modalities.split(",")
|
||||
else:
|
||||
@@ -411,6 +384,36 @@ def run_multimodal_generation(args, client: OpenAI) -> None:
|
||||
|
||||
# Test multiple concurrent completions
|
||||
num_concurrent_requests = args.num_concurrent_requests
|
||||
prompt_list = _parse_csv_arg(getattr(args, "prompts", None))
|
||||
speaker_list = _parse_csv_arg(getattr(args, "speakers", None))
|
||||
|
||||
request_payloads = []
|
||||
for idx in range(num_concurrent_requests):
|
||||
per_req_prompt = (
|
||||
prompt_list[idx]
|
||||
if idx < len(prompt_list)
|
||||
else (custom_prompt if idx == 0 or not prompt_list else prompt_list[-1])
|
||||
)
|
||||
per_req_speaker = (
|
||||
speaker_list[idx]
|
||||
if idx < len(speaker_list)
|
||||
else (args.speaker if idx == 0 or not speaker_list else speaker_list[-1])
|
||||
)
|
||||
prompt = _build_prompt_for_query_type(
|
||||
query_type=args.query_type,
|
||||
custom_prompt=per_req_prompt,
|
||||
video_path=video_path,
|
||||
image_path=image_path,
|
||||
audio_path=audio_path,
|
||||
)
|
||||
extra_body = {
|
||||
# Optional, it has default settings in stage configs. you can override them here.
|
||||
}
|
||||
if args.query_type == "use_audio_in_video":
|
||||
extra_body["mm_processor_kwargs"] = {"use_audio_in_video": True}
|
||||
if per_req_speaker and per_req_speaker.strip():
|
||||
extra_body["speaker"] = per_req_speaker.strip()
|
||||
request_payloads.append({"prompt": prompt, "extra_body": extra_body})
|
||||
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=num_concurrent_requests) as executor:
|
||||
# Submit multiple completion requests concurrently
|
||||
@@ -419,14 +422,14 @@ def run_multimodal_generation(args, client: OpenAI) -> None:
|
||||
client.chat.completions.create,
|
||||
messages=[
|
||||
get_system_prompt(),
|
||||
prompt,
|
||||
payload["prompt"],
|
||||
],
|
||||
model=model_name,
|
||||
modalities=output_modalities,
|
||||
extra_body=extra_body,
|
||||
extra_body=payload["extra_body"],
|
||||
stream=args.stream,
|
||||
)
|
||||
for _ in range(num_concurrent_requests)
|
||||
for payload in request_payloads
|
||||
]
|
||||
|
||||
# Wait for all requests to complete and collect results
|
||||
@@ -437,10 +440,11 @@ def run_multimodal_generation(args, client: OpenAI) -> None:
|
||||
if not args.stream:
|
||||
# Verify all completions succeeded
|
||||
for chat_completion in chat_completions:
|
||||
request_id = getattr(chat_completion, "id", None)
|
||||
for choice in chat_completion.choices:
|
||||
if choice.message.audio:
|
||||
audio_data = base64.b64decode(choice.message.audio.data)
|
||||
audio_file_path = f"audio_{count}.wav"
|
||||
audio_file_path = make_audio_output_filename(request_id=request_id, index=count)
|
||||
with open(audio_file_path, "wb") as f:
|
||||
f.write(audio_data)
|
||||
print(f"Audio saved to {audio_file_path}")
|
||||
@@ -459,7 +463,8 @@ def run_multimodal_generation(args, client: OpenAI) -> None:
|
||||
|
||||
if getattr(chunk, "modality", None) == "audio" and content:
|
||||
audio_data = base64.b64decode(content)
|
||||
audio_file_path = f"audio_{count}.wav"
|
||||
request_id = getattr(chunk, "id", None)
|
||||
audio_file_path = make_audio_output_filename(request_id=request_id, index=count)
|
||||
with open(audio_file_path, "wb") as f:
|
||||
f.write(audio_data)
|
||||
print(f"\nAudio saved to {audio_file_path}")
|
||||
@@ -546,6 +551,30 @@ def parse_args():
|
||||
default="localhost",
|
||||
help="Host/IP of the vLLM Omni API server.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--speaker",
|
||||
type=str,
|
||||
default=None,
|
||||
help="TTS speaker/voice for audio output (e.g. Ethan, Vivian). Passed via extra_body to the talker stage.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--speakers",
|
||||
type=str,
|
||||
default=None,
|
||||
help=(
|
||||
"Comma-separated speakers for concurrent requests, e.g. "
|
||||
"'Ethan,Vivian,Ryan'. Overrides --speaker per request."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompts",
|
||||
type=str,
|
||||
default=None,
|
||||
help=(
|
||||
"Comma-separated prompts for concurrent requests. "
|
||||
"If fewer than --num-concurrent-requests, the last prompt is reused."
|
||||
),
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
@@ -27,7 +27,7 @@ cd examples/online_serving/qwen2_5_omni
|
||||
#### Send request via python
|
||||
|
||||
```bash
|
||||
python openai_chat_completion_client_for_multimodal_generation.py --query-type mixed_modalities --port 8091 --host "localhost"
|
||||
python examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py --model Qwen/Qwen2.5-Omni-7B --query-type use_mixed_modalities --port 8091 --host "localhost"
|
||||
```
|
||||
|
||||
The Python client supports the following command-line arguments:
|
||||
@@ -42,11 +42,12 @@ The Python client supports the following command-line arguments:
|
||||
For example, to use mixed modalities with all local files:
|
||||
|
||||
```bash
|
||||
python openai_chat_completion_client_for_multimodal_generation.py \
|
||||
--query-type mixed_modalities \
|
||||
python examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py \
|
||||
--query-type use_mixed_modalities \
|
||||
--video-path /path/to/your/video.mp4 \
|
||||
--image-path /path/to/your/image.jpg \
|
||||
--audio-path /path/to/your/audio.wav \
|
||||
--model Qwen/Qwen2.5-Omni-7B \
|
||||
--prompt "Analyze all the media content and provide a comprehensive summary."
|
||||
```
|
||||
|
||||
@@ -97,8 +98,9 @@ curl http://localhost:8091/v1/chat/completions \
|
||||
### Using Python client
|
||||
|
||||
```bash
|
||||
python openai_chat_completion_client_for_multimodal_generation.py \
|
||||
--query-type mixed_modalities \
|
||||
python examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py \
|
||||
--query-type use_mixed_modalities \
|
||||
--model Qwen/Qwen2.5-Omni-7B \
|
||||
--modalities text
|
||||
```
|
||||
|
||||
@@ -139,8 +141,9 @@ print(response.choices[1].message.audio) # Audio response
|
||||
## Streaming Output
|
||||
If you want to enable streaming output, please set the argument as below. The final output will be obtained just after generated by corresponding stage. Now we only support text streaming output. Other modalities can output normally.
|
||||
```bash
|
||||
python openai_chat_completion_client_for_multimodal_generation.py \
|
||||
--query-type mixed_modalities \
|
||||
python examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py \
|
||||
--query-type use_mixed_modalities \
|
||||
--model Qwen/Qwen2.5-Omni-7B \
|
||||
--stream
|
||||
```
|
||||
|
||||
|
||||
-459
@@ -1,459 +0,0 @@
|
||||
import base64
|
||||
import os
|
||||
|
||||
import requests
|
||||
from openai import OpenAI
|
||||
from vllm.assets.audio import AudioAsset
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
SEED = 42
|
||||
|
||||
|
||||
def encode_base64_content_from_url(content_url: str) -> str:
|
||||
"""Encode a content retrieved from a remote url to base64 format."""
|
||||
|
||||
with requests.get(content_url) as response:
|
||||
response.raise_for_status()
|
||||
result = base64.b64encode(response.content).decode("utf-8")
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def encode_base64_content_from_file(file_path: str) -> str:
|
||||
"""Encode a local file to base64 format."""
|
||||
with open(file_path, "rb") as f:
|
||||
content = f.read()
|
||||
result = base64.b64encode(content).decode("utf-8")
|
||||
return result
|
||||
|
||||
|
||||
def get_video_url_from_path(video_path: str | None) -> str:
|
||||
"""Convert a video path (local file or URL) to a video URL format for the API.
|
||||
|
||||
If video_path is None or empty, returns the default URL.
|
||||
If video_path is a local file path, encodes it to base64 data URL.
|
||||
If video_path is a URL, returns it as-is.
|
||||
"""
|
||||
if not video_path:
|
||||
# Default video URL
|
||||
return "https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4"
|
||||
|
||||
# Check if it's a URL (starts with http:// or https://)
|
||||
if video_path.startswith(("http://", "https://")):
|
||||
return video_path
|
||||
|
||||
# Otherwise, treat it as a local file path
|
||||
if not os.path.exists(video_path):
|
||||
raise FileNotFoundError(f"Video file not found: {video_path}")
|
||||
|
||||
# Detect video MIME type from file extension
|
||||
video_path_lower = video_path.lower()
|
||||
if video_path_lower.endswith(".mp4"):
|
||||
mime_type = "video/mp4"
|
||||
elif video_path_lower.endswith(".webm"):
|
||||
mime_type = "video/webm"
|
||||
elif video_path_lower.endswith(".mov"):
|
||||
mime_type = "video/quicktime"
|
||||
elif video_path_lower.endswith(".avi"):
|
||||
mime_type = "video/x-msvideo"
|
||||
elif video_path_lower.endswith(".mkv"):
|
||||
mime_type = "video/x-matroska"
|
||||
else:
|
||||
# Default to mp4 if extension is unknown
|
||||
mime_type = "video/mp4"
|
||||
|
||||
video_base64 = encode_base64_content_from_file(video_path)
|
||||
return f"data:{mime_type};base64,{video_base64}"
|
||||
|
||||
|
||||
def get_image_url_from_path(image_path: str | None) -> str:
|
||||
"""Convert an image path (local file or URL) to an image URL format for the API.
|
||||
|
||||
If image_path is None or empty, returns the default URL.
|
||||
If image_path is a local file path, encodes it to base64 data URL.
|
||||
If image_path is a URL, returns it as-is.
|
||||
"""
|
||||
if not image_path:
|
||||
# Default image URL
|
||||
return "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg"
|
||||
|
||||
# Check if it's a URL (starts with http:// or https://)
|
||||
if image_path.startswith(("http://", "https://")):
|
||||
return image_path
|
||||
|
||||
# Otherwise, treat it as a local file path
|
||||
if not os.path.exists(image_path):
|
||||
raise FileNotFoundError(f"Image file not found: {image_path}")
|
||||
|
||||
# Detect image MIME type from file extension
|
||||
image_path_lower = image_path.lower()
|
||||
if image_path_lower.endswith((".jpg", ".jpeg")):
|
||||
mime_type = "image/jpeg"
|
||||
elif image_path_lower.endswith(".png"):
|
||||
mime_type = "image/png"
|
||||
elif image_path_lower.endswith(".gif"):
|
||||
mime_type = "image/gif"
|
||||
elif image_path_lower.endswith(".webp"):
|
||||
mime_type = "image/webp"
|
||||
else:
|
||||
# Default to jpeg if extension is unknown
|
||||
mime_type = "image/jpeg"
|
||||
|
||||
image_base64 = encode_base64_content_from_file(image_path)
|
||||
return f"data:{mime_type};base64,{image_base64}"
|
||||
|
||||
|
||||
def get_audio_url_from_path(audio_path: str | None) -> str:
|
||||
"""Convert an audio path (local file or URL) to an audio URL format for the API.
|
||||
|
||||
If audio_path is None or empty, returns the default URL.
|
||||
If audio_path is a local file path, encodes it to base64 data URL.
|
||||
If audio_path is a URL, returns it as-is.
|
||||
"""
|
||||
if not audio_path:
|
||||
# Default audio URL
|
||||
return AudioAsset("mary_had_lamb").url
|
||||
|
||||
# Check if it's a URL (starts with http:// or https://)
|
||||
if audio_path.startswith(("http://", "https://")):
|
||||
return audio_path
|
||||
|
||||
# Otherwise, treat it as a local file path
|
||||
if not os.path.exists(audio_path):
|
||||
raise FileNotFoundError(f"Audio file not found: {audio_path}")
|
||||
|
||||
# Detect audio MIME type from file extension
|
||||
audio_path_lower = audio_path.lower()
|
||||
if audio_path_lower.endswith((".mp3", ".mpeg")):
|
||||
mime_type = "audio/mpeg"
|
||||
elif audio_path_lower.endswith(".wav"):
|
||||
mime_type = "audio/wav"
|
||||
elif audio_path_lower.endswith(".ogg"):
|
||||
mime_type = "audio/ogg"
|
||||
elif audio_path_lower.endswith(".flac"):
|
||||
mime_type = "audio/flac"
|
||||
elif audio_path_lower.endswith(".m4a"):
|
||||
mime_type = "audio/mp4"
|
||||
else:
|
||||
# Default to wav if extension is unknown
|
||||
mime_type = "audio/wav"
|
||||
|
||||
audio_base64 = encode_base64_content_from_file(audio_path)
|
||||
return f"data:{mime_type};base64,{audio_base64}"
|
||||
|
||||
|
||||
def get_system_prompt():
|
||||
return {
|
||||
"role": "system",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": (
|
||||
"You are Qwen, a virtual human developed by the Qwen Team, "
|
||||
"Alibaba Group, capable of perceiving auditory and visual inputs, "
|
||||
"as well as generating text and speech."
|
||||
),
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def get_text_query(custom_prompt: str | None = None):
|
||||
question = (
|
||||
custom_prompt or "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words."
|
||||
)
|
||||
prompt = {
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": f"{question}",
|
||||
}
|
||||
],
|
||||
}
|
||||
return prompt
|
||||
|
||||
|
||||
def get_mixed_modalities_query(
|
||||
video_path: str | None = None,
|
||||
image_path: str | None = None,
|
||||
audio_path: str | None = None,
|
||||
custom_prompt: str | None = None,
|
||||
):
|
||||
question = (
|
||||
custom_prompt or "What is recited in the audio? What is the content of this image? Why is this video funny?"
|
||||
)
|
||||
video_url = get_video_url_from_path(video_path)
|
||||
image_url = get_image_url_from_path(image_path)
|
||||
audio_url = get_audio_url_from_path(audio_path)
|
||||
prompt = {
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "audio_url",
|
||||
"audio_url": {"url": audio_url},
|
||||
},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": image_url},
|
||||
},
|
||||
{
|
||||
"type": "video_url",
|
||||
"video_url": {"url": video_url},
|
||||
},
|
||||
{
|
||||
"type": "text",
|
||||
"text": f"{question}",
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
return prompt
|
||||
|
||||
|
||||
def get_use_audio_in_video_query(video_path: str | None = None, custom_prompt: str | None = None):
|
||||
question = custom_prompt or "Describe the content of the video, then convert what the baby say into text."
|
||||
video_url = get_video_url_from_path(video_path)
|
||||
|
||||
prompt = {
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "video_url",
|
||||
"video_url": {
|
||||
"url": video_url,
|
||||
"num_frames": 16,
|
||||
},
|
||||
},
|
||||
{
|
||||
"type": "text",
|
||||
"text": f"{question}",
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
return prompt
|
||||
|
||||
|
||||
def get_multi_audios_query(audio_path: str | None = None, custom_prompt: str | None = None):
|
||||
question = custom_prompt or "Are these two audio clips the same?"
|
||||
audio_url = get_audio_url_from_path(audio_path)
|
||||
prompt = {
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "audio_url",
|
||||
"audio_url": {"url": audio_url},
|
||||
},
|
||||
{
|
||||
"type": "audio_url",
|
||||
"audio_url": {"url": AudioAsset("winning_call").url},
|
||||
},
|
||||
{
|
||||
"type": "text",
|
||||
"text": f"{question}",
|
||||
},
|
||||
],
|
||||
}
|
||||
return prompt
|
||||
|
||||
|
||||
query_map = {
|
||||
"mixed_modalities": get_mixed_modalities_query,
|
||||
"use_audio_in_video": get_use_audio_in_video_query,
|
||||
"multi_audios": get_multi_audios_query,
|
||||
"text": get_text_query,
|
||||
}
|
||||
|
||||
|
||||
def run_multimodal_generation(args, client: OpenAI) -> None:
|
||||
model_name = "Qwen/Qwen2.5-Omni-7B"
|
||||
thinker_sampling_params = {
|
||||
"temperature": 0.0, # Deterministic - no randomness
|
||||
"top_p": 1.0, # Disable nucleus sampling
|
||||
"top_k": -1, # Disable top-k sampling
|
||||
"max_tokens": 2048,
|
||||
"seed": SEED, # Fixed seed for sampling
|
||||
"detokenize": True,
|
||||
"repetition_penalty": 1.1,
|
||||
}
|
||||
talker_sampling_params = {
|
||||
"temperature": 0.9,
|
||||
"top_p": 0.8,
|
||||
"top_k": 40,
|
||||
"max_tokens": 2048,
|
||||
"seed": SEED, # Fixed seed for sampling
|
||||
"detokenize": True,
|
||||
"repetition_penalty": 1.05,
|
||||
"stop_token_ids": [8294],
|
||||
}
|
||||
code2wav_sampling_params = {
|
||||
"temperature": 0.0, # Deterministic - no randomness
|
||||
"top_p": 1.0, # Disable nucleus sampling
|
||||
"top_k": -1, # Disable top-k sampling
|
||||
"max_tokens": 2048,
|
||||
"seed": SEED, # Fixed seed for sampling
|
||||
"detokenize": True,
|
||||
"repetition_penalty": 1.1,
|
||||
}
|
||||
|
||||
sampling_params_list = [
|
||||
thinker_sampling_params,
|
||||
talker_sampling_params,
|
||||
code2wav_sampling_params,
|
||||
]
|
||||
|
||||
# Get paths and custom prompt from args
|
||||
video_path = getattr(args, "video_path", None)
|
||||
image_path = getattr(args, "image_path", None)
|
||||
audio_path = getattr(args, "audio_path", None)
|
||||
custom_prompt = getattr(args, "prompt", None)
|
||||
|
||||
# Get the query function and call it with appropriate parameters
|
||||
query_func = query_map[args.query_type]
|
||||
if args.query_type == "mixed_modalities":
|
||||
prompt = query_func(
|
||||
video_path=video_path, image_path=image_path, audio_path=audio_path, custom_prompt=custom_prompt
|
||||
)
|
||||
elif args.query_type == "use_audio_in_video":
|
||||
prompt = query_func(video_path=video_path, custom_prompt=custom_prompt)
|
||||
elif args.query_type == "multi_audios":
|
||||
prompt = query_func(audio_path=audio_path, custom_prompt=custom_prompt)
|
||||
elif args.query_type == "text":
|
||||
prompt = query_func(custom_prompt=custom_prompt)
|
||||
else:
|
||||
prompt = query_func()
|
||||
|
||||
extra_body = {
|
||||
"sampling_params_list": sampling_params_list # Optional, it has a default setting in stage_configs of the corresponding model.
|
||||
}
|
||||
|
||||
if args.query_type == "use_audio_in_video":
|
||||
extra_body["mm_processor_kwargs"] = {"use_audio_in_video": True}
|
||||
|
||||
if args.modalities is not None:
|
||||
output_modalities = args.modalities.split(",")
|
||||
else:
|
||||
output_modalities = None
|
||||
|
||||
chat_completion = client.chat.completions.create(
|
||||
messages=[
|
||||
get_system_prompt(),
|
||||
prompt,
|
||||
],
|
||||
model=model_name,
|
||||
modalities=output_modalities,
|
||||
extra_body=extra_body,
|
||||
stream=args.stream,
|
||||
)
|
||||
|
||||
count = 0
|
||||
if not args.stream:
|
||||
for choice in chat_completion.choices:
|
||||
if choice.message.audio:
|
||||
audio_data = base64.b64decode(choice.message.audio.data)
|
||||
audio_file_path = f"audio_{count}.wav"
|
||||
with open(audio_file_path, "wb") as f:
|
||||
f.write(audio_data)
|
||||
print(f"Audio saved to {audio_file_path}")
|
||||
count += 1
|
||||
elif choice.message.content:
|
||||
print("Chat completion output from text:", choice.message.content)
|
||||
else:
|
||||
printed_content = False
|
||||
for chunk in chat_completion:
|
||||
for choice in chunk.choices:
|
||||
if hasattr(choice, "delta"):
|
||||
content = getattr(choice.delta, "content", None)
|
||||
else:
|
||||
content = None
|
||||
|
||||
if getattr(chunk, "modality", None) == "audio" and content:
|
||||
audio_data = base64.b64decode(content)
|
||||
audio_file_path = f"audio_{count}.wav"
|
||||
with open(audio_file_path, "wb") as f:
|
||||
f.write(audio_data)
|
||||
print(f"\nAudio saved to {audio_file_path}")
|
||||
count += 1
|
||||
|
||||
elif getattr(chunk, "modality", None) == "text":
|
||||
if not printed_content:
|
||||
printed_content = True
|
||||
print("\ncontent:", end="", flush=True)
|
||||
print(content, end="", flush=True)
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
|
||||
parser.add_argument(
|
||||
"--query-type",
|
||||
"-q",
|
||||
type=str,
|
||||
default="mixed_modalities",
|
||||
choices=query_map.keys(),
|
||||
help="Query type.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--video-path",
|
||||
"-v",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to local video file or URL. If not provided and query-type uses video, uses default video URL.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--image-path",
|
||||
"-i",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to local image file or URL. If not provided and query-type uses image, uses default image URL.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--audio-path",
|
||||
"-a",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to local audio file or URL. If not provided and query-type uses audio, uses default audio URL.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompt",
|
||||
"-p",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Custom text prompt/question to use instead of the default prompt for the selected query type.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--modalities",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Output modalities to use for the prompts.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stream",
|
||||
action="store_true",
|
||||
help="Stream the response.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--port",
|
||||
type=int,
|
||||
default=8091,
|
||||
help="Port of the vLLM Omni API server.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--host",
|
||||
type=str,
|
||||
default="localhost",
|
||||
help="Host/IP of the vLLM Omni API server.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
host = args.host
|
||||
port = args.port
|
||||
openai_api_base = f"http://{host}:{port}/v1"
|
||||
client = OpenAI(
|
||||
api_key="EMPTY",
|
||||
base_url=openai_api_base,
|
||||
)
|
||||
run_multimodal_generation(args, client)
|
||||
@@ -15,7 +15,7 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
|
||||
If you want to open async chunking for qwen3-omni, launch the server with command below
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --stage-configs-path /vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml
|
||||
```
|
||||
|
||||
If you have custom stage configs file, launch the server with command below
|
||||
@@ -33,7 +33,7 @@ cd examples/online_serving/qwen3_omni
|
||||
#### Send request via python
|
||||
|
||||
```bash
|
||||
python openai_chat_completion_client_for_multimodal_generation.py --query-type use_image --port 8091 --host "localhost"
|
||||
python examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py --model Qwen/Qwen3-Omni-30B-A3B-Instruct --query-type use_image --port 8091 --host "localhost"
|
||||
```
|
||||
|
||||
The Python client supports the following command-line arguments:
|
||||
@@ -44,14 +44,16 @@ The Python client supports the following command-line arguments:
|
||||
- `--image-path` (or `-i`): Path to local image file or URL. If not provided and query-type is `use_image`, uses default image URL. Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs and common image formats: JPEG, PNG, GIF, WebP. Example: `--image-path /path/to/image.jpg` or `--image-path https://example.com/image.png`
|
||||
- `--audio-path` (or `-a`): Path to local audio file or URL. If not provided and query-type is `use_audio`, uses default audio URL. Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs and common audio formats: MP3, WAV, OGG, FLAC, M4A. Example: `--audio-path /path/to/audio.wav` or `--audio-path https://example.com/audio.mp3`
|
||||
- `--prompt` (or `-p`): Custom text prompt/question. If not provided, uses default prompt for the selected query type. Example: `--prompt "What are the main activities shown in this video?"`
|
||||
- `--speaker`: TTS speaker/voice for audio output when requesting audio (e.g. `ethan`, `chelsie`, `aiden`). Omit to use the model default. Example: `--speaker "chelsie"`
|
||||
|
||||
|
||||
For example, to use a local video file with custom prompt:
|
||||
|
||||
```bash
|
||||
python openai_chat_completion_client_for_multimodal_generation.py \
|
||||
python examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py \
|
||||
--query-type use_video \
|
||||
--video-path /path/to/your/video.mp4 \
|
||||
--model Qwen/Qwen3-Omni-30B-A3B-Instruct \
|
||||
--prompt "What are the main activities shown in this video?"
|
||||
```
|
||||
|
||||
@@ -78,7 +80,7 @@ You can control output modalities to specify which types of output the model sho
|
||||
| Modalities | Output |
|
||||
|------------|--------|
|
||||
| `["text"]` | Text only |
|
||||
| `["audio"]` | Text + Audio |
|
||||
| `["audio"]` | Audio only |
|
||||
| `["text", "audio"]` | Text + Audio |
|
||||
| Not specified | Text + Audio (default) |
|
||||
|
||||
@@ -99,20 +101,21 @@ curl http://localhost:8091/v1/chat/completions \
|
||||
#### Text + Audio
|
||||
|
||||
```bash
|
||||
curl http://localhost:8091/v1/chat/completions \
|
||||
curl -s http://localhost:8091/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
|
||||
"messages": [{"role": "user", "content": "Describe vLLM in brief."}],
|
||||
"modalities": ["audio"]
|
||||
}'
|
||||
}' | jq -r '.choices[0].message.audio.data' | base64 -d > output.wav
|
||||
```
|
||||
|
||||
### Using Python client
|
||||
|
||||
```bash
|
||||
python openai_chat_completion_client_for_multimodal_generation.py \
|
||||
python examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py \
|
||||
--query-type use_image \
|
||||
--model Qwen/Qwen3-Omni-30B-A3B-Instruct \
|
||||
--modalities text
|
||||
```
|
||||
|
||||
@@ -136,6 +139,7 @@ print(response.choices[0].message.content)
|
||||
#### Text + Audio
|
||||
|
||||
```python
|
||||
import base64
|
||||
from openai import OpenAI
|
||||
|
||||
client = OpenAI(base_url="http://localhost:8091/v1", api_key="EMPTY")
|
||||
@@ -143,18 +147,75 @@ client = OpenAI(base_url="http://localhost:8091/v1", api_key="EMPTY")
|
||||
response = client.chat.completions.create(
|
||||
model="Qwen/Qwen3-Omni-30B-A3B-Instruct",
|
||||
messages=[{"role": "user", "content": "Describe vLLM in brief."}],
|
||||
modalities=["audio"]
|
||||
modalities=["text", "audio"]
|
||||
)
|
||||
# Response contains two choices: one with text, one with audio
|
||||
print(response.choices[0].message.content) # Text response
|
||||
print(response.choices[1].message.audio) # Audio response
|
||||
|
||||
# Save audio to file
|
||||
audio_data = base64.b64decode(response.choices[1].message.audio.data)
|
||||
with open("output.wav", "wb") as f:
|
||||
f.write(audio_data)
|
||||
```
|
||||
|
||||
## Speaker selection
|
||||
|
||||
When requesting audio output, you can choose the TTS speaker (voice) used for synthesis. If not specified, the model uses its default speaker.
|
||||
|
||||
### Using curl
|
||||
|
||||
Pass a `speaker` field in the request body:
|
||||
|
||||
```bash
|
||||
curl http://localhost:8091/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
|
||||
"messages": [{"role": "user", "content": "Say hello in one sentence."}],
|
||||
"modalities": ["audio"],
|
||||
"speaker": "chelsie"
|
||||
}'
|
||||
```
|
||||
|
||||
### Using Python client
|
||||
|
||||
Use the `--speaker` argument when generating audio:
|
||||
|
||||
```bash
|
||||
python examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py \
|
||||
--query-type use_image \
|
||||
--modalities audio \
|
||||
--model Qwen/Qwen3-Omni-30B-A3B-Instruct \
|
||||
--speaker "chelsie"
|
||||
```
|
||||
|
||||
### Using OpenAI Python SDK
|
||||
|
||||
Pass `speaker` in `extra_body`:
|
||||
|
||||
```python
|
||||
from openai import OpenAI
|
||||
|
||||
client = OpenAI(base_url="http://localhost:8091/v1", api_key="EMPTY")
|
||||
|
||||
response = client.chat.completions.create(
|
||||
model="Qwen/Qwen3-Omni-30B-A3B-Instruct",
|
||||
messages=[{"role": "user", "content": "Say hello in one sentence."}],
|
||||
modalities=["audio"],
|
||||
extra_body={"speaker": "chelsie"}
|
||||
)
|
||||
# Audio uses the specified speaker
|
||||
print(response.choices[1].message.audio)
|
||||
```
|
||||
|
||||
Supported speaker names depend on the model (e.g. `Ethan`, `Chelsie`, `Aiden`). Omit `speaker` to use the default.
|
||||
|
||||
## Streaming Output
|
||||
If you want to enable streaming output, please set the argument as below. The final output will be obtained just after generated by corresponding stage. Now we only support text streaming output. Other modalities can output normally.
|
||||
```bash
|
||||
python openai_chat_completion_client_for_multimodal_generation.py \
|
||||
python examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py \
|
||||
--query-type use_image \
|
||||
--model Qwen/Qwen3-Omni-30B-A3B-Instruct \
|
||||
--stream
|
||||
```
|
||||
|
||||
|
||||
@@ -18,10 +18,10 @@ Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/
|
||||
|
||||
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`) |
|
||||
| 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
|
||||
@@ -43,30 +43,23 @@ Then open http://localhost:7860 in your browser.
|
||||
|
||||
### Launch the Server
|
||||
|
||||
The default stage config is located at `vllm_omni/model_executor/stage_configs/qwen3_tts.yaml`. For other platforms (e.g., NPU), refer to `vllm_omni/platforms/npu/stage_configs/qwen3_tts.yaml`.
|
||||
|
||||
```bash
|
||||
# CustomVoice model (predefined speakers)
|
||||
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts.yaml \
|
||||
--omni \
|
||||
--port 8091 \
|
||||
--trust-remote-code \
|
||||
--enforce-eager
|
||||
--port 8091
|
||||
|
||||
# VoiceDesign model
|
||||
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts.yaml \
|
||||
--omni \
|
||||
--port 8091 \
|
||||
--trust-remote-code \
|
||||
--enforce-eager
|
||||
--port 8091
|
||||
|
||||
# Base model (voice cloning)
|
||||
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts.yaml \
|
||||
--omni \
|
||||
--port 8091 \
|
||||
--trust-remote-code \
|
||||
--enforce-eager
|
||||
--port 8091
|
||||
```
|
||||
|
||||
If you have custom stage configs file, launch the server with command below
|
||||
@@ -74,9 +67,7 @@ If you have custom stage configs file, launch the server with command below
|
||||
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
|
||||
--port 8091
|
||||
```
|
||||
|
||||
Alternatively, use the convenience script:
|
||||
@@ -99,12 +90,14 @@ cd examples/online_serving/qwen3_tts
|
||||
```bash
|
||||
# CustomVoice: Use predefined speaker
|
||||
python openai_speech_client.py \
|
||||
--model Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
|
||||
--text "你好,我是通义千问" \
|
||||
--voice vivian \
|
||||
--language Chinese
|
||||
|
||||
# CustomVoice with style instruction
|
||||
python openai_speech_client.py \
|
||||
--model Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
|
||||
--text "今天天气真好" \
|
||||
--voice ryan \
|
||||
--instructions "用开心的语气说"
|
||||
@@ -134,8 +127,8 @@ The Python client supports the following command-line arguments:
|
||||
- `--voice`: Speaker/voice 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)
|
||||
- `--ref-text`: Reference audio transcript for voice cloning (Base task)
|
||||
- `--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`)
|
||||
|
||||
@@ -239,6 +232,7 @@ Upload a new voice sample for voice cloning in Base task TTS requests.
|
||||
- `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).
|
||||
|
||||
**Response Example:**
|
||||
```json
|
||||
@@ -259,7 +253,8 @@ Upload a new voice sample for voice cloning in Base task TTS requests.
|
||||
curl -X POST http://localhost:8000/v1/audio/voices \
|
||||
-F "audio_sample=@/path/to/voice_sample.wav" \
|
||||
-F "consent=user_consent_id" \
|
||||
-F "name=custom_voice_1"
|
||||
-F "name=custom_voice_1" \
|
||||
-F "ref_text=The exact transcript of the audio sample."
|
||||
```
|
||||
|
||||
### Endpoint
|
||||
@@ -281,7 +276,7 @@ This endpoint follows the [OpenAI Audio Speech API](https://platform.openai.com/
|
||||
"task_type": "CustomVoice",
|
||||
"language": "Auto",
|
||||
"instructions": "Optional style instructions",
|
||||
"ref_audio": "URL or base64 for voice cloning",
|
||||
"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
|
||||
@@ -308,24 +303,24 @@ Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/w
|
||||
|
||||
### 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"`) |
|
||||
| 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 (URL or base64 data URL) |
|
||||
| `ref_text` | string | No | Transcript of reference audio (for ICL mode) |
|
||||
| `x_vector_only_mode` | bool | No | Use speaker embedding only (no ICL) |
|
||||
| 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
|
||||
|
||||
|
||||
@@ -0,0 +1,161 @@
|
||||
"""Batch speech client for Qwen3-TTS via /v1/audio/speech/batch endpoint.
|
||||
|
||||
This script demonstrates how to synthesize multiple texts in a single request.
|
||||
A particularly useful scenario is voice cloning: set ref_audio once at the
|
||||
batch level and generate many utterances in the cloned voice without repeating
|
||||
the reference for each item.
|
||||
|
||||
Start the server (with batch-optimized config for best throughput):
|
||||
|
||||
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts_batch.yaml \
|
||||
--trust-remote-code
|
||||
|
||||
Examples:
|
||||
# Batch with a predefined voice
|
||||
python batch_speech_client.py \
|
||||
--texts "Hello, how are you?" "Goodbye, see you later!"
|
||||
|
||||
# Voice cloning: one ref_audio, many outputs
|
||||
python batch_speech_client.py \
|
||||
--task-type Base \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--ref-text "Transcript of the reference audio" \
|
||||
--texts "First cloned sentence." "Second cloned sentence." \
|
||||
"Third cloned sentence."
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import os
|
||||
|
||||
import httpx
|
||||
|
||||
DEFAULT_API_BASE = "http://localhost:8091"
|
||||
DEFAULT_API_KEY = "EMPTY"
|
||||
|
||||
|
||||
def encode_audio_to_base64(audio_path: str) -> str:
|
||||
"""Encode a local audio file to a base64 data URL."""
|
||||
if not os.path.exists(audio_path):
|
||||
raise FileNotFoundError(f"Audio file not found: {audio_path}")
|
||||
|
||||
ext = os.path.splitext(audio_path)[1].lower()
|
||||
mime_map = {".wav": "audio/wav", ".mp3": "audio/mpeg", ".flac": "audio/flac", ".ogg": "audio/ogg"}
|
||||
mime_type = mime_map.get(ext, "audio/wav")
|
||||
|
||||
with open(audio_path, "rb") as f:
|
||||
audio_b64 = base64.b64encode(f.read()).decode("utf-8")
|
||||
return f"data:{mime_type};base64,{audio_b64}"
|
||||
|
||||
|
||||
def run_batch(args) -> None:
|
||||
"""Send a batch TTS request and save each result to a file."""
|
||||
items = [{"input": text} for text in args.texts]
|
||||
|
||||
payload: dict = {
|
||||
"items": items,
|
||||
"response_format": args.response_format,
|
||||
}
|
||||
if args.voice:
|
||||
payload["voice"] = args.voice
|
||||
if args.language:
|
||||
payload["language"] = args.language
|
||||
if args.task_type:
|
||||
payload["task_type"] = args.task_type
|
||||
if args.instructions:
|
||||
payload["instructions"] = args.instructions
|
||||
if args.max_new_tokens:
|
||||
payload["max_new_tokens"] = args.max_new_tokens
|
||||
|
||||
# Voice cloning parameters (shared across all items)
|
||||
if args.ref_audio:
|
||||
if args.ref_audio.startswith(("http://", "https://")):
|
||||
payload["ref_audio"] = args.ref_audio
|
||||
else:
|
||||
payload["ref_audio"] = encode_audio_to_base64(args.ref_audio)
|
||||
if args.ref_text:
|
||||
payload["ref_text"] = args.ref_text
|
||||
|
||||
print(f"Sending batch of {len(items)} item(s) to {args.api_base}")
|
||||
if args.ref_audio:
|
||||
print("Voice cloning mode — ref_audio applied to all items")
|
||||
|
||||
url = f"{args.api_base}/v1/audio/speech/batch"
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {args.api_key}",
|
||||
}
|
||||
|
||||
with httpx.Client(timeout=300.0) as client:
|
||||
response = client.post(url, json=payload, headers=headers)
|
||||
|
||||
if response.status_code != 200:
|
||||
print(f"Error {response.status_code}: {response.text}")
|
||||
return
|
||||
|
||||
data = response.json()
|
||||
print(f"Total: {data['total']} Succeeded: {data['succeeded']} Failed: {data['failed']}")
|
||||
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
for result in data["results"]:
|
||||
idx = result["index"]
|
||||
if result["status"] == "success":
|
||||
audio_bytes = base64.b64decode(result["audio_data"])
|
||||
out_path = os.path.join(args.output_dir, f"batch_{idx}.{args.response_format}")
|
||||
with open(out_path, "wb") as f:
|
||||
f.write(audio_bytes)
|
||||
print(f" [{idx}] saved {len(audio_bytes)} bytes -> {out_path}")
|
||||
else:
|
||||
print(f" [{idx}] FAILED: {result['error']}")
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Batch speech client for /v1/audio/speech/batch",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog=__doc__,
|
||||
)
|
||||
|
||||
parser.add_argument("--api-base", default=DEFAULT_API_BASE, help="API base URL")
|
||||
parser.add_argument("--api-key", default=DEFAULT_API_KEY, help="API key")
|
||||
|
||||
# Texts to synthesize
|
||||
parser.add_argument(
|
||||
"--texts",
|
||||
nargs="+",
|
||||
required=True,
|
||||
help="One or more texts to synthesize",
|
||||
)
|
||||
|
||||
# Shared voice settings
|
||||
parser.add_argument("--voice", default="vivian", help="Speaker name (default: vivian)")
|
||||
parser.add_argument("--language", default=None, help="Language: Auto, Chinese, English, etc.")
|
||||
parser.add_argument("--instructions", default=None, help="Voice style/emotion instructions")
|
||||
parser.add_argument(
|
||||
"--task-type",
|
||||
default=None,
|
||||
choices=["CustomVoice", "VoiceDesign", "Base"],
|
||||
help="TTS task type (default: CustomVoice)",
|
||||
)
|
||||
|
||||
# Voice cloning (Base task)
|
||||
parser.add_argument("--ref-audio", default=None, help="Reference audio path or URL for voice cloning")
|
||||
parser.add_argument("--ref-text", default=None, help="Reference audio transcript for voice cloning")
|
||||
|
||||
# Generation
|
||||
parser.add_argument("--max-new-tokens", type=int, default=None, help="Max new tokens per item")
|
||||
parser.add_argument(
|
||||
"--response-format",
|
||||
default="wav",
|
||||
choices=["wav", "mp3", "flac", "pcm", "aac", "opus"],
|
||||
help="Audio format (default: wav)",
|
||||
)
|
||||
parser.add_argument("--output-dir", "-o", default="batch_output", help="Output directory (default: batch_output)")
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
run_batch(args)
|
||||
@@ -1,9 +1,13 @@
|
||||
"""Gradio demo for Qwen3-TTS online serving via /v1/audio/speech API.
|
||||
"""Gradio demo for Qwen3-TTS with gapless streaming audio playback.
|
||||
|
||||
Uses Gradio's built-in audio component for playback. Supports streaming
|
||||
(progressive PCM chunks) and non-streaming (full audio download) modes.
|
||||
Uses a custom AudioWorklet-based player for gap-free streaming,
|
||||
inspired by github.com/KoljaB/RealtimeVoiceChat. Audio is streamed
|
||||
from the vLLM server through a same-origin proxy and played via the
|
||||
Web Audio API's AudioWorklet, which maintains a FIFO buffer queue
|
||||
and plays samples at the audio clock rate — eliminating inter-chunk
|
||||
gaps inherent in Gradio's built-in streaming audio component.
|
||||
|
||||
For gapless streaming playback, see gradio_fastrtc_demo.py which uses WebRTC.
|
||||
Also supports non-streaming mode (full audio download) via gr.Audio.
|
||||
|
||||
Supports all 3 task types:
|
||||
- CustomVoice: Predefined speaker with optional style instructions
|
||||
@@ -11,19 +15,21 @@ Supports all 3 task types:
|
||||
- Base: Voice cloning from reference audio (upload or URL)
|
||||
|
||||
Usage:
|
||||
# Start the server first (see run_server.sh), then:
|
||||
# Start the vLLM server first (see run_server.sh), then:
|
||||
python gradio_demo.py --api-base http://localhost:8000
|
||||
|
||||
# Or use run_gradio_demo.sh to start both server and demo together.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import io
|
||||
import json
|
||||
import logging
|
||||
|
||||
import gradio as gr
|
||||
import httpx
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
from fastapi import FastAPI, Request
|
||||
from fastapi.responses import Response, StreamingResponse
|
||||
from tts_common import (
|
||||
PCM_SAMPLE_RATE,
|
||||
SUPPORTED_LANGUAGES,
|
||||
@@ -31,10 +37,253 @@ from tts_common import (
|
||||
add_common_args,
|
||||
build_payload,
|
||||
fetch_voices,
|
||||
on_task_type_change,
|
||||
stream_pcm_chunks,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# ── AudioWorklet processor (loaded in browser via Blob URL) ──────────
|
||||
WORKLET_JS = r"""
|
||||
class TTSPlaybackProcessor extends AudioWorkletProcessor {
|
||||
constructor() {
|
||||
super();
|
||||
this.queue = [];
|
||||
this.buf = null;
|
||||
this.pos = 0;
|
||||
this.playing = false;
|
||||
this.played = 0;
|
||||
this.port.onmessage = (e) => {
|
||||
if (e.data && e.data.type === 'clear') {
|
||||
this.queue = []; this.buf = null; this.pos = 0; this.played = 0;
|
||||
if (this.playing) { this.playing = false; this.port.postMessage({type:'stopped'}); }
|
||||
return;
|
||||
}
|
||||
this.queue.push(e.data);
|
||||
};
|
||||
}
|
||||
process(inputs, outputs) {
|
||||
const out = outputs[0][0];
|
||||
for (let i = 0; i < out.length; i++) {
|
||||
if (!this.buf || this.pos >= this.buf.length) {
|
||||
if (this.queue.length > 0) {
|
||||
this.buf = this.queue.shift(); this.pos = 0;
|
||||
} else {
|
||||
for (let j = i; j < out.length; j++) out[j] = 0;
|
||||
if (this.playing) { this.playing = false; this.port.postMessage({type:'stopped', played:this.played}); }
|
||||
return true;
|
||||
}
|
||||
}
|
||||
out[i] = this.buf[this.pos++] / 32768;
|
||||
this.played++;
|
||||
}
|
||||
if (!this.playing) { this.playing = true; this.port.postMessage({type:'started'}); }
|
||||
return true;
|
||||
}
|
||||
}
|
||||
registerProcessor('tts-playback-processor', TTSPlaybackProcessor);
|
||||
"""
|
||||
|
||||
# ── Player HTML (container with metric cards) ────────────────────────
|
||||
PLAYER_HTML = """
|
||||
<div id="tts-player">
|
||||
<div style="display:flex; align-items:center; gap:10px;">
|
||||
<div id="tts-status-dot" style="width:10px;height:10px;border-radius:50%;background:#ccc;flex-shrink:0;"></div>
|
||||
<span id="tts-status" style="font-weight:600;font-size:1.05em;">Ready</span>
|
||||
<button id="tts-stop-btn" onclick="window.ttsStop()"
|
||||
style="display:none; margin-left:auto; padding:5px 16px; border-radius:6px; border:1px solid #EF5552;
|
||||
background:#fff; color:#EF5552; cursor:pointer; font-size:0.85em; transition:all 0.15s;">Stop</button>
|
||||
</div>
|
||||
<div id="tts-metrics" style="display:none; grid-template-columns:repeat(4,1fr); gap:10px; margin-top:12px;">
|
||||
<div style="background:#f8f9fa;border-radius:6px;padding:8px 10px;text-align:center;">
|
||||
<div style="font-size:0.7em;text-transform:uppercase;color:#888;letter-spacing:0.5px;margin-bottom:2px;">TTFP</div>
|
||||
<div id="tts-m-ttfp" style="font-size:1.2em;font-weight:700;color:#333;">—</div>
|
||||
</div>
|
||||
<div style="background:#f8f9fa;border-radius:6px;padding:8px 10px;text-align:center;">
|
||||
<div style="font-size:0.7em;text-transform:uppercase;color:#888;letter-spacing:0.5px;margin-bottom:2px;">RTF</div>
|
||||
<div id="tts-m-rtf" style="font-size:1.2em;font-weight:700;color:#333;">—</div>
|
||||
</div>
|
||||
<div style="background:#f8f9fa;border-radius:6px;padding:8px 10px;text-align:center;">
|
||||
<div style="font-size:0.7em;text-transform:uppercase;color:#888;letter-spacing:0.5px;margin-bottom:2px;">Audio</div>
|
||||
<div id="tts-m-dur" style="font-size:1.2em;font-weight:700;color:#333;">—</div>
|
||||
</div>
|
||||
<div style="background:#f8f9fa;border-radius:6px;padding:8px 10px;text-align:center;">
|
||||
<div style="font-size:0.7em;text-transform:uppercase;color:#888;letter-spacing:0.5px;margin-bottom:2px;">Speed</div>
|
||||
<div id="tts-m-speed" style="font-size:1.2em;font-weight:700;color:#333;">—</div>
|
||||
</div>
|
||||
</div>
|
||||
<div id="tts-rtf-bar-wrap" style="display:none; background:#e8ecf1; border-radius:4px; height:20px; overflow:hidden; position:relative; margin-top:10px;">
|
||||
<div id="tts-rtf-bar" style="height:100%; border-radius:4px; transition:width 0.3s ease, background 0.3s ease; width:0%;"></div>
|
||||
<span id="tts-rtf-label" style="position:absolute;top:50%;left:50%;transform:translate(-50%,-50%);font-size:0.75em;font-weight:600;color:#444;"></span>
|
||||
</div>
|
||||
<div id="tts-elapsed" style="display:none; margin-top:6px; font-size:0.8em; color:#999; text-align:right;"></div>
|
||||
</div>
|
||||
"""
|
||||
|
||||
|
||||
def _build_player_js(sample_rate: int) -> str:
|
||||
"""Build the JavaScript that powers the AudioWorklet player."""
|
||||
return f"""
|
||||
() => {{
|
||||
const SR = {sample_rate};
|
||||
const WC = {json.dumps(WORKLET_JS)};
|
||||
let ctx = null, node = null, abort = null, gen = false, st = {{}};
|
||||
|
||||
async function init() {{
|
||||
if (ctx) return;
|
||||
ctx = new AudioContext({{ sampleRate: SR }});
|
||||
const b = new Blob([WC], {{ type: 'application/javascript' }});
|
||||
const u = URL.createObjectURL(b);
|
||||
await ctx.audioWorklet.addModule(u);
|
||||
URL.revokeObjectURL(u);
|
||||
node = new AudioWorkletNode(ctx, 'tts-playback-processor');
|
||||
node.connect(ctx.destination);
|
||||
node.port.onmessage = (e) => {{
|
||||
if (e.data.type === 'started') setStatus('Playing...', '#64dd17');
|
||||
else if (e.data.type === 'stopped' && !gen) {{
|
||||
setStatus('Done', '#64dd17'); showStats(true);
|
||||
const btn = document.getElementById('tts-stop-btn');
|
||||
if (btn) btn.style.display = 'none';
|
||||
}}
|
||||
}};
|
||||
}}
|
||||
|
||||
function setStatus(text, color) {{
|
||||
const s = document.getElementById('tts-status');
|
||||
const d = document.getElementById('tts-status-dot');
|
||||
if (s) s.textContent = text;
|
||||
if (d) d.style.background = color || '#ccc';
|
||||
}}
|
||||
|
||||
function showStats(fin) {{
|
||||
if (!st.t0) return;
|
||||
const elapsed = (fin && st.streamEnd ? (st.streamEnd - st.t0) : (performance.now() - st.t0)) / 1000;
|
||||
const dur = st.samples / SR;
|
||||
const mTtfp = document.getElementById('tts-m-ttfp');
|
||||
const mRtf = document.getElementById('tts-m-rtf');
|
||||
const mDur = document.getElementById('tts-m-dur');
|
||||
const mSpeed = document.getElementById('tts-m-speed');
|
||||
const bar = document.getElementById('tts-rtf-bar');
|
||||
const barLabel = document.getElementById('tts-rtf-label');
|
||||
const elapsedEl = document.getElementById('tts-elapsed');
|
||||
|
||||
if (mTtfp && st.ttfp != null) mTtfp.textContent = st.ttfp.toFixed(0) + 'ms';
|
||||
if (mDur) mDur.textContent = dur.toFixed(1) + 's';
|
||||
|
||||
if (dur > 0 && elapsed > 0) {{
|
||||
const rtf = elapsed / dur;
|
||||
const speed = 1 / rtf;
|
||||
if (mRtf) {{
|
||||
mRtf.textContent = rtf.toFixed(2) + 'x';
|
||||
mRtf.style.color = rtf < 1 ? '#64dd17' : rtf < 1.5 ? '#e8a317' : '#EF5552';
|
||||
}}
|
||||
if (mSpeed) {{
|
||||
mSpeed.textContent = speed.toFixed(1) + 'x';
|
||||
mSpeed.style.color = speed > 1 ? '#64dd17' : speed > 0.7 ? '#e8a317' : '#EF5552';
|
||||
}}
|
||||
if (bar) {{
|
||||
const pct = Math.min(speed / 3 * 100, 100);
|
||||
bar.style.width = pct + '%';
|
||||
bar.style.background = speed > 1 ? 'linear-gradient(90deg,#4A90D9,#64dd17)' : speed > 0.7 ? 'linear-gradient(90deg,#e8a317,#f0b866)' : 'linear-gradient(90deg,#EF5552,#f87171)';
|
||||
}}
|
||||
if (barLabel) barLabel.textContent = speed.toFixed(1) + 'x realtime';
|
||||
}}
|
||||
if (elapsedEl) {{
|
||||
elapsedEl.style.display = 'block';
|
||||
elapsedEl.textContent = fin ? 'Completed in ' + elapsed.toFixed(1) + 's (' + st.chunks + ' chunks)' : elapsed.toFixed(1) + 's elapsed (' + st.chunks + ' chunks)';
|
||||
}}
|
||||
}}
|
||||
|
||||
window.ttsStop = function() {{
|
||||
if (abort) abort.abort();
|
||||
if (node) node.port.postMessage({{ type: 'clear' }});
|
||||
gen = false;
|
||||
setStatus('Stopped', '#999');
|
||||
const btn = document.getElementById('tts-stop-btn');
|
||||
if (btn) btn.style.display = 'none';
|
||||
}};
|
||||
|
||||
window.ttsGenerate = async function(payload) {{
|
||||
try {{ await init(); if (ctx.state === 'suspended') await ctx.resume(); }}
|
||||
catch (e) {{ const s = document.getElementById('tts-status'); if (s) s.textContent = 'Audio init error: ' + e.message; return; }}
|
||||
|
||||
// Abort previous request and clear worklet buffer
|
||||
if (abort) abort.abort();
|
||||
node.port.postMessage({{ type: 'clear' }});
|
||||
// Wait for worklet to process clear before sending new data
|
||||
await new Promise(r => setTimeout(r, 50));
|
||||
node.port.postMessage({{ type: 'clear' }});
|
||||
|
||||
gen = true;
|
||||
st = {{ t0: null, chunks: 0, samples: 0, ttfp: null }};
|
||||
setStatus('Connecting...', '#4A90D9');
|
||||
const bEl = document.getElementById('tts-stop-btn');
|
||||
if (bEl) bEl.style.display = 'inline-block';
|
||||
const mp = document.getElementById('tts-metrics');
|
||||
if (mp) {{ mp.style.display = 'grid'; ['tts-m-ttfp','tts-m-rtf','tts-m-dur','tts-m-speed'].forEach(id => {{ const e = document.getElementById(id); if(e) {{ e.textContent = '—'; e.style.color = '#333'; }} }}); }}
|
||||
const bw = document.getElementById('tts-rtf-bar-wrap');
|
||||
if (bw) bw.style.display = 'block';
|
||||
const bar = document.getElementById('tts-rtf-bar');
|
||||
if (bar) bar.style.width = '0%';
|
||||
const bl = document.getElementById('tts-rtf-label');
|
||||
if (bl) bl.textContent = '';
|
||||
const ee = document.getElementById('tts-elapsed');
|
||||
if (ee) {{ ee.style.display = 'none'; ee.textContent = ''; }}
|
||||
abort = new AbortController();
|
||||
|
||||
try {{
|
||||
console.log('fetch payload:', JSON.stringify({{input: payload.input?.slice(0,30), task: payload.task_type, has_ref: !!payload.ref_audio, stream: payload.stream}}));
|
||||
st.t0 = performance.now();
|
||||
const r = await fetch('/proxy/v1/audio/speech', {{
|
||||
method: 'POST',
|
||||
headers: {{ 'Content-Type': 'application/json' }},
|
||||
body: JSON.stringify(payload),
|
||||
signal: abort.signal,
|
||||
}});
|
||||
if (!r.ok) {{ const t = await r.text(); throw new Error('Server ' + r.status + ': ' + t.slice(0, 200)); }}
|
||||
setStatus('Streaming...', '#4A90D9');
|
||||
|
||||
const reader = r.body.getReader();
|
||||
let left = new Uint8Array(0);
|
||||
while (true) {{
|
||||
const {{ done, value }} = await reader.read();
|
||||
if (done) break;
|
||||
let raw;
|
||||
if (left.length > 0) {{
|
||||
raw = new Uint8Array(left.length + value.length);
|
||||
raw.set(left); raw.set(value, left.length);
|
||||
}} else {{ raw = value; }}
|
||||
const usable = raw.length - (raw.length % 2);
|
||||
left = usable < raw.length ? raw.slice(usable) : new Uint8Array(0);
|
||||
if (usable > 0) {{
|
||||
const ab = new ArrayBuffer(usable);
|
||||
new Uint8Array(ab).set(raw.subarray(0, usable));
|
||||
const pcm = new Int16Array(ab);
|
||||
node.port.postMessage(pcm);
|
||||
st.chunks++;
|
||||
st.samples += pcm.length;
|
||||
if (st.ttfp == null) st.ttfp = performance.now() - st.t0;
|
||||
showStats(false);
|
||||
}}
|
||||
}}
|
||||
}} catch (e) {{
|
||||
if (e.name !== 'AbortError') {{
|
||||
setStatus('Error: ' + e.message, '#EF5552');
|
||||
console.error('TTS error:', e);
|
||||
}}
|
||||
}} finally {{
|
||||
// Freeze RTF at stream-end time (before playback finishes)
|
||||
st.streamEnd = performance.now();
|
||||
showStats(true);
|
||||
gen = false;
|
||||
if (st.samples > 0) setStatus('Finishing playback...', '#64dd17');
|
||||
else {{
|
||||
setStatus('No audio received', '#999');
|
||||
if (bEl) bEl.style.display = 'none';
|
||||
}}
|
||||
}}
|
||||
}};
|
||||
}}
|
||||
"""
|
||||
|
||||
|
||||
def generate_speech(
|
||||
api_base: str,
|
||||
@@ -50,7 +299,7 @@ def generate_speech(
|
||||
response_format: str,
|
||||
speed: float,
|
||||
):
|
||||
"""Non-streaming: call /v1/audio/speech and return audio."""
|
||||
"""Non-streaming: call /v1/audio/speech and return complete audio."""
|
||||
payload = build_payload(
|
||||
text,
|
||||
task_type,
|
||||
@@ -65,21 +314,17 @@ def generate_speech(
|
||||
speed,
|
||||
stream=False,
|
||||
)
|
||||
|
||||
try:
|
||||
with httpx.Client(timeout=300.0) as client:
|
||||
resp = client.post(
|
||||
f"{api_base}/v1/audio/speech",
|
||||
json=payload,
|
||||
headers={
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": "Bearer EMPTY",
|
||||
},
|
||||
headers={"Content-Type": "application/json", "Authorization": "Bearer EMPTY"},
|
||||
)
|
||||
except httpx.TimeoutException:
|
||||
raise gr.Error("Request timed out. The server may be busy.")
|
||||
except httpx.ConnectError:
|
||||
raise gr.Error(f"Cannot connect to server at {api_base}. Make sure the vLLM server is running.")
|
||||
raise gr.Error(f"Cannot connect to server at {api_base}.")
|
||||
|
||||
if resp.status_code != 200:
|
||||
raise gr.Error(f"Server error ({resp.status_code}): {resp.text}")
|
||||
@@ -87,8 +332,7 @@ def generate_speech(
|
||||
content_type = resp.headers.get("content-type", "")
|
||||
if "application/json" in content_type:
|
||||
try:
|
||||
error_data = resp.json()
|
||||
raise gr.Error(f"Server error: {error_data}")
|
||||
raise gr.Error(f"Server error: {resp.json()}")
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
@@ -101,113 +345,120 @@ def generate_speech(
|
||||
audio_np = audio_np[:, 0]
|
||||
return (sample_rate, audio_np.astype(np.float32))
|
||||
except Exception as e:
|
||||
raise gr.Error(f"Failed to decode audio response: {e}")
|
||||
raise gr.Error(f"Failed to decode audio: {e}")
|
||||
|
||||
|
||||
def generate_speech_stream(
|
||||
api_base: str,
|
||||
text: str,
|
||||
task_type: str,
|
||||
voice: str,
|
||||
language: str,
|
||||
instructions: str,
|
||||
ref_audio: tuple | None,
|
||||
ref_audio_url: str,
|
||||
ref_text: str,
|
||||
x_vector_only: bool,
|
||||
response_format: str,
|
||||
speed: float,
|
||||
chunk_seconds: float = 0.25,
|
||||
):
|
||||
"""Streaming: yield individual PCM chunks for progressive playback."""
|
||||
payload = build_payload(
|
||||
text,
|
||||
task_type,
|
||||
voice,
|
||||
language,
|
||||
instructions,
|
||||
ref_audio,
|
||||
ref_audio_url,
|
||||
ref_text,
|
||||
x_vector_only,
|
||||
response_format,
|
||||
speed,
|
||||
stream=True,
|
||||
)
|
||||
def create_app(api_base: str):
|
||||
"""Create the FastAPI app with streaming proxy + Gradio UI."""
|
||||
fastapi_app = FastAPI()
|
||||
|
||||
min_chunk_samples = int(PCM_SAMPLE_RATE * chunk_seconds)
|
||||
# Pre-buffer 2 chunks before first playback to avoid the gap
|
||||
# between chunk 1 and chunk 2 during progressive streaming.
|
||||
prebuffer_chunks = 2
|
||||
pending: list[np.ndarray] = []
|
||||
pending_len = 0
|
||||
chunks_yielded = 0
|
||||
prebuffer: list[np.ndarray] = []
|
||||
# Server-side payload store (avoids sending large base64 through Gradio textbox)
|
||||
_pending_payloads: dict[str, dict] = {}
|
||||
|
||||
try:
|
||||
for samples in stream_pcm_chunks(api_base, payload):
|
||||
float_samples = samples.astype(np.float32) / 32767.0
|
||||
pending.append(float_samples)
|
||||
pending_len += len(float_samples)
|
||||
if pending_len >= min_chunk_samples:
|
||||
audio_chunk = np.concatenate(pending)
|
||||
pending.clear()
|
||||
pending_len = 0
|
||||
|
||||
if chunks_yielded < prebuffer_chunks:
|
||||
prebuffer.append(audio_chunk)
|
||||
chunks_yielded += 1
|
||||
if chunks_yielded == prebuffer_chunks:
|
||||
yield (PCM_SAMPLE_RATE, np.concatenate(prebuffer))
|
||||
prebuffer.clear()
|
||||
else:
|
||||
yield (PCM_SAMPLE_RATE, audio_chunk)
|
||||
# Flush: prebuffer (if stream was shorter than 2 chunks) + remaining pending
|
||||
remaining = prebuffer + pending
|
||||
if remaining:
|
||||
yield (PCM_SAMPLE_RATE, np.concatenate(remaining))
|
||||
except httpx.TimeoutException:
|
||||
raise gr.Error("Request timed out. The server may be busy.")
|
||||
except httpx.ConnectError:
|
||||
raise gr.Error(f"Cannot connect to server at {api_base}. Make sure the vLLM server is running.")
|
||||
|
||||
|
||||
def on_stream_change(stream: bool):
|
||||
"""When streaming is enabled, lock format to PCM and disable speed."""
|
||||
if stream:
|
||||
return (
|
||||
gr.update(value="pcm", interactive=False),
|
||||
gr.update(interactive=False),
|
||||
# ── Streaming proxy (same-origin, no CORS issues) ────────────
|
||||
@fastapi_app.post("/proxy/v1/audio/speech")
|
||||
async def proxy_speech(request: Request):
|
||||
body = await request.json()
|
||||
# Check if this is a request ID referencing a stored payload
|
||||
req_id = body.get("_req_id")
|
||||
if req_id and req_id in _pending_payloads:
|
||||
body = _pending_payloads.pop(req_id)
|
||||
# Pass ref_audio URL directly to vLLM server (it handles URL resolution).
|
||||
# Pre-downloading and re-encoding adds ~2-3s to TTFP for large files.
|
||||
logger.info(
|
||||
"Proxy request: %s",
|
||||
{k: (f"<{len(str(v))} chars>" if k == "ref_audio" else v) for k, v in body.items()},
|
||||
)
|
||||
return (
|
||||
gr.update(value="wav", interactive=True),
|
||||
gr.update(interactive=True),
|
||||
)
|
||||
try:
|
||||
client = httpx.AsyncClient(timeout=300)
|
||||
resp = await client.send(
|
||||
client.build_request(
|
||||
"POST",
|
||||
f"{api_base}/v1/audio/speech",
|
||||
json=body,
|
||||
headers={"Authorization": "Bearer EMPTY", "Content-Type": "application/json"},
|
||||
),
|
||||
stream=True,
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.exception("Proxy connection error")
|
||||
await client.aclose()
|
||||
return Response(content=str(exc), status_code=502)
|
||||
|
||||
if resp.status_code != 200:
|
||||
content = await resp.aread()
|
||||
logger.error("Proxy upstream error %d: %s", resp.status_code, content[:200])
|
||||
await resp.aclose()
|
||||
await client.aclose()
|
||||
return Response(content=content, status_code=resp.status_code)
|
||||
|
||||
def build_interface(api_base: str, stream_chunk_seconds: float = 0.25):
|
||||
"""Build the Gradio interface."""
|
||||
async def relay():
|
||||
total = 0
|
||||
try:
|
||||
async for chunk in resp.aiter_bytes():
|
||||
total += len(chunk)
|
||||
yield chunk
|
||||
except Exception:
|
||||
logger.exception("Proxy relay error after %d bytes", total)
|
||||
finally:
|
||||
logger.info("Proxy relay done: %d bytes", total)
|
||||
await resp.aclose()
|
||||
await client.aclose()
|
||||
|
||||
return StreamingResponse(relay(), media_type="application/octet-stream")
|
||||
|
||||
# ── Gradio UI ────────────────────────────────────────────────
|
||||
voices = fetch_voices(api_base)
|
||||
|
||||
css = """
|
||||
#generate-btn button { width: 100%; }
|
||||
.task-info { padding: 8px 12px; border-radius: 6px;
|
||||
background: #f0f4ff; margin-bottom: 8px; }
|
||||
#streaming-player { border: 1px solid var(--border-color-primary) !important; border-radius: var(--block-radius) !important; padding: var(--block-padding) !important; }
|
||||
"""
|
||||
|
||||
with gr.Blocks(css=css, title="Qwen3-TTS Demo") as demo:
|
||||
gr.Markdown("# Qwen3-TTS Online Serving Demo")
|
||||
gr.Markdown(f"**Server:** `{api_base}`")
|
||||
theme = gr.themes.Default(
|
||||
primary_hue=gr.themes.Color(
|
||||
c50="#f0f5ff",
|
||||
c100="#dce6f9",
|
||||
c200="#b8cef3",
|
||||
c300="#8eb2eb",
|
||||
c400="#6496e0",
|
||||
c500="#4A90D9",
|
||||
c600="#3a7bc8",
|
||||
c700="#2d66b0",
|
||||
c800="#1f4f8f",
|
||||
c900="#163a6e",
|
||||
c950="#0e2650",
|
||||
),
|
||||
)
|
||||
|
||||
with gr.Blocks(
|
||||
css=css,
|
||||
title="Qwen3-TTS Demo",
|
||||
js=_build_player_js(PCM_SAMPLE_RATE),
|
||||
theme=theme,
|
||||
) as demo:
|
||||
gr.HTML(f"""
|
||||
<div style="display:flex; align-items:center; gap:16px; margin-bottom:8px;">
|
||||
<img src="https://raw.githubusercontent.com/vllm-project/vllm-omni/main/docs/source/logos/vllm-omni-logo.png"
|
||||
alt="vLLM-Omni" style="height:42px;">
|
||||
<div>
|
||||
<h1 style="margin:0; font-size:1.5em;">Qwen3-TTS Streaming Demo</h1>
|
||||
<span style="font-size:0.85em; color:#666;">
|
||||
Served by <a href="https://github.com/vllm-project/vllm-omni" target="_blank"
|
||||
style="color:#4A90D9; text-decoration:none; font-weight:600;">vLLM-Omni</a>
|
||||
· <code style="background:#eef2f7; padding:2px 6px; border-radius:4px; font-size:0.9em;">{api_base}</code>
|
||||
</span>
|
||||
</div>
|
||||
</div>
|
||||
""")
|
||||
|
||||
with gr.Row():
|
||||
# Left column: inputs
|
||||
with gr.Column(scale=3):
|
||||
text_input = gr.Textbox(
|
||||
label="Text to Synthesize",
|
||||
placeholder="Enter text here, e.g., Hello, how are you?",
|
||||
lines=4,
|
||||
)
|
||||
|
||||
with gr.Row():
|
||||
task_type = gr.Radio(
|
||||
choices=TASK_TYPES,
|
||||
@@ -221,53 +472,49 @@ def build_interface(api_base: str, stream_chunk_seconds: float = 0.25):
|
||||
label="Language",
|
||||
scale=1,
|
||||
)
|
||||
|
||||
voice = gr.Dropdown(
|
||||
choices=voices,
|
||||
value=voices[0] if voices else None,
|
||||
label="Speaker",
|
||||
label="speaker",
|
||||
visible=True,
|
||||
allow_custom_value=True,
|
||||
)
|
||||
|
||||
instructions = gr.Textbox(
|
||||
label="Instructions",
|
||||
placeholder=("e.g., Speak with excitement / A warm, friendly female voice"),
|
||||
placeholder="e.g., Speak with excitement / A warm, friendly female voice",
|
||||
lines=2,
|
||||
visible=True,
|
||||
info="Optional style/emotion instructions",
|
||||
)
|
||||
|
||||
ref_audio = gr.Audio(
|
||||
label="Reference Audio (upload for voice cloning)",
|
||||
type="numpy",
|
||||
sources=["upload", "microphone"],
|
||||
visible=False,
|
||||
)
|
||||
ref_audio_url = gr.Textbox(
|
||||
label="Reference Audio URL",
|
||||
placeholder=("https://example.com/reference.wav (alternative to uploading)"),
|
||||
lines=1,
|
||||
visible=False,
|
||||
)
|
||||
ref_text = gr.Textbox(
|
||||
label="Reference Audio Transcript",
|
||||
placeholder=("Transcript of the reference audio (optional, improves quality)"),
|
||||
lines=2,
|
||||
visible=False,
|
||||
)
|
||||
x_vector_only = gr.Checkbox(
|
||||
label="Use x-vector only",
|
||||
value=False,
|
||||
visible=False,
|
||||
info=("Skip reference transcript, use speaker embedding only (lower quality)"),
|
||||
)
|
||||
with gr.Column(visible=False) as ref_group:
|
||||
ref_audio = gr.Audio(
|
||||
label="Reference Audio (upload for voice cloning)",
|
||||
type="numpy",
|
||||
sources=["upload", "microphone"],
|
||||
)
|
||||
ref_audio_url = gr.Textbox(
|
||||
label="Reference Audio URL",
|
||||
placeholder="https://example.com/reference.wav",
|
||||
lines=1,
|
||||
)
|
||||
ref_text = gr.Textbox(
|
||||
label="Reference Audio Transcript",
|
||||
placeholder="Transcript of reference audio (optional, improves quality)",
|
||||
lines=2,
|
||||
)
|
||||
x_vector_only = gr.Checkbox(
|
||||
label="Use x-vector only",
|
||||
value=False,
|
||||
info="Skip reference transcript, use speaker embedding only",
|
||||
)
|
||||
|
||||
with gr.Row():
|
||||
response_format = gr.Dropdown(
|
||||
choices=["wav", "mp3", "flac", "pcm", "aac", "opus"],
|
||||
value="wav",
|
||||
value="pcm",
|
||||
label="Audio Format",
|
||||
scale=1,
|
||||
interactive=False,
|
||||
)
|
||||
speed = gr.Slider(
|
||||
minimum=0.25,
|
||||
@@ -276,60 +523,179 @@ def build_interface(api_base: str, stream_chunk_seconds: float = 0.25):
|
||||
step=0.05,
|
||||
label="Speed",
|
||||
scale=1,
|
||||
interactive=False,
|
||||
)
|
||||
stream_checkbox = gr.Checkbox(
|
||||
label="Stream output",
|
||||
value=False,
|
||||
info=(
|
||||
"Enable streaming (PCM format, speed disabled). "
|
||||
"For gapless streaming, use gradio_fastrtc_demo.py"
|
||||
),
|
||||
label="Stream (gapless)",
|
||||
value=True,
|
||||
info="AudioWorklet streaming (recommended)",
|
||||
scale=1,
|
||||
)
|
||||
|
||||
generate_btn = gr.Button(
|
||||
"Generate Speech",
|
||||
variant="primary",
|
||||
size="lg",
|
||||
elem_id="generate-btn",
|
||||
)
|
||||
with gr.Row():
|
||||
generate_btn = gr.Button(
|
||||
"Generate Speech",
|
||||
variant="primary",
|
||||
size="lg",
|
||||
elem_id="generate-btn",
|
||||
scale=3,
|
||||
)
|
||||
reset_btn = gr.Button(
|
||||
"Reset",
|
||||
variant="secondary",
|
||||
size="lg",
|
||||
scale=1,
|
||||
)
|
||||
|
||||
# Right column: output
|
||||
with gr.Column(scale=2):
|
||||
player_html = gr.HTML(
|
||||
value=PLAYER_HTML,
|
||||
visible=True,
|
||||
label="streaming player",
|
||||
elem_id="streaming-player",
|
||||
)
|
||||
audio_output = gr.Audio(
|
||||
label="Generated Audio",
|
||||
label="generated audio",
|
||||
interactive=False,
|
||||
streaming=True,
|
||||
autoplay=True,
|
||||
visible=False,
|
||||
)
|
||||
gr.Markdown(
|
||||
"### Task Types\n"
|
||||
"- **CustomVoice**: Use a predefined speaker "
|
||||
"(Vivian, Ryan, etc.) with optional style instructions\n"
|
||||
"- **VoiceDesign**: Describe the desired voice in natural "
|
||||
"language (instructions required)\n"
|
||||
"- **Base**: Clone a voice from reference audio "
|
||||
"(upload a file or provide a URL)"
|
||||
)
|
||||
with gr.Column(visible=True) as examples_cv:
|
||||
gr.Examples(
|
||||
examples=[
|
||||
[
|
||||
"Have you ever wondered what it would be like to travel through time and visit ancient civilizations? The possibilities are endless, from witnessing the construction of the pyramids to experiencing the Renaissance firsthand.",
|
||||
"ryan",
|
||||
"English",
|
||||
"",
|
||||
],
|
||||
[
|
||||
"其实我真的有发现,我是一个特别善于观察别人情绪的人。比如说在一个聚会上,我总能第一时间察觉到谁不太开心,然后想办法让大家都能融入到欢乐的氛围中来。",
|
||||
"vivian",
|
||||
"Chinese",
|
||||
"用特别愤怒的语气说",
|
||||
],
|
||||
[
|
||||
"It was a dark and stormy night when the old lighthouse keeper heard a knock at the door. He set down his cup of tea, adjusted his glasses, and walked slowly toward the entrance.",
|
||||
"aiden",
|
||||
"English",
|
||||
"Speak in a mysterious, suspenseful tone",
|
||||
],
|
||||
],
|
||||
inputs=[text_input, voice, language, instructions],
|
||||
label="examples",
|
||||
)
|
||||
with gr.Column(visible=False) as examples_vd:
|
||||
gr.Examples(
|
||||
examples=[
|
||||
[
|
||||
"It's in the top drawer... wait, it's empty? No way, that's impossible! I'm sure I put it there! Someone must have moved it while I was gone.",
|
||||
"English",
|
||||
"Speak in an incredulous tone, but with a hint of panic beginning to creep into your voice.",
|
||||
],
|
||||
[
|
||||
"哥哥,你回来啦,人家等了你好久好久了,要抱抱!你去哪里了呀,都不跟人家说一声,人家好担心你哦。",
|
||||
"Chinese",
|
||||
"体现撒娇稚嫩的萝莉女声,音调偏高且起伏明显,营造出黏人、做作又刻意卖萌的听觉效果。",
|
||||
],
|
||||
],
|
||||
inputs=[text_input, language, instructions],
|
||||
label="examples",
|
||||
)
|
||||
with gr.Column(visible=False) as examples_base:
|
||||
gr.Examples(
|
||||
examples=[
|
||||
[
|
||||
"Good one. Okay, fine, I'm just gonna leave this sock monkey here. Goodbye. I hope you enjoy the rest of your evening and have a wonderful time.",
|
||||
"English",
|
||||
"https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-TTS-Repo/clone_2.wav",
|
||||
"Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you.",
|
||||
],
|
||||
[
|
||||
"其实我真的有发现,我是一个特别善于观察别人情绪的人。比如说在一个聚会上,我总能第一时间察觉到谁不太开心。",
|
||||
"Chinese",
|
||||
"https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-TTS-Repo/clone_2.wav",
|
||||
"Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you.",
|
||||
],
|
||||
],
|
||||
inputs=[text_input, language, ref_audio_url, ref_text],
|
||||
label="examples",
|
||||
)
|
||||
gr.HTML("""
|
||||
<div style="text-align:center; padding:8px 0; margin-top:4px;">
|
||||
<a href="https://github.com/vllm-project/vllm-omni" target="_blank">
|
||||
<img src="https://raw.githubusercontent.com/vllm-project/vllm-omni/main/docs/source/logos/vllm-omni-logo.png"
|
||||
alt="vLLM-Omni" style="height:28px; opacity:0.7;">
|
||||
</a>
|
||||
</div>
|
||||
""")
|
||||
|
||||
# Hidden textbox to pass payload from Python → JavaScript
|
||||
hidden_payload = gr.Textbox(visible=False, elem_id="tts-payload")
|
||||
|
||||
# ── Event wiring ─────────────────────────────────────────
|
||||
def on_task_change(tt):
|
||||
is_cv = tt == "CustomVoice"
|
||||
is_vd = tt == "VoiceDesign"
|
||||
is_base = tt == "Base"
|
||||
return (
|
||||
gr.update(visible=is_cv), # voice
|
||||
gr.update(
|
||||
visible=(is_cv or is_vd),
|
||||
info="Required: describe the voice style" if is_vd else "Optional style/emotion instructions",
|
||||
), # instructions
|
||||
gr.update(visible=is_base), # ref_group
|
||||
gr.update(visible=is_cv), # examples_cv
|
||||
gr.update(visible=is_vd), # examples_vd
|
||||
gr.update(visible=is_base), # examples_base
|
||||
)
|
||||
|
||||
# Dynamic UI updates
|
||||
task_type.change(
|
||||
fn=on_task_type_change,
|
||||
fn=on_task_change,
|
||||
inputs=[task_type],
|
||||
outputs=[
|
||||
voice,
|
||||
instructions,
|
||||
ref_audio,
|
||||
ref_audio_url,
|
||||
ref_text,
|
||||
x_vector_only,
|
||||
ref_group,
|
||||
examples_cv,
|
||||
examples_vd,
|
||||
examples_base,
|
||||
],
|
||||
)
|
||||
|
||||
def on_stream_change(stream: bool):
|
||||
if stream:
|
||||
return (
|
||||
gr.update(value="pcm", interactive=False),
|
||||
gr.update(interactive=False),
|
||||
gr.update(visible=True), # player
|
||||
gr.update(visible=False), # audio
|
||||
)
|
||||
return (
|
||||
gr.update(value="wav", interactive=True),
|
||||
gr.update(interactive=True),
|
||||
gr.update(visible=False),
|
||||
gr.update(visible=True),
|
||||
)
|
||||
|
||||
stream_checkbox.change(
|
||||
fn=on_stream_change,
|
||||
inputs=[stream_checkbox],
|
||||
outputs=[response_format, speed],
|
||||
outputs=[response_format, speed, player_html, audio_output],
|
||||
)
|
||||
|
||||
def on_reset():
|
||||
return (
|
||||
"", # text
|
||||
None, # audio_output
|
||||
"", # hidden_payload
|
||||
PLAYER_HTML, # reset player
|
||||
)
|
||||
|
||||
reset_btn.click(
|
||||
fn=on_reset,
|
||||
outputs=[text_input, audio_output, hidden_payload, player_html],
|
||||
js="() => { if (window.ttsStop) window.ttsStop(); }",
|
||||
)
|
||||
|
||||
all_inputs = [
|
||||
@@ -346,58 +712,85 @@ def build_interface(api_base: str, stream_chunk_seconds: float = 0.25):
|
||||
speed,
|
||||
]
|
||||
|
||||
def dispatch(stream_enabled, *args):
|
||||
def on_generate(stream_enabled, *args):
|
||||
if stream_enabled:
|
||||
yield from generate_speech_stream(
|
||||
api_base,
|
||||
*args,
|
||||
chunk_seconds=stream_chunk_seconds,
|
||||
import time as _time
|
||||
|
||||
text, task_type_v, voice_v, lang_v, instr_v, ref_a, ref_url, ref_t, xvec, _fmt, _spd = args
|
||||
# For streaming, use URL for ref_audio (base64 is too large
|
||||
# for the Gradio textbox → JS → fetch pipeline).
|
||||
# If user uploaded audio but no URL, encode and store server-side.
|
||||
if ref_a is not None and not (ref_url and ref_url.strip()):
|
||||
req_id = f"req-{int(_time.time() * 1000)}"
|
||||
full_payload = build_payload(
|
||||
text,
|
||||
task_type_v,
|
||||
voice_v,
|
||||
lang_v,
|
||||
instr_v,
|
||||
ref_a,
|
||||
ref_url,
|
||||
ref_t,
|
||||
xvec,
|
||||
stream=True,
|
||||
)
|
||||
_pending_payloads[req_id] = full_payload
|
||||
browser_payload = {"_req_id": req_id, "_nonce": int(_time.time() * 1000)}
|
||||
return json.dumps(browser_payload), gr.update()
|
||||
|
||||
# URL-only path: payload is small, pass directly to browser
|
||||
payload = build_payload(
|
||||
text,
|
||||
task_type_v,
|
||||
voice_v,
|
||||
lang_v,
|
||||
instr_v,
|
||||
None,
|
||||
ref_url,
|
||||
ref_t,
|
||||
xvec,
|
||||
stream=True,
|
||||
)
|
||||
payload["_nonce"] = int(_time.time() * 1000)
|
||||
return json.dumps(payload), gr.update()
|
||||
else:
|
||||
yield generate_speech(api_base, *args)
|
||||
audio = generate_speech(api_base, *args)
|
||||
return "", audio
|
||||
|
||||
generate_btn.click(
|
||||
fn=dispatch,
|
||||
fn=on_generate,
|
||||
inputs=[stream_checkbox] + all_inputs,
|
||||
outputs=[audio_output],
|
||||
outputs=[hidden_payload, audio_output],
|
||||
).then(
|
||||
fn=lambda p: p,
|
||||
inputs=[hidden_payload],
|
||||
outputs=[hidden_payload],
|
||||
js="(p) => { if (p && p.trim()) { const d = JSON.parse(p); delete d._nonce; window.ttsGenerate(d); } return p; }",
|
||||
)
|
||||
|
||||
demo.queue()
|
||||
return demo
|
||||
|
||||
return gr.mount_gradio_app(fastapi_app, demo, path="/")
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Gradio demo for Qwen3-TTS online serving.")
|
||||
add_common_args(parser)
|
||||
parser.add_argument(
|
||||
"--stream-chunk-seconds",
|
||||
type=float,
|
||||
default=0.25,
|
||||
help=(
|
||||
"Seconds of audio to buffer per streaming chunk (default: 0.25). "
|
||||
"Lower = smoother but more overhead, higher = chunkier."
|
||||
),
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Gradio demo for Qwen3-TTS with gapless AudioWorklet streaming.",
|
||||
)
|
||||
add_common_args(parser)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
print(f"Connecting to vLLM server at: {args.api_base}")
|
||||
demo = build_interface(
|
||||
args.api_base,
|
||||
stream_chunk_seconds=args.stream_chunk_seconds,
|
||||
)
|
||||
|
||||
try:
|
||||
demo.launch(
|
||||
server_name=args.host,
|
||||
server_port=args.port,
|
||||
share=args.share,
|
||||
)
|
||||
except KeyboardInterrupt:
|
||||
print("\nShutting down...")
|
||||
app = create_app(args.api_base)
|
||||
|
||||
import uvicorn
|
||||
|
||||
uvicorn.run(app, host=args.host, port=args.port)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -1,274 +0,0 @@
|
||||
"""Gradio demo for Qwen3-TTS using FastRTC for gapless streaming audio.
|
||||
|
||||
Uses WebRTC via FastRTC to deliver audio chunks as a continuous real-time
|
||||
stream, avoiding the inter-chunk gaps present in Gradio's built-in streaming.
|
||||
|
||||
Requires: pip install fastrtc
|
||||
|
||||
Supports all 3 task types:
|
||||
- CustomVoice: Predefined speaker with optional style instructions
|
||||
- VoiceDesign: Natural language voice description
|
||||
- Base: Voice cloning from reference audio (URL only in this demo)
|
||||
|
||||
Usage:
|
||||
# Start the server first (see run_server.sh), then:
|
||||
python gradio_fastrtc_demo.py --api-base http://localhost:8000
|
||||
|
||||
# With public sharing:
|
||||
python gradio_fastrtc_demo.py --api-base http://localhost:8000 --share
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import logging
|
||||
from collections.abc import Generator
|
||||
|
||||
import gradio as gr
|
||||
import numpy as np
|
||||
from fastrtc import WebRTC
|
||||
from tts_common import (
|
||||
PCM_SAMPLE_RATE,
|
||||
SUPPORTED_LANGUAGES,
|
||||
TASK_TYPES,
|
||||
add_common_args,
|
||||
build_payload,
|
||||
fetch_voices,
|
||||
stream_pcm_chunks,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def make_tts_generator(api_base: str):
|
||||
"""Create a generator function for FastRTC's mode='receive'.
|
||||
|
||||
FastRTC calls this with the additional_inputs values whenever the user
|
||||
clicks "Start Stream". It must be a synchronous generator yielding
|
||||
(sample_rate, int16_ndarray) tuples.
|
||||
"""
|
||||
|
||||
def tts_generate(
|
||||
text: str,
|
||||
task_type: str,
|
||||
voice: str,
|
||||
language: str,
|
||||
instructions: str,
|
||||
ref_audio_url: str,
|
||||
ref_text: str,
|
||||
x_vector_only: bool,
|
||||
) -> Generator[tuple[int, np.ndarray], None, None]:
|
||||
if not text or not text.strip():
|
||||
return
|
||||
|
||||
payload = build_payload(
|
||||
text,
|
||||
task_type,
|
||||
voice,
|
||||
language,
|
||||
instructions,
|
||||
ref_audio=None, # WebRTC UI doesn't support file upload
|
||||
ref_audio_url=ref_audio_url,
|
||||
ref_text=ref_text,
|
||||
x_vector_only=x_vector_only,
|
||||
stream=True,
|
||||
)
|
||||
|
||||
try:
|
||||
for samples in stream_pcm_chunks(api_base, payload):
|
||||
yield (PCM_SAMPLE_RATE, samples)
|
||||
except Exception:
|
||||
logger.exception("TTS streaming error")
|
||||
|
||||
return tts_generate
|
||||
|
||||
|
||||
def build_interface(api_base: str):
|
||||
"""Build a custom Gradio Blocks UI matching gradio_demo.py layout."""
|
||||
voices = fetch_voices(api_base)
|
||||
handler = make_tts_generator(api_base)
|
||||
|
||||
css = """
|
||||
#generate-btn button { width: 100%; }
|
||||
.task-info { padding: 8px 12px; border-radius: 6px;
|
||||
background: #f0f4ff; margin-bottom: 8px; }
|
||||
"""
|
||||
|
||||
with gr.Blocks(css=css, title="Qwen3-TTS (FastRTC)") as demo:
|
||||
gr.Markdown("# Qwen3-TTS Online Serving Demo (FastRTC)")
|
||||
gr.Markdown(f"**Server:** `{api_base}` | **Transport:** WebRTC (gapless streaming)")
|
||||
|
||||
with gr.Row():
|
||||
# Left column: inputs
|
||||
with gr.Column(scale=3):
|
||||
text_input = gr.Textbox(
|
||||
label="Text to Synthesize",
|
||||
placeholder="Enter text here, e.g., Hello, how are you?",
|
||||
lines=4,
|
||||
)
|
||||
|
||||
with gr.Row():
|
||||
task_type = gr.Radio(
|
||||
choices=TASK_TYPES,
|
||||
value="CustomVoice",
|
||||
label="Task Type",
|
||||
scale=2,
|
||||
)
|
||||
language = gr.Dropdown(
|
||||
choices=SUPPORTED_LANGUAGES,
|
||||
value="Auto",
|
||||
label="Language",
|
||||
scale=1,
|
||||
)
|
||||
|
||||
voice = gr.Dropdown(
|
||||
choices=voices,
|
||||
value=voices[0] if voices else None,
|
||||
label="Speaker",
|
||||
visible=True,
|
||||
)
|
||||
|
||||
instructions = gr.Textbox(
|
||||
label="Instructions",
|
||||
placeholder=("e.g., Speak with excitement / A warm, friendly female voice"),
|
||||
lines=2,
|
||||
visible=True,
|
||||
info="Optional style/emotion instructions",
|
||||
)
|
||||
|
||||
# Base (voice clone) controls
|
||||
# Note: file upload is not supported with WebRTC,
|
||||
# use a URL instead.
|
||||
ref_audio_url = gr.Textbox(
|
||||
label="Reference Audio URL",
|
||||
placeholder=("https://example.com/reference.wav (alternative to uploading)"),
|
||||
lines=1,
|
||||
visible=False,
|
||||
)
|
||||
ref_text = gr.Textbox(
|
||||
label="Reference Audio Transcript",
|
||||
placeholder=("Transcript of the reference audio (optional, improves quality)"),
|
||||
lines=2,
|
||||
visible=False,
|
||||
)
|
||||
x_vector_only = gr.Checkbox(
|
||||
label="Use x-vector only",
|
||||
value=False,
|
||||
visible=False,
|
||||
info=("Skip reference transcript, use speaker embedding only (lower quality)"),
|
||||
)
|
||||
|
||||
generate_btn = gr.Button(
|
||||
"Generate Speech",
|
||||
variant="primary",
|
||||
size="lg",
|
||||
elem_id="generate-btn",
|
||||
)
|
||||
|
||||
# Right column: output
|
||||
with gr.Column(scale=2):
|
||||
webrtc_output = WebRTC(
|
||||
label="Generated Audio (WebRTC)",
|
||||
mode="receive",
|
||||
modality="audio",
|
||||
)
|
||||
gr.Markdown(
|
||||
"### Task Types\n"
|
||||
"- **CustomVoice**: Use a predefined speaker "
|
||||
"(Vivian, Ryan, etc.) with optional style instructions\n"
|
||||
"- **VoiceDesign**: Describe the desired voice in natural "
|
||||
"language (instructions required)\n"
|
||||
"- **Base**: Clone a voice from reference audio "
|
||||
"(provide a URL)"
|
||||
)
|
||||
|
||||
# Dynamic UI updates for task type
|
||||
# FastRTC doesn't support file upload, so we only toggle
|
||||
# 5 components (no ref_audio)
|
||||
def on_task_type_change_fastrtc(task_type: str):
|
||||
if task_type == "CustomVoice":
|
||||
return (
|
||||
gr.update(visible=True), # voice
|
||||
gr.update(visible=True, info="Optional style/emotion instructions"),
|
||||
gr.update(visible=False), # ref_audio_url
|
||||
gr.update(visible=False), # ref_text
|
||||
gr.update(visible=False), # x_vector_only
|
||||
)
|
||||
elif task_type == "VoiceDesign":
|
||||
return (
|
||||
gr.update(visible=False),
|
||||
gr.update(visible=True, info="Required: describe the voice style"),
|
||||
gr.update(visible=False),
|
||||
gr.update(visible=False),
|
||||
gr.update(visible=False),
|
||||
)
|
||||
elif task_type == "Base":
|
||||
return (
|
||||
gr.update(visible=False),
|
||||
gr.update(visible=False),
|
||||
gr.update(visible=True),
|
||||
gr.update(visible=True),
|
||||
gr.update(visible=True),
|
||||
)
|
||||
return (
|
||||
gr.update(visible=True),
|
||||
gr.update(visible=True),
|
||||
gr.update(visible=False),
|
||||
gr.update(visible=False),
|
||||
gr.update(visible=False),
|
||||
)
|
||||
|
||||
task_type.change(
|
||||
fn=on_task_type_change_fastrtc,
|
||||
inputs=[task_type],
|
||||
outputs=[
|
||||
voice,
|
||||
instructions,
|
||||
ref_audio_url,
|
||||
ref_text,
|
||||
x_vector_only,
|
||||
],
|
||||
)
|
||||
|
||||
# Wire up WebRTC streaming
|
||||
all_inputs = [
|
||||
text_input,
|
||||
task_type,
|
||||
voice,
|
||||
language,
|
||||
instructions,
|
||||
ref_audio_url,
|
||||
ref_text,
|
||||
x_vector_only,
|
||||
]
|
||||
|
||||
webrtc_output.stream(
|
||||
fn=handler,
|
||||
inputs=all_inputs,
|
||||
outputs=[webrtc_output],
|
||||
trigger=generate_btn.click,
|
||||
)
|
||||
|
||||
demo.queue()
|
||||
return demo
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Gradio FastRTC demo for Qwen3-TTS (gapless streaming).")
|
||||
add_common_args(parser)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
print(f"Connecting to vLLM server at: {args.api_base}")
|
||||
|
||||
demo = build_interface(args.api_base)
|
||||
demo.launch(
|
||||
server_name=args.host,
|
||||
server_port=args.port,
|
||||
share=args.share,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Loaded 100 of 346 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user