Compare commits
8
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
18252b9d8d | ||
|
|
a326cd0d57 | ||
|
|
c15f51a2f6 | ||
|
|
7c934b6188 | ||
|
|
74050cdacd | ||
|
|
f0cec3ecd4 | ||
|
|
7b885f9d13 | ||
|
|
69bcc6416d |
No files matched your search
@@ -10,7 +10,7 @@ steps:
|
|||||||
DOCKER_BUILDKIT: "1"
|
DOCKER_BUILDKIT: "1"
|
||||||
# Buildkite will automatically replace this with the actual commit hash
|
# Buildkite will automatically replace this with the actual commit hash
|
||||||
VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
|
VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
|
||||||
VLLM_VERSION: "v0.22.0"
|
VLLM_VERSION: "v0.21.0"
|
||||||
priority: 100
|
priority: 100
|
||||||
timeout_in_minutes: 60
|
timeout_in_minutes: 60
|
||||||
soft_fail: false
|
soft_fail: false
|
||||||
@@ -1,28 +1,15 @@
|
|||||||
steps:
|
steps:
|
||||||
|
|
||||||
- label: "Simple · Diffusion & Model Executor Test"
|
- label: "Simple Unit Test"
|
||||||
agent_pool: mi325_1
|
agent_pool: mi325_1
|
||||||
depends_on: amd-build
|
depends_on: amd-build
|
||||||
timeout_in_minutes: 50
|
|
||||||
mirror_hardwares: [amdproduction]
|
mirror_hardwares: [amdproduction]
|
||||||
grade: Blocking
|
grade: Blocking
|
||||||
commands:
|
commands:
|
||||||
- export VLLM_ROCM_USE_AITER=0
|
- export VLLM_ROCM_USE_AITER=0
|
||||||
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
|
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
|
||||||
# that is not supported on CPU
|
# that is not supported on CPU
|
||||||
- "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
- "timeout 60m pytest -v -s -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||||
|
|
||||||
- label: "Simple · Other Test"
|
|
||||||
agent_pool: mi325_1
|
|
||||||
depends_on: amd-build
|
|
||||||
timeout_in_minutes: 50
|
|
||||||
mirror_hardwares: [amdproduction]
|
|
||||||
grade: Blocking
|
|
||||||
commands:
|
|
||||||
- export VLLM_ROCM_USE_AITER=0
|
|
||||||
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
|
|
||||||
# that is not supported on CPU
|
|
||||||
- "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
|
||||||
|
|
||||||
- label: "Diffusion Images API LoRA E2E"
|
- label: "Diffusion Images API LoRA E2E"
|
||||||
agent_pool: mi325_1
|
agent_pool: mi325_1
|
||||||
@@ -62,6 +49,16 @@ steps:
|
|||||||
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||||
# - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_expansion.py -m "advanced_model and diffusion and L4" --run-level advanced_model
|
# - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_expansion.py -m "advanced_model and diffusion and L4" --run-level advanced_model
|
||||||
|
|
||||||
|
- label: "Diffusion Cache Backend Test"
|
||||||
|
agent_pool: mi325_1
|
||||||
|
depends_on: amd-build
|
||||||
|
mirror_hardwares: [amdproduction]
|
||||||
|
grade: Blocking
|
||||||
|
commands:
|
||||||
|
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||||
|
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||||
|
- timeout 15m pytest -s -v -m "core_model and cache and diffusion and not distributed_cuda and L4"
|
||||||
|
|
||||||
- label: "Diffusion Sequence Parallelism Test (Need 4 GPUs)"
|
- label: "Diffusion Sequence Parallelism Test (Need 4 GPUs)"
|
||||||
agent_pool: mi325_4
|
agent_pool: mi325_4
|
||||||
depends_on: amd-build
|
depends_on: amd-build
|
||||||
@@ -126,7 +123,7 @@ steps:
|
|||||||
grade: Blocking
|
grade: Blocking
|
||||||
commands:
|
commands:
|
||||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||||
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py tests/e2e/online_serving/test_mimo_audio.py -m "advanced_model" --run-level "advanced_model"
|
||||||
|
|
||||||
- label: "Qwen3-TTS CustomVoice E2E Test"
|
- label: "Qwen3-TTS CustomVoice E2E Test"
|
||||||
agent_pool: mi325_1
|
agent_pool: mi325_1
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
steps:
|
steps:
|
||||||
|
|
||||||
- label: "Simple · Diffusion & Model Executor Test"
|
- label: "Simple Unit Test"
|
||||||
agent_pool: mi325_1
|
agent_pool: mi325_1
|
||||||
depends_on: amd-build
|
depends_on: amd-build
|
||||||
mirror_hardwares: [amdproduction]
|
mirror_hardwares: [amdproduction]
|
||||||
@@ -9,18 +9,7 @@ steps:
|
|||||||
- export VLLM_ROCM_USE_AITER=0
|
- export VLLM_ROCM_USE_AITER=0
|
||||||
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
|
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
|
||||||
# that is not supported on CPU
|
# that is not supported on CPU
|
||||||
- "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
- "timeout 60m pytest -v -s -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||||
|
|
||||||
- label: "Simple · Other Test"
|
|
||||||
agent_pool: mi325_1
|
|
||||||
depends_on: amd-build
|
|
||||||
mirror_hardwares: [amdproduction]
|
|
||||||
grade: Blocking
|
|
||||||
commands:
|
|
||||||
- export VLLM_ROCM_USE_AITER=0
|
|
||||||
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
|
|
||||||
# that is not supported on CPU
|
|
||||||
- "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
|
||||||
|
|
||||||
- label: "Voxtral TTS CUDA Unit Test"
|
- label: "Voxtral TTS CUDA Unit Test"
|
||||||
agent_pool: mi325_1
|
agent_pool: mi325_1
|
||||||
@@ -30,13 +19,13 @@ steps:
|
|||||||
commands:
|
commands:
|
||||||
- timeout 10m pytest -s -v tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py
|
- timeout 10m pytest -s -v tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py
|
||||||
|
|
||||||
- label: "Diffusion · Batch Test"
|
- label: "Diffusion Batching Test"
|
||||||
agent_pool: mi325_1
|
agent_pool: mi325_1
|
||||||
depends_on: amd-build
|
depends_on: amd-build
|
||||||
mirror_hardwares: [amdproduction]
|
mirror_hardwares: [amdproduction]
|
||||||
grade: Blocking
|
grade: Blocking
|
||||||
commands:
|
commands:
|
||||||
- timeout 20m pytest -sv tests/diffusion/batching -m 'core_model and cuda' --run-level "core_model"
|
- timeout 20m pytest -s -v tests/diffusion/batching/test_diffusion_batching.py -m "core_model and diffusion" --run-level "core_model"
|
||||||
|
|
||||||
- label: "Custom Pipeline Test"
|
- label: "Custom Pipeline Test"
|
||||||
agent_pool: mi325_1
|
agent_pool: mi325_1
|
||||||
|
|||||||
@@ -107,7 +107,7 @@ steps:
|
|||||||
depends_on: upload-merge-pipeline
|
depends_on: upload-merge-pipeline
|
||||||
steps:
|
steps:
|
||||||
- label: "Entrypoint Test with H100"
|
- label: "Entrypoint Test with H100"
|
||||||
timeout_in_minutes: 60
|
timeout_in_minutes: 35
|
||||||
commands:
|
commands:
|
||||||
- pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and H100" --run-level "advanced_model"
|
- pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and H100" --run-level "advanced_model"
|
||||||
agents:
|
agents:
|
||||||
@@ -389,29 +389,6 @@ steps:
|
|||||||
volumes:
|
volumes:
|
||||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||||
|
|
||||||
- label: "TTS · VoxCPM2 Online Test"
|
|
||||||
timeout_in_minutes: 30
|
|
||||||
commands:
|
|
||||||
- |
|
|
||||||
timeout 30m bash -c "
|
|
||||||
export VLLM_LOGGING_LEVEL=DEBUG
|
|
||||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
|
||||||
pytest -s -v tests/e2e/online_serving/test_voxcpm2_tts.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
|
||||||
"
|
|
||||||
agents:
|
|
||||||
queue: "gpu_1_queue"
|
|
||||||
plugins:
|
|
||||||
- docker#v5.2.0:
|
|
||||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
|
||||||
always-pull: true
|
|
||||||
propagate-environment: true
|
|
||||||
shm-size: "8gb"
|
|
||||||
environment:
|
|
||||||
- "HF_HOME=/fsx/hf_cache"
|
|
||||||
- "HF_TOKEN"
|
|
||||||
volumes:
|
|
||||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
|
||||||
|
|
||||||
- label: "Omni · Higgs-Audio V2 Test"
|
- label: "Omni · Higgs-Audio V2 Test"
|
||||||
timeout_in_minutes: 30
|
timeout_in_minutes: 30
|
||||||
commands:
|
commands:
|
||||||
|
|||||||
+4
-188
@@ -394,11 +394,11 @@ steps:
|
|||||||
depends_on: upload-nightly-pipeline
|
depends_on: upload-nightly-pipeline
|
||||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "tts-test"
|
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "tts-test"
|
||||||
steps:
|
steps:
|
||||||
- label: ":full_moon: TTS · Function Test with L4"
|
- label: ":full_moon: TTS · Function Test"
|
||||||
timeout_in_minutes: 90
|
timeout_in_minutes: 90
|
||||||
commands:
|
commands:
|
||||||
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||||
- pytest -s -v tests/e2e/ -m "full_model and L4 and tts" --run-level "full_model" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_voxcpm2_tts_expansion.py
|
- pytest -s -v tests/e2e/ -m "full_model and L4 and tts" --run-level "full_model" --ignore=tests/e2e/accuracy
|
||||||
agents:
|
agents:
|
||||||
queue: "gpu_1_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
|
queue: "gpu_1_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
|
||||||
plugins:
|
plugins:
|
||||||
@@ -413,69 +413,6 @@ steps:
|
|||||||
volumes:
|
volumes:
|
||||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||||
|
|
||||||
# - label: ":full_moon: TTS · Function Test with H100"
|
|
||||||
# timeout_in_minutes: 90
|
|
||||||
# commands:
|
|
||||||
# - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
|
||||||
# - pytest -s -v tests/e2e/ -m "full_model and H100 and tts" --run-level "full_model" --ignore=tests/e2e/accuracy
|
|
||||||
# agents:
|
|
||||||
# queue: "mithril-h100-pool"
|
|
||||||
# plugins:
|
|
||||||
# - kubernetes:
|
|
||||||
# podSpec:
|
|
||||||
# containers:
|
|
||||||
# - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
|
||||||
# resources:
|
|
||||||
# limits:
|
|
||||||
# nvidia.com/gpu: 1
|
|
||||||
# volumeMounts:
|
|
||||||
# - name: devshm
|
|
||||||
# mountPath: /dev/shm
|
|
||||||
# - name: hf-cache
|
|
||||||
# mountPath: /root/.cache/huggingface
|
|
||||||
# env:
|
|
||||||
# - name: HF_HOME
|
|
||||||
# value: /root/.cache/huggingface
|
|
||||||
# - name: HF_TOKEN
|
|
||||||
# valueFrom:
|
|
||||||
# secretKeyRef:
|
|
||||||
# name: hf-token-secret
|
|
||||||
# key: token
|
|
||||||
# nodeSelector:
|
|
||||||
# node.kubernetes.io/instance-type: gpu-h100-sxm
|
|
||||||
# volumes:
|
|
||||||
# - name: devshm
|
|
||||||
# emptyDir:
|
|
||||||
# medium: Memory
|
|
||||||
# - name: hf-cache
|
|
||||||
# hostPath:
|
|
||||||
# path: /mnt/hf-cache
|
|
||||||
# type: DirectoryOrCreate
|
|
||||||
|
|
||||||
- label: ":full_moon: TTS · Accuracy Test · VoxCPM2"
|
|
||||||
timeout_in_minutes: 30
|
|
||||||
commands:
|
|
||||||
- |
|
|
||||||
timeout 30m bash -c "
|
|
||||||
export VLLM_LOGGING_LEVEL=DEBUG
|
|
||||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
|
||||||
export VLLM_ALLOW_LONG_MAX_MODEL_LEN=1
|
|
||||||
pytest -s -v tests/e2e/online_serving/test_voxcpm2_tts_expansion.py -m 'full_model and cuda and tts' --run-level 'full_model'
|
|
||||||
"
|
|
||||||
agents:
|
|
||||||
queue: "gpu_1_queue"
|
|
||||||
plugins:
|
|
||||||
- docker#v5.2.0:
|
|
||||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
|
||||||
always-pull: true
|
|
||||||
propagate-environment: true
|
|
||||||
shm-size: "8gb"
|
|
||||||
environment:
|
|
||||||
- "HF_HOME=/fsx/hf_cache"
|
|
||||||
- "HF_TOKEN"
|
|
||||||
volumes:
|
|
||||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
|
||||||
|
|
||||||
- label: ":full_moon: TTS · Perf Test"
|
- label: ":full_moon: TTS · Perf Test"
|
||||||
key: nightly-tts-performance
|
key: nightly-tts-performance
|
||||||
timeout_in_minutes: 180
|
timeout_in_minutes: 180
|
||||||
@@ -522,8 +459,8 @@ steps:
|
|||||||
path: /mnt/hf-cache
|
path: /mnt/hf-cache
|
||||||
type: DirectoryOrCreate
|
type: DirectoryOrCreate
|
||||||
|
|
||||||
- label: ":full_moon: TTS · VoxCPM2 · Perf Test"
|
- label: ":full_moon: VoxCPM2 · Perf Test"
|
||||||
key: nightly-tts-performance-voxcpm2
|
key: nightly-voxcpm2-performance
|
||||||
timeout_in_minutes: 180
|
timeout_in_minutes: 180
|
||||||
commands:
|
commands:
|
||||||
- export BENCHMARK_DIR=tests/dfx/perf/results
|
- export BENCHMARK_DIR=tests/dfx/perf/results
|
||||||
@@ -758,87 +695,6 @@ steps:
|
|||||||
path: /mnt/hf-cache
|
path: /mnt/hf-cache
|
||||||
type: DirectoryOrCreate
|
type: DirectoryOrCreate
|
||||||
|
|
||||||
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100 · Multi-GPU DreamZero"
|
|
||||||
timeout_in_minutes: 120
|
|
||||||
commands:
|
|
||||||
- >-
|
|
||||||
pytest -sv
|
|
||||||
tests/e2e/online_serving/test_dreamzero_expansion.py
|
|
||||||
-m "full_model and diffusion and H100 and distributed_cuda"
|
|
||||||
--run-level "full_model"
|
|
||||||
agents:
|
|
||||||
queue: "mithril-h100-pool"
|
|
||||||
plugins:
|
|
||||||
- kubernetes:
|
|
||||||
podSpec:
|
|
||||||
containers:
|
|
||||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
|
||||||
resources:
|
|
||||||
limits:
|
|
||||||
nvidia.com/gpu: 2
|
|
||||||
volumeMounts:
|
|
||||||
- name: devshm
|
|
||||||
mountPath: /dev/shm
|
|
||||||
- name: hf-cache
|
|
||||||
mountPath: /root/.cache/huggingface
|
|
||||||
env:
|
|
||||||
- name: HF_HOME
|
|
||||||
value: /root/.cache/huggingface
|
|
||||||
- name: HF_TOKEN
|
|
||||||
valueFrom:
|
|
||||||
secretKeyRef:
|
|
||||||
name: hf-token-secret
|
|
||||||
key: token
|
|
||||||
nodeSelector:
|
|
||||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
|
||||||
volumes:
|
|
||||||
- name: devshm
|
|
||||||
emptyDir:
|
|
||||||
medium: Memory
|
|
||||||
- name: hf-cache
|
|
||||||
hostPath:
|
|
||||||
path: /mnt/hf-cache
|
|
||||||
type: DirectoryOrCreate
|
|
||||||
|
|
||||||
|
|
||||||
- label: ":full_moon: Diffusion X2I(&A&T) · GLM-Image Function Test with H100"
|
|
||||||
timeout_in_minutes: 120
|
|
||||||
commands:
|
|
||||||
- pytest -sv tests/e2e/online_serving/test_glm_image_expansion.py -m "full_model and diffusion and H100" --run-level "full_model"
|
|
||||||
agents:
|
|
||||||
queue: "mithril-h100-pool"
|
|
||||||
plugins:
|
|
||||||
- kubernetes:
|
|
||||||
podSpec:
|
|
||||||
containers:
|
|
||||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
|
||||||
resources:
|
|
||||||
limits:
|
|
||||||
nvidia.com/gpu: 2
|
|
||||||
volumeMounts:
|
|
||||||
- name: devshm
|
|
||||||
mountPath: /dev/shm
|
|
||||||
- name: hf-cache
|
|
||||||
mountPath: /root/.cache/huggingface
|
|
||||||
env:
|
|
||||||
- name: HF_HOME
|
|
||||||
value: /root/.cache/huggingface
|
|
||||||
- name: HF_TOKEN
|
|
||||||
valueFrom:
|
|
||||||
secretKeyRef:
|
|
||||||
name: hf-token-secret
|
|
||||||
key: token
|
|
||||||
nodeSelector:
|
|
||||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
|
||||||
volumes:
|
|
||||||
- name: devshm
|
|
||||||
emptyDir:
|
|
||||||
medium: Memory
|
|
||||||
- name: hf-cache
|
|
||||||
hostPath:
|
|
||||||
path: /mnt/hf-cache
|
|
||||||
type: DirectoryOrCreate
|
|
||||||
|
|
||||||
|
|
||||||
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with L4"
|
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with L4"
|
||||||
timeout_in_minutes: 60
|
timeout_in_minutes: 60
|
||||||
@@ -937,46 +793,6 @@ steps:
|
|||||||
path: /mnt/hf-cache
|
path: /mnt/hf-cache
|
||||||
type: DirectoryOrCreate
|
type: DirectoryOrCreate
|
||||||
|
|
||||||
- label: ":full_moon: HunyuanImage3-DIT · Accuracy Test"
|
|
||||||
timeout_in_minutes: 180
|
|
||||||
commands:
|
|
||||||
- export HUNYUAN_IMAGE3_MODEL="tencent/HunyuanImage-3.0-Instruct"
|
|
||||||
- export HUNYUAN_IMAGE3_DEVICES="0,1,2,3"
|
|
||||||
- pytest -s -v tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py -m full_model --run-level full_model
|
|
||||||
agents:
|
|
||||||
queue: "mithril-h100-pool"
|
|
||||||
plugins:
|
|
||||||
- kubernetes:
|
|
||||||
podSpec:
|
|
||||||
containers:
|
|
||||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
|
||||||
resources:
|
|
||||||
limits:
|
|
||||||
nvidia.com/gpu: 4
|
|
||||||
volumeMounts:
|
|
||||||
- name: devshm
|
|
||||||
mountPath: /dev/shm
|
|
||||||
- name: hf-cache
|
|
||||||
mountPath: /root/.cache/huggingface
|
|
||||||
env:
|
|
||||||
- name: HF_HOME
|
|
||||||
value: /root/.cache/huggingface
|
|
||||||
- name: HF_TOKEN
|
|
||||||
valueFrom:
|
|
||||||
secretKeyRef:
|
|
||||||
name: hf-token-secret
|
|
||||||
key: token
|
|
||||||
nodeSelector:
|
|
||||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
|
||||||
volumes:
|
|
||||||
- name: devshm
|
|
||||||
emptyDir:
|
|
||||||
medium: Memory
|
|
||||||
- name: hf-cache
|
|
||||||
hostPath:
|
|
||||||
path: /mnt/hf-cache
|
|
||||||
type: DirectoryOrCreate
|
|
||||||
|
|
||||||
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image"
|
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image"
|
||||||
key: nightly-diffusion-x2iat-performance-qwen-image
|
key: nightly-diffusion-x2iat-performance-qwen-image
|
||||||
timeout_in_minutes: 180
|
timeout_in_minutes: 180
|
||||||
|
|||||||
@@ -183,7 +183,7 @@ steps:
|
|||||||
depends_on: upload-ready-pipeline
|
depends_on: upload-ready-pipeline
|
||||||
commands:
|
commands:
|
||||||
- |
|
- |
|
||||||
timeout 20m bash -c "
|
timeout 15m bash -c "
|
||||||
pytest -sv tests/distributed/ -m 'core_model and cuda' --run-level 'core_model'
|
pytest -sv tests/distributed/ -m 'core_model and cuda' --run-level 'core_model'
|
||||||
"
|
"
|
||||||
agents:
|
agents:
|
||||||
@@ -273,42 +273,6 @@ steps:
|
|||||||
path: /mnt/hf-cache
|
path: /mnt/hf-cache
|
||||||
type: DirectoryOrCreate
|
type: DirectoryOrCreate
|
||||||
|
|
||||||
- label: "Omni · Qwen3-Omni Realtime WebSocket"
|
|
||||||
commands:
|
|
||||||
- |
|
|
||||||
timeout 30m bash -c "
|
|
||||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
|
||||||
pytest -s -v tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py::TestQwen3OmniRealtimeWebSocket::test_streaming_audio_input_pcm_output_async_chunk_smoke -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: "Omni · VoxCPM2 Test"
|
- label: "Omni · VoxCPM2 Test"
|
||||||
timeout_in_minutes: 20
|
timeout_in_minutes: 20
|
||||||
commands:
|
commands:
|
||||||
@@ -332,29 +296,6 @@ steps:
|
|||||||
volumes:
|
volumes:
|
||||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||||
|
|
||||||
- label: "TTS · VoxCPM2 Online Test"
|
|
||||||
timeout_in_minutes: 30
|
|
||||||
commands:
|
|
||||||
- |
|
|
||||||
timeout 30m bash -c "
|
|
||||||
export VLLM_LOGGING_LEVEL=DEBUG
|
|
||||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
|
||||||
pytest -s -v tests/e2e/online_serving/test_voxcpm2_tts.py -m 'core_model and tts' --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
|
|
||||||
shm-size: "8gb"
|
|
||||||
environment:
|
|
||||||
- "HF_HOME=/fsx/hf_cache"
|
|
||||||
- "HF_TOKEN"
|
|
||||||
volumes:
|
|
||||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
|
||||||
|
|
||||||
- label: "TTS · Qwen3-TTS CustomVoice Test"
|
- label: "TTS · Qwen3-TTS CustomVoice Test"
|
||||||
commands:
|
commands:
|
||||||
- |
|
- |
|
||||||
@@ -399,29 +340,6 @@ steps:
|
|||||||
volumes:
|
volumes:
|
||||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||||
|
|
||||||
- label: "TTS · Higgs-Audio-V3 Test"
|
|
||||||
commands:
|
|
||||||
- |
|
|
||||||
timeout 20m bash -c "
|
|
||||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
|
||||||
export VLLM_USE_DEEP_GEMM=0
|
|
||||||
export VLLM_MOE_USE_DEEP_GEMM=0
|
|
||||||
pytest -s -v tests/e2e/online_serving/test_higgs_audio_v3.py -m 'core_model' --run-level 'core_model'
|
|
||||||
"
|
|
||||||
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"
|
|
||||||
- "HF_TOKEN"
|
|
||||||
volumes:
|
|
||||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
|
||||||
|
|
||||||
- label: "Diffusion · Qwen Image Edit Test"
|
- label: "Diffusion · Qwen Image Edit Test"
|
||||||
commands:
|
commands:
|
||||||
- |
|
- |
|
||||||
|
|||||||
@@ -53,9 +53,6 @@
|
|||||||
{% endif %}
|
{% endif %}
|
||||||
{{ indented_cmd | safe }}
|
{{ indented_cmd | safe }}
|
||||||
priority: 100
|
priority: 100
|
||||||
{% if step.timeout_in_minutes %}
|
|
||||||
timeout_in_minutes: {{ step.timeout_in_minutes }}
|
|
||||||
{% endif %}
|
|
||||||
{% if step.grade and step.grade == "Blocking" %}
|
{% if step.grade and step.grade == "Blocking" %}
|
||||||
soft_fail: false
|
soft_fail: false
|
||||||
{% else %}
|
{% else %}
|
||||||
|
|||||||
@@ -16,7 +16,6 @@ Easy, fast, and cheap omni-modality model serving for everyone
|
|||||||
---
|
---
|
||||||
|
|
||||||
*Latest News* 🔥
|
*Latest News* 🔥
|
||||||
- [2026/06] We released [0.22.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.22.0) - an **omnimodal world-model** release aligned with vLLM 0.22, featuring [Nvidia Cosmos3](recipes/cosmos3/Cosmos3-Nano.md)/DreamZero world model support, expanded quantization coverage across Blackwell/NPU/XPU, TTS production improvements, new models including MiniCPM-o 4.5, MOSS-TTS, and Lance, plus RL integration with veRL-Omni.
|
|
||||||
- [2026/05] We released [0.20.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.20.0) - refreshes the serving/runtime stack for large-scale omni workloads, and improves diffusion model performance, quantization, and hardware readiness across CUDA, ROCm, MUSA, NPU, and XPU backends.
|
- [2026/05] We released [0.20.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.20.0) - refreshes the serving/runtime stack for large-scale omni workloads, and improves diffusion model performance, quantization, and hardware readiness across CUDA, ROCm, MUSA, NPU, and XPU backends.
|
||||||
- [2026/03] We released [0.18.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.18.0) - strengthens the core runtime through a large entrypoint refactor and scheduler/runtime cleanups, expands unified quantization and diffusion execution, broadens multimodal model coverage, and improves production readiness across audio, omni, image, video, RL, and multi-platform deployments.
|
- [2026/03] We released [0.18.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.18.0) - strengthens the core runtime through a large entrypoint refactor and scheduler/runtime cleanups, expands unified quantization and diffusion execution, broadens multimodal model coverage, and improves production readiness across audio, omni, image, video, RL, and multi-platform deployments.
|
||||||
- [2026/03] Check out our first public [project deepdive](https://youtu.be/sgwNfsNnR9I) at the vLLM Hong Kong Meetup!
|
- [2026/03] Check out our first public [project deepdive](https://youtu.be/sgwNfsNnR9I) at the vLLM Hong Kong Meetup!
|
||||||
|
|||||||
@@ -151,9 +151,6 @@ class VllmOmniImageClient:
|
|||||||
seed: int | None = None,
|
seed: int | None = None,
|
||||||
negative_prompt: str | None = None,
|
negative_prompt: str | None = None,
|
||||||
output_compression: int | None = None,
|
output_compression: int | None = None,
|
||||||
bot_task: str | None = None,
|
|
||||||
sys_type: str | None = None,
|
|
||||||
system_prompt: str | None = None,
|
|
||||||
) -> Image.Image:
|
) -> Image.Image:
|
||||||
if not isinstance(images, list):
|
if not isinstance(images, list):
|
||||||
images = [images]
|
images = [images]
|
||||||
@@ -173,12 +170,6 @@ class VllmOmniImageClient:
|
|||||||
data["negative_prompt"] = negative_prompt
|
data["negative_prompt"] = negative_prompt
|
||||||
if output_compression is not None:
|
if output_compression is not None:
|
||||||
data["output_compression"] = str(output_compression)
|
data["output_compression"] = str(output_compression)
|
||||||
if bot_task is not None:
|
|
||||||
data["bot_task"] = bot_task
|
|
||||||
if sys_type is not None:
|
|
||||||
data["sys_type"] = sys_type
|
|
||||||
if system_prompt is not None:
|
|
||||||
data["system_prompt"] = system_prompt
|
|
||||||
|
|
||||||
files = [
|
files = [
|
||||||
(
|
(
|
||||||
|
|||||||
@@ -64,47 +64,6 @@ def resolve_model_name(*, model_name: str | None, model: str | None = None, outp
|
|||||||
raise ValueError("model-name is required when it cannot be inferred from model or output-root")
|
raise ValueError("model-name is required when it cannot be inferred from model or output-root")
|
||||||
|
|
||||||
|
|
||||||
def _coerce_score_int(value: Any) -> int:
|
|
||||||
if isinstance(value, bool):
|
|
||||||
raise TypeError(f"Expected numeric score, got bool: {value!r}")
|
|
||||||
if isinstance(value, (int, float)):
|
|
||||||
return int(value)
|
|
||||||
if isinstance(value, str):
|
|
||||||
stripped = value.strip()
|
|
||||||
if stripped.lstrip("-").isdigit():
|
|
||||||
return int(stripped)
|
|
||||||
raise TypeError(f"Expected numeric score, got {type(value).__name__}: {value!r}")
|
|
||||||
|
|
||||||
|
|
||||||
def _expand_score_entry(entry: Any) -> list[int]:
|
|
||||||
if isinstance(entry, dict):
|
|
||||||
if "score" in entry:
|
|
||||||
return [_coerce_score_int(entry["score"])]
|
|
||||||
if "value" in entry:
|
|
||||||
return [_coerce_score_int(entry["value"])]
|
|
||||||
numeric_values = [_coerce_score_int(value) for value in entry.values()]
|
|
||||||
if numeric_values:
|
|
||||||
return numeric_values
|
|
||||||
raise TypeError(f"Score dict has no numeric values: {entry!r}")
|
|
||||||
return [_coerce_score_int(entry)]
|
|
||||||
|
|
||||||
|
|
||||||
def _normalize_score_list(score: Any) -> list[int]:
|
|
||||||
if isinstance(score, dict):
|
|
||||||
entries: list[Any] = [score]
|
|
||||||
elif isinstance(score, list):
|
|
||||||
entries = score
|
|
||||||
else:
|
|
||||||
entries = [score]
|
|
||||||
|
|
||||||
normalized: list[int] = []
|
|
||||||
for entry in entries:
|
|
||||||
normalized.extend(_expand_score_entry(entry))
|
|
||||||
if not normalized:
|
|
||||||
raise ValueError(f"Empty score list after normalization: {score!r}")
|
|
||||||
return normalized
|
|
||||||
|
|
||||||
|
|
||||||
def parse_score_payload(raw_text: str) -> dict[str, Any]:
|
def parse_score_payload(raw_text: str) -> dict[str, Any]:
|
||||||
try:
|
try:
|
||||||
parsed = extract_json_object(raw_text)
|
parsed = extract_json_object(raw_text)
|
||||||
@@ -114,7 +73,10 @@ def parse_score_payload(raw_text: str) -> dict[str, Any]:
|
|||||||
parsed = {"score": json.loads(stripped), "reasoning": ""}
|
parsed = {"score": json.loads(stripped), "reasoning": ""}
|
||||||
else:
|
else:
|
||||||
raise
|
raise
|
||||||
parsed["score"] = _normalize_score_list(parsed.get("score", []))
|
score = parsed.get("score", [])
|
||||||
|
if not isinstance(score, list):
|
||||||
|
score = [score]
|
||||||
|
parsed["score"] = [int(value) for value in score]
|
||||||
return parsed
|
return parsed
|
||||||
|
|
||||||
|
|
||||||
@@ -429,9 +391,6 @@ class GEditBenchRunner:
|
|||||||
num_inference_steps: int = 20,
|
num_inference_steps: int = 20,
|
||||||
guidance_scale: float | None = None,
|
guidance_scale: float | None = None,
|
||||||
seed: int | None = 42,
|
seed: int | None = 42,
|
||||||
bot_task: str | None = None,
|
|
||||||
sys_type: str | None = None,
|
|
||||||
system_prompt: str | None = None,
|
|
||||||
):
|
):
|
||||||
self.dataset_ref = dataset_ref
|
self.dataset_ref = dataset_ref
|
||||||
self.output_root = output_root
|
self.output_root = output_root
|
||||||
@@ -441,9 +400,6 @@ class GEditBenchRunner:
|
|||||||
self.num_inference_steps = num_inference_steps
|
self.num_inference_steps = num_inference_steps
|
||||||
self.guidance_scale = guidance_scale
|
self.guidance_scale = guidance_scale
|
||||||
self.seed = seed
|
self.seed = seed
|
||||||
self.bot_task = bot_task
|
|
||||||
self.sys_type = sys_type
|
|
||||||
self.system_prompt = system_prompt
|
|
||||||
self.client = VllmOmniImageClient(base_url=base_url, api_key=api_key)
|
self.client = VllmOmniImageClient(base_url=base_url, api_key=api_key)
|
||||||
|
|
||||||
def generate(
|
def generate(
|
||||||
@@ -522,9 +478,6 @@ class GEditBenchRunner:
|
|||||||
num_inference_steps=self.num_inference_steps,
|
num_inference_steps=self.num_inference_steps,
|
||||||
guidance_scale=self.guidance_scale,
|
guidance_scale=self.guidance_scale,
|
||||||
seed=self.seed,
|
seed=self.seed,
|
||||||
bot_task=self.bot_task,
|
|
||||||
sys_type=self.sys_type,
|
|
||||||
system_prompt=self.system_prompt,
|
|
||||||
)
|
)
|
||||||
edited_image.save(output_path)
|
edited_image.save(output_path)
|
||||||
return {
|
return {
|
||||||
@@ -680,25 +633,6 @@ def build_parser() -> argparse.ArgumentParser:
|
|||||||
generate.add_argument("--num-inference-steps", type=int, default=20)
|
generate.add_argument("--num-inference-steps", type=int, default=20)
|
||||||
generate.add_argument("--guidance-scale", type=float, default=None)
|
generate.add_argument("--guidance-scale", type=float, default=None)
|
||||||
generate.add_argument("--seed", type=int, default=42)
|
generate.add_argument("--seed", type=int, default=42)
|
||||||
generate.add_argument(
|
|
||||||
"--bot-task",
|
|
||||||
type=str,
|
|
||||||
default=None,
|
|
||||||
choices=["think", "recaption", "think_recaption", "vanilla"],
|
|
||||||
help="HunyuanImage-3 prompting mode for /v1/images/edits (optional).",
|
|
||||||
)
|
|
||||||
generate.add_argument(
|
|
||||||
"--sys-type",
|
|
||||||
type=str,
|
|
||||||
default=None,
|
|
||||||
help="Hunyuan system-prompt type (e.g. en_unified). Defaults from --bot-task when omitted.",
|
|
||||||
)
|
|
||||||
generate.add_argument(
|
|
||||||
"--system-prompt",
|
|
||||||
type=str,
|
|
||||||
default=None,
|
|
||||||
help="Custom Hunyuan system prompt body (sys_type=custom).",
|
|
||||||
)
|
|
||||||
generate.add_argument("--workers", type=int, default=1)
|
generate.add_argument("--workers", type=int, default=1)
|
||||||
generate.add_argument("--max-samples", type=int, default=None)
|
generate.add_argument("--max-samples", type=int, default=None)
|
||||||
generate.add_argument("--samples-per-group", type=int, default=None)
|
generate.add_argument("--samples-per-group", type=int, default=None)
|
||||||
@@ -741,9 +675,6 @@ def main(argv: list[str] | None = None) -> int:
|
|||||||
num_inference_steps=args.num_inference_steps,
|
num_inference_steps=args.num_inference_steps,
|
||||||
guidance_scale=args.guidance_scale,
|
guidance_scale=args.guidance_scale,
|
||||||
seed=args.seed,
|
seed=args.seed,
|
||||||
bot_task=args.bot_task,
|
|
||||||
sys_type=args.sys_type,
|
|
||||||
system_prompt=args.system_prompt,
|
|
||||||
)
|
)
|
||||||
records = runner.generate(
|
records = runner.generate(
|
||||||
model_name=model_name,
|
model_name=model_name,
|
||||||
|
|||||||
@@ -1,10 +0,0 @@
|
|||||||
se001|ocean waves crashing on a rocky beach|10.0
|
|
||||||
se002|heavy rain on a metal roof at night|8.0
|
|
||||||
se003|a dog barking in the distance|6.0
|
|
||||||
se004|a crackling campfire in a quiet forest|10.0
|
|
||||||
se005|footsteps on dry leaves|6.0
|
|
||||||
se006|thunder rolling across a wide valley|8.0
|
|
||||||
se007|a busy coffee shop with chatter and espresso machines|10.0
|
|
||||||
se008|wind howling through tall pine trees|8.0
|
|
||||||
se009|a small stream flowing over pebbles|10.0
|
|
||||||
se010|a creaking wooden door slowly opening|6.0
|
|
||||||
@@ -1,10 +0,0 @@
|
|||||||
ttsd001|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Hello and welcome to the show today. [S2] Thanks for having me, it's a pleasure to be here. [S1] Let's dive right in.
|
|
||||||
ttsd002|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Did you finish the report? [S2] Almost — I'm just polishing the conclusion. [S1] Great, send it over when you're done.
|
|
||||||
ttsd003|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Have you tried that new restaurant downtown? [S2] Yes, the pasta is amazing. [S1] I'll book a table for Friday.
|
|
||||||
ttsd004|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] What did you think of yesterday's meeting? [S2] Productive, but a bit long. [S1] Agreed, we should set an hour cap next time.
|
|
||||||
ttsd005|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] The weather is beautiful today. [S2] Perfect for a hike. [S1] Want to drive out to the trails this afternoon?
|
|
||||||
ttsd006|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] How did the demo go? [S2] Better than expected — the client signed off on the prototype. [S1] Fantastic news!
|
|
||||||
ttsd007|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Coffee or tea this morning? [S2] Coffee please, with a splash of milk. [S1] Coming right up.
|
|
||||||
ttsd008|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Did you read the email about the deadline change? [S2] I did, we now have until next Friday. [S1] That should give us enough buffer.
|
|
||||||
ttsd009|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Are you joining the team dinner tonight? [S2] I'll be there around seven. [S1] See you then.
|
|
||||||
ttsd010|prompt-wavs/spk1.wav|prompt-wavs/spk2.wav|[S1] Have we backed up the database? [S2] Last backup ran two hours ago. [S1] Good, I'll trigger a fresh snapshot before the migration.
|
|
||||||
Whitespace-only changes.
@@ -53,8 +53,6 @@ _TASK_TO_DATASET: dict[str, str] = {
|
|||||||
"voice_clone": "seed-tts",
|
"voice_clone": "seed-tts",
|
||||||
"default_voice": "seed-tts-text",
|
"default_voice": "seed-tts-text",
|
||||||
"voice_design": "seed-tts-design",
|
"voice_design": "seed-tts-design",
|
||||||
"dialogue": "ttsd",
|
|
||||||
"sound_effect": "sound-effect",
|
|
||||||
}
|
}
|
||||||
|
|
||||||
# Default design dataset path (bundled with the repo)
|
# Default design dataset path (bundled with the repo)
|
||||||
@@ -83,7 +81,6 @@ def build_bench_args(
|
|||||||
output_dir: str | None,
|
output_dir: str | None,
|
||||||
result_filename: str | None,
|
result_filename: str | None,
|
||||||
extra_cli_args: list[str],
|
extra_cli_args: list[str],
|
||||||
output_len: int | None = None,
|
|
||||||
) -> list[str]:
|
) -> list[str]:
|
||||||
"""Build the ``vllm bench serve --omni`` command for one (task, concurrency) run."""
|
"""Build the ``vllm bench serve --omni`` command for one (task, concurrency) run."""
|
||||||
dataset_name = _TASK_TO_DATASET[task]
|
dataset_name = _TASK_TO_DATASET[task]
|
||||||
@@ -142,9 +139,6 @@ def build_bench_args(
|
|||||||
if wer_eval:
|
if wer_eval:
|
||||||
cmd.append("--seed-tts-wer-eval")
|
cmd.append("--seed-tts-wer-eval")
|
||||||
|
|
||||||
if output_len is not None:
|
|
||||||
cmd += ["--hf-output-len", str(output_len)]
|
|
||||||
|
|
||||||
if output_dir or result_filename:
|
if output_dir or result_filename:
|
||||||
out_dir = output_dir or "."
|
out_dir = output_dir or "."
|
||||||
os.makedirs(out_dir, exist_ok=True)
|
os.makedirs(out_dir, exist_ok=True)
|
||||||
@@ -244,15 +238,6 @@ def main() -> None:
|
|||||||
parser.add_argument("--host", default="localhost")
|
parser.add_argument("--host", default="localhost")
|
||||||
parser.add_argument("--port", type=int, default=8000)
|
parser.add_argument("--port", type=int, default=8000)
|
||||||
parser.add_argument("--model-configs", default=str(_DEFAULT_MODEL_CONFIGS), help="Path to model_configs.yaml")
|
parser.add_argument("--model-configs", default=str(_DEFAULT_MODEL_CONFIGS), help="Path to model_configs.yaml")
|
||||||
parser.add_argument(
|
|
||||||
"--output-len",
|
|
||||||
type=int,
|
|
||||||
default=None,
|
|
||||||
help="Override per-request max_new_tokens forwarded to the server "
|
|
||||||
"(seed-tts datasets default to 2048; the MOSS-TTS full talker takes "
|
|
||||||
"this as 'max_new_frames' so smaller values cap audio length per "
|
|
||||||
"request and keep bench runtime tractable).",
|
|
||||||
)
|
|
||||||
parser.add_argument("extra", nargs=argparse.REMAINDER, help="Extra args passed directly to vllm bench serve")
|
parser.add_argument("extra", nargs=argparse.REMAINDER, help="Extra args passed directly to vllm bench serve")
|
||||||
args = parser.parse_args()
|
args = parser.parse_args()
|
||||||
|
|
||||||
@@ -307,7 +292,6 @@ def main() -> None:
|
|||||||
output_dir=args.output_dir,
|
output_dir=args.output_dir,
|
||||||
result_filename=result_filename,
|
result_filename=result_filename,
|
||||||
extra_cli_args=args.extra or [],
|
extra_cli_args=args.extra or [],
|
||||||
output_len=args.output_len,
|
|
||||||
)
|
)
|
||||||
result = run_one_benchmark(cmd)
|
result = run_one_benchmark(cmd)
|
||||||
if result is not None:
|
if result is not None:
|
||||||
|
|||||||
@@ -14,7 +14,6 @@ models:
|
|||||||
supported_tasks: [default_voice, voice_design]
|
supported_tasks: [default_voice, voice_design]
|
||||||
backend: openai-audio-speech
|
backend: openai-audio-speech
|
||||||
endpoint: /v1/audio/speech
|
endpoint: /v1/audio/speech
|
||||||
trust_remote_code: true
|
|
||||||
task_extra_body:
|
task_extra_body:
|
||||||
default_voice:
|
default_voice:
|
||||||
voice: Vivian
|
voice: Vivian
|
||||||
@@ -28,7 +27,6 @@ models:
|
|||||||
supported_tasks: [voice_clone]
|
supported_tasks: [voice_clone]
|
||||||
backend: openai-audio-speech
|
backend: openai-audio-speech
|
||||||
endpoint: /v1/audio/speech
|
endpoint: /v1/audio/speech
|
||||||
trust_remote_code: true
|
|
||||||
task_extra_body:
|
task_extra_body:
|
||||||
voice_clone:
|
voice_clone:
|
||||||
task_type: Base
|
task_type: Base
|
||||||
@@ -38,66 +36,7 @@ models:
|
|||||||
trust_remote_code: true
|
trust_remote_code: true
|
||||||
backend: openai-audio-speech
|
backend: openai-audio-speech
|
||||||
endpoint: /v1/audio/speech
|
endpoint: /v1/audio/speech
|
||||||
trust_remote_code: true
|
|
||||||
task_extra_body:
|
task_extra_body:
|
||||||
voice_clone: {}
|
voice_clone: {}
|
||||||
default_voice:
|
default_voice:
|
||||||
voice: default
|
voice: default
|
||||||
|
|
||||||
# MOSS-TTS family. After the serving_speech.py variant-detection refactor
|
|
||||||
# all five variants now go through OpenAI /v1/audio/speech with the
|
|
||||||
# variant-specific request fields (ambient_sound, ref_audio_2, instructions).
|
|
||||||
OpenMOSS-Team/MOSS-TTS:
|
|
||||||
supported_tasks: [voice_clone]
|
|
||||||
backend: openai-audio-speech
|
|
||||||
endpoint: /v1/audio/speech
|
|
||||||
task_extra_body:
|
|
||||||
voice_clone: {}
|
|
||||||
|
|
||||||
# MOSS-TTS-v1.5: continued-training upgrade of MOSS-TTS 1.0 — same
|
|
||||||
# MossTTSDelay-8B architecture and API, so identical wiring. Adds 31-language
|
|
||||||
# synthesis via the `language` field and inline `[pause Xs]` markers.
|
|
||||||
OpenMOSS-Team/MOSS-TTS-v1.5:
|
|
||||||
supported_tasks: [voice_clone]
|
|
||||||
backend: openai-audio-speech
|
|
||||||
endpoint: /v1/audio/speech
|
|
||||||
task_extra_body:
|
|
||||||
voice_clone: {}
|
|
||||||
|
|
||||||
OpenMOSS-Team/MOSS-TTS-Realtime:
|
|
||||||
supported_tasks: [voice_clone]
|
|
||||||
backend: openai-audio-speech
|
|
||||||
endpoint: /v1/audio/speech
|
|
||||||
task_extra_body:
|
|
||||||
voice_clone: {}
|
|
||||||
|
|
||||||
# MOSS-TTSD synthesizes [S1]/[S2] dialogue. The seed-tts voice_clone dataset
|
|
||||||
# has single-speaker rows; bench piggybacks on it (both speakers share the
|
|
||||||
# same ref_audio). For real dialogue evaluation use the ttsd dataset.
|
|
||||||
OpenMOSS-Team/MOSS-TTSD-v1.0:
|
|
||||||
supported_tasks: [voice_clone, dialogue]
|
|
||||||
backend: openai-audio-speech
|
|
||||||
endpoint: /v1/audio/speech
|
|
||||||
task_extra_body:
|
|
||||||
voice_clone: {}
|
|
||||||
dialogue: {}
|
|
||||||
|
|
||||||
# MOSS-SoundEffect ignores text/ref_audio; expects ambient_sound +
|
|
||||||
# duration_seconds. The sound_effect dataset module emits those.
|
|
||||||
OpenMOSS-Team/MOSS-SoundEffect:
|
|
||||||
supported_tasks: [sound_effect]
|
|
||||||
backend: openai-audio-speech
|
|
||||||
endpoint: /v1/audio/speech
|
|
||||||
task_extra_body:
|
|
||||||
sound_effect:
|
|
||||||
duration_seconds: 10
|
|
||||||
|
|
||||||
# MOSS-VoiceGenerator: instructions-driven (no ref_audio). seed-tts-design
|
|
||||||
# dataset already emits {instructions, target_text}, so voice_design works
|
|
||||||
# out of the box once the serving validator accepts instructions.
|
|
||||||
OpenMOSS-Team/MOSS-VoiceGenerator:
|
|
||||||
supported_tasks: [voice_design]
|
|
||||||
backend: openai-audio-speech
|
|
||||||
endpoint: /v1/audio/speech
|
|
||||||
task_extra_body:
|
|
||||||
voice_design: {}
|
|
||||||
@@ -1,8 +1,5 @@
|
|||||||
# Target vllm version: v0.21.0 (releases/v0.21.0 branch).
|
|
||||||
# The pinned VLLM_PRECOMPILED_WHEEL_COMMIT and VLLM_BASE_TAG below must
|
|
||||||
# be updated on every rebase to match the new upstream vllm commit/tag.
|
|
||||||
ARG VLLM_BASE_IMAGE=vllm/vllm-openai
|
ARG VLLM_BASE_IMAGE=vllm/vllm-openai
|
||||||
ARG VLLM_BASE_TAG=v0.22.0
|
ARG VLLM_BASE_TAG=v0.21.0
|
||||||
FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
|
FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
|
||||||
ARG APP_DIR=/workspace/vllm-omni
|
ARG APP_DIR=/workspace/vllm-omni
|
||||||
WORKDIR ${APP_DIR}
|
WORKDIR ${APP_DIR}
|
||||||
|
|||||||
@@ -1,4 +1,4 @@
|
|||||||
ARG BASE_IMAGE=vllm/vllm-openai:v0.22.0
|
ARG BASE_IMAGE=vllm/vllm-openai:v0.21.0
|
||||||
FROM ${BASE_IMAGE}
|
FROM ${BASE_IMAGE}
|
||||||
|
|
||||||
ARG COMMON_WORKDIR=/app
|
ARG COMMON_WORKDIR=/app
|
||||||
|
|||||||
@@ -1,4 +1,4 @@
|
|||||||
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.22.0
|
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.21.0
|
||||||
FROM ${BASE_IMAGE} AS base
|
FROM ${BASE_IMAGE} AS base
|
||||||
|
|
||||||
# Declare a variable to know if we want to use the nightly build or the stable build.
|
# Declare a variable to know if we want to use the nightly build or the stable build.
|
||||||
|
|||||||
@@ -88,7 +88,7 @@ ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
|
|||||||
ENV UV_INDEX_STRATEGY="unsafe-best-match"
|
ENV UV_INDEX_STRATEGY="unsafe-best-match"
|
||||||
ENV UV_LINK_MODE="copy"
|
ENV UV_LINK_MODE="copy"
|
||||||
|
|
||||||
ARG VLLM_VERSION=v0.22.0
|
ARG VLLM_VERSION=v0.21.0
|
||||||
RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm
|
RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm
|
||||||
WORKDIR /workspace/vllm
|
WORKDIR /workspace/vllm
|
||||||
|
|
||||||
|
|||||||
+3
-111
@@ -2,7 +2,7 @@
|
|||||||
The vllm bench command launches the vLLM-Omni benchmark to evaluate the performance of multimodal models.
|
The vllm bench command launches the vLLM-Omni benchmark to evaluate the performance of multimodal models.
|
||||||
|
|
||||||
## Notes
|
## Notes
|
||||||
We currently only support using the "openai-chat-omni" and "openai-image-edits-omni" backend.
|
We currently only support using the "openai-chat-omni" backend.
|
||||||
|
|
||||||
## Basic Parameter Description
|
## Basic Parameter Description
|
||||||
You can use `vllm bench serve --omni --help=all` to get descriptions of all parameters. The commonly used parameters are described below:
|
You can use `vllm bench serve --omni --help=all` to get descriptions of all parameters. The commonly used parameters are described below:
|
||||||
@@ -52,11 +52,7 @@ You can use `vllm bench serve --omni --help=all` to get descriptions of all para
|
|||||||
- `--percentile-metrics`
|
- `--percentile-metrics`
|
||||||
"Comma-separated list of selected metrics to report percentiles."
|
"Comma-separated list of selected metrics to report percentiles."
|
||||||
"This argument specifies the metrics to report percentiles."
|
"This argument specifies the metrics to report percentiles."
|
||||||
'Allowed metric names are "ttft", "tpot", "itl", "ttfc", "tpoc", "icl", '
|
'Allowed metric names are "ttft", "tpot", "itl", "e2el", "audio_ttfp", "audio_rtf", "audio_duration". '
|
||||||
'"tpop", "e2el", "audio_ttfp", "audio_rtf", "audio_duration". '
|
|
||||||
|
|
||||||
- `-print-stage`
|
|
||||||
Print per-stage benchmark metrics for --omni serving when stage metrics are returned by the server. Disabled by default.
|
|
||||||
|
|
||||||
- `--save-result`
|
- `--save-result`
|
||||||
Specify to save benchmark results to a json file
|
Specify to save benchmark results to a json file
|
||||||
@@ -280,9 +276,7 @@ Generate synthetic image、video、audio inputs alongside random text prompts to
|
|||||||
|
|
||||||
Notes:
|
Notes:
|
||||||
|
|
||||||
- Works only with online benchmark via:
|
- Works only with online benchmark via the OpenAI backend (`--backend openai-chat-omni`) and endpoint `/v1/chat/completions`.
|
||||||
- the OpenAI chat backend (`--backend openai-chat-omni`) and endpoint `/v1/chat/completions`.
|
|
||||||
- the OpenAI edit image backend (`--backend openai-image-edits-omni`) and endpoint `/v1/images/edits`.
|
|
||||||
|
|
||||||
Start the server (example):
|
Start the server (example):
|
||||||
|
|
||||||
@@ -363,105 +357,3 @@ How sampling works:
|
|||||||
This should be seen as an edge case, and if this behavior can be avoided by setting `--random-mm-limit-mm-per-prompt` to a large number. Note that this might result in errors due to engine config `--limit-mm-per-prompt`.
|
This should be seen as an edge case, and if this behavior can be avoided by setting `--random-mm-limit-mm-per-prompt` to a large number. Note that this might result in errors due to engine config `--limit-mm-per-prompt`.
|
||||||
- The resulting request contains synthetic image data in `multi_modal_data` (OpenAI Chat format). When `random-mm` is used with the OpenAI Chat backend, prompts remain text and MM content is attached via `multi_modal_data`.
|
- The resulting request contains synthetic image data in `multi_modal_data` (OpenAI Chat format). When `random-mm` is used with the OpenAI Chat backend, prompts remain text and MM content is attached via `multi_modal_data`.
|
||||||
</details>
|
</details>
|
||||||
|
|
||||||
### Multi-Stage Benchmark
|
|
||||||
|
|
||||||
<details class="admonition abstract" markdown="1">
|
|
||||||
<summary>Show more</summary>
|
|
||||||
|
|
||||||
Start the server (change the port and model):
|
|
||||||
|
|
||||||
```bash
|
|
||||||
vllm serve --omni --port 29999 --model ~/Qwen3-Omni-30B-A3B-Instruct
|
|
||||||
```
|
|
||||||
|
|
||||||
Then run the benchmarking script (remember to add the flag `--print-stage`):
|
|
||||||
```bash
|
|
||||||
vllm bench serve --omni \
|
|
||||||
--dataset-name random \
|
|
||||||
--port 29999 \
|
|
||||||
--max-concurrency 1 \
|
|
||||||
--num-warmups 2 \
|
|
||||||
--model ~/Qwen3-Omni-30B-A3B-Instruct \
|
|
||||||
--endpoint /v1/chat/completions \
|
|
||||||
--backend openai-chat-omni \
|
|
||||||
--num-prompts 3 \
|
|
||||||
--random-input-len 2500 \
|
|
||||||
--ignore-eos \
|
|
||||||
--percentile-metrics ttft,tpot,itl,e2el,audio_ttfp,audio_rtf,ttfc,tpoc,icl \
|
|
||||||
--random-output-len 900 \
|
|
||||||
--extra_body '{"modalities": ["text","audio"]}' \
|
|
||||||
--print-stage
|
|
||||||
```
|
|
||||||
|
|
||||||
Besides the "Serving Benchmark Result", there will be "Stage Benchmark Result" below:
|
|
||||||
```text
|
|
||||||
============= Stage Benchmark Result =============
|
|
||||||
=============== Stage 0 (thinker) ================
|
|
||||||
-------------------Stage Timing-------------------
|
|
||||||
Mean stage_gen_time (ms): 18243.16
|
|
||||||
Median stage_gen_time (ms): 18220.83
|
|
||||||
P99 stage_gen_time (ms): 18296.75
|
|
||||||
================== Text Result ===================
|
|
||||||
Stage generated tokens: 2700
|
|
||||||
-----------Serving Time to First Token------------
|
|
||||||
Mean Serving TTFT (ms): 816.31
|
|
||||||
Median Serving TTFT (ms): 818.94
|
|
||||||
P99 Serving TTFT (ms): 829.76
|
|
||||||
-----Time per Output Token (excl. 1st token)------
|
|
||||||
Mean TPOP (ms): 19.41
|
|
||||||
Median TPOP (ms): 19.40
|
|
||||||
P99 TPOP (ms): 19.47
|
|
||||||
---------------Inter-token Latency----------------
|
|
||||||
Mean ITL (ms): 19.41
|
|
||||||
Median ITL (ms): 19.12
|
|
||||||
P99 ITL (ms): 23.58
|
|
||||||
================ Stage 1 (talker) ================
|
|
||||||
-------------------Stage Timing-------------------
|
|
||||||
Mean stage_gen_time (ms): 18849.64
|
|
||||||
Median stage_gen_time (ms): 13641.15
|
|
||||||
P99 stage_gen_time (ms): 33655.58
|
|
||||||
============= Internal Stream Result =============
|
|
||||||
-----------Serving Time to First Chunk------------
|
|
||||||
Mean Serving TTFC (ms): 1029.14
|
|
||||||
Median Serving TTFC (ms): 1027.98
|
|
||||||
P99 Serving TTFC (ms): 1032.59
|
|
||||||
-----Time per Output Chunk (excl. 1st chunk)------
|
|
||||||
Mean TPOP (ms): 17.53
|
|
||||||
Median TPOP (ms): 18.70
|
|
||||||
P99 TPOP (ms): 18.88
|
|
||||||
---------------Inter-chunk Latency----------------
|
|
||||||
Mean ICL (ms): 16.27
|
|
||||||
Median ICL (ms): 18.19
|
|
||||||
P99 ICL (ms): 25.86
|
|
||||||
=============== Stage 2 (code2wav) ===============
|
|
||||||
-------------------Stage Timing-------------------
|
|
||||||
Mean stage_gen_time (ms): 18860.74
|
|
||||||
Median stage_gen_time (ms): 13655.84
|
|
||||||
P99 stage_gen_time (ms): 33659.79
|
|
||||||
================== Audio Result ==================
|
|
||||||
Stage audio duration generated(s): 260.11
|
|
||||||
Stage audio frames generated: 6242730
|
|
||||||
-----------Serving Time to First Packet-----------
|
|
||||||
Mean Serving AUDIO_TTFP (ms): 1366.16
|
|
||||||
Median Serving AUDIO_TTFP (ms): 1367.20
|
|
||||||
P99 Serving AUDIO_TTFP (ms): 1368.26
|
|
||||||
-----------------Real Time Factor-----------------
|
|
||||||
Mean AUDIO_RTF: 0.24
|
|
||||||
Median AUDIO_RTF: 0.26
|
|
||||||
P99 AUDIO_RTF: 0.27
|
|
||||||
```
|
|
||||||
Explanation:
|
|
||||||
- stage_gen_time: Time from submitting a request to a specific stage to that stage finishing generation.
|
|
||||||
|
|
||||||
- Serving TTFC (Time to First Chunk): Time from the HTTP request being accepted by the serving frontend to
|
|
||||||
the stage producing its first non-empty output chunk. Keep the name TTFT in text stage, TTFP in audio stage
|
|
||||||
for easier mapping.
|
|
||||||
|
|
||||||
- TPOP (Time per Output Chunk): Average time from the first output chunk to stage completion, divided by the
|
|
||||||
number of remaining output chunks. The TPOP abbreviation follows Qwen3.5-Omni Technical Report.
|
|
||||||
|
|
||||||
- ICL (Inter-Chunk Latency): Time between two consecutive output chunks produced by the same stage. Keep the
|
|
||||||
name ITL in text stage for easier mapping.
|
|
||||||
|
|
||||||
</details>
|
|
||||||
@@ -13,10 +13,6 @@ This section contains design documents and architecture specifications for vLLM-
|
|||||||
- [Adding Step Execution Support for Diffusion Pipelines](feature/diffusion_step_execution.md)
|
- [Adding Step Execution Support for Diffusion Pipelines](feature/diffusion_step_execution.md)
|
||||||
- [Continuous Batching for Step-Wise Diffusion](feature/diffusion_continuous_batching.md)
|
- [Continuous Batching for Step-Wise Diffusion](feature/diffusion_continuous_batching.md)
|
||||||
|
|
||||||
## Infrastructure Design Documents
|
|
||||||
|
|
||||||
- [Prometheus Metrics](metrics.md)
|
|
||||||
|
|
||||||
## Module Design Documents
|
## Module Design Documents
|
||||||
|
|
||||||
- [AR Module](module/ar_module.md)
|
- [AR Module](module/ar_module.md)
|
||||||
|
|||||||
@@ -1,208 +0,0 @@
|
|||||||
# Prometheus Metrics Design
|
|
||||||
|
|
||||||
This document describes how vLLM-Omni exposes Prometheus metrics for multi-stage pipelines, the constraints that shaped the design, and how the pipeline-level metrics coexist with upstream vLLM per-engine metrics. This iteration covers the text and audio surface; diffusion / image / video families land in a follow-up.
|
|
||||||
|
|
||||||
## Objectives
|
|
||||||
|
|
||||||
- Expose pipeline-level request and latency metrics that span the full multi-stage execution (orchestrator scope).
|
|
||||||
- Preserve all upstream vLLM per-engine metrics (`vllm:*`) for stages backed by an AR LLM engine, and reshape their `engine` label into `stage` + `replica` so multi-replica deployments gain per-replica visibility automatically.
|
|
||||||
- Expose per-modality SLO metrics that the upstream `vllm:*` families do not capture — audio TTFP / RTF / duration / frames / streaming continuity / silent-loss.
|
|
||||||
- Expose per-replica-edge cross-stage transfer metrics so the slack between E2E latency and the sum of per-stage `gen_time` (queueing, serialization, network) becomes attributable.
|
|
||||||
- Keep the metrics collection overhead low enough that it does not regress TTFA or throughput.
|
|
||||||
|
|
||||||
## Background
|
|
||||||
|
|
||||||
### Upstream vLLM Metrics
|
|
||||||
|
|
||||||
Upstream vLLM defines ~37 Prometheus metric families under the `vllm:` prefix. These are registered by `PrometheusStatLogger` and cover engine-level state: KV cache usage, running/waiting request counts, token throughput, TTFT, inter-token latency, e2e latency, and so on. They are served via the `/metrics` HTTP endpoint provided by `prometheus_fastapi_instrumentator` and the default `prometheus_client` WSGI handler.
|
|
||||||
|
|
||||||
vLLM's `unregister_vllm_metrics()` function strips every `prometheus_client` collector whose `_name` attribute contains the substring `"vllm"`. This runs during engine initialization to clean up stale collectors from prior instantiations within the same process.
|
|
||||||
|
|
||||||
### The Problem
|
|
||||||
|
|
||||||
vLLM-Omni runs multiple engine instances (stages × replicas) within a single process, coordinated by an Orchestrator. The pipeline needs its own metrics — aggregate request counts, end-to-end latency across all stages, per-modality SLO signals, and cross-stage transfer attribution — that do not exist in upstream vLLM. All pipeline-level metrics use the `vllm:omni_` prefix to distinguish them from upstream per-engine metrics. At import time (see `vllm_omni/patch.py`) we replace `unregister_vllm_metrics()` with a scoped version that still strips upstream `vllm:*` collectors before each new `PrometheusStatLogger` registers (so multi-engine processes don't crash on duplicate registration), but preserves anything prefixed `vllm:omni_` / `vllm_omni`.
|
|
||||||
|
|
||||||
Upstream per-engine metrics retain the `vllm:` prefix but are now registered by `OmniPrometheusStatLogger`, a thin subclass of upstream's `PrometheusStatLogger` that reshapes the single `engine` label into a `stage` + `replica` pair (see "OmniPrometheusStatLogger wrap" below).
|
|
||||||
|
|
||||||
## Architecture
|
|
||||||
|
|
||||||
### Component Overview
|
|
||||||
|
|
||||||
```text
|
|
||||||
+------------------------+
|
|
||||||
| API Server (FastAPI) |
|
|
||||||
| GET /metrics |
|
|
||||||
+-----------+------------+
|
|
||||||
|
|
|
||||||
prometheus_client default registry
|
|
||||||
|
|
|
||||||
+--------+--------+--------+--------+--------+
|
|
||||||
| |
|
|
||||||
vllm:omni_* vllm:*
|
|
||||||
collectors collectors
|
|
||||||
| |
|
|
||||||
+----+--------+ +-----------+ +----------+ +-----------+
|
|
||||||
| OmniPromet- | | OmniMod- | | OmniTra- | | OmniProm- |
|
|
||||||
| heusMetrics | | alityMet- | | nsferMe- | | etheusSt- |
|
|
||||||
| | | rics | | trics | | atLogger |
|
|
||||||
+----+--------+ +-----+-----+ +----+-----+ +----+------+
|
|
||||||
| | | |
|
|
||||||
OmniBase OmniBase Orchestrator Orchestrator
|
|
||||||
(request life- (finalize + (record_trans- (per-(stage,
|
|
||||||
cycle, success/ streaming fer_tx/rx replica)
|
|
||||||
fail counter) hooks via hooks via scheduler/
|
|
||||||
observe_* emit hook in iteration
|
|
||||||
APIs) OrchestratorAg- stats)
|
|
||||||
gregator)
|
|
||||||
```
|
|
||||||
|
|
||||||
### Data Flow
|
|
||||||
|
|
||||||
There are four independent paths for metric collection.
|
|
||||||
|
|
||||||
**Path 1: Pipeline-level metrics (`vllm:omni_*`)**
|
|
||||||
|
|
||||||
`OmniPrometheusMetrics` registers the Gauge / Counter / Histogram collectors at import time. It is instantiated once per entrypoint, labeled with the model name. The entrypoint calls its methods as requests progress:
|
|
||||||
|
|
||||||
- `set_running(n)` / `set_waiting(n)` — updated after each request completes. The running count comes from `OmniRequestCounter`, a simple counter incremented/decremented by the Orchestrator as it tracks requests. Waiting is derived as `total - running`.
|
|
||||||
- `request_succeeded(e2e_seconds, finished_reason="stop")` — recorded when a request finishes at the final stage. `finished_reason` is extracted from `engine_outputs.outputs[0].finish_reason` (vLLM `CompletionOutput` convention) and increments `vllm:omni_requests_success_total{finished_reason}`.
|
|
||||||
- `request_failed()` — recorded by the cleanup path when a request exits without natural completion. Internally maps to `finished_reason="abort"` so a single Counter family covers both natural and aborted completion.
|
|
||||||
|
|
||||||
**Path 2: Audio modality metrics (`vllm:omni_audio_*`)**
|
|
||||||
|
|
||||||
`OmniModalityMetrics` registers seven audio families with `{model_name, stage, replica}` (plus an extra `threshold_ms` / `reason` label on the two extra-cardinality Counters). Three observation sites:
|
|
||||||
|
|
||||||
- `observe_modality_at_finalize(...)` — called from `omni_base._process_single_result` inside the existing `e2e_done` finalize guard. For `output_type == "audio"` it emits `audio_frames_total`, `audio_duration_s`, `audio_rtf` (or `audio_skipped_requests_total{reason="no_audio_data"}` when no audio was produced). Sample rate is resolved from `engine_outputs.multimodal_output` via `definitions.resolve_audio_sample_rate(...)` (fallback chain mirrors `serving_chat.py`'s audio response path).
|
|
||||||
- `observe_audio_first_packet(...)` — called from the OpenAI SSE audio branch in `serving_chat.py` on the first audio packet for a request. The once-per-request guard is held by `ClientRequestState.first_audio_ts`. The `request_arrival_ts` anchor is stored in `ClientRequestState` by `async_omni.generate()`, computed at request entry.
|
|
||||||
- `observe_audio_streaming_finalize(...)` — called from `serving_chat.py` after the streaming chunk loop exhausts. It runs the per-chunk player simulation from `vllm_omni/benchmarks/audio_continuity.py` to compute the worst-case underrun and emits `audio_underrun_s` plus (when the request stayed below the threshold) `audio_continuity_ok_total{threshold_ms}`. Per-chunk PCM byte counts and arrival timestamps are recorded by the same audio branch that updates `first_audio_ts`.
|
|
||||||
|
|
||||||
**Path 3: Cross-stage transfer metrics (`vllm:omni_transfer_*`)**
|
|
||||||
|
|
||||||
`OmniTransferMetrics` registers four Histogram families with `{model_name, from_stage, from_replica, to_stage, to_replica}` labels. Each observation corresponds to one physical transfer hop (one chunk between adjacent stages), not the per-request accumulated total — so the histograms track per-transfer distribution.
|
|
||||||
|
|
||||||
The hook lives in `OrchestratorAggregator.record_transfer_tx` and `record_transfer_rx`. After the existing `TransferEdgeStats` accumulation, the aggregator calls `_emit_transfer_tx` / `_emit_transfer_rx`. Those:
|
|
||||||
|
|
||||||
1. Resolve `from_replica` / `to_replica` via a `replica_resolver` callback supplied by `async_omni.py`. The resolver delegates to `stage_pool.get_bound_replica_id(request_id)` — i.e. the orchestrator's existing sticky-routing binding is the source of truth.
|
|
||||||
2. Convert the underlying `_ms` accumulators to seconds and call the `_s`-suffixed observe methods on `OmniTransferMetrics`.
|
|
||||||
|
|
||||||
Defensive fail-safe: if `transfer_emitter` or `replica_resolver` is missing, or the resolver returns `None` for either side, the emit is skipped silently (the underlying `TransferEdgeStats` accumulation is unaffected).
|
|
||||||
|
|
||||||
**Path 4: Per-engine metrics (`vllm:*`, stage/replica wrap)**
|
|
||||||
|
|
||||||
The Orchestrator instantiates `OmniPrometheusStatLogger` (a thin subclass of upstream `vllm.v1.metrics.loggers.PrometheusStatLogger`) and feeds it scheduler stats and iteration stats after processing each batch of engine outputs. This populates the standard ~37 vLLM metric families (TTFT, ITL, TPOT, KV cache usage, etc.) using the same upstream code path — but with the `engine` label reshaped into `stage` + `replica` so multi-replica deployments produce distinct series per replica. See the next section for the wrap mechanics.
|
|
||||||
|
|
||||||
### Shared State Between Threads
|
|
||||||
|
|
||||||
The Orchestrator runs in a background thread. The API server (OmniBase) runs in the asyncio event loop thread. `OmniRequestCounter` bridges them — a plain Python object with an `int` field. The Orchestrator increments/decrements it; the entrypoint reads it for gauge updates. No lock is needed because the counter is advisory (a stale read by one Prometheus scrape interval is acceptable). It is created by `AsyncOmniEngine.__init__()` and passed to the Orchestrator at construction time.
|
|
||||||
|
|
||||||
### Metric Registration and Lifecycle
|
|
||||||
|
|
||||||
All `vllm:omni_*` collectors are registered once when their owning class (`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics`) is imported. Per-`(stage, replica)` labels are bound lazily on first observation to avoid registering label sets for combinations that never produce data.
|
|
||||||
|
|
||||||
The `prometheus_client` default registry holds all collectors. FastAPI's `/metrics` endpoint serves the default registry, so `vllm:omni_*` and the wrapped `vllm:*` metrics appear in the same scrape response alongside `http_*` and `process_*` metrics from the instrumentator and the Python client runtime.
|
|
||||||
|
|
||||||
## OmniPrometheusStatLogger Wrap
|
|
||||||
|
|
||||||
Upstream `PrometheusStatLogger.__init__` hard-codes `labelnames = ["model_name", "engine"]` as a local variable, references it across ~37 metric-family construction sites, and uses the `engine` label value in five different `.labels()` call shapes (kwarg with int engine, kwarg with str engine, positional with str engine in the middle, plus a `metrics_info["engine"] = str(...)` dict pattern). To reshape `engine` into `stage` + `replica` without forking the entire upstream `__init__`, the wrap uses three coordinated mechanisms:
|
|
||||||
|
|
||||||
1. **Class-level metric class slot overrides.** `OmniPrometheusStatLogger` overrides `_gauge_cls`, `_counter_cls`, `_histogram_cls` (which upstream calls via `self._gauge_cls(...)` etc.) with `_RelabelGauge` / `_RelabelCounter` / `_RelabelHistogram` wrapper classes. These intercept the `labelnames` kwarg at metric family creation time and replace `engine` with `("stage", "replica")`.
|
|
||||||
2. **Property descriptor for `per_engine_labelvalues`.** Upstream builds `self.per_engine_labelvalues = {idx: [model_name, str(idx)]}` inside `__init__` and then captures it into a local variable for `create_metric_per_engine` calls. By making `per_engine_labelvalues` a Python property on the subclass, the setter intercepts upstream's assignment and rewrites each 2-tuple into a 3-tuple `[model_name, stage, replica]` using the `stage_replica_map` supplied at construction time. The captured local then sees the rewritten dict.
|
|
||||||
3. **Override of `.labels()` on the wrapper classes.** For the five call sites that pass `engine` directly (kwarg or positional, int or str), `_RelabelMixin.labels()` translates the engine value back to `(stage, replica)` via a process-level `_ENGINE_INDEX_MAP` populated by `OmniPrometheusStatLogger.__init__`. This handles `gauge_engine_sleep_state.labels(engine=idx, ...)`, `counter_request_success_base.labels(model_name, str(idx), str(reason))`, `info_gauge.labels(**metrics_info)`, etc.
|
|
||||||
|
|
||||||
The three sub-helpers that upstream `PrometheusStatLogger.__init__` constructs (`spec_decoding_prom` / `kv_connector_prom` / `perf_metrics_prom`) use their own `_counter_cls` / `_gauge_cls` / `_histogram_cls` slots and would otherwise build families with the raw 2-element labelnames. `_OmniPerfMetricsProm` / `_OmniSpecDecodingProm` / `_OmniKVConnectorProm` subclass each helper to route the same relabel mixin through their internal family construction.
|
|
||||||
|
|
||||||
The `Orchestrator` constructs `stage_replica_map` from the static `stage_pools` configuration at startup:
|
|
||||||
|
|
||||||
```python
|
|
||||||
stage_replica_map = {
|
|
||||||
flat_idx: (str(stage_id), str(replica_id))
|
|
||||||
for flat_idx, (stage_id, replica_id) in enumerate(
|
|
||||||
(s, r)
|
|
||||||
for s, pool in enumerate(stage_pools)
|
|
||||||
for r in range(pool.num_replicas)
|
|
||||||
)
|
|
||||||
}
|
|
||||||
```
|
|
||||||
|
|
||||||
A reverse map `(stage_id, replica_id) -> flat_idx` is maintained on the Orchestrator so the per-replica `record(engine_idx=...)` call site can look up the right flat index.
|
|
||||||
|
|
||||||
> Dynamic add/remove of replicas at runtime is intentionally out of scope — the upstream `PrometheusStatLogger` materializes per-engine_idx child metrics at init time, and supporting hot-add would require non-trivial intervention into upstream's per-family child dictionaries.
|
|
||||||
|
|
||||||
## Gating: `--log-stats`
|
|
||||||
|
|
||||||
All metrics — both the 15 `vllm:omni_*` families and the ~65 upstream `vllm:*` wrap families — are gated by the user's `--log-stats` CLI flag (default off). The flag is plumbed from `OmniBase.__init__(log_stats=...)` through `AsyncOmniEngine` to the stage-spawn helpers and to the three Prometheus metric classes (`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics`), and is forwarded to `Orchestrator._init_metrics_state(...)`.
|
|
||||||
|
|
||||||
Behavior with `--log-stats=off` (default):
|
|
||||||
|
|
||||||
- The three `Omni*Metrics` classes register their module-level Gauge / Counter / Histogram families at import time (prometheus_client requires up-front registration), but each `observe / inc / set` method early-returns. The per-label child series for `vllm:omni_*` stay materialized but never have data written to them.
|
|
||||||
- `OmniPrometheusStatLogger` is not constructed in `_init_metrics_state`, so the ~65 upstream `vllm:*` wrap families are not registered in the default registry at all.
|
|
||||||
- The engine core's `Scheduler.make_stats()` also short-circuits inside upstream (`if not self.log_stats: return None`), so no `SchedulerStats` is produced per step — the per-iteration cost is bounded by the existing upstream gate.
|
|
||||||
|
|
||||||
Behavior with `--log-stats=on`: all metric paths fire normally; the orchestrator's per-replica recording is bounded only by `OmniSchedulerMixin.make_stats()`'s per-scheduler 1 Hz throttle (see next section).
|
|
||||||
|
|
||||||
The overhead with the flag on is small enough that an A/B benchmark on Qwen3-Omni-30B single replica (30 sequential audio requests) showed a mean latency delta of +0.6% (Welch's t = 0.318, n=30, not statistically significant at α=0.05); the `/metrics` line count drops from 1358 to 124 lines when the flag is off.
|
|
||||||
|
|
||||||
## Throttling: `make_stats()` Override
|
|
||||||
|
|
||||||
Upstream vLLM's `Scheduler.make_stats()` runs on every AR generation step, returning a SchedulerStats object for the orchestrator. Under vLLM's architecture, this is fine. But since vLLM-Omni requires that the object be serialized and transferred over ZMQ, receiving a SchedulerStats object on every step can introduce unacceptable overhead to the system.
|
|
||||||
|
|
||||||
`OmniSchedulerMixin.make_stats()` (in `vllm_omni/core/sched/omni_scheduler_mixin.py`) throttles stats emission to at most once per second **per scheduler** — i.e. per `(stage, replica)` since each replica owns its own scheduler instance. Between intervals it returns `None`, which the engine core skips serializing. This keeps gauges fresh enough for Prometheus scrapes (typically 15-30s intervals) while eliminating the per-step overhead.
|
|
||||||
|
|
||||||
The orchestrator side does not add its own throttle on top: the per-replica recording loop gates only on `raw_outputs.scheduler_stats is not None` (i.e. this replica's scheduler passed its own 1 Hz gate). A previous global `_last_stats_ts` on the orchestrator was removed because it starved every replica other than the first to emit within each second.
|
|
||||||
|
|
||||||
## Metric Definitions
|
|
||||||
|
|
||||||
### Pipeline (4)
|
|
||||||
|
|
||||||
| Metric | Type | Labels | Description |
|
|
||||||
|--------|------|--------|-------------|
|
|
||||||
| `vllm:omni_num_requests_running` | Gauge | `model_name` | Requests currently executing across all stages |
|
|
||||||
| `vllm:omni_num_requests_waiting` | Gauge | `model_name` | Requests queued but not yet scheduled |
|
|
||||||
| `vllm:omni_requests_success_total` | Counter | `model_name`, `finished_reason` | Total requests by completion reason ({stop, length, abort, ...}); aborts cover client-disconnect / cancellation paths in addition to upstream `FinishReason.ABORT` |
|
|
||||||
| `vllm:omni_e2e_request_latency_s` | Histogram | `model_name` | Pipeline-global end-to-end latency in seconds |
|
|
||||||
|
|
||||||
### Audio (7)
|
|
||||||
|
|
||||||
Labels: `{model_name, stage, replica}` plus the listed extra label.
|
|
||||||
|
|
||||||
| Metric | Type | Extra label | Description |
|
|
||||||
|--------|------|-------------|-------------|
|
|
||||||
| `vllm:omni_audio_ttfp_s` | Histogram | — | Time from request arrival to first audio packet/frame |
|
|
||||||
| `vllm:omni_audio_duration_s` | Histogram | — | Audio content duration (`audio_frames / sample_rate`) |
|
|
||||||
| `vllm:omni_audio_rtf` | Histogram | — | Real-time factor `stage_gen_time_s / audio_duration_s` (SLO `< 1`); uses `RTF_BUCKETS` |
|
|
||||||
| `vllm:omni_audio_frames_total` | Counter | — | Cumulative audio frames generated |
|
|
||||||
| `vllm:omni_audio_underrun_s` | Histogram | — | Per-request worst-case player deficit; `> 0` indicates listener heard silent gaps |
|
|
||||||
| `vllm:omni_audio_continuity_ok_total` | Counter | `threshold_ms` | Incremented when the request's worst underrun stayed below `threshold_ms` |
|
|
||||||
| `vllm:omni_audio_skipped_requests_total` | Counter | `reason` | Silent-loss counter — code2wav rejected malformed codec input and returned `200 OK` with empty audio |
|
|
||||||
|
|
||||||
### Cross-stage transfer (4)
|
|
||||||
|
|
||||||
Labels: `{model_name, from_stage, from_replica, to_stage, to_replica}`.
|
|
||||||
|
|
||||||
| Metric | Type | Description |
|
|
||||||
|--------|------|-------------|
|
|
||||||
| `vllm:omni_transfer_size_bytes` | Histogram | Per-transfer payload size in bytes |
|
|
||||||
| `vllm:omni_transfer_tx_s` | Histogram | Sender-side time (serialize + submit to connector) |
|
|
||||||
| `vllm:omni_transfer_rx_s` | Histogram | Receiver-side time (recv + deserialize) |
|
|
||||||
| `vllm:omni_transfer_in_flight_s` | Histogram | Network in-flight time (TX done → RX recv start) |
|
|
||||||
|
|
||||||
### LLM stage-level (wrapped `vllm:*`)
|
|
||||||
|
|
||||||
After the wrap, every upstream `vllm:*` family — TTFT, ITL, TPOT, e2e latency, KV cache usage, scheduler running/waiting, request success counts, etc. — carries `{model_name, stage, replica}` labels. For the full upstream catalog see [the vLLM docs](https://github.com/vllm-project/vllm/blob/main/docs/usage/metrics.md); note that metrics depending on features unsupported in vLLM-Omni (e.g. speculative decoding, LoRA) will not be available.
|
|
||||||
|
|
||||||
## Naming Convention
|
|
||||||
|
|
||||||
- All time-bearing metrics use the `_s` suffix (values in seconds). Two bucket families are used:
|
|
||||||
- `SECONDS_BUCKETS` (0.05 s – 300 s) for e2e / generation / TTFP style values.
|
|
||||||
- `SECONDS_FAST_BUCKETS` (0.001 s – 60 s) for fine-grained cross-stage transfer and audio-underrun values that need millisecond-level resolution.
|
|
||||||
- Counters use the `_total` suffix (auto-appended by `prometheus_client`).
|
|
||||||
- Sizes use the `_bytes` suffix.
|
|
||||||
- All omni-specific families are prefixed `vllm:omni_`. The upstream `unregister_vllm_metrics()` function is monkey-patched to a scoped version that still strips upstream `vllm:*` collectors (so multi-engine init within one process does not crash on duplicate registration) but preserves anything prefixed `vllm:omni_` / `vllm_omni`.
|
|
||||||
|
|
||||||
## Logging vs. Prometheus
|
|
||||||
|
|
||||||
`OrchestratorAggregator` (in `vllm_omni/metrics/stats.py`) is the logging-oriented metrics path. It collects detailed per-request, per-stage, and per-transfer statistics and prints formatted tables to the `INFO` log. This is designed for development and debugging — individual request traces, transfer bandwidth, inter-stage timing.
|
|
||||||
|
|
||||||
`OmniPrometheusMetrics` / `OmniModalityMetrics` / `OmniTransferMetrics` form the Prometheus-oriented path. They record aggregate counters, gauges, and histograms suitable for time-series monitoring and alerting. Both paths share the same source data (`StageRequestStats`, `TransferEdgeStats`) — `OrchestratorAggregator.record_transfer_tx/rx` in particular calls both the existing accumulator code and the Prometheus emit hook in the same method body. The two consumption models can run simultaneously without coupling.
|
|
||||||
|
|
||||||
The separation follows upstream vLLM's pattern of `LoggingStatLogger` vs. `PrometheusStatLogger` — same underlying data, different consumption models.
|
|
||||||
@@ -411,8 +411,6 @@ Notes:
|
|||||||
- `runtime.max_batch_size` controls stage-level batching.
|
- `runtime.max_batch_size` controls stage-level batching.
|
||||||
- Thinker/Talker commonly use `enforce_eager: false` for CUDA Graph paths.
|
- Thinker/Talker commonly use `enforce_eager: false` for CUDA Graph paths.
|
||||||
- Code2Wav often remains eager (`enforce_eager: true`) depending on runtime behavior.
|
- Code2Wav often remains eager (`enforce_eager: true`) depending on runtime behavior.
|
||||||
- Qwen3-Omni defaults `VLLM_USE_FLASHINFER_MOE_FP16=0`. The Triton has been more stable & faster
|
|
||||||
than the FlashInfer CUTLASS unquantized MoE backend on recent vLLM rebases.
|
|
||||||
|
|
||||||
#### 2) Enable async chunk
|
#### 2) Enable async chunk
|
||||||
|
|
||||||
|
|||||||
@@ -34,7 +34,7 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
|
|||||||
|
|
||||||
### Pre-built wheels
|
### Pre-built wheels
|
||||||
|
|
||||||
Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, and 0.22.0. If you need a newer unreleased revision, 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 available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, and 0.21.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
|
||||||
|
|
||||||
=== "NVIDIA CUDA"
|
=== "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.
|
vLLM-Omni is built based on vLLM. Please install it with command below.
|
||||||
```bash
|
```bash
|
||||||
uv pip install vllm==0.22.0 --torch-backend=auto
|
uv pip install vllm==0.21.0 --torch-backend=auto
|
||||||
```
|
```
|
||||||
|
|
||||||
#### Installation of vLLM-Omni
|
#### Installation of vLLM-Omni
|
||||||
@@ -39,13 +39,13 @@ uv pip install 'vllm-omni[demo]'
|
|||||||
# --8<-- [start:build-wheel-from-source]
|
# --8<-- [start:build-wheel-from-source]
|
||||||
|
|
||||||
#### Installation of vLLM
|
#### Installation of vLLM
|
||||||
If you do not need to modify source code of vLLM, you can directly install the stable 0.22.0 release version of the library
|
If you do not need to modify source code of vLLM, you can directly install the stable 0.21.0 release version of the library
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
uv pip install vllm==0.22.0 --torch-backend=auto
|
uv pip install vllm==0.21.0 --torch-backend=auto
|
||||||
```
|
```
|
||||||
|
|
||||||
The 0.22.0 release of vLLM ships CUDA 12.9-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
|
The 0.21.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
|
||||||
|
|
||||||
#### Installation of vLLM-Omni
|
#### Installation of vLLM-Omni
|
||||||
Since vllm-omni is rapidly evolving, it's recommended to install it from source
|
Since vllm-omni is rapidly evolving, it's recommended to install it from source
|
||||||
@@ -66,12 +66,12 @@ If you want to check, modify or debug with source code of vLLM, install the libr
|
|||||||
```bash
|
```bash
|
||||||
git clone https://github.com/vllm-project/vllm.git
|
git clone https://github.com/vllm-project/vllm.git
|
||||||
cd vllm
|
cd vllm
|
||||||
git checkout v0.22.0
|
git checkout v0.21.0
|
||||||
```
|
```
|
||||||
Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file.
|
Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file.
|
||||||
```bash
|
```bash
|
||||||
#For CUDA 12.9 (the default for v0.22.0; the wheel filename has no `+cu129` suffix)
|
#For CUDA 13.0 (the default for v0.21.0; the wheel filename has no `+cu130` suffix)
|
||||||
export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.22.0/vllm-0.22.0-cp38-abi3-manylinux_2_28_x86_64.whl
|
export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.21.0/vllm-0.21.0-cp38-abi3-manylinux_2_35_x86_64.whl
|
||||||
```
|
```
|
||||||
Install vllm with command below (If you have no existing PyTorch).
|
Install vllm with command below (If you have no existing PyTorch).
|
||||||
```bash
|
```bash
|
||||||
@@ -102,7 +102,7 @@ docker run --runtime nvidia --gpus 2 \
|
|||||||
--env "HF_TOKEN=$HF_TOKEN" \
|
--env "HF_TOKEN=$HF_TOKEN" \
|
||||||
-p 8091:8091 \
|
-p 8091:8091 \
|
||||||
--ipc=host \
|
--ipc=host \
|
||||||
vllm/vllm-omni:v0.22.0 \
|
vllm/vllm-omni:v0.21.0 \
|
||||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
|
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --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.
|
vLLM-Omni is built based on vLLM. Please install it with command below.
|
||||||
```bash
|
```bash
|
||||||
uv pip install vllm==0.22.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.22.0/rocm722
|
uv pip install vllm==0.21.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.21.0/rocm722
|
||||||
```
|
```
|
||||||
|
|
||||||
#### Installation of vLLM-Omni
|
#### Installation of vLLM-Omni
|
||||||
@@ -34,13 +34,13 @@ uv pip install onnxruntime-rocm
|
|||||||
# --8<-- [start:build-wheel-from-source]
|
# --8<-- [start:build-wheel-from-source]
|
||||||
|
|
||||||
#### Installation of vLLM
|
#### Installation of vLLM
|
||||||
If you do not need to modify source code of vLLM, you can directly install the stable 0.22.0 release version of the library
|
If you do not need to modify source code of vLLM, you can directly install the stable 0.21.0 release version of the library
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
uv pip install vllm==0.22.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.22.0/rocm722
|
uv pip install vllm==0.21.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.21.0/rocm722
|
||||||
```
|
```
|
||||||
|
|
||||||
The pre-built 0.22.0 vLLM wheel targets ROCm 7.2.2. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
|
The pre-built 0.21.0 vLLM wheel targets ROCm 7.2.2. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
|
||||||
|
|
||||||
#### Installation of vLLM-Omni
|
#### Installation of vLLM-Omni
|
||||||
Since vllm-omni is rapidly evolving, it's recommended to install it from source
|
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
|
```bash
|
||||||
git clone https://github.com/vllm-project/vllm.git
|
git clone https://github.com/vllm-project/vllm.git
|
||||||
cd vllm
|
cd vllm
|
||||||
git checkout v0.22.0
|
git checkout v0.21.0
|
||||||
python3 -m pip install -r requirements/rocm.txt
|
python3 -m pip install -r requirements/rocm.txt
|
||||||
python3 setup.py develop
|
python3 setup.py develop
|
||||||
```
|
```
|
||||||
@@ -130,7 +130,7 @@ docker run --rm \
|
|||||||
-v ~/.cache/huggingface:/root/.cache/huggingface \
|
-v ~/.cache/huggingface:/root/.cache/huggingface \
|
||||||
--env "HF_TOKEN=$HF_TOKEN" \
|
--env "HF_TOKEN=$HF_TOKEN" \
|
||||||
-p 8091:8091 \
|
-p 8091:8091 \
|
||||||
vllm/vllm-omni-rocm:v0.22.0 \
|
vllm/vllm-omni-rocm:v0.21.0 \
|
||||||
--model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
|
--model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
|
||||||
```
|
```
|
||||||
|
|
||||||
@@ -149,7 +149,7 @@ docker run --rm -it \
|
|||||||
-v ~/.cache/huggingface:/root/.cache/huggingface \
|
-v ~/.cache/huggingface:/root/.cache/huggingface \
|
||||||
--env "HF_TOKEN=$HF_TOKEN" \
|
--env "HF_TOKEN=$HF_TOKEN" \
|
||||||
--entrypoint bash \
|
--entrypoint bash \
|
||||||
vllm/vllm-omni-rocm:v0.22.0
|
vllm/vllm-omni-rocm:v0.21.0
|
||||||
```
|
```
|
||||||
|
|
||||||
# --8<-- [end:pre-built-images]
|
# --8<-- [end:pre-built-images]
|
||||||
@@ -19,10 +19,10 @@ uv venv --python 3.12 --seed
|
|||||||
source .venv/bin/activate
|
source .venv/bin/activate
|
||||||
|
|
||||||
# On CUDA
|
# On CUDA
|
||||||
uv pip install vllm==0.22.0 --torch-backend=auto
|
uv pip install vllm==0.21.0 --torch-backend=auto
|
||||||
|
|
||||||
# On ROCm
|
# On ROCm
|
||||||
uv pip install vllm==0.22.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.22.0/rocm722
|
uv pip install vllm==0.21.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.21.0/rocm721
|
||||||
|
|
||||||
git clone https://github.com/vllm-project/vllm-omni.git
|
git clone https://github.com/vllm-project/vllm-omni.git
|
||||||
cd vllm-omni
|
cd vllm-omni
|
||||||
@@ -35,7 +35,7 @@ For additional installation methods — please see the [installation guide](inst
|
|||||||
!!! note
|
!!! note
|
||||||
It is important to install the same major & minor version of vLLM and vLLM Omni, otherwise things may not work as expected. If the versions are misaligned, you will see a warning when you import vLLM Omni.
|
It is important to install the same major & minor version of vLLM and vLLM Omni, otherwise things may not work as expected. If the versions are misaligned, you will see a warning when you import vLLM Omni.
|
||||||
|
|
||||||
If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.22.0` and vLLM Omni `0.22.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.
|
If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.21.0` and vLLM Omni `0.21.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.
|
||||||
|
|
||||||
## Offline Inference
|
## Offline Inference
|
||||||
|
|
||||||
|
|||||||
@@ -119,10 +119,9 @@ def extract_omni_serve_subparser_init():
|
|||||||
exec_globals = {
|
exec_globals = {
|
||||||
"_FlexibleArgumentParser": _FlexibleArgumentParser,
|
"_FlexibleArgumentParser": _FlexibleArgumentParser,
|
||||||
"FlexibleArgumentParser": _FlexibleArgumentParser,
|
"FlexibleArgumentParser": _FlexibleArgumentParser,
|
||||||
"TrackingArgumentParser": _FlexibleArgumentParser, # only needs add_argument, not tracking
|
|
||||||
"TrackingNamespace": None, # type annotation only
|
|
||||||
"make_arg_parser": lambda parser: parser, # no-op for doc
|
"make_arg_parser": lambda parser: parser, # no-op for doc
|
||||||
"_ensure_vllm_platform": lambda: None, # no-op for doc
|
"_ensure_vllm_platform": lambda: None, # no-op for doc
|
||||||
|
"nullify_stage_engine_defaults": lambda parser: None, # no-op for doc
|
||||||
"VLLM_SUBCMD_PARSER_EPILOG": "",
|
"VLLM_SUBCMD_PARSER_EPILOG": "",
|
||||||
"logger": logger,
|
"logger": logger,
|
||||||
"DummySubparsers": DummySubparsers,
|
"DummySubparsers": DummySubparsers,
|
||||||
|
|||||||
@@ -33,7 +33,6 @@ th {
|
|||||||
| `ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
| `ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||||
| `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
| `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||||
| `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
| `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||||
| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano` | ✅︎ | | | |
|
|
||||||
| `WanSpeechToVideoPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
| `WanSpeechToVideoPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||||
| `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
| `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||||
| `LTX2Pipeline` | LTX-2-T2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
|
| `LTX2Pipeline` | LTX-2-T2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
|
||||||
@@ -42,7 +41,6 @@ th {
|
|||||||
| `LTX2ImageToVideoTwoStagesPipeline` | LTX-2-I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
|
| `LTX2ImageToVideoTwoStagesPipeline` | LTX-2-I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
|
||||||
| `LTX23Pipeline` | LTX-2.3-T2V | `dg845/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
|
| `LTX23Pipeline` | LTX-2.3-T2V | `dg845/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
|
||||||
| `LTX23ImageToVideoPipeline` | LTX-2.3-I2V | `dg845/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
|
| `LTX23ImageToVideoPipeline` | LTX-2.3-I2V | `dg845/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
|
||||||
| `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | | | |
|
|
||||||
| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | |
|
| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | |
|
||||||
| `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | |
|
| `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | |
|
||||||
| `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ |
|
| `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ |
|
||||||
@@ -50,8 +48,6 @@ th {
|
|||||||
| `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
| `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||||
| `StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | ✅︎ | ✅︎ | | ✅︎ |
|
| `StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | ✅︎ | ✅︎ | | ✅︎ |
|
||||||
| `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ |
|
| `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ |
|
||||||
| `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | |
|
|
||||||
| `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | |
|
|
||||||
| `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | |
|
| `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | |
|
||||||
| `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
| `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||||
| `FluxKontextPipeline` | FLUX.1-Kontext-dev | `black-forest-labs/FLUX.1-Kontext-dev` | ✅︎ | ✅︎ | | |
|
| `FluxKontextPipeline` | FLUX.1-Kontext-dev | `black-forest-labs/FLUX.1-Kontext-dev` | ✅︎ | ✅︎ | | |
|
||||||
@@ -64,10 +60,6 @@ th {
|
|||||||
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
| `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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||||
| `GLMTTSForConditionalGeneration` | GLM-TTS | `zai-org/GLM-TTS` | ✅︎ | | | |
|
| `GLMTTSForConditionalGeneration` | GLM-TTS | `zai-org/GLM-TTS` | ✅︎ | | | |
|
||||||
| `MossTTSNanoForCausalLM` | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✅︎ | | | |
|
|
||||||
| `MossTTSDelayModel` | MOSS-TTS, MOSS-TTSD, MOSS-SoundEffect, MOSS-VoiceGenerator | `OpenMOSS-Team/MOSS-VoiceGenerator` | ✅︎ | | | |
|
|
||||||
| `MossTTSRealtime` | MOSS-TTS-Realtime | `OpenMOSS-Team/MOSS-TTS-Realtime` | ✅︎ | | | |
|
|
||||||
| `HiggsAudioV2ForConditionalGeneration` | Higgs-Audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` | ✅︎ | | | |
|
|
||||||
| `NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` | ✅︎ | ✅︎ | | ✅︎ |
|
| `NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` | ✅︎ | ✅︎ | | ✅︎ |
|
||||||
| `MiMoAudioModel` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` | ✅︎ | ✅︎ | | |
|
| `MiMoAudioModel` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` | ✅︎ | ✅︎ | | |
|
||||||
| `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | |
|
| `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | |
|
||||||
@@ -75,13 +67,11 @@ th {
|
|||||||
| `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | |
|
| `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | |
|
||||||
| `DreamIDOmniPipeline` | DreamID-Omni | `XuGuo699/DreamID-Omni` | ✅︎ | ✅︎ | | |
|
| `DreamIDOmniPipeline` | DreamID-Omni | `XuGuo699/DreamID-Omni` | ✅︎ | ✅︎ | | |
|
||||||
| `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only) | `SenseNova/SenseNova-U1-8B-MoT` | ✅︎ | | | |
|
| `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only) | `SenseNova/SenseNova-U1-8B-MoT` | ✅︎ | | | |
|
||||||
| `LancePipeline` | Lance | `bytedance-research/Lance` | ✅︎ | | | |
|
|
||||||
| `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | |
|
| `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | |
|
||||||
| `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | |
|
| `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | |
|
||||||
| `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | |
|
| `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | |
|
||||||
| `CovoAudioForConditionalGeneration` | Covo-Audio-Chat | `tencent/Covo-Audio-Chat` | ✅︎ | | | |
|
| `CovoAudioForConditionalGeneration` | Covo-Audio-Chat | `tencent/Covo-Audio-Chat` | ✅︎ | | | |
|
||||||
|`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | |
|
|`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | |
|
||||||
| `MiniCPMO45OmniForConditionalGeneration` | MiniCPM-o 4.5 | `openbmb/MiniCPM-o-4_5` | ✅︎ | | ✅︎ | |
|
|
||||||
| `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
| `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||||
|`HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | |
|
|`HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | |
|
||||||
|
|
||||||
|
|||||||
@@ -381,44 +381,6 @@ are cached in-process with a shared LRU so repeated requests with the same
|
|||||||
all TTS model types; deleting a voice invalidates every model-type slot at
|
all TTS model types; deleting a voice invalidates every model-type slot at
|
||||||
once.
|
once.
|
||||||
|
|
||||||
### Precomputed Custom Voices
|
|
||||||
|
|
||||||
Qwen3-TTS Base and VoxCPM2 can load offline-precomputed voices at startup.
|
|
||||||
Generate a directory containing `custom_voice_manifest.json` plus one
|
|
||||||
`.safetensors` file per voice, then set the pipeline-wide deploy config field:
|
|
||||||
|
|
||||||
```yaml
|
|
||||||
custom_voice_dir: /path/to/custom_voices
|
|
||||||
```
|
|
||||||
|
|
||||||
Qwen3-TTS profiles are created with:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
python examples/online_serving/text_to_speech/qwen3_tts/precompute_custom_voice.py \
|
|
||||||
--model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
|
|
||||||
--voice-name alice \
|
|
||||||
--ref-audio /path/to/reference.wav \
|
|
||||||
--ref-text "Original transcript of the reference audio" \
|
|
||||||
--mode icl \
|
|
||||||
--output-dir /path/to/custom_voices
|
|
||||||
```
|
|
||||||
|
|
||||||
VoxCPM2 profiles are created with:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
python examples/online_serving/text_to_speech/voxcpm2/precompute_custom_voice.py \
|
|
||||||
--model openbmb/VoxCPM2 \
|
|
||||||
--voice-name alice \
|
|
||||||
--ref-audio /path/to/reference.wav \
|
|
||||||
--mode ref_continuation \
|
|
||||||
--prompt-text "Original transcript of the reference audio" \
|
|
||||||
--output-dir /path/to/custom_voices
|
|
||||||
```
|
|
||||||
|
|
||||||
Only profiles whose safetensors payload can be loaded and validated are exposed
|
|
||||||
by `GET /v1/audio/voices`. Valid precomputed voices can be used in
|
|
||||||
`POST /v1/audio/speech` by passing `voice="alice"` without `ref_audio`.
|
|
||||||
|
|
||||||
**Configuration (environment variables):**
|
**Configuration (environment variables):**
|
||||||
|
|
||||||
| Variable | Default | Description |
|
| Variable | Default | Description |
|
||||||
|
|||||||
@@ -1,98 +0,0 @@
|
|||||||
# Production Metrics
|
|
||||||
|
|
||||||
vLLM-Omni exposes Prometheus metrics via the `/metrics` endpoint on the OpenAI-compatible API server. This page covers the text and audio surface; diffusion / image / video metrics are tracked in a follow-up PR.
|
|
||||||
|
|
||||||
```bash
|
|
||||||
vllm-omni serve Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8000 --log-stats
|
|
||||||
curl http://localhost:8000/metrics
|
|
||||||
```
|
|
||||||
|
|
||||||
**`--log-stats` is required to populate metric data.** Without the flag, the endpoint still returns `200 OK` but the upstream `vllm:*` wrap is not registered at all, and the 15 `vllm:omni_*` families are registered as placeholders with no data written to them. This keeps the runtime cost essentially zero for deployments that don't need monitoring. With the flag, all ~80 families populate.
|
|
||||||
|
|
||||||
## Metric Namespaces
|
|
||||||
|
|
||||||
| Prefix | Source | Present when |
|
|
||||||
|--------|--------|--------------|
|
|
||||||
| `vllm:omni_` | vLLM-Omni orchestrator / audio modality / cross-stage transfer | Pipeline-dependent |
|
|
||||||
| `vllm:` | Upstream vLLM engine, wrapped by `OmniPrometheusStatLogger` to expose `{stage, replica}` | Pipeline includes an LLM (AR) stage |
|
|
||||||
| `http_` / `process_` | Uvicorn / Python runtime | Always |
|
|
||||||
|
|
||||||
## Pipeline-Level Metrics (`vllm:omni_`)
|
|
||||||
|
|
||||||
Defined in `vllm_omni/metrics/prometheus.py`. Track request lifecycle across the full multi-stage pipeline.
|
|
||||||
|
|
||||||
### Request counts
|
|
||||||
|
|
||||||
| Metric | Type | Labels | Description |
|
|
||||||
|--------|------|--------|-------------|
|
|
||||||
| `vllm:omni_num_requests_running` | Gauge | `model_name` | Pipeline-global in-flight requests (dispatched to engine, not yet finalized) |
|
|
||||||
| `vllm:omni_num_requests_waiting` | Gauge | `model_name` | Requests waiting in the Orchestrator queue |
|
|
||||||
| `vllm:omni_requests_success_total` | Counter | `model_name`, `finished_reason` | Total requests by completion reason. `finished_reason` ∈ {`stop`, `length`, `abort`, ...} mirroring upstream `vllm:request_success_total`; aborts cover client disconnect / cancellation paths in addition to upstream `FinishReason.ABORT` |
|
|
||||||
|
|
||||||
### Latency
|
|
||||||
|
|
||||||
| Metric | Type | Labels | Description |
|
|
||||||
|--------|------|--------|-------------|
|
|
||||||
| `vllm:omni_e2e_request_latency_s` | Histogram | `model_name` | Pipeline-global end-to-end request latency in seconds |
|
|
||||||
|
|
||||||
## Audio Modality Metrics (`vllm:omni_`)
|
|
||||||
|
|
||||||
Emitted at request finalize, except for `audio_ttfp_s` (streaming-hook at the first audio packet) and `audio_underrun_s` / `audio_continuity_ok_total` (streaming finalize, after the chunk stream is exhausted). All carry `{model_name, stage, replica}` plus the listed extra label.
|
|
||||||
|
|
||||||
| Metric | Type | Extra label | Description |
|
|
||||||
|--------|------|-------------|-------------|
|
|
||||||
| `vllm:omni_audio_ttfp_s` | Histogram | — | Time from request arrival to first audio packet/frame |
|
|
||||||
| `vllm:omni_audio_duration_s` | Histogram | — | Audio content duration (`audio_frames / sample_rate`) |
|
|
||||||
| `vllm:omni_audio_rtf` | Histogram | — | Real-time factor (`stage_gen_time_s / audio_duration_s`); streaming TTS SLO red line `< 1`; uses `RTF_BUCKETS` |
|
|
||||||
| `vllm:omni_audio_frames_total` | Counter | — | Cumulative audio frame count; throughput via `rate()` |
|
|
||||||
| `vllm:omni_audio_underrun_s` | Histogram | — | Per-request worst-case player deficit; `> 0` indicates listener heard silent gaps |
|
|
||||||
| `vllm:omni_audio_continuity_ok_total` | Counter | `threshold_ms` | Incremented when the request's worst underrun stayed below `threshold_ms` |
|
|
||||||
| `vllm:omni_audio_skipped_requests_total` | Counter | `reason` | Silent-loss counter — code2wav rejected malformed codec input and returned `200 OK` with empty audio |
|
|
||||||
|
|
||||||
The continuity math comes from `vllm_omni/benchmarks/audio_continuity.py::compute_continuity_stats` so the server-side observation aligns with the bench-side definition.
|
|
||||||
|
|
||||||
## Cross-Stage Transfer Metrics (`vllm:omni_`)
|
|
||||||
|
|
||||||
Per-physical-transfer histograms tracking the data hop between adjacent stages. Labels `{model_name, from_stage, from_replica, to_stage, to_replica}` let dashboards attribute latency to specific replica edges. `from_replica` / `to_replica` are resolved from the orchestrator's sticky-routing binding (`stage_pool.get_bound_replica_id(request_id)`), so no extra plumbing through `TransferEdgeStats` is needed.
|
|
||||||
|
|
||||||
| Metric | Type | Description |
|
|
||||||
|--------|------|-------------|
|
|
||||||
| `vllm:omni_transfer_size_bytes` | Histogram | Per-transfer payload size in bytes |
|
|
||||||
| `vllm:omni_transfer_tx_s` | Histogram | Sender-side time (serialize + submit to connector) |
|
|
||||||
| `vllm:omni_transfer_rx_s` | Histogram | Receiver-side time (recv + deserialize) |
|
|
||||||
| `vllm:omni_transfer_in_flight_s` | Histogram | Network in-flight time (TX done → RX recv start) |
|
|
||||||
|
|
||||||
## vLLM Engine Metrics (`vllm:`)
|
|
||||||
|
|
||||||
When the pipeline includes an LLM stage, the upstream vLLM engine exposes its full set of ~37 metric families under the `vllm:` prefix.
|
|
||||||
|
|
||||||
vLLM-Omni wraps the upstream `vllm.v1.metrics.loggers.PrometheusStatLogger` with `OmniPrometheusStatLogger` so that the original `engine` single label is reshaped into `stage` + `replica`. Every `vllm:*` family — TTFT, ITL, TPOT, e2e latency, KV cache usage, scheduler running/waiting, request success counts, etc. — therefore gains per-`(stage, replica)` visibility automatically. No omni-side duplicate is needed for the text path.
|
|
||||||
|
|
||||||
```text
|
|
||||||
# Before wrap:
|
|
||||||
vllm:num_requests_running{model_name="...", engine="1"} 3.0
|
|
||||||
|
|
||||||
# After wrap:
|
|
||||||
vllm:num_requests_running{model_name="...", stage="1", replica="0"} 2.0
|
|
||||||
vllm:num_requests_running{model_name="...", stage="1", replica="1"} 1.0
|
|
||||||
```
|
|
||||||
|
|
||||||
For the full list of upstream metrics, see [the vLLM docs](https://github.com/vllm-project/vllm/blob/main/docs/usage/metrics.md).
|
|
||||||
|
|
||||||
## Metric Availability by Pipeline Type
|
|
||||||
|
|
||||||
| Metric group | Multi-stage LLM (Qwen3-Omni) |
|
|
||||||
|---|---|
|
|
||||||
| `vllm:omni_` request tracking + latency | With `--log-stats` |
|
|
||||||
| `vllm:omni_` audio modality | With `--log-stats`, if pipeline has a talker stage |
|
|
||||||
| `vllm:omni_` transfer | With `--log-stats`, if pipeline has ≥ 2 stages |
|
|
||||||
| `vllm:` engine metrics (per `(stage, replica)`) | With `--log-stats` |
|
|
||||||
| `vllm:` MFU metrics | With `--log-stats --enable-mfu-metrics` |
|
|
||||||
|
|
||||||
## Naming Convention
|
|
||||||
|
|
||||||
- All time-bearing metrics use the `_s` suffix (values in seconds). Buckets are `SECONDS_BUCKETS` for e2e / generation-style values and `SECONDS_FAST_BUCKETS` (1 ms → 60 s) for the fine-grained transfer and audio-underrun values.
|
|
||||||
- Counters use the `_total` suffix (auto-appended by `prometheus_client`).
|
|
||||||
- Sizes use the `_bytes` suffix.
|
|
||||||
- All omni-specific families are prefixed `vllm:omni_`. The upstream `unregister_vllm_metrics()` function is monkey-patched (see `vllm_omni/patch.py`) to a scoped version that still strips upstream `vllm:*` collectors so multi-engine init within one process does not crash on duplicate registration, but preserves anything prefixed `vllm:omni_` / `vllm_omni`.
|
|
||||||
- Text and audio first-output use distinct families (`vllm:time_to_first_token_seconds` reused from upstream for text; `vllm:omni_audio_ttfp_s` for audio) rather than a single metric with a `modality` label.
|
|
||||||
@@ -133,27 +133,22 @@ The following tables show which models support each feature:
|
|||||||
| **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ |
|
| **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||||
| **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ❌ | ✅ | ❌ | ✅ (decode) | ✅ | ❌ |
|
| **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ❌ | ✅ | ❌ | ✅ (decode) | ✅ | ❌ |
|
||||||
| **ERNIE-Image** | ❌ | ✅ | ✅ | ❓ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
| **ERNIE-Image** | ❌ | ✅ | ✅ | ❓ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||||
| **Cosmos3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ |
|
|
||||||
|
|
||||||
> Notes:
|
> Notes:
|
||||||
> 1. Nextstep_1(T2I) does not support cache acceleration methods such as TeaCache or Cache-DiT.
|
> 1. Nextstep_1(T2I) does not support cache acceleration methods such as TeaCache or Cache-DiT.
|
||||||
> 2. `Tongyi-MAI/Z-Image-Turbo` and `SII-GAIR/daVinci-MagiHuman-Base-1080p` are distilled models with minimal NFEs; CFG-Parallel is not necessary.
|
> 2. `Tongyi-MAI/Z-Image-Turbo` and `SII-GAIR/daVinci-MagiHuman-Base-1080p` are distilled models with minimal NFEs; CFG-Parallel is not necessary.
|
||||||
> 3. Cosmos3 T2I uses `Cosmos3OmniDiffusersPipeline` with `modalities=["image"]`. Model-level CPU offload is not supported; use layerwise offload.
|
|
||||||
|
|
||||||
### VideoGen
|
### VideoGen
|
||||||
|
|
||||||
| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | Pipeline-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
|
| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | Pipeline-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
|
||||||
|------------------------------|:---------:|:----------:|:---------------------:|:--------------:|:-----------------:|:-----------------:|:------:|:-------------------------:|:--------------------:|:--------------:|:----------------:|
|
|------------------------------|:---------:|:----------:|:---------------------:|:--------------:|:-----------------:|:-----------------:|:------:|:-------------------------:|:--------------------:|:--------------:|:----------------:|
|
||||||
| **Wan2.2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (encode/decode) | ❌ | ❌ |
|
| **Wan2.2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (encode/decode) | ❌ | ❌ |
|
||||||
| **Wan2.2-S2V** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ (encode/decode) | ❌ | ❌ |
|
|
||||||
| **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
|
| **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||||
| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||||
| **LTX-2.3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
| **LTX-2.3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||||
| **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
| **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||||
| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ |
|
| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ |
|
||||||
| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||||
| **Cosmos3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ |
|
|
||||||
|
|
||||||
|
|
||||||
**Frame Interpolation Support**
|
**Frame Interpolation Support**
|
||||||
|
|
||||||
|
|||||||
@@ -17,9 +17,6 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
|
|||||||
|
|
||||||
The default deployment configuration situated at `vllm_omni/deploy/qwen3_omni_moe.yaml` is resolved and loaded
|
The default deployment configuration situated at `vllm_omni/deploy/qwen3_omni_moe.yaml` is resolved and loaded
|
||||||
automatically via the model registry, obviating the necessity for the `--deploy-config` flag in standard deployment topologies.
|
automatically via the model registry, obviating the necessity for the `--deploy-config` flag in standard deployment topologies.
|
||||||
The bundled Qwen3-Omni setup defaults `VLLM_USE_FLASHINFER_MOE_FP16=0`. This keeps the Thinker & Talker on vLLM's
|
|
||||||
Triton unquantized MoE path and avoids the performance regression observed with the FlashInfer CUTLASS unquantized MoE
|
|
||||||
backend.
|
|
||||||
Asynchronous chunk streaming is **enabled by default** within the bundled configuration.
|
Asynchronous chunk streaming is **enabled by default** within the bundled configuration.
|
||||||
|
|
||||||
To explicitly utilize a custom deployment YAML, specify the configuration path:
|
To explicitly utilize a custom deployment YAML, specify the configuration path:
|
||||||
@@ -75,12 +72,6 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
|
|||||||
--stage-overrides '{"1": {"gpu_memory_utilization": 0.5}}'
|
--stage-overrides '{"1": {"gpu_memory_utilization": 0.5}}'
|
||||||
```
|
```
|
||||||
|
|
||||||
To experiment with the FlashInfer FP16 MoE path, set `VLLM_USE_FLASHINFER_MOE_FP16=1` before launching the server:
|
|
||||||
```bash
|
|
||||||
VLLM_USE_FLASHINFER_MOE_FP16=1 \
|
|
||||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
|
|
||||||
```
|
|
||||||
|
|
||||||
For the stage-based CLI, you usually do **not** need `--stage-overrides` for
|
For the stage-based CLI, you usually do **not** need `--stage-overrides` for
|
||||||
that kind of change. Since each command launches one stage, just pass the knob
|
that kind of change. Since each command launches one stage, just pass the knob
|
||||||
directly on that stage command:
|
directly on that stage command:
|
||||||
|
|||||||
@@ -380,31 +380,6 @@ curl -X POST http://localhost:8091/v1/audio/voices \
|
|||||||
```
|
```
|
||||||
Uploaded voices are then usable as `voice="custom_voice_1"` on subsequent requests.
|
Uploaded voices are then usable as `voice="custom_voice_1"` on subsequent requests.
|
||||||
|
|
||||||
### Precomputed custom voices
|
|
||||||
For reused Base voice-cloning speakers, precompute the reference artifacts once and load them at server startup:
|
|
||||||
```bash
|
|
||||||
python qwen3_tts/precompute_custom_voice.py \
|
|
||||||
--model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
|
|
||||||
--voice-name alice \
|
|
||||||
--ref-audio /path/to/reference.wav \
|
|
||||||
--ref-text "Original transcript of the reference audio" \
|
|
||||||
--mode icl \
|
|
||||||
--output-dir /path/to/custom_voices
|
|
||||||
```
|
|
||||||
`--mode icl` stores both `speaker_embedding` and `ref_code`; `--mode xvec` stores only the speaker embedding. Add the output directory to a deploy config:
|
|
||||||
```yaml
|
|
||||||
custom_voice_dir: /path/to/custom_voices
|
|
||||||
```
|
|
||||||
Then start the server with that config and call the Speech API with only the voice name:
|
|
||||||
```bash
|
|
||||||
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base --omni --deploy-config /path/to/qwen3_tts_custom_voice.yaml
|
|
||||||
|
|
||||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
|
||||||
-H "Content-Type: application/json" \
|
|
||||||
-d '{"input":"Hello from a precomputed voice.","voice":"alice","task_type":"Base"}' \
|
|
||||||
--output alice.wav
|
|
||||||
```
|
|
||||||
|
|
||||||
### Streaming PCM
|
### Streaming PCM
|
||||||
```bash
|
```bash
|
||||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
curl -X POST http://localhost:8091/v1/audio/speech \
|
||||||
@@ -444,10 +419,8 @@ python qwen3_tts/gradio_fastrtc_demo.py --api-base http://localhost:8000
|
|||||||
`qwen3_tts/batch_speech_client.py` issues many concurrent requests for throughput measurement.
|
`qwen3_tts/batch_speech_client.py` issues many concurrent requests for throughput measurement.
|
||||||
|
|
||||||
### Notes
|
### Notes
|
||||||
- Base voice cloning has uniproc-vs-mp tradeoffs depending on per-request reference audio cost; see the executor-backend section above.
|
- Base voice cloning has per-request reference-audio cost; the uniproc default keeps IPC overhead off the critical path. See the executor-backend section above for background.
|
||||||
- With async chunking, Qwen3-TTS Base voice cloning sends the full reference context in the first Code2Wav packet, then caches that prefix on the Code2Wav stage for follow-up chunks in the same request.
|
|
||||||
- `vllm_omni/deploy/qwen3_tts.yaml` is the default deploy config (loaded by HF `model_type`); per-stage runtime overrides are available via `--stage-N-<field> <value>`.
|
- `vllm_omni/deploy/qwen3_tts.yaml` is the default deploy config (loaded by HF `model_type`); per-stage runtime overrides are available via `--stage-N-<field> <value>`.
|
||||||
- Under vocoder-bound overload (single-stream `rtf_p99 ≥ 1` at the target concurrency), set `active_stream_window: 2` at the top of the deploy yaml to cap simultaneously active Stage 1 streams. Off by default; trades TTFP for streaming continuity. See [#3592](https://github.com/vllm-project/vllm-omni/pull/3592) for the mechanism and tradeoff numbers.
|
|
||||||
|
|
||||||
---
|
---
|
||||||
|
|
||||||
@@ -475,23 +448,6 @@ python voxcpm2/openai_speech_client.py \
|
|||||||
```
|
```
|
||||||
The `ref_audio` field accepts local file paths (auto-base64), HTTP URLs, or `data:audio/wav;base64,...` data URIs.
|
The `ref_audio` field accepts local file paths (auto-base64), HTTP URLs, or `data:audio/wav;base64,...` data URIs.
|
||||||
|
|
||||||
### Precomputed custom voices
|
|
||||||
For repeated VoxCPM2 speakers, precompute the prompt cache and load it through `custom_voice_dir`:
|
|
||||||
```bash
|
|
||||||
python voxcpm2/precompute_custom_voice.py \
|
|
||||||
--model openbmb/VoxCPM2 \
|
|
||||||
--voice-name alice \
|
|
||||||
--ref-audio /path/to/reference.wav \
|
|
||||||
--mode ref_continuation \
|
|
||||||
--prompt-text "Original transcript of the reference audio" \
|
|
||||||
--output-dir /path/to/custom_voices
|
|
||||||
```
|
|
||||||
Add the output directory to the deploy config:
|
|
||||||
```yaml
|
|
||||||
custom_voice_dir: /path/to/custom_voices
|
|
||||||
```
|
|
||||||
After startup, `/v1/audio/voices` lists `alice`, and `/v1/audio/speech` can use `voice="alice"` without sending `ref_audio`.
|
|
||||||
|
|
||||||
### Gradio demo (gapless streaming via AudioWorklet)
|
### Gradio demo (gapless streaming via AudioWorklet)
|
||||||
```bash
|
```bash
|
||||||
python voxcpm2/gradio_demo.py
|
python voxcpm2/gradio_demo.py
|
||||||
|
|||||||
@@ -74,7 +74,7 @@ Online:
|
|||||||
vllm serve Qwen/Qwen-Image \
|
vllm serve Qwen/Qwen-Image \
|
||||||
--omni \
|
--omni \
|
||||||
--port 8000 \
|
--port 8000 \
|
||||||
--diffusion-quantization-config '{"method":"gguf","gguf_model":"QuantStack/Qwen-Image-GGUF/Qwen_Image-Q4_K_M.gguf"}'
|
--quantization-config '{"method":"gguf","gguf_model":"QuantStack/Qwen-Image-GGUF/Qwen_Image-Q4_K_M.gguf"}'
|
||||||
```
|
```
|
||||||
|
|
||||||
## Parameters
|
## Parameters
|
||||||
|
|||||||
@@ -195,101 +195,6 @@ stages:
|
|||||||
devices: "0,1"
|
devices: "0,1"
|
||||||
```
|
```
|
||||||
|
|
||||||
### VAE Patch Parallelism
|
|
||||||
|
|
||||||
[VAE Patch Parallelism](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism/vae_patch_parallel.html) splits Bagel VAE **decode/encode** tiles across multiple GPUs on the **DiT stage**, reducing **per-GPU peak memory during VAE decode**. Use it when high-resolution `text2img` or `img2img` hits VAE OOM or large decode spikes.
|
|
||||||
|
|
||||||
**Bagel-specific notes:**
|
|
||||||
|
|
||||||
- Implemented in `BagelPipeline` via `DistributedAutoEncoder` (DiT stage only).
|
|
||||||
- **Single-stage** is the simplest path: one DiT process with TP + VAE patch parallel.
|
|
||||||
- **Two-stage**: enable on **stage 1 (DiT)** only; stage 0 (Thinker) keeps encoder-only `VAEEncoder` and does not use VAE patch parallel.
|
|
||||||
- You need a DiT `world_size` ≥ `vae_patch_parallel_size` (typically `tensor_parallel_size=2` on that stage). VAE PP reuses the DiT process group; it is not a standalone second-GPU VAE worker.
|
|
||||||
|
|
||||||
**Single-stage via deploy YAML** (recommended for `end2end.py`):
|
|
||||||
|
|
||||||
```yaml
|
|
||||||
pipeline: bagel_single_stage
|
|
||||||
async_chunk: false
|
|
||||||
|
|
||||||
stages:
|
|
||||||
- stage_id: 0
|
|
||||||
max_num_batched_tokens: 32768
|
|
||||||
max_num_seqs: 1
|
|
||||||
enforce_eager: true
|
|
||||||
trust_remote_code: true
|
|
||||||
enable_prefix_caching: false
|
|
||||||
devices: "0,1"
|
|
||||||
vae_use_tiling: true
|
|
||||||
parallel_config:
|
|
||||||
tensor_parallel_size: 2
|
|
||||||
vae_patch_parallel_size: 2
|
|
||||||
default_sampling_params:
|
|
||||||
seed: 52
|
|
||||||
```
|
|
||||||
|
|
||||||
```bash
|
|
||||||
cd examples/offline_inference/bagel
|
|
||||||
|
|
||||||
CUDA_VISIBLE_DEVICES=0,1 python end2end.py \
|
|
||||||
--model /path/to/BAGEL-7B-MoT \
|
|
||||||
--deploy-config /path/to/bagel_single_stage_vae_pp.yaml \
|
|
||||||
--modality text2img \
|
|
||||||
--prompts "A cute cat" \
|
|
||||||
--steps 10 \
|
|
||||||
--output ./out_vae_pp
|
|
||||||
```
|
|
||||||
|
|
||||||
**Single-stage via `Omni` kwargs** (same flags as online serving):
|
|
||||||
|
|
||||||
```python
|
|
||||||
from vllm_omni.entrypoints.omni import Omni
|
|
||||||
|
|
||||||
omni = Omni(
|
|
||||||
model="ByteDance-Seed/BAGEL-7B-MoT",
|
|
||||||
deploy_config="vllm_omni/deploy/bagel_single_stage.yaml",
|
|
||||||
tensor_parallel_size=2,
|
|
||||||
vae_patch_parallel_size=2,
|
|
||||||
vae_use_tiling=True,
|
|
||||||
)
|
|
||||||
# Then call omni.generate(...) as in end2end.py
|
|
||||||
```
|
|
||||||
|
|
||||||
**Two-stage (VAE PP on DiT only):**
|
|
||||||
|
|
||||||
```yaml
|
|
||||||
stages:
|
|
||||||
- stage_id: 0
|
|
||||||
devices: "0"
|
|
||||||
# AR Thinker — no vae_patch_parallel here
|
|
||||||
|
|
||||||
- stage_id: 1
|
|
||||||
devices: "0,1"
|
|
||||||
vae_use_tiling: true
|
|
||||||
parallel_config:
|
|
||||||
tensor_parallel_size: 2
|
|
||||||
vae_patch_parallel_size: 2
|
|
||||||
```
|
|
||||||
|
|
||||||
```bash
|
|
||||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
|
||||||
--deploy-config /path/to/bagel_vae_pp.yaml \
|
|
||||||
--modality text2img \
|
|
||||||
--prompts "A cute cat"
|
|
||||||
```
|
|
||||||
|
|
||||||
**Startup log checks:**
|
|
||||||
|
|
||||||
```text
|
|
||||||
INFO ... vae_patch_parallel_size=2 requires vae_use_tiling; automatically enabling it.
|
|
||||||
```
|
|
||||||
|
|
||||||
| Setting | Role |
|
|
||||||
| :------ | :--- |
|
|
||||||
| `parallel_config.tensor_parallel_size` | DiT world size / TP (must be ≥ `vae_patch_parallel_size`) |
|
|
||||||
| `parallel_config.vae_patch_parallel_size` | Number of ranks for distributed VAE tiles (`1` = off) |
|
|
||||||
| `vae_use_tiling` | Enable spatial tiling (auto-enabled when `vae_patch_parallel_size > 1`) |
|
|
||||||
|
|
||||||
#### Hybrid Sharded Data Parallel (HSDP)
|
#### Hybrid Sharded Data Parallel (HSDP)
|
||||||
|
|
||||||
For larger Bagel deployments on multiple GPUs, you can enable HSDP (Hybrid Sharded Data Parallel) by modifying the stage configuration (for example, [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml)). HSDP shards transformer weights across GPUs to reduce per-GPU memory usage.
|
For larger Bagel deployments on multiple GPUs, you can enable HSDP (Hybrid Sharded Data Parallel) by modifying the stage configuration (for example, [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml)). HSDP shards transformer weights across GPUs to reduce per-GPU memory usage.
|
||||||
|
|||||||
@@ -121,6 +121,9 @@ def parse_args():
|
|||||||
help="Temperature for text generation sampling (default: 0.3).",
|
help="Temperature for text generation sampling (default: 0.3).",
|
||||||
)
|
)
|
||||||
|
|
||||||
|
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||||
|
|
||||||
|
nullify_stage_engine_defaults(parser)
|
||||||
args = parser.parse_args()
|
args = parser.parse_args()
|
||||||
return args
|
return args
|
||||||
|
|
||||||
@@ -151,7 +154,6 @@ def main():
|
|||||||
from vllm_omni.entrypoints.omni import Omni
|
from vllm_omni.entrypoints.omni import Omni
|
||||||
|
|
||||||
omni_kwargs = vars(args).copy()
|
omni_kwargs = vars(args).copy()
|
||||||
omni_kwargs.pop("model", None)
|
|
||||||
deploy_config = args.deploy_config
|
deploy_config = args.deploy_config
|
||||||
if args.think and deploy_config is None:
|
if args.think and deploy_config is None:
|
||||||
deploy_config = "vllm_omni/deploy/bagel_think.yaml"
|
deploy_config = "vllm_omni/deploy/bagel_think.yaml"
|
||||||
@@ -162,7 +164,9 @@ def main():
|
|||||||
if args.quantization:
|
if args.quantization:
|
||||||
omni_kwargs["quantization_config"] = args.quantization
|
omni_kwargs["quantization_config"] = args.quantization
|
||||||
|
|
||||||
omni = Omni(model=model_name, **omni_kwargs)
|
# Override CLI --model with the derived model_name.
|
||||||
|
omni_kwargs["model"] = model_name
|
||||||
|
omni = Omni(**omni_kwargs)
|
||||||
|
|
||||||
formatted_prompts = []
|
formatted_prompts = []
|
||||||
for p in prompts:
|
for p in prompts:
|
||||||
|
|||||||
@@ -1,304 +0,0 @@
|
|||||||
#!/usr/bin/env python3
|
|
||||||
# SPDX-License-Identifier: Apache-2.0
|
|
||||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
||||||
|
|
||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
import argparse
|
|
||||||
import uuid
|
|
||||||
from pathlib import Path
|
|
||||||
|
|
||||||
import cv2
|
|
||||||
import numpy as np
|
|
||||||
import torch
|
|
||||||
from PIL import Image
|
|
||||||
|
|
||||||
from vllm_omni import Omni
|
|
||||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
|
||||||
from vllm_omni.outputs import OmniRequestOutput
|
|
||||||
|
|
||||||
WORKER_EXTENSION = "vllm_omni.diffusion.models.dreamzero.video_export_worker.DreamZeroVideoExportWorkerExtension"
|
|
||||||
ASSET_REPO_ID = "YangshenDeng/vllm-omni-dreamzero-assets"
|
|
||||||
DEFAULT_SESSION_PREFIX = "dreamzero-export"
|
|
||||||
RELATIVE_OFFSETS = [-23, -16, -8, 0]
|
|
||||||
ACTION_HORIZON = 24
|
|
||||||
CAMERA_FILES = {
|
|
||||||
"observation/exterior_image_0_left": "exterior_image_1_left.mp4",
|
|
||||||
"observation/exterior_image_1_left": "exterior_image_2_left.mp4",
|
|
||||||
"observation/wrist_image_left": "wrist_image_left.mp4",
|
|
||||||
}
|
|
||||||
|
|
||||||
|
|
||||||
def _parse_args() -> argparse.Namespace:
|
|
||||||
repo_root = Path(__file__).resolve().parents[3]
|
|
||||||
parser = argparse.ArgumentParser(description="Export DreamZero prediction video from downloaded example inputs.")
|
|
||||||
parser.add_argument("--model", default="GEAR-Dreams/DreamZero-DROID")
|
|
||||||
parser.add_argument("--deploy-config", type=Path, required=True)
|
|
||||||
parser.add_argument(
|
|
||||||
"--video-dir",
|
|
||||||
type=Path,
|
|
||||||
default=repo_root / "outputs" / "dreamzero" / "assets",
|
|
||||||
help="Directory containing the three camera MP4 files.",
|
|
||||||
)
|
|
||||||
parser.add_argument(
|
|
||||||
"--output-dir", type=Path, default=repo_root / "outputs" / "dreamzero" / "generated_predictions"
|
|
||||||
)
|
|
||||||
parser.add_argument("--output-stem", default="dreamzero_prediction")
|
|
||||||
parser.add_argument(
|
|
||||||
"--prompt",
|
|
||||||
default="Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan",
|
|
||||||
)
|
|
||||||
parser.add_argument("--session-id", default=None)
|
|
||||||
parser.add_argument("--save-input-video", action="store_true")
|
|
||||||
parser.add_argument("--save-gif", action="store_true")
|
|
||||||
parser.add_argument("--save-actions", action="store_true")
|
|
||||||
parser.add_argument("--fps", type=int, default=5)
|
|
||||||
return parser.parse_args()
|
|
||||||
|
|
||||||
|
|
||||||
def _load_all_frames(video_path: Path) -> np.ndarray:
|
|
||||||
cap = cv2.VideoCapture(str(video_path))
|
|
||||||
frames = []
|
|
||||||
while True:
|
|
||||||
ok, frame = cap.read()
|
|
||||||
if not ok:
|
|
||||||
break
|
|
||||||
frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
|
|
||||||
cap.release()
|
|
||||||
if not frames:
|
|
||||||
raise RuntimeError(f"No frames loaded from {video_path}")
|
|
||||||
return np.stack(frames, axis=0)
|
|
||||||
|
|
||||||
|
|
||||||
def _load_camera_frames(video_dir: Path) -> dict[str, np.ndarray]:
|
|
||||||
camera_frames: dict[str, np.ndarray] = {}
|
|
||||||
for camera_key, file_name in CAMERA_FILES.items():
|
|
||||||
video_path = video_dir / file_name
|
|
||||||
if not video_path.exists():
|
|
||||||
raise FileNotFoundError(
|
|
||||||
f"Missing DreamZero example asset: {video_path}. "
|
|
||||||
"Download the example videos with: "
|
|
||||||
f"`hf download {ASSET_REPO_ID} --repo-type dataset --local-dir {video_dir}`"
|
|
||||||
)
|
|
||||||
camera_frames[camera_key] = _load_all_frames(video_path)
|
|
||||||
return camera_frames
|
|
||||||
|
|
||||||
|
|
||||||
def _build_frame_schedule(total_frames: int, num_chunks: int) -> list[list[int]]:
|
|
||||||
chunks: list[list[int]] = []
|
|
||||||
current_frame = 23
|
|
||||||
for _ in range(num_chunks):
|
|
||||||
indices = [max(current_frame + offset, 0) for offset in RELATIVE_OFFSETS]
|
|
||||||
if indices[-1] >= total_frames:
|
|
||||||
break
|
|
||||||
chunks.append(indices)
|
|
||||||
current_frame += ACTION_HORIZON
|
|
||||||
return chunks
|
|
||||||
|
|
||||||
|
|
||||||
def _make_obs_from_video(
|
|
||||||
camera_frames: dict[str, np.ndarray],
|
|
||||||
frame_indices: list[int],
|
|
||||||
*,
|
|
||||||
prompt: str,
|
|
||||||
session_id: str,
|
|
||||||
) -> dict:
|
|
||||||
obs: dict = {}
|
|
||||||
for camera_key, all_frames in camera_frames.items():
|
|
||||||
selected = all_frames[frame_indices]
|
|
||||||
obs[camera_key] = selected[0] if len(frame_indices) == 1 else selected
|
|
||||||
|
|
||||||
obs["observation/joint_position"] = np.zeros(7, dtype=np.float32)
|
|
||||||
obs["observation/cartesian_position"] = np.zeros(6, dtype=np.float32)
|
|
||||||
obs["observation/gripper_position"] = np.zeros(1, dtype=np.float32)
|
|
||||||
obs["prompt"] = prompt
|
|
||||||
obs["session_id"] = session_id
|
|
||||||
return obs
|
|
||||||
|
|
||||||
|
|
||||||
def _build_observations(video_dir: Path, prompt: str, session_id: str) -> tuple[dict[str, np.ndarray], list[dict]]:
|
|
||||||
camera_frames = _load_camera_frames(video_dir)
|
|
||||||
total_frames = min(frames.shape[0] for frames in camera_frames.values())
|
|
||||||
chunks = _build_frame_schedule(total_frames, 1)
|
|
||||||
observations = [
|
|
||||||
_make_obs_from_video(camera_frames, [0], prompt=prompt, session_id=session_id),
|
|
||||||
]
|
|
||||||
if chunks:
|
|
||||||
observations.append(
|
|
||||||
_make_obs_from_video(
|
|
||||||
camera_frames,
|
|
||||||
chunks[0],
|
|
||||||
prompt=prompt,
|
|
||||||
session_id=session_id,
|
|
||||||
)
|
|
||||||
)
|
|
||||||
if len(observations) < 2:
|
|
||||||
raise RuntimeError("Need at least two DreamZero example observations to export a prediction video.")
|
|
||||||
return camera_frames, observations[:2]
|
|
||||||
|
|
||||||
|
|
||||||
def _extract_latents(output: OmniRequestOutput) -> torch.Tensor:
|
|
||||||
if not isinstance(output, OmniRequestOutput):
|
|
||||||
raise TypeError(f"Expected OmniRequestOutput, got {type(output)!r}")
|
|
||||||
if not output.images:
|
|
||||||
raise RuntimeError("DreamZero output does not contain video latents in `images`.")
|
|
||||||
|
|
||||||
latents = output.images[0]
|
|
||||||
if not isinstance(latents, torch.Tensor):
|
|
||||||
raise TypeError(f"Expected tensor latents, got {type(latents)!r}")
|
|
||||||
|
|
||||||
latents = latents.detach().cpu()
|
|
||||||
if latents.dim() == 4:
|
|
||||||
latents = latents.unsqueeze(0)
|
|
||||||
if latents.dim() != 5:
|
|
||||||
raise ValueError(f"Unexpected latent shape: {tuple(latents.shape)}")
|
|
||||||
|
|
||||||
if latents.shape[1] < latents.shape[2]:
|
|
||||||
latents = latents.transpose(1, 2).contiguous()
|
|
||||||
return latents
|
|
||||||
|
|
||||||
|
|
||||||
def _write_mp4(path: Path, frames: np.ndarray, fps: int) -> None:
|
|
||||||
path.parent.mkdir(parents=True, exist_ok=True)
|
|
||||||
height, width = frames.shape[1:3]
|
|
||||||
writer = cv2.VideoWriter(
|
|
||||||
str(path),
|
|
||||||
cv2.VideoWriter_fourcc(*"mp4v"),
|
|
||||||
float(fps),
|
|
||||||
(width, height),
|
|
||||||
)
|
|
||||||
if not writer.isOpened():
|
|
||||||
raise RuntimeError(f"Failed to open video writer for {path}")
|
|
||||||
try:
|
|
||||||
for frame in frames:
|
|
||||||
writer.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
|
|
||||||
finally:
|
|
||||||
writer.release()
|
|
||||||
|
|
||||||
|
|
||||||
def _write_gif(path: Path, frames: np.ndarray, fps: int) -> None:
|
|
||||||
path.parent.mkdir(parents=True, exist_ok=True)
|
|
||||||
images = [Image.fromarray(frame) for frame in frames]
|
|
||||||
duration_ms = max(int(round(1000 / max(fps, 1))), 1)
|
|
||||||
images[0].save(
|
|
||||||
path,
|
|
||||||
save_all=True,
|
|
||||||
append_images=images[1:],
|
|
||||||
duration=duration_ms,
|
|
||||||
loop=0,
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
def _stitch_input_frames(camera_frames: dict[str, np.ndarray]) -> np.ndarray:
|
|
||||||
total_frames = min(frames.shape[0] for frames in camera_frames.values())
|
|
||||||
stitched = []
|
|
||||||
for frame_index in range(total_frames):
|
|
||||||
left = camera_frames["observation/exterior_image_0_left"][frame_index]
|
|
||||||
right = camera_frames["observation/exterior_image_1_left"][frame_index]
|
|
||||||
wrist = camera_frames["observation/wrist_image_left"][frame_index]
|
|
||||||
pad = np.zeros((left.shape[0], left.shape[1], 3), dtype=np.uint8)
|
|
||||||
canvas = np.concatenate([left, right], axis=1)
|
|
||||||
bottom = np.concatenate([wrist, pad], axis=1)
|
|
||||||
stitched.append(np.concatenate([canvas, bottom], axis=0))
|
|
||||||
return np.stack(stitched, axis=0)
|
|
||||||
|
|
||||||
|
|
||||||
def _run_generation(
|
|
||||||
model: str, deploy_config_path: Path, observations: list[dict]
|
|
||||||
) -> tuple[Omni, list[OmniRequestOutput]]:
|
|
||||||
omni = Omni(
|
|
||||||
model=model,
|
|
||||||
deploy_config=str(deploy_config_path),
|
|
||||||
enforce_eager=True,
|
|
||||||
worker_extension_cls=WORKER_EXTENSION,
|
|
||||||
)
|
|
||||||
|
|
||||||
outputs: list[OmniRequestOutput] = []
|
|
||||||
for index, obs in enumerate(observations):
|
|
||||||
sampling_params = OmniDiffusionSamplingParams(
|
|
||||||
extra_args={
|
|
||||||
"reset": index == 0,
|
|
||||||
"session_id": obs["session_id"],
|
|
||||||
"robot_obs": obs,
|
|
||||||
}
|
|
||||||
)
|
|
||||||
result = omni.generate(obs["prompt"], sampling_params_list=[sampling_params])
|
|
||||||
if not result:
|
|
||||||
raise RuntimeError(f"No output returned for DreamZero request {index}")
|
|
||||||
outputs.append(result[0])
|
|
||||||
return omni, outputs
|
|
||||||
|
|
||||||
|
|
||||||
def _decode_with_worker(omni: Omni, full_latents: torch.Tensor) -> np.ndarray:
|
|
||||||
stage_client = omni.engine.stage_clients[0]
|
|
||||||
engine = getattr(stage_client, "_engine", None)
|
|
||||||
if engine is None:
|
|
||||||
raise RuntimeError("DreamZero export requires inline diffusion stage access.")
|
|
||||||
|
|
||||||
decoded = engine.executor.collective_rpc(
|
|
||||||
"decode_video_latents_to_uint8",
|
|
||||||
args=(full_latents,),
|
|
||||||
unique_reply_rank=0,
|
|
||||||
exec_all_ranks=True,
|
|
||||||
)
|
|
||||||
if isinstance(decoded, torch.Tensor):
|
|
||||||
decoded = decoded.numpy()
|
|
||||||
if not isinstance(decoded, np.ndarray):
|
|
||||||
raise TypeError(f"Unexpected decoded output type: {type(decoded)!r}")
|
|
||||||
return decoded
|
|
||||||
|
|
||||||
|
|
||||||
def main() -> None:
|
|
||||||
args = _parse_args()
|
|
||||||
session_id = args.session_id or f"{DEFAULT_SESSION_PREFIX}-{uuid.uuid4()}"
|
|
||||||
|
|
||||||
camera_frames, observations = _build_observations(
|
|
||||||
video_dir=args.video_dir,
|
|
||||||
prompt=args.prompt,
|
|
||||||
session_id=session_id,
|
|
||||||
)
|
|
||||||
|
|
||||||
args.output_dir.mkdir(parents=True, exist_ok=True)
|
|
||||||
|
|
||||||
if args.save_input_video:
|
|
||||||
input_frames = _stitch_input_frames(camera_frames)
|
|
||||||
_write_mp4(args.output_dir / f"{args.output_stem}_input.mp4", input_frames, fps=15)
|
|
||||||
if args.save_gif:
|
|
||||||
_write_gif(args.output_dir / f"{args.output_stem}_input.gif", input_frames[::3], fps=5)
|
|
||||||
|
|
||||||
omni = None
|
|
||||||
try:
|
|
||||||
omni, outputs = _run_generation(
|
|
||||||
model=args.model,
|
|
||||||
deploy_config_path=args.deploy_config,
|
|
||||||
observations=observations,
|
|
||||||
)
|
|
||||||
latent_steps = [_extract_latents(output) for output in outputs]
|
|
||||||
full_latents = torch.cat(latent_steps, dim=2)
|
|
||||||
frames = _decode_with_worker(omni, full_latents)
|
|
||||||
finally:
|
|
||||||
if omni is not None:
|
|
||||||
omni.close()
|
|
||||||
|
|
||||||
mp4_path = args.output_dir / f"{args.output_stem}.mp4"
|
|
||||||
|
|
||||||
_write_mp4(mp4_path, frames, fps=args.fps)
|
|
||||||
print(f"SAVED_MP4={mp4_path}")
|
|
||||||
|
|
||||||
if args.save_gif:
|
|
||||||
gif_path = args.output_dir / f"{args.output_stem}.gif"
|
|
||||||
_write_gif(gif_path, frames, fps=args.fps)
|
|
||||||
print(f"SAVED_GIF={gif_path}")
|
|
||||||
|
|
||||||
if args.save_actions:
|
|
||||||
npz_path = args.output_dir / f"{args.output_stem}_actions.npz"
|
|
||||||
np.savez(
|
|
||||||
npz_path,
|
|
||||||
step0=np.asarray(outputs[0].multimodal_output.get("actions")),
|
|
||||||
step1=np.asarray(outputs[1].multimodal_output.get("actions")),
|
|
||||||
)
|
|
||||||
print(f"SAVED_ACTIONS={npz_path}")
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
main()
|
|
||||||
@@ -42,7 +42,7 @@ pip install \
|
|||||||
## 3. Hardware and VRAM Requirements
|
## 3. Hardware and VRAM Requirements
|
||||||
|
|
||||||
This example uses a 3-stage pipeline on one GPU by default
|
This example uses a 3-stage pipeline on one GPU by default
|
||||||
([`dynin_omni.yaml`](../../../vllm_omni/deploy/dynin_omni.yaml)):
|
([`dynin_omni.yaml`](../../../vllm_omni/model_executor/stage_configs/dynin_omni.yaml)):
|
||||||
|
|
||||||
- Stage-0 (`token2text`): `gpu_memory_utilization: 0.5`
|
- Stage-0 (`token2text`): `gpu_memory_utilization: 0.5`
|
||||||
- Stage-1 (`token2image`): `gpu_memory_utilization: 0.1`
|
- Stage-1 (`token2image`): `gpu_memory_utilization: 0.1`
|
||||||
|
|||||||
@@ -18,6 +18,8 @@ import numpy as np
|
|||||||
import torch
|
import torch
|
||||||
from PIL import Image
|
from PIL import Image
|
||||||
|
|
||||||
|
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||||
|
|
||||||
TASK_CHOICES = ("t2t", "t2i", "t2s", "i2i", "i2t", "s2t", "v2t")
|
TASK_CHOICES = ("t2t", "t2i", "t2s", "i2i", "i2t", "s2t", "v2t")
|
||||||
|
|
||||||
TASK_DEFAULT_RUNTIME = {
|
TASK_DEFAULT_RUNTIME = {
|
||||||
@@ -878,8 +880,8 @@ def parse_args(repo_root: Path) -> argparse.Namespace:
|
|||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
"--stage-config-path",
|
"--stage-config-path",
|
||||||
type=str,
|
type=str,
|
||||||
default=str(repo_root / "vllm_omni/deploy/dynin_omni.yaml"),
|
default=str(repo_root / "vllm_omni/model_executor/stage_configs/dynin_omni.yaml"),
|
||||||
help="Path to deploy config yaml.",
|
help="Path to stage config yaml.",
|
||||||
)
|
)
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
"--dynin-config-path",
|
"--dynin-config-path",
|
||||||
@@ -970,6 +972,8 @@ def parse_args(repo_root: Path) -> argparse.Namespace:
|
|||||||
parser.add_argument("--vq-model-audio-local-files-only", action=argparse.BooleanOptionalAction, default=None)
|
parser.add_argument("--vq-model-audio-local-files-only", action=argparse.BooleanOptionalAction, default=None)
|
||||||
|
|
||||||
parser.add_argument("--disable-hf-xet", action=argparse.BooleanOptionalAction, default=True)
|
parser.add_argument("--disable-hf-xet", action=argparse.BooleanOptionalAction, default=True)
|
||||||
|
|
||||||
|
nullify_stage_engine_defaults(parser)
|
||||||
return parser.parse_args()
|
return parser.parse_args()
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -107,6 +107,9 @@ def parse_args():
|
|||||||
),
|
),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||||
|
|
||||||
|
nullify_stage_engine_defaults(parser)
|
||||||
return parser.parse_args()
|
return parser.parse_args()
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -210,12 +210,6 @@ def main():
|
|||||||
args = parse_args()
|
args = parse_args()
|
||||||
os.makedirs(args.output, exist_ok=True)
|
os.makedirs(args.output, exist_ok=True)
|
||||||
|
|
||||||
# Lance aligns more naturally with upstream when using FLASH_ATTN
|
|
||||||
# (upstream uses flash_attn_varlen_func; SDPA accumulates ~5x more
|
|
||||||
# numerical drift through the 36-layer Qwen2 stack on B300). Default
|
|
||||||
# to flash_attn when the user hasn't explicitly set a backend.
|
|
||||||
os.environ.setdefault("DIFFUSION_ATTENTION_BACKEND", "FLASH_ATTN")
|
|
||||||
|
|
||||||
if args.txt_prompts:
|
if args.txt_prompts:
|
||||||
with open(args.txt_prompts, encoding="utf-8") as f:
|
with open(args.txt_prompts, encoding="utf-8") as f:
|
||||||
prompts = [ln.strip() for ln in f if ln.strip()]
|
prompts = [ln.strip() for ln in f if ln.strip()]
|
||||||
|
|||||||
@@ -288,19 +288,11 @@ def _init_models(
|
|||||||
f" replicas={replicas_per_omni}...",
|
f" replicas={replicas_per_omni}...",
|
||||||
flush=True,
|
flush=True,
|
||||||
)
|
)
|
||||||
# Lance single-stage engine defaults — mirrors the offline ``end2end.py``
|
# Lance needs ``vllm_omni/deploy/lance.yaml`` as the pipeline-selector
|
||||||
# flat-kwargs path so this demo does not depend on ``deploy/lance.yaml``.
|
# (Lance HF config has no ``model_type`` field for auto-detect). Pass
|
||||||
# ``pipeline="lance"`` is the explicit selector (Lance HF config has no
|
# via ``deploy_config`` — flat kwargs alone don't trigger pipeline
|
||||||
# ``model_type`` field, so auto-detect from the model dir alone fails).
|
# resolution.
|
||||||
lance_engine_defaults = dict(
|
lance_engine_defaults = dict(deploy_config="vllm_omni/deploy/lance.yaml")
|
||||||
pipeline="lance",
|
|
||||||
max_num_batched_tokens=32768,
|
|
||||||
max_num_seqs=1,
|
|
||||||
enforce_eager=True,
|
|
||||||
trust_remote_code=True,
|
|
||||||
enable_prefix_caching=False,
|
|
||||||
async_chunk=False,
|
|
||||||
)
|
|
||||||
|
|
||||||
img_kwargs: dict = {"model": img_ckpt, **lance_engine_defaults}
|
img_kwargs: dict = {"model": img_ckpt, **lance_engine_defaults}
|
||||||
if img_devices != "0":
|
if img_devices != "0":
|
||||||
@@ -1245,11 +1237,6 @@ def build_ui(
|
|||||||
|
|
||||||
def main():
|
def main():
|
||||||
args = parse_args()
|
args = parse_args()
|
||||||
# Lance aligns better with upstream when using FLASH_ATTN
|
|
||||||
# (upstream uses flash_attn_varlen_func; SDPA accumulates ~5x more
|
|
||||||
# numerical drift through the 36-layer Qwen2 stack). Default to
|
|
||||||
# flash_attn unless the user explicitly set a backend.
|
|
||||||
os.environ.setdefault("DIFFUSION_ATTENTION_BACKEND", "FLASH_ATTN")
|
|
||||||
print("=" * 60, flush=True)
|
print("=" * 60, flush=True)
|
||||||
print("[unified] Lance × vllm-omni demo", flush=True)
|
print("[unified] Lance × vllm-omni demo", flush=True)
|
||||||
print(f" snapshot: {args.model}", flush=True)
|
print(f" snapshot: {args.model}", flush=True)
|
||||||
|
|||||||
@@ -19,6 +19,7 @@ from vllm import SamplingParams
|
|||||||
from vllm.multimodal.image import convert_image_mode
|
from vllm.multimodal.image import convert_image_mode
|
||||||
|
|
||||||
from vllm_omni import Omni
|
from vllm_omni import Omni
|
||||||
|
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||||
|
|
||||||
DEFAULT_SYSTEM = "You are a helpful assistant."
|
DEFAULT_SYSTEM = "You are a helpful assistant."
|
||||||
DEFAULT_QUESTION = "Please summarize the content of this image."
|
DEFAULT_QUESTION = "Please summarize the content of this image."
|
||||||
@@ -48,6 +49,7 @@ def parse_args() -> argparse.Namespace:
|
|||||||
action="store_true",
|
action="store_true",
|
||||||
help="Enable diffusion pipeline profiler to display stage durations.",
|
help="Enable diffusion pipeline profiler to display stage durations.",
|
||||||
)
|
)
|
||||||
|
nullify_stage_engine_defaults(parser)
|
||||||
return parser.parse_args()
|
return parser.parse_args()
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -29,6 +29,7 @@ from PIL import Image
|
|||||||
from vllm.sampling_params import SamplingParams
|
from vllm.sampling_params import SamplingParams
|
||||||
|
|
||||||
from vllm_omni import Omni
|
from vllm_omni import Omni
|
||||||
|
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||||
|
|
||||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
|
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
@@ -117,6 +118,7 @@ def parse_args() -> argparse.Namespace:
|
|||||||
)
|
)
|
||||||
p.add_argument("--out", type=str, default="output.png", help="Path to save the generated image.")
|
p.add_argument("--out", type=str, default="output.png", help="Path to save the generated image.")
|
||||||
p.add_argument("--trust-remote-code", action="store_true", help="Trust remote code when loading the model.")
|
p.add_argument("--trust-remote-code", action="store_true", help="Trust remote code when loading the model.")
|
||||||
|
nullify_stage_engine_defaults(p)
|
||||||
args = p.parse_args()
|
args = p.parse_args()
|
||||||
if not args.prompt:
|
if not args.prompt:
|
||||||
args.prompt = ["A stylish woman with sunglasses riding a motorcycle in NYC."]
|
args.prompt = ["A stylish woman with sunglasses riding a motorcycle in NYC."]
|
||||||
|
|||||||
@@ -21,10 +21,11 @@ from message_convert import (
|
|||||||
to_prompt,
|
to_prompt,
|
||||||
)
|
)
|
||||||
from vllm import SamplingParams
|
from vllm import SamplingParams
|
||||||
|
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||||
|
|
||||||
|
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||||
from vllm_omni.entrypoints.omni import Omni
|
from vllm_omni.entrypoints.omni import Omni
|
||||||
from vllm_omni.inputs.data import OmniTokensPrompt
|
from vllm_omni.inputs.data import OmniTokensPrompt
|
||||||
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
|
|
||||||
|
|
||||||
SEED = 42
|
SEED = 42
|
||||||
|
|
||||||
@@ -342,7 +343,7 @@ def main(args):
|
|||||||
|
|
||||||
|
|
||||||
def parse_args():
|
def parse_args():
|
||||||
parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
|
parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
"--model-name",
|
"--model-name",
|
||||||
"-m",
|
"-m",
|
||||||
@@ -438,6 +439,7 @@ def parse_args():
|
|||||||
"vllm_omni/deploy/mimo_audio.yaml based on the HF model_type.",
|
"vllm_omni/deploy/mimo_audio.yaml based on the HF model_type.",
|
||||||
)
|
)
|
||||||
|
|
||||||
|
nullify_stage_engine_defaults(parser)
|
||||||
return parser.parse_args()
|
return parser.parse_args()
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -53,7 +53,7 @@ python examples/offline_inference/ming_flash_omni/end2end.py \
|
|||||||
--image-path ./3_0.png
|
--image-path ./3_0.png
|
||||||
```
|
```
|
||||||
|
|
||||||
### Omni-Speech (Thinker + Talker) {#omni-speech-thinker--talker}
|
### Omni-Speech (Thinker + Talker)
|
||||||
|
|
||||||
The default deploy YAML already runs thinker+talker, so spoken output only requires requesting `audio` (or `text,audio`) modalities.
|
The default deploy YAML already runs thinker+talker, so spoken output only requires requesting `audio` (or `text,audio`) modalities.
|
||||||
The thinker processes your multimodal input, generates text, then the talker synthesises the response as speech.
|
The thinker processes your multimodal input, generates text, then the talker synthesises the response as speech.
|
||||||
|
|||||||
@@ -17,13 +17,14 @@ from vllm.assets.image import ImageAsset
|
|||||||
from vllm.assets.video import VideoAsset, video_to_ndarrays
|
from vllm.assets.video import VideoAsset, video_to_ndarrays
|
||||||
from vllm.multimodal.image import convert_image_mode
|
from vllm.multimodal.image import convert_image_mode
|
||||||
from vllm.multimodal.media.audio import load_audio
|
from vllm.multimodal.media.audio import load_audio
|
||||||
|
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||||
|
|
||||||
import vllm_omni
|
import vllm_omni
|
||||||
|
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||||
from vllm_omni.entrypoints.omni import Omni
|
from vllm_omni.entrypoints.omni import Omni
|
||||||
|
|
||||||
# Imports the processor also registers itself
|
# Imports the processor also registers itself
|
||||||
from vllm_omni.transformers_utils.processors.ming import MingFlashOmniProcessor # noqa: F401
|
from vllm_omni.transformers_utils.processors.ming import MingFlashOmniProcessor # noqa: F401
|
||||||
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
|
|
||||||
|
|
||||||
SEED = 42
|
SEED = 42
|
||||||
MODEL_NAME = "Jonathan1909/Ming-flash-omni-2.0"
|
MODEL_NAME = "Jonathan1909/Ming-flash-omni-2.0"
|
||||||
@@ -396,7 +397,7 @@ def main(args):
|
|||||||
|
|
||||||
|
|
||||||
def parse_args():
|
def parse_args():
|
||||||
parser = TrackingArgumentParser(description="Ming-flash-omni 2.0 offline inference example")
|
parser = FlexibleArgumentParser(description="Ming-flash-omni 2.0 offline inference example")
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
"--query-type",
|
"--query-type",
|
||||||
"-q",
|
"-q",
|
||||||
@@ -495,6 +496,7 @@ def parse_args():
|
|||||||
help="Output directory for results.",
|
help="Output directory for results.",
|
||||||
)
|
)
|
||||||
|
|
||||||
|
nullify_stage_engine_defaults(parser)
|
||||||
return parser.parse_args()
|
return parser.parse_args()
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -19,9 +19,10 @@ from vllm.assets.video import VideoAsset, video_to_ndarrays
|
|||||||
from vllm.multimodal.image import convert_image_mode
|
from vllm.multimodal.image import convert_image_mode
|
||||||
from vllm.multimodal.media.audio import load_audio
|
from vllm.multimodal.media.audio import load_audio
|
||||||
from vllm.sampling_params import SamplingParams
|
from vllm.sampling_params import SamplingParams
|
||||||
|
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||||
|
|
||||||
|
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||||
from vllm_omni.entrypoints.omni import Omni
|
from vllm_omni.entrypoints.omni import Omni
|
||||||
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
|
|
||||||
|
|
||||||
SEED = 42
|
SEED = 42
|
||||||
|
|
||||||
@@ -431,7 +432,7 @@ def main(args):
|
|||||||
|
|
||||||
|
|
||||||
def parse_args():
|
def parse_args():
|
||||||
parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
|
parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
"--model",
|
"--model",
|
||||||
type=str,
|
type=str,
|
||||||
@@ -553,6 +554,7 @@ def parse_args():
|
|||||||
default=False,
|
default=False,
|
||||||
help="Use py_generator mode. The returned type of Omni.generate() is a Python Generator object.",
|
help="Use py_generator mode. The returned type of Omni.generate() is a Python Generator object.",
|
||||||
)
|
)
|
||||||
|
nullify_stage_engine_defaults(parser)
|
||||||
return parser.parse_args()
|
return parser.parse_args()
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -19,9 +19,10 @@ from vllm.assets.image import ImageAsset
|
|||||||
from vllm.assets.video import VideoAsset, video_to_ndarrays
|
from vllm.assets.video import VideoAsset, video_to_ndarrays
|
||||||
from vllm.multimodal.image import convert_image_mode
|
from vllm.multimodal.image import convert_image_mode
|
||||||
from vllm.multimodal.media.audio import load_audio
|
from vllm.multimodal.media.audio import load_audio
|
||||||
|
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||||
|
|
||||||
|
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||||
from vllm_omni.entrypoints.omni import Omni
|
from vllm_omni.entrypoints.omni import Omni
|
||||||
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
|
|
||||||
|
|
||||||
SEED = 42
|
SEED = 42
|
||||||
|
|
||||||
@@ -421,7 +422,7 @@ def main(args):
|
|||||||
|
|
||||||
|
|
||||||
def parse_args():
|
def parse_args():
|
||||||
parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
|
parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
"--model",
|
"--model",
|
||||||
type=str,
|
type=str,
|
||||||
@@ -560,6 +561,7 @@ def parse_args():
|
|||||||
help="Model dtype (auto, half, float16, bfloat16, float, float32).",
|
help="Model dtype (auto, half, float16, bfloat16, float, float32).",
|
||||||
)
|
)
|
||||||
|
|
||||||
|
nullify_stage_engine_defaults(parser)
|
||||||
return parser.parse_args()
|
return parser.parse_args()
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -39,9 +39,10 @@ from vllm.assets.image import ImageAsset
|
|||||||
from vllm.assets.video import VideoAsset, video_to_ndarrays
|
from vllm.assets.video import VideoAsset, video_to_ndarrays
|
||||||
from vllm.multimodal.image import convert_image_mode
|
from vllm.multimodal.image import convert_image_mode
|
||||||
from vllm.multimodal.media.audio import load_audio
|
from vllm.multimodal.media.audio import load_audio
|
||||||
|
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||||
|
|
||||||
|
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||||
from vllm_omni.entrypoints.async_omni import AsyncOmni
|
from vllm_omni.entrypoints.async_omni import AsyncOmni
|
||||||
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
|
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
@@ -382,9 +383,7 @@ async def run_all(args):
|
|||||||
print(f"[Info] Creating AsyncOmni with deploy_config={args.deploy_config}")
|
print(f"[Info] Creating AsyncOmni with deploy_config={args.deploy_config}")
|
||||||
async_omni = None
|
async_omni = None
|
||||||
try:
|
try:
|
||||||
# ``from_cli_args`` forwards only explicitly-passed CLI args so
|
async_omni = AsyncOmni(**vars(args))
|
||||||
# argparse defaults do not silently override deploy YAML values.
|
|
||||||
async_omni = AsyncOmni.from_cli_args(args, model=args.model)
|
|
||||||
|
|
||||||
# Use default sampling params from stage config (they are pre-configured
|
# Use default sampling params from stage config (they are pre-configured
|
||||||
# in the YAML for each stage).
|
# in the YAML for each stage).
|
||||||
@@ -454,7 +453,7 @@ async def run_all(args):
|
|||||||
|
|
||||||
|
|
||||||
def parse_args():
|
def parse_args():
|
||||||
parser = TrackingArgumentParser(
|
parser = FlexibleArgumentParser(
|
||||||
description=(
|
description=(
|
||||||
"Offline inference with async_chunk enabled via AsyncOmni. "
|
"Offline inference with async_chunk enabled via AsyncOmni. "
|
||||||
"Downstream stages start before upstream stages finish, "
|
"Downstream stages start before upstream stages finish, "
|
||||||
@@ -590,6 +589,7 @@ def parse_args():
|
|||||||
default=16000,
|
default=16000,
|
||||||
help="Sampling rate for audio loading.",
|
help="Sampling rate for audio loading.",
|
||||||
)
|
)
|
||||||
|
nullify_stage_engine_defaults(parser)
|
||||||
return parser.parse_args()
|
return parser.parse_args()
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -29,6 +29,7 @@ import os
|
|||||||
|
|
||||||
from PIL import Image
|
from PIL import Image
|
||||||
|
|
||||||
|
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||||
from vllm_omni.entrypoints.omni import Omni
|
from vllm_omni.entrypoints.omni import Omni
|
||||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||||
|
|
||||||
@@ -167,6 +168,7 @@ def parse_args():
|
|||||||
help="Enable cache-dit summary logging after diffusion forward passes.",
|
help="Enable cache-dit summary logging after diffusion forward passes.",
|
||||||
)
|
)
|
||||||
|
|
||||||
|
nullify_stage_engine_defaults(parser)
|
||||||
return parser.parse_args()
|
return parser.parse_args()
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -16,10 +16,9 @@ Usage:
|
|||||||
"""
|
"""
|
||||||
|
|
||||||
import argparse
|
import argparse
|
||||||
import json
|
import os
|
||||||
import time
|
import time
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import Any
|
|
||||||
|
|
||||||
import torch
|
import torch
|
||||||
|
|
||||||
@@ -30,16 +29,6 @@ from vllm_omni.outputs import OmniRequestOutput
|
|||||||
from vllm_omni.platforms import current_omni_platform
|
from vllm_omni.platforms import current_omni_platform
|
||||||
|
|
||||||
|
|
||||||
def parse_profiler_config(value: str) -> dict[str, Any]:
|
|
||||||
try:
|
|
||||||
config = json.loads(value)
|
|
||||||
except json.JSONDecodeError as e:
|
|
||||||
raise argparse.ArgumentTypeError(f"--profiler-config must be valid JSON: {e}") from e
|
|
||||||
if not isinstance(config, dict):
|
|
||||||
raise argparse.ArgumentTypeError("--profiler-config must be a JSON object")
|
|
||||||
return config
|
|
||||||
|
|
||||||
|
|
||||||
def parse_args() -> argparse.Namespace:
|
def parse_args() -> argparse.Namespace:
|
||||||
parser = argparse.ArgumentParser(
|
parser = argparse.ArgumentParser(
|
||||||
description="Generate a talking-head video from a reference image and audio (Wan2.2 S2V)."
|
description="Generate a talking-head video from a reference image and audio (Wan2.2 S2V)."
|
||||||
@@ -98,12 +87,6 @@ def parse_args() -> argparse.Namespace:
|
|||||||
default=3.0,
|
default=3.0,
|
||||||
help="Scheduler flow shift (default: 3.0).",
|
help="Scheduler flow shift (default: 3.0).",
|
||||||
)
|
)
|
||||||
parser.add_argument(
|
|
||||||
"--boundary-ratio",
|
|
||||||
type=float,
|
|
||||||
default=None,
|
|
||||||
help="Boundary split ratio for low/high DiT stages (S2V). Default varies by model.",
|
|
||||||
)
|
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
"--output",
|
"--output",
|
||||||
type=str,
|
type=str,
|
||||||
@@ -195,10 +178,30 @@ def parse_args() -> argparse.Namespace:
|
|||||||
help="Enable cache-dit summary logging after diffusion forward passes.",
|
help="Enable cache-dit summary logging after diffusion forward passes.",
|
||||||
)
|
)
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
"--profiler-config",
|
"--profile-dir",
|
||||||
type=parse_profiler_config,
|
type=str,
|
||||||
default=None,
|
default=None,
|
||||||
help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
|
help="Enable torch profiler and save traces to this directory.",
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--profile-record-shapes",
|
||||||
|
action="store_true",
|
||||||
|
help="Record tensor shapes in profiler (increases trace size).",
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--profile-with-stack",
|
||||||
|
action="store_true",
|
||||||
|
help="Record stack traces in profiler (increases overhead).",
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--profile-with-memory",
|
||||||
|
action="store_true",
|
||||||
|
help="Profile memory usage.",
|
||||||
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--profile-with-flops",
|
||||||
|
action="store_true",
|
||||||
|
help="Estimate FLOPs for operations.",
|
||||||
)
|
)
|
||||||
return parser.parse_args()
|
return parser.parse_args()
|
||||||
|
|
||||||
@@ -236,9 +239,23 @@ def main():
|
|||||||
ulysses_degree=args.ulysses_degree,
|
ulysses_degree=args.ulysses_degree,
|
||||||
)
|
)
|
||||||
|
|
||||||
profiler_enabled = args.profiler_config is not None
|
# Check if profiling is requested via CLI args or environment variable
|
||||||
|
profile_dir = args.profile_dir or os.getenv("VLLM_TORCH_PROFILER_DIR")
|
||||||
|
profiler_enabled = bool(profile_dir)
|
||||||
|
profiler_config = None
|
||||||
|
if profiler_enabled:
|
||||||
|
from vllm.config import ProfilerConfig
|
||||||
|
|
||||||
omni_kwargs = dict(
|
profiler_config = ProfilerConfig(
|
||||||
|
profiler="torch",
|
||||||
|
torch_profiler_dir=profile_dir,
|
||||||
|
torch_profiler_record_shapes=args.profile_record_shapes,
|
||||||
|
torch_profiler_with_stack=args.profile_with_stack,
|
||||||
|
torch_profiler_with_memory=args.profile_with_memory,
|
||||||
|
torch_profiler_with_flops=args.profile_with_flops,
|
||||||
|
)
|
||||||
|
|
||||||
|
omni = Omni(
|
||||||
model=args.model,
|
model=args.model,
|
||||||
model_class_name="WanS2VPipeline",
|
model_class_name="WanS2VPipeline",
|
||||||
flow_shift=args.flow_shift,
|
flow_shift=args.flow_shift,
|
||||||
@@ -252,12 +269,8 @@ def main():
|
|||||||
cache_backend=args.cache_backend,
|
cache_backend=args.cache_backend,
|
||||||
cache_config=cache_config,
|
cache_config=cache_config,
|
||||||
enable_cache_dit_summary=args.enable_cache_dit_summary,
|
enable_cache_dit_summary=args.enable_cache_dit_summary,
|
||||||
profiler_config=args.profiler_config,
|
profiler_config=profiler_config,
|
||||||
)
|
)
|
||||||
if args.boundary_ratio is not None:
|
|
||||||
omni_kwargs["boundary_ratio"] = args.boundary_ratio
|
|
||||||
|
|
||||||
omni = Omni(**omni_kwargs)
|
|
||||||
|
|
||||||
# Print generation configuration
|
# Print generation configuration
|
||||||
print(f"\n{'=' * 60}")
|
print(f"\n{'=' * 60}")
|
||||||
@@ -312,16 +325,13 @@ def main():
|
|||||||
if profiler_enabled:
|
if profiler_enabled:
|
||||||
print("\n[Profiler] Stopping profiler and collecting results...")
|
print("\n[Profiler] Stopping profiler and collecting results...")
|
||||||
profile_results = omni.stop_profile()
|
profile_results = omni.stop_profile()
|
||||||
if profile_results and isinstance(profile_results, dict):
|
if profile_results and isinstance(profile_results, list):
|
||||||
traces = profile_results.get("traces", [])
|
|
||||||
print("\n" + "=" * 60)
|
print("\n" + "=" * 60)
|
||||||
print("PROFILING RESULTS:")
|
print("PROFILING RESULTS:")
|
||||||
for rank, trace in enumerate(traces):
|
for i, result in enumerate(profile_results):
|
||||||
print(f"\nRank {rank}:")
|
print(f"\nStage {i}:")
|
||||||
if trace:
|
if result:
|
||||||
print(f" • Trace: {trace}")
|
print(f" {result}")
|
||||||
if not traces:
|
|
||||||
print(" No traces collected.")
|
|
||||||
print("=" * 60)
|
print("=" * 60)
|
||||||
else:
|
else:
|
||||||
print("[Profiler] No valid profiling data returned.")
|
print("[Profiler] No valid profiling data returned.")
|
||||||
|
|||||||
@@ -307,6 +307,9 @@ def parse_args() -> argparse.Namespace:
|
|||||||
help="Supplementary auxiliary text encoder parameters model name or path (especially for Hidream-l1-full).",
|
help="Supplementary auxiliary text encoder parameters model name or path (especially for Hidream-l1-full).",
|
||||||
)
|
)
|
||||||
current_omni_platform.pre_register_and_update(parser)
|
current_omni_platform.pre_register_and_update(parser)
|
||||||
|
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||||
|
|
||||||
|
nullify_stage_engine_defaults(parser)
|
||||||
return parser.parse_args()
|
return parser.parse_args()
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -10,6 +10,7 @@ import soundfile as sf
|
|||||||
from vllm import SamplingParams
|
from vllm import SamplingParams
|
||||||
from vllm.multimodal.media.audio import load_audio
|
from vllm.multimodal.media.audio import load_audio
|
||||||
|
|
||||||
|
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||||
from vllm_omni.entrypoints.omni import Omni
|
from vllm_omni.entrypoints.omni import Omni
|
||||||
from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer
|
from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer
|
||||||
from vllm_omni.model_executor.models.cosyvoice3.utils import extract_text_token
|
from vllm_omni.model_executor.models.cosyvoice3.utils import extract_text_token
|
||||||
@@ -64,6 +65,7 @@ def run_e2e():
|
|||||||
required=True,
|
required=True,
|
||||||
help="Path to tokenizer directory (e.g., <model_path>/CosyVoice-BlankEN).",
|
help="Path to tokenizer directory (e.g., <model_path>/CosyVoice-BlankEN).",
|
||||||
)
|
)
|
||||||
|
nullify_stage_engine_defaults(parser)
|
||||||
args = parser.parse_args()
|
args = parser.parse_args()
|
||||||
# Ensure tokenizer directory exists
|
# Ensure tokenizer directory exists
|
||||||
if not os.path.exists(args.tokenizer):
|
if not os.path.exists(args.tokenizer):
|
||||||
|
|||||||
@@ -26,14 +26,16 @@ import torch
|
|||||||
|
|
||||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||||
|
|
||||||
|
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||||
|
|
||||||
from vllm_omni import AsyncOmni, Omni
|
from vllm_omni import AsyncOmni, Omni
|
||||||
|
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||||
from vllm_omni.model_executor.models.fish_speech.dac_utils import DAC_HOP_LENGTH, DAC_SAMPLE_RATE
|
from vllm_omni.model_executor.models.fish_speech.dac_utils import DAC_HOP_LENGTH, DAC_SAMPLE_RATE
|
||||||
from vllm_omni.model_executor.models.fish_speech.prompt_utils import (
|
from vllm_omni.model_executor.models.fish_speech.prompt_utils import (
|
||||||
build_fish_text_only_prompt_ids,
|
build_fish_text_only_prompt_ids,
|
||||||
estimate_fish_voice_clone_prompt_len_from_normalized,
|
estimate_fish_voice_clone_prompt_len_from_normalized,
|
||||||
normalize_fish_voice_clone_texts,
|
normalize_fish_voice_clone_texts,
|
||||||
)
|
)
|
||||||
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
|
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
@@ -209,7 +211,7 @@ async def main_streaming(args):
|
|||||||
|
|
||||||
|
|
||||||
def parse_args():
|
def parse_args():
|
||||||
parser = TrackingArgumentParser(
|
parser = FlexibleArgumentParser(
|
||||||
description="Fish Speech S2 Pro offline inference with vLLM Omni",
|
description="Fish Speech S2 Pro offline inference with vLLM Omni",
|
||||||
)
|
)
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
@@ -265,6 +267,7 @@ def parse_args():
|
|||||||
default=False,
|
default=False,
|
||||||
help="Stream audio chunks as they arrive via AsyncOmni.",
|
help="Stream audio chunks as they arrive via AsyncOmni.",
|
||||||
)
|
)
|
||||||
|
nullify_stage_engine_defaults(parser)
|
||||||
return parser.parse_args()
|
return parser.parse_args()
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -30,9 +30,10 @@ from pathlib import Path
|
|||||||
import numpy as np
|
import numpy as np
|
||||||
import soundfile as sf
|
import soundfile as sf
|
||||||
import torch
|
import torch
|
||||||
|
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||||
|
|
||||||
from vllm_omni import Omni
|
from vllm_omni import Omni
|
||||||
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
|
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||||
|
|
||||||
SAMPLE_RATE = 24_000
|
SAMPLE_RATE = 24_000
|
||||||
DEFAULT_TEXTS = (
|
DEFAULT_TEXTS = (
|
||||||
@@ -42,7 +43,7 @@ DEFAULT_TEXTS = (
|
|||||||
|
|
||||||
|
|
||||||
def parse_args():
|
def parse_args():
|
||||||
parser = TrackingArgumentParser(description="Offline higgs-audio v2 inference")
|
parser = FlexibleArgumentParser(description="Offline higgs-audio v2 inference")
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
"--model",
|
"--model",
|
||||||
type=str,
|
type=str,
|
||||||
@@ -87,6 +88,7 @@ def parse_args():
|
|||||||
default=None,
|
default=None,
|
||||||
help="Transcript of the reference clip. Required when --ref-audio is set.",
|
help="Transcript of the reference clip. Required when --ref-audio is set.",
|
||||||
)
|
)
|
||||||
|
nullify_stage_engine_defaults(parser)
|
||||||
return parser.parse_args()
|
return parser.parse_args()
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -1,204 +0,0 @@
|
|||||||
# SPDX-License-Identifier: Apache-2.0
|
|
||||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
||||||
"""Offline higgs-audio v3 inference example.
|
|
||||||
|
|
||||||
Runs Stage 0 (Qwen3 talker) + Stage 1 (HiggsAudio codec) end-to-end
|
|
||||||
through the vLLM-Omni engine without going through the HTTP server, and
|
|
||||||
saves a 24 kHz mono WAV per prompt.
|
|
||||||
|
|
||||||
Example (plain TTS):
|
|
||||||
|
|
||||||
python examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py \\
|
|
||||||
--texts "Hello world." \\
|
|
||||||
--output-dir results/higgs_v3_wavs
|
|
||||||
|
|
||||||
Example (voice clone):
|
|
||||||
|
|
||||||
python examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py \\
|
|
||||||
--texts "Hello world." \\
|
|
||||||
--ref-audio path/to/reference.wav \\
|
|
||||||
--ref-text "Transcript of the reference clip." \\
|
|
||||||
--output-dir results/higgs_v3_wavs
|
|
||||||
"""
|
|
||||||
|
|
||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
import os
|
|
||||||
|
|
||||||
os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0")
|
|
||||||
os.environ.setdefault("VLLM_MOE_USE_DEEP_GEMM", "0")
|
|
||||||
|
|
||||||
import re
|
|
||||||
import time
|
|
||||||
from pathlib import Path
|
|
||||||
|
|
||||||
import numpy as np
|
|
||||||
import soundfile as sf
|
|
||||||
import torch
|
|
||||||
|
|
||||||
from vllm_omni import Omni
|
|
||||||
|
|
||||||
SAMPLE_RATE = 24_000
|
|
||||||
DEFAULT_TEXTS = (
|
|
||||||
"Hello world.",
|
|
||||||
"The quick brown fox jumps over the lazy dog.",
|
|
||||||
"Today is a beautiful day for a walk in the park.",
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
def parse_args():
|
|
||||||
import argparse
|
|
||||||
|
|
||||||
parser = argparse.ArgumentParser(description="Offline higgs-audio v3 inference")
|
|
||||||
parser.add_argument(
|
|
||||||
"--model",
|
|
||||||
type=str,
|
|
||||||
default="bosonai/higgs-audio-v3-tts-4b",
|
|
||||||
help="Stage-0 talker model id or path.",
|
|
||||||
)
|
|
||||||
parser.add_argument(
|
|
||||||
"--texts",
|
|
||||||
type=str,
|
|
||||||
nargs="+",
|
|
||||||
default=list(DEFAULT_TEXTS),
|
|
||||||
help="One or more plain-text prompts to synthesize.",
|
|
||||||
)
|
|
||||||
parser.add_argument(
|
|
||||||
"--output-dir",
|
|
||||||
type=str,
|
|
||||||
default="results/higgs_v3_wavs",
|
|
||||||
help="Directory to write per-prompt WAV files.",
|
|
||||||
)
|
|
||||||
parser.add_argument(
|
|
||||||
"--deploy-config",
|
|
||||||
type=str,
|
|
||||||
default=None,
|
|
||||||
help="Override the deploy config path.",
|
|
||||||
)
|
|
||||||
parser.add_argument(
|
|
||||||
"--max-new-tokens",
|
|
||||||
type=int,
|
|
||||||
default=2048,
|
|
||||||
help="Cap on Stage-0 codec frames per request.",
|
|
||||||
)
|
|
||||||
parser.add_argument(
|
|
||||||
"--ref-audio",
|
|
||||||
type=str,
|
|
||||||
default=None,
|
|
||||||
help="Reference audio for voice clone (WAV/FLAC/MP3 path).",
|
|
||||||
)
|
|
||||||
parser.add_argument(
|
|
||||||
"--ref-text",
|
|
||||||
type=str,
|
|
||||||
default=None,
|
|
||||||
help="Transcript of the reference audio. Optional but improves fidelity.",
|
|
||||||
)
|
|
||||||
return parser.parse_args()
|
|
||||||
|
|
||||||
|
|
||||||
def _slugify(text: str) -> str:
|
|
||||||
slug = re.sub(r"\s+", "_", text.strip().lower())
|
|
||||||
slug = re.sub(r"[^a-z0-9_]+", "", slug)
|
|
||||||
return slug[:48] or "prompt"
|
|
||||||
|
|
||||||
|
|
||||||
def _extract_pcm(multimodal_output: dict) -> torch.Tensor:
|
|
||||||
audio = multimodal_output.get("model_outputs")
|
|
||||||
if audio is None:
|
|
||||||
audio = multimodal_output.get("audio")
|
|
||||||
if audio is None:
|
|
||||||
raise ValueError(f"no audio key in multimodal_output: {list(multimodal_output.keys())}")
|
|
||||||
if isinstance(audio, list):
|
|
||||||
valid = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio if a is not None]
|
|
||||||
if not valid:
|
|
||||||
raise ValueError("audio list is empty")
|
|
||||||
return torch.cat(valid, dim=0) if len(valid) > 1 else valid[0]
|
|
||||||
return torch.as_tensor(audio).float().cpu().reshape(-1)
|
|
||||||
|
|
||||||
|
|
||||||
def _pcm_to_int16(pcm: torch.Tensor) -> np.ndarray:
|
|
||||||
arr = pcm.numpy()
|
|
||||||
if arr.dtype.kind == "f":
|
|
||||||
arr = np.clip(arr, -1.0, 1.0)
|
|
||||||
arr = (arr * 32767.0).astype(np.int16)
|
|
||||||
else:
|
|
||||||
arr = arr.astype(np.int16)
|
|
||||||
return arr
|
|
||||||
|
|
||||||
|
|
||||||
def main():
|
|
||||||
args = parse_args()
|
|
||||||
output_dir = Path(args.output_dir)
|
|
||||||
output_dir.mkdir(parents=True, exist_ok=True)
|
|
||||||
|
|
||||||
engine = Omni(model=args.model, deploy_config=args.deploy_config, trust_remote_code=True)
|
|
||||||
|
|
||||||
from transformers import AutoTokenizer
|
|
||||||
|
|
||||||
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_tokenizer import (
|
|
||||||
HiggsAudioV3TokenizerAdapter,
|
|
||||||
apply_delay_pattern,
|
|
||||||
encode_reference_audio,
|
|
||||||
)
|
|
||||||
|
|
||||||
tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
|
|
||||||
adapter = HiggsAudioV3TokenizerAdapter(tokenizer)
|
|
||||||
|
|
||||||
# Load and encode reference audio once (if voice cloning)
|
|
||||||
ref_codes_delayed: torch.Tensor | None = None
|
|
||||||
if args.ref_audio is not None:
|
|
||||||
ref_wav, ref_sr = sf.read(args.ref_audio, always_2d=False)
|
|
||||||
if ref_wav.ndim == 2:
|
|
||||||
ref_wav = ref_wav.mean(axis=1)
|
|
||||||
ref_codes_raw = encode_reference_audio(ref_wav, int(ref_sr))
|
|
||||||
ref_codes_delayed = apply_delay_pattern(ref_codes_raw)
|
|
||||||
print(f"Reference : {args.ref_audio}")
|
|
||||||
print(f"Ref codes : {ref_codes_raw.shape[0]} frames -> {ref_codes_delayed.shape[0]} delayed")
|
|
||||||
if args.ref_text:
|
|
||||||
print(f"Ref text : {args.ref_text}")
|
|
||||||
|
|
||||||
print(f"Model : {args.model}")
|
|
||||||
print(f"Prompts : {len(args.texts)}")
|
|
||||||
print(f"Output dir : {output_dir}")
|
|
||||||
print(f"Voice clone : {'yes' if ref_codes_delayed is not None else 'no'}")
|
|
||||||
|
|
||||||
total_elapsed = 0.0
|
|
||||||
total_dur = 0.0
|
|
||||||
for text in args.texts:
|
|
||||||
if ref_codes_delayed is not None:
|
|
||||||
prompt_ids = adapter.build_prompt(
|
|
||||||
text,
|
|
||||||
num_ref_tokens=int(ref_codes_delayed.shape[0]),
|
|
||||||
reference_text=args.ref_text,
|
|
||||||
)
|
|
||||||
prompt = {
|
|
||||||
"prompt_token_ids": prompt_ids,
|
|
||||||
"additional_information": {
|
|
||||||
"audio_input_ids": ref_codes_delayed.to(torch.long),
|
|
||||||
"audio_input_ids_mask": torch.ones(ref_codes_delayed.shape[0], dtype=torch.bool),
|
|
||||||
},
|
|
||||||
}
|
|
||||||
else:
|
|
||||||
prompt_ids = adapter.build_prompt(text)
|
|
||||||
prompt = {"prompt_token_ids": prompt_ids}
|
|
||||||
|
|
||||||
t_start = time.perf_counter()
|
|
||||||
outputs = engine.generate([prompt])
|
|
||||||
elapsed = time.perf_counter() - t_start
|
|
||||||
total_elapsed += elapsed
|
|
||||||
|
|
||||||
mm = outputs[0].outputs[0].multimodal_output
|
|
||||||
pcm = _extract_pcm(mm)
|
|
||||||
slug = _slugify(text)
|
|
||||||
out_path = output_dir / f"{slug}.wav"
|
|
||||||
sf.write(str(out_path), _pcm_to_int16(pcm), SAMPLE_RATE, format="WAV", subtype="PCM_16")
|
|
||||||
dur = pcm.numel() / SAMPLE_RATE
|
|
||||||
total_dur += dur
|
|
||||||
print(f" {slug:<50} dur={dur:5.2f}s -> {out_path}")
|
|
||||||
|
|
||||||
rtf = total_elapsed / total_dur if total_dur > 0 else float("inf")
|
|
||||||
print(f"Total infer : {total_elapsed:.2f}s total audio: {total_dur:.2f}s RTF: {rtf:.3f}")
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
main()
|
|
||||||
@@ -5,11 +5,11 @@ from typing import Any
|
|||||||
|
|
||||||
import soundfile as sf
|
import soundfile as sf
|
||||||
import torch
|
import torch
|
||||||
|
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||||
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
|
|
||||||
|
|
||||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||||
|
|
||||||
|
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||||
from vllm_omni.entrypoints.omni import Omni
|
from vllm_omni.entrypoints.omni import Omni
|
||||||
from vllm_omni.inputs.data import OmniTokensPrompt
|
from vllm_omni.inputs.data import OmniTokensPrompt
|
||||||
from vllm_omni.model_executor.models.ming_flash_omni.prompt_utils import (
|
from vllm_omni.model_executor.models.ming_flash_omni.prompt_utils import (
|
||||||
@@ -68,7 +68,7 @@ def save_audio(mm: dict[str, Any], output_path: str) -> None:
|
|||||||
|
|
||||||
|
|
||||||
def parse_args():
|
def parse_args():
|
||||||
parser = TrackingArgumentParser(description="Ming-flash-omni standalone talker offline e2e example")
|
parser = FlexibleArgumentParser(description="Ming-flash-omni standalone talker offline e2e example")
|
||||||
parser.add_argument("--model", type=str, default=MODEL_NAME, help="Model name or local path.")
|
parser.add_argument("--model", type=str, default=MODEL_NAME, help="Model name or local path.")
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
"--deploy-config",
|
"--deploy-config",
|
||||||
@@ -89,6 +89,7 @@ def parse_args():
|
|||||||
parser.add_argument("--init-timeout", type=int, default=600, help="Engine init timeout in seconds.")
|
parser.add_argument("--init-timeout", type=int, default=600, help="Engine init timeout in seconds.")
|
||||||
parser.add_argument("--stage-init-timeout", type=int, default=300, help="Single stage init timeout in seconds.")
|
parser.add_argument("--stage-init-timeout", type=int, default=300, help="Single stage init timeout in seconds.")
|
||||||
|
|
||||||
|
nullify_stage_engine_defaults(parser)
|
||||||
return parser.parse_args()
|
return parser.parse_args()
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -1,377 +0,0 @@
|
|||||||
"""Offline inference example for MOSS-TTS variants.
|
|
||||||
|
|
||||||
Covers all five models:
|
|
||||||
- OpenMOSS-Team/MOSS-TTS (8B, general TTS)
|
|
||||||
- OpenMOSS-Team/MOSS-TTSD-v1.0 (8B, dialogue TTS)
|
|
||||||
- OpenMOSS-Team/MOSS-SoundEffect (8B, sound effect synthesis)
|
|
||||||
- OpenMOSS-Team/MOSS-TTS-Realtime (1.7B, low-latency streaming)
|
|
||||||
- OpenMOSS-Team/MOSS-VoiceGenerator(1.7B, voice design)
|
|
||||||
|
|
||||||
Usage
|
|
||||||
-----
|
|
||||||
# Standard TTS with voice cloning:
|
|
||||||
python end2end.py \\
|
|
||||||
--model OpenMOSS-Team/MOSS-TTS \\
|
|
||||||
--text "Hello, this is a MOSS-TTS test." \\
|
|
||||||
--ref-audio path/to/reference.wav \\
|
|
||||||
--output output.wav
|
|
||||||
"""
|
|
||||||
|
|
||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
import argparse
|
|
||||||
import gc
|
|
||||||
import sys
|
|
||||||
|
|
||||||
import soundfile as sf
|
|
||||||
import torch
|
|
||||||
|
|
||||||
|
|
||||||
def _load_ref_audio(path: str, target_sr: int = 24000) -> torch.Tensor:
|
|
||||||
"""Load and resample reference audio to (1, T) at target_sr."""
|
|
||||||
data, sr = sf.read(path, always_2d=True) # (T, C)
|
|
||||||
wav = torch.from_numpy(data.T.astype("float32")) # (C, T)
|
|
||||||
if wav.shape[0] > 1:
|
|
||||||
wav = wav.mean(dim=0, keepdim=True)
|
|
||||||
if sr != target_sr:
|
|
||||||
import torchaudio
|
|
||||||
|
|
||||||
wav = torchaudio.functional.resample(wav, sr, target_sr)
|
|
||||||
return wav # (1, T)
|
|
||||||
|
|
||||||
|
|
||||||
def _build_unified_codes(
|
|
||||||
model_id: str,
|
|
||||||
text: str,
|
|
||||||
ref_audio: torch.Tensor | None,
|
|
||||||
*,
|
|
||||||
instruction: str | None = None,
|
|
||||||
tokens: int | None = None,
|
|
||||||
quality: str | None = None,
|
|
||||||
sound_event: str | None = None,
|
|
||||||
ambient_sound: str | None = None,
|
|
||||||
language: str | None = None,
|
|
||||||
) -> tuple[list[int], torch.Tensor, int]:
|
|
||||||
"""Build the upstream MossTTSDelay unified-codes prompt.
|
|
||||||
|
|
||||||
Returns (text_token_ids, audio_codes_LxNQ, n_vq) where
|
|
||||||
``text_token_ids`` is the prefill prompt for the talker and
|
|
||||||
``audio_codes_LxNQ`` is the matching delay-pattern audio code grid
|
|
||||||
(``audio_pad_code`` everywhere except inside the reference-audio block).
|
|
||||||
"""
|
|
||||||
from transformers import AutoProcessor
|
|
||||||
|
|
||||||
proc = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
|
|
||||||
|
|
||||||
reference: list[torch.Tensor] | None = None
|
|
||||||
if ref_audio is not None:
|
|
||||||
# Pre-encode so the processor doesn't need torchaudio.load (which
|
|
||||||
# now requires torchcodec on recent torchaudio releases).
|
|
||||||
# encode_audios_from_wav expects (C, T) per item.
|
|
||||||
wav_2d = ref_audio if ref_audio.dim() == 2 else ref_audio.unsqueeze(0)
|
|
||||||
codes_list = proc.encode_audios_from_wav([wav_2d], sampling_rate=24000, n_vq=proc.model_config.n_vq)
|
|
||||||
reference = [codes_list[0]]
|
|
||||||
|
|
||||||
user_msg_kwargs: dict = {"text": text}
|
|
||||||
if reference is not None:
|
|
||||||
user_msg_kwargs["reference"] = reference
|
|
||||||
if instruction is not None:
|
|
||||||
user_msg_kwargs["instruction"] = instruction
|
|
||||||
if tokens is not None:
|
|
||||||
user_msg_kwargs["tokens"] = tokens
|
|
||||||
if quality is not None:
|
|
||||||
user_msg_kwargs["quality"] = quality
|
|
||||||
if sound_event is not None:
|
|
||||||
user_msg_kwargs["sound_event"] = sound_event
|
|
||||||
if ambient_sound is not None:
|
|
||||||
user_msg_kwargs["ambient_sound"] = ambient_sound
|
|
||||||
if language is not None:
|
|
||||||
user_msg_kwargs["language"] = language
|
|
||||||
|
|
||||||
user_msg = proc.build_user_message(**user_msg_kwargs)
|
|
||||||
batch = proc(conversations=[[user_msg]], mode="generation")
|
|
||||||
unified = batch["input_ids"][0] # (L, 1+n_vq)
|
|
||||||
n_vq = int(unified.shape[-1] - 1)
|
|
||||||
text_ids = unified[:, 0].tolist()
|
|
||||||
audio_codes = unified[:, 1:].contiguous().to(torch.int64) # (L, n_vq)
|
|
||||||
|
|
||||||
# Free the processor (incl. ~1.6B-param audio_tokenizer on CPU) before
|
|
||||||
# we instantiate Omni — we only needed it for the one-shot encode.
|
|
||||||
del proc
|
|
||||||
gc.collect()
|
|
||||||
|
|
||||||
return text_ids, audio_codes, n_vq
|
|
||||||
|
|
||||||
|
|
||||||
def _build_realtime_prompt(
|
|
||||||
model_id: str,
|
|
||||||
text: str,
|
|
||||||
ref_audio: torch.Tensor | None,
|
|
||||||
) -> tuple[list[int], torch.Tensor, int]:
|
|
||||||
"""Build the prefill prompt for MOSS-TTS-Realtime.
|
|
||||||
|
|
||||||
Realtime uses a different prompt format than the delay variants: a system
|
|
||||||
prompt + (optional) voice-clone context with the encoded reference audio
|
|
||||||
inlined, followed by the user text and an assistant header. The format is
|
|
||||||
a ``(L, channels+1=17)`` int grid where col 0 is text/special tokens and
|
|
||||||
cols 1..16 are RVQ codebook entries (``audio_channel_pad`` outside the
|
|
||||||
ref-audio block).
|
|
||||||
|
|
||||||
The upstream ``MossTTSRealtimeProcessor`` is shipped via
|
|
||||||
``trust_remote_code`` but is not auto-discovered by ``AutoProcessor``
|
|
||||||
(no ``processor_config.json``) — we import its module directly from the
|
|
||||||
snapshot directory and use its ``make_ensemble``/``make_user_prompt``.
|
|
||||||
"""
|
|
||||||
import importlib.util
|
|
||||||
import sys as _sys
|
|
||||||
from pathlib import Path
|
|
||||||
|
|
||||||
import numpy as np
|
|
||||||
import torch
|
|
||||||
from huggingface_hub import snapshot_download
|
|
||||||
from transformers import AutoModel, AutoTokenizer
|
|
||||||
|
|
||||||
snap_dir = Path(snapshot_download(repo_id=model_id))
|
|
||||||
proc_module_path = snap_dir / "processing_mossttsrealtime.py"
|
|
||||||
if not proc_module_path.exists():
|
|
||||||
raise FileNotFoundError(f"realtime processor module missing at {proc_module_path}")
|
|
||||||
|
|
||||||
# Load the processor module without polluting sys.path globally.
|
|
||||||
spec = importlib.util.spec_from_file_location("_moss_tts_realtime_proc", proc_module_path)
|
|
||||||
mod = importlib.util.module_from_spec(spec) # type: ignore[arg-type]
|
|
||||||
_sys.modules[spec.name] = mod # type: ignore[union-attr]
|
|
||||||
spec.loader.exec_module(mod) # type: ignore[union-attr]
|
|
||||||
|
|
||||||
tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
|
|
||||||
realtime_processor = mod.MossTTSRealtimeProcessor(tokenizer=tokenizer)
|
|
||||||
|
|
||||||
# Encode reference audio via the standalone MOSS Audio Tokenizer (CPU is
|
|
||||||
# fine — it runs once per request).
|
|
||||||
audio_tokens = None
|
|
||||||
if ref_audio is not None:
|
|
||||||
codec = AutoModel.from_pretrained("OpenMOSS-Team/MOSS-Audio-Tokenizer", trust_remote_code=True).eval()
|
|
||||||
wav_2d = ref_audio if ref_audio.dim() == 2 else ref_audio.unsqueeze(0)
|
|
||||||
with torch.no_grad():
|
|
||||||
enc = codec.batch_encode([wav_2d.squeeze(0)], num_quantizers=16)
|
|
||||||
# enc.audio_codes: (NQ, B, T); we want (T, NQ) numpy.
|
|
||||||
codes = enc.audio_codes[:, 0, : int(enc.audio_codes_lengths[0].item())]
|
|
||||||
audio_tokens = codes.transpose(0, 1).contiguous().cpu().numpy() # (T, 16)
|
|
||||||
del codec
|
|
||||||
gc.collect()
|
|
||||||
|
|
||||||
# Build the inference-style prefill (mirrors upstream
|
|
||||||
# ``MossTTSRealtimeInference._build_prefill_batch``):
|
|
||||||
# system_prompt [+ voice_clone_block]
|
|
||||||
# <|im_start|>assistant\n
|
|
||||||
# <text_tokens...> (col 0 = text, cols 1..16 = audio_channel_pad)
|
|
||||||
# <last text token> (col 1 = audio_bos at the last text position)
|
|
||||||
# Generation continues from here; the local transformer emits real audio
|
|
||||||
# codes for codebook 0 directly (no BOS frame at gen_step=0).
|
|
||||||
system_grid = realtime_processor.make_ensemble(prompt_audio_tokens=audio_tokens) # (L_sys, 17)
|
|
||||||
|
|
||||||
# Append "<|im_start|>assistant\n" header (the snapshot's make_ensemble
|
|
||||||
# omits it; inferencer.py's version includes it. We add it here.)
|
|
||||||
assistant_header_ids = tokenizer.encode("<|im_start|>assistant\n", add_special_tokens=False)
|
|
||||||
assistant_grid = np.full((len(assistant_header_ids), 17), 1024, dtype=np.int64)
|
|
||||||
assistant_grid[:, 0] = assistant_header_ids
|
|
||||||
|
|
||||||
# Append the FIRST 12 text tokens with audio_bos (1025) at the last
|
|
||||||
# included text position (audio column 0 == grid col 1). The remaining
|
|
||||||
# text tokens are fed one-per-step during generation via the talker's
|
|
||||||
# compute_logits → vLLM sampler pipeline (see ``ids.all`` in the
|
|
||||||
# additional_information). This matches upstream
|
|
||||||
# ``MossTTSRealtimeInference._build_prefill_batch`` (prefill_max_text=12).
|
|
||||||
text_ids_only = tokenizer.encode(text, add_special_tokens=False)
|
|
||||||
if not text_ids_only:
|
|
||||||
raise ValueError("Empty text after tokenisation")
|
|
||||||
PREFILL_MAX_TEXT = 12
|
|
||||||
cur_len = min(len(text_ids_only), PREFILL_MAX_TEXT)
|
|
||||||
text_grid = np.full((cur_len, 17), 1024, dtype=np.int64)
|
|
||||||
text_grid[:, 0] = text_ids_only[:cur_len]
|
|
||||||
text_grid[-1, 1] = 1025 # audio_bos_token at last *prefilled* text token
|
|
||||||
|
|
||||||
grid = np.concatenate([system_grid, assistant_grid, text_grid], axis=0)
|
|
||||||
text_ids = grid[:, 0].tolist()
|
|
||||||
audio_codes = torch.from_numpy(grid[:, 1:].astype(np.int64).copy()) # (L, 16)
|
|
||||||
|
|
||||||
# ``remaining_text_ids`` will be fed one-per-step by the talker.
|
|
||||||
remaining_text_ids = list(text_ids_only[cur_len:])
|
|
||||||
return text_ids, audio_codes, 16, remaining_text_ids
|
|
||||||
|
|
||||||
|
|
||||||
def run_tts(args: argparse.Namespace) -> None:
|
|
||||||
from vllm import SamplingParams
|
|
||||||
|
|
||||||
from vllm_omni import Omni
|
|
||||||
|
|
||||||
deploy_map = {
|
|
||||||
"OpenMOSS-Team/MOSS-TTS": "moss_tts",
|
|
||||||
"OpenMOSS-Team/MOSS-TTSD-v1.0": "moss_ttsd",
|
|
||||||
"OpenMOSS-Team/MOSS-SoundEffect": "moss_sound_effect",
|
|
||||||
"OpenMOSS-Team/MOSS-TTS-Realtime": "moss_tts_realtime",
|
|
||||||
"OpenMOSS-Team/MOSS-VoiceGenerator": "moss_voice_generator",
|
|
||||||
}
|
|
||||||
deploy_config = deploy_map.get(args.model, "moss_tts")
|
|
||||||
|
|
||||||
is_sound_effect = "SoundEffect" in args.model
|
|
||||||
is_voice_gen = "VoiceGenerator" in args.model
|
|
||||||
is_realtime = "Realtime" in args.model
|
|
||||||
|
|
||||||
ref_audio = None
|
|
||||||
builder_kwargs: dict = {}
|
|
||||||
|
|
||||||
remaining_text_ids: list[int] | None = None
|
|
||||||
if is_realtime:
|
|
||||||
# MOSS-TTS-Realtime uses a different prompt format and has its own
|
|
||||||
# per-step depth transformer; the delay-style ``MossTTSDelayProcessor``
|
|
||||||
# path doesn't apply.
|
|
||||||
if not args.ref_audio:
|
|
||||||
sys.exit("MOSS-TTS-Realtime currently requires --ref-audio (voice clone).")
|
|
||||||
ref_audio = _load_ref_audio(args.ref_audio)
|
|
||||||
text = args.text or ""
|
|
||||||
text_ids, audio_codes, n_vq, remaining_text_ids = _build_realtime_prompt(args.model, text, ref_audio)
|
|
||||||
print(
|
|
||||||
f"Prefill prompt: {len(text_ids)} tokens, {n_vq}-quantizer audio block (realtime); "
|
|
||||||
f"{len(remaining_text_ids)} text tokens to stream"
|
|
||||||
)
|
|
||||||
elif is_sound_effect:
|
|
||||||
if not args.ambient_sound:
|
|
||||||
sys.exit("MOSS-SoundEffect requires --ambient-sound")
|
|
||||||
builder_kwargs["ambient_sound"] = args.ambient_sound
|
|
||||||
# MOSS-SoundEffect emits ~12.5 codec frames per second; ``tokens``
|
|
||||||
# tells the AR loop the target duration.
|
|
||||||
if args.duration_tokens:
|
|
||||||
builder_kwargs["tokens"] = args.duration_tokens
|
|
||||||
elif args.duration_seconds:
|
|
||||||
builder_kwargs["tokens"] = max(1, int(float(args.duration_seconds) * 12.5))
|
|
||||||
text = args.text or ""
|
|
||||||
elif is_voice_gen:
|
|
||||||
# MOSS-VoiceGenerator needs an instruction (voice description) plus
|
|
||||||
# the text to synthesise; no reference audio.
|
|
||||||
if not args.instruction:
|
|
||||||
sys.exit("MOSS-VoiceGenerator requires --instruction (voice description)")
|
|
||||||
builder_kwargs["instruction"] = args.instruction
|
|
||||||
text = args.text or ""
|
|
||||||
else:
|
|
||||||
if not args.ref_audio:
|
|
||||||
sys.exit("Voice cloning requires --ref-audio")
|
|
||||||
ref_audio = _load_ref_audio(args.ref_audio)
|
|
||||||
text = args.text or ""
|
|
||||||
|
|
||||||
if not is_realtime:
|
|
||||||
text_ids, audio_codes, n_vq = _build_unified_codes(args.model, text, ref_audio, **builder_kwargs)
|
|
||||||
print(f"Prefill prompt: {len(text_ids)} tokens, {n_vq}-quantizer audio block")
|
|
||||||
|
|
||||||
omni = Omni(model=args.model, deploy_config=deploy_config, stage_init_timeout=600)
|
|
||||||
|
|
||||||
talker_sp = SamplingParams(
|
|
||||||
temperature=1.7,
|
|
||||||
top_p=0.8,
|
|
||||||
top_k=25,
|
|
||||||
max_tokens=2048,
|
|
||||||
seed=42,
|
|
||||||
detokenize=False,
|
|
||||||
)
|
|
||||||
codec_sp = SamplingParams(
|
|
||||||
temperature=0.0,
|
|
||||||
top_p=1.0,
|
|
||||||
top_k=-1,
|
|
||||||
max_tokens=65536,
|
|
||||||
seed=42,
|
|
||||||
detokenize=True,
|
|
||||||
)
|
|
||||||
|
|
||||||
additional_info: dict = {
|
|
||||||
# Reference-audio codes consumed by the talker prefill (additive
|
|
||||||
# embedding); shape is (L, n_vq) where pad_code marks non-audio
|
|
||||||
# positions. The OmniPayload schema reserves ``codes.ref`` for
|
|
||||||
# exactly this kind of conditioning input.
|
|
||||||
"codes": {"ref": audio_codes},
|
|
||||||
}
|
|
||||||
if remaining_text_ids:
|
|
||||||
# Realtime variant: feed the remaining text tokens (everything past
|
|
||||||
# the first ``PREFILL_MAX_TEXT``) one-per-step during decode via the
|
|
||||||
# talker's compute_logits → vLLM sampler one-hot trick.
|
|
||||||
additional_info["ids"] = {"all": remaining_text_ids}
|
|
||||||
|
|
||||||
request = {
|
|
||||||
"prompt_token_ids": text_ids,
|
|
||||||
"additional_information": additional_info,
|
|
||||||
}
|
|
||||||
|
|
||||||
import time as _time
|
|
||||||
|
|
||||||
chunks: list[torch.Tensor] = []
|
|
||||||
t_gen_start = _time.monotonic()
|
|
||||||
t_first_audio: float | None = None
|
|
||||||
for stage_outputs in omni.generate(request, [talker_sp, codec_sp]):
|
|
||||||
# Each yielded ``stage_outputs`` is an OmniRequestOutput; the codec
|
|
||||||
# stage attaches its waveform under ``multimodal_output["model_outputs"]``.
|
|
||||||
# Avoid ``a or b`` when ``a`` may be a tensor (truthiness raises).
|
|
||||||
mm = stage_outputs.multimodal_output or {}
|
|
||||||
wav = mm.get("audio")
|
|
||||||
if wav is None:
|
|
||||||
wav = mm.get("model_outputs")
|
|
||||||
if wav is None:
|
|
||||||
continue
|
|
||||||
if isinstance(wav, list):
|
|
||||||
non_empty = [w.reshape(-1) for w in wav if isinstance(w, torch.Tensor) and w.numel() > 0]
|
|
||||||
if non_empty and t_first_audio is None:
|
|
||||||
t_first_audio = _time.monotonic()
|
|
||||||
chunks.extend(non_empty)
|
|
||||||
elif isinstance(wav, torch.Tensor) and wav.numel() > 0:
|
|
||||||
if t_first_audio is None:
|
|
||||||
t_first_audio = _time.monotonic()
|
|
||||||
chunks.append(wav.reshape(-1))
|
|
||||||
|
|
||||||
t_gen_end = _time.monotonic()
|
|
||||||
if not chunks:
|
|
||||||
print("No audio in output — check model loading logs.")
|
|
||||||
return
|
|
||||||
|
|
||||||
wav_t = torch.cat(chunks).float().cpu()
|
|
||||||
sf.write(args.output, wav_t.numpy(), 24000)
|
|
||||||
print(f"Saved {len(wav_t) / 24000:.2f}s of audio to {args.output}")
|
|
||||||
ttfa_ms = (t_first_audio - t_gen_start) * 1000.0 if t_first_audio is not None else float("nan")
|
|
||||||
gen_total_ms = (t_gen_end - t_gen_start) * 1000.0
|
|
||||||
print(f"TTFA_MS={ttfa_ms:.1f} GEN_TOTAL_MS={gen_total_ms:.1f} AUDIO_S={len(wav_t) / 24000:.2f}")
|
|
||||||
|
|
||||||
|
|
||||||
def main() -> None:
|
|
||||||
parser = argparse.ArgumentParser(description="MOSS-TTS offline inference example")
|
|
||||||
parser.add_argument(
|
|
||||||
"--model",
|
|
||||||
default="OpenMOSS-Team/MOSS-TTS",
|
|
||||||
help="HuggingFace model ID",
|
|
||||||
)
|
|
||||||
parser.add_argument("--text", default="", help="Text to synthesise (TTS variants)")
|
|
||||||
parser.add_argument("--ref-audio", default=None, help="Reference audio for voice cloning")
|
|
||||||
parser.add_argument(
|
|
||||||
"--ambient-sound",
|
|
||||||
default=None,
|
|
||||||
help="Sound description for MOSS-SoundEffect",
|
|
||||||
)
|
|
||||||
parser.add_argument(
|
|
||||||
"--duration-tokens",
|
|
||||||
type=int,
|
|
||||||
default=None,
|
|
||||||
help="Target duration in tokens (1 s ≈ 12.5 tokens) for MOSS-SoundEffect",
|
|
||||||
)
|
|
||||||
parser.add_argument(
|
|
||||||
"--duration-seconds",
|
|
||||||
type=float,
|
|
||||||
default=None,
|
|
||||||
help="Target duration in seconds (converted via ~12.5 tokens/s) for MOSS-SoundEffect",
|
|
||||||
)
|
|
||||||
parser.add_argument(
|
|
||||||
"--instruction",
|
|
||||||
default=None,
|
|
||||||
help="Voice description for MOSS-VoiceGenerator (e.g. 'a young woman with a clear voice')",
|
|
||||||
)
|
|
||||||
parser.add_argument("--output", default="output.wav", help="Output WAV path")
|
|
||||||
args = parser.parse_args()
|
|
||||||
run_tts(args)
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
main()
|
|
||||||
@@ -38,8 +38,7 @@ from pathlib import Path
|
|||||||
import soundfile as sf
|
import soundfile as sf
|
||||||
import torch
|
import torch
|
||||||
from vllm import SamplingParams
|
from vllm import SamplingParams
|
||||||
|
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||||
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
|
|
||||||
|
|
||||||
# Prevent multiprocessing from re-importing CUDA in the wrong context.
|
# Prevent multiprocessing from re-importing CUDA in the wrong context.
|
||||||
os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
|
os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
|
||||||
@@ -156,7 +155,7 @@ def main(args) -> None:
|
|||||||
|
|
||||||
|
|
||||||
def parse_args():
|
def parse_args():
|
||||||
parser = TrackingArgumentParser(description="MOSS-TTS-Nano offline inference")
|
parser = FlexibleArgumentParser(description="MOSS-TTS-Nano offline inference")
|
||||||
parser.add_argument("--text", default="Hello, this is MOSS-TTS-Nano speaking.", help="Text to synthesize.")
|
parser.add_argument("--text", default="Hello, this is MOSS-TTS-Nano speaking.", help="Text to synthesize.")
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
"--ref-audio",
|
"--ref-audio",
|
||||||
|
|||||||
@@ -21,6 +21,7 @@ import os
|
|||||||
import numpy as np
|
import numpy as np
|
||||||
import soundfile as sf
|
import soundfile as sf
|
||||||
|
|
||||||
|
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||||
from vllm_omni.entrypoints.omni import Omni
|
from vllm_omni.entrypoints.omni import Omni
|
||||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||||
|
|
||||||
@@ -85,6 +86,7 @@ def run_e2e():
|
|||||||
default=None,
|
default=None,
|
||||||
help="Random seed for generation",
|
help="Random seed for generation",
|
||||||
)
|
)
|
||||||
|
nullify_stage_engine_defaults(parser)
|
||||||
args = parser.parse_args()
|
args = parser.parse_args()
|
||||||
|
|
||||||
if not os.path.exists(args.stage_config):
|
if not os.path.exists(args.stage_config):
|
||||||
|
|||||||
@@ -15,8 +15,10 @@ import torch
|
|||||||
|
|
||||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||||
|
|
||||||
|
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||||
|
|
||||||
from vllm_omni import AsyncOmni, Omni
|
from vllm_omni import AsyncOmni, Omni
|
||||||
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
|
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
@@ -433,7 +435,7 @@ async def main_streaming(args):
|
|||||||
|
|
||||||
|
|
||||||
def parse_args():
|
def parse_args():
|
||||||
parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
|
parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
"--query-type",
|
"--query-type",
|
||||||
"-q",
|
"-q",
|
||||||
@@ -546,6 +548,7 @@ def parse_args():
|
|||||||
help="Number of prompts per batch (default: 1, sequential).",
|
help="Number of prompts per batch (default: 1, sequential).",
|
||||||
)
|
)
|
||||||
|
|
||||||
|
nullify_stage_engine_defaults(parser)
|
||||||
return parser.parse_args()
|
return parser.parse_args()
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -12,16 +12,17 @@ from pathlib import Path
|
|||||||
|
|
||||||
import soundfile as sf
|
import soundfile as sf
|
||||||
import torch
|
import torch
|
||||||
|
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||||
|
|
||||||
from vllm_omni import Omni
|
from vllm_omni import Omni
|
||||||
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
|
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||||
|
|
||||||
REPO_ROOT = Path(__file__).resolve().parents[4]
|
REPO_ROOT = Path(__file__).resolve().parents[4]
|
||||||
SAMPLE_RATE = 48_000
|
SAMPLE_RATE = 48_000
|
||||||
|
|
||||||
|
|
||||||
def parse_args():
|
def parse_args():
|
||||||
parser = TrackingArgumentParser(description="Offline VoxCPM2 native AR inference")
|
parser = FlexibleArgumentParser(description="Offline VoxCPM2 native AR inference")
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
"--model",
|
"--model",
|
||||||
type=str,
|
type=str,
|
||||||
@@ -59,6 +60,7 @@ def parse_args():
|
|||||||
default=None,
|
default=None,
|
||||||
help="Optional transcript of --ref-audio (enables continuation mode).",
|
help="Optional transcript of --ref-audio (enables continuation mode).",
|
||||||
)
|
)
|
||||||
|
nullify_stage_engine_defaults(parser)
|
||||||
return parser.parse_args()
|
return parser.parse_args()
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -27,10 +27,11 @@ except ImportError:
|
|||||||
"pip install -e ."
|
"pip install -e ."
|
||||||
)
|
)
|
||||||
from vllm import SamplingParams
|
from vllm import SamplingParams
|
||||||
|
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||||
|
|
||||||
from vllm_omni import AsyncOmni
|
from vllm_omni import AsyncOmni
|
||||||
|
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||||
from vllm_omni.entrypoints.omni import Omni
|
from vllm_omni.entrypoints.omni import Omni
|
||||||
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
|
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
@@ -229,7 +230,7 @@ def run_non_streaming(inputs, sampling_params_list, model_name, args, output_dir
|
|||||||
|
|
||||||
|
|
||||||
def parse_args() -> Namespace:
|
def parse_args() -> Namespace:
|
||||||
parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with Voxtral TTS")
|
parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with Voxtral TTS")
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
"--model",
|
"--model",
|
||||||
type=str,
|
type=str,
|
||||||
@@ -312,6 +313,7 @@ def parse_args() -> Namespace:
|
|||||||
default=None,
|
default=None,
|
||||||
help="Quantization method (e.g. 'fp8'). Applied to the language model only.",
|
help="Quantization method (e.g. 'fp8'). Applied to the language model only.",
|
||||||
)
|
)
|
||||||
|
nullify_stage_engine_defaults(parser)
|
||||||
return parser.parse_args()
|
return parser.parse_args()
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -68,67 +68,6 @@ vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --tensor-parallel-size
|
|||||||
|
|
||||||
Or set `tensor_parallel_size` per stage in a custom deploy YAML.
|
Or set `tensor_parallel_size` per stage in a custom deploy YAML.
|
||||||
|
|
||||||
### VAE Patch Parallelism
|
|
||||||
|
|
||||||
[VAE Patch Parallelism](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism/vae_patch_parallel.html) distributes Bagel VAE **decode/encode** across multiple GPUs by splitting latent tiles. It lowers **per-GPU peak memory during VAE decode**, which helps high-resolution `text2img` / `img2img` when VAE becomes a bottleneck.
|
|
||||||
|
|
||||||
**Scope for Bagel:**
|
|
||||||
|
|
||||||
| Topology | VAE patch parallel |
|
|
||||||
| :------- | :----------------- |
|
|
||||||
| **Single-stage** (DiT only) | Supported on stage 0 (`BagelPipeline` + `DistributedAutoEncoder`) |
|
|
||||||
| **Two-stage** | Supported on **stage 1 (DiT)** only; stage 0 (Thinker) uses encoder-only VAE and is unrelated |
|
|
||||||
|
|
||||||
**Requirements:**
|
|
||||||
|
|
||||||
- `vae_patch_parallel_size > 1` and a distributed VAE (`DistributedAutoEncoder` on the DiT pipeline).
|
|
||||||
- The DiT process group must have at least `vae_patch_parallel_size` ranks. In practice this means the diffusion stage `world_size` must be ≥ 2 (commonly `tensor_parallel_size=2` on that stage).
|
|
||||||
- `vae_use_tiling` must be enabled. If you set `vae_patch_parallel_size > 1` and omit tiling, the registry auto-enables `vae_use_tiling` at startup.
|
|
||||||
|
|
||||||
VAE patch parallel **reuses the DiT process group** (`dit_group`); it does not create a separate VAE-only worker pool. It is not a substitute for single-GPU VAE tiling (`vae_pp=1`).
|
|
||||||
|
|
||||||
**Online serving (single-stage, 2 GPUs):**
|
|
||||||
|
|
||||||
```bash
|
|
||||||
CUDA_VISIBLE_DEVICES=0,1 vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \
|
|
||||||
--deploy-config vllm_omni/deploy/bagel_single_stage.yaml \
|
|
||||||
--tensor-parallel-size 2 \
|
|
||||||
--vae-patch-parallel-size 2 \
|
|
||||||
--vae-use-tiling
|
|
||||||
```
|
|
||||||
|
|
||||||
**Online serving (two-stage, VAE PP on DiT stage 1):** use a custom deploy YAML, for example:
|
|
||||||
|
|
||||||
```yaml
|
|
||||||
stages:
|
|
||||||
- stage_id: 0
|
|
||||||
devices: "0"
|
|
||||||
# Thinker (AR) — no VAE patch parallel here
|
|
||||||
|
|
||||||
- stage_id: 1
|
|
||||||
devices: "0,1"
|
|
||||||
vae_use_tiling: true
|
|
||||||
parallel_config:
|
|
||||||
tensor_parallel_size: 2
|
|
||||||
vae_patch_parallel_size: 2
|
|
||||||
```
|
|
||||||
|
|
||||||
```bash
|
|
||||||
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \
|
|
||||||
--deploy-config /path/to/bagel_vae_pp.yaml
|
|
||||||
```
|
|
||||||
|
|
||||||
**Verify it is active** (check server logs at startup):
|
|
||||||
|
|
||||||
```text
|
|
||||||
INFO ... vae_patch_parallel_size=2 requires vae_use_tiling; automatically enabling it.
|
|
||||||
```
|
|
||||||
|
|
||||||
| CLI flag | Default | Description |
|
|
||||||
| :------- | :------ | :---------- |
|
|
||||||
| `--vae-patch-parallel-size` | `1` | Number of DiT ranks used for VAE tile parallelism. Set to `2` or higher to enable. Should be ≤ DiT process group size (typically match `--tensor-parallel-size` on the diffusion stage). |
|
|
||||||
| `--vae-use-tiling` | off | Enable VAE spatial tiling. Required for VAE patch parallel (auto-enabled when `vae_patch_parallel_size > 1`). |
|
|
||||||
|
|
||||||
#### Hybrid Sharded Data Parallel (HSDP)
|
#### Hybrid Sharded Data Parallel (HSDP)
|
||||||
|
|
||||||
For larger Bagel deployments on multiple GPUs, you can enable HSDP (Hybrid Sharded Data Parallel) by modifying the stage configuration (for example, [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml)). HSDP shards transformer weights across GPUs to reduce per-GPU memory usage.
|
For larger Bagel deployments on multiple GPUs, you can enable HSDP (Hybrid Sharded Data Parallel) by modifying the stage configuration (for example, [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml)). HSDP shards transformer weights across GPUs to reduce per-GPU memory usage.
|
||||||
@@ -384,8 +323,6 @@ python openai_chat_client.py \
|
|||||||
| `stages[].gpu_memory_utilization` | per-stage | Fraction of GPU memory to use |
|
| `stages[].gpu_memory_utilization` | per-stage | Fraction of GPU memory to use |
|
||||||
| `stages[].enforce_eager` | per-stage | Disable CUDA graphs |
|
| `stages[].enforce_eager` | per-stage | Disable CUDA graphs |
|
||||||
| `stages[].tensor_parallel_size` | per-stage | TP degree for this stage |
|
| `stages[].tensor_parallel_size` | per-stage | TP degree for this stage |
|
||||||
| `stages[].parallel_config.vae_patch_parallel_size` | per-stage (DiT) | VAE tile parallelism degree (DiT stage only) |
|
|
||||||
| `stages[].vae_use_tiling` | per-stage (DiT) | Enable VAE tiling (required for VAE patch parallel) |
|
|
||||||
| `connectors` | top-level | Define available connector instances (SHM, Mooncake) |
|
| `connectors` | top-level | Define available connector instances (SHM, Mooncake) |
|
||||||
| `platforms` | top-level | Platform-specific overrides (e.g. `xpu`) |
|
| `platforms` | top-level | Platform-specific overrides (e.g. `xpu`) |
|
||||||
|
|
||||||
|
|||||||
@@ -1,132 +0,0 @@
|
|||||||
# DreamZero
|
|
||||||
|
|
||||||
DreamZero is the robot-policy serving example for the OpenPI-compatible `/v1/realtime/robot/openpi` endpoint.
|
|
||||||
|
|
||||||
## Supported checkpoints
|
|
||||||
|
|
||||||
| Model | HuggingFace repo |
|
|
||||||
|---|---|
|
|
||||||
| DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` |
|
|
||||||
|
|
||||||
## Quick start
|
|
||||||
|
|
||||||
### Start the server
|
|
||||||
|
|
||||||
```bash
|
|
||||||
bash examples/online_serving/dreamzero/run_server.sh
|
|
||||||
```
|
|
||||||
|
|
||||||
By default this launches:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
vllm serve GEAR-Dreams/DreamZero-DROID --omni --port 8000 \
|
|
||||||
--served-model-name dreamzero-droid \
|
|
||||||
--deploy-config vllm_omni/deploy/dreamzero_tp1_cfg2.yaml \
|
|
||||||
--enforce-eager --disable-log-stats
|
|
||||||
```
|
|
||||||
|
|
||||||
Override `MODEL`, `PORT`, `HOST`, `DEPLOY_CONFIG`, or `SERVED_MODEL_NAME` through the script environment if needed.
|
|
||||||
|
|
||||||
### Install example dependencies
|
|
||||||
|
|
||||||
The core `pip install -e .` setup does not include the extra packages used by the DreamZero example scripts.
|
|
||||||
|
|
||||||
- `openpi_client.py`:
|
|
||||||
`openpi-client`, `websockets`, `opencv-python`
|
|
||||||
- `droid_sim_eval_client.py`:
|
|
||||||
`mediapy`, `websockets`, `openpi-client`
|
|
||||||
|
|
||||||
The DROID sim-eval script also needs an Isaac Lab environment that provides `isaaclab`, `isaaclab_tasks`, `sim_evals`, and `gymnasium`.
|
|
||||||
|
|
||||||
If you run the DROID client on Python < 3.12, also install `typing-extensions`.
|
|
||||||
|
|
||||||
### Configure TP and CFG parallelism
|
|
||||||
|
|
||||||
The bundled DreamZero configs intentionally keep only:
|
|
||||||
|
|
||||||
| Config | Purpose |
|
|
||||||
|---|---|
|
|
||||||
| `vllm_omni/deploy/dreamzero.yaml` | Default TP=1, CFG parallel disabled |
|
|
||||||
| `vllm_omni/deploy/dreamzero_tp1_cfg2.yaml` | TP=1, CFG parallel size=2 |
|
|
||||||
|
|
||||||
For other topologies, use CLI parallelism flags and update stage 0 `devices` with `--stage-overrides`. The number of listed devices must match `tensor_parallel_size * cfg_parallel_size`.
|
|
||||||
|
|
||||||
TP=2 with CFG parallel disabled:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
vllm serve GEAR-Dreams/DreamZero-DROID --omni --port 8000 \
|
|
||||||
--served-model-name dreamzero-droid \
|
|
||||||
--deploy-config vllm_omni/deploy/dreamzero.yaml \
|
|
||||||
--tensor-parallel-size 2 \
|
|
||||||
--stage-overrides '{"0": {"devices": "0,1"}}' \
|
|
||||||
--enforce-eager --disable-log-stats
|
|
||||||
```
|
|
||||||
|
|
||||||
TP=2 with CFG parallel size=2:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
vllm serve GEAR-Dreams/DreamZero-DROID --omni --port 8000 \
|
|
||||||
--served-model-name dreamzero-droid \
|
|
||||||
--deploy-config vllm_omni/deploy/dreamzero.yaml \
|
|
||||||
--tensor-parallel-size 2 \
|
|
||||||
--cfg-parallel-size 2 \
|
|
||||||
--stage-overrides '{"0": {"devices": "0,1,2,3"}}' \
|
|
||||||
--enforce-eager --disable-log-stats
|
|
||||||
```
|
|
||||||
|
|
||||||
### Download example assets
|
|
||||||
|
|
||||||
The OpenPI client and DROID sim-eval example expect the three camera MP4 files in `outputs/dreamzero/assets`.
|
|
||||||
|
|
||||||
```bash
|
|
||||||
hf download YangshenDeng/vllm-omni-dreamzero-assets --repo-type dataset --local-dir outputs/dreamzero/assets
|
|
||||||
```
|
|
||||||
|
|
||||||
### Run the OpenPI client
|
|
||||||
|
|
||||||
```bash
|
|
||||||
python examples/online_serving/dreamzero/openpi_client.py \
|
|
||||||
--host 127.0.0.1 \
|
|
||||||
--port 8000 \
|
|
||||||
--video-dir outputs/dreamzero/assets
|
|
||||||
```
|
|
||||||
|
|
||||||
This client uses downloaded example videos and talks to the OpenPI websocket server.
|
|
||||||
|
|
||||||
### Run DROID sim eval
|
|
||||||
|
|
||||||
```bash
|
|
||||||
${ISAACLAB_LAUNCHER} -p examples/online_serving/dreamzero/droid_sim_eval_client.py \
|
|
||||||
--host 127.0.0.1 \
|
|
||||||
--port 8000 \
|
|
||||||
--scene 1 \
|
|
||||||
--episodes 1 \
|
|
||||||
--headless \
|
|
||||||
--device cuda:0
|
|
||||||
```
|
|
||||||
|
|
||||||
Set `ISAACLAB_LAUNCHER=path/to/isaaclab.sh` from the vLLM-Omni repository root before running the command.
|
|
||||||
This launches Isaac Lab / sim-evals and runs the DROID benchmark loop against the same websocket endpoint.
|
|
||||||
|
|
||||||
### Export comparison videos offline
|
|
||||||
|
|
||||||
```bash
|
|
||||||
python examples/offline_inference/dreamzero/export_prediction_video.py \
|
|
||||||
--deploy-config vllm_omni/deploy/dreamzero_tp1_cfg2.yaml \
|
|
||||||
--save-input-video \
|
|
||||||
--save-gif
|
|
||||||
```
|
|
||||||
|
|
||||||
The export script uses local `Omni` inference, not the websocket server. It writes the input rollout video and the predicted output video artifacts for side-by-side inspection.
|
|
||||||
|
|
||||||
### MolmoSpace demo
|
|
||||||
|
|
||||||
```bash
|
|
||||||
python examples/online_serving/dreamzero/molmospace_dreamzero_eval_demo.py \
|
|
||||||
--host 127.0.0.1 \
|
|
||||||
--port 8000 \
|
|
||||||
--benchmark_dir /path/to/benchmark \
|
|
||||||
--output_dir /path/to/output
|
|
||||||
```
|
|
||||||
|
|
||||||
This demo adapts DreamZero to the MolmoSpace-style remote policy eval loop.
|
|
||||||
@@ -1,819 +0,0 @@
|
|||||||
#!/usr/bin/env python3
|
|
||||||
# SPDX-License-Identifier: Apache-2.0
|
|
||||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
||||||
|
|
||||||
"""Run one or more DROID sim-eval rollouts against the vLLM DreamZero server.
|
|
||||||
|
|
||||||
This script is the vLLM/OpenPI adaptation of the upstream DreamZero
|
|
||||||
`eval_utils/run_sim_eval.py` sim-eval client.
|
|
||||||
|
|
||||||
Behavior intentionally kept close to upstream:
|
|
||||||
- same DROID observation extraction (`external_cam`, `external_cam_2`,
|
|
||||||
`wrist_cam`, joint position, gripper position)
|
|
||||||
- same resize-with-pad preprocessing to `(180, 320)`
|
|
||||||
- same `open_loop_horizon=8`
|
|
||||||
- same gripper binarization rule (`> 0.5 -> 1`, else `0`)
|
|
||||||
- same per-scene language prompts
|
|
||||||
|
|
||||||
Unlike upstream DreamZero, vLLM serves the compatible websocket policy endpoint
|
|
||||||
at `/v1/realtime/robot/openpi`, so this script includes the path suffix in the
|
|
||||||
client URI.
|
|
||||||
|
|
||||||
Run this script through Isaac Lab's launcher from the vLLM-Omni repository
|
|
||||||
root, for example:
|
|
||||||
|
|
||||||
"${ISAACLAB_LAUNCHER}" -p \
|
|
||||||
examples/online_serving/dreamzero/droid_sim_eval_client.py \
|
|
||||||
--host 127.0.0.1 \
|
|
||||||
--port 8000 \
|
|
||||||
--scene 1 \
|
|
||||||
--episodes 1 \
|
|
||||||
--headless \
|
|
||||||
--device cuda:1
|
|
||||||
"""
|
|
||||||
|
|
||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
import argparse
|
|
||||||
import json
|
|
||||||
import logging
|
|
||||||
import sys
|
|
||||||
import time
|
|
||||||
import uuid
|
|
||||||
from dataclasses import dataclass
|
|
||||||
from datetime import UTC, datetime
|
|
||||||
from pathlib import Path
|
|
||||||
from typing import Any
|
|
||||||
|
|
||||||
import numpy as np
|
|
||||||
import torch
|
|
||||||
|
|
||||||
try:
|
|
||||||
import mediapy
|
|
||||||
except ImportError as exc: # pragma: no cover - runtime dependency guard
|
|
||||||
raise ImportError("Optional DROID sim-eval client requires `mediapy`.") from exc
|
|
||||||
|
|
||||||
try:
|
|
||||||
from typing import override
|
|
||||||
except ImportError:
|
|
||||||
try:
|
|
||||||
from typing_extensions import override
|
|
||||||
except ImportError as exc: # pragma: no cover - runtime dependency guard
|
|
||||||
raise ImportError("Optional DROID sim-eval client requires `typing-extensions` on Python < 3.12.") from exc
|
|
||||||
|
|
||||||
try:
|
|
||||||
import websockets.sync.client
|
|
||||||
except ImportError as exc: # pragma: no cover - runtime dependency guard
|
|
||||||
raise ImportError("Optional DROID sim-eval client requires `websockets`.") from exc
|
|
||||||
|
|
||||||
# NOTE:
|
|
||||||
# This directory already contains a local file named `openpi_client.py`.
|
|
||||||
# However, what we want here is the *installed* `openpi_client` package from
|
|
||||||
# upstream OpenPI, not the sibling example file. When a script is executed
|
|
||||||
# directly, Python often puts the script directory at `sys.path[0]`, which
|
|
||||||
# would cause `import openpi_client` to resolve to the local example file and
|
|
||||||
# create a circular import.
|
|
||||||
#
|
|
||||||
# To avoid that ambiguity, temporarily remove the current example directory
|
|
||||||
# from the front of `sys.path`, import the real package, and then restore the
|
|
||||||
# path afterwards.
|
|
||||||
example_dir = str(Path(__file__).resolve().parent)
|
|
||||||
removed_path = False
|
|
||||||
if sys.path and sys.path[0] == example_dir:
|
|
||||||
sys.path.pop(0)
|
|
||||||
removed_path = True
|
|
||||||
try:
|
|
||||||
from openpi_client import image_tools, msgpack_numpy
|
|
||||||
from openpi_client.base_policy import BasePolicy
|
|
||||||
except ImportError as exc: # pragma: no cover - runtime dependency guard
|
|
||||||
raise ImportError("Optional DROID sim-eval client requires the `openpi-client` package.") from exc
|
|
||||||
finally:
|
|
||||||
if removed_path:
|
|
||||||
sys.path.insert(0, example_dir)
|
|
||||||
|
|
||||||
# -----------------------------------------------------------------------------
|
|
||||||
# Constants
|
|
||||||
# -----------------------------------------------------------------------------
|
|
||||||
#
|
|
||||||
# These values intentionally mirror the upstream DreamZero sim-eval client
|
|
||||||
# where possible. The important distinction is:
|
|
||||||
#
|
|
||||||
# - ACTION_HORIZON = 24
|
|
||||||
# The model returns 24 future actions per inference call.
|
|
||||||
# - DEFAULT_OPEN_LOOP_HORIZON = 8
|
|
||||||
# The sim client only executes the first 8 actions locally before asking
|
|
||||||
# the server to replan from a fresh observation.
|
|
||||||
#
|
|
||||||
# So a single server call predicts 24x8 actions, but the rollout consumes only
|
|
||||||
# 8 of them before replanning.
|
|
||||||
PING_INTERVAL_SECS = 60
|
|
||||||
PING_TIMEOUT_SECS = 600
|
|
||||||
DEFAULT_PATH = "/v1/realtime/robot/openpi"
|
|
||||||
DEFAULT_OPEN_LOOP_HORIZON = 8
|
|
||||||
ACTION_HORIZON = 24
|
|
||||||
ACTION_DIM = 8
|
|
||||||
DEFAULT_OUTPUT_ROOT = Path("runs") / "dreamzero_sim_eval"
|
|
||||||
SCENE_PROMPTS = {
|
|
||||||
1: "put the cube in the bowl",
|
|
||||||
2: "pick up the can and put it in the mug",
|
|
||||||
3: "put the banana in the bin",
|
|
||||||
}
|
|
||||||
|
|
||||||
|
|
||||||
def _decode_action_response(response: bytes | str) -> np.ndarray:
|
|
||||||
if isinstance(response, str):
|
|
||||||
raise RuntimeError(f"Error in inference server:\n{response}")
|
|
||||||
decoded = msgpack_numpy.unpackb(response)
|
|
||||||
if isinstance(decoded, dict) and decoded.get("type") == "error":
|
|
||||||
message = decoded.get("message", decoded)
|
|
||||||
raise RuntimeError(f"Error in inference server:\n{message}")
|
|
||||||
return np.asarray(decoded, dtype=np.float32)
|
|
||||||
|
|
||||||
|
|
||||||
@dataclass(frozen=True)
|
|
||||||
class StepRecord:
|
|
||||||
"""One fully materialized rollout step for later JSON export.
|
|
||||||
|
|
||||||
The `episode_00.json` artifact is intended to be human-readable and
|
|
||||||
post-process-friendly. Instead of keeping raw tensors around, each step is
|
|
||||||
flattened into plain Python types so it can be serialized directly.
|
|
||||||
"""
|
|
||||||
|
|
||||||
# Index of this control step within the episode.
|
|
||||||
step_index: int
|
|
||||||
# Whether this step triggered a fresh model call (as opposed to reusing the
|
|
||||||
# cached open-loop chunk from the previous server response).
|
|
||||||
used_server_call: bool
|
|
||||||
# End-to-end latency of the server call that produced the current chunk.
|
|
||||||
# This is `None` on steps that only reuse cached actions.
|
|
||||||
chunk_latency_s: float | None
|
|
||||||
# The concrete 8-D action sent into the simulator at this step.
|
|
||||||
action: list[float]
|
|
||||||
# Observed 7-DoF arm joint positions before the next environment step.
|
|
||||||
joint_position: list[float]
|
|
||||||
# Observed gripper scalar before the next environment step.
|
|
||||||
gripper_position: list[float]
|
|
||||||
# Reward and termination signals directly returned by the simulator.
|
|
||||||
reward: float
|
|
||||||
terminated: bool
|
|
||||||
truncated: bool
|
|
||||||
# Optional scene object positions for downstream debugging / success
|
|
||||||
# heuristics. This may be empty if the environment does not expose them.
|
|
||||||
scene_objects: dict[str, list[float]]
|
|
||||||
|
|
||||||
|
|
||||||
class OpenPIWebsocketClientPolicy(BasePolicy):
|
|
||||||
"""Minimal websocket client for the DreamZero/OpenPI policy protocol.
|
|
||||||
|
|
||||||
Protocol shape:
|
|
||||||
- connect -> server immediately sends a metadata dict
|
|
||||||
- infer -> send msgpack observation, receive action chunk
|
|
||||||
- reset -> send msgpack reset command, receive confirmation string
|
|
||||||
|
|
||||||
This class intentionally stays very small because the more interesting
|
|
||||||
DreamZero-specific behavior lives one layer above, in
|
|
||||||
`DreamZeroJointPosClient`.
|
|
||||||
"""
|
|
||||||
|
|
||||||
def __init__(
|
|
||||||
self,
|
|
||||||
host: str = "127.0.0.1",
|
|
||||||
port: int = 8000,
|
|
||||||
path: str = DEFAULT_PATH,
|
|
||||||
) -> None:
|
|
||||||
# vLLM serves the robot endpoint under `/v1/realtime/robot/openpi`.
|
|
||||||
self._uri = f"ws://{host}:{port}{path}"
|
|
||||||
# Upstream protocol uses msgpack with numpy support, not JSON.
|
|
||||||
self._packer = msgpack_numpy.Packer()
|
|
||||||
# Connect immediately and cache the server handshake metadata.
|
|
||||||
self._ws, self._server_metadata = self._wait_for_server()
|
|
||||||
|
|
||||||
def get_server_metadata(self) -> dict[str, Any]:
|
|
||||||
"""Return a copy of the server handshake metadata."""
|
|
||||||
|
|
||||||
return dict(self._server_metadata)
|
|
||||||
|
|
||||||
def _wait_for_server(self):
|
|
||||||
"""Connect to the websocket server and read the initial metadata frame."""
|
|
||||||
|
|
||||||
logging.info("Connecting to %s", self._uri)
|
|
||||||
conn = websockets.sync.client.connect(
|
|
||||||
self._uri,
|
|
||||||
compression=None,
|
|
||||||
max_size=None,
|
|
||||||
ping_interval=PING_INTERVAL_SECS,
|
|
||||||
ping_timeout=PING_TIMEOUT_SECS,
|
|
||||||
)
|
|
||||||
metadata = msgpack_numpy.unpackb(conn.recv())
|
|
||||||
if not isinstance(metadata, dict):
|
|
||||||
raise TypeError(f"Expected dict metadata from server, got {type(metadata)!r}")
|
|
||||||
return conn, metadata
|
|
||||||
|
|
||||||
@override
|
|
||||||
def infer(self, obs: dict[str, Any]) -> np.ndarray:
|
|
||||||
"""Send an inference request and return the decoded action chunk."""
|
|
||||||
|
|
||||||
# Keep the upstream DreamZero/OpenPI convention that the request itself
|
|
||||||
# tells the server which logical endpoint is being called.
|
|
||||||
payload = dict(obs)
|
|
||||||
payload["endpoint"] = "infer"
|
|
||||||
self._ws.send(self._packer.pack(payload))
|
|
||||||
response = self._ws.recv()
|
|
||||||
return _decode_action_response(response)
|
|
||||||
|
|
||||||
@override
|
|
||||||
def reset(self, reset_info: dict[str, Any] | None = None) -> str:
|
|
||||||
"""Tell the server to reset its session-side state."""
|
|
||||||
|
|
||||||
payload = dict(reset_info or {})
|
|
||||||
payload["endpoint"] = "reset"
|
|
||||||
self._ws.send(self._packer.pack(payload))
|
|
||||||
response = self._ws.recv()
|
|
||||||
if isinstance(response, str):
|
|
||||||
return response
|
|
||||||
decoded = msgpack_numpy.unpackb(response)
|
|
||||||
if not isinstance(decoded, dict) or decoded.get("status") != "reset successful":
|
|
||||||
raise RuntimeError(f"Unexpected reset response: {decoded!r}")
|
|
||||||
return str(decoded["status"])
|
|
||||||
|
|
||||||
def close(self) -> None:
|
|
||||||
"""Close the websocket connection explicitly."""
|
|
||||||
|
|
||||||
self._ws.close()
|
|
||||||
|
|
||||||
|
|
||||||
class DreamZeroJointPosClient:
|
|
||||||
"""DROID sim-eval client that talks to the vLLM OpenPI websocket server.
|
|
||||||
|
|
||||||
This is the main compatibility layer between:
|
|
||||||
- Isaac Lab DROID observations (`obs["policy"][...]`)
|
|
||||||
- DreamZero/OpenPI websocket payloads
|
|
||||||
- local open-loop action reuse across several simulator steps
|
|
||||||
|
|
||||||
In other words:
|
|
||||||
simulator obs -> websocket request -> action chunk -> one action per step
|
|
||||||
"""
|
|
||||||
|
|
||||||
def __init__(
|
|
||||||
self,
|
|
||||||
remote_host: str = "127.0.0.1",
|
|
||||||
remote_port: int = 8000,
|
|
||||||
path: str = DEFAULT_PATH,
|
|
||||||
open_loop_horizon: int = DEFAULT_OPEN_LOOP_HORIZON,
|
|
||||||
) -> None:
|
|
||||||
# Low-level transport client.
|
|
||||||
self.client = OpenPIWebsocketClientPolicy(remote_host, remote_port, path=path)
|
|
||||||
# Number of actions to execute locally before replanning.
|
|
||||||
self.open_loop_horizon = open_loop_horizon
|
|
||||||
# Cursor into the currently cached action chunk.
|
|
||||||
self.actions_from_chunk_completed = 0
|
|
||||||
# Most recent `(ACTION_HORIZON, ACTION_DIM)` server response.
|
|
||||||
self.pred_action_chunk: np.ndarray | None = None
|
|
||||||
# Session id is part of the DreamZero serving contract. Changing it
|
|
||||||
# causes the server side to treat the rollout as a fresh episode.
|
|
||||||
self.session_id = str(uuid.uuid4())
|
|
||||||
# Simple runtime stats for reporting.
|
|
||||||
self.server_calls = 0
|
|
||||||
self.last_chunk_latency_s: float | None = None
|
|
||||||
self.last_used_server_call = False
|
|
||||||
|
|
||||||
def metadata(self) -> dict[str, Any]:
|
|
||||||
"""Expose the server metadata to callers / logs."""
|
|
||||||
|
|
||||||
return self.client.get_server_metadata()
|
|
||||||
|
|
||||||
def reset(self) -> str:
|
|
||||||
"""Reset local chunk state and remote session state.
|
|
||||||
|
|
||||||
Local reset:
|
|
||||||
- drop cached action chunk
|
|
||||||
- rewind chunk cursor
|
|
||||||
- allocate a fresh session id
|
|
||||||
|
|
||||||
Remote reset:
|
|
||||||
- send a websocket `reset` message so the server can clear any
|
|
||||||
request/session-side state it associates with this client
|
|
||||||
"""
|
|
||||||
|
|
||||||
self.actions_from_chunk_completed = 0
|
|
||||||
self.pred_action_chunk = None
|
|
||||||
self.session_id = str(uuid.uuid4())
|
|
||||||
self.last_chunk_latency_s = None
|
|
||||||
self.last_used_server_call = False
|
|
||||||
return self.client.reset({})
|
|
||||||
|
|
||||||
def infer(self, obs: dict[str, Any], instruction: str) -> dict[str, Any]:
|
|
||||||
"""Turn one simulator observation into one executable 8-D action.
|
|
||||||
|
|
||||||
Key behavior:
|
|
||||||
- call the server only when the local chunk cache is empty/exhausted
|
|
||||||
- otherwise, keep consuming the cached chunk open-loop
|
|
||||||
- always return exactly one 8-D action for the current simulator step
|
|
||||||
"""
|
|
||||||
|
|
||||||
# Convert Isaac Lab observation structure into a plain numpy-friendly
|
|
||||||
# record that is easier to serialize and visualize.
|
|
||||||
curr_obs = self._extract_observation(obs)
|
|
||||||
self.last_used_server_call = False
|
|
||||||
|
|
||||||
# Replan if:
|
|
||||||
# 1. this is the first step of a rollout / chunk
|
|
||||||
# 2. we already consumed `open_loop_horizon` actions from the current chunk
|
|
||||||
# 3. no cached chunk is currently available
|
|
||||||
if (
|
|
||||||
self.actions_from_chunk_completed == 0
|
|
||||||
or self.actions_from_chunk_completed >= self.open_loop_horizon
|
|
||||||
or self.pred_action_chunk is None
|
|
||||||
):
|
|
||||||
self.actions_from_chunk_completed = 0
|
|
||||||
# Build the exact DreamZero/OpenPI payload expected by the server.
|
|
||||||
#
|
|
||||||
# Notes:
|
|
||||||
# - images are resized/padded to the serving contract's 180x320
|
|
||||||
# - proprio is cast to float64 to match upstream client behavior
|
|
||||||
# - cartesian_position is currently unused by DreamZero DROID, so
|
|
||||||
# a dummy zero vector is sent for protocol completeness
|
|
||||||
request_data = {
|
|
||||||
"observation/exterior_image_0_left": image_tools.resize_with_pad(curr_obs["right_image"], 180, 320),
|
|
||||||
"observation/exterior_image_1_left": image_tools.resize_with_pad(curr_obs["left_image"], 180, 320),
|
|
||||||
"observation/wrist_image_left": image_tools.resize_with_pad(curr_obs["wrist_image"], 180, 320),
|
|
||||||
"observation/joint_position": curr_obs["joint_position"].astype(np.float64),
|
|
||||||
"observation/cartesian_position": np.zeros((6,), dtype=np.float64),
|
|
||||||
"observation/gripper_position": curr_obs["gripper_position"].astype(np.float64),
|
|
||||||
"prompt": instruction,
|
|
||||||
"session_id": self.session_id,
|
|
||||||
}
|
|
||||||
|
|
||||||
# Measure end-to-end server latency for this chunk request.
|
|
||||||
start = time.perf_counter()
|
|
||||||
actions = self.client.infer(request_data)
|
|
||||||
self.last_chunk_latency_s = time.perf_counter() - start
|
|
||||||
self.last_used_server_call = True
|
|
||||||
self.server_calls += 1
|
|
||||||
|
|
||||||
# DreamZero DROID serving is expected to return an action chunk with
|
|
||||||
# 24 future actions, each action being 8-D.
|
|
||||||
if actions.ndim != 2:
|
|
||||||
raise AssertionError(f"Expected 2D action array, got shape {actions.shape}")
|
|
||||||
if actions.shape != (ACTION_HORIZON, ACTION_DIM):
|
|
||||||
raise AssertionError(f"Expected action shape {(ACTION_HORIZON, ACTION_DIM)}, got {actions.shape}")
|
|
||||||
self.pred_action_chunk = actions
|
|
||||||
|
|
||||||
# Consume exactly one action row from the cached chunk for this
|
|
||||||
# simulator step.
|
|
||||||
action = np.array(self.pred_action_chunk[self.actions_from_chunk_completed], copy=True)
|
|
||||||
self.actions_from_chunk_completed += 1
|
|
||||||
|
|
||||||
# Upstream DreamZero sim-eval binarizes the gripper command.
|
|
||||||
action[-1] = 1.0 if action[-1].item() > 0.5 else 0.0
|
|
||||||
|
|
||||||
# Produce a human-friendly visualization strip for videos:
|
|
||||||
# right external | wrist | left external
|
|
||||||
img1 = image_tools.resize_with_pad(curr_obs["right_image"], 224, 224)
|
|
||||||
img2 = image_tools.resize_with_pad(curr_obs["wrist_image"], 224, 224)
|
|
||||||
img3 = image_tools.resize_with_pad(curr_obs["left_image"], 224, 224)
|
|
||||||
viz = np.concatenate([img1, img2, img3], axis=1)
|
|
||||||
|
|
||||||
# Return both the executable action and auxiliary debug info.
|
|
||||||
return {
|
|
||||||
"action": action,
|
|
||||||
"viz": viz,
|
|
||||||
"joint_position": curr_obs["joint_position"],
|
|
||||||
"gripper_position": curr_obs["gripper_position"],
|
|
||||||
"used_server_call": self.last_used_server_call,
|
|
||||||
"chunk_latency_s": self.last_chunk_latency_s if self.last_used_server_call else None,
|
|
||||||
}
|
|
||||||
|
|
||||||
@staticmethod
|
|
||||||
def _extract_observation(obs_dict: dict[str, Any]) -> dict[str, np.ndarray]:
|
|
||||||
"""Extract the pieces DreamZero cares about from Isaac Lab observations.
|
|
||||||
|
|
||||||
`sim-evals` exposes camera frames and robot state inside the
|
|
||||||
`obs["policy"]` group. This helper converts those tensors into numpy
|
|
||||||
arrays so they can be fed into image preprocessing / websocket packing.
|
|
||||||
"""
|
|
||||||
|
|
||||||
policy = obs_dict["policy"]
|
|
||||||
# Isaac Lab stores camera observations as batched tensors; use env 0.
|
|
||||||
right_image = policy["external_cam"][0].clone().detach().cpu().numpy()
|
|
||||||
left_image = policy["external_cam_2"][0].clone().detach().cpu().numpy()
|
|
||||||
wrist_image = policy["wrist_cam"][0].clone().detach().cpu().numpy()
|
|
||||||
# Robot proprioception.
|
|
||||||
joint_position = policy["arm_joint_pos"].clone().detach().cpu().numpy()
|
|
||||||
gripper_position = policy["gripper_pos"].clone().detach().cpu().numpy()
|
|
||||||
|
|
||||||
return {
|
|
||||||
"right_image": right_image,
|
|
||||||
"left_image": left_image,
|
|
||||||
"wrist_image": wrist_image,
|
|
||||||
"joint_position": joint_position,
|
|
||||||
"gripper_position": gripper_position,
|
|
||||||
}
|
|
||||||
|
|
||||||
|
|
||||||
def _scene_instruction(scene: int) -> str:
|
|
||||||
"""Map a numeric scene id onto the fixed language prompt used for rollout."""
|
|
||||||
|
|
||||||
try:
|
|
||||||
return SCENE_PROMPTS[scene]
|
|
||||||
except KeyError as exc:
|
|
||||||
raise ValueError(f"Unsupported scene {scene}. Available scenes: {sorted(SCENE_PROMPTS)}") from exc
|
|
||||||
|
|
||||||
|
|
||||||
def _capture_scene_objects(env: Any) -> dict[str, list[float]]:
|
|
||||||
"""Best-effort extraction of scene object root positions.
|
|
||||||
|
|
||||||
This is only for debugging / reporting. The rollout logic itself does not
|
|
||||||
depend on these positions.
|
|
||||||
"""
|
|
||||||
|
|
||||||
objects: dict[str, list[float]] = {}
|
|
||||||
scene = getattr(env, "scene", None)
|
|
||||||
if scene is None:
|
|
||||||
return objects
|
|
||||||
|
|
||||||
# Skip non-task entities such as cameras, the robot, and lighting.
|
|
||||||
for name in scene.keys():
|
|
||||||
if name in {"robot", "external_cam", "external_cam_2", "wrist_cam", "sphere_light", "scene"}:
|
|
||||||
continue
|
|
||||||
entity = scene[name]
|
|
||||||
data = getattr(entity, "data", None)
|
|
||||||
root_pos_w = getattr(data, "root_pos_w", None)
|
|
||||||
if root_pos_w is None:
|
|
||||||
continue
|
|
||||||
# Convert tensors to plain lists for JSON serialization.
|
|
||||||
value = root_pos_w[0].detach().cpu().to(torch.float32).tolist()
|
|
||||||
objects[str(name)] = [float(x) for x in value]
|
|
||||||
return objects
|
|
||||||
|
|
||||||
|
|
||||||
def _maybe_infer_success(scene: int, final_objects: dict[str, list[float]]) -> dict[str, Any]:
|
|
||||||
"""Best-effort geometric heuristic.
|
|
||||||
|
|
||||||
The simulator itself does not expose a built-in success term; this function
|
|
||||||
provides a transparent fallback for human-readable reporting only.
|
|
||||||
"""
|
|
||||||
|
|
||||||
task_pairs = {
|
|
||||||
1: ("cube", "bowl"),
|
|
||||||
2: ("can", "mug"),
|
|
||||||
3: ("banana", "bin"),
|
|
||||||
}
|
|
||||||
source_name, target_name = task_pairs.get(scene, (None, None))
|
|
||||||
if source_name not in final_objects or target_name not in final_objects:
|
|
||||||
return {
|
|
||||||
"has_builtin_success": False,
|
|
||||||
"heuristic_success": None,
|
|
||||||
"reason": "scene object names unavailable for heuristic",
|
|
||||||
}
|
|
||||||
|
|
||||||
source = np.asarray(final_objects[source_name], dtype=np.float32)
|
|
||||||
target = np.asarray(final_objects[target_name], dtype=np.float32)
|
|
||||||
xy_distance = float(np.linalg.norm(source[:2] - target[:2]))
|
|
||||||
z_delta = float(source[2] - target[2])
|
|
||||||
heuristic_success = bool(xy_distance < 0.12 and z_delta > -0.08)
|
|
||||||
return {
|
|
||||||
"has_builtin_success": False,
|
|
||||||
"heuristic_success": heuristic_success,
|
|
||||||
"xy_distance": xy_distance,
|
|
||||||
"z_delta": z_delta,
|
|
||||||
"source_object": source_name,
|
|
||||||
"target_object": target_name,
|
|
||||||
"reason": "no built-in env success flag; using final object pose heuristic",
|
|
||||||
}
|
|
||||||
|
|
||||||
|
|
||||||
def _dump_json(path: Path, payload: dict[str, Any]) -> None:
|
|
||||||
"""Write a JSON file with stable UTF-8 formatting."""
|
|
||||||
|
|
||||||
path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
|
|
||||||
|
|
||||||
|
|
||||||
def _scalar_from_env_value(value: Any) -> float:
|
|
||||||
"""Normalize simulator scalar outputs into a plain Python float.
|
|
||||||
|
|
||||||
Isaac Lab / Gym values may come back as tensors, numpy arrays, tuples, or
|
|
||||||
direct Python scalars depending on the wrapper stack. Centralizing the
|
|
||||||
conversion here makes the rollout loop cleaner and more robust.
|
|
||||||
"""
|
|
||||||
|
|
||||||
if isinstance(value, torch.Tensor):
|
|
||||||
return float(value.reshape(-1)[0].detach().cpu().item())
|
|
||||||
if isinstance(value, np.ndarray):
|
|
||||||
return float(value.reshape(-1)[0])
|
|
||||||
if isinstance(value, (list, tuple)):
|
|
||||||
return float(np.asarray(value).reshape(-1)[0])
|
|
||||||
return float(value)
|
|
||||||
|
|
||||||
|
|
||||||
def _bool_from_env_value(value: Any) -> bool:
|
|
||||||
"""Normalize simulator boolean-like outputs into a plain Python bool."""
|
|
||||||
|
|
||||||
if isinstance(value, torch.Tensor):
|
|
||||||
return bool(value.reshape(-1)[0].detach().cpu().item())
|
|
||||||
if isinstance(value, np.ndarray):
|
|
||||||
return bool(value.reshape(-1)[0])
|
|
||||||
if isinstance(value, (list, tuple)):
|
|
||||||
return bool(np.asarray(value).reshape(-1)[0])
|
|
||||||
return bool(value)
|
|
||||||
|
|
||||||
|
|
||||||
def _make_output_dir(output_root: Path, scene: int) -> Path:
|
|
||||||
"""Create a timestamped output directory for one sim-eval run."""
|
|
||||||
|
|
||||||
timestamp = datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ")
|
|
||||||
output_dir = output_root / f"scene{scene}_{timestamp}"
|
|
||||||
output_dir.mkdir(parents=True, exist_ok=True)
|
|
||||||
return output_dir
|
|
||||||
|
|
||||||
|
|
||||||
def main() -> None:
|
|
||||||
"""Entry point for one or more DROID simulation rollouts.
|
|
||||||
|
|
||||||
High-level flow:
|
|
||||||
1. parse command-line flags
|
|
||||||
2. bootstrap Isaac Lab / sim-evals imports
|
|
||||||
3. create the DROID environment
|
|
||||||
4. connect the DreamZero websocket client
|
|
||||||
5. run `episodes` rollouts
|
|
||||||
6. export videos + JSON summaries
|
|
||||||
"""
|
|
||||||
|
|
||||||
# Script-level arguments.
|
|
||||||
parser = argparse.ArgumentParser(description=__doc__)
|
|
||||||
parser.add_argument("--episodes", type=int, default=1, help="Number of episodes to run.")
|
|
||||||
parser.add_argument("--scene", type=int, default=1, help="DROID scene id (1/2/3).")
|
|
||||||
parser.add_argument("--host", type=str, default="127.0.0.1", help="vLLM DreamZero server host.")
|
|
||||||
parser.add_argument("--port", type=int, default=8000, help="vLLM DreamZero server port.")
|
|
||||||
parser.add_argument("--path", type=str, default=DEFAULT_PATH, help="Websocket path suffix.")
|
|
||||||
parser.add_argument(
|
|
||||||
"--open-loop-horizon",
|
|
||||||
type=int,
|
|
||||||
default=DEFAULT_OPEN_LOOP_HORIZON,
|
|
||||||
help="How many actions to consume locally before requesting the next chunk.",
|
|
||||||
)
|
|
||||||
parser.add_argument(
|
|
||||||
"--output-root",
|
|
||||||
type=Path,
|
|
||||||
default=DEFAULT_OUTPUT_ROOT,
|
|
||||||
help="Directory where videos and trajectory logs are stored.",
|
|
||||||
)
|
|
||||||
|
|
||||||
try:
|
|
||||||
from isaaclab.app import AppLauncher
|
|
||||||
except ImportError as exc: # pragma: no cover - runtime dependency guard
|
|
||||||
raise ImportError(
|
|
||||||
"Optional DROID sim-eval client requires Isaac Lab (`isaaclab`). "
|
|
||||||
"Launch it from an Isaac Lab environment, e.g. via `isaaclab.sh -p`."
|
|
||||||
) from exc
|
|
||||||
|
|
||||||
# Let Isaac Lab inject its own runtime flags (e.g. headless, device).
|
|
||||||
AppLauncher.add_app_launcher_args(parser)
|
|
||||||
args = parser.parse_args()
|
|
||||||
|
|
||||||
# DreamZero sim-eval always needs camera observations enabled.
|
|
||||||
args.enable_cameras = True
|
|
||||||
# Boot Isaac Sim / Isaac Lab.
|
|
||||||
app_launcher = AppLauncher(args)
|
|
||||||
simulation_app = app_launcher.app
|
|
||||||
# Set defaults so the `finally` block can clean up safely even if an
|
|
||||||
# earlier step fails.
|
|
||||||
env = None
|
|
||||||
client = None
|
|
||||||
|
|
||||||
# Import simulator modules only *after* the app is launched. This matches
|
|
||||||
# Isaac Lab's required import ordering.
|
|
||||||
try:
|
|
||||||
import gymnasium as gym
|
|
||||||
except ImportError as exc: # pragma: no cover - runtime dependency guard
|
|
||||||
raise ImportError("Optional DROID sim-eval client requires `gymnasium`.") from exc
|
|
||||||
|
|
||||||
try:
|
|
||||||
import sim_evals.environments # noqa: F401
|
|
||||||
except ImportError as exc: # pragma: no cover - runtime dependency guard
|
|
||||||
raise ImportError(
|
|
||||||
"Optional DROID sim-eval client requires the external `sim-evals` package or checkout to be importable."
|
|
||||||
) from exc
|
|
||||||
|
|
||||||
try:
|
|
||||||
from isaaclab_tasks.utils import parse_env_cfg
|
|
||||||
except ImportError as exc: # pragma: no cover - runtime dependency guard
|
|
||||||
raise ImportError("Optional DROID sim-eval client requires `isaaclab_tasks`.") from exc
|
|
||||||
|
|
||||||
# Resolve output location and scene prompt.
|
|
||||||
output_dir = _make_output_dir(args.output_root.expanduser().resolve(), args.scene)
|
|
||||||
instruction = _scene_instruction(args.scene)
|
|
||||||
|
|
||||||
# Build the DROID environment configuration from `sim-evals`.
|
|
||||||
env_cfg = parse_env_cfg(
|
|
||||||
"DROID",
|
|
||||||
device=args.device,
|
|
||||||
num_envs=1,
|
|
||||||
use_fabric=True,
|
|
||||||
)
|
|
||||||
# Select one of the pre-authored scenes/tasks.
|
|
||||||
env_cfg.set_scene(args.scene)
|
|
||||||
env = gym.make("DROID", cfg=env_cfg)
|
|
||||||
|
|
||||||
# Upstream sim-evals resets twice so materials / cameras are fully ready.
|
|
||||||
obs, _ = env.reset()
|
|
||||||
obs, _ = env.reset()
|
|
||||||
|
|
||||||
# Connect the websocket policy client.
|
|
||||||
client = DreamZeroJointPosClient(
|
|
||||||
remote_host=args.host,
|
|
||||||
remote_port=args.port,
|
|
||||||
path=args.path,
|
|
||||||
open_loop_horizon=args.open_loop_horizon,
|
|
||||||
)
|
|
||||||
|
|
||||||
# Aggregated per-run results.
|
|
||||||
all_episode_summaries: list[dict[str, Any]] = []
|
|
||||||
max_steps = int(env.env.max_episode_length)
|
|
||||||
logging.info("DreamZero metadata: %s", client.metadata())
|
|
||||||
logging.info("Scene %s prompt: %s", args.scene, instruction)
|
|
||||||
logging.info("Writing outputs to %s", output_dir)
|
|
||||||
|
|
||||||
try:
|
|
||||||
# No gradients are needed in inference-only rollout mode.
|
|
||||||
with torch.no_grad():
|
|
||||||
for episode_index in range(args.episodes):
|
|
||||||
# Per-episode collectors.
|
|
||||||
frames: list[np.ndarray] = []
|
|
||||||
step_records: list[StepRecord] = []
|
|
||||||
episode_start = time.perf_counter()
|
|
||||||
server_time_s = 0.0
|
|
||||||
final_reward = 0.0
|
|
||||||
terminated = False
|
|
||||||
truncated = False
|
|
||||||
|
|
||||||
for step_index in range(max_steps):
|
|
||||||
# Ask the policy for the next action. Internally this may or
|
|
||||||
# may not trigger a real server request depending on whether
|
|
||||||
# the local chunk cache has been exhausted.
|
|
||||||
logging.debug("Episode %d step %d: requesting action", episode_index, step_index)
|
|
||||||
result = client.infer(obs, instruction)
|
|
||||||
logging.debug(
|
|
||||||
"Episode %d step %d: got action (server_call=%s latency=%s)",
|
|
||||||
episode_index,
|
|
||||||
step_index,
|
|
||||||
result["used_server_call"],
|
|
||||||
result["chunk_latency_s"],
|
|
||||||
)
|
|
||||||
# Save one visualization frame per simulator step.
|
|
||||||
frames.append(result["viz"])
|
|
||||||
|
|
||||||
# Isaac Lab expects batched actions, hence `[None]`.
|
|
||||||
action_tensor = torch.tensor(result["action"], dtype=torch.float32)[None]
|
|
||||||
logging.debug("Episode %d step %d: stepping env", episode_index, step_index)
|
|
||||||
obs, reward, term, trunc, info = env.step(action_tensor)
|
|
||||||
logging.debug("Episode %d step %d: env.step returned", episode_index, step_index)
|
|
||||||
logging.debug("Episode %d step %d: parsing reward/flags", episode_index, step_index)
|
|
||||||
logging.debug(
|
|
||||||
"Episode %d step %d: raw types reward=%s term=%s trunc=%s",
|
|
||||||
episode_index,
|
|
||||||
step_index,
|
|
||||||
type(reward).__name__,
|
|
||||||
type(term).__name__,
|
|
||||||
type(trunc).__name__,
|
|
||||||
)
|
|
||||||
# Normalize environment outputs into plain Python scalars so
|
|
||||||
# the rest of the code does not depend on wrapper-specific types.
|
|
||||||
reward_value = _scalar_from_env_value(reward)
|
|
||||||
term_value = _bool_from_env_value(term)
|
|
||||||
trunc_value = _bool_from_env_value(trunc)
|
|
||||||
logging.debug(
|
|
||||||
"Episode %d step %d: parsed reward=%s term=%s trunc=%s",
|
|
||||||
episode_index,
|
|
||||||
step_index,
|
|
||||||
reward_value,
|
|
||||||
term_value,
|
|
||||||
trunc_value,
|
|
||||||
)
|
|
||||||
# Keep scene-object capture optional. It is useful for
|
|
||||||
# debugging / success heuristics, but the rollout should not
|
|
||||||
# fail if the environment does not expose object roots.
|
|
||||||
scene_objects = _capture_scene_objects(env)
|
|
||||||
|
|
||||||
# Accumulate total server-side time only on steps that
|
|
||||||
# triggered a fresh chunk inference.
|
|
||||||
if result["chunk_latency_s"] is not None:
|
|
||||||
server_time_s += float(result["chunk_latency_s"])
|
|
||||||
|
|
||||||
# Materialize one JSON-serializable trajectory record.
|
|
||||||
logging.debug("Episode %d step %d: appending trajectory", episode_index, step_index)
|
|
||||||
step_records.append(
|
|
||||||
StepRecord(
|
|
||||||
step_index=step_index,
|
|
||||||
used_server_call=bool(result["used_server_call"]),
|
|
||||||
chunk_latency_s=(
|
|
||||||
float(result["chunk_latency_s"]) if result["chunk_latency_s"] is not None else None
|
|
||||||
),
|
|
||||||
action=[float(x) for x in np.asarray(result["action"], dtype=np.float32).tolist()],
|
|
||||||
joint_position=[
|
|
||||||
float(x) for x in np.asarray(result["joint_position"], dtype=np.float32).tolist()
|
|
||||||
],
|
|
||||||
gripper_position=[
|
|
||||||
float(x) for x in np.asarray(result["gripper_position"], dtype=np.float32).tolist()
|
|
||||||
],
|
|
||||||
reward=reward_value,
|
|
||||||
terminated=term_value,
|
|
||||||
truncated=trunc_value,
|
|
||||||
scene_objects=scene_objects,
|
|
||||||
)
|
|
||||||
)
|
|
||||||
logging.debug("Episode %d step %d: trajectory appended", episode_index, step_index)
|
|
||||||
|
|
||||||
# Track final status for summary export.
|
|
||||||
final_reward = reward_value
|
|
||||||
terminated = term_value
|
|
||||||
truncated = trunc_value
|
|
||||||
if term_value or trunc_value:
|
|
||||||
# End the rollout early if the environment terminates.
|
|
||||||
break
|
|
||||||
|
|
||||||
# Episode-level timing and video export.
|
|
||||||
episode_wall_time_s = time.perf_counter() - episode_start
|
|
||||||
video_path = output_dir / f"episode_{episode_index:02d}.mp4"
|
|
||||||
logging.info("Episode %d: writing video to %s", episode_index, video_path)
|
|
||||||
mediapy.write_video(video_path, frames, fps=15)
|
|
||||||
|
|
||||||
# Reset the policy server between episodes.
|
|
||||||
logging.info("Episode %d: sending reset", episode_index)
|
|
||||||
reset_response = client.reset()
|
|
||||||
final_objects = step_records[-1].scene_objects if step_records else {}
|
|
||||||
success_report = _maybe_infer_success(args.scene, final_objects)
|
|
||||||
|
|
||||||
# Assemble the per-episode summary that is written to
|
|
||||||
# `episode_XX.json`.
|
|
||||||
episode_summary = {
|
|
||||||
"episode_index": episode_index,
|
|
||||||
"prompt": instruction,
|
|
||||||
"video_path": str(video_path),
|
|
||||||
"steps_executed": len(step_records),
|
|
||||||
"max_steps": max_steps,
|
|
||||||
"terminated": terminated,
|
|
||||||
"truncated": truncated,
|
|
||||||
"final_reward": final_reward,
|
|
||||||
"server_calls": client.server_calls,
|
|
||||||
"server_time_s": server_time_s,
|
|
||||||
"episode_wall_time_s": episode_wall_time_s,
|
|
||||||
"avg_server_time_per_call_s": (
|
|
||||||
server_time_s / client.server_calls if client.server_calls else None
|
|
||||||
),
|
|
||||||
"reset_response": reset_response,
|
|
||||||
"success_report": success_report,
|
|
||||||
"server_metadata": client.metadata(),
|
|
||||||
"trajectory": [record.__dict__ for record in step_records],
|
|
||||||
}
|
|
||||||
_dump_json(output_dir / f"episode_{episode_index:02d}.json", episode_summary)
|
|
||||||
all_episode_summaries.append(episode_summary)
|
|
||||||
|
|
||||||
logging.info(
|
|
||||||
"Episode %d done: steps=%d wall=%.2fs server_calls=%d heuristic_success=%s",
|
|
||||||
episode_index,
|
|
||||||
len(step_records),
|
|
||||||
episode_wall_time_s,
|
|
||||||
client.server_calls,
|
|
||||||
success_report.get("heuristic_success"),
|
|
||||||
)
|
|
||||||
|
|
||||||
# Reset per-episode counters while keeping the client alive.
|
|
||||||
client.server_calls = 0
|
|
||||||
|
|
||||||
# Top-level run summary across all episodes.
|
|
||||||
summary = {
|
|
||||||
"scene": args.scene,
|
|
||||||
"prompt": instruction,
|
|
||||||
"episodes": args.episodes,
|
|
||||||
"host": args.host,
|
|
||||||
"port": args.port,
|
|
||||||
"path": args.path,
|
|
||||||
"device": args.device,
|
|
||||||
"output_dir": str(output_dir),
|
|
||||||
"server_metadata": client.metadata(),
|
|
||||||
"episodes_summary": all_episode_summaries,
|
|
||||||
}
|
|
||||||
_dump_json(output_dir / "summary.json", summary)
|
|
||||||
# Also print the summary to stdout so the caller can capture it in logs.
|
|
||||||
print(json.dumps(summary, ensure_ascii=False, indent=2))
|
|
||||||
finally:
|
|
||||||
# Best-effort cleanup. Avoid masking the main error if cleanup fails.
|
|
||||||
try:
|
|
||||||
if client is not None:
|
|
||||||
client.client.close()
|
|
||||||
except Exception:
|
|
||||||
pass
|
|
||||||
if env is not None:
|
|
||||||
env.close()
|
|
||||||
simulation_app.close()
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
# Keep the script-level logs readable. Per-step rollout details are still
|
|
||||||
# available via `DEBUG` if needed, but websocket / asyncio internals are
|
|
||||||
# usually too noisy for normal usage.
|
|
||||||
logging.basicConfig(level=logging.INFO)
|
|
||||||
logging.getLogger("websockets").setLevel(logging.WARNING)
|
|
||||||
logging.getLogger("asyncio").setLevel(logging.WARNING)
|
|
||||||
main()
|
|
||||||
@@ -1,86 +0,0 @@
|
|||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
import argparse
|
|
||||||
import os
|
|
||||||
import sys
|
|
||||||
from pathlib import Path
|
|
||||||
|
|
||||||
os.environ.setdefault("MUJOCO_GL", "egl")
|
|
||||||
os.environ.setdefault("PYOPENGL_PLATFORM", "egl")
|
|
||||||
|
|
||||||
# Import base configs at module top level so the subclasses below are pickle-
|
|
||||||
# resolvable (worker processes import this module fresh via __main__).
|
|
||||||
from molmo_spaces.configs.policy_configs_baselines import ( # noqa: E402
|
|
||||||
DreamZeroPolicyConfig,
|
|
||||||
)
|
|
||||||
from molmo_spaces.evaluation.configs.evaluation_configs import ( # noqa: E402
|
|
||||||
DreamZeroPolicyEvalConfig,
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
# We only need to change the backend host and port to the vllm-host!
|
|
||||||
class DreamZeroVllmOmniPolicyConfig(DreamZeroPolicyConfig):
|
|
||||||
remote_config: dict = dict(host="127.0.0.1", port=8000)
|
|
||||||
|
|
||||||
|
|
||||||
class DreamZeroVllmOmniEvalConfig(DreamZeroPolicyEvalConfig):
|
|
||||||
policy_config: DreamZeroVllmOmniPolicyConfig = DreamZeroVllmOmniPolicyConfig()
|
|
||||||
|
|
||||||
|
|
||||||
def main() -> int:
|
|
||||||
parser = argparse.ArgumentParser()
|
|
||||||
parser.add_argument("--host", default="127.0.0.1")
|
|
||||||
parser.add_argument("--port", type=int, default=8000)
|
|
||||||
parser.add_argument(
|
|
||||||
"--benchmark_dir",
|
|
||||||
required=True,
|
|
||||||
help=(
|
|
||||||
"Path to a MolmoSpaces benchmark directory, for example "
|
|
||||||
"$MOLMOSPACES_BENCHMARK_DIR/20260327/ithor/FrankaCloseHardBench/"
|
|
||||||
"FrankaCloseHardBench_20260206_json_benchmark"
|
|
||||||
),
|
|
||||||
)
|
|
||||||
parser.add_argument("--max_episodes", type=int, default=1)
|
|
||||||
parser.add_argument("--task_horizon_steps", type=int, default=80)
|
|
||||||
parser.add_argument(
|
|
||||||
"--output_dir",
|
|
||||||
required=True,
|
|
||||||
help="Directory to write evaluation outputs (created if missing).",
|
|
||||||
)
|
|
||||||
parser.add_argument("--episode_idx", type=int, default=None)
|
|
||||||
args = parser.parse_args()
|
|
||||||
|
|
||||||
policy_config = DreamZeroVllmOmniPolicyConfig(remote_config=dict(host=args.host, port=args.port))
|
|
||||||
DreamZeroVllmOmniPolicyConfig.model_fields["remote_config"].default = policy_config.remote_config
|
|
||||||
DreamZeroVllmOmniEvalConfig.model_fields["policy_config"].default = policy_config
|
|
||||||
|
|
||||||
# Import after env vars are set so MuJoCo picks EGL.
|
|
||||||
from molmo_spaces.evaluation import run_evaluation
|
|
||||||
|
|
||||||
cfg_cls = DreamZeroVllmOmniEvalConfig
|
|
||||||
|
|
||||||
output_dir = args.output_dir
|
|
||||||
Path(output_dir).mkdir(parents=True, exist_ok=True)
|
|
||||||
|
|
||||||
print(f"[eval] benchmark_dir={args.benchmark_dir}")
|
|
||||||
print(f"[eval] max_episodes={args.max_episodes} task_horizon_steps={args.task_horizon_steps}")
|
|
||||||
print(f"[eval] remote policy: ws://{args.host}:{args.port}/v1/realtime/robot/openpi")
|
|
||||||
|
|
||||||
results = run_evaluation(
|
|
||||||
eval_config_cls=cfg_cls,
|
|
||||||
benchmark_dir=Path(args.benchmark_dir),
|
|
||||||
max_episodes=args.max_episodes,
|
|
||||||
task_horizon_steps=args.task_horizon_steps,
|
|
||||||
num_workers=1,
|
|
||||||
use_wandb=False,
|
|
||||||
output_dir=output_dir,
|
|
||||||
episode_idx=args.episode_idx,
|
|
||||||
)
|
|
||||||
|
|
||||||
print(f"[eval] success={results.success_count}/{results.total_count} ({results.success_rate:.1%})")
|
|
||||||
print(f"[eval] output_dir={results.output_dir}")
|
|
||||||
return 0
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
sys.exit(main())
|
|
||||||
@@ -1,311 +0,0 @@
|
|||||||
#!/usr/bin/env python3
|
|
||||||
# SPDX-License-Identifier: Apache-2.0
|
|
||||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
||||||
|
|
||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
import argparse
|
|
||||||
import json
|
|
||||||
import logging
|
|
||||||
import sys
|
|
||||||
import uuid
|
|
||||||
from dataclasses import dataclass
|
|
||||||
from pathlib import Path
|
|
||||||
from typing import Any
|
|
||||||
|
|
||||||
import numpy as np
|
|
||||||
|
|
||||||
try:
|
|
||||||
import cv2
|
|
||||||
except ImportError as exc: # pragma: no cover - runtime dependency guard
|
|
||||||
raise ImportError("DreamZero OpenPI example requires `opencv-python`.") from exc
|
|
||||||
|
|
||||||
try:
|
|
||||||
example_dir = str(Path(__file__).resolve().parent)
|
|
||||||
removed_path = False
|
|
||||||
if sys.path and sys.path[0] == example_dir:
|
|
||||||
sys.path.pop(0)
|
|
||||||
removed_path = True
|
|
||||||
try:
|
|
||||||
from openpi_client.websocket_client_policy import WebsocketClientPolicy
|
|
||||||
finally:
|
|
||||||
if removed_path:
|
|
||||||
sys.path.insert(0, example_dir)
|
|
||||||
except ImportError as exc: # pragma: no cover - runtime dependency guard
|
|
||||||
raise ImportError("DreamZero OpenPI example requires `openpi-client`.") from exc
|
|
||||||
|
|
||||||
DEFAULT_HOST = "127.0.0.1"
|
|
||||||
DEFAULT_PORT = 8000
|
|
||||||
DEFAULT_PATH = "/v1/realtime/robot/openpi"
|
|
||||||
DEFAULT_PROMPT = "Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan"
|
|
||||||
ACTION_HORIZON = 24
|
|
||||||
DEFAULT_ACTION_DIM = 8
|
|
||||||
RELATIVE_OFFSETS = [-23, -16, -8, 0]
|
|
||||||
REPO_ROOT = Path(__file__).resolve().parents[3]
|
|
||||||
ASSET_REPO_ID = "YangshenDeng/vllm-omni-dreamzero-assets"
|
|
||||||
DEFAULT_VIDEO_DIR = REPO_ROOT / "outputs" / "dreamzero" / "assets"
|
|
||||||
CAMERA_FILES = {
|
|
||||||
"observation/exterior_image_0_left": "exterior_image_1_left.mp4",
|
|
||||||
"observation/exterior_image_1_left": "exterior_image_2_left.mp4",
|
|
||||||
"observation/wrist_image_left": "wrist_image_left.mp4",
|
|
||||||
}
|
|
||||||
|
|
||||||
|
|
||||||
def _as_action_array(response: Any) -> np.ndarray:
|
|
||||||
if isinstance(response, dict) and response.get("type") == "error":
|
|
||||||
message = response.get("message", response)
|
|
||||||
raise RuntimeError(f"Inference failed: {message}")
|
|
||||||
return np.asarray(response, dtype=np.float32)
|
|
||||||
|
|
||||||
|
|
||||||
@dataclass(frozen=True)
|
|
||||||
class DreamZeroServerMetadata:
|
|
||||||
image_resolution: tuple[int, int]
|
|
||||||
n_external_cameras: int
|
|
||||||
needs_wrist_camera: bool
|
|
||||||
needs_stereo_camera: bool
|
|
||||||
needs_session_id: bool
|
|
||||||
action_space: str
|
|
||||||
|
|
||||||
@classmethod
|
|
||||||
def from_dict(cls, payload: dict[str, Any]) -> DreamZeroServerMetadata:
|
|
||||||
required_keys = (
|
|
||||||
"image_resolution",
|
|
||||||
"n_external_cameras",
|
|
||||||
"needs_wrist_camera",
|
|
||||||
"needs_stereo_camera",
|
|
||||||
"needs_session_id",
|
|
||||||
"action_space",
|
|
||||||
)
|
|
||||||
missing_keys = [key for key in required_keys if key not in payload]
|
|
||||||
if missing_keys:
|
|
||||||
raise ValueError(f"Missing DreamZero metadata keys: {missing_keys}")
|
|
||||||
|
|
||||||
image_resolution = payload["image_resolution"]
|
|
||||||
if not isinstance(image_resolution, (list, tuple)) or len(image_resolution) != 2:
|
|
||||||
raise ValueError(f"Invalid image_resolution: {image_resolution!r}")
|
|
||||||
|
|
||||||
return cls(
|
|
||||||
image_resolution=(int(image_resolution[0]), int(image_resolution[1])),
|
|
||||||
n_external_cameras=int(payload["n_external_cameras"]),
|
|
||||||
needs_wrist_camera=bool(payload["needs_wrist_camera"]),
|
|
||||||
needs_stereo_camera=bool(payload["needs_stereo_camera"]),
|
|
||||||
needs_session_id=bool(payload["needs_session_id"]),
|
|
||||||
action_space=str(payload["action_space"]),
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
def _openpi_uri(host: str, port: int, path: str) -> str:
|
|
||||||
if host.startswith(("ws://", "wss://")):
|
|
||||||
return host
|
|
||||||
return f"ws://{host}:{port}{path}"
|
|
||||||
|
|
||||||
|
|
||||||
def _close_policy_client(client: WebsocketClientPolicy) -> None:
|
|
||||||
ws = getattr(client, "_ws", None)
|
|
||||||
if ws is not None:
|
|
||||||
ws.close()
|
|
||||||
|
|
||||||
|
|
||||||
def load_all_frames(video_path: Path) -> np.ndarray:
|
|
||||||
cap = cv2.VideoCapture(str(video_path))
|
|
||||||
frames = []
|
|
||||||
while True:
|
|
||||||
ok, frame = cap.read()
|
|
||||||
if not ok:
|
|
||||||
break
|
|
||||||
frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
|
|
||||||
cap.release()
|
|
||||||
if not frames:
|
|
||||||
raise RuntimeError(f"No frames loaded from {video_path}")
|
|
||||||
return np.stack(frames, axis=0)
|
|
||||||
|
|
||||||
|
|
||||||
def load_camera_frames(video_dir: Path) -> dict[str, np.ndarray]:
|
|
||||||
camera_frames: dict[str, np.ndarray] = {}
|
|
||||||
for camera_key, file_name in CAMERA_FILES.items():
|
|
||||||
video_path = video_dir / file_name
|
|
||||||
if not video_path.exists():
|
|
||||||
raise FileNotFoundError(
|
|
||||||
f"Missing DreamZero example asset: {video_path}. "
|
|
||||||
"Download the example videos with: "
|
|
||||||
f"`hf download {ASSET_REPO_ID} --repo-type dataset --local-dir {video_dir}`"
|
|
||||||
)
|
|
||||||
camera_frames[camera_key] = load_all_frames(video_path)
|
|
||||||
return camera_frames
|
|
||||||
|
|
||||||
|
|
||||||
def build_frame_schedule(total_frames: int, num_chunks: int) -> list[list[int]]:
|
|
||||||
chunks: list[list[int]] = []
|
|
||||||
current_frame = 23
|
|
||||||
for _ in range(num_chunks):
|
|
||||||
indices = [max(current_frame + offset, 0) for offset in RELATIVE_OFFSETS]
|
|
||||||
if indices[-1] >= total_frames:
|
|
||||||
break
|
|
||||||
chunks.append(indices)
|
|
||||||
current_frame += ACTION_HORIZON
|
|
||||||
return chunks
|
|
||||||
|
|
||||||
|
|
||||||
def make_obs_from_video(
|
|
||||||
camera_frames: dict[str, np.ndarray],
|
|
||||||
frame_indices: list[int],
|
|
||||||
*,
|
|
||||||
prompt: str,
|
|
||||||
session_id: str,
|
|
||||||
) -> dict[str, Any]:
|
|
||||||
obs: dict[str, Any] = {}
|
|
||||||
for camera_key, all_frames in camera_frames.items():
|
|
||||||
selected = all_frames[frame_indices]
|
|
||||||
obs[camera_key] = selected[0] if len(frame_indices) == 1 else selected
|
|
||||||
|
|
||||||
obs["observation/joint_position"] = np.zeros(7, dtype=np.float32)
|
|
||||||
obs["observation/cartesian_position"] = np.zeros(6, dtype=np.float32)
|
|
||||||
obs["observation/gripper_position"] = np.zeros(1, dtype=np.float32)
|
|
||||||
obs["prompt"] = prompt
|
|
||||||
obs["session_id"] = session_id
|
|
||||||
return obs
|
|
||||||
|
|
||||||
|
|
||||||
def build_demo_observations(
|
|
||||||
camera_frames: dict[str, np.ndarray],
|
|
||||||
*,
|
|
||||||
prompt: str,
|
|
||||||
session_id: str,
|
|
||||||
num_chunks: int = 2,
|
|
||||||
) -> list[dict[str, Any]]:
|
|
||||||
if num_chunks < 1:
|
|
||||||
raise ValueError("num_chunks must be at least 1")
|
|
||||||
|
|
||||||
total_frames = min(frames.shape[0] for frames in camera_frames.values())
|
|
||||||
observations = [
|
|
||||||
make_obs_from_video(
|
|
||||||
camera_frames,
|
|
||||||
[0],
|
|
||||||
prompt=prompt,
|
|
||||||
session_id=session_id,
|
|
||||||
)
|
|
||||||
]
|
|
||||||
for indices in build_frame_schedule(total_frames, num_chunks - 1):
|
|
||||||
observations.append(
|
|
||||||
make_obs_from_video(
|
|
||||||
camera_frames,
|
|
||||||
indices,
|
|
||||||
prompt=prompt,
|
|
||||||
session_id=session_id,
|
|
||||||
)
|
|
||||||
)
|
|
||||||
return observations
|
|
||||||
|
|
||||||
|
|
||||||
def validate_session_result(
|
|
||||||
result: dict[str, Any],
|
|
||||||
*,
|
|
||||||
expected_action_horizon: int = ACTION_HORIZON,
|
|
||||||
expected_action_dim: int = DEFAULT_ACTION_DIM,
|
|
||||||
expected_action_count: int | None = None,
|
|
||||||
) -> None:
|
|
||||||
metadata = DreamZeroServerMetadata.from_dict(result["metadata"])
|
|
||||||
if metadata.image_resolution != (180, 320):
|
|
||||||
raise AssertionError(f"Unexpected image_resolution: {metadata.image_resolution}")
|
|
||||||
if metadata.n_external_cameras != 2:
|
|
||||||
raise AssertionError(f"Unexpected n_external_cameras: {metadata.n_external_cameras}")
|
|
||||||
if not metadata.needs_wrist_camera:
|
|
||||||
raise AssertionError("DreamZero example expects wrist camera metadata")
|
|
||||||
if metadata.action_space != "joint_position":
|
|
||||||
raise AssertionError(f"Unexpected action_space: {metadata.action_space}")
|
|
||||||
|
|
||||||
actions = result["actions"]
|
|
||||||
if expected_action_count is not None and len(actions) != expected_action_count:
|
|
||||||
raise AssertionError(f"Expected {expected_action_count} action tensors, got {len(actions)}")
|
|
||||||
if not actions:
|
|
||||||
raise AssertionError("Expected at least one action tensor")
|
|
||||||
for index, action in enumerate(actions):
|
|
||||||
if action.shape != (expected_action_horizon, expected_action_dim):
|
|
||||||
raise AssertionError(
|
|
||||||
f"Action {index} shape mismatch: expected "
|
|
||||||
f"{(expected_action_horizon, expected_action_dim)}, got {action.shape}"
|
|
||||||
)
|
|
||||||
if not np.isfinite(action).all():
|
|
||||||
raise AssertionError(f"Action {index} contains non-finite values")
|
|
||||||
|
|
||||||
|
|
||||||
def run_policy_session(
|
|
||||||
*,
|
|
||||||
host: str = DEFAULT_HOST,
|
|
||||||
port: int = DEFAULT_PORT,
|
|
||||||
path: str = DEFAULT_PATH,
|
|
||||||
video_dir: Path = DEFAULT_VIDEO_DIR,
|
|
||||||
prompt: str = DEFAULT_PROMPT,
|
|
||||||
session_id: str | None = None,
|
|
||||||
num_chunks: int = 2,
|
|
||||||
) -> dict[str, Any]:
|
|
||||||
session_id = session_id or str(uuid.uuid4())
|
|
||||||
camera_frames = load_camera_frames(video_dir)
|
|
||||||
observations = build_demo_observations(
|
|
||||||
camera_frames,
|
|
||||||
prompt=prompt,
|
|
||||||
session_id=session_id,
|
|
||||||
num_chunks=num_chunks,
|
|
||||||
)
|
|
||||||
|
|
||||||
uri = _openpi_uri(host, port, path)
|
|
||||||
logging.info("Connecting to %s", uri)
|
|
||||||
client = WebsocketClientPolicy(host=uri)
|
|
||||||
try:
|
|
||||||
metadata = dict(client.get_server_metadata())
|
|
||||||
actions = [_as_action_array(client.infer(obs)) for obs in observations]
|
|
||||||
return {
|
|
||||||
"metadata": metadata,
|
|
||||||
"actions": actions,
|
|
||||||
"session_id": session_id,
|
|
||||||
}
|
|
||||||
finally:
|
|
||||||
_close_policy_client(client)
|
|
||||||
|
|
||||||
|
|
||||||
def format_action_summary(index: int, action: np.ndarray) -> str:
|
|
||||||
return (
|
|
||||||
f"Action {index}: shape={tuple(action.shape)} dtype={action.dtype} "
|
|
||||||
f"min={action.min():.6f} max={action.max():.6f}"
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
def parse_args() -> argparse.Namespace:
|
|
||||||
parser = argparse.ArgumentParser(description="DreamZero OpenPI client example with downloaded real videos.")
|
|
||||||
parser.add_argument("--host", default=DEFAULT_HOST)
|
|
||||||
parser.add_argument("--port", type=int, default=DEFAULT_PORT)
|
|
||||||
parser.add_argument("--path", default=DEFAULT_PATH)
|
|
||||||
parser.add_argument(
|
|
||||||
"--video-dir", type=Path, default=DEFAULT_VIDEO_DIR, help="Directory containing the three camera MP4 files."
|
|
||||||
)
|
|
||||||
parser.add_argument("--prompt", default=DEFAULT_PROMPT)
|
|
||||||
parser.add_argument("--session-id", default=None)
|
|
||||||
parser.add_argument("--num-chunks", type=int, default=2)
|
|
||||||
return parser.parse_args()
|
|
||||||
|
|
||||||
|
|
||||||
def main() -> int:
|
|
||||||
args = parse_args()
|
|
||||||
logging.basicConfig(level=logging.INFO)
|
|
||||||
|
|
||||||
result = run_policy_session(
|
|
||||||
host=args.host,
|
|
||||||
port=args.port,
|
|
||||||
path=args.path,
|
|
||||||
video_dir=args.video_dir,
|
|
||||||
prompt=args.prompt,
|
|
||||||
session_id=args.session_id,
|
|
||||||
num_chunks=args.num_chunks,
|
|
||||||
)
|
|
||||||
validate_session_result(result, expected_action_count=args.num_chunks)
|
|
||||||
|
|
||||||
print("Server metadata:", json.dumps(result["metadata"], sort_keys=True))
|
|
||||||
for index, action in enumerate(result["actions"]):
|
|
||||||
print(format_action_summary(index, action))
|
|
||||||
print("Session ID:", result["session_id"])
|
|
||||||
return 0
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
raise SystemExit(main())
|
|
||||||
@@ -1,24 +0,0 @@
|
|||||||
#!/usr/bin/env bash
|
|
||||||
set -euo pipefail
|
|
||||||
|
|
||||||
MODEL="${MODEL:-GEAR-Dreams/DreamZero-DROID}"
|
|
||||||
HOST="${HOST:-127.0.0.1}"
|
|
||||||
PORT="${PORT:-8000}"
|
|
||||||
DEPLOY_CONFIG="${DEPLOY_CONFIG:-vllm_omni/deploy/dreamzero_tp1_cfg2.yaml}"
|
|
||||||
SERVED_MODEL_NAME="${SERVED_MODEL_NAME:-dreamzero-droid}"
|
|
||||||
|
|
||||||
args=(
|
|
||||||
serve
|
|
||||||
"$MODEL"
|
|
||||||
--omni
|
|
||||||
--host "$HOST"
|
|
||||||
--port "$PORT"
|
|
||||||
--served-model-name "$SERVED_MODEL_NAME"
|
|
||||||
--deploy-config "$DEPLOY_CONFIG"
|
|
||||||
--enforce-eager
|
|
||||||
--disable-log-stats
|
|
||||||
)
|
|
||||||
|
|
||||||
ATTENTION_BACKEND="${ATTENTION_BACKEND:-torch}" \
|
|
||||||
DIFFUSION_ATTENTION_BACKEND="${DIFFUSION_ATTENTION_BACKEND:-TORCH_SDPA}" \
|
|
||||||
vllm "${args[@]}"
|
|
||||||
@@ -22,14 +22,13 @@ export PYTHONPATH=<transformers_modules_path>:$PYTHONPATH
|
|||||||
export HF_MODULES_CACHE=<transformers_modules_path>
|
export HF_MODULES_CACHE=<transformers_modules_path>
|
||||||
```
|
```
|
||||||
|
|
||||||
Run from repository root. The 3-stage pipeline is resolved automatically from
|
Run from repository root:
|
||||||
the model type (`vllm_omni/deploy/dynin_omni.yaml`), so no config path is
|
|
||||||
required. Pass `--deploy-config` only to override the defaults:
|
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
vllm-omni serve snu-aidas/Dynin-Omni \
|
vllm-omni serve snu-aidas/Dynin-Omni \
|
||||||
--omni \
|
--omni \
|
||||||
--port 8091
|
--port 8091 \
|
||||||
|
--stage-configs-path "$(pwd)/vllm_omni/model_executor/stage_configs/dynin_omni.yaml"
|
||||||
```
|
```
|
||||||
|
|
||||||
If `vllm-omni` is not in PATH, run:
|
If `vllm-omni` is not in PATH, run:
|
||||||
@@ -37,7 +36,8 @@ If `vllm-omni` is not in PATH, run:
|
|||||||
```bash
|
```bash
|
||||||
PYTHONPATH="$(pwd)" python -m vllm_omni.entrypoints.cli.main serve snu-aidas/Dynin-Omni \
|
PYTHONPATH="$(pwd)" python -m vllm_omni.entrypoints.cli.main serve snu-aidas/Dynin-Omni \
|
||||||
--omni \
|
--omni \
|
||||||
--port 8091
|
--port 8091 \
|
||||||
|
--stage-configs-path "$(pwd)/vllm_omni/model_executor/stage_configs/dynin_omni.yaml"
|
||||||
```
|
```
|
||||||
|
|
||||||
Wait until the server logs show both `All stages initialized successfully` and
|
Wait until the server logs show both `All stages initialized successfully` and
|
||||||
|
|||||||
+2
-3
@@ -12,8 +12,7 @@ from typing import Any
|
|||||||
import requests
|
import requests
|
||||||
from openai import OpenAI
|
from openai import OpenAI
|
||||||
from vllm.assets.audio import AudioAsset
|
from vllm.assets.audio import AudioAsset
|
||||||
|
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||||
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
|
|
||||||
|
|
||||||
# Modify OpenAI's API key and API base to use vLLM's API server.
|
# Modify OpenAI's API key and API base to use vLLM's API server.
|
||||||
openai_api_key = "EMPTY"
|
openai_api_key = "EMPTY"
|
||||||
@@ -494,7 +493,7 @@ def run_multimodal_generation(args) -> None:
|
|||||||
|
|
||||||
|
|
||||||
def parse_args():
|
def parse_args():
|
||||||
parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
|
parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
"--query-type",
|
"--query-type",
|
||||||
"-q",
|
"-q",
|
||||||
|
|||||||
@@ -1,65 +0,0 @@
|
|||||||
# vLLM-Omni · MiniCPM-o 4.5 Online Demo
|
|
||||||
|
|
||||||
Gradio-based web UI for **MiniCPM-o 4.5** served via `vllm-omni`'s
|
|
||||||
OpenAI-compatible endpoints.
|
|
||||||
|
|
||||||
The UI supports:
|
|
||||||
|
|
||||||
- **Inputs**: text prompt + optional image, audio (file or mic), video.
|
|
||||||
- **Outputs**: text + speech (WAV player).
|
|
||||||
|
|
||||||
## 1. Start the backend server
|
|
||||||
|
|
||||||
The deploy config auto-loads via `--omni`; the default
|
|
||||||
`vllm_omni/deploy/minicpmo_4_5.yaml` targets a 2-GPU layout (thinker on
|
|
||||||
GPU 0, talker + t2w sharing GPU 1). For other hardware layouts pick
|
|
||||||
one of the deploy variants below.
|
|
||||||
|
|
||||||
| deploy config | GPUs | Notes |
|
|
||||||
|---|---|---|
|
|
||||||
| `minicpmo_4_5.yaml` (default) | 2 | Thinker on GPU0, talker+t2w on GPU1. |
|
|
||||||
| `minicpmo_4_5_3gpu.yaml` | 3 | Thinker 2-way TP on GPU0/1, talker+t2w share GPU2. |
|
|
||||||
| `minicpmo_4_5_8x4090.yaml` | 8 | Full 8x4090 layout. |
|
|
||||||
|
|
||||||
Default (2-GPU):
|
|
||||||
|
|
||||||
```bash
|
|
||||||
vllm serve openbmb/MiniCPM-o-4_5 --omni \
|
|
||||||
--trust-remote-code \
|
|
||||||
--host 0.0.0.0 --port 8099
|
|
||||||
```
|
|
||||||
|
|
||||||
Other layouts via `--deploy-config`:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
vllm serve openbmb/MiniCPM-o-4_5 --omni \
|
|
||||||
--deploy-config vllm_omni/deploy/minicpmo_4_5_8x4090.yaml \
|
|
||||||
--trust-remote-code \
|
|
||||||
--host 0.0.0.0 --port 8099
|
|
||||||
```
|
|
||||||
|
|
||||||
## 2. Launch the Gradio demo
|
|
||||||
|
|
||||||
```bash
|
|
||||||
bash examples/online_serving/minicpmo/run_gradio_demo.sh
|
|
||||||
|
|
||||||
# Or run the python entry point directly:
|
|
||||||
python examples/online_serving/minicpmo/gradio_demo.py \
|
|
||||||
--minicpmo45-api-base http://localhost:8099/v1 \
|
|
||||||
--minicpmo45-model openbmb/MiniCPM-o-4_5 \
|
|
||||||
--port 7862
|
|
||||||
```
|
|
||||||
|
|
||||||
Open `http://<host>:7862` in a browser.
|
|
||||||
|
|
||||||
## Notes
|
|
||||||
|
|
||||||
- **TTS trigger**: the demo sets
|
|
||||||
`extra_body.chat_template_kwargs.use_tts_template=True`, which appends
|
|
||||||
`<|tts_bos|>` to the assistant prefix.
|
|
||||||
- Uncheck **"Generate speech output (TTS)"** to get text-only responses
|
|
||||||
(faster).
|
|
||||||
- The audio output is the raw WAV returned by the stage-1 talker +
|
|
||||||
Token2Wav; sample rate is 24 kHz.
|
|
||||||
- Video input is forwarded as a base64 `video_url` entry; the server needs
|
|
||||||
decord/torchvision to decode it.
|
|
||||||
@@ -1,580 +0,0 @@
|
|||||||
"""Gradio demo for MiniCPM-o 4.5 online serving through vllm-omni.
|
|
||||||
|
|
||||||
Supports:
|
|
||||||
- Text / Image / Audio / Video inputs (OpenAI chat-completions multimodal format).
|
|
||||||
- Text + audio outputs, returned side-by-side in the UI.
|
|
||||||
|
|
||||||
Usage:
|
|
||||||
# run a 4.5 server on :8099
|
|
||||||
python gradio_demo.py \
|
|
||||||
--minicpmo45-api-base http://localhost:8099/v1 \
|
|
||||||
--minicpmo45-model openbmb/MiniCPM-o-4_5 \
|
|
||||||
--port 7862
|
|
||||||
"""
|
|
||||||
|
|
||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
import argparse
|
|
||||||
import base64
|
|
||||||
import io
|
|
||||||
import os
|
|
||||||
from collections.abc import Iterable
|
|
||||||
from pathlib import Path
|
|
||||||
from typing import Any
|
|
||||||
|
|
||||||
import gradio as gr
|
|
||||||
import numpy as np
|
|
||||||
import soundfile as sf
|
|
||||||
from openai import OpenAI
|
|
||||||
from PIL import Image
|
|
||||||
|
|
||||||
MINICPMO45 = "MiniCPM-o-4.5"
|
|
||||||
|
|
||||||
# Default reference audio per model, relative to the model path. Used to build
|
|
||||||
# the audio_assistant system prompt so that the LLM adopts the "speak in this
|
|
||||||
# voice" persona (matches the official MiniCPM-o demo).
|
|
||||||
_DEFAULT_REF_AUDIO = {
|
|
||||||
MINICPMO45: "assets/HT_ref_audio.wav",
|
|
||||||
}
|
|
||||||
|
|
||||||
# audio_assistant prompt text, mirrors MiniCPMO.get_sys_prompt("audio_assistant")
|
|
||||||
# in the MiniCPM-o-Demo reference implementation.
|
|
||||||
_AUDIO_ASSISTANT_PROMPT = {
|
|
||||||
"zh": {
|
|
||||||
"prefix": "模仿音频样本的音色并生成新的内容。",
|
|
||||||
"suffix": (
|
|
||||||
"你的任务是用这种声音模式来当一个助手。请认真、高质量地回复用户的问题。"
|
|
||||||
"请用高自然度的方式和用户聊天。你是由面壁智能开发的人工智能助手:面壁小钢炮。"
|
|
||||||
),
|
|
||||||
},
|
|
||||||
"en": {
|
|
||||||
"prefix": "Use the voice in the audio prompt to synthesize new content.",
|
|
||||||
"suffix": "You are a helpful assistant with the above voice style.",
|
|
||||||
},
|
|
||||||
}
|
|
||||||
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
# Media helpers
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
|
|
||||||
|
|
||||||
def image_to_base64_data_url(image: Image.Image) -> str:
|
|
||||||
buf = io.BytesIO()
|
|
||||||
if image.mode != "RGB":
|
|
||||||
image = image.convert("RGB")
|
|
||||||
image.save(buf, format="JPEG", quality=92)
|
|
||||||
b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
|
|
||||||
return f"data:image/jpeg;base64,{b64}"
|
|
||||||
|
|
||||||
|
|
||||||
def audio_to_base64_data_url(audio_np: np.ndarray, sample_rate: int) -> str:
|
|
||||||
if audio_np.dtype != np.int16:
|
|
||||||
if audio_np.dtype in (np.float32, np.float64):
|
|
||||||
audio_np = np.clip(audio_np, -1.0, 1.0)
|
|
||||||
audio_np = (audio_np * 32767).astype(np.int16)
|
|
||||||
else:
|
|
||||||
audio_np = audio_np.astype(np.int16)
|
|
||||||
buf = io.BytesIO()
|
|
||||||
sf.write(buf, audio_np, sample_rate, format="WAV")
|
|
||||||
b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
|
|
||||||
return f"data:audio/wav;base64,{b64}"
|
|
||||||
|
|
||||||
|
|
||||||
def ref_audio_to_data_url(path: str, target_sr: int = 16000, max_s: float | None = 8.0) -> str:
|
|
||||||
"""Load a file, downmix to mono, resample to 16kHz, optionally truncate, and
|
|
||||||
return a data-URL usable as audio_url content in the chat API.
|
|
||||||
|
|
||||||
MiniCPM-o expects 16kHz mono audio input; long reference audios can waste
|
|
||||||
tokens so we cap at max_s seconds by default (3-10s is the recommended range).
|
|
||||||
"""
|
|
||||||
data, sr = sf.read(path)
|
|
||||||
if data.ndim > 1:
|
|
||||||
data = data[:, 0]
|
|
||||||
data = data.astype(np.float32)
|
|
||||||
if sr != target_sr:
|
|
||||||
try:
|
|
||||||
# Demo client only; librosa here resamples audio before sending to
|
|
||||||
# the server. The fallback below covers environments without it.
|
|
||||||
import librosa # noqa: TID251
|
|
||||||
|
|
||||||
data = librosa.resample(data, orig_sr=sr, target_sr=target_sr)
|
|
||||||
except ImportError:
|
|
||||||
from math import gcd
|
|
||||||
|
|
||||||
g = gcd(sr, target_sr)
|
|
||||||
up, down = target_sr // g, sr // g
|
|
||||||
idx = (np.arange(len(data) * up) / up).astype(np.int64)
|
|
||||||
data = data[idx][::down]
|
|
||||||
sr = target_sr
|
|
||||||
if max_s is not None and len(data) > int(max_s * sr):
|
|
||||||
data = data[: int(max_s * sr)]
|
|
||||||
return audio_to_base64_data_url(data, sr)
|
|
||||||
|
|
||||||
|
|
||||||
def video_to_base64_data_url(path: str) -> str:
|
|
||||||
p = Path(path)
|
|
||||||
if not p.exists():
|
|
||||||
raise FileNotFoundError(f"Video not found: {path}")
|
|
||||||
ext = p.suffix.lower()
|
|
||||||
mime = {
|
|
||||||
".mp4": "video/mp4",
|
|
||||||
".webm": "video/webm",
|
|
||||||
".mov": "video/quicktime",
|
|
||||||
".avi": "video/x-msvideo",
|
|
||||||
".mkv": "video/x-matroska",
|
|
||||||
}.get(ext, "video/mp4")
|
|
||||||
b64 = base64.b64encode(p.read_bytes()).decode("utf-8")
|
|
||||||
return f"data:{mime};base64,{b64}"
|
|
||||||
|
|
||||||
|
|
||||||
def process_audio_input(audio_input: Any | None) -> tuple[np.ndarray, int] | None:
|
|
||||||
"""Normalize Gradio audio input to (np.ndarray mono float32, sample_rate).
|
|
||||||
|
|
||||||
Gradio's `gr.Audio(type="numpy")` yields (sr, ndarray); we also accept
|
|
||||||
filesystem paths and a few nested tuple variants that appear in different
|
|
||||||
Gradio versions.
|
|
||||||
"""
|
|
||||||
if audio_input is None:
|
|
||||||
return None
|
|
||||||
|
|
||||||
def _from_path(p: str) -> tuple[np.ndarray, int] | None:
|
|
||||||
if not p:
|
|
||||||
return None
|
|
||||||
fp = Path(p)
|
|
||||||
if not fp.exists():
|
|
||||||
return None
|
|
||||||
data, sr = sf.read(fp)
|
|
||||||
if data.ndim > 1:
|
|
||||||
data = data[:, 0]
|
|
||||||
return data.astype(np.float32), int(sr)
|
|
||||||
|
|
||||||
audio_np: np.ndarray | None = None
|
|
||||||
sr: int | None = None
|
|
||||||
|
|
||||||
if isinstance(audio_input, tuple) and len(audio_input) == 2:
|
|
||||||
a, b = audio_input
|
|
||||||
if isinstance(a, int | float) and isinstance(b, np.ndarray):
|
|
||||||
sr, audio_np = int(a), b
|
|
||||||
elif isinstance(a, str):
|
|
||||||
loaded = _from_path(a)
|
|
||||||
if loaded is not None:
|
|
||||||
audio_np, sr = loaded
|
|
||||||
elif isinstance(audio_input, str):
|
|
||||||
loaded = _from_path(audio_input)
|
|
||||||
if loaded is not None:
|
|
||||||
audio_np, sr = loaded
|
|
||||||
|
|
||||||
if audio_np is None or sr is None:
|
|
||||||
return None
|
|
||||||
if audio_np.ndim > 1:
|
|
||||||
audio_np = audio_np[:, 0]
|
|
||||||
return audio_np.astype(np.float32), sr
|
|
||||||
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
# Per-model inference config
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
|
|
||||||
|
|
||||||
class ModelEndpoint:
|
|
||||||
"""Endpoint + TTS-trigger convention for one MiniCPM-o variant."""
|
|
||||||
|
|
||||||
def __init__(self, name: str, api_base: str, model_path: str):
|
|
||||||
self.name = name
|
|
||||||
self.api_base = api_base
|
|
||||||
self.model_path = model_path
|
|
||||||
self.client = OpenAI(api_key="EMPTY", base_url=api_base)
|
|
||||||
self._ref_audio_data_url: str | None = None
|
|
||||||
|
|
||||||
@property
|
|
||||||
def ref_audio_path(self) -> str | None:
|
|
||||||
rel = _DEFAULT_REF_AUDIO.get(self.name)
|
|
||||||
if not rel:
|
|
||||||
return None
|
|
||||||
full = os.path.join(self.model_path, rel)
|
|
||||||
return full if os.path.exists(full) else None
|
|
||||||
|
|
||||||
def get_ref_audio_data_url(self) -> str | None:
|
|
||||||
"""Lazily load and cache the default reference audio as a 16kHz mono
|
|
||||||
data URL. Returns None if the reference audio file is missing."""
|
|
||||||
if self._ref_audio_data_url is not None:
|
|
||||||
return self._ref_audio_data_url
|
|
||||||
p = self.ref_audio_path
|
|
||||||
if not p:
|
|
||||||
print(f"[{self.name}] no default reference audio found; falling back to text-only system prompt")
|
|
||||||
return None
|
|
||||||
try:
|
|
||||||
self._ref_audio_data_url = ref_audio_to_data_url(p)
|
|
||||||
print(f"[{self.name}] loaded reference audio: {p}")
|
|
||||||
except Exception as e:
|
|
||||||
print(f"[{self.name}] failed to load reference audio {p}: {e}")
|
|
||||||
return None
|
|
||||||
return self._ref_audio_data_url
|
|
||||||
|
|
||||||
def build_extras(self) -> dict[str, Any]:
|
|
||||||
"""Extra kwargs passed to chat.completions.create for this model.
|
|
||||||
|
|
||||||
4.5: use_tts_template=True chat-template flag appends <|tts_bos|>.
|
|
||||||
"""
|
|
||||||
if self.name == MINICPMO45:
|
|
||||||
return {
|
|
||||||
"extra_body": {
|
|
||||||
"chat_template_kwargs": {"use_tts_template": True},
|
|
||||||
"modalities": ["text", "audio"],
|
|
||||||
},
|
|
||||||
}
|
|
||||||
return {}
|
|
||||||
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
# System prompt (matches MiniCPMO.get_sys_prompt("audio_assistant") in the
|
|
||||||
# official MiniCPM-o-Demo). When TTS is disabled we fall back to a plain
|
|
||||||
# text-only system prompt.
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
|
|
||||||
|
|
||||||
def build_audio_assistant_system(
|
|
||||||
ref_audio_data_url: str | None,
|
|
||||||
language: str = "zh",
|
|
||||||
) -> dict[str, Any]:
|
|
||||||
"""Return an OpenAI-style system message carrying the reference audio."""
|
|
||||||
txt = _AUDIO_ASSISTANT_PROMPT.get(language, _AUDIO_ASSISTANT_PROMPT["zh"])
|
|
||||||
if ref_audio_data_url is None:
|
|
||||||
return {
|
|
||||||
"role": "system",
|
|
||||||
"content": (
|
|
||||||
"You are MiniCPM-o, a helpful multimodal assistant that can "
|
|
||||||
"understand images, audio and video, and respond in text and speech."
|
|
||||||
),
|
|
||||||
}
|
|
||||||
return {
|
|
||||||
"role": "system",
|
|
||||||
"content": [
|
|
||||||
{"type": "text", "text": txt["prefix"]},
|
|
||||||
{"type": "audio_url", "audio_url": {"url": ref_audio_data_url}},
|
|
||||||
{"type": "text", "text": txt["suffix"]},
|
|
||||||
],
|
|
||||||
}
|
|
||||||
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
# Core inference function
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
|
|
||||||
|
|
||||||
def run_inference(
|
|
||||||
endpoint: ModelEndpoint,
|
|
||||||
user_prompt: str,
|
|
||||||
audio_tuple: tuple[np.ndarray, int] | None,
|
|
||||||
image_pil: Image.Image | None,
|
|
||||||
video_path: str | None,
|
|
||||||
enable_tts: bool,
|
|
||||||
max_tokens: int,
|
|
||||||
temperature: float,
|
|
||||||
) -> Iterable[tuple[str, tuple[int, np.ndarray] | None]]:
|
|
||||||
user_prompt = user_prompt or ""
|
|
||||||
if not user_prompt.strip() and not audio_tuple and not image_pil and not video_path:
|
|
||||||
yield "Please provide a text prompt or at least one multimodal input.", None
|
|
||||||
return
|
|
||||||
|
|
||||||
try:
|
|
||||||
content: list[dict[str, Any]] = []
|
|
||||||
|
|
||||||
if audio_tuple is not None:
|
|
||||||
audio_np, sr = audio_tuple
|
|
||||||
content.append(
|
|
||||||
{"type": "audio_url", "audio_url": {"url": audio_to_base64_data_url(audio_np, sr)}},
|
|
||||||
)
|
|
||||||
|
|
||||||
if image_pil is not None:
|
|
||||||
img = image_pil if image_pil.mode == "RGB" else image_pil.convert("RGB")
|
|
||||||
content.append(
|
|
||||||
{"type": "image_url", "image_url": {"url": image_to_base64_data_url(img)}},
|
|
||||||
)
|
|
||||||
|
|
||||||
if video_path:
|
|
||||||
content.append(
|
|
||||||
{"type": "video_url", "video_url": {"url": video_to_base64_data_url(video_path)}},
|
|
||||||
)
|
|
||||||
|
|
||||||
if user_prompt.strip():
|
|
||||||
content.append({"type": "text", "text": user_prompt})
|
|
||||||
|
|
||||||
if enable_tts:
|
|
||||||
ref_url = endpoint.get_ref_audio_data_url()
|
|
||||||
system_msg = build_audio_assistant_system(ref_url, language="zh")
|
|
||||||
else:
|
|
||||||
system_msg = {
|
|
||||||
"role": "system",
|
|
||||||
"content": (
|
|
||||||
"You are MiniCPM-o, a helpful multimodal assistant that can "
|
|
||||||
"understand images, audio and video, and respond in text and speech."
|
|
||||||
),
|
|
||||||
}
|
|
||||||
|
|
||||||
messages = [
|
|
||||||
system_msg,
|
|
||||||
{"role": "user", "content": content},
|
|
||||||
]
|
|
||||||
|
|
||||||
kwargs: dict[str, Any] = {
|
|
||||||
"model": endpoint.model_path,
|
|
||||||
"messages": messages,
|
|
||||||
"max_tokens": int(max_tokens),
|
|
||||||
"temperature": float(temperature),
|
|
||||||
}
|
|
||||||
if enable_tts:
|
|
||||||
extras = endpoint.build_extras()
|
|
||||||
# Merge extra_body dicts cleanly.
|
|
||||||
extra_body = dict(extras.get("extra_body", {}))
|
|
||||||
kwargs["extra_body"] = extra_body
|
|
||||||
else:
|
|
||||||
# Text-only: suppress audio output.
|
|
||||||
kwargs["extra_body"] = {"modalities": ["text"]}
|
|
||||||
|
|
||||||
completion = endpoint.client.chat.completions.create(**kwargs)
|
|
||||||
|
|
||||||
text_parts: list[str] = []
|
|
||||||
audio_out: tuple[int, np.ndarray] | None = None
|
|
||||||
|
|
||||||
for choice in completion.choices:
|
|
||||||
msg = choice.message
|
|
||||||
if getattr(msg, "content", None):
|
|
||||||
text_parts.append(msg.content)
|
|
||||||
audio_obj = getattr(msg, "audio", None)
|
|
||||||
if audio_obj and getattr(audio_obj, "data", None):
|
|
||||||
raw = base64.b64decode(audio_obj.data)
|
|
||||||
if len(raw) > 80:
|
|
||||||
data, sr = sf.read(io.BytesIO(raw))
|
|
||||||
if data.ndim > 1:
|
|
||||||
data = data[:, 0]
|
|
||||||
audio_out = (int(sr), data.astype(np.float32))
|
|
||||||
|
|
||||||
text_response = "\n\n".join(t for t in text_parts if t) or "(no text returned)"
|
|
||||||
yield text_response, audio_out
|
|
||||||
except Exception as exc: # noqa: BLE001
|
|
||||||
yield f"Inference failed: {type(exc).__name__}: {exc}", None
|
|
||||||
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
# UI
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
|
|
||||||
DEFAULT_PROMPTS = [
|
|
||||||
"你好,请用一句话介绍北京。",
|
|
||||||
"Describe what you see in the image.",
|
|
||||||
"Transcribe the audio and translate it into English.",
|
|
||||||
"What is happening in this video?",
|
|
||||||
]
|
|
||||||
|
|
||||||
|
|
||||||
def build_interface(endpoints: dict[str, ModelEndpoint], default_model: str) -> gr.Blocks:
|
|
||||||
model_choices = list(endpoints.keys())
|
|
||||||
|
|
||||||
css = """
|
|
||||||
#generate-btn button { width: 100%; font-size: 1.05rem; }
|
|
||||||
.media-row > div { min-height: 260px; }
|
|
||||||
.model-banner { color: #555; font-size: 0.9rem; }
|
|
||||||
"""
|
|
||||||
|
|
||||||
with gr.Blocks(css=css, title="vLLM-Omni · MiniCPM-o Demo") as demo:
|
|
||||||
gr.Markdown(
|
|
||||||
"# vLLM-Omni · MiniCPM-o Online Demo\n"
|
|
||||||
"Multimodal chat (text / image / audio / video in; text + speech out) for "
|
|
||||||
"MiniCPM-o 4.5 served by vllm-omni."
|
|
||||||
)
|
|
||||||
|
|
||||||
with gr.Row():
|
|
||||||
model_dd = gr.Dropdown(
|
|
||||||
label="Model",
|
|
||||||
choices=model_choices,
|
|
||||||
value=default_model,
|
|
||||||
scale=2,
|
|
||||||
)
|
|
||||||
endpoint_info = gr.Markdown(
|
|
||||||
elem_classes="model-banner",
|
|
||||||
value=f"`{endpoints[default_model].api_base}` → `{endpoints[default_model].model_path}`",
|
|
||||||
)
|
|
||||||
|
|
||||||
with gr.Row():
|
|
||||||
prompt_box = gr.Textbox(
|
|
||||||
label="Text Prompt",
|
|
||||||
placeholder="Ask the model anything. Combine with image / audio / video below.",
|
|
||||||
lines=3,
|
|
||||||
value=DEFAULT_PROMPTS[0],
|
|
||||||
scale=3,
|
|
||||||
)
|
|
||||||
with gr.Column(scale=1):
|
|
||||||
tts_checkbox = gr.Checkbox(label="Generate speech output (TTS)", value=True)
|
|
||||||
max_tokens_slider = gr.Slider(label="Max tokens", minimum=32, maximum=4096, value=1024, step=32)
|
|
||||||
temperature_slider = gr.Slider(label="Temperature", minimum=0.0, maximum=1.5, value=0.7, step=0.05)
|
|
||||||
|
|
||||||
with gr.Row(elem_classes="media-row"):
|
|
||||||
image_input = gr.Image(label="Image (optional)", type="pil", sources=["upload"])
|
|
||||||
audio_input = gr.Audio(
|
|
||||||
label="Audio (optional)",
|
|
||||||
type="numpy",
|
|
||||||
sources=["upload", "microphone"],
|
|
||||||
)
|
|
||||||
video_input = gr.Video(label="Video (optional)", sources=["upload"])
|
|
||||||
|
|
||||||
with gr.Row():
|
|
||||||
generate_btn = gr.Button("Generate", variant="primary", elem_id="generate-btn")
|
|
||||||
clear_btn = gr.Button("Clear")
|
|
||||||
|
|
||||||
with gr.Row():
|
|
||||||
text_output = gr.Textbox(label="Text Output", lines=10, scale=2)
|
|
||||||
audio_output = gr.Audio(label="Audio Output", interactive=False, scale=1)
|
|
||||||
|
|
||||||
gr.Examples(
|
|
||||||
examples=[
|
|
||||||
["你好,请用一句话介绍北京。"],
|
|
||||||
["Describe what you see in detail."],
|
|
||||||
["请先识别这段语音的内容,再用中文回复说话人。"],
|
|
||||||
["What is happening in this video?"],
|
|
||||||
],
|
|
||||||
inputs=[prompt_box],
|
|
||||||
label="Quick prompts",
|
|
||||||
)
|
|
||||||
|
|
||||||
def _on_model_change(name: str) -> str:
|
|
||||||
ep = endpoints[name]
|
|
||||||
return f"`{ep.api_base}` → `{ep.model_path}`"
|
|
||||||
|
|
||||||
model_dd.change(_on_model_change, inputs=[model_dd], outputs=[endpoint_info])
|
|
||||||
|
|
||||||
def _run(
|
|
||||||
model_name: str,
|
|
||||||
prompt: str,
|
|
||||||
image_pil: Image.Image | None,
|
|
||||||
audio_in: Any,
|
|
||||||
video_path: str | None,
|
|
||||||
enable_tts: bool,
|
|
||||||
max_tokens: int,
|
|
||||||
temperature: float,
|
|
||||||
):
|
|
||||||
ep = endpoints[model_name]
|
|
||||||
audio_tuple = process_audio_input(audio_in)
|
|
||||||
yield from run_inference(
|
|
||||||
ep,
|
|
||||||
prompt,
|
|
||||||
audio_tuple,
|
|
||||||
image_pil,
|
|
||||||
video_path,
|
|
||||||
enable_tts,
|
|
||||||
max_tokens,
|
|
||||||
temperature,
|
|
||||||
)
|
|
||||||
|
|
||||||
generate_btn.click(
|
|
||||||
_run,
|
|
||||||
inputs=[
|
|
||||||
model_dd,
|
|
||||||
prompt_box,
|
|
||||||
image_input,
|
|
||||||
audio_input,
|
|
||||||
video_input,
|
|
||||||
tts_checkbox,
|
|
||||||
max_tokens_slider,
|
|
||||||
temperature_slider,
|
|
||||||
],
|
|
||||||
outputs=[text_output, audio_output],
|
|
||||||
)
|
|
||||||
|
|
||||||
def _clear():
|
|
||||||
return "", None, None, None, "", None
|
|
||||||
|
|
||||||
clear_btn.click(
|
|
||||||
_clear,
|
|
||||||
outputs=[
|
|
||||||
prompt_box,
|
|
||||||
image_input,
|
|
||||||
audio_input,
|
|
||||||
video_input,
|
|
||||||
text_output,
|
|
||||||
audio_output,
|
|
||||||
],
|
|
||||||
)
|
|
||||||
|
|
||||||
demo.queue()
|
|
||||||
return demo
|
|
||||||
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
# Entry point
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
|
|
||||||
|
|
||||||
def parse_args() -> argparse.Namespace:
|
|
||||||
p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
|
||||||
p.add_argument("--minicpmo45-api-base", default=os.environ.get("MINICPMO45_API_BASE", "http://localhost:8099/v1"))
|
|
||||||
p.add_argument(
|
|
||||||
"--minicpmo45-model",
|
|
||||||
default=os.environ.get("MINICPMO45_MODEL", "openbmb/MiniCPM-o-4_5"),
|
|
||||||
)
|
|
||||||
p.add_argument("--host", default="0.0.0.0")
|
|
||||||
p.add_argument("--port", type=int, default=7862)
|
|
||||||
p.add_argument("--share", action="store_true")
|
|
||||||
p.add_argument(
|
|
||||||
"--ssl-certfile",
|
|
||||||
default=os.environ.get("GRADIO_SSL_CERTFILE", ""),
|
|
||||||
help="Path to TLS certificate (PEM). Browsers require HTTPS to access the microphone.",
|
|
||||||
)
|
|
||||||
p.add_argument(
|
|
||||||
"--ssl-keyfile",
|
|
||||||
default=os.environ.get("GRADIO_SSL_KEYFILE", ""),
|
|
||||||
help="Path to TLS private key (PEM).",
|
|
||||||
)
|
|
||||||
p.add_argument(
|
|
||||||
"--ssl-verify",
|
|
||||||
action="store_true",
|
|
||||||
default=os.environ.get("GRADIO_SSL_VERIFY", "0") == "1",
|
|
||||||
help="Require browser to verify the certificate (off by default; self-signed certs should stay off).",
|
|
||||||
)
|
|
||||||
return p.parse_args()
|
|
||||||
|
|
||||||
|
|
||||||
def _ping(api_base: str, timeout: float = 3.0) -> bool:
|
|
||||||
import urllib.request
|
|
||||||
|
|
||||||
try:
|
|
||||||
with urllib.request.urlopen(api_base.rstrip("/").replace("/v1", "/health"), timeout=timeout):
|
|
||||||
return True
|
|
||||||
except Exception:
|
|
||||||
return False
|
|
||||||
|
|
||||||
|
|
||||||
def main() -> None:
|
|
||||||
args = parse_args()
|
|
||||||
endpoints: dict[str, ModelEndpoint] = {}
|
|
||||||
|
|
||||||
if args.minicpmo45_api_base and args.minicpmo45_model:
|
|
||||||
ok = _ping(args.minicpmo45_api_base)
|
|
||||||
status = "reachable" if ok else "not reachable (will error on generate)"
|
|
||||||
print(f"[4.5] {args.minicpmo45_api_base} ({status})")
|
|
||||||
endpoints[MINICPMO45] = ModelEndpoint(MINICPMO45, args.minicpmo45_api_base, args.minicpmo45_model)
|
|
||||||
|
|
||||||
if not endpoints:
|
|
||||||
raise SystemExit("No endpoints configured. Pass --minicpmo45-api-base/--minicpmo45-model.")
|
|
||||||
|
|
||||||
default_model = next(iter(endpoints.keys()))
|
|
||||||
demo = build_interface(endpoints, default_model)
|
|
||||||
|
|
||||||
launch_kwargs: dict = {
|
|
||||||
"server_name": args.host,
|
|
||||||
"server_port": args.port,
|
|
||||||
"share": args.share,
|
|
||||||
}
|
|
||||||
if args.ssl_certfile and args.ssl_keyfile:
|
|
||||||
launch_kwargs.update(
|
|
||||||
ssl_certfile=args.ssl_certfile,
|
|
||||||
ssl_keyfile=args.ssl_keyfile,
|
|
||||||
ssl_verify=args.ssl_verify,
|
|
||||||
)
|
|
||||||
print(f"[tls] cert={args.ssl_certfile} key={args.ssl_keyfile} verify={args.ssl_verify}")
|
|
||||||
demo.launch(**launch_kwargs)
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
main()
|
|
||||||
@@ -1,36 +0,0 @@
|
|||||||
#!/bin/bash
|
|
||||||
# Launch the MiniCPM-o 4.5 gradio demo.
|
|
||||||
#
|
|
||||||
# Prereq:
|
|
||||||
# Start a vllm-omni OpenAI server for MiniCPM-o 4.5 on :8099 (see the
|
|
||||||
# 8x4090 stage config under vllm_omni/model_executor/stage_configs).
|
|
||||||
set -e
|
|
||||||
|
|
||||||
HERE="$(cd "$(dirname "$0")" && pwd)"
|
|
||||||
|
|
||||||
: "${MINICPMO45_API_BASE:=http://localhost:8099/v1}"
|
|
||||||
: "${MINICPMO45_MODEL:=openbmb/MiniCPM-o-4_5}"
|
|
||||||
: "${GRADIO_HOST:=0.0.0.0}"
|
|
||||||
: "${GRADIO_PORT:=7862}"
|
|
||||||
# HTTPS (browsers require a secure context for microphone access).
|
|
||||||
# Set GRADIO_SSL_CERTFILE / GRADIO_SSL_KEYFILE to enable TLS.
|
|
||||||
: "${GRADIO_SSL_CERTFILE:=}"
|
|
||||||
: "${GRADIO_SSL_KEYFILE:=}"
|
|
||||||
|
|
||||||
export MINICPMO45_API_BASE MINICPMO45_MODEL
|
|
||||||
|
|
||||||
SSL_ARGS=()
|
|
||||||
if [ -n "$GRADIO_SSL_CERTFILE" ] && [ -n "$GRADIO_SSL_KEYFILE" ] \
|
|
||||||
&& [ -f "$GRADIO_SSL_CERTFILE" ] && [ -f "$GRADIO_SSL_KEYFILE" ]; then
|
|
||||||
SSL_ARGS=(--ssl-certfile "$GRADIO_SSL_CERTFILE" --ssl-keyfile "$GRADIO_SSL_KEYFILE")
|
|
||||||
echo "HTTPS enabled: cert=$GRADIO_SSL_CERTFILE key=$GRADIO_SSL_KEYFILE"
|
|
||||||
else
|
|
||||||
echo "HTTPS disabled (cert/key not found). Microphone won't work on remote browsers."
|
|
||||||
fi
|
|
||||||
|
|
||||||
exec python "$HERE/gradio_demo.py" \
|
|
||||||
--minicpmo45-api-base "$MINICPMO45_API_BASE" \
|
|
||||||
--minicpmo45-model "$MINICPMO45_MODEL" \
|
|
||||||
--host "$GRADIO_HOST" \
|
|
||||||
--port "$GRADIO_PORT" \
|
|
||||||
"${SSL_ARGS[@]}"
|
|
||||||
@@ -6,8 +6,7 @@ from typing import NamedTuple
|
|||||||
import requests
|
import requests
|
||||||
from openai import OpenAI
|
from openai import OpenAI
|
||||||
from vllm.assets.audio import AudioAsset
|
from vllm.assets.audio import AudioAsset
|
||||||
|
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||||
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
|
|
||||||
|
|
||||||
SEED = 42
|
SEED = 42
|
||||||
|
|
||||||
@@ -198,14 +197,7 @@ def _build_prompt_for_query_type(
|
|||||||
return query_func(custom_prompt=custom_prompt)
|
return query_func(custom_prompt=custom_prompt)
|
||||||
if query_type == "use_audio_in_video":
|
if query_type == "use_audio_in_video":
|
||||||
return query_func(video_path=video_path, custom_prompt=custom_prompt)
|
return query_func(video_path=video_path, custom_prompt=custom_prompt)
|
||||||
if query_type == "use_mixed_modalities":
|
# use_mixed_modalities / use_multi_audios
|
||||||
return query_func(
|
|
||||||
video_path=video_path,
|
|
||||||
image_path=image_path,
|
|
||||||
audio_path=audio_path,
|
|
||||||
custom_prompt=custom_prompt,
|
|
||||||
)
|
|
||||||
# use_multi_audios
|
|
||||||
return query_func(custom_prompt=custom_prompt)
|
return query_func(custom_prompt=custom_prompt)
|
||||||
|
|
||||||
|
|
||||||
@@ -486,7 +478,7 @@ def run_multimodal_generation(args, client: OpenAI) -> None:
|
|||||||
|
|
||||||
|
|
||||||
def parse_args():
|
def parse_args():
|
||||||
parser = TrackingArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
|
parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
"--query-type",
|
"--query-type",
|
||||||
"-q",
|
"-q",
|
||||||
|
|||||||
@@ -412,31 +412,6 @@ curl -X POST http://localhost:8091/v1/audio/voices \
|
|||||||
```
|
```
|
||||||
Uploaded voices are then usable as `voice="custom_voice_1"` on subsequent requests.
|
Uploaded voices are then usable as `voice="custom_voice_1"` on subsequent requests.
|
||||||
|
|
||||||
### Precomputed custom voices
|
|
||||||
For reused Base voice-cloning speakers, precompute the reference artifacts once and load them at server startup:
|
|
||||||
```bash
|
|
||||||
python qwen3_tts/precompute_custom_voice.py \
|
|
||||||
--model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
|
|
||||||
--voice-name alice \
|
|
||||||
--ref-audio /path/to/reference.wav \
|
|
||||||
--ref-text "Original transcript of the reference audio" \
|
|
||||||
--mode icl \
|
|
||||||
--output-dir /path/to/custom_voices
|
|
||||||
```
|
|
||||||
`--mode icl` stores both `speaker_embedding` and `ref_code`; `--mode xvec` stores only the speaker embedding. Add the output directory to a deploy config:
|
|
||||||
```yaml
|
|
||||||
custom_voice_dir: /path/to/custom_voices
|
|
||||||
```
|
|
||||||
Then start the server with that config and call the Speech API with only the voice name:
|
|
||||||
```bash
|
|
||||||
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base --omni --deploy-config /path/to/qwen3_tts_custom_voice.yaml
|
|
||||||
|
|
||||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
|
||||||
-H "Content-Type: application/json" \
|
|
||||||
-d '{"input":"Hello from a precomputed voice.","voice":"alice","task_type":"Base"}' \
|
|
||||||
--output alice.wav
|
|
||||||
```
|
|
||||||
|
|
||||||
### Streaming PCM
|
### Streaming PCM
|
||||||
```bash
|
```bash
|
||||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
curl -X POST http://localhost:8091/v1/audio/speech \
|
||||||
@@ -472,8 +447,7 @@ python qwen3_tts/streaming_speech_client.py --text "..." --simulate-stt --stt-de
|
|||||||
`qwen3_tts/batch_speech_client.py` issues many concurrent requests for throughput measurement.
|
`qwen3_tts/batch_speech_client.py` issues many concurrent requests for throughput measurement.
|
||||||
|
|
||||||
### Notes
|
### Notes
|
||||||
- Base voice cloning has uniproc-vs-mp tradeoffs depending on per-request reference audio cost; see the executor-backend section above.
|
- Base voice cloning has per-request reference-audio cost; the uniproc default keeps IPC overhead off the critical path. See the executor-backend section above for background.
|
||||||
- With async chunking, Qwen3-TTS Base voice cloning sends the full reference context in the first Code2Wav packet, then caches that prefix on the Code2Wav stage for follow-up chunks in the same request.
|
|
||||||
- `vllm_omni/deploy/qwen3_tts.yaml` is the default deploy config (loaded by HF `model_type`); per-stage runtime overrides are available via `--stage-N-<field> <value>`.
|
- `vllm_omni/deploy/qwen3_tts.yaml` is the default deploy config (loaded by HF `model_type`); per-stage runtime overrides are available via `--stage-N-<field> <value>`.
|
||||||
|
|
||||||
---
|
---
|
||||||
@@ -500,23 +474,6 @@ python voxcpm2/openai_speech_client.py \
|
|||||||
```
|
```
|
||||||
The `ref_audio` field accepts local file paths (auto-base64), HTTP URLs, or `data:audio/wav;base64,...` data URIs.
|
The `ref_audio` field accepts local file paths (auto-base64), HTTP URLs, or `data:audio/wav;base64,...` data URIs.
|
||||||
|
|
||||||
### Precomputed custom voices
|
|
||||||
For repeated VoxCPM2 speakers, precompute the prompt cache and load it through `custom_voice_dir`:
|
|
||||||
```bash
|
|
||||||
python voxcpm2/precompute_custom_voice.py \
|
|
||||||
--model openbmb/VoxCPM2 \
|
|
||||||
--voice-name alice \
|
|
||||||
--ref-audio /path/to/reference.wav \
|
|
||||||
--mode ref_continuation \
|
|
||||||
--prompt-text "Original transcript of the reference audio" \
|
|
||||||
--output-dir /path/to/custom_voices
|
|
||||||
```
|
|
||||||
Add the output directory to the deploy config:
|
|
||||||
```yaml
|
|
||||||
custom_voice_dir: /path/to/custom_voices
|
|
||||||
```
|
|
||||||
After startup, `/v1/audio/voices` lists `alice`, and `/v1/audio/speech` can use `voice="alice"` without sending `ref_audio`.
|
|
||||||
|
|
||||||
### Gradio demo (gapless streaming via AudioWorklet)
|
### Gradio demo (gapless streaming via AudioWorklet)
|
||||||
```bash
|
```bash
|
||||||
python voxcpm2/gradio_demo.py
|
python voxcpm2/gradio_demo.py
|
||||||
|
|||||||
@@ -1,40 +0,0 @@
|
|||||||
# Higgs-Audio V3 Online Serving
|
|
||||||
|
|
||||||
## Start the server
|
|
||||||
|
|
||||||
```bash
|
|
||||||
# Default: GPU 0, port 8095
|
|
||||||
./examples/online_serving/text_to_speech/higgs_audio_v3/run_server.sh
|
|
||||||
|
|
||||||
# Custom GPU / port
|
|
||||||
PORT=8096 GPUS=0,1 ./examples/online_serving/text_to_speech/higgs_audio_v3/run_server.sh
|
|
||||||
```
|
|
||||||
|
|
||||||
## Plain text TTS
|
|
||||||
|
|
||||||
```bash
|
|
||||||
python examples/online_serving/text_to_speech/higgs_audio_v3/batch_speech_client.py \
|
|
||||||
--base-url http://localhost:8095 \
|
|
||||||
--output-dir /tmp/higgs_v3_batch \
|
|
||||||
--prompts "Hello world." "The quick brown fox jumps over the lazy dog."
|
|
||||||
```
|
|
||||||
|
|
||||||
## Voice clone
|
|
||||||
|
|
||||||
```bash
|
|
||||||
python examples/online_serving/text_to_speech/higgs_audio_v3/batch_speech_client.py \
|
|
||||||
--base-url http://localhost:8095 \
|
|
||||||
--output-dir /tmp/higgs_v3_clone \
|
|
||||||
--ref-audio path/to/reference.wav \
|
|
||||||
--ref-text "transcript of the reference clip" \
|
|
||||||
--prompts "Text to synthesize in the cloned voice."
|
|
||||||
```
|
|
||||||
|
|
||||||
## curl example
|
|
||||||
|
|
||||||
```bash
|
|
||||||
curl -X POST http://localhost:8095/v1/audio/speech \
|
|
||||||
-H "Content-Type: application/json" \
|
|
||||||
-d '{"model": "higgs_audio_v3", "input": "Hello world."}' \
|
|
||||||
--output hello.wav
|
|
||||||
```
|
|
||||||
@@ -1,115 +0,0 @@
|
|||||||
# SPDX-License-Identifier: Apache-2.0
|
|
||||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
||||||
"""Batch client for the higgs-audio v3 online server.
|
|
||||||
|
|
||||||
Sends prompts to ``/v1/audio/speech`` and saves the returned WAV files.
|
|
||||||
|
|
||||||
Usage (plain text -> speech):
|
|
||||||
|
|
||||||
python examples/online_serving/text_to_speech/higgs_audio_v3/batch_speech_client.py \
|
|
||||||
--base-url http://localhost:8095 \
|
|
||||||
--output-dir /tmp/higgs_v3_batch \
|
|
||||||
--prompts "Hello world." "The quick brown fox jumps over the lazy dog."
|
|
||||||
|
|
||||||
Usage (voice clone):
|
|
||||||
|
|
||||||
python examples/online_serving/text_to_speech/higgs_audio_v3/batch_speech_client.py \
|
|
||||||
--base-url http://localhost:8095 \
|
|
||||||
--output-dir /tmp/higgs_v3_clone \
|
|
||||||
--ref-audio path/to/reference.wav \
|
|
||||||
--ref-text "the transcript of the reference clip" \
|
|
||||||
--prompts "Hello world."
|
|
||||||
"""
|
|
||||||
|
|
||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
import argparse
|
|
||||||
import base64
|
|
||||||
import sys
|
|
||||||
from pathlib import Path
|
|
||||||
|
|
||||||
DEFAULT_PROMPTS = (
|
|
||||||
"Hello world.",
|
|
||||||
"The quick brown fox jumps over the lazy dog.",
|
|
||||||
"Today is a beautiful day for a walk in the park.",
|
|
||||||
"Innovation distinguishes between a leader and a follower.",
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
def _slug(text: str) -> str:
|
|
||||||
import re
|
|
||||||
|
|
||||||
s = re.sub(r"\s+", "_", text.strip().lower())
|
|
||||||
return re.sub(r"[^a-z0-9_]+", "", s)[:32] or "prompt"
|
|
||||||
|
|
||||||
|
|
||||||
def main() -> int:
|
|
||||||
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
|
||||||
parser.add_argument("--base-url", default="http://localhost:8095")
|
|
||||||
parser.add_argument("--model", default="higgs_audio_v3")
|
|
||||||
parser.add_argument("--prompts", nargs="+", default=list(DEFAULT_PROMPTS))
|
|
||||||
parser.add_argument("--output-dir", type=Path, default=Path("/tmp/higgs_v3_batch"))
|
|
||||||
parser.add_argument("--format", choices=("wav", "pcm"), default="wav")
|
|
||||||
parser.add_argument("--max-new-tokens", type=int, default=2048)
|
|
||||||
parser.add_argument("--seed", type=int, default=42)
|
|
||||||
parser.add_argument("--timeout-s", type=float, default=120.0)
|
|
||||||
parser.add_argument(
|
|
||||||
"--ref-audio",
|
|
||||||
type=Path,
|
|
||||||
default=None,
|
|
||||||
help="Reference clip for voice clone (WAV/FLAC/MP3 path). Pair with --ref-text.",
|
|
||||||
)
|
|
||||||
parser.add_argument(
|
|
||||||
"--ref-text",
|
|
||||||
type=str,
|
|
||||||
default=None,
|
|
||||||
help="Transcript of the reference clip. Optional but improves fidelity.",
|
|
||||||
)
|
|
||||||
args = parser.parse_args()
|
|
||||||
|
|
||||||
ref_audio_data_url: str | None = None
|
|
||||||
if args.ref_audio is not None:
|
|
||||||
if not args.ref_audio.exists():
|
|
||||||
print(f"ref-audio file not found: {args.ref_audio}", file=sys.stderr)
|
|
||||||
return 2
|
|
||||||
mime = "audio/wav" if args.ref_audio.suffix.lower() == ".wav" else "audio/mpeg"
|
|
||||||
ref_b64 = base64.b64encode(args.ref_audio.read_bytes()).decode("ascii")
|
|
||||||
ref_audio_data_url = f"data:{mime};base64,{ref_b64}"
|
|
||||||
|
|
||||||
try:
|
|
||||||
import httpx
|
|
||||||
except ImportError:
|
|
||||||
print("this client needs `httpx`. Install with `pip install httpx`.", file=sys.stderr)
|
|
||||||
return 2
|
|
||||||
|
|
||||||
args.output_dir.mkdir(parents=True, exist_ok=True)
|
|
||||||
url = args.base_url.rstrip("/") + "/v1/audio/speech"
|
|
||||||
failures = 0
|
|
||||||
with httpx.Client(timeout=args.timeout_s) as client:
|
|
||||||
for prompt in args.prompts:
|
|
||||||
payload = {
|
|
||||||
"model": args.model,
|
|
||||||
"input": prompt,
|
|
||||||
"response_format": args.format,
|
|
||||||
"max_new_tokens": args.max_new_tokens,
|
|
||||||
"seed": args.seed,
|
|
||||||
}
|
|
||||||
if ref_audio_data_url is not None:
|
|
||||||
payload["ref_audio"] = ref_audio_data_url
|
|
||||||
if args.ref_text:
|
|
||||||
payload["ref_text"] = args.ref_text
|
|
||||||
resp = client.post(url, json=payload)
|
|
||||||
if resp.status_code != 200:
|
|
||||||
print(f"[FAIL] {prompt!r} -> {resp.status_code}: {resp.text[:200]}", file=sys.stderr)
|
|
||||||
failures += 1
|
|
||||||
continue
|
|
||||||
suffix = ".wav" if args.format == "wav" else ".pcm"
|
|
||||||
out = args.output_dir / f"{_slug(prompt)}{suffix}"
|
|
||||||
out.write_bytes(resp.content)
|
|
||||||
print(f"[ ok ] {prompt!r} -> {out} ({len(resp.content)} bytes)")
|
|
||||||
|
|
||||||
return 1 if failures else 0
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
sys.exit(main())
|
|
||||||
@@ -1,32 +0,0 @@
|
|||||||
#!/bin/bash
|
|
||||||
# Launch vLLM-Omni server for higgs-audio v3.
|
|
||||||
#
|
|
||||||
# Supports plain text TTS and voice cloning via /v1/audio/speech.
|
|
||||||
#
|
|
||||||
# Usage:
|
|
||||||
# ./run_server.sh # default port 8095, GPU 0
|
|
||||||
# PORT=8096 GPUS=0,1 ./run_server.sh
|
|
||||||
# MODEL=/path/to/local/checkpoint ./run_server.sh
|
|
||||||
|
|
||||||
set -e
|
|
||||||
|
|
||||||
MODEL="${MODEL:-bosonai/higgs-audio-v3-tts-4b}"
|
|
||||||
PORT="${PORT:-8095}"
|
|
||||||
GPUS="${GPUS:-0}"
|
|
||||||
GPU_UTIL="${GPU_UTIL:-0.6}"
|
|
||||||
|
|
||||||
echo "Starting higgs-audio v3 server"
|
|
||||||
echo " MODEL=$MODEL"
|
|
||||||
echo " PORT=$PORT"
|
|
||||||
echo " CUDA_VISIBLE_DEVICES=$GPUS"
|
|
||||||
|
|
||||||
CUDA_VISIBLE_DEVICES="$GPUS" \
|
|
||||||
VLLM_USE_DEEP_GEMM=0 \
|
|
||||||
VLLM_MOE_USE_DEEP_GEMM=0 \
|
|
||||||
vllm-omni serve "$MODEL" \
|
|
||||||
--deploy-config vllm_omni/deploy/higgs_multimodal_qwen3.yaml \
|
|
||||||
--host 0.0.0.0 \
|
|
||||||
--port "$PORT" \
|
|
||||||
--gpu-memory-utilization "$GPU_UTIL" \
|
|
||||||
--trust-remote-code \
|
|
||||||
--omni
|
|
||||||
@@ -1,243 +0,0 @@
|
|||||||
"""Pre-compute Qwen3-TTS custom voice profiles.
|
|
||||||
|
|
||||||
The generated directory can be passed to the server via
|
|
||||||
``custom_voice_dir`` in ``vllm_omni/deploy/qwen3_tts.yaml``. Requests can then
|
|
||||||
use ``/v1/audio/speech`` with ``voice="<name>"`` and no per-request ref_audio.
|
|
||||||
"""
|
|
||||||
|
|
||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
import argparse
|
|
||||||
import json
|
|
||||||
import os
|
|
||||||
import sys
|
|
||||||
from pathlib import Path
|
|
||||||
from typing import Any
|
|
||||||
|
|
||||||
import numpy as np
|
|
||||||
import torch
|
|
||||||
from safetensors import safe_open
|
|
||||||
from safetensors.torch import save_file
|
|
||||||
|
|
||||||
REPO_ROOT = Path(__file__).resolve().parents[4]
|
|
||||||
if str(REPO_ROOT) not in sys.path:
|
|
||||||
sys.path.insert(0, str(REPO_ROOT))
|
|
||||||
|
|
||||||
from vllm_omni.utils.custom_voice_io import safe_voice_stem # noqa: E402
|
|
||||||
|
|
||||||
MANIFEST_NAME = "custom_voice_manifest.json"
|
|
||||||
_SPEAKER_ENCODER_SAMPLE_RATE = 24000
|
|
||||||
_SPEAKER_ENCODER_NUM_MELS = 128
|
|
||||||
_SPEAKER_ENCODER_N_FFT = 1024
|
|
||||||
_SPEAKER_ENCODER_HOP_SIZE = 256
|
|
||||||
_SPEAKER_ENCODER_WIN_SIZE = 1024
|
|
||||||
_SPEAKER_ENCODER_FMIN = 0
|
|
||||||
_SPEAKER_ENCODER_FMAX = 12000
|
|
||||||
|
|
||||||
|
|
||||||
def _resolve_model_dir(model: str) -> str:
|
|
||||||
if os.path.isdir(model):
|
|
||||||
return model
|
|
||||||
from transformers.utils.hub import cached_file
|
|
||||||
|
|
||||||
config_path = cached_file(model, "config.json")
|
|
||||||
if config_path is None:
|
|
||||||
raise FileNotFoundError(f"Could not resolve config.json for {model}")
|
|
||||||
return os.path.dirname(config_path)
|
|
||||||
|
|
||||||
|
|
||||||
def _read_audio_mono(path: str) -> tuple[np.ndarray, int]:
|
|
||||||
import soundfile as sf
|
|
||||||
|
|
||||||
wav, sr = sf.read(path, dtype="float32")
|
|
||||||
if wav.ndim > 1:
|
|
||||||
wav = wav.mean(axis=-1)
|
|
||||||
return np.asarray(wav, dtype=np.float32), int(sr)
|
|
||||||
|
|
||||||
|
|
||||||
def _resample(wav: np.ndarray, sr: int, target_sr: int) -> np.ndarray:
|
|
||||||
if sr == target_sr:
|
|
||||||
return wav.astype(np.float32)
|
|
||||||
from vllm.multimodal.audio import AudioResampler
|
|
||||||
|
|
||||||
return AudioResampler(target_sr=target_sr).resample(wav.astype(np.float32), orig_sr=int(sr))
|
|
||||||
|
|
||||||
|
|
||||||
def _load_speaker_encoder(model_dir: str, device: torch.device):
|
|
||||||
from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import Qwen3TTSConfig
|
|
||||||
from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import Qwen3TTSSpeakerEncoder
|
|
||||||
|
|
||||||
config = Qwen3TTSConfig.from_pretrained(model_dir)
|
|
||||||
encoder = Qwen3TTSSpeakerEncoder(config.speaker_encoder_config)
|
|
||||||
index_path = Path(model_dir) / "model.safetensors.index.json"
|
|
||||||
state: dict[str, torch.Tensor] = {}
|
|
||||||
if index_path.exists():
|
|
||||||
index = json.loads(index_path.read_text(encoding="utf-8"))
|
|
||||||
weight_map = index.get("weight_map", {})
|
|
||||||
shard_names = sorted({v for k, v in weight_map.items() if k.startswith("speaker_encoder.")})
|
|
||||||
shard_paths = [Path(model_dir) / name for name in shard_names]
|
|
||||||
else:
|
|
||||||
shard_paths = sorted(Path(model_dir).glob("model*.safetensors"))
|
|
||||||
for shard in shard_paths:
|
|
||||||
with safe_open(str(shard), framework="pt", device="cpu") as f:
|
|
||||||
for key in f.keys():
|
|
||||||
if key.startswith("speaker_encoder."):
|
|
||||||
state[key.removeprefix("speaker_encoder.")] = f.get_tensor(key)
|
|
||||||
if not state:
|
|
||||||
raise RuntimeError(
|
|
||||||
f"No speaker_encoder.* tensors found under {model_dir}. "
|
|
||||||
"Use a Qwen3-TTS Base checkpoint for pre-computing custom voice profiles."
|
|
||||||
)
|
|
||||||
encoder.load_state_dict(state)
|
|
||||||
encoder.to(device=device, dtype=torch.bfloat16).eval()
|
|
||||||
return config, encoder
|
|
||||||
|
|
||||||
|
|
||||||
def _speaker_embedding(
|
|
||||||
encoder: torch.nn.Module,
|
|
||||||
speaker_encoder_config: object,
|
|
||||||
wav: np.ndarray,
|
|
||||||
sr: int,
|
|
||||||
device: torch.device,
|
|
||||||
) -> torch.Tensor:
|
|
||||||
from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import mel_spectrogram
|
|
||||||
|
|
||||||
target_sr = int(getattr(speaker_encoder_config, "sample_rate", _SPEAKER_ENCODER_SAMPLE_RATE))
|
|
||||||
mel_dim = int(getattr(speaker_encoder_config, "mel_dim", _SPEAKER_ENCODER_NUM_MELS))
|
|
||||||
if target_sr != _SPEAKER_ENCODER_SAMPLE_RATE or mel_dim != _SPEAKER_ENCODER_NUM_MELS:
|
|
||||||
raise ValueError(
|
|
||||||
"Unsupported Qwen3-TTS speaker encoder mel config for precompute: "
|
|
||||||
f"sample_rate={target_sr}, mel_dim={mel_dim}. "
|
|
||||||
f"Expected sample_rate={_SPEAKER_ENCODER_SAMPLE_RATE}, mel_dim={_SPEAKER_ENCODER_NUM_MELS} "
|
|
||||||
"to match the serving-time speaker embedding path."
|
|
||||||
)
|
|
||||||
wav = _resample(wav, sr, target_sr)
|
|
||||||
wav_t = torch.from_numpy(wav).to(device=device, dtype=torch.float32)
|
|
||||||
mels = mel_spectrogram(
|
|
||||||
wav_t.unsqueeze(0),
|
|
||||||
n_fft=_SPEAKER_ENCODER_N_FFT,
|
|
||||||
num_mels=_SPEAKER_ENCODER_NUM_MELS,
|
|
||||||
sampling_rate=_SPEAKER_ENCODER_SAMPLE_RATE,
|
|
||||||
hop_size=_SPEAKER_ENCODER_HOP_SIZE,
|
|
||||||
win_size=_SPEAKER_ENCODER_WIN_SIZE,
|
|
||||||
fmin=_SPEAKER_ENCODER_FMIN,
|
|
||||||
fmax=_SPEAKER_ENCODER_FMAX,
|
|
||||||
).transpose(1, 2)
|
|
||||||
with torch.inference_mode():
|
|
||||||
return encoder(mels.to(device=device, dtype=torch.bfloat16))[0].float().cpu().contiguous()
|
|
||||||
|
|
||||||
|
|
||||||
def _ref_code(model_dir: str, wav: np.ndarray, sr: int, device: torch.device) -> torch.Tensor:
|
|
||||||
from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_tokenizer import Qwen3TTSTokenizer
|
|
||||||
|
|
||||||
tok = Qwen3TTSTokenizer.from_pretrained(
|
|
||||||
str(Path(model_dir) / "speech_tokenizer"),
|
|
||||||
torch_dtype=torch.bfloat16,
|
|
||||||
)
|
|
||||||
try:
|
|
||||||
del tok.model.decoder
|
|
||||||
tok.model.decoder = None
|
|
||||||
tok.model.encoder.to(device)
|
|
||||||
tok.device = device
|
|
||||||
except Exception:
|
|
||||||
tok.device = device
|
|
||||||
with torch.inference_mode():
|
|
||||||
enc = tok.encode(wav, sr=int(sr), return_dict=True)
|
|
||||||
codes = getattr(enc, "audio_codes", None)
|
|
||||||
if isinstance(codes, list):
|
|
||||||
codes = codes[0] if codes else None
|
|
||||||
if not isinstance(codes, torch.Tensor):
|
|
||||||
raise RuntimeError("Speech tokenizer did not return audio_codes")
|
|
||||||
if codes.ndim == 3:
|
|
||||||
codes = codes[0]
|
|
||||||
return codes.to(dtype=torch.int32, device="cpu").contiguous()
|
|
||||||
|
|
||||||
|
|
||||||
def _load_manifest(output_dir: Path, model: str, hidden_size: int) -> dict[str, Any]:
|
|
||||||
path = output_dir / MANIFEST_NAME
|
|
||||||
if path.exists():
|
|
||||||
return json.loads(path.read_text(encoding="utf-8"))
|
|
||||||
return {
|
|
||||||
"schema_version": 1,
|
|
||||||
"model_type": "qwen3_tts",
|
|
||||||
"model": model,
|
|
||||||
"hidden_size": hidden_size,
|
|
||||||
"voices": {},
|
|
||||||
}
|
|
||||||
|
|
||||||
|
|
||||||
def _write_voice(
|
|
||||||
*,
|
|
||||||
model: str,
|
|
||||||
output_dir: Path,
|
|
||||||
voice_name: str,
|
|
||||||
ref_audio: str,
|
|
||||||
ref_text: str | None,
|
|
||||||
mode: str,
|
|
||||||
speaker_description: str | None,
|
|
||||||
device: torch.device,
|
|
||||||
) -> None:
|
|
||||||
model_dir = _resolve_model_dir(model)
|
|
||||||
config, encoder = _load_speaker_encoder(model_dir, device)
|
|
||||||
wav, sr = _read_audio_mono(ref_audio)
|
|
||||||
if wav.size < 1024:
|
|
||||||
raise ValueError(f"Reference audio too short: {wav.size} samples")
|
|
||||||
|
|
||||||
tensors: dict[str, torch.Tensor] = {
|
|
||||||
"speaker_embedding": _speaker_embedding(encoder, config.speaker_encoder_config, wav, sr, device),
|
|
||||||
}
|
|
||||||
if mode == "icl":
|
|
||||||
if not ref_text or not ref_text.strip():
|
|
||||||
raise ValueError("--ref-text is required for --mode icl")
|
|
||||||
tensors["ref_code"] = _ref_code(model_dir, wav, sr, device)
|
|
||||||
|
|
||||||
output_dir.mkdir(parents=True, exist_ok=True)
|
|
||||||
filename = f"{safe_voice_stem(voice_name)}.safetensors"
|
|
||||||
save_file(tensors, str(output_dir / filename))
|
|
||||||
|
|
||||||
hidden_size = int(getattr(config.talker_config, "hidden_size", tensors["speaker_embedding"].numel()))
|
|
||||||
manifest = _load_manifest(output_dir, model, hidden_size)
|
|
||||||
entry: dict[str, Any] = {
|
|
||||||
"name": voice_name,
|
|
||||||
"file": filename,
|
|
||||||
"mode": mode,
|
|
||||||
"embedding_dim": int(tensors["speaker_embedding"].numel()),
|
|
||||||
}
|
|
||||||
if "ref_code" in tensors:
|
|
||||||
entry["ref_code_length"] = int(tensors["ref_code"].shape[0])
|
|
||||||
if ref_text:
|
|
||||||
entry["ref_text"] = ref_text
|
|
||||||
if speaker_description:
|
|
||||||
entry["speaker_description"] = speaker_description
|
|
||||||
manifest.setdefault("voices", {})[voice_name] = entry
|
|
||||||
(output_dir / MANIFEST_NAME).write_text(json.dumps(manifest, indent=2, ensure_ascii=False), encoding="utf-8")
|
|
||||||
print(f"Wrote {output_dir / filename}")
|
|
||||||
print(f"Updated {output_dir / MANIFEST_NAME}")
|
|
||||||
|
|
||||||
|
|
||||||
def main() -> None:
|
|
||||||
parser = argparse.ArgumentParser(description="Pre-compute Qwen3-TTS custom voice profile")
|
|
||||||
parser.add_argument("--model", required=True, help="Qwen3-TTS Base model path or Hugging Face ID")
|
|
||||||
parser.add_argument("--voice-name", required=True)
|
|
||||||
parser.add_argument("--ref-audio", required=True)
|
|
||||||
parser.add_argument("--ref-text", default=None)
|
|
||||||
parser.add_argument("--mode", choices=["xvec", "icl"], default="xvec")
|
|
||||||
parser.add_argument("--speaker-description", default=None)
|
|
||||||
parser.add_argument("--output-dir", required=True)
|
|
||||||
parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
|
|
||||||
args = parser.parse_args()
|
|
||||||
|
|
||||||
_write_voice(
|
|
||||||
model=args.model,
|
|
||||||
output_dir=Path(args.output_dir),
|
|
||||||
voice_name=args.voice_name,
|
|
||||||
ref_audio=args.ref_audio,
|
|
||||||
ref_text=args.ref_text,
|
|
||||||
mode=args.mode,
|
|
||||||
speaker_description=args.speaker_description,
|
|
||||||
device=torch.device(args.device),
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
main()
|
|
||||||
@@ -1,128 +0,0 @@
|
|||||||
"""Pre-compute VoxCPM2 custom voice profiles.
|
|
||||||
|
|
||||||
The generated directory can be passed to the server via
|
|
||||||
``custom_voice_dir`` in ``vllm_omni/deploy/voxcpm2.yaml``. Requests can then
|
|
||||||
use ``/v1/audio/speech`` with ``voice="<name>"`` and no per-request ref_audio.
|
|
||||||
"""
|
|
||||||
|
|
||||||
from __future__ import annotations
|
|
||||||
|
|
||||||
import argparse
|
|
||||||
import json
|
|
||||||
import sys
|
|
||||||
from pathlib import Path
|
|
||||||
from typing import Any
|
|
||||||
|
|
||||||
import torch
|
|
||||||
from safetensors.torch import save_file
|
|
||||||
|
|
||||||
REPO_ROOT = Path(__file__).resolve().parents[4]
|
|
||||||
if str(REPO_ROOT) not in sys.path:
|
|
||||||
sys.path.insert(0, str(REPO_ROOT))
|
|
||||||
|
|
||||||
from vllm_omni.utils.custom_voice_io import safe_voice_stem # noqa: E402
|
|
||||||
|
|
||||||
MANIFEST_NAME = "custom_voice_manifest.json"
|
|
||||||
|
|
||||||
|
|
||||||
def _load_tts(model: str, device: torch.device):
|
|
||||||
from vllm_omni.model_executor.models.voxcpm2.voxcpm2_import_utils import import_voxcpm2_core
|
|
||||||
|
|
||||||
VoxCPM = import_voxcpm2_core()
|
|
||||||
native = VoxCPM.from_pretrained(model, load_denoiser=False, optimize=False)
|
|
||||||
return native.tts_model.to(device).eval()
|
|
||||||
|
|
||||||
|
|
||||||
def _load_manifest(output_dir: Path, model: str) -> dict[str, Any]:
|
|
||||||
path = output_dir / MANIFEST_NAME
|
|
||||||
if path.exists():
|
|
||||||
return json.loads(path.read_text(encoding="utf-8"))
|
|
||||||
return {
|
|
||||||
"schema_version": 1,
|
|
||||||
"model_type": "voxcpm2",
|
|
||||||
"model": model,
|
|
||||||
"voices": {},
|
|
||||||
}
|
|
||||||
|
|
||||||
|
|
||||||
def _write_voice(
|
|
||||||
*,
|
|
||||||
model: str,
|
|
||||||
output_dir: Path,
|
|
||||||
voice_name: str,
|
|
||||||
ref_audio: str,
|
|
||||||
prompt_text: str | None,
|
|
||||||
mode: str,
|
|
||||||
speaker_description: str | None,
|
|
||||||
device: torch.device,
|
|
||||||
) -> None:
|
|
||||||
if mode in ("continuation", "ref_continuation") and not prompt_text:
|
|
||||||
raise ValueError("--prompt-text is required for continuation/ref_continuation modes")
|
|
||||||
|
|
||||||
tts = _load_tts(model, device)
|
|
||||||
tensors: dict[str, torch.Tensor] = {}
|
|
||||||
with torch.inference_mode():
|
|
||||||
if mode in ("reference", "ref_continuation"):
|
|
||||||
tensors["ref_audio_feat"] = tts._encode_wav(ref_audio, padding_mode="right").float().cpu().contiguous()
|
|
||||||
if mode in ("continuation", "ref_continuation"):
|
|
||||||
tensors["audio_feat"] = tts._encode_wav(ref_audio, padding_mode="left").float().cpu().contiguous()
|
|
||||||
|
|
||||||
output_dir.mkdir(parents=True, exist_ok=True)
|
|
||||||
filename = f"{safe_voice_stem(voice_name)}.safetensors"
|
|
||||||
save_file(tensors, str(output_dir / filename))
|
|
||||||
|
|
||||||
manifest = _load_manifest(output_dir, model)
|
|
||||||
entry: dict[str, Any] = {
|
|
||||||
"name": voice_name,
|
|
||||||
"file": filename,
|
|
||||||
"mode": mode,
|
|
||||||
}
|
|
||||||
if "ref_audio_feat" in tensors:
|
|
||||||
entry["ref_audio_feat_len"] = int(tensors["ref_audio_feat"].shape[0])
|
|
||||||
if "audio_feat" in tensors:
|
|
||||||
entry["audio_feat_len"] = int(tensors["audio_feat"].shape[0])
|
|
||||||
if prompt_text:
|
|
||||||
entry["prompt_text"] = prompt_text
|
|
||||||
if speaker_description:
|
|
||||||
entry["speaker_description"] = speaker_description
|
|
||||||
|
|
||||||
manifest.setdefault("voices", {})[voice_name] = entry
|
|
||||||
(output_dir / MANIFEST_NAME).write_text(json.dumps(manifest, indent=2, ensure_ascii=False), encoding="utf-8")
|
|
||||||
print(f"Wrote {output_dir / filename}")
|
|
||||||
print(f"Updated {output_dir / MANIFEST_NAME}")
|
|
||||||
|
|
||||||
|
|
||||||
def main() -> None:
|
|
||||||
parser = argparse.ArgumentParser(description="Pre-compute VoxCPM2 custom voice profile")
|
|
||||||
parser.add_argument("--model", default="openbmb/VoxCPM2", help="VoxCPM2 model path or Hugging Face ID")
|
|
||||||
parser.add_argument("--voice-name", required=True)
|
|
||||||
parser.add_argument("--ref-audio", required=True)
|
|
||||||
parser.add_argument(
|
|
||||||
"--prompt-text",
|
|
||||||
default=None,
|
|
||||||
help="Transcript of ref audio for continuation/ref_continuation modes",
|
|
||||||
)
|
|
||||||
parser.add_argument(
|
|
||||||
"--mode",
|
|
||||||
choices=["reference", "continuation", "ref_continuation"],
|
|
||||||
default="reference",
|
|
||||||
)
|
|
||||||
parser.add_argument("--speaker-description", default=None)
|
|
||||||
parser.add_argument("--output-dir", required=True)
|
|
||||||
parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
|
|
||||||
args = parser.parse_args()
|
|
||||||
|
|
||||||
_write_voice(
|
|
||||||
model=args.model,
|
|
||||||
output_dir=Path(args.output_dir),
|
|
||||||
voice_name=args.voice_name,
|
|
||||||
ref_audio=args.ref_audio,
|
|
||||||
prompt_text=args.prompt_text,
|
|
||||||
mode=args.mode,
|
|
||||||
speaker_description=args.speaker_description,
|
|
||||||
device=torch.device(args.device),
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
main()
|
|
||||||
+1
-26
@@ -64,9 +64,7 @@ dev = [
|
|||||||
"lpips>=0.1.4", # for quantization fp8
|
"lpips>=0.1.4", # for quantization fp8
|
||||||
"orjson>=3.11.7", # to make v1/embeddings API fast
|
"orjson>=3.11.7", # to make v1/embeddings API fast
|
||||||
"voxcpm>=2.0.2", # for VoxCPM2 path-locating in tests
|
"voxcpm>=2.0.2", # for VoxCPM2 path-locating in tests
|
||||||
"torchdiffeq>=0.2.5", # for covo_audio Test
|
"torchdiffeq>=0.2.5" # for covo_audio Test
|
||||||
"FlagEmbedding", # for hunyuan_image3 tests
|
|
||||||
"kernels<0.15.0" # for accuracy tests comparison with diffusers output (0.15.0 incompatible with diffusers 0.38.0)
|
|
||||||
]
|
]
|
||||||
|
|
||||||
demo = [
|
demo = [
|
||||||
@@ -78,18 +76,6 @@ local = [
|
|||||||
"FlagEmbedding",
|
"FlagEmbedding",
|
||||||
]
|
]
|
||||||
|
|
||||||
# Optional runtime dependencies for the MiniCPM-o 4.5 talker pipeline.
|
|
||||||
# The talker loads MiniCPMTTS via `trust_remote_code` from the model repo,
|
|
||||||
# and that module performs `from stepaudio2 import Token2wav` against the
|
|
||||||
# MiniCPM-o-flavored Token2wav vocoder (PyPI: `stepaudio2-minicpmo`, NOT
|
|
||||||
# upstream `stepfun-ai/Step-Audio2` — the upstream signature rejects
|
|
||||||
# `n_timesteps` and crashes at construction). A clean `vllm-omni` install
|
|
||||||
# does not pull this in by default to avoid forcing the dep onto users
|
|
||||||
# who only need other models; opt in with `pip install vllm-omni[minicpmo]`.
|
|
||||||
minicpmo = [
|
|
||||||
"stepaudio2-minicpmo",
|
|
||||||
]
|
|
||||||
|
|
||||||
docs = [
|
docs = [
|
||||||
"mkdocs>=1.5.0",
|
"mkdocs>=1.5.0",
|
||||||
"mkdocs-api-autonav",
|
"mkdocs-api-autonav",
|
||||||
@@ -115,13 +101,6 @@ Documentation = "https://vllm-omni.readthedocs.io"
|
|||||||
[project.scripts]
|
[project.scripts]
|
||||||
vllm-omni = "vllm_omni.entrypoints.cli.main:main"
|
vllm-omni = "vllm_omni.entrypoints.cli.main:main"
|
||||||
|
|
||||||
# Register vllm-omni's model registrations on every vllm subprocess. vllm
|
|
||||||
# auto-loads anything advertised under `vllm.general_plugins` in its own
|
|
||||||
# `load_general_plugins()`, so worker subprocesses spawned by vllm (which
|
|
||||||
# only import vllm, not vllm_omni) will still pick up the omni archs.
|
|
||||||
[project.entry-points."vllm.general_plugins"]
|
|
||||||
vllm_omni_register_models = "vllm_omni.engine.arg_utils:register_omni_models_to_vllm"
|
|
||||||
|
|
||||||
|
|
||||||
[tool.setuptools.packages.find]
|
[tool.setuptools.packages.find]
|
||||||
where = ["."]
|
where = ["."]
|
||||||
@@ -170,10 +149,6 @@ ignore = [
|
|||||||
[tool.ruff.lint.per-file-ignores]
|
[tool.ruff.lint.per-file-ignores]
|
||||||
"examples/**" = ["E501"] # Allow long lines in examples
|
"examples/**" = ["E501"] # Allow long lines in examples
|
||||||
"tests/**" = ["E501"] # Allow long lines in tests
|
"tests/**" = ["E501"] # Allow long lines in tests
|
||||||
# MiniCPM-o reference modeling code vendored from upstream HF/transformers/MiniCPM-o.
|
|
||||||
# These files keep upstream conventions (star imports, long lines, math-symbol
|
|
||||||
# argument names) so they stay easy to diff against upstream releases.
|
|
||||||
"vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py" = ["E402", "E501", "F403", "F405", "N803"]
|
|
||||||
|
|
||||||
[tool.ruff.lint.flake8-tidy-imports.banned-api]
|
[tool.ruff.lint.flake8-tidy-imports.banned-api]
|
||||||
"librosa".msg = "The librosa module is banned, use vllm.multimodal helpers instead"
|
"librosa".msg = "The librosa module is banned, use vllm.multimodal helpers instead"
|
||||||
|
|||||||
@@ -1,110 +0,0 @@
|
|||||||
# Higgs-Audio V3 TTS
|
|
||||||
|
|
||||||
> Multilingual text-to-speech with voice cloning on 1×H100
|
|
||||||
|
|
||||||
## Summary
|
|
||||||
|
|
||||||
- Vendor: Boson AI
|
|
||||||
- Model: `bosonai/higgs-audio-v3-tts-4b`
|
|
||||||
- Task: Text-to-speech synthesis with optional voice cloning (100+ languages)
|
|
||||||
- Mode: Online serving with the OpenAI-compatible `/v1/audio/speech` API; offline batch inference
|
|
||||||
- Maintainer: @yuekaiz
|
|
||||||
|
|
||||||
## When to use this recipe
|
|
||||||
|
|
||||||
Use this recipe to serve `bosonai/higgs-audio-v3-tts-4b` for high-quality
|
|
||||||
multilingual TTS. The model generates 24 kHz speech, supports zero-shot voice
|
|
||||||
cloning from a reference clip, and handles 100+ languages with inline control
|
|
||||||
tokens for emotion, style, and prosody. The architecture is a ~4B Qwen3
|
|
||||||
backbone with fused multi-codebook embedding/head (8 codebooks × 1026 vocab,
|
|
||||||
MusicGen-style delay pattern).
|
|
||||||
|
|
||||||
## References
|
|
||||||
|
|
||||||
- Model card: [bosonai/higgs-audio-v3-tts-4b](https://huggingface.co/bosonai/higgs-audio-v3-tts-4b)
|
|
||||||
- Offline example: [`examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py`](../../examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py)
|
|
||||||
- Online example: [`examples/online_serving/text_to_speech/higgs_audio_v3/`](../../examples/online_serving/text_to_speech/higgs_audio_v3/)
|
|
||||||
- Benchmark results: see Performance section below
|
|
||||||
|
|
||||||
## Hardware Support
|
|
||||||
|
|
||||||
## GPU
|
|
||||||
|
|
||||||
### 1×H100 80GB
|
|
||||||
|
|
||||||
#### Environment
|
|
||||||
|
|
||||||
- OS: Linux
|
|
||||||
- Python: 3.12+
|
|
||||||
- CUDA: 12.x
|
|
||||||
- vLLM version: 0.22.0
|
|
||||||
- vLLM-Omni version or commit: `36e048fd` (branch `higgs-v3`)
|
|
||||||
|
|
||||||
#### Command
|
|
||||||
|
|
||||||
**Online serving:**
|
|
||||||
|
|
||||||
```bash
|
|
||||||
vllm-omni serve bosonai/higgs-audio-v3-tts-4b \
|
|
||||||
--host 0.0.0.0 --port 8095 \
|
|
||||||
--trust-remote-code --omni
|
|
||||||
```
|
|
||||||
|
|
||||||
The default deploy config `vllm_omni/deploy/higgs_multimodal_qwen3.yaml` is
|
|
||||||
loaded automatically by model registry (HF `model_type=higgs_multimodal_qwen3`).
|
|
||||||
|
|
||||||
**Offline batch inference:**
|
|
||||||
|
|
||||||
```bash
|
|
||||||
python examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py \
|
|
||||||
--texts "Hello world." "The quick brown fox jumps over the lazy dog." \
|
|
||||||
--output-dir results/higgs_v3_wavs
|
|
||||||
```
|
|
||||||
|
|
||||||
**Offline voice clone:**
|
|
||||||
|
|
||||||
```bash
|
|
||||||
python examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py \
|
|
||||||
--texts "Text to synthesize in the cloned voice." \
|
|
||||||
--ref-audio path/to/reference.wav \
|
|
||||||
--ref-text "Transcript of the reference clip." \
|
|
||||||
--output-dir results/higgs_v3_clone
|
|
||||||
```
|
|
||||||
|
|
||||||
#### Verification
|
|
||||||
|
|
||||||
Basic TTS via curl:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
curl -X POST http://localhost:8095/v1/audio/speech \
|
|
||||||
-H "Content-Type: application/json" \
|
|
||||||
-d '{
|
|
||||||
"model": "bosonai/higgs-audio-v3-tts-4b",
|
|
||||||
"input": "Hello, how are you?"
|
|
||||||
}' --output hello.wav
|
|
||||||
```
|
|
||||||
|
|
||||||
Voice clone via Python client:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
python examples/online_serving/text_to_speech/higgs_audio_v3/batch_speech_client.py \
|
|
||||||
--base-url http://localhost:8095 \
|
|
||||||
--model bosonai/higgs-audio-v3-tts-4b \
|
|
||||||
--ref-audio path/to/reference.wav \
|
|
||||||
--ref-text "Transcript of the reference." \
|
|
||||||
--prompts "Text to clone."
|
|
||||||
```
|
|
||||||
|
|
||||||
#### Notes
|
|
||||||
|
|
||||||
- Memory usage: Stage 0 (talker, ~4B) uses ~60% GPU memory; Stage 1 (codec decoder) uses ~25%.
|
|
||||||
- Key flags: `--trust-remote-code` and `--omni` are required.
|
|
||||||
- Output: 24 kHz mono WAV.
|
|
||||||
- Voice cloning: `ref_audio` accepts WAV/FLAC/MP3; `ref_text` is optional but improves fidelity.
|
|
||||||
- Deploy config: `vllm_omni/deploy/higgs_multimodal_qwen3.yaml` (auto-discovered from `model_type`).
|
|
||||||
- `max_num_seqs=16` for both stages.
|
|
||||||
- `enforce_eager=true` for both stages (CUDA graph available for stage 0 but no throughput gain at batch>1).
|
|
||||||
- Known limitations:
|
|
||||||
- CUDA graph for stage 0 helps batch=1 (-13% RTF) but not batch>1 due to model-owned sampler running outside graph.
|
|
||||||
- Stage 1 (code2wav) must use `enforce_eager=true` (`@torch.inference_mode` incompatible with graph capture).
|
|
||||||
- Async streaming (chunk-based) not yet implemented; sync-only pipeline.
|
|
||||||
@@ -1,229 +0,0 @@
|
|||||||
# MiniCPM-o 4.5
|
|
||||||
|
|
||||||
> Online serving for omni multimodal chat (text / image / audio / video → text + 24 kHz speech)
|
|
||||||
|
|
||||||
## Summary
|
|
||||||
|
|
||||||
- Vendor: OpenBMB
|
|
||||||
- Model: [`openbmb/MiniCPM-o-4_5`](https://huggingface.co/openbmb/MiniCPM-o-4_5)
|
|
||||||
- Task: Omni multimodal chat — accepts text / image / audio / video input;
|
|
||||||
emits text and 24 kHz mono speech in the same response
|
|
||||||
- Mode: Online serving via the OpenAI-compatible `/v1/chat/completions`
|
|
||||||
API, plus a bundled Gradio demo (text + speech UI)
|
|
||||||
- Maintainer: [`@tc-mb`](https://github.com/tc-mb) (MiniCPM-V / MiniCPM-o team)
|
|
||||||
|
|
||||||
## When to use this recipe
|
|
||||||
|
|
||||||
Use this recipe as a known-good starting point for serving
|
|
||||||
`openbmb/MiniCPM-o-4_5` on vLLM-Omni. MiniCPM-o 4.5 is the omni member
|
|
||||||
of the MiniCPM-o family — it pairs a multimodal-understanding thinker
|
|
||||||
LLM with a streaming `MiniCPMTTS + Token2Wav` talker so a single
|
|
||||||
`/v1/chat/completions` call can return text and 24 kHz speech in one
|
|
||||||
shot. The recipe covers three shipped GPU layouts (2 / 3 / 8 GPUs)
|
|
||||||
selected via `--deploy-config`.
|
|
||||||
|
|
||||||
## References
|
|
||||||
|
|
||||||
- Default deploy configs (auto-loaded by HF `model_type=minicpmo` +
|
|
||||||
`hf_config.version="4.5"`):
|
|
||||||
- 2-GPU layout (default):
|
|
||||||
[`vllm_omni/deploy/minicpmo_4_5.yaml`](../../vllm_omni/deploy/minicpmo_4_5.yaml)
|
|
||||||
- 3-GPU layout (thinker TP=2):
|
|
||||||
[`vllm_omni/deploy/minicpmo_4_5_3gpu.yaml`](../../vllm_omni/deploy/minicpmo_4_5_3gpu.yaml)
|
|
||||||
- 8x RTX 4090 layout:
|
|
||||||
[`vllm_omni/deploy/minicpmo_4_5_8x4090.yaml`](../../vllm_omni/deploy/minicpmo_4_5_8x4090.yaml)
|
|
||||||
- Online example + Gradio demo:
|
|
||||||
[`examples/online_serving/minicpmo/`](../../examples/online_serving/minicpmo/)
|
|
||||||
- Pipeline / talker source:
|
|
||||||
[`vllm_omni/model_executor/models/minicpmo_4_5/`](../../vllm_omni/model_executor/models/minicpmo_4_5/)
|
|
||||||
- Stage-input processor (thinker → talker bridge):
|
|
||||||
[`vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py`](../../vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py)
|
|
||||||
- Upstream model card:
|
|
||||||
[`openbmb/MiniCPM-o-4_5`](https://huggingface.co/openbmb/MiniCPM-o-4_5)
|
|
||||||
- Integration PR:
|
|
||||||
[vllm-project/vllm-omni#3642](https://github.com/vllm-project/vllm-omni/pull/3642)
|
|
||||||
|
|
||||||
## Hardware Support
|
|
||||||
|
|
||||||
Three GPU layouts ship with default deploy configs. Pick the layout that
|
|
||||||
matches your hardware and pass it via `--deploy-config`; the talker
|
|
||||||
(`MiniCPMTTS + Token2Wav`) always lives on its own GPU because of the
|
|
||||||
in-process vocoder, and the thinker is the part that scales out via TP.
|
|
||||||
|
|
||||||
| Layout | Thinker | Talker + Token2Wav | Typical hardware |
|
|
||||||
| --- | --- | --- | --- |
|
|
||||||
| 2-GPU (default) | GPU 0 | GPU 1 | 2x A100/H100/H200 80GB |
|
|
||||||
| 3-GPU (thinker TP=2) | GPU 0,1 (TP=2) | GPU 2 | 3x mid-tier GPUs |
|
|
||||||
| 8x RTX 4090 24GB | GPU 0–3 (TP=4) | GPU 4 | 8x RTX 4090 consumer |
|
|
||||||
|
|
||||||
## GPU
|
|
||||||
|
|
||||||
### 2 x GPU (default — single command)
|
|
||||||
|
|
||||||
The default
|
|
||||||
[`vllm_omni/deploy/minicpmo_4_5.yaml`](../../vllm_omni/deploy/minicpmo_4_5.yaml)
|
|
||||||
puts the thinker on GPU 0 (`~70 %` memory, `enforce_eager: true`,
|
|
||||||
`max_num_seqs: 1`) and the talker + Token2Wav vocoder on GPU 1
|
|
||||||
(`~75 %` memory). This is the recommended starting layout — works on
|
|
||||||
any pair of 80GB-class GPUs (A100, H100, H200) and on most 40GB+
|
|
||||||
pairs as long as the thinker model weights fit.
|
|
||||||
|
|
||||||
#### Environment
|
|
||||||
|
|
||||||
- OS: Linux
|
|
||||||
- Python: 3.10+
|
|
||||||
- vLLM / vLLM-Omni: >= 0.21.0 (or current `main`)
|
|
||||||
- Optional Talker dep: `stepaudio2-minicpmo` (see Notes for why this is
|
|
||||||
required and how to install it)
|
|
||||||
|
|
||||||
#### Command
|
|
||||||
|
|
||||||
```bash
|
|
||||||
vllm serve openbmb/MiniCPM-o-4_5 --omni \
|
|
||||||
--trust-remote-code \
|
|
||||||
--host 0.0.0.0 --port 8099
|
|
||||||
```
|
|
||||||
|
|
||||||
The deploy config is auto-loaded by the model registry — no
|
|
||||||
`--deploy-config` flag needed for this default 2-GPU layout.
|
|
||||||
|
|
||||||
#### Verification
|
|
||||||
|
|
||||||
**Quick smoke test (text-only output)**:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
curl http://localhost:8099/v1/chat/completions \
|
|
||||||
-H "Content-Type: application/json" \
|
|
||||||
-d '{
|
|
||||||
"model": "openbmb/MiniCPM-o-4_5",
|
|
||||||
"messages": [{"role": "user", "content": "Briefly introduce yourself."}],
|
|
||||||
"modalities": ["text"]
|
|
||||||
}'
|
|
||||||
```
|
|
||||||
|
|
||||||
**Text + speech in one response** (the headline 4.5 feature). The TTS
|
|
||||||
path is gated by a Jinja flag on the chat template, passed through
|
|
||||||
`extra_body.chat_template_kwargs.use_tts_template=true`:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
curl http://localhost:8099/v1/chat/completions \
|
|
||||||
-H "Content-Type: application/json" \
|
|
||||||
-d '{
|
|
||||||
"model": "openbmb/MiniCPM-o-4_5",
|
|
||||||
"messages": [{"role": "user", "content": "Say hello, then introduce vLLM in one sentence."}],
|
|
||||||
"modalities": ["text", "audio"],
|
|
||||||
"extra_body": {"chat_template_kwargs": {"use_tts_template": true}}
|
|
||||||
}'
|
|
||||||
```
|
|
||||||
|
|
||||||
Response carries text in `choices[0].message.content` and base64 WAV
|
|
||||||
in `choices[0].message.audio.data` (24 kHz mono, see Notes).
|
|
||||||
|
|
||||||
**Gradio demo (text + image + audio + video UI)**:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
bash examples/online_serving/minicpmo/run_gradio_demo.sh
|
|
||||||
# or run the python entry point directly:
|
|
||||||
python examples/online_serving/minicpmo/gradio_demo.py \
|
|
||||||
--minicpmo45-api-base http://localhost:8099/v1 \
|
|
||||||
--minicpmo45-model openbmb/MiniCPM-o-4_5 \
|
|
||||||
--port 7862
|
|
||||||
```
|
|
||||||
|
|
||||||
Open `http://<host>:7862` and try a text prompt with the **"Generate
|
|
||||||
speech output (TTS)"** checkbox on / off.
|
|
||||||
|
|
||||||
#### Notes
|
|
||||||
|
|
||||||
- Memory budget: thinker weights occupy GPU 0 at `gpu_memory_utilization:
|
|
||||||
0.7`; talker + Token2Wav vocoder share GPU 1 at `0.75`.
|
|
||||||
- `--trust-remote-code` is required — the HF repo ships a custom
|
|
||||||
`MiniCPMO` config / model class.
|
|
||||||
- Pin: `enforce_eager: true` on both stages (CUDA graph capture is off
|
|
||||||
by design for the talker's Token2Wav path).
|
|
||||||
- Stage 1 (talker) is hard-capped to `max_num_seqs: 1`: the talker
|
|
||||||
only consumes `runtime_additional_information[0]`, so any value > 1
|
|
||||||
makes concurrent requests share request-0's audio. This is the same
|
|
||||||
cap baked into the deploy config.
|
|
||||||
|
|
||||||
### 3 x GPU (thinker TP=2)
|
|
||||||
|
|
||||||
Use
|
|
||||||
[`vllm_omni/deploy/minicpmo_4_5_3gpu.yaml`](../../vllm_omni/deploy/minicpmo_4_5_3gpu.yaml)
|
|
||||||
when you have a third GPU available and want the thinker on 2-way
|
|
||||||
tensor parallel for higher throughput; the talker stays on its own
|
|
||||||
GPU (talker has its own in-process Token2Wav vocoder, so co-locating
|
|
||||||
it with the thinker risks OOM under load).
|
|
||||||
|
|
||||||
#### Command
|
|
||||||
|
|
||||||
```bash
|
|
||||||
vllm serve openbmb/MiniCPM-o-4_5 --omni \
|
|
||||||
--deploy-config vllm_omni/deploy/minicpmo_4_5_3gpu.yaml \
|
|
||||||
--trust-remote-code \
|
|
||||||
--host 0.0.0.0 --port 8099
|
|
||||||
```
|
|
||||||
|
|
||||||
Verification and Notes mirror the 2-GPU section; thinker latency
|
|
||||||
roughly halves under load thanks to TP=2.
|
|
||||||
|
|
||||||
### 8 x RTX 4090 24GB (consumer-GPU layout)
|
|
||||||
|
|
||||||
Use
|
|
||||||
[`vllm_omni/deploy/minicpmo_4_5_8x4090.yaml`](../../vllm_omni/deploy/minicpmo_4_5_8x4090.yaml)
|
|
||||||
on an 8x RTX 4090 host. Thinker uses 4-way TP across GPUs 0–3
|
|
||||||
(`~85 %` mem each ≈ 20.4 GiB/card), talker + Token2Wav lives on GPU 4
|
|
||||||
(`~90 %` mem). GPUs 5–7 are left free.
|
|
||||||
|
|
||||||
#### Command
|
|
||||||
|
|
||||||
```bash
|
|
||||||
vllm serve openbmb/MiniCPM-o-4_5 --omni \
|
|
||||||
--deploy-config vllm_omni/deploy/minicpmo_4_5_8x4090.yaml \
|
|
||||||
--trust-remote-code \
|
|
||||||
--host 0.0.0.0 --port 8099
|
|
||||||
```
|
|
||||||
|
|
||||||
#### Notes
|
|
||||||
|
|
||||||
- `max_model_len` is capped at 4096 in this layout — 8192 still OOMs on
|
|
||||||
4090s. Raise it if your cards have more headroom (e.g. 4090 D /
|
|
||||||
custom 32 GB SKUs), but verify with a long-prompt run before
|
|
||||||
promoting.
|
|
||||||
- All other knobs match the 2-GPU section; the only difference is the
|
|
||||||
per-card memory pressure on the thinker shards.
|
|
||||||
|
|
||||||
## Notes (applies to all layouts)
|
|
||||||
|
|
||||||
- **Talker dependency**: the `MiniCPM-o 4.5` talker calls
|
|
||||||
`from stepaudio2 import Token2wav` against the MiniCPM-o-flavored
|
|
||||||
vocoder (PyPI package `stepaudio2-minicpmo` — NOT the upstream
|
|
||||||
`stepfun-ai/Step-Audio2`, whose `Token2wav.__init__` signature
|
|
||||||
rejects `n_timesteps`). Install via the published extra:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
pip install 'vllm-omni[minicpmo]'
|
|
||||||
```
|
|
||||||
|
|
||||||
Equivalent direct install: `pip install stepaudio2-minicpmo`. A
|
|
||||||
missing dep raises `ImportError` at first request with the same
|
|
||||||
install hint instead of silently emitting empty audio.
|
|
||||||
|
|
||||||
- **TTS trigger**: speech output is only emitted when the client passes
|
|
||||||
`extra_body.chat_template_kwargs.use_tts_template=true`. Without it,
|
|
||||||
the response is text-only (which is also faster).
|
|
||||||
|
|
||||||
- **Output audio**: 24 kHz mono WAV inside the OpenAI-style
|
|
||||||
`message.audio.data` (base64). The Gradio demo's WAV player decodes
|
|
||||||
this automatically.
|
|
||||||
|
|
||||||
- **Routing**: MiniCPM-o 4.5 and 2.6 both ship `architectures=
|
|
||||||
["MiniCPMO"]` in HF config; routing is disambiguated by
|
|
||||||
`hf_config.version == "4.5"` via the
|
|
||||||
`hf_config_predicate` on the 4.5 pipeline. A 2.6 checkpoint loaded
|
|
||||||
with this recipe's `--deploy-config` will be rejected at startup
|
|
||||||
rather than silently misrouted.
|
|
||||||
|
|
||||||
- **Async chunking**: disabled in all three deploy configs
|
|
||||||
(`async_chunk: false`) — the talker batches a single full thinker
|
|
||||||
output, not chunks.
|
|
||||||
@@ -1,249 +0,0 @@
|
|||||||
# VoxCPM2
|
|
||||||
|
|
||||||
> Online + offline TTS with native AR pipeline (48 kHz, 30+ languages)
|
|
||||||
|
|
||||||
## Summary
|
|
||||||
|
|
||||||
- Vendor: OpenBMB
|
|
||||||
- Model: `openbmb/VoxCPM2`
|
|
||||||
- Task: Text-to-speech (zero-shot synthesis and reference-audio voice cloning)
|
|
||||||
- Mode: Online serving via the OpenAI-compatible `/v1/audio/speech` API, plus
|
|
||||||
an offline end-to-end example
|
|
||||||
- Maintainer: Community
|
|
||||||
|
|
||||||
## When to use this recipe
|
|
||||||
|
|
||||||
Use this recipe as a known-good starting point for serving `openbmb/VoxCPM2`
|
|
||||||
on vLLM-Omni on consumer-class GPUs. VoxCPM2 is a 2B-parameter native AR
|
|
||||||
TTS model that emits 48 kHz mono audio and supports 30+ languages out of
|
|
||||||
the box. It runs as a single-stage pipeline (`MiniCPM4 base LM → FSQ →
|
|
||||||
MiniCPM4 residual LM → LocDiT CFM solver → AudioVAE`) and supports both
|
|
||||||
zero-shot synthesis and reference-audio voice cloning, with optional
|
|
||||||
gapless streaming via the bundled Gradio AudioWorklet demo.
|
|
||||||
|
|
||||||
Because the model is only 2B params, a single 24 GB consumer GPU is
|
|
||||||
sufficient — no tensor parallelism is required.
|
|
||||||
|
|
||||||
## References
|
|
||||||
|
|
||||||
- Example guide (online):
|
|
||||||
[`examples/online_serving/text_to_speech/README.md`](../../examples/online_serving/text_to_speech/README.md#voxcpm2)
|
|
||||||
- Online client / Gradio demo:
|
|
||||||
[`examples/online_serving/text_to_speech/voxcpm2/`](../../examples/online_serving/text_to_speech/voxcpm2/)
|
|
||||||
- Offline end-to-end script:
|
|
||||||
[`examples/offline_inference/text_to_speech/voxcpm2/end2end.py`](../../examples/offline_inference/text_to_speech/voxcpm2/end2end.py)
|
|
||||||
- Default deploy config:
|
|
||||||
[`vllm_omni/deploy/voxcpm2.yaml`](../../vllm_omni/deploy/voxcpm2.yaml)
|
|
||||||
- Talker / pipeline source:
|
|
||||||
[`vllm_omni/model_executor/models/voxcpm2/`](../../vllm_omni/model_executor/models/voxcpm2/)
|
|
||||||
- Speech API reference:
|
|
||||||
[`docs/serving/speech_api.md`](../../docs/serving/speech_api.md#voxcpm2)
|
|
||||||
- Related issue or discussion:
|
|
||||||
[RFC: add recipes folder](https://github.com/vllm-project/vllm-omni/issues/2645)
|
|
||||||
|
|
||||||
## Hardware Support
|
|
||||||
|
|
||||||
This recipe documents one tested 24 GB consumer-GPU configuration.
|
|
||||||
Larger-VRAM (H20 / H100 / A100) and other vendor sections (ROCm, NPU) are
|
|
||||||
welcome as community validation lands.
|
|
||||||
|
|
||||||
## GPU
|
|
||||||
|
|
||||||
### 1 x RTX 4090 24GB (Single GPU, Minimum Recommended)
|
|
||||||
|
|
||||||
VoxCPM2 (2B params in bfloat16) fits comfortably on a single 24 GB GPU.
|
|
||||||
The bundled default config at
|
|
||||||
[`vllm_omni/deploy/voxcpm2.yaml`](../../vllm_omni/deploy/voxcpm2.yaml)
|
|
||||||
(`gpu_memory_utilization: 0.9`, `max_num_seqs: 4`, `enforce_eager: true`)
|
|
||||||
loads cleanly on a 4090 with ~4.9 GiB for model weights, ~15.2 GiB for
|
|
||||||
KV cache, and ~2 GiB for the talker's CUDA-Graph / CFM / VAE buffers,
|
|
||||||
for a total resident footprint of about **22 GiB / 24 GiB** in steady
|
|
||||||
state.
|
|
||||||
|
|
||||||
#### Environment
|
|
||||||
|
|
||||||
- OS: Linux (Ubuntu 22.04)
|
|
||||||
- Python: 3.12
|
|
||||||
- Driver / runtime: NVIDIA driver 575.57.08, CUDA 12.9 runtime
|
|
||||||
- torch: 2.11.0+cu128
|
|
||||||
- vLLM: 0.21.0
|
|
||||||
- vLLM-Omni: 0.20.1.dev (current `main`)
|
|
||||||
- Native VoxCPM2 source: `voxcpm` >= 2.0
|
|
||||||
|
|
||||||
#### Command
|
|
||||||
|
|
||||||
Start the server from the repository root. On a clean 4090 with no other
|
|
||||||
tenant on the GPU, the defaults from `voxcpm2.yaml` work as-is:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
vllm serve openbmb/VoxCPM2 --omni \
|
|
||||||
--host 0.0.0.0 --port 8000
|
|
||||||
```
|
|
||||||
|
|
||||||
The deploy config at
|
|
||||||
[`vllm_omni/deploy/voxcpm2.yaml`](../../vllm_omni/deploy/voxcpm2.yaml)
|
|
||||||
is loaded automatically by the model registry (HF `model_type=voxcpm2`).
|
|
||||||
Pass `--deploy-config <path>` to override.
|
|
||||||
|
|
||||||
#### Verification
|
|
||||||
|
|
||||||
**Server cold-start**: ~60 s from `vllm serve` to `Application startup
|
|
||||||
complete` (subprocess fork + vLLM 0.21 init + model load + flashinfer JIT
|
|
||||||
+ torch.compile of LocDiT / feat_encoder / AudioVAE + CUDA-Graph warmup).
|
|
||||||
The first request after startup pays a small additional cost; steady-state
|
|
||||||
requests are much faster.
|
|
||||||
|
|
||||||
**T1 — zero-shot synthesis (curl, WAV)**:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
curl -X POST http://127.0.0.1:8000/v1/audio/speech \
|
|
||||||
-H "Content-Type: application/json" \
|
|
||||||
-d '{
|
|
||||||
"model": "voxcpm2",
|
|
||||||
"input": "Hello, this is VoxCPM2 served by vLLM-Omni on a single RTX 4090.",
|
|
||||||
"voice": "default",
|
|
||||||
"response_format": "wav"
|
|
||||||
}' --output output.wav
|
|
||||||
```
|
|
||||||
|
|
||||||
Observed: `output.wav` 491 KB, 5.12 s @ 48 kHz mono, peak amplitude 0.996.
|
|
||||||
Server-to-client wall time: 2.6 s (server-side RTF ~0.5 after warmup).
|
|
||||||
|
|
||||||
**T2 — bundled Python client (zero-shot)**:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
python examples/online_serving/text_to_speech/voxcpm2/openai_speech_client.py \
|
|
||||||
--text "This audio was generated by the bundled openai_speech_client.py script against a vLLM-Omni server." \
|
|
||||||
--api-base http://127.0.0.1:8000 \
|
|
||||||
--output output.wav
|
|
||||||
```
|
|
||||||
|
|
||||||
Observed: `output.wav` 968 KB, 10.08 s @ 48 kHz mono, peak amplitude 0.949.
|
|
||||||
|
|
||||||
**T3 — voice cloning from a reference WAV**:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
python examples/online_serving/text_to_speech/voxcpm2/openai_speech_client.py \
|
|
||||||
--text "This should sound like the reference speaker, but saying a different sentence on a single RTX 4090." \
|
|
||||||
--ref-audio /path/to/reference.wav \
|
|
||||||
--api-base http://127.0.0.1:8000 \
|
|
||||||
--output cloned.wav
|
|
||||||
```
|
|
||||||
|
|
||||||
Observed (using the T2 output `output.wav` as `--ref-audio`):
|
|
||||||
`cloned.wav` 568 KB, 5.92 s @ 48 kHz mono, peak amplitude 0.990;
|
|
||||||
end-to-end wall time 3.16 s.
|
|
||||||
|
|
||||||
`--ref-audio` accepts local file paths (auto-base64), HTTP/HTTPS URLs, or
|
|
||||||
`data:audio/wav;base64,...` data URIs. The `voice` field is required by
|
|
||||||
the OpenAI schema but ignored by VoxCPM2 unless it refers to an
|
|
||||||
uploaded voice — voice cloning is driven entirely by `ref_audio`.
|
|
||||||
|
|
||||||
**T4 — streaming PCM (raw, 48 kHz s16le)**:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
curl -X POST http://127.0.0.1:8000/v1/audio/speech \
|
|
||||||
-H "Content-Type: application/json" \
|
|
||||||
-d '{
|
|
||||||
"model": "voxcpm2",
|
|
||||||
"input": "Streaming PCM test on a single RTX 4090 with VoxCPM2.",
|
|
||||||
"voice": "default",
|
|
||||||
"stream": true,
|
|
||||||
"response_format": "pcm"
|
|
||||||
}' --no-buffer | play -t raw -r 48000 -e signed -b 16 -c 1 -
|
|
||||||
```
|
|
||||||
|
|
||||||
Observed: 230400 int16 samples = 4.80 s of PCM at 48 kHz, peak 32525 /
|
|
||||||
32767 (~99 % full scale). The player sample rate is **48 kHz**, not
|
|
||||||
24 kHz.
|
|
||||||
|
|
||||||
**T5 — gapless streaming (Gradio AudioWorklet demo)**:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
python examples/online_serving/text_to_speech/voxcpm2/gradio_demo.py \
|
|
||||||
--api-base http://127.0.0.1:8000
|
|
||||||
```
|
|
||||||
|
|
||||||
Observed: the demo binds to `http://0.0.0.0:7860` in ~4 s after launch.
|
|
||||||
The UI ("VoxCPM2 Streaming Demo") exposes three modes — **Voice Design**
|
|
||||||
(control instruction only), **Controllable Cloning** (ref audio +
|
|
||||||
optional style control), and **Ultimate Cloning** (ref audio +
|
|
||||||
transcript for audio continuation) — plus a `Stream (gapless)`
|
|
||||||
checkbox (default on) that drives the AudioWorklet streaming path. The
|
|
||||||
right pane includes runnable example prompts and a `Ready` status
|
|
||||||
indicator. End-to-end browser playback through `Generate Speech` was
|
|
||||||
verified working on the same vLLM-Omni server backing T1–T4 above.
|
|
||||||
|
|
||||||
**T6 — offline inference (no server, single process per request)**:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
python examples/offline_inference/text_to_speech/voxcpm2/end2end.py \
|
|
||||||
--text "Hello, this is VoxCPM2 native AR synthesis running on vLLM-Omni and a single RTX 4090." \
|
|
||||||
--output-dir output_audio
|
|
||||||
```
|
|
||||||
|
|
||||||
`end2end.py` is a one-shot script (init → one `generate()` → exit), so
|
|
||||||
its reported RTF includes per-process compile + CUDA-Graph capture +
|
|
||||||
first-request runtime warmup. Note that the default `voxcpm2.yaml` sets
|
|
||||||
`enforce_eager: true`, so vLLM's engine-level CUDA graphs are off by
|
|
||||||
design; the "CUDA-Graph capture" referenced throughout this recipe is the
|
|
||||||
talker's own `_CapturedGraph` over the CFM solver / AudioVAE, not the LM
|
|
||||||
decode loop. Flipping `enforce_eager=false` does **not** remove this cost.
|
|
||||||
To separate the warmup cost from the
|
|
||||||
talker's true per-step throughput, the table below reports **5
|
|
||||||
consecutive `engine.generate([prompt])` calls in the same process**
|
|
||||||
after a one-off ~28 s engine init:
|
|
||||||
|
|
||||||
| call | duration | inference | RTF | notes |
|
|
||||||
| ---- | -------- | --------- | ----- | ------------------------------------------------ |
|
|
||||||
| #1 | 6.72 s | 11.97 s | 1.782 | cold: torch.compile + CUDA-Graph capture |
|
|
||||||
| #2 | 6.24 s | 11.43 s | 1.831 | still runtime warmup |
|
|
||||||
| #3 | 6.88 s | 0.82 s | 0.120 | ⚡ steady-state |
|
|
||||||
| #4 | 6.56 s | 0.78 s | 0.119 | steady-state |
|
|
||||||
| #5 | 5.76 s | 0.70 s | 0.121 | steady-state |
|
|
||||||
|
|
||||||
**Steady-state RTF (mean of calls #3–#5): 0.120**, i.e. ~8× real-time on
|
|
||||||
a single 4090. The high call-#1 / call-#2 numbers are one-time warmup
|
|
||||||
costs that **amortise to zero after the second request** — they are not
|
|
||||||
the model's per-step throughput.
|
|
||||||
|
|
||||||
A literal one-shot `python end2end.py` invocation in a fresh process
|
|
||||||
only sees call #1 (~RTF 1.8). For repeated workloads, either reuse the
|
|
||||||
same engine across calls (as in the table above) or use the online
|
|
||||||
serving path (T1–T5); a long-lived server amortises the warmup across
|
|
||||||
all subsequent requests (T1 wall-time RTF was 0.51 with a warm server).
|
|
||||||
|
|
||||||
Voice cloning + continuation mode (`ref_audio` + transcript):
|
|
||||||
|
|
||||||
```bash
|
|
||||||
python examples/offline_inference/text_to_speech/voxcpm2/end2end.py \
|
|
||||||
--text "Hello, this is the offline voice cloning test on a single RTX 4090." \
|
|
||||||
--ref-audio /path/to/reference.wav \
|
|
||||||
--ref-text "Original transcript of the reference audio."
|
|
||||||
```
|
|
||||||
|
|
||||||
Same 5-call methodology as the zero-shot table, this time with a
|
|
||||||
10.08 s reference WAV supplied via `--ref-audio` (and matching
|
|
||||||
`--ref-text`):
|
|
||||||
|
|
||||||
| call | duration | inference | RTF | notes |
|
|
||||||
| ---- | -------- | --------- | ----- | ------------------------------------------------ |
|
|
||||||
| #1 | 5.44 s | 12.38 s | 2.276 | cold: compile + CUDA-Graph capture + ref encode |
|
|
||||||
| #2 | 5.12 s | 2.47 s | 0.482 | most warmup done |
|
|
||||||
| #3 | 5.44 s | 0.76 s | 0.139 | ⚡ steady-state |
|
|
||||||
| #4 | 4.96 s | 0.68 s | 0.137 | steady-state |
|
|
||||||
| #5 | 5.28 s | 0.71 s | 0.134 | steady-state |
|
|
||||||
|
|
||||||
**Steady-state cloning RTF (mean of calls #3–#5): 0.137**, i.e. only
|
|
||||||
**+0.017 RTF** above the zero-shot steady-state of 0.120. That delta is
|
|
||||||
the per-request `AudioVAE.encode` over the reference clip (~85 ms for a
|
|
||||||
10 s ref WAV) plus the ref-continuation prompt build; the AR decode
|
|
||||||
steps themselves run at the same per-step rate as zero-shot.
|
|
||||||
|
|
||||||
#### Notes
|
|
||||||
|
|
||||||
- Prerequisites: `pip install voxcpm soundfile httpx ninja`. `ninja` must be on `PATH` because `flashinfer` JIT-compiles a top-k sampling kernel during the first profile pass (without it, startup dies with `FileNotFoundError: 'ninja'`).
|
|
||||||
- Memory usage: ~4.9 GiB model weights + ~15.2 GiB KV cache ≈ 22 GiB / 24 GiB resident with the default `voxcpm2.yaml`. On a shared GPU, pass `--gpu-memory-utilization 0.75` (or lower) if startup fails the free-memory check.
|
|
||||||
- Key flags: `--omni` is required. `--trust-remote-code` is not needed — the HF config registers under `model_type=voxcpm2`.
|
|
||||||
- Output: 48 kHz mono. When streaming PCM to a player, use `-r 48000`.
|
|
||||||
- Cold start: ~60 s server boot + ~25 s first-inference overhead (torch.compile + flashinfer JIT + the talker's own `_CapturedGraph` capture over CFM / AudioVAE — not vLLM engine CUDA graphs, which `enforce_eager: true` disables). Steady-state inference RTF is **~0.12** (~8× real-time on a single 4090; see T6 table for the per-call breakdown); online server requests run at ~RTF 0.5 wall-time (includes HTTP round-trip). Keep the server warm for interactive use.
|
|
||||||
@@ -1,145 +0,0 @@
|
|||||||
# MOSS-TTS
|
|
||||||
|
|
||||||
## Summary
|
|
||||||
|
|
||||||
- Vendor: OpenMOSS
|
|
||||||
- Models: `OpenMOSS-Team/MOSS-TTS` (8B), `OpenMOSS-Team/MOSS-TTS-v1.5` (8B),
|
|
||||||
`OpenMOSS-Team/MOSS-TTS-Realtime` (1.7B), `OpenMOSS-Team/MOSS-TTSD-v1.0` (8B),
|
|
||||||
`OpenMOSS-Team/MOSS-SoundEffect` (8B), `OpenMOSS-Team/MOSS-VoiceGenerator` (1.7B)
|
|
||||||
- Task: Text-to-speech synthesis, sound effect generation, zero-shot voice design
|
|
||||||
- Mode: Online serving via the OpenAI-compatible `/v1/audio/speech` API; offline inference
|
|
||||||
- Maintainer: Community
|
|
||||||
|
|
||||||
## When to use this recipe
|
|
||||||
|
|
||||||
Use this recipe for 24 kHz multilingual TTS with voice cloning (20 languages including
|
|
||||||
Chinese and English). Choose a variant based on your latency and quality requirements:
|
|
||||||
|
|
||||||
| Model | Params | Use case |
|
|
||||||
|---|---|---|
|
|
||||||
| MOSS-TTS | 8B | General TTS, highest quality |
|
|
||||||
| MOSS-TTS-v1.5 | 8B | General TTS upgrade of 1.0: 31 languages, steadier cloning, `[pause Xs]` markers (set `language` for best results); same `MossTTSDelay` API |
|
|
||||||
| MOSS-TTS-Realtime | 1.7B | Lowest latency (TTFB ~180 ms), streaming-first |
|
|
||||||
| MOSS-TTSD-v1.0 | 8B | Multi-turn dialogue TTS |
|
|
||||||
| MOSS-SoundEffect | 8B | Sound effect synthesis from text description |
|
|
||||||
| MOSS-VoiceGenerator | 1.7B | Zero-shot voice design |
|
|
||||||
|
|
||||||
All variants share the same codec (`OpenMOSS-Team/MOSS-Audio-Tokenizer`, ~7 GB) and
|
|
||||||
output 24 kHz mono audio.
|
|
||||||
|
|
||||||
## References
|
|
||||||
|
|
||||||
- Offline inference example: [`examples/offline_inference/text_to_speech/moss_tts/`](../../examples/offline_inference/text_to_speech/moss_tts/)
|
|
||||||
- Deploy configs: [`vllm_omni/deploy/moss_tts.yaml`](../../vllm_omni/deploy/moss_tts.yaml) and variants
|
|
||||||
- HuggingFace org: <https://huggingface.co/OpenMOSS-Team>
|
|
||||||
|
|
||||||
## Hardware Support
|
|
||||||
|
|
||||||
### GPU
|
|
||||||
|
|
||||||
#### 1x H100 80GB — MOSS-TTS (8B)
|
|
||||||
|
|
||||||
##### Environment
|
|
||||||
|
|
||||||
- OS: Linux
|
|
||||||
- Python: 3.11+
|
|
||||||
- CUDA 12.8
|
|
||||||
- vLLM-Omni version: see `vllm_omni/__version__.py`
|
|
||||||
|
|
||||||
##### Command
|
|
||||||
|
|
||||||
```bash
|
|
||||||
# The codec is loaded automatically from OpenMOSS-Team/MOSS-Audio-Tokenizer.
|
|
||||||
# Override the path with MOSS_TTS_CODEC_PATH if you have a local copy.
|
|
||||||
vllm serve OpenMOSS-Team/MOSS-TTS --omni --port 8091
|
|
||||||
```
|
|
||||||
|
|
||||||
##### Verification
|
|
||||||
|
|
||||||
Voice cloning (provide a reference audio clip):
|
|
||||||
|
|
||||||
```bash
|
|
||||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
|
||||||
-H "Content-Type: application/json" \
|
|
||||||
-d '{
|
|
||||||
"model": "OpenMOSS-Team/MOSS-TTS",
|
|
||||||
"input": "Hello, this is a voice cloning test.",
|
|
||||||
"voice": "default",
|
|
||||||
"ref_audio": "https://raw.githubusercontent.com/OpenMOSS/MOSS-TTS/main/assets/audio/zh_1.wav",
|
|
||||||
"response_format": "wav"
|
|
||||||
}' --output output.wav
|
|
||||||
```
|
|
||||||
|
|
||||||
##### Notes
|
|
||||||
|
|
||||||
- Peak GPU memory: ~18 GB for the talker (8B) + ~8 GB for the codec decoder on the same device.
|
|
||||||
Use `gpu_memory_utilization: 0.85` in `moss_tts.yaml` (default).
|
|
||||||
- Output: 24 kHz mono WAV.
|
|
||||||
- The `MOSS_TTS_CODEC_PATH` environment variable overrides the codec checkpoint location.
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
#### 1x A10G 24GB — MOSS-TTS-Realtime (1.7B)
|
|
||||||
|
|
||||||
##### Environment
|
|
||||||
|
|
||||||
- OS: Linux
|
|
||||||
- Python: 3.11+
|
|
||||||
- CUDA 12.8
|
|
||||||
|
|
||||||
##### Command
|
|
||||||
|
|
||||||
```bash
|
|
||||||
vllm serve OpenMOSS-Team/MOSS-TTS-Realtime --omni --port 8091
|
|
||||||
```
|
|
||||||
|
|
||||||
##### Verification
|
|
||||||
|
|
||||||
```bash
|
|
||||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
|
||||||
-H "Content-Type: application/json" \
|
|
||||||
-d '{
|
|
||||||
"model": "OpenMOSS-Team/MOSS-TTS-Realtime",
|
|
||||||
"input": "This is a low-latency streaming TTS test.",
|
|
||||||
"voice": "default",
|
|
||||||
"ref_audio": "https://raw.githubusercontent.com/OpenMOSS/MOSS-TTS/main/assets/audio/zh_1.wav",
|
|
||||||
"response_format": "wav",
|
|
||||||
"stream": true
|
|
||||||
}' --output output.wav
|
|
||||||
```
|
|
||||||
|
|
||||||
##### Notes
|
|
||||||
|
|
||||||
- Peak GPU memory: ~6 GB for the talker (1.7B) + ~8 GB for the codec decoder.
|
|
||||||
- First-audio latency (TTFB): ~180 ms on A10G.
|
|
||||||
- `codec_chunk_frames: 15` in `moss_tts_realtime.yaml` for lower TTFA than the 8B variant.
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
#### 1x A10G 24GB — MOSS-SoundEffect (8B, sound effect synthesis)
|
|
||||||
|
|
||||||
##### Command
|
|
||||||
|
|
||||||
```bash
|
|
||||||
vllm serve OpenMOSS-Team/MOSS-SoundEffect --omni --port 8091
|
|
||||||
```
|
|
||||||
|
|
||||||
##### Verification
|
|
||||||
|
|
||||||
Sound effect synthesis takes a text description instead of reference audio:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
|
||||||
-H "Content-Type: application/json" \
|
|
||||||
-d '{
|
|
||||||
"model": "OpenMOSS-Team/MOSS-SoundEffect",
|
|
||||||
"input": "Thunder rumbling, rain pattering on a tin roof.",
|
|
||||||
"response_format": "wav"
|
|
||||||
}' --output thunder.wav
|
|
||||||
```
|
|
||||||
|
|
||||||
##### Notes
|
|
||||||
|
|
||||||
- No `ref_audio` required or accepted for MOSS-SoundEffect.
|
|
||||||
- Input field maps to the `ambient_sound` parameter in the upstream processor.
|
|
||||||
- Rate: ~12.5 tokens per second; longer descriptions produce longer audio.
|
|
||||||
@@ -36,10 +36,6 @@ recipes/
|
|||||||
| [`LTX/LTX-2.md`](./LTX/LTX-2.md) | Text-to-video and image-to-video serving | 1x H200 141GB |
|
| [`LTX/LTX-2.md`](./LTX/LTX-2.md) | Text-to-video and image-to-video serving | 1x H200 141GB |
|
||||||
| [`LTX/LTX-2.3.md`](./LTX/LTX-2.3.md) | Text-to-video with audio generation (22B) | 1x GPU (96GB VRAM) |
|
| [`LTX/LTX-2.3.md`](./LTX/LTX-2.3.md) | Text-to-video with audio generation (22B) | 1x GPU (96GB VRAM) |
|
||||||
| [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB |
|
| [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB |
|
||||||
| [`cosmos3/Cosmos3-Nano.md`](./cosmos3/Cosmos3-Nano.md) | Text-to-image, text-to-video, image-to-video generation, text to video with sound, action policy | 1x H200 141GB / B300 |
|
|
||||||
| [`cosmos3/Cosmos3-Super.md`](./cosmos3/Cosmos3-Super.md) | 64B T2I / T2V / I2V generation (+ optional audio) / Action policy | 8x H200/H100/A100 / 2x H200 / B300 |
|
|
||||||
| [`OpenBMB/MiniCPM-o-4_5.md`](./OpenBMB/MiniCPM-o-4_5.md) | Online serving for omni multimodal chat (text / image / audio / video → text + 24 kHz speech) | 2x A100/H100 80GB / 3x mid-tier GPU / 8x RTX 4090 24GB |
|
|
||||||
| [`OpenBMB/VoxCPM2.md`](./OpenBMB/VoxCPM2.md) | Online + offline TTS with native AR pipeline (48 kHz, 30+ languages) | 1x RTX 4090 24GB |
|
|
||||||
| [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt mixed FP8/NVFP4 | 1x A100 80GB / 2x B200 |
|
| [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt mixed FP8/NVFP4 | 1x A100 80GB / 2x B200 |
|
||||||
| [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay | 1x A100 80GB |
|
| [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay | 1x A100 80GB |
|
||||||
| [`Qwen/Qwen-Image-Edit.md`](./Qwen/Qwen-Image-Edit.md) | Text-guided single-image editing | 1x or 2x H200 141GB |
|
| [`Qwen/Qwen-Image-Edit.md`](./Qwen/Qwen-Image-Edit.md) | Text-guided single-image editing | 1x or 2x H200 141GB |
|
||||||
@@ -52,7 +48,6 @@ recipes/
|
|||||||
| [`Wan-AI/Wan2.2-S2V.md`](./Wan-AI/Wan2.2-S2V.md) | Speech-to-video serving (Wan2.2 14B) | 2x A100/H100 80GB |
|
| [`Wan-AI/Wan2.2-S2V.md`](./Wan-AI/Wan2.2-S2V.md) | Speech-to-video serving (Wan2.2 14B) | 2x A100/H100 80GB |
|
||||||
| [`zai-org/GLM-TTS.md`](./zai-org/GLM-TTS.md) | Online serving for Chinese/English zero-shot voice-cloned TTS | 1x A40 48GB |
|
| [`zai-org/GLM-TTS.md`](./zai-org/GLM-TTS.md) | Online serving for Chinese/English zero-shot voice-cloned TTS | 1x A40 48GB |
|
||||||
| [`GLM/GLM-Image.md`](./GLM/GLM-Image.md) | Online serving for image generation | 1x A800 80GB / 2x A800 80GB |
|
| [`GLM/GLM-Image.md`](./GLM/GLM-Image.md) | Online serving for image generation | 1x A800 80GB / 2x A800 80GB |
|
||||||
|
|
|
||||||
|
|
||||||
Within a single recipe file, include different hardware support sections such
|
Within a single recipe file, include different hardware support sections such
|
||||||
as `GPU`, `ROCm`, and `NPU`, and add concrete tested configurations like
|
as `GPU`, `ROCm`, and `NPU`, and add concrete tested configurations like
|
||||||
|
|||||||
@@ -1,294 +0,0 @@
|
|||||||
# Cosmos3-Nano
|
|
||||||
|
|
||||||
> Text-to-image, text-to-video, and image-to-video serving
|
|
||||||
|
|
||||||
## Summary
|
|
||||||
|
|
||||||
- Vendor: NVIDIA
|
|
||||||
- Model: `nvidia/Cosmos3-Nano`
|
|
||||||
- Task: Text-to-image (T2I), text-to-video (T2V), and image-to-video (I2V) generation, with optional synchronized audio (video + sound), action policy
|
|
||||||
- Mode: Online serving with the OpenAI-compatible image/video APIs, plus offline generation via the `Omni` API
|
|
||||||
- Maintainer: Community
|
|
||||||
|
|
||||||
## When to use this recipe
|
|
||||||
|
|
||||||
Use this recipe to deploy `nvidia/Cosmos3-Nano` for image and video generation.
|
|
||||||
A single pipeline class (`Cosmos3OmniDiffusersPipeline`) serves all three modes;
|
|
||||||
the mode is selected per request:
|
|
||||||
|
|
||||||
- **T2I** — `POST /v1/images/generations` (or a prompt carrying `modalities=["image"]`).
|
|
||||||
- **T2V** — `POST /v1/videos/sync` with `num_frames > 1` and no reference image.
|
|
||||||
- **I2V** — `POST /v1/videos/sync` with a reference image (`input_reference` file
|
|
||||||
upload, or `image_reference` JSON).
|
|
||||||
- **T2VS / I2VS** — add `generate_sound=true` (and optional `sound_duration`) to a
|
|
||||||
T2V/I2V `/v1/videos/sync` request to also generate synchronized audio, muxed into
|
|
||||||
the mp4 as AAC 48 kHz stereo. See the official model card's "Video + Audio" examples.
|
|
||||||
- **Action** — pass `extra_params={"action_mode": ...}` to drive Physical-AI tasks:
|
|
||||||
- `forward_dynamics` — given a first frame **and** an action trajectory, roll out
|
|
||||||
the resulting video. Synchronous: `POST /v1/videos/sync`.
|
|
||||||
- `policy` — given a first frame and a language instruction, **predict** the action
|
|
||||||
trajectory (and a rollout video). Use the async `POST /v1/videos` endpoint and
|
|
||||||
read the predicted action from the top-level `action` field
|
|
||||||
(`{data, shape, dtype, raw_action_dim, domain_id}`).
|
|
||||||
|
|
||||||
Action requests also take `domain_name` (e.g. `av`, `bridge_orig_lerobot`,
|
|
||||||
`droid_lerobot`, `agibotworld`, …; or a numeric `domain_id`), `raw_action_dim`,
|
|
||||||
and `action_chunk_size` (must equal `num_frames` or `num_frames - 1`). For
|
|
||||||
`forward_dynamics` also pass the `action` array. The dedicated policy checkpoint
|
|
||||||
**`nvidia/Cosmos3-Nano-Policy-DROID`** is served the same way
|
|
||||||
(`domain_name=droid_lerobot`).
|
|
||||||
|
|
||||||
`inverse_dynamics` (recover the action from a given video) is supported by the
|
|
||||||
pipeline; **online inference of inverse dynamics will be added in a follow-up MR.**
|
|
||||||
|
|
||||||
## References
|
|
||||||
|
|
||||||
- Model card (authoritative usage + example assets): <https://huggingface.co/nvidia/Cosmos3-Nano>
|
|
||||||
- Example inputs/outputs live in the repo's `assets/` (`example_t2v_prompt.json`,
|
|
||||||
`example_i2v_prompt.json`, `example_i2v_input.jpg`, `negative_prompt.json`;
|
|
||||||
audio examples: `example_t2vs_prompt.json`, `example_t2vs_output.mp4`,
|
|
||||||
`example_i2vs_output.mp4`).
|
|
||||||
- Prompt upsampling (recommended for quality): the model expects JSON-upsampled
|
|
||||||
structured prompts; see NVIDIA's `cosmos-framework` prompt-upsampling docs.
|
|
||||||
- Pipeline: [`vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py`](../../vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py)
|
|
||||||
- Smoke tests (canonical request formats): [`tests/e2e/accuracy/test_cosmos3_similarity.py`](../../tests/e2e/accuracy/test_cosmos3_similarity.py)
|
|
||||||
|
|
||||||
## Hardware Support
|
|
||||||
|
|
||||||
## GPU
|
|
||||||
|
|
||||||
### 1x H200 141GB / B300 (Online serving)
|
|
||||||
|
|
||||||
#### Environment
|
|
||||||
|
|
||||||
- OS: Ubuntu 22.04+
|
|
||||||
- Python: 3.12+
|
|
||||||
- Driver / runtime: NVIDIA CUDA environment
|
|
||||||
- vLLM version: match the repository requirements from your current checkout
|
|
||||||
- vLLM-Omni version or commit: use the commit you are deploying from
|
|
||||||
|
|
||||||
#### Command
|
|
||||||
|
|
||||||
Requires the `vllm-omni` package (or the `vllm/vllm-omni:cosmos3` container),
|
|
||||||
which provides the `vllm serve … --omni` entrypoint used below.
|
|
||||||
|
|
||||||
Safety guardrails are **on by default** (NVIDIA Open Model License). They load
|
|
||||||
the **gated** `nvidia/Cosmos-1.0-Guardrail` model, so to keep them on you must:
|
|
||||||
|
|
||||||
1. `pip install cosmos-guardrail`
|
|
||||||
2. Accept the license at <https://huggingface.co/nvidia/Cosmos-1.0-Guardrail>
|
|
||||||
3. Export a token with access: `export HF_TOKEN=hf_...`
|
|
||||||
|
|
||||||
Then launch the recommended server:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
vllm serve nvidia/Cosmos3-Nano \
|
|
||||||
--omni \
|
|
||||||
--host 0.0.0.0 --port 8000 \
|
|
||||||
--init-timeout 1800
|
|
||||||
```
|
|
||||||
|
|
||||||
To run **without** guardrails (you are responsible for license compliance),
|
|
||||||
add `--no-guardrails` (no token/`cosmos-guardrail` needed). For extra GPUs use
|
|
||||||
`--ulysses-degree N` (context parallel) or `--tensor-parallel-size N`;
|
|
||||||
`--enable-layerwise-offload` reduces VRAM on smaller GPUs. The pipeline
|
|
||||||
auto-resolves from `model_index.json`; pass
|
|
||||||
`--model-class-name Cosmos3OmniDiffusersPipeline` to force it explicitly.
|
|
||||||
|
|
||||||
#### Verification
|
|
||||||
|
|
||||||
Best quality uses the JSON-upsampled prompts from `assets/` (download with
|
|
||||||
`hf download nvidia/Cosmos3-Nano assets/ --local-dir Cosmos3-Nano`). Minimal
|
|
||||||
self-contained examples:
|
|
||||||
|
|
||||||
```bash
|
|
||||||
curl http://localhost:8000/v1/models
|
|
||||||
|
|
||||||
# Text-to-image -> /v1/images/generations (1024x1024, 50 steps; base64 PNG)
|
|
||||||
curl -sS -X POST http://localhost:8000/v1/images/generations \
|
|
||||||
-H "Content-Type: application/json" \
|
|
||||||
-d '{
|
|
||||||
"model": "nvidia/Cosmos3-Nano",
|
|
||||||
"prompt": "A photorealistic red sports car on a city street at golden hour, cinematic lighting.",
|
|
||||||
"negative_prompt": "blurry, distorted, low quality",
|
|
||||||
"size": "1024x1024", "n": 1, "response_format": "b64_json",
|
|
||||||
"num_inference_steps": 50, "guidance_scale": 7.0, "seed": 42
|
|
||||||
}' | python -c "import sys,json,base64; open('cosmos3_t2i.png','wb').write(base64.b64decode(json.load(sys.stdin)['data'][0]['b64_json']))"
|
|
||||||
|
|
||||||
# Text-to-video -> /v1/videos/sync (720p, 189 frames @ 24fps; official params)
|
|
||||||
curl -sS -X POST http://localhost:8000/v1/videos/sync \
|
|
||||||
-H "Accept: video/mp4" \
|
|
||||||
-F "model=nvidia/Cosmos3-Nano" \
|
|
||||||
-F "prompt=A robot arm is cleaning a plate in the kitchen" \
|
|
||||||
-F "negative_prompt=blurry, distorted, low quality, jittery, deformed" \
|
|
||||||
-F "size=1280x720" -F "num_frames=189" -F "fps=24" \
|
|
||||||
-F "num_inference_steps=35" -F "guidance_scale=6.0" \
|
|
||||||
-F "max_sequence_length=4096" -F "flow_shift=10.0" \
|
|
||||||
-F 'extra_params={"use_resolution_template":false,"use_duration_template":false,"guardrails":true}' \
|
|
||||||
-F "seed=123" \
|
|
||||||
-o cosmos3_t2v.mp4
|
|
||||||
|
|
||||||
# Image-to-video -> /v1/videos/sync with an uploaded reference image
|
|
||||||
curl -sS -X POST http://localhost:8000/v1/videos/sync \
|
|
||||||
-H "Accept: video/mp4" \
|
|
||||||
-F "model=nvidia/Cosmos3-Nano" \
|
|
||||||
-F "prompt=The scene comes to life with smooth, natural motion." \
|
|
||||||
-F "negative_prompt=blurry, distorted, low quality" \
|
|
||||||
-F "size=1280x720" -F "num_frames=189" -F "fps=24" \
|
|
||||||
-F "num_inference_steps=35" -F "guidance_scale=6.0" \
|
|
||||||
-F "max_sequence_length=4096" -F "flow_shift=10.0" \
|
|
||||||
-F 'extra_params={"use_resolution_template":false,"use_duration_template":false,"guardrails":true}' \
|
|
||||||
-F "seed=1111" \
|
|
||||||
-F "input_reference=@/path/to/reference.jpg;type=image/jpeg" \
|
|
||||||
-o cosmos3_i2v.mp4
|
|
||||||
|
|
||||||
|
|
||||||
# Text-to-video-with-sound
|
|
||||||
curl -sS -X POST http://localhost:8000/v1/videos/sync \
|
|
||||||
-H "Accept: video/mp4" \
|
|
||||||
-F "model=nvidia/Cosmos3-Nano" \
|
|
||||||
-F "prompt=The video opens with a view of a well-lit indoor fruit display. A robotic arm picks up a pear, an orange, and a carambola one by one, placing each into a plastic bag in a shopping cart with red handles. The video is 7.875 seconds long, 24 FPS, and 1280x720. Audio description: soft servo whirs, gentle fruit thuds, plastic bag rustling, and a faint refrigeration hum." \
|
|
||||||
-F "negative_prompt=blurry, distorted, low quality" \
|
|
||||||
-F "size=1280x720" \
|
|
||||||
-F "num_frames=189" \
|
|
||||||
-F "fps=24" \
|
|
||||||
-F "num_inference_steps=35" \
|
|
||||||
-F "guidance_scale=6.0" \
|
|
||||||
-F "max_sequence_length=4096" \
|
|
||||||
-F "flow_shift=10.0" \
|
|
||||||
-F "seed=0" \
|
|
||||||
-F "generate_sound=true" \
|
|
||||||
-F "sound_duration=7.875" \
|
|
||||||
-F 'extra_params={"use_resolution_template":false,"use_duration_template":false,"guardrails":true}' \
|
|
||||||
-o cosmos3_t2v_with_sound.mp4
|
|
||||||
|
|
||||||
# Action — forward dynamics (first frame + action trajectory -> rollout video).
|
|
||||||
# Synchronous; `action` is a JSON array shaped [action_chunk_size, raw_action_dim].
|
|
||||||
curl -sS -X POST http://localhost:8000/v1/videos/sync \
|
|
||||||
-H "Accept: video/mp4" \
|
|
||||||
--form-string "model=nvidia/Cosmos3-Nano" \
|
|
||||||
--form-string "prompt=You are an autonomous vehicle. This video is captured from a first-person perspective." \
|
|
||||||
-F "input_reference=@first_frame.jpg;type=image/jpeg" \
|
|
||||||
-F "size=640x480" -F "num_frames=61" -F "fps=10" \
|
|
||||||
-F "num_inference_steps=30" -F "guidance_scale=1.0" -F "flow_shift=5.0" \
|
|
||||||
--form-string "extra_params={\"action_mode\":\"forward_dynamics\",\"domain_name\":\"av\",\"raw_action_dim\":9,\"action_chunk_size\":60,\"action\":$(cat action.json)}" \
|
|
||||||
-F "seed=0" \
|
|
||||||
-o cosmos3_forward_dynamics.mp4
|
|
||||||
|
|
||||||
# Action — policy (first frame + instruction -> predicted action trajectory + video).
|
|
||||||
# Asynchronous: POST returns a job id; poll, then read the predicted action from
|
|
||||||
# the top-level `action` field ({data, shape, dtype, raw_action_dim, domain_id}).
|
|
||||||
VIDEO_ID=$(curl -sS -X POST http://localhost:8000/v1/videos \
|
|
||||||
-H "Accept: application/json" \
|
|
||||||
--form-string "model=nvidia/Cosmos3-Nano" \
|
|
||||||
--form-string "prompt=Pick up the banana and place it in the bowl." \
|
|
||||||
-F "input_reference=@first_frame.jpg;type=image/jpeg" \
|
|
||||||
-F "size=640x480" -F "num_frames=17" -F "fps=5" \
|
|
||||||
-F "num_inference_steps=30" -F "guidance_scale=1.0" -F "flow_shift=5.0" \
|
|
||||||
--form-string 'extra_params={"action_mode":"policy","domain_name":"bridge_orig_lerobot","raw_action_dim":10,"action_chunk_size":16}' \
|
|
||||||
-F "seed=0" | jq -r '.id')
|
|
||||||
# poll until status == completed, then:
|
|
||||||
curl -sS "http://localhost:8000/v1/videos/$VIDEO_ID" | jq '.action | {shape, dtype, raw_action_dim, domain_id}'
|
|
||||||
curl -sS -L "http://localhost:8000/v1/videos/$VIDEO_ID/content" -o cosmos3_policy.mp4
|
|
||||||
```
|
|
||||||
|
|
||||||
#### Notes
|
|
||||||
|
|
||||||
- **Measured latency (1x B300, bf16, guardrails off):**
|
|
||||||
- T2I 1024² — 10 / 25 / 50 steps → ~0.4 / 0.7 / **1.3 s**
|
|
||||||
- T2V 1280×720 @ 35 steps — 25 / 49 / 93 / **189** frames → ~7 / 15 / 33 / **~93 s**
|
|
||||||
- I2V 1280×720, 189 frames @ 35 steps → ~**99 s**
|
|
||||||
- Action 640×480 @ 30 steps — forward-dynamics 61f ~**4 s**, policy 17f ~**1–3 s**.
|
|
||||||
- Guardrails-on overhead: ~8% on T2I, negligible on video.
|
|
||||||
- **Memory:** transformer ~17 GiB (bf16); peak ~46 GiB for 720p video on 1 GPU;
|
|
||||||
full repo (transformer + Wan VAE + Qwen3-VL vision encoder + audio tokenizer)
|
|
||||||
~33 GB on disk.
|
|
||||||
- **Determinism:** identical seed reproduces identical output on the same
|
|
||||||
hardware; outputs are not bit-identical across different GPU types.
|
|
||||||
- **Supported sizes (per model card):** 256p / 480p / 720p at 16:9, 4:3, 1:1,
|
|
||||||
3:4, 9:16. Defaults: T2I 1024², 50 steps, guidance 7.0; T2V/I2V 1280×720,
|
|
||||||
189 frames, 35 steps, guidance 6.0, `flow_shift=10.0`.
|
|
||||||
- **Key flags / params:** `--no-guardrails` (server) or
|
|
||||||
`extra_params={"guardrails":false}` (per request) toggles safety. The
|
|
||||||
per-request flag only takes effect when the server was launched **with**
|
|
||||||
guardrails enabled (it cannot re-enable them on a `--no-guardrails` server).
|
|
||||||
`use_resolution_template` / `use_duration_template` are off by default and only
|
|
||||||
needed when not using upsampled prompts that already encode resolution/duration.
|
|
||||||
- **Known limitations:**
|
|
||||||
- Guardrails-on requires `cosmos-guardrail` **and** access to the gated
|
|
||||||
`nvidia/Cosmos-1.0-Guardrail` repo (accept license + `HF_TOKEN`); otherwise
|
|
||||||
the server fails at pipeline build with a gated-repo / safety-checker error.
|
|
||||||
- A guardrail-blocked prompt currently returns HTTP 500
|
|
||||||
(`"Guardrail blocked prompt"`).
|
|
||||||
- Action `forward_dynamics` (sync `/v1/videos/sync`) and `policy` (async
|
|
||||||
`/v1/videos`, returns the predicted action under the top-level `action`
|
|
||||||
field) are supported online. **Online inference of inverse dynamics will be
|
|
||||||
added in a follow-up MR.**
|
|
||||||
|
|
||||||
### 1x GPU (Offline generation)
|
|
||||||
|
|
||||||
#### Environment
|
|
||||||
|
|
||||||
- OS: Ubuntu 22.04+
|
|
||||||
- Python: 3.12+
|
|
||||||
- Driver / runtime: NVIDIA CUDA environment
|
|
||||||
- vLLM-Omni version or commit: use the commit you are deploying from
|
|
||||||
|
|
||||||
#### Command
|
|
||||||
|
|
||||||
```python
|
|
||||||
# cosmos3_offline.py — run with: python cosmos3_offline.py
|
|
||||||
import torch
|
|
||||||
from vllm_omni.entrypoints.omni import Omni
|
|
||||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
|
||||||
|
|
||||||
|
|
||||||
def main():
|
|
||||||
omni = Omni(
|
|
||||||
model="nvidia/Cosmos3-Nano",
|
|
||||||
model_class_name="Cosmos3OmniDiffusersPipeline",
|
|
||||||
trust_remote_code=True,
|
|
||||||
enforce_eager=True,
|
|
||||||
# Guardrails are disabled here for a quick local run; install
|
|
||||||
# cosmos-guardrail + gated-repo access and drop this to enable them.
|
|
||||||
model_config={"guardrails": False},
|
|
||||||
)
|
|
||||||
gen = torch.Generator(device="cpu").manual_seed(42)
|
|
||||||
|
|
||||||
# Text-to-image (modalities=["image"]). For T2V use modalities=["video"]
|
|
||||||
# plus num_frames/fps; for I2V add multi_modal_data={"image": <PIL.Image>}.
|
|
||||||
outputs = omni.generate(
|
|
||||||
{
|
|
||||||
"prompt": "A photorealistic red sports car at golden hour, cinematic lighting.",
|
|
||||||
"negative_prompt": "blurry, distorted, low quality",
|
|
||||||
"modalities": ["image"],
|
|
||||||
},
|
|
||||||
OmniDiffusionSamplingParams(
|
|
||||||
height=1024, width=1024, generator=gen,
|
|
||||||
guidance_scale=7.0, num_inference_steps=50, num_outputs_per_prompt=1,
|
|
||||||
),
|
|
||||||
)
|
|
||||||
outputs[0].request_output.images[0].save("cosmos3_t2i.png")
|
|
||||||
omni.close()
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
main()
|
|
||||||
```
|
|
||||||
|
|
||||||
#### Verification
|
|
||||||
|
|
||||||
```bash
|
|
||||||
python cosmos3_offline.py
|
|
||||||
python -c "from PIL import Image; im=Image.open('cosmos3_t2i.png'); print('image', im.size, im.mode)"
|
|
||||||
```
|
|
||||||
|
|
||||||
#### Notes
|
|
||||||
|
|
||||||
- Same `Cosmos3OmniDiffusersPipeline` as online; mode is chosen by
|
|
||||||
`prompt["modalities"]` (`["image"]` → T2I, `["video"]` → T2V) plus
|
|
||||||
`num_frames`/`fps`, and `multi_modal_data={"image": ...}` for I2V. For video,
|
|
||||||
frames are returned in `outputs[0].request_output.images` as an
|
|
||||||
`(B, F, H, W, 3)` array.
|
|
||||||
- The offline entry must be guarded by `if __name__ == "__main__":` — the engine
|
|
||||||
spawns workers with the `spawn` start method.
|
|
||||||
@@ -1,112 +0,0 @@
|
|||||||
# Cosmos3-Super
|
|
||||||
|
|
||||||
> Frontier 64B world model: text-to-image, text-to-video, image-to-video (+ optional audio)
|
|
||||||
|
|
||||||
## Summary
|
|
||||||
|
|
||||||
- Vendor: NVIDIA
|
|
||||||
- Model: `nvidia/Cosmos3-Super` (64B; also `Cosmos3-Super-Text2Image`, `Cosmos3-Super-Image2Video`)
|
|
||||||
- Task: T2I, T2V, I2V generation, with optional synchronized audio (video + sound)
|
|
||||||
- Mode: Online serving with the OpenAI-compatible image/video APIs
|
|
||||||
- Maintainer: Community
|
|
||||||
|
|
||||||
## When to use this recipe
|
|
||||||
|
|
||||||
Use this recipe to deploy the 64B `nvidia/Cosmos3-Super` for the highest-quality
|
|
||||||
Cosmos3 generation. It shares the same `Cosmos3OmniDiffusersPipeline` and request
|
|
||||||
formats as [Cosmos3-Nano](./Cosmos3-Nano.md) — only the checkpoint size and the
|
|
||||||
recommended parallelism differ. Mode is selected per request (T2I →
|
|
||||||
`/v1/images/generations`; T2V/I2V → `/v1/videos/sync`; add `generate_sound=true`
|
|
||||||
for audio).
|
|
||||||
|
|
||||||
## References
|
|
||||||
|
|
||||||
- Model card (authoritative usage + example assets): <https://huggingface.co/nvidia/Cosmos3-Super>
|
|
||||||
- Nano recipe (same APIs/params): [`Cosmos3-Nano.md`](./Cosmos3-Nano.md)
|
|
||||||
- Pipeline: [`vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py`](../../vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py)
|
|
||||||
|
|
||||||
## Hardware Support
|
|
||||||
|
|
||||||
## GPU
|
|
||||||
|
|
||||||
Requires the `vllm-omni` package (or the `vllm/vllm-omni:cosmos3` container),
|
|
||||||
which provides the `vllm serve … --omni` entrypoint used below.
|
|
||||||
|
|
||||||
### 8x H200/H100/A100 (recommended, per model card)
|
|
||||||
|
|
||||||
```bash
|
|
||||||
vllm serve nvidia/Cosmos3-Super \
|
|
||||||
--omni \
|
|
||||||
--host 0.0.0.0 --port 8000 \
|
|
||||||
--cfg-parallel-size 2 \
|
|
||||||
--ulysses-degree 4 \
|
|
||||||
--use-hsdp --hsdp-shard-size 8 \
|
|
||||||
--init-timeout 1800
|
|
||||||
```
|
|
||||||
|
|
||||||
### 2x H200 / B300 (minimum)
|
|
||||||
|
|
||||||
```bash
|
|
||||||
vllm serve nvidia/Cosmos3-Super \
|
|
||||||
--omni \
|
|
||||||
--host 0.0.0.0 --port 8000 \
|
|
||||||
--cfg-parallel-size 2 \
|
|
||||||
--use-hsdp --hsdp-shard-size 2 \
|
|
||||||
--init-timeout 1800
|
|
||||||
```
|
|
||||||
|
|
||||||
Guardrails are on by default (gated `nvidia/Cosmos-1.0-Guardrail` — `pip install
|
|
||||||
cosmos-guardrail`, accept the license, set `HF_TOKEN`); add `--no-guardrails` to
|
|
||||||
disable. `--enable-layerwise-offload` reduces VRAM on smaller GPUs.
|
|
||||||
|
|
||||||
#### Verification
|
|
||||||
|
|
||||||
Requests are identical to Nano (see [`Cosmos3-Nano.md`](./Cosmos3-Nano.md) for full
|
|
||||||
T2I/T2V/I2V/T2VS curls); official params: `size=1280x720, num_frames=189, fps=24,
|
|
||||||
num_inference_steps=35, guidance_scale=6.0, flow_shift=10.0, max_sequence_length=4096`.
|
|
||||||
|
|
||||||
```bash
|
|
||||||
curl http://localhost:8000/v1/models
|
|
||||||
# T2V (official prompt assets give best quality)
|
|
||||||
curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \
|
|
||||||
-F "model=nvidia/Cosmos3-Super" -F "prompt=A robot arm is cleaning a plate in the kitchen" \
|
|
||||||
-F "size=1280x720" -F "num_frames=189" -F "fps=24" -F "num_inference_steps=35" \
|
|
||||||
-F "guidance_scale=6.0" -F "max_sequence_length=4096" -F "flow_shift=10.0" \
|
|
||||||
-F 'extra_params={"use_resolution_template":false,"use_duration_template":false,"guardrails":true}' \
|
|
||||||
-F "seed=17" -o cosmos3_super_t2v.mp4
|
|
||||||
|
|
||||||
# I2V — add an uploaded reference image
|
|
||||||
curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \
|
|
||||||
-F "model=nvidia/Cosmos3-Super" -F "prompt=The scene comes to life with smooth, natural motion." \
|
|
||||||
-F "size=1280x720" -F "num_frames=189" -F "fps=24" -F "num_inference_steps=35" \
|
|
||||||
-F "guidance_scale=6.0" -F "max_sequence_length=4096" -F "flow_shift=10.0" \
|
|
||||||
-F 'extra_params={"use_resolution_template":false,"use_duration_template":false,"guardrails":true}' \
|
|
||||||
-F "seed=1111" -F "input_reference=@/path/to/reference.jpg;type=image/jpeg" \
|
|
||||||
-o cosmos3_super_i2v.mp4
|
|
||||||
|
|
||||||
# T2V + sound — add generate_sound/sound_duration (output muxes AAC 48 kHz stereo)
|
|
||||||
curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \
|
|
||||||
-F "model=nvidia/Cosmos3-Super" -F "prompt=A robot arm is cleaning a plate in the kitchen" \
|
|
||||||
-F "size=1280x720" -F "num_frames=189" -F "fps=24" -F "num_inference_steps=35" \
|
|
||||||
-F "guidance_scale=6.0" -F "max_sequence_length=4096" -F "flow_shift=10.0" \
|
|
||||||
-F "generate_sound=true" -F "sound_duration=7.875" \
|
|
||||||
-F 'extra_params={"use_resolution_template":false,"use_duration_template":false,"guardrails":true}' \
|
|
||||||
-F "seed=17" -o cosmos3_super_t2vs.mp4
|
|
||||||
```
|
|
||||||
|
|
||||||
#### Notes
|
|
||||||
|
|
||||||
- **Measured (2x B300, bf16, guardrails off, official 2-GPU config above):**
|
|
||||||
- T2I 1024², 50 steps → **~6 s**
|
|
||||||
- T2V 1280×720, 189 frames, 35 steps → **~197 s**
|
|
||||||
- I2V 1280×720, 189 frames, 35 steps → **~200 s**
|
|
||||||
- T2V + sound (189 frames, 35 steps) → **~198 s**, output muxes **AAC 48 kHz stereo**
|
|
||||||
- (NVIDIA's reference: 8×H200 @ 50 steps ≈ 55 s/video; 2×H200 @ 35 steps ≈ 3 min/video.)
|
|
||||||
- **Memory:** ~61.5 GiB per GPU when sharded across 2 GPUs (HSDP shard 2); repo ~135 GB on disk.
|
|
||||||
- Same generation defaults, supported sizes, and `generate_sound`/`sound_duration`
|
|
||||||
semantics as Nano, including the **action** modality: `forward_dynamics`
|
|
||||||
(sync `/v1/videos/sync`) and `policy` (async `/v1/videos`, predicted action under
|
|
||||||
the top-level `action` field) — see the Cosmos3-Nano recipe for the request shape.
|
|
||||||
Online inference of inverse dynamics will be added in a follow-up MR. Verified on
|
|
||||||
the 64B Super under `--cfg-parallel-size 2`: async `policy` returns the predicted
|
|
||||||
action (`[16, 10]`) and the rollout video reliably.
|
|
||||||
@@ -57,6 +57,12 @@ Adjust `devices` in the YAML to match your hardware.
|
|||||||
|
|
||||||
#### Command
|
#### Command
|
||||||
|
|
||||||
|
Thinker only (text output):
|
||||||
|
|
||||||
|
```bash
|
||||||
|
vllm serve Jonathan1909/Ming-flash-omni-2.0 --omni --port 8091
|
||||||
|
```
|
||||||
|
|
||||||
Thinker + talker (text and/or audio output):
|
Thinker + talker (text and/or audio output):
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
@@ -66,15 +72,6 @@ vllm serve Jonathan1909/Ming-flash-omni-2.0 \
|
|||||||
--log-stats
|
--log-stats
|
||||||
```
|
```
|
||||||
|
|
||||||
Thinker only (text-only output):
|
|
||||||
|
|
||||||
```bash
|
|
||||||
vllm serve Jonathan1909/Ming-flash-omni-2.0 \
|
|
||||||
--omni \
|
|
||||||
--deploy-config vllm_omni/deploy/ming_flash_omni_thinker_only.yaml \
|
|
||||||
--port 8091
|
|
||||||
```
|
|
||||||
|
|
||||||
`--log-stats` is optional but recommended while validating the deployment.
|
`--log-stats` is optional but recommended while validating the deployment.
|
||||||
|
|
||||||
#### Verification
|
#### Verification
|
||||||
|
|||||||
@@ -1,10 +1,4 @@
|
|||||||
# Common dependencies for all platforms
|
# Common dependencies for all platforms
|
||||||
# transformers: keep aligned with upstream vLLM constraint.
|
|
||||||
# vLLM pins ">= 4.56.0, != 5.0.*, != 5.1.*, != 5.2.*, != 5.3.*, != 5.4.*, != 5.5.0".
|
|
||||||
# 5.9.0 breaks models that define _keys_to_ignore_on_load_unexpected as a list
|
|
||||||
# (e.g. OpenMOSS-Team/MOSS-TTS-Nano) because it changed
|
|
||||||
# _adjust_missing_and_unexpected_keys from list+list to set|set.
|
|
||||||
transformers >= 4.56.0, != 5.0.*, != 5.1.*, != 5.2.*, != 5.3.*, != 5.4.*, != 5.5.0, < 5.9.0
|
|
||||||
av>=14.0.0
|
av>=14.0.0
|
||||||
omegaconf>=2.3.0
|
omegaconf>=2.3.0
|
||||||
diffusers==0.38.0
|
diffusers==0.38.0
|
||||||
|
|||||||
Binary file not shown.
|
Before Width: | Height: | Size: 1.1 MiB |
@@ -345,14 +345,6 @@ def test_parse_score_payload_handles_raw_json_and_delimited_json():
|
|||||||
assert parse_score_payload(wrapped)["score"] == [6]
|
assert parse_score_payload(wrapped)["score"] == [6]
|
||||||
|
|
||||||
|
|
||||||
def test_parse_score_payload_handles_qwen_vl_nested_score_dicts():
|
|
||||||
qwen_style = '{"score": [{"naturalness": 8, "artifact_free": 7}], "reasoning": "good quality"}'
|
|
||||||
assert parse_score_payload(qwen_style)["score"] == [8, 7]
|
|
||||||
|
|
||||||
nested_values = '{"score": [{"score": 9}, {"score": 6}], "reasoning": "ok"}'
|
|
||||||
assert parse_score_payload(nested_values)["score"] == [9, 6]
|
|
||||||
|
|
||||||
|
|
||||||
def test_summarize_gedit_generated_records_groups_by_task_and_language():
|
def test_summarize_gedit_generated_records_groups_by_task_and_language():
|
||||||
records = []
|
records = []
|
||||||
for group in GEDIT_GROUPS[:2]:
|
for group in GEDIT_GROUPS[:2]:
|
||||||
|
|||||||
@@ -30,12 +30,12 @@ from PIL import Image
|
|||||||
from pytest_mock import MockerFixture
|
from pytest_mock import MockerFixture
|
||||||
from vllm import SamplingParams
|
from vllm import SamplingParams
|
||||||
from vllm.outputs import CompletionOutput, RequestOutput
|
from vllm.outputs import CompletionOutput, RequestOutput
|
||||||
|
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||||
|
|
||||||
from vllm_omni.entrypoints.async_omni import AsyncOmni as RealAsyncOmni
|
from vllm_omni.entrypoints.async_omni import AsyncOmni as RealAsyncOmni
|
||||||
from vllm_omni.entrypoints.cli.serve import OmniServeCommand
|
from vllm_omni.entrypoints.cli.serve import OmniServeCommand
|
||||||
from vllm_omni.inputs.data import OmniSamplingParams
|
from vllm_omni.inputs.data import OmniSamplingParams
|
||||||
from vllm_omni.outputs import OmniRequestOutput
|
from vllm_omni.outputs import OmniRequestOutput
|
||||||
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
|
|
||||||
|
|
||||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||||
|
|
||||||
@@ -444,7 +444,7 @@ def mock_async_omni(
|
|||||||
@pytest.fixture
|
@pytest.fixture
|
||||||
def api_server(unused_tcp_port_factory, server_case: ServerCase, mock_async_omni):
|
def api_server(unused_tcp_port_factory, server_case: ServerCase, mock_async_omni):
|
||||||
"""Set up a API server in background process from command line with parametrized model name and mocked AsyncOmni."""
|
"""Set up a API server in background process from command line with parametrized model name and mocked AsyncOmni."""
|
||||||
parser = TrackingArgumentParser()
|
parser = FlexibleArgumentParser()
|
||||||
subparsers = parser.add_subparsers(dest="command")
|
subparsers = parser.add_subparsers(dest="command")
|
||||||
cmd = OmniServeCommand()
|
cmd = OmniServeCommand()
|
||||||
cmd.subparser_init(subparsers)
|
cmd.subparser_init(subparsers)
|
||||||
|
|||||||
@@ -61,7 +61,7 @@ def test_stage0_streaming_update_discards_outstanding_async_placeholder_token()
|
|||||||
|
|
||||||
sched._update_request_as_session(session, _make_update([10, 20]))
|
sched._update_request_as_session(session, _make_update([10, 20]))
|
||||||
|
|
||||||
assert session.async_tokens_to_discard == 1
|
assert session.discard_latest_async_tokens is True
|
||||||
assert session.num_output_placeholders == 0
|
assert session.num_output_placeholders == 0
|
||||||
assert session.spec_token_ids == []
|
assert session.spec_token_ids == []
|
||||||
# The async placeholder makes token 9 unconfirmed, so only 7 and 8 are
|
# The async placeholder makes token 9 unconfirmed, so only 7 and 8 are
|
||||||
@@ -83,7 +83,7 @@ def test_stage0_streaming_update_keeps_all_computed_tokens_without_placeholder()
|
|||||||
|
|
||||||
sched._update_request_as_session(session, _make_update([10, 20]))
|
sched._update_request_as_session(session, _make_update([10, 20]))
|
||||||
|
|
||||||
assert getattr(session, "async_tokens_to_discard", 0) == 0
|
assert session.discard_latest_async_tokens is False
|
||||||
assert session.num_output_placeholders == 0
|
assert session.num_output_placeholders == 0
|
||||||
assert session.prompt_token_ids == [1, 2, 3, 7, 8, 9, 10, 20]
|
assert session.prompt_token_ids == [1, 2, 3, 7, 8, 9, 10, 20]
|
||||||
assert list(session._all_token_ids) == [1, 2, 3, 7, 8, 9, 10, 20]
|
assert list(session._all_token_ids) == [1, 2, 3, 7, 8, 9, 10, 20]
|
||||||
|
|||||||
@@ -157,75 +157,6 @@ def test_update_no_multimodal():
|
|||||||
assert torch.all(slot_states == new_states)
|
assert torch.all(slot_states == new_states)
|
||||||
|
|
||||||
|
|
||||||
def test_update_uses_precomputed_hidden_states_cpu():
|
|
||||||
"""Precomputed CPU staging should be the only hidden-state cache source."""
|
|
||||||
cache = get_omni_pcache()
|
|
||||||
|
|
||||||
num_tokens_unpadded = 8
|
|
||||||
slot_offset = 8
|
|
||||||
slot_mapping = torch.arange(slot_offset, slot_offset + num_tokens_unpadded)
|
|
||||||
hidden_states = torch.zeros((num_tokens_unpadded, HIDDEN_SIZE), dtype=DTYPE)
|
|
||||||
hidden_states_cpu = torch.rand((num_tokens_unpadded, HIDDEN_SIZE), dtype=DTYPE)
|
|
||||||
|
|
||||||
cache.update_omni_tensor_prefix_cache(
|
|
||||||
hidden_states=hidden_states,
|
|
||||||
hidden_states_cpu=hidden_states_cpu,
|
|
||||||
multimodal_outputs=None,
|
|
||||||
num_tokens_unpadded=num_tokens_unpadded,
|
|
||||||
slot_mapping=slot_mapping,
|
|
||||||
)
|
|
||||||
|
|
||||||
hs_rows = cache.hidden_states_cache.view(NUM_BLOCKS * BLOCK_SIZE, HIDDEN_SIZE)
|
|
||||||
for slot_idx, new_states in zip(slot_mapping, hidden_states_cpu):
|
|
||||||
slot_states = hs_rows[slot_idx]
|
|
||||||
assert torch.all(slot_states == new_states)
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.parametrize(
|
|
||||||
"hidden_states_cpu",
|
|
||||||
[
|
|
||||||
torch.rand((4, HIDDEN_SIZE), dtype=DTYPE),
|
|
||||||
torch.rand((8, HIDDEN_SIZE), dtype=OTHER_DTYPE),
|
|
||||||
torch.rand((8, HIDDEN_SIZE + 1), dtype=DTYPE),
|
|
||||||
torch.rand((HIDDEN_SIZE, 8), dtype=DTYPE).t(),
|
|
||||||
],
|
|
||||||
)
|
|
||||||
def test_update_rejects_invalid_hidden_states_cpu(hidden_states_cpu):
|
|
||||||
cache = get_omni_pcache()
|
|
||||||
|
|
||||||
num_tokens_unpadded = 8
|
|
||||||
slot_mapping = torch.arange(num_tokens_unpadded)
|
|
||||||
hidden_states = torch.rand((num_tokens_unpadded, HIDDEN_SIZE), dtype=DTYPE)
|
|
||||||
|
|
||||||
with pytest.raises(RuntimeError):
|
|
||||||
cache.update_omni_tensor_prefix_cache(
|
|
||||||
hidden_states=hidden_states,
|
|
||||||
hidden_states_cpu=hidden_states_cpu,
|
|
||||||
multimodal_outputs=None,
|
|
||||||
num_tokens_unpadded=num_tokens_unpadded,
|
|
||||||
slot_mapping=slot_mapping,
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
|
|
||||||
def test_update_rejects_gpu_hidden_states_cpu():
|
|
||||||
cache = get_omni_pcache()
|
|
||||||
|
|
||||||
num_tokens_unpadded = 8
|
|
||||||
slot_mapping = torch.arange(num_tokens_unpadded)
|
|
||||||
hidden_states = torch.rand((num_tokens_unpadded, HIDDEN_SIZE), dtype=DTYPE)
|
|
||||||
hidden_states_cpu = torch.rand((num_tokens_unpadded, HIDDEN_SIZE), dtype=DTYPE, device="cuda")
|
|
||||||
|
|
||||||
with pytest.raises(RuntimeError):
|
|
||||||
cache.update_omni_tensor_prefix_cache(
|
|
||||||
hidden_states=hidden_states,
|
|
||||||
hidden_states_cpu=hidden_states_cpu,
|
|
||||||
multimodal_outputs=None,
|
|
||||||
num_tokens_unpadded=num_tokens_unpadded,
|
|
||||||
slot_mapping=slot_mapping,
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.parametrize(
|
@pytest.mark.parametrize(
|
||||||
"feat_dims",
|
"feat_dims",
|
||||||
[
|
[
|
||||||
@@ -428,76 +359,6 @@ def test_get_merged_hidden_states(num_tokens_padded, mocker):
|
|||||||
assert torch.all(req2_merged_states == req2_new_states)
|
assert torch.all(req2_merged_states == req2_new_states)
|
||||||
|
|
||||||
|
|
||||||
def test_get_merged_hidden_states_uses_precomputed_hidden_states_cpu(mocker):
|
|
||||||
cache = get_omni_pcache()
|
|
||||||
|
|
||||||
orig_num_tokens_unpadded = 8
|
|
||||||
slot_offset = 8
|
|
||||||
orig_slot_mapping = torch.arange(slot_offset, slot_offset + orig_num_tokens_unpadded)
|
|
||||||
orig_hidden_states = torch.rand((orig_num_tokens_unpadded, HIDDEN_SIZE), dtype=DTYPE)
|
|
||||||
|
|
||||||
cache.update_omni_tensor_prefix_cache(
|
|
||||||
hidden_states=orig_hidden_states,
|
|
||||||
multimodal_outputs=None,
|
|
||||||
num_tokens_unpadded=orig_num_tokens_unpadded,
|
|
||||||
slot_mapping=orig_slot_mapping,
|
|
||||||
)
|
|
||||||
|
|
||||||
num_new_toks_req1 = 3
|
|
||||||
num_new_toks_req2 = 2
|
|
||||||
cache.add_prefix_cached_new_req_id("req1")
|
|
||||||
num_scheduled_tokens = {
|
|
||||||
"req1": num_new_toks_req1,
|
|
||||||
"req2": num_new_toks_req2,
|
|
||||||
}
|
|
||||||
new_hidden_states = torch.zeros(
|
|
||||||
(num_new_toks_req1 + num_new_toks_req2, HIDDEN_SIZE),
|
|
||||||
dtype=DTYPE,
|
|
||||||
)
|
|
||||||
hidden_states_cpu = torch.rand_like(new_hidden_states)
|
|
||||||
req1_new_states = hidden_states_cpu[:num_new_toks_req1]
|
|
||||||
req2_new_states = hidden_states_cpu[-num_new_toks_req2:]
|
|
||||||
input_batch = MockInputBatch(num_computed_tokens_cpu=torch.Tensor([orig_num_tokens_unpadded, 0]))
|
|
||||||
|
|
||||||
mocker.patch(
|
|
||||||
"vllm_omni.core.prefix_cache.OmniTensorPrefixCache._get_cached_block_ids",
|
|
||||||
new=fake_get_cached_block_ids,
|
|
||||||
)
|
|
||||||
merged_states = cache.get_merged_hidden_states(
|
|
||||||
query_start_loc=[0, num_new_toks_req1],
|
|
||||||
input_batch=input_batch,
|
|
||||||
hidden_states=new_hidden_states,
|
|
||||||
hidden_states_cpu=hidden_states_cpu,
|
|
||||||
num_scheduled_tokens=num_scheduled_tokens,
|
|
||||||
)
|
|
||||||
|
|
||||||
assert torch.all(merged_states["req1"][:orig_num_tokens_unpadded] == orig_hidden_states)
|
|
||||||
assert torch.all(merged_states["req1"][-num_new_toks_req1:] == req1_new_states)
|
|
||||||
assert torch.all(merged_states["req2"] == req2_new_states)
|
|
||||||
|
|
||||||
|
|
||||||
def test_get_merged_hidden_states_rejects_short_hidden_states_cpu(mocker):
|
|
||||||
cache = get_omni_pcache()
|
|
||||||
num_scheduled_tokens = {
|
|
||||||
"req1": 3,
|
|
||||||
"req2": 2,
|
|
||||||
}
|
|
||||||
input_batch = MockInputBatch(num_computed_tokens_cpu=torch.Tensor([0, 0]))
|
|
||||||
|
|
||||||
mocker.patch(
|
|
||||||
"vllm_omni.core.prefix_cache.OmniTensorPrefixCache._get_cached_block_ids",
|
|
||||||
new=fake_get_cached_block_ids,
|
|
||||||
)
|
|
||||||
with pytest.raises(RuntimeError):
|
|
||||||
cache.get_merged_hidden_states(
|
|
||||||
query_start_loc=[0, 3],
|
|
||||||
input_batch=input_batch,
|
|
||||||
hidden_states=torch.rand((5, HIDDEN_SIZE), dtype=DTYPE),
|
|
||||||
hidden_states_cpu=torch.rand((4, HIDDEN_SIZE), dtype=DTYPE),
|
|
||||||
num_scheduled_tokens=num_scheduled_tokens,
|
|
||||||
)
|
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.parametrize("num_tokens_padded", [None, 16])
|
@pytest.mark.parametrize("num_tokens_padded", [None, 16])
|
||||||
@pytest.mark.parametrize(
|
@pytest.mark.parametrize(
|
||||||
"feat_dims",
|
"feat_dims",
|
||||||
|
|||||||
@@ -223,12 +223,8 @@ def _get_open_port() -> int:
|
|||||||
return s.getsockname()[1]
|
return s.getsockname()[1]
|
||||||
|
|
||||||
|
|
||||||
def _wait_for_port(host: str, port: int, timeout: int = 1200, proc: subprocess.Popen | None = None) -> None:
|
def _wait_for_port(host: str, port: int, timeout: int = 1200) -> None:
|
||||||
"""Block until the given host:port accepts connections or timeout expires.
|
"""Block until the given host:port accepts connections or timeout expires."""
|
||||||
|
|
||||||
If *proc* is provided, also monitors the process; raises RuntimeError
|
|
||||||
immediately if the server process exits before the port becomes available.
|
|
||||||
"""
|
|
||||||
start = time.time()
|
start = time.time()
|
||||||
while time.time() - start < timeout:
|
while time.time() - start < timeout:
|
||||||
try:
|
try:
|
||||||
@@ -238,10 +234,6 @@ def _wait_for_port(host: str, port: int, timeout: int = 1200, proc: subprocess.P
|
|||||||
return
|
return
|
||||||
except Exception:
|
except Exception:
|
||||||
pass
|
pass
|
||||||
if proc is not None:
|
|
||||||
ret = proc.poll()
|
|
||||||
if ret is not None:
|
|
||||||
raise RuntimeError(f"Server process exited with code {ret} before port {host}:{port} became ready")
|
|
||||||
time.sleep(2)
|
time.sleep(2)
|
||||||
raise RuntimeError(f"Server did not start on {host}:{port} within {timeout}s")
|
raise RuntimeError(f"Server did not start on {host}:{port} within {timeout}s")
|
||||||
|
|
||||||
@@ -339,7 +331,7 @@ class DiffusionServer:
|
|||||||
env=env,
|
env=env,
|
||||||
cwd=str(Path(__file__).parent.parent.parent.parent),
|
cwd=str(Path(__file__).parent.parent.parent.parent),
|
||||||
)
|
)
|
||||||
_wait_for_port(self.host, self.port, proc=self.proc)
|
_wait_for_port(self.host, self.port)
|
||||||
print(f"DiffusionServer ready on {self.host}:{self.port}")
|
print(f"DiffusionServer ready on {self.host}:{self.port}")
|
||||||
|
|
||||||
def __enter__(self):
|
def __enter__(self):
|
||||||
|
|||||||
@@ -1,7 +1,7 @@
|
|||||||
[
|
[
|
||||||
{
|
{
|
||||||
"test_name": "test_hunyuan_image3_ti2i_inline",
|
"test_name": "test_hunyuan_image3_ti2i_tp4_tp4_inline",
|
||||||
"description": "AR (TP=2) + 2 * DiT (TP=2) - using inline dataset and deploy-config",
|
"description": "AR (TP=4) + DiT (TP=4) - using inline dataset and deploy-config",
|
||||||
"server_type": "vllm-omni",
|
"server_type": "vllm-omni",
|
||||||
"benchmark_endpoint": "/v1/images/edits",
|
"benchmark_endpoint": "/v1/images/edits",
|
||||||
"server_params": {
|
"server_params": {
|
||||||
@@ -21,7 +21,7 @@
|
|||||||
"zmq_port": 50051,
|
"zmq_port": 50051,
|
||||||
"protocol": "rdma",
|
"protocol": "rdma",
|
||||||
"device_name": "",
|
"device_name": "",
|
||||||
"memory_pool_size": 17179869184,
|
"memory_pool_size": 4294967296,
|
||||||
"memory_pool_device": "cpu"
|
"memory_pool_device": "cpu"
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -32,13 +32,13 @@
|
|||||||
"is_comprehension": true,
|
"is_comprehension": true,
|
||||||
"final_output": true,
|
"final_output": true,
|
||||||
"final_output_type": "text",
|
"final_output_type": "text",
|
||||||
"max_num_seqs": 16,
|
"max_num_seqs": 4,
|
||||||
"gpu_memory_utilization": 0.95,
|
"gpu_memory_utilization": 0.95,
|
||||||
"enforce_eager": false,
|
"enforce_eager": false,
|
||||||
"enable_prefix_caching": false,
|
"enable_prefix_caching": false,
|
||||||
"max_num_batched_tokens": 32768,
|
"max_num_batched_tokens": 32768,
|
||||||
"devices": "0,1",
|
"devices": "0,1,2,3",
|
||||||
"tensor_parallel_size": 2,
|
"tensor_parallel_size": 4,
|
||||||
"hf_overrides": {
|
"hf_overrides": {
|
||||||
"rope_parameters": {
|
"rope_parameters": {
|
||||||
"mrope_section": [0, 32, 32],
|
"mrope_section": [0, 32, 32],
|
||||||
@@ -65,14 +65,13 @@
|
|||||||
"stage_id": 1,
|
"stage_id": 1,
|
||||||
"max_num_seqs": 1,
|
"max_num_seqs": 1,
|
||||||
"enforce_eager": true,
|
"enforce_eager": true,
|
||||||
"devices": "2,3",
|
"devices": "4,5,6,7",
|
||||||
"distributed_executor_backend": "mp",
|
"distributed_executor_backend": "mp",
|
||||||
"omni_kv_config": {
|
"omni_kv_config": {
|
||||||
"need_recv_cache": true
|
"need_recv_cache": true
|
||||||
},
|
},
|
||||||
"parallel_config": {
|
"parallel_config": {
|
||||||
"tensor_parallel_size": 2,
|
"tensor_parallel_size": 4,
|
||||||
"vae_patch_parallel_size": 2,
|
|
||||||
"enable_expert_parallel": true
|
"enable_expert_parallel": true
|
||||||
},
|
},
|
||||||
"input_connectors": {
|
"input_connectors": {
|
||||||
@@ -115,8 +114,8 @@
|
|||||||
},
|
},
|
||||||
"task": "ti2i",
|
"task": "ti2i",
|
||||||
"bot-task": "think_recaption",
|
"bot-task": "think_recaption",
|
||||||
"max-concurrency": 16,
|
"max-concurrency": 4,
|
||||||
"num-prompts": 96,
|
"num-prompts": 16,
|
||||||
"skip-performance-assertion": true,
|
"skip-performance-assertion": true,
|
||||||
"baseline": {
|
"baseline": {
|
||||||
"throughput_qps": 0.1,
|
"throughput_qps": 0.1,
|
||||||
|
|||||||
Loaded 100 of 393 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user