Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
04d7fdbdb2 | ||
|
|
e3ba91a29e | ||
|
|
86bc94c07e | ||
|
|
62c67efcc5 | ||
|
|
ae1d0098ce | ||
|
|
c27111d345 | ||
|
|
cf081eb149 | ||
|
|
86e3e4e64d | ||
|
|
1f953c114c | ||
|
|
de3a9918e8 | ||
|
|
a6171ff44b | ||
|
|
e5e6562fed | ||
|
|
af012babf8 | ||
|
|
4de0c8d41b | ||
|
|
365dfd5f8e | ||
|
|
b39228d32e | ||
|
|
5377619c0b | ||
|
|
dc72aede40 | ||
|
|
c107f0b336 | ||
|
|
e17130090c | ||
|
|
cebf78ba3e | ||
|
|
3e8cb6c631 | ||
|
|
92f771d455 | ||
|
|
e0b71056ba | ||
|
|
26b3ddee47 | ||
|
|
53a3df7309 | ||
|
|
75056b2db2 | ||
|
|
717dbee5a6 | ||
|
|
c8cf8a799d | ||
|
|
0540af327e | ||
|
|
54a999d031 | ||
|
|
366b3368ea | ||
|
|
36c6cb9634 | ||
|
|
9c6fbd76be | ||
|
|
8ad8988e2b | ||
|
|
b3cced310a | ||
|
|
21e621f84c | ||
|
|
c8c60e1970 | ||
|
|
28dd1a699d | ||
|
|
ff3365f6a2 | ||
|
|
fe7e8770db | ||
|
|
a52867abde | ||
|
|
7c5159037f | ||
|
|
1c7f963a2d | ||
|
|
4dbaa74f2b | ||
|
|
11d9ca8f41 | ||
|
|
3306fa6611 | ||
|
|
0e005817f7 | ||
|
|
c501c2d512 | ||
|
|
5d28b6a40c | ||
|
|
09b66bf6db | ||
|
|
ac17431745 | ||
|
|
e1efcb7560 | ||
|
|
da4a077155 | ||
|
|
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 |
No files matched your search
@@ -0,0 +1,56 @@
|
||||
#!/bin/bash
|
||||
|
||||
# This script build the XPU docker image and run the offline inference inside the container.
|
||||
set -ex
|
||||
|
||||
omni_source_dir=$(git rev-parse --show-toplevel)
|
||||
|
||||
base_image_name="xpu/vllm-omni-ci-base:${VLLM_VERSION:?VLLM_VERSION must be set}"
|
||||
image_name="xpu/vllm-omni-ci:${BUILDKITE_COMMIT:?BUILDKITE_COMMIT must be set}"
|
||||
container_name="xpu_${BUILDKITE_COMMIT}_$(
|
||||
tr -dc A-Za-z0-9 </dev/urandom | head -c 10
|
||||
echo
|
||||
)"
|
||||
|
||||
cd "${omni_source_dir}"
|
||||
if [ -z "$(docker images -q "${base_image_name}")" ]; then
|
||||
docker build --target vllm-base -t "${base_image_name}" --build-arg "VLLM_VERSION=${VLLM_VERSION}" -f docker/Dockerfile.xpu .
|
||||
fi
|
||||
|
||||
# Try building the docker image
|
||||
docker build --build-arg "VLLM_BASE=${base_image_name}" --build-arg "VLLM_VERSION=${VLLM_VERSION}" -t "${image_name}" -f docker/Dockerfile.xpu .
|
||||
|
||||
# Setup cleanup
|
||||
remove_docker_container() {
|
||||
docker rm -f "${container_name}" || true
|
||||
docker image rm -f "${image_name}" || true
|
||||
docker system prune -f || true
|
||||
}
|
||||
trap remove_docker_container EXIT
|
||||
|
||||
HF_CACHE="${HF_CACHE:-$(realpath ~)/.cache/huggingface}"
|
||||
mkdir -p "${HF_CACHE}"
|
||||
HF_MOUNT="/root/.cache/huggingface"
|
||||
|
||||
time timeout -k 30 30m docker run \
|
||||
--device /dev/dri:/dev/dri \
|
||||
--net=host \
|
||||
--ipc=host \
|
||||
-v /dev/dri/by-path:/dev/dri/by-path \
|
||||
-v "${HF_CACHE}:${HF_MOUNT}" \
|
||||
--security-opt seccomp=unconfined \
|
||||
--entrypoint="" \
|
||||
-e VLLM_LOGGING_LEVEL \
|
||||
-e VLLM_OMNI_LOGGING_LEVEL \
|
||||
-e HF_TOKEN \
|
||||
-e ZE_AFFINITY_MASK \
|
||||
--name "${container_name}" \
|
||||
"${image_name}" \
|
||||
bash -c '
|
||||
set -e
|
||||
echo $ZE_AFFINITY_MASK
|
||||
pip install tblib==3.1.0
|
||||
cd /workspace/vllm-omni
|
||||
pytest -v -s -m "core_model and xpu and B60"
|
||||
pytest -v -s -m "advanced_model and xpu and B60"
|
||||
'
|
||||
@@ -54,7 +54,6 @@ 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:
|
||||
@@ -62,38 +61,13 @@ steps:
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: ":full_moon: Qwen3-TTS Non-Async-Chunk E2E Test"
|
||||
timeout_in_minutes: 30
|
||||
- label: ":full_moon: Diffusion Model Test with H100"
|
||||
timeout_in_minutes: 60
|
||||
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"
|
||||
# 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"
|
||||
- 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:
|
||||
@@ -123,17 +97,82 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":email: Nightly Perf Collection & Email"
|
||||
- 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 & Test Case Statistics"
|
||||
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"
|
||||
- python tools/nightly/buildkite_testcase_statistics.py -o tests/buildkite_testcase_statistics.html
|
||||
- buildkite-agent artifact upload "tests/*.html"
|
||||
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 Collection & Email"
|
||||
key: nightly-perf-distribution
|
||||
depends_on: nightly-performance
|
||||
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
|
||||
- buildkite-agent artifact download "tests/*.html" . --step nightly-performance
|
||||
- python tools/nightly/generate_nightly_perf_excel.py
|
||||
- python tools/nightly/send_nightly_perf_email.py
|
||||
- python tools/nightly/send_nightly_email.py --report-file "tests/*.xlsx, tests/*.html"
|
||||
- buildkite-agent artifact upload "tests/*.xlsx"
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
@@ -34,15 +34,7 @@ steps:
|
||||
- label: "Diffusion Model CPU offloading Test"
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
set +e
|
||||
pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
|
||||
EXIT1=$$?
|
||||
pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
|
||||
EXIT2=$$?
|
||||
exit $$((EXIT1 | EXIT2))
|
||||
'
|
||||
- timeout 10m pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
|
||||
agents:
|
||||
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
|
||||
plugins:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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://drive.google.com/file/d/12k_DOo2xoEROvd4xQxYNuEi6C2GECoRx/view?usp=drive_link"><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?slide=id.p1#slide=id.p1"><b>Slides</b></a> |
|
||||
</p>
|
||||
|
||||
|
||||
@@ -17,6 +17,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
|
||||
|
||||
*Latest News* 🔥
|
||||
|
||||
- [2026/03] Joined vLLM HK First Meetup! Video on [youtube](https://youtu.be/sgwNfsNnR9I)
|
||||
- [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.
|
||||
|
||||
@@ -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,8 @@ 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 asyncio
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
import aiohttp
|
||||
@@ -18,6 +20,7 @@ from .format import (
|
||||
base64_to_audio,
|
||||
base64_to_image_tensor,
|
||||
bytes_to_audio,
|
||||
bytes_to_video,
|
||||
image_tensor_to_base64,
|
||||
image_tensor_to_png_bytes,
|
||||
video_to_base64,
|
||||
@@ -29,10 +32,43 @@ from .types import AudioFormat
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
async def url_json(session: aiohttp.ClientSession, url: str, verb: str = "get", **kwargs) -> dict[str, Any]:
|
||||
try:
|
||||
async with getattr(session, verb)(url, **kwargs) 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:
|
||||
return 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}")
|
||||
|
||||
|
||||
async def url_bytes(session: aiohttp.ClientSession, url: str, verb: str = "get", **kwargs) -> bytes:
|
||||
try:
|
||||
async with getattr(session, verb)(url, **kwargs) 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}"
|
||||
)
|
||||
return await response.read()
|
||||
except aiohttp.ClientError as e:
|
||||
raise RuntimeError(f"Network error connecting to vLLM-Omni at {url}: {e}")
|
||||
|
||||
|
||||
class VLLMOmniClient:
|
||||
def __init__(self, base_url: str, timeout: float = 300.0):
|
||||
def __init__(
|
||||
self, base_url: str, timeout: float | None = None, poll_interval: float = 5.0, max_poll_duration: float = 60 * 5
|
||||
):
|
||||
self.base_url = base_url
|
||||
self.timeout = aiohttp.ClientTimeout(total=timeout)
|
||||
self.poll_interval = poll_interval
|
||||
self.max_poll_duration = max_poll_duration
|
||||
|
||||
async def generate_image(
|
||||
self,
|
||||
@@ -43,12 +79,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 +95,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 +149,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 +171,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 +225,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 +236,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 +251,87 @@ 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",
|
||||
)
|
||||
|
||||
async with aiohttp.ClientSession(timeout=self.timeout) as session:
|
||||
# Start the video generation job
|
||||
url = f"{self.base_url}/videos"
|
||||
data = await url_json(session, url, "post", data=form)
|
||||
if (job_id := data.get("id", None)) is None:
|
||||
raise RuntimeError("API response missing job 'id' field - expected OpenAI compliant format")
|
||||
if (job_status := data.get("status", None)) is None:
|
||||
raise RuntimeError("API response missing job 'status' field - expected OpenAI compliant format")
|
||||
|
||||
# Poll for video generation job completion
|
||||
deadline = asyncio.get_running_loop().time() + self.max_poll_duration
|
||||
url = f"{self.base_url}/videos/{job_id}"
|
||||
while job_status not in {"completed", "failed"}:
|
||||
await asyncio.sleep(self.poll_interval)
|
||||
|
||||
data = await url_json(session, url)
|
||||
if (job_status := data.get("status", None)) is None:
|
||||
raise RuntimeError("API response missing job 'status' field - expected OpenAI compliant format")
|
||||
if asyncio.get_running_loop().time() >= deadline:
|
||||
raise RuntimeError(f"Timed out waiting for video job {job_id} to complete")
|
||||
|
||||
if job_status == "failed":
|
||||
raise RuntimeError(f"Video job failed: {data}")
|
||||
|
||||
# Retrieve completed content
|
||||
video_bytes = await url_bytes(session, f"{url}/content")
|
||||
|
||||
# Decode video and make a best effort at cleaning up server resources
|
||||
try:
|
||||
return bytes_to_video(video_bytes)
|
||||
finally:
|
||||
try:
|
||||
await url_json(session, url, "delete")
|
||||
except Exception as exc:
|
||||
logger.warning("Failed to clean up video job %s: %s", job_id, exc)
|
||||
|
||||
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,75 @@ def video_to_base64(video: VideoInput, filename: str = "video.mp4") -> str:
|
||||
return f"data:{mime_type};base64,{base64_str}"
|
||||
|
||||
|
||||
def bytes_to_video(video_bytes: bytes) -> VideoInput:
|
||||
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 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}")
|
||||
|
||||
return bytes_to_video(video_bytes)
|
||||
|
||||
|
||||
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
|
||||
@@ -70,13 +70,17 @@ def add_sampling_parameters_to_stage(
|
||||
|
||||
stages = pipeline_spec["stages"]
|
||||
if isinstance(sampling_param_list, dict):
|
||||
sampling_param_list = sampling_param_list.__class__(sampling_param_list)
|
||||
sampling_param_list.update(params_to_add)
|
||||
elif sampling_param_list is None:
|
||||
sampling_param_list = params_to_add.copy()
|
||||
else:
|
||||
for i, stage in enumerate(stages):
|
||||
if stage == stage_type:
|
||||
sampling_param_list[i].update(params_to_add)
|
||||
stage_param = sampling_param_list[i]
|
||||
stage_param = stage_param.__class__(stage_param)
|
||||
stage_param.update(params_to_add)
|
||||
sampling_param_list[i] = stage_param
|
||||
|
||||
return sampling_param_list
|
||||
|
||||
|
||||
|
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
|
||||
}
|
||||
@@ -43,6 +43,27 @@ The benchmark supports three dataset modes via `--dataset`:
|
||||
|
||||
### VBench dataset
|
||||
|
||||
`vbench` only provides prompt data (and image paths for i2v/i2i); it does not carry
|
||||
per-request generation fields. In this mode, all requests share CLI values:
|
||||
`--width --height --num-frames --fps --num-inference-steps`
|
||||
(pass `--width` and `--height` together).
|
||||
|
||||
Example (`t2v`):
|
||||
|
||||
```bash
|
||||
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
--base-url http://localhost:8099 \
|
||||
--model Wan-AI/Wan2.2-T2V-A14B-Diffusers \
|
||||
--task t2v \
|
||||
--dataset vbench \
|
||||
--num-prompts 50 \
|
||||
--width 640 --height 480 \
|
||||
--num-frames 81 --fps 16 \
|
||||
--num-inference-steps 40
|
||||
```
|
||||
|
||||
Note: `vbench` can also be used for other tasks such as `t2i` / `i2v` (and `i2i`). For `t2i`, the loader reuses VBench t2v text prompts; for `i2v` / `i2i`, it loads the VBench i2v dataset (with image paths).
|
||||
|
||||
If you use i2v/i2i bench datasets and need auto-download support, you may need:
|
||||
|
||||
```bash
|
||||
@@ -90,7 +111,7 @@ Common optional flags:
|
||||
Related flags: `--width`, `--height`, `--num-frames`, `--fps`, `--num-inference-steps`.
|
||||
|
||||
- For `vbench` / `random`: these CLI flags act as global defaults for all generated requests.
|
||||
- For `trace`: each request can carry its own fields (e.g., `width/height/num_frames/num_inference_steps`).
|
||||
- For `trace`: requests can carry their own fields (e.g., `width/height/num_frames/num_inference_steps`), with overrides/fallbacks as below.
|
||||
|
||||
Precedence rules for `trace` (i.e., what actually gets sent):
|
||||
|
||||
|
||||
@@ -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,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
|
||||
@@ -7,7 +7,7 @@ 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/*
|
||||
|
||||
|
||||
@@ -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/*
|
||||
|
||||
|
||||
@@ -14,6 +14,7 @@ RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRO
|
||||
RUN apt clean && apt-get update -y && \
|
||||
apt-get install -y --no-install-recommends --fix-missing \
|
||||
curl \
|
||||
espeak-ng \
|
||||
ffmpeg \
|
||||
git \
|
||||
libsndfile1 \
|
||||
@@ -75,7 +76,7 @@ ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
|
||||
ENV UV_INDEX_STRATEGY="unsafe-best-match"
|
||||
ENV UV_LINK_MODE="copy"
|
||||
|
||||
ARG VLLM_VERSION=v0.16.0
|
||||
ARG VLLM_VERSION=v0.17.0
|
||||
RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm
|
||||
WORKDIR /workspace/vllm
|
||||
|
||||
|
||||
@@ -13,9 +13,14 @@ nav:
|
||||
- examples/README.md
|
||||
- Offline Inference:
|
||||
- BAGEL-7B-MoT: user_guide/examples/offline_inference/bagel.md
|
||||
- CosyVoice3: user_guide/examples/offline_inference/cosyvoice3.md
|
||||
- Fish Speech S2 Pro: user_guide/examples/offline_inference/fish_speech.md
|
||||
- GLM-Image Multistage End-to-End Inference: user_guide/examples/offline_inference/glm_image.md
|
||||
- Helios Video Generation: user_guide/examples/offline_inference/helios.md
|
||||
- HunyuanImage-3.0 Image-to-Text Inference: user_guide/examples/offline_inference/hunyuan_image3.md
|
||||
- Image-To-Image: user_guide/examples/offline_inference/image_to_image.md
|
||||
- Image-To-Video: user_guide/examples/offline_inference/image_to_video.md
|
||||
- MammothModa2-Preview: user_guide/examples/offline_inference/mammothmodal2_preview.md
|
||||
- MiMo-Audio Offline Inference: user_guide/examples/offline_inference/mimo_audio.md
|
||||
- Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md
|
||||
- Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md
|
||||
@@ -23,8 +28,11 @@ nav:
|
||||
- Text-To-Audio: user_guide/examples/offline_inference/text_to_audio.md
|
||||
- Text-To-Image: user_guide/examples/offline_inference/text_to_image.md
|
||||
- Text-To-Video: user_guide/examples/offline_inference/text_to_video.md
|
||||
- X-To-Video-Audio: user_guide/examples/offline_inference/x_to_video_audio.md
|
||||
- Online Serving:
|
||||
- BAGEL-7B-MoT: user_guide/examples/online_serving/bagel.md
|
||||
- vLLM-Omni Helm Chart: user_guide/examples/online_serving/chart-helm.md
|
||||
- Fish Speech S2 Pro: user_guide/examples/online_serving/fish_speech.md
|
||||
- GLM-Image Online Serving: user_guide/examples/online_serving/glm_image.md
|
||||
- Image-To-Image: user_guide/examples/online_serving/image_to_image.md
|
||||
- Image-To-Video: user_guide/examples/online_serving/image_to_video.md
|
||||
@@ -92,6 +100,7 @@ nav:
|
||||
- api/vllm_omni
|
||||
- CLI Reference: cli
|
||||
- Community:
|
||||
- Governance: community/governance.md
|
||||
- community/*
|
||||
- Slack: https://slack.vllm.ai
|
||||
- Blog: https://blog.vllm.ai
|
||||
|
||||
@@ -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?slide=id.p1#slide=id.p1)
|
||||
- [vllm-omni announcement blogpost](https://blog.vllm.ai/2025/11/30/vllm-omni.html)
|
||||
|
Before Width: | Height: | Size: 125 KiB After Width: | Height: | Size: 98 KiB |
@@ -0,0 +1,52 @@
|
||||
# Governance
|
||||
|
||||
vLLM-Omni's governance is inspired by the [vLLM governance process](https://docs.vllm.ai/en/latest/governance/process/). We share the same commitment to open source community and meritocratic norms.
|
||||
|
||||
## Values
|
||||
|
||||
vLLM-Omni aims to be the fastest and easiest-to-use omni-modality inference and serving engine. Our values are aligned with [vLLM's values](https://docs.vllm.ai/en/latest/governance/process/#values):
|
||||
|
||||
### Design Values
|
||||
|
||||
1. **Top performance**: System performance is our top priority. We monitor overheads, optimize kernels, and publish benchmarks. We never leave performance on the table.
|
||||
2. **Ease of use**: vLLM-Omni must be simple to install, configure, and operate. We provide clear documentation, fast startup, clean logs, helpful error messages, and monitoring guides. Many users fork our code or study it deeply, so we keep it readable and modular.
|
||||
3. **Wide coverage**: vLLM-Omni supports frontier models and high-performance accelerators. We make it easy to add new models and hardware. vLLM-Omni + PyTorch form a simple interface that avoids complexity.
|
||||
4. **Production ready**: vLLM-Omni runs 24/7 in production. It must be easy to operate and monitor for health issues.
|
||||
5. **Extensibility**: vLLM-Omni serves as fundamental omni-modality infrastructure. Our codebase cannot cover every use case, so we design for easy forking and customization.
|
||||
|
||||
### Collaboration Values
|
||||
|
||||
1. **Tightly Knit and Fast-Moving**: Our maintainer team is aligned on vision, philosophy, and roadmap. We work closely to unblock each other and move quickly.
|
||||
2. **Individual Merit**: No one buys their way into governance. Committer status belongs to individuals, not companies. We reward contribution, maintenance, and project stewardship.
|
||||
|
||||
## Project Maintainers
|
||||
|
||||
### Lead Maintainers
|
||||
|
||||
Lead maintainers are responsible for the overall direction and strategy of the project:
|
||||
|
||||
- [@Gaohan123](https://github.com/Gaohan123)
|
||||
- [@hsliuustc0106](https://github.com/hsliuustc0106)
|
||||
- [@ywang96](https://github.com/ywang96)
|
||||
|
||||
### Active Committers
|
||||
|
||||
Committers have write access and merge rights. They typically have deep expertise in specific areas of this project and shepherd the community contributions:
|
||||
|
||||
- [@david6666666](https://github.com/david6666666): Quantization and Community Relationship
|
||||
- [@gcanlin](https://github.com/gcanlin): Hardware plugin and NPU integration
|
||||
- [@Isotr0py](https://github.com/Isotr0py): Diffusion and Quantization
|
||||
- [@linyueqian](https://github.com/linyueqian): TTS and Omni Support
|
||||
- [@princepride](https://github.com/princepride): Diffusion and Omni Support
|
||||
- [@SamitHuang](https://github.com/SamitHuang): RL and Diffusion
|
||||
- [@tzhouam](https://github.com/tzhouam): Engine and New Model Support
|
||||
- [@wtomin](https://github.com/wtomin):
|
||||
- [@ZJY0516](https://github.com/ZJY0516): Diffusion and CustomOp
|
||||
|
||||
## Meetings
|
||||
|
||||
Committers hold **bi-weekly meetings** to discuss future directions and collaborations of the project.
|
||||
|
||||
## Committer Nomination Process
|
||||
|
||||
Every month, any active committer can nominate new committer(s) to the project. Up to **two new committers** will be admitted per month based on the quality and impact of their contributions.
|
||||
@@ -18,7 +18,7 @@ If users want to modify some part of it. The custom stage_configs file can be in
|
||||
For offline (Assume necessary dependencies have ben imported):
|
||||
```python
|
||||
model_name = "Qwen/Qwen2.5-Omni-7B"
|
||||
omni_llm = OmniLLM(model=model_name, stage_configs_path="/path/to/custom_stage_configs.yaml")
|
||||
omni_llm = Omni(model=model_name, stage_configs_path="/path/to/custom_stage_configs.yaml")
|
||||
```
|
||||
|
||||
For online serving:
|
||||
|
||||
@@ -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`
|
||||
|
||||
|
||||
@@ -17,6 +17,7 @@ Defined in `pyproject.toml`:
|
||||
| `gpu` | Tests that run on GPU * |
|
||||
| `cuda` | Tests that run on CUDA * |
|
||||
| `rocm` | Tests that run on AMD/ROCm * |
|
||||
| `xpu` | Tests that run on Intel XPU * |
|
||||
| `npu` | Tests that run on NPU/Ascend * |
|
||||
| `H100` | Tests that require H100 GPU * |
|
||||
| `L4` | Tests that require L4 GPU * |
|
||||
@@ -32,7 +33,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,12 +50,13 @@ 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:
|
||||
|
||||
1. **Applies platform and resource markers**
|
||||
Adds the appropriate pytest markers for each specified hardware platform (e.g., `cuda`, `rocm`, `npu`) and resource type (e.g., `L4`, `H100`, `MI325`, `A2`, `A3`).
|
||||
Adds the appropriate pytest markers for each specified hardware platform (e.g., `cuda`, `rocm`, `xpu`, `npu`) and resource type (e.g., `L4`, `H100`, `MI325`, `B60`, `A2`, `A3`).
|
||||
```
|
||||
@pytest.mark.cuda
|
||||
@pytest.mark.L4
|
||||
@@ -79,14 +81,14 @@ This decorator is intended to make hardware-aware, cross-platform test authoring
|
||||
- Single call for multiple platforms:
|
||||
```python
|
||||
@hardware_test(
|
||||
res={"cuda": "L4", "rocm": "MI325", "npu": "A2"},
|
||||
num_cards={"cuda": 2, "rocm": 2, "npu": 2},
|
||||
res={"cuda": "L4", "rocm": "MI325", "xpu": "B60", "npu": "A2"},
|
||||
num_cards={"cuda": 2, "rocm": 2, "xpu": 2, "npu": 2},
|
||||
)
|
||||
```
|
||||
or
|
||||
```python
|
||||
@hardware_test(
|
||||
res={"cuda": "L4", "rocm": "MI325", "npu": "A2"},
|
||||
res={"cuda": "L4", "rocm": "MI325", "xpu": "B60", "npu": "A2"},
|
||||
num_cards=2,
|
||||
)
|
||||
```
|
||||
@@ -98,6 +100,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 +152,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 +172,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).
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -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 `codec_chunk_frames` (default=25) before sending. Set `initial_codec_chunk_frames` to emit smaller chunks during the initial phase for reduced TTFA
|
||||
- **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
|
||||
|
||||
|
||||
@@ -72,3 +72,241 @@ Parameters:
|
||||
|
||||
For more details, refer to the
|
||||
[Mooncake repository](https://github.com/kvcache-ai/Mooncake).
|
||||
|
||||
---
|
||||
|
||||
## Design
|
||||
|
||||
### 1. Overview
|
||||
|
||||
`MooncakeStoreConnector` is the storage-oriented remote connector in `vllm_omni/distributed/omni_connectors`. It is built on top of Mooncake distributed store APIs and provides a uniform `put()` / `get()` abstraction for transferring stage payloads across nodes.
|
||||
|
||||
Compared with `SharedMemoryConnector`, this connector is intended for multi-node deployments. Compared with `MooncakeTransferEngineConnector`, it is simpler and more object-oriented: data is written into a distributed store and fetched back by key, rather than transferred through an explicit remote-memory write protocol.
|
||||
|
||||
Its primary role is to provide a general-purpose remote connector for stage payloads when:
|
||||
|
||||
- the pipeline spans multiple hosts
|
||||
- a shared-memory connector is not possible
|
||||
- a simpler remote backend is preferred over the RDMA-oriented transfer engine
|
||||
|
||||
### 2. Relationship with the OmniConnector System
|
||||
|
||||
`MooncakeStoreConnector` implements `OmniConnectorBase`, so it integrates with the same connector lifecycle as the other backends:
|
||||
|
||||
- `OmniConnectorFactory` instantiates it from `ConnectorSpec`
|
||||
- `load_omni_transfer_config()` maps edge config to the connector spec
|
||||
- All callers (batch forwarding, chunk transfer, KV transfer, etc.) interact with it through the same `put()` / `get()` contract
|
||||
|
||||
This means the rest of the pipeline does not need store-specific logic. It only relies on the generic connector contract.
|
||||
|
||||
### 3. Design Goals
|
||||
|
||||
The connector is designed around the following goals:
|
||||
|
||||
- **Cross-node object transfer** using a single key-based abstraction
|
||||
- **Minimal connector-specific control plane**, since the store itself is the shared medium
|
||||
- **Compatibility with arbitrary Python payloads** through the Omni serializer
|
||||
- **Operational simplicity** compared with the transfer-engine-based remote connector
|
||||
|
||||
The connector is intentionally not optimized for zero-copy tensor movement or direct remote-memory access.
|
||||
|
||||
### 4. Core Design
|
||||
|
||||
#### 4.1 Object-Oriented Transfer Model
|
||||
|
||||
`MooncakeStoreConnector` treats the transport backend as a distributed object store:
|
||||
|
||||
1. Serialize the Python object into bytes.
|
||||
2. Generate a unique connector key.
|
||||
3. Store the bytes in Mooncake.
|
||||
4. Fetch the bytes from Mooncake on the receiver side.
|
||||
5. Deserialize the bytes back into the original object.
|
||||
|
||||
This design makes the connector conceptually close to a distributed key-value transport.
|
||||
|
||||
#### 4.2 Key Construction
|
||||
|
||||
The connector uses `OmniConnectorBase._make_key()` to derive the internal store key:
|
||||
|
||||
```text
|
||||
{request_key}@{from_stage}_{to_stage}
|
||||
```
|
||||
|
||||
This adds stage routing information to the logical request key and avoids collisions between different pipeline edges.
|
||||
|
||||
#### 4.3 No Extra Metadata Path
|
||||
|
||||
Unlike `SharedMemoryConnector` and `MooncakeTransferEngineConnector`, this connector returns:
|
||||
|
||||
```python
|
||||
metadata = None
|
||||
```
|
||||
|
||||
from `put()`.
|
||||
|
||||
This is a meaningful design difference:
|
||||
|
||||
- the store itself is the shared rendezvous point
|
||||
- consumers only need the same key
|
||||
- no transport handle, address, or side-channel metadata needs to be propagated
|
||||
|
||||
As a result, the control plane is simpler for this connector than for the SHM and RDMA variants.
|
||||
|
||||
### 5. Initialization
|
||||
|
||||
#### 5.1 Required Mooncake Components
|
||||
|
||||
The connector requires the Mooncake Python bindings to expose:
|
||||
|
||||
- `MooncakeDistributedStore`
|
||||
- `ReplicateConfig`
|
||||
|
||||
If these symbols are unavailable, construction fails immediately with `ImportError`. This keeps startup failures explicit and avoids silent fallback behavior.
|
||||
|
||||
#### 5.2 Store Setup
|
||||
|
||||
During `_init_store()`, the connector:
|
||||
|
||||
1. creates `MooncakeDistributedStore()`
|
||||
2. calls `store.setup(...)`
|
||||
3. validates the return code
|
||||
4. creates a `ReplicateConfig`
|
||||
5. enables `with_soft_pin = True`
|
||||
|
||||
The setup step is completed during connector construction, not lazily at first use.
|
||||
|
||||
### 6. Put / Get Flow
|
||||
|
||||
#### 6.1 Producer Flow: `put()`
|
||||
|
||||
The producer-side flow is:
|
||||
|
||||
1. Validate that the store has been initialized.
|
||||
2. Serialize the payload into bytes.
|
||||
3. Build the internal key with stage routing information.
|
||||
4. Call `store.put(key, serialized_data, self.pin)`.
|
||||
5. Update metrics and return success.
|
||||
|
||||
The returned tuple is:
|
||||
|
||||
```python
|
||||
(True, len(serialized_data), None)
|
||||
```
|
||||
|
||||
#### 6.2 Consumer Flow: `get()`
|
||||
|
||||
The consumer-side flow is:
|
||||
|
||||
1. Validate that the store has been initialized.
|
||||
2. Build the same internal key.
|
||||
3. Poll `store.get(key)` for up to 20 retries.
|
||||
4. Sleep 50 ms between retries.
|
||||
5. If data is found, deserialize it and return `(data, payload_size)`.
|
||||
6. If all retries are exhausted, record a timeout and return `None`.
|
||||
|
||||
This design gives the connector a bounded waiting model rather than an indefinite blocking get.
|
||||
|
||||
### 7. Integration with Stage Communication
|
||||
|
||||
All callers use the connector through the same `put()` / `get()` contract:
|
||||
|
||||
- the sender calls `put()` to serialize and store the payload
|
||||
- the receiver calls `get()` to retrieve and deserialize it
|
||||
- no connector-specific metadata is required, since the store key is the rendezvous point
|
||||
|
||||
Because `put()` returns `metadata=None`, the connector is naturally compatible with callers that do not forward metadata (e.g. polling-based flows). The trade-off is that all payloads incur full serialization and deserialization costs, which makes the connector functional but not the highest-performance option for large payloads such as KV cache blocks.
|
||||
|
||||
### 8. Data Flow in the Pipeline
|
||||
|
||||
The end-to-end transfer model is:
|
||||
|
||||
```mermaid
|
||||
sequenceDiagram
|
||||
participant SenderStage
|
||||
participant MooncakeStoreConnector
|
||||
participant MooncakeStore
|
||||
participant ReceiverStage
|
||||
|
||||
SenderStage->>MooncakeStoreConnector: put(from_stage, to_stage, put_key, data)
|
||||
MooncakeStoreConnector->>MooncakeStoreConnector: serialize object
|
||||
MooncakeStoreConnector->>MooncakeStore: put(internal_key, bytes)
|
||||
MooncakeStoreConnector-->>SenderStage: (success, size, None)
|
||||
|
||||
ReceiverStage->>MooncakeStoreConnector: get(from_stage, to_stage, get_key)
|
||||
MooncakeStoreConnector->>MooncakeStore: get(internal_key)
|
||||
MooncakeStoreConnector->>MooncakeStoreConnector: deserialize bytes
|
||||
MooncakeStoreConnector-->>ReceiverStage: (data, size)
|
||||
```
|
||||
|
||||
This is a store-mediated remote transfer model rather than a direct peer-to-peer transport model.
|
||||
|
||||
### 9. Strengths and Trade-offs
|
||||
|
||||
#### Strengths
|
||||
|
||||
- Simple conceptual model: store by key, fetch by key.
|
||||
- No connector-specific metadata handoff is required.
|
||||
- Works naturally for remote multi-node stage transfer.
|
||||
- Easy to integrate into the existing connector abstraction.
|
||||
|
||||
#### Trade-offs
|
||||
|
||||
- Full serialization and deserialization are always required.
|
||||
- Large payloads are more expensive than in direct-memory transports.
|
||||
- Runtime behavior depends on external Mooncake services being available.
|
||||
- Cleanup semantics are weaker than request-scoped local buffer management.
|
||||
|
||||
### 10. Important Implementation Characteristics
|
||||
|
||||
#### 10.1 Cleanup Is a No-op
|
||||
|
||||
`cleanup()` only logs a debug message and does not actively delete remote data. The current design assumes Mooncake-side lifecycle management rather than explicit per-request removal.
|
||||
|
||||
This keeps the connector implementation simple but means request-level reclamation is not modeled inside the connector itself.
|
||||
|
||||
#### 10.2 Close Releases the Store Handle
|
||||
|
||||
Unlike `SharedMemoryConnector`, where `close()` is a no-op, `MooncakeStoreConnector.close()` performs a meaningful shutdown:
|
||||
|
||||
1. Calls `self.store.close()` to release the Mooncake store handle.
|
||||
2. Sets `self.store = None` to prevent further operations.
|
||||
|
||||
This ensures the connector releases its connection to the external Mooncake service on shutdown. Errors during close are logged but do not propagate.
|
||||
|
||||
#### 10.3 Health Output Uses a Shared Schema
|
||||
|
||||
`health()` reports:
|
||||
|
||||
- `host`
|
||||
- `metadata_server`
|
||||
- `master`
|
||||
- metrics
|
||||
|
||||
and also includes placeholder fields such as `protocol`, `pool_device`, and `pool_size` to keep a more uniform shape across connector health outputs. This reflects a system-level consistency choice rather than a store-specific design need.
|
||||
|
||||
#### 10.4 Get Is Retry-Based, Not Event-Driven
|
||||
|
||||
The receiver side polls the store with a short retry loop. This is simple and robust, but it means the connector is latency-sensitive to:
|
||||
|
||||
- store visibility delay
|
||||
- network jitter
|
||||
- payload size
|
||||
|
||||
If the payload is not visible within the retry window, the connector reports a timeout.
|
||||
|
||||
### 11. Summary
|
||||
|
||||
`MooncakeStoreConnector` is the remote, store-based connector in the OmniConnector stack. Its design is straightforward:
|
||||
|
||||
- serialize the payload
|
||||
- store it in Mooncake under a stage-qualified key
|
||||
- fetch it by the same key on the receiver side
|
||||
- deserialize it back into the original object
|
||||
|
||||
This connector fills an important gap in the system:
|
||||
|
||||
1. It enables cross-node transfer without requiring shared memory.
|
||||
2. It keeps the stage communication model uniform.
|
||||
3. It provides a simpler operational alternative to the transfer-engine-based remote connector.
|
||||
|
||||
Its simplicity is also its main trade-off: it favors a clean object-transfer model over the fast-path and peer-to-peer optimizations implemented by `MooncakeTransferEngineConnector`.
|
||||
@@ -205,3 +205,589 @@ for test-specific setup instructions.
|
||||
|
||||
For more details on the underlying engine, refer to the
|
||||
[Mooncake repository](https://github.com/kvcache-ai/Mooncake).
|
||||
|
||||
---
|
||||
|
||||
## Design
|
||||
|
||||
### 1. Overview
|
||||
|
||||
`MooncakeTransferEngineConnector` is the high-performance remote connector in `vllm_omni/distributed/omni_connectors`. It is built on top of Mooncake `TransferEngine` and combines:
|
||||
|
||||
- a **direct data plane** for remote memory writes
|
||||
- a **ZMQ side channel** for metadata lookup, handshake, and completion signaling
|
||||
- a **managed local memory pool** for both send and receive buffers
|
||||
|
||||
Unlike `MooncakeStoreConnector`, which treats the backend as a distributed store, `MooncakeTransferEngineConnector` is designed as a peer-to-peer transport. Its goal is to move large stage payloads efficiently while still fitting the common `put()` / `get()` API defined by `OmniConnectorBase`.
|
||||
|
||||
It is the most performance-oriented connector in the current OmniConnector family and is intended for large remote payloads such as:
|
||||
|
||||
- KV cache transfer
|
||||
- stage hidden-state payloads
|
||||
- streaming chunk payloads
|
||||
- other binary-heavy inter-stage artifacts
|
||||
|
||||
### 2. Relationship with the OmniConnector System
|
||||
|
||||
`MooncakeTransferEngineConnector` implements the same connector contract as the other backends:
|
||||
|
||||
- `put(from_stage, to_stage, put_key, data)`
|
||||
- `get(from_stage, to_stage, get_key, metadata=None)`
|
||||
- `cleanup(request_id, ...)`
|
||||
- `health()`
|
||||
- `close()`
|
||||
|
||||
It is integrated into the system through the standard connector plumbing:
|
||||
|
||||
- `OmniConnectorFactory` constructs the connector from `ConnectorSpec`
|
||||
- `load_omni_transfer_config()` resolves the edge-level connector configuration
|
||||
- `get_connectors_config_for_stage()` and `resolve_omni_kv_config_for_stage()` inject the connector role
|
||||
- All callers (batch forwarding, chunk transfer, KV transfer, etc.) interact with it through the same `put()` / `get()` contract
|
||||
|
||||
The key system-level distinction is that this connector is **role-aware**:
|
||||
|
||||
- sender instances expose data and listen for pull requests
|
||||
- receiver instances allocate buffers and actively pull data from the sender
|
||||
|
||||
### 3. Design Goals
|
||||
|
||||
The connector is designed around four primary goals:
|
||||
|
||||
1. **High-throughput remote transfer**
|
||||
Avoid store-mediated round trips and write directly into the receiver memory region.
|
||||
|
||||
2. **Fast path for raw payloads**
|
||||
Support `torch.Tensor`, `bytes`, and `ManagedBuffer` without forcing all traffic through full object serialization.
|
||||
|
||||
3. **Unified connector abstraction**
|
||||
Preserve the same `put()` / `get()` interface used by the rest of the OmniConnector stack.
|
||||
|
||||
4. **Safe lifecycle management**
|
||||
Manage allocation, reuse, cleanup, and failure recovery for a registered memory pool.
|
||||
|
||||
### 4. Architecture Overview
|
||||
|
||||
At a high level, the connector is composed of four main subsystems:
|
||||
|
||||
```mermaid
|
||||
classDiagram
|
||||
class OmniConnectorBase {
|
||||
<<abstract>>
|
||||
+put(from_stage, to_stage, put_key, data)
|
||||
+get(from_stage, to_stage, get_key, metadata)
|
||||
+cleanup(request_id)
|
||||
+health()
|
||||
+close()
|
||||
}
|
||||
|
||||
class MooncakeTransferEngineConnector {
|
||||
+supports_raw_data: bool
|
||||
-engine: TransferEngine
|
||||
-allocator: BufferAllocator
|
||||
-pool: torch.Tensor
|
||||
-zmq_ctx: zmq.Context
|
||||
-_local_buffers: dict
|
||||
-_sender_executor: ThreadPoolExecutor
|
||||
-_listener_thread: threading.Thread
|
||||
+put(...)
|
||||
+get(...)
|
||||
+update_sender_info(sender_host, sender_zmq_port)
|
||||
+get_connection_info()
|
||||
+cleanup(request_id, from_stage, to_stage)
|
||||
+close()
|
||||
}
|
||||
|
||||
class BufferAllocator {
|
||||
-total_size: int
|
||||
-alignment: int
|
||||
-free_blocks: list
|
||||
+alloc(size) int
|
||||
+free(offset, size)
|
||||
}
|
||||
|
||||
class ManagedBuffer {
|
||||
-allocator: BufferAllocator
|
||||
-offset: int
|
||||
-size: int
|
||||
-pool_tensor: torch.Tensor
|
||||
+tensor
|
||||
+as_tensor(dtype, shape) torch.Tensor
|
||||
+to_bytes() bytes
|
||||
+release()
|
||||
}
|
||||
|
||||
class TransferEngine {
|
||||
+initialize(host, handshake, protocol, device_name)
|
||||
+register_memory(base_ptr, size)
|
||||
+batch_transfer_sync_write(remote_session, src_addrs, dst_addrs, lengths)
|
||||
+unregister_memory(base_ptr)
|
||||
+get_rpc_port() int
|
||||
}
|
||||
|
||||
class QueryRequest {
|
||||
+request_id: str
|
||||
}
|
||||
|
||||
class QueryResponse {
|
||||
+request_id: str
|
||||
+data_size: int
|
||||
+is_fast_path: bool
|
||||
}
|
||||
|
||||
class MooncakeAgentMetadata {
|
||||
+remote_hostname: str
|
||||
+remote_port: int
|
||||
+request_id: str
|
||||
+dst_addrs: list[int]
|
||||
+lengths: list[int]
|
||||
}
|
||||
|
||||
OmniConnectorBase <|-- MooncakeTransferEngineConnector
|
||||
MooncakeTransferEngineConnector *-- BufferAllocator
|
||||
MooncakeTransferEngineConnector *-- TransferEngine
|
||||
ManagedBuffer --> BufferAllocator : releases to
|
||||
ManagedBuffer --> "1" torch.Tensor : views
|
||||
MooncakeTransferEngineConnector ..> ManagedBuffer : returns / retains
|
||||
MooncakeTransferEngineConnector ..> QueryRequest : decodes
|
||||
MooncakeTransferEngineConnector ..> QueryResponse : encodes
|
||||
MooncakeTransferEngineConnector ..> MooncakeAgentMetadata : exchanges
|
||||
```
|
||||
|
||||
#### 4.1 Transfer Engine
|
||||
|
||||
Mooncake `TransferEngine` is responsible for the actual data-plane transfer. It registers local memory and performs synchronous remote writes through:
|
||||
|
||||
```python
|
||||
batch_transfer_sync_write(...)
|
||||
```
|
||||
|
||||
#### 4.2 Managed Memory Pool
|
||||
|
||||
Each connector instance owns a large pre-registered memory pool:
|
||||
|
||||
- CPU pinned memory when `memory_pool_device == "cpu"`
|
||||
- GPU memory when `memory_pool_device == "cuda"`
|
||||
|
||||
This avoids repeated memory registration and allows each transfer to allocate subranges from one long-lived pool.
|
||||
|
||||
#### 4.3 Buffer Manager
|
||||
|
||||
Two helper classes control local memory ownership:
|
||||
|
||||
- `BufferAllocator`
|
||||
Manages aligned subrange allocation and free-list merging.
|
||||
|
||||
- `ManagedBuffer`
|
||||
Represents one live slice of the pool and exposes:
|
||||
- `.tensor`
|
||||
- `.as_tensor(dtype, shape)`
|
||||
- `.to_bytes()`
|
||||
- `.release()`
|
||||
|
||||
#### 4.4 ZMQ Side Channel
|
||||
|
||||
ZMQ is used for transport coordination, not for the data payload itself. It handles:
|
||||
|
||||
- metadata query from receiver to sender
|
||||
- pull request submission
|
||||
- completion or error signaling
|
||||
- internal notification from worker threads back to the listener thread
|
||||
|
||||
This split makes the control plane lightweight while keeping the bulk payload on the transfer engine data plane.
|
||||
|
||||
### 5. Role Model
|
||||
|
||||
#### 5.1 Sender Role
|
||||
|
||||
A sender connector:
|
||||
|
||||
- accepts `put()` calls
|
||||
- stores live transfer-ready buffers in `_local_buffers`
|
||||
- starts a ZMQ listener thread
|
||||
- responds to metadata queries and pull requests from receivers
|
||||
|
||||
#### 5.2 Receiver Role
|
||||
|
||||
A receiver connector:
|
||||
|
||||
- does not bind the sender-side ZMQ listener
|
||||
- accepts `get()` calls
|
||||
- allocates receive buffers from its own pool
|
||||
- requests metadata or transfer service from the sender
|
||||
|
||||
The role is not inferred dynamically. It is injected by the stage configuration layer:
|
||||
|
||||
- incoming edge for a stage -> `role="receiver"`
|
||||
- outgoing edge for a stage -> `role="sender"`
|
||||
|
||||
This is important because incorrect role assignment would break initialization semantics.
|
||||
|
||||
#### 5.3 Host Auto-Detection
|
||||
|
||||
The `host` configuration field supports the special value `"auto"`. When set, the connector auto-detects the local IP address that would be used for external communication (via a UDP socket probe to `8.8.8.8`). If that fails, it falls back to hostname resolution, and ultimately to `127.0.0.1`.
|
||||
|
||||
This is useful in environments where the operator does not want to hard-code IP addresses in the connector config.
|
||||
|
||||
#### 5.4 RDMA Device Filtering
|
||||
|
||||
The `device_name` configuration field allows the operator to specify which RDMA NICs to use (comma-separated, e.g. `"mlx5_0,mlx5_1"`). If not set in config, the connector also checks the `RDMA_DEVICE_NAME` environment variable.
|
||||
|
||||
This is important in environments with mixed InfiniBand/RoCE NICs, where not all devices are suitable for the transfer engine.
|
||||
|
||||
### 6. Local Memory Management
|
||||
|
||||
#### 6.1 Memory Pool Registration
|
||||
|
||||
During initialization, the connector:
|
||||
|
||||
1. allocates a large pool tensor
|
||||
2. records its base pointer
|
||||
3. registers that memory with Mooncake
|
||||
4. creates a `BufferAllocator` for subrange management
|
||||
|
||||
This means every later transfer only allocates offsets inside the pre-registered pool rather than registering memory per request.
|
||||
|
||||
#### 6.2 BufferAllocator
|
||||
|
||||
`BufferAllocator` maintains a sorted free list of `(offset, size)` blocks and enforces alignment. Its responsibilities include:
|
||||
|
||||
- aligned allocation
|
||||
- freeing previously allocated blocks
|
||||
- adjacent block merging
|
||||
- double-free detection
|
||||
- overlap detection to catch corruption
|
||||
|
||||
This is a critical piece of the connector because both sender and receiver depend on long-lived pool reuse.
|
||||
|
||||
#### 6.3 ManagedBuffer
|
||||
|
||||
`ManagedBuffer` is the main fast-path data wrapper. It can:
|
||||
|
||||
- expose the pool slice as a zero-copy 1D `uint8` tensor
|
||||
- reinterpret that slice as a typed tensor
|
||||
- copy out the contents as Python `bytes`
|
||||
- release the slice back to the allocator
|
||||
|
||||
The connector uses `ManagedBuffer` in two different ways:
|
||||
|
||||
- as a send-side holder to keep the pool slice alive
|
||||
- as a receive-side return type when `is_fast_path=True`
|
||||
|
||||
### 7. Put Flow
|
||||
|
||||
#### 7.1 High-Level Behavior
|
||||
|
||||
`put()` is only valid in sender mode. Its job is to expose a payload for later remote pull by the receiver.
|
||||
|
||||
The high-level flow is:
|
||||
|
||||
1. validate connector state and role
|
||||
2. convert the input into a pool-backed transferable representation
|
||||
3. store the transfer metadata in `_local_buffers`
|
||||
4. return lightweight metadata describing how the receiver can fetch the data
|
||||
|
||||
#### 7.2 Payload Type Handling
|
||||
|
||||
`put()` supports three payload classes:
|
||||
|
||||
**A. `ManagedBuffer`**
|
||||
|
||||
If the buffer belongs to the same pool, the connector can use it directly without copying. This is the most efficient path.
|
||||
|
||||
If the buffer comes from a different pool, the connector falls back to a copy path.
|
||||
|
||||
**B. `torch.Tensor` or `bytes`**
|
||||
|
||||
These payloads are copied into the local pool and marked as fast-path data:
|
||||
|
||||
- no Omni object serialization is required
|
||||
- receiver can get a `ManagedBuffer` back
|
||||
|
||||
**C. Generic Python object**
|
||||
|
||||
Any other payload is serialized via `OmniSerializer.serialize(...)` and then copied into the pool.
|
||||
|
||||
In this case:
|
||||
|
||||
- `is_fast_path=False`
|
||||
- the receiver will deserialize back into a Python object
|
||||
|
||||
#### 7.3 Sender Metadata
|
||||
|
||||
The sender returns:
|
||||
|
||||
```python
|
||||
{
|
||||
"source_host": self.host,
|
||||
"source_port": self.zmq_port,
|
||||
"data_size": size,
|
||||
"is_fast_path": is_fast_path,
|
||||
}
|
||||
```
|
||||
|
||||
This metadata is intentionally lightweight. It tells the receiver:
|
||||
|
||||
- where the sender-side control plane lives
|
||||
- how large the remote transfer will be
|
||||
- whether the payload should be returned as a `ManagedBuffer` or a deserialized object
|
||||
|
||||
#### 7.4 Sender Buffer Table
|
||||
|
||||
The sender stores each live payload in `_local_buffers` under the stage-qualified key. Each entry contains:
|
||||
|
||||
- source addresses
|
||||
- lengths
|
||||
- the holder object
|
||||
- ownership information (`should_release`)
|
||||
- `is_fast_path`
|
||||
- creation time
|
||||
|
||||
This table is the sender-side truth source for both metadata queries and pull requests.
|
||||
|
||||
### 8. Get Flow
|
||||
|
||||
#### 8.1 High-Level Behavior
|
||||
|
||||
`get()` runs on the receiver side and performs four steps:
|
||||
|
||||
1. resolve metadata
|
||||
2. allocate a destination buffer in the local pool
|
||||
3. request the sender to write into that destination buffer
|
||||
4. return either a `ManagedBuffer` or a deserialized object
|
||||
|
||||
#### 8.2 Metadata Resolution Paths
|
||||
|
||||
The `metadata` parameter in `get()` is optional. The connector supports two resolution modes depending on whether the caller supplies it.
|
||||
|
||||
**With metadata**
|
||||
|
||||
When the caller passes metadata, the connector uses it directly. The metadata carries:
|
||||
|
||||
- `source_host` / `source_port` — sender ZMQ endpoint
|
||||
- `data_size` — payload byte count
|
||||
- `is_fast_path` — whether the receiver gets a `ManagedBuffer` or a deserialized object
|
||||
|
||||
This mode is suitable when the control plane already forwards the sender's `put()` output to the receiver.
|
||||
|
||||
**Without metadata**
|
||||
|
||||
When `get(metadata=None)` is called, the connector queries the sender over ZMQ to discover the same fields (`data_size`, `is_fast_path`). The caller must first call:
|
||||
|
||||
```python
|
||||
update_sender_info(sender_host, sender_zmq_port)
|
||||
```
|
||||
|
||||
so that the connector knows where to send the query.
|
||||
|
||||
This mode is suitable for polling-based flows (e.g. KV transfer, async chunk transfer) where the receiver does not have metadata from the control plane.
|
||||
|
||||
#### 8.3 Destination Allocation
|
||||
|
||||
Once metadata is resolved, the receiver:
|
||||
|
||||
1. allocates a subrange from its own local pool
|
||||
2. wraps it in a `ManagedBuffer`
|
||||
3. builds a `MooncakeAgentMetadata` request containing:
|
||||
- receiver hostname
|
||||
- receiver RPC port
|
||||
- request ID
|
||||
- destination addresses
|
||||
- transfer lengths
|
||||
|
||||
This tells the sender exactly where to write the incoming data.
|
||||
|
||||
#### 8.4 Transfer Completion
|
||||
|
||||
The receiver then sends the pull request over ZMQ and waits for:
|
||||
|
||||
- `TRANS_DONE`
|
||||
- or `TRANS_ERROR`
|
||||
|
||||
If the transfer succeeds:
|
||||
|
||||
- for `is_fast_path=True`, the receiver returns `(ManagedBuffer, size)`
|
||||
- for `is_fast_path=False`, the receiver copies to bytes, deserializes, releases the buffer, and returns `(object, size)`
|
||||
|
||||
### 9. Sender-Side Listener Design
|
||||
|
||||
#### 9.1 Listener Thread
|
||||
|
||||
In sender mode, the connector starts `_zmq_listener_loop()` after initialization. The listener:
|
||||
|
||||
- binds `tcp://{host}:{zmq_port}`
|
||||
- receives incoming requests
|
||||
- uses a poller for socket events and internal notifications
|
||||
- periodically reclaims stale buffers
|
||||
|
||||
If the bind fails, initialization fails immediately. The code does not silently downgrade the connector role.
|
||||
|
||||
#### 9.2 Worker Thread Pool
|
||||
|
||||
The listener hands work to `_sender_executor` so that the listener thread does not block on transfer work.
|
||||
|
||||
There are two request types:
|
||||
|
||||
- metadata query -> `_handle_query_request(...)`
|
||||
- transfer request -> `_handle_pull_request(...)`
|
||||
|
||||
#### 9.3 Query Handling
|
||||
|
||||
For metadata queries, the sender looks up the request ID in `_local_buffers` and returns:
|
||||
|
||||
- data size
|
||||
- fast-path flag
|
||||
|
||||
This supports consumers that only know the sender endpoint but not the original sender metadata.
|
||||
|
||||
#### 9.4 Pull Handling
|
||||
|
||||
For a transfer request, the sender:
|
||||
|
||||
1. locates the source addresses in `_local_buffers`
|
||||
2. constructs the remote session identifier
|
||||
3. calls `batch_transfer_sync_write(...)`
|
||||
4. replies `TRANS_DONE` or `TRANS_ERROR`
|
||||
|
||||
On success, the sender immediately calls `cleanup(meta.request_id)` and frees the producer-side buffer if it owns it.
|
||||
|
||||
This choice is important: it makes the connector effectively a **single-consumer transfer model** for each successful put/get pair.
|
||||
|
||||
### 10. Fast Path Semantics
|
||||
|
||||
This connector explicitly advertises:
|
||||
|
||||
```python
|
||||
supports_raw_data = True
|
||||
```
|
||||
|
||||
That means it can move raw payloads without forcing everything through the Omni object serializer.
|
||||
|
||||
#### Fast Path
|
||||
|
||||
For `torch.Tensor`, `bytes`, or pool-local `ManagedBuffer`:
|
||||
|
||||
- sender returns `is_fast_path=True`
|
||||
- receiver returns a `ManagedBuffer`
|
||||
- caller is responsible for calling `release()`
|
||||
|
||||
This avoids an unnecessary copy on the receiver side.
|
||||
|
||||
#### Serialized Object Path
|
||||
|
||||
For arbitrary Python objects:
|
||||
|
||||
- sender serializes the payload
|
||||
- receiver converts the receive buffer to bytes
|
||||
- receiver deserializes the object
|
||||
- receive buffer is released internally
|
||||
|
||||
This preserves a uniform object-oriented API while still allowing optimized raw-data transport when possible.
|
||||
|
||||
### 11. Failure Handling and Cleanup
|
||||
|
||||
#### 11.1 Timeouts and Socket Recovery
|
||||
|
||||
The receiver caches ZMQ REQ sockets per thread, but invalidates them after failures. This avoids reusing sockets that may be stuck in a bad state after timeout or receive errors.
|
||||
|
||||
Timeout is scaled based on payload size:
|
||||
|
||||
- a base timeout
|
||||
- plus additional time for large payloads
|
||||
|
||||
This is intended to reduce false timeouts for large remote writes.
|
||||
|
||||
#### 11.2 Stale Buffer Reclamation
|
||||
|
||||
The sender periodically reclaims old entries from `_local_buffers` using a TTL policy. This protects the memory pool from permanent leaks if a receiver crashes or never consumes a prepared payload.
|
||||
|
||||
This is a practical recovery mechanism, although the code notes that TTL cleanup can still race with very long-running in-flight transfers.
|
||||
|
||||
#### 11.3 Connector Shutdown
|
||||
|
||||
`close()` is a full resource teardown routine. It:
|
||||
|
||||
- stops the listener thread
|
||||
- shuts down the worker executor
|
||||
- releases all pending buffers
|
||||
- closes cached sockets
|
||||
- unregisters memory from the engine when supported
|
||||
- terminates the ZMQ context
|
||||
- drops the pool reference
|
||||
|
||||
This makes `MooncakeTransferEngineConnector` the most lifecycle-aware connector in the current connector family.
|
||||
|
||||
### 12. Current Implementation Constraints
|
||||
|
||||
The current code documents several important topology constraints.
|
||||
|
||||
#### 12.1 One Sender to One Receiver per Successful Transfer
|
||||
|
||||
After a successful transfer, the sender-side buffer is cleaned up immediately. This means the same prepared payload is not retained for multiple independent receivers.
|
||||
|
||||
#### 12.2 One Receiver to One Active Sender Endpoint
|
||||
|
||||
The receiver only stores one `(sender_host, sender_zmq_port)` pair through `update_sender_info(...)`. So the metadata-query mode is currently single-sender at a time.
|
||||
|
||||
#### 12.3 Explicit Buffer Ownership Matters
|
||||
|
||||
When the connector allocates a pool slice internally, it is responsible for releasing it. When a caller passes an externally owned `ManagedBuffer`, the connector keeps it alive for transfer but does not assume ownership of its eventual release.
|
||||
|
||||
These constraints are consistent with the current implementation and should be treated as design assumptions rather than incidental behavior.
|
||||
|
||||
### 13. Data Flow in the Pipeline
|
||||
|
||||
The end-to-end sender/receiver interaction is:
|
||||
|
||||
```mermaid
|
||||
sequenceDiagram
|
||||
participant SenderStage
|
||||
participant SenderConnector
|
||||
participant ReceiverConnector
|
||||
participant ReceiverStage
|
||||
|
||||
SenderStage->>SenderConnector: put(from_stage, to_stage, put_key, data)
|
||||
SenderConnector->>SenderConnector: place payload in local memory pool
|
||||
SenderConnector-->>SenderStage: metadata(source_host, source_port, data_size, is_fast_path)
|
||||
|
||||
ReceiverStage->>ReceiverConnector: get(..., metadata)
|
||||
ReceiverConnector->>ReceiverConnector: allocate destination buffer
|
||||
ReceiverConnector->>SenderConnector: ZMQ pull request with dst addr
|
||||
SenderConnector->>ReceiverConnector: TransferEngine remote write
|
||||
SenderConnector-->>ReceiverConnector: TRANS_DONE
|
||||
ReceiverConnector-->>ReceiverStage: ManagedBuffer or deserialized object
|
||||
```
|
||||
|
||||
For metadata-less polling, the flow simply adds a metadata query step before the pull request.
|
||||
|
||||
### 14. Strengths and Trade-offs
|
||||
|
||||
#### Strengths
|
||||
|
||||
- Best remote-transfer design in the current connector stack for large payloads.
|
||||
- Supports raw-data fast path.
|
||||
- Keeps stage communication under the same connector abstraction.
|
||||
- Includes real lifecycle and memory-pool management.
|
||||
- Works for both stage payload transfer and KV transfer scenarios.
|
||||
|
||||
#### Trade-offs
|
||||
|
||||
- More complex than the store-based connector.
|
||||
- Correctness depends on role injection and endpoint coordination.
|
||||
- Caller must release fast-path receive buffers.
|
||||
- Current implementation is optimized for single-consumer transfer semantics.
|
||||
|
||||
### 15. Summary
|
||||
|
||||
`MooncakeTransferEngineConnector` is the high-performance peer-to-peer transport in the OmniConnector system. Its design combines:
|
||||
|
||||
- a registered memory pool
|
||||
- a safe subrange allocator
|
||||
- a ZMQ control plane
|
||||
- a Mooncake transfer-engine data plane
|
||||
|
||||
This allows the connector to support both:
|
||||
|
||||
1. a **fast path** for raw tensors and bytes
|
||||
2. a **generic object path** for arbitrary Python payloads
|
||||
|
||||
Within vLLM-Omni, it is the connector that most directly targets performance-sensitive remote transfer, especially for large payloads and KV cache movement. Its additional complexity is deliberate: it is the connector that turns the generic OmniConnector abstraction into a transport capable of efficient remote memory movement rather than simple object storage.
|
||||
@@ -7,9 +7,10 @@ auto-configured when no explicit connector is specified for an edge.
|
||||
|
||||
## How It Works
|
||||
|
||||
- Small payloads (< threshold): serialized and passed inline in metadata.
|
||||
- Large payloads (>= threshold): stored in shared memory; the SHM name is
|
||||
returned in metadata.
|
||||
All payloads are serialized and stored in shared memory (`/dev/shm`); the SHM
|
||||
segment name is returned in metadata. The configuration exposes a
|
||||
`shm_threshold_bytes` field for a future inline-vs-SHM split, but the current
|
||||
implementation always uses shared memory regardless of payload size.
|
||||
|
||||
## Configuration
|
||||
|
||||
@@ -25,3 +26,234 @@ runtime:
|
||||
## Notes
|
||||
|
||||
- Auto-mode uses SharedMemoryConnector if no connector is declared for an edge.
|
||||
|
||||
---
|
||||
|
||||
## Design
|
||||
|
||||
### 1. Overview
|
||||
|
||||
`SharedMemoryConnector` is the default same-node connector in `vllm_omni/distributed/omni_connectors`. It is designed for stage-to-stage transfer when producer and consumer processes run on the same host and can share `/dev/shm`.
|
||||
|
||||
The connector provides a unified `put()` / `get()` API for arbitrary Python objects while keeping the control plane lightweight:
|
||||
|
||||
- The payload is serialized by the connector.
|
||||
- The serialized bytes are placed in shared memory.
|
||||
- The queue/control plane only carries a small metadata handle.
|
||||
|
||||
This makes `SharedMemoryConnector` the simplest connector in the OmniConnector family and the default fallback when an edge does not explicitly configure another backend.
|
||||
|
||||
### 2. Relationship with the OmniConnector System
|
||||
|
||||
`SharedMemoryConnector` implements `OmniConnectorBase`, so it follows the same lifecycle and API contract as the other connectors:
|
||||
|
||||
- `put(from_stage, to_stage, put_key, data)`
|
||||
- `get(from_stage, to_stage, get_key, metadata=None)`
|
||||
- `cleanup(request_id)`
|
||||
- `health()`
|
||||
- `close()`
|
||||
|
||||
Within the larger system:
|
||||
|
||||
- `load_omni_transfer_config()` automatically fills missing edges with `SharedMemoryConnector`.
|
||||
- Callers interact with the connector exclusively through the `put()` / `get()` / `cleanup()` contract — the connector does not require caller-specific logic.
|
||||
|
||||
Compared with the remote Mooncake-based connectors, `SharedMemoryConnector` is intentionally minimal and local-only.
|
||||
|
||||
### 3. Design Goals
|
||||
|
||||
The connector is built around the following goals:
|
||||
|
||||
- **Low-friction local transfer** for single-node multi-process pipelines.
|
||||
- **Unified object semantics** for arbitrary Python payloads.
|
||||
- **Small control-plane overhead** by passing only metadata through queues.
|
||||
- **Zero external dependencies** beyond Python shared memory and the existing stage utilities.
|
||||
|
||||
It is not intended to provide cross-node transfer, RDMA, or raw tensor zero-copy semantics across processes.
|
||||
|
||||
### 4. Core Design
|
||||
|
||||
#### 4.1 Serialization Model
|
||||
|
||||
`SharedMemoryConnector` always starts from a Python object and serializes it through the shared Omni serializer:
|
||||
|
||||
```python
|
||||
payload = self.serialize_obj(data)
|
||||
```
|
||||
|
||||
This keeps the connector behavior consistent with the rest of the connector stack:
|
||||
|
||||
- producer code does not need connector-specific serialization logic
|
||||
- consumer code always receives the original object after deserialization
|
||||
- the connector can reuse the same serializer used by other backends
|
||||
|
||||
#### 4.2 Shared Memory as the Data Plane
|
||||
|
||||
The actual data plane is a shared-memory segment created by:
|
||||
|
||||
- `shm_write_bytes(...)`
|
||||
- `shm_read_bytes(...)`
|
||||
|
||||
The connector stores a small metadata object such as:
|
||||
|
||||
```python
|
||||
{
|
||||
"shm": {"name": ..., "size": ...},
|
||||
"size": ...
|
||||
}
|
||||
```
|
||||
|
||||
This metadata is passed over the control plane and allows the downstream stage to locate the shared-memory segment.
|
||||
|
||||
#### 4.3 Locking Model
|
||||
|
||||
To avoid races between the producer and consumer, the connector uses a lock file per request:
|
||||
|
||||
```text
|
||||
/dev/shm/shm_{put_key}_lockfile.lock
|
||||
```
|
||||
|
||||
Locking is done with `fcntl.flock`:
|
||||
|
||||
- producer uses `LOCK_EX`
|
||||
- consumer uses `LOCK_EX`
|
||||
|
||||
Both sides acquire an exclusive lock. This ensures that the shared-memory segment is not read while it is still being written and makes the handoff safer in a multi-process environment.
|
||||
|
||||
### 5. Put / Get Flow
|
||||
|
||||
#### 5.1 Producer Flow: `put()`
|
||||
|
||||
The producer-side flow is:
|
||||
|
||||
1. Serialize the input object to bytes.
|
||||
2. Compute the payload size.
|
||||
3. Acquire the per-request lock file.
|
||||
4. Write the bytes into shared memory.
|
||||
5. Return lightweight metadata to the caller.
|
||||
|
||||
The returned tuple is:
|
||||
|
||||
```python
|
||||
(success, serialized_size, metadata)
|
||||
```
|
||||
|
||||
where `metadata` contains the shared-memory handle needed by the consumer.
|
||||
|
||||
#### 5.2 Consumer Flow: `get(metadata=...)`
|
||||
|
||||
The primary consumer path is metadata-driven:
|
||||
|
||||
1. Extract the shared-memory handle from `metadata`.
|
||||
2. Acquire the exclusive lock.
|
||||
3. Read the raw bytes from shared memory.
|
||||
4. Deserialize the bytes back into the original Python object.
|
||||
5. Remove the lock file if it still exists.
|
||||
|
||||
This is the path used by the current stage-to-stage connector flow.
|
||||
|
||||
#### 5.3 Compatibility Flow: `get(metadata=None)`
|
||||
|
||||
The connector also keeps a compatibility path for callers that only know the key:
|
||||
|
||||
1. Attempt to open the shared-memory segment by name via `SharedMemory(name=get_key)`.
|
||||
2. If the segment exists and has non-zero size, acquire the exclusive lock and read the bytes.
|
||||
3. Deserialize the bytes and return the object.
|
||||
|
||||
If the segment does not exist or any exception occurs, the call returns `None` immediately. There is no retry loop in this path -- it is a single-attempt open.
|
||||
|
||||
This path is mainly for older code paths and is not the preferred mode for the current connector pipeline.
|
||||
|
||||
### 6. Key Implementation Characteristics
|
||||
|
||||
#### 6.1 Threshold Exists, but the Current Code Always Uses SHM
|
||||
|
||||
The class keeps a `shm_threshold_bytes` field and still exposes metrics for inline writes. However, the current implementation uses:
|
||||
|
||||
```python
|
||||
if True:
|
||||
...
|
||||
```
|
||||
|
||||
inside `put()`, which means the current code path always writes to shared memory.
|
||||
|
||||
So the design still suggests a future split between:
|
||||
|
||||
- small payloads inline
|
||||
- large payloads in shared memory
|
||||
|
||||
but the current behavior is effectively:
|
||||
|
||||
- all payloads go through shared memory
|
||||
|
||||
This should be documented because it affects real runtime behavior.
|
||||
|
||||
#### 6.2 Cleanup Is Currently Passive
|
||||
|
||||
`cleanup()` is currently a no-op. The intended assumption is:
|
||||
|
||||
- the consumer reads the segment
|
||||
- the underlying shared-memory helpers unlink it
|
||||
|
||||
If the consumer never executes `get()`, the shared-memory segment may remain allocated. This means the connector relies on the normal success path for resource reclamation.
|
||||
|
||||
#### 6.3 Close Is Currently Minimal
|
||||
|
||||
`close()` is also a no-op. There is no connector-owned background thread, socket, or memory pool to tear down, so the lifecycle is simple. The trade-off is that `close()` does not scan or recover leaked shared-memory resources.
|
||||
|
||||
### 7. Data Flow in the Pipeline
|
||||
|
||||
The typical flow with `SharedMemoryConnector` is:
|
||||
|
||||
```mermaid
|
||||
sequenceDiagram
|
||||
participant SenderStage
|
||||
participant SharedMemoryConnector
|
||||
participant QueueOrControlPlane
|
||||
participant ReceiverStage
|
||||
|
||||
SenderStage->>SharedMemoryConnector: put(from_stage, to_stage, put_key, data)
|
||||
SharedMemoryConnector->>SharedMemoryConnector: serialize object
|
||||
SharedMemoryConnector->>SharedMemoryConnector: write bytes to /dev/shm
|
||||
SharedMemoryConnector-->>SenderStage: metadata {shm, size}
|
||||
SenderStage->>QueueOrControlPlane: forward connector metadata
|
||||
QueueOrControlPlane->>ReceiverStage: task + connector metadata
|
||||
ReceiverStage->>SharedMemoryConnector: get(from_stage, to_stage, get_key, metadata)
|
||||
SharedMemoryConnector->>SharedMemoryConnector: read bytes from /dev/shm
|
||||
SharedMemoryConnector->>SharedMemoryConnector: deserialize object
|
||||
SharedMemoryConnector-->>ReceiverStage: (data, size)
|
||||
```
|
||||
|
||||
This is a classic split-control-plane / data-plane design, but constrained to a single host.
|
||||
|
||||
### 8. Strengths and Trade-offs
|
||||
|
||||
#### Strengths
|
||||
|
||||
- Very simple deployment model.
|
||||
- No external service dependency.
|
||||
- Fits naturally into the existing queue-driven orchestration flow.
|
||||
- Good default for local multi-process pipelines.
|
||||
|
||||
#### Trade-offs
|
||||
|
||||
- Same-node only.
|
||||
- Full object serialization and deserialization are still required.
|
||||
- Resource cleanup depends on the normal consumer path.
|
||||
- Shared memory capacity is limited by host configuration.
|
||||
|
||||
### 9. Summary
|
||||
|
||||
`SharedMemoryConnector` is the baseline local transport for the OmniConnector system. Its design is intentionally straightforward:
|
||||
|
||||
- serialize object
|
||||
- place bytes in shared memory
|
||||
- pass metadata through the control plane
|
||||
- deserialize on the receiving side
|
||||
|
||||
It plays two important roles in vLLM-Omni:
|
||||
|
||||
1. It is the simplest production-ready connector for same-node stage pipelines.
|
||||
2. It serves as the automatic fallback connector when no explicit edge transport is configured.
|
||||
|
||||
Although the current implementation is deliberately minimal, it provides the foundation for reliable local connector semantics and keeps the stage communication model uniform across the system.
|
||||
@@ -94,8 +94,265 @@ stage_args:
|
||||
Parameters:
|
||||
|
||||
- host: datasystem worker host.
|
||||
- port: datasystem worker port.
|
||||
- port: datasystem worker port (default: `35001` if omitted; the example above uses `31501` to match the worker startup command).
|
||||
- get_sub_timeout_ms: get timeout in milliseconds (0 for no timeout).
|
||||
|
||||
For more details, refer to the
|
||||
[Yuanrong Datasystem repository](https://atomgit.com/openeuler/yuanrong-datasystem).
|
||||
|
||||
---
|
||||
|
||||
## Design
|
||||
|
||||
### 1. Overview
|
||||
|
||||
`YuanrongConnector` is the Datasystem-based remote connector in `vllm_omni/distributed/omni_connectors`. It uses Yuanrong Datasystem's distributed key-value client as the transport backend and exposes the same `put()` / `get()` interface as the other OmniConnectors.
|
||||
|
||||
Like `MooncakeStoreConnector`, it is a store-oriented remote connector rather than a direct peer-to-peer transport. Its role is to let stage payloads move across nodes through a deterministic key-based storage abstraction while keeping the rest of the pipeline on the common connector API.
|
||||
|
||||
It is intended for deployments that already use Yuanrong Datasystem and want a remote connector that integrates with the existing OmniConnector configuration and orchestration model.
|
||||
|
||||
### 2. Relationship with the OmniConnector System
|
||||
|
||||
`YuanrongConnector` implements `OmniConnectorBase`, so it participates in the same connector lifecycle as the other backends:
|
||||
|
||||
- `OmniConnectorFactory` constructs it from a `ConnectorSpec`
|
||||
- stage edge configuration is resolved by `load_omni_transfer_config()`
|
||||
- All callers (batch forwarding, chunk transfer, KV transfer, etc.) interact with it through the same `put()` / `get()` contract
|
||||
|
||||
This means the connector is not exposed directly to stage logic. Stages only interact with the generic connector contract, and the backend choice remains a configuration concern.
|
||||
|
||||
### 3. Design Goals
|
||||
|
||||
The connector is built around the following goals:
|
||||
|
||||
1. **Cross-node payload transfer through Datasystem**
|
||||
Reuse Yuanrong Datasystem as the remote exchange medium for stage data.
|
||||
|
||||
2. **Uniform object transfer semantics**
|
||||
Allow arbitrary Python objects to be transmitted through the shared Omni serializer.
|
||||
|
||||
3. **Minimal connector-specific control plane**
|
||||
Use deterministic keys so that consumers can retrieve data without an extra transport metadata handoff.
|
||||
|
||||
4. **Operational reuse of existing infrastructure**
|
||||
Fit into environments that already deploy Yuanrong Datasystem workers and etcd.
|
||||
|
||||
The connector is not designed for direct remote-memory writes or tensor-specific fast-path transfer.
|
||||
|
||||
### 4. Core Design
|
||||
|
||||
#### 4.1 Store-Oriented Transfer Model
|
||||
|
||||
`YuanrongConnector` treats the transport backend as a distributed object store:
|
||||
|
||||
1. serialize the Python payload
|
||||
2. build a deterministic connector key
|
||||
3. write the serialized bytes into Datasystem
|
||||
4. read the bytes back on the receiver side
|
||||
5. deserialize them into the original object
|
||||
|
||||
This is the same broad architectural class as `MooncakeStoreConnector`, but implemented on top of Yuanrong Datasystem APIs instead of Mooncake store APIs.
|
||||
|
||||
#### 4.2 Deterministic Keying
|
||||
|
||||
Unlike the default `_make_key()` in `OmniConnectorBase`, `YuanrongConnector` defines its own key format:
|
||||
|
||||
```text
|
||||
{request_id}:{from_stage}_{to_stage}
|
||||
```
|
||||
|
||||
This has two design implications:
|
||||
|
||||
- the request identifier remains the primary lookup handle
|
||||
- stage routing information is embedded in the key so that the same logical request ID can safely appear on different edges
|
||||
|
||||
The explicit override also makes the key format easier to align with Datasystem-side debugging and operational inspection.
|
||||
|
||||
#### 4.3 No Extra Metadata Hand-off
|
||||
|
||||
`put()` returns:
|
||||
|
||||
```python
|
||||
(success, serialized_size, None)
|
||||
```
|
||||
|
||||
and does not generate connector-specific metadata.
|
||||
|
||||
This design works because the receiver can reconstruct the exact same key from:
|
||||
|
||||
- `get_key`
|
||||
- `from_stage`
|
||||
- `to_stage`
|
||||
|
||||
As a result, the connector does not require a separate side-channel metadata handoff.
|
||||
|
||||
### 5. Initialization
|
||||
|
||||
#### 5.1 Datasystem Client Dependency
|
||||
|
||||
The connector requires the Datasystem Python bindings to expose:
|
||||
|
||||
- `KVClient`
|
||||
- `SetParam`
|
||||
- `WriteMode`
|
||||
|
||||
If any of these symbols are unavailable, connector construction fails immediately with `ImportError`. This keeps configuration errors explicit and avoids a partially initialized runtime.
|
||||
|
||||
#### 5.2 Client Setup
|
||||
|
||||
During `_init_client()`, the connector:
|
||||
|
||||
1. reads `host` and `port`
|
||||
2. creates `KVClient(host, port)`
|
||||
3. calls `client.init()`
|
||||
|
||||
At construction time it also creates a `SetParam` and fixes:
|
||||
|
||||
```python
|
||||
self.set_param.write_mode = WriteMode.NONE_L2_CACHE_EVICT
|
||||
```
|
||||
|
||||
This means the connector has a stable write policy for all writes and does not currently expose write-mode selection as a higher-level connector option.
|
||||
|
||||
### 6. Put / Get Flow
|
||||
|
||||
#### 6.1 Producer Flow: `put()`
|
||||
|
||||
The producer-side flow is:
|
||||
|
||||
1. verify that the Datasystem client has been initialized
|
||||
2. serialize the input object with the shared Omni serializer
|
||||
3. build the Datasystem key using the connector-specific `_make_key()`
|
||||
4. call `client.set(key, serialized_data, self.set_param.write_mode)`
|
||||
5. update metrics and return success
|
||||
|
||||
The returned metadata is always `None`, because the Datasystem key itself is the lookup contract between producer and consumer.
|
||||
|
||||
#### 6.2 Consumer Flow: `get()`
|
||||
|
||||
The consumer-side flow is:
|
||||
|
||||
1. verify that the Datasystem client has been initialized
|
||||
2. rebuild the same key with `from_stage`, `to_stage`, and `get_key`
|
||||
3. call:
|
||||
|
||||
```python
|
||||
client.get([key], False, self.get_sub_timeout_ms)
|
||||
```
|
||||
|
||||
4. take the first returned element if present
|
||||
5. deserialize it and return `(data, payload_size)`
|
||||
|
||||
If the returned list is empty or contains no data for the key, `get()` returns `None`.
|
||||
|
||||
### 7. Timeout and Retrieval Semantics
|
||||
|
||||
The connector uses `get_sub_timeout_ms` as its read timeout. Unlike `MooncakeStoreConnector`, which performs an explicit retry loop in Python, `YuanrongConnector` delegates waiting behavior more directly to the Datasystem client call.
|
||||
|
||||
This leads to a slightly different retrieval model:
|
||||
|
||||
- `MooncakeStoreConnector`: retry-oriented polling in connector code
|
||||
- `YuanrongConnector`: single client call with backend-managed timeout behavior
|
||||
|
||||
From the connector API perspective the result is the same, but operationally the waiting behavior is more dependent on Datasystem client semantics.
|
||||
|
||||
### 8. Integration with Stage Communication
|
||||
|
||||
All callers use the connector through the same `put()` / `get()` contract:
|
||||
|
||||
- the sender calls `put()` to serialize and store the payload
|
||||
- the receiver calls `get()` to retrieve and deserialize it
|
||||
- no connector-specific metadata is required, since the Datasystem key is the rendezvous point
|
||||
|
||||
Because `put()` returns `metadata=None`, the connector is naturally compatible with callers that do not forward metadata (e.g. polling-based flows). The trade-off is that all payloads incur full serialization and deserialization costs, and there is no raw tensor fast path, which makes the connector functional but not optimized for the largest payloads.
|
||||
|
||||
### 9. Data Flow in the Pipeline
|
||||
|
||||
The end-to-end transfer model is:
|
||||
|
||||
```mermaid
|
||||
sequenceDiagram
|
||||
participant SenderStage
|
||||
participant YuanrongConnector
|
||||
participant Datasystem
|
||||
participant ReceiverStage
|
||||
|
||||
SenderStage->>YuanrongConnector: put(from_stage, to_stage, put_key, data)
|
||||
YuanrongConnector->>YuanrongConnector: serialize object
|
||||
YuanrongConnector->>Datasystem: set(key, bytes)
|
||||
YuanrongConnector-->>SenderStage: (success, size, None)
|
||||
|
||||
ReceiverStage->>YuanrongConnector: get(from_stage, to_stage, get_key)
|
||||
YuanrongConnector->>Datasystem: get([key], timeout)
|
||||
YuanrongConnector->>YuanrongConnector: deserialize bytes
|
||||
YuanrongConnector-->>ReceiverStage: (data, size)
|
||||
```
|
||||
|
||||
This is a store-mediated remote connector design with deterministic key lookup and no explicit side-channel metadata exchange.
|
||||
|
||||
### 10. Strengths and Trade-offs
|
||||
|
||||
#### Strengths
|
||||
|
||||
- Reuses existing Yuanrong Datasystem infrastructure.
|
||||
- Keeps the connector contract simple and uniform.
|
||||
- Requires no connector-specific metadata handoff.
|
||||
- Suitable for remote stage transfer in Datasystem-based deployments.
|
||||
|
||||
#### Trade-offs
|
||||
|
||||
- Always pays full serialization and deserialization cost.
|
||||
- Does not support raw bytes / tensor fast-path semantics.
|
||||
- Depends on external Datasystem worker availability.
|
||||
- Retrieval and timeout behavior are partly delegated to the backend client.
|
||||
|
||||
### 11. Important Implementation Characteristics
|
||||
|
||||
#### 11.1 Cleanup Is a No-op
|
||||
|
||||
`cleanup()` only logs a debug message and does not explicitly remove data from Datasystem. The current design assumes backend-side lifecycle or garbage collection rather than request-scoped delete semantics inside the connector.
|
||||
|
||||
This mirrors the same design trade-off seen in other store-based connectors: simplicity over explicit per-request data reclamation.
|
||||
|
||||
#### 11.2 Close Only Releases the Local Client Handle
|
||||
|
||||
`close()` does not perform a remote shutdown. It simply clears the local `client` reference and marks the connector as closed from the process perspective:
|
||||
|
||||
```python
|
||||
self.client = None
|
||||
```
|
||||
|
||||
This is appropriate for a client-based store connector, but it also means that backend resource lifecycle remains outside the connector's control.
|
||||
|
||||
#### 11.3 Error and Timeout Metrics Are Coarse-Grained
|
||||
|
||||
The connector tracks:
|
||||
|
||||
- `puts`
|
||||
- `gets`
|
||||
- `bytes_transferred`
|
||||
- `errors`
|
||||
- `timeouts`
|
||||
|
||||
In the current code, failed `put()` increments `errors`, while `get()` exceptions increment `timeouts`. This is operationally useful, but it does not distinguish between:
|
||||
|
||||
- backend timeout
|
||||
- not-found result
|
||||
- transport failure
|
||||
- deserialization failure
|
||||
|
||||
So the metrics should be read as high-level health indicators, not detailed root-cause diagnostics.
|
||||
|
||||
### 12. Summary
|
||||
|
||||
`YuanrongConnector` is the Datasystem-backed remote connector in the OmniConnector family. Its design is straightforward:
|
||||
|
||||
- serialize payloads
|
||||
- store them under a deterministic stage-qualified key
|
||||
- retrieve them by the same key
|
||||
- deserialize them on the receiving side
|
||||
|
||||
Within vLLM-Omni, it provides a clean integration point for Yuanrong Datasystem-based deployments while preserving the same connector abstraction used by the rest of the pipeline.
|
||||
|
||||
Its design priorities are simplicity, infrastructure reuse, and API consistency rather than direct-memory transport optimization.
|
||||
@@ -55,6 +55,8 @@ The `_sp_plan` mechanism allows SP **without modifying `forward()` logic**. The
|
||||
- Tensor operations happen at `nn.Module` boundaries
|
||||
- Predictable sharding/gathering patterns
|
||||
|
||||
This is the ideal approach for integrating sequence parallelism into new models, as it is easier to maintain and ensure compatibility with other types of acceleration.
|
||||
|
||||
**How it works:**
|
||||
1. Declare `_sp_plan` dict in your transformer class
|
||||
2. Framework automatically applies hooks when `sequence_parallel_size > 1`
|
||||
@@ -201,6 +203,36 @@ class TransformerWithRoPE(nn.Module):
|
||||
}
|
||||
```
|
||||
|
||||
**Pattern 3: Shard RoPE for Dual Stream Attention**
|
||||
In some cases, different streams in attention may need to handle sequence parallelism differently. For example, we may want to shard the image embeddings, while replicating the text embeddings to correctly configure joint attention.
|
||||
|
||||
```python
|
||||
class DualStreamTransformer(nn.Module):
|
||||
"""
|
||||
Dual-stream model where we need to replicate the text components, but shard
|
||||
the image components to correctly handle sequence parallelism.
|
||||
"""
|
||||
_sp_plan = {
|
||||
# In this case, the rope_preparer returns a tuple of len 4, where the
|
||||
# first 2 items correspond to the text, and the second 2 correspond to
|
||||
# visual inputs, so we only shard the second.
|
||||
"rope_preparer": {
|
||||
# Outputs 0, 1 (text) - NOT sharded (replicated)
|
||||
# Outputs 2, 3 (image) - sharded
|
||||
2: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True), # img_cos
|
||||
3: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True), # img_sin
|
||||
},
|
||||
# Shard transformer block INPUT
|
||||
"transformer_blocks.0": {
|
||||
"hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
|
||||
},
|
||||
# Gather at output
|
||||
"proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
|
||||
}
|
||||
```
|
||||
|
||||
NOTE: be careful to test adequately when refactoring classes that take this style of plan, as changing the order of the return values will break sequence parallelism.
|
||||
|
||||
### API Reference
|
||||
|
||||
**SequenceParallelInput Parameters:**
|
||||
@@ -240,7 +272,7 @@ 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. 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.
|
||||
For models with dynamic sharding logic that cannot be expressed via `_sp_plan`, manually insert shard/gather calls.
|
||||
|
||||
|
||||
**When to use:**
|
||||
@@ -253,18 +285,15 @@ 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.
|
||||
|
||||
---
|
||||
|
||||
@@ -439,6 +468,7 @@ Complete examples in the codebase:
|
||||
|
||||
| Model | Path | Pattern | Notes |
|
||||
|-------|------|---------|-------|
|
||||
| **LongCat** | `vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py` | Dual-stream | Text components replicated, image components sharded |
|
||||
| **Qwen-Image** | `vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py` | Dual-stream + preprocessing | auto_pad, separate RoPE |
|
||||
| **Wan2.2** | `vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py` | Dual-Transformer + RoPE | Video transformer |
|
||||
| **Z-Image** | `vllm_omni/diffusion/models/z_image/z_image_transformer.py` | Unified sequence | Concatenated input |
|
||||
@@ -453,7 +483,7 @@ Complete examples in the codebase:
|
||||
Adding Sequence Parallel support to a transformer:
|
||||
|
||||
1. ✅ **Choose approach** - Use `_sp_plan` for standard cases, intrusive modification for complex cases
|
||||
2. ✅ **Identify sharding boundaries** - Where should tensors be split/gathered?
|
||||
2. ✅ **Identify sharding boundaries** - Where should tensors be split/gathered? And which module boundaries need to be moved to facilitate this?
|
||||
3. ✅ **Extract inline operations** - Move `torch.cat`, `pad_sequence`, etc. to submodules
|
||||
4. ✅ **Define `_sp_plan`** - Declare shard/gather points as class attribute
|
||||
5. ✅ **Use `auto_pad` for variable lengths** - Support non-uniform sequences
|
||||
|
||||
@@ -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)]
|
||||
|
||||
@@ -49,6 +49,9 @@ 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` |
|
||||
|`FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` |
|
||||
|`DreamIDOmniPipeline`| DreamID-Omni | `XuGuo699/DreamID-Omni` |
|
||||
|
||||
|
||||
## List of Supported Models for NPU
|
||||
@@ -64,6 +67,7 @@ th {
|
||||
|--------------|--------|-------------------|
|
||||
| `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` |
|
||||
| `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B`|
|
||||
| `HunyuanImage3ForCausalMM` | HunyuanImage3.0 (DiT-only) | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` |
|
||||
| `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` |
|
||||
| `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` |
|
||||
| `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` |
|
||||
|
||||
@@ -1,6 +1,9 @@
|
||||
# Speech API
|
||||
|
||||
vLLM-Omni provides an OpenAI-compatible API for text-to-speech (TTS) generation using Qwen3-TTS models.
|
||||
vLLM-Omni provides an OpenAI-compatible API for text-to-speech (TTS) generation. Supported TTS models include:
|
||||
|
||||
- **Qwen3-TTS** (`Qwen/Qwen3-TTS-12Hz-*`) -- Qwen3-based TTS with CustomVoice, VoiceDesign, and Base (voice cloning) task types. Output: 24 kHz.
|
||||
- **Fish Speech S2 Pro** (`fishaudio/s2-pro`) -- Dual-AR TTS with DAC codec. Supports text-to-speech and voice cloning via reference audio. Output: 44.1 kHz.
|
||||
|
||||
Each server instance runs a single model (specified at startup via `vllm serve <model> --omni`).
|
||||
|
||||
@@ -9,13 +12,22 @@ Each server instance runs a single model (specified at startup via `vllm serve <
|
||||
### Start the Server
|
||||
|
||||
```bash
|
||||
# CustomVoice model (predefined speakers)
|
||||
# Qwen3-TTS: CustomVoice model (predefined speakers)
|
||||
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts.yaml \
|
||||
--omni \
|
||||
--port 8091 \
|
||||
--trust-remote-code \
|
||||
--enforce-eager
|
||||
|
||||
# Fish Speech S2 Pro
|
||||
vllm-omni serve fishaudio/s2-pro \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/fish_speech_s2_pro.yaml \
|
||||
--omni \
|
||||
--port 8091 \
|
||||
--trust-remote-code \
|
||||
--enforce-eager \
|
||||
--gpu-memory-utilization 0.9
|
||||
```
|
||||
|
||||
### Generate Speech
|
||||
@@ -96,7 +108,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 | Initial chunk size for reduced TTFA (overrides stage config) |
|
||||
| `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
|
||||
|
||||
@@ -164,6 +176,42 @@ curl -X POST http://localhost:8091/v1/audio/voices \
|
||||
-F "name=custom_voice_1"
|
||||
```
|
||||
|
||||
## Streaming Text Input (WebSocket)
|
||||
|
||||
The `/v1/audio/speech/stream` WebSocket endpoint accepts text incrementally and generates audio per sentence as boundaries are detected.
|
||||
|
||||
> Note: text input is always streamed incrementally. Audio output remains sentence-scoped:
|
||||
> use `stream_audio=false` for one binary frame per sentence, or `stream_audio=true` for one or more PCM chunks per sentence.
|
||||
|
||||
### WebSocket Protocol
|
||||
|
||||
Client -> Server:
|
||||
|
||||
| Message | Description |
|
||||
|---------|-------------|
|
||||
| `{"type": "session.config", ...}` | Session configuration (sent once, first message) |
|
||||
| `{"type": "input.text", "text": "..."}` | Text chunk |
|
||||
| `{"type": "input.done"}` | End of input, flushes remaining buffer |
|
||||
|
||||
Server -> Client:
|
||||
|
||||
| Message | Description |
|
||||
|---------|-------------|
|
||||
| `{"type": "audio.start", "sentence_index": 0, "sentence_text": "...", "format": "pcm", "sample_rate": 24000}` | Audio generation starting for a sentence |
|
||||
| Binary frame | Raw audio bytes (one or more PCM chunks when `stream_audio=true`) |
|
||||
| `{"type": "audio.done", "sentence_index": 0, "total_bytes": 96000, "error": false}` | Audio complete for a sentence |
|
||||
| `{"type": "session.done", "total_sentences": N}` | Session complete |
|
||||
| `{"type": "error", "message": "..."}` | Non-fatal error |
|
||||
|
||||
### Session Config Parameters
|
||||
|
||||
All REST API parameters are supported, plus:
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| `stream_audio` | bool | false | Stream one or more PCM chunks per sentence over WebSocket |
|
||||
| `split_granularity` | string | "sentence" | Text splitting granularity |
|
||||
|
||||
|
||||
```bash
|
||||
DELETE /v1/audio/voices/{name}
|
||||
@@ -281,6 +329,8 @@ curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
|
||||
## Supported Models
|
||||
|
||||
### Qwen3-TTS
|
||||
|
||||
| Model | Task Type | Description |
|
||||
|-------|-----------|-------------|
|
||||
| `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | CustomVoice | Predefined speaker voices with optional style control |
|
||||
@@ -289,6 +339,14 @@ curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
| `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 |
|
||||
|
||||
### Fish Speech S2 Pro
|
||||
|
||||
| Model | Description |
|
||||
|-------|-------------|
|
||||
| `fishaudio/s2-pro` | 4B dual-AR TTS with DAC codec (44.1 kHz). Supports text-to-speech and voice cloning. |
|
||||
|
||||
Fish Speech uses `ref_audio` and `ref_text` for voice cloning (no `task_type` needed). The `voice` field should be set to `"default"`. See the [Fish Speech online serving example](../user_guide/examples/online_serving/fish_speech.md) for details.
|
||||
|
||||
## Error Responses
|
||||
|
||||
### 400 Bad Request
|
||||
|
||||
|
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,28 @@ 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` | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ |
|
||||
| **DreamID-Omni** | `XuGuo699/DreamID-Omni` | ❌ | ❌ | ✅ | ❌ | ❌ | N/A |
|
||||
|
||||
!!! note "TP Limitations for Diffusion Models"
|
||||
We currently implement Tensor Parallelism (TP) only for the DiT (Diffusion Transformer) blocks. This is because the `text_encoder` component in vLLM-Omni uses the original Transformers implementation, which does not yet support TP.
|
||||
@@ -414,3 +419,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,
|
||||
),
|
||||
)
|
||||
```
|
||||
@@ -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
|
||||
```
|
||||
|
||||
|
||||
@@ -69,6 +69,7 @@ The following table shows which models are currently supported by each accelerat
|
||||
| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ |
|
||||
| **NextStep-1.1** | `stepfun-ai/NextStep-1.1` | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ |
|
||||
| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ |
|
||||
| **FLUX.2-dev** | `black-forest-labs/FLUX.2-dev` | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ |
|
||||
|
||||
### VideoGen
|
||||
|
||||
@@ -76,6 +77,7 @@ The following table shows which models are currently supported by each accelerat
|
||||
|-------|------------------|:--------:|:---------:|:----------:|:--------------:|:------------:|:----:|:----:|
|
||||
| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
|
||||
| **LTX-2** | `Lightricks/LTX-2` | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
|
||||
| **DreamID-Omni** | `XuGuo699/DreamID-Omni` | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ |
|
||||
|
||||
### Quantization
|
||||
|
||||
|
||||
@@ -0,0 +1,74 @@
|
||||
# CosyVoice3
|
||||
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/cosyvoice3>.
|
||||
|
||||
|
||||
## Setup
|
||||
|
||||
Install dependencies:
|
||||
```
|
||||
uv pip install -e .
|
||||
```
|
||||
|
||||
> **Note:** This includes required libraries such as `librosa`, `soundfile`,
|
||||
> `onnxruntime`, `x-transformers`, and `einops` via
|
||||
> `requirements/common.txt` and platform-specific requirements files.
|
||||
|
||||
Download the model snapshot:
|
||||
```
|
||||
from huggingface_hub import snapshot_download
|
||||
snapshot_download('FunAudioLLM/Fun-CosyVoice3-0.5B-2512', local_dir='pretrained_models/Fun-CosyVoice3-0.5B')
|
||||
```
|
||||
|
||||
Add `config.json` in `pretrained_models/Fun-CosyVoice3-0.5B/`:
|
||||
```json
|
||||
{
|
||||
"model_type": "cosyvoice3",
|
||||
"architectures": [
|
||||
"CosyVoice3Model"
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
> **Why `config.json` is required:**
|
||||
> `AutoConfig.register("cosyvoice3", CosyVoice3Config)` only registers a class mapping.
|
||||
> The loader still needs `model_type: "cosyvoice3"` from `config.json` to select that class.
|
||||
> If no `config.json` is present, model type cannot be inferred automatically.
|
||||
> If your downloaded checkpoint already includes a valid `config.json` with
|
||||
> `model_type: "cosyvoice3"`, this manual step can be skipped.
|
||||
|
||||
Run the offline verification script:
|
||||
```
|
||||
python examples/offline_inference/cosyvoice3/verify_e2e_cosyvoice.py \
|
||||
--model pretrained_models/Fun-CosyVoice3-0.5B \
|
||||
--tokenizer pretrained_models/Fun-CosyVoice3-0.5B/CosyVoice-BlankEN
|
||||
```
|
||||
|
||||
## Implementation Overview
|
||||
|
||||
CosyVoice3 runs as a **2-stage Omni pipeline**.
|
||||
|
||||
- **Stage 0** (`talker`) converts text + prompt audio to speech tokens.
|
||||
- **Stage 1** (`code2wav`) converts speech tokens + prompt features to waveform via flow matching and HiFiGAN.
|
||||
|
||||
Key components live in `vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py`.
|
||||
|
||||
- `CosyVoice3MultiModalProcessor` builds the multimodal inputs:
|
||||
- Tokenizes `prompt` and `prompt_text`.
|
||||
- Extracts speech tokens and mel features from the prompt audio.
|
||||
- Extracts a speaker embedding.
|
||||
- `CosyVoice3Model` implements both stages:
|
||||
- Stage 0 uses `CosyVoice3LM` and outputs speech tokens + conditioning features.
|
||||
- Stage 1 runs the flow model (DiT-based CFM) and HiFiGAN to synthesize waveform.
|
||||
|
||||
Stage wiring is configured in `vllm_omni/model_executor/stage_configs/cosyvoice3.yaml`.
|
||||
|
||||
- Stage 0 emits latent speech tokens.
|
||||
- Stage 1 consumes them via `custom_process_input_func` and outputs audio.
|
||||
|
||||
## Example materials
|
||||
|
||||
??? abstract "verify_e2e_cosyvoice.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/cosyvoice3/verify_e2e_cosyvoice.py"
|
||||
``````
|
||||
@@ -0,0 +1,97 @@
|
||||
# Fish Speech S2 Pro
|
||||
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/fish_speech>.
|
||||
|
||||
|
||||
This directory contains an offline demo for running Fish Speech S2 Pro (`fishaudio/s2-pro`) with vLLM Omni. It supports text-only synthesis and voice cloning with reference audio.
|
||||
|
||||
## Model Overview
|
||||
|
||||
[Fish Speech S2 Pro](https://huggingface.co/fishaudio/s2-pro) is a 4B dual-AR text-to-speech model by FishAudio. It uses:
|
||||
|
||||
- **Slow AR**: Generates semantic tokens autoregressively (Qwen3-based backbone)
|
||||
- **Fast AR**: Predicts residual codebook tokens from semantic tokens
|
||||
- **DAC Decoder**: Converts codec codes to 44.1 kHz audio waveform
|
||||
|
||||
The model produces high-quality speech with voice cloning capabilities.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
Install the `fish-speech` package for the DAC codec:
|
||||
|
||||
```bash
|
||||
pip install fish-speech
|
||||
```
|
||||
|
||||
## 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.
|
||||
|
||||
## Quick Start
|
||||
|
||||
### Text-Only Synthesis
|
||||
|
||||
```bash
|
||||
python end2end.py --text "Hello, this is a test of the Fish Speech text to speech system."
|
||||
```
|
||||
|
||||
Generated audio files are saved to `output_audio/` by default.
|
||||
|
||||
### Voice Cloning
|
||||
|
||||
Provide a reference audio file and its transcript to clone a voice:
|
||||
|
||||
```bash
|
||||
python end2end.py \
|
||||
--text "Hello, this is a cloned voice." \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--ref-text "Transcript of the reference audio."
|
||||
```
|
||||
|
||||
## Streaming Mode
|
||||
|
||||
Add `--streaming` to stream audio chunks progressively via `AsyncOmni` (requires `async_chunk: true` in the stage config):
|
||||
|
||||
```bash
|
||||
python end2end.py --text "Hello, this is a test." --streaming
|
||||
```
|
||||
|
||||
Each DAC decoder chunk is logged as it arrives. The final WAV file is written once generation completes.
|
||||
|
||||
## Configuration
|
||||
|
||||
| Argument | Default | Description |
|
||||
|----------|---------|-------------|
|
||||
| `--text` | `"Hello, this is a test..."` | Text to synthesize |
|
||||
| `--ref-audio` | None | Path to reference audio for voice cloning |
|
||||
| `--ref-text` | None | Transcript of the reference audio |
|
||||
| `--model` | `fishaudio/s2-pro` | HuggingFace model path |
|
||||
| `--stage-configs-path` | `fish_speech_s2_pro.yaml` | Path to stage configs YAML |
|
||||
| `--output-dir` | `output_audio` | Output directory for WAV files |
|
||||
| `--streaming` | False | Enable streaming via AsyncOmni |
|
||||
| `--stage-init-timeout` | 600 | Stage initialization timeout (seconds) |
|
||||
|
||||
## Notes
|
||||
|
||||
- Output audio is 44.1 kHz mono WAV.
|
||||
- The DAC codec weights (`codec.pth`) are loaded lazily from the model directory.
|
||||
- Voice cloning encodes the reference audio using the DAC codec to extract semantic codes, then prepends them as a system message.
|
||||
|
||||
## Example materials
|
||||
|
||||
??? abstract "diag_codebook.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/fish_speech/diag_codebook.py"
|
||||
``````
|
||||
??? abstract "diag_reference_compare.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/fish_speech/diag_reference_compare.py"
|
||||
``````
|
||||
??? abstract "end2end.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/fish_speech/end2end.py"
|
||||
``````
|
||||
??? abstract "standalone_test.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/fish_speech/standalone_test.py"
|
||||
``````
|
||||
@@ -99,7 +99,8 @@ python end2end.py --query-type CustomVoice --streaming --output-dir /tmp/out_str
|
||||
```
|
||||
|
||||
Each Code2Wav chunk is logged as it arrives (default 25 frames; configurable via `codec_chunk_frames`
|
||||
and `initial_codec_chunk_frames` in the stage config). The final WAV file is written once generation
|
||||
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.
|
||||
|
||||
@@ -0,0 +1,58 @@
|
||||
# X-To-Video-Audio
|
||||
|
||||
The `DreamID-Omni` pipeline generates short videos from text, image and video.
|
||||
|
||||
## Local CLI Usage
|
||||
### Download the Model locally
|
||||
Since DreamID-Omni combine multiple models, and without any config, so we need to download them locally.
|
||||
|
||||
```bash
|
||||
python download_dreamid_omni.py --output-dir ./dreamid_omni
|
||||
```
|
||||
After download, the model directory will look like this:
|
||||
|
||||
```
|
||||
dreamid_omni/
|
||||
├── DreamID-Omni/
|
||||
│ ├── dreamid_omni.safetensors
|
||||
├── MMAudio/
|
||||
│ ├── ext_weights/
|
||||
│ │ ├── best_netG.pt
|
||||
│ │ ├── v1-16.pth
|
||||
├── Wan2.2-TI2V-5B/
|
||||
│ ├── google/*
|
||||
│ ├── models_t5_umt5-xxl-enc-bf16.pth
|
||||
│ ├── Wan2.2_VAE.pth
|
||||
│
|
||||
├── model_index.json # create by download_dreamid_omni.py
|
||||
```
|
||||
|
||||
### Run the Inference
|
||||
```
|
||||
python x_to_video_audio.py \
|
||||
--model /xx/dreamid_omni \
|
||||
--prompt "Two people walking together and singing happily" \
|
||||
--image-path ./example0.png ./example1.png \
|
||||
--audio-path ./example0.wav ./example1.wav \
|
||||
--video-negative-prompt "jitter, bad hands, blur, distortion" \
|
||||
--audio-negative-prompt "robotic, muffled, echo, distorted" \
|
||||
--cfg-parallel-size 2 \
|
||||
--num-inference-steps 45 \
|
||||
--height 704 \
|
||||
--width 1280 \
|
||||
--output dreamid_omni.mp4
|
||||
```
|
||||
In the current test scenario (2 images + 2 audio inputs), the VRAM requirement is 72GB, regardless of whether cfg-parallel is enabled or disabled.
|
||||
The VRAM usage can be reduced by enabling CPU offload via --enable-cpu-offload.
|
||||
|
||||
Key arguments:
|
||||
- `--prompt`: text description (string).
|
||||
- `--model`: path to the model local directory.
|
||||
- `--height/--width`: output resolution (defaults 704 * 1024).
|
||||
- `--image-path`: path to the input image list.
|
||||
- `--audio-path`: path to the input audio list, indicate the timbre of the output video.
|
||||
- `--cfg-parallel-size`: number of parallel cfg parallel (defaults 1).
|
||||
- `--num-inference-steps`: number of denoising steps (defaults 45).
|
||||
- `--video-negative-prompt`: negative prompt for video generation.
|
||||
- `--audio-negative-prompt`: negative prompt for audio generation.
|
||||
- `--enable-cpu-offload`: enable CPU offload (defaults False).
|
||||
@@ -0,0 +1,194 @@
|
||||
# Fish Speech S2 Pro
|
||||
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/fish_speech>.
|
||||
|
||||
|
||||
## Model
|
||||
|
||||
| Model | Description |
|
||||
|-------|-------------|
|
||||
| `fishaudio/s2-pro` | Fish Speech S2 Pro -- 4B dual-AR TTS model with DAC codec (44.1 kHz) |
|
||||
|
||||
## Gradio Demo
|
||||
|
||||
An interactive Gradio demo is available with text-to-speech synthesis, voice cloning, and streaming support.
|
||||
|
||||
```bash
|
||||
# Option 1: Launch server + Gradio together
|
||||
./run_gradio_demo.sh
|
||||
|
||||
# Option 2: If server is already running
|
||||
python gradio_demo.py --api-base http://localhost:8091
|
||||
```
|
||||
|
||||
Then open http://localhost:7860 in your browser.
|
||||
|
||||
Features:
|
||||
|
||||
- Text-to-speech synthesis
|
||||
- Voice cloning from uploaded audio or URL
|
||||
- Streaming mode (progressive PCM playback)
|
||||
|
||||
## Launch the Server
|
||||
|
||||
```bash
|
||||
vllm-omni serve fishaudio/s2-pro \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/fish_speech_s2_pro.yaml \
|
||||
--omni \
|
||||
--port 8091 \
|
||||
--trust-remote-code \
|
||||
--enforce-eager \
|
||||
--gpu-memory-utilization 0.9
|
||||
```
|
||||
|
||||
Or use the convenience script:
|
||||
|
||||
```bash
|
||||
./run_server.sh
|
||||
```
|
||||
|
||||
## Send TTS Request
|
||||
|
||||
### Using curl
|
||||
|
||||
```bash
|
||||
# Basic TTS
|
||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "Hello, how are you?",
|
||||
"voice": "default",
|
||||
"response_format": "wav"
|
||||
}' --output output.wav
|
||||
```
|
||||
|
||||
### Voice Cloning
|
||||
|
||||
Provide a reference audio (URL or base64 data URL) and its transcript:
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "Hello, this is a cloned voice.",
|
||||
"voice": "default",
|
||||
"ref_audio": "https://example.com/reference.wav",
|
||||
"ref_text": "Transcript of the reference audio."
|
||||
}' --output cloned.wav
|
||||
```
|
||||
|
||||
### Using Python
|
||||
|
||||
```python
|
||||
import httpx
|
||||
|
||||
# Basic TTS
|
||||
response = httpx.post(
|
||||
"http://localhost:8091/v1/audio/speech",
|
||||
json={
|
||||
"input": "Hello, how are you?",
|
||||
"voice": "default",
|
||||
"response_format": "wav",
|
||||
},
|
||||
timeout=300.0,
|
||||
)
|
||||
|
||||
with open("output.wav", "wb") as f:
|
||||
f.write(response.content)
|
||||
```
|
||||
|
||||
### Using the CLI Client
|
||||
|
||||
```bash
|
||||
cd examples/online_serving/fish_speech
|
||||
|
||||
# Basic TTS
|
||||
python speech_client.py --text "Hello, how are you?"
|
||||
|
||||
# Voice cloning
|
||||
python speech_client.py \
|
||||
--text "Hello, this is a cloned voice." \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--ref-text "Transcript of the reference audio."
|
||||
|
||||
# Streaming PCM output
|
||||
python speech_client.py --text "Hello world" --stream --output output.pcm
|
||||
```
|
||||
|
||||
The CLI client supports:
|
||||
|
||||
- `--api-base`: API base URL (default: `http://localhost:8091`)
|
||||
- `--model` (or `-m`): Model name (default: `fishaudio/s2-pro`)
|
||||
- `--text`: Text to synthesize (required)
|
||||
- `--ref-audio`: Reference audio for voice cloning (local path or URL)
|
||||
- `--ref-text`: Transcript of the reference audio
|
||||
- `--stream`: Enable streaming (PCM output)
|
||||
- `--response-format`: Audio format: wav, mp3, flac, pcm, aac, opus (default: wav)
|
||||
- `--output` (or `-o`): Output file path
|
||||
|
||||
## Streaming
|
||||
|
||||
Set `stream=true` with `response_format="pcm"` to receive raw PCM audio chunks as they are decoded:
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "Hello, how are you?",
|
||||
"voice": "default",
|
||||
"stream": true,
|
||||
"response_format": "pcm"
|
||||
}' --no-buffer | play -t raw -r 44100 -e signed -b 16 -c 1 -
|
||||
```
|
||||
|
||||
**Note:** Fish Speech outputs at 44.1 kHz (unlike Qwen3-TTS which outputs at 24 kHz).
|
||||
|
||||
## API Parameters
|
||||
|
||||
Fish Speech uses the same `/v1/audio/speech` endpoint as Qwen3-TTS. See the [Speech API reference](https://docs.vllm.ai/projects/vllm-omni/en/latest/serving/speech_api/) for full parameter documentation.
|
||||
|
||||
Key parameters for Fish Speech:
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| `input` | string | **required** | Text to synthesize |
|
||||
| `voice` | string | "default" | Voice name (use "default" for Fish Speech) |
|
||||
| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
|
||||
| `ref_audio` | string | null | Reference audio URL or base64 data URL for voice cloning |
|
||||
| `ref_text` | string | null | Transcript of reference audio (required for voice cloning) |
|
||||
| `max_new_tokens` | int | 4096 | Maximum tokens to generate |
|
||||
| `stream` | bool | false | Stream raw PCM chunks |
|
||||
|
||||
## Prerequisites
|
||||
|
||||
Install the `fish-speech` package for the DAC codec:
|
||||
|
||||
```bash
|
||||
pip install fish-speech
|
||||
```
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
1. **No audio output**: Make sure the `fish-speech` package is installed for the DAC decoder
|
||||
2. **Connection refused**: Ensure the server is running on the correct port
|
||||
3. **Flashinfer version mismatch**: Set `FLASHINFER_DISABLE_VERSION_CHECK=1` if you see version warnings
|
||||
4. **Out of memory**: Reduce `--gpu-memory-utilization` or use a GPU with more VRAM
|
||||
|
||||
## Example materials
|
||||
|
||||
??? abstract "gradio_demo.py"
|
||||
``````py
|
||||
--8<-- "examples/online_serving/fish_speech/gradio_demo.py"
|
||||
``````
|
||||
??? abstract "run_gradio_demo.sh"
|
||||
``````sh
|
||||
--8<-- "examples/online_serving/fish_speech/run_gradio_demo.sh"
|
||||
``````
|
||||
??? abstract "run_server.sh"
|
||||
``````sh
|
||||
--8<-- "examples/online_serving/fish_speech/run_server.sh"
|
||||
``````
|
||||
??? abstract "speech_client.py"
|
||||
``````py
|
||||
--8<-- "examples/online_serving/fish_speech/speech_client.py"
|
||||
``````
|
||||
@@ -29,6 +29,35 @@ The script allows overriding:
|
||||
- `CACHE_BACKEND` (default: `none`)
|
||||
- `ENABLE_CACHE_DIT_SUMMARY` (default: `0`)
|
||||
|
||||
## Async Job Behavior
|
||||
|
||||
`POST /v1/videos` is asynchronous. It creates a video job and immediately
|
||||
returns metadata like the job ID and initial `queued` status. To get the final
|
||||
artifact, poll the job status and then download the completed file from the
|
||||
content endpoint.
|
||||
|
||||
The main endpoints are:
|
||||
- `POST /v1/videos`: create a video generation job
|
||||
- `GET /v1/videos/{video_id}`: retrieve the current job status and metadata
|
||||
- `GET /v1/videos`: list stored video jobs
|
||||
- `GET /v1/videos/{video_id}/content`: download the generated video file
|
||||
- `DELETE /v1/videos/{video_id}`: delete the job and any stored output
|
||||
|
||||
## Storage
|
||||
|
||||
Generated video files are stored on local disk by the async video API.
|
||||
Local file storage behavior can be controlled via the following environment variables:
|
||||
|
||||
- `VLLM_OMNI_STORAGE_PATH`: directory used for generated files (default: `/tmp/storage`)
|
||||
- `VLLM_OMNI_STORAGE_MAX_CONCURRENCY`: max concurrent save/delete operations (default: `4`)
|
||||
|
||||
Example:
|
||||
|
||||
```bash
|
||||
export VLLM_OMNI_STORAGE_PATH=/var/tmp/vllm-omni-videos
|
||||
export VLLM_OMNI_STORAGE_MAX_CONCURRENCY=8
|
||||
```
|
||||
|
||||
## API Calls
|
||||
|
||||
### Method 1: Using curl
|
||||
@@ -38,7 +67,7 @@ The script allows overriding:
|
||||
bash run_curl_image_to_video.sh
|
||||
|
||||
# Or execute directly (OpenAI-style multipart)
|
||||
curl -X POST http://localhost:8091/v1/videos \
|
||||
create_response=$(curl -s http://localhost:8091/v1/videos \
|
||||
-H "Accept: application/json" \
|
||||
-F "prompt=A bear playing with yarn, smooth motion" \
|
||||
-F "negative_prompt=low quality, blurry, static" \
|
||||
@@ -52,7 +81,23 @@ curl -X POST http://localhost:8091/v1/videos \
|
||||
-F "guidance_scale_2=1.0" \
|
||||
-F "boundary_ratio=0.875" \
|
||||
-F "flow_shift=12.0" \
|
||||
-F "seed=42" | jq -r '.data[0].b64_json' | base64 -d > wan22_i2v_output.mp4
|
||||
-F "seed=42")
|
||||
|
||||
video_id=$(echo "$create_response" | jq -r '.id')
|
||||
while true; do
|
||||
status=$(curl -s "http://localhost:8091/v1/videos/${video_id}" | jq -r '.status')
|
||||
if [ "$status" = "completed" ]; then
|
||||
break
|
||||
fi
|
||||
if [ "$status" = "failed" ]; then
|
||||
echo "Video generation failed"
|
||||
exit 1
|
||||
fi
|
||||
sleep 2
|
||||
done
|
||||
|
||||
curl -s "http://localhost:8091/v1/videos/${video_id}" | jq .
|
||||
curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o wan22_i2v_output.mp4
|
||||
```
|
||||
|
||||
## Request Format
|
||||
@@ -66,6 +111,18 @@ curl -X POST http://localhost:8091/v1/videos \
|
||||
-F "input_reference=@/path/to/qwen-bear.png"
|
||||
```
|
||||
|
||||
### Alternative JSON-Safe Reference Input
|
||||
|
||||
Use `image_reference` when you want to pass a URL or JSON-safe image reference
|
||||
instead of uploading a file. Do not send `input_reference` and
|
||||
`image_reference` together.
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/videos \
|
||||
-F "prompt=A bear playing with yarn, smooth motion" \
|
||||
-F 'image_reference={"image_url":"https://example.com/qwen-bear.png"}'
|
||||
```
|
||||
|
||||
### Generation with Parameters
|
||||
|
||||
```bash
|
||||
@@ -85,6 +142,63 @@ curl -X POST http://localhost:8091/v1/videos \
|
||||
-F "seed=42"
|
||||
```
|
||||
|
||||
## Create Response Format
|
||||
|
||||
`POST /v1/videos` returns a job record.
|
||||
|
||||
```json
|
||||
{
|
||||
"id": "video_gen_123",
|
||||
"object": "video",
|
||||
"status": "queued",
|
||||
"model": "Wan-AI/Wan2.2-I2V-A14B-Diffusers",
|
||||
"prompt": "A bear playing with yarn, smooth motion",
|
||||
"created_at": 1234567890
|
||||
}
|
||||
```
|
||||
|
||||
## Retrieve, List, Download, and Delete
|
||||
|
||||
### Retrieve a job
|
||||
|
||||
```bash
|
||||
curl -s http://localhost:8091/v1/videos/${video_id} | jq .
|
||||
```
|
||||
|
||||
### List jobs
|
||||
|
||||
```bash
|
||||
curl -s http://localhost:8091/v1/videos | jq .
|
||||
```
|
||||
|
||||
### Download the completed video
|
||||
|
||||
```bash
|
||||
curl -L http://localhost:8091/v1/videos/${video_id}/content -o wan22_i2v_output.mp4
|
||||
```
|
||||
|
||||
### Delete a job and its stored file
|
||||
|
||||
```bash
|
||||
curl -X DELETE http://localhost:8091/v1/videos/${video_id} | jq .
|
||||
```
|
||||
|
||||
## Poll Until Complete
|
||||
|
||||
```bash
|
||||
while true; do
|
||||
status=$(curl -s http://localhost:8091/v1/videos/${video_id} | jq -r '.status')
|
||||
if [ "$status" = "completed" ]; then
|
||||
break
|
||||
fi
|
||||
if [ "$status" = "failed" ]; then
|
||||
echo "Video generation failed"
|
||||
exit 1
|
||||
fi
|
||||
sleep 2
|
||||
done
|
||||
```
|
||||
|
||||
## Example materials
|
||||
|
||||
??? abstract "run_curl_image_to_video.sh"
|
||||
|
||||
@@ -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
|
||||
@@ -112,7 +137,7 @@ The Python client supports the following command-line arguments:
|
||||
- `--voice`: Speaker/voice name (default: `vivian`). Options: `vivian`, `ryan`, `aiden`, etc.
|
||||
- `--language`: Language. Options: `Auto`, `Chinese`, `English`, `Japanese`, `Korean`, `German`, `French`, `Russian`, `Portuguese`, `Spanish`, `Italian`
|
||||
- `--instructions`: Voice style/emotion instructions
|
||||
- `--ref-audio`: Reference audio file path or URL for voice cloning (Base task). Local paths are automatically base64-encoded by the client before sending to the server.
|
||||
- `--ref-audio`: Reference audio file path or URL for voice cloning (Base task)
|
||||
- `--ref-text`: Reference audio transcript for voice cloning (Base task)
|
||||
- `--response-format`: Audio output format (default: `wav`). Options: `wav`, `mp3`, `flac`, `pcm`, `aac`, `opus`
|
||||
- `--output` (or `-o`): Output audio file path (default: `tts_output.wav`)
|
||||
@@ -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
|
||||
@@ -220,7 +284,7 @@ Lists available voices for the loaded model:
|
||||
"task_type": "CustomVoice",
|
||||
"language": "Auto",
|
||||
"instructions": "Optional style instructions",
|
||||
"ref_audio": "HTTP URL, base64 data URL, or file:// URI for voice cloning",
|
||||
"ref_audio": "URL or base64 for voice cloning",
|
||||
"ref_text": "Reference audio transcript",
|
||||
"x_vector_only_mode": false,
|
||||
"max_new_tokens": 2048
|
||||
@@ -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
|
||||
@@ -261,7 +326,7 @@ Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/w
|
||||
|
||||
| Parameter | Type | Required | Description |
|
||||
| -------------------- | ------ | -------- | -------------------------------------------- |
|
||||
| `ref_audio` | string | **Yes** | Reference audio (HTTP URL, base64 data URL, or `file://` URI with `--allowed-local-media-path`) |
|
||||
| `ref_audio` | string | **Yes** | Reference audio (URL or base64 data URL) |
|
||||
| `ref_text` | string | No | Transcript of reference audio (for ICL mode) |
|
||||
| `x_vector_only_mode` | bool | No | Use speaker embedding only (no ICL) |
|
||||
|
||||
@@ -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"
|
||||
``````
|
||||
@@ -29,6 +29,34 @@ The script allows overriding:
|
||||
- `CACHE_BACKEND` (default: `none`)
|
||||
- `ENABLE_CACHE_DIT_SUMMARY` (default: `0`)
|
||||
|
||||
## Async Job Behavior
|
||||
|
||||
`POST /v1/videos` is asynchronous. It creates a video job and immediately
|
||||
returns metadata like the job ID and initial `queued` status. To get the final
|
||||
artifact, poll the job status and then download the completed file from the
|
||||
content endpoint.
|
||||
|
||||
The main endpoints are:
|
||||
- `POST /v1/videos`: create a video generation job
|
||||
- `GET /v1/videos/{video_id}`: retrieve the current job status and metadata
|
||||
- `GET /v1/videos`: list stored video jobs
|
||||
- `GET /v1/videos/{video_id}/content`: download the generated video file
|
||||
- `DELETE /v1/videos/{video_id}`: delete the job and any stored output
|
||||
|
||||
## Storage
|
||||
|
||||
Generated video files are stored on local disk by the async video API.
|
||||
|
||||
- `VLLM_OMNI_STORAGE_PATH`: directory used for generated files (default: `/tmp/storage`)
|
||||
- `VLLM_OMNI_STORAGE_MAX_CONCURRENCY`: max concurrent save/delete operations (default: `4`)
|
||||
|
||||
Example:
|
||||
|
||||
```bash
|
||||
export VLLM_OMNI_STORAGE_PATH=/var/tmp/vllm-omni-videos
|
||||
export VLLM_OMNI_STORAGE_MAX_CONCURRENCY=8
|
||||
```
|
||||
|
||||
## API Calls
|
||||
|
||||
### Method 1: Using curl
|
||||
@@ -38,7 +66,7 @@ The script allows overriding:
|
||||
bash run_curl_text_to_video.sh
|
||||
|
||||
# Or execute directly (OpenAI-style multipart)
|
||||
curl -s http://localhost:8091/v1/videos \
|
||||
create_response=$(curl -s http://localhost:8091/v1/videos \
|
||||
-H "Accept: application/json" \
|
||||
-F "prompt=Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \
|
||||
-F "width=832" \
|
||||
@@ -50,7 +78,24 @@ curl -s http://localhost:8091/v1/videos \
|
||||
-F "guidance_scale=4.0" \
|
||||
-F "guidance_scale_2=4.0" \
|
||||
-F "boundary_ratio=0.875" \
|
||||
-F "seed=42" | jq -r '.data[0].b64_json' | base64 -d > wan22_output.mp4
|
||||
-F "flow_shift=5.0" \
|
||||
-F "seed=42")
|
||||
|
||||
video_id=$(echo "$create_response" | jq -r '.id')
|
||||
while true; do
|
||||
status=$(curl -s "http://localhost:8091/v1/videos/${video_id}" | jq -r '.status')
|
||||
if [ "$status" = "completed" ]; then
|
||||
break
|
||||
fi
|
||||
if [ "$status" = "failed" ]; then
|
||||
echo "Video generation failed"
|
||||
exit 1
|
||||
fi
|
||||
sleep 2
|
||||
done
|
||||
|
||||
curl -s "http://localhost:8091/v1/videos/${video_id}" | jq .
|
||||
curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o wan22_output.mp4
|
||||
```
|
||||
|
||||
## Request Format
|
||||
@@ -85,8 +130,9 @@ curl -X POST http://localhost:8091/v1/videos \
|
||||
| Parameter | Type | Default | Description |
|
||||
| --------------------- | ------ | ------- | ------------------------------------------------ |
|
||||
| `prompt` | str | - | Text description of the desired video |
|
||||
| `seconds` | str | None | Clip duration in seconds |
|
||||
| `size` | str | None | Output size in `WIDTHxHEIGHT` format |
|
||||
| `negative_prompt` | str | None | Negative prompt |
|
||||
| `n` | int | 1 | Number of videos to generate |
|
||||
| `width` | int | None | Video width in pixels |
|
||||
| `height` | int | None | Video height in pixels |
|
||||
| `num_frames` | int | None | Number of frames to generate |
|
||||
@@ -98,24 +144,62 @@ curl -X POST http://localhost:8091/v1/videos \
|
||||
| `flow_shift` | float | None | Scheduler flow shift (Wan2.2) |
|
||||
| `seed` | int | None | Random seed (reproducible) |
|
||||
| `lora` | object | None | LoRA configuration |
|
||||
| `extra_body` | object | None | Model-specific extra parameters |
|
||||
|
||||
## Response Format
|
||||
## Create Response Format
|
||||
|
||||
`POST /v1/videos` returns a job record.
|
||||
|
||||
```json
|
||||
{
|
||||
"created": 1234567890,
|
||||
"data": [
|
||||
{ "b64_json": "<base64-mp4>" }
|
||||
]
|
||||
"id": "video_gen_123",
|
||||
"object": "video",
|
||||
"status": "queued",
|
||||
"model": "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
|
||||
"prompt": "A cinematic view of a futuristic city at sunset",
|
||||
"created_at": 1234567890
|
||||
}
|
||||
```
|
||||
|
||||
## Extract Video
|
||||
## Retrieve, List, Download, and Delete
|
||||
|
||||
### Retrieve a job
|
||||
|
||||
```bash
|
||||
# Extract base64 from response and decode to video
|
||||
cat response.json | jq -r '.data[0].b64_json' | base64 -d > wan22_output.mp4
|
||||
curl -s http://localhost:8091/v1/videos/${video_id} | jq .
|
||||
```
|
||||
|
||||
### List jobs
|
||||
|
||||
```bash
|
||||
curl -s http://localhost:8091/v1/videos | jq .
|
||||
```
|
||||
|
||||
### Download the completed video
|
||||
|
||||
```bash
|
||||
curl -L http://localhost:8091/v1/videos/${video_id}/content -o wan22_output.mp4
|
||||
```
|
||||
|
||||
### Delete a job and its stored file
|
||||
|
||||
```bash
|
||||
curl -X DELETE http://localhost:8091/v1/videos/${video_id} | jq .
|
||||
```
|
||||
|
||||
## Poll Until Complete
|
||||
|
||||
```bash
|
||||
while true; do
|
||||
status=$(curl -s http://localhost:8091/v1/videos/${video_id} | jq -r '.status')
|
||||
if [ "$status" = "completed" ]; then
|
||||
break
|
||||
fi
|
||||
if [ "$status" = "failed" ]; then
|
||||
echo "Video generation failed"
|
||||
exit 1
|
||||
fi
|
||||
sleep 2
|
||||
done
|
||||
```
|
||||
|
||||
## Example materials
|
||||
|
||||
@@ -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():
|
||||
@@ -58,6 +57,7 @@ def parse_args():
|
||||
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
|
||||
@@ -88,108 +88,80 @@ def main():
|
||||
|
||||
from PIL import Image
|
||||
|
||||
if args.modality == "img2img":
|
||||
from vllm_omni.entrypoints.omni_diffusion import OmniDiffusion
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
print(f"[Info] Running in {args.modality} mode (Stage 1 only, cfg_parallel_size={args.cfg_parallel_size})")
|
||||
omni_kwargs = {}
|
||||
if args.stage_configs_path:
|
||||
omni_kwargs["stage_configs_path"] = args.stage_configs_path
|
||||
|
||||
client = OmniDiffusion(
|
||||
model=model_name,
|
||||
parallel_config={"cfg_parallel_size": args.cfg_parallel_size},
|
||||
)
|
||||
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"]}
|
||||
if args.negative_prompt is not None:
|
||||
prompt_dict["negative_prompt"] = args.negative_prompt
|
||||
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
|
||||
extra = {
|
||||
"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"],
|
||||
}
|
||||
if args.negative_prompt is not None:
|
||||
extra["negative_prompt"] = args.negative_prompt
|
||||
diffusion_params.extra_args = extra # type: ignore
|
||||
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)
|
||||
|
||||
@@ -0,0 +1,64 @@
|
||||
# CosyVoice3
|
||||
|
||||
## Setup
|
||||
|
||||
Install dependencies:
|
||||
```
|
||||
uv pip install -e .
|
||||
```
|
||||
|
||||
> **Note:** This includes required libraries such as `librosa`, `soundfile`,
|
||||
> `onnxruntime`, `x-transformers`, and `einops` via
|
||||
> `requirements/common.txt` and platform-specific requirements files.
|
||||
|
||||
Download the model snapshot:
|
||||
```
|
||||
from huggingface_hub import snapshot_download
|
||||
snapshot_download('FunAudioLLM/Fun-CosyVoice3-0.5B-2512', local_dir='pretrained_models/Fun-CosyVoice3-0.5B')
|
||||
```
|
||||
|
||||
Add `config.json` in `pretrained_models/Fun-CosyVoice3-0.5B/`:
|
||||
```json
|
||||
{
|
||||
"model_type": "cosyvoice3",
|
||||
"architectures": [
|
||||
"CosyVoice3Model"
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
> **Why `config.json` is required:**
|
||||
> `AutoConfig.register("cosyvoice3", CosyVoice3Config)` only registers a class mapping.
|
||||
> The loader still needs `model_type: "cosyvoice3"` from `config.json` to select that class.
|
||||
> If no `config.json` is present, model type cannot be inferred automatically.
|
||||
> If your downloaded checkpoint already includes a valid `config.json` with
|
||||
> `model_type: "cosyvoice3"`, this manual step can be skipped.
|
||||
|
||||
Run the offline verification script:
|
||||
```
|
||||
python examples/offline_inference/cosyvoice3/verify_e2e_cosyvoice.py \
|
||||
--model pretrained_models/Fun-CosyVoice3-0.5B \
|
||||
--tokenizer pretrained_models/Fun-CosyVoice3-0.5B/CosyVoice-BlankEN
|
||||
```
|
||||
|
||||
## Implementation Overview
|
||||
|
||||
CosyVoice3 runs as a **2-stage Omni pipeline**.
|
||||
|
||||
- **Stage 0** (`talker`) converts text + prompt audio to speech tokens.
|
||||
- **Stage 1** (`code2wav`) converts speech tokens + prompt features to waveform via flow matching and HiFiGAN.
|
||||
|
||||
Key components live in `vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py`.
|
||||
|
||||
- `CosyVoice3MultiModalProcessor` builds the multimodal inputs:
|
||||
- Tokenizes `prompt` and `prompt_text`.
|
||||
- Extracts speech tokens and mel features from the prompt audio.
|
||||
- Extracts a speaker embedding.
|
||||
- `CosyVoice3Model` implements both stages:
|
||||
- Stage 0 uses `CosyVoice3LM` and outputs speech tokens + conditioning features.
|
||||
- Stage 1 runs the flow model (DiT-based CFM) and HiFiGAN to synthesize waveform.
|
||||
|
||||
Stage wiring is configured in `vllm_omni/model_executor/stage_configs/cosyvoice3.yaml`.
|
||||
|
||||
- Stage 0 emits latent speech tokens.
|
||||
- Stage 1 consumes them via `custom_process_input_func` and outputs audio.
|
||||
@@ -0,0 +1,197 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import argparse
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import librosa
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
from vllm import SamplingParams
|
||||
from vllm.assets.audio import AudioAsset
|
||||
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.model_executor.models.cosyvoice3.config import CosyVoice3Config
|
||||
from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer
|
||||
from vllm_omni.model_executor.models.cosyvoice3.utils import extract_text_token
|
||||
|
||||
|
||||
def _ensure_mel_filters_asset() -> None:
|
||||
repo_root = Path(__file__).resolve().parents[3]
|
||||
filters_path = repo_root / "vllm_omni" / "model_executor" / "models" / "cosyvoice3" / "assets" / "mel_filters.npz"
|
||||
if filters_path.exists():
|
||||
return
|
||||
|
||||
source_url = "https://raw.githubusercontent.com/openai/whisper/main/whisper/assets/mel_filters.npz"
|
||||
raise FileNotFoundError(
|
||||
"Missing CosyVoice3 mel filter asset:\n"
|
||||
f" {filters_path}\n"
|
||||
"Download it with:\n"
|
||||
f" mkdir -p {filters_path.parent} && "
|
||||
f"curl -L {source_url} -o {filters_path}"
|
||||
)
|
||||
|
||||
|
||||
def run_e2e():
|
||||
parser = argparse.ArgumentParser()
|
||||
# ""FunAudioLLM/Fun-CosyVoice3-0.5B-2512
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Path to CosyVoice3 model directory (e.g., pretrained_models/Fun-CosyVoice3-0.5B/).",
|
||||
)
|
||||
parser.add_argument("--stage-config", type=str, default="vllm_omni/model_executor/stage_configs/cosyvoice3.yaml")
|
||||
parser.add_argument("--prompt", type=str, default="Hello, this is a test of the CosyVoice system capability.")
|
||||
parser.add_argument(
|
||||
"--prompt-text",
|
||||
type=str,
|
||||
default="You are a helpful assistant.<|endofprompt|>Testing my voices. Why should I not?",
|
||||
)
|
||||
parser.add_argument("--audio-path", type=str, default="prompt.wav")
|
||||
parser.add_argument(
|
||||
"--tokenizer",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Path to tokenizer directory (e.g., <model_path>/CosyVoice-BlankEN).",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
_ensure_mel_filters_asset()
|
||||
# Ensure tokenizer directory exists
|
||||
if not os.path.exists(args.tokenizer):
|
||||
raise FileNotFoundError(f"{args.tokenizer} does not exist!")
|
||||
|
||||
# Ensure stage config exists
|
||||
if not os.path.exists(args.stage_config):
|
||||
raise FileNotFoundError(f"{args.stage_config} does not exist!")
|
||||
|
||||
print(f"Initializing cosyvoice E2E with model={args.model}")
|
||||
|
||||
# Initialize Omni
|
||||
# This spins up the engine(s) based on the stage config
|
||||
# We pass trust_remote_code=True same as Qwen examples
|
||||
omni = Omni(
|
||||
model=args.model,
|
||||
stage_configs_path=args.stage_config,
|
||||
trust_remote_code=True,
|
||||
tokenizer=args.tokenizer,
|
||||
log_stats=True,
|
||||
)
|
||||
|
||||
sampling_cfg = {"top_p": 0.8, "top_k": 25, "eos_token_id": 6561 + 1}
|
||||
|
||||
print("Model initialized. Preparing inputs...")
|
||||
if args.audio_path:
|
||||
if not os.path.exists(args.audio_path):
|
||||
raise FileNotFoundError(f"Audio file not found: {args.audio_path}")
|
||||
# Load at native sample rate
|
||||
audio_signal, sr = librosa.load(args.audio_path, sr=None)
|
||||
|
||||
# Validate sample rate before processing (similar to original CosyVoice)
|
||||
min_sr = 16000
|
||||
if sr < min_sr:
|
||||
raise ValueError(
|
||||
f"Audio sample rate {sr} Hz is too low. "
|
||||
f"Minimum required: {min_sr} Hz. "
|
||||
f"Please provide audio with sample rate >= {min_sr} Hz."
|
||||
)
|
||||
|
||||
audio_data = (audio_signal.astype(np.float32), sr)
|
||||
else:
|
||||
audio_data = AudioAsset("mary_had_lamb").audio_and_sample_rate
|
||||
|
||||
prompts = {
|
||||
"prompt": args.prompt,
|
||||
"multi_modal_data": {
|
||||
"audio": audio_data,
|
||||
},
|
||||
"mm_processor_kwargs": {
|
||||
"prompt_text": args.prompt_text,
|
||||
"sample_rate": audio_data[1],
|
||||
},
|
||||
}
|
||||
|
||||
print(f"Generating for prompt: {args.prompt}")
|
||||
|
||||
config = CosyVoice3Config()
|
||||
tokenizer = get_qwen_tokenizer(
|
||||
token_path=args.tokenizer,
|
||||
skip_special_tokens=config.skip_special_tokens,
|
||||
version=config.version,
|
||||
)
|
||||
_, text_token_len = extract_text_token(args.prompt, tokenizer, config.allowed_special)
|
||||
base_len = int(text_token_len)
|
||||
min_len = int(base_len * config.min_token_text_ratio)
|
||||
max_len = int(base_len * config.max_token_text_ratio)
|
||||
|
||||
# Build SamplingParams for each stage (GPT, S2Mel, Vocoder)
|
||||
gpt_sampling = SamplingParams(
|
||||
temperature=1.0,
|
||||
top_p=sampling_cfg["top_p"],
|
||||
top_k=sampling_cfg["top_k"],
|
||||
repetition_penalty=2.0,
|
||||
min_tokens=min_len,
|
||||
max_tokens=max_len,
|
||||
stop_token_ids=[sampling_cfg["eos_token_id"]],
|
||||
# allowed_token_ids=list(range(6561+3)),
|
||||
detokenize=False,
|
||||
)
|
||||
# Not used
|
||||
s2mel_sampling = SamplingParams(
|
||||
temperature=1.0,
|
||||
top_p=1.0,
|
||||
top_k=-1,
|
||||
repetition_penalty=2.0,
|
||||
max_tokens=256,
|
||||
detokenize=False,
|
||||
)
|
||||
|
||||
sampling_params_list = [gpt_sampling, s2mel_sampling]
|
||||
|
||||
# Start profiling (requires VLLM_TORCH_PROFILER_DIR env var)
|
||||
if os.environ.get("VLLM_TORCH_PROFILER_DIR"):
|
||||
print("Starting profiler...")
|
||||
omni.start_profile()
|
||||
|
||||
# Generate (Omni orchestrator requires a per-stage SamplingParams list)
|
||||
outputs = list(omni.generate(prompts, sampling_params_list=sampling_params_list[:2]))
|
||||
|
||||
# Stop profiling and get results
|
||||
if os.environ.get("VLLM_TORCH_PROFILER_DIR"):
|
||||
print("Stopping profiler...")
|
||||
profile_results = omni.stop_profile()
|
||||
print(f"Profile traces saved to: {profile_results}")
|
||||
|
||||
print(outputs)
|
||||
# Verify outputs
|
||||
print(f"Received {len(outputs)} outputs.")
|
||||
for i, output in enumerate(outputs):
|
||||
try:
|
||||
ro_list = output.request_output or []
|
||||
if not ro_list:
|
||||
print("No request_output found.")
|
||||
continue
|
||||
|
||||
for ro in ro_list:
|
||||
# Multimodal output may be attached to RequestOutput or CompletionOutput.
|
||||
mm = getattr(ro, "multimodal_output", None)
|
||||
if not mm and ro.outputs:
|
||||
mm = getattr(ro.outputs[0], "multimodal_output", None)
|
||||
|
||||
if mm:
|
||||
print(f"Multimodal output keys: {mm.keys()}")
|
||||
if "audio" in mm:
|
||||
audio_out = mm["audio"]
|
||||
print(f"Generated Audio Shape: {audio_out.shape}")
|
||||
out_path = f"output_{i}.wav"
|
||||
sf.write(out_path, audio_out.cpu().numpy().squeeze(), 22050)
|
||||
print(f"Saved audio to {out_path}")
|
||||
else:
|
||||
print("No multimodal output found.")
|
||||
except Exception as e:
|
||||
print(f"Error inspecting output: {e}")
|
||||
omni.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
run_e2e()
|
||||
@@ -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
|
||||
@@ -0,0 +1,75 @@
|
||||
# Fish Speech S2 Pro
|
||||
|
||||
This directory contains an offline demo for running Fish Speech S2 Pro (`fishaudio/s2-pro`) with vLLM Omni. It supports text-only synthesis and voice cloning with reference audio.
|
||||
|
||||
## Model Overview
|
||||
|
||||
[Fish Speech S2 Pro](https://huggingface.co/fishaudio/s2-pro) is a 4B dual-AR text-to-speech model by FishAudio. It uses:
|
||||
|
||||
- **Slow AR**: Generates semantic tokens autoregressively (Qwen3-based backbone)
|
||||
- **Fast AR**: Predicts residual codebook tokens from semantic tokens
|
||||
- **DAC Decoder**: Converts codec codes to 44.1 kHz audio waveform
|
||||
|
||||
The model produces high-quality speech with voice cloning capabilities.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
Install the `fish-speech` package for the DAC codec:
|
||||
|
||||
```bash
|
||||
pip install fish-speech
|
||||
```
|
||||
|
||||
## 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.
|
||||
|
||||
## Quick Start
|
||||
|
||||
### Text-Only Synthesis
|
||||
|
||||
```bash
|
||||
python end2end.py --text "Hello, this is a test of the Fish Speech text to speech system."
|
||||
```
|
||||
|
||||
Generated audio files are saved to `output_audio/` by default.
|
||||
|
||||
### Voice Cloning
|
||||
|
||||
Provide a reference audio file and its transcript to clone a voice:
|
||||
|
||||
```bash
|
||||
python end2end.py \
|
||||
--text "Hello, this is a cloned voice." \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--ref-text "Transcript of the reference audio."
|
||||
```
|
||||
|
||||
## Streaming Mode
|
||||
|
||||
Add `--streaming` to stream audio chunks progressively via `AsyncOmni` (requires `async_chunk: true` in the stage config):
|
||||
|
||||
```bash
|
||||
python end2end.py --text "Hello, this is a test." --streaming
|
||||
```
|
||||
|
||||
Each DAC decoder chunk is logged as it arrives. The final WAV file is written once generation completes.
|
||||
|
||||
## Configuration
|
||||
|
||||
| Argument | Default | Description |
|
||||
|----------|---------|-------------|
|
||||
| `--text` | `"Hello, this is a test..."` | Text to synthesize |
|
||||
| `--ref-audio` | None | Path to reference audio for voice cloning |
|
||||
| `--ref-text` | None | Transcript of the reference audio |
|
||||
| `--model` | `fishaudio/s2-pro` | HuggingFace model path |
|
||||
| `--stage-configs-path` | `fish_speech_s2_pro.yaml` | Path to stage configs YAML |
|
||||
| `--output-dir` | `output_audio` | Output directory for WAV files |
|
||||
| `--streaming` | False | Enable streaming via AsyncOmni |
|
||||
| `--stage-init-timeout` | 600 | Stage initialization timeout (seconds) |
|
||||
|
||||
## Notes
|
||||
|
||||
- Output audio is 44.1 kHz mono WAV.
|
||||
- The DAC codec weights (`codec.pth`) are loaded lazily from the model directory.
|
||||
- Voice cloning encodes the reference audio using the DAC codec to extract semantic codes, then prepends them as a system message.
|
||||
@@ -0,0 +1,252 @@
|
||||
"""Offline inference demo for Fish Speech S2 Pro via vLLM Omni.
|
||||
|
||||
Generates speech from text using the fishaudio/s2-pro model.
|
||||
Supports both text-only synthesis and voice cloning with reference audio.
|
||||
|
||||
Usage:
|
||||
# Text-only synthesis
|
||||
python end2end.py --text "Hello, this is a test."
|
||||
|
||||
# Voice cloning with reference audio
|
||||
python end2end.py --text "Hello, this is a test." --ref-audio ref.wav --ref-text "Reference text."
|
||||
|
||||
# Streaming mode
|
||||
python end2end.py --text "Hello, this is a test." --streaming
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
|
||||
import soundfile as sf
|
||||
import torch
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
from vllm_omni import AsyncOmni, Omni
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
DEFAULT_MODEL = "fishaudio/s2-pro"
|
||||
DEFAULT_STAGE_CONFIG = os.path.join(
|
||||
os.path.dirname(__file__),
|
||||
"..",
|
||||
"..",
|
||||
"..",
|
||||
"vllm_omni",
|
||||
"model_executor",
|
||||
"stage_configs",
|
||||
"fish_speech_s2_pro.yaml",
|
||||
)
|
||||
|
||||
|
||||
def build_prompt(
|
||||
text: str,
|
||||
ref_audio_path: str | None = None,
|
||||
ref_text: str | None = None,
|
||||
model_name: str = DEFAULT_MODEL,
|
||||
) -> dict:
|
||||
"""Build a prompt for Fish Speech S2 Pro.
|
||||
|
||||
Uses Qwen3 chat template format with <|voice|> modality marker:
|
||||
<|im_start|>user\n<|speaker:0|>text<|im_end|>\n<|im_start|>assistant\n<|voice|>
|
||||
|
||||
The model then generates semantic tokens autoregressively.
|
||||
"""
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
|
||||
|
||||
# Build user message with speaker tag prefix.
|
||||
user_text = f"<|speaker:0|>{text}"
|
||||
|
||||
# Use chat template for user turn + assistant generation prompt.
|
||||
messages = [{"role": "user", "content": user_text}]
|
||||
prompt_ids = tokenizer.apply_chat_template(
|
||||
messages,
|
||||
tokenize=True,
|
||||
add_generation_prompt=True,
|
||||
)
|
||||
|
||||
# Append <|voice|> token (151673) to signal voice generation mode.
|
||||
voice_token_id = tokenizer.encode("<|voice|>", add_special_tokens=False)
|
||||
prompt_ids = prompt_ids + voice_token_id
|
||||
|
||||
additional_information = {
|
||||
"text": [text],
|
||||
"max_new_tokens": [4096],
|
||||
}
|
||||
|
||||
if ref_audio_path:
|
||||
additional_information["ref_audio"] = [ref_audio_path]
|
||||
if ref_text:
|
||||
additional_information["ref_text"] = [ref_text]
|
||||
|
||||
return {
|
||||
"prompt_token_ids": prompt_ids,
|
||||
"additional_information": additional_information,
|
||||
}
|
||||
|
||||
|
||||
def _save_wav(output_dir: str, request_id: str, mm: dict) -> None:
|
||||
"""Concatenate audio chunks and write to a wav file."""
|
||||
audio_data = mm["audio"]
|
||||
sr_raw = mm["sr"]
|
||||
sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw
|
||||
sr = sr_val.item() if hasattr(sr_val, "item") else int(sr_val)
|
||||
audio_tensor = torch.cat(audio_data, dim=-1) if isinstance(audio_data, list) else audio_data
|
||||
out_wav = os.path.join(output_dir, f"output_{request_id}.wav")
|
||||
sf.write(out_wav, audio_tensor.float().cpu().numpy().flatten(), samplerate=sr, format="WAV")
|
||||
logger.info("Request %s: saved audio to %s (sr=%d)", request_id, out_wav, sr)
|
||||
|
||||
|
||||
def main(args):
|
||||
"""Run offline inference with Omni."""
|
||||
output_dir = args.output_dir
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
stage_configs_path = args.stage_configs_path or DEFAULT_STAGE_CONFIG
|
||||
model_name = args.model or DEFAULT_MODEL
|
||||
|
||||
inputs = [
|
||||
build_prompt(
|
||||
text=args.text,
|
||||
ref_audio_path=args.ref_audio,
|
||||
ref_text=args.ref_text,
|
||||
model_name=model_name,
|
||||
)
|
||||
]
|
||||
|
||||
omni = Omni(
|
||||
model=model_name,
|
||||
stage_configs_path=stage_configs_path,
|
||||
log_stats=args.log_stats,
|
||||
stage_init_timeout=args.stage_init_timeout,
|
||||
)
|
||||
|
||||
t_start = time.perf_counter()
|
||||
for stage_outputs in omni.generate(inputs):
|
||||
for output in stage_outputs.request_output:
|
||||
_save_wav(output_dir, output.request_id, output.outputs[0].multimodal_output)
|
||||
t_end = time.perf_counter()
|
||||
logger.info("Total inference time: %.1f ms", (t_end - t_start) * 1000)
|
||||
|
||||
|
||||
async def main_streaming(args):
|
||||
"""Run offline inference with AsyncOmni for streaming audio chunks."""
|
||||
output_dir = args.output_dir
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
stage_configs_path = args.stage_configs_path or DEFAULT_STAGE_CONFIG
|
||||
model_name = args.model or DEFAULT_MODEL
|
||||
|
||||
prompt = build_prompt(
|
||||
text=args.text,
|
||||
ref_audio_path=args.ref_audio,
|
||||
ref_text=args.ref_text,
|
||||
model_name=model_name,
|
||||
)
|
||||
|
||||
omni = AsyncOmni(
|
||||
model=model_name,
|
||||
stage_configs_path=stage_configs_path,
|
||||
log_stats=args.log_stats,
|
||||
stage_init_timeout=args.stage_init_timeout,
|
||||
)
|
||||
|
||||
request_id = "0"
|
||||
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)
|
||||
dt_ms = (t_now - t_prev) * 1000
|
||||
ttfa_ms = (t_now - t_start) * 1000
|
||||
if chunk_idx == 0:
|
||||
logger.info("Request %s: chunk %d samples=%d TTFA=%.1fms", request_id, chunk_idx, n, ttfa_ms)
|
||||
else:
|
||||
logger.info("Request %s: chunk %d samples=%d inter_chunk=%.1fms", request_id, chunk_idx, n, dt_ms)
|
||||
t_prev = t_now
|
||||
chunk_idx += 1
|
||||
else:
|
||||
t_end = time.perf_counter()
|
||||
total_ms = (t_end - t_start) * 1000
|
||||
logger.info("Request %s: done total=%.1fms chunks=%d", request_id, total_ms, chunk_idx)
|
||||
_save_wav(output_dir, request_id, mm)
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = FlexibleArgumentParser(
|
||||
description="Fish Speech S2 Pro offline inference with vLLM Omni",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--text",
|
||||
type=str,
|
||||
default="Hello, this is a test of the Fish Speech text to speech system.",
|
||||
help="Text to synthesize.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ref-audio",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to reference audio for voice cloning.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ref-text",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Reference text corresponding to the reference audio.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
default=None,
|
||||
help=f"HuggingFace model path (default: {DEFAULT_MODEL}).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stage-configs-path",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to stage configs YAML.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-dir",
|
||||
default="output_audio",
|
||||
help="Output directory for generated wav files.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--log-stats",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Enable writing detailed statistics.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stage-init-timeout",
|
||||
type=int,
|
||||
default=600,
|
||||
help="Timeout for initializing stages in seconds.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--streaming",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Stream audio chunks as they arrive via AsyncOmni.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
args = parse_args()
|
||||
if args.streaming:
|
||||
asyncio.run(main_streaming(args))
|
||||
else:
|
||||
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")
|
||||
|
||||
|
||||
@@ -22,7 +22,7 @@ Usage:
|
||||
python image_to_video.py --model /path/to/LTX-2 \
|
||||
--model-class-name LTX2ImageToVideoPipeline \
|
||||
--image input.jpg --prompt "A cinematic dolly shot of a boat" \
|
||||
--num-frames 121 --num_inference_steps 40 --guidance_scale 4.0 \
|
||||
--num-frames 121 --num-inference-steps 40 --guidance-scale 4.0 \
|
||||
--frame-rate 24 --fps 24 --output ltx2_i2v.mp4
|
||||
"""
|
||||
|
||||
|
||||
@@ -96,7 +96,8 @@ python end2end.py --query-type CustomVoice --streaming --output-dir /tmp/out_str
|
||||
```
|
||||
|
||||
Each Code2Wav chunk is logged as it arrives (default 25 frames; configurable via `codec_chunk_frames`
|
||||
and `initial_codec_chunk_frames` in the stage config). The final WAV file is written once generation
|
||||
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.
|
||||
|
||||
@@ -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}")
|
||||
|
||||
@@ -127,6 +127,11 @@ def parse_args() -> argparse.Namespace:
|
||||
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()
|
||||
|
||||
|
||||
@@ -160,6 +165,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
|
||||
@@ -193,7 +199,7 @@ def main():
|
||||
print(
|
||||
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}"
|
||||
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")
|
||||
|
||||
@@ -0,0 +1,103 @@
|
||||
import argparse
|
||||
import fcntl
|
||||
import os
|
||||
import site
|
||||
import subprocess
|
||||
import tempfile
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
DEPENDENCY_REPO = "https://github.com/bytedance/DreamID-V.git"
|
||||
DEPENDENCY_BRANCH = "omni"
|
||||
CACHE_DIR = Path(tempfile.gettempdir()) / "vllm-omni-dependency"
|
||||
LOCK_FILE = CACHE_DIR / ".install.lock"
|
||||
DEPENDENCY_DIR = CACHE_DIR / "DreamID-Omni"
|
||||
|
||||
|
||||
def download_dependency():
|
||||
CACHE_DIR.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
with open(LOCK_FILE, "w") as f:
|
||||
fcntl.flock(f, fcntl.LOCK_EX)
|
||||
if not DEPENDENCY_DIR.exists():
|
||||
print(f"Downloading DreamID-Omni to {DEPENDENCY_DIR} ...")
|
||||
subprocess.run(
|
||||
["git", "clone", "--depth", "1", DEPENDENCY_REPO, "--branch", DEPENDENCY_BRANCH, str(DEPENDENCY_DIR)],
|
||||
check=True,
|
||||
)
|
||||
print("Download finished.")
|
||||
fcntl.flock(f, fcntl.LOCK_UN)
|
||||
|
||||
# write .pth to site-packages
|
||||
site_packages = Path(site.getsitepackages()[0])
|
||||
pth_file = site_packages / "vllm_omni_dependency.pth"
|
||||
pth_file.write_text(str(DEPENDENCY_DIR))
|
||||
print(f"Added {DEPENDENCY_DIR} to site-packages via {pth_file}")
|
||||
|
||||
|
||||
def timed_download(repo_id: str, local_dir: str, allow_patterns: list | None = None):
|
||||
"""Download files from HF repo and log time + destination."""
|
||||
if os.path.exists(local_dir):
|
||||
print(f"Directory {local_dir} already exists. Skipping download.")
|
||||
return
|
||||
print(f"Starting download from {repo_id} into {local_dir}")
|
||||
start_time = time.time()
|
||||
|
||||
snapshot_download(
|
||||
repo_id=repo_id,
|
||||
local_dir=local_dir,
|
||||
local_dir_use_symlinks=False,
|
||||
allow_patterns=allow_patterns,
|
||||
)
|
||||
|
||||
elapsed = time.time() - start_time
|
||||
print(f"✅ Finished downloading {repo_id} in {elapsed:.2f} seconds. Files saved at: {local_dir}")
|
||||
|
||||
|
||||
def main(output_dir: str):
|
||||
# Wan2.2
|
||||
wan_dir = os.path.join(output_dir, "Wan2.2-TI2V-5B")
|
||||
timed_download(
|
||||
repo_id="Wan-AI/Wan2.2-TI2V-5B",
|
||||
local_dir=wan_dir,
|
||||
allow_patterns=["google/*", "models_t5_umt5-xxl-enc-bf16.pth", "Wan2.2_VAE.pth"],
|
||||
)
|
||||
|
||||
# MMAudio
|
||||
mm_audio_dir = os.path.join(output_dir, "MMAudio")
|
||||
timed_download(
|
||||
repo_id="hkchengrex/MMAudio",
|
||||
local_dir=mm_audio_dir,
|
||||
allow_patterns=["ext_weights/best_netG.pt", "ext_weights/v1-16.pth"],
|
||||
)
|
||||
|
||||
dreamid_dir = os.path.join(output_dir, "DreamID-Omni")
|
||||
|
||||
timed_download(repo_id="XuGuo699/DreamID-Omni", local_dir=dreamid_dir)
|
||||
|
||||
# Now we construct the config file
|
||||
import json
|
||||
|
||||
data = {
|
||||
"_class_name": "DreamIDOmniPipeline",
|
||||
"fusion": "DreamID-Omni/dreamid_omni.safetensors",
|
||||
}
|
||||
|
||||
with open(os.path.join(output_dir, "model_index.json"), "w", encoding="utf-8") as f:
|
||||
json.dump(data, f, indent=2)
|
||||
|
||||
print(f"model_index.json created at {os.path.join(output_dir, 'model_index.json')}")
|
||||
|
||||
# now we download the dependency code
|
||||
download_dependency()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="Download models from Hugging Face")
|
||||
parser.add_argument(
|
||||
"--output-dir", type=str, default="./dreamid_omni", help="Base directory to save downloaded models"
|
||||
)
|
||||
args = parser.parse_args()
|
||||
main(args.output_dir)
|
||||
@@ -0,0 +1,58 @@
|
||||
# X-To-Video-Audio
|
||||
|
||||
The `DreamID-Omni` pipeline generates short videos from text, image and video.
|
||||
|
||||
## Local CLI Usage
|
||||
### Download the Model locally
|
||||
Since DreamID-Omni combine multiple models, and without any config, so we need to download them locally.
|
||||
|
||||
```bash
|
||||
python download_dreamid_omni.py --output-dir ./dreamid_omni
|
||||
```
|
||||
After download, the model directory will look like this:
|
||||
|
||||
```
|
||||
dreamid_omni/
|
||||
├── DreamID-Omni/
|
||||
│ ├── dreamid_omni.safetensors
|
||||
├── MMAudio/
|
||||
│ ├── ext_weights/
|
||||
│ │ ├── best_netG.pt
|
||||
│ │ ├── v1-16.pth
|
||||
├── Wan2.2-TI2V-5B/
|
||||
│ ├── google/*
|
||||
│ ├── models_t5_umt5-xxl-enc-bf16.pth
|
||||
│ ├── Wan2.2_VAE.pth
|
||||
│
|
||||
├── model_index.json # create by download_dreamid_omni.py
|
||||
```
|
||||
|
||||
### Run the Inference
|
||||
```
|
||||
python x_to_video_audio.py \
|
||||
--model /xx/dreamid_omni \
|
||||
--prompt "Two people walking together and singing happily" \
|
||||
--image-path ./example0.png ./example1.png \
|
||||
--audio-path ./example0.wav ./example1.wav \
|
||||
--video-negative-prompt "jitter, bad hands, blur, distortion" \
|
||||
--audio-negative-prompt "robotic, muffled, echo, distorted" \
|
||||
--cfg-parallel-size 2 \
|
||||
--num-inference-steps 45 \
|
||||
--height 704 \
|
||||
--width 1280 \
|
||||
--output dreamid_omni.mp4
|
||||
```
|
||||
In the current test scenario (2 images + 2 audio inputs), the VRAM requirement is 72GB, regardless of whether cfg-parallel is enabled or disabled.
|
||||
The VRAM usage can be reduced by enabling CPU offload via --enable-cpu-offload.
|
||||
|
||||
Key arguments:
|
||||
- `--prompt`: text description (string).
|
||||
- `--model`: path to the model local directory.
|
||||
- `--height/--width`: output resolution (defaults 704 * 1024).
|
||||
- `--image-path`: path to the input image list.
|
||||
- `--audio-path`: path to the input audio list, indicate the timbre of the output video.
|
||||
- `--cfg-parallel-size`: number of parallel cfg parallel (defaults 1).
|
||||
- `--num-inference-steps`: number of denoising steps (defaults 45).
|
||||
- `--video-negative-prompt`: negative prompt for video generation.
|
||||
- `--audio-negative-prompt`: negative prompt for audio generation.
|
||||
- `--enable-cpu-offload`: enable CPU offload (defaults False).
|
||||
@@ -0,0 +1,148 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import argparse
|
||||
import re
|
||||
import time
|
||||
|
||||
import librosa
|
||||
from PIL import Image
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".gif"}
|
||||
AUDIO_EXTENSIONS = {".wav", ".mp3", ".flac", ".m4a", ".aac", ".ogg"}
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Offline inference for DreamID-Omni (video + audio).")
|
||||
parser.add_argument("--model", required=True, help="DreamID ckpt root directory.")
|
||||
parser.add_argument("--model-type", default="dreamid-omni", help="Model type.")
|
||||
parser.add_argument("--prompt", default=None, help="Text prompt.")
|
||||
|
||||
parser.add_argument("--image-path", type=str, nargs="+", help="list of image-path")
|
||||
parser.add_argument("--audio-path", type=str, nargs="+", help="list of audio-path")
|
||||
parser.add_argument("--prompt-file", type=str, default=None, help="Text prompt in json format.")
|
||||
|
||||
parser.add_argument("--height", type=int, default=704, help="Video height.")
|
||||
parser.add_argument("--width", type=int, default=1280, help="Video width.")
|
||||
parser.add_argument("--num-inference-steps", type=int, default=45, help="Sampling steps.")
|
||||
parser.add_argument("--solver-name", default="unipc", help="Solver name: unipc|dpm++|euler.")
|
||||
parser.add_argument("--shift", type=float, default=5.0, help="Scheduler shift.")
|
||||
parser.add_argument("--seed", type=int, default=103, help="Random seed for reproducible generation.")
|
||||
parser.add_argument(
|
||||
"--cfg-parallel-size",
|
||||
type=int,
|
||||
default=1,
|
||||
choices=[1, 2],
|
||||
help="Number of GPUs used for classifier free guidance parallel size.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--video-negative-prompt",
|
||||
default="jitter, bad hands, blur, distortion",
|
||||
help="Negative prompt for video.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--audio-negative-prompt",
|
||||
default="robotic, muffled, echo, distorted",
|
||||
help="Negative prompt for audio.",
|
||||
)
|
||||
parser.add_argument("--output", default="dreamid_output.mp4", help="Output video path.")
|
||||
parser.add_argument(
|
||||
"--enable-cpu-offload",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Enable CPU offloading for diffusion models.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def load_image_and_audio(image_paths, audio_paths):
|
||||
image = []
|
||||
audio = []
|
||||
|
||||
for path in image_paths:
|
||||
with Image.open(path) as img:
|
||||
img = img.convert("RGB")
|
||||
image.append(img)
|
||||
|
||||
for path in audio_paths:
|
||||
audio_array, sr = librosa.load(path, sr=16000)
|
||||
audio_array = audio_array[int(sr * 1) : int(sr * 3)]
|
||||
audio.append(audio_array)
|
||||
return image, audio
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
if args.prompt is None and args.prompt_file is None:
|
||||
raise ValueError("Either --prompt or --prompt-file must be provided.")
|
||||
|
||||
text_prompt = args.prompt
|
||||
if args.prompt_file:
|
||||
import json
|
||||
|
||||
with open(args.prompt_file) as f:
|
||||
text_prompt = json.load(f)
|
||||
text_prompt = re.sub(
|
||||
r"\[SPEAKER_TIMESTAMPS_START\].*?\[SPEAKER_TIMESTAMPS_END\]", "", text_prompt, flags=re.DOTALL
|
||||
).strip()
|
||||
text_prompt = re.sub(
|
||||
r"\[AUDIO_DESCRIPTION_START].*?\[AUDIO_DESCRIPTION_END]", "", text_prompt, flags=re.DOTALL
|
||||
).strip()
|
||||
text_prompt = re.sub(r"\[[A-Z_]+\]", "", text_prompt)
|
||||
text_prompt = re.sub(r"\n\s*\n", "\n", text_prompt).strip()
|
||||
|
||||
image, audio = load_image_and_audio(args.image_path, args.audio_path)
|
||||
|
||||
prompt = {
|
||||
"prompt": text_prompt,
|
||||
"video_negative_prompt": args.video_negative_prompt,
|
||||
"audio_negative_prompt": args.audio_negative_prompt,
|
||||
"multi_modal_data": {"image": image, "audio": audio},
|
||||
}
|
||||
|
||||
sampling_params = OmniDiffusionSamplingParams(
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
seed=args.seed,
|
||||
extra_args={
|
||||
"solver_name": args.solver_name,
|
||||
"shift": args.shift,
|
||||
},
|
||||
)
|
||||
|
||||
parallel_config = DiffusionParallelConfig(
|
||||
cfg_parallel_size=args.cfg_parallel_size,
|
||||
)
|
||||
|
||||
omni = Omni(
|
||||
model=args.model,
|
||||
parallel_config=parallel_config,
|
||||
model_type=args.model_type,
|
||||
enable_cpu_offload=args.enable_cpu_offload,
|
||||
)
|
||||
start = time.perf_counter()
|
||||
outputs = omni.generate(prompt, sampling_params)
|
||||
elapsed = time.perf_counter() - start
|
||||
|
||||
if not outputs:
|
||||
raise RuntimeError("No output returned from DreamID-Omni.")
|
||||
output = outputs[0].request_output
|
||||
generated_video = output[0].images[0][0]
|
||||
generated_audio = output[0].images[0][1]
|
||||
try:
|
||||
from dreamid_omni.utils.io_utils import save_video
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to extract video and audio from DreamID-Omni output. Error: {e}")
|
||||
output_path = args.output
|
||||
save_video(output_path, generated_video, generated_audio, fps=24, sample_rate=16000)
|
||||
print(f"Saved generated video to {output_path}")
|
||||
print(f"Total time: {elapsed:.2f}s")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -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 }}
|
||||
@@ -0,0 +1,16 @@
|
||||
{{- if .Values.modelCache.enabled }}
|
||||
apiVersion: v1
|
||||
kind: PersistentVolumeClaim
|
||||
metadata:
|
||||
name: "{{ .Release.Name }}-model-cache"
|
||||
namespace: {{ .Release.Namespace }}
|
||||
spec:
|
||||
accessModes:
|
||||
{{- toYaml .Values.modelCache.accessModes | nindent 4 }}
|
||||
{{- if .Values.modelCache.storageClassName }}
|
||||
storageClassName: {{ .Values.modelCache.storageClassName }}
|
||||
{{- end }}
|
||||
resources:
|
||||
requests:
|
||||
storage: {{ .Values.modelCache.storageSize }}
|
||||
{{- end }}
|
||||
@@ -0,0 +1,23 @@
|
||||
{{- if .Values.secrets }}
|
||||
apiVersion: v1
|
||||
kind: Secret
|
||||
metadata:
|
||||
name: "{{ .Release.Name }}-secrets"
|
||||
namespace: {{ .Release.Namespace }}
|
||||
type: Opaque
|
||||
data:
|
||||
{{- range $key, $val := .Values.secrets }}
|
||||
{{ $key }}: {{ $val | b64enc | quote }}
|
||||
{{- end }}
|
||||
---
|
||||
{{- end }}
|
||||
{{- if .Values.hfToken }}
|
||||
apiVersion: v1
|
||||
kind: Secret
|
||||
metadata:
|
||||
name: "{{ .Release.Name }}-hf-token"
|
||||
namespace: {{ .Release.Namespace }}
|
||||
type: Opaque
|
||||
data:
|
||||
token: {{ .Values.hfToken | b64enc | quote }}
|
||||
{{- end }}
|
||||
@@ -0,0 +1,14 @@
|
||||
apiVersion: v1
|
||||
kind: Service
|
||||
metadata:
|
||||
name: "{{ include "chart.service-name" . }}"
|
||||
namespace: {{ .Release.Namespace }}
|
||||
spec:
|
||||
type: ClusterIP
|
||||
ports:
|
||||
- name: {{ include "chart.service-port-name" . }}
|
||||
port: {{ include "chart.service-port" . }}
|
||||
targetPort: {{ include "chart.container-port-name" . }}
|
||||
protocol: TCP
|
||||
selector:
|
||||
{{- include "chart.labels" . | nindent 4 }}
|
||||
@@ -0,0 +1,181 @@
|
||||
suite: test deployment
|
||||
templates:
|
||||
- deployment.yaml
|
||||
tests:
|
||||
- it: should create deployment with default omni command
|
||||
asserts:
|
||||
- hasDocuments:
|
||||
count: 1
|
||||
- isKind:
|
||||
of: Deployment
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].name
|
||||
value: vllm-omni
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].command[0]
|
||||
value: vllm
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].command[1]
|
||||
value: serve
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].command[2]
|
||||
value: Tongyi-MAI/Z-Image-Turbo
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].command[3]
|
||||
value: "--omni"
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].command[4]
|
||||
value: "--host"
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].command[5]
|
||||
value: "0.0.0.0"
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].command[6]
|
||||
value: "--port"
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].command[7]
|
||||
value: "8000"
|
||||
|
||||
- it: should use custom model when set
|
||||
set:
|
||||
model: "Qwen/Qwen2.5-Omni-7B"
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].command[2]
|
||||
value: Qwen/Qwen2.5-Omni-7B
|
||||
|
||||
- it: should include omniArgs flags in command
|
||||
set:
|
||||
omniArgs:
|
||||
vaeUseSlicing: true
|
||||
vaeUseTiling: true
|
||||
enableCpuOffload: true
|
||||
numGpus: 2
|
||||
cacheBackend: "tea_cache"
|
||||
workerBackend: "ray"
|
||||
extraArgs:
|
||||
- "--enable-layerwise-offload"
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].command
|
||||
content: "--vae-use-slicing"
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].command
|
||||
content: "--vae-use-tiling"
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].command
|
||||
content: "--enable-cpu-offload"
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].command
|
||||
content: "--num-gpus"
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].command
|
||||
content: "--cache-backend"
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].command
|
||||
content: "--worker-backend"
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].command
|
||||
content: "--enable-layerwise-offload"
|
||||
|
||||
- it: should use full command override when image.command is set
|
||||
set:
|
||||
image:
|
||||
command:
|
||||
- "custom-binary"
|
||||
- "--flag"
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].command[0]
|
||||
value: custom-binary
|
||||
- equal:
|
||||
path: spec.template.spec.containers[0].command[1]
|
||||
value: "--flag"
|
||||
|
||||
- it: should mount shm and model-cache volumes
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].volumeMounts
|
||||
content:
|
||||
name: shm
|
||||
mountPath: /dev/shm
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].volumeMounts
|
||||
content:
|
||||
name: model-cache
|
||||
mountPath: /cache
|
||||
- contains:
|
||||
path: spec.template.spec.volumes
|
||||
content:
|
||||
name: shm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
sizeLimit: 8Gi
|
||||
|
||||
- it: should use PVC for model-cache when modelCache is enabled
|
||||
set:
|
||||
modelCache:
|
||||
enabled: true
|
||||
storageSize: "50Gi"
|
||||
accessModes:
|
||||
- ReadWriteOnce
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.volumes
|
||||
content:
|
||||
name: model-cache
|
||||
persistentVolumeClaim:
|
||||
claimName: RELEASE-NAME-model-cache
|
||||
|
||||
- it: should use emptyDir for model-cache when modelCache is disabled
|
||||
set:
|
||||
modelCache:
|
||||
enabled: false
|
||||
storageSize: "50Gi"
|
||||
accessModes:
|
||||
- ReadWriteOnce
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.volumes
|
||||
content:
|
||||
name: model-cache
|
||||
emptyDir: {}
|
||||
|
||||
- it: should set HF_HOME environment variable
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: HF_HOME
|
||||
value: /cache/huggingface
|
||||
|
||||
- it: should set HF_TOKEN env when hfToken is provided
|
||||
set:
|
||||
hfToken: "hf_test_token_123"
|
||||
asserts:
|
||||
- contains:
|
||||
path: spec.template.spec.containers[0].env
|
||||
content:
|
||||
name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: RELEASE-NAME-hf-token
|
||||
key: token
|
||||
|
||||
- it: should create custom init containers when specified
|
||||
set:
|
||||
extraInit:
|
||||
initContainers:
|
||||
- name: my-init
|
||||
image: busybox:latest
|
||||
command: ["echo", "init"]
|
||||
asserts:
|
||||
- lengthEqual:
|
||||
path: spec.template.spec.initContainers
|
||||
count: 1
|
||||
- equal:
|
||||
path: spec.template.spec.initContainers[0].name
|
||||
value: my-init
|
||||
- equal:
|
||||
path: spec.template.spec.initContainers[0].image
|
||||
value: busybox:latest
|
||||
@@ -0,0 +1,86 @@
|
||||
suite: test ingress
|
||||
templates:
|
||||
- ingress.yaml
|
||||
tests:
|
||||
- it: should not create Ingress when disabled
|
||||
set:
|
||||
ingress:
|
||||
enabled: false
|
||||
asserts:
|
||||
- hasDocuments:
|
||||
count: 0
|
||||
|
||||
- it: should create Ingress when enabled
|
||||
set:
|
||||
ingress:
|
||||
enabled: true
|
||||
host: "vllm-omni.example.com"
|
||||
asserts:
|
||||
- hasDocuments:
|
||||
count: 1
|
||||
- isKind:
|
||||
of: Ingress
|
||||
- equal:
|
||||
path: spec.rules[0].host
|
||||
value: vllm-omni.example.com
|
||||
- equal:
|
||||
path: spec.rules[0].http.paths[0].path
|
||||
value: /
|
||||
- equal:
|
||||
path: spec.rules[0].http.paths[0].pathType
|
||||
value: Prefix
|
||||
- equal:
|
||||
path: spec.rules[0].http.paths[0].backend.service.name
|
||||
value: RELEASE-NAME-service
|
||||
- equal:
|
||||
path: spec.rules[0].http.paths[0].backend.service.port.number
|
||||
value: 80
|
||||
|
||||
- it: should set ingressClassName when specified
|
||||
set:
|
||||
ingress:
|
||||
enabled: true
|
||||
host: "vllm-omni.example.com"
|
||||
ingressClassName: "nginx"
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.ingressClassName
|
||||
value: nginx
|
||||
|
||||
- it: should not set ingressClassName when empty
|
||||
set:
|
||||
ingress:
|
||||
enabled: true
|
||||
host: "vllm-omni.example.com"
|
||||
ingressClassName: ""
|
||||
asserts:
|
||||
- notExists:
|
||||
path: spec.ingressClassName
|
||||
|
||||
- it: should include annotations when specified
|
||||
set:
|
||||
ingress:
|
||||
enabled: true
|
||||
host: "vllm-omni.example.com"
|
||||
annotations:
|
||||
nginx.ingress.kubernetes.io/proxy-body-size: "100m"
|
||||
asserts:
|
||||
- isNotEmpty:
|
||||
path: metadata.annotations
|
||||
|
||||
- it: should include TLS configuration when specified
|
||||
set:
|
||||
ingress:
|
||||
enabled: true
|
||||
host: "vllm-omni.example.com"
|
||||
tls:
|
||||
- secretName: vllm-omni-tls
|
||||
hosts:
|
||||
- vllm-omni.example.com
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.tls[0].secretName
|
||||
value: vllm-omni-tls
|
||||
- equal:
|
||||
path: spec.tls[0].hosts[0]
|
||||
value: vllm-omni.example.com
|
||||
@@ -0,0 +1,73 @@
|
||||
suite: test pvc
|
||||
templates:
|
||||
- pvc.yaml
|
||||
tests:
|
||||
- it: should create PVC when modelCache is enabled
|
||||
set:
|
||||
modelCache:
|
||||
enabled: true
|
||||
storageSize: "50Gi"
|
||||
storageClassName: ""
|
||||
accessModes:
|
||||
- ReadWriteOnce
|
||||
asserts:
|
||||
- hasDocuments:
|
||||
count: 1
|
||||
- isKind:
|
||||
of: PersistentVolumeClaim
|
||||
- equal:
|
||||
path: spec.accessModes[0]
|
||||
value: ReadWriteOnce
|
||||
- equal:
|
||||
path: spec.resources.requests.storage
|
||||
value: 50Gi
|
||||
|
||||
- it: should not create PVC when modelCache is disabled
|
||||
set:
|
||||
modelCache:
|
||||
enabled: false
|
||||
storageSize: "50Gi"
|
||||
accessModes:
|
||||
- ReadWriteOnce
|
||||
asserts:
|
||||
- hasDocuments:
|
||||
count: 0
|
||||
|
||||
- it: should use custom storage size
|
||||
set:
|
||||
modelCache:
|
||||
enabled: true
|
||||
storageSize: "100Gi"
|
||||
storageClassName: ""
|
||||
accessModes:
|
||||
- ReadWriteOnce
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.resources.requests.storage
|
||||
value: 100Gi
|
||||
|
||||
- it: should set storageClassName when specified
|
||||
set:
|
||||
modelCache:
|
||||
enabled: true
|
||||
storageSize: "50Gi"
|
||||
storageClassName: "fast-ssd"
|
||||
accessModes:
|
||||
- ReadWriteOnce
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.storageClassName
|
||||
value: fast-ssd
|
||||
|
||||
- it: should support ReadWriteMany access mode
|
||||
set:
|
||||
modelCache:
|
||||
enabled: true
|
||||
storageSize: "50Gi"
|
||||
storageClassName: ""
|
||||
accessModes:
|
||||
- ReadWriteMany
|
||||
asserts:
|
||||
- equal:
|
||||
path: spec.accessModes[0]
|
||||
value: ReadWriteMany
|
||||
@@ -0,0 +1,48 @@
|
||||
suite: test secrets
|
||||
templates:
|
||||
- secrets.yaml
|
||||
tests:
|
||||
- it: should create HF token secret when hfToken is provided
|
||||
set:
|
||||
hfToken: "hf_test_token_123"
|
||||
asserts:
|
||||
- hasDocuments:
|
||||
count: 1
|
||||
- isKind:
|
||||
of: Secret
|
||||
- equal:
|
||||
path: metadata.name
|
||||
value: RELEASE-NAME-hf-token
|
||||
- exists:
|
||||
path: data.token
|
||||
|
||||
- it: should not create any secrets when hfToken is empty and no secrets defined
|
||||
set:
|
||||
hfToken: ""
|
||||
secrets: {}
|
||||
asserts:
|
||||
- hasDocuments:
|
||||
count: 0
|
||||
|
||||
- it: should create both generic and HF token secrets
|
||||
set:
|
||||
hfToken: "hf_test_token_123"
|
||||
secrets:
|
||||
mykey: myvalue
|
||||
asserts:
|
||||
- hasDocuments:
|
||||
count: 2
|
||||
|
||||
- it: should create only generic secrets when no hfToken
|
||||
set:
|
||||
hfToken: ""
|
||||
secrets:
|
||||
mykey: myvalue
|
||||
asserts:
|
||||
- hasDocuments:
|
||||
count: 1
|
||||
- isKind:
|
||||
of: Secret
|
||||
- equal:
|
||||
path: metadata.name
|
||||
value: RELEASE-NAME-secrets
|
||||
@@ -0,0 +1,197 @@
|
||||
# -- Default values for chart vllm-omni
|
||||
# -- Declare variables to be passed into your templates.
|
||||
|
||||
# -- HuggingFace model ID to serve
|
||||
model: "Tongyi-MAI/Z-Image-Turbo"
|
||||
|
||||
# -- Image configuration
|
||||
image:
|
||||
# -- Image repository
|
||||
repository: "vllm/vllm-omni"
|
||||
# -- Image tag
|
||||
tag: "v0.16.0"
|
||||
# -- Override the container command entirely. If empty, the chart constructs
|
||||
# -- it automatically from model + omniArgs.
|
||||
command: []
|
||||
# -- Optional environment variables for the container
|
||||
env: []
|
||||
# -- Security context override
|
||||
securityContext: {}
|
||||
|
||||
# -- Container port
|
||||
containerPort: 8000
|
||||
# -- Service name (auto-generated from release name if empty)
|
||||
serviceName:
|
||||
# -- Service port
|
||||
servicePort: 80
|
||||
# -- Additional ports configuration
|
||||
extraPorts: []
|
||||
|
||||
# -- Number of replicas
|
||||
replicaCount: 1
|
||||
|
||||
# -- Deployment strategy configuration
|
||||
deploymentStrategy: {}
|
||||
|
||||
# -- Resource configuration
|
||||
resources:
|
||||
requests:
|
||||
# -- Number of CPUs
|
||||
cpu: 4
|
||||
# -- CPU memory
|
||||
memory: 24Gi
|
||||
# -- Number of GPUs
|
||||
nvidia.com/gpu: 1
|
||||
limits:
|
||||
# -- Number of CPUs
|
||||
cpu: 4
|
||||
# -- CPU memory
|
||||
memory: 24Gi
|
||||
# -- Number of GPUs
|
||||
nvidia.com/gpu: 1
|
||||
|
||||
# -- GPU model types for node affinity scheduling (optional)
|
||||
# -- Leave empty to schedule on any GPU node. Set to restrict to specific GPU types.
|
||||
# -- Example: ["NVIDIA-A100-SXM4-40GB"]
|
||||
gpuModels: []
|
||||
|
||||
# -- vllm-omni specific CLI arguments
|
||||
omniArgs:
|
||||
# -- Additional raw CLI flags appended to the serve command (list of strings)
|
||||
extraArgs: []
|
||||
# -- Enable VAE slicing for memory optimization
|
||||
vaeUseSlicing: false
|
||||
# -- Enable VAE tiling for memory optimization
|
||||
vaeUseTiling: false
|
||||
# -- Enable CPU offloading for diffusion models
|
||||
enableCpuOffload: false
|
||||
# -- Number of GPUs for diffusion inference (empty = auto)
|
||||
numGpus:
|
||||
# -- Path to stage configs file (empty = auto-detected from model)
|
||||
stageConfigsPath:
|
||||
# -- Cache backend: none, tea_cache, or cache_dit
|
||||
cacheBackend: "none"
|
||||
# -- Default sampling params as a JSON string
|
||||
defaultSamplingParams:
|
||||
# -- Worker backend: multi_process or ray
|
||||
workerBackend: "multi_process"
|
||||
|
||||
# -- HuggingFace token for gated models (optional)
|
||||
hfToken: ""
|
||||
|
||||
# -- Shared memory size for PyTorch multiprocessing
|
||||
shmSize: "8Gi"
|
||||
|
||||
# -- Model cache configuration (HuggingFace downloads)
|
||||
modelCache:
|
||||
# -- Use a PersistentVolumeClaim for model cache (recommended for production)
|
||||
enabled: true
|
||||
# -- Storage size for the PVC
|
||||
storageSize: "50Gi"
|
||||
# -- Storage class name (empty = cluster default)
|
||||
storageClassName: ""
|
||||
# -- Access modes
|
||||
accessModes:
|
||||
- ReadWriteOnce
|
||||
|
||||
# -- Autoscaling configuration
|
||||
autoscaling:
|
||||
# -- Enable autoscaling
|
||||
enabled: false
|
||||
# -- Minimum replicas
|
||||
minReplicas: 1
|
||||
# -- Maximum replicas
|
||||
maxReplicas: 10
|
||||
# -- Target CPU utilization for autoscaling
|
||||
targetCPUUtilizationPercentage: 80
|
||||
# targetMemoryUtilizationPercentage: 80
|
||||
|
||||
# -- ConfigMap data (key-value pairs injected as environment variables)
|
||||
configs: {}
|
||||
|
||||
# -- Secrets data (key-value pairs, base64-encoded automatically)
|
||||
secrets: {}
|
||||
|
||||
# -- External ConfigMaps/Secrets references
|
||||
externalConfigs: []
|
||||
|
||||
# -- Custom Kubernetes objects
|
||||
customObjects: []
|
||||
|
||||
# -- PodDisruptionBudget max unavailable
|
||||
maxUnavailablePodDisruptionBudget: ""
|
||||
|
||||
# -- Additional init containers
|
||||
extraInit:
|
||||
initContainers: []
|
||||
|
||||
# -- Additional sidecar containers
|
||||
extraContainers: []
|
||||
|
||||
# -- Startup probe configuration
|
||||
# -- Protects slow-starting containers. Liveness and readiness probes
|
||||
# -- do not start until the startup probe succeeds.
|
||||
# -- failureThreshold * periodSeconds = max startup time (40 * 30 = 1200s = 20 min)
|
||||
startupProbe:
|
||||
# -- HTTP check configuration
|
||||
httpGet:
|
||||
# -- Path to check
|
||||
path: /health
|
||||
# -- Port to check (uses named port to stay in sync with containerPort)
|
||||
port: "container-port"
|
||||
# -- Failures before killing the container
|
||||
failureThreshold: 40
|
||||
# -- How often to perform the probe
|
||||
periodSeconds: 30
|
||||
|
||||
# -- Readiness probe configuration
|
||||
readinessProbe:
|
||||
# -- How often to perform the probe
|
||||
periodSeconds: 10
|
||||
# -- Failures before marking unready
|
||||
failureThreshold: 3
|
||||
# -- HTTP check configuration
|
||||
httpGet:
|
||||
# -- Path to check
|
||||
path: /health
|
||||
# -- Port to check (uses named port to stay in sync with containerPort)
|
||||
port: "container-port"
|
||||
|
||||
# -- Liveness probe configuration
|
||||
livenessProbe:
|
||||
# -- Failures before restarting
|
||||
failureThreshold: 3
|
||||
# -- How often to perform the probe
|
||||
periodSeconds: 15
|
||||
# -- HTTP check configuration
|
||||
httpGet:
|
||||
# -- Path to check
|
||||
path: /health
|
||||
# -- Port to check (uses named port to stay in sync with containerPort)
|
||||
port: "container-port"
|
||||
|
||||
# -- Ingress configuration
|
||||
ingress:
|
||||
# -- Enable Ingress
|
||||
enabled: false
|
||||
# -- Ingress class name (e.g., nginx, traefik)
|
||||
ingressClassName: ""
|
||||
# -- Hostname for the Ingress rule
|
||||
host: "vllm-omni.example.com"
|
||||
# -- Additional annotations
|
||||
annotations: {}
|
||||
# -- TLS configuration
|
||||
tls: []
|
||||
# - secretName: vllm-omni-tls
|
||||
# hosts:
|
||||
# - vllm-omni.example.com
|
||||
|
||||
# -- Node selector for pod scheduling
|
||||
nodeSelector: {}
|
||||
|
||||
# -- Tolerations for pod scheduling
|
||||
tolerations: []
|
||||
|
||||
# -- Labels applied to all resources
|
||||
labels:
|
||||
app: "vllm-omni"
|
||||
@@ -0,0 +1,172 @@
|
||||
# Fish Speech S2 Pro
|
||||
|
||||
## Model
|
||||
|
||||
| Model | Description |
|
||||
|-------|-------------|
|
||||
| `fishaudio/s2-pro` | Fish Speech S2 Pro -- 4B dual-AR TTS model with DAC codec (44.1 kHz) |
|
||||
|
||||
## Gradio Demo
|
||||
|
||||
An interactive Gradio demo is available with text-to-speech synthesis, voice cloning, and streaming support.
|
||||
|
||||
```bash
|
||||
# Option 1: Launch server + Gradio together
|
||||
./run_gradio_demo.sh
|
||||
|
||||
# Option 2: If server is already running
|
||||
python gradio_demo.py --api-base http://localhost:8091
|
||||
```
|
||||
|
||||
Then open http://localhost:7860 in your browser.
|
||||
|
||||
Features:
|
||||
|
||||
- Text-to-speech synthesis
|
||||
- Voice cloning from uploaded audio or URL
|
||||
- Streaming mode (progressive PCM playback)
|
||||
|
||||
## Launch the Server
|
||||
|
||||
```bash
|
||||
vllm-omni serve fishaudio/s2-pro \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/fish_speech_s2_pro.yaml \
|
||||
--omni \
|
||||
--port 8091 \
|
||||
--trust-remote-code \
|
||||
--enforce-eager \
|
||||
--gpu-memory-utilization 0.9
|
||||
```
|
||||
|
||||
Or use the convenience script:
|
||||
|
||||
```bash
|
||||
./run_server.sh
|
||||
```
|
||||
|
||||
## Send TTS Request
|
||||
|
||||
### Using curl
|
||||
|
||||
```bash
|
||||
# Basic TTS
|
||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "Hello, how are you?",
|
||||
"voice": "default",
|
||||
"response_format": "wav"
|
||||
}' --output output.wav
|
||||
```
|
||||
|
||||
### Voice Cloning
|
||||
|
||||
Provide a reference audio (URL or base64 data URL) and its transcript:
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "Hello, this is a cloned voice.",
|
||||
"voice": "default",
|
||||
"ref_audio": "https://example.com/reference.wav",
|
||||
"ref_text": "Transcript of the reference audio."
|
||||
}' --output cloned.wav
|
||||
```
|
||||
|
||||
### Using Python
|
||||
|
||||
```python
|
||||
import httpx
|
||||
|
||||
# Basic TTS
|
||||
response = httpx.post(
|
||||
"http://localhost:8091/v1/audio/speech",
|
||||
json={
|
||||
"input": "Hello, how are you?",
|
||||
"voice": "default",
|
||||
"response_format": "wav",
|
||||
},
|
||||
timeout=300.0,
|
||||
)
|
||||
|
||||
with open("output.wav", "wb") as f:
|
||||
f.write(response.content)
|
||||
```
|
||||
|
||||
### Using the CLI Client
|
||||
|
||||
```bash
|
||||
cd examples/online_serving/fish_speech
|
||||
|
||||
# Basic TTS
|
||||
python speech_client.py --text "Hello, how are you?"
|
||||
|
||||
# Voice cloning
|
||||
python speech_client.py \
|
||||
--text "Hello, this is a cloned voice." \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--ref-text "Transcript of the reference audio."
|
||||
|
||||
# Streaming PCM output
|
||||
python speech_client.py --text "Hello world" --stream --output output.pcm
|
||||
```
|
||||
|
||||
The CLI client supports:
|
||||
|
||||
- `--api-base`: API base URL (default: `http://localhost:8091`)
|
||||
- `--model` (or `-m`): Model name (default: `fishaudio/s2-pro`)
|
||||
- `--text`: Text to synthesize (required)
|
||||
- `--ref-audio`: Reference audio for voice cloning (local path or URL)
|
||||
- `--ref-text`: Transcript of the reference audio
|
||||
- `--stream`: Enable streaming (PCM output)
|
||||
- `--response-format`: Audio format: wav, mp3, flac, pcm, aac, opus (default: wav)
|
||||
- `--output` (or `-o`): Output file path
|
||||
|
||||
## Streaming
|
||||
|
||||
Set `stream=true` with `response_format="pcm"` to receive raw PCM audio chunks as they are decoded:
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "Hello, how are you?",
|
||||
"voice": "default",
|
||||
"stream": true,
|
||||
"response_format": "pcm"
|
||||
}' --no-buffer | play -t raw -r 44100 -e signed -b 16 -c 1 -
|
||||
```
|
||||
|
||||
**Note:** Fish Speech outputs at 44.1 kHz (unlike Qwen3-TTS which outputs at 24 kHz).
|
||||
|
||||
## API Parameters
|
||||
|
||||
Fish Speech uses the same `/v1/audio/speech` endpoint as Qwen3-TTS. See the [Speech API reference](https://docs.vllm.ai/projects/vllm-omni/en/latest/serving/speech_api/) for full parameter documentation.
|
||||
|
||||
Key parameters for Fish Speech:
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| `input` | string | **required** | Text to synthesize |
|
||||
| `voice` | string | "default" | Voice name (use "default" for Fish Speech) |
|
||||
| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
|
||||
| `ref_audio` | string | null | Reference audio URL or base64 data URL for voice cloning |
|
||||
| `ref_text` | string | null | Transcript of reference audio (required for voice cloning) |
|
||||
| `max_new_tokens` | int | 4096 | Maximum tokens to generate |
|
||||
| `stream` | bool | false | Stream raw PCM chunks |
|
||||
|
||||
## Prerequisites
|
||||
|
||||
Install the `fish-speech` package for the DAC codec:
|
||||
|
||||
```bash
|
||||
pip install fish-speech
|
||||
```
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
1. **No audio output**: Make sure the `fish-speech` package is installed for the DAC decoder
|
||||
2. **Connection refused**: Ensure the server is running on the correct port
|
||||
3. **Flashinfer version mismatch**: Set `FLASHINFER_DISABLE_VERSION_CHECK=1` if you see version warnings
|
||||
4. **Out of memory**: Reduce `--gpu-memory-utilization` or use a GPU with more VRAM
|
||||