Compare commits

..
Author SHA1 Message Date
Hongsheng Liuandhsliu_ustc f5f78e46ff [skip ci]Update slides links in README&Docs (#1828)
Signed-off-by: hsliu_ustc <hsliu_ustc@noreply.gitcode.com>
Co-authored-by: hsliu_ustc <hsliu_ustc@noreply.gitcode.com>
2026-03-12 10:00:44 +08:00
166 changed files with 1210 additions and 16310 deletions

No files matched your search

@@ -1,56 +0,0 @@
#!/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"
'
+4 -8
View File
@@ -119,7 +119,7 @@ steps:
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: ":full_moon: Omni Model Perf Test & Test Case Statistics"
- label: ":full_moon: Omni Model Perf Test with H100"
key: nightly-performance
timeout_in_minutes: 180
depends_on: upload-nightly-pipeline
@@ -129,8 +129,6 @@ steps:
- 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:
@@ -160,19 +158,17 @@ steps:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: ":email: Nightly Collection & Email"
- label: ":email: Nightly Perf 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_email.py --report-file "tests/*.xlsx, tests/*.html"
- python tools/nightly/send_nightly_perf_email.py
- buildkite-agent artifact upload "tests/*.xlsx"
agents:
queue: "cpu_queue_premerge"
+45 -115
View File
@@ -34,7 +34,15 @@ steps:
- label: "Diffusion Model CPU offloading Test"
depends_on: upload-ready-pipeline
commands:
- timeout 10m pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
- |
timeout 20m bash -c '
set +e
pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
EXIT1=$$?
pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
EXIT2=$$?
exit $$((EXIT1 | EXIT2))
'
agents:
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
plugins:
@@ -261,117 +269,39 @@ steps:
# path: /mnt/hf-cache
# type: DirectoryOrCreate
- label: "Bagel Text2Img Model Test with H100"
depends_on: upload-ready-pipeline
commands:
- |
timeout 30m bash -c '
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_TEST_CLEAN_GPU_MEMORY=1
pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py
'
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 1
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: "Bagel Img2Img Model Test with H100"
depends_on: upload-ready-pipeline
commands:
- |
timeout 30m bash -c '
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_TEST_CLEAN_GPU_MEMORY=1
pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py
'
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 1
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: "Bagel Online Serving Test with H100"
depends_on: upload-ready-pipeline
commands:
- |
timeout 40m bash -c '
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_TEST_CLEAN_GPU_MEMORY=1
export VLLM_IMAGE_FETCH_TIMEOUT=60
pytest -s -v tests/e2e/online_serving/test_bagel_online.py
'
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 1
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
# - label: "Bagel Text2Img Model Test with H100"
# depends_on: upload-ready-pipeline
# commands:
# - |
# timeout 30m bash -c '
# export VLLM_WORKER_MULTIPROC_METHOD=spawn
# pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py
# '
# agents:
# queue: "mithril-h100-pool"
# plugins:
# - kubernetes:
# podSpec:
# containers:
# - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
# resources:
# limits:
# nvidia.com/gpu: 1
# volumeMounts:
# - name: devshm
# mountPath: /dev/shm
# - name: hf-cache
# mountPath: /root/.cache/huggingface
# env:
# - name: HF_HOME
# value: /root/.cache/huggingface
# nodeSelector:
# node.kubernetes.io/instance-type: gpu-h100-sxm
# volumes:
# - name: devshm
# emptyDir:
# medium: Memory
# - name: hf-cache
# hostPath:
# path: /mnt/hf-cache
# type: DirectoryOrCreate
+1 -2
View File
@@ -9,7 +9,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
</h3>
<p align="center">
| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://docs.google.com/presentation/d/1XJWgv79lORl8rbaVvp2d5Sqs6ZEBgAgj/edit?slide=id.p1#slide=id.p1"><b>Slides</b></a> |
| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://docs.google.com/presentation/d/1XJWgv79lORl8rbaVvp2d5Sqs6ZEBgAgj/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true"><b>Slides</b></a> |
</p>
@@ -17,7 +17,6 @@ Easy, fast, and cheap omni-modality model serving for everyone
*Latest News* 🔥
- [2026/03] Check out our first public [project deepdive](https://youtu.be/sgwNfsNnR9I) at the vLLM Hong Kong Meetup!
- [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.
@@ -6,7 +6,6 @@ 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
@@ -19,8 +18,8 @@ from .format import (
audio_to_base64,
base64_to_audio,
base64_to_image_tensor,
base64_to_video,
bytes_to_audio,
bytes_to_video,
image_tensor_to_base64,
image_tensor_to_png_bytes,
video_to_base64,
@@ -32,43 +31,10 @@ 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 | None = None, poll_interval: float = 5.0, max_poll_duration: float = 60 * 5
):
def __init__(self, base_url: str, timeout: float | None = None):
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,
@@ -296,41 +262,33 @@ class VLLMOmniClient:
content_type="image/png",
)
url = self.base_url + "/videos"
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 with session.post(url, data=form) as response:
if not response.ok:
error_text = await response.text()
raise (ValueError if response.status < 500 else RuntimeError)(
f"vLLM-Omni API returned status {response.status}: {error_text}"
)
try:
data = await response.json()
except aiohttp.ContentTypeError as e:
raise RuntimeError(f"Invalid JSON response from vLLM-Omni: {e}")
except aiohttp.ClientError as e:
raise RuntimeError(f"Network error connecting to vLLM-Omni at {url}: {e}")
if "data" not in data:
raise RuntimeError("API response missing 'data' field - expected OpenAI DALL-E format")
if not data["data"]:
raise RuntimeError("API returned empty data array")
try:
base64_str = data["data"][0]["b64_json"]
except (KeyError, IndexError):
raise RuntimeError("API response missing 'b64_json' field in first data item")
return base64_to_video(base64_str)
async def generate_understanding_chat_completion(
self,
@@ -143,7 +143,15 @@ 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:
def base64_to_video(base64_str: str) -> VideoInput:
if base64_str.startswith("data:video"):
_, base64_str = base64_str.split(",", 1)
try:
video_bytes = base64.b64decode(base64_str)
except Exception as e:
raise ValueError(f"Invalid base64 string: {e}")
video_buffer = BytesIO(video_bytes)
try:
@@ -200,18 +208,6 @@ def bytes_to_video(video_bytes: bytes) -> VideoInput:
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"]
@@ -70,17 +70,13 @@ 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:
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
sampling_param_list[i].update(params_to_add)
return sampling_param_list
+1 -22
View File
@@ -43,27 +43,6 @@ 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
@@ -111,7 +90,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`: requests can carry their own fields (e.g., `width/height/num_frames/num_inference_steps`), with overrides/fallbacks as below.
- For `trace`: each request can carry its own fields (e.g., `width/height/num_frames/num_inference_steps`).
Precedence rules for `trace` (i.e., what actually gets sent):
+1 -1
View File
@@ -76,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.17.0
ARG VLLM_VERSION=v0.16.0
RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm
WORKDIR /workspace/vllm
-9
View File
@@ -13,14 +13,9 @@ 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
@@ -28,11 +23,8 @@ 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
@@ -100,7 +92,6 @@ nav:
- api/vllm_omni
- CLI Reference: cli
- Community:
- Governance: community/governance.md
- community/*
- Slack: https://slack.vllm.ai
- Blog: https://blog.vllm.ai
+1 -1
View File
@@ -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/1XJWgv79lORl8rbaVvp2d5Sqs6ZEBgAgj/edit?slide=id.p1#slide=id.p1)
- [vllm-omni architecture design and recent roadmaps](https://docs.google.com/presentation/d/1XJWgv79lORl8rbaVvp2d5Sqs6ZEBgAgj/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true)
- [vllm-omni announcement blogpost](https://blog.vllm.ai/2025/11/30/vllm-omni.html)
-52
View File
@@ -1,52 +0,0 @@
# 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.
+1 -1
View File
@@ -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 = Omni(model=model_name, stage_configs_path="/path/to/custom_stage_configs.yaml")
omni_llm = OmniLLM(model=model_name, stage_configs_path="/path/to/custom_stage_configs.yaml")
```
For online serving:
+4 -5
View File
@@ -17,7 +17,6 @@ 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 * |
@@ -56,7 +55,7 @@ def test_video_to_audio()
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`, `xpu`, `npu`) and resource type (e.g., `L4`, `H100`, `MI325`, `B60`, `A2`, `A3`).
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`).
```
@pytest.mark.cuda
@pytest.mark.L4
@@ -81,14 +80,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", "xpu": "B60", "npu": "A2"},
num_cards={"cuda": 2, "rocm": 2, "xpu": 2, "npu": 2},
res={"cuda": "L4", "rocm": "MI325", "npu": "A2"},
num_cards={"cuda": 2, "rocm": 2, "npu": 2},
)
```
or
```python
@hardware_test(
res={"cuda": "L4", "rocm": "MI325", "xpu": "B60", "npu": "A2"},
res={"cuda": "L4", "rocm": "MI325", "npu": "A2"},
num_cards=2,
)
```
@@ -72,241 +72,3 @@ 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,589 +205,3 @@ 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,10 +7,9 @@ auto-configured when no explicit connector is specified for an edge.
## How It Works
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.
- Small payloads (< threshold): serialized and passed inline in metadata.
- Large payloads (>= threshold): stored in shared memory; the SHM name is
returned in metadata.
## Configuration
@@ -26,234 +25,3 @@ 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,265 +94,8 @@ stage_args:
Parameters:
- host: datasystem worker host.
- port: datasystem worker port (default: `35001` if omitted; the example above uses `31501` to match the worker startup command).
- port: datasystem worker port.
- 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.
+5 -35
View File
@@ -55,8 +55,6 @@ 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`
@@ -203,36 +201,6 @@ 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:**
@@ -272,7 +240,7 @@ NOTE: be careful to test adequately when refactoring classes that take this styl
## Approach 2: Intrusive Modification (For Complex Cases)
For models with dynamic sharding logic that cannot be expressed via `_sp_plan`, manually insert shard/gather calls.
For models with dynamic sharding logic that cannot be expressed via `_sp_plan`, manually insert shard/gather calls. Importantly, when taking this approach, be careful to ensure that you correctly manage the `_sp_shard_depth`; if the sequence parallel shard depth is 0, Ulysses will not be used.
**When to use:**
@@ -285,15 +253,18 @@ from vllm_omni.diffusion.distributed.sp_sharding import sp_shard, sp_gather
def forward(self, hidden_states, ...):
if self.parallel_config.sequence_parallel_size > 1:
# <Increment the _sp_shard_depth on the fwd context>
hidden_states = sp_shard(hidden_states, dim=1)
# ... computation ...
if self.parallel_config.sequence_parallel_size > 1:
output = sp_gather(output, dim=1)
# <Decrement the _sp_shard_depth on the fwd context>
return output
```
Note that currently, `sp_shard` / `sp_gather` do *not* automatically manage the `_sp_shard_depth`; you need to be careful to manage it yourself.
---
@@ -468,7 +439,6 @@ 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 |
@@ -483,7 +453,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? And which module boundaries need to be moved to facilitate this?
2. ✅ **Identify sharding boundaries** - Where should tensors be split/gathered?
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
-3
View File
@@ -50,8 +50,6 @@ th {
|`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
@@ -67,7 +65,6 @@ 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` |
+2 -60
View File
@@ -1,9 +1,6 @@
# Speech API
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.
vLLM-Omni provides an OpenAI-compatible API for text-to-speech (TTS) generation using Qwen3-TTS models.
Each server instance runs a single model (specified at startup via `vllm serve <model> --omni`).
@@ -12,22 +9,13 @@ Each server instance runs a single model (specified at startup via `vllm serve <
### Start the Server
```bash
# Qwen3-TTS: CustomVoice model (predefined speakers)
# 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
@@ -176,42 +164,6 @@ 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}
@@ -329,8 +281,6 @@ 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 |
@@ -339,14 +289,6 @@ 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
@@ -24,22 +24,21 @@ The following table shows which models are currently supported by parallelism me
### ImageGen
| Model | Model Identifier | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel | VAE-Patch-Parallel | Expert-Parallel | HSDP |
|--------------------------|--------------------------------------|:----------:|:-------:|:------------:|:---------------:|:------------------:|:---------------:|:----:|
| **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 | ❌ |
| Model | Model Identifier | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel | VAE-Patch-Parallel | Expert-Parallel |
|--------------------------|--------------------------------------|:----------:|:-------:|:------------:|:---------------:|:------------------:|:---------------:|
| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ✅ | ✅ | ❌ | ✅ | ❌ | N/A |
| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ✅ | ✅ | ❌ | ✅ | ❌ | N/A |
| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ❌ | ❌ | ❌ | ❌ | N/A |
| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | ✅ | N/A |
| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ | ❌ | N/A |
| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ | ❌ | N/A |
| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ✅ | ✅ | ✅ | ✅ | ❌ | N/A |
| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ✅ (TP=2 only) | ✅ | N/A |
| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ❌ | ❌ | ✅ | ✅ | N/A |
| **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ❌ | ❌ | ❌ | ✅ | ❌ | N/A |
| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ❌ | ✅ | ✅ | ❌ | N/A |
| **FLUX.2-dev** | `black-forest-labs/FLUX.2-dev` | ❌ | ❌ | ❌ | ✅ | ❌ | N/A |
| **HunyuanImage3.0** | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | ❌ | ❌ | ❌ | ✅ | ❌ | ✅ |
!!! note "TP Limitations for Diffusion Models"
We currently implement Tensor Parallelism (TP) only for the DiT (Diffusion Transformer) blocks. This is because the `text_encoder` component in vLLM-Omni uses the original Transformers implementation, which does not yet support TP.
+1 -2
View File
@@ -69,7 +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` | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ |
| **FLUX.2-dev** | `black-forest-labs/FLUX.2-dev` | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ |
### VideoGen
@@ -77,7 +77,6 @@ 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
@@ -1,74 +0,0 @@
# 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"
``````
@@ -1,97 +0,0 @@
# 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"
``````
@@ -1,58 +0,0 @@
# 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).
@@ -1,194 +0,0 @@
# 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,35 +29,6 @@ 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
@@ -67,7 +38,7 @@ export VLLM_OMNI_STORAGE_MAX_CONCURRENCY=8
bash run_curl_image_to_video.sh
# Or execute directly (OpenAI-style multipart)
create_response=$(curl -s http://localhost:8091/v1/videos \
curl -X POST http://localhost:8091/v1/videos \
-H "Accept: application/json" \
-F "prompt=A bear playing with yarn, smooth motion" \
-F "negative_prompt=low quality, blurry, static" \
@@ -81,23 +52,7 @@ create_response=$(curl -s http://localhost:8091/v1/videos \
-F "guidance_scale_2=1.0" \
-F "boundary_ratio=0.875" \
-F "flow_shift=12.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_i2v_output.mp4
-F "seed=42" | jq -r '.data[0].b64_json' | base64 -d > wan22_i2v_output.mp4
```
## Request Format
@@ -111,18 +66,6 @@ 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
@@ -142,63 +85,6 @@ 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"
@@ -29,34 +29,6 @@ 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
@@ -66,7 +38,7 @@ export VLLM_OMNI_STORAGE_MAX_CONCURRENCY=8
bash run_curl_text_to_video.sh
# Or execute directly (OpenAI-style multipart)
create_response=$(curl -s http://localhost:8091/v1/videos \
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" \
@@ -78,24 +50,7 @@ create_response=$(curl -s http://localhost:8091/v1/videos \
-F "guidance_scale=4.0" \
-F "guidance_scale_2=4.0" \
-F "boundary_ratio=0.875" \
-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
-F "seed=42" | jq -r '.data[0].b64_json' | base64 -d > wan22_output.mp4
```
## Request Format
@@ -130,9 +85,8 @@ 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 |
@@ -144,62 +98,24 @@ 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 |
## Create Response Format
`POST /v1/videos` returns a job record.
## Response Format
```json
{
"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
"created": 1234567890,
"data": [
{ "b64_json": "<base64-mp4>" }
]
}
```
## Retrieve, List, Download, and Delete
### Retrieve a job
## Extract Video
```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_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
# Extract base64 from response and decode to video
cat response.json | jq -r '.data[0].b64_json' | base64 -d > wan22_output.mp4
```
## Example materials
@@ -1,64 +0,0 @@
# 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.
@@ -1,197 +0,0 @@
# 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()
@@ -1,75 +0,0 @@
# 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.
@@ -1,252 +0,0 @@
"""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)
@@ -1,103 +0,0 @@
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)
@@ -1,58 +0,0 @@
# 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).
@@ -1,148 +0,0 @@
# 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()
@@ -1,172 +0,0 @@
# 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
@@ -1,272 +0,0 @@
"""Gradio demo for Fish Speech S2 Pro online serving via /v1/audio/speech API.
Supports:
- Text-to-speech synthesis
- Voice cloning from reference audio (upload or URL)
- Streaming (progressive PCM chunks) and non-streaming modes
Usage:
# Start the server first (see run_server.sh), then:
python gradio_demo.py --api-base http://localhost:8091
# Or use run_gradio_demo.sh to start both server and demo together.
"""
import argparse
import base64
import io
import gradio as gr
import httpx
import numpy as np
import soundfile as sf
PCM_SAMPLE_RATE = 44100
DEFAULT_API_BASE = "http://localhost:8091"
def encode_audio_to_base64(audio_data: tuple) -> str:
"""Encode Gradio audio input (sample_rate, numpy_array) to base64 data URL."""
sample_rate, audio_np = audio_data
if audio_np.dtype != np.int16:
if audio_np.dtype in (np.float32, np.float64):
audio_np = np.clip(audio_np, -1.0, 1.0)
audio_np = (audio_np * 32767).astype(np.int16)
else:
audio_np = audio_np.astype(np.int16)
buf = io.BytesIO()
sf.write(buf, audio_np, sample_rate, format="WAV")
wav_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
return f"data:audio/wav;base64,{wav_b64}"
def build_payload(
text: str,
ref_audio: tuple | None,
ref_audio_url: str,
ref_text: str,
response_format: str = "wav",
stream: bool = False,
) -> dict:
"""Build the /v1/audio/speech request payload."""
if not text or not text.strip():
raise gr.Error("Please enter text to synthesize.")
payload: dict = {
"input": text.strip(),
"voice": "default",
"response_format": "pcm" if stream else response_format,
"stream": stream,
}
# Voice cloning: ref_audio takes priority over URL.
ref_url_stripped = ref_audio_url.strip() if ref_audio_url else ""
if ref_audio is not None:
payload["ref_audio"] = encode_audio_to_base64(ref_audio)
elif ref_url_stripped:
payload["ref_audio"] = ref_url_stripped
if "ref_audio" in payload:
if not ref_text or not ref_text.strip():
raise gr.Error("Voice cloning requires a transcript of the reference audio.")
payload["ref_text"] = ref_text.strip()
return payload
def generate_speech(api_base: str, text: str, ref_audio, ref_audio_url, ref_text, response_format):
"""Non-streaming: call /v1/audio/speech and return full audio."""
payload = build_payload(text, ref_audio, ref_audio_url, ref_text, response_format, stream=False)
try:
with httpx.Client(timeout=300.0) as client:
resp = client.post(
f"{api_base}/v1/audio/speech",
json=payload,
headers={"Content-Type": "application/json", "Authorization": "Bearer EMPTY"},
)
except httpx.TimeoutException:
raise gr.Error("Request timed out. The server may be busy.")
except httpx.ConnectError:
raise gr.Error(f"Cannot connect to server at {api_base}. Is the server running?")
if resp.status_code != 200:
raise gr.Error(f"Server error ({resp.status_code}): {resp.text}")
content_type = resp.headers.get("content-type", "")
if "application/json" in content_type:
try:
raise gr.Error(f"Server error: {resp.json()}")
except ValueError:
pass
try:
if response_format == "pcm":
audio_np = np.frombuffer(resp.content, dtype=np.int16).astype(np.float32) / 32767.0
return (PCM_SAMPLE_RATE, audio_np)
audio_np, sample_rate = sf.read(io.BytesIO(resp.content))
if audio_np.ndim > 1:
audio_np = audio_np[:, 0]
return (sample_rate, audio_np.astype(np.float32))
except Exception as e:
raise gr.Error(f"Failed to decode audio: {e}")
def generate_speech_stream(api_base, text, ref_audio, ref_audio_url, ref_text, response_format, chunk_seconds=0.5):
"""Streaming: yield PCM chunks for progressive playback."""
payload = build_payload(text, ref_audio, ref_audio_url, ref_text, response_format, stream=True)
min_chunk_samples = int(PCM_SAMPLE_RATE * chunk_seconds)
prebuffer_chunks = 2
pending: list[np.ndarray] = []
pending_len = 0
chunks_yielded = 0
prebuffer: list[np.ndarray] = []
try:
leftover = b""
with httpx.Client(timeout=300.0) as client:
with client.stream(
"POST",
f"{api_base}/v1/audio/speech",
json=payload,
headers={"Content-Type": "application/json", "Authorization": "Bearer EMPTY"},
) as resp:
if resp.status_code != 200:
resp.read()
raise gr.Error(f"Server error ({resp.status_code}): {resp.text}")
for chunk in resp.iter_bytes():
if not chunk:
continue
raw = leftover + chunk
usable = len(raw) - (len(raw) % 2)
leftover = raw[usable:]
if usable == 0:
continue
float_samples = np.frombuffer(raw[:usable], dtype=np.int16).copy().astype(np.float32) / 32767.0
pending.append(float_samples)
pending_len += len(float_samples)
if pending_len >= min_chunk_samples:
audio_chunk = np.concatenate(pending)
pending.clear()
pending_len = 0
if chunks_yielded < prebuffer_chunks:
prebuffer.append(audio_chunk)
chunks_yielded += 1
if chunks_yielded == prebuffer_chunks:
yield (PCM_SAMPLE_RATE, np.concatenate(prebuffer))
prebuffer.clear()
else:
yield (PCM_SAMPLE_RATE, audio_chunk)
remaining = prebuffer + pending
if remaining:
yield (PCM_SAMPLE_RATE, np.concatenate(remaining))
except httpx.TimeoutException:
raise gr.Error("Request timed out. The server may be busy.")
except httpx.ConnectError:
raise gr.Error(f"Cannot connect to server at {api_base}. Is the server running?")
def on_stream_change(stream: bool):
"""Lock format to PCM when streaming."""
if stream:
return gr.update(value="pcm", interactive=False)
return gr.update(value="wav", interactive=True)
def build_interface(api_base: str, stream_chunk_seconds: float = 0.5):
"""Build the Gradio interface."""
with gr.Blocks(title="Fish Speech S2 Pro Demo") as demo:
gr.Markdown("# Fish Speech S2 Pro - Text to Speech")
gr.Markdown(f"**Server:** `{api_base}` | **Model:** fishaudio/s2-pro | **Output:** 44.1kHz")
with gr.Row():
with gr.Column(scale=3):
text_input = gr.Textbox(
label="Text to Synthesize",
placeholder="Enter text here...",
lines=4,
)
with gr.Accordion("Voice Cloning (optional)", open=False):
gr.Markdown(
"Upload or link a short reference audio (10-30s) and provide its transcript to clone the voice."
)
ref_audio = gr.Audio(
label="Reference Audio",
type="numpy",
sources=["upload", "microphone"],
)
ref_audio_url = gr.Textbox(
label="Reference Audio URL (alternative to upload)",
placeholder="https://example.com/reference.wav",
lines=1,
)
ref_text = gr.Textbox(
label="Reference Audio Transcript (required for cloning)",
placeholder="Exact transcript of the reference audio...",
lines=2,
)
with gr.Row():
response_format = gr.Dropdown(
choices=["wav", "mp3", "flac", "pcm"],
value="wav",
label="Audio Format",
scale=1,
)
stream_checkbox = gr.Checkbox(
label="Stream output",
value=False,
info="Progressive PCM streaming",
scale=1,
)
generate_btn = gr.Button("Generate Speech", variant="primary", size="lg")
with gr.Column(scale=2):
audio_output = gr.Audio(
label="Generated Audio",
interactive=False,
streaming=True,
autoplay=True,
)
gr.Markdown(
"### About\n"
"- **Fish Speech S2 Pro** by FishAudio: 4B dual-AR model\n"
"- **Voice cloning**: upload 10-30s reference + transcript\n"
"- **Streaming**: real-time PCM output\n"
"- **44.1kHz** output via DAC codec"
)
stream_checkbox.change(fn=on_stream_change, inputs=[stream_checkbox], outputs=[response_format])
all_inputs = [text_input, ref_audio, ref_audio_url, ref_text, response_format]
def dispatch(stream_enabled, *args):
if stream_enabled:
yield from generate_speech_stream(api_base, *args, chunk_seconds=stream_chunk_seconds)
else:
yield generate_speech(api_base, *args)
generate_btn.click(fn=dispatch, inputs=[stream_checkbox] + all_inputs, outputs=[audio_output])
demo.queue()
return demo
def main():
parser = argparse.ArgumentParser(description="Gradio demo for Fish Speech S2 Pro")
parser.add_argument("--api-base", default=DEFAULT_API_BASE, help=f"API base URL (default: {DEFAULT_API_BASE})")
parser.add_argument("--host", default="0.0.0.0", help="Gradio host (default: 0.0.0.0)")
parser.add_argument("--port", type=int, default=7860, help="Gradio port (default: 7860)")
parser.add_argument("--share", action="store_true", help="Share publicly")
parser.add_argument("--stream-chunk-seconds", type=float, default=0.5, help="Seconds per streaming chunk")
args = parser.parse_args()
print(f"Connecting to vLLM server at: {args.api_base}")
demo = build_interface(args.api_base, stream_chunk_seconds=args.stream_chunk_seconds)
demo.launch(server_name=args.host, server_port=args.port, share=args.share)
if __name__ == "__main__":
main()
@@ -1,48 +0,0 @@
#!/bin/bash
# Launch Fish Speech S2 Pro server + Gradio demo together.
#
# Usage:
# ./run_gradio_demo.sh
# CUDA_VISIBLE_DEVICES=0 PORT=8091 GRADIO_PORT=7860 ./run_gradio_demo.sh
set -e
MODEL="${MODEL:-fishaudio/s2-pro}"
PORT="${PORT:-8091}"
GRADIO_PORT="${GRADIO_PORT:-7860}"
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
REPO_ROOT="$(cd "$SCRIPT_DIR/../../.." && pwd)"
echo "Starting Fish Speech S2 Pro server (port $PORT)..."
FLASHINFER_DISABLE_VERSION_CHECK=1 \
vllm-omni serve "$MODEL" \
--stage-configs-path "$REPO_ROOT/vllm_omni/model_executor/stage_configs/fish_speech_s2_pro.yaml" \
--host 0.0.0.0 \
--port "$PORT" \
--gpu-memory-utilization 0.9 \
--trust-remote-code \
--enforce-eager \
--omni &
SERVER_PID=$!
cleanup() {
echo "Stopping server (PID $SERVER_PID)..."
kill $SERVER_PID 2>/dev/null
wait $SERVER_PID 2>/dev/null
}
trap cleanup EXIT
# Wait for server to be ready.
echo "Waiting for server to start..."
for i in $(seq 1 120); do
if curl -s "http://localhost:$PORT/health" > /dev/null 2>&1; then
echo "Server ready."
break
fi
sleep 2
done
echo "Starting Gradio demo (port $GRADIO_PORT)..."
python "$SCRIPT_DIR/gradio_demo.py" \
--api-base "http://localhost:$PORT" \
--port "$GRADIO_PORT"
@@ -1,23 +0,0 @@
#!/bin/bash
# Launch vLLM-Omni server for Fish Speech S2 Pro
#
# Usage:
# ./run_server.sh
# CUDA_VISIBLE_DEVICES=0 ./run_server.sh
set -e
MODEL="${MODEL:-fishaudio/s2-pro}"
PORT="${PORT:-8091}"
echo "Starting Fish Speech S2 Pro server with model: $MODEL"
FLASHINFER_DISABLE_VERSION_CHECK=1 \
vllm-omni serve "$MODEL" \
--stage-configs-path vllm_omni/model_executor/stage_configs/fish_speech_s2_pro.yaml \
--host 0.0.0.0 \
--port "$PORT" \
--gpu-memory-utilization 0.9 \
--trust-remote-code \
--enforce-eager \
--omni
@@ -1,129 +0,0 @@
"""Client for Fish Speech S2 Pro via /v1/audio/speech endpoint.
Examples:
# Basic TTS
python speech_client.py --text "Hello, how are you?"
# Voice cloning
python speech_client.py --text "Hello, how are you?" \
--ref-audio ref.wav --ref-text "This is the reference transcript."
# Streaming PCM output
python speech_client.py --text "Hello world" --stream --output output.pcm
"""
import argparse
import base64
import os
import httpx
DEFAULT_API_BASE = "http://localhost:8091"
DEFAULT_API_KEY = "EMPTY"
def encode_audio_to_base64(audio_path: str) -> str:
"""Encode a local audio file to base64 data URL."""
if not os.path.exists(audio_path):
raise FileNotFoundError(f"Audio file not found: {audio_path}")
ext = audio_path.lower().rsplit(".", 1)[-1]
mime_map = {"wav": "audio/wav", "mp3": "audio/mpeg", "flac": "audio/flac", "ogg": "audio/ogg"}
mime_type = mime_map.get(ext, "audio/wav")
with open(audio_path, "rb") as f:
audio_b64 = base64.b64encode(f.read()).decode("utf-8")
return f"data:{mime_type};base64,{audio_b64}"
def run_tts(args) -> None:
"""Generate speech via /v1/audio/speech API."""
payload = {
"model": args.model,
"input": args.text,
"voice": "default",
"response_format": args.response_format,
}
# Voice cloning parameters.
if args.ref_audio:
if args.ref_audio.startswith(("http://", "https://")):
payload["ref_audio"] = args.ref_audio
else:
payload["ref_audio"] = encode_audio_to_base64(args.ref_audio)
if args.ref_text:
payload["ref_text"] = args.ref_text
if args.stream:
payload["stream"] = True
payload["response_format"] = "pcm"
print(f"Model: {args.model}")
print(f"Text: {args.text}")
if args.ref_audio:
print(f"Voice cloning: ref_audio={args.ref_audio}, ref_text={args.ref_text}")
print("Generating audio...")
api_url = f"{args.api_base}/v1/audio/speech"
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {args.api_key}",
}
if args.stream:
output_path = args.output or "output.pcm"
with httpx.Client(timeout=300.0) as client:
with client.stream("POST", api_url, json=payload, headers=headers) as resp:
if resp.status_code != 200:
print(f"Error: {resp.status_code}")
print(resp.read().decode())
return
total_bytes = 0
with open(output_path, "wb") as f:
for chunk in resp.iter_bytes():
f.write(chunk)
total_bytes += len(chunk)
print(f"Streamed {total_bytes} bytes to: {output_path}")
else:
with httpx.Client(timeout=300.0) as client:
response = client.post(api_url, json=payload, headers=headers)
if response.status_code != 200:
print(f"Error: {response.status_code}")
print(response.text)
return
try:
text = response.content.decode("utf-8")
if text.startswith('{"error"'):
print(f"Error: {text}")
return
except UnicodeDecodeError:
pass
output_path = args.output or "output.wav"
with open(output_path, "wb") as f:
f.write(response.content)
print(f"Audio saved to: {output_path}")
def main():
parser = argparse.ArgumentParser(description="Fish Speech S2 Pro TTS client")
parser.add_argument("--api-base", default=DEFAULT_API_BASE, help="API base URL")
parser.add_argument("--api-key", default=DEFAULT_API_KEY, help="API key")
parser.add_argument("--model", "-m", default="fishaudio/s2-pro", help="Model name")
parser.add_argument("--text", required=True, help="Text to synthesize")
parser.add_argument("--ref-audio", default=None, help="Reference audio for voice cloning (path or URL)")
parser.add_argument("--ref-text", default=None, help="Transcript of reference audio")
parser.add_argument("--stream", action="store_true", help="Enable streaming (PCM output)")
parser.add_argument(
"--response-format",
default="wav",
choices=["wav", "mp3", "flac", "pcm", "aac", "opus"],
help="Audio format (default: wav)",
)
parser.add_argument("--output", "-o", default=None, help="Output file path")
args = parser.parse_args()
run_tts(args)
if __name__ == "__main__":
main()
@@ -26,35 +26,6 @@ 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
@@ -64,7 +35,7 @@ export VLLM_OMNI_STORAGE_MAX_CONCURRENCY=8
bash run_curl_image_to_video.sh
# Or execute directly (OpenAI-style multipart)
create_response=$(curl -s http://localhost:8091/v1/videos \
curl -X POST http://localhost:8091/v1/videos \
-H "Accept: application/json" \
-F "prompt=A bear playing with yarn, smooth motion" \
-F "negative_prompt=low quality, blurry, static" \
@@ -78,23 +49,7 @@ create_response=$(curl -s http://localhost:8091/v1/videos \
-F "guidance_scale_2=1.0" \
-F "boundary_ratio=0.875" \
-F "flow_shift=12.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_i2v_output.mp4
-F "seed=42" | jq -r '.data[0].b64_json' | base64 -d > wan22_i2v_output.mp4
```
## Request Format
@@ -108,18 +63,6 @@ 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
@@ -138,60 +81,3 @@ curl -X POST http://localhost:8091/v1/videos \
-F "flow_shift=12.0" \
-F "seed=42"
```
## Create Response Format
`POST /v1/videos` returns a job record, not inline base64 video data.
```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
```
@@ -1,75 +1,33 @@
#!/bin/bash
# Wan2.2 image-to-video curl example using the async video job API.
set -euo pipefail
# Wan2.2 image-to-video curl example (OpenAI-style multipart)
INPUT_IMAGE="${INPUT_IMAGE:-../../offline_inference/image_to_video/qwen-bear.png}"
BASE_URL="${BASE_URL:-http://localhost:8099}"
OUTPUT_PATH="${OUTPUT_PATH:-wan22_i2v_output.mp4}"
NEGATIVE_PROMPT="${NEGATIVE_PROMPT:-}"
POLL_INTERVAL="${POLL_INTERVAL:-2}"
if [ ! -f "$INPUT_IMAGE" ]; then
echo "Input image not found: $INPUT_IMAGE"
exit 1
fi
create_cmd=(
curl -sS -X POST "${BASE_URL}/v1/videos"
-H "Accept: application/json"
-F "prompt=A bear playing with yarn, smooth motion"
-F "input_reference=@${INPUT_IMAGE}"
-F "seconds=2"
-F "size=832x480"
-F "fps=16"
-F "num_inference_steps=40"
-F "guidance_scale=1.0"
-F "guidance_scale_2=1.0"
-F "boundary_ratio=0.875"
-F "flow_shift=12.0"
-F "seed=42"
)
if [ -n "${NEGATIVE_PROMPT}" ]; then
create_cmd+=(-F "negative_prompt=${NEGATIVE_PROMPT}")
NEGATIVE_PROMPT_FLAG=""
if [ -n "$NEGATIVE_PROMPT" ]; then
NEGATIVE_PROMPT_FLAG="-F negative_prompt=${NEGATIVE_PROMPT}"
fi
create_response="$("${create_cmd[@]}")"
video_id="$(echo "${create_response}" | jq -r '.id')"
if [ -z "${video_id}" ] || [ "${video_id}" = "null" ]; then
echo "Failed to create video job:"
echo "${create_response}" | jq .
exit 1
fi
curl -X POST http://localhost:8099/v1/videos \
-H "Accept: application/json" \
-F "prompt=A bear playing with yarn, smooth motion" \
$NEGATIVE_PROMPT_FLAG \
-F "input_reference=@${INPUT_IMAGE}" \
-F "seconds=2" \
-F "size=832x480" \
-F "fps=16" \
-F "num_inference_steps=40" \
-F "guidance_scale=1.0" \
-F "guidance_scale_2=1.0" \
-F "boundary_ratio=0.875" \
-F "flow_shift=12.0" \
-F "seed=42" | jq -r '.data[0].b64_json' | base64 -d > "${OUTPUT_PATH}"
echo "Created video job ${video_id}"
echo "${create_response}" | jq .
while true; do
status_response="$(curl -sS "${BASE_URL}/v1/videos/${video_id}")"
status="$(echo "${status_response}" | jq -r '.status')"
case "${status}" in
queued|in_progress)
echo "Video job ${video_id} status: ${status}"
sleep "${POLL_INTERVAL}"
;;
completed)
echo "${status_response}" | jq .
break
;;
failed)
echo "Video generation failed:"
echo "${status_response}" | jq .
exit 1
;;
*)
echo "Unexpected status response:"
echo "${status_response}" | jq .
exit 1
;;
esac
done
curl -sS -L "${BASE_URL}/v1/videos/${video_id}/content" -o "${OUTPUT_PATH}"
echo "Saved video to ${OUTPUT_PATH}"
@@ -349,42 +349,6 @@ curl -X POST http://localhost:8091/v1/audio/speech \
- `speed` adjustment is not supported when streaming.
- Requires the server stage config to have `async_chunk: true` (default in `qwen3_tts.yaml`).
## Streaming Text Input (WebSocket)
The `/v1/audio/speech/stream` WebSocket endpoint accepts text incrementally, buffers and splits it at sentence boundaries, and generates audio per sentence.
When `stream_audio=true`, each sentence is emitted as `audio.start`, one or more binary PCM frames, and `audio.done`.
### Quick Start
```bash
python streaming_speech_client.py \
--text "Hello world. How are you? I am fine."
python streaming_speech_client.py \
--text "Hello world. How are you? I am fine." \
--simulate-stt --stt-delay 0.1
```
### WebSocket Protocol
Client -> Server:
```jsonc
{"type": "session.config", "voice": "Vivian", "task_type": "CustomVoice", "language": "Auto", "split_granularity": "sentence", "stream_audio": true, "response_format": "pcm"}
{"type": "input.text", "text": "Hello, how are you? "}
{"type": "input.done"}
```
Server -> Client:
```jsonc
{"type": "audio.start", "sentence_index": 0, "sentence_text": "Hello, how are you?", "format": "pcm", "sample_rate": 24000}
// binary PCM frame(s)
{"type": "audio.done", "sentence_index": 0, "total_bytes": 96000, "error": false}
{"type": "session.done", "total_sentences": 1}
```
## Limitations
- **Single request**: Batch processing is not yet optimized for online serving.
@@ -1,249 +0,0 @@
"""WebSocket client for streaming text-input TTS.
Connects to the /v1/audio/speech/stream endpoint, sends text incrementally
(simulating real-time STT output), and saves per-sentence audio files.
Usage:
# Send full text at once
python streaming_speech_client.py --text "Hello world. How are you? I am fine."
# Simulate STT: send text word-by-word with delay
python streaming_speech_client.py \
--text "Hello world. How are you? I am fine." \
--simulate-stt --stt-delay 0.1
# VoiceDesign task
python streaming_speech_client.py \
--text "Today is a great day. The weather is nice." \
--task-type VoiceDesign \
--instructions "A cheerful young female voice"
# Base task (voice cloning)
python streaming_speech_client.py \
--text "Hello world. How are you?" \
--task-type Base \
--ref-audio /path/to/reference.wav \
--ref-text "Transcript of reference audio"
Requirements:
pip install websockets
"""
import argparse
import asyncio
import json
import os
try:
import websockets
except ImportError:
print("Please install websockets: pip install websockets")
raise SystemExit(1)
async def stream_tts(
url: str,
text: str,
config: dict,
output_dir: str,
simulate_stt: bool = False,
stt_delay: float = 0.1,
) -> None:
"""Connect to the streaming TTS endpoint and process audio responses."""
os.makedirs(output_dir, exist_ok=True)
async with websockets.connect(url) as ws:
# 1. Send session config
config_msg = {"type": "session.config", **config}
await ws.send(json.dumps(config_msg))
print(f"Sent session config: {config}")
# 2. Send text (either all at once or word-by-word)
async def send_text():
if simulate_stt:
words = text.split(" ")
for i, word in enumerate(words):
chunk = word + (" " if i < len(words) - 1 else "")
await ws.send(
json.dumps(
{
"type": "input.text",
"text": chunk,
}
)
)
print(f" Sent: {chunk!r}")
await asyncio.sleep(stt_delay)
else:
await ws.send(
json.dumps(
{
"type": "input.text",
"text": text,
}
)
)
print(f"Sent full text: {text!r}")
# 3. Signal end of input
await ws.send(json.dumps({"type": "input.done"}))
print("Sent input.done")
# Run sender and receiver concurrently
sender_task = asyncio.create_task(send_text())
response_format = config.get("response_format", "wav")
current_sentence_index = 0
current_chunks: list[bytes] = []
try:
while True:
message = await ws.recv()
if isinstance(message, bytes):
current_chunks.append(message)
print(f" Received audio chunk for sentence {current_sentence_index}: {len(message)} bytes")
else:
# JSON frame
msg = json.loads(message)
msg_type = msg.get("type")
if msg_type == "audio.start":
current_sentence_index = msg["sentence_index"]
current_chunks = []
print(f" [sentence {msg['sentence_index']}] Generating: {msg['sentence_text']!r}")
elif msg_type == "audio.done":
filename = os.path.join(
output_dir,
f"sentence_{msg['sentence_index']:03d}.{response_format}",
)
with open(filename, "wb") as f:
f.write(b"".join(current_chunks))
print(
f" [sentence {msg['sentence_index']}] Done"
f" bytes={msg.get('total_bytes', len(b''.join(current_chunks)))}"
f" error={msg.get('error', False)}"
f" -> {filename}"
)
current_chunks = []
elif msg_type == "session.done":
print(f"\nSession complete: {msg['total_sentences']} sentence(s) generated")
break
elif msg_type == "error":
print(f" ERROR: {msg['message']}")
else:
print(f" Unknown message: {msg}")
finally:
sender_task.cancel()
try:
await sender_task
except asyncio.CancelledError:
pass # Task cancellation is expected during shutdown
print(f"\nAudio files saved to: {output_dir}/")
def main():
parser = argparse.ArgumentParser(description="Streaming text-input TTS client")
parser.add_argument(
"--url",
default="ws://localhost:8000/v1/audio/speech/stream",
help="WebSocket endpoint URL",
)
parser.add_argument(
"--text",
required=True,
help="Text to synthesize",
)
parser.add_argument(
"--output-dir",
default="streaming_tts_output",
help="Directory to save audio files (default: streaming_tts_output)",
)
# Session config options
parser.add_argument("--model", default=None, help="Model name")
parser.add_argument("--voice", default="Vivian", help="Speaker voice")
parser.add_argument(
"--task-type",
default="CustomVoice",
choices=["CustomVoice", "VoiceDesign", "Base"],
help="TTS task type",
)
parser.add_argument("--language", default="Auto", help="Language")
parser.add_argument("--instructions", default=None, help="Voice style instructions")
parser.add_argument(
"--response-format",
default="wav",
choices=["wav", "pcm", "flac", "mp3", "aac", "opus"],
help="Audio format",
)
parser.add_argument(
"--stream-audio",
action="store_true",
help="Receive one or more PCM chunks per sentence (requires --response-format pcm)",
)
parser.add_argument("--speed", type=float, default=1.0, help="Playback speed (0.25-4.0)")
parser.add_argument("--max-new-tokens", type=int, default=None, help="Max tokens")
# Base task options
parser.add_argument("--ref-audio", default=None, help="Reference audio")
parser.add_argument("--ref-text", default=None, help="Reference text")
parser.add_argument(
"--x-vector-only-mode",
action="store_true",
default=False,
help="Speaker embedding only mode",
)
# STT simulation
parser.add_argument(
"--simulate-stt",
action="store_true",
help="Simulate STT by sending text word-by-word",
)
parser.add_argument(
"--stt-delay",
type=float,
default=0.1,
help="Delay between words in STT simulation (seconds)",
)
args = parser.parse_args()
# Build session config (only include non-None values)
config = {}
for key in [
"model",
"voice",
"task_type",
"language",
"instructions",
"response_format",
"speed",
"max_new_tokens",
"ref_audio",
"ref_text",
]:
val = getattr(args, key.replace("-", "_"), None)
if val is not None:
config[key] = val
if args.stream_audio:
config["stream_audio"] = True
if args.x_vector_only_mode:
config["x_vector_only_mode"] = True
asyncio.run(
stream_tts(
url=args.url,
text=args.text,
config=config,
output_dir=args.output_dir,
simulate_stt=args.simulate_stt,
stt_delay=args.stt_delay,
)
)
if __name__ == "__main__":
main()
+12 -97
View File
@@ -26,35 +26,6 @@ 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
@@ -64,7 +35,7 @@ export VLLM_OMNI_STORAGE_MAX_CONCURRENCY=8
bash run_curl_text_to_video.sh
# Or execute directly (OpenAI-style multipart)
create_response=$(curl -s http://localhost:8091/v1/videos \
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" \
@@ -76,24 +47,7 @@ create_response=$(curl -s http://localhost:8091/v1/videos \
-F "guidance_scale=4.0" \
-F "guidance_scale_2=4.0" \
-F "boundary_ratio=0.875" \
-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
-F "seed=42" | jq -r '.data[0].b64_json' | base64 -d > wan22_output.mp4
```
## Request Format
@@ -128,9 +82,8 @@ 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 |
@@ -142,60 +95,22 @@ 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 |
## Create Response Format
`POST /v1/videos` returns a job record, not inline base64 video data.
## Response Format
```json
{
"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
"created": 1234567890,
"data": [
{ "b64_json": "<base64-mp4>" }
]
}
```
## Retrieve, List, Download, and Delete
### Retrieve a job
## Extract Video
```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_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
# Extract base64 from response and decode to video
cat response.json | jq -r '.data[0].b64_json' | base64 -d > wan22_output.mp4
```
@@ -1,63 +1,20 @@
#!/bin/bash
# Wan2.2 text-to-video curl example using the async video job API.
# Wan2.2 text-to-video curl example
set -euo pipefail
OUTPUT_PATH="wan22_output.mp4"
BASE_URL="${BASE_URL:-http://localhost:8098}"
OUTPUT_PATH="${OUTPUT_PATH:-wan22_output.mp4}"
POLL_INTERVAL="${POLL_INTERVAL:-2}"
curl -X POST http://localhost:8098/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 "seconds=2" \
-F "size=832x480" \
-F "negative_prompt=色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,静止,整体发灰,最差质量,低质量,JPEG压缩残留,丑陋的,残缺的,多余的手指,画得不好的手部,画得不好的脸部,畸形的,毁容的,形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,背景人很多,倒着走" \
-F "fps=16" \
-F "num_inference_steps=40" \
-F "guidance_scale=4.0" \
-F "guidance_scale_2=4.0" \
-F "boundary_ratio=0.875" \
-F "flow_shift=5.0" \
-F "seed=42" | jq -r '.data[0].b64_json' | base64 -d > "${OUTPUT_PATH}"
create_response=$(
curl -sS -X POST "${BASE_URL}/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 "seconds=2" \
-F "size=832x480" \
-F "negative_prompt=色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,静止,整体发灰,最差质量,低质量,JPEG压缩残留,丑陋的,残缺的,多余的手指,画得不好的手部,画得不好的脸部,畸形的,毁容的,形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,背景人很多,倒着走" \
-F "fps=16" \
-F "num_inference_steps=40" \
-F "guidance_scale=4.0" \
-F "guidance_scale_2=4.0" \
-F "boundary_ratio=0.875" \
-F "flow_shift=5.0" \
-F "seed=42"
)
video_id="$(echo "${create_response}" | jq -r '.id')"
if [ -z "${video_id}" ] || [ "${video_id}" = "null" ]; then
echo "Failed to create video job:"
echo "${create_response}" | jq .
exit 1
fi
echo "Created video job ${video_id}"
echo "${create_response}" | jq .
while true; do
status_response="$(curl -sS "${BASE_URL}/v1/videos/${video_id}")"
status="$(echo "${status_response}" | jq -r '.status')"
case "${status}" in
queued|in_progress)
echo "Video job ${video_id} status: ${status}"
sleep "${POLL_INTERVAL}"
;;
completed)
echo "${status_response}" | jq .
break
;;
failed)
echo "Video generation failed:"
echo "${status_response}" | jq .
exit 1
;;
*)
echo "Unexpected status response:"
echo "${status_response}" | jq .
exit 1
;;
esac
done
curl -sS -L "${BASE_URL}/v1/videos/${video_id}/content" -o "${OUTPUT_PATH}"
echo "Saved video to ${OUTPUT_PATH}"
-3
View File
@@ -167,7 +167,6 @@ markers = [
"gpu: Tests that run on GPU (auto-added)",
"cuda: Tests that run on CUDA (auto-added)",
"rocm: Tests that run on AMD/ROCm (auto-added)",
"xpu: Tests that run on XPU (auto-added)",
"npu: Tests that run on NPU/Ascend (auto-added)",
# specified computation resources marks (auto-added)
"H100: Tests that require H100 GPU",
@@ -177,11 +176,9 @@ markers = [
"A3: Tests that require A3 NPU",
"distributed_cuda: Tests that require multi cards on CUDA platform",
"distributed_rocm: Tests that require multi cards on ROCm platform",
"distributed_xpu: Tests that require multi cards on XPU platform",
"distributed_npu: Tests that require multi cards on NPU platform",
"skipif_cuda: Skip if the num of CUDA cards is less than the required",
"skipif_rocm: Skip if the num of ROCm cards is less than the required",
"skipif_xpu: Skip if the num of XPU cards is less than the required",
"skipif_npu: Skip if the num of NPU cards is less than the required",
# more detailed markers
"slow: Slow tests (may skip in quick CI)",
+1 -1
View File
@@ -6,7 +6,7 @@ diffusers>=0.36.0
accelerate==1.12.0
gradio==5.50
soundfile>=0.13.1
cache-dit==1.3.0
cache-dit==1.2.0
tqdm>=4.66.0
torchsde>=0.2.6
openai-whisper>=20250625
+1 -5
View File
@@ -5,14 +5,10 @@ import pytest
from tests.conftest import OmniServerParams
from tests.utils import hardware_test
from vllm_omni.platforms import current_omni_platform
models = ["Qwen/Qwen2.5-Omni-7B"]
stage_configs = [str(Path(__file__).parent.parent / "e2e" / "stage_configs" / "qwen2_5_omni_ci.yaml")]
if current_omni_platform.is_xpu():
stage_configs = [str(Path(__file__).parent.parent / "e2e" / "stage_configs" / "xpu" / "qwen2_5_omni_ci.yaml")]
# Create parameter combinations for model and stage config
test_params = [
OmniServerParams(model=model, stage_config_path=stage_config) for model in models for stage_config in stage_configs
@@ -21,7 +17,7 @@ test_params = [
@pytest.mark.core_model
@pytest.mark.benchmark
@hardware_test(res={"cuda": "L4", "xpu": "B60"}, num_cards=3)
@hardware_test(res={"cuda": "L4"}, num_cards=3)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
def test_bench_serve_chat(omni_server):
command = [
+3 -12
View File
@@ -82,15 +82,6 @@ DIFFUSION_SINGLE_SAMPLING_PARAMS = DiffusionSamplingParams(
}
)
DIFFUSION_VIDEO_SINGLE_SAMPLING_PARAMS = DiffusionSamplingParams(
{
"num_inference_steps": 30,
"guidance_scale": 6.0,
"true_cfg_scale": 1.5,
}
)
AR_LIST_SAMPLING_PARAMS = [
AutoregressionSamplingParams(
{
@@ -319,8 +310,8 @@ def _build_mock_outputs(outputs: Iterable[OmniRequestOutput], sampling_case: Sam
)
elif sampling_case.kind is SamplingKind.VIDEO_DIFFUSION_SINGLE:
assert len(received_sampling_params_list) == 1
expected = DIFFUSION_VIDEO_SINGLE_SAMPLING_PARAMS.copy()
# expected["num_outputs_per_prompt"] = expected.pop("n") # convert from n to num_outputs_per_prompt
expected = DIFFUSION_SINGLE_SAMPLING_PARAMS.copy()
expected["num_outputs_per_prompt"] = expected.pop("n") # convert from n to num_outputs_per_prompt
_assert_sampling_param_values(
received_sampling_params_list[0],
{
@@ -699,7 +690,7 @@ async def test_tts_nodes(api_server: str, node_cls, call_kwargs: dict, sampling_
pytest.param(
SamplingCase(
kind=SamplingKind.VIDEO_DIFFUSION_SINGLE,
sampling_params=DIFFUSION_VIDEO_SINGLE_SAMPLING_PARAMS,
sampling_params=DIFFUSION_SINGLE_SAMPLING_PARAMS,
lora=LORA_PARAMS,
),
id="single-diffusion-sampling-params",
-47
View File
@@ -508,50 +508,3 @@ def test_scale_keys_are_rounded():
# Currently we round keys to 3 decimal places
manager._update_adapter_scale(adapter_id, 0.0031)
assert manager._adapter_scales[adapter_id] == 0.003
def test_lora_manager_uses_valid_max_rank(monkeypatch):
"""Ensure that the LoRA manager uses a valid max rank for vLLM."""
manager = DiffusionLoRAManager(
pipeline=torch.nn.Module(),
device=torch.device("cpu"),
dtype=torch.bfloat16,
)
# Ensure that the rank is correctly adjusted to the smallest valid max rank
supported_max_rank = 64
unsupported_max_rank = 63
assert supported_max_rank in DiffusionLoRAManager._VALID_MAX_RANKS
assert unsupported_max_rank not in DiffusionLoRAManager._VALID_MAX_RANKS
def _fake_load(_req: LoRARequest):
lora_model = type("LM", (), {"id": _req.lora_int_id})()
peft_helper = type("PH", (), {"r": unsupported_max_rank})()
return lora_model, peft_helper
monkeypatch.setattr(manager, "_load_adapter", _fake_load)
req1 = _dummy_lora_request(1)
manager.add_adapter(req1)
assert manager._max_lora_rank == supported_max_rank
@pytest.mark.parametrize("rank", [-1, 0, DiffusionLoRAManager._VALID_MAX_RANKS[-1] + 1])
def test_lora_manager_max_rank_validation(monkeypatch, rank):
"""Check that invalid max ranks are handled correctly."""
manager = DiffusionLoRAManager(
pipeline=torch.nn.Module(),
device=torch.device("cpu"),
dtype=torch.bfloat16,
)
lora_rank = rank
def _fake_load(_req: LoRARequest):
lora_model = type("LM", (), {"id": _req.lora_int_id})()
peft_helper = type("PH", (), {"r": lora_rank})()
return lora_model, peft_helper
monkeypatch.setattr(manager, "_load_adapter", _fake_load)
req1 = _dummy_lora_request(1)
with pytest.raises(ValueError):
manager.add_adapter(req1)
@@ -1,112 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for HunyuanFusedMoE (Support HunyuanImage3 Diffusion Model, 5a779b4)."""
import pytest
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
class TestHunyuanFusedMoEPlatformDispatch:
"""Test platform dispatch via platform qualname hooks."""
def test_default_platform_uses_default_impl_qualname(self, mocker):
"""HunyuanFusedMoE should resolve the impl class from the platform hook."""
import vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe as hunyuan_moe
mock_platform = mocker.MagicMock()
mock_platform.get_diffusion_model_impl_qualname.return_value = (
"vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe.HunyuanFusedMoEDefault"
)
mocker.patch.object(
hunyuan_moe,
"current_omni_platform",
mock_platform,
)
mock_resolve = mocker.patch.object(hunyuan_moe, "resolve_obj_by_qualname")
mock_impl = mocker.MagicMock()
mock_resolve.return_value = mock_impl
from vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe import (
HunyuanFusedMoE,
)
HunyuanFusedMoE(prefix="")
mock_platform.prepare_diffusion_op_runtime.assert_called_once_with("hunyuan_fused_moe")
mock_platform.get_diffusion_model_impl_qualname.assert_called_once_with("hunyuan_fused_moe")
mock_resolve.assert_called_once_with(
"vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe.HunyuanFusedMoEDefault"
)
mock_impl.assert_called_once_with(prefix="")
class TestHunyuanFusedMoEFactory:
"""Test HunyuanFusedMoE factory __new__ and make_expert_params_mapping delegation."""
def test_new_delegates_to_impl_class(self, mocker):
"""HunyuanFusedMoE(prefix=..., **kwargs) should instantiate and return impl instance."""
import vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe as hunyuan_moe
class MockImpl:
def __init__(self, *, prefix: str = "", **kwargs):
self.prefix = prefix
self.kwargs = kwargs
mock_platform = mocker.MagicMock()
mock_platform.get_diffusion_model_impl_qualname.return_value = "mock.impl.Qualname"
mocker.patch.object(hunyuan_moe, "current_omni_platform", mock_platform)
mock_impl_class = mocker.MagicMock(return_value=MockImpl(prefix="test", a=1))
mocker.patch.object(hunyuan_moe, "resolve_obj_by_qualname", return_value=mock_impl_class)
from vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe import (
HunyuanFusedMoE,
)
result = HunyuanFusedMoE(prefix="test", a=1)
assert isinstance(result, MockImpl)
assert result.prefix == "test"
assert result.kwargs == {"a": 1}
mock_platform.prepare_diffusion_op_runtime.assert_called_once_with("hunyuan_fused_moe")
mock_platform.get_diffusion_model_impl_qualname.assert_called_once_with("hunyuan_fused_moe")
mock_impl_class.assert_called_once_with(prefix="test", a=1)
def test_make_expert_params_mapping_delegates_to_impl(self, mocker):
"""make_expert_params_mapping should delegate to impl class method."""
import vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe as hunyuan_moe
expected_mapping = [("a", "b", 0, "c")]
mock_platform = mocker.MagicMock()
mock_platform.get_diffusion_model_impl_qualname.return_value = "mock.impl.Qualname"
mocker.patch.object(hunyuan_moe, "current_omni_platform", mock_platform)
mock_impl_class = mocker.MagicMock()
mock_impl_class.make_expert_params_mapping = mocker.MagicMock(return_value=expected_mapping)
mocker.patch.object(hunyuan_moe, "resolve_obj_by_qualname", return_value=mock_impl_class)
from vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe import (
HunyuanFusedMoE,
)
result = HunyuanFusedMoE.make_expert_params_mapping(
model=None,
ckpt_gate_proj_name="gate",
ckpt_down_proj_name="down",
ckpt_up_proj_name="up",
num_experts=4,
num_redundant_experts=0,
)
assert result == expected_mapping
mock_platform.get_diffusion_model_impl_qualname.assert_called_once_with("hunyuan_fused_moe")
mock_impl_class.make_expert_params_mapping.assert_called_once_with(
None,
ckpt_gate_proj_name="gate",
ckpt_down_proj_name="down",
ckpt_up_proj_name="up",
num_experts=4,
num_redundant_experts=0,
)
-78
View File
@@ -1,78 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for OmniDiffusionConfig.is_moe (fix is_moe type and threshold, 6663c0b)."""
import pytest
from vllm_omni.diffusion.data import OmniDiffusionConfig, TransformerConfig
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
class TestOmniDiffusionConfigIsMoE:
"""Tests for OmniDiffusionConfig.is_moe property.
Covers commit 6663c0b: fix is_moe type and threshold
- num_experts must be (list, tuple, int); otherwise return False.
- Threshold: is_moe is True when num_experts > 0 (not > 1).
"""
def test_is_moe_missing_num_experts_returns_false(self):
"""When num_experts is absent, is_moe should be False."""
tf_config = TransformerConfig.from_dict({})
config = OmniDiffusionConfig(model="test", tf_model_config=tf_config)
assert config.is_moe is False
def test_is_moe_none_num_experts_returns_false(self):
"""When num_experts is explicitly None (e.g. in params), is_moe should be False."""
tf_config = TransformerConfig.from_dict({"num_experts": None})
config = OmniDiffusionConfig(model="test", tf_model_config=tf_config)
assert config.is_moe is False
def test_is_moe_non_allowed_type_returns_false(self):
"""When num_experts is not int/list/tuple (e.g. str), is_moe should be False."""
tf_config = TransformerConfig.from_dict({"num_experts": "2"})
config = OmniDiffusionConfig(model="test", tf_model_config=tf_config)
assert config.is_moe is False
def test_is_moe_int_zero_returns_false(self):
"""num_experts int 0 should yield is_moe False (threshold > 0)."""
tf_config = TransformerConfig.from_dict({"num_experts": 0})
config = OmniDiffusionConfig(model="test", tf_model_config=tf_config)
assert config.is_moe is False
def test_is_moe_int_one_returns_true(self):
"""num_experts int 1 should yield is_moe True (threshold > 0, not > 1)."""
tf_config = TransformerConfig.from_dict({"num_experts": 1})
config = OmniDiffusionConfig(model="test", tf_model_config=tf_config)
assert config.is_moe is True
def test_is_moe_int_gt_one_returns_true(self):
"""num_experts int > 1 should yield is_moe True."""
tf_config = TransformerConfig.from_dict({"num_experts": 2})
config = OmniDiffusionConfig(model="test", tf_model_config=tf_config)
assert config.is_moe is True
def test_is_moe_list_all_zero_returns_false(self):
"""num_experts list with all <= 0 should yield is_moe False."""
tf_config = TransformerConfig.from_dict({"num_experts": [0]})
config = OmniDiffusionConfig(model="test", tf_model_config=tf_config)
assert config.is_moe is False
def test_is_moe_list_has_positive_returns_true(self):
"""num_experts list with any int > 0 should yield is_moe True."""
tf_config = TransformerConfig.from_dict({"num_experts": [0, 1]})
config = OmniDiffusionConfig(model="test", tf_model_config=tf_config)
assert config.is_moe is True
def test_is_moe_tuple_has_positive_returns_true(self):
"""num_experts tuple with any int > 0 should yield is_moe True."""
tf_config = TransformerConfig.from_dict({"num_experts": (0, 2)})
config = OmniDiffusionConfig(model="test", tf_model_config=tf_config)
assert config.is_moe is True
def test_is_moe_list_non_int_ignored(self):
"""num_experts list with only non-int entries should yield is_moe False."""
tf_config = TransformerConfig.from_dict({"num_experts": ["a", 0.0]})
config = OmniDiffusionConfig(model="test", tf_model_config=tf_config)
assert config.is_moe is False
@@ -10,7 +10,7 @@ stage_args:
max_batch_size: 1
engine_args:
model_stage: thinker
model_arch: OmniBagelForConditionalGeneration
model_arch: BagelForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
gpu_memory_utilization: 0.45
@@ -1,184 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
End-to-end test for Bagel img2img generation.
This test validates that the Bagel model generates images from an input image
and text prompt that match expected reference pixel values within a ±5 tolerance.
Equivalent to running:
python3 examples/offline_inference/bagel/end2end.py \
--prompts "Change the grass color to red" \
--modality img2img --step 15 \
--image-path 2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg
"""
import socket
from pathlib import Path
from typing import Any
import pytest
from PIL import Image
from vllm.assets.image import ImageAsset
from tests.utils import hardware_test
from vllm_omni.entrypoints.omni import Omni
# Reference pixel data extracted from the known-good output image
# Generated with seed=52, num_inference_steps=15,
# prompt='Change the grass color to red',
# input image: 2560px-Gfp-wisconsin-madison-the-nature-boardwalk.jpg
REFERENCE_PIXELS = [
{"position": (100, 100), "rgb": (157, 172, 217)},
{"position": (400, 50), "rgb": (105, 144, 218)},
{"position": (700, 100), "rgb": (118, 159, 233)},
{"position": (150, 400), "rgb": (195, 34, 60)},
{"position": (512, 336), "rgb": (222, 214, 193)},
{"position": (700, 400), "rgb": (197, 15, 43)},
{"position": (100, 600), "rgb": (105, 13, 18)},
{"position": (400, 600), "rgb": (169, 33, 44)},
{"position": (700, 600), "rgb": (101, 86, 93)},
{"position": (256, 256), "rgb": (181, 202, 222)},
]
PIXEL_TOLERANCE = 5
DEFAULT_PROMPT = "<|fim_middle|><|im_start|>Change the grass color to red<|im_end|>"
EXPECTED_OUTPUT_SIZE = (1024, 672)
def _load_input_image() -> Image.Image:
"""Load the test input image via vllm's ImageAsset."""
return ImageAsset("2560px-Gfp-wisconsin-madison-the-nature-boardwalk").pil_image.convert("RGB")
def _find_free_port() -> int:
"""Find and return a free ephemeral port by binding to port 0."""
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.bind(("127.0.0.1", 0))
s.listen(1)
port = s.getsockname()[1]
return port
def _configure_sampling_params(omni: Omni, max_tokens: int = 1, num_inference_steps: int = 15) -> list:
"""Configure sampling parameters for Bagel img2img generation.
Args:
omni: The Omni instance to get default params from.
max_tokens: Maximum tokens for the first stage.
num_inference_steps: Number of inference steps for the diffusion stage.
Returns:
Configured sampling params list.
"""
params_list = omni.default_sampling_params_list
params_list[0].max_tokens = max_tokens # type: ignore
if len(params_list) > 1:
params_list[1].num_inference_steps = num_inference_steps # type: ignore
params_list[1].extra_args = { # type: ignore
"cfg_text_scale": 4.0,
"cfg_img_scale": 1.5,
}
return params_list
def _extract_generated_image(omni_outputs: list) -> Image.Image | None:
"""Extract the generated image from Omni outputs.
Args:
omni_outputs: List of outputs from omni.generate().
Returns:
The first generated PIL Image, or None if no image found.
"""
for req_output in omni_outputs:
if images := getattr(req_output, "images", None):
return images[0]
if hasattr(req_output, "request_output") and req_output.request_output:
for stage_out in req_output.request_output:
if hasattr(stage_out, "images") and stage_out.images:
return stage_out.images[0]
return None
def _validate_pixels(
image: Image.Image,
reference_pixels: list[dict[str, Any]] = REFERENCE_PIXELS,
tolerance: int = PIXEL_TOLERANCE,
) -> None:
"""Validate that image pixels match expected reference values.
Args:
image: The PIL Image to validate.
reference_pixels: List of dicts with 'position' (x, y) and 'rgb' (R, G, B).
tolerance: Maximum allowed difference per color channel.
Raises:
AssertionError: If any pixel differs beyond tolerance.
"""
for ref in reference_pixels:
x, y = ref["position"]
expected = ref["rgb"]
actual = image.getpixel((x, y))[:3]
assert all(abs(a - e) <= tolerance for a, e in zip(actual, expected)), (
f"Pixel mismatch at ({x}, {y}): expected {expected}, got {actual}"
)
def _generate_bagel_img2img(
omni: Omni,
input_image: Image.Image,
prompt: str = DEFAULT_PROMPT,
) -> Image.Image:
"""Generate an image using Bagel model with img2img pipeline.
Args:
omni: The Omni instance to use for generation.
input_image: The input PIL Image for img2img.
prompt: The text prompt for image editing.
Returns:
The generated PIL Image.
Raises:
AssertionError: If no image is generated or size is incorrect.
"""
params_list = _configure_sampling_params(omni)
omni_outputs = list(
omni.generate(
prompts=[
{
"prompt": prompt,
"multi_modal_data": {"img2img": input_image},
"modalities": ["img2img"],
}
],
sampling_params_list=params_list,
)
)
generated_image = _extract_generated_image(omni_outputs)
assert generated_image is not None, "No images generated"
assert generated_image.size == EXPECTED_OUTPUT_SIZE, f"Expected {EXPECTED_OUTPUT_SIZE}, got {generated_image.size}"
return generated_image
@pytest.mark.core_model
@pytest.mark.diffusion
@hardware_test(res={"cuda": "H100"})
def test_bagel_img2img_shared_memory_connector():
"""Test Bagel img2img with shared memory connector."""
input_image = _load_input_image()
config_path = str(Path(__file__).parent / "stage_configs" / "bagel_sharedmemory_ci.yaml")
omni = Omni(model="ByteDance-Seed/BAGEL-7B-MoT", stage_configs_path=config_path, stage_init_timeout=300)
try:
generated_image = _generate_bagel_img2img(omni, input_image)
_validate_pixels(generated_image)
finally:
omni.close()
@@ -14,6 +14,10 @@ Equivalent to running:
"""
import os
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
import signal
import socket
import subprocess
@@ -33,23 +37,23 @@ from vllm_omni.entrypoints.omni import Omni
# "Generated with seed=52, num_inference_steps=15,
# prompt='A futuristic city skyline at twilight, cyberpunk style'"
REFERENCE_PIXELS = [
{"position": (100, 100), "rgb": (121, 118, 100)},
{"position": (400, 50), "rgb": (163, 162, 143)},
{"position": (700, 100), "rgb": (170, 156, 127)},
{"position": (150, 400), "rgb": (129, 127, 112)},
{"position": (512, 512), "rgb": (135, 61, 59)},
{"position": (700, 400), "rgb": (205, 107, 43)},
{"position": (100, 700), "rgb": (197, 177, 157)},
{"position": (400, 700), "rgb": (139, 107, 86)},
{"position": (700, 700), "rgb": (247, 205, 146)},
{"position": (256, 256), "rgb": (171, 160, 153)},
{"position": (100, 100), "rgb": (49, 96, 134)},
{"position": (400, 50), "rgb": (63, 127, 167)},
{"position": (700, 100), "rgb": (70, 101, 141)},
{"position": (150, 400), "rgb": (115, 90, 150)},
{"position": (512, 512), "rgb": (98, 86, 119)},
{"position": (700, 400), "rgb": (29, 42, 91)},
{"position": (100, 700), "rgb": (47, 50, 88)},
{"position": (400, 700), "rgb": (36, 52, 91)},
{"position": (700, 700), "rgb": (45, 58, 99)},
{"position": (256, 256), "rgb": (62, 94, 135)},
]
# Maximum allowed difference per color channel
PIXEL_TOLERANCE = 5
# Default test prompt
DEFAULT_PROMPT = "<|im_start|>A cute cat<|im_end|>"
DEFAULT_PROMPT = "<|im_start|>A futuristic city skyline at twilight, cyberpunk style<|im_end|>"
def _find_free_port() -> int:
@@ -36,7 +36,7 @@ models = ["riverclouds/qwen_image_random"]
@pytest.mark.core_model
@pytest.mark.diffusion
@pytest.mark.cache
@hardware_test(res={"cuda": "L4", "rocm": "MI325", "xpu": "B60"})
@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
@pytest.mark.parametrize("model_name", models)
def test_cache_dit(model_name: str):
"""Test cache-dit backend with diffusion model."""
@@ -70,11 +70,6 @@ def test_cpu_offload_diffusion_model(model_name: str):
pytest.fail("Inference failed")
print(f"Offload peak memory: {offload_peak_memory} MB")
print(f"No offload peak memory: {no_offload_peak_memory} MB")
# Set platform-specific VRAM saving thresholds to account
# for varying runtime memory overhead and fragmentation between CUDA and ROCm.
is_rocm = torch.version.hip is not None
threshold = 2500 if not is_rocm else 2100
assert offload_peak_memory + threshold < no_offload_peak_memory, (
f"Offload peak memory {offload_peak_memory} MB should be less than "
f"no offload peak memory {no_offload_peak_memory} MB by {threshold} MB"
assert offload_peak_memory + 2500 < no_offload_peak_memory, (
f"Offload peak memory {offload_peak_memory} MB should be less than no offload peak memory {no_offload_peak_memory} MB"
)
@@ -39,9 +39,6 @@ if current_omni_platform.is_npu():
elif current_omni_platform.is_rocm():
# ROCm stage config optimized for MI325 GPU
stage_config = str(Path(__file__).parent.parent / "stage_configs" / "rocm" / "qwen2_5_omni_ci.yaml")
elif current_omni_platform.is_xpu():
# Intel XPU stage config optimized for B60 GPU
stage_config = str(Path(__file__).parent.parent / "stage_configs" / "xpu" / "qwen2_5_omni_ci.yaml")
else:
stage_config = get_cuda_graph_config()
@@ -59,7 +56,7 @@ def get_question(prompt_type="mix"):
@pytest.mark.core_model
@pytest.mark.omni
@hardware_test(res={"cuda": "L4", "rocm": "MI325", "xpu": "B60"}, num_cards={"cuda": 4, "rocm": 2, "xpu": 3})
@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2})
@pytest.mark.parametrize("omni_runner", test_params, indirect=True)
def test_mix_to_audio(omni_runner, omni_runner_handler) -> None:
"""
@@ -90,7 +87,7 @@ def test_mix_to_audio(omni_runner, omni_runner_handler) -> None:
@pytest.mark.core_model
@pytest.mark.omni
@hardware_test(res={"cuda": "L4", "rocm": "MI325", "xpu": "B60"}, num_cards={"cuda": 4, "rocm": 2, "xpu": 3})
@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2})
@pytest.mark.parametrize("omni_runner", test_params, indirect=True)
def test_text_to_text(omni_runner, omni_runner_handler) -> None:
"""
@@ -40,8 +40,6 @@ def get_cuda_graph_config():
if current_omni_platform.is_rocm():
# ROCm stage config optimized for MI325 GPU
stage_configs = [str(Path(__file__).parent.parent / "stage_configs" / "rocm" / "qwen3_omni_ci.yaml")]
elif current_omni_platform.is_xpu():
stage_configs = [str(Path(__file__).parent.parent / "stage_configs" / "xpu" / "qwen3_omni_ci.yaml")]
else:
stage_configs = [get_cuda_graph_config()]
@@ -146,6 +146,9 @@ def _run_inference(
# - warmup: whether to run warmup for this SP config
# - is_perf_test: whether this is a performance test (show speedup metrics)
SP_CONFIGS_L2 = [
# Ulysses-2 - performance test
(2, 1, DEFAULT_HEIGHT, DEFAULT_WIDTH, True, True),
(1, 2, DEFAULT_HEIGHT, DEFAULT_WIDTH, True, True), # Ring-2 - performance test
# Hybrid - correctness only
(2, 2, DEFAULT_HEIGHT, DEFAULT_WIDTH, False, False),
]
@@ -23,7 +23,7 @@ models = ["linyueqian/stable_audio_random"]
@pytest.mark.core_model
@pytest.mark.diffusion
@hardware_test(res={"cuda": "L4", "xpu": "B60"})
@hardware_test(res={"cuda": "L4"})
@pytest.mark.parametrize("model_name", models)
def test_stable_audio_model(model_name: str):
m = Omni(model=model_name)
@@ -36,7 +36,7 @@ elif current_omni_platform.is_rocm():
@pytest.mark.core_model
@pytest.mark.advanced_model
@pytest.mark.diffusion
@hardware_test(res={"cuda": "L4", "rocm": "MI325", "xpu": "B60"}, num_cards={"cuda": 1, "rocm": 2, "xpu": 2})
@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 1, "rocm": 2})
@pytest.mark.parametrize("model_name", models)
def test_diffusion_model(model_name: str, run_level):
if run_level == "core_model" and model_name != "riverclouds/qwen_image_random":
+1 -1
View File
@@ -36,7 +36,7 @@ models = ["riverclouds/qwen_image_random"]
@pytest.mark.core_model
@pytest.mark.diffusion
@pytest.mark.cache
@hardware_test(res={"cuda": "L4", "rocm": "MI325", "xpu": "B60"})
@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
@pytest.mark.parametrize("model_name", models)
def test_teacache(model_name: str):
"""Test TeaCache backend with diffusion model."""
@@ -35,6 +35,7 @@ if str(REPO_ROOT) not in sys.path:
sys.path.insert(0, str(REPO_ROOT))
# os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1"
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
PROMPT = "a photo of a cat sitting on a laptop keyboard"
@@ -1,247 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
End-to-end online serving test for Bagel text2img and img2img generation.
This test validates that the Bagel model can serve image generation requests
via the OpenAI-compatible chat completions API.
Equivalent to running:
vllm-omni serve "ByteDance-Seed/BAGEL-7B-MoT" --omni --port 8091
# text2img
python3 examples/online_serving/bagel/openai_chat_client.py \\
--prompt "A cute cat" --modality text2img
# img2img
python3 examples/online_serving/bagel/openai_chat_client.py \\
--prompt "Let the woman wear a blue dress" --modality img2img \\
--image-url women.jpg
"""
import base64
import os
import signal
import socket
import subprocess
import sys
import time
from io import BytesIO
from pathlib import Path
from typing import Any
import pytest
import requests
from PIL import Image
from vllm.assets.image import ImageAsset
from tests.utils import hardware_test
MODEL = "ByteDance-Seed/BAGEL-7B-MoT"
STAGE_CONFIGS_PATH = str(
Path(__file__).parent.parent / "offline_inference" / "stage_configs" / "bagel_sharedmemory_ci.yaml"
)
TEXT2IMG_PROMPT = "A cute cat"
IMG2IMG_PROMPT = "Change the grass color to red"
class BagelOmniServer:
"""Context manager to start/stop a vLLM-Omni server for Bagel model tests."""
def __init__(
self,
model: str = MODEL,
stage_configs_path: str = STAGE_CONFIGS_PATH,
env_dict: dict[str, str] | None = None,
) -> None:
self.model = model
self.stage_configs_path = stage_configs_path
self.env_dict = env_dict
self.proc: subprocess.Popen | None = None
self.host = "127.0.0.1"
self.port = _find_free_port()
@property
def base_url(self) -> str:
return f"http://{self.host}:{self.port}"
def _start_server(self) -> None:
env = os.environ.copy()
if self.env_dict is not None:
env.update(self.env_dict)
cmd = [
sys.executable,
"-m",
"vllm_omni.entrypoints.cli.main",
"serve",
self.model,
"--omni",
"--host",
self.host,
"--port",
str(self.port),
"--stage-configs-path",
self.stage_configs_path,
"--stage-init-timeout",
"300",
]
self.proc = subprocess.Popen(
cmd,
env=env,
start_new_session=True,
)
try:
if not _wait_for_port(self.host, self.port, timeout=600, proc=self.proc):
self.terminate()
raise RuntimeError(f"Server failed to start within 600 seconds on {self.host}:{self.port}")
except Exception:
self.terminate()
raise
def __enter__(self):
self._start_server()
return self
def terminate(self) -> None:
if self.proc:
try:
os.killpg(os.getpgid(self.proc.pid), signal.SIGTERM)
except ProcessLookupError:
pass
try:
self.proc.wait(timeout=30)
except subprocess.TimeoutExpired:
try:
os.killpg(os.getpgid(self.proc.pid), signal.SIGKILL)
except ProcessLookupError:
pass
self.proc.wait()
self.proc = None
def __exit__(self, exc_type, exc_val, exc_tb):
self.terminate()
def _find_free_port() -> int:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.bind(("127.0.0.1", 0))
s.listen(1)
return s.getsockname()[1]
def _wait_for_port(host: str, port: int, timeout: int = 600, proc: subprocess.Popen | None = None) -> bool:
start = time.time()
while time.time() - start < timeout:
if proc is not None and proc.poll() is not None:
# Server process exited early
return False
try:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
sock.settimeout(1)
if sock.connect_ex((host, port)) == 0:
return True
except Exception:
pass
time.sleep(2)
return False
def _send_chat_request(
server_url: str,
prompt: str,
*,
modality: str = "text2img",
image: Image.Image | None = None,
timeout: int = 300,
) -> dict[str, Any]:
"""Send a chat completion request matching the openai_chat_client.py format."""
content: list[dict[str, Any]] = [{"type": "text", "text": f"<|im_start|>{prompt}<|im_end|>"}]
if image is not None:
buffer = BytesIO()
image.save(buffer, format="JPEG")
b64_data = base64.b64encode(buffer.getvalue()).decode("utf-8")
content.append(
{
"type": "image_url",
"image_url": {"url": f"data:image/jpeg;base64,{b64_data}"},
}
)
payload: dict[str, Any] = {
"messages": [{"role": "user", "content": content}],
}
if modality in ("text2img", "img2img"):
payload["modalities"] = ["image"]
resp = requests.post(
f"{server_url}/v1/chat/completions",
headers={"Content-Type": "application/json"},
json=payload,
timeout=timeout,
)
resp.raise_for_status()
return resp.json()
def _extract_image_from_response(data: dict[str, Any]) -> Image.Image | None:
"""Extract the generated PIL Image from a chat completion response."""
for choice in data.get("choices", []):
content = choice.get("message", {}).get("content")
if isinstance(content, list) and content:
first_item = content[0]
if isinstance(first_item, dict) and "image_url" in first_item:
url = first_item["image_url"].get("url", "")
if url.startswith("data:image"):
_, b64 = url.split(",", 1)
return Image.open(BytesIO(base64.b64decode(b64)))
return None
@pytest.mark.core_model
@pytest.mark.diffusion
@hardware_test(res={"cuda": "H100"})
def test_bagel_text2img_online():
"""Test Bagel text2img via OpenAI-compatible chat completions API."""
with BagelOmniServer() as server:
response_data = _send_chat_request(
server.base_url,
TEXT2IMG_PROMPT,
modality="text2img",
)
image = _extract_image_from_response(response_data)
assert image is not None, f"No image in response: {response_data}"
image.load()
w, h = image.size
assert w > 0 and h > 0, f"Invalid image size: {image.size}"
@pytest.mark.core_model
@pytest.mark.diffusion
@hardware_test(res={"cuda": "H100"})
def test_bagel_img2img_online():
"""Test Bagel img2img via OpenAI-compatible chat completions API."""
input_image = ImageAsset("2560px-Gfp-wisconsin-madison-the-nature-boardwalk").pil_image.convert("RGB")
with BagelOmniServer() as server:
response_data = _send_chat_request(
server.base_url,
IMG2IMG_PROMPT,
modality="img2img",
image=input_image,
)
image = _extract_image_from_response(response_data)
assert image is not None, f"No image in response: {response_data}"
image.load()
w, h = image.size
assert w > 0 and h > 0, f"Invalid image size: {image.size}"
@@ -147,7 +147,7 @@ def _basic_payload() -> dict:
@pytest.mark.advanced_model
@pytest.mark.diffusion
@hardware_test(res={"cuda": "L4", "rocm": "MI325", "xpu": "B60"})
@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
def test_images_generations_per_request_lora_switching(omni_server: OmniServer, tmp_path: Path) -> None:
# Base generation.
base_img = _post_images(omni_server, _basic_payload())
+1 -2
View File
@@ -48,8 +48,6 @@ def get_chunk_config():
if current_omni_platform.is_rocm():
# ROCm stage config optimized for MI325 GPU
stage_configs = [str(Path(__file__).parent.parent / "stage_configs" / "rocm" / "qwen3_omni_ci.yaml")]
elif current_omni_platform.is_xpu():
stage_configs = [str(Path(__file__).parent.parent / "stage_configs" / "xpu" / "qwen3_omni_ci.yaml")]
else:
stage_configs = [get_chunk_config()]
@@ -120,6 +118,7 @@ def test_mix_to_text_audio_001(omni_server, openai_client) -> None:
"stream": True,
"key_words": {
"audio": ["test"],
"image": ["square", "quadrate"],
},
}
@@ -6,8 +6,6 @@ E2E Online tests for Qwen3-Omni model.
import os
from vllm_omni.platforms import current_omni_platform
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
from pathlib import Path
@@ -52,10 +50,6 @@ def get_chunk_config(default_path):
# CI stage config for 2*H100-80G GPUs
default_path = str(Path(__file__).parent.parent / "stage_configs" / "qwen3_omni_ci.yaml")
stage_configs = [default_path, get_chunk_config(default_path)]
if current_omni_platform.is_xpu():
stage_configs = [str(Path(__file__).parent.parent / "stage_configs" / "xpu" / "qwen3_omni_ci.yaml")]
# Create parameter combinations for model and stage config
test_params = [
OmniServerParams(model=model, stage_config_path=stage_config) for model in models for stage_config in stage_configs
@@ -1,146 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
E2E online tests for Qwen3-TTS WebSocket streaming speech.
"""
import asyncio
import json
import os
from pathlib import Path
import pytest
import websockets
from tests.conftest import OmniServer
from tests.utils import hardware_test
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
MODEL = "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice"
def get_stage_config() -> str:
return str(
Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / "qwen3_tts.yaml"
)
@pytest.fixture(scope="module")
def omni_server():
stage_config_path = get_stage_config()
with OmniServer(
MODEL,
[
"--stage-configs-path",
stage_config_path,
"--stage-init-timeout",
"120",
"--trust-remote-code",
"--enforce-eager",
"--disable-log-stats",
],
env_dict={"VLLM_DISABLE_COMPILE_CACHE": "1"},
) as server:
yield server
async def _run_ws_session(host: str, port: int) -> dict:
uri = f"ws://{host}:{port}/v1/audio/speech/stream"
starts: list[dict] = []
dones: list[dict] = []
chunk_lengths: dict[int, list[int]] = {}
session_done: dict | None = None
async with websockets.connect(uri, max_size=None) as ws:
await ws.send(
json.dumps(
{
"type": "session.config",
"model": MODEL,
"voice": "vivian",
"language": "English",
"response_format": "pcm",
"stream_audio": True,
}
)
)
await ws.send(
json.dumps(
{
"type": "input.text",
"text": (
"Hello, this is a websocket streaming test for Qwen three TTS, "
"and this sentence is intentionally long enough to produce audio chunks. "
"This is the second sentence."
),
}
)
)
await ws.send(json.dumps({"type": "input.done"}))
while True:
message = await asyncio.wait_for(ws.recv(), timeout=180)
if isinstance(message, bytes):
if not starts:
raise AssertionError("Received audio bytes before audio.start")
sentence_index = starts[-1]["sentence_index"]
chunk_lengths.setdefault(sentence_index, []).append(len(message))
continue
payload = json.loads(message)
msg_type = payload.get("type")
if msg_type == "audio.start":
starts.append(payload)
chunk_lengths.setdefault(payload["sentence_index"], [])
elif msg_type == "audio.done":
dones.append(payload)
elif msg_type == "session.done":
session_done = payload
break
elif msg_type == "error":
raise AssertionError(f"WebSocket error: {payload['message']}")
else:
raise AssertionError(f"Unexpected WebSocket message: {payload}")
return {
"starts": starts,
"dones": dones,
"chunk_lengths": chunk_lengths,
"session_done": session_done,
}
class TestQwen3TTSWebSocket:
@pytest.mark.core_model
@pytest.mark.omni
@hardware_test(res={"cuda": "L4"}, num_cards=4)
def test_streaming_pcm_output(self, omni_server) -> None:
result = asyncio.run(_run_ws_session(omni_server.host, omni_server.port))
starts = result["starts"]
dones = result["dones"]
chunk_lengths = result["chunk_lengths"]
session_done = result["session_done"]
assert session_done is not None
assert session_done["total_sentences"] == 2
assert len(starts) == 2
assert len(dones) == 2
for idx, start in enumerate(starts):
assert start["type"] == "audio.start"
assert start["sentence_index"] == idx
assert start["format"] == "pcm"
assert start["sample_rate"] == 24000
assert start["sentence_text"]
for done in dones:
sentence_index = done["sentence_index"]
total_bytes = done["total_bytes"]
assert done["error"] is False
assert total_bytes > 0
assert chunk_lengths[sentence_index], f"Expected binary PCM frames for sentence {sentence_index}"
assert sum(chunk_lengths[sentence_index]) == total_bytes
@@ -1,260 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
E2E online serving tests for the async video generation API.
These tests exercise the real `/v1/videos` job lifecycle against a running
diffusion model, including create, retrieve, list, download, and delete flows.
"""
import os
import time
import uuid
from pathlib import Path
from typing import Any
import pytest
import requests
from tests.conftest import OmniServer
from tests.utils import hardware_test
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
MODEL = "Wan-AI/Wan2.2-T2V-A14B-Diffusers"
VIDEO_POLL_INTERVAL_S = 2.0
VIDEO_TIMEOUT_S = 900.0
@pytest.fixture(scope="module")
def video_server():
"""Start a real video generation server for async video API tests."""
with OmniServer(
MODEL,
[
"--num-gpus",
"1",
"--boundary-ratio",
"0.875",
"--flow-shift",
"5.0",
"--disable-log-stats",
],
) as server:
yield server
def _video_api_url(server: OmniServer, suffix: str = "") -> str:
return f"http://{server.host}:{server.port}/v1/videos{suffix}"
def _multipart_fields(payload: dict[str, Any]) -> list[tuple[str, tuple[None, str]]]:
return [(key, (None, str(value))) for key, value in payload.items() if value is not None]
def _create_video_job(
server: OmniServer,
*,
prompt: str | None = None,
**overrides: Any,
) -> requests.Response:
payload: dict[str, Any] = {
"prompt": prompt or f"video api e2e {uuid.uuid4().hex[:8]}",
"width": 640,
"height": 480,
"num_frames": 5,
"fps": 8,
"num_inference_steps": 2,
"guidance_scale": 1.0,
"guidance_scale_2": 1.0,
"boundary_ratio": 0.875,
"flow_shift": 5.0,
"seed": 42,
}
payload.update(overrides)
return requests.post(
_video_api_url(server),
files=_multipart_fields(payload),
timeout=VIDEO_TIMEOUT_S,
)
def _retrieve_video_job(server: OmniServer, video_id: str) -> requests.Response:
return requests.get(_video_api_url(server, f"/{video_id}"), timeout=VIDEO_TIMEOUT_S)
def _wait_for_video_status(server: OmniServer, video_id: str, expected_status: str) -> dict[str, Any]:
deadline = time.time() + VIDEO_TIMEOUT_S
last_payload: dict[str, Any] | None = None
while time.time() < deadline:
response = _retrieve_video_job(server, video_id)
assert response.status_code == 200, response.text
last_payload = response.json()
status = last_payload["status"]
if status == expected_status:
return last_payload
if status == "failed":
raise AssertionError(f"Video job {video_id} failed unexpectedly: {last_payload}")
time.sleep(VIDEO_POLL_INTERVAL_S)
raise AssertionError(
f"Timed out waiting for video job {video_id} to reach status={expected_status}. Last payload: {last_payload}"
)
def _wait_for_video_missing(server: OmniServer, video_id: str) -> None:
deadline = time.time() + 30.0
last_status: int | None = None
while time.time() < deadline:
response = _retrieve_video_job(server, video_id)
last_status = response.status_code
if response.status_code == 404:
return
time.sleep(0.5)
raise AssertionError(f"Timed out waiting for video job {video_id} to disappear. Last status={last_status}")
def _delete_video_job(server: OmniServer, video_id: str) -> requests.Response:
return requests.delete(_video_api_url(server, f"/{video_id}"), timeout=VIDEO_TIMEOUT_S)
def _delete_video_job_with_retry(server: OmniServer, video_id: str) -> requests.Response:
deadline = time.time() + 30.0
last_response: requests.Response | None = None
while time.time() < deadline:
response = _delete_video_job(server, video_id)
last_response = response
if response.status_code != 409:
return response
time.sleep(1.0)
raise AssertionError(
f"Timed out waiting to delete video job {video_id}. "
f"Last response: {None if last_response is None else last_response.text}"
)
def _best_effort_delete(server: OmniServer, video_id: str) -> None:
try:
response = _delete_video_job_with_retry(server, video_id)
if response.status_code not in (200, 404):
print(f"Cleanup delete for {video_id} returned {response.status_code}: {response.text}")
except Exception as exc:
print(f"Cleanup delete for {video_id} failed: {exc}")
def _assert_mp4_payload(content: bytes) -> None:
assert len(content) > 32, f"Downloaded video payload is unexpectedly small: {len(content)} bytes"
assert content[4:8] == b"ftyp", "Downloaded payload does not look like an MP4 file."
@pytest.mark.core_model
@pytest.mark.diffusion
@hardware_test(res={"cuda": "H100", "rocm": "MI325"})
def test_create_list_retrieve_download_video(video_server: OmniServer, tmp_path: Path) -> None:
video_id: str | None = None
try:
create_response = _create_video_job(video_server, prompt="A small paper airplane gliding over a lake at dawn.")
assert create_response.status_code == 200, create_response.text
created = create_response.json()
video_id = created["id"]
assert created["object"] == "video"
assert created["status"] == "queued"
assert created["model"] == MODEL
completed = _wait_for_video_status(video_server, video_id, "completed")
assert completed["file_name"] is not None
assert completed["completed_at"] is not None
assert completed["progress"] == 100
retrieve_response = _retrieve_video_job(video_server, video_id)
assert retrieve_response.status_code == 200, retrieve_response.text
retrieved = retrieve_response.json()
assert retrieved["id"] == video_id
assert retrieved["status"] == "completed"
list_response = requests.get(_video_api_url(video_server), timeout=VIDEO_TIMEOUT_S)
assert list_response.status_code == 200, list_response.text
listed = list_response.json()
assert listed["object"] == "list"
assert any(item["id"] == video_id for item in listed["data"])
download_response = requests.get(
_video_api_url(video_server, f"/{video_id}/content"),
timeout=VIDEO_TIMEOUT_S,
)
assert download_response.status_code == 200, download_response.text
assert download_response.headers["content-type"].startswith("video/mp4")
assert completed["file_name"] in download_response.headers.get("content-disposition", "")
_assert_mp4_payload(download_response.content)
output_path = tmp_path / completed["file_name"]
output_path.write_bytes(download_response.content)
assert output_path.stat().st_size == len(download_response.content)
finally:
if video_id is not None:
_best_effort_delete(video_server, video_id)
@pytest.mark.core_model
@pytest.mark.diffusion
@hardware_test(res={"cuda": "H100", "rocm": "MI325"})
def test_delete_video_while_queued_or_in_progress(video_server: OmniServer) -> None:
create_response = _create_video_job(
video_server,
prompt="A toy boat sailing through fog.",
num_frames=13,
num_inference_steps=4,
)
assert create_response.status_code == 200, create_response.text
created = create_response.json()
video_id = created["id"]
assert created["status"] == "queued"
# Delete immediately so the job is still queued or just entering execution.
delete_response = _delete_video_job_with_retry(video_server, video_id)
assert delete_response.status_code == 200, delete_response.text
deleted = delete_response.json()
assert deleted["id"] == video_id
assert deleted["deleted"] is True
assert deleted["object"] == "video.deleted"
_wait_for_video_missing(video_server, video_id)
@pytest.mark.core_model
@pytest.mark.diffusion
@hardware_test(res={"cuda": "H100", "rocm": "MI325"})
def test_delete_video_after_processing(video_server: OmniServer) -> None:
create_response = _create_video_job(
video_server,
prompt="A lantern drifting above a quiet forest.",
seed=43,
)
assert create_response.status_code == 200, create_response.text
video_id = create_response.json()["id"]
completed = _wait_for_video_status(video_server, video_id, "completed")
assert completed["file_name"] is not None
delete_response = _delete_video_job_with_retry(video_server, video_id)
assert delete_response.status_code == 200, delete_response.text
deleted = delete_response.json()
assert deleted["id"] == video_id
assert deleted["deleted"] is True
assert deleted["object"] == "video.deleted"
_wait_for_video_missing(video_server, video_id)
download_response = requests.get(
_video_api_url(video_server, f"/{video_id}/content"),
timeout=VIDEO_TIMEOUT_S,
)
assert download_response.status_code == 404, download_response.text
@@ -1,110 +0,0 @@
# stage config for running qwen2.5-omni with architecture of OmniLLM.
# The following config is verified with 2 * Intel Arc Pro B60 XPU.
stage_args:
- stage_id: 0
stage_type: llm # Use llm stage type to launch OmniLLM
runtime:
process: true # Run this stage in a separate process
devices: "0" # Visible devices for this stage
max_batch_size: 1
engine_args:
model_stage: thinker
model_arch: Qwen2_5OmniForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
max_model_len: 16384
max_num_batched_tokens: 16384
max_num_seqs: 1
gpu_memory_utilization: 0.9 # thinker weight is around 16.74GB for Qwen2.5-Omni-7B
skip_mm_profiling: true
enforce_eager: true
trust_remote_code: true
engine_output_type: latent
enable_prefix_caching: false
mm_processor_cache_gb: 0
is_comprehension: true
final_output: true
final_output_type: text
default_sampling_params:
temperature: 0.0
top_p: 1.0
top_k: -1
max_tokens: 128
seed: 42
detokenize: True
repetition_penalty: 1.1
- stage_id: 1
stage_type: llm # Use llm stage type to launch OmniLLM
runtime:
process: true
devices: "1"
max_batch_size: 1
engine_args:
model_stage: talker
model_arch: Qwen2_5OmniForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
max_model_len: 16384
max_num_batched_tokens: 16384
max_num_seqs: 1
gpu_memory_utilization: 0.5 # talker weight is 6.03GB for Qwen2.5-Omni-7B
skip_mm_profiling: true
enforce_eager: true
trust_remote_code: true
enable_prefix_caching: false
engine_output_type: latent
engine_input_source: [0]
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen2_5_omni.thinker2talker
default_sampling_params:
temperature: 0.9
top_p: 0.8
top_k: 40
max_tokens: 4096
seed: 42
detokenize: True
repetition_penalty: 1.05
stop_token_ids: [8294]
- stage_id: 2
stage_type: llm # Use llm stage type to launch OmniLLM
runtime:
process: true
devices: "2"
max_batch_size: 1
engine_args:
model_stage: code2wav
model_arch: Qwen2_5OmniForConditionalGeneration
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
gpu_memory_utilization: 0.3 # code2wav weight is around 1.46GB for Qwen2.5-Omni-7B
enforce_eager: true
trust_remote_code: true
enable_prefix_caching: false
engine_output_type: audio
engine_input_source: [1]
final_output: true
final_output_type: audio
default_sampling_params:
temperature: 0.0
top_p: 1.0
top_k: -1
max_tokens: 2048
seed: 42
detokenize: True
repetition_penalty: 1.1
# Top-level runtime config (concise): default windows and stage edges
runtime:
enabled: true
defaults:
window_size: -1 # Simplified: trigger downstream only after full upstream completion
max_inflight: 1 # Simplified: process serially within each stage
edges:
- from: 0 # thinker → talker: trigger only after receiving full input (-1)
to: 1
window_size: -1
- from: 1 # talker → code2wav: trigger only after receiving full input (-1)
to: 2
window_size: -1
@@ -1,109 +0,0 @@
# Stage config for running Qwen3-Omni-MoE with 3-stage architecture
# Stage 0: Thinker (multimodal understanding + text generation)
# Stage 1: Talker (text embeddings → 8-layer RVQ codec codes)
# Stage 2: Code2Wav (8-layer RVQ codes → audio waveform)
# The following config is verified with 8 * Intel Arc Pro B60 XPU.
stage_args:
- stage_id: 0
stage_type: llm # Use llm stage type to launch OmniLLM
runtime:
devices: "0,1,2,3"
max_batch_size: 1
engine_args:
model_stage: thinker
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
gpu_memory_utilization: 0.85 # thinker weight is around 61.08GB for Qwen3-Omni-30B-A3B-Instruct
skip_mm_profiling: true
enforce_eager: true
trust_remote_code: true
engine_output_type: latent # Output hidden states for talker
distributed_executor_backend: "mp"
max_num_batched_tokens: 4096
max_model_len: 4096
enable_prefix_caching: false
hf_config_name: thinker_config
tensor_parallel_size: 4
max_cudagraph_capture_size: 0
final_output: true
final_output_type: text
is_comprehension: true
default_sampling_params:
temperature: 0.4
top_p: 0.9
top_k: 1
max_tokens: 100
seed: 42
ignore_eos: False
detokenize: True
repetition_penalty: 1.05
- stage_id: 1
stage_type: llm # Use llm stage type to launch OmniLLM
runtime:
devices: "4"
max_batch_size: 1
engine_args:
model_stage: talker
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
gpu_memory_utilization: 0.6 # talker weight is around 8.5GB for Qwen3-Omni-30B-A3B-Instruct
skip_mm_profiling: true
enforce_eager: true
trust_remote_code: true
engine_output_type: latent # Output codec codes for code2wav
enable_prefix_caching: false
max_num_batched_tokens: 4096
max_model_len: 4096
distributed_executor_backend: "mp"
hf_config_name: talker_config
max_cudagraph_capture_size: 0
engine_input_source: [0]
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker
# final_output: true
# final_output_type: text
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: 2
stage_type: llm # Use llm stage type to launch OmniLLM
runtime:
devices: "5"
max_batch_size: 1
engine_args:
model_stage: code2wav
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
enforce_eager: true
trust_remote_code: true
enable_prefix_caching: false
engine_output_type: audio # Final output: audio waveform
gpu_memory_utilization: 0.3 # code2wav weight is around 0.4GB for Qwen3-Omni-30B-A3B-Instruct
skip_mm_profiling: true
distributed_executor_backend: "mp"
max_num_batched_tokens: 100000
hf_config_name: thinker_config
async_scheduling: false
max_cudagraph_capture_size: 0
engine_input_source: [1]
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
final_output: true
final_output_type: audio
default_sampling_params:
temperature: 0.0
top_p: 1.0
top_k: -1
max_tokens: 2000
seed: 42
detokenize: True
repetition_penalty: 1.1
@@ -1,29 +1,23 @@
# tests/entrypoints/openai/test_serving_speech.py
import asyncio
import logging
import os
import struct
from inspect import Signature, signature
from pathlib import Path
from unittest.mock import AsyncMock, MagicMock, patch
from unittest.mock import MagicMock, patch
import numpy as np
import pytest
import torch
from fastapi import FastAPI, HTTPException, Request, UploadFile
from fastapi import FastAPI, HTTPException, UploadFile
from fastapi.params import File, Form
from fastapi.responses import JSONResponse
from fastapi.testclient import TestClient
from pydantic import ValidationError
from pytest_mock import MockerFixture
from vllm.entrypoints.openai.engine.protocol import ErrorInfo, ErrorResponse
from vllm_omni.entrypoints.openai import api_server as api_server_module
from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin
from vllm_omni.entrypoints.openai.protocol.audio import CreateAudio, OpenAICreateSpeechRequest
from vllm_omni.entrypoints.openai.serving_speech import (
OmniOpenAIServingSpeech,
_create_wav_header,
)
from vllm_omni.outputs import OmniRequestOutput
@@ -922,20 +916,17 @@ class TestStreamingProtocolValidation:
"""Unit tests for the stream field validators in OpenAICreateSpeechRequest."""
def test_stream_validation_errors(self):
"""stream=True requires response_format not in ('pcm', 'wav') and speed=1.0."""
with pytest.raises(ValidationError, match="requires response_format not in \\('pcm', 'wav'\\)"):
OpenAICreateSpeechRequest(input="Hello", stream=True, response_format="mp3")
"""stream=True requires response_format='pcm' and speed=1.0."""
with pytest.raises(ValidationError, match="response_format='pcm'"):
OpenAICreateSpeechRequest(input="Hello", stream=True, response_format="wav")
with pytest.raises(ValidationError, match="Speed adjustment is not supported"):
OpenAICreateSpeechRequest(input="Hello", stream=True, response_format="pcm", speed=2.0)
def test_stream_valid(self):
"""stream=True + response_format in ('pcm', 'wav') + speed=1.0 is accepted."""
"""stream=True + response_format='pcm' + speed=1.0 is accepted."""
req = OpenAICreateSpeechRequest(input="Hello", stream=True, response_format="pcm")
assert req.stream is True
req = OpenAICreateSpeechRequest(input="Hello", stream=True, response_format="wav")
assert req.stream is True
def test_sse_stream_format_is_blocked(self):
"""stream_format='sse' is blocked."""
with pytest.raises(ValidationError, match="sse"):
@@ -1034,248 +1025,19 @@ class TestAsyncOmniSupportedTasks:
@pytest.mark.asyncio
async def test_tts_only_no_generate_task(self):
"""TTS-only models (audio output, no text) should not include 'generate'."""
from unittest.mock import MagicMock
from vllm_omni.entrypoints.async_omni import AsyncOmni
omni = AsyncOmni.__new__(AsyncOmni)
omni.output_modalities = [None, "audio"]
stage = MagicMock()
stage.is_comprehension = False
omni.stage_list = [stage]
tasks = await omni.get_supported_tasks()
assert "generate" not in tasks
assert "speech" in tasks
@pytest.mark.asyncio
async def test_omni_model_includes_generate(self):
"""Models with text output (e.g. Qwen3-Omni) should include 'generate'."""
from unittest.mock import MagicMock
from vllm_omni.entrypoints.async_omni import AsyncOmni
omni = AsyncOmni.__new__(AsyncOmni)
omni.output_modalities = ["text", None, "audio"]
stage = MagicMock()
stage.is_comprehension = True
omni.stage_list = [stage]
tasks = await omni.get_supported_tasks()
assert "generate" in tasks
def test_api_server_create_speech_wraps_error_response_status():
handler = MagicMock()
handler.create_speech = AsyncMock(
return_value=ErrorResponse(
error=ErrorInfo(message="bad request", type="BadRequestError", param=None, code=400),
)
)
app = FastAPI()
app.state.openai_serving_speech = handler
scope = {
"type": "http",
"app": app,
"method": "POST",
"path": "/v1/audio/speech",
"headers": [],
"query_string": b"",
"client": ("127.0.0.1", 12345),
"server": ("testserver", 80),
"scheme": "http",
}
raw_request = Request(scope)
request = OpenAICreateSpeechRequest(input="Hello")
response = asyncio.run(api_server_module.create_speech(request, raw_request))
assert isinstance(response, JSONResponse)
assert response.status_code == 400
class TestWAVHeaderGeneration:
"""Unit tests for WAV header generation with placeholder values."""
def test_wav_header_basic_structure(self):
"""Test basic WAV header structure with default parameters."""
header = _create_wav_header(sample_rate=24000, num_channels=1, bits_per_sample=16)
# Verify header length (should be 44 bytes)
assert len(header) == 44, f"Expected 44 bytes, got {len(header)}"
# Parse and verify header structure
(
chunk_id,
chunk_size,
format_type,
subchunk1_id,
subchunk1_size,
audio_format,
num_channels,
sample_rate,
byte_rate,
block_align,
bits_per_sample,
subchunk2_id,
subchunk2_size,
) = struct.unpack("<4sI4s4sIHHIIHH4sI", header)
# Verify RIFF header
assert chunk_id == b"RIFF", f"Expected RIFF, got {chunk_id}"
assert chunk_size == 0xFFFFFFFF, f"Expected placeholder 0xFFFFFFFF, got {chunk_size:#x}"
assert format_type == b"WAVE", f"Expected WAVE, got {format_type}"
# Verify fmt chunk
assert subchunk1_id == b"fmt ", f"Expected 'fmt ', got {subchunk1_id}"
assert subchunk1_size == 16, f"Expected 16, got {subchunk1_size}"
assert audio_format == 1, f"Expected PCM (1), got {audio_format}"
assert num_channels == 1, f"Expected 1 channel, got {num_channels}"
assert sample_rate == 24000, f"Expected 24000 Hz, got {sample_rate}"
assert byte_rate == 48000, f"Expected 48000 byte/s, got {byte_rate}"
assert block_align == 2, f"Expected 2 bytes block align, got {block_align}"
assert bits_per_sample == 16, f"Expected 16 bits, got {bits_per_sample}"
# Verify data chunk
assert subchunk2_id == b"data", f"Expected 'data', got {subchunk2_id}"
assert subchunk2_size == 0xFFFFFFFF, f"Expected placeholder 0xFFFFFFFF, got {subchunk2_size:#x}"
def test_wav_header_different_sample_rates(self):
"""Test WAV header with different sample rates."""
test_cases = [
(16000, 1, 16),
(22050, 1, 16),
(24000, 1, 16),
(44100, 1, 16),
(48000, 1, 16),
]
for sample_rate, num_channels, bits_per_sample in test_cases:
header = _create_wav_header(sample_rate, num_channels, bits_per_sample)
assert len(header) == 44, f"Header length mismatch for {sample_rate} Hz"
# Parse sample rate from header
parsed_sample_rate = struct.unpack("<I", header[24:28])[0]
assert parsed_sample_rate == sample_rate, (
f"Sample rate mismatch: expected {sample_rate}, got {parsed_sample_rate}"
)
def test_wav_header_stereo(self):
"""Test WAV header with stereo audio."""
header = _create_wav_header(sample_rate=44100, num_channels=2, bits_per_sample=16)
# Parse header
parsed = struct.unpack("<4sI4s4sIHHIIHH4sI", header)
num_channels = parsed[6]
byte_rate = parsed[8]
block_align = parsed[9]
assert num_channels == 2, f"Expected 2 channels, got {num_channels}"
assert byte_rate == 44100 * 2 * 16 // 8, "Byte rate mismatch"
assert block_align == 2 * 16 // 8, "Block align mismatch"
def test_wav_header_placeholder_values(self):
"""Test that placeholder values are correctly set to 0xFFFFFFFF."""
header = _create_wav_header(sample_rate=24000)
# Extract size fields
chunk_size = struct.unpack("<I", header[4:8])[0]
subchunk2_size = struct.unpack("<I", header[40:44])[0]
assert chunk_size == 0xFFFFFFFF, "ChunkSize should be 0xFFFFFFFF for streaming"
assert subchunk2_size == 0xFFFFFFFF, "Subchunk2Size should be 0xFFFFFFFF for streaming"
class TestWAVStreaming:
"""Integration tests for WAV format streaming."""
@pytest.fixture
def wav_streaming_app(self, mocker: MockerFixture):
"""Test app configured for WAV streaming."""
mock_engine_client = mocker.MagicMock()
mock_engine_client.errored = False
def _make_output(finished: bool) -> OmniRequestOutput:
chunk = torch.sin(torch.linspace(0, 440 * 2 * torch.pi, 24000))
class MockCompletionOutput:
def __init__(self, index: int = 0):
self.index = index
self.text = ""
self.token_ids = []
self.finish_reason = "stop"
self.stop_reason = None
self.logprobs = None
class MockRequestOutput:
def __init__(self, audio_tensor: torch.Tensor):
self.request_id = "speech-wav-stream-test"
self.outputs = [MockCompletionOutput(index=0)]
self.multimodal_output = {"audio": audio_tensor, "sr": 24000}
self.finished = finished
self.prompt_token_ids = None
self.encoder_prompt_token_ids = None
self.num_cached_tokens = None
self.prompt_logprobs = None
self.kv_transfer_params = None
return OmniRequestOutput(
stage_id=0,
final_output_type="audio",
request_output=MockRequestOutput(audio_tensor=chunk),
finished=finished,
)
async def mock_generate_streaming(*args, **kwargs):
yield _make_output(finished=False)
yield _make_output(finished=True)
mock_engine_client.generate = mocker.MagicMock(side_effect=mock_generate_streaming)
mock_engine_client.default_sampling_params_list = [{}]
mock_models = mocker.MagicMock()
mock_models.is_base_model.return_value = True
speech_server = OmniOpenAIServingSpeech(
engine_client=mock_engine_client,
models=mock_models,
request_logger=mocker.MagicMock(),
)
original_create_speech = speech_server.create_speech
sig = signature(original_create_speech)
new_parameters = [p for name, p in sig.parameters.items() if name != "raw_request"]
new_sig = Signature(parameters=new_parameters, return_annotation=sig.return_annotation)
async def awaitable_create_speech(*args, **kwargs):
return await original_create_speech(*args, **kwargs)
awaitable_create_speech.__signature__ = new_sig
speech_server.create_speech = awaitable_create_speech
app = FastAPI()
app.add_api_route("/v1/audio/speech", speech_server.create_speech, methods=["POST"], response_model=None)
return app
def test_wav_streaming_success(self, wav_streaming_app):
"""Test WAV format streaming returns correct content type and includes WAV header."""
client = TestClient(wav_streaming_app)
response = client.post("/v1/audio/speech", json={"input": "Hello", "stream": True, "response_format": "wav"})
assert response.status_code == 200
assert "audio/wav" in response.headers["content-type"]
assert len(response.content) > 44 # Should have WAV header + audio data
# Verify WAV header is present
header = response.content[:44]
chunk_id = header[0:4]
format_type = header[8:12]
assert chunk_id == b"RIFF", "Should start with RIFF"
assert format_type == b"WAVE", "Should contain WAVE format"
def test_streaming_unsupported_format_rejected(self, wav_streaming_app):
"""Test that unsupported formats are rejected for streaming."""
client = TestClient(wav_streaming_app)
unsupported_formats = ["mp3"]
for fmt in unsupported_formats:
response = client.post("/v1/audio/speech", json={"input": "Hello", "stream": True, "response_format": fmt})
assert response.status_code == 422
@@ -1,387 +0,0 @@
import asyncio
from unittest.mock import AsyncMock, MagicMock
import pytest
from fastapi import FastAPI, WebSocket
from starlette.testclient import TestClient
from starlette.websockets import WebSocketDisconnect
from vllm_omni.entrypoints.openai import serving_speech_stream as streaming_speech_module
from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
from vllm_omni.entrypoints.openai.serving_speech_stream import OmniStreamingSpeechHandler
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def _build_test_app(speech_service=None, *, idle_timeout=30.0, config_timeout=10.0):
if speech_service is None:
speech_service = MagicMock(spec=OmniOpenAIServingSpeech)
speech_service._generate_audio_bytes = AsyncMock(return_value=(b"RIFF" + b"\x00" * 32, "audio/wav"))
speech_service._prepare_speech_generation = AsyncMock(return_value=("req-1", object(), {}))
async def mock_generate_pcm_chunks(_generator, _request_id):
for chunk in (b"\x01\x02", b"\x03\x04\x05"):
yield chunk
speech_service._generate_pcm_chunks = mock_generate_pcm_chunks
speech_service.engine_client = MagicMock()
speech_service.engine_client.abort = AsyncMock()
handler = OmniStreamingSpeechHandler(
speech_service=speech_service,
idle_timeout=idle_timeout,
config_timeout=config_timeout,
)
app = FastAPI()
@app.websocket("/v1/audio/speech/stream")
async def ws_endpoint(websocket: WebSocket):
await handler.handle_session(websocket)
return app, speech_service
class TestStreamingSpeechWebSocket:
def test_non_streaming_single_frame(self):
app, speech_service = _build_test_app()
with TestClient(app) as client:
with client.websocket_connect("/v1/audio/speech/stream") as ws:
ws.send_json({"type": "session.config", "voice": "Vivian"})
ws.send_json({"type": "input.text", "text": "Hello world. "})
start = ws.receive_json()
assert start["type"] == "audio.start"
assert start["sentence_index"] == 0
assert start["sentence_text"] == "Hello world."
assert start["format"] == "wav"
audio = ws.receive_bytes()
assert audio.startswith(b"RIFF")
done = ws.receive_json()
assert done == {"type": "audio.done", "sentence_index": 0, "total_bytes": len(audio), "error": False}
ws.send_json({"type": "input.done"})
session_done = ws.receive_json()
assert session_done == {"type": "session.done", "total_sentences": 1}
assert speech_service._generate_audio_bytes.await_count == 1
def test_streaming_multiple_binary_frames(self):
captured_requests = []
speech_service = MagicMock(spec=OmniOpenAIServingSpeech)
speech_service._generate_audio_bytes = AsyncMock(return_value=(b"", "audio/wav"))
speech_service.engine_client = MagicMock()
speech_service.engine_client.abort = AsyncMock()
async def mock_prepare_speech_generation(request):
captured_requests.append(request)
return "req-stream", object(), {}
speech_service._prepare_speech_generation = mock_prepare_speech_generation
async def mock_generate_pcm_chunks(_generator, _request_id):
for chunk in (b"\x01\x02", b"\x03\x04\x05", b"\x06"):
yield chunk
speech_service._generate_pcm_chunks = mock_generate_pcm_chunks
app, _ = _build_test_app(speech_service)
with TestClient(app) as client:
with client.websocket_connect("/v1/audio/speech/stream") as ws:
ws.send_json(
{
"type": "session.config",
"voice": "Vivian",
"stream_audio": True,
"response_format": "pcm",
"initial_codec_chunk_frames": 12,
}
)
ws.send_json({"type": "input.text", "text": "Hello world. "})
start = ws.receive_json()
assert start["type"] == "audio.start"
assert start["format"] == "pcm"
assert start["sample_rate"] == 24000
assert ws.receive_bytes() == b"\x01\x02"
assert ws.receive_bytes() == b"\x03\x04\x05"
assert ws.receive_bytes() == b"\x06"
done = ws.receive_json()
assert done == {"type": "audio.done", "sentence_index": 0, "total_bytes": 6, "error": False}
ws.send_json({"type": "input.done"})
assert ws.receive_json() == {"type": "session.done", "total_sentences": 1}
assert len(captured_requests) == 1
assert captured_requests[0].stream is True
assert captured_requests[0].response_format == "pcm"
assert captured_requests[0].initial_codec_chunk_frames == 12
assert speech_service._generate_audio_bytes.await_count == 0
def test_flush_on_input_done(self):
app, _ = _build_test_app()
with TestClient(app) as client:
with client.websocket_connect("/v1/audio/speech/stream") as ws:
ws.send_json({"type": "session.config", "voice": "Vivian"})
ws.send_json({"type": "input.text", "text": "Hello world without punctuation"})
ws.send_json({"type": "input.done"})
assert ws.receive_json()["type"] == "audio.start"
assert ws.receive_bytes()
assert ws.receive_json() == {
"type": "audio.done",
"sentence_index": 0,
"total_bytes": 36,
"error": False,
}
assert ws.receive_json() == {"type": "session.done", "total_sentences": 1}
def test_invalid_streaming_config(self):
app, _ = _build_test_app()
with TestClient(app) as client:
with client.websocket_connect("/v1/audio/speech/stream") as ws:
ws.send_json(
{
"type": "session.config",
"voice": "Vivian",
"stream_audio": True,
"response_format": "wav",
}
)
error = ws.receive_json()
assert error["type"] == "error"
assert "response_format='pcm'" in error["message"]
def test_empty_input_text_emits_no_audio(self):
app, speech_service = _build_test_app()
with TestClient(app) as client:
with client.websocket_connect("/v1/audio/speech/stream") as ws:
ws.send_json({"type": "session.config", "voice": "Vivian"})
ws.send_json({"type": "input.text", "text": ""})
ws.send_json({"type": "input.done"})
assert ws.receive_json() == {"type": "session.done", "total_sentences": 0}
assert speech_service._generate_audio_bytes.await_count == 0
def test_multiple_sentences_increment_indices(self):
app, _ = _build_test_app()
with TestClient(app) as client:
with client.websocket_connect("/v1/audio/speech/stream") as ws:
ws.send_json({"type": "session.config", "voice": "Vivian"})
ws.send_json({"type": "input.text", "text": "First sentence. Second sentence. "})
first_start = ws.receive_json()
assert first_start["sentence_index"] == 0
ws.receive_bytes()
assert ws.receive_json() == {
"type": "audio.done",
"sentence_index": 0,
"total_bytes": 36,
"error": False,
}
second_start = ws.receive_json()
assert second_start["sentence_index"] == 1
ws.receive_bytes()
assert ws.receive_json() == {
"type": "audio.done",
"sentence_index": 1,
"total_bytes": 36,
"error": False,
}
ws.send_json({"type": "input.done"})
assert ws.receive_json() == {"type": "session.done", "total_sentences": 2}
def test_unknown_message_type_keeps_session_open(self):
app, _ = _build_test_app()
with TestClient(app) as client:
with client.websocket_connect("/v1/audio/speech/stream") as ws:
ws.send_json({"type": "session.config", "voice": "Vivian"})
ws.send_json({"type": "unknown"})
error = ws.receive_json()
assert error == {"type": "error", "message": "Unknown message type: unknown"}
ws.send_json({"type": "input.text", "text": "Hello world. "})
assert ws.receive_json()["type"] == "audio.start"
ws.receive_bytes()
assert ws.receive_json() == {
"type": "audio.done",
"sentence_index": 0,
"total_bytes": 36,
"error": False,
}
ws.send_json({"type": "input.done"})
assert ws.receive_json() == {"type": "session.done", "total_sentences": 1}
def test_config_timeout_closes_session(self):
app, _ = _build_test_app(config_timeout=0.01)
with TestClient(app) as client:
with client.websocket_connect("/v1/audio/speech/stream") as ws:
error = ws.receive_json()
assert error == {"type": "error", "message": "Timeout waiting for session.config"}
def test_generation_error_marks_audio_done(self):
speech_service = MagicMock(spec=OmniOpenAIServingSpeech)
speech_service._generate_audio_bytes = AsyncMock(side_effect=RuntimeError("boom"))
speech_service._prepare_speech_generation = AsyncMock(return_value=("req-err", object(), {}))
speech_service._generate_pcm_chunks = AsyncMock()
speech_service.engine_client = MagicMock()
speech_service.engine_client.abort = AsyncMock()
app, _ = _build_test_app(speech_service)
with TestClient(app) as client:
with client.websocket_connect("/v1/audio/speech/stream") as ws:
ws.send_json({"type": "session.config", "voice": "Vivian"})
ws.send_json({"type": "input.text", "text": "Hello world. "})
assert ws.receive_json()["type"] == "audio.start"
assert ws.receive_json() == {"type": "error", "message": "Generation failed for sentence 0: boom"}
assert ws.receive_json() == {"type": "audio.done", "sentence_index": 0, "total_bytes": 0, "error": True}
ws.send_json({"type": "input.done"})
assert ws.receive_json() == {"type": "session.done", "total_sentences": 1}
def test_streaming_generation_error_marks_audio_done(self):
speech_service = MagicMock(spec=OmniOpenAIServingSpeech)
speech_service._generate_audio_bytes = AsyncMock(return_value=(b"", "audio/wav"))
speech_service._prepare_speech_generation = AsyncMock(return_value=("req-stream-err", object(), {}))
speech_service.engine_client = MagicMock()
speech_service.engine_client.abort = AsyncMock()
async def mock_generate_pcm_chunks(_generator, _request_id):
yield b"\x01\x02"
raise RuntimeError("stream boom")
speech_service._generate_pcm_chunks = mock_generate_pcm_chunks
app, _ = _build_test_app(speech_service)
with TestClient(app) as client:
with client.websocket_connect("/v1/audio/speech/stream") as ws:
ws.send_json(
{
"type": "session.config",
"voice": "Vivian",
"stream_audio": True,
"response_format": "pcm",
}
)
ws.send_json({"type": "input.text", "text": "Hello world. "})
assert ws.receive_json()["type"] == "audio.start"
assert ws.receive_bytes() == b"\x01\x02"
assert ws.receive_json() == {
"type": "error",
"message": "Generation failed for sentence 0: stream boom",
}
assert ws.receive_json() == {
"type": "audio.done",
"sentence_index": 0,
"total_bytes": 2,
"error": True,
}
ws.send_json({"type": "input.done"})
assert ws.receive_json() == {"type": "session.done", "total_sentences": 1}
def test_invalid_input_text_type_returns_validation_error(self):
app, speech_service = _build_test_app()
with TestClient(app) as client:
with client.websocket_connect("/v1/audio/speech/stream") as ws:
ws.send_json({"type": "session.config", "voice": "Vivian"})
ws.send_json({"type": "input.text", "text": 123})
assert ws.receive_json() == {
"type": "error",
"message": "input.text requires a string value",
}
ws.send_json({"type": "input.done"})
assert ws.receive_json() == {"type": "session.done", "total_sentences": 0}
assert speech_service._generate_audio_bytes.await_count == 0
def test_input_text_message_too_large(self, monkeypatch):
monkeypatch.setattr(streaming_speech_module, "_MAX_INPUT_TEXT_MESSAGE_SIZE", 32)
app, speech_service = _build_test_app()
with TestClient(app) as client:
with client.websocket_connect("/v1/audio/speech/stream") as ws:
ws.send_json({"type": "session.config", "voice": "Vivian"})
ws.send_json({"type": "input.text", "text": "x" * 128})
assert ws.receive_json() == {
"type": "error",
"message": "input.text message too large",
}
ws.send_json({"type": "input.done"})
assert ws.receive_json() == {"type": "session.done", "total_sentences": 0}
assert speech_service._generate_audio_bytes.await_count == 0
def test_session_config_message_too_large(self, monkeypatch):
monkeypatch.setattr(streaming_speech_module, "_MAX_CONFIG_MESSAGE_SIZE", 64)
app, _ = _build_test_app()
with TestClient(app) as client:
with client.websocket_connect("/v1/audio/speech/stream") as ws:
ws.send_json({"type": "session.config", "voice": "Vivian", "ref_audio": "x" * 512})
assert ws.receive_json() == {
"type": "error",
"message": "session.config message too large",
}
def test_disconnect_aborts_streaming_request(self):
speech_service = MagicMock(spec=OmniOpenAIServingSpeech)
speech_service._generate_audio_bytes = AsyncMock(return_value=(b"", "audio/wav"))
speech_service._prepare_speech_generation = AsyncMock(return_value=("req-abort", object(), {}))
speech_service.engine_client = MagicMock()
speech_service.engine_client.abort = AsyncMock()
async def mock_generate_pcm_chunks(_generator, _request_id):
yield b"\x01\x02"
speech_service._generate_pcm_chunks = mock_generate_pcm_chunks
handler = OmniStreamingSpeechHandler(speech_service=speech_service)
websocket = MagicMock()
websocket.send_json = AsyncMock(side_effect=[None, WebSocketDisconnect()])
websocket.send_bytes = AsyncMock(side_effect=WebSocketDisconnect())
config = MagicMock()
config.model = None
config.voice = "Vivian"
config.task_type = None
config.language = None
config.instructions = None
config.response_format = "pcm"
config.speed = 1.0
config.max_new_tokens = None
config.initial_codec_chunk_frames = None
config.ref_audio = None
config.ref_text = None
config.x_vector_only_mode = None
config.stream_audio = True
with pytest.raises(WebSocketDisconnect):
asyncio.run(handler._generate_and_send(websocket, config, "Hello world.", 0))
speech_service.engine_client.abort.assert_awaited_once_with("req-abort")
assert websocket.send_json.await_count == 2
@@ -1,253 +0,0 @@
"""Tests for SentenceSplitter used in streaming TTS input."""
import pytest
from vllm_omni.entrypoints.openai.text_splitter import SentenceSplitter
pytestmark = [pytest.mark.openai, pytest.mark.speech]
class TestSentenceSplitterEnglish:
"""Tests for English sentence splitting."""
def test_single_sentence_no_boundary(self):
splitter = SentenceSplitter()
result = splitter.add_text("Hello world")
assert result == []
assert splitter.buffer == "Hello world"
def test_single_sentence_with_boundary(self):
splitter = SentenceSplitter()
result = splitter.add_text("Hello world. How are you?")
assert len(result) == 1
assert result[0] == "Hello world."
def test_multiple_sentences(self):
splitter = SentenceSplitter()
result = splitter.add_text("Hello. How are you? I am fine! ")
assert len(result) == 3
assert result[0] == "Hello."
assert result[1] == "How are you?"
assert result[2] == "I am fine!"
def test_exclamation_mark(self):
splitter = SentenceSplitter()
result = splitter.add_text("Wow, that is great! Tell me more.")
assert len(result) == 1
assert result[0] == "Wow, that is great!"
def test_question_mark(self):
splitter = SentenceSplitter()
result = splitter.add_text("Can you hear me? I hope so.")
assert len(result) == 1
assert result[0] == "Can you hear me?"
class TestSentenceSplitterChinese:
"""Tests for CJK sentence splitting."""
def test_chinese_period(self):
splitter = SentenceSplitter()
result = splitter.add_text("你好世界。你好吗")
assert len(result) == 1
assert result[0] == "你好世界。"
def test_chinese_exclamation(self):
splitter = SentenceSplitter()
result = splitter.add_text("太好了!谢谢你")
assert len(result) == 1
assert result[0] == "太好了!"
def test_chinese_question(self):
splitter = SentenceSplitter()
result = splitter.add_text("你是谁?我是小明")
assert len(result) == 1
assert result[0] == "你是谁?"
def test_chinese_comma_no_split(self):
"""Chinese commas are clause-level and should not trigger a split."""
splitter = SentenceSplitter()
result = splitter.add_text("你好,世界")
assert result == []
assert splitter.buffer == "你好,世界"
def test_chinese_semicolon_no_split(self):
"""Chinese semicolons are clause-level and should not trigger a split."""
splitter = SentenceSplitter()
result = splitter.add_text("第一点;第二点")
assert result == []
assert splitter.buffer == "第一点;第二点"
def test_chinese_multiple(self):
splitter = SentenceSplitter()
result = splitter.add_text("你好!你好吗?我很好。")
assert len(result) == 3
assert result[0] == "你好!"
assert result[1] == "你好吗?"
assert result[2] == "我很好。"
class TestSentenceSplitterMixed:
"""Tests for mixed-language sentence splitting."""
def test_mixed_english_chinese(self):
splitter = SentenceSplitter()
result = splitter.add_text("Hello世界。How are you? ")
assert len(result) == 2
assert result[0] == "Hello世界。"
assert result[1] == "How are you?"
class TestSentenceSplitterIncremental:
"""Tests for incremental (multi-chunk) text input."""
def test_accumulation_across_chunks(self):
splitter = SentenceSplitter()
# First chunk: no boundary
result1 = splitter.add_text("Hello ")
assert result1 == []
# Second chunk: completes a sentence
result2 = splitter.add_text("world. How")
assert len(result2) == 1
assert result2[0] == "Hello world."
assert splitter.buffer == "How"
def test_word_by_word(self):
splitter = SentenceSplitter()
words = ["Hello, ", "how ", "are ", "you? ", "I ", "am ", "fine."]
all_sentences = []
for word in words:
all_sentences.extend(splitter.add_text(word))
assert len(all_sentences) == 1
assert all_sentences[0] == "Hello, how are you?"
# "I am fine." stays in buffer (no trailing whitespace after period)
def test_three_chunks(self):
splitter = SentenceSplitter()
splitter.add_text("The quick brown ")
splitter.add_text("fox jumps. ")
result = splitter.add_text("Over the lazy dog. ")
# "The quick brown fox jumps." should have been returned on second chunk
# "Over the lazy dog." on third chunk
assert len(result) == 1
assert result[0] == "Over the lazy dog."
class TestSentenceSplitterFlush:
"""Tests for flush behavior."""
def test_flush_returns_remaining(self):
splitter = SentenceSplitter()
splitter.add_text("Hello world")
result = splitter.flush()
assert result == "Hello world"
assert splitter.buffer == ""
def test_flush_empty_buffer(self):
splitter = SentenceSplitter()
result = splitter.flush()
assert result is None
def test_flush_after_sentence(self):
splitter = SentenceSplitter()
splitter.add_text("Hello world. Remaining text")
result = splitter.flush()
assert result == "Remaining text"
def test_flush_whitespace_only(self):
splitter = SentenceSplitter()
splitter.add_text("Hello. ")
# "Hello." extracted, buffer is " "
result = splitter.flush()
# Whitespace-only should return None
assert result is None
def test_flush_clears_buffer(self):
splitter = SentenceSplitter()
splitter.add_text("some text")
splitter.flush()
assert splitter.buffer == ""
# Second flush should return None
assert splitter.flush() is None
class TestSentenceSplitterEdgeCases:
"""Edge case tests."""
def test_empty_input(self):
splitter = SentenceSplitter()
result = splitter.add_text("")
assert result == []
assert splitter.buffer == ""
def test_none_like_empty(self):
"""Empty string should not affect buffer."""
splitter = SentenceSplitter()
splitter.add_text("Hello")
splitter.add_text("")
assert splitter.buffer == "Hello"
def test_only_punctuation(self):
splitter = SentenceSplitter()
result = splitter.add_text(". ")
# "." is 1 char, below default min_sentence_length of 2
# It will be carried forward
assert result == []
def test_min_sentence_length(self):
splitter = SentenceSplitter(min_sentence_length=10)
result = splitter.add_text("Hi. Hello world. ")
# "Hi." is 3 chars (< 10), so it gets carried to "Hello world."
assert len(result) == 1
assert "Hi." in result[0]
assert "Hello world." in result[0]
def test_short_segments_are_carried_until_long_enough(self):
splitter = SentenceSplitter(min_sentence_length=10)
result = splitter.add_text("Hi. Ok. Hello there. ")
assert result == ["Hi.Ok.Hello there."]
assert splitter.buffer == ""
def test_min_sentence_length_zero(self):
splitter = SentenceSplitter(min_sentence_length=0)
result = splitter.add_text("A. B. ")
assert len(result) == 2
def test_no_boundary_then_flush(self):
splitter = SentenceSplitter()
result = splitter.add_text("Hello world how are you")
assert result == []
flushed = splitter.flush()
assert flushed == "Hello world how are you"
def test_consecutive_punctuation(self):
splitter = SentenceSplitter()
result = splitter.add_text("Really?! Yes, really. ")
assert len(result) >= 1
def test_reuse_after_flush(self):
"""Splitter can be reused after flush."""
splitter = SentenceSplitter()
splitter.add_text("First session.")
splitter.flush()
result = splitter.add_text("Second session. More text")
assert len(result) == 1
assert result[0] == "Second session."
assert splitter.buffer == "More text"
class TestSentenceSplitterBufferLimit:
"""Tests for buffer overflow protection."""
def test_buffer_overflow_raises(self):
from vllm_omni.entrypoints.openai.text_splitter import _MAX_BUFFER_SIZE
splitter = SentenceSplitter()
# Fill buffer just under the limit
splitter.add_text("x" * (_MAX_BUFFER_SIZE - 1))
# One more char should exceed the limit
with pytest.raises(ValueError, match="exceeded maximum size"):
splitter.add_text("xx")
+80 -333
View File
@@ -4,31 +4,19 @@
Unit tests for OpenAI-compatible video generation endpoints.
"""
import asyncio
import base64
import io
import json
import os
import threading
import time
from types import SimpleNamespace
import pytest
from fastapi import FastAPI
from fastapi.testclient import TestClient
from PIL import Image
from pydantic import ValidationError
from pytest_mock import MockerFixture
from vllm_omni.entrypoints.openai import api_server
from vllm_omni.entrypoints.openai.api_server import router
from vllm_omni.entrypoints.openai.protocol.videos import (
VideoGenerationRequest,
VideoGenerationStatus,
VideoResponse,
)
from vllm_omni.entrypoints.openai.protocol.videos import VideoGenerationRequest, VideoResponseFormat
from vllm_omni.entrypoints.openai.serving_video import OmniOpenAIServingVideo
from vllm_omni.entrypoints.openai.storage import LocalStorageManager
from vllm_omni.entrypoints.openai.stores import AsyncDictStore, TaskRegistry
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -56,38 +44,6 @@ class FakeAsyncOmni:
yield MockVideoResult(videos)
class BlockingVideoHandler:
def __init__(self):
self.model_name = "Wan-AI/Wan2.2-T2V-A14B-Diffusers"
self.stage_configs = None
self.started = threading.Event()
self.cancelled = threading.Event()
def set_stage_configs_if_missing(self, stage_configs):
if self.stage_configs is None:
self.stage_configs = stage_configs
async def generate_videos(self, request, reference_id, *, reference_image=None):
self.started.set()
try:
await asyncio.Future()
except asyncio.CancelledError:
self.cancelled.set()
raise
@pytest.fixture(autouse=True)
def isolated_video_backends(tmp_path, monkeypatch):
"""Use isolated in-memory metadata and local storage for each test."""
store: AsyncDictStore[VideoResponse] = AsyncDictStore()
tasks = TaskRegistry()
storage = LocalStorageManager(storage_path=str(tmp_path / "storage"))
monkeypatch.setattr(api_server, "VIDEO_STORE", store)
monkeypatch.setattr(api_server, "VIDEO_TASKS", tasks)
monkeypatch.setattr(api_server, "STORAGE_MANAGER", storage)
return store, tasks, storage
@pytest.fixture
def test_client():
app = FastAPI()
@@ -96,8 +52,7 @@ def test_client():
diffusion_engine=FakeAsyncOmni(),
model_name="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
)
with TestClient(app) as client:
yield client
return TestClient(app)
def _make_test_image_bytes(size=(64, 64)) -> bytes:
@@ -107,34 +62,6 @@ def _make_test_image_bytes(size=(64, 64)) -> bytes:
return buf.getvalue()
def _make_test_image_data_url(size=(64, 64)) -> str:
image_bytes = _make_test_image_bytes(size)
encoded = base64.b64encode(image_bytes).decode("utf-8")
return f"data:image/png;base64,{encoded}"
def _wait_for_status(client: TestClient, video_id: str, status: str, timeout_s: float = 2.0):
deadline = time.time() + timeout_s
last_payload = None
while time.time() < deadline:
response = client.get(f"/v1/videos/{video_id}")
assert response.status_code == 200
last_payload = response.json()
if last_payload["status"] == status:
return last_payload
time.sleep(0.02)
raise AssertionError(f"Timed out waiting for status={status}. Last payload: {last_payload}")
def _wait_until(predicate, timeout_s: float = 2.0, interval_s: float = 0.02):
deadline = time.time() + timeout_s
while time.time() < deadline:
if predicate():
return
time.sleep(interval_s)
raise AssertionError("Timed out waiting for condition")
def test_t2v_video_generation_form(test_client, mocker: MockerFixture):
fps_values = []
@@ -153,22 +80,24 @@ def test_t2v_video_generation_form(test_client, mocker: MockerFixture):
"size": "640x360",
"seconds": "2",
"fps": "12",
"n": "2",
},
)
assert response.status_code == 200
video_id = response.json()["id"]
_wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
data = response.json()
assert "data" in data and len(data["data"]) == 2
assert all(item["b64_json"] == "Zg==" for item in data["data"])
engine = test_client.app.state.openai_serving_video._engine_client
captured = engine.captured_sampling_params_list[0]
assert captured.num_outputs_per_prompt == 1
assert captured.num_outputs_per_prompt == 2
assert captured.width == 640
assert captured.height == 360
assert captured.num_frames == 24
assert captured.fps == 12
assert captured.frame_rate == 12.0
assert fps_values == [12]
assert fps_values == [12, 12]
def test_i2v_video_generation_form(test_client, mocker: MockerFixture):
@@ -185,8 +114,9 @@ def test_i2v_video_generation_form(test_client, mocker: MockerFixture):
)
assert response.status_code == 200
video_id = response.json()["id"]
_wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
data = response.json()
assert "data" in data and len(data["data"]) == 1
assert data["data"][0]["b64_json"] == "Zg=="
engine = test_client.app.state.openai_serving_video._engine_client
prompt = engine.captured_prompt
@@ -215,8 +145,6 @@ def test_i2v_video_generation_resizes_input_to_requested_dimensions(test_client,
)
assert response.status_code == 200
video_id = response.json()["id"]
_wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
engine = test_client.app.state.openai_serving_video._engine_client
prompt = engine.captured_prompt
@@ -225,30 +153,6 @@ def test_i2v_video_generation_resizes_input_to_requested_dimensions(test_client,
assert input_image.size == (96, 64)
def test_i2v_video_generation_with_image_reference_form(test_client, mocker: MockerFixture):
mocker.patch(
"vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
return_value="Zg==",
)
response = test_client.post(
"/v1/videos",
data={
"prompt": "A fox running through snow.",
"image_reference": json.dumps({"image_url": _make_test_image_data_url((40, 24))}),
},
)
assert response.status_code == 200
video_id = response.json()["id"]
_wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
engine = test_client.app.state.openai_serving_video._engine_client
prompt = engine.captured_prompt
input_image = prompt["multi_modal_data"]["image"]
assert isinstance(input_image, Image.Image)
assert input_image.size == (40, 24)
def test_seconds_defaults_fps_and_frames(test_client, mocker: MockerFixture):
fps_values = []
@@ -269,13 +173,10 @@ def test_seconds_defaults_fps_and_frames(test_client, mocker: MockerFixture):
)
assert response.status_code == 200
video_id = response.json()["id"]
_wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
engine = test_client.app.state.openai_serving_video._engine_client
captured = engine.captured_sampling_params_list[0]
assert captured.num_frames == 72
assert captured.fps == 24
assert captured.frame_rate == 24.0
assert fps_values == [24]
@@ -293,50 +194,16 @@ def test_size_param_sets_width_height(test_client, mocker: MockerFixture):
)
assert response.status_code == 200
video_id = response.json()["id"]
_wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
engine = test_client.app.state.openai_serving_video._engine_client
captured = engine.captured_sampling_params_list[0]
assert captured.width == 320
assert captured.height == 240
def test_sampling_params_pass_through(test_client, mocker: MockerFixture):
mocker.patch(
"vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
return_value="Zg==",
)
response = test_client.post(
"/v1/videos",
data={
"prompt": "param pass",
"num_inference_steps": "30",
"guidance_scale": "6.5",
"guidance_scale_2": "8.0",
"true_cfg_scale": "4.0",
"boundary_ratio": "0.7",
"flow_shift": "0.25",
},
)
assert response.status_code == 200
video_id = response.json()["id"]
_wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
engine = test_client.app.state.openai_serving_video._engine_client
captured = engine.captured_sampling_params_list[0]
assert captured.num_inference_steps == 30
assert captured.guidance_scale == 6.5
assert captured.guidance_scale_2 == 8.0
assert captured.true_cfg_scale == 4.0
assert captured.boundary_ratio == 0.7
assert captured.extra_args["flow_shift"] == 0.25
def test_audio_sample_rate_comes_from_model_config(test_client, mocker: MockerFixture):
audio_sample_rates = []
def _fake_encode(video, fps, audio=None, audio_sample_rate=None):
del video, fps, audio
audio_sample_rates.append(audio_sample_rate)
return "Zg=="
@@ -366,11 +233,38 @@ def test_audio_sample_rate_comes_from_model_config(test_client, mocker: MockerFi
)
assert response.status_code == 200
video_id = response.json()["id"]
_wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
assert audio_sample_rates == [16000]
def test_sampling_params_pass_through(test_client, mocker: MockerFixture):
mocker.patch(
"vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
return_value="Zg==",
)
response = test_client.post(
"/v1/videos",
data={
"prompt": "param pass",
"num_inference_steps": "30",
"guidance_scale": "6.5",
"guidance_scale_2": "8.0",
"true_cfg_scale": "4.0",
"boundary_ratio": "0.7",
"flow_shift": "0.25",
},
)
assert response.status_code == 200
engine = test_client.app.state.openai_serving_video._engine_client
captured = engine.captured_sampling_params_list[0]
assert captured.num_inference_steps == 30
assert captured.guidance_scale == 6.5
assert captured.guidance_scale_2 == 8.0
assert captured.true_cfg_scale == 4.0
assert captured.boundary_ratio == 0.7
assert captured.extra_args["flow_shift"] == 0.25
def test_missing_handler_returns_503():
app = FastAPI()
app.include_router(router)
@@ -393,29 +287,29 @@ def test_missing_prompt_returns_422(test_client):
assert response.status_code == 422
def test_invalid_size_parse_returns_422(test_client):
def test_invalid_size_format_raises_validation_error(test_client):
with pytest.raises(ValidationError):
test_client.post(
"/v1/videos",
data={"prompt": "bad size", "size": "invalid"},
)
def test_invalid_size_parse_returns_500(test_client):
response = test_client.post(
"/v1/videos",
data={"prompt": "bad size", "size": "640x"},
)
assert response.status_code == 422
body = response.json()
assert body["detail"][0]["loc"] == ["body", "size"]
assert body["detail"][0]["type"] == "string_pattern_mismatch"
assert body["detail"][0]["input"] == "640x"
assert response.status_code == 500
assert "video generation failed" in response.json()["detail"].lower()
def test_rejects_input_reference_and_image_reference_together(test_client):
response = test_client.post(
"/v1/videos",
data={
"prompt": "bad refs",
"image_reference": '{"image_url": "https://example.com/cat.png"}',
},
files={"input_reference": ("input.png", _make_test_image_bytes(), "image/png")},
)
assert response.status_code == 400
assert "either input_reference or image_reference" in response.json()["detail"].lower()
def test_invalid_response_format_raises_validation_error(test_client):
with pytest.raises(ValidationError):
test_client.post(
"/v1/videos",
data={"prompt": "bad format", "response_format": "url"},
)
def test_invalid_seconds_returns_422(test_client):
@@ -426,6 +320,20 @@ def test_invalid_seconds_returns_422(test_client):
assert response.status_code == 422
def test_invalid_n_raises_validation_error(test_client):
with pytest.raises(ValidationError):
test_client.post(
"/v1/videos",
data={"prompt": "bad n", "n": "0"},
)
with pytest.raises(ValidationError):
test_client.post(
"/v1/videos",
data={"prompt": "bad n", "n": "5"},
)
def test_negative_prompt_and_seed_pass_through(test_client, mocker: MockerFixture):
mocker.patch(
"vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
@@ -441,8 +349,6 @@ def test_negative_prompt_and_seed_pass_through(test_client, mocker: MockerFixtur
)
assert response.status_code == 200
video_id = response.json()["id"]
_wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
engine = test_client.app.state.openai_serving_video._engine_client
captured_prompt = engine.captured_prompt
captured_params = engine.captured_sampling_params_list[0]
@@ -458,38 +364,19 @@ def test_invalid_lora_returns_400(test_client):
"lora": '{"name": "bad-lora"}',
},
)
assert response.status_code == 200
video_id = response.json()["id"]
failed = _wait_for_status(test_client, video_id, VideoGenerationStatus.FAILED.value)
assert failed["error"]["code"] == "HTTPException"
assert "lora object" in failed["error"]["message"].lower()
def test_unsupported_image_reference_file_id_returns_400(test_client):
response = test_client.post(
"/v1/videos",
data={
"prompt": "unsupported ref",
"image_reference": '{"file_id": "file-123"}',
},
)
assert response.status_code == 400
assert response.json()["detail"] == "Invalid image_reference: file_id is not supported yet."
def test_invalid_uploaded_input_reference_returns_400(test_client):
response = test_client.post(
"/v1/videos",
data={"prompt": "bad upload"},
files={"input_reference": ("input.png", b"not-an-image", "image/png")},
)
assert response.status_code == 400
assert response.json()["detail"] == "Invalid input_reference: provided content is not a valid image."
assert "lora" in response.json()["detail"].lower()
def test_video_request_validation():
req = VideoGenerationRequest(prompt="test")
assert req.prompt == "test"
assert req.n == 1
assert req.response_format == VideoResponseFormat.B64_JSON
with pytest.raises(ValueError):
VideoGenerationRequest(prompt="test", response_format="url")
with pytest.raises(ValueError):
VideoGenerationRequest(prompt="test", size="invalid")
@@ -497,144 +384,4 @@ def test_video_request_validation():
VideoGenerationRequest(prompt="test", seconds="abc")
with pytest.raises(ValueError):
VideoGenerationRequest(prompt="test", image_reference={"file_id": "file-1", "image_url": "https://example.com"})
def test_list_videos_supports_order_after_and_limit(test_client, mocker: MockerFixture):
mocker.patch(
"vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
return_value="Zg==",
)
ids = []
for i in range(3):
create_resp = test_client.post("/v1/videos", data={"prompt": f"video-{i}"})
assert create_resp.status_code == 200
video_id = create_resp.json()["id"]
_wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
ids.append(video_id)
asyncio.run(api_server.VIDEO_STORE.update_fields(ids[0], {"created_at": 100}))
asyncio.run(api_server.VIDEO_STORE.update_fields(ids[1], {"created_at": 200}))
asyncio.run(api_server.VIDEO_STORE.update_fields(ids[2], {"created_at": 300}))
asc_resp = test_client.get("/v1/videos", params={"order": "asc"})
assert asc_resp.status_code == 200
asc_body = asc_resp.json()
asc_ids = [item["id"] for item in asc_body["data"]]
assert asc_ids == [ids[0], ids[1], ids[2]]
assert asc_body["object"] == "list"
assert asc_body["first_id"] == ids[0]
assert asc_body["last_id"] == ids[2]
assert asc_body["has_more"] is False
desc_resp = test_client.get("/v1/videos", params={"order": "desc", "limit": 2})
assert desc_resp.status_code == 200
desc_body = desc_resp.json()
desc_ids = [item["id"] for item in desc_body["data"]]
assert desc_ids == [ids[2], ids[1]]
assert desc_body["object"] == "list"
assert desc_body["first_id"] == ids[2]
assert desc_body["last_id"] == ids[1]
assert desc_body["has_more"] is True
after_resp = test_client.get("/v1/videos", params={"order": "asc", "after": ids[0]})
assert after_resp.status_code == 200
after_body = after_resp.json()
after_ids = [item["id"] for item in after_body["data"]]
assert after_ids == [ids[1], ids[2]]
assert after_body["object"] == "list"
assert after_body["first_id"] == ids[1]
assert after_body["last_id"] == ids[2]
assert after_body["has_more"] is False
zero_limit_resp = test_client.get("/v1/videos", params={"order": "asc", "limit": 0})
assert zero_limit_resp.status_code == 200
zero_limit_body = zero_limit_resp.json()
assert zero_limit_body["data"] == []
assert zero_limit_body["object"] == "list"
assert zero_limit_body["first_id"] is None
assert zero_limit_body["last_id"] is None
assert zero_limit_body["has_more"] is True
zero_limit_after_resp = test_client.get(
"/v1/videos",
params={"order": "asc", "after": ids[2], "limit": 0},
)
assert zero_limit_after_resp.status_code == 200
zero_limit_after_body = zero_limit_after_resp.json()
assert zero_limit_after_body["data"] == []
assert zero_limit_after_body["object"] == "list"
assert zero_limit_after_body["first_id"] is None
assert zero_limit_after_body["last_id"] is None
assert zero_limit_after_body["has_more"] is False
def test_delete_completed_job_removes_file_and_metadata(test_client, mocker: MockerFixture):
mocker.patch(
"vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
return_value="Zg==",
)
create_resp = test_client.post("/v1/videos", data={"prompt": "Delete this video"})
assert create_resp.status_code == 200
video_id = create_resp.json()["id"]
final = _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
file_name = final["file_name"]
assert file_name is not None
file_path = os.path.join(api_server.STORAGE_MANAGER.storage_path, file_name)
assert os.path.exists(file_path)
delete_resp = test_client.delete(f"/v1/videos/{video_id}")
assert delete_resp.status_code == 200
assert delete_resp.json()["id"] == video_id
assert delete_resp.json()["deleted"] is True
assert delete_resp.json()["object"] == "video.deleted"
assert not os.path.exists(file_path)
def test_delete_in_progress_job_cancels_task_and_removes_metadata(test_client):
handler = BlockingVideoHandler()
test_client.app.state.openai_serving_video = handler
create_resp = test_client.post("/v1/videos", data={"prompt": "Cancel this video"})
assert create_resp.status_code == 200
video_id = create_resp.json()["id"]
assert handler.started.wait(timeout=2.0)
delete_resp = test_client.delete(f"/v1/videos/{video_id}")
assert delete_resp.status_code == 200
assert delete_resp.json()["id"] == video_id
assert delete_resp.json()["deleted"] is True
assert delete_resp.json()["object"] == "video.deleted"
assert handler.cancelled.wait(timeout=2.0)
_wait_until(lambda: asyncio.run(api_server.VIDEO_TASKS.get(video_id)) is None)
assert asyncio.run(api_server.VIDEO_STORE.get(video_id)) is None
retrieve_resp = test_client.get(f"/v1/videos/{video_id}")
assert retrieve_resp.status_code == 404
def test_video_response_file_extension_is_robust():
response = VideoResponse(model="test-model", prompt="Make something beautiful")
assert response.file_extension == "mp4"
with_params = VideoResponse.model_construct(
model="test-model",
media_type="video/mp4; charset=binary",
)
assert with_params.file_extension == "mp4"
webm = VideoResponse.model_construct(
model="test-model",
media_type="video/webm",
)
assert webm.file_extension == "webm"
with pytest.raises(ValueError):
unknown = VideoResponse.model_construct(
model="test-model",
media_type="application/x-custom-video",
)
_ = unknown.file_extension
VideoGenerationRequest(prompt="test", n=0)
+1 -1
View File
@@ -345,7 +345,7 @@ def _setup_multiprocessing_mocks(monkeypatch: pytest.MonkeyPatch, mocker: Mocker
fake_process_instance = mocker.MagicMock()
fake_process_instance.start = mocker.MagicMock()
fake_process_instance.join = mocker.MagicMock()
fake_process_instance.is_alive = mocker.MagicMock(return_value=True)
fake_process_instance.is_alive = mocker.MagicMock(return_value=False)
fake_process_instance.terminate = mocker.MagicMock()
fake_process_class.return_value = fake_process_instance
@@ -4,8 +4,6 @@ Example online tests for Qwen2.5-Omni-7B model.
import os
from vllm_omni.platforms import current_omni_platform
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
import re
@@ -22,11 +20,6 @@ models = ["Qwen/Qwen2.5-Omni-7B"]
stage_configs = [str(Path(__file__).parent.parent.parent / "e2e" / "stage_configs" / "qwen2_5_omni_ci.yaml")]
if current_omni_platform.is_xpu():
stage_configs = [
str(Path(__file__).parent.parent.parent / "e2e" / "stage_configs" / "xpu" / "qwen2_5_omni_ci.yaml")
]
example_dir = str(Path(__file__).parent.parent.parent.parent / "examples" / "online_serving" / "qwen2_5_omni")
# Create parameter combinations for model and stage config
test_params = [
@@ -4,8 +4,6 @@ Example Online tests for Qwen3-Omni model.
import os
from vllm_omni.platforms import current_omni_platform
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
import re
@@ -22,10 +20,6 @@ models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
stage_configs = [str(Path(__file__).parent.parent.parent / "e2e" / "stage_configs" / "qwen3_omni_ci.yaml")]
if current_omni_platform.is_xpu():
stage_configs = [str(Path(__file__).parent.parent.parent / "e2e" / "stage_configs" / "xpu" / "qwen3_omni_ci.yaml")]
example_dir = str(Path(__file__).parent.parent.parent.parent / "examples" / "online_serving" / "qwen3_omni")
# Create parameter combinations for model and stage config
test_params = [
-179
View File
@@ -1,179 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
Test metadata for Omni models.
Mirrors vllm's ``tests/models/registry.py`` (_HfExamplesInfo) but adapted for
the Omni multi-stage architecture.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Literal
import pytest
from packaging.version import Version
from transformers import __version__ as TRANSFORMERS_VERSION
@dataclass(frozen=True)
class _OmniExamplesInfo:
"""Metadata used by processing-correctness tests for Omni models."""
default: str
"""The default HF model ID to use for testing this architecture."""
model_stage: str = "thinker"
"""Stage type identifier, e.g. ``"thinker"`` or ``"talker"``."""
hf_config_name: str | None = None
"""
Name of the sub-config inside the top-level HF config
(e.g. ``"thinker_config"``). ``None`` means use the root config.
"""
has_multimodal_processor: bool = False
"""Whether the architecture registers a multimodal processor."""
min_transformers_version: str | None = None
"""Minimum ``transformers`` version required to load this model."""
max_transformers_version: str | None = None
"""Maximum ``transformers`` version that this model supports."""
is_available_online: bool = True
"""Set to ``False`` if the model is no longer hosted on HF Hub."""
trust_remote_code: bool = False
"""Whether ``trust_remote_code=True`` is needed to load the model."""
max_model_len: int | None = None
"""
An explicit ``max_model_len`` override. Useful when the default value
from the HF config is too large for CI.
"""
# ----- helper methods (same API as _HfExamplesInfo) --------------------
def check_transformers_version(
self,
*,
on_fail: Literal["error", "skip", "return"],
) -> str | None:
"""Skip / error / return a message when the installed
``transformers`` does not satisfy the version constraints."""
if self.min_transformers_version is None and self.max_transformers_version is None:
return None
current = TRANSFORMERS_VERSION
cur_base = Version(current).base_version
min_ver = self.min_transformers_version
max_ver = self.max_transformers_version
msg = f"`transformers=={current}` installed, but `transformers"
if min_ver and Version(cur_base) < Version(min_ver):
msg += f">={min_ver}` is required to run this model."
elif max_ver and Version(cur_base) > Version(max_ver):
msg += f"<={max_ver}` is required to run this model."
else:
return None
if on_fail == "error":
raise RuntimeError(msg)
elif on_fail == "skip":
pytest.skip(msg)
return msg
def check_available_online(
self,
*,
on_fail: Literal["error", "skip"],
) -> None:
"""Skip / error when the model is not available online."""
if not self.is_available_online:
msg = "Model is not available online"
if on_fail == "error":
raise RuntimeError(msg)
else:
pytest.skip(msg)
# ---------------------------------------------------------------------------
# Example model registry for tests
# ---------------------------------------------------------------------------
_OMNI_EXAMPLE_MODELS: dict[str, _OmniExamplesInfo] = {
# ---- Qwen2.5-Omni ----
"Qwen2_5OmniForConditionalGeneration": _OmniExamplesInfo(
default="Qwen/Qwen2.5-Omni-7B",
model_stage="thinker",
has_multimodal_processor=True,
),
"Qwen2_5OmniThinkerModel": _OmniExamplesInfo(
default="Qwen/Qwen2.5-Omni-7B",
model_stage="thinker",
hf_config_name="thinker_config",
has_multimodal_processor=True,
),
"Qwen2_5OmniTalkerModel": _OmniExamplesInfo(
default="Qwen/Qwen2.5-Omni-7B",
model_stage="talker",
hf_config_name="talker_config",
has_multimodal_processor=True,
),
"Qwen2_5OmniToken2WavModel": _OmniExamplesInfo(
default="Qwen/Qwen2.5-Omni-7B",
model_stage="token2wav",
),
"Qwen2_5OmniToken2WavDiTModel": _OmniExamplesInfo(
default="Qwen/Qwen2.5-Omni-7B",
model_stage="token2wav",
),
# Internal sub-model key hardcoded by Qwen2_5OmniTalkerForConditionalGeneration
# to load its language model component (see qwen2_5_omni_talker.py).
"Qwen2ForCausalLM_old": _OmniExamplesInfo(
default="Qwen/Qwen2.5-Omni-7B",
model_stage="thinker",
),
# ---- Qwen3-Omni MoE ----
"Qwen3OmniMoeForConditionalGeneration": _OmniExamplesInfo(
default="Qwen/Qwen3-Omni-30B-A3B-Instruct",
model_stage="thinker",
has_multimodal_processor=True,
),
"Qwen3OmniMoeThinkerForConditionalGeneration": _OmniExamplesInfo(
default="Qwen/Qwen3-Omni-30B-A3B-Instruct",
model_stage="thinker",
hf_config_name="thinker_config",
has_multimodal_processor=True,
),
"Qwen3OmniMoeTalkerForConditionalGeneration": _OmniExamplesInfo(
default="Qwen/Qwen3-Omni-30B-A3B-Instruct",
model_stage="talker",
hf_config_name="talker_config",
),
"Qwen3OmniMoeCode2Wav": _OmniExamplesInfo(
default="Qwen/Qwen3-Omni-30B-A3B-Instruct",
model_stage="code2wav",
),
# ---- Qwen3-TTS ----
"Qwen3TTSForConditionalGeneration": _OmniExamplesInfo(
default="Qwen/Qwen3-TTS",
model_stage="talker",
),
"Qwen3TTSTalkerForConditionalGeneration": _OmniExamplesInfo(
default="Qwen/Qwen3-TTS",
model_stage="talker",
),
"Qwen3TTSCode2Wav": _OmniExamplesInfo(
default="Qwen/Qwen3-TTS",
model_stage="code2wav",
),
}
# Only architectures that register a multimodal processor are relevant for
# processing-correctness tests.
_MULTIMODAL_OMNI_EXAMPLE_MODELS: dict[str, _OmniExamplesInfo] = {
arch: info for arch, info in _OMNI_EXAMPLE_MODELS.items() if info.has_multimodal_processor
}
@@ -1,327 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""
Test multimodal processing correctness for Omni models.
Tests that:
1. Cached vs non-cached processor outputs are identical.
2. Text prompt vs token prompt inputs produce identical results after processing.
Adapted from vllm/tests/models/multimodal/processing/test_common.py
"""
from functools import partial
import numpy as np
import pytest
from PIL import Image
from vllm.config.multimodal import (
AudioDummyOptions,
BaseDummyOptions,
ImageDummyOptions,
VideoDummyOptions,
)
from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalDataDict
from vllm.multimodal.cache import MultiModalProcessorOnlyCache
from vllm.multimodal.inputs import MultiModalInputs, batched_tensors_equal
from vllm.multimodal.processing import BaseMultiModalProcessor, InputProcessingContext
from vllm.tokenizers import TokenizerLike, cached_tokenizer_from_config
from tests.model_executor.models.registry import (
_MULTIMODAL_OMNI_EXAMPLE_MODELS,
_OmniExamplesInfo,
)
from vllm_omni.config import OmniModelConfig
def random_image(rng: np.random.RandomState, min_wh: int, max_wh: int):
w, h = rng.randint(min_wh, max_wh, size=(2,))
arr = rng.randint(0, 255, size=(w, h, 3), dtype=np.uint8)
return Image.fromarray(arr)
def random_video(
rng: np.random.RandomState,
min_frames: int,
max_frames: int,
min_wh: int,
max_wh: int,
):
num_frames = rng.randint(min_frames, max_frames)
w, h = rng.randint(min_wh, max_wh, size=(2,))
return rng.randint(0, 255, size=(num_frames, w, h, 3), dtype=np.uint8)
def random_audio(
rng: np.random.RandomState,
min_len: int,
max_len: int,
sr: int,
):
audio_len = rng.randint(min_len, max_len)
return rng.rand(audio_len), sr
_IGNORE_MM_KEYS: dict[str, set[str]] = {}
def _get_model_archs_to_test() -> list[str]:
"""Return architecture strings for parametrization."""
return list(_MULTIMODAL_OMNI_EXAMPLE_MODELS.keys())
def get_text_token_prompts(
processor: BaseMultiModalProcessor,
mm_data: MultiModalDataDict,
):
"""Return ``(text_prompt, token_prompt)`` for the given multimodal data."""
dummy_inputs = processor.dummy_inputs
tokenizer: TokenizerLike = processor.info.get_tokenizer()
model_config = processor.info.ctx.model_config
parsed_data = processor.info.parse_mm_data(mm_data)
mm_counts = {k: len(vs) for k, vs in parsed_data.items()}
inputs = dummy_inputs.get_dummy_processor_inputs(
model_config.max_model_len,
mm_counts,
mm_options={},
)
text_prompt: str | None
token_prompt: list[int]
if isinstance(inputs.prompt, list):
text_prompt = None
token_prompt = inputs.prompt
elif isinstance(inputs.prompt, str):
text_prompt = inputs.prompt
token_prompt = tokenizer.encode(
text_prompt,
**processor.info.get_default_tok_params().get_encode_kwargs(),
)
else:
raise TypeError(type(inputs.prompt))
return text_prompt, token_prompt
def _build_model_config(model_arch: str, info: _OmniExamplesInfo) -> OmniModelConfig:
"""Create an ``OmniModelConfig`` suitable for processor testing."""
kwargs: dict = dict(
model=info.default,
tokenizer=info.default,
tokenizer_mode="auto",
trust_remote_code=info.trust_remote_code,
model_arch=model_arch,
model_stage=info.model_stage,
max_model_len=info.max_model_len,
enforce_eager=True,
dtype="auto",
)
if info.hf_config_name is not None:
kwargs["hf_config_name"] = info.hf_config_name
model_config = OmniModelConfig(**kwargs)
# Ensure cache is effectively unlimited to prevent eviction during test
model_config.multimodal_config.mm_processor_cache_gb = 2048
return model_config
def _test_processing_correctness(
model_arch: str,
hit_rate: float,
num_batches: int,
simplify_rate: float,
):
"""Run the full processing-correctness test for one Omni model."""
info = _MULTIMODAL_OMNI_EXAMPLE_MODELS[model_arch]
info.check_transformers_version(on_fail="skip")
info.check_available_online(on_fail="skip")
model_config = _build_model_config(model_arch, info)
model_cls = MULTIMODAL_REGISTRY._get_model_cls(model_config)
factories = model_cls._processor_factory
ctx = InputProcessingContext(
model_config,
tokenizer=cached_tokenizer_from_config(model_config),
)
cache = MultiModalProcessorOnlyCache(model_config)
processing_info = factories.info(ctx)
supported_mm_limits = processing_info.get_supported_mm_limits()
limit_mm_per_prompt_ints = {
modality: 3 if limit is None else limit for modality, limit in supported_mm_limits.items()
}
def _to_dummy_options(modality: str, count: int) -> BaseDummyOptions:
if modality == "video":
return VideoDummyOptions(count=count)
if modality == "image":
return ImageDummyOptions(count=count)
if modality == "audio":
return AudioDummyOptions(count=count)
return BaseDummyOptions(count=count)
model_config.get_multimodal_config().limit_per_prompt = {
modality: _to_dummy_options(modality, count) for modality, count in limit_mm_per_prompt_ints.items()
}
baseline_processor = factories.build_processor(ctx, cache=None)
cached_processor = factories.build_processor(ctx, cache=cache)
rng = np.random.RandomState(0)
input_to_hit = {
"image": Image.new("RGB", size=(128, 128)),
"video": np.zeros((4, 128, 128, 3), dtype=np.uint8),
"audio": (np.zeros((512,)), 16000),
}
input_factory = {
"image": partial(random_image, rng, min_wh=128, max_wh=256),
"video": partial(random_video, rng, min_frames=2, max_frames=16, min_wh=128, max_wh=256),
"audio": partial(random_audio, rng, min_len=512, max_len=1024, sr=16000),
}
for batch_idx in range(num_batches):
mm_data = {
k: [
(input_to_hit[k] if rng.rand() < hit_rate else input_factory[k]())
for _ in range(rng.randint(limit + 1))
]
for k, limit in limit_mm_per_prompt_ints.items()
}
# Drop unnecessary keys and test single -> multi conversion
if rng.rand() < simplify_rate:
for k in list(mm_data.keys()):
if not mm_data[k]:
del mm_data[k]
elif len(mm_data[k]) == 1:
mm_data[k] = mm_data[k][0]
_test_processing_correctness_one(
model_config,
mm_data,
baseline_processor,
cached_processor,
batch_idx,
)
def _test_processing_correctness_one(
model_config: OmniModelConfig,
mm_data: MultiModalDataDict,
baseline_processor: BaseMultiModalProcessor,
cached_processor: BaseMultiModalProcessor,
batch_idx: int,
):
model_type = model_config.hf_config.model_type
text_prompt, token_prompt = get_text_token_prompts(baseline_processor, mm_data)
mm_items = baseline_processor.info.parse_mm_data(mm_data)
ignore_mm_keys = _IGNORE_MM_KEYS.get(model_type, set[str]())
baseline_tokenized_result = baseline_processor(
token_prompt,
mm_items=mm_items,
hf_processor_mm_kwargs={},
)
cached_tokenized_result = cached_processor(
token_prompt,
mm_items=mm_items,
hf_processor_mm_kwargs={},
)
_assert_inputs_equal(
baseline_tokenized_result,
cached_tokenized_result,
ignore_mm_keys=ignore_mm_keys,
msg=f"Failed ({batch_idx=}, {token_prompt=}, {mm_data=})",
)
if text_prompt is not None:
baseline_text_result = baseline_processor(
text_prompt,
mm_items=mm_items,
hf_processor_mm_kwargs={},
)
cached_text_result = cached_processor(
text_prompt,
mm_items=mm_items,
hf_processor_mm_kwargs={},
)
_assert_inputs_equal(
baseline_text_result,
cached_text_result,
ignore_mm_keys=ignore_mm_keys,
msg=f"Failed ({batch_idx=}, {text_prompt=}, {mm_data=})",
)
_assert_inputs_equal(
baseline_text_result,
baseline_tokenized_result,
ignore_mm_keys=ignore_mm_keys,
msg=f"Failed ({batch_idx=}, {text_prompt=}, {token_prompt=}, {mm_data=})",
)
_assert_inputs_equal(
cached_text_result,
cached_tokenized_result,
ignore_mm_keys=ignore_mm_keys,
msg=f"Failed ({batch_idx=}, {text_prompt=}, {token_prompt=}, {mm_data=})",
)
@pytest.mark.core_model
@pytest.mark.omni
@pytest.mark.cpu
@pytest.mark.parametrize("model_arch", _get_model_archs_to_test())
@pytest.mark.parametrize("hit_rate", [0.3, 0.5, 1.0])
@pytest.mark.parametrize("num_batches", [32])
@pytest.mark.parametrize("simplify_rate", [1.0])
def test_omni_processing_correctness(
model_arch: str,
hit_rate: float,
num_batches: int,
simplify_rate: float,
):
"""
For each Omni model verify that:
- cached and non-cached processor outputs are identical.
- text-prompt and token-prompt inputs yield the same result.
"""
_test_processing_correctness(
model_arch,
hit_rate=hit_rate,
num_batches=num_batches,
simplify_rate=simplify_rate,
)
def _assert_inputs_equal(
a: MultiModalInputs,
b: MultiModalInputs,
*,
ignore_mm_keys: set[str] | None = None,
msg: str = "",
):
if ignore_mm_keys is None:
ignore_mm_keys = set()
ignore_prompt_keys = ("prompt", "mm_kwargs")
a_rest = {k: v for k, v in a.items() if k not in ignore_prompt_keys}
b_rest = {k: v for k, v in b.items() if k not in ignore_prompt_keys}
assert a_rest == b_rest, msg
a_data = a["mm_kwargs"].get_data()
b_data = b["mm_kwargs"].get_data()
for key in ignore_mm_keys:
a_data.pop(key, None)
b_data.pop(key, None)
assert batched_tensors_equal(a_data, b_data), msg
-570
View File
@@ -1,570 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Unit tests for StageConfigFactory and related classes.
"""
from vllm_omni.config.stage_config import (
ModelPipeline,
StageConfig,
StageConfigFactory,
StageType,
)
class TestStageType:
"""Tests for StageType enum."""
def test_stage_type_values(self):
"""Test StageType enum values."""
assert StageType.LLM.value == "llm"
assert StageType.DIFFUSION.value == "diffusion"
def test_stage_type_from_string(self):
"""Test creating StageType from string."""
assert StageType("llm") == StageType.LLM
assert StageType("diffusion") == StageType.DIFFUSION
class TestStageConfig:
"""Tests for StageConfig dataclass."""
def test_minimal_config(self):
"""Test creating StageConfig with minimal required fields."""
config = StageConfig(stage_id=0, model_stage="thinker")
assert config.stage_id == 0
assert config.model_stage == "thinker"
assert config.stage_type == StageType.LLM
assert config.input_sources == []
assert config.final_output is False
assert config.worker_type is None
def test_full_config(self):
"""Test creating StageConfig with all fields."""
config = StageConfig(
stage_id=1,
model_stage="talker",
stage_type=StageType.LLM,
input_sources=[0],
custom_process_input_func="module.path.func",
final_output=True,
final_output_type="audio",
worker_type="ar",
scheduler_cls="path.to.Scheduler",
hf_config_name="talker_config",
is_comprehension=False,
)
assert config.stage_id == 1
assert config.model_stage == "talker"
assert config.input_sources == [0]
assert config.final_output_type == "audio"
assert config.worker_type == "ar"
def test_to_omegaconf_basic(self):
"""Test converting StageConfig to OmegaConf format."""
config = StageConfig(
stage_id=0,
model_stage="thinker",
stage_type=StageType.LLM,
worker_type="ar",
final_output=True,
final_output_type="text",
)
omega_config = config.to_omegaconf()
assert omega_config.stage_id == 0
assert omega_config.stage_type == "llm"
assert omega_config.engine_args.model_stage == "thinker"
assert omega_config.engine_args.worker_type == "ar"
assert omega_config.final_output is True
assert omega_config.final_output_type == "text"
# Legacy field name for backward compatibility
assert omega_config.engine_input_source == []
def test_to_omegaconf_with_runtime_overrides(self):
"""Test that runtime overrides are applied to OmegaConf output."""
config = StageConfig(
stage_id=0,
model_stage="thinker",
runtime_overrides={
"gpu_memory_utilization": 0.9,
"tensor_parallel_size": 2,
"devices": "0,1",
"max_batch_size": 64,
},
)
omega_config = config.to_omegaconf()
assert omega_config.engine_args.gpu_memory_utilization == 0.9
assert omega_config.engine_args.tensor_parallel_size == 2
assert omega_config.runtime.devices == "0,1"
assert omega_config.runtime.max_batch_size == 64
class TestModelPipeline:
"""Tests for ModelPipeline class."""
def test_valid_linear_dag(self):
"""Test validation of a valid linear DAG."""
stages = [
StageConfig(stage_id=0, model_stage="thinker", input_sources=[]),
StageConfig(stage_id=1, model_stage="talker", input_sources=[0]),
StageConfig(stage_id=2, model_stage="code2wav", input_sources=[1]),
]
pipeline = ModelPipeline(model_type="test", stages=stages)
errors = pipeline.validate_pipeline()
assert errors == [], f"Unexpected errors: {errors}"
def test_valid_branching_dag(self):
"""Test validation of a valid branching DAG."""
stages = [
StageConfig(stage_id=0, model_stage="input", input_sources=[]),
StageConfig(stage_id=1, model_stage="branch_a", input_sources=[0]),
StageConfig(stage_id=2, model_stage="branch_b", input_sources=[0]),
]
pipeline = ModelPipeline(model_type="test", stages=stages)
errors = pipeline.validate_pipeline()
assert errors == [], f"Unexpected errors: {errors}"
def test_missing_entry_point(self):
"""Test that missing entry point is detected."""
stages = [
StageConfig(stage_id=0, model_stage="stage_a", input_sources=[1]),
StageConfig(stage_id=1, model_stage="stage_b", input_sources=[0]),
]
pipeline = ModelPipeline(model_type="test", stages=stages)
errors = pipeline.validate_pipeline()
assert any("entry point" in e.lower() for e in errors)
def test_missing_dependency(self):
"""Test that missing stage reference is detected."""
stages = [
StageConfig(stage_id=0, model_stage="input", input_sources=[]),
StageConfig(stage_id=1, model_stage="output", input_sources=[99]), # Invalid
]
pipeline = ModelPipeline(model_type="test", stages=stages)
errors = pipeline.validate_pipeline()
assert any("non-existent" in e.lower() for e in errors)
def test_duplicate_stage_ids(self):
"""Test that duplicate stage IDs are detected."""
stages = [
StageConfig(stage_id=0, model_stage="stage_a", input_sources=[]),
StageConfig(stage_id=0, model_stage="stage_b", input_sources=[]), # Duplicate
]
pipeline = ModelPipeline(model_type="test", stages=stages)
errors = pipeline.validate_pipeline()
assert any("duplicate" in e.lower() for e in errors)
def test_self_reference(self):
"""Test that self-references are detected."""
stages = [
StageConfig(stage_id=0, model_stage="entry", input_sources=[]),
StageConfig(stage_id=1, model_stage="self_ref", input_sources=[1]), # Self
]
pipeline = ModelPipeline(model_type="test", stages=stages)
errors = pipeline.validate_pipeline()
assert any("itself" in e.lower() for e in errors)
def test_get_stage_by_id(self):
"""Test getting stage by ID."""
stages = [
StageConfig(stage_id=0, model_stage="thinker", input_sources=[]),
StageConfig(stage_id=1, model_stage="talker", input_sources=[0]),
]
pipeline = ModelPipeline(model_type="test", stages=stages)
stage = pipeline.get_stage(1)
assert stage is not None
assert stage.model_stage == "talker"
missing = pipeline.get_stage(99)
assert missing is None
def test_empty_pipeline(self):
"""Test validation of empty pipeline."""
pipeline = ModelPipeline(model_type="test", stages=[])
errors = pipeline.validate_pipeline()
assert any("no stages" in e.lower() for e in errors)
class TestStageConfigFactory:
"""Tests for StageConfigFactory class."""
def test_default_diffusion_no_yaml(self):
"""Test single-stage diffusion works without YAML config (@ZJY0516)."""
kwargs = {
"cache_backend": "none",
"cache_config": None,
"dtype": "bfloat16",
}
configs = StageConfigFactory.create_default_diffusion(kwargs)
assert len(configs) == 1
cfg = configs[0]
assert cfg["stage_id"] == 0
assert cfg["stage_type"] == "diffusion"
assert cfg["final_output"] is True
assert cfg["final_output_type"] == "image"
def test_default_diffusion_with_parallel_config(self):
"""Test diffusion config calculates devices from parallel_config."""
class MockParallelConfig:
world_size = 4
kwargs = {
"parallel_config": MockParallelConfig(),
"cache_backend": "tea_cache",
}
configs = StageConfigFactory.create_default_diffusion(kwargs)
assert configs[0]["runtime"]["devices"] == "0,1,2,3"
def test_per_stage_override_precedence(self):
"""Test that --stage-0-gpu-memory-utilization overrides global."""
stage = StageConfig(stage_id=0, model_stage="thinker", input_sources=[])
cli_overrides = {
"gpu_memory_utilization": 0.5, # Global
"stage_0_gpu_memory_utilization": 0.9, # Per-stage override
}
overrides = StageConfigFactory._merge_cli_overrides(stage, cli_overrides)
# Per-stage should override global
assert overrides["gpu_memory_utilization"] == 0.9
def test_cli_override_forwards_engine_registered_args(self):
"""Test that any engine-registered CLI arg is forwarded (@wuhang2014)."""
stage = StageConfig(stage_id=0, model_stage="thinker", input_sources=[])
cli_overrides = {
"gpu_memory_utilization": 0.9, # Well-known param
"custom_engine_flag": True, # Not in _INTERNAL_KEYS, so forwarded
}
overrides = StageConfigFactory._merge_cli_overrides(stage, cli_overrides)
assert overrides["gpu_memory_utilization"] == 0.9
assert overrides["custom_engine_flag"] is True
def test_cli_override_excludes_internal_keys(self):
"""Test that internal/orchestrator keys are not forwarded."""
stage = StageConfig(stage_id=0, model_stage="thinker", input_sources=[])
cli_overrides = {
"gpu_memory_utilization": 0.9,
"model": "some_model", # Internal
"stage_configs_path": "/path", # Internal
"batch_timeout": 10, # Internal
}
overrides = StageConfigFactory._merge_cli_overrides(stage, cli_overrides)
assert overrides["gpu_memory_utilization"] == 0.9
assert "model" not in overrides
assert "stage_configs_path" not in overrides
assert "batch_timeout" not in overrides
def test_per_stage_override_excludes_internal_keys(self):
"""Test that per-stage overrides also skip internal keys."""
stage = StageConfig(stage_id=0, model_stage="thinker", input_sources=[])
cli_overrides = {
"stage_0_gpu_memory_utilization": 0.9,
"stage_0_model": "override_model", # Internal, should be skipped
"stage_0_batch_timeout": 5, # Internal, should be skipped
}
overrides = StageConfigFactory._merge_cli_overrides(stage, cli_overrides)
assert overrides["gpu_memory_utilization"] == 0.9
assert "model" not in overrides
assert "batch_timeout" not in overrides
class TestPipelineYamlParsing:
"""Tests for pipeline YAML file parsing (@ZJY0516)."""
def test_parse_qwen3_omni_moe_yaml(self, tmp_path):
"""Test parsing the qwen3_omni_moe pipeline YAML."""
yaml_content = """\
model_type: qwen3_omni_moe
async_chunk: false
stages:
- stage_id: 0
model_stage: thinker
stage_type: llm
input_sources: []
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
hf_config_name: thinker_config
final_output: true
final_output_type: text
is_comprehension: true
- stage_id: 1
model_stage: talker
stage_type: llm
input_sources: [0]
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
hf_config_name: talker_config
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker
- stage_id: 2
model_stage: code2wav
stage_type: llm
input_sources: [1]
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
hf_config_name: thinker_config
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
final_output: true
final_output_type: audio
"""
yaml_file = tmp_path / "qwen3_omni_moe.yaml"
yaml_file.write_text(yaml_content)
pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "qwen3_omni_moe")
assert pipeline.model_type == "qwen3_omni_moe"
assert len(pipeline.stages) == 3
# Stage 0: thinker
s0 = pipeline.stages[0]
assert s0.stage_id == 0
assert s0.model_stage == "thinker"
assert s0.stage_type == StageType.LLM
assert s0.input_sources == []
assert s0.worker_type == "ar"
assert s0.final_output is True
assert s0.final_output_type == "text"
assert s0.is_comprehension is True
# Stage 1: talker
s1 = pipeline.stages[1]
assert s1.stage_id == 1
assert s1.input_sources == [0]
assert s1.custom_process_input_func == (
"vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker"
)
assert s1.final_output is False
# Stage 2: code2wav
s2 = pipeline.stages[2]
assert s2.stage_id == 2
assert s2.input_sources == [1]
assert s2.worker_type == "generation"
assert s2.final_output_type == "audio"
def test_parse_yaml_with_legacy_engine_input_source(self, tmp_path):
"""Test backward compatibility with engine_input_source field."""
yaml_content = """\
model_type: legacy_model
stages:
- stage_id: 0
model_stage: entry
stage_type: llm
- stage_id: 1
model_stage: downstream
stage_type: llm
engine_input_source: [0]
"""
yaml_file = tmp_path / "legacy.yaml"
yaml_file.write_text(yaml_content)
pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "legacy_model")
assert pipeline.stages[1].input_sources == [0]
def test_parse_yaml_with_connectors_and_edges(self, tmp_path):
"""Test parsing pipeline with optional connectors and edges."""
yaml_content = """\
model_type: test_model
stages:
- stage_id: 0
model_stage: entry
stage_type: llm
input_sources: []
connectors:
type: ray
edges:
- from: 0
to: 1
"""
yaml_file = tmp_path / "with_connectors.yaml"
yaml_file.write_text(yaml_content)
pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "test_model")
assert pipeline.connectors == {"type": "ray"}
assert pipeline.edges == [{"from": 0, "to": 1}]
def test_parsed_pipeline_passes_validation(self, tmp_path):
"""Test that a well-formed YAML produces a valid pipeline."""
yaml_content = """\
model_type: valid_model
stages:
- stage_id: 0
model_stage: entry
stage_type: llm
input_sources: []
final_output: true
final_output_type: text
- stage_id: 1
model_stage: next
stage_type: llm
input_sources: [0]
"""
yaml_file = tmp_path / "valid.yaml"
yaml_file.write_text(yaml_content)
pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "valid_model")
errors = pipeline.validate_pipeline()
assert errors == [], f"Unexpected validation errors: {errors}"
def test_parse_diffusion_stage_type(self, tmp_path):
"""Test parsing a diffusion stage type from YAML."""
yaml_content = """\
model_type: diff_model
stages:
- stage_id: 0
model_stage: dit
stage_type: diffusion
input_sources: []
final_output: true
final_output_type: image
"""
yaml_file = tmp_path / "diffusion.yaml"
yaml_file.write_text(yaml_content)
pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "diff_model")
assert pipeline.stages[0].stage_type == StageType.DIFFUSION
def test_parse_mimo_audio_yaml(self, tmp_path):
"""Test parsing the MiMo Audio pipeline YAML."""
yaml_content = """\
model_type: mimo_audio
stages:
- stage_id: 0
model_stage: fused_thinker_talker
stage_type: llm
input_sources: []
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
is_comprehension: true
final_output: true
final_output_type: text
- stage_id: 1
model_stage: code2wav
stage_type: llm
input_sources: [0]
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.mimo_audio.llm2code2wav
final_output: true
final_output_type: audio
"""
yaml_file = tmp_path / "mimo_audio.yaml"
yaml_file.write_text(yaml_content)
pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "mimo_audio")
assert pipeline.model_type == "mimo_audio"
assert len(pipeline.stages) == 2
# Stage 0: fused_thinker_talker
s0 = pipeline.stages[0]
assert s0.stage_id == 0
assert s0.model_stage == "fused_thinker_talker"
assert s0.input_sources == []
assert s0.worker_type == "ar"
assert s0.is_comprehension is True
assert s0.final_output is True
assert s0.final_output_type == "text"
# Stage 1: code2wav
s1 = pipeline.stages[1]
assert s1.stage_id == 1
assert s1.input_sources == [0]
assert s1.worker_type == "generation"
assert s1.custom_process_input_func == (
"vllm_omni.model_executor.stage_input_processors.mimo_audio.llm2code2wav"
)
assert s1.final_output is True
assert s1.final_output_type == "audio"
# Pipeline validation
errors = pipeline.validate_pipeline()
assert errors == [], f"Unexpected errors: {errors}"
class TestAsyncChunk:
"""Tests for async_chunk pipeline flag."""
def test_model_pipeline_async_chunk_default(self):
"""Test ModelPipeline defaults async_chunk to False."""
stages = [StageConfig(stage_id=0, model_stage="entry", input_sources=[])]
pipeline = ModelPipeline(model_type="test", stages=stages)
assert pipeline.async_chunk is False
def test_model_pipeline_async_chunk_set(self):
"""Test ModelPipeline with async_chunk=True."""
stages = [StageConfig(stage_id=0, model_stage="entry", input_sources=[])]
pipeline = ModelPipeline(model_type="test", stages=stages, async_chunk=True)
assert pipeline.async_chunk is True
def test_parse_async_chunk_from_yaml(self, tmp_path):
"""Test that async_chunk is parsed from pipeline YAML."""
yaml_content = """\
model_type: test_async
async_chunk: true
stages:
- stage_id: 0
model_stage: entry
stage_type: llm
input_sources: []
"""
yaml_file = tmp_path / "async.yaml"
yaml_file.write_text(yaml_content)
pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "test_async")
assert pipeline.async_chunk is True
def test_parse_missing_async_chunk_defaults_false(self, tmp_path):
"""Test that missing async_chunk defaults to False."""
yaml_content = """\
model_type: test_no_async
stages:
- stage_id: 0
model_stage: entry
stage_type: llm
input_sources: []
"""
yaml_file = tmp_path / "no_async.yaml"
yaml_file.write_text(yaml_content)
pipeline = StageConfigFactory._parse_pipeline_yaml(yaml_file, "test_no_async")
assert pipeline.async_chunk is False
class TestArchitectureFallback:
"""Tests for architecture-based model detection fallback."""
def test_architecture_models_mapping_exists(self):
"""Test that _ARCHITECTURE_MODELS contains expected entries."""
assert "MiMoAudioForConditionalGeneration" in StageConfigFactory._ARCHITECTURE_MODELS
assert StageConfigFactory._ARCHITECTURE_MODELS["MiMoAudioForConditionalGeneration"] == "mimo_audio"
assert "HunyuanImage3ForCausalMM" in StageConfigFactory._ARCHITECTURE_MODELS
assert StageConfigFactory._ARCHITECTURE_MODELS["HunyuanImage3ForCausalMM"] == "hunyuan_image3"
def test_mimo_audio_in_pipeline_models(self):
"""Test that mimo_audio is registered in PIPELINE_MODELS."""
assert "mimo_audio" in StageConfigFactory.PIPELINE_MODELS
-108
View File
@@ -1,108 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests that parallel_config survives the create_default_diffusion roundtrip.
Regression tests for https://github.com/vllm-project/vllm-omni/issues/1862
"""
from collections.abc import Mapping
import torch
from vllm_omni.config.stage_config import StageConfigFactory
from vllm_omni.diffusion.data import (
DiffusionParallelConfig,
OmniDiffusionConfig,
)
def _roundtrip_diffusion_config(**kwargs) -> OmniDiffusionConfig:
"""Simulate the real path: create_default_diffusion → OmniDiffusionConfig.
Does NOT manually reconstruct parallel_config — relies on
OmniDiffusionConfig.__post_init__ to handle the dict, just like
the production code path does.
"""
stages = StageConfigFactory.create_default_diffusion(kwargs)
engine_args = dict(stages[0]["engine_args"])
return OmniDiffusionConfig.from_kwargs(**engine_args)
class TestParallelConfigPropagation:
"""Core regression tests: parallel_config must survive serialization."""
def test_tp2_roundtrip(self):
pc = DiffusionParallelConfig(tensor_parallel_size=2)
od = _roundtrip_diffusion_config(model="test-model", parallel_config=pc)
assert od.parallel_config.tensor_parallel_size == 2
assert od.parallel_config.world_size == 2
def test_tp4_devices_and_config(self):
pc = DiffusionParallelConfig(tensor_parallel_size=4)
stages = StageConfigFactory.create_default_diffusion({"parallel_config": pc, "model": "x"})
assert stages[0]["runtime"]["devices"] == "0,1,2,3"
# Let __post_init__ reconstruct from dict (real code path)
ea = dict(stages[0]["engine_args"])
od = OmniDiffusionConfig.from_kwargs(**ea)
assert od.parallel_config.tensor_parallel_size == 4
assert od.parallel_config.world_size == 4
def test_sp_config_roundtrip(self):
pc = DiffusionParallelConfig(
tensor_parallel_size=2,
ulysses_degree=2,
ring_degree=1,
)
od = _roundtrip_diffusion_config(model="x", parallel_config=pc)
assert od.parallel_config.ulysses_degree == 2
assert od.parallel_config.ring_degree == 1
def test_cfg_parallel_roundtrip(self):
pc = DiffusionParallelConfig(cfg_parallel_size=2)
od = _roundtrip_diffusion_config(model="x", parallel_config=pc)
assert od.parallel_config.cfg_parallel_size == 2
assert od.parallel_config.world_size == 2
def test_no_parallel_config_defaults_to_tp1(self):
od = _roundtrip_diffusion_config(model="x")
assert od.parallel_config.tensor_parallel_size == 1
assert od.parallel_config.world_size == 1
def test_num_gpus_derived_from_world_size(self):
pc = DiffusionParallelConfig(tensor_parallel_size=2)
od = _roundtrip_diffusion_config(model="x", parallel_config=pc)
assert od.num_gpus == 2
class TestCreateDefaultDiffusion:
"""Verify engine_args structure from create_default_diffusion."""
def test_parallel_config_serialized_as_dict(self):
"""The key fix: parallel_config must appear in engine_args as a dict."""
pc = DiffusionParallelConfig(tensor_parallel_size=2)
stages = StageConfigFactory.create_default_diffusion({"model": "x", "parallel_config": pc})
ea = stages[0]["engine_args"]
assert "parallel_config" in ea
assert isinstance(ea["parallel_config"], Mapping)
assert ea["parallel_config"]["tensor_parallel_size"] == 2
def test_dtype_serialized_as_string(self):
stages = StageConfigFactory.create_default_diffusion({"dtype": torch.float16, "model": "x"})
assert stages[0]["engine_args"]["dtype"] == "torch.float16"
def test_cache_backend_defaults_to_none(self):
stages = StageConfigFactory.create_default_diffusion({"model": "x"})
assert stages[0]["engine_args"]["cache_backend"] == "none"
def test_single_gpu_default_devices(self):
stages = StageConfigFactory.create_default_diffusion({"model": "x"})
assert stages[0]["runtime"]["devices"] == "0"
def test_extra_kwargs_forwarded(self):
stages = StageConfigFactory.create_default_diffusion(
{"model": "x", "enforce_eager": True, "lora_path": "/tmp/lora"}
)
ea = stages[0]["engine_args"]
assert ea["enforce_eager"] is True
assert ea["lora_path"] == "/tmp/lora"
+4 -34
View File
@@ -402,41 +402,13 @@ def rocm_marks(*, res: str, num_cards: int):
return marks + [test_distributed]
def xpu_marks(*, res: str, num_cards: int):
"""
Get a collection of pytest marks to apply for `@xpu_test`.
Args:
res: Resource type, e.g., "B60".
num_cards: Number of GPU cards required.
Returns:
List of pytest marks to apply.
"""
test_platform_detail = pytest.mark.xpu
if res == "B60":
test_resource = pytest.mark.B60
else:
raise ValueError(f"Invalid XPU resource type: {res}. Supported: B60")
marks = [test_resource, test_platform_detail]
if num_cards == 1:
return marks
else:
test_distributed = pytest.mark.distributed_rocm(num_cards=num_cards)
# TODO: add XPU support for `skipif_xpu` marker
return marks + [test_distributed]
def gpu_marks(*, res: str, num_cards: int):
"""
Get a collection of pytest marks to apply for `@gpu_test`.
Platform is automatically determined based on resource type.
Args:
res: Resource type, e.g., "L4", "H100" for CUDA, or "MI325" for ROCm, or "B60" for XPU.
res: Resource type, e.g., "L4", "H100" for CUDA, or "MI325" for ROCm.
num_cards: Number of GPU cards required.
Returns:
@@ -447,8 +419,6 @@ def gpu_marks(*, res: str, num_cards: int):
return [test_platform] + cuda_marks(res=res, num_cards=num_cards)
if res == "MI325":
return [test_platform] + rocm_marks(res=res, num_cards=num_cards)
if res == "B60":
return [test_platform] + xpu_marks(res=res, num_cards=num_cards)
raise ValueError(f"Invalid resource type: {res}. Supported: L4, H100, MI325")
@@ -476,13 +446,13 @@ def npu_marks(*, res: str, num_cards: int):
def hardware_marks(*, res: dict[str, str], num_cards: int | dict[str, int] = 1):
"""
Get a collection of pytest marks to apply for `@hardware_test`,
including CUDA, ROCm, XPU, and NPU,
including CUDA, ROCm, and NPU,
based on the specified platforms and resources.
"""
# Validate platforms
# Don't validate platform details in this decorator
for platform, _ in res.items():
if platform not in ("cuda", "rocm", "xpu", "npu"):
if platform not in ("cuda", "rocm", "npu"):
raise ValueError(f"Unsupported platform: {platform}")
# Normalize num_cards
@@ -503,7 +473,7 @@ def hardware_marks(*, res: dict[str, str], num_cards: int | dict[str, int] = 1):
all_marks: list[pytest.MarkDecorator] = []
for platform, resource in res.items():
cards = num_cards_dict[platform]
if platform == "cuda" or platform == "rocm" or platform == "xpu":
if platform == "cuda" or platform == "rocm":
marks = gpu_marks(res=resource, num_cards=cards)
elif platform == "npu":
marks = npu_marks(res=resource, num_cards=cards)
@@ -1,808 +0,0 @@
#!/usr/bin/env python3
"""
Parse pytest commands from Buildkite test-ready.yml, test-merge.yml, test-nightly.yml;
collect test cases (including parametrized) via pytest --collect-only -q and produce an HTML report.
Usage (run from repo root):
python scripts/buildkite_testcase_statistics.py -o buildkite_testcase_statistics.html
Requires: PyYAML (pip install pyyaml)
"""
from __future__ import annotations
import argparse
import ast
import html
import re
import subprocess
import sys
from pathlib import Path
import yaml
# Repo root (parent of the directory containing this script)
REPO_ROOT = Path(__file__).resolve().parent.parent.parent
BUILDKITE_DIR = REPO_ROOT / ".buildkite"
PIPELINE_FILES = ["test-ready.yml", "test-merge.yml", "test-nightly.yml"]
def load_yaml(path: Path) -> dict:
with open(path, encoding="utf-8") as f:
return yaml.safe_load(f)
PYTEST_CMD_RE = re.compile(r"(?:timeout\s+\S+\s+)?(?:python3? -m\s+)?pytest\s+[^\n&|;]*")
def normalize_commands(step: dict) -> list[str]:
raw = step.get("commands") or []
if isinstance(raw, str):
raw = [raw]
text = "\n".join((c.strip() if isinstance(c, str) else "") for c in raw if c)
text = text.replace("$$", "$")
result = []
for m in PYTEST_CMD_RE.finditer(text):
line = m.group(0).strip()
line_start = text.rfind("\n", 0, m.start()) + 1
before_match = text[line_start : m.start()]
if before_match.lstrip().startswith("#"):
continue
if line:
result.append(line)
return result
def extract_pytest_targets_from_line(line: str) -> list[str]:
"""
Parse pytest test targets (file path or marker expression) from a single command line.
Returns e.g. ['tests/xxx.py'].
When there is no explicit test file but only -m/--run-level, returns
['-m expr'] or ['-m expr --run-level level'] with the full marker args.
"""
if "pytest" not in line:
return []
# Strip leading timeout and anything before pytest
idx = line.find("pytest")
if idx == -1:
return []
rest = line[idx:].strip()
if not rest.startswith("pytest"):
return []
# Simple split on whitespace (no handling of quoted spaces)
parts = re.split(r"\s+", rest, maxsplit=1)
args = parts[1].strip() if len(parts) > 1 else ""
if not args:
return []
targets = []
# Match tests/.../*.py
for m in re.finditer(r"tests/[^\s'\"#]+\.py", args):
path = m.group(0)
if path not in targets:
targets.append(path)
if not targets:
# May be -m/--run-level only; record full marker arguments instead of a generic marker.
extra = _parse_extra_args_from_line(line)
if extra:
# Join into a single string, e.g. "-m core and cpu --run-level nightly"
targets.append(" ".join(extra))
return targets
def get_pytest_targets_from_step(step: dict) -> list[tuple[str, str]]:
"""
Extract all pytest-related targets from a step.
Returns [(target, raw_line), ...] where target is tests/xxx.py or '(marker/run-level)'.
"""
lines = normalize_commands(step)
out = []
seen = set()
for line in lines:
for t in extract_pytest_targets_from_line(line):
if t not in seen:
seen.add(t)
out.append((t, line))
return out
def _parse_extra_args_from_line(raw_line: str) -> list[str]:
"""Extract -m and --run-level args from a raw pytest command line. Strips extra spaces and quotes.
Correctly handles quoted expressions with spaces, e.g. -m 'core_model and cpu'.
"""
extra: list[str] = []
# -m: match -m followed by '...', "..." or a single token (so -m expression is one argv)
m = re.search(r"-m\s+(?:'([^']*)'|\"([^\"]*)\"|(\S+))", raw_line)
if m:
value = (m.group(1) or m.group(2) or m.group(3) or "").strip()
if value:
extra.extend(["-m", value])
r = re.search(r"--run-level\s+(?:'([^']*)'|\"([^\"]*)\"|(\S+))", raw_line)
if r:
value = (r.group(1) or r.group(2) or r.group(3) or "").strip()
if value:
extra.extend(["--run-level", value])
return extra
def _parse_collect_only_stdout(stdout: str, *, raise_on_empty: bool = True, stderr: str = "") -> list[str]:
"""
Parse pytest --collect-only -q stdout into a list of node ids.
Normalizes \\r\\n and \\r to \\n before splitting.
If raise_on_empty and no node ids found, raises RuntimeError; otherwise returns [].
"""
raw = (stdout or "").replace("\r\n", "\n").replace("\r", "\n")
out: list[str] = []
for line in raw.splitlines():
line = line.strip()
if not line or "::" not in line or "test_" not in line:
continue
out.append(line)
if not out and raise_on_empty:
raise RuntimeError(
f"Failed to parse pytest --collect-only output: no node ids found.\nstdout:\n{stdout}\n\nstderr:\n{stderr}"
)
return out
def _resolve_pytest_target(target: str, repo_root: Path, extra: list[str], raw_line: str) -> tuple[list[str], str, int]:
"""
Resolve target + repo_root + extra to (path_args, fallback_path, timeout).
Raises FileNotFoundError/RuntimeError on invalid config.
"""
if target.endswith(".py"):
# Target extracted from buildkite yaml may include glob patterns like
# tests/e2e/online_serving/test_*_expansion.py. Pytest does not reliably
# expand globs across platforms, so expand them here.
rel = Path(target)
if any(ch in target for ch in ["*", "?", "["]):
matches = sorted(repo_root.glob(rel.as_posix()))
if not matches:
raise FileNotFoundError(f"Pytest target glob did not match any files: {repo_root / target}")
return [str(p.resolve()) for p in matches], target.replace("\\", "/"), 90
path = (repo_root / rel).resolve()
if not path.exists():
raise FileNotFoundError(f"Pytest target path does not exist: {path}")
return [str(path)], target.replace("\\", "/"), 60
if not extra:
raise RuntimeError(f"Failed to parse -m/--run-level from pytest line: {raw_line!r}")
if not (repo_root / "tests").exists():
raise FileNotFoundError("tests/ directory not found under repo root")
return ["tests/"], "tests/", 120
def collect_test_names(target: str, repo_root: Path, raw_line: str) -> list[str]:
"""
Run pytest --collect-only -q for a target and return test node id list (incl. parametrized).
target: path like tests/foo.py or marker string (e.g. "-m expr"). raw_line used for -m/--run-level.
"""
extra = _parse_extra_args_from_line(raw_line)
path_args, _fallback, timeout_quiet = _resolve_pytest_target(target, repo_root, extra, raw_line)
cmd_quiet = [sys.executable, "-m", "pytest", *path_args, "--collect-only", "-q"]
cmd_quiet.extend(extra)
result = subprocess.run(
cmd_quiet,
cwd=repo_root,
capture_output=True,
text=True,
timeout=timeout_quiet,
)
if result.returncode != 0:
raise RuntimeError(f"pytest --collect-only failed for {path_args} with args {extra}: {result.stderr.strip()}")
print(f"pytest --collect-only success for {path_args} with args {extra}: {result.stdout.strip()}")
return _parse_collect_only_stdout(result.stdout or "", stderr=result.stderr or "")
def _extra_with_skip_marker(extra: list[str]) -> list[str]:
"""Build extra args that add -m skip: if extra has -m expr then use 'expr and skip', else -m skip."""
out: list[str] = []
i = 0
has_m = False
while i < len(extra):
if extra[i] == "-m" and i + 1 < len(extra):
out.extend(["-m", extra[i + 1] + " and skip"])
has_m = True
i += 2
elif extra[i] == "--run-level" and i + 1 < len(extra):
out.extend(["--run-level", extra[i + 1]])
i += 2
else:
i += 1
if not has_m:
out = ["-m", "skip"] + out
return out
def get_skip_status(target: str, raw_line: str, repo_root: Path) -> set[str]:
"""
Use pytest -m skip --collect-only -q to collect tests marked with skip; return their node ids.
Does not run tests; only collects. Supports both pytest xxx.py and pytest -m "marker".
"""
extra = _parse_extra_args_from_line(raw_line)
try:
path_args, _, timeout_quiet = _resolve_pytest_target(target, repo_root, extra, raw_line)
except (FileNotFoundError, RuntimeError):
return set()
extra_skip = _extra_with_skip_marker(extra)
cmd = [sys.executable, "-m", "pytest", *path_args, "--collect-only", "-q"]
cmd.extend(extra_skip)
try:
result = subprocess.run(
cmd,
cwd=repo_root,
capture_output=True,
text=True,
timeout=timeout_quiet,
)
except (subprocess.TimeoutExpired, FileNotFoundError):
return set()
if result.returncode != 0:
return set()
ids = _parse_collect_only_stdout(result.stdout or "", raise_on_empty=False)
return set(ids)
def get_docstring_for_node_id(repo_root: Path, node_id: str) -> str:
"""
Parse the docstring of the test_ function/method from source by pytest node id.
Supports path::test_func, path::TestClass::test_method, path::test_func[param], etc.
"""
if not node_id or "::" not in node_id or node_id.startswith("["):
return ""
parts = node_id.split("::")
# Path accepts forward slashes; no platform-specific replacement needed
path = (repo_root / parts[0]).resolve()
if not path.exists():
return ""
# Strip parametrize suffix, e.g. test_foo[param] -> test_foo
func_part = parts[-1]
if "[" in func_part:
func_part = func_part.split("[", 1)[0]
try:
source = path.read_text(encoding="utf-8", errors="replace")
tree = ast.parse(source)
except Exception:
return ""
# parts: [file, maybe_class, func] or [file, func]
if len(parts) == 2:
# path::test_func
target_func = func_part
target_class = None
else:
# path::TestClass::test_method
target_class = parts[1]
target_func = func_part
for node in tree.body:
if isinstance(node, ast.FunctionDef) and node.name == target_func and target_class is None:
doc = ast.get_docstring(node)
return (doc or "").strip()
if isinstance(node, ast.ClassDef) and node.name == target_class:
for item in node.body:
if isinstance(item, ast.FunctionDef) and item.name == target_func:
doc = ast.get_docstring(item)
return (doc or "").strip()
break
return ""
def get_test_file_for_node_id(node_id: str) -> str:
"""Extract test file path from a pytest node id."""
if not node_id:
return ""
if "::" in node_id:
return node_id.split("::", 1)[0]
if node_id.endswith(".py"):
return node_id
return "(marker-only / no collection)"
def write_html(all_stats: list[tuple], out_path: Path, total: int, repo_root: Path | None = None) -> None:
"""Write stats to HTML with styling and responsive tables."""
root = repo_root if repo_root is not None else REPO_ROOT
pipeline_badges = {
"test-ready": "ready",
"test-merge": "merge",
"test-nightly": "nightly",
}
def esc(s: str) -> str:
return html.escape(str(s))
def tr_class(i: int) -> str:
return "odd" if i % 2 == 0 else "even"
pipeline_totals: dict[str, int] = {}
for item in all_stats:
_, pipeline_name, _, _, _, _, count, _, _, _ = item
pipeline_totals[pipeline_name] = pipeline_totals.get(pipeline_name, 0) + count
pipeline_summary_rows = []
for idx, (pipeline_name, pipeline_count) in enumerate(sorted(pipeline_totals.items())):
badge = pipeline_badges.get(pipeline_name, "")
badge_span = f'<span class="badge {badge}">{esc(pipeline_name)}</span>' if badge else esc(pipeline_name)
pipeline_summary_rows.append(
f'<tr class="{tr_class(idx)}"><td>{badge_span}</td><td class="num">{pipeline_count}</td></tr>'
)
# Summary rows by module (item: ..., names, docstrings, skipped_ids)
summary_rows = []
for i, item in enumerate(all_stats):
_, pipeline_name, label, test_files, _, marker_targets, count, _, _, _ = item
if test_files:
files_html = "<br>".join(esc(f) for f in test_files)
elif marker_targets:
# Display full marker / run-level arguments when no explicit test files.
files_html = "<br>".join(esc(m) for m in marker_targets)
else:
files_html = "<em>No explicit tests or markers</em>"
badge = pipeline_badges.get(pipeline_name, "")
badge_span = f'<span class="badge {badge}">{esc(pipeline_name)}</span>' if badge else esc(pipeline_name)
summary_rows.append(
f'<tr class="{tr_class(i)}">'
f'<td>{badge_span}</td><td>{esc(label)}</td><td class="num">{count}</td>'
f'<td class="files">{files_html}</td></tr>'
)
# Detail rows grouped by Pipeline, then by Test Suite (two-level collapsible)
case_details_thead = """
<tr>
<th>Pipeline</th>
<th>Test Suite</th>
<th>Test File</th>
<th>No.</th>
<th>Test Name</th>
<th>Status</th>
<th>Description</th>
</tr>"""
# Per-item: (pipeline_name, label, count, badge_span, suite_rows_html)
pipeline_suites: dict[str, list[tuple[str, int, str, str]]] = {} # pipeline -> [(label, count, summary, body)]
for item in all_stats:
_, pipeline_name, label, _, _, _, count, names, docstrings, skipped_ids = item
badge = pipeline_badges.get(pipeline_name, "")
badge_span = f'<span class="badge {badge}">{esc(pipeline_name)}</span>' if badge else esc(pipeline_name)
suite_summary = f"{esc(label)} ({count} cases)"
file_sections: list[str] = []
if not names:
suite_rows = [
f'<tr class="{tr_class(0)}">'
f"<td>{badge_span}</td><td>{esc(label)}</td>"
f'<td class="files"><em>(marker-only / no collection)</em></td><td class="num">0</td>'
f'<td class="name"><em>No collection or only -m/--run-level</em></td>'
f'<td class="status"></td><td class="desc"></td></tr>'
]
file_sections.append(
'<details class="file-details" open>'
"<summary>(marker-only / no collection)</summary>"
f'<div class="table-wrap"><table><thead>{case_details_thead}</thead>'
f"<tbody>{''.join(suite_rows)}</tbody></table></div></details>"
)
else:
file_groups: dict[str, list[tuple[str, str, bool]]] = {}
for idx, name in enumerate(names):
desc = (docstrings[idx] if idx < len(docstrings) else "") or ""
if not desc and "::" in name and not name.startswith("["):
desc = get_docstring_for_node_id(root, name)
file_key = get_test_file_for_node_id(name)
file_groups.setdefault(file_key, []).append((name, desc, name in skipped_ids))
for file_name in sorted(file_groups.keys()):
suite_rows = []
for i, (name, desc, is_skipped) in enumerate(file_groups[file_name], 1):
desc_html = esc(desc).replace("\n", "<br>") if desc else ""
status_html = '<span class="badge skip">Skipped</span>' if is_skipped else "—"
suite_rows.append(
f'<tr class="{tr_class(i - 1)}">'
f'<td>{badge_span}</td><td>{esc(label)}</td><td class="files">{esc(file_name)}</td>'
f'<td class="num">{i}</td>'
f'<td class="name"><code>{esc(name)}</code></td>'
f'<td class="status">{status_html}</td><td class="desc">{desc_html}</td></tr>'
)
file_sections.append(
f'<details class="file-details" data-test-file="{esc(file_name.lower())}">'
f"<summary>{esc(file_name)} ({len(file_groups[file_name])} cases)</summary>"
f'<div class="table-wrap"><table><thead>{case_details_thead}</thead>'
f"<tbody>{''.join(suite_rows)}</tbody></table></div></details>"
)
body = "".join(file_sections)
if pipeline_name not in pipeline_suites:
pipeline_suites[pipeline_name] = []
pipeline_suites[pipeline_name].append((label, count, suite_summary, body))
# Build two-level HTML: Pipeline (outer) -> Test Suite (inner) -> table
detail_sections_html_parts = []
for pipeline_name in sorted(pipeline_badges.keys()):
if pipeline_name not in pipeline_suites:
continue
suites = pipeline_suites[pipeline_name]
pipeline_total = sum(c for _, c, _, _ in suites)
badge = pipeline_badges.get(pipeline_name, "")
badge_span = f'<span class="badge {badge}">{esc(pipeline_name)}</span>' if badge else esc(pipeline_name)
pipeline_summary = f"{badge_span} ({pipeline_total} cases)"
inner_html = "".join(
f'<details class="suite-details"><summary>{suite_summary}</summary>'
f'<div class="case-files-inner">{body}</div></details>'
for _, _, suite_summary, body in suites
)
detail_sections_html_parts.append(
f'<details class="pipeline-details"><summary>{pipeline_summary}</summary>'
f'<div class="case-details-inner">{inner_html}</div></details>'
)
# Include pipelines that might be in all_stats but not in pipeline_badges (e.g. custom names)
for pipeline_name in sorted(pipeline_suites.keys()):
if pipeline_name in pipeline_badges:
continue
suites = pipeline_suites[pipeline_name]
pipeline_total = sum(c for _, c, _, _ in suites)
pipeline_summary = f"{esc(pipeline_name)} ({pipeline_total} cases)"
inner_html = "".join(
f'<details class="suite-details"><summary>{suite_summary}</summary>'
f'<div class="case-files-inner">{body}</div></details>'
for _, _, suite_summary, body in suites
)
detail_sections_html_parts.append(
f'<details class="pipeline-details"><summary>{pipeline_summary}</summary>'
f'<div class="case-details-inner">{inner_html}</div></details>'
)
detail_sections_html = "".join(detail_sections_html_parts)
html_content = f"""<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1">
<title>Buildkite Test Case Statistics</title>
<style>
:root {{
--bg: #0f1419;
--card: #1a2332;
--border: #2d3a4f;
--text: #e6edf3;
--muted: #8b949e;
--accent: #58a6ff;
--ready: #3fb950;
--merge: #a371f7;
--nightly: #f0883e;
}}
* {{ box-sizing: border-box; }}
body {{
font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", "Noto Sans SC", sans-serif;
background: var(--bg);
color: var(--text);
line-height: 1.6;
margin: 0;
padding: 2rem;
}}
.container {{ max-width: 1200px; margin: 0 auto; }}
h1 {{
font-size: 1.75rem;
font-weight: 600;
margin-bottom: 0.5rem;
}}
.meta {{ color: var(--muted); font-size: 0.9rem; margin-bottom: 2rem; }}
section {{
background: var(--card);
border: 1px solid var(--border);
border-radius: 12px;
padding: 1.25rem 1.5rem;
margin-bottom: 1.5rem;
}}
section h2 {{
font-size: 1.15rem;
font-weight: 600;
margin: 0 0 1rem 0;
padding-bottom: 0.5rem;
border-bottom: 1px solid var(--border);
color: var(--accent);
}}
.filter-bar {{
display: flex;
gap: 0.75rem;
align-items: center;
flex-wrap: wrap;
margin-bottom: 1rem;
}}
.filter-bar label {{
color: var(--muted);
font-size: 0.9rem;
}}
.filter-input {{
min-width: 320px;
max-width: 520px;
width: 100%;
padding: 0.7rem 0.9rem;
border-radius: 8px;
border: 1px solid var(--border);
background: rgba(255,255,255,0.04);
color: var(--text);
outline: none;
}}
.filter-input:focus {{
border-color: var(--accent);
box-shadow: 0 0 0 2px rgba(88, 166, 255, 0.18);
}}
.filter-hint {{
color: var(--muted);
font-size: 0.82rem;
}}
.table-wrap {{ overflow-x: auto; }}
table {{
width: 100%;
border-collapse: collapse;
font-size: 0.9rem;
}}
th {{
text-align: left;
padding: 0.65rem 0.85rem;
background: rgba(88, 166, 255, 0.12);
color: var(--accent);
font-weight: 600;
border: 1px solid var(--border);
}}
th:first-child {{ border-radius: 8px 0 0 0; }}
th:last-child {{ border-radius: 0 8px 0 0; }}
td {{
padding: 0.6rem 0.85rem;
border: 1px solid var(--border);
vertical-align: top;
}}
tr.odd td {{ background: rgba(255,255,255,0.02); }}
tr.even td {{ background: transparent; }}
tr:hover td {{ background: rgba(88, 166, 255, 0.06); }}
td.num {{ text-align: right; font-variant-numeric: tabular-nums; }}
td.files {{ font-size: 0.85rem; color: var(--muted); }}
td.name code {{ font-size: 0.8rem; background: rgba(0,0,0,0.3); padding: 0.15em 0.4em; border-radius: 4px; }}
td.desc {{ font-size: 0.85rem; color: var(--muted); max-width: 320px; }}
.badge {{
display: inline-block;
padding: 0.2em 0.5em;
border-radius: 6px;
font-size: 0.8rem;
font-weight: 500;
}}
.badge.ready {{ background: rgba(63, 185, 80, 0.2); color: var(--ready); }}
.badge.merge {{ background: rgba(163, 113, 247, 0.2); color: var(--merge); }}
.badge.nightly {{ background: rgba(240, 136, 62, 0.2); color: var(--nightly); }}
.badge.skip {{ background: rgba(139, 148, 158, 0.25); color: var(--muted); }}
td.status {{ color: var(--muted); font-size: 0.85rem; }}
.case-details details {{
margin-bottom: 0.5rem;
border: 1px solid var(--border);
border-radius: 8px;
overflow: hidden;
}}
.case-details details summary {{
padding: 0.65rem 0.85rem;
cursor: pointer;
font-weight: 500;
list-style: none;
user-select: none;
}}
.case-details details summary::-webkit-details-marker {{ display: none; }}
.case-details details summary::before {{
content: "▶ ";
font-size: 0.7rem;
color: var(--muted);
margin-right: 0.35rem;
}}
.case-details details[open] summary::before {{ content: "▼ "; }}
.case-details details summary:hover {{ background: rgba(88, 166, 255, 0.06); }}
.case-details details .table-wrap {{ margin: 0; border-top: 1px solid var(--border); }}
.case-details details table {{ margin: 0; }}
.case-details details th {{ border-top: 0; }}
.case-details .pipeline-details {{ margin-bottom: 0.5rem; }}
.case-details .pipeline-details .case-details-inner {{ padding-left: 0.5rem; margin-top: 0.25rem; }}
.case-details .suite-details {{ margin-bottom: 0.35rem; margin-left: 0.5rem; }}
.case-details .suite-details summary::before {{
content: "▶ "; font-size: 0.7rem; color: var(--muted); margin-right: 0.35rem;
}}
.case-details .suite-details[open] summary::before {{ content: "▼ "; }}
.case-details .case-files-inner {{ padding-left: 0.5rem; margin-top: 0.35rem; }}
.case-details .file-details {{
margin-bottom: 0.35rem;
margin-left: 0.75rem;
border: 1px solid var(--border);
border-radius: 8px;
overflow: hidden;
}}
.case-details .file-details summary {{
padding: 0.55rem 0.8rem;
cursor: pointer;
user-select: none;
background: rgba(255,255,255,0.015);
}}
.case-details .file-details summary::-webkit-details-marker {{ display: none; }}
.case-details .file-details summary::before {{
content: "▶ ";
font-size: 0.7rem;
color: var(--muted);
margin-right: 0.35rem;
}}
.case-details .file-details[open] summary::before {{ content: "▼ "; }}
.back-to-top {{
position: fixed;
right: 24px;
bottom: 24px;
display: inline-block;
padding: 0.7rem 0.9rem;
border-radius: 999px;
background: rgba(88, 166, 255, 0.18);
border: 1px solid var(--border);
color: var(--text);
text-decoration: none;
font-size: 0.9rem;
backdrop-filter: blur(6px);
}}
.back-to-top:hover {{ background: rgba(88, 166, 255, 0.28); }}
</style>
</head>
<body id="top">
<div class="container">
<h1>Buildkite Pytest Case Statistics</h1>
<p class="meta">Test Suites: {len(all_stats)} · Total cases (sum of steps): {total} · </p>
<section>
<h2>Pipeline Summary</h2>
<div class="table-wrap">
<table>
<thead>
<tr><th>Pipeline</th><th>Total Cases</th></tr>
</thead>
<tbody>
{"".join(pipeline_summary_rows)}
</tbody>
</table>
</div>
</section>
<section>
<h2>Summary by Test Suite</h2>
<div class="table-wrap">
<table>
<thead>
<tr><th>Pipeline</th><th>Test Suite</th><th>Cases</th><th>Test Files</th></tr>
</thead>
<tbody>
{"".join(summary_rows)}
</tbody>
</table>
</div>
</section>
<section class="case-details">
<h2>Case Details</h2>
<p class="meta">Expand by Pipeline, then by Test Suite, then by Test File.</p>
<div class="filter-bar">
<label for="test-file-filter">Filter by Test File</label>
<input
id="test-file-filter"
class="filter-input"
type="text"
placeholder="Type part of a test file path, e.g. tests/e2e/online_serving"
/>
<span class="filter-hint">Filter applies to the Test File fold sections below.</span>
</div>
{detail_sections_html}
</section>
</div>
<a href="#top" class="back-to-top" aria-label="Back to top">↑ Top</a>
<script>
(() => {{
const input = document.getElementById("test-file-filter");
if (!input) return;
const applyFilter = () => {{
const query = input.value.trim().toLowerCase();
document.querySelectorAll(".pipeline-details").forEach((pipeline) => {{
let pipelineVisible = false;
pipeline.querySelectorAll(".suite-details").forEach((suite) => {{
let suiteVisible = false;
suite.querySelectorAll(".file-details").forEach((fileDetail) => {{
const fileName = (fileDetail.dataset.testFile || fileDetail.textContent || "").toLowerCase();
const matched = !query || fileName.includes(query);
fileDetail.style.display = matched ? "" : "none";
if (matched) suiteVisible = true;
}});
suite.style.display = suiteVisible ? "" : "none";
if (suiteVisible) pipelineVisible = true;
}});
pipeline.style.display = pipelineVisible ? "" : "none";
}});
}};
input.addEventListener("input", applyFilter);
applyFilter();
}})();
</script>
</body>
</html>
"""
out_path.write_text(html_content, encoding="utf-8")
def main() -> None:
parser = argparse.ArgumentParser(description="Count pytest cases from Buildkite pipelines")
parser.add_argument(
"--buildkite-dir",
type=Path,
default=BUILDKITE_DIR,
help="Path to .buildkite directory",
)
default_out = REPO_ROOT / "buildkite_testcase_statistics.html"
parser.add_argument(
"-o",
"--output",
type=Path,
default=default_out,
help=f"Output HTML file path (default: {default_out})",
)
args = parser.parse_args()
all_stats = []
for filename in PIPELINE_FILES:
yml_path = args.buildkite_dir / filename
if not yml_path.exists():
print(f"Skipped (not found): {yml_path}", file=sys.stderr)
continue
data = load_yaml(yml_path)
steps = data.get("steps") or []
pipeline_name = filename.replace(".yml", "")
for step in steps:
if not isinstance(step, dict):
continue
label = step.get("label") or "(no label)"
# Skip commented-out steps (usually no label or commands)
if not step.get("commands"):
continue
targets = get_pytest_targets_from_step(step)
if not targets:
continue
module_key = f"{pipeline_name} / {label}"
# test_files: only real .py files; marker-only steps will have an empty list.
test_files = [t for t, _ in targets if t.endswith(".py")]
# marker_targets: synthetic marker targets (e.g. "-m expr", "-m expr --run-level level").
marker_targets = [t for t, _ in targets if not t.endswith(".py")]
# marker_only: all targets are marker expressions (no explicit .py files).
marker_only = bool(targets) and all(not t.endswith(".py") for t, _ in targets)
names = []
skipped_ids: set[str] = set()
for target, raw in targets:
collected = collect_test_names(target, REPO_ROOT, raw)
skipped_ids.update(get_skip_status(target, raw, REPO_ROOT))
names.extend(collected)
count = len(names) - len(skipped_ids)
docstrings = [
get_docstring_for_node_id(REPO_ROOT, n) if "::" in n and not n.startswith("[") else "" for n in names
]
all_stats.append(
(
module_key,
pipeline_name,
label,
test_files,
marker_only,
marker_targets,
count,
names,
docstrings,
skipped_ids,
)
)
out_path = args.output.resolve()
out_path.parent.mkdir(parents=True, exist_ok=True)
total = sum(item[6] for item in all_stats)
write_html(all_stats, out_path, total, REPO_ROOT)
print(f"Written: {out_path}")
print(f"Test Suites: {len(all_stats)}, total cases (sum of steps): {total}")
if __name__ == "__main__":
main()
-351
View File
@@ -1,351 +0,0 @@
#!/usr/bin/env python3
"""
Send the nightly report by email to DAILY_EMAIL_LIST.
Reads SMTP and recipient config from environment variables. Use --dry-run to
validate config and print subject/body without sending.
"""
from __future__ import annotations
import argparse
import logging
import mimetypes
import os
import smtplib
import sys
from datetime import date
from email.header import Header
from email.mime.application import MIMEApplication
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
LOGGER = logging.getLogger(__name__)
# Env keys for SMTP and recipients.
ENV_SMTP_HOST = "SMTP_HOST"
ENV_SMTP_PORT = "SMTP_PORT"
ENV_SMTP_USERNAME = "SMTP_USERNAME"
ENV_SMTP_PASSWORD = "SMTP_PASSWORD"
ENV_DAILY_EMAIL_LIST = "DAILY_EMAIL_LIST"
ENV_EMAIL_SENDER = "EMAIL_SENDER"
ENV_EMAIL_SUBJECT_PREFIX = "EMAIL_SUBJECT_PREFIX"
ENV_BUILD_URL = "BUILDKITE_BUILD_URL"
ENV_COMMIT = "BUILDKITE_COMMIT"
DEFAULT_OUTPUT_DIR = os.getenv("DEFAULT_OUTPUT_DIR")
# Do not attach file if size >= this (bytes); body will suggest downloading from build URL.
MAX_ATTACHMENT_BYTES = 20 * 1024 * 1024
SMTP_RETRIES = 3
SMTP_RETRY_DELAY_SEC = 5
def _strip_quoted(s: str) -> str:
"""Remove one layer of surrounding quote characters from env values.
Handles both ASCII quotes ('", ''') and common Unicode quotes (“ ” ‘ ’).
"""
if not s:
return s
s = s.strip()
if not s:
return s
quotes = {'"', "'", "“", "”", "‘", "’"}
if len(s) >= 2 and s[0] in quotes and s[-1] in quotes:
s = s[1:-1].strip()
return s
def _get_required_env() -> dict[str, str]:
"""Read required env vars; raise SystemExit with clear message if any missing."""
required = {
ENV_SMTP_HOST: os.environ.get(ENV_SMTP_HOST),
ENV_SMTP_PORT: os.environ.get(ENV_SMTP_PORT),
ENV_SMTP_USERNAME: os.environ.get(ENV_SMTP_USERNAME),
ENV_SMTP_PASSWORD: os.environ.get(ENV_SMTP_PASSWORD),
ENV_DAILY_EMAIL_LIST: os.environ.get(ENV_DAILY_EMAIL_LIST),
}
missing = [k for k, v in required.items() if not (v and str(v).strip())]
if missing:
raise SystemExit(f"Missing required env vars: {', '.join(missing)}. Set them (e.g. in Buildkite secrets).")
return {k: _strip_quoted(str(v)) for k, v in required.items()}
def _get_latest_file(path: str) -> str | list[str]:
"""
Resolve report file patterns into concrete files.
- If path is a directory: return all files in that directory (list).
- If path contains wildcards (e.g. *.xlsx, *.html): expand using glob and
return all matching files (list).
- If path is a plain filename or file path: treat as (dir, filename),
search that directory for files with the same basename and return the
latest modified one (str).
Note: higher-level callers will turn a single str into [str] when they
want to attach multiple files.
"""
import glob
# Support comma-/semicolon-separated list: e.g. "*.xlsx, *.html; reports/foo.xlsx"
parts = [p.strip() for p in path.split(";") for p in p.split(",") if p and p.strip()]
if len(parts) > 1:
files: list[str] = []
for part in parts:
files_part = _get_latest_file(part)
if isinstance(files_part, str):
files.append(files_part)
else:
files.extend(files_part)
if not files:
raise SystemExit(f"No files resolved from list: {path}")
return sorted(set(files))
# Single expression after splitting
path = parts[0] if parts else path
# Directory: return all files inside
if os.path.isdir(path):
files = [os.path.join(path, f) for f in os.listdir(path) if os.path.isfile(os.path.join(path, f))]
if not files:
raise SystemExit(f"No files found in {path}")
return sorted(files)
# Wildcard pattern: expand via glob
if any(ch in path for ch in ("*", "?", "[")) and not os.path.exists(path):
matches = [p for p in glob.glob(path) if os.path.isfile(p)]
if not matches:
raise SystemExit(f"No files match pattern {path!r}")
return sorted(matches)
# Plain filename or explicit file path: search its directory for same basename
parent = os.path.dirname(path) or "."
base = os.path.basename(path)
if not base:
raise SystemExit(f"Invalid path (no filename): {path}")
if not os.path.isdir(parent):
raise SystemExit(f"Directory not found: {parent}")
candidates = [
p
for p in (os.path.join(parent, f) for f in os.listdir(parent))
if os.path.basename(p) == base and os.path.isfile(p)
]
if not candidates:
raise SystemExit(f"No file named {base!r} found in {parent}")
return max(candidates, key=os.path.getmtime)
def _recipients_list(comma_separated: str) -> list[str]:
"""Parse DAILY_EMAIL_LIST into a list of addresses."""
return [_strip_quoted(a.strip()) for a in comma_separated.split(",") if a and a.strip()]
def _build_body(
date_str: str,
commit_sha: str | None,
build_url: str | None,
attachment_skipped: bool,
) -> str:
"""Plain-text body with build metadata; note when attachment was skipped."""
lines = [
f"Nightly report date: {date_str}",
"",
f"Commit: {commit_sha or 'N/A'}",
f"Build: {build_url or 'N/A'}",
"",
]
if attachment_skipped:
lines.append(
"Some report file(s) were too large to attach. Please download from the build artifacts (Build URL above)."
)
lines.append("")
return "\n".join(lines)
def _build_subject(prefix: str | None, date_str: str) -> str:
"""Subject line: optional prefix + date."""
base = f"Nightly Report {date_str}"
if prefix and prefix.strip():
return f"{prefix.strip()} {base}"
return base
def _mime_type_and_subtype(file_path: str) -> tuple[str, str]:
"""Return (maintype, subtype) for the file; fallback to application/octet-stream."""
full_type, _ = mimetypes.guess_type(file_path, strict=False)
if not full_type:
return "application", "octet-stream"
parts = full_type.split("/", 1)
if len(parts) != 2:
return "application", "octet-stream"
return parts[0], parts[1]
def _attach_file(msg: MIMEMultipart, file_path: str, max_bytes: int) -> bool:
"""
Attach one file to msg with correct MIME type. Return True if attached, False if skipped (e.g. size).
"""
if not os.path.isfile(file_path):
return False
size = os.path.getsize(file_path)
if size <= 0 or size >= max_bytes:
return False
maintype, subtype = _mime_type_and_subtype(file_path)
name = os.path.basename(file_path)
try:
if maintype == "text":
with open(file_path, encoding="utf-8", errors="replace") as f:
part = MIMEText(f.read(), subtype, "utf-8")
else:
with open(file_path, "rb") as f:
part = MIMEApplication(f.read(), _subtype=subtype)
part.add_header("Content-Disposition", "attachment", filename=name)
msg.attach(part)
return True
except OSError:
return False
def _send_mail(
report_files: list[str],
date_str: str,
dry_run: bool,
) -> None:
"""Load config, build message (attach each file with correct MIME type), and send (or dry-run)."""
cfg = _get_required_env()
recipients = _recipients_list(cfg[ENV_DAILY_EMAIL_LIST])
if not recipients:
raise SystemExit("DAILY_EMAIL_LIST is empty after parsing.")
commit_sha = os.environ.get(ENV_COMMIT)
build_url = os.environ.get(ENV_BUILD_URL)
# Allow EMAIL_SENDER / EMAIL_SUBJECT_PREFIX to be quoted; strip quotes if present.
raw_sender = os.environ.get(ENV_EMAIL_SENDER)
sender = _strip_quoted(raw_sender) if raw_sender else cfg[ENV_SMTP_USERNAME]
raw_prefix = os.environ.get(ENV_EMAIL_SUBJECT_PREFIX)
prefix = _strip_quoted(raw_prefix) if raw_prefix else None
attachment_skipped = any(os.path.isfile(p) and os.path.getsize(p) >= MAX_ATTACHMENT_BYTES for p in report_files)
body = _build_body(
commit_sha=commit_sha, build_url=build_url, date_str=date_str, attachment_skipped=attachment_skipped
)
subject = _build_subject(prefix=prefix, date_str=date_str)
msg = MIMEMultipart()
msg["Subject"] = Header(subject, "utf-8")
msg["From"] = sender
msg["To"] = ", ".join(recipients)
msg.attach(MIMEText(body, "plain", "utf-8"))
attached_count = 0
for file_path in report_files:
if os.path.isfile(file_path) and os.path.getsize(file_path) >= MAX_ATTACHMENT_BYTES:
attachment_skipped = True
else:
if _attach_file(msg, file_path, MAX_ATTACHMENT_BYTES):
attached_count += 1
elif os.path.isfile(file_path):
attachment_skipped = True
if dry_run:
LOGGER.info("dry-run: not sending mail")
print("To:", recipients, file=sys.stderr)
print("Subject:", subject, file=sys.stderr)
print("Attachments:", attached_count, "(skipped large:", attachment_skipped, ")", file=sys.stderr)
print("Body preview:", body[:300] + ("..." if len(body) > 300 else ""), file=sys.stderr)
return
host = cfg[ENV_SMTP_HOST]
port = int(cfg[ENV_SMTP_PORT], 10)
use_ssl = port == 465
last_err: Exception | None = None
for attempt in range(SMTP_RETRIES):
try:
if use_ssl:
with smtplib.SMTP_SSL(host, port=port, timeout=30) as smtp:
smtp.login(cfg[ENV_SMTP_USERNAME], cfg[ENV_SMTP_PASSWORD])
smtp.sendmail(sender, recipients, msg.as_string())
else:
with smtplib.SMTP(host, port=port, timeout=30) as smtp:
smtp.starttls()
smtp.login(cfg[ENV_SMTP_USERNAME], cfg[ENV_SMTP_PASSWORD])
smtp.sendmail(sender, recipients, msg.as_string())
LOGGER.info("sent nightly email to %d recipient(s)", len(recipients))
return
except Exception as e:
last_err = e
LOGGER.warning("SMTP attempt %d/%d failed: %s", attempt + 1, SMTP_RETRIES, e)
if attempt < SMTP_RETRIES - 1:
import time
time.sleep(SMTP_RETRY_DELAY_SEC)
hint = ""
if last_err and "getaddrinfo" in str(last_err).lower():
hint = f" (check SMTP_HOST={host!r} is correct and reachable, e.g. nslookup {host})"
raise SystemExit(f"Failed to send email after {SMTP_RETRIES} attempts.{hint}") from last_err
def _vllm_omni_root() -> str:
"""Resolve vllm-omni repo root: directory that contains a 'tests' subdir (and usually 'tools')."""
path = os.path.dirname(os.path.abspath(__file__))
while path and path != os.path.dirname(path):
if os.path.isdir(os.path.join(path, "tests")):
return path
path = os.path.dirname(path)
return os.path.normpath(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", ".."))
def _default_output_dir() -> str:
"""Default: vllm-omni root / DEFAULT_OUTPUT_DIR (where nightly report files live)."""
root = _vllm_omni_root()
subdir = DEFAULT_OUTPUT_DIR if (DEFAULT_OUTPUT_DIR and DEFAULT_OUTPUT_DIR.strip()) else "."
return os.path.join(root, subdir)
def parse_args() -> argparse.Namespace:
"""Parse command-line arguments."""
parser = argparse.ArgumentParser(
description="Send nightly report by email (config from env).",
)
parser.add_argument(
"--report-file",
type=str,
default=_default_output_dir(),
help="Folder/file path to the sent email file; default is DEFAULT_OUTPUT_DIR.",
)
parser.add_argument(
"--date",
type=str,
default=None,
help="Date string for subject/body (default: inferred from report filename).",
)
parser.add_argument(
"--dry-run",
action="store_true",
help="Print recipient, subject, and body; do not send.",
)
return parser.parse_args()
def main() -> None:
"""Entrypoint."""
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
)
args = parse_args()
result = _get_latest_file(args.report_file)
report_files = result if isinstance(result, list) else [result]
report_files = [p for p in report_files if os.path.isfile(p)]
if not report_files:
raise SystemExit("No report file(s) found.")
date_str = args.date or date.today().strftime("%Y-%m-%d")
_send_mail(report_files=report_files, date_str=date_str, dry_run=args.dry_run)
if __name__ == "__main__":
main()
+234
View File
@@ -0,0 +1,234 @@
#!/usr/bin/env python3
"""
Send the nightly performance Excel report by email to DAILY_EMAIL_LIST.
Reads SMTP and recipient config from environment variables. Use --dry-run to
validate config and print subject/body without sending.
"""
from __future__ import annotations
import argparse
import logging
import os
import smtplib
import sys
from email.header import Header
from email.mime.application import MIMEApplication
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
LOGGER = logging.getLogger(__name__)
# Env keys for SMTP and recipients.
ENV_SMTP_HOST = "SMTP_HOST"
ENV_SMTP_PORT = "SMTP_PORT"
ENV_SMTP_USERNAME = "SMTP_USERNAME"
ENV_SMTP_PASSWORD = "SMTP_PASSWORD"
ENV_DAILY_EMAIL_LIST = "DAILY_EMAIL_LIST"
ENV_EMAIL_SENDER = "EMAIL_SENDER"
ENV_EMAIL_SUBJECT_PREFIX = "EMAIL_SUBJECT_PREFIX"
ENV_BUILD_URL = "BUILDKITE_BUILD_URL"
ENV_COMMIT = "BUILDKITE_COMMIT"
DEFAULT_OUTPUT_DIR = os.getenv("DEFAULT_OUTPUT_DIR")
# Do not attach Excel if size >= this (bytes); body will suggest downloading from build URL.
MAX_ATTACHMENT_BYTES = 20 * 1024 * 1024
SMTP_RETRIES = 3
SMTP_RETRY_DELAY_SEC = 5
def _get_required_env() -> dict[str, str]:
"""Read required env vars; raise SystemExit with clear message if any missing."""
required = {
ENV_SMTP_HOST: os.environ.get(ENV_SMTP_HOST),
ENV_SMTP_PORT: os.environ.get(ENV_SMTP_PORT),
ENV_SMTP_USERNAME: os.environ.get(ENV_SMTP_USERNAME),
ENV_SMTP_PASSWORD: os.environ.get(ENV_SMTP_PASSWORD),
ENV_DAILY_EMAIL_LIST: os.environ.get(ENV_DAILY_EMAIL_LIST),
}
missing = [k for k, v in required.items() if not (v and str(v).strip())]
if missing:
raise SystemExit(f"Missing required env vars: {', '.join(missing)}. Set them (e.g. in Buildkite secrets).")
return {k: str(v).strip() for k, v in required.items()}
def _get_latest_file(folder_path: str) -> str:
"""Get the latest modified file from the folder path."""
files = [os.path.join(folder_path, f) for f in os.listdir(folder_path) if f.endswith(".xlsx")]
if not files:
raise SystemExit(f"No Excel files found in {folder_path}")
return max(files, key=os.path.getmtime)
def _recipients_list(comma_separated: str) -> list[str]:
"""Parse DAILY_EMAIL_LIST into a list of addresses."""
return [a.strip() for a in comma_separated.split(",") if a.strip()]
def _build_body(
date_str: str,
commit_sha: str | None,
build_url: str | None,
attachment_skipped: bool,
) -> str:
"""Plain-text body with build metadata; note when attachment was skipped."""
lines = [
f"Nightly performance report date: {date_str}",
"",
f"Commit: {commit_sha or 'N/A'}",
f"Build: {build_url or 'N/A'}",
"",
]
if attachment_skipped:
lines.append(
"Report file was too large to attach. Please download the Excel from the build artifacts (Build URL above)."
)
lines.append("")
return "\n".join(lines)
def _build_subject(prefix: str | None, date_str: str) -> str:
"""Subject line: optional prefix + date."""
base = f"Nightly Perf {date_str}"
if prefix and prefix.strip():
return f"{prefix.strip()} {base}"
return base
def _send_mail(
report_file: str,
date_str: str,
dry_run: bool,
) -> None:
"""Load config, build message, and send (or dry-run)."""
cfg = _get_required_env()
recipients = _recipients_list(cfg[ENV_DAILY_EMAIL_LIST])
if not recipients:
raise SystemExit("DAILY_EMAIL_LIST is empty after parsing.")
commit_sha = os.environ.get(ENV_COMMIT)
build_url = os.environ.get(ENV_BUILD_URL)
sender = os.environ.get(ENV_EMAIL_SENDER) or cfg[ENV_SMTP_USERNAME]
prefix = os.environ.get(ENV_EMAIL_SUBJECT_PREFIX)
size = os.path.getsize(report_file) if os.path.isfile(report_file) else 0
attach_excel = size < MAX_ATTACHMENT_BYTES and size > 0
attachment_skipped = os.path.isfile(report_file) and not attach_excel
body = _build_body(
commit_sha=commit_sha, build_url=build_url, date_str=date_str, attachment_skipped=attachment_skipped
)
subject = _build_subject(prefix=prefix, date_str=date_str)
msg = MIMEMultipart()
msg["Subject"] = Header(subject, "utf-8")
msg["From"] = sender
msg["To"] = ", ".join(recipients)
msg.attach(MIMEText(body, "plain", "utf-8"))
if attach_excel:
with open(report_file, "rb") as f:
part = MIMEApplication(f.read(), _subtype="vnd.openxmlformats-officedocument.spreadsheetml.sheet")
part.add_header("Content-Disposition", "attachment", filename=os.path.basename(report_file))
msg.attach(part)
if dry_run:
LOGGER.info("dry-run: not sending mail")
print("To:", recipients, file=sys.stderr)
print("Subject:", subject, file=sys.stderr)
print("Attachment:", "yes" if attach_excel else "no (size limit)", file=sys.stderr)
print("Body preview:", body[:300] + ("..." if len(body) > 300 else ""), file=sys.stderr)
return
port = int(cfg[ENV_SMTP_PORT], 10)
last_err: Exception | None = None
for attempt in range(SMTP_RETRIES):
try:
with smtplib.SMTP(cfg[ENV_SMTP_HOST], port=port, timeout=30) as smtp:
smtp.starttls()
smtp.login(cfg[ENV_SMTP_USERNAME], cfg[ENV_SMTP_PASSWORD])
smtp.sendmail(sender, recipients, msg.as_string())
LOGGER.info("sent nightly perf email to %d recipient(s)", len(recipients))
return
except Exception as e:
last_err = e
LOGGER.warning("SMTP attempt %d/%d failed: %s", attempt + 1, SMTP_RETRIES, e)
if attempt < SMTP_RETRIES - 1:
import time
time.sleep(SMTP_RETRY_DELAY_SEC)
raise SystemExit(f"Failed to send email after {SMTP_RETRIES} attempts.") from last_err
def _date_from_filename(path: str) -> str:
"""Try to derive a date string from report filename (e.g. nightly_perf_20260211-020704.xlsx -> 2026-02-11)."""
base = os.path.splitext(os.path.basename(path))[0]
if base.startswith("nightly_perf_") and len(base) > 13:
raw = base.replace("nightly_perf_", "")[:8]
if len(raw) == 8 and raw.isdigit():
return f"{raw[:4]}-{raw[4:6]}-{raw[6:8]}"
return base or "unknown"
def _vllm_omni_root() -> str:
"""Resolve vllm-omni repo root: directory that contains a 'tests' subdir (and usually 'tools')."""
path = os.path.dirname(os.path.abspath(__file__))
while path and path != os.path.dirname(path):
if os.path.isdir(os.path.join(path, "tests")):
return path
path = os.path.dirname(path)
return os.path.normpath(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", ".."))
def _default_output_dir() -> str:
"""Default: vllm-omni root / DEFAULT_OUTPUT_DIR (where performance .xlsx files live)."""
root = _vllm_omni_root()
return os.path.join(root, DEFAULT_OUTPUT_DIR)
def parse_args() -> argparse.Namespace:
"""Parse command-line arguments."""
parser = argparse.ArgumentParser(
description="Send nightly performance Excel report by email (config from env).",
)
parser.add_argument(
"--report-file",
type=str,
default=_default_output_dir(),
help="Folder/file path to the nightly_perf_*.xlsx file; default is DEFAULT_OUTPUT_DIR.",
)
parser.add_argument(
"--date",
type=str,
default=None,
help="Date string for subject/body (default: inferred from report filename).",
)
parser.add_argument(
"--dry-run",
action="store_true",
help="Print recipient, subject, and body; do not send.",
)
return parser.parse_args()
def main() -> None:
"""Entrypoint."""
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
)
args = parse_args()
report_file = _get_latest_file(args.report_file) if os.path.isdir(args.report_file) else args.report_file
if not os.path.isfile(report_file):
raise SystemExit(f"Report file not found: {report_file}")
date_str = args.date or _date_from_filename(report_file)
_send_mail(report_file=report_file, date_str=date_str, dry_run=args.dry_run)
if __name__ == "__main__":
main()
+2 -2
View File
@@ -228,8 +228,8 @@ def calculate_metrics(
good_completed += 1
if completed == 0:
warnings.formatwarning = lambda msg, category, filename, lineno, line=None: (
f"{filename}:{lineno}: {category.__name__}: {msg}\n"
warnings.formatwarning = (
lambda msg, category, filename, lineno, line=None: f"{filename}:{lineno}: {category.__name__}: {msg}\n"
)
warnings.warn(
"All requests failed. This is likely due to a misconfiguration on the benchmark arguments.",
-20
View File
@@ -4,28 +4,8 @@ Configuration module for vLLM-Omni.
from vllm_omni.config.lora import LoRAConfig
from vllm_omni.config.model import OmniModelConfig
from vllm_omni.config.stage_config import (
ModelPipeline,
StageConfig,
StageConfigFactory,
StageType,
)
from vllm_omni.config.yaml_util import (
create_config,
load_yaml_config,
merge_configs,
to_dict,
)
__all__ = [
"OmniModelConfig",
"LoRAConfig",
"StageConfig",
"StageConfigFactory",
"ModelPipeline",
"StageType",
"create_config",
"load_yaml_config",
"merge_configs",
"to_dict",
]
+2 -4
View File
@@ -48,7 +48,7 @@ class OmniModelConfig(ModelConfig):
stage_id: int = 0
async_chunk: bool = False
model_stage: str = "thinker"
model_arch: str | None = None
model_arch: str = "Qwen2_5OmniForConditionalGeneration"
worker_type: str | None = None
engine_output_type: str | None = None
hf_config_name: str | None = None
@@ -69,9 +69,7 @@ class OmniModelConfig(ModelConfig):
@property
def architectures(self) -> list[str]:
if self.model_arch is not None:
return [self.model_arch]
return super().architectures
return [self.model_arch]
@property
def embedding_size(self):
-490
View File
@@ -1,490 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Stage Configuration System for vLLM-Omni.
Pipeline structure (stages, types, data-flow) is defined in per-model YAML
files and is set by model developers at integration time.
Runtime parameters (gpu_memory_utilization, tp_size, etc.) come from CLI.
"""
from __future__ import annotations
import re
from dataclasses import asdict, dataclass, field
from enum import Enum
from pathlib import Path
from typing import Any
from vllm.logger import init_logger
from vllm_omni.config.yaml_util import create_config, load_yaml_config, to_dict
# Pipeline YAMLs live alongside model code in model_executor/models/<model>/
_MODELS_DIR = Path(__file__).resolve().parent.parent / "model_executor" / "models"
def get_pipeline_path(model_dir: str, filename: str) -> Path:
"""Return the full path to a pipeline YAML file.
Args:
model_dir: Model subdirectory name (e.g., "qwen3_omni").
filename: Name of the YAML file (e.g., "pipeline.yaml").
Returns:
Absolute path to the file.
"""
return _MODELS_DIR / model_dir / filename
logger = init_logger(__name__)
class StageType(str, Enum):
"""Type of processing stage in the Omni pipeline."""
LLM = "llm"
DIFFUSION = "diffusion"
@dataclass
class StageConfig:
"""Per-stage configuration — pipeline-structure fields only.
Engine params (gpu_memory_utilization, tp_size, etc.) come from CLI,
NOT from this class.
"""
# Identity
stage_id: int
model_stage: str
# Stage type
stage_type: StageType = StageType.LLM
input_sources: list[int] = field(default_factory=list)
custom_process_input_func: str | None = None
final_output: bool = False
final_output_type: str | None = None # "text", "audio", "image"
worker_type: str | None = None # "ar" or "generation"
scheduler_cls: str | None = None
hf_config_name: str | None = None
is_comprehension: bool = False
# Runtime overrides (populated from CLI, not from pipeline YAML)
runtime_overrides: dict[str, Any] = field(default_factory=dict)
def to_omegaconf(self) -> Any:
"""Convert to OmegaConf for backward compatibility with OmniStage.
Returns:
OmegaConf DictConfig with stage configuration in legacy format.
"""
# Build engine_args dict with required fields
engine_args: dict[str, Any] = {
"model_stage": self.model_stage,
}
if self.worker_type:
engine_args["worker_type"] = self.worker_type
if self.scheduler_cls:
engine_args["scheduler_cls"] = self.scheduler_cls
if self.hf_config_name:
engine_args["hf_config_name"] = self.hf_config_name
# Apply runtime overrides (CLI args)
for key, value in self.runtime_overrides.items():
if key not in ("devices", "max_batch_size"):
engine_args[key] = value
# Build runtime config
runtime: dict[str, Any] = {
"process": True,
"max_batch_size": self.runtime_overrides.get("max_batch_size", 1),
}
if "devices" in self.runtime_overrides:
runtime["devices"] = self.runtime_overrides["devices"]
# Build full config dict
config_dict: dict[str, Any] = {
"stage_id": self.stage_id,
"stage_type": StageType(self.stage_type).value,
"engine_args": create_config(engine_args),
"runtime": create_config(runtime),
"engine_input_source": self.input_sources, # Legacy field name
"final_output": self.final_output,
"final_output_type": self.final_output_type,
"is_comprehension": self.is_comprehension,
}
if self.custom_process_input_func:
config_dict["custom_process_input_func"] = self.custom_process_input_func
return create_config(config_dict)
@dataclass
class ModelPipeline:
"""Complete pipeline definition for a multi-stage model.
Defined by model developers, bundled with the model, not user-editable.
"""
model_type: str
stages: list[StageConfig]
# Pipeline-wide behavior flags
async_chunk: bool = False
# Optional distributed configuration
connectors: dict[str, Any] | None = None
edges: list[dict[str, Any]] | None = None
def get_stage(self, stage_id: int) -> StageConfig | None:
"""Look up a stage by its ID.
Args:
stage_id: The stage ID to search for.
Returns:
The matching StageConfig, or None if not found.
"""
for stage in self.stages:
if stage.stage_id == stage_id:
return stage
return None
def validate_pipeline(self) -> list[str]:
"""Validate pipeline topology at model integration time (not runtime).
Checks:
- All stage IDs are unique
- All input_sources reference valid stage IDs
- At least one entry point (stage with empty input_sources)
Returns:
List of validation error messages. Empty list if valid.
"""
errors: list[str] = []
if not self.stages:
errors.append("Topology has no stages defined")
return errors
# Check for unique stage IDs
stage_ids = [s.stage_id for s in self.stages]
if len(stage_ids) != len(set(stage_ids)):
errors.append("Duplicate stage IDs found")
stage_id_set = set(stage_ids)
# Check input_sources reference valid stages
for stage in self.stages:
for source_id in stage.input_sources:
if source_id not in stage_id_set:
errors.append(f"Stage {stage.stage_id} references non-existent input source {source_id}")
if source_id == stage.stage_id:
errors.append(f"Stage {stage.stage_id} references itself as input source")
# Check for at least one entry point
entry_points = [s for s in self.stages if not s.input_sources]
if not entry_points:
errors.append("No entry point found (stage with empty input_sources)")
return errors
class StageConfigFactory:
"""Factory that loads pipeline YAML and merges CLI overrides.
Handles both single-stage and multi-stage models.
"""
# Mapping of model types to directories under model_executor/models/.
PIPELINE_MODELS: dict[str, str] = {
"qwen3_omni_moe": "qwen3_omni",
"qwen2_5_omni": "qwen2_5_omni",
"bagel": "bagel",
"qwen3_tts": "qwen3_tts",
"mimo_audio": "mimo_audio",
"glm-image": "glm_image",
"cosyvoice3": "cosyvoice3",
"mammothmoda2": "mammoth_moda2",
}
# Fallback: map HF architecture class names to pipeline dirs.
# Used when model_type collides with another model (e.g. MiMo Audio
# reports model_type="qwen2" which matches plain Qwen2, not our pipeline).
_ARCHITECTURE_MODELS: dict[str, str] = {
"MiMoAudioForConditionalGeneration": "mimo_audio",
"HunyuanImage3ForCausalMM": "hunyuan_image3",
}
@classmethod
def create_from_model(
cls,
model: str,
cli_overrides: dict[str, Any] | None = None,
) -> list[StageConfig] | None:
"""Load pipeline YAML, merge with CLI overrides.
Args:
model: Model name or path.
cli_overrides: CLI overrides from VllmConfig/OmniDiffusionConfig.
Returns:
List of StageConfig objects with CLI overrides applied,
or None if no pipeline definition was found for this model.
"""
if cli_overrides is None:
cli_overrides = {}
trust_remote_code = cli_overrides.get("trust_remote_code", True)
pipeline = cls._load_pipeline(model, trust_remote_code=trust_remote_code)
if pipeline is None:
return None
errors = pipeline.validate_pipeline()
if errors:
logger.warning(f"Pipeline validation warnings for {model}: {errors}")
# Apply CLI overrides
result: list[StageConfig] = []
for stage in pipeline.stages:
# Merge global CLI overrides
stage.runtime_overrides = cls._merge_cli_overrides(stage, cli_overrides)
result.append(stage)
return result
@classmethod
def create_default_diffusion(cls, kwargs: dict[str, Any]) -> list[dict[str, Any]]:
"""Single-stage diffusion - no YAML needed.
Creates a default diffusion stage configuration for single-stage
diffusion models. Returns a legacy OmegaConf-compatible dict for
backward compatibility with OmniStage.
Args:
kwargs: Engine arguments from CLI/API.
Returns:
List containing a single config dict for the diffusion stage.
"""
# Calculate devices based on parallel config
devices = "0"
if "parallel_config" in kwargs:
num_devices = kwargs["parallel_config"].world_size
for i in range(1, num_devices):
devices += f",{i}"
engine_args: dict[str, Any] = {}
for key, value in kwargs.items():
if key in ("parallel_config",):
continue
engine_args[key] = value
# Serialize parallel_config as dict for OmegaConf compatibility
if "parallel_config" in kwargs:
engine_args["parallel_config"] = asdict(kwargs["parallel_config"])
engine_args.setdefault("cache_backend", "none")
engine_args["model_stage"] = "diffusion"
# Convert dtype to string for OmegaConf
if "dtype" in engine_args:
engine_args["dtype"] = str(engine_args["dtype"])
config_dict: dict[str, Any] = {
"stage_id": 0,
"stage_type": StageType.DIFFUSION.value,
"runtime": {
"process": True,
"devices": devices,
"max_batch_size": 1,
},
"engine_args": create_config(engine_args),
"final_output": True,
"final_output_type": "image",
}
return [config_dict]
@classmethod
def _load_pipeline(cls, model: str, trust_remote_code: bool = True) -> ModelPipeline | None:
"""Load pipeline YAML for the model.
Args:
model: Model name or path.
trust_remote_code: Whether to trust remote code for HF config loading.
Returns:
ModelPipeline if found, None otherwise.
"""
model_type, hf_config = cls._auto_detect_model_type(model, trust_remote_code=trust_remote_code)
if model_type is None:
return None
pipeline_dir = cls.PIPELINE_MODELS.get(model_type)
# Fallback: check HF architectures when model_type doesn't match
if pipeline_dir is None and hf_config is not None:
for arch in getattr(hf_config, "architectures", []) or []:
pipeline_dir = cls._ARCHITECTURE_MODELS.get(arch)
if pipeline_dir is not None:
model_type = pipeline_dir
break
if pipeline_dir is None:
logger.debug(f"No pipeline mapping for model_type: {model_type}")
return None
pipeline_path = get_pipeline_path(pipeline_dir, "pipeline.yaml")
if not pipeline_path.exists():
logger.debug(f"Pipeline file not found: {pipeline_path}")
return None
return cls._parse_pipeline_yaml(pipeline_path, model_type)
@classmethod
def _parse_pipeline_yaml(cls, path: Path, model_type: str) -> ModelPipeline:
"""Parse a pipeline YAML file.
Args:
path: Path to the YAML file.
model_type: Model type identifier.
Returns:
ModelPipeline object.
"""
config_data = load_yaml_config(path)
stages: list[StageConfig] = []
for stage_data in config_data.stages:
# Use .get() for optional fields — idiomatic for OmegaConf DictConfig
stage_type_str = stage_data.get("stage_type", "llm")
stage_type = StageType(stage_type_str) if stage_type_str else StageType.LLM
# Handle both 'input_sources' (new) and 'engine_input_source' (legacy)
input_sources = stage_data.get("input_sources", None)
if input_sources is None:
input_sources = stage_data.get("engine_input_source", [])
if input_sources is None:
input_sources = []
input_sources = list(input_sources)
stage = StageConfig(
stage_id=stage_data.stage_id,
model_stage=stage_data.model_stage,
stage_type=stage_type,
input_sources=input_sources,
custom_process_input_func=stage_data.get("custom_process_input_func", None),
final_output=stage_data.get("final_output", False),
final_output_type=stage_data.get("final_output_type", None),
worker_type=stage_data.get("worker_type", None),
scheduler_cls=stage_data.get("scheduler_cls", None),
hf_config_name=stage_data.get("hf_config_name", None),
is_comprehension=stage_data.get("is_comprehension", False),
)
stages.append(stage)
# Get pipeline-wide flags
async_chunk = config_data.get("async_chunk", False)
# Get optional connector config
connectors = to_dict(config_data.connectors) if hasattr(config_data, "connectors") else None
edges = to_dict(config_data.edges) if hasattr(config_data, "edges") else None
return ModelPipeline(
model_type=getattr(config_data, "model_type", model_type),
stages=stages,
async_chunk=async_chunk,
connectors=connectors,
edges=edges,
)
@classmethod
def _auto_detect_model_type(cls, model: str, trust_remote_code: bool = True) -> tuple[str | None, Any]:
"""Auto-detect model_type from model directory.
Args:
model: Model name or path.
trust_remote_code: Whether to trust remote code for HF config loading.
Returns:
Tuple of (model_type, hf_config). Both may be None on failure.
"""
try:
from vllm.transformers_utils.config import get_config
hf_config = get_config(model, trust_remote_code=trust_remote_code)
return hf_config.model_type, hf_config
except Exception as e:
logger.debug(f"Failed to auto-detect model type for {model}: {e}")
return None, None
# Keys that should never be forwarded as engine overrides (internal /
# orchestrator-only knobs, complex objects, etc.).
_INTERNAL_KEYS: set[str] = {
"model",
"stage_configs_path",
"stage_id",
"stage_init_timeout",
"init_timeout",
"shm_threshold_bytes",
"worker_backend",
"ray_address",
"batch_timeout",
"log_stats",
"tokenizer",
"parallel_config",
}
@classmethod
def _merge_cli_overrides(
cls,
stage: StageConfig,
cli_overrides: dict[str, Any],
) -> dict[str, Any]:
"""Merge CLI overrides into stage runtime config.
All CLI arguments registered by engine config classes (e.g.
EngineArgs / OmniDiffusionConfig) are accepted as overrides
unless they appear in ``_INTERNAL_KEYS``.
Handles:
- Global overrides (apply to all stages)
- Per-stage overrides (--stage-N-* format, take precedence)
Args:
stage: The stage to merge overrides into.
cli_overrides: CLI arguments from VllmConfig/OmniDiffusionConfig.
Returns:
Dict of runtime overrides for this stage.
"""
result: dict[str, Any] = {}
# Apply global overrides – any key not in the internal blocklist
# is forwarded so that engine-registered params work out of the box.
for key, value in cli_overrides.items():
if key in cls._INTERNAL_KEYS:
continue
if re.match(r"stage_\d+_", key):
# Per-stage keys handled below
continue
if value is not None:
result[key] = value
# Apply per-stage overrides (--stage-N-* format, take precedence)
stage_prefix = f"stage_{stage.stage_id}_"
for key, value in cli_overrides.items():
if key.startswith(stage_prefix) and value is not None:
param_name = key[len(stage_prefix) :]
if param_name in cls._INTERNAL_KEYS:
continue
result[param_name] = value
return result
-65
View File
@@ -1,65 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Centralized OmegaConf wrapper for vLLM-Omni.
All OmegaConf usage in the project MUST go through this module.
Other modules should import these wrapper functions instead of
using OmegaConf directly.
"""
from __future__ import annotations
from typing import Any
from omegaconf import DictConfig, OmegaConf
def load_yaml_config(path: str | Any) -> DictConfig:
"""Load a YAML file and return it as a DictConfig.
Args:
path: Path to the YAML file.
Returns:
OmegaConf DictConfig with attribute-style access.
"""
return OmegaConf.load(path) # type: ignore[return-value]
def create_config(data: Any) -> DictConfig:
"""Wrap a dict (or list) into a DictConfig.
Args:
data: Dict, list, or other structure to wrap.
Returns:
OmegaConf DictConfig / ListConfig.
"""
return OmegaConf.create(data)
def merge_configs(*cfgs: Any) -> dict:
"""Deep-merge multiple configs and return a plain dict.
Args:
*cfgs: DictConfig or dict objects to merge (left to right).
Returns:
Plain dict with merged, resolved values.
"""
merged = OmegaConf.merge(*cfgs)
return OmegaConf.to_container(merged, resolve=True) # type: ignore[return-value]
def to_dict(obj: Any, *, resolve: bool = True) -> Any:
"""Convert a DictConfig (or similar) to a plain dict.
Args:
obj: OmegaConf container to convert.
resolve: Whether to resolve interpolations (default True).
Returns:
Plain dict.
"""
return OmegaConf.to_container(obj, resolve=resolve) # type: ignore[return-value]
-75
View File
@@ -984,80 +984,6 @@ def enable_cache_for_bagel(pipeline: Any, cache_config: Any) -> Callable[[int],
return refresh_cache_context
def enable_cache_for_flux2(pipeline: Any, cache_config: Any) -> Callable[[int], None]:
"""Enable cache-dit for Flux.2-dev pipeline.
Args:
pipeline: The Flux2 pipeline instance.
cache_config: DiffusionCacheConfig instance with cache configuration.
Returns:
A refresh function that can be called with a new ``num_inference_steps``
to update the cache context for the pipeline.
"""
# Build DBCacheConfig for transformer
db_cache_config = _build_db_cache_config(cache_config)
calibrator = None
if cache_config.enable_taylorseer:
taylorseer_order = cache_config.taylorseer_order
calibrator = TaylorSeerCalibratorConfig(taylorseer_order=taylorseer_order)
logger.info(f"TaylorSeer enabled with order={taylorseer_order}")
# Build ParamsModifier for transformer
modifier = ParamsModifier(
cache_config=db_cache_config,
calibrator_config=calibrator,
)
logger.info(
f"Enabling cache-dit on Flux transformer with BlockAdapter: "
f"Fn={db_cache_config.Fn_compute_blocks}, "
f"Bn={db_cache_config.Bn_compute_blocks}, "
f"W={db_cache_config.max_warmup_steps}, "
)
# Enable cache-dit using BlockAdapter for transformer
cache_dit.enable_cache(
(
BlockAdapter(
transformer=pipeline.transformer,
blocks=[
pipeline.transformer.transformer_blocks,
pipeline.transformer.single_transformer_blocks,
],
forward_pattern=[ForwardPattern.Pattern_1, ForwardPattern.Pattern_2],
params_modifiers=[modifier],
)
),
cache_config=db_cache_config,
)
def refresh_cache_context(pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None:
"""Refresh cache context for the transformer with new num_inference_steps.
Args:
pipeline: The Flux2 pipeline instance.
num_inference_steps: New number of inference steps.
"""
if cache_config.scm_steps_mask_policy is None:
cache_dit.refresh_context(pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose)
else:
cache_dit.refresh_context(
pipeline.transformer,
cache_config=DBCacheConfig().reset(
num_inference_steps=num_inference_steps,
steps_computation_mask=cache_dit.steps_mask(
mask_policy=cache_config.scm_steps_mask_policy,
total_steps=num_inference_steps,
),
steps_computation_policy=cache_config.scm_steps_policy,
),
verbose=verbose,
)
return refresh_cache_context
# Register custom cache-dit enablers after function definitions
CUSTOM_DIT_ENABLERS.update(
{
@@ -1072,7 +998,6 @@ CUSTOM_DIT_ENABLERS.update(
"LTX2Pipeline": enable_cache_for_ltx2,
"LTX2ImageToVideoPipeline": enable_cache_for_ltx2,
"BagelPipeline": enable_cache_for_bagel,
"Flux2Pipeline": enable_cache_for_flux2,
}
)
+7 -9
View File
@@ -322,7 +322,6 @@ class OmniDiffusionConfig:
dtype: torch.dtype = torch.bfloat16
model_config: dict[str, Any] = field(default_factory=dict)
tf_model_config: TransformerConfig = field(default_factory=TransformerConfig)
# Attention
@@ -464,13 +463,11 @@ class OmniDiffusionConfig:
@property
def is_moe(self) -> bool:
num_experts = self.tf_model_config.get("num_experts", None)
if not isinstance(num_experts, (list, tuple, int)):
return False
if isinstance(num_experts, int):
return num_experts > 0
return num_experts > 1
if isinstance(num_experts, (list, tuple)):
return any(isinstance(n, int) and n > 0 for n in num_experts)
return any(isinstance(n, int) and n > 1 for n in num_experts)
return False
@@ -514,11 +511,12 @@ class OmniDiffusionConfig:
initial_master_port = (self.master_port or 30005) + random.randint(0, 100)
self.master_port = self.settle_port(initial_master_port, 37)
# Convert parallel_config dict/DictConfig to DiffusionParallelConfig
# Use Mapping to handle both plain dicts and OmegaConf DictConfig
if isinstance(self.parallel_config, Mapping):
self.parallel_config = DiffusionParallelConfig.from_dict(dict(self.parallel_config))
# Convert parallel_config dict to DiffusionParallelConfig if needed
# This must be done before accessing parallel_config.world_size
if isinstance(self.parallel_config, dict):
self.parallel_config = DiffusionParallelConfig.from_dict(self.parallel_config)
elif not isinstance(self.parallel_config, DiffusionParallelConfig):
# If it's neither dict nor DiffusionParallelConfig, use default config
self.parallel_config = DiffusionParallelConfig()
if self.num_gpus is None:
+1 -21
View File
@@ -6,7 +6,6 @@ import time
from collections.abc import Iterable
from typing import Any
import numpy as np
import PIL.Image
from vllm.logger import init_logger
@@ -31,13 +30,6 @@ def supports_image_input(model_class_name: str) -> bool:
return bool(getattr(model_cls, "support_image_input", False))
def supports_audio_input(model_class_name: str) -> bool:
model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name)
if model_cls is None:
return False
return bool(getattr(model_cls, "support_audio_input", False))
def image_color_format(model_class_name: str) -> str:
model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name)
return getattr(model_cls, "color_format", "RGB")
@@ -377,19 +369,7 @@ class DiffusionEngine:
dummy_image = PIL.Image.new(color_format, (width, height))
else:
dummy_image = None
if supports_audio_input(self.od_config.model_class_name):
audio_sr = 16000
audio_duration_sec = 4
audio_array = np.random.randn(audio_sr * audio_duration_sec).astype(np.float32)
dummy_audio = audio_array[audio_sr * 1 : audio_sr * 3]
else:
dummy_audio = None
prompt: OmniTextPrompt = {
"prompt": "dummy run",
"multi_modal_data": {"image": dummy_image, "audio": dummy_audio},
}
prompt: OmniTextPrompt = {"prompt": "dummy run", "multi_modal_data": {"image": dummy_image}}
req = OmniDiffusionRequest(
prompts=[prompt],
sampling_params=OmniDiffusionSamplingParams(
+1 -1
View File
@@ -20,7 +20,7 @@ environment_variables: dict[str, Callable[[], Any]] = {
# used in distributed environment to determine the master address
"MASTER_ADDR": lambda: os.getenv("MASTER_ADDR", ""),
# used in distributed environment to manually set the communication port
"MASTER_PORT": lambda: int(os.getenv("MASTER_PORT", "0")) if "MASTER_PORT" in os.environ else None,
"MASTER_PORT": lambda: (int(os.getenv("MASTER_PORT", "0")) if "MASTER_PORT" in os.environ else None),
# path to cudatoolkit home directory, under which should be bin, include,
# and lib directories.
"CUDA_HOME": lambda: os.environ.get("CUDA_HOME", None),
+1 -4
View File
@@ -8,7 +8,6 @@ and modifying model forward passes without invasive changes to model code.
from __future__ import annotations
import inspect
from collections.abc import Callable
from dataclasses import dataclass
from typing import Any
@@ -167,9 +166,7 @@ class HookRegistry:
# `_original_forward` attribute on the same module.
if not hasattr(module, "_omni_original_forward"):
module._omni_original_forward = module.forward # type: ignore[attr-defined]
wrapped = _WrappedForward(module)
wrapped.__signature__ = inspect.signature(module.forward) # type: ignore[attr-defined]
module.forward = wrapped # type: ignore[assignment]
module.forward = _WrappedForward(module) # type: ignore[assignment]
return registry
-26
View File
@@ -157,29 +157,3 @@ class RotaryEmbedding(CustomOp):
sin,
interleaved=self.interleaved,
)
def apply_rope_to_qk(
rope: RotaryEmbedding,
query: torch.Tensor,
key: torch.Tensor,
image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Apply rotary positional embeddings to query and key tensors.
Args:
rope: RotaryEmbedding instance for applying position embeddings
query: Query tensor [B, S, H, D]
key: Key tensor [B, S, H, D]
image_rotary_emb: Tuple of (cos, sin) tensors or None
Returns:
Tuple of (query, key) with RoPE applied if rotary embeddings provided
"""
if image_rotary_emb is not None:
cos, sin = image_rotary_emb
cos = cos.to(query.dtype)
sin = sin.to(query.dtype)
query = rope(query, cos, sin)
key = rope(key, cos, sin)
return query, key
Loaded 100 of 166 files, more files were not shown because too many files have changed in this diff. Show more