Compare commits
81
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
7a2007cc17 | ||
|
|
76de642c54 | ||
|
|
a4ea67a21b | ||
|
|
e0c2b9d73f | ||
|
|
24741961eb | ||
|
|
727f90958e | ||
|
|
67c54777bb | ||
|
|
2c4e9a945b | ||
|
|
4b1547b4d8 | ||
|
|
6e8c279a68 | ||
|
|
6f98c807d2 | ||
|
|
185e3d6df6 | ||
|
|
67aa6c55a4 | ||
|
|
9d67f8e02f | ||
|
|
cd5daaea6a | ||
|
|
8cb79308c6 | ||
|
|
0a03359224 | ||
|
|
fe02e8b81e | ||
|
|
896f05e7e6 | ||
|
|
d33c905d80 | ||
|
|
4715570acc | ||
|
|
4956a41f12 | ||
|
|
2739f0776b | ||
|
|
2d8ef4754c | ||
|
|
f8a8fbf72e | ||
|
|
226cb4a38c | ||
|
|
0a51a18570 | ||
|
|
499d930fe5 | ||
|
|
273fc8eb1d | ||
|
|
a951dd0cb4 | ||
|
|
4563999ec0 | ||
|
|
86d7864588 | ||
|
|
3b0718f22e | ||
|
|
ce3d8e4bac | ||
|
|
81cd7964d4 | ||
|
|
09c6ad4204 | ||
|
|
f7e87e7cfd | ||
|
|
744c65b730 | ||
|
|
b3f4fbf9e2 | ||
|
|
1752e8ee5b | ||
|
|
089bb88706 | ||
|
|
04b408ffd2 | ||
|
|
0c88f8a0b6 | ||
|
|
1d6c0816a4 | ||
|
|
d50ea9fccc | ||
|
|
daec3a6a49 | ||
|
|
0461da345a | ||
|
|
aa5364f705 | ||
|
|
1b718f7b15 | ||
|
|
dba05c3517 | ||
|
|
f2d9eb1630 | ||
|
|
fa0658eabb | ||
|
|
7c17111ff5 | ||
|
|
925347332f | ||
|
|
7be7e0d249 | ||
|
|
18be2ba3df | ||
|
|
1c6e731339 | ||
|
|
b8e7b21188 | ||
|
|
39663d5287 | ||
|
|
ab0da662c0 | ||
|
|
fa41fdf252 | ||
|
|
7fa40681ff | ||
|
|
23bb4395bb | ||
|
|
0feb92bdfa | ||
|
|
f6e59c329b | ||
|
|
c96fe54bd8 | ||
|
|
c2b55de5b9 | ||
|
|
61324e65d7 | ||
|
|
f775e80732 | ||
|
|
c96f0bc773 | ||
|
|
f4640d9a2e | ||
|
|
12d75cf63a | ||
|
|
66a9b0c84f | ||
|
|
cee2e40c33 | ||
|
|
5014fea4a0 | ||
|
|
f2c5152df0 | ||
|
|
d3e0171cd7 | ||
|
|
807db6efd7 | ||
|
|
d4a5eaac90 | ||
|
|
8f115164c2 | ||
|
|
cbab1d4d8f |
No files matched your search
@@ -89,7 +89,7 @@ steps:
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- pytest -s -v tests/diffusion/distributed/test_zimage_parallelism.py
|
||||
- pytest -s -v tests/diffusion/distributed/test_tensor_parallel.py
|
||||
|
||||
- label: "Diffusion Model Test"
|
||||
agent_pool: mi325_1
|
||||
|
||||
@@ -35,7 +35,7 @@ steps:
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- timeout 20m pytest -s -v tests/e2e/offline_inference/custom_pipeline/ -m "core_model"
|
||||
- timeout 20m pytest -s -v tests/e2e/features/custom_pipeline/ -m "core_model"
|
||||
|
||||
|
||||
- label: "Voxtral TTS CUDA Unit Test"
|
||||
|
||||
@@ -8,14 +8,24 @@ steps:
|
||||
- group: ":card_index_dividers: Simple Test"
|
||||
depends_on: upload-merge-pipeline
|
||||
steps:
|
||||
- label: "Simple · Diffusion & Model Executor Test"
|
||||
- label: "Simple · Model Executor Test"
|
||||
commands:
|
||||
- "timeout 20m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||
- "timeout 20m pytest -sv tests/model_executor -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
- label: "Simple · Diffusion Test"
|
||||
commands:
|
||||
- "timeout 20m pytest -sv tests/diffusion -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
- label: "Simple · Engine&Entrypoints Test"
|
||||
commands:
|
||||
- "timeout 20m pytest -sv tests/entrypoints tests/engine -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
- label: "Simple · Other Test"
|
||||
commands:
|
||||
- "timeout 20m pytest -sv -m 'core_model and cpu' --ignore=tests/diffusion --ignore=tests/model_executor --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||
- "timeout 20m pytest -sv tests/ -m 'core_model and cpu' --ignore=tests/diffusion --ignore=tests/model_executor --ignore=tests/entrypoints --ignore=tests/engine --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
- group: ":card_index_dividers: Diffusion Test"
|
||||
@@ -27,12 +37,6 @@ steps:
|
||||
- pytest -sv tests/diffusion/lora/ -m 'advanced_model and cuda' --run-level "advanced_model"
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
- label: "Diffusion · Offloader Test"
|
||||
timeout_in_minutes: 40
|
||||
commands:
|
||||
- timeout 40m pytest -s -v tests/diffusion/offloader -m "advanced_model and cuda"
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
- label: "Distributed Test"
|
||||
depends_on: upload-merge-pipeline
|
||||
timeout_in_minutes: 60
|
||||
@@ -49,13 +53,17 @@ steps:
|
||||
depends_on: upload-merge-pipeline
|
||||
steps:
|
||||
- label: "Entrypoint Test with H100"
|
||||
timeout_in_minutes: 60
|
||||
# TEMPORARY 60 -> 90: engine startup/sleep-wake is ~13-15 min slower per
|
||||
# run on the vLLM 0.26 image (model loading ~2x, init ~2x; see builds
|
||||
# 2915 vs 2916), which pushed this suite past 60 min. Revert once the
|
||||
# 0.26 startup regression is understood and fixed.
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and H100" --run-level "advanced_model"
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: "Entrypoint Test with L4"
|
||||
timeout_in_minutes: 30
|
||||
timeout_in_minutes: 60
|
||||
commands:
|
||||
- pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and L4" --run-level "advanced_model"
|
||||
mirror_hardwares: l4_1
|
||||
@@ -131,39 +139,32 @@ steps:
|
||||
timeout_in_minutes: 50
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/offline_inference/test_minicpmo_4_5.py tests/e2e/online_serving/test_minicpmo_4_5.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 2
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: "Omni · MiniCPM-o 4.5 Duplex Test"
|
||||
source_file_dependencies:
|
||||
- tests/assets/minicpmo_4_5/response_required_16k.wav
|
||||
- tests/assets/minicpmo_4_5/soft_interrupt_16k.wav
|
||||
- tests/e2e/online_serving/test_minicpmo_4_5_duplex.py
|
||||
- tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py
|
||||
- tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py
|
||||
- tests/e2e/online_serving/minicpmo_realtime_duplex_scenarios.py
|
||||
- tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py
|
||||
- tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py
|
||||
- tests/helpers/minicpmo_4_5_duplex.py
|
||||
- vllm_omni/deploy/minicpmo_4_5_duplex.yaml
|
||||
- vllm_omni/experimental/fullduplex/
|
||||
- vllm_omni/model_executor/models/minicpmo_4_5/
|
||||
- vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py
|
||||
- vllm_omni/core/sched/
|
||||
- vllm_omni/outputs/multimodal_accumulation.py
|
||||
- vllm_omni/engine/stage_init_utils.py
|
||||
- vllm_omni/engine/async_omni_engine.py
|
||||
- vllm_omni/entrypoints/cli/serve.py
|
||||
timeout_in_minutes: 50
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: "TTS · VoxCPM2 Test"
|
||||
source_file_dependencies:
|
||||
@@ -181,22 +182,6 @@ steps:
|
||||
"
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
- label: "TTS · SoulX-Singer Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/offline_inference/test_soulxsinger.py
|
||||
- vllm_omni/diffusion/models/soulx_singer/
|
||||
- vllm_omni/deploy/soulxsinger_svc.yaml
|
||||
- vllm_omni/deploy/soulxsinger_svs.yaml
|
||||
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/offline_inference/test_soulxsinger.py -k 'svc or precomputed' -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
"
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
- label: "Diffusion · Qwen Image Edit Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_qwen_image_edit.py
|
||||
|
||||
@@ -15,7 +15,13 @@ steps:
|
||||
- label: ":full_moon: Omni · Function Test with H100"
|
||||
timeout_in_minutes: 120
|
||||
commands:
|
||||
- pytest -sv tests/e2e --run-level "full_model" -m "full_model and H100 and omni" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py
|
||||
- pytest -sv tests/e2e --run-level "full_model" -m "full_model and H100 and omni" --ignore=tests/e2e/accuracy --ignore=tests/e2e/online_serving/test_qwen3_omni_multi_replicas.py --ignore=tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: ":full_moon: Omni · MiniCPM-o 4.5 Duplex Test"
|
||||
timeout_in_minutes: 50
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py -m "full_model and cuda and H100 and omni" --run-level "full_model"
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: ":full_moon: Omni · Doc Test with L4"
|
||||
@@ -426,7 +432,7 @@ steps:
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/accuracy/hunyuanvideo15_t2v/test_hunyuanvideo15_t2v_video_similarity.py -m full_model --run-level full_model
|
||||
- pytest -s -v tests/e2e/accuracy/hunyuanvideo15_i2v/test_hunyuanvideo15_i2v_video_similarity.py -m full_model --run-level full_model
|
||||
- pytest -s -v tests/e2e/accuracy/test_ltx2_3_video_similarity.py -m full_model --run-level full_model
|
||||
- pytest -s -v 'tests/e2e/accuracy/ltx/test_ltx_official_similarity.py::test_ltx_one_stage_matches_official[ltx2_3_i2v]' -m full_model --run-level full_model
|
||||
- pytest -s -v tests/e2e/accuracy/test_video_streaming_output_similarity.py -m full_model --run-level full_model
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
@@ -449,36 +455,41 @@ steps:
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- group: ":card_index_dividers: Diffusion Test"
|
||||
key: nightly-diffusion-test-group
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "diffusion-x2iat-test"
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
|
||||
steps:
|
||||
- label: ":full_moon: Diffusion · Distributed Test with L4"
|
||||
timeout_in_minutes: 60
|
||||
commands:
|
||||
- pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and L4' --run-level "full_model"
|
||||
mirror_hardwares: l4_4
|
||||
|
||||
- label: ":full_moon: Diffusion · Distributed Test with H100"
|
||||
timeout_in_minutes: 60
|
||||
commands:
|
||||
- pytest -sv tests/diffusion/distributed/ -m 'full_model and cuda and H100' --run-level "full_model"
|
||||
mirror_hardwares: h100_4
|
||||
|
||||
- label: ":full_moon: Diffusion · Offloader Test"
|
||||
timeout_in_minutes: 40
|
||||
timeout_in_minutes: 60
|
||||
commands:
|
||||
- timeout 40m pytest -s -v tests/diffusion/offloader -m "full_model and cuda"
|
||||
- pytest -sv tests/diffusion/offloader -m "full_model and cuda"
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
- label: ":full_moon: Diffusion · Quantization Test with H100"
|
||||
timeout_in_minutes: 60
|
||||
commands:
|
||||
- timeout 60m pytest -sv tests/diffusion/quantization -m 'full_model and cuda and H100' --run-level "full_model"
|
||||
- pytest -sv tests/diffusion/quantization -m 'full_model and cuda and H100' --run-level "full_model"
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
- label: ":full_moon: Diffusion · Quantization Test with L4"
|
||||
timeout_in_minutes: 60
|
||||
commands:
|
||||
- pip install "vllm-gguf-plugin>=0.0.3"
|
||||
- pytest -sv tests/diffusion/quantization -m 'full_model and cuda and L4' --run-level "full_model"
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
- label: ":full_moon: Diffusion · Wan VAE Spatial Shard Correctness Test"
|
||||
timeout_in_minutes: 60
|
||||
commands:
|
||||
- >-
|
||||
pytest -sv tests/diffusion/distributed/test_wan_spatial_shard.py
|
||||
-m "full_model and diffusion and H100 and distributed_cuda"
|
||||
--run-level "full_model"
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: ":bar_chart: Testcase Statistics"
|
||||
key: nightly-testcase-statistics
|
||||
timeout_in_minutes: 120
|
||||
|
||||
@@ -8,22 +8,26 @@ steps:
|
||||
- group: ":card_index_dividers: Simple Test"
|
||||
depends_on: upload-ready-pipeline
|
||||
steps:
|
||||
- label: "Simple · Diffusion & Model Executor Test"
|
||||
- label: "Simple · Model Executor Test"
|
||||
commands:
|
||||
- "timeout 20m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||
- "timeout 20m pytest -sv tests/model_executor -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
- label: "Simple · Diffusion Test"
|
||||
commands:
|
||||
- "timeout 20m pytest -sv tests/diffusion -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
- label: "Simple · Engine&Entrypoints Test"
|
||||
commands:
|
||||
- "timeout 20m pytest -sv tests/entrypoints tests/engine -m 'core_model and cpu' --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
- label: "Simple · Other Test"
|
||||
commands:
|
||||
- "timeout 20m pytest -sv -m 'core_model and cpu' --ignore=tests/diffusion --ignore=tests/model_executor --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||
- "timeout 20m pytest -sv tests/ -m 'core_model and cpu' --ignore=tests/diffusion --ignore=tests/model_executor --ignore=tests/entrypoints --ignore=tests/engine --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
- label: "Custom Pipeline Test"
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
- timeout 20m pytest -s -v tests/e2e/offline_inference/custom_pipeline/ -m "core_model"
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
- group: ":card_index_dividers: Diffusion Test"
|
||||
depends_on: upload-ready-pipeline
|
||||
steps:
|
||||
@@ -104,11 +108,34 @@ steps:
|
||||
- group: ":card_index_dividers: E2E Test"
|
||||
depends_on: upload-ready-pipeline
|
||||
steps:
|
||||
- label: "Custom Pipeline Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/features/custom_pipeline/
|
||||
- tests/e2e/features/helpers/custom_pipeline.py
|
||||
- vllm_omni/diffusion/
|
||||
- vllm_omni/engine/
|
||||
- vllm_omni/entrypoints/
|
||||
- vllm_omni/inputs/
|
||||
- vllm_omni/outputs/
|
||||
commands:
|
||||
- timeout 20m pytest -s -v tests/e2e/features/custom_pipeline/ -m "core_model"
|
||||
agents:
|
||||
queue: "gpu_1_queue"
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "RLHF · VeRL-Omni E2E Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/offline_inference/rlhf_test/test_verl_omni_e2e.py
|
||||
- tests/e2e/offline_inference/custom_pipeline/
|
||||
- tests/e2e/features/rlhf_test/test_verl_omni_e2e.py
|
||||
- tests/e2e/features/helpers/verl_omni_server.py
|
||||
- tests/e2e/features/helpers/custom_pipeline.py
|
||||
- vllm_omni/diffusion/
|
||||
- vllm_omni/engine/
|
||||
- vllm_omni/entrypoints/
|
||||
@@ -117,7 +144,7 @@ steps:
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c "
|
||||
pytest -s -v tests/e2e/offline_inference/rlhf_test/test_verl_omni_e2e.py
|
||||
pytest -s -v tests/e2e/features/rlhf_test/test_verl_omni_e2e.py
|
||||
"
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
@@ -143,34 +170,7 @@ steps:
|
||||
timeout 20m bash -c "
|
||||
pytest -s -v tests/e2e/offline_inference/test_minicpmo_4_5.py -m 'core_model' --run-level 'core_model'
|
||||
"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 2
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: "Omni · MiniCPM-o 4.5 Online Test"
|
||||
source_file_dependencies:
|
||||
@@ -181,34 +181,34 @@ steps:
|
||||
timeout 20m bash -c "
|
||||
pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5.py -m 'core_model' --run-level 'core_model'
|
||||
"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 2
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: "Omni · MiniCPM-o 4.5 Duplex Test"
|
||||
source_file_dependencies:
|
||||
- tests/assets/minicpmo_4_5/response_required_16k.wav
|
||||
- tests/assets/minicpmo_4_5/soft_interrupt_16k.wav
|
||||
- tests/e2e/online_serving/test_minicpmo_4_5_duplex.py
|
||||
- tests/e2e/online_serving/test_minicpmo_4_5_duplex_expansion.py
|
||||
- tests/e2e/online_serving/test_minicpmo_realtime_duplex_drivers.py
|
||||
- tests/e2e/online_serving/minicpmo_realtime_duplex_scenarios.py
|
||||
- tests/e2e/online_serving/run_minicpmo_realtime_duplex_multi_session.py
|
||||
- tests/e2e/online_serving/run_minicpmo_realtime_duplex_soft_interrupt.py
|
||||
- tests/helpers/minicpmo_4_5_duplex.py
|
||||
- vllm_omni/deploy/minicpmo_4_5_duplex.yaml
|
||||
- vllm_omni/experimental/fullduplex/
|
||||
- vllm_omni/model_executor/models/minicpmo_4_5/
|
||||
- vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py
|
||||
- vllm_omni/core/sched/
|
||||
- vllm_omni/outputs/multimodal_accumulation.py
|
||||
- vllm_omni/engine/stage_init_utils.py
|
||||
- vllm_omni/engine/async_omni_engine.py
|
||||
- vllm_omni/entrypoints/cli/serve.py
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c "
|
||||
pytest -s -v tests/e2e/online_serving/test_minicpmo_4_5_duplex.py -m 'core_model and cuda' --run-level 'core_model'
|
||||
"
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: "TTS · VoxCPM2 Test"
|
||||
source_file_dependencies:
|
||||
@@ -344,13 +344,3 @@ steps:
|
||||
pytest -s -v tests/e2e/online_serving/test_qwen_image_layered.py -m 'core_model' --run-level 'core_model'
|
||||
"
|
||||
mirror_hardwares: h100_1
|
||||
|
||||
- label: "Diffusion · GGUF Plugin Test"
|
||||
soft_fail: true
|
||||
source_file_dependencies:
|
||||
- tests/plugin/test_gguf_diffusion.py
|
||||
- vllm_omni/diffusion/model_loader
|
||||
commands:
|
||||
- pip install "vllm-gguf-plugin>=0.0.3"
|
||||
- timeout 30m pytest -sv tests/plugin/test_gguf_diffusion.py
|
||||
mirror_hardwares: l4_1
|
||||
@@ -10,7 +10,7 @@ steps:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
# Buildkite will automatically replace this with the actual commit hash
|
||||
VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
|
||||
VLLM_VERSION: "v0.25.0"
|
||||
VLLM_VERSION: "v0.26.0"
|
||||
priority: 100
|
||||
timeout_in_minutes: 60
|
||||
soft_fail: false
|
||||
@@ -1,25 +1,59 @@
|
||||
steps:
|
||||
- group: ":card_index_dividers: Diffusion X2V Model Test"
|
||||
key: nightly-diffusion-x2v-group
|
||||
- group: ":card_index_dividers: Omni Model Test"
|
||||
key: nightly-omni-test-group
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: >-
|
||||
build.env("NIGHTLY") == "1" ||
|
||||
build.pull_request.labels includes "nightly-test" ||
|
||||
build.pull_request.labels includes "diffusion-x2v-test"
|
||||
build.pull_request.labels includes "omni-test"
|
||||
steps:
|
||||
- label: ":full_moon: Diffusion X2V · Function Test"
|
||||
mirror_hardwares: a2b3_npu_4
|
||||
- label: ":full_moon: Omni · Perf Test · No Async Chunk"
|
||||
key: nightly-omni-performance-no-async-chunk
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: a3_npu_4
|
||||
env:
|
||||
VLLM_WORKER_MULTIPROC_METHOD: spawn
|
||||
VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
|
||||
HF_TOKEN: "${HF_TOKEN}"
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k npu_i2v --run-level "full_model"
|
||||
- export BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py -m npu --test-config-file tests/dfx/perf/tests/test_qwen3_omni_no_async_chunk.json
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
|
||||
exit $$EXIT
|
||||
- label: ":full_moon: Omni · Perf Test · Async Chunk"
|
||||
key: nightly-omni-performance-async-chunk
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: a3_npu_4
|
||||
env:
|
||||
VLLM_WORKER_MULTIPROC_METHOD: spawn
|
||||
VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
|
||||
HF_TOKEN: "${HF_TOKEN}"
|
||||
commands:
|
||||
- export BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py -m npu --test-config-file tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
|
||||
exit $$EXIT
|
||||
- label: ":robot: Omni · MiniCPM-o 4.5 Test"
|
||||
mirror_hardwares: a3_npu_2
|
||||
env:
|
||||
VLLM_WORKER_MULTIPROC_METHOD: spawn
|
||||
VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
|
||||
HF_TOKEN: "${HF_TOKEN}"
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/offline_inference/test_minicpmo_4_5.py tests/e2e/online_serving/test_minicpmo_4_5.py tests/e2e/online_serving/test_minicpmo_4_5_expansion.py -m 'full_model' --run-level 'full_model'
|
||||
- group: ":card_index_dividers: TTS Model Test"
|
||||
key: nightly-tts-test-group
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "tts-test"
|
||||
steps:
|
||||
- label: ":full_moon: TTS · Perf Test"
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: a3_npu_2
|
||||
env:
|
||||
VLLM_WORKER_MULTIPROC_METHOD: spawn
|
||||
@@ -34,25 +68,138 @@ steps:
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
|
||||
exit $$EXIT
|
||||
|
||||
- group: ":card_index_dividers: Omni Model Test"
|
||||
key: nightly-omni-group
|
||||
- group: ":card_index_dividers: Diffusion X2I(&A&T) Model Test"
|
||||
key: nightly-diffusion-x2iat-group
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "omni-test"
|
||||
if: >-
|
||||
build.env("NIGHTLY") == "1" ||
|
||||
build.pull_request.labels includes "nightly-test" ||
|
||||
build.pull_request.labels includes "diffusion-x2iat-test"
|
||||
steps:
|
||||
- label: ":robot: Omni · MiniCPM-o 4.5 Test"
|
||||
agents:
|
||||
queue: "ascend-a3"
|
||||
resource_class: "npu-2"
|
||||
image: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/vllm-omni-ci-npu:${BUILDKITE_COMMIT}-a3"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpecPatch:
|
||||
imagePullSecrets:
|
||||
- name: swr-secret
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Edit-2511"
|
||||
key: nightly-diffusion-x2iat-performance-qwen-image-edit-2511
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: a3_npu_8
|
||||
env:
|
||||
VLLM_WORKER_MULTIPROC_METHOD: spawn
|
||||
VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
|
||||
HF_TOKEN: "${HF_TOKEN}"
|
||||
DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA
|
||||
commands:
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export CACHE_DIT_VERSION=1.3.0
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m npu --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
|
||||
exit $$EXIT
|
||||
- group: ":card_index_dividers: Diffusion I2I Model Test"
|
||||
key: nightly-HunyuanImage3-x2i-test-group
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: >-
|
||||
build.env("NIGHTLY") == "1" ||
|
||||
build.pull_request.labels includes "nightly-test" ||
|
||||
build.pull_request.labels includes "diffusion-x2iat-test"
|
||||
steps:
|
||||
- label: ":full_moon: HunyuanImage3-DIT · Accuracy Test"
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: a3_npu_4
|
||||
env:
|
||||
VLLM_WORKER_MULTIPROC_METHOD: spawn
|
||||
HUNYUAN_IMAGE3_MODEL: "tencent/HunyuanImage-3.0-Instruct"
|
||||
HUNYUAN_IMAGE3_DEVICES: "0,1,2,3"
|
||||
DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/accuracy/test_hunyuan_image3_pixel_accuracy.py -m "full_model and A3" --run-level full_model
|
||||
- group: ":card_index_dividers: Diffusion Benchmark Perf Test"
|
||||
key: nightly-diffusion-benchmark-group
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: >-
|
||||
build.env("NIGHTLY") == "1" ||
|
||||
build.pull_request.labels includes "nightly-test" ||
|
||||
build.pull_request.labels includes "diffusion-x2iat-test"
|
||||
steps:
|
||||
- label: ":full_moon: Diffusion Perf · HunyuanImage3 TP=2 CFGP=2"
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: a3_npu_4
|
||||
env:
|
||||
VLLM_WORKER_MULTIPROC_METHOD: spawn
|
||||
HF_TOKEN: "${HF_TOKEN}"
|
||||
DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA
|
||||
commands:
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m A3 --test-config-file tests/dfx/perf/tests/test_hunyuan_image_tp2_cfgp2.json
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
|
||||
exit $$EXIT
|
||||
- label: ":full_moon: Diffusion Perf · HunyuanImage3 TP=2 SP=2"
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: a3_npu_4
|
||||
env:
|
||||
VLLM_WORKER_MULTIPROC_METHOD: spawn
|
||||
HF_TOKEN: "${HF_TOKEN}"
|
||||
DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA
|
||||
commands:
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m A3 --test-config-file tests/dfx/perf/tests/test_hunyuan_image_tp2_sp2.json
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
|
||||
exit $$EXIT
|
||||
- label: ":full_moon: Diffusion Perf · HunyuanImage3 TP=4"
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: a3_npu_4
|
||||
env:
|
||||
VLLM_WORKER_MULTIPROC_METHOD: spawn
|
||||
HF_TOKEN: "${HF_TOKEN}"
|
||||
DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA
|
||||
commands:
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m A3 --test-config-file tests/dfx/perf/tests/test_hunyuan_image_tp4.json
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
|
||||
exit $$EXIT
|
||||
- group: ":card_index_dividers: Diffusion X2V Model Test"
|
||||
key: nightly-diffusion-x2v-group
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: >-
|
||||
build.env("NIGHTLY") == "1" ||
|
||||
build.pull_request.labels includes "nightly-test" ||
|
||||
build.pull_request.labels includes "diffusion-x2v-test"
|
||||
steps:
|
||||
- label: ":full_moon: Diffusion X2V · Function Test"
|
||||
mirror_hardwares: a2b3_npu_4
|
||||
env:
|
||||
VLLM_WORKER_MULTIPROC_METHOD: spawn
|
||||
VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
|
||||
HF_TOKEN: "${HF_TOKEN}"
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/offline_inference/test_minicpmo_4_5.py tests/e2e/online_serving/test_minicpmo_4_5.py tests/e2e/online_serving/test_minicpmo_4_5_expansion.py -m 'full_model' --run-level 'full_model'
|
||||
- pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k npu_i2v --run-level "full_model"
|
||||
- label: ":full_moon: Diffusion X2V · Perf Test"
|
||||
key: nightly-diffusion-x2v-performance
|
||||
timeout_in_minutes: 180
|
||||
mirror_hardwares: a3_npu_4
|
||||
env:
|
||||
VLLM_WORKER_MULTIPROC_METHOD: spawn
|
||||
VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
|
||||
HF_TOKEN: "${HF_TOKEN}"
|
||||
DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA
|
||||
commands:
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m npu --test-config-file tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
|
||||
exit $$EXIT
|
||||
@@ -31,4 +31,4 @@ steps:
|
||||
VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
|
||||
HF_TOKEN: "${HF_TOKEN}"
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m 'core_model' --run-level 'core_model'
|
||||
- pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m 'core_model and npu' --run-level 'core_model'
|
||||
@@ -1,12 +1,8 @@
|
||||
from collections.abc import Callable
|
||||
from enum import Enum, auto
|
||||
from typing import (
|
||||
Any,
|
||||
Literal,
|
||||
NotRequired,
|
||||
TypeAlias,
|
||||
TypedDict,
|
||||
)
|
||||
from typing import Any, Literal, TypeAlias
|
||||
|
||||
from typing_extensions import NotRequired, TypedDict
|
||||
|
||||
AudioFormat: TypeAlias = Literal["mp3", "opus", "aac", "flac", "wav", "pcm"]
|
||||
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
# Target vllm version: v0.25.0 (releases/v0.25.0 branch).
|
||||
# Target vllm version: v0.26.0 (releases/v0.26.0 branch).
|
||||
# Keep VLLM_BASE_TAG aligned with the target vLLM release on every rebase.
|
||||
ARG VLLM_BASE_IMAGE=vllm/vllm-openai
|
||||
ARG VLLM_BASE_TAG=v0.25.0
|
||||
ARG VLLM_BASE_TAG=v0.26.0
|
||||
|
||||
# We use registry-based layer caching for the dependencies,
|
||||
# where the tag is determined based on the hash of the contents
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
ARG BASE_IMAGE=vllm/vllm-openai:v0.25.0
|
||||
ARG BASE_IMAGE=vllm/vllm-openai:v0.26.0
|
||||
FROM ${BASE_IMAGE}
|
||||
|
||||
ARG COMMON_WORKDIR=/app
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
|
||||
ARG VLLM_ASCEND_TAG=v0.25.0
|
||||
ARG VLLM_ASCEND_TAG=v0.26.0
|
||||
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
|
||||
|
||||
# WORKDIR /vllm-workspace/vllm
|
||||
@@ -26,6 +26,10 @@ RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi
|
||||
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
|
||||
python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation
|
||||
|
||||
RUN python3 -m pip install -v --no-deps \
|
||||
"stepaudio2-minicpmo>=0.1.1" \
|
||||
"step-audio2>=1.0.0"
|
||||
|
||||
ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
|
||||
ENTRYPOINT []
|
||||
@@ -1,5 +1,5 @@
|
||||
ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
|
||||
ARG VLLM_ASCEND_TAG=v0.25.0-a3
|
||||
ARG VLLM_ASCEND_TAG=v0.26.0-a3
|
||||
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
|
||||
|
||||
# WORKDIR /vllm-workspace/vllm
|
||||
@@ -26,6 +26,10 @@ RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi
|
||||
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
|
||||
python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation
|
||||
|
||||
RUN python3 -m pip install -v --no-deps \
|
||||
"stepaudio2-minicpmo>=0.1.1" \
|
||||
"step-audio2>=1.0.0"
|
||||
|
||||
ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
|
||||
ENTRYPOINT []
|
||||
@@ -1,5 +1,5 @@
|
||||
ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
|
||||
ARG VLLM_ASCEND_TAG=v0.25.0
|
||||
ARG VLLM_ASCEND_TAG=v0.26.0
|
||||
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
|
||||
|
||||
ARG APP_DIR=/vllm-workspace/vllm-omni
|
||||
@@ -35,8 +35,13 @@ RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
|
||||
"psutil>=7.2.0" \
|
||||
"soundfile>=0.13.1" \
|
||||
"pyttsx3>=2.99" \
|
||||
"opencc==1.4.1" \
|
||||
"opencv-python>=4.12.0.88" \
|
||||
"imageio[ffmpeg]>=0.6.0" \
|
||||
"jiwer>=3.0.0" \
|
||||
"zhon>=2.0.0" \
|
||||
"zhconv>=1.4.2" \
|
||||
"funasr>=1.0.0" \
|
||||
"hyperpyyaml>=1.2.0" \
|
||||
"onnx>=1.16.0" \
|
||||
"librosa>=0.10.0" \
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
|
||||
ARG VLLM_ASCEND_TAG=v0.25.0-a3
|
||||
ARG VLLM_ASCEND_TAG=v0.26.0-a3
|
||||
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
|
||||
|
||||
ARG APP_DIR=/vllm-workspace/vllm-omni
|
||||
@@ -35,11 +35,17 @@ RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
|
||||
"psutil>=7.2.0" \
|
||||
"soundfile>=0.13.1" \
|
||||
"pyttsx3>=2.99" \
|
||||
"opencc==1.4.1" \
|
||||
"opencv-python>=4.12.0.88" \
|
||||
"imageio[ffmpeg]>=0.6.0" \
|
||||
"jiwer>=3.0.0" \
|
||||
"zhon>=2.0.0" \
|
||||
"zhconv>=1.4.2" \
|
||||
"funasr>=1.0.0" \
|
||||
"hyperpyyaml>=1.2.0" \
|
||||
"onnx>=1.16.0" \
|
||||
"librosa>=0.10.0" \
|
||||
"torchmetrics>=1.9.0" \
|
||||
--index-url http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple \
|
||||
--trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \
|
||||
--default-timeout=60 && \
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.25.0
|
||||
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.26.0
|
||||
FROM ${BASE_IMAGE} AS base
|
||||
|
||||
# Declare a variable to know if we want to use the nightly build or the stable build.
|
||||
@@ -10,7 +10,7 @@ FROM ${BASE_IMAGE} AS base
|
||||
# we should swap over to use stable release ASAP.
|
||||
# We should avoid relying on custom commits.
|
||||
ARG USE_NIGHTLY_BUILD=0
|
||||
ARG VLLM_VERSION_OR_COMMIT_HASH=89138b21cc246ae944c741d5c399c148e2b770ab
|
||||
ARG VLLM_VERSION_OR_COMMIT_HASH=v0.26.0
|
||||
ARG ARG_PYTORCH_ROCM_ARCH
|
||||
ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}}
|
||||
|
||||
|
||||
@@ -88,7 +88,7 @@ ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
|
||||
ENV UV_INDEX_STRATEGY="unsafe-best-match"
|
||||
ENV UV_LINK_MODE="copy"
|
||||
|
||||
ARG VLLM_VERSION=v0.25.0
|
||||
ARG VLLM_VERSION=v0.26.0
|
||||
RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm
|
||||
WORKDIR /workspace/vllm
|
||||
|
||||
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 281 KiB After Width: | Height: | Size: 274 KiB |
+1
-7
@@ -18,18 +18,12 @@ Specify the port:
|
||||
vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091
|
||||
```
|
||||
|
||||
For a migrated model, load a custom deploy configuration with `--deploy-config`:
|
||||
Load a custom deploy configuration with `--deploy-config`:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen2.5-Omni-7B --omni --deploy-config /path/to/deploy_config.yaml
|
||||
```
|
||||
|
||||
The deprecated `--stage-configs-path` flag is retained for models that still use the legacy `stage_args` schema:
|
||||
|
||||
```bash
|
||||
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --stage-configs-path /path/to/legacy_stage_config.yaml
|
||||
```
|
||||
|
||||
## bench
|
||||
|
||||
Run benchmark tests for online serving throughput.
|
||||
|
||||
+1
-2
@@ -8,7 +8,6 @@ The stage-based CLI is designed for deployments that require launching each pipe
|
||||
- For **migrated models** that utilize the bundled deployment YAML configurations located in
|
||||
`vllm_omni/deploy/`, the `--deploy-config` flag is only required to override the default configuration. By default, executing `vllm serve MODEL --omni ...`
|
||||
automatically loads the bundled deployment configuration.
|
||||
- For custom legacy `stage_args` YAMLs, pass the file with `--stage-configs-path`.
|
||||
|
||||
Example: Initializing Stage 0 (Orchestrator and API Server):
|
||||
The commands below show a common device mapping where Stage 0 uses GPU 0 and
|
||||
@@ -32,7 +31,7 @@ CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
--omni-master-port 26000
|
||||
```
|
||||
|
||||
When utilizing a custom deployment YAML based on the new schema, append `--deploy-config /path/to/override.yaml` to each command execution. Conversely, for legacy models, substitute this parameter with `--stage-configs-path /path/to/stage_configs.yaml`.
|
||||
When utilizing a custom deployment YAML, append `--deploy-config /path/to/override.yaml` to each command execution.
|
||||
|
||||
In the standard execution paradigm, the `--stage-overrides` argument is utilized to apply stage-specific configurations from a single CLI command.
|
||||
However, under the **stage-based CLI** paradigm, where each process strictly encapsulates a single stage, it is recommended to specify tuning parameters directly via discrete command-line flags for the respective stage, rather than constructing a composite `--stage-overrides` JSON string.
|
||||
|
||||
@@ -34,12 +34,8 @@ Two additional `StrategySpec` slots — `layer_hook_specs` and `kernel_specs`
|
||||
| Flag | Description |
|
||||
|------|-------------|
|
||||
| `--strategy-config PATH` | Path to a `strategy.yaml`. Only takes effect on registry-based deploy paths (`--deploy-config` or the bundled default deploy YAML). Its derived sizing is overlaid onto the merged stages *before* CLI overrides. |
|
||||
| `--stage-configs-path PATH` | **Legacy.** Loads the old `stage_args` YAML schema for un-migrated models. Passing `--strategy-config` together with this path is silently inapplicable — vLLM-Omni emits a warning and the strategy is ignored (see `vllm_omni/entrypoints/utils.py`, lines 372–380). |
|
||||
| `--omni-lb-policy POLICY` | Orchestrator-wide load-balancer policy. Strategy-derived from any `stage_replica` axis; an explicit CLI value wins and must match the derived one or `AsyncOmniEngine` raises `Conflicting load-balancer policy` at construction. |
|
||||
|
||||
!!! warning
|
||||
`--strategy-config` cannot be used with `--stage-configs-path` (the legacy deploy path). If your model has not been migrated to the registry-based pipeline yet, the strategy is silently dropped with a warning. Migrate to a registry-based model or use `--deploy-config` to apply a strategy.
|
||||
|
||||
## Sub-schemas
|
||||
|
||||
### `axis`
|
||||
@@ -216,17 +212,6 @@ Strategy keys MUST be `model_stage` *names* (strings), the same labels stage con
|
||||
!!! warning "Mixing CLI overrides with a strategy on the same axis"
|
||||
A strategy is meant to be the single writer for the axes it declares. If a CLI override (`--tensor-parallel-size 2`, `--stage-overrides '{"0": {"tensor_parallel_size": 2}}'`) sets the same axis, the **CLI value wins** and a warning is emitted from `_reconcile_strategy_with_cli`. The device-layout guard then re-runs against the effective layout, so a CLI override that breaks `tp * dp * pp * num_replicas` will still fail fast at config time.
|
||||
|
||||
!!! warning "Passing `--strategy-config` against the legacy `--stage-configs-path`"
|
||||
Composable parallel only applies on the registry-based deploy path. If a model still resolves through the legacy `stage_args` YAML (loaded via `--stage-configs-path`), `--strategy-config` is silently inapplicable and emits:
|
||||
|
||||
```text
|
||||
--strategy-config (<path>) was provided but model '<model>' resolves via the
|
||||
legacy stage_configs YAML path, which does not support composable-parallel
|
||||
strategies; the strategy is ignored. Use a registry-based model to apply it.
|
||||
```
|
||||
|
||||
Migrate the model to the new deploy schema (see `configuration/stage_configs.md`) before applying a strategy.
|
||||
|
||||
!!! warning "Conflicting `omni_lb_policy` across stages"
|
||||
`omni_lb_policy` is a single pipeline-wide knob. If two `stage_replica` axes in different stages derive different policies, `apply_strategy_specs` raises:
|
||||
|
||||
|
||||
@@ -158,7 +158,7 @@ runtime:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
|
||||
--stage-configs-path /path/to/qwen3_omni_pd.yaml
|
||||
--deploy-config /path/to/qwen3_omni_pd.yaml
|
||||
```
|
||||
|
||||
## Operational Notes
|
||||
|
||||
@@ -1,9 +1,9 @@
|
||||
# Deploy and legacy stage configurations
|
||||
# Pipeline and deploy configurations
|
||||
|
||||
In vLLM-Omni, a model's `PipelineConfig` defines its fixed stage topology, while a deploy configuration controls how those stages run. Models that have not migrated to this split still use legacy stage configurations with a `stage_args` schema.
|
||||
In vLLM-Omni, a model's `PipelineConfig` defines its fixed stage topology, while a deploy configuration controls how those stages run.
|
||||
|
||||
!!! note
|
||||
Default deploy config YAMLs (for example, `vllm_omni/deploy/qwen2_5_omni.yaml`, `vllm_omni/deploy/qwen3_omni_moe.yaml`, and `vllm_omni/deploy/qwen3_tts.yaml`) are bundled and loaded automatically when neither `--stage-configs-path` nor `--deploy-config` is provided — the model registry resolves the right pipeline + deploy YAML by `model_type`. Custom legacy `stage_args` YAMLs are still accepted via `--stage-configs-path`.
|
||||
Default deploy config YAMLs (for example, `vllm_omni/deploy/qwen2_5_omni.yaml`, `vllm_omni/deploy/qwen3_omni_moe.yaml`, and `vllm_omni/deploy/qwen3_tts.yaml`) are bundled and loaded automatically when `--deploy-config` is omitted. The model registry resolves the right pipeline and deploy YAML by `model_type`.
|
||||
|
||||
## Deploy configuration schema
|
||||
|
||||
@@ -86,10 +86,9 @@ stages:
|
||||
|
||||
| Flag | Description |
|
||||
|------|-------------|
|
||||
| `--deploy-config PATH` | Load a new-schema deploy YAML. It is mutually exclusive with `--stage-configs-path`. **Optional** — when omitted, the bundled `vllm_omni/deploy/<model_type>.yaml` is auto-loaded by the model registry. |
|
||||
| `--deploy-config PATH` | Load a deploy YAML. **Optional** — when omitted, the bundled `vllm_omni/deploy/<model_type>.yaml` is auto-loaded by the model registry. |
|
||||
| `--stage-overrides JSON` | Per-stage JSON overrides, e.g. `'{"0":{"gpu_memory_utilization":0.5}}'`. Per-stage values always win over global flags. |
|
||||
| `--async-chunk` / `--no-async-chunk` | Flip the deploy YAML's `async_chunk:` bool. Unset (default) leaves the YAML value in force. |
|
||||
| `--stage-configs-path` | **Deprecated.** Load a legacy `stage_args` YAML. New-schema YAMLs passed to this flag are auto-detected for compatibility, but should use `--deploy-config`. It is mutually exclusive with `--deploy-config` and will be removed in a future release. |
|
||||
|
||||
### Stage-Based CLI Paradigm
|
||||
|
||||
@@ -119,8 +118,7 @@ CUDA_VISIBLE_DEVICES=1 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
|
||||
--omni-master-port 26000
|
||||
```
|
||||
|
||||
When instantiating a custom deployment YAML conforming to the updated schema, append the `--deploy-config /path/to/override.yaml` directive
|
||||
to all node invocations. For legacy architectures (e.g., BAGEL) configured via deprecated `stage_args:` schemas, continue to specify the relevant configuration via `--stage-configs-path /path/to/config.yaml`.
|
||||
When instantiating a custom deployment YAML, append the `--deploy-config /path/to/override.yaml` directive to all node invocations.
|
||||
|
||||
In the context of standard initialization architectures, utilizing the `--stage-overrides` parameter operates as the optimal methodology
|
||||
for delineating stage-specific tuning from the CLI interface:
|
||||
@@ -220,8 +218,7 @@ Therefore, as a core part of vLLM-Omni, a model's pipeline and deployment config
|
||||
|
||||
To override specific parameters, explicitly inject the customized configuration schema
|
||||
in both online and offline instantiation flows. Use the `--deploy-config` flag
|
||||
when loading a deploy configuration, reserving the `--stage-configs-path` parameter
|
||||
exclusively to maintain compatibility with legacy `stage_args` YAML constructs.
|
||||
when loading a deploy configuration.
|
||||
|
||||
Examples:
|
||||
|
||||
@@ -236,11 +233,6 @@ For online serving:
|
||||
vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --deploy-config /path/to/deploy_config.yaml
|
||||
```
|
||||
|
||||
Legacy online serving:
|
||||
|
||||
```bash
|
||||
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
|
||||
```
|
||||
!!! important
|
||||
We are actively iterating on the definition of deployment configurations, and we welcome feedback from users and developers.
|
||||
|
||||
|
||||
@@ -56,7 +56,11 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
|
||||
<td>/</td>
|
||||
<td><code>core_model and cpu</code></td>
|
||||
<td rowspan="2"><15min</td>
|
||||
<td>/tests/{component_name}/test_xxx</td>
|
||||
<td>
|
||||
<code>tests/{component}/…</code> mirroring <code>vllm_omni/{component}/</code><br>
|
||||
(e.g. <code>tests/diffusion/</code>, <code>tests/engine/</code>, <code>tests/entrypoints/</code>)<br>
|
||||
Do <strong>not</strong> add new top-level <code>tests/</code> dirs unrelated to a component.
|
||||
</td>
|
||||
<td>
|
||||
<a href="test_writing_guide.md#l1--l2-level-testing-unit-testing-and-basic-end-to-end-verification">L1 & L2</a><br>
|
||||
Section 1 L1&L2: Purpose, Test Content, Directory Location, Example
|
||||
@@ -68,16 +72,18 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
|
||||
<td><strong>L2</strong><br>(E2E across models & GPU-required UT)</td>
|
||||
<td>Online (basic deployment scenarios):<br>dummy, normal inference function (output format, stream), some instance startup UT</td>
|
||||
<td>High-priority models + online basic scenarios; request success, non-empty output, format match (no Whisper/accuracy)</td>
|
||||
<td>High-priority features (using random lightweight models)</td>
|
||||
<td>High-priority features and component tests that require launching instances</td>
|
||||
<td>High-priority interfaces (using random lightweight models)</td>
|
||||
<td><code>core_model and hardware_test(H100, L4, etc.) and omni/tts/diffusion</code></td>
|
||||
<td>
|
||||
<strong>Model tests:</strong><br>
|
||||
/tests/e2e/online_serving/test_{model_name}.py<br>
|
||||
<strong>Feature tests:</strong><br>
|
||||
/tests/{component_name}/test_xxx<br>
|
||||
<strong>Interface tests:</strong><br>
|
||||
/tests/entrypoints/test_xxx
|
||||
<strong>Model E2E:</strong><br>
|
||||
<code>tests/e2e/online_serving/test_{model}.py</code><br>
|
||||
<code>tests/e2e/offline_inference/test_{model}.py</code><br>
|
||||
<strong>Feature integration:</strong><br>
|
||||
<code>tests/e2e/features/<feature>/</code><br>
|
||||
(e.g. <code>fullduplex/</code>, <code>custom_pipeline/</code>, <code>rlhf_test/</code>)<br>
|
||||
<strong>Component / interface Test (GPU):</strong><br>
|
||||
<code>tests/{component}/…</code>, <code>tests/entrypoints/…</code>
|
||||
</td>
|
||||
<td>
|
||||
<a href="test_writing_guide.md#l1--l2-level-testing-unit-testing-and-basic-end-to-end-verification">L1 & L2</a><br>
|
||||
@@ -90,18 +96,19 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
|
||||
<td><strong>L3</strong><br>(Important Perf & Integration & Accuracy)</td>
|
||||
<td>Online & Offline (multiple deployment scenarios):<br>real model, normal inference function, normal accuracy</td>
|
||||
<td>High/medium-priority models + key online/offline scenarios; real weights, Whisper/similarity, preset voice gender, basic accuracy</td>
|
||||
<td>Medium-priority features (using random lightweight models)</td>
|
||||
<td>Medium-priority features and component tests that require launching instances</td>
|
||||
<td>Medium-priority interfaces (using random lightweight models)</td>
|
||||
<td><code>advanced_model and hardware_test(H100, L4, etc.) and omni/tts/diffusion</code></td>
|
||||
<td><30min</td>
|
||||
<td>
|
||||
<strong>Model tests:</strong><br>
|
||||
/tests/e2e/online_serving/test_{model_name}.py<br>
|
||||
/tests/e2e/offline_inference/test_{model_name}.py<br>
|
||||
<strong>Feature tests:</strong><br>
|
||||
/tests/{component_name}/test_xxx<br>
|
||||
<strong>Interface tests:</strong><br>
|
||||
/tests/entrypoints/test_xxx
|
||||
<strong>Model E2E:</strong><br>
|
||||
<code>tests/e2e/online_serving/test_{model}.py</code><br>
|
||||
<code>tests/e2e/offline_inference/test_{model}.py</code><br>
|
||||
<code>tests/e2e/accuracy/</code><br>
|
||||
<strong>Feature integration:</strong><br>
|
||||
<code>tests/e2e/features/<feature>/</code><br>
|
||||
<strong>Component / interface Test:</strong><br>
|
||||
<code>tests/{component}/…</code>, <code>tests/entrypoints/…</code>
|
||||
</td>
|
||||
<td>
|
||||
<a href="test_writing_guide.md#l3-level-testing-core-integration-performance-and-accuracy-verification">L3</a><br>
|
||||
@@ -114,26 +121,23 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
|
||||
<td><strong>L4</strong><br>(Perf & Integration & Accuracy)</td>
|
||||
<td>Online: full functional scenarios + performance test + doc test + accuracy test</td>
|
||||
<td>High-priority models: function, performance, accuracy, and doc testing<br>Medium-priority models: function and doc testing</td>
|
||||
<td>Low-priority features (using real weights)</td>
|
||||
<td>Low-priority features and component tests that require launching instances</td>
|
||||
<td>Low-priority interfaces (using real weights)</td>
|
||||
<td><code>full_model and hardware_test(H100, L4, etc.) and omni/tts/diffusion</code></td>
|
||||
<td><3 hour</td>
|
||||
<td>
|
||||
<strong>Model tests:</strong><br>
|
||||
/tests/e2e/online_serving/test_{model_name}_expansion.py<br>
|
||||
<strong>Feature tests:</strong><br>
|
||||
/tests/{component_name}/test_xxx<br>
|
||||
<strong>Interface tests:</strong><br>
|
||||
/tests/entrypoints/test_xxx<br>
|
||||
<strong>Model E2E:</strong><br>
|
||||
<code>tests/e2e/online_serving/test_{model}_expansion.py</code><br>
|
||||
<code>tests/e2e/offline_inference/test_{model}_expansion.py</code><br>
|
||||
<code>tests/e2e/accuracy/test_{model}.py</code><br>
|
||||
<strong>Feature integration:</strong><br>
|
||||
<code>tests/e2e/features/<feature>/</code><br>
|
||||
<strong>Component / interface Test:</strong><br>
|
||||
<code>tests/{component}/…</code>, <code>tests/entrypoints/…</code><br>
|
||||
<strong>Performance:</strong><br>
|
||||
/tests/dfx/perf/tests/test_qwen3_omni_*.json (Omni), test_tts.json (TTS),<br>
|
||||
test_voxcpm2.json, test_higgs_audio_v3.json, and<br>
|
||||
/tests/dfx/perf/tests/test_{diffusion_model}_vllm_omni.json (Diffusion)<br>
|
||||
<strong>Doc Test:</strong><br>
|
||||
tests/examples/online_serving/test_{model_name}.py<br>
|
||||
tests/examples/offline_inference/test_{model_name}.py<br>
|
||||
<strong>Accuracy Test:</strong><br>
|
||||
/tests/e2e/accuracy/test_{model_name}.py
|
||||
<code>tests/dfx/perf/tests/</code><br>
|
||||
<strong>Doc examples:</strong><br>
|
||||
<code>tests/examples/online_serving/</code>, <code>tests/examples/offline_inference/</code>
|
||||
</td>
|
||||
<td>
|
||||
<a href="test_writing_guide.md#l4-level-testing-full-functionality-performance-and-documentation-testing">L4</a><br>
|
||||
@@ -152,11 +156,9 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
|
||||
<td> Depends on reality </td>
|
||||
<td>
|
||||
<strong>Stability:</strong><br>
|
||||
/tests/dfx/stability/tests/test_qwen3_omni.json<br>
|
||||
/tests/dfx/stability/tests/test_wan22.json<br>
|
||||
<code>tests/dfx/stability/tests/</code><br>
|
||||
<strong>Reliability:</strong><br>
|
||||
tests/dfx/reliability/test_reliability_{model_key}.py<br>
|
||||
(e.g. <code>test_reliability_qwen3_omni.py</code>, <code>test_reliability_wan22.py</code>, <code>test_reliability_hunyuan_image.py</code>, <code>test_reliability_voxcpm2.py</code>)
|
||||
<code>tests/dfx/reliability/test_reliability_{model}.py</code>
|
||||
</td>
|
||||
<td>
|
||||
<a href="test_writing_guide.md#l5-level-testing-stability-and-reliability-testing">L5</a><br>
|
||||
@@ -168,7 +170,9 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
|
||||
</tbody>
|
||||
</table>
|
||||
|
||||
For per-level test authoring (directories, markers, examples), see [Test Writing Guide](./test_writing_guide.md).
|
||||
**Test Dir placement (summary):** component / unit under `tests/{component}/` mirroring `vllm_omni/`; model E2E under `tests/e2e/online_serving/`, `tests/e2e/offline_inference/`, and `tests/e2e/accuracy/`; feature integration under `tests/e2e/features/<feature>/`; doc example tests under `tests/examples/online_serving/` and `tests/examples/offline_inference/`; performance under `tests/dfx/perf/`; stability under `tests/dfx/stability/`; reliability under `tests/dfx/reliability/`. Do **not** add new top-level directories under `tests/` that are unrelated to a `vllm_omni` component (or to the established `e2e` / `dfx` / `helpers` / `examples` / `buildkite` layout).
|
||||
|
||||
For per-level test authoring (markers, examples), see [Test Writing Guide](./test_writing_guide.md).
|
||||
|
||||
## Common Specifications
|
||||
|
||||
|
||||
@@ -4,137 +4,20 @@ This document describes how to author tests for each CI level (L1–L5). For the
|
||||
|
||||
## Test Directory Structure
|
||||
|
||||
Legend: `✅` = test exists, `⬜` = suggested to add.
|
||||
```
|
||||
vllm_omni/ tests/
|
||||
├── config/ → ├── config/
|
||||
│ ├── model.py │ └── test_model.py ⬜
|
||||
│ └── lora.py │ └── test_lora.py ⬜
|
||||
│
|
||||
├── core/ → ├── core/
|
||||
│ └── sched/ │ └── sched/
|
||||
│ ├── omni_ar_scheduler.py │ ├── test_omni_ar_scheduler.py ⬜
|
||||
│ ├── omni_generation_scheduler.py │ ├── test_omni_generation_scheduler.py ⬜
|
||||
│ └── output.py │ └── test_output.py ✅ currently in entrypoints/test_omni_new_request_data.py (tests output.OmniNewRequestData)
|
||||
│
|
||||
├── diffusion/ → ├── diffusion/
|
||||
│ ├── diffusion_engine.py │ ├── test_diffusion_engine.py ⬜
|
||||
│ ├── attention/ │ ├── attention/
|
||||
│ │ ├── layer.py │ │ ├── test_attention_sp.py ✅
|
||||
│ │ └── backends/ │ │ └── test_flash_attn.py ✅
|
||||
│ ├── distributed/ │ ├── distributed/
|
||||
│ │ └── ... │ │ ├── test_comm.py ✅
|
||||
│ │ │ │ ├── test_cfg_parallel.py ✅
|
||||
│ │ │ │ └── test_sp_plan_hooks.py ✅
|
||||
│ ├── lora/ │ ├── lora/
|
||||
│ │ └── ... │ │ ├── test_base_linear.py ✅
|
||||
│ │ │ │ └── test_lora_manager.py ✅
|
||||
│ ├── models/ │ ├── models/
|
||||
│ │ ├── qwen_image/ │ │ ├── qwen_image/ (e2e coverage)
|
||||
│ │ ├── ovis_image/ │ │ ├── ovis_image/
|
||||
│ │ │ └── ... │ │ │ └── test_ovis_image.py ✅
|
||||
│ │ ├── z_image/ │ │ └── z_image/
|
||||
│ │ └── ... │ │ └── test_zimage_tp_constraints.py ✅
|
||||
│ └── worker/ │ └── worker/
|
||||
│ ├── diffusion_worker.py │ └── test_diffusion_worker.py ✅ file at diffusion/test_diffusion_worker.py
|
||||
│ └── diffusion_model_runner.py │
|
||||
│
|
||||
├── distributed/ → ├── distributed/
|
||||
│ └── omni_connectors/ │ └── omni_connectors/
|
||||
│ ├── adapter.py │ ├── test_adapter_and_flow.py ✅
|
||||
│ ├── kv_transfer_manager.py │ ├── test_basic_connectors.py ✅
|
||||
│ ├── connectors/ │ ├── test_kv_flow.py ✅
|
||||
│ └── utils/ │ └── test_omni_connector_configs.py ✅
|
||||
│
|
||||
├── engine/ → ├── engine/
|
||||
│ ├── input_processor.py │ ├── test_input_processor.py ⬜ (no processor.py in source)
|
||||
│ ├── output_processor.py │ └── test_output_processor.py ⬜
|
||||
│ └── arg_utils.py │ └── test_arg_utils.py ⬜
|
||||
│
|
||||
├── entrypoints/ → ├── entrypoints/
|
||||
│ ├── stage_utils.py │ ├── test_stage_utils.py ✅
|
||||
│ ├── cli/ │ ├── cli/ (benchmarks/test_serve_cli.py covers CLI serve)
|
||||
│ │ └── ... │ │ └── test_*.py ⬜
|
||||
│ └── openai/ │ └── openai_api/ # maps to entrypoints/openai/
|
||||
│ ├── api_server.py │ ├── test_api_server.py ⬜ (e2e indirect coverage)
|
||||
│ ├── serving_chat.py │ ├── test_serving_chat_sampling_params.py ✅
|
||||
│ ├── serving_speech.py │ ├── test_serving_speech.py ✅
|
||||
│ └── image_api_utils.py │ └── test_image_server.py ✅
|
||||
│
|
||||
├── inputs/ → ├── inputs/
|
||||
│ ├── data.py │ ├── test_data.py ⬜
|
||||
│ ├── parse.py │ ├── test_parse.py ⬜
|
||||
│ └── preprocess.py │ └── test_preprocess.py ✅ currently in entrypoints/test_omni_input_preprocessor.py
|
||||
│
|
||||
├── model_executor/ → ├── model_executor/
|
||||
│ ├── layers/ │ ├── layers/
|
||||
│ │ └── mrope.py │ │ └── test_mrope.py ⬜
|
||||
│ ├── model_loader/ │ ├── model_loader/
|
||||
│ │ └── weight_utils.py │ │ └── test_weight_utils.py ⬜
|
||||
│ ├── models/ │ ├── models/
|
||||
│ │ ├── qwen2_5_omni/ │ │ ├── qwen2_5_omni/
|
||||
│ │ │ ├── qwen2_5_omni_thinker.py │ │ │ ├── test_audio_length.py ✅
|
||||
│ │ │ ├── qwen2_5_omni_talker.py │ │ │ ├── test_qwen2_5_omni_thinker.py ⬜
|
||||
│ │ │ └── qwen2_5_omni_token2wav.py │ │ │ ├── test_qwen2_5_omni_talker.py ⬜
|
||||
│ │ └── qwen3_omni/ │ │ │ └── test_qwen2_5_omni_token2wav.py ⬜
|
||||
│ │ └── ... │ │ └── qwen3_omni/
|
||||
│ ├── stage_configs/ │ │ └── test_*.py ⬜
|
||||
│ │ └── *.yaml │ └── stage_configs/ (used by e2e, test_*.py can be added) ⬜
|
||||
│ └── stage_input_processors/ │ └── stage_input_processors/
|
||||
│ └── ... │ └── test_*.py ⬜
|
||||
│
|
||||
├── sample/ → ├── sample/
|
||||
│ └── __init__.py │ └── test_*.py ⬜
|
||||
│
|
||||
├── utils/ → ├── utils/
|
||||
│ └── __init__.py │ └── test_*.py ⬜ (no platform_utils.py currently)
|
||||
│
|
||||
├── worker/ → ├── worker/
|
||||
│ ├── gpu_ar_model_runner.py │ ├── test_gpu_ar_model_runner.py ⬜
|
||||
│ ├── gpu_ar_worker.py │ ├── test_gpu_ar_worker.py ⬜
|
||||
│ ├── gpu_generation_model_runner.py │ ├── test_gpu_generation_model_runner.py ✅
|
||||
│ ├── gpu_generation_worker.py │ ├── test_gpu_generation_worker.py ⬜
|
||||
│ ├── gpu_model_runner.py │ ├── test_omni_gpu_model_runner.py ✅
|
||||
│ └── mixins.py │ └── (npu under platforms/npu/worker/) # not worker/npu/
|
||||
│
|
||||
├── platforms/ → (no tests/platforms/, e2e and stage_configs provide indirect coverage)
|
||||
│ ├── cuda/
|
||||
│ ├── npu/worker/ # NPU worker here, not vllm_omni/worker/npu/
|
||||
│ ├── rocm/
|
||||
│ └── xpu/worker/
|
||||
│
|
||||
├── outputs.py → test_outputs.py ✅ (at tests root)
|
||||
├── (logger, patch, request, version) → (no corresponding unit test)
|
||||
│
|
||||
└── e2e (tests side only) → ├── e2e/
|
||||
├── online_serving/ ✅ non-empty
|
||||
│ ├── test_qwen2_5_omni.py
|
||||
│ ├── test_async_omni.py
|
||||
│ ├── test_qwen3_omni.py
|
||||
│ ├── test_qwen3_omni_expansion.py
|
||||
│ ├── test_mimo_audio.py
|
||||
│ └── test_images_generations_lora.py
|
||||
└── offline_inference/ ✅
|
||||
├── test_qwen2_5_omni.py
|
||||
├── test_qwen3_omni.py
|
||||
├── test_bagel_text2img.py
|
||||
├── test_z_image.py
|
||||
├── test_wan22.py
|
||||
├── test_zimage_tensor_parallel.py
|
||||
├── test_cache_dit.py
|
||||
├── test_teacache.py
|
||||
├── test_stable_audio_expansion.py
|
||||
├── test_diffusion_cpu_offload.py
|
||||
├── test_diffusion_layerwise_offload.py
|
||||
├── test_diffusion_lora.py
|
||||
├── test_sequence_parallel.py
|
||||
└── stage_configs/ (legacy schema, still
|
||||
├── bagel_*.yaml present for unmigrated
|
||||
└── npu/, rocm/, etc. models)
|
||||
Canonical placement by CI level is in the **Test Dir** column of the level matrix in [Test System Overview](./test_system_overview.md).
|
||||
|
||||
# Migrated models (qwen3_omni_moe, qwen2_5_omni, qwen3_tts) live under
|
||||
# vllm_omni/deploy/ instead — see docs/configuration/stage_configs.md.
|
||||
```
|
||||
Summary:
|
||||
|
||||
- **Component / unit:** `tests/{component}/…` mirroring `vllm_omni/{component}/`
|
||||
- **Model E2E:** `tests/e2e/online_serving/`, `tests/e2e/offline_inference/`, `tests/e2e/accuracy/`
|
||||
- **Feature integration:** `tests/e2e/features/<feature>/`
|
||||
- **Doc example tests:** `tests/examples/online_serving/`, `tests/examples/offline_inference/`
|
||||
- **Performance:** `tests/dfx/perf/tests/` (JSON configs; runners under `tests/dfx/perf/scripts/`)
|
||||
- **Stability:** `tests/dfx/stability/tests/`
|
||||
- **Reliability:** `tests/dfx/reliability/`
|
||||
- Do **not** add new top-level `tests/` directories unrelated to a component (or to `e2e` / `dfx` / `helpers` / `examples` / `buildkite`)
|
||||
|
||||
Deploy YAMLs: `vllm_omni/deploy/` via `tests.helpers.stage_config.get_deploy_config_path` (see [stage configs](../../configuration/stage_configs.md)).
|
||||
|
||||
## Markers for Tests
|
||||
|
||||
@@ -373,13 +256,13 @@ L2 level testing builds upon L1 by introducing GPU resources and verifying that
|
||||
|
||||
#### 1.3 Test Directory and Execution Files
|
||||
|
||||
A clear directory structure is key to managing test cases efficiently.
|
||||
A clear directory structure is key to managing test cases efficiently. See the **Test Dir** column in [Test System Overview](./test_system_overview.md) for the canonical placement rules.
|
||||
|
||||
- ***L1 Test Directory***: `/tests/{component_name}/test_xxx.py`
|
||||
- - Here, `{component_name}` corresponds to modules in the source code, such as `distributed`, `entrypoints`, etc., and `test_xxx.py` is the specific test file.
|
||||
- ***L2 Test Directory***:
|
||||
- - Online Serving: `/tests/e2e/online_serving/test_{model_name}.py`
|
||||
- Offline Inference: `/tests/e2e/offline_inference/test_{model_name}.py`
|
||||
- ***L1 (component / unit)***: `tests/{component_name}/…` mirroring `vllm_omni/{component_name}/` (e.g. `tests/diffusion/`, `tests/entrypoints/`). Do **not** invent new top-level `tests/` folders unrelated to a component.
|
||||
- ***L2 (basic model E2E)***:
|
||||
- Online serving: `tests/e2e/online_serving/test_{model_name}.py`
|
||||
- Offline inference: `tests/e2e/offline_inference/test_{model_name}.py`
|
||||
- ***Feature-level integration*** (cross-cutting, not a single-model smoke): `tests/e2e/features/<feature>/` (e.g. `fullduplex/`, `custom_pipeline/`, `rlhf_test/`).
|
||||
|
||||
#### 1.4 Execution Method and Example
|
||||
|
||||
@@ -485,10 +368,12 @@ L3 level testing executes after code is merged into the main branch. Its core pu
|
||||
|
||||
#### 2.3 Test Directory and Execution Files
|
||||
|
||||
- ***Functional Testing***:
|
||||
- - Online Serving: `/tests/e2e/online_serving/test_{model_name}_expansion.py`
|
||||
- Offline Inference: `/tests/e2e/offline_inference/test_{model_name}_expansion.py`
|
||||
- (Note: `_expansion.py` likely means it contains more comprehensive scenario cases compared to L2 tests).
|
||||
- ***Functional / model E2E***:
|
||||
- Online serving: `tests/e2e/online_serving/test_{model_name}_expansion.py`
|
||||
- Offline inference: `tests/e2e/offline_inference/test_{model_name}_expansion.py`
|
||||
- Accuracy / similarity: `tests/e2e/accuracy/`
|
||||
- (`_expansion.py` holds broader scenario coverage than the L2 smoke file.)
|
||||
- ***Feature-level integration***: `tests/e2e/features/<feature>/` (not under a new top-level `tests/` directory).
|
||||
|
||||
#### 2.4 Execution Method and Example
|
||||
|
||||
|
||||
@@ -653,7 +653,7 @@ python end2end.py --model your-org/your-model-name --modality text2img --prompts
|
||||
Mirror BAGEL’s online serving setup:
|
||||
|
||||
- Server launcher: `examples/online_serving/your_model_name/run_server.sh`
|
||||
- Wrap `vllm serve ... --omni --port ...` (and `--stage-configs-path ...` if needed)
|
||||
- Wrap `vllm serve ... --omni --port ...` (and `--deploy-config ...` if needed)
|
||||
- Client: `examples/online_serving/your_model_name/openai_chat_client.py`
|
||||
- Send requests to `POST /v1/chat/completions`
|
||||
- Support multimodal inputs (e.g., base64 image) if your model needs it
|
||||
|
||||
@@ -48,8 +48,8 @@ vllm_omni/deploy/
|
||||
```
|
||||
|
||||
New in-tree models should register their `PipelineConfig` in
|
||||
`vllm_omni/config/pipeline_registry.py`; use `--stage-configs-path` only for
|
||||
custom legacy `stage_args` YAMLs.
|
||||
`vllm_omni/config/pipeline_registry.py`; use `--deploy-config` for custom
|
||||
deployment YAMLs.
|
||||
|
||||
## Step-by-Step Implementation
|
||||
|
||||
|
||||
@@ -449,9 +449,8 @@ Key configuration fields:
|
||||
!!! note
|
||||
New in-tree models should define frozen topology in `models/<model>/pipeline.py`,
|
||||
register it in `vllm_omni/config/pipeline_registry.py`, and put deployment
|
||||
knobs in `vllm_omni/deploy/<model>.yaml`. Legacy `stage_args` YAMLs are still
|
||||
accepted for custom configs via `--stage-configs-path`, but should not be used
|
||||
for new bundled defaults.
|
||||
knobs in `vllm_omni/deploy/<model>.yaml`. Use `--deploy-config` to load a
|
||||
custom deployment YAML.
|
||||
|
||||
### Batch mode
|
||||
|
||||
|
||||
@@ -0,0 +1,174 @@
|
||||
# Async Diffusion Output
|
||||
|
||||
## Table of Contents
|
||||
|
||||
1. [Overview](#overview)
|
||||
2. [Performance](#performance)
|
||||
3. [Architecture](#architecture)
|
||||
4. [Model Coverage](#model-coverage)
|
||||
5. [Configuration](#configuration)
|
||||
6. [Related Files](#related-files)
|
||||
|
||||
## Overview
|
||||
|
||||
The async diffusion output feature moves the D2H (Device→Host) copy and SHM packing from the Worker main thread to a background daemon thread, so the GPU default stream can immediately start the next request's forward pass. This eliminates the GPU bubble caused by synchronous D2H/packing.
|
||||
|
||||
**Core benefit**: the previous request's output D2H/packing overlaps with the next request's forward pass.
|
||||
|
||||
**Automatically enabled** when `step_execution=False` (default, request-mode). No extra configuration needed. When `step_execution=True` (step-mode), the original synchronous path is used — pump thread and Worker background thread are not started.
|
||||
|
||||
## Performance
|
||||
|
||||
### HunyuanImage-3.0 (TP4)
|
||||
|
||||
| Resolution | Async Off (QPS) | Async On (QPS) | Change |
|
||||
|---|---|---|---|
|
||||
| 1024×1024 | 0.4773 | 0.4802 | +0.60% |
|
||||
| 768×768 | 0.8370 | 0.8533 | +1.95% |
|
||||
|
||||
**Analysis**: smaller resolution → larger QPS gain. D2H/packing is relatively fixed overhead; forward is variable (scales with resolution). Smaller images → shorter forward → D2H/packing occupies a larger fraction of total time → more benefit from overlapping it.
|
||||
|
||||
## Architecture
|
||||
|
||||
### Execution Timeline
|
||||
|
||||
```
|
||||
Before (synchronous D2H):
|
||||
[forward req1] [D2H+SHM pack req1] [forward req2] [D2H+SHM pack req2]
|
||||
^^^^^^^^^^^^^^^^ GPU bubble
|
||||
|
||||
After (async D2H):
|
||||
[forward req1] [forward req2]
|
||||
[D2H+SHM pack req1 (side stream)] [D2H+SHM pack req2 (side stream)]
|
||||
^^^^^^^^^^^^^^^^ bubble eliminated
|
||||
```
|
||||
|
||||
### Data Flow
|
||||
|
||||
```
|
||||
┌─────────────────────────────────────────────────────┐
|
||||
│ Worker Process │
|
||||
│ │
|
||||
execute_model ──> │ WorkerBusyLoop ──> forward() ──> DiffusionOutput │
|
||||
│ │ │ (GPU tensor)│
|
||||
│ │ ┌─────────┴──────────┐ │
|
||||
│ │ │ async output path │ │
|
||||
│ │ │ │ │
|
||||
│ │ compute_done ──> result_mq │ │
|
||||
│ │ │ │ │
|
||||
│ │ ▼ │ │
|
||||
│ │ AsyncOutputThread │ │
|
||||
│ │ (side CUDA stream) │ │
|
||||
│ │ wait_event(gpu_event) │ │
|
||||
│ │ pack_diffusion_output_shm() │ │
|
||||
│ │ D2H + SHM write │ │
|
||||
│ │ │ │ │
|
||||
│ │ output_ready ──> result_mq │ │
|
||||
│ │ │ │ │
|
||||
│ ▼ ▼ │ │
|
||||
│ (dequeue next request) │ │
|
||||
└─────────────────────────────────────────────────────┘
|
||||
│
|
||||
▼
|
||||
┌─────────────────────────────────────────────────────┐
|
||||
│ Executor (main process) │
|
||||
│ │
|
||||
│ ResultPumpThread (sole reader of result_mq) │
|
||||
│ │ │
|
||||
│ ├── AsyncDiffusionOutput(COMPUTE_DONE) │
|
||||
│ │ → resolve _rpc_futures[rpc_id] │
|
||||
│ │ → RunnerOutput(async_output_id=...) │
|
||||
│ │ │
|
||||
│ ├── AsyncDiffusionOutput(OUTPUT_READY) │
|
||||
│ │ → resolve _output_futures[async_output_id]│
|
||||
│ │ → or batch split via _batch_split_map │
|
||||
│ │ │
|
||||
│ └── Non-async message │
|
||||
│ → _sync_result_buffer (for other RPCs) │
|
||||
│ │
|
||||
│ wait_output_ready(async_output_id) → Future │
|
||||
└─────────────────────────────────────────────────────┘
|
||||
```
|
||||
|
||||
### Key Components
|
||||
|
||||
1. **AsyncDiffusionOutput** (`data.py`): Protocol envelope for `result_mq`. `kind` field routes messages:
|
||||
- `COMPUTE_DONE` — forward finished, GPU can start next request
|
||||
- `OUTPUT_READY` — D2H/SHM packing finished, final output available
|
||||
- `RPC_RESULT` — ordinary RPC return (including error propagation)
|
||||
|
||||
2. **Worker AsyncOutputThread** (`diffusion_worker.py`): Background daemon thread that:
|
||||
- Waits for `gpu_event` (cross-stream sync ensuring forward wrote the tensor)
|
||||
- Runs `pack_diffusion_output_shm()` on a side CUDA stream
|
||||
- Enqueues `OUTPUT_READY` to `result_mq` when packing completes
|
||||
|
||||
3. **ResultPumpThread** (`multiproc_executor.py`): Sole reader of `result_mq` in request-mode. Dispatches messages:
|
||||
- `COMPUTE_DONE` / `RPC_RESULT` → resolves `_rpc_futures[rpc_id]`
|
||||
- `OUTPUT_READY` → resolves `_output_futures[async_output_id]` (with batch split via `_batch_split_map`)
|
||||
- Non-async messages → `_sync_result_buffer` for other RPCs
|
||||
|
||||
4. **collective_rpc() two-path dispatch** (`multiproc_executor.py`):
|
||||
- **Path 1**: `execute_model` / `execute_model_batch` → generates `rpc_id`, registers Future, waits for pump to deliver `compute_done`
|
||||
- **Path 2**: All other RPCs → `_sync_result_buffer` (pump-fed) or `_result_mq` (step-mode, no pump)
|
||||
|
||||
5. **record_device_event** (`platforms/`): Cross-platform GPU event recording for side-stream synchronization. Implemented for CUDA, NPU, ROCm, XPU, MUSA; base class returns `None` (safe no-op).
|
||||
|
||||
6. **IPC side-stream D2H** (`ipc.py`): `pack_diffusion_output_shm()` and helpers accept `d2h_stream` parameter; use `pin_memory` + `copy_(non_blocking=True)` on the side stream instead of synchronous `.cpu()`.
|
||||
|
||||
### Batch Split
|
||||
|
||||
When `execute_model_batch` returns a `COMPUTE_DONE` with a single `async_output_id` for the entire batch, the executor splits it into per-request `async_output_id`s (format `{batch_id}/{request_id}`) via `_batch_split_map`. When `OUTPUT_READY` arrives, the pump extracts per-request results from the batch output and resolves each request's output future independently.
|
||||
|
||||
## Model Coverage
|
||||
|
||||
### Currently Supported (request-mode, `step_execution=False`)
|
||||
|
||||
All models running in request-mode (`step_execution=False`, the default) automatically use async output.
|
||||
|
||||
**Verified models**:
|
||||
|
||||
| Model | Type | `supports_request_batch` |
|
||||
|---|---|---|
|
||||
| **HunyuanImage-3.0** | Image | `False` |
|
||||
| **Qwen-Image** | Image | `True` |
|
||||
|
||||
Other models with `step_execution=False` are also supported but not yet verified.
|
||||
|
||||
### Not Yet Supported (step-mode, `step_execution=True`)
|
||||
|
||||
When `step_execution=True` (or `streaming_output=True`, which auto-enables step-mode), models use `execute_stepwise` instead of `execute_model`, which is not in the async Path 1 whitelist. Async output is not applicable.
|
||||
|
||||
**Verified**: Helios (`step_execution=True`, does not benefit from this feature).
|
||||
|
||||
Other models with `step_execution=True` are also not applicable.
|
||||
|
||||
Adapting step-mode requires additional design because:
|
||||
|
||||
- Each step's `RunnerOutput` must be available synchronously for the scheduler to advance
|
||||
- Intermediate latent tensors may not need D2H (only the final step needs D2H + postprocess)
|
||||
- Steps are serial (step N+1 depends on step N output), so the overlap benefit is limited to "final step's D2H overlaps with the next request's forward"
|
||||
|
||||
**Planned approach**: first cover only the final step's async D2H, keep intermediate steps synchronous. This matches the request-mode benefit pattern.
|
||||
|
||||
## Configuration
|
||||
|
||||
No configuration needed. Async output is automatically enabled when `step_execution=False` (default).
|
||||
|
||||
| `step_execution` | Mode | Async Output | Pump Thread | Worker BG Thread |
|
||||
|---|---|---|---|---|
|
||||
| `False` (default) | request-mode | ✅ enabled | ✅ started | ✅ started |
|
||||
| `True` | step-mode | ❌ disabled | ❌ not started | ❌ not started |
|
||||
|
||||
## Related Files
|
||||
|
||||
- `vllm_omni/diffusion/data.py`: `AsyncDiffusionOutput`, `AsyncOutputKind`, `DiffusionOutput.async_output_id`
|
||||
- `vllm_omni/diffusion/worker/diffusion_worker.py`: `WorkerProc._return_result()`, `_async_output_loop()`, `_generate_async_output_id()`
|
||||
- `vllm_omni/diffusion/executor/multiproc_executor.py`: `ResultPumpThread`, `collective_rpc()` two-path dispatch, `wait_output_ready()`, `_batch_split_map`
|
||||
- `vllm_omni/diffusion/diffusion_engine.py`: `step()` / `step_streaming()` / `add_req_and_wait_for_response()` async output waiting
|
||||
- `vllm_omni/diffusion/sched/request_scheduler.py`: `update_from_output()` — `async_output_id` → `FINISHED_COMPLETED`
|
||||
- `vllm_omni/diffusion/ipc.py`: `pack_diffusion_output_shm()` side-stream D2H path
|
||||
- `vllm_omni/platforms/interface.py`: `OmniPlatform.record_device_event()` base class
|
||||
- `vllm_omni/platforms/cuda/platform.py`: `CudaOmniPlatform.record_device_event()`
|
||||
- `vllm_omni/platforms/npu/platform.py`: `NPUOmniPlatform.record_device_event()`
|
||||
- `vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py`: `get_hunyuan_image3_post_process_func()`, postprocess removal
|
||||
- `vllm_omni/diffusion/registry.py`: `_DIFFUSION_POST_PROCESS_FUNCS` postprocess registry
|
||||
@@ -264,7 +264,7 @@ Complete examples in the codebase:
|
||||
| **Z-Image** | `vllm_omni/diffusion/models/z_image/z_image_transformer.py` | Standard TP | Full implementation with validation |
|
||||
| **FLUX** | `vllm_omni/diffusion/models/flux/flux_transformer.py` | Dual-stream | Image + text streams |
|
||||
| **Qwen-Image** | `vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py` | Standard TP | With RoPE |
|
||||
| **TP Tests** | `tests/e2e/offline_inference/test_zimage_parallelism.py` | E2E testing | TP correctness and performance |
|
||||
| **TP Tests** | `tests/diffusion/distributed/test_tensor_parallel.py` | E2E testing | TP correctness and performance |
|
||||
| **Constraint Tests** | `tests/diffusion/models/z_image/test_zimage_tp_constraints.py` | Unit testing | Validation logic |
|
||||
|
||||
---
|
||||
|
||||
@@ -148,7 +148,13 @@ Upstream vLLM's `Scheduler.make_stats()` runs on every AR generation step, retur
|
||||
|
||||
`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.
|
||||
The orchestrator side does not add its own throttle on top. Scheduler gauges are
|
||||
recorded when `raw_outputs.scheduler_stats is not None` (i.e. this replica's
|
||||
scheduler passed its own 1 Hz gate), while output-batch request/token metrics
|
||||
are recorded from `IterationStats` whenever a non-empty output batch is
|
||||
processed. 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
|
||||
|
||||
|
||||
@@ -419,8 +419,7 @@ Notes:
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct \
|
||||
--omni \
|
||||
--port 8091 \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml
|
||||
--port 8091
|
||||
```
|
||||
|
||||
#### 3) Key config knobs
|
||||
@@ -494,7 +493,7 @@ The default config (`qwen3_tts.yaml`) enables the full optimization stack:
|
||||
vllm-omni serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
|
||||
--omni \
|
||||
--port 8000 \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts_no_async_chunk.yaml
|
||||
--no-async-chunk
|
||||
```
|
||||
|
||||
#### 3) Key config knobs
|
||||
|
||||
@@ -34,7 +34,7 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
|
||||
|
||||
### 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, 0.22.0, 0.23.0, 0.24.0, and 0.25.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, 0.21.0, 0.22.0, 0.23.0, 0.24.0, 0.25.0, and 0.26.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"
|
||||
|
||||
|
||||
@@ -20,7 +20,7 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new**
|
||||
|
||||
vLLM-Omni is built based on vLLM. Please install it with command below.
|
||||
```bash
|
||||
uv pip install vllm==0.25.0 --torch-backend=auto
|
||||
uv pip install vllm==0.26.0 --torch-backend=auto
|
||||
```
|
||||
|
||||
#### Installation of vLLM-Omni
|
||||
@@ -39,13 +39,13 @@ uv pip install 'vllm-omni[demo]'
|
||||
# --8<-- [start:build-wheel-from-source]
|
||||
|
||||
#### Installation of vLLM
|
||||
If you do not need to modify source code of vLLM, you can directly install the stable 0.25.0 release version of the library
|
||||
If you do not need to modify source code of vLLM, you can directly install the stable 0.26.0 release version of the library
|
||||
|
||||
```bash
|
||||
uv pip install vllm==0.25.0 --torch-backend=auto
|
||||
uv pip install vllm==0.26.0 --torch-backend=auto
|
||||
```
|
||||
|
||||
The 0.25.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
|
||||
The 0.26.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
|
||||
|
||||
#### Installation of vLLM-Omni
|
||||
Since vllm-omni is rapidly evolving, it's recommended to install it from source
|
||||
@@ -66,12 +66,12 @@ If you want to check, modify or debug with source code of vLLM, install the libr
|
||||
```bash
|
||||
git clone https://github.com/vllm-project/vllm.git
|
||||
cd vllm
|
||||
git checkout v0.25.0
|
||||
git checkout v0.26.0
|
||||
```
|
||||
Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file.
|
||||
```bash
|
||||
#For CUDA 13.0 (the default for v0.25.0; the wheel filename has no `+cu130` suffix)
|
||||
export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.25.0/vllm-0.25.0-cp38-abi3-manylinux_2_28_x86_64.whl
|
||||
#For CUDA 13.0 (the default for v0.26.0; the wheel filename has no `+cu130` suffix)
|
||||
export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.26.0/vllm-0.26.0-cp38-abi3-manylinux_2_28_x86_64.whl
|
||||
```
|
||||
Install vllm with command below (If you have no existing PyTorch).
|
||||
```bash
|
||||
@@ -102,7 +102,7 @@ docker run --runtime nvidia --gpus 2 \
|
||||
--env "HF_TOKEN=$HF_TOKEN" \
|
||||
-p 8091:8091 \
|
||||
--ipc=host \
|
||||
vllm/vllm-omni:v0.25.0 \
|
||||
vllm/vllm-omni:v0.26.0 \
|
||||
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.
|
||||
```bash
|
||||
uv pip install vllm==0.25.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.25.0/rocm723
|
||||
uv pip install vllm==0.26.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.26.0/rocm723
|
||||
```
|
||||
|
||||
#### Installation of vLLM-Omni
|
||||
@@ -34,13 +34,13 @@ uv pip install onnxruntime-rocm
|
||||
# --8<-- [start:build-wheel-from-source]
|
||||
|
||||
#### Installation of vLLM
|
||||
If you do not need to modify source code of vLLM, you can directly install the stable 0.25.0 release version of the library
|
||||
If you do not need to modify source code of vLLM, you can directly install the stable 0.26.0 release version of the library
|
||||
|
||||
```bash
|
||||
uv pip install vllm==0.25.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.25.0/rocm723
|
||||
uv pip install vllm==0.26.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.26.0/rocm723
|
||||
```
|
||||
|
||||
The pre-built 0.25.0 vLLM wheel targets ROCm 7.2.3. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
|
||||
The pre-built 0.26.0 vLLM wheel targets ROCm 7.2.3. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
|
||||
|
||||
#### Installation of vLLM-Omni
|
||||
Since vllm-omni is rapidly evolving, it's recommended to install it from source
|
||||
@@ -58,7 +58,7 @@ If you want to check, modify or debug with source code of vLLM, install the libr
|
||||
```bash
|
||||
git clone https://github.com/vllm-project/vllm.git
|
||||
cd vllm
|
||||
git checkout v0.25.0
|
||||
git checkout v0.26.0
|
||||
python3 -m pip install -r requirements/rocm.txt
|
||||
python3 setup.py develop
|
||||
```
|
||||
@@ -130,7 +130,7 @@ docker run --rm \
|
||||
-v ~/.cache/huggingface:/root/.cache/huggingface \
|
||||
--env "HF_TOKEN=$HF_TOKEN" \
|
||||
-p 8091:8091 \
|
||||
vllm/vllm-omni-rocm:v0.25.0 \
|
||||
vllm/vllm-omni-rocm:v0.26.0 \
|
||||
--model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
|
||||
```
|
||||
|
||||
@@ -149,7 +149,7 @@ docker run --rm -it \
|
||||
-v ~/.cache/huggingface:/root/.cache/huggingface \
|
||||
--env "HF_TOKEN=$HF_TOKEN" \
|
||||
--entrypoint bash \
|
||||
vllm/vllm-omni-rocm:v0.25.0
|
||||
vllm/vllm-omni-rocm:v0.26.0
|
||||
```
|
||||
|
||||
# --8<-- [end:pre-built-images]
|
||||
@@ -9,8 +9,8 @@ The recommended way to use vLLM-Omni on NPU is through the vLLM-Ascend pre-built
|
||||
|
||||
```bash
|
||||
# vLLM-Ascend image
|
||||
# Atlas A2: quay.io/atlas-ci/vllm-ascend:v0.25.0
|
||||
# Atlas A3: quay.io/atlas-ci/vllm-ascend:v0.25.0-a3
|
||||
# Atlas A2: quay.io/atlas-ci/vllm-ascend:v0.26.0
|
||||
# Atlas A3: quay.io/atlas-ci/vllm-ascend:v0.26.0-a3
|
||||
docker run --rm \
|
||||
--name vllm-omni-npu \
|
||||
--shm-size=1g \
|
||||
@@ -28,11 +28,11 @@ docker run --rm \
|
||||
-v /etc/ascend_install.info:/etc/ascend_install.info \
|
||||
-v /root/.cache:/root/.cache \
|
||||
-p 8000:8000 \
|
||||
-it quay.io/atlas-ci/vllm-ascend:0.25.0-a3 bash
|
||||
-it quay.io/atlas-ci/vllm-ascend:v0.26.0-a3 bash
|
||||
|
||||
# Inside the container, install vLLM-Omni from source
|
||||
cd /vllm-workspace
|
||||
git clone -b v0.25.0 https://github.com/vllm-project/vllm-omni.git
|
||||
git clone -b v0.26.0 https://github.com/vllm-project/vllm-omni.git
|
||||
cd vllm-omni
|
||||
pip install -v -e . --no-build-isolation
|
||||
# or VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e .
|
||||
@@ -84,8 +84,8 @@ Supported images as following.
|
||||
Here's an example deployment command that has been verified on 4 x NPUs:
|
||||
|
||||
```bash
|
||||
# Atlas A2: quay.io/ascend/vllm-omni:v0.25.0
|
||||
# Atlas A3: quay.io/ascend/vllm-omni:v0.25.0-a3
|
||||
# Atlas A2: quay.io/ascend/vllm-omni:v0.26.0
|
||||
# Atlas A3: quay.io/ascend/vllm-omni:v0.26.0-a3
|
||||
docker run --rm \
|
||||
--name vllm-omni-a3 \
|
||||
--shm-size=4g \
|
||||
@@ -110,7 +110,7 @@ docker run --rm \
|
||||
-v /etc/ascend_install.info:/etc/ascend_install.info \
|
||||
-v ~/.cache:/root/.cache \
|
||||
-p 8091:8091 \
|
||||
-it quay.io/ascend/vllm-omni:v0.25.0-a3 bash
|
||||
-it quay.io/ascend/vllm-omni:v0.26.0-a3 bash
|
||||
```
|
||||
|
||||
!!! tip
|
||||
@@ -122,9 +122,9 @@ Or build IMAGE from **source code**:
|
||||
git clone https://github.com/vllm-project/vllm-omni.git
|
||||
cd vllm-omni
|
||||
# A2
|
||||
docker build -t vllm-omni-dev-image:latest -f ./docker/Dockerfile.npu .
|
||||
# docker build -t vllm-omni-dev-image:latest -f ./docker/Dockerfile.npu .
|
||||
# A3
|
||||
# docker build -t vllm-omni-dev-image:latest -f ./docker/Dockerfile.npu.a3 .
|
||||
docker build -t vllm-omni-dev-image:latest -f ./docker/Dockerfile.npu.a3 .
|
||||
```
|
||||
|
||||
# --8<-- [end:pre-built-images]
|
||||
@@ -19,10 +19,10 @@ uv venv --python 3.12 --seed
|
||||
source .venv/bin/activate
|
||||
|
||||
# On CUDA
|
||||
uv pip install vllm==0.25.0 --torch-backend=auto
|
||||
uv pip install vllm==0.26.0 --torch-backend=auto
|
||||
|
||||
# On ROCm
|
||||
uv pip install vllm==0.25.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.25.0/rocm723
|
||||
uv pip install vllm==0.26.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.26.0/rocm723
|
||||
|
||||
git clone https://github.com/vllm-project/vllm-omni.git
|
||||
cd vllm-omni
|
||||
@@ -34,7 +34,7 @@ For additional installation methods — please see the [installation guide](inst
|
||||
!!! note
|
||||
It is important to install the same major & minor version of vLLM and vLLM Omni, otherwise things may not work as expected. If the versions are misaligned, you will see a warning when you import vLLM Omni.
|
||||
|
||||
If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.25.0` and vLLM Omni `0.25.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.26.0` and vLLM Omni `0.26.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.
|
||||
|
||||
## Offline Inference
|
||||
|
||||
|
||||
@@ -41,13 +41,10 @@ th {
|
||||
| `Wan22S2VPipeline` | 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.2-VACE | `Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers` | ✅︎ | | | |
|
||||
| `LTX2Pipeline` | LTX-2-T2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
|
||||
| `LTX2ImageToVideoPipeline` | LTX-2-I2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
|
||||
| `LTX2TwoStagesPipeline` | LTX-2-T2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
|
||||
| `LTX2ImageToVideoTwoStagesPipeline` | LTX-2-I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
|
||||
| `LTX23Pipeline` | LTX-2.3-T2V | `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
|
||||
| `LTX23ImageToVideoPipeline` | LTX-2.3-I2V | `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
|
||||
| `LTX2Pipeline` | LTX-2 / LTX-2.3 one-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
|
||||
| `LTX2DistilledPipeline` | LTX-2 distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
|
||||
| `LingBotVideoPipeline` | LingBot-Video dense and MoE T2V | `robbyant/lingbot-video-dense-1.3b`, `robbyant/lingbot-video-moe-30b-a3b` | ✅︎ | | | |
|
||||
| `MiniMaxH3Pipeline` | MiniMax H3 T2VA, FL2VA, and Ref2VA | `MiniMaxAI/MiniMax-H3` | ✅︎ | | | |
|
||||
| `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | |
|
||||
| `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | |
|
||||
|
||||
@@ -276,20 +276,42 @@ Client -> Server:
|
||||
|
||||
| Message | Description |
|
||||
|---------|-------------|
|
||||
| `{"type": "session.config", ...}` | Session configuration (sent once, first message) |
|
||||
| `{"type": "session.config", ...}` | Session configuration (first message; may be resent between utterances to change it) |
|
||||
| `{"type": "input.text", "text": "..."}` | Text chunk |
|
||||
| `{"type": "input.done"}` | End of input, flushes remaining buffer |
|
||||
| `{"type": "input.done"}` | End of utterance: flushes the buffer and keeps the connection open |
|
||||
| `{"type": "session.close"}` | End of connection |
|
||||
|
||||
Server -> Client:
|
||||
|
||||
| Message | Description |
|
||||
|---------|-------------|
|
||||
| `{"type": "audio.start", "sentence_index": 0, "sentence_text": "...", "format": "pcm", "sample_rate": 24000}` | Audio generation starting for the buffered input |
|
||||
| `{"type": "audio.start", "utterance_index": 0, "sentence_index": 0, "sentence_text": "...", "format": "pcm", "sample_rate": 24000}` | Audio generation starting for the buffered input |
|
||||
| Binary frame | Raw audio bytes (one or more PCM chunks when `stream_audio=true`) |
|
||||
| `{"type": "audio.done", "sentence_index": 0, "total_bytes": 96000, "error": false}` | Audio complete for the buffered input |
|
||||
| `{"type": "session.done", "total_sentences": N}` | Session complete |
|
||||
| `{"type": "audio.done", "utterance_index": 0, "sentence_index": 0, "total_bytes": 96000, "error": false}` | Audio complete for the buffered input |
|
||||
| `{"type": "session.done", "utterance_index": 0, "total_sentences": N}` | Flushed utterance complete |
|
||||
| `{"type": "error", "message": "..."}` | Non-fatal error |
|
||||
|
||||
### Flushing vs. Closing
|
||||
|
||||
`input.done` is a flush, not a disconnect. The server synthesizes the buffered
|
||||
text, emits `session.done`, and then waits on the same connection for the next
|
||||
utterance, so a client that speaks repeatedly (for example one driven by an
|
||||
upstream LLM) pays the WebSocket handshake once instead of once per utterance.
|
||||
|
||||
* The session config is sticky. Send `input.text` again straight after
|
||||
`session.done` to reuse it, or send another `session.config` first to change
|
||||
voice, format, or reference audio. A `session.config` sent while text is
|
||||
still buffered is rejected so no pending input is silently dropped.
|
||||
* An utterance is the flush unit, not a linguistic one: it is whatever text was
|
||||
buffered when `input.done` arrived, of any length, synthesized as one request.
|
||||
`utterance_index` counts those flushes across the connection, so it tells you
|
||||
which `input.done` a frame belongs to. `sentence_index` counts within one
|
||||
flush and so pairs with `total_sentences`, which means every utterance reports
|
||||
`sentence_index: 0` of `total_sentences: 1` (or `0` for an empty buffer).
|
||||
* End the connection with `session.close`, or by closing the socket. An idle
|
||||
connection is still closed after the server's idle timeout, which now also
|
||||
applies to the gap between utterances.
|
||||
|
||||
### Session Config Parameters
|
||||
|
||||
All REST API parameters are supported, plus:
|
||||
|
||||
@@ -73,7 +73,7 @@ curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o output.mp4
|
||||
| `audio_reference` | string | null | JSON-encoded audio reference for speech-to-video: `{"audio_url": "..."}` — supports HTTP(s) URLs or base64 data URLs |
|
||||
| `width` | integer | model default | Output video width |
|
||||
| `height` | integer | model default | Output video height |
|
||||
| `num_frames` | integer | model default | Number of generated frames |
|
||||
| `num_frames` | integer | 1 | Number of generated frames |
|
||||
| `fps` | integer | model default | Output frames per second |
|
||||
| `num_inference_steps` | integer | model default | Number of diffusion steps |
|
||||
| `guidance_scale` | number | null | CFG guidance scale for the low-noise stage |
|
||||
|
||||
@@ -15,7 +15,7 @@ The full set of backends and their platform defaults is in the **Backend Options
|
||||
| Value | Notes |
|
||||
|---|---|
|
||||
| `TRTLLM_ATTN` | FlashInfer's trtllm-gen FMHA (TensorRT-LLM's generated kernels, vendored by FlashInfer). BF16, GQA native, `head_dim=128`. Datacenter Blackwell only (sm_100 / sm_103). Supports **Skip-Softmax** sparse attention — see below. Requires `flashinfer`. |
|
||||
| `FLASH_ATTN` | Wraps FlashAttention 2. Default on Hopper / Ada / Ampere when `flash-attn` is installed. |
|
||||
| `FLASH_ATTN` | Wraps FlashAttention 4 on Blackwell when `flash-attn-4` is installed, then falls back to FlashAttention 3/2. Default on Hopper / Ada / Ampere when a compatible FlashAttention package is installed. |
|
||||
| `CUDNN_ATTN` | Pins `sdpa_kernel([CUDNN_ATTENTION])`. Default on Blackwell (sm_10x / sm_12x) with cuDNN ≥ 9.5. Wins on mask-heavy DiTs (HunyuanVideo-1.5: 2× e2e vs SDPA). |
|
||||
| `FLASHINFER_ATTN` | Calls FlashInfer's dense `single_prefill_with_kv_cache` directly with `custom_mask` for non-causal masked attention. Used as Blackwell fallback when cuDNN is unavailable. Requires `flashinfer`. |
|
||||
| `TORCH_SDPA` | PyTorch `scaled_dot_product_attention` with the default backend dispatcher. Most conservative; always available. |
|
||||
@@ -167,6 +167,48 @@ sequence length, since Skip-Softmax only accelerates attention. Dense BF16 is `T
|
||||
Requires datacenter Blackwell with `head_dim == 128`; elsewhere, or without FlashInfer, selecting
|
||||
`TRTLLM_ATTN` raises rather than silently degrading.
|
||||
|
||||
## TRTLLM_ATTN SAGE Quantization
|
||||
|
||||
On top of dense/Skip-Softmax, `TRTLLM_ATTN` can run **SAGE** (SageAttention) — the trtllm-gen
|
||||
FMHA kernel with per-block quantized Q/K and per-channel fp8 V, so the two attention matmuls run
|
||||
in low precision. Enable it through the typed `quant` block on the attention spec:
|
||||
|
||||
| Key | Valid values | Meaning |
|
||||
|---|---|---|
|
||||
| `dtype_qk` | `int8`, `fp8_e4m3` | Quantization dtype for Q and K. Absent ⇒ dense (SAGE off). |
|
||||
| `q_block_size` | `1`, `4`, `16` | Per-token block size for Q scales (default `1`). |
|
||||
| `k_block_size` | `1`, `4`, `16` | Per-token block size for K scales (default `16`). |
|
||||
|
||||
V is always quantized per-channel to fp8_e4m3 and K-smoothing is applied inside the routine, so
|
||||
neither is a user knob. Only block sizes `1`, `4`, `16` have compiled kernels.
|
||||
|
||||
The `quant` block is a shared `AttnQuantSpec` consumed by more than one backend; each reads the
|
||||
fields that apply and rejects the rest. `TRTLLM_ATTN` (SAGE) reads `dtype_qk` (`int8` / `fp8_e4m3`),
|
||||
`q_block_size`, and `k_block_size`; `FLASHINFER_ATTN` reads `dtype_qk` (`float16` / `bfloat16`),
|
||||
`dtype_vo`, and `flashinfer_backend` (see its section). A `dtype_qk` from the wrong set raises.
|
||||
|
||||
```bash
|
||||
vllm-omni serve Wan-AI/Wan2.2-T2V-A14B-Diffusers \
|
||||
--diffusion-attention-config '{"default": {"backend": "TRTLLM_ATTN",
|
||||
"quant": {"dtype_qk": "fp8_e4m3", "q_block_size": 1, "k_block_size": 16}}}'
|
||||
```
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.data import AttentionConfig, AttentionSpec, AttnQuantSpec
|
||||
|
||||
AttentionConfig(
|
||||
default=AttentionSpec(
|
||||
backend="TRTLLM_ATTN",
|
||||
quant=AttnQuantSpec(dtype_qk="fp8_e4m3", q_block_size=1, k_block_size=16),
|
||||
),
|
||||
)
|
||||
```
|
||||
|
||||
|
||||
**Requirements.** Needs FlashInfer ≥ 0.6.16rc1. Kernel
|
||||
availability is arch-dependent: `fp8_e4m3` QK has kernels on both **SM100** (B200) and **SM103**
|
||||
(B300); `int8` QK kernels are compiled for **SM100 only**.
|
||||
|
||||
## End-to-End Benchmark (BF16, sm_120 RTX Pro 6000 Blackwell)
|
||||
|
||||
Same prompt and seed across runs. `Total generation time` from `text_to_video.py` / `text_to_image.py`.
|
||||
@@ -193,9 +235,17 @@ DIFFUSION_ATTENTION_BACKEND=FLASHINFER_ATTN python examples/offline_inference/te
|
||||
--model Lightricks/LTX-2 ...
|
||||
```
|
||||
|
||||
### FA4 not yet integrated
|
||||
### FlashAttention-4 on Blackwell
|
||||
|
||||
FlashAttention-4 (released March 2026) targets Blackwell natively and reportedly beats cuDNN by ~20% on B200. As of this writing the `flash-attn-4 4.0.0b10` wheel crashes with `AttributeError: 'NoneType' object has no attribute '_trait'` during JIT on sm_120. Not yet wired into vLLM-Omni; revisit when stable lands.
|
||||
Install the optional CUDA 13 extra to use the CuTe-based FA4 path:
|
||||
|
||||
```bash
|
||||
pip install 'vllm-omni[fa4]'
|
||||
```
|
||||
|
||||
`FLASH_ATTN` prefers `flash_attn.cute` on Blackwell and falls back to FA3/FA2
|
||||
when it is unavailable. Version `4.0.0b18` is required; earlier beta
|
||||
wheels had known JIT failures on Blackwell.
|
||||
|
||||
## Choosing a Backend Manually
|
||||
|
||||
@@ -353,64 +403,3 @@ DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN_3 python examples/offline_inference/text_t
|
||||
--tensor-parallel-size 2 \
|
||||
--output outputs/hv15_sage3.mp4
|
||||
```
|
||||
|
||||
### Mixed backends across roles
|
||||
|
||||
Use `FLASH_ATTN` for self-attention and `TORCH_SDPA` for cross-attention:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
--model Wan-AI/Wan2.2-TI2V-5B-Diffusers \
|
||||
--prompt "A dog running across a field of golden wheat." \
|
||||
--diffusion-attention-config.per_role.self.backend FLASH_ATTN \
|
||||
--diffusion-attention-config.per_role.cross.backend TORCH_SDPA \
|
||||
--tensor-parallel-size 2 \
|
||||
--output outputs/wan22_mixed.mp4
|
||||
```
|
||||
|
||||
### Compare against FlashAttention
|
||||
|
||||
Unset the backend override, or explicitly use `FLASH_ATTN`:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_video/text_to_video.py \
|
||||
--model Wan-AI/Wan2.2-TI2V-5B-Diffusers \
|
||||
--prompt "A dog running across a field of golden wheat." \
|
||||
--height 704 --width 1280 --num-frames 49 \
|
||||
--num-inference-steps 30 --seed 42 --guidance-scale 5.0 \
|
||||
--tensor-parallel-size 2 \
|
||||
--output outputs/wan22_fa3.mp4
|
||||
```
|
||||
|
||||
## Validation Guidance
|
||||
|
||||
Don't assume a faster attention backend is numerically interchangeable with `TORCH_SDPA`.
|
||||
|
||||
Always compare:
|
||||
|
||||
- End-to-end runtime
|
||||
- Diffusion-stage runtime (`add_req_and_wait` line in DiffusionEngine.step breakdown)
|
||||
- Output quality against a known-good baseline (CLIP similarity, frame-level diff, or visual review)
|
||||
|
||||
At minimum, keep the same:
|
||||
|
||||
- model
|
||||
- prompt
|
||||
- seed
|
||||
- resolution
|
||||
- frame count / step count
|
||||
- parallel config (TP / CFG-parallel / Ulysses degrees)
|
||||
|
||||
## Reproducing the Benchmark Table
|
||||
|
||||
The end-to-end numbers above were collected by running `text_to_video.py` /
|
||||
`text_to_image.py` with the same prompt and seed while varying
|
||||
`DIFFUSION_ATTENTION_BACKEND`. For a quick kernel-level comparison of the
|
||||
backends without loading a model:
|
||||
|
||||
```bash
|
||||
python benchmarks/diffusion/bench_attention_backends.py --preset hv15
|
||||
```
|
||||
|
||||
It runs all three BF16 backends on representative DiT attention shapes and
|
||||
prints a ranking table at the end.
|
||||
@@ -179,7 +179,134 @@ class Flux2Transformer2DModel(nn.Module):
|
||||
- Support single GPU only for now
|
||||
|
||||
|
||||
### Implementation Notes
|
||||
## Distributed Layerwise Offloading
|
||||
|
||||
### How It Works
|
||||
|
||||
Distributed layerwise offloading extends single-GPU layerwise offloading to
|
||||
multi-device deployments. Each DP rank stores only **1/dp_size** of the model
|
||||
weights in host memory; full layer weights are reconstructed at runtime via
|
||||
**AllGather** on a dedicated communication stream, overlapped with computation
|
||||
via a fixed double-buffer scheme.
|
||||
|
||||
**Key features:**
|
||||
- **Weight sharding + AllGather**: each rank stores 1/dp_size of weights,
|
||||
reconstructed per-layer via `all_gather_into_tensor`
|
||||
- **Fixed double-buffer**: exactly 2 transformer blocks on each device at any
|
||||
time, regardless of model size
|
||||
- **DP multi-concurrency**: N concurrent requests processed in parallel
|
||||
(AllGather only gathers weight shards, which are request-independent)
|
||||
- **mmap weight loading**: weights loaded as mmap views pointing to shared OS
|
||||
page cache, eliminating O(dp_size × model_size) RSS during model creation
|
||||
- **Platform-agnostic**: works on NVIDIA GPU (CUDA/NCCL) and Ascend NPU
|
||||
(CANN/HCCL) via vLLM-Omni's platform abstraction
|
||||
|
||||
**Execution Flow (per device):**
|
||||
|
||||
```
|
||||
Compute Stream: [Layer N] [Layer N+1] [Layer N+2]
|
||||
H2D Stream: [H2D Shard N+1] [H2D Shard N+2]
|
||||
AllGather: [AG N+1] [AG N+2]
|
||||
Slot Usage: Slot0: Layer N Slot1: Layer N+1
|
||||
```
|
||||
|
||||
### Usage
|
||||
|
||||
**CLI:**
|
||||
```bash
|
||||
# 4× GPU/NPU with AllGather (recommended)
|
||||
vllm serve /path/to/model --omni \
|
||||
--enable-distributed-layerwise-offload \
|
||||
--data-parallel-size 4
|
||||
|
||||
# Without AllGather (each rank loads full weights, no sharding)
|
||||
vllm serve /path/to/model --omni \
|
||||
--enable-distributed-layerwise-offload \
|
||||
--data-parallel-size 4 \
|
||||
--dlo-no-use-allgather
|
||||
|
||||
# With SP instead of DP (long sequences)
|
||||
vllm serve /path/to/model --omni \
|
||||
--enable-distributed-layerwise-offload \
|
||||
--usp 4
|
||||
```
|
||||
|
||||
**Python API:**
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
|
||||
m = Omni(
|
||||
model="/path/to/model",
|
||||
enable_distributed_layerwise_offload=True,
|
||||
dlo_use_allgather=True, # default
|
||||
)
|
||||
```
|
||||
|
||||
### CLI Flags
|
||||
|
||||
| Flag | Description | Default |
|
||||
|------|-------------|---------|
|
||||
| `--enable-distributed-layerwise-offload` | Enable DLO | `false` |
|
||||
| `--data-parallel-size N` | Number of DP ranks for weight sharding | `1` |
|
||||
| `--dlo-use-allgather` | Shard + AllGather (saves CPU, requires concurrent requests) | `true` |
|
||||
| `--dlo-no-use-allgather` | Full weights per rank (no sharding, no AllGather) | `false` |
|
||||
|
||||
### How model weights are loaded (mmap path)
|
||||
|
||||
When DLO + AllGather is active, the offloader:
|
||||
|
||||
1. **Saves non-persistent buffers** (e.g. RoPE `inv_freq`, timestep `freqs`)
|
||||
from the normally-created transformer
|
||||
2. **Converts to meta device** via `to_empty(device="meta")`, releasing random
|
||||
initialization weights
|
||||
3. **Loads checkpoint weights as mmap views** via `safe_open().get_tensor()`,
|
||||
which return views into the OS page cache (shared across all ranks, 0 RSS)
|
||||
4. **Calls `post_load_weights()`** to apply model-specific dtype conversions
|
||||
(e.g. Cosmos3's `time_embedder` → FP32)
|
||||
5. **Restores non-persistent buffers** from saved copies
|
||||
|
||||
This approach requires **zero model-specific code changes** — no pipeline or
|
||||
transformer modifications are needed.
|
||||
|
||||
### OffloadPlan (declarative topology metadata)
|
||||
|
||||
Models can optionally declare an `OffloadPlan` class variable to provide
|
||||
topology metadata (block attribute names, submodules to offload) without
|
||||
any offload-specific logic:
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.offloader import OffloadPlan
|
||||
|
||||
class MyPipeline(nn.Module):
|
||||
_dit_modules = ["transformer"]
|
||||
_offload_plan = OffloadPlan(
|
||||
block_attrs={"transformer": ("blocks",)},
|
||||
offload_submodules={"context_encoder": "layers"},
|
||||
)
|
||||
```
|
||||
|
||||
When not declared, the offloader falls back to `_layerwise_offload_blocks_attrs`
|
||||
and heuristic attribute search. This is backward-compatible — existing models
|
||||
work without any changes.
|
||||
|
||||
### DP Multi-concurrency
|
||||
|
||||
When `--data-parallel-size > 1` and AllGather is enabled, the scheduler batches
|
||||
up to `dp_size` requests per denoise step. Each DP rank processes a different
|
||||
request while AllGather synchronizes only weight shards (request-independent).
|
||||
|
||||
**Requirements for concurrent requests:**
|
||||
- `num_inference_steps` must be specified explicitly (None is not allowed)
|
||||
- All concurrent requests must have the same `num_inference_steps` value
|
||||
- These constraints exist because AllGather is a collective that requires
|
||||
every rank to participate at each denoise step
|
||||
|
||||
### Limitations
|
||||
- Online quantization (FP8) is incompatible with mmap loading — falls back to
|
||||
regular `load_weights()` automatically
|
||||
- Tensor Parallel is not supported (DLO uses DP-based sharding)
|
||||
- HSDP + AllGather is rejected (would double-shard weights)
|
||||
- `num_inference_steps=None` is not allowed in DP multi-concurrency mode
|
||||
|
||||
**Module Discovery**
|
||||
|
||||
@@ -208,7 +335,10 @@ OffloadBackend (base class)
|
||||
├── ModelLevelOffloadBackend → uses SequentialOffloadHook (.to() swap)
|
||||
│ (delegates to a pipeline's enable_omni_model_cpu_offload
|
||||
│ for split models like Cosmos3)
|
||||
└── LayerWiseOffloadBackend → uses LayerwiseOffloadHook
|
||||
├── LayerWiseOffloadBackend → uses LayerwiseOffloadHook
|
||||
│ (single-GPU, full weights on host)
|
||||
└── DistributedLayerwiseOffloadBackend → uses DistributedLayerwiseOffloadHook
|
||||
(multi-GPU, 1/dp_size sharded weights + AllGather)
|
||||
```
|
||||
|
||||
Factory function `get_offload_backend()` selects the appropriate backend based on
|
||||
@@ -221,20 +351,22 @@ reasoner/generator components inside the model forward pass.
|
||||
|
||||
## Supported Models
|
||||
|
||||
| Architecture | Example Models | DiT Class | Model-Level Offload | Layerwise Offload | Blocks Attrs (Layerwise specific) |
|
||||
|--------------|----------------|-----------|---------------------|-------------------|-----------------------------------|
|
||||
| Flux2Pipeline | `black-forest-labs/FLUX.2-dev` | `Flux2Transformer2DModel` | ✓ | ✓ | `"transformer_blocks"`, `"single_transformer_blocks"` |
|
||||
| LongCatImagePipeline | `meituan-longcat/LongCat-Image` | `LongCatImageTransformer2DModel` | - | ✓ | `"transformer_blocks"`, `"single_transformer_blocks"` |
|
||||
| NextStep11Pipeline | `stepfun-ai/NextStep-1.1` | `NextStepModel` | - | ✓ | `"layers"` |
|
||||
| OvisImagePipeline | `AIDC-AI/Ovis-Image-7B` | `OvisImageTransformer2DModel` | - | ✓ | `"transformer"` |
|
||||
| QwenImagePipeline | `Qwen/Qwen-Image` | `QwenImageTransformer2DModel` | ✓ | ✓ | `"transformer_blocks"` |
|
||||
| StableDiffusionXLPipeline | `stabilityai/stable-diffusion-xl-base-1.0` | `SDXLUNet2DConditionModel` | ✓ | ✓ | `"down_blocks"`, `"up_blocks"` |
|
||||
| StableDiffusion3Pipeline | `stabilityai/stable-diffusion-3.5-medium` | `SD3Transformer2DModel` | - | ✓ | `"transformer_blocks"` |
|
||||
| Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` |
|
||||
| Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` |
|
||||
| SoulXSingerPipeline / SoulXSingerSVCPipeline | `Soul-AILab/SoulX-Singer` | `DiffLlama` (`cfm_decoder.model.diff_estimator`) | ✓ | ✓ | `"layers"` |
|
||||
| BagelPipeline | `ByteDance-Seed/BAGEL-7B-MoT` | `Qwen2MoTModel` | - | ✓ | `"layers"`, `"customized modules"` |
|
||||
| Architecture | Example Models | DiT Class | Model-Level Offload | Layerwise Offload | Distributed Layerwise Offload | Blocks Attrs (Layerwise specific) |
|
||||
|--------------|----------------|-----------|---------------------|-------------------|-------------------------------|-----------------------------------|
|
||||
| Flux2Pipeline | `black-forest-labs/FLUX.2-dev` | `Flux2Transformer2DModel` | ✓ | ✓ | - | `"transformer_blocks"`, `"single_transformer_blocks"` |
|
||||
| LongCatImagePipeline | `meituan-longcat/LongCat-Image` | `LongCatImageTransformer2DModel` | - | ✓ | - | `"transformer_blocks"`, `"single_transformer_blocks"` |
|
||||
| NextStep11Pipeline | `stepfun-ai/NextStep-1.1` | `NextStepModel` | - | ✓ | - | `"layers"` |
|
||||
| OvisImagePipeline | `AIDC-AI/Ovis-Image-7B` | `OvisImageTransformer2DModel` | - | ✓ | - | `"transformer"` |
|
||||
| QwenImagePipeline | `Qwen/Qwen-Image` | `QwenImageTransformer2DModel` | ✓ | ✓ | - | `"transformer_blocks"` |
|
||||
| StableDiffusionXLPipeline | `stabilityai/stable-diffusion-xl-base-1.0` | `SDXLUNet2DConditionModel` | ✓ | ✓ | - | `"down_blocks"`, `"up_blocks"` |
|
||||
| StableDiffusion3Pipeline | `stabilityai/stable-diffusion-3.5-medium` | `SD3Transformer2DModel` | - | ✓ | - | `"transformer_blocks"` |
|
||||
| Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | - | `"blocks"` |
|
||||
| Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | - | `"blocks"` |
|
||||
| SoulXSingerPipeline / SoulXSingerSVCPipeline | `Soul-AILab/SoulX-Singer` | `DiffLlama` (`cfm_decoder.model.diff_estimator`) | ✓ | ✓ | - | `"layers"` |
|
||||
| BagelPipeline | `ByteDance-Seed/BAGEL-7B-MoT` | `Qwen2MoTModel` | - | ✓ | - | `"layers"`, `"customized modules"` |
|
||||
| Cosmos3OmniDiffusersPipeline | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | `Cosmos3VFMTransformer`, `Cosmos3LanguageModel` | ✓ | ✓ | ✓ | `"layers"`, `"gen_layers"` |
|
||||
|
||||
**Notes:**
|
||||
- Model-Level Offloading is expected to be supported by all common diffusion models (DiT and encoders) naturally
|
||||
- Layerwise Offloading requires DiT class to define `_layerwise_offload_blocks_attrs` pointing to transformer blocks
|
||||
- Distributed Layerwise Offloading works with any model that supports Layerwise Offloading — no additional model changes required. See [Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md) for usage examples.
|
||||
@@ -4,7 +4,7 @@ Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inf
|
||||
|
||||
|
||||
This example demonstrates how to generate videos from images using Wan2.2,
|
||||
LTX-2/LTX-2.3, and HunyuanVideo-1.5 Image-to-Video models with vLLM-Omni's
|
||||
LTX-2, and HunyuanVideo-1.5 Image-to-Video models with vLLM-Omni's
|
||||
offline inference API.
|
||||
|
||||
## Local CLI Usage
|
||||
@@ -53,34 +53,22 @@ python image_to_video.py \
|
||||
--output i2v_output.mp4
|
||||
```
|
||||
|
||||
### LTX-2.3
|
||||
### LTX-2
|
||||
|
||||
```bash
|
||||
python image_to_video.py \
|
||||
--model diffusers/LTX-2.3-Diffusers \
|
||||
--model-class-name LTX23ImageToVideoPipeline \
|
||||
--model Lightricks/LTX-2 \
|
||||
--image cherry_blossom.jpg \
|
||||
--prompt "Cherry blossoms swaying gently in the breeze with synchronized ambient sound" \
|
||||
--negative-prompt "worst quality, inconsistent motion, blurry, jittery, distorted" \
|
||||
--height 384 \
|
||||
--width 512 \
|
||||
--num-frames 25 \
|
||||
--guidance-scale 4.0 \
|
||||
--num-inference-steps 20 \
|
||||
--frame-rate 24 \
|
||||
--fps 24 \
|
||||
--output ltx23_i2v_output.mp4
|
||||
--output ltx2_i2v_output.mp4
|
||||
```
|
||||
|
||||
Use the Diffusers-format checkpoint `diffusers/LTX-2.3-Diffusers`; the
|
||||
upstream `Lightricks/LTX-2.3` raw safetensors repo is not directly loadable by
|
||||
this pipeline. Pass `--model-class-name LTX23ImageToVideoPipeline` to select
|
||||
the LTX-2.3 image-to-video pipeline.
|
||||
See the [LTX-2 recipe](../../../../recipes/LTX/LTX-2.md) for all checkpoints,
|
||||
pipeline selection, T2V, defaults, and advanced options.
|
||||
|
||||
Key arguments:
|
||||
|
||||
- `--model`: Model ID (I2V-A14B for MoE, TI2V-5B for unified T2V+I2V, or
|
||||
LTX-2/LTX-2.3).
|
||||
- `--model`: Model ID (I2V-A14B for MoE, TI2V-5B for unified T2V+I2V, or LTX-2).
|
||||
- `--image`: Path to input image (required).
|
||||
- `--prompt`: Text description of desired motion/animation.
|
||||
- `--height/--width`: Output resolution (auto-calculated from image if not set).
|
||||
|
||||
@@ -12,7 +12,7 @@ For backend selection and SageAttention usage, see the [Diffusion Attention Back
|
||||
| Model | Default Resolution | Default Frames | Default Steps | Guidance | VRAM (BF16) |
|
||||
|---|---|---|---|---|---|
|
||||
| `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 720x1280 | 81 | 40 | 4.0 | ~60 GiB |
|
||||
| `diffusers/LTX-2.3-Diffusers` | 384x512 | 25 | 20 | 4.0 | 96GB-class GPU |
|
||||
| `Lightricks/LTX-2` | 512x768 | 121 | 40 | video 3.0 / audio 7.0 | Model-dependent |
|
||||
| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v` | 480x832 | 121 | 50 | 6.0 | 1×A100 80GB |
|
||||
| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required |
|
||||
|
||||
@@ -35,45 +35,17 @@ python text_to_video.py \
|
||||
--output t2v_out.mp4
|
||||
```
|
||||
|
||||
LTX2 example:
|
||||
### LTX-2
|
||||
|
||||
```bash
|
||||
python text_to_video.py \
|
||||
--model "Lightricks/LTX-2" \
|
||||
--prompt "A cinematic close-up of ocean waves at golden hour." \
|
||||
--negative-prompt "worst quality, inconsistent motion, blurry, jittery, distorted" \
|
||||
--height 512 \
|
||||
--width 768 \
|
||||
--num-frames 121 \
|
||||
--num-inference-steps 40 \
|
||||
--guidance-scale 4.0 \
|
||||
--frame-rate 24 \
|
||||
--output ltx2_out.mp4
|
||||
```
|
||||
|
||||
### LTX-2.3
|
||||
|
||||
```bash
|
||||
python text_to_video.py \
|
||||
--model diffusers/LTX-2.3-Diffusers \
|
||||
--model-class-name LTX23Pipeline \
|
||||
--model Lightricks/LTX-2 \
|
||||
--prompt "Cherry blossoms swaying gently in the breeze with synchronized ambient sound" \
|
||||
--negative-prompt "worst quality, inconsistent motion, blurry, jittery, distorted" \
|
||||
--height 384 \
|
||||
--width 512 \
|
||||
--num-frames 25 \
|
||||
--num-inference-steps 20 \
|
||||
--guidance-scale 4.0 \
|
||||
--frame-rate 24 \
|
||||
--fps 24 \
|
||||
--audio-sample-rate 48000 \
|
||||
--output ltx23_t2v_output.mp4
|
||||
--output ltx2_output.mp4
|
||||
```
|
||||
|
||||
Use the Diffusers-format checkpoint `diffusers/LTX-2.3-Diffusers`; the
|
||||
upstream `Lightricks/LTX-2.3` raw safetensors repo is not directly loadable by
|
||||
this pipeline. Pass `--model-class-name LTX23Pipeline` to select the LTX-2.3
|
||||
text-to-video pipeline explicitly.
|
||||
See the [LTX-2 recipe](../../../../recipes/LTX/LTX-2.md) for all checkpoints,
|
||||
pipeline selection, I2V, defaults, and advanced options.
|
||||
|
||||
### HunyuanVideo-1.5 (480p)
|
||||
|
||||
@@ -135,8 +107,7 @@ python text_to_video.py \
|
||||
### Common
|
||||
|
||||
- `--model`: Diffusers model ID or local path.
|
||||
- `--model-class-name`: Optional explicit pipeline class. Use `LTX23Pipeline`
|
||||
for LTX-2.3 text-to-video.
|
||||
- `--model-class-name`: Optional explicit pipeline override.
|
||||
- `--prompt`: text description (string).
|
||||
- `--height/--width`: output resolution. Default depends on model.
|
||||
- `--num-frames`: number of frames. Default depends on model.
|
||||
@@ -151,8 +122,7 @@ python text_to_video.py \
|
||||
- `--enable-cpu-offload`: enable CPU offloading for diffusion models.
|
||||
- `--enable-layerwise-offload`: enable layerwise offloading on DiT modules.
|
||||
- `--frame-rate`: generation FPS for pipelines that require it (e.g., LTX2).
|
||||
- `--audio-sample-rate`: audio sample rate for embedded audio (when the
|
||||
pipeline returns audio; LTX-2.3 outputs 48kHz audio).
|
||||
- `--audio-sample-rate`: fallback audio sample rate when the pipeline returns audio.
|
||||
- `--quantization`: quantization method (`fp8` for FP8, `gguf` for GGUF).
|
||||
- `--flow-shift`: scheduler flow_shift parameter.
|
||||
|
||||
|
||||
@@ -40,7 +40,7 @@ bash run_server_stage_cli.sh --stage 0
|
||||
bash run_server_stage_cli.sh --stage 1
|
||||
```
|
||||
|
||||
To use a custom deploy YAML (note: `--stage-configs-path` is deprecated in favor of `--deploy-config`):
|
||||
To use a custom deploy YAML:
|
||||
|
||||
```bash
|
||||
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \
|
||||
|
||||
@@ -11,8 +11,8 @@ memory, sampling params) live in `vllm_omni/deploy/glm_image.yaml`.
|
||||
vllm serve zai-org/GLM-Image --omni --port 8091
|
||||
```
|
||||
|
||||
The config system auto-detects the pipeline from the model's `model_index.json` — no
|
||||
manual `--stage-configs-path` or `--deploy-config` needed.
|
||||
The config system auto-detects the pipeline from the model's `model_index.json`, so
|
||||
no manual `--deploy-config` is needed.
|
||||
|
||||
By default, stage 0 (AR) runs on GPU 0 and stage 1 (Diffusion) on GPU 1. To colocate
|
||||
both stages on a single GPU, override per stage:
|
||||
|
||||
@@ -15,9 +15,9 @@ Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/
|
||||
vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091
|
||||
```
|
||||
|
||||
If you have custom stage configs file, launch the server with command below
|
||||
To use a custom deploy config, launch the server with the command below:
|
||||
```bash
|
||||
vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
|
||||
vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --deploy-config /path/to/deploy_config.yaml
|
||||
```
|
||||
|
||||
### Send Multi-modal Request
|
||||
@@ -182,7 +182,7 @@ The script supports the following arguments:
|
||||
- `--model`: Model name/path (default: Qwen/Qwen2.5-Omni-7B)
|
||||
- `--server-port`: Port for vLLM server (default: 8091)
|
||||
- `--gradio-port`: Port for Gradio demo (default: 7861)
|
||||
- `--stage-configs-path`: Path to custom stage configs YAML file (optional)
|
||||
- `--deploy-config`: Path to a custom deploy config YAML file (optional)
|
||||
- `--server-host`: Host for vLLM server (default: 0.0.0.0)
|
||||
- `--gradio-ip`: IP for Gradio demo (default: 127.0.0.1)
|
||||
- `--share`: Share Gradio demo publicly (creates a public link)
|
||||
@@ -195,9 +195,9 @@ The script supports the following arguments:
|
||||
vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091
|
||||
```
|
||||
|
||||
If you have custom stage configs file:
|
||||
To use a custom deploy config:
|
||||
```bash
|
||||
vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
|
||||
vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --deploy-config /path/to/deploy_config.yaml
|
||||
```
|
||||
|
||||
**Step 2: Run the Gradio demo**
|
||||
|
||||
@@ -306,7 +306,7 @@ The gradio script supports the following arguments:
|
||||
``````
|
||||
??? abstract "qwen3_omni_moe_thinking.yaml"
|
||||
``````yaml
|
||||
--8<-- "examples/online_serving/qwen3_omni/qwen3_omni_moe_thinking.yaml"
|
||||
--8<-- "vllm_omni/deploy/qwen3_omni_moe_thinking.yaml"
|
||||
``````
|
||||
??? abstract "run_curl_multimodal_generation.sh"
|
||||
``````sh
|
||||
|
||||
@@ -422,12 +422,19 @@ curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
Raw PCM streaming requires `stream_format="audio"`, `response_format="pcm"`, and `async_chunk: true` on the stage config (default in `qwen3_tts.yaml`). `speed` is not supported when streaming.
|
||||
|
||||
### Streaming WebSocket
|
||||
The `/v1/audio/speech/stream` endpoint accepts text incrementally, splits it at sentence boundaries, and emits one PCM stream per sentence:
|
||||
The `/v1/audio/speech/stream` endpoint accepts text incrementally and synthesizes the buffered text as one continuous request on `input.done`:
|
||||
```bash
|
||||
python qwen3_tts/streaming_speech_client.py --text "Hello world. How are you? I am fine."
|
||||
python qwen3_tts/streaming_speech_client.py --text "..." --simulate-stt --stt-delay 0.1
|
||||
```
|
||||
|
||||
`input.done` flushes without closing, so repeating `--text` synthesizes several utterances over one connection:
|
||||
```bash
|
||||
python qwen3_tts/streaming_speech_client.py \
|
||||
--text "First utterance." \
|
||||
--text "Second utterance, same connection."
|
||||
```
|
||||
|
||||
### Gradio demos
|
||||
```bash
|
||||
./qwen3_tts/run_gradio_demo.sh # CustomVoice (default)
|
||||
|
||||
@@ -13,7 +13,6 @@ This example demonstrates how to deploy text-to-video models for online video ge
|
||||
| Wan2.1 T2V (14B) | `Wan-AI/Wan2.1-T2V-14B-Diffusers` |
|
||||
| Wan2.2 T2V | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` |
|
||||
| LTX-2 | `Lightricks/LTX-2` |
|
||||
| LTX-2.3 | `diffusers/LTX-2.3-Diffusers` |
|
||||
|
||||
## Wan2.2 T2V
|
||||
|
||||
@@ -192,7 +191,7 @@ curl -X POST http://localhost:8091/v1/videos \
|
||||
| `negative_prompt` | str | None | Negative prompt |
|
||||
| `width` | int | None | Video width in pixels |
|
||||
| `height` | int | None | Video height in pixels |
|
||||
| `num_frames` | int | None | Number of frames to generate |
|
||||
| `num_frames` | int | 1 | Number of frames to generate; set explicitly for video generation |
|
||||
| `fps` | int | None | Frames per second for output video |
|
||||
| `num_inference_steps` | int | None | Number of denoising steps |
|
||||
| `guidance_scale` | float | None | CFG guidance scale (low-noise stage) |
|
||||
@@ -290,101 +289,12 @@ done
|
||||
|
||||
## LTX-2
|
||||
|
||||
### Start Server
|
||||
|
||||
#### Basic Start
|
||||
|
||||
```bash
|
||||
vllm serve Lightricks/LTX-2 --omni --port 8098 \
|
||||
--enforce-eager --boundary-ratio 1.0
|
||||
vllm serve Lightricks/LTX-2 --omni --port 8098
|
||||
```
|
||||
|
||||
For multi-GPU memory reduction, you can enable HSDP:
|
||||
|
||||
```bash
|
||||
vllm serve Lightricks/LTX-2 --omni --port 8098 \
|
||||
--enforce-eager --boundary-ratio 1.0 \
|
||||
--use-hsdp --hsdp-shard-size 2
|
||||
```
|
||||
|
||||
#### Start with Optimization Presets
|
||||
|
||||
Use the LTX-2 startup script with built-in optimization presets:
|
||||
|
||||
```bash
|
||||
# Baseline (1 GPU, eager)
|
||||
bash run_server_ltx2.sh baseline
|
||||
|
||||
# 4-GPU Ulysses sequence parallelism (lossless)
|
||||
bash run_server_ltx2.sh ulysses4
|
||||
|
||||
# Cache-DiT lossy acceleration (1 GPU, ~1.4× speedup)
|
||||
bash run_server_ltx2.sh cache-dit
|
||||
|
||||
# Best combo: 4-GPU Ulysses SP + Cache-DiT (~2.2× speedup)
|
||||
bash run_server_ltx2.sh best-combo
|
||||
```
|
||||
|
||||
#### Optimization Benchmarks
|
||||
|
||||
Benchmarked on H800, online serving (480×768, 41 frames, 20 steps, `seed=42`).
|
||||
"Inference" is the server-reported inference time; excludes HTTP/poll overhead.
|
||||
|
||||
| Preset | Server Command | Inference (s) | Speedup | Type |
|
||||
|--------|---------------|---------------|---------|------|
|
||||
| `baseline` | `--enforce-eager` | 10.3 | 1.00× | — |
|
||||
| `compile` | *(default, no --enforce-eager)* | ~10.3 (warm) | ~1.00× | Lossless |
|
||||
| `ulysses4` | `--enforce-eager --usp 4` | ~10.3 | ~1.00× | Lossless |
|
||||
| `cache-dit` | `--enforce-eager --cache-backend cache_dit` | 7.4 avg | ~1.4× | Lossy |
|
||||
| `best-combo` | `--enforce-eager --usp 4 --cache-backend cache_dit` | 4.7 avg | **~2.2×** | Lossless + Lossy |
|
||||
|
||||
**Observations**:
|
||||
- **torch.compile**: On H800, warm-request inference time matches the eager baseline (~10.3s).
|
||||
The first request pays ~6s compilation overhead. Benefit depends on model architecture and GPU.
|
||||
- **Ulysses SP (4 GPU)**: No measurable speedup alone for 41-frame generation at this resolution.
|
||||
Communication overhead outweighs gains at this sequence length.
|
||||
- **Cache-DiT**: Inference varies per request (6–10s) due to dynamic caching decisions.
|
||||
Average is ~7.4s (~1.4× speedup) with slight quality tradeoff.
|
||||
- **Best combo**: 4-GPU Ulysses SP + Cache-DiT synergize well — Cache-DiT reduces per-step
|
||||
computation, making the communication overhead of Ulysses SP worthwhile. Average ~4.7s
|
||||
(~2.2× speedup).
|
||||
- **FP8 quantization**: Reduces VRAM but does not speed up LTX-2 on H800 (compute-bound).
|
||||
|
||||
**Deployment Recommendations**:
|
||||
- For **production with quality priority**: use `baseline` with `--enforce-eager`
|
||||
- For **maximum throughput** (4 GPUs, quality tradeoff): use `best-combo` (~2.2× speedup)
|
||||
- For **single-GPU throughput**: use `cache-dit` (~1.4× speedup)
|
||||
- `--enforce-eager` is recommended to avoid torch.compile warmup latency on first request
|
||||
|
||||
### Send Requests (curl)
|
||||
|
||||
```bash
|
||||
# Using the provided script
|
||||
bash run_curl_ltx2.sh
|
||||
|
||||
# Or directly
|
||||
curl -sS -X POST http://localhost:8098/v1/videos \
|
||||
-H "Accept: application/json" \
|
||||
-F "prompt=A serene lakeside sunrise with mist over the water." \
|
||||
-F "width=768" \
|
||||
-F "height=480" \
|
||||
-F "num_frames=41" \
|
||||
-F "fps=24" \
|
||||
-F "num_inference_steps=20" \
|
||||
-F "guidance_scale=3.0" \
|
||||
-F "seed=42"
|
||||
```
|
||||
|
||||
## LTX-2.3
|
||||
|
||||
LTX-2.3 uses a Diffusers-format checkpoint and a dedicated pipeline class:
|
||||
|
||||
```bash
|
||||
vllm serve diffusers/LTX-2.3-Diffusers --omni --port 8098 \
|
||||
--model-class-name LTX23Pipeline
|
||||
```
|
||||
|
||||
Send requests through the same `/v1/videos` API shown in the LTX-2 example.
|
||||
See the [LTX-2 recipe](../../../../recipes/LTX/LTX-2.md) for all checkpoints,
|
||||
pipeline selection, requests, defaults, and advanced options.
|
||||
|
||||
## Example materials
|
||||
|
||||
|
||||
@@ -69,19 +69,19 @@ The behaviors and conclusions above are summarized from current fault injection
|
||||
|
||||
### Example: `SIGTERM` fault injection log (Qwen3-Omni)
|
||||
|
||||
The following full log is from a real run where a `SIGTERM` signal was sent to the serve root process.
|
||||
The following log is based on a real run where a `SIGTERM` signal was sent to the serve root process. Its command line has been normalized to the current `--deploy-config` option.
|
||||
|
||||
```text
|
||||
Launching OmniServer with: /workspace/.venv/bin/python3 -m vllm_omni.entrypoints.cli.main serve Qwen/Qwen3-Omni-30B-A3B-Instruct --host 127.0.0.1 --port 60675 --omni --async-chunk --stage-init-timeout 600 --init-timeout 900 --log-stats --stage-configs-path vllm-omni/vllm_omni/deploy/qwen3_omni_moe.yaml
|
||||
Launching OmniServer with: /workspace/.venv/bin/python3 -m vllm_omni.entrypoints.cli.main serve Qwen/Qwen3-Omni-30B-A3B-Instruct --host 127.0.0.1 --port 60675 --omni --async-chunk --stage-init-timeout 600 --init-timeout 900 --log-stats --deploy-config vllm-omni/vllm_omni/deploy/qwen3_omni_moe.yaml
|
||||
Server ready on 127.0.0.1:60675 (OmniServer startup took 206.023s)
|
||||
OmniServer started successfully
|
||||
[reliability][process-kill] current_server_proc pid=556855 name=python3 cmdline=/workspace/.venv/bin/python3 -m vllm_omni.entrypoints.cli.main serve Qwen/Qwen3-Omni-30B-A3B-Instruct --host 127.0.0.1 --port 60675 --omni --async-chunk --stage-init-timeout 600 --init-timeout 900 --log-stats --stage-configs-path vllm-omni/vllm_omni/deploy/qwen3_omni_moe.yaml
|
||||
[reliability][process-kill] current_server_proc pid=556855 name=python3 cmdline=/workspace/.venv/bin/python3 -m vllm_omni.entrypoints.cli.main serve Qwen/Qwen3-Omni-30B-A3B-Instruct --host 127.0.0.1 --port 60675 --omni --async-chunk --stage-init-timeout 600 --init-timeout 900 --log-stats --deploy-config vllm-omni/vllm_omni/deploy/qwen3_omni_moe.yaml
|
||||
[reliability][process-kill] current_server_proc pid=557118 name=python3 cmdline=/workspace/.venv/bin/python3 -c from multiprocessing.resource_tracker import main;main(57)
|
||||
[reliability][process-kill] current_server_proc pid=557119 name=VLLM::StageEngineCoreProc_noid_replica0_DP0 cmdline=VLLM::StageEngineCoreProc_noid_replica0_DP0
|
||||
[reliability][process-kill] current_server_proc pid=557122 name=VLLM::StageEngineCoreProc_noid_replica0_DP0 cmdline=VLLM::StageEngineCoreProc_noid_replica0_DP0
|
||||
[reliability][process-kill] current_server_proc pid=558201 name=VLLM::StageEngineCoreProc_noid_replica0_DP0 cmdline=VLLM::StageEngineCoreProc_noid_replica0_DP0
|
||||
|
||||
[reliability][process-kill] root-kill pid=556855 name=python3 signal=SIGTERM cmdline=/workspace/.venv/bin/python3 -m vllm_omni.entrypoints.cli.main serve Qwen/Qwen3-Omni-30B-A3B-Instruct --host 127.0.0.1 --port 60675 --omni --async-chunk --stage-init-timeout 600 --init-timeout 900 --log-stats --stage-configs-path vllm-omni/vllm_omni/deploy/qwen3_omni_moe.yaml
|
||||
[reliability][process-kill] root-kill pid=556855 name=python3 signal=SIGTERM cmdline=/workspace/.venv/bin/python3 -m vllm_omni.entrypoints.cli.main serve Qwen/Qwen3-Omni-30B-A3B-Instruct --host 127.0.0.1 --port 60675 --omni --async-chunk --stage-init-timeout 600 --init-timeout 900 --log-stats --deploy-config vllm-omni/vllm_omni/deploy/qwen3_omni_moe.yaml
|
||||
[0;36m(APIServer pid=556855)[0;0m INFO 05-19 12:24:51 [omni_base.py:463] [AsyncOmni] Shutting down
|
||||
[0;36m(APIServer pid=556855)[0;0m INFO 05-19 12:24:51 [async_omni_engine.py:2397] [AsyncOmniEngine] Shutting down Orchestrator
|
||||
[0;36m(APIServer pid=556855)[0;0m INFO 05-19 12:24:51 [orchestrator.py:351] [Orchestrator] Received shutdown signal
|
||||
|
||||
@@ -8,5 +8,4 @@ mimo_audio: MiMo-Audio
|
||||
ming_flash_omni: Ming-flash-omni 2.0
|
||||
qwen2_5_omni: Qwen2.5-Omni
|
||||
qwen3_omni: Qwen3-Omni
|
||||
sensenova_u1: SenseNova-U1
|
||||
step_audio2: Step-Audio2
|
||||
@@ -0,0 +1,114 @@
|
||||
# Audex (Nemotron-Labs-Audex-2B) offline inference
|
||||
|
||||
One checkpoint, four tasks — each script builds its own `Omni` engine with
|
||||
the right deploy yaml, so they run as plain `python <script>` with no extra
|
||||
flags. Pass the HF repo ROOT as `--model` (default
|
||||
`nvidia/Nemotron-Labs-Audex-2B`); per-stage subfolders and download subsets
|
||||
resolve automatically.
|
||||
|
||||
| script | pipeline (`vllm_omni/deploy/<name>.yaml`) | audio in | text out | speech out | general audio out |
|
||||
|---|---|---|---|---|---|
|
||||
| `text_to_speech.py` | `audex_tts` | ❌ | ❌ | ✅ | ❌ |
|
||||
| `text_to_audio.py` | `audex_tta` | ❌ | ❌ | ❌ | ✅ |
|
||||
| `audio_qa.py` | `audex_thinker_only` | ✅ | ✅ | ❌ | ❌ |
|
||||
| `speech_to_speech.py` | `audex_s2s` | ✅ | ✅ | ✅ | ❌ |
|
||||
|
||||
Shared plumbing (corpus/WAV IO, tokenizer loading, the CFG cond/uncond
|
||||
pair contract) lives in `common.py`; each task script keeps only its own
|
||||
prompt recipe and flow. Note `text_to_audio.py` here drives the Audex
|
||||
two-stage AR pipeline — for diffusion text-to-audio models (Stable Audio
|
||||
Open, AudioX) use the generic
|
||||
`examples/offline_inference/text_to_audio/text_to_audio.py` instead.
|
||||
|
||||
## text_to_speech.py — text → speech
|
||||
|
||||
Thinker generates `<speechcodec_N>` tokens; the streaming causal decoder
|
||||
turns them into 16 kHz WAVs (one file per prompt, slugified filename).
|
||||
|
||||
```bash
|
||||
# three built-in demo sentences -> results/audex_wavs/
|
||||
python examples/offline_inference/audex/text_to_speech.py
|
||||
|
||||
# your own texts, with classifier-free guidance (official quality setting)
|
||||
python examples/offline_inference/audex/text_to_speech.py \
|
||||
--texts "Hello there." "Nice to meet you." --cfg-scale 1.5 \
|
||||
--output-dir results/my_tts
|
||||
```
|
||||
|
||||
Key flags: `--texts` / `--texts-file` (TSV `utt_id<TAB>text`),
|
||||
`--cfg-scale` (default 1.0 = off; 1.5 recommended — guided requests run
|
||||
one at a time, each paired with a length-matched null prompt),
|
||||
`--temperature` (overrides the deploy yaml's stage-0 default).
|
||||
|
||||
## text_to_audio.py — caption → general audio (sound effects)
|
||||
|
||||
Same thinker over the interleaved 4-codebook `<audiocodec_N>` RVQ block,
|
||||
decoded by the external XCodec1 checkpoint. CFG is effectively mandatory
|
||||
(default scale 3.0). Clips are capped at `--codec-cap 4000` codec tokens
|
||||
(10 s).
|
||||
|
||||
```bash
|
||||
# three built-in demo captions -> results/audex_tta_wavs/
|
||||
python examples/offline_inference/audex/text_to_audio.py
|
||||
|
||||
python examples/offline_inference/audex/text_to_audio.py \
|
||||
--captions "Thunder rolling across a valley." --output-dir results/my_tta
|
||||
```
|
||||
|
||||
XCodec1 resolves from `--xcodec1-path`, the `XCODEC1_PATH` env var, or the
|
||||
default `hf-audio/xcodec-hubert-general-balanced` repo (downloaded on
|
||||
first use).
|
||||
|
||||
## audio_qa.py — audio (+ instruction) → text
|
||||
|
||||
Single-stage audio understanding on the full checkpoint (NV-Whisper
|
||||
encoder + projector + LM). The default question is ASR
|
||||
("Transcribe the input speech."); without `--audio-files`, vLLM's public
|
||||
`mary_had_lamb` asset is transcribed.
|
||||
|
||||
```bash
|
||||
# transcribe the built-in demo asset
|
||||
python examples/offline_inference/audex/audio_qa.py
|
||||
|
||||
# your own clips / free-form QA
|
||||
python examples/offline_inference/audex/audio_qa.py \
|
||||
--audio-files a.wav b.wav --question "What language is being spoken?"
|
||||
```
|
||||
|
||||
## speech_to_speech.py — spoken question → spoken answer
|
||||
|
||||
The official three-pass cascade over ONE `audex_s2s` deployment: ASR
|
||||
(audio → transcript, text-final) → chat (transcript → answer, text-final)
|
||||
→ TTS (answer → speech through the streaming decoder). Only the TTS pass
|
||||
touches the codec path. Without `--audio-file`, the `mary_had_lamb` asset
|
||||
is used as the spoken input.
|
||||
|
||||
```bash
|
||||
# full cascade on the built-in demo asset -> results/audex_s2s_answer.wav
|
||||
python examples/offline_inference/audex/speech_to_speech.py
|
||||
|
||||
python examples/offline_inference/audex/speech_to_speech.py \
|
||||
--audio-file question.wav --output results/answer.wav --cfg-scale 1.5
|
||||
```
|
||||
|
||||
`--cfg-scale` (default 1.5) applies to the TTS pass only; 1.0 disables it.
|
||||
|
||||
Online-serving counterparts (server + HTTP clients) live in
|
||||
`examples/online_serving/audex/`.
|
||||
|
||||
## 30B-A3B (nvidia/Nemotron-Labs-Audex-30B-A3B)
|
||||
|
||||
The same four scripts serve the 30B MoE checkpoint (hybrid Mamba + MoE
|
||||
NemotronH thinker; identical token space and decoder). The 30B REQUIRES an
|
||||
explicit 30B deploy yaml — the model root's default resolution lands on the
|
||||
2B-tuned configs:
|
||||
|
||||
python examples/offline_inference/audex/text_to_speech.py \
|
||||
--model nvidia/Nemotron-Labs-Audex-30B-A3B \
|
||||
--deploy-config vllm_omni/deploy/audex_tts_30b.yaml
|
||||
|
||||
Per-mode yamls: `audex_{tts,tta,thinker_only,s2s}_30b.yaml`. Defaults fit a
|
||||
single H100 80 GB (thinker + decoder share the card; prefix caching is off —
|
||||
hybrid Mamba); if long sequences OOM, set `tensor_parallel_size: 2` on
|
||||
stage 0. First run downloads ~60 GB. Verified on a single H100 80 GB: ~61 GiB
|
||||
weights + healthy KV headroom; TP2 was not needed.
|
||||
@@ -0,0 +1,93 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Offline Audex (Nemotron-Labs-Audex-2B) audio understanding example.
|
||||
|
||||
Speech (or general audio) + text instruction → text, through the
|
||||
single-stage ``audex_thinker_only`` pipeline: NV-Whisper
|
||||
encoder + relu2 projector + 2B dense LM on ``checkpoint_folder_full``.
|
||||
|
||||
The prompt format mirrors the official audioqa script: a ChatML user turn
|
||||
holding the ``<so_embedding>`` placeholder (expanded by the processor to
|
||||
750 embedding positions per 30 s clip) plus the instruction, with a closed
|
||||
``<think></think>`` priming.
|
||||
|
||||
Examples (without --audio-files, vLLM's public ``mary_had_lamb`` asset is
|
||||
transcribed):
|
||||
|
||||
# Transcribe WAVs (ASR):
|
||||
python examples/offline_inference/audex/audio_qa.py \\
|
||||
--audio-files a.wav b.wav
|
||||
|
||||
# Free-form audio QA:
|
||||
python examples/offline_inference/audex/audio_qa.py \\
|
||||
--audio-files a.wav --question "What language is being spoken?"
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
|
||||
from common import default_deploy_config, load_audio_file, request_index
|
||||
from vllm.assets.audio import AudioAsset
|
||||
|
||||
from vllm_omni import Omni
|
||||
|
||||
ASR_QUESTION = "Transcribe the input speech."
|
||||
|
||||
|
||||
def build_prompt(question: str) -> str:
|
||||
return f"<|im_start|>user\n<so_embedding>\n{question}<|im_end|>\n<|im_start|>assistant\n<think></think>"
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Offline Audex audio understanding")
|
||||
parser.add_argument("--model", type=str, default="nvidia/Nemotron-Labs-Audex-2B")
|
||||
parser.add_argument(
|
||||
"--audio-files",
|
||||
type=str,
|
||||
nargs="+",
|
||||
default=None,
|
||||
help="Input clips. Defaults to vLLM's mary_had_lamb audio asset.",
|
||||
)
|
||||
parser.add_argument("--question", type=str, default=ASR_QUESTION)
|
||||
parser.add_argument(
|
||||
"--deploy-config",
|
||||
type=str,
|
||||
# The model root's default deploy yaml is the TTS pipeline; audio
|
||||
# understanding needs the single-stage thinker-only pipeline.
|
||||
default=default_deploy_config("audex_thinker_only.yaml"),
|
||||
help="Deploy yaml (defaults to the audex_thinker_only pipeline).",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
engine = Omni(model=args.model, deploy_config=args.deploy_config, trust_remote_code=True)
|
||||
|
||||
prompt_text = build_prompt(args.question)
|
||||
if args.audio_files:
|
||||
clips = [load_audio_file(path) for path in args.audio_files]
|
||||
else:
|
||||
clips = [AudioAsset("mary_had_lamb").audio_and_sample_rate]
|
||||
prompts = []
|
||||
for audio, sr in clips:
|
||||
prompts.append(
|
||||
{
|
||||
"prompt": prompt_text,
|
||||
"multi_modal_data": {"audio": (audio, sr)},
|
||||
}
|
||||
)
|
||||
|
||||
outputs = engine.generate(prompts)
|
||||
|
||||
ordered = sorted(outputs, key=request_index)
|
||||
names = [Path(p).name for p in args.audio_files] if args.audio_files else ["mary_had_lamb"]
|
||||
for name, req_output in zip(names, ordered):
|
||||
text = req_output.outputs[0].text if req_output.outputs else ""
|
||||
print(f"{name}: {text.strip()}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,149 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Shared helpers for the offline Audex examples in this directory.
|
||||
|
||||
Import style follows the repo's example convention (same-directory module,
|
||||
see ``text_to_speech/qwen3_tts/tts_common.py``): scripts run as files, so
|
||||
``from common import ...`` resolves via the script's own directory.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import re
|
||||
from collections.abc import Callable
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
import torch
|
||||
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.model_executor.models.audex.prompt import build_null_prompt
|
||||
|
||||
SAMPLE_RATE = 16_000
|
||||
_DEPLOY_DIR = Path(__file__).resolve().parents[3] / "vllm_omni" / "deploy"
|
||||
|
||||
# Guided decoding sharpens the distribution, so the unguided temperature
|
||||
# (0.1) adds excess sampling noise under CFG. Measured on the en-24 gate
|
||||
# corpus at cfg 1.5: temp 0.1 -> CER 7.31%, temp 0.05 -> CER 6.87% (vs the
|
||||
# unguided baseline 7.24%).
|
||||
GUIDED_TEMPERATURE = 0.05
|
||||
|
||||
|
||||
def default_deploy_config(name: str) -> str:
|
||||
"""Absolute path of a deploy yaml bundled with vllm_omni."""
|
||||
return str(_DEPLOY_DIR / name)
|
||||
|
||||
|
||||
def slugify(text: str, fallback: str = "prompt") -> str:
|
||||
slug = re.sub(r"\s+", "_", text.strip().lower())
|
||||
slug = re.sub(r"[^a-z0-9_]+", "", slug)
|
||||
return slug[:48] or fallback
|
||||
|
||||
|
||||
def load_corpus(
|
||||
corpus_file: str | None,
|
||||
texts: list[str] | None,
|
||||
default_texts: tuple[str, ...],
|
||||
) -> list[tuple[str, str]]:
|
||||
"""(utt_id, text) pairs from a TSV corpus file, CLI texts, or defaults."""
|
||||
if corpus_file:
|
||||
corpus = []
|
||||
for line in Path(corpus_file).read_text().splitlines():
|
||||
if not line.strip():
|
||||
continue
|
||||
utt, text = line.split("\t", 1)
|
||||
corpus.append((utt, text.strip()))
|
||||
return corpus
|
||||
resolved = texts if texts else list(default_texts)
|
||||
return [(slugify(t), t) for t in resolved]
|
||||
|
||||
|
||||
def load_audio_file(path: str) -> tuple[np.ndarray, int]:
|
||||
"""Mono float32 samples + native sample rate (processors resample)."""
|
||||
audio, sr = sf.read(path, dtype="float32")
|
||||
if audio.ndim == 2:
|
||||
audio = audio.mean(axis=1)
|
||||
return audio, sr
|
||||
|
||||
|
||||
def extract_pcm(multimodal_output: dict) -> torch.Tensor:
|
||||
"""Flatten one request's audio payload to a 1-D float32 CPU 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 write_wav(path: Path, pcm: torch.Tensor, sample_rate: int = SAMPLE_RATE) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
arr = (np.clip(pcm.numpy(), -1.0, 1.0) * 32767.0).astype(np.int16)
|
||||
sf.write(str(path), arr, sample_rate, format="WAV", subtype="PCM_16")
|
||||
|
||||
|
||||
def load_audex_tokenizer(model: str, profile: str = "tts", folder: str = "checkpoint_folder_audiogen"):
|
||||
"""Thinker tokenizer for CFG/TTA prompt building (resolves the snapshot)."""
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
from vllm_omni.model_executor.models.audex.checkpoint import ensure_audex_snapshot
|
||||
|
||||
root = ensure_audex_snapshot(model, profile=profile)
|
||||
return AutoTokenizer.from_pretrained(str(Path(root) / folder), trust_remote_code=True)
|
||||
|
||||
|
||||
def request_index(req_output) -> int:
|
||||
"""Submission index from the numeric request-id prefix.
|
||||
|
||||
``Omni.generate`` may return outputs out of submission order, and a
|
||||
lexicographic sort would misplace "10_..." before "2_...".
|
||||
"""
|
||||
match = re.search(r"(\d+)", str(req_output.request_id))
|
||||
return int(match.group(1)) if match else 0
|
||||
|
||||
|
||||
def guided_stage_params(
|
||||
engine: Omni,
|
||||
cond_prompt: str,
|
||||
tokenizer,
|
||||
cfg_scale: float,
|
||||
pair_id: str,
|
||||
*,
|
||||
temperature: float | None = None,
|
||||
extra_args: dict[str, Any] | None = None,
|
||||
null_prompt_builder: Callable[[str, Any], str] = build_null_prompt,
|
||||
) -> list:
|
||||
"""Stage sampling params carrying the CFG pair contract for ONE request.
|
||||
|
||||
Clones the engine's resolved defaults (never mutates the shared list),
|
||||
applies ``extra_args`` unconditionally (e.g. the TTA RVQ contract), and
|
||||
attaches the cond/uncond pair contract when ``cfg_scale > 1.0``. Guided
|
||||
requests must go one at a time: each needs its own pair id and a
|
||||
length-matched null prompt built from its own text.
|
||||
"""
|
||||
params = copy.deepcopy(engine.resolve_sampling_params_list(None))
|
||||
stage0 = params[0]
|
||||
if temperature is not None:
|
||||
stage0.temperature = temperature
|
||||
if stage0.extra_args is None:
|
||||
stage0.extra_args = {}
|
||||
if extra_args:
|
||||
stage0.extra_args.update(extra_args)
|
||||
if cfg_scale > 1.0:
|
||||
stage0.extra_args.update(
|
||||
{
|
||||
"cfg_scale": float(cfg_scale),
|
||||
"cfg_role": "cond",
|
||||
"cfg_pair_id": pair_id,
|
||||
"cfg_null_prompt": null_prompt_builder(cond_prompt, tokenizer),
|
||||
}
|
||||
)
|
||||
return params
|
||||
@@ -0,0 +1,277 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Offline Audex (Nemotron-Labs-Audex-2B) cascaded speech-to-speech example.
|
||||
|
||||
The official three-pass cascade over ONE ``audex_s2s``
|
||||
deployment:
|
||||
|
||||
1. ASR pass (audio + "Transcribe the input speech." → transcript, text)
|
||||
2. chat pass (transcript → answer, text)
|
||||
3. TTS pass (answer → speech, audio via the streaming causal decoder)
|
||||
|
||||
Text passes carry ``modalities: ["text"]`` and finish at stage 0; only the
|
||||
TTS pass (``modalities: ["audio"]``) streams codec frames into code2wav.
|
||||
|
||||
Example (without --audio-file, vLLM's public ``mary_had_lamb`` asset is
|
||||
used as the spoken input):
|
||||
|
||||
python examples/offline_inference/audex/speech_to_speech.py \\
|
||||
--audio-file question.wav --output results/answer.wav
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import copy
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
from common import (
|
||||
SAMPLE_RATE,
|
||||
default_deploy_config,
|
||||
guided_stage_params,
|
||||
load_audex_tokenizer,
|
||||
load_audio_file,
|
||||
write_wav,
|
||||
)
|
||||
from vllm.assets.audio import AudioAsset
|
||||
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.model_executor.models.audex.prompt import build_cond_prompt
|
||||
|
||||
ASR_QUESTION = "Transcribe the input speech."
|
||||
SYSTEM_PROMPT = "You are a helpful and harmless assistant.\n\nYou are not allowed to use any tools."
|
||||
# The official chat-pass recipe below (SYSTEM_PROMPT + TEXT_INSTRUCTION
|
||||
# prefix, open <think> priming, temp 1.0/top_p 0.95, seed retry) applies to
|
||||
# the 30B-A3B ONLY (see is_30b()/_chat_prompt below and
|
||||
# tests/e2e/offline_inference/test_audex_s2s.py, which mirrors this same
|
||||
# split). It is load-bearing on the 30B: without the formatting-instruction
|
||||
# prefix a bare transcript turn is treated as a TTS request and the model
|
||||
# emits literal <speechcodec_N> tokens instead of a text answer (the system
|
||||
# prompt alone does not prevent this). The 2B must NOT use this recipe: it
|
||||
# is long-verified stable on a bare prompt + engine-default sampling, while
|
||||
# the official temp-1.0 seed-0 draw hallucinates off-language text (yet
|
||||
# closes </think> cleanly, so the seed retry never fires) and fails the
|
||||
# round-trip WER gate — measured, reproduced byte-identically.
|
||||
TEXT_INSTRUCTION = (
|
||||
"SYSTEM & FORMATTING INSTRUCTION: You are Audex, created by NVIDIA based on the Nemotron-Cascade-2 architecture. "
|
||||
"You may output your reasoning, followed by your final response. "
|
||||
"You may format your reasoning block however you like. "
|
||||
"However, your final response must strictly follow these rules: "
|
||||
"* Write in plain, unformatted prose like a book or newspaper article. "
|
||||
"* Do not use markdown, bullet points, lists, or headers. "
|
||||
"* Standard numbers, abbreviations, and symbols are acceptable. "
|
||||
"* [CRITICAL] You must press enter after every single sentence, placing each sentence on its own separate line."
|
||||
)
|
||||
|
||||
|
||||
def is_30b(model: str) -> bool:
|
||||
# Users pass arbitrary repo IDs or local snapshot paths, so this is a
|
||||
# heuristic (not an exact match against the official 30B-A3B repo ID)
|
||||
# — good enough for an example script; keep in sync with the same
|
||||
# branch in examples/online_serving/audex/client.py and
|
||||
# tests/e2e/offline_inference/test_audex_s2s.py.
|
||||
return "30b" in model.lower()
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Offline Audex cascaded S2S")
|
||||
parser.add_argument("--model", type=str, default="nvidia/Nemotron-Labs-Audex-2B")
|
||||
parser.add_argument(
|
||||
"--audio-file",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Spoken question (WAV). Defaults to vLLM's mary_had_lamb audio asset.",
|
||||
)
|
||||
parser.add_argument("--output", type=str, default="results/audex_s2s_answer.wav")
|
||||
parser.add_argument(
|
||||
"--deploy-config",
|
||||
type=str,
|
||||
# The model root's default deploy yaml is the TTS-only pipeline; the
|
||||
# cascade needs the audio-capable 2-stage full pipeline.
|
||||
default=default_deploy_config("audex_s2s.yaml"),
|
||||
help="Deploy yaml (defaults to the audex_s2s pipeline).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--cfg-scale",
|
||||
type=float,
|
||||
default=1.5,
|
||||
help="CFG strength for the TTS pass (1.0 disables; official setting 1.5).",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def _chat_prompt(user_text: str) -> str:
|
||||
# Official chat-pass recipe: system prompt + formatting instruction
|
||||
# prefixed to the user text, thinking ENABLED via an open <think> block
|
||||
# (the official demo's web UI ships "Enable reasoning" checked, i.e.
|
||||
# build_chat_prompt(..., enable_thinking=True); chat_template.jinja then
|
||||
# renders "<|im_start|>assistant\n<think>\n"). The instruction invites
|
||||
# the model to reason before answering — priming the block open keeps
|
||||
# that reasoning inside <think>...</think> so _extract_final_response
|
||||
# below can drop it; with a closed block the 30B leaks the reasoning as
|
||||
# plain text in front of the answer.
|
||||
return (
|
||||
f"<|im_start|>system\n{SYSTEM_PROMPT}<|im_end|>\n"
|
||||
f"<|im_start|>user\n{TEXT_INSTRUCTION}\n\n{user_text}<|im_end|>\n"
|
||||
"<|im_start|>assistant\n<think>\n"
|
||||
)
|
||||
|
||||
|
||||
def _asr_prompt() -> str:
|
||||
return f"<|im_start|>user\n<so_embedding>\n{ASR_QUESTION}<|im_end|>\n<|im_start|>assistant\n<think></think>"
|
||||
|
||||
|
||||
def _extract_final_response(text: str) -> str:
|
||||
# Mirror the official demo's extract_spoken_answer
|
||||
# (cascaded_s2s_web_server.py): keep only what follows the LAST closing
|
||||
# think tag, falling back to the full text when the model skipped
|
||||
# reasoning or never closed the block. Benign no-op for the 2B, whose
|
||||
# answers carry no think block.
|
||||
if "</think>" in text:
|
||||
answer = text.rsplit("</think>", 1)[-1].strip()
|
||||
if answer:
|
||||
return answer
|
||||
return text.strip()
|
||||
|
||||
|
||||
def _clean(text: str) -> str:
|
||||
# The full checkpoint prefixes transcripts with a language tag sentence
|
||||
# and quotes the content; extract the quoted payload when present. Span
|
||||
# first-to-last quote (not split on every apostrophe) so contractions
|
||||
# inside the transcript (e.g. "don't") survive intact.
|
||||
text = text.strip()
|
||||
match = re.search(r"'(.*)'", text, flags=re.S)
|
||||
if match and match.group(1).strip():
|
||||
return match.group(1).strip()
|
||||
return text
|
||||
|
||||
|
||||
def _text_of(outputs) -> str:
|
||||
(req_output,) = outputs
|
||||
return (req_output.outputs[0].text or "").strip()
|
||||
|
||||
|
||||
CHAT_MAX_TOKENS = 2048 # official --text-max-tokens default
|
||||
# The official demo samples this pass unseeded; some trajectories deliberate
|
||||
# past the token budget without ever closing the think block, so retry a few
|
||||
# deterministic seeds until one closes (see the chat pass below).
|
||||
CHAT_SEEDS = (0, 1, 2)
|
||||
|
||||
|
||||
def _chat_params(engine: Omni, seed: int):
|
||||
# Official demo text-pass sampling (cascaded_s2s_web_server.py defaults:
|
||||
# --text-temperature 1.0, --text-top-p 0.95, --text-max-tokens 2048).
|
||||
# The yaml stage-0 defaults (temperature 0.1, top_k 80) are TTS-thinker
|
||||
# settings; under near-greedy decoding the 30B loops inside its
|
||||
# reasoning block until the token budget runs out and never emits
|
||||
# </think>, so the extracted "answer" would be truncated reasoning.
|
||||
params = copy.deepcopy(engine.resolve_sampling_params_list(None))
|
||||
stage0 = params[0]
|
||||
stage0.temperature = 1.0
|
||||
stage0.top_p = 0.95
|
||||
stage0.top_k = -1
|
||||
stage0.max_tokens = CHAT_MAX_TOKENS
|
||||
stage0.seed = seed
|
||||
return params
|
||||
|
||||
|
||||
def _chat_once(engine: Omni, prompt: str, seed: int) -> tuple[str, bool]:
|
||||
(req_output,) = engine.generate([{"prompt": prompt, "modalities": ["text"]}], _chat_params(engine, seed))
|
||||
out = req_output.outputs[0]
|
||||
text = (out.text or "").strip()
|
||||
truncated = getattr(out, "finish_reason", None) == "length" or len(out.token_ids or []) >= CHAT_MAX_TOKENS
|
||||
return text, truncated
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
engine = Omni(model=args.model, deploy_config=args.deploy_config, trust_remote_code=True)
|
||||
|
||||
tokenizer = None
|
||||
if args.cfg_scale > 1.0:
|
||||
tokenizer = load_audex_tokenizer(args.model, profile="full", folder="checkpoint_folder_full")
|
||||
|
||||
if args.audio_file:
|
||||
audio, sr = load_audio_file(args.audio_file)
|
||||
else:
|
||||
audio, sr = AudioAsset("mary_had_lamb").audio_and_sample_rate
|
||||
|
||||
# Pass 1 — ASR (text-final; never touches the speech decoder).
|
||||
transcript = _clean(
|
||||
_text_of(
|
||||
engine.generate(
|
||||
[
|
||||
{
|
||||
"prompt": _asr_prompt(),
|
||||
"multi_modal_data": {"audio": (audio, sr)},
|
||||
"modalities": ["text"],
|
||||
}
|
||||
]
|
||||
)
|
||||
)
|
||||
)
|
||||
print(f"[1/3] transcript : {transcript}")
|
||||
if not transcript:
|
||||
raise SystemExit("ASR pass produced an empty transcript; aborting cascade")
|
||||
|
||||
# Pass 2 — chat answer (text-final). The recipe is size-conditional (see
|
||||
# the TEXT_INSTRUCTION comment above): the 30B requires the official
|
||||
# recipe, the 2B regresses under it.
|
||||
if is_30b(args.model):
|
||||
# Official recipe. The raw completion is
|
||||
# <reasoning...></think><answer>; speak only the final response.
|
||||
# Retry with the next seed only when the budget ran out before the
|
||||
# reasoning block closed (a natural stop without </think> is a
|
||||
# complete answer); if every seed spins, fall back to the official
|
||||
# demo's behavior of keeping the full text.
|
||||
raw = ""
|
||||
for seed in CHAT_SEEDS:
|
||||
raw, truncated = _chat_once(engine, _chat_prompt(transcript), seed)
|
||||
if "</think>" in raw or not truncated:
|
||||
break
|
||||
print(f"[2/3] chat seed {seed} hit the token budget mid-reasoning; retrying")
|
||||
answer = _extract_final_response(raw)
|
||||
else:
|
||||
# 2B: bare prompt with a closed think block, engine-default sampling
|
||||
# (yaml near-greedy), no retry; the answer is used as-is.
|
||||
chat_prompt = f"<|im_start|>user\n{transcript}<|im_end|>\n<|im_start|>assistant\n<think></think>"
|
||||
answer = _text_of(engine.generate([{"prompt": chat_prompt, "modalities": ["text"]}]))
|
||||
print(f"[2/3] answer : {answer}")
|
||||
if not answer:
|
||||
raise SystemExit("Chat pass produced an empty answer; aborting cascade")
|
||||
|
||||
# Pass 3 — TTS (audio-final; streams through code2wav). Both final
|
||||
# stages of the pipeline may emit an output record for the request;
|
||||
# collect the audio-bearing one.
|
||||
# The formatting instruction puts every sentence on its own line;
|
||||
# flatten to plain prose for the TTS thinker.
|
||||
tts_prompt = build_cond_prompt(re.sub(r"\s+", " ", answer).strip())
|
||||
outputs = engine.generate(
|
||||
[{"prompt": tts_prompt, "modalities": ["audio"]}],
|
||||
guided_stage_params(engine, tts_prompt, tokenizer, args.cfg_scale, "s2s-tts-pass"),
|
||||
)
|
||||
chunks: list[torch.Tensor] = []
|
||||
for req_output in outputs:
|
||||
mm = getattr(req_output, "multimodal_output", None) or {}
|
||||
if not mm and req_output.outputs:
|
||||
mm = getattr(req_output.outputs[0], "multimodal_output", None) or {}
|
||||
audio_val = mm.get("model_outputs", mm.get("audio"))
|
||||
if audio_val is None:
|
||||
continue
|
||||
if isinstance(audio_val, list):
|
||||
chunks.extend(torch.as_tensor(a).float().cpu().reshape(-1) for a in audio_val if a is not None)
|
||||
else:
|
||||
chunks.append(torch.as_tensor(audio_val).float().cpu().reshape(-1))
|
||||
pcm = torch.cat(chunks) if chunks else torch.empty(0)
|
||||
if pcm.numel() == 0:
|
||||
raise SystemExit("TTS pass produced empty audio")
|
||||
|
||||
out_path = Path(args.output)
|
||||
write_wav(out_path, pcm)
|
||||
print(f"[3/3] speech : {pcm.numel() / SAMPLE_RATE:.2f}s -> {out_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,157 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Offline Audex (Nemotron-Labs-Audex-2B) text-to-audio inference example.
|
||||
|
||||
Caption → general audio through the vLLM-Omni TTA pipeline: the audiogen
|
||||
thinker generates interleaved 4-codebook <audiocodec_N> RVQ tokens under an
|
||||
RVQ phase mask, and the external XCodec1 checkpoint decodes them to a 16 kHz
|
||||
mono WAV.
|
||||
|
||||
(For DIFFUSION text-to-audio models — Stable Audio Open, AudioX — use
|
||||
``examples/offline_inference/text_to_audio/text_to_audio.py`` instead; Audex
|
||||
TTA is a two-stage autoregressive pipeline with a different interface.)
|
||||
|
||||
Classifier-free guidance is effectively mandatory for TTA quality (official
|
||||
default scale 3.0), so every request submits a cond/uncond pair; requests
|
||||
run one at a time.
|
||||
|
||||
Pass the HF repo ROOT as --model. XCodec1 resolves from --xcodec1-path /
|
||||
XCODEC1_PATH / the default hf-audio repo.
|
||||
|
||||
Example:
|
||||
|
||||
python examples/offline_inference/audex/text_to_audio.py \\
|
||||
--captions "Heavy rain falling on a tin roof." \\
|
||||
--output-dir results/audex_tta_wavs
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
from common import (
|
||||
SAMPLE_RATE,
|
||||
default_deploy_config,
|
||||
extract_pcm,
|
||||
guided_stage_params,
|
||||
load_audex_tokenizer,
|
||||
load_corpus,
|
||||
write_wav,
|
||||
)
|
||||
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.model_executor.models.audex.prompt import build_tta_cond_prompt, build_tta_null_prompt
|
||||
from vllm_omni.model_executor.models.audex.tta import build_tta_phase_token_ids
|
||||
|
||||
DEFAULT_CAPTIONS = (
|
||||
"Heavy rain falling on a tin roof.",
|
||||
"A dog barking in the distance while birds chirp.",
|
||||
"Ocean waves crashing on a rocky shore.",
|
||||
)
|
||||
DEFAULT_CODEC_CAP = 4000
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Offline Audex TTA inference")
|
||||
parser.add_argument("--model", type=str, default="nvidia/Nemotron-Labs-Audex-2B")
|
||||
parser.add_argument("--captions", type=str, nargs="+", default=None, help="Audio captions.")
|
||||
parser.add_argument(
|
||||
"--captions-file",
|
||||
type=str,
|
||||
default=None,
|
||||
help="TSV corpus: one 'utt_id<TAB>caption' per line (overrides --captions).",
|
||||
)
|
||||
parser.add_argument("--output-dir", type=str, default="results/audex_tta_wavs")
|
||||
parser.add_argument(
|
||||
"--deploy-config",
|
||||
type=str,
|
||||
# The model root's default deploy yaml is the TTS pipeline; TTA
|
||||
# prompts and RVQ sampling params require the dedicated TTA pipeline.
|
||||
default=default_deploy_config("audex_tta.yaml"),
|
||||
help="Deploy yaml (defaults to the audex_tta pipeline).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--xcodec1-path",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Local XCodec1 checkpoint (defaults to $XCODEC1_PATH or the hf-audio repo).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--cfg-scale",
|
||||
type=float,
|
||||
default=3.0,
|
||||
help="Classifier-free guidance strength (official TTA default 3.0; 1.0 disables).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--codec-cap",
|
||||
type=int,
|
||||
default=DEFAULT_CODEC_CAP,
|
||||
help="Max generated codec tokens before the phase mask forces <audiogen_end>.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
if args.xcodec1_path:
|
||||
os.environ["XCODEC1_PATH"] = args.xcodec1_path
|
||||
output_dir = Path(args.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
corpus = load_corpus(args.captions_file, args.captions, DEFAULT_CAPTIONS)
|
||||
|
||||
tokenizer = load_audex_tokenizer(args.model, profile="tta")
|
||||
phase_token_ids, start_tid, end_tid = build_tta_phase_token_ids(tokenizer)
|
||||
tta_rvq = {
|
||||
"phase_token_ids": phase_token_ids,
|
||||
"start_tid": start_tid,
|
||||
"end_tid": end_tid,
|
||||
"codec_cap": args.codec_cap,
|
||||
# The prompt ends with <audiogen_start>, so generation begins at
|
||||
# RVQ phase 0 immediately.
|
||||
"start_in_prompt": True,
|
||||
}
|
||||
|
||||
engine = Omni(model=args.model, deploy_config=args.deploy_config, trust_remote_code=True)
|
||||
|
||||
print(f"Model : {args.model}")
|
||||
print(f"Captions : {len(corpus)}")
|
||||
print(f"CFG scale : {args.cfg_scale}")
|
||||
print(f"Codec cap : {args.codec_cap}")
|
||||
print(f"Output dir : {output_dir}")
|
||||
|
||||
total_elapsed = 0.0
|
||||
total_dur = 0.0
|
||||
for utt, caption in corpus:
|
||||
cond_prompt = build_tta_cond_prompt(caption)
|
||||
sampling = guided_stage_params(
|
||||
engine,
|
||||
cond_prompt,
|
||||
tokenizer,
|
||||
args.cfg_scale,
|
||||
f"tta-{utt}",
|
||||
extra_args={"tta_rvq": tta_rvq},
|
||||
null_prompt_builder=build_tta_null_prompt,
|
||||
)
|
||||
t_start = time.perf_counter()
|
||||
outputs = engine.generate([cond_prompt], sampling)
|
||||
total_elapsed += time.perf_counter() - t_start
|
||||
|
||||
(req_output,) = outputs
|
||||
pcm = extract_pcm(req_output.outputs[0].multimodal_output)
|
||||
dur = pcm.numel() / SAMPLE_RATE
|
||||
if dur <= 0:
|
||||
raise RuntimeError(f"empty audio for {utt}")
|
||||
out_path = output_dir / f"{utt}.wav"
|
||||
write_wav(out_path, pcm)
|
||||
total_dur += dur
|
||||
print(f" {utt:<48} dur={dur:6.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()
|
||||
@@ -0,0 +1,162 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Offline Audex (Nemotron-Labs-Audex-2B) TTS inference example.
|
||||
|
||||
Runs Stage 0 (thinker) + Stage 1 (streaming causal speech decoder) end-to-end
|
||||
through the vLLM-Omni engine and saves a 16 kHz mono WAV per prompt.
|
||||
|
||||
Pass the HF repo ROOT as --model; per-stage subfolders resolve automatically.
|
||||
|
||||
Examples:
|
||||
|
||||
python examples/offline_inference/audex/text_to_speech.py \\
|
||||
--texts "The weather is so good today." \\
|
||||
--output-dir results/audex_wavs
|
||||
|
||||
# TSV corpus (utt_id<TAB>text per line), all requests in flight at once:
|
||||
python examples/offline_inference/audex/text_to_speech.py \\
|
||||
--texts-file corpus.tsv --output-dir results/audex_wavs
|
||||
|
||||
# One request at a time:
|
||||
python examples/offline_inference/audex/text_to_speech.py \\
|
||||
--texts-file corpus.tsv --sequential --output-dir results/audex_wavs
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
from common import (
|
||||
GUIDED_TEMPERATURE,
|
||||
SAMPLE_RATE,
|
||||
extract_pcm,
|
||||
guided_stage_params,
|
||||
load_audex_tokenizer,
|
||||
load_corpus,
|
||||
request_index,
|
||||
write_wav,
|
||||
)
|
||||
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.model_executor.models.audex.prompt import build_cond_prompt
|
||||
|
||||
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():
|
||||
parser = argparse.ArgumentParser(description="Offline Audex TTS inference")
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
default="nvidia/Nemotron-Labs-Audex-2B",
|
||||
help="HF repo root (or local snapshot root).",
|
||||
)
|
||||
parser.add_argument("--texts", type=str, nargs="+", default=None, help="Plain-text prompts.")
|
||||
parser.add_argument(
|
||||
"--texts-file",
|
||||
type=str,
|
||||
default=None,
|
||||
help="TSV corpus: one 'utt_id<TAB>text' per line (overrides --texts).",
|
||||
)
|
||||
parser.add_argument("--output-dir", type=str, default="results/audex_wavs")
|
||||
parser.add_argument("--deploy-config", type=str, default=None, help="Override the deploy config path.")
|
||||
parser.add_argument(
|
||||
"--sequential",
|
||||
action="store_true",
|
||||
help="Submit one request at a time instead of all at once.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--cfg-scale",
|
||||
type=float,
|
||||
default=1.0,
|
||||
help=(
|
||||
"Classifier-free guidance strength. 1.0 disables guidance (default, "
|
||||
"matches the v1 baseline); 1.5 is the official quality setting. "
|
||||
"Guided runs submit one request at a time (each carries its own "
|
||||
"length-matched null prompt and pair id)."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--temperature",
|
||||
type=float,
|
||||
default=None,
|
||||
help="Override the deploy yaml's stage-0 sampling temperature.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def _write_outputs(outputs, batch: list[tuple[str, str]], output_dir: Path) -> float:
|
||||
ordered = sorted(outputs, key=request_index)
|
||||
if len(ordered) != len(batch):
|
||||
raise RuntimeError(f"expected {len(batch)} outputs, got {len(ordered)}")
|
||||
|
||||
total_dur = 0.0
|
||||
for (utt, _text), req_output in zip(batch, ordered):
|
||||
pcm = extract_pcm(req_output.outputs[0].multimodal_output)
|
||||
dur = pcm.numel() / SAMPLE_RATE
|
||||
if dur <= 0:
|
||||
raise RuntimeError(f"empty audio for {utt}")
|
||||
out_path = output_dir / f"{utt}.wav"
|
||||
write_wav(out_path, pcm)
|
||||
total_dur += dur
|
||||
print(f" {utt:<40} dur={dur:6.2f}s -> {out_path}")
|
||||
return total_dur
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
output_dir = Path(args.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
corpus = load_corpus(args.texts_file, args.texts, DEFAULT_TEXTS)
|
||||
|
||||
engine = Omni(model=args.model, deploy_config=args.deploy_config, trust_remote_code=True)
|
||||
|
||||
cfg_enabled = args.cfg_scale > 1.0
|
||||
tokenizer = load_audex_tokenizer(args.model) if cfg_enabled else None
|
||||
|
||||
print(f"Model : {args.model}")
|
||||
print(f"Prompts : {len(corpus)}")
|
||||
print(f"Mode : {'sequential' if args.sequential or cfg_enabled else 'batched'}")
|
||||
print(f"CFG scale : {args.cfg_scale}" + (" (guidance off)" if not cfg_enabled else ""))
|
||||
print(f"Output dir : {output_dir}")
|
||||
|
||||
total_elapsed = 0.0
|
||||
total_dur = 0.0
|
||||
# CFG runs one request at a time: each request needs its own pair id and
|
||||
# text-specific null prompt in the stage-0 sampling params.
|
||||
batches = [[item] for item in corpus] if (args.sequential or cfg_enabled) else [corpus]
|
||||
for batch in batches:
|
||||
prompts = [build_cond_prompt(text) for _utt, text in batch]
|
||||
sampling_params_list = None
|
||||
if cfg_enabled:
|
||||
sampling_params_list = guided_stage_params(
|
||||
engine,
|
||||
prompts[0],
|
||||
tokenizer,
|
||||
args.cfg_scale,
|
||||
f"cfg-{batch[0][0]}",
|
||||
temperature=GUIDED_TEMPERATURE,
|
||||
)
|
||||
if args.temperature is not None:
|
||||
import copy
|
||||
|
||||
if sampling_params_list is None:
|
||||
sampling_params_list = copy.deepcopy(engine.resolve_sampling_params_list(None))
|
||||
sampling_params_list[0].temperature = args.temperature
|
||||
t_start = time.perf_counter()
|
||||
outputs = engine.generate(prompts, sampling_params_list)
|
||||
total_elapsed += time.perf_counter() - t_start
|
||||
total_dur += _write_outputs(outputs, batch, output_dir)
|
||||
|
||||
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()
|
||||
@@ -1,7 +1,7 @@
|
||||
# Image-To-Video
|
||||
|
||||
This shared example generates videos from images with VACE, Wan2.2,
|
||||
LTX-2/LTX-2.3, HunyuanVideo-1.5, Cosmos3, and other compatible pipelines.
|
||||
This shared example generates videos from images with VACE, Wan2.2, LTX-2,
|
||||
HunyuanVideo-1.5, Cosmos3, and other compatible pipelines.
|
||||
|
||||
- `image_to_video.py`: command-line script for single video generation with advanced options.
|
||||
|
||||
@@ -26,7 +26,7 @@ This folder provides a unified CLI script for image-to-video generation using vL
|
||||
| `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | 480 x 832 | 81 | 50 | 5.0 | Around 60 GiB BF16 for basic single-card usage |
|
||||
| `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | 480 x 832 | 81 | 50 | 4.0 | Around 20–25 GiB BF16, smallest I2V model |
|
||||
| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v` | 480 x 832 | 121 | 50 | 6.0 | Around 100 GiB at default settings; the example enables `--enable-cpu-offload` + VAE tiling/slicing to fit an 80 GiB card |
|
||||
| LTX2 (local path + `--model-class-name LTX2ImageToVideoPipeline`) | 512 x 768 | 121 | 40 | 4.0 | Memory use depends on frame count and tensor parallelism |
|
||||
| `Lightricks/LTX-2` | 512 x 768 | 121 | 40 | video 3.0 / audio 7.0 | Memory use depends on frame count and tensor parallelism |
|
||||
|
||||
!!! info
|
||||
Peak VRAM: based on basic single-card usage, batch size = 1, without any acceleration/optimization features. Some model weights cannot fit into one card with 80 GiB VRAM, which may need to use CPU offloading.
|
||||
@@ -105,7 +105,7 @@ python image_to_video.py \
|
||||
| Argument | Type | Default | Description |
|
||||
| -------- | ---- | ------- | ----------- |
|
||||
| `--model` | str | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | Diffusers I2V model ID or local path |
|
||||
| `--model-class-name` | str | `None` | Override model class name (e.g., `LTX2ImageToVideoPipeline`) |
|
||||
| `--model-class-name` | str | `None` | Optional pipeline override |
|
||||
| `--image` | str | (required) | Path to input image |
|
||||
| `--prompt` | str | `""` | Text description of desired motion/animation |
|
||||
| `--negative-prompt` | str | `""` | Optional list of artifacts to suppress |
|
||||
@@ -195,24 +195,19 @@ python image_to_video.py \
|
||||
--output hunyuan_i2v.mp4
|
||||
```
|
||||
|
||||
### LTX2 Image-to-Video
|
||||
### LTX-2
|
||||
|
||||
```bash
|
||||
python image_to_video.py \
|
||||
--model /path/to/LTX-2 \
|
||||
--model-class-name LTX2ImageToVideoPipeline \
|
||||
--model Lightricks/LTX-2 \
|
||||
--image cherry_blossom.jpg \
|
||||
--prompt "A cinematic dolly shot of cherry blossoms" \
|
||||
--height 512 \
|
||||
--width 768 \
|
||||
--num-frames 121 \
|
||||
--num-inference-steps 40 \
|
||||
--guidance-scale 4.0 \
|
||||
--frame-rate 24 \
|
||||
--fps 24 \
|
||||
--output ltx2_i2v.mp4
|
||||
```
|
||||
|
||||
See the [LTX-2 recipe](../../../recipes/LTX/LTX-2.md) for all checkpoints,
|
||||
pipeline selection, T2V, defaults, and advanced options.
|
||||
|
||||
## Advanced Features
|
||||
|
||||
### CFG Parallel
|
||||
@@ -233,30 +228,6 @@ python image_to_video.py \
|
||||
--output i2v_cached.mp4
|
||||
```
|
||||
|
||||
### LTX-2.3
|
||||
|
||||
```bash
|
||||
python image_to_video.py \
|
||||
--model diffusers/LTX-2.3-Diffusers \
|
||||
--model-class-name LTX23ImageToVideoPipeline \
|
||||
--image cherry_blossom.jpg \
|
||||
--prompt "Cherry blossoms swaying gently in the breeze with synchronized ambient sound" \
|
||||
--negative-prompt "worst quality, inconsistent motion, blurry, jittery, distorted" \
|
||||
--height 384 \
|
||||
--width 512 \
|
||||
--num-frames 25 \
|
||||
--guidance-scale 4.0 \
|
||||
--num-inference-steps 20 \
|
||||
--frame-rate 24 \
|
||||
--fps 24 \
|
||||
--output ltx23_i2v_output.mp4
|
||||
```
|
||||
|
||||
Use the Diffusers-format checkpoint `diffusers/LTX-2.3-Diffusers`; the
|
||||
upstream `Lightricks/LTX-2.3` raw safetensors repo is not directly loadable by
|
||||
this pipeline. Pass `--model-class-name LTX23ImageToVideoPipeline` to select
|
||||
the LTX-2.3 image-to-video pipeline.
|
||||
|
||||
### Cosmos3
|
||||
|
||||
```bash
|
||||
@@ -274,8 +245,8 @@ python image_to_video.py \
|
||||
|
||||
Key arguments:
|
||||
|
||||
- `--model`: Model ID (I2V-A14B for MoE, TI2V-5B for unified T2V+I2V,
|
||||
LTX-2/LTX-2.3, Cosmos3, or VACE).
|
||||
- `--model`: Model ID (I2V-A14B for MoE, TI2V-5B for unified T2V+I2V, LTX-2,
|
||||
Cosmos3, or VACE).
|
||||
- `--image`: Path to the first-frame or source image.
|
||||
- `--last-image`: Optional last-frame condition for models such as VACE.
|
||||
- `--mask-image`: Optional inpainting mask. White pixels are regenerated and black pixels are preserved.
|
||||
@@ -290,7 +261,7 @@ Key arguments:
|
||||
- `--guidance-scale` and `--guidance-scale-high`: CFG scale (applied to low/high-noise stages for MoE).
|
||||
- `--negative-prompt`: Optional list of artifacts to suppress.
|
||||
- `--boundary-ratio`: Boundary split ratio for two-stage MoE models.
|
||||
- `--flow-shift`: Scheduler flow shift (default: model-specific — Wan/LTX2 5.0, Cosmos3 10.0).
|
||||
- `--flow-shift`: Scheduler flow shift (default: Wan 5.0, Cosmos3 10.0). LTX does not consume this option.
|
||||
- `--sample-solver`: Wan2.2 sampling solver. Use `unipc` for the default multistep solver, or `euler` for Lightning/Distill checkpoints.
|
||||
- `--num-inference-steps`: Number of denoising steps (default: model-specific — Wan 50, LTX2 40, Cosmos3 35).
|
||||
- `--fps`: Frames per second for the saved MP4 (requires `diffusers` export_to_video).
|
||||
|
||||
@@ -23,16 +23,14 @@ Usage:
|
||||
|
||||
# LTX2 image-to-video
|
||||
python image_to_video.py --model /path/to/LTX-2 \
|
||||
--model-class-name LTX2ImageToVideoPipeline \
|
||||
--image input.jpg --prompt "A cinematic dolly shot of a boat" \
|
||||
--num-frames 121 --num-inference-steps 40 --guidance-scale 4.0 \
|
||||
--num-frames 121 --num-inference-steps 40 \
|
||||
--frame-rate 24 --fps 24 --output ltx2_i2v.mp4
|
||||
|
||||
# LTX-2.3 image-to-video
|
||||
python image_to_video.py --model diffusers/LTX-2.3-Diffusers \
|
||||
--model-class-name LTX23ImageToVideoPipeline \
|
||||
--image input.jpg --prompt "A cinematic dolly shot of a boat" \
|
||||
--height 384 --width 512 --num-frames 25 --num-inference-steps 20 \
|
||||
--height 512 --width 768 --num-frames 121 --num-inference-steps 30 \
|
||||
--frame-rate 24 --fps 24 --output ltx23_i2v.mp4
|
||||
|
||||
# HunyuanVideo-1.5 I2V (480p)
|
||||
@@ -100,7 +98,7 @@ def parse_args() -> argparse.Namespace:
|
||||
parser.add_argument(
|
||||
"--model-class-name",
|
||||
default=None,
|
||||
help="Override model class name (e.g., LTX2ImageToVideoPipeline or LTX23ImageToVideoPipeline).",
|
||||
help="Override model class name (LTX checkpoints default to LTX2Pipeline).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--deploy-config",
|
||||
@@ -117,7 +115,7 @@ def parse_args() -> argparse.Namespace:
|
||||
help="Path to a reference image. Repeat to provide multiple references.",
|
||||
)
|
||||
parser.add_argument("--prompt", default="", help="Text prompt describing the desired motion.")
|
||||
parser.add_argument("--negative-prompt", default="", help="Negative prompt.")
|
||||
parser.add_argument("--negative-prompt", default=None, help="Negative prompt. Default: model-specific.")
|
||||
parser.add_argument("--seed", type=int, default=42, help="Random seed.")
|
||||
parser.add_argument("--guidance-scale", type=float, default=None, help="CFG scale. Default: model-specific.")
|
||||
parser.add_argument(
|
||||
@@ -331,8 +329,12 @@ def main():
|
||||
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
|
||||
model_name = str(args.model).lower() if args.model is not None else ""
|
||||
model_class_name = args.model_class_name
|
||||
is_ltx2 = model_class_name in {"LTX2ImageToVideoPipeline", "LTX23ImageToVideoPipeline"}
|
||||
model_class_name_lower = (model_class_name or "").lower()
|
||||
is_ltx2_distilled = "distilled" in model_class_name_lower or "distilled" in model_name
|
||||
is_ltx23 = "ltx23" in model_class_name_lower or "ltx-2.3" in model_name
|
||||
is_ltx2 = is_ltx2_distilled or is_ltx23 or "ltx2" in model_class_name_lower or "ltx-2" in model_name
|
||||
is_cosmos = "cosmos" in model_name or (model_class_name is not None and "cosmos" in model_class_name.lower())
|
||||
is_cosmos_edge = is_cosmos and ("edge" in model_name or "edge" in model_class_name_lower)
|
||||
|
||||
image = PIL.Image.open(args.image).convert("RGB") if args.image else None
|
||||
last_image = PIL.Image.open(args.last_image).convert("RGB") if args.last_image else None
|
||||
@@ -346,7 +348,19 @@ def main():
|
||||
|
||||
# Per-model generation defaults, applied only when the matching flag is omitted.
|
||||
# Cosmos3 would otherwise silently inherit the Wan2.2 defaults (wrong size/steps/shift).
|
||||
if is_cosmos:
|
||||
if is_cosmos_edge:
|
||||
# Cosmos3-Edge native defaults: 480x832, gs 5.0, flow_shift 3.0 — NOT the Nano/Super
|
||||
# 720p / gs 6.0 / flow_shift 10.0 below, which produce degenerate output on Edge.
|
||||
d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = (
|
||||
24,
|
||||
5.0,
|
||||
189,
|
||||
35,
|
||||
3.0,
|
||||
480 * 832,
|
||||
16,
|
||||
)
|
||||
elif is_cosmos:
|
||||
d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = (
|
||||
24,
|
||||
6.0,
|
||||
@@ -356,13 +370,23 @@ def main():
|
||||
1280 * 720,
|
||||
16,
|
||||
)
|
||||
elif is_ltx2_distilled:
|
||||
d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = (
|
||||
24,
|
||||
None,
|
||||
121,
|
||||
8,
|
||||
None,
|
||||
1024 * 1536,
|
||||
64,
|
||||
)
|
||||
elif is_ltx2:
|
||||
d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = (
|
||||
24,
|
||||
4.0,
|
||||
None,
|
||||
121,
|
||||
40,
|
||||
5.0,
|
||||
30 if is_ltx23 else 40,
|
||||
None,
|
||||
512 * 768,
|
||||
32,
|
||||
)
|
||||
@@ -445,7 +469,6 @@ def main():
|
||||
vae_use_slicing=args.vae_use_slicing,
|
||||
vae_use_tiling=args.vae_use_tiling,
|
||||
boundary_ratio=args.boundary_ratio,
|
||||
flow_shift=flow_shift,
|
||||
diffusion_kv_cache_dtype=args.diffusion_kv_cache_dtype,
|
||||
diffusion_kv_cache_skip_steps=args.diffusion_kv_cache_skip_steps,
|
||||
diffusion_kv_cache_skip_layers=args.diffusion_kv_cache_skip_layers,
|
||||
@@ -460,6 +483,8 @@ def main():
|
||||
)
|
||||
if args.deploy_config:
|
||||
omni_kwargs["deploy_config"] = args.deploy_config
|
||||
if flow_shift is not None:
|
||||
omni_kwargs["flow_shift"] = flow_shift
|
||||
if args.quantization is not None:
|
||||
omni_kwargs["quantization"] = args.quantization
|
||||
# Cosmos3 loads its (gated) guardrail models at build time, so the guardrails
|
||||
@@ -492,10 +517,14 @@ def main():
|
||||
print(f" Video size: {width}x{height}")
|
||||
print(f"{'=' * 60}\n")
|
||||
|
||||
negative_prompt = args.negative_prompt
|
||||
if negative_prompt is None and not is_ltx2:
|
||||
# Preserve the historical empty-prompt behavior for non-LTX examples.
|
||||
negative_prompt = ""
|
||||
prompt_dict = build_image_to_video_prompt(
|
||||
model_class_name=model_class_name,
|
||||
prompt=args.prompt,
|
||||
negative_prompt=args.negative_prompt,
|
||||
negative_prompt=negative_prompt,
|
||||
media_inputs=media_inputs,
|
||||
height=height,
|
||||
width=width,
|
||||
@@ -511,11 +540,10 @@ def main():
|
||||
num_inference_steps=num_inference_steps,
|
||||
num_frames=num_frames,
|
||||
frame_rate=frame_rate,
|
||||
extra_args={
|
||||
"sample_solver": args.sample_solver,
|
||||
"flow_shift": flow_shift,
|
||||
},
|
||||
extra_args={"sample_solver": args.sample_solver},
|
||||
)
|
||||
if flow_shift is not None:
|
||||
sampling_params.extra_args["flow_shift"] = flow_shift
|
||||
|
||||
# Route model-specific knobs through extra_body, filtered against the model's
|
||||
# declared extra_body_params. Models without a declaration only forward explicit
|
||||
|
||||
@@ -182,10 +182,13 @@ def run_e2e():
|
||||
print(f"Multimodal output keys: {mm.keys()}")
|
||||
if "audio" in mm:
|
||||
audio_out = mm["audio"]
|
||||
# The model emits its own rate (24 kHz for CosyVoice3); take
|
||||
# it from the payload so the file is not mislabeled.
|
||||
sample_rate = int(mm["sr"]) if "sr" in mm.keys() else 24000
|
||||
print(f"Generated Audio Shape: {audio_out.shape}")
|
||||
out_path = f"output_{i}.wav"
|
||||
sf.write(out_path, audio_out.cpu().numpy().squeeze(), 22050)
|
||||
print(f"Saved audio to {out_path}")
|
||||
sf.write(out_path, audio_out.cpu().numpy().squeeze(), sample_rate)
|
||||
print(f"Saved audio to {out_path} (sr={sample_rate})")
|
||||
else:
|
||||
print("No multimodal output found.")
|
||||
except Exception as e:
|
||||
|
||||
@@ -8,7 +8,7 @@ A unified script for text-to-video generation. Supports multiple models with mod
|
||||
|---|---|---|---|---|---|
|
||||
| `Wan-AI/Wan2.1-VACE-1.3B-diffusers` | 480x832 | 81 | 30 | 5.0 | ~20 GiB (RTX 5090, VAE tiling) |
|
||||
| `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 720x1280 | 81 | 40 | 4.0 | ~60 GiB |
|
||||
| `diffusers/LTX-2.3-Diffusers` | 384x512 | 25 | 20 | 4.0 | 96GB-class GPU |
|
||||
| `Lightricks/LTX-2` | 512x768 | 121 | 40 | video 3.0 / audio 7.0 | Model-dependent |
|
||||
| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v` | 480x832 | 121 | 50 | 6.0 | 1×A100 80GB |
|
||||
| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required |
|
||||
| `nvidia/Cosmos3-Nano` | 720x1280 | 189 | 35 | 6.0 | ~46 GiB (peak, 720p) |
|
||||
@@ -55,45 +55,17 @@ python text_to_video.py \
|
||||
--output vace_t2v_output.mp4
|
||||
```
|
||||
|
||||
LTX2 example:
|
||||
### LTX-2
|
||||
|
||||
```bash
|
||||
python text_to_video.py \
|
||||
--model "Lightricks/LTX-2" \
|
||||
--prompt "A cinematic close-up of ocean waves at golden hour." \
|
||||
--negative-prompt "worst quality, inconsistent motion, blurry, jittery, distorted" \
|
||||
--height 512 \
|
||||
--width 768 \
|
||||
--num-frames 121 \
|
||||
--num-inference-steps 40 \
|
||||
--guidance-scale 4.0 \
|
||||
--frame-rate 24 \
|
||||
--output ltx2_out.mp4
|
||||
```
|
||||
|
||||
### LTX-2.3
|
||||
|
||||
```bash
|
||||
python text_to_video.py \
|
||||
--model diffusers/LTX-2.3-Diffusers \
|
||||
--model-class-name LTX23Pipeline \
|
||||
--model Lightricks/LTX-2 \
|
||||
--prompt "Cherry blossoms swaying gently in the breeze with synchronized ambient sound" \
|
||||
--negative-prompt "worst quality, inconsistent motion, blurry, jittery, distorted" \
|
||||
--height 384 \
|
||||
--width 512 \
|
||||
--num-frames 25 \
|
||||
--num-inference-steps 20 \
|
||||
--guidance-scale 4.0 \
|
||||
--frame-rate 24 \
|
||||
--fps 24 \
|
||||
--audio-sample-rate 48000 \
|
||||
--output ltx23_t2v_output.mp4
|
||||
--output ltx2_output.mp4
|
||||
```
|
||||
|
||||
Use the Diffusers-format checkpoint `diffusers/LTX-2.3-Diffusers`; the
|
||||
upstream `Lightricks/LTX-2.3` raw safetensors repo is not directly loadable by
|
||||
this pipeline. Pass `--model-class-name LTX23Pipeline` to select the LTX-2.3
|
||||
text-to-video pipeline explicitly.
|
||||
See the [LTX-2 recipe](../../../recipes/LTX/LTX-2.md) for all checkpoints,
|
||||
pipeline selection, I2V, defaults, and advanced options.
|
||||
|
||||
### HunyuanVideo-1.5 (480p)
|
||||
|
||||
@@ -212,8 +184,7 @@ python text_to_video.py \
|
||||
### Common
|
||||
|
||||
- `--model`: Diffusers model ID or local path.
|
||||
- `--model-class-name`: Optional explicit pipeline class. Use `LTX23Pipeline`
|
||||
for LTX-2.3 text-to-video.
|
||||
- `--model-class-name`: Optional explicit pipeline override.
|
||||
- `--prompt`: text description (string).
|
||||
- `--height/--width`: output resolution. Default depends on model.
|
||||
- `--num-frames`: number of frames. Default depends on model.
|
||||
@@ -229,8 +200,7 @@ python text_to_video.py \
|
||||
- `--enable-cpu-offload`: enable CPU offloading for diffusion models.
|
||||
- `--enable-layerwise-offload`: enable layerwise offloading on DiT modules.
|
||||
- `--frame-rate`: generation FPS for pipelines that require it (e.g., LTX2).
|
||||
- `--audio-sample-rate`: audio sample rate for embedded audio (when the
|
||||
pipeline returns audio; LTX-2.3 outputs 48kHz audio).
|
||||
- `--audio-sample-rate`: fallback audio sample rate when the pipeline returns audio.
|
||||
- `--quantization`: quantization method (such as `fp8` for FP8).
|
||||
- `--flow-shift`: scheduler flow_shift parameter.
|
||||
- `--extra-body`: JSON object of model-specific generation params, filtered against the model's declared `extra_body_params` (see [`vllm_omni/model_extras`](../../../vllm_omni/model_extras)). Used by Cosmos3 (see above).
|
||||
|
||||
@@ -56,6 +56,19 @@ _MODEL_PRESETS = {
|
||||
"flow_shift": 10.0,
|
||||
"output": "cosmos3_t2v_output.mp4",
|
||||
},
|
||||
# Cosmos3-Edge is a compact checkpoint with its own native resolution / guidance /
|
||||
# flow_shift. It must NOT inherit the Nano/Super "cosmos" numbers above (720p / gs 6.0 /
|
||||
# flow_shift 10.0), which produce degenerate output on Edge.
|
||||
"cosmos3_edge": {
|
||||
"height": 480,
|
||||
"width": 832,
|
||||
"num_frames": 189,
|
||||
"num_inference_steps": 35,
|
||||
"guidance_scale": 5.0,
|
||||
"fps": 24,
|
||||
"flow_shift": 3.0,
|
||||
"output": "cosmos3_edge_t2v_output.mp4",
|
||||
},
|
||||
"helios": {
|
||||
"height": 384,
|
||||
"width": 640,
|
||||
@@ -65,13 +78,48 @@ _MODEL_PRESETS = {
|
||||
"fps": 16,
|
||||
"output": "helios_output.mp4",
|
||||
},
|
||||
"ltx2": {
|
||||
"height": 512,
|
||||
"width": 768,
|
||||
"num_frames": 121,
|
||||
"num_inference_steps": 40,
|
||||
"fps": 24,
|
||||
"output": "ltx2_output.mp4",
|
||||
},
|
||||
"ltx2_distilled": {
|
||||
"height": 1024,
|
||||
"width": 1536,
|
||||
"num_frames": 121,
|
||||
"num_inference_steps": 8,
|
||||
"fps": 24,
|
||||
"output": "ltx2_distilled_output.mp4",
|
||||
},
|
||||
"ltx23": {
|
||||
"height": 512,
|
||||
"width": 768,
|
||||
"num_frames": 121,
|
||||
"num_inference_steps": 30,
|
||||
"fps": 24,
|
||||
"output": "ltx23_output.mp4",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _detect_preset(model: str) -> dict:
|
||||
def _detect_preset(model: str, model_class_name: str | None = None) -> dict:
|
||||
model_lower = model.lower()
|
||||
class_lower = (model_class_name or "").lower()
|
||||
if "distilled" in class_lower or "distilled" in model_lower:
|
||||
return _MODEL_PRESETS["ltx2_distilled"]
|
||||
if "ltx23" in class_lower or "ltx-2.3" in model_lower or "ltx_2.3" in model_lower:
|
||||
return _MODEL_PRESETS["ltx23"]
|
||||
if "ltx2" in class_lower or "ltx-2" in model_lower or "ltx_2" in model_lower:
|
||||
return _MODEL_PRESETS["ltx2"]
|
||||
if "vace" in model_lower:
|
||||
return _MODEL_PRESETS["vace"]
|
||||
# Edge must be matched before the generic cosmos branch (its "cosmos" substring would
|
||||
# otherwise pick up the Nano/Super 720p / gs 6.0 / flow_shift 10.0 preset).
|
||||
if ("cosmos" in model_lower or "cosmos" in class_lower) and ("edge" in model_lower or "edge" in class_lower):
|
||||
return _MODEL_PRESETS["cosmos3_edge"]
|
||||
if "cosmos" in model_lower:
|
||||
return _MODEL_PRESETS["cosmos"]
|
||||
if "hunyuan" in model_lower:
|
||||
@@ -122,7 +170,7 @@ def parse_args() -> argparse.Namespace:
|
||||
parser.add_argument(
|
||||
"--model-class-name",
|
||||
default=None,
|
||||
help="Override model class name (e.g., LTX2TwoStagesVideoPipeline).",
|
||||
help="Override model class name (e.g., LTX2Pipeline).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--deploy-config",
|
||||
@@ -130,7 +178,7 @@ def parse_args() -> argparse.Namespace:
|
||||
help="Optional deploy config YAML to use for pipeline-backed runs.",
|
||||
)
|
||||
parser.add_argument("--prompt", default="A serene lakeside sunrise with mist over the water.", help="Text prompt.")
|
||||
parser.add_argument("--negative-prompt", default="", help="Negative prompt.")
|
||||
parser.add_argument("--negative-prompt", default=None, help="Negative prompt. Default: model-specific.")
|
||||
parser.add_argument(
|
||||
"--extra-body",
|
||||
type=parse_extra_body,
|
||||
@@ -331,7 +379,7 @@ def main():
|
||||
args = parse_args()
|
||||
model_class_name = args.model_class_name
|
||||
|
||||
preset = _detect_preset(args.model)
|
||||
preset = _detect_preset(args.model, model_class_name)
|
||||
for key, default_val in preset.items():
|
||||
if getattr(args, key.replace("-", "_"), None) is None:
|
||||
setattr(args, key.replace("-", "_"), default_val)
|
||||
@@ -422,8 +470,11 @@ def main():
|
||||
print(f"{'=' * 60}\n")
|
||||
|
||||
prompt_dict = {"prompt": args.prompt}
|
||||
if args.negative_prompt:
|
||||
if args.negative_prompt is not None:
|
||||
prompt_dict["negative_prompt"] = args.negative_prompt
|
||||
elif preset not in (_MODEL_PRESETS["ltx2"], _MODEL_PRESETS["ltx23"]):
|
||||
# Preserve the historical empty-prompt behavior for non-LTX examples.
|
||||
prompt_dict["negative_prompt"] = ""
|
||||
|
||||
sampling_kwargs = dict(
|
||||
height=args.height,
|
||||
|
||||
@@ -0,0 +1,89 @@
|
||||
# Audex (Nemotron-Labs-Audex-2B) online serving
|
||||
|
||||
One checkpoint, four deployment modes. `run_server.sh` starts the server
|
||||
for a given `MODE`; `client.py --mode <mode>` tests it. The capability
|
||||
matrix:
|
||||
|
||||
| mode (`vllm_omni/deploy/audex_<mode>.yaml`) | audio in | text out | speech out | general audio out | endpoint |
|
||||
|---|---|---|---|---|---|
|
||||
| `tts` | ❌ | ❌ | ✅ | ❌ | `/v1/audio/speech` |
|
||||
| `tta` | ❌ | ❌ | ❌ | ✅ | `/v1/audio/speech` |
|
||||
| `thinker_only` | ✅ | ✅ | ❌ | ❌ | `/v1/chat/completions` |
|
||||
| `s2s` | ✅ | ✅ | ✅ | ❌ | both |
|
||||
|
||||
Audio-input modes of `client.py` fall back to vLLM's public
|
||||
`mary_had_lamb` asset when `--audio-file` is omitted; CFG defaults follow
|
||||
the official settings per mode (tts/s2s 1.5, tta 3.0; `--cfg-scale 1.0`
|
||||
disables).
|
||||
|
||||
## tts — text → speech
|
||||
|
||||
Text-only thinker (`checkpoint_folder_audiogen`) emits `<speechcodec_N>`
|
||||
tokens; the streaming causal decoder produces 16 kHz WAVs. This is also
|
||||
the DEFAULT pipeline when the repo root is served without a deploy config.
|
||||
|
||||
```bash
|
||||
./run_server.sh # MODE=tts, port 8097
|
||||
|
||||
python client.py --mode tts --text "Hello there." --output hello.wav
|
||||
# or raw curl:
|
||||
curl -s http://localhost:8097/v1/audio/speech \
|
||||
-H 'Content-Type: application/json' \
|
||||
-d '{"model": "nvidia/Nemotron-Labs-Audex-2B", "input": "Hello there.", "response_format": "wav", "extra_params": {"cfg_scale": 1.5}}' \
|
||||
-o hello.wav
|
||||
```
|
||||
|
||||
## tta — caption → general audio
|
||||
|
||||
Same thinker over the 4-codebook `<audiocodec_N>` RVQ block, decoded by
|
||||
the external XCodec1 checkpoint (auto-downloaded; override with
|
||||
`XCODEC1_PATH`). Clips are capped at 10 s.
|
||||
|
||||
```bash
|
||||
MODE=tta ./run_server.sh
|
||||
|
||||
python client.py --mode tta --caption "Heavy rain falling on a tin roof." \
|
||||
--output rain.wav
|
||||
```
|
||||
|
||||
## thinker_only — audio (+ instruction) → text
|
||||
|
||||
Single-stage audio understanding on the full checkpoint (NV-Whisper
|
||||
encoder + projector + LM). Send the clip as `input_audio` chat content.
|
||||
|
||||
```bash
|
||||
MODE=thinker_only ./run_server.sh
|
||||
|
||||
python client.py --mode thinker_only # transcribe the demo asset
|
||||
python client.py --mode thinker_only --audio-file a.wav \
|
||||
--question "What language is being spoken?"
|
||||
```
|
||||
|
||||
## s2s — spoken question → spoken answer
|
||||
|
||||
The audio-capable thinker plus the speech decoder in one deployment.
|
||||
`client.py --mode s2s` runs the official three passes: ASR (chat with
|
||||
`input_audio`) → chat on the transcript → TTS on the answer. Text-final
|
||||
passes carry `"modalities": ["text"]`; only the TTS pass streams through
|
||||
the decoder.
|
||||
|
||||
```bash
|
||||
MODE=s2s PORT=8098 ./run_server.sh
|
||||
|
||||
python client.py --mode s2s --port 8098 \
|
||||
--audio-file question.wav --output answer.wav
|
||||
```
|
||||
|
||||
Offline counterparts (no HTTP, one `Omni` engine per script) live in
|
||||
`examples/offline_inference/audex/`.
|
||||
|
||||
## 30B-A3B (nvidia/Nemotron-Labs-Audex-30B-A3B)
|
||||
|
||||
`SIZE=30b` switches run_server.sh to the 30B yamls and model id:
|
||||
|
||||
SIZE=30b ./run_server.sh # 30B tts on port 8097
|
||||
SIZE=30b MODE=s2s PORT=8098 ./run_server.sh
|
||||
|
||||
Pass `--model nvidia/Nemotron-Labs-Audex-30B-A3B` to the client — the server
|
||||
validates the model id and 404s on a mismatch. Single-H100 defaults; see the
|
||||
yaml comments for the TP2 fallback. First launch downloads ~60 GB.
|
||||
@@ -0,0 +1,319 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Audex online client for all four deployment modes.
|
||||
|
||||
Pick ``--mode`` to match the server's deploy yaml (see run_server.sh):
|
||||
|
||||
tts /v1/audio/speech: text -> speech WAV
|
||||
tta /v1/audio/speech: caption -> general-audio WAV
|
||||
thinker_only /v1/chat/completions: audio (+ question) -> text
|
||||
s2s the official three-pass cascade against one audex_s2s
|
||||
server: ASR (chat, input_audio) -> chat -> TTS
|
||||
|
||||
Audio-input modes fall back to vLLM's public ``mary_had_lamb`` asset when
|
||||
``--audio-file`` is omitted.
|
||||
|
||||
Examples:
|
||||
|
||||
python examples/online_serving/audex/client.py --mode tts \\
|
||||
--text "Hello there." --output hello.wav
|
||||
python examples/online_serving/audex/client.py --mode tta \\
|
||||
--caption "Heavy rain falling on a tin roof." --output rain.wav
|
||||
python examples/online_serving/audex/client.py --mode thinker_only
|
||||
python examples/online_serving/audex/client.py --mode s2s \\
|
||||
--audio-file question.wav --output answer.wav
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
import requests
|
||||
|
||||
ASR_QUESTION = "Transcribe the input speech."
|
||||
SYSTEM_PROMPT = "You are a helpful and harmless assistant.\n\nYou are not allowed to use any tools."
|
||||
# The official chat-pass recipe below (SYSTEM_PROMPT + TEXT_INSTRUCTION
|
||||
# prefix, reasoning enabled, temp 1.0/top_p 0.95, seed retry) applies to the
|
||||
# 30B-A3B ONLY (see is_30b() below and
|
||||
# tests/e2e/offline_inference/test_audex_s2s.py, which mirrors this same
|
||||
# split). It is load-bearing on the 30B: without the formatting-instruction
|
||||
# prefix a bare transcript turn is treated as a TTS request and the model
|
||||
# emits literal <speechcodec_N> tokens instead of a text answer (the system
|
||||
# prompt alone does not prevent this). The 2B must NOT use this recipe: the
|
||||
# official temp-1.0 seed-0 draw measurably hallucinates off-language text on
|
||||
# the 2B (yet closes </think> cleanly, so the seed retry never fires) and
|
||||
# fails the round-trip WER gate — measured, reproduced byte-identically. The
|
||||
# 2B instead keeps its original bare-message / enable_thinking=False /
|
||||
# temperature=0.0 recipe.
|
||||
TEXT_INSTRUCTION = (
|
||||
"SYSTEM & FORMATTING INSTRUCTION: You are Audex, created by NVIDIA based on the Nemotron-Cascade-2 architecture. "
|
||||
"You may output your reasoning, followed by your final response. "
|
||||
"You may format your reasoning block however you like. "
|
||||
"However, your final response must strictly follow these rules: "
|
||||
"* Write in plain, unformatted prose like a book or newspaper article. "
|
||||
"* Do not use markdown, bullet points, lists, or headers. "
|
||||
"* Standard numbers, abbreviations, and symbols are acceptable. "
|
||||
"* [CRITICAL] You must press enter after every single sentence, placing each sentence on its own separate line."
|
||||
)
|
||||
DEFAULT_TEXT = "Hello there! This is the Audex text to speech pipeline."
|
||||
DEFAULT_CAPTION = "Heavy rain falling on a tin roof."
|
||||
# Official quality settings; 1.0 disables guidance.
|
||||
DEFAULT_CFG = {"tts": 1.5, "tta": 3.0, "s2s": 1.5}
|
||||
|
||||
|
||||
def is_30b(model: str) -> bool:
|
||||
# Users pass arbitrary repo IDs or local snapshot paths, so this is a
|
||||
# heuristic (not an exact match against the official 30B-A3B repo ID)
|
||||
# — good enough for an example script; keep in sync with the same
|
||||
# branch in examples/offline_inference/audex/speech_to_speech.py and
|
||||
# tests/e2e/offline_inference/test_audex_s2s.py.
|
||||
return "30b" in model.lower()
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Audex online client (tts / tta / thinker_only / s2s)")
|
||||
parser.add_argument("--mode", choices=("tts", "tta", "thinker_only", "s2s"), default="s2s")
|
||||
parser.add_argument("--host", type=str, default="127.0.0.1")
|
||||
parser.add_argument("--port", type=int, default=8097)
|
||||
parser.add_argument("--model", type=str, default="nvidia/Nemotron-Labs-Audex-2B")
|
||||
parser.add_argument("--text", type=str, default=DEFAULT_TEXT, help="TTS input text (mode=tts).")
|
||||
parser.add_argument("--caption", type=str, default=DEFAULT_CAPTION, help="TTA caption (mode=tta).")
|
||||
parser.add_argument(
|
||||
"--audio-file",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Input WAV for thinker_only/s2s. Defaults to vLLM's mary_had_lamb asset.",
|
||||
)
|
||||
parser.add_argument("--question", type=str, default=ASR_QUESTION, help="Instruction (mode=thinker_only).")
|
||||
parser.add_argument("--output", type=str, default="results/audex_online.wav")
|
||||
parser.add_argument(
|
||||
"--cfg-scale",
|
||||
type=float,
|
||||
default=None,
|
||||
help="CFG strength; defaults per mode (tts/s2s 1.5, tta 3.0). 1.0 disables.",
|
||||
)
|
||||
parser.add_argument("--max-tokens", type=int, default=512)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def _audio_b64(audio_file: str | None) -> str:
|
||||
if audio_file:
|
||||
return base64.b64encode(Path(audio_file).read_bytes()).decode()
|
||||
# Fall back to vLLM's public asset, re-encoded as WAV for the chat API.
|
||||
import io
|
||||
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
from vllm.assets.audio import AudioAsset
|
||||
|
||||
audio, sr = AudioAsset("mary_had_lamb").audio_and_sample_rate
|
||||
buf = io.BytesIO()
|
||||
sf.write(buf, np.asarray(audio, dtype="float32"), int(sr), format="WAV")
|
||||
return base64.b64encode(buf.getvalue()).decode()
|
||||
|
||||
|
||||
def _chat(
|
||||
base_url: str,
|
||||
model: str,
|
||||
messages: list[dict],
|
||||
max_tokens: int,
|
||||
enable_thinking: bool = False,
|
||||
temperature: float = 0.0,
|
||||
top_p: float = 1.0,
|
||||
seed: int = 0,
|
||||
) -> tuple[str, str]:
|
||||
resp = requests.post(
|
||||
f"{base_url}/v1/chat/completions",
|
||||
json={
|
||||
"model": model,
|
||||
"messages": messages,
|
||||
"max_tokens": max_tokens,
|
||||
"temperature": temperature,
|
||||
"top_p": top_p,
|
||||
"seed": seed,
|
||||
# Text-final passes must say so explicitly: without "modalities"
|
||||
# the server falls back to the deployment's configured output
|
||||
# modalities, which for the s2s pipeline include the audio final
|
||||
# stage — routing ASR/chat through code2wav instead of stopping
|
||||
# at stage 0.
|
||||
"modalities": ["text"],
|
||||
# The checkpoint's chat template thinking toggle. The official
|
||||
# cascaded demo transcribes with enable_thinking=False but runs
|
||||
# the chat pass with reasoning ENABLED (its web UI ships the
|
||||
# "Enable reasoning" checkbox checked), priming the generation
|
||||
# prompt with an open "<think>\n" so the model's reasoning stays
|
||||
# inside the think block; _extract_final_response then drops it.
|
||||
# With thinking disabled the 30B leaks its reasoning as plain
|
||||
# text in front of the final response.
|
||||
"chat_template_kwargs": {"enable_thinking": enable_thinking},
|
||||
},
|
||||
timeout=300,
|
||||
)
|
||||
resp.raise_for_status()
|
||||
choice = resp.json()["choices"][0]
|
||||
return choice["message"]["content"].strip(), choice.get("finish_reason") or ""
|
||||
|
||||
|
||||
def _speech(base_url: str, model: str, text: str, cfg_scale: float, out_path: Path) -> None:
|
||||
payload: dict = {"model": model, "input": text, "response_format": "wav"}
|
||||
if cfg_scale > 1.0:
|
||||
payload["extra_params"] = {"cfg_scale": cfg_scale}
|
||||
resp = requests.post(f"{base_url}/v1/audio/speech", json=payload, timeout=600)
|
||||
resp.raise_for_status()
|
||||
out_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
out_path.write_bytes(resp.content)
|
||||
|
||||
|
||||
def _audio_question(audio_b64: str, question: str) -> list[dict]:
|
||||
return [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "input_audio", "input_audio": {"data": audio_b64, "format": "wav"}},
|
||||
{"type": "text", "text": question},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def _tts_text(answer: str) -> str:
|
||||
# The formatting instruction's CRITICAL rule puts each final-response
|
||||
# sentence on its own line, and the official demo TTS-es those sentence
|
||||
# segments (split_sentence_segments). Mirror that per-line segmentation
|
||||
# but keep a single /v1/audio/speech request: speak the leading whole
|
||||
# sentences of the CLEANED final response up to ~200 chars, never
|
||||
# cutting mid-sentence. Markdown markup is stripped as before (it makes
|
||||
# the TTS thinker stop almost immediately).
|
||||
sentences = []
|
||||
for line in answer.splitlines():
|
||||
line = re.sub(r"[*#_`]+", "", line)
|
||||
line = re.sub(r"\s+", " ", line).strip()
|
||||
if line:
|
||||
sentences.append(line)
|
||||
picked: list[str] = []
|
||||
total = 0
|
||||
for sentence in sentences:
|
||||
if picked and total + len(sentence) > 200:
|
||||
break
|
||||
picked.append(sentence)
|
||||
total += len(sentence) + 1
|
||||
return " ".join(picked)
|
||||
|
||||
|
||||
def _extract_final_response(text: str) -> str:
|
||||
# Mirror the official demo's extract_spoken_answer
|
||||
# (cascaded_s2s_web_server.py): the chat pass reasons inside the primed
|
||||
# <think> block, so keep only what follows the LAST closing tag; fall
|
||||
# back to the full text when the model skipped reasoning or never
|
||||
# closed the block. Benign no-op for the 2B's clean answers.
|
||||
if "</think>" in text:
|
||||
answer = text.rsplit("</think>", 1)[-1].strip()
|
||||
if answer:
|
||||
return answer
|
||||
return text.strip()
|
||||
|
||||
|
||||
def _clean_transcript(text: str) -> str:
|
||||
# The full checkpoint prefixes transcripts with a language-tag sentence
|
||||
# and quotes the content; extract the quoted payload when present. Span
|
||||
# first-to-last quote (not split on every apostrophe) so contractions
|
||||
# inside the transcript (e.g. "don't") survive intact.
|
||||
text = text.strip()
|
||||
match = re.search(r"'(.*)'", text, flags=re.S)
|
||||
if match and match.group(1).strip():
|
||||
return match.group(1).strip()
|
||||
return text
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
base_url = f"http://{args.host}:{args.port}"
|
||||
cfg_scale = args.cfg_scale if args.cfg_scale is not None else DEFAULT_CFG.get(args.mode, 1.0)
|
||||
out_path = Path(args.output)
|
||||
|
||||
if args.mode == "tts":
|
||||
_speech(base_url, args.model, args.text, cfg_scale, out_path)
|
||||
print(f"speech: {out_path.stat().st_size} bytes -> {out_path}")
|
||||
return
|
||||
|
||||
if args.mode == "tta":
|
||||
_speech(base_url, args.model, args.caption, cfg_scale, out_path)
|
||||
print(f"audio : {out_path.stat().st_size} bytes -> {out_path}")
|
||||
return
|
||||
|
||||
audio_b64 = _audio_b64(args.audio_file)
|
||||
|
||||
if args.mode == "thinker_only":
|
||||
content, _ = _chat(base_url, args.model, _audio_question(audio_b64, args.question), args.max_tokens)
|
||||
print(f"answer: {_extract_final_response(content)}")
|
||||
return
|
||||
|
||||
# s2s — the official three-pass cascade.
|
||||
asr_content, _ = _chat(base_url, args.model, _audio_question(audio_b64, ASR_QUESTION), args.max_tokens)
|
||||
transcript = _clean_transcript(asr_content)
|
||||
print(f"[1/3] transcript : {transcript}")
|
||||
if not transcript:
|
||||
raise SystemExit("ASR pass produced an empty transcript; aborting cascade")
|
||||
|
||||
# The chat-pass recipe is size-conditional (see the TEXT_INSTRUCTION
|
||||
# comment above): the 30B requires the official recipe, the 2B keeps its
|
||||
# original bare-message recipe.
|
||||
if is_30b(args.model):
|
||||
# Official recipe: system prompt + formatting instruction prefixed
|
||||
# to the transcript, reasoning enabled, and the official demo's
|
||||
# text-pass sampling (its defaults: --text-temperature 1.0,
|
||||
# --text-top-p 0.95, --text-max-tokens 2048). Near-greedy decoding
|
||||
# makes the 30B loop inside its reasoning block until the token
|
||||
# budget runs out without ever closing </think>. Even at the
|
||||
# official temperature some trajectories spin past the budget, so
|
||||
# retry a few deterministic seeds; a natural stop without </think>
|
||||
# is a complete answer, and if every seed spins we keep the full
|
||||
# text like the official demo does.
|
||||
chat_messages = [
|
||||
{"role": "system", "content": SYSTEM_PROMPT},
|
||||
{"role": "user", "content": f"{TEXT_INSTRUCTION}\n\n{transcript}"},
|
||||
]
|
||||
raw = ""
|
||||
for seed in (0, 1, 2):
|
||||
raw, finish_reason = _chat(
|
||||
base_url,
|
||||
args.model,
|
||||
chat_messages,
|
||||
max(args.max_tokens, 2048),
|
||||
enable_thinking=True,
|
||||
temperature=1.0,
|
||||
top_p=0.95,
|
||||
seed=seed,
|
||||
)
|
||||
if "</think>" in raw or finish_reason != "length":
|
||||
break
|
||||
print(f"[2/3] chat seed {seed} hit the token budget mid-reasoning; retrying")
|
||||
answer = _extract_final_response(raw)
|
||||
else:
|
||||
# 2B: bare user message (no system/instruction prefix), reasoning
|
||||
# disabled, temperature 0.0, no retry — its original recipe. The
|
||||
# official temp-1.0 seed-0 recipe measurably makes the 2B
|
||||
# hallucinate off-language answers (see the module-level comment).
|
||||
chat_messages = [{"role": "user", "content": transcript}]
|
||||
raw, _ = _chat(
|
||||
base_url,
|
||||
args.model,
|
||||
chat_messages,
|
||||
args.max_tokens,
|
||||
enable_thinking=False,
|
||||
temperature=0.0,
|
||||
)
|
||||
answer = _extract_final_response(raw)
|
||||
print(f"[2/3] answer : {answer}")
|
||||
if not answer:
|
||||
raise SystemExit("Chat pass produced an empty answer; aborting cascade")
|
||||
|
||||
_speech(base_url, args.model, _tts_text(answer), cfg_scale, out_path)
|
||||
print(f"[3/3] speech : {out_path.stat().st_size} bytes -> {out_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Executable
+57
@@ -0,0 +1,57 @@
|
||||
#!/bin/bash
|
||||
# Launch the vLLM-Omni server for Audex (Nemotron-Labs-Audex-2B).
|
||||
#
|
||||
# MODE picks the deployment (see README.md for the capability matrix):
|
||||
# tts text -> speech /v1/audio/speech (default)
|
||||
# tta caption -> general audio /v1/audio/speech
|
||||
# thinker_only audio -> text /v1/chat/completions
|
||||
# s2s cascaded speech-to-speech (both endpoints)
|
||||
#
|
||||
# Pass the HF repo ROOT as MODEL — per-stage subfolders resolve
|
||||
# automatically.
|
||||
#
|
||||
# Usage:
|
||||
# ./run_server.sh # 2B tts on port 8097, GPU 0
|
||||
# SIZE=30b ./run_server.sh # 30B-A3B tts (explicit 30B yaml)
|
||||
# SIZE=30b MODE=s2s PORT=8098 ./run_server.sh
|
||||
# MODE=tta ./run_server.sh # needs XCodec1 (auto-downloaded)
|
||||
# MODEL=/path/to/local/snapshot ./run_server.sh
|
||||
|
||||
set -e
|
||||
|
||||
MODE="${MODE:-${1:-tts}}"
|
||||
SIZE="${SIZE:-2b}"
|
||||
if [ "$SIZE" = "30b" ]; then
|
||||
MODEL="${MODEL:-nvidia/Nemotron-Labs-Audex-30B-A3B}"
|
||||
YAML_SUFFIX="_30b"
|
||||
elif [ "$SIZE" = "2b" ]; then
|
||||
MODEL="${MODEL:-nvidia/Nemotron-Labs-Audex-2B}"
|
||||
YAML_SUFFIX=""
|
||||
else
|
||||
echo "Unknown SIZE '$SIZE'; expected 2b|30b" >&2
|
||||
exit 1
|
||||
fi
|
||||
PORT="${PORT:-8097}"
|
||||
GPUS="${GPUS:-0}"
|
||||
|
||||
REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../.." && pwd)"
|
||||
DEPLOY_YAML="$REPO_ROOT/vllm_omni/deploy/audex_${MODE}${YAML_SUFFIX}.yaml"
|
||||
if [ ! -f "$DEPLOY_YAML" ]; then
|
||||
echo "Unknown MODE '$MODE' / SIZE '$SIZE' (no $DEPLOY_YAML); expected MODE=tts|tta|thinker_only|s2s, SIZE=2b|30b" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "Starting Audex server"
|
||||
echo " MODE=$MODE ($DEPLOY_YAML)"
|
||||
echo " SIZE=$SIZE"
|
||||
echo " MODEL=$MODEL"
|
||||
echo " PORT=$PORT"
|
||||
echo " CUDA_VISIBLE_DEVICES=$GPUS"
|
||||
|
||||
CUDA_VISIBLE_DEVICES="$GPUS" \
|
||||
vllm-omni serve "$MODEL" \
|
||||
--host 0.0.0.0 \
|
||||
--port "$PORT" \
|
||||
--trust-remote-code \
|
||||
--deploy-config "$DEPLOY_YAML" \
|
||||
--omni
|
||||
@@ -28,8 +28,11 @@ If image.command is set, uses that as a full override.
|
||||
- {{ .Values.omniArgs.numGpus | quote }}
|
||||
{{- end }}
|
||||
{{- if .Values.omniArgs.stageConfigsPath }}
|
||||
- "--stage-configs-path"
|
||||
- {{ .Values.omniArgs.stageConfigsPath | quote }}
|
||||
{{- fail "omniArgs.stageConfigsPath has been removed; use omniArgs.deployConfigPath" }}
|
||||
{{- end }}
|
||||
{{- if .Values.omniArgs.deployConfigPath }}
|
||||
- "--deploy-config"
|
||||
- {{ .Values.omniArgs.deployConfigPath | quote }}
|
||||
{{- end }}
|
||||
{{- if and .Values.omniArgs.cacheBackend (ne .Values.omniArgs.cacheBackend "none") }}
|
||||
- "--cache-backend"
|
||||
|
||||
@@ -78,6 +78,26 @@ tests:
|
||||
path: spec.template.spec.containers[0].command
|
||||
content: "--enable-layerwise-offload"
|
||||
|
||||
- it: should include deploy config in command
|
||||
set:
|
||||
omniArgs:
|
||||
deployConfigPath: "/configs/deploy.yaml"
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].command[8]
|
||||
value: "--deploy-config"
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].command[9]
|
||||
value: "/configs/deploy.yaml"
|
||||
|
||||
- it: should reject removed stage config value
|
||||
set:
|
||||
omniArgs:
|
||||
stageConfigsPath: "/configs/legacy.yaml"
|
||||
asserts:
|
||||
- failedTemplate:
|
||||
errorMessage: "omniArgs.stageConfigsPath has been removed; use omniArgs.deployConfigPath"
|
||||
|
||||
- it: should use full command override when image.command is set
|
||||
set:
|
||||
image:
|
||||
|
||||
@@ -67,8 +67,8 @@ omniArgs:
|
||||
enableCpuOffload: false
|
||||
# -- Number of GPUs for diffusion inference (empty = auto)
|
||||
numGpus:
|
||||
# -- Path to stage configs file (empty = auto-detected from model)
|
||||
stageConfigsPath:
|
||||
# -- Path to deploy config file (empty = auto-detected from model)
|
||||
deployConfigPath:
|
||||
# -- Cache backend: none, tea_cache, or cache_dit
|
||||
cacheBackend: "none"
|
||||
# -- Default sampling params as a JSON string
|
||||
|
||||
@@ -23,17 +23,18 @@ including `librosa`.
|
||||
|
||||
The deploy config auto-loads via `--omni`.
|
||||
The default `vllm_omni/deploy/minicpmo_4_5.yaml` keeps all three stages on
|
||||
logical device 0 with memory budgets of 55%, 22%, and 22%. The profile admits
|
||||
at most four sequences per stage and bounds startup video profiling to 32
|
||||
frames per video. For throughput,`minicpmo_4_5_batching.yaml` gives the Thinker
|
||||
logical device 0 with memory budgets of 55%, 15%, and 15%, leaving headroom
|
||||
for runtime kernels such as the HiFi-GAN vocoder's cuDNN workspace. The
|
||||
profile admits at most four sequences per stage and bounds startup video
|
||||
profiling to 32 frames per video. For throughput,
|
||||
`minicpmo_4_5_2gpu.yaml` gives the Thinker
|
||||
GPU 0 (90%) and colocates the Talker (55%) and Code2Wav (35%) on GPU 1. That
|
||||
profile admits at most four concurrent sequences per stage.
|
||||
|
||||
| deploy config | GPUs | Notes |
|
||||
|---|---|---|
|
||||
| `minicpmo_4_5.yaml` (default) | 1 | Memory-constrained compatibility layout. |
|
||||
| `minicpmo_4_5_batching.yaml` | 2 | Recommended continuous-batching layout; Talker and Code2Wav share GPU 1. |
|
||||
| `minicpmo_4_5_2gpu.yaml` | 2 | Backward-compatible alias for `minicpmo_4_5_batching.yaml`. |
|
||||
| `minicpmo_4_5_2gpu.yaml` | 2 | Recommended continuous-batching layout; Talker and Code2Wav share GPU 1. |
|
||||
| `minicpmo_4_5_3gpu.yaml` | 3 | One GPU per stage. |
|
||||
| `minicpmo_4_5_8x4090.yaml` | 8 | Full 8x4090 layout. |
|
||||
| `minicpmo_4_5_duplex.yaml` | 1 | Experimental native full-duplex overlay. |
|
||||
@@ -89,17 +90,18 @@ python openai_chat_completion_client_for_multimodal_generation.py \
|
||||
--prompt "Briefly introduce yourself."
|
||||
```
|
||||
|
||||
Streaming text + audio:
|
||||
Streaming text + audio (use `--stream`):
|
||||
|
||||
```bash
|
||||
python streaming_chat_completion.py \
|
||||
--base-url http://localhost:8099/v1 \
|
||||
--output minicpmo_stream.wav
|
||||
python openai_chat_completion_client_for_multimodal_generation.py \
|
||||
--query-type text \
|
||||
--prompt "Briefly introduce yourself." \
|
||||
--port 8099 \
|
||||
--stream
|
||||
```
|
||||
|
||||
The example prints text deltas immediately and joins the independently encoded
|
||||
24 kHz WAV audio deltas into one valid WAV file. Add `--text-only` to skip
|
||||
Talker and Code2Wav.
|
||||
The client prints text deltas as they arrive and saves streamed audio chunks
|
||||
to WAV files.
|
||||
|
||||
Shared helpers also work if you pass MiniCPM defaults yourself:
|
||||
|
||||
|
||||
@@ -6,6 +6,7 @@ class FullDuplexPcmPlayback extends AudioWorkletProcessor {
|
||||
this.playedFrames = 0;
|
||||
this.underrunFrames = 0;
|
||||
this.drain = null;
|
||||
this.terminalFadeFrames = 0;
|
||||
this.started = false;
|
||||
this.activeResponseId = null;
|
||||
this.initialBufferFrames = Math.round(sampleRate * 0.2);
|
||||
@@ -31,6 +32,7 @@ class FullDuplexPcmPlayback extends AudioWorkletProcessor {
|
||||
}
|
||||
} else if (message.type === 'drain') {
|
||||
this.drain = { responseId: message.responseId || null };
|
||||
this.terminalFadeFrames = Math.min(this.fadeFrames, this.bufferedFrames());
|
||||
if (!this.started && this.bufferedFrames() > 0) {
|
||||
this.bufferWaitFrames = 0;
|
||||
this.startPlayback();
|
||||
@@ -42,6 +44,7 @@ class FullDuplexPcmPlayback extends AudioWorkletProcessor {
|
||||
this.playedFrames = 0;
|
||||
this.underrunFrames = 0;
|
||||
this.drain = null;
|
||||
this.terminalFadeFrames = 0;
|
||||
this.started = false;
|
||||
this.activeResponseId = null;
|
||||
this.bufferWaitFrames = this.initialBufferFrames;
|
||||
@@ -76,6 +79,7 @@ class FullDuplexPcmPlayback extends AudioWorkletProcessor {
|
||||
this.playedFrames = 0;
|
||||
this.underrunFrames = 0;
|
||||
this.drain = null;
|
||||
this.terminalFadeFrames = 0;
|
||||
this.started = false;
|
||||
this.activeResponseId = null;
|
||||
this.bufferWaitFrames = this.initialBufferFrames;
|
||||
@@ -116,6 +120,7 @@ class FullDuplexPcmPlayback extends AudioWorkletProcessor {
|
||||
while (target < output.length && this.queue.length > 0) {
|
||||
const pcm = this.queue[0];
|
||||
const count = Math.min(output.length - target, pcm.length - this.offset);
|
||||
const remainingBeforeChunk = this.bufferedFrames();
|
||||
for (let index = 0; index < count; index += 1) {
|
||||
let sample = pcm[this.offset + index] / 32768;
|
||||
if (this.fadeInFrames > 0) {
|
||||
@@ -123,6 +128,16 @@ class FullDuplexPcmPlayback extends AudioWorkletProcessor {
|
||||
sample *= elapsed / this.fadeFrames;
|
||||
this.fadeInFrames -= 1;
|
||||
}
|
||||
const remainingFrames = remainingBeforeChunk - index;
|
||||
if (
|
||||
this.drain
|
||||
&& this.terminalFadeFrames > 0
|
||||
&& remainingFrames <= this.terminalFadeFrames
|
||||
) {
|
||||
sample *= this.terminalFadeFrames === 1
|
||||
? 0
|
||||
: (remainingFrames - 1) / (this.terminalFadeFrames - 1);
|
||||
}
|
||||
output[target + index] = sample;
|
||||
}
|
||||
target += count;
|
||||
|
||||
@@ -3,7 +3,7 @@
|
||||
#
|
||||
# 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).
|
||||
# 8x4090 deploy config at vllm_omni/deploy/minicpmo_4_5_8x4090.yaml).
|
||||
set -e
|
||||
|
||||
HERE="$(cd "$(dirname "$0")" && pwd)"
|
||||
|
||||
@@ -1,71 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Stream MiniCPM-o 4.5 text and audio from /v1/chat/completions."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import io
|
||||
import wave
|
||||
from pathlib import Path
|
||||
|
||||
from openai import OpenAI
|
||||
|
||||
|
||||
def append_wav_chunk(encoded: str, pcm_parts: list[bytes]) -> tuple[int, int, int]:
|
||||
"""Decode one base64 WAV delta and append its PCM frames."""
|
||||
with wave.open(io.BytesIO(base64.b64decode(encoded)), "rb") as wav:
|
||||
channels = wav.getnchannels()
|
||||
sample_width = wav.getsampwidth()
|
||||
sample_rate = wav.getframerate()
|
||||
pcm_parts.append(wav.readframes(wav.getnframes()))
|
||||
return channels, sample_width, sample_rate
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--base-url", default="http://localhost:8099/v1")
|
||||
parser.add_argument("--model", default="openbmb/MiniCPM-o-4_5")
|
||||
parser.add_argument("--prompt", default="Say hello, then introduce vLLM in one sentence.")
|
||||
parser.add_argument("--output", type=Path, default=Path("minicpmo_stream.wav"))
|
||||
parser.add_argument("--text-only", action="store_true")
|
||||
args = parser.parse_args()
|
||||
|
||||
modalities = ["text"] if args.text_only else ["text", "audio"]
|
||||
stream = OpenAI(base_url=args.base_url, api_key="EMPTY").chat.completions.create(
|
||||
model=args.model,
|
||||
messages=[{"role": "user", "content": args.prompt}],
|
||||
modalities=modalities,
|
||||
stream=True,
|
||||
extra_body={"chat_template_kwargs": {"use_tts_template": not args.text_only}},
|
||||
)
|
||||
|
||||
pcm_parts: list[bytes] = []
|
||||
audio_format: tuple[int, int, int] | None = None
|
||||
for chunk in stream:
|
||||
for choice in chunk.choices:
|
||||
content = getattr(choice.delta, "content", None)
|
||||
if not content:
|
||||
continue
|
||||
if getattr(chunk, "modality", None) == "audio":
|
||||
current_format = append_wav_chunk(content, pcm_parts)
|
||||
if audio_format is not None and current_format != audio_format:
|
||||
raise ValueError(f"Audio format changed: {audio_format} -> {current_format}")
|
||||
audio_format = current_format
|
||||
elif getattr(chunk, "modality", None) == "text":
|
||||
print(content, end="", flush=True)
|
||||
|
||||
print()
|
||||
if pcm_parts and audio_format is not None:
|
||||
channels, sample_width, sample_rate = audio_format
|
||||
with wave.open(str(args.output), "wb") as wav:
|
||||
wav.setnchannels(channels)
|
||||
wav.setsampwidth(sample_width)
|
||||
wav.setframerate(sample_rate)
|
||||
wav.writeframes(b"".join(pcm_parts))
|
||||
print(f"Audio saved to {args.output} ({sample_rate} Hz)")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -12,9 +12,9 @@ Please refer to [README.md](../../../README.md)
|
||||
vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091
|
||||
```
|
||||
|
||||
If you have custom stage configs file, launch the server with command below
|
||||
To use a custom deploy config, launch the server with the command below:
|
||||
```bash
|
||||
vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
|
||||
vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --deploy-config /path/to/deploy_config.yaml
|
||||
```
|
||||
|
||||
### Send Multi-modal Request
|
||||
@@ -173,7 +173,7 @@ The script supports the following arguments:
|
||||
- `--model`: Model name/path (default: Qwen/Qwen2.5-Omni-7B)
|
||||
- `--server-port`: Port for vLLM server (default: 8091)
|
||||
- `--gradio-port`: Port for Gradio demo (default: 7861)
|
||||
- `--stage-configs-path`: Path to custom stage configs YAML file (optional)
|
||||
- `--deploy-config`: Path to a custom deploy config YAML file (optional)
|
||||
- `--server-host`: Host for vLLM server (default: 0.0.0.0)
|
||||
- `--gradio-ip`: IP for Gradio demo (default: 127.0.0.1)
|
||||
- `--share`: Share Gradio demo publicly (creates a public link)
|
||||
@@ -186,9 +186,9 @@ The script supports the following arguments:
|
||||
vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091
|
||||
```
|
||||
|
||||
If you have custom stage configs file:
|
||||
To use a custom deploy config:
|
||||
```bash
|
||||
vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
|
||||
vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --deploy-config /path/to/deploy_config.yaml
|
||||
```
|
||||
|
||||
**Step 2: Run the Gradio demo**
|
||||
|
||||
@@ -13,7 +13,7 @@ set -e
|
||||
MODEL="Qwen/Qwen2.5-Omni-7B"
|
||||
SERVER_PORT=8091
|
||||
GRADIO_PORT=7861
|
||||
STAGE_CONFIGS_PATH=""
|
||||
DEPLOY_CONFIG_PATH=""
|
||||
SERVER_HOST="0.0.0.0"
|
||||
GRADIO_IP="127.0.0.1"
|
||||
GRADIO_SHARE=false
|
||||
@@ -33,8 +33,8 @@ while [[ $# -gt 0 ]]; do
|
||||
GRADIO_PORT="$2"
|
||||
shift 2
|
||||
;;
|
||||
--stage-configs-path)
|
||||
STAGE_CONFIGS_PATH="$2"
|
||||
--deploy-config)
|
||||
DEPLOY_CONFIG_PATH="$2"
|
||||
shift 2
|
||||
;;
|
||||
--server-host)
|
||||
@@ -56,7 +56,7 @@ while [[ $# -gt 0 ]]; do
|
||||
echo " --model MODEL Model name/path (default: Qwen/Qwen2.5-Omni-7B)"
|
||||
echo " --server-port PORT Port for vLLM server (default: 8091)"
|
||||
echo " --gradio-port PORT Port for Gradio demo (default: 7861)"
|
||||
echo " --stage-configs-path PATH Path to custom stage configs YAML file (optional)"
|
||||
echo " --deploy-config PATH Path to custom deploy config YAML file (optional)"
|
||||
echo " --server-host HOST Host for vLLM server (default: 0.0.0.0)"
|
||||
echo " --gradio-ip IP IP for Gradio demo (default: 127.0.0.1)"
|
||||
echo " --share Share Gradio demo publicly"
|
||||
@@ -87,8 +87,8 @@ echo "=========================================="
|
||||
|
||||
# Build vLLM server command
|
||||
SERVER_CMD=("vllm" "serve" "$MODEL" "--omni" "--port" "$SERVER_PORT" "--host" "$SERVER_HOST")
|
||||
if [ -n "$STAGE_CONFIGS_PATH" ]; then
|
||||
SERVER_CMD+=("--stage-configs-path" "$STAGE_CONFIGS_PATH")
|
||||
if [ -n "$DEPLOY_CONFIG_PATH" ]; then
|
||||
SERVER_CMD+=("--deploy-config" "$DEPLOY_CONFIG_PATH")
|
||||
fi
|
||||
|
||||
# Function to cleanup on exit
|
||||
|
||||
@@ -1,36 +0,0 @@
|
||||
# Stage config for running Qwen3-Omni-MoE-Thinking (text-only output)
|
||||
# This config is for models like Qwen3-Omni-30B-A3B-Thinking that only have the
|
||||
# thinker component and do not support audio output.
|
||||
#
|
||||
# Single stage: Thinker (multimodal understanding + text generation)
|
||||
|
||||
# The following config has been verified on 2x H100-80G GPUs.
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
runtime:
|
||||
devices: "0,1"
|
||||
engine_args:
|
||||
model_stage: thinker
|
||||
max_num_seqs: 1
|
||||
model_arch: Qwen3OmniMoeForConditionalGeneration
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
gpu_memory_utilization: 0.9
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
engine_output_type: text
|
||||
distributed_executor_backend: "mp"
|
||||
enable_prefix_caching: true
|
||||
hf_config_name: thinker_config
|
||||
tensor_parallel_size: 2
|
||||
final_output: true
|
||||
final_output_type: text
|
||||
is_comprehension: true
|
||||
default_sampling_params:
|
||||
temperature: 0.4
|
||||
top_p: 0.9
|
||||
top_k: 1
|
||||
max_tokens: 2048
|
||||
seed: 42
|
||||
detokenize: True
|
||||
repetition_penalty: 1.05
|
||||
@@ -13,7 +13,7 @@ set -e
|
||||
MODEL="Qwen/Qwen3-Omni-30B-A3B-Instruct"
|
||||
SERVER_PORT=8091
|
||||
GRADIO_PORT=7861
|
||||
STAGE_CONFIGS_PATH=""
|
||||
DEPLOY_CONFIG=""
|
||||
SERVER_HOST="0.0.0.0"
|
||||
GRADIO_IP="127.0.0.1"
|
||||
GRADIO_SHARE=false
|
||||
@@ -33,8 +33,8 @@ while [[ $# -gt 0 ]]; do
|
||||
GRADIO_PORT="$2"
|
||||
shift 2
|
||||
;;
|
||||
--stage-configs-path)
|
||||
STAGE_CONFIGS_PATH="$2"
|
||||
--deploy-config)
|
||||
DEPLOY_CONFIG="$2"
|
||||
shift 2
|
||||
;;
|
||||
--server-host)
|
||||
@@ -56,7 +56,7 @@ while [[ $# -gt 0 ]]; do
|
||||
echo " --model MODEL Model name/path (default: Qwen/Qwen3-Omni-30B-A3B-Instruct)"
|
||||
echo " --server-port PORT Port for vLLM server (default: 8091)"
|
||||
echo " --gradio-port PORT Port for Gradio demo (default: 7861)"
|
||||
echo " --stage-configs-path PATH Path to custom stage configs YAML file (optional)"
|
||||
echo " --deploy-config PATH Path to custom deploy config YAML file (optional)"
|
||||
echo " --server-host HOST Host for vLLM server (default: 0.0.0.0)"
|
||||
echo " --gradio-ip IP IP for Gradio demo (default: 127.0.0.1)"
|
||||
echo " --share Share Gradio demo publicly"
|
||||
@@ -87,8 +87,8 @@ echo "=========================================="
|
||||
|
||||
# Build vLLM server command
|
||||
SERVER_CMD=("vllm" "serve" "$MODEL" "--omni" "--port" "$SERVER_PORT" "--host" "$SERVER_HOST")
|
||||
if [ -n "$STAGE_CONFIGS_PATH" ]; then
|
||||
SERVER_CMD+=("--stage-configs-path" "$STAGE_CONFIGS_PATH")
|
||||
if [ -n "$DEPLOY_CONFIG" ]; then
|
||||
SERVER_CMD+=("--deploy-config" "$DEPLOY_CONFIG")
|
||||
fi
|
||||
|
||||
# Function to cleanup on exit
|
||||
|
||||
@@ -336,7 +336,7 @@ vllm serve OpenMOSS-Team/MOSS-TTS-Nano --omni --port 8091
|
||||
# or:
|
||||
./moss_tts_nano/run_server.sh
|
||||
```
|
||||
The deploy config at `vllm_omni/deploy/moss_tts_nano.yaml` auto-loads; no `--stage-configs-path`, `--trust-remote-code`, or `--enforce-eager` flags are needed.
|
||||
The deploy config at `vllm_omni/deploy/moss_tts_nano.yaml` auto-loads; no `--deploy-config`, `--trust-remote-code`, or `--enforce-eager` flags are needed.
|
||||
|
||||
### Sending requests
|
||||
```bash
|
||||
@@ -537,12 +537,19 @@ curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
Raw PCM streaming requires `stream_format="audio"`, `response_format="pcm"`, and `async_chunk: true` on the stage config (default in `qwen3_tts.yaml`). `speed` is not supported when streaming.
|
||||
|
||||
### Streaming WebSocket
|
||||
The `/v1/audio/speech/stream` endpoint accepts text incrementally, splits it at sentence boundaries, and emits one PCM stream per sentence:
|
||||
The `/v1/audio/speech/stream` endpoint accepts text incrementally and synthesizes the buffered text as one continuous request on `input.done`:
|
||||
```bash
|
||||
python qwen3_tts/streaming_speech_client.py --text "Hello world. How are you? I am fine."
|
||||
python qwen3_tts/streaming_speech_client.py --text "..." --simulate-stt --stt-delay 0.1
|
||||
```
|
||||
|
||||
`input.done` flushes without closing, so repeating `--text` synthesizes several utterances over one connection:
|
||||
```bash
|
||||
python qwen3_tts/streaming_speech_client.py \
|
||||
--text "First utterance." \
|
||||
--text "Second utterance, same connection."
|
||||
```
|
||||
|
||||
To receive word-level timestamps, launch the server with a forced aligner:
|
||||
```bash
|
||||
vllm-omni serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
|
||||
|
||||
@@ -7,6 +7,12 @@ Usage:
|
||||
# Send full text at once
|
||||
python streaming_speech_client.py --text "Hello world. How are you? I am fine."
|
||||
|
||||
# Several utterances over one connection: input.done flushes each of them
|
||||
# and the connection stays open, so only the first pays the handshake
|
||||
python streaming_speech_client.py \
|
||||
--text "First utterance." \
|
||||
--text "Second utterance, same connection."
|
||||
|
||||
# Pick a built-in speaker for CustomVoice models
|
||||
python streaming_speech_client.py \
|
||||
--text "打开电子枪,关闭扫描,切换到低倍模式。" \
|
||||
@@ -64,18 +70,27 @@ except ImportError:
|
||||
raise SystemExit(1)
|
||||
|
||||
|
||||
def save_audio_file(output_dir: str, sentence_index: int, response_format: str, chunks: list[bytes]) -> str:
|
||||
def frame_basename(msg: dict) -> str:
|
||||
"""Name a file after the utterance and sentence a frame belongs to.
|
||||
|
||||
Every utterance is sentence 0, so the utterance index is what keeps the
|
||||
files of one connection from overwriting each other.
|
||||
"""
|
||||
return f"utterance_{msg['utterance_index']:03d}_sentence_{msg['sentence_index']:03d}"
|
||||
|
||||
|
||||
def save_audio_file(output_dir: str, basename: str, response_format: str, chunks: list[bytes]) -> str:
|
||||
audio_bytes = b"".join(chunks)
|
||||
if response_format == "pcm":
|
||||
raw_path = os.path.join(output_dir, f"sentence_{sentence_index:03d}.pcm")
|
||||
raw_path = os.path.join(output_dir, f"{basename}.pcm")
|
||||
with open(raw_path, "wb") as f:
|
||||
f.write(audio_bytes)
|
||||
|
||||
wav_path = os.path.join(output_dir, f"sentence_{sentence_index:03d}.wav")
|
||||
wav_path = os.path.join(output_dir, f"{basename}.wav")
|
||||
save_pcm_wav(wav_path, audio_bytes)
|
||||
return wav_path
|
||||
|
||||
filename = os.path.join(output_dir, f"sentence_{sentence_index:03d}.{response_format}")
|
||||
filename = os.path.join(output_dir, f"{basename}.{response_format}")
|
||||
with open(filename, "wb") as f:
|
||||
f.write(audio_bytes)
|
||||
return filename
|
||||
@@ -95,183 +110,207 @@ def timestamp_words_text(timestamps: list[dict] | None) -> str:
|
||||
return " ".join(str(item.get("word", "")) for item in timestamps if item.get("word"))
|
||||
|
||||
|
||||
async def send_utterance(ws, text: str, simulate_stt: bool, stt_delay: float) -> None:
|
||||
"""Send one utterance's text, then input.done to flush it."""
|
||||
if simulate_stt:
|
||||
words = text.split(" ")
|
||||
for i, word in enumerate(words):
|
||||
chunk = word + (" " if i < len(words) - 1 else "")
|
||||
await ws.send(
|
||||
json.dumps(
|
||||
{
|
||||
"type": "input.text",
|
||||
"text": chunk,
|
||||
}
|
||||
)
|
||||
)
|
||||
print(f" Sent: {chunk!r}")
|
||||
await asyncio.sleep(stt_delay)
|
||||
else:
|
||||
await ws.send(
|
||||
json.dumps(
|
||||
{
|
||||
"type": "input.text",
|
||||
"text": text,
|
||||
}
|
||||
)
|
||||
)
|
||||
print(f"Sent full text: {text!r}")
|
||||
|
||||
await ws.send(json.dumps({"type": "input.done"}))
|
||||
print("Sent input.done")
|
||||
|
||||
|
||||
async def receive_utterance(
|
||||
ws,
|
||||
output_dir: str,
|
||||
response_format: str,
|
||||
word_timestamps: bool,
|
||||
save_chunks: bool,
|
||||
) -> None:
|
||||
"""Consume frames until session.done marks the flushed utterance complete."""
|
||||
current_label = "utterance 0"
|
||||
current_sentence_text = ""
|
||||
current_chunks: list[bytes] = []
|
||||
current_timestamps: list[dict] = []
|
||||
# Distinguish `null` (aligner failed) from `[]` (silence / no tokens):
|
||||
# the sentence is a failure only if no frame ever carried timestamps.
|
||||
alignment_frame_seen = False
|
||||
|
||||
while True:
|
||||
message = await ws.recv()
|
||||
|
||||
if isinstance(message, bytes):
|
||||
current_chunks.append(message)
|
||||
print(f" Received audio chunk for {current_label}: {len(message)} bytes")
|
||||
else:
|
||||
# JSON frame
|
||||
msg = json.loads(message)
|
||||
msg_type = msg.get("type")
|
||||
|
||||
if msg_type == "audio.start":
|
||||
current_label = f"utterance {msg['utterance_index']}"
|
||||
current_sentence_text = msg.get("sentence_text", "")
|
||||
current_chunks = []
|
||||
current_timestamps = []
|
||||
alignment_frame_seen = False
|
||||
print(f" [{current_label}] Generating: {current_sentence_text!r}")
|
||||
elif msg_type == "audio.chunk":
|
||||
audio = base64.b64decode(msg["audio_b64"])
|
||||
current_chunks.append(audio)
|
||||
chunk_timestamps = msg.get("timestamps")
|
||||
# Sentence-level alignment sends a trailing timestamp-only
|
||||
# frame with empty audio; don't write an empty chunk file.
|
||||
if save_chunks and audio:
|
||||
chunk_base = os.path.join(
|
||||
output_dir,
|
||||
f"{frame_basename(msg)}_chunk_{msg['chunk_id']:03d}",
|
||||
)
|
||||
chunk_pcm = f"{chunk_base}.pcm"
|
||||
chunk_wav = f"{chunk_base}.wav"
|
||||
chunk_json = f"{chunk_base}_timestamps.json"
|
||||
with open(chunk_pcm, "wb") as f:
|
||||
f.write(audio)
|
||||
save_pcm_wav(chunk_wav, audio, int(msg.get("sample_rate", 24000)))
|
||||
with open(chunk_json, "w", encoding="utf-8") as f:
|
||||
json.dump(
|
||||
{
|
||||
"utterance_index": msg["utterance_index"],
|
||||
"sentence_index": msg["sentence_index"],
|
||||
"chunk_id": msg["chunk_id"],
|
||||
"chunk_start_ms": msg.get("chunk_start_ms"),
|
||||
"chunk_end_ms": msg.get("chunk_end_ms"),
|
||||
"sentence_text": current_sentence_text,
|
||||
"covered_text": timestamp_words_text(chunk_timestamps),
|
||||
"timestamps": chunk_timestamps,
|
||||
},
|
||||
f,
|
||||
ensure_ascii=False,
|
||||
indent=2,
|
||||
)
|
||||
if chunk_timestamps is None:
|
||||
print(f" [{current_label}] chunk {msg['chunk_id']}: {len(audio)} bytes, timestamps unavailable")
|
||||
else:
|
||||
alignment_frame_seen = True
|
||||
current_timestamps.extend(chunk_timestamps)
|
||||
print(
|
||||
f" [{current_label}] chunk {msg['chunk_id']}: "
|
||||
f"{len(audio)} bytes, {len(chunk_timestamps)} timestamp(s)"
|
||||
)
|
||||
if save_chunks and audio:
|
||||
print(f" saved chunk -> {chunk_wav}")
|
||||
elif msg_type == "audio.done":
|
||||
filename = save_audio_file(
|
||||
output_dir,
|
||||
frame_basename(msg),
|
||||
response_format,
|
||||
current_chunks,
|
||||
)
|
||||
if word_timestamps:
|
||||
# `null` on failure, else the accumulated list (`[]` for silence).
|
||||
sentence_timestamps = current_timestamps if alignment_frame_seen else None
|
||||
ts_filename = os.path.join(
|
||||
output_dir,
|
||||
f"{frame_basename(msg)}_timestamps.json",
|
||||
)
|
||||
with open(ts_filename, "w", encoding="utf-8") as f:
|
||||
json.dump(
|
||||
{
|
||||
"utterance_index": msg["utterance_index"],
|
||||
"sentence_index": msg["sentence_index"],
|
||||
"sentence_text": current_sentence_text,
|
||||
"covered_text": timestamp_words_text(sentence_timestamps),
|
||||
"timestamps": sentence_timestamps,
|
||||
},
|
||||
f,
|
||||
ensure_ascii=False,
|
||||
indent=2,
|
||||
)
|
||||
print(
|
||||
f" [{current_label}] Done"
|
||||
f" bytes={msg.get('total_bytes', len(b''.join(current_chunks)))}"
|
||||
f" error={msg.get('error', False)}"
|
||||
f" -> {filename}"
|
||||
)
|
||||
if word_timestamps:
|
||||
print(f" [{current_label}] Timestamps -> {ts_filename}")
|
||||
current_chunks = []
|
||||
current_timestamps = []
|
||||
alignment_frame_seen = False
|
||||
elif msg_type == "session.done":
|
||||
print(f"\nUtterance {msg['utterance_index']} complete: {msg['total_sentences']} sentence(s) generated")
|
||||
return
|
||||
elif msg_type == "error":
|
||||
print(f" ERROR: {msg['message']}")
|
||||
else:
|
||||
print(f" Unknown message: {msg}")
|
||||
|
||||
|
||||
async def stream_tts(
|
||||
url: str,
|
||||
text: str,
|
||||
texts: list[str],
|
||||
config: dict,
|
||||
output_dir: str,
|
||||
simulate_stt: bool = False,
|
||||
stt_delay: float = 0.1,
|
||||
save_chunks: bool = False,
|
||||
) -> None:
|
||||
"""Connect to the streaming TTS endpoint and process audio responses."""
|
||||
"""Connect to the streaming TTS endpoint and process audio responses.
|
||||
|
||||
Every text is synthesized over the same connection: input.done flushes an
|
||||
utterance without closing, so only the first one pays the handshake.
|
||||
"""
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
response_format = config.get("response_format", "wav")
|
||||
word_timestamps = bool(config.get("word_timestamps", False))
|
||||
|
||||
async with websockets.connect(url) as ws:
|
||||
# 1. Send session config
|
||||
# 1. Send session config; it stays in effect for every utterance below.
|
||||
config_msg = {"type": "session.config", **config}
|
||||
await ws.send(json.dumps(config_msg))
|
||||
print(f"Sent session config: {config}")
|
||||
|
||||
# 2. Send text (either all at once or word-by-word)
|
||||
async def send_text():
|
||||
if simulate_stt:
|
||||
words = text.split(" ")
|
||||
for i, word in enumerate(words):
|
||||
chunk = word + (" " if i < len(words) - 1 else "")
|
||||
await ws.send(
|
||||
json.dumps(
|
||||
{
|
||||
"type": "input.text",
|
||||
"text": chunk,
|
||||
}
|
||||
)
|
||||
)
|
||||
print(f" Sent: {chunk!r}")
|
||||
await asyncio.sleep(stt_delay)
|
||||
else:
|
||||
await ws.send(
|
||||
json.dumps(
|
||||
{
|
||||
"type": "input.text",
|
||||
"text": text,
|
||||
}
|
||||
)
|
||||
)
|
||||
print(f"Sent full text: {text!r}")
|
||||
|
||||
# 3. Signal end of input
|
||||
await ws.send(json.dumps({"type": "input.done"}))
|
||||
print("Sent input.done")
|
||||
|
||||
# Run sender and receiver concurrently
|
||||
sender_task = asyncio.create_task(send_text())
|
||||
|
||||
response_format = config.get("response_format", "wav")
|
||||
word_timestamps = bool(config.get("word_timestamps", False))
|
||||
current_sentence_index = 0
|
||||
current_sentence_text = ""
|
||||
current_chunks: list[bytes] = []
|
||||
current_timestamps: list[dict] = []
|
||||
# Distinguish `null` (aligner failed) from `[]` (silence / no tokens):
|
||||
# the sentence is a failure only if no frame ever carried timestamps.
|
||||
alignment_frame_seen = False
|
||||
|
||||
try:
|
||||
while True:
|
||||
message = await ws.recv()
|
||||
|
||||
if isinstance(message, bytes):
|
||||
current_chunks.append(message)
|
||||
print(f" Received audio chunk for sentence {current_sentence_index}: {len(message)} bytes")
|
||||
else:
|
||||
# JSON frame
|
||||
msg = json.loads(message)
|
||||
msg_type = msg.get("type")
|
||||
|
||||
if msg_type == "audio.start":
|
||||
current_sentence_index = msg["sentence_index"]
|
||||
current_sentence_text = msg.get("sentence_text", "")
|
||||
current_chunks = []
|
||||
current_timestamps = []
|
||||
alignment_frame_seen = False
|
||||
print(f" [sentence {msg['sentence_index']}] Generating: {current_sentence_text!r}")
|
||||
elif msg_type == "audio.chunk":
|
||||
audio = base64.b64decode(msg["audio_b64"])
|
||||
current_chunks.append(audio)
|
||||
chunk_timestamps = msg.get("timestamps")
|
||||
# Sentence-level alignment sends a trailing timestamp-only
|
||||
# frame with empty audio; don't write an empty chunk file.
|
||||
if save_chunks and audio:
|
||||
chunk_base = os.path.join(
|
||||
output_dir,
|
||||
f"sentence_{msg['sentence_index']:03d}_chunk_{msg['chunk_id']:03d}",
|
||||
)
|
||||
chunk_pcm = f"{chunk_base}.pcm"
|
||||
chunk_wav = f"{chunk_base}.wav"
|
||||
chunk_json = f"{chunk_base}_timestamps.json"
|
||||
with open(chunk_pcm, "wb") as f:
|
||||
f.write(audio)
|
||||
save_pcm_wav(chunk_wav, audio, int(msg.get("sample_rate", 24000)))
|
||||
with open(chunk_json, "w", encoding="utf-8") as f:
|
||||
json.dump(
|
||||
{
|
||||
"sentence_index": msg["sentence_index"],
|
||||
"chunk_id": msg["chunk_id"],
|
||||
"chunk_start_ms": msg.get("chunk_start_ms"),
|
||||
"chunk_end_ms": msg.get("chunk_end_ms"),
|
||||
"sentence_text": current_sentence_text,
|
||||
"covered_text": timestamp_words_text(chunk_timestamps),
|
||||
"timestamps": chunk_timestamps,
|
||||
},
|
||||
f,
|
||||
ensure_ascii=False,
|
||||
indent=2,
|
||||
)
|
||||
if chunk_timestamps is None:
|
||||
print(
|
||||
f" [sentence {msg['sentence_index']}] chunk {msg['chunk_id']}: "
|
||||
f"{len(audio)} bytes, timestamps unavailable"
|
||||
)
|
||||
else:
|
||||
alignment_frame_seen = True
|
||||
current_timestamps.extend(chunk_timestamps)
|
||||
print(
|
||||
f" [sentence {msg['sentence_index']}] chunk {msg['chunk_id']}: "
|
||||
f"{len(audio)} bytes, {len(chunk_timestamps)} timestamp(s)"
|
||||
)
|
||||
if save_chunks and audio:
|
||||
print(f" saved chunk -> {chunk_wav}")
|
||||
elif msg_type == "audio.done":
|
||||
filename = save_audio_file(
|
||||
output_dir,
|
||||
msg["sentence_index"],
|
||||
response_format,
|
||||
current_chunks,
|
||||
)
|
||||
if word_timestamps:
|
||||
# `null` on failure, else the accumulated list (`[]` for silence).
|
||||
sentence_timestamps = current_timestamps if alignment_frame_seen else None
|
||||
ts_filename = os.path.join(
|
||||
output_dir,
|
||||
f"sentence_{msg['sentence_index']:03d}_timestamps.json",
|
||||
)
|
||||
with open(ts_filename, "w", encoding="utf-8") as f:
|
||||
json.dump(
|
||||
{
|
||||
"sentence_index": msg["sentence_index"],
|
||||
"sentence_text": current_sentence_text,
|
||||
"covered_text": timestamp_words_text(sentence_timestamps),
|
||||
"timestamps": sentence_timestamps,
|
||||
},
|
||||
f,
|
||||
ensure_ascii=False,
|
||||
indent=2,
|
||||
)
|
||||
print(
|
||||
f" [sentence {msg['sentence_index']}] Done"
|
||||
f" bytes={msg.get('total_bytes', len(b''.join(current_chunks)))}"
|
||||
f" error={msg.get('error', False)}"
|
||||
f" -> {filename}"
|
||||
)
|
||||
if word_timestamps:
|
||||
print(f" [sentence {msg['sentence_index']}] Timestamps -> {ts_filename}")
|
||||
current_chunks = []
|
||||
current_timestamps = []
|
||||
alignment_frame_seen = False
|
||||
elif msg_type == "session.done":
|
||||
print(f"\nSession complete: {msg['total_sentences']} sentence(s) generated")
|
||||
break
|
||||
elif msg_type == "error":
|
||||
print(f" ERROR: {msg['message']}")
|
||||
else:
|
||||
print(f" Unknown message: {msg}")
|
||||
finally:
|
||||
sender_task.cancel()
|
||||
for text in texts:
|
||||
# 2. Send text (either all at once or word-by-word), then flush.
|
||||
sender_task = asyncio.create_task(send_utterance(ws, text, simulate_stt, stt_delay))
|
||||
try:
|
||||
await sender_task
|
||||
except asyncio.CancelledError:
|
||||
pass # Task cancellation is expected during shutdown
|
||||
await receive_utterance(
|
||||
ws,
|
||||
output_dir=output_dir,
|
||||
response_format=response_format,
|
||||
word_timestamps=word_timestamps,
|
||||
save_chunks=save_chunks,
|
||||
)
|
||||
finally:
|
||||
sender_task.cancel()
|
||||
try:
|
||||
await sender_task
|
||||
except asyncio.CancelledError:
|
||||
pass # Task cancellation is expected during shutdown
|
||||
|
||||
# 3. Release the connection once there is nothing left to say.
|
||||
await ws.send(json.dumps({"type": "session.close"}))
|
||||
print("Sent session.close")
|
||||
|
||||
print(f"\nAudio files saved to: {output_dir}/")
|
||||
|
||||
@@ -286,7 +325,8 @@ def main():
|
||||
parser.add_argument(
|
||||
"--text",
|
||||
required=True,
|
||||
help="Text to synthesize",
|
||||
action="append",
|
||||
help="Text to synthesize; repeat to send several utterances over one connection",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-dir",
|
||||
@@ -385,7 +425,7 @@ def main():
|
||||
asyncio.run(
|
||||
stream_tts(
|
||||
url=args.url,
|
||||
text=args.text,
|
||||
texts=args.text,
|
||||
config=config,
|
||||
output_dir=args.output_dir,
|
||||
simulate_stt=args.simulate_stt,
|
||||
|
||||
@@ -248,82 +248,12 @@ done
|
||||
|
||||
## LTX-2
|
||||
|
||||
### Start Server
|
||||
|
||||
#### Basic Start
|
||||
|
||||
```bash
|
||||
vllm serve Lightricks/LTX-2 --omni --port 8098 \
|
||||
--enforce-eager --boundary-ratio 1.0
|
||||
vllm serve Lightricks/LTX-2 --omni --port 8098
|
||||
```
|
||||
|
||||
#### Start with Optimization Presets
|
||||
|
||||
Use the LTX-2 startup script with built-in optimization presets:
|
||||
|
||||
```bash
|
||||
# Baseline (1 GPU, eager)
|
||||
bash run_server_ltx2.sh baseline
|
||||
|
||||
# 4-GPU Ulysses sequence parallelism (lossless)
|
||||
bash run_server_ltx2.sh ulysses4
|
||||
|
||||
# Cache-DiT lossy acceleration (1 GPU, ~1.4× speedup)
|
||||
bash run_server_ltx2.sh cache-dit
|
||||
|
||||
# Best combo: 4-GPU Ulysses SP + Cache-DiT (~2.2× speedup)
|
||||
bash run_server_ltx2.sh best-combo
|
||||
```
|
||||
|
||||
#### Optimization Benchmarks
|
||||
|
||||
Benchmarked on H800, online serving (480×768, 41 frames, 20 steps, `seed=42`).
|
||||
"Inference" is the server-reported inference time; excludes HTTP/poll overhead.
|
||||
|
||||
| Preset | Server Command | Inference (s) | Speedup | Type |
|
||||
|--------|---------------|---------------|---------|------|
|
||||
| `baseline` | `--enforce-eager` | 10.3 | 1.00× | — |
|
||||
| `compile` | *(default, no --enforce-eager)* | ~10.3 (warm) | ~1.00× | Lossless |
|
||||
| `ulysses4` | `--enforce-eager --usp 4` | ~10.3 | ~1.00× | Lossless |
|
||||
| `cache-dit` | `--enforce-eager --cache-backend cache_dit` | 7.4 avg | ~1.4× | Lossy |
|
||||
| `best-combo` | `--enforce-eager --usp 4 --cache-backend cache_dit` | 4.7 avg | **~2.2×** | Lossless + Lossy |
|
||||
|
||||
**Observations**:
|
||||
- **torch.compile**: On H800, warm-request inference time matches the eager baseline (~10.3s).
|
||||
The first request pays ~6s compilation overhead. Benefit depends on model architecture and GPU.
|
||||
- **Ulysses SP (4 GPU)**: No measurable speedup alone for 41-frame generation at this resolution.
|
||||
Communication overhead outweighs gains at this sequence length.
|
||||
- **Cache-DiT**: Inference varies per request (6–10s) due to dynamic caching decisions.
|
||||
Average is ~7.4s (~1.4× speedup) with slight quality tradeoff.
|
||||
- **Best combo**: 4-GPU Ulysses SP + Cache-DiT synergize well — Cache-DiT reduces per-step
|
||||
computation, making the communication overhead of Ulysses SP worthwhile. Average ~4.7s
|
||||
(~2.2× speedup).
|
||||
- **FP8 quantization**: Reduces VRAM but does not speed up LTX-2 on H800 (compute-bound).
|
||||
|
||||
**Deployment Recommendations**:
|
||||
- For **production with quality priority**: use `baseline` with `--enforce-eager`
|
||||
- For **maximum throughput** (4 GPUs, quality tradeoff): use `best-combo` (~2.2× speedup)
|
||||
- For **single-GPU throughput**: use `cache-dit` (~1.4× speedup)
|
||||
- `--enforce-eager` is recommended to avoid torch.compile warmup latency on first request
|
||||
|
||||
### Send Requests (curl)
|
||||
|
||||
```bash
|
||||
# Using the provided script
|
||||
bash run_curl_ltx2.sh
|
||||
|
||||
# Or directly
|
||||
curl -sS -X POST http://localhost:8098/v1/videos \
|
||||
-H "Accept: application/json" \
|
||||
-F "prompt=A serene lakeside sunrise with mist over the water." \
|
||||
-F "width=768" \
|
||||
-F "height=480" \
|
||||
-F "num_frames=41" \
|
||||
-F "fps=24" \
|
||||
-F "num_inference_steps=20" \
|
||||
-F "guidance_scale=3.0" \
|
||||
-F "seed=42"
|
||||
```
|
||||
See the [LTX-2 recipe](../../../recipes/LTX/LTX-2.md) for all checkpoints,
|
||||
pipeline selection, requests, defaults, and advanced options.
|
||||
|
||||
## Helios
|
||||
|
||||
|
||||
@@ -22,7 +22,6 @@ create_response=$(
|
||||
-F "num_frames=41" \
|
||||
-F "fps=24" \
|
||||
-F "num_inference_steps=20" \
|
||||
-F "guidance_scale=3.0" \
|
||||
-F "seed=42"
|
||||
)
|
||||
|
||||
|
||||
@@ -21,8 +21,6 @@ set -euo pipefail
|
||||
|
||||
MODEL="${MODEL:-Lightricks/LTX-2}"
|
||||
PORT="${PORT:-8098}"
|
||||
FLOW_SHIFT="${FLOW_SHIFT:-1.0}"
|
||||
BOUNDARY_RATIO="${BOUNDARY_RATIO:-1.0}"
|
||||
|
||||
PRESET="${1:-baseline}"
|
||||
|
||||
@@ -66,19 +64,13 @@ case "$PRESET" in
|
||||
echo "Environment variables:"
|
||||
echo " MODEL - Model path (default: Lightricks/LTX-2)"
|
||||
echo " PORT - Server port (default: 8098)"
|
||||
echo " FLOW_SHIFT - Scheduler flow shift (default: 1.0)"
|
||||
echo " BOUNDARY_RATIO - Boundary ratio (default: 1.0)"
|
||||
exit 1
|
||||
;;
|
||||
esac
|
||||
|
||||
echo "Model: $MODEL"
|
||||
echo "Port: $PORT"
|
||||
echo "Flow shift: $FLOW_SHIFT"
|
||||
echo "Boundary ratio: $BOUNDARY_RATIO"
|
||||
|
||||
vllm serve "$MODEL" --omni \
|
||||
--port "$PORT" \
|
||||
--flow-shift "$FLOW_SHIFT" \
|
||||
--boundary-ratio "$BOUNDARY_RATIO" \
|
||||
"${EXTRA_ARGS[@]}"
|
||||
@@ -75,6 +75,7 @@ dev = [
|
||||
]
|
||||
|
||||
demo = [
|
||||
"gradio>=6.7.0",
|
||||
"opencv-python>=4.12.0.88", # cli/stream_client frame reading
|
||||
"requests>=2.28.0", # cli/stream_client HTTP
|
||||
"cn2an>=0.5.22", # Chinese number normalization for TTS (e.g. "1次" -> "一次")
|
||||
@@ -95,6 +96,12 @@ quack = [
|
||||
"quack-kernels>=0.3.11",
|
||||
]
|
||||
|
||||
# Optional Blackwell-native FlashAttention-4 backend. Keep this opt-in because
|
||||
# the CuTe/CUTLASS dependency is CUDA 13 and Blackwell specific.
|
||||
fa4 = [
|
||||
"flash-attn-4[cu13]==4.0.0b18",
|
||||
]
|
||||
|
||||
docs = [
|
||||
"mkdocs>=1.5.0",
|
||||
"mkdocs-api-autonav",
|
||||
|
||||
@@ -1,138 +0,0 @@
|
||||
# Realtime Video Prompt Interaction Protocol
|
||||
|
||||
This document specifies the prompt-update subset of the realtime video
|
||||
interaction protocol exposed by:
|
||||
|
||||
```text
|
||||
WebSocket /v1/realtime/video
|
||||
```
|
||||
|
||||
A WebSocket connection contains exactly one video-generation request. Clients
|
||||
therefore do not send a request or session ID with an interaction. The server
|
||||
assigns a `request_id` in `video.start` and includes it in subsequent
|
||||
server-to-client messages.
|
||||
|
||||
## Interaction Request
|
||||
|
||||
After `video.start`, a client may submit a prompt update:
|
||||
|
||||
```json
|
||||
{
|
||||
"type": "session.interaction",
|
||||
"interaction": {
|
||||
"event_id": "ui-update-42",
|
||||
"event": {
|
||||
"prompt": "A snowy street illuminated by warm lanterns"
|
||||
},
|
||||
"transition_chunks": 3
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
`interaction` has the following fields:
|
||||
|
||||
- `event_id` (optional string): session-unique idempotency key. The
|
||||
server preserves a client-provided value or generates a UUID when omitted.
|
||||
- `event.prompt` (required string): non-empty target prompt.
|
||||
over which the prompt reaches its target.
|
||||
- `transition_chunks` (optional non-negative integer): model-generation chunks
|
||||
over which the prompt reaches its target. If omitted, the model's default transition is used.
|
||||
|
||||
Note: Although `["interaction"]["event"]` only contains `prompt`, still consider using `OmniTextPrompt` for this part (with total=false) so that other fields will be supported in the future. Currently, `vllm_omni/diffusion/worker/diffusion_model_runner.py` already takes care of enforcing subset format, routing to meaningful input handing routines, and raising `NotImplementedError`. Only need to adapt this design to the input data structure above.
|
||||
|
||||
### Transition Semantics
|
||||
|
||||
Prompt transitions are target updates and begin at the next eligible model
|
||||
generation boundary after the interaction is processed.
|
||||
|
||||
- `transition_chunks = 0` is a hard update: the first affected generation
|
||||
chunk uses the target prompt.
|
||||
- `transition_chunks = N`, where `N > 0`, reaches the target in the `N`th
|
||||
affected generation chunk. For example, a three-chunk linear transition uses
|
||||
progress values `1/3`, `2/3`, and `1`.
|
||||
|
||||
### Idempotency and Replacement
|
||||
|
||||
Within a session:
|
||||
|
||||
- repeating an `event_id` with the same payload is a retry and must not
|
||||
apply the update twice;
|
||||
- repeating an `event_id` with a different payload is rejected with
|
||||
`event_id_conflict`;
|
||||
- a newer prompt target supersedes any older pending prompt target;
|
||||
- when a new target interrupts an active transition, its transition starts
|
||||
from the prompt state resolved at its effective boundary.
|
||||
|
||||
## Queue Acknowledgement
|
||||
|
||||
The server quickly acknowledges admission to the interaction queue:
|
||||
|
||||
```json
|
||||
{
|
||||
"type": "session.interaction.queued",
|
||||
"request_id": "video_stream-...",
|
||||
"event_id": "ui-update-42"
|
||||
}
|
||||
```
|
||||
|
||||
`queued` confirms syntax validation, idempotency handling, and queue admission.
|
||||
It does not claim that the model has scheduled or applied the update. A later
|
||||
failure is reported as a correlated error:
|
||||
|
||||
```json
|
||||
{
|
||||
"type": "error",
|
||||
"request_id": "video_stream-...",
|
||||
"event_id": "ui-update-42",
|
||||
"code": "interaction_failed",
|
||||
"message": "The active model does not support prompt updates"
|
||||
}
|
||||
```
|
||||
|
||||
## Video Chunk Metadata
|
||||
|
||||
Each emitted binary payload is preceded by a JSON metadata message. The
|
||||
metadata message and its binary payload form one ordered pair and must not be
|
||||
interleaved with another outbound message.
|
||||
|
||||
```json
|
||||
{
|
||||
"type": "video.chunk_metadata",
|
||||
"request_id": "video_stream-...",
|
||||
"kind": "media",
|
||||
"transport_chunk_index": 7,
|
||||
"generation_chunk_index": 6,
|
||||
"num_frames": 9,
|
||||
"byte_length": 184320,
|
||||
"started_event_ids": ["ui-update-42"],
|
||||
"active_event_ids": ["ui-update-42"],
|
||||
"completed_event_ids": []
|
||||
}
|
||||
```
|
||||
|
||||
- `transport_chunk_index` identifies the binary WebSocket payload.
|
||||
- `generation_chunk_index` identifies the model-generation chunk represented
|
||||
by the payload.
|
||||
- `started_event_ids` lists interactions that first affect this
|
||||
generation chunk.
|
||||
- `active_event_ids` lists interactions contributing to this generation
|
||||
chunk, including an ongoing transition.
|
||||
- `completed_event_ids` lists interactions that reach their target in
|
||||
this generation chunk.
|
||||
|
||||
Encoder trailer bytes are emitted as:
|
||||
|
||||
```json
|
||||
{
|
||||
"type": "video.chunk_metadata",
|
||||
"request_id": "video_stream-...",
|
||||
"kind": "trailer",
|
||||
"transport_chunk_index": 8,
|
||||
"generation_chunk_index": null,
|
||||
"num_frames": 0,
|
||||
"byte_length": 824,
|
||||
"started_event_ids": [],
|
||||
"active_event_ids": [],
|
||||
"completed_event_ids": []
|
||||
}
|
||||
```
|
||||
@@ -168,7 +168,7 @@ the model + orchestrator + WebUI together — point `WEBUI_DIR` at your clone.
|
||||
python examples/online_serving/joyvl_interaction/cli/run_cli_demo.py \
|
||||
path/to/video.mp4 --query "Alert me if a fire breaks out"
|
||||
|
||||
pytest tests/fullduplex # framework + JoyVL unit tests
|
||||
pytest tests/e2e/features/fullduplex # framework + JoyVL unit tests
|
||||
```
|
||||
|
||||
## Testing with an RTSP stream (optional)
|
||||
|
||||
@@ -1,154 +0,0 @@
|
||||
# LTX-2.3
|
||||
|
||||
> 22B parameter text-to-video + audio generation model served via vLLM-Omni
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: Lightricks
|
||||
- Model: `diffusers/LTX-2.3-Diffusers`
|
||||
- Task: Text-to-video and image-to-video with synchronized audio generation
|
||||
- Mode: Online serving (pure diffusion)
|
||||
- Maintainer: @oglok
|
||||
|
||||
## When to use this recipe
|
||||
|
||||
Use this recipe when you want to serve LTX-2.3 for text-to-video or
|
||||
image-to-video generation with audio. The model generates videos up to 20+
|
||||
seconds at 768x512 resolution with 48kHz audio from text prompts, and can
|
||||
condition video generation on an initial image. Start with a 96GB-class GPU for
|
||||
validation because the model combines a 22B parameter transformer with text
|
||||
encoder, VAE, and vocoder components.
|
||||
|
||||
## References
|
||||
|
||||
- Upstream raw checkpoints: <https://huggingface.co/Lightricks/LTX-2.3>
|
||||
- Diffusers-format checkpoint: <https://huggingface.co/diffusers/LTX-2.3-Diffusers>
|
||||
- Requires `diffusers >= 0.38.0` (install from git: `pip install git+https://github.com/huggingface/diffusers.git`)
|
||||
|
||||
## Serving
|
||||
|
||||
### Command
|
||||
|
||||
```bash
|
||||
vllm serve diffusers/LTX-2.3-Diffusers \
|
||||
--omni \
|
||||
--model-class-name LTX23Pipeline \
|
||||
--stage-init-timeout 600
|
||||
```
|
||||
|
||||
For image-to-video:
|
||||
|
||||
```bash
|
||||
vllm serve diffusers/LTX-2.3-Diffusers \
|
||||
--omni \
|
||||
--model-class-name LTX23ImageToVideoPipeline \
|
||||
--stage-init-timeout 600
|
||||
```
|
||||
|
||||
### T2V Verification
|
||||
|
||||
```bash
|
||||
# Health check
|
||||
curl http://localhost:8000/health
|
||||
|
||||
# Generate a 3-second video (81 frames at 24fps)
|
||||
curl -X POST http://localhost:8000/v1/videos \
|
||||
-F "prompt=A majestic bald eagle soaring over a misty mountain valley at dawn, golden sunlight breaking through clouds" \
|
||||
-F "negative_prompt=blurry, low quality, distorted, watermark" \
|
||||
-F "model=diffusers/LTX-2.3-Diffusers" \
|
||||
-F "num_frames=81" \
|
||||
-F "fps=24" \
|
||||
-F "size=768x512" \
|
||||
-F "num_inference_steps=30" \
|
||||
-F "guidance_scale=4.0" \
|
||||
-F "seed=42"
|
||||
|
||||
# Generate a 10-second video (241 frames)
|
||||
curl -X POST http://localhost:8000/v1/videos \
|
||||
-F "prompt=A cozy Japanese ramen shop at night in the rain, steam rising from bowls, neon signs reflecting on wet cobblestone streets" \
|
||||
-F "model=diffusers/LTX-2.3-Diffusers" \
|
||||
-F "num_frames=241" \
|
||||
-F "fps=24" \
|
||||
-F "size=768x512" \
|
||||
-F "num_inference_steps=30" \
|
||||
-F "guidance_scale=4.0"
|
||||
|
||||
# Generate a 20-second video (481 frames)
|
||||
curl -X POST http://localhost:8000/v1/videos \
|
||||
-F "prompt=An underwater coral reef teeming with tropical fish, sea turtles gliding gracefully, National Geographic documentary style" \
|
||||
-F "model=diffusers/LTX-2.3-Diffusers" \
|
||||
-F "num_frames=481" \
|
||||
-F "fps=24" \
|
||||
-F "size=768x512" \
|
||||
-F "num_inference_steps=30" \
|
||||
-F "guidance_scale=4.0"
|
||||
```
|
||||
|
||||
### I2V Verification
|
||||
|
||||
Start the server with `--model-class-name LTX23ImageToVideoPipeline`, then pass
|
||||
an image reference. The I2V pipeline requires `input_reference` or
|
||||
`image_reference` because the first frame is the conditioning input.
|
||||
|
||||
```bash
|
||||
curl http://localhost:8000/health
|
||||
|
||||
curl -X POST http://localhost:8000/v1/videos/sync \
|
||||
-F "prompt=A plush toy astronaut gently waving while the camera slowly pushes in." \
|
||||
-F "input_reference=@/absolute/path/to/reference.png" \
|
||||
-F "model=diffusers/LTX-2.3-Diffusers" \
|
||||
-F "num_frames=81" \
|
||||
-F "fps=24" \
|
||||
-F "size=768x512" \
|
||||
-F "num_inference_steps=30" \
|
||||
-F "guidance_scale=4.0" \
|
||||
-F "seed=42" \
|
||||
-o ltx23_i2v.mp4
|
||||
```
|
||||
|
||||
### Notes
|
||||
|
||||
- Memory usage: run a latest-head serving smoke or DFX benchmark on the target
|
||||
GPU before publishing load/peak VRAM numbers.
|
||||
- Checkpoint format: use a Diffusers-format checkpoint such as
|
||||
`diffusers/LTX-2.3-Diffusers`; the upstream `Lightricks/LTX-2.3` repository ships
|
||||
raw safetensors and does not contain the subfolder configs required by this
|
||||
pipeline loader.
|
||||
- Key flags:
|
||||
- `--stage-init-timeout 600`: Gives the initial `torch.compile` warmup enough time on large checkpoints
|
||||
- `--model-class-name LTX23Pipeline`: Selects the LTX-2.3 text-to-video pipeline (not LTX-2)
|
||||
- `--model-class-name LTX23ImageToVideoPipeline`: Selects the LTX-2.3 image-to-video pipeline
|
||||
- `input_reference=@...`: Uploads the I2V reference image directly
|
||||
- Audio: 48kHz AAC via BWE vocoder, automatically synced with video
|
||||
- CPU offloading: Text encoder (Gemma-3-12B), connectors, VAE, audio VAE, and vocoder stay on CPU and are moved to GPU only when needed
|
||||
- Supported resolutions: 768x512, 512x384 (must be divisible by 32)
|
||||
- Frame rate: 24 fps
|
||||
- Duration: Controlled by `num_frames` (frames = duration_seconds * 24 + 1)
|
||||
- Known limitations:
|
||||
- LTX-2.3 I2V currently supports first-frame image conditioning, matching the LTX image-conditioning contract.
|
||||
- Requires `diffusers >= 0.38.0` (not yet on PyPI, install from git)
|
||||
|
||||
## Hardware Support
|
||||
|
||||
## GPU
|
||||
|
||||
### 1x NVIDIA RTX PRO 6000 Blackwell (96GB)
|
||||
|
||||
#### Environment
|
||||
|
||||
- OS: Ubuntu 22.04
|
||||
- Python: 3.10+
|
||||
- Driver / runtime: CUDA 13.0, Driver 580.126.09
|
||||
- vLLM version: Match the PR or release checkout being validated
|
||||
- vLLM-Omni version or commit: Use the commit you are deploying from
|
||||
|
||||
### Benchmarking
|
||||
|
||||
Do not copy latency or VRAM numbers between commits or machines. Record them
|
||||
only after running a latest-head sweep on the target hardware.
|
||||
|
||||
For formal PR benchmarking, reuse
|
||||
`tests/dfx/perf/tests/test_ltx2_vllm_omni.json` with
|
||||
`tests/dfx/perf/scripts/run_diffusion_benchmark.py`. Each case in that file includes a JSON **`mark`** (`hardware_marks` + `full_model` + `diffusion`) for local `-m` filtering. That config captures the
|
||||
single-device eager baseline and CFG-parallel=2 case for a small 384x512,
|
||||
25-frame, 20-step workload.
|
||||
+329
-113
@@ -1,121 +1,337 @@
|
||||
# LTX-2
|
||||
# LTX-2 Family
|
||||
|
||||
> Text-to-video and image-to-video serving
|
||||
> LTX-2 and LTX-2.3 text-to-video and image-to-video generation with synchronized audio
|
||||
|
||||
## Summary
|
||||
## Pipelines
|
||||
|
||||
- Vendor: Lightricks
|
||||
- Model: `Lightricks/LTX-2`
|
||||
- Task: Text-to-video generation and image-to-video generation
|
||||
- Mode: Online serving with the OpenAI-compatible video API
|
||||
- Maintainer: Community
|
||||
| Model | Checkpoint | Task | `--model-class-name` | Batching |
|
||||
|---|---|---|---|---|
|
||||
| LTX-2 | `Lightricks/LTX-2` | One-stage T2V/I2V | `LTX2Pipeline` | Yes |
|
||||
| LTX-2 distilled | `rootonchair/LTX-2-19b-distilled` | Two-stage T2V/I2V | `LTX2DistilledPipeline` | No |
|
||||
| LTX-2.3 | `diffusers/LTX-2.3-Diffusers` | One-stage T2V/I2V | `LTX2Pipeline` | Yes |
|
||||
|
||||
## When to use this recipe
|
||||
`LTX2Pipeline` is the unified one-stage entry. Checkpoint metadata selects the
|
||||
LTX-2 or LTX-2.3 profile; omitting an image selects T2V, while one initial
|
||||
image selects I2V. Both one-stage repositories declare this class, so
|
||||
`--model-class-name` is optional. LTX-2.3 requires the Diffusers checkpoint;
|
||||
the raw `Lightricks/LTX-2.3` safetensors repository is not directly loadable.
|
||||
|
||||
Use this recipe when you want to deploy the base `Lightricks/LTX-2` model
|
||||
for both text-to-video and image-to-video generation. The same model ID is used for
|
||||
both tasks; switch the pipeline class with `--model-class-name` depending on
|
||||
whether you want T2V or I2V.
|
||||
`LTX2DistilledPipeline` is the unified distilled entry: omit `image` for T2V
|
||||
or provide one initial image for I2V. As in the official LTX pipelines,
|
||||
distilled mode is selected explicitly rather than inferred from the checkpoint
|
||||
name or metadata; always pass `--model-class-name LTX2DistilledPipeline`.
|
||||
|
||||
## API Migration
|
||||
|
||||
Only `req` may be passed positionally to `LTX2Pipeline`; every optional
|
||||
`forward` argument is keyword-only:
|
||||
|
||||
```python
|
||||
# No longer supported
|
||||
pipe(req, prompt)
|
||||
pipe(req, image, prompt)
|
||||
|
||||
# Supported
|
||||
pipe(req, prompt=prompt)
|
||||
pipe(req, image=image, prompt=prompt)
|
||||
```
|
||||
|
||||
The consolidation also removes these registry names without aliases:
|
||||
|
||||
| Removed name | Replacement |
|
||||
|---|---|
|
||||
| `LTX23Pipeline` | `LTX2Pipeline`; checkpoint metadata selects LTX-2.3 |
|
||||
| `LTX2ImageToVideoPipeline` | `LTX2Pipeline` with `image=` |
|
||||
| `LTX23ImageToVideoPipeline` | `LTX2Pipeline` with `image=`; checkpoint metadata selects LTX-2.3 |
|
||||
| `LTX2TwoStagesPipeline` | `LTX2DistilledPipeline` |
|
||||
| `LTX2ImageToVideoTwoStagesPipeline` | `LTX2DistilledPipeline` with `image=` |
|
||||
|
||||
Passing any second positional argument now raises `TypeError`. These changes
|
||||
affect direct Python callers and explicit `--model-class-name` overrides;
|
||||
offline and serving entrypoints already use named fields and are unaffected.
|
||||
|
||||
## One-Stage Defaults
|
||||
|
||||
| Parameter | LTX-2 | LTX-2.3 |
|
||||
|---|---:|---:|
|
||||
| Width × height | 768 × 512 | 768 × 512 |
|
||||
| Frames / frame rate | 121 / 24 | 121 / 24 |
|
||||
| Denoise steps | 40 | 30 |
|
||||
| Video/audio CFG | 3.0 / 7.0 | 3.0 / 7.0 |
|
||||
| Video/audio STG | 1.0 / 1.0 | 1.0 / 1.0 |
|
||||
| Video/audio modality guidance | 3.0 / 3.0 | 3.0 / 3.0 |
|
||||
| Video/audio rescale | 0.7 / 0.7 | 0.7 / 0.7 |
|
||||
| Video/audio STG blocks | `[29]` / `[29]` | `[28]` / `[28]` |
|
||||
|
||||
The model recipe also supplies the default negative prompt. Top-level
|
||||
`guidance_scale` overrides both video and audio CFG values. The online video
|
||||
API defaults `num_frames` to `1`, so set it explicitly; the offline LTX scripts
|
||||
default to `121`.
|
||||
|
||||
## Distilled Two-Stage Defaults
|
||||
|
||||
| Parameter | Value |
|
||||
|---|---:|
|
||||
| Final width × height | 1536 × 1024 |
|
||||
| Stage 1 width × height | 768 × 512 |
|
||||
| Frames / frame rate | 121 / 24 |
|
||||
| Stage 1 / Stage 2 denoise steps | 8 / 3 |
|
||||
| Guidance | Fixed positive-only |
|
||||
|
||||
The dimensions passed to the API are the final output dimensions and must be
|
||||
divisible by 64. Both one- and two-stage requests require `num_frames = 8k+1`.
|
||||
The distilled recipe fixes both sigma schedules; `num_inference_steps` must be
|
||||
`8` when supplied.
|
||||
|
||||
## Serving
|
||||
|
||||
Start either checkpoint:
|
||||
|
||||
```bash
|
||||
vllm serve Lightricks/LTX-2 --omni --stage-init-timeout 600
|
||||
```
|
||||
|
||||
```bash
|
||||
vllm serve diffusers/LTX-2.3-Diffusers --omni --stage-init-timeout 600
|
||||
```
|
||||
|
||||
Start the distilled checkpoint with its two-stage entry:
|
||||
|
||||
```bash
|
||||
vllm serve rootonchair/LTX-2-19b-distilled --omni \
|
||||
--model-class-name LTX2DistilledPipeline --stage-init-timeout 600
|
||||
```
|
||||
|
||||
The same server handles T2V and I2V. A T2V request using the selected recipe's
|
||||
default guidance is:
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8000/v1/videos/sync \
|
||||
-F "prompt=A cinematic close-up of ocean waves at golden hour." \
|
||||
-F "negative_prompt=worst quality, inconsistent motion, blurry, jittery, distorted" \
|
||||
-F "size=768x512" \
|
||||
-F "num_frames=121" \
|
||||
-F "fps=24" \
|
||||
-F "seed=42" \
|
||||
-o ltx_t2v.mp4
|
||||
```
|
||||
|
||||
For I2V, provide exactly one initial image:
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8000/v1/videos/sync \
|
||||
-F "prompt=A plush toy astronaut gently waving while the camera slowly pushes in." \
|
||||
-F "negative_prompt=worst quality, inconsistent motion, blurry, jittery, distorted" \
|
||||
-F "input_reference=@/absolute/path/to/reference.png" \
|
||||
-F "size=768x512" \
|
||||
-F "num_frames=121" \
|
||||
-F "fps=24" \
|
||||
-F "seed=42" \
|
||||
-o ltx_i2v.mp4
|
||||
```
|
||||
|
||||
Use `image_reference` for a URL or JSON-safe image reference. Do not provide it
|
||||
together with `input_reference`.
|
||||
|
||||
## Guidance
|
||||
|
||||
LTX supports independent video/audio CFG, spatio-temporal guidance (STG),
|
||||
cross-modality guidance, and rescaling:
|
||||
|
||||
| Parameter | Default | Effect | Alias |
|
||||
|---|---:|---|---|
|
||||
| `video_cfg_scale` | 3.0 | Video text CFG; `1.0` disables it | `video_cfg_guidance_scale` |
|
||||
| `audio_cfg_scale` | 7.0 | Audio text CFG; `1.0` disables it | `audio_cfg_guidance_scale` |
|
||||
| `video_stg_scale` | 1.0 | Video STG; `0.0` disables it | `video_stg_guidance_scale` |
|
||||
| `audio_stg_scale` | 1.0 | Audio STG; `0.0` disables it | `audio_stg_guidance_scale` |
|
||||
| `video_modality_scale` | 3.0 | Audio-to-video guidance; `1.0` disables it | `a2v_guidance_scale` |
|
||||
| `audio_modality_scale` | 3.0 | Video-to-audio guidance; `1.0` disables it | `v2a_guidance_scale` |
|
||||
| `video_rescale_scale` | 0.7 | Video guidance rescale; `0.0` disables it | — |
|
||||
| `audio_rescale_scale` | 0.7 | Audio guidance rescale; `0.0` disables it | — |
|
||||
| `video_stg_blocks` | `[29]` / `[28]` | Perturbed video transformer blocks | — |
|
||||
| `audio_stg_blocks` | `[29]` / `[28]` | Perturbed audio transformer blocks | — |
|
||||
|
||||
The STG block defaults are shown as LTX-2 / LTX-2.3. A canonical name wins
|
||||
over its alias. Explicit top-level `guidance_scale` wins over both CFG fields;
|
||||
omit it for independent video/audio CFG.
|
||||
|
||||
Pass these fields in online `extra_params` or offline `--extra-body`:
|
||||
|
||||
```bash
|
||||
# Add to the curl request above
|
||||
-F 'extra_params={"video_cfg_scale":3.0,"audio_cfg_scale":7.0}'
|
||||
|
||||
# Add to text_to_video.py or image_to_video.py
|
||||
--extra-body '{"video_cfg_scale":3.0,"audio_cfg_scale":7.0}'
|
||||
```
|
||||
|
||||
### Guidance Parallelism
|
||||
|
||||
For LTX, `cfg_parallel_size` is the number of ranks used to execute the
|
||||
complete guidance plan in parallel. Despite the legacy `cfg` name, it covers
|
||||
text CFG, STG, and cross-modality guidance passes; guidance rescaling is
|
||||
applied after all pass predictions are gathered.
|
||||
|
||||
The default one-stage recipe has four Transformer passes per denoise step:
|
||||
`cond`, `uncond`, `ptb` (STG), and `mod` (cross-modality). The useful balanced
|
||||
configurations are therefore:
|
||||
|
||||
| `--cfg-parallel-size` | Passes per rank | Guidance-slot utilization | Notes |
|
||||
|---:|---:|---:|---|
|
||||
| `1` | 4 | 100% | Single-rank fused guidance batch |
|
||||
| `2` | 2 | 100% | Recommended two-rank configuration |
|
||||
| `4` | 1 | 100% | One guidance pass per rank |
|
||||
|
||||
Other positive sizes are accepted. When the pass count is not divisible by
|
||||
`cfg_parallel_size`, ranks are padded to an equal number of execution slots
|
||||
and LTX emits a warning with the expected utilization. For example, four
|
||||
passes on three ranks use six slots and have 66.7% guidance-slot utilization.
|
||||
|
||||
Start an LTX-2.3 server with two-way guidance parallelism:
|
||||
|
||||
```bash
|
||||
vllm serve diffusers/LTX-2.3-Diffusers --omni \
|
||||
--cfg-parallel-size 2 --stage-init-timeout 600
|
||||
```
|
||||
|
||||
The total device count is the product of `cfg_parallel_size` and the other
|
||||
configured parallel dimensions. Distilled pipelines use positive-only
|
||||
guidance, so `cfg_parallel_size > 1` does not provide guidance speedup for
|
||||
them.
|
||||
|
||||
### Python API
|
||||
|
||||
Normal per-request values belong in `OmniDiffusionSamplingParams`. Pass the
|
||||
resulting `DiffusionRequestBatch` as the only positional argument:
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
||||
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
params = OmniDiffusionSamplingParams(
|
||||
width=768,
|
||||
height=512,
|
||||
num_frames=121,
|
||||
frame_rate=24.0,
|
||||
num_inference_steps=30,
|
||||
seed=42,
|
||||
extra_args={"video_cfg_scale": 3.0, "audio_cfg_scale": 7.0},
|
||||
)
|
||||
req = DiffusionRequestBatch([
|
||||
OmniDiffusionRequest(
|
||||
prompt={"prompt": "Cherry blossoms moving in a light breeze"},
|
||||
sampling_params=params,
|
||||
request_id="ltx-example",
|
||||
)
|
||||
])
|
||||
|
||||
t2v_output = pipe(req)
|
||||
i2v_output = pipe(req, image=image)
|
||||
```
|
||||
|
||||
Direct keyword arguments are low-level fallbacks, for example
|
||||
`pipe(req, prompt=prompt, width=768)`. Request fields take precedence unless
|
||||
noted below.
|
||||
|
||||
### Complete `forward` Surface
|
||||
|
||||
| Argument | Type/default | Meaning and constraints |
|
||||
|---|---|---|
|
||||
| `req` | `DiffusionRequestBatch`, required | Only positional argument; contains prompts and per-request sampling parameters. |
|
||||
| `image` | image or batch, `None` | Direct value wins over request images; no image selects T2V. I2V accepts one image per prompt, and a batch cannot mix T2V/I2V. |
|
||||
| `prompt` | string or list, `None` | Positive-text fallback; request prompts win. Mutually exclusive with `prompt_embeds`. |
|
||||
| `negative_prompt` | string or list, `None` | Fallback after request values and before the recipe default; mutually exclusive with negative embeddings. |
|
||||
| `height` | `int`, `None` | Request → direct value → recipe default; divisible by 32 one-stage or 64 distilled. |
|
||||
| `width` | `int`, `None` | Same precedence and alignment as `height`. |
|
||||
| `num_frames` | `int`, `None` | Request → direct value → recipe default; must be `8k+1` and also determines audio duration with `frame_rate`. |
|
||||
| `frame_rate` | `float`, `None` | Request `frame_rate` → request `fps` → direct value → recipe default. |
|
||||
| `num_inference_steps` | `int`, `None` | Request → direct value → recipe default; minimum 2 one-stage, fixed at 8 for distilled Stage 1. Custom one-stage `sigmas` determine actual steps. |
|
||||
| `sigmas` | list of float, `None` | One-stage only. Request values win; every request in a fused batch must use the same schedule. |
|
||||
| `timesteps` | list of int, `None` | Compatibility slot; LTX accepts only `None`. Use `sigmas`. |
|
||||
| `guidance_scale` | `float`, `None` | One-stage common video/audio CFG fallback; an explicit request value wins. Distilled two-stage uses fixed positive-only guidance. |
|
||||
| `guidance_rescale` | `float`, `None` | Accepts only `None` or `0.0`; use the modality rescale fields. |
|
||||
| `noise_scale` | `float`, `0.0` | Compatibility slot; LTX accepts only `0.0`. |
|
||||
| `num_videos_per_prompt` | `int`, `1` | Output-count fallback; positive request `num_outputs_per_prompt` wins. |
|
||||
| `generator` | generator or list, `None` | Explicit RNG; otherwise request generators/seeds are collated. Lists must match the effective output batch. |
|
||||
| `latents` | tensor, `None` | Request tensors win. One-stage accepts packed `[B, S, C]` and validated unpacked video layouts. |
|
||||
| `audio_latents` | tensor, `None` | One-stage initial audio latents; request tensors win and are collated. |
|
||||
| `prompt_embeds` | tensor, `None` | Precomputed positive conditioning; requires `prompt_attention_mask` and cannot accompany `prompt`. |
|
||||
| `negative_prompt_embeds` | tensor, `None` | One-stage precomputed negative conditioning; requires its mask and cannot accompany `negative_prompt`. |
|
||||
| `prompt_attention_mask` | tensor, `None` | Mask for positive embeddings; request `prompt_attention_mask` or `attention_mask` wins. |
|
||||
| `negative_prompt_attention_mask` | tensor, `None` | Mask for negative embeddings; request negative mask fields win. |
|
||||
| `decode_timestep` | float or list, `0.0` | Video-VAE decode timestep; request value wins. Lists may match 1, prompt batch, or output batch. |
|
||||
| `decode_noise_scale` | float or list, `None` | Same list rules; request value wins. Defaults to `decode_timestep`. |
|
||||
| `output_type` | string, `"np"` | Request value wins. `"np"` decodes output; `"latent"` skips VAE/vocoder decode. |
|
||||
| `return_dict` | `bool`, `True` | Compatibility slot; only `True` is accepted. |
|
||||
| `attention_kwargs` | dict, `None` | Per-call values are unsupported; configure attention at engine startup. |
|
||||
| `max_sequence_length` | `int`, `None` | Request → direct value → tokenizer limit. |
|
||||
|
||||
Request-object naming differs only in a few places: `num_videos_per_prompt`
|
||||
maps to `num_outputs_per_prompt`; images and prompt text/embeddings live in the
|
||||
request prompt payload; LTX guidance fields live in sampling `extra_args`.
|
||||
|
||||
### Recipe-Specific Request Capabilities
|
||||
|
||||
| Override | One-stage | Distilled two-stage |
|
||||
|---|---|---|
|
||||
| Guidance | Supported | Fixed positive-only |
|
||||
| Negative prompt/embeddings | Supported | Rejected |
|
||||
| `num_inference_steps` | Supported | Fixed at 8 for Stage 1; Stage 2 uses 3 |
|
||||
| Custom `sigmas` | Supported | Rejected; both phases use fixed schedules |
|
||||
| Video/audio latents | Supported | Rejected |
|
||||
|
||||
These capability checks apply equally to direct `forward` keywords and
|
||||
values in `OmniDiffusionSamplingParams`; unsupported values fail instead of
|
||||
being silently ignored.
|
||||
|
||||
### Custom Sigma Schedules
|
||||
|
||||
One-stage Python requests may set final scheduler boundaries directly:
|
||||
|
||||
```python
|
||||
params = OmniDiffusionSamplingParams(sigmas=[1.0, 0.75, 0.5, 0.25])
|
||||
```
|
||||
|
||||
Each nonterminal value produces one denoise step, and a terminal `0.0` is
|
||||
appended when omitted. This schedule overrides `num_inference_steps`. All
|
||||
requests in a fused batch must use the same list. The video form API and
|
||||
bundled offline CLI do not currently expose `sigmas`.
|
||||
|
||||
### Constraints
|
||||
|
||||
- LTX rejects non-default `timesteps`, `flow_shift`, `guidance_rescale`,
|
||||
`noise_scale`, `attention_kwargs`, and `return_dict=False`. Use final
|
||||
`sigmas`, modality rescale fields, startup attention configuration, and the
|
||||
standard `DiffusionOutput` respectively.
|
||||
- Fused one-stage requests must resolve to identical LTX guidance and sigma
|
||||
schedules. Keep concurrent requests guidance-homogeneous or use
|
||||
`--max-num-seqs 1`.
|
||||
- `--cfg-parallel-size` shards the complete LTX guidance plan, including STG,
|
||||
modality guidance, and rescale-compatible prediction gathering.
|
||||
|
||||
## Operational Notes
|
||||
|
||||
- LTX-2 one-stage previously loaded and peaked at about 73.5 GiB on one H200
|
||||
141GB; remeasure on the deployed commit and hardware.
|
||||
- LTX-2.3 includes a 22B transformer, Gemma encoder, two VAEs, and a vocoder.
|
||||
Start on a 96GB-class GPU or use CPU/layerwise offload on smaller devices.
|
||||
- The distilled upsampler uses the configured pipeline dtype and participates
|
||||
in component discovery/device placement; it remains resident during denoise
|
||||
offload because it is used only at the phase boundary.
|
||||
- The output audio sample rate comes from the loaded components and is not a
|
||||
request parameter.
|
||||
- For benchmarks, use `tests/dfx/perf/tests/test_ltx2_vllm_omni.json` with
|
||||
`tests/dfx/perf/scripts/run_diffusion_benchmark.py`.
|
||||
- Current two-stage support is limited to the distilled checkpoint; official
|
||||
LTX two-stage/HQ execution remains out of scope.
|
||||
|
||||
## References
|
||||
|
||||
- Model: <https://huggingface.co/Lightricks/LTX-2>
|
||||
- User guide:
|
||||
[`docs/user_guide/examples/online_serving/text_to_video.md`](../../docs/user_guide/examples/online_serving/text_to_video.md)
|
||||
- Example guide:
|
||||
[`examples/online_serving/text_to_video/README.md`](../../examples/online_serving/text_to_video/README.md)
|
||||
[`examples/online_serving/image_to_video/README.md`](../../examples/online_serving/image_to_video/README.md)
|
||||
|
||||
## Hardware Support
|
||||
|
||||
## GPU
|
||||
|
||||
### 1x H200 141GB (Text-to-Video)
|
||||
|
||||
#### Environment
|
||||
|
||||
- OS: Ubuntu 22.04.5 LTS
|
||||
- Python: 3.12+
|
||||
- Driver / runtime: NVIDIA CUDA environment with H200 141GB GPU
|
||||
- vLLM version: Match the repository requirements from your current checkout
|
||||
- vLLM-Omni version or commit: Use the commit you are deploying from
|
||||
|
||||
#### Command
|
||||
|
||||
```bash
|
||||
vllm serve Lightricks/LTX-2 \
|
||||
--omni \
|
||||
--model-class-name LTX2Pipeline
|
||||
```
|
||||
|
||||
#### Verification
|
||||
|
||||
```bash
|
||||
curl http://localhost:8000/health
|
||||
|
||||
curl -X POST http://localhost:8000/v1/videos \
|
||||
-F "prompt=A cinematic close-up of ocean waves at golden hour." \
|
||||
-F "negative_prompt=worst quality, inconsistent motion, blurry, jittery, distorted" \
|
||||
-F "width=512" \
|
||||
-F "height=768" \
|
||||
-F "num_frames=121" \
|
||||
-F "fps=24" \
|
||||
-F "num_inference_steps=40" \
|
||||
-F "guidance_scale=4.0" \
|
||||
-o ltx2_t2v_h200.mp4
|
||||
```
|
||||
|
||||
#### Notes
|
||||
|
||||
- Memory usage: Model loads at ~73.5GB, peaks at ~73.5 GiB during inference.
|
||||
- Key flags:
|
||||
- `--model-class-name LTX2Pipeline` selects the text-to-video pipeline explicitly.
|
||||
- Known limitations:
|
||||
- This recipe is tested on 1x H200 validation; add measured latency and VRAM after running it on your target machine.
|
||||
|
||||
### 1x H200 141GB (Image-to-Video)
|
||||
|
||||
#### Environment
|
||||
|
||||
- OS: Ubuntu 22.04.5 LTS
|
||||
- Python: 3.12+
|
||||
- Driver / runtime: NVIDIA CUDA environment with H200 141GB GPU
|
||||
- vLLM version: Match the repository requirements from your current checkout
|
||||
- vLLM-Omni version or commit: Use the commit you are deploying from
|
||||
|
||||
#### Command
|
||||
|
||||
```bash
|
||||
vllm serve Lightricks/LTX-2 \
|
||||
--omni \
|
||||
--model-class-name LTX2ImageToVideoPipeline
|
||||
```
|
||||
|
||||
#### Verification
|
||||
|
||||
```bash
|
||||
curl http://localhost:8000/health
|
||||
|
||||
curl -X POST http://localhost:8000/v1/videos/sync \
|
||||
-F "prompt=A plush toy astronaut gently waving while the camera slowly pushes in." \
|
||||
-F "input_reference=@/absolute/path/to/reference.png" \
|
||||
-F "width=512" \
|
||||
-F "height=768" \
|
||||
-F "num_frames=121" \
|
||||
-F "fps=24" \
|
||||
-F "num_inference_steps=40" \
|
||||
-F "guidance_scale=4.0" \
|
||||
-o ltx2_i2v_h200.mp4
|
||||
```
|
||||
|
||||
#### Notes
|
||||
|
||||
- Memory usage: Model loads at ~73.5GB, peaks at ~73.5 GiB during inference.
|
||||
- Key flags:
|
||||
- `--model-class-name LTX2ImageToVideoPipeline` switches the same `Lightricks/LTX-2` weights into image-to-video mode.
|
||||
- `input_reference=@...` uploads the reference image directly
|
||||
- Known limitations:
|
||||
- Use `image_reference` when you want to pass a URL or JSON-safe image reference
|
||||
instead of uploading a file.
|
||||
- Do not send `input_reference` and `image_reference` in the same request.
|
||||
- This recipe is tested on 1x H200 validation; add measured latency and VRAM after running it on your target machine.
|
||||
- <https://huggingface.co/Lightricks/LTX-2>
|
||||
- <https://huggingface.co/Lightricks/LTX-2.3>
|
||||
- <https://huggingface.co/diffusers/LTX-2.3-Diffusers>
|
||||
- [Online video generation](../../docs/user_guide/examples/online_serving/text_to_video.md)
|
||||
- [Request batching](../../docs/user_guide/diffusion/request_batching.md)
|
||||
- [T2V offline example](../../examples/offline_inference/text_to_video/text_to_video.md)
|
||||
- [I2V offline example](../../examples/offline_inference/image_to_video/README.md)
|
||||
Executable
+162
@@ -0,0 +1,162 @@
|
||||
# MiniMax H3 on Ascend NPU
|
||||
|
||||
> Joint video and audio generation with text, first-frame, image/audio, or
|
||||
> multi-video conditions — Ascend NPU deployment guide
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: MiniMaxAI
|
||||
- Model: [`MiniMaxAI/MiniMax-H3`](https://huggingface.co/MiniMaxAI/MiniMax-H3)
|
||||
- Tasks: T2VA, FL2VA, and Ref2VA
|
||||
- Mode: OpenAI-compatible `/v1/videos` HTTP serving
|
||||
- Hardware: Atlas 800I A3
|
||||
- Maintainer: Community
|
||||
|
||||
This recipe adapts [MiniMax-H3.md](MiniMax-H3.md) for Ascend NPU
|
||||
environments. Differences from the GPU path:
|
||||
|
||||
- Runs on a CPU-only (aarch64) PyTorch build with `torch_npu`; no CUDA
|
||||
runtime is present.
|
||||
- Audio loading does **not** require TorchCodec (whose aarch64 wheels are
|
||||
built against CUDA torch and fail to load on CPU-only builds). vLLM-Omni
|
||||
automatically falls back to soundfile / ffmpeg for wav/mp3/m4a/mp4 audio
|
||||
inputs.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
### Checkpoint
|
||||
|
||||
Same as the GPU recipe — Hugging Face access approval is required:
|
||||
|
||||
```bash
|
||||
hf auth login
|
||||
export MODEL_ROOT=/path/to/MiniMax-H3
|
||||
hf download MiniMaxAI/MiniMax-H3 --local-dir "${MODEL_ROOT}"
|
||||
```
|
||||
|
||||
### Environment
|
||||
|
||||
- Ascend driver & firmware: 25.5.1
|
||||
- CANN toolkit: 9.0.1
|
||||
- Python: 3.12
|
||||
- PyTorch: 2.10.0+cpu
|
||||
- torch_npu: 2.10.0.post2
|
||||
- Install vLLM-Omni from a checkout with MiniMax H3 support:
|
||||
|
||||
```bash
|
||||
uv venv
|
||||
source .venv/bin/activate
|
||||
uv pip install -e .
|
||||
```
|
||||
|
||||
Install the **mindie-sd** operator library to enable Ascend-optimized fused
|
||||
operators (`adalayernorm`, etc.):
|
||||
|
||||
```bash
|
||||
git clone https://gitcode.com/Ascend/MindIE-SD.git && cd MindIE-SD
|
||||
|
||||
# Comment out the tik_ops build step (not needed for this use case)
|
||||
sed -i 's|^\(\s*\)source ${current_script_dir}/build_tik_ops.sh|\1# source ${current_script_dir}/build_tik_ops.sh|' build/build_ops.sh
|
||||
|
||||
python setup.py bdist_wheel
|
||||
cd dist
|
||||
pip install mindiesd-*.whl
|
||||
```
|
||||
|
||||
- `ffmpeg` and `ffprobe` must be available on `PATH`. They are used for
|
||||
reference-video preparation and MP4 output.
|
||||
- Reference-video decoding uses `decord` when available and falls back to
|
||||
PyAV otherwise.
|
||||
- Audio inputs do not require TorchCodec: wav/mp3/m4a/mp4 files are loaded
|
||||
at native sample rate through the soundfile / ffmpeg fallback.
|
||||
|
||||
## Start a server
|
||||
|
||||
One server loads one checkpoint partition. Set `MODEL` to `FL2VA` for T2VA
|
||||
and FL2VA requests, or to `Ref2VA` for Ref2VA requests.
|
||||
|
||||
### Multi-NPU: 768P validated configuration
|
||||
|
||||
Validated on eight NPUs of an Atlas 800I A3 server with Ulysses sequence
|
||||
parallelism degree 8, text-encoder tensor parallelism degree 8, native tiled
|
||||
VAE patch parallelism degree 8, and layerwise offload:
|
||||
|
||||
```bash
|
||||
export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
|
||||
export PORT=9098
|
||||
export MODEL="${MODEL_ROOT}/FL2VA"
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
export VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800
|
||||
export PYTHONDONTWRITEBYTECODE=1
|
||||
export HF_HUB_OFFLINE=1
|
||||
export TRANSFORMERS_OFFLINE=1
|
||||
|
||||
vllm serve "${MODEL}" \
|
||||
--omni \
|
||||
--host 0.0.0.0 \
|
||||
--port "${PORT}" \
|
||||
--trust-remote-code \
|
||||
--num-gpus 8 \
|
||||
--usp 8 \
|
||||
--ring 1 \
|
||||
--text-encoder-tp-size 8 \
|
||||
--enable-layerwise-offload \
|
||||
--vae-parallel-mode tile \
|
||||
--vae-use-tiling \
|
||||
--vae-patch-parallel-size 8 \
|
||||
--diffusion-attention-backend FLASH_ATTN
|
||||
```
|
||||
|
||||
Do not add `--enforce-eager`. The first request includes regional
|
||||
compilation; warm the server once before measuring steady-state latency.
|
||||
H3 is CFG-distilled, so `--cfg-parallel-size` must remain 1.
|
||||
|
||||
To serve Ref2VA, stop the FL2VA server and restart with:
|
||||
|
||||
```bash
|
||||
export MODEL="${MODEL_ROOT}/Ref2VA"
|
||||
```
|
||||
|
||||
## HTTP API examples
|
||||
|
||||
The request format is identical to the GPU recipe; see
|
||||
[MiniMax-H3.md § HTTP API examples](MiniMax-H3.md#http-api-examples).
|
||||
Use the validated 768P shapes (e.g. `width=1344 height=768`) on NPU.
|
||||
|
||||
## Key parameters
|
||||
|
||||
Same as the GPU recipe; see
|
||||
[MiniMax-H3.md § Key parameters](MiniMax-H3.md#key-parameters).
|
||||
The validated resolution on NPU is 768P (e.g. 1344x768).
|
||||
|
||||
## Validated NPU evidence
|
||||
|
||||
Measured on an Atlas 800I A3 server (8x NPU) with CANN 9.0.1,
|
||||
PyTorch 2.10.0+cpu, and torch_npu 2.10.0.post2, using the multi-NPU
|
||||
configuration above:
|
||||
|
||||
| Workload | Configuration |
|
||||
|----------|---------------|
|
||||
| T2VA, 209 frames, 1344x768 | TE TP8, layerwise offload, Ulysses 8, VPP8 tile, regional compile |
|
||||
| Ref2VA (prompt + video), 124 frames, 1344x768 | TE TP8, layerwise offload, Ulysses 8, VPP8 tile, regional compile |
|
||||
|
||||
These measurements describe the validated shapes rather than a general
|
||||
throughput guarantee.
|
||||
|
||||
## Known limitations
|
||||
|
||||
- Each server process loads only one checkpoint partition.
|
||||
- H3 currently executes one generation request per diffusion batch.
|
||||
- The first regional-compile request is a warmup and should not be included
|
||||
in steady-state performance measurements.
|
||||
- Image+audio Ref2VA accepts exactly one image and one audio reference.
|
||||
- Video Ref2VA accepts one or more video files, but not an additional
|
||||
standalone audio reference.
|
||||
- VAE patch parallelism requires size 1 or the full DiT group size and
|
||||
supports the H3 native `tile` mode only.
|
||||
|
||||
## Additional resources
|
||||
|
||||
- [MiniMax-H3.md](MiniMax-H3.md) — full GPU guide
|
||||
- [Supported models](../../docs/models/supported_models.md)
|
||||
- [Video API](../../docs/serving/videos_api.md)
|
||||
@@ -0,0 +1,329 @@
|
||||
# MiniMax H3
|
||||
|
||||
> Joint video and audio generation with text, first-frame, image/audio, or
|
||||
> multi-video conditions
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: MiniMaxAI
|
||||
- Model: [`MiniMaxAI/MiniMax-H3`](https://huggingface.co/MiniMaxAI/MiniMax-H3)
|
||||
- Tasks: T2VA, FL2VA, and Ref2VA
|
||||
- Mode: OpenAI-compatible `/v1/videos` HTTP serving
|
||||
- Maintainer: Community
|
||||
|
||||
MiniMax H3 is a CFG-distilled joint video/audio diffusion transformer. Its
|
||||
checkpoint has two independently served partitions:
|
||||
|
||||
- `FL2VA`: text-to-video+audio (`t2va`) and first-frame-to-video+audio
|
||||
(`fl2va`)
|
||||
- `Ref2VA`: image+audio or one-or-more-video reference-to-video+audio
|
||||
(`ref2va`)
|
||||
|
||||
The generated MP4 contains H.264 video and synchronized stereo audio.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
The checkpoint requires Hugging Face access approval. Download it and point
|
||||
`MODEL_ROOT` at the directory containing the `FL2VA` and `Ref2VA`
|
||||
subdirectories:
|
||||
|
||||
```bash
|
||||
hf auth login
|
||||
export MODEL_ROOT=/path/to/MiniMax-H3
|
||||
hf download MiniMaxAI/MiniMax-H3 --local-dir "${MODEL_ROOT}"
|
||||
```
|
||||
|
||||
Install vLLM-Omni from the checkout containing MiniMax H3 support. On
|
||||
Blackwell, install the optional FlashAttention-4 dependency:
|
||||
|
||||
```bash
|
||||
uv venv
|
||||
source .venv/bin/activate
|
||||
uv pip install -e '.[fa4]'
|
||||
```
|
||||
|
||||
`ffmpeg` and `ffprobe` must be available on `PATH`. They are used for
|
||||
reference-video preparation and MP4 output.
|
||||
|
||||
## Start a server
|
||||
|
||||
One server loads one checkpoint partition. Set `MODEL` to `FL2VA` for T2VA
|
||||
and FL2VA requests, or to `Ref2VA` for either Ref2VA request.
|
||||
|
||||
### Single GPU: accuracy and memory first
|
||||
|
||||
The single-GPU configuration uses model-level CPU offload.
|
||||
This matches the accuracy-qualified reference path and prevents the Qwen3-VL
|
||||
encoder and DiT from being resident on the GPU at the same time.
|
||||
|
||||
```bash
|
||||
export MODEL="${MODEL_ROOT}/FL2VA"
|
||||
export PORT=8091
|
||||
|
||||
CUDA_VISIBLE_DEVICES=0 \
|
||||
VLLM_WORKER_MULTIPROC_METHOD=spawn \
|
||||
VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800 \
|
||||
vllm serve "${MODEL}" \
|
||||
--omni \
|
||||
--host 0.0.0.0 \
|
||||
--port "${PORT}" \
|
||||
--trust-remote-code \
|
||||
--num-gpus 1 \
|
||||
--enable-cpu-offload \
|
||||
--diffusion-attention-backend FLASH_ATTN
|
||||
```
|
||||
|
||||
Use a GPU with enough memory for the active H3 component and enough system RAM
|
||||
for the offloaded components. CPU offload reduces GPU memory pressure but adds
|
||||
PCIe/NVLink transfer latency.
|
||||
|
||||
### Four GPUs: measured best-practice throughput
|
||||
|
||||
The best validated four-GPU configuration on four NVIDIA B300 GPUs is:
|
||||
|
||||
- no CPU or layerwise offload;
|
||||
- Ulysses sequence parallelism degree 4;
|
||||
- native tiled VAE patch parallelism degree 4;
|
||||
- regional `torch.compile` for the repeated DiT blocks;
|
||||
- FlashAttention, with Ring and TP left at 1.
|
||||
|
||||
```bash
|
||||
export MODEL="${MODEL_ROOT}/FL2VA"
|
||||
export PORT=8091
|
||||
|
||||
CUDA_VISIBLE_DEVICES=0,1,2,3 \
|
||||
FLASHINFER_DISABLE_VERSION_CHECK=1 \
|
||||
VLLM_WORKER_MULTIPROC_METHOD=spawn \
|
||||
VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800 \
|
||||
vllm serve "${MODEL}" \
|
||||
--omni \
|
||||
--host 0.0.0.0 \
|
||||
--port "${PORT}" \
|
||||
--trust-remote-code \
|
||||
--num-gpus 4 \
|
||||
--usp 4 \
|
||||
--ring 1 \
|
||||
--vae-patch-parallel-size 4 \
|
||||
--vae-parallel-mode tile \
|
||||
--vae-use-tiling \
|
||||
--diffusion-attention-backend FLASH_ATTN
|
||||
```
|
||||
|
||||
Do not add `--enforce-eager` to this performance configuration. The first
|
||||
request includes regional compilation; warm the server once before measuring
|
||||
steady-state latency. H3 is CFG-distilled, so `--cfg-parallel-size` must remain
|
||||
1. The H3 VAE supports its native `tile` mode, not
|
||||
`spatial_shard_height` or `spatial_shard_width`.
|
||||
|
||||
### Text encoder tensor parallelism
|
||||
|
||||
The Qwen3-VL text encoder (~51.5 GB in BF16 for the retained 50 layers) is by
|
||||
default fully resident on the DiT main rank. On multi-GPU no-offload runs that
|
||||
rank becomes the peak-memory hotspot. Add `--text-encoder-tp-size N` to shard
|
||||
the encoder across the first `N` DiT ranks (the encoder is implemented with
|
||||
vLLM-style tensor-parallel layers and runs with distributed collectives over
|
||||
its own encoder process group):
|
||||
|
||||
```bash
|
||||
export MODEL="${MODEL_ROOT}/FL2VA"
|
||||
export PORT=8091
|
||||
|
||||
CUDA_VISIBLE_DEVICES=0,1,2,3 \
|
||||
FLASHINFER_DISABLE_VERSION_CHECK=1 \
|
||||
VLLM_WORKER_MULTIPROC_METHOD=spawn \
|
||||
VLLM_OMNI_VIDEO_SYNC_TIMEOUT=1800 \
|
||||
vllm serve "${MODEL}" \
|
||||
--omni \
|
||||
--host 0.0.0.0 \
|
||||
--port "${PORT}" \
|
||||
--trust-remote-code \
|
||||
--num-gpus 4 \
|
||||
--usp 4 \
|
||||
--ring 1 \
|
||||
--text-encoder-tp-size 4 \
|
||||
--vae-patch-parallel-size 4 \
|
||||
--vae-parallel-mode tile \
|
||||
--vae-use-tiling \
|
||||
--diffusion-attention-backend FLASH_ATTN
|
||||
```
|
||||
|
||||
`N` must divide the Qwen3-VL head counts (64 attention heads / 8 KV heads), so
|
||||
valid values on a 4-GPU server are 1, 2, and 4 (1, 2, 4, 8 on 8 GPUs). The
|
||||
encoder TP rank set is the first `N` DiT ranks; on 4 GPUs
|
||||
`--text-encoder-tp-size 4` shards the encoder 4-way, dropping the DiT main
|
||||
rank's no-offload peak by roughly `(N-1)/N` of the ~51.5 GB encoder while the
|
||||
other ranks each gain `~51.5/N` GB. The encoder output remains identical to
|
||||
the reference path within bf16 rounding: every encoder rank all-reduces the
|
||||
row-parallel projections, so the full `[seq, 5120]` layer-50 hidden state is
|
||||
replicated on every rank.
|
||||
|
||||
To serve Ref2VA, stop the FL2VA server and restart the same one- or four-GPU
|
||||
command with:
|
||||
|
||||
```bash
|
||||
export MODEL="${MODEL_ROOT}/Ref2VA"
|
||||
```
|
||||
|
||||
## HTTP API examples
|
||||
|
||||
The following requests use the synchronous endpoint so the returned body can
|
||||
be saved directly as an MP4. The asynchronous `POST /v1/videos` endpoint can
|
||||
also be used when job polling is preferred.
|
||||
|
||||
All four tasks use 24 FPS, 50 sigma points, seed values from the validated
|
||||
workloads, and the checkpoint-reference video/audio flow shifts of 12 and 3.
|
||||
Decimal durations are passed through `extra_params`.
|
||||
|
||||
Set the endpoint once:
|
||||
|
||||
```bash
|
||||
export API_URL="http://127.0.0.1:${PORT}/v1/videos/sync"
|
||||
```
|
||||
|
||||
### 1. T2VA: text to video and audio
|
||||
|
||||
Run this request against the `FL2VA` partition:
|
||||
|
||||
```bash
|
||||
curl -sS -X POST "${API_URL}" \
|
||||
-F 'prompt=In a snowy blue-purple forest, Ori carefully walks past a sleeping giant; footsteps crunch in the snow while the creature breathes and softly snorts.' \
|
||||
-F 'width=1344' \
|
||||
-F 'height=768' \
|
||||
-F 'fps=24' \
|
||||
-F 'num_inference_steps=50' \
|
||||
-F 'flow_shift=12' \
|
||||
-F 'seed=1101' \
|
||||
-F 'extra_params={"task":"t2va","duration":8.7,"audio_flow_shift":3.0}' \
|
||||
-o t2va.mp4
|
||||
```
|
||||
|
||||
### 2. FL2VA: first frame to video and audio
|
||||
|
||||
Run this request against the `FL2VA` partition. When width and height are
|
||||
omitted, H3 preserves the first-frame aspect ratio and uses a 768-pixel short
|
||||
edge.
|
||||
|
||||
```bash
|
||||
export FIRST_FRAME=/path/to/fl2va_first_frame.png
|
||||
|
||||
curl -sS -X POST "${API_URL}" \
|
||||
-F 'prompt=A man stands beside a yellow car at night. The car drives away; he follows it with his eyes and begins singing sadly, with synchronized voice and city ambience.' \
|
||||
-F 'fps=24' \
|
||||
-F 'num_inference_steps=50' \
|
||||
-F 'flow_shift=12' \
|
||||
-F 'seed=2101' \
|
||||
-F 'extra_params={"task":"fl2va","duration":8.7,"audio_flow_shift":3.0}' \
|
||||
-F "input_reference=@${FIRST_FRAME};type=image/png" \
|
||||
-o fl2va.mp4
|
||||
```
|
||||
|
||||
### 3. Ref2VA: image and audio reference
|
||||
|
||||
Run this request against the `Ref2VA` partition. `audio_reference` accepts an
|
||||
HTTP(S) URL or a `data:` URL. In one terminal, expose the local reference
|
||||
assets to the serving host:
|
||||
|
||||
```bash
|
||||
python -m http.server 8092 \
|
||||
--bind 127.0.0.1 \
|
||||
--directory /path/to/reference_assets
|
||||
```
|
||||
|
||||
Then submit the image and audio request from another terminal:
|
||||
|
||||
```bash
|
||||
export REF_IMAGE=/path/to/reference_assets/ref2va_image.png
|
||||
export AUDIO_URL=http://127.0.0.1:8092/ref2va_audio.mp3
|
||||
|
||||
curl -sS -X POST "${API_URL}" \
|
||||
-F 'prompt=A white cat with black mustache and eyebrow markings sits on a beige couch, lip-syncing precisely to the complete reference audio before shifting from confusion to deadpan speechlessness.' \
|
||||
-F 'width=1344' \
|
||||
-F 'height=768' \
|
||||
-F 'fps=24' \
|
||||
-F 'num_inference_steps=50' \
|
||||
-F 'flow_shift=12' \
|
||||
-F 'seed=3101' \
|
||||
-F 'extra_params={"task":"ref2va","duration":15.0,"audio_flow_shift":3.0}' \
|
||||
-F "input_reference=@${REF_IMAGE};type=image/png" \
|
||||
-F "audio_reference={\"audio_url\":\"${AUDIO_URL}\"}" \
|
||||
-o ref2va_image_audio.mp4
|
||||
```
|
||||
|
||||
The requested duration should cover the complete audio. If `duration` is
|
||||
shorter, the reference soundtrack is truncated to the generated clip.
|
||||
|
||||
### 4. Ref2VA: two video references
|
||||
|
||||
Run this request against the `Ref2VA` partition. Repeat the
|
||||
`input_references` multipart field once per source video. H3 consumes the
|
||||
videos in form order and preserves their original soundtracks during
|
||||
conditioning.
|
||||
|
||||
```bash
|
||||
export SUBJECT_VIDEO=/path/to/green_screen_subject.mp4
|
||||
export BACKGROUND_VIDEO=/path/to/fairytale_background.mov
|
||||
|
||||
curl -sS -X POST "${API_URL}" \
|
||||
-F 'prompt=Remove the green screen background of Video 1 and replace it with the fairytale environment from Video 2. Match the background motion to the character actions and relight the character to fit the scene.' \
|
||||
-F 'width=1344' \
|
||||
-F 'height=768' \
|
||||
-F 'fps=24' \
|
||||
-F 'num_inference_steps=50' \
|
||||
-F 'flow_shift=12' \
|
||||
-F 'seed=3101' \
|
||||
-F 'extra_params={"task":"ref2va","duration":15.0,"audio_flow_shift":3.0}' \
|
||||
-F "input_references=@${SUBJECT_VIDEO};type=video/mp4" \
|
||||
-F "input_references=@${BACKGROUND_VIDEO};type=video/quicktime" \
|
||||
-o ref2va_video_video.mp4
|
||||
```
|
||||
|
||||
The server stores uploaded multi-video references only for the lifetime of the
|
||||
request and deletes the temporary files after generation. Video Ref2VA uses
|
||||
the source-video soundtracks and does not accept a separate
|
||||
`audio_reference`.
|
||||
|
||||
## Key parameters
|
||||
|
||||
| Parameter | Recommended value | Notes |
|
||||
|-----------|-------------------|-------|
|
||||
| `task` | `t2va`, `fl2va`, or `ref2va` | Passed in `extra_params`; must match the served partition |
|
||||
| `duration` | Workload-specific | Decimal seconds in `extra_params`; converted to H3-compatible frame count |
|
||||
| `fps` | `24` | H3 output FPS is fixed |
|
||||
| `num_inference_steps` | `50` | Matches the reference accuracy workloads |
|
||||
| `flow_shift` | `12` | Video sigma shift |
|
||||
| `audio_flow_shift` | `3` | Audio sigma shift, passed in `extra_params` |
|
||||
| `seed` | Task-specific | Use a fixed value for reproducibility |
|
||||
| `width`, `height` | Multiples of 32 | Aspect ratio must be between 1:4 and 4:1 |
|
||||
|
||||
## Validated four-GPU evidence
|
||||
|
||||
The four-GPU recommendation was measured on four NVIDIA B300 GPUs with one
|
||||
excluded warmup followed by three requests.
|
||||
|
||||
| Workload | Configuration | Observed result |
|
||||
|----------|---------------|-----------------|
|
||||
| FL2VA, 209 frames, 1248x768 | no offload, U4, VPP4 tile, regional compile | 86.964 s mean HTTP client latency |
|
||||
| Two-video Ref2VA, 362 frames, 1344x768 | no offload, U4, VPP4 tile, regional compile | 784.394 s accounted model-stage mean |
|
||||
|
||||
These measurements describe the validated shapes rather than a general
|
||||
throughput guarantee. Multi-video Ref2VA is much slower because the two
|
||||
reference videos expand both the Qwen3-VL vision sequence and the packed DiT
|
||||
attention sequence.
|
||||
|
||||
## Known limitations
|
||||
|
||||
- Each server process loads only one checkpoint partition.
|
||||
- H3 currently executes one generation request per diffusion batch.
|
||||
- The first regional-compile request is a warmup and should not be included in
|
||||
steady-state performance measurements.
|
||||
- Image+audio Ref2VA accepts exactly one image and one audio reference.
|
||||
- Video Ref2VA accepts one or more video files, but not an additional standalone
|
||||
audio reference.
|
||||
- VAE patch parallelism requires size 1 or the full DiT group size and supports
|
||||
the H3 native `tile` mode only.
|
||||
|
||||
## Additional resources
|
||||
|
||||
- [Supported models](../../docs/models/supported_models.md)
|
||||
- [Video API](../../docs/serving/videos_api.md)
|
||||
- [Diffusion parallelism](../../docs/user_guide/diffusion/parallelism/overview.md)
|
||||
@@ -69,8 +69,13 @@ Notes:
|
||||
## Verification
|
||||
|
||||
```python
|
||||
from tests.gr00t.openpi_client_helper import run_policy_session, validate_session_result
|
||||
validate_session_result(run_policy_session(host="127.0.0.1", port=8000))
|
||||
from tests.helpers.runtime import OpenAIClientHandler
|
||||
|
||||
handler = OpenAIClientHandler(host="127.0.0.1", port=8000, log_stats=False)
|
||||
response = handler.send_robot_openpi_ws_request(
|
||||
{"run_default_policy_session": True, "session_id": "gr00t-smoke"}
|
||||
)[0]
|
||||
assert response.operation_responses[-1]["status"] == "reset successful"
|
||||
```
|
||||
|
||||
Or run the e2e test suite:
|
||||
|
||||
@@ -0,0 +1,142 @@
|
||||
# Nemotron-Labs Audex (2B / 30B-A3B)
|
||||
|
||||
> Speech suite on 1×H100: TTS, text-to-audio, audio understanding, and cascaded speech-to-speech
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: NVIDIA
|
||||
- Model: `nvidia/Nemotron-Labs-Audex-2B`, `nvidia/Nemotron-Labs-Audex-30B-A3B`
|
||||
- Task: Text-to-speech (English), text-to-audio (sound effects), audio understanding (ASR / audio QA), cascaded speech-to-speech
|
||||
- Mode: Online serving (`/v1/audio/speech`, `/v1/chat/completions`) and offline batch inference
|
||||
- Maintainer: Community
|
||||
|
||||
## When to use this recipe
|
||||
|
||||
Use this recipe to serve either Audex checkpoint through one of its four
|
||||
deployment pipelines. Both sizes share the same repo layout, codec token
|
||||
space, and streaming causal speech decoder (16 kHz output); they differ in
|
||||
the thinker: the 2B is a dense LM, the 30B-A3B is a hybrid Mamba + MoE
|
||||
NemotronH (~3B active parameters). Pick the pipeline by task:
|
||||
|
||||
| pipeline (`vllm_omni/deploy/<name>[_30b].yaml`) | audio in | text out | speech out | general audio out | endpoint |
|
||||
|---|---|---|---|---|---|
|
||||
| `audex_tts` | ❌ | ❌ | ✅ | ❌ | `/v1/audio/speech` |
|
||||
| `audex_tta` | ❌ | ❌ | ❌ | ✅ | `/v1/audio/speech` |
|
||||
| `audex_thinker_only` | ✅ | ✅ | ❌ | ❌ | `/v1/chat/completions` |
|
||||
| `audex_s2s` | ✅ | ✅ | ✅ | ❌ | both |
|
||||
|
||||
## References
|
||||
|
||||
- Model cards: [Nemotron-Labs-Audex-2B](https://huggingface.co/nvidia/Nemotron-Labs-Audex-2B), [Nemotron-Labs-Audex-30B-A3B](https://huggingface.co/nvidia/Nemotron-Labs-Audex-30B-A3B)
|
||||
- Offline examples: [`examples/offline_inference/audex/`](../../examples/offline_inference/audex/)
|
||||
- Online examples: [`examples/online_serving/audex/`](../../examples/online_serving/audex/)
|
||||
|
||||
## Hardware Support
|
||||
|
||||
## GPU
|
||||
|
||||
### 1×H100 80GB — Audex-2B
|
||||
|
||||
#### Environment
|
||||
|
||||
- OS: Linux
|
||||
- Python: 3.12+
|
||||
- CUDA: 12.x
|
||||
- vLLM version: 0.24.0
|
||||
- vLLM-Omni version or commit: PR #4976 (branch `audex`)
|
||||
|
||||
#### Command
|
||||
|
||||
**Online serving** (TTS is the default pipeline; the repo root's
|
||||
`model_type` auto-resolves `audex_tts.yaml`):
|
||||
|
||||
```bash
|
||||
vllm-omni serve nvidia/Nemotron-Labs-Audex-2B \
|
||||
--host 0.0.0.0 --port 8097 \
|
||||
--trust-remote-code --omni
|
||||
# other modes: add --stage-configs-path vllm_omni/deploy/audex_{tta,thinker_only,s2s}.yaml
|
||||
# or use the launcher: MODE=s2s PORT=8098 examples/online_serving/audex/run_server.sh
|
||||
```
|
||||
|
||||
**Offline batch inference** (one script per task, each defaults to its
|
||||
correct deploy yaml):
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/audex/text_to_speech.py \
|
||||
--texts "Hello world." --cfg-scale 1.5 --output-dir results/audex_wavs
|
||||
python examples/offline_inference/audex/text_to_audio.py # sound effects
|
||||
python examples/offline_inference/audex/audio_qa.py # ASR / audio QA
|
||||
python examples/offline_inference/audex/speech_to_speech.py # 3-pass cascade
|
||||
```
|
||||
|
||||
#### Verification
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8097/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "nvidia/Nemotron-Labs-Audex-2B",
|
||||
"input": "Hello, how are you?",
|
||||
"response_format": "wav",
|
||||
"extra_params": {"cfg_scale": 1.5}
|
||||
}' --output hello.wav
|
||||
# or the universal client (all four modes):
|
||||
python examples/online_serving/audex/client.py --mode tts --output hello.wav
|
||||
```
|
||||
|
||||
#### Notes
|
||||
|
||||
- Memory usage: TTS deploy — stage 0 (thinker) 0.4, stage 1 (decoder) 0.25
|
||||
GPU-memory fraction; both stages share GPU 0.
|
||||
- Key flags: `--trust-remote-code` and `--omni` are required.
|
||||
- Classifier-free guidance: `extra_params.cfg_scale` (1.0 = off; 1.5 is the
|
||||
official quality setting; TTA effectively requires its default 3.0).
|
||||
Measured en-24 self-transcribed CER: 6.87% guided vs 7.24% unguided.
|
||||
- Output: 16 kHz mono WAV. English, single built-in voice (no cloning).
|
||||
- TTA decodes through the external XCodec1 checkpoint
|
||||
(`hf-audio/xcodec-hubert-general-balanced`, auto-downloaded; override
|
||||
with `XCODEC1_PATH`).
|
||||
|
||||
### 1×H100 80GB — Audex-30B-A3B
|
||||
|
||||
#### Environment
|
||||
|
||||
Same as the 2B row above.
|
||||
|
||||
#### Command
|
||||
|
||||
The 30B REQUIRES its explicit deploy yaml (the shared HF `model_type`
|
||||
auto-resolves to the 2B-tuned config otherwise):
|
||||
|
||||
```bash
|
||||
vllm-omni serve nvidia/Nemotron-Labs-Audex-30B-A3B \
|
||||
--host 0.0.0.0 --port 8097 \
|
||||
--trust-remote-code --omni \
|
||||
--stage-configs-path vllm_omni/deploy/audex_tts_30b.yaml
|
||||
# or: SIZE=30b MODE=tts examples/online_serving/audex/run_server.sh
|
||||
```
|
||||
|
||||
Offline: same four scripts with
|
||||
`--model nvidia/Nemotron-Labs-Audex-30B-A3B --deploy-config vllm_omni/deploy/audex_<mode>_30b.yaml`.
|
||||
|
||||
#### Verification
|
||||
|
||||
Same curl / `client.py` as the 2B with
|
||||
`"model": "nvidia/Nemotron-Labs-Audex-30B-A3B"` (the server validates the
|
||||
model id). First launch downloads ~60 GB.
|
||||
|
||||
#### Notes
|
||||
|
||||
- Memory usage: ~61 GiB weights on one H100 80 GB (stage 0 at 0.85 GPU
|
||||
fraction + decoder at 0.08 on the same card); verified with healthy KV
|
||||
headroom. If long sequences OOM, set `tensor_parallel_size: 2` on
|
||||
stage 0 (documented in the yaml headers).
|
||||
- `enable_prefix_caching: false` on stage 0 — hybrid Mamba constraint.
|
||||
- CUDA graphs stay ON for the thinker: the NemotronH stage captures
|
||||
cleanly and eager mode costs ~7× single-stream TTS RTF (review-measured
|
||||
0.30 vs 2.08).
|
||||
- Quality (en-24 self-transcribed CER): 0.72% unguided baseline; 0.72% at
|
||||
cfg 1.5 with guided temperature 0.1.
|
||||
- Known limitation: the S2S chat pass needs the official
|
||||
formatting-instruction prompt (the bundled examples/client send it);
|
||||
a bare transcript turn makes the 30B answer in speech-codec tokens.
|
||||
@@ -30,8 +30,6 @@ also provided.
|
||||
`hf_config.version="4.5"`):
|
||||
- Default single-GPU compatibility layout (auto-loaded):
|
||||
[`vllm_omni/deploy/minicpmo_4_5.yaml`](../../vllm_omni/deploy/minicpmo_4_5.yaml)
|
||||
- Recommended 2-GPU continuous-batching layout:
|
||||
[`vllm_omni/deploy/minicpmo_4_5_batching.yaml`](../../vllm_omni/deploy/minicpmo_4_5_batching.yaml)
|
||||
- 2-GPU and 3-GPU layouts:
|
||||
[`vllm_omni/deploy/minicpmo_4_5_2gpu.yaml`](../../vllm_omni/deploy/minicpmo_4_5_2gpu.yaml),
|
||||
[`vllm_omni/deploy/minicpmo_4_5_3gpu.yaml`](../../vllm_omni/deploy/minicpmo_4_5_3gpu.yaml)
|
||||
@@ -68,7 +66,7 @@ Code2Wav consumes them through a shared-memory async connector.
|
||||
MiniCPM-o 4.5 now requires the three-stage topology: the Talker owns
|
||||
request-local codec generation and Code2Wav owns waveform state and
|
||||
reference-voice prompt features. `minicpmo_4_5.yaml` remains the stable
|
||||
single-GPU entry point; `minicpmo_4_5_batching.yaml` is the recommended
|
||||
single-GPU entry point; `minicpmo_4_5_2gpu.yaml` is the recommended
|
||||
two-GPU profile. The removed fused two-stage implementation is not retained as
|
||||
a fallback because it would duplicate state machines and correctness paths.
|
||||
|
||||
@@ -79,7 +77,9 @@ a fallback because it would duplicate state machines and correctness paths.
|
||||
The default
|
||||
[`vllm_omni/deploy/minicpmo_4_5.yaml`](../../vllm_omni/deploy/minicpmo_4_5.yaml)
|
||||
co-locates Thinker, codec-only Talker, and Code2Wav on GPU 0. Their
|
||||
`gpu_memory_utilization` budgets are 0.55, 0.22, and 0.22. All stages admit
|
||||
`gpu_memory_utilization` budgets are 0.55, 0.15, and 0.15. The remaining
|
||||
device memory is left available for runtime workspaces such as the HiFi-GAN
|
||||
vocoder's cuDNN kernels. All stages admit
|
||||
up to four sequences. Startup video profiling is bounded to 32 frames per
|
||||
video. Use the two-GPU profile for production throughput.
|
||||
|
||||
@@ -105,7 +105,7 @@ The deploy config is auto-loaded by the model registry — no
|
||||
For the recommended two-GPU layout, add:
|
||||
|
||||
```bash
|
||||
--deploy-config vllm_omni/deploy/minicpmo_4_5_batching.yaml
|
||||
--deploy-config vllm_omni/deploy/minicpmo_4_5_2gpu.yaml
|
||||
```
|
||||
|
||||
#### Performance comparison
|
||||
@@ -175,16 +175,18 @@ in another choice's `message.audio.data` (24 kHz mono, see Notes). With
|
||||
`modalities: ["text", "audio"]` you typically get two `choices` entries
|
||||
(one text, one audio).
|
||||
|
||||
**Streaming text + speech**:
|
||||
**Streaming text + speech** (use `--stream`):
|
||||
|
||||
```bash
|
||||
python examples/online_serving/minicpmo/streaming_chat_completion.py \
|
||||
--base-url http://localhost:8099/v1 \
|
||||
--output minicpmo_stream.wav
|
||||
python examples/online_serving/minicpmo/openai_chat_completion_client_for_multimodal_generation.py \
|
||||
--query-type text \
|
||||
--prompt "Say hello, then introduce vLLM in one sentence." \
|
||||
--port 8099 \
|
||||
--stream
|
||||
```
|
||||
|
||||
The client prints text deltas as they arrive and reconstructs one valid WAV
|
||||
from the independently encoded audio deltas.
|
||||
The client prints text deltas as they arrive and saves streamed audio chunks
|
||||
to WAV files.
|
||||
|
||||
**Gradio demo (text + image + audio + video UI)**:
|
||||
|
||||
@@ -202,9 +204,10 @@ speech output (TTS)"** checkbox on / off.
|
||||
|
||||
#### Notes
|
||||
|
||||
- Memory budget: Thinker, Talker, and Code2Wav reserve 0.55, 0.22, and
|
||||
0.22 of GPU 0. The larger Thinker share protects its multimodal KV cache;
|
||||
all three model processes still share one CUDA device.
|
||||
- Memory budget: Thinker, Talker, and Code2Wav reserve 0.55, 0.15, and
|
||||
0.15 of GPU 0. The larger Thinker share protects its multimodal KV cache,
|
||||
while the unreserved memory remains available to runtime kernels; all three
|
||||
model processes still share one CUDA device.
|
||||
- `--trust-remote-code` is required — the HF repo ships a custom
|
||||
`MiniCPMO` config / model class.
|
||||
- Stage 0 Thinker and Stage 1 Talker enable vLLM CUDA Graphs. Stage 2 remains
|
||||
|
||||
+3
-2
@@ -45,9 +45,9 @@ recipes/
|
||||
| [`inclusionAI/Ming-flash-omni-2.0.md`](./inclusionAI/Ming-flash-omni-2.0.md) | Online serving for multimodal chat + standalone TTS | 4x H100 / 1x H100 80GB |
|
||||
| [`inclusionAI/Ming-omni-tts.md`](./inclusionAI/Ming-omni-tts.md) | Offline + online dense Ming TTS/audio generation | 1x H100 80GB / 1x AMD MI300X (ROCm 7.2) |
|
||||
| [`IndexTeam/IndexTTS-2.md`](./IndexTeam/IndexTTS-2.md) | Online serving for voice-cloned TTS with optional emotion control | 1x L4 24GB or larger CUDA GPU |
|
||||
| [`MiniMaxAI/MiniMax-H3.md`](./MiniMaxAI/MiniMax-H3.md) | T2VA, FL2VA, image+audio Ref2VA, and multi-video Ref2VA serving | 1x GPU with CPU offload / 4x B300 |
|
||||
| [`krea/Krea-2.md`](./krea/Krea-2.md) | Text-to-image (Turbo + Raw), offline + online, with LoRA | 1x H100 80GB |
|
||||
| [`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.md`](./LTX/LTX-2.md) | LTX-2/LTX-2.3 text-to-video and image-to-video with synchronized audio | H200 141GB / 96GB-class GPU |
|
||||
| [`MammothModa2/MammothModa2.md`](./MammothModa2/MammothModa2.md) | Preview and Dev text-to-image (AR → DiT); Dev text/image understanding | Preview: 1x L40S 48GB / 1x ≥40GB GPU; Dev: 1x NVIDIA GPU with sufficient cache headroom |
|
||||
| [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB |
|
||||
| [`Robbyant/LingBot-Video.md`](./Robbyant/LingBot-Video.md) | Native dense and MoE text-to-video serving | Dense: 1x L20X; MoE: 1x L20X (~67.7 GiB peak) |
|
||||
@@ -57,6 +57,7 @@ recipes/
|
||||
| [`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 |
|
||||
| [`OpenMOSS/MOSS-TTS.md`](./OpenMOSS/MOSS-TTS.md) | Online + offline multilingual TTS (MOSS-TTS family, 8B) | 1x H100 80GB |
|
||||
| [`NVIDIA/Nemotron-Labs-Audex.md`](./NVIDIA/Nemotron-Labs-Audex.md) | TTS / text-to-audio / audio understanding / cascaded S2S (Audex 2B + 30B-A3B) | 1x H100 80GB |
|
||||
| [`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-2512.md`](./Qwen/Qwen-Image-2512.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt FP8 / mixed FP8/NVFP4 | 1x A800 80GB / 2x B200 |
|
||||
| [`Qwen/Qwen-Image-Edit.md`](./Qwen/Qwen-Image-Edit.md) | Text-guided single-image editing | 1x or 2x H200 141GB |
|
||||
|
||||
@@ -0,0 +1,191 @@
|
||||
# Cosmos3 — Distributed Layerwise Offload
|
||||
|
||||
> Sharding weights across DP ranks with H2D + AllGather overlap for multi-NPU/GPU deployment
|
||||
|
||||
## When to use this recipe
|
||||
|
||||
Use this recipe when deploying Cosmos3-Super (64B / 124 GB) or Cosmos3-Nano (17B / 33 GB) across multiple devices where the model does not fit on a single device's HBM, or when you need higher throughput via DP multi-concurrency.
|
||||
|
||||
**Key benefits:**
|
||||
- Each rank stores only 1/dp_size of the model weights on host memory
|
||||
- Only 2 transformer blocks reside on each device at any time (fixed double-buffer)
|
||||
- DP multi-concurrency: N concurrent requests with near-linear throughput scaling
|
||||
- mmap weight loading: shared page cache eliminates O(dp_size × model_size) RSS
|
||||
|
||||
## Prerequisites
|
||||
|
||||
- 2+ NPU or GPU devices (Ascend 910B3, NVIDIA B300, etc.)
|
||||
- vLLM-Omni with distributed layerwise offload support
|
||||
- Cosmos3-Nano or Cosmos3-Super checkpoint (local path or HuggingFace repo ID)
|
||||
|
||||
## Serve commands
|
||||
|
||||
### Cosmos3-Nano with DP=4 + AllGather (recommended for max throughput)
|
||||
|
||||
```bash
|
||||
vllm serve /path/to/Cosmos3-Nano \
|
||||
--omni \
|
||||
--host 0.0.0.0 --port 8000 \
|
||||
--no-guardrails --init-timeout 3600 \
|
||||
--vae-use-tiling --vae-patch-parallel-size 1 \
|
||||
--enable-distributed-layerwise-offload \
|
||||
--data-parallel-size 4
|
||||
```
|
||||
|
||||
### Cosmos3-Super with DP=2 + AllGather (for models > single-card HBM)
|
||||
|
||||
```bash
|
||||
vllm serve /path/to/Cosmos3-Super \
|
||||
--omni \
|
||||
--host 0.0.0.0 --port 8000 \
|
||||
--no-guardrails --init-timeout 3600 \
|
||||
--vae-use-tiling --vae-patch-parallel-size 1 \
|
||||
--enable-distributed-layerwise-offload \
|
||||
--data-parallel-size 2
|
||||
```
|
||||
|
||||
### Cosmos3-Super with DP=4 + AllGather
|
||||
|
||||
```bash
|
||||
vllm serve /path/to/Cosmos3-Super \
|
||||
--omni \
|
||||
--host 0.0.0.0 --port 8000 \
|
||||
--no-guardrails --init-timeout 3600 \
|
||||
--vae-use-tiling --vae-patch-parallel-size 1 \
|
||||
--enable-distributed-layerwise-offload \
|
||||
--data-parallel-size 4
|
||||
```
|
||||
|
||||
### Without AllGather (full weights per rank, no sharding)
|
||||
|
||||
Use `--dlo-no-use-allgather` when you want each rank to load full weights independently —
|
||||
no AllGather synchronization, no concurrent request requirement, but N× host memory.
|
||||
|
||||
```bash
|
||||
vllm serve /path/to/Cosmos3-Nano \
|
||||
--omni \
|
||||
--host 0.0.0.0 --port 8000 \
|
||||
--no-guardrails --init-timeout 3600 \
|
||||
--vae-use-tiling --vae-patch-parallel-size 1 \
|
||||
--enable-distributed-layerwise-offload \
|
||||
--data-parallel-size 4 \
|
||||
--dlo-no-use-allgather
|
||||
```
|
||||
|
||||
### With SP (sequence parallel) instead of DP
|
||||
|
||||
For long-sequence workloads where SP parallelism is preferred:
|
||||
|
||||
```bash
|
||||
vllm serve /path/to/Cosmos3-Super \
|
||||
--omni \
|
||||
--host 0.0.0.0 --port 8000 \
|
||||
--no-guardrails --init-timeout 3600 \
|
||||
--vae-use-tiling --vae-patch-parallel-size 1 \
|
||||
--enable-distributed-layerwise-offload \
|
||||
--usp 4
|
||||
```
|
||||
|
||||
## Request examples
|
||||
|
||||
### Text-to-Image (T2I)
|
||||
|
||||
```bash
|
||||
curl -s -o output.png -X POST "http://localhost:8000/v1/images/generations" \
|
||||
-H "Content-Type: multipart/form-data" \
|
||||
-F "model=/path/to/Cosmos3-Nano" \
|
||||
-F "prompt=A robot arm cleaning a plate, cinematic" \
|
||||
-F "size=1024x1024" \
|
||||
-F "num_inference_steps=50" \
|
||||
-F "guidance_scale=6.0" \
|
||||
-F "seed=42"
|
||||
```
|
||||
|
||||
### Text-to-Video (T2V)
|
||||
|
||||
```bash
|
||||
curl -s -o output.mp4 -X POST "http://localhost:8000/v1/videos/sync" \
|
||||
-H "Accept: video/mp4" \
|
||||
-F "model=/path/to/Cosmos3-Nano" \
|
||||
-F "prompt=A robot arm cleaning a plate, cinematic shot" \
|
||||
-F "negative_prompt=blurry" \
|
||||
-F "size=832x480" \
|
||||
-F "num_frames=29" \
|
||||
-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":false}' \
|
||||
-F "seed=17"
|
||||
```
|
||||
|
||||
> **Important:** When using DP multi-concurrency (AllGather mode), `num_inference_steps`
|
||||
> must be specified explicitly and must be the same for all concurrent requests.
|
||||
> `num_inference_steps=None` is not allowed because it may resolve differently
|
||||
> per request mode (e.g., T2V=35 vs action_mode=30), which would deadlock AllGather.
|
||||
|
||||
### Multiple concurrent requests (DP multi-concurrency)
|
||||
|
||||
With `--data-parallel-size 4`, send 4 concurrent requests to maximize throughput:
|
||||
|
||||
```bash
|
||||
for i in 1 2 3 4; do
|
||||
curl -s -o output_${i}.mp4 -X POST "http://localhost:8000/v1/videos/sync" \
|
||||
-H "Accept: video/mp4" \
|
||||
-F "model=/path/to/Cosmos3-Nano" \
|
||||
-F "prompt=A robot arm cleaning a plate, shot ${i}" \
|
||||
-F "negative_prompt=blurry" \
|
||||
-F "size=832x480" \
|
||||
-F "num_frames=29" \
|
||||
-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":false}' \
|
||||
-F "seed=17" &
|
||||
done
|
||||
wait
|
||||
```
|
||||
|
||||
## CLI flags
|
||||
|
||||
| Flag | Description | Default |
|
||||
|------|-------------|---------|
|
||||
| `--enable-distributed-layerwise-offload` | Enable DLO with H2D + AllGather overlap | `false` |
|
||||
| `--data-parallel-size N` | Number of DP ranks (weight sharding + concurrent requests) | `1` |
|
||||
| `--dlo-use-allgather` | Use shard + AllGather for weight reconstruction (recommended) | `true` |
|
||||
| `--dlo-no-use-allgather` | Each rank loads full weights independently (no sharding, no AllGather) | `false` |
|
||||
| `--usp N` | Use SP instead of DP for weight sharding (long sequences) | `1` |
|
||||
|
||||
## Mode comparison
|
||||
|
||||
| Mode | CLI flags | CPU/rank | HBM/card | Throughput | Use case |
|
||||
|------|----------|---------|---------|-----------|----------|
|
||||
| DLO + AllGather (DP4) | `--enable-distributed-layerwise-offload --data-parallel-size 4` | 1/4 model | 2 blocks | **3.3× HSDP** | Max throughput, short sequences |
|
||||
| DLO + AllGather (DP2) | `--enable-distributed-layerwise-offload --data-parallel-size 2` | 1/2 model | 2 blocks | 2× HSDP | Balanced throughput + memory |
|
||||
| DLO + AllGather (SP4) | `--enable-distributed-layerwise-offload --usp 4` | 1/4 model | 2 blocks | 1× (single req) | Long sequences (720p+) |
|
||||
| DLO no-AllGather | `--enable-distributed-layerwise-offload --dlo-no-use-allgather` | Full model | 2 blocks | 0.4× HSDP | No AllGather overhead, high CPU |
|
||||
| HSDP | `--use-hsdp --hsdp-shard-size 4` | 0 | 1/4 model | 1× (baseline) | Weights fit in HBM |
|
||||
|
||||
## Memory expectations
|
||||
|
||||
| Model | Config | Host RAM (cgroup) | HBM/card |
|
||||
|-------|--------|-------------------|----------|
|
||||
| Nano (33 GB) | DP2 + AllGather | ~38 GB | ~10 GB |
|
||||
| Nano (33 GB) | DP4 + AllGather | ~47 GB | ~10 GB |
|
||||
| Super (124 GB) | DP2 + AllGather | ~157 GB | ~15 GB |
|
||||
| Super (124 GB) | DP4 + AllGather | ~172 GB | ~10 GB |
|
||||
|
||||
> Host RAM is dominated by page cache (1× model_size, shared across ranks).
|
||||
> HBM only holds 2 transformer blocks + framework + activations, independent of model size.
|
||||
|
||||
## Limitations
|
||||
|
||||
- `--dlo-use-allgather` (default) requires all concurrent requests to have the same
|
||||
`num_inference_steps` (AllGather is a collective that needs all ranks synchronized)
|
||||
- Online quantization (FP8) is incompatible with mmap loading — falls back to regular
|
||||
`load_weights()` automatically
|
||||
- Tensor Parallel is not supported (DLO uses DP-based sharding, a different dimension)
|
||||
- HSDP + AllGather is rejected (would double-shard weights)
|
||||
@@ -17,7 +17,9 @@ The small Cosmos3 checkpoint, for constrained deployments. It runs on the shared
|
||||
|
||||
- **Transformer**: `Cosmos3EdgeVFMTransformer` (Nemotron-based), auto-selected from
|
||||
`backbone_type: cosmos3_edge_nemotron_dense` — nothing to pass.
|
||||
- **Defaults** differ; omit `--height` / `--width` to get them:
|
||||
- **Defaults** differ. The offline examples auto-detect Edge (from `edge` in the model id /
|
||||
transformer class) and apply the Edge column below, so you can omit `--height` / `--width` /
|
||||
`--guidance-scale` / `--flow-shift` and still get native Edge output:
|
||||
|
||||
| | Edge | Nano / Super |
|
||||
|---|---|---|
|
||||
@@ -26,6 +28,10 @@ The small Cosmos3 checkpoint, for constrained deployments. It runs on the shared
|
||||
| Video `guidance_scale` | 5.0 | 6.0 |
|
||||
| Video `flow_shift` | 3.0 | 10.0 |
|
||||
|
||||
> **Important:** these auto-defaults live in the example CLIs
|
||||
> (`text_to_video.py` / `image_to_video.py`). Do **not** pass the Nano/Super numbers
|
||||
> (720p / `guidance_scale` 6.0 / `flow_shift` 10.0) to Edge — they produce degenerate output.
|
||||
|
||||
For the larger checkpoints see [Cosmos3-Nano](./Cosmos3-Nano.md) and [Cosmos3-Super](./Cosmos3-Super.md).
|
||||
|
||||
## Support notes
|
||||
|
||||
@@ -125,3 +125,134 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \
|
||||
under the top-level `action` field. Verified on the 64B Super under
|
||||
`--cfg-parallel-size 2`: async `policy` returns the predicted action (`[16, 10]`)
|
||||
and the rollout video reliably.
|
||||
|
||||
## NPU
|
||||
|
||||
### 8× Ascend910 (A2, A3)
|
||||
|
||||
#### Environment
|
||||
|
||||
- OS: Linux
|
||||
- Python: 3.10+
|
||||
- Driver / runtime: Ascend NPU driver with CANN toolkit
|
||||
- Recommended operator library: **mindie-sd** (Ascend high-performance fused
|
||||
operators — enables `adalayernorm` and other fused kernels automatically upon
|
||||
installation)
|
||||
- vLLM version: Match the repository requirements for your checkout
|
||||
- vLLM-Omni version or commit: Use the commit you are deploying from
|
||||
|
||||
A pre-built Docker image is available on
|
||||
[Docker Hub](https://hub.docker.com/r/vllm/vllm-omni) and
|
||||
[Quay.io](https://quay.io/ascend/vllm-omni). Ensure the image tag matches your
|
||||
vLLM-Omni checkout so that NPU-specific code is in sync with the container.
|
||||
|
||||
#### Prerequisites
|
||||
|
||||
Install the **mindie-sd** operator library to enable Ascend-optimized fused
|
||||
operators (`adalayernorm`, etc.):
|
||||
|
||||
```bash
|
||||
git clone https://gitcode.com/Ascend/MindIE-SD.git && cd MindIE-SD
|
||||
|
||||
# Comment out the tik_ops build step (not needed for this use case)
|
||||
sed -i 's|^\(\s*\)source ${current_script_dir}/build_tik_ops.sh|\1# source ${current_script_dir}/build_tik_ops.sh|' build/build_ops.sh
|
||||
|
||||
python setup.py bdist_wheel
|
||||
cd dist
|
||||
pip install mindiesd-*.whl
|
||||
```
|
||||
|
||||
After installation, enable the Laser Attention kernel for significant
|
||||
long-sequence speedups:
|
||||
|
||||
```bash
|
||||
export MINDIE_SD_FA_TYPE=ascend_laser_attention
|
||||
```
|
||||
|
||||
#### Command
|
||||
|
||||
```bash
|
||||
export MINDIE_SD_FA_TYPE=ascend_laser_attention
|
||||
|
||||
vllm serve nvidia/Cosmos3-Super \
|
||||
--omni \
|
||||
--host 0.0.0.0 --port 8000 \
|
||||
--tensor-parallel-size 8 \
|
||||
--model-class-name Cosmos3OmniDiffusersPipeline \
|
||||
--no-guardrails \
|
||||
--init-timeout 1800
|
||||
```
|
||||
|
||||
#### Verification
|
||||
|
||||
Same requests as the GPU section above — all modes (T2I, T2V, I2V, V2V,
|
||||
T2VS, I2VS) work identically on NPU. Quick reference with
|
||||
`--no-guardrails`:
|
||||
|
||||
```bash
|
||||
curl http://localhost:8000/v1/models
|
||||
|
||||
# T2I (1024x1024, 50 steps)
|
||||
curl -sS -X POST http://localhost:8000/v1/images/generations \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "nvidia/Cosmos3-Super",
|
||||
"prompt": "A photorealistic red sports car on a city street at golden hour, cinematic lighting.",
|
||||
"size": "1024x1024", "n": 1, "response_format": "b64_json",
|
||||
"num_inference_steps": 50, "guidance_scale": 7.0, "seed": 42
|
||||
}' | python3 -c "import sys,json,base64; open('cosmos3_super_t2i.png','wb').write(base64.b64decode(json.load(sys.stdin)['data'][0]['b64_json']))"
|
||||
|
||||
# T2V (1280×720, 189 frames, 35 steps — official params)
|
||||
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":false}' \
|
||||
-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":false}' \
|
||||
-F "seed=1111" -F "input_reference=@/path/to/reference.jpg;type=image/jpeg" \
|
||||
-o cosmos3_super_i2v.mp4
|
||||
|
||||
# V2V — add an uploaded reference video. condition_video_keep can be "first" or "last".
|
||||
curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \
|
||||
-F "model=nvidia/Cosmos3-Super" -F "prompt=Continue the same scene 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={"condition_frame_indexes_vision":[0,1],"condition_video_keep":"first","guardrails":false}' \
|
||||
-F "seed=2222" -F "input_reference=@/path/to/reference.mp4;type=video/mp4" \
|
||||
-o cosmos3_super_v2v.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":false}' \
|
||||
-F "seed=17" -o cosmos3_super_t2vs.mp4
|
||||
|
||||
# I2V + sound — reference image with synchronized audio
|
||||
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 and ambient sound." \
|
||||
-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":false}' \
|
||||
-F "seed=1111" -F "input_reference=@/path/to/reference.jpg;type=image/jpeg" \
|
||||
-o cosmos3_super_i2vs.mp4
|
||||
```
|
||||
|
||||
#### Notes
|
||||
|
||||
- **Parallelism:** `--tensor-parallel-size 8` matches the model's 8 KV heads (`num_key_value_heads: 8`, `num_attention_heads: 64` → GQA). This uses 8 davinci devices (NPU 0–3, both cores each). The remaining NPU 4–7 stay idle.
|
||||
- **Model config:** 64 layers × 5120 hidden size × 64 attention heads, ~120 GB transformer on disk (27 shards). Each NPU device loads ~15 GB of model weights.
|
||||
- **Peak HBM:** ~22.7 GB per device at startup (bf16 weights). Additional memory is allocated per-generation for KV cache and activations — resolution and frame count drive the peak.
|
||||
- **Performance (verified on 8× Ascend910):** T2I 256² / 2 steps / guidance 1.0 → ~1.5 s.
|
||||
- **Guardrails** are disabled with `--no-guardrails` (guards are on by default). The gated `nvidia/Cosmos-1.0-Guardrail` model and `cosmos-guardrail` package are not shipped. Add `guardrails: false` in `extra_params` for per-request overrides when the server has guardrails enabled.
|
||||
- **Known limitations:** FP8 quantization not yet validated on Ascend NPU. `--enable-layerwise-offload` is available but untested on NPU.
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 1.1 MiB |
File renamed without changes.
Binary file not shown.
Binary file not shown.
Loaded 100 of 590 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user