Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
f5f78e46ff | ||
|
|
acde431668 | ||
|
|
f144f5effe | ||
|
|
a6b5519616 | ||
|
|
0f7c9c755d | ||
|
|
af9784c03f | ||
|
|
74eedd3ef1 | ||
|
|
a19098233b | ||
|
|
b313279af7 | ||
|
|
4d89ebaa07 | ||
|
|
617b1d6053 | ||
|
|
ca37f221a8 | ||
|
|
2b5a3dd6b3 | ||
|
|
7543f2ff3b | ||
|
|
6ac85d4607 | ||
|
|
4725593a4b | ||
|
|
c51b4fd7b8 | ||
|
|
2ee4a07adb | ||
|
|
1d0a97f0d1 | ||
|
|
2fdc647ed7 | ||
|
|
4027534dbf | ||
|
|
e449cc411e | ||
|
|
5d62fcd717 | ||
|
|
745ef312a7 | ||
|
|
7a560492a2 | ||
|
|
761eff930a | ||
|
|
112a9aa237 | ||
|
|
e79ad8d59a | ||
|
|
82f63a4403 | ||
|
|
7271116973 | ||
|
|
5628d41242 | ||
|
|
155856ff30 | ||
|
|
89b527e533 | ||
|
|
4db10d6dec | ||
|
|
29357027b0 | ||
|
|
d6871db6a3 | ||
|
|
fb717a4071 | ||
|
|
41964cb90a | ||
|
|
382bbd0a1f | ||
|
|
16248e97ef | ||
|
|
92e1a54419 | ||
|
|
a9f694b08a | ||
|
|
0e432e524c | ||
|
|
bb051ab1df | ||
|
|
c3a59f10f3 | ||
|
|
b255654a76 | ||
|
|
074af30a9b | ||
|
|
30191ed112 | ||
|
|
b7fcc9d07e | ||
|
|
411fe05bd4 | ||
|
|
92adcdeab4 | ||
|
|
846167b237 | ||
|
|
3bc89de0c0 | ||
|
|
32fd084ae2 | ||
|
|
90ec499bd5 | ||
|
|
fd51841230 | ||
|
|
128e6f4b97 | ||
|
|
8536dce759 | ||
|
|
c0b714a8ac | ||
|
|
e643aaec15 | ||
|
|
d5ddc190f9 | ||
|
|
ea317e4552 | ||
|
|
28c2200b47 | ||
|
|
83fe7d005b | ||
|
|
6a45efb287 | ||
|
|
070ea0ddc3 | ||
|
|
e674d60008 | ||
|
|
562136d390 | ||
|
|
fa944bd8e8 | ||
|
|
24c80070d6 | ||
|
|
58eb1c68db | ||
|
|
5493e75877 | ||
|
|
c316153b49 | ||
|
|
b63e7c7c9a | ||
|
|
1612948062 | ||
|
|
599045a94c | ||
|
|
11a5842bd1 | ||
|
|
46b7e0cf9f | ||
|
|
84338494d2 | ||
|
|
b84d73be4b | ||
|
|
0c6a10ebb6 | ||
|
|
29eecc2da8 | ||
|
|
35bd3ff7d2 | ||
|
|
4082202484 | ||
|
|
f9ec07f0e2 | ||
|
|
675c35db97 | ||
|
|
a3240fd484 | ||
|
|
742d845b8e | ||
|
|
b13761ed26 | ||
|
|
d275c0c993 | ||
|
|
e7f2c88701 | ||
|
|
3b2b5012eb | ||
|
|
90a48f4444 | ||
|
|
80c89e00cb | ||
|
|
a09b6499f8 | ||
|
|
22a51a7cdb | ||
|
|
a4b2e48718 | ||
|
|
f371f7bf05 | ||
|
|
f082f6eedb | ||
|
|
1ca198eaff | ||
|
|
e37a89ffcb | ||
|
|
68764cc4f4 | ||
|
|
fec0182ac2 |
No files matched your search
@@ -0,0 +1,38 @@
|
||||
steps:
|
||||
- label: ":docker: Build image"
|
||||
key: image-build
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "docker build --progress=plain --file docker/Dockerfile.ci -t vllm-omni-ci ."
|
||||
- "docker tag vllm-omni-ci public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
# L2 Test
|
||||
- label: "Upload Ready Pipeline"
|
||||
depends_on: image-build
|
||||
key: upload-ready-pipeline
|
||||
commands:
|
||||
- buildkite-agent pipeline upload .buildkite/test-ready.yml
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
# L3 Test
|
||||
- label: "Upload Merge Pipeline"
|
||||
depends_on: image-build
|
||||
key: upload-merge-pipeline
|
||||
commands:
|
||||
- buildkite-agent pipeline upload .buildkite/test-merge.yml
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
# L4 Test
|
||||
- label: "Upload Nightly Pipeline"
|
||||
depends_on: image-build
|
||||
key: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
commands:
|
||||
- buildkite-agent pipeline upload .buildkite/test-nightly.yml
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
@@ -1,5 +1,5 @@
|
||||
steps:
|
||||
- label: "Omni Model Test with H100"
|
||||
- label: ":full_moon: Omni Model Test with H100"
|
||||
timeout_in_minutes: 90
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
@@ -41,7 +41,7 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Omni Model Test"
|
||||
- label: ":full_moon: Omni Model Test"
|
||||
timeout_in_minutes: 60
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
@@ -54,20 +54,20 @@ steps:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
shm-size: "8gb"
|
||||
propagate-environment: true
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Omni Model Perf Test"
|
||||
timeout_in_minutes: 120
|
||||
- label: ":full_moon: Diffusion Model Test with H100"
|
||||
timeout_in_minutes: 60
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
# if: build.env("NIGHTLY") == "1"
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/perf/scripts/run_benchmark.py
|
||||
- pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and diffusion and H100" --run-level "advanced_model"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
@@ -96,3 +96,79 @@ steps:
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: Qwen3-TTS Non-Async-Chunk E2E Test"
|
||||
timeout_in_minutes: 30
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
commands:
|
||||
- |
|
||||
huggingface-cli download Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py::TestQwen3TTSNoAsyncChunk
|
||||
agents:
|
||||
queue: "gpu_4_queue"
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: ":full_moon: Omni Model Perf Test with H100"
|
||||
key: nightly-performance
|
||||
timeout_in_minutes: 180
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export BENCHMARK_DIR=tests
|
||||
- pytest -s -v tests/perf/scripts/run_benchmark.py
|
||||
- buildkite-agent artifact upload "tests/*.json"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 2
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":email: Nightly Perf Collection & Email"
|
||||
key: nightly-perf-distribution
|
||||
depends_on: nightly-performance
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
commands:
|
||||
- pip install openpyxl
|
||||
- export DEFAULT_INPUT_DIR=tests
|
||||
- export DEFAULT_OUTPUT_DIR=tests
|
||||
- buildkite-agent artifact download "tests/*.json" . --step nightly-performance
|
||||
- python tools/nightly/generate_nightly_perf_excel.py
|
||||
- python tools/nightly/send_nightly_perf_email.py
|
||||
- buildkite-agent artifact upload "tests/*.xlsx"
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
@@ -214,9 +214,10 @@ steps:
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
huggingface-cli download Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py
|
||||
pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py -k "not NoAsyncChunk"
|
||||
'
|
||||
agents:
|
||||
queue: "gpu_4_queue"
|
||||
|
||||
@@ -8,6 +8,7 @@ repos:
|
||||
# list of supported hooks: https://pre-commit.com/hooks.html
|
||||
- id: check-yaml
|
||||
args: ["--unsafe"]
|
||||
exclude: "examples/online_serving/chart-helm/templates/"
|
||||
- id: debug-statements
|
||||
- id: end-of-file-fixer
|
||||
- id: mixed-line-ending
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
<p align="center">
|
||||
<picture>
|
||||
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/logos/vllm-omni-logo.png">
|
||||
<source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/logos/vllm-omni-logo.png">
|
||||
<img alt="vllm-omni" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/logos/vllm-omni-logo.png" width=55%>
|
||||
</picture>
|
||||
</p>
|
||||
@@ -9,7 +9,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
|
||||
</h3>
|
||||
|
||||
<p align="center">
|
||||
| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://docs.google.com/presentation/d/1qv4qMW1rKAqDREMXiUDLIgqqHQe7TDPj/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true"><b>Slides</b></a> |
|
||||
| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://docs.google.com/presentation/d/1XJWgv79lORl8rbaVvp2d5Sqs6ZEBgAgj/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true"><b>Slides</b></a> |
|
||||
</p>
|
||||
|
||||
|
||||
@@ -17,6 +17,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
|
||||
|
||||
*Latest News* 🔥
|
||||
|
||||
- [2026/03] **[vllm-omni-skills](https://github.com/hsliuustc0106/vllm-omni-skills)** is a community-driven collection of AI assistant skills that help developers work with vLLM-Omni more effectively. These skills can be used with popular agentic AI coding assistants like **Cursor IDE**, **Claude**, **Codex**, and more.
|
||||
- [2026/02] We released [0.16.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.16.0) - A major alignment + capability release that rebases onto **upstream vLLM v0.16.0** and significantly expands performance, distributed execution, and production readiness across **Qwen3-Omni / Qwen3-TTS**, **Bagel**, **MiMo-Audio**, **GLM-Image** and the **Diffusion (DiT) image/video stack**—while also improving platform coverage (CUDA / ROCm / NPU / XPU), CI quality, and documentation.
|
||||
- [2026/02] We released [0.14.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.14.0) - This is the first **stable release** of vLLM-Omni that expands Omni’s diffusion / image-video generation and audio / TTS stack, improves distributed execution and memory efficiency, and broadens platform/backend coverage (GPU/ROCm/NPU/XPU). It also brings meaningful upgrades to serving APIs, profiling & benchmarking, and overall stability. Please check our latest [paper](https://arxiv.org/abs/2602.02204) for architecture design and performance results.
|
||||
- [2026/01] We released [0.12.0rc1](https://github.com/vllm-project/vllm-omni/releases/tag/v0.12.0rc1) - a major RC milestone focused on maturing the diffusion stack, strengthening OpenAI-compatible serving, expanding omni-model coverage, and improving stability across platforms (GPU/NPU/ROCm).
|
||||
|
||||
@@ -47,6 +47,7 @@ If no, check your shell running the ComfyUI process. There may be some error mes
|
||||
This extension offers the following nodes based on the output modalities (at **ComfyUI sidebar -> Node Library**):
|
||||
|
||||
- **Generate Image** for text-to-image and image-to-image tasks
|
||||
- **Generate Video** for text-to-video and image-to-video tasks
|
||||
- **Multimodality Understanding** for multimodality-to-text and multimodality-to-audio tasks
|
||||
- **TTS** and **TTS Voice Clone** for TTS tasks
|
||||
|
||||
@@ -59,7 +60,7 @@ To build a simple workflow yourself,
|
||||
|
||||
- Drag a generation node onto the canvas.
|
||||
- Depending on your need, grab built-in multimedia file loader nodes, such as **image->Load Image**, **image->video->Load Video**, **audio->Load Audio**
|
||||
- Depending on your need, grab built-in multimedia file preview nodes, such as **image->Preview Image**, **image->video->Save Video**, **audio->Preview Audio**. For text output, you can install [ComfyUI-Custom-Scripts plugin](https://github.com/pythongosssss/ComfyUI-Custom-Scripts/) and grab its **utils->Show Text 🐍** node.
|
||||
- Depending on your need, grab built-in multimedia file preview nodes, such as **image->Preview Image**, **image->video->Save Video**, **audio->Preview Audio**, **utils->Preview as Text**.
|
||||
- If you want to tune sampling parameters, grab corresponding nodes from **vLLM-Omni-> Sampling Params**.
|
||||
- For multi-stage models, you can connect multiple **AR Sampling Params** and **Diffusion Sampling Params** nodes to a **Multi-Stage Sampling Params List** node, and connect this node to the generation node.
|
||||
- For some multi-stage models like BAGEL, [only one stage's sampling parameters are exposed and tunable via vLLM-Omni's online serving API](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/examples/online_serving/bagel/). Thus, these models are treated as single-stage ones. Please check the vLLM-Omni documentation on how to correctly set each model's sampling parameters.
|
||||
@@ -104,6 +105,20 @@ You can configure per-stage sampling parameters for multi-stage models.
|
||||
> [!TIP]
|
||||
> The node automatically choose text-to-image or image-to-image API endpoints depending on whether you connect an image input or not.
|
||||
|
||||
### Text-to-video and image-to-video generation (e.g., Wan)
|
||||
|
||||
(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Video Generation**)
|
||||
|
||||
<p align="center">
|
||||
<picture>
|
||||
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-video-generation.jpg">
|
||||
<img alt="vLLM-Omni video generation" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-video-generation.jpg" width=55%>
|
||||
</picture>
|
||||
</p>
|
||||
|
||||
> [!TIP]
|
||||
> The node automatically choose text-to-video or image-to-video API endpoints depending on whether you connect an image input or not.
|
||||
|
||||
### TTS (e.g., Qwen TTS series)
|
||||
|
||||
(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni TTS**)
|
||||
|
||||
@@ -14,17 +14,21 @@ from .comfyui_vllm_omni.nodes import (
|
||||
VLLMOmniARSampling,
|
||||
VLLMOmniDiffusionSampling,
|
||||
VLLMOmniGenerateImage,
|
||||
VLLMOmniGenerateVideo,
|
||||
VLLMOmniQwenTTSParams,
|
||||
VLLMOmniRemoteLoRA,
|
||||
VLLMOmniSamplingParamsList,
|
||||
VLLMOmniTTS,
|
||||
VLLMOmniUnderstanding,
|
||||
VLLMOmniVoiceClone,
|
||||
VLLMOmniWanParams,
|
||||
)
|
||||
|
||||
# A dictionary that contains all nodes you want to export with their names
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
# === Generation ===
|
||||
"VLLMOmniGenerateImage": VLLMOmniGenerateImage,
|
||||
"VLLMOmniGenerateVideo": VLLMOmniGenerateVideo,
|
||||
"VLLMOmniUnderstanding": VLLMOmniUnderstanding,
|
||||
"VLLMOmniTTS": VLLMOmniTTS,
|
||||
"VLLMOmniVoiceClone": VLLMOmniVoiceClone,
|
||||
@@ -32,13 +36,16 @@ NODE_CLASS_MAPPINGS = {
|
||||
"VLLMOmniARSampling": VLLMOmniARSampling,
|
||||
"VLLMOmniDiffusionSampling": VLLMOmniDiffusionSampling,
|
||||
"VLLMOmniSamplingParamsList": VLLMOmniSamplingParamsList,
|
||||
"VLLMOmniRemoteLoRA": VLLMOmniRemoteLoRA,
|
||||
"VLLMOmniQwenTTSParams": VLLMOmniQwenTTSParams,
|
||||
"VLLMOmniWanParams": VLLMOmniWanParams,
|
||||
}
|
||||
|
||||
# A dictionary that contains the friendly/humanly readable titles for the nodes
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
# === Generation ===
|
||||
"VLLMOmniGenerateImage": "Generate Image",
|
||||
"VLLMOmniGenerateVideo": "Generate Video",
|
||||
"VLLMOmniUnderstanding": "Multimodality Understanding",
|
||||
"VLLMOmniTTS": "TTS (Text to Speech)",
|
||||
"VLLMOmniVoiceClone": "TTS Voice Cloning",
|
||||
@@ -46,7 +53,9 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"VLLMOmniARSampling": "AR Sampling Params",
|
||||
"VLLMOmniDiffusionSampling": "Diffusion Sampling Params",
|
||||
"VLLMOmniSamplingParamsList": "Multi-Stage Sampling Params List",
|
||||
"VLLMOmniRemoteLoRA": "LoRA",
|
||||
"VLLMOmniQwenTTSParams": "Qwen TTS Params",
|
||||
"VLLMOmniWanParams": "Wan Video Params",
|
||||
}
|
||||
|
||||
WEB_DIRECTORY = "./web"
|
||||
@@ -11,6 +11,7 @@ from .utils.types import (
|
||||
AutoregressionSamplingParams,
|
||||
DiffusionSamplingParams,
|
||||
QwenTTSModelSpecificParams,
|
||||
WanModelSpecificParams,
|
||||
)
|
||||
from .utils.validators import (
|
||||
add_sampling_parameters_to_stage,
|
||||
@@ -56,6 +57,7 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
|
||||
# "video": ("VIDEO",),
|
||||
# "audio": ("AUDIO",),
|
||||
"sampling_params": ("SAMPLING_PARAMS",),
|
||||
"lora": ("REMOTE_LORA",),
|
||||
},
|
||||
}
|
||||
|
||||
@@ -76,9 +78,11 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
|
||||
audio: AudioInput | None = None, # Hidden & unused
|
||||
video: VideoInput | None = None, # Hidden & unused
|
||||
sampling_params: dict | list[dict] | None = None,
|
||||
lora: dict | None = None,
|
||||
**kwargs,
|
||||
):
|
||||
logger.info("Uncaught kwargs: %s", kwargs)
|
||||
if kwargs:
|
||||
logger.info("Uncaught kwargs: %s", kwargs)
|
||||
logger.debug("Got sampling params: %s", sampling_params)
|
||||
validate_model_and_sampling_params_types(model, sampling_params)
|
||||
if image is None and mask is not None:
|
||||
@@ -105,6 +109,7 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
|
||||
height=height,
|
||||
negative_prompt=negative_prompt,
|
||||
sampling_params=sampling_params,
|
||||
lora=lora,
|
||||
)
|
||||
return (output,)
|
||||
elif image is not None and audio is None and video is None:
|
||||
@@ -119,6 +124,7 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
|
||||
negative_prompt=negative_prompt,
|
||||
mask=mask,
|
||||
sampling_params=sampling_params,
|
||||
lora=lora,
|
||||
)
|
||||
return (output,)
|
||||
|
||||
@@ -136,11 +142,91 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
|
||||
audio=audio,
|
||||
video=video,
|
||||
sampling_params=sampling_params,
|
||||
lora=lora,
|
||||
)
|
||||
|
||||
return (output,)
|
||||
|
||||
|
||||
class VLLMOmniGenerateVideo(_VLLMOmniGenerateBase):
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"url": ("STRING", {"default": "http://localhost:8000/v1"}),
|
||||
"model": ("STRING", {"default": "Wan-AI/Wan2.2-T2V-A14B-Diffusers"}),
|
||||
"prompt": ("STRING", {"multiline": True}),
|
||||
"negative_prompt": ("STRING", {"multiline": True, "default": ""}),
|
||||
"width": ("INT", {"default": 832, "min": 1}),
|
||||
"height": ("INT", {"default": 480, "min": 1}),
|
||||
"fps": ("INT", {"default": 16, "min": 1}),
|
||||
"num_frames": ("INT", {"default": 41, "min": 1}),
|
||||
},
|
||||
"optional": {
|
||||
"image": ("IMAGE",),
|
||||
"sampling_params": ("SAMPLING_PARAMS",),
|
||||
"lora": ("REMOTE_LORA",),
|
||||
"model_params": ("VIDEO_PARAMS",),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("VIDEO",)
|
||||
RETURN_NAMES = ("video",)
|
||||
FUNCTION = "generate"
|
||||
|
||||
async def generate(
|
||||
self,
|
||||
url: str,
|
||||
model: str,
|
||||
prompt: str,
|
||||
width: int,
|
||||
height: int,
|
||||
fps: int,
|
||||
num_frames: int,
|
||||
negative_prompt: str | None = None,
|
||||
image: torch.Tensor | None = None,
|
||||
sampling_params: dict | list[dict] | None = None,
|
||||
model_params: dict | None = None,
|
||||
lora: dict | None = None,
|
||||
**kwargs,
|
||||
):
|
||||
if kwargs:
|
||||
logger.info("Uncaught kwargs: %s", kwargs)
|
||||
logger.debug("Got sampling params: %s", sampling_params)
|
||||
logger.debug("Got model params: %s", model_params)
|
||||
validate_model_and_sampling_params_types(model, sampling_params)
|
||||
|
||||
# Currently, all video generation models are single-stage diffusion models
|
||||
if isinstance(sampling_params, list):
|
||||
if len(sampling_params) != 1:
|
||||
raise ValueError(
|
||||
"Video generation expects a single sampling params group. "
|
||||
"Please provide one Diffusion sampling node."
|
||||
)
|
||||
sampling_params = sampling_params[0]
|
||||
|
||||
if sampling_params is not None:
|
||||
sampling_params.pop("type", None) # internal fields
|
||||
if model_params is not None:
|
||||
model_params.pop("type", None) # internal fields
|
||||
|
||||
client = VLLMOmniClient(url)
|
||||
output = await client.generate_video(
|
||||
model=model,
|
||||
prompt=prompt,
|
||||
image=image, # image present => i2v, absent => t2v
|
||||
width=width,
|
||||
height=height,
|
||||
num_frames=num_frames,
|
||||
fps=fps,
|
||||
negative_prompt=negative_prompt,
|
||||
sampling_params=sampling_params,
|
||||
lora=lora,
|
||||
model_params=model_params,
|
||||
)
|
||||
return (output,)
|
||||
|
||||
|
||||
class VLLMOmniUnderstanding(_VLLMOmniGenerateBase):
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
@@ -188,7 +274,8 @@ class VLLMOmniUnderstanding(_VLLMOmniGenerateBase):
|
||||
use_audio_in_video: bool = True,
|
||||
**kwargs,
|
||||
) -> tuple[str, AudioInput]:
|
||||
logger.info("Uncaught kwargs: %s", kwargs)
|
||||
if kwargs:
|
||||
logger.info("Uncaught kwargs: %s", kwargs)
|
||||
logger.debug("Got sampling params: %s", sampling_params)
|
||||
validate_model_and_sampling_params_types(model, sampling_params)
|
||||
|
||||
@@ -547,6 +634,52 @@ class VLLMOmniSamplingParamsList:
|
||||
return (params,)
|
||||
|
||||
|
||||
class VLLMOmniRemoteLoRA:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"local_path": ("STRING", {"default": ""}),
|
||||
"name": ("STRING", {"default": ""}),
|
||||
"scale": (
|
||||
"FLOAT",
|
||||
{"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.1},
|
||||
),
|
||||
"int_id": (
|
||||
"INT",
|
||||
{
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"step": 1,
|
||||
"tooltip": "0 means it is not set and the server can derive it.",
|
||||
},
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("REMOTE_LORA",)
|
||||
RETURN_NAMES = ("lora",)
|
||||
FUNCTION = "get_lora"
|
||||
CATEGORY = "vLLM-Omni"
|
||||
|
||||
@classmethod
|
||||
def VALIDATE_INPUTS(cls, local_path, name) -> str | Literal[True]:
|
||||
if not local_path.strip() or not name.strip():
|
||||
return "Both local_path and name must be provided."
|
||||
return True
|
||||
|
||||
def get_lora(self, local_path: str, name: str, scale: float, int_id: int):
|
||||
local_path = local_path.strip()
|
||||
name = name.strip()
|
||||
lora = {
|
||||
"local_path": local_path or None,
|
||||
"name": name or None,
|
||||
"scale": float(scale),
|
||||
"int_id": int(int_id) if int_id > 0 else None,
|
||||
}
|
||||
return (lora,)
|
||||
|
||||
|
||||
class VLLMOmniQwenTTSParams:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
@@ -572,3 +705,32 @@ class VLLMOmniQwenTTSParams:
|
||||
|
||||
def get_params(self, **kwargs):
|
||||
return (QwenTTSModelSpecificParams(kwargs),)
|
||||
|
||||
|
||||
class VLLMOmniWanParams:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"guidance_scale_2": (
|
||||
"FLOAT",
|
||||
{"default": 4.0, "min": 0.0, "max": 20.0, "step": 0.1},
|
||||
),
|
||||
"boundary_ratio": (
|
||||
"FLOAT",
|
||||
{"default": 0.875, "min": 0.0, "max": 1.0, "step": 0.001},
|
||||
),
|
||||
"flow_shift": (
|
||||
"FLOAT",
|
||||
{"default": 5.0, "min": 0.0, "max": 100.0, "step": 0.1},
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("VIDEO_PARAMS",)
|
||||
RETURN_NAMES = ("Wan video params",)
|
||||
FUNCTION = "get_params"
|
||||
CATEGORY = "vLLM-Omni/Video Params"
|
||||
|
||||
def get_params(self, **kwargs):
|
||||
return (WanModelSpecificParams(kwargs),)
|
||||
@@ -6,6 +6,7 @@ The image generation part is derived from dougbtv/comfyui-vllm-omni by Doug (@do
|
||||
Original source at https://github.com/dougbtv/comfyui-vllm-omni, distributed under the MIT License.
|
||||
"""
|
||||
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
import aiohttp
|
||||
@@ -17,6 +18,7 @@ from .format import (
|
||||
audio_to_base64,
|
||||
base64_to_audio,
|
||||
base64_to_image_tensor,
|
||||
base64_to_video,
|
||||
bytes_to_audio,
|
||||
image_tensor_to_base64,
|
||||
image_tensor_to_png_bytes,
|
||||
@@ -30,7 +32,7 @@ logger = get_logger(__name__)
|
||||
|
||||
|
||||
class VLLMOmniClient:
|
||||
def __init__(self, base_url: str, timeout: float = 300.0):
|
||||
def __init__(self, base_url: str, timeout: float | None = None):
|
||||
self.base_url = base_url
|
||||
self.timeout = aiohttp.ClientTimeout(total=timeout)
|
||||
|
||||
@@ -43,12 +45,13 @@ class VLLMOmniClient:
|
||||
height: int,
|
||||
negative_prompt: str | None = None,
|
||||
sampling_params: dict | None = None,
|
||||
lora: dict | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Run text-to-image generation via DALLE API"""
|
||||
await self._check_model_exist(model)
|
||||
|
||||
size = f"{width}x{height}"
|
||||
payload = {
|
||||
payload: dict[str, Any] = {
|
||||
"model": model,
|
||||
"prompt": prompt,
|
||||
"size": size,
|
||||
@@ -58,6 +61,8 @@ class VLLMOmniClient:
|
||||
payload["negative_prompt"] = negative_prompt
|
||||
if sampling_params is not None:
|
||||
payload.update(sampling_params)
|
||||
if lora is not None:
|
||||
payload["lora"] = lora
|
||||
logger.debug("img gen payload: %s", payload)
|
||||
|
||||
url = self.base_url + "/images/generations"
|
||||
@@ -110,6 +115,7 @@ class VLLMOmniClient:
|
||||
negative_prompt: str | None = None,
|
||||
mask: torch.Tensor | None = None,
|
||||
sampling_params: dict | None = None,
|
||||
lora: dict | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Run image editing via DALLE API"""
|
||||
await self._check_model_exist(model)
|
||||
@@ -131,6 +137,8 @@ class VLLMOmniClient:
|
||||
if sampling_params is not None:
|
||||
for k, v in sampling_params.items():
|
||||
form.add_field(k, str(v))
|
||||
if lora is not None:
|
||||
form.add_field("lora", json.dumps(lora, ensure_ascii=False))
|
||||
if mask is not None:
|
||||
mask_filename = "mask.png"
|
||||
form.add_field(
|
||||
@@ -183,6 +191,7 @@ class VLLMOmniClient:
|
||||
audio: AudioInput | None = None,
|
||||
video: VideoInput | None = None,
|
||||
sampling_params: dict | list[dict] | None = None,
|
||||
lora: dict | None = None,
|
||||
) -> torch.Tensor:
|
||||
payload = VLLMOmniClient._prepare_chat_completion_messages(
|
||||
model=model,
|
||||
@@ -193,6 +202,8 @@ class VLLMOmniClient:
|
||||
video=video,
|
||||
sampling_params=sampling_params,
|
||||
modalities=["image"],
|
||||
# === below are additional `extra_body` fields, handled by **kwargs ===
|
||||
lora=lora,
|
||||
)
|
||||
choices = await self._generate_base_chat_completion(model, payload)
|
||||
|
||||
@@ -206,6 +217,79 @@ class VLLMOmniClient:
|
||||
|
||||
return torch.stack(image_tensors, dim=0)
|
||||
|
||||
async def generate_video(
|
||||
self,
|
||||
*,
|
||||
model: str,
|
||||
prompt: str,
|
||||
width: int,
|
||||
height: int,
|
||||
num_frames: int,
|
||||
fps: int,
|
||||
negative_prompt: str | None = None,
|
||||
image: torch.Tensor | None = None,
|
||||
sampling_params: dict | None = None,
|
||||
model_params: dict | None = None,
|
||||
lora: dict | None = None,
|
||||
**extra_body,
|
||||
) -> VideoInput:
|
||||
form = aiohttp.FormData()
|
||||
form.add_field("model", model)
|
||||
form.add_field("prompt", prompt)
|
||||
form.add_field("width", str(width))
|
||||
form.add_field("height", str(height))
|
||||
form.add_field("num_frames", str(num_frames))
|
||||
form.add_field("fps", str(fps))
|
||||
if negative_prompt:
|
||||
form.add_field("negative_prompt", negative_prompt)
|
||||
if sampling_params is not None:
|
||||
for k, v in sampling_params.items():
|
||||
form.add_field(k, str(v))
|
||||
if model_params is not None:
|
||||
for k, v in model_params.items():
|
||||
form.add_field(k, str(v))
|
||||
if lora is not None:
|
||||
form.add_field("lora", json.dumps(lora, ensure_ascii=False))
|
||||
if extra_body:
|
||||
form.add_field("extra_body", json.dumps(extra_body, ensure_ascii=False))
|
||||
|
||||
if image is not None:
|
||||
image_filename = "image.png" # Required for multipart form
|
||||
form.add_field(
|
||||
"input_reference",
|
||||
image_tensor_to_png_bytes(image, image_filename),
|
||||
filename=image_filename,
|
||||
content_type="image/png",
|
||||
)
|
||||
|
||||
url = self.base_url + "/videos"
|
||||
async with aiohttp.ClientSession(timeout=self.timeout) as session:
|
||||
try:
|
||||
async with session.post(url, data=form) as response:
|
||||
if not response.ok:
|
||||
error_text = await response.text()
|
||||
raise (ValueError if response.status < 500 else RuntimeError)(
|
||||
f"vLLM-Omni API returned status {response.status}: {error_text}"
|
||||
)
|
||||
|
||||
try:
|
||||
data = await response.json()
|
||||
except aiohttp.ContentTypeError as e:
|
||||
raise RuntimeError(f"Invalid JSON response from vLLM-Omni: {e}")
|
||||
except aiohttp.ClientError as e:
|
||||
raise RuntimeError(f"Network error connecting to vLLM-Omni at {url}: {e}")
|
||||
|
||||
if "data" not in data:
|
||||
raise RuntimeError("API response missing 'data' field - expected OpenAI DALL-E format")
|
||||
if not data["data"]:
|
||||
raise RuntimeError("API returned empty data array")
|
||||
try:
|
||||
base64_str = data["data"][0]["b64_json"]
|
||||
except (KeyError, IndexError):
|
||||
raise RuntimeError("API response missing 'b64_json' field in first data item")
|
||||
|
||||
return base64_to_video(base64_str)
|
||||
|
||||
async def generate_understanding_chat_completion(
|
||||
self,
|
||||
*,
|
||||
|
||||
@@ -6,12 +6,17 @@ Original source at https://github.com/dougbtv/comfyui-vllm-omni, distributed und
|
||||
|
||||
import base64
|
||||
import mimetypes
|
||||
from fractions import Fraction
|
||||
from io import BytesIO
|
||||
|
||||
import av
|
||||
import numpy as np
|
||||
import torch
|
||||
from av.audio.frame import AudioFrame
|
||||
from av.audio.resampler import AudioResampler
|
||||
from av.video.frame import VideoFrame
|
||||
from comfy_api.input import AudioInput, VideoInput
|
||||
from comfy_api.latest import InputImpl, Types
|
||||
from comfy_extras import nodes_audio
|
||||
from PIL import Image
|
||||
|
||||
@@ -138,6 +143,71 @@ def video_to_base64(video: VideoInput, filename: str = "video.mp4") -> str:
|
||||
return f"data:{mime_type};base64,{base64_str}"
|
||||
|
||||
|
||||
def base64_to_video(base64_str: str) -> VideoInput:
|
||||
if base64_str.startswith("data:video"):
|
||||
_, base64_str = base64_str.split(",", 1)
|
||||
|
||||
try:
|
||||
video_bytes = base64.b64decode(base64_str)
|
||||
except Exception as e:
|
||||
raise ValueError(f"Invalid base64 string: {e}")
|
||||
|
||||
video_buffer = BytesIO(video_bytes)
|
||||
|
||||
try:
|
||||
with av.open(video_buffer, mode="r") as container:
|
||||
video_stream = next((s for s in container.streams if s.type == "video"), None)
|
||||
if video_stream is None:
|
||||
raise ValueError("No video stream found in decoded payload.")
|
||||
|
||||
frames = []
|
||||
for frame in container.decode(video_stream):
|
||||
if not isinstance(frame, VideoFrame):
|
||||
continue
|
||||
image = frame.to_ndarray(format="rgb24")
|
||||
frames.append(torch.from_numpy(image).float() / 255.0)
|
||||
|
||||
if len(frames) == 0:
|
||||
raise ValueError("No video frames found in decoded payload.")
|
||||
|
||||
images = torch.stack(frames, dim=0)
|
||||
frame_rate = Fraction(video_stream.average_rate) if video_stream.average_rate else Fraction(1)
|
||||
|
||||
audio: AudioInput | None = None
|
||||
if len(container.streams.audio):
|
||||
audio_stream = container.streams.audio[-1]
|
||||
audio_frames = []
|
||||
resampler = AudioResampler(format="fltp")
|
||||
for frame in container.decode(audio_stream):
|
||||
if not isinstance(frame, AudioFrame):
|
||||
continue
|
||||
resampled = resampler.resample(frame)
|
||||
if not isinstance(resampled, list):
|
||||
resampled = [resampled]
|
||||
for audio_frame in resampled:
|
||||
if audio_frame is not None:
|
||||
audio_frames.append(audio_frame.to_ndarray())
|
||||
|
||||
if len(audio_frames) > 0:
|
||||
audio_data = np.concatenate(audio_frames, axis=1)
|
||||
sample_rate = int(audio_stream.sample_rate) if audio_stream.sample_rate else 1
|
||||
audio = {
|
||||
"waveform": torch.from_numpy(audio_data).unsqueeze(0),
|
||||
"sample_rate": sample_rate,
|
||||
}
|
||||
|
||||
components = Types.VideoComponents(
|
||||
images=images,
|
||||
frame_rate=frame_rate,
|
||||
audio=audio,
|
||||
metadata=container.metadata if container.metadata else None,
|
||||
)
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to decode video: {e}")
|
||||
|
||||
return InputImpl.VideoFromComponents(components)
|
||||
|
||||
|
||||
def audio_to_bytes(audio: AudioInput, filename: str = "audio.mp3", quality: str = "128k") -> BytesIO:
|
||||
waveform = audio["waveform"][0] # Shape: (C, T)
|
||||
sample_rate = audio["sample_rate"]
|
||||
|
||||
@@ -23,6 +23,10 @@ class QwenTTSModelSpecificParams(dict):
|
||||
pass
|
||||
|
||||
|
||||
class WanModelSpecificParams(dict):
|
||||
pass
|
||||
|
||||
|
||||
class ModelMode(Enum):
|
||||
IMAGE_GENERATION = auto()
|
||||
VIDEO_GENERATION = auto()
|
||||
|
||||
@@ -16,7 +16,7 @@ def validate_model_and_sampling_params_types(
|
||||
# Skip if no spec or no sampling params
|
||||
pipeline_spec, _ = lookup_model_spec(model_name)
|
||||
if pipeline_spec is None:
|
||||
logger.info("skipping sampling params check because spec is not found")
|
||||
logger.info(f"skipping sampling params check because spec for {model_name} is not found")
|
||||
return
|
||||
if sampling_param_list is None:
|
||||
return
|
||||
|
||||
|
After Width: | Height: | Size: 129 KiB |
@@ -0,0 +1,513 @@
|
||||
{
|
||||
"id": "3aea1d08-7de5-4997-ba71-b2f66ad72a21",
|
||||
"revision": 0,
|
||||
"last_node_id": 14,
|
||||
"last_link_id": 28,
|
||||
"nodes": [
|
||||
{
|
||||
"id": 10,
|
||||
"type": "MarkdownNote",
|
||||
"pos": [
|
||||
227.99306819575278,
|
||||
-231.24143306069843
|
||||
],
|
||||
"size": [
|
||||
240.3326478922474,
|
||||
136.3505820791881
|
||||
],
|
||||
"flags": {},
|
||||
"order": 0,
|
||||
"mode": 0,
|
||||
"inputs": [],
|
||||
"outputs": [],
|
||||
"title": "Note: Task and Input",
|
||||
"properties": {},
|
||||
"widgets_values": [
|
||||
"vLLM-Omni nodes are categorized based on the output modality. The \"Generate Video\" node supports both text-to-video generation and image-to-video generation. The node will create corresponding payload depending on whether an input image is present or not."
|
||||
],
|
||||
"color": "#432",
|
||||
"bgcolor": "#000"
|
||||
},
|
||||
{
|
||||
"id": 4,
|
||||
"type": "LoadImage",
|
||||
"pos": [
|
||||
496.31859627609606,
|
||||
-229.71277089860084
|
||||
],
|
||||
"size": [
|
||||
270,
|
||||
314
|
||||
],
|
||||
"flags": {},
|
||||
"order": 1,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"localized_name": "image",
|
||||
"name": "image",
|
||||
"type": "COMBO",
|
||||
"widget": {
|
||||
"name": "image"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "choose file to upload",
|
||||
"name": "upload",
|
||||
"type": "IMAGEUPLOAD",
|
||||
"widget": {
|
||||
"name": "upload"
|
||||
},
|
||||
"link": null
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"localized_name": "IMAGE",
|
||||
"name": "IMAGE",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
25
|
||||
]
|
||||
},
|
||||
{
|
||||
"localized_name": "MASK",
|
||||
"name": "MASK",
|
||||
"type": "MASK",
|
||||
"links": null
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "LoadImage"
|
||||
},
|
||||
"widgets_values": [
|
||||
"cute-cat.jpg",
|
||||
"image"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 8,
|
||||
"type": "VLLMOmniDiffusionSampling",
|
||||
"pos": [
|
||||
465.5140218220927,
|
||||
148.80072646828987
|
||||
],
|
||||
"size": [
|
||||
284.205078125,
|
||||
226
|
||||
],
|
||||
"flags": {},
|
||||
"order": 2,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"localized_name": "n",
|
||||
"name": "n",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "n"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "num_inference_steps",
|
||||
"name": "num_inference_steps",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "num_inference_steps"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "guidance_scale",
|
||||
"name": "guidance_scale",
|
||||
"type": "FLOAT",
|
||||
"widget": {
|
||||
"name": "guidance_scale"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "true_cfg_scale",
|
||||
"name": "true_cfg_scale",
|
||||
"type": "FLOAT",
|
||||
"widget": {
|
||||
"name": "true_cfg_scale"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "vae_use_slicing",
|
||||
"name": "vae_use_slicing",
|
||||
"type": "BOOLEAN",
|
||||
"widget": {
|
||||
"name": "vae_use_slicing"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "vae_use_tiling",
|
||||
"name": "vae_use_tiling",
|
||||
"type": "BOOLEAN",
|
||||
"widget": {
|
||||
"name": "vae_use_tiling"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "seed",
|
||||
"name": "seed",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "seed"
|
||||
},
|
||||
"link": null
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"localized_name": "diffusion sampling params",
|
||||
"name": "diffusion sampling params",
|
||||
"type": "SAMPLING_PARAMS",
|
||||
"links": [
|
||||
27
|
||||
]
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "VLLMOmniDiffusionSampling"
|
||||
},
|
||||
"widgets_values": [
|
||||
1,
|
||||
50,
|
||||
1,
|
||||
1,
|
||||
false,
|
||||
false,
|
||||
1992,
|
||||
"randomize"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 14,
|
||||
"type": "VLLMOmniWanParams",
|
||||
"pos": [
|
||||
518.6005521845747,
|
||||
428.4455663719664
|
||||
],
|
||||
"size": [
|
||||
270,
|
||||
106
|
||||
],
|
||||
"flags": {},
|
||||
"order": 3,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"localized_name": "guidance_scale_2",
|
||||
"name": "guidance_scale_2",
|
||||
"type": "FLOAT",
|
||||
"widget": {
|
||||
"name": "guidance_scale_2"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "boundary_ratio",
|
||||
"name": "boundary_ratio",
|
||||
"type": "FLOAT",
|
||||
"widget": {
|
||||
"name": "boundary_ratio"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "flow_shift",
|
||||
"name": "flow_shift",
|
||||
"type": "FLOAT",
|
||||
"widget": {
|
||||
"name": "flow_shift"
|
||||
},
|
||||
"link": null
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"localized_name": "Wan video params",
|
||||
"name": "Wan video params",
|
||||
"type": "VIDEO_PARAMS",
|
||||
"links": [
|
||||
28
|
||||
]
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "VLLMOmniWanParams"
|
||||
},
|
||||
"widgets_values": [
|
||||
4,
|
||||
0.875,
|
||||
5
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 12,
|
||||
"type": "VLLMOmniGenerateVideo",
|
||||
"pos": [
|
||||
827.2566336087867,
|
||||
-73.77449831827545
|
||||
],
|
||||
"size": [
|
||||
400,
|
||||
346
|
||||
],
|
||||
"flags": {},
|
||||
"order": 4,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"localized_name": "image",
|
||||
"name": "image",
|
||||
"shape": 7,
|
||||
"type": "IMAGE",
|
||||
"link": 25
|
||||
},
|
||||
{
|
||||
"localized_name": "sampling_params",
|
||||
"name": "sampling_params",
|
||||
"shape": 7,
|
||||
"type": "SAMPLING_PARAMS",
|
||||
"link": 27
|
||||
},
|
||||
{
|
||||
"localized_name": "lora",
|
||||
"name": "lora",
|
||||
"shape": 7,
|
||||
"type": "REMOTE_LORA",
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "model_params",
|
||||
"name": "model_params",
|
||||
"shape": 7,
|
||||
"type": "VIDEO_PARAMS",
|
||||
"link": 28
|
||||
},
|
||||
{
|
||||
"localized_name": "url",
|
||||
"name": "url",
|
||||
"type": "STRING",
|
||||
"widget": {
|
||||
"name": "url"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "model",
|
||||
"name": "model",
|
||||
"type": "STRING",
|
||||
"widget": {
|
||||
"name": "model"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "prompt",
|
||||
"name": "prompt",
|
||||
"type": "STRING",
|
||||
"widget": {
|
||||
"name": "prompt"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "negative_prompt",
|
||||
"name": "negative_prompt",
|
||||
"type": "STRING",
|
||||
"widget": {
|
||||
"name": "negative_prompt"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "width",
|
||||
"name": "width",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "width"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "height",
|
||||
"name": "height",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "height"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "fps",
|
||||
"name": "fps",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "fps"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "num_frames",
|
||||
"name": "num_frames",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "num_frames"
|
||||
},
|
||||
"link": null
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"localized_name": "video",
|
||||
"name": "video",
|
||||
"type": "VIDEO",
|
||||
"links": [
|
||||
26
|
||||
]
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "VLLMOmniGenerateVideo"
|
||||
},
|
||||
"widgets_values": [
|
||||
"http://localhost:8000/v1",
|
||||
"Wan-AI/Wan2.2-TI2V-5B-Diffusers",
|
||||
"Make the cat in the attached video yawning. Eye closed and mouth wide open. Keep the original realistic style.",
|
||||
"Cartoonish",
|
||||
900,
|
||||
1200,
|
||||
12,
|
||||
41
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 13,
|
||||
"type": "SaveVideo",
|
||||
"pos": [
|
||||
1257.5441369501657,
|
||||
-44.12956394419863
|
||||
],
|
||||
"size": [
|
||||
374,
|
||||
372.34896826553796
|
||||
],
|
||||
"flags": {},
|
||||
"order": 5,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"localized_name": "video",
|
||||
"name": "video",
|
||||
"type": "VIDEO",
|
||||
"link": 26
|
||||
},
|
||||
{
|
||||
"localized_name": "filename_prefix",
|
||||
"name": "filename_prefix",
|
||||
"type": "STRING",
|
||||
"widget": {
|
||||
"name": "filename_prefix"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "format",
|
||||
"name": "format",
|
||||
"type": "COMBO",
|
||||
"widget": {
|
||||
"name": "format"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "codec",
|
||||
"name": "codec",
|
||||
"type": "COMBO",
|
||||
"widget": {
|
||||
"name": "codec"
|
||||
},
|
||||
"link": null
|
||||
}
|
||||
],
|
||||
"outputs": [],
|
||||
"properties": {},
|
||||
"widgets_values": [
|
||||
"video/ComfyUI",
|
||||
"auto",
|
||||
"auto"
|
||||
]
|
||||
}
|
||||
],
|
||||
"links": [
|
||||
[
|
||||
25,
|
||||
4,
|
||||
0,
|
||||
12,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
26,
|
||||
12,
|
||||
0,
|
||||
13,
|
||||
0,
|
||||
"VIDEO"
|
||||
],
|
||||
[
|
||||
27,
|
||||
8,
|
||||
0,
|
||||
12,
|
||||
1,
|
||||
"SAMPLING_PARAMS"
|
||||
],
|
||||
[
|
||||
28,
|
||||
14,
|
||||
0,
|
||||
12,
|
||||
3,
|
||||
"VIDEO_PARAMS"
|
||||
]
|
||||
],
|
||||
"groups": [
|
||||
{
|
||||
"id": 1,
|
||||
"title": "Input",
|
||||
"bounding": [
|
||||
213.1706010087147,
|
||||
-313.20095750667554,
|
||||
560.8246471437027,
|
||||
407.4246532472182
|
||||
],
|
||||
"color": "#3f789e",
|
||||
"font_size": 24,
|
||||
"flags": {}
|
||||
}
|
||||
],
|
||||
"config": {},
|
||||
"extra": {
|
||||
"workflowRendererVersion": "LG",
|
||||
"ds": {
|
||||
"scale": 1.1469075819486894,
|
||||
"offset": [
|
||||
101.7632950847089,
|
||||
492.29121889347806
|
||||
]
|
||||
}
|
||||
},
|
||||
"version": 0.4
|
||||
}
|
||||
@@ -0,0 +1,100 @@
|
||||
# Qwen3-TTS Benchmark
|
||||
|
||||
Benchmarks for Qwen3-TTS text-to-speech models, comparing vLLM-Omni streaming serving against HuggingFace Transformers offline inference.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
```bash
|
||||
pip install matplotlib aiohttp soundfile numpy tqdm
|
||||
pip install qwen_tts # for HF baseline
|
||||
```
|
||||
|
||||
## Quick Start
|
||||
|
||||
Run the full benchmark (vllm-omni + HF baseline) with a single command:
|
||||
|
||||
```bash
|
||||
cd benchmarks/qwen3-tts
|
||||
bash run_benchmark.sh
|
||||
```
|
||||
|
||||
Results (JSON + PNG plots) are saved to `results/`.
|
||||
|
||||
### Common options
|
||||
|
||||
```bash
|
||||
# Only vllm-omni (skip HF baseline)
|
||||
bash run_benchmark.sh --async-only
|
||||
|
||||
# Only HF baseline
|
||||
bash run_benchmark.sh --hf-only
|
||||
|
||||
# Use a different model (e.g. 1.7B)
|
||||
MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice bash run_benchmark.sh --async-only
|
||||
|
||||
# Use batch_size=4 config for higher throughput
|
||||
STAGE_CONFIG=vllm_omni/configs/qwen3_tts_bs4.yaml bash run_benchmark.sh --async-only
|
||||
|
||||
# Custom GPU, prompt count, concurrency levels
|
||||
GPU_DEVICE=1 NUM_PROMPTS=20 CONCURRENCY="1 4" bash run_benchmark.sh
|
||||
```
|
||||
|
||||
## Manual Steps
|
||||
|
||||
### 1) Start the vLLM-Omni server
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=0 python -m vllm_omni.entrypoints.cli.main serve \
|
||||
"Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice" \
|
||||
--omni --host 127.0.0.1 --port 8000 \
|
||||
--stage-configs-path benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs1.yaml \
|
||||
--trust-remote-code
|
||||
```
|
||||
|
||||
### 2) Run online serving benchmark
|
||||
|
||||
```bash
|
||||
python vllm_omni/bench_tts_serve.py \
|
||||
--port 8000 \
|
||||
--num-prompts 50 \
|
||||
--max-concurrency 1 4 10 \
|
||||
--config-name "async_chunk" \
|
||||
--result-dir results/
|
||||
```
|
||||
|
||||
### 3) Run HuggingFace baseline
|
||||
|
||||
```bash
|
||||
python transformers/bench_tts_hf.py \
|
||||
--model "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice" \
|
||||
--num-prompts 50 \
|
||||
--gpu-device 0 \
|
||||
--result-dir results/
|
||||
```
|
||||
|
||||
### 4) Generate comparison plots
|
||||
|
||||
```bash
|
||||
python plot_results.py \
|
||||
--results results/bench_async_chunk_*.json results/bench_hf_transformers_*.json \
|
||||
--labels "vllm-omni" "hf_transformers" \
|
||||
--output results/comparison.png
|
||||
```
|
||||
|
||||
## Stage Configs
|
||||
|
||||
| Config | Batch Size | Description |
|
||||
|--------|:----------:|-------------|
|
||||
| `vllm_omni/configs/qwen3_tts_bs1.yaml` | 1 | Single-request processing (lower latency) |
|
||||
| `vllm_omni/configs/qwen3_tts_bs4.yaml` | 4 | Concurrent request processing (higher throughput) |
|
||||
|
||||
Both configs use a 2-stage pipeline (Talker -> Code2Wav) with `async_chunk` streaming enabled. The `SharedMemoryConnector` streams codec frames (25-frame chunks with 25-frame context overlap) between stages.
|
||||
|
||||
The model is specified via the CLI `--model` flag (or `MODEL` env var in `run_benchmark.sh`), so the same configs work for both the 0.6B and 1.7B model variants.
|
||||
|
||||
## Metrics
|
||||
|
||||
- **TTFP (Time to First Audio Packet)**: Time from request to first audio chunk (streaming latency)
|
||||
- **E2E (End-to-End Latency)**: Total time from request to complete audio response
|
||||
- **RTF (Real-Time Factor)**: E2E latency / audio duration. RTF < 1.0 means faster-than-real-time synthesis
|
||||
- **Throughput**: Total audio seconds generated per wall-clock second
|
||||
@@ -0,0 +1,254 @@
|
||||
"""Plot Qwen3-TTS benchmark results.
|
||||
|
||||
Generates comparison bar charts similar to the async_chunk design doc:
|
||||
- TTFP (Time-to-First-Packet) across concurrency levels
|
||||
- E2E latency across concurrency levels
|
||||
- RTF (Real-Time Factor) across concurrency levels
|
||||
|
||||
Usage:
|
||||
# Compare two configs (async_chunk vs no_async_chunk):
|
||||
python plot_results.py \
|
||||
--results results/bench_async_chunk_*.json results/bench_no_async_chunk_*.json \
|
||||
--labels "async_chunk" "no_async_chunk" \
|
||||
--output results/qwen3_tts_benchmark.png
|
||||
|
||||
# Single config:
|
||||
python plot_results.py \
|
||||
--results results/bench_async_chunk_*.json \
|
||||
--labels "async_chunk" \
|
||||
--output results/qwen3_tts_benchmark.png
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
|
||||
|
||||
def load_results(result_files: list[str]) -> list[list[dict]]:
|
||||
"""Load benchmark results from JSON files."""
|
||||
all_results = []
|
||||
for f in result_files:
|
||||
with open(f) as fh:
|
||||
data = json.load(fh)
|
||||
all_results.append(data)
|
||||
return all_results
|
||||
|
||||
|
||||
def plot_comparison(
|
||||
all_results: list[list[dict]],
|
||||
labels: list[str],
|
||||
output_path: str,
|
||||
title_prefix: str = "Qwen3-TTS",
|
||||
):
|
||||
"""Generate comparison bar charts."""
|
||||
n_configs = len(all_results)
|
||||
|
||||
# Collect concurrency levels present in ALL configs (skip missing data)
|
||||
all_concurrencies = [set(r["concurrency"] for r in results) for results in all_results]
|
||||
concurrencies = sorted(set.union(*all_concurrencies))
|
||||
|
||||
# Build data arrays, using None for missing concurrency levels
|
||||
ttfp_data = {label: [] for label in labels}
|
||||
e2e_data = {label: [] for label in labels}
|
||||
rtf_data = {label: [] for label in labels}
|
||||
throughput_data = {label: [] for label in labels}
|
||||
|
||||
for results, label in zip(all_results, labels):
|
||||
conc_map = {r["concurrency"]: r for r in results}
|
||||
for c in concurrencies:
|
||||
r = conc_map.get(c)
|
||||
ttfp_data[label].append(r["mean_ttfp_ms"] if r else None)
|
||||
e2e_data[label].append(r["mean_e2e_ms"] if r else None)
|
||||
rtf_data[label].append(r["mean_rtf"] if r else None)
|
||||
throughput_data[label].append(r["audio_throughput"] if r else None)
|
||||
|
||||
fig, axes = plt.subplots(2, 2, figsize=(14, 10))
|
||||
fig.suptitle(f"{title_prefix} Performance Benchmark", fontsize=16, fontweight="bold")
|
||||
|
||||
x = np.arange(len(concurrencies))
|
||||
width = 0.35 if n_configs == 2 else 0.5
|
||||
if n_configs > 1:
|
||||
offsets = np.linspace(-width / 2 * (n_configs - 1), width / 2 * (n_configs - 1), n_configs)
|
||||
else:
|
||||
offsets = [0]
|
||||
|
||||
colors = ["#2196F3", "#FF5722", "#4CAF50", "#FFC107"]
|
||||
|
||||
def plot_metric(ax, data_dict, ylabel, title, fmt=".1f"):
|
||||
bars = []
|
||||
for i, (label, values) in enumerate(data_dict.items()):
|
||||
# Replace None with 0 for plotting, but track which are missing
|
||||
plot_values = [v if v is not None else 0 for v in values]
|
||||
color = colors[i % len(colors)]
|
||||
bar = ax.bar(x + offsets[i], plot_values, width, label=label, color=color, alpha=0.85)
|
||||
bars.append(bar)
|
||||
# Add value labels on bars (skip None/missing data)
|
||||
max_val = max((v for v in values if v is not None), default=1)
|
||||
for rect, val in zip(bar, values):
|
||||
if val is not None and val > 0:
|
||||
ax.text(
|
||||
rect.get_x() + rect.get_width() / 2,
|
||||
rect.get_height() + max_val * 0.02,
|
||||
f"{val:{fmt}}",
|
||||
ha="center",
|
||||
va="bottom",
|
||||
fontsize=9,
|
||||
fontweight="bold",
|
||||
)
|
||||
ax.set_xlabel("Concurrency", fontsize=12)
|
||||
ax.set_ylabel(ylabel, fontsize=12)
|
||||
ax.set_title(title, fontsize=13, fontweight="bold")
|
||||
ax.set_xticks(x)
|
||||
ax.set_xticklabels([str(c) for c in concurrencies])
|
||||
ax.legend(fontsize=10)
|
||||
ax.grid(axis="y", alpha=0.3)
|
||||
ax.set_axisbelow(True)
|
||||
|
||||
plot_metric(axes[0, 0], ttfp_data, "TTFP (ms)", "Time to First Audio Packet (TTFP)")
|
||||
plot_metric(axes[0, 1], e2e_data, "E2E Latency (ms)", "End-to-End Latency (E2E)")
|
||||
plot_metric(axes[1, 0], rtf_data, "RTF", "Real-Time Factor (RTF)", fmt=".3f")
|
||||
plot_metric(axes[1, 1], throughput_data, "Audio-sec / Wall-sec", "Audio Throughput", fmt=".2f")
|
||||
|
||||
plt.tight_layout()
|
||||
plt.savefig(output_path, dpi=150, bbox_inches="tight")
|
||||
print(f"Plot saved to {output_path}")
|
||||
plt.close()
|
||||
|
||||
|
||||
def plot_single_summary(results: list[dict], label: str, output_path: str):
|
||||
"""Generate a single-config summary with percentile breakdown."""
|
||||
concurrencies = [r["concurrency"] for r in results]
|
||||
|
||||
fig, axes = plt.subplots(1, 3, figsize=(16, 5))
|
||||
fig.suptitle(f"Qwen3-TTS Benchmark - {label}", fontsize=15, fontweight="bold")
|
||||
|
||||
# TTFP breakdown
|
||||
ax = axes[0]
|
||||
means = [r["mean_ttfp_ms"] for r in results]
|
||||
medians = [r["median_ttfp_ms"] for r in results]
|
||||
p90s = [r["p90_ttfp_ms"] for r in results]
|
||||
p99s = [r["p99_ttfp_ms"] for r in results]
|
||||
x = np.arange(len(concurrencies))
|
||||
w = 0.2
|
||||
ax.bar(x - 1.5 * w, means, w, label="mean", color="#2196F3")
|
||||
ax.bar(x - 0.5 * w, medians, w, label="median", color="#4CAF50")
|
||||
ax.bar(x + 0.5 * w, p90s, w, label="p90", color="#FF9800")
|
||||
ax.bar(x + 1.5 * w, p99s, w, label="p99", color="#F44336")
|
||||
ax.set_xticks(x)
|
||||
ax.set_xticklabels([str(c) for c in concurrencies])
|
||||
ax.set_xlabel("Concurrency")
|
||||
ax.set_ylabel("TTFP (ms)")
|
||||
ax.set_title("Time to First Audio Packet")
|
||||
ax.legend(fontsize=9)
|
||||
ax.grid(axis="y", alpha=0.3)
|
||||
|
||||
# E2E breakdown
|
||||
ax = axes[1]
|
||||
means = [r["mean_e2e_ms"] for r in results]
|
||||
medians = [r["median_e2e_ms"] for r in results]
|
||||
p90s = [r["p90_e2e_ms"] for r in results]
|
||||
p99s = [r["p99_e2e_ms"] for r in results]
|
||||
ax.bar(x - 1.5 * w, means, w, label="mean", color="#2196F3")
|
||||
ax.bar(x - 0.5 * w, medians, w, label="median", color="#4CAF50")
|
||||
ax.bar(x + 0.5 * w, p90s, w, label="p90", color="#FF9800")
|
||||
ax.bar(x + 1.5 * w, p99s, w, label="p99", color="#F44336")
|
||||
ax.set_xticks(x)
|
||||
ax.set_xticklabels([str(c) for c in concurrencies])
|
||||
ax.set_xlabel("Concurrency")
|
||||
ax.set_ylabel("E2E Latency (ms)")
|
||||
ax.set_title("End-to-End Latency")
|
||||
ax.legend(fontsize=9)
|
||||
ax.grid(axis="y", alpha=0.3)
|
||||
|
||||
# RTF
|
||||
ax = axes[2]
|
||||
means = [r["mean_rtf"] for r in results]
|
||||
medians = [r["median_rtf"] for r in results]
|
||||
ax.bar(x - 0.15, means, 0.3, label="mean", color="#2196F3")
|
||||
ax.bar(x + 0.15, medians, 0.3, label="median", color="#4CAF50")
|
||||
ax.set_xticks(x)
|
||||
ax.set_xticklabels([str(c) for c in concurrencies])
|
||||
ax.set_xlabel("Concurrency")
|
||||
ax.set_ylabel("RTF")
|
||||
ax.set_title("Real-Time Factor")
|
||||
ax.legend(fontsize=9)
|
||||
ax.grid(axis="y", alpha=0.3)
|
||||
|
||||
plt.tight_layout()
|
||||
plt.savefig(output_path, dpi=150, bbox_inches="tight")
|
||||
print(f"Plot saved to {output_path}")
|
||||
plt.close()
|
||||
|
||||
|
||||
def print_comparison_table(all_results: list[list[dict]], labels: list[str]):
|
||||
"""Print a markdown-formatted comparison table."""
|
||||
concurrencies = sorted(set(r["concurrency"] for r in all_results[0]))
|
||||
|
||||
print("\n## Benchmark Results\n")
|
||||
header = "| Metric | Concurrency |"
|
||||
sep = "| --- | --- |"
|
||||
for label in labels:
|
||||
header += f" {label} |"
|
||||
sep += " --- |"
|
||||
print(header)
|
||||
print(sep)
|
||||
|
||||
for metric, key, fmt in [
|
||||
("TTFP (ms)", "mean_ttfp_ms", ".1f"),
|
||||
("E2E (ms)", "mean_e2e_ms", ".1f"),
|
||||
("RTF", "mean_rtf", ".3f"),
|
||||
("Throughput (audio-s/s)", "audio_throughput", ".2f"),
|
||||
]:
|
||||
for c in concurrencies:
|
||||
row = f"| {metric} | {c} |"
|
||||
for results in all_results:
|
||||
conc_map = {r["concurrency"]: r for r in results}
|
||||
val = conc_map.get(c, {}).get(key, 0)
|
||||
row += f" {val:{fmt}} |"
|
||||
print(row)
|
||||
|
||||
# Improvement calculation (only if 2 configs)
|
||||
if len(all_results) == 2:
|
||||
print(f"\n## Improvement ({labels[0]} vs {labels[1]})\n")
|
||||
print("| Metric | Concurrency | Improvement |")
|
||||
print("| --- | --- | --- |")
|
||||
for metric, key in [("TTFP", "mean_ttfp_ms"), ("E2E", "mean_e2e_ms"), ("RTF", "mean_rtf")]:
|
||||
for c in concurrencies:
|
||||
m0 = {r["concurrency"]: r for r in all_results[0]}
|
||||
m1 = {r["concurrency"]: r for r in all_results[1]}
|
||||
v0 = m0.get(c, {}).get(key, 0)
|
||||
v1 = m1.get(c, {}).get(key, 0)
|
||||
if v1 > 0:
|
||||
pct = (v1 - v0) / v1 * 100
|
||||
print(f"| {metric} | {c} | {pct:+.1f}% |")
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Plot Qwen3-TTS benchmark results")
|
||||
parser.add_argument(
|
||||
"--results", type=str, nargs="+", required=True, help="Path(s) to result JSON files (one per config)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--labels", type=str, nargs="+", required=True, help="Labels for each config (must match --results count)"
|
||||
)
|
||||
parser.add_argument("--output", type=str, default="results/qwen3_tts_benchmark.png", help="Output image path")
|
||||
parser.add_argument("--title", type=str, default="Qwen3-TTS", help="Title prefix for the plot")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
assert len(args.results) == len(args.labels), "--results and --labels must have the same count"
|
||||
|
||||
all_results = load_results(args.results)
|
||||
print_comparison_table(all_results, args.labels)
|
||||
|
||||
Path(args.output).parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
if len(all_results) == 1:
|
||||
plot_single_summary(all_results[0], args.labels[0], args.output)
|
||||
else:
|
||||
plot_comparison(all_results, args.labels, args.output, title_prefix=args.title)
|
||||
@@ -0,0 +1,3 @@
|
||||
# Benchmark results are machine-specific - do not commit
|
||||
*
|
||||
!.gitignore
|
||||
@@ -0,0 +1,272 @@
|
||||
#!/bin/bash
|
||||
# Qwen3-TTS Benchmark Runner
|
||||
#
|
||||
# Compares vllm-omni streaming serving vs HuggingFace transformers offline inference.
|
||||
# Produces JSON results and comparison plots.
|
||||
#
|
||||
# Usage:
|
||||
# # Full comparison (vllm-omni + HF):
|
||||
# bash run_benchmark.sh
|
||||
#
|
||||
# # Only vllm-omni async_chunk config:
|
||||
# bash run_benchmark.sh --async-only
|
||||
#
|
||||
# # Only HuggingFace baseline:
|
||||
# bash run_benchmark.sh --hf-only
|
||||
#
|
||||
# # vllm-omni only (skip HF):
|
||||
# bash run_benchmark.sh --skip-hf
|
||||
#
|
||||
# # Custom settings:
|
||||
# GPU_DEVICE=1 NUM_PROMPTS=20 CONCURRENCY="1 4" bash run_benchmark.sh
|
||||
#
|
||||
# # Use 1.7B model:
|
||||
# MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice bash run_benchmark.sh --async-only
|
||||
#
|
||||
# # Use batch_size=4 config:
|
||||
# STAGE_CONFIG=vllm_omni/configs/qwen3_tts_bs4.yaml bash run_benchmark.sh --async-only
|
||||
#
|
||||
# Environment variables:
|
||||
# GPU_DEVICE - GPU index to use (default: 0)
|
||||
# NUM_PROMPTS - Number of prompts per concurrency level (default: 50)
|
||||
# CONCURRENCY - Space-separated concurrency levels (default: "1 4 10")
|
||||
# MODEL - Model name (default: Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice)
|
||||
# PORT - Server port (default: 8000)
|
||||
# GPU_MEM_TALKER - gpu_memory_utilization for talker stage (default: 0.3)
|
||||
# GPU_MEM_CODE2WAV - gpu_memory_utilization for code2wav stage (default: 0.2)
|
||||
# STAGE_CONFIG - Path to stage config YAML (default: configs/qwen3_tts_bs1.yaml)
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
PROJECT_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)"
|
||||
|
||||
# Defaults
|
||||
GPU_DEVICE="${GPU_DEVICE:-0}"
|
||||
NUM_PROMPTS="${NUM_PROMPTS:-50}"
|
||||
CONCURRENCY="${CONCURRENCY:-1 4 10}"
|
||||
MODEL="${MODEL:-Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice}"
|
||||
PORT="${PORT:-8000}"
|
||||
GPU_MEM_TALKER="${GPU_MEM_TALKER:-0.3}"
|
||||
GPU_MEM_CODE2WAV="${GPU_MEM_CODE2WAV:-0.2}"
|
||||
NUM_WARMUPS="${NUM_WARMUPS:-3}"
|
||||
STAGE_CONFIG="${STAGE_CONFIG:-vllm_omni/configs/qwen3_tts_bs1.yaml}"
|
||||
RESULT_DIR="${SCRIPT_DIR}/results"
|
||||
TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
|
||||
|
||||
# Parse args
|
||||
RUN_ASYNC=true
|
||||
RUN_HF=true
|
||||
for arg in "$@"; do
|
||||
case "$arg" in
|
||||
--async-only) RUN_HF=false ;;
|
||||
--hf-only) RUN_ASYNC=false ;;
|
||||
--skip-hf) RUN_HF=false ;;
|
||||
esac
|
||||
done
|
||||
|
||||
mkdir -p "${RESULT_DIR}"
|
||||
|
||||
echo "============================================================"
|
||||
echo " Qwen3-TTS Benchmark"
|
||||
echo "============================================================"
|
||||
echo " GPU: ${GPU_DEVICE}"
|
||||
echo " Model: ${MODEL}"
|
||||
echo " Prompts: ${NUM_PROMPTS}"
|
||||
echo " Concurrency: ${CONCURRENCY}"
|
||||
echo " Port: ${PORT}"
|
||||
echo " Stage config: ${STAGE_CONFIG}"
|
||||
echo " Results: ${RESULT_DIR}"
|
||||
echo "============================================================"
|
||||
|
||||
# Prepare stage config with correct GPU device and memory settings
|
||||
prepare_config() {
|
||||
local config_template="$1"
|
||||
local config_name="$2"
|
||||
local output_path="${RESULT_DIR}/${config_name}_stage_config.yaml"
|
||||
|
||||
# Use sed to patch GPU device and memory utilization
|
||||
sed \
|
||||
-e "s/devices: \"0\"/devices: \"${GPU_DEVICE}\"/g" \
|
||||
-e "s/gpu_memory_utilization: 0.3/gpu_memory_utilization: ${GPU_MEM_TALKER}/g" \
|
||||
-e "s/gpu_memory_utilization: 0.2/gpu_memory_utilization: ${GPU_MEM_CODE2WAV}/g" \
|
||||
"${config_template}" > "${output_path}"
|
||||
|
||||
echo "${output_path}"
|
||||
}
|
||||
|
||||
# Start server and wait for it to be ready
|
||||
start_server() {
|
||||
local stage_config="$1"
|
||||
local config_name="$2"
|
||||
local log_file="${RESULT_DIR}/server_${config_name}_${TIMESTAMP}.log"
|
||||
|
||||
echo ""
|
||||
echo "Starting server with config: ${config_name}"
|
||||
echo " Stage config: ${stage_config}"
|
||||
echo " Log file: ${log_file}"
|
||||
|
||||
VLLM_WORKER_MULTIPROC_METHOD=spawn \
|
||||
CUDA_VISIBLE_DEVICES="${GPU_DEVICE}" \
|
||||
python -m vllm_omni.entrypoints.cli.main serve "${MODEL}" \
|
||||
--omni \
|
||||
--host 127.0.0.1 \
|
||||
--port "${PORT}" \
|
||||
--stage-configs-path "${stage_config}" \
|
||||
--stage-init-timeout 120 \
|
||||
--trust-remote-code \
|
||||
--disable-log-stats \
|
||||
> "${log_file}" 2>&1 &
|
||||
|
||||
SERVER_PID=$!
|
||||
echo " Server PID: ${SERVER_PID}"
|
||||
|
||||
# Wait for server to be ready
|
||||
echo " Waiting for server to be ready..."
|
||||
local max_wait=300
|
||||
local waited=0
|
||||
while [ ${waited} -lt ${max_wait} ]; do
|
||||
if curl -sf "http://127.0.0.1:${PORT}/v1/models" > /dev/null 2>&1; then
|
||||
echo " Server is ready! (waited ${waited}s)"
|
||||
return 0
|
||||
fi
|
||||
# Check if process is still alive
|
||||
if ! kill -0 ${SERVER_PID} 2>/dev/null; then
|
||||
echo " ERROR: Server process died. Check log: ${log_file}"
|
||||
tail -20 "${log_file}"
|
||||
return 1
|
||||
fi
|
||||
sleep 2
|
||||
waited=$((waited + 2))
|
||||
done
|
||||
|
||||
echo " ERROR: Server did not start within ${max_wait}s. Check log: ${log_file}"
|
||||
kill ${SERVER_PID} 2>/dev/null || true
|
||||
return 1
|
||||
}
|
||||
|
||||
# Stop the server
|
||||
stop_server() {
|
||||
if [ -n "${SERVER_PID:-}" ]; then
|
||||
echo " Stopping server (PID: ${SERVER_PID})..."
|
||||
kill ${SERVER_PID} 2>/dev/null || true
|
||||
wait ${SERVER_PID} 2>/dev/null || true
|
||||
# Kill any remaining child processes on the port
|
||||
local pids
|
||||
pids=$(lsof -ti:${PORT} 2>/dev/null || true)
|
||||
if [ -n "${pids}" ]; then
|
||||
echo " Cleaning up remaining processes on port ${PORT}..."
|
||||
echo "${pids}" | xargs kill -9 2>/dev/null || true
|
||||
fi
|
||||
echo " Server stopped."
|
||||
SERVER_PID=""
|
||||
fi
|
||||
}
|
||||
|
||||
# Cleanup on exit
|
||||
trap 'stop_server' EXIT
|
||||
|
||||
# Run benchmark for a given config
|
||||
run_bench() {
|
||||
local config_name="$1"
|
||||
local config_template="$2"
|
||||
|
||||
echo ""
|
||||
echo "============================================================"
|
||||
echo " Benchmarking: ${config_name}"
|
||||
echo "============================================================"
|
||||
|
||||
local stage_config
|
||||
stage_config=$(prepare_config "${config_template}" "${config_name}")
|
||||
|
||||
start_server "${stage_config}" "${config_name}"
|
||||
|
||||
# Convert concurrency string to args
|
||||
local conc_args=""
|
||||
for c in ${CONCURRENCY}; do
|
||||
conc_args="${conc_args} ${c}"
|
||||
done
|
||||
|
||||
cd "${PROJECT_ROOT}"
|
||||
python "${SCRIPT_DIR}/vllm_omni/bench_tts_serve.py" \
|
||||
--host 127.0.0.1 \
|
||||
--port "${PORT}" \
|
||||
--num-prompts "${NUM_PROMPTS}" \
|
||||
--max-concurrency ${conc_args} \
|
||||
--num-warmups "${NUM_WARMUPS}" \
|
||||
--config-name "${config_name}" \
|
||||
--result-dir "${RESULT_DIR}"
|
||||
|
||||
stop_server
|
||||
|
||||
# Allow GPU memory to settle
|
||||
sleep 5
|
||||
}
|
||||
|
||||
# Run vllm-omni benchmark
|
||||
if [ "${RUN_ASYNC}" = true ]; then
|
||||
run_bench "async_chunk" "${SCRIPT_DIR}/${STAGE_CONFIG}"
|
||||
fi
|
||||
|
||||
# Run HuggingFace baseline benchmark
|
||||
if [ "${RUN_HF}" = true ]; then
|
||||
echo ""
|
||||
echo "============================================================"
|
||||
echo " Benchmarking: HuggingFace transformers (offline)"
|
||||
echo "============================================================"
|
||||
|
||||
cd "${PROJECT_ROOT}"
|
||||
python "${SCRIPT_DIR}/transformers/bench_tts_hf.py" \
|
||||
--model "${MODEL}" \
|
||||
--num-prompts "${NUM_PROMPTS}" \
|
||||
--num-warmups "${NUM_WARMUPS}" \
|
||||
--gpu-device "${GPU_DEVICE}" \
|
||||
--config-name "hf_transformers" \
|
||||
--result-dir "${RESULT_DIR}"
|
||||
|
||||
# Allow GPU memory to settle
|
||||
sleep 5
|
||||
fi
|
||||
|
||||
# Plot results
|
||||
echo ""
|
||||
echo "============================================================"
|
||||
echo " Generating plots..."
|
||||
echo "============================================================"
|
||||
|
||||
RESULT_FILES=""
|
||||
LABELS=""
|
||||
|
||||
if [ "${RUN_ASYNC}" = true ]; then
|
||||
ASYNC_FILE=$(ls -t "${RESULT_DIR}"/bench_async_chunk_*.json 2>/dev/null | head -1)
|
||||
if [ -n "${ASYNC_FILE}" ]; then
|
||||
RESULT_FILES="${ASYNC_FILE}"
|
||||
LABELS="async_chunk"
|
||||
fi
|
||||
fi
|
||||
|
||||
if [ "${RUN_HF}" = true ]; then
|
||||
HF_FILE=$(ls -t "${RESULT_DIR}"/bench_hf_transformers_*.json 2>/dev/null | head -1)
|
||||
if [ -n "${HF_FILE}" ]; then
|
||||
if [ -n "${RESULT_FILES}" ]; then
|
||||
RESULT_FILES="${RESULT_FILES} ${HF_FILE}"
|
||||
LABELS="${LABELS} hf_transformers"
|
||||
else
|
||||
RESULT_FILES="${HF_FILE}"
|
||||
LABELS="hf_transformers"
|
||||
fi
|
||||
fi
|
||||
fi
|
||||
|
||||
if [ -n "${RESULT_FILES}" ]; then
|
||||
python "${SCRIPT_DIR}/plot_results.py" \
|
||||
--results ${RESULT_FILES} \
|
||||
--labels ${LABELS} \
|
||||
--output "${RESULT_DIR}/qwen3_tts_benchmark_${TIMESTAMP}.png"
|
||||
fi
|
||||
|
||||
echo ""
|
||||
echo "============================================================"
|
||||
echo " Benchmark complete!"
|
||||
echo " Results: ${RESULT_DIR}"
|
||||
echo "============================================================"
|
||||
@@ -0,0 +1,281 @@
|
||||
"""Benchmark Qwen3-TTS using HuggingFace transformers (qwen_tts library).
|
||||
|
||||
Measures E2E latency, RTF, and audio duration for offline (non-serving) inference.
|
||||
Results are saved in the same JSON format as bench_tts_serve.py for unified plotting.
|
||||
|
||||
Usage:
|
||||
python bench_tts_hf.py \
|
||||
--model Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice \
|
||||
--num-prompts 50 \
|
||||
--num-warmups 3 \
|
||||
--gpu-device 0 \
|
||||
--result-dir results/
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import time
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
import torch
|
||||
|
||||
PROMPTS = [
|
||||
"Hello, welcome to the voice synthesis benchmark test.",
|
||||
"She said she would be here by noon, but nobody showed up.",
|
||||
"The quick brown fox jumps over the lazy dog near the riverbank.",
|
||||
"I can't believe how beautiful the sunset looks from up here on the mountain.",
|
||||
"Please remember to bring your identification documents to the appointment tomorrow morning.",
|
||||
"Have you ever wondered what it would be like to travel through time and visit ancient civilizations?",
|
||||
"The restaurant on the corner serves the best pasta I have ever tasted in my entire life.",
|
||||
"After the meeting, we should discuss the quarterly results and plan for the next phase.",
|
||||
"Learning a new language takes patience, practice, and a genuine curiosity about other cultures.",
|
||||
"The train leaves at half past seven, so we need to arrive at the station before then.",
|
||||
"Could you please turn down the music a little bit, I'm trying to concentrate on my work.",
|
||||
"It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.",
|
||||
]
|
||||
|
||||
|
||||
@dataclass
|
||||
class BenchmarkResult:
|
||||
config_name: str = ""
|
||||
concurrency: int = 1 # always 1 for offline
|
||||
num_prompts: int = 0
|
||||
completed: int = 0
|
||||
failed: int = 0
|
||||
duration_s: float = 0.0
|
||||
# TTFP stats - not applicable for HF offline, set to E2E for compatibility
|
||||
mean_ttfp_ms: float = 0.0
|
||||
median_ttfp_ms: float = 0.0
|
||||
std_ttfp_ms: float = 0.0
|
||||
p90_ttfp_ms: float = 0.0
|
||||
p95_ttfp_ms: float = 0.0
|
||||
p99_ttfp_ms: float = 0.0
|
||||
# E2E stats (ms)
|
||||
mean_e2e_ms: float = 0.0
|
||||
median_e2e_ms: float = 0.0
|
||||
std_e2e_ms: float = 0.0
|
||||
p90_e2e_ms: float = 0.0
|
||||
p95_e2e_ms: float = 0.0
|
||||
p99_e2e_ms: float = 0.0
|
||||
# RTF stats
|
||||
mean_rtf: float = 0.0
|
||||
median_rtf: float = 0.0
|
||||
std_rtf: float = 0.0
|
||||
p99_rtf: float = 0.0
|
||||
# Audio stats
|
||||
mean_audio_duration_s: float = 0.0
|
||||
total_audio_duration_s: float = 0.0
|
||||
audio_throughput: float = 0.0
|
||||
request_throughput: float = 0.0
|
||||
# Per-request details
|
||||
per_request: list = field(default_factory=list)
|
||||
|
||||
|
||||
def run_benchmark(args):
|
||||
from qwen_tts import Qwen3TTSModel
|
||||
|
||||
device = f"cuda:{args.gpu_device}"
|
||||
print(f"Loading model: {args.model} on {device}")
|
||||
model = Qwen3TTSModel.from_pretrained(
|
||||
args.model,
|
||||
device_map=device,
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
print("Model loaded.")
|
||||
|
||||
# Build prompt list
|
||||
prompts = [PROMPTS[i % len(PROMPTS)] for i in range(args.num_prompts)]
|
||||
|
||||
# Warmup
|
||||
if args.num_warmups > 0:
|
||||
print(f"Warming up with {args.num_warmups} requests...")
|
||||
for i in range(args.num_warmups):
|
||||
p = PROMPTS[i % len(PROMPTS)]
|
||||
wavs, sr = model.generate_custom_voice(
|
||||
text=p,
|
||||
language=args.language,
|
||||
speaker=args.voice,
|
||||
)
|
||||
# Sync GPU
|
||||
torch.cuda.synchronize(device)
|
||||
print("Warmup done.")
|
||||
|
||||
# Benchmark
|
||||
print(f"Running {args.num_prompts} requests sequentially...")
|
||||
e2e_times = []
|
||||
rtfs = []
|
||||
audio_durations = []
|
||||
per_request = []
|
||||
failed = 0
|
||||
|
||||
audio_dir = None
|
||||
if args.save_audio:
|
||||
audio_dir = Path(args.result_dir) / "audio_hf"
|
||||
audio_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
total_start = time.perf_counter()
|
||||
|
||||
for i, prompt in enumerate(prompts):
|
||||
try:
|
||||
torch.cuda.synchronize(device)
|
||||
st = time.perf_counter()
|
||||
|
||||
wavs, sr = model.generate_custom_voice(
|
||||
text=prompt,
|
||||
language=args.language,
|
||||
speaker=args.voice,
|
||||
)
|
||||
|
||||
torch.cuda.synchronize(device)
|
||||
elapsed = time.perf_counter() - st
|
||||
|
||||
# Compute audio duration
|
||||
audio_samples = wavs[0]
|
||||
if isinstance(audio_samples, torch.Tensor):
|
||||
audio_samples = audio_samples.cpu().numpy()
|
||||
audio_dur = len(audio_samples) / sr
|
||||
|
||||
rtf = elapsed / audio_dur if audio_dur > 0 else 0.0
|
||||
|
||||
e2e_times.append(elapsed)
|
||||
rtfs.append(rtf)
|
||||
audio_durations.append(audio_dur)
|
||||
per_request.append(
|
||||
{
|
||||
"e2e_ms": elapsed * 1000,
|
||||
"ttfp_ms": elapsed * 1000, # no streaming, TTFP = E2E
|
||||
"rtf": rtf,
|
||||
"audio_duration_s": audio_dur,
|
||||
"prompt": prompt,
|
||||
}
|
||||
)
|
||||
|
||||
if audio_dir:
|
||||
sf.write(str(audio_dir / f"output_{i:04d}.wav"), audio_samples, sr)
|
||||
|
||||
if (i + 1) % 10 == 0 or i == 0:
|
||||
print(
|
||||
f" [{i + 1}/{args.num_prompts}] e2e={elapsed * 1000:.0f}ms rtf={rtf:.3f} audio={audio_dur:.2f}s"
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
print(f" [{i + 1}/{args.num_prompts}] FAILED: {e}")
|
||||
failed += 1
|
||||
|
||||
total_duration = time.perf_counter() - total_start
|
||||
completed = len(e2e_times)
|
||||
|
||||
# Compute stats
|
||||
result = BenchmarkResult(
|
||||
config_name=args.config_name,
|
||||
concurrency=1,
|
||||
num_prompts=args.num_prompts,
|
||||
completed=completed,
|
||||
failed=failed,
|
||||
duration_s=total_duration,
|
||||
)
|
||||
|
||||
if e2e_times:
|
||||
e2e_ms = [t * 1000 for t in e2e_times]
|
||||
|
||||
result.mean_e2e_ms = float(np.mean(e2e_ms))
|
||||
result.median_e2e_ms = float(np.median(e2e_ms))
|
||||
result.std_e2e_ms = float(np.std(e2e_ms))
|
||||
result.p90_e2e_ms = float(np.percentile(e2e_ms, 90))
|
||||
result.p95_e2e_ms = float(np.percentile(e2e_ms, 95))
|
||||
result.p99_e2e_ms = float(np.percentile(e2e_ms, 99))
|
||||
|
||||
# For HF offline, TTFP = E2E (no streaming)
|
||||
result.mean_ttfp_ms = result.mean_e2e_ms
|
||||
result.median_ttfp_ms = result.median_e2e_ms
|
||||
result.std_ttfp_ms = result.std_e2e_ms
|
||||
result.p90_ttfp_ms = result.p90_e2e_ms
|
||||
result.p95_ttfp_ms = result.p95_e2e_ms
|
||||
result.p99_ttfp_ms = result.p99_e2e_ms
|
||||
|
||||
result.mean_rtf = float(np.mean(rtfs))
|
||||
result.median_rtf = float(np.median(rtfs))
|
||||
result.std_rtf = float(np.std(rtfs))
|
||||
result.p99_rtf = float(np.percentile(rtfs, 99))
|
||||
|
||||
result.mean_audio_duration_s = float(np.mean(audio_durations))
|
||||
result.total_audio_duration_s = float(np.sum(audio_durations))
|
||||
result.audio_throughput = result.total_audio_duration_s / total_duration
|
||||
result.request_throughput = completed / total_duration
|
||||
result.per_request = per_request
|
||||
|
||||
# Print summary in standardized performance template
|
||||
W = 50
|
||||
print("")
|
||||
print(f"{'=' * W}")
|
||||
print(f"{'Serving Benchmark Result':^{W}}")
|
||||
print(f"{'=' * W}")
|
||||
print(f"{'Successful requests:':<40}{completed:<10}")
|
||||
print(f"{'Failed requests:':<40}{failed:<10}")
|
||||
print(f"{'Maximum request concurrency:':<40}{1:<10}")
|
||||
print(f"{'Benchmark duration (s):':<40}{total_duration:<10.2f}")
|
||||
print(f"{'Request throughput (req/s):':<40}{result.request_throughput:<10.2f}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'End-to-end Latency':^{W}}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Mean E2EL (ms):':<40}{result.mean_e2e_ms:<10.2f}")
|
||||
print(f"{'Median E2EL (ms):':<40}{result.median_e2e_ms:<10.2f}")
|
||||
print(f"{'P99 E2EL (ms):':<40}{result.p99_e2e_ms:<10.2f}")
|
||||
print(f"{'=' * W}")
|
||||
print(f"{'Audio Result':^{W}}")
|
||||
print(f"{'=' * W}")
|
||||
print(f"{'Total audio duration generated (s):':<40}{result.total_audio_duration_s:<10.2f}")
|
||||
print(f"{'Audio throughput (audio duration/s):':<40}{result.audio_throughput:<10.2f}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Time to First Packet':^{W}}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Mean AUDIO_TTFP (ms):':<40}{result.mean_ttfp_ms:<10.2f}")
|
||||
print(f"{'Median AUDIO_TTFP (ms):':<40}{result.median_ttfp_ms:<10.2f}")
|
||||
print(f"{'P99 AUDIO_TTFP (ms):':<40}{result.p99_ttfp_ms:<10.2f}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Real Time Factor':^{W}}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Mean AUDIO_RTF:':<40}{result.mean_rtf:<10.3f}")
|
||||
print(f"{'Median AUDIO_RTF:':<40}{result.median_rtf:<10.3f}")
|
||||
print(f"{'P99 AUDIO_RTF:':<40}{result.p99_rtf:<10.3f}")
|
||||
print(f"{'=' * W}")
|
||||
print("")
|
||||
|
||||
# Save results (as a list with single concurrency=1 entry, matching serve format)
|
||||
result_dir = Path(args.result_dir)
|
||||
result_dir.mkdir(parents=True, exist_ok=True)
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
result_file = result_dir / f"bench_{args.config_name}_{timestamp}.json"
|
||||
|
||||
with open(result_file, "w") as f:
|
||||
json.dump([asdict(result)], f, indent=2)
|
||||
print(f"Results saved to {result_file}")
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Qwen3-TTS HuggingFace Benchmark")
|
||||
parser.add_argument(
|
||||
"--model", type=str, default="Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice", help="HuggingFace model name or path"
|
||||
)
|
||||
parser.add_argument("--num-prompts", type=int, default=50)
|
||||
parser.add_argument("--num-warmups", type=int, default=3)
|
||||
parser.add_argument("--gpu-device", type=int, default=0)
|
||||
parser.add_argument("--voice", type=str, default="Vivian")
|
||||
parser.add_argument("--language", type=str, default="English")
|
||||
parser.add_argument(
|
||||
"--config-name", type=str, default="hf_transformers", help="Label for this config (used in filenames)"
|
||||
)
|
||||
parser.add_argument("--result-dir", type=str, default="results")
|
||||
parser.add_argument("--save-audio", action="store_true", help="Save generated audio files")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
run_benchmark(args)
|
||||
@@ -0,0 +1,301 @@
|
||||
"""Benchmark comparing async_chunk on vs off for Qwen3-TTS.
|
||||
|
||||
Measures TTFP (Time-to-First-Packet), E2E latency, and RTF across
|
||||
concurrency levels for both async_chunk modes. Saves results as JSON.
|
||||
|
||||
Usage:
|
||||
# Run against a server already serving with a given config:
|
||||
python bench_async_chunk.py \
|
||||
--host 127.0.0.1 --port 8000 \
|
||||
--config-name async_chunk_on \
|
||||
--num-prompts 50 \
|
||||
--max-concurrency 1 10 \
|
||||
--result-dir results/
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import time
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
import aiohttp
|
||||
import numpy as np
|
||||
from tqdm.asyncio import tqdm
|
||||
|
||||
PROMPTS = [
|
||||
"Hello, welcome to the voice synthesis benchmark test.",
|
||||
"She said she would be here by noon, but nobody showed up.",
|
||||
"The quick brown fox jumps over the lazy dog near the riverbank.",
|
||||
"I can't believe how beautiful the sunset looks from up here on the mountain.",
|
||||
"Please remember to bring your identification documents to the appointment tomorrow morning.",
|
||||
"Have you ever wondered what it would be like to travel through time and visit ancient civilizations?",
|
||||
"The restaurant on the corner serves the best pasta I have ever tasted in my entire life.",
|
||||
"After the meeting, we should discuss the quarterly results and plan for the next phase.",
|
||||
"Learning a new language takes patience, practice, and a genuine curiosity about other cultures.",
|
||||
"The train leaves at half past seven, so we need to arrive at the station before then.",
|
||||
"Could you please turn down the music a little bit, I'm trying to concentrate on my work.",
|
||||
"It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.",
|
||||
]
|
||||
|
||||
|
||||
@dataclass
|
||||
class RequestResult:
|
||||
success: bool = False
|
||||
ttfp: float = 0.0
|
||||
e2e: float = 0.0
|
||||
audio_bytes: int = 0
|
||||
audio_duration: float = 0.0
|
||||
rtf: float = 0.0
|
||||
prompt: str = ""
|
||||
error: str = ""
|
||||
|
||||
|
||||
@dataclass
|
||||
class BenchmarkResult:
|
||||
config_name: str = ""
|
||||
concurrency: int = 0
|
||||
num_prompts: int = 0
|
||||
completed: int = 0
|
||||
failed: int = 0
|
||||
duration_s: float = 0.0
|
||||
mean_ttfp_ms: float = 0.0
|
||||
median_ttfp_ms: float = 0.0
|
||||
std_ttfp_ms: float = 0.0
|
||||
p90_ttfp_ms: float = 0.0
|
||||
p95_ttfp_ms: float = 0.0
|
||||
p99_ttfp_ms: float = 0.0
|
||||
mean_e2e_ms: float = 0.0
|
||||
median_e2e_ms: float = 0.0
|
||||
std_e2e_ms: float = 0.0
|
||||
p90_e2e_ms: float = 0.0
|
||||
p95_e2e_ms: float = 0.0
|
||||
p99_e2e_ms: float = 0.0
|
||||
mean_rtf: float = 0.0
|
||||
median_rtf: float = 0.0
|
||||
std_rtf: float = 0.0
|
||||
mean_audio_duration_s: float = 0.0
|
||||
total_audio_duration_s: float = 0.0
|
||||
audio_throughput: float = 0.0
|
||||
request_throughput: float = 0.0
|
||||
per_request: list = field(default_factory=list)
|
||||
|
||||
|
||||
def pcm_bytes_to_duration(num_bytes: int, sample_rate: int = 24000, sample_width: int = 2) -> float:
|
||||
return num_bytes / sample_width / sample_rate
|
||||
|
||||
|
||||
async def send_tts_request(
|
||||
session: aiohttp.ClientSession,
|
||||
api_url: str,
|
||||
prompt: str,
|
||||
voice: str = "vivian",
|
||||
language: str = "English",
|
||||
stream: bool = True,
|
||||
pbar: tqdm | None = None,
|
||||
) -> RequestResult:
|
||||
payload = {
|
||||
"input": prompt,
|
||||
"voice": voice,
|
||||
"language": language,
|
||||
"stream": stream,
|
||||
"response_format": "pcm",
|
||||
}
|
||||
|
||||
result = RequestResult(prompt=prompt)
|
||||
st = time.perf_counter()
|
||||
|
||||
try:
|
||||
async with session.post(api_url, json=payload) as response:
|
||||
if response.status != 200:
|
||||
result.error = f"HTTP {response.status}: {await response.text()}"
|
||||
return result
|
||||
|
||||
first_chunk = True
|
||||
total_bytes = 0
|
||||
|
||||
async for chunk in response.content.iter_any():
|
||||
if first_chunk and len(chunk) > 0:
|
||||
result.ttfp = time.perf_counter() - st
|
||||
first_chunk = False
|
||||
total_bytes += len(chunk)
|
||||
|
||||
result.e2e = time.perf_counter() - st
|
||||
result.audio_bytes = total_bytes
|
||||
result.audio_duration = pcm_bytes_to_duration(total_bytes)
|
||||
if result.audio_duration > 0:
|
||||
result.rtf = result.e2e / result.audio_duration
|
||||
result.success = True
|
||||
|
||||
except Exception as e:
|
||||
result.error = str(e)
|
||||
result.e2e = time.perf_counter() - st
|
||||
|
||||
if pbar:
|
||||
pbar.update(1)
|
||||
return result
|
||||
|
||||
|
||||
async def run_benchmark(
|
||||
host: str,
|
||||
port: int,
|
||||
num_prompts: int,
|
||||
max_concurrency: int,
|
||||
num_warmups: int = 3,
|
||||
voice: str = "vivian",
|
||||
language: str = "English",
|
||||
stream: bool = True,
|
||||
) -> BenchmarkResult:
|
||||
api_url = f"http://{host}:{port}/v1/audio/speech"
|
||||
|
||||
connector = aiohttp.TCPConnector(limit=max_concurrency, limit_per_host=max_concurrency, keepalive_timeout=60)
|
||||
session = aiohttp.ClientSession(connector=connector, timeout=aiohttp.ClientTimeout(total=600))
|
||||
|
||||
if num_warmups > 0:
|
||||
print(f" Warming up with {num_warmups} requests...")
|
||||
warmup_tasks = [
|
||||
send_tts_request(session, api_url, PROMPTS[i % len(PROMPTS)], voice, language, stream)
|
||||
for i in range(num_warmups)
|
||||
]
|
||||
await asyncio.gather(*warmup_tasks)
|
||||
print(" Warmup done.")
|
||||
|
||||
request_prompts = [PROMPTS[i % len(PROMPTS)] for i in range(num_prompts)]
|
||||
|
||||
print(f" Running {num_prompts} requests with concurrency={max_concurrency}...")
|
||||
semaphore = asyncio.Semaphore(max_concurrency)
|
||||
pbar = tqdm(total=num_prompts, desc=f" concurrency={max_concurrency}")
|
||||
|
||||
async def limited_request(prompt):
|
||||
async with semaphore:
|
||||
return await send_tts_request(session, api_url, prompt, voice, language, stream, pbar)
|
||||
|
||||
start_time = time.perf_counter()
|
||||
tasks = [asyncio.create_task(limited_request(p)) for p in request_prompts]
|
||||
results: list[RequestResult] = await asyncio.gather(*tasks)
|
||||
duration = time.perf_counter() - start_time
|
||||
pbar.close()
|
||||
|
||||
await session.close()
|
||||
|
||||
successful = [r for r in results if r.success]
|
||||
failed = [r for r in results if not r.success]
|
||||
|
||||
bench = BenchmarkResult(
|
||||
concurrency=max_concurrency,
|
||||
num_prompts=num_prompts,
|
||||
completed=len(successful),
|
||||
failed=len(failed),
|
||||
duration_s=duration,
|
||||
)
|
||||
|
||||
if successful:
|
||||
ttfps = [r.ttfp * 1000 for r in successful]
|
||||
e2es = [r.e2e * 1000 for r in successful]
|
||||
rtfs = [r.rtf for r in successful]
|
||||
audio_durs = [r.audio_duration for r in successful]
|
||||
|
||||
bench.mean_ttfp_ms = float(np.mean(ttfps))
|
||||
bench.median_ttfp_ms = float(np.median(ttfps))
|
||||
bench.std_ttfp_ms = float(np.std(ttfps))
|
||||
bench.p90_ttfp_ms = float(np.percentile(ttfps, 90))
|
||||
bench.p95_ttfp_ms = float(np.percentile(ttfps, 95))
|
||||
bench.p99_ttfp_ms = float(np.percentile(ttfps, 99))
|
||||
|
||||
bench.mean_e2e_ms = float(np.mean(e2es))
|
||||
bench.median_e2e_ms = float(np.median(e2es))
|
||||
bench.std_e2e_ms = float(np.std(e2es))
|
||||
bench.p90_e2e_ms = float(np.percentile(e2es, 90))
|
||||
bench.p95_e2e_ms = float(np.percentile(e2es, 95))
|
||||
bench.p99_e2e_ms = float(np.percentile(e2es, 99))
|
||||
|
||||
bench.mean_rtf = float(np.mean(rtfs))
|
||||
bench.median_rtf = float(np.median(rtfs))
|
||||
bench.std_rtf = float(np.std(rtfs))
|
||||
|
||||
bench.mean_audio_duration_s = float(np.mean(audio_durs))
|
||||
bench.total_audio_duration_s = float(np.sum(audio_durs))
|
||||
bench.audio_throughput = bench.total_audio_duration_s / duration
|
||||
bench.request_throughput = len(successful) / duration
|
||||
|
||||
bench.per_request = [
|
||||
{
|
||||
"ttfp_ms": r.ttfp * 1000,
|
||||
"e2e_ms": r.e2e * 1000,
|
||||
"rtf": r.rtf,
|
||||
"audio_duration_s": r.audio_duration,
|
||||
"prompt": r.prompt,
|
||||
}
|
||||
for r in successful
|
||||
]
|
||||
|
||||
print(f"\n{'=' * 60}")
|
||||
print(f" Concurrency: {max_concurrency} | Completed: {bench.completed} | Failed: {bench.failed}")
|
||||
print(f" Duration: {duration:.2f}s | Throughput: {bench.request_throughput:.2f} req/s")
|
||||
print(
|
||||
f" TTFP (ms): mean={bench.mean_ttfp_ms:.1f} median={bench.median_ttfp_ms:.1f}"
|
||||
f" p90={bench.p90_ttfp_ms:.1f} p99={bench.p99_ttfp_ms:.1f}"
|
||||
)
|
||||
print(
|
||||
f" E2E (ms): mean={bench.mean_e2e_ms:.1f} median={bench.median_e2e_ms:.1f}"
|
||||
f" p90={bench.p90_e2e_ms:.1f} p99={bench.p99_e2e_ms:.1f}"
|
||||
)
|
||||
print(f" RTF: mean={bench.mean_rtf:.3f} median={bench.median_rtf:.3f}")
|
||||
print(f" Throughput: {bench.audio_throughput:.2f} audio-sec/wall-sec")
|
||||
print(f"{'=' * 60}\n")
|
||||
|
||||
if failed:
|
||||
for r in failed[:3]:
|
||||
print(f" [ERROR] {r.error[:200]}")
|
||||
|
||||
return bench
|
||||
|
||||
|
||||
async def main(args):
|
||||
all_results = []
|
||||
|
||||
for concurrency in args.max_concurrency:
|
||||
result = await run_benchmark(
|
||||
host=args.host,
|
||||
port=args.port,
|
||||
num_prompts=args.num_prompts,
|
||||
max_concurrency=concurrency,
|
||||
num_warmups=args.num_warmups,
|
||||
voice=args.voice,
|
||||
language=args.language,
|
||||
stream=args.stream,
|
||||
)
|
||||
result.config_name = args.config_name
|
||||
all_results.append(asdict(result))
|
||||
|
||||
result_dir = Path(args.result_dir)
|
||||
result_dir.mkdir(parents=True, exist_ok=True)
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
result_file = result_dir / f"bench_{args.config_name}_{timestamp}.json"
|
||||
|
||||
with open(result_file, "w") as f:
|
||||
json.dump(all_results, f, indent=2)
|
||||
print(f"Results saved to {result_file}")
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Qwen3-TTS async_chunk benchmark client")
|
||||
parser.add_argument("--host", type=str, default="127.0.0.1")
|
||||
parser.add_argument("--port", type=int, default=8000)
|
||||
parser.add_argument("--num-prompts", type=int, default=50)
|
||||
parser.add_argument("--max-concurrency", type=int, nargs="+", default=[1, 10])
|
||||
parser.add_argument("--num-warmups", type=int, default=3)
|
||||
parser.add_argument("--voice", type=str, default="vivian")
|
||||
parser.add_argument("--language", type=str, default="English")
|
||||
parser.add_argument("--stream", action="store_true", default=True)
|
||||
parser.add_argument("--no-stream", dest="stream", action="store_false")
|
||||
parser.add_argument("--config-name", type=str, default="async_chunk_on")
|
||||
parser.add_argument("--result-dir", type=str, default="results")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
asyncio.run(main(args))
|
||||
@@ -0,0 +1,348 @@
|
||||
"""Benchmark client for Qwen3-TTS via /v1/audio/speech endpoint.
|
||||
|
||||
Measures TTFP (Time-to-First-Packet), E2E latency, and RTF (Real-Time Factor)
|
||||
across configurable concurrency levels. Saves results as JSON for plotting.
|
||||
|
||||
Usage:
|
||||
python bench_tts_serve.py \
|
||||
--host 127.0.0.1 --port 8000 \
|
||||
--num-prompts 50 \
|
||||
--max-concurrency 1 4 10 \
|
||||
--result-dir results/
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import time
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
import aiohttp
|
||||
import numpy as np
|
||||
from tqdm.asyncio import tqdm
|
||||
|
||||
PROMPTS = [
|
||||
"Hello, welcome to the voice synthesis benchmark test.",
|
||||
"She said she would be here by noon, but nobody showed up.",
|
||||
"The quick brown fox jumps over the lazy dog near the riverbank.",
|
||||
"I can't believe how beautiful the sunset looks from up here on the mountain.",
|
||||
"Please remember to bring your identification documents to the appointment tomorrow morning.",
|
||||
"Have you ever wondered what it would be like to travel through time and visit ancient civilizations?",
|
||||
"The restaurant on the corner serves the best pasta I have ever tasted in my entire life.",
|
||||
"After the meeting, we should discuss the quarterly results and plan for the next phase.",
|
||||
"Learning a new language takes patience, practice, and a genuine curiosity about other cultures.",
|
||||
"The train leaves at half past seven, so we need to arrive at the station before then.",
|
||||
"Could you please turn down the music a little bit, I'm trying to concentrate on my work.",
|
||||
"It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.",
|
||||
]
|
||||
|
||||
|
||||
@dataclass
|
||||
class RequestResult:
|
||||
success: bool = False
|
||||
ttfp: float = 0.0 # Time to first audio packet (seconds)
|
||||
e2e: float = 0.0 # End-to-end latency (seconds)
|
||||
audio_bytes: int = 0 # Total audio bytes received
|
||||
audio_duration: float = 0.0 # Audio duration in seconds (estimated from PCM)
|
||||
rtf: float = 0.0 # Real-time factor = e2e / audio_duration
|
||||
prompt: str = ""
|
||||
error: str = ""
|
||||
|
||||
|
||||
@dataclass
|
||||
class BenchmarkResult:
|
||||
config_name: str = ""
|
||||
concurrency: int = 0
|
||||
num_prompts: int = 0
|
||||
completed: int = 0
|
||||
failed: int = 0
|
||||
duration_s: float = 0.0
|
||||
# TTFP stats (ms)
|
||||
mean_ttfp_ms: float = 0.0
|
||||
median_ttfp_ms: float = 0.0
|
||||
std_ttfp_ms: float = 0.0
|
||||
p90_ttfp_ms: float = 0.0
|
||||
p95_ttfp_ms: float = 0.0
|
||||
p99_ttfp_ms: float = 0.0
|
||||
# E2E stats (ms)
|
||||
mean_e2e_ms: float = 0.0
|
||||
median_e2e_ms: float = 0.0
|
||||
std_e2e_ms: float = 0.0
|
||||
p90_e2e_ms: float = 0.0
|
||||
p95_e2e_ms: float = 0.0
|
||||
p99_e2e_ms: float = 0.0
|
||||
# RTF stats
|
||||
mean_rtf: float = 0.0
|
||||
median_rtf: float = 0.0
|
||||
std_rtf: float = 0.0
|
||||
p99_rtf: float = 0.0
|
||||
# Audio stats
|
||||
mean_audio_duration_s: float = 0.0
|
||||
total_audio_duration_s: float = 0.0
|
||||
audio_throughput: float = 0.0 # audio_duration / wall_time
|
||||
request_throughput: float = 0.0 # requests / second
|
||||
# Per-request details
|
||||
per_request: list = field(default_factory=list)
|
||||
|
||||
|
||||
def pcm_bytes_to_duration(num_bytes: int, sample_rate: int = 24000, sample_width: int = 2) -> float:
|
||||
"""Convert raw PCM byte count to duration in seconds."""
|
||||
num_samples = num_bytes / sample_width
|
||||
return num_samples / sample_rate
|
||||
|
||||
|
||||
async def send_tts_request(
|
||||
session: aiohttp.ClientSession,
|
||||
api_url: str,
|
||||
prompt: str,
|
||||
voice: str = "vivian",
|
||||
language: str = "English",
|
||||
pbar: tqdm | None = None,
|
||||
) -> RequestResult:
|
||||
"""Send a streaming TTS request and measure latency metrics."""
|
||||
payload = {
|
||||
"input": prompt,
|
||||
"voice": voice,
|
||||
"language": language,
|
||||
"stream": True,
|
||||
"response_format": "pcm",
|
||||
}
|
||||
|
||||
result = RequestResult(prompt=prompt)
|
||||
st = time.perf_counter()
|
||||
|
||||
try:
|
||||
async with session.post(api_url, json=payload) as response:
|
||||
if response.status != 200:
|
||||
result.error = f"HTTP {response.status}: {await response.text()}"
|
||||
result.success = False
|
||||
return result
|
||||
|
||||
first_chunk = True
|
||||
total_bytes = 0
|
||||
|
||||
async for chunk in response.content.iter_any():
|
||||
if first_chunk and len(chunk) > 0:
|
||||
result.ttfp = time.perf_counter() - st
|
||||
first_chunk = False
|
||||
total_bytes += len(chunk)
|
||||
|
||||
result.e2e = time.perf_counter() - st
|
||||
result.audio_bytes = total_bytes
|
||||
result.audio_duration = pcm_bytes_to_duration(total_bytes)
|
||||
|
||||
if result.audio_duration > 0:
|
||||
result.rtf = result.e2e / result.audio_duration
|
||||
result.success = True
|
||||
|
||||
except Exception as e:
|
||||
result.error = str(e)
|
||||
result.success = False
|
||||
result.e2e = time.perf_counter() - st
|
||||
|
||||
if pbar:
|
||||
pbar.update(1)
|
||||
return result
|
||||
|
||||
|
||||
async def run_benchmark(
|
||||
host: str,
|
||||
port: int,
|
||||
num_prompts: int,
|
||||
max_concurrency: int,
|
||||
num_warmups: int = 3,
|
||||
voice: str = "vivian",
|
||||
language: str = "English",
|
||||
) -> BenchmarkResult:
|
||||
"""Run benchmark at a given concurrency level."""
|
||||
api_url = f"http://{host}:{port}/v1/audio/speech"
|
||||
|
||||
connector = aiohttp.TCPConnector(
|
||||
limit=max_concurrency,
|
||||
limit_per_host=max_concurrency,
|
||||
keepalive_timeout=60,
|
||||
)
|
||||
session = aiohttp.ClientSession(
|
||||
connector=connector,
|
||||
timeout=aiohttp.ClientTimeout(total=600),
|
||||
)
|
||||
|
||||
# Warmup
|
||||
if num_warmups > 0:
|
||||
print(f" Warming up with {num_warmups} requests...")
|
||||
warmup_tasks = []
|
||||
for i in range(num_warmups):
|
||||
prompt = PROMPTS[i % len(PROMPTS)]
|
||||
warmup_tasks.append(send_tts_request(session, api_url, prompt, voice, language))
|
||||
await asyncio.gather(*warmup_tasks)
|
||||
print(" Warmup done.")
|
||||
|
||||
# Build request list
|
||||
request_prompts = [PROMPTS[i % len(PROMPTS)] for i in range(num_prompts)]
|
||||
|
||||
# Run benchmark
|
||||
print(f" Running {num_prompts} requests with concurrency={max_concurrency}...")
|
||||
semaphore = asyncio.Semaphore(max_concurrency)
|
||||
pbar = tqdm(total=num_prompts, desc=f" concurrency={max_concurrency}")
|
||||
|
||||
async def limited_request(prompt):
|
||||
async with semaphore:
|
||||
return await send_tts_request(session, api_url, prompt, voice, language, pbar)
|
||||
|
||||
start_time = time.perf_counter()
|
||||
tasks = [asyncio.create_task(limited_request(p)) for p in request_prompts]
|
||||
results: list[RequestResult] = await asyncio.gather(*tasks)
|
||||
duration = time.perf_counter() - start_time
|
||||
pbar.close()
|
||||
|
||||
await session.close()
|
||||
|
||||
# Compute stats
|
||||
successful = [r for r in results if r.success]
|
||||
failed = [r for r in results if not r.success]
|
||||
|
||||
bench = BenchmarkResult(
|
||||
concurrency=max_concurrency,
|
||||
num_prompts=num_prompts,
|
||||
completed=len(successful),
|
||||
failed=len(failed),
|
||||
duration_s=duration,
|
||||
)
|
||||
|
||||
if successful:
|
||||
ttfps = [r.ttfp * 1000 for r in successful] # convert to ms
|
||||
e2es = [r.e2e * 1000 for r in successful]
|
||||
rtfs = [r.rtf for r in successful]
|
||||
audio_durs = [r.audio_duration for r in successful]
|
||||
|
||||
bench.mean_ttfp_ms = float(np.mean(ttfps))
|
||||
bench.median_ttfp_ms = float(np.median(ttfps))
|
||||
bench.std_ttfp_ms = float(np.std(ttfps))
|
||||
bench.p90_ttfp_ms = float(np.percentile(ttfps, 90))
|
||||
bench.p95_ttfp_ms = float(np.percentile(ttfps, 95))
|
||||
bench.p99_ttfp_ms = float(np.percentile(ttfps, 99))
|
||||
|
||||
bench.mean_e2e_ms = float(np.mean(e2es))
|
||||
bench.median_e2e_ms = float(np.median(e2es))
|
||||
bench.std_e2e_ms = float(np.std(e2es))
|
||||
bench.p90_e2e_ms = float(np.percentile(e2es, 90))
|
||||
bench.p95_e2e_ms = float(np.percentile(e2es, 95))
|
||||
bench.p99_e2e_ms = float(np.percentile(e2es, 99))
|
||||
|
||||
bench.mean_rtf = float(np.mean(rtfs))
|
||||
bench.median_rtf = float(np.median(rtfs))
|
||||
bench.std_rtf = float(np.std(rtfs))
|
||||
bench.p99_rtf = float(np.percentile(rtfs, 99))
|
||||
|
||||
bench.mean_audio_duration_s = float(np.mean(audio_durs))
|
||||
bench.total_audio_duration_s = float(np.sum(audio_durs))
|
||||
bench.audio_throughput = bench.total_audio_duration_s / duration
|
||||
bench.request_throughput = len(successful) / duration
|
||||
|
||||
bench.per_request = [
|
||||
{
|
||||
"ttfp_ms": r.ttfp * 1000,
|
||||
"e2e_ms": r.e2e * 1000,
|
||||
"rtf": r.rtf,
|
||||
"audio_duration_s": r.audio_duration,
|
||||
"prompt": r.prompt,
|
||||
}
|
||||
for r in successful
|
||||
]
|
||||
|
||||
# Print summary in standardized performance template
|
||||
W = 50
|
||||
print("")
|
||||
print(f"{'=' * W}")
|
||||
print(f"{'Serving Benchmark Result':^{W}}")
|
||||
print(f"{'=' * W}")
|
||||
print(f"{'Successful requests:':<40}{bench.completed:<10}")
|
||||
print(f"{'Failed requests:':<40}{bench.failed:<10}")
|
||||
print(f"{'Maximum request concurrency:':<40}{max_concurrency:<10}")
|
||||
print(f"{'Benchmark duration (s):':<40}{duration:<10.2f}")
|
||||
print(f"{'Request throughput (req/s):':<40}{bench.request_throughput:<10.2f}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'End-to-end Latency':^{W}}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Mean E2EL (ms):':<40}{bench.mean_e2e_ms:<10.2f}")
|
||||
print(f"{'Median E2EL (ms):':<40}{bench.median_e2e_ms:<10.2f}")
|
||||
print(f"{'P99 E2EL (ms):':<40}{bench.p99_e2e_ms:<10.2f}")
|
||||
print(f"{'=' * W}")
|
||||
print(f"{'Audio Result':^{W}}")
|
||||
print(f"{'=' * W}")
|
||||
print(f"{'Total audio duration generated (s):':<40}{bench.total_audio_duration_s:<10.2f}")
|
||||
print(f"{'Audio throughput (audio duration/s):':<40}{bench.audio_throughput:<10.2f}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Time to First Packet':^{W}}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Mean AUDIO_TTFP (ms):':<40}{bench.mean_ttfp_ms:<10.2f}")
|
||||
print(f"{'Median AUDIO_TTFP (ms):':<40}{bench.median_ttfp_ms:<10.2f}")
|
||||
print(f"{'P99 AUDIO_TTFP (ms):':<40}{bench.p99_ttfp_ms:<10.2f}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Real Time Factor':^{W}}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Mean AUDIO_RTF:':<40}{bench.mean_rtf:<10.3f}")
|
||||
print(f"{'Median AUDIO_RTF:':<40}{bench.median_rtf:<10.3f}")
|
||||
print(f"{'P99 AUDIO_RTF:':<40}{bench.p99_rtf:<10.3f}")
|
||||
print(f"{'=' * W}")
|
||||
print("")
|
||||
|
||||
if failed:
|
||||
for r in failed[:3]:
|
||||
print(f" [ERROR] {r.error[:200]}")
|
||||
|
||||
return bench
|
||||
|
||||
|
||||
async def main(args):
|
||||
all_results = []
|
||||
|
||||
for concurrency in args.max_concurrency:
|
||||
result = await run_benchmark(
|
||||
host=args.host,
|
||||
port=args.port,
|
||||
num_prompts=args.num_prompts,
|
||||
max_concurrency=concurrency,
|
||||
num_warmups=args.num_warmups,
|
||||
voice=args.voice,
|
||||
language=args.language,
|
||||
)
|
||||
result.config_name = args.config_name
|
||||
all_results.append(asdict(result))
|
||||
|
||||
# Save results
|
||||
result_dir = Path(args.result_dir)
|
||||
result_dir.mkdir(parents=True, exist_ok=True)
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
result_file = result_dir / f"bench_{args.config_name}_{timestamp}.json"
|
||||
|
||||
with open(result_file, "w") as f:
|
||||
json.dump(all_results, f, indent=2)
|
||||
print(f"Results saved to {result_file}")
|
||||
|
||||
return all_results
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Qwen3-TTS Benchmark Client")
|
||||
parser.add_argument("--host", type=str, default="127.0.0.1")
|
||||
parser.add_argument("--port", type=int, default=8000)
|
||||
parser.add_argument("--num-prompts", type=int, default=50, help="Number of prompts per concurrency level")
|
||||
parser.add_argument( # noqa: E501
|
||||
"--max-concurrency", type=int, nargs="+", default=[1, 4, 10], help="Concurrency levels to test"
|
||||
)
|
||||
parser.add_argument("--num-warmups", type=int, default=3)
|
||||
parser.add_argument("--voice", type=str, default="vivian")
|
||||
parser.add_argument("--language", type=str, default="English")
|
||||
parser.add_argument(
|
||||
"--config-name", type=str, default="async_chunk", help="Label for this config (used in filenames)"
|
||||
)
|
||||
parser.add_argument("--result-dir", type=str, default="results")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
asyncio.run(main(args))
|
||||
@@ -0,0 +1,96 @@
|
||||
# Qwen3-TTS batch_size=1 config (streaming with async_chunk)
|
||||
# 2-stage pipeline: Talker -> Code2Wav
|
||||
async_chunk: true
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
stage_type: llm
|
||||
runtime:
|
||||
devices: "0"
|
||||
max_batch_size: 1
|
||||
engine_args:
|
||||
model_stage: qwen3_tts
|
||||
model_arch: Qwen3TTSTalkerForConditionalGeneration
|
||||
hf_overrides:
|
||||
architectures: [Qwen3TTSTalkerForConditionalGeneration]
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
enforce_eager: false
|
||||
trust_remote_code: true
|
||||
async_scheduling: false
|
||||
enable_prefix_caching: false
|
||||
engine_output_type: latent
|
||||
gpu_memory_utilization: 0.3
|
||||
distributed_executor_backend: "mp"
|
||||
max_num_batched_tokens: 512
|
||||
max_model_len: 4096
|
||||
custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_tts.talker2code2wav_async_chunk
|
||||
output_connectors:
|
||||
to_stage_1: connector_of_shared_memory
|
||||
default_sampling_params:
|
||||
temperature: 0.9
|
||||
top_k: 50
|
||||
max_tokens: 4096
|
||||
seed: 42
|
||||
detokenize: false
|
||||
repetition_penalty: 1.05
|
||||
stop_token_ids: [2150]
|
||||
|
||||
- stage_id: 1
|
||||
stage_type: llm
|
||||
runtime:
|
||||
devices: "0"
|
||||
max_batch_size: 1
|
||||
engine_args:
|
||||
model_stage: code2wav
|
||||
model_arch: Qwen3TTSCode2Wav
|
||||
hf_overrides:
|
||||
architectures: [Qwen3TTSCode2Wav]
|
||||
worker_type: generation
|
||||
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
async_scheduling: false
|
||||
enable_prefix_caching: false
|
||||
engine_output_type: audio
|
||||
gpu_memory_utilization: 0.2
|
||||
distributed_executor_backend: "mp"
|
||||
max_num_batched_tokens: 8192
|
||||
max_model_len: 32768
|
||||
engine_input_source: [0]
|
||||
final_output: true
|
||||
final_output_type: audio
|
||||
input_connectors:
|
||||
from_stage_0: connector_of_shared_memory
|
||||
tts_args:
|
||||
max_instructions_length: 500
|
||||
default_sampling_params:
|
||||
temperature: 0.0
|
||||
top_p: 1.0
|
||||
top_k: -1
|
||||
max_tokens: 65536
|
||||
seed: 42
|
||||
detokenize: true
|
||||
repetition_penalty: 1.0
|
||||
|
||||
runtime:
|
||||
enabled: true
|
||||
defaults:
|
||||
window_size: -1
|
||||
max_inflight: 1
|
||||
|
||||
connectors:
|
||||
connector_of_shared_memory:
|
||||
name: SharedMemoryConnector
|
||||
extra:
|
||||
shm_threshold_bytes: 65536
|
||||
codec_streaming: true
|
||||
connector_get_sleep_s: 0.01
|
||||
connector_get_max_wait_first_chunk: 3000
|
||||
connector_get_max_wait: 300
|
||||
codec_chunk_frames: 25
|
||||
codec_left_context_frames: 25
|
||||
|
||||
edges:
|
||||
- from: 0
|
||||
to: 1
|
||||
window_size: -1
|
||||
@@ -0,0 +1,97 @@
|
||||
# Qwen3-TTS batch_size=4 config (streaming with async_chunk)
|
||||
# Enables concurrent request processing with max_inflight=4
|
||||
# 2-stage pipeline: Talker -> Code2Wav
|
||||
async_chunk: true
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
stage_type: llm
|
||||
runtime:
|
||||
devices: "0"
|
||||
max_batch_size: 4
|
||||
engine_args:
|
||||
model_stage: qwen3_tts
|
||||
model_arch: Qwen3TTSTalkerForConditionalGeneration
|
||||
hf_overrides:
|
||||
architectures: [Qwen3TTSTalkerForConditionalGeneration]
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
enforce_eager: false
|
||||
trust_remote_code: true
|
||||
async_scheduling: false
|
||||
enable_prefix_caching: false
|
||||
engine_output_type: latent
|
||||
gpu_memory_utilization: 0.3
|
||||
distributed_executor_backend: "mp"
|
||||
max_num_batched_tokens: 512
|
||||
max_model_len: 4096
|
||||
custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_tts.talker2code2wav_async_chunk
|
||||
output_connectors:
|
||||
to_stage_1: connector_of_shared_memory
|
||||
default_sampling_params:
|
||||
temperature: 0.9
|
||||
top_k: 50
|
||||
max_tokens: 4096
|
||||
seed: 42
|
||||
detokenize: false
|
||||
repetition_penalty: 1.05
|
||||
stop_token_ids: [2150]
|
||||
|
||||
- stage_id: 1
|
||||
stage_type: llm
|
||||
runtime:
|
||||
devices: "0"
|
||||
max_batch_size: 4
|
||||
engine_args:
|
||||
model_stage: code2wav
|
||||
model_arch: Qwen3TTSCode2Wav
|
||||
hf_overrides:
|
||||
architectures: [Qwen3TTSCode2Wav]
|
||||
worker_type: generation
|
||||
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
async_scheduling: false
|
||||
enable_prefix_caching: false
|
||||
engine_output_type: audio
|
||||
gpu_memory_utilization: 0.2
|
||||
distributed_executor_backend: "mp"
|
||||
max_num_batched_tokens: 8192
|
||||
max_model_len: 32768
|
||||
engine_input_source: [0]
|
||||
final_output: true
|
||||
final_output_type: audio
|
||||
input_connectors:
|
||||
from_stage_0: connector_of_shared_memory
|
||||
tts_args:
|
||||
max_instructions_length: 500
|
||||
default_sampling_params:
|
||||
temperature: 0.0
|
||||
top_p: 1.0
|
||||
top_k: -1
|
||||
max_tokens: 65536
|
||||
seed: 42
|
||||
detokenize: true
|
||||
repetition_penalty: 1.0
|
||||
|
||||
runtime:
|
||||
enabled: true
|
||||
defaults:
|
||||
window_size: -1
|
||||
max_inflight: 4
|
||||
|
||||
connectors:
|
||||
connector_of_shared_memory:
|
||||
name: SharedMemoryConnector
|
||||
extra:
|
||||
shm_threshold_bytes: 65536
|
||||
codec_streaming: true
|
||||
connector_get_sleep_s: 0.01
|
||||
connector_get_max_wait_first_chunk: 3000
|
||||
connector_get_max_wait: 300
|
||||
codec_chunk_frames: 25
|
||||
codec_left_context_frames: 25
|
||||
|
||||
edges:
|
||||
- from: 0
|
||||
to: 1
|
||||
window_size: -1
|
||||
@@ -0,0 +1,249 @@
|
||||
"""Plot TTFP comparison: async_chunk off vs on.
|
||||
|
||||
Generates a bar chart with improvement arrows, matching the Qwen3-Omni
|
||||
async_chunk benchmark figure style.
|
||||
|
||||
Usage:
|
||||
python plot_async_chunk.py \
|
||||
--off results/bench_async_chunk_off_*.json \
|
||||
--on results/bench_async_chunk_on_*.json \
|
||||
--output results/qwen3_tts_async_chunk_ttfp.png
|
||||
|
||||
# Also supports E2E and RTF metrics:
|
||||
python plot_async_chunk.py \
|
||||
--off results/bench_async_chunk_off_*.json \
|
||||
--on results/bench_async_chunk_on_*.json \
|
||||
--metric e2e \
|
||||
--output results/qwen3_tts_async_chunk_e2e.png
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
|
||||
METRIC_CONFIG = {
|
||||
"ttfp": {
|
||||
"key": "mean_ttfp_ms",
|
||||
"ylabel": "TTFP (s)",
|
||||
"title": "TTFP (Time to First Audio Packet) - Qwen3-TTS, by concurrency",
|
||||
"to_seconds": True,
|
||||
},
|
||||
"e2e": {
|
||||
"key": "mean_e2e_ms",
|
||||
"ylabel": "E2E (s)",
|
||||
"title": "E2E Latency - Qwen3-TTS, by concurrency",
|
||||
"to_seconds": True,
|
||||
},
|
||||
"rtf": {
|
||||
"key": "mean_rtf",
|
||||
"ylabel": "RTF",
|
||||
"title": "Real-Time Factor - Qwen3-TTS, by concurrency",
|
||||
"to_seconds": False,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def load_results(path: str) -> list[dict]:
|
||||
with open(path) as f:
|
||||
return json.load(f)
|
||||
|
||||
|
||||
def plot_ttfp_comparison(
|
||||
off_results: list[dict],
|
||||
on_results: list[dict],
|
||||
metric: str,
|
||||
output_path: str,
|
||||
title_override: str | None = None,
|
||||
):
|
||||
cfg = METRIC_CONFIG[metric]
|
||||
key = cfg["key"]
|
||||
to_seconds = cfg["to_seconds"]
|
||||
|
||||
off_map = {r["concurrency"]: r for r in off_results}
|
||||
on_map = {r["concurrency"]: r for r in on_results}
|
||||
concurrencies = sorted(set(off_map.keys()) & set(on_map.keys()))
|
||||
|
||||
off_vals = []
|
||||
on_vals = []
|
||||
for c in concurrencies:
|
||||
v_off = off_map[c][key]
|
||||
v_on = on_map[c][key]
|
||||
if to_seconds:
|
||||
v_off /= 1000.0
|
||||
v_on /= 1000.0
|
||||
off_vals.append(v_off)
|
||||
on_vals.append(v_on)
|
||||
|
||||
fig, ax = plt.subplots(figsize=(8, 6))
|
||||
|
||||
x = np.arange(len(concurrencies))
|
||||
width = 0.3
|
||||
|
||||
ax.bar(x - width / 2, off_vals, width, label="async_chunk off", color="#87CEEB", edgecolor="none")
|
||||
ax.bar(x + width / 2, on_vals, width, label="async_chunk on", color="#FFF8DC", edgecolor="#DDD8B8")
|
||||
|
||||
# Draw improvement arrows and labels
|
||||
for i in range(len(concurrencies)):
|
||||
v_off = off_vals[i]
|
||||
v_on = on_vals[i]
|
||||
if v_on > 0:
|
||||
improvement = v_off / v_on
|
||||
else:
|
||||
improvement = float("inf")
|
||||
|
||||
# Arrow from top of off-bar to top of on-bar
|
||||
arrow_start_x = x[i] - width / 2
|
||||
arrow_start_y = v_off * 0.95
|
||||
arrow_end_x = x[i] + width / 2
|
||||
arrow_end_y = v_on * 1.05
|
||||
|
||||
ax.annotate(
|
||||
"",
|
||||
xy=(arrow_end_x, arrow_end_y),
|
||||
xytext=(arrow_start_x, arrow_start_y),
|
||||
arrowprops=dict(arrowstyle="->", color="red", lw=1.5),
|
||||
)
|
||||
|
||||
# Improvement label
|
||||
label_x = (arrow_start_x + arrow_end_x) / 2
|
||||
label_y = arrow_start_y + (v_off - v_on) * 0.15
|
||||
ax.text(
|
||||
label_x,
|
||||
label_y,
|
||||
f"{improvement:.1f}x improvement",
|
||||
ha="center",
|
||||
va="bottom",
|
||||
fontsize=10,
|
||||
color="red",
|
||||
fontweight="bold",
|
||||
)
|
||||
|
||||
title = title_override or cfg["title"]
|
||||
ax.set_title(title, fontsize=13, fontweight="bold")
|
||||
ax.set_ylabel(cfg["ylabel"], fontsize=12)
|
||||
ax.set_xlabel("Max concurrency", fontsize=12)
|
||||
ax.set_xticks(x)
|
||||
ax.set_xticklabels([str(c) for c in concurrencies])
|
||||
ax.set_yscale("log")
|
||||
ax.legend(loc="upper left", fontsize=11)
|
||||
ax.grid(axis="y", alpha=0.3, linestyle="--")
|
||||
ax.set_axisbelow(True)
|
||||
|
||||
plt.tight_layout()
|
||||
Path(output_path).parent.mkdir(parents=True, exist_ok=True)
|
||||
plt.savefig(output_path, dpi=150, bbox_inches="tight")
|
||||
print(f"Plot saved to {output_path}")
|
||||
plt.close()
|
||||
|
||||
|
||||
def plot_all_metrics(off_results: list[dict], on_results: list[dict], output_path: str):
|
||||
"""Generate a 1x3 subplot with TTFP, E2E, and RTF comparisons."""
|
||||
off_map = {r["concurrency"]: r for r in off_results}
|
||||
on_map = {r["concurrency"]: r for r in on_results}
|
||||
concurrencies = sorted(set(off_map.keys()) & set(on_map.keys()))
|
||||
|
||||
fig, axes = plt.subplots(1, 3, figsize=(18, 6))
|
||||
fig.suptitle("Qwen3-TTS: async_chunk on vs off", fontsize=15, fontweight="bold")
|
||||
|
||||
for ax, metric in zip(axes, ["ttfp", "e2e", "rtf"]):
|
||||
cfg = METRIC_CONFIG[metric]
|
||||
key = cfg["key"]
|
||||
to_seconds = cfg["to_seconds"]
|
||||
|
||||
off_vals = []
|
||||
on_vals = []
|
||||
for c in concurrencies:
|
||||
v_off = off_map[c][key]
|
||||
v_on = on_map[c][key]
|
||||
if to_seconds:
|
||||
v_off /= 1000.0
|
||||
v_on /= 1000.0
|
||||
off_vals.append(v_off)
|
||||
on_vals.append(v_on)
|
||||
|
||||
x = np.arange(len(concurrencies))
|
||||
width = 0.3
|
||||
ax.bar(x - width / 2, off_vals, width, label="async_chunk off", color="#87CEEB")
|
||||
ax.bar(x + width / 2, on_vals, width, label="async_chunk on", color="#FFF8DC", edgecolor="#DDD8B8")
|
||||
|
||||
for i in range(len(concurrencies)):
|
||||
if on_vals[i] > 0:
|
||||
improvement = off_vals[i] / on_vals[i]
|
||||
ax.annotate(
|
||||
"",
|
||||
xy=(x[i] + width / 2, on_vals[i] * 1.05),
|
||||
xytext=(x[i] - width / 2, off_vals[i] * 0.95),
|
||||
arrowprops=dict(arrowstyle="->", color="red", lw=1.5),
|
||||
)
|
||||
label_y = off_vals[i] * 0.85
|
||||
ax.text(x[i], label_y, f"{improvement:.1f}x", ha="center", fontsize=10, color="red", fontweight="bold")
|
||||
|
||||
ax.set_title(cfg["title"].split(" - ")[0], fontsize=12, fontweight="bold")
|
||||
ax.set_ylabel(cfg["ylabel"], fontsize=11)
|
||||
ax.set_xlabel("Max concurrency", fontsize=11)
|
||||
ax.set_xticks(x)
|
||||
ax.set_xticklabels([str(c) for c in concurrencies])
|
||||
if metric != "rtf":
|
||||
ax.set_yscale("log")
|
||||
ax.legend(fontsize=9)
|
||||
ax.grid(axis="y", alpha=0.3, linestyle="--")
|
||||
ax.set_axisbelow(True)
|
||||
|
||||
plt.tight_layout()
|
||||
Path(output_path).parent.mkdir(parents=True, exist_ok=True)
|
||||
plt.savefig(output_path, dpi=150, bbox_inches="tight")
|
||||
print(f"Plot saved to {output_path}")
|
||||
plt.close()
|
||||
|
||||
|
||||
def print_table(off_results: list[dict], on_results: list[dict]):
|
||||
off_map = {r["concurrency"]: r for r in off_results}
|
||||
on_map = {r["concurrency"]: r for r in on_results}
|
||||
concurrencies = sorted(set(off_map.keys()) & set(on_map.keys()))
|
||||
|
||||
print("\n## Benchmark Results: async_chunk off vs on\n")
|
||||
print("| Metric | Concurrency | async_chunk off | async_chunk on | Improvement |")
|
||||
print("| --- | --- | --- | --- | --- |")
|
||||
|
||||
for name, key, fmt in [
|
||||
("TTFP (ms)", "mean_ttfp_ms", ".1f"),
|
||||
("E2E (ms)", "mean_e2e_ms", ".1f"),
|
||||
("RTF", "mean_rtf", ".3f"),
|
||||
("Throughput", "audio_throughput", ".2f"),
|
||||
]:
|
||||
for c in concurrencies:
|
||||
v_off = off_map[c].get(key, 0)
|
||||
v_on = on_map[c].get(key, 0)
|
||||
if v_on > 0 and key != "audio_throughput":
|
||||
ratio = f"{v_off / v_on:.1f}x"
|
||||
elif v_off > 0 and key == "audio_throughput":
|
||||
ratio = f"{v_on / v_off:.1f}x"
|
||||
else:
|
||||
ratio = "N/A"
|
||||
print(f"| {name} | {c} | {v_off:{fmt}} | {v_on:{fmt}} | {ratio} |")
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Plot async_chunk comparison for Qwen3-TTS")
|
||||
parser.add_argument("--off", type=str, required=True, help="JSON results for async_chunk off")
|
||||
parser.add_argument("--on", type=str, required=True, help="JSON results for async_chunk on")
|
||||
parser.add_argument("--metric", type=str, default="ttfp", choices=["ttfp", "e2e", "rtf", "all"])
|
||||
parser.add_argument("--output", type=str, default="results/qwen3_tts_async_chunk.png")
|
||||
parser.add_argument("--title", type=str, default=None, help="Custom title override")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
off_results = load_results(args.off)
|
||||
on_results = load_results(args.on)
|
||||
|
||||
print_table(off_results, on_results)
|
||||
|
||||
if args.metric == "all":
|
||||
plot_all_metrics(off_results, on_results, args.output)
|
||||
else:
|
||||
plot_ttfp_comparison(off_results, on_results, args.metric, args.output, args.title)
|
||||
@@ -0,0 +1,167 @@
|
||||
#!/bin/bash
|
||||
# Qwen3-TTS async_chunk on vs off Benchmark
|
||||
#
|
||||
# Starts two servers (async_chunk on and off), benchmarks both,
|
||||
# and generates comparison plots.
|
||||
#
|
||||
# Usage:
|
||||
# bash run_async_chunk_benchmark.sh
|
||||
#
|
||||
# Environment variables:
|
||||
# GPU_DEVICE - GPU index (default: 0)
|
||||
# MODEL - Model path (default: Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice)
|
||||
# NUM_PROMPTS - Prompts per concurrency level (default: 50)
|
||||
# CONCURRENCY - Space-separated concurrency levels (default: "1 10")
|
||||
# PORT_ON - Port for async_chunk on server (default: 8000)
|
||||
# PORT_OFF - Port for async_chunk off server (default: 8001)
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
PROJECT_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)"
|
||||
cd "$PROJECT_ROOT"
|
||||
|
||||
GPU_DEVICE="${GPU_DEVICE:-0}"
|
||||
MODEL="${MODEL:-Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice}"
|
||||
NUM_PROMPTS="${NUM_PROMPTS:-50}"
|
||||
CONCURRENCY="${CONCURRENCY:-1 10}"
|
||||
NUM_WARMUPS="${NUM_WARMUPS:-3}"
|
||||
PORT_ON="${PORT_ON:-8000}"
|
||||
PORT_OFF="${PORT_OFF:-8001}"
|
||||
RESULT_DIR="${SCRIPT_DIR}/results"
|
||||
TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
|
||||
|
||||
STAGE_CONFIG_ON="vllm_omni/model_executor/stage_configs/qwen3_tts.yaml"
|
||||
STAGE_CONFIG_OFF="vllm_omni/model_executor/stage_configs/qwen3_tts_no_async_chunk.yaml"
|
||||
|
||||
mkdir -p "${RESULT_DIR}"
|
||||
|
||||
echo "============================================================"
|
||||
echo " Qwen3-TTS async_chunk Benchmark"
|
||||
echo "============================================================"
|
||||
echo " GPU: ${GPU_DEVICE}"
|
||||
echo " Model: ${MODEL}"
|
||||
echo " Prompts: ${NUM_PROMPTS}"
|
||||
echo " Concurrency: ${CONCURRENCY}"
|
||||
echo " Port (on/off): ${PORT_ON} / ${PORT_OFF}"
|
||||
echo " Results: ${RESULT_DIR}"
|
||||
echo "============================================================"
|
||||
|
||||
cleanup() {
|
||||
echo "Cleaning up servers..."
|
||||
kill "$PID_ON" 2>/dev/null || true
|
||||
kill "$PID_OFF" 2>/dev/null || true
|
||||
wait "$PID_ON" 2>/dev/null || true
|
||||
wait "$PID_OFF" 2>/dev/null || true
|
||||
}
|
||||
trap cleanup EXIT
|
||||
|
||||
wait_for_server() {
|
||||
local port=$1
|
||||
local name=$2
|
||||
local max_wait=300
|
||||
local elapsed=0
|
||||
echo "Waiting for ${name} server on port ${port}..."
|
||||
while ! curl -s "http://localhost:${port}/health" >/dev/null 2>&1; do
|
||||
sleep 5
|
||||
elapsed=$((elapsed + 5))
|
||||
if [ $elapsed -ge $max_wait ]; then
|
||||
echo "ERROR: ${name} server failed to start within ${max_wait}s"
|
||||
exit 1
|
||||
fi
|
||||
done
|
||||
echo "${name} server ready (${elapsed}s)"
|
||||
}
|
||||
|
||||
# ---- Phase 1: Start async_chunk ON server ----
|
||||
echo ""
|
||||
echo "[Phase 1] Starting async_chunk ON server on port ${PORT_ON}..."
|
||||
CUDA_VISIBLE_DEVICES=${GPU_DEVICE} vllm-omni serve "${MODEL}" \
|
||||
--stage-configs-path "${STAGE_CONFIG_ON}" \
|
||||
--host 0.0.0.0 --port "${PORT_ON}" \
|
||||
--trust-remote-code --enforce-eager --omni \
|
||||
> "${RESULT_DIR}/server_on_${TIMESTAMP}.log" 2>&1 &
|
||||
PID_ON=$!
|
||||
|
||||
wait_for_server "${PORT_ON}" "async_chunk_on"
|
||||
|
||||
echo "[Phase 1] Benchmarking async_chunk ON..."
|
||||
# shellcheck disable=SC2086
|
||||
python "${SCRIPT_DIR}/bench_async_chunk.py" \
|
||||
--host 127.0.0.1 --port "${PORT_ON}" \
|
||||
--config-name "async_chunk_on" \
|
||||
--num-prompts "${NUM_PROMPTS}" \
|
||||
--max-concurrency ${CONCURRENCY} \
|
||||
--num-warmups "${NUM_WARMUPS}" \
|
||||
--result-dir "${RESULT_DIR}"
|
||||
|
||||
echo "[Phase 1] Stopping async_chunk ON server..."
|
||||
kill "$PID_ON" 2>/dev/null || true
|
||||
wait "$PID_ON" 2>/dev/null || true
|
||||
sleep 5
|
||||
|
||||
# ---- Phase 2: Start async_chunk OFF server ----
|
||||
echo ""
|
||||
echo "[Phase 2] Starting async_chunk OFF server on port ${PORT_OFF}..."
|
||||
CUDA_VISIBLE_DEVICES=${GPU_DEVICE} vllm-omni serve "${MODEL}" \
|
||||
--stage-configs-path "${STAGE_CONFIG_OFF}" \
|
||||
--host 0.0.0.0 --port "${PORT_OFF}" \
|
||||
--trust-remote-code --enforce-eager --omni \
|
||||
> "${RESULT_DIR}/server_off_${TIMESTAMP}.log" 2>&1 &
|
||||
PID_OFF=$!
|
||||
|
||||
wait_for_server "${PORT_OFF}" "async_chunk_off"
|
||||
|
||||
echo "[Phase 2] Benchmarking async_chunk OFF (non-streaming)..."
|
||||
# shellcheck disable=SC2086
|
||||
python "${SCRIPT_DIR}/bench_async_chunk.py" \
|
||||
--host 127.0.0.1 --port "${PORT_OFF}" \
|
||||
--config-name "async_chunk_off" \
|
||||
--num-prompts "${NUM_PROMPTS}" \
|
||||
--max-concurrency ${CONCURRENCY} \
|
||||
--num-warmups "${NUM_WARMUPS}" \
|
||||
--no-stream \
|
||||
--result-dir "${RESULT_DIR}"
|
||||
|
||||
echo "[Phase 2] Stopping async_chunk OFF server..."
|
||||
kill "$PID_OFF" 2>/dev/null || true
|
||||
wait "$PID_OFF" 2>/dev/null || true
|
||||
|
||||
# ---- Phase 3: Plot results ----
|
||||
echo ""
|
||||
echo "[Phase 3] Generating plots..."
|
||||
|
||||
# Find the latest result files
|
||||
RESULT_ON=$(ls -t "${RESULT_DIR}"/bench_async_chunk_on_*.json 2>/dev/null | head -1)
|
||||
RESULT_OFF=$(ls -t "${RESULT_DIR}"/bench_async_chunk_off_*.json 2>/dev/null | head -1)
|
||||
|
||||
if [ -z "$RESULT_ON" ] || [ -z "$RESULT_OFF" ]; then
|
||||
echo "ERROR: Could not find result files. Check logs in ${RESULT_DIR}/"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo " ON results: ${RESULT_ON}"
|
||||
echo " OFF results: ${RESULT_OFF}"
|
||||
|
||||
# TTFP comparison (main figure)
|
||||
python "${SCRIPT_DIR}/plot_async_chunk.py" \
|
||||
--off "${RESULT_OFF}" \
|
||||
--on "${RESULT_ON}" \
|
||||
--metric ttfp \
|
||||
--output "${RESULT_DIR}/qwen3_tts_async_chunk_ttfp.png"
|
||||
|
||||
# All metrics comparison
|
||||
python "${SCRIPT_DIR}/plot_async_chunk.py" \
|
||||
--off "${RESULT_OFF}" \
|
||||
--on "${RESULT_ON}" \
|
||||
--metric all \
|
||||
--output "${RESULT_DIR}/qwen3_tts_async_chunk_all.png"
|
||||
|
||||
echo ""
|
||||
echo "============================================================"
|
||||
echo " Benchmark complete!"
|
||||
echo " Results: ${RESULT_DIR}/"
|
||||
echo " Plots:"
|
||||
echo " - ${RESULT_DIR}/qwen3_tts_async_chunk_ttfp.png"
|
||||
echo " - ${RESULT_DIR}/qwen3_tts_async_chunk_all.png"
|
||||
echo "============================================================"
|
||||
@@ -1,14 +1,13 @@
|
||||
ARG VLLM_BASE_IMAGE=vllm/vllm-openai
|
||||
ARG VLLM_BASE_TAG=v0.16.0
|
||||
ARG VLLM_BASE_TAG=v0.17.0
|
||||
FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
|
||||
ARG APP_DIR=/workspace/vllm-omni
|
||||
WORKDIR ${APP_DIR}
|
||||
|
||||
COPY . .
|
||||
|
||||
# Install system dependencies
|
||||
RUN apt-get update && \
|
||||
apt-get install -y ffmpeg git sox libsox-fmt-all jq && \
|
||||
apt-get install -y espeak-ng ffmpeg git sox libsox-fmt-all jq && \
|
||||
apt-get clean && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.16.0
|
||||
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.17.0
|
||||
FROM ${BASE_IMAGE} AS base
|
||||
|
||||
# Declare a variable to know if we want to use the nightly build or the stable build.
|
||||
@@ -19,7 +19,7 @@ WORKDIR ${COMMON_WORKDIR}
|
||||
|
||||
# Step 1: Setup - Install system dependencies
|
||||
RUN apt-get update && \
|
||||
apt-get install -y ffmpeg git sox libsox-fmt-all jq && \
|
||||
apt-get install -y espeak-ng ffmpeg git sox libsox-fmt-all jq && \
|
||||
apt-get clean && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
|
||||
@@ -9,12 +9,12 @@ ARG PYTHON_VERSION=3.12
|
||||
ARG PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/xpu"
|
||||
|
||||
RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | gpg --dearmor | tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null && \
|
||||
echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list && \
|
||||
add-apt-repository -y ppa:kobuk-team/intel-graphics
|
||||
echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list
|
||||
|
||||
RUN apt clean && apt-get update -y && \
|
||||
apt-get install -y --no-install-recommends --fix-missing \
|
||||
curl \
|
||||
espeak-ng \
|
||||
ffmpeg \
|
||||
git \
|
||||
libsndfile1 \
|
||||
@@ -31,9 +31,22 @@ RUN apt clean && apt-get update -y && \
|
||||
python3-pip
|
||||
|
||||
RUN apt update && apt upgrade -y && \
|
||||
apt install -y libze1 libze-dev libze-intel-gpu1 intel-opencl-icd libze-intel-gpu-raytracing intel-ocloc && \
|
||||
apt install -y intel-oneapi-compiler-dpcpp-cpp-2025.3
|
||||
|
||||
# Install UMD with fixed version
|
||||
RUN mkdir neo && \
|
||||
cd neo && \
|
||||
wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.24.8/intel-igc-core-2_2.24.8+20344_amd64.deb && \
|
||||
wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.24.8/intel-igc-opencl-2_2.24.8+20344_amd64.deb && \
|
||||
wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/intel-ocloc_25.48.36300.8-0_amd64.deb && \
|
||||
wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/intel-opencl-icd_25.48.36300.8-0_amd64.deb && \
|
||||
wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/libigdgmm12_22.8.2_amd64.deb && \
|
||||
wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/libze-intel-gpu1_25.48.36300.8-0_amd64.deb && \
|
||||
wget https://github.com/oneapi-src/level-zero/releases/download/v1.26.0/level-zero_1.26.0+u24.04_amd64.deb && \
|
||||
dpkg -i *.deb && \
|
||||
cd .. && \
|
||||
rm -rf neo
|
||||
|
||||
ENV PATH="/root/.local/bin:$PATH"
|
||||
ENV VIRTUAL_ENV="/opt/venv"
|
||||
ENV UV_PYTHON_INSTALL_DIR=/opt/uv/python
|
||||
|
||||
@@ -66,6 +66,7 @@ nav:
|
||||
- Model Implementation:
|
||||
- contributing/model/README.md
|
||||
- contributing/model/adding_omni_model.md
|
||||
- contributing/model/adding_tts_model.md
|
||||
- contributing/model/adding_diffusion_model.md
|
||||
- CI: contributing/ci
|
||||
- Design Documents:
|
||||
|
||||
@@ -60,5 +60,5 @@ vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, in
|
||||
|
||||
For more information, checkout the following:
|
||||
|
||||
- [vllm-omni architecture design and recent roadmaps](https://docs.google.com/presentation/d/1qv4qMW1rKAqDREMXiUDLIgqqHQe7TDPj/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true)
|
||||
- [vllm-omni architecture design and recent roadmaps](https://docs.google.com/presentation/d/1XJWgv79lORl8rbaVvp2d5Sqs6ZEBgAgj/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true)
|
||||
- [vllm-omni announcement blogpost](https://blog.vllm.ai/2025/11/30/vllm-omni.html)
|
||||
|
Before Width: | Height: | Size: 127 KiB After Width: | Height: | Size: 98 KiB |
@@ -2,7 +2,7 @@
|
||||
|
||||
This section lists the most common options for running vLLM-Omni.
|
||||
|
||||
For options within a vLLM Engine. Please refer to [vLLM Configuration](https://docs.vllm.ai/en/v0.14.0/configuration/index.html)
|
||||
For options within a vLLM Engine. Please refer to [vLLM Configuration](https://docs.vllm.ai/en/v0.16.0/configuration/index.html)
|
||||
|
||||
Currently, the main options are maintained by stage configs for each model.
|
||||
|
||||
@@ -19,3 +19,4 @@ For introduction, please check [Introduction for stage config](./stage_configs.m
|
||||
- **[TeaCache Configuration](../user_guide/diffusion/teacache.md)** - Enable TeaCache adaptive caching for DiT models to achieve 1.5x-2.0x speedup with minimal quality loss
|
||||
- **[Cache-DiT Configuration](../user_guide/diffusion/cache_dit_acceleration.md)** - Enable Cache-DiT as cache acceleration backends for DiT models
|
||||
- **[Parallelism Configuration](../user_guide/diffusion/parallelism_acceleration.md)** - Enable parallelism (e.g., sequence parallelism) for for DiT models
|
||||
- **[CPU Offloading](../user_guide/diffusion/cpu_offload_diffusion.md)** - Enable CPU offloading (model-level and layerwise) for for DiT models
|
||||
@@ -135,6 +135,14 @@ Each stage in the `stage_args` list contains the following configuration options
|
||||
|
||||
A unique identifier for each stage in the multi-stage pipeline. Stages are numbered sequentially starting from 0, and this ID is used to reference stages in inter-stage dependencies (e.g., `engine_input_source`).
|
||||
|
||||
### `prompt_expand_func` (Optional)
|
||||
|
||||
A custom Python function hook for the LLM stage (Stage 0) that expands a single incoming prompt object into multiple prompts. This is primarily used for multi-modal Classifier-Free Guidance (CFG), where it generates the necessary companion requests (like a negative text prompt) and tags them with internal roles (e.g., `cfg_text`). This ensures the upstream LLM generates the needed contextual hidden states for both the conditional and unconditional generations simultaneously.
|
||||
|
||||
### `cfg_kv_collect_func` (Optional)
|
||||
|
||||
A custom Python function hook for downstream diffusion stages (Stage 1+) to collect, map, and process the KV caches transferred from the companion requests fired by `prompt_expand_func`. It aggregates the hidden condition states cleanly (e.g., binding them as `cfg_text_past_key_values` and `cfg_text_kv_metadata`), allowing the diffusion runtime to perform CFG smoothly without redundantly evaluating text paths on the DiT workers.
|
||||
|
||||
### `runtime`
|
||||
|
||||
Configuration for disaggregated execution of the stage, controlling how the stage is deployed and executed.
|
||||
|
||||
@@ -71,14 +71,7 @@ For additional features and advanced configurations, refer to the:
|
||||
### Testing
|
||||
|
||||
vLLM-Omni uses `pytest` to test the codebase.
|
||||
|
||||
```bash
|
||||
# Run all tests
|
||||
pytest tests/
|
||||
|
||||
# Run tests for a single test file with detailed output
|
||||
pytest -s -v tests/test_omni_llm.py
|
||||
```
|
||||
Please refer to the [test instructions](./ci/test_guide.md) for detailed testing information.
|
||||
|
||||
!!! warning
|
||||
Currently, not all unit tests pass when run on CPU platforms. If you don't have access to a GPU platform to run unit tests locally, rely on the continuous integration system to run the tests for now.
|
||||
@@ -123,6 +116,9 @@ Only specific types of PRs will be reviewed. The PR title is prefixed appropriat
|
||||
!!! note
|
||||
If the PR spans more than one category, please include all relevant prefixes.
|
||||
|
||||
### Local Test
|
||||
Please run the L1 and L2 test cases locally first and attach the results before contacting us to add the "ready" label. Please refer to the [test instructions](./ci/test_guide.md) for running the test cases.
|
||||
|
||||
### Code Quality
|
||||
|
||||
The PR needs to meet the following code quality standards:
|
||||
|
||||
@@ -67,8 +67,8 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
|
||||
<td>Online & Offline (multiple deployment scenarios):<br>real model, normal inference function, normal accuracy</td>
|
||||
<td><30min</td>
|
||||
<td>
|
||||
/tests/e2e/online_serving/test_{model_name}_expansion.py<br>
|
||||
/tests/e2e/offline_inference/test_{model_name}_expansion.py
|
||||
/tests/e2e/online_serving/test_{model_name}.py<br>
|
||||
/tests/e2e/offline_inference/test_{model_name}.py
|
||||
</td>
|
||||
<td>
|
||||
<a href="#chapter-2-l3-level-testing-core-integration-performance-and-accuracy-verification">Chapter 2</a><br>
|
||||
@@ -508,7 +508,9 @@ L3 level testing executes after code is merged into the main branch. Its core pu
|
||||
**2.4.5 Request Execution**
|
||||
|
||||
```python
|
||||
openai_client.send_request(request_config, request_num=1)
|
||||
openai_client.send_omni_request(request_config, request_num=1) # for omni-understanding models
|
||||
# or
|
||||
openai_client.send_diffusion_request(request_config, request_num=1) # for diffusion models
|
||||
```
|
||||
|
||||
**Explanation**:
|
||||
@@ -517,7 +519,7 @@ L3 level testing executes after code is merged into the main branch. Its core pu
|
||||
|
||||
**Single Request**: The comment clearly states this is a single-request completion test. For concurrent testing, it can be extended to multiple requests using request_num = n.
|
||||
|
||||
**Implicit Validation**: The send_request method internally includes validation logic dynamically selected based on the --run-level parameter: core_model performs basic validation, while advanced_model performs deep validation.
|
||||
**Implicit Validation**: The `send_omni_request` and `send_diffusion_request` methods internally includes validation logic dynamically selected based on the --run-level parameter: core_model performs basic validation, while advanced_model performs deep validation.
|
||||
|
||||
- ***Run Command***: `pytest -s -v /tests/e2e/online_serving/test_{model_name}.py -m advanced_model --run-level=advanced_model`
|
||||
|
||||
|
||||
@@ -0,0 +1,75 @@
|
||||
# Test Guide
|
||||
## Setting Up the Test Environment
|
||||
### Creating a Container
|
||||
vLLM-Omni provides an official Docker image for deployment. These images are built upon vLLM Docker images and are available on [Docker Hub](https://hub.docker.com/r/vllm/vllm-omni/tags). The version of vLLM-Omni indicates which vLLM release it is based on.
|
||||
For a local test environment, you can follow the steps below to create a container:
|
||||
## Installing Dependencies
|
||||
### vLLM & vLLM-Omni
|
||||
vLLM-Omni is built based on vLLM. You can follow [install guide](../../getting_started/installation/README.md) to build your local environment.
|
||||
|
||||
### Test Case Dependencies
|
||||
When running test cases, you may need to install the following dependencies:
|
||||
|
||||
```bash
|
||||
uv pip install ".[dev]"
|
||||
apt-get install -y ffmpeg
|
||||
```
|
||||
|
||||
## Running Tests
|
||||
Our test scripts use the pytest framework. First, please use `git clone https://github.com/vllm-project/vllm-omni.git` to download the vllm-omni source code. Then, in the root directory of vllm-omni, you can run the following commands in your local test environment to execute the corresponding test cases.
|
||||
|
||||
=== "L1 level"
|
||||
|
||||
```bash
|
||||
cd tests
|
||||
pytest -s -v -m "core_model and cpu"
|
||||
```
|
||||
The latest test command is available in the "Simple Unit Test" step of this [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-ready.yml).
|
||||
|
||||
=== "L2 level"
|
||||
|
||||
```bash
|
||||
cd tests
|
||||
pytest -s -v -m "core_model and not cpu" --run-level=core_model
|
||||
```
|
||||
If you only want to run a specific test case, you can use:
|
||||
```bash
|
||||
pytest -s -v test_xxxx.py --run-level=core_model
|
||||
```
|
||||
If you only want to run specific test cases on a particular platform, you can use:
|
||||
```bash
|
||||
pytest -s -v -m "core_model and distributed_cuda and L4" --run-level=core_model
|
||||
```
|
||||
The latest test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-ready.yml).
|
||||
|
||||
=== "L3 level & L4 level"
|
||||
|
||||
```bash
|
||||
cd tests
|
||||
pytest -s -v -m "advanced_model" --run-level=advanced_model
|
||||
```
|
||||
If you only want to run L3 test case, you can use:
|
||||
```bash
|
||||
pytest -s -v e2e/ --ignore-glob='*expansion.py' -m "advanced_model" --run-level=advanced_model
|
||||
```
|
||||
If you only want to run a specific test case, you can use:
|
||||
```bash
|
||||
pytest -s -v test_xxxx.py --run-level=advanced_model
|
||||
```
|
||||
If you only want to run specific test cases on a particular platform, you can use:
|
||||
```bash
|
||||
pytest -s -v -m "core_model and distributed_cuda and L4" --run-level=core_model
|
||||
```
|
||||
Note: To run performance tests, use:
|
||||
```bash
|
||||
pytest -s -v perf/scripts/run_benchmark.py
|
||||
```
|
||||
|
||||
The latest L3 test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-merge.yml).
|
||||
|
||||
The latest L4 test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-nightly.yml).
|
||||
|
||||
You can find more information about markers in the documentation: [marker doc](./tests_markers.md)
|
||||
|
||||
## Adding New Test Cases
|
||||
Please refer to the [L5 Layering Specification document](./CI_5levels.md).
|
||||
@@ -32,7 +32,7 @@ Defined in `pyproject.toml`:
|
||||
| `slow` | Slow tests (may skip in quick CI) |
|
||||
| `benchmark` | Benchmark tests |
|
||||
|
||||
\* Means those markers are auto-added, and they will be added by the `@hardware_test` decorator.
|
||||
\* Means those markers are auto-added by `@hardware_test` (parametrization decorator) or `hardware_marks` (only returning the list of marks for flexibility).
|
||||
|
||||
### Example usage for markers
|
||||
|
||||
@@ -49,6 +49,7 @@ from tests.utils import hardware_test
|
||||
def test_video_to_audio()
|
||||
...
|
||||
```
|
||||
|
||||
### Decorator: `@hardware_test`
|
||||
|
||||
This decorator is intended to make hardware-aware, cross-platform test authoring easier and more robust for CI/CD environments. The `hardware_test` decorator in `vllm-omni/tests/utils.py` performs the following actions:
|
||||
@@ -98,6 +99,24 @@ This decorator is intended to make hardware-aware, cross-platform test authoring
|
||||
- ROCm only: `pytest -m "distributed_rocm and MI325"`
|
||||
- NPU only: `pytest -m "distributed_npu"`
|
||||
|
||||
### Function: `hardware_marks`
|
||||
|
||||
`hardware_marks` returns a list of pytest mark objects with the same signature as `@hardware_test`. Use it when you need more flexibility, such as attaching hardware marks to individual `pytest.param` entries rather than an entire test function.
|
||||
|
||||
```python
|
||||
from tests.utils import hardware_marks
|
||||
|
||||
MULTI_CARD_MARKS = hardware_marks(
|
||||
res={"cuda": "H100", "rocm": "MI325", "npu": "A2"}, num_cards=2
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize("omni_server", [
|
||||
pytest.param(OmniServerParams(...), id="case_001", marks=MULTI_CARD_MARKS),
|
||||
], indirect=True)
|
||||
def test_feature(omni_server):
|
||||
...
|
||||
```
|
||||
|
||||
## Add Support for a New Platform
|
||||
|
||||
If you want to add support for a new platform (e.g., "tpu" for a new accelerator), follow these steps:
|
||||
@@ -132,12 +151,13 @@ If you want to add support for a new platform (e.g., "tpu" for a new accelerator
|
||||
# Optionally: add skipif_tpu when implemented
|
||||
return [test_platform, test_resource, test_distributed]
|
||||
```
|
||||
3. **Update `hardware_test` to recognize your new platform**:
|
||||
In the relevant place (see the `hardware_test` implementation), add:
|
||||
3. **Update `hardware_marks` to recognize your new platform**:
|
||||
In the relevant place (see the `hardware_marks` implementation), add:
|
||||
```python
|
||||
if platform == "tpu":
|
||||
marks = tpu_marks(res=resource, num_cards=cards)
|
||||
```
|
||||
(`hardware_test` calls `hardware_marks` internally, so both will pick up the change.)
|
||||
4. **(Recommended) Add a test using your new markers**:
|
||||
```python
|
||||
@hardware_test(
|
||||
@@ -151,7 +171,7 @@ If you want to add support for a new platform (e.g., "tpu" for a new accelerator
|
||||
**Summary**:
|
||||
- Add pytest markers for your new platform/resources
|
||||
- Implement a marker function (`xxx_marks`)
|
||||
- Plug into `hardware_test`
|
||||
- You're done: tests decorated with `@hardware_test` using your platform now automatically get the correct markers, distribution, and isolation!
|
||||
- Plug into `hardware_marks`
|
||||
- You're done: tests using `@hardware_test` or `hardware_marks` with your platform now automatically get the correct markers, distribution, and isolation!
|
||||
|
||||
See code in `vllm-omni/tests/utils.py` for existing examples (`cuda_marks`, `rocm_marks`, `npu_marks`).
|
||||
@@ -158,6 +158,7 @@ vllm_omni/ tests/
|
||||
├── test_diffusion_layerwise_offload.py
|
||||
├── test_diffusion_lora.py
|
||||
├── test_sequence_parallel.py
|
||||
├── test_qwen_image_edit_expansion.py
|
||||
└── stage_configs/
|
||||
├── qwen2_5_omni_ci.yaml
|
||||
├── qwen3_omni_ci.yaml
|
||||
|
||||
@@ -651,8 +651,27 @@ For a fair comparison, keep the same **prompt**, **seed**, **resolution**, **num
|
||||
|
||||
#### 5.2 Functionality Check in CI
|
||||
|
||||
To ensure project maintainability and sustainable development, we encourage contributors to submit test code (unit tests, system tests, or end-to-end tests) alongside their code changes.
|
||||
For comprehensive testing guidelines, please refer to the [Test File Structure and Style Guide](../ci/tests_style.md).
|
||||
To ensure project maintainability and sustainable development, please submit test code (unit tests, system tests, or end-to-end tests) alongside their code changes.
|
||||
|
||||
For comprehensive testing guidelines and the definition of test levels (L1-L5), please refer to the [Test File Structure and Style Guide](../ci/tests_style.md).
|
||||
The following tests are required to add:
|
||||
|
||||
- L4 test of the model's full *functionality* (i.e., all the *diffusion features* that are supported by this model), including several [parallelism acceleration methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism_acceleration/), [CPU offloading](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cpu_offload_diffusion/), [TeaCache](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/teacache/) and [Cache-DiT](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cache_dit_acceleration/) cache backends, [quantization methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/quantization/overview/).
|
||||
- Test cases: Currently all the features are available in online serving mode. Hence, only need to add `tests/e2e/online_serving/test_{model}_expansion.py`. The following test cases shall cover all features:
|
||||
- 1 GPU: TeaCache & GGUF (or fallback to FP8, or disable it) & Layer-wise CPU offloading (or fallback to Module-wise)
|
||||
- 2 GPUs: Cache-DiT & FP8 (or fallback to GGUF, or disable it) & Ulysses = 2
|
||||
- 2 GPUs: Cache-DiT & GGUF (or fallback to FP8, or disable it) & Ring = 2
|
||||
- 2 GPUs: TeaCache & FP8 (or fallback to GGUF, or disable it) & CFG Parallel = 2
|
||||
- 2 GPUs: Cache-DiT & FP8 (or fallback to GGUF, or disable it) & Tensor Parallel = 2 & VAE Patch Parallel = 2
|
||||
- 2 GPUs: Cache-DiT & GGUF (or fallback to FP8, or disable it) & HSDP = 2 & VAE Patch Parallel = 2
|
||||
- Validation: test that the multimodal output files of your model have the correct shapes.
|
||||
- Test marks: always add `advanced_model` and `diffusion`. Add `parallel` and GPU-related marks if needed. Ref: [Markers for Tests](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/tests_markers/)
|
||||
- To maximize code reuse, you may refer to
|
||||
- `tests/conftest.py` for `omni_server` and `openai_client` fixtures, `generate_synthetic_image` and `assert_XXX_valid` helper.
|
||||
- `tests/utils.py` for `@hardware_test(...)` and `hardware_marks`.
|
||||
- [Parametrizing tests (pytest doc)](https://docs.pytest.org/en/stable/example/parametrize.html) to reuse test function implementation for different cases.
|
||||
- Doc: add a concise dostring for each test function.
|
||||
- Reference L4 test implementation: [tests/e2e/online_serving/test_qwen_image_edit_expansion.py](https://github.com/vllm-project/vllm-omni/blob/main/tests/e2e/online_serving/test_qwen_image_edit_expansion.py).
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -460,7 +460,7 @@ def forward(self, ...):
|
||||
These keys are then accessible in your stage transition function:
|
||||
```python
|
||||
# In stage_input_processors/qwen3_omni.py
|
||||
thinker_embeddings = output.multimodal_output["0"] # Access by key
|
||||
thinker_prefill_embeddings = output.multimodal_output["0"] # Access by key
|
||||
thinker_hidden_states = output.multimodal_output["24"]
|
||||
```
|
||||
|
||||
@@ -513,11 +513,11 @@ def thinker2talker(
|
||||
for thinker_output in thinker_outputs:
|
||||
output = thinker_output.outputs[0]
|
||||
# Extract thinker embeddings and hidden states
|
||||
thinker_embeddings = output.multimodal_output["0"].float().clone().detach().cuda()
|
||||
thinker_prefill_embeddings = output.multimodal_output["0"].float().clone().detach().cuda()
|
||||
thinker_hidden_states = output.multimodal_output["24"].float().clone().detach().cuda()
|
||||
|
||||
info = {
|
||||
"thinker_embeddings": thinker_embeddings,
|
||||
"thinker_prefill_embeddings": thinker_prefill_embeddings,
|
||||
"thinker_hidden_states": thinker_hidden_states,
|
||||
"thinker_sequences": thinker_output.prompt_token_ids + output.token_ids,
|
||||
"thinker_input_ids": thinker_output.prompt_token_ids,
|
||||
|
||||
@@ -0,0 +1,584 @@
|
||||
# Adding a TTS Model
|
||||
|
||||
This guide walks through adding a new TTS model to vLLM-Omni, using **Qwen3-TTS**
|
||||
as a reference. Qwen3-TTS demonstrates the standard two-stage TTS pipeline and the
|
||||
key optimizations all TTS models in this repo should follow.
|
||||
|
||||
## Table of Contents
|
||||
|
||||
1. [Overview](#overview)
|
||||
2. [Directory Structure](#directory-structure)
|
||||
3. [Step-by-Step Implementation](#step-by-step-implementation)
|
||||
4. [Key Components](#key-components)
|
||||
5. [Model Registration](#model-registration)
|
||||
6. [Stage Configuration](#stage-configuration)
|
||||
7. [Stage Input Processors](#stage-input-processors)
|
||||
8. [Testing](#testing)
|
||||
9. [Summary](#summary)
|
||||
|
||||
## Overview
|
||||
|
||||
vLLM-Omni supports TTS models as multi-stage pipelines where each stage runs independently
|
||||
and can be placed on different devices. Qwen3-TTS has two stages:
|
||||
|
||||
| Stage | Name | Input | Output |
|
||||
|-------|------|-------|--------|
|
||||
| 0 | Code Predictor (AR) | Text tokens | Discrete RVQ codec codes |
|
||||
| 1 | Code2Wav (Decoder) | RVQ codec codes | Audio waveform |
|
||||
|
||||
Each stage is a separate model class configured independently via YAML. The two stages
|
||||
are connected by the `async_chunk` framework, which enables inter-stage streaming for
|
||||
low first-packet latency (see [Async Chunk Design](../../design/feature/async_chunk_design.md)).
|
||||
|
||||
### Without async_chunk (batch mode)
|
||||
|
||||
Stage 0 runs to completion before Stage 1 starts, resulting in long first-packet latency:
|
||||
|
||||
```mermaid
|
||||
flowchart TB
|
||||
subgraph stage0["Stage 0: AR Code Predictor"]
|
||||
direction LR
|
||||
P[Prefill] --> D1[Decode 1]
|
||||
D1 --> D2[Decode 2]
|
||||
D2 --> Dots1["..."]
|
||||
Dots1 --> DN[Decode N]
|
||||
end
|
||||
|
||||
subgraph stage1["Stage 1: Code2Wav"]
|
||||
direction LR
|
||||
DEC[Decode all codes at once]
|
||||
end
|
||||
|
||||
stage0 -- "all N codes" --> stage1
|
||||
stage1 --> FPL["First Packet Latency = Stage 0 + Stage 1"]
|
||||
|
||||
style stage0 fill:#dae8fc,stroke:#6c8ebf
|
||||
style stage1 fill:#f8d7c8,stroke:#d4856a
|
||||
style FPL fill:#e8f0fe,stroke:#3366CC,stroke-width:2px
|
||||
```
|
||||
|
||||
### With async_chunk (streaming mode)
|
||||
|
||||
Stage 0 sends codec codes to Stage 1 every `chunk_size=25` tokens. Stage 1 begins decoding
|
||||
immediately, reducing first-packet latency from the full AR time to just the first chunk:
|
||||
|
||||
```mermaid
|
||||
flowchart TB
|
||||
subgraph stage0["Stage 0: Code Predictor (AR)"]
|
||||
direction LR
|
||||
P[Prefill] --> D1["Decode 1-25"]
|
||||
D1 --> D2["Decode 26-50"]
|
||||
D2 --> Dots1["..."]
|
||||
Dots1 --> DN["Decode N"]
|
||||
end
|
||||
|
||||
subgraph stage1["Stage 1: Code2Wav"]
|
||||
direction LR
|
||||
C1["Chunk 1\n(25 frames)"] --> C2["Chunk 2\n(context + 25)"]
|
||||
C2 --> Dots2["..."]
|
||||
Dots2 --> CN["Final chunk"]
|
||||
end
|
||||
|
||||
D1 -. "chunk 1 (25 codes)" .-> C1
|
||||
D2 -. "chunk 2 (context + 25)" .-> C2
|
||||
DN -. "final" .-> CN
|
||||
|
||||
stage0 --> FPL["⏱ First Packet Latency = Prefill + 25 decode steps only"]
|
||||
|
||||
style stage0 fill:#dae8fc,stroke:#6c8ebf
|
||||
style stage1 fill:#e8d4f8,stroke:#8a6cad
|
||||
style FPL fill:#e8f0fe,stroke:#3366CC,stroke-width:2px
|
||||
```
|
||||
|
||||
Key parameters: `chunk_size=25`, `left_context_size=25` (validated defaults from Qwen3-TTS
|
||||
and Qwen3-Omni).
|
||||
|
||||
## Directory Structure
|
||||
|
||||
When adding a new TTS model, create the following structure:
|
||||
|
||||
```
|
||||
vllm_omni/model_executor/models/
|
||||
your_model_name/
|
||||
__init__.py
|
||||
your_model.py # Unified class (stage dispatch)
|
||||
your_model_ar_stage.py # Stage 0: AR stage
|
||||
your_model_decoder.py # Stage 1: audio decoder
|
||||
|
||||
vllm_omni/model_executor/stage_input_processors/
|
||||
your_model_name.py # Stage 0 -> Stage 1 transition
|
||||
|
||||
vllm_omni/model_executor/stage_configs/
|
||||
your_model_name.yaml # Batch mode config
|
||||
your_model_name_async_chunk.yaml # Streaming mode config
|
||||
```
|
||||
|
||||
**Qwen3-TTS reference files:**
|
||||
|
||||
| File | Purpose |
|
||||
|------|---------|
|
||||
| `models/qwen3_tts/qwen3_tts.py` | Unified model class |
|
||||
| `models/qwen3_tts/qwen3_tts_code_predictor_vllm.py` | Stage 0 - optimized AR |
|
||||
| `models/qwen3_tts/qwen3_tts_code2wav.py` | Stage 1 - decoder |
|
||||
| `stage_configs/qwen3_tts.yaml` | Stage config (async_chunk enabled) |
|
||||
| `stage_configs/qwen3_tts_batch.yaml` | Batch mode config |
|
||||
| `stage_input_processors/qwen3_tts.py` | Stage transition processors |
|
||||
|
||||
## Step-by-Step Implementation
|
||||
|
||||
### Step 1: Implement Stage 0 - AR Stage
|
||||
|
||||
Stage 0 is the autoregressive stage that generates intermediate audio representations.
|
||||
**It must use vLLM's native decoder layers with fused ops and PagedAttention** for the LLM
|
||||
backbone - this is the primary source of speedup over HuggingFace inference.
|
||||
|
||||
#### 1.1 Use vLLM Decoder Layers Directly
|
||||
|
||||
Build your transformer layers from the corresponding vLLM decoder layer class (e.g.
|
||||
`Qwen3DecoderLayer` for Qwen3-based backbones, or the equivalent for LLaMA, Qwen2, etc.).
|
||||
Do not wrap the HuggingFace model directly - that bypasses PagedAttention and fused kernels.
|
||||
|
||||
```python
|
||||
# your_model_ar_stage.py
|
||||
|
||||
from vllm.model_executor.models.qwen3 import Qwen3DecoderLayer
|
||||
|
||||
class YourTTSARStage(nn.Module):
|
||||
|
||||
def __init__(self, config, vllm_config, prefix):
|
||||
self.layers = nn.ModuleList([
|
||||
Qwen3DecoderLayer(
|
||||
config, vllm_config=vllm_config, prefix=f"{prefix}.layers.{i}"
|
||||
)
|
||||
for i in range(config.num_hidden_layers)
|
||||
])
|
||||
self.lm_head = ParallelLMHead(config.codec_size, config.hidden_size)
|
||||
```
|
||||
|
||||
See `qwen3_tts_code_predictor_vllm.py` for the full implementation.
|
||||
|
||||
#### 1.2 Forward Pass
|
||||
|
||||
Implement `forward()` to return an `OmniOutput` with intermediate data for Stage 1:
|
||||
|
||||
```python
|
||||
def forward(self, input_ids, positions, intermediate_tensors=None,
|
||||
inputs_embeds=None, **kwargs) -> OmniOutput:
|
||||
hidden_states = self.run_layers(input_ids, positions, intermediate_tensors, inputs_embeds)
|
||||
logits = self.lm_head(hidden_states)
|
||||
|
||||
return OmniOutput(
|
||||
text_hidden_states=hidden_states,
|
||||
multimodal_outputs={
|
||||
"audio_codes": self.extract_codes(logits),
|
||||
},
|
||||
)
|
||||
```
|
||||
|
||||
The keys in `multimodal_outputs` are what your stage input processor will read to build
|
||||
Stage 1 inputs.
|
||||
|
||||
#### 1.3 Weight Loading with Fused QKV
|
||||
|
||||
When using vLLM's fused `QKVParallelLinear`, pack the HF `q_proj`/`k_proj`/`v_proj` weights
|
||||
into `qkv_proj` using `stacked_params_mapping`. See the `load_weights()` method in
|
||||
`qwen3_tts_code_predictor_vllm.py` for the standard pattern - it can be reused as-is
|
||||
for any Qwen-family backbone.
|
||||
|
||||
#### 1.4 Custom Stop Condition (if needed)
|
||||
|
||||
Some TTS models use a learned stop head rather than an EOS token. If your model does this,
|
||||
implement it inside `sample()`:
|
||||
|
||||
```python
|
||||
def sample(self, logits, sampling_metadata) -> SamplerOutput | None:
|
||||
output = self.sampler(logits, sampling_metadata)
|
||||
if self._stop_head_fired():
|
||||
output = mark_as_finished(output)
|
||||
return output
|
||||
```
|
||||
|
||||
### Step 2: Implement Stage 1 - Decoder
|
||||
|
||||
Stage 1 decodes Stage 0 output into audio. It runs outside the scheduler (no PagedAttention
|
||||
needed). Implement `chunked_decode_streaming()` to support async_chunk streaming:
|
||||
|
||||
```python
|
||||
# your_model_decoder.py
|
||||
|
||||
class YourTTSDecoder(nn.Module):
|
||||
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
||||
super().__init__()
|
||||
# Initialize your audio decoder (SpeechTokenizer, HiFiGAN, etc.)
|
||||
|
||||
def forward(self, codes: torch.Tensor, **kwargs) -> torch.Tensor:
|
||||
return self.decoder(codes)
|
||||
|
||||
def chunked_decode_streaming(self, codes, chunk_size=25,
|
||||
left_context_size=25) -> torch.Tensor:
|
||||
"""Decode with a sliding context window for smooth chunk boundaries."""
|
||||
end_index = codes.shape[-1]
|
||||
context_size = 0 if end_index <= chunk_size else left_context_size
|
||||
wav_chunk = self(codes)
|
||||
# Trim left context to avoid duplicate audio
|
||||
return wav_chunk[..., context_size * self.total_upsample:]
|
||||
```
|
||||
|
||||
### Step 3: Implement the Unified Model Class
|
||||
|
||||
The unified class dispatches to the correct stage based on `model_stage` in the config:
|
||||
|
||||
```python
|
||||
# your_model.py
|
||||
|
||||
class YourTTSModelForConditionalGeneration(nn.Module, SupportsPP):
|
||||
|
||||
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
|
||||
super().__init__()
|
||||
self.model_stage = vllm_config.model_config.model_stage
|
||||
|
||||
if self.model_stage == "ar_stage":
|
||||
ar_vllm_config = vllm_config.with_hf_config(
|
||||
vllm_config.model_config.hf_config.ar_config,
|
||||
architectures=["YourTTSARStageForConditionalGeneration"],
|
||||
)
|
||||
self.ar_stage = init_vllm_registered_model(
|
||||
vllm_config=ar_vllm_config,
|
||||
prefix=maybe_prefix(prefix, "ar"),
|
||||
hf_config=ar_vllm_config.model_config.hf_config,
|
||||
architectures=["YourTTSARStageForConditionalGeneration"],
|
||||
)
|
||||
self.model = self.ar_stage
|
||||
|
||||
elif self.model_stage == "decoder":
|
||||
self.decoder = YourTTSDecoder(vllm_config=vllm_config, prefix=prefix)
|
||||
self.model = self.decoder
|
||||
```
|
||||
|
||||
### Step 4: Create `__init__.py`
|
||||
|
||||
```python
|
||||
# vllm_omni/model_executor/models/your_model_name/__init__.py
|
||||
from .your_model import YourTTSModelForConditionalGeneration
|
||||
|
||||
__all__ = ["YourTTSModelForConditionalGeneration"]
|
||||
```
|
||||
|
||||
## Key Components
|
||||
|
||||
### Model Interfaces
|
||||
|
||||
Your unified model class should implement the appropriate interfaces:
|
||||
|
||||
- **`SupportsPP`**: Required for pipeline parallelism support (all models should implement this)
|
||||
- **`SupportsMultiModal`**: Only if your model accepts multimodal inputs (e.g. reference audio for voice cloning)
|
||||
|
||||
### Output Format
|
||||
|
||||
Use `OmniOutput` so the orchestrator can route intermediate data between stages:
|
||||
|
||||
```python
|
||||
from vllm_omni.model_executor.models.output_templates import OmniOutput
|
||||
|
||||
return OmniOutput(
|
||||
text_hidden_states=hidden_states,
|
||||
multimodal_outputs={
|
||||
"audio_codes": codec_codes,
|
||||
},
|
||||
)
|
||||
```
|
||||
|
||||
### Weight Loading from a Single Checkpoint
|
||||
|
||||
If both stages load from one checkpoint, separate them by prefix in the unified class:
|
||||
|
||||
```python
|
||||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||||
ar_weights, decoder_weights = [], []
|
||||
for name, tensor in weights:
|
||||
if name.startswith("decoder."):
|
||||
decoder_weights.append((name, tensor))
|
||||
else:
|
||||
ar_weights.append((name, tensor))
|
||||
|
||||
if self.model_stage == "ar_stage":
|
||||
return self.ar_stage.load_weights(ar_weights)
|
||||
elif self.model_stage == "decoder":
|
||||
return self.decoder.load_weights(decoder_weights)
|
||||
```
|
||||
|
||||
## Model Registration
|
||||
|
||||
Register all stage classes in `vllm_omni/model_executor/models/registry.py`:
|
||||
|
||||
```python
|
||||
_OMNI_MODELS = {
|
||||
# (package_name, module_name, class_name)
|
||||
"YourTTSModelForConditionalGeneration": (
|
||||
"your_model_name", "your_model",
|
||||
"YourTTSModelForConditionalGeneration",
|
||||
),
|
||||
"YourTTSARStageForConditionalGeneration": (
|
||||
"your_model_name", "your_model_ar_stage",
|
||||
"YourTTSARStageForConditionalGeneration",
|
||||
),
|
||||
"YourTTSDecoder": (
|
||||
"your_model_name", "your_model_decoder",
|
||||
"YourTTSDecoder",
|
||||
),
|
||||
}
|
||||
```
|
||||
|
||||
The registry uses lazy loading - model classes are only imported when needed.
|
||||
|
||||
## Stage Configuration
|
||||
|
||||
Each stage has a `worker_type` that determines how it is scheduled:
|
||||
|
||||
- `worker_type: ar` - autoregressive stage, uses `OmniARScheduler` with PagedAttention
|
||||
- `worker_type: generation` - non-AR stage (e.g. decoder), uses `OmniGenerationScheduler`
|
||||
|
||||
Key configuration fields:
|
||||
|
||||
| Field | Description |
|
||||
|-------|-------------|
|
||||
| `model_stage` | Which stage to initialize (`ar_stage`, `decoder`, etc.) |
|
||||
| `model_arch` | Architecture name, must match `registry.py` |
|
||||
| `engine_input_source` | List of upstream stage IDs that provide input (e.g. `[0]`) |
|
||||
| `engine_output_type` | Output type: `latent` for intermediate, `audio` for final |
|
||||
| `custom_process_next_stage_input_func` | Async chunk processor function path (streaming only) |
|
||||
| `final_output` | Whether this stage produces the final user-facing output |
|
||||
| `final_output_type` | Type of final output (`audio`, `text`, etc.) |
|
||||
|
||||
### Batch mode
|
||||
|
||||
```yaml
|
||||
# stage_configs/your_model_name.yaml
|
||||
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
stage_type: llm
|
||||
runtime:
|
||||
devices: "0"
|
||||
max_batch_size: 64
|
||||
engine_args:
|
||||
model_stage: ar_stage
|
||||
model_arch: YourTTSModelForConditionalGeneration
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
engine_output_type: latent
|
||||
default_sampling_params:
|
||||
temperature: 0.9
|
||||
top_k: 50
|
||||
max_tokens: 2048
|
||||
|
||||
- stage_id: 1
|
||||
stage_type: llm
|
||||
runtime:
|
||||
devices: "0"
|
||||
engine_args:
|
||||
model_stage: decoder
|
||||
model_arch: YourTTSModelForConditionalGeneration
|
||||
worker_type: generation
|
||||
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
|
||||
engine_output_type: audio
|
||||
engine_input_source: [0]
|
||||
final_output: true
|
||||
final_output_type: audio
|
||||
```
|
||||
|
||||
### Streaming mode (async_chunk)
|
||||
|
||||
Add `async_chunk: true` at the top level and specify `custom_process_next_stage_input_func`
|
||||
on Stage 0 to define how intermediate outputs are chunked and forwarded:
|
||||
|
||||
```yaml
|
||||
# stage_configs/your_model_name_async_chunk.yaml
|
||||
|
||||
async_chunk: true
|
||||
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
stage_type: llm
|
||||
runtime:
|
||||
devices: "0"
|
||||
max_batch_size: 64
|
||||
engine_args:
|
||||
model_stage: ar_stage
|
||||
model_arch: YourTTSModelForConditionalGeneration
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
engine_output_type: latent
|
||||
custom_process_next_stage_input_func: >
|
||||
vllm_omni.model_executor.stage_input_processors.your_model_name.ar2decoder_async_chunk
|
||||
default_sampling_params:
|
||||
temperature: 0.9
|
||||
top_k: 50
|
||||
max_tokens: 2048
|
||||
|
||||
- stage_id: 1
|
||||
stage_type: llm
|
||||
runtime:
|
||||
devices: "0"
|
||||
engine_args:
|
||||
model_stage: decoder
|
||||
model_arch: YourTTSModelForConditionalGeneration
|
||||
worker_type: generation
|
||||
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
|
||||
engine_output_type: audio
|
||||
engine_input_source: [0]
|
||||
final_output: true
|
||||
final_output_type: audio
|
||||
```
|
||||
|
||||
## Stage Input Processors
|
||||
|
||||
Stage input processors convert Stage 0 outputs into Stage 1 inputs. Create yours in
|
||||
`vllm_omni/model_executor/stage_input_processors/your_model_name.py`.
|
||||
|
||||
See `stage_input_processors/qwen3_tts.py` for the full reference implementation.
|
||||
|
||||
### Data structures
|
||||
|
||||
Understanding what's available in stage outputs:
|
||||
|
||||
- `stage_list[source_id].engine_outputs` - list of `EngineCoreOutput` objects
|
||||
- Each `EngineCoreOutput` has `outputs` - list of `RequestOutput` objects
|
||||
- Each `RequestOutput` has:
|
||||
- `token_ids` - generated token IDs
|
||||
- `multimodal_output` - dict with keys matching your model's `OmniOutput.multimodal_outputs`
|
||||
- `prompt_token_ids` - original prompt token IDs
|
||||
|
||||
### Batch mode (non-streaming)
|
||||
|
||||
Collects all Stage 0 outputs and forwards them to Stage 1 in one shot:
|
||||
|
||||
```python
|
||||
def ar2decoder(
|
||||
stage_list: list[Any],
|
||||
engine_input_source: list[int],
|
||||
prompt: OmniTokensPrompt | TextPrompt | None = None,
|
||||
requires_multimodal_data: bool = False,
|
||||
) -> list[OmniTokensPrompt]:
|
||||
source_id = engine_input_source[0]
|
||||
decoder_inputs = []
|
||||
|
||||
for output in stage_list[source_id].engine_outputs:
|
||||
result = output.outputs[0]
|
||||
codes = result.multimodal_output["audio_codes"].cpu()
|
||||
decoder_inputs.append(
|
||||
OmniTokensPrompt(prompt_token_ids=codes.reshape(-1).tolist())
|
||||
)
|
||||
|
||||
return decoder_inputs
|
||||
```
|
||||
|
||||
### Streaming mode (async_chunk)
|
||||
|
||||
Buffers Stage 0 outputs and forwards a chunk to Stage 1 once `chunk_size` frames
|
||||
have accumulated. The function signature follows the `OmniChunkTransferAdapter` protocol:
|
||||
|
||||
```python
|
||||
def ar2decoder_async_chunk(
|
||||
transfer_manager: Any,
|
||||
pooling_output: dict[str, Any] | None,
|
||||
request: Any,
|
||||
is_finished: bool = False,
|
||||
) -> dict[str, Any] | None:
|
||||
"""Forward chunks of AR output to the decoder stage."""
|
||||
request_id = request.external_req_id
|
||||
finished = bool(is_finished or request.is_finished())
|
||||
|
||||
# Extract and buffer the latest frame
|
||||
if isinstance(pooling_output, dict):
|
||||
frame = extract_frame(pooling_output)
|
||||
if frame is not None:
|
||||
transfer_manager.code_prompt_token_ids[request_id].append(
|
||||
frame.cpu().tolist()
|
||||
)
|
||||
elif not finished:
|
||||
return None
|
||||
|
||||
# Read chunk config from connector
|
||||
chunk_size = 25
|
||||
left_context_size = 25
|
||||
|
||||
length = len(transfer_manager.code_prompt_token_ids[request_id])
|
||||
if length <= 0:
|
||||
if finished:
|
||||
return {"codes": [], "finished": torch.tensor(True, dtype=torch.bool)}
|
||||
return None
|
||||
|
||||
# Wait until a full chunk is ready (or request is finished)
|
||||
chunk_length = length % chunk_size
|
||||
if chunk_length != 0 and not finished:
|
||||
return None
|
||||
|
||||
# Build context window: left_context + chunk
|
||||
context_length = chunk_length if chunk_length != 0 else chunk_size
|
||||
end_index = min(length, left_context_size + context_length)
|
||||
window = transfer_manager.code_prompt_token_ids[request_id][-end_index:]
|
||||
|
||||
return {
|
||||
"codes": torch.tensor(window).transpose(0, 1).reshape(-1).tolist(),
|
||||
"left_context_size": max(0, int(end_index - context_length)),
|
||||
"finished": torch.tensor(finished, dtype=torch.bool),
|
||||
}
|
||||
```
|
||||
|
||||
Key points:
|
||||
- `transfer_manager` is the `OmniChunkTransferAdapter` that owns the chunk lifecycle
|
||||
- Each call appends one AR decode step's output; a chunk is emitted every `chunk_size` steps
|
||||
- The final (possibly partial) chunk is flushed when `is_finished` is true
|
||||
- `left_context_size` frames of overlap are included for smooth audio boundaries
|
||||
|
||||
## Testing
|
||||
|
||||
For general testing conventions, see [tests_style.md](../ci/tests_style.md).
|
||||
|
||||
Recommended test cases for a new TTS model:
|
||||
|
||||
1. **Single request** - verify waveform output shape and sample rate
|
||||
2. **Batched requests** - verify each request in the batch finishes independently
|
||||
3. **async_chunk streaming** - verify audio chunks arrive incrementally and decode correctly
|
||||
4. **Speaker conditioning** (if applicable) - verify different speaker inputs produce different outputs
|
||||
|
||||
Reference test: `tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py`
|
||||
|
||||
## Adding a Model Recipe
|
||||
|
||||
After implementing and testing your model, add a model recipe to the
|
||||
[vllm-project/recipes](https://github.com/vllm-project/recipes) repository so users can
|
||||
get started quickly. See [Adding an Omni-Modality Model](./adding_omni_model.md#adding-a-model-recipe)
|
||||
for the expected format.
|
||||
|
||||
## Summary
|
||||
|
||||
Adding a TTS model to vLLM-Omni involves:
|
||||
|
||||
1. **Create model directory** with AR stage, decoder stage, and unified class
|
||||
2. **AR stage** - use vLLM's native decoder layers with fused QKV; do not wrap HF directly
|
||||
3. **Decoder stage** - thin wrapper around your audio decoder; implement `chunked_decode_streaming()`
|
||||
4. **Unified class** - dispatches on `model_stage`; same structure as `Qwen3TTSModelForGeneration`
|
||||
5. **Register** all stage classes in `registry.py`
|
||||
6. **YAML configs** - provide both batch and `async_chunk` variants
|
||||
7. **Stage input processor** - buffer Stage 0 outputs and forward in chunks of 25
|
||||
8. **Tests** - cover single request, batching, and async_chunk streaming
|
||||
9. **Model recipe** - add to [vllm-project/recipes](https://github.com/vllm-project/recipes)
|
||||
|
||||
### Qwen3-TTS Reference Files
|
||||
|
||||
| File | Purpose |
|
||||
|------|---------|
|
||||
| `models/qwen3_tts/qwen3_tts.py` | Unified model class |
|
||||
| `models/qwen3_tts/qwen3_tts_code_predictor_vllm.py` | AR stage with vLLM fused ops |
|
||||
| `models/qwen3_tts/qwen3_tts_code2wav.py` | Decoder stage with `chunked_decode_streaming()` |
|
||||
| `stage_configs/qwen3_tts.yaml` | Stage configuration |
|
||||
| `stage_input_processors/qwen3_tts.py` | Stage transition processors |
|
||||
|
||||
For more information, see:
|
||||
|
||||
- [Architecture Overview](../../design/architecture_overview.md)
|
||||
- [Async Chunk Design](../../design/feature/async_chunk_design.md)
|
||||
- [Stage Configuration Guide](../../configuration/stage_configs.md)
|
||||
@@ -92,6 +92,13 @@ The framework achieves high performance through several optimization techniques:
|
||||
* **Quantization:** Supports various quantization implementations including FP8 and AWQ.
|
||||
* **FusedOps:** Allows for custom and third-party integration.
|
||||
|
||||
### Classifier-Free Guidance (CFG) Companion Flow
|
||||
|
||||
vLLM-Omni natively models Classifier-Free Guidance (CFG) across disaggregated multi-stage setups via a "companion request" paradigm, eliminating redundant textual/multimodal context computation boundaries:
|
||||
1. **Prompt Expansion:** In the initial autoregressive (AR) stage, a customized `prompt_expand_func` hook intercepts incoming generation prompts and pairs them directly with negative companion prompts (e.g., a default negative prompt) on the fly, tagging the secondary prompt with a specific internal role (`cfg_text`).
|
||||
2. **Synchronized KV Cache Transfer:** The AR stage evaluates both the primary and companion sequence batches concurrently. The `OmniConnector` captures these specific structural dependencies and reliably passes the positive and negative outcome KV caches seamlessly across stage boundaries via shared memory or network protocols.
|
||||
3. **KV Cache Collection & Injection:** Upon reaching the downstream Diffusion (DiT) Engine, an assigned `cfg_kv_collect_func` automatically intercepts the mapped companion caches (`cfg_text_past_key_values`). These auxiliary dependencies are natively gathered and seamlessly bound to the primary generation sequence variables, enabling the DiT Engine to cleanly implement cross-attention CFG guidance over accurate conditioning and unconditioning structures in parallel.
|
||||
|
||||
### Flexibility and Usability
|
||||
|
||||
vLLM-Omni is designed to be flexible and straightforward for users:
|
||||
|
||||
@@ -19,7 +19,7 @@ The `async_chunk` feature enables asynchronous, chunked processing of data acros
|
||||
|
||||
For qwen3-omni:
|
||||
- **Thinker → Talker**: Per decode step (typically chunk_size=1)
|
||||
- **Talker → Code2Wav**: Accumulated to code2wav chunk_size(default=25, currently only support default, will support chunk_size soon) before sending
|
||||
- **Talker → Code2Wav**: Accumulated to `codec_chunk_frames` (default=25) before sending. During the initial phase, a dynamic initial chunk size (IC) is automatically selected based on server load to reduce TTFA. Use the per-request `initial_codec_chunk_frames` API field to override.
|
||||
- **Code2Wav**: Streaming decode with code2wav chunk_size
|
||||
|
||||
With `async_chunk`:
|
||||
@@ -67,6 +67,13 @@ Enabling **async_chunk** (False→True) sharply reduces time-to-first-audio (TTF
|
||||
</picture>
|
||||
</p>
|
||||
|
||||
<p align="center">
|
||||
<picture>
|
||||
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/performance/qwen3-omni_e2e_performance.png">
|
||||
<img alt="E2E Performance Data Comparison" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/performance/qwen3-omni_e2e_performance.png" width=100%>
|
||||
</picture>
|
||||
</p>
|
||||
|
||||
## Architecture
|
||||
### Data Flow
|
||||
|
||||
|
||||
@@ -240,7 +240,8 @@ class TransformerWithRoPE(nn.Module):
|
||||
|
||||
## Approach 2: Intrusive Modification (For Complex Cases)
|
||||
|
||||
For models with dynamic sharding logic that cannot be expressed via `_sp_plan`, manually insert shard/gather calls.
|
||||
For models with dynamic sharding logic that cannot be expressed via `_sp_plan`, manually insert shard/gather calls. Importantly, when taking this approach, be careful to ensure that you correctly manage the `_sp_shard_depth`; if the sequence parallel shard depth is 0, Ulysses will not be used.
|
||||
|
||||
|
||||
**When to use:**
|
||||
- Dynamic/conditional sharding logic
|
||||
@@ -252,15 +253,18 @@ from vllm_omni.diffusion.distributed.sp_sharding import sp_shard, sp_gather
|
||||
|
||||
def forward(self, hidden_states, ...):
|
||||
if self.parallel_config.sequence_parallel_size > 1:
|
||||
# <Increment the _sp_shard_depth on the fwd context>
|
||||
hidden_states = sp_shard(hidden_states, dim=1)
|
||||
|
||||
# ... computation ...
|
||||
|
||||
if self.parallel_config.sequence_parallel_size > 1:
|
||||
output = sp_gather(output, dim=1)
|
||||
# <Decrement the _sp_shard_depth on the fwd context>
|
||||
|
||||
return output
|
||||
```
|
||||
Note that currently, `sp_shard` / `sp_gather` do *not* automatically manage the `_sp_shard_depth`; you need to be careful to manage it yourself.
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -0,0 +1,289 @@
|
||||
# VAE Patch Parallelism
|
||||
|
||||
This document describes how to add **VAE Patch Parallelism** support to a diffusion model.
|
||||
We use **Qwen-Image** as the reference implementation.
|
||||
|
||||
---
|
||||
|
||||
## Table of Contents
|
||||
|
||||
- [Overview](#overview)
|
||||
- [Step-by-Step Implementation](#step-by-step-implementation)
|
||||
- [Testing](#testing)
|
||||
- [Reference Implementations](#reference-implementations)
|
||||
- [Summary](#summary)
|
||||
|
||||
---
|
||||
|
||||
## Overview
|
||||
|
||||
### What is Vae Patch parallel?
|
||||
|
||||
**VAE Patch Parallelism** is a decoding acceleration technique. Instead of decoding the entire latent tensor at once, the latent tensor is:
|
||||
|
||||
+ Split into multiple spatial tiles
|
||||
|
||||
+ Distributed across multiple ranks
|
||||
|
||||
+ Decoded in parallel
|
||||
|
||||
+ Merged to reconstruct the final output
|
||||
|
||||
This approach:
|
||||
|
||||
+ Distributes computation across multiple devices
|
||||
|
||||
+ Reduces peak memory usage per device
|
||||
|
||||
+ Accelerates decoding latency
|
||||
|
||||
### Architecture
|
||||
We introduce **DistributedVaeExecutor** as the core component responsible for distributed VAE decoding.
|
||||
|
||||
The executor is model-agnostic and accepts three function parameters:
|
||||
|
||||
+ split – Partition the latent into tiles
|
||||
|
||||
+ exec – Decode a single tile
|
||||
|
||||
+ merge – Combine decoded tiles into the final output
|
||||
|
||||
#### Execution Flow
|
||||
|
||||
+ Call split(z) to generate a list of TileTask and a GridSpec
|
||||
|
||||
+ Dispatch tasks across ranks using workload-based balancing
|
||||
|
||||
+ Each rank executes exec(task) on its assigned tiles
|
||||
|
||||
+ Gather decoded tile results to rank 0
|
||||
|
||||
+ Rank 0 performs merge(...)
|
||||
|
||||
+ (Optional) Broadcast final result to all ranks
|
||||
|
||||
This design separates:
|
||||
|
||||
+ Distributed execution logic
|
||||
|
||||
+ Model-specific tiling and merging logic
|
||||
|
||||
#### Why split / exec / merge is necessary?
|
||||
|
||||
The latent tensor cannot be arbitrarily partitioned.
|
||||
|
||||
During decoding:
|
||||
|
||||
+ Each output pixel may depend on neighboring pixels
|
||||
|
||||
+ The receptive field is model-dependent
|
||||
|
||||
Therefore:
|
||||
|
||||
+ Tiles must include overlap
|
||||
|
||||
+ Merge must perform blending to avoid seams
|
||||
|
||||
## Step-by-Step Implementation
|
||||
|
||||
### Step 1: Implement DistributedAutoencoderKLQwenImage
|
||||
`QwenImagePipeline` use `AutoencoderKLQwenImage` for vae, so implement a distributed version:
|
||||
|
||||
|
||||
```
|
||||
class DistributedAutoencoderKLQwenImage(AutoencoderKLQwenImage, DistributedVaeMixin):
|
||||
@classmethod
|
||||
def from_pretrained(cls, *args: Any, **kwargs: Any):
|
||||
model = super().from_pretrained(*args, **kwargs)
|
||||
model.init_distributed()
|
||||
return model
|
||||
```
|
||||
**Key points**:
|
||||
+ Inherit both AutoencoderKLQwenImage and DistributedVaeMixin
|
||||
+ Call init_distributed() after loading weights
|
||||
|
||||
### Step 2: Implement split/exec/merge
|
||||
Reuse `AutoencoderKLQwenImage.tiled_decode` logic and divide it into three stages. And we need return tiles with `GridSpec` and `TileTask`:
|
||||
```
|
||||
class GridSpec:
|
||||
split_dims: tuple[int, ...] # Tensor dimensions being split (e.g., (2, 3) for (B, C, H, W))
|
||||
grid_shape: tuple[int, ...] # Tile grid layout (num_rows, num_cols)
|
||||
tile_spec: dict = field(default_factory=dict) # Metadata required for merging
|
||||
output_dtype: torch.dtype | None = None # Final output dtype
|
||||
```
|
||||
```
|
||||
class TileTask:
|
||||
tile_id: int # task id
|
||||
grid_coord: tuple[int, ...] # Tile position in grid
|
||||
tensor: torch.Tensor | list[torch.Tensor] # The tile tensor
|
||||
workload: int | float = 1 # Used for load balancing (e.g., tile area)
|
||||
```
|
||||
And tiled base split/exec/merge as follow:
|
||||
```
|
||||
def tile_split(self, z: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
|
||||
# mostly copy from AutoencoderKL
|
||||
_, _, num_frames, height, width = z.shape
|
||||
sample_height = height * self.spatial_compression_ratio
|
||||
sample_width = width * self.spatial_compression_ratio
|
||||
|
||||
tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
|
||||
tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
|
||||
tile_latent_stride_height = self.tile_sample_stride_height // self.spatial_compression_ratio
|
||||
tile_latent_stride_width = self.tile_sample_stride_width // self.spatial_compression_ratio
|
||||
|
||||
blend_height = self.tile_sample_min_height - self.tile_sample_stride_height
|
||||
blend_width = self.tile_sample_min_width - self.tile_sample_stride_width
|
||||
|
||||
# Split z into overlapping tiles and decode them separately.
|
||||
# The tiles have an overlap to avoid seams between tiles.
|
||||
tiletask_list = []
|
||||
for i in range(0, height, tile_latent_stride_height):
|
||||
for j in range(0, width, tile_latent_stride_width):
|
||||
time_list = []
|
||||
for k in range(num_frames):
|
||||
self._conv_idx = [0]
|
||||
tile = z[:, :, k : k + 1, i : i + tile_latent_min_height, j : j + tile_latent_min_width]
|
||||
time_list.append(tile)
|
||||
tiletask_list.append(
|
||||
TileTask(
|
||||
len(tiletask_list),
|
||||
(i // tile_latent_stride_height, j // tile_latent_stride_width),
|
||||
time_list,
|
||||
workload=time_list[0].shape[3] * time_list[0].shape[4],
|
||||
)
|
||||
)
|
||||
tile_spec = {
|
||||
"sample_height": sample_height,
|
||||
"sample_width": sample_width,
|
||||
"blend_height": blend_height,
|
||||
"blend_width": blend_width,
|
||||
}
|
||||
grid_spec = GridSpec(
|
||||
split_dims=(3, 4),
|
||||
grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1),
|
||||
tile_spec=tile_spec,
|
||||
output_dtype=self.dtype,
|
||||
)
|
||||
return tiletask_list, grid_spec
|
||||
|
||||
def tile_exec(self, task: TileTask) -> torch.Tensor:
|
||||
"""Decode a single latent tile into RGB space."""
|
||||
self.clear_cache()
|
||||
time = []
|
||||
for k in range(len(task.tensor)):
|
||||
self._conv_idx = [0]
|
||||
tile = self.post_quant_conv(task.tensor[k])
|
||||
decoded = self.decoder(tile, feat_cache=self._feat_map, feat_idx=self._conv_idx)
|
||||
time.append(decoded)
|
||||
result = torch.cat(time, dim=2)
|
||||
return result
|
||||
|
||||
def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec) -> torch.Tensor:
|
||||
"""Merge decoded tiles into a full image."""
|
||||
grid_h, grid_w = grid_spec.grid_shape
|
||||
result_rows = []
|
||||
self.clear_cache()
|
||||
|
||||
result_rows = []
|
||||
for i in range(grid_h):
|
||||
result_row = []
|
||||
for j in range(grid_w):
|
||||
tile = coord_tensor_map[(i, j)]
|
||||
if i > 0:
|
||||
tile = self.blend_v(coord_tensor_map[(i - 1, j)], tile, grid_spec.tile_spec["blend_height"])
|
||||
if j > 0:
|
||||
tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_width"])
|
||||
result_row.append(tile[:, :, :, : self.tile_sample_stride_height, : self.tile_sample_stride_width])
|
||||
result_rows.append(torch.cat(result_row, dim=-1))
|
||||
dec = torch.cat(result_rows, dim=3)[
|
||||
:, :, :, : grid_spec.tile_spec["sample_height"], : grid_spec.tile_spec["sample_width"]
|
||||
]
|
||||
return dec
|
||||
```
|
||||
|
||||
### Step 3: Override tiled_decode
|
||||
We need to override tiled_decode, the main logic is:
|
||||
+ check distributed is enabled
|
||||
+ select split/exec/merge
|
||||
+ Invoke self.distributed_decoder.execute to decode
|
||||
```
|
||||
def tiled_decode(self, z: torch.Tensor, return_dict: bool = True):
|
||||
if not self.is_distributed_enabled():
|
||||
return super().tiled_decode(z, return_dict=return_dict)
|
||||
|
||||
logger.info("Decode run with distributed executor")
|
||||
result = self.distributed_decoder.execute(
|
||||
z,
|
||||
DistributedOperator(split=self.tile_split, exec=self.tile_exec, merge=self.tile_merge),
|
||||
broadcast_result=True,
|
||||
)
|
||||
if not return_dict:
|
||||
return (result,)
|
||||
|
||||
return DecoderOutput(sample=result)
|
||||
```
|
||||
`broadcast_result` is set to True or False depending on the model; when enabled, the result will be used even on ranks other than 0.
|
||||
|
||||
### Step 4: Modify Pipeline
|
||||
Change vae model from AutoencoderKLQwenImage to DistributedAutoencoderKLQwenImage
|
||||
```
|
||||
class YourModelPipeline(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
od_config: OmniDiffusionConfig,
|
||||
prefix: str = "",
|
||||
):
|
||||
super().__init__()
|
||||
...
|
||||
- self.vae = AutoencoderKL.from_pretrained(
|
||||
- model, subfolder="vae", local_files_only=local_files_only).to(self.device)
|
||||
+ self.vae = DistributedAutoencoderKL.from_pretrained(
|
||||
+ model, subfolder="vae", local_files_only=local_files_only
|
||||
+ ).to(self.device)
|
||||
```
|
||||
|
||||
## Testing
|
||||
Verify numerical consistency between:
|
||||
+ vae_patch_parallel_size = 1
|
||||
|
||||
+ vae_patch_parallel_size = N
|
||||
|
||||
Example:
|
||||
torch.allclose(output_1, output_n, atol=1e-5)
|
||||
|
||||
Testing requirements:
|
||||
+ Fix random seed
|
||||
+ Use identical tiling strategy
|
||||
|
||||
```python
|
||||
m = Omni(
|
||||
model=model_name,
|
||||
vae_use_tiling=True,
|
||||
parallel_config=DiffusionParallelConfig(
|
||||
tensor_parallel_size=2,
|
||||
vae_patch_parallel_size=1, # or 2
|
||||
),
|
||||
)
|
||||
```
|
||||
When vae_patch_parallel_size is larger than the DiT world size, it will automatically fall back to using the DiT world size instead.
|
||||
|
||||
## Reference Implementations
|
||||
|
||||
Complete examples in the codebase:
|
||||
|
||||
| Model | Path | Notes |
|
||||
|-------|------|-------|
|
||||
| **Z-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl.py` | Distributed AutoencoderKL |
|
||||
| **Wan2.2** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py` | Distributed AutoencoderKLWan |
|
||||
| **Qwen-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py` | Distributed AutoencoderKLQwenImage |
|
||||
|
||||
---
|
||||
|
||||
## Summary
|
||||
|
||||
Adding Vae Patch Parallel support to diffusion model:
|
||||
|
||||
1. **Implement Distributed Vae** - mainly copy from `diffusers` tiled_decode, and refactor into split/exec/merge
|
||||
2. **Change vae model in pipeline to Distributed Vae**
|
||||
3. **Test** - Verify with `tensor_parallel_size=N` quality
|
||||
@@ -47,7 +47,8 @@ If no, check your shell running the ComfyUI process. There may be some error mes
|
||||
This extension offers the following nodes based on the output modalities:
|
||||
|
||||
- **Generate Image** for text-to-image and image-to-image tasks
|
||||
- **Multimodality Comprehension** for multimodality-to-text and multimodality-to-audio tasks
|
||||
- **Generate Video** for text-to-video and image-to-video tasks
|
||||
- **Multimodality Understanding** for multimodality-to-text and multimodality-to-audio tasks
|
||||
- **TTS** and **TTS Voice Clone** for TTS tasks
|
||||
|
||||
This extension also offers example workflows (at **ComfyUI sidebar -> Templates -> vLLM-Omni**)
|
||||
@@ -59,7 +60,7 @@ To build a simple workflow yourself,
|
||||
|
||||
- Drag a generation node onto the canvas.
|
||||
- Depending on your need, grab built-in multimedia file loader nodes, such as **image->Load Image**, **image->video->Load Video**, **audio->Load Audio**
|
||||
- Depending on your need, grab built-in multimedia file preview nodes, such as **image->Preview Image**, **image->video->Save Video**, **audio->Preview Audio**. For text output, you can install [ComfyUI-Custom-Scripts plugin](https://github.com/pythongosssss/ComfyUI-Custom-Scripts/) and grab its **utils->Show Text 🐍** node.
|
||||
- Depending on your need, grab built-in multimedia file preview nodes, such as **image->Preview Image**, **image->video->Save Video**, **audio->Preview Audio**, **utils->Preview as Text**.
|
||||
- If you want to tune sampling parameters, grab corresponding nodes from **vLLM-Omni-> Sampling Params**.
|
||||
- For multi-stage models, you can connect multiple **AR Sampling Params** and **Diffusion Sampling Params** nodes to a **Multi-Stage Sampling Params List** node, and connect this node to the generation node.
|
||||
- For some multi-stage models like BAGEL, [only one stage's sampling parameters are exposed and tunable via vLLM-Omni's online serving API](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/examples/online_serving/bagel/). Thus, these models are treated as single-stage ones. Please check the vLLM-Omni documentation on how to correctly set each model's sampling parameters.
|
||||
|
||||
@@ -115,7 +115,7 @@ python image_edit.py \
|
||||
--image cherry_blossom.jpg \
|
||||
--prompt "Let this mascot dance under the moon, surrounded by floating stars" \
|
||||
--output output_image_edit.png \
|
||||
--num_inference_steps 10
|
||||
--num-inference-steps 10
|
||||
```
|
||||
|
||||
## Advanced Usage
|
||||
|
||||
@@ -34,13 +34,13 @@ uv pip install vllm-omni
|
||||
# --8<-- [start:build-wheel-from-source]
|
||||
|
||||
#### Installation of vLLM
|
||||
If you do not need to modify source code of vLLM, you can directly install the stable 0.16.0 release version of the library
|
||||
If you do not need to modify source code of vLLM, you can directly install the stable 0.17.0 release version of the library
|
||||
|
||||
```bash
|
||||
uv pip install vllm==0.16.0 --torch-backend=auto
|
||||
uv pip install vllm==0.17.0 --torch-backend=auto
|
||||
```
|
||||
|
||||
The release 0.16.0 of vLLM is based on PyTorch 2.9.0 which requires CUDA 12.9 environment.
|
||||
The release 0.17.0 of vLLM is based on PyTorch 2.10.0 which requires CUDA 12.9 environment.
|
||||
|
||||
#### Installation of vLLM-Omni
|
||||
Since vllm-omni is rapidly evolving, it's recommended to install it from source
|
||||
@@ -56,11 +56,11 @@ If you want to check, modify or debug with source code of vLLM, install the libr
|
||||
```bash
|
||||
git clone https://github.com/vllm-project/vllm.git
|
||||
cd vllm
|
||||
git checkout v0.16.0
|
||||
git checkout v0.17.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
|
||||
export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.16.0/vllm-0.16.0-cp38-abi3-manylinux_2_31_x86_64.whl
|
||||
export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.17.0/vllm-0.17.0-cp38-abi3-manylinux_2_31_x86_64.whl
|
||||
```
|
||||
Install vllm with command below (If you have no existing PyTorch).
|
||||
```bash
|
||||
|
||||
@@ -13,7 +13,7 @@ vLLM-Omni current recommends the steps in under setup through Docker Images.
|
||||
|
||||
vLLM-Omni is built based on vLLM. Please install it with command below.
|
||||
```bash
|
||||
uv pip install vllm==0.16.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.16.0/rocm700
|
||||
uv pip install vllm==0.17.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.17.0/rocm700
|
||||
```
|
||||
|
||||
#### Installation of vLLM-Omni
|
||||
@@ -34,13 +34,13 @@ uv pip install onnxruntime-rocm sox
|
||||
# --8<-- [start:build-wheel-from-source]
|
||||
|
||||
#### Installation of vLLM
|
||||
If you do not need to modify source code of vLLM, you can directly install the stable 0.16.0 release version of the library
|
||||
If you do not need to modify source code of vLLM, you can directly install the stable 0.17.0 release version of the library
|
||||
|
||||
```bash
|
||||
uv pip install vllm==0.16.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.16.0/rocm700
|
||||
uv pip install vllm==0.17.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.17.0/rocm700
|
||||
```
|
||||
|
||||
The release 0.16.0 of vLLM requires ROCm 7.0 environment.
|
||||
The release 0.17.0 of vLLM requires ROCm 7.0 environment.
|
||||
|
||||
#### Installation of vLLM-Omni
|
||||
Since vllm-omni is rapidly evolving, it's recommended to install it from source
|
||||
@@ -58,7 +58,7 @@ If you want to check, modify or debug with source code of vLLM, install the libr
|
||||
```bash
|
||||
git clone https://github.com/vllm-project/vllm.git
|
||||
cd vllm
|
||||
git checkout v0.16.0
|
||||
git checkout v0.17.0
|
||||
python3 -m pip install -r requirements/rocm.txt
|
||||
python3 setup.py develop
|
||||
```
|
||||
|
||||
@@ -19,10 +19,10 @@ uv venv --python 3.12 --seed
|
||||
source .venv/bin/activate
|
||||
|
||||
# On CUDA
|
||||
uv pip install vllm==0.16.0 --torch-backend=auto
|
||||
uv pip install vllm==0.17.0 --torch-backend=auto
|
||||
|
||||
# On ROCm
|
||||
uv pip install vllm==0.16.0 --extra-index-url https://wheels.vllm.ai/rocm/0.16.0/rocm700
|
||||
uv pip install vllm==0.17.0 --extra-index-url https://wheels.vllm.ai/rocm/0.17.0/rocm700
|
||||
|
||||
git clone https://github.com/vllm-project/vllm-omni.git
|
||||
cd vllm-omni
|
||||
|
||||
@@ -132,9 +132,19 @@ class Example:
|
||||
".dat",
|
||||
".db",
|
||||
".sqlite", # Data files
|
||||
".pyc",
|
||||
".pyo",
|
||||
".pyd", # Python compiled
|
||||
".npy",
|
||||
".npz",
|
||||
".pkl",
|
||||
".pickle", # Serialized data
|
||||
}
|
||||
excluded_dirs = {"__pycache__", ".git", "node_modules", ".tox", ".mypy_cache"}
|
||||
|
||||
def is_other_file(file: Path) -> bool:
|
||||
if any(part in excluded_dirs for part in file.parts):
|
||||
return False
|
||||
return file.is_file() and file != self.main_file and file.suffix.lower() not in binary_extensions
|
||||
|
||||
return [file for file in self.path.rglob("*") if is_other_file(file)]
|
||||
|
||||
@@ -30,10 +30,15 @@ th {
|
||||
|`ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` |
|
||||
| `WanPipeline` | Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` |
|
||||
| `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` |
|
||||
| `LTX2Pipeline` | LTX-2-T2V | `Lightricks/LTX-2` |
|
||||
| `LTX2ImageToVideoPipeline` | LTX-2-I2V | `Lightricks/LTX-2` |
|
||||
| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` |
|
||||
| `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` |
|
||||
|`LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` |
|
||||
|`LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` |
|
||||
|`StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` |
|
||||
|`CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` |
|
||||
|`MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` |
|
||||
|`Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` |
|
||||
|`FluxPipeline` | FLUX.1-dev | `black-forest-labs/FLUX.1-dev` |
|
||||
|`OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` |
|
||||
@@ -44,6 +49,7 @@ th {
|
||||
|`GlmImageForConditionalGeneration` | GLM-Image | `zai-org/GLM-Image` |
|
||||
|`NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` |
|
||||
|`MiMoAudioForConditionalGeneration` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` |
|
||||
|`Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` |
|
||||
|
||||
|
||||
## List of Supported Models for NPU
|
||||
|
||||
@@ -96,6 +96,7 @@ Content-Type: application/json
|
||||
| `language` | string | "Auto" | Language (see supported languages below) |
|
||||
| `instructions` | string | "" | Voice style/emotion instructions |
|
||||
| `max_new_tokens` | integer | 2048 | Maximum tokens to generate |
|
||||
| `initial_codec_chunk_frames` | integer | null | Per-request initial chunk size override for TTFA tuning. When null, IC is computed dynamically based on server load. |
|
||||
|
||||
**Supported languages:** Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian
|
||||
|
||||
@@ -103,7 +104,7 @@ Content-Type: application/json
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| `ref_audio` | string | null | Reference audio (URL or base64 data URL) |
|
||||
| `ref_audio` | string | null | Reference audio (HTTP URL, base64 data URL, or `file://` URI with `--allowed-local-media-path`) |
|
||||
| `ref_text` | string | null | Transcript of reference audio |
|
||||
| `x_vector_only_mode` | bool | null | Use speaker embedding only (no ICL) |
|
||||
|
||||
@@ -124,6 +125,81 @@ Lists available voices for the loaded model.
|
||||
"voices": ["aiden", "dylan", "eric", "ono_anna", "ryan", "serena", "sohee", "uncle_fu", "vivian"]
|
||||
}
|
||||
```
|
||||
```
|
||||
POST /v1/audio/voices
|
||||
Content-Type: multipart/form-data
|
||||
```
|
||||
|
||||
Upload a new voice sample for voice cloning in Base task TTS requests.
|
||||
|
||||
**Form Parameters:**
|
||||
|
||||
| Parameter | Type | Required | Description |
|
||||
|-----------|------|----------|-------------|
|
||||
| `audio_sample` | file | Yes | Audio file (max 10MB, supported formats: wav, mp3, flac, ogg, aac, webm, mp4) |
|
||||
| `consent` | string | Yes | Consent recording ID |
|
||||
| `name` | string | Yes | Name for the new voice |
|
||||
|
||||
**Response Example:**
|
||||
|
||||
```json
|
||||
{
|
||||
"success": true,
|
||||
"voice": {
|
||||
"name": "custom_voice_1",
|
||||
"consent": "user_consent_id",
|
||||
"created_at": 1738660000,
|
||||
"mime_type": "audio/wav",
|
||||
"file_size": 1024000
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
**Usage Example:**
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/voices \
|
||||
-F "audio_sample=@/path/to/voice_sample.wav" \
|
||||
-F "consent=user_consent_id" \
|
||||
-F "name=custom_voice_1"
|
||||
```
|
||||
|
||||
|
||||
```bash
|
||||
DELETE /v1/audio/voices/{name}
|
||||
```
|
||||
|
||||
Delete an uploaded voice sample.
|
||||
|
||||
**Path Parameters:**
|
||||
|
||||
| Parameter | Type | Required | Description |
|
||||
|-----------|------|----------|-------------|
|
||||
| `name` | string | Yes | Name of the voice to delete |
|
||||
|
||||
**Response Example:**
|
||||
|
||||
```json
|
||||
{
|
||||
"success": true,
|
||||
"message": "Voice 'custom_voice_1' deleted successfully"
|
||||
}
|
||||
```
|
||||
|
||||
**Error Response (404 Not Found):**
|
||||
|
||||
```json
|
||||
{
|
||||
"success": false,
|
||||
"error": "Voice 'unknown_voice' not found"
|
||||
}
|
||||
```
|
||||
|
||||
**Usage Example:**
|
||||
|
||||
```bash
|
||||
curl -X DELETE http://localhost:8091/v1/audio/voices/custom_voice_1
|
||||
```
|
||||
|
||||
## Examples
|
||||
|
||||
@@ -184,6 +260,25 @@ curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
}' --output cloned.wav
|
||||
```
|
||||
|
||||
upload voice
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/voices \
|
||||
-F "audio_sample=@/path/to/voice_sample.wav" \
|
||||
-F "consent=user_consent_id" \
|
||||
-F "name=custom_voice_1"
|
||||
```
|
||||
|
||||
use upload voice
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "Hello, this is a cloned voice",
|
||||
"task_type": "Base",
|
||||
"voice": "custom_voice_1"
|
||||
}' --output cloned.wav
|
||||
```
|
||||
|
||||
## Supported Models
|
||||
|
||||
| Model | Task Type | Description |
|
||||
|
||||
|
After Width: | Height: | Size: 34 KiB |
|
Before Width: | Height: | Size: 36 KiB After Width: | Height: | Size: 37 KiB |
|
Before Width: | Height: | Size: 33 KiB After Width: | Height: | Size: 45 KiB |
@@ -18,23 +18,27 @@ The following parallelism methods are currently supported in vLLM-Omni:
|
||||
|
||||
6. [HSDP](#hsdp): Hybrid Sharded Data Parallel shards model weights across GPUs using PyTorch FSDP2. This reduces per-GPU memory usage, enabling inference of large models on GPUs with limited memory.
|
||||
|
||||
7. [Expert Parallel](#expert-parallelism): Expert Parallelism shards the Experts of a Mixture-of-Experts (MoE) layer across multiple devices. During the forward, a gating mechanism routes tokens to their designated experts, necessitating cross-cards communication(all-to-all) to dispatch tokens to the correct ranks and combine the results. This parallelism allows for massive scaling of model parameters without a proportional increase in the computational load per device.
|
||||
|
||||
The following table shows which models are currently supported by parallelism method:
|
||||
|
||||
### ImageGen
|
||||
|
||||
| Model | Model Identifier | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel | VAE-Patch-Parallel |
|
||||
|--------------------------|--------------------------------------|:----------:|:-------:|:------------:|:---------------:|:------------------:|
|
||||
| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ✅ | ✅ | ❌ | ✅ | ❌ |
|
||||
| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ✅ | ✅ | ❌ | ✅ | ❌ |
|
||||
| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | ❌ |
|
||||
| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ | ❌ |
|
||||
| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ | ❌ |
|
||||
| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ✅ | ✅ | ✅ | ✅ | ❌ |
|
||||
| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ✅ (TP=2 only) | ✅ |
|
||||
| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ❌ | ❌ | ✅ | ✅ |
|
||||
| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ❌ | ❌ | ❌ | ✅ | ❌ |
|
||||
| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ❌ | ✅ | ✅ | ❌ |
|
||||
| Model | Model Identifier | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel | VAE-Patch-Parallel | Expert-Parallel |
|
||||
|--------------------------|--------------------------------------|:----------:|:-------:|:------------:|:---------------:|:------------------:|:---------------:|
|
||||
| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ✅ | ✅ | ❌ | ✅ | ❌ | N/A |
|
||||
| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ✅ | ✅ | ❌ | ✅ | ❌ | N/A |
|
||||
| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ❌ | ❌ | ❌ | ❌ | N/A |
|
||||
| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | ✅ | N/A |
|
||||
| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ | ❌ | N/A |
|
||||
| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ | ❌ | N/A |
|
||||
| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ✅ | ✅ | ✅ | ✅ | ❌ | N/A |
|
||||
| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ✅ (TP=2 only) | ✅ | N/A |
|
||||
| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ❌ | ❌ | ✅ | ✅ | N/A |
|
||||
| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ❌ | ❌ | ❌ | ✅ | ❌ | N/A |
|
||||
| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ❌ | ✅ | ✅ | ❌ | N/A |
|
||||
| **FLUX.2-dev** | `black-forest-labs/FLUX.2-dev` | ❌ | ❌ | ❌ | ✅ | ❌ | N/A |
|
||||
| **HunyuanImage3.0** | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ |
|
||||
|
||||
!!! note "TP Limitations for Diffusion Models"
|
||||
We currently implement Tensor Parallelism (TP) only for the DiT (Diffusion Transformer) blocks. This is because the `text_encoder` component in vLLM-Omni uses the original Transformers implementation, which does not yet support TP.
|
||||
@@ -51,9 +55,10 @@ The following table shows which models are currently supported by parallelism me
|
||||
|
||||
### VideoGen
|
||||
|
||||
| Model | Model Identifier | Ulysses-SP | Ring-Attention | Tensor-Parallel | HSDP |
|
||||
|-------|------------------|:----------:|:--------------:|:---------------:|:----:|
|
||||
| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ✅ | ✅ | ✅ | ✅ |
|
||||
| Model | Model Identifier | Ulysses-SP | Ring-Attention | Tensor-Parallel | HSDP | VAE-Patch-Parallel |
|
||||
|-------|------------------|:----------:|:--------------:|:---------------:|:----:| :----:|
|
||||
| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ✅ | ✅ | ✅ | ✅ | ✅ |
|
||||
| **LTX-2** | `Lightricks/LTX-2` | ✅ | ✅ | ✅ | ❌ | ❌ |
|
||||
|
||||
### Tensor Parallelism
|
||||
|
||||
@@ -85,7 +90,7 @@ outputs = omni.generate(
|
||||
VAE patch parallelism distributes the VAE decode workload across multiple ranks by splitting the latent spatially. It is configured via `DiffusionParallelConfig.vae_patch_parallel_size` and can be combined with other parallelism methods (e.g., TP).
|
||||
|
||||
!!! note "Enablement and feature gate"
|
||||
- VAE patch parallelism is currently **enabled only for validated pipelines** (currently: `Tongyi-MAI/Z-Image-Turbo`, `stabilityai/stable-diffusion-3.5`).
|
||||
- VAE patch parallelism is currently **enabled only for validated pipelines** (check [ImageGen](#imagegen) and [VideoGen](#videogen) for more information).
|
||||
- If `vae_patch_parallel_size > 1` is set for a validated pipeline, vLLM-Omni will automatically enable `vae_use_tiling` as a safety gate. (We use `vae_use_tiling` because it indicates the VAE supports diffusers tiling parameters like `tile_latent_min_size` and `tile_overlap_factor`.)
|
||||
|
||||
#### Offline Inference
|
||||
@@ -413,3 +418,28 @@ vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --port 8091 --use-hsdp --usp
|
||||
#### Adding HSDP Support to New Models
|
||||
|
||||
For detailed instructions on adding HSDP support to new models, see the [HSDP Contributing Guide](../../design/feature/hsdp.md).
|
||||
|
||||
### Expert Parallelism
|
||||
|
||||
Unlike Tensor Parallelism which shards every layer's weights, EP only shards the MoE expert MLP blocks. This significantly reduces the memory footprint of MoE models (e.g., HunyuanImage3.0) while maintaining constant dense-equivalent compute efficiency. Expert Parallelism is enabled via `DiffusionParallelConfig.enable_expert_parallel`. And `self.ep = tp * sp * cfg * dp` for now, so at least one of TP/SP/CFG/DP should set when EP enabled.
|
||||
|
||||
#### Offline Inference
|
||||
|
||||
```python
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
|
||||
omni = Omni(
|
||||
model="tencent/HunyuanImage-3.0",
|
||||
parallel_config=DiffusionParallelConfig(tensor_parallel_size=8, enable_expert_parallel=True),
|
||||
)
|
||||
|
||||
outputs = omni.generate(
|
||||
"A brown and white dog is running on the grass",
|
||||
OmniDiffusionSamplingParams(
|
||||
num_inference_steps=50,
|
||||
width=1024,
|
||||
height=1024,
|
||||
),
|
||||
)
|
||||
```
|
||||
@@ -62,6 +62,7 @@ The available `ignored_layers` names depend on the model architecture (e.g., `to
|
||||
|-------|-----------|---------------|------------------|
|
||||
| Z-Image | `Tongyi-MAI/Z-Image-Turbo` | All layers | None |
|
||||
| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Skip sensitive layers | `img_mlp` |
|
||||
| Flux | `black-forest-labs/FLUX.1-dev` | All layers | None |
|
||||
|
||||
## Combining with Other Features
|
||||
|
||||
|
||||
@@ -136,9 +136,9 @@ python examples/offline_inference/text_to_image/text_to_image.py \
|
||||
--height 768 \
|
||||
--width 1360 \
|
||||
--seed 42 \
|
||||
--cfg_scale 4.0 \
|
||||
--num_images_per_prompt 1 \
|
||||
--num_inference_steps 4 \
|
||||
--cfg-scale 4.0 \
|
||||
--num-images-per-prompt 1 \
|
||||
--num-inference-steps 4 \
|
||||
--output outputs/flux2_klein_4b.png
|
||||
```
|
||||
|
||||
@@ -152,9 +152,9 @@ python examples/offline_inference/text_to_image/text_to_image.py \
|
||||
--height 768 \
|
||||
--width 1360 \
|
||||
--seed 42 \
|
||||
--cfg_scale 4.0 \
|
||||
--num_images_per_prompt 1 \
|
||||
--num_inference_steps 4 \
|
||||
--cfg-scale 4.0 \
|
||||
--num-images-per-prompt 1 \
|
||||
--num-inference-steps 4 \
|
||||
--output outputs/flux2_klein_4b_gguf.png
|
||||
```
|
||||
|
||||
|
||||
@@ -58,8 +58,8 @@ The following table shows which models are currently supported by each accelerat
|
||||
| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ |
|
||||
| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ |
|
||||
| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ |
|
||||
| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
|
||||
| **Qwen-Image-2512** | `Qwen/Qwen-Image-2512` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
|
||||
| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
|
||||
| **Qwen-Image-2512** | `Qwen/Qwen-Image-2512` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
|
||||
| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
|
||||
| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
|
||||
| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
|
||||
@@ -67,13 +67,16 @@ The following table shows which models are currently supported by each accelerat
|
||||
| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ |
|
||||
| **Bagel** | `ByteDance-Seed/BAGEL-7B-MoT` | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ |
|
||||
| **NextStep-1.1** | `stepfun-ai/NextStep-1.1` | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ |
|
||||
| **NextStep-1.1** | `stepfun-ai/NextStep-1.1` | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ |
|
||||
| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ |
|
||||
| **FLUX.2-dev** | `black-forest-labs/FLUX.2-dev` | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ |
|
||||
|
||||
### VideoGen
|
||||
|
||||
| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention | CFG-Parallel | HSDP |
|
||||
|-------|------------------|:--------:|:---------:|:----------:|:--------------:|:------------:|:----:|
|
||||
| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ |
|
||||
| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention | CFG-Parallel | HSDP | VAE-Patch-Parallel |
|
||||
|-------|------------------|:--------:|:---------:|:----------:|:--------------:|:------------:|:----:|:----:|
|
||||
| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
|
||||
| **LTX-2** | `Lightricks/LTX-2` | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
|
||||
|
||||
### Quantization
|
||||
|
||||
|
||||
@@ -11,7 +11,7 @@ Qwen3 TTS provides multiple task variants for speech generation:
|
||||
|
||||
- **CustomVoice**: Generate speech with a known speaker identity (speaker ID) and optional instruction.
|
||||
- **VoiceDesign**: Generate speech from text plus a descriptive instruction that designs a new voice.
|
||||
- **Base**: Voice cloning using a reference audio + reference transcript, with optional mode selection.
|
||||
- **Base**: Voice cloning using a reference audio + reference transcript, with optional mode selection. The `ref_audio` field accepts a local file path, HTTP URL, or base64 data URL.
|
||||
|
||||
## Setup
|
||||
Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
|
||||
@@ -98,7 +98,9 @@ Add `--streaming` to stream audio chunks progressively via `AsyncOmni` (requires
|
||||
python end2end.py --query-type CustomVoice --streaming --output-dir /tmp/out_stream
|
||||
```
|
||||
|
||||
Each 25-frame Code2Wav chunk is logged as it arrives. The final WAV file is written once generation
|
||||
Each Code2Wav chunk is logged as it arrives (default 25 frames; configurable via `codec_chunk_frames`
|
||||
in the stage config). The initial chunk size is dynamically selected based on server load for reduced
|
||||
TTFA, and can be overridden per-request via the `initial_codec_chunk_frames` API field. The final WAV file is written once generation
|
||||
completes. This demonstrates that audio data is available progressively rather than only at the end.
|
||||
|
||||
> **Note:** Streaming uses `AsyncOmni` internally. The non-streaming path (`Omni`) is unchanged.
|
||||
|
||||
@@ -17,6 +17,31 @@ Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/
|
||||
| `Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice` | CustomVoice | Smaller/faster variant |
|
||||
| `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | Base | Smaller/faster variant for voice cloning |
|
||||
|
||||
## Gradio Demo
|
||||
|
||||
Two interactive Gradio demos are available, both supporting all 3 task types:
|
||||
|
||||
| Demo | File | Transport | Streaming Quality |
|
||||
| ---- | ---- | --------- | ----------------- |
|
||||
| Standard | `gradio_demo.py` | HTTP chunked | May have small gaps between chunks |
|
||||
| FastRTC | `gradio_fastrtc_demo.py` | WebRTC | Gapless streaming (requires `pip install fastrtc`) |
|
||||
|
||||
```bash
|
||||
# Option 1: Launch server + Standard Gradio together
|
||||
./run_gradio_demo.sh # CustomVoice (default)
|
||||
./run_gradio_demo.sh --task-type VoiceDesign # VoiceDesign
|
||||
./run_gradio_demo.sh --task-type Base # Voice cloning
|
||||
|
||||
# Option 2: If server is already running
|
||||
python gradio_demo.py --api-base http://localhost:8000
|
||||
|
||||
# Option 3: FastRTC demo (gapless streaming)
|
||||
pip install fastrtc
|
||||
python gradio_fastrtc_demo.py --api-base http://localhost:8000
|
||||
```
|
||||
|
||||
Then open http://localhost:7860 in your browser.
|
||||
|
||||
## Run examples (Qwen3-TTS)
|
||||
|
||||
### Launch the Server
|
||||
@@ -187,6 +212,59 @@ sudo apt install ffmpeg
|
||||
|
||||
## API Reference
|
||||
|
||||
### Voices Endpoint
|
||||
|
||||
#### GET /v1/audio/voices
|
||||
|
||||
List all available voices/speakers from the loaded model, including both built-in model voices and uploaded custom voices.
|
||||
|
||||
**Response Example:**
|
||||
```json
|
||||
{
|
||||
"voices": ["vivian", "ryan", "custom_voice_1"],
|
||||
"uploaded_voices": [
|
||||
{
|
||||
"name": "custom_voice_1",
|
||||
"consent": "user_consent_id",
|
||||
"created_at": 1738660000,
|
||||
"file_size": 1024000,
|
||||
"mime_type": "audio/wav"
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
#### POST /v1/audio/voices
|
||||
|
||||
Upload a new voice sample for voice cloning in Base task TTS requests.
|
||||
|
||||
**Form Parameters:**
|
||||
- `audio_sample` (required): Audio file (max 10MB, supported formats: wav, mp3, flac, ogg, aac, webm, mp4)
|
||||
- `consent` (required): Consent recording ID
|
||||
- `name` (required): Name for the new voice
|
||||
|
||||
**Response Example:**
|
||||
```json
|
||||
{
|
||||
"success": true,
|
||||
"voice": {
|
||||
"name": "custom_voice_1",
|
||||
"consent": "user_consent_id",
|
||||
"created_at": 1738660000,
|
||||
"mime_type": "audio/wav",
|
||||
"file_size": 1024000
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
**Usage Example:**
|
||||
```bash
|
||||
curl -X POST http://localhost:8000/v1/audio/voices \
|
||||
-F "audio_sample=@/path/to/voice_sample.wav" \
|
||||
-F "consent=user_consent_id" \
|
||||
-F "name=custom_voice_1"
|
||||
```
|
||||
|
||||
### Endpoint
|
||||
|
||||
```
|
||||
@@ -196,20 +274,6 @@ Content-Type: application/json
|
||||
|
||||
This endpoint follows the [OpenAI Audio Speech API](https://platform.openai.com/docs/api-reference/audio/createSpeech) format with additional Qwen3-TTS parameters.
|
||||
|
||||
### Voices Endpoint
|
||||
|
||||
```
|
||||
GET /v1/audio/voices
|
||||
```
|
||||
|
||||
Lists available voices for the loaded model:
|
||||
|
||||
```json
|
||||
{
|
||||
"voices": ["aiden", "dylan", "eric", "one_anna", "ryan", "serena", "sohee", "uncle_fu", "vivian"]
|
||||
}
|
||||
```
|
||||
|
||||
### Request Body
|
||||
|
||||
```json
|
||||
@@ -253,6 +317,7 @@ Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/w
|
||||
| `language` | string | "Auto" | Language (see supported languages below) |
|
||||
| `instructions` | string | "" | Voice style/emotion instructions |
|
||||
| `max_new_tokens` | int | 2048 | Maximum tokens to generate |
|
||||
| `initial_codec_chunk_frames` | int | null | Per-request initial chunk size override for TTFA tuning. When null, IC is computed dynamically based on server load. |
|
||||
| `stream` | bool | false | Stream raw PCM chunks as they are decoded (requires `response_format="pcm"`) |
|
||||
|
||||
**Supported languages:** Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian
|
||||
@@ -268,7 +333,7 @@ Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/w
|
||||
## Streaming
|
||||
|
||||
Set `stream=true` with `response_format="pcm"` to receive raw PCM audio chunks as they are decoded
|
||||
(one chunk per 25-frame Code2Wav window):
|
||||
(one chunk per Code2Wav window, default 25 frames; configurable in the stage config):
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
@@ -301,11 +366,27 @@ curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
|
||||
## Example materials
|
||||
|
||||
??? abstract "gradio_demo.py"
|
||||
``````py
|
||||
--8<-- "examples/online_serving/qwen3_tts/gradio_demo.py"
|
||||
``````
|
||||
??? abstract "gradio_fastrtc_demo.py"
|
||||
``````py
|
||||
--8<-- "examples/online_serving/qwen3_tts/gradio_fastrtc_demo.py"
|
||||
``````
|
||||
??? abstract "openai_speech_client.py"
|
||||
``````py
|
||||
--8<-- "examples/online_serving/qwen3_tts/openai_speech_client.py"
|
||||
``````
|
||||
??? abstract "run_gradio_demo.sh"
|
||||
``````sh
|
||||
--8<-- "examples/online_serving/qwen3_tts/run_gradio_demo.sh"
|
||||
``````
|
||||
??? abstract "run_server.sh"
|
||||
``````sh
|
||||
--8<-- "examples/online_serving/qwen3_tts/run_server.sh"
|
||||
``````
|
||||
??? abstract "tts_common.py"
|
||||
``````py
|
||||
--8<-- "examples/online_serving/qwen3_tts/tts_common.py"
|
||||
``````
|
||||
@@ -23,6 +23,29 @@ Or use the startup script:
|
||||
bash run_server.sh
|
||||
```
|
||||
|
||||
### Start with Parallelism Acceleration
|
||||
|
||||
Enable Tensor Parallelism and VAE Patch Parallelism for faster inference:
|
||||
|
||||
```bash
|
||||
# With Tensor Parallelism (requires >= 2 GPUs)
|
||||
vllm serve Qwen/Qwen-Image --omni --port 8091 --tensor-parallel-size 2
|
||||
|
||||
# With Tensor Parallelism and VAE Patch Parallelism (requires >= 2 GPUs)
|
||||
vllm serve Qwen/Qwen-Image --omni --port 8091 --tensor-parallel-size 2 --vae-patch-parallel-size 2 --vae-use-tiling
|
||||
|
||||
# With Sequence Parallelism (Ulysses-SP, requires >= 2 GPUs)
|
||||
vllm serve Qwen/Qwen-Image --omni --port 8091 --usp 2
|
||||
|
||||
# With Ring-Attention (requires >= 2 GPUs)
|
||||
vllm serve Qwen/Qwen-Image --omni --port 8091 --ring 2
|
||||
|
||||
# Combined: Ulysses + Ring (requires >= 4 GPUs)
|
||||
vllm serve Qwen/Qwen-Image --omni --port 8091 --usp 2 --ring 2
|
||||
```
|
||||
|
||||
For more details on parallelism acceleration, see the [Parallelism Acceleration Guide](../../diffusion/parallelism_acceleration.md).
|
||||
|
||||
## API Calls
|
||||
|
||||
### Method 1: Using curl
|
||||
|
||||
@@ -1,8 +1,7 @@
|
||||
import argparse
|
||||
import os
|
||||
from typing import cast
|
||||
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
|
||||
from vllm_omni.inputs.data import OmniPromptType
|
||||
|
||||
|
||||
def parse_args():
|
||||
@@ -49,8 +48,16 @@ def parse_args():
|
||||
parser.add_argument("--cfg-text-scale", type=float, default=4.0, help="Text CFG scale (default: 4.0)")
|
||||
parser.add_argument("--cfg-img-scale", type=float, default=1.5, help="Image CFG scale (default: 1.5)")
|
||||
parser.add_argument(
|
||||
"--negative-prompt", type=str, default=None, help="Negative prompt (not yet supported, reserved for future)"
|
||||
"--negative-prompt", type=str, default=None, help="Negative prompt for CFG (default: empty prompt)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--cfg-parallel-size",
|
||||
type=int,
|
||||
default=1,
|
||||
choices=[1, 2, 3],
|
||||
help="CFG parallel size: 1=batched (single GPU), 2=parallel with 2 branches (text CFG only), 3=parallel (3 GPUs).",
|
||||
)
|
||||
parser.add_argument("--seed", type=int, default=None, help="Random seed for generation.")
|
||||
|
||||
args = parser.parse_args()
|
||||
return args
|
||||
@@ -81,101 +88,80 @@ def main():
|
||||
|
||||
from PIL import Image
|
||||
|
||||
if args.modality == "img2img":
|
||||
from PIL import Image
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
from vllm_omni.entrypoints.omni_diffusion import OmniDiffusion
|
||||
omni_kwargs = {}
|
||||
if args.stage_configs_path:
|
||||
omni_kwargs["stage_configs_path"] = args.stage_configs_path
|
||||
|
||||
print("[Info] Running in img2img mode (Stage 1 only)")
|
||||
client = OmniDiffusion(model=model_name)
|
||||
omni_kwargs.update(
|
||||
{
|
||||
"log_stats": args.log_stats,
|
||||
"init_sleep_seconds": args.init_sleep_seconds,
|
||||
"batch_timeout": args.batch_timeout,
|
||||
"init_timeout": args.init_timeout,
|
||||
"shm_threshold_bytes": args.shm_threshold_bytes,
|
||||
"worker_backend": args.worker_backend,
|
||||
"ray_address": args.ray_address,
|
||||
}
|
||||
)
|
||||
|
||||
if args.image_path:
|
||||
if os.path.exists(args.image_path):
|
||||
loaded_image = Image.open(args.image_path).convert("RGB")
|
||||
prompts = [
|
||||
{
|
||||
"prompt": cast(str, p),
|
||||
"multi_modal_data": {"image": loaded_image},
|
||||
}
|
||||
for p in prompts
|
||||
]
|
||||
else:
|
||||
print(f"[Warning] Image path {args.image_path} does not exist.")
|
||||
omni = Omni(model=model_name, **omni_kwargs)
|
||||
|
||||
result = client.generate(
|
||||
prompts,
|
||||
OmniDiffusionSamplingParams(
|
||||
seed=52,
|
||||
need_kv_receive=False,
|
||||
num_inference_steps=args.steps,
|
||||
extra_args={
|
||||
"cfg_text_scale": args.cfg_text_scale,
|
||||
"cfg_img_scale": args.cfg_img_scale,
|
||||
},
|
||||
),
|
||||
)
|
||||
|
||||
# Ensure result is a list for iteration
|
||||
if not isinstance(result, list):
|
||||
omni_outputs = [result]
|
||||
else:
|
||||
omni_outputs = result
|
||||
|
||||
else:
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
omni_kwargs = {}
|
||||
if args.stage_configs_path:
|
||||
omni_kwargs["stage_configs_path"] = args.stage_configs_path
|
||||
|
||||
omni_kwargs.update(
|
||||
{
|
||||
"log_stats": args.log_stats,
|
||||
"init_sleep_seconds": args.init_sleep_seconds,
|
||||
"batch_timeout": args.batch_timeout,
|
||||
"init_timeout": args.init_timeout,
|
||||
"shm_threshold_bytes": args.shm_threshold_bytes,
|
||||
"worker_backend": args.worker_backend,
|
||||
"ray_address": args.ray_address,
|
||||
formatted_prompts = []
|
||||
for p in prompts:
|
||||
if args.modality == "img2img":
|
||||
if not args.image_path or not os.path.exists(args.image_path):
|
||||
raise ValueError(f"img2img requires --image-path pointing to an existing file, got: {args.image_path}")
|
||||
loaded_image = Image.open(args.image_path).convert("RGB")
|
||||
final_prompt_text = f"<|fim_middle|><|im_start|>{p}<|im_end|>"
|
||||
prompt_dict = {
|
||||
"prompt": final_prompt_text,
|
||||
"multi_modal_data": {"img2img": loaded_image},
|
||||
"modalities": ["img2img"],
|
||||
}
|
||||
)
|
||||
|
||||
omni = Omni(model=model_name, **omni_kwargs)
|
||||
|
||||
formatted_prompts = []
|
||||
for p in args.prompts:
|
||||
if args.modality == "img2text":
|
||||
if args.image_path:
|
||||
loaded_image = Image.open(args.image_path).convert("RGB")
|
||||
final_prompt_text = f"<|im_start|>user\n<|image_pad|>\n{p}<|im_end|>\n<|im_start|>assistant\n"
|
||||
prompt_dict = {
|
||||
"prompt": final_prompt_text,
|
||||
"multi_modal_data": {"image": loaded_image},
|
||||
"modalities": ["text"],
|
||||
}
|
||||
formatted_prompts.append(prompt_dict)
|
||||
elif args.modality == "text2text":
|
||||
final_prompt_text = f"<|im_start|>user\n{p}<|im_end|>\n<|im_start|>assistant\n"
|
||||
prompt_dict = {"prompt": final_prompt_text, "modalities": ["text"]}
|
||||
formatted_prompts.append(prompt_dict)
|
||||
else:
|
||||
# text2img
|
||||
final_prompt_text = f"<|im_start|>{p}<|im_end|>"
|
||||
prompt_dict = {"prompt": final_prompt_text, "modalities": ["image"]}
|
||||
formatted_prompts.append(prompt_dict)
|
||||
|
||||
params_list = omni.default_sampling_params_list
|
||||
if args.modality == "text2img":
|
||||
params_list[0].max_tokens = 1 # type: ignore # The first stage is a SamplingParam (vllm)
|
||||
if len(params_list) > 1:
|
||||
diffusion_params = params_list[1]
|
||||
diffusion_params.num_inference_steps = args.steps # type: ignore
|
||||
diffusion_params.extra_args = { # type: ignore
|
||||
"cfg_text_scale": args.cfg_text_scale,
|
||||
"cfg_img_scale": args.cfg_img_scale,
|
||||
if args.negative_prompt is not None:
|
||||
prompt_dict["negative_prompt"] = args.negative_prompt
|
||||
formatted_prompts.append(prompt_dict)
|
||||
elif args.modality == "img2text":
|
||||
if args.image_path:
|
||||
loaded_image = Image.open(args.image_path).convert("RGB")
|
||||
final_prompt_text = f"<|im_start|>user\n<|image_pad|>\n{p}<|im_end|>\n<|im_start|>assistant\n"
|
||||
prompt_dict = {
|
||||
"prompt": final_prompt_text,
|
||||
"multi_modal_data": {"image": loaded_image},
|
||||
"modalities": ["text"],
|
||||
}
|
||||
formatted_prompts.append(prompt_dict)
|
||||
elif args.modality == "text2text":
|
||||
final_prompt_text = f"<|im_start|>user\n{p}<|im_end|>\n<|im_start|>assistant\n"
|
||||
prompt_dict = {"prompt": final_prompt_text, "modalities": ["text"]}
|
||||
formatted_prompts.append(prompt_dict)
|
||||
else:
|
||||
final_prompt_text = f"<|im_start|>{p}<|im_end|>"
|
||||
prompt_dict = {"prompt": final_prompt_text, "modalities": ["image"]}
|
||||
if args.negative_prompt is not None:
|
||||
prompt_dict["negative_prompt"] = args.negative_prompt
|
||||
formatted_prompts.append(prompt_dict)
|
||||
|
||||
omni_outputs = list(omni.generate(prompts=formatted_prompts, sampling_params_list=params_list))
|
||||
params_list = omni.default_sampling_params_list
|
||||
if args.modality in ("text2img", "img2img"):
|
||||
params_list[0].max_tokens = 1 # type: ignore
|
||||
if len(params_list) > 1:
|
||||
diffusion_params = params_list[1]
|
||||
diffusion_params.num_inference_steps = args.steps # type: ignore
|
||||
diffusion_params.cfg_parallel_size = args.cfg_parallel_size # type: ignore
|
||||
if args.seed is not None:
|
||||
diffusion_params.seed = args.seed # type: ignore
|
||||
extra = {
|
||||
"cfg_text_scale": args.cfg_text_scale,
|
||||
"cfg_img_scale": args.cfg_img_scale,
|
||||
}
|
||||
if args.negative_prompt is not None:
|
||||
extra["negative_prompt"] = args.negative_prompt
|
||||
diffusion_params.extra_args = extra # type: ignore
|
||||
|
||||
omni_outputs = list(omni.generate(prompts=formatted_prompts, sampling_params_list=params_list))
|
||||
|
||||
for i, req_output in enumerate(omni_outputs):
|
||||
images = getattr(req_output, "images", None)
|
||||
|
||||
@@ -29,8 +29,8 @@ Accelerated with cache_dit:
|
||||
python image_edit.py \
|
||||
--image input.png \
|
||||
--prompt "Edit description" \
|
||||
--cache_backend cache_dit \
|
||||
--cache_dit_enable_taylorseer
|
||||
--cache-backend cache_dit \
|
||||
--cache-dit-enable-taylorseer
|
||||
|
||||
Layered RGBA output:
|
||||
python image_edit.py \
|
||||
@@ -66,38 +66,38 @@ def parse_args() -> argparse.Namespace:
|
||||
parser.add_argument("--model", default="Qwen/Qwen-Image-Edit", help="Model name or local path.")
|
||||
parser.add_argument("--image", type=str, nargs="+", required=True, help="Input image file(s).")
|
||||
parser.add_argument("--prompt", type=str, required=True, help="Edit description prompt.")
|
||||
parser.add_argument("--negative_prompt", type=str, default=None)
|
||||
parser.add_argument("--negative-prompt", type=str, default=None)
|
||||
parser.add_argument("--seed", type=int, default=0)
|
||||
parser.add_argument("--cfg_scale", type=float, default=4.0)
|
||||
parser.add_argument("--guidance_scale", type=float, default=1.0)
|
||||
parser.add_argument("--cfg-scale", type=float, default=4.0)
|
||||
parser.add_argument("--guidance-scale", type=float, default=1.0)
|
||||
parser.add_argument("--output", type=str, default="output.png")
|
||||
parser.add_argument("--num_outputs_per_prompt", type=int, default=1)
|
||||
parser.add_argument("--num_inference_steps", type=int, default=50)
|
||||
parser.add_argument("--cache_backend", type=str, default=None, choices=["cache_dit", "tea_cache"])
|
||||
parser.add_argument("--ulysses_degree", type=int, default=1)
|
||||
parser.add_argument("--ring_degree", type=int, default=1)
|
||||
parser.add_argument("--tensor_parallel_size", type=int, default=1)
|
||||
parser.add_argument("--num-outputs-per-prompt", type=int, default=1)
|
||||
parser.add_argument("--num-inference-steps", type=int, default=50)
|
||||
parser.add_argument("--cache-backend", type=str, default=None, choices=["cache_dit", "tea_cache"])
|
||||
parser.add_argument("--ulysses-degree", type=int, default=1)
|
||||
parser.add_argument("--ring-degree", type=int, default=1)
|
||||
parser.add_argument("--tensor-parallel-size", type=int, default=1)
|
||||
parser.add_argument("--layers", type=int, default=4)
|
||||
parser.add_argument("--resolution", type=int, default=640)
|
||||
parser.add_argument("--color-format", type=str, default="RGB")
|
||||
|
||||
# Acceleration + Optimization Options
|
||||
parser.add_argument("--cache_dit_fn_compute_blocks", type=int, default=1)
|
||||
parser.add_argument("--cache_dit_bn_compute_blocks", type=int, default=0)
|
||||
parser.add_argument("--cache_dit_max_warmup_steps", type=int, default=4)
|
||||
parser.add_argument("--cache_dit_residual_diff_threshold", type=float, default=0.24)
|
||||
parser.add_argument("--cache_dit_max_continuous_cached_steps", type=int, default=3)
|
||||
parser.add_argument("--cache_dit_enable_taylorseer", action="store_true", default=False)
|
||||
parser.add_argument("--cache_dit_taylorseer_order", type=int, default=1)
|
||||
parser.add_argument("--cache-dit-fn-compute-blocks", type=int, default=1)
|
||||
parser.add_argument("--cache-dit-bn-compute-blocks", type=int, default=0)
|
||||
parser.add_argument("--cache-dit-max-warmup-steps", type=int, default=4)
|
||||
parser.add_argument("--cache-dit-residual-diff-threshold", type=float, default=0.24)
|
||||
parser.add_argument("--cache-dit-max-continuous-cached-steps", type=int, default=3)
|
||||
parser.add_argument("--cache-dit-enable-taylorseer", action="store_true", default=False)
|
||||
parser.add_argument("--cache-dit-taylorseer-order", type=int, default=1)
|
||||
parser.add_argument(
|
||||
"--cache_dit_scm_steps_mask_policy", type=str, default=None, choices=[None, "slow", "medium", "fast", "ultra"]
|
||||
"--cache-dit-scm-steps-mask-policy", type=str, default=None, choices=[None, "slow", "medium", "fast", "ultra"]
|
||||
)
|
||||
parser.add_argument("--cache_dit_scm_steps_policy", type=str, default="dynamic", choices=["dynamic", "static"])
|
||||
parser.add_argument("--tea_cache_rel_l1_thresh", type=float, default=0.2)
|
||||
parser.add_argument("--cfg_parallel_size", type=int, default=1, choices=[1, 2])
|
||||
parser.add_argument("--enforce_eager", action="store_true")
|
||||
parser.add_argument("--vae_use_slicing", action="store_true")
|
||||
parser.add_argument("--vae_use_tiling", action="store_true")
|
||||
parser.add_argument("--cache-dit-scm-steps-policy", type=str, default="dynamic", choices=["dynamic", "static"])
|
||||
parser.add_argument("--tea-cache-rel-l1-thresh", type=float, default=0.2)
|
||||
parser.add_argument("--cfg-parallel-size", type=int, default=1, choices=[1, 2])
|
||||
parser.add_argument("--enforce-eager", action="store_true")
|
||||
parser.add_argument("--vae-use-slicing", action="store_true")
|
||||
parser.add_argument("--vae-use-tiling", action="store_true")
|
||||
parser.add_argument("--enable-cpu-offload", action="store_true")
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
@@ -3,4 +3,4 @@ python image_edit.py \
|
||||
--image cherry_blossom.jpg \
|
||||
--prompt "Let this mascot dance under the moon, surrounded by floating stars and poetic bubbles such as 'Be Kind'" \
|
||||
--output output_image_edit.png \
|
||||
--num_inference_steps 10
|
||||
--num-inference-steps 10
|
||||
@@ -40,8 +40,7 @@ GLM-Image uses a 2-stage pipeline:
|
||||
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model-path /path/to/glm-image \
|
||||
--config-path ../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
|
||||
--config-path ../../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
|
||||
--prompt "A beautiful sunset over the ocean with sailing boats" \
|
||||
--height 1024 \
|
||||
--width 1024 \
|
||||
@@ -52,8 +51,7 @@ python end2end.py \
|
||||
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model-path /path/to/glm-image \
|
||||
--config-path ../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
|
||||
--config-path ../../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
|
||||
--prompt "Transform this scene into a winter wonderland" \
|
||||
--image input.png \
|
||||
--output output_i2i.png
|
||||
@@ -64,7 +62,7 @@ python end2end.py \
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model-path /path/to/glm-image \
|
||||
--config-path ../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
|
||||
--config-path ../../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
|
||||
--prompt "A photorealistic cat sitting on a window sill" \
|
||||
--height 1024 \
|
||||
--width 1024 \
|
||||
|
||||
@@ -178,9 +178,6 @@ def main(args: argparse.Namespace) -> None:
|
||||
print("=" * 60)
|
||||
|
||||
# Validate arguments
|
||||
if not args.model_path:
|
||||
raise ValueError("--model-path is required")
|
||||
|
||||
if not args.prompt:
|
||||
raise ValueError("--prompt is required")
|
||||
|
||||
@@ -365,7 +362,7 @@ def parse_args() -> argparse.Namespace:
|
||||
parser.add_argument(
|
||||
"--model-path",
|
||||
type=str,
|
||||
required=True,
|
||||
default="zai-org/GLM-Image",
|
||||
help="Path to GLM-Image model directory or HuggingFace model ID",
|
||||
)
|
||||
parser.add_argument(
|
||||
|
||||
@@ -0,0 +1,193 @@
|
||||
# Helios Video Generation
|
||||
|
||||
Helios is a text-to-video (T2V), image-to-video (I2V), and video-to-video (V2V) diffusion model. This example demonstrates end-to-end video generation using vLLM-Omni with three model variants:
|
||||
|
||||
| Variant | Description | Key Features |
|
||||
|---------|-------------|--------------|
|
||||
| **Helios-Base** | Base model, Stage 1 only | Single-stage denoising, `guidance_scale=5.0` |
|
||||
| **Helios-Mid** | Mid model, Stage 2 pyramid | Multi-stage pyramid denoising, CFG-Zero* support |
|
||||
| **Helios-Distilled** | Distilled model, Stage 2+3 | Few-step inference with DMD, `guidance_scale=1.0` |
|
||||
|
||||
## Setup
|
||||
|
||||
Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
|
||||
|
||||
## Run Examples
|
||||
|
||||
Get into the example folder:
|
||||
```bash
|
||||
cd examples/offline_inference/helios
|
||||
```
|
||||
|
||||
### Text-to-Video (T2V)
|
||||
|
||||
**Helios-Base** (Stage 1 only):
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Base \
|
||||
--sample-type t2v \
|
||||
--prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \
|
||||
--guidance-scale 5.0 \
|
||||
--output helios_t2v_base.mp4
|
||||
```
|
||||
|
||||
**Helios-Mid** (Stage 2 + CFG-Zero*):
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Mid \
|
||||
--sample-type t2v \
|
||||
--prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \
|
||||
--guidance-scale 5.0 \
|
||||
--is-enable-stage2 \
|
||||
--pyramid-num-inference-steps-list 20 20 20 \
|
||||
--use-cfg-zero-star --use-zero-init --zero-steps 1 \
|
||||
--output helios_t2v_mid.mp4
|
||||
```
|
||||
|
||||
**Helios-Distilled** (Stage 2 pyramid + DMD):
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Distilled \
|
||||
--sample-type t2v \
|
||||
--prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \
|
||||
--num-frames 240 \
|
||||
--guidance-scale 1.0 \
|
||||
--is-enable-stage2 \
|
||||
--pyramid-num-inference-steps-list 2 2 2 \
|
||||
--is-amplify-first-chunk \
|
||||
--output helios_t2v_distilled.mp4
|
||||
```
|
||||
|
||||
### Image-to-Video (I2V)
|
||||
|
||||
**Helios-Base**:
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Base \
|
||||
--sample-type i2v \
|
||||
--image-path /path/to/image.jpg \
|
||||
--prompt "A towering emerald wave surges forward, its crest curling with raw power and energy." \
|
||||
--guidance-scale 5.0 \
|
||||
--output helios_i2v_base.mp4
|
||||
```
|
||||
|
||||
**Helios-Mid** (Stage 2 + CFG-Zero*):
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Mid \
|
||||
--sample-type i2v \
|
||||
--image-path /path/to/image.jpg \
|
||||
--prompt "A towering emerald wave surges forward, its crest curling with raw power and energy." \
|
||||
--guidance-scale 5.0 \
|
||||
--is-enable-stage2 \
|
||||
--pyramid-num-inference-steps-list 20 20 20 \
|
||||
--use-cfg-zero-star --use-zero-init --zero-steps 1 \
|
||||
--output helios_i2v_mid.mp4
|
||||
```
|
||||
|
||||
**Helios-Distilled**:
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Distilled \
|
||||
--sample-type i2v \
|
||||
--image-path /path/to/image.jpg \
|
||||
--prompt "A towering emerald wave surges forward, its crest curling with raw power and energy." \
|
||||
--num-frames 240 \
|
||||
--guidance-scale 1.0 \
|
||||
--is-enable-stage2 \
|
||||
--pyramid-num-inference-steps-list 2 2 2 \
|
||||
--is-amplify-first-chunk \
|
||||
--output helios_i2v_distilled.mp4
|
||||
```
|
||||
|
||||
### Video-to-Video (V2V)
|
||||
|
||||
**Helios-Base**:
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Base \
|
||||
--sample-type v2v \
|
||||
--video-path /path/to/video.mp4 \
|
||||
--prompt "A bright yellow Lamborghini speeds along a curving mountain road." \
|
||||
--guidance-scale 5.0 \
|
||||
--output helios_v2v_base.mp4
|
||||
```
|
||||
|
||||
**Helios-Mid** (Stage 2 + CFG-Zero*):
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Mid \
|
||||
--sample-type v2v \
|
||||
--video-path /path/to/video.mp4 \
|
||||
--prompt "A bright yellow Lamborghini speeds along a curving mountain road." \
|
||||
--guidance-scale 5.0 \
|
||||
--is-enable-stage2 \
|
||||
--pyramid-num-inference-steps-list 20 20 20 \
|
||||
--use-cfg-zero-star --use-zero-init --zero-steps 1 \
|
||||
--output helios_v2v_mid.mp4
|
||||
```
|
||||
|
||||
**Helios-Distilled**:
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Distilled \
|
||||
--sample-type v2v \
|
||||
--video-path /path/to/video.mp4 \
|
||||
--prompt "A bright yellow Lamborghini speeds along a curving mountain road." \
|
||||
--num-frames 240 \
|
||||
--guidance-scale 1.0 \
|
||||
--is-enable-stage2 \
|
||||
--pyramid-num-inference-steps-list 2 2 2 \
|
||||
--is-amplify-first-chunk \
|
||||
--output helios_v2v_distilled.mp4
|
||||
```
|
||||
|
||||
## Common Parameters
|
||||
|
||||
| Parameter | Default | Description |
|
||||
|-----------|---------|-------------|
|
||||
| `--model` | `BestWishYsh/Helios-Base` | Model ID or local path |
|
||||
| `--sample-type` | `t2v` | Generation mode: `t2v`, `i2v`, or `v2v` |
|
||||
| `--prompt` | — | Text prompt describing the video |
|
||||
| `--negative-prompt` | *(see source)* | Negative prompt for CFG (includes anti-static terms) |
|
||||
| `--image-path` | — | Input image (required for `i2v`) |
|
||||
| `--video-path` | — | Input video (required for `v2v`) |
|
||||
| `--height` | `384` | Video height in pixels |
|
||||
| `--width` | `640` | Video width in pixels |
|
||||
| `--num-frames` | `99` | Number of output frames |
|
||||
| `--num-inference-steps` | `50` | Denoising steps (Stage 1 only) |
|
||||
| `--guidance-scale` | `5.0` | CFG scale (`1.0` for Distilled) |
|
||||
| `--seed` | `42` | Random seed |
|
||||
| `--fps` | `16` | Output video frame rate |
|
||||
| `--output` | `helios_output.mp4` | Output file path |
|
||||
|
||||
### Stage 2 / Pyramid Parameters
|
||||
|
||||
| Parameter | Default | Description |
|
||||
|-----------|---------|-------------|
|
||||
| `--is-enable-stage2` | off | Enable pyramid multi-stage denoising |
|
||||
| `--pyramid-num-stages` | `3` | Number of pyramid stages |
|
||||
| `--pyramid-num-inference-steps-list` | `10 10 10` | Steps per pyramid stage |
|
||||
| `--is-amplify-first-chunk` | off | DMD amplification (Distilled) |
|
||||
|
||||
### CFG-Zero* Parameters
|
||||
|
||||
| Parameter | Default | Description |
|
||||
|-----------|---------|-------------|
|
||||
| `--use-cfg-zero-star` | off | Enable CFG-Zero* guidance (Mid) |
|
||||
| `--use-zero-init` | off | Zero init for first steps |
|
||||
| `--zero-steps` | `1` | Number of zero-init steps |
|
||||
|
||||
### Memory & Parallelism
|
||||
|
||||
| Parameter | Default | Description |
|
||||
|-----------|---------|-------------|
|
||||
| `--vae-use-slicing` | off | Enable VAE slicing |
|
||||
| `--vae-use-tiling` | off | Enable VAE tiling |
|
||||
| `--enforce-eager` | off | Disable torch.compile |
|
||||
| `--enable-cpu-offload` | off | CPU offloading |
|
||||
| `--enable-layerwise-offload` | off | Layerwise offloading |
|
||||
| `--ulysses-degree` | `1` | Ulysses SP degree |
|
||||
| `--ring-degree` | `1` | Ring SP degree |
|
||||
| `--cfg-parallel-size` | `1` | CFG parallel size (1 or 2) |
|
||||
| `--tensor-parallel-size` | `1` | Tensor parallelism size |
|
||||
@@ -0,0 +1,343 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
Helios video generation example supporting T2V, I2V, and V2V.
|
||||
|
||||
Usage (T2V, Helios-Base, Stage 1 only):
|
||||
python end2end.py \
|
||||
--model /path/to/Helios-Base --sample-type t2v \
|
||||
--prompt "A serene lakeside sunrise with mist over the water." \
|
||||
--height 384 --width 640 --num-frames 99 \
|
||||
--num-inference-steps 50 --guidance-scale 5.0
|
||||
|
||||
Usage (I2V, Helios-Base):
|
||||
python end2end.py \
|
||||
--model /path/to/Helios-Base --sample-type i2v \
|
||||
--image-path /path/to/image.jpg \
|
||||
--prompt "Description of desired animation." \
|
||||
--guidance-scale 5.0
|
||||
|
||||
Usage (V2V, Helios-Base):
|
||||
python end2end.py \
|
||||
--model /path/to/Helios-Base --sample-type v2v \
|
||||
--video-path /path/to/video.mp4 \
|
||||
--prompt "Description of desired transformation." \
|
||||
--guidance-scale 5.0
|
||||
|
||||
Usage (Helios-Mid, Stage 2 + CFG-Zero*):
|
||||
python end2end.py \
|
||||
--model /path/to/Helios-Mid --sample-type t2v \
|
||||
--prompt "A serene lakeside sunrise with mist over the water." \
|
||||
--guidance-scale 5.0 \
|
||||
--is-enable-stage2 \
|
||||
--pyramid-num-inference-steps-list 20 20 20 \
|
||||
--use-cfg-zero-star --use-zero-init --zero-steps 1
|
||||
|
||||
Usage (Helios-Distilled, Stage 2 pyramid + DMD):
|
||||
python end2end.py \
|
||||
--model /path/to/Helios-Distilled --sample-type t2v \
|
||||
--prompt "A serene lakeside sunrise with mist over the water." \
|
||||
--guidance-scale 1.0 \
|
||||
--is-enable-stage2 \
|
||||
--pyramid-num-inference-steps-list 2 2 2 \
|
||||
--is-amplify-first-chunk
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
|
||||
def load_image_as_tensor(image_path: str, height: int, width: int) -> torch.Tensor:
|
||||
"""Load an image and return a (1, C, H, W) tensor normalized to [-1, 1]."""
|
||||
from PIL import Image
|
||||
from torchvision import transforms
|
||||
|
||||
image = Image.open(image_path).convert("RGB").resize((width, height))
|
||||
transform = transforms.Compose(
|
||||
[
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
|
||||
]
|
||||
)
|
||||
return transform(image).unsqueeze(0)
|
||||
|
||||
|
||||
def load_video_as_tensor(video_path: str, height: int, width: int) -> torch.Tensor:
|
||||
"""Load a video and return a (1, C, T, H, W) tensor normalized to [-1, 1]."""
|
||||
from torchvision import transforms
|
||||
from torchvision.io import read_video
|
||||
|
||||
video_frames, _, _ = read_video(video_path, output_format="TCHW")
|
||||
video_frames = video_frames.float() / 255.0
|
||||
|
||||
transform = transforms.Compose(
|
||||
[
|
||||
transforms.Resize((height, width), antialias=True),
|
||||
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
|
||||
]
|
||||
)
|
||||
video_frames = torch.stack([transform(f) for f in video_frames])
|
||||
return video_frames.permute(1, 0, 2, 3).unsqueeze(0)
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Generate a video with Helios (T2V / I2V / V2V).")
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
default="BestWishYsh/Helios-Base",
|
||||
help="Helios model ID or local path (e.g. Helios-Base, Helios-Mid, Helios-Distilled).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--sample-type",
|
||||
choices=["t2v", "i2v", "v2v"],
|
||||
default="t2v",
|
||||
help="Generation mode: t2v (text-to-video), i2v (image-to-video), v2v (video-to-video).",
|
||||
)
|
||||
parser.add_argument("--prompt", default="A serene lakeside sunrise with mist over the water.", help="Text prompt.")
|
||||
parser.add_argument(
|
||||
"--negative-prompt",
|
||||
default=(
|
||||
"Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, "
|
||||
"images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, "
|
||||
"incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, "
|
||||
"misshapen limbs, fused fingers, still picture, messy background, three legs, many people "
|
||||
"in the background, walking backwards"
|
||||
),
|
||||
help="Negative prompt.",
|
||||
)
|
||||
parser.add_argument("--seed", type=int, default=42, help="Random seed.")
|
||||
parser.add_argument("--guidance-scale", type=float, default=5.0, help="CFG scale.")
|
||||
parser.add_argument("--height", type=int, default=384, help="Video height.")
|
||||
parser.add_argument("--width", type=int, default=640, help="Video width.")
|
||||
parser.add_argument("--num-frames", type=int, default=99, help="Number of video frames.")
|
||||
parser.add_argument("--num-inference-steps", type=int, default=50, help="Sampling steps (Stage 1 only).")
|
||||
parser.add_argument("--output", type=str, default="helios_output.mp4", help="Output video path.")
|
||||
parser.add_argument("--fps", type=int, default=16, help="Frames per second for the output video.")
|
||||
|
||||
# I2V / V2V
|
||||
parser.add_argument("--image-path", type=str, default=None, help="Input image path for I2V mode.")
|
||||
parser.add_argument("--video-path", type=str, default=None, help="Input video path for V2V mode.")
|
||||
|
||||
# Stage 2 (pyramid multi-stage denoising)
|
||||
parser.add_argument(
|
||||
"--is-enable-stage2",
|
||||
action="store_true",
|
||||
help="Enable pyramid multi-stage denoising (Stage 2). Required for Helios-Distilled.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--pyramid-num-stages",
|
||||
type=int,
|
||||
default=3,
|
||||
help="Number of pyramid stages for Stage 2.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--pyramid-num-inference-steps-list",
|
||||
type=int,
|
||||
nargs="+",
|
||||
default=[10, 10, 10],
|
||||
help="Inference steps per pyramid stage.",
|
||||
)
|
||||
|
||||
# DMD
|
||||
parser.add_argument(
|
||||
"--is-amplify-first-chunk",
|
||||
action="store_true",
|
||||
help="Enable DMD amplification for the first chunk (Helios-Distilled).",
|
||||
)
|
||||
|
||||
# CFG Zero Star
|
||||
parser.add_argument(
|
||||
"--use-cfg-zero-star",
|
||||
action="store_true",
|
||||
help="Enable CFG Zero Star guidance (recommended for Helios-Mid).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--use-zero-init",
|
||||
action="store_true",
|
||||
help="Use zero initialization for the first denoising steps with CFG-Zero*.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--zero-steps",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of initial denoising steps using zero prediction (default: 1).",
|
||||
)
|
||||
|
||||
# Memory & parallelism
|
||||
parser.add_argument("--vae-use-slicing", action="store_true", help="Enable VAE slicing.")
|
||||
parser.add_argument("--vae-use-tiling", action="store_true", help="Enable VAE tiling.")
|
||||
parser.add_argument("--enforce-eager", action="store_true", help="Disable torch.compile.")
|
||||
parser.add_argument("--enable-cpu-offload", action="store_true", help="Enable CPU offloading.")
|
||||
parser.add_argument("--enable-layerwise-offload", action="store_true", help="Enable layerwise offloading.")
|
||||
parser.add_argument("--ulysses-degree", type=int, default=1, help="Ulysses SP degree.")
|
||||
parser.add_argument("--ring-degree", type=int, default=1, help="Ring SP degree.")
|
||||
parser.add_argument("--cfg-parallel-size", type=int, default=1, choices=[1, 2], help="CFG parallel size.")
|
||||
parser.add_argument("--tensor-parallel-size", type=int, default=1, help="Tensor parallelism size.")
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
|
||||
|
||||
parallel_config = DiffusionParallelConfig(
|
||||
ulysses_degree=args.ulysses_degree,
|
||||
ring_degree=args.ring_degree,
|
||||
cfg_parallel_size=args.cfg_parallel_size,
|
||||
tensor_parallel_size=args.tensor_parallel_size,
|
||||
)
|
||||
|
||||
omni = Omni(
|
||||
model=args.model,
|
||||
enable_layerwise_offload=args.enable_layerwise_offload,
|
||||
vae_use_slicing=args.vae_use_slicing,
|
||||
vae_use_tiling=args.vae_use_tiling,
|
||||
enable_cpu_offload=args.enable_cpu_offload,
|
||||
parallel_config=parallel_config,
|
||||
enforce_eager=args.enforce_eager,
|
||||
)
|
||||
|
||||
# Validate I2V / V2V arguments
|
||||
if args.sample_type == "i2v" and not args.image_path:
|
||||
raise ValueError("--image-path is required for I2V mode.")
|
||||
if args.sample_type == "v2v" and not args.video_path:
|
||||
raise ValueError("--video-path is required for V2V mode.")
|
||||
|
||||
# Build extra_args for Helios-specific parameters
|
||||
extra_args = {}
|
||||
if args.is_enable_stage2:
|
||||
extra_args["is_enable_stage2"] = True
|
||||
extra_args["pyramid_num_stages"] = args.pyramid_num_stages
|
||||
extra_args["pyramid_num_inference_steps_list"] = args.pyramid_num_inference_steps_list
|
||||
if args.is_amplify_first_chunk:
|
||||
extra_args["is_amplify_first_chunk"] = True
|
||||
if args.use_cfg_zero_star:
|
||||
extra_args["use_cfg_zero_star"] = True
|
||||
if args.use_zero_init:
|
||||
extra_args["use_zero_init"] = True
|
||||
extra_args["zero_steps"] = args.zero_steps
|
||||
|
||||
if args.sample_type == "i2v":
|
||||
image_tensor = load_image_as_tensor(args.image_path, args.height, args.width)
|
||||
extra_args["image"] = image_tensor
|
||||
elif args.sample_type == "v2v":
|
||||
video_tensor = load_video_as_tensor(args.video_path, args.height, args.width)
|
||||
extra_args["video"] = video_tensor
|
||||
|
||||
# Print generation configuration
|
||||
print(f"\n{'=' * 60}")
|
||||
print("Helios Generation Configuration:")
|
||||
print(f" Model: {args.model}")
|
||||
print(f" Sample type: {args.sample_type.upper()}")
|
||||
if args.sample_type == "i2v":
|
||||
print(f" Image: {args.image_path}")
|
||||
elif args.sample_type == "v2v":
|
||||
print(f" Video: {args.video_path}")
|
||||
print(f" Prompt: {args.prompt[:80]}{'...' if len(args.prompt) > 80 else ''}")
|
||||
print(f" Video size: {args.width}x{args.height}, {args.num_frames} frames")
|
||||
print(f" Inference steps: {args.num_inference_steps}")
|
||||
print(f" Guidance scale: {args.guidance_scale}")
|
||||
if args.is_enable_stage2:
|
||||
print(f" Stage 2: enabled (stages={args.pyramid_num_stages}, steps={args.pyramid_num_inference_steps_list})")
|
||||
if args.is_amplify_first_chunk:
|
||||
print(" DMD amplify first chunk: enabled")
|
||||
if args.use_cfg_zero_star:
|
||||
print(f" CFG Zero Star: enabled (zero_init={args.use_zero_init}, zero_steps={args.zero_steps})")
|
||||
else:
|
||||
if args.use_cfg_zero_star:
|
||||
print(f" CFG Zero Star: enabled (zero_init={args.use_zero_init}, zero_steps={args.zero_steps})")
|
||||
print(" Stage 2: disabled (Stage 1 only)")
|
||||
print(f"{'=' * 60}\n")
|
||||
|
||||
generation_start = time.perf_counter()
|
||||
frames = omni.generate(
|
||||
{
|
||||
"prompt": args.prompt,
|
||||
"negative_prompt": args.negative_prompt,
|
||||
},
|
||||
OmniDiffusionSamplingParams(
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
generator=generator,
|
||||
guidance_scale=args.guidance_scale,
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
num_frames=args.num_frames,
|
||||
extra_args=extra_args,
|
||||
),
|
||||
)
|
||||
generation_end = time.perf_counter()
|
||||
generation_time = generation_end - generation_start
|
||||
|
||||
print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)")
|
||||
|
||||
# Extract video frames from OmniRequestOutput
|
||||
if isinstance(frames, list) and len(frames) > 0:
|
||||
first_item = frames[0]
|
||||
|
||||
if hasattr(first_item, "final_output_type"):
|
||||
if first_item.final_output_type != "image":
|
||||
raise ValueError(
|
||||
f"Unexpected output type '{first_item.final_output_type}', expected 'image' for video generation."
|
||||
)
|
||||
|
||||
if hasattr(first_item, "is_pipeline_output") and first_item.is_pipeline_output:
|
||||
if isinstance(first_item.request_output, list) and len(first_item.request_output) > 0:
|
||||
inner_output = first_item.request_output[0]
|
||||
if isinstance(inner_output, OmniRequestOutput) and hasattr(inner_output, "images"):
|
||||
frames = inner_output.images[0] if inner_output.images else None
|
||||
if frames is None:
|
||||
raise ValueError("No video frames found in output.")
|
||||
elif hasattr(first_item, "images") and first_item.images:
|
||||
frames = first_item.images
|
||||
else:
|
||||
raise ValueError("No video frames found in OmniRequestOutput.")
|
||||
|
||||
# Unwrap batch list from postprocess_video: [numpy(T,H,W,C)] -> numpy(T,H,W,C)
|
||||
if isinstance(frames, list) and len(frames) > 0 and isinstance(frames[0], np.ndarray):
|
||||
frames = frames[0]
|
||||
|
||||
output_path = Path(args.output)
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
try:
|
||||
from diffusers.utils import export_to_video
|
||||
except ImportError:
|
||||
raise ImportError("diffusers is required for export_to_video.")
|
||||
|
||||
if isinstance(frames, torch.Tensor):
|
||||
video_tensor = frames.detach().cpu()
|
||||
if video_tensor.dim() == 5:
|
||||
if video_tensor.shape[1] in (3, 4):
|
||||
video_tensor = video_tensor[0].permute(1, 2, 3, 0)
|
||||
else:
|
||||
video_tensor = video_tensor[0]
|
||||
elif video_tensor.dim() == 4 and video_tensor.shape[0] in (3, 4):
|
||||
video_tensor = video_tensor.permute(1, 2, 3, 0)
|
||||
if video_tensor.is_floating_point():
|
||||
video_tensor = video_tensor.clamp(-1, 1) * 0.5 + 0.5
|
||||
video_array = video_tensor.float().numpy()
|
||||
elif isinstance(frames, np.ndarray):
|
||||
video_array = frames
|
||||
else:
|
||||
video_array = frames
|
||||
|
||||
if isinstance(video_array, np.ndarray) and video_array.ndim == 5:
|
||||
video_array = video_array[0]
|
||||
|
||||
export_to_video(video_array, str(output_path), fps=args.fps)
|
||||
print(f"Saved generated video to {output_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,25 @@
|
||||
# HunyuanImage-3.0 Image-to-Text Inference
|
||||
|
||||
This example demonstrates how to run HunyuanImage-3.0 Image-to-Text with the vLLM-Omni.
|
||||
|
||||
## Local CLI Usage
|
||||
|
||||
Download the example image:
|
||||
|
||||
```bash
|
||||
wget https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg
|
||||
```
|
||||
|
||||
Run example:
|
||||
|
||||
```bash
|
||||
python image_to_text.py \
|
||||
--image cherry_blossom.jpg \
|
||||
--prompt "<|startoftext|>You are an assistant that understands images and outputs text.<img>Describe the content of the picture."
|
||||
```
|
||||
|
||||
Key arguments:
|
||||
|
||||
- `--model`: Model used. Default is: tencent/HunyuanImage-3.0-Instruct (Optional).
|
||||
- `--image`: Path to input image (required).
|
||||
- `--prompt`: Text description used to guide image understanding (required).
|
||||
@@ -0,0 +1,81 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import argparse
|
||||
import os
|
||||
|
||||
from PIL import Image
|
||||
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
"""
|
||||
The tencent/HunyuanImage-3.0-Instruct base model is built on the Hunyuan v1 architecture, specifically the tencent/Hunyuan-A13B-Instruct model. It utilizes two tokenizer delimiter templates:
|
||||
|
||||
1) Pretrained template (default for gen_text mode), which concatenates system, image
|
||||
tokens, and user question WITHOUT role delimiters:
|
||||
"<|startoftext|>{system_prompt}{image_tokens}{user_question}"
|
||||
|
||||
Example (before image token expansion):
|
||||
"<|startoftext|>You are an assistant that understands images and outputs text.<img>Describe the content of the picture."
|
||||
|
||||
2) Instruct template, which uses explicit role prefixes and separators.
|
||||
"""
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Generate text from image using HunyuanImage-3.0-Instruct.")
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
default="tencent/HunyuanImage-3.0-Instruct",
|
||||
help="Model name or local path.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--image",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Path to input image file (PNG, JPG, etc.).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompt",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Pretrain template prompt: <|startoftext|>{system}<img>{question}",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def load_image(image_path: str) -> Image.Image:
|
||||
"""Load an image from file path."""
|
||||
if not os.path.exists(image_path):
|
||||
raise FileNotFoundError(f"Image file not found: {image_path}")
|
||||
return Image.open(image_path).convert("RGB")
|
||||
|
||||
|
||||
def main(args: argparse.Namespace) -> None:
|
||||
omni = Omni(model=args.model)
|
||||
|
||||
prompt_dict = {
|
||||
"prompt": args.prompt,
|
||||
"modalities": ["text"],
|
||||
}
|
||||
|
||||
# Add image input if provided
|
||||
if args.image:
|
||||
if not os.path.exists(args.image):
|
||||
raise FileNotFoundError(f"Input image not found: {args.image}")
|
||||
|
||||
input_image = load_image(args.image)
|
||||
prompt_dict["multi_modal_data"] = {"image": input_image}
|
||||
|
||||
prompts = [prompt_dict]
|
||||
omni_outputs = omni.generate(prompts=prompts)
|
||||
|
||||
prompt_text = omni_outputs[0].request_output[0].prompt
|
||||
generated_text = omni_outputs[0].request_output[0].outputs[0].text
|
||||
print(f"Prompt: {prompt_text}")
|
||||
print(f"Text: {generated_text}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
main(args)
|
||||
@@ -208,6 +208,11 @@ def parse_args() -> argparse.Namespace:
|
||||
default=1,
|
||||
help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enable-expert-parallel",
|
||||
action="store_true",
|
||||
help="Enable expert parallelism for MoE layers.",
|
||||
)
|
||||
parser.add_argument("--layers", type=int, default=4, help="Number of layers to decompose the input image into.")
|
||||
parser.add_argument(
|
||||
"--resolution",
|
||||
@@ -348,6 +353,7 @@ def main():
|
||||
ring_degree=args.ring_degree,
|
||||
cfg_parallel_size=args.cfg_parallel_size,
|
||||
tensor_parallel_size=args.tensor_parallel_size,
|
||||
enable_expert_parallel=args.enable_expert_parallel,
|
||||
)
|
||||
|
||||
# Configure cache based on backend type
|
||||
@@ -402,7 +408,7 @@ def main():
|
||||
else:
|
||||
print(f" Input image size: {input_image.size}")
|
||||
print(
|
||||
f" Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}, tensor_parallel_size={args.tensor_parallel_size}"
|
||||
f" Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}, tensor_parallel_size={args.tensor_parallel_size}, enable_expert_parallel: {args.enable_expert_parallel}"
|
||||
)
|
||||
print(f"{'=' * 60}\n")
|
||||
|
||||
|
||||
@@ -65,6 +65,7 @@ Key arguments:
|
||||
- `--vae-use-slicing`: Enable VAE slicing for memory optimization.
|
||||
- `--vae-use-tiling`: Enable VAE tiling for memory optimization.
|
||||
- `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
|
||||
- `--tensor-parallel-size`: tensor parallel size (effective for models that support TP, e.g. LTX2).
|
||||
- `--enable-cpu-offload`: enable CPU offloading for diffusion models.
|
||||
- `--use-hsdp`: Enable Hybrid Sharded Data Parallel to shard model weights across GPUs.
|
||||
- `--hsdp-shard-size`: Number of GPUs to shard model weights across within each replica group. -1 (default) auto-calculates as world_size / replicate_size.
|
||||
|
||||
@@ -2,11 +2,12 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
Image-to-Video generation example using Wan2.2 I2V/TI2V models.
|
||||
Image-to-Video generation example using Wan2.2 I2V/TI2V models or LTX2.
|
||||
|
||||
Supports:
|
||||
- Wan2.2-I2V-A14B-Diffusers: MoE model with CLIP image encoder
|
||||
- Wan2.2-TI2V-5B-Diffusers: Unified T2V+I2V model (dense 5B)
|
||||
- LTX2 image-to-video pipeline
|
||||
|
||||
Usage:
|
||||
# I2V-A14B (MoE)
|
||||
@@ -16,10 +17,18 @@ Usage:
|
||||
# TI2V-5B (unified)
|
||||
python image_to_video.py --model Wan-AI/Wan2.2-TI2V-5B-Diffusers \
|
||||
--image input.jpg --prompt "A cat playing with yarn"
|
||||
|
||||
# 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 \
|
||||
--frame-rate 24 --fps 24 --output ltx2_i2v.mp4
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
@@ -34,12 +43,17 @@ from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Generate a video from an image with Wan2.2 I2V/TI2V.")
|
||||
parser = argparse.ArgumentParser(description="Generate a video from an image with Wan2.2 or LTX2.")
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
default="Wan-AI/Wan2.2-I2V-A14B-Diffusers",
|
||||
help="Diffusers Wan2.2 I2V model ID or local path.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model-class-name",
|
||||
default=None,
|
||||
help="Override model class name (e.g., LTX2ImageToVideoPipeline).",
|
||||
)
|
||||
parser.add_argument("--image", required=True, help="Path to input image.")
|
||||
parser.add_argument("--prompt", default="", help="Text prompt describing the desired motion.")
|
||||
parser.add_argument("--negative-prompt", default="", help="Negative prompt.")
|
||||
@@ -55,11 +69,17 @@ def parse_args() -> argparse.Namespace:
|
||||
parser.add_argument("--num-frames", type=int, default=81, help="Number of frames.")
|
||||
parser.add_argument("--num-inference-steps", type=int, default=50, help="Sampling steps.")
|
||||
parser.add_argument("--boundary-ratio", type=float, default=0.875, help="Boundary split ratio for MoE models.")
|
||||
parser.add_argument(
|
||||
"--frame-rate",
|
||||
type=float,
|
||||
default=None,
|
||||
help="Optional generation frame rate (used by models like LTX2). Defaults to --fps.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--flow-shift", type=float, default=5.0, help="Scheduler flow_shift (5.0 for 720p, 12.0 for 480p)."
|
||||
)
|
||||
parser.add_argument("--output", type=str, default="i2v_output.mp4", help="Path to save the video (mp4).")
|
||||
parser.add_argument("--fps", type=int, default=16, help="Frames per second for the output video.")
|
||||
parser.add_argument("--fps", type=int, default=None, help="Frames per second for the output video.")
|
||||
parser.add_argument(
|
||||
"--vae-use-slicing",
|
||||
action="store_true",
|
||||
@@ -105,11 +125,34 @@ def parse_args() -> argparse.Namespace:
|
||||
choices=[1, 2],
|
||||
help="Number of GPUs used for classifier free guidance parallel size.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--vae-patch-parallel-size",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of GPUs used for VAE patch/tile parallelism (decode).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enforce-eager",
|
||||
action="store_true",
|
||||
help="Disable torch.compile and force eager execution.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--audio-sample-rate",
|
||||
type=int,
|
||||
default=24000,
|
||||
help="Sample rate for audio output when saved (default: 24000 for LTX2).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--cache-backend",
|
||||
type=str,
|
||||
default=None,
|
||||
choices=["cache_dit", "tea_cache"],
|
||||
help=(
|
||||
"Cache backend to use for acceleration. "
|
||||
"Options: 'cache_dit' (DBCache + SCM + TaylorSeer), 'tea_cache' (Timestep Embedding Aware Cache). "
|
||||
"Default: None (no cache acceleration)."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--use-hsdp",
|
||||
action="store_true",
|
||||
@@ -136,10 +179,13 @@ def parse_args() -> argparse.Namespace:
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def calculate_dimensions(image: PIL.Image.Image, max_area: int = 480 * 832) -> tuple[int, int]:
|
||||
def calculate_dimensions(
|
||||
image: PIL.Image.Image,
|
||||
max_area: int = 480 * 832,
|
||||
mod_value: int = 16,
|
||||
) -> tuple[int, int]:
|
||||
"""Calculate output dimensions maintaining aspect ratio."""
|
||||
aspect_ratio = image.height / image.width
|
||||
mod_value = 16 # Must be divisible by 16
|
||||
|
||||
height = round(np.sqrt(max_area * aspect_ratio)) // mod_value * mod_value
|
||||
width = round(np.sqrt(max_area / aspect_ratio)) // mod_value * mod_value
|
||||
@@ -150,22 +196,66 @@ def calculate_dimensions(image: PIL.Image.Image, max_area: int = 480 * 832) -> t
|
||||
def main():
|
||||
args = parse_args()
|
||||
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 = "ltx2" in model_name or (model_class_name and "ltx2" in model_class_name.lower())
|
||||
if model_class_name is None and is_ltx2:
|
||||
model_class_name = "LTX2ImageToVideoPipeline"
|
||||
|
||||
# Load input image
|
||||
image = PIL.Image.open(args.image).convert("RGB")
|
||||
|
||||
fps = args.fps if args.fps is not None else (24 if is_ltx2 else 16)
|
||||
frame_rate = args.frame_rate if args.frame_rate is not None else float(fps)
|
||||
guidance_scale = args.guidance_scale if args.guidance_scale is not None else (4.0 if is_ltx2 else 5.0)
|
||||
num_frames = args.num_frames if args.num_frames is not None else (121 if is_ltx2 else 81)
|
||||
num_inference_steps = args.num_inference_steps if args.num_inference_steps is not None else (40 if is_ltx2 else 50)
|
||||
|
||||
# Calculate dimensions if not provided
|
||||
height = args.height
|
||||
width = args.width
|
||||
if height is None or width is None:
|
||||
# Default to 480P area for I2V
|
||||
calc_height, calc_width = calculate_dimensions(image, max_area=480 * 832)
|
||||
# Default to 480P area for Wan2.2 I2V, 512x768 area for LTX2
|
||||
max_area = 512 * 768 if is_ltx2 else 480 * 832
|
||||
mod_value = 32 if is_ltx2 else 16
|
||||
calc_height, calc_width = calculate_dimensions(image, max_area=max_area, mod_value=mod_value)
|
||||
height = height or calc_height
|
||||
width = width or calc_width
|
||||
|
||||
# Resize image to target dimensions
|
||||
image = image.resize((width, height), PIL.Image.Resampling.LANCZOS)
|
||||
|
||||
# Configure cache based on backend type
|
||||
cache_config = None
|
||||
if args.cache_backend == "cache_dit":
|
||||
if is_ltx2:
|
||||
cache_config = {
|
||||
"Fn_compute_blocks": 2,
|
||||
"Bn_compute_blocks": 0,
|
||||
"max_warmup_steps": 8,
|
||||
"residual_diff_threshold": 0.12,
|
||||
"max_continuous_cached_steps": 1,
|
||||
"max_cached_steps": 20,
|
||||
"enable_taylorseer": False,
|
||||
"scm_steps_mask_policy": None,
|
||||
}
|
||||
else:
|
||||
cache_config = {
|
||||
"Fn_compute_blocks": 1,
|
||||
"Bn_compute_blocks": 0,
|
||||
"max_warmup_steps": 4,
|
||||
"residual_diff_threshold": 0.24,
|
||||
"max_continuous_cached_steps": 3,
|
||||
"enable_taylorseer": False,
|
||||
"taylorseer_order": 1,
|
||||
"scm_steps_mask_policy": None,
|
||||
"scm_steps_policy": "dynamic",
|
||||
}
|
||||
elif args.cache_backend == "tea_cache":
|
||||
cache_config = {
|
||||
"rel_l1_thresh": 0.2,
|
||||
}
|
||||
|
||||
# Check if profiling is requested via environment variable
|
||||
profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
|
||||
parallel_config = DiffusionParallelConfig(
|
||||
@@ -173,6 +263,7 @@ def main():
|
||||
ring_degree=args.ring_degree,
|
||||
cfg_parallel_size=args.cfg_parallel_size,
|
||||
tensor_parallel_size=args.tensor_parallel_size,
|
||||
vae_patch_parallel_size=args.vae_patch_parallel_size,
|
||||
use_hsdp=args.use_hsdp,
|
||||
hsdp_shard_size=args.hsdp_shard_size,
|
||||
hsdp_replicate_size=args.hsdp_replicate_size,
|
||||
@@ -187,6 +278,9 @@ def main():
|
||||
enable_cpu_offload=args.enable_cpu_offload,
|
||||
parallel_config=parallel_config,
|
||||
enforce_eager=args.enforce_eager,
|
||||
model_class_name=model_class_name,
|
||||
cache_backend=args.cache_backend,
|
||||
cache_config=cache_config,
|
||||
)
|
||||
|
||||
if profiler_enabled:
|
||||
@@ -199,10 +293,14 @@ def main():
|
||||
print(f" Model: {args.model}")
|
||||
print(f" Inference steps: {args.num_inference_steps}")
|
||||
print(f" Frames: {args.num_frames}")
|
||||
print(f" Parallel configuration: cfg_parallel_size={args.cfg_parallel_size}")
|
||||
print(
|
||||
f" Parallel configuration: cfg_parallel_size={args.cfg_parallel_size},"
|
||||
f" tensor_parallel_size={args.tensor_parallel_size}, vae_patch_parallel_size={args.vae_patch_parallel_size}"
|
||||
)
|
||||
print(f" Video size: {args.width}x{args.height}")
|
||||
print(f"{'=' * 60}\n")
|
||||
|
||||
generation_start = time.perf_counter()
|
||||
# omni.generate() returns Generator[OmniRequestOutput, None, None]
|
||||
frames = omni.generate(
|
||||
{
|
||||
@@ -214,38 +312,69 @@ def main():
|
||||
height=height,
|
||||
width=width,
|
||||
generator=generator,
|
||||
guidance_scale=args.guidance_scale,
|
||||
guidance_scale=guidance_scale,
|
||||
guidance_scale_2=args.guidance_scale_high,
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
num_frames=args.num_frames,
|
||||
num_inference_steps=num_inference_steps,
|
||||
num_frames=num_frames,
|
||||
frame_rate=frame_rate,
|
||||
),
|
||||
)
|
||||
generation_end = time.perf_counter()
|
||||
generation_time = generation_end - generation_start
|
||||
|
||||
# Extract video frames from OmniRequestOutput
|
||||
if isinstance(frames, list) and len(frames) > 0:
|
||||
first_item = frames[0]
|
||||
# Print profiling results
|
||||
print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)")
|
||||
|
||||
# Check if it's an OmniRequestOutput
|
||||
if hasattr(first_item, "final_output_type"):
|
||||
if first_item.final_output_type != "image":
|
||||
raise ValueError(
|
||||
f"Unexpected output type '{first_item.final_output_type}', expected 'image' for video generation."
|
||||
)
|
||||
audio = None
|
||||
if isinstance(frames, list):
|
||||
frames = frames[0] if frames else None
|
||||
|
||||
# Pipeline mode: extract from nested request_output
|
||||
if hasattr(first_item, "is_pipeline_output") and first_item.is_pipeline_output:
|
||||
if isinstance(first_item.request_output, list) and len(first_item.request_output) > 0:
|
||||
inner_output = first_item.request_output[0]
|
||||
if isinstance(inner_output, OmniRequestOutput) and hasattr(inner_output, "images"):
|
||||
frames = inner_output.images[0] if inner_output.images else None
|
||||
if frames is None:
|
||||
raise ValueError("No video frames found in output.")
|
||||
# Diffusion mode: use direct images field
|
||||
elif hasattr(first_item, "images") and first_item.images:
|
||||
frames = first_item.images
|
||||
if isinstance(frames, OmniRequestOutput):
|
||||
if frames.final_output_type != "image":
|
||||
raise ValueError(
|
||||
f"Unexpected output type '{frames.final_output_type}', expected 'image' for video generation."
|
||||
)
|
||||
if frames.multimodal_output and "audio" in frames.multimodal_output:
|
||||
audio = frames.multimodal_output["audio"]
|
||||
if frames.is_pipeline_output and frames.request_output is not None:
|
||||
inner_output = frames.request_output
|
||||
if isinstance(inner_output, list):
|
||||
inner_output = inner_output[0] if inner_output else None
|
||||
if isinstance(inner_output, OmniRequestOutput):
|
||||
if inner_output.multimodal_output and "audio" in inner_output.multimodal_output:
|
||||
audio = inner_output.multimodal_output["audio"]
|
||||
frames = inner_output
|
||||
if isinstance(frames, OmniRequestOutput):
|
||||
if frames.images:
|
||||
if len(frames.images) == 1 and isinstance(frames.images[0], tuple) and len(frames.images[0]) == 2:
|
||||
frames, audio = frames.images[0]
|
||||
elif len(frames.images) == 1 and isinstance(frames.images[0], dict):
|
||||
audio = frames.images[0].get("audio")
|
||||
frames = frames.images[0].get("frames") or frames.images[0].get("video")
|
||||
else:
|
||||
frames = frames.images
|
||||
else:
|
||||
raise ValueError("No video frames found in OmniRequestOutput.")
|
||||
|
||||
if isinstance(frames, list) and frames:
|
||||
first_item = frames[0]
|
||||
if isinstance(first_item, tuple) and len(first_item) == 2:
|
||||
frames, audio = first_item
|
||||
elif isinstance(first_item, dict):
|
||||
audio = first_item.get("audio")
|
||||
frames = first_item.get("frames") or first_item.get("video")
|
||||
elif isinstance(first_item, list):
|
||||
frames = first_item
|
||||
|
||||
if isinstance(frames, tuple) and len(frames) == 2:
|
||||
frames, audio = frames
|
||||
elif isinstance(frames, dict):
|
||||
audio = frames.get("audio")
|
||||
frames = frames.get("frames") or frames.get("video")
|
||||
|
||||
if frames is None:
|
||||
raise ValueError("No video frames found in output.")
|
||||
|
||||
output_path = Path(args.output)
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
@@ -254,32 +383,128 @@ def main():
|
||||
except ImportError:
|
||||
raise ImportError("diffusers is required for export_to_video.")
|
||||
|
||||
# frames may be np.ndarray (preferred) or torch.Tensor
|
||||
def _normalize_frame(frame):
|
||||
if isinstance(frame, torch.Tensor):
|
||||
frame_tensor = frame.detach().cpu()
|
||||
if frame_tensor.dim() == 4 and frame_tensor.shape[0] == 1:
|
||||
frame_tensor = frame_tensor[0]
|
||||
if frame_tensor.dim() == 3 and frame_tensor.shape[0] in (3, 4):
|
||||
frame_tensor = frame_tensor.permute(1, 2, 0)
|
||||
if frame_tensor.is_floating_point():
|
||||
frame_tensor = frame_tensor.clamp(-1, 1) * 0.5 + 0.5
|
||||
return frame_tensor.float().numpy()
|
||||
if isinstance(frame, np.ndarray):
|
||||
frame_array = frame
|
||||
if frame_array.ndim == 4 and frame_array.shape[0] == 1:
|
||||
frame_array = frame_array[0]
|
||||
if np.issubdtype(frame_array.dtype, np.integer):
|
||||
frame_array = frame_array.astype(np.float32) / 255.0
|
||||
return frame_array
|
||||
try:
|
||||
from PIL import Image
|
||||
except ImportError:
|
||||
Image = None
|
||||
if Image is not None and isinstance(frame, Image.Image):
|
||||
return np.asarray(frame).astype(np.float32) / 255.0
|
||||
return frame
|
||||
|
||||
def _ensure_frame_list(video_array):
|
||||
if isinstance(video_array, list):
|
||||
if len(video_array) == 0:
|
||||
return video_array
|
||||
first_item = video_array[0]
|
||||
if isinstance(first_item, np.ndarray):
|
||||
if first_item.ndim == 5:
|
||||
return list(first_item[0])
|
||||
if first_item.ndim == 4:
|
||||
if len(video_array) == 1:
|
||||
return list(first_item)
|
||||
return list(first_item)
|
||||
if first_item.ndim == 3:
|
||||
return video_array
|
||||
return video_array
|
||||
if isinstance(video_array, np.ndarray):
|
||||
if video_array.ndim == 5:
|
||||
return list(video_array[0])
|
||||
if video_array.ndim == 4:
|
||||
return list(video_array)
|
||||
if video_array.ndim == 3:
|
||||
return [video_array]
|
||||
return video_array
|
||||
|
||||
# frames may be np.ndarray, torch.Tensor, or list of tensors/arrays/images
|
||||
# export_to_video expects a list of frames with values in [0, 1]
|
||||
if isinstance(frames, torch.Tensor):
|
||||
video_tensor = frames.detach().cpu()
|
||||
if video_tensor.dim() == 5:
|
||||
# [B, C, F, H, W] or [B, F, H, W, C]
|
||||
if video_tensor.shape[1] in (3, 4):
|
||||
video_tensor = video_tensor[0].permute(1, 2, 3, 0)
|
||||
else:
|
||||
video_tensor = video_tensor[0]
|
||||
elif video_tensor.dim() == 4 and video_tensor.shape[0] in (3, 4):
|
||||
video_tensor = video_tensor.permute(1, 2, 3, 0)
|
||||
# If float, assume [-1,1] and normalize to [0,1]
|
||||
if video_tensor.is_floating_point():
|
||||
video_tensor = video_tensor.clamp(-1, 1) * 0.5 + 0.5
|
||||
video_array = video_tensor.float().numpy()
|
||||
elif isinstance(frames, np.ndarray):
|
||||
video_array = frames
|
||||
if video_array.ndim == 5:
|
||||
video_array = video_array[0]
|
||||
if np.issubdtype(video_array.dtype, np.integer):
|
||||
video_array = video_array.astype(np.float32) / 255.0
|
||||
elif isinstance(frames, list):
|
||||
if len(frames) == 0:
|
||||
raise ValueError("No video frames found in output.")
|
||||
video_array = [_normalize_frame(frame) for frame in frames]
|
||||
else:
|
||||
video_array = frames
|
||||
if hasattr(video_array, "shape") and video_array.ndim == 5:
|
||||
video_array = video_array[0]
|
||||
|
||||
# Convert 4D array (frames, H, W, C) to list of frames for export_to_video
|
||||
if isinstance(video_array, np.ndarray) and video_array.ndim == 4:
|
||||
video_array = list(video_array)
|
||||
video_array = _ensure_frame_list(video_array)
|
||||
|
||||
export_to_video(video_array, str(output_path), fps=args.fps)
|
||||
use_ltx2_export = is_ltx2
|
||||
encode_video = None
|
||||
if use_ltx2_export:
|
||||
try:
|
||||
from diffusers.pipelines.ltx2.export_utils import encode_video
|
||||
except ImportError:
|
||||
encode_video = None
|
||||
|
||||
if use_ltx2_export and encode_video is not None:
|
||||
if isinstance(video_array, list):
|
||||
frames_np = np.stack(video_array, axis=0)
|
||||
elif isinstance(video_array, np.ndarray):
|
||||
frames_np = video_array
|
||||
else:
|
||||
frames_np = np.asarray(video_array)
|
||||
|
||||
if frames_np.ndim == 4 and frames_np.shape[-1] == 4:
|
||||
frames_np = frames_np[..., :3]
|
||||
|
||||
frames_np = np.clip(frames_np, 0.0, 1.0)
|
||||
frames_u8 = (frames_np * 255).round().clip(0, 255).astype("uint8")
|
||||
video_tensor = torch.from_numpy(frames_u8)
|
||||
|
||||
audio_out = None
|
||||
if audio is not None:
|
||||
if isinstance(audio, list):
|
||||
audio = audio[0] if audio else None
|
||||
if isinstance(audio, np.ndarray):
|
||||
audio = torch.from_numpy(audio)
|
||||
if isinstance(audio, torch.Tensor):
|
||||
audio_out = audio
|
||||
if audio_out.dim() > 1:
|
||||
audio_out = audio_out[0]
|
||||
audio_out = audio_out.float().cpu()
|
||||
|
||||
encode_video(
|
||||
video_tensor,
|
||||
fps=fps,
|
||||
audio=audio_out,
|
||||
audio_sample_rate=args.audio_sample_rate if audio_out is not None else None,
|
||||
output_path=str(output_path),
|
||||
)
|
||||
else:
|
||||
export_to_video(video_array, str(output_path), fps=fps)
|
||||
print(f"Saved generated video to {output_path}")
|
||||
|
||||
if profiler_enabled:
|
||||
|
||||
@@ -0,0 +1,32 @@
|
||||
# MammothModa2-Preview
|
||||
|
||||
## Run examples (MammothModa2-Preview)
|
||||
|
||||
Download model
|
||||
```bash
|
||||
hf bytedance-research/MammothModa2-Preview --local-dir ./MammothModa2-Preview
|
||||
```
|
||||
|
||||
### Text-to-Image (T2I)
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py \
|
||||
--model ./MammothModa2-Preview \
|
||||
--stage-config ./vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml \
|
||||
--prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \
|
||||
--height 1024 \
|
||||
--width 1024 \
|
||||
--num-inference-steps 50 \
|
||||
--text-guidance-scale 4.0 \
|
||||
--out output.png
|
||||
```
|
||||
|
||||
### Image Summary
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py \
|
||||
--model ./MammothModa2-Preview \
|
||||
--stage-config ./vllm_omni/model_executor/stage_configs/mammoth_moda2_ar.yaml \
|
||||
--question "Summarize this image." \
|
||||
--image ./image.png
|
||||
```
|
||||
@@ -0,0 +1,115 @@
|
||||
"""
|
||||
Offline inference example: MammothModa2 image summarization (single AR stage).
|
||||
|
||||
Example:
|
||||
uv run python examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py \
|
||||
--model path/to/MammothModa2-Preview \
|
||||
--stage-config vllm_omni/model_executor/stage_configs/mammoth_moda2_ar.yaml \
|
||||
--image /path/to/input.jpg \
|
||||
--question "Please summarize the content of this image."
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import os
|
||||
|
||||
from PIL import Image
|
||||
from vllm import SamplingParams
|
||||
from vllm.multimodal.image import convert_image_mode
|
||||
|
||||
from vllm_omni import Omni
|
||||
|
||||
DEFAULT_SYSTEM = "You are a helpful assistant."
|
||||
DEFAULT_QUESTION = "Please summarize the content of this image."
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="MammothModa2 image summarization (offline, AR only).")
|
||||
parser.add_argument("--model", type=str, required=True, help="Path to model directory or model id.")
|
||||
parser.add_argument(
|
||||
"--stage-config", type=str, required=True, help="Path to stage config yaml (single-stage AR->text)."
|
||||
)
|
||||
parser.add_argument("--image", type=str, required=True, help="Path to input image.")
|
||||
parser.add_argument("--question", type=str, default=DEFAULT_QUESTION, help="Question/instruction for the model.")
|
||||
parser.add_argument("--system", type=str, default=DEFAULT_SYSTEM, help="System prompt.")
|
||||
parser.add_argument(
|
||||
"--max-tokens",
|
||||
type=int,
|
||||
default=512,
|
||||
help="Max new tokens to generate.",
|
||||
)
|
||||
parser.add_argument("--temperature", type=float, default=0.2)
|
||||
parser.add_argument("--top-p", type=float, default=0.9)
|
||||
parser.add_argument("--seed", type=int, default=42)
|
||||
parser.add_argument("--trust-remote-code", action="store_true")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def build_prompt(system: str, question: str) -> str:
|
||||
return (
|
||||
f"<|im_start|>system\n{system}<|im_end|>\n"
|
||||
"<|im_start|>user\n"
|
||||
"<|vision_start|><|image_pad|><|vision_end|>"
|
||||
f"{question}<|im_end|>\n"
|
||||
"<|im_start|>assistant\n"
|
||||
)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
|
||||
if not os.path.exists(args.image):
|
||||
raise FileNotFoundError(f"Image file not found: {args.image}")
|
||||
|
||||
pil_image = Image.open(args.image)
|
||||
image_data = convert_image_mode(pil_image, "RGB")
|
||||
prompt = build_prompt(args.system, args.question)
|
||||
|
||||
omni = Omni(
|
||||
model=args.model,
|
||||
stage_configs_path=args.stage_config,
|
||||
trust_remote_code=args.trust_remote_code,
|
||||
)
|
||||
try:
|
||||
sp = SamplingParams(
|
||||
temperature=float(args.temperature),
|
||||
top_p=float(args.top_p),
|
||||
top_k=-1,
|
||||
max_tokens=int(args.max_tokens),
|
||||
seed=int(args.seed),
|
||||
detokenize=True,
|
||||
)
|
||||
# NOTE: omni.generate() returns a Generator[OmniRequestOutput, None, None].
|
||||
# Consume it inside the try block so the worker isn't closed early.
|
||||
outputs = list(
|
||||
omni.generate(
|
||||
[
|
||||
{
|
||||
"prompt": prompt,
|
||||
"multi_modal_data": {"image": image_data},
|
||||
"additional_information": {"omni_task": ["chat"]},
|
||||
}
|
||||
],
|
||||
[sp],
|
||||
)
|
||||
)
|
||||
finally:
|
||||
omni.close()
|
||||
|
||||
lines: list[str] = []
|
||||
for stage_outputs in outputs:
|
||||
req_outputs = getattr(stage_outputs, "request_output", stage_outputs)
|
||||
req_outputs = req_outputs if isinstance(req_outputs, list) else [req_outputs]
|
||||
for ro in req_outputs:
|
||||
text = ro.outputs[0].text if getattr(ro, "outputs", None) else str(ro)
|
||||
lines.append(f"request_id: {getattr(ro, 'request_id', 'unknown')}\n")
|
||||
lines.append("answer:\n")
|
||||
lines.append(text.strip() + "\n")
|
||||
lines.append("\n")
|
||||
|
||||
print("\n".join(lines))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,251 @@
|
||||
"""
|
||||
Offline inference example for MammothModa2 Text-to-Image (T2I) generation.
|
||||
This script uses the vllm_omni.Omni pipeline with a multi-stage configuration.
|
||||
|
||||
Workflow:
|
||||
1. Stage 0 (AR): Generates visual tokens and their corresponding hidden states.
|
||||
2. Stage 1 (DiT): Consumes the hidden states as conditions to perform diffusion
|
||||
and VAE decoding to produce the final image.
|
||||
|
||||
Example Usage:
|
||||
uv run python examples/offline_inference/run_mammothmoda2_t2i.py \
|
||||
--model path/to/MammothModa2-Preview \
|
||||
--stage-config vllm_omni/model_executor/stage_configs/mammoth_moda2.yaml \
|
||||
--prompt "A stylish woman riding a motorcycle in NYC, movie poster style" \
|
||||
--out output.png
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import NamedTuple
|
||||
|
||||
import torch
|
||||
from PIL import Image
|
||||
from vllm.sampling_params import SamplingParams
|
||||
|
||||
from vllm_omni import Omni
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Constants
|
||||
# ---------------------------------------------------------------------------
|
||||
_PATCH_SIZE = 16 # AR image grid patch size (pixels per token)
|
||||
|
||||
|
||||
class T2IGenConfig(NamedTuple):
|
||||
eol_token_id: int
|
||||
visual_token_start_id: int
|
||||
visual_token_end_id: int
|
||||
top_k: int # AR sampling top-k (covers the full visual generation vocabulary)
|
||||
# Qwen2.5-VL special vision tokens: <|image_pad|>, <|video_pad|>, <|vision_start|>, <|vision_end|>
|
||||
visual_ids: list[int]
|
||||
|
||||
|
||||
def load_t2i_generation_config(model_dir: str) -> T2IGenConfig:
|
||||
"""Load T2I token IDs from t2i_generation_config.json and config.json.
|
||||
|
||||
Supports both local directory paths and HuggingFace Hub model IDs.
|
||||
"""
|
||||
model_path = Path(model_dir)
|
||||
|
||||
def _read_json(filename: str) -> dict:
|
||||
local = model_path / filename
|
||||
if local.exists():
|
||||
with local.open(encoding="utf-8") as f:
|
||||
return json.load(f)
|
||||
# Fall back to HuggingFace Hub when model_dir is a repo ID.
|
||||
try:
|
||||
from huggingface_hub import hf_hub_download
|
||||
except ImportError as exc:
|
||||
raise ImportError(
|
||||
"huggingface_hub is required to load configs from HF Hub. Install it with: pip install huggingface_hub"
|
||||
) from exc
|
||||
cached = hf_hub_download(repo_id=model_dir, filename=filename)
|
||||
with open(cached, encoding="utf-8") as f:
|
||||
return json.load(f)
|
||||
|
||||
gen_cfg = _read_json("t2i_generation_config.json")
|
||||
llm_cfg = _read_json("config.json").get("llm_config", {})
|
||||
|
||||
return T2IGenConfig(
|
||||
eol_token_id=int(gen_cfg["eol_token_id"]),
|
||||
visual_token_start_id=int(gen_cfg["visual_token_start_id"]),
|
||||
visual_token_end_id=int(gen_cfg["visual_token_end_id"]),
|
||||
top_k=int(gen_cfg["top_k"]),
|
||||
visual_ids=[
|
||||
int(llm_cfg["image_token_id"]),
|
||||
int(llm_cfg["video_token_id"]),
|
||||
int(llm_cfg["vision_start_token_id"]),
|
||||
int(llm_cfg["vision_end_token_id"]),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
p = argparse.ArgumentParser(description="Run MammothModa2 T2I (AR -> DiT) with vLLM-Omni.")
|
||||
p.add_argument("--model", type=str, required=True, help="Path to the model directory.")
|
||||
p.add_argument("--stage-config", type=str, required=True, help="Path to the multi-stage YAML configuration.")
|
||||
p.add_argument(
|
||||
"--prompt",
|
||||
type=str,
|
||||
action="append",
|
||||
default=None,
|
||||
help=(
|
||||
"Text prompt for image generation. Can be provided multiple times "
|
||||
"to generate multiple images with shared height/width/CFG settings."
|
||||
),
|
||||
)
|
||||
p.add_argument("--height", type=int, default=1024, help="Output image height (must be a multiple of 16).")
|
||||
p.add_argument("--width", type=int, default=1024, help="Output image width (must be a multiple of 16).")
|
||||
p.add_argument("--num-inference-steps", type=int, default=50, help="Number of diffusion steps for the DiT stage.")
|
||||
p.add_argument(
|
||||
"--text-guidance-scale", type=float, default=9.0, help="Classifier-Free Guidance (CFG) scale for DiT."
|
||||
)
|
||||
p.add_argument(
|
||||
"--cfg-range",
|
||||
type=float,
|
||||
nargs=2,
|
||||
default=(0.0, 1.0),
|
||||
help="Relative step range [start, end] where CFG is active.",
|
||||
)
|
||||
p.add_argument("--out", type=str, default="output.png", help="Path to save the generated image.")
|
||||
p.add_argument("--trust-remote-code", action="store_true", help="Trust remote code when loading the model.")
|
||||
args = p.parse_args()
|
||||
if not args.prompt:
|
||||
args.prompt = ["A stylish woman with sunglasses riding a motorcycle in NYC."]
|
||||
return args
|
||||
|
||||
|
||||
def tensor_to_pil(image: torch.Tensor) -> Image.Image:
|
||||
"""Convert a normalized torch tensor [-1, 1] to a PIL Image."""
|
||||
if image.ndim == 4:
|
||||
image = image[0]
|
||||
image = image.detach().to("cpu")
|
||||
image = (image / 2 + 0.5).clamp(0, 1)
|
||||
image = (image * 255).to(torch.uint8)
|
||||
image = image.permute(1, 2, 0).contiguous().numpy()
|
||||
return Image.fromarray(image)
|
||||
|
||||
|
||||
def _format_prompt(user_prompt: str, ar_width: int, ar_height: int) -> str:
|
||||
"""Build the AR-stage prompt string including the image grid header."""
|
||||
return (
|
||||
"<|im_start|>system\nYou are a helpful image generator.<|im_end|>\n"
|
||||
f"<|im_start|>user\n{user_prompt}<|im_end|>\n"
|
||||
"<|im_start|>assistant\n"
|
||||
f"<|image start|>{ar_width}*{ar_height}<|image token|>"
|
||||
)
|
||||
|
||||
|
||||
def _collect_images(outputs: list) -> list[torch.Tensor]:
|
||||
"""Extract all image tensors produced by the final (DiT) stage."""
|
||||
images: list[torch.Tensor] = []
|
||||
for out in outputs:
|
||||
ro_list = getattr(out, "request_output", out)
|
||||
if not isinstance(ro_list, list):
|
||||
ro_list = [ro_list]
|
||||
for ro_item in ro_list:
|
||||
for completion in getattr(ro_item, "outputs", None) or []:
|
||||
mm = getattr(completion, "multimodal_output", None)
|
||||
if not isinstance(mm, dict) or "image" not in mm:
|
||||
raise RuntimeError(f"Missing image in multimodal output: {mm}")
|
||||
payload = mm["image"]
|
||||
for tensor in payload if isinstance(payload, list) else [payload]:
|
||||
if not isinstance(tensor, torch.Tensor):
|
||||
raise TypeError(f"Expected image tensor, got {type(tensor)}")
|
||||
images.append(tensor)
|
||||
return images
|
||||
|
||||
|
||||
def _save_images(images: list[torch.Tensor], out_path: str) -> list[str]:
|
||||
"""Save image tensors to disk.
|
||||
|
||||
Single image: written to *out_path* exactly.
|
||||
Multiple images: suffixed as ``<base>_0<ext>``, ``<base>_1<ext>``, …
|
||||
"""
|
||||
if not images:
|
||||
raise RuntimeError("No images to save.")
|
||||
base, ext = os.path.splitext(out_path)
|
||||
ext = ext or ".png"
|
||||
paths = []
|
||||
for i, tensor in enumerate(images):
|
||||
path = out_path if len(images) == 1 else f"{base}_{i}{ext}"
|
||||
tensor_to_pil(tensor).save(path)
|
||||
paths.append(path)
|
||||
return paths
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True)
|
||||
|
||||
if args.height <= 0 or args.width <= 0:
|
||||
raise ValueError(f"Height and width must be positive, got {args.height}x{args.width}")
|
||||
if args.height % _PATCH_SIZE != 0 or args.width % _PATCH_SIZE != 0:
|
||||
raise ValueError(f"Height and width must be multiples of {_PATCH_SIZE}, got {args.height}x{args.width}")
|
||||
|
||||
ar_height = args.height // _PATCH_SIZE
|
||||
ar_width = args.width // _PATCH_SIZE
|
||||
gen_cfg = load_t2i_generation_config(args.model)
|
||||
expected_grid_tokens = ar_height * (ar_width + 1)
|
||||
|
||||
logger.info("Initializing Omni pipeline...")
|
||||
omni = Omni(model=args.model, stage_configs_path=args.stage_config, trust_remote_code=args.trust_remote_code)
|
||||
try:
|
||||
ar_sampling = SamplingParams(
|
||||
temperature=1.0,
|
||||
top_p=1.0,
|
||||
top_k=gen_cfg.top_k,
|
||||
max_tokens=max(1, expected_grid_tokens + 1), # +1 for hidden state of eoi
|
||||
detokenize=False,
|
||||
)
|
||||
dit_sampling = SamplingParams(
|
||||
temperature=0.0,
|
||||
top_p=1.0,
|
||||
top_k=-1,
|
||||
max_tokens=1,
|
||||
detokenize=False,
|
||||
)
|
||||
|
||||
additional_information = {
|
||||
"omni_task": ["t2i"],
|
||||
"ar_width": [ar_width],
|
||||
"ar_height": [ar_height],
|
||||
"eol_token_id": [gen_cfg.eol_token_id],
|
||||
"visual_token_start_id": [gen_cfg.visual_token_start_id],
|
||||
"visual_token_end_id": [gen_cfg.visual_token_end_id],
|
||||
"image_height": [args.height],
|
||||
"image_width": [args.width],
|
||||
"num_inference_steps": [args.num_inference_steps],
|
||||
"text_guidance_scale": [args.text_guidance_scale],
|
||||
"cfg_range": [args.cfg_range[0], args.cfg_range[1]],
|
||||
"visual_ids": gen_cfg.visual_ids,
|
||||
}
|
||||
inputs = [
|
||||
{
|
||||
"prompt": _format_prompt(p, ar_width, ar_height),
|
||||
"additional_information": dict(additional_information),
|
||||
}
|
||||
for p in args.prompt
|
||||
]
|
||||
|
||||
logger.info("Starting generation...")
|
||||
# omni.generate() returns a Generator; consume it to run the full pipeline.
|
||||
outputs = list(omni.generate(inputs, [ar_sampling, dit_sampling]))
|
||||
|
||||
logger.info("Post-processing and saving image(s)...")
|
||||
for path in _save_images(_collect_images(outputs), args.out):
|
||||
logger.info(f"Saved: {path}")
|
||||
finally:
|
||||
omni.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -64,6 +64,51 @@ If media file paths are not provided, the script will use default assets. Suppor
|
||||
- `multi_audios`: Multiple audio inputs
|
||||
- `mixed_modalities`: Combination of video, image, and audio inputs
|
||||
|
||||
### Async-chunk (offline)
|
||||
|
||||
For true stage-level concurrency -- where downstream stages (Talker, Code2Wav)
|
||||
start **before** the upstream stage (Thinker) finishes -- use the async_chunk
|
||||
example. This requires:
|
||||
|
||||
1. A stage config YAML with ``async_chunk: true`` (e.g.
|
||||
``qwen3_omni_moe_async_chunk.yaml``).
|
||||
2. Hardware that matches the config (e.g. 2x H100 for the default 3-stage
|
||||
config).
|
||||
|
||||
The async_chunk example uses ``AsyncOmni`` instead of the synchronous ``Omni``
|
||||
class, which enables the async orchestrator to receive stage-0 intermediate
|
||||
outputs and trigger downstream stages early. Chunk data flows directly between
|
||||
stage workers via the in-worker ``OmniChunkTransferAdapter`` / connector,
|
||||
**not** through the orchestrator.
|
||||
|
||||
#### Single prompt
|
||||
```bash
|
||||
cd examples/offline_inference/qwen3_omni
|
||||
bash run_single_prompt_async_chunk.sh
|
||||
```
|
||||
|
||||
#### Multiple prompts with concurrency control
|
||||
```bash
|
||||
bash run_multiple_prompts_async_chunk.sh --max-in-flight 4
|
||||
```
|
||||
|
||||
#### Text-only output (skip audio generation)
|
||||
```bash
|
||||
python end2end_async_chunk.py --query-type text --modalities text
|
||||
```
|
||||
|
||||
#### Custom stage config
|
||||
```bash
|
||||
python end2end_async_chunk.py \
|
||||
--query-type use_audio \
|
||||
--stage-configs-path /path/to/your_async_chunk.yaml
|
||||
```
|
||||
|
||||
> **Note**: The synchronous ``end2end.py`` (using ``Omni``) is still the
|
||||
> recommended entry point for non-async-chunk workflows. Only use the
|
||||
> async_chunk example when you need the stage-level concurrency semantics
|
||||
> described in PR #962 / #1151.
|
||||
|
||||
### FAQ
|
||||
|
||||
If you encounter error about backend of librosa, try to install ffmpeg with command below.
|
||||
|
||||
@@ -0,0 +1,614 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Offline inference with async_chunk enabled via AsyncOmni.
|
||||
|
||||
This script uses AsyncOmni (the async orchestrator) to run offline inference
|
||||
with async_chunk semantics: downstream stages (Talker, Code2Wav) start
|
||||
*before* upstream stages finish, consuming chunks as they arrive via
|
||||
the in-worker OmniChunkTransferAdapter / connector.
|
||||
|
||||
Compared to the synchronous ``end2end.py`` (which uses ``Omni``), this
|
||||
entry point achieves true stage-level concurrency -- stage-1/2 are
|
||||
actively processing while stage-0 is still generating.
|
||||
|
||||
Usage
|
||||
-----
|
||||
python end2end_async_chunk.py --query-type use_audio \
|
||||
--stage-configs-path <path-to-async-chunk-yaml>
|
||||
|
||||
See ``--help`` for all options.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
import uuid
|
||||
from typing import NamedTuple
|
||||
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
import torch
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
import librosa
|
||||
from PIL import Image
|
||||
from vllm import SamplingParams
|
||||
from vllm.assets.audio import AudioAsset
|
||||
from vllm.assets.image import ImageAsset
|
||||
from vllm.assets.video import VideoAsset, video_to_ndarrays
|
||||
from vllm.multimodal.image import convert_image_mode
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
from vllm_omni.entrypoints.async_omni import AsyncOmni
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Query builders (reuse the patterns from end2end.py)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
default_system = (
|
||||
"You are Qwen, a virtual human developed by the Qwen Team, Alibaba "
|
||||
"Group, capable of perceiving auditory and visual inputs, as well as "
|
||||
"generating text and speech."
|
||||
)
|
||||
|
||||
|
||||
class QueryResult(NamedTuple):
|
||||
inputs: dict
|
||||
limit_mm_per_prompt: dict[str, int]
|
||||
|
||||
|
||||
def get_text_query(question: str = None) -> QueryResult:
|
||||
if question is None:
|
||||
question = "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words."
|
||||
prompt = (
|
||||
f"<|im_start|>system\n{default_system}<|im_end|>\n"
|
||||
"<|im_start|>user\n"
|
||||
f"{question}<|im_end|>\n"
|
||||
"<|im_start|>assistant\n"
|
||||
)
|
||||
return QueryResult(inputs={"prompt": prompt}, limit_mm_per_prompt={})
|
||||
|
||||
|
||||
def get_audio_query(
|
||||
question: str = None,
|
||||
audio_path: str | None = None,
|
||||
sampling_rate: int = 16000,
|
||||
) -> QueryResult:
|
||||
if question is None:
|
||||
question = "What is the content of this audio?"
|
||||
prompt = (
|
||||
f"<|im_start|>system\n{default_system}<|im_end|>\n"
|
||||
"<|im_start|>user\n<|audio_start|><|audio_pad|><|audio_end|>"
|
||||
f"{question}<|im_end|>\n"
|
||||
"<|im_start|>assistant\n"
|
||||
)
|
||||
if audio_path:
|
||||
if not os.path.exists(audio_path):
|
||||
raise FileNotFoundError(f"Audio file not found: {audio_path}")
|
||||
audio_signal, sr = librosa.load(audio_path, sr=sampling_rate)
|
||||
audio_data = (audio_signal.astype(np.float32), sr)
|
||||
else:
|
||||
audio_data = AudioAsset("mary_had_lamb").audio_and_sample_rate
|
||||
return QueryResult(
|
||||
inputs={
|
||||
"prompt": prompt,
|
||||
"multi_modal_data": {"audio": audio_data},
|
||||
},
|
||||
limit_mm_per_prompt={"audio": 1},
|
||||
)
|
||||
|
||||
|
||||
def get_image_query(question: str = None, image_path: str | None = None) -> QueryResult:
|
||||
if question is None:
|
||||
question = "What is the content of this image?"
|
||||
prompt = (
|
||||
f"<|im_start|>system\n{default_system}<|im_end|>\n"
|
||||
"<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>"
|
||||
f"{question}<|im_end|>\n"
|
||||
"<|im_start|>assistant\n"
|
||||
)
|
||||
if image_path:
|
||||
if not os.path.exists(image_path):
|
||||
raise FileNotFoundError(f"Image file not found: {image_path}")
|
||||
pil_image = Image.open(image_path)
|
||||
image_data = convert_image_mode(pil_image, "RGB")
|
||||
else:
|
||||
image_data = convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB")
|
||||
return QueryResult(
|
||||
inputs={
|
||||
"prompt": prompt,
|
||||
"multi_modal_data": {"image": image_data},
|
||||
},
|
||||
limit_mm_per_prompt={"image": 1},
|
||||
)
|
||||
|
||||
|
||||
def get_video_query(
|
||||
question: str = None,
|
||||
video_path: str | None = None,
|
||||
num_frames: int = 16,
|
||||
) -> QueryResult:
|
||||
if question is None:
|
||||
question = "Why is this video funny?"
|
||||
prompt = (
|
||||
f"<|im_start|>system\n{default_system}<|im_end|>\n"
|
||||
"<|im_start|>user\n<|vision_start|><|video_pad|><|vision_end|>"
|
||||
f"{question}<|im_end|>\n"
|
||||
"<|im_start|>assistant\n"
|
||||
)
|
||||
if video_path:
|
||||
if not os.path.exists(video_path):
|
||||
raise FileNotFoundError(f"Video file not found: {video_path}")
|
||||
video_frames = video_to_ndarrays(video_path, num_frames=num_frames)
|
||||
else:
|
||||
video_frames = VideoAsset(name="baby_reading", num_frames=num_frames).np_ndarrays
|
||||
return QueryResult(
|
||||
inputs={
|
||||
"prompt": prompt,
|
||||
"multi_modal_data": {"video": video_frames},
|
||||
},
|
||||
limit_mm_per_prompt={"video": 1},
|
||||
)
|
||||
|
||||
|
||||
query_map = {
|
||||
"text": get_text_query,
|
||||
"use_audio": get_audio_query,
|
||||
"use_image": get_image_query,
|
||||
"use_video": get_video_query,
|
||||
}
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Core async routine
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def clone_prompt_for_request(template: dict) -> dict:
|
||||
"""Shallow-clone prompt dict so concurrent requests own independent containers."""
|
||||
cloned = dict(template)
|
||||
for key in ("multi_modal_data", "mm_processor_kwargs", "additional_information"):
|
||||
value = template.get(key)
|
||||
if isinstance(value, dict):
|
||||
cloned[key] = dict(value)
|
||||
elif isinstance(value, list):
|
||||
cloned[key] = list(value)
|
||||
return cloned
|
||||
|
||||
|
||||
def _default_async_chunk_stage_configs_path() -> str | None:
|
||||
"""Best-effort default stage config for running Qwen3-Omni with async_chunk.
|
||||
|
||||
When this example is executed from within the repository, we resolve the
|
||||
default YAML path relative to this file. When installed elsewhere, the
|
||||
file may not exist and callers should pass --stage-configs-path explicitly.
|
||||
"""
|
||||
repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../.."))
|
||||
candidate = os.path.join(
|
||||
repo_root,
|
||||
"vllm_omni",
|
||||
"model_executor",
|
||||
"stage_configs",
|
||||
"qwen3_omni_moe_async_chunk.yaml",
|
||||
)
|
||||
return candidate if os.path.exists(candidate) else None
|
||||
|
||||
|
||||
async def run_single_request(
|
||||
async_omni: AsyncOmni,
|
||||
prompt: dict,
|
||||
request_id: str,
|
||||
sampling_params_list: list[SamplingParams] | None,
|
||||
output_dir: str,
|
||||
output_modalities: list[str] | None = None,
|
||||
stream_audio_to_disk: bool = False,
|
||||
) -> dict:
|
||||
"""Run one request through AsyncOmni and collect outputs.
|
||||
|
||||
Returns a dict with timing information and saved file paths.
|
||||
"""
|
||||
t_start = time.perf_counter()
|
||||
text_parts: list[str] = []
|
||||
audio_chunks: list[torch.Tensor] = []
|
||||
audio_sr: int | None = None
|
||||
first_audio_ts: float | None = None
|
||||
audio_list_consumed: int = 0
|
||||
audio_last_tensor: torch.Tensor | None = None
|
||||
stage_0_first_output_ts: float | None = None
|
||||
|
||||
samplerate = 24000
|
||||
wav_file = os.path.join(output_dir, f"output_{request_id}.wav")
|
||||
sf_writer: sf.SoundFile | None = None
|
||||
audio_samples_written: int = 0
|
||||
|
||||
try:
|
||||
async for omni_output in async_omni.generate(
|
||||
prompt=prompt,
|
||||
request_id=request_id,
|
||||
sampling_params_list=sampling_params_list,
|
||||
output_modalities=output_modalities,
|
||||
):
|
||||
if not isinstance(omni_output.request_output, list):
|
||||
outputs_list = [omni_output.request_output]
|
||||
else:
|
||||
outputs_list = omni_output.request_output
|
||||
|
||||
for output in outputs_list:
|
||||
if omni_output.final_output_type == "text":
|
||||
if stage_0_first_output_ts is None:
|
||||
stage_0_first_output_ts = time.perf_counter()
|
||||
text_output = output.outputs[0].text
|
||||
if output.finished:
|
||||
text_parts.append(text_output)
|
||||
elif omni_output.final_output_type == "audio":
|
||||
mm_out = output.outputs[0].multimodal_output
|
||||
if mm_out and "audio" in mm_out:
|
||||
if first_audio_ts is None:
|
||||
first_audio_ts = time.perf_counter()
|
||||
if audio_sr is None and "sr" in mm_out:
|
||||
sr_val = mm_out["sr"]
|
||||
audio_sr = sr_val.item() if hasattr(sr_val, "item") else int(sr_val)
|
||||
samplerate = audio_sr
|
||||
audio_data = mm_out["audio"]
|
||||
if isinstance(audio_data, list):
|
||||
new_chunks = audio_data[audio_list_consumed:]
|
||||
audio_list_consumed = len(audio_data)
|
||||
elif isinstance(audio_data, torch.Tensor):
|
||||
new_chunks = [audio_data]
|
||||
audio_last_tensor = audio_data
|
||||
else:
|
||||
new_chunks = []
|
||||
|
||||
if stream_audio_to_disk and new_chunks:
|
||||
if sf_writer is None:
|
||||
sf_writer = sf.SoundFile(
|
||||
wav_file,
|
||||
mode="w",
|
||||
samplerate=samplerate,
|
||||
channels=1,
|
||||
subtype="FLOAT",
|
||||
)
|
||||
for chunk in new_chunks:
|
||||
chunk_np = chunk.float().detach().cpu().numpy().flatten()
|
||||
sf_writer.write(chunk_np)
|
||||
audio_samples_written += len(chunk_np)
|
||||
else:
|
||||
audio_chunks.extend(new_chunks)
|
||||
finally:
|
||||
if sf_writer is not None:
|
||||
sf_writer.close()
|
||||
|
||||
t_end = time.perf_counter()
|
||||
result = {
|
||||
"request_id": request_id,
|
||||
"e2e_latency_s": t_end - t_start,
|
||||
"saved_files": [],
|
||||
}
|
||||
|
||||
# Save text output
|
||||
if text_parts:
|
||||
text_file = os.path.join(output_dir, f"{request_id}.txt")
|
||||
with open(text_file, "w", encoding="utf-8") as f:
|
||||
f.write("\n".join(text_parts))
|
||||
result["saved_files"].append(text_file)
|
||||
print(
|
||||
f"[Request {request_id}] Text saved to {text_file} "
|
||||
f"(stage-0 first output at {stage_0_first_output_ts - t_start:.3f}s)"
|
||||
)
|
||||
|
||||
# Save audio output
|
||||
if stream_audio_to_disk and audio_samples_written > 0:
|
||||
result["saved_files"].append(wav_file)
|
||||
result["audio_duration_s"] = audio_samples_written / samplerate
|
||||
result["num_audio_chunks"] = audio_list_consumed
|
||||
ttfa = (first_audio_ts - t_start) if first_audio_ts else None
|
||||
result["time_to_first_audio_s"] = ttfa
|
||||
ttfa_str = f"{ttfa:.3f}s" if ttfa is not None else "N/A"
|
||||
print(
|
||||
f"[Request {request_id}] Audio streamed to {wav_file} "
|
||||
f"(duration={result['audio_duration_s']:.2f}s, "
|
||||
f"TTFA={ttfa_str}, "
|
||||
f"e2e={result['e2e_latency_s']:.3f}s)"
|
||||
)
|
||||
elif audio_chunks or audio_last_tensor is not None:
|
||||
if audio_chunks:
|
||||
if len(audio_chunks) > 1:
|
||||
audio_tensor = torch.cat(audio_chunks, dim=-1)
|
||||
else:
|
||||
audio_tensor = audio_chunks[0]
|
||||
else:
|
||||
audio_tensor = audio_last_tensor
|
||||
audio_numpy = audio_tensor.float().detach().cpu().numpy()
|
||||
if audio_numpy.ndim > 1:
|
||||
audio_numpy = audio_numpy.flatten()
|
||||
sf.write(wav_file, audio_numpy, samplerate=samplerate, format="WAV")
|
||||
result["saved_files"].append(wav_file)
|
||||
result["audio_duration_s"] = len(audio_numpy) / samplerate
|
||||
result["num_audio_chunks"] = len(audio_chunks)
|
||||
ttfa = (first_audio_ts - t_start) if first_audio_ts else None
|
||||
result["time_to_first_audio_s"] = ttfa
|
||||
ttfa_str = f"{ttfa:.3f}s" if ttfa is not None else "N/A"
|
||||
print(
|
||||
f"[Request {request_id}] Audio saved to {wav_file} "
|
||||
f"({len(audio_chunks)} chunks, "
|
||||
f"duration={result['audio_duration_s']:.2f}s, "
|
||||
f"TTFA={ttfa_str}, "
|
||||
f"e2e={result['e2e_latency_s']:.3f}s)"
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
async def run_all(args):
|
||||
"""Main async entry: build prompts, create AsyncOmni, run requests."""
|
||||
# Build query
|
||||
query_func = query_map[args.query_type]
|
||||
if args.query_type == "use_video":
|
||||
query_result = query_func(
|
||||
video_path=getattr(args, "video_path", None),
|
||||
num_frames=getattr(args, "num_frames", 16),
|
||||
)
|
||||
elif args.query_type == "use_image":
|
||||
query_result = query_func(image_path=getattr(args, "image_path", None))
|
||||
elif args.query_type == "use_audio":
|
||||
query_result = query_func(
|
||||
audio_path=getattr(args, "audio_path", None),
|
||||
sampling_rate=getattr(args, "sampling_rate", 16000),
|
||||
)
|
||||
else:
|
||||
query_result = query_func()
|
||||
|
||||
# Build prompt list
|
||||
if args.txt_prompts is not None:
|
||||
assert args.query_type == "text", "txt-prompts is only supported for text query type"
|
||||
with open(args.txt_prompts, encoding="utf-8") as f:
|
||||
lines = [ln.strip() for ln in f if ln.strip()]
|
||||
prompts = [get_text_query(ln).inputs for ln in lines]
|
||||
print(f"[Info] Loaded {len(prompts)} prompts from {args.txt_prompts}")
|
||||
else:
|
||||
prompts = [clone_prompt_for_request(query_result.inputs) for _ in range(args.num_prompts)]
|
||||
|
||||
# Inject output modalities if specified
|
||||
output_modalities = None
|
||||
if args.modalities is not None:
|
||||
output_modalities = args.modalities.split(",")
|
||||
for prompt in prompts:
|
||||
prompt["modalities"] = output_modalities
|
||||
|
||||
# Create AsyncOmni
|
||||
print(f"[Info] Creating AsyncOmni with stage_configs_path={args.stage_configs_path}")
|
||||
async_omni = None
|
||||
try:
|
||||
async_omni = AsyncOmni(
|
||||
model=args.model,
|
||||
stage_configs_path=args.stage_configs_path,
|
||||
log_stats=args.log_stats,
|
||||
stage_init_timeout=args.stage_init_timeout,
|
||||
)
|
||||
|
||||
# Use default sampling params from stage config (they are pre-configured
|
||||
# in the YAML for each stage).
|
||||
sampling_params_list = None
|
||||
|
||||
output_dir = args.output_dir
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
# Run requests with concurrency control
|
||||
semaphore = asyncio.Semaphore(args.max_in_flight)
|
||||
request_timeout = getattr(args, "request_timeout_s", None)
|
||||
stream_audio = getattr(args, "stream_audio_to_disk", False)
|
||||
|
||||
async def _run_one(idx: int, prompt: dict):
|
||||
async with semaphore:
|
||||
request_id = f"req_{idx}_{uuid.uuid4().hex[:8]}"
|
||||
coro = run_single_request(
|
||||
async_omni=async_omni,
|
||||
prompt=prompt,
|
||||
request_id=request_id,
|
||||
sampling_params_list=sampling_params_list,
|
||||
output_dir=output_dir,
|
||||
output_modalities=output_modalities,
|
||||
stream_audio_to_disk=stream_audio,
|
||||
)
|
||||
if request_timeout and request_timeout > 0:
|
||||
return await asyncio.wait_for(coro, timeout=request_timeout)
|
||||
return await coro
|
||||
|
||||
wall_start = time.perf_counter()
|
||||
tasks = [_run_one(i, p) for i, p in enumerate(prompts)]
|
||||
all_results = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
wall_end = time.perf_counter()
|
||||
|
||||
# Print summary
|
||||
print("\n" + "=" * 60)
|
||||
print("Summary")
|
||||
print("=" * 60)
|
||||
success_count = 0
|
||||
total_audio_dur = 0.0
|
||||
for r in all_results:
|
||||
if isinstance(r, Exception):
|
||||
print(f" [ERROR] {type(r).__name__}: {r}")
|
||||
else:
|
||||
success_count += 1
|
||||
total_audio_dur += r.get("audio_duration_s", 0.0)
|
||||
print(f" [{r['request_id']}] e2e={r['e2e_latency_s']:.3f}s files={r['saved_files']}")
|
||||
wall_time = wall_end - wall_start
|
||||
print(f"\nTotal: {success_count}/{len(prompts)} succeeded")
|
||||
print(f"Wall time: {wall_time:.3f}s")
|
||||
if total_audio_dur > 0:
|
||||
print(f"Total audio duration: {total_audio_dur:.2f}s")
|
||||
print(f"Real-time factor: {total_audio_dur / wall_time:.2f}x")
|
||||
print("=" * 60)
|
||||
finally:
|
||||
if async_omni is not None:
|
||||
async_omni.shutdown()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CLI
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = FlexibleArgumentParser(
|
||||
description=(
|
||||
"Offline inference with async_chunk enabled via AsyncOmni. "
|
||||
"Downstream stages start before upstream stages finish, "
|
||||
"achieving true stage-level concurrency."
|
||||
)
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
default="Qwen/Qwen3-Omni-30B-A3B-Instruct",
|
||||
help="Model name or path.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--query-type",
|
||||
"-q",
|
||||
type=str,
|
||||
default="use_audio",
|
||||
choices=query_map.keys(),
|
||||
help="Query type.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stage-configs-path",
|
||||
type=str,
|
||||
default=_default_async_chunk_stage_configs_path(),
|
||||
help=(
|
||||
"Path to an async_chunk stage config YAML. "
|
||||
"If not set, uses the model's default config "
|
||||
"(make sure it has async_chunk: true)."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--log-stats",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Enable writing detailed statistics.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stage-init-timeout",
|
||||
type=int,
|
||||
default=300,
|
||||
help="Timeout for initializing a single stage (seconds).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-dir",
|
||||
type=str,
|
||||
default="output_audio_async_chunk",
|
||||
help="Directory to save output files.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-prompts",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of prompts to generate (duplicated from query).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--txt-prompts",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to a .txt file with one prompt per line.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-in-flight",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Maximum concurrent requests (default: 1).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--request-timeout-s",
|
||||
type=float,
|
||||
default=None,
|
||||
help=(
|
||||
"Per-request timeout in seconds. When set, a request that "
|
||||
"exceeds this duration is cancelled and reported as an error. "
|
||||
"Default: None (no timeout)."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--batch-timeout-s",
|
||||
type=float,
|
||||
default=None,
|
||||
help=(
|
||||
"Global timeout for the entire batch in seconds. When set, "
|
||||
"the whole run_all() is cancelled if it exceeds this duration. "
|
||||
"Default: None (no global timeout)."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stream-audio-to-disk",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help=(
|
||||
"Write audio chunks to WAV incrementally instead of "
|
||||
"accumulating in memory. Useful for very long audio or "
|
||||
"high --max-in-flight to reduce memory footprint."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--modalities",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Comma-separated output modalities filter (e.g. 'text', 'audio', 'text,audio').",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--audio-path",
|
||||
"-a",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to local audio file.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--image-path",
|
||||
"-i",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to local image file.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--video-path",
|
||||
"-v",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to local video file.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-frames",
|
||||
type=int,
|
||||
default=16,
|
||||
help="Number of frames to extract from video.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--sampling-rate",
|
||||
type=int,
|
||||
default=16000,
|
||||
help="Sampling rate for audio loading.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
|
||||
async def _main():
|
||||
batch_timeout = getattr(args, "batch_timeout_s", None)
|
||||
if batch_timeout and batch_timeout > 0:
|
||||
await asyncio.wait_for(run_all(args), timeout=batch_timeout)
|
||||
else:
|
||||
await run_all(args)
|
||||
|
||||
try:
|
||||
asyncio.run(_main())
|
||||
except asyncio.TimeoutError:
|
||||
print(
|
||||
f"\n[TIMEOUT] Batch exceeded --batch-timeout-s="
|
||||
f"{args.batch_timeout_s}s. AsyncOmni shutdown was handled "
|
||||
f"by finally block."
|
||||
)
|
||||
except KeyboardInterrupt:
|
||||
print("\nInterrupted by user. AsyncOmni shutdown was handled by finally block.")
|
||||
@@ -0,0 +1,23 @@
|
||||
#!/bin/bash
|
||||
# Run multiple Qwen3-Omni requests with async_chunk enabled.
|
||||
#
|
||||
# Uses AsyncOmni with --max-in-flight to control request-level
|
||||
# concurrency (each request still gets true stage-level concurrency
|
||||
# via async_chunk).
|
||||
#
|
||||
# Usage:
|
||||
# bash run_multiple_prompts_async_chunk.sh
|
||||
# bash run_multiple_prompts_async_chunk.sh --max-in-flight 4
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)"
|
||||
|
||||
python "${SCRIPT_DIR}/end2end_async_chunk.py" \
|
||||
--query-type text \
|
||||
--txt-prompts "${SCRIPT_DIR}/text_prompts_10.txt" \
|
||||
--stage-configs-path "${REPO_ROOT}/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml" \
|
||||
--output-dir output_audio_async_chunk \
|
||||
--max-in-flight 2 \
|
||||
"$@"
|
||||
@@ -0,0 +1,26 @@
|
||||
#!/bin/bash
|
||||
# Run a single Qwen3-Omni request with async_chunk enabled.
|
||||
#
|
||||
# This uses AsyncOmni (async orchestrator) so that downstream stages
|
||||
# (Talker, Code2Wav) start *before* stage-0 (Thinker) finishes,
|
||||
# achieving true stage-level concurrency via chunk-level streaming.
|
||||
#
|
||||
# Prerequisites:
|
||||
# - An async_chunk stage config YAML (e.g. qwen3_omni_moe_async_chunk.yaml)
|
||||
# - Hardware matching the config (e.g. 2x H100 for the default 3-stage config)
|
||||
#
|
||||
# Usage:
|
||||
# bash run_single_prompt_async_chunk.sh
|
||||
# bash run_single_prompt_async_chunk.sh --query-type text --modalities text
|
||||
# bash run_single_prompt_async_chunk.sh --stage-configs-path /path/to/custom.yaml
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)"
|
||||
|
||||
python "${SCRIPT_DIR}/end2end_async_chunk.py" \
|
||||
--query-type use_audio \
|
||||
--stage-configs-path "${REPO_ROOT}/vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml" \
|
||||
--output-dir output_audio_async_chunk \
|
||||
"$@"
|
||||
@@ -95,7 +95,9 @@ Add `--streaming` to stream audio chunks progressively via `AsyncOmni` (requires
|
||||
python end2end.py --query-type CustomVoice --streaming --output-dir /tmp/out_stream
|
||||
```
|
||||
|
||||
Each 25-frame Code2Wav chunk is logged as it arrives. The final WAV file is written once generation
|
||||
Each Code2Wav chunk is logged as it arrives (default 25 frames; configurable via `codec_chunk_frames`
|
||||
in the stage config). The initial chunk size is dynamically selected based on server load for reduced
|
||||
TTFA, and can be overridden per-request via the `initial_codec_chunk_frames` API field. The final WAV file is written once generation
|
||||
completes. This demonstrates that audio data is available progressively rather than only at the end.
|
||||
|
||||
> **Note:** Streaming uses `AsyncOmni` internally. The non-streaming path (`Omni`) is unchanged.
|
||||
|
||||
@@ -7,6 +7,7 @@ tasks, then runs Omni generation and saves output wav files.
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from typing import Any, NamedTuple
|
||||
|
||||
import soundfile as sf
|
||||
@@ -78,10 +79,16 @@ def get_custom_voice_query(use_batch_sample: bool = False) -> QueryResult:
|
||||
task_type = "CustomVoice"
|
||||
model_name = "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"
|
||||
if use_batch_sample:
|
||||
texts = ["其实我真的有发现,我是一个特别善于观察别人情绪的人。", "She said she would be here by noon."]
|
||||
instructs = ["", "Very happy."]
|
||||
languages = ["Chinese", "English"]
|
||||
speakers = ["Vivian", "Ryan"]
|
||||
texts = [
|
||||
"其实我真的有发现,我是一个特别善于观察别人情绪的人。",
|
||||
"She said she would be here by noon.",
|
||||
"I like you very much.",
|
||||
"Really, you do?",
|
||||
"Yes, absolutely.",
|
||||
]
|
||||
instructs = ["", "Very happy.", "Very happy.", "Very happy.", "Very happy."]
|
||||
languages = ["Chinese", "English", "English", "English", "English"]
|
||||
speakers = ["Vivian", "Ryan", "Ryan", "Ryan", "Ryan"]
|
||||
inputs = []
|
||||
for text, instruct, language, speaker in zip(texts, instructs, languages, speakers):
|
||||
additional_information = {
|
||||
@@ -337,13 +344,27 @@ async def main_streaming(args):
|
||||
|
||||
for i, prompt in enumerate(inputs):
|
||||
request_id = str(i)
|
||||
t_start = time.perf_counter()
|
||||
t_prev = t_start
|
||||
chunk_idx = 0
|
||||
async for stage_output in omni.generate(prompt, request_id=request_id):
|
||||
mm = stage_output.request_output.outputs[0].multimodal_output
|
||||
if not stage_output.finished:
|
||||
t_now = time.perf_counter()
|
||||
audio = mm.get("audio")
|
||||
n = len(audio) if isinstance(audio, list) else (0 if audio is None else 1)
|
||||
logger.info(f"Request {request_id}: received chunk {n}")
|
||||
dt_ms = (t_now - t_prev) * 1000
|
||||
ttfa_ms = (t_now - t_start) * 1000
|
||||
if chunk_idx == 0:
|
||||
logger.info(f"Request {request_id}: chunk {chunk_idx} samples={n} TTFA={ttfa_ms:.1f}ms")
|
||||
else:
|
||||
logger.info(f"Request {request_id}: chunk {chunk_idx} samples={n} inter_chunk={dt_ms:.1f}ms")
|
||||
t_prev = t_now
|
||||
chunk_idx += 1
|
||||
else:
|
||||
t_end = time.perf_counter()
|
||||
total_ms = (t_end - t_start) * 1000
|
||||
logger.info(f"Request {request_id}: done total={total_ms:.1f}ms chunks={chunk_idx}")
|
||||
_save_wav(output_dir, request_id, mm)
|
||||
|
||||
|
||||
|
||||
@@ -106,6 +106,23 @@ python text_to_image.py \
|
||||
--seed 42
|
||||
```
|
||||
|
||||
### Flux.2-dev Models
|
||||
To start Flux.2-dev with a single GPU, cpu-offload must be enabled because the total size of its weights exceeds the 80GB memory capacity of the GPU.
|
||||
```bash
|
||||
python examples/offline_inference/text_to_image/text_to_image.py \
|
||||
--model black-forest-labs/FLUX.2-dev \
|
||||
--prompt "a lovely bunny holding a sign that says 'vllm-omni'" \
|
||||
--seed 42 \
|
||||
--tensor-parallel-size 1 \
|
||||
--num-images-per-prompt 1 \
|
||||
--num-inference-steps 50 \
|
||||
--guidance-scale 4.0 \
|
||||
--height 1024 \
|
||||
--width 1024 \
|
||||
--enable-cpu-offload \
|
||||
--output flux2-dev.png
|
||||
```
|
||||
|
||||
### Key Arguments
|
||||
|
||||
**Common arguments:**
|
||||
|
||||
@@ -170,6 +170,11 @@ def parse_args() -> argparse.Namespace:
|
||||
default=1,
|
||||
help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enable-expert-parallel",
|
||||
action="store_true",
|
||||
help="Enable expert parallelism for MoE layers.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--lora-path",
|
||||
type=str,
|
||||
@@ -258,6 +263,7 @@ def main():
|
||||
cfg_parallel_size=args.cfg_parallel_size,
|
||||
tensor_parallel_size=args.tensor_parallel_size,
|
||||
vae_patch_parallel_size=args.vae_patch_parallel_size,
|
||||
enable_expert_parallel=args.enable_expert_parallel,
|
||||
)
|
||||
|
||||
# Check if profiling is requested via environment variable
|
||||
@@ -323,7 +329,7 @@ def main():
|
||||
print(
|
||||
f" Parallel configuration: tensor_parallel_size={args.tensor_parallel_size}, "
|
||||
f"ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}, "
|
||||
f"vae_patch_parallel_size={args.vae_patch_parallel_size}"
|
||||
f"vae_patch_parallel_size={args.vae_patch_parallel_size}, enable_expert_parallel={args.enable_expert_parallel}."
|
||||
)
|
||||
print(f" CPU offload: {args.enable_cpu_offload}")
|
||||
print(f" Image size: {args.width}x{args.height}")
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
# Text-To-Video
|
||||
|
||||
The `Wan-AI/Wan2.2-T2V-A14B-Diffusers` pipeline generates short videos from text prompts.
|
||||
The `Wan-AI/Wan2.2-T2V-A14B-Diffusers` pipeline generates short videos from text prompts. This script can also be used
|
||||
for `Lightricks/LTX-2` to generate video+audio.
|
||||
|
||||
## Local CLI Usage
|
||||
|
||||
@@ -19,6 +20,22 @@ python text_to_video.py \
|
||||
--output t2v_out.mp4
|
||||
```
|
||||
|
||||
LTX2 example:
|
||||
|
||||
```bash
|
||||
python text_to_video.py \
|
||||
--model "Lightricks/LTX-2" \
|
||||
--prompt "A cinematic close-up of ocean waves at golden hour." \
|
||||
--negative-prompt "worst quality, inconsistent motion, blurry, jittery, distorted" \
|
||||
--height 512 \
|
||||
--width 768 \
|
||||
--num-frames 121 \
|
||||
--num-inference-steps 40 \
|
||||
--guidance-scale 4.0 \
|
||||
--frame-rate 24 \
|
||||
--output ltx2_out.mp4
|
||||
```
|
||||
|
||||
Key arguments:
|
||||
|
||||
- `--prompt`: text description (string).
|
||||
@@ -32,6 +49,9 @@ Key arguments:
|
||||
- `--vae-use-slicing`: enable VAE slicing for memory optimization.
|
||||
- `--vae-use-tiling`: enable VAE tiling for memory optimization.
|
||||
- `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
|
||||
- `--tensor-parallel-size`: tensor parallel size (effective for models that support TP, e.g. LTX2).
|
||||
- `--enable-cpu-offload`: enable CPU offloading for diffusion models.
|
||||
- `--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).
|
||||
|
||||
> ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.
|
||||
@@ -31,6 +31,12 @@ def parse_args() -> argparse.Namespace:
|
||||
parser.add_argument("--height", type=int, default=720, help="Video height.")
|
||||
parser.add_argument("--width", type=int, default=1280, help="Video width.")
|
||||
parser.add_argument("--num-frames", type=int, default=81, help="Number of frames (Wan default is 81).")
|
||||
parser.add_argument(
|
||||
"--frame-rate",
|
||||
type=float,
|
||||
default=None,
|
||||
help="Optional generation frame rate (used by models like LTX2). Defaults to --fps.",
|
||||
)
|
||||
parser.add_argument("--num-inference-steps", type=int, default=40, help="Sampling steps.")
|
||||
parser.add_argument(
|
||||
"--boundary-ratio",
|
||||
@@ -109,12 +115,30 @@ def parse_args() -> argparse.Namespace:
|
||||
default=1,
|
||||
help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--audio-sample-rate",
|
||||
type=int,
|
||||
default=24000,
|
||||
help="Sample rate for audio output when saved (default: 24000 for LTX2).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--vae-patch-parallel-size",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of GPUs used for VAE patch/tile parallelism (decode).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enable-expert-parallel",
|
||||
action="store_true",
|
||||
help="Enable expert parallelism for MoE layers.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
|
||||
frame_rate = args.frame_rate if args.frame_rate is not None else float(args.fps)
|
||||
|
||||
# Wan2.2 cache-dit tuning (from cache-dit examples and cache_alignment).
|
||||
cache_config = None
|
||||
@@ -134,13 +158,14 @@ def main():
|
||||
"scm_steps_mask_policy": None, # SCM mask policy: None (disabled), "slow", "medium", "fast", "ultra"
|
||||
"scm_steps_policy": "dynamic", # SCM steps policy: "dynamic" or "static"
|
||||
}
|
||||
# Configure parallel settings (only SP is supported for Wan)
|
||||
# Note: cfg_parallel and tensor_parallel are not implemented for Wan models
|
||||
# Configure parallel settings
|
||||
parallel_config = DiffusionParallelConfig(
|
||||
ulysses_degree=args.ulysses_degree,
|
||||
ring_degree=args.ring_degree,
|
||||
cfg_parallel_size=args.cfg_parallel_size,
|
||||
tensor_parallel_size=args.tensor_parallel_size,
|
||||
vae_patch_parallel_size=args.vae_patch_parallel_size,
|
||||
enable_expert_parallel=args.enable_expert_parallel,
|
||||
)
|
||||
|
||||
# Check if profiling is requested via environment variable
|
||||
@@ -153,12 +178,12 @@ def main():
|
||||
vae_use_tiling=args.vae_use_tiling,
|
||||
boundary_ratio=args.boundary_ratio,
|
||||
flow_shift=args.flow_shift,
|
||||
cache_backend=args.cache_backend,
|
||||
cache_config=cache_config,
|
||||
enable_cache_dit_summary=args.enable_cache_dit_summary,
|
||||
enable_cpu_offload=args.enable_cpu_offload,
|
||||
parallel_config=parallel_config,
|
||||
enforce_eager=args.enforce_eager,
|
||||
cache_backend=args.cache_backend,
|
||||
cache_config=cache_config,
|
||||
)
|
||||
|
||||
if profiler_enabled:
|
||||
@@ -172,7 +197,9 @@ def main():
|
||||
print(f" Inference steps: {args.num_inference_steps}")
|
||||
print(f" Frames: {args.num_frames}")
|
||||
print(
|
||||
f" Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}, tensor_parallel_size={args.tensor_parallel_size}"
|
||||
f" Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree},"
|
||||
f" cfg_parallel_size={args.cfg_parallel_size}, tensor_parallel_size={args.tensor_parallel_size},"
|
||||
f" vae_patch_parallel_size={args.vae_patch_parallel_size}, enable_expert_parallel={args.enable_expert_parallel}"
|
||||
)
|
||||
print(f" Video size: {args.width}x{args.height}")
|
||||
print(f"{'=' * 60}\n")
|
||||
@@ -191,6 +218,7 @@ def main():
|
||||
guidance_scale_2=args.guidance_scale_high,
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
num_frames=args.num_frames,
|
||||
frame_rate=frame_rate,
|
||||
),
|
||||
)
|
||||
generation_end = time.perf_counter()
|
||||
@@ -199,31 +227,56 @@ def main():
|
||||
# Print profiling results
|
||||
print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)")
|
||||
|
||||
# Extract video frames from OmniRequestOutput
|
||||
if isinstance(frames, list) and len(frames) > 0:
|
||||
first_item = frames[0]
|
||||
audio = None
|
||||
if isinstance(frames, list):
|
||||
frames = frames[0] if frames else None
|
||||
|
||||
# Check if it's an OmniRequestOutput
|
||||
if hasattr(first_item, "final_output_type"):
|
||||
if first_item.final_output_type != "image":
|
||||
raise ValueError(
|
||||
f"Unexpected output type '{first_item.final_output_type}', expected 'image' for video generation."
|
||||
)
|
||||
|
||||
# Pipeline mode: extract from nested request_output
|
||||
if hasattr(first_item, "is_pipeline_output") and first_item.is_pipeline_output:
|
||||
if isinstance(first_item.request_output, list) and len(first_item.request_output) > 0:
|
||||
inner_output = first_item.request_output[0]
|
||||
if isinstance(inner_output, OmniRequestOutput) and hasattr(inner_output, "images"):
|
||||
frames = inner_output.images[0] if inner_output.images else None
|
||||
if frames is None:
|
||||
raise ValueError("No video frames found in output.")
|
||||
# Diffusion mode: use direct images field
|
||||
elif hasattr(first_item, "images") and first_item.images:
|
||||
frames = first_item.images
|
||||
if isinstance(frames, OmniRequestOutput):
|
||||
if frames.final_output_type != "image":
|
||||
raise ValueError(
|
||||
f"Unexpected output type '{frames.final_output_type}', expected 'image' for video generation."
|
||||
)
|
||||
if frames.multimodal_output and "audio" in frames.multimodal_output:
|
||||
audio = frames.multimodal_output["audio"]
|
||||
if frames.is_pipeline_output and frames.request_output is not None:
|
||||
inner_output = frames.request_output
|
||||
if isinstance(inner_output, list):
|
||||
inner_output = inner_output[0] if inner_output else None
|
||||
if isinstance(inner_output, OmniRequestOutput):
|
||||
if inner_output.multimodal_output and "audio" in inner_output.multimodal_output:
|
||||
audio = inner_output.multimodal_output["audio"]
|
||||
frames = inner_output
|
||||
if isinstance(frames, OmniRequestOutput):
|
||||
if frames.images:
|
||||
if len(frames.images) == 1 and isinstance(frames.images[0], tuple) and len(frames.images[0]) == 2:
|
||||
frames, audio = frames.images[0]
|
||||
elif len(frames.images) == 1 and isinstance(frames.images[0], dict):
|
||||
audio = frames.images[0].get("audio")
|
||||
frames = frames.images[0].get("frames") or frames.images[0].get("video")
|
||||
else:
|
||||
frames = frames.images
|
||||
else:
|
||||
raise ValueError("No video frames found in OmniRequestOutput.")
|
||||
|
||||
if isinstance(frames, list) and frames:
|
||||
first_item = frames[0]
|
||||
if isinstance(first_item, tuple) and len(first_item) == 2:
|
||||
frames, audio = first_item
|
||||
elif isinstance(first_item, dict):
|
||||
audio = first_item.get("audio")
|
||||
frames = first_item.get("frames") or first_item.get("video")
|
||||
elif isinstance(first_item, list):
|
||||
frames = first_item
|
||||
|
||||
if isinstance(frames, tuple) and len(frames) == 2:
|
||||
frames, audio = frames
|
||||
elif isinstance(frames, dict):
|
||||
audio = frames.get("audio")
|
||||
frames = frames.get("frames") or frames.get("video")
|
||||
|
||||
if frames is None:
|
||||
raise ValueError("No video frames found in output.")
|
||||
|
||||
output_path = Path(args.output)
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
try:
|
||||
@@ -231,32 +284,127 @@ def main():
|
||||
except ImportError:
|
||||
raise ImportError("diffusers is required for export_to_video.")
|
||||
|
||||
# frames may be np.ndarray (preferred) or torch.Tensor
|
||||
def _normalize_frame(frame):
|
||||
if isinstance(frame, torch.Tensor):
|
||||
frame_tensor = frame.detach().cpu()
|
||||
if frame_tensor.dim() == 4 and frame_tensor.shape[0] == 1:
|
||||
frame_tensor = frame_tensor[0]
|
||||
if frame_tensor.dim() == 3 and frame_tensor.shape[0] in (3, 4):
|
||||
frame_tensor = frame_tensor.permute(1, 2, 0)
|
||||
if frame_tensor.is_floating_point():
|
||||
frame_tensor = frame_tensor.clamp(-1, 1) * 0.5 + 0.5
|
||||
return frame_tensor.float().numpy()
|
||||
if isinstance(frame, np.ndarray):
|
||||
frame_array = frame
|
||||
if frame_array.ndim == 4 and frame_array.shape[0] == 1:
|
||||
frame_array = frame_array[0]
|
||||
if np.issubdtype(frame_array.dtype, np.integer):
|
||||
frame_array = frame_array.astype(np.float32) / 255.0
|
||||
return frame_array
|
||||
try:
|
||||
from PIL import Image
|
||||
except ImportError:
|
||||
Image = None
|
||||
if Image is not None and isinstance(frame, Image.Image):
|
||||
return np.asarray(frame).astype(np.float32) / 255.0
|
||||
return frame
|
||||
|
||||
def _ensure_frame_list(video_array):
|
||||
if isinstance(video_array, list):
|
||||
if len(video_array) == 0:
|
||||
return video_array
|
||||
first_item = video_array[0]
|
||||
if isinstance(first_item, np.ndarray):
|
||||
if first_item.ndim == 5:
|
||||
return list(first_item[0])
|
||||
if first_item.ndim == 4:
|
||||
if len(video_array) == 1:
|
||||
return list(first_item)
|
||||
return list(first_item)
|
||||
if first_item.ndim == 3:
|
||||
return video_array
|
||||
return video_array
|
||||
if isinstance(video_array, np.ndarray):
|
||||
if video_array.ndim == 5:
|
||||
return list(video_array[0])
|
||||
if video_array.ndim == 4:
|
||||
return list(video_array)
|
||||
if video_array.ndim == 3:
|
||||
return [video_array]
|
||||
return video_array
|
||||
|
||||
# frames may be np.ndarray, torch.Tensor, or list of tensors/arrays/images
|
||||
# export_to_video expects a list of frames with values in [0, 1]
|
||||
if isinstance(frames, torch.Tensor):
|
||||
video_tensor = frames.detach().cpu()
|
||||
if video_tensor.dim() == 5:
|
||||
# [B, C, F, H, W] or [B, F, H, W, C]
|
||||
if video_tensor.shape[1] in (3, 4):
|
||||
video_tensor = video_tensor[0].permute(1, 2, 3, 0)
|
||||
else:
|
||||
video_tensor = video_tensor[0]
|
||||
elif video_tensor.dim() == 4 and video_tensor.shape[0] in (3, 4):
|
||||
video_tensor = video_tensor.permute(1, 2, 3, 0)
|
||||
# If float, assume [-1,1] and normalize to [0,1]
|
||||
if video_tensor.is_floating_point():
|
||||
video_tensor = video_tensor.clamp(-1, 1) * 0.5 + 0.5
|
||||
video_array = video_tensor.float().numpy()
|
||||
elif isinstance(frames, np.ndarray):
|
||||
video_array = frames
|
||||
if video_array.ndim == 5:
|
||||
video_array = video_array[0]
|
||||
if np.issubdtype(video_array.dtype, np.integer):
|
||||
video_array = video_array.astype(np.float32) / 255.0
|
||||
elif isinstance(frames, list):
|
||||
if len(frames) == 0:
|
||||
raise ValueError("No video frames found in output.")
|
||||
video_array = [_normalize_frame(frame) for frame in frames]
|
||||
else:
|
||||
video_array = frames
|
||||
if hasattr(video_array, "shape") and video_array.ndim == 5:
|
||||
video_array = video_array[0]
|
||||
|
||||
# Convert 4D array (frames, H, W, C) to list of frames for export_to_video
|
||||
if isinstance(video_array, np.ndarray) and video_array.ndim == 4:
|
||||
video_array = list(video_array)
|
||||
video_array = _ensure_frame_list(video_array)
|
||||
|
||||
export_to_video(video_array, str(output_path), fps=args.fps)
|
||||
use_ltx2_export = False
|
||||
if args.model and "ltx" in str(args.model).lower():
|
||||
use_ltx2_export = True
|
||||
if audio is not None:
|
||||
use_ltx2_export = True
|
||||
|
||||
if use_ltx2_export:
|
||||
try:
|
||||
from diffusers.pipelines.ltx2.export_utils import encode_video
|
||||
except ImportError:
|
||||
raise ImportError("diffusers is required for LTX2 encode_video.")
|
||||
|
||||
if isinstance(video_array, list):
|
||||
frames_np = np.stack(video_array, axis=0)
|
||||
elif isinstance(video_array, np.ndarray):
|
||||
frames_np = video_array
|
||||
else:
|
||||
frames_np = np.asarray(video_array)
|
||||
|
||||
frames_u8 = (frames_np * 255).round().clip(0, 255).astype("uint8")
|
||||
video_tensor = torch.from_numpy(frames_u8)
|
||||
|
||||
audio_out = None
|
||||
if audio is not None:
|
||||
if isinstance(audio, list):
|
||||
audio = audio[0] if audio else None
|
||||
if isinstance(audio, np.ndarray):
|
||||
audio = torch.from_numpy(audio)
|
||||
if isinstance(audio, torch.Tensor):
|
||||
audio_out = audio
|
||||
if audio_out.dim() > 1:
|
||||
audio_out = audio_out[0]
|
||||
audio_out = audio_out.float().cpu()
|
||||
|
||||
encode_video(
|
||||
video_tensor,
|
||||
fps=args.fps,
|
||||
audio=audio_out,
|
||||
audio_sample_rate=args.audio_sample_rate if audio_out is not None else None,
|
||||
output_path=str(output_path),
|
||||
)
|
||||
else:
|
||||
export_to_video(video_array, str(output_path), fps=args.fps)
|
||||
print(f"Saved generated video to {output_path}")
|
||||
|
||||
if profiler_enabled:
|
||||
|
||||
@@ -0,0 +1,6 @@
|
||||
*.png
|
||||
.git/
|
||||
ct.yaml
|
||||
lintconf.yaml
|
||||
values.schema.json
|
||||
/workflows
|
||||
@@ -0,0 +1,15 @@
|
||||
apiVersion: v2
|
||||
name: chart-vllm-omni
|
||||
description: A Helm chart for deploying vLLM-Omni on Kubernetes for omni-modality model serving
|
||||
|
||||
# Application chart that can be packaged and deployed
|
||||
type: application
|
||||
|
||||
# Chart version - increment on each change to the chart
|
||||
version: 0.1.0
|
||||
|
||||
# vllm-omni application version — keep in sync with image.tag in values.yaml
|
||||
appVersion: "0.16.0"
|
||||
|
||||
maintainers:
|
||||
- name: vllm-omni-team
|
||||
@@ -0,0 +1,144 @@
|
||||
# vLLM-Omni Helm Chart
|
||||
|
||||
Helm chart for deploying [vLLM-Omni](https://github.com/vllm-project/vllm-omni) on Kubernetes. vLLM-Omni extends vLLM with omni-modality model serving, supporting text-to-image, multimodal chat, text-to-speech, and more.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
- Kubernetes 1.24+
|
||||
- Helm 3.x
|
||||
- NVIDIA GPU nodes with [NVIDIA Device Plugin](https://github.com/NVIDIA/k8s-device-plugin)
|
||||
|
||||
## Quick Start
|
||||
|
||||
```bash
|
||||
helm install my-release ./chart-helm \
|
||||
--set model=Tongyi-MAI/Z-Image-Turbo
|
||||
```
|
||||
|
||||
## Configuration
|
||||
|
||||
### Model Selection
|
||||
|
||||
Set the `model` value to any supported HuggingFace model ID:
|
||||
|
||||
| Model | Type | GPUs | Notes |
|
||||
|-------|------|------|-------|
|
||||
| `Tongyi-MAI/Z-Image-Turbo` | text-to-image | 1 | Small, fast (default) |
|
||||
| `stabilityai/stable-diffusion-3.5-medium` | text-to-image | 1 | ~6GB VRAM |
|
||||
| `Qwen/Qwen-Image` | text-to-image | 1 | Large, ~40GB+ VRAM |
|
||||
| `Qwen/Qwen2.5-Omni-7B` | multimodal | 2 | Text + audio + image + video |
|
||||
| `Qwen/Qwen3-Omni-7B-Chat` | multimodal | 2 | Latest omni model |
|
||||
| `Qwen/Qwen3-TTS` | text-to-speech | 1 | TTS |
|
||||
|
||||
### HuggingFace Token
|
||||
|
||||
For gated models that require authentication:
|
||||
|
||||
```bash
|
||||
helm install my-release ./chart-helm \
|
||||
--set model=Qwen/Qwen2.5-Omni-7B \
|
||||
--set hfToken=hf_xxxxx \
|
||||
--set resources.requests."nvidia\.com/gpu"=2 \
|
||||
--set resources.limits."nvidia\.com/gpu"=2
|
||||
```
|
||||
|
||||
### Omni-Specific Flags
|
||||
|
||||
Enable VAE memory optimizations for diffusion models:
|
||||
|
||||
```bash
|
||||
helm install my-release ./chart-helm \
|
||||
--set model=Qwen/Qwen-Image \
|
||||
--set omniArgs.vaeUseSlicing=true \
|
||||
--set omniArgs.vaeUseTiling=true
|
||||
```
|
||||
|
||||
Enable CPU offloading:
|
||||
|
||||
```bash
|
||||
helm install my-release ./chart-helm \
|
||||
--set model=Qwen/Qwen-Image \
|
||||
--set omniArgs.enableCpuOffload=true
|
||||
```
|
||||
|
||||
Pass additional raw CLI flags:
|
||||
|
||||
```bash
|
||||
helm install my-release ./chart-helm \
|
||||
--set model=Qwen/Qwen-Image \
|
||||
--set omniArgs.extraArgs[0]="--enable-layerwise-offload"
|
||||
```
|
||||
|
||||
### Model Cache
|
||||
|
||||
By default, a PersistentVolumeClaim is created for the HuggingFace model cache to avoid re-downloading models on pod restarts:
|
||||
|
||||
```yaml
|
||||
modelCache:
|
||||
enabled: true
|
||||
storageSize: "50Gi"
|
||||
storageClassName: ""
|
||||
```
|
||||
|
||||
To use an ephemeral volume instead:
|
||||
|
||||
```bash
|
||||
helm install my-release ./chart-helm \
|
||||
--set modelCache.enabled=false
|
||||
```
|
||||
|
||||
### Custom Command Override
|
||||
|
||||
To fully override the container command:
|
||||
|
||||
```bash
|
||||
helm install my-release ./chart-helm \
|
||||
--set image.command[0]=vllm \
|
||||
--set image.command[1]=serve \
|
||||
--set image.command[2]=my-model \
|
||||
--set image.command[3]=--omni \
|
||||
--set image.command[4]=--host \
|
||||
--set image.command[5]=0.0.0.0
|
||||
```
|
||||
|
||||
## API Endpoints
|
||||
|
||||
Once deployed, vLLM-Omni exposes the following OpenAI-compatible endpoints:
|
||||
|
||||
| Endpoint | Method | Description |
|
||||
|----------|--------|-------------|
|
||||
| `/health` | GET | Health check |
|
||||
| `/v1/models` | GET | List available models |
|
||||
| `/v1/chat/completions` | POST | Chat completions (text/multimodal) |
|
||||
| `/v1/images/generations` | POST | Image generation |
|
||||
| `/v1/images/edits` | POST | Image editing |
|
||||
| `/v1/audio/speech` | POST | Text-to-speech |
|
||||
|
||||
## Files
|
||||
|
||||
| File | Description |
|
||||
|------|-------------|
|
||||
| `Chart.yaml` | Chart metadata (name, version, maintainers) |
|
||||
| `values.yaml` | Default configuration values |
|
||||
| `values.schema.json` | JSON schema for validating values |
|
||||
| `templates/_helpers.tpl` | Helper templates for common configurations |
|
||||
| `templates/deployment.yaml` | Kubernetes Deployment |
|
||||
| `templates/service.yaml` | Kubernetes Service (ClusterIP) |
|
||||
| `templates/secrets.yaml` | Secrets (generic + HuggingFace token) |
|
||||
| `templates/pvc.yaml` | PersistentVolumeClaim for model cache |
|
||||
| `templates/configmap.yaml` | Optional ConfigMap |
|
||||
| `templates/hpa.yaml` | HorizontalPodAutoscaler |
|
||||
| `templates/poddisruptionbudget.yaml` | PodDisruptionBudget |
|
||||
| `templates/custom-objects.yaml` | Custom Kubernetes objects |
|
||||
|
||||
## Running Tests
|
||||
|
||||
This chart includes unit tests using [helm-unittest](https://github.com/helm-unittest/helm-unittest). Install the plugin and run tests:
|
||||
|
||||
```bash
|
||||
# Install plugin
|
||||
helm plugin install https://github.com/helm-unittest/helm-unittest
|
||||
|
||||
# Run tests
|
||||
helm unittest .
|
||||
```
|
||||
@@ -0,0 +1,3 @@
|
||||
chart-dirs:
|
||||
- charts
|
||||
validate-maintainers: false
|
||||
@@ -0,0 +1,42 @@
|
||||
---
|
||||
rules:
|
||||
braces:
|
||||
min-spaces-inside: 0
|
||||
max-spaces-inside: 0
|
||||
min-spaces-inside-empty: -1
|
||||
max-spaces-inside-empty: -1
|
||||
brackets:
|
||||
min-spaces-inside: 0
|
||||
max-spaces-inside: 0
|
||||
min-spaces-inside-empty: -1
|
||||
max-spaces-inside-empty: -1
|
||||
colons:
|
||||
max-spaces-before: 0
|
||||
max-spaces-after: 1
|
||||
commas:
|
||||
max-spaces-before: 0
|
||||
min-spaces-after: 1
|
||||
max-spaces-after: 1
|
||||
comments:
|
||||
require-starting-space: true
|
||||
min-spaces-from-content: 2
|
||||
document-end: disable
|
||||
document-start: disable
|
||||
empty-lines:
|
||||
max: 2
|
||||
max-start: 0
|
||||
max-end: 0
|
||||
hyphens:
|
||||
max-spaces-after: 1
|
||||
indentation:
|
||||
spaces: consistent
|
||||
indent-sequences: whatever
|
||||
check-multi-line-strings: false
|
||||
key-duplicates: enable
|
||||
line-length: disable
|
||||
new-line-at-end-of-file: disable
|
||||
new-lines:
|
||||
type: unix
|
||||
trailing-spaces: enable
|
||||
truthy:
|
||||
level: warning
|
||||
@@ -0,0 +1,203 @@
|
||||
{{/*
|
||||
Define the vllm-omni serve command from model + omniArgs.
|
||||
If image.command is set, uses that as a full override.
|
||||
*/}}
|
||||
{{- define "chart.omni-command" -}}
|
||||
{{- if .Values.image.command }}
|
||||
{{- toYaml .Values.image.command }}
|
||||
{{- else }}
|
||||
- "vllm"
|
||||
- "serve"
|
||||
- {{ .Values.model | quote }}
|
||||
- "--omni"
|
||||
- "--host"
|
||||
- "0.0.0.0"
|
||||
- "--port"
|
||||
- {{ include "chart.container-port" . | quote }}
|
||||
{{- if .Values.omniArgs.vaeUseSlicing }}
|
||||
- "--vae-use-slicing"
|
||||
{{- end }}
|
||||
{{- if .Values.omniArgs.vaeUseTiling }}
|
||||
- "--vae-use-tiling"
|
||||
{{- end }}
|
||||
{{- if .Values.omniArgs.enableCpuOffload }}
|
||||
- "--enable-cpu-offload"
|
||||
{{- end }}
|
||||
{{- if .Values.omniArgs.numGpus }}
|
||||
- "--num-gpus"
|
||||
- {{ .Values.omniArgs.numGpus | quote }}
|
||||
{{- end }}
|
||||
{{- if .Values.omniArgs.stageConfigsPath }}
|
||||
- "--stage-configs-path"
|
||||
- {{ .Values.omniArgs.stageConfigsPath | quote }}
|
||||
{{- end }}
|
||||
{{- if and .Values.omniArgs.cacheBackend (ne .Values.omniArgs.cacheBackend "none") }}
|
||||
- "--cache-backend"
|
||||
- {{ .Values.omniArgs.cacheBackend | quote }}
|
||||
{{- end }}
|
||||
{{- if .Values.omniArgs.defaultSamplingParams }}
|
||||
- "--default-sampling-params"
|
||||
- {{ .Values.omniArgs.defaultSamplingParams | quote }}
|
||||
{{- end }}
|
||||
{{- if and .Values.omniArgs.workerBackend (ne .Values.omniArgs.workerBackend "multi_process") }}
|
||||
- "--worker-backend"
|
||||
- {{ .Values.omniArgs.workerBackend | quote }}
|
||||
{{- end }}
|
||||
{{- range .Values.omniArgs.extraArgs }}
|
||||
- {{ . | quote }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
|
||||
{{/*
|
||||
Define HuggingFace environment variables
|
||||
*/}}
|
||||
{{- define "chart.hf-env" -}}
|
||||
- name: HF_HOME
|
||||
value: "/cache/huggingface"
|
||||
- name: HOME
|
||||
value: "/cache"
|
||||
{{- if .Values.hfToken }}
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: {{ .Release.Name }}-hf-token
|
||||
key: token
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
|
||||
{{/*
|
||||
Define ports for the pods
|
||||
*/}}
|
||||
{{- define "chart.container-port" -}}
|
||||
{{- default "8000" .Values.containerPort }}
|
||||
{{- end }}
|
||||
|
||||
{{/*
|
||||
Define service name
|
||||
*/}}
|
||||
{{- define "chart.service-name" -}}
|
||||
{{- if .Values.serviceName -}}
|
||||
{{ .Values.serviceName | lower | trim }}
|
||||
{{- else -}}
|
||||
{{ printf "%s-service" .Release.Name }}
|
||||
{{- end -}}
|
||||
{{- end }}
|
||||
|
||||
{{/*
|
||||
Define service port
|
||||
*/}}
|
||||
{{- define "chart.service-port" -}}
|
||||
{{- if .Values.servicePort }}
|
||||
{{- .Values.servicePort }}
|
||||
{{- else }}
|
||||
{{- include "chart.container-port" . }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
|
||||
{{/*
|
||||
Define service port name
|
||||
*/}}
|
||||
{{- define "chart.service-port-name" -}}
|
||||
"service-port"
|
||||
{{- end }}
|
||||
|
||||
{{/*
|
||||
Define container port name
|
||||
*/}}
|
||||
{{- define "chart.container-port-name" -}}
|
||||
"container-port"
|
||||
{{- end }}
|
||||
|
||||
{{/*
|
||||
Define deployment strategy
|
||||
*/}}
|
||||
{{- define "chart.strategy" -}}
|
||||
strategy:
|
||||
{{- if not .Values.deploymentStrategy }}
|
||||
type: Recreate
|
||||
{{- else }}
|
||||
{{ toYaml .Values.deploymentStrategy | indent 2 }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
|
||||
{{/*
|
||||
Define additional ports
|
||||
*/}}
|
||||
{{- define "chart.extraPorts" }}
|
||||
{{- with .Values.extraPorts }}
|
||||
{{ toYaml . }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
|
||||
{{/*
|
||||
Define chart external ConfigMaps and Secrets
|
||||
*/}}
|
||||
{{- define "chart.externalConfigs" -}}
|
||||
{{- with .Values.externalConfigs -}}
|
||||
{{ toYaml . }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
|
||||
{{/*
|
||||
Define startup, liveness and readiness probes
|
||||
*/}}
|
||||
{{- define "chart.probes" -}}
|
||||
{{- if .Values.startupProbe }}
|
||||
startupProbe:
|
||||
{{- with .Values.startupProbe }}
|
||||
{{- toYaml . | nindent 2 }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
{{- if .Values.readinessProbe }}
|
||||
readinessProbe:
|
||||
{{- with .Values.readinessProbe }}
|
||||
{{- toYaml . | nindent 2 }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
{{- if .Values.livenessProbe }}
|
||||
livenessProbe:
|
||||
{{- with .Values.livenessProbe }}
|
||||
{{- toYaml . | nindent 2 }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
|
||||
{{/*
|
||||
Define resources
|
||||
*/}}
|
||||
{{- define "chart.resources" -}}
|
||||
requests:
|
||||
memory: {{ required "Value 'resources.requests.memory' must be defined !" .Values.resources.requests.memory | quote }}
|
||||
cpu: {{ required "Value 'resources.requests.cpu' must be defined !" .Values.resources.requests.cpu | quote }}
|
||||
{{- if and (gt (int (index .Values.resources.requests "nvidia.com/gpu")) 0) (gt (int (index .Values.resources.limits "nvidia.com/gpu")) 0) }}
|
||||
nvidia.com/gpu: {{ required "Value 'resources.requests.nvidia.com/gpu' must be defined !" (index .Values.resources.requests "nvidia.com/gpu") }}
|
||||
{{- end }}
|
||||
limits:
|
||||
memory: {{ required "Value 'resources.limits.memory' must be defined !" .Values.resources.limits.memory | quote }}
|
||||
cpu: {{ required "Value 'resources.limits.cpu' must be defined !" .Values.resources.limits.cpu | quote }}
|
||||
{{- if and (gt (int (index .Values.resources.requests "nvidia.com/gpu")) 0) (gt (int (index .Values.resources.limits "nvidia.com/gpu")) 0) }}
|
||||
nvidia.com/gpu: {{ required "Value 'resources.limits.nvidia.com/gpu' must be defined !" (index .Values.resources.limits "nvidia.com/gpu") }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
|
||||
{{/*
|
||||
Define user for the main container
|
||||
*/}}
|
||||
{{- define "chart.user" }}
|
||||
{{- if .Values.image.runAsUser }}
|
||||
runAsUser:
|
||||
{{- with .Values.image.runAsUser }}
|
||||
{{- toYaml . | nindent 2 }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
|
||||
{{/*
|
||||
Define chart labels
|
||||
*/}}
|
||||
{{- define "chart.labels" -}}
|
||||
{{- with .Values.labels -}}
|
||||
{{ toYaml . }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
@@ -0,0 +1,11 @@
|
||||
{{- if .Values.configs -}}
|
||||
apiVersion: v1
|
||||
kind: ConfigMap
|
||||
metadata:
|
||||
name: "{{ .Release.Name }}-configs"
|
||||
namespace: {{ .Release.Namespace }}
|
||||
data:
|
||||
{{- with .Values.configs }}
|
||||
{{- toYaml . | nindent 2 }}
|
||||
{{- end }}
|
||||
{{- end -}}
|
||||
@@ -0,0 +1,6 @@
|
||||
{{- if .Values.customObjects }}
|
||||
{{- range .Values.customObjects }}
|
||||
{{- tpl (. | toYaml) $ }}
|
||||
---
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
@@ -0,0 +1,106 @@
|
||||
apiVersion: apps/v1
|
||||
kind: Deployment
|
||||
metadata:
|
||||
name: "{{ .Release.Name }}-deployment-vllm-omni"
|
||||
namespace: {{ .Release.Namespace }}
|
||||
labels:
|
||||
{{- include "chart.labels" . | nindent 4 }}
|
||||
spec:
|
||||
{{- if not .Values.autoscaling.enabled }}
|
||||
replicas: {{ .Values.replicaCount }}
|
||||
{{- end }}
|
||||
{{- include "chart.strategy" . | nindent 2 }}
|
||||
selector:
|
||||
matchLabels:
|
||||
{{- include "chart.labels" . | nindent 6 }}
|
||||
progressDeadlineSeconds: 1800
|
||||
template:
|
||||
metadata:
|
||||
labels:
|
||||
{{- include "chart.labels" . | nindent 8 }}
|
||||
spec:
|
||||
containers:
|
||||
- name: "vllm-omni"
|
||||
image: "{{ required "Required value 'image.repository' must be defined !" .Values.image.repository }}:{{ required "Required value 'image.tag' must be defined !" .Values.image.tag }}"
|
||||
command:
|
||||
{{- include "chart.omni-command" . | nindent 12 }}
|
||||
securityContext:
|
||||
{{- if .Values.image.securityContext }}
|
||||
{{- with .Values.image.securityContext }}
|
||||
{{- toYaml . | nindent 12 }}
|
||||
{{- end }}
|
||||
{{- else }}
|
||||
runAsNonRoot: false
|
||||
{{- include "chart.user" . | indent 12 }}
|
||||
{{- end }}
|
||||
imagePullPolicy: IfNotPresent
|
||||
env:
|
||||
{{- include "chart.hf-env" . | nindent 12 }}
|
||||
{{- if .Values.image.env }}
|
||||
{{- toYaml .Values.image.env | nindent 12 }}
|
||||
{{- end }}
|
||||
{{- if or .Values.externalConfigs .Values.configs .Values.secrets }}
|
||||
envFrom:
|
||||
{{- if .Values.configs }}
|
||||
- configMapRef:
|
||||
name: "{{ .Release.Name }}-configs"
|
||||
{{- end }}
|
||||
{{- if .Values.secrets }}
|
||||
- secretRef:
|
||||
name: "{{ .Release.Name }}-secrets"
|
||||
{{- end }}
|
||||
{{- include "chart.externalConfigs" . | nindent 12 }}
|
||||
{{- end }}
|
||||
ports:
|
||||
- name: {{ include "chart.container-port-name" . }}
|
||||
containerPort: {{ include "chart.container-port" . }}
|
||||
{{- include "chart.extraPorts" . | nindent 12 }}
|
||||
{{- include "chart.probes" . | indent 10 }}
|
||||
resources: {{- include "chart.resources" . | nindent 12 }}
|
||||
volumeMounts:
|
||||
- name: shm
|
||||
mountPath: /dev/shm
|
||||
- name: model-cache
|
||||
mountPath: /cache
|
||||
|
||||
{{- with .Values.extraContainers }}
|
||||
{{ toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
|
||||
{{- if .Values.extraInit.initContainers }}
|
||||
initContainers:
|
||||
{{- toYaml .Values.extraInit.initContainers | nindent 8 }}
|
||||
{{- end }}
|
||||
|
||||
volumes:
|
||||
- name: shm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
sizeLimit: {{ .Values.shmSize }}
|
||||
- name: model-cache
|
||||
{{- if .Values.modelCache.enabled }}
|
||||
persistentVolumeClaim:
|
||||
claimName: {{ .Release.Name }}-model-cache
|
||||
{{- else }}
|
||||
emptyDir: {}
|
||||
{{- end }}
|
||||
|
||||
{{- with .Values.nodeSelector }}
|
||||
nodeSelector:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- with .Values.tolerations }}
|
||||
tolerations:
|
||||
{{- toYaml . | nindent 8 }}
|
||||
{{- end }}
|
||||
{{- if .Values.gpuModels }}
|
||||
affinity:
|
||||
nodeAffinity:
|
||||
requiredDuringSchedulingIgnoredDuringExecution:
|
||||
nodeSelectorTerms:
|
||||
- matchExpressions:
|
||||
- key: nvidia.com/gpu.product
|
||||
operator: In
|
||||
values:
|
||||
{{- toYaml .Values.gpuModels | nindent 20 }}
|
||||
{{- end }}
|
||||
@@ -0,0 +1,31 @@
|
||||
{{- if .Values.autoscaling.enabled }}
|
||||
apiVersion: autoscaling/v2
|
||||
kind: HorizontalPodAutoscaler
|
||||
metadata:
|
||||
name: "{{ .Release.Name }}-hpa"
|
||||
namespace: {{ .Release.Namespace }}
|
||||
spec:
|
||||
scaleTargetRef:
|
||||
apiVersion: apps/v1
|
||||
kind: Deployment
|
||||
name: "{{ .Release.Name }}-deployment-vllm-omni"
|
||||
minReplicas: {{ .Values.autoscaling.minReplicas }}
|
||||
maxReplicas: {{ .Values.autoscaling.maxReplicas }}
|
||||
metrics:
|
||||
{{- if .Values.autoscaling.targetCPUUtilizationPercentage }}
|
||||
- type: Resource
|
||||
resource:
|
||||
name: cpu
|
||||
target:
|
||||
type: Utilization
|
||||
averageUtilization: {{ .Values.autoscaling.targetCPUUtilizationPercentage }}
|
||||
{{- end }}
|
||||
{{- if .Values.autoscaling.targetMemoryUtilizationPercentage }}
|
||||
- type: Resource
|
||||
resource:
|
||||
name: memory
|
||||
target:
|
||||
type: Utilization
|
||||
averageUtilization: {{ .Values.autoscaling.targetMemoryUtilizationPercentage }}
|
||||
{{- end }}
|
||||
{{- end }}
|
||||
@@ -0,0 +1,32 @@
|
||||
{{- if .Values.ingress.enabled }}
|
||||
apiVersion: networking.k8s.io/v1
|
||||
kind: Ingress
|
||||
metadata:
|
||||
name: "{{ .Release.Name }}-ingress"
|
||||
namespace: {{ .Release.Namespace }}
|
||||
labels:
|
||||
{{- include "chart.labels" . | nindent 4 }}
|
||||
{{- with .Values.ingress.annotations }}
|
||||
annotations:
|
||||
{{- toYaml . | nindent 4 }}
|
||||
{{- end }}
|
||||
spec:
|
||||
{{- if .Values.ingress.ingressClassName }}
|
||||
ingressClassName: {{ .Values.ingress.ingressClassName }}
|
||||
{{- end }}
|
||||
{{- if .Values.ingress.tls }}
|
||||
tls:
|
||||
{{- toYaml .Values.ingress.tls | nindent 4 }}
|
||||
{{- end }}
|
||||
rules:
|
||||
- host: {{ .Values.ingress.host }}
|
||||
http:
|
||||
paths:
|
||||
- path: /
|
||||
pathType: Prefix
|
||||
backend:
|
||||
service:
|
||||
name: "{{ include "chart.service-name" . }}"
|
||||
port:
|
||||
number: {{ include "chart.service-port" . }}
|
||||
{{- end }}
|
||||
@@ -0,0 +1,12 @@
|
||||
{{- if gt (int .Values.replicaCount) 1 }}
|
||||
apiVersion: policy/v1
|
||||
kind: PodDisruptionBudget
|
||||
metadata:
|
||||
name: "{{ .Release.Name }}-pdb"
|
||||
namespace: {{ .Release.Namespace }}
|
||||
spec:
|
||||
maxUnavailable: {{ default 1 .Values.maxUnavailablePodDisruptionBudget }}
|
||||
selector:
|
||||
matchLabels:
|
||||
{{- include "chart.labels" . | nindent 6 }}
|
||||
{{- end }}
|
||||