Compare commits

..
Author SHA1 Message Date
Hongsheng Liu 77f8ebdca2 Revert "[misc] Feature/pr reviewer auto trigger&update model (#1431)"
This reverts commit 267170115b.
2026-02-23 07:51:33 +08:00
305 changed files with 5358 additions and 36887 deletions

No files matched your search

+306 -24
View File
@@ -3,38 +3,320 @@ steps:
key: image-build
commands:
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
- "docker build --progress=plain --file docker/Dockerfile.ci -t vllm-omni-ci ."
- "docker build --file docker/Dockerfile.ci -t vllm-omni-ci ."
- "docker tag vllm-omni-ci public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
- "docker push public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
agents:
queue: "cpu_queue_premerge"
# L2 Test
- label: "Upload Ready Pipeline"
depends_on: image-build
key: upload-ready-pipeline
if: build.branch != "main"
commands:
- buildkite-agent pipeline upload .buildkite/test-ready.yml
agents:
queue: "cpu_queue_premerge"
# L3 Test
- label: "Upload Merge Pipeline"
depends_on: image-build
key: upload-merge-pipeline
if: build.branch == "main"
commands:
- buildkite-agent pipeline upload .buildkite/test-merge.yml
agents:
queue: "cpu_queue_premerge"
# L4 Test
- label: "Upload Nightly Pipeline"
depends_on: image-build
key: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1"
commands:
- buildkite-agent pipeline upload .buildkite/test-nightly.yml
- buildkite-agent pipeline upload .buildkite/test-nightly.yaml
agents:
queue: "cpu_queue_premerge"
- label: "Simple Unit Test"
depends_on: image-build
commands:
- "pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
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 Test"
timeout_in_minutes: 20
depends_on: image-build
commands:
- pytest -s -v tests/e2e/offline_inference/test_t2i_model.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"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Diffusion Model CPU offloading Test"
timeout_in_minutes: 20
depends_on: image-build
commands:
- pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
- pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
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"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Audio Generation Model Test"
timeout_in_minutes: 20
depends_on: image-build
commands:
- pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.py
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"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Diffusion Cache Backend Test"
timeout_in_minutes: 15
depends_on: image-build
commands:
- pytest -s -v -m 'core_model and cache and diffusion and not distributed_cuda and L4'
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"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Diffusion Sequence Parallelism Test"
timeout_in_minutes: 20
depends_on: image-build
commands:
- pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py -m core_model
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
shm-size: "8gb"
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Diffusion GPU Worker Test"
timeout_in_minutes: 20
depends_on: image-build
commands:
- pytest -s -v tests/diffusion/test_diffusion_worker.py
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
shm-size: "8gb"
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
# - label: "Benchmark&Engine Test"
# timeout_in_minutes: 15
# depends_on: image-build
# commands:
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
# - pytest -s -v tests/benchmarks/test_serve_cli.py
# - pytest -s -v tests/engine/test_async_omni_engine_abort.py
# agents:
# queue: "mithril-h100-pool"
# plugins:
# - kubernetes:
# podSpec:
# containers:
# - image: public.ecr.aws/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
# 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: "Omni Model Test"
timeout_in_minutes: 20
depends_on: image-build
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
# - label: "Omni Model Test with H100"
# timeout_in_minutes: 20
# depends_on: image-build
# commands:
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
# - 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 "core_model" --run-level "core_model"
# agents:
# queue: "mithril-h100-pool"
# plugins:
# - kubernetes:
# podSpec:
# containers:
# - image: public.ecr.aws/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
# 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"
timeout_in_minutes: 20
depends_on: image-build
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py
agents:
queue: "gpu_4_queue"
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
shm-size: "8gb"
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
# - label: "Diffusion Image Edit Test with H100 (1 GPU)"
# timeout_in_minutes: 20
# depends_on: image-build
# commands:
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
# - pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
# agents:
# queue: "mithril-h100-pool"
# plugins:
# - kubernetes:
# podSpec:
# containers:
# - image: public.ecr.aws/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: "Bagel Text2Img Model Test with H100"
# timeout_in_minutes: 30
# depends_on: image-build
# commands:
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
# - pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py
# agents:
# queue: "mithril-h100-pool"
# plugins:
# - kubernetes:
# podSpec:
# containers:
# - image: public.ecr.aws/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
+1 -3
View File
@@ -33,7 +33,6 @@ steps:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
- pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
- label: "Diffusion Cache Backend Test"
timeout_in_minutes: 15
@@ -89,7 +88,6 @@ steps:
- 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/engine/test_async_omni_engine_abort.py
- label: "Omni Model Test Qwen3-Omni"
timeout_in_minutes: 15
@@ -103,7 +101,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
- pytest -s -v tests/engine/test_async_omni_engine_abort.py
- label: "Diffusion Image Edit Test"
timeout_in_minutes: 15
+86 -65
View File
@@ -1,6 +1,16 @@
steps:
- label: ":docker: Build image"
key: image-build
commands:
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
- "docker build --file docker/Dockerfile.ci -t vllm-omni-ci ."
- "docker tag vllm-omni-ci public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
- "docker push public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
agents:
queue: "cpu_queue_premerge"
- label: "Simple Unit Test"
depends_on: upload-merge-pipeline
depends_on: image-build
commands:
- "pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
agents:
@@ -17,7 +27,7 @@ steps:
- label: "Diffusion Model Test"
timeout_in_minutes: 20
depends_on: upload-merge-pipeline
depends_on: image-build
commands:
- pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "advanced_model and diffusion" --run-level "advanced_model"
agents:
@@ -34,7 +44,7 @@ steps:
- label: "Diffusion Images API LoRA E2E"
timeout_in_minutes: 20
depends_on: upload-merge-pipeline
depends_on: image-build
commands:
- pytest -s -v tests/e2e/online_serving/test_images_generations_lora.py
agents:
@@ -51,7 +61,7 @@ steps:
- label: "Diffusion Model CPU offloading Test"
timeout_in_minutes: 20
depends_on: upload-merge-pipeline
depends_on: image-build
commands:
- pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
- pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
@@ -69,7 +79,7 @@ steps:
- label: "Audio Generation Model Test"
timeout_in_minutes: 20
depends_on: upload-merge-pipeline
depends_on: image-build
commands:
- pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.py
agents:
@@ -86,9 +96,9 @@ steps:
- label: "Diffusion Cache Backend Test"
timeout_in_minutes: 15
depends_on: upload-merge-pipeline
depends_on: image-build
commands:
- pytest -s -v -m 'core_model and cache and diffusion and not distributed_cuda and L4'
- pytest -s -v tests/e2e/offline_inference/test_cache_dit.py tests/e2e/offline_inference/test_teacache.py
agents:
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
plugins:
@@ -103,7 +113,7 @@ steps:
- label: "Diffusion Sequence Parallelism Test"
timeout_in_minutes: 20
depends_on: upload-merge-pipeline
depends_on: image-build
commands:
- pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py
agents:
@@ -121,9 +131,9 @@ steps:
- label: "Diffusion Tensor Parallelism Test"
timeout_in_minutes: 20
depends_on: upload-merge-pipeline
depends_on: image-build
commands:
- pytest -s -v tests/e2e/offline_inference/test_zimage_parallelism.py
- pytest -s -v tests/e2e/offline_inference/test_zimage_tensor_parallel.py
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
plugins:
@@ -139,7 +149,7 @@ steps:
- label: "Diffusion GPU Worker Test"
timeout_in_minutes: 20
depends_on: upload-merge-pipeline
depends_on: image-build
commands:
- pytest -s -v tests/diffusion/test_diffusion_worker.py
agents:
@@ -155,66 +165,21 @@ steps:
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Benchmark & Engine Test"
depends_on: upload-merge-pipeline
commands:
- |
timeout 15m bash -c '
export VLLM_WORKER_MULTIPROC_METHOD=spawn
set +e
pytest -s -v tests/benchmarks/test_serve_cli.py
EXIT1=$$?
pytest -s -v tests/engine/test_async_omni_engine_abort.py
EXIT2=$$?
exit $$((EXIT1 | EXIT2))
'
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"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Omni Model Test"
- label: "Benchmark&Engine Test"
timeout_in_minutes: 15
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
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Omni Model Test with H100"
timeout_in_minutes: 30
depends_on: upload-merge-pipeline
depends_on: image-build
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- 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/benchmarks/test_serve_cli.py
- pytest -s -v tests/engine/test_async_omni_engine_abort.py
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
- image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 2
@@ -237,9 +202,65 @@ steps:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: "Omni Model Test"
timeout_in_minutes: 15
depends_on: image-build
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Omni Model Test with H100"
timeout_in_minutes: 30
depends_on: image-build
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- 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"
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: public.ecr.aws/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
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: "Diffusion Image Edit Test with H100 (1 GPU)"
timeout_in_minutes: 20
depends_on: upload-merge-pipeline
depends_on: image-build
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
@@ -249,7 +270,7 @@ steps:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
- image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 1
@@ -274,7 +295,7 @@ steps:
# - label: "Bagel Text2Img Model Test with H100"
# timeout_in_minutes: 30
# depends_on: upload-merge-pipeline
# depends_on: image-build
# commands:
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
# - pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py
@@ -284,7 +305,7 @@ steps:
# - kubernetes:
# podSpec:
# containers:
# - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
# - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
# resources:
# limits:
# nvidia.com/gpu: 1
@@ -1,24 +1,19 @@
steps:
- label: "Omni Model Test with H100"
timeout_in_minutes: 90
depends_on: upload-nightly-pipeline
timeout_in_minutes: 60
depends_on: image-build
if: build.env("NIGHTLY") == "1"
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- |
set +e
pytest -s -v tests/e2e/online_serving/test_qwen3_omni_expansion.py -m "advanced_model" --run-level "advanced_model"
EXIT1=$$?
pytest -s -v tests/examples/online_serving/test_qwen3_omni.py -m "advanced_model" --run-level "advanced_model"
EXIT2=$$?
exit $$((EXIT1 | EXIT2))
- pytest -s -v tests/e2e/online_serving/test_qwen3_omni_expansion.py -m "advanced_model" --run-level "advanced_model"
- pytest -s -v tests/examples/online_serving/test_qwen3_omni.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
- image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 2
@@ -43,9 +38,10 @@ steps:
- label: "Omni Model Test"
timeout_in_minutes: 60
depends_on: upload-nightly-pipeline
depends_on: image-build
if: build.env("NIGHTLY") == "1"
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- 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"
agents:
@@ -54,16 +50,16 @@ steps:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
shm-size: "8gb"
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Omni Model Perf Test"
timeout_in_minutes: 120
depends_on: upload-nightly-pipeline
depends_on: image-build
if: build.env("NIGHTLY") == "1"
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
@@ -74,7 +70,7 @@ steps:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
- image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 2
-306
View File
@@ -1,306 +0,0 @@
steps:
- label: "Simple Unit Test"
depends_on: upload-ready-pipeline
commands:
- "timeout 20m pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
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 Test"
depends_on: upload-ready-pipeline
commands:
- timeout 20m pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "core_model and diffusion" --run-level "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:
- |
timeout 20m bash -c '
set +e
pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
EXIT1=$$?
pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
EXIT2=$$?
exit $$((EXIT1 | EXIT2))
'
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"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Audio Generation Model Test"
depends_on: upload-ready-pipeline
commands:
- timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.py
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"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Diffusion Cache Backend Test"
depends_on: upload-ready-pipeline
commands:
- timeout 15m pytest -s -v -m 'core_model and cache and diffusion and not distributed_cuda and L4'
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"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Diffusion Sequence Parallelism Test"
depends_on: upload-ready-pipeline
commands:
- timeout 20m pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py -m core_model
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
shm-size: "8gb"
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Diffusion GPU Worker Test"
depends_on: upload-ready-pipeline
commands:
- timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
shm-size: "8gb"
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Benchmark & Engine Test"
depends_on: upload-ready-pipeline
commands:
- |
timeout 15m bash -c '
export VLLM_WORKER_MULTIPROC_METHOD=spawn
set +e
pytest -s -v tests/benchmarks/test_serve_cli.py
EXIT1=$$?
pytest -s -v tests/engine/test_async_omni_engine_abort.py
EXIT2=$$?
exit $$((EXIT1 | EXIT2))
'
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"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Omni Model Test"
depends_on: upload-ready-pipeline
commands:
- |
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
'
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
# - label: "Omni Model Test with H100"
# depends_on: upload-ready-pipeline
# commands:
# - |
# timeout 20m bash -c '
# export VLLM_WORKER_MULTIPROC_METHOD=spawn
# 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 "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
# 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:
- |
timeout 20m bash -c '
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py
'
agents:
queue: "gpu_4_queue"
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
shm-size: "8gb"
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
# - label: "Diffusion Image Edit Test with H100 (1 GPU)"
# depends_on: upload-ready-pipeline
# commands:
# - |
# timeout 20m bash -c '
# export VLLM_WORKER_MULTIPROC_METHOD=spawn
# pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
# '
# 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: "Bagel Text2Img Model Test with H100"
# depends_on: upload-ready-pipeline
# commands:
# - |
# timeout 30m bash -c '
# export VLLM_WORKER_MULTIPROC_METHOD=spawn
# pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py
# '
# 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
+1 -1
View File
@@ -14,7 +14,7 @@
soft_fail: false
commands:
- "source .buildkite/scripts/docker_login_ecr_public.sh && safe_docker_login_ecr_public"
- "docker build -f docker/Dockerfile.rocm -t {{ docker_image_amd }} --target test --build-arg ARG_PYTORCH_ROCM_ARCH=gfx942 --progress plain ."
- "docker build -f docker/Dockerfile.rocm -t {{ docker_image_amd }} --target final --progress plain ."
- "docker push {{ docker_image_amd }}"
key: "amd-build"
env:
+3 -3
View File
@@ -12,10 +12,10 @@ PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTT
<summary> Essential Elements of an Effective PR Description Checklist </summary>
- [ ] The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
- [ ] The test plan. Please provide the test scripts & test commands. Please state the reasons if your codes don't require additional test scripts. For test file guidelines, please check the [test style doc](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/tests_style/)
- [ ] The test results. Please paste the results comparison before and after, or the e2e results.
- [ ] The test plan. Please providing the test scripts & test commands. Please state the reasons if your codes don't require additional test scripts. For test file guidelines, please check the [test style doc](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/tests_style/)
- [ ] The test results. Please pasting the results comparison before and after, or e2e results.
- [ ] (Optional) The necessary documentation update, such as updating `supported_models.md` and `examples` for a new model. **Please run `mkdocs serve` to sync the documentation editions to `./docs`.**
- [ ] (Optional) Release notes update. If your change is user-facing, please update the release notes draft.
- [ ] (Optional) Release notes update. If your change is user facing, please update the release notes draft.
</details>
**BEFORE SUBMITTING, PLEASE READ <https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md>** (anything written below this line will be removed by GitHub Actions)
+27
View File
@@ -0,0 +1,27 @@
name: ComfyUI Integration -- Publish to Comfy registry
on:
workflow_dispatch:
push:
tags:
- "*"
paths:
- "apps/ComfyUI-vLLM-Omni/**"
permissions:
issues: write
defaults:
run:
working-directory: apps/ComfyUI-vLLM-Omni
jobs:
publish-node:
name: Publish Custom Node to registry
runs-on: ubuntu-latest
steps:
- name: ♻️ Check out code
uses: actions/checkout@v4
- name: 📦 Publish Custom Node
uses: Comfy-Org/publish-node-action@main
with:
personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }}
+40
View File
@@ -0,0 +1,40 @@
name: ComfyUI Integration -- Validate backwards compatibility
on:
pull_request:
branches:
- master
- main
paths:
- "apps/ComfyUI-vLLM-Omni/**"
defaults:
run:
working-directory: apps/ComfyUI-vLLM-Omni
jobs:
check-base-path:
runs-on: ubuntu-latest
outputs:
exists: ${{ steps.check.outputs.exists }}
defaults:
run:
working-directory: .
steps:
- uses: actions/checkout@v4
with:
ref: ${{ github.event.pull_request.base.sha }}
fetch-depth: 1
- id: check
run: |
if [ -d "apps/ComfyUI-vLLM-Omni" ]; then
echo "exists=true" >> "$GITHUB_OUTPUT"
else
echo "exists=false" >> "$GITHUB_OUTPUT"
fi
validate:
needs: check-base-path
if: needs.check-base-path.outputs.exists == 'true'
runs-on: ubuntu-latest
steps:
- uses: comfy-org/node-diff@main
-2
View File
@@ -243,7 +243,5 @@ Dockerfile.dev
discussion
tmp_test
# Auto-generated version file (created by setuptools_scm during build)
vllm_omni/_version.py
# output files
*.wav
-1
View File
@@ -17,7 +17,6 @@ Easy, fast, and cheap omni-modality model serving for everyone
*Latest News* 🔥
- [2026/02] We released [0.16.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.16.0) - A major alignment + capability release that rebases onto **upstream vLLM v0.16.0** and significantly expands performance, distributed execution, and production readiness across **Qwen3-Omni / Qwen3-TTS**, **Bagel**, **MiMo-Audio**, **GLM-Image** and the **Diffusion (DiT) image/video stack**—while also improving platform coverage (CUDA / ROCm / NPU / XPU), CI quality, and documentation.
- [2026/02] We released [0.14.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.14.0) - This is the first **stable release** of vLLM-Omni that expands Omni’s diffusion / image-video generation and audio / TTS stack, improves distributed execution and memory efficiency, and broadens platform/backend coverage (GPU/ROCm/NPU/XPU). It also brings meaningful upgrades to serving APIs, profiling & benchmarking, and overall stability. Please check our latest [paper](https://arxiv.org/abs/2602.02204) for architecture design and performance results.
- [2026/01] We released [0.12.0rc1](https://github.com/vllm-project/vllm-omni/releases/tag/v0.12.0rc1) - a major RC milestone focused on maturing the diffusion stack, strengthening OpenAI-compatible serving, expanding omni-model coverage, and improving stability across platforms (GPU/NPU/ROCm).
- [2025/11] vLLM community officially released [vllm-project/vllm-omni](https://github.com/vllm-project/vllm-omni) in order to support omni-modality models serving.
+21 -20
View File
@@ -47,12 +47,12 @@ If no, check your shell running the ComfyUI process. There may be some error mes
This extension offers the following nodes based on the output modalities (at **ComfyUI sidebar -> Node Library**):
- **Generate Image** for text-to-image and image-to-image tasks
- **Multimodality Understanding** for multimodality-to-text and multimodality-to-audio tasks
- **Multimodality Comprehension** for multimodality-to-text and multimodality-to-audio tasks
- **TTS** and **TTS Voice Clone** for TTS tasks
This extension also offers example workflows (at **ComfyUI sidebar -> Templates -> vLLM-Omni**)
> [!NOTE]
> [!INFO]
> The node UI and feature designs are intended to match vLLM-Omni online serving interfaces. It cannot offer more than what the interfaces support.
To build a simple workflow yourself,
@@ -65,16 +65,28 @@ To build a simple workflow yourself,
- For some multi-stage models like BAGEL, [only one stage's sampling parameters are exposed and tunable via vLLM-Omni's online serving API](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/examples/online_serving/bagel/). Thus, these models are treated as single-stage ones. Please check the vLLM-Omni documentation on how to correctly set each model's sampling parameters.
- For multi-stage models where all stages are either autoregression or diffusion, you can also connect only a single Sampling Params node, indicating that this set of sampling parameters will be used for all stages.
**The following features are tested**:
- Single-node workflows for
- Multimodal Comprehension (e.g., Qwen Omni, BAGEL)
- Text-to-Image Generation (e.g., Qwen-Image)
- Image-to-Image Generation (e.g., Qwen-Image-Edit)
- TTS (e.g., Qwen TTS, including VoiceDesign, VoiceClone, CustomVoice)
**The following features are not currently tested**. They may work or break. You are welcomed to test it out and offer comments.
- Multi-node workflow that connects multiple model services together.
## Screenshots and Examples
### Multimodal understanding (e.g., Qwen Omni series, BAGEL)
### Multimodal comprehension (e.g., Qwen Omni series, BAGEL)
(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Multimodal Understanding**)
(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Annotated Example**)
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-understanding.jpg">
<img alt="vLLM-Omni multimodal understanding" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-understanding.jpg" width=55%>
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-comprehension.jpg">
<img alt="vLLM-Omni Main Architecture" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-comprehension.jpg" width=55%>
</picture>
</p>
@@ -86,7 +98,7 @@ You can configure per-stage sampling parameters for multi-stage models.
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-multi-stage.jpg">
<img alt="vLLM-Omni multiple stages" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-multi-stage.jpg" width=55%>
<img alt="vLLM-Omni Main Architecture" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-multi-stage.jpg" width=55%>
</picture>
</p>
@@ -97,7 +109,7 @@ You can configure per-stage sampling parameters for multi-stage models.
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-image-generation.jpg">
<img alt="vLLM-Omni image generation" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-image-generation.jpg" width=55%>
<img alt="vLLM-Omni Main Architecture" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-image-generation.jpg" width=55%>
</picture>
</p>
@@ -111,24 +123,13 @@ You can configure per-stage sampling parameters for multi-stage models.
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-tts.jpg">
<img alt="vLLM-Omni TTS" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-tts.jpg" width=55%>
<img alt="vLLM-Omni Main Architecture" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-tts.jpg" width=55%>
</picture>
</p>
> [!TIP]
> There is a dedicated node for VoiceClone tasks with reference audio input. Other simple text-to-speech tasks should use the regular TTS node.
### Chaining multiple model services
(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Chaining Services**)
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-chaining-services.jpg">
<img alt="vLLM-Omni TTS" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-chaining-services.jpg" width=55%>
</picture>
</p>
## Develop
Follow the [development convention and rules of vLLM-Omni](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/).
+5 -5
View File
@@ -6,18 +6,18 @@ __all__ = [
"WEB_DIRECTORY",
]
__author__ = """vLLM-Omni Team"""
__email__ = "vllm-omni@vllm.ai"
__author__ = """Zeyu Huang"""
__email__ = "11222265+fhfuih@users.noreply.github.com"
__version__ = "0.0.1"
from .comfyui_vllm_omni.nodes import (
VLLMOmniARSampling,
VLLMOmniComprehension,
VLLMOmniDiffusionSampling,
VLLMOmniGenerateImage,
VLLMOmniQwenTTSParams,
VLLMOmniSamplingParamsList,
VLLMOmniTTS,
VLLMOmniUnderstanding,
VLLMOmniVoiceClone,
)
@@ -25,7 +25,7 @@ from .comfyui_vllm_omni.nodes import (
NODE_CLASS_MAPPINGS = {
# === Generation ===
"VLLMOmniGenerateImage": VLLMOmniGenerateImage,
"VLLMOmniUnderstanding": VLLMOmniUnderstanding,
"VLLMOmniComprehension": VLLMOmniComprehension,
"VLLMOmniTTS": VLLMOmniTTS,
"VLLMOmniVoiceClone": VLLMOmniVoiceClone,
# === Params ===
@@ -39,7 +39,7 @@ NODE_CLASS_MAPPINGS = {
NODE_DISPLAY_NAME_MAPPINGS = {
# === Generation ===
"VLLMOmniGenerateImage": "Generate Image",
"VLLMOmniUnderstanding": "Multimodality Understanding",
"VLLMOmniComprehension": "Multimodality Comprehension",
"VLLMOmniTTS": "TTS (Text to Speech)",
"VLLMOmniVoiceClone": "TTS Voice Cloning",
# === Params ===
@@ -1,4 +1,4 @@
from typing import Literal
from typing import Literal, cast
import torch
from comfy_api.input import AudioInput, VideoInput
@@ -6,12 +6,7 @@ from comfy_api.input import AudioInput, VideoInput
from .utils.api_client import VLLMOmniClient
from .utils.logger import get_logger
from .utils.models import lookup_model_spec
from .utils.types import (
AudioFormat,
AutoregressionSamplingParams,
DiffusionSamplingParams,
QwenTTSModelSpecificParams,
)
from .utils.types import AudioFormat
from .utils.validators import (
add_sampling_parameters_to_stage,
validate_model_and_sampling_params_types,
@@ -91,10 +86,7 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
# Prefer DALL-E compatible API for simple (one-stage) diffusion models
if (spec is None or spec["stages"] == ["diffusion"]) and not is_bagel:
# The number of sampling parameter groups should have been validated.
# Now, simply convert single-item list to dict.
if isinstance(sampling_params, list):
sampling_params = sampling_params[0]
sampling_params = cast(dict | None, sampling_params)
if audio is None and image is None and video is None:
# No multimodal input --- use DALL-E image generation
logger.info("Using DALL-E image generation endpoint")
@@ -141,7 +133,7 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
return (output,)
class VLLMOmniUnderstanding(_VLLMOmniGenerateBase):
class VLLMOmniComprehension(_VLLMOmniGenerateBase):
@classmethod
def INPUT_TYPES(cls):
return {
@@ -205,7 +197,7 @@ class VLLMOmniUnderstanding(_VLLMOmniGenerateBase):
(
text_response,
_,
) = await client.generate_understanding_chat_completion(
) = await client.generate_comprehension_chat_completion(
model=model,
prompt=prompt,
image=image,
@@ -229,7 +221,7 @@ class VLLMOmniUnderstanding(_VLLMOmniGenerateBase):
(
text_response,
audio,
) = await client.generate_understanding_chat_completion(
) = await client.generate_comprehension_chat_completion(
model=model,
prompt=prompt,
image=image,
@@ -295,13 +287,15 @@ class VLLMOmniTTS(_VLLMOmniGenerateBase):
logger.info("Got extra kwargs in TTS: %s", kwargs)
is_qwen_tts = "qwen3-tts" in model.lower()
if not is_qwen_tts and isinstance(model_specific_params, QwenTTSModelSpecificParams):
extra_params_type = None if model_specific_params is None else model_specific_params["type"]
if not is_qwen_tts and extra_params_type == "qwen-tts":
raise ValueError(
"You have provided Qwen-specific TTS params."
"However, the model appears to not be a Qwen TTS model (no 'Qwen3-TTS' in model name)."
)
combined_params = {**kwargs, **(model_specific_params or {})}
combined_params.pop("type", None) # Internal fields in model_specific_params
client = VLLMOmniClient(url)
@@ -358,7 +352,8 @@ class VLLMOmniVoiceClone(_VLLMOmniGenerateBase):
**kwargs,
):
is_qwen_tts = "qwen3-tts" in model.lower()
if not is_qwen_tts and isinstance(model_specific_params, QwenTTSModelSpecificParams):
extra_params_type = None if model_specific_params is None else model_specific_params["type"]
if not is_qwen_tts and extra_params_type == "qwen-tts":
raise ValueError(
"You have provided Qwen-specific TTS params."
"However, the model appears to not be a Qwen TTS model (no 'Qwen3-TTS' in model name)."
@@ -371,6 +366,7 @@ class VLLMOmniVoiceClone(_VLLMOmniGenerateBase):
**kwargs,
**(model_specific_params or {}),
}
combined_params.pop("type", None) # Internal fields in model_specific_params
client = VLLMOmniClient(url)
@@ -423,7 +419,10 @@ class VLLMOmniARSampling:
CATEGORY = "vLLM-Omni/Sampling Params"
def get_params(self, seed, **kwargs):
params = AutoregressionSamplingParams(kwargs)
params = {
"type": "autoregression", # for internal use, removed before sending the request
**kwargs,
}
if seed >= 0:
params["seed"] = seed
return (params,)
@@ -480,13 +479,6 @@ class VLLMOmniDiffusionSampling:
"tooltip": "Enable VAE slicing for reduced memory usage (slight quality trade-off)",
},
),
"vae_use_tiling": (
"BOOLEAN",
{
"default": False,
"tooltip": "Enable VAE tiling for reduced memory usage (slight quality trade-off)",
},
),
# === Put seed at last. ===
# Whenever a field named "seed" is present, ComfyUI adds another field called "control after generate"
"seed": (
@@ -507,7 +499,10 @@ class VLLMOmniDiffusionSampling:
CATEGORY = "vLLM-Omni/Sampling Params"
def get_params(self, seed, **kwargs):
params = DiffusionSamplingParams(kwargs)
params = {
"type": "diffusion", # for internal use, removed before sending the request
**kwargs,
}
if seed >= 0:
params["seed"] = seed
return (params,)
@@ -571,4 +566,4 @@ class VLLMOmniQwenTTSParams:
CATEGORY = "vLLM-Omni/TTS Params"
def get_params(self, **kwargs):
return (QwenTTSModelSpecificParams(kwargs),)
return ({"type": "qwen-tts", **kwargs},)
@@ -57,7 +57,16 @@ class VLLMOmniClient:
if negative_prompt:
payload["negative_prompt"] = negative_prompt
if sampling_params is not None:
payload.update(sampling_params)
# Only select specific sampling params
for k in (
"n",
"num_inference_steps",
"guidance_scale",
"true_cfg_scale",
"vae_use_slicing",
):
if k in sampling_params and sampling_params[k] is not None:
payload[k] = sampling_params[k]
logger.debug("img gen payload: %s", payload)
url = self.base_url + "/images/generations"
@@ -129,8 +138,10 @@ class VLLMOmniClient:
if negative_prompt:
form.add_field("negative_prompt", negative_prompt)
if sampling_params is not None:
for k, v in sampling_params.items():
form.add_field(k, str(v))
# Only select specific sampling params
for k in ("n", "num_inference_steps", "guidance_scale", "true_cfg_scale"):
if k in sampling_params and sampling_params[k] is not None:
form.add_field(k, str(sampling_params[k]))
if mask is not None:
mask_filename = "mask.png"
form.add_field(
@@ -206,7 +217,7 @@ class VLLMOmniClient:
return torch.stack(image_tensors, dim=0)
async def generate_understanding_chat_completion(
async def generate_comprehension_chat_completion(
self,
*,
model: str,
@@ -416,28 +427,30 @@ class VLLMOmniClient:
message_content.append({"type": "video_url", "video_url": {"url": video_to_base64(video)}})
messages = [{"role": "user", "content": message_content}]
payload: dict[str, Any] = {"messages": messages, "model": model}
if modalities:
payload["modalities"] = modalities
combined_extra_body: dict[str, Any] = {}
if sampling_params is not None:
spec, _ = lookup_model_spec(model)
is_single_sampling_param = isinstance(sampling_params, dict) or len(sampling_params) == 1
# Exclude internal key
if isinstance(sampling_params, dict):
sampling_params = {k: v for k, v in sampling_params.items() if k != "type"}
else:
sampling_params = [{k: v for k, v in sp.items() if k != "type"} for sp in sampling_params]
if (spec is None and is_single_sampling_param) or (spec is not None and spec["stages"] == ["diffusion"]):
# Diffusion format: extra_body directly contains sampling params.
# Validation has already taken care of matching sampling params' types and length. Safe to take [0].
# Validation should have taken care of matching sampling params' types.
# * Use this mode if the model is a simple one-stage diffusion model.
# * Fallback to this mode if model is not registered and a single sampling param is provided.
sampling_params = sampling_params if isinstance(sampling_params, dict) else sampling_params[0]
combined_extra_body: dict[str, Any] = sampling_params.copy()
combined_extra_body: dict[str, Any] = {**sampling_params}
if "n" in combined_extra_body:
combined_extra_body["num_outputs_per_prompt"] = combined_extra_body.pop("n")
combined_extra_body["num_outputs_per_prompt"] = combined_extra_body["n"]
del combined_extra_body["n"]
else:
# AR format: the payload has a sampling_params_list field, containing a list.
sampling_params_list = sampling_params if isinstance(sampling_params, list) else [sampling_params]
payload["sampling_params_list"] = sampling_params_list
# Use AR style payload, extra_body has a sampling_params_list field
combined_extra_body: dict[str, Any] = {"sampling_params_list": sampling_params}
if negative_prompt:
combined_extra_body["negative_prompt"] = negative_prompt
@@ -445,11 +458,12 @@ class VLLMOmniClient:
if extra_body:
combined_extra_body.update(extra_body)
# Add extra_body only if it has any content.
payload: dict[str, Any] = {"messages": messages, "model": model}
if combined_extra_body:
payload["extra_body"] = combined_extra_body
if modalities:
payload["modalities"] = modalities
# Place to inject any model-specific payload adjustment
spec, _ = lookup_model_spec(model)
if spec:
preprocessor = spec.get("payload_preprocessor", None)
@@ -11,7 +11,8 @@ def _bagel_payload_preprocessor(payload: dict) -> dict:
content["text"] = "<|im_start|>" + content["text"] + "<|im_end|>"
except (KeyError, TypeError):
raise RuntimeError("Internal Error: malformatted BAGEL payload")
return payload
extra_body = payload.pop("extra_body", {})
return {**payload, **extra_body}
def _qwen25_payload_preprocessor(payload: dict) -> dict:
@@ -36,7 +37,7 @@ _MODEL_PIPELINE_SPECS: dict[str, Spec] = {
],
"modes": [
{
"mode": ModelMode.UNDERSTANDING,
"mode": ModelMode.COMPREHENSION,
"input_modalities": [Modality.TEXT, Modality.IMAGE],
}
],
@@ -47,7 +48,7 @@ _MODEL_PIPELINE_SPECS: dict[str, Spec] = {
"payload_preprocessor": _qwen25_payload_preprocessor,
"modes": [
{
"mode": ModelMode.UNDERSTANDING,
"mode": ModelMode.COMPREHENSION,
"input_modalities": [
Modality.TEXT,
Modality.IMAGE,
@@ -61,7 +62,7 @@ _MODEL_PIPELINE_SPECS: dict[str, Spec] = {
"stages": ["autoregression", "autoregression", "autoregression"],
"modes": [
{
"mode": ModelMode.UNDERSTANDING,
"mode": ModelMode.COMPREHENSION,
"input_modalities": [
Modality.TEXT,
Modality.IMAGE,
@@ -11,23 +11,11 @@ from typing import (
AudioFormat: TypeAlias = Literal["mp3", "opus", "aac", "flac", "wav", "pcm"]
class AutoregressionSamplingParams(dict):
pass
class DiffusionSamplingParams(dict):
pass
class QwenTTSModelSpecificParams(dict):
pass
class ModelMode(Enum):
IMAGE_GENERATION = auto()
VIDEO_GENERATION = auto()
AUDIO_GENERATION = auto()
UNDERSTANDING = auto()
COMPREHENSION = auto()
class Modality(Enum):
@@ -1,6 +1,5 @@
from .logger import get_logger
from .models import lookup_model_spec
from .types import AutoregressionSamplingParams, DiffusionSamplingParams
logger = get_logger(__name__)
@@ -32,20 +31,24 @@ def validate_model_and_sampling_params_types(
)
# Check that each stage's type match
for i, sp in enumerate(sampling_param_list):
if not _check_sampling_param_matches_stage(sp, stages[i]):
if "type" not in sp:
raise RuntimeError("Internal error: unknown sampling parameter type")
if sp["type"] != stages[i]:
raise ValueError(
f"Sampling parameter type ({sp.__class__.__name__}) does not match "
f"Sampling parameter type ({sp['type']}) does not match "
f"stage type ({stages[i]}) at index {i} for model {model_name}."
)
elif isinstance(sampling_param_list, dict):
if "type" not in sampling_param_list:
raise RuntimeError("Internal error: unknown sampling parameter type")
# Check that the provided single sampling param matches all stages
for i, stage in enumerate(stages):
if not _check_sampling_param_matches_stage(sampling_param_list, stage):
raise ValueError(
f"Provided single sampling parameter type ({sampling_param_list.__class__.__name__}) must match "
f"the types of all stages of the model. "
f"However, stage {i} of model {model_name} is of type {stage}."
)
elif any(stage != sampling_param_list["type"] for stage in stages):
raise ValueError(
f"When passing a single sampling parameter node, all stages of the model must match "
f"the provided sampling parameter's type. "
f"However, the stages of model {model_name} are: {stages}. "
f"The provided sampling parameter is {sampling_param_list['type']}"
)
def add_sampling_parameters_to_stage(
@@ -79,11 +82,3 @@ def add_sampling_parameters_to_stage(
sampling_param_list[i].update(params_to_add)
return sampling_param_list
def _check_sampling_param_matches_stage(sampling_param: dict, stage_type: str) -> bool:
if stage_type == "autoregression":
return isinstance(sampling_param, AutoregressionSamplingParams)
if stage_type == "diffusion":
return isinstance(sampling_param, DiffusionSamplingParams)
raise RuntimeError(f"Internal error: unknown stage type {stage_type}.")
Binary file not shown.

Before

Width:  |  Height:  |  Size: 260 KiB

@@ -0,0 +1 @@
{"id":"1c99f525-0a37-45ba-a28a-7df7c3af66b4","revision":0,"last_node_id":12,"last_link_id":14,"nodes":[{"id":1,"type":"VLLMOmniComprehension","pos":[1191.2177053682556,144.66829928181377],"size":[400,268],"flags":{},"order":8,"mode":0,"inputs":[{"localized_name":"image","name":"image","shape":7,"type":"IMAGE","link":1},{"localized_name":"video","name":"video","shape":7,"type":"VIDEO","link":2},{"localized_name":"audio","name":"audio","shape":7,"type":"AUDIO","link":3},{"localized_name":"sampling_params","name":"sampling_params","shape":7,"type":"SAMPLING_PARAMS","link":14},{"localized_name":"url","name":"url","type":"STRING","widget":{"name":"url"},"link":null},{"localized_name":"model","name":"model","type":"STRING","widget":{"name":"model"},"link":null},{"localized_name":"prompt","name":"prompt","type":"STRING","widget":{"name":"prompt"},"link":null},{"localized_name":"output_text","name":"output_text","type":"BOOLEAN","widget":{"name":"output_text"},"link":null},{"localized_name":"output_audio","name":"output_audio","type":"BOOLEAN","widget":{"name":"output_audio"},"link":null},{"localized_name":"use_audio_in_video","name":"use_audio_in_video","type":"BOOLEAN","widget":{"name":"use_audio_in_video"},"link":null}],"outputs":[{"localized_name":"text_response","name":"text_response","type":"STRING","links":[8]},{"localized_name":"audio_response","name":"audio_response","type":"AUDIO","links":[9]}],"properties":{"Node name for S&R":"VLLMOmniComprehension"},"widgets_values":["http://localhost:8000/v1","Qwen/Qwen2.5-Omni-7B","",true,true,true]},{"id":3,"type":"LoadVideo","pos":[729.5984141255855,-198.631920454299],"size":[282.798828125,233.0743408203125],"flags":{},"order":0,"mode":0,"inputs":[{"localized_name":"file","name":"file","type":"COMBO","widget":{"name":"file"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"IMAGEUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"VIDEO","name":"VIDEO","type":"VIDEO","links":[2]}],"properties":{"Node name for S&R":"LoadVideo"},"widgets_values":["draw.mp4","image"]},{"id":4,"type":"LoadAudio","pos":[729.8037086965753,99.86963519703949],"size":[282.798828125,136],"flags":{},"order":1,"mode":0,"inputs":[{"localized_name":"audio","name":"audio","type":"COMBO","widget":{"name":"audio"},"link":null},{"localized_name":"audioUI","name":"audioUI","type":"AUDIO_UI","widget":{"name":"audioUI"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"AUDIOUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"AUDIO","name":"AUDIO","type":"AUDIO","links":[3]}],"properties":{"Node name for S&R":"LoadAudio"},"widgets_values":["Megan-Fox.mp3",null,null]},{"id":5,"type":"VLLMOmniARSampling","pos":[510.3517536642828,658.073751009259],"size":[270,178],"flags":{},"order":2,"mode":0,"inputs":[{"localized_name":"max_tokens","name":"max_tokens","type":"INT","widget":{"name":"max_tokens"},"link":null},{"localized_name":"temperature","name":"temperature","type":"FLOAT","widget":{"name":"temperature"},"link":null},{"localized_name":"top_p","name":"top_p","type":"FLOAT","widget":{"name":"top_p"},"link":null},{"localized_name":"repetition_penalty","name":"repetition_penalty","type":"FLOAT","widget":{"name":"repetition_penalty"},"link":null},{"localized_name":"seed","name":"seed","type":"INT","widget":{"name":"seed"},"link":null}],"outputs":[{"localized_name":"AR sampling params","name":"AR sampling params","type":"SAMPLING_PARAMS","links":[12]}],"properties":{"Node name for S&R":"VLLMOmniARSampling"},"widgets_values":[100,1,1,1,-1,"randomize"]},{"id":7,"type":"VLLMOmniARSampling","pos":[503.33235181647115,419.34158016181806],"size":[270,178],"flags":{},"order":3,"mode":0,"inputs":[{"localized_name":"max_tokens","name":"max_tokens","type":"INT","widget":{"name":"max_tokens"},"link":null},{"localized_name":"temperature","name":"temperature","type":"FLOAT","widget":{"name":"temperature"},"link":null},{"localized_name":"top_p","name":"top_p","type":"FLOAT","widget":{"name":"top_p"},"link":null},{"localized_name":"repetition_penalty","name":"repetition_penalty","type":"FLOAT","widget":{"name":"repetition_penalty"},"link":null},{"localized_name":"seed","name":"seed","type":"INT","widget":{"name":"seed"},"link":null}],"outputs":[{"localized_name":"AR sampling params","name":"AR sampling params","type":"SAMPLING_PARAMS","links":[5,13]}],"properties":{"Node name for S&R":"VLLMOmniARSampling"},"widgets_values":[100,1,1,1,-1,"randomize"]},{"id":8,"type":"VLLMOmniSamplingParamsList","pos":[820.6056617389042,426.38372037182273],"size":[263.066015625,66],"flags":{},"order":7,"mode":0,"inputs":[{"localized_name":"param1","name":"param1","type":"SAMPLING_PARAMS","link":5},{"localized_name":"param2","name":"param2","shape":7,"type":"SAMPLING_PARAMS","link":12},{"localized_name":"param3","name":"param3","shape":7,"type":"SAMPLING_PARAMS","link":13}],"outputs":[{"localized_name":"param list","name":"param list"Line truncated
@@ -1,552 +0,0 @@
{
"id": "6643e5fd-fa2a-4f25-935a-483173a8097c",
"revision": 0,
"last_node_id": 8,
"last_link_id": 6,
"nodes": [
{
"id": 1,
"type": "PreviewImage",
"pos": [
1446.2005205859512,
-316.7359686049902
],
"size": [
305.6628685610174,
307.0336884172169
],
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{
"localized_name": "images",
"name": "images",
"type": "IMAGE",
"link": 1
}
],
"outputs": [],
"properties": {
"Node name for S&R": "PreviewImage"
},
"widgets_values": []
},
{
"id": 8,
"type": "PreviewImage",
"pos": [
1010.1848219273943,
-311.54216706114363
],
"size": [
305.6628685610174,
307.0336884172169
],
"flags": {},
"order": 4,
"mode": 0,
"inputs": [
{
"localized_name": "images",
"name": "images",
"type": "IMAGE",
"link": 6
}
],
"outputs": [],
"properties": {
"Node name for S&R": "PreviewImage"
},
"widgets_values": []
},
{
"id": 7,
"type": "VLLMOmniDiffusionSampling",
"pos": [
234.61958293819296,
9.405681270043619
],
"size": [
284.205078125,
226
],
"flags": {},
"order": 0,
"mode": 0,
"inputs": [
{
"localized_name": "n",
"name": "n",
"type": "INT",
"widget": {
"name": "n"
},
"link": null
},
{
"localized_name": "num_inference_steps",
"name": "num_inference_steps",
"type": "INT",
"widget": {
"name": "num_inference_steps"
},
"link": null
},
{
"localized_name": "guidance_scale",
"name": "guidance_scale",
"type": "FLOAT",
"widget": {
"name": "guidance_scale"
},
"link": null
},
{
"localized_name": "true_cfg_scale",
"name": "true_cfg_scale",
"type": "FLOAT",
"widget": {
"name": "true_cfg_scale"
},
"link": null
},
{
"localized_name": "vae_use_slicing",
"name": "vae_use_slicing",
"type": "BOOLEAN",
"widget": {
"name": "vae_use_slicing"
},
"link": null
},
{
"localized_name": "vae_use_tiling",
"name": "vae_use_tiling",
"type": "BOOLEAN",
"widget": {
"name": "vae_use_tiling"
},
"link": null
},
{
"localized_name": "seed",
"name": "seed",
"type": "INT",
"widget": {
"name": "seed"
},
"link": null
}
],
"outputs": [
{
"localized_name": "diffusion sampling params",
"name": "diffusion sampling params",
"type": "SAMPLING_PARAMS",
"links": [
5
]
}
],
"properties": {
"Node name for S&R": "VLLMOmniDiffusionSampling"
},
"widgets_values": [
1,
50,
1,
1,
false,
false,
1525,
"randomize"
]
},
{
"id": 4,
"type": "VLLMOmniDiffusionSampling",
"pos": [
666.8380026154548,
268.86271068330126
],
"size": [
284.205078125,
226
],
"flags": {},
"order": 1,
"mode": 0,
"inputs": [
{
"localized_name": "n",
"name": "n",
"type": "INT",
"widget": {
"name": "n"
},
"link": null
},
{
"localized_name": "num_inference_steps",
"name": "num_inference_steps",
"type": "INT",
"widget": {
"name": "num_inference_steps"
},
"link": null
},
{
"localized_name": "guidance_scale",
"name": "guidance_scale",
"type": "FLOAT",
"widget": {
"name": "guidance_scale"
},
"link": null
},
{
"localized_name": "true_cfg_scale",
"name": "true_cfg_scale",
"type": "FLOAT",
"widget": {
"name": "true_cfg_scale"
},
"link": null
},
{
"localized_name": "vae_use_slicing",
"name": "vae_use_slicing",
"type": "BOOLEAN",
"widget": {
"name": "vae_use_slicing"
},
"link": null
},
{
"localized_name": "vae_use_tiling",
"name": "vae_use_tiling",
"type": "BOOLEAN",
"widget": {
"name": "vae_use_tiling"
},
"link": null
},
{
"localized_name": "seed",
"name": "seed",
"type": "INT",
"widget": {
"name": "seed"
},
"link": null
}
],
"outputs": [
{
"localized_name": "diffusion sampling params",
"name": "diffusion sampling params",
"type": "SAMPLING_PARAMS",
"links": [
2
]
}
],
"properties": {
"Node name for S&R": "VLLMOmniDiffusionSampling"
},
"widgets_values": [
4,
50,
7,
1,
false,
false,
42,
"fixed"
]
},
{
"id": 5,
"type": "VLLMOmniGenerateImage",
"pos": [
984.723613585788,
63.376900027553276
],
"size": [
416.56628685610167,
372.1662621294205
],
"flags": {},
"order": 3,
"mode": 0,
"inputs": [
{
"localized_name": "image",
"name": "image",
"shape": 7,
"type": "IMAGE",
"link": 4
},
{
"localized_name": "mask",
"name": "mask",
"shape": 7,
"type": "MASK",
"link": null
},
{
"localized_name": "sampling_params",
"name": "sampling_params",
"shape": 7,
"type": "SAMPLING_PARAMS",
"link": 2
},
{
"localized_name": "url",
"name": "url",
"type": "STRING",
"widget": {
"name": "url"
},
"link": null
},
{
"localized_name": "model",
"name": "model",
"type": "STRING",
"widget": {
"name": "model"
},
"link": null
},
{
"localized_name": "prompt",
"name": "prompt",
"type": "STRING",
"widget": {
"name": "prompt"
},
"link": null
},
{
"localized_name": "negative_prompt",
"name": "negative_prompt",
"type": "STRING",
"widget": {
"name": "negative_prompt"
},
"link": null
},
{
"localized_name": "width",
"name": "width",
"type": "INT",
"widget": {
"name": "width"
},
"link": null
},
{
"localized_name": "height",
"name": "height",
"type": "INT",
"widget": {
"name": "height"
},
"link": null
}
],
"outputs": [
{
"localized_name": "image",
"name": "image",
"type": "IMAGE",
"links": [
1
]
}
],
"properties": {
"Node name for S&R": "VLLMOmniGenerateImage"
},
"widgets_values": [
"http://localhost:8001/v1",
"/home/models/Qwen/Qwen-Image-Edit",
"A high-quality, high contrast, stylized portrait of the object in the uploaded reference image. Pop art and doodle style, with abundant scribbling patterns, such as a teal crown, orange lightning bolts, colorful handwritten scripts.",
"Realistic",
800,
800
]
},
{
"id": 6,
"type": "VLLMOmniGenerateImage",
"pos": [
541.2702717429818,
-160.70392744708548
],
"size": [
416.56628685610167,
372.1662621294205
],
"flags": {},
"order": 2,
"mode": 0,
"inputs": [
{
"localized_name": "image",
"name": "image",
"shape": 7,
"type": "IMAGE",
"link": null
},
{
"localized_name": "mask",
"name": "mask",
"shape": 7,
"type": "MASK",
"link": null
},
{
"localized_name": "sampling_params",
"name": "sampling_params",
"shape": 7,
"type": "SAMPLING_PARAMS",
"link": 5
},
{
"localized_name": "url",
"name": "url",
"type": "STRING",
"widget": {
"name": "url"
},
"link": null
},
{
"localized_name": "model",
"name": "model",
"type": "STRING",
"widget": {
"name": "model"
},
"link": null
},
{
"localized_name": "prompt",
"name": "prompt",
"type": "STRING",
"widget": {
"name": "prompt"
},
"link": null
},
{
"localized_name": "negative_prompt",
"name": "negative_prompt",
"type": "STRING",
"widget": {
"name": "negative_prompt"
},
"link": null
},
{
"localized_name": "width",
"name": "width",
"type": "INT",
"widget": {
"name": "width"
},
"link": null
},
{
"localized_name": "height",
"name": "height",
"type": "INT",
"widget": {
"name": "height"
},
"link": null
}
],
"outputs": [
{
"localized_name": "image",
"name": "image",
"type": "IMAGE",
"links": [
4,
6
]
}
],
"properties": {
"Node name for S&R": "VLLMOmniGenerateImage"
},
"widgets_values": [
"http://localhost:8000/v1",
"/home/models/Tongyi-MAI/Z-Image-Turbo",
"A headshot of a cute Siamese kitty. Blurred background due to wide aperture. Close-up look. Realistic.",
"Cartoonish.",
800,
800
]
}
],
"links": [
[
1,
5,
0,
1,
0,
"IMAGE"
],
[
2,
4,
0,
5,
2,
"SAMPLING_PARAMS"
],
[
4,
6,
0,
5,
0,
"IMAGE"
],
[
5,
7,
0,
6,
2,
"SAMPLING_PARAMS"
],
[
6,
6,
0,
8,
0,
"IMAGE"
]
],
"groups": [],
"config": {},
"extra": {
"workflowRendererVersion": "LG",
"ds": {
"scale": 1.0426432563169903,
"offset": [
-90.4071374799346,
575.2948920069742
]
}
},
"version": 0.4
}
@@ -1 +1 @@
{"id":"91f75acc-8040-40f6-865a-2e8a7cfd6672","revision":0,"last_node_id":11,"last_link_id":23,"nodes":[{"id":3,"type":"PreviewImage","pos":[1281.8455167767304,-69.02638461454333],"size":[305.6628685610174,307.0336884172169],"flags":{},"order":4,"mode":0,"inputs":[{"localized_name":"图像","name":"images","type":"IMAGE","link":22}],"outputs":[],"properties":{"Node name for S&R":"PreviewImage"},"widgets_values":[]},{"id":11,"type":"VLLMOmniGenerateImage","pos":[816.0962515873642,-137.51112584854445],"size":[400,278],"flags":{},"order":3,"mode":0,"inputs":[{"localized_name":"image","name":"image","shape":7,"type":"IMAGE","link":21},{"localized_name":"mask","name":"mask","shape":7,"type":"MASK","link":null},{"localized_name":"sampling_params","name":"sampling_params","shape":7,"type":"SAMPLING_PARAMS","link":23},{"localized_name":"url","name":"url","type":"STRING","widget":{"name":"url"},"link":null},{"localized_name":"model","name":"model","type":"STRING","widget":{"name":"model"},"link":null},{"localized_name":"prompt","name":"prompt","type":"STRING","widget":{"name":"prompt"},"link":null},{"localized_name":"negative_prompt","name":"negative_prompt","type":"STRING","widget":{"name":"negative_prompt"},"link":null},{"localized_name":"width","name":"width","type":"INT","widget":{"name":"width"},"link":null},{"localized_name":"height","name":"height","type":"INT","widget":{"name":"height"},"link":null}],"outputs":[{"localized_name":"image","name":"image","type":"IMAGE","links":[22]}],"properties":{"Node name for S&R":"VLLMOmniGenerateImage"},"widgets_values":["http://localhost:8000/v1","Qwen/Qwen-Image-Edit","Put this figure in a realistic mountain view","",512,512]},{"id":10,"type":"MarkdownNote","pos":[227.99306819575278,-231.24143306069843],"size":[240.3326478922474,136.3505820791881],"flags":{},"order":1,"mode":0,"inputs":[],"outputs":[],"title":"Note: Task and Input","properties":{},"widgets_values":["vLLM-Omni nodes are categorized based on the output modality. The \"Generate Image\" node supports both text-to-image generation or image-to-image generation (a.k.a. image editing). The node will route to the correct endpoint depending on whether an input image is present or not."],"color":"#432","bgcolor":"#000"},{"id":4,"type":"LoadImage","pos":[496.31859627609606,-229.71277089860084],"size":[270,314],"flags":{},"order":2,"mode":0,"inputs":[{"localized_name":"图像","name":"image","type":"COMBO","widget":{"name":"image"},"link":null},{"localized_name":"选择文件上传","name":"upload","type":"IMAGEUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"图像","name":"IMAGE","type":"IMAGE","links":[21]},{"localized_name":"遮罩","name":"MASK","type":"MASK","links":null}],"properties":{"Node name for S&R":"LoadImage"},"widgets_values":["example.png","image"]},{"id":8,"type":"VLLMOmniDiffusionSampling","pos":[478.59266934006774,183.67711984955648],"size":[284.205078125,226],"flags":{},"order":0,"mode":0,"inputs":[{"localized_name":"n","name":"n","type":"INT","widget":{"name":"n"},"link":null},{"localized_name":"num_inference_steps","name":"num_inference_steps","type":"INT","widget":{"name":"num_inference_steps"},"link":null},{"localized_name":"guidance_scale","name":"guidance_scale","type":"FLOAT","widget":{"name":"guidance_scale"},"link":null},{"localized_name":"true_cfg_scale","name":"true_cfg_scale","type":"FLOAT","widget":{"name":"true_cfg_scale"},"link":null},{"localized_name":"vae_use_slicing","name":"vae_use_slicing","type":"BOOLEAN","widget":{"name":"vae_use_slicing"},"link":null},{"localized_name":"vae_use_tiling","name":"vae_use_tiling","type":"BOOLEAN","widget":{"name":"vae_use_tiling"},"link":null},{"localized_name":"seed","name":"seed","type":"INT","widget":{"name":"seed"},"link":null}],"outputs":[{"localized_name":"diffusion sampling params","name":"diffusion sampling params","type":"SAMPLING_PARAMS","links":[23]}],"properties":{"Node name for S&R":"VLLMOmniDiffusionSampling"},"widgets_values":[4,50,1,1,false,false,42,"randomize"]}],"links":[[21,4,0,11,0,"IMAGE"],[22,11,0,3,0,"IMAGE"],[23,8,0,11,2,"SAMPLING_PARAMS"]],"groups":[{"id":1,"title":"Input","bounding":[213.1706010087147,-313.20095750667554,560.8246471437027,407.4246532472182],"color":"#3f789e","font_size":24,"flags":{}}],"config":{},"extra":{"workflowRendererVersion":"LG","ds":{"scale":1.1469075819486894,"offset":[-55.18047513099182,220.2553505195962]}},"version":0.4}
{"id":"91f75acc-8040-40f6-865a-2e8a7cfd6672","revision":0,"last_node_id":11,"last_link_id":23,"nodes":[{"id":3,"type":"PreviewImage","pos":[1281.8455167767304,-69.02638461454333],"size":[305.6628685610174,307.0336884172169],"flags":{},"order":4,"mode":0,"inputs":[{"localized_name":"images","name":"images","type":"IMAGE","link":22}],"outputs":[],"properties":{"Node name for S&R":"PreviewImage"},"widgets_values":[]},{"id":11,"type":"VLLMOmniGenerateImage","pos":[816.0962515873642,-137.51112584854445],"size":[400,278],"flags":{},"order":3,"mode":0,"inputs":[{"localized_name":"image","name":"image","shape":7,"type":"IMAGE","link":21},{"localized_name":"mask","name":"mask","shape":7,"type":"MASK","link":null},{"localized_name":"sampling_params","name":"sampling_params","shape":7,"type":"SAMPLING_PARAMS","link":23},{"localized_name":"url","name":"url","type":"STRING","widget":{"name":"url"},"link":null},{"localized_name":"model","name":"model","type":"STRING","widget":{"name":"model"},"link":null},{"localized_name":"prompt","name":"prompt","type":"STRING","widget":{"name":"prompt"},"link":null},{"localized_name":"negative_prompt","name":"negative_prompt","type":"STRING","widget":{"name":"negative_prompt"},"link":null},{"localized_name":"width","name":"width","type":"INT","widget":{"name":"width"},"link":null},{"localized_name":"height","name":"height","type":"INT","widget":{"name":"height"},"link":null}],"outputs":[{"localized_name":"image","name":"image","type":"IMAGE","links":[22]}],"properties":{"Node name for S&R":"VLLMOmniGenerateImage"},"widgets_values":["http://localhost:8000/v1","Qwen/Qwen-Image-Edit","Put this figure in a realistic mountain view","",512,512]},{"id":8,"type":"VLLMOmniDiffusionSampling","pos":[478.59266934006774,183.67711984955648],"size":[284.205078125,202],"flags":{},"order":0,"mode":0,"inputs":[{"localized_name":"n","name":"n","type":"INT","widget":{"name":"n"},"link":null},{"localized_name":"num_inference_steps","name":"num_inference_steps","type":"INT","widget":{"name":"num_inference_steps"},"link":null},{"localized_name":"guidance_scale","name":"guidance_scale","type":"FLOAT","widget":{"name":"guidance_scale"},"link":null},{"localized_name":"true_cfg_scale","name":"true_cfg_scale","type":"FLOAT","widget":{"name":"true_cfg_scale"},"link":null},{"localized_name":"vae_use_slicing","name":"vae_use_slicing","type":"BOOLEAN","widget":{"name":"vae_use_slicing"},"link":null},{"localized_name":"seed","name":"seed","type":"INT","widget":{"name":"seed"},"link":null}],"outputs":[{"localized_name":"diffusion sampling params","name":"diffusion sampling params","type":"SAMPLING_PARAMS","links":[23]}],"properties":{"Node name for S&R":"VLLMOmniDiffusionSampling"},"widgets_values":[4,50,1,1,false,42,"fixed"]},{"id":10,"type":"MarkdownNote","pos":[227.99306819575278,-231.24143306069843],"size":[240.3326478922474,136.3505820791881],"flags":{},"order":2,"mode":0,"inputs":[],"outputs":[],"title":"Note: Task and Input","properties":{},"widgets_values":["vLLM-Omni nodes are categorized based on the output modality. The \"Generate Image\" node supports both text-to-image generation or image-to-image generation (a.k.a. image editing). The node will route to the correct endpoint depending on whether an input image is present or not."],"color":"#432","bgcolor":"#000"},{"id":4,"type":"LoadImage","pos":[496.31859627609606,-229.71277089860084],"size":[270,314],"flags":{},"order":1,"mode":0,"inputs":[{"localized_name":"image","name":"image","type":"COMBO","widget":{"name":"image"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"IMAGEUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"IMAGE","name":"IMAGE","type":"IMAGE","links":[21]},{"localized_name":"MASK","name":"MASK","type":"MASK","links":null}],"properties":{"Node name for S&R":"LoadImage"},"widgets_values":["example.png","image"]}],"links":[[21,4,0,11,0,"IMAGE"],[22,11,0,3,0,"IMAGE"],[23,8,0,11,2,"SAMPLING_PARAMS"]],"groups":[{"id":1,"title":"Input","bounding":[213.1706010087147,-313.20095750667554,560.8246471437027,407.4246532472182],"color":"#3f789e","font_size":24,"flags":{}}],"config":{},"extra":{"workflowRendererVersion":"LG","ds":{"scale":1.1469075819486894,"offset":[276.8427259980847,549.4885226647582]}},"version":0.4}
@@ -1,761 +0,0 @@
{
"id": "1c99f525-0a37-45ba-a28a-7df7c3af66b4",
"revision": 0,
"last_node_id": 12,
"last_link_id": 14,
"nodes": [
{
"id": 1,
"type": "VLLMOmniUnderstanding",
"pos": [
1191.2177053682556,
144.66829928181377
],
"size": [
400,
268
],
"flags": {},
"order": 8,
"mode": 0,
"inputs": [
{
"localized_name": "image",
"name": "image",
"shape": 7,
"type": "IMAGE",
"link": 1
},
{
"localized_name": "video",
"name": "video",
"shape": 7,
"type": "VIDEO",
"link": 2
},
{
"localized_name": "audio",
"name": "audio",
"shape": 7,
"type": "AUDIO",
"link": 3
},
{
"localized_name": "sampling_params",
"name": "sampling_params",
"shape": 7,
"type": "SAMPLING_PARAMS",
"link": 14
},
{
"localized_name": "url",
"name": "url",
"type": "STRING",
"widget": {
"name": "url"
},
"link": null
},
{
"localized_name": "model",
"name": "model",
"type": "STRING",
"widget": {
"name": "model"
},
"link": null
},
{
"localized_name": "prompt",
"name": "prompt",
"type": "STRING",
"widget": {
"name": "prompt"
},
"link": null
},
{
"localized_name": "output_text",
"name": "output_text",
"type": "BOOLEAN",
"widget": {
"name": "output_text"
},
"link": null
},
{
"localized_name": "output_audio",
"name": "output_audio",
"type": "BOOLEAN",
"widget": {
"name": "output_audio"
},
"link": null
},
{
"localized_name": "use_audio_in_video",
"name": "use_audio_in_video",
"type": "BOOLEAN",
"widget": {
"name": "use_audio_in_video"
},
"link": null
}
],
"outputs": [
{
"localized_name": "text_response",
"name": "text_response",
"type": "STRING",
"links": [
8
]
},
{
"localized_name": "audio_response",
"name": "audio_response",
"type": "AUDIO",
"links": [
9
]
}
],
"properties": {
"Node name for S&R": "VLLMOmniUnderstanding"
},
"widgets_values": [
"http://localhost:8000/v1",
"Qwen/Qwen2.5-Omni-7B",
"",
true,
true,
true
]
},
{
"id": 3,
"type": "LoadVideo",
"pos": [
729.5984141255855,
-198.631920454299
],
"size": [
282.798828125,
233.0743408203125
],
"flags": {},
"order": 0,
"mode": 0,
"inputs": [
{
"localized_name": "file",
"name": "file",
"type": "COMBO",
"widget": {
"name": "file"
},
"link": null
},
{
"localized_name": "choose file to upload",
"name": "upload",
"type": "IMAGEUPLOAD",
"widget": {
"name": "upload"
},
"link": null
}
],
"outputs": [
{
"localized_name": "VIDEO",
"name": "VIDEO",
"type": "VIDEO",
"links": [
2
]
}
],
"properties": {
"Node name for S&R": "LoadVideo"
},
"widgets_values": [
"draw.mp4",
"image"
]
},
{
"id": 4,
"type": "LoadAudio",
"pos": [
729.8037086965753,
99.86963519703949
],
"size": [
282.798828125,
136
],
"flags": {},
"order": 1,
"mode": 0,
"inputs": [
{
"localized_name": "audio",
"name": "audio",
"type": "COMBO",
"widget": {
"name": "audio"
},
"link": null
},
{
"localized_name": "audioUI",
"name": "audioUI",
"type": "AUDIO_UI",
"widget": {
"name": "audioUI"
},
"link": null
},
{
"localized_name": "choose file to upload",
"name": "upload",
"type": "AUDIOUPLOAD",
"widget": {
"name": "upload"
},
"link": null
}
],
"outputs": [
{
"localized_name": "AUDIO",
"name": "AUDIO",
"type": "AUDIO",
"links": [
3
]
}
],
"properties": {
"Node name for S&R": "LoadAudio"
},
"widgets_values": [
"Megan-Fox.mp3",
null,
null
]
},
{
"id": 5,
"type": "VLLMOmniARSampling",
"pos": [
510.3517536642828,
658.073751009259
],
"size": [
270,
178
],
"flags": {},
"order": 2,
"mode": 0,
"inputs": [
{
"localized_name": "max_tokens",
"name": "max_tokens",
"type": "INT",
"widget": {
"name": "max_tokens"
},
"link": null
},
{
"localized_name": "temperature",
"name": "temperature",
"type": "FLOAT",
"widget": {
"name": "temperature"
},
"link": null
},
{
"localized_name": "top_p",
"name": "top_p",
"type": "FLOAT",
"widget": {
"name": "top_p"
},
"link": null
},
{
"localized_name": "repetition_penalty",
"name": "repetition_penalty",
"type": "FLOAT",
"widget": {
"name": "repetition_penalty"
},
"link": null
},
{
"localized_name": "seed",
"name": "seed",
"type": "INT",
"widget": {
"name": "seed"
},
"link": null
}
],
"outputs": [
{
"localized_name": "AR sampling params",
"name": "AR sampling params",
"type": "SAMPLING_PARAMS",
"links": [
12
]
}
],
"properties": {
"Node name for S&R": "VLLMOmniARSampling"
},
"widgets_values": [
100,
1,
1,
1,
-1,
"randomize"
]
},
{
"id": 7,
"type": "VLLMOmniARSampling",
"pos": [
503.33235181647115,
419.34158016181806
],
"size": [
270,
178
],
"flags": {},
"order": 3,
"mode": 0,
"inputs": [
{
"localized_name": "max_tokens",
"name": "max_tokens",
"type": "INT",
"widget": {
"name": "max_tokens"
},
"link": null
},
{
"localized_name": "temperature",
"name": "temperature",
"type": "FLOAT",
"widget": {
"name": "temperature"
},
"link": null
},
{
"localized_name": "top_p",
"name": "top_p",
"type": "FLOAT",
"widget": {
"name": "top_p"
},
"link": null
},
{
"localized_name": "repetition_penalty",
"name": "repetition_penalty",
"type": "FLOAT",
"widget": {
"name": "repetition_penalty"
},
"link": null
},
{
"localized_name": "seed",
"name": "seed",
"type": "INT",
"widget": {
"name": "seed"
},
"link": null
}
],
"outputs": [
{
"localized_name": "AR sampling params",
"name": "AR sampling params",
"type": "SAMPLING_PARAMS",
"links": [
5,
13
]
}
],
"properties": {
"Node name for S&R": "VLLMOmniARSampling"
},
"widgets_values": [
100,
1,
1,
1,
-1,
"randomize"
]
},
{
"id": 8,
"type": "VLLMOmniSamplingParamsList",
"pos": [
820.6056617389042,
426.38372037182273
],
"size": [
263.066015625,
66
],
"flags": {},
"order": 7,
"mode": 0,
"inputs": [
{
"localized_name": "param1",
"name": "param1",
"type": "SAMPLING_PARAMS",
"link": 5
},
{
"localized_name": "param2",
"name": "param2",
"shape": 7,
"type": "SAMPLING_PARAMS",
"link": 12
},
{
"localized_name": "param3",
"name": "param3",
"shape": 7,
"type": "SAMPLING_PARAMS",
"link": 13
}
],
"outputs": [
{
"localized_name": "param list",
"name": "param list",
"type": "SAMPLING_PARAMS",
"links": [
14
]
}
],
"properties": {
"Node name for S&R": "VLLMOmniSamplingParamsList"
},
"widgets_values": []
},
{
"id": 11,
"type": "MarkdownNote",
"pos": [
826.2328280438272,
569.1890318701705
],
"size": [
333.8220435590464,
261.63596728060656
],
"flags": {},
"order": 4,
"mode": 0,
"inputs": [],
"outputs": [],
"title": "Note: Sampling Parameters",
"properties": {},
"widgets_values": [
"## Sampling Parameter Types\n\nThere are two types of sampling parameters: one for autoregression and one for diffusion.\nYou should ensure that you have chosen the correct type of sampling parameters for the model you request.\n\n## Stages & Shorthand\n\nFor multi-stage models such as Qwen Omni, you can either\n- connect one sampling parameter node, which is applied to all stages.\n- connect exactly the same number of sampling parameter nodes to a \"Multi-Stage Sampling Parameter List\" node, then connect this node to the primary request node.\n\nNote that this shorthand is intended to stay consistent with the [online serving API](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/examples/online_serving/qwen2_5_omni/)"
],
"color": "#432",
"bgcolor": "#000"
},
{
"id": 12,
"type": "MarkdownNote",
"pos": [
378.9866207003777,
152.59550252215752
],
"size": [
319.7287574247016,
107.15904081906785
],
"flags": {},
"order": 6,
"mode": 0,
"inputs": [],
"outputs": [],
"title": "Note: Input",
"properties": {},
"widgets_values": [
"Note that not all models support every modality as input. For example, `ByteDance-Seed/BAGEL-7B-MoT` in Multimodality Understanding mode only support text and image input.\n\nYou should ensure that the input are supported by the model. You can check the corresponding [online serving documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/examples/online_serving/bagel/) for confirmation."
],
"color": "#432",
"bgcolor": "#000"
},
{
"id": 2,
"type": "LoadImage",
"pos": [
394.4674804308822,
-207.6987397548834
],
"size": [
282.798828125,
314
],
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{
"localized_name": "image",
"name": "image",
"type": "COMBO",
"widget": {
"name": "image"
},
"link": null
},
{
"localized_name": "choose file to upload",
"name": "upload",
"type": "IMAGEUPLOAD",
"widget": {
"name": "upload"
},
"link": null
}
],
"outputs": [
{
"localized_name": "IMAGE",
"name": "IMAGE",
"type": "IMAGE",
"links": [
1
]
},
{
"localized_name": "MASK",
"name": "MASK",
"type": "MASK",
"links": null
}
],
"properties": {
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"example.png",
"image"
]
},
{
"id": 10,
"type": "PreviewAudio",
"pos": [
1664.548345556043,
297.5921292054054
],
"size": [
270,
88
],
"flags": {},
"order": 10,
"mode": 0,
"inputs": [
{
"localized_name": "audio",
"name": "audio",
"type": "AUDIO",
"link": 9
},
{
"localized_name": "audioUI",
"name": "audioUI",
"type": "AUDIO_UI",
"widget": {
"name": "audioUI"
},
"link": null
}
],
"outputs": [],
"properties": {
"Node name for S&R": "PreviewAudio"
},
"widgets_values": []
},
{
"id": 9,
"type": "ShowText|pysssss",
"pos": [
1649.2506875091847,
66.22823888292349
],
"size": [
318.7188464232943,
173.38502269972975
],
"flags": {},
"order": 9,
"mode": 0,
"inputs": [
{
"localized_name": "text",
"name": "text",
"type": "STRING",
"link": 8
}
],
"outputs": [
{
"localized_name": "STRING",
"name": "STRING",
"shape": 6,
"type": "STRING",
"links": null
}
],
"properties": {
"Node name for S&R": "ShowText|pysssss"
},
"widgets_values": []
}
],
"links": [
[
1,
2,
0,
1,
0,
"IMAGE"
],
[
2,
3,
0,
1,
1,
"VIDEO"
],
[
3,
4,
0,
1,
2,
"AUDIO"
],
[
5,
7,
0,
8,
0,
"SAMPLING_PARAMS"
],
[
8,
1,
0,
9,
0,
"STRING"
],
[
9,
1,
1,
10,
0,
"AUDIO"
],
[
12,
5,
0,
8,
1,
"SAMPLING_PARAMS"
],
[
13,
7,
0,
8,
2,
"SAMPLING_PARAMS"
],
[
14,
8,
0,
1,
3,
"SAMPLING_PARAMS"
]
],
"groups": [
{
"id": 1,
"title": "Sampling Parameters",
"bounding": [
480.1649301181556,
341.08402513937995,
692.4113568972277,
510.48648853403665
],
"color": "#3f789e",
"font_size": 24,
"flags": {}
},
{
"id": 2,
"title": "Input",
"bounding": [
344.2364817813528,
-287.5850484183313,
704.8832238768634,
559.5124009832894
],
"color": "#3f789e",
"font_size": 24,
"flags": {}
}
],
"config": {},
"extra": {
"workflowRendererVersion": "LG",
"ds": {
"scale": 0.9478575057427204,
"offset": [
33.81037136029557,
307.1974296197726
]
}
},
"version": 0.4
}
-198
View File
@@ -1,198 +0,0 @@
import base64
import mimetypes
import os
import time
import uuid
from dataclasses import dataclass, field
from typing import Any
import aiohttp
from tqdm import tqdm
@dataclass
class RequestFuncInput:
prompt: str
api_url: str
model: str
width: int | None = None
height: int | None = None
num_frames: int | None = None
num_inference_steps: int | None = None
seed: int | None = None
fps: int | None = None
timestamp: float | None = None
slo_ms: float | None = None
extra_body: dict[str, Any] = field(default_factory=dict)
image_paths: list[str] | None = None
request_id: str = field(default_factory=lambda: str(uuid.uuid4()))
@dataclass
class RequestFuncOutput:
success: bool = False
latency: float = 0.0
error: str = ""
start_time: float = 0.0
response_body: dict[str, Any] = field(default_factory=dict)
peak_memory_mb: float = 0.0
slo_achieved: bool | None = None
def _guess_mime_type(path: str) -> str:
mime, _ = mimetypes.guess_type(path)
return mime or "application/octet-stream"
def _encode_image_as_data_url(path: str) -> str:
with open(path, "rb") as f:
encoded = base64.b64encode(f.read()).decode("utf-8")
mime = _guess_mime_type(path)
return f"data:{mime};base64,{encoded}"
async def async_request_chat_completions(
input: RequestFuncInput,
session: aiohttp.ClientSession,
pbar: tqdm | None = None,
) -> RequestFuncOutput:
output = RequestFuncOutput()
output.start_time = time.perf_counter()
extra_body = dict(input.extra_body)
if input.width and input.height:
extra_body.setdefault("height", input.height)
extra_body.setdefault("width", input.width)
if input.num_frames:
extra_body.setdefault("num_frames", input.num_frames)
if input.num_inference_steps:
extra_body.setdefault("num_inference_steps", input.num_inference_steps)
if input.seed is not None:
extra_body.setdefault("seed", input.seed)
if input.fps:
extra_body.setdefault("fps", input.fps)
if input.image_paths and len(input.image_paths) > 0:
content = []
if input.prompt:
content.append({"type": "text", "text": input.prompt})
for img_path in input.image_paths:
if not os.path.exists(img_path):
output.error = f"Image file not found: {img_path}"
output.success = False
if pbar:
pbar.update(1)
return output
content.append(
{
"type": "image_url",
"image_url": {"url": _encode_image_as_data_url(img_path)},
}
)
messages = [{"role": "user", "content": content}]
else:
messages = [{"role": "user", "content": input.prompt}]
payload = {
"model": input.model,
"messages": messages,
}
if extra_body:
payload["extra_body"] = extra_body
try:
async with session.post(input.api_url, json=payload) as response:
if response.status == 200:
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"]
else:
output.error = f"HTTP {response.status}: {await response.text()}"
output.success = False
except Exception as e:
output.error = str(e)
output.success = False
output.latency = time.perf_counter() - output.start_time
if output.success and input.slo_ms is not None:
output.slo_achieved = (output.latency * 1000.0) <= float(input.slo_ms)
if pbar:
pbar.update(1)
return output
async def async_request_openai_images(
input: RequestFuncInput,
session: aiohttp.ClientSession,
pbar: tqdm | None = None,
) -> RequestFuncOutput:
"""
Send request to OpenAI's /v1/images/generations endpoint.
"""
output = RequestFuncOutput()
output.start_time = time.perf_counter()
# Build size string from width/height
width = input.width or 1024
height = input.height or 1024
size = f"{width}x{height}"
payload: dict[str, Any] = {
"model": input.model,
"prompt": input.prompt,
"n": 1,
"size": size,
"response_format": "b64_json",
}
# Add optional parameters
if input.seed is not None:
payload["seed"] = input.seed
if input.num_inference_steps is not None:
payload["num_inference_steps"] = input.num_inference_steps
# Add any extra body parameters
if input.extra_body:
for key, value in input.extra_body.items():
if key not in payload:
payload[key] = value
headers = {
"Content-Type": "application/json",
"Authorization": "Bearer EMPTY",
}
try:
async with session.post(input.api_url, json=payload, headers=headers) as response:
if response.status == 200:
resp_json = await response.json()
output.response_body = resp_json
output.success = True
# Check for usage/memory info if available
if "usage" in resp_json and "peak_memory_mb" in resp_json.get("usage", {}):
output.peak_memory_mb = resp_json["usage"]["peak_memory_mb"]
else:
output.error = f"HTTP {response.status}: {await response.text()}"
output.success = False
except Exception as e:
output.error = str(e)
output.success = False
output.latency = time.perf_counter() - output.start_time
if output.success and input.slo_ms is not None:
output.slo_achieved = (output.latency * 1000.0) <= float(input.slo_ms)
if pbar:
pbar.update(1)
return output
backends_function_mapping = {
"vllm-omni": (async_request_chat_completions, "/v1/chat/completions"),
"openai": (async_request_openai_images, "/v1/images/generations"),
}
@@ -6,60 +6,80 @@
Benchmark online serving for diffusion models (Image/Video Generation).
If you want to use i2v, i2i dataset, you should `uv pip install gdown` first
Supports multiple backends:
- vllm-omni: Uses /v1/chat/completions endpoint (default)
- openai: Uses /v1/images/generations endpoint
Usage:
# Video (vllm-omni backend)
# Video
t2v:
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
--backend vllm-omni --dataset vbench --task t2v --num-prompts 10 \
--dataset vbench --task t2v --num-prompts 10 \
--height 480 --width 640 --fps 16 --num-frames 80
i2v:
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
--backend vllm-omni --dataset vbench --task i2v --num-prompts 10
--dataset vbench --task i2v --num-prompts 10
# Image (vllm-omni backend)
# Image
t2i:
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
--backend vllm-omni --dataset vbench --task t2i --num-prompts 10 \
--dataset vbench --task t2i --num-prompts 10 \
--height 1024 --width 1024
i2i:
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
--backend vllm-omni --dataset vbench --task i2i --num-prompts 10
# Image (openai backend)
t2i:
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
--backend openai --dataset vbench --task t2i --num-prompts 10 \
--height 1024 --width 1024 --port 3000
--dataset vbench --task i2i --num-prompts 10
"""
import argparse
import ast
import asyncio
import base64
import glob
import json
import mimetypes
import os
import time
import uuid
from abc import ABC, abstractmethod
from collections.abc import AsyncGenerator
from dataclasses import replace
from dataclasses import dataclass, field, replace
from typing import Any
import aiohttp
import numpy as np
import requests
from backends import RequestFuncInput, RequestFuncOutput, backends_function_mapping
from tqdm.asyncio import tqdm
@dataclass
class RequestFuncInput:
prompt: str
api_url: str
model: str
width: int | None = None
height: int | None = None
num_frames: int | None = None
num_inference_steps: int | None = None
seed: int | None = None
fps: int | None = None
timestamp: float | None = None
slo_ms: float | None = None
extra_body: dict[str, Any] = field(default_factory=dict)
image_paths: list[str] | None = None
request_id: str = field(default_factory=lambda: str(uuid.uuid4()))
@dataclass
class RequestFuncOutput:
success: bool = False
latency: float = 0.0
error: str = ""
start_time: float = 0.0
response_body: dict[str, Any] = field(default_factory=dict)
peak_memory_mb: float = 0.0
slo_achieved: bool | None = None
class BaseDataset(ABC):
def __init__(self, args, api_url: str, model: str):
self.args = args
@@ -664,6 +684,92 @@ async def iter_requests(
yield req
def _guess_mime_type(path: str) -> str:
mime, _ = mimetypes.guess_type(path)
return mime or "application/octet-stream"
def _encode_image_as_data_url(path: str) -> str:
with open(path, "rb") as f:
encoded = base64.b64encode(f.read()).decode("utf-8")
mime = _guess_mime_type(path)
return f"data:{mime};base64,{encoded}"
async def async_request_chat_completions(
input: RequestFuncInput,
session: aiohttp.ClientSession,
pbar: tqdm | None = None,
) -> RequestFuncOutput:
output = RequestFuncOutput()
output.start_time = time.perf_counter()
extra_body = dict(input.extra_body)
if input.width and input.height:
extra_body.setdefault("height", input.height)
extra_body.setdefault("width", input.width)
if input.num_frames:
extra_body.setdefault("num_frames", input.num_frames)
if input.num_inference_steps:
extra_body.setdefault("num_inference_steps", input.num_inference_steps)
if input.seed is not None:
extra_body.setdefault("seed", input.seed)
if input.fps:
extra_body.setdefault("fps", input.fps)
if input.image_paths and len(input.image_paths) > 0:
content = []
if input.prompt:
content.append({"type": "text", "text": input.prompt})
for img_path in input.image_paths:
if not os.path.exists(img_path):
output.error = f"Image file not found: {img_path}"
output.success = False
if pbar:
pbar.update(1)
return output
content.append(
{
"type": "image_url",
"image_url": {"url": _encode_image_as_data_url(img_path)},
}
)
messages = [{"role": "user", "content": content}]
else:
messages = [{"role": "user", "content": input.prompt}]
payload = {
"model": input.model,
"messages": messages,
}
if extra_body:
payload["extra_body"] = extra_body
try:
async with session.post(input.api_url, json=payload) as response:
if response.status == 200:
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"]
else:
output.error = f"HTTP {response.status}: {await response.text()}"
output.success = False
except Exception as e:
output.error = str(e)
output.success = False
output.latency = time.perf_counter() - output.start_time
if output.success and input.slo_ms is not None:
output.slo_achieved = (output.latency * 1000.0) <= float(input.slo_ms)
if pbar:
pbar.update(1)
return output
def calculate_metrics(
outputs: list[RequestFuncOutput],
total_duration: float,
@@ -742,9 +848,9 @@ async def benchmark(args):
if args.base_url is None:
args.base_url = f"http://{args.host}:{args.port}"
# Setup API URL and request function based on backend
request_func, api_url = backends_function_mapping[args.backend]
api_url = f"{args.base_url}{api_url}"
# Setup dataset (vLLM-Omni supports diffusion via /v1/chat/completions)
api_url = f"{args.base_url}/v1/chat/completions"
request_func = async_request_chat_completions
if args.dataset == "vbench":
dataset = VBenchDataset(args, api_url, args.model)
@@ -814,16 +920,9 @@ async def benchmark(args):
# Calculate metrics
metrics = calculate_metrics(outputs, total_duration, requests_list, args, args.slo)
# Add configuration info to metrics for JSON output
metrics["backend"] = args.backend
metrics["model"] = args.model
metrics["dataset"] = args.dataset
metrics["task"] = args.task
print("\n{s:{c}^{n}}".format(s=" Serving Benchmark Result ", n=60, c="="))
# Section 1: Configuration
print("{:<40} {:<15}".format("Backend:", args.backend))
print("{:<40} {:<15}".format("Model:", args.model))
print("{:<40} {:<15}".format("Dataset:", args.dataset))
print("{:<40} {:<15}".format("Task:", args.task))
@@ -879,13 +978,6 @@ if __name__ == "__main__":
parser.add_argument("--host", type=str, default="localhost", help="Server host.")
parser.add_argument("--port", type=int, default=8091, help="Server port.")
parser.add_argument("--model", type=str, default="default", help="Model name.")
parser.add_argument(
"--backend",
type=str,
default="vllm-omni",
choices=["vllm-omni", "openai"],
help="Backend to target the benchmark to.",
)
parser.add_argument(
"--dataset",
type=str,
@@ -1,397 +0,0 @@
# RDMA Test Configuration Guide
This document explains how to configure the RDMA environment and run tests for `MooncakeTransferEngineConnector`.
## Table of Contents
- [Docker Container Permissions](#docker-container-permissions)
- [Single-Node Testing](#single-node-testing)
- [Multi-Node Testing](#multi-node-testing)
- [Running Tests](#running-tests)
- [Cross-Node Testing](#cross-node-testing)
- [Troubleshooting](#troubleshooting)
---
## Docker Container Permissions
RDMA tests require access to InfiniBand/RoCE devices and system topology. Add the following permissions when running `docker run`.
### Option 1: Minimal Permissions (Recommended)
```bash
docker run -it \
--cap-add=SYS_PTRACE \
--cap-add=IPC_LOCK \
--security-opt seccomp=unconfined \
--network=host \
--device=/dev/infiniband \
-v /sys/class/infiniband:/sys/class/infiniband:ro \
your-image:tag
```
Parameter explanation:
- `--cap-add=SYS_PTRACE`: Allow reading system topology information
- `--cap-add=IPC_LOCK`: Allow memory locking (required for RDMA memory registration)
- `--security-opt seccomp=unconfined`: Disable seccomp restrictions
- `--network=host`: Use host network (required for RDMA)
- `--device=/dev/infiniband`: Mount InfiniBand devices
- `-v /sys/class/infiniband`: Mount IB device info (read-only)
### Option 2: Full Permissions (Quick but not recommended for production)
```bash
docker run -it \
--privileged \
--network=host \
your-image:tag
```
`--privileged` grants full host permissions. Suitable for quick testing but not recommended for production.
---
## Single-Node Testing
When running single-node tests (producer and consumer on the same machine), ensure they use the **same RDMA device**.
### Problem Background
InfiniBand devices use LID (Local Identifier) for routing. Different devices have different LIDs and cannot communicate directly. If no device is specified, Mooncake may assign different devices to connectors, causing handshake failures.
Common error:
```
[Handshake] Failed to modify QP to RTR, check mtu, gid, peer lid, peer qp num: Invalid argument [22]
```
### Solution
**Method 1: Set Environment Variable (Recommended)**
```bash
# List available RDMA devices
ibstat
# Select a device (e.g., mlx5_0)
export RDMA_DEVICE_NAME='mlx5_0'
# Run tests
pytest test_mooncake_transfer_engine_rdma.py -v -s
```
**Method 2: Use RoCE Devices**
If the system has RoCE devices (using IPv4 routing), the test code will automatically detect and prefer them. RoCE device GIDs start with `00:00:00:00:00:00:00:00:00:00:ff:ff` (IPv4-mapped).
**Method 3: Ensure MTU Consistency**
Make sure both endpoints use the same MTU:
```bash
# Check device MTU
ibstatus mlx5_0
```
---
## Multi-Node Testing
For multi-node tests, producer and consumer run on different machines connected via InfiniBand switch.
### Prerequisites
1. Both machines have Mooncake and RDMA drivers installed
2. Both machines are in the same InfiniBand subnet
3. Switch is properly configured
### Configuration
**Machine A (Producer):**
```bash
# Set RDMA host IP (InfiniBand interface IP)
export RDMA_TEST_HOST='10.0.0.1'
# Optional: Specify device
export RDMA_DEVICE_NAME='mlx5_0'
```
**Machine B (Consumer):**
```bash
# Set RDMA host IP
export RDMA_TEST_HOST='10.0.0.2'
# Optional: Specify device
export RDMA_DEVICE_NAME='mlx5_0'
```
### Verify Connectivity
```bash
# Ping IB interface
ping 10.0.0.2
# Test RDMA connectivity with ibping
# On Machine B (server)
ibping -S
# On Machine A (client)
ibping -G <Machine_B_GID>
```
---
## Running Tests
### Run All RDMA Tests (Single-Node, fast suite)
Slow tests (large payloads, stress, concurrency integrity) are marked `@pytest.mark.slow`. Use `-m "not slow"` to skip them in quick CI or local fast iteration.
```bash
cd tests/distributed/omni_connectors
# Fast suite only (excludes slow/stress tests)
pytest test_mooncake_transfer_engine_rdma.py test_mooncake_transfer_engine_buffer.py -v -s -m "not slow"
```
### Run Including Slow Tests
```bash
# Run ALL tests including slow/stress tests
pytest test_mooncake_transfer_engine_rdma.py test_mooncake_transfer_engine_buffer.py -v -s
# Run ONLY the slow/stress tests
pytest test_mooncake_transfer_engine_rdma.py test_mooncake_transfer_engine_buffer.py -v -s -m slow
```
### Run Buffer Management Tests
```bash
# Fast only
pytest test_mooncake_transfer_engine_buffer.py -v -s -m "not slow"
# Including allocator invariant tests (double-free, overlap, merge)
pytest test_mooncake_transfer_engine_buffer.py -v -s
```
### Run Specific Test Classes
```bash
# Basic connector tests
pytest test_mooncake_transfer_engine_rdma.py::TestBasicConnector -v -s
# End-to-end RDMA transfer tests
pytest test_mooncake_transfer_engine_rdma.py::TestEndToEnd -v -s
# Lifecycle & resource management tests
pytest test_mooncake_transfer_engine_rdma.py::TestLifecycle -v -s
# GPU memory pool tests (requires CUDA)
pytest test_mooncake_transfer_engine_rdma.py::TestGPUPool -v -s
# Stress / correctness tests (slow)
pytest test_mooncake_transfer_engine_rdma.py::TestStressCorrectness -v -s
```
### RDMA Environment Diagnostics
For quick diagnostics (device status, Mooncake availability, env vars, etc.),
see the [Troubleshooting section](../../../docs/design/feature/omni_connectors/mooncake_transfer_engine_connector.md#troubleshooting)
in the connector documentation.
---
## Cross-Node Testing
The `cross_node_mooncake_transfer_engine.py` script enables testing RDMA transfers between two separate physical machines. This script is **not** auto-discovered by `pytest` (it does not start with `test_`) — it must be run manually on each node.
### Prerequisites
1. Both machines have Mooncake installed
2. Both machines are connected via InfiniBand/RoCE switch
3. Firewall allows ZMQ ports (default: 15500, 15501)
4. Same RDMA device name on both nodes (if multiple devices exist)
### Running Cross-Node Tests
**On Machine A (Producer) — start first:**
```bash
cd benchmarks/distributed/omni_connectors/
# Optional: specify device if multiple exist
export RDMA_DEVICE_NAME='mlx5_0'
python cross_node_mooncake_transfer_engine.py \
--role producer \
--local-host <PRODUCER_IP> \
--remote-host <CONSUMER_IP> \
--tensor-size-mb 100 \
--num-transfers 3
```
**On Machine B (Consumer) — start after producer:**
```bash
cd benchmarks/distributed/omni_connectors/
export RDMA_DEVICE_NAME='mlx5_0'
python cross_node_mooncake_transfer_engine.py \
--role consumer \
--local-host <CONSUMER_IP> \
--remote-host <PRODUCER_IP> \
--tensor-size-mb 100 \
--num-transfers 3
```
### Transfer Modes
| Mode | Description | Example |
|------|-------------|---------|
| `copy` | Normal path — tensor copied to RDMA pool (default) | `--mode copy` |
| `zerocopy` | Zero-copy path — data created directly in RDMA pool | `--mode zerocopy` |
| `gpu` | GPU transfer — RDMA pool on GPU, uses GPUDirect | `--mode gpu --gpu-id 0` |
### Benchmark Mode
Skip MD5 verification and measure pure RDMA throughput:
```bash
# Producer
python cross_node_mooncake_transfer_engine.py \
--role producer \
--local-host <PRODUCER_IP> \
--remote-host <CONSUMER_IP> \
--tensor-size-mb 1024 \
--num-transfers 20 \
--benchmark
# Consumer
python cross_node_mooncake_transfer_engine.py \
--role consumer \
--local-host <CONSUMER_IP> \
--remote-host <PRODUCER_IP> \
--tensor-size-mb 1024 \
--num-transfers 20 \
--benchmark
```
### Cross-Node Test Options
| Option | Description | Default |
|--------|-------------|---------|
| `--role` | `producer` or `consumer` | Required |
| `--local-host` | Local RDMA IP address | Required |
| `--remote-host` | Remote RDMA IP address | Required |
| `--local-port` | Local ZMQ port for RDMA data | 15500 |
| `--remote-port` | Remote ZMQ port for RDMA data | 15500 |
| `--ctrl-port` | Control channel port | 15501 |
| `--tensor-size-mb` | Tensor size in MB | 100 |
| `--num-transfers` | Number of transfers | 3 |
| `--mode` | `copy`, `zerocopy`, or `gpu` | `copy` |
| `--gpu-id` | GPU ID for GPU mode | 0 |
| `--benchmark` | Skip MD5, pure performance test | off |
---
## Troubleshooting
### 1. "Failed to modify QP to RTR" Error
**Cause**: QP handshake failed, usually due to device configuration mismatch.
**Solution**:
```bash
# Force using the same device
export RDMA_DEVICE_NAME='mlx5_0'
```
### 2. "Mooncake TransferEngine is not available"
**Cause**: Mooncake not installed or import failed.
**Solution**:
```bash
# Check Mooncake installation
python -c "from mooncake.engine import TransferEngine; print('OK')"
# Reinstall if needed
pip install mooncake-transfer-engine
# Or using uv
uv pip install mooncake-transfer-engine
```
### 3. "Permission denied" accessing /dev/infiniband
**Cause**: Container lacks IB device access permissions.
**Solution**:
```bash
docker run --device=/dev/infiniband --cap-add=IPC_LOCK ...
```
### 4. Test Timeout
**Cause**: RDMA connection establishment failed or network latency.
**Solution**:
```bash
# Check network status
ibstat
ibstatus
```
### 5. GPU Test Failed "CUDA is not available"
**Cause**: CUDA environment not configured or GPU unavailable.
**Solution**:
```bash
# Check CUDA
python -c "import torch; print(torch.cuda.is_available())"
# Docker needs NVIDIA runtime
docker run --gpus all ...
```
---
## Environment Variables Reference
| Variable | Description | Example |
|----------|-------------|---------|
| `RDMA_DEVICE_NAME` | Specify RDMA device name | `mlx5_0` |
| `RDMA_TEST_HOST` | Specify test host IP | `10.0.0.1` |
| `MC_TE_METRIC` | Enable Mooncake metrics | `1` |
| `MC_IB_PCI_RELAXED_ORDERING` | Enable PCIe relaxed ordering | `1` |
---
## Test Files Overview
| File | Description | Auto-discovered by pytest |
|------|-------------|--------------------------|
| `test_mooncake_transfer_engine_rdma.py` | Integration tests for MooncakeTransferEngineConnector (basic, E2E, lifecycle, GPU) | Yes |
| `test_mooncake_transfer_engine_buffer.py` | Memory pool and buffer management unit tests | Yes |
| `cross_node_mooncake_transfer_engine.py` | Cross-node (multi-machine) testing script — run manually | No (filename does not start with `test_`) |
### test_mooncake_transfer_engine_rdma.py — Test Classes
| Test Class | Memory Pool | Marker | Description |
|------------|-------------|--------|-------------|
| `TestBasicConnector` | CPU | — | Initialization, put tensor/bytes/object, cleanup, pool exhaustion |
| `TestEndToEnd` | CPU | — | E2E RDMA transfer: tensor, bytes, object, zero-copy, large payload (100MB), mixed types, concurrency |
| `TestLifecycle` | CPU | — | Close, context manager, double-close safety |
| `TestGPUPool` | GPU | — | GPU pool init, put CPU/GPU tensor, GPU E2E transfer |
| `TestStressCorrectness` | CPU | `slow` | Concurrent put+get with MD5 integrity, bidirectional concurrency, edge cases (1-element tensor, empty bytes), 500MB payload, rapid alloc/free cycles |
### test_mooncake_transfer_engine_buffer.py — Test Classes
| Test Class | Marker | Description |
|------------|--------|-------------|
| `TestBufferAllocator` | — | Basic alloc/free, alignment, exhaustion/recovery, thread safety |
| `TestAllocatorInvariants` | `slow` | Double-free safety, overlap corruption detection, adjacent-block merging, fragmentation/defrag |
| `TestManagedBuffer` | — | Tensor views, context manager |
@@ -1,644 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Cross-Node RDMA Test Script (Automated Version)
This script enables testing RDMA transfers between two separate machines.
Supports three transfer modes:
- copy: Normal path - tensor copied to RDMA pool (default)
- zerocopy: Zero-copy path - data created directly in RDMA pool
- gpu: GPU transfer - RDMA pool on GPU, uses GPUDirect
Usage:
# On Machine A (Producer) - start first:
python cross_node_mooncake_transfer_engine.py --role producer --local-host hostname_A --remote-host hostname_B
# On Machine B (Consumer) - start after producer:
python cross_node_mooncake_transfer_engine.py --role consumer --local-host hostname_B --remote-host hostname_A
# Zero-copy mode:
python cross_node_mooncake_transfer_engine.py --role producer ... --mode zerocopy
# GPU mode (requires GPUDirect RDMA support):
python cross_node_mooncake_transfer_engine.py --role producer ... --mode gpu --gpu-id 0
# Benchmark mode (skip random data generation and MD5 verification,
# measures pure RDMA throughput):
python cross_node_mooncake_transfer_engine.py --role producer ... --benchmark
Environment Variables:
RDMA_DEVICE_NAME: Specify RDMA device (e.g., mlx5_0)
MC_IB_PCI_RELAXED_ORDERING: Set to 1 to enable PCIe relaxed ordering
for higher RDMA throughput
"""
import argparse
import hashlib
import os
import sys
import time
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Any
import msgspec
import torch
import zmq
# Add parent path for imports
sys.path.insert(
0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))))
)
from vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector import (
ManagedBuffer,
MooncakeTransferEngineConnector,
TransferEngine,
)
def compute_md5(tensor: torch.Tensor) -> str:
"""Compute MD5 checksum of a tensor."""
if tensor.is_cuda:
tensor = tensor.cpu()
data = tensor.contiguous().view(torch.uint8).numpy().tobytes()
return hashlib.md5(data).hexdigest()
# Control channel message types
class CtrlMsg(msgspec.Struct):
"""Control channel message."""
msg_type: str # "READY", "TRANSFER", "ACK", "DONE", "ERROR"
request_id: str = ""
md5: str = ""
data_size: int = 0
error: str = ""
@dataclass
class TransferConfig:
"""Configuration for cross-node transfer test."""
local_host: str
remote_host: str
local_port: int
remote_port: int
ctrl_port: int
num_transfers: int
tensor_size_mb: int
mode: str # "copy", "zerocopy", "gpu"
gpu_id: int = 0
benchmark: bool = False # Skip MD5 verification for pure performance test
@dataclass
class TransferStats:
"""Statistics for transfer operations."""
success_count: int = 0
fail_count: int = 0
total_bytes: int = 0
elapsed_time: float = 0.0
@property
def throughput_mbps(self) -> float:
if self.elapsed_time > 0:
return (self.total_bytes / (1024 * 1024)) / self.elapsed_time
return 0.0
def print_summary(self, role: str):
print(f"\n{'=' * 60}")
print(f" {role.upper()} SUMMARY")
print(f" Successful: {self.success_count}/{self.success_count + self.fail_count}")
print(f" Failed: {self.fail_count}/{self.success_count + self.fail_count}")
print(f" Total: {self.total_bytes / (1024 * 1024):.2f} MB")
print(f" Time: {self.elapsed_time:.2f} s")
print(f" Throughput: {self.throughput_mbps:.2f} MB/s")
print(f"{'=' * 60}")
class CrossNodeTester(ABC):
"""Abstract base class for cross-node RDMA testing."""
def __init__(self, config: TransferConfig):
self.config = config
self.connector: MooncakeTransferEngineConnector | None = None
self.zmq_ctx: zmq.Context | None = None
self.ctrl_socket: zmq.Socket | None = None
self.stats = TransferStats()
def get_connector_config(self) -> dict:
"""Get connector configuration based on mode."""
pool_size = int(self.config.tensor_size_mb * 1.5) * 1024 * 1024
pool_size = max(pool_size, 128 * 1024 * 1024)
conn_config = {
"host": self.config.local_host,
"zmq_port": self.config.local_port,
"protocol": "rdma",
"memory_pool_size": pool_size,
}
# Set device based on mode
if self.config.mode == "gpu":
conn_config["memory_pool_device"] = f"cuda:{self.config.gpu_id}"
else:
conn_config["memory_pool_device"] = "cpu"
# RDMA device name from environment
device_name = os.environ.get("RDMA_DEVICE_NAME")
if device_name:
conn_config["device_name"] = device_name
print(f"[CONFIG] Using RDMA device: {device_name}")
return conn_config
def initialize(self):
"""Initialize connector and ZMQ context."""
print(f"[{self.role}] Initializing connector...")
conn_config = self.get_connector_config()
self.connector = MooncakeTransferEngineConnector(conn_config)
self.zmq_ctx = zmq.Context()
print(f"[{self.role}] Ready at {self.config.local_host}:{self.config.local_port}")
def cleanup(self):
"""Cleanup resources."""
if self.ctrl_socket:
self.ctrl_socket.close()
if self.zmq_ctx:
self.zmq_ctx.term()
if self.connector:
self.connector.close()
print(f"[{self.role}] Closed.")
@property
@abstractmethod
def role(self) -> str:
pass
@abstractmethod
def run(self):
pass
class Producer(CrossNodeTester):
"""Producer node - sends data to consumer."""
@property
def role(self) -> str:
return "PRODUCER"
def print_header(self):
print(f"\n{'=' * 60}")
print(f" PRODUCER MODE ({self.config.mode.upper()})")
print(f" Local: {self.config.local_host}:{self.config.local_port}")
print(f" Remote: {self.config.remote_host}:{self.config.remote_port}")
print(f" Control Port: {self.config.ctrl_port}")
print(f" Transfer Mode: {self.config.mode}")
if self.config.mode == "gpu":
print(f" GPU ID: {self.config.gpu_id}")
print(f"{'=' * 60}\n")
def setup_control_channel(self):
"""Setup ZMQ control channel as server (REP socket)."""
self.ctrl_socket = self.zmq_ctx.socket(zmq.REP)
self.ctrl_socket.bind(f"tcp://*:{self.config.ctrl_port}")
print(f"[PRODUCER] Control channel listening on port {self.config.ctrl_port}")
print("[PRODUCER] Waiting for consumer to connect...")
def wait_for_consumer(self) -> bool:
"""Wait for consumer READY signal."""
msg_data = self.ctrl_socket.recv()
msg = msgspec.msgpack.decode(msg_data, type=CtrlMsg)
if msg.msg_type != "READY":
print(f"[PRODUCER] Unexpected message: {msg.msg_type}")
return False
print("[PRODUCER] Consumer connected!")
self.ctrl_socket.send(msgspec.msgpack.encode(CtrlMsg(msg_type="ACK")))
return True
def create_test_data(self, transfer_idx: int) -> tuple[Any, str, int]:
"""
Create test data based on transfer mode.
Returns:
(data, md5, size) tuple
"""
num_elements = (self.config.tensor_size_mb * 1024 * 1024) // 4
data_size = num_elements * 4
if self.config.mode == "zerocopy":
# Zero-Copy Path: Allocate directly from connector's pool
offset = self.connector.allocator.alloc(data_size)
managed_buf = ManagedBuffer(self.connector.allocator, offset, data_size, self.connector.pool)
if self.config.benchmark:
# In benchmark mode, skip random data generation (use uninitialized memory)
return managed_buf, "", data_size
else:
# Fill buffer with random data using tensor view
tensor_view = managed_buf.as_tensor(dtype=torch.float32, shape=(num_elements,))
random_data = torch.randn(num_elements, dtype=torch.float32)
if tensor_view.is_cuda:
tensor_view.copy_(random_data.to(tensor_view.device))
else:
tensor_view.copy_(random_data)
md5 = compute_md5(tensor_view)
return managed_buf, md5, data_size
elif self.config.mode == "gpu":
# GPU Path: Create tensor on GPU
device = f"cuda:{self.config.gpu_id}"
if self.config.benchmark:
# In benchmark mode, use empty tensor (no random generation)
gpu_tensor = torch.empty(num_elements, dtype=torch.float32, device=device)
return gpu_tensor, "", data_size
else:
cpu_tensor = torch.randn(num_elements, dtype=torch.float32)
md5 = compute_md5(cpu_tensor)
gpu_tensor = cpu_tensor.to(device)
return gpu_tensor, md5, data_size
else:
# Copy Path (default): Create regular CPU tensor
if self.config.benchmark:
# In benchmark mode, use empty tensor (no random generation)
tensor = torch.empty(num_elements, dtype=torch.float32)
return tensor, "", data_size
else:
tensor = torch.randn(num_elements, dtype=torch.float32)
md5 = compute_md5(tensor)
return tensor, md5, data_size
def do_transfer(self, transfer_idx: int) -> bool:
"""Perform a single transfer."""
req_id = f"cross_node_transfer_{transfer_idx}"
if not self.config.benchmark:
print(f"\n[PRODUCER] Transfer {transfer_idx + 1}/{self.config.num_transfers}")
# Create test data
t0 = time.time()
data, md5, data_size = self.create_test_data(transfer_idx)
t_create = time.time() - t0
if not self.config.benchmark:
print(f" Mode: {self.config.mode}")
print(f" Size: {self.config.tensor_size_mb} MB")
if md5:
print(f" MD5: {md5[:16]}...")
print(f" Create time: {t_create * 1000:.1f} ms")
# Put data
t1 = time.time()
success, size, metadata = self.connector.put("producer", "consumer", req_id, data)
t_put = time.time() - t1
if not success:
print(" [FAIL] Put failed")
return False
if not self.config.benchmark:
print(f" [OK] Put successful, {size} bytes ({t_put * 1000:.1f} ms)")
# Wait for consumer to request transfer info
msg_data = self.ctrl_socket.recv()
msg = msgspec.msgpack.decode(msg_data, type=CtrlMsg)
if msg.msg_type != "READY":
print(f" [ERROR] Unexpected message: {msg.msg_type}")
return False
# Send transfer metadata to consumer
transfer_msg = CtrlMsg(
msg_type="TRANSFER",
request_id=req_id,
md5=md5,
data_size=data_size,
)
self.ctrl_socket.send(msgspec.msgpack.encode(transfer_msg))
# Wait for consumer ACK (this includes RDMA transfer time)
t2 = time.time()
msg_data = self.ctrl_socket.recv()
t_rdma = time.time() - t2
msg = msgspec.msgpack.decode(msg_data, type=CtrlMsg)
success = msg.msg_type == "ACK"
if success:
if not self.config.benchmark:
print(f" [OK] RDMA transfer complete ({t_rdma * 1000:.1f} ms)")
self.stats.success_count += 1
self.stats.total_bytes += size
else:
print(f" [WARN] Consumer reported error: {msg.error}")
self.stats.fail_count += 1
# Send ACK to allow consumer to continue
self.ctrl_socket.send(msgspec.msgpack.encode(CtrlMsg(msg_type="ACK")))
# Cleanup buffer
self.connector.cleanup(req_id)
return success
def run(self):
"""Run the producer."""
self.print_header()
self.initialize()
self.setup_control_channel()
try:
if not self.wait_for_consumer():
return
if self.config.benchmark:
print(
f"[BENCHMARK] Running {self.config.num_transfers} "
f"transfers of {self.config.tensor_size_mb} MB each..."
)
start_time = time.time()
for i in range(self.config.num_transfers):
self.do_transfer(i)
if self.config.benchmark and (i + 1) % 10 == 0:
elapsed = time.time() - start_time
current_throughput = (self.stats.total_bytes / (1024 * 1024)) / elapsed
print(f" Progress: {i + 1}/{self.config.num_transfers}, Throughput: {current_throughput:.2f} MB/s")
self.stats.elapsed_time = time.time() - start_time
self.stats.print_summary("PRODUCER")
# Wait for final consumer message and send DONE
self.ctrl_socket.recv()
self.ctrl_socket.send(msgspec.msgpack.encode(CtrlMsg(msg_type="DONE")))
finally:
self.cleanup()
class Consumer(CrossNodeTester):
"""Consumer node - receives data from producer."""
@property
def role(self) -> str:
return "CONSUMER"
def print_header(self):
print(f"\n{'=' * 60}")
print(f" CONSUMER MODE ({self.config.mode.upper()})")
print(f" Local: {self.config.local_host}:{self.config.local_port}")
print(f" Remote: {self.config.remote_host}:{self.config.remote_port}")
print(f" Control Port: {self.config.ctrl_port}")
print(f" Transfer Mode: {self.config.mode}")
if self.config.mode == "gpu":
print(f" GPU ID: {self.config.gpu_id}")
print(f"{'=' * 60}\n")
def setup_control_channel(self):
"""Setup ZMQ control channel as client (REQ socket)."""
self.ctrl_socket = self.zmq_ctx.socket(zmq.REQ)
ctrl_addr = f"tcp://{self.config.remote_host}:{self.config.ctrl_port}"
print(f"[CONSUMER] Connecting to producer control channel at {ctrl_addr}...")
self.ctrl_socket.connect(ctrl_addr)
def connect_to_producer(self) -> bool:
"""Connect to producer and send READY signal."""
self.ctrl_socket.send(msgspec.msgpack.encode(CtrlMsg(msg_type="READY")))
msg_data = self.ctrl_socket.recv()
msg = msgspec.msgpack.decode(msg_data, type=CtrlMsg)
if msg.msg_type != "ACK":
print(f"[CONSUMER] Unexpected response: {msg.msg_type}")
return False
print("[CONSUMER] Connected to producer! Starting transfers...")
return True
def do_transfer(self, transfer_idx: int) -> bool:
"""Perform a single transfer."""
if not self.config.benchmark:
print(f"\n[CONSUMER] Transfer {transfer_idx + 1}/{self.config.num_transfers}")
# Request next transfer info
self.ctrl_socket.send(msgspec.msgpack.encode(CtrlMsg(msg_type="READY")))
msg_data = self.ctrl_socket.recv()
msg = msgspec.msgpack.decode(msg_data, type=CtrlMsg)
if msg.msg_type == "DONE":
print("[CONSUMER] Producer signaled completion")
return False
if msg.msg_type != "TRANSFER":
print(f"[CONSUMER] Unexpected message: {msg.msg_type}")
return False
req_id = msg.request_id
expected_md5 = msg.md5
data_size = msg.data_size
num_elements = data_size // 4
if not self.config.benchmark:
print(f" Request ID: {req_id}")
if expected_md5:
print(f" Expected MD5: {expected_md5[:16]}...")
print(f" Data Size: {data_size / (1024 * 1024):.2f} MB")
# Build metadata for get
metadata = {
"request_id": req_id,
"source_host": self.config.remote_host,
"source_port": self.config.remote_port,
"data_size": data_size,
"dtype": "float32",
"shape": [num_elements],
"is_fast_path": True,
}
if not self.config.benchmark:
print(f" [INFO] Requesting from {self.config.remote_host}:{self.config.remote_port}")
# Get data with timing
t0 = time.time()
result = self.connector.get("producer", "consumer", req_id, metadata)
t_get = time.time() - t0
response_msg = CtrlMsg(msg_type="ERROR", error="Get failed")
if result is not None:
recv_buffer, recv_size = result
if not self.config.benchmark:
print(f" [OK] Get successful, {recv_size} bytes ({t_get * 1000:.1f} ms)")
if isinstance(recv_buffer, ManagedBuffer):
# In benchmark mode, skip MD5 verification
if self.config.benchmark or not expected_md5:
response_msg = CtrlMsg(msg_type="ACK")
self.stats.success_count += 1
self.stats.total_bytes += recv_size
else:
# Verify data
t1 = time.time()
reconstructed = recv_buffer.as_tensor(dtype=torch.float32, shape=(num_elements,))
recv_md5 = compute_md5(reconstructed)
t_md5 = time.time() - t1
print(f" MD5: {recv_md5[:16]}... ({t_md5 * 1000:.1f} ms)")
if recv_md5 == expected_md5:
print(" [PASS] MD5 checksum verified!")
response_msg = CtrlMsg(msg_type="ACK")
self.stats.success_count += 1
self.stats.total_bytes += recv_size
else:
print(" [FAIL] MD5 mismatch!")
response_msg = CtrlMsg(msg_type="ERROR", error="MD5 mismatch")
self.stats.fail_count += 1
recv_buffer.release()
else:
response_msg = CtrlMsg(msg_type="ACK")
self.stats.success_count += 1
self.stats.total_bytes += recv_size
else:
print(" [FAIL] Get failed")
self.stats.fail_count += 1
# Send response to producer
self.ctrl_socket.send(msgspec.msgpack.encode(response_msg))
# Wait for ACK
self.ctrl_socket.recv()
return response_msg.msg_type == "ACK"
def run(self):
"""Run the consumer."""
self.print_header()
self.initialize()
self.setup_control_channel()
try:
if not self.connect_to_producer():
return
if self.config.benchmark:
print(
f"[BENCHMARK] Running {self.config.num_transfers} "
f"transfers of {self.config.tensor_size_mb} MB each..."
)
start_time = time.time()
for i in range(self.config.num_transfers):
if not self.do_transfer(i):
break
if self.config.benchmark and (i + 1) % 10 == 0:
elapsed = time.time() - start_time
current_throughput = (self.stats.total_bytes / (1024 * 1024)) / elapsed
print(f" Progress: {i + 1}/{self.config.num_transfers}, Throughput: {current_throughput:.2f} MB/s")
self.stats.elapsed_time = time.time() - start_time
self.stats.print_summary("CONSUMER")
# Send final READY and wait for DONE
self.ctrl_socket.send(msgspec.msgpack.encode(CtrlMsg(msg_type="READY")))
self.ctrl_socket.recv()
finally:
self.cleanup()
def main():
parser = argparse.ArgumentParser(
description="Cross-Node RDMA Test (Automated)",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog="""
Transfer Modes:
copy - Normal path: tensor copied to RDMA pool (default)
zerocopy - Zero-copy path: data created directly in RDMA pool
gpu - GPU transfer: RDMA pool on GPU, uses GPUDirect
Examples:
# Copy mode (default):
python cross_node_mooncake_transfer_engine.py --role producer \
--local-host hostA --remote-host hostB
# Zero-copy mode:
python cross_node_mooncake_transfer_engine.py --role producer \
--local-host hostA --remote-host hostB --mode zerocopy
# GPU mode:
python cross_node_mooncake_transfer_engine.py --role producer \
--local-host hostA --remote-host hostB --mode gpu --gpu-id 0
# Benchmark mode (skip MD5, measure pure RDMA performance):
python cross_node_mooncake_transfer_engine.py --role producer \
--local-host hostA --remote-host hostB --benchmark
# With specific RDMA device:
RDMA_DEVICE_NAME=mlx5_0 python cross_node_mooncake_transfer_engine.py --role producer ...
""",
)
parser.add_argument(
"--role", required=True, choices=["producer", "consumer"], help="Role: producer (sends) or consumer (receives)"
)
parser.add_argument("--local-host", required=True, help="Local hostname or IP address")
parser.add_argument("--remote-host", required=True, help="Remote hostname or IP address")
parser.add_argument("--local-port", type=int, default=15500, help="Local ZMQ port for RDMA data (default: 15500)")
parser.add_argument("--remote-port", type=int, default=15500, help="Remote ZMQ port for RDMA data (default: 15500)")
parser.add_argument("--ctrl-port", type=int, default=15501, help="Control channel port (default: 15501)")
parser.add_argument("--num-transfers", type=int, default=20, help="Number of transfers to perform (default: 3)")
parser.add_argument("--tensor-size-mb", type=int, default=100, help="Tensor size in MB (default: 100)")
parser.add_argument(
"--mode",
choices=["copy", "zerocopy", "gpu"],
default="copy",
help="Transfer mode: copy, zerocopy, or gpu (default: copy)",
)
parser.add_argument("--gpu-id", type=int, default=0, help="GPU ID for GPU mode (default: 0)")
parser.add_argument(
"--benchmark", action="store_true", help="Benchmark mode: skip MD5 verification for pure performance test"
)
args = parser.parse_args()
# Check Mooncake
if TransferEngine is None:
print("[ERROR] Mooncake TransferEngine is not available.")
print("Install with: pip install mooncake")
sys.exit(1)
# Check CUDA for GPU mode
if args.mode == "gpu":
if not torch.cuda.is_available():
print("[ERROR] CUDA is not available but GPU mode was requested.")
sys.exit(1)
if args.gpu_id >= torch.cuda.device_count():
print(f"[ERROR] GPU {args.gpu_id} not available. Found {torch.cuda.device_count()} GPUs.")
sys.exit(1)
print(f"[INFO] Using GPU {args.gpu_id}: {torch.cuda.get_device_name(args.gpu_id)}")
config = TransferConfig(
local_host=args.local_host,
remote_host=args.remote_host,
local_port=args.local_port,
remote_port=args.remote_port,
ctrl_port=args.ctrl_port,
num_transfers=args.num_transfers,
tensor_size_mb=args.tensor_size_mb,
mode=args.mode,
gpu_id=args.gpu_id,
benchmark=args.benchmark,
)
if args.role == "producer":
tester = Producer(config)
else:
tester = Consumer(config)
tester.run()
if __name__ == "__main__":
main()
+26 -2
View File
@@ -1,5 +1,5 @@
ARG VLLM_BASE_IMAGE=vllm/vllm-openai
ARG VLLM_BASE_TAG=v0.16.0
ARG VLLM_BASE_TAG=v0.15.0
FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
ARG APP_DIR=/workspace/vllm-omni
WORKDIR ${APP_DIR}
@@ -8,14 +8,38 @@ COPY . .
# Install system dependencies
RUN apt-get update && \
apt-get install -y ffmpeg git sox libsox-fmt-all jq && \
apt-get install -y ffmpeg sox libsox-fmt-all && \
apt-get clean && \
rm -rf /var/lib/apt/lists/*
# Install vLLM wheel for the specific commit
# URL format: https://wheels.vllm.ai/<commit_hash>/
# Current vLLM wheel commit: 2d5be1dd5ce2e44dfea53ea03ff61143da5137eb
# Note: --prerelease=allow is required because nightly/commit wheels use
# pre-release version numbers (e.g. 2.10.0.dev...) which uv ignores by default,
# causing it to fall back to the older stable release from PyPI.
RUN uv pip install --system --upgrade --force-reinstall --prerelease=allow vllm --extra-index-url https://wheels.vllm.ai/2d5be1dd5ce2e44dfea53ea03ff61143da5137eb
RUN uv pip uninstall --system nvidia-cublas-cu12
# Install vllm-omni into the same uv-managed Python environment used by the base image.
# Use bash -c so that $(python3 -c ...) is expanded inside the container.
RUN uv pip install --system --no-cache-dir ".[dev]"
# pplx_kernels is optional; the base image may already ship a version.
# Reinstall attempt may fail if no compatible wheel exists for the current
# PyTorch/CUDA combination – allow the build to continue regardless.
RUN pip install --upgrade --force-reinstall "pplx_kernels" || true
RUN /bin/bash -c 'FLASHINFER_CUDA_TAG="$(python3 -c "import torch; print((torch.version.cuda or \"12.4\").replace(\".\", \"\"))")" && \
uv pip install --system --upgrade --force-reinstall \
"flashinfer-python==0.6.3" \
"flashinfer-cubin==0.6.3" \
"flashinfer-jit-cache==0.6.3" \
--extra-index-url "https://flashinfer.ai/whl/cu${FLASHINFER_CUDA_TAG}"'
# Pin nvidia-cublas-cu12 after vllm-omni install so dependency resolution does not overwrite it
RUN uv pip install --system --upgrade --force-reinstall "nvidia-cublas-cu12==12.9.1.4"
RUN uv pip install --system --upgrade --force-reinstall "numpy==2.2.6"
RUN ln -sf /usr/bin/python3 /usr/bin/python
ENTRYPOINT []
+2 -5
View File
@@ -2,16 +2,13 @@ ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
ARG VLLM_ASCEND_TAG=v0.14.0rc1
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
WORKDIR /vllm-workspace/vllm-ascend
RUN git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
RUN pip install -v -e .
ARG APP_DIR=/vllm-workspace/vllm-omni
WORKDIR ${APP_DIR}
COPY . .
RUN pip install -v -e .
# Install vllm-omni with dev dependencies
RUN pip install --no-cache-dir -e . --no-build-isolation
ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
+2 -5
View File
@@ -2,16 +2,13 @@ ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
ARG VLLM_ASCEND_TAG=v0.14.0rc1-a3
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
WORKDIR /vllm-workspace/vllm-ascend
RUN git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
RUN pip install -v -e .
ARG APP_DIR=/vllm-workspace/vllm-omni
WORKDIR ${APP_DIR}
COPY . .
RUN pip install -v -e .
# Install vllm-omni with dev dependencies
RUN pip install --no-cache-dir -e . --no-build-isolation
ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
+9 -34
View File
@@ -1,53 +1,28 @@
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.16.0
FROM ${BASE_IMAGE} AS base
# Declare a variable to know if we want to use the nightly build or the stable build.
# NOTE: REMINDER to vLLM-Omni rebase maintainer
# Remember to set `USE_NIGHTLY_BUILD` to 0, when switching back to stable vLLM docker image
# 1. If vLLM-omni maintainer is forced to use custom commits
# during rebasing, they can change the two variables
# 2. Whenever vLLM upstream has released a stable version,
# we should swap over to use stable release ASAP.
# We should avoid relying on custom commits.
ARG USE_NIGHTLY_BUILD=0
ARG VLLM_VERSION_OR_COMMIT_HASH=2d5be1dd5ce2e44dfea53ea03ff61143da5137eb
ARG ARG_PYTORCH_ROCM_ARCH
ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}}
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.15.0
FROM ${BASE_IMAGE} AS final
ARG COMMON_WORKDIR=/app
WORKDIR ${COMMON_WORKDIR}
# Step 1: Setup - Install system dependencies
RUN apt-get update && \
apt-get install -y ffmpeg git sox libsox-fmt-all jq && \
apt-get install -y ffmpeg && \
apt-get clean && \
rm -rf /var/lib/apt/lists/*
# Step 2: Conditionally reinstall vllm from source for nightly builds
RUN if [ "${USE_NIGHTLY_BUILD}" = "1" ]; then \
python3 -m pip uninstall -y vllm && rm -rf vllm && \
git clone https://github.com/vllm-project/vllm.git && \
cd vllm && \
git checkout ${VLLM_VERSION_OR_COMMIT_HASH} && \
python3 -m pip install -r requirements/rocm.txt && \
python3 setup.py clean --all && \
PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} python3 setup.py develop && \
cd ../ && \
rm -rf vllm/.git; \
fi
# Step 3: Copy vllm-omni code and install without uv
RUN mkdir -p ${COMMON_WORKDIR}/vllm-omni
# Step 2: Copy vllm-omni code and install without uv
COPY . ${COMMON_WORKDIR}/vllm-omni
RUN cd ${COMMON_WORKDIR}/vllm-omni && uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" --no-cache-dir ".[dev]" --no-build-isolation
RUN ln -sf /usr/bin/python3 /usr/bin/python
FROM base AS test
CMD ["/bin/bash"]
ENTRYPOINT []
# Set entrypoint for vllm-openai official images
FROM base AS vllm-openai
#Set entrypoint for vllm-openai official images
FROM final AS vllm-openai
ENTRYPOINT ["vllm", "serve", "--omni"]
+9 -18
View File
@@ -13,10 +13,9 @@ nav:
- examples/README.md
- Offline Inference:
- BAGEL-7B-MoT: user_guide/examples/offline_inference/bagel.md
- GLM-Image Multistage End-to-End Inference: user_guide/examples/offline_inference/glm_image.md
- Image-To-Image: user_guide/examples/offline_inference/image_to_image.md
- Image-To-Video: user_guide/examples/offline_inference/image_to_video.md
- MiMo-Audio Offline Inference: user_guide/examples/offline_inference/mimo_audio.md
- LoRA Inference(Diffusion): user_guide/examples/offline_inference/lora_inference.md
- Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md
- Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md
- Qwen3-TTS: user_guide/examples/offline_inference/qwen3_tts.md
@@ -25,15 +24,12 @@ nav:
- Text-To-Video: user_guide/examples/offline_inference/text_to_video.md
- Online Serving:
- BAGEL-7B-MoT: user_guide/examples/online_serving/bagel.md
- GLM-Image Online Serving: user_guide/examples/online_serving/glm_image.md
- Image-To-Image: user_guide/examples/online_serving/image_to_image.md
- Image-To-Video: user_guide/examples/online_serving/image_to_video.md
- Online serving Example of vLLM-Omni for MiMo-Audio: user_guide/examples/online_serving/mimo_audio.md
- LoRA Inference(Diffusion): user_guide/examples/online_serving/lora_inference.md
- Qwen2.5-Omni: user_guide/examples/online_serving/qwen2_5_omni.md
- Qwen3-Omni: user_guide/examples/online_serving/qwen3_omni.md
- Qwen3-TTS: user_guide/examples/online_serving/qwen3_tts.md
- Text-To-Image: user_guide/examples/online_serving/text_to_image.md
- Text-To-Video: user_guide/examples/online_serving/text_to_video.md
- General:
- usage/*
- Configuration:
@@ -44,29 +40,30 @@ nav:
- Features:
- Sleep Mode: features/sleep_mode.md
- Diffusion Features:
- Acceleration Overview: user_guide/diffusion_acceleration.md
- Overview: user_guide/diffusion_acceleration.md
- TeaCache: user_guide/diffusion/teacache.md
- Cache-DiT: user_guide/diffusion/cache_dit_acceleration.md
- Quantization:
- Overview: user_guide/diffusion/quantization/overview.md
- FP8: user_guide/diffusion/quantization/fp8.md
- GGUF: user_guide/diffusion/quantization/gguf.md
- Parallelism Acceleration: user_guide/diffusion/parallelism_acceleration.md
- CPU Offloading: user_guide/diffusion/cpu_offload_diffusion.md
- LoRA: user_guide/diffusion/lora.md
- Hybrid Sharded Data Parallel: design/feature/hsdp.md
- Custom Pipeline: features/custom_pipeline.md
- ComfyUI: features/comfyui.md
- Developer Guide:
- General:
- contributing/README.md
- pr_reviewer.md
- glob: contributing/*
flatten_single_child_sections: true
- Model Implementation:
- contributing/model/README.md
- contributing/model/adding_omni_model.md
- contributing/model/adding_diffusion_model.md
- Advanced Features:
- contributing/features/cfg_parallel.md
- contributing/features/sequence_parallel.md
- contributing/features/tensor_parallel.md
- contributing/features/cache_dit.md
- contributing/features/teacache.md
- CI: contributing/ci
- Design Documents:
- design/index.md
@@ -75,12 +72,6 @@ nav:
- design/feature/disaggregated_inference.md
- design/feature/ray_based_execution.md
- design/feature/omni_connectors/
- design/feature/cfg_parallel.md
- design/feature/sequence_parallel.md
- design/feature/tensor_parallel.md
- design/feature/cache_dit.md
- design/feature/teacache.md
- design/feature/async_chunk_design.md
- Module Design:
- design/module/ar_module.md
- design/module/dit_module.md
+1 -11
View File
@@ -7,6 +7,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.async_omni_llm.AsyncOmniLLM][]
- [vllm_omni.entrypoints.chat_utils.extract_audio_from_video_async][]
- [vllm_omni.entrypoints.cli.benchmark.base.OmniBenchmarkSubcommandBase][]
- [vllm_omni.entrypoints.cli.benchmark.main.OmniBenchmarkSubcommand][]
- [vllm_omni.entrypoints.cli.benchmark.serve.OmniBenchmarkServingSubcommand][]
@@ -18,7 +19,6 @@ Main entry points for vLLM-Omni inference and serving.
- [vllm_omni.entrypoints.omni_llm.OmniLLM][]
- [vllm_omni.entrypoints.omni_stage.OmniStage][]
- [vllm_omni.entrypoints.stage_utils.OmniStageTaskType][]
- [vllm_omni.entrypoints.zmq_utils.ZmqQueue][]
## Inputs
@@ -36,12 +36,6 @@ Input data structures for multi-modal inputs.
Engine classes for offline and online inference.
- [vllm_omni.diffusion.diffusion_engine.DiffusionEngine][]
- [vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector.BufferAllocator][]
- [vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector.ManagedBuffer][]
- [vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector.MooncakeAgentMetadata][]
- [vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector.MooncakeTransferEngineConnector][]
- [vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector.QueryRequest][]
- [vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector.QueryResponse][]
- [vllm_omni.engine.AdditionalInformationEntry][]
- [vllm_omni.engine.AdditionalInformationPayload][]
- [vllm_omni.engine.OmniEngineCoreOutput][]
@@ -78,8 +72,6 @@ Configuration classes.
- [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.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][]
- [vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts.Qwen3TTSSpeakerEncoderConfig][]
- [vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts.Qwen3TTSTalkerCodePredictorConfig][]
@@ -97,10 +89,8 @@ Configuration classes.
Worker classes and model runners for distributed inference.
- [vllm_omni.diffusion.worker.diffusion_model_runner.DiffusionModelRunner][]
- [vllm_omni.diffusion.worker.diffusion_worker.CustomPipelineWorkerExtension][]
- [vllm_omni.diffusion.worker.diffusion_worker.DiffusionWorker][]
- [vllm_omni.diffusion.worker.diffusion_worker.WorkerProc][]
- [vllm_omni.diffusion.worker.diffusion_worker.WorkerWrapperBase][]
- [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][]
Binary file not shown.

Before

Width:  |  Height:  |  Size: 127 KiB

After

Width:  |  Height:  |  Size: 123 KiB

-9
View File
@@ -1,9 +0,0 @@
# vllm-omni serve
## JSON CLI Arguments
--8<-- "docs/cli/json_tip.inc.md"
## Arguments
--8<-- "docs/generated/argparse_omni/omni_serve.inc.md"
-12
View File
@@ -273,15 +273,3 @@ Default: `True`
Penalty applied to tokens that have already appeared in the generated sequence. Values greater than 1.0 discourage repetition, while values less than 1.0 encourage it. A value of 1.0 applies no penalty.
Default: `1.1`
### `tts_args` (TTS stages only)
Configuration for Text-to-Speech specific parameters. This section is only applicable to TTS model stages (e.g., `qwen3_tts`).
#### `tts_args.max_instructions_length`
Maximum character length for voice style/emotion instructions. Instructions exceeding this limit will be rejected with a validation error.
Default: `500`
This value can be overridden at runtime using the `--tts-max-instructions-length` CLI parameter when starting the server.
@@ -1,4 +1,4 @@
# Cache-DiT
# Support Cache-DiT
This section describes how to add cache-dit acceleration to a new diffusion pipeline. We use the Qwen-Image pipeline and LongCat-Image pipeline as reference implementations.
@@ -1,4 +1,4 @@
# CFG-Parallel
# Support CFG-Parallel
This section describes how to add CFG-Parallel (Classifier-Free Guidance Parallel) to a diffusion pipeline. We use the Qwen-Image pipeline as the reference implementation.
@@ -1,4 +1,4 @@
# Sequence Parallel
# Support Sequence Parallel
This section describes how to add Sequence Parallel (SP) to a diffusion transformer model. We use the Qwen-Image transformer and Wan2.2 transformer as reference implementations.
@@ -333,64 +333,13 @@ _sp_plan = {
**Problem:** `SequenceParallelInput(auto_pad=False)` - auto_pad should be True to enable automatic sequence padding.
**Solution:** In `SequenceParallelInput`, set `auto_pad=True` and add attention mask support.
> **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).
**Constraints of auto_pad:**
| Constraint | Description |
|------------|-------------|
| **Attention Backend Compatibility** | The attention backends must support `attention_mask`. Currently only `TORCH_SDPA` and `FLASH_ATTN` (default for diffusion models) are supported. |
| **Ring Attention Limitation** | Ring attention does not support `attention_mask`. Therefore, when using `auto_pad=True`, the combination of Ulysses + Ring attention is not feasible. |
1. Enable `auto_pad=True` for all sequence-dimension inputs in `_sp_plan`:
**Solution:** In `SequenceParallelInput`, set `auto_pad=True`:
```python
_sp_plan = {
"rope": {
0: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True, auto_pad=True),
1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True, auto_pad=True),
},
"blocks.0": {
"hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, auto_pad=True)
},
...
"blocks.0": {
"hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, auto_pad=True)
}
```
2. Create attention mask dynamically when padding is applied:
```python
from vllm_omni.diffusion.forward_context import get_forward_context
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
# In model forward(), before transformer blocks:
hidden_states_mask = None
ctx = get_forward_context()
if ctx.sp_original_seq_len is not None and ctx.sp_padding_size > 0:
batch_size = hidden_states.shape[0]
padded_seq_len = ctx.sp_original_seq_len + ctx.sp_padding_size
hidden_states_mask = torch.ones(batch_size, padded_seq_len, dtype=torch.bool, device=hidden_states.device)
hidden_states_mask[:, ctx.sp_original_seq_len:] = False
# Pass mask to attention layers
attn_metadata = AttentionMetadata(attn_mask=hidden_states_mask) if hidden_states_mask is not None else None
output = self.attn(query, key, value, attn_metadata)
```
**Important Quality Considerations:**
While `auto_pad` enables generation for irregular resolutions, be aware of potential quality impacts:
| Aspect | Impact |
|--------|--------|
| **Training Distribution** | Models perform best on aspect ratios seen during training (e.g., 1:1, 16:9, 4:3). Unusual ratios like 700x400 (1.75:1) may produce lower quality results. |
| **Padding Overhead** | Padded positions consume compute even when masked. For best efficiency, prefer resolutions divisible by `sp_size`. |
**Recommendations for users:**
- Use standard aspect ratios when possible (e.g., 768x432 for 16:9 instead of 700x400)
- Ensure post-patch dimensions are divisible by `sp_size` for optimal quality
- Test generation quality when using unusual resolutions
### Issue: Inline operations not sharded
**Symptoms:** Some tensors remain full-sized, not sharded.
@@ -1,4 +1,4 @@
# TeaCache
# Support TeaCache
This section describes how to add TeaCache to a diffusion transformer model. We use the Qwen-Image transformer as the reference implementation.
@@ -1,4 +1,4 @@
# Tensor Parallel
# Support Tensor Parallel
This section describes how to add Tensor Parallel (TP) to a diffusion transformer model. We use the Z-Image transformer as the reference implementation.
+18 -26
View File
@@ -1,27 +1,25 @@
# Metrics
# Metrics vLLM-Omni:
You can use these metrics in production to monitor the health and performance of the vLLM-omni system. Typical scenarios include:
- **Performance Monitoring**: Track throughput (e.g., `e2e_avg_tokens_per_s`), latency (e.g., `e2e_total_ms`), and resource utilization to verify that the system meets expected standards.
- **Debugging and Troubleshooting**: Use detailed per-request metrics to diagnose issues, such as high transfer times or unexpected token counts.
## How to Enable and View Metrics
### Start the Service with Metrics Logging
### 1. Start the Service with Metrics Logging
```bash
vllm serve /workspace/models/Qwen3-Omni-30B-A3B-Instruct --omni --port 8014 --log-stats
```
### Send a Request
### 2. Send a Request
```bash
python openai_chat_completion_client_for_multimodal_generation.py --query-type use_image
```
### What You Will See
### 3. What You Will See
With `--log-stats` enabled, the server will output detailed metrics logs after each request. Example output:
@@ -71,13 +69,9 @@ With `--log-stats` enabled, the server will output detailed metrics logs after e
These logs include:
- **Overall summary**: total requests, wall time, average tokens/sec, etc.
- **E2E table**: per-request latency and token counts.
- **Stage table**: per-stage batch and timing details.
- **Transfer table**: data transfer and timing for each edge.
You can use these logs to monitor system health, debug performance, and analyze request-level metrics as described above.
@@ -93,8 +87,6 @@ For **online inference** (serving mode), the summary is always per-request. `e2e
## Parameter Details
### Summary Metrics
| Field | Meaning |
|---------------------------|----------------------------------------------------------------------------------------------|
| `e2e_requests` | Number of completed requests. |
@@ -106,7 +98,7 @@ For **online inference** (serving mode), the summary is always per-request. `e2e
---
### E2E Table (per request)
## E2E Table (per request)
| Field | Meaning |
|---------------------------|-----------------------------------------------------------------------|
@@ -118,7 +110,7 @@ For **online inference** (serving mode), the summary is always per-request. `e2e
---
### Stage Table (per stage event / request)
## Stage Table (per stage event / request)
| Field | Meaning |
|---------------------------|-------------------------------------------------------------------------------------------------|
@@ -133,7 +125,7 @@ For **online inference** (serving mode), the summary is always per-request. `e2e
---
### Transfer Table (per edge / request)
## Transfer Table (per edge / request)
| Field | Meaning |
|----------------------|---------------------------------------------------------------------------|
@@ -143,31 +135,31 @@ For **online inference** (serving mode), the summary is always per-request. `e2e
| `in_flight_time_ms` | In-flight time in ms. |
### Expectation of the Numbers (Verification)
## Expectation of the Numbers (Verification)
**Formulas:**
- `e2e_total_tokens = Stage0's num_tokens_in + sum(all stages' num_tokens_out)`
- `transfers_total_time_ms = sum(tx_time_ms + rx_decode_time_ms + in_flight_time_ms)` for every edge
**Using the example above:**
**e2e_total_tokens**
### e2e_total_tokens
- Stage0's `num_tokens_in`: **4,860**
- Stage0's `num_tokens_out`: **67**
- Stage1's `num_tokens_out`: **275**
- Stage2's `num_tokens_out`: **0**
so `e2e_total_tokens = 4,860 + 67 + 275 + 0 = 5,202`, which matches the table value `e2e_total_tokens`.
**transfers_total_time_ms**
So,
```
e2e_total_tokens = 4,860 + 67 + 275 + 0 = 5,202
```
This matches the table value: `e2e_total_tokens = 5,202`.
### transfers_total_time_ms
For each edge:
- 0->1: tx_time_ms (**78.701**) + rx_decode_time_ms (**111.865**) + in_flight_time_ms (**2.015**) = **192.581**
- 1->2: tx_time_ms (**18.790**) + rx_decode_time_ms (**31.706**) + in_flight_time_ms (**2.819**) = **53.315**
192.581 + 53.315 = **245.896** = transfers_total_time_ms, which matches the calculation (difference is due to rounding)
Sum: 192.581 + 53.315 = **245.896**
The table shows `transfers_total_time_ms = 245.895`, which matches the calculation (difference is due to rounding).
@@ -1,4 +1,4 @@
# Adding a Diffusion Model
# Adding a Diffusion Model to vLLM-Omni
This guide walks you through adding a new diffusion model to vLLM-Omni. We use **Qwen-Image** as the primary example, with references to other models (LongCat, Flux, Wan2.2) to illustrate different patterns.
@@ -680,7 +680,7 @@ vLLM-Omni automatically compiles blocks in `_repeated_blocks` when `torch.compil
### Tensor Parallelism
See detailed guide: [How to add Tensor Parallel support](../../design/feature/tensor_parallel.md)
See detailed guide: [How to add Tensor Parallel support](../features/tensor_parallel.md)
**Quick setup:**
@@ -694,7 +694,7 @@ omni = Omni(model="your-model", tensor_parallel_size=2)
### CFG Parallelism
See detailed guide: [How to add CFG-Parallel support](../../design/feature/cfg_parallel.md)
See detailed guide: [How to add CFG-Parallel support](../features/cfg_parallel.md)
**Quick setup:**
@@ -708,7 +708,7 @@ omni = Omni(model="your-model", cfg_parallel_size=2)
### Sequence Parallelism
See detailed guide: [How to add Sequence Parallel support](../../design/feature/sequence_parallel.md)
See detailed guide: [How to add Sequence Parallel support](../features/sequence_parallel.md)
**Quick setup:**
@@ -724,7 +724,7 @@ omni = Omni(model="your-model", ulysses_degree=2, ring_degree=2)
#### TeaCache
See detailed guide: [How to add TeaCache support](../../design/feature/teacache.md)
See detailed guide: [How to add TeaCache support](../features/teacache.md)
**Quick setup:**
@@ -744,7 +744,7 @@ omni = Omni(model="your-model",
#### Cache-DiT
See detailed guide: [How to add Cache-DiT support](../../design/feature/cache_dit.md)
See detailed guide: [How to add Cache-DiT support](../features/cache_dit.md)
**Quick setup:**
+3
View File
@@ -131,6 +131,9 @@ python image_to_video.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)
> **Note:**
As of now, asynchronous (online) profiling is not fully supported in vLLM-Omni. While start_profile() and stop_profile() methods exist, they are only reliable in offline inference scripts (e.g., the provided end2end.py examples). Do not use them in server-mode or streaming scenarios—traces may be incomplete or fail to flush.
### 4. Analyzing Omni Traces
Output files are saved to your configured ```VLLM_TORCH_PROFILER_DIR```.
-209
View File
@@ -1,209 +0,0 @@
# Async Chunk Design
## Table of Contents
1. [Overview](#overview)
2. [Performance](#performance)
3. [Architecture](#architecture)
4. [Configuration](#configuration)
5. [Related Files](#related-files)
## Overview
The `async_chunk` feature enables asynchronous, chunked processing of data across multiple stages in a multi-stage pipeline (e.g., Qwen3-Omni with Thinker → Talker → Code2Wav stages). Instead of waiting for a complete stage output before forwarding to the next stage, this feature allows stages to process and forward data in chunks as it becomes available, significantly reducing latency and improving throughput.
**Chunk Size Definition**
- **Prefill Phase**: `chunk_size = num_scheduled_tokens` for chunked prefill processing
- **Decode Phase**: `chunk_size = num_scheduled_tokens = 1 ` for per-token streaming
For qwen3-omni:
- **Thinker → Talker**: Per decode step (typically chunk_size=1)
- **Talker → Code2Wav**: Accumulated to code2wav chunk_size(default=25, currently only support default, will support chunk_size soon) before sending
- **Code2Wav**: Streaming decode with code2wav chunk_size
With `async_chunk`:
- Stages can start processing as soon as chunks are available
- Overlapping execution across stages
- Reduced latency and improved throughput
- Better resource utilization
- Async scheduling: Chunk IO (get/put) overlaps with compute via background threads so the scheduler is not blocked waiting for chunks
## Performance
1. **Reduced Latency**: Next stage can start processing immediately
2. **Streaming Support**: Enables streaming for audio generation
3. **IO-Compute Overlap**: Chunk retrieval happens asynchronously while other requests compute
4. **Non-blocking Scheduler**: Requests waiting for chunks don't block the entire scheduler
5. **Code2Wav Batch Inference**: Supports batched processing in code2wav stage
| Input | Output | Async_chunk enabled | Code2Wav batch size | Max_Concurrency | Prompts | Mean E2E | Mean TTFT | Mean TPOT | Mean TTFP | Mean RTF | Mean ITL |
|-----------|------------------|---------------------|---------------------|----------------|---------|------------|------------|-----------|-------------|----------|----------|
| text 100 | text 100+audio | False | 1 | 1 | 50 | 6581.80 | 43.22 | 8.31 | 6459.34 | 0.24 | 8.22 |
| text 100 | text 100+audio | False | 1 | 4 | 50 | 7398.63 | 67.57 | 9.14 | 7285.35 | 0.27 | 9.05 |
| text 100 | text 100+audio | False | 1 | 10 | 50 | 13522.99 | 131.82 | 12.72 | 13410.44 | 0.49 | 12.60 |
| text 100 | text 100+audio | False | 64 | 1 | 50 | 6505.13 | 43.14 | 8.52 | 6395.40 | 0.24 | 8.44 |
| text 100 | text 100+audio | False | 64 | 4 | 50 | 7668.15 | 51.15 | 9.36 | 7562.37 | 0.28 | 9.27 |
| text 100 | text 100+audio | False | 64 | 10 | 50 | 9516.18 | 138.06 | 14.75 | 9409.26 | 0.34 | 14.60 |
| text 100 | text 100+audio | True | 1 | 1 | 50 | 6179.79 | 44.58 | 8.69 | 522.99 | 0.22 | 8.60 |
| text 100 | text 100+audio | True | 1 | 4 | 50 | 7692.69 | 103.96 | 10.22 | 785.85 | 0.29 | 10.12 |
| text 100 | text 100+audio | True | 1 | 10 | 50 | 11152.71 | 685.60 | 17.64 | 1628.88 | 0.41 | 17.62 |
Performance data collected on H800 GPUs through comprehensive benchmarking with cudagraph enabled. text input uses random dataset.
Enabling **async_chunk** (False→True) sharply reduces time-to-first-audio (TTFP)—e.g. ~92% at concurrency 1 (6.5s→0.52s)—and improves E2E latency (e.g. ~6% at conc 1, ~17% at conc 10). RTF (Real Time Factor) also improves with async_chunk on (e.g. ~8% at conc 1: 0.24→0.22, ~16% at conc 10: 0.49→0.41). Enabling **Code2Wav batch size 64** (vs 1) improves E2E and TTFP at higher concurrency when async_chunk is off (e.g. ~30% at conc 10: 13.5s→9.5s E2E, 13.4s→9.4s TTFP).
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/performance/qwen3-omni_ttfp_performance.png">
<img alt="TTFP Performance Data Comparison" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/performance/qwen3-omni_ttfp_performance.png" width=100%>
</picture>
</p>
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/performance/qwen3-omni_rtf_performance.png">
<img alt="RTF Performance Data Comparison" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/performance/qwen3-omni_rtf_performance.png" width=100%>
</picture>
</p>
## Architecture
### Data Flow
#### Sequential Flow
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/qwen3-omni-non-async-chunk.png">
<img alt="Data Flow between stages" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/qwen3-omni-non-async-chunk.png" width=100%>
</picture>
</p>
#### Async Chunk Flow
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/qwen3-omni-async-chunk.png">
<img alt="Data Flow between stages" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/qwen3-omni-async-chunk.png" width=100%>
</picture>
</p>
### Async Chunk architecture
<p align="center">
<picture>
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/async-chunk-architecture.png">
<img alt="Async Chunk Architecture" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/architecture/async-chunk-architecture.png" width=100%>
</picture>
</p>
### Key Components
1. **OmniConnector**: Inter-stage data transport only
- Shared memory or other IPC mechanisms
- **Transport-only API**: `put(from_stage, to_stage, put_key, data)` and `get(from_stage, to_stage, get_key)` (optionally with timeout)
- **No request-specific state**: Connector does not track put_requests, get_requests, request_payload, finished_requests, or other request-bound metadata; it only performs put/get operations
- Chunk keys and request/chunk lifecycle are managed by **OmniChunkTransferAdapter**
2. **Transfer Adapter Layer**: Extensible abstraction for managing data transfer via connectors
- **OmniTransferAdapterBase**: Base class with background **recv_loop** and **save_loop** threads;
- **OmniChunkTransferAdapter**: Chunk-specific implementation that owns the full chunk lifecycle when async_chunk is enabled
- **Chunk ID and key construction**: Builds keys like `{req_id}_{stage_id}_{chunk_id}` for put/get
- **Async get**: `load_async(request)` enqueues the request; background **recv_loop** polls the connector (non-blocking); when data is available, updates the request and marks it in `_finished_load_reqs`; scheduler calls `get_finished_requests()` to learn which requests have chunks ready
- **Async put**: `save_async(pooling_output, request)` invokes `custom_process_next_stage_input_func` in the main thread to build the payload, then enqueues a save task; background **save_loop** performs `connector.put()`; payload processing and chunk accumulation (e.g. code2wav chunk_size) remain in the main thread
3. **Stage Input Processors**: Custom functions that process stage outputs into chunks for different models
- Receive **transfer_manager** (OmniChunkTransferAdapter)
- Qwen3-omni reference: `thinker2talker_async_chunk`, `talker2code2wav_async_chunk`
4. **Schedulers**: Modified to handle chunk-based scheduling with async IO-compute overlap
- `OmniARScheduler`: For autoregressive stages
- `OmniGenerationScheduler`: For generation stages
- Both schedulers use **OmniChunkTransferAdapter** and **before/after** hooks around `super().schedule()`:
- **Before** `super().schedule()`: `process_pending_chunks(waiting, running)` moves requests waiting for chunks to `WAITING_FOR_CHUNK`, enqueues load tasks for background polling
- **After** `super().schedule()`: `restore_queues(waiting, running)` restores requests with ready chunks back to waiting/running, `postprocess_scheduler_output(scheduler_output)` attaches cached additional_information, clears chunk-ready flags
- **put_chunk** `save_async(pooler_output, request)`; **get_chunk** / **get_chunk_for_generation** `load_async(request)`
5. **Model Runners**: Handle chunk processing
- `OmniGPUModelRunner`: Processes chunks in AR stages
- `GPUGenerationModelRunner`: Processes chunks in generation stages
- Uses `ubatch_slices` from `get_forward_context()` to track per-request sequence lengths in batched inference
- Reuses `ubatch_slices_padded` for code2wav batching to properly split batch outputs
- Handles list-type multimodal outputs: iterates through requests and assigns corresponding tensor to each
- Improved request state management: removes unscheduled and finished requests from input batch
6. **Model Implementation**: Model-specific chunk handling
- `Qwen3OmniMoeForConditionalGeneration`: Main model with async_chunk support
- **Code2Wav stage batching**: Uses `ubatch_slices` to construct batched codec codes tensor `[batch_size, 16, max_seq_len]`
- **Batch output handling**: `generate_audio()` returns `list[torch.Tensor]`, one audio tensor per request
- **Multimodal outputs**: Returns list of audio tensors for batch processing instead of single concatenated tensor
- `Qwen3OmniCode2WavDecoder`: Audio generation model
- `chunked_decode()` and `chunked_decode_streaming()`: Return `list[torch.Tensor]` (one per request)
- Uses `ubatch_slices` to split batched waveform output into per-request audio chunks
- Each request gets correctly sized audio based on its code sequence length: `waveform[:, :, :code_seq_len * total_upsample]`
7. **Request status**: `RequestStatus.WAITING_FOR_CHUNK` is added via patch (e.g. in `vllm_omni/patch.py`) so requests waiting for a chunk are not scheduled by the base vLLM scheduler until the chunk is ready.
## Configuration
Enable async_chunk in stage configuration YAML:
```yaml
async_chunk: true
stage_args:
- stage_id: 0
engine_args:
custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_async_chunk
- stage_id: 1
engine_args:
custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav_async_chunk
```
### Stage Configuration
- `async_chunk: bool`: Enable/disable async chunk mode
- `custom_process_next_stage_input_func: str`: Path to custom chunk processing function; receives `(transfer_manager, pooling_output, request)`. For qwen3-omni: `thinker2talker_async_chunk`, `talker2code2wav_async_chunk`
- `stage_connector_config: dict`: Connector configuration
- `worker_type: str`: Model type, e.g. `"ar"` or `"generation"` (used by OmniChunkTransferAdapter for mode-specific payload handling)
- `max_batch_size: int`: Maximum batch size for the stage
### Connector Configuration
```yaml
connectors:
- from_stage: 0
to_stage: 1
spec:
name: SharedMemoryConnector
extra:
stage_id: 0
```
### Code2Wav Batch Configuration
For optimal performance with async_chunk, the code2wav stage should be configured with batching:
```yaml
stage_args:
- stage_id: 2 # code2wav stage
runtime:
devices: "1"
max_batch_size: 64 # Enables batched audio generation
engine_args:
model_stage: code2wav
```
## Related Files
- `vllm_omni/model_executor/stage_input_processors/qwen3_omni.py`: Chunk processing functions (receive `transfer_manager` as first param)
- `vllm_omni/distributed/omni_connectors/transfer_adapter/base.py`: OmniTransferAdapterBase (recv_loop, save_loop, load_async, save_async)
- `vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py`: OmniChunkTransferAdapter (process_pending_chunks, restore_queues, postprocess_scheduler_output)
- `vllm_omni/distributed/omni_connectors/connectors/shm_connector.py`: SharedMemoryConnector (transport-only put/get)
- `vllm_omni/core/sched/omni_ar_scheduler.py`: AR scheduler with chunk_transfer_adapter
- `vllm_omni/core/sched/omni_generation_scheduler.py`: Generation scheduler with same async chunk pattern
- `vllm_omni/worker/gpu_model_runner.py`: Model runner with chunk handling
- `vllm_omni/worker/gpu_generation_model_runner.py`: Generation model runner with batch output handling and ubatch_slices support
- `vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py`: Model implementation with code2wav batching
- `vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py`: Code2wav decoder with batch support
- `vllm_omni/engine/arg_utils.py`: Configuration definitions (async_chunk, worker_type)
- `vllm_omni/config/model.py`: Model config with async_chunk field
@@ -6,8 +6,7 @@ This guide explains how to configure and use distributed connectors
Backend-specific setup lives in separate docs:
- [SharedMemoryConnector](omni_connectors/shared_memory_connector.md)
- [MooncakeStoreConnector](omni_connectors/mooncake_store_connector.md)
- [MooncakeTransferEngineConnector](omni_connectors/mooncake_transfer_engine_connector.md)
- [MooncakeConnector](omni_connectors/mooncake_connector.md)
- [YuanrongConnector](omni_connectors/yuanrong_connector.md)
## Overview
@@ -20,8 +19,7 @@ Current connectors operate in D2H2D (device to host to device) mode.
| Use Case | Recommended Connector | Notes |
| :--- | :--- | :--- |
| Single node | SharedMemoryConnector | Auto-configured if no connector is specified. |
| Multi node (Mooncake Store) | MooncakeStoreConnector | TCP-based, requires Mooncake Master + metadata server. |
| Multi node (Mooncake RDMA) | MooncakeTransferEngineConnector | RDMA/TCP direct transfer with managed memory pool. Fastest. |
| Multi node (Mooncake) | MooncakeConnector | Requires Mooncake Master + metadata server. |
| Multi node (Yuanrong) | YuanrongConnector | Requires Yuanrong Datasystem + etcd. |
## Core API
-146
View File
@@ -1,146 +0,0 @@
# HSDP
This section describes how to add HSDP (Hybrid Sharded Data Parallel) support to a diffusion transformer model. We use the Wan2.2 transformer as the reference implementation.
---
## Table of Contents
- [Overview](#overview)
- [Step-by-Step Implementation](#step-by-step-implementation)
- [Testing](#testing)
- [Reference Implementations](#reference-implementations)
---
## Overview
### What is HSDP?
HSDP (Hybrid Sharded Data Parallel) is a memory optimization technique that **shards model weights** across multiple GPUs using PyTorch's FSDP2. Unlike Tensor Parallelism which splits computation, HSDP:
- Shards weights across GPUs to reduce per-GPU memory usage
- Gathers weights on-demand during forward passes
- Can work standalone or combined with other parallelism (e.g., Sequence Parallel)
This enables inference of large models (e.g., Wan2.2 14B) on GPUs with limited memory.
**Important constraints:**
- HSDP cannot be used with Tensor Parallelism
- For standalone HSDP (no other parallelism), `hsdp_shard_size` must be specified explicitly
### Architecture
HSDP implementation relies on:
1. **`_hsdp_shard_conditions`**: Model attribute specifying which modules to shard
2. **`apply_hsdp_to_model`**: Function that applies FSDP2 sharding
3. **`HSDPInferenceConfig`**: Runtime configuration for HSDP
---
## Step-by-Step Implementation
### Step 1: Identify Modules to Shard
Determine which modules in your transformer should be sharded. Typically, these are:
- Transformer blocks (e.g., `blocks.0`, `blocks.1`, ...)
- Large submodules with significant weight memory
**Key questions:**
- Which modules have the largest weights?
- Which modules are repeated (like transformer blocks)?
### Step 2: Define Shard Conditions
Add `_hsdp_shard_conditions` to your model class. This is a list of functions that return `True` for modules that should be sharded.
**Example (Transformer Blocks):**
```python
class MyTransformerModel(nn.Module):
@staticmethod
def _is_transformer_block(name: str, module) -> bool:
"""Match transformer blocks for HSDP sharding.
Args:
name: Module name from named_modules() (e.g., "blocks.0", "blocks.0.attn")
module: The module instance
Returns:
True if this module should be sharded
"""
return "blocks" in name and name.split(".")[-1].isdigit()
_hsdp_shard_conditions = [_is_transformer_block]
```
**Multiple Conditions Example:**
```python
class MyModel(nn.Module):
@staticmethod
def _is_transformer_block(name: str, module) -> bool:
return "blocks" in name and name.split(".")[-1].isdigit()
@staticmethod
def _is_moe_expert(name: str, module) -> bool:
# Also shard MoE expert layers
return "experts" in name and name.split(".")[-1].isdigit()
# Module is sharded if ANY condition returns True
_hsdp_shard_conditions = [_is_transformer_block, _is_moe_expert]
```
---
## Testing
After adding HSDP support, test with:
```python
from vllm_omni import Omni
from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
parallel_config = DiffusionParallelConfig(
use_hsdp=True,
hsdp_shard_size=8, # Shard across 8 GPUs
)
omni = Omni(model="your-model-name", parallel_config=parallel_config)
output = omni.generate(
"a cup of coffee on the table",
OmniDiffusionSamplingParams(num_inference_steps=50),
)
```
**Or via command line:**
```bash
vllm serve Your-org/your-model --omni --port 8091 --use-hsdp
```
**Verify:**
1. Check logs for "HSDP Inference: replicate_size=..., shard_size=..."
2. Check logs for "Sharded N modules + root"
3. Verify memory usage is reduced proportionally
4. Compare generated output quality with HSDP disabled
---
## Reference Implementations
Complete examples in the codebase:
| Model | Path | Notes |
|-------|------|-------|
| **Wan2.2** | `vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py` | Reference implementation |
| **HSDP Core** | `vllm_omni/diffusion/distributed/hsdp.py` | `apply_hsdp_to_model`, `shard_model` |
| **HSDP Tests** | `tests/diffusion/distributed/test_hsdp.py` | Unit tests |
---
@@ -1,4 +1,4 @@
# MooncakeStoreConnector
# MooncakeConnector
## When to Use
@@ -38,7 +38,7 @@ Define the connector in runtime:
runtime:
connectors:
connector_of_mooncake:
name: MooncakeStoreConnector
name: MooncakeConnector
extra:
host: "127.0.0.1"
metadata_server: "http://<MASTER_IP>:8080/metadata"
@@ -1,207 +0,0 @@
# MooncakeTransferEngineConnector
## When to Use
Best for high-performance multi-node data transfer between stages using Mooncake
Transfer Engine. Supports both RDMA and TCP protocols with a managed memory pool,
zero-copy deserialization, and optional GPUDirect RDMA. Applicable to any
inter-stage data (KV caches, request payloads, etc.), not limited to KV cache transfer.
Compared to `MooncakeStoreConnector` (TCP key-value store), this connector
provides **~60x faster** data transfer via RDMA direct memory access.
## Installation
```bash
pip install mooncake-transfer-engine
```
Ensure RDMA drivers are installed on all nodes (e.g., Mellanox OFED for
InfiniBand/RoCE NICs).
## Configuration
Define the connector in runtime:
```yaml
runtime:
connectors:
rdma_connector:
name: MooncakeTransferEngineConnector
extra:
host: "auto" # Auto-detect local RDMA IP
zmq_port: 50051 # ZMQ base port (see "Port Offset Scheme" below)
protocol: "rdma" # "rdma" or "tcp"
device_name: "" # RDMA device (e.g., "mlx5_0"), empty for auto-detect
memory_pool_size: 2147483648 # 2GB memory pool
memory_pool_device: "cpu" # "cpu" for pinned memory, "cuda" for GPUDirect RDMA
```
Wire stages to the connector:
```yaml
stage_args:
- stage_id: 0
output_connectors:
to_stage_1: rdma_connector
- stage_id: 1
input_connectors:
from_stage_0: rdma_connector
```
## Parameters
### Required
| Parameter | Description |
|---|---|
| `role` | **Internal, do not set manually.** Auto-injected by the orchestration layer (`"sender"` for `output_connectors`, `"receiver"` for `input_connectors`). Defaults to `"sender"` if omitted. |
| `host` | Local IP address for RDMA. `"auto"` detects from network interfaces. |
| `protocol` | Transport protocol: `"rdma"` (InfiniBand/RoCE) or `"tcp"`. |
### Memory Pool
| Parameter | Default | Description |
|---|---|---|
| `memory_pool_size` | 1 GB | Total size of the RDMA-registered memory pool in bytes. |
| `memory_pool_device` | `"cpu"` | `"cpu"`: pinned host memory (recommended). `"cuda"`: GPU VRAM for GPUDirect RDMA (requires NIC-GPU direct PCIe connectivity). |
### Networking
| Parameter | Default | Description |
|---|---|---|
| `zmq_port` | 50051 | ZMQ **base** port. The orchestration layer computes the actual port as `base + purpose_offset + stage_offset` (see table below). Users only set this base value. |
| `sender_host` | `None` | **Internal.** Receiver-side only — dynamically resolved via `update_sender_info()`. Not needed in YAML. |
| `sender_zmq_port` | `None` | **Internal.** Receiver-side only — defaults to the sender's adjusted port. Not needed in YAML. |
| `device_name` | `""` | RDMA device name (e.g., `"mlx5_0"`). Empty for auto-detect. Can also be set via `RDMA_DEVICE_NAME` env var. |
#### ZMQ Port Offset Scheme
To avoid port conflicts when multiple edges, purposes, DP replicas, or TP ranks share the same node, the actual ZMQ port is computed as:
```
side_channel_port = zmq_port + purpose_offset + stage_offset + dp_index * tp_size
sender_listen = side_channel_port + tp_rank
receiver_connect = remote_side_channel_port + tp_rank
```
| Component | Value | Description |
|---|---|---|
| `zmq_port` | 50051 (default) | Base port from YAML config |
| `purpose_offset` | `request_forwarding` = 0, `kv_transfer` = 100 | Separates control-plane vs KV-cache connections |
| `stage_offset` | `int(from_stage)` (0, 1, 2...) | Separates edges from different source stages |
| `dp_index * tp_size` | e.g., DP1 × TP2 = 2 | Each DP replica reserves a port range of size `tp_size` (following vLLM convention: `VLLM_MOONCAKE_BOOTSTRAP_PORT + dp_index * tp_size`) |
| `tp_rank` | 0, 1, 2... | Each TP rank within a DP replica uses its own port |
| orchestrator | +200 | Extra offset when caller is the orchestrator (avoids collision with stage workers on the same node) |
**Example** (base=50051, stage 0→1, DP=2, TP=2, kv_transfer):
| Caller | DP | TP rank | Port |
|---|---|---|---|
| Stage worker | DP0 | rank 0 | `50051 + 100 + 0 + 0×2 + 0 = 50151` |
| Stage worker | DP0 | rank 1 | `50051 + 100 + 0 + 0×2 + 1 = 50152` |
| Stage worker | DP1 | rank 0 | `50051 + 100 + 0 + 1×2 + 0 = 50153` |
| Stage worker | DP1 | rank 1 | `50051 + 100 + 0 + 1×2 + 1 = 50154` |
| Orchestrator | — | — | `50051 + 200 + 0 = 50251` |
## Memory Pool Modes
| Mode | Config | Data Flow | Best For |
|---|---|---|---|
| CPU Pinned | `memory_pool_device: "cpu"` | GPU → CPU pool → RDMA → CPU pool → GPU | Most hardware topologies (recommended) |
| GPUDirect | `memory_pool_device: "cuda"` | GPU → GPU pool → RDMA (NIC reads GPU BAR1) → GPU pool | NIC-GPU direct PCIe (PIX topology) |
> **Note**: GPUDirect RDMA requires the NIC and GPU to share a direct PCIe
> switch (PIX topology). On systems where they are connected via PXB or NODE,
> CPU pinned memory is faster due to GPU BAR1 bandwidth limitations.
## Environment Variables
| Variable | Description |
|---|---|
| `RDMA_DEVICE_NAME` | Override RDMA device name (e.g., `mlx5_0`). |
| `MC_IB_PCI_RELAXED_ORDERING` | Set to `1` to enable PCIe relaxed ordering for GPUDirect. |
## Docker / Container Setup
RDMA requires host-level device access:
```bash
docker run -it \
--cap-add=SYS_PTRACE \
--cap-add=IPC_LOCK \
--security-opt seccomp=unconfined \
--network=host \
--device=/dev/infiniband \
-v /sys/class/infiniband:/sys/class/infiniband:ro \
your-image:tag
```
## Performance
Benchmark results on H800 GPUs with mlx5_0 RDMA NIC (~186 MB KV cache):
| Metric | MooncakeStoreConnector | MooncakeTransferEngineConnector (CPU) |
|---|---|---|
| KV transfer wall time | ~810 ms | **~14 ms** |
| RDMA throughput | N/A (TCP) | ~22 GB/s |
| Speedup | 1x | **58x** |
## Troubleshooting
### Quick Diagnostics
```bash
# 1. Check RDMA devices and link status
ibdev2netdev
# Expected: "mlx5_X port 1 ==> <iface> (Up)"
# RoCE devices map to Ethernet interfaces (e.g., enp75s0f0)
# IB devices map to ib0, ib1, etc.
# 2. Check InfiniBand device details
ibstat
# 3. Verify /dev/infiniband is accessible (required in containers)
ls /dev/infiniband/
# 4. Check Mooncake installation
python -c "from mooncake.engine import TransferEngine; print('OK')"
# 5. Check environment variables
echo "RDMA_DEVICE_NAME=${RDMA_DEVICE_NAME:-<not set>}"
echo "MC_IB_PCI_RELAXED_ORDERING=${MC_IB_PCI_RELAXED_ORDERING:-<not set>}"
```
### Common Issues
| Symptom | Cause | Fix |
|---|---|---|
| `Failed to modify QP to RTR` | Cross-NIC QP handshake failure (multi-NIC DGX) | Set `device_name` to a single RoCE NIC (e.g., `mlx5_2`) or set `RDMA_DEVICE_NAME` env var |
| `transport retry counter exceeded` | RDMA path between incompatible NICs | Same as above — restrict to one NIC |
| `zmq.error.Again: Resource temporarily unavailable` | ZMQ recv timeout (transfer took too long) | Check NIC selection; increase data may need longer timeout |
| `Mooncake Engine initialization failed` | Missing RDMA drivers or `/dev/infiniband` | Install Mellanox OFED; in Docker add `--device=/dev/infiniband` |
| `MemoryError` in allocator | Memory pool too small for payload | Increase `memory_pool_size` |
| GPU transfer slower than CPU | GPU BAR1 bandwidth limitation (PXB/NODE topology) | Use `memory_pool_device: "cpu"` instead of `"cuda"` |
### Multi-NIC Environments (DGX)
On DGX machines with 12+ RDMA NICs, only RoCE NICs (with a bound network
interface) reliably support loopback. IB-only NICs may fail cross-NIC QP
handshakes. To identify RoCE NICs:
```bash
ibdev2netdev | grep -v "ib[0-9]"
# RoCE devices show Ethernet interface names like enp75s0f0
```
Then configure the connector:
```yaml
device_name: "mlx5_2" # or set RDMA_DEVICE_NAME=mlx5_2
```
See the RDMA Test README in tests/distributed/omni_connectors/README.md
for test-specific setup instructions.
For more details on the underlying engine, refer to the
[Mooncake repository](https://github.com/kvcache-ai/Mooncake).
+1 -1
View File
@@ -45,7 +45,7 @@ ray start --address=<HEAD_NODE_IP>:6399
When running on Ray, the system automatically adapts its communication strategy:
* **Cross-Node**: Recommended to use `MooncakeTransferEngineConnector` (RDMA, fastest) or `MooncakeStoreConnector` (TCP fallback).
* **Cross-Node**: Recommended to use `MooncakeConnector` (requires separate configuration).
* **Same-Node**: Can still use `SharedMemoryConnector` for efficiency, or Ray's native object store (plasma).
* **SHM threshold default differs**: when `worker_backend="ray"`, the SharedMemoryConnector default threshold is set to `sys.maxsize`, which forces payloads to go inline (no SHM). Override `shm_threshold_bytes` in the connector config if you want SHM for Ray runs.
+1 -1
View File
@@ -30,7 +30,7 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
### Pre-built wheels
Note: Pre-built wheels are currently only available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0. For the latest version, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
Note: Pre-built wheels are currently only available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0. For the latest version, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
=== "NVIDIA CUDA"
@@ -20,7 +20,7 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new**
vLLM-Omni is built based on vLLM. Please install it with command below.
```bash
uv pip install vllm --torch-backend=auto
uv pip install vllm==0.16.0 --torch-backend=auto
```
#### Installation of vLLM-Omni
@@ -40,7 +40,7 @@ If you do not need to modify source code of vLLM, you can directly install the s
uv pip install vllm==0.16.0 --torch-backend=auto
```
The release 0.16.0 of vLLM is based on PyTorch 2.9.0 which requires CUDA 12.9 environment.
The release 0.14.0 of vLLM is based on PyTorch 2.9.0 which requires CUDA 12.9 environment.
#### Installation of vLLM-Omni
Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -91,7 +91,7 @@ docker run --runtime nvidia --gpus 2 \
--env "HF_TOKEN=$HF_TOKEN" \
-p 8091:8091 \
--ipc=host \
vllm/vllm-omni:v0.16.0 \
vllm/vllm-omni:v0.14.0 \
--model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8091
```
@@ -13,7 +13,7 @@ vLLM-Omni current recommends the steps in under setup through Docker Images.
vLLM-Omni is built based on vLLM. Please install it with command below.
```bash
uv pip install vllm==0.16.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.16.0/rocm700
uv pip install vllm==0.14.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.14.0/rocm700
```
#### Installation of vLLM-Omni
@@ -34,13 +34,13 @@ uv pip install onnxruntime-rocm sox
# --8<-- [start:build-wheel-from-source]
#### Installation of vLLM
If you do not need to modify source code of vLLM, you can directly install the stable 0.16.0 release version of the library
If you do not need to modify source code of vLLM, you can directly install the stable 0.14.0 release version of the library
```bash
uv pip install vllm==0.16.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.16.0/rocm700
uv pip install vllm==0.14.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.14.0/rocm700
```
The release 0.16.0 of vLLM requires ROCm 7.0 environment.
The release 0.14.0 of vLLM requires ROCm 7.0 environment.
#### Installation of vLLM-Omni
Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -58,7 +58,7 @@ If you want to check, modify or debug with source code of vLLM, install the libr
```bash
git clone https://github.com/vllm-project/vllm.git
cd vllm
git checkout v0.16.0
git checkout v0.14.0
python3 -m pip install -r requirements/rocm.txt
python3 setup.py develop
```
@@ -130,7 +130,7 @@ docker run --rm \
-v ~/.cache/huggingface:/root/.cache/huggingface \
--env "HF_TOKEN=$HF_TOKEN" \
-p 8091:8091 \
vllm/vllm-omni-rocm:v0.16.0 \
vllm/vllm-omni-rocm:v0.14.0 \
--model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
```
@@ -149,7 +149,7 @@ docker run --rm -it \
-v ~/.cache/huggingface:/root/.cache/huggingface \
--env "HF_TOKEN=$HF_TOKEN" \
--entrypoint bash \
vllm/vllm-omni-rocm:v0.16.0
vllm/vllm-omni-rocm:v0.14.0
```
# --8<-- [end:pre-built-images]
+2 -6
View File
@@ -24,10 +24,6 @@ vLLM-Omni supports NPU through the vLLM Ascend Plugin (vllm-ascend). This is a c
### Build wheel from source
=== "NPU release"
=== "NPU"
--8<-- "docs/getting_started/installation/npu/npu.inc.md:installation-release"
=== "NPU from main"
--8<-- "docs/getting_started/installation/npu/npu.inc.md:installation-main"
--8<-- "docs/getting_started/installation/npu/npu.inc.md:installation"
@@ -3,11 +3,14 @@
For detailed hardware and software requirements, please refer to the [vllm-ascend installation documentation](https://docs.vllm.ai/projects/ascend/en/latest/installation.html).
# --8<-- [end:requirements]
# --8<-- [start:installation-release]
# --8<-- [start:installation]
The recommended way to use vLLM-Omni on NPU is through the vllm-ascend pre-built Docker images:
```bash
# Update DEVICE according to your NPUs (/dev/davinci[0-7])
export DEVICE0=/dev/davinci0
export DEVICE1=/dev/davinci1
# Update the vllm-ascend image
# Atlas A2:
# export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0rc1
@@ -17,10 +20,8 @@ export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0rc1
docker run --rm \
--name vllm-omni-npu \
--shm-size=1g \
--device /dev/davinci0 \
--device /dev/davinci1 \
--device /dev/davinci2 \
--device /dev/davinci3 \
--device $DEVICE0 \
--device $DEVICE1 \
--device /dev/davinci_manager \
--device /dev/devmm_svm \
--device /dev/hisi_hdc \
@@ -33,18 +34,13 @@ docker run --rm \
-p 8000:8000 \
-it $IMAGE bash
# Because vllm-ascend will release v0.16.0rc1 after vllm-omni 0.16.0,
# we have to pin vllm-ascend at the current commit.
cd /vllm-workspace/vllm-ascend
git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
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.14.0 https://github.com/vllm-project/vllm-omni.git
cd vllm-omni
pip install -v -e .
VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e .
# OR pip install -v -e . --no-build-isolation
export VLLM_WORKER_MULTIPROC_METHOD=spawn
```
@@ -52,42 +48,13 @@ The default workdir is `/workspace`, with vLLM, vLLM-Ascend and vLLM-Omni code p
For other installation methods (pip installation, building from source, custom Docker builds), please refer to the [vllm-ascend installation guide](https://docs.vllm.ai/projects/ascend/en/latest/installation.html).
We are keeping [issue #886](https://github.com/vllm-project/vllm-omni/issues/886) up to date with the aligned versions of vLLM, vLLM-Ascend, and vLLM-Omni, and also outlining the Q1 roadmap there.
We are keeping [issue #997](https://github.com/vllm-project/vllm-omni/issues/997) up to date with the aligned versions of vLLM, vLLM-Ascend, and vLLM-Omni, and also outlining the Q1 roadmap there.
# --8<-- [end:installation-release]
# --8<-- [start:installation-main]
You can also build vLLM-Omni from the latest main branch if you want to use the latest features or bug fixes. (But sometimes it will break for a while. You can check [issue #886](https://github.com/vllm-project/vllm-omni/issues/886) for the status of the latest commit of vLLM-Omni main branch on NPU.)
```bash
# Pin vLLM version to 0.16.0
cd /vllm-workspace/vllm
git pull origin main
git fetch origin --tags
git checkout v0.16.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-ascend
git pull origin main
git fetch origin --tags
git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
pip install -v -e .
# Install vLLM-Omni from the latest main branch
cd ../vllm-omni
git clone https://github.com/vllm-project/vllm-omni.git
pip install -v -e . --no-build-isolation
# or VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e .
export VLLM_WORKER_MULTIPROC_METHOD=spawn
```
# --8<-- [end:installation-main]
# --8<-- [end:installation]
# --8<-- [start:pre-built-images]
vLLM-Omni offers Docker images for Ascend NPU deployment. You can just pull the **prebuilt image** from the image repository [ascend/vllm-omni](https://quay.io/repository/ascend/vllm-omni?tab=tags) and run it with bash.
`vllm-ascend` offers Docker images for deployment. You can just pull the **prebuilt image** from the image repository [ascend/vllm-ascend](https://quay.io/repository/ascend/vllm-ascend?tab=tags) and run it with bash.
Supported images as following.
@@ -96,14 +63,16 @@ Supported images as following.
| image-tag | Atlas A2 | Ubuntu |
| image-tag-a3 | Atlas A3 | Ubuntu |
Here's an example deployment command that has been verified on 4 x NPUs:
vLLM-Omni offers Docker images for Ascend NPU deployment. You can just pull the **prebuilt image** from the image repository [ascend/vllm-ascend](https://quay.io/repository/ascend/vllm-ascend?tab=tags) and run it with bash.
Here's an example deployment command that has been verified on 6 x NPUs:
```bash
# Atlas A2:
# export IMAGE=quay.io/ascend/vllm-omni:v0.16.0
# export IMAGE=quay.io/ascend/vllm-omni:v0.14.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.14.0-a3
export IMAGE=quay.io/ascend/vllm-omni:v0.14.0
docker run --rm \
--name vllm-omni-npu \
--shm-size=1g \
@@ -111,6 +80,8 @@ docker run --rm \
--device /dev/davinci1 \
--device /dev/davinci2 \
--device /dev/davinci3 \
--device /dev/davinci4 \
--device /dev/davinci5 \
--device /dev/davinci_manager \
--device /dev/devmm_svm \
--device /dev/hisi_hdc \
-258
View File
@@ -1,258 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
import ast
import logging
import re
import sys
from argparse import SUPPRESS, Action, ArgumentParser, HelpFormatter, _ArgumentGroup
from collections.abc import Iterable
from importlib.machinery import ModuleSpec
from pathlib import Path
from typing import Literal
from unittest.mock import MagicMock
from pydantic_core import core_schema
class _FlexibleArgumentParser(ArgumentParser):
"""Fallback parser for docs when vllm is unavailable.
Accepts the 'deprecated' kwarg used by vllm CLI and emits warnings
if a deprecated argument is actually provided.
"""
_deprecated: set[Action] = set()
_deprecated_warned: set[str] = set()
if sys.version_info < (3, 13):
def parse_known_args(self, args=None, namespace=None):
namespace, args = super().parse_known_args(args, namespace)
for action in _FlexibleArgumentParser._deprecated:
if (
hasattr(namespace, dest := action.dest)
and getattr(namespace, dest) != action.default
and dest not in _FlexibleArgumentParser._deprecated_warned
):
_FlexibleArgumentParser._deprecated_warned.add(dest)
logger.warning("argument '%s' is deprecated", dest)
return namespace, args
def add_argument(self, *args, **kwargs):
deprecated = kwargs.pop("deprecated", False)
action = super().add_argument(*args, **kwargs)
if deprecated:
_FlexibleArgumentParser._deprecated.add(action)
return action
class _FlexibleArgumentGroup(_ArgumentGroup):
def add_argument(self, *args, **kwargs):
deprecated = kwargs.pop("deprecated", False)
action = super().add_argument(*args, **kwargs)
if deprecated:
_FlexibleArgumentParser._deprecated.add(action)
return action
def add_argument_group(self, *args, **kwargs):
group = self._FlexibleArgumentGroup(self, *args, **kwargs)
self._action_groups.append(group)
return group
logger = logging.getLogger("mkdocs")
# Define root and doc output directories
ROOT_DIR = Path(__file__).parent.parent.parent.parent
ARGPARSE_DOC_DIR = ROOT_DIR / "docs/generated/argparse_omni"
# Ensure the repo root is in sys.path for dynamic imports
sys.path.insert(0, str(ROOT_DIR))
class PydanticMagicMock(MagicMock):
"""`MagicMock` that's able to generate pydantic-core schemas and avoids _mock_name=None errors."""
def __init__(self, *args, **kwargs):
name = kwargs.pop("name", None)
super().__init__(*args, **kwargs)
self.__spec__ = ModuleSpec(name, None)
self._mock_name = name or "mock"
def __get_pydantic_core_schema__(self, source_type, handler):
return core_schema.any_schema()
# --- Static extraction for CLI argument docs ---
def extract_omni_serve_subparser_init():
"""
Statically parse vllm_omni/entrypoints/cli/serve.py to extract the subparser_init method
and return a callable that adds arguments to a parser. This avoids import and mock issues.
"""
serve_path = ROOT_DIR / "vllm_omni" / "entrypoints" / "cli" / "serve.py"
with open(serve_path, encoding="utf-8") as f:
source = f.read()
tree = ast.parse(source, filename=str(serve_path))
# Find class OmniServeCommand
for node in tree.body:
if isinstance(node, ast.ClassDef) and node.name == "OmniServeCommand":
for item in node.body:
if isinstance(item, ast.FunctionDef) and item.name == "subparser_init":
# We'll exec this function body in a dummy context
func_code = ast.Module(body=[item], type_ignores=[])
code = compile(func_code, filename=str(serve_path), mode="exec")
# Prepare dummy context
local_vars = {}
# Provide a dummy subparsers with add_parser returning our parser
class DummySubparsers:
def add_parser(self, name, **kwargs):
return _FlexibleArgumentParser(prog=name)
dummy_subparsers = DummySubparsers()
# Provide globals for exec
# Extract DESCRIPTION from the source file (assume it's a triple-quoted string at the top)
m = re.search(r'DESCRIPTION\s*=\s*([ru]?""".*?""")', source, re.DOTALL)
DESCRIPTION = m.group(1)[3:-3] if m else ""
exec_globals = {
"_FlexibleArgumentParser": _FlexibleArgumentParser,
"FlexibleArgumentParser": _FlexibleArgumentParser,
"make_arg_parser": lambda parser: parser, # no-op for doc
"VLLM_SUBCMD_PARSER_EPILOG": "",
"logger": logger,
"DummySubparsers": DummySubparsers,
"argparse": __import__("argparse"),
"json": __import__("json"),
"DESCRIPTION": DESCRIPTION,
}
exec(code, exec_globals, local_vars)
# Get the function
subparser_init = local_vars["subparser_init"]
# Return a callable that mimics OmniServeCommand().subparser_init
def parser_factory():
class DummySelf:
name = "serve"
return subparser_init(DummySelf(), dummy_subparsers)
return parser_factory
raise RuntimeError("Could not statically extract OmniServeCommand.subparser_init")
OmniServeCommand = type("OmniServeCommand", (), {"subparser_init": staticmethod(extract_omni_serve_subparser_init())})
class MarkdownFormatter(HelpFormatter):
"""Custom formatter that generates markdown for argument groups."""
def __init__(self, prog: str, starting_heading_level: int = 3):
super().__init__(prog, max_help_position=sys.maxsize, width=sys.maxsize)
self._section_heading_prefix = "#" * starting_heading_level
self._argument_heading_prefix = "#" * (starting_heading_level + 1)
self._markdown_output = []
def start_section(self, heading: str):
if heading not in {"positional arguments", "options"}:
heading_md = f"\n{self._section_heading_prefix} {heading}\n\n"
self._markdown_output.append(heading_md)
def end_section(self):
pass
def add_text(self, text: str):
if text:
self._markdown_output.append(f"{text.strip()}\n\n")
def add_usage(self, usage, actions, groups, prefix=None):
pass
def add_arguments(self, actions: Iterable[Action]):
for action in actions:
if len(action.option_strings) == 0 or "--help" in action.option_strings:
continue
option_strings = f"`{'`, `'.join(action.option_strings)}`"
heading_md = f"{self._argument_heading_prefix} {option_strings}\n\n"
self._markdown_output.append(heading_md)
if choices := action.choices:
choices = f"`{'`, `'.join(str(c) for c in choices)}`"
self._markdown_output.append(f"Possible choices: {choices}\n\n")
elif (metavar := action.metavar) and isinstance(metavar, (list, tuple)):
metavar = f"`{'`, `'.join(str(m) for m in metavar)}`"
self._markdown_output.append(f"Possible choices: {metavar}\n\n")
if action.help:
self._markdown_output.append(f"{action.help}\n\n")
if (default := action.default) != SUPPRESS:
# Make empty string defaults visible
if default == "":
default = '""'
self._markdown_output.append(f"Default: `{default}`\n\n")
def format_help(self):
"""Return the formatted help as markdown."""
return "".join(self._markdown_output)
# Function to create parser using subparser_init style CLI class
def create_parser_subparser_init(subcmd_class):
"""
Create an argparse parser using subparser_init style CLI class, with MarkdownFormatter.
Compatible with both AST-extracted static parser_factory and real class instance method.
"""
class DummySubparsers:
def add_parser(self, name, **kwargs):
return _FlexibleArgumentParser(prog=name)
dummy_subparsers = DummySubparsers()
# If subparser_init is a zero-arg static function (AST-extracted), call it directly
subparser_init = getattr(subcmd_class, "subparser_init", None)
if subparser_init is not None:
import inspect
sig = inspect.signature(subparser_init)
if len(sig.parameters) == 0:
parser = subparser_init()
else:
parser = subcmd_class().subparser_init(dummy_subparsers)
else:
parser = subcmd_class().subparser_init(dummy_subparsers)
parser.formatter_class = MarkdownFormatter
return parser
def on_startup(command: Literal["build", "gh-deploy", "serve"], dirty: bool):
"""
Entry point for doc generation. Builds doc directory and outputs markdown for each CLI command.
"""
logger.info("Generating vllm-omni argparse documentation")
logger.debug("Root directory: %s", ROOT_DIR.resolve())
logger.debug("Output directory: %s", ARGPARSE_DOC_DIR.resolve())
if not ARGPARSE_DOC_DIR.exists():
ARGPARSE_DOC_DIR.mkdir(parents=True)
# Register all CLI parsers; you can easily add more commands here
parsers = {
"omni_serve": create_parser_subparser_init(OmniServeCommand),
# "another_cmd": create_parser_subparser_init(AnotherCommandClass),
}
for stem, parser in parsers.items():
doc_path = ARGPARSE_DOC_DIR / f"{stem}.inc.md"
with open(doc_path, "w", encoding="utf-8") as f:
f.write(super(type(parser), parser).format_help())
logger.info("Argparse generated: %s", doc_path.relative_to(ROOT_DIR))
if __name__ == "__main__":
logging.basicConfig(level=logging.INFO)
on_startup("build", False)
-7
View File
@@ -36,14 +36,10 @@ th {
|`StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` |
|`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` |
## List of Supported Models for NPU
@@ -66,10 +62,7 @@ th {
| `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.2-T2V, Wan2.2-TI2V | `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` |
Binary file not shown.

Before

Width:  |  Height:  |  Size: 134 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 194 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 257 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 36 KiB

Binary file not shown.

Before

Width:  |  Height:  |  Size: 33 KiB

-63
View File
@@ -1,63 +0,0 @@
# LoRA (Low-Rank Adaptation) Guide
LoRA (Low-Rank Adaptation) enables fine-tuning diffusion models by adding trainable low-rank matrices to existing model weights. vLLM-Omni currently supports PEFT-style LoRA adapters, allowing you to customize model behavior without modifying the base model weights.
## Overview
LoRA adapters are lightweight, model-specific fine-tuning weights that can be dynamically loaded and applied to diffusion models. vLLM-Omni uses a unified LoRA handling mechanism similar to vLLM with LRU cache management.
## LoRA Adapter Format
LoRA adapters must be in **PEFT (Parameter-Efficient Fine-Tuning)** format. A typical LoRA adapter directory structure:
```
lora_adapter/
├── adapter_config.json
└── adapter_model.safetensors
```
The `adapter_config.json` file contains metadata about the LoRA adapter, including:
- `r`: LoRA rank
- `lora_alpha`: LoRA alpha scaling factor
- `target_modules`: List of module names to apply LoRA to
## Quick Start
### Offline Inference
#### Pre-loaded LoRA
Load a LoRA adapter at initialization. This adapter is pre-loaded into the cache and can be activated by requests:
```python
from vllm_omni import Omni
from vllm_omni.lora.request import LoRARequest
lora_path="/path/to/lora_adapter"
omni = Omni(
model="stabilityai/stable-diffusion-3.5-medium",
lora_path=lora_path
)
lora_request = LoRARequest(
lora_name="preloaded",
lora_int_id=1,
lora_path=lora_path
)
outputs = omni.generate(
prompt="A piece of cheesecake",
lora_request=lora_request,
lora_scale=2.0, # optional arg, default 1.0
)
```
!!! note "Server-side Path Requirement"
The LoRA adapter path (`local_path`) must be readable on the **server** machine. If your client and server are on different machines, ensure the LoRA adapter is accessible via a shared mount or copied to the server.
## See Also
- [Text-to-Image Offline Example](../examples/offline_inference/text_to_image.md#lora) - Complete offline LoRA example
- [Text-to-Image Online Example](../examples/online_serving/text_to_image.md#lora) - Complete online LoRA example
@@ -16,8 +16,6 @@ The following parallelism methods are currently supported in vLLM-Omni:
5. [VAE Patch Parallelism](#vae-patch-parallelism): VAE patch parallelism shards VAE decode spatially across ranks. This can reduce the peak memory of VAE decode and (depending on resolution and communication overhead) speed up VAE decode.
6. [HSDP](#hsdp): Hybrid Sharded Data Parallel shards model weights across GPUs using PyTorch FSDP2. This reduces per-GPU memory usage, enabling inference of large models on GPUs with limited memory.
The following table shows which models are currently supported by parallelism method:
### ImageGen
@@ -32,7 +30,7 @@ The following table shows which models are currently supported by parallelism me
| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ | ❌ |
| **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` | ❌ | ❌ | ❌ | ✅ | ✅ |
| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ❌ | ❌ | ✅ | ❌ |
| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ❌ | ❌ | ❌ | ✅ | ❌ |
| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ❌ | ✅ | ✅ | ❌ |
@@ -51,9 +49,9 @@ The following table shows which models are currently supported by parallelism me
### VideoGen
| Model | Model Identifier | Ulysses-SP | Ring-Attention | Tensor-Parallel | HSDP |
|-------|------------------|:----------:|:--------------:|:---------------:|:----:|
| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ✅ | ✅ | ✅ | ✅ |
| Model | Model Identifier | Ulysses-SP | Ring-Attention | Tensor-Parallel |
|-------|------------------|------------|---------|--------------------------|
| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ✅ | ✅ | ✅ |
### Tensor Parallelism
@@ -85,7 +83,7 @@ outputs = omni.generate(
VAE patch parallelism distributes the VAE decode workload across multiple ranks by splitting the latent spatially. It is configured via `DiffusionParallelConfig.vae_patch_parallel_size` and can be combined with other parallelism methods (e.g., TP).
!!! note "Enablement and feature gate"
- VAE patch parallelism is currently **enabled only for validated pipelines** (currently: `Tongyi-MAI/Z-Image-Turbo`, `stabilityai/stable-diffusion-3.5`).
- VAE patch parallelism is currently **enabled only for validated pipelines** (currently: `Tongyi-MAI/Z-Image-Turbo`).
- If `vae_patch_parallel_size > 1` is set for a validated pipeline, vLLM-Omni will automatically enable `vae_use_tiling` as a safety gate. (We use `vae_use_tiling` because it indicates the VAE supports diffusers tiling parameters like `tile_latent_min_size` and `tile_overlap_factor`.)
#### Offline Inference
@@ -332,84 +330,3 @@ You can enable CFG-Parallel in online serving for diffusion models via `--cfg-pa
```bash
vllm serve Qwen/Qwen-Image-Edit --omni --port 8091 --cfg-parallel-size 2
```
### HSDP
HSDP (Hybrid Sharded Data Parallel) shards model weights across GPUs to reduce per-GPU memory usage. This enables inference of large models (e.g., Wan2.2 14B) on GPUs with limited memory.
Unlike Tensor Parallelism which splits computation, HSDP uses PyTorch's FSDP2 to shard and redistribute weights at runtime. Each GPU only holds a fraction of the model weights, and weights are gathered on-demand during forward passes.
#### Configuration
HSDP is configured via `DiffusionParallelConfig`:
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `use_hsdp` | bool | False | Enable HSDP |
| `hsdp_shard_size` | int | -1 | Number of GPUs to shard weights across. -1 = auto (requires other parallelism > 1) |
| `hsdp_replicate_size` | int | 1 | Number of replica groups. Each group holds a full sharded copy |
**Constraints:**
- `hsdp_replicate_size × hsdp_shard_size == world_size`
- HSDP cannot be used with Tensor Parallelism (`tensor_parallel_size` must be 1)
#### Operating Modes
HSDP can work in two modes:
- **Standalone Mode**: HSDP alone without other parallelism. Must specify `hsdp_shard_size` explicitly.
- **Combined Mode**: HSDP overlays on top of other parallelism (Ulysses Sequence Parallel, CFG Parallel). HSDP dimensions must match world_size.
#### Offline Inference
**Standalone HSDP** (shard across 4 GPUs, no other parallelism):
```python
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.diffusion.data import DiffusionParallelConfig
omni = Omni(
model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
parallel_config=DiffusionParallelConfig(
use_hsdp=True,
hsdp_shard_size=4, # Shard across 4 GPUs
),
)
outputs = omni.generate(
"A cat playing piano",
OmniDiffusionSamplingParams(num_inference_steps=50),
)
```
**Combined HSDP + Sequence Parallel**:
```python
omni = Omni(
model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
parallel_config=DiffusionParallelConfig(
ulysses_degree=4, # Sequence parallel
use_hsdp=True, # HSDP overlays on SP
),
)
```
#### Online Serving
**Standalone HSDP** (shard model across 4 GPUs):
```bash
vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --port 8091 --use-hsdp --hsdp-shard-size 4
```
**Combined with Sequence Parallel**:
```bash
vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --port 8091 --use-hsdp --usp 4
```
#### Adding HSDP Support to New Models
For detailed instructions on adding HSDP support to new models, see the [HSDP Contributing Guide](../../design/feature/hsdp.md).
@@ -1,185 +0,0 @@
# GGUF Quantization
## Goals
1. Reuse vLLM quantization configs and weight loaders as much as possible.
2. Add native GGUF support to diffusion transformers without changing model definitions.
3. Keep user-facing knobs minimal and consistent across offline and online flows.
## Scope
1. Models: Z-Image, and Flux2-klein.
2. Components: diffusion transformer weights, loader paths, and quantization configs.
3. Modes: native GGUF (transformer-only weights).
## Architecture Overview
1. `OmniDiffusionConfig` accepts `quantization` or `quantization_config`.
2. Diffusion quantization wrapper (`DiffusionGgufConfig`) produces vLLM `QuantizationConfig` objects for linear layers.
3. `DiffusersPipelineLoader` branches on quantization method and loads either HF weights or GGUF weights for the transformer.
4. GGUF transformer loading is routed through model-specific adapters (e.g., Flux2Klein).
5. vLLM GGUF path uses `GGUFConfig` and `GGUFLinearMethod` for matmul.
## Call Chain (Offline)
```
CLI (examples/offline_inference/text_to_image/text_to_image.py)
|
v
Omni (vllm_omni/entrypoints/omni.py)
|
v
OmniStage (diffusion)
|
v
DiffusionWorker
|
v
DiffusionModelRunner
|
v
DiffusersPipelineLoader
|
v
Pipeline.forward (Flux2/Qwen/Z-Image)
|
v
DiffusionEngine
|
v
OmniRequestOutput
|
v
Client (saved PNG)
```
## Call Chain (Online)
```
Client
|
| POST /v1/images/generations
v
APIServer (vllm_omni/entrypoints/openai/api_server.py)
|
v
_generate_with_async_omni
|
v
AsyncOmni
|
v
DiffusionEngine
|
v
OmniRequestOutput
|
v
encode_image_base64
|
v
ImageGenerationResponse
|
v
Client
```
## Call Chain (GGUF Operator Path)
```
Pipeline.forward (Flux2/Qwen/Z-Image)
|
v
Transformer blocks
|
v
QKVParallelLinear / ColumnParallelLinear / RowParallelLinear
|
v
LinearBase.forward
|
v
QuantMethod.apply (GGUFLinearMethod.apply)
|
v
fused_mul_mat_gguf
|
v
_fused_mul_mat_gguf (custom op)
|
v
ops.ggml_dequantize
|
v
x @ weight.T
```
## GGUF Weight Loading Path (Transformer-Only)
1. `DiffusersPipelineLoader.load_model` detects `quantization_config.method == "gguf"`.
2. `gguf_model` is resolved as one of: local file, `repo/file.gguf`, or `repo:quant_type`.
3. GGUF weights are routed through adapters in `vllm_omni/diffusion/model_loader/gguf_adapters/`.
4. Name mapping is applied per-architecture (Z-Image, Flux2Klein).
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).
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`
## User Usage (Offline)
### Baseline BF16
```bash
python examples/offline_inference/text_to_image/text_to_image.py \
--model /workspace/models/black-forest-labs/FLUX.2-klein-4B \
--prompt "a photo of a forest with mist swirling around the tree trunks. The word 'FLUX.2' is painted over it in big, red brush strokes with visible texture" \
--height 768 \
--width 1360 \
--seed 42 \
--cfg_scale 4.0 \
--num_images_per_prompt 1 \
--num_inference_steps 4 \
--output outputs/flux2_klein_4b.png
```
### Native GGUF (Transformer Only)
```bash
python examples/offline_inference/text_to_image/text_to_image.py \
--model /workspace/models/black-forest-labs/FLUX.2-klein-4B \
--gguf-model "/workspace/models/unsloth/FLUX.2-klein-4B-GGUF/flux-2-klein-4b-Q8_0.gguf" \
--quantization gguf \
--prompt "a photo of a forest with mist swirling around the tree trunks. The word 'FLUX.2' is painted over it in big, red brush strokes with visible texture" \
--height 768 \
--width 1360 \
--seed 42 \
--cfg_scale 4.0 \
--num_images_per_prompt 1 \
--num_inference_steps 4 \
--output outputs/flux2_klein_4b_gguf.png
```
Notes for GGUF:
1. Many GGUF repos do not ship `model_index.json` and configs. Use the base repo for `--model` and only pass the GGUF file via `--gguf-model`.
2. `gguf_model` supports local path, `repo/file.gguf`, or `repo:quant_type`.
## User Usage (Online)
### Start Server (Native GGUF via CLI)
```bash
vllm serve /workspace/models/black-forest-labs/FLUX.2-klein-4B \
--omni \
--port 8000 \
--quantization-config '{"method":"gguf","gguf_model":"/workspace/models/unsloth/FLUX.2-klein-4B-GGUF/flux-2-klein-4b-Q8_0.gguf"}'
```
### Online Request (Images API)
```bash
curl -X POST http://localhost:8000/v1/images/generations \
-H "Content-Type: application/json" \
-d '{
"prompt": "a dragon laying over the spine of the Green Mountains of Vermont",
"size": "1024x1024",
"seed": 42,
"num_inference_steps": 4
}'
```
@@ -7,7 +7,6 @@ vLLM-Omni supports quantization of DiT linear layers to reduce memory usage and
| Method | Guide |
|--------|-------|
| FP8 | [FP8](fp8.md) |
| GGUF | [GGUF](gguf.md) |
## Device Compatibility
+3 -31
View File
@@ -30,8 +30,6 @@ vLLM-Omni also supports parallelism methods for diffusion models, including:
5. [VAE Patch Parallelism](diffusion/parallelism_acceleration.md#vae-patch-parallelism) - shards VAE decode/encode spatially across ranks to reduce VAE peak memory (and can speed up VAE decode).
6. [HSDP](diffusion/parallelism_acceleration.md#hsdp) - Hybrid Sharded Data Parallel shards model weights across GPUs to reduce per-GPU memory usage, enabling inference of large models on limited GPU memory.
## Quick Comparison
### Cache Methods
@@ -67,13 +65,12 @@ The following table shows which models are currently supported by each accelerat
| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ |
| **Bagel** | `ByteDance-Seed/BAGEL-7B-MoT` | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ |
| **NextStep-1.1** | `stepfun-ai/NextStep-1.1` | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ |
### VideoGen
| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention | CFG-Parallel | HSDP |
|-------|------------------|:--------:|:---------:|:----------:|:--------------:|:------------:|:----:|
| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ |
| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention | CFG-Parallel |
|-------|------------------|:--------:|:---------:|:----------:|:--------------:|:----------------:|
| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ |
### Quantization
@@ -267,30 +264,6 @@ outputs = omni.generate(
)
```
### Using HSDP
HSDP (Hybrid Sharded Data Parallel) shards model weights across GPUs to reduce per-GPU memory usage. This enables inference of large models (e.g., Wan2.2 14B) on GPUs with limited memory.
```python
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.diffusion.data import DiffusionParallelConfig
omni = Omni(
model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
parallel_config=DiffusionParallelConfig(
use_hsdp=True,
hsdp_replicate_size=1, # Number of replica groups
hsdp_shard_size=8, # Number of GPUs to shard across (or -1 for auto)
),
)
outputs = omni.generate(
"A cat playing piano",
OmniDiffusionSamplingParams(num_inference_steps=50),
)
```
### Using CFG-Parallel
Run image-to-image:
@@ -345,4 +318,3 @@ For detailed information on each acceleration method:
- **[Sequence Parallelism](diffusion/parallelism_acceleration.md#sequence-parallelism)** - Guidance on how to set sequence parallelism with configuration.
- **[CFG-Parallel](diffusion/parallelism_acceleration.md#cfg-parallel)** - Guidance on how to set CFG-Parallel to run positive/negative branches across ranks.
- **[VAE Patch Parallelism](diffusion/parallelism_acceleration.md#vae-patch-parallelism)** - Guidance on how to reduce VAE memory via patch/tile parallelism.
- **[HSDP](diffusion/parallelism_acceleration.md#hsdp)** - Hybrid Sharded Data Parallel for memory-efficient inference of large models.
@@ -154,24 +154,6 @@ The default yaml configuration deploys Thinker and DiT on the same GPU. You can
------
#### Tensor Parallelism (TP)
For larger models or multi-GPU environments, you can enable Tensor Parallelism (TP) by modifying the stage configuration (e.g., [`bagel.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/model_executor/stage_configs/bagel.yaml)).
1. **Set `tensor_parallel_size`**: Increase this value (e.g., to `2` or `4`).
2. **Set `devices`**: Specify the comma-separated GPU IDs to be used for the stage (e.g., `"0,1"`).
Example configuration for TP=2 on GPUs 0 and 1:
```yaml
engine_args:
tensor_parallel_size: 2
...
runtime:
devices: "0,1"
```
------
#### 🔗 Runtime Configuration
| Parameter | Value | Description |
@@ -180,74 +162,6 @@ Example configuration for TP=2 on GPUs 0 and 1:
| `max_inflight` | `1` | Maximum inflight requests |
| `shm_threshold_bytes` | `65536` | Shared memory threshold (64KB) |
## Using Mooncake Connector
[Mooncake](https://github.com/kvcache-ai/Mooncake) is a high-performance distributed KV cache transfer engine that enables efficient cross-node data movement via TCP or RDMA, making it ideal for multi-node disaggregated inference.
By default, BAGEL uses `SharedMemoryConnector` for inter-stage communication. You can switch to the Mooncake connector for better performance on multi-GPU setups and to enable multi-node deployment.
### Prerequisites
Install the Mooncake transfer engine:
```bash
# For CUDA-enabled systems (recommended)
pip install mooncake-transfer-engine
# For non-CUDA systems
pip install mooncake-transfer-engine-non-cuda
```
### Step 1: Start the Mooncake Master
On the **primary node**, start the Mooncake master service (run in a separate terminal or background with `&`):
```bash
# Optional: enable disk-backed storage by creating a directory and passing --root_fs_dir.
# Without it, Mooncake runs in memory-only mode, which is sufficient for KV cache transfer.
mkdir -p ./mc_storage
mooncake_master \
--rpc_port=50051 \
--enable_http_metadata_server=true \
--http_metadata_server_host=0.0.0.0 \
--http_metadata_server_port=8080 \
--metrics_port=9003 \
--root_fs_dir=./mc_storage/ \
--cluster_id=mc-local-1 &
```
### Step 2: Run Offline Inference with Mooncake
Use the provided Mooncake stage config [`bagel_multiconnector.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml). Before launching, update the `metadata_server` and `master` addresses in the YAML to match your Mooncake master node's IP (use `127.0.0.1` for single-node testing).
```bash
cd examples/offline_inference/bagel
# Text to Image with Mooncake
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2img \
--prompts "A cute cat" \
--stage-configs-path ../../../vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
# Image to Text with Mooncake
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality img2text \
--image-path /path/to/image.jpg \
--prompts "Describe this image" \
--stage-configs-path ../../../vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
# Text to Text with Mooncake
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2text \
--prompts "What is the capital of France?" \
--stage-configs-path ../../../vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
```
For more details on the Mooncake connector and multi-node setup, see the [Mooncake Store Connector documentation](https://github.com/vllm-project/vllm-omni/tree/main/docs/design/feature/omni_connectors/mooncake_store_connector.md).
------
## FAQ
- If you encounter an error about the backend of librosa, try to install ffmpeg with the command below.
@@ -1,156 +0,0 @@
# GLM-Image Multistage End-to-End Inference
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/glm_image>.
This example demonstrates how to run GLM-Image with the vLLM-Omni multistage architecture.
## Architecture
GLM-Image uses a 2-stage pipeline:
```
┌─────────────────────────────────────────────────────────────┐
│ GLM-Image Pipeline │
├─────────────────────────────────────────────────────────────┤
│ │
│ Stage 0 (AR Model) Stage 1 (Diffusion) │
│ ┌─────────────────┐ ┌─────────────────────┐ │
│ │ vLLM-optimized │ │ GlmImagePipeline │ │
│ │ GlmImageFor │ prior │ ┌───────────────┐ │ │
│ │ Conditional │──tokens───►│ │ DiT Denoiser │ │ │
│ │ Generation │ │ └───────────────┘ │ │
│ │ (9B AR model) │ │ │ │ │
│ └─────────────────┘ │ ▼ │ │
│ ▲ │ ┌───────────────┐ │ │
│ │ │ │ VAE Decode │──┼──► Image
│ Text/Image │ └───────────────┘ │ │
│ Input └─────────────────────┘ │
│ │
└─────────────────────────────────────────────────────────────┘
```
## Features
- **vLLM-optimized AR**: Uses PagedAttention and tensor parallelism for faster prior token generation
- **Flexible deployment**: AR and Diffusion stages can run on different GPUs
- **Text-to-Image**: Generate images from text descriptions
- **Image-to-Image**: Edit existing images with text prompts
## Usage
### Text-to-Image
```bash
python end2end.py \
--model-path /path/to/glm-image \
--config-path ../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
--prompt "A beautiful sunset over the ocean with sailing boats" \
--height 1024 \
--width 1024 \
--output output_t2i.png
```
### Image-to-Image (Image Editing)
```bash
python end2end.py \
--model-path /path/to/glm-image \
--config-path ../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
--prompt "Transform this scene into a winter wonderland" \
--image input.png \
--output output_i2i.png
```
### With Custom Parameters
```bash
python end2end.py \
--model-path /path/to/glm-image \
--config-path ../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
--prompt "A photorealistic cat sitting on a window sill" \
--height 1024 \
--width 1024 \
--num-inference-steps 50 \
--guidance-scale 1.5 \
--seed 42 \
--output output.png
```
## Shell Scripts
### Run Text-to-Image
```bash
./run_t2i.sh
```
### Run Image-to-Image
```bash
./run_i2i.sh --image /path/to/input.png
```
## Stage Configuration
The stage config (`glm_image.yaml`) defines:
- **Stage 0 (AR)**: Uses `GPUARWorker` with vLLM engine
- Model: `GlmImageForConditionalGeneration`
- Output: `token_ids` (prior tokens)
- **Stage 1 (Diffusion)**: Uses diffusion engine
- Model: `GlmImagePipeline`
- Output: Generated image
See `vllm_omni/model_executor/stage_configs/glm_image.yaml` for full configuration.
## Comparison with Single-Stage
| Aspect | Single-Stage (transformers) | Multistage (vLLM) |
| ----------- | --------------------------- | ------------------- |
| AR Model | transformers native | vLLM PagedAttention |
| Memory | Higher (no KV cache opt) | Lower (optimized) |
| Throughput | Lower | Higher |
| Flexibility | Single GPU | Multi-GPU support |
## Troubleshooting
### OOM Error
Try reducing memory usage:
```bash
# In glm_image.yaml, adjust:
gpu_memory_utilization: 0.5 # Reduce from 0.6
```
### Slow Initialization
The first run loads model weights. Subsequent runs are faster:
```bash
--stage-init-timeout 900 # Increase timeout for slow storage
```
## Requirements
- vLLM-Omni with GLM-Image support
- CUDA-capable GPU (recommended: H100/A100 with 80GB)
- GLM-Image model weights
## Example materials
??? abstract "end2end.py"
``````py
--8<-- "examples/offline_inference/glm_image/end2end.py"
``````
??? abstract "run_i2i.sh"
``````sh
--8<-- "examples/offline_inference/glm_image/run_i2i.sh"
``````
??? abstract "run_t2i.sh"
``````sh
--8<-- "examples/offline_inference/glm_image/run_t2i.sh"
``````
@@ -69,11 +69,6 @@ Key arguments:
- `--vae-use-tiling`: Enable VAE tiling for memory optimization.
- `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](https://github.com/vllm-project/vllm-omni/tree/main/docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
- `--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.
- `--hsdp-replicate-size`: Number of replica groups for HSDP. Each replica holds a full sharded copy. Default 1 means pure sharding (no replication).
> ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.
@@ -0,0 +1,108 @@
# LoRA Inference(Diffusion)
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/lora_inference>.
This directory contains examples for using LoRA (Low-Rank Adaptation) adapters with vLLM-omni diffusion models for offline inference.
The example uses the `stabilityai/stable-diffusion-3.5-medium` as the default model, but you can replace it with other models in vLLM-omni.
## Overview
Similar to vLLM, vLLM-omni uses a unified LoRA handling mechanism:
- **Pre-loaded LoRA**: Loaded at initialization via `--lora-path` (pre-loaded into cache)
- **Per-request LoRA**: Loaded on-demand. In the example, the LoRA is loaded via `--lora-request-path` in each request
Both approaches use the same underlying mechanism - all LoRA adapters are handled uniformly through `set_active_adapter()`. If no LoRA request is provided in a request, all adapters are deactivated.
## Usage
### Pre-loaded LoRA (via --lora-path)
Load a LoRA adapter at initialization. This adapter is pre-loaded into the cache and can be activated by requests:
```bash
python -m examples.offline_inference.lora_inference.lora_inference \
--prompt "A piece of cheesecake" \
--lora-path /path/to/lora/ \
--lora-scale 1.0 \
--num-inference-steps 50 \
--height 1024 \
--width 1024 \
--output output_preloaded.png
```
**Note**: When using `--lora-path`, the adapter is loaded at init time with a stable ID derived from the adapter path. This example activates it automatically for the request.
### Per-request LoRA (via --lora-request-path)
Load a LoRA adapter on-demand for each request:
```bash
python -m examples.offline_inference.lora_inference.lora_inference \
--prompt "A piece of cheesecake" \
--lora-request-path /path/to/lora/ \
--lora-scale 1.0 \
--num-inference-steps 50 \
--height 1024 \
--width 1024 \
--output output_per_request.png
```
### No LoRA
If no LoRA request is provided, we will use the base model without any LoRA adapters:
```bash
python -m examples.offline_inference.lora_inference.lora_inference \
--prompt "A piece of cheesecake" \
--num-inference-steps 50 \
--height 1024 \
--width 1024 \
--output output_no_lora.png
```
## Parameters
### LoRA Parameters
- `--lora-path`: Path to LoRA adapter folder to pre-load at initialization (loads into cache with a stable ID derived from the path)
- `--lora-request-path`: Path to LoRA adapter folder for per-request loading
- `--lora-request-id`: Integer ID for the LoRA adapter (optional). If not provided and `--lora-request-path` is set, will derive a stable ID from the path.
- `--lora-scale`: Scale factor for LoRA weights (default: 1.0). Higher values increase the influence of the LoRA adapter.
### Standard Parameters
- `--prompt`: Text prompt for image generation (required)
- `--seed`: Random seed for reproducibility (default: 42)
- `--height`: Image height in pixels (default: 1024)
- `--width`: Image width in pixels (default: 1024)
- `--num-inference-steps`: Number of denoising steps (default: 50)
- `--output`: Output file path (default: `lora_output.png`)
## How LoRA Works
All LoRA adapters are handled uniformly:
1. **Initialization**: If `--lora-path` is provided, the adapter is loaded into cache with a stable ID derived from the adapter path
2. **Per-request**: If `--lora-request-path` is provided, the adapter is loaded/activated for that request
3. **No LoRA**: If no LoRA request is provided (`req.lora_request` is None), all adapters are deactivated
The system uses LRU cache management - adapters are cached and evicted when the cache is full (unless pinned).
## LoRA Adapter Format
LoRA adapters must be in PEFT (Parameter-Efficient Fine-Tuning) format. A typical LoRA adapter directory structure:
```
lora_adapter/
├── adapter_config.json
└── adapter_model.safetensors
```
## Example materials
??? abstract "lora_inference.py"
``````py
--8<-- "examples/offline_inference/lora_inference/lora_inference.py"
``````
@@ -90,43 +90,13 @@ Examples:
python end2end.py --query-type Base --mode-tag icl
```
## Streaming Mode
Add `--streaming` to stream audio chunks progressively via `AsyncOmni` (requires `async_chunk: true` in the stage config):
```bash
python end2end.py --query-type CustomVoice --streaming --output-dir /tmp/out_stream
```
Each 25-frame Code2Wav chunk is logged as it arrives. The final WAV file is written once generation
completes. This demonstrates that audio data is available progressively rather than only at the end.
> **Note:** Streaming uses `AsyncOmni` internally. The non-streaming path (`Omni`) is unchanged.
## Batched Decoding
The Code2Wav stage (stage 1) supports batched decoding, where multiple requests are decoded in a single forward pass through the SpeechTokenizer. To use it, provide a stage config with `max_batch_size > 1` and pass multiple prompts via `--txt-prompts` with a matching `--batch-size`.
```
python end2end.py --query-type CustomVoice \
--txt-prompts benchmark_prompts.txt \
--batch-size 4 \
--stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts_batch.yaml
```
**Important:** `--batch-size` must match a CUDA graph capture size (1, 2, 4, 8, 16...) because the Talker's code predictor KV cache is sized to `max_num_seqs`, and CUDA graphs pad the batch to the next capture size. Both stages need `max_batch_size >= batch_size` in the stage config for batching to take effect. If only stage 1 has a higher `max_batch_size`, it won't help — stage 1 can only batch chunks from requests that are in-flight simultaneously, which requires stage 0 to also process multiple requests concurrently.
## Notes
- The script uses the model paths embedded in `end2end.py`. Update them if your local cache path differs.
- Use `--output-dir` to change the output folder.
- Use `--output-dir` (preferred) or `--output-wav` to change the output folder.
## Example materials
??? abstract "benchmark_prompts.txt"
``````txt
--8<-- "examples/offline_inference/qwen3_tts/benchmark_prompts.txt"
``````
??? abstract "end2end.py"
``````py
--8<-- "examples/offline_inference/qwen3_tts/end2end.py"
@@ -3,7 +3,7 @@
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.
This folder provides several entrypoints for experimenting with `Qwen/Qwen-Image` `Qwen/Qwen-Image-2512` `Tongyi-MAI/Z-Image-Turbo` using vLLM-Omni:
- `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.
@@ -77,8 +77,6 @@ if __name__ == "__main__":
## Local CLI Usage
### Qwen/Tongyi Models
```bash
python text_to_image.py \
--model Tongyi-MAI/Z-Image-Turbo \
@@ -92,26 +90,7 @@ python text_to_image.py \
--output outputs/coffee.png
```
### NextStep Models
NextStep-1.1 can have extra arguments
```bash
python text_to_image.py \
--model stepfun-ai/NextStep-1.1 \
--prompt "A baby panda wearing an Iron Man mask, holding a board with 'NextStep-1' written on it" \
--height 512 \
--width 512 \
--num-inference-steps 28 \
--guidance-scale 7.5 \
--guidance-scale-2 1.0 \
--cfg-schedule constant \
--output nextstep_output.png \
--seed 42
```
### Key Arguments
**Common arguments:**
Key arguments:
- `--prompt`: text description (string).
- `--seed`: integer seed for deterministic sampling.
@@ -122,50 +101,13 @@ python text_to_image.py \
- `--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).
- `--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).
- `--enable-cpu-offload`: enable CPU offloading for diffusion models.
- `--guidance-scale`: classifier-free guidance scale.
**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 you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.
> ℹ️ 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.
## 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.
### Basic usage with LoRA
```bash
python text_to_image.py \
--model Tongyi-MAI/Z-Image-Turbo \
--prompt "A piece of cheesecake" \
--lora-path /path/to/lora/ \
--lora-scale 1.0 \
--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_adapter/
├── adapter_config.json
└── adapter_model.safetensors
```
## Web UI Demo
Launch the gradio demo:
@@ -25,152 +25,12 @@ cd /workspace/vllm-omni/examples/online_serving/bagel
bash run_server.sh
```
If you have a custom stage configs file, launch the server with the command below:
```bash
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
```
#### 🚀 Tensor Parallelism (TP)
For larger models or multi-GPU environments, you can enable Tensor Parallelism (TP) for the server.
1. **Modify Stage Config**: Create or modify a stage configuration yaml (e.g., [`bagel.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/model_executor/stage_configs/bagel.yaml)). Set `tensor_parallel_size` to `2` (or more) and update `devices` to include multiple GPU IDs (e.g., `"0,1"`).
```yaml
engine_args:
tensor_parallel_size: 2
...
runtime:
devices: "0,1"
```
2. **Launch Server**:
```bash
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --stage-configs-path /path/to/your/custom_bagel.yaml
```
#### Using Mooncake Connector
By default, BAGEL uses `SharedMemoryConnector` for inter-stage communication. You can use the [Mooncake](https://github.com/kvcache-ai/Mooncake) connector to transfer KV cache between stages, which also enables multi-node deployment.
**1. Install Mooncake**
```bash
# For CUDA-enabled systems (recommended)
pip install mooncake-transfer-engine
# For non-CUDA systems
pip install mooncake-transfer-engine-non-cuda
```
**2. Start Mooncake Master** on the primary node:
```bash
# Optional: enable disk-backed storage by creating a directory and passing --root_fs_dir.
# Without it, Mooncake runs in memory-only mode, which is sufficient for KV cache transfer.
mkdir -p ./mc_storage
mooncake_master \
--rpc_port=50051 \
--enable_http_metadata_server=true \
--http_metadata_server_host=0.0.0.0 \
--http_metadata_server_port=8080 \
--metrics_port=9003 \
--root_fs_dir=./mc_storage/ \
--cluster_id=mc-local-1 &
```
**3. Launch the server** with the Mooncake stage config:
```bash
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \
--stage-configs-path vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
```
> **Note**: Before launching, edit [`bagel_multiconnector.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml) and replace the `metadata_server` and `master` addresses with your Mooncake master node's actual IP. For single-node testing, `127.0.0.1` works.
The client-side usage is identical to the default setup -- the Mooncake connector is transparent to the API. See the requests section below.
For more details on the Mooncake connector configuration, see the [Mooncake Store Connector documentation](https://github.com/vllm-project/vllm-omni/tree/main/docs/design/feature/omni_connectors/mooncake_store_connector.md).
#### Multi-Node Deployment
You can deploy each stage on a **separate node** for better resource utilization. In this example, the orchestrator (Stage 0 / Thinker) and Stage 1 (DiT) run on different machines, connected via Mooncake.
Replace `<ORCHESTRATOR_IP>` below with the actual IP address of your orchestrator node (e.g., `10.244.227.244`).
> [!WARNING]
> **Before launching**, edit [`bagel_multiconnector.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml) and replace the `metadata_server` and `master` addresses with your Mooncake master node's actual IP. Mismatched addresses will cause silent connection failures.
**1. Start Mooncake Master** (on the orchestrator node):
```bash
mooncake_master \
--rpc_port=50051 \
--enable_http_metadata_server=true \
--http_metadata_server_host=<ORCHESTRATOR_IP> \
--http_metadata_server_port=8080 \
--metrics_port=9003
```
**2. Launch Stage 0 (Thinker / Orchestrator)** on the orchestrator node:
```bash
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni \
--port 8000 \ # API server port for client requests
--stage-configs-path vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml \
--stage-id 0 \
-oma <ORCHESTRATOR_IP> \
-omp 8091
```
**3. Launch Stage 1 (DiT)** on the remote node in headless mode:
```bash
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni \
--stage-configs-path vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml \
--stage-id 1 \
--headless \
-oma <ORCHESTRATOR_IP> \
-omp 8091
```
**Mooncake Master arguments:**
| Argument | Description |
| :------- | :---------- |
| `--rpc_port` | Mooncake RPC port for control-plane coordination between stages |
| `--enable_http_metadata_server` | Enable the HTTP metadata server for service discovery |
| `--http_metadata_server_host` | IP address to bind the metadata server (use the orchestrator node's IP) |
| `--http_metadata_server_port` | Port for the HTTP metadata server |
| `--metrics_port` | Port for Prometheus-compatible metrics endpoint |
**vllm serve arguments:**
| Argument | Description |
| :------- | :---------- |
| `--stage-id` | Which stage this process runs (0 = Thinker, 1 = DiT) |
| `--headless` | Run without the API server (worker-only mode) |
| `-oma` | Orchestrator master address |
| `-omp` | Orchestrator master port for Stage 1 to connect to Stage 0 for task coordination |
> [!IMPORTANT]
> **Startup Order**: Stage 0 (orchestrator) must be launched **before** Stage 1 (headless).
> Stage 0 will appear to hang on startup until Stage 1 (worker) connects — this is expected behavior.
**Network Requirements**
All nodes must have network connectivity to each other. Ensure the following ports are open **between all participating nodes**:
| Port | Protocol | Service | Direction |
| :--- | :------- | :------ | :-------- |
| 50051 | TCP | Mooncake Master RPC | Worker → Orchestrator |
| 8080 | TCP | Mooncake HTTP Metadata Server | Worker → Orchestrator |
| 8091 | TCP | Orchestrator Master (`-omp`) | Worker → Orchestrator |
| 8000 | TCP | API Server (`--port`) | Client → Orchestrator |
| 9003 | TCP | Metrics (optional) | Monitoring → Orchestrator |
> **Tip**: If nodes are behind a firewall or in different VPCs/security groups, make sure the above ports are allowed in ingress/egress rules. All nodes should be reachable via their IP addresses (no NAT). Using nodes on the same subnet or VPC is recommended to minimize latency for Mooncake KV cache transfers.
### Send Multi-modal Request
Get into the bagel folder:
@@ -382,7 +242,3 @@ sudo apt install ffmpeg
``````sh
--8<-- "examples/online_serving/bagel/run_server.sh"
``````
??? abstract "run_server_stage_cli.sh"
``````sh
--8<-- "examples/online_serving/bagel/run_server_stage_cli.sh"
``````
@@ -1,269 +0,0 @@
# GLM-Image Online Serving
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/glm_image>.
This example demonstrates how to deploy GLM-Image for online image generation using vLLM-Omni.
## 🛠️ Installation
Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/README.md)
## Run examples (GLM-Image)
**Note**: These examples work with the default configuration on **2× NVIDIA A100 (80GB)** or equivalent. Stage 0 (AR) and Stage 1 (Diffusion) each use one GPU by default. For single-GPU setups, modify the stage configuration to share the same device.
### Launch the Server
```bash
# Use default configuration
vllm serve zai-org/GLM-Image --omni --port 8091
```
Or use the convenience script:
```bash
cd examples/online_serving/glm_image
bash run_server.sh
```
If you have a custom stage configs file:
```bash
vllm serve zai-org/GLM-Image --omni --port 8091 --stage-configs-path /path/to/glm_image.yaml
```
### Send Requests
Get into the glm_image folder:
```bash
cd examples/online_serving/glm_image
```
Send request via Python:
```bash
python openai_chat_client.py --prompt "A cute cat sitting on a window sill"
```
The Python client supports the following command-line arguments:
- `--prompt` (or `-p`): Text prompt for generation (default: `A beautiful sunset over the ocean with sailing boats`)
- `--output` (or `-o`): Output file path (default: `glm_image_output.png`)
- `--server` (or `-s`): Server URL (default: `http://localhost:8091`)
- `--image` (or `-i`): Input image path (for image-to-image editing)
- `--height`: Image height in pixels (default: 1024)
- `--width`: Image width in pixels (default: 1024)
- `--steps`: Number of inference steps (default: 50)
- `--guidance-scale`: Classifier-free guidance scale (default: 1.5)
- `--seed`: Random seed (default: 42)
- `--negative`: Negative prompt
## Modality Control
GLM-Image supports **text-to-image** and **image-to-image** modes.
The default yaml configuration deploys AR on GPU 0 and DiT on GPU 1. You can use the default configuration file: [`glm_image.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/model_executor/stage_configs/glm_image.yaml)
| Mode | Input | Output | Description |
| -------------- | ------------ | ------ | ---------------------------------- |
| Text-to-Image | Text | Image | Generate images from text prompts |
| Image-to-Image | Image + Text | Image | Edit images with text instructions |
### 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
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
bash run_curl_image_edit.sh input.png "Convert to watercolor style"
```
## Generation Parameters (extra_body)
| Parameter | Type | Default | Description |
| --------------------- | ----- | ------- | ----------------------------------- |
| `height` | int | 1024 | Image height in pixels |
| `width` | int | 1024 | Image width in pixels |
| `num_inference_steps` | int | 50 | Number of diffusion denoising steps |
| `guidance_scale` | float | 1.5 | Classifier-free guidance scale |
| `seed` | int | 42 | Random seed for reproducibility |
| `negative_prompt` | str | None | Negative prompt |
## Response Format
```json
{
"id": "chatcmpl-xxx",
"created": 1234567890,
"model": "zai-org/GLM-Image",
"choices": [
{
"index": 0,
"message": {
"role": "assistant",
"content": [
{
"type": "image_url",
"image_url": {
"url": "data:image/png;base64,..."
}
}
]
},
"finish_reason": "stop"
}
],
"usage": {}
}
```
## Extract Image
```bash
# From a saved JSON response
cat response.json | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > output.png
```
## Architecture
GLM-Image uses a 2-stage multistage pipeline:
```
Stage 0 (AR Model) Stage 1 (Diffusion)
┌───────────────────┐ ┌─────────────────────┐
│ vLLM-optimized │ prior │ GlmImagePipeline │
│ GlmImageFor │──tokens──►│ ┌───────────────┐ │
│ Conditional │ │ │ DiT Denoiser │ │
│ Generation │ │ └───────┬───────┘ │
│ (9B AR model) │ │ ▼ │
└───────────────────┘ │ ┌───────────────┐ │
▲ │ │ VAE Decode │──┼──► Image
│ │ └───────────────┘ │
Text / Image └─────────────────────┘
Input
```
## VRAM Requirements
| Stage | VRAM |
| :---------------- | :--------------------- |
| Stage-0 (AR) | **~18 GiB + KV Cache** |
| Stage-1 (DiT+VAE) | **~20 GiB** |
| Total | **~38 GiB + KV Cache** |
## File Description
| File | Description |
| --------------------------- | ------------------------------------- |
| `run_server.sh` | Server startup script |
| `run_curl_text_to_image.sh` | Text-to-image curl example |
| `run_curl_image_edit.sh` | Image-to-image (editing) curl example |
| `openai_chat_client.py` | Python client (t2i + i2i) |
## FAQ
- If you encounter OOM errors, adjust `gpu_memory_utilization` in the stage config:
```yaml
# In glm_image.yaml, reduce from default 0.6:
gpu_memory_utilization: 0.5
```
- The first request may be slow due to model warmup. Subsequent requests will be faster.
## Example materials
??? abstract "openai_chat_client.py"
``````py
--8<-- "examples/online_serving/glm_image/openai_chat_client.py"
``````
??? abstract "run_curl_image_edit.sh"
``````sh
--8<-- "examples/online_serving/glm_image/run_curl_image_edit.sh"
``````
??? abstract "run_curl_text_to_image.sh"
``````sh
--8<-- "examples/online_serving/glm_image/run_curl_text_to_image.sh"
``````
??? abstract "run_server.sh"
``````sh
--8<-- "examples/online_serving/glm_image/run_server.sh"
``````
@@ -1,97 +0,0 @@
# Image-To-Video
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/image_to_video>.
This example demonstrates how to deploy the Wan2.2 image-to-video model for online video generation using vLLM-Omni.
## Start Server
### Basic Start
```bash
vllm serve Wan-AI/Wan2.2-I2V-A14B-Diffusers --omni --port 8091
```
### Start with Parameters
Or use the startup script:
```bash
bash run_server.sh
```
The script allows overriding:
- `MODEL` (default: `Wan-AI/Wan2.2-I2V-A14B-Diffusers`)
- `PORT` (default: `8091`)
- `BOUNDARY_RATIO` (default: `0.875`)
- `FLOW_SHIFT` (default: `12.0`)
- `CACHE_BACKEND` (default: `none`)
- `ENABLE_CACHE_DIT_SUMMARY` (default: `0`)
## API Calls
### Method 1: Using curl
```bash
# Basic image-to-video generation
bash run_curl_image_to_video.sh
# Or execute directly (OpenAI-style multipart)
curl -X POST http://localhost:8091/v1/videos \
-H "Accept: application/json" \
-F "prompt=A bear playing with yarn, smooth motion" \
-F "negative_prompt=low quality, blurry, static" \
-F "input_reference=@/path/to/qwen-bear.png" \
-F "width=832" \
-F "height=480" \
-F "num_frames=33" \
-F "fps=16" \
-F "num_inference_steps=40" \
-F "guidance_scale=1.0" \
-F "guidance_scale_2=1.0" \
-F "boundary_ratio=0.875" \
-F "flow_shift=12.0" \
-F "seed=42" | jq -r '.data[0].b64_json' | base64 -d > wan22_i2v_output.mp4
```
## Request Format
### Required Fields
```bash
curl -X POST http://localhost:8091/v1/videos \
-F "prompt=A bear playing with yarn, smooth motion" \
-F "negative_prompt=low quality, blurry, static" \
-F "input_reference=@/path/to/qwen-bear.png"
```
### Generation with Parameters
```bash
curl -X POST http://localhost:8091/v1/videos \
-F "prompt=A bear playing with yarn, smooth motion" \
-F "negative_prompt=low quality, blurry, static" \
-F "input_reference=@/path/to/qwen-bear.png" \
-F "width=832" \
-F "height=480" \
-F "num_frames=33" \
-F "fps=16" \
-F "num_inference_steps=40" \
-F "guidance_scale=1.0" \
-F "guidance_scale_2=1.0" \
-F "boundary_ratio=0.875" \
-F "flow_shift=12.0" \
-F "seed=42"
```
## Example materials
??? abstract "run_curl_image_to_video.sh"
``````sh
--8<-- "examples/online_serving/image_to_video/run_curl_image_to_video.sh"
``````
??? abstract "run_server.sh"
``````sh
--8<-- "examples/online_serving/image_to_video/run_server.sh"
``````
@@ -0,0 +1,72 @@
# LoRA Inference(Diffusion)
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/lora_inference>.
This example shows how to use **per-request LoRA** with vLLM-Omni diffusion models via the OpenAI-compatible Chat Completions API.
> Note: The LoRA adapter path must be readable on the **server** machine (usually a local path or a mounted directory).
> Note: This example uses `/v1/chat/completions`. LoRA payloads for other OpenAI endpoints are not implemented here.
## Start Server
```bash
# Pick a diffusion model (examples)
# export MODEL=stabilityai/stable-diffusion-3.5-medium
# export MODEL=Qwen/Qwen-Image
bash run_server.sh
```
## Call API (curl)
```bash
# Required: local LoRA folder on the server
export LORA_PATH=/path/to/lora_adapter
# Optional
export SERVER=http://localhost:8091
export PROMPT="A piece of cheesecake"
export LORA_NAME=my_lora
export LORA_SCALE=1.0
# Optional: if omitted, the server derives a stable id from LORA_PATH.
# export LORA_INT_ID=123
bash run_curl_lora_inference.sh
```
## Call API (Python)
```bash
python openai_chat_client.py \
--prompt "A piece of cheesecake" \
--lora-path /path/to/lora_adapter \
--lora-name my_lora \
--lora-scale 1.0 \
--output output.png
```
## LoRA Format
LoRA adapters should be in PEFT format, for example:
```
lora_adapter/
├── adapter_config.json
└── adapter_model.safetensors
```
## Example materials
??? abstract "openai_chat_client.py"
``````py
--8<-- "examples/online_serving/lora_inference/openai_chat_client.py"
``````
??? abstract "run_curl_lora_inference.sh"
``````sh
--8<-- "examples/online_serving/lora_inference/run_curl_lora_inference.sh"
``````
??? abstract "run_server.sh"
``````sh
--8<-- "examples/online_serving/lora_inference/run_server.sh"
``````
@@ -30,7 +30,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 openai_chat_completion_client_for_multimodal_generation.py --query-type mixed_modalities
```
The Python client supports the following command-line arguments:
@@ -36,7 +36,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 openai_chat_completion_client_for_multimodal_generation.py --query-type use_image
```
The Python client supports the following command-line arguments:
@@ -9,13 +9,13 @@ Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/
## Supported Models
| Model | Task Type | Description |
| -------------------------------------- | ----------- | ----------------------------------------------------- |
| Model | Task Type | Description |
|-------|-----------|-------------|
| `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | CustomVoice | Predefined speaker voices with optional style control |
| `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | VoiceDesign | Natural language voice style description |
| `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | Base | Voice cloning from reference audio |
| `Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice` | CustomVoice | Smaller/faster variant |
| `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | Base | Smaller/faster variant for voice cloning |
| `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | VoiceDesign | Natural language voice style description |
| `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | Base | Voice cloning from reference audio |
| `Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice` | CustomVoice | Smaller/faster variant |
| `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | Base | Smaller/faster variant for voice cloning |
## Run examples (Qwen3-TTS)
@@ -237,58 +237,36 @@ Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/w
### OpenAI Standard Parameters
| Parameter | Type | Default | Description |
| ----------------- | ------ | -------------- | ----------------------------------------------------------- |
| `input` | string | **required** | Text to synthesize |
| `model` | string | server's model | Model to use (optional, should match server if specified) |
| `voice` | string | "vivian" | Speaker name (e.g., vivian, ryan, aiden) |
| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
| `speed` | float | 1.0 | Playback speed (0.25-4.0, not supported with `stream=true`) |
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `input` | string | **required** | Text to synthesize |
| `model` | string | server's model | Model to use (optional, should match server if specified) |
| `voice` | string | "vivian" | Speaker name (e.g., vivian, ryan, aiden) |
| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
| `speed` | float | 1.0 | Playback speed (0.25-4.0) |
### 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 |
| `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 |
**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) |
## Streaming
Set `stream=true` with `response_format="pcm"` to receive raw PCM audio chunks as they are decoded
(one chunk per 25-frame Code2Wav window):
```bash
curl -X POST http://localhost:8091/v1/audio/speech \
-H "Content-Type: application/json" \
-d '{
"input": "Hello, how are you?",
"voice": "vivian",
"language": "English",
"stream": true,
"response_format": "pcm"
}' --no-buffer | play -t raw -r 24000 -e signed -b 16 -c 1 -
```
**Constraints:**
- `stream=true` requires `response_format="pcm"` (raw 16-bit signed PCM, 24 kHz mono).
- `speed` adjustment is not supported when streaming.
- Requires the server stage config to have `async_chunk: true` (default in `qwen3_tts.yaml`).
| 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) |
## Limitations
- **No streaming**: Audio is generated completely before being returned. Streaming will be supported after the pipeline is disaggregated (see RFC #938).
- **Single request**: Batch processing is not yet optimized for online serving.
## Troubleshooting
@@ -61,59 +61,6 @@ python gradio_demo.py
# Visit http://localhost:7860
```
## LoRA
This example supports Peft-compatible LoRA (Low-Rank Adaptation) adapters for diffusion models. The LoRA adapter path must be readable on the **server** machine (usually a local path or a mounted directory).
### Using Python Client with LoRA
```bash
python openai_chat_client.py \
--prompt "A piece of cheesecake" \
--lora-path /path/to/lora_adapter \
--lora-name my_lora \
--lora-scale 1.0 \
--output output.png
```
### Using curl with LoRA (Images API)
The `/v1/images/generations` endpoint supports a `lora` field in the request body:
```bash
curl -X POST http://localhost:8091/v1/images/generations \
-H "Content-Type: application/json" \
-d '{
"prompt": "A piece of cheesecake",
"size": "1024x1024",
"seed": 42,
"lora": {
"name": "my_lora",
"local_path": "/path/to/lora_adapter",
"scale": 1.0
}
}' | jq -r '.data[0].b64_json' | base64 -d > output.png
```
### 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) |
### LoRA Adapter Format
LoRA adapters must be in PEFT (Parameter-Efficient Fine-Tuning) format. A typical LoRA adapter directory structure:
```
lora_adapter/
├── adapter_config.json
└── adapter_model.safetensors
```
## Request Format
### Simple Text Generation
@@ -1,130 +0,0 @@
# Text-To-Video
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/text_to_video>.
This example demonstrates how to deploy the Wan2.2 text-to-video model for online video generation using vLLM-Omni.
## Start Server
### Basic Start
```bash
vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --port 8091
```
### Start with Parameters
Or use the startup script:
```bash
bash run_server.sh
```
The script allows overriding:
- `MODEL` (default: `Wan-AI/Wan2.2-T2V-A14B-Diffusers`)
- `PORT` (default: `8091`)
- `BOUNDARY_RATIO` (default: `0.875`)
- `FLOW_SHIFT` (default: `5.0`)
- `CACHE_BACKEND` (default: `none`)
- `ENABLE_CACHE_DIT_SUMMARY` (default: `0`)
## API Calls
### Method 1: Using curl
```bash
# Basic text-to-video generation
bash run_curl_text_to_video.sh
# Or execute directly (OpenAI-style multipart)
curl -s http://localhost:8091/v1/videos \
-H "Accept: application/json" \
-F "prompt=Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \
-F "width=832" \
-F "height=480" \
-F "num_frames=33" \
-F "negative_prompt=色调艳丽 ,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,静止,整体发灰,最差质量,低质量,JPEG压缩残留,丑陋的,残缺的,多余的手指,画得不好的手部,画得不好的脸部,畸形的,毁容的,形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,背景人很多,倒着走" \
-F "fps=16" \
-F "num_inference_steps=40" \
-F "guidance_scale=4.0" \
-F "guidance_scale_2=4.0" \
-F "boundary_ratio=0.875" \
-F "seed=42" | jq -r '.data[0].b64_json' | base64 -d > wan22_output.mp4
```
## Request Format
### Simple Text-to-Video Generation
```bash
curl -X POST http://localhost:8091/v1/videos \
-F "prompt=A cinematic view of a futuristic city at sunset"
```
### Generation with Parameters
```bash
curl -X POST http://localhost:8091/v1/videos \
-F "prompt=A cinematic view of a futuristic city at sunset" \
-F "width=832" \
-F "height=480" \
-F "num_frames=33" \
-F "negative_prompt=low quality, blurry, static" \
-F "fps=16" \
-F "num_inference_steps=40" \
-F "guidance_scale=4.0" \
-F "guidance_scale_2=4.0" \
-F "boundary_ratio=0.875" \
-F "flow_shift=5.0" \
-F "seed=42"
```
## Generation Parameters
| Parameter | Type | Default | Description |
| --------------------- | ------ | ------- | ------------------------------------------------ |
| `prompt` | str | - | Text description of the desired video |
| `negative_prompt` | str | None | Negative prompt |
| `n` | int | 1 | Number of videos to generate |
| `width` | int | None | Video width in pixels |
| `height` | int | None | Video height in pixels |
| `num_frames` | int | None | Number of frames to generate |
| `fps` | int | None | Frames per second for output video |
| `num_inference_steps` | int | None | Number of denoising steps |
| `guidance_scale` | float | None | CFG guidance scale (low-noise stage) |
| `guidance_scale_2` | float | None | CFG guidance scale (high-noise stage, Wan2.2) |
| `boundary_ratio` | float | None | Boundary split ratio for low/high DiT (Wan2.2) |
| `flow_shift` | float | None | Scheduler flow shift (Wan2.2) |
| `seed` | int | None | Random seed (reproducible) |
| `lora` | object | None | LoRA configuration |
| `extra_body` | object | None | Model-specific extra parameters |
## Response Format
```json
{
"created": 1234567890,
"data": [
{ "b64_json": "<base64-mp4>" }
]
}
```
## Extract Video
```bash
# Extract base64 from response and decode to video
cat response.json | jq -r '.data[0].b64_json' | base64 -d > wan22_output.mp4
```
## Example materials
??? abstract "run_curl_text_to_video.sh"
``````sh
--8<-- "examples/online_serving/text_to_video/run_curl_text_to_video.sh"
``````
??? abstract "run_server.sh"
``````sh
--8<-- "examples/online_serving/text_to_video/run_server.sh"
``````
@@ -177,74 +177,6 @@ Example configuration for TP=2 on GPUs 0 and 1:
| `max_inflight` | `1` | Maximum inflight requests |
| `shm_threshold_bytes` | `65536` | Shared memory threshold (64KB) |
## Using Mooncake Connector
[Mooncake](https://github.com/kvcache-ai/Mooncake) is a high-performance distributed KV cache transfer engine that enables efficient cross-node data movement via TCP or RDMA, making it ideal for multi-node disaggregated inference.
By default, BAGEL uses `SharedMemoryConnector` for inter-stage communication. You can switch to the Mooncake connector for better performance on multi-GPU setups and to enable multi-node deployment.
### Prerequisites
Install the Mooncake transfer engine:
```bash
# For CUDA-enabled systems (recommended)
pip install mooncake-transfer-engine
# For non-CUDA systems
pip install mooncake-transfer-engine-non-cuda
```
### Step 1: Start the Mooncake Master
On the **primary node**, start the Mooncake master service (run in a separate terminal or background with `&`):
```bash
# Optional: enable disk-backed storage by creating a directory and passing --root_fs_dir.
# Without it, Mooncake runs in memory-only mode, which is sufficient for KV cache transfer.
mkdir -p ./mc_storage
mooncake_master \
--rpc_port=50051 \
--enable_http_metadata_server=true \
--http_metadata_server_host=0.0.0.0 \
--http_metadata_server_port=8080 \
--metrics_port=9003 \
--root_fs_dir=./mc_storage/ \
--cluster_id=mc-local-1 &
```
### Step 2: Run Offline Inference with Mooncake
Use the provided Mooncake stage config [`bagel_multiconnector.yaml`](../../../vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml). Before launching, update the `metadata_server` and `master` addresses in the YAML to match your Mooncake master node's IP (use `127.0.0.1` for single-node testing).
```bash
cd examples/offline_inference/bagel
# Text to Image with Mooncake
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2img \
--prompts "A cute cat" \
--stage-configs-path ../../../vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
# Image to Text with Mooncake
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality img2text \
--image-path /path/to/image.jpg \
--prompts "Describe this image" \
--stage-configs-path ../../../vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
# Text to Text with Mooncake
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2text \
--prompts "What is the capital of France?" \
--stage-configs-path ../../../vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
```
For more details on the Mooncake connector and multi-node setup, see the [Mooncake Store Connector documentation](../../../docs/design/feature/omni_connectors/mooncake_store_connector.md).
------
## FAQ
- If you encounter an error about the backend of librosa, try to install ffmpeg with the command below.
@@ -1,138 +0,0 @@
# GLM-Image Multistage End-to-End Inference
This example demonstrates how to run GLM-Image with the vLLM-Omni multistage architecture.
## Architecture
GLM-Image uses a 2-stage pipeline:
```
┌─────────────────────────────────────────────────────────────┐
│ GLM-Image Pipeline │
├─────────────────────────────────────────────────────────────┤
│ │
│ Stage 0 (AR Model) Stage 1 (Diffusion) │
│ ┌─────────────────┐ ┌─────────────────────┐ │
│ │ vLLM-optimized │ │ GlmImagePipeline │ │
│ │ GlmImageFor │ prior │ ┌───────────────┐ │ │
│ │ Conditional │──tokens───►│ │ DiT Denoiser │ │ │
│ │ Generation │ │ └───────────────┘ │ │
│ │ (9B AR model) │ │ │ │ │
│ └─────────────────┘ │ ▼ │ │
│ ▲ │ ┌───────────────┐ │ │
│ │ │ │ VAE Decode │──┼──► Image
│ Text/Image │ └───────────────┘ │ │
│ Input └─────────────────────┘ │
│ │
└─────────────────────────────────────────────────────────────┘
```
## Features
- **vLLM-optimized AR**: Uses PagedAttention and tensor parallelism for faster prior token generation
- **Flexible deployment**: AR and Diffusion stages can run on different GPUs
- **Text-to-Image**: Generate images from text descriptions
- **Image-to-Image**: Edit existing images with text prompts
## Usage
### Text-to-Image
```bash
python end2end.py \
--model-path /path/to/glm-image \
--config-path ../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
--prompt "A beautiful sunset over the ocean with sailing boats" \
--height 1024 \
--width 1024 \
--output output_t2i.png
```
### Image-to-Image (Image Editing)
```bash
python end2end.py \
--model-path /path/to/glm-image \
--config-path ../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
--prompt "Transform this scene into a winter wonderland" \
--image input.png \
--output output_i2i.png
```
### With Custom Parameters
```bash
python end2end.py \
--model-path /path/to/glm-image \
--config-path ../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
--prompt "A photorealistic cat sitting on a window sill" \
--height 1024 \
--width 1024 \
--num-inference-steps 50 \
--guidance-scale 1.5 \
--seed 42 \
--output output.png
```
## Shell Scripts
### Run Text-to-Image
```bash
./run_t2i.sh
```
### Run Image-to-Image
```bash
./run_i2i.sh --image /path/to/input.png
```
## Stage Configuration
The stage config (`glm_image.yaml`) defines:
- **Stage 0 (AR)**: Uses `GPUARWorker` with vLLM engine
- Model: `GlmImageForConditionalGeneration`
- Output: `token_ids` (prior tokens)
- **Stage 1 (Diffusion)**: Uses diffusion engine
- Model: `GlmImagePipeline`
- Output: Generated image
See `vllm_omni/model_executor/stage_configs/glm_image.yaml` for full configuration.
## Comparison with Single-Stage
| Aspect | Single-Stage (transformers) | Multistage (vLLM) |
| ----------- | --------------------------- | ------------------- |
| AR Model | transformers native | vLLM PagedAttention |
| Memory | Higher (no KV cache opt) | Lower (optimized) |
| Throughput | Lower | Higher |
| Flexibility | Single GPU | Multi-GPU support |
## Troubleshooting
### OOM Error
Try reducing memory usage:
```bash
# In glm_image.yaml, adjust:
gpu_memory_utilization: 0.5 # Reduce from 0.6
```
### Slow Initialization
The first run loads model weights. Subsequent runs are faster:
```bash
--stage-init-timeout 900 # Increase timeout for slow storage
```
## Requirements
- vLLM-Omni with GLM-Image support
- CUDA-capable GPU (recommended: H100/A100 with 80GB)
- GLM-Image model weights
@@ -1,476 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
End-to-end offline inference example for GLM-Image with multistage architecture.
This script tests the multistage pipeline where:
- Stage 0 (AR): vLLM-optimized GlmImageForConditionalGeneration generates prior_token_ids
- Stage 1 (Diffusion): GlmImagePipeline performs DiT denoising + VAE decode
Usage (text-to-image):
python end2end.py \
--model-path /path/to/glm-image \
--config-path /path/to/glm_image.yaml \
--prompt "A beautiful sunset over the ocean" \
--output output_t2i.png
Usage (image-to-image / image edit):
python end2end.py \
--model-path /path/to/glm-image \
--config-path /path/to/glm_image.yaml \
--prompt "Make it look like winter" \
--image input.png \
--output output_i2i.png
Usage (with custom parameters):
python end2end.py \
--model-path /path/to/glm-image \
--config-path /path/to/glm_image.yaml \
--prompt "A cat sitting on a window sill" \
--height 1024 \
--width 1024 \
--num-inference-steps 50 \
--guidance-scale 1.5 \
--seed 42
For more options, run:
python end2end.py --help
"""
import argparse
import os
import time
from pathlib import Path
from PIL import Image
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
# Default stage config path (relative to vllm_omni package)
DEFAULT_CONFIG_PATH = "vllm_omni/model_executor/stage_configs/glm_image.yaml"
SEED = 42
# GLM-Image special tokens
GLM_IMAGE_EOS_TOKEN_ID = 16385 # eos_token_id from generation_config.json
GLM_IMAGE_VISION_VOCAB_SIZE = 16512 # top_k should be vision_vocab_size
def compute_max_tokens(height: int, width: int, factor: int = 32) -> int:
"""
Compute max_new_tokens for GLM-Image AR generation.
GLM-Image generates tokens in this order for text-to-image:
1. Small preview image (half resolution in each dimension)
2. Large target image (full resolution)
3. EOS token
Args:
height: Target image height in pixels
width: Target image width in pixels
factor: Downsampling factor (32 for GLM-Image AR output)
Returns:
Total number of tokens to generate (small + large + EOS)
"""
# Large image tokens (target resolution)
token_h = height // factor
token_w = width // factor
large_tokens = token_h * token_w
# Small preview tokens (half resolution in each dimension)
small_h = token_h // 2
small_w = token_w // 2
small_tokens = small_h * small_w
# Total: small + large + EOS
return small_tokens + large_tokens + 1
def load_image(image_path: str) -> Image.Image:
"""Load an image from file path."""
if not os.path.exists(image_path):
raise FileNotFoundError(f"Image file not found: {image_path}")
return Image.open(image_path).convert("RGB")
def save_image(image: Image.Image, output_path: str) -> None:
"""Save an image to file path."""
output_dir = os.path.dirname(output_path)
if output_dir:
os.makedirs(output_dir, exist_ok=True)
image.save(output_path)
print(f"Image saved to: {output_path}")
def build_prompt_for_t2i(
prompt: str,
height: int = 1024,
width: int = 1024,
) -> dict:
"""
Build prompt dict for text-to-image generation.
Args:
prompt: Text description for image generation
height: Target image height
width: Target image width
Returns:
Dict containing prompt and generation parameters
"""
return {
"prompt": prompt,
"height": height,
"width": width,
# Pass target dimensions to AR processor for proper grid token generation
"mm_processor_kwargs": {
"target_h": height,
"target_w": width,
},
}
def build_prompt_for_i2i(
prompt: str,
image: Image.Image,
height: int | None = None,
width: int | None = None,
) -> dict:
"""
Build prompt dict for image-to-image generation.
Args:
prompt: Text description for image editing
image: Source image for editing
height: Target image height (default: use source image size)
width: Target image width (default: use source image size)
Returns:
Dict containing prompt, image, and generation parameters
"""
# Use source image dimensions if not specified
if height is None:
height = image.height
if width is None:
width = image.width
return {
"prompt": prompt,
"multi_modal_data": {
"image": image,
},
"height": height,
"width": width,
# Pass target dimensions to AR processor for proper grid token generation
"mm_processor_kwargs": {
"target_h": height,
"target_w": width,
},
}
def main(args: argparse.Namespace) -> None:
"""Main entry point for GLM-Image end-to-end inference."""
print("=" * 60)
print("GLM-Image Multistage End-to-End Inference")
print("=" * 60)
# Validate arguments
if not args.model_path:
raise ValueError("--model-path is required")
if not args.prompt:
raise ValueError("--prompt is required")
# Determine config path
config_path = args.config_path
if config_path is None:
# Try to find default config
if os.path.exists(DEFAULT_CONFIG_PATH):
config_path = DEFAULT_CONFIG_PATH
else:
# Try relative to script location
script_dir = Path(__file__).parent.parent.parent.parent
config_path = script_dir / "vllm_omni/model_executor/stage_configs/glm_image.yaml"
if not config_path.exists():
raise FileNotFoundError(
f"Stage config not found. Please specify --config-path. Tried: {DEFAULT_CONFIG_PATH}"
)
config_path = str(config_path)
print(f"Model path: {args.model_path}")
print(f"Config path: {config_path}")
print(f"Prompt: {args.prompt}")
# Load source image for image-to-image mode
source_image = None
if args.image:
print(f"Source image: {args.image}")
source_image = load_image(args.image)
print(f" Image size: {source_image.width}x{source_image.height}")
# Build prompt based on mode
if source_image is not None:
# Image-to-image mode
prompt_dict = build_prompt_for_i2i(
prompt=args.prompt,
image=source_image,
height=args.height,
width=args.width,
)
mode = "image-to-image"
else:
# Text-to-image mode
prompt_dict = build_prompt_for_t2i(
prompt=args.prompt,
height=args.height or 1024,
width=args.width or 1024,
)
mode = "text-to-image"
print(f"Mode: {mode}")
print(f"Target size: {prompt_dict.get('height', 1024)}x{prompt_dict.get('width', 1024)}")
# Add generation parameters to prompt
prompt_dict["seed"] = args.seed
prompt_dict["num_inference_steps"] = args.num_inference_steps
prompt_dict["guidance_scale"] = args.guidance_scale
if args.negative_prompt:
prompt_dict["negative_prompt"] = args.negative_prompt
# Initialize Omni with multistage config
print("\nInitializing Omni with multistage pipeline...")
start_time = time.time()
omni = Omni(
model=args.model_path,
stage_configs_path=config_path,
log_stats=args.enable_stats,
stage_init_timeout=args.stage_init_timeout,
)
init_time = time.time() - start_time
print(f"Initialization completed in {init_time:.2f}s")
# Prepare prompts (support batch generation)
prompts = [prompt_dict for _ in range(args.num_prompts)]
# No explicit sampling_params for diffusion - parameters are in prompt_dict
# For multistage, the AR stage may need sampling params
from vllm import SamplingParams
# Compute max_tokens dynamically based on target image size
target_height = prompt_dict.get("height", 1024)
target_width = prompt_dict.get("width", 1024)
calculated_max_tokens = compute_max_tokens(target_height, target_width)
# Use calculated value unless user explicitly specified a different value
# Default args.max_tokens is 16384 (very large), so prefer calculated value
effective_max_tokens = calculated_max_tokens if args.max_tokens == 16384 else args.max_tokens
if args.verbose:
print(f"AR max_tokens: {effective_max_tokens} (calculated: {calculated_max_tokens}, arg: {args.max_tokens})")
# IMPORTANT: GLM-Image AR model requires these exact sampling parameters
# from generation_config.json for proper image token generation.
# - temperature=0.9, top_p=0.75, top_k=16512 (vision_vocab_size)
# - stop_token_ids=[16385] (eos_token_id) is CRITICAL to stop generation
ar_sampling_params = SamplingParams(
temperature=0.9, # From generation_config.json
top_p=0.75, # From generation_config.json
top_k=GLM_IMAGE_VISION_VOCAB_SIZE, # 16512, vision vocabulary size
max_tokens=effective_max_tokens,
stop_token_ids=[GLM_IMAGE_EOS_TOKEN_ID], # 16385, CRITICAL for stopping
seed=args.seed,
detokenize=False,
)
# For diffusion stage, sampling_params contains diffusion-specific parameters
# These are passed as kwargs to the diffusion engine
diffusion_sampling_params = OmniDiffusionSamplingParams(
num_inference_steps=args.num_inference_steps,
guidance_scale=args.guidance_scale,
height=prompt_dict.get("height", 1024),
width=prompt_dict.get("width", 1024),
seed=args.seed,
)
# For multistage, we need sampling_params for each stage
# Stage 0 (AR): SamplingParams for vLLM
# Stage 1 (Diffusion): dict with diffusion kwargs
sampling_params_list = [ar_sampling_params, diffusion_sampling_params]
# Run generation
print(f"\nGenerating {args.num_prompts} image(s)...")
gen_start_time = time.time()
output_dir = os.path.dirname(args.output) if args.output else "outputs"
if output_dir:
os.makedirs(output_dir, exist_ok=True)
output_count = 0
for stage_outputs in omni.generate(prompts, sampling_params_list, py_generator=True):
if stage_outputs.final_output_type == "image":
for output in stage_outputs.request_output:
request_id = output.request_id
# Get generated images
images = output.images if hasattr(output, "images") else []
if not images and hasattr(output, "multimodal_output"):
images = output.multimodal_output.get("images", [])
# Save each generated image
for idx, img in enumerate(images):
if args.num_prompts == 1 and len(images) == 1:
output_path = args.output
else:
base, ext = os.path.splitext(args.output)
output_path = f"{base}_{request_id}_{idx}{ext}"
if isinstance(img, Image.Image):
save_image(img, output_path)
else:
print(f"Warning: Unexpected image type for request {request_id}: {type(img)}")
output_count += 1
elif stage_outputs.final_output_type == "text":
# AR stage output (intermediate, for debugging)
if args.verbose:
for output in stage_outputs.request_output:
print(f"AR output for request {output.request_id}:")
print(f" Token count: {len(output.outputs[0].token_ids)}")
gen_time = time.time() - gen_start_time
print(f"\nGeneration completed in {gen_time:.2f}s")
print(f"Generated {output_count} image(s)")
# Cleanup
omni.close()
print("\nDone!")
def parse_args() -> argparse.Namespace:
"""Parse command line arguments."""
parser = argparse.ArgumentParser(
description="GLM-Image Multistage End-to-End Inference",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog=__doc__,
)
# Required arguments
parser.add_argument(
"--model-path",
type=str,
required=True,
help="Path to GLM-Image model directory or HuggingFace model ID",
)
parser.add_argument(
"--prompt",
type=str,
required=True,
help="Text prompt for image generation",
)
# Optional arguments
parser.add_argument(
"--config-path",
type=str,
default=None,
help="Path to stage config YAML file (default: auto-detect)",
)
parser.add_argument(
"--image",
type=str,
default=None,
help="Path to source image for image-to-image mode",
)
parser.add_argument(
"--output",
type=str,
default="output_glm_image.png",
help="Output image path (default: output_glm_image.png)",
)
parser.add_argument(
"--negative-prompt",
type=str,
default=None,
help="Negative prompt for classifier-free guidance",
)
# Generation parameters
parser.add_argument(
"--height",
type=int,
default=None,
help="Output image height (default: 1024 for t2i, source size for i2i)",
)
parser.add_argument(
"--width",
type=int,
default=None,
help="Output image width (default: 1024 for t2i, source size for i2i)",
)
parser.add_argument(
"--num-inference-steps",
type=int,
default=50,
help="Number of diffusion denoising steps (default: 50)",
)
parser.add_argument(
"--guidance-scale",
type=float,
default=1.5,
help="Classifier-free guidance scale (default: 1.5)",
)
parser.add_argument(
"--seed",
type=int,
default=SEED,
help=f"Random seed for reproducibility (default: {SEED})",
)
parser.add_argument(
"--max-tokens",
type=int,
default=16384,
help="Maximum tokens for AR generation (default: 16384)",
)
# Batch processing
parser.add_argument(
"--num-prompts",
type=int,
default=1,
help="Number of images to generate (default: 1)",
)
# Runtime options
parser.add_argument(
"--enable-stats",
action="store_true",
default=False,
help="Enable statistics logging",
)
parser.add_argument(
"--stage-init-timeout",
type=int,
default=600,
help="Timeout for stage initialization in seconds (default: 600)",
)
parser.add_argument(
"--verbose",
"-v",
action="store_true",
default=False,
help="Enable verbose output",
)
return parser.parse_args()
if __name__ == "__main__":
args = parse_args()
main(args)
@@ -1,93 +0,0 @@
#!/bin/bash
# SPDX-License-Identifier: Apache-2.0
# Run GLM-Image image-to-image (editing) with multistage pipeline
set -e
# Default values
MODEL_PATH="${MODEL_PATH:-/path/to/glm-image}"
CONFIG_PATH="${CONFIG_PATH:-vllm_omni/model_executor/stage_configs/glm_image.yaml}"
PROMPT="${PROMPT:-Transform this image into an oil painting style}"
INPUT_IMAGE=""
OUTPUT="${OUTPUT:-output_i2i.png}"
NUM_STEPS="${NUM_STEPS:-50}"
GUIDANCE="${GUIDANCE:-1.5}"
SEED="${SEED:-42}"
# Parse command line arguments
while [[ $# -gt 0 ]]; do
case $1 in
--model-path)
MODEL_PATH="$2"
shift 2
;;
--config-path)
CONFIG_PATH="$2"
shift 2
;;
--prompt)
PROMPT="$2"
shift 2
;;
--image)
INPUT_IMAGE="$2"
shift 2
;;
--output)
OUTPUT="$2"
shift 2
;;
--num-steps)
NUM_STEPS="$2"
shift 2
;;
--guidance)
GUIDANCE="$2"
shift 2
;;
--seed)
SEED="$2"
shift 2
;;
*)
echo "Unknown option: $1"
exit 1
;;
esac
done
# Check if input image is provided
if [ -z "${INPUT_IMAGE}" ]; then
echo "Error: --image is required for image-to-image mode"
echo "Usage: ./run_i2i.sh --image /path/to/input.png [--prompt \"edit instruction\"]"
exit 1
fi
if [ ! -f "${INPUT_IMAGE}" ]; then
echo "Error: Input image not found: ${INPUT_IMAGE}"
exit 1
fi
echo "=============================================="
echo "GLM-Image Image-to-Image Generation"
echo "=============================================="
echo "Model: ${MODEL_PATH}"
echo "Config: ${CONFIG_PATH}"
echo "Input: ${INPUT_IMAGE}"
echo "Prompt: ${PROMPT}"
echo "Output: ${OUTPUT}"
echo "Steps: ${NUM_STEPS}"
echo "Guidance: ${GUIDANCE}"
echo "Seed: ${SEED}"
echo "=============================================="
python end2end.py \
--model-path "${MODEL_PATH}" \
--config-path "${CONFIG_PATH}" \
--prompt "${PROMPT}" \
--image "${INPUT_IMAGE}" \
--output "${OUTPUT}" \
--num-inference-steps "${NUM_STEPS}" \
--guidance-scale "${GUIDANCE}" \
--seed "${SEED}" \
--verbose
@@ -1,87 +0,0 @@
#!/bin/bash
# SPDX-License-Identifier: Apache-2.0
# Run GLM-Image text-to-image generation with multistage pipeline
set -e
# Default values
MODEL_PATH="${MODEL_PATH:-/path/to/glm-image}"
CONFIG_PATH="${CONFIG_PATH:-vllm_omni/model_executor/stage_configs/glm_image.yaml}"
PROMPT="${PROMPT:-A beautiful mountain landscape with snow-capped peaks and a clear blue lake}"
OUTPUT="${OUTPUT:-output_t2i.png}"
HEIGHT="${HEIGHT:-1024}"
WIDTH="${WIDTH:-1024}"
NUM_STEPS="${NUM_STEPS:-50}"
GUIDANCE="${GUIDANCE:-1.5}"
SEED="${SEED:-42}"
# Parse command line arguments
while [[ $# -gt 0 ]]; do
case $1 in
--model-path)
MODEL_PATH="$2"
shift 2
;;
--config-path)
CONFIG_PATH="$2"
shift 2
;;
--prompt)
PROMPT="$2"
shift 2
;;
--output)
OUTPUT="$2"
shift 2
;;
--height)
HEIGHT="$2"
shift 2
;;
--width)
WIDTH="$2"
shift 2
;;
--num-steps)
NUM_STEPS="$2"
shift 2
;;
--guidance)
GUIDANCE="$2"
shift 2
;;
--seed)
SEED="$2"
shift 2
;;
*)
echo "Unknown option: $1"
exit 1
;;
esac
done
echo "=============================================="
echo "GLM-Image Text-to-Image Generation"
echo "=============================================="
echo "Model: ${MODEL_PATH}"
echo "Config: ${CONFIG_PATH}"
echo "Prompt: ${PROMPT}"
echo "Output: ${OUTPUT}"
echo "Size: ${WIDTH}x${HEIGHT}"
echo "Steps: ${NUM_STEPS}"
echo "Guidance: ${GUIDANCE}"
echo "Seed: ${SEED}"
echo "=============================================="
python end2end.py \
--model-path "${MODEL_PATH}" \
--config-path "${CONFIG_PATH}" \
--prompt "${PROMPT}" \
--output "${OUTPUT}" \
--height "${HEIGHT}" \
--width "${WIDTH}" \
--num-inference-steps "${NUM_STEPS}" \
--guidance-scale "${GUIDANCE}" \
--seed "${SEED}" \
--verbose
@@ -2,23 +2,6 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Example script for image editing with OmniGen2.
python image_edit.py \
--image input.png \
--model "OmniGen2/OmniGen2" \
--prompt "Change the background to classroom." \
--negative-prompt "(((deformed))), blurry, over saturation, bad anatomy, disfigured, poorly drawn face, mutation, mutated, (extra_limb), (ugly), (poorly drawn hands), fused fingers, messy drawing, broken legs censor, censored, censor_bar" \
--num-inference-steps 50 \
--seed 0 \
--guidance-scale 5.0 \
--guidance-scale-2 2.0 \
--output outputs/image_edit.png \
--num-outputs-per-prompt 2
Note: For OmniGen2, `guidance_scale` works as `text_guidance_scale`,
and `guidance_scale_2` works as `image_guidance_scale`.
Example script for image editing with Qwen-Image-Edit.
Usage (single image):
@@ -158,9 +141,6 @@ def parse_args() -> argparse.Namespace:
"directly as an input parameter. Enabled when guidance_scale > 1. Ignored when not using guidance-distilled models."
),
)
parser.add_argument(
"--guidance-scale-2", type=float, default=None, help="image guidance scale for image-to-image generation."
)
parser.add_argument(
"--output",
type=str,
@@ -423,7 +403,6 @@ def main():
generator=generator,
true_cfg_scale=args.cfg_scale,
guidance_scale=args.guidance_scale,
guidance_scale_2=args.guidance_scale_2,
num_inference_steps=args.num_inference_steps,
num_outputs_per_prompt=args.num_outputs_per_prompt,
layers=args.layers,
@@ -66,10 +66,5 @@ Key arguments:
- `--vae-use-tiling`: Enable VAE tiling for memory optimization.
- `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
- `--enable-cpu-offload`: enable CPU offloading for diffusion models.
- `--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.
- `--hsdp-replicate-size`: Number of replica groups for HSDP. Each replica holds a full sharded copy. Default 1 means pure sharding (no replication).
> ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.
@@ -80,24 +80,6 @@ def parse_args() -> argparse.Namespace:
action="store_true",
help="Enable layerwise (blockwise) offloading on DiT modules.",
)
parser.add_argument(
"--ulysses-degree",
type=int,
default=1,
help="Number of GPUs used for ulysses sequence parallelism.",
)
parser.add_argument(
"--ring-degree",
type=int,
default=1,
help="Number of GPUs used for ring sequence parallelism.",
)
parser.add_argument(
"--tensor-parallel-size",
type=int,
default=1,
help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
)
parser.add_argument(
"--cfg-parallel-size",
type=int,
@@ -110,29 +92,6 @@ def parse_args() -> argparse.Namespace:
action="store_true",
help="Disable torch.compile and force eager execution.",
)
parser.add_argument(
"--use-hsdp",
action="store_true",
help=("Enable Hybrid Sharded Data Parallel to shard model weights across GPUs. "),
)
parser.add_argument(
"--hsdp-shard-size",
type=int,
default=-1,
help=(
"Number of GPUs to shard model weights across within each replica group. "
"-1 (default) auto-calculates as world_size / replicate_size. "
),
)
parser.add_argument(
"--hsdp-replicate-size",
type=int,
default=1,
help=(
"Number of replica groups for HSDP. Each replica holds a full sharded copy. "
"Default 1 means pure sharding (no replication). "
),
)
return parser.parse_args()
@@ -169,13 +128,7 @@ def main():
# Check if profiling is requested via environment variable
profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
parallel_config = DiffusionParallelConfig(
ulysses_degree=args.ulysses_degree,
ring_degree=args.ring_degree,
cfg_parallel_size=args.cfg_parallel_size,
tensor_parallel_size=args.tensor_parallel_size,
use_hsdp=args.use_hsdp,
hsdp_shard_size=args.hsdp_shard_size,
hsdp_replicate_size=args.hsdp_replicate_size,
)
omni = Omni(
model=args.model,
@@ -0,0 +1,98 @@
# LoRA Inference(Diffusion)
This directory contains examples for using LoRA (Low-Rank Adaptation) adapters with vLLM-omni diffusion models for offline inference.
The example uses the `stabilityai/stable-diffusion-3.5-medium` as the default model, but you can replace it with other models in vLLM-omni.
## Overview
Similar to vLLM, vLLM-omni uses a unified LoRA handling mechanism:
- **Pre-loaded LoRA**: Loaded at initialization via `--lora-path` (pre-loaded into cache)
- **Per-request LoRA**: Loaded on-demand. In the example, the LoRA is loaded via `--lora-request-path` in each request
Both approaches use the same underlying mechanism - all LoRA adapters are handled uniformly through `set_active_adapter()`. If no LoRA request is provided in a request, all adapters are deactivated.
## Usage
### Pre-loaded LoRA (via --lora-path)
Load a LoRA adapter at initialization. This adapter is pre-loaded into the cache and can be activated by requests:
```bash
python -m examples.offline_inference.lora_inference.lora_inference \
--prompt "A piece of cheesecake" \
--lora-path /path/to/lora/ \
--lora-scale 1.0 \
--num-inference-steps 50 \
--height 1024 \
--width 1024 \
--output output_preloaded.png
```
**Note**: When using `--lora-path`, the adapter is loaded at init time with a stable ID derived from the adapter path. This example activates it automatically for the request.
### Per-request LoRA (via --lora-request-path)
Load a LoRA adapter on-demand for each request:
```bash
python -m examples.offline_inference.lora_inference.lora_inference \
--prompt "A piece of cheesecake" \
--lora-request-path /path/to/lora/ \
--lora-scale 1.0 \
--num-inference-steps 50 \
--height 1024 \
--width 1024 \
--output output_per_request.png
```
### No LoRA
If no LoRA request is provided, we will use the base model without any LoRA adapters:
```bash
python -m examples.offline_inference.lora_inference.lora_inference \
--prompt "A piece of cheesecake" \
--num-inference-steps 50 \
--height 1024 \
--width 1024 \
--output output_no_lora.png
```
## Parameters
### LoRA Parameters
- `--lora-path`: Path to LoRA adapter folder to pre-load at initialization (loads into cache with a stable ID derived from the path)
- `--lora-request-path`: Path to LoRA adapter folder for per-request loading
- `--lora-request-id`: Integer ID for the LoRA adapter (optional). If not provided and `--lora-request-path` is set, will derive a stable ID from the path.
- `--lora-scale`: Scale factor for LoRA weights (default: 1.0). Higher values increase the influence of the LoRA adapter.
### Standard Parameters
- `--prompt`: Text prompt for image generation (required)
- `--seed`: Random seed for reproducibility (default: 42)
- `--height`: Image height in pixels (default: 1024)
- `--width`: Image width in pixels (default: 1024)
- `--num-inference-steps`: Number of denoising steps (default: 50)
- `--output`: Output file path (default: `lora_output.png`)
## How LoRA Works
All LoRA adapters are handled uniformly:
1. **Initialization**: If `--lora-path` is provided, the adapter is loaded into cache with a stable ID derived from the adapter path
2. **Per-request**: If `--lora-request-path` is provided, the adapter is loaded/activated for that request
3. **No LoRA**: If no LoRA request is provided (`req.lora_request` is None), all adapters are deactivated
The system uses LRU cache management - adapters are cached and evicted when the cache is full (unless pinned).
## LoRA Adapter Format
LoRA adapters must be in PEFT (Parameter-Efficient Fine-Tuning) format. A typical LoRA adapter directory structure:
```
lora_adapter/
├── adapter_config.json
└── adapter_model.safetensors
```
@@ -0,0 +1,147 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import argparse
from pathlib import Path
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.lora.request import LoRARequest
from vllm_omni.lora.utils import stable_lora_int_id
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Generate images with LoRA adapters.")
parser.add_argument("--model", default="stabilityai/stable-diffusion-3.5-medium", help="Model name or path.")
parser.add_argument("--prompt", required=True, help="Text prompt for image generation.")
parser.add_argument("--seed", type=int, default=42, help="Random seed for deterministic results.")
parser.add_argument("--height", type=int, default=1024, help="Height of generated image.")
parser.add_argument("--width", type=int, default=1024, help="Width of generated image.")
parser.add_argument(
"--num-inference-steps",
type=int,
default=50,
help="Number of denoising steps for the diffusion sampler.",
)
parser.add_argument(
"--output",
type=str,
default="lora_output.png",
help="Path to save the generated image (PNG).",
)
parser.add_argument(
"--lora-path",
type=str,
default=None,
help="Path to LoRA adapter folder to pre-load at initialization (PEFT format). "
"Note: pre-loading populates the cache; you still need to pass a lora_request to activate it.",
)
parser.add_argument(
"--lora-request-path",
type=str,
default=None,
help="Path to LoRA adapter folder for per-request activation (dynamic LoRA). "
"If --lora-request-id is not provided, a stable ID will be derived from this path.",
)
parser.add_argument(
"--lora-request-id",
type=int,
default=None,
help="Integer ID for the LoRA adapter (for dynamic LoRA). "
"If not provided and --lora-request-path is set, will derive a stable ID from the path.",
)
parser.add_argument(
"--lora-scale",
type=float,
default=1.0,
help="Scale factor for LoRA weights (default: 1.0).",
)
return parser.parse_args()
def main():
args = parse_args()
model = args.model
omni_kwargs = {}
if args.lora_path:
omni_kwargs["lora_path"] = args.lora_path
print(f"Using static LoRA from: {args.lora_path}")
omni = Omni(model=model, **omni_kwargs)
lora_request = None
if args.lora_request_path:
if args.lora_request_id is None:
lora_request_id = stable_lora_int_id(args.lora_request_path)
else:
lora_request_id = args.lora_request_id
lora_name = Path(args.lora_request_path).stem
lora_request = LoRARequest(
lora_name=lora_name,
lora_int_id=lora_request_id,
lora_path=args.lora_request_path,
)
print(f"Using per-request LoRA: name={lora_name}, id={lora_request_id}, scale={args.lora_scale}")
elif args.lora_path:
# pre-loaded LoRA
lora_request_id = stable_lora_int_id(args.lora_path)
lora_request = LoRARequest(
lora_name="preloaded",
lora_int_id=lora_request_id,
lora_path=args.lora_path,
)
print(f"Activating pre-loaded LoRA: id={lora_request_id}, scale={args.lora_scale}")
sampling_params = OmniDiffusionSamplingParams(
height=args.height,
width=args.width,
num_inference_steps=args.num_inference_steps,
)
if lora_request:
sampling_params.lora_request = lora_request
sampling_params.lora_scale = args.lora_scale
outputs = omni.generate(args.prompt, sampling_params)
if not outputs or len(outputs) == 0:
raise ValueError("No output generated from omni.generate()")
if isinstance(outputs, list):
first_output = outputs[0]
else:
first_output = outputs
images = None
if hasattr(first_output, "images") and first_output.images:
images = first_output.images
elif hasattr(first_output, "request_output") and first_output.request_output:
req_out = first_output.request_output
if isinstance(req_out, list) and len(req_out) > 0:
req_out = req_out[0]
if hasattr(req_out, "images") and req_out.images:
images = req_out.images
if not images:
raise ValueError("No images found in request_output")
output_path = Path(args.output)
output_path.parent.mkdir(parents=True, exist_ok=True)
suffix = output_path.suffix or ".png"
stem = output_path.stem or "lora_output"
if len(images) <= 1:
images[0].save(output_path)
print(f"Saved generated image to {output_path}")
else:
for idx, img in enumerate(images):
save_path = output_path.parent / f"{stem}_{idx}{suffix}"
img.save(save_path)
print(f"Saved generated image to {save_path}")
if __name__ == "__main__":
main()
@@ -1,232 +0,0 @@
# MiMo-Audio Offline Inference
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).
- **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.
@@ -1,443 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
This example shows how to use vLLM for running offline inference
with the correct prompt format on MiMo-Audio-Omni.
"""
import copy
import json
import os
from typing import NamedTuple
import soundfile as sf
from message_convert import (
get_audio_data,
get_audio_understanding_sft_prompt,
get_s2t_dialogue_sft_multiturn_prompt,
get_spoken_dialogue_sft_multiturn_prompt,
get_text_dialogue_sft_multiturn_prompt,
get_tts_sft_prompt,
to_prompt,
)
from vllm import SamplingParams
from vllm.utils.argparse_utils import FlexibleArgumentParser
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniTokensPrompt
SEED = 42
MAX_CODE2WAV_TOKENS = 18192 # Maximum tokens supported by code2wav model
class QueryResult(NamedTuple):
inputs: dict
limit_mm_per_prompt: dict[str, int]
def get_codes_query_from_json(codes_path: str) -> QueryResult:
with open(codes_path, encoding="utf-8") as f:
data = json.load(f)
if isinstance(data, list):
code_final = data
elif isinstance(data, dict) and "code_final" in data:
code_final = data["code_final"]
else:
raise ValueError(
f"Unsupported codes json format in {codes_path}.\n"
"Expect a JSON list[int] or {{'code_final': list[int]}}."
)
if not isinstance(code_final, list) or not all(isinstance(x, int) for x in code_final):
raise ValueError("code_final must be a list[int].")
if len(code_final) > MAX_CODE2WAV_TOKENS:
print(f"[Warn] code_final len={len(code_final)} > {MAX_CODE2WAV_TOKENS}, truncating.")
code_final = code_final[:MAX_CODE2WAV_TOKENS]
return QueryResult(
inputs=OmniTokensPrompt(
prompt_token_ids=code_final,
multi_modal_data=None,
mm_processor_kwargs=None,
),
limit_mm_per_prompt={},
)
def get_tts_sft(
text="The weather is so nice today.",
instruct=None,
read_text_only=True,
prompt_speech=None,
audio_list=None,
):
res = get_tts_sft_prompt(
text,
instruct=instruct,
read_text_only=read_text_only,
prompt_speech=prompt_speech,
)
prompt = to_prompt(res)
final_prompt = {
"prompt": prompt,
}
if audio_list is not None:
final_prompt.update(
{
"multi_modal_data": {
"audio": audio_list,
},
}
)
return final_prompt
def get_audio_understanding_sft(audio_path, text="", thinking=False, use_sostm=False):
audio_list = []
audio_list.append(get_audio_data(audio_path))
res = get_audio_understanding_sft_prompt(
input_speech=audio_path, input_text=text, thinking=thinking, use_sostm=use_sostm
)
prompt = to_prompt(res)
final_prompt = {
"prompt": prompt,
"multi_modal_data": {
"audio": audio_list,
},
}
return final_prompt
def get_spoken_dialogue_sft_multiturn(message_list, system_prompt=None, ref_audio_path=None, audio_list=None):
res = get_spoken_dialogue_sft_multiturn_prompt(
message_list, system_prompt=system_prompt, prompt_speech=ref_audio_path
)
prompt = to_prompt(res)
final_prompt = {
"prompt": prompt,
"multi_modal_data": {
"audio": audio_list,
},
}
return final_prompt
def get_speech2text_dialogue_sft_multiturn(message_list, thinking=False, audio_list=None):
res = get_s2t_dialogue_sft_multiturn_prompt(
message_list,
thinking=thinking,
)
prompt = to_prompt(res)
final_prompt = {
"prompt": prompt,
"multi_modal_data": {
"audio": audio_list,
},
}
return final_prompt
def get_text_dialogue_sft_multiturn(
message_list,
):
res = get_text_dialogue_sft_multiturn_prompt(
message_list,
)
prompt = to_prompt(res)
final_prompt = {
"prompt": prompt,
}
return final_prompt
query_map = {
"tts_sft": get_tts_sft,
"tts_sft_with_instruct": get_tts_sft,
"tts_sft_with_audio": get_tts_sft,
"tts_sft_with_natural_instruction": get_tts_sft,
"audio_trancribing_sft": get_audio_understanding_sft,
"audio_understanding_sft": get_audio_understanding_sft,
"audio_understanding_sft_with_thinking": get_audio_understanding_sft,
"spoken_dialogue_sft_multiturn": get_spoken_dialogue_sft_multiturn,
"speech2text_dialogue_sft_multiturn": get_speech2text_dialogue_sft_multiturn,
"text_dialogue_sft_multiturn": get_text_dialogue_sft_multiturn,
}
def main(args):
model_name = args.model_name
# Get paths from args
text = getattr(args, "text", None)
audio_path = getattr(args, "audio_path", None)
instruct = getattr(args, "instruct", None)
# Get the query function and call it with appropriate parameters
query_func = query_map[args.query_type]
omni_llm = Omni(
model=model_name,
stage_configs_path=args.stage_configs_path,
log_stats=args.enable_stats,
log_file=("omni_llm_pipeline.log" if args.enable_stats else None),
init_sleep_seconds=args.init_sleep_seconds,
batch_timeout=args.batch_timeout,
init_timeout=args.init_timeout,
shm_threshold_bytes=args.shm_threshold_bytes,
)
thinker_sampling_params = SamplingParams(
temperature=0.0,
top_p=1.0,
top_k=-1,
max_tokens=2048,
seed=SEED,
logit_bias={},
repetition_penalty=1.1,
)
code2wav_sampling_params = SamplingParams(
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,
code2wav_sampling_params,
]
# Build query result based on query type
# Notice: The audio files used in this example are available at: https://github.com/XiaomiMiMo/MiMo-Audio/tree/main/examples
if args.query_type == "tts_sft":
# python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH} --query-type tts_sft
query_result = query_func(text=text, read_text_only=True)
elif args.query_type == "tts_sft_with_instruct":
# python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH} --query-type tts_sft_with_instruct --instruct "Speak happily in a child's voice"
query_result = query_func(text=text, instruct=instruct, read_text_only=True)
elif args.query_type == "tts_sft_with_audio":
# python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH} --query-type tts_sft_with_audio --audio_path "./spoken_dialogue_assistant_turn_1.wav"
audio_list = [get_audio_data(audio_path)]
query_result = query_func(text=text, read_text_only=True, prompt_speech=audio_path, audio_list=audio_list)
elif args.query_type == "tts_sft_with_natural_instruction":
# python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH} --query-type tts_sft_with_natural_instruction --text "In a panting young male voice, he said: I can't run anymore, wait for me!"
query_result = query_func(text=text, read_text_only=False)
elif args.query_type == "audio_trancribing_sft":
# python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH} --query-type audio_trancribing_sft --audio_path "./spoken_dialogue_assistant_turn_1.wav"
audio_path = "spoken_dialogue_assistant_turn_1.wav"
text = "Please transcribe this audio and repeat it once."
query_result = query_func(text=text, audio_path=audio_path, use_sostm=True)
elif args.query_type == "audio_understanding_sft":
# python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH} --query-type audio_understanding_sft --text "Summarize the audio." --audio_path "./spoken_dialogue_assistant_turn_1.wav"
query_result = query_func(text=text, audio_path=audio_path)
elif args.query_type == "audio_understanding_sft_with_thinking":
# python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH} --query-type audio_understanding_sft_with_thinking --text "Summarize the audio." --audio_path "./spoken_dialogue_assistant_turn_1.wav"
query_result = query_func(text=text, audio_path=audio_path, thinking=True)
elif args.query_type == "spoken_dialogue_sft_multiturn":
# python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH} --query-type spoken_dialogue_sft_multiturn --audio_path "./prompt_speech_zh_m.wav"
first_turn_text_response = "我没办法获取实时的天气信息。不过呢,你可以试试几个方法来查看今天的天气。首先,你可以用手机自带的天气功能,比如苹果手机的天气应用,或者直接在系统设置里查看。其次,你也可以用一些专业的天气服务,像是国外的AccuWeather、Weather.com,或者国内的中国天气网、墨迹天气等等。再有就是,你还可以在谷歌或者百度里直接搜索你所在的城市加上天气这两个字。如果你能告诉我你所在的城市,我也可以帮你分析一下历史天气趋势,不过最新的数据还是需要你通过官方渠道去获取哦。"
audio_list = []
s1_audio_path = "weather_of_today.mp3"
s2_audio_path = "spoken_dialogue_assistant_turn_1.wav"
s3_audio_path = "beijing.mp3"
audio_list.append(get_audio_data(audio_path))
audio_list.append(get_audio_data(s1_audio_path))
audio_list.append(get_audio_data(s2_audio_path))
audio_list.append(get_audio_data(s3_audio_path))
message_list = [
{"role": "user", "content": s1_audio_path},
{"role": "assistant", "content": {"text": first_turn_text_response, "audio": s2_audio_path}},
{"role": "user", "content": s3_audio_path},
]
query_result = query_func(message_list, system_prompt=None, ref_audio_path=audio_path, audio_list=audio_list)
elif args.query_type == "speech2text_dialogue_sft_multiturn":
# python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH} --query-type speech2text_dialogue_sft_multiturn
s1_audio_path = "weather_of_today.mp3"
s2_audio_path = "beijing.mp3"
audio_list = []
audio_list.append(get_audio_data(s1_audio_path))
audio_list.append(get_audio_data(s2_audio_path))
message_list = [
{"role": "user", "content": s1_audio_path},
{
"role": "assistant",
"content": "Hello, I can't get real-time weather information. If you can tell me your city, I can help you analyze historical weather trends, but for the latest data, you'll need to get it through official channels.",
},
{"role": "user", "content": s2_audio_path},
]
query_result = query_func(message_list, thinking=True, audio_list=audio_list)
elif args.query_type == "text_dialogue_sft_multiturn":
# python3 -u end2end.py --stage-configs-path ${config_file} --model ${MODEL_PATH} --query-type text_dialogue_sft_multiturn
message_list = [
{"role": "user", "content": "Could you recommend some tourist attractions in China?"},
{"role": "assistant", "content": "Hello, which city would you like to travel to?"},
{"role": "user", "content": "Beijing"},
]
query_result = query_func(message_list=message_list)
else:
raise ValueError(f"Invalid query type: {args.query_type}")
prompts = [copy.deepcopy(query_result) for _ in range(args.num_prompts)]
print("prompts", prompts)
omni_outputs = omni_llm.generate(prompts, sampling_params_list)
output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav
if args.query_type is not None:
output_dir = os.path.join(output_dir, args.query_type)
os.makedirs(output_dir, exist_ok=True)
for stage_outputs in omni_outputs:
if stage_outputs.final_output_type == "text":
for output in stage_outputs.request_output:
request_id = output.request_id
text_output = output.outputs[0].text
# Save aligned text file per request
prompt_text = output.prompt
out_txt = os.path.join(output_dir, f"{request_id}.txt")
lines = []
lines.append("Prompt:\n")
lines.append(str(prompt_text) + "\n")
lines.append("vllm_text_output:\n")
lines.append(str(text_output).strip() + "\n")
try:
with open(out_txt, "w", encoding="utf-8") as f:
print("lines", lines)
f.writelines(lines)
except Exception as e:
print(f"[Warn] Failed writing text file {out_txt}: {e}")
print(f"Request ID: {request_id}, Text saved to {out_txt}\n")
elif stage_outputs.final_output_type == "audio":
for output in stage_outputs.request_output:
request_id = output.request_id
audio_tensor = output.outputs[0].multimodal_output.get("audio")
if audio_tensor is None:
continue
output_wav = os.path.join(output_dir, f"{request_id}.wav")
# Convert to numpy array and ensure correct format
audio_numpy = audio_tensor.float().detach().cpu().numpy()
# Ensure audio is 1D (flatten if needed)
if audio_numpy.ndim > 1:
audio_numpy = audio_numpy.flatten()
# Save audio file with explicit WAV format
sf.write(output_wav, audio_numpy, samplerate=24000, format="WAV")
print(f"Request ID: {request_id}, Audio saved to {output_wav}")
def parse_args():
parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
parser.add_argument(
"--model-name",
"-m",
type=str,
default="XiaomiMiMo/MiMo-Audio-7B-Instruct",
help="Backbone LLM path.",
)
parser.add_argument(
"--text",
"-t",
type=str,
default="The weather is so nice today.",
help="input text",
)
parser.add_argument(
"--query-type",
"-q",
type=str,
default="tts_sft",
choices=query_map.keys(),
help="Query type.",
)
parser.add_argument(
"--audio-path",
"-a",
type=str,
default=None,
help="Path to local audio file. If not provided, uses default audio asset.",
)
parser.add_argument(
"--instruct",
type=str,
default=None,
help="instruct",
)
parser.add_argument(
"--enable-stats",
action="store_true",
help="Enable writing detailed statistics (default: disabled)",
)
parser.add_argument(
"--init-sleep-seconds",
type=int,
default=20,
help="Sleep seconds after starting each stage process to allow initialization (default: 20)",
)
parser.add_argument(
"--batch-timeout",
type=int,
default=5,
help="Timeout for batching in seconds (default: 5)",
)
parser.add_argument(
"--init-timeout",
type=int,
default=5000,
help="Timeout for initializing stages in seconds (default: 5000)",
)
parser.add_argument(
"--shm-threshold-bytes",
type=int,
default=65536,
help="Threshold for using shared memory in bytes (default: 65536)",
)
parser.add_argument(
"--output-dir",
default="./output_audio",
help="Output audio wav directory.",
)
parser.add_argument(
"--output-wav",
default="output_audio",
help="[Deprecated] Output wav directory (use --output-dir).",
)
parser.add_argument(
"--num-prompts",
type=int,
default=1,
help="Number of prompts to generate.",
)
parser.add_argument(
"--sampling-rate",
type=int,
default=24000,
help="Sampling rate for audio.",
)
parser.add_argument(
"--stage-configs-path",
type=str,
default="../../../model_executor/stage_configs/mimo_audio.yaml",
help="Path to a stage configs file.",
)
return parser.parse_args()
if __name__ == "__main__":
args = parse_args()
main(args)
@@ -1,56 +0,0 @@
{
"messages": [
{
"role": "system",
"content": [
{
"type": "text",
"text": "Your voice should be:"
},
{
"type": "audio_url",
"audio_url": {
"url": "data:audio/wav;base64,UklGRiRKAwBXQVZFZm10IBAAAAABAAEAwF0AAIC7AAACABAAZGF0YQBKAwAFAP3/CAD7/wkAAAD9/wgA+v8JAP//AQADAP//CAD+/wEABAABAP7/BgD8//7/BgD0/wYA/f/8/wMA9f8GAP///P8DAPz/AAAFAP3//v8CAP7////+////+/8IAP//AQAAAAMABgD6/wUA+v8DAAcA+/8AAAMAAAAAAP3//v8DAAMACgD8/wQACQD9/wUA/f8BAAcABQADAAAABAADAAEA+v8AAAIAAAAFAPz//f8DAAAA/f8AAAMA/v8FAP7/9v8GAP7//f8EAAAAAQAIAPv/+/8AAP3////+//3/AAAHAPr//v8BAPz/AwD///n/AAACAPr//f/8//7/AAD//wAA+f8GAPz/9v8CAPX/BgACAPv/AQD//wQA///7/wAAAAAHAAQA/f8FAAAAAgAAAPv/BAAGAAAABgACAP7/DQD7//3/BAD//wkA//8DAPz/AwAHAPf/AwAAAAAABAD7//f/AwD8/wAA/f/5/wQA/P8AAPr//P8AAAAA//8AAP////8DAPL/AQD9//r/CwD4/wEAAgD9//7//P/9//7/AAD//wAA/v8DAPr/+v8CAPr///8CAPz///8GAPr/AAAEAP//AgADAP//AAAHAPz/BgD//wUACwD//wYAAgABAAQABQD8/wgABgACAAUA//8CAAQAAAADAP7/AQAIAPr/AwD+//v/AgACAPv/AQAJAPf/AQD+//n/AgAAAP///f8CAAAA/P8DAPf/+/8HAPv///////7//f8AAP3//P8BAPr/AAD6//z/AAD5/wIAAAD7/wIAAgD4/wAA/v/6/wQAAQAAAAAABwAAAAAAAwD+/wIABQADAAEABAADAAAAAwABAAAABAAGAAEAAgAGAAAABQD///7/BQAAAAMAAQAAAAIABAD8/wMA/////wQA/f8AAAEAAQD//wMA+/8FAPr///////P/CwDz/wQA/P/6/wIA8/8AAPH//P/+//f//v/+//n//P/8//L//v/7//r////7//3//v/5//v/AAD5/wMAAAD8/wYA/P///////////wAABwD6/wwAAQD+/wUA/P8GAP7/AwD+/wUABwD+/wUA//8FAAMAAgD//wUAAAAEAAcA+/8KAAUAAgAEAAAA//8HAAAA+v8EAAIAAQAEAPz/+/8EAP7//////wAAAQABAP//+///////AAD9/wEA/f8DAAAA+P8AAPj/AQD+//v/AQD8/wIA/v/6/wAA/////wIA/P///wMA+/8BAPj///8CAAAAAwD+/wcAAAAFAP3/AAAFAP7/AwABAP7/AAACAPv/AwD+/wMABAD9/wUA/v8AAAEA/v8AAAIAAQAAAAAAAAD8/wIA/v/8/wMAAAD///3/AAD5//7/AwD2/wgAAAD8/wQA+f8CAP7/AAAAAP//AwD//wAA/v//////AAAAAAAAAwD9/wIA+//9/wQA+/8DAAAA//8AAP///f/+/wAAAAABAP//BgD8/wIAAgD3/wwA//8AAAkA/f8EAAAA//8BAAAAAwADAP//BQD///z/AwD8/wEAAwAAAAEAAQAAAAAA/v8BAAAAAQAGAP7/AAD///3//v8AAAAAAgAGAP//AgAAAP3////9/wAAAQACAAAAAQD///3/AQD9/wAABwD//wMABAD8/wEAAwD+/wIABQD+/wEAAwD6////BAD5/wYAAAD+/wgA+f8CAP3///8DAP7/AgAAAAAA/v8EAPv/AAAHAPv/CQACAP7/BgD//wIAAAADAAAAAgAEAPz/BQD8/wQA/P/+/wMA+f8EAP3//P8AAP////8AAPv/AgD8/wAAAgDz/wgA/f/7/wQA+P8BAAAA+v8BAP7/AAACAP3/AAACAP7/AQAEAPj/BwAAAAAABgD7/wcA+v8CAP7/+/8IAP3/AQACAAAAAQACAAAA/v8CAAcA/f8HAAIA/v8IAPv/AQAAAAAABAD+/wQAAgD//wMA/P8AAAAA//8BAPr/BQD9/wAA///4/wIA+/8AAP3//v8EAPr/AAD9//r/BQD3/wEA+v/6/wYA8/8BAPz/9v8CAPn/+f/+//3//v/+//3//v///wAA/P/7/wAAAAD//wAA//8AAP7/AgD4////BAD7/wYAAAADAAAAAwAAAPr/BQD6/wUABAD//wcA/f8EAPz//P8FAPv/BQADAP//BgAAAP//AQAAAAAAAgAAAAEAAwAAAAAA//8AAAAAAAAAAP7/BAD///7/AgD8/wAAAQD8////AQD9/wQA/v///wIA//8BAPv/AQD9/wMA/v/8/wQA+f8DAP3/9/8EAAAA/P8GAPz/AwAEAP//AAD8/wQAAAAGAAAAAgAFAAEABAD8/wEAAQADAAAAAgADAAAACAD7//7/BQD9/wUAAAD+/wcAAAABAAAA/f8DAAMAAAAAAAEA/v8AAAAA+f8AAAEA//8AAP//AAD7/wUA9//6/wQA9/8BAP3//f8AAP7/AAD7//3/AwD+//v//v/7//3/BAD3//r/CAD8/wIA/v/6//3/AQD+//j/AAAAAAgAAAD///v//f8BAP7//v/9/wEA//8EAAAA//8BAAEAAAD+/////f8AAP7/AgD9/wQABQD4/wIA/f/5/wAA/P/9/wEAAQABAPf/AAD8//7/AQD5/wAA/f8BAPf/+f/7//b/AwADAPz///8BAPn//P/4/////P8AAAAA+P8JAAAA/f/8//3/AAABAAAA8v8EAAMAAAACAPn/BAAOAAcA+//t////DwAEAAIA/f8HAAYADAD7//f/BgD//wwA//8BAAEADwAEAPv/BQD1/xEA//8DAP//9/8LAPn/AgABAAMABAAEAPz//f8FAPb/BAD+//r/CwAAAP7/BQD4////BAD4/wEA/v8AAAAA/v/+//z/BAABAP//BQD+//z/AQD7/wAA/P8DAAUAAAAHAP//AQD9////AAABAP//AAACAAQACAD+/wYA//8DAAcA/v8CAAEABQADAAIAAgACAAUABQAAAAAAAgAAAAAAAwABAAAABQAAAAYABgAAAAEA/v8EAAAA/P8GAP7/BAAGAPn/AQABAAIA/f8EAPn///8GAPT/BgD5/wAAAgD8/wQA+v8BAPv/9/8AAP3/AAAAAPz/AQD9//r////6////AAD9/////f8AAPz/AAD9//z/CQD5/wUA+v/7/wIA9P8FAPn/AQADAP3//v8AAP3/+/8DAPr/AQABAAAA/f8BAAAA9/8EAP7///8GAPr///8BAPz/AQD9/wAABQD//wEAAAD9/wMAAQD7/wAAAwAAAAMAAAD+/wMAAQAAAP3/AAABAP//AgD+/wEAAQAHAPv///8EAPb/CQD//wAAAAAAAAUA9v8JAPf///8FAPr/AwD9/wIA+v8DAP//+P8GAP////8EAPv/+/8EAPv/AAD+//7/AQD//wEA+/8AAAAA/v8CAP///P8GAP////8DAPv/BgAAAAAAAgD//wEA//8CAP7/AgAAAP//BgAAAAEAAgAAAP//BwD6/wIABAD6/woA+/8EAAAAAAAEAPv/BgD//wIAAAABAAMAAAAHAPf/BwAAAPr/BwD4/wQAAAAAAAAAAAD//wAAAAD7/wcA/v8AAAAA///+/wIA///7/wkA+/8GAAAA/f8BAP//AgD9/wYA//8FAAQA//8CAPz/AQAFAPr/AwAAAP//BgD5//7//v8CAPv/AAD9//b/CAD4//z//v/6/wAAAAD+//z/AAABAP///v////3/BwABAAAAAAAAAAUA/f8DAPb/AAAIAPv/AwD9////BAAAAPf/+/8AAPv/AQD7//v/AAABAP7/+v8BAPn/AQACAPr/AQACAAEA/f8FAPv/AAAFAPj/AAABAPv//P/+//n//v////r//f/+//r/AAD0//b//f/0/wMA9//+//7//f8BAPP/BAD6////AwD8/wEAAAAKAPf/BwACAPz/DAD9//3/AwAGAPn/BwD+//3/DgD7/wAAAAD7/wYA/f/5/wQA+v8GAAgA+P8IAAMAAwAKAP//AQAFAAMABQAAAAUABQAGAAcA//8IAPr/AgAAAPX/BAD+//3/BAAAAPX/AwD5//T/AwD3//n/BgD6//z/BAD3/wAAAAD+/wAA//8AAAAA/v/+//7///8BAAAAAAABAP7//v/+//b/AAD8/wAAAAD+/wEA/f8CAPr/AQD7/wAABQD7/wcA//8FAAcAAgAIAAUABAAGAAcAAwAGAAQABAAKAAkABQAIAAUABwADAAMAAAACAAMAAgAFAP7/DAABAAgA/v/+/wYA/P8KAPj/BAAGAP7/CQAAAAEACAAEAAUA/P///wIA+/8DAP7///8FAP//AAD7/wAA/f/2/wIA9P/9/wMA8v///wAA+/8DAAEA/P8CAP7//v/5//r/CAAAAAkABgAAAAUABwD+//3/CwD//w4ABwAAAAkAAgAJAPz/BwALAAEAEgD//wAACgABAAcAAgAFAAcABwAGAAAABAAEAAoAAwAIAAsABQAMAAAAAQACAAkABwAGAA4ACwAKABMAAgAAAAoA/f8JAAAABAADAAsACQD9/wQA+/8AAPf//v/z//r/AQDv//3/8f/w//r/8f/0/+//9P/u/+j/8v/g/+7/8f/r//X/8f/x/+7/7v/m/+v/7v/3//D/8//4/+3//v/z//b/9//7//n/9f8BAPT//P8AAPj/AAADAAAAAAD+/wIA+f/8/wQA+v8DAAcA/v8AAA8AAAAFABEA//8NAA4ABwAHABAAEgAKABgAFAARABYAFAARABIAHAAXABIAIAAYABQAIgAUABEAIQAXABUAGgAVABMAFwAVAAgADQAJAAkABAAJAAkAAAAKAPz/+f/3//z/8P/2/wAA8/8CAP7/8//0//j/7v/t/Line truncated
}
}
]
},
{
"role": "user",
"content": [
{
"type": "audio_url",
"audio_url": {
"url": "data:audio/wav;base64,UklGRgrVAABXQVZFZm10IBAAAAABAAEAwF0AAIC7AAACABAAZGF0YebUAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAEAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA//8AAAAAAAAAAAAAAAAAAAAAAAAAAP////8AAAAAAAAAAAAAAAD//wAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAP////8AAAAAAAAAAAAA///p/9T/xv/G/8b/yf/K/8r/zP/M/83/zf/O/83/zf/P/8//z//O/87/z//P/8//z//P/8//z//P/8//z//P/8//z//Q/9D/0P/Q/9D/0P/R/9H/0f/R/9H/0f/R/9H/0f/R/9L/0v/S/9L/0v/S/9P/0//T/9P/0//T/9P/0//U/9T/1P/U/9T/1P/V/9T/1f/V/9X/1f/W/9b/1v/W/9b/1v/W/9b/1//X/9f/1//Y/9j/2P/Y/9j/2P/Z/9n/2f/Z/9n/2f/a/9r/2v/a/9v/2//b/9v/2//b/9v/3P/c/9z/3f/c/93/3f/d/93/3f/e/97/3v/e/9//3//f/9//4P/g/+D/4P/g/+D/4P/h/+H/4f/h/+H/4f/i/+L/4v/j/+P/4//j/+P/4//k/+P/5P/k/+T/5P/l/+X/5f/l/+X/5f/l/+X/5f/l/+X/5v/n/+f/5//o/+j/6P/o/+j/5v/n/+f/5//o/+j/6P/o/+n/6f/p/+j/6P/o/+j/6P/p/+n/6f/p/+n/6v/q/+r/6v/q/+r/6v/q/+r/6v/q/+r/6//r/+r/6v/r/+r/6v/p/+r/6v/r/+r/6v/r/+z/7P/s/+v/7P/r/+v/6//s/+v/6//r/+v/6//t/+z/7f/t/+3/7f/r/+z/6//s/+v/7P/s/+z/7P/s/+z/7P/t/+3/7f/s/+z/6//s/+r/6//q/+z/7P/t/+3/7v/u/+3/7f/t/+z/7P/s/+z/7P/t/+3/7f/t/+3/7f/t/+3/7f/t/+z/7P/s/+z/7f/t/+3/7f/u/+3/7f/t/+3/7f/t/+3/7f/u/+3/7f/t/+3/7f/t/+3/7f/t/+3/7P/r/+z/6//t/+z/7f/t/+7/7f/u/+7/7f/t/+z/7P/t/+z/7P/t/+3/7P/t/+7/7v/v/+7/7v/t/+z/7P/t/+3/7v/u/+7/7v/v/+//7v/u/+3/7P/u/+z/7P/s/+z/7f/u/+7/7v/v/+7/7v/s/+z/7P/t/+3/7f/u/+7/7f/u/+3/7v/t/+3/7f/t/+3/7P/t/+z/7f/u/+3/7//u/+7/7f/t/+3/7f/u/+//7v/v/+7/7v/t/+7/7f/v/+7/7v/t/+7/7f/t/+3/7v/u/+7/7v/v/+3/7v/t/+z/7P/r/+z/6//t/+r/7P/t/+7/7v/u/+7/7f/s/+z/6//s/+v/7P/t/+3/7//w//D/7v/w/+7/7v/s/+3/6//t/+v/7f/u/+//7//u/+//7v/u/+r/6//q/+v/6//s/+7/7v/w/+//7//v/+3/7f/s/+v/7P/s/+3/7f/u/+7/7v/u/+3/7P/r/+z/7f/s/+z/7P/t/+3/7P/t/+3/7v/s/+3/7f/t/+z/6//r/+z/7P/s/+z/7P/t/+z/7f/t/+z/7v/v/+7/7P/s/+z/6//u/+3/7v/u/+7/7f/u/+3/7f/t/+v/7f/r/+v/6//s/+3/7P/u/+7/7v/v/+7/7v/u/+z/7P/r/+3/6//t/+3/7v/t/+7/7v/v/+7/7v/t/+3/7P/t/+7/7v/v/+3/7v/t/+3/7f/s/+z/7f/s/+3/7f/t/+7/7v/u/+3/7//u/+7/7f/s/+z/7P/s/+v/7P/t/+3/7f/v/+7/7v/u/+3/7f/t/+3/7f/u/+7/7v/u/+7/7//u/+7/7v/t/+3/7f/t/+3/7f/u/+//7//v/+//7v/u/+7/7v/u/+7/7v/u/+7/7v/t/+7/7v/u/+7/7v/u/+7/7f/t/+3/7f/t/+7/7f/t/+3/7P/u/+3/7v/t/+7/7v/u/+7/7v/v/+7/7v/v/+7/7v/u/+7/7v/u/+3/7f/u/+3/7v/u/+7/7f/u/+3/7v/u/+7/7f/t/+7/7v/v/+7/7v/u/+3/7v/t/+3/7f/t/+3/7f/u/+7/7v/u/+3/7v/t/+3/7P/t/+3/7f/u/+3/7v/u/+//7v/u/+7/7v/u/+7/7v/u/+7/7f/u/+3/7f/u/+7/7f/u/+3/7f/t/+3/7f/t/+3/7f/s/+z/7f/t/+3/7f/t/+3/7P/t/+3/7f/t/+z/7P/s/+3/7P/t/+3/7f/u/+3/7f/t/+3/7f/t/+7/7P/t/+3/7f/t/+3/7f/t/+3/7f/t/+3/7f/t/+3/7f/t/+3/7f/s/+z/7f/t/+3/7f/t/+3/7f/t/+z/7P/t/+3/7f/t/+3/7f/t/+3/7f/t/+3/7f/t/+z/7P/s/+3/7f/t/+3/7f/t/+3/7P/s/+z/7P/t/+z/7P/s/+z/7P/s/+3/7P/t/+3/7f/t/+3/7f/t/+z/7P/t/+z/7P/s/+z/7P/s/+z/7P/t/+z/7P/t/+z/7P/s/+z/7P/s/+z/7P/t/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+v/6//s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7f/s/+z/7P/s/+z/7P/s/+v/6//r/+z/6//s/+z/7P/s/+z/7P/s/+z/7P/s/+z/6//s/+z/7P/s/+v/6//s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/t/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+3/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/t/+z/7P/s/+z/7P/t/+z/7P/s/+z/7P/t/+z/7P/s/+z/7P/s/+z/7P/s/+z/7f/t/+3/7f/t/+3/7f/t/+3/7f/t/+3/7f/t/+3/7f/s/+3/7P/t/+3/7P/s/+z/7P/t/+z/7f/t/+3/7f/t/+3/7f/u/+z/7f/t/+3/7f/t/+3/7f/u/+3/7f/t/+z/7P/s/+z/7P/t/+z/7v/t/+7/7f/t/+7/7f/v/+z/7f/r/+3/7P/t/+3/7f/u/+7/7v/v/+7/7v/u/+z/7f/s/+7/7f/u/+3/7f/v/+7/7//u/+7/7//t/+//7v/u/+z/7P/s/+7/7f/w/+3/7f/u/+//7f/u/+7/7//s/+//6//w/+3/7//u/+3/7v/s/+7/7P/u/+//7v/x/+z/7//t/+3/7f/u/+//7//u/+7/6//w/+3/8f/t/+//7v/v/+7/7v/u/+//7v/w/+3/7v/t/+//7//t/+7/7v/w/+3/7P/v/+3/7//t/+//7v/w/+7/7//y/+3/7v/v/+3/7//t/+7/7P/v/+3/9P/w/+3/6//r//D/7//w/+3/8v/t/+7/7f/v/+3/7f/v//D/7f/t//D/7v/u/+3/8P/v/+7/7v/w/+z/7//r//L/7v/x/+3/7v/u/+7/7//t/+//7f/v/+7/7f/v/+7/7f/u/+7/7f/v/+//8P/u/+3/7//u/+//7f/v/+z/7f/v/+3/7v/p//D/7f/x/+7/7v/s/+3/8P/v/+//7f/r/+7/7P/t/+3/7v/t/+3/7f/u/+7/7//t/+7/7f/u/+3/6//s/+z/7f/t/+3/7//s/+//7f/y/+z/7v/s/+3/7P/v/+3/6//r/+3/7P/u/+//7f/t/+7/7v/v/+z/7f/r/+//7f/w/+z/7f/v/+v/7f/s/+7/7v/v/+z/7P/u/+z/7//s/+3/7f/u/+7/7f/w/+z/7v/t/+3/7v/t/+3/7P/u/+7/7v/s/+z/7f/v/+3/7v/t/+3/7f/t/+7/7P/t/+3/7v/v/+z/7f/u/+3/7v/v/+z/7v/t/+//7f/u/Line truncated
}
}
]
},
{
"role": "assistant",
"content": [
{
"type": "text",
"text": "我没办法获取实时的天气信息。不过呢,你可以试试几个方法来查看今天的天气。首先,你可以用手机自带的天气功能,比如苹果手机的天气应用,或者直接在系统设置里查看。其次,你也可以用一些专业的天气服务,像是国外的AccuWeather、Weather.com,或者国内的中国天气网、墨迹天气等等。再有就是,你还可以在谷歌或者百度里直接搜索你所在的城市加上天气这两个字。如果你能告诉我你所在的城市,我也可以帮你分析一下历史天气趋势,不过最新的数据还是需要你通过官方渠道去获取哦。"
},
{
"type": "audio_url",
"audio_url": {
"url": "data:audio/wav;base64,UklGRiQiKQBXQVZFZm10IBAAAAABAAEAwF0AAIC7AAACABAAZGF0YQAiKQAGAAcABwAIAAkACQAJAAoACwAMAAwADQAOAA4ADwAPABAAEAAQABAAEQAQABAAEQARABAAEQARABEAEAAQABEAEQAQABAAEAAPAA4ADwAOAA0ADQANAAwACwAJAAYABQAEAAIAAQAAAP7//f/7//r/+v/5//b/9v/0//X/9f/0//X/9P/0//L/9P/y//D/9P/z//T/8v/x//L/8v/x//H/8P/v/+//7v/u/+3/7P/t/+3/7P/t/+7/7v/u//D/8f/z//X/9f/3//n/+//7//3///8AAAEABAAGAAgACgAMAA4AEAASABMAFAATABQAFAAUABUAFQATABIAEgARABEAEAAOAAwACwAJAAcABgAEAAMAAgABAAEAAAAAAAAA//////7//f/8//v/+f/4//j/9//3//f/9v/2//b/9f/1//b/9f/1//b/9v/2//b/9f/1//X/9P/z//P/9P/z//P/8//0//T/9P/0//X/9v/2//f/+P/4//n/+v/7//z//f/+//7///8AAAEAAgACAAMABAAFAAYABwAJAAoADAANAA4ADwAPAA8AEAAQABEAEgASABMAFAAVABYAFwAXABYAFgAWABYAFgAVABQAEwATABMAEgASABIAEQARABAADwAOAA4ADQANAAsACgAIAAcABAACAAEA///9//v/+f/4//b/9f/z//H/8f/v/+3/7P/r/+n/5//m/+X/5f/k/+T/5P/l/+X/5//o/+r/7P/u//D/8v/1//f/+f/7//z//v8AAAEAAgAEAAQABQAHAAUABQAGAAQAAwACAAEAAgAAAAEAAQD+/wAA/v8AAP7//P////7//f/9/wAA/v/9//7//v///wAA//////////8AAAEAAwACAAIAAwADAAMAAgACAAIAAwACAAEAAAABAAEAAAD////////////////+/////v/9//3//P/8//z//f/9//3//v/+//7///8AAP////////////////////////7//v/9//3//f/8//r/+P/3//j/9//3//b/9f/2//b/9f/1//b/9//3//f/9//3//f/9//3//f/9//3//f/9//3//f/9//3//f/9//3//f/9//4//j/+P/5//n/+f/5//n/+f/5//n/+f/6//r/+//7//z//P/+////AAACAAMABQAGAAcACAAIAAkACQAJAAkACQAKAAoACQAJAAcABgAEAAMAAQD///3/+//6//n/+P/3//b/9v/1//b/9v/2//b/9v/2//b/9//4//n/+v/8//7//v///wEAAwAEAAUABwAJAAoACgALAA4ADAANAA0ADwAOAA0ADwAPAA8ADgARABEAEQARABEAEAAQABEAEQARABAAEAAQABEAEQAQABAADwAPAA4ADQAOAAwADQANAAwADQANAAsADAALAAsACgAJAAkACAAHAAcABQAFAAUAAgABAAEAAAD///7//P/6//n/9v/z//P/8v/w/+//7f/r/+v/6v/p/+r/6f/o/+j/6P/o/+j/6f/p/+n/6v/r/+z/7P/t/+7/7//w//H/8//1//b/+P/5//r//P/9////AAABAAIAAwAFAAYABwAIAAkACwALAAwADQAOAA4ADwAQABEAEgASABMAFAAVABUAFQAVABYAFQAVABUAFAATABIAEQAPAA4ADQAMAAsACgAJAAcABQADAAAA///8//r/+P/2//X/8//x/+//7v/t/+z/6//r/+r/6v/q/+v/6//s/+z/7f/t/+7/7//w//H/8v/z//T/9f/2//f/9//4//j/+f/6//r/+//7//z//P/9//7//v///wAAAAAAAAEAAQACAAMABAAGAAcACQAKAAwADQAPABAAEQASABIAEgATABUAFwAYABkAGgAbAB0AHgAfAB4AHgAdABwAGwAaABkAGAAWABQAEgAQAA0ACwAJAAcABQADAAEA/v/8//r/+P/2//T/8//x//D/7//t/+z/6//r/+v/6//r/+v/6//s/+3/7f/t/+7/7v/v/+//8P/w//D/8f/w//D/7//w//H/8f/w//D/8f/y//L/8v/y//P/8//1//f/+P/5//z//f/8/wAA//8DAAEAAAAEAAIAAQADAAMAAQAAAAEAAAABAAEA/////////v/9//7////9//z//P/9//3//P/8//z//f/+//7//v////////8AAAAAAQACAAIAAgACAAIAAgABAAEAAQABAAIAAQAAAAAAAAD///7//v/9//3/+//6//r/+v/6//r/+f/6//r/+v/5//n/+v/7//v/+v/6//v//f/+//7//v//////AAD///////////7//f/8//v/+//6//r/+f/5//n/+v/6//r/+v/6//z//P/9//3///8AAAEAAgADAAUABQAGAAYABwAIAAkACQAKAAsADAAMAAwADAAMAA0ADAAMAAsACwAKAAkACAAHAAUAAwACAAEAAAD+//3/+//6//j/9//1//T/8//y//L/8//z//T/9f/1//b/9v/2//f/+P/4//n/+f/5//r//P/9////AAACAAMABAAFAAYACAAJAAoACwANAA4ADQAOABEAEAARABEAFAATABMAFQAVABYAFQAXABYAFgAWABUAFAATABQAEwASABIAEQAQABAADwAPAA4ADQAMAAsACwAMAAoACwAKAAkACgAJAAgABwAFAAUAAwACAAEA///9//3/+v/5//j/9P/y//D/7v/s/+v/6P/m/+X/4v/h/+H/4f/g/+H/4P/g/+H/4v/j/+T/5f/m/+f/6P/q/+v/7f/u/+//8f/y//T/9f/3//j/+v/7//z//v///wEAAgADAAQABQAGAAgACQAKAAsADAAOAA8AEAASABMAFAAVABYAFgAXABgAGAAZABkAGgAaABoAGQAZABgAFgAVABMAEgAQAA8ADgAMAAoACQAHAAUABAACAAAA///9//v/+f/3//X/8//x//D/7//v/+7/7v/t/+3/7f/t/+7/7v/v//D/8P/x//L/8//0//X/9f/2//f/9//4//j/+f/5//r/+v/6//r/+v/6//r/+f/6//r/+v/7//v//P/8//3//v///wEAAgADAAUABgAIAAsADQAPABEAEwAUABYAFwAYABgAGQAZABkAGAAZABkAGQAZABgAFwAVABQAEgARAA4ACwAIAAUAAwABAP///v/8//v/+v/4//b/9f/0//P/8v/y//H/8P/w/+//7v/u/+3/7f/t/+3/7f/s/+z/7P/t/+3/7v/u/+7/7v/u/+//8P/v//D/8P/x//H/8f/y//P/9P/0//X/9f/1//f/+P/5//n/+//8//z//f/9//7//f///wAAAAAAAAEAAQD//wAA/f8AAPz/+//+//v/+v/8//z/+//7//z//f/+/////v///wEAAQAAAAIABAADAAIAAwADAAQAAwADAAMAAwADAAIAAgACAAEAAAAAAAAAAAD//////v/+//7//v/9//7//v/+//////////////////////8AAAAA////////AAAAAAAAAAAAAAEAAAD///7//////////v/9//3//////////v/+//7//v/+//7//v////////8AAAAAAQACAAQABQAGAAcACAAKAAoACgALAAwADAAMAAwADAAMAAwADAAMAAwACwAKAAkACQAJAAgACAAIAAgABwAHAAYABQAEAAMAAgAAAP///v/9//z/+v/5//f/9f/0//P/8v/x//D/7//u/+3/7f/s/+z/7P/s/+7/7//w//L/8//0//X/9v/3//j/+v/8//3///8BAAMABgAIAAoADAAOABAAEQARABIAEwAUABQAFQAVABYAFAAUABYAFAAUABQAFgAVABMAFQAVABUAFAAVABUAFQAVABQAFAATABQAFAAUABQAFAATABQAFAAUABMAEgASABAAEAAQAA4ADgAMAAkACAAGAAIAAAD8//r/9//0//L/7//t/+v/6f/n/+b/4//h/+D/3v/e/9//3v/e/97/3v/e/+H/4v/j/+b/5v/n/+n/6v/s/+7/7//w//L/8//0//b/9//4//n/+v/7//z//f/+////AAABAAIABAAFAAcACAAJAAoACwANAA4ADwARABIAEwAUABUAFgAXABcAFwAXABcAFgAWABUAFAATABIAEAAPAA0ACwAJAAYABAABAP///f/7//n/9//1//T/8v/w/+//7f/s/+r/6f/o/+f/5v/l/+T/5P/j/+T/5P/l/+X/5v/n/+j/6P/p/+r/6//s/+z/7f/u/+7/7v/v/+//7//v/+//8P/w//H/8f/y//P/9P/1//f/+P/5//v//P/+/wEAAwAFAAcACgALAA4AEAATABUAFwAZABsAHQAfACEAIwAkACUAJgAmACUAJQAkACMAIQAgAB0AGwAaABkAGAAWABQAEgAPAA4ADAALAAgABgAEAAIAAQAAAAAA///+//7//f/8//v/+f/5//j/+P/4//f/9//2//X/9f/0//T/9P/0//X/9f/1//X/9f/2//f/9//4//j/+f/6//v//P/8//3//v///wAAAAABAAIAAwAEAAUABQAGAAcACAAIAAcACAAJAAgACAAIAAgABwAIAAkACQAIAAoACgAIAAkABwAKAAcABgAKAAgABwAKAAoACQAJAAoACwALAAwACgAKAAsACQAIAAgACAAGAAQAAwACAAEA///+//3//P/7//r/+f/4//j/9//3//f/9//3//f/9//3//j/+P/4//j/+P/5//n/+f/5//n/+f/5//j/+f/5//n/+f/4//n/+f/6//r/+f/6//v/+//6//r/+//8//3//P/8//7/AAABAAIAAgADAAQABQAGAAcACAAJALine truncated
}
}
]
},
{
"role": "user",
"content": [
{
"type": "audio_url",
"audio_url": {
"url": "data:audio/wav;base64,UklGRkqYAABXQVZFZm10IBAAAAABAAEAwF0AAIC7AAACABAAZGF0YSaYAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAABAAAAAAAAAAAAAAAAAAAAAAAAAP//AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAEAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAD//wAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAP//AAAAAP//AAAAAAAA///o/9T/xv/G/8b/yP/K/8r/y//M/8z/zf/O/83/zf/P/87/zv/O/87/zv/P/8//z//P/8//z//P/8//z//P/8//z//P/8//0P/Q/8//z//Q/9D/0f/R/9D/0f/R/9H/0f/R/9H/0f/R/9H/0f/S/9L/0v/S/9L/0v/S/9L/0v/T/9P/0//T/9P/0//T/9P/0//T/9P/0//T/9P/0//U/9T/1f/V/9T/1P/U/9X/1P/V/9X/1f/V/9b/1v/W/9b/1v/W/9b/1v/W/9f/1//X/9f/1//X/9f/1//X/9j/2P/Y/9j/2P/Y/9j/2P/Y/9j/2f/Z/9r/2v/a/9r/2v/b/9v/2//b/9v/2//c/9v/3P/c/9z/3P/d/9z/3P/c/93/3f/e/93/3f/e/97/3v/e/97/3//f/+D/4P/g/9//3//g/+D/4P/g/+H/4f/g/+H/4P/h/+H/4f/h/+L/4v/i/+H/4f/i/+L/4v/i/+L/4v/j/+L/4//j/+P/4//j/+P/4//j/+P/5P/k/+T/4//l/+T/5P/k/+T/5P/l/+T/5f/l/+X/5f/m/+X/5v/m/+b/5v/m/+b/5v/m/+b/5v/n/+b/5//n/+b/5v/o/+f/5//n/+j/5//p/+f/5//o/+j/6P/n/+j/6P/p/+n/6f/o/+j/6P/p/+j/6P/o/+n/6P/p/+j/6P/p/+n/6f/q/+n/6v/p/+n/6f/q/+n/6f/p/+n/6f/p/+n/6f/p/+n/6v/p/+n/6f/p/+n/6f/p/+n/6f/r/+r/6v/q/+r/6v/r/+n/6v/p/+r/6v/r/+n/6v/q/+r/6v/q/+r/6v/q/+r/6v/q/+r/6//q/+v/6//r/+r/6//r/+r/6v/q/+r/6v/q/+r/6v/q/+r/6//q/+v/6//s/+r/6//q/+r/6v/s/+n/6v/q/+v/6v/r/+v/6//s/+v/6//s/+v/7P/r/+r/6//s/+z/6//r/+v/7P/r/+z/6//r/+v/6//r/+r/6//r/+v/6//r/+v/6//s/+v/6//r/+z/7P/r/+v/6//s/+v/7P/r/+v/6//r/+v/6//s/+v/6//r/+z/7P/s/+z/7P/r/+z/6//s/+v/6//r/+v/6//s/+v/7P/r/+z/7P/s/+z/6//t/+v/7P/s/+v/7P/r/+z/6//s/+z/7P/r/+z/6//t/+v/7P/s/+v/7P/s/+z/6//s/+z/7P/s/+z/7P/s/+v/6//s/+z/7P/r/+z/7P/r/+z/7P/t/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/t/+z/7P/s/+z/7P/s/+z/7P/s/+v/6//s/+z/7P/s/+z/6//r/+z/7P/s/+v/7P/s/+v/7P/s/+v/6//s/+v/7P/s/+z/6//s/+v/7P/q/+v/6//s/+v/6//r/+z/6//r/+z/7P/r/+z/6//s/+z/7P/r/+z/7P/s/+v/6//t/+r/7P/r/+v/6//s/+v/7P/s/+v/7P/r/+z/6//s/+v/6//r/+z/6//r/+v/6//r/+v/6//t/+v/7P/r/+z/6//r/+z/7P/s/+v/6//r/+v/6//s/+v/7P/r/+z/7P/r/+v/6//s/+v/6//r/+z/6//s/+v/6//s/+v/7P/s/+z/7P/r/+v/6//s/+v/7P/r/+z/6//r/+v/6//s/+v/6//r/+z/7P/r/+z/6//s/+v/6//r/+v/6//s/+v/6//s/+r/6//r/+v/6//r/+v/6//r/+v/7P/r/+v/6//s/+v/6//r/+z/6//s/+v/6//q/+v/6//s/+v/6//s/+v/6//s/+v/6//r/+v/6//s/+v/6//r/+z/6//s/+v/6//s/+v/7P/q/+z/6//s/+v/6//s/+v/6//r/+v/6//r/+v/6//r/+v/6//s/+v/6//r/+v/6//s/+v/7P/r/+v/6//r/+v/6//r/+v/6//r/+v/6//s/+v/6//r/+v/6//r/+v/6//s/+r/6//r/+v/6//r/+r/6//r/+v/6//s/+v/6//q/+v/6v/s/+r/6//r/+v/6//r/+v/6//s/+v/6//r/+v/6//r/+v/6//r/+v/6//s/+v/6//r/+r/6//q/+v/6//r/+v/6//r/+v/6//r/+r/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//s/+v/6//r/+v/6//r/+v/6//s/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//s/+v/6//r/+v/6//r/+v/6//s/+v/6//r/+v/6//s/+v/6//r/+z/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//s/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//s/+v/6//r/+v/6//s/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//s/+v/6//s/+v/6//r/+v/6//s/+v/6//r/+v/6//s/+v/6//s/+v/6//r/+z/6//r/+v/7P/r/+v/6//s/+v/6//q/+v/6//s/+r/6//r/+v/6//s/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+v/6//s/+v/6//r/+v/6//s/+v/6//s/+v/7P/r/+z/6//s/+v/6//r/+v/6//r/+v/6//r/+v/6//s/+v/6//r/+v/6//r/+v/6//r/+v/6//r/+z/6//r/+v/6//s/+v/6//r/+v/6//r/+v/6//s/+v/6//s/+v/7P/r/+v/6//r/+v/6//s/+v/6//r/+v/6//r/+v/6v/s/+v/6//r/+v/6//r/+v/6//s/+v/6//r/+v/6//s/+v/6//s/+z/7P/r/+v/7P/s/+v/6//s/+v/6//r/+v/6//r/+z/7P/s/+v/6//r/+z/6//s/+v/7P/r/+z/6//s/+v/6//s/+v/6//r/+z/6//s/+v/6//s/+v/7P/r/+v/6//s/+v/7P/r/+v/6//s/+v/7P/s/+v/7P/r/+z/7P/r/+z/6//s/+v/6//r/+z/6//s/+z/7P/s/+v/6//r/+v/6//r/+z/6//s/+v/6//r/+v/6//r/+v/6//s/+v/6//r/+v/6//s/+z/6//s/+v/6//s/+v/7P/s/+z/6//s/+v/6//s/+v/7P/s/+z/7P/s/+z/6//r/+v/6//s/+v/7P/s/+z/6//s/+v/6//r/+z/6//s/+v/6//r/+z/6//r/+v/6//s/+v/6//s/+v/6//r/+v/6//s/+v/7P/r/+z/7P/r/+v/7P/s/+z/7P/s/+v/7P/s/+z/6//s/+v/7P/s/+v/7P/s/+z/7P/s/+v/6//r/+v/6//s/+v/7P/r/+z/6//r/+z/6//s/+v/7P/r/+v/7P/r/+z/6//s/+v/6//s/+v/7P/r/+z/6//s/+z/7P/s/+v/7P/r/+z/7P/s/+v/6//r/+z/7P/s/+z/6//s/+v/6//s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/7P/s/+z/6//t/+v/7P/s/+z/7P/s/Line truncated
}
}
]
}
]
}
@@ -1,714 +0,0 @@
import base64
import io
import os
import random
import re
from collections.abc import Callable
import librosa
import numpy as np
import torch
import torchaudio
from process_speechdata import InputSegment, StreamingInputSegment
from torchaudio.transforms import MelSpectrogram
speech_zeroemb_idx = 151667
empty_token = "<|empty|>"
mimo_audio_tokenizer = None
device = "cpu"
asr_zh_templates = [
"请将这段语音转换为文字",
"帮我识别这个音频文件中的内容",
"把这段录音转成文本",
"请转录这段语音",
"将音频内容转换成文字格式",
"识别并转写这段语音",
"把语音内容写成文字",
"转录这个音频片段",
"将这段对话转换为文本",
"麻烦帮我把这段录音整理成详细的文字记录",
]
asr_en_templates = [
"Please transcribe this audio file",
"Convert this speech recording to text",
"Transcribe the following voice message",
"Turn this audio into readable text",
"Please convert the recording to written format",
"Transcribe what you hear in this audio",
"Convert this spoken content to text",
"Please write down what is said in this recording",
"Transcribe this voice recording",
"Could you please help me transcribe this important recording?",
"Would you mind converting this voice message into a readable text format?",
"I'd really appreciate it if you could turn this audio file into a written document",
]
tts_zh_templates = [
"请将这段文字转换为语音",
"帮我把这个文本读出来",
"将这些文字生成音频",
"请朗读这段内容",
"把这段话转换成语音文件",
"生成这段文字的语音版本",
"请用语音播报这些内容",
"将文本转换为可听的音频",
"帮我朗读这段文字",
"把这些内容念出来",
]
tts_en_templates = [
"Please convert this text to speech",
"Turn this writing into audio",
"Generate speech from this text",
"Read this content out loud",
"Convert these words to voice",
"Create an audio version of this text",
"Please vocalize this content",
"Turn this text into audible format",
"Help me convert this writing to speech",
"Make this text into spoken audio",
]
def detect_language(text):
if re.search(r"[\u4e00-\u9fff]", text):
return "zh"
else:
return "en"
# ============================================
# Common Helper Functions - InputSegment Creation
# ============================================
def create_segment(text: str = "", audio=None) -> InputSegment:
"""Create a standard InputSegment with default zeroemb parameters"""
return InputSegment(
text=text,
audio=audio,
speech_zeroemb_idx=speech_zeroemb_idx,
text_zeroemb_idx=empty_token,
)
def create_streaming_segment(
text: str, audio, tokenizer, group_size: int, audio_channels: int
) -> StreamingInputSegment:
"""Create a StreamingInputSegment"""
return StreamingInputSegment(
text=text,
audio=audio,
tokenizer=tokenizer,
group_size=group_size,
audio_channels=audio_channels,
speech_zeroemb_idx=speech_zeroemb_idx,
text_zeroemb_idx=empty_token,
)
# ============================================
# Common Helper Functions - Common Token Segment Creation
# ============================================
def create_user_start() -> InputSegment:
"""Create user role start token"""
return create_segment(text="<|im_start|>user\n")
def create_user_end() -> InputSegment:
"""Create role end token"""
return create_segment(text="<|im_end|>\n")
def create_assistant_start() -> InputSegment:
"""Create assistant role start token"""
return create_segment(text="<|im_start|>assistant\n")
def create_system_start() -> InputSegment:
"""Create system role start token"""
return create_segment(text="<|im_start|>system\n")
def create_thinking_segment(thinking: bool = False) -> InputSegment:
"""Create thinking token, closed or open based on thinking parameter"""
if thinking:
return create_segment(text="<think>\n")
else:
return create_segment(text="<think>\n\n</think>\n")
def create_sostm_segment() -> InputSegment:
"""Create streaming output start token"""
return create_segment(text="<|sostm|>")
def create_assistant_start_with_sostm() -> InputSegment:
"""Create assistant start token with sostm"""
return create_segment(text="<|im_start|>assistant\n<|sostm|>")
def create_assistant_start_with_think() -> InputSegment:
"""Create assistant start token with think opening"""
return create_segment(text="<|im_start|>assistant\n<think>\n")
# ============================================
# Common Helper Functions - Composite Segment Creation
# ============================================
def create_user_turn_with_audio(audio_tokenized, extra_text: str = None) -> list[InputSegment]:
"""Create a user turn containing audio"""
segments = [
create_user_start(),
create_segment(audio=audio_tokenized),
]
if extra_text:
segments.append(create_segment(text=extra_text))
segments.append(create_user_end())
return segments
def create_user_turn_with_text(text: str) -> list[InputSegment]:
"""Create a text-only user turn"""
return [
create_user_start(),
create_segment(text=text),
create_user_end(),
]
def create_system_turn_with_voice_prompt(prompt_text: str, audio_token) -> list[InputSegment]:
"""Create a system turn with voice prompt"""
return [
create_system_start(),
create_segment(text=prompt_text),
create_segment(text="", audio=audio_token),
create_user_end(),
]
def create_system_turn_text_only(system_text: str) -> list[InputSegment]:
"""Create a text-only system turn"""
return [
create_system_start(),
create_segment(text=system_text),
create_user_end(),
]
# ============================================
# Common Helper Functions - Multi-turn Dialogue Processing
# ============================================
def process_multiturn_messages(
message_list: list[dict],
user_processor: Callable[[dict], list[InputSegment]],
assistant_processor: Callable[[dict], list[InputSegment]],
) -> list[InputSegment]:
"""
Generic multi-turn dialogue message processing function
Args:
message_list: List of messages, each containing 'role' and 'content'
user_processor: Function to process user messages
assistant_processor: Function to process assistant messages
Returns:
Processed list of InputSegments
"""
lm_prompt = []
for message in message_list:
role = message["role"]
if role == "user":
lm_prompt.extend(user_processor(message))
elif role == "assistant":
lm_prompt.extend(assistant_processor(message))
else:
raise ValueError(f"Invalid role: {role}")
return lm_prompt
def create_text_user_message(message: dict) -> list[InputSegment]:
"""Process a text-only user message"""
return [
create_user_start(),
create_segment(text=message["content"]),
create_user_end(),
]
def create_text_assistant_message(message: dict) -> list[InputSegment]:
"""Process a text-only assistant message"""
return [
create_assistant_start(),
create_segment(text=message["content"]),
create_user_end(),
]
def create_audio_user_message(message: dict) -> list[InputSegment]:
"""Process an audio user message"""
return [
create_user_start(),
create_segment(audio=preprocess_input(message["content"])),
create_user_end(),
]
def append_assistant_ending(
lm_prompt: list[InputSegment], thinking: bool = False, use_sostm: bool = False
) -> list[InputSegment]:
"""
Append assistant ending to the prompt
Args:
lm_prompt: Existing prompt list
thinking: Whether to use open thinking token
use_sostm: Whether to use sostm token (for speech output)
"""
if use_sostm:
lm_prompt.append(create_assistant_start_with_sostm())
else:
lm_prompt.append(create_assistant_start())
lm_prompt.append(create_thinking_segment(thinking))
return lm_prompt
def get_asr_sft_prompt(
input: None | str = None,
):
"""Build prompt for ASR (Automatic Speech Recognition) task"""
audio_tokenized = preprocess_input(input)
template = random.choice(asr_zh_templates + asr_en_templates)
lm_prompt = create_user_turn_with_audio(audio_tokenized, extra_text=template)
lm_prompt = append_assistant_ending(lm_prompt, thinking=False)
return lm_prompt
def resample_audio_if_needed(wav_tensor: torch.Tensor, original_sr: int):
target_sr = 24000
if original_sr != target_sr:
wav_tensor = torchaudio.functional.resample(wav_tensor, original_sr, target_sr)
return wav_tensor
def wav2mel(wav, device="cpu"):
mel_transform = MelSpectrogram(
sample_rate=mimo_audio_tokenizer.config.sampling_rate,
n_fft=mimo_audio_tokenizer.config.nfft,
hop_length=mimo_audio_tokenizer.config.hop_length,
win_length=mimo_audio_tokenizer.config.window_size,
f_min=mimo_audio_tokenizer.config.fmin,
f_max=mimo_audio_tokenizer.config.fmax,
n_mels=mimo_audio_tokenizer.config.n_mels,
power=1.0,
center=True,
).to(device)
spec = mel_transform(wav[None, :])
return torch.log(torch.clip(spec, min=1e-7)).squeeze()
def group_by_length(features: torch.Tensor, lengths: torch.Tensor, max_length: int):
if features.size(0) != lengths.sum().item():
raise ValueError(f"Feature size mismatch: {features.size(0)} vs {lengths.sum().item()}")
split_points = []
current_sum = 0
for i, seq_len in enumerate(lengths):
if current_sum + seq_len > max_length and current_sum > 0:
split_points.append(i)
current_sum = seq_len.item()
else:
current_sum += seq_len.item()
# Convert split points to group sizes
group_sizes = []
prev = 0
for point in split_points:
group_sizes.append(point - prev)
prev = point
if prev < len(lengths):
group_sizes.append(len(lengths) - prev)
len_groups = torch.split(lengths, group_sizes)
feature_sizes = [group.sum().item() for group in len_groups]
feature_groups = torch.split(features, feature_sizes)
return feature_groups, len_groups
def encode_batch(input_features: torch.Tensor, input_lens: torch.Tensor, max_length: int = 256000):
feature_groups, len_groups = group_by_length(input_features, input_lens, max_length)
encoded_parts = []
for features, lengths in zip(feature_groups, len_groups):
with torch.no_grad():
codes, _ = mimo_audio_tokenizer.encoder.encode(
input_features=features.to(device), input_lens=lengths.to(device), return_codes_only=True
)
encoded_parts.append(codes)
return torch.cat(encoded_parts, dim=-1)
def preprocess_input(input: None | str | torch.Tensor = None, device="cpu", audio_channels=4, group_size=8):
if isinstance(input, torch.Tensor) or (isinstance(input, str) and os.path.isfile(input)):
return "<|sosp|><|empty|><|eosp|>"
else:
text = input
if (
text.isupper() or text.islower()
): # If the text only contains upper-case or lower-case letters, capitalize it.
text = text.capitalize()
return text
def _build_tts_system_prompt(has_voice_prompt: bool, voice_audio_token=None) -> list[InputSegment]:
"""Build system prompt for TTS task"""
if has_voice_prompt and voice_audio_token is not None:
return [
create_system_start(),
create_segment(
# text="You need to generate a speech with the same timbre as the speech prompt, based on the specified style instructions and text content. Your timbre should be: "
text="你需要根据指定的风格指令和文本内容来生成和语音prompt具有相同音色的语音。你的音色应该是:"
),
create_segment(text="", audio=voice_audio_token),
create_user_end(),
]
else:
return create_system_turn_text_only("你需要根据指定的风格指令和文本内容来生成语音。")
# return create_system_turn_text_only(
# "You need to generate speech based on the specified style instructions and text content."
# )
def _build_tts_system_prompt_no_instruct(has_voice_prompt: bool, voice_audio_token=None) -> list[InputSegment]:
"""Build system prompt for TTS task"""
if has_voice_prompt and voice_audio_token is not None:
return [
create_system_start(),
create_segment(
# text="You need to generate a speech with the same timbre as the speech prompt, based on the specified style instructions and text content. Your timbre should be:"
text="你需要根据指定的风格指令和文本内容来生成和语音prompt具有相同音色的语音。你的音色应该是:"
),
create_segment(text="", audio=voice_audio_token),
create_user_end(),
]
else:
return []
def get_tts_sft_prompt(
input: None | str = None,
instruct=None,
read_text_only=True,
prompt_speech=None,
):
"""
Build prompt for TTS (Text-to-Speech) task
Args:
input: Input text
instruct: Style instruction (e.g., "speak happily in a child's voice")
read_text_only: Whether to read only plain text (False means text contains template)
prompt_speech: Reference audio (for voice cloning)
"""
assistant_prompt_audio_token = preprocess_input(prompt_speech) if prompt_speech is not None else None
if not read_text_only:
# Not just reading text, text contains template (template:text format)
text = preprocess_input(input)
lm_prompt = _build_tts_system_prompt(
has_voice_prompt=assistant_prompt_audio_token is not None,
voice_audio_token=assistant_prompt_audio_token,
)
lm_prompt.append(create_segment(text=f"<|im_start|>user\n{text}<|im_end|>\n"))
lm_prompt.append(create_assistant_start_with_think())
else:
# Plain text (no instruction inside)
language = detect_language(input)
template = random.choice(tts_zh_templates if language == "zh" else tts_en_templates)
text = preprocess_input(input)
if instruct is None:
# No instruct instruction
lm_prompt = _build_tts_system_prompt_no_instruct(
has_voice_prompt=assistant_prompt_audio_token is not None,
voice_audio_token=assistant_prompt_audio_token,
)
lm_prompt.extend(
[
create_segment(text=f"<|im_start|>user\n{template}: {text}<|im_end|>\n"),
create_assistant_start_with_sostm(),
]
)
else:
# Has instruct instruction
lm_prompt = _build_tts_system_prompt(
has_voice_prompt=assistant_prompt_audio_token is not None,
voice_audio_token=assistant_prompt_audio_token,
)
lm_prompt.append(create_segment(text=f"<|im_start|>user\n{template}: {text}({instruct})<|im_end|>\n"))
lm_prompt.append(create_assistant_start_with_think())
return lm_prompt
def get_audio_understanding_sft_prompt(
input_speech,
input_text,
thinking=False,
use_sostm=False,
):
"""Build prompt for audio understanding task"""
audio_tokenized = preprocess_input(input_speech)
lm_prompt = create_user_turn_with_audio(audio_tokenized, extra_text=input_text)
lm_prompt = append_assistant_ending(lm_prompt, thinking=thinking, use_sostm=use_sostm)
return lm_prompt
def _build_voice_prompt_system(prompt_speech) -> list[InputSegment]:
"""Build system prompt with voice prompt"""
return create_system_turn_with_voice_prompt(
prompt_text="Your voice should be:", audio_token=preprocess_input(prompt_speech)
)
def get_spoken_dialogue_sft_prompt(
input_speech,
system_prompt=None,
prompt_speech=None,
add_history=False,
):
"""
Build prompt for spoken dialogue task
Args:
input_speech: Input speech
system_prompt: System prompt text
prompt_speech: Reference audio (for voice cloning)
add_history: Whether to add history (Note: history variable is undefined in original code)
"""
audio_tokenized = preprocess_input(input_speech)
lm_prompt = []
# Note: history variable is undefined in original code, this branch may never execute
# To use history feature, history should be passed as a parameter
if add_history:
# Simplified form of adding history
lm_prompt = create_user_turn_with_audio(audio_tokenized)
lm_prompt.append(create_assistant_start_with_sostm())
else:
# Add voice prompt (if available)
if prompt_speech:
lm_prompt.extend(_build_voice_prompt_system(prompt_speech))
# Add user turn
lm_prompt.append(create_user_start())
if system_prompt:
lm_prompt.append(create_segment(text=system_prompt))
lm_prompt.append(create_segment(audio=audio_tokenized))
lm_prompt.append(create_user_end())
lm_prompt.append(create_assistant_start_with_sostm())
return lm_prompt
def get_spoken_dialogue_sft_multiturn_prompt(
message_list,
system_prompt=None,
prompt_speech=None,
tokenizer=None,
group_size=8,
audio_channels=4,
):
"""
Build prompt for multi-turn spoken dialogue task
Args:
message_list: List of messages containing role and content
system_prompt: System prompt text
prompt_speech: Reference audio (for voice cloning)
tokenizer: Tokenizer
group_size: Group size
audio_channels: Number of audio channels
"""
lm_prompt = []
# Add voice prompt (if available)
if prompt_speech:
lm_prompt.extend(
create_system_turn_with_voice_prompt(
prompt_text="Your voice should be:", audio_token=preprocess_input(prompt_speech)
)
)
# Define message processors
def user_processor(msg):
segments = [create_user_start()]
if system_prompt:
segments.append(create_segment(text=system_prompt))
segments.append(create_segment(audio=preprocess_input(msg["content"])))
segments.append(create_user_end())
return segments
def assistant_processor(msg):
return [
create_assistant_start(),
create_streaming_segment(
text=msg["content"]["text"],
audio=preprocess_input(msg["content"]["audio"]),
tokenizer=tokenizer,
group_size=group_size,
audio_channels=audio_channels,
),
create_user_end(),
]
# Process message list
lm_prompt.extend(process_multiturn_messages(message_list, user_processor, assistant_processor))
lm_prompt.append(create_assistant_start_with_sostm())
return lm_prompt
def get_s2t_dialogue_sft_prompt(
input_speech,
thinking=False,
):
"""Build prompt for speech-to-text dialogue task"""
audio_tokenized = preprocess_input(input_speech)
lm_prompt = create_user_turn_with_audio(audio_tokenized)
lm_prompt = append_assistant_ending(lm_prompt, thinking=thinking)
return lm_prompt
def get_s2t_dialogue_sft_multiturn_prompt(message_list, thinking=False):
"""Build prompt for multi-turn speech-to-text dialogue task"""
lm_prompt = process_multiturn_messages(
message_list, user_processor=create_audio_user_message, assistant_processor=create_text_assistant_message
)
lm_prompt = append_assistant_ending(lm_prompt, thinking=thinking)
return lm_prompt
def get_text_dialogue_sft_prompt(
input_text,
thinking=False,
):
"""Build prompt for text-only dialogue task"""
lm_prompt = create_user_turn_with_text(input_text)
lm_prompt = append_assistant_ending(lm_prompt, thinking=thinking)
return lm_prompt
def get_text_dialogue_sft_multiturn_prompt(
message_list,
thinking=False,
):
"""Build prompt for multi-turn text-only dialogue task"""
lm_prompt = process_multiturn_messages(
message_list, user_processor=create_text_user_message, assistant_processor=create_text_assistant_message
)
lm_prompt = append_assistant_ending(lm_prompt, thinking=thinking)
return lm_prompt
def get_in_context_learning_s2s_prompt(
instruction,
prompt_examples,
audio,
tokenizer=None,
group_size=8,
audio_channels=4,
):
"""
Build prompt for In-Context Learning speech-to-speech task
Args:
instruction: Instruction text
prompt_examples: List of examples, each containing input_audio, output_transcription, output_audio
audio: Input audio to be processed
tokenizer: Tokenizer
group_size: Group size
audio_channels: Number of audio channels
"""
prompt = [create_segment(text=f"[Int]:{instruction}\n")]
# Add examples
for example in prompt_examples:
prompt.extend(
[
create_segment(audio=preprocess_input(example["input_audio"])),
create_segment(text="\n"),
create_streaming_segment(
text=example["output_transcription"],
audio=preprocess_input(example["output_audio"]),
tokenizer=tokenizer,
group_size=group_size,
audio_channels=audio_channels,
),
create_segment(text=" \n\n"),
]
)
# Add input audio to be processed
prompt.extend(
[
create_segment(audio=preprocess_input(audio)),
create_segment(text="\n"),
create_sostm_segment(),
]
)
return prompt
def get_audio_data(audio_url):
if audio_url.startswith("data:"):
header, b64_data = audio_url.split(",", 1)
audio_bytes = base64.b64decode(b64_data.strip())
audio_file = io.BytesIO(audio_bytes)
else:
# File path
audio_file = audio_url
audio_signal, sr = librosa.load(audio_file, sr=24000)
audio_data = (audio_signal.astype(np.float32), sr)
return audio_data
def to_prompt(input_segs):
out_put = []
for input_seg in input_segs:
if isinstance(input_seg, StreamingInputSegment) and input_seg.text:
out_put.append("<|sostm|>")
if input_seg.audio is not None and isinstance(input_seg.audio, str):
out_put.append(input_seg.text)
out_put.append("<|eot|>")
out_put.append("<|empty|>")
else:
out_put.append(input_seg.text)
out_put.append("<|eot|>")
out_put.append("<|eostm|>")
else:
out_put.append(input_seg.text)
if input_seg.audio is not None:
out_put.append(input_seg.audio)
prompt = "".join(out_put)
return prompt
@@ -1,261 +0,0 @@
#!/usr/bin/env python3
# Copyright 2025 Xiaomi Corporation.
# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
import torch
class InputSegment:
def __init__(
self,
text: str = "",
audio: torch.Tensor = None,
tokenized_text: torch.Tensor = None,
speech_zeroemb_idx: int | list[int] = 1024,
text_zeroemb_idx: int = 152067,
add_sosp_eosp=True,
) -> None:
has_text = text is not None
has_tokenized_text = tokenized_text is not None
assert has_text or has_tokenized_text, "Text or tokenized text must be provided"
self.audio = audio
self.text = text
self.tokenized_text = tokenized_text
self.speech_zeroemb_idx = speech_zeroemb_idx
self.text_zeroemb_idx = text_zeroemb_idx
self.add_sosp_eosp = add_sosp_eosp
@staticmethod
def insert_between(tensor, i, value=-1):
return torch.scatter(
torch.full(
(1, tensor.shape[1] + (tensor.shape[1] - 1) * i + i),
value,
dtype=tensor.dtype,
),
1,
torch.arange(0, tensor.shape[1], dtype=torch.int64)[None] * (i + 1),
tensor,
)
def to_input_id(
self,
tokenizer,
group_size: int,
audio_channels: int = 8,
) -> tuple[torch.Tensor, torch.Tensor]:
if (
self.audio is None
): ## If no audio is provided, tokenize the text with (group_size-1) tokens inserted between each, fill with -100; and set audio to empty indices
if self.tokenized_text is None:
tokenized_text = tokenizer(
self.text,
return_tensors="pt",
truncation=True,
max_length=999999,
padding=False,
add_special_tokens=False,
)["input_ids"].int()
else:
tokenized_text = self.tokenized_text.unsqueeze(0)
if group_size > 1:
tokenized_text = self.insert_between(tokenized_text, group_size - 1, value=-100)
if isinstance(self.speech_zeroemb_idx, list):
audio_part_input_id = torch.zeros((audio_channels, tokenized_text.shape[1]), dtype=torch.int)
for i, idx in enumerate(self.speech_zeroemb_idx):
audio_part_input_id[i, :] = idx
else:
audio_part_input_id = torch.full(
(audio_channels, tokenized_text.shape[1]), self.speech_zeroemb_idx, dtype=torch.int
)
else: # If audio is provided, add sosp/eosp markers and convert audio: fill with empty tokens based on audio length
sosp_token = tokenizer.convert_tokens_to_ids("<|sosp|>") if self.add_sosp_eosp else None
eosp_token = tokenizer.convert_tokens_to_ids("<|eosp|>") if self.add_sosp_eosp else None
audio_part = self.audio.reshape(-1, audio_channels).T # [audio_channels, seqlen]
assert audio_part.shape[1] % group_size == 0, (
f"Audio shape {audio_part.shape} is not divisible by group_size {group_size}"
)
text_len = audio_part.shape[1] // group_size
empty_token = self.text_zeroemb_idx
if empty_token is None:
empty_token = tokenizer.eod
tokenized_text = torch.full((1, text_len), empty_token, dtype=torch.int)
tokenized_text = (
torch.cat(
[
torch.tensor([[sosp_token]], dtype=torch.int),
tokenized_text,
torch.tensor([[eosp_token]], dtype=torch.int),
],
dim=1,
)
if self.add_sosp_eosp
else tokenized_text
)
tokenized_text = self.insert_between(tokenized_text, group_size - 1, value=-100)
if self.add_sosp_eosp:
if isinstance(self.speech_zeroemb_idx, list):
sosp_part = torch.zeros((audio_channels, group_size), dtype=torch.int)
eosp_part = torch.zeros((audio_channels, group_size), dtype=torch.int)
for i, idx in enumerate(self.speech_zeroemb_idx):
sosp_part[i, :] = idx
eosp_part[i, :] = idx
audio_part_input_id = torch.cat([sosp_part, audio_part, eosp_part], dim=1)
else:
audio_part_input_id = torch.cat(
[
torch.full((audio_channels, group_size), self.speech_zeroemb_idx, dtype=torch.int),
audio_part,
torch.full((audio_channels, group_size), self.speech_zeroemb_idx, dtype=torch.int),
],
dim=1,
)
else:
audio_part_input_id = audio_part
input_ids = torch.cat([tokenized_text, audio_part_input_id], dim=0) # [n_rvq + 1, seqlen]
return input_ids
class StreamingInputSegment:
def __init__(
self,
text: str = "",
audio: torch.Tensor = None,
tokenized_text: torch.Tensor = None,
speech_zeroemb_idx: int | list[int] = 1024,
text_zeroemb_idx: int = 152067,
text_segment_size: int = 5,
audio_segment_size: int = 5,
tokenizer=None,
group_size=None,
audio_channels=None,
) -> None:
has_text = text is not None
has_tokenized_text = tokenized_text is not None
assert has_text or has_tokenized_text, "Text or tokenized text must be provided"
self.audio = audio
self.text = text
self.tokenized_text = tokenized_text
self.speech_zeroemb_idx = speech_zeroemb_idx
self.text_zeroemb_idx = text_zeroemb_idx
self.text_segment_size = text_segment_size
self.audio_segment_size = audio_segment_size
self.tokenizer = tokenizer
self.group_size = group_size
self.audio_channels = audio_channels
def to_input_id(
self,
tokenizer,
group_size: int,
audio_channels: int = 8,
):
if self.tokenized_text is None:
tokenized_text = tokenizer(
self.text,
return_tensors="pt",
truncation=True,
max_length=999999,
padding=False,
add_special_tokens=False,
)["input_ids"].int() # [1, seqlen]
else:
tokenized_text = self.tokenized_text.unsqueeze(0)
tokenized_text = tokenized_text.squeeze(0)
text_segments = tokenized_text.split(self.text_segment_size, dim=0)
audio_segments = self.audio.split(self.audio_segment_size * group_size * audio_channels, dim=0)
tokenized_segments = []
tokenized_segments.append(
InputSegment(
text="<|sostm|>",
speech_zeroemb_idx=self.speech_zeroemb_idx,
text_zeroemb_idx=self.text_zeroemb_idx,
),
)
eot_tokens = tokenizer(
"<|eot|>",
return_tensors="pt",
truncation=True,
max_length=999999,
padding=False,
add_special_tokens=False,
)["input_ids"][0].to(text_segments[-1])
text_segments = text_segments[:-1] + (torch.cat([text_segments[-1], eot_tokens], dim=0),)
length = min(len(text_segments), len(audio_segments)) # The number of empty tokens here may vary
for i in range(length):
text_segment = text_segments[i]
audio_segment = audio_segments[i]
tokenized_segments.append(
InputSegment(
tokenized_text=text_segment,
speech_zeroemb_idx=self.speech_zeroemb_idx,
text_zeroemb_idx=self.text_zeroemb_idx,
),
)
tokenized_segments.append(
InputSegment(
audio=audio_segment,
add_sosp_eosp=False,
speech_zeroemb_idx=self.speech_zeroemb_idx,
text_zeroemb_idx=self.text_zeroemb_idx,
),
)
for j in range(length, len(text_segments)):
tokenized_segments.append(
InputSegment(
tokenized_text=text_segments[j],
speech_zeroemb_idx=self.speech_zeroemb_idx,
text_zeroemb_idx=self.text_zeroemb_idx,
),
)
for j in range(length, len(audio_segments)):
tokenized_segments.append(
InputSegment(
audio=audio_segments[j],
add_sosp_eosp=False,
speech_zeroemb_idx=self.speech_zeroemb_idx,
text_zeroemb_idx=self.text_zeroemb_idx,
),
)
tokenized_segments.append(
InputSegment(
text="<|eostm|>",
speech_zeroemb_idx=self.speech_zeroemb_idx,
text_zeroemb_idx=self.text_zeroemb_idx,
),
)
input_ids = [
seg.to_input_id(
self.tokenizer,
self.group_size,
self.audio_channels,
)
for seg in tokenized_segments
]
input_ids = torch.cat(input_ids, dim=1).type(torch.int64) # [n_rvq + 1, seqlen]
return input_ids
+1 -27
View File
@@ -87,33 +87,7 @@ Examples:
python end2end.py --query-type Base --mode-tag icl
```
## Streaming Mode
Add `--streaming` to stream audio chunks progressively via `AsyncOmni` (requires `async_chunk: true` in the stage config):
```bash
python end2end.py --query-type CustomVoice --streaming --output-dir /tmp/out_stream
```
Each 25-frame Code2Wav chunk is logged as it arrives. The final WAV file is written once generation
completes. This demonstrates that audio data is available progressively rather than only at the end.
> **Note:** Streaming uses `AsyncOmni` internally. The non-streaming path (`Omni`) is unchanged.
## Batched Decoding
The Code2Wav stage (stage 1) supports batched decoding, where multiple requests are decoded in a single forward pass through the SpeechTokenizer. To use it, provide a stage config with `max_batch_size > 1` and pass multiple prompts via `--txt-prompts` with a matching `--batch-size`.
```
python end2end.py --query-type CustomVoice \
--txt-prompts benchmark_prompts.txt \
--batch-size 4 \
--stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts_batch.yaml
```
**Important:** `--batch-size` must match a CUDA graph capture size (1, 2, 4, 8, 16...) because the Talker's code predictor KV cache is sized to `max_num_seqs`, and CUDA graphs pad the batch to the next capture size. Both stages need `max_batch_size >= batch_size` in the stage config for batching to take effect. If only stage 1 has a higher `max_batch_size`, it won't help — stage 1 can only batch chunks from requests that are in-flight simultaneously, which requires stage 0 to also process multiple requests concurrently.
## Notes
- The script uses the model paths embedded in `end2end.py`. Update them if your local cache path differs.
- Use `--output-dir` to change the output folder.
- Use `--output-dir` (preferred) or `--output-wav` to change the output folder.
@@ -1,12 +0,0 @@
Hello, welcome to the voice synthesis benchmark test.
She said she would be here by noon, but nobody showed up.
The quick brown fox jumps over the lazy dog near the riverbank.
I can't believe how beautiful the sunset looks from up here on the mountain.
Please remember to bring your identification documents to the appointment tomorrow morning.
Have you ever wondered what it would be like to travel through time and visit ancient civilizations?
The restaurant on the corner serves the best pasta I have ever tasted in my entire life.
After the meeting, we should discuss the quarterly results and plan for the next phase.
Learning a new language takes patience, practice, and a genuine curiosity about other cultures.
The train leaves at half past seven, so we need to arrive at the station before then.
Could you please turn down the music a little bit, I'm trying to concentrate on my work.
It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.
Loaded 100 of 305 files, more files were not shown because too many files have changed in this diff. Show more