anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

Merge branch 'main' into feat/moss-tts-nano

Browse Source 
Signed-off-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
This commit is contained in:
Yueqian Lin authored and GitHub committed 2026-04-24 11:00:13 -04:00
2 parents 0a43fb94ac 7b876324c2
commit 74522ca3ec
307 files changed
+20557 -5581
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

.buildkite/scripts/hardware_ci/run-amd-test.sh
-21
View File
@@ -11,15 +11,6 @@ set -o pipefail
export PYTHONPATH=".."
# Print ROCm version
echo "--- Confirming Clean Initial State"
while true; do
sleep 3
if grep -q clean /opt/amdgpu/etc/gpu_state; then
echo "GPUs state is \"clean\""
break
fi
done
echo "--- ROCm info"
rocminfo
@@ -51,18 +42,6 @@ cleanup_docker() {
# Call the cleanup docker function
cleanup_docker
echo "--- Resetting GPUs"
echo "reset" > /opt/amdgpu/etc/gpu_state
while true; do
sleep 3
if grep -q clean /opt/amdgpu/etc/gpu_state; then
echo "GPUs state is \"clean\""
break
fi
done
echo "--- Pulling container"
## Temporary change to use AMD Docker Hub to store the vllm-omni image
# to bypass the rate limit issue with ECR Public Gallery.
.buildkite/test-merge.yml
+40
View File
@@ -297,6 +297,46 @@ steps:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: "Audio Streaming Input Test with H100"
timeout_in_minutes: 30
depends_on: upload-merge-pipeline
commands:
- export VLLM_TEST_CLEAN_GPU_MEMORY="1"
- pytest -s -v tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py -m "advanced_model" --run-level "advanced_model"
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 2
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
- name: HF_TOKEN
valueFrom:
secretKeyRef:
name: hf-token-secret
key: token
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: "Diffusion Image Edit Test with H100 (1 GPU)"
timeout_in_minutes: 20
depends_on: upload-merge-pipeline
.buildkite/test-nightly.yml
+107 -12
View File
@@ -13,7 +13,7 @@ steps:
- label: ":full_moon: Omni · Function Test"
timeout_in_minutes: 90
commands:
- pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and H100 and omni" --run-level "advanced_model"
- pytest -s -v tests/e2e/ -m "full_model and H100 and omni" --run-level "full_model" --ignore=tests/e2e/accuracy
agents:
queue: "mithril-h100-pool"
plugins:
@@ -52,7 +52,7 @@ steps:
timeout_in_minutes: 90
commands:
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
- pytest -s -v tests/examples/ -m "advanced_model and omni and L4" --run-level "advanced_model"
- pytest -s -v tests/examples/ -m "full_model and omni and L4" --run-level "full_model"
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
plugins:
@@ -70,7 +70,52 @@ steps:
- label: ":full_moon: Omni · Doc Test with H100"
timeout_in_minutes: 90
commands:
- pytest -s -v tests/examples/ -m "advanced_model and omni and H100" --run-level "advanced_model"
- pytest -s -v tests/examples/ -m "full_model and omni and H100" --run-level "full_model"
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 2
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
- name: HF_TOKEN
valueFrom:
secretKeyRef:
name: hf-token-secret
key: token
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: ":full_moon: Omni · Accuracy Test"
timeout_in_minutes: 180
commands:
- export SEED_TTS_WER_EVAL=1
- export SEED_TTS_EVAL_DEVICE=cuda:1
- |
set +e
pytest -s -v tests/e2e/accuracy/qwen3_omni/test_qwen3_omni.py -m "full_model" --run-level full_model
EXIT=$$?
buildkite-agent artifact upload "tests/e2e/accuracy/qwen3_omni/results/qwen_omni_acc/*.json"
exit $$EXIT
agents:
queue: "mithril-h100-pool"
plugins:
@@ -160,7 +205,7 @@ steps:
timeout_in_minutes: 90
commands:
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
- pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and L4 and omni" --run-level "advanced_model"
- pytest -s -v tests/e2e/ -m "full_model and L4 and omni" --run-level "full_model" --ignore=tests/e2e/accuracy
agents:
queue: "gpu_1_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
plugins:
@@ -233,7 +278,7 @@ steps:
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100"
timeout_in_minutes: 120
commands:
- pytest -s -v tests/e2e/online_serving/test_*_expansion.py -k "not test_wan22_expansion and not test_wan_2_1_vace_expansion and not hunyuan" -m "advanced_model and diffusion and H100" --run-level "advanced_model"
- pytest -sv tests/e2e/ -k "not test_wan and not hunyuan" -m "full_model and diffusion and H100" --run-level "full_model" --ignore=tests/e2e/accuracy
agents:
queue: "mithril-h100-pool"
plugins:
@@ -271,7 +316,7 @@ steps:
- label: ":full_moon: Diffusion X2I(&A&T) · Function Test with L4"
timeout_in_minutes: 60
commands:
- pytest -s -v tests/e2e/online_serving/test_*_expansion.py -k "not test_wan22_expansion and not test_wan_2_1_vace_expansion and not hunyuan" -m "advanced_model and diffusion and L4" --run-level "advanced_model"
- pytest -sv tests/e2e/ -k "not test_wan and not hunyuan" -m "full_model and diffusion and L4" --run-level "full_model" --ignore=tests/e2e/accuracy
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
plugins:
@@ -290,7 +335,7 @@ steps:
timeout_in_minutes: 60
commands:
- export VLLM_TEST_CLEAN_GPU_MEMORY="1"
- pytest -s -v tests/examples/*/test_text_to_image.py -m "advanced_model and example and H100" --run-level "advanced_model"
- pytest -s -v tests/examples/*/test_text_to_image.py -m "full_model and example and H100" --run-level "full_model"
agents:
queue: "mithril-h100-pool"
plugins:
@@ -328,7 +373,7 @@ steps:
- label: ":full_moon: Diffusion X2I(&A&T) · GEBench Accuracy Test"
timeout_in_minutes: 60
commands:
- pytest -s -v tests/e2e/accuracy/test_gebench_h100_smoke.py --run-level advanced_model --gebench-model Qwen/Qwen-Image-2512 --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gebench-port 8093 --accuracy-workers 1
- pytest -s -v tests/e2e/accuracy/test_gebench_h100_smoke.py --run-level full_model --gebench-model Qwen/Qwen-Image-2512 --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gebench-port 8093 --accuracy-workers 1
- buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gebench_qwen-image-2512/summary*.json"
agents:
queue: "mithril-h100-pool"
@@ -367,7 +412,7 @@ steps:
- label: ":full_moon: Diffusion X2I(&A&T) · GEdit-Bench Accuracy Test"
timeout_in_minutes: 60
commands:
- pytest -s -v tests/e2e/accuracy/test_gedit_bench_h100_smoke.py --run-level advanced_model --gedit-model Qwen/Qwen-Image-Edit --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gedit-port 8093 --gedit-samples-per-group 20 --accuracy-workers 1
- pytest -s -v tests/e2e/accuracy/test_gedit_bench_h100_smoke.py --run-level full_model --gedit-model Qwen/Qwen-Image-Edit --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gedit-port 8093 --gedit-samples-per-group 20 --accuracy-workers 1
- buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gedit_scores_qwen-image-edit/qwen-image-edit_all_all_vie_score_*.csv"
- buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gedit_scores_qwen-image-edit/qwen-image-edit_all_all_summary_*.json"
agents:
@@ -409,7 +454,7 @@ steps:
- label: ":full_moon: Diffusion X2I(&A&T) · Accuracy Test"
timeout_in_minutes: 180
commands:
- pytest -s -v tests/e2e/accuracy/test_qwen_image*.py --run-level advanced_model
- pytest -s -v tests/e2e/accuracy/test_qwen_image*.py --run-level full_model
agents:
queue: "mithril-h100-pool"
plugins:
@@ -514,7 +559,7 @@ steps:
- label: ":full_moon: Diffusion X2V · Function Test"
timeout_in_minutes: 90
commands:
- pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_hunyuan_video_15_expansion.py -m "advanced_model" --run-level "advanced_model"
- pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_hunyuan_video_15_expansion.py --run-level "full_model"
agents:
queue: "mithril-h100-pool"
plugins:
@@ -552,7 +597,55 @@ steps:
- label: ":full_moon: Diffusion X2V · Accuracy Test"
timeout_in_minutes: 180
commands:
- pytest -s -v tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py -m advanced_model --run-level advanced_model
- pytest -s -v tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py -m full_model --run-level full_model
- pytest -s -v tests/e2e/accuracy/test_ltx2_3_video_similarity.py -m advanced_model --run-level advanced_model
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 2
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
- name: HF_TOKEN
valueFrom:
secretKeyRef:
name: hf-token-secret
key: token
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: ":full_moon: Diffusion X2V · Perf Test"
key: nightly-diffusion-x2v-performance
timeout_in_minutes: 180
commands:
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
- |
set +e
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json
EXIT1=$$?
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
exit $$EXIT1
agents:
queue: "mithril-h100-pool"
plugins:
@@ -636,6 +729,7 @@ steps:
- nightly-omni-performance
- nightly-tts-performance
- nightly-diffusion-x2iat-performance
- nightly-diffusion-x2v-performance
- nightly-testcase-statistics
if: build.env("NIGHTLY") == "1"
commands:
@@ -645,6 +739,7 @@ steps:
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-tts-performance
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2v-performance
- buildkite-agent artifact download "tests/dfx/perf/results/*.html" . --step nightly-testcase-statistics
- python tools/nightly/generate_nightly_perf_excel.py
- python tools/nightly/generate_nightly_perf_html.py
.buildkite/test-ready.yml
-17
View File
@@ -120,23 +120,6 @@ steps:
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Audio Generation Model Test"
depends_on: upload-ready-pipeline
commands:
- timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_expansion.py -m "advanced_model and diffusion and L4" --run-level advanced_model
agents:
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
environment:
- "HF_HOME=/fsx/hf_cache"
- "HF_TOKEN"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Diffusion Cache Backend Test"
depends_on: upload-ready-pipeline
commands:
.claude/skills/add-diffusion-model/SKILL.md
+36 -4
View File
@@ -282,10 +282,11 @@ For Omni or custom models, create:
Required updates:
1. `docs/user_guide/diffusion/parallelism_acceleration.md` — parallelism support table
2. `docs/user_guide/diffusion/teacache.md` — if TeaCache supported
3. `docs/user_guide/diffusion/cache_dit_acceleration.md` — if Cache-DiT supported
4. `examples/offline_inference/xxx/README.md` — offline example docs
5. `examples/online_serve/xxx/README.md` — online serve docs
2. `docs/user_guide/diffusion/cpu_offload_diffusion.md` — if CPU offload supported (add to supported models table)
3. `docs/user_guide/diffusion/teacache.md` — if TeaCache supported
4. `docs/user_guide/diffusion/cache_dit_acceleration.md` — if Cache-DiT supported
5. `examples/offline_inference/xxx/README.md` — offline example docs
6. `examples/online_serve/xxx/README.md` — online serve docs
### Step 8: Add E2E Tests (Recommended)
@@ -512,6 +513,37 @@ After adding parallelism support, update:
1. `docs/user_guide/diffusion/parallelism_acceleration.md` — add your model to the support table
2. Record which parallelism methods are supported (USP, Ring, CFG, TP, HSDP, VAE-Patch)
### Step 11: Add CPU Offload Support
Implement `SupportsModuleOffload` on your pipeline class to enable
`--enable-cpu-offload` and `--enable-layerwise-offload`. The protocol
declares which submodules the offloader should manage:
```python
from typing import ClassVar
from vllm_omni.diffusion.models.interface import SupportsModuleOffload
class YourPipeline(nn.Module, SupportsModuleOffload):
_dit_modules: ClassVar[list[str]] = ["transformer"]
_encoder_modules: ClassVar[list[str]] = ["text_encoder"]
_vae_modules: ClassVar[list[str]] = ["vae"]
_resident_modules: ClassVar[list[str]] = [] # optional
```
- `_dit_modules`: denoising submodules (kept on GPU during diffusion loop)
- `_encoder_modules`: encoder/vision submodules (offloaded to CPU during diffusion loop)
- `_vae_modules`: VAE(s) (handled by both sequential and layerwise backends)
- `_resident_modules`: additional modules to pin on GPU during layerwise
offloading (e.g. embedders, connectors). Only used by the layerwise
backend. Optional — defaults to `[]`.
All attribute names support dotted paths for nested submodules
(e.g. `"pipe.transformer"`, `"bagel.time_embedder"`).
Pipelines without `SupportsModuleOffload` fall back to scanning
well-known attribute names (`transformer`, `text_encoder`, `vae`,
etc.), which fails for non-standard names.
---
## Iterative Development Tips
.claude/skills/add-tts-model/SKILL.md
+261 -41
View File
@@ -8,10 +8,90 @@ description: "Integrate a new text-to-speech model into vLLM-Omni from HuggingFa
## Overview
```
HF Reference -> Stage Separation -> Online Serving -> Async Chunk -> CUDA Graph
(Phase 1) (Phase 2) (Phase 3) (Phase 4) (Phase 5)
HF Reference -> Stage Separation -> Online Serving -> Async Chunk -> CUDA Graph -> Pre-commit/DCO
(Phase 1) (Phase 2) (Phase 3) (Phase 4) (Phase 5) (Phase 6)
```
Three architecture patterns are supported:
- **Two-stage pipeline** (e.g. Qwen3-TTS, Fish Speech, CosyVoice3): AR
code-predictor → audio decoder, connected via async_chunk for low-latency
streaming. Use this for maximum performance.
- **Single-stage AR via generator** (e.g. MOSS-TTS-Nano): entire model runs
inside one AR worker, streaming audio chunks through a per-request
`inference_stream()` generator. Use this when the upstream model bundles AR
+ codec inseparably. See [references/single-stage-ar.md](references/single-stage-ar.md).
- **Single-stage, vLLM-native base LM + side computation** (e.g. VoxCPM2):
the base language model runs under vLLM's PagedAttention as a normal AR
model; diffusion / VAE / side computations run outside vLLM and are
attached via the runner post-processing hook. This is a distinct pattern
from the generator approach above — do not confuse the two.
The single-stage variants skip Phase 4 (async_chunk) but Phase 5 (CUDA graph)
is still encouraged for the inner AR loop.
## Cross-Cutting Invariants
These rules apply to every TTS model regardless of architecture (AR vs AR+diffusion, single-stage vs two-stage, codec-based vs VAE-based). They surface repeatedly across PRs — check them at the end of every phase.
### I1. Streaming output contract
Pick exactly one per-step semantics for `forward()` and document it in the docstring:
- **Delta**: yield only new audio samples produced this step. Preferred — linear cost, low memory.
- **Cumulative**: re-decode from step 0 every call. O(N²); only acceptable if the codec has no streaming decode path.
If you choose **delta**, verify the full emit→consolidate→consume chain:
1. `forward()` returns `{"model_outputs": <new_chunk_only>, ...}`
2. `_consolidate_multimodal_tensors()` in `vllm_omni/engine/output_processor.py` concatenates the audio key into one tensor at finish. If it skips the key (`continue`), offline consumers receive only the final chunk. See `output_processor.py` for the concrete list of handled modality keys.
3. Streaming consumers (SSE, Gradio) receive per-step deltas; offline consumers (`engine.generate()`) receive a single concatenated tensor.
Cumulative-vs-delta mismatch is the most common silent bug — offline RTF benchmarks pass, but users hear replays or truncation.
### I2. Multimodal output consumer hygiene
`outputs[0].outputs[0].multimodal_output[<key>]` can be any of `Tensor`, `list[Tensor]` (pre-consolidation snapshot), `np.ndarray`, or scalar. When writing tests, examples, and benchmarks:
- **Never** use `dict.get("a") or dict.get("b")` on tensor values — Python evaluates the tensor's boolean, raising `RuntimeError: Boolean value of Tensor with more than one value is ambiguous`. Use explicit `if x is None` chains.
- Always defensively handle the list form: `if isinstance(x, list): x = torch.cat([t.reshape(-1) for t in x], dim=0)`.
- Assert `shape` / `dtype` / `duration` explicitly; do not rely on truthiness for presence checks.
### I3. Hot-loop GPU discipline
Inside any per-step model loop (AR decode, diffusion solver, CFM Euler, vocoder block loop):
- No `tensor.item()`, `.cpu()`, or `.tolist()` — each triggers a GPU→CPU sync; at 10 steps × 60 frames × 4 ops that is 2400 syncs per request.
- Prefer `dst.copy_(src)` over `dst.fill_(src.item())` when writing a scalar tensor into a buffer.
- Prefer `torch.compile(Model.forward, fullgraph=False)` on the whole forward over per-submodule compile — fewer dispatch boundaries, larger fusion regions. Measure before choosing granularity.
- No Python-side control flow that depends on tensor values; use `torch.where` / masking instead.
Profile first, optimize second. See the profiling docs / project memory for the trace-analysis workflow.
### I4. Validation pyramid
Offline RTF alone is necessary but not sufficient. Every new TTS model must pass all three:
| Layer | Catches | Tool |
|-------|---------|------|
| Offline RTF / duration check | Throughput regressions, missing audio, wrong sample rate | `end2end.py`, pytest e2e |
| Browser streaming playback | Delta/cumulative bugs, chunk boundary glitches, TTFP regressions | Gradio demo over `/v1/audio/speech?stream=true` |
| Concurrent requests | Per-request state leaks, codec window round-robin gaps | `max_num_seqs>1` smoke test with 4+ parallel prompts |
Declaring a model "done" without all three has shipped regressions more than once.
### I5. Per-request state is owned by the request, not the model
If the model caches *anything* across `forward()` calls (streaming generators, codec buffers, sliding-window pads, CUDA graph state), key it by request ID:
```python
self._state: dict[str, YourState] = {} # request_key → state
# fetch: request_key = str(info.get("_omni_req_id", "0"))
# free on finish: del self._state[request_key]
```
A shared buffer silently corrupts audio across concurrent requests — the symptom is crosstalk or truncation only under load.
## Phase 1: HuggingFace Reference
**Goal**: Understand the reference implementation and verify it produces correct audio.
@@ -85,6 +165,48 @@ When audio output is wrong, check in this order:
6. **Tensor layout**: Codebook-major vs frame-major ordering
7. **dtype**: Float32 for codec decoders (autocast can corrupt audio)
### Streaming Correctness Rules (single-stage and two-stage)
These bugs appear in almost every new TTS PR. Check all before the first push. See also the cross-cutting invariants I1 (output contract) and I5 (per-request state) above — the rules below are the Phase 2-specific instances of those invariants:
- **Accumulate codes across AR steps** — each `forward()` appends new codes; do not reset between steps or audio will be truncated (fish speech: `fix: accumulate audio_codes across steps`)
- **Emit delta audio, not full waveform** — in streaming mode yield only the new chunk per step, not the re-decoded full waveform from step 0 (fish speech: `fix: emit delta audio not full waveform`)
- **All return paths must emit `model_outputs`** — if any early-return branch skips setting `model_outputs`, the serving layer silently drops that step's audio (fish speech: `fix: ensure ALL return paths emit model_outputs`)
- **Per-request state isolation** — for batched concurrent requests, key all state by request ID; a shared buffer corrupts audio across requests (fish speech: `fix: per-request vocode + delta emission`)
- **Codec tensor device** — move codec codes to the codec decoder's device before calling decode; mismatches cause silent CPU fallback or crashes (fish speech: `fix: use model device for CUDA stream`)
- **AR stage `max_num_seqs`** — set to at least 4 in production configs; for single-stage models this is the only stage. For two-stage models, Stage 0 (AR) needs `max_num_seqs ≥ 4` to pipeline concurrent requests; Stage 1 (codec decoder) typically uses `max_num_seqs: 1` intentionally. Default of 1 everywhere causes audio gaps under concurrency because the codec window round-robins across requests (RFC #2568)
### Optional Dependency Handling
Patch optional dependencies (`torchaudio` / `torchcodec` / `soundfile`) at
the top of `load_weights()`, not at module import. Failures to do so cause
cryptic errors only on environments missing the optional package — after
the model is already deployed. See
[references/optional-deps.md](references/optional-deps.md) for the full
pattern, signature constraints, and MOSS-TTS-Nano reference.
### Single-Stage AR Pattern (alternative to two-stage)
When the upstream model cannot be cleanly split into an AR stage and a
separate decoder, run the full pipeline inside a single AR worker and
stream audio through a per-request `inference_stream()` generator keyed by
`_omni_req_id`. Stage config must set `worker_type: ar`,
`engine_output_type: audio`, `final_output: true`, `is_comprehension: true`,
and `async_chunk: false` at the top level. Only extract params from
`additional_information` that you actually forward, or pre-commit fails
`ruff F841`.
Full walkthrough with the complete `forward()` / `_create_stream_gen()`
skeleton and stage-config fields:
[references/single-stage-ar.md](references/single-stage-ar.md). For an
in-tree reference, look for any single-stage AR model under
`vllm_omni/model_executor/models/` — e.g. the MOSS-TTS-Nano integration when
it lands.
**VoxCPM2 is a different pattern** and should not reuse this skeleton — it
runs the base LM under vLLM PagedAttention with external side-computation.
See `plan/voxcpm2_native_ar_design.md`.
### Deliverables
- Model files in `vllm_omni/model_executor/models/<model_name>/`
@@ -98,10 +220,66 @@ When audio output is wrong, check in this order:
### Steps
1. **Register in `serving_speech.py`**:
- Add model stage name to `_TTS_MODEL_STAGES` set
- Add model detection flag (e.g., `_is_fish_speech`)
- Implement prompt builder method (e.g., `_build_fish_speech_prompt()`)
1. **Register in `serving_speech.py`** — add all 5 points in a **single commit**;
partial integration causes hard-to-debug failures. This file is modified by every
model PR and is the most common source of rebase conflicts — see conflict note below.
**Point 1** — stage constant (near the top, alongside the other `_*_TTS_MODEL_STAGES` sets):
```python
_YOUR_MODEL_TTS_MODEL_STAGES = {"your_stage_key"}
```
**Point 2** — union into `_TTS_MODEL_STAGES`:
```python
_TTS_MODEL_STAGES: set[str] = (
...
| _YOUR_MODEL_TTS_MODEL_STAGES
)
```
**Point 3** — model type detection in `_detect_tts_model_type()`:
```python
if model_stage in _YOUR_MODEL_TTS_MODEL_STAGES:
return "your_model"
```
**Point 4** — validation dispatch in `_validate_tts_request()`:
```python
if self._tts_model_type == "your_model":
return self._validate_your_model_request(request)
```
**Point 5** — validation + parameter-builder methods:
```python
def _validate_your_model_request(self, request) -> str | None:
if not request.input or not request.input.strip():
return "Input text cannot be empty"
return None
def _build_your_model_params(self, request) -> dict:
params = {"text": [request.input]}
if request.voice is not None:
params["voice"] = [request.voice]
return params
```
Wire `_build_your_model_params` into `_create_tts_request()` alongside the other
model-specific param builders.
> **Two dispatch patterns coexist**: Fish Speech uses a `self._is_fish_speech` boolean
> instance attribute checked before `elif self._is_tts`, while all newer models
> (CosyVoice3, MOSS-TTS-Nano) use the `_tts_model_type` string returned by
> `_detect_tts_model_type()`. For new models, always use the `_tts_model_type` string
> pattern — do not add new `_is_*` flags.
> **Unused variable rule**: only extract fields in `_build_your_model_params` that
> are actually forwarded to the model. Unused extractions fail `ruff F841`.
> For voice-cloning fields (`ref_audio` → `prompt_audio_path`, `ref_text` →
> `prompt_text`), add them to the param builder and verify they reach the model call.
**Rebase conflict note**: when rebasing onto `main` after another model was merged,
`serving_speech.py` will conflict. Resolution: always keep *both* the upstream
model's additions and your own — never discard either side.
2. **Handle model-specific parameters**:
- Voice cloning: `ref_audio` encoding and prompt injection
- `max_new_tokens` override in sampling params
@@ -129,11 +307,29 @@ def build_voice_clone_prompt(ref_audio_path: str, text: str, codec) -> list:
### Deliverables
- Updated `serving_speech.py` with model-specific prompt builder
- Updated `serving_speech.py` with all 5 integration points (single commit)
- Client scripts and server launcher
- Gradio demo with streaming and voice cloning UI
- E2E online serving test (`tests/e2e/online_serving/test_<model>.py`)
- Buildkite CI entry in `.buildkite/test-merge.yml`
- Documentation (offline + online serving docs)
### E2E test pitfalls to avoid
- **One `OmniServerParams` set per file.** `omni_server` is module-scoped; a second
id in the same file forces mid-module teardown/restart and exposes startup
races (`APIConnectionError` on the first request post-restart). Split variants
into separate files instead.
- **No external URL fetches from the server.** CI and some dev hosts can't
reach `raw.githubusercontent.com` over TLS. Inline ref audio as
`data:audio/wav;base64,...`; the serving layer accepts both URL and data URL.
- **Use the harness readiness gate.** The fixture waits for HTTP 200 on
`/health`; don't add `time.sleep` in tests. If warmup is incomplete, make
`/health` return non-200 until you're actually ready.
- **Mark with `@pytest.mark.core_model` + `hardware_test(res={"cuda": "H100"})`**
so the test lands in `test-ready.yml` (triggered by the `ready` label) rather
than only nightly.
## Phase 4: Async Chunk (Streaming)
**Goal**: Enable inter-stage streaming so audio chunks are produced while AR generation continues.
@@ -199,35 +395,10 @@ Stage 0 (AR) Stage 1 (Decoder)
- Fixed batch size (fall back to eager for other sizes)
- No dynamic control flow inside the graph
### Example: Code Predictor CUDA Graph (Qwen3-TTS)
```python
import torch
class CodePredictorGraph:
"""Captures the 16-step code predictor AR loop as a single CUDA graph."""
def setup_graph(self, device: torch.device, kv_heads: int = 4, head_dim: int = 64):
self.num_steps = 16
self.kv_cache = torch.zeros(1, kv_heads, self.num_steps, head_dim, device=device)
self.positions = torch.arange(self.num_steps, device=device)
self.causal_mask = torch.tril(torch.ones(self.num_steps, self.num_steps, device=device))
self.input_buf = torch.zeros(1, 1, kv_heads * head_dim, device=device)
self.output_buf = torch.zeros(1, self.num_steps, device=device, dtype=torch.long)
# Warm up, then: self.graph = torch.cuda.CUDAGraph(); self.graph.capture(...)
def run_graph(self, initial_input: torch.Tensor) -> torch.Tensor:
self.input_buf.copy_(initial_input)
self.graph.replay()
return self.output_buf.clone()
```
### Performance Expectations
Based on Qwen3-TTS code predictor experience:
- **3-5x speedup** for the graphed component
- Only effective for fixed batch sizes (typically batch_size=1)
- Falls back to eager mode for unsupported configurations
See [references/cuda-graph-example.md](references/cuda-graph-example.md) for
a worked skeleton (Qwen3-TTS code predictor, 16-step AR loop), performance
expectations (3–5× on the graphed component for fixed batch_size=1), and the
graph-safety constraints you must honor inside the captured region.
### Deliverables
@@ -235,10 +406,33 @@ Based on Qwen3-TTS code predictor experience:
- Benchmark script comparing eager vs graph performance
- Documentation of constraints and fallback behavior
## Phase 6: Pre-commit and DCO
**Goal**: Every commit passes `pre-commit` lint and carries a DCO
`Signed-off-by` line that matches the author email.
- Install hooks once: `pre-commit install`.
- Run `pre-commit run --files <changed-files>` before every push; accept any
auto-fixes, stage, re-commit.
- Sign every commit with `git commit -s`. DCO checks that author email and
`Signed-off-by` email match — `git config user.email` must match your
GitHub account email.
Common pre-commit failures, recovery commands for missing sign-off, and the
full `pre-commit run` invocation for a TTS model:
[references/precommit-dco.md](references/precommit-dco.md).
## Integration Checklist
Use this checklist when integrating a new TTS model:
### Cross-Cutting Invariants (verify at end of every phase)
- [ ] I1: `forward()` docstring states cumulative vs delta; consolidation path audited end-to-end
- [ ] I2: Tests / examples / benchmarks never use `dict.get(a) or dict.get(b)` on tensor values; list form handled
- [ ] I3: No `.item()` / `.cpu()` / Python branch on tensor values inside per-step loops
- [ ] I4: Offline RTF, browser streaming playback, and concurrent-request smoke test all pass
- [ ] I5: Any cross-step cache keyed by `_omni_req_id`; entries freed when the request finishes
### Phase 1: HF Reference
- [ ] Reference model runs and produces correct audio
- [ ] Architecture documented (stages, codebooks, tokens, sample rate)
@@ -248,17 +442,27 @@ Use this checklist when integrating a new TTS model:
- [ ] Model registered in `registry.py`
- [ ] Config classes created with `model_type` registration
- [ ] Stage 0 (AR) implemented and generates correct tokens
- [ ] Stage 1 (Decoder) produces correct audio from tokens
- [ ] Stage 1 (Decoder) produces correct audio from tokens — dtype float32 for codec decoder
- [ ] Stage 1 `max_num_seqs` ≥ 4 in production config (default 1 causes gaps under concurrency)
- [ ] Optional dependency fallbacks handled at `load_weights()` time (torchaudio/soundfile/etc.)
- [ ] Streaming: codec codes accumulated across AR steps (not reset per step)
- [ ] Streaming: delta audio emitted per chunk, not full re-decoded waveform
- [ ] Streaming: all `forward()` return paths emit `model_outputs`
- [ ] Streaming: per-request state keyed by request ID (not shared across requests)
- [ ] Streaming: codec tensors moved to codec decoder device before decode
- [ ] Stage config YAML created
- [ ] `end2end.py` produces audio matching reference quality
- [ ] README.md written
### Phase 3: Online Serving
- [ ] Model added to `serving_speech.py`
- [ ] All 5 `serving_speech.py` integration points added in one commit
- [ ] Only extract params in `_build_*_params` that are forwarded to the model call (ruff F841)
- [ ] Prompt builder handles text input correctly
- [ ] Voice cloning works (if supported)
- [ ] All response formats work (wav, mp3, flac, pcm)
- [ ] Client scripts and server launcher created
- [ ] E2E online serving test written (`tests/e2e/online_serving/test_<model>.py`)
- [ ] Buildkite CI entry added to `.buildkite/test-merge.yml`
- [ ] Gradio demo working
- [ ] Documentation added (offline + online docs, nav, supported models)
@@ -276,9 +480,25 @@ Use this checklist when integrating a new TTS model:
- [ ] Benchmark shows meaningful speedup
- [ ] Fallback to eager works for unsupported configs
### Phase 6: Pre-commit and DCO
- [ ] `pre-commit run --files <changed>` passes before every push
- [ ] Every commit has `Signed-off-by` matching the author email (`git commit -s`)
- [ ] `git config user.email` matches the email registered on your GitHub account
- [ ] Details and failure-recovery commands: [references/precommit-dco.md](references/precommit-dco.md)
## References
- [TTS audio skill](../vllm-omni-audio-tts/SKILL.md) -- supported models and usage
- [Fish Speech integration](../vllm-omni-audio-tts/references/fish-speech.md) -- complete example of Phases 1-3
- [Qwen3-TTS reference](../vllm-omni-audio-tts/references/qwen-tts.md) -- complete example of all 5 phases
In-skill references (details split out of the main body):
- [references/single-stage-ar.md](references/single-stage-ar.md) — full `forward()` / generator skeleton for the MOSS-TTS-Nano-style pattern
- [references/optional-deps.md](references/optional-deps.md) — torchaudio / torchcodec fallback pattern
- [references/cuda-graph-example.md](references/cuda-graph-example.md) — Qwen3-TTS code-predictor CUDA graph skeleton
- [references/precommit-dco.md](references/precommit-dco.md) — full pre-commit invocation, failure table, DCO recovery
Project docs and adjacent skills:
- [TTS audio skill](../vllm-omni-audio-tts/SKILL.md) — supported models and usage
- [Fish Speech integration](../vllm-omni-audio-tts/references/fish-speech.md) — complete example of Phases 1–3
- [Qwen3-TTS reference](../vllm-omni-audio-tts/references/qwen-tts.md) — complete example of all 5 phases
- [Adding a TTS model (developer guide)](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/model/adding_tts_model.md)
- `plan/voxcpm2_native_ar_design.md` — VoxCPM2's vLLM-native AR + side-computation pattern (distinct from the generator-based single-stage described above)
.claude/skills/add-tts-model/references/cuda-graph-example.md
+42
View File
@@ -0,0 +1,42 @@
# CUDA Graph Example: Qwen3-TTS Code Predictor
Reference sketch for capturing the 16-step code-predictor AR loop as a single
CUDA graph. Adapt the shapes, number of steps, and KV-head layout to your
model.
```python
import torch
class CodePredictorGraph:
"""Captures the 16-step code predictor AR loop as a single CUDA graph."""
def setup_graph(self, device: torch.device, kv_heads: int = 4, head_dim: int = 64):
self.num_steps = 16
self.kv_cache = torch.zeros(1, kv_heads, self.num_steps, head_dim, device=device)
self.positions = torch.arange(self.num_steps, device=device)
self.causal_mask = torch.tril(torch.ones(self.num_steps, self.num_steps, device=device))
self.input_buf = torch.zeros(1, 1, kv_heads * head_dim, device=device)
self.output_buf = torch.zeros(1, self.num_steps, device=device, dtype=torch.long)
# Warm up, then: self.graph = torch.cuda.CUDAGraph(); self.graph.capture(...)
def run_graph(self, initial_input: torch.Tensor) -> torch.Tensor:
self.input_buf.copy_(initial_input)
self.graph.replay()
return self.output_buf.clone()
```
## Performance expectations (Qwen3-TTS code predictor)
- **3–5× speedup** on the graphed component.
- Effective only for fixed batch sizes (typically `batch_size=1`).
- Fall back to eager for any shape/config that wasn't captured — do not try
to recapture per request.
## Graph-safety constraints
- `torch.argmax` instead of `torch.multinomial`.
- Fixed batch size.
- No Python control flow that branches on tensor values inside the captured
region (use `torch.where` / masks).
- No `.item()`, `.cpu()`, `.tolist()` — each would break the capture or
cause a GPU→CPU sync during replay.
.claude/skills/add-tts-model/references/optional-deps.md
+47
View File
@@ -0,0 +1,47 @@
# Optional Dependency Handling
Models that rely on `torchaudio`, `torchcodec`, `soundfile`, or other optional
packages must handle the missing-package case at import time, not at call
time. Failing to do this causes cryptic errors only on environments without
the optional package — after the model is already deployed.
## Pattern (used in MOSS-TTS-Nano)
```python
def _patch_torchaudio_load() -> None:
"""Fallback torchaudio.load/save to soundfile if torchcodec is unavailable."""
try:
import torchcodec # noqa: F401
return # torchcodec present, torchaudio works as-is
except ImportError:
pass
import soundfile as sf
import torchaudio
def _sf_load(path, **kwargs):
data, sr = sf.read(str(path), dtype="float32", always_2d=True)
return torch.from_numpy(data).T, sr
torchaudio.load = _sf_load
# patch .save similarly if needed
```
## Rules
- Mirror the full signature of the replaced function. `torchaudio.load`
accepts `frame_offset`, `num_frames`, `normalize`, `channels_first`,
`format` — missing any of them causes `TypeError` from calling code.
- Catch `except Exception`, not just `ImportError`. `import torchaudio`
itself can fail with non-`ImportError` errors on broken installs. Log the
exception type and message (`logger.warning("torchaudio probe failed: %s: %s",
type(exc).__name__, exc)`) before falling back, so unrelated errors are not
silently swallowed.
- Call the patch function at the top of `load_weights()` before loading any
audio assets. Do not call it at module import time.
## Reference implementation
Any in-tree model that patches `torchaudio.load` in its `load_weights()` —
e.g. MOSS-TTS-Nano's `modeling_moss_tts_nano.py` once that integration
lands.
.claude/skills/add-tts-model/references/precommit-dco.md
+54
View File
@@ -0,0 +1,54 @@
# Pre-commit and DCO
Every commit must pass `pre-commit` lint and carry a `Signed-off-by` line
that matches the commit author email.
## Pre-commit
Install hooks once:
```bash
pre-commit install
```
Run before every push on the files you changed:
```bash
pre-commit run --files \
vllm_omni/model_executor/models/<model_name>/*.py \
vllm_omni/entrypoints/openai/serving_speech.py \
vllm_omni/model_executor/models/registry.py \
tests/e2e/offline_inference/test_<model_name>.py \
tests/e2e/online_serving/test_<model_name>.py
```
When pre-commit **modifies files** (ruff format auto-fix), it exits non-zero
but the changes are correct — stage the modified files and re-commit.
| Failure | Root cause | Fix |
|---------|-----------|-----|
| `ruff F841` | Variable extracted but never forwarded to model call | Remove the extraction or wire it through |
| `ruff E402` | Import added below function definitions | Move to top-level import block |
| `ruff format` | Line length, spacing, quote style | Accept auto-fix, stage, re-commit |
## DCO sign-off
Every commit must carry `Signed-off-by: Your Name <your@email.com>`. Use
`-s`:
```bash
git commit -s -m "feat(<model>): add <Model> TTS support"
```
Or set it permanently: `git config format.signOff true`.
The DCO check verifies that the commit author email matches the
`Signed-off-by` line. Confirm `git config user.email` matches your GitHub
account email before committing.
Fix a missing or mismatched sign-off on the latest commit:
```bash
git commit --amend -s --no-edit
git push origin <branch> --force-with-lease
```
.claude/skills/add-tts-model/references/single-stage-ar.md
+108
View File
@@ -0,0 +1,108 @@
# Single-Stage AR Pattern
When the upstream model cannot be cleanly split into an AR stage and a separate
decoder (e.g. MOSS-TTS-Nano, or any model that bundles AR + codec via an
`inference_stream()` generator), run the whole pipeline inside a single AR
worker that yields audio chunks per request.
This is distinct from VoxCPM2's pattern, which also runs in a single stage but
uses vLLM's native PagedAttention on the base language model with diffusion /
VAE side-computation outside vLLM — see
`plan/voxcpm2_native_ar_design.md` for that variant.
## Implementation
1. **Single model file** — load both AR LM and codec inside
`modeling_<model>.py`.
2. **Load weights in `load_weights()`**, not `__init__()` — vLLM initializes
distributed state before any CUDA allocations.
3. **Stream via a per-request generator** stored in `self._stream_gens`:
```python
class YourModelForCausalLM(nn.Module):
def __init__(self, *, vllm_config, prefix=""):
super().__init__()
self._lm = None # populated in load_weights()
self._stream_gens: dict = {} # request_key → generator
def load_weights(self, weights):
# Load self._lm here, after vLLM distributed init
...
def forward(
self,
input_ids,
positions,
intermediate_tensors=None,
inputs_embeds=None,
runtime_additional_information: list[dict] | None = None, # one dict per request
**kwargs,
) -> OmniOutput:
infos = runtime_additional_information or [{}]
# Skip dummy/profiling calls
if not runtime_additional_information or all(i.get("_is_dummy") for i in infos):
self._ar_emit_stop_token = True
return OmniOutput(...) # return empty outputs
outputs, last_flags = [], []
for info in infos:
request_key = str(info.get("_omni_req_id", "0")) # per-request ID from vLLM
if request_key not in self._stream_gens:
self._stream_gens[request_key] = self._create_stream_gen(info)
try:
chunk, is_last = next(self._stream_gens[request_key])
except StopIteration:
chunk, is_last = torch.zeros(0), True
if is_last:
del self._stream_gens[request_key]
outputs.append(chunk)
last_flags.append(is_last)
self._ar_emit_stop_token = all(last_flags)
return OmniOutput(multimodal_outputs={"model_outputs": outputs, ...})
def _create_stream_gen(self, info: dict):
"""Yield (waveform_tensor, is_last) tuples from inference_stream()."""
for event in self._lm.inference_stream(...):
if event["type"] == "audio":
yield event["waveform"], False
elif event["type"] == "result":
# Fallback: some models emit a single "result" event instead of
# incremental "audio" events — handle both paths
yield event.get("waveform", torch.zeros(0)), True
return
yield torch.zeros(0), True
def compute_logits(self, hidden_states, sampling_metadata):
# Emit EOS only after the last chunk so the AR scheduler ends the request
...
```
## Key points
- `runtime_additional_information` is the correct parameter name (not
`**kwargs`) — it carries one dict per request in the batch.
- The request ID is `info.get("_omni_req_id")` — set by vLLM, not by user code.
- Handle both `"audio"` (incremental) and `"result"` (final combined) event
types from upstream models.
## Stage config
Single stage with `worker_type: ar`, `engine_output_type: audio`,
`final_output: true`, `is_comprehension: true`, and `async_chunk: false` at
the top level. Omitting any of these causes silent misclassification in the
serving layer.
## Lint discipline
Only extract variables from `additional_information` that you actually
forward to the model call — unused extractions trip `ruff F841` in
pre-commit.
## Reference implementation
Look for any single-stage AR model under
`vllm_omni/model_executor/models/` — e.g. `moss_tts_nano/` when its
integration lands. If none is in tree yet, follow the skeleton above and
cross-check against the `is_comprehension: true` / `async_chunk: false`
dispatch in `vllm_omni/entrypoints/openai/serving_speech.py`.
benchmarks/accuracy/README.md
+2 -2
View File
@@ -23,5 +23,5 @@ Test guidance:
- Local static/self-checks live in `tests/benchmarks/test_accuracy_bench_utils.py`.
- End-to-end generation/evaluation should be validated in a remote GPU
environment. In the current repo marker system there is `L4` but no `L5`
marker, so benchmark smoke tests should be wired as `advanced_model +
benchmark + L4` when GPU capacity is available.
marker, so benchmark smoke tests should be wired as `full_model +
benchmark + L4` for nightly when GPU capacity is available.
benchmarks/accuracy/image_to_image/README.md
+2 -2
View File
@@ -99,5 +99,5 @@ Notes:
- This flow requires the optional Hugging Face `datasets` package.
- `generate` writes `generation_manifest.json` with local output coverage.
- The current repo marker set exposes `L4` but not `L5`, so if you promote an
end-to-end smoke test into CI, use the existing `advanced_model`, `benchmark`,
and `L4` markers or introduce a new repo-wide marker explicitly first.
end-to-end smoke test into CI, use the `full_model`, `benchmark`,
and `L4` markers for nightly (or `advanced_model` for merge) or introduce a new repo-wide marker explicitly first.
benchmarks/diffusion/backends.py
+12
View File
@@ -122,6 +122,18 @@ async def async_request_chat_completions(
output.peak_memory_mb = first_item.get("peak_memory_mb", 0.0)
except (IndexError, TypeError, AttributeError):
pass
if (not output.stage_durations or output.peak_memory_mb == 0.0) and isinstance(
resp_json.get("metrics"), dict
):
m = resp_json["metrics"]
if not output.stage_durations and isinstance(m.get("stage_durations"), dict):
output.stage_durations = m.get("stage_durations") or {}
if output.peak_memory_mb == 0.0 and m.get("peak_memory_mb") is not None:
try:
output.peak_memory_mb = float(m.get("peak_memory_mb") or 0.0)
except (TypeError, ValueError):
pass
else:
output.error = f"HTTP {response.status}: {await response.text()}"
output.success = False
benchmarks/diffusion/diffusion_benchmark_serving.py
+4 -4
View File
@@ -12,15 +12,15 @@ Supports multiple backends:
- v1/videos: Use /v1/videos endpoint
Usage:
# Video (vllm-omni backend)
# Video (v1/videos backend)
t2v:
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
--backend vllm-omni --dataset vbench --task t2v --num-prompts 10 \
--backend v1/videos --dataset vbench --task t2v --num-prompts 10 \
--height 480 --width 640 --fps 16 --num-frames 80
i2v:
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
--backend vllm-omni --dataset vbench --task i2v --num-prompts 10
--backend v1/videos --dataset vbench --task i2v --num-prompts 10
# Image (vllm-omni backend)
@@ -49,7 +49,7 @@ Usage:
--backend openai --dataset vbench --task t2i --num-prompts 10 \
--height 1024 --width 1024 --port 3000
# Video (v1/vedeos)
# Video (v1/videos)
t2v:
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
--backend v1/videos --dataset random --task t2v --num-prompts 1 \
benchmarks/qwen3-omni/README.md
-86
View File
@@ -1,86 +0,0 @@
# Benchmarks Guide
This README explains how to (1) prepare benchmark datasets and (2) run the provided Qwen3-Omni benchmarks.
## 1) Prepare the dataset (SeedTTS top100)
```bash
cd benchmarks/build_dataset
pip install gdown
# Download SeedTTS test set from Google Drive
gdown 1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP
# Extract
tar -xf seedtts_testset.tar
# Copy metadata and extract top-100 prompts
cp seedtts_testset/en/meta.lst meta.lst
python extract_prompts.py -i meta.lst -o top100.txt -n 100
# (Optional) clean up to save space
rm -rf seedtts_testset seedtts_testset.tar meta.lst
```
Artifacts:
- `benchmarks/build_dataset/top100.txt` — 100 text prompts (one per line).
## 2) Run benchmarks
All commands assume repo root (`vllm-omni`).
### A. Transformers benchmark (offline, HF Transformers)
```
bash benchmarks/qwen3-omni/transformers/eval_qwen3_moe_omni_transformers.sh
```
What it does:
- Runs `qwen3_omni_moe_transformers.py` over `top100.txt` with `--num-prompts 100`.
- Outputs to `benchmarks/qwen3-omni/transformers/benchmark_results/`:
- `perf_stats.json` — aggregated & per-prompt TPS/latency (thinker/talker/code2wav/overall).
- `results.json` — per-prompt outputs and audio paths.
- `audio/` — ~100 generated `.wav` files.
Key checks:
- `overall_tps` and `*_tps_avg` should be non-zero and reasonably stable.
- Investigate any 0/NaN or unusually low TPS / long-tail latency.
### B. vLLM Omni end-to-end benchmark (pipeline)
```
bash benchmarks/qwen3-omni/vllm_omni/eval_qwen3_moe_omni.sh
```
What it does:
- Runs `examples/offline_inference/qwen3_omni/end2end.py` with `--log-stats`.
- Uses `benchmarks/build_dataset/top100.txt` and writes to:
- Logs: `benchmarks/qwen3-omni/vllm_omni/logs/`
- `omni_pipeline_text.orchestrator.stats.jsonl` — per-stage latency stats.
- `omni_pipeline_text.overall.stats.jsonl` — end-to-end latency/TPS.
- `omni_pipeline_text.stage{0,1,2}.log` — per-stage detailed logs/errors.
- Outputs: `benchmarks/qwen3-omni/vllm_omni/outputs/` — ~100 text and `.wav` files.
Key checks:
- Overall stats: end-to-end latency/TPS should be reasonable.
- Orchestrator stats: per-stage latency should be stable; investigate long tails.
- Stage logs: ensure no errors and no unusually slow stages.
## Performance snapshot
The chart below summarizes our measured Qwen3-Omni MoE end-to-end benchmark, comparing vLLM-Omni against HF Transformers. It shows the overall throughput advantage for vLLM-Omni. These are actual experiment results—please refer to this performance when evaluating or reproducing the benchmark.
![vLLM-Omni vs HF](./vllm-omni-vs-hf.png)
## Directory layout
- `benchmarks/build_dataset/` — dataset prep utilities (e.g., SeedTTS top100).
- `benchmarks/<model>/vllm_omni/` — vLLM-Omni pipeline benchmarks, logs, outputs.
- Add new tasks under `benchmarks/<model>/...` with the same pattern: `transformers/`, `vllm_omni/`, task-specific README, and (optionally) dataset prep notes.
- `benchmarks/<model>/vllm-omni-vs-hf.png` — current performance snapshot (overall throughput comparison).
- `benchmarks/<model>/transformers/` — HF Transformers benchmarks (offline reference).
## Troubleshooting
- Make sure GPU/driver/FlashAttention2 requirements are met for the chosen model.
- If downloads fail, confirm network access to Google Drive (`gdown`) and Hugging Face.
- If audio files are missing, check for errors in stage logs or model generation.***
benchmarks/qwen3-omni/transformers/eval_qwen3_moe_omni_transformers.sh
-26
View File
@@ -1,26 +0,0 @@
#!/bin/bash
# Qwen3-Omni Transformers Benchmark Evaluation Script
# This script must be run from the vllm-omni root directory
# Get the directory where this script is located
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
# Navigate to vllm-omni root directory (4 levels up from script location)
VLLM_OMNI_ROOT="$(cd "$SCRIPT_DIR/../../.." && pwd)"
cd "$VLLM_OMNI_ROOT" || { echo "Error: Failed to navigate to vllm-omni directory"; exit 1; }
echo "Working directory: $(pwd)"
# Verify we're in the correct directory and run benchmark
if [[ ! -f "benchmarks/qwen3-omni/transformers/qwen3_omni_moe_transformers.py" ]]; then
echo "Error: Not in vllm-omni root directory. Please run from vllm-omni folder."
else
cd benchmarks/qwen3-omni/transformers
python qwen3_omni_moe_transformers.py --prompts-file ../../build_dataset/top100.txt --num-prompts 100
echo "Logs and outputs are saved to $(pwd)/benchmark_results:"
echo " - perf_stats.json Aggregated/per-prompt TPS and latency (thinker/talker/code2wav/overall)"
echo " - results.json Per-prompt outputs and audio paths"
echo " - audio/ Generated wav files, there should be 100 wav file generated"
echo "Key checks: overall_tps and *_tps_avg should be non-zero and stable; investigate 0/NaN or unusually low TPS/long-tail latency."
fi
benchmarks/qwen3-omni/transformers/qwen3_omni_moe_model.py
-265
View File
@@ -1,265 +0,0 @@
import time
import torch
from transformers import Qwen3OmniMoeForConditionalGeneration
class Qwen3OmniMoeForConditionalGenerationWithLogging(Qwen3OmniMoeForConditionalGeneration):
@torch.no_grad()
def generate(
self,
input_ids: torch.Tensor | None = None,
speaker: str = "Ethan",
use_audio_in_video: bool = False,
return_audio: bool | None = None,
thinker_max_new_tokens: int = 1024,
thinker_eos_token_id: int = 151645,
talker_max_new_tokens: int = 4096,
talker_do_sample: bool = True,
talker_top_k: int = 50,
talker_top_p: float = 1.0,
talker_temperature: float = 0.9,
talker_repetition_penalty: float = 1.05,
**kwargs,
):
total_t0 = time.time()
perf_stats = {
"thinker_tokens": 0,
"thinker_time_s": 0.0,
"thinker_tps": 0.0,
"talker_tokens": 0,
"talker_time_s": 0.0,
"talker_tps": 0.0,
"code2wav_tokens": 0,
"code2wav_time_s": 0.0,
"code2wav_tps": 0.0,
"total_tokens": 0,
"total_time_s": 0.0,
"total_tps": 0.0,
}
if return_audio and not self.has_talker:
raise ValueError(
"Cannot use talker when talker module not initialized. "
"Use `enable_talker` method or set enable_talker in config "
"to enable talker."
)
if return_audio is None:
return_audio = self.has_talker
shared_kwargs = {"use_audio_in_video": use_audio_in_video}
thinker_kwargs = {
"max_new_tokens": thinker_max_new_tokens,
"eos_token_id": thinker_eos_token_id,
}
talker_kwargs = {}
token2wav_kwargs = {}
if return_audio:
speaker_id = self.config.talker_config.speaker_id.get(speaker.lower())
if speaker_id is None:
raise NotImplementedError(f"Speaker {speaker} not implemented")
if input_ids.shape[0] != 1:
raise NotImplementedError("Qwen3-Omni currently does not support batched inference with audio output")
talker_suppressed_tokens = [
i
for i in range(
self.config.talker_config.text_config.vocab_size - 1024,
self.config.talker_config.text_config.vocab_size,
)
if i != self.config.talker_config.codec_eos_token_id
] # Suppress additional special tokens, should not be predicted
talker_kwargs = {
"max_new_tokens": talker_max_new_tokens,
"do_sample": talker_do_sample,
"top_k": talker_top_k,
"top_p": talker_top_p,
"temperature": talker_temperature,
"eos_token_id": self.config.talker_config.codec_eos_token_id,
"repetition_penalty": talker_repetition_penalty,
"suppress_tokens": talker_suppressed_tokens,
"output_hidden_states": True,
"return_dict_in_generate": True,
}
token2wav_kwargs = {}
for key, value in kwargs.items():
if key.startswith("thinker_"):
thinker_kwargs[key[len("thinker_") :]] = value
elif key.startswith("talker_"):
talker_kwargs[key[len("talker_") :]] = value
elif key.startswith("token2wav_"):
token2wav_kwargs[key[len("token2wav_") :]] = value
# Process special input values
elif key == "feature_attention_mask":
thinker_kwargs[key] = value
talker_kwargs["audio_feature_lengths"] = torch.sum(value, dim=1)
elif key in ("input_features", "attention_mask"):
thinker_kwargs[key] = value
# Put other key to shared kwargs
else:
shared_kwargs[key] = value
# Merge kwargs
for key, value in shared_kwargs.items():
if key not in thinker_kwargs:
thinker_kwargs[key] = value
if key not in talker_kwargs and key in ["image_grid_thw", "video_grid_thw", "video_second_per_grid"]:
talker_kwargs[key] = value
if key not in token2wav_kwargs:
token2wav_kwargs[key] = value
# 1. Generate from thinker module
generate_audio = return_audio and self.has_talker
if generate_audio:
thinker_kwargs["output_hidden_states"] = True
thinker_kwargs["return_dict_in_generate"] = True
t0 = time.time()
thinker_result = self.thinker.generate(input_ids=input_ids, **thinker_kwargs)
t1 = time.time()
perf_stats["thinker_time_s"] = max(0.0, t1 - t0)
try:
prompt_len = int(input_ids.shape[1]) if input_ids is not None else 0
total_len = int(thinker_result.sequences.shape[-1])
thinker_out_len = max(0, total_len - prompt_len)
except Exception:
thinker_out_len = 0
perf_stats["thinker_tokens"] = thinker_out_len
perf_stats["thinker_tps"] = (
(thinker_out_len / perf_stats["thinker_time_s"]) if perf_stats["thinker_time_s"] > 0 else 0.0
)
if not generate_audio:
perf_stats["total_tokens"] = perf_stats["thinker_tokens"]
perf_stats["total_time_s"] = time.time() - total_t0
perf_stats["total_tps"] = (
(perf_stats["total_tokens"] / perf_stats["total_time_s"]) if perf_stats["total_time_s"] > 0 else 0.0
)
# attach stats to self
setattr(self, "_perf_stats_last", perf_stats)
if not hasattr(self, "_perf_stats_history"):
setattr(self, "_perf_stats_history", [])
self._perf_stats_history.append(perf_stats)
return thinker_result, None
# 2. Prepare talker input
thinker_embed = torch.cat([hidden_states[0] for hidden_states in thinker_result.hidden_states], dim=1).to(
self.talker.device
) # [1 t d]
thinker_hidden = torch.cat(
[
hidden_states[self.config.talker_config.accept_hidden_layer]
for hidden_states in thinker_result.hidden_states
],
dim=1,
).to(self.talker.device) # [1 t d]
im_start_indexes = torch.cat(
(
torch.nonzero(input_ids[0] == self.config.im_start_token_id).squeeze(),
torch.tensor([thinker_result.sequences.shape[-1]], device=input_ids.device, dtype=input_ids.dtype),
),
dim=-1,
).to(self.talker.device) # Shape [n_starts + 1]; Take batch 0 since batched inference is not supported here.
multimodal_mask = (
(thinker_result.sequences == self.config.thinker_config.audio_token_id) |
(thinker_result.sequences == self.config.thinker_config.image_token_id) |
(thinker_result.sequences == self.config.thinker_config.video_token_id)
).to(self.talker.device) # [1 t] # fmt: skip
talker_special_tokens = torch.tensor(
[[self.config.tts_bos_token_id, self.config.tts_eos_token_id, self.config.tts_pad_token_id]],
device=self.thinker.device,
dtype=input_ids.dtype,
)
tts_bos_embed, tts_eos_embed, tts_pad_embed = (
self.talker.text_projection(self.thinker.get_input_embeddings()(talker_special_tokens))
.to(self.talker.device)
.chunk(3, dim=1)
) # 3 * [1 1 d]
talker_input_embeds = [] # [1 t d]
talker_input_ids = []
# For every chatml parts
for i in range(len(im_start_indexes) - 1):
im_start_index = im_start_indexes[i]
segment_end_index = im_start_indexes[i + 1]
role_token = input_ids[0][im_start_index + 1]
# Talker should ignore thinker system prompt
if role_token == self.config.system_token_id:
continue
# Talker takes word embeddings for tokens and hidden state from `accept_hidden_layer` for multimodal inputs
elif role_token == self.config.user_token_id:
talker_user_part = self._get_talker_user_parts(
im_start_index, segment_end_index, multimodal_mask, thinker_hidden, thinker_embed
)
talker_input_embeds.append(talker_user_part)
talker_input_ids.append(thinker_result.sequences[:, im_start_index:segment_end_index])
# Take assistant output (for now)
elif role_token == self.config.assistant_token_id and i == len(im_start_indexes) - 2:
talker_assistant_embeds, talker_assistant_ids, trailing_text_hidden = self._get_talker_assistant_parts(
im_start_index,
segment_end_index,
speaker_id,
thinker_embed,
tts_pad_embed,
tts_bos_embed,
tts_eos_embed,
)
talker_input_embeds.append(talker_assistant_embeds)
talker_input_ids.append(talker_assistant_ids)
# History assistant output (ignore for now)
elif role_token == self.config.assistant_token_id and i != len(im_start_indexes) - 2:
continue
else:
raise AssertionError("Expect role id after <|im_start|> (assistant, user, system)")
talker_input_embed = torch.cat([embed.to(self.talker.device) for embed in talker_input_embeds], dim=1)
talker_input_id = torch.cat([embed.to(self.talker.device) for embed in talker_input_ids], dim=1)
t2 = time.time()
talker_result = self.talker.generate(
inputs_embeds=talker_input_embed,
trailing_text_hidden=trailing_text_hidden,
tts_pad_embed=tts_pad_embed,
talker_input_ids=talker_input_id, # Not use input_ids to prevent repetition penalty out of bound
**talker_kwargs,
)
t3 = time.time()
perf_stats["talker_time_s"] = max(0.0, t3 - t2)
talker_codes = (
torch.stack([hid[-1] for hid in talker_result.hidden_states if hid[-1] is not None], dim=1)
.transpose(1, 2)
.to(self.code2wav.device)
)
try:
# codes shape: (B, num_quantizers, T). We log T as token length.
perf_stats["talker_tokens"] = int(talker_codes.shape[-1])
except Exception:
perf_stats["talker_tokens"] = 0
perf_stats["talker_tps"] = (
(perf_stats["talker_tokens"] / perf_stats["talker_time_s"]) if perf_stats["talker_time_s"] > 0 else 0.0
)
t4 = time.time()
talker_wavs = self.code2wav.chunked_decode(talker_codes, chunk_size=300, left_context_size=25).float()
t5 = time.time()
perf_stats["code2wav_time_s"] = max(0.0, t5 - t4)
perf_stats["code2wav_tokens"] = perf_stats["talker_tokens"] # same T, not times 16
perf_stats["code2wav_tps"] = (
(perf_stats["code2wav_tokens"] / perf_stats["code2wav_time_s"])
if perf_stats["code2wav_time_s"] > 0
else 0.0
)
perf_stats["total_tokens"] = perf_stats["thinker_tokens"] + perf_stats["talker_tokens"]
perf_stats["total_time_s"] = time.time() - total_t0
perf_stats["total_tps"] = (
(perf_stats["total_tokens"] / perf_stats["total_time_s"]) if perf_stats["total_time_s"] > 0 else 0.0
)
setattr(self, "_perf_stats_last", perf_stats)
if not hasattr(self, "_perf_stats_history"):
setattr(self, "_perf_stats_history", [])
self._perf_stats_history.append(perf_stats)
return thinker_result, talker_wavs.float()
__all__ = [
"Qwen3OmniMoeForConditionalGenerationWithLogging",
]
benchmarks/qwen3-omni/transformers/qwen3_omni_moe_transformers.py
-275
View File
@@ -1,275 +0,0 @@
import argparse
import json
import os
import soundfile as sf
from qwen3_omni_moe_model import Qwen3OmniMoeForConditionalGenerationWithLogging
from qwen_omni_utils import process_mm_info
from tqdm import tqdm
from transformers import Qwen3OmniMoeProcessor
MODEL_PATH = "Qwen/Qwen3-Omni-30B-A3B-Instruct"
# MODEL_PATH = "Qwen/Qwen3-Omni-30B-A3B-Thinking"
def load_prompts(prompts_file: str) -> list[str]:
"""Load prompts from a text file, one prompt per line."""
prompts = []
with open(prompts_file, encoding="utf-8") as f:
for line in f:
line = line.strip()
if line:
prompts.append(line)
return prompts
def run_benchmark(
model,
processor,
prompts: list[str],
output_dir: str = "benchmark_results",
speaker: str = "Ethan",
use_audio_in_video: bool = True,
):
"""
Run benchmark on a list of prompts and collect performance stats.
Args:
model: The Qwen3OmniMoe model
processor: The Qwen3OmniMoe processor
prompts: List of text prompts to process
output_dir: Directory to save results
speaker: Speaker voice for audio output
use_audio_in_video: Whether to use audio in video
Returns:
tuple: (aggregated_stats, results, audio_outputs)
- aggregated_stats: dict with aggregated performance statistics
- results: list of dicts with per-prompt results
- audio_outputs: list of audio tensors/arrays (or None if no audio)
"""
os.makedirs(output_dir, exist_ok=True)
audio_dir = os.path.join(output_dir, "audio")
os.makedirs(audio_dir, exist_ok=True)
all_stats = []
results = []
audio_outputs = []
for idx, prompt in enumerate(tqdm(prompts, desc="Processing prompts")):
conversation = [
{
"role": "user",
"content": [{"type": "text", "text": prompt}],
},
]
# Preparation for inference
text = processor.apply_chat_template(conversation, add_generation_prompt=True, tokenize=False)
audios, images, videos = process_mm_info(conversation, use_audio_in_video=use_audio_in_video)
inputs = processor(
text=text,
audio=audios,
images=images,
videos=videos,
return_tensors="pt",
padding=True,
use_audio_in_video=use_audio_in_video,
)
inputs = inputs.to(model.device).to(model.dtype)
# Inference: Generation of the output text and audio
text_ids, audio = model.generate(
**inputs, speaker=speaker, thinker_return_dict_in_generate=True, use_audio_in_video=use_audio_in_video
)
# Decode output text
output_text = processor.batch_decode(
text_ids.sequences[:, inputs["input_ids"].shape[1] :],
skip_special_tokens=True,
clean_up_tokenization_spaces=False,
)[0]
# Collect performance stats
perf_stats = None
if hasattr(model, "_perf_stats_last"):
perf_stats = model._perf_stats_last.copy()
perf_stats["prompt_idx"] = idx
perf_stats["prompt"] = prompt
all_stats.append(perf_stats)
# Save audio and collect audio output
audio_path = None
audio_data = None
if audio is not None:
audio_data = audio.reshape(-1).detach().cpu().numpy()
audio_path = os.path.join(audio_dir, f"output_{idx:04d}.wav")
sf.write(
audio_path,
audio_data,
samplerate=24000,
)
audio_outputs.append(audio_data)
else:
audio_outputs.append(None)
# Save result
result = {
"idx": idx,
"prompt": prompt,
"output": output_text,
"audio_path": audio_path,
"perf_stats": perf_stats,
}
results.append(result)
# Aggregate statistics
aggregated_stats = aggregate_stats(all_stats)
# Save all results
results_path = os.path.join(output_dir, "results.json")
with open(results_path, "w", encoding="utf-8") as f:
json.dump(results, f, ensure_ascii=False, indent=2)
# Save aggregated stats
stats_path = os.path.join(output_dir, "perf_stats.json")
with open(stats_path, "w", encoding="utf-8") as f:
json.dump({"aggregated": aggregated_stats, "per_prompt": all_stats}, f, ensure_ascii=False, indent=2)
# Count saved audio files
num_audio_saved = sum(1 for a in audio_outputs if a is not None)
print(f"\nSaved {num_audio_saved} audio files to {audio_dir}/")
return aggregated_stats, results, audio_outputs
def aggregate_stats(all_stats: list[dict]) -> dict:
"""Aggregate performance statistics from multiple runs."""
if not all_stats:
return {}
keys = [
"thinker_tokens",
"thinker_time_s",
"thinker_tps",
"talker_tokens",
"talker_time_s",
"talker_tps",
"code2wav_tokens",
"code2wav_time_s",
"code2wav_tps",
"total_tokens",
"total_time_s",
"total_tps",
]
aggregated = {
"num_samples": len(all_stats),
}
for key in keys:
values = [s.get(key, 0) for s in all_stats if key in s]
if values:
aggregated[f"{key}_sum"] = sum(values)
aggregated[f"{key}_avg"] = sum(values) / len(values)
aggregated[f"{key}_min"] = min(values)
aggregated[f"{key}_max"] = max(values)
# Calculate overall throughput
total_tokens = aggregated.get("total_tokens_sum", 0)
total_time = aggregated.get("total_time_s_sum", 0)
if total_time > 0:
aggregated["overall_tps"] = total_tokens / total_time
return aggregated
def print_stats(stats: dict):
"""Print performance statistics in a formatted way."""
print("\n" + "=" * 60)
print("Performance Statistics Summary")
print("=" * 60)
print(f"\nNumber of samples: {stats.get('num_samples', 0)}")
print("\n--- Thinker ---")
print(f" Total tokens: {stats.get('thinker_tokens_sum', 0):.0f}")
print(f" Total time: {stats.get('thinker_time_s_sum', 0):.2f}s")
print(f" Avg TPS: {stats.get('thinker_tps_avg', 0):.2f}")
print(f" Min TPS: {stats.get('thinker_tps_min', 0):.2f}")
print(f" Max TPS: {stats.get('thinker_tps_max', 0):.2f}")
print("\n--- Talker ---")
print(f" Total tokens: {stats.get('talker_tokens_sum', 0):.0f}")
print(f" Total time: {stats.get('talker_time_s_sum', 0):.2f}s")
print(f" Avg TPS: {stats.get('talker_tps_avg', 0):.2f}")
print(f" Min TPS: {stats.get('talker_tps_min', 0):.2f}")
print(f" Max TPS: {stats.get('talker_tps_max', 0):.2f}")
print("\n--- Code2Wav ---")
print(f" Total tokens: {stats.get('code2wav_tokens_sum', 0):.0f}")
print(f" Total time: {stats.get('code2wav_time_s_sum', 0):.2f}s")
print(f" Avg TPS: {stats.get('code2wav_tps_avg', 0):.2f}")
print(f" Min TPS: {stats.get('code2wav_tps_min', 0):.2f}")
print(f" Max TPS: {stats.get('code2wav_tps_max', 0):.2f}")
print("\n--- Overall ---")
print(f" Total tokens: {stats.get('total_tokens_sum', 0):.0f}")
print(f" Total time: {stats.get('total_time_s_sum', 0):.2f}s")
print(f" Overall TPS: {stats.get('overall_tps', 0):.2f}")
print(f" Avg TPS: {stats.get('total_tps_avg', 0):.2f}")
print(f" Min TPS: {stats.get('total_tps_min', 0):.2f}")
print(f" Max TPS: {stats.get('total_tps_max', 0):.2f}")
print("=" * 60 + "\n")
def main():
parser = argparse.ArgumentParser(description="Qwen3-Omni Benchmark Script")
parser.add_argument(
"--prompts-file",
type=str,
default="benchmark/build_dataset/top100.txt",
help="Path to the prompts file (one prompt per line)",
)
parser.add_argument(
"--output-dir", type=str, default="benchmark_results", help="Directory to save benchmark results"
)
parser.add_argument("--model-path", type=str, default=MODEL_PATH, help="Path to the model")
parser.add_argument("--speaker", type=str, default="Ethan", help="Speaker voice for audio output")
parser.add_argument("--num-prompts", type=int, default=None, help="Number of prompts to process (default: all)")
args = parser.parse_args()
# Load model and processor
print(f"Loading model from {args.model_path}...")
model = Qwen3OmniMoeForConditionalGenerationWithLogging.from_pretrained(
args.model_path,
dtype="auto",
device_map="auto",
attn_implementation="flash_attention_2",
)
processor = Qwen3OmniMoeProcessor.from_pretrained(args.model_path)
# Benchmark mode
print(f"Loading prompts from {args.prompts_file}...")
prompts = load_prompts(args.prompts_file)
if args.num_prompts:
prompts = prompts[: args.num_prompts]
print(f"Running benchmark on {len(prompts)} prompts...")
aggregated_stats, results, audio_outputs = run_benchmark(
model=model,
processor=processor,
prompts=prompts,
output_dir=args.output_dir,
speaker=args.speaker,
)
print_stats(aggregated_stats)
print(f"\nResults saved to {args.output_dir}/")
if __name__ == "__main__":
main()
benchmarks/qwen3-omni/vllm-omni-vs-hf.png
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 115 KiB

benchmarks/qwen3-omni/vllm_omni/eval_qwen3_moe_omni.sh
-37
View File
@@ -1,37 +0,0 @@
#!/bin/bash
# Qwen3-Omni Benchmark Evaluation Script
# This script must be run from the vllm-omni root directory
# Get the directory where this script is located
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
# Navigate to vllm-omni root directory (4 levels up from script location)
VLLM_OMNI_ROOT="$(cd "$SCRIPT_DIR/../../.." && pwd)"
cd "$VLLM_OMNI_ROOT" || { echo "Error: Failed to navigate to vllm-omni directory"; exit 1; }
echo "Working directory: $(pwd)"
# Verify we're in the correct directory and run benchmark
if [[ ! -d "benchmarks/qwen3-omni/vllm_omni" ]]; then
echo "Error: Not in vllm-omni root directory. Please run from vllm-omni folder."
else
log_dir=benchmarks/qwen3-omni/vllm_omni/logs
outputs_dir=benchmarks/qwen3-omni/vllm_omni/outputs
end2end_script_path=examples/offline_inference/qwen3_omni/end2end.py
build_dataset_path=benchmarks/build_dataset/top100.txt
python $end2end_script_path --output-wav $outputs_dir \
--query-type text \
--txt-prompts $build_dataset_path \
--log-stats \
--log-dir $log_dir
echo "Logs and outputs are saved in ${log_dir} and ${outputs_dir} respectively:"
echo " - omni_pipeline_text run dir/base name"
echo " - omni_pipeline_text.orchestrator.stats.jsonl orchestrator-stage latency stats"
echo " - omni_pipeline_text.overall.stats.jsonl overall latency/TPS stats"
echo " - omni_pipeline_text.stage0.log per-stage detailed logs"
echo " - omni_pipeline_text.stage1.log"
echo " - omni_pipeline_text.stage2.log"
echo "Key checks: overall.stats.jsonl for end-to-end latency/TPS; orchestrator.stats.jsonl for stable per-stage latency; stage*.log for errors or long tails."
echo " - outputs/ Generated txt and wav files, there should be 100 text and wav files generated respectively"
fi
docs/.nav.yml
+3
View File
@@ -10,6 +10,7 @@ nav:
- Image Generation: serving/image_generation_api.md
- Image Edit: serving/image_edit_api.md
- Text to Speech: serving/speech_api.md
- Streaming Video Input: serving/video_stream_api.md
- Examples:
- examples/README.md
- Offline Inference:
@@ -34,6 +35,7 @@ nav:
- Online Serving:
- BAGEL-7B-MoT: user_guide/examples/online_serving/bagel.md
- vLLM-Omni Helm Chart: user_guide/examples/online_serving/chart-helm.md
- Diffusers Backend Adapter Example: user_guide/examples/online_serving/diffusers_pipeline_adapter.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
@@ -64,6 +66,7 @@ nav:
- FP8: user_guide/diffusion/quantization/fp8.md
- Int8: user_guide/diffusion/quantization/int8.md
- GGUF: user_guide/diffusion/quantization/gguf.md
- Attention Backends: user_guide/diffusion/attention_backends.md
- Frame Interpolation: user_guide/diffusion/frame_interpolation.md
- Parallelism:
- Overview: user_guide/diffusion/parallelism/overview.md
docs/cli/serve.md
+54
View File
@@ -1,5 +1,59 @@
# vllm-omni serve
## Stage-based CLI quickstart
The stage-based CLI is designed for deployments that require launching each pipeline stage in an isolated process
(e.g., across separate operating system processes, distinct GPUs, or distributed hosts).
- For **migrated models** that utilize the bundled deployment YAML configurations located in
`vllm_omni/deploy/`, the `--deploy-config` flag is only required to override the default configuration. By default, executing `vllm serve MODEL --omni ...`
automatically loads the bundled deployment configuration.
- For **legacy models** utilizing configuration files located in
`vllm_omni/model_executor/stage_configs/`, the `--stage-configs-path` parameter remains mandatory.
Example: Initializing Stage 0 (Orchestrator and API Server):
```bash
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
--port 8091 \
--stage-id 0 \
--omni-master-address 127.0.0.1 \
--omni-master-port 26000
```
Example: Initializing a Headless Worker Stage (Stage 1):
```bash
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
--stage-id 1 \
--headless \
--omni-master-address 127.0.0.1 \
--omni-master-port 26000
```
When utilizing a custom deployment YAML based on the new schema, append `--deploy-config /path/to/override.yaml` to each command execution. Conversely, for legacy models, substitute this parameter with `--stage-configs-path /path/to/stage_configs.yaml`.
In the standard execution paradigm, the `--stage-overrides` argument is utilized to apply stage-specific configurations from a single CLI command.
However, under the **stage-based CLI** paradigm, where each process strictly encapsulates a single stage, it is recommended to specify tuning parameters directly via discrete command-line flags for the respective stage, rather than constructing a composite `--stage-overrides` JSON string.
For example, as an alternative to the following composite configuration:
```bash
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
--stage-overrides '{"1": {"gpu_memory_utilization": 0.5}}'
```
the stage-based CLI permits the direct initialization of Stage 1 with explicit parameters:
```bash
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
--stage-id 1 \
--headless \
--gpu-memory-utilization 0.5 \
--omni-master-address 127.0.0.1 \
--omni-master-port 26000
```
## JSON CLI Arguments
--8<-- "docs/cli/json_tip.inc.md"
docs/configuration/stage_configs.md
+74 -3
View File
@@ -88,6 +88,55 @@ stages:
| `--async-chunk` / `--no-async-chunk` | Flip the deploy YAML's `async_chunk:` bool. Unset (default) leaves the YAML value in force. |
| `--stage-configs-path` | **Deprecated.** Accepts legacy `stage_args` yamls and (auto-detected) new deploy yamls; emits a deprecation warning. Migrate to `--deploy-config`. To be removed in a follow-up PR. |
### Stage-Based CLI Paradigm
The stage-based CLI paradigm facilitates the execution of discrete pipeline stages within isolated processes:
- **Stage 0** typically encapsulates the orchestrator and the primary API server. Invocation requires `--stage-id 0`,
`--omni-master-address`, `--omni-master-port`, and standard port declarations (e.g., `--port`).
- **Worker Stages** operate without a distinct API server (i.e., using `--headless`), are assigned sequential `--stage-id` identifiers, and must reference the corresponding
`--omni-master-address` and `--omni-master-port` parameters to successfully register with Stage 0.
For migrated architectures, the system automatically resolves and loads the bundled deployment YAML. Consequently, the primary execution path
does **not** necessitate the explicit definition of `--deploy-config`:
```bash
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
--port 8091 \
--stage-id 0 \
--omni-master-address 127.0.0.1 \
--omni-master-port 26000
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
--stage-id 1 \
--headless \
--omni-master-address 127.0.0.1 \
--omni-master-port 26000
```
When instantiating a custom deployment YAML conforming to the updated schema, append the `--deploy-config /path/to/override.yaml` directive
to all node invocations. For legacy architectures (e.g., BAGEL) configured via deprecated `stage_args:` schemas, continue to specify the relevant configuration via `--stage-configs-path /path/to/config.yaml`.
In the context of standard initialization architectures, utilizing the `--stage-overrides` parameter operates as the optimal methodology
for delineating stage-specific tuning from the CLI interface:
```bash
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
--stage-overrides '{"1": {"gpu_memory_utilization": 0.5}}'
```
Conversely, in the context of the **stage-based CLI** paradigm, given that each execution process exclusively instantiates a single pipeline stage, configuration override attributes
can be defined uniformly via explicit CLI flags on the corresponding instantiation command, rendering composite `--stage-overrides` JSON strings unnecessary:
```bash
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
--stage-id 1 \
--headless \
--gpu-memory-utilization 0.5 \
--omni-master-address 127.0.0.1 \
--omni-master-port 26000
```
### Precedence
From highest to lowest:
@@ -133,6 +182,17 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
--stage-overrides '{"0": {"max_num_seqs": 8}}'
```
Within the stage-based CLI paradigm, equivalent configuration parameters can inherently be passed directly
as command-line arguments to the designated single-stage process instantiation:
```bash
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
--stage-id 0 \
--max-num-seqs 8 \
--omni-master-address 127.0.0.1 \
--omni-master-port 26000
```
Effective config per stage after the merge:
| Stage | Field | Final value | Source |
@@ -153,9 +213,14 @@ Therefore, as a core part of vLLM-Omni, the stage configs for a model have sever
- Input and output dependencies for each stage.
- Default input parameters.
If users want to modify some part of it. The custom stage_configs file can be input as input argument in both online and offline. Just like examples below:
To override specific parameters, explicitly inject the customized configuration schema
in both online and offline instantiation flows. Prioritize the `--deploy-config` flag
when loading the new-schema deploy YAML schemas, reserving the `--stage-configs-path` parameter
exclusively to maintain compatibility with legacy `stage_args` YAML constructs.
For offline (Assume necessary dependencies have ben imported):
Examples:
For offline (Assume necessary dependencies have been imported):
```python
model_name = "Qwen/Qwen2.5-Omni-7B"
omni = Omni(model=model_name, stage_configs_path="/path/to/custom_stage_configs.yaml")
@@ -163,7 +228,13 @@ omni = Omni(model=model_name, stage_configs_path="/path/to/custom_stage_configs.
For online serving:
```bash
vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --deploy-config /path/to/deploy_config.yaml
```
Legacy online serving:
```bash
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
```
!!! important
We are actively iterating on the definition of stage configs, and we welcome all feedbacks from both community users and developers to help us shape the development!
docs/contributing/ci/CI_5levels.md
+12 -9
View File
@@ -105,7 +105,8 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
<td> Depends on reality </td>
<td>
<strong>Stability:</strong><br>
/tests/dfx/stability/tests/test.json<br>
/tests/dfx/stability/tests/test_qwen3_omni.json<br>
/tests/dfx/stability/tests/test_wan22.json<br>
<strong>Reliability:</strong><br>
tests/e2e/reliability/test_{model_name}.py
</td>
@@ -418,13 +419,13 @@ L3 level testing executes after code is merged into the main branch. Its core pu
**Explanation**:
@pytest.mark.advanced_model: Marks the test as L3 or L4 level, indicating that this test case performs deep validation, using real models for performance, integration, and accuracy testing. This forms a "basic-advanced" correspondence with the core_model mark at the L2 level.
@pytest.mark.advanced_model: Marks the test as L3 merge level, indicating deep validation with real models. @pytest.mark.full_model: Marks L4 nightly-only suites (e.g. `test_*_expansion.py`, doc examples).
@pytest.mark.core_model: Marks the test as L1 or L2 level, indicating that this test case validates the basic functionality of the core model. It uses mock weights and only checks if the relevant interface functions correctly.
@pytest.mark.parametrize: A parameterization decorator that allows abstracting test data into parameters, enabling reuse of the same test logic across different data configurations. indirect=True indicates that parameters will be passed to the fixture for processing.
**Notes**: If you believe the test case only needs to execute basic run logic at the PR-level CI, you can mark it only with @pytest.mark.core_model. If you believe it only needs to execute deep validation run logic at the merge or nightly level, you can mark it only with @pytest.mark.advanced_model. If you believe the test case needs to accommodate both basic run and deep validation test logic, you should mark it with both @pytest.mark.core_model and @pytest.mark.advanced_model.
**Notes**: If you believe the test case only needs to execute basic run logic at the PR-level CI, you can mark it only with @pytest.mark.core_model. If you believe it only needs to execute deep validation at merge (L3), use @pytest.mark.advanced_model. For L4 nightly-only expansion and doc-example tests, use @pytest.mark.full_model with `--run-level full_model`. If the test case needs both basic run and deep validation, mark with @pytest.mark.core_model and the appropriate L3/L4 marker (`advanced_model` and/or `full_model`).
**2.4.2 Test Function Definition and Documentation**
@@ -516,9 +517,11 @@ L3 level testing executes after code is merged into the main branch. Its core pu
**Single Request**: The comment clearly states this is a single-request completion test. For concurrent testing, it can be extended to multiple requests using request_num = n.
**Implicit Validation**: The `send_omni_request` and `send_diffusion_request` methods internally includes validation logic dynamically selected based on the --run-level parameter: core_model performs basic validation, while advanced_model performs deep validation.
**Implicit Validation**: The `send_omni_request` and `send_diffusion_request` methods internally includes validation logic dynamically selected based on the --run-level parameter: core_model performs basic validation, while advanced_model and full_model perform deep validation.
- ***Run Command***: `pytest -s -v /tests/e2e/online_serving/test_{model_name}.py -m advanced_model --run-level=advanced_model`
- ***Run Command (L3 merge)***: `pytest -s -v /tests/e2e/online_serving/test_{model_name}.py -m advanced_model --run-level=advanced_model`
- ***Run Command (L4 nightly expansion)***: `pytest -s -v /tests/e2e/online_serving/test_{model_name}_expansion.py -m full_model --run-level=full_model`
## Chapter 3: L4 Level Testing - Full Functionality, Performance, and Documentation Testing
@@ -572,12 +575,12 @@ L5 level testing focuses on the performance of model services under ***long-runn
### 4.2 Testing Content and Scope
- ***Long-term Stability (Stability) Testing***: Uses the `tests/dfx/stability/tests/test.json` configuration to run the service under moderate load for an extended period (e.g., over 12 hours), monitoring whether metrics like memory/VRAM usage, response time, and throughput degrade over time, and whether the service process remains stable.
- ***Long-term Stability (Stability) Testing***: Uses JSON under `tests/dfx/stability/tests/` (for example `test_qwen3_omni.json` and `test_wan22.json`) to run the service under moderate load for an extended period (e.g., over 12 hours), monitoring whether metrics like memory/VRAM usage, response time, and throughput degrade over time, and whether the service process remains stable.
- ***Reliability Testing***: Uses `tests/e2e/reliability/test_{model_name}.py` to actively simulate various fault and abnormal scenarios, such as: dependent service interruption, abnormal input data, network flicker, hardware resource preemption, etc., to verify the system's fault tolerance, self-healing, and graceful degradation capabilities.
### 4.3 Test Directory and Execution Files
- ***Stability Test Configuration***: `tests/dfx/stability/tests/test.json`
- ***Stability Test Configuration***: `tests/dfx/stability/tests/test_qwen3_omni.json`, `tests/dfx/stability/tests/test_wan22.json` (one JSON per model / runner family)
- ***Reliability Test Suite***: `tests/e2e/reliability/test_{model_name}.py`
### 4.4 Execution Method and Example
@@ -588,7 +591,7 @@ L5 level testing focuses on the performance of model services under ***long-runn
<details>
<summary> Test Examples</summary>
When you want to add L5-level stability test cases, you can refer to the following format for case addition in `tests/dfx/stability/tests/test.json`:
When you want to add L5-level stability test cases, add or extend the appropriate JSON file under `tests/dfx/stability/tests/` (for example `test_qwen3_omni.json` for Omni bench traffic, or `test_wan22.json` for diffusion `/v1/videos` workloads). The following illustrates the Qwen3-Omni shape:
```json
{
@@ -659,7 +662,7 @@ All other optional parameters follow the same rules as the in Chapter 3.4.
</details>
- - ***Stability***: `pytest -s -v tests/dfx/stability/scripts/test_{model_name}.py`
- - ***Stability***: `pytest -s -v tests/dfx/stability/scripts/test_stability_qwen3_omni.py` or `pytest -s -v tests/dfx/stability/scripts/test_stability_wan22.py` (or add `test_stability_<model>.py` alongside a matching JSON config)
- ***Reliability***: `pytest -s -v tests/e2e/reliability/test_{model_name}.py`
## Summary
docs/contributing/ci/test_examples/l4_functionality_tests.inc.md
+1 -1
View File
@@ -37,7 +37,7 @@ Currently all the features are available in online serving mode. Hence, only nee
**Code Style**
- Validation: test that the multimodal output files of your model have the correct shapes. `OpenAIClientHandler.send_diffusion_request` should have taken care of this.
- Test marks: always add `advanced_model` and `diffusion`. Add GPU-related marks if needed. Ref: [Markers for Tests](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/tests_markers/).
- Test marks: always add `full_model` and `diffusion` for L4 nightly `test_*_expansion.py` cases. Add GPU-related marks if needed. Ref: [Markers for Tests](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/tests_markers/).
- To maximize code reuse, you may refer to
- `tests/conftest.py` for `omni_server` (running server in subprocess) and `openai_client` fixtures (sending requests and validating output), `generate_synthetic_image` and `assert_XXX_valid` helper.
- `tests/helpers/mark.py` for `@hardware_test(...)` and `hardware_marks`.
docs/contributing/ci/test_guide.md
+40 -8
View File
@@ -42,31 +42,63 @@ Our test scripts use the pytest framework. First, please use `git clone https://
```
The latest test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-ready.yml).
=== "L3 level & L4 level"
=== "L3 level"
```bash
pytest -s -v -m "advanced_model" --run-level=advanced_model
```
If you only want to run L3 test case, you can use:
```bash
pytest -s -v e2e/ --ignore-glob='*expansion.py' -m "advanced_model" --run-level=advanced_model
```
If you only want to run a specific test case, you can use:
```bash
pytest -s -v test_xxxx.py --run-level=advanced_model
```
If you only want to run specific test cases on a particular platform, you can use:
```bash
pytest -s -v -m "core_model and distributed_cuda and L4" --run-level=core_model
pytest -s -v -m "advanced_model and distributed_cuda and L4" --run-level=advanced_model
```
The latest L3 test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-merge.yml).
=== "L4 level"
```bash
cd tests
pytest -s -v -m "full_model" --run-level=full_model
```
If you only want to run a specific test case, you can use:
```bash
pytest -s -v test_xxxx.py --run-level=full_model
```
If you only want to run specific test cases on a particular platform, you can use:
```bash
pytest -s -v -m "full_model and distributed_cuda and L4" --run-level=full_model
```
Note: To run performance tests (defaults to ``test_qwen_omni.json``; use ``--test-config-file tests/dfx/perf/tests/test_tts.json`` for TTS):
```bash
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py
```
The latest L4 (nightly) test commands use the `full_model` marker and `--run-level full_model` (see [test-nightly.yml](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-nightly.yml) and [test-nightly-diffusion.yml](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-nightly-diffusion.yml)). Example:
The latest L3 test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-merge.yml).
```bash
cd tests
pytest -s -v -m "full_model and omni and H100" --run-level=full_model
```
The latest L4 test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-nightly.yml).
=== "L5 level"
L5 includes stability and reliability testing. Typical commands:
```bash
cd tests
# Stability: Qwen3-Omni
pytest -s -v dfx/stability/scripts/test_stability_qwen3_omni.py
# Stability: Wan2.2 (v1/videos diffusion benchmark loop)
pytest -s -v dfx/stability/scripts/test_stability_wan22.py
```
The latest L5 commands for CI can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-ready.yml).
You can find more information about markers in the documentation: [marker doc](./tests_markers.md)
docs/contributing/ci/tests_markers.md
+2 -1
View File
@@ -8,7 +8,8 @@ Defined in `pyproject.toml`:
| Marker | Description |
| ------------------ | --------------------------------------------------------- |
| `core_model` | L1&L2 tests (run in each PR) |
| `advanced_model` | L3&L4 level tests (run in each merge or nightly) |
| `advanced_model` | L3 tests (run on each merge to main) |
| `full_model` | L4 tests (run nightly) |
| `diffusion` | Diffusion model tests |
| `omni` | Omni model tests |
| `cache` | Cache backend tests |
docs/contributing/model/adding_diffusion_model.md
+10 -8
View File
@@ -813,16 +813,16 @@ class WanTransformer3DModel(nn.Module):
---
### Diffusion Timing (Performance Profiling)
### Diffusion Pipeline Profiler (Performance Profiling)
When adapting a new diffusion model, it is often useful to analyze the latency of key components such as text encoding, diffusion denoising, and VAE decoding.
vLLM-Omni provides a timing utility via `DiffusionPipelineProfilerMixin` to help developers quickly identify performance bottlenecks.
!!! info
`DiffusionPipelineProfilerMixin` is different from using `torch.profiler` for diffusion models, as introduced in this [tutorial](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/profiling.md#3-profiling-diffusion-models). `DiffusionPipelineProfilerMixin` only prints the timing information of multiple functions (such as `vae.decode`), while `torch.profiler` saves detailed GPU/CPU computation time, call/execution steps.
`DiffusionPipelineProfilerMixin` is different from using `torch.profiler` for diffusion models, as introduced in this [tutorial](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/profiling.md). `DiffusionPipelineProfilerMixin` only prints the timing information of multiple functions (such as `vae.decode`), while `torch.profiler` saves detailed GPU/CPU computation time, call/execution steps.
This tool automatically measures the execution time of selected pipeline modules and prints the results in the logs.
**Enabling Diffusion Timing**
**Enabling Diffusion Pipeline Profiler**
Enable timing by setting:
@@ -843,7 +843,7 @@ If not specified, the default targets are used:
**Adding DiffusionPipelineProfilerMixin to a Pipeline**
To enable timing support in your pipeline, inherit from DiffusionPipelineProfilerMixin.
```python
from vllm_omni.diffusion.utils.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
from vllm_omni.diffusion.profiler import DiffusionPipelineProfilerMixin
class YourModelPipeline(nn.Module, DiffusionPipelineProfilerMixin):
# Optional: Specify custom timing targets
@@ -862,7 +862,9 @@ class YourModelPipeline(nn.Module, DiffusionPipelineProfilerMixin):
...
# initialize timing profiler
self.setup_diffusion_pipeline_profiler(enable_diffusion_pipeline_profiler)
self.setup_diffusion_pipeline_profiler(
enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler
)
```
The mixin dynamically wraps selected methods and records their execution time during inference.
@@ -906,9 +908,9 @@ tokenizer.forward
When enabled, timing logs appear like this:
```
[DiffusionTiming] text_encoder.forward took 0.018s
[DiffusionTiming] diffuse took 2.412s
[DiffusionTiming] vae.decode took 0.063s
[DiffusionPipelineProfiler] text_encoder.forward took 0.018s
[DiffusionPipelineProfiler] diffuse took 2.412s
[DiffusionPipelineProfiler] vae.decode took 0.063s
```
These measurements help identify bottlenecks during model adaptation and optimization
docs/contributing/model/adding_tts_model.md
+389 -16
View File
@@ -1,20 +1,93 @@
# Adding a TTS Model
This guide walks through adding a new TTS model to vLLM-Omni, using **Qwen3-TTS**
as a reference. Qwen3-TTS demonstrates the standard two-stage TTS pipeline and the
key optimizations all TTS models in this repo should follow.
This guide walks through adding a new TTS model to vLLM-Omni. Two patterns are
supported:
- **Two-stage pipeline** (e.g. Qwen3-TTS, Fish Speech): an AR code-predictor stage
feeds an audio decoder stage via the `async_chunk` framework. This is the standard
pattern for maximum streaming performance.
- **Single-stage AR model** (e.g. MOSS-TTS-Nano): the model runs entirely inside one
AR worker and streams audio chunks directly from its own `inference_stream()` generator.
Qwen3-TTS is used as the reference for the two-stage pattern. For the single-stage
pattern, refer to MOSS-TTS-Nano.
## Table of Contents
1. [Overview](#overview)
2. [Directory Structure](#directory-structure)
3. [Step-by-Step Implementation](#step-by-step-implementation)
4. [Key Components](#key-components)
5. [Model Registration](#model-registration)
6. [Stage Configuration](#stage-configuration)
7. [Stage Input Processors](#stage-input-processors)
8. [Testing](#testing)
9. [Summary](#summary)
2. [Cross-Cutting Invariants](#cross-cutting-invariants)
3. [Directory Structure](#directory-structure)
4. [Step-by-Step Implementation](#step-by-step-implementation)
5. [Key Components](#key-components)
6. [Model Registration](#model-registration)
7. [Stage Configuration](#stage-configuration)
8. [Stage Input Processors](#stage-input-processors)
9. [Online Serving Integration](#online-serving-integration)
10. [Single-Stage Models](#single-stage-models)
11. [Testing](#testing)
12. [Pre-commit and DCO](#pre-commit-and-dco)
13. [Summary](#summary)
## Cross-Cutting Invariants
These rules apply to every TTS model regardless of architecture (AR vs AR+diffusion,
single-stage vs two-stage, codec-based vs VAE-based). Each has surfaced as a silent
bug in a shipped PR — check them at the end of every phase, not just at the start.
**I1. Streaming output contract.** Pick one per-step semantics for `forward()` and
document it in the docstring:
- *Delta*: yield only new audio samples produced this step. Preferred — linear cost.
- *Cumulative*: re-decode from step 0 every call. O(N²); only acceptable when the
codec exposes no streaming decode.
If you choose delta, audit the full chain: `forward()` returns the new chunk →
`_consolidate_multimodal_tensors()` in `vllm_omni/engine/output_processor.py`
concatenates the audio key into a single tensor at finish → streaming consumers
receive per-step chunks, offline consumers receive the concatenated tensor. A
mismatch (consolidator skips the key with `continue`, or consumers expect a list
but receive a tensor) is invisible in offline RTF benchmarks — users hear replays
or truncation only under live playback.
**I2. Multimodal output consumer hygiene.** `outputs[0].outputs[0].multimodal_output[key]`
can be `Tensor`, `list[Tensor]` (pre-consolidation snapshot), `np.ndarray`, or
scalar. In every test, example, and benchmark:
- Never write `dict.get("a") or dict.get("b")` on tensor values — Python evaluates
the tensor's truthiness and raises `Boolean value of Tensor with more than one
value is ambiguous`. Use explicit `if x is None` chains.
- Defensively handle the list form:
`if isinstance(x, list): x = torch.cat([t.reshape(-1) for t in x], dim=0)`.
- Assert `shape` / `dtype` / `duration` explicitly — do not rely on truthiness for
presence checks.
**I3. Hot-loop GPU discipline.** Inside any per-step model loop (AR decode,
diffusion solver, CFM Euler step, per-frame vocoder):
- No `tensor.item()`, `.cpu()`, or `.tolist()` — each triggers a GPU→CPU sync; a
10-step × 60-frame × 4-op loop creates 2400 syncs per request.
- Prefer `dst.copy_(src)` over `dst.fill_(src.item())` for scalar-into-buffer writes.
- Whole-model `torch.compile(Model.forward, fullgraph=False)` usually outperforms
per-submodule compile — fewer dispatch boundaries, larger fusion regions. Measure
before choosing granularity.
- No Python control flow that depends on tensor values; use `torch.where` or masking.
Profile before optimizing.
**I4. Validation pyramid.** Offline RTF alone is necessary but not sufficient. A
new TTS model must pass all three levels:
| Layer | Catches | Tool |
|-------|---------|------|
| Offline RTF / duration | Throughput regressions, missing audio, wrong sample rate | `end2end.py`, pytest e2e |
| Browser streaming playback | Delta-vs-cumulative bugs, chunk boundary glitches, TTFP regressions | Gradio demo over `/v1/audio/speech?stream=true` |
| Concurrent requests | Per-request state leaks, codec window round-robin gaps | `max_num_seqs>1` smoke with 4+ parallel prompts |
**I5. Per-request state belongs to the request.** If the model caches anything
across `forward()` calls (streaming generators, codec buffers, sliding-window pads,
CUDA graph state), key it by `info.get("_omni_req_id")` and free the entry on
request finish. A shared buffer silently corrupts audio across concurrent requests —
the symptom is crosstalk or truncation under load, nothing in single-request tests.
## Overview
@@ -556,6 +629,302 @@ Recommended test cases for a new TTS model:
Reference test: `tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py`
### E2E Online Serving Tests (`tests/e2e/online_serving/test_<your_model>.py`)
The `omni_server` fixture in `tests/conftest.py` is **module-scoped**. Each distinct
`OmniServerParams` id in the same test file forces the fixture to tear the server
down and spawn a new one mid-module. A few rules that save real CI debugging time:
- **Prefer a single `OmniServerParams` set per file.** If you need to exercise two
deploy variants (e.g. `model.yaml` and `model_async_chunk.yaml`), either use one
variant and exercise streaming via request args, or split into two test files so
each file does exactly one server lifecycle. Mid-module teardown/restart is the
fragile path and surfaces startup races first.
- **Never depend on server-side fetching of external URLs** for reference audio or
other fixture data. CI runners (and China-hosted dev boxes) routinely fail on
SSL/DNS for public URLs. Inline the payload as a `data:audio/wav;base64,...`
ref_audio value — the serving layer accepts both forms.
- **Don't roll your own readiness probe.** The harness already waits for HTTP 200
on `/health` before releasing the server to the test. If your model needs extra
warmup signals, expose them through `/health` rather than adding `time.sleep(...)`
inside the test. (Bare TCP `connect_ex` probes were insufficient; see
`tests/conftest.py` `OmniServer.wait_for_ready`.)
- **Use `core_model` marker + H100 hardware_test** to match the `test-ready.yml`
pipeline so your test is picked up by the `ready` label, not only nightly.
## Online Serving Integration
To expose your model through the `/v1/audio/speech` OpenAI-compatible endpoint, add
**all five** of the following integration points to
`vllm_omni/entrypoints/openai/serving_speech.py` in a **single commit**. Adding them
piecemeal causes partial-integration failures that are hard to debug.
### 1. Stage constant
Near the top of the file, alongside the other `_*_TTS_MODEL_STAGES` constants:
```python
_YOUR_MODEL_TTS_MODEL_STAGES = {"your_model_stage_key"}
```
### 2. Union into `_TTS_MODEL_STAGES`
Add to the `_TTS_MODEL_STAGES` set union:
```python
_TTS_MODEL_STAGES: set[str] = (
...
| _YOUR_MODEL_TTS_MODEL_STAGES
)
```
### 3. Model type detection
In `_detect_tts_model_type()`, add before the final `return None`:
```python
if model_stage in _YOUR_MODEL_TTS_MODEL_STAGES:
return "your_model"
```
### 4. Request validation dispatch
In `_validate_tts_request()`, add before the fallback `return`:
```python
if self._tts_model_type == "your_model":
return self._validate_your_model_request(request)
```
### 5. Validation and parameter-builder methods
Add two new methods:
```python
def _validate_your_model_request(
self, request: OpenAICreateSpeechRequest
) -> str | None:
"""Validate YourModel request. Returns an error string or None."""
if not request.input or not request.input.strip():
return "Input text cannot be empty"
return None
def _build_your_model_params(
self, request: OpenAICreateSpeechRequest
) -> dict[str, Any]:
"""Build additional_information dict for YourModel."""
params: dict[str, Any] = {"text": [request.input]}
if request.voice is not None:
params["voice"] = [request.voice]
# Add any other model-specific fields here
return params
```
Then wire `_build_your_model_params` into the request-dispatch block in
`_create_tts_request()` (search for the equivalent `_build_*_params` call for an
existing model to find the right location). If the model supports voice cloning
(`ref_audio` → `prompt_audio_path`, `ref_text` → `prompt_text`), add those mappings
here too — follow any existing `_build_<model>_params` in `serving_speech.py` (e.g.
`_build_moss_tts_params` for the voice-cloning variant) for the pattern.
> **Two dispatch patterns coexist:** Fish Speech uses a `self._is_fish_speech` boolean
> checked *before* `elif self._is_tts`. All newer models use the `_tts_model_type`
> string pattern shown above. For new models, always use the string pattern — do not
> add new `_is_*` boolean flags.
> **Note on unused variables:** Only extract parameters in `_build_your_model_params`
> that you actually pass to the model's generate / `inference_stream` call. Extracting
> a variable without forwarding it will trigger a `ruff F841` pre-commit failure.
### Merge conflicts
`serving_speech.py` is modified by every new model PR and is the most common source of
rebase conflicts. When rebasing onto `main` and a conflict appears here, the resolution
is always to **keep both** the upstream model's additions and your own — never discard
either side. After resolving:
```bash
git add vllm_omni/entrypoints/openai/serving_speech.py
git rebase --continue
```
## Single-Stage Models
Some TTS models (e.g. MOSS-TTS-Nano) do not use a two-stage pipeline. Instead the
entire AR LM and audio decoder run inside a single AR worker, streaming audio chunks
directly from the model's own generator.
### Directory structure
```
vllm_omni/model_executor/models/your_model_name/
__init__.py
modeling_your_model_name.py # unified class: load_weights + forward + streaming
vllm_omni/model_executor/stage_configs/your_model_name.yaml
```
No stage input processor is needed.
### Stage config
Use a single stage with `worker_type: ar`. The `is_comprehension: true` field and the
top-level `async_chunk: false` are required — omitting them causes silent
misclassification in the serving layer. Set `max_num_seqs` to at least 4 for
concurrent production use.
```yaml
# stage_configs/your_model_name.yaml
async_chunk: false
stage_args:
- stage_id: 0
stage_type: llm
is_comprehension: true # required for serving_speech.py dispatch
runtime:
devices: "0"
engine_args:
model_stage: your_model_stage_key
model_arch: YourModelForCausalLM
worker_type: ar
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
engine_output_type: audio
max_num_seqs: 4 # min 4 for concurrent requests; default 1 causes gaps
final_output: true
final_output_type: audio
```
### Generator-based streaming pattern
This is the MOSS-TTS-Nano pattern, distinct from VoxCPM2's vLLM-native AR pattern
(see `plan/voxcpm2_native_ar_design.md` for that variant). Load model weights in
`load_weights()` (not `__init__`) so vLLM finishes distributed initialisation before
any CUDA allocations. Stream via a per-request generator stored in an instance dict:
```python
class YourModelForCausalLM(nn.Module):
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
super().__init__()
self._lm: nn.Module | None = None # populated in load_weights()
self._stream_gens: dict[str, Any] = {} # request_key → generator
def load_weights(self, weights):
# Load self._lm here, after vLLM distributed init
...
def forward(
self,
input_ids,
positions,
intermediate_tensors=None,
inputs_embeds=None,
runtime_additional_information: list[dict] | None = None, # one dict per request
**kwargs,
) -> OmniOutput:
infos = runtime_additional_information or [{}]
# Return empty output during dummy/profiling calls
if not runtime_additional_information or all(i.get("_is_dummy") for i in infos):
self._ar_emit_stop_token = True
return OmniOutput(...)
outputs, last_flags = [], []
for info in infos:
request_key = str(info.get("_omni_req_id", "0")) # set by vLLM, not user code
if request_key not in self._stream_gens:
self._stream_gens[request_key] = self._create_stream_gen(info)
try:
chunk, is_last = next(self._stream_gens[request_key])
except StopIteration:
chunk, is_last = torch.zeros(0), True
if is_last:
del self._stream_gens[request_key]
outputs.append(chunk)
last_flags.append(is_last)
self._ar_emit_stop_token = all(last_flags)
return OmniOutput(multimodal_outputs={"model_outputs": outputs, "is_last": last_flags})
def _create_stream_gen(self, info: dict):
"""Yield (waveform_tensor, is_last) from the model's inference_stream().
Handle both incremental ("audio" events) and batch ("result" event) models:
some upstream implementations emit one "result" event with the full waveform
instead of incremental "audio" events. Both paths must be covered.
"""
for event in self._lm.inference_stream(...):
if event["type"] == "audio":
yield event["waveform"], False
elif event["type"] == "result":
# Fallback for models that don't emit incremental audio events
yield event.get("waveform", torch.zeros(0)), True
return
yield torch.zeros(0), True
def compute_logits(self, hidden_states, sampling_metadata):
# Emit EOS only when the last chunk has been yielded so the AR
# scheduler ends the request at the right time.
...
```
For an in-tree reference, look for any single-stage AR model under
`vllm_omni/model_executor/models/` (for example
`moss_tts_nano/modeling_moss_tts_nano.py` once its integration has landed).
## Pre-commit and DCO
All contributions must pass the pre-commit checks and the Developer Certificate of
Origin (DCO) sign-off before merging.
### Running pre-commit
Install the hooks once with `pre-commit install`. Then run before committing:
```bash
pre-commit run --files \
vllm_omni/model_executor/models/your_model_name/*.py \
vllm_omni/entrypoints/openai/serving_speech.py \
vllm_omni/model_executor/models/registry.py \
tests/e2e/offline_inference/test_your_model_name.py \
tests/e2e/online_serving/test_your_model_name.py
```
When pre-commit **modifies files**, it exits with a non-zero code but the reformatting
is correct. Stage the modified files and commit again — do not revert the changes.
Common failures and fixes:
| Check | Cause | Fix |
|-------|-------|-----|
| `ruff F841` | Local variable assigned but never used | Remove the extraction or forward it to the model call |
| `ruff E402` | Module-level import not at top of file | Move import to the top-level import block |
| `ruff format` | Line length, spacing, or quote style | Accept the auto-fix, stage, and re-commit |
### DCO sign-off
Every commit must carry a `Signed-off-by` trailer. Use the `-s` flag when committing:
```bash
git commit -s -m "feat(your-model): add YourModel TTS support"
```
Or configure git to add it automatically:
```bash
git config format.signOff true
```
To fix a missing sign-off on the most recent commit:
```bash
git commit --amend -s --no-edit
git push origin your-branch --force-with-lease
```
> The DCO check verifies that the commit author email matches the `Signed-off-by` email.
> Make sure `git config user.email` is set to the address associated with your GitHub
> account before committing.
## Adding a Model Recipe
After implementing and testing your model, add a model recipe to the
@@ -567,15 +936,19 @@ for the expected format.
Adding a TTS model to vLLM-Omni involves:
1. **Create model directory** with AR stage, decoder stage, and unified class
1. **Create model directory** with AR stage, decoder stage, and unified class (two-stage)
or a single unified class with generator-based streaming (single-stage)
2. **AR stage** - use vLLM's native decoder layers with fused QKV; do not wrap HF directly
3. **Decoder stage** - thin wrapper around your audio decoder; implement `chunked_decode_streaming()`
4. **Unified class** - dispatches on `model_stage`; same structure as `Qwen3TTSModelForGeneration`
5. **Register** all stage classes in `registry.py`
6. **YAML configs** - provide both batch and `async_chunk` variants
7. **Stage input processor** - buffer Stage 0 outputs and forward in chunks of 25
8. **Tests** - cover single request, batching, and async_chunk streaming
9. **Model recipe** - add to [vllm-project/recipes](https://github.com/vllm-project/recipes)
6. **YAML configs** - provide both batch and `async_chunk` variants (two-stage), or a single-stage AR config
7. **Stage input processor** - buffer Stage 0 outputs and forward in chunks of 25 (two-stage only)
8. **Online serving** - add all 5 integration points to `serving_speech.py` in one commit
9. **Tests** - cover single request, batching, and streaming
10. **Pre-commit + DCO** - run `pre-commit` before pushing; sign every commit with `git commit -s`
11. **Model recipe** - add to [vllm-project/recipes](https://github.com/vllm-project/recipes)
12. **Invariants** - re-check I1–I5 (streaming contract, consumer hygiene, hot-loop discipline, validation pyramid, per-request state) at the end of every phase
### Qwen3-TTS Reference Files
docs/contributing/profiling.md
+147 -75
View File
@@ -1,15 +1,17 @@
# Profiling vLLM-Omni
# Profiling Diffusion Models
> **Warning:** Profiling is for development and debugging only. It adds significant overhead and should not be enabled in production.
vLLM-Omni supports two profiler backends through `profiler_config`:
Diffusion profiling supports two backends through `profiler_config`:
- `torch`: detailed CPU/CUDA traces written to `torch_profiler_dir`
- `torch`: detailed CPU/CUDA traces, operator tables, and optional memory snapshots
- `cuda`: low-overhead CUDA range control for NVIDIA Nsight Systems (`nsys`)
## 1. Configure Profiling
## 1. Configure `profiler_config`
Use the same `profiler_config` shape everywhere:
Use `profiler_config` to enable profiling for a diffusion model. For diffusion usage, pass it directly to `Omni(...)` or `vllm serve`.
Minimal torch-profiler config:
```yaml
profiler_config:
@@ -21,54 +23,114 @@ Supported fields:
| Field | Description |
|---|---|
| `profiler` | Profiler backend. Supported values: `torch`, `cuda`. |
| `torch_profiler_dir` | Output directory for torch traces. Required when `profiler: torch`. |
| `profiler` | Profiler backend. Supported values: `torch`, `cuda`. Use `torch` for `trace.json`, Excel operator tables, and optional memory snapshots. Use `cuda` for Nsight Systems only. |
| `torch_profiler_dir` | Output directory for torch-profiler artifacts. Required when `profiler: torch`. |
| `torch_profiler_use_gzip` | Compress `trace_rank*.json` into `trace_rank*.json.gz`. |
| `torch_profiler_record_shapes` | Record input shapes and add a `by_shape` sheet to `ops_rank*.xlsx`. |
| `torch_profiler_with_stack` | Record call stacks, add a `by_stack` sheet to `ops_rank*.xlsx`, and export `stacks_cpu_rank*.txt` and `stacks_cuda_rank*.txt`. |
| `torch_profiler_with_memory` | Enable memory profiling and attempt to dump `memory_snapshot_rank*.pickle`. The pickle is only generated when the current backend supports memory history and snapshot APIs. |
| `torch_profiler_with_flops` | Enable FLOPs collection in `torch.profiler`. This does not add a separate output file. |
| `torch_profiler_dump_cuda_time_total` | Export an additional text summary `profiler_out_<rank>.txt` sorted by `self_cuda_time_total`. |
| `delay_iterations` | Number of worker iterations to skip before profiling starts. |
| `max_iterations` | Maximum number of worker iterations to capture before auto-stop. |
| `wait_iterations` | Torch-profiler wait iterations before warmup. |
| `warmup_iterations` | Torch-profiler warmup iterations. |
| `active_iterations` | Torch-profiler active iterations. |
| `wait_iterations` | Torch-profiler wait iterations before warmup. |
For multi-stage omni pipelines, put `profiler_config` under the target stage's `engine_args`.
### Minimal configurations by output
```yaml
stage_args:
- stage_id: 0
stage_type: llm
engine_args:
profiler_config:
profiler: torch
torch_profiler_dir: ./perf
```
For single-stage diffusion usage, pass `profiler_config` directly to `Omni(...)` or `vllm serve`.
## 2. Profiling Omni Pipelines
It is usually best to profile only the stages you need.
Only collect trace output:
```python
# Profile all stages.
omni.start_profile()
# Profile selected stages only.
omni.start_profile(stages=[0, 2])
...
omni.stop_profile(stages=[0, 2])
profiler_config = {
"profiler": "torch",
"torch_profiler_dir": "./perf",
}
```
Always stop the same stage set that you started. If only some stages have `profiler_config`, pass an explicit `stages=[...]` list instead of relying on the default "all stages" behavior.
Outputs:
Examples:
- `trace_rank*.json`
- `ops_rank*.xlsx` with a `summary` sheet
1. [Qwen2.5-Omni end2end](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen2_5_omni/end2end.py)
2. [Qwen3-Omni end2end](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen3_omni/end2end.py)
Collect compressed trace output:
## 3. Profiling Single-Stage Diffusion
```python
profiler_config = {
"profiler": "torch",
"torch_profiler_dir": "./perf",
"torch_profiler_use_gzip": True,
}
```
Single-stage diffusion models use the same `start_profile()` / `stop_profile()` controls, but you must provide `profiler_config` explicitly.
Outputs:
### PyTorch profiler
- `trace_rank*.json.gz`
- `ops_rank*.xlsx` with a `summary` sheet
Collect trace and full operator tables:
```python
profiler_config = {
"profiler": "torch",
"torch_profiler_dir": "./perf",
"torch_profiler_record_shapes": True,
"torch_profiler_with_stack": True,
}
```
Outputs:
- `trace_rank*.json`
- `ops_rank*.xlsx` with `summary`, `by_shape`, and `by_stack`
- `stacks_cpu_rank*.txt`
- `stacks_cuda_rank*.txt`
Collect trace, operator tables, and memory snapshots:
```python
profiler_config = {
"profiler": "torch",
"torch_profiler_dir": "./perf",
"torch_profiler_record_shapes": True,
"torch_profiler_with_stack": True,
"torch_profiler_with_memory": True,
}
```
Outputs:
- `trace_rank*.json`
- `ops_rank*.xlsx` with `summary`, `by_shape`, and `by_stack`
- `stacks_cpu_rank*.txt`
- `stacks_cuda_rank*.txt`
- `memory_snapshot_rank*.pickle` when supported by the current backend
### Full torch-profiler configuration
If you want to enable the commonly used torch-profiler options together:
```python
profiler_config = {
"profiler": "torch",
"torch_profiler_dir": "./perf",
"torch_profiler_use_gzip": False,
"torch_profiler_record_shapes": True,
"torch_profiler_with_stack": True,
"torch_profiler_with_memory": True,
"torch_profiler_with_flops": False,
"torch_profiler_dump_cuda_time_total": False,
"delay_iterations": 0,
"max_iterations": 0,
"wait_iterations": 0,
"warmup_iterations": 0,
"active_iterations": 0,
}
```
## 2. Profiling Diffusion with PyTorch Profiler
Single-stage diffusion models use `start_profile()` / `stop_profile()` controls. The profiler only writes artifacts after profiling has been started and then stopped.
```python
from vllm_omni import Omni
@@ -86,7 +148,29 @@ omni.start_profile()
omni.stop_profile()
```
### Nsight Systems (`nsys`)
For diffusion offline example scripts under `examples/offline_inference/`, pass `--profiler-config` as a JSON object. The script enables profiling when this argument is set and wraps generation with `start_profile()` / `stop_profile()`.
Example:
```bash
python examples/offline_inference/image_to_video/image_to_video.py \
--model Wan-AI/Wan2.2-I2V-A14B-Diffusers \
--image input.jpg \
--prompt "A cat playing with yarn" \
--profiler-config '{
"profiler": "torch",
"torch_profiler_dir": "./perf",
"torch_profiler_record_shapes": true,
"torch_profiler_with_stack": true
}'
```
Examples:
1. [Image edit example](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/image_to_image/image_edit.py)
2. [Image to video example](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_video)
## 3. Profiling Diffusion with Nsight Systems (`nsys`)
For Nsight Systems, use `profiler: cuda` and wrap the process with `nsys profile`.
@@ -103,35 +187,20 @@ nsys profile \
The Python process being profiled must create the diffusion engine with:
```python
profiler_config={"profiler": "cuda"}
profiler_config = {"profiler": "cuda"}
```
Then call `start_profile()` before the requests you want to capture and `stop_profile()` after them. The diffusion worker processes open and close the CUDA capture range themselves, so `nsys` sees the actual GPU work instead of only the parent process.
Examples:
1. [Image edit example](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/image_to_image/image_edit.py)
2. [Image to video example](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_video)
## 4. Profiling Online Serving
When any stage has `profiler_config.profiler` set, the server exposes:
When `profiler_config.profiler` is set for a diffusion model, the server exposes:
- `POST /start_profile`
- `POST /stop_profile`
### Start the server
Multi-stage omni serving:
```bash
vllm serve Qwen/Qwen2.5-Omni-7B \
--omni \
--port 8091
```
(The default deploy config at `vllm_omni/deploy/qwen2_5_omni.yaml` is loaded automatically. Pass `--deploy-config /path/to/custom.yaml` to override.)
Single-stage diffusion serving with torch profiler:
```bash
@@ -144,21 +213,17 @@ vllm serve Wan-AI/Wan2.2-I2V-A14B-Diffusers \
"torch_profiler_with_stack": true,
"torch_profiler_with_flops": false,
"torch_profiler_use_gzip": true,
"torch_profiler_dump_cuda_time_total": true,
"torch_profiler_dump_cuda_time_total": false,
"torch_profiler_record_shapes": true,
"torch_profiler_with_memory": true,
"delay_iterations": 0,
"max_iterations": 0
"max_iterations": 0,
"wait_iterations": 0,
"warmup_iterations": 0,
"active_iterations": 0
}'
```
Useful optional fields:
- `torch_profiler_with_stack`: export `by_stack` operator views and stack text files.
- `torch_profiler_record_shapes`: export `by_shape` operator views.
- `torch_profiler_with_memory`: dump `memory_snapshot_rank*.pickle` when the backend supports memory history.
- `torch_profiler_use_gzip`: write the trace as `trace_rank*.json.gz`.
Single-stage diffusion serving with Nsight Systems:
```bash
@@ -176,30 +241,37 @@ nsys profile \
### Control capture
Example profiling flow for an online Qwen-Image request:
```bash
# Start profiling on all profiled stages.
# Start profiling.
curl -X POST http://localhost:8091/start_profile
# Start profiling on selected stages.
curl -X POST http://localhost:8091/start_profile \
# Send a Qwen-Image generation request while profiling is active.
curl http://localhost:8091/v1/images/generations \
-H "Content-Type: application/json" \
-d '{"stages": [0]}'
-d '{
"model": "Qwen/Qwen-Image",
"prompt": "A red vintage bicycle parked beside a quiet canal at sunset"
}'
# Stop profiling.
# Stop profiling and flush profiler artifacts.
curl -X POST http://localhost:8091/stop_profile
```
For mixed-stage pipelines, use explicit `stages` and pass the same stage list to both endpoints.
## 5. Diffusion Pipeline Profiler
## 5. Analyze Results
For lightweight per-stage pipeline timing such as `vae.decode` or `diffuse`, see [Diffusion Pipeline Profiler](model/adding_diffusion_model.md#diffusion-pipeline-profiler-performance-profiling). That utility logs stage durations only and does not generate torch-profiler artifacts such as `trace.json`, Excel tables, or memory snapshots.
Torch profiler output:
## 6. Analyze Results
Torch-profiler output:
- Chrome/Perfetto trace: `trace_rank*.json` or `trace_rank*.json.gz`
- Excel workbook: `ops_rank*.xlsx` with `summary`, and optional `by_shape` / `by_stack` sheets
- Stack exports: `stacks_cpu_rank*.txt` and `stacks_cuda_rank*.txt` when stack capture is enabled
- Memory snapshot: `memory_snapshot_rank*.pickle` when memory capture is enabled and supported by the backend
- Optional aggregated CUDA-time tables under the same session directory
- Optional CUDA-time text summary: `profiler_out_<rank>.txt` when `torch_profiler_dump_cuda_time_total` is enabled
CUDA profiler / Nsight Systems output:
@@ -211,4 +283,4 @@ Recommended viewers:
- `nsys stats <report>.nsys-rep` for CLI summaries
- Nsight Systems GUI for CUDA kernel timelines
vLLM-Omni reuses the vLLM profiling infrastructure where possible. For the upstream reference, see the [vLLM profiling guide](https://docs.vllm.ai/en/stable/contributing/profiling/).
For upstream background on the underlying vLLM profiling infrastructure, see the [vLLM profiling guide](https://docs.vllm.ai/en/stable/contributing/profiling/).
docs/mkdocs/hooks/generate_argparse.py
+1
View File
@@ -121,6 +121,7 @@ def extract_omni_serve_subparser_init():
"FlexibleArgumentParser": _FlexibleArgumentParser,
"make_arg_parser": lambda parser: parser, # no-op for doc
"_ensure_vllm_platform": lambda: None, # no-op for doc
"nullify_stage_engine_defaults": lambda parser: None, # no-op for doc
"VLLM_SUBCMD_PARSER_EPILOG": "",
"logger": logger,
"DummySubparsers": DummySubparsers,
docs/models/supported_models.md
+4 -1
View File
@@ -36,6 +36,8 @@ th {
| `LTX2ImageToVideoPipeline` | LTX-2-I2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
| `LTX2TwoStagesPipeline` | LTX-2-T2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
| `LTX2ImageToVideoTwoStagesPipeline` | LTX-2-I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
| `LTX23Pipeline` | LTX-2.3-T2V | `dg845/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
| `LTX23ImageToVideoPipeline` | LTX-2.3-I2V | `dg845/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | |
| `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | |
| `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ |
@@ -54,7 +56,8 @@ th {
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` | ✅︎ | ✅︎ | | ✅︎ |
| `MiMoAudioForConditionalGeneration` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` | ✅︎ | ✅︎ | | |
| `MiMoAudioModel` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` | ✅︎ | ✅︎ | | |
| `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | |
| `Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` | ✅︎ | ✅︎ | | |
| `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | |
| `DreamIDOmniPipeline` | DreamID-Omni | `XuGuo699/DreamID-Omni` | ✅︎ | ✅︎ | | |
docs/serving/video_stream_api.md
+93
View File
@@ -0,0 +1,93 @@
# Streaming Video Input API
vLLM-Omni provides a WebSocket API for streaming video frames and optional audio chunks into Qwen3-Omni, then asking questions over the buffered session context.
Each server instance runs a single model specified at startup with `vllm serve <model> --omni`.
## Quick Start
### Start the Server
```bash
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct \
--deploy-config vllm_omni/deploy/qwen3_omni.yaml \
--omni \
--port 8000 \
--trust-remote-code
```
### Run the Example Client
```bash
python examples/online_serving/qwen3_omni/streaming_video_client.py \
--url ws://localhost:8000/v1/video/chat/stream \
--video /path/to/video.mp4 \
--query "Describe what is happening in the video."
```
## API Reference
### Endpoint
```text
WebSocket /v1/video/chat/stream
```
### Protocol
| Direction | Type | Required fields | Description |
|-----------|------|-----------------|-------------|
| Client -> Server | `session.config` | none | First message. Configures output modalities, frame sampling, EVS, and prompts. |
| Client -> Server | `video.frame` | `data` | Base64 JPEG/PNG frame. |
| Client -> Server | `audio.chunk` | `data` | Base64 PCM16 16 kHz mono audio bytes. |
| Client -> Server | `video.query` | `text` | Ask a question over the buffered frames and audio. |
| Client -> Server | `video.done` | none | End the WebSocket session. |
| Server -> Client | `response.start` | none | Query generation started. |
| Server -> Client | `response.text.delta` | `delta` | Incremental text output. |
| Server -> Client | `response.text.done` | `text` | Final text output for the query. |
| Server -> Client | `response.audio.delta` | `data`, `format` | Incremental generated audio, base64 WAV. |
| Server -> Client | `response.audio.done` | none | Audio output finished. |
| Server -> Client | `session.done` | none | Session closed. |
| Server -> Client | `error` | `message` | Recoverable protocol or generation error. |
### `session.config` Fields
| Field | Type | Default | Description |
|-------|------|---------|-------------|
| `model` | string or null | null | Optional model name. Usually omitted because the server hosts one model. |
| `modalities` | list[string] | `["text", "audio"]` | Output modalities. Use `["text"]`, `["audio"]`, or both. |
| `num_frames` | integer, 1-128 | `4` | Number of buffered frames sampled for each query. |
| `max_frames` | integer, 1-256 | `50` | Maximum retained frame buffer size. Oldest frames are evicted first. |
| `system_prompt` | string or null | null | Optional custom system prompt. |
| `use_audio_in_video` | bool | `true` | Include streamed audio chunks in multimodal video understanding when audio is present. |
| `sampling_params_list` | list or null | null | Optional per-stage sampling parameter overrides. |
| `enable_frame_filter` | bool | `true` | Enable EVS near-duplicate frame filtering. |
| `frame_filter_threshold` | float, 0.0-1.0 | `0.95` | EVS similarity threshold. Higher keeps more frames; lower drops more near-duplicates. |
### Legacy Aliases
The server accepts these legacy field names and rewrites them before validation. New clients should send the canonical names above.
| Legacy field | Canonical field |
|--------------|-----------------|
| `num_sample_frames` | `num_frames` |
| `evs_enabled` | `enable_frame_filter` |
| `evs_threshold` | `frame_filter_threshold` |
### Environment Variables
| Variable | Values | Default | Description |
|----------|--------|---------|-------------|
| `VLLM_VIDEO_ASYNC_CHUNK` | `on`, `off` | `on` | Wire-level streaming switch. `off` buffers server-side deltas and emits coalesced outputs at the end of a query. |
| `VLLM_VIDEO_AUDIO_DELTA_MODE` | `fast`, `slow` | `fast` | Audio delta extraction strategy. `fast` emits only newly produced chunks; `slow` recomputes from accumulated audio and exists for A/B verification. |
## EVS Semantics
EVS compares downsampled frames and drops near-duplicate frames before they enter the session frame buffer. `frame_filter_threshold` controls retention: higher values are more permissive and keep more frames; lower values are more aggressive and drop more similar frames.
## Known Limitations
- Session KV reuse and incremental prefill are not implemented in this PR. Each `video.query` rebuilds the model prompt from the retained frame and audio buffers.
- Back-to-back short replies can still expose an engine-layer scheduler race. The PR notes an observed workaround of at least 200 ms idle between turns when clients repeatedly see idle timeouts.
- If the audio buffer exceeds the server limit, the server emits `Audio buffer overflow` and clears the currently buffered audio for the session.
- The API is intended for Qwen3-Omni streaming video understanding; other models may not support the same multimodal processor arguments.
docs/user_guide/diffusion/attention_backends.md
+120
View File
@@ -0,0 +1,120 @@
# Diffusion Attention Backends
This document describes the diffusion attention backends available in vLLM-Omni, how to select them, and how to use SageAttention.
## Overview
Diffusion attention backend selection is controlled by the `DIFFUSION_ATTENTION_BACKEND` environment variable and resolved in `vllm_omni.diffusion.attention.selector`.
This backend is used by diffusion attention layers such as the DiT attention in video and image generation models.
On CUDA, the practical choices today are:
- `FLASH_ATTN`: FlashAttention backend. This is the default on supported CUDA systems when FlashAttention is installed.
- `TORCH_SDPA`: PyTorch `scaled_dot_product_attention`.
- `SAGE_ATTN`: SageAttention backend, if `sageattention` is installed.
If `DIFFUSION_ATTENTION_BACKEND` is unset, vLLM-Omni asks the current platform to choose the default backend. On CUDA, that normally means `FLASH_ATTN` when available, otherwise `TORCH_SDPA`.
## Backend Options
| Value | Notes |
|---|---|
| `FLASH_ATTN` | Default on CUDA when FlashAttention is available. Good default for most diffusion workloads. |
| `TORCH_SDPA` | Most conservative fallback. Useful for debugging or compatibility. |
| `SAGE_ATTN` | Requires `sageattention`. Can improve performance on some workloads, but output quality must be validated model-by-model. |
## Selection Priority
Diffusion attention backend selection follows this order:
1. `DIFFUSION_ATTENTION_BACKEND`
2. Platform default
Example:
```bash
export DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN
```
## SageAttention Installation
vLLM-Omni expects SageAttention to be installed into the same Python environment as vLLM-Omni.
Build from source:
```bash
git clone https://github.com/thu-ml/SageAttention.git
cd SageAttention
export EXT_PARALLEL=4 NVCC_APPEND_FLAGS="--threads 8" MAX_JOBS=32
pip install . --no-build-isolation
```
Quick check:
```bash
python -c "import sageattention; print(sageattention.__file__)"
```
## Usage
### Enable SageAttention
Example: HunyuanVideo-1.5 text-to-video
```bash
DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
--model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \
--prompt "A dog running across a field of golden wheat." \
--height 480 --width 832 --num-frames 33 \
--num-inference-steps 30 --seed 42 --guidance-scale 6.0 \
--tensor-parallel-size 2 \
--output ../tmp/hv15_modelopt_sage.mp4
```
Example: Wan2.2 TI2V 5B
```bash
DIFFUSION_ATTENTION_BACKEND=SAGE_ATTN python examples/offline_inference/text_to_video/text_to_video.py \
--model Wan-AI/Wan2.2-TI2V-5B-Diffusers \
--prompt "A dog running across a field of golden wheat." \
--height 704 --width 1280 --num-frames 49 \
--num-inference-steps 30 --seed 42 --guidance-scale 5.0 \
--tensor-parallel-size 2 \
--output outputs/wan22_sage.mp4
```
### Compare Against FlashAttention
Unset the backend override, or explicitly use `FLASH_ATTN`:
```bash
python examples/offline_inference/text_to_video/text_to_video.py \
--model Wan-AI/Wan2.2-TI2V-5B-Diffusers \
--prompt "A dog running across a field of golden wheat." \
--height 704 --width 1280 --num-frames 49 \
--num-inference-steps 30 --seed 42 --guidance-scale 5.0 \
--tensor-parallel-size 2 \
--output outputs/wan22_fa3.mp4
```
## Validation Guidance
Do not assume that a faster attention backend is numerically interchangeable with `FLASH_ATTN`.
Always compare:
- End-to-end runtime
- DiT / diffusion stage runtime
- Output quality against a known-good baseline
At minimum, keep the same:
- model
- prompt
- seed
- resolution
- frame count
- inference steps
- parallel config
docs/user_guide/diffusion/cpu_offload_diffusion.md
+51 -4
View File
@@ -36,6 +36,45 @@ m = Omni(model="Wan-AI/Wan2.2-T2V-A14B-Diffusers", enable_cpu_offload=True)
vllm-omni serve diffusion Wan-AI/Wan2.2-T2V-A14B-Diffusers --enable-cpu-offload
```
### To Support a Model
Implement the `SupportsModuleOffload` protocol to declare which
submodules participate in offloading:
```python
from typing import ClassVar
from vllm_omni.diffusion.models.interface import SupportsModuleOffload
class MyPipeline(nn.Module, SupportsModuleOffload):
_dit_modules: ClassVar[list[str]] = ["transformer"]
_encoder_modules: ClassVar[list[str]] = ["text_encoder", "vision_model"]
_vae_modules: ClassVar[list[str]] = ["vae"]
_resident_modules: ClassVar[list[str]] = [] # optional
def __init__(self):
super().__init__()
self.transformer = ... # DiT — stays on GPU during denoising
self.text_encoder = ... # Encoder — offloaded to CPU during denoising
self.vision_model = ... # Encoder — offloaded to CPU during denoising
self.vae = ... # VAE — always on GPU
```
- `_dit_modules`: attribute names of denoising submodules (kept on GPU
during the diffusion loop).
- `_encoder_modules`: attribute names of encoder/vision submodules
(offloaded to CPU during the diffusion loop).
- `_vae_modules`: attribute names of VAE(s) (always kept on GPU, not
part of the mutual exclusion hooks).
- `_resident_modules`: attribute names of small submodules that must
stay on GPU during layerwise offloading (e.g. embedders, connectors).
Optional — defaults to `[]`.
All attribute names support dotted paths for nested submodules
(e.g. `"pipe.transformer"`, `"bagel.time_embedder"`).
Both DiT and encoder lists are needed because the offload hooks use
mutual exclusion: when one group runs, the other moves to CPU.
### Limitations
- Cold start latency increases
- Adds overhead from CPU-GPU transfers between encoder and denoising phases
@@ -116,11 +155,19 @@ class Flux2Transformer2DModel(nn.Module):
**Module Discovery**
The offloader automatically discovers pipeline components:
The offloader discovers pipeline components in two ways:
- **DiT modules**: `transformer`, `transformer_2`, `dit`
- **Encoders**: `text_encoder`, `text_encoder_2`, `text_encoder_3`, `image_encoder`
- **VAE**: `vae`
1. **Protocol-based** (preferred): If the pipeline implements
`SupportsModuleOffload`, its `_dit_modules`, `_encoder_modules`,
`_vae_modules`, and `_resident_modules` class variables are used
directly. All attribute names support dotted paths (e.g.
`"pipe.transformer"`, `"bagel.time_embedder"`) for nested submodules.
2. **Fallback attribute scan**: Otherwise, the offloader scans for
well-known attribute names:
- **DiT modules**: `transformer`, `transformer_2`, `dit`, `sr_dit`, `language_model`, `transformer_blocks`, `model`
- **Encoders**: `text_encoder`, `text_encoder_2`, `text_encoder_3`, `image_encoder`
- **VAE**: `vae`, `audio_vae`
**Hook System**
docs/user_guide/diffusion/quantization/autoround.md
+2
View File
@@ -72,6 +72,8 @@ At load time:
| Model | HF Checkpoint | Scheme | Group Size | Backend |
|-------|--------------|--------|------------|---------|
| FLUX.1-dev | `vllm-project-org/FLUX.1-dev-AutoRound-w4a16` | W4A16 | 128 | GPTQ-Marlin |
| Qwen2.5-Omni-7B | `Intel/Qwen2.5-Omni-7B-int4-AutoRound` | W4A16 | 128 | GPTQ-Marlin |
| Qwen3-Omni-30B-A3B-Instruct | `Intel/Qwen3-Omni-30B-A3B-Instruct-int4-AutoRound` | W4A16 | 128 | GPTQ-Marlin |
## Creating a Quantized Checkpoint
docs/user_guide/diffusion/quantization/fp8.md
+1
View File
@@ -65,6 +65,7 @@ The available `ignored_layers` names depend on the model architecture (e.g., `to
| Flux | `black-forest-labs/FLUX.1-dev` | All layers | None |
| HunyuanImage-3 | `tencent/HunyuanImage3` | All layers | None |
| HunyuanVideo-1.5 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `720p_t2v`, `480p_i2v` | All layers | None |
| GLM-Image | `zai-org/GLM-Image` | All layers | None |
| Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | All layers | None |
## Combining with Other Features
docs/user_guide/diffusion_features.md
+2 -1
View File
@@ -141,6 +141,7 @@ The following tables show which models support each feature:
| **Wan2.2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (encode/decode) | ❌ | ❌ |
| **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
| **LTX-2.3** | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **Helios** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ |
| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
@@ -148,7 +149,7 @@ The following tables show which models support each feature:
**Frame Interpolation Support**
- **Supported**: Wan2.2 text-to-video, image-to-video, and TI2V pipelines
- **Not supported**: Wan2.1-VACE, LTX-2, Helios, HunyuanVideo-1.5, DreamID-Omni
- **Not supported**: Wan2.1-VACE, LTX-2, LTX-2.3, Helios, HunyuanVideo-1.5, DreamID-Omni
### AudioGen
docs/user_guide/examples/offline_inference/bagel.md
+208 -185
View File
@@ -2,46 +2,61 @@
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/bagel>.
## Set up
## Setup
Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
## Run examples
## Architecture
**Note**: These examples work with the default configuration on an **NVIDIA A100 (80GB)**. We also tested on dual **NVIDIA RTX 5000 Ada (32GB each)**. For dual-GPU setups, please modify the stage configuration to distribute the model across devices.
BAGEL-7B-MoT is a Mixture-of-Transformers (MoT) model supporting both image generation and understanding. It offers two deployment topologies:
Get into the bagel folder
| Topology | Stages | Description |
| :------- | :----- | :---------- |
| **Two-stage** (default) | Stage 0 (Thinker, AR) + Stage 1 (DiT, Diffusion) | Thinker handles text/understanding via vLLM AR engine; DiT handles image generation. KV cache is transferred between stages. |
| **Single-stage** | Stage 0 (DiT, Diffusion) only | The DiT stage contains a full LLM, ViT, VAE, and tokenizer internally. All modalities are handled within a single diffusion process. |
Both topologies support all four modalities: `text2img`, `img2img`, `img2text`, `text2text`.
## Quick Start
```bash
cd examples/offline_inference/bagel
# Default two-stage mode (auto-detected)
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2img \
--prompts "A cute cat"
# Single-stage mode
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2img \
--prompts "A cute cat" \
--deploy-config vllm_omni/deploy/bagel_single_stage.yaml
```
### Modality Control
> **Note**: These examples work with the default configuration on an **NVIDIA A100 (80GB)**. For dual-GPU setups, modify the deploy YAML to distribute stages across devices.
BAGEL-7B-MoT supports multiple modality modes. You can control the mode using the `--modality` argument:
## Modality Control
#### Text to Image (text2img)
Control the mode using the `--modality` argument:
- **Pipeline**: Text → Thinker → DiT → VAE Decode → Image
- **Stages Used**: Stage 0 (Thinker) + Stage 1 (DiT)
- **KV Transfer**: Thinker sends KV cache to DiT for conditioned generation
| Modality | Input | Output | Description |
| :------- | :---- | :----- | :---------- |
| `text2img` | Text | Image | Generate images from text prompts |
| `img2img` | Image + Text | Image | Transform images using text guidance |
| `img2text` | Image + Text | Text | Generate text descriptions from images |
| `text2text` | Text | Text | Pure text generation (language model mode) |
Generate images from text prompts:
### Text to Image (text2img)
```bash
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2img \
--prompts "A cute cat"
--prompts "A cute cat" \
--steps 50
```
#### Image to Image (img2img)
- **Pipeline**: Image → VAE Encode → DiT → VAE Decode → New Image
- **Stages Used**: Stage 1 (DiT) only
- **Special**: Bypasses the Thinker stage, direct image-to-image transformation
Transform images based on text prompts:
### Image to Image (img2img)
```bash
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
@@ -50,13 +65,7 @@ python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--prompts "Let the woman wear a blue dress"
```
#### Image to Text (img2text)
- **Pipeline**: Image → ViT + VAE Encode → Thinker → Text Output
- **Stages Used**: Stage 0 (Thinker) only
- **Special**: Uses both VAE latent encoding AND ViT semantic encoding for comprehensive image understanding
Generate text descriptions from images:
### Image to Text (img2text)
```bash
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
@@ -65,196 +74,210 @@ python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--prompts "Describe this image in detail"
```
#### Text to Text (text2text)
- **Pipeline**: Text → Thinker → Text Output
- **Stages Used**: Stage 0 (Thinker) only
- **Special**: No visual components involved, operates as pure language model
Pure text generation:
### Text to Text (text2text)
```bash
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2text \
--prompts "What is the capital of France?"
# You can load prompts from a text file (one prompt per line):
# Load prompts from a text file (one prompt per line):
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2text \
--txt-prompts /path/to/prompts.txt
```
### Inference Steps
## Think Mode
Control the number of inference steps for image generation:
Think mode enables the model to generate `<think>...</think>` planning/reasoning tokens before producing the final output. This improves generation quality for complex prompts.
- **Two-stage**: The Thinker (AR) stage decodes think tokens, then transfers the augmented KV cache to the DiT stage for image generation.
- **Single-stage**: The DiT's internal LLM generates think tokens in-place before proceeding to denoise.
```bash
# You can adjust steps to 100 to improve image quality
# Think + text2img: plan before generating
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2img \
--steps 50 \
--prompts "A cute cat"
```
--prompts "A futuristic city with flying cars" \
--think \
--max-think-tokens 1000
### Key arguments
BAGEL-7B-MoT supports **multiple modality modes** for different use cases.
The default yaml configuration deploys Thinker and DiT on the same GPU. You can use the default configuration file: [`bagel.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/model_executor/stage_configs/bagel.yaml)
#### 📌 Command Line Arguments (end2end.py)
| Argument | Type | Default | Description |
| :--------------------- | :----- | :---------------------------- | :----------------------------------------------------------- |
| `--model` | string | `ByteDance-Seed/BAGEL-7B-MoT` | Model path or name |
| `--modality` | choice | `text2img` | Modality mode: `text2img`, `img2img`, `img2text`, `text2text` |
| `--prompts` | list | `None` | Input text prompts directly |
| `--txt-prompts` | string | `None` | Path to txt file with one prompt per line |
| `--image-path` | string | `None` | Input image path (for `img2img`/`img2text`) |
| `--steps` | int | `50` | Number of inference steps |
| `--stage-configs-path` | string | `None` | Custom stage config file path |
| `--worker-backend` | choice | `process` | Worker backend: `process` or `ray` |
| `--ray-address` | string | `None` | Ray cluster address |
| `--enable-stats` | flag | `False` | Enable statistics logging |
| `--init-sleep-seconds` | int | `20` | Initialization sleep time |
| `--batch-timeout` | int | `5` | Batch timeout |
| `--init-timeout` | int | `300` | Initialization timeout |
------
#### ⚙️ Stage Configuration Parameters (bagel.yaml)
**Stage 0 - Thinker (LLM Stage)**
| Parameter | Value | Description |
| :------------------------------- | :------------------------------ | :----------------------- |
| `stage_type` | `llm` | Stage type |
| `devices` | `"0"` | GPU device ID |
| `max_num_seqs` | `1` | Maximum batch size |
| `model_stage` | `thinker` | Model stage identifier |
| `model_arch` | `BagelForConditionalGeneration` | Model architecture |
| `gpu_memory_utilization` | `0.4` | GPU memory utilization |
| `tensor_parallel_size` | `1` | Tensor parallel size |
| `max_num_batched_tokens` | `32768` | Maximum batched tokens |
| `omni_kv_config.need_send_cache` | `true` | Whether to send KV cache |
------
**Stage 1 - DiT (Diffusion Stage)**
| Parameter | Value | Description |
| :------------------------------- | :---------- | :-------------------------- |
| `stage_type` | `diffusion` | Stage type |
| `devices` | `"0"` | GPU device ID |
| `max_num_seqs` | `1` | Maximum batch size |
| `model_stage` | `dit` | Model stage identifier |
| `gpu_memory_utilization` | `0.4` | GPU memory utilization |
| `omni_kv_config.need_recv_cache` | `true` | Whether to receive KV cache |
| `engine_input_source` | `[0]` | Input source from Stage 0 |
------
#### Tensor Parallelism (TP)
For larger models or multi-GPU environments, you can enable Tensor Parallelism (TP) by modifying the stage configuration (e.g., [`bagel.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/model_executor/stage_configs/bagel.yaml)).
1. **Set `tensor_parallel_size`**: Increase this value (e.g., to `2` or `4`).
2. **Set `devices`**: Specify the comma-separated GPU IDs to be used for the stage (e.g., `"0,1"`).
Example configuration for TP=2 on GPUs 0 and 1:
```yaml
engine_args:
tensor_parallel_size: 2
...
runtime:
devices: "0,1"
```
------
#### 🔗 Runtime Configuration
| Parameter | Value | Description |
| :-------------------- | :------ | :------------------------------- |
| `shm_threshold_bytes` | `65536` | Shared memory threshold (64KB) |
## Using Mooncake Connector
[Mooncake](https://github.com/kvcache-ai/Mooncake) is a high-performance distributed KV cache transfer engine that enables efficient cross-node data movement via TCP or RDMA, making it ideal for multi-node disaggregated inference.
By default, BAGEL uses `SharedMemoryConnector` for inter-stage communication. You can switch to the Mooncake connector for better performance on multi-GPU setups and to enable multi-node deployment.
### Prerequisites
Install the Mooncake transfer engine:
```bash
# For CUDA-enabled systems (recommended)
pip install mooncake-transfer-engine
# For non-CUDA systems
pip install mooncake-transfer-engine-non-cuda
```
### Step 1: Start the Mooncake Master
On the **primary node**, start the Mooncake master service (run in a separate terminal or background with `&`):
```bash
# Optional: enable disk-backed storage by creating a directory and passing --root_fs_dir.
# Without it, Mooncake runs in memory-only mode, which is sufficient for KV cache transfer.
mkdir -p ./mc_storage
mooncake_master \
--rpc_port=50051 \
--enable_http_metadata_server=true \
--http_metadata_server_host=0.0.0.0 \
--http_metadata_server_port=8080 \
--metrics_port=9003 \
--root_fs_dir=./mc_storage/ \
--cluster_id=mc-local-1 &
```
### Step 2: Run Offline Inference with Mooncake
Use the provided Mooncake stage config [`bagel_multiconnector.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml). Before launching, update the `metadata_server` and `master` addresses in the YAML to match your Mooncake master node's IP (use `127.0.0.1` for single-node testing).
```bash
cd examples/offline_inference/bagel
# Text to Image with Mooncake
# Think + img2img: reason about the edit
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2img \
--prompts "A cute cat" \
--stage-configs-path ../../../vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
--modality img2img \
--image-path /path/to/image.jpg \
--prompts "Make it look like a watercolor painting" \
--think
# Image to Text with Mooncake
# Think + img2text: reason before describing
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality img2text \
--image-path /path/to/image.jpg \
--prompts "Describe this image" \
--stage-configs-path ../../../vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
--prompts "What is happening in this image?" \
--think
# Text to Text with Mooncake
# Think + text2text: chain-of-thought reasoning
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2text \
--prompts "What is the capital of France?" \
--stage-configs-path ../../../vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
--prompts "Solve: 23 * 47" \
--think
```
For more details on the Mooncake connector and multi-node setup, see the [Mooncake Store Connector documentation](https://github.com/vllm-project/vllm-omni/tree/main/docs/design/feature/omni_connectors/mooncake_store_connector.md).
Think mode parameters:
------
| Argument | Default | Description |
| :------- | :------ | :---------- |
| `--think` | `False` | Enable thinking mode |
| `--max-think-tokens` | `1000` | Maximum tokens for think generation |
| `--do-sample` | `False` | Enable sampling (vs. greedy) for text generation |
| `--text-temperature` | `0.3` | Temperature for text generation sampling |
## Classifier-Free Guidance (CFG)
CFG controls the trade-off between prompt fidelity and diversity. These parameters apply to image generation modalities (`text2img`, `img2img`).
```bash
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2img \
--prompts "A photorealistic portrait" \
--cfg-text-scale 6.0 \
--cfg-img-scale 2.0 \
--negative-prompt "blurry, low quality, distorted" \
--cfg-interval 0.4 1.0 \
--cfg-renorm-type global \
--cfg-renorm-min 0.0
```
| Argument | Default | Description |
| :------- | :------ | :---------- |
| `--cfg-text-scale` | `4.0` | Text CFG scale (higher = more prompt-adherent) |
| `--cfg-img-scale` | `1.5` | Image CFG scale (for img2img) |
| `--negative-prompt` | `None` | Negative prompt for CFG conditioning |
| `--cfg-interval` | pipeline default | CFG active interval `[start, end]` as fractions of total timesteps |
| `--cfg-renorm-type` | `None` | Renormalization type: `global`, `text_channel`, `channel` |
| `--cfg-renorm-min` | `None` | Minimum renormalization value |
| `--cfg-parallel-size` | `1` | CFG parallel size: `1` = batched (single GPU), `2` = 2-branch parallel, `3` = full 3-GPU parallel |
## Deployment Topologies
### Two-Stage (Default)
The default topology auto-detected from the model. No extra flags needed.
```bash
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2img \
--prompts "A cute cat"
```
The pipeline is defined in [`bagel.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/deploy/bagel.yaml). Stage 0 (Thinker) and Stage 1 (DiT) share GPU 0 by default. For dual-GPU setups, customize the deploy YAML and set `devices: "1"` for stage 1.
### Single-Stage
Pass the single-stage deploy config via `--deploy-config`:
```bash
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2img \
--prompts "A cute cat" \
--deploy-config vllm_omni/deploy/bagel_single_stage.yaml
```
See [`bagel_single_stage.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/deploy/bagel_single_stage.yaml) for configuration details. The `pipeline: bagel_single_stage` field selects the single-stage topology from the pipeline registry.
### Tensor Parallelism (TP)
For larger models or multi-GPU environments, customize the deploy YAML (see [`bagel.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/deploy/bagel.yaml)) and set per-stage `tensor_parallel_size` and `devices`:
```yaml
# Example: TP=2 on GPUs 0,1 for the Thinker stage
stages:
- stage_id: 0
tensor_parallel_size: 2
devices: "0,1"
```
Then pass the custom deploy YAML:
```bash
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2img \
--prompts "A cute cat" \
--deploy-config /path/to/custom_bagel.yaml
```
### FP8 Quantization
```bash
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2img \
--prompts "A cute cat" \
--quantization fp8
```
## Command Line Reference
### Core Arguments
| Argument | Type | Default | Description |
| :------- | :--- | :------ | :---------- |
| `--model` | string | `ByteDance-Seed/BAGEL-7B-MoT` | Model path or HuggingFace name |
| `--modality` | choice | `text2img` | `text2img`, `img2img`, `img2text`, `text2text` |
| `--prompts` | list | `None` | Input text prompts |
| `--txt-prompts` | string | `None` | Path to text file with one prompt per line |
| `--image-path` | string | `None` | Input image path (required for `img2img`/`img2text`) |
| `--output` | string | `.` | Output directory for saved images |
| `--steps` | int | `50` | Number of diffusion inference steps |
| `--seed` | int | `None` | Random seed for reproducibility |
### Think Mode Arguments
| Argument | Type | Default | Description |
| :------- | :--- | :------ | :---------- |
| `--think` | flag | `False` | Enable `<think>...</think>` planning/reasoning |
| `--max-think-tokens` | int | `1000` | Maximum tokens for think generation |
| `--do-sample` | flag | `False` | Use sampling instead of greedy decoding |
| `--text-temperature` | float | `0.3` | Sampling temperature for text generation |
### CFG Arguments
| Argument | Type | Default | Description |
| :------- | :--- | :------ | :---------- |
| `--cfg-text-scale` | float | `4.0` | Text CFG guidance scale |
| `--cfg-img-scale` | float | `1.5` | Image CFG guidance scale |
| `--negative-prompt` | string | `None` | Negative prompt for CFG |
| `--cfg-parallel-size` | int | `1` | CFG parallel GPU count (1, 2, or 3) |
| `--cfg-interval` | float[2] | pipeline default | CFG active window `[start, end]` |
| `--cfg-renorm-type` | string | `None` | `global`, `text_channel`, or `channel` |
| `--cfg-renorm-min` | float | `None` | Minimum renormalization value |
### Engine Arguments
| Argument | Type | Default | Description |
| :------- | :--- | :------ | :---------- |
| `--deploy-config` | string | `None` | Path to deploy YAML (auto-detected if omitted) |
| `--stage-configs-path` | string | `None` | [Deprecated] Legacy path to `stage_args` YAML; prefer `--deploy-config` |
| `--worker-backend` | choice | `process` | `process` or `ray` |
| `--ray-address` | string | `None` | Ray cluster address |
| `--quantization` | string | `None` | Quantization method (e.g. `fp8`) |
| `--log-stats` | flag | `False` | Enable statistics logging |
| `--init-timeout` | int | `300` | Initialization timeout (seconds) |
| `--batch-timeout` | int | `5` | Batch timeout (seconds) |
| `--enable-diffusion-pipeline-profiler` | flag | `False` | Profile diffusion stage durations |
## FAQ
- If you don’t know how much VRAM is needed for the model or encounter the OOM error, you can try to decrease the max_model_len.
- If you encounter OOM errors, try decreasing `max_model_len` or `gpu_memory_utilization` in the deploy YAML.
| Stage | VRAM |
| :------------------ | :--------------------------- |
| Stage-0 (Thinker) | **15.04 GiB** **+ KV Cache** |
| Stage-1 (DiT) | **26.50 GiB** |
| Total | **~42 GiB + KV Cache** |
**Two-stage VRAM usage:**
| Stage | VRAM |
| :---- | :--- |
| Stage 0 (Thinker) | **15.04 GiB + KV Cache** |
| Stage 1 (DiT) | **26.50 GiB** |
| Total | **~42 GiB + KV Cache** |
**Single-stage VRAM usage:** The DiT loads the full model (~42 GiB) in one process.
## Example materials
docs/user_guide/examples/offline_inference/text_to_video.md
+2
View File
@@ -5,6 +5,8 @@ Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inf
A unified script for text-to-video generation. Supports multiple models with model-aware defaults.
For backend selection and SageAttention usage, see the [Diffusion Attention Backends](../../diffusion/attention_backends.md) guide.
## Supported Models
| Model | Default Resolution | Default Frames | Default Steps | Guidance | VRAM (BF16) |
docs/user_guide/examples/online_serving/bagel.md
+143 -185
View File
@@ -2,147 +2,112 @@
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/bagel>.
## 🛠️ Installation
## Installation
Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/README.md)
## Run examples (BAGEL-7B-MoT)
## Architecture
**Note**: These examples work with the default configuration on an **NVIDIA A100 (80GB)**. We also tested on dual **NVIDIA RTX 5000 Ada (32GB each)**. For dual-GPU setups, please modify the stage configuration to distribute the model across devices.
BAGEL-7B-MoT is a Mixture-of-Transformers (MoT) model supporting both image generation and understanding. It offers two deployment topologies:
### Launch the Server
| Topology | Stages | Description |
| :------- | :----- | :---------- |
| **Two-stage** (default) | Stage 0 (Thinker, AR) + Stage 1 (DiT, Diffusion) | Thinker handles text/understanding via vLLM AR engine; DiT handles image generation. KV cache is transferred between stages. |
| **Single-stage** | Stage 0 (DiT, Diffusion) only | The DiT stage contains a full LLM, ViT, VAE, and tokenizer internally. All modalities are handled within a single diffusion process. |
Both topologies support all four modalities: `text2img`, `img2img`, `img2text`, `text2text`.
> **Note**: These examples work with the default configuration on an **NVIDIA A100 (80GB)**. We also tested on dual **NVIDIA RTX 5000 Ada (32GB each)**. For dual-GPU setups, modify the deploy YAML to distribute stages across devices.
## Launch the Server
### Two-Stage (Default)
The default pipeline is auto-detected from the model. No extra flags needed:
```bash
# Use default configuration
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091
```
Or use the convenience script:
```bash
cd /workspace/vllm-omni/examples/online_serving/bagel
cd examples/online_serving/bagel
bash run_server.sh
# Launch a single stage per terminal
bash run_server_stage_cli.sh --stage 0
bash run_server_stage_cli.sh --stage 1
```
```bash
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
```
#### 🚀 Tensor Parallelism (TP)
For larger models or multi-GPU environments, you can enable Tensor Parallelism (TP) for the server.
1. **Modify Stage Config**: Create or modify a stage configuration yaml (e.g., [`bagel.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/model_executor/stage_configs/bagel.yaml)). Set `tensor_parallel_size` to `2` (or more) and update `devices` to include multiple GPU IDs (e.g., `"0,1"`).
```yaml
engine_args:
tensor_parallel_size: 2
...
runtime:
devices: "0,1"
```
2. **Launch Server**:
```bash
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --stage-configs-path /path/to/your/custom_bagel.yaml
```
#### Using Mooncake Connector
By default, BAGEL uses `SharedMemoryConnector` for inter-stage communication. You can use the [Mooncake](https://github.com/kvcache-ai/Mooncake) connector to transfer KV cache between stages, which also enables multi-node deployment.
**1. Install Mooncake**
```bash
# For CUDA-enabled systems (recommended)
pip install mooncake-transfer-engine
# For non-CUDA systems
pip install mooncake-transfer-engine-non-cuda
```
**2. Start Mooncake Master** on the primary node:
```bash
# Optional: enable disk-backed storage by creating a directory and passing --root_fs_dir.
# Without it, Mooncake runs in memory-only mode, which is sufficient for KV cache transfer.
mkdir -p ./mc_storage
mooncake_master \
--rpc_port=50051 \
--enable_http_metadata_server=true \
--http_metadata_server_host=0.0.0.0 \
--http_metadata_server_port=8080 \
--metrics_port=9003 \
--root_fs_dir=./mc_storage/ \
--cluster_id=mc-local-1 &
```
**3. Launch the server** with the Mooncake stage config:
To use a custom deploy YAML (note: `--stage-configs-path` is deprecated in favor of `--deploy-config`):
```bash
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \
--stage-configs-path vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
--deploy-config /path/to/deploy_config.yaml
```
> **Note**: Before launching, edit [`bagel_multiconnector.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml) and replace the `metadata_server` and `master` addresses with your Mooncake master node's actual IP. For single-node testing, `127.0.0.1` works.
See [`bagel.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/deploy/bagel.yaml) for the default two-stage deploy configuration.
The client-side usage is identical to the default setup -- the Mooncake connector is transparent to the API. See the requests section below.
### Single-Stage
For more details on the Mooncake connector configuration, see the [Mooncake Store Connector documentation](https://github.com/vllm-project/vllm-omni/tree/main/docs/design/feature/omni_connectors/mooncake_store_connector.md).
#### Multi-Node Deployment
You can deploy each stage on a **separate node** for better resource utilization. In this example, the orchestrator (Stage 0 / Thinker) and Stage 1 (DiT) run on different machines, connected via Mooncake.
Replace `<ORCHESTRATOR_IP>` below with the actual IP address of your orchestrator node (e.g., `10.244.227.244`).
> [!WARNING]
> **Before launching**, edit [`bagel_multiconnector.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml) and replace the `metadata_server` and `master` addresses with your Mooncake master node's actual IP. Mismatched addresses will cause silent connection failures.
**1. Start Mooncake Master** (on the orchestrator node):
The DiT stage contains a full LLM, ViT, VAE, and tokenizer, so it can handle all modalities (text2img, img2img, img2text, text2text, think) without a separate Thinker stage:
```bash
mooncake_master \
--rpc_port=50051 \
--enable_http_metadata_server=true \
--http_metadata_server_host=<ORCHESTRATOR_IP> \
--http_metadata_server_port=8080 \
--metrics_port=9003
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \
--deploy-config vllm_omni/deploy/bagel_single_stage.yaml
```
**2. Launch Stage 0 (Thinker / Orchestrator)** on the orchestrator node:
See [`bagel_single_stage.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/deploy/bagel_single_stage.yaml) for configuration. The `pipeline: bagel_single_stage` field selects the single-stage topology from the pipeline registry.
### Tensor Parallelism (TP)
For larger models or multi-GPU environments, enable TP via CLI:
```bash
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --tensor-parallel-size 2
```
Or set `tensor_parallel_size` per stage in a custom deploy YAML.
### Multi-Node Deployment
Deploy each stage on a **separate node** for better resource utilization. Replace `<ORCHESTRATOR_IP>` with the actual IP address of your orchestrator node.
**1. Launch Stage 0 (Thinker / Orchestrator)** on the orchestrator node:
```bash
# API server port for client requests: 8000
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni \
--port 8000 \ # API server port for client requests
--stage-configs-path vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml \
--port 8000 \
--stage-id 0 \
-oma <ORCHESTRATOR_IP> \
-omp 8091
--omni-master-address <ORCHESTRATOR_IP> \
--omni-master-port 8091
```
**3. Launch Stage 1 (DiT)** on the remote node in headless mode:
**2. Launch Stage 1 (DiT)** on the remote node in headless mode:
```bash
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni \
--stage-configs-path vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml \
--stage-id 1 \
--headless \
-oma <ORCHESTRATOR_IP> \
-omp 8091
--omni-master-address <ORCHESTRATOR_IP> \
--omni-master-port 8091
```
**Mooncake Master arguments:**
Or use the convenience script:
| Argument | Description |
| :------- | :---------- |
| `--rpc_port` | Mooncake RPC port for control-plane coordination between stages |
| `--enable_http_metadata_server` | Enable the HTTP metadata server for service discovery |
| `--http_metadata_server_host` | IP address to bind the metadata server (use the orchestrator node's IP) |
| `--http_metadata_server_port` | Port for the HTTP metadata server |
| `--metrics_port` | Port for Prometheus-compatible metrics endpoint |
```bash
# Terminal 1: Stage 0
bash run_server_stage_cli.sh --stage 0
# Terminal 2: Stage 1
bash run_server_stage_cli.sh --stage 1
# With extra args
bash run_server_stage_cli.sh --stage 0 -- --tensor-parallel-size 2
bash run_server_stage_cli.sh --stage 1 -- --gpu-memory-utilization 0.9
```
**vllm serve arguments:**
@@ -150,85 +115,31 @@ vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni \
| :------- | :---------- |
| `--stage-id` | Which stage this process runs (0 = Thinker, 1 = DiT) |
| `--headless` | Run without the API server (worker-only mode) |
| `-oma` | Orchestrator master address |
| `-omp` | Orchestrator master port for Stage 1 to connect to Stage 0 for task coordination |
| `-oma` / `--omni-master-address` | Orchestrator master address |
| `-omp` / `--omni-master-port` | Orchestrator master port |
> [!IMPORTANT]
> **Startup Order**: Stage 0 (orchestrator) must be launched **before** Stage 1 (headless).
> Stage 0 will appear to hang on startup until Stage 1 (worker) connects — this is expected behavior.
**Network Requirements**
### Inter-Stage Connectors
All nodes must have network connectivity to each other. Ensure the following ports are open **between all participating nodes**:
When deploying stages across nodes, configure the connector type in the deploy YAML:
| Port | Protocol | Service | Direction |
| :--- | :------- | :------ | :-------- |
| 50051 | TCP | Mooncake Master RPC | Worker → Orchestrator |
| 8080 | TCP | Mooncake HTTP Metadata Server | Worker → Orchestrator |
| 8091 | TCP | Orchestrator Master (`-omp`) | Worker → Orchestrator |
| 8000 | TCP | API Server (`--port`) | Client → Orchestrator |
| 9003 | TCP | Metrics (optional) | Monitoring → Orchestrator |
- **SharedMemoryConnector** (default): Used for single-node deployments. No explicit configuration needed.
- **MooncakeTransferEngineConnector**: For multi-node setups with RDMA hardware. Defined in [`bagel.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/deploy/bagel.yaml) under `connectors.rdma_connector`.
> **Tip**: If nodes are behind a firewall or in different VPCs/security groups, make sure the above ports are allowed in ingress/egress rules. All nodes should be reachable via their IP addresses (no NAT). Using nodes on the same subnet or VPC is recommended to minimize latency for Mooncake KV cache transfers.
To use Mooncake, create a custom deploy YAML that binds `output_connectors` / `input_connectors` on each stage to the `rdma_connector` defined in the `connectors` section.
### Send Multi-modal Request
Get into the bagel folder:
## Send Requests
```bash
cd examples/online_serving/bagel
```
Send request via Python
```bash
python openai_chat_client.py --prompt "A cute cat" --modality text2img
```
The Python client supports the following command-line arguments:
- `--prompt` (or `-p`): Text prompt for generation (default: `A cute cat`)
- `--output` (or `-o`): Output file path for image results (default: `bagel_output.png`)
- `--server` (or `-s`): Server URL (default: `http://localhost:8091`)
- `--image-url` (or `-i`): Input image URL or local file path (for img2img/img2text modes)
- `--modality` (or `-m`): Task modality (default: `text2img`). Options: `text2img`, `img2img`, `img2text`, `text2text`
- `--height`: Image height in pixels (default: 512)
- `--width`: Image width in pixels (default: 512)
- `--steps`: Number of inference steps (default: 25)
- `--seed`: Random seed (default: 42)
- `--negative`: Negative prompt for image generation
Example with custom parameters:
```bash
python openai_chat_client.py \
--prompt "A futuristic city" \
--modality text2img \
--height 768 \
--width 768 \
--steps 50 \
--seed 42 \
--negative "blurry, low quality"
```
## Modality Control
BAGEL-7B-MoT supports **multiple modality modes** for different use cases.
The default yaml configuration deploys Thinker and DiT on the same GPU. You can use the default configuration file: [`bagel.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/model_executor/stage_configs/bagel.yaml)
| Modality | Input | Output | Description |
| ----------- | ------------ | ------ | -------------------------------------- |
| `text2img` | Text | Image | Generate images from text prompts |
| `img2img` | Image + Text | Image | Transform images using text guidance |
| `img2text` | Image + Text | Text | Generate text descriptions from images |
| `text2text` | Text | Text | Pure text generation |
### Text to Image (text2img)
Generate images from text prompts:
**Using Python client**
**Python client:**
```bash
python openai_chat_client.py \
@@ -238,7 +149,7 @@ python openai_chat_client.py \
--steps 50
```
**Using curl**
**curl:**
```bash
curl http://localhost:8091/v1/chat/completions \
@@ -253,12 +164,9 @@ curl http://localhost:8091/v1/chat/completions \
}'
```
### Image to Image (img2img)
Transform images based on text prompts:
**Using Python client**
**Python client:**
```bash
python openai_chat_client.py \
@@ -268,7 +176,7 @@ python openai_chat_client.py \
--output transformed.png
```
**Using curl**
**curl:**
```bash
IMAGE_BASE64=$(base64 -w 0 cat.jpg)
@@ -293,14 +201,11 @@ EOF
curl http://localhost:8091/v1/chat/completions \
-H "Content-Type: application/json" \
-d @payload.json
```
### Image to Text (img2text)
Generate text descriptions from images:
**Using Python client**
**Python client:**
```bash
python openai_chat_client.py \
@@ -309,7 +214,7 @@ python openai_chat_client.py \
--image-url /path/to/image.jpg
```
**Using curl**
**curl:**
```bash
IMAGE_BASE64=$(base64 -w 0 cat.jpg)
@@ -334,9 +239,7 @@ curl http://localhost:8091/v1/chat/completions \
### Text to Text (text2text)
Pure text generation:
**Using Python client**
**Python client:**
```bash
python openai_chat_client.py \
@@ -344,26 +247,81 @@ python openai_chat_client.py \
--modality text2text
```
**Using curl**
**curl:**
```bash
curl http://localhost:8091/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{
"messages": [{"role": "user", "content": [{"type": "text", "text": "<|im_start|>user\nWhat is the capital of France?<|im_end|>\n<|im_start|>assistant\n"}]}]
"messages": [{"role": "user", "content": [{"type": "text", "text": "<|im_start|>user\nWhat is the capital of France?<|im_end|>\n<|im_start|>assistant\n"}]}],
"modalities": ["text"]
}'
```
### Python Client Arguments
| Argument | Default | Description |
| :------- | :------ | :---------- |
| `--prompt` / `-p` | `A cute cat` | Text prompt |
| `--output` / `-o` | `bagel_output.png` | Output file path |
| `--server` / `-s` | `http://localhost:8091` | Server URL |
| `--image-url` / `-i` | `None` | Input image URL or local path (img2img/img2text) |
| `--modality` / `-m` | `text2img` | `text2img`, `img2img`, `img2text`, `text2text` |
| `--height` | `512` | Image height in pixels |
| `--width` | `512` | Image width in pixels |
| `--steps` | `25` | Number of inference steps |
| `--seed` | `42` | Random seed |
| `--negative` | `None` | Negative prompt for CFG |
Example with custom parameters:
```bash
python openai_chat_client.py \
--prompt "A futuristic city" \
--modality text2img \
--height 768 \
--width 768 \
--steps 50 \
--seed 42 \
--negative "blurry, low quality"
```
## Configuration Reference
### Deploy YAML Files
| File | Description |
| :--- | :---------- |
| [`bagel.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/deploy/bagel.yaml) | Two-stage default (Thinker + DiT on GPU 0) |
| [`bagel_single_stage.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/deploy/bagel_single_stage.yaml) | Single-stage (DiT only) |
### Key Deploy YAML Fields
| Field | Scope | Description |
| :---- | :---- | :---------- |
| `pipeline` | top-level | Override auto-detected pipeline (e.g. `bagel_single_stage`) |
| `stages[].stage_id` | per-stage | Stage identifier (0, 1, ...) |
| `stages[].devices` | per-stage | GPU device IDs (e.g. `"0"`, `"0,1"`) |
| `stages[].max_num_seqs` | per-stage | Maximum concurrent sequences |
| `stages[].gpu_memory_utilization` | per-stage | Fraction of GPU memory to use |
| `stages[].enforce_eager` | per-stage | Disable CUDA graphs |
| `stages[].tensor_parallel_size` | per-stage | TP degree for this stage |
| `connectors` | top-level | Define available connector instances (SHM, Mooncake) |
| `platforms` | top-level | Platform-specific overrides (e.g. `xpu`) |
## FAQ
- If you don’t know how much VRAM is needed for the model or encounter the OOM error, you can try to decrease the max_model_len.
- If you encounter OOM errors, try decreasing `max_model_len` or `gpu_memory_utilization` in the deploy YAML.
| Stage | VRAM |
| :------------------ | :--------------------------- |
| Stage-0 (Thinker) | **15.04 GiB** **+ KV Cache** |
| Stage-1 (DiT) | **26.50 GiB** |
| Total | **~42 GiB + KV Cache** |
**Two-stage VRAM usage:**
| Stage | VRAM |
| :---- | :--- |
| Stage 0 (Thinker) | **15.04 GiB + KV Cache** |
| Stage 1 (DiT) | **26.50 GiB** |
| Total | **~42 GiB + KV Cache** |
**Single-stage VRAM usage:** The DiT loads the full model (~42 GiB) in one process.
## Example materials
docs/user_guide/examples/online_serving/diffusers_pipeline_adapter.md
+93
View File
@@ -0,0 +1,93 @@
# Diffusers Backend Adapter Example
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/diffusers_pipeline_adapter>.
This example demonstrates how to serve any 🤗 Diffusers pipeline through vLLM-Omni
using the `diffusers` load format.
## Supported Models
Any model loadable via `DiffusionPipeline.from_pretrained()` should be supported, including text-to-image, image-to-image, text-to-video, image-to-video, and text-to-audio.
## Limitations
The diffusers backend is a black-box adapter. The following features are NOT yet supported.
It is not guaranteed whether they will be supported in the future.
- CFG parallel execution
- Sequence parallel execution
- TeaCache / Cache-DiT acceleration
- Step-wise execution (continuous batching)
For these features, it is recommended to use natively supported pipelines instead.
## Usage
### Option 1: CLI arguments
```bash
vllm serve "stable-diffusion-v1-5/stable-diffusion-v1-5" \
--omni \
--diffusion-load-format diffusers \
--diffusers-load-kwargs '{"use_safetensors": true}' \
--diffusers-call-kwargs '{"num_inference_steps": 30, "guidance_scale": 7.5}'
```
`--diffusers-load-kwargs` and `--diffusers-call-kwargs` are only valid together with `--diffusion-load-format diffusers`.
### Option 2: Stage config YAML
```bash
vllm serve stable-diffusion-v1-5/stable-diffusion-v1-5 --stage-configs-path examples/online_serving/diffusers_pipeline_adapter/stage_config.yaml --omni
```
The particular fields of interest are `model`, `diffusion_load_format`, `diffusers_load_kwargs`, and `diffusers_call_kwargs` under `engine_args`. They are the same as the CLI arguments.
## Send a Request
```bash
curl http://localhost:8000/v1/images/generations \
-H "Content-Type: application/json" \
-d '{
"model": "stable-diffusion-v1-5/stable-diffusion-v1-5",
"prompt": "a photo of an astronaut riding a horse on mars",
"n": 1,
"size": "512x512"
}'
```
Or refer to other documentation pages on how to request a particular input/output modality, such as `examples/online_serving/text_to_image/openai_chat_client.py`.
## Configuration Reference
For the diffusers adapter, set options under **`engine_args`**:
### `diffusion_load_format: "diffusers"`
This field selects the Hugging Face diffusers adapter path (see `DiffusersPipelineLoader`).
### `diffusers_load_kwargs`
Passed to `DiffusionPipeline.from_pretrained()`.
This is suitable for model-specific configurations not available through the vLLM-Omni interface (such as `Omni.__init__()`, `vllm serve` CLI arguments, and stage config YAML fields outside `diffusers_load_kwargs`).
When a parameter is available in the vLLM-Omni interface, it will be adapted here.
But if that parameter is simultaneously set in both the vLLM-Omni interface and `diffusers_load_kwargs`, the **latter** will take precedence.
### `diffusers_call_kwargs`
Passed to `pipeline.__call__()`.
This is suitable for sampling parameters not available through the vLLM-Omni interface (such as `Omni.generate()` and online serving payloads).
When a parameter is available in the vLLM-Omni interface, it will be adapted here.
But if that parameter is simultaneously set in both the vLLM-Omni interface and `diffusers_call_kwargs`, the **former** will take precedence (because it is set at request time).
## Example materials
??? abstract "stage_config.yaml"
``````yaml
--8<-- "examples/online_serving/diffusers_pipeline_adapter/stage_config.yaml"
``````
docs/user_guide/examples/online_serving/glm_image.md
+3 -3
View File
@@ -104,9 +104,9 @@ guides.
When using `/v1/chat/completions`, pass these inside `extra_body` in the curl
JSON, or via the `extra_body` keyword argument in the OpenAI Python SDK (see the
[Diffusion Chat API guide](../../../../serving/diffusion_chat_api.md)).
When using the dedicated [`/v1/images/generations`](../../../../serving/image_generation_api.md)
or [`/v1/images/edits`](../../../../serving/image_edit_api.md) endpoints, pass
[Diffusion Chat API guide](../../../serving/diffusion_chat_api.md)).
When using the dedicated [`/v1/images/generations`](../../../serving/image_generation_api.md)
or [`/v1/images/edits`](../../../serving/image_edit_api.md) endpoints, pass
the supported generation controls as top-level fields directly. For image
dimensions and count, use `size` and `n` rather than `height` or `width`.
docs/user_guide/examples/online_serving/qwen3_omni.md
+61 -4
View File
@@ -15,15 +15,72 @@ Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
```
If you want to open async chunking for qwen3-omni, launch the server with command below
The default deployment configuration situated at `vllm_omni/deploy/qwen3_omni_moe.yaml` is resolved and loaded
automatically via the model registry, obviating the necessity for the `--deploy-config` flag in standard deployment topologies.
Asynchronous chunk streaming is **enabled by default** within the bundled configuration.
To explicitly utilize a custom deployment YAML, specify the configuration path:
```bash
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --deploy-config /vllm_omni/deploy/qwen3_omni_moe.yaml
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
--deploy-config /path/to/deploy_config_file
```
If you have custom stage configs file, launch the server with command below
### Launch individual stages (stage-based CLI)
Adopt the stage-based CLI architecture to independently instantiate execution processes per functional stage.
**1. Stage 0 (Thinker + API server)**
```bash
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --deploy-config /path/to/deploy_config_file
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
--port 8091 \
--stage-id 0 \
--omni-master-address 127.0.0.1 \
--omni-master-port 26000
```
**2. Stage 1 (Talker)**
```bash
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
--stage-id 1 \
--headless \
--omni-master-address 127.0.0.1 \
--omni-master-port 26000
```
**3. Stage 2 (Code2Wav)**
```bash
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
--stage-id 2 \
--headless \
--omni-master-address 127.0.0.1 \
--omni-master-port 26000
```
Add `--deploy-config /path/to/deploy_config_file` to every command if you want
to override the bundled deploy YAML.
For the regular one-process launch, stage-specific CLI tuning is usually done
with `--stage-overrides`, for example:
```bash
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
--stage-overrides '{"1": {"gpu_memory_utilization": 0.5}}'
```
For the stage-based CLI, you usually do **not** need `--stage-overrides` for
that kind of change. Since each command launches one stage, just pass the knob
directly on that stage command:
```bash
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
--stage-id 1 \
--headless \
--gpu-memory-utilization 0.5 \
--omni-master-address 127.0.0.1 \
--omni-master-port 26000
```
### Send Multi-modal Request
examples/offline_inference/bagel/README.md
+207 -184
View File
@@ -1,44 +1,60 @@
# BAGEL-7B-MoT
## Set up
## Setup
Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
## Run examples
## Architecture
**Note**: These examples work with the default configuration on an **NVIDIA A100 (80GB)**. We also tested on dual **NVIDIA RTX 5000 Ada (32GB each)**. For dual-GPU setups, please modify the stage configuration to distribute the model across devices.
BAGEL-7B-MoT is a Mixture-of-Transformers (MoT) model supporting both image generation and understanding. It offers two deployment topologies:
Get into the bagel folder
| Topology | Stages | Description |
| :------- | :----- | :---------- |
| **Two-stage** (default) | Stage 0 (Thinker, AR) + Stage 1 (DiT, Diffusion) | Thinker handles text/understanding via vLLM AR engine; DiT handles image generation. KV cache is transferred between stages. |
| **Single-stage** | Stage 0 (DiT, Diffusion) only | The DiT stage contains a full LLM, ViT, VAE, and tokenizer internally. All modalities are handled within a single diffusion process. |
Both topologies support all four modalities: `text2img`, `img2img`, `img2text`, `text2text`.
## Quick Start
```bash
cd examples/offline_inference/bagel
# Default two-stage mode (auto-detected)
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2img \
--prompts "A cute cat"
# Single-stage mode
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2img \
--prompts "A cute cat" \
--deploy-config vllm_omni/deploy/bagel_single_stage.yaml
```
### Modality Control
> **Note**: These examples work with the default configuration on an **NVIDIA A100 (80GB)**. For dual-GPU setups, modify the deploy YAML to distribute stages across devices.
BAGEL-7B-MoT supports multiple modality modes. You can control the mode using the `--modality` argument:
## Modality Control
#### Text to Image (text2img)
Control the mode using the `--modality` argument:
- **Pipeline**: Text → Thinker → DiT → VAE Decode → Image
- **Stages Used**: Stage 0 (Thinker) + Stage 1 (DiT)
- **KV Transfer**: Thinker sends KV cache to DiT for conditioned generation
| Modality | Input | Output | Description |
| :------- | :---- | :----- | :---------- |
| `text2img` | Text | Image | Generate images from text prompts |
| `img2img` | Image + Text | Image | Transform images using text guidance |
| `img2text` | Image + Text | Text | Generate text descriptions from images |
| `text2text` | Text | Text | Pure text generation (language model mode) |
Generate images from text prompts:
### Text to Image (text2img)
```bash
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2img \
--prompts "A cute cat"
--prompts "A cute cat" \
--steps 50
```
#### Image to Image (img2img)
- **Pipeline**: Image → VAE Encode → DiT → VAE Decode → New Image
- **Stages Used**: Stage 1 (DiT) only
- **Special**: Bypasses the Thinker stage, direct image-to-image transformation
Transform images based on text prompts:
### Image to Image (img2img)
```bash
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
@@ -47,13 +63,7 @@ python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--prompts "Let the woman wear a blue dress"
```
#### Image to Text (img2text)
- **Pipeline**: Image → ViT + VAE Encode → Thinker → Text Output
- **Stages Used**: Stage 0 (Thinker) only
- **Special**: Uses both VAE latent encoding AND ViT semantic encoding for comprehensive image understanding
Generate text descriptions from images:
### Image to Text (img2text)
```bash
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
@@ -62,193 +72,206 @@ python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--prompts "Describe this image in detail"
```
#### Text to Text (text2text)
- **Pipeline**: Text → Thinker → Text Output
- **Stages Used**: Stage 0 (Thinker) only
- **Special**: No visual components involved, operates as pure language model
Pure text generation:
### Text to Text (text2text)
```bash
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2text \
--prompts "What is the capital of France?"
# You can load prompts from a text file (one prompt per line):
# Load prompts from a text file (one prompt per line):
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2text \
--txt-prompts /path/to/prompts.txt
```
### Inference Steps
## Think Mode
Control the number of inference steps for image generation:
Think mode enables the model to generate `<think>...</think>` planning/reasoning tokens before producing the final output. This improves generation quality for complex prompts.
- **Two-stage**: The Thinker (AR) stage decodes think tokens, then transfers the augmented KV cache to the DiT stage for image generation.
- **Single-stage**: The DiT's internal LLM generates think tokens in-place before proceeding to denoise.
```bash
# You can adjust steps to 100 to improve image quality
# Think + text2img: plan before generating
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2img \
--steps 50 \
--prompts "A cute cat"
```
--prompts "A futuristic city with flying cars" \
--think \
--max-think-tokens 1000
### Key arguments
BAGEL-7B-MoT supports **multiple modality modes** for different use cases.
The default yaml configuration deploys Thinker and DiT on the same GPU. You can use the default configuration file: [`bagel.yaml`](../../../vllm_omni/model_executor/stage_configs/bagel.yaml)
#### 📌 Command Line Arguments (end2end.py)
| Argument | Type | Default | Description |
| :--------------------- | :----- | :---------------------------- | :----------------------------------------------------------- |
| `--model` | string | `ByteDance-Seed/BAGEL-7B-MoT` | Model path or name |
| `--modality` | choice | `text2img` | Modality mode: `text2img`, `img2img`, `img2text`, `text2text` |
| `--prompts` | list | `None` | Input text prompts directly |
| `--txt-prompts` | string | `None` | Path to txt file with one prompt per line |
| `--image-path` | string | `None` | Input image path (for `img2img`/`img2text`) |
| `--steps` | int | `50` | Number of inference steps |
| `--stage-configs-path` | string | `None` | Custom stage config file path |
| `--worker-backend` | choice | `process` | Worker backend: `process` or `ray` |
| `--ray-address` | string | `None` | Ray cluster address |
| `--enable-stats` | flag | `False` | Enable statistics logging |
| `--init-sleep-seconds` | int | `20` | Initialization sleep time |
| `--batch-timeout` | int | `5` | Batch timeout |
| `--init-timeout` | int | `300` | Initialization timeout |
------
#### ⚙️ Stage Configuration Parameters (bagel.yaml)
**Stage 0 - Thinker (LLM Stage)**
| Parameter | Value | Description |
| :------------------------------- | :------------------------------ | :----------------------- |
| `stage_type` | `llm` | Stage type |
| `devices` | `"0"` | GPU device ID |
| `max_num_seqs` | `1` | Maximum batch size |
| `model_stage` | `thinker` | Model stage identifier |
| `model_arch` | `BagelForConditionalGeneration` | Model architecture |
| `gpu_memory_utilization` | `0.4` | GPU memory utilization |
| `tensor_parallel_size` | `1` | Tensor parallel size |
| `max_num_batched_tokens` | `32768` | Maximum batched tokens |
| `omni_kv_config.need_send_cache` | `true` | Whether to send KV cache |
------
**Stage 1 - DiT (Diffusion Stage)**
| Parameter | Value | Description |
| :------------------------------- | :---------- | :-------------------------- |
| `stage_type` | `diffusion` | Stage type |
| `devices` | `"0"` | GPU device ID |
| `max_num_seqs` | `1` | Maximum batch size |
| `model_stage` | `dit` | Model stage identifier |
| `gpu_memory_utilization` | `0.4` | GPU memory utilization |
| `omni_kv_config.need_recv_cache` | `true` | Whether to receive KV cache |
| `engine_input_source` | `[0]` | Input source from Stage 0 |
------
#### Tensor Parallelism (TP)
For larger models or multi-GPU environments, you can enable Tensor Parallelism (TP) by modifying the stage configuration (e.g., [`bagel.yaml`](../../../vllm_omni/model_executor/stage_configs/bagel.yaml)).
1. **Set `tensor_parallel_size`**: Increase this value (e.g., to `2` or `4`).
2. **Set `devices`**: Specify the comma-separated GPU IDs to be used for the stage (e.g., `"0,1"`).
Example configuration for TP=2 on GPUs 0 and 1:
```yaml
engine_args:
tensor_parallel_size: 2
...
runtime:
devices: "0,1"
```
------
#### 🔗 Runtime Configuration
| Parameter | Value | Description |
| :-------------------- | :------ | :------------------------------- |
| `shm_threshold_bytes` | `65536` | Shared memory threshold (64KB) |
## Using Mooncake Connector
[Mooncake](https://github.com/kvcache-ai/Mooncake) is a high-performance distributed KV cache transfer engine that enables efficient cross-node data movement via TCP or RDMA, making it ideal for multi-node disaggregated inference.
By default, BAGEL uses `SharedMemoryConnector` for inter-stage communication. You can switch to the Mooncake connector for better performance on multi-GPU setups and to enable multi-node deployment.
### Prerequisites
Install the Mooncake transfer engine:
```bash
# For CUDA-enabled systems (recommended)
pip install mooncake-transfer-engine
# For non-CUDA systems
pip install mooncake-transfer-engine-non-cuda
```
### Step 1: Start the Mooncake Master
On the **primary node**, start the Mooncake master service (run in a separate terminal or background with `&`):
```bash
# Optional: enable disk-backed storage by creating a directory and passing --root_fs_dir.
# Without it, Mooncake runs in memory-only mode, which is sufficient for KV cache transfer.
mkdir -p ./mc_storage
mooncake_master \
--rpc_port=50051 \
--enable_http_metadata_server=true \
--http_metadata_server_host=0.0.0.0 \
--http_metadata_server_port=8080 \
--metrics_port=9003 \
--root_fs_dir=./mc_storage/ \
--cluster_id=mc-local-1 &
```
### Step 2: Run Offline Inference with Mooncake
Use the provided Mooncake stage config [`bagel_multiconnector.yaml`](../../../vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml). Before launching, update the `metadata_server` and `master` addresses in the YAML to match your Mooncake master node's IP (use `127.0.0.1` for single-node testing).
```bash
cd examples/offline_inference/bagel
# Text to Image with Mooncake
# Think + img2img: reason about the edit
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2img \
--prompts "A cute cat" \
--stage-configs-path ../../../vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
--modality img2img \
--image-path /path/to/image.jpg \
--prompts "Make it look like a watercolor painting" \
--think
# Image to Text with Mooncake
# Think + img2text: reason before describing
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality img2text \
--image-path /path/to/image.jpg \
--prompts "Describe this image" \
--stage-configs-path ../../../vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
--prompts "What is happening in this image?" \
--think
# Text to Text with Mooncake
# Think + text2text: chain-of-thought reasoning
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2text \
--prompts "What is the capital of France?" \
--stage-configs-path ../../../vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
--prompts "Solve: 23 * 47" \
--think
```
For more details on the Mooncake connector and multi-node setup, see the [Mooncake Store Connector documentation](../../../docs/design/feature/omni_connectors/mooncake_store_connector.md).
Think mode parameters:
------
| Argument | Default | Description |
| :------- | :------ | :---------- |
| `--think` | `False` | Enable thinking mode |
| `--max-think-tokens` | `1000` | Maximum tokens for think generation |
| `--do-sample` | `False` | Enable sampling (vs. greedy) for text generation |
| `--text-temperature` | `0.3` | Temperature for text generation sampling |
## Classifier-Free Guidance (CFG)
CFG controls the trade-off between prompt fidelity and diversity. These parameters apply to image generation modalities (`text2img`, `img2img`).
```bash
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2img \
--prompts "A photorealistic portrait" \
--cfg-text-scale 6.0 \
--cfg-img-scale 2.0 \
--negative-prompt "blurry, low quality, distorted" \
--cfg-interval 0.4 1.0 \
--cfg-renorm-type global \
--cfg-renorm-min 0.0
```
| Argument | Default | Description |
| :------- | :------ | :---------- |
| `--cfg-text-scale` | `4.0` | Text CFG scale (higher = more prompt-adherent) |
| `--cfg-img-scale` | `1.5` | Image CFG scale (for img2img) |
| `--negative-prompt` | `None` | Negative prompt for CFG conditioning |
| `--cfg-interval` | pipeline default | CFG active interval `[start, end]` as fractions of total timesteps |
| `--cfg-renorm-type` | `None` | Renormalization type: `global`, `text_channel`, `channel` |
| `--cfg-renorm-min` | `None` | Minimum renormalization value |
| `--cfg-parallel-size` | `1` | CFG parallel size: `1` = batched (single GPU), `2` = 2-branch parallel, `3` = full 3-GPU parallel |
## Deployment Topologies
### Two-Stage (Default)
The default topology auto-detected from the model. No extra flags needed.
```bash
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2img \
--prompts "A cute cat"
```
The pipeline is defined in [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml). Stage 0 (Thinker) and Stage 1 (DiT) share GPU 0 by default. For dual-GPU setups, customize the deploy YAML and set `devices: "1"` for stage 1.
### Single-Stage
Pass the single-stage deploy config via `--deploy-config`:
```bash
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2img \
--prompts "A cute cat" \
--deploy-config vllm_omni/deploy/bagel_single_stage.yaml
```
See [`bagel_single_stage.yaml`](../../../vllm_omni/deploy/bagel_single_stage.yaml) for configuration details. The `pipeline: bagel_single_stage` field selects the single-stage topology from the pipeline registry.
### Tensor Parallelism (TP)
For larger models or multi-GPU environments, customize the deploy YAML (see [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml)) and set per-stage `tensor_parallel_size` and `devices`:
```yaml
# Example: TP=2 on GPUs 0,1 for the Thinker stage
stages:
- stage_id: 0
tensor_parallel_size: 2
devices: "0,1"
```
Then pass the custom deploy YAML:
```bash
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2img \
--prompts "A cute cat" \
--deploy-config /path/to/custom_bagel.yaml
```
### FP8 Quantization
```bash
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
--modality text2img \
--prompts "A cute cat" \
--quantization fp8
```
## Command Line Reference
### Core Arguments
| Argument | Type | Default | Description |
| :------- | :--- | :------ | :---------- |
| `--model` | string | `ByteDance-Seed/BAGEL-7B-MoT` | Model path or HuggingFace name |
| `--modality` | choice | `text2img` | `text2img`, `img2img`, `img2text`, `text2text` |
| `--prompts` | list | `None` | Input text prompts |
| `--txt-prompts` | string | `None` | Path to text file with one prompt per line |
| `--image-path` | string | `None` | Input image path (required for `img2img`/`img2text`) |
| `--output` | string | `.` | Output directory for saved images |
| `--steps` | int | `50` | Number of diffusion inference steps |
| `--seed` | int | `None` | Random seed for reproducibility |
### Think Mode Arguments
| Argument | Type | Default | Description |
| :------- | :--- | :------ | :---------- |
| `--think` | flag | `False` | Enable `<think>...</think>` planning/reasoning |
| `--max-think-tokens` | int | `1000` | Maximum tokens for think generation |
| `--do-sample` | flag | `False` | Use sampling instead of greedy decoding |
| `--text-temperature` | float | `0.3` | Sampling temperature for text generation |
### CFG Arguments
| Argument | Type | Default | Description |
| :------- | :--- | :------ | :---------- |
| `--cfg-text-scale` | float | `4.0` | Text CFG guidance scale |
| `--cfg-img-scale` | float | `1.5` | Image CFG guidance scale |
| `--negative-prompt` | string | `None` | Negative prompt for CFG |
| `--cfg-parallel-size` | int | `1` | CFG parallel GPU count (1, 2, or 3) |
| `--cfg-interval` | float[2] | pipeline default | CFG active window `[start, end]` |
| `--cfg-renorm-type` | string | `None` | `global`, `text_channel`, or `channel` |
| `--cfg-renorm-min` | float | `None` | Minimum renormalization value |
### Engine Arguments
| Argument | Type | Default | Description |
| :------- | :--- | :------ | :---------- |
| `--deploy-config` | string | `None` | Path to deploy YAML (auto-detected if omitted) |
| `--worker-backend` | choice | `process` | `process` or `ray` |
| `--ray-address` | string | `None` | Ray cluster address |
| `--quantization` | string | `None` | Quantization method (e.g. `fp8`) |
| `--log-stats` | flag | `False` | Enable statistics logging |
| `--init-timeout` | int | `300` | Initialization timeout (seconds) |
| `--batch-timeout` | int | `5` | Batch timeout (seconds) |
| `--enable-diffusion-pipeline-profiler` | flag | `False` | Profile diffusion stage durations |
## FAQ
- If you don’t know how much VRAM is needed for the model or encounter the OOM error, you can try to decrease the max_model_len.
- If you encounter OOM errors, try decreasing `max_model_len` or `gpu_memory_utilization` in the deploy YAML.
| Stage | VRAM |
| :------------------ | :--------------------------- |
| Stage-0 (Thinker) | **15.04 GiB** **+ KV Cache** |
| Stage-1 (DiT) | **26.50 GiB** |
| Total | **~42 GiB + KV Cache** |
**Two-stage VRAM usage:**
| Stage | VRAM |
| :---- | :--- |
| Stage 0 (Thinker) | **15.04 GiB + KV Cache** |
| Stage 1 (DiT) | **26.50 GiB** |
| Total | **~42 GiB + KV Cache** |
**Single-stage VRAM usage:** The DiT loads the full model (~42 GiB) in one process.
examples/offline_inference/bagel/end2end.py
+20 -12
View File
@@ -53,7 +53,12 @@ def parse_args():
parser.add_argument("--shm-threshold-bytes", type=int, default=65536)
parser.add_argument("--worker-backend", type=str, default="process", choices=["process", "ray"])
parser.add_argument("--ray-address", type=str, default=None)
parser.add_argument("--stage-configs-path", type=str, default=None)
parser.add_argument(
"--deploy-config",
type=str,
default=None,
help="Path to deploy YAML. If unset, auto-loads vllm_omni/deploy/bagel.yaml based on the HF model_type.",
)
parser.add_argument("--steps", type=int, default=50, help="Number of inference steps.")
parser.add_argument("--cfg-text-scale", type=float, default=4.0, help="Text CFG scale (default: 4.0)")
@@ -116,6 +121,9 @@ def parse_args():
help="Temperature for text generation sampling (default: 0.3).",
)
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
nullify_stage_engine_defaults(parser)
args = parser.parse_args()
return args
@@ -146,14 +154,12 @@ def main():
from vllm_omni.entrypoints.omni import Omni
omni_kwargs = {}
stage_configs_path = args.stage_configs_path
is_single_stage = stage_configs_path and "single_stage" in stage_configs_path
if args.think and stage_configs_path is None:
stage_configs_path = "vllm_omni/model_executor/stage_configs/bagel_think.yaml"
print(f"[Info] Think mode enabled, using stage config: {stage_configs_path}")
if stage_configs_path:
omni_kwargs["stage_configs_path"] = stage_configs_path
is_single_stage = "single_stage" in stage_configs_path
deploy_config = args.deploy_config
if args.think and deploy_config is None:
deploy_config = "vllm_omni/deploy/bagel_think.yaml"
print(f"[Info] Think mode enabled, using deploy config: {deploy_config}")
if deploy_config:
omni_kwargs["deploy_config"] = deploy_config
omni_kwargs.update(
{
@@ -170,7 +176,7 @@ def main():
if args.quantization:
omni_kwargs["quantization_config"] = args.quantization
omni = Omni(model=model_name, **omni_kwargs)
omni = Omni.from_cli_args(args, model=model_name, **omni_kwargs)
formatted_prompts = []
for p in prompts:
@@ -215,9 +221,11 @@ def main():
formatted_prompts.append(prompt_dict)
params_list = omni.default_sampling_params_list
# Bagel exposes 1 sampling param set for single-stage (DiT-only) and
# 2 for two-stage (Thinker + DiT). This heuristic may need updating
# if future pipelines break that 1:1 mapping.
is_single_stage = len(params_list) == 1
# For single-stage DiT, think/text params go into the diffusion sampling params extra_args.
# For 2-stage, diffusion params are at index 1.
diffusion_params_idx = 0 if is_single_stage else (1 if len(params_list) > 1 else 0)
diffusion_params = params_list[diffusion_params_idx]
examples/offline_inference/cosyvoice3/verify_e2e_cosyvoice.py
-1
View File
@@ -57,7 +57,6 @@ def run_e2e():
omni = Omni(
model=args.model,
deploy_config=args.deploy_config,
trust_remote_code=True,
tokenizer=args.tokenizer,
log_stats=True,
)
examples/offline_inference/custom_pipeline/image_to_image/image_edit.py
+23 -1
View File
@@ -44,9 +44,11 @@ Layered RGBA output:
import argparse
import asyncio
import json
import os
import time
from pathlib import Path
from typing import Any
import torch
from PIL import Image
@@ -58,6 +60,16 @@ from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.platforms import current_omni_platform
def parse_profiler_config(value: str) -> dict[str, Any]:
try:
config = json.loads(value)
except json.JSONDecodeError as e:
raise argparse.ArgumentTypeError(f"--profiler-config must be valid JSON: {e}") from e
if not isinstance(config, dict):
raise argparse.ArgumentTypeError("--profiler-config must be a JSON object")
return config
# ===========================
# Argument Parser
# ===========================
@@ -99,7 +111,16 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--vae-use-slicing", action="store_true")
parser.add_argument("--vae-use-tiling", action="store_true")
parser.add_argument("--enable-cpu-offload", action="store_true")
parser.add_argument(
"--profiler-config",
type=parse_profiler_config,
default=None,
help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
)
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
nullify_stage_engine_defaults(parser)
return parser.parse_args()
@@ -158,12 +179,13 @@ async def main():
enable_cpu_offload=args.enable_cpu_offload,
diffusion_load_format="dummy",
custom_pipeline_args={"pipeline_class": "custom_pipeline.CustomPipeline"},
profiler_config=args.profiler_config,
)
print(">>> Pipeline loaded successfully")
# ---- Profiling + Info ----
profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
profiler_enabled = args.profiler_config is not None
print(f"\n{'=' * 60}")
print("Generation Configuration")
print(f"Model: {args.model}")
examples/offline_inference/dynin_omni/end2end.py
+3
View File
@@ -970,6 +970,9 @@ def parse_args(repo_root: Path) -> argparse.Namespace:
parser.add_argument("--vq-model-audio-local-files-only", action=argparse.BooleanOptionalAction, default=None)
parser.add_argument("--disable-hf-xet", action=argparse.BooleanOptionalAction, default=True)
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
nullify_stage_engine_defaults(parser)
return parser.parse_args()
examples/offline_inference/glm_image/README.md
-145
View File
@@ -1,145 +0,0 @@
# GLM-Image Multistage End-to-End Inference
This example demonstrates how to run GLM-Image with the vLLM-Omni multistage architecture.
## Architecture
GLM-Image uses a 2-stage pipeline:
```
┌─────────────────────────────────────────────────────────────┐
│ GLM-Image Pipeline │
├─────────────────────────────────────────────────────────────┤
│ │
│ Stage 0 (AR Model) Stage 1 (Diffusion) │
│ ┌─────────────────┐ ┌─────────────────────┐ │
│ │ vLLM-optimized │ │ GlmImagePipeline │ │
│ │ GlmImageFor │ prior │ ┌───────────────┐ │ │
│ │ Conditional │──tokens───►│ │ DiT Denoiser │ │ │
│ │ Generation │ │ └───────────────┘ │ │
│ │ (9B AR model) │ │ │ │ │
│ └─────────────────┘ │ ▼ │ │
│ ▲ │ ┌───────────────┐ │ │
│ │ │ │ VAE Decode │──┼──► Image
│ Text/Image │ └───────────────┘ │ │
│ Input └─────────────────────┘ │
│ │
└─────────────────────────────────────────────────────────────┘
```
## Features
- **vLLM-optimized AR**: Uses PagedAttention and tensor parallelism for faster prior token generation
- **Flexible deployment**: AR and Diffusion stages can run on different GPUs
- **Text-to-Image**: Generate images from text descriptions
- **Image-to-Image**: Edit existing images with text prompts
## Usage
### Text-to-Image
```bash
python end2end.py \
--config-path ../../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
--prompt "A beautiful sunset over the ocean with sailing boats" \
--height 1024 \
--width 1024 \
--output output_t2i.png
```
### Image-to-Image (Image Editing)
```bash
python end2end.py \
--config-path ../../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
--prompt "Transform this scene into a winter wonderland" \
--image input.png \
--output output_i2i.png
```
### With Custom Parameters
```bash
python end2end.py \
--model-path /path/to/glm-image \
--config-path ../../../vllm_omni/model_executor/stage_configs/glm_image.yaml \
--prompt "A photorealistic cat sitting on a window sill" \
--height 1024 \
--width 1024 \
--num-inference-steps 50 \
--guidance-scale 1.5 \
--seed 42 \
--output output.png
```
## Shell Scripts
### Run Text-to-Image
```bash
./run_t2i.sh
```
### Run Image-to-Image
```bash
./run_i2i.sh --image /path/to/input.png
```
## Stage Configuration
The stage config (`glm_image.yaml`) defines:
- **Stage 0 (AR)**: Uses `GPUARWorker` with vLLM engine
- Model: `GlmImageForConditionalGeneration`
- Output: `token_ids` (prior tokens)
- **Stage 1 (Diffusion)**: Uses diffusion engine
- Model: `GlmImagePipeline`
- Output: Generated image
See `vllm_omni/model_executor/stage_configs/glm_image.yaml` for full configuration.
## Comparison with Single-Stage
| Aspect | Single-Stage (transformers) | Multistage (vLLM) |
| ----------- | --------------------------- | ------------------- |
| AR Model | transformers native | vLLM PagedAttention |
| Memory | Higher (no KV cache opt) | Lower (optimized) |
| Throughput | Lower | Higher |
| Flexibility | Single GPU | Multi-GPU support |
## Troubleshooting
### OOM Error
Try reducing memory usage:
```bash
# In glm_image.yaml, adjust:
gpu_memory_utilization: 0.5 # Reduce from 0.6
```
### Slow Initialization
The first run loads model weights. Subsequent runs are faster:
```bash
--stage-init-timeout 900 # Increase timeout for slow storage
```
### `Transformers does not recognize this architecture` Error
Your have to upgrade `transformers` package to `5.3.0` or above:
```
pip install --upgrade transformers
```
## Requirements
- vLLM-Omni with GLM-Image support
- CUDA-capable GPU (recommended: H100/A100 with 80GB)
- GLM-Image model weights
- `transformers` v5.3.0 or above
examples/offline_inference/glm_image/end2end.py
-525
View File
@@ -1,525 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
End-to-end offline inference example for GLM-Image with multistage architecture.
This script tests the multistage pipeline where:
- Stage 0 (AR): vLLM-optimized GlmImageForConditionalGeneration generates prior_token_ids
- Stage 1 (Diffusion): GlmImagePipeline performs DiT denoising + VAE decode
Usage (text-to-image):
python end2end.py \
--model-path /path/to/glm-image \
--config-path /path/to/glm_image.yaml \
--prompt "A beautiful sunset over the ocean" \
--output output_t2i.png
Usage (image-to-image / image edit):
python end2end.py \
--model-path /path/to/glm-image \
--config-path /path/to/glm_image.yaml \
--prompt "Make it look like winter" \
--image input.png \
--output output_i2i.png
Usage (with custom parameters):
python end2end.py \
--model-path /path/to/glm-image \
--config-path /path/to/glm_image.yaml \
--prompt "A cat sitting on a window sill" \
--height 1024 \
--width 1024 \
--num-inference-steps 50 \
--guidance-scale 1.5 \
--seed 42
For more options, run:
python end2end.py --help
"""
import argparse
import os
import time
from pathlib import Path
from PIL import Image
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
# Default stage config path (relative to vllm_omni package)
DEFAULT_CONFIG_PATH = "vllm_omni/model_executor/stage_configs/glm_image.yaml"
SEED = 42
# GLM-Image special tokens
GLM_IMAGE_EOS_TOKEN_ID = 16385 # eos_token_id from generation_config.json
GLM_IMAGE_VISION_VOCAB_SIZE = 16512 # top_k should be vision_vocab_size
def compute_max_tokens(height: int, width: int, factor: int = 32, is_i2i: bool = False) -> int:
"""
Compute max_new_tokens for GLM-Image AR generation.
GLM-Image generation differs by mode:
- text-to-image (t2i): small preview + large target + EOS
- image-to-image (i2i): large target + EOS
Args:
height: Target image height in pixels
width: Target image width in pixels
factor: Downsampling factor (32 for GLM-Image AR output)
is_i2i: Whether the request is image-to-image mode
Returns:
Total number of tokens to generate for the specified mode
"""
# Large image tokens (target resolution)
token_h = height // factor
token_w = width // factor
large_tokens = token_h * token_w
# Small preview tokens (half resolution in each dimension)
import math
ratio = token_h / token_w if token_w > 0 else 1.0
small_token_h = max(1, int(math.sqrt(ratio) * (factor // 2)))
small_token_w = max(1, int(math.sqrt(1 / ratio) * (factor // 2)))
small_tokens = small_token_h * small_token_w
if is_i2i:
return large_tokens + 1
return small_tokens + large_tokens + 1
def load_image(image_path: str) -> Image.Image:
"""Load an image from file path."""
if not os.path.exists(image_path):
raise FileNotFoundError(f"Image file not found: {image_path}")
return Image.open(image_path).convert("RGB")
def save_image(image: Image.Image, output_path: str) -> None:
"""Save an image to file path."""
output_dir = os.path.dirname(output_path)
if output_dir:
os.makedirs(output_dir, exist_ok=True)
image.save(output_path)
print(f"Image saved to: {output_path}")
def build_prompt_for_t2i(
prompt: str,
height: int = 1024,
width: int = 1024,
) -> dict:
"""
Build prompt dict for text-to-image generation.
Args:
prompt: Text description for image generation
height: Target image height
width: Target image width
Returns:
Dict containing prompt and generation parameters
"""
return {
"prompt": prompt,
"height": height,
"width": width,
# Pass target dimensions to AR processor for proper grid token generation
"mm_processor_kwargs": {
"target_h": height,
"target_w": width,
},
}
def build_prompt_for_i2i(
prompt: str,
image: Image.Image,
height: int | None = None,
width: int | None = None,
) -> dict:
"""
Build prompt dict for image-to-image generation.
Args:
prompt: Text description for image editing
image: Source image for editing
height: Target image height (default: use source image size)
width: Target image width (default: use source image size)
Returns:
Dict containing prompt, image, and generation parameters
"""
# Use source image dimensions if not specified
if height is None:
height = image.height
if width is None:
width = image.width
return {
"prompt": prompt,
"multi_modal_data": {
"image": image,
},
"height": height,
"width": width,
# Pass target dimensions to AR processor for proper grid token generation
"mm_processor_kwargs": {
"target_h": height,
"target_w": width,
},
}
def main(args: argparse.Namespace) -> None:
"""Main entry point for GLM-Image end-to-end inference."""
print("=" * 60)
print("GLM-Image Multistage End-to-End Inference")
print("=" * 60)
# Validate arguments
if not args.prompt:
raise ValueError("--prompt is required")
# Determine config path
config_path = args.config_path
if config_path is None:
# Try to find default config
if os.path.exists(DEFAULT_CONFIG_PATH):
config_path = DEFAULT_CONFIG_PATH
else:
# Try relative to script location
script_dir = Path(__file__).parent.parent.parent.parent
config_path = script_dir / "vllm_omni/model_executor/stage_configs/glm_image.yaml"
if not config_path.exists():
raise FileNotFoundError(
f"Stage config not found. Please specify --config-path. Tried: {DEFAULT_CONFIG_PATH}"
)
config_path = str(config_path)
print(f"Model path: {args.model_path}")
print(f"Config path: {config_path}")
print(f"Prompt: {args.prompt}")
# Load source image for image-to-image mode
source_image = None
if args.image:
print(f"Source image: {args.image}")
source_image = load_image(args.image)
print(f" Image size: {source_image.width}x{source_image.height}")
# Build prompt based on mode
if source_image is not None:
# Image-to-image mode
prompt_dict = build_prompt_for_i2i(
prompt=args.prompt,
image=source_image,
height=args.height,
width=args.width,
)
mode = "image-to-image"
else:
# Text-to-image mode
prompt_dict = build_prompt_for_t2i(
prompt=args.prompt,
height=args.height or 1024,
width=args.width or 1024,
)
mode = "text-to-image"
print(f"Mode: {mode}")
print(f"Target size: {prompt_dict.get('height', 1024)}x{prompt_dict.get('width', 1024)}")
# Add generation parameters to prompt
prompt_dict["seed"] = args.seed
prompt_dict["num_inference_steps"] = args.num_inference_steps
prompt_dict["guidance_scale"] = args.guidance_scale
if args.negative_prompt:
prompt_dict["negative_prompt"] = args.negative_prompt
# Build cache-dit config if requested
cache_config = None
if args.cache_backend == "cache_dit":
cache_config = {
"Fn_compute_blocks": 1,
"Bn_compute_blocks": 0,
"max_warmup_steps": 4,
"residual_diff_threshold": 0.24,
"max_continuous_cached_steps": 3,
"enable_taylorseer": False,
"taylorseer_order": 1,
"scm_steps_mask_policy": None,
"scm_steps_policy": "dynamic",
}
# Initialize Omni with multistage config
print("\nInitializing Omni with multistage pipeline...")
print(f" Cache backend: {args.cache_backend or 'None (no acceleration)'}")
start_time = time.time()
omni = Omni(
model=args.model_path,
stage_configs_path=config_path,
log_stats=args.enable_stats,
stage_init_timeout=args.stage_init_timeout,
cache_backend=args.cache_backend,
cache_config=cache_config,
enable_cache_dit_summary=getattr(args, "enable_cache_dit_summary", False),
enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
)
init_time = time.time() - start_time
print(f"Initialization completed in {init_time:.2f}s")
# Prepare prompts (support batch generation)
prompts = [prompt_dict for _ in range(args.num_prompts)]
# No explicit sampling_params for diffusion - parameters are in prompt_dict
# For multistage, the AR stage may need sampling params
from vllm import SamplingParams
# Compute max_tokens dynamically based on target image size
target_height = prompt_dict.get("height", 1024)
target_width = prompt_dict.get("width", 1024)
is_i2i = source_image is not None
calculated_max_tokens = compute_max_tokens(target_height, target_width, is_i2i=is_i2i)
# Use calculated value unless user explicitly specified a different value
# Default args.max_tokens is 16384 (very large), so prefer calculated value
effective_max_tokens = calculated_max_tokens if args.max_tokens == 16384 else args.max_tokens
if args.verbose:
print(
f"AR max_tokens: {effective_max_tokens} "
f"(calculated: {calculated_max_tokens}, arg: {args.max_tokens}, mode: {'i2i' if is_i2i else 't2i'})"
)
# IMPORTANT: GLM-Image AR model requires these exact sampling parameters
# from generation_config.json for proper image token generation.
# - temperature=0.9, top_p=0.75, top_k=16512 (vision_vocab_size)
# - stop_token_ids=[16385] (eos_token_id) is CRITICAL to stop generation
ar_sampling_params = SamplingParams(
temperature=0.9, # From generation_config.json
top_p=0.75, # From generation_config.json
top_k=GLM_IMAGE_VISION_VOCAB_SIZE, # 16512, vision vocabulary size
max_tokens=effective_max_tokens,
stop_token_ids=[GLM_IMAGE_EOS_TOKEN_ID], # 16385, CRITICAL for stopping
seed=args.seed,
detokenize=False,
# Keep target size available in runner/model for deterministic M-RoPE
# decode grids in t2i (no mm_features available in this path).
extra_args={
"target_h": int(target_height),
"target_w": int(target_width),
},
)
# For diffusion stage, sampling_params contains diffusion-specific parameters
# These are passed as kwargs to the diffusion engine
diffusion_sampling_params = OmniDiffusionSamplingParams(
num_inference_steps=args.num_inference_steps,
guidance_scale=args.guidance_scale,
height=prompt_dict.get("height", 1024),
width=prompt_dict.get("width", 1024),
seed=args.seed,
)
# For multistage, we need sampling_params for each stage
# Stage 0 (AR): SamplingParams for vLLM
# Stage 1 (Diffusion): dict with diffusion kwargs
sampling_params_list = [ar_sampling_params, diffusion_sampling_params]
# Run generation
print(f"\nGenerating {args.num_prompts} image(s)...")
gen_start_time = time.time()
output_dir = os.path.dirname(args.output) if args.output else "outputs"
if output_dir:
os.makedirs(output_dir, exist_ok=True)
output_count = 0
for stage_outputs in omni.generate(prompts, sampling_params_list, py_generator=True):
output = stage_outputs.request_output
if stage_outputs.final_output_type == "image":
request_id = output.request_id
# Get generated images
images = output.images if hasattr(output, "images") else []
if not images and hasattr(output, "multimodal_output"):
images = output.multimodal_output.get("images", [])
# Save each generated image
for idx, img in enumerate(images):
if args.num_prompts == 1 and len(images) == 1:
output_path = args.output
else:
base, ext = os.path.splitext(args.output)
output_path = f"{base}_{request_id}_{idx}{ext}"
if isinstance(img, Image.Image):
save_image(img, output_path)
else:
print(f"Warning: Unexpected image type for request {request_id}: {type(img)}")
output_count += 1
elif stage_outputs.final_output_type == "text":
# AR stage output (intermediate, for debugging)
if args.verbose:
print(f"AR output for request {output.request_id}:")
print(f" Token count: {len(output.outputs[0].token_ids)}")
gen_time = time.time() - gen_start_time
print(f"\nGeneration completed in {gen_time:.2f}s")
print(f"Generated {output_count} image(s)")
# Cleanup
omni.close()
print("\nDone!")
def parse_args() -> argparse.Namespace:
"""Parse command line arguments."""
parser = argparse.ArgumentParser(
description="GLM-Image Multistage End-to-End Inference",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog=__doc__,
)
# Required arguments
parser.add_argument(
"--model-path",
type=str,
default="zai-org/GLM-Image",
help="Path to GLM-Image model directory or HuggingFace model ID",
)
parser.add_argument(
"--prompt",
type=str,
required=True,
help="Text prompt for image generation",
)
# Optional arguments
parser.add_argument(
"--config-path",
type=str,
default=None,
help="Path to stage config YAML file (default: auto-detect)",
)
parser.add_argument(
"--image",
type=str,
default=None,
help="Path to source image for image-to-image mode",
)
parser.add_argument(
"--output",
type=str,
default="output_glm_image.png",
help="Output image path (default: output_glm_image.png)",
)
parser.add_argument(
"--negative-prompt",
type=str,
default=None,
help="Negative prompt for classifier-free guidance",
)
# Generation parameters
parser.add_argument(
"--height",
type=int,
default=None,
help="Output image height (default: 1024 for t2i, source size for i2i)",
)
parser.add_argument(
"--width",
type=int,
default=None,
help="Output image width (default: 1024 for t2i, source size for i2i)",
)
parser.add_argument(
"--num-inference-steps",
type=int,
default=50,
help="Number of diffusion denoising steps (default: 50)",
)
parser.add_argument(
"--guidance-scale",
type=float,
default=1.5,
help="Classifier-free guidance scale (default: 1.5)",
)
parser.add_argument(
"--seed",
type=int,
default=SEED,
help=f"Random seed for reproducibility (default: {SEED})",
)
parser.add_argument(
"--max-tokens",
type=int,
default=16384,
help="Maximum tokens for AR generation (default: 16384)",
)
# Batch processing
parser.add_argument(
"--num-prompts",
type=int,
default=1,
help="Number of images to generate (default: 1)",
)
# Cache acceleration
parser.add_argument(
"--cache-backend",
type=str,
default=None,
choices=["cache_dit"],
help="Cache backend for DiT acceleration. Default: None (no cache).",
)
parser.add_argument(
"--enable-cache-dit-summary",
action="store_true",
help="Enable cache-dit summary logging after diffusion forward passes.",
)
# Runtime options
parser.add_argument(
"--enable-stats",
action="store_true",
default=False,
help="Enable statistics logging",
)
parser.add_argument(
"--stage-init-timeout",
type=int,
default=600,
help="Timeout for stage initialization in seconds (default: 600)",
)
parser.add_argument(
"--verbose",
"-v",
action="store_true",
default=False,
help="Enable verbose output",
)
parser.add_argument(
"--enable-diffusion-pipeline-profiler",
action="store_true",
help="Enable diffusion pipeline profiler to display stage durations.",
)
return parser.parse_args()
if __name__ == "__main__":
args = parse_args()
main(args)
examples/offline_inference/glm_image/run_i2i.sh
-93
View File
@@ -1,93 +0,0 @@
#!/bin/bash
# SPDX-License-Identifier: Apache-2.0
# Run GLM-Image image-to-image (editing) with multistage pipeline
set -e
# Default values
MODEL_PATH="${MODEL_PATH:-/path/to/glm-image}"
CONFIG_PATH="${CONFIG_PATH:-vllm_omni/model_executor/stage_configs/glm_image.yaml}"
PROMPT="${PROMPT:-Transform this image into an oil painting style}"
INPUT_IMAGE=""
OUTPUT="${OUTPUT:-output_i2i.png}"
NUM_STEPS="${NUM_STEPS:-50}"
GUIDANCE="${GUIDANCE:-1.5}"
SEED="${SEED:-42}"
# Parse command line arguments
while [[ $# -gt 0 ]]; do
case $1 in
--model-path)
MODEL_PATH="$2"
shift 2
;;
--config-path)
CONFIG_PATH="$2"
shift 2
;;
--prompt)
PROMPT="$2"
shift 2
;;
--image)
INPUT_IMAGE="$2"
shift 2
;;
--output)
OUTPUT="$2"
shift 2
;;
--num-steps)
NUM_STEPS="$2"
shift 2
;;
--guidance)
GUIDANCE="$2"
shift 2
;;
--seed)
SEED="$2"
shift 2
;;
*)
echo "Unknown option: $1"
exit 1
;;
esac
done
# Check if input image is provided
if [ -z "${INPUT_IMAGE}" ]; then
echo "Error: --image is required for image-to-image mode"
echo "Usage: ./run_i2i.sh --image /path/to/input.png [--prompt \"edit instruction\"]"
exit 1
fi
if [ ! -f "${INPUT_IMAGE}" ]; then
echo "Error: Input image not found: ${INPUT_IMAGE}"
exit 1
fi
echo "=============================================="
echo "GLM-Image Image-to-Image Generation"
echo "=============================================="
echo "Model: ${MODEL_PATH}"
echo "Config: ${CONFIG_PATH}"
echo "Input: ${INPUT_IMAGE}"
echo "Prompt: ${PROMPT}"
echo "Output: ${OUTPUT}"
echo "Steps: ${NUM_STEPS}"
echo "Guidance: ${GUIDANCE}"
echo "Seed: ${SEED}"
echo "=============================================="
python end2end.py \
--model-path "${MODEL_PATH}" \
--config-path "${CONFIG_PATH}" \
--prompt "${PROMPT}" \
--image "${INPUT_IMAGE}" \
--output "${OUTPUT}" \
--num-inference-steps "${NUM_STEPS}" \
--guidance-scale "${GUIDANCE}" \
--seed "${SEED}" \
--verbose
examples/offline_inference/glm_image/run_t2i.sh
-87
View File
@@ -1,87 +0,0 @@
#!/bin/bash
# SPDX-License-Identifier: Apache-2.0
# Run GLM-Image text-to-image generation with multistage pipeline
set -e
# Default values
MODEL_PATH="${MODEL_PATH:-/path/to/glm-image}"
CONFIG_PATH="${CONFIG_PATH:-vllm_omni/model_executor/stage_configs/glm_image.yaml}"
PROMPT="${PROMPT:-A beautiful mountain landscape with snow-capped peaks and a clear blue lake}"
OUTPUT="${OUTPUT:-output_t2i.png}"
HEIGHT="${HEIGHT:-1024}"
WIDTH="${WIDTH:-1024}"
NUM_STEPS="${NUM_STEPS:-50}"
GUIDANCE="${GUIDANCE:-1.5}"
SEED="${SEED:-42}"
# Parse command line arguments
while [[ $# -gt 0 ]]; do
case $1 in
--model-path)
MODEL_PATH="$2"
shift 2
;;
--config-path)
CONFIG_PATH="$2"
shift 2
;;
--prompt)
PROMPT="$2"
shift 2
;;
--output)
OUTPUT="$2"
shift 2
;;
--height)
HEIGHT="$2"
shift 2
;;
--width)
WIDTH="$2"
shift 2
;;
--num-steps)
NUM_STEPS="$2"
shift 2
;;
--guidance)
GUIDANCE="$2"
shift 2
;;
--seed)
SEED="$2"
shift 2
;;
*)
echo "Unknown option: $1"
exit 1
;;
esac
done
echo "=============================================="
echo "GLM-Image Text-to-Image Generation"
echo "=============================================="
echo "Model: ${MODEL_PATH}"
echo "Config: ${CONFIG_PATH}"
echo "Prompt: ${PROMPT}"
echo "Output: ${OUTPUT}"
echo "Size: ${WIDTH}x${HEIGHT}"
echo "Steps: ${NUM_STEPS}"
echo "Guidance: ${GUIDANCE}"
echo "Seed: ${SEED}"
echo "=============================================="
python end2end.py \
--model-path "${MODEL_PATH}" \
--config-path "${CONFIG_PATH}" \
--prompt "${PROMPT}" \
--output "${OUTPUT}" \
--height "${HEIGHT}" \
--width "${WIDTH}" \
--num-inference-steps "${NUM_STEPS}" \
--guidance-scale "${GUIDANCE}" \
--seed "${SEED}" \
--verbose
examples/offline_inference/helios/end2end.py
+3
View File
@@ -196,6 +196,9 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--cfg-parallel-size", type=int, default=1, choices=[1, 2], help="CFG parallel size.")
parser.add_argument("--tensor-parallel-size", type=int, default=1, help="Tensor parallelism size.")
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
nullify_stage_engine_defaults(parser)
return parser.parse_args()
examples/offline_inference/hunyuan_image3/end2end.py
+3
View File
@@ -135,6 +135,9 @@ def parse_args():
parser.add_argument("--init-timeout", type=int, default=300, help="Initialization timeout in seconds.")
parser.add_argument("--enforce-eager", action="store_true", help="Disable torch.compile.")
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
nullify_stage_engine_defaults(parser)
return parser.parse_args()
examples/offline_inference/image_to_image/image_edit.py
+53 -2
View File
@@ -87,9 +87,11 @@ For more options, run:
"""
import argparse
import json
import os
import time
from pathlib import Path
from typing import Any
import torch
from PIL import Image
@@ -101,6 +103,16 @@ from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.platforms import current_omni_platform
def parse_profiler_config(value: str) -> dict[str, Any]:
try:
config = json.loads(value)
except json.JSONDecodeError as e:
raise argparse.ArgumentTypeError(f"--profiler-config must be valid JSON: {e}") from e
if not isinstance(config, dict):
raise argparse.ArgumentTypeError("--profiler-config must be a JSON object")
return config
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Edit an image with Qwen-Image-Edit.")
parser.add_argument(
@@ -196,6 +208,13 @@ def parse_args() -> argparse.Namespace:
default=1,
help="Number of GPUs used for ulysses sequence parallelism.",
)
parser.add_argument(
"--ulysses-mode",
type=str,
default="strict",
choices=["strict", "advanced_uaa"],
help="Ulysses sequence-parallel mode: 'strict' (divisibility required) or 'advanced_uaa' (UAA).",
)
parser.add_argument(
"--ring-degree",
type=int,
@@ -325,11 +344,43 @@ def parse_args() -> argparse.Namespace:
action="store_true",
help="Enable layerwise (blockwise) offloading on DiT modules.",
)
parser.add_argument(
"--vae-patch-parallel-size",
type=int,
default=1,
help="Number of GPUs used for VAE patch/tile parallelism (decode).",
)
parser.add_argument(
"--use-hsdp",
action="store_true",
help="Enable HSDP (Hybrid Sharded Data Parallel) for diffusion models.",
)
parser.add_argument(
"--hsdp-shard-size",
type=int,
default=1,
help="Number of GPUs to shard weights across for HSDP.",
)
parser.add_argument(
"--hsdp-replicate-size",
type=int,
default=1,
help="Number of HSDP replica groups.",
)
parser.add_argument(
"--enable-diffusion-pipeline-profiler",
action="store_true",
help="Enable diffusion pipeline profiler to display stage durations.",
)
parser.add_argument(
"--profiler-config",
type=parse_profiler_config,
default=None,
help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
)
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
nullify_stage_engine_defaults(parser)
return parser.parse_args()
@@ -395,11 +446,11 @@ def main():
enforce_eager=args.enforce_eager,
enable_cpu_offload=args.enable_cpu_offload,
enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
profiler_config=args.profiler_config,
)
print("Pipeline loaded")
# Check if profiling is requested via environment variable
profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
profiler_enabled = args.profiler_config is not None
# Time profiling for generation
print(f"\n{'=' * 60}")
examples/offline_inference/image_to_video/image_to_video.py
+23 -3
View File
@@ -33,9 +33,10 @@ Usage:
"""
import argparse
import os
import json
import time
from pathlib import Path
from typing import Any
import numpy as np
import PIL.Image
@@ -48,6 +49,16 @@ from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.platforms import current_omni_platform
def parse_profiler_config(value: str) -> dict[str, Any]:
try:
config = json.loads(value)
except json.JSONDecodeError as e:
raise argparse.ArgumentTypeError(f"--profiler-config must be valid JSON: {e}") from e
if not isinstance(config, dict):
raise argparse.ArgumentTypeError("--profiler-config must be a JSON object")
return config
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Generate a video from an image (Wan2.2, LTX2, HunyuanVideo-1.5).")
parser.add_argument(
@@ -194,6 +205,15 @@ def parse_args() -> argparse.Namespace:
action="store_true",
help="Enable diffusion pipeline profiler to display stage durations.",
)
parser.add_argument(
"--profiler-config",
type=parse_profiler_config,
default=None,
help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
)
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
nullify_stage_engine_defaults(parser)
return parser.parse_args()
@@ -274,8 +294,7 @@ def main():
"rel_l1_thresh": 0.2,
}
# Check if profiling is requested via environment variable
profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
profiler_enabled = args.profiler_config is not None
parallel_config = DiffusionParallelConfig(
ulysses_degree=args.ulysses_degree,
ring_degree=args.ring_degree,
@@ -300,6 +319,7 @@ def main():
cache_backend=args.cache_backend,
cache_config=cache_config,
enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
profiler_config=args.profiler_config,
)
if profiler_enabled:
examples/offline_inference/magi_human/end2end.py
+3
View File
@@ -24,6 +24,9 @@ def parse_args():
parser.add_argument("--width", type=int, default=448, help="Video width.")
parser.add_argument("--num-inference-steps", type=int, default=8, help="Number of denoising steps.")
parser.add_argument("--seed", type=int, default=52, help="Random seed for generation.")
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
nullify_stage_engine_defaults(parser)
return parser.parse_args()
examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py
+2
View File
@@ -19,6 +19,7 @@ from vllm import SamplingParams
from vllm.multimodal.image import convert_image_mode
from vllm_omni import Omni
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
DEFAULT_SYSTEM = "You are a helpful assistant."
DEFAULT_QUESTION = "Please summarize the content of this image."
@@ -48,6 +49,7 @@ def parse_args() -> argparse.Namespace:
action="store_true",
help="Enable diffusion pipeline profiler to display stage durations.",
)
nullify_stage_engine_defaults(parser)
return parser.parse_args()
examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py
+2
View File
@@ -29,6 +29,7 @@ from PIL import Image
from vllm.sampling_params import SamplingParams
from vllm_omni import Omni
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)
@@ -117,6 +118,7 @@ def parse_args() -> argparse.Namespace:
)
p.add_argument("--out", type=str, default="output.png", help="Path to save the generated image.")
p.add_argument("--trust-remote-code", action="store_true", help="Trust remote code when loading the model.")
nullify_stage_engine_defaults(p)
args = p.parse_args()
if not args.prompt:
args.prompt = ["A stylish woman with sunglasses riding a motorcycle in NYC."]
examples/offline_inference/mimo_audio/end2end.py
+5 -2
View File
@@ -309,7 +309,10 @@ def main(args):
lines.append("Prompt:\n")
lines.append(str(prompt_text) + "\n")
lines.append("vllm_text_output:\n")
lines.append(str(text_output).strip() + "\n")
output_text = str(text_output)
if "<chinese>" in output_text or "<english>" in output_text:
output_text = output_text.replace("<chinese>", "").replace("<english>", "").strip()
lines.append(output_text + "\n")
try:
with open(out_txt, "w", encoding="utf-8") as f:
print("lines", lines)
@@ -351,7 +354,7 @@ def parse_args():
"--text",
"-t",
type=str,
default="The weather is so nice today.",
default="",
help="input text",
)
parser.add_argument(
examples/offline_inference/ming_flash_omni/README.md
+57 -41
View File
@@ -1,75 +1,91 @@
# Ming-flash-omni 2.0
[Ming-flash-omni-2.0](https://github.com/inclusionAI/Ming) is an omni-modal model supporting text, image, video, and audio understanding, with outputs in text, image, and audio. For now, Ming-flash-omni-2.0 in vLLM-Omni is supported with thinker stage (multi-modal understanding).
[Ming-flash-omni-2.0](https://github.com/inclusionAI/Ming) is an omni-modal model supporting text, image, video, and audio understanding, with text and speech outputs.
vLLM-Omni supports two deployment modes:
| Mode | Stage config | Output |
|------|-------------|--------|
| Thinker only (multimodal understanding) | `ming_flash_omni_thinker.yaml` (default `--omni`) | Text |
| Thinker + Talker (omni-speech) | `ming_flash_omni.yaml` | Text + Audio |
For standalone TTS (talker only), see [`examples/offline_inference/ming_flash_omni_tts/`](../ming_flash_omni_tts/).
## 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.
The default `--omni` flag runs thinker only. For omni-speech, pass the two-stage config explicitly:
```bash
--stage-configs-path vllm_omni/model_executor/stage_configs/ming_flash_omni.yaml
```
## Run examples
### Text-only
The end-to-end script defaults to built-in assets; pass `--image-path`,
`--audio-path`, or `--video-path` to override.
```bash
# Text-only
python examples/offline_inference/ming_flash_omni/end2end.py --query-type text
# Image / audio / video / mixed understanding
python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_image
python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_audio
python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_video --num-frames 16
python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_mixed_modalities \
--image-path /path/to/image.jpg --audio-path /path/to/audio.wav
```
#### Reasoning (Thinking Mode)
Reasoning (Thinking) mode is enabled via applying "detailed thinking on" when building the system prompt template (in `apply_chat_template`).
In the end2end example, a default problem for thinking mode is provided, as referred to the example usage of Ming's cookbook;
To utilize it, you have to download the example figure from https://github.com/inclusionAI/Ming/blob/3954fcb880ff5e61ff128bcf7f1ec344d46a6fe3/figures/cases/3_0.png
Reasoning ("detailed thinking on") is applied by the script when
`--query-type reasoning` is set. The default prompt matches Ming's cookbook
and expects the reference figure from the upstream repo — see
`get_reasoning_query` in `end2end.py`.
```bash
python examples/offline_inference/ming_flash_omni/end2end.py -q reasoning --image-path ./3_0.png
```
### Image understanding
```bash
python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_image
### Omni-speech (thinker + talker)
# With a local image
python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_image --image-path /path/to/image.jpg
To enable spoken output, use the two-stage config and request `audio` (or `text,audio`) modalities.
The thinker processes your multimodal input, generates text, then the talker synthesises the response as speech.
**Audio-only output** (speech response, no text):
```bash
python examples/offline_inference/ming_flash_omni/end2end.py \
--query-type text \
--stage-configs-path vllm_omni/model_executor/stage_configs/ming_flash_omni.yaml \
--modalities audio \
--output-dir output_ming_omni_speech
```
### Audio understanding
**Both text and audio output**:
```bash
python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_audio
# With a local audio file
python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_audio --audio-path /path/to/audio.wav
python examples/offline_inference/ming_flash_omni/end2end.py \
--query-type use_audio \
--stage-configs-path vllm_omni/model_executor/stage_configs/ming_flash_omni.yaml \
--modalities text,audio \
--output-dir output_ming_omni_speech
```
### Video understanding
```bash
python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_video
Generated `.wav` files are saved to `--output-dir` (default `output_ming`), one per request.
# With a local video and custom frame count
python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_video --video-path /path/to/video.mp4 --num-frames 16
```
### Mixed modalities (image + audio)
```bash
python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_mixed_modalities \
--image-path /path/to/image.jpg \
--audio-path /path/to/audio.wav
```
If media file paths are not provided, the script uses built-in default assets.
The stage config allocates thinker on GPUs 0–3 and talker on GPU 3 by default. Adjust `devices` in the YAML to match your hardware.
### Modality control
To control output modalities (e.g. text-only output):
```bash
python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_audio --modalities text
```
*For now, only text output is supported*
| `--modalities` | Thinker output | Talker | Saved files |
|---------------|----------------|--------|-------------|
| `text` (default) | Text | Not run | `<id>.txt` |
| `audio` | Text (internal) | Runs | `<id>.wav` |
| `text,audio` | Text | Runs | `<id>.txt` + `<id>.wav` |
### Custom stage config
```bash
python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_image \
--stage-configs-path /path/to/your_config.yaml
```
Pass `--stage-configs-path /path/to/your_config.yaml` to any of the commands
above to override the stage config.
## Online serving
examples/offline_inference/ming_flash_omni/end2end.py
+24 -2
View File
@@ -7,6 +7,7 @@ import time
from typing import NamedTuple
import numpy as np
import soundfile as sf
import vllm
from PIL import Image
from transformers import AutoProcessor
@@ -319,7 +320,16 @@ def main(args):
seed=SEED,
detokenize=True,
)
sampling_params_list = [thinker_sampling_params]
# Talker (ming_tts) uses a custom generation loop (CFM + AudioVAE);
# vLLM sampling is a no-op here — max_tokens=1 just satisfies the scheduler.
talker_sampling_params = SamplingParams(
temperature=0.0,
max_tokens=1,
)
all_sampling_params = [thinker_sampling_params, talker_sampling_params]
# Match sampling params to the number of configured stages
# (thinker-only yaml → 1, thinker+talker yaml → 2).
sampling_params_list = all_sampling_params[: omni.num_stages]
prompts = [query_result.inputs for _ in range(args.num_prompts)]
@@ -362,7 +372,19 @@ def main(args):
print(f"Failed to write output file {out_txt}: {e}")
elif stage_outputs.final_output_type == "audio":
raise NotImplementedError("Add audio example after talker supported.")
request_id = output.request_id
mm = output.outputs[0].multimodal_output
if mm and "audio" in mm:
audio = mm["audio"]
sr_raw = mm.get("sr", 44100)
sample_rate = int(sr_raw.item() if hasattr(sr_raw, "item") else sr_raw)
audio_numpy = audio.float().squeeze().cpu().numpy()
output_wav = os.path.join(output_dir, f"{request_id}.wav")
sf.write(output_wav, audio_numpy, samplerate=sample_rate, format="WAV")
print(
f"Request ID: {request_id}, audio saved to {output_wav} "
f"({len(audio_numpy) / sample_rate:.2f}s, {sample_rate}Hz)"
)
processed_count += 1
if profiler_enabled and processed_count >= total_requests:
examples/offline_inference/ming_flash_omni_tts/README.md
+47
View File
@@ -0,0 +1,47 @@
# Ming-flash-omni Standalone TTS (Offline)
This example runs **Ming-flash-omni-2.0 talker-only** offline inference with:
- `model`: `Jonathan1909/Ming-flash-omni-2.0`
- `stage config`: `vllm_omni/model_executor/stage_configs/ming_flash_omni_tts.yaml`
It follows the Ming cookbook parameter style:
- `prompt`: `"Please generate speech based on the following description.\n"`
- `max_decode_steps`: `200`
- `cfg`: `2.0`
- `sigma`: `0.25`
- `temperature`: `0.0`
## Quick Start
```bash
python examples/offline_inference/ming_flash_omni_tts/end2end.py --case style
```
## Cases
```bash
# Style
python examples/offline_inference/ming_flash_omni_tts/end2end.py --case style
# IP
python examples/offline_inference/ming_flash_omni_tts/end2end.py --case ip
# Basic (speed/pitch/volume control)
python examples/offline_inference/ming_flash_omni_tts/end2end.py --case basic
```
## Useful Arguments
- `--text`: override default text in the selected case
- `--output`: custom output wav path
- `--model`: local model path or HF repo id
- `--stage-configs-path`: custom talker stage config path
- `--log-stats`: enable runtime stats logs
## Notes
- This directory is for **standalone talker deployment (TTS)**.
- For Ming thinker multimodal understanding examples, see:
`examples/offline_inference/ming_flash_omni/`.
examples/offline_inference/ming_flash_omni_tts/end2end.py
+128
View File
@@ -0,0 +1,128 @@
"""Offline e2e example for Ming-flash-omni-2.0 standalone talker (TTS)."""
import os
from typing import Any
import soundfile as sf
import torch
from vllm.utils.argparse_utils import FlexibleArgumentParser
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniTokensPrompt
from vllm_omni.model_executor.models.ming_flash_omni.prompt_utils import (
DEFAULT_PROMPT,
create_instruction,
)
MODEL_NAME = "Jonathan1909/Ming-flash-omni-2.0"
DEFAULT_STAGE_CONFIG = "vllm_omni/model_executor/stage_configs/ming_flash_omni_tts.yaml"
def get_messages(case: str, text_override: str | None) -> dict[str, Any]:
if case == "style":
text = text_override or "我会一直在这里陪着你，直到你慢慢、慢慢地沉入那个最温柔的梦里……好吗？"
instruction = create_instruction(
{
"风格": "这是一种ASMR耳语，属于一种旨在引发特殊感官体验的创意风格。这个女性使用轻柔的普通话进行耳语，声音气音成分重。音量极低，紧贴麦克风，语速极慢，旨在制造触发听者颅内快感的声学刺激。",
}
)
return {
"prompt": DEFAULT_PROMPT,
"text": text,
"instruction": instruction,
"use_zero_spk_emb": True,
}
if case == "ip":
text = text_override or "这款产品的名字，叫变态坑爹牛肉丸。"
return {
"prompt": DEFAULT_PROMPT,
"text": text,
"instruction": create_instruction({"IP": "灵小甄"}),
"use_zero_spk_emb": True,
}
if case == "basic":
text = text_override or "我们当迎着阳光辛勤耕作，去摘取，去制作，去品尝，去馈赠。"
return {
"prompt": DEFAULT_PROMPT,
"text": text,
"instruction": create_instruction({"语速": "快速", "基频": "中", "音量": "中"}),
"use_zero_spk_emb": True,
}
raise ValueError(f"Unknown case: {case}")
def save_audio(mm: dict[str, Any], output_path: str) -> None:
if not mm or "audio" not in mm:
raise RuntimeError("No audio found in model output")
audio = mm["audio"]
sr_raw = mm.get("sr", 44100)
if isinstance(sr_raw, torch.Tensor):
sample_rate = int(sr_raw.item())
else:
sample_rate = int(sr_raw)
waveform = audio.squeeze().float().cpu().numpy()
sf.write(output_path, waveform, sample_rate)
print(f"Saved {output_path} ({len(waveform) / sample_rate:.2f}s, {sample_rate}Hz)")
def parse_args():
parser = FlexibleArgumentParser(description="Ming-flash-omni standalone talker offline e2e example")
parser.add_argument("--model", type=str, default=MODEL_NAME, help="Model name or local path.")
parser.add_argument(
"--stage-configs-path",
type=str,
default=DEFAULT_STAGE_CONFIG,
help="Path to stage configs yaml for standalone talker deployment.",
)
parser.add_argument(
"--case",
type=str,
default="style",
choices=["style", "ip", "basic"],
help="Example case.",
)
parser.add_argument("--text", type=str, default=None, help="Override default text for the selected case.")
parser.add_argument("--output", type=str, default=None, help="Output wav path.")
parser.add_argument("--log-stats", action="store_true", default=False, help="Enable stats logging.")
parser.add_argument("--init-timeout", type=int, default=600, help="Engine init timeout in seconds.")
parser.add_argument("--stage-init-timeout", type=int, default=300, help="Single stage init timeout in seconds.")
return parser.parse_args()
def main():
args = parse_args()
omni = Omni(
model=args.model,
stage_configs_path=args.stage_configs_path,
log_stats=args.log_stats,
init_timeout=args.init_timeout,
stage_init_timeout=args.stage_init_timeout,
)
messages = get_messages(args.case, args.text)
decode_args = {
# Standalone TTS deployment
"ming_task": "instruct",
"max_decode_steps": 200,
"cfg": 2.0,
"sigma": 0.25,
"temperature": 0.0,
}
req = OmniTokensPrompt(
prompt_token_ids=[0],
additional_information={**messages, **decode_args},
)
outputs = omni.generate(req)
mm = outputs[0].outputs[0].multimodal_output
output_path = args.output or f"output_{args.case}.wav"
save_audio(mm, output_path)
omni.close()
if __name__ == "__main__":
main()
examples/offline_inference/omnivoice/end2end.py
-1
View File
@@ -89,7 +89,6 @@ def run_e2e():
omni = Omni(
model=args.model,
stage_configs_path=args.stage_config,
trust_remote_code=True,
log_stats=True,
)
examples/offline_inference/qwen2_5_omni/end2end.py
+18 -1
View File
@@ -5,6 +5,7 @@ This example shows how to use vLLM-Omni for running offline inference
with the correct prompt format on Qwen2.5-Omni
"""
import json
import os
import time
from typing import NamedTuple
@@ -289,7 +290,10 @@ query_map = {
def main(args):
model_name = "Qwen/Qwen2.5-Omni-7B"
model_name = args.model
quantization_config = None
if args.quantization_config is not None:
quantization_config = json.loads(args.quantization_config)
# Get paths from args
video_path = getattr(args, "video_path", None)
@@ -320,6 +324,7 @@ def main(args):
query_result = query_func(audio_path=audio_path, sampling_rate=sampling_rate)
else:
query_result = query_func()
args.quantization_config = quantization_config
omni = Omni.from_cli_args(args, model=model_name)
thinker_sampling_params = SamplingParams(
temperature=0.0, # Deterministic - no randomness
@@ -424,6 +429,18 @@ def main(args):
def parse_args():
parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
parser.add_argument(
"--model",
type=str,
default="Qwen/Qwen2.5-Omni-7B",
help="Model name or local path.",
)
parser.add_argument(
"--quantization-config",
type=str,
default=None,
help="Optional JSON string forwarded to Omni(quantization_config=...).",
)
parser.add_argument(
"--query-type",
"-q",
examples/offline_inference/qwen3_omni/end2end.py
+2
View File
@@ -21,6 +21,7 @@ from vllm.multimodal.image import convert_image_mode
from vllm.multimodal.media.audio import load_audio
from vllm.utils.argparse_utils import FlexibleArgumentParser
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
from vllm_omni.entrypoints.omni import Omni
SEED = 42
@@ -550,6 +551,7 @@ def parse_args():
help="Model dtype (auto, half, float16, bfloat16, float, float32).",
)
nullify_stage_engine_defaults(parser)
return parser.parse_args()
examples/offline_inference/text_to_audio/text_to_audio.py
+38
View File
@@ -131,6 +131,44 @@ def parse_args() -> argparse.Namespace:
default=1,
help="Number of HSDP replica groups. Default 1 means pure sharding.",
)
parser.add_argument(
"--tensor-parallel-size",
type=int,
default=1,
help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
)
parser.add_argument(
"--ulysses-degree",
type=int,
default=1,
help="Number of GPUs used for ulysses sequence parallelism.",
)
parser.add_argument(
"--ulysses-mode",
type=str,
default="strict",
choices=["strict", "advanced_uaa"],
help="Ulysses sequence-parallel mode: 'strict' (divisibility required) or 'advanced_uaa' (UAA).",
)
parser.add_argument(
"--ring-degree",
type=int,
default=1,
help="Number of GPUs used for ring sequence parallelism.",
)
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(
"--vae-patch-parallel-size",
type=int,
default=1,
help="Number of GPUs used for VAE patch/tile parallelism (decode).",
)
return parser.parse_args()
examples/offline_inference/text_to_image/text_to_image.py
+39 -3
View File
@@ -2,7 +2,7 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import argparse
import os
import json
import time
from pathlib import Path
from typing import Any
@@ -30,6 +30,16 @@ def is_nextstep_model(model_name: str) -> bool:
return False
def parse_profiler_config(value: str) -> dict[str, Any]:
try:
config = json.loads(value)
except json.JSONDecodeError as e:
raise argparse.ArgumentTypeError(f"--profiler-config must be valid JSON: {e}") from e
if not isinstance(config, dict):
raise argparse.ArgumentTypeError("--profiler-config must be a JSON object")
return config
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Generate an image with supported diffusion models.")
parser.add_argument(
@@ -144,6 +154,23 @@ def parse_args() -> argparse.Namespace:
action="store_true",
help="Enable layerwise (blockwise) offloading on DiT modules.",
)
parser.add_argument(
"--use-hsdp",
action="store_true",
help="Enable HSDP (Hybrid Sharded Data Parallel) for diffusion models.",
)
parser.add_argument(
"--hsdp-shard-size",
type=int,
default=1,
help="Number of GPUs to shard weights across for HSDP.",
)
parser.add_argument(
"--hsdp-replicate-size",
type=int,
default=1,
help="Number of HSDP replica groups.",
)
parser.add_argument(
"--quantization",
type=str,
@@ -237,6 +264,12 @@ def parse_args() -> argparse.Namespace:
action="store_true",
help="Enable diffusion pipeline profiler to display stage durations.",
)
parser.add_argument(
"--profiler-config",
type=parse_profiler_config,
default=None,
help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
)
parser.add_argument(
"--log-stats",
action="store_true",
@@ -267,6 +300,9 @@ def parse_args() -> argparse.Namespace:
default=None,
help=("Custom system prompt. Used when --use-system-prompt is custom. "),
)
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
nullify_stage_engine_defaults(parser)
return parser.parse_args()
@@ -316,8 +352,7 @@ def main():
enable_expert_parallel=args.enable_expert_parallel,
)
# Check if profiling is requested via environment variable
profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
profiler_enabled = args.profiler_config is not None
# Prepare LoRA kwargs for Omni initialization
lora_args: dict[str, Any] = {}
@@ -358,6 +393,7 @@ def main():
"mode": "text-to-image",
"log_stats": args.log_stats,
"enable_diffusion_pipeline_profiler": args.enable_diffusion_pipeline_profiler,
"profiler_config": args.profiler_config,
"init_timeout": args.init_timeout,
"stage_init_timeout": args.stage_init_timeout,
**lora_args,
examples/offline_inference/text_to_video/text_to_video.py
+47 -3
View File
@@ -2,9 +2,10 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import argparse
import os
import json
import time
from pathlib import Path
from typing import Any
import numpy as np
import torch
@@ -44,6 +45,16 @@ def _detect_preset(model: str) -> dict:
return _MODEL_PRESETS["wan"]
def parse_profiler_config(value: str) -> dict[str, Any]:
try:
config = json.loads(value)
except json.JSONDecodeError as e:
raise argparse.ArgumentTypeError(f"--profiler-config must be valid JSON: {e}") from e
if not isinstance(config, dict):
raise argparse.ArgumentTypeError("--profiler-config must be a JSON object")
return config
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description="Generate a video from a text prompt. "
@@ -131,6 +142,13 @@ def parse_args() -> argparse.Namespace:
action="store_true",
help="Enable layerwise (blockwise) offloading on DiT modules.",
)
parser.add_argument(
"--ulysses-mode",
type=str,
default="strict",
choices=["strict", "advanced_uaa"],
help="Ulysses sequence-parallel mode: 'strict' (divisibility required) or 'advanced_uaa' (UAA).",
)
parser.add_argument(
"--ulysses-degree",
type=int,
@@ -178,6 +196,12 @@ def parse_args() -> argparse.Namespace:
action="store_true",
help="Enable diffusion pipeline profiler to display stage durations.",
)
parser.add_argument(
"--profiler-config",
type=parse_profiler_config,
default=None,
help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
)
parser.add_argument(
"--quantization",
type=str,
@@ -185,6 +209,26 @@ def parse_args() -> argparse.Namespace:
choices=["fp8", "gguf"],
help="Quantization method for the transformer (fp8 for online FP8 quantization).",
)
parser.add_argument(
"--use-hsdp",
action="store_true",
help="Enable HSDP (Hybrid Sharded Data Parallel) for diffusion models.",
)
parser.add_argument(
"--hsdp-shard-size",
type=int,
default=1,
help="Number of GPUs to shard weights across for HSDP.",
)
parser.add_argument(
"--hsdp-replicate-size",
type=int,
default=1,
help="Number of HSDP replica groups.",
)
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
nullify_stage_engine_defaults(parser)
return parser.parse_args()
@@ -224,8 +268,7 @@ def main():
enable_expert_parallel=args.enable_expert_parallel,
)
# Check if profiling is requested via environment variable
profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
profiler_enabled = args.profiler_config is not None
omni_kwargs = dict(
model=args.model,
@@ -239,6 +282,7 @@ def main():
cache_backend=args.cache_backend,
cache_config=cache_config,
enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
profiler_config=args.profiler_config,
)
if args.boundary_ratio is not None:
omni_kwargs["boundary_ratio"] = args.boundary_ratio
examples/offline_inference/vace/vace_video_generation.py
+3
View File
@@ -71,6 +71,9 @@ def parse_args() -> argparse.Namespace:
parser.add_argument("--ulysses-degree", type=int, default=1, help="Ulysses SP degree.")
parser.add_argument("--ring-degree", type=int, default=1, help="Ring attention degree.")
parser.add_argument("--cfg-parallel-size", type=int, default=1, choices=[1, 2], help="CFG parallel size.")
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
nullify_stage_engine_defaults(parser)
return parser.parse_args()
examples/offline_inference/voxtral_tts/end2end.py
+11
View File
@@ -298,6 +298,12 @@ def parse_args() -> Namespace:
default=None,
help="Voice to use instead of audio file.",
)
parser.add_argument(
"--cfg-alpha",
type=float,
default=None,
help="CFG alpha for flow-matching guidance (default: use value from stage config, typically 1.2).",
)
return parser.parse_args()
@@ -349,8 +355,13 @@ def main(args: Any) -> None:
inputs = compose_request(model_name, text_chunk, audio_prompt_file, args)
extra_args = {}
if args.cfg_alpha is not None:
extra_args["cfg_alpha"] = args.cfg_alpha
sampling_params = SamplingParams(
max_tokens=max_num_tokens,
extra_args=extra_args if extra_args else None,
)
sampling_params_list = [
sampling_params,
examples/online_serving/bagel/README.md
+143 -184
View File
@@ -1,145 +1,111 @@
# BAGEL-7B-MoT
## 🛠️ Installation
## Installation
Please refer to [README.md](../../../README.md)
## Run examples (BAGEL-7B-MoT)
## Architecture
**Note**: These examples work with the default configuration on an **NVIDIA A100 (80GB)**. We also tested on dual **NVIDIA RTX 5000 Ada (32GB each)**. For dual-GPU setups, please modify the stage configuration to distribute the model across devices.
BAGEL-7B-MoT is a Mixture-of-Transformers (MoT) model supporting both image generation and understanding. It offers two deployment topologies:
### Launch the Server
| Topology | Stages | Description |
| :------- | :----- | :---------- |
| **Two-stage** (default) | Stage 0 (Thinker, AR) + Stage 1 (DiT, Diffusion) | Thinker handles text/understanding via vLLM AR engine; DiT handles image generation. KV cache is transferred between stages. |
| **Single-stage** | Stage 0 (DiT, Diffusion) only | The DiT stage contains a full LLM, ViT, VAE, and tokenizer internally. All modalities are handled within a single diffusion process. |
Both topologies support all four modalities: `text2img`, `img2img`, `img2text`, `text2text`.
> **Note**: These examples work with the default configuration on an **NVIDIA A100 (80GB)**. We also tested on dual **NVIDIA RTX 5000 Ada (32GB each)**. For dual-GPU setups, modify the deploy YAML to distribute stages across devices.
## Launch the Server
### Two-Stage (Default)
The default pipeline is auto-detected from the model. No extra flags needed:
```bash
# Use default configuration
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091
```
Or use the convenience script:
```bash
cd /workspace/vllm-omni/examples/online_serving/bagel
cd examples/online_serving/bagel
bash run_server.sh
# Initialize each stage in a discrete isolated process terminal
bash run_server_stage_cli.sh --stage 0
bash run_server_stage_cli.sh --stage 1
```
```bash
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
```
#### 🚀 Tensor Parallelism (TP)
For larger models or multi-GPU environments, you can enable Tensor Parallelism (TP) for the server.
1. **Modify Stage Config**: Create or modify a stage configuration yaml (e.g., [`bagel.yaml`](../../../vllm_omni/model_executor/stage_configs/bagel.yaml)). Set `tensor_parallel_size` to `2` (or more) and update `devices` to include multiple GPU IDs (e.g., `"0,1"`).
```yaml
engine_args:
tensor_parallel_size: 2
...
runtime:
devices: "0,1"
```
2. **Launch Server**:
```bash
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --stage-configs-path /path/to/your/custom_bagel.yaml
```
#### Using Mooncake Connector
By default, BAGEL uses `SharedMemoryConnector` for inter-stage communication. You can use the [Mooncake](https://github.com/kvcache-ai/Mooncake) connector to transfer KV cache between stages, which also enables multi-node deployment.
**1. Install Mooncake**
```bash
# For CUDA-enabled systems (recommended)
pip install mooncake-transfer-engine
# For non-CUDA systems
pip install mooncake-transfer-engine-non-cuda
```
**2. Start Mooncake Master** on the primary node:
```bash
# Optional: enable disk-backed storage by creating a directory and passing --root_fs_dir.
# Without it, Mooncake runs in memory-only mode, which is sufficient for KV cache transfer.
mkdir -p ./mc_storage
mooncake_master \
--rpc_port=50051 \
--enable_http_metadata_server=true \
--http_metadata_server_host=0.0.0.0 \
--http_metadata_server_port=8080 \
--metrics_port=9003 \
--root_fs_dir=./mc_storage/ \
--cluster_id=mc-local-1 &
```
**3. Launch the server** with the Mooncake stage config:
To use a custom deploy YAML, pass it via `--deploy-config`:
```bash
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \
--stage-configs-path vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml
--deploy-config /path/to/deploy_config.yaml
```
> **Note**: Before launching, edit [`bagel_multiconnector.yaml`](../../../vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml) and replace the `metadata_server` and `master` addresses with your Mooncake master node's actual IP. For single-node testing, `127.0.0.1` works.
See [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml) for the default two-stage deploy configuration.
The client-side usage is identical to the default setup -- the Mooncake connector is transparent to the API. See the requests section below.
### Single-Stage
For more details on the Mooncake connector configuration, see the [Mooncake Store Connector documentation](../../../docs/design/feature/omni_connectors/mooncake_store_connector.md).
#### Multi-Node Deployment
You can deploy each stage on a **separate node** for better resource utilization. In this example, the orchestrator (Stage 0 / Thinker) and Stage 1 (DiT) run on different machines, connected via Mooncake.
Replace `<ORCHESTRATOR_IP>` below with the actual IP address of your orchestrator node (e.g., `10.244.227.244`).
> [!WARNING]
> **Before launching**, edit [`bagel_multiconnector.yaml`](../../../vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml) and replace the `metadata_server` and `master` addresses with your Mooncake master node's actual IP. Mismatched addresses will cause silent connection failures.
**1. Start Mooncake Master** (on the orchestrator node):
The DiT stage contains a full LLM, ViT, VAE, and tokenizer, so it can handle all modalities (text2img, img2img, img2text, text2text, think) without a separate Thinker stage:
```bash
mooncake_master \
--rpc_port=50051 \
--enable_http_metadata_server=true \
--http_metadata_server_host=<ORCHESTRATOR_IP> \
--http_metadata_server_port=8080 \
--metrics_port=9003
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \
--deploy-config vllm_omni/deploy/bagel_single_stage.yaml
```
**2. Launch Stage 0 (Thinker / Orchestrator)** on the orchestrator node:
See [`bagel_single_stage.yaml`](../../../vllm_omni/deploy/bagel_single_stage.yaml) for configuration. The `pipeline: bagel_single_stage` field selects the single-stage topology from the pipeline registry.
### Tensor Parallelism (TP)
For larger models or multi-GPU environments, enable TP via CLI:
```bash
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --tensor-parallel-size 2
```
Or set `tensor_parallel_size` per stage in a custom deploy YAML.
### Multi-Node Deployment
Deploy each stage on a **separate node** for better resource utilization. Replace `<ORCHESTRATOR_IP>` with the actual IP address of your orchestrator node.
**1. Launch Stage 0 (Thinker / Orchestrator)** on the orchestrator node:
```bash
# API server port for client requests: 8000
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni \
--port 8000 \ # API server port for client requests
--stage-configs-path vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml \
--port 8000 \
--stage-id 0 \
-oma <ORCHESTRATOR_IP> \
-omp 8091
--omni-master-address <ORCHESTRATOR_IP> \
--omni-master-port 8091
```
**3. Launch Stage 1 (DiT)** on the remote node in headless mode:
**2. Launch Stage 1 (DiT)** on the remote node in headless mode:
```bash
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni \
--stage-configs-path vllm_omni/model_executor/stage_configs/bagel_multiconnector.yaml \
--stage-id 1 \
--headless \
-oma <ORCHESTRATOR_IP> \
-omp 8091
--omni-master-address <ORCHESTRATOR_IP> \
--omni-master-port 8091
```
**Mooncake Master arguments:**
Or use the convenience script:
| Argument | Description |
| :------- | :---------- |
| `--rpc_port` | Mooncake RPC port for control-plane coordination between stages |
| `--enable_http_metadata_server` | Enable the HTTP metadata server for service discovery |
| `--http_metadata_server_host` | IP address to bind the metadata server (use the orchestrator node's IP) |
| `--http_metadata_server_port` | Port for the HTTP metadata server |
| `--metrics_port` | Port for Prometheus-compatible metrics endpoint |
```bash
# Terminal 1: Stage 0
bash run_server_stage_cli.sh --stage 0
# Terminal 2: Stage 1
bash run_server_stage_cli.sh --stage 1
# With extra args
bash run_server_stage_cli.sh --stage 0 -- --tensor-parallel-size 2
bash run_server_stage_cli.sh --stage 1 -- --gpu-memory-utilization 0.9
```
**vllm serve arguments:**
@@ -147,85 +113,31 @@ vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni \
| :------- | :---------- |
| `--stage-id` | Which stage this process runs (0 = Thinker, 1 = DiT) |
| `--headless` | Run without the API server (worker-only mode) |
| `-oma` | Orchestrator master address |
| `-omp` | Orchestrator master port for Stage 1 to connect to Stage 0 for task coordination |
| `-oma` / `--omni-master-address` | Orchestrator master address |
| `-omp` / `--omni-master-port` | Orchestrator master port |
> [!IMPORTANT]
> **Startup Order**: Stage 0 (orchestrator) must be launched **before** Stage 1 (headless).
> Stage 0 will appear to hang on startup until Stage 1 (worker) connects — this is expected behavior.
**Network Requirements**
### Inter-Stage Connectors
All nodes must have network connectivity to each other. Ensure the following ports are open **between all participating nodes**:
When deploying stages across nodes, configure the connector type in the deploy YAML:
| Port | Protocol | Service | Direction |
| :--- | :------- | :------ | :-------- |
| 50051 | TCP | Mooncake Master RPC | Worker → Orchestrator |
| 8080 | TCP | Mooncake HTTP Metadata Server | Worker → Orchestrator |
| 8091 | TCP | Orchestrator Master (`-omp`) | Worker → Orchestrator |
| 8000 | TCP | API Server (`--port`) | Client → Orchestrator |
| 9003 | TCP | Metrics (optional) | Monitoring → Orchestrator |
- **SharedMemoryConnector** (default): Used for single-node deployments. No explicit configuration needed.
- **MooncakeTransferEngineConnector**: For multi-node setups with RDMA hardware. Defined in [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml) under `connectors.rdma_connector`.
> **Tip**: If nodes are behind a firewall or in different VPCs/security groups, make sure the above ports are allowed in ingress/egress rules. All nodes should be reachable via their IP addresses (no NAT). Using nodes on the same subnet or VPC is recommended to minimize latency for Mooncake KV cache transfers.
To use Mooncake, create a custom deploy YAML that binds `output_connectors` / `input_connectors` on each stage to the `rdma_connector` defined in the `connectors` section.
### Send Multi-modal Request
Get into the bagel folder:
## Send Requests
```bash
cd examples/online_serving/bagel
```
Send request via Python
```bash
python openai_chat_client.py --prompt "A cute cat" --modality text2img
```
The Python client supports the following command-line arguments:
- `--prompt` (or `-p`): Text prompt for generation (default: `A cute cat`)
- `--output` (or `-o`): Output file path for image results (default: `bagel_output.png`)
- `--server` (or `-s`): Server URL (default: `http://localhost:8091`)
- `--image-url` (or `-i`): Input image URL or local file path (for img2img/img2text modes)
- `--modality` (or `-m`): Task modality (default: `text2img`). Options: `text2img`, `img2img`, `img2text`, `text2text`
- `--height`: Image height in pixels (default: 512)
- `--width`: Image width in pixels (default: 512)
- `--steps`: Number of inference steps (default: 25)
- `--seed`: Random seed (default: 42)
- `--negative`: Negative prompt for image generation
Example with custom parameters:
```bash
python openai_chat_client.py \
--prompt "A futuristic city" \
--modality text2img \
--height 768 \
--width 768 \
--steps 50 \
--seed 42 \
--negative "blurry, low quality"
```
## Modality Control
BAGEL-7B-MoT supports **multiple modality modes** for different use cases.
The default yaml configuration deploys Thinker and DiT on the same GPU. You can use the default configuration file: [`bagel.yaml`](../../../vllm_omni/model_executor/stage_configs/bagel.yaml)
| Modality | Input | Output | Description |
| ----------- | ------------ | ------ | -------------------------------------- |
| `text2img` | Text | Image | Generate images from text prompts |
| `img2img` | Image + Text | Image | Transform images using text guidance |
| `img2text` | Image + Text | Text | Generate text descriptions from images |
| `text2text` | Text | Text | Pure text generation |
### Text to Image (text2img)
Generate images from text prompts:
**Using Python client**
**Python client:**
```bash
python openai_chat_client.py \
@@ -235,7 +147,7 @@ python openai_chat_client.py \
--steps 50
```
**Using curl**
**curl:**
```bash
curl http://localhost:8091/v1/chat/completions \
@@ -250,12 +162,9 @@ curl http://localhost:8091/v1/chat/completions \
}'
```
### Image to Image (img2img)
Transform images based on text prompts:
**Using Python client**
**Python client:**
```bash
python openai_chat_client.py \
@@ -265,7 +174,7 @@ python openai_chat_client.py \
--output transformed.png
```
**Using curl**
**curl:**
```bash
IMAGE_BASE64=$(base64 -w 0 cat.jpg)
@@ -290,14 +199,11 @@ EOF
curl http://localhost:8091/v1/chat/completions \
-H "Content-Type: application/json" \
-d @payload.json
```
### Image to Text (img2text)
Generate text descriptions from images:
**Using Python client**
**Python client:**
```bash
python openai_chat_client.py \
@@ -306,7 +212,7 @@ python openai_chat_client.py \
--image-url /path/to/image.jpg
```
**Using curl**
**curl:**
```bash
IMAGE_BASE64=$(base64 -w 0 cat.jpg)
@@ -331,9 +237,7 @@ curl http://localhost:8091/v1/chat/completions \
### Text to Text (text2text)
Pure text generation:
**Using Python client**
**Python client:**
```bash
python openai_chat_client.py \
@@ -341,23 +245,78 @@ python openai_chat_client.py \
--modality text2text
```
**Using curl**
**curl:**
```bash
curl http://localhost:8091/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{
"messages": [{"role": "user", "content": [{"type": "text", "text": "<|im_start|>user\nWhat is the capital of France?<|im_end|>\n<|im_start|>assistant\n"}]}]
"messages": [{"role": "user", "content": [{"type": "text", "text": "<|im_start|>user\nWhat is the capital of France?<|im_end|>\n<|im_start|>assistant\n"}]}],
"modalities": ["text"]
}'
```
### Python Client Arguments
| Argument | Default | Description |
| :------- | :------ | :---------- |
| `--prompt` / `-p` | `A cute cat` | Text prompt |
| `--output` / `-o` | `bagel_output.png` | Output file path |
| `--server` / `-s` | `http://localhost:8091` | Server URL |
| `--image-url` / `-i` | `None` | Input image URL or local path (img2img/img2text) |
| `--modality` / `-m` | `text2img` | `text2img`, `img2img`, `img2text`, `text2text` |
| `--height` | `512` | Image height in pixels |
| `--width` | `512` | Image width in pixels |
| `--steps` | `25` | Number of inference steps |
| `--seed` | `42` | Random seed |
| `--negative` | `None` | Negative prompt for CFG |
Example with custom parameters:
```bash
python openai_chat_client.py \
--prompt "A futuristic city" \
--modality text2img \
--height 768 \
--width 768 \
--steps 50 \
--seed 42 \
--negative "blurry, low quality"
```
## Configuration Reference
### Deploy YAML Files
| File | Description |
| :--- | :---------- |
| [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml) | Two-stage default (Thinker + DiT on GPU 0) |
| [`bagel_single_stage.yaml`](../../../vllm_omni/deploy/bagel_single_stage.yaml) | Single-stage (DiT only) |
### Key Deploy YAML Fields
| Field | Scope | Description |
| :---- | :---- | :---------- |
| `pipeline` | top-level | Override auto-detected pipeline (e.g. `bagel_single_stage`) |
| `stages[].stage_id` | per-stage | Stage identifier (0, 1, ...) |
| `stages[].devices` | per-stage | GPU device IDs (e.g. `"0"`, `"0,1"`) |
| `stages[].max_num_seqs` | per-stage | Maximum concurrent sequences |
| `stages[].gpu_memory_utilization` | per-stage | Fraction of GPU memory to use |
| `stages[].enforce_eager` | per-stage | Disable CUDA graphs |
| `stages[].tensor_parallel_size` | per-stage | TP degree for this stage |
| `connectors` | top-level | Define available connector instances (SHM, Mooncake) |
| `platforms` | top-level | Platform-specific overrides (e.g. `xpu`) |
## FAQ
- If you don’t know how much VRAM is needed for the model or encounter the OOM error, you can try to decrease the max_model_len.
- If you encounter OOM errors, try decreasing `max_model_len` or `gpu_memory_utilization` in the deploy YAML.
| Stage | VRAM |
| :------------------ | :--------------------------- |
| Stage-0 (Thinker) | **15.04 GiB** **+ KV Cache** |
| Stage-1 (DiT) | **26.50 GiB** |
| Total | **~42 GiB + KV Cache** |
**Two-stage VRAM usage:**
| Stage | VRAM |
| :---- | :--- |
| Stage 0 (Thinker) | **15.04 GiB + KV Cache** |
| Stage 1 (DiT) | **26.50 GiB** |
| Total | **~42 GiB + KV Cache** |
**Single-stage VRAM usage:** The DiT loads the full model (~42 GiB) in one process.
examples/online_serving/bagel/run_server_stage_cli.sh
+11 -11
View File
@@ -20,7 +20,7 @@ MASTER_ADDRESS="${MASTER_ADDRESS:-127.0.0.1}"
MASTER_PORT="${MASTER_PORT:-8092}"
STAGE="all"
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
STAGE_CONFIGS_PATH="${STAGE_CONFIGS_PATH:-$SCRIPT_DIR/../../../vllm_omni/model_executor/stage_configs/bagel.yaml}"
DEPLOY_CONFIG="${DEPLOY_CONFIG:-$SCRIPT_DIR/../../../vllm_omni/deploy/bagel.yaml}"
EXTRA_ARGS=()
usage() {
@@ -33,7 +33,7 @@ Options:
--port PORT API port for stage 0 (default: $PORT)
--master-address ADDRESS Master/orchestrator address (default: $MASTER_ADDRESS)
--master-port PORT Master/orchestrator port (default: $MASTER_PORT)
--stage-configs-path PATH Stage config YAML path (default: $STAGE_CONFIGS_PATH)
--deploy-config PATH Deploy config YAML path (default: $DEPLOY_CONFIG)
--help Show this help message
Examples:
@@ -71,8 +71,8 @@ while [[ $# -gt 0 ]]; do
MASTER_PORT="$2"
shift 2
;;
--stage-configs-path)
STAGE_CONFIGS_PATH="$2"
--deploy-config)
DEPLOY_CONFIG="$2"
shift 2
;;
--help|-h)
@@ -103,7 +103,7 @@ print_config() {
echo "API Port: $PORT"
echo "Master Address: $MASTER_ADDRESS"
echo "Master Port: $MASTER_PORT"
echo "Stage Configs: $STAGE_CONFIGS_PATH"
echo "Deploy Config: $DEPLOY_CONFIG"
echo "Selected Stage: $STAGE"
if [[ ${#EXTRA_ARGS[@]} -gt 0 ]]; then
echo "Extra Args: ${EXTRA_ARGS[*]}"
@@ -114,21 +114,21 @@ run_stage_0() {
echo "Starting Stage 0 (Thinker) as master..."
vllm serve "$MODEL" --omni \
--port "$PORT" \
--stage-configs-path "$STAGE_CONFIGS_PATH" \
--deploy-config "$DEPLOY_CONFIG" \
--stage-id 0 \
-oma "$MASTER_ADDRESS" \
-omp "$MASTER_PORT" \
--omni-master-address "$MASTER_ADDRESS" \
--omni-master-port "$MASTER_PORT" \
"${EXTRA_ARGS[@]}"
}
run_stage_1() {
echo "Starting Stage 1 (DiT) in headless mode..."
vllm serve "$MODEL" --omni \
--stage-configs-path "$STAGE_CONFIGS_PATH" \
--deploy-config "$DEPLOY_CONFIG" \
--stage-id 1 \
--headless \
-oma "$MASTER_ADDRESS" \
-omp "$MASTER_PORT" \
--omni-master-address "$MASTER_ADDRESS" \
--omni-master-port "$MASTER_PORT" \
"${EXTRA_ARGS[@]}"
}
examples/online_serving/diffusers_pipeline_adapter/README.md
+83
View File
@@ -0,0 +1,83 @@
# Diffusers Backend Adapter Example
This example demonstrates how to serve any 🤗 Diffusers pipeline through vLLM-Omni
using the `diffusers` load format.
## Supported Models
Any model loadable via `DiffusionPipeline.from_pretrained()` should be supported, including text-to-image, image-to-image, text-to-video, image-to-video, and text-to-audio.
## Limitations
The diffusers backend is a black-box adapter. The following features are NOT yet supported.
It is not guaranteed whether they will be supported in the future.
- CFG parallel execution
- Sequence parallel execution
- TeaCache / Cache-DiT acceleration
- Step-wise execution (continuous batching)
For these features, it is recommended to use natively supported pipelines instead.
## Usage
### Option 1: CLI arguments
```bash
vllm serve "stable-diffusion-v1-5/stable-diffusion-v1-5" \
--omni \
--diffusion-load-format diffusers \
--diffusers-load-kwargs '{"use_safetensors": true}' \
--diffusers-call-kwargs '{"num_inference_steps": 30, "guidance_scale": 7.5}'
```
`--diffusers-load-kwargs` and `--diffusers-call-kwargs` are only valid together with `--diffusion-load-format diffusers`.
### Option 2: Stage config YAML
```bash
vllm serve stable-diffusion-v1-5/stable-diffusion-v1-5 --stage-configs-path examples/online_serving/diffusers_pipeline_adapter/stage_config.yaml --omni
```
The particular fields of interest are `model`, `diffusion_load_format`, `diffusers_load_kwargs`, and `diffusers_call_kwargs` under `engine_args`. They are the same as the CLI arguments.
## Send a Request
```bash
curl http://localhost:8000/v1/images/generations \
-H "Content-Type: application/json" \
-d '{
"model": "stable-diffusion-v1-5/stable-diffusion-v1-5",
"prompt": "a photo of an astronaut riding a horse on mars",
"n": 1,
"size": "512x512"
}'
```
Or refer to other documentation pages on how to request a particular input/output modality, such as `examples/online_serving/text_to_image/openai_chat_client.py`.
## Configuration Reference
For the diffusers adapter, set options under **`engine_args`**:
### `diffusion_load_format: "diffusers"`
This field selects the Hugging Face diffusers adapter path (see `DiffusersPipelineLoader`).
### `diffusers_load_kwargs`
Passed to `DiffusionPipeline.from_pretrained()`.
This is suitable for model-specific configurations not available through the vLLM-Omni interface (such as `Omni.__init__()`, `vllm serve` CLI arguments, and stage config YAML fields outside `diffusers_load_kwargs`).
When a parameter is available in the vLLM-Omni interface, it will be adapted here.
But if that parameter is simultaneously set in both the vLLM-Omni interface and `diffusers_load_kwargs`, the **latter** will take precedence.
### `diffusers_call_kwargs`
Passed to `pipeline.__call__()`.
This is suitable for sampling parameters not available through the vLLM-Omni interface (such as `Omni.generate()` and online serving payloads).
When a parameter is available in the vLLM-Omni interface, it will be adapted here.
But if that parameter is simultaneously set in both the vLLM-Omni interface and `diffusers_call_kwargs`, the **former** will take precedence (because it is set at request time).
examples/online_serving/diffusers_pipeline_adapter/stage_config.yaml
+31
View File
@@ -0,0 +1,31 @@
# Example stage config for diffusers backend
# This config demonstrates serving Stable Diffusion 1.5 via the diffusers adapter.
# Users should copy and modify this for their own models.
model_type: diffusion
stage_args:
- stage_id: 0
stage_type: diffusion
engine_args:
model_stage: diffusion
model: "stable-diffusion-v1-5/stable-diffusion-v1-5"
distributed_executor_backend: "mp"
# gpu_memory_utilization: 0.9
engine_output_type: image
# Select the HF diffusers adapter
diffusion_load_format: "diffusers"
# model_class_name: "DiffusersAdapterPipeline" # default when diffusion_load_format is diffusers
diffusers_load_kwargs:
# Passed to DiffusionPipeline.from_pretrained().
# Good for model-specific loading parameters not covered by OmniDiffusionConfig.
# During model load time, parameters here override their counterparts in the vLLM-Omni interface.
use_safetensors: true
diffusers_call_kwargs:
# Passed to pipeline.__call__().
# Good for model-specific sampling parameters not covered by OmniDiffusionSamplingParams.
# During request time, parameters here are overridden by the counterparts in OmniDiffusionSamplingParams.
num_inference_steps: 30
guidance_scale: 7.5
final_output: true
final_output_type: image
examples/online_serving/glm_image/README.md
-204
View File
@@ -1,204 +0,0 @@
# GLM-Image Online Serving
This example demonstrates how to deploy GLM-Image for online image generation using vLLM-Omni.
## 🛠️ Installation
Please refer to [README.md](../../../README.md)
## Run examples (GLM-Image)
**Note**: These examples work with the default configuration on **2× NVIDIA A100 (80GB)** or equivalent. Stage 0 (AR) and Stage 1 (Diffusion) each use one GPU by default. For single-GPU setups, modify the stage configuration to share the same device.
### Launch the Server
```bash
# Use default configuration
vllm serve zai-org/GLM-Image --omni --port 8091
```
Or use the convenience script:
```bash
cd examples/online_serving/glm_image
bash run_server.sh
```
If you have a custom stage configs file:
```bash
vllm serve zai-org/GLM-Image --omni --port 8091 --stage-configs-path /path/to/glm_image.yaml
```
### Send Requests
Get into the glm_image folder:
```bash
cd examples/online_serving/glm_image
```
Send request via Python:
```bash
python openai_chat_client.py --prompt "A cute cat sitting on a window sill"
```
The Python client supports the following command-line arguments:
- `--prompt` (or `-p`): Text prompt for generation (default: `A beautiful sunset over the ocean with sailing boats`)
- `--output` (or `-o`): Output file path (default: `glm_image_output.png`)
- `--server` (or `-s`): Server URL (default: `http://localhost:8091`)
- `--image` (or `-i`): Input image path (for image-to-image editing)
- `--height`: Image height in pixels (default: 1024)
- `--width`: Image width in pixels (default: 1024)
- `--steps`: Number of inference steps (default: 50)
- `--guidance-scale`: Classifier-free guidance scale (default: 1.5)
- `--seed`: Random seed (default: 42)
- `--negative`: Negative prompt
## Modality Control
GLM-Image supports **text-to-image** and **image-to-image** modes.
The default yaml configuration deploys AR on GPU 0 and DiT on GPU 1. You can use the default configuration file: [`glm_image.yaml`](../../../vllm_omni/model_executor/stage_configs/glm_image.yaml)
| Mode | Input | Output | Description |
| -------------- | ------------ | ------ | ---------------------------------- |
| Text-to-Image | Text | Image | Generate images from text prompts |
| Image-to-Image | Image + Text | Image | Edit images with text instructions |
### Text-to-Image
```bash
python openai_chat_client.py \
--prompt "A photorealistic mountain landscape at sunset" \
--height 1024 \
--width 1024 \
--output landscape.png
# Or use the curl script:
bash run_curl_text_to_image.sh "A futuristic city skyline at night"
```
### Image-to-Image (Image Editing)
```bash
python openai_chat_client.py \
--prompt "Convert this image to watercolor style" \
--image input.png \
--output watercolor.png
# Or use the curl script:
bash run_curl_image_edit.sh input.png "Convert to watercolor style"
```
For general-purpose request methods (curl, OpenAI SDK, Python `requests`), see
the [Text-to-Image](../text_to_image/README.md) and
[Image-to-Image](../image_to_image/README.md) guides.
## Generation Parameters
When using `/v1/chat/completions`, pass these inside `extra_body` in the curl
JSON, or via the `extra_body` keyword argument in the OpenAI Python SDK.
When using the dedicated `/v1/images/generations` or `/v1/images/edits`
endpoints, pass the supported generation controls as top-level fields directly.
For image dimensions and count, use `size` and `n` rather than `height` or
`width`.
| Parameter | Type | Default | Description |
| --------------------- | ----- | ------- | ----------------------------------- |
| `height` | int | 1024 | Image height in pixels |
| `width` | int | 1024 | Image width in pixels |
| `num_inference_steps` | int | 50 | Number of diffusion denoising steps |
| `guidance_scale` | float | 1.5 | Classifier-free guidance scale |
| `seed` | int | None | Optional random seed; `/v1/images/*` generates one server-side if omitted |
| `negative_prompt` | str | None | Negative prompt |
## Response Format
```json
{
"id": "chatcmpl-xxx",
"created": 1234567890,
"model": "zai-org/GLM-Image",
"choices": [
{
"index": 0,
"message": {
"role": "assistant",
"content": [
{
"type": "image_url",
"image_url": {
"url": "data:image/png;base64,..."
}
}
]
},
"finish_reason": "stop"
}
],
"usage": {}
}
```
## Extract Image
```bash
# From a saved JSON response
cat response.json | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > output.png
```
## Architecture
GLM-Image uses a 2-stage multistage pipeline:
```
Stage 0 (AR Model) Stage 1 (Diffusion)
┌───────────────────┐ ┌─────────────────────┐
│ vLLM-optimized │ prior │ GlmImagePipeline │
│ GlmImageFor │──tokens──►│ ┌───────────────┐ │
│ Conditional │ │ │ DiT Denoiser │ │
│ Generation │ │ └───────┬───────┘ │
│ (9B AR model) │ │ ▼ │
└───────────────────┘ │ ┌───────────────┐ │
▲ │ │ VAE Decode │──┼──► Image
│ │ └───────────────┘ │
Text / Image └─────────────────────┘
Input
```
## VRAM Requirements
| Stage | VRAM |
| :---------------- | :--------------------- |
| Stage-0 (AR) | **~18 GiB + KV Cache** |
| Stage-1 (DiT+VAE) | **~20 GiB** |
| Total | **~38 GiB + KV Cache** |
## File Description
| File | Description |
| --------------------------- | ------------------------------------- |
| `run_server.sh` | Server startup script |
| `run_curl_text_to_image.sh` | Text-to-image curl example |
| `run_curl_image_edit.sh` | Image-to-image (editing) curl example |
| `openai_chat_client.py` | Python client (t2i + i2i) |
## FAQ
- If you encounter OOM errors, adjust `gpu_memory_utilization` in the stage config:
```yaml
# In glm_image.yaml, reduce from default 0.6:
gpu_memory_utilization: 0.5
```
- The first request may be slow due to model warmup. Subsequent requests will be faster.
- If you encounter `Transformers does not recognize this architecture` error, your have to upgrade `transformers` package to `5.3.0` or above:
```
pip install --upgrade transformers
```
examples/online_serving/glm_image/openai_chat_client.py
-172
View File
@@ -1,172 +0,0 @@
#!/usr/bin/env python3
"""
GLM-Image OpenAI-compatible chat client for text-to-image and image-to-image.
Usage:
# Text-to-image
python openai_chat_client.py --prompt "A cute cat" --output output.png
# Image-to-image (image editing)
python openai_chat_client.py --prompt "Convert to watercolor style" --image input.png --output output.png
"""
import argparse
import base64
from pathlib import Path
import requests
def generate_image(
prompt: str,
server_url: str = "http://localhost:8091",
image_path: str | None = None,
height: int = 1024,
width: int = 1024,
steps: int = 50,
guidance_scale: float = 1.5,
seed: int | None = None,
negative_prompt: str | None = None,
) -> bytes | None:
"""Generate or edit an image using the chat completions API.
Args:
prompt: Text description or editing instruction
server_url: Server URL
image_path: Path to input image (for image-to-image editing)
height: Image height in pixels
width: Image width in pixels
steps: Number of inference steps
guidance_scale: Classifier-free guidance scale
seed: Random seed for reproducibility
negative_prompt: Negative prompt
Returns:
Image bytes or None if failed
"""
# Build message content
content: list[dict] = [{"type": "text", "text": prompt}]
if image_path:
img_path = Path(image_path)
if not img_path.exists():
print(f"Error: Image file not found: {image_path}")
return None
b64_data = base64.b64encode(img_path.read_bytes()).decode("utf-8")
suffix = img_path.suffix.lstrip(".").lower()
mime = {"jpg": "jpeg", "jpeg": "jpeg", "png": "png", "webp": "webp"}.get(suffix, "png")
content.append(
{
"type": "image_url",
"image_url": {"url": f"data:image/{mime};base64,{b64_data}"},
}
)
messages = [{"role": "user", "content": content}]
# Build request payload
extra_body: dict = {
"height": height,
"width": width,
"num_inference_steps": steps,
"guidance_scale": guidance_scale,
}
if seed is not None:
extra_body["seed"] = seed
if negative_prompt:
extra_body["negative_prompt"] = negative_prompt
payload = {"messages": messages, "extra_body": extra_body}
# Send request
try:
mode = "image-to-image" if image_path else "text-to-image"
print(f"Sending {mode} request to {server_url}...")
response = requests.post(
f"{server_url}/v1/chat/completions",
headers={"Content-Type": "application/json"},
json=payload,
timeout=600,
)
response.raise_for_status()
data = response.json()
# Extract image from response
choices = data.get("choices", [])
for choice in choices:
choice_content = choice.get("message", {}).get("content")
if isinstance(choice_content, list):
for item in choice_content:
if isinstance(item, dict) and "image_url" in item:
img_url = item["image_url"].get("url", "")
if img_url.startswith("data:image"):
_, b64 = img_url.split(",", 1)
return base64.b64decode(b64)
print(f"Unexpected response format: {data}")
return None
except Exception as e:
print(f"Error: {e}")
return None
def main():
parser = argparse.ArgumentParser(description="GLM-Image chat client")
parser.add_argument(
"--prompt",
"-p",
default="A beautiful sunset over the ocean with sailing boats",
help="Text prompt",
)
parser.add_argument("--output", "-o", default="glm_image_output.png", help="Output file")
parser.add_argument("--server", "-s", default="http://localhost:8091", help="Server URL")
# Image-to-image
parser.add_argument(
"--image",
"-i",
type=str,
help="Input image path (for image-to-image editing)",
)
# Generation parameters
parser.add_argument("--height", type=int, default=1024, help="Image height")
parser.add_argument("--width", type=int, default=1024, help="Image width")
parser.add_argument("--steps", type=int, default=50, help="Inference steps")
parser.add_argument("--guidance-scale", type=float, default=1.5, help="CFG guidance scale")
parser.add_argument("--seed", type=int, default=42, help="Random seed")
parser.add_argument("--negative", help="Negative prompt")
args = parser.parse_args()
mode = "image-to-image" if args.image else "text-to-image"
print(f"Mode: {mode}")
print(f"Prompt: {args.prompt}")
if args.image:
print(f"Input image: {args.image}")
image_bytes = generate_image(
prompt=args.prompt,
server_url=args.server,
image_path=args.image,
height=args.height,
width=args.width,
steps=args.steps,
guidance_scale=args.guidance_scale,
seed=args.seed,
negative_prompt=args.negative,
)
if image_bytes:
output_path = Path(args.output)
output_path.write_bytes(image_bytes)
print(f"Image saved to: {output_path}")
print(f"Size: {len(image_bytes) / 1024:.1f} KB")
else:
print("Failed to generate image")
exit(1)
if __name__ == "__main__":
main()
examples/online_serving/glm_image/run_curl_image_edit.sh
-61
View File
@@ -1,61 +0,0 @@
#!/bin/bash
# GLM-Image image-edit (image-to-image) curl example
set -euo pipefail
if [[ $# -lt 2 ]]; then
echo "Usage: $0 <input_image> \"<edit_prompt>\" [output_file]" >&2
exit 1
fi
INPUT_IMG=$1
PROMPT=$2
SERVER="${SERVER:-http://localhost:8091}"
CURRENT_TIME=$(date +%Y%m%d%H%M%S)
OUTPUT="${3:-glm_image_i2i_${CURRENT_TIME}.png}"
if [[ ! -f "$INPUT_IMG" ]]; then
echo "Input image not found: $INPUT_IMG" >&2
exit 1
fi
# base64 encode (macOS uses -i, Linux uses -w0)
if [[ "$(uname)" == "Darwin" ]]; then
IMG_B64=$(base64 < "$INPUT_IMG" | tr -d '\n')
else
IMG_B64=$(base64 -w0 "$INPUT_IMG")
fi
REQUEST_JSON=$(
jq -n --arg prompt "$PROMPT" --arg img "$IMG_B64" '{
messages: [{
role: "user",
content: [
{"type": "text", "text": $prompt},
{"type": "image_url", "image_url": {"url": ("data:image/png;base64," + $img)}}
]
}],
extra_body: {
height: 1024,
width: 1024,
num_inference_steps: 50,
guidance_scale: 1.5,
seed: 42
}
}'
)
echo "Generating edited image..."
echo "Server: $SERVER"
echo "Prompt: $PROMPT"
echo "Input : $INPUT_IMG"
echo "Output: $OUTPUT"
curl -s "$SERVER/v1/chat/completions" \
-H "Content-Type: application/json" \
-d "$REQUEST_JSON" \
| jq -r '.choices[0].message.content[0].image_url.url' \
| cut -d',' -f2- \
| base64 -d > "$OUTPUT"
echo "Image saved to: $OUTPUT"
examples/online_serving/glm_image/run_curl_text_to_image.sh
-30
View File
@@ -1,30 +0,0 @@
#!/bin/bash
# GLM-Image text-to-image curl example
set -euo pipefail
PROMPT="${1:-A beautiful sunset over the ocean with sailing boats}"
SERVER="${SERVER:-http://localhost:8091}"
OUTPUT="${OUTPUT:-glm_image_t2i_output.png}"
echo "Generating image..."
echo "Server: $SERVER"
echo "Prompt: $PROMPT"
echo "Output: $OUTPUT"
curl -s "$SERVER/v1/chat/completions" \
-H "Content-Type: application/json" \
-d "{
\"messages\": [
{\"role\": \"user\", \"content\": \"$PROMPT\"}
],
\"extra_body\": {
\"height\": 1024,
\"width\": 1024,
\"num_inference_steps\": 50,
\"guidance_scale\": 1.5,
\"seed\": 42
}
}" | jq -r '.choices[0].message.content[0].image_url.url' | cut -d',' -f2- | base64 -d > "$OUTPUT"
echo "Image saved to: $OUTPUT"
examples/online_serving/glm_image/run_server.sh
-12
View File
@@ -1,12 +0,0 @@
#!/bin/bash
# GLM-Image online serving startup script
MODEL="${MODEL:-zai-org/GLM-Image}"
PORT="${PORT:-8091}"
echo "Starting GLM-Image server..."
echo "Model: $MODEL"
echo "Port: $PORT"
vllm serve "$MODEL" --omni \
--port "$PORT"
examples/online_serving/ming_flash_omni/README.md
+37 -146
View File
@@ -4,38 +4,49 @@
Please refer to [README.md](../../../README.md)
## Deployment modes
| Mode | Launch command | Output |
|------|---------------|--------|
| Thinker only (multimodal understanding) | `vllm serve ... --omni` | Text |
| Thinker + Talker (omni-speech) | `vllm serve ... --omni --stage-configs-path ming_flash_omni.yaml` | Text + Audio |
For standalone TTS (talker only), see [`examples/online_serving/ming_flash_omni_tts/`](../ming_flash_omni_tts/).
## Run examples (Ming-flash-omni 2.0)
### Launch the Server
**Thinker only (text output):**
```bash
vllm serve Jonathan1909/Ming-flash-omni-2.0 --omni --port 8091
```
If you have custom stage configs file, launch the server with command below
**Thinker + Talker (omni-speech, text + audio output):**
```bash
vllm serve Jonathan1909/Ming-flash-omni-2.0 --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
vllm serve Jonathan1909/Ming-flash-omni-2.0 --omni --port 8091 \
--stage-configs-path vllm_omni/model_executor/stage_configs/ming_flash_omni.yaml
```
Pass `--stage-configs-path /path/to/your_config.yaml` to use a custom stage
config.
### Send Multi-modal Request
#### Send request via python
Shared Python client (supports `text | use_image | use_audio | use_video |
use_mixed_modalities`; pass `--image-path` / `--audio-path` / `--video-path`
for local files or URLs, `--modalities text` for output, `--help` for the
full flag list):
```bash
python examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py --model Jonathan1909/Ming-flash-omni-2.0 --query-type use_mixed_modalities --port 8091 --host "localhost" --modalities text
python examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py \
--model Jonathan1909/Ming-flash-omni-2.0 \
--query-type use_mixed_modalities \
--port 8091 --host localhost \
--modalities text
```
The Python client supports the following command-line arguments:
- `--query-type` (or `-q`): Query type. Options: `text`, `use_audio`, `use_image`, `use_video`, `use_mixed_modalities`
- `--video-path` (or `-v`): Path to local video file or URL. If not provided and query-type uses video, uses default video URL. Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs. Example: `--video-path /path/to/video.mp4` or `--video-path https://example.com/video.mp4`
- `--image-path` (or `-i`): Path to local image file or URL. If not provided and query-type uses image, uses default image URL. Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs and common image formats: JPEG, PNG, GIF, WebP. Example: `--image-path /path/to/image.jpg` or `--image-path https://example.com/image.png`
- `--audio-path` (or `-a`): Path to local audio file or URL. If not provided and query-type uses audio, uses default audio URL. Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs and common audio formats: MP3, WAV, OGG, FLAC, M4A. Example: `--audio-path /path/to/audio.wav` or `--audio-path https://example.com/audio.mp3`
- `--prompt` (or `-p`): Custom text prompt/question. If not provided, uses default prompt for the selected query type. Example: `--prompt "What are the main activities shown in this video?"`
- `--modalities`: Output modalities. For now, only `text` is supported. Example: `--modalities text`
#### Send request via curl
Parameterized curl wrapper in this directory:
```bash
bash run_curl_multimodal_generation.sh text
@@ -47,83 +58,17 @@ bash run_curl_multimodal_generation.sh use_mixed_modalities
## Modality control
Ming-flash-omni 2.0 currently supports text output only (thinker stage).
| `modalities` | Server config | Output |
|-------------|--------------|--------|
| `["text"]` or omitted | Thinker only | Text |
| `["audio"]` | Thinker + Talker | Audio (speech) |
| `["text", "audio"]` | Thinker + Talker | Text + Audio |
| Modalities | Output |
|------------|--------|
| `["text"]` | Text only |
| Not specified | Text only (default) |
For ready-to-copy curl examples (text / audio / multimodal input, SSE
streaming, reasoning mode), see the recipe at
[`recipes/inclusionAI/Ming-flash-omni-2.0.md`](../../../recipes/inclusionAI/Ming-flash-omni-2.0.md).
### Using curl
```bash
curl http://localhost:8091/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{
"model": "Jonathan1909/Ming-flash-omni-2.0",
"messages": [
{"role": "system", "content": [{"type": "text", "text": "你是一个友好的AI助手。\n\ndetailed thinking off"}]},
{"role": "user", "content": "请详细介绍鹦鹉的生活习性。"}
],
"modalities": ["text"]
}'
```
### Using OpenAI Python SDK
```python
from openai import OpenAI
client = OpenAI(base_url="http://localhost:8091/v1", api_key="EMPTY")
response = client.chat.completions.create(
model="Jonathan1909/Ming-flash-omni-2.0",
messages=[
{"role": "system", "content": [{"type": "text", "text": "你是一个友好的AI助手。\n\ndetailed thinking off"}]},
{"role": "user", "content": "请详细介绍鹦鹉的生活习性。"},
],
modalities=["text"],
)
print(response.choices[0].message.content)
```
### Multi-modal input with OpenAI Python SDK
```python
from openai import OpenAI
client = OpenAI(base_url="http://localhost:8091/v1", api_key="EMPTY")
response = client.chat.completions.create(
model="Jonathan1909/Ming-flash-omni-2.0",
messages=[
{"role": "system", "content": [{"type": "text", "text": "你是一个友好的AI助手。\n\ndetailed thinking off"}]},
{
"role": "user",
"content": [
{"type": "image_url", "image_url": {"url": "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg"}},
{"type": "text", "text": "Describe this image in detail."},
],
},
],
modalities=["text"],
)
print(response.choices[0].message.content)
```
## Streaming Output
To enable streaming output:
```bash
python examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py \
--query-type use_image \
--model Jonathan1909/Ming-flash-omni-2.0 \
--modalities text \
--stream
```
Or with the OpenAI Python SDK:
## OpenAI Python SDK — streaming
```python
from openai import OpenAI
@@ -146,59 +91,5 @@ for chunk in response:
print()
```
Or using curl:
```bash
curl http://localhost:8091/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{
"model": "Jonathan1909/Ming-flash-omni-2.0",
"messages": [
{"role": "system", "content": [{"type": "text", "text": "你是一个友好的AI助手。\n\ndetailed thinking off"}]},
{"role": "user", "content": "请详细介绍鹦鹉的生活习性。"}
],
"modalities": ["text"],
"stream": true,
}'
```
## Reasoning (Thinking Mode)
To enable reasoning/thinking mode, change `detailed thinking off` to `detailed thinking on` in the system prompt:
### Using curl
```bash
curl http://localhost:8091/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{
"model": "Jonathan1909/Ming-flash-omni-2.0",
"messages": [
{"role": "system", "content": [{"type": "text", "text": "你是一个友好的AI助手。\n\ndetailed thinking on"}]},
{"role": "user", "content": [
{"type": "image_url", "image_url": {"url": "https://example.com/math_problem.png"}},
{"type": "text", "text": "Solve this math problem step by step."}
]}
],
"modalities": ["text"]
}'
```
### Using OpenAI Python SDK
```python
from openai import OpenAI
client = OpenAI(base_url="http://localhost:8091/v1", api_key="EMPTY")
response = client.chat.completions.create(
model="Jonathan1909/Ming-flash-omni-2.0",
messages=[
{"role": "system", "content": [{"type": "text", "text": "你是一个友好的AI助手。\n\ndetailed thinking on"}]},
{"role": "user", "content": "If a train travels 120 km in 2 hours, what is its average speed?"},
],
modalities=["text"],
)
print(response.choices[0].message.content)
```
The `--stream` flag on the Python client script above shows the same pattern
driven by the shared multimodal client.
examples/online_serving/ming_flash_omni_tts/README.md
+54
View File
@@ -0,0 +1,54 @@
# Ming-flash-omni Standalone TTS (Online Serving)
This directory contains online e2e examples for **Ming-flash-omni-2.0 standalone talker deployment**.
Server uses:
- `model`: `Jonathan1909/Ming-flash-omni-2.0`
- `stage config`: `vllm_omni/model_executor/stage_configs/ming_flash_omni_tts.yaml`
## Launch the Server
```bash
# from repo root
bash examples/online_serving/ming_flash_omni_tts/run_server.sh
```
Equivalent manual command:
```bash
vllm serve Jonathan1909/Ming-flash-omni-2.0 \
--stage-configs-path vllm_omni/model_executor/stage_configs/ming_flash_omni_tts.yaml \
--host 0.0.0.0 \
--port 8091 \
--trust-remote-code \
--omni
```
## Send TTS Request
Python client:
```bash
python examples/online_serving/ming_flash_omni_tts/speech_client.py \
--text "我们当迎着阳光辛勤耕作，去摘取，去制作，去品尝，去馈赠。" \
--output ming_online.wav
```
Long-form `instructions` (e.g. ASMR whisper style) via the client:
```bash
python examples/online_serving/ming_flash_omni_tts/speech_client.py \
--text "我会一直在这里陪着你，直到你慢慢、慢慢地沉入那个最温柔的梦里……好吗？" \
--instructions "这是一种ASMR耳语，属于一种旨在引发特殊感官体验的创意风格。这个女性使用轻柔的普通话进行耳语，声音气音成分重。音量极低，紧贴麦克风，语速极慢，旨在制造触发听者颅内快感的声学刺激。" \
--output ming_online_asmr.wav
```
## Notes
- This is the **online serving** counterpart of `examples/offline_inference/ming_flash_omni_tts/`.
- The server uses `use_zero_spk_emb=True` and the default decode args
(`max_decode_steps=200`, `cfg=2.0`, `sigma=0.25`, `temperature=0.0`).
For other caption fields (`语速`, `基频`, `IP`, BGM, etc.) or overriding
decode args, use the offline e2e example where `additional_information`
is set explicitly.
examples/online_serving/ming_flash_omni_tts/run_server.sh
Executable
+25
View File
@@ -0,0 +1,25 @@
#!/bin/bash
# Launch vLLM-Omni server for Ming-flash-omni-2.0 standalone talker (TTS).
#
# Usage:
# ./run_server.sh
# MODEL=/path/to/local/model ./run_server.sh
# PORT=8091 ./run_server.sh
# HOST=127.0.0.1 ./run_server.sh # bind only to loopback
set -e
MODEL="${MODEL:-Jonathan1909/Ming-flash-omni-2.0}"
HOST="${HOST:-0.0.0.0}"
PORT="${PORT:-8091}"
STAGE_CONFIG="${STAGE_CONFIG:-vllm_omni/model_executor/stage_configs/ming_flash_omni_tts.yaml}"
echo "Starting Ming standalone TTS server with model: $MODEL"
echo "Stage config: $STAGE_CONFIG"
vllm serve "$MODEL" \
--stage-configs-path "$STAGE_CONFIG" \
--host "$HOST" \
--port "$PORT" \
--trust-remote-code \
--omni
examples/online_serving/ming_flash_omni_tts/speech_client.py
+93
View File
@@ -0,0 +1,93 @@
"""Client for Ming standalone TTS via /v1/audio/speech endpoint."""
import argparse
import json
import sys
import httpx
DEFAULT_API_BASE = "http://localhost:8091"
DEFAULT_API_KEY = "EMPTY"
DEFAULT_MODEL = "Jonathan1909/Ming-flash-omni-2.0"
def run_tts(args) -> None:
payload = {
"model": args.model,
"input": args.text,
"response_format": args.response_format,
}
instructions = args.instructions
if args.instruction_json:
if instructions:
sys.exit("--instructions and --instruction-json are mutually exclusive")
try:
parsed = json.loads(args.instruction_json)
except json.JSONDecodeError as exc:
sys.exit(f"--instruction-json must be valid JSON: {exc}")
if not isinstance(parsed, dict):
sys.exit("--instruction-json must decode to a JSON object")
# Re-encode with ensure_ascii=False so UTF-8 Chinese keys/values
# arrive at the server intact rather than as \\uXXXX escapes.
instructions = json.dumps(parsed, ensure_ascii=False)
if instructions:
payload["instructions"] = instructions
print(f"Model: {args.model}")
print(f"Text: {args.text}")
print("Generating audio...")
api_url = f"{args.api_base}/v1/audio/speech"
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {args.api_key}",
}
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
output_path = args.output or "ming_tts_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="Ming standalone TTS speech 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=DEFAULT_MODEL, help="Model name or local path")
parser.add_argument("--text", required=True, help="Text to synthesize")
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")
parser.add_argument(
"--instructions",
default=None,
help="Free-form style description (mapped to caption 风格 on the server).",
)
parser.add_argument(
"--instruction-json",
default=None,
help=(
"Structured caption JSON forwarded as `instructions`. Accepts Ming "
"caption keys: 方言, 风格, 语速, 基频, 音量, 情感, IP, 说话人, BGM. "
),
)
args = parser.parse_args()
run_tts(args)
if __name__ == "__main__":
main()
examples/online_serving/qwen3_omni/README.md
+70 -8
View File
@@ -12,21 +12,80 @@ Please refer to [README.md](../../../README.md)
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
```
The default deploy config at `vllm_omni/deploy/qwen3_omni_moe.yaml` is loaded
automatically by the model registry — no `--deploy-config` flag needed for the
common case. Async-chunk streaming is **enabled by default** in the bundled config.
NPU / ROCm / XPU per-platform deltas are merged in automatically from the
`platforms:` section of the same YAML.
The default deployment configuration, situated at `vllm_omni/deploy/qwen3_omni_moe.yaml`, is resolved and loaded
automatically via the model registry, obviating the `--deploy-config` flag in standard deployment topologies.
Asynchronous chunk streaming operates as **enabled by default** within this bundled configuration.
Additionally, NPU, ROCm, and XPU per-platform configuration deltas are deterministically merged from the
`platforms`: section of the corresponding YAML.
**Note:** The OpenAI-style **`/v1/realtime`** WebSocket (streaming PCM audio in, audio + transcription out) is **not supported** when `async_chunk` is enabled. Use the default omni layout or a stage config with `async_chunk: false` for realtime sessions.
If you have a custom deploy YAML, point at it explicitly:
**Note:** The OpenAI-style **`/v1/realtime`** WebSocket interface (facilitating streaming PCM audio input alongside audio and transcription output)
is currently **unsupported** while the `async_chunk` configuration attribute is enabled.
It is requisite to instantiate the default omni architecture or utilize a deployment configuration specifying `async_chunk: false` to facilitate real-time streaming sessions.
To explicitly utilize a custom deployment YAML, mandate the configuration path accordingly:
```bash
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
--deploy-config /path/to/your_deploy_config.yaml
```
### Launch individual stages (stage-based CLI)
Use the stage-based CLI when you want to run one stage per process.
**1. Stage 0 (Thinker + API server)**
```bash
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
--port 8091 \
--stage-id 0 \
--omni-master-address 127.0.0.1 \
--omni-master-port 26000
```
**2. Stage 1 (Talker)**
```bash
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
--stage-id 1 \
--headless \
--omni-master-address 127.0.0.1 \
--omni-master-port 26000
```
**3. Stage 2 (Code2Wav)**
```bash
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
--stage-id 2 \
--headless \
--omni-master-address 127.0.0.1 \
--omni-master-port 26000
```
Append `--deploy-config /path/to/your_deploy_config.yaml` to each node invocation if it is necessary
to explicitly override the bundled deployment YAML schema.
For standard **unified-process** launcher, stage-specific CLI configuration tuning is conventionally implemented
via the `--stage-overrides` directive, as demonstrated below:
```bash
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
--stage-overrides '{"1": {"gpu_memory_utilization": 0.5}}'
```
Conversely, within the stage-based CLI paradigm, `--stage-overrides` modifiers are typically **unnecessary**
for this category of optimization. Given that each instantiation strictly initiates a single functional stage,
parameter flags can be systematically assigned directly onto that specific stage's command sequence:
```bash
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni \
--stage-id 1 \
--headless \
--gpu-memory-utilization 0.5 \
--omni-master-address 127.0.0.1 \
--omni-master-port 26000
```
### Tuning deployment parameters
Most engine knobs (`max_num_batched_tokens`, `max_model_len`, `enforce_eager`,
@@ -93,6 +152,9 @@ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
Per-stage values are always treated as explicit and beat YAML defaults for
the named stage. Other stages keep their YAML values.
If you switch to the stage-based CLI, the same per-stage tuning can usually be
passed directly on that stage's command instead of using `--stage-overrides`.
#### 3. Custom deploy YAML
When per-stage overrides get long, write a small overlay YAML that inherits
examples/online_serving/qwen3_omni/streaming_video_client.py
+208
View File
@@ -0,0 +1,208 @@
#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Example WebSocket client for the /v1/video/chat/stream endpoint.
Sends video frames from a local file (or generates synthetic ones), submits a
query, and prints the streamed text response.
Requirements:
pip install websockets pillow
Usage:
# With a video file (requires opencv-python):
python streaming_video_client.py --video my_clip.mp4 \\
--query "What is happening in this video?"
# Synthetic frames (no extra deps):
python streaming_video_client.py \\
--query "Describe what you see." \\
--synthetic-frames 10
# With audio (Phase 3):
python streaming_video_client.py --video my_clip.mp4 \\
--audio my_audio.pcm \\
--query "What is the person saying and doing?"
"""
from __future__ import annotations
import argparse
import asyncio
import base64
import io
import json
import sys
try:
import websockets
except ImportError:
print("Please install websockets: pip install websockets")
sys.exit(1)
from PIL import Image
def _generate_synthetic_frame(index: int, width: int = 320, height: int = 240) -> bytes:
"""Generate a simple synthetic JPEG frame with a colour gradient."""
r = (index * 37) % 256
g = (index * 73) % 256
b = (index * 113) % 256
img = Image.new("RGB", (width, height), (r, g, b))
buf = io.BytesIO()
img.save(buf, format="JPEG", quality=80)
return buf.getvalue()
def _load_video_frames(path: str, max_frames: int = 64, fps: int = 2) -> list[bytes]:
"""Extract frames from a video file using OpenCV."""
try:
import cv2
except ImportError:
print("opencv-python is required to read video files: pip install opencv-python")
sys.exit(1)
cap = cv2.VideoCapture(path)
if not cap.isOpened():
print(f"Cannot open video: {path}")
sys.exit(1)
video_fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
frame_interval = max(1, int(video_fps / fps))
frames: list[bytes] = []
idx = 0
while len(frames) < max_frames:
ret, frame = cap.read()
if not ret:
break
if idx % frame_interval == 0:
_, buf = cv2.imencode(".jpg", frame, [cv2.IMWRITE_JPEG_QUALITY, 80])
frames.append(buf.tobytes())
idx += 1
cap.release()
print(f"Loaded {len(frames)} frames from {path} (interval={frame_interval})")
return frames
async def run(args: argparse.Namespace) -> None:
uri = f"ws://{args.host}:{args.port}/v1/video/chat/stream"
# Prepare frames
if args.video:
frames = _load_video_frames(args.video, max_frames=args.max_frames, fps=args.fps)
else:
frames = [_generate_synthetic_frame(i) for i in range(args.synthetic_frames)]
print(f"Generated {len(frames)} synthetic frames")
# Prepare audio (optional, Phase 3)
audio_data: bytes | None = None
if args.audio:
with open(args.audio, "rb") as f:
audio_data = f.read()
print(f"Loaded audio: {len(audio_data)} bytes")
async with websockets.connect(uri, max_size=16 * 1024 * 1024) as ws:
# 1. Send session.config
config = {
"type": "session.config",
"model": args.model,
"modalities": ["text", "audio"] if audio_data else ["text"],
"max_frames": args.max_frames,
"num_frames": args.num_sample_frames,
"enable_frame_filter": args.evs,
"frame_filter_threshold": args.evs_threshold,
"use_audio_in_video": bool(audio_data),
}
await ws.send(json.dumps(config))
print(f"Sent session.config: model={args.model} evs={args.evs}")
# 2. Send frames
for i, frame in enumerate(frames):
msg = {
"type": "video.frame",
"data": base64.b64encode(frame).decode(),
}
await ws.send(json.dumps(msg))
if (i + 1) % 10 == 0:
print(f" Sent {i + 1}/{len(frames)} frames")
print(f"Sent all {len(frames)} frames")
# 3. Send audio chunks (Phase 3)
if audio_data:
chunk_size = 16000 * 2 # 1 second of 16 kHz 16-bit PCM
for offset in range(0, len(audio_data), chunk_size):
chunk = audio_data[offset : offset + chunk_size]
msg = {
"type": "audio.chunk",
"data": base64.b64encode(chunk).decode(),
}
await ws.send(json.dumps(msg))
print(f"Sent audio in {(len(audio_data) + chunk_size - 1) // chunk_size} chunks")
# 4. Send query, then immediately send video.done so the server
# knows the session is complete (avoids deadlock where client
# waits for session.done while server waits for video.done).
await ws.send(json.dumps({"type": "video.query", "text": args.query}))
print(f"\nQuery: {args.query}")
print("Response: ", end="", flush=True)
# Signal end of session right after the query. The server will
# process the query first (it's already queued), then handle
# video.done and reply with session.done.
await ws.send(json.dumps({"type": "video.done"}))
# 5. Receive response until session.done
recv_timeout = 120 # seconds — avoid infinite hang if server stalls
while True:
raw = await asyncio.wait_for(ws.recv(), timeout=recv_timeout)
data = json.loads(raw)
msg_type = data.get("type")
if msg_type == "response.text.delta":
print(data.get("delta", ""), end="", flush=True)
elif msg_type == "response.text.done":
print() # newline
elif msg_type == "response.evs_stats":
retained = data.get("retained_count", 0)
dropped = data.get("dropped_count", 0)
rate = data.get("drop_rate", 0)
print(f"\nEVS stats: retained={retained} dropped={dropped} drop_rate={rate:.1%}")
elif msg_type == "session.done":
print("Session complete.")
break
elif msg_type == "error":
print(f"\nError: {data.get('message')}")
break
elif msg_type == "response.start":
pass # expected
else:
print(f"\n[unknown message] {data}")
def main() -> None:
parser = argparse.ArgumentParser(description="Streaming video chat client")
parser.add_argument("--host", default="localhost")
parser.add_argument("--port", type=int, default=8000)
parser.add_argument("--model", default="Qwen/Qwen3-Omni-MoE")
parser.add_argument("--video", help="Path to video file (requires opencv-python)")
parser.add_argument("--audio", help="Path to raw PCM 16kHz audio file (Phase 3)")
parser.add_argument("--query", default="What do you see in this video?")
parser.add_argument(
"--synthetic-frames", type=int, default=10, help="Number of synthetic frames if --video is not set"
)
parser.add_argument("--max-frames", type=int, default=64)
parser.add_argument("--num-sample-frames", type=int, default=16)
parser.add_argument("--fps", type=int, default=2, help="Frame extraction rate from video")
parser.add_argument(
"--no-evs", dest="evs", action="store_false", help="Disable EVS frame filtering (enabled by default)"
)
parser.set_defaults(evs=True)
parser.add_argument("--evs-threshold", type=float, default=0.95)
args = parser.parse_args()
asyncio.run(run(args))
if __name__ == "__main__":
main()
examples/online_serving/voxtral_tts/gradio_demo.py
+23 -4
View File
@@ -216,6 +216,7 @@ def make_update_voice_dropdown(language_voices: dict[str, list[str]]):
def run_inference(
voice_name: str,
text_prompt: str,
cfg_alpha: float,
base_url: str,
model: str,
) -> tuple[int, np.ndarray]:
@@ -233,6 +234,7 @@ def run_inference(
"model": model,
"response_format": "wav",
"voice": voice_name,
"extra_params": {"cfg_alpha": cfg_alpha},
}
response = httpx.post(
@@ -377,6 +379,14 @@ def main(
placeholder="Enter the text you want to synthesize...",
lines=4,
)
cfg_alpha_slider = gr.Slider(
minimum=1.0,
maximum=2.0,
step=0.1,
value=1.2,
label="CFG Alpha",
info="Flow-matching guidance strength (default: 1.2)",
)
with gr.Row():
reset_btn = gr.Button("Clear")
submit_btn = gr.Button("Generate audio", interactive=False)
@@ -415,9 +425,9 @@ def main(
)
# --- Wiring inference + persistence to the button ---
def _on_submit(voice: str, text: str):
def _on_submit(voice: str, text: str, cfg_alpha: float):
assert text.strip() != ""
sr, audio_array = run_inference(voice, text, base_url, model)
sr, audio_array = run_inference(voice, text, cfg_alpha, base_url, model)
if outputs_dir is not None:
share_id, saved_audio_path = _save_example(
outputs_dir,
@@ -432,7 +442,7 @@ def main(
submit_btn.click(
fn=_on_submit,
inputs=[voice_name, text_prompt],
inputs=[voice_name, text_prompt, cfg_alpha_slider],
outputs=[output_audio, share_link_box],
)
@@ -446,6 +456,7 @@ def main(
language, # language_dropdown
voice, # voice_name
"", # text_prompt
1.2, # cfg_alpha_slider
None, # output_audio
gr.update(interactive=False), # submit_btn
"", # share_link_box
@@ -456,7 +467,15 @@ def main(
reset_btn.click(
fn=make_on_reset(languages, language_voices),
inputs=[],
outputs=[language_dropdown, voice_name, text_prompt, output_audio, submit_btn, share_link_box],
outputs=[
language_dropdown,
voice_name,
text_prompt,
cfg_alpha_slider,
output_audio,
submit_btn,
share_link_box,
],
)
def make_load_from_share(outputs_dir: Path | None, languages: list[str], language_voices: dict[str, list[str]]):
pyproject.toml
+7 -11
View File
@@ -56,23 +56,17 @@ dev = [
"opencc>=1.2.0",
"mistune>=3.2.0", # for example tests
"torchmetrics>=1.4.0", # for accuracy similarity metrics
]
demo = [
"gradio>=6.7.0",
]
# Seed-TTS serve benchmark WER (BytedanceSpeech/seed-tts-eval run_wer.py protocol).
seed-tts-eval = [
"jiwer>=3.0.0",
"zhon>=2.0.0",
"zhconv>=1.4.2",
"scipy>=1.10.0",
"soundfile>=0.12.0",
"transformers>=4.36.0",
"funasr>=1.0.0",
]
demo = [
"gradio>=6.7.0",
]
docs = [
"mkdocs>=1.5.0",
"mkdocs-api-autonav",
@@ -179,7 +173,8 @@ addopts = [
markers = [
# ci/cd required
"core_model: L1&L2 tests (run in each PR)",
"advanced_model: L3&L4 level tests (run in each merge or nightly)",
"advanced_model: L3 level tests (run on each merge)",
"full_model: L4 level tests (run nightly)",
# function module markers
"diffusion: Diffusion model tests",
"omni: Omni model tests",
@@ -241,4 +236,5 @@ ue = "ue"
semantics = "semantics"
fullset = "fullset"
Vai = "Vai"
tockens = "tockens"
CANN = "CANN"
recipes/LTX/LTX-2.3.md
+112
View File
@@ -0,0 +1,112 @@
# LTX-2.3 Text-to-Video with Audio on 1x GPU (96GB VRAM)
> 22B parameter text-to-video + audio generation model served via vLLM-Omni
## Summary
- Vendor: Lightricks
- Model: `dg845/LTX-2.3-Diffusers`
- Task: Text-to-video with synchronized audio generation
- Mode: Online serving (pure diffusion)
- Maintainer: @oglok
## When to use this recipe
Use this recipe when you want to serve LTX-2.3 for text-to-video generation
with audio. The model generates videos up to 20+ seconds at 768x512 resolution
with 48kHz audio, all from a single text prompt. Requires a GPU with at least
96GB VRAM due to the 22B parameter transformer (~44GB weights) plus text
encoder, VAE, and vocoder components.
## References
- Model: <https://huggingface.co/dg845/LTX-2.3-Diffusers>
- Requires `diffusers >= 0.38.0` (install from git: `pip install git+https://github.com/huggingface/diffusers.git`)
## Serving
### Command
```bash
vllm serve dg845/LTX-2.3-Diffusers \
--omni \
--model-class-name LTX23Pipeline \
--stage-init-timeout 600
```
### Verification
```bash
# Health check
curl http://localhost:8000/health
# Generate a 3-second video (81 frames at 24fps)
curl -X POST http://localhost:8000/v1/videos \
-F "prompt=A majestic bald eagle soaring over a misty mountain valley at dawn, golden sunlight breaking through clouds" \
-F "negative_prompt=blurry, low quality, distorted, watermark" \
-F "model=dg845/LTX-2.3-Diffusers" \
-F "num_frames=81" \
-F "fps=24" \
-F "size=768x512" \
-F "num_inference_steps=30" \
-F "guidance_scale=4.0" \
-F "seed=42"
# Generate a 10-second video (241 frames)
curl -X POST http://localhost:8000/v1/videos \
-F "prompt=A cozy Japanese ramen shop at night in the rain, steam rising from bowls, neon signs reflecting on wet cobblestone streets" \
-F "model=dg845/LTX-2.3-Diffusers" \
-F "num_frames=241" \
-F "fps=24" \
-F "size=768x512" \
-F "num_inference_steps=30" \
-F "guidance_scale=4.0"
# Generate a 20-second video (481 frames)
curl -X POST http://localhost:8000/v1/videos \
-F "prompt=An underwater coral reef teeming with tropical fish, sea turtles gliding gracefully, National Geographic documentary style" \
-F "model=dg845/LTX-2.3-Diffusers" \
-F "num_frames=481" \
-F "fps=24" \
-F "size=768x512" \
-F "num_inference_steps=30" \
-F "guidance_scale=4.0"
```
### Notes
- Memory usage: Model loads at ~36 GiB, peaks at ~62 GiB during inference
- Key flags:
- `--stage-init-timeout 600`: Required for the initial `torch.compile` warmup (~90-140 seconds on first request)
- `--model-class-name LTX23Pipeline`: Selects the LTX-2.3 pipeline (not LTX-2)
- Audio: 48kHz AAC via BWE vocoder, automatically synced with video
- CPU offloading: Text encoder (Gemma-3-12B), connectors, VAE, audio VAE, and vocoder stay on CPU and are moved to GPU only when needed
- Supported resolutions: 768x512, 512x384 (must be divisible by 32)
- Frame rate: 24 fps
- Duration: Controlled by `num_frames` (frames = duration_seconds * 24 + 1)
- Known limitations:
- No image-to-video support yet (LTX23ImageToVideoPipeline is a placeholder)
- No CFG-parallel support (single-GPU only)
- Requires `diffusers >= 0.38.0` (not yet on PyPI, install from git)
## Hardware Support
## GPU
### 1x NVIDIA RTX PRO 6000 Blackwell (96GB)
#### Environment
- OS: Ubuntu 22.04
- Python: 3.10+
- Driver / runtime: CUDA 13.0, Driver 580.126.09
- vLLM version: 0.19.x
- vLLM-Omni version: 0.19.x
### Validated configurations
| Duration | Frames | Resolution | Steps | Guidance | Inference Time | Peak VRAM |
|----------|--------|------------|-------|----------|----------------|-----------|
| 3s | 81 | 768x512 | 30 | 4.0 | ~110s | ~62 GB |
| 10s | 241 | 768x512 | 30 | 4.0 | ~130s | ~62 GB |
| 20s | 481 | 768x512 | 30 | 4.0 | ~420s | ~62 GB |
recipes/Qwen/Qwen3-Omni.md
+17 -8
View File
@@ -50,13 +50,22 @@ Start the server from the repository root:
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
```
To enable async chunking, use the bundled stage config:
Async chunking is enabled by default in the bundled deployment config. For
common runtime tuning, prefer CLI overrides instead of editing or passing a
custom YAML file:
```bash
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct \
--omni \
--port 8091 \
--stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml
# Disable async chunking for /v1/realtime sessions
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
--no-async-chunk
```
Use a custom deploy config only for advanced cases such as custom topology,
connector wiring, or a larger overlay of stage defaults:
```bash
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
--deploy-config /path/to/your_qwen3_omni_overrides.yaml
```
#### Verification
@@ -85,6 +94,6 @@ curl http://localhost:8091/v1/chat/completions \
#### Notes
- Memory usage: Size depends on runtime options and output modalities; leave headroom for multimodal workloads.
- Key flags: `--omni` is required; `--stage-configs-path` is optional for custom or async-chunk stage configs.
- Known limitations: This starter recipe is intentionally narrow and focuses on the single-GPU online-serving path already documented in the repo examples.
- Memory usage: Size depends on runtime options and output modalities; leave headroom for multimodal workloads. Prefer CLI overrides such as `--gpu-memory-utilization` for routine tuning.
- Key flags: `--omni` is required; async chunking is enabled by default; use `--no-async-chunk` for realtime sessions and `--deploy-config` only for advanced custom deployments.
- Known limitations: The `/v1/realtime` WebSocket flow is currently unsupported while async chunking is enabled. This starter recipe is intentionally narrow and focuses on the single-GPU online-serving path already documented in the repo examples.
recipes/README.md
+2
View File
@@ -29,6 +29,8 @@ recipes/
multimodal chat on `1x A100 80GB`
- [`Wan-AI/Wan2.2-I2V.md`](./Wan-AI/Wan2.2-I2V.md): image-to-video serving
recipe for Wan2.2 14B on `8x Ascend NPU (A2/A3)`
- [`inclusionAI/Ming-flash-omni-2.0.md`](./inclusionAI/Ming-flash-omni-2.0.md):
online serving recipe for multimodal chat (`4x H100 80GB`) and standalone TTS (`1x H100 80GB`)
Within a single recipe file, include different hardware support sections such
as `GPU`, `ROCm`, and `NPU`, and add concrete tested configurations like
recipes/inclusionAI/Ming-flash-omni-2.0.md
+210
View File
@@ -0,0 +1,210 @@
# Ming-flash-omni 2.0 for omni-speech chat and standalone TTS
## Summary
- Vendor: inclusionAI
- Model: `Jonathan1909/Ming-flash-omni-2.0`
- Task: Multimodal chat with text, image, audio, or video input; standalone text-to-speech (TTS);
and image generation
- Mode: Online serving with the OpenAI-compatible API
- Maintainer: Community
## When to use this recipe
Use this recipe when you want a known-good starting point for serving
`Jonathan1909/Ming-flash-omni-2.0` with vLLM-Omni in one of three modes:
- **Thinker only** — multimodal understanding with text output.
- **Thinker + Talker (omni-speech)** — multimodal understanding with text and spoken output.
- **Talker only (TTS)** — standalone text-to-speech via the OpenAI `/v1/audio/speech` endpoint.
## References
- Upstream model:
[`inclusionAI/Ming`](https://github.com/inclusionAI/Ming)
- For offline inference and additional client variants, see
`examples/offline_inference/ming_flash_omni{,_tts}/` and
`examples/online_serving/ming_flash_omni{,_tts}/`.
## Hardware Support
This recipe documents reference GPU configurations for the two-stage
omni-speech deployment and the standalone TTS deployment.
Other hardware and configurations are welcome as community validation lands.
## GPU
### 4x H100 80GB — omni-speech/chat (thinker + talker)
The bundled `ming_flash_omni.yaml` runs the thinker with tensor parallel size
4 on GPUs 0–3 and the talker on GPU 3.
Adjust `devices` in the YAML to match your hardware.
#### Environment
- OS: Linux
- Python: 3.10+
- CUDA Driver Version: 590.48.01
- CUDA 12.5
- vLLM version: 0.19.0
- vLLM-Omni version or commit: 0.19.0rc1
#### Command
Thinker only (text output):
```bash
vllm serve Jonathan1909/Ming-flash-omni-2.0 --omni --port 8091
```
Thinker + talker (text and/or audio output):
```bash
vllm serve Jonathan1909/Ming-flash-omni-2.0 \
--omni \
--port 8091 \
--stage-configs-path vllm_omni/model_executor/stage_configs/ming_flash_omni.yaml \
--log-stats
```
`--log-stats` is optional but recommended while validating the deployment.
#### Verification
Text output from a multimodal (image) input:
```bash
curl http://localhost:8091/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{
"model": "Jonathan1909/Ming-flash-omni-2.0",
"messages": [
{"role": "system", "content": [{"type": "text", "text": "你是一个友好的AI助手。\n\ndetailed thinking off"}]},
{"role": "user", "content": [
{"type": "image_url", "image_url": {"url": "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg"}},
{"type": "text", "text": "Describe this image in detail."}
]}
],
"modalities": ["text"]
}'
```
Spoken response from a text query (save the WAV bytes):
```bash
curl http://localhost:8091/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{
"model": "Jonathan1909/Ming-flash-omni-2.0",
"messages": [
{"role": "system", "content": [{"type": "text", "text": "你是一个友好的AI助手。\n\ndetailed thinking off"}]},
{"role": "user", "content": "请详细介绍鹦鹉的生活习性。"}
],
"modalities": ["audio"]
}' | jq -r '.choices[0].message.audio.data' | base64 -d > ming_omni_parrot.wav
```
Text + audio output from an audio input (swap `audio_url` for `video_url`
or `image_url` to exercise the other multimodal input paths):
```bash
curl http://localhost:8091/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{
"model": "Jonathan1909/Ming-flash-omni-2.0",
"messages": [
{"role": "system", "content": [{"type": "text", "text": "你是一个友好的AI助手。\n\ndetailed thinking off"}]},
{"role": "user", "content": [
{"type": "audio_url", "audio_url": {"url": "https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/mary_had_lamb.ogg"}},
{"type": "text", "text": "Please recognize the language of this speech and transcribe it. Format: oral."}
]}
],
"modalities": ["text", "audio"]
}' | jq -r '.choices[0].message.content'
```
Streaming text output via SSE (set `"stream": true`):
```bash
curl -N http://localhost:8091/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{
"model": "Jonathan1909/Ming-flash-omni-2.0",
"messages": [
{"role": "system", "content": [{"type": "text", "text": "你是一个友好的AI助手。\n\ndetailed thinking off"}]},
{"role": "user", "content": "请详细介绍鹦鹉的生活习性。"}
],
"modalities": ["text"],
"stream": true
}'
```
Each SSE event carries a `data:` line with a chat-completion chunk; text
deltas appear at `choices[0].delta.content`.
#### Notes
- Output modality is selected by the request body: `"modalities": ["text"]`,
`["audio"]`, or `["text", "audio"]`. The two-stage omni-speech server must be launched
for any request containing `audio`.
- Reasoning mode: flip the system prompt suffix from `detailed thinking off`
to `detailed thinking on` in any request above.
- Memory usage: size depends on output modalities and multimodal input; leave
headroom for video frames and audio caches.
### 1x H100 80GB — standalone TTS (talker only)
The bundled `ming_flash_omni_tts.yaml` runs the talker on a single GPU and exposes the OpenAI `/v1/audio/speech` endpoint.
#### Environment
- OS: Linux
- Python: 3.10+
- CUDA Driver Version: 590.48.01
- CUDA 12.5
- vLLM version: 0.19.0
- vLLM-Omni version or commit: 0.19.0rc1
#### Command
```bash
vllm serve Jonathan1909/Ming-flash-omni-2.0 \
--omni \
--stage-configs-path vllm_omni/model_executor/stage_configs/ming_flash_omni_tts.yaml \
--port 8091 \
--log-stats
```
`--log-stats` is optional but recommended while validating the deployment.
#### Verification
Basic curl:
```bash
curl -X POST http://localhost:8091/v1/audio/speech \
-H "Content-Type: application/json" \
-d '{
"model": "Jonathan1909/Ming-flash-omni-2.0",
"input": "我会一直在这里陪着你。",
"response_format": "wav"
}' --output ming_online.wav
```
Speaker selection (e.g. `lingguang`):
```bash
curl -X POST http://localhost:8091/v1/audio/speech \
-H "Content-Type: application/json" \
-d '{
"model": "Jonathan1909/Ming-flash-omni-2.0",
"input": "春天来了，万物复苏，大地一片生机盎然。田野里的油菜花开得金灿灿的，蜜蜂在花丛中忙碌地采蜜。远处的山坡上，桃花和杏花竞相绽放，粉的白的交织在一起，美不胜收。清晨的微风带着泥土的芬芳，轻轻拂过脸颊，让人感到无比惬意。孩子们在田间小路上追逐嬉戏，老人们坐在门前晒太阳，享受着这份宁静与美好。",
"speaker": "lingguang",
"response_format": "wav"
}' --output ming_online_lingguang.wav
```
#### Notes
- The OpenAI `instructions` field is forwarded to the talker as the caption JSON — pass a raw string for `风格` (style) only, or a JSON-encoded object for multiple entries such as `方言` (dialect) and `情感` (emotion).
requirements/musa.txt
+3 -1
View File
@@ -1,4 +1,6 @@
-r common.txt
# MUSA platform dependencies
torchada>=0.1.46
torchada>=0.1.50
onnxruntime>=1.23.2
mate>=0.2.0
flash_attn_3>=0.1.4
tests/benchmarks/test_accuracy_bench_utils.py
+63 -1
View File
@@ -1,11 +1,17 @@
# ruff: noqa: E402, I001
import argparse
import math
import os
import sys
import types
from pathlib import Path
import pytest
from PIL import Image
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
REPO_ROOT = Path(__file__).resolve().parents[2]
if str(REPO_ROOT) not in sys.path:
sys.path.insert(0, str(REPO_ROOT))
@@ -36,8 +42,64 @@ from benchmarks.accuracy.text_to_image.gbench import (
summarize_generated_records as summarize_gebench_generated_records,
summarize_gebench_results,
)
from tests.e2e.accuracy.qwen3_omni.qwen3_omni_acc_bench_core import seed_tts_bench_argv
from tests.e2e.accuracy.qwen3_omni.run_qwen_omni_acc_benchmark import sync_dataset_env_from_ns
from vllm_omni.benchmarks.data_modules.seed_tts_dataset import resolve_seed_tts_root
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
def test_seed_tts_bench_argv_preserves_hf_repo_id_from_env(monkeypatch):
monkeypatch.setenv("VLLM_SEED_TTS_DATASET_PATH", "zhaochenyang20/seed-tts-eval")
monkeypatch.delenv("VLLM_SEED_TTS_REPO", raising=False)
argv = seed_tts_bench_argv(locale="en")
dataset_idx = argv.index("--dataset-path")
assert argv[dataset_idx + 1] == "zhaochenyang20/seed-tts-eval"
def test_sync_dataset_env_preserves_seed_tts_hf_repo_id(monkeypatch):
ns = argparse.Namespace(
daily_omni_repo=None,
daily_omni_qa_json=None,
daily_omni_video_dir=None,
seed_tts_dataset_path="zhaochenyang20/seed-tts-eval",
seed_tts_root=None,
)
monkeypatch.delenv("VLLM_SEED_TTS_DATASET_PATH", raising=False)
sync_dataset_env_from_ns(ns)
assert os.environ["VLLM_SEED_TTS_DATASET_PATH"] == "zhaochenyang20/seed-tts-eval"
def test_resolve_seed_tts_root_downloads_only_requested_locale(monkeypatch, tmp_path: Path):
downloaded_root = tmp_path / "seed_tts_cache"
(downloaded_root / "zh" / "prompt-wavs").mkdir(parents=True)
(downloaded_root / "zh" / "meta.lst").write_text("", encoding="utf-8")
captured: dict[str, object] = {}
def fake_snapshot_download(*, repo_id, repo_type, allow_patterns):
captured["repo_id"] = repo_id
captured["repo_type"] = repo_type
captured["allow_patterns"] = allow_patterns
return str(downloaded_root)
monkeypatch.setitem(
sys.modules,
"huggingface_hub",
types.SimpleNamespace(snapshot_download=fake_snapshot_download),
)
resolved = resolve_seed_tts_root(
"zhaochenyang20/seed-tts-eval",
explicit_root=None,
locale="zh",
)
assert resolved == downloaded_root.resolve()
assert captured["repo_id"] == "zhaochenyang20/seed-tts-eval"
assert captured["repo_type"] == "dataset"
assert captured["allow_patterns"] == ["zh/**"]
def test_summarize_gebench_generated_records_groups_by_type():
tests/benchmarks/test_diffusion_backends_metrics.py
+107
View File
@@ -0,0 +1,107 @@
import pytest
from benchmarks.diffusion.backends import RequestFuncInput, async_request_chat_completions
class _MockResponse:
def __init__(self, payload: dict, status: int = 200):
self._payload = payload
self.status = status
async def __aenter__(self):
return self
async def __aexit__(self, exc_type, exc, tb):
return False
async def json(self):
return self._payload
async def text(self):
return str(self._payload)
class _MockSession:
def __init__(self, payload: dict):
self._payload = payload
def post(self, *args, **kwargs):
return _MockResponse(self._payload)
@pytest.mark.core_model
@pytest.mark.benchmark
@pytest.mark.cpu
@pytest.mark.asyncio
async def test_chat_completions_metrics_fallback_to_top_level():
payload = {
"choices": [
{
"message": {
"content": [
{
"type": "image_url",
"image_url": {"url": "data:image/png;base64,abc"},
}
]
}
}
],
"metrics": {
"stage_durations": {"diffusion": 1.25},
"peak_memory_mb": 4096.0,
},
}
output = await async_request_chat_completions(
RequestFuncInput(
prompt="draw a cat",
api_url="http://test.local/v1/chat/completions",
model="ByteDance-Seed/BAGEL-7B-MoT",
),
session=_MockSession(payload),
)
assert output.success is True
assert output.stage_durations == {"diffusion": 1.25}
assert output.peak_memory_mb == 4096.0
@pytest.mark.core_model
@pytest.mark.benchmark
@pytest.mark.cpu
@pytest.mark.asyncio
async def test_chat_completions_metrics_message_level_takes_precedence():
payload = {
"choices": [
{
"message": {
"content": [
{
"type": "image_url",
"image_url": {"url": "data:image/png;base64,abc"},
"stage_durations": {"message_stage": 0.7},
"peak_memory_mb": 1234.0,
}
]
}
}
],
"metrics": {
"stage_durations": {"top_level_stage": 9.9},
"peak_memory_mb": 9999.0,
},
}
output = await async_request_chat_completions(
RequestFuncInput(
prompt="draw a dog",
api_url="http://test.local/v1/chat/completions",
model="ByteDance-Seed/BAGEL-7B-MoT",
),
session=_MockSession(payload),
)
assert output.success is True
assert output.stage_durations == {"message_stage": 0.7}
assert output.peak_memory_mb == 1234.0
tests/dfx/conftest.py
+174 -3
View File
@@ -1,4 +1,7 @@
import json
import os
import subprocess
from datetime import datetime
from pathlib import Path
from typing import Any
@@ -37,12 +40,45 @@ def modify_stage(default_path, updates, deletes):
return path
def _build_serve_args(serve_args: Any) -> list[str]:
"""Convert server_params.serve_args to a flat CLI args list."""
if serve_args is None:
return []
if isinstance(serve_args, list):
return [str(item) for item in serve_args]
if not isinstance(serve_args, dict):
raise TypeError(f"serve_args must be dict/list/None, got {type(serve_args).__name__}")
args: list[str] = []
for key, value in serve_args.items():
flag = f"--{str(key).replace('_', '-')}"
if isinstance(value, bool):
if value:
args.append(flag)
continue
if value is None:
continue
if isinstance(value, (dict, list)):
args.extend([flag, json.dumps(value, ensure_ascii=False, separators=(",", ":"))])
continue
args.extend([flag, str(value)])
return args
def create_unique_server_params(
configs: list[dict[str, Any]],
stage_configs_dir: Path,
) -> list[tuple[str, str, str | None, str | None, tuple[str, ...]]]:
unique_params = []
seen = set()
"""Return one row per unique server configuration (same 5-tuple shape as upstream).
``(test_name, model, deploy_yaml_path, stage_overrides_json, extra_cli_args)``.
JSON ``server_params.serve_args`` (dict/list) is expanded via ``_build_serve_args``
and **prepended** to ``extra_cli_args`` so perf / stability ``omni_server`` fixtures
stay identical to main while still honoring ``serve_args`` in benchmark JSON.
"""
unique_params: list[tuple[str, str, str | None, str | None, tuple[str, ...]]] = []
seen: set[tuple[str, str, str | None, str | None, tuple[str, ...]]] = set()
for config in configs:
test_name = config["test_name"]
server_params = config["server_params"]
@@ -63,7 +99,9 @@ def create_unique_server_params(
# ``vllm_omni.entrypoints.cli.main serve`` — used for flags that
# don't map to stage-level overrides, e.g. ``--async-chunk`` /
# ``--no-async-chunk`` toggling the deploy-level async_chunk bool.
extra_cli_args = tuple(server_params.get("extra_cli_args") or ())
serve_flat = _build_serve_args(server_params.get("serve_args"))
raw_extra = tuple(server_params.get("extra_cli_args") or ())
extra_cli_args = tuple(serve_flat) + raw_extra
server_param = (test_name, model, stage_config_path, stage_overrides_json, extra_cli_args)
if server_param not in seen:
@@ -109,6 +147,139 @@ def create_benchmark_indices(
return indices
def _safe_filename_token(value: Any | None, *, default: str = "na") -> str:
"""Make a single path segment safe for result filenames on common filesystems."""
if value is None:
return default
s = str(value).strip()
for bad in ("/", "\\", ":", "*", "?", '"', "<", ">", "|"):
s = s.replace(bad, "_")
return s if s else default
def _resolve_baseline_value(
baseline_raw: Any,
*,
sweep_index: int | None,
max_concurrency: Any = None,
request_rate: Any = None,
) -> Any:
"""Pick the baseline threshold for this sweep step."""
if baseline_raw is None:
return 100000
if isinstance(baseline_raw, dict):
if max_concurrency is not None:
for key in (max_concurrency, str(max_concurrency)):
if key in baseline_raw:
return baseline_raw[key]
if request_rate is not None:
for key in (request_rate, str(request_rate)):
if key in baseline_raw:
return baseline_raw[key]
raise KeyError(
f"baseline dict has no key for max_concurrency={max_concurrency!r} "
f"or request_rate={request_rate!r}; keys={list(baseline_raw.keys())!r}"
)
if isinstance(baseline_raw, (list, tuple)):
if sweep_index is None:
raise ValueError("list baseline requires sweep_index")
if not (0 <= sweep_index < len(baseline_raw)):
raise IndexError(f"baseline list len={len(baseline_raw)} has no index {sweep_index}")
return baseline_raw[sweep_index]
return baseline_raw
def _baseline_thresholds_for_step(
baseline_data: dict[str, Any],
*,
sweep_index: int | None = None,
max_concurrency: Any = None,
request_rate: Any = None,
) -> dict[str, Any]:
"""Resolve baseline config to one threshold per metric for this iteration."""
return {
metric_name: _resolve_baseline_value(
baseline_raw,
sweep_index=sweep_index,
max_concurrency=max_concurrency,
request_rate=request_rate,
)
for metric_name, baseline_raw in baseline_data.items()
}
def run_benchmark(
args: list[str],
test_name: str,
flow: Any,
dataset_name: str,
num_prompt: int,
*,
baseline_config: dict[str, Any] | None = None,
sweep_index: int | None = None,
request_rate: Any | None = None,
max_concurrency: Any | None = None,
random_input_len: Any | None = None,
random_output_len: Any | None = None,
) -> dict[str, Any]:
"""Run one ``vllm bench serve --omni`` iteration and return parsed metrics."""
current_dt = datetime.now().strftime("%Y%m%d-%H%M%S")
ri = _safe_filename_token(random_input_len)
ro = _safe_filename_token(random_output_len)
result_filename = f"result_{test_name}_{dataset_name}_{flow}_{num_prompt}_in{ri}_out{ro}_{current_dt}.json"
if "--result-filename" in args:
print(f"The result file will be overwritten by {result_filename}")
command = (
["vllm", "bench", "serve", "--omni"]
+ args
+ [
"--num-warmups",
"2",
"--save-result",
"--result-dir",
os.environ.get("BENCHMARK_DIR", "tests"),
"--result-filename",
result_filename,
]
)
process = subprocess.Popen(
command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, bufsize=1, universal_newlines=True
)
for line in iter(process.stdout.readline, ""):
print(line, end=" ")
for line in iter(process.stderr.readline, ""):
print(line, end=" ")
if "--result-dir" in command:
index = command.index("--result-dir")
result_dir = command[index + 1]
else:
result_dir = "./"
result_path = os.path.join(result_dir, result_filename)
with open(result_path, encoding="utf-8") as f:
result = json.load(f)
if baseline_config:
result["baseline"] = _baseline_thresholds_for_step(
baseline_config,
sweep_index=sweep_index,
request_rate=request_rate,
max_concurrency=max_concurrency,
)
else:
result["baseline"] = {}
if random_input_len is not None:
result["random_input_len"] = random_input_len
if random_output_len is not None:
result["random_output_len"] = random_output_len
with open(result_path, "w", encoding="utf-8") as f:
json.dump(result, f, ensure_ascii=False, indent=2)
return result
def pytest_addoption(parser: pytest.Parser) -> None:
"""Register shared CLI options for DFX benchmark suites."""
parser.addoption(
tests/dfx/perf/scripts/run_benchmark.py
+4
View File
@@ -17,6 +17,9 @@ from tests.dfx.conftest import (
)
from tests.helpers.runtime import OmniServer
pytestmark = [pytest.mark.full_model, pytest.mark.omni]
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
@@ -293,6 +296,7 @@ def assert_result(
print(f"ERROR: Test results exceeded baseline: {metric_name}: {current_value} < {baseline_value}")
@pytest.mark.benchmark
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
@pytest.mark.parametrize("benchmark_params", benchmark_indices, indirect=True)
def test_performance_benchmark(omni_server, benchmark_params):
tests/dfx/perf/scripts/run_diffusion_benchmark.py
+45 -21
View File
@@ -1,16 +1,17 @@
"""
Performance benchmark CI runner for diffusion models.
Supports vLLM-Omni server backend:
- vllm-omni (default): starts DiffusionServer via vllm_omni.entrypoints.cli.main,
benchmarks with diffusion_benchmark_serving.py --backend vllm-omni
This runner separates two concepts:
1. ``server_type``: how the serving process is started.
Currently only ``vllm-omni`` is supported here.
2. ``benchmark_backend``: which serving API the benchmark client calls.
Examples: ``vllm-omni`` for ``/v1/chat/completions`` and ``v1/videos``
for async video jobs.
A config JSON file is REQUIRED via --test-config-file:
pytest run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
JSON config entries use a "server_type" field, and this runner executes
the vllm-omni path.
All benchmark results for a session are consolidated into a single JSON file under
BENCHMARK_RESULT_DIR (override via the DIFFUSION_BENCHMARK_DIR environment variable).
Each entry in the file contains the test metadata (test_name, backend, benchmark_params,
@@ -32,6 +33,8 @@ from typing import Any, cast
import psutil
import pytest
pytestmark = [pytest.mark.diffusion, pytest.mark.full_model]
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
os.environ.setdefault("DIFFUSION_ATTENTION_BACKEND", "FLASH_ATTN")
@@ -387,17 +390,18 @@ def _unique_server_params(configs: list[dict[str, Any]]) -> list[dict[str, Any]]
if test_name in seen:
continue
seen.add(test_name)
if cfg.get("server_type", "vllm-omni") != "vllm-omni":
raise ValueError(f"Unsupported server_type in config: {cfg.get('server_type')}")
server_type = cfg.get("server_type", "vllm-omni")
if server_type != "vllm-omni":
raise ValueError(f"Unsupported server_type in config: {server_type}")
serve_args_dict = cfg["server_params"].get("serve_args", {})
result.append(
{
"test_name": test_name,
"server_type": "vllm-omni",
"server_type": server_type,
"model": cfg["server_params"]["model"],
"serve_args_dict": serve_args_dict,
"serve_args": _build_serve_args(serve_args_dict),
"benchmark_backend": "vllm-omni",
"benchmark_backend": cfg.get("benchmark_backend"),
"server_params": cfg["server_params"],
}
)
@@ -497,7 +501,8 @@ def run_benchmark(
log_dir = BENCHMARK_RESULT_DIR / "logs"
log_dir.mkdir(parents=True, exist_ok=True)
log_file = log_dir / f"{test_name}_{backend}_{timestamp}.log"
backend_label = backend.replace("/", "_")
log_file = log_dir / f"{test_name}_{backend_label}_{timestamp}.log"
with tempfile.NamedTemporaryFile(mode="w", suffix=".json", prefix="diffusion_bench_tmp_", delete=False) as tmp:
tmp_result_file = Path(tmp.name)
@@ -583,6 +588,7 @@ def run_benchmark(
tmp_result_file.unlink(missing_ok=True)
server_cfg = server_cfg or {}
server_type = cast(str, server_cfg.get("server_type", "vllm-omni"))
serve_args_dict = server_cfg.get("serve_args_dict", {})
if not isinstance(serve_args_dict, dict):
serve_args_dict = {}
@@ -599,18 +605,19 @@ def run_benchmark(
"result": metrics,
"log_file": str(log_file),
"Model": model,
"Framework": backend,
"Framework": server_type,
"API Backend": backend,
"Hardware": "",
"Deployment": "",
"Task": params.get("task", "t2i"),
"Dataset": params.get("dataset", "random"),
"resolution": _to_resolution_string(params),
"Parallelism": _to_parallelism_string(backend, serve_args_dict),
"Parallelism": _to_parallelism_string(server_type, serve_args_dict),
"max_concurrency": params.get("max-concurrency", ""),
"Cache": _to_cache_string(backend, serve_args_dict),
"Quantization": _to_quantization_value(backend, serve_args_dict),
"offload": _to_offload_string(backend, serve_args_dict),
"compile": _to_compile_value(backend, serve_args_dict),
"Cache": _to_cache_string(server_type, serve_args_dict),
"Quantization": _to_quantization_value(server_type, serve_args_dict),
"offload": _to_offload_string(server_type, serve_args_dict),
"compile": _to_compile_value(server_type, serve_args_dict),
"Attn_backend": os.environ.get("DIFFUSION_ATTENTION_BACKEND", ""),
"num_inference_steps": params.get("num-inference-steps", ""),
"completed": completed,
@@ -659,11 +666,27 @@ def assert_result(result: dict[str, Any], params: dict[str, Any]) -> None:
assert current <= threshold, f"{metric}: {current:.4f} > baseline {threshold}"
def _default_benchmark_backend_for_task(task: str) -> str:
"""Return the default client-side benchmark backend for a diffusion task."""
if task in {"t2v", "i2v", "ti2v"}:
return "v1/videos"
if task in {"t2i", "i2i", "ti2i"}:
return "vllm-omni"
raise ValueError(f"Unsupported task for benchmark backend resolution: {task}")
def _resolve_benchmark_backend(server_cfg: dict[str, Any], params: dict[str, Any]) -> str:
"""Resolve which serving API the benchmark client should call."""
configured = server_cfg.get("benchmark_backend")
if configured:
return cast(str, configured)
return _default_benchmark_backend_for_task(cast(str, params.get("task", "t2i")))
# ---------------------------------------------------------------------------
# Test entry point
# ---------------------------------------------------------------------------
@pytest.mark.benchmark
@pytest.mark.parametrize(
"diffusion_server",
server_params,
@@ -684,7 +707,8 @@ def test_diffusion_performance_benchmark(diffusion_server, benchmark_params):
"""
test_name = benchmark_params["test_name"]
params = benchmark_params["params"]
backend = diffusion_server.server_type # "vllm-omni"
server_cfg = getattr(diffusion_server, "server_cfg", {})
backend = _resolve_benchmark_backend(server_cfg, params)
result = run_benchmark(
host=diffusion_server.host,
@@ -693,7 +717,7 @@ def test_diffusion_performance_benchmark(diffusion_server, benchmark_params):
params=params,
test_name=test_name,
backend=backend,
server_cfg=getattr(diffusion_server, "server_cfg", {}),
server_cfg=server_cfg,
source_file=cast(str, CONFIG_FILE_PATH),
)
Loaded 100 of 307 files, more files were not shown because too many files have changed in this diff. Show more