diff --git a/.buildkite/scripts/hardware_ci/run-amd-test.sh b/.buildkite/scripts/hardware_ci/run-amd-test.sh index 2707dfefe..96c139c8f 100755 --- a/.buildkite/scripts/hardware_ci/run-amd-test.sh +++ b/.buildkite/scripts/hardware_ci/run-amd-test.sh @@ -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. diff --git a/.buildkite/test-merge.yml b/.buildkite/test-merge.yml index cbb44ac69..848297cec 100644 --- a/.buildkite/test-merge.yml +++ b/.buildkite/test-merge.yml @@ -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 diff --git a/.buildkite/test-nightly.yml b/.buildkite/test-nightly.yml index c33d7b4d1..ce67b76d9 100644 --- a/.buildkite/test-nightly.yml +++ b/.buildkite/test-nightly.yml @@ -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 diff --git a/.buildkite/test-ready.yml b/.buildkite/test-ready.yml index 3ca1747fe..080f18885 100644 --- a/.buildkite/test-ready.yml +++ b/.buildkite/test-ready.yml @@ -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: diff --git a/.claude/skills/add-diffusion-model/SKILL.md b/.claude/skills/add-diffusion-model/SKILL.md index a7e0bbf9a..0b979e1a9 100644 --- a/.claude/skills/add-diffusion-model/SKILL.md +++ b/.claude/skills/add-diffusion-model/SKILL.md @@ -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 diff --git a/.claude/skills/add-tts-model/SKILL.md b/.claude/skills/add-tts-model/SKILL.md index e64e7e763..963ffb4f6 100644 --- a/.claude/skills/add-tts-model/SKILL.md +++ b/.claude/skills/add-tts-model/SKILL.md @@ -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": , ...}` +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[]` 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//` @@ -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_.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 ` 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_.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 ` 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) diff --git a/.claude/skills/add-tts-model/references/cuda-graph-example.md b/.claude/skills/add-tts-model/references/cuda-graph-example.md new file mode 100644 index 000000000..6f4993b5c --- /dev/null +++ b/.claude/skills/add-tts-model/references/cuda-graph-example.md @@ -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. diff --git a/.claude/skills/add-tts-model/references/optional-deps.md b/.claude/skills/add-tts-model/references/optional-deps.md new file mode 100644 index 000000000..0a55f30f0 --- /dev/null +++ b/.claude/skills/add-tts-model/references/optional-deps.md @@ -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. diff --git a/.claude/skills/add-tts-model/references/precommit-dco.md b/.claude/skills/add-tts-model/references/precommit-dco.md new file mode 100644 index 000000000..86a1f42ce --- /dev/null +++ b/.claude/skills/add-tts-model/references/precommit-dco.md @@ -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//*.py \ + vllm_omni/entrypoints/openai/serving_speech.py \ + vllm_omni/model_executor/models/registry.py \ + tests/e2e/offline_inference/test_.py \ + tests/e2e/online_serving/test_.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 `. Use +`-s`: + +```bash +git commit -s -m "feat(): add 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 --force-with-lease +``` diff --git a/.claude/skills/add-tts-model/references/single-stage-ar.md b/.claude/skills/add-tts-model/references/single-stage-ar.md new file mode 100644 index 000000000..ed53d3026 --- /dev/null +++ b/.claude/skills/add-tts-model/references/single-stage-ar.md @@ -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_.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`. diff --git a/benchmarks/accuracy/README.md b/benchmarks/accuracy/README.md index 0d73215b6..dbe20916a 100644 --- a/benchmarks/accuracy/README.md +++ b/benchmarks/accuracy/README.md @@ -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. diff --git a/benchmarks/accuracy/image_to_image/README.md b/benchmarks/accuracy/image_to_image/README.md index ee1d58f10..86e7b0cf3 100644 --- a/benchmarks/accuracy/image_to_image/README.md +++ b/benchmarks/accuracy/image_to_image/README.md @@ -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. diff --git a/benchmarks/diffusion/backends.py b/benchmarks/diffusion/backends.py index 13ce7c830..d33160f13 100644 --- a/benchmarks/diffusion/backends.py +++ b/benchmarks/diffusion/backends.py @@ -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 diff --git a/benchmarks/diffusion/diffusion_benchmark_serving.py b/benchmarks/diffusion/diffusion_benchmark_serving.py index 32ec48a69..77b36b3d9 100644 --- a/benchmarks/diffusion/diffusion_benchmark_serving.py +++ b/benchmarks/diffusion/diffusion_benchmark_serving.py @@ -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 \ diff --git a/benchmarks/qwen3-omni/README.md b/benchmarks/qwen3-omni/README.md deleted file mode 100644 index dc282d052..000000000 --- a/benchmarks/qwen3-omni/README.md +++ /dev/null @@ -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//vllm_omni/` — vLLM-Omni pipeline benchmarks, logs, outputs. -- Add new tasks under `benchmarks//...` with the same pattern: `transformers/`, `vllm_omni/`, task-specific README, and (optionally) dataset prep notes. -- `benchmarks//vllm-omni-vs-hf.png` — current performance snapshot (overall throughput comparison). -- `benchmarks//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.*** diff --git a/benchmarks/qwen3-omni/transformers/eval_qwen3_moe_omni_transformers.sh b/benchmarks/qwen3-omni/transformers/eval_qwen3_moe_omni_transformers.sh deleted file mode 100644 index bae514fab..000000000 --- a/benchmarks/qwen3-omni/transformers/eval_qwen3_moe_omni_transformers.sh +++ /dev/null @@ -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 diff --git a/benchmarks/qwen3-omni/transformers/qwen3_omni_moe_model.py b/benchmarks/qwen3-omni/transformers/qwen3_omni_moe_model.py deleted file mode 100644 index 43b56f3e9..000000000 --- a/benchmarks/qwen3-omni/transformers/qwen3_omni_moe_model.py +++ /dev/null @@ -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", -] diff --git a/benchmarks/qwen3-omni/transformers/qwen3_omni_moe_transformers.py b/benchmarks/qwen3-omni/transformers/qwen3_omni_moe_transformers.py deleted file mode 100644 index 87d3de797..000000000 --- a/benchmarks/qwen3-omni/transformers/qwen3_omni_moe_transformers.py +++ /dev/null @@ -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() diff --git a/benchmarks/qwen3-omni/vllm-omni-vs-hf.png b/benchmarks/qwen3-omni/vllm-omni-vs-hf.png deleted file mode 100644 index e47079335..000000000 Binary files a/benchmarks/qwen3-omni/vllm-omni-vs-hf.png and /dev/null differ diff --git a/benchmarks/qwen3-omni/vllm_omni/eval_qwen3_moe_omni.sh b/benchmarks/qwen3-omni/vllm_omni/eval_qwen3_moe_omni.sh deleted file mode 100644 index e4c83e975..000000000 --- a/benchmarks/qwen3-omni/vllm_omni/eval_qwen3_moe_omni.sh +++ /dev/null @@ -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 diff --git a/docs/.nav.yml b/docs/.nav.yml index b0022965a..b44e8e6b5 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -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 diff --git a/docs/cli/serve.md b/docs/cli/serve.md index 47a873b72..035fa0567 100644 --- a/docs/cli/serve.md +++ b/docs/cli/serve.md @@ -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" diff --git a/docs/configuration/stage_configs.md b/docs/configuration/stage_configs.md index 55b4053cc..4a7c9cc67 100644 --- a/docs/configuration/stage_configs.md +++ b/docs/configuration/stage_configs.md @@ -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! diff --git a/docs/contributing/ci/CI_5levels.md b/docs/contributing/ci/CI_5levels.md index 2452ef5d4..3baa7ff88 100644 --- a/docs/contributing/ci/CI_5levels.md +++ b/docs/contributing/ci/CI_5levels.md @@ -105,7 +105,8 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari Depends on reality Stability:
- /tests/dfx/stability/tests/test.json
+ /tests/dfx/stability/tests/test_qwen3_omni.json
+ /tests/dfx/stability/tests/test_wan22.json
Reliability:
tests/e2e/reliability/test_{model_name}.py @@ -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
Test Examples -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.
-- - ***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_.py` alongside a matching JSON config) - ***Reliability***: `pytest -s -v tests/e2e/reliability/test_{model_name}.py` ## Summary diff --git a/docs/contributing/ci/test_examples/l4_functionality_tests.inc.md b/docs/contributing/ci/test_examples/l4_functionality_tests.inc.md index ab4deecd6..e1309b1ad 100644 --- a/docs/contributing/ci/test_examples/l4_functionality_tests.inc.md +++ b/docs/contributing/ci/test_examples/l4_functionality_tests.inc.md @@ -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`. diff --git a/docs/contributing/ci/test_guide.md b/docs/contributing/ci/test_guide.md index 08b2e3b4e..018c47b05 100644 --- a/docs/contributing/ci/test_guide.md +++ b/docs/contributing/ci/test_guide.md @@ -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) diff --git a/docs/contributing/ci/tests_markers.md b/docs/contributing/ci/tests_markers.md index 7628db284..6130541a6 100644 --- a/docs/contributing/ci/tests_markers.md +++ b/docs/contributing/ci/tests_markers.md @@ -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 | diff --git a/docs/contributing/model/adding_diffusion_model.md b/docs/contributing/model/adding_diffusion_model.md index 8d85eb4f6..6d5782a6e 100644 --- a/docs/contributing/model/adding_diffusion_model.md +++ b/docs/contributing/model/adding_diffusion_model.md @@ -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 diff --git a/docs/contributing/model/adding_tts_model.md b/docs/contributing/model/adding_tts_model.md index 622064173..34fd2dbb5 100644 --- a/docs/contributing/model/adding_tts_model.md +++ b/docs/contributing/model/adding_tts_model.md @@ -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_.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__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 diff --git a/docs/contributing/profiling.md b/docs/contributing/profiling.md index a7b508302..e1dbc8234 100644 --- a/docs/contributing/profiling.md +++ b/docs/contributing/profiling.md @@ -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_.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_.txt` when `torch_profiler_dump_cuda_time_total` is enabled CUDA profiler / Nsight Systems output: @@ -211,4 +283,4 @@ Recommended viewers: - `nsys stats .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/). diff --git a/docs/mkdocs/hooks/generate_argparse.py b/docs/mkdocs/hooks/generate_argparse.py index 0aed44a0c..e3cfb1b6a 100644 --- a/docs/mkdocs/hooks/generate_argparse.py +++ b/docs/mkdocs/hooks/generate_argparse.py @@ -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, diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index b4976b09c..feea969e5 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -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` | ✅︎ | ✅︎ | | | diff --git a/docs/serving/video_stream_api.md b/docs/serving/video_stream_api.md new file mode 100644 index 000000000..88f74affc --- /dev/null +++ b/docs/serving/video_stream_api.md @@ -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 --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. diff --git a/docs/user_guide/diffusion/attention_backends.md b/docs/user_guide/diffusion/attention_backends.md new file mode 100644 index 000000000..692bcc0f9 --- /dev/null +++ b/docs/user_guide/diffusion/attention_backends.md @@ -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 diff --git a/docs/user_guide/diffusion/cpu_offload_diffusion.md b/docs/user_guide/diffusion/cpu_offload_diffusion.md index f80005ccb..b85c7e408 100644 --- a/docs/user_guide/diffusion/cpu_offload_diffusion.md +++ b/docs/user_guide/diffusion/cpu_offload_diffusion.md @@ -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** diff --git a/docs/user_guide/diffusion/quantization/autoround.md b/docs/user_guide/diffusion/quantization/autoround.md index 48df176b0..d06627d40 100644 --- a/docs/user_guide/diffusion/quantization/autoround.md +++ b/docs/user_guide/diffusion/quantization/autoround.md @@ -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 diff --git a/docs/user_guide/diffusion/quantization/fp8.md b/docs/user_guide/diffusion/quantization/fp8.md index 9906631b6..ceb3d006c 100644 --- a/docs/user_guide/diffusion/quantization/fp8.md +++ b/docs/user_guide/diffusion/quantization/fp8.md @@ -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 diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 6e4a54a65..11ed1330c 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -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 diff --git a/docs/user_guide/examples/offline_inference/bagel.md b/docs/user_guide/examples/offline_inference/bagel.md index 1fb4d4045..0d3498b28 100644 --- a/docs/user_guide/examples/offline_inference/bagel.md +++ b/docs/user_guide/examples/offline_inference/bagel.md @@ -2,46 +2,61 @@ Source . - -## 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 `...` 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 `...` 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 diff --git a/docs/user_guide/examples/offline_inference/text_to_video.md b/docs/user_guide/examples/offline_inference/text_to_video.md index 4288c089c..a09dbfc97 100644 --- a/docs/user_guide/examples/offline_inference/text_to_video.md +++ b/docs/user_guide/examples/offline_inference/text_to_video.md @@ -5,6 +5,8 @@ Source . - -## 🛠️ 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 `` 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= \ - --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 `` 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 \ - -omp 8091 + --omni-master-address \ + --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 \ - -omp 8091 + --omni-master-address \ + --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 diff --git a/docs/user_guide/examples/online_serving/diffusers_pipeline_adapter.md b/docs/user_guide/examples/online_serving/diffusers_pipeline_adapter.md new file mode 100644 index 000000000..ac88071d5 --- /dev/null +++ b/docs/user_guide/examples/online_serving/diffusers_pipeline_adapter.md @@ -0,0 +1,93 @@ +# Diffusers Backend Adapter Example + +Source . + + +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" + `````` diff --git a/docs/user_guide/examples/online_serving/glm_image.md b/docs/user_guide/examples/online_serving/glm_image.md index f7027b906..b67eed1e0 100644 --- a/docs/user_guide/examples/online_serving/glm_image.md +++ b/docs/user_guide/examples/online_serving/glm_image.md @@ -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`. diff --git a/docs/user_guide/examples/online_serving/qwen3_omni.md b/docs/user_guide/examples/online_serving/qwen3_omni.md index 611eb6fd3..22d89ee80 100644 --- a/docs/user_guide/examples/online_serving/qwen3_omni.md +++ b/docs/user_guide/examples/online_serving/qwen3_omni.md @@ -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 diff --git a/examples/offline_inference/bagel/README.md b/examples/offline_inference/bagel/README.md index 3e653d0e3..9955fd90d 100644 --- a/examples/offline_inference/bagel/README.md +++ b/examples/offline_inference/bagel/README.md @@ -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 `...` 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 `...` 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. diff --git a/examples/offline_inference/bagel/end2end.py b/examples/offline_inference/bagel/end2end.py index ed5fa57e8..a6ce1f131 100644 --- a/examples/offline_inference/bagel/end2end.py +++ b/examples/offline_inference/bagel/end2end.py @@ -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] diff --git a/examples/offline_inference/cosyvoice3/verify_e2e_cosyvoice.py b/examples/offline_inference/cosyvoice3/verify_e2e_cosyvoice.py index de2712891..a5dc564ec 100644 --- a/examples/offline_inference/cosyvoice3/verify_e2e_cosyvoice.py +++ b/examples/offline_inference/cosyvoice3/verify_e2e_cosyvoice.py @@ -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, ) diff --git a/examples/offline_inference/custom_pipeline/image_to_image/image_edit.py b/examples/offline_inference/custom_pipeline/image_to_image/image_edit.py index 8ab5e0d9a..dc1085c28 100644 --- a/examples/offline_inference/custom_pipeline/image_to_image/image_edit.py +++ b/examples/offline_inference/custom_pipeline/image_to_image/image_edit.py @@ -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}") diff --git a/examples/offline_inference/dynin_omni/end2end.py b/examples/offline_inference/dynin_omni/end2end.py index 66047934d..82cff0c00 100644 --- a/examples/offline_inference/dynin_omni/end2end.py +++ b/examples/offline_inference/dynin_omni/end2end.py @@ -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() diff --git a/examples/offline_inference/glm_image/README.md b/examples/offline_inference/glm_image/README.md deleted file mode 100644 index c3c7c2916..000000000 --- a/examples/offline_inference/glm_image/README.md +++ /dev/null @@ -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 diff --git a/examples/offline_inference/glm_image/end2end.py b/examples/offline_inference/glm_image/end2end.py deleted file mode 100644 index 7ae9478ba..000000000 --- a/examples/offline_inference/glm_image/end2end.py +++ /dev/null @@ -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) diff --git a/examples/offline_inference/glm_image/run_i2i.sh b/examples/offline_inference/glm_image/run_i2i.sh deleted file mode 100755 index f81b157b0..000000000 --- a/examples/offline_inference/glm_image/run_i2i.sh +++ /dev/null @@ -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 diff --git a/examples/offline_inference/glm_image/run_t2i.sh b/examples/offline_inference/glm_image/run_t2i.sh deleted file mode 100755 index 5d249960b..000000000 --- a/examples/offline_inference/glm_image/run_t2i.sh +++ /dev/null @@ -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 diff --git a/examples/offline_inference/helios/end2end.py b/examples/offline_inference/helios/end2end.py index 88c3b865d..6cf7dfdcb 100644 --- a/examples/offline_inference/helios/end2end.py +++ b/examples/offline_inference/helios/end2end.py @@ -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() diff --git a/examples/offline_inference/hunyuan_image3/end2end.py b/examples/offline_inference/hunyuan_image3/end2end.py index 3c1ae3866..2cea30388 100644 --- a/examples/offline_inference/hunyuan_image3/end2end.py +++ b/examples/offline_inference/hunyuan_image3/end2end.py @@ -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() diff --git a/examples/offline_inference/image_to_image/image_edit.py b/examples/offline_inference/image_to_image/image_edit.py index 1a7e86f13..b857bb22d 100644 --- a/examples/offline_inference/image_to_image/image_edit.py +++ b/examples/offline_inference/image_to_image/image_edit.py @@ -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}") diff --git a/examples/offline_inference/image_to_video/image_to_video.py b/examples/offline_inference/image_to_video/image_to_video.py index 53319c822..fea5178d8 100644 --- a/examples/offline_inference/image_to_video/image_to_video.py +++ b/examples/offline_inference/image_to_video/image_to_video.py @@ -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: diff --git a/examples/offline_inference/magi_human/end2end.py b/examples/offline_inference/magi_human/end2end.py index 64f11c465..7ea816138 100644 --- a/examples/offline_inference/magi_human/end2end.py +++ b/examples/offline_inference/magi_human/end2end.py @@ -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() diff --git a/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py b/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py index ca87b9e9a..a742b535f 100644 --- a/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py +++ b/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py @@ -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() diff --git a/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py b/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py index a4c41fee1..bd1282a21 100644 --- a/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py +++ b/examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py @@ -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."] diff --git a/examples/offline_inference/mimo_audio/end2end.py b/examples/offline_inference/mimo_audio/end2end.py index 2c67d66fa..9c652fe2b 100644 --- a/examples/offline_inference/mimo_audio/end2end.py +++ b/examples/offline_inference/mimo_audio/end2end.py @@ -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 "" in output_text or "" in output_text: + output_text = output_text.replace("", "").replace("", "").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( diff --git a/examples/offline_inference/ming_flash_omni/README.md b/examples/offline_inference/ming_flash_omni/README.md index 7414163fc..be90b408d 100644 --- a/examples/offline_inference/ming_flash_omni/README.md +++ b/examples/offline_inference/ming_flash_omni/README.md @@ -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 | `.txt` | +| `audio` | Text (internal) | Runs | `.wav` | +| `text,audio` | Text | Runs | `.txt` + `.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 diff --git a/examples/offline_inference/ming_flash_omni/end2end.py b/examples/offline_inference/ming_flash_omni/end2end.py index 8f8730131..e00dcea7b 100644 --- a/examples/offline_inference/ming_flash_omni/end2end.py +++ b/examples/offline_inference/ming_flash_omni/end2end.py @@ -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: diff --git a/examples/offline_inference/ming_flash_omni_tts/README.md b/examples/offline_inference/ming_flash_omni_tts/README.md new file mode 100644 index 000000000..15b84041d --- /dev/null +++ b/examples/offline_inference/ming_flash_omni_tts/README.md @@ -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/`. diff --git a/examples/offline_inference/ming_flash_omni_tts/end2end.py b/examples/offline_inference/ming_flash_omni_tts/end2end.py new file mode 100644 index 000000000..928994510 --- /dev/null +++ b/examples/offline_inference/ming_flash_omni_tts/end2end.py @@ -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() diff --git a/examples/offline_inference/omnivoice/end2end.py b/examples/offline_inference/omnivoice/end2end.py index 9371c9514..cc6f585c5 100644 --- a/examples/offline_inference/omnivoice/end2end.py +++ b/examples/offline_inference/omnivoice/end2end.py @@ -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, ) diff --git a/examples/offline_inference/qwen2_5_omni/end2end.py b/examples/offline_inference/qwen2_5_omni/end2end.py index dfe124700..a65c554a9 100644 --- a/examples/offline_inference/qwen2_5_omni/end2end.py +++ b/examples/offline_inference/qwen2_5_omni/end2end.py @@ -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", diff --git a/examples/offline_inference/qwen3_omni/end2end.py b/examples/offline_inference/qwen3_omni/end2end.py index f028c32aa..04aa7914d 100644 --- a/examples/offline_inference/qwen3_omni/end2end.py +++ b/examples/offline_inference/qwen3_omni/end2end.py @@ -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() diff --git a/examples/offline_inference/text_to_audio/text_to_audio.py b/examples/offline_inference/text_to_audio/text_to_audio.py index 53f3daf8d..2a1613e5e 100644 --- a/examples/offline_inference/text_to_audio/text_to_audio.py +++ b/examples/offline_inference/text_to_audio/text_to_audio.py @@ -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() diff --git a/examples/offline_inference/text_to_image/text_to_image.py b/examples/offline_inference/text_to_image/text_to_image.py index 3b3f8e77c..0f9f10dce 100644 --- a/examples/offline_inference/text_to_image/text_to_image.py +++ b/examples/offline_inference/text_to_image/text_to_image.py @@ -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, diff --git a/examples/offline_inference/text_to_video/text_to_video.py b/examples/offline_inference/text_to_video/text_to_video.py index 83925cc45..a96949fa5 100644 --- a/examples/offline_inference/text_to_video/text_to_video.py +++ b/examples/offline_inference/text_to_video/text_to_video.py @@ -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 diff --git a/examples/offline_inference/vace/vace_video_generation.py b/examples/offline_inference/vace/vace_video_generation.py index 6ca0d74c5..5fad37366 100644 --- a/examples/offline_inference/vace/vace_video_generation.py +++ b/examples/offline_inference/vace/vace_video_generation.py @@ -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() diff --git a/examples/offline_inference/voxtral_tts/end2end.py b/examples/offline_inference/voxtral_tts/end2end.py index cf28f917e..0a6f88715 100644 --- a/examples/offline_inference/voxtral_tts/end2end.py +++ b/examples/offline_inference/voxtral_tts/end2end.py @@ -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, diff --git a/examples/online_serving/bagel/README.md b/examples/online_serving/bagel/README.md index 0939bc5f3..763927222 100644 --- a/examples/online_serving/bagel/README.md +++ b/examples/online_serving/bagel/README.md @@ -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 `` 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= \ - --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 `` 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 \ - -omp 8091 + --omni-master-address \ + --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 \ - -omp 8091 + --omni-master-address \ + --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. diff --git a/examples/online_serving/bagel/run_server_stage_cli.sh b/examples/online_serving/bagel/run_server_stage_cli.sh index 2d0b4bc36..912e212f9 100644 --- a/examples/online_serving/bagel/run_server_stage_cli.sh +++ b/examples/online_serving/bagel/run_server_stage_cli.sh @@ -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[@]}" } diff --git a/examples/online_serving/diffusers_pipeline_adapter/README.md b/examples/online_serving/diffusers_pipeline_adapter/README.md new file mode 100644 index 000000000..8dbf9369a --- /dev/null +++ b/examples/online_serving/diffusers_pipeline_adapter/README.md @@ -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). diff --git a/examples/online_serving/diffusers_pipeline_adapter/stage_config.yaml b/examples/online_serving/diffusers_pipeline_adapter/stage_config.yaml new file mode 100644 index 000000000..7c96eb6c1 --- /dev/null +++ b/examples/online_serving/diffusers_pipeline_adapter/stage_config.yaml @@ -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 diff --git a/examples/online_serving/glm_image/README.md b/examples/online_serving/glm_image/README.md deleted file mode 100644 index 7ba4b501c..000000000 --- a/examples/online_serving/glm_image/README.md +++ /dev/null @@ -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 -``` diff --git a/examples/online_serving/glm_image/openai_chat_client.py b/examples/online_serving/glm_image/openai_chat_client.py deleted file mode 100644 index e142b0719..000000000 --- a/examples/online_serving/glm_image/openai_chat_client.py +++ /dev/null @@ -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() diff --git a/examples/online_serving/glm_image/run_curl_image_edit.sh b/examples/online_serving/glm_image/run_curl_image_edit.sh deleted file mode 100755 index bb1e851ba..000000000 --- a/examples/online_serving/glm_image/run_curl_image_edit.sh +++ /dev/null @@ -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 \"\" [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" diff --git a/examples/online_serving/glm_image/run_curl_text_to_image.sh b/examples/online_serving/glm_image/run_curl_text_to_image.sh deleted file mode 100755 index aecb6953c..000000000 --- a/examples/online_serving/glm_image/run_curl_text_to_image.sh +++ /dev/null @@ -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" diff --git a/examples/online_serving/glm_image/run_server.sh b/examples/online_serving/glm_image/run_server.sh deleted file mode 100755 index b47d9f885..000000000 --- a/examples/online_serving/glm_image/run_server.sh +++ /dev/null @@ -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" diff --git a/examples/online_serving/ming_flash_omni/README.md b/examples/online_serving/ming_flash_omni/README.md index 502232725..8b7d03e21 100644 --- a/examples/online_serving/ming_flash_omni/README.md +++ b/examples/online_serving/ming_flash_omni/README.md @@ -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. diff --git a/examples/online_serving/ming_flash_omni_tts/README.md b/examples/online_serving/ming_flash_omni_tts/README.md new file mode 100644 index 000000000..e56d9ea12 --- /dev/null +++ b/examples/online_serving/ming_flash_omni_tts/README.md @@ -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. diff --git a/examples/online_serving/ming_flash_omni_tts/run_server.sh b/examples/online_serving/ming_flash_omni_tts/run_server.sh new file mode 100755 index 000000000..91be35c6c --- /dev/null +++ b/examples/online_serving/ming_flash_omni_tts/run_server.sh @@ -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 diff --git a/examples/online_serving/ming_flash_omni_tts/speech_client.py b/examples/online_serving/ming_flash_omni_tts/speech_client.py new file mode 100644 index 000000000..b3fceba25 --- /dev/null +++ b/examples/online_serving/ming_flash_omni_tts/speech_client.py @@ -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() diff --git a/examples/online_serving/qwen3_omni/README.md b/examples/online_serving/qwen3_omni/README.md index 32722b3db..c85970555 100644 --- a/examples/online_serving/qwen3_omni/README.md +++ b/examples/online_serving/qwen3_omni/README.md @@ -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 diff --git a/examples/online_serving/qwen3_omni/streaming_video_client.py b/examples/online_serving/qwen3_omni/streaming_video_client.py new file mode 100644 index 000000000..58f26d245 --- /dev/null +++ b/examples/online_serving/qwen3_omni/streaming_video_client.py @@ -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() diff --git a/examples/online_serving/voxtral_tts/gradio_demo.py b/examples/online_serving/voxtral_tts/gradio_demo.py index 35d6b590c..7905c6261 100644 --- a/examples/online_serving/voxtral_tts/gradio_demo.py +++ b/examples/online_serving/voxtral_tts/gradio_demo.py @@ -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]]): diff --git a/pyproject.toml b/pyproject.toml index b6f4092fd..8346693f1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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" diff --git a/recipes/LTX/LTX-2.3.md b/recipes/LTX/LTX-2.3.md new file mode 100644 index 000000000..8d92562fb --- /dev/null +++ b/recipes/LTX/LTX-2.3.md @@ -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: +- 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 | diff --git a/recipes/Qwen/Qwen3-Omni.md b/recipes/Qwen/Qwen3-Omni.md index 081e1453d..f78e4dda2 100644 --- a/recipes/Qwen/Qwen3-Omni.md +++ b/recipes/Qwen/Qwen3-Omni.md @@ -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. diff --git a/recipes/README.md b/recipes/README.md index 01ecc41f1..69ce4d750 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -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 diff --git a/recipes/inclusionAI/Ming-flash-omni-2.0.md b/recipes/inclusionAI/Ming-flash-omni-2.0.md new file mode 100644 index 000000000..873158c8a --- /dev/null +++ b/recipes/inclusionAI/Ming-flash-omni-2.0.md @@ -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). diff --git a/requirements/musa.txt b/requirements/musa.txt index 112f32604..c100c70cf 100644 --- a/requirements/musa.txt +++ b/requirements/musa.txt @@ -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 diff --git a/tests/benchmarks/test_accuracy_bench_utils.py b/tests/benchmarks/test_accuracy_bench_utils.py index a0479fb1b..6ceebb11b 100644 --- a/tests/benchmarks/test_accuracy_bench_utils.py +++ b/tests/benchmarks/test_accuracy_bench_utils.py @@ -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(): diff --git a/tests/benchmarks/test_diffusion_backends_metrics.py b/tests/benchmarks/test_diffusion_backends_metrics.py new file mode 100644 index 000000000..2d51d0f1d --- /dev/null +++ b/tests/benchmarks/test_diffusion_backends_metrics.py @@ -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 diff --git a/tests/dfx/conftest.py b/tests/dfx/conftest.py index c3f6d0a15..607570443 100644 --- a/tests/dfx/conftest.py +++ b/tests/dfx/conftest.py @@ -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( diff --git a/tests/dfx/perf/scripts/run_benchmark.py b/tests/dfx/perf/scripts/run_benchmark.py index 13011b4bd..f14638c15 100644 --- a/tests/dfx/perf/scripts/run_benchmark.py +++ b/tests/dfx/perf/scripts/run_benchmark.py @@ -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): diff --git a/tests/dfx/perf/scripts/run_diffusion_benchmark.py b/tests/dfx/perf/scripts/run_diffusion_benchmark.py index 23efa8bb0..7513c2d3f 100644 --- a/tests/dfx/perf/scripts/run_diffusion_benchmark.py +++ b/tests/dfx/perf/scripts/run_diffusion_benchmark.py @@ -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), ) diff --git a/tests/dfx/perf/tests/test_qwen_omni.json b/tests/dfx/perf/tests/test_qwen_omni.json index ca3eb5557..eda9720c4 100644 --- a/tests/dfx/perf/tests/test_qwen_omni.json +++ b/tests/dfx/perf/tests/test_qwen_omni.json @@ -10,16 +10,16 @@ "dataset_name": "random", "backend": "openai-chat-omni", "endpoint": "/v1/chat/completions", - "num_prompts": [4, 16, 40], - "max_concurrency": [1, 4, 10], + "num_prompts": [4, 16, 32, 64, 128], + "max_concurrency": [1, 4, 8, 16, 32], "random_input_len": 2500, "random_output_len": 900, "ignore_eos": true, "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "mean_ttft_ms": [1000, 3000, 5000], - "mean_audio_ttfp_ms": [30000, 60000, 90000], - "mean_audio_rtf": [0.35, 0.45, 0.55] + "mean_ttft_ms": [1000, 3000, 5000, 7000, 9000], + "mean_audio_ttfp_ms": [30000, 60000, 90000, 120000, 150000], + "mean_audio_rtf": [0.35, 0.45, 0.55, 0.65, 0.75] } }, { @@ -52,7 +52,7 @@ "backend": "openai-chat-omni", "endpoint": "/v1/chat/completions", "num_prompts": [40], - "request_rate": [0.3], + "request_rate": [0.5], "random_input_len": 100, "random_output_len": 100, "random_range_ratio": 0.0, @@ -69,9 +69,9 @@ }, "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "mean_ttft_ms": [4000], - "mean_audio_ttfp_ms": [13000], - "mean_audio_rtf": [0.35] + "mean_ttft_ms": [6000], + "mean_audio_ttfp_ms": [15000], + "mean_audio_rtf": [0.45] } }, { @@ -79,7 +79,7 @@ "backend": "openai-chat-omni", "endpoint": "/v1/chat/completions", "num_prompts": [100], - "request_rate": [0.5], + "request_rate": [1.0], "random_input_len": 100, "random_output_len": 100, "random_range_ratio": 0.0, @@ -98,9 +98,9 @@ }, "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "mean_ttft_ms": [6000], - "mean_audio_ttfp_ms": [15000], - "mean_audio_rtf": [0.45] + "mean_ttft_ms": [12000], + "mean_audio_ttfp_ms": [18000], + "mean_audio_rtf": [0.9] } } ] @@ -116,16 +116,16 @@ "dataset_name": "random", "backend": "openai-chat-omni", "endpoint": "/v1/chat/completions", - "num_prompts": [4, 16, 40], - "max_concurrency": [1, 4, 10], + "num_prompts": [4, 16, 32, 64], + "max_concurrency": [1, 4, 8, 16], "random_input_len": 2500, "random_output_len": 900, "ignore_eos": true, "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "mean_ttft_ms": [1000, 3000, 5000], - "mean_audio_ttfp_ms": [1000, 3000, 5000], - "mean_audio_rtf": [0.2, 0.35, 0.6] + "mean_ttft_ms": [1000, 3000, 5000, 7000], + "mean_audio_ttfp_ms": [1000, 3000, 5000, 7000], + "mean_audio_rtf": [0.2, 0.35, 0.6, 0.85] } }, { @@ -158,7 +158,7 @@ "backend": "openai-chat-omni", "endpoint": "/v1/chat/completions", "num_prompts": [40], - "request_rate": [0.3], + "request_rate": [0.5], "random_input_len": 100, "random_output_len": 100, "random_range_ratio": 0.0, @@ -175,9 +175,9 @@ }, "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "mean_ttft_ms": [4000], - "mean_audio_ttfp_ms": [4000], - "mean_audio_rtf": [0.4] + "mean_ttft_ms": [6000], + "mean_audio_ttfp_ms": [6000], + "mean_audio_rtf": [0.7] } }, { @@ -185,7 +185,7 @@ "backend": "openai-chat-omni", "endpoint": "/v1/chat/completions", "num_prompts": [100], - "request_rate": [0.5], + "request_rate": [1.0], "random_input_len": 100, "random_output_len": 100, "random_range_ratio": 0.0, @@ -204,17 +204,17 @@ }, "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration", "baseline": { - "mean_ttft_ms": [6000], - "mean_audio_ttfp_ms": [6000], - "mean_audio_rtf": [0.7] + "mean_ttft_ms": [12000], + "mean_audio_ttfp_ms": [12000], + "mean_audio_rtf": [1.0] } }, { "dataset_name": "random", "backend": "openai-chat-omni", "endpoint": "/v1/chat/completions", - "num_prompts": [4, 16, 40], - "max_concurrency": [1, 4, 10], + "num_prompts": [4, 16, 32, 64, 128], + "max_concurrency": [1, 4, 8, 16, 32], "random_input_len": 2500, "random_output_len": 900, "ignore_eos": true, @@ -223,7 +223,7 @@ }, "percentile-metrics": "ttft,tpot,itl,e2el", "baseline": { - "mean_ttft_ms": [1000, 3000, 5000] + "mean_ttft_ms": [1000, 3000, 5000, 7000, 9000] } }, { @@ -257,7 +257,7 @@ "backend": "openai-chat-omni", "endpoint": "/v1/chat/completions", "num_prompts": [40], - "request_rate": [0.3], + "request_rate": [0.5], "random_input_len": 100, "random_output_len": 100, "random_range_ratio": 0.0, @@ -277,7 +277,7 @@ }, "percentile-metrics": "ttft,tpot,itl,e2el", "baseline": { - "mean_ttft_ms": [4000] + "mean_ttft_ms": [6000] } }, { @@ -285,7 +285,7 @@ "backend": "openai-chat-omni", "endpoint": "/v1/chat/completions", "num_prompts": [100], - "request_rate": [0.5], + "request_rate": [1.0], "random_input_len": 100, "random_output_len": 100, "random_range_ratio": 0.0, diff --git a/tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json b/tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json new file mode 100644 index 000000000..58a17c980 --- /dev/null +++ b/tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json @@ -0,0 +1,107 @@ +[ + { + "test_name": "test_wan22_i2v_single_device", + "description": "Single-device baseline", + "server_type": "vllm-omni", + "server_params": { + "model": "Wan-AI/Wan2.2-I2V-A14B-Diffusers", + "serve_args": { + "enable-diffusion-pipeline-profiler": true + } + }, + "benchmark_params": [ + { + "name": "832x480_frames81_steps4", + "dataset": "random", + "task": "i2v", + "num-prompts": 10, + "max-concurrency": 1, + "num-input-images": 1, + "seed": 42, + "enable-negative-prompt": true, + "random-request-config": [ + { + "width": 832, + "height": 480, + "num_inference_steps": 4, + "num_frames": 81, + "fps": 16, + "weight": 1 + } + ], + "baseline": { + "throughput_qps": 0.034, + "latency_mean": 26.0, + "peak_memory_mb_mean": 80000 + } + } + ] + }, + { + "test_name": "test_wan22_i2v_usp2_vae_patch2_hsdp_slicing", + "description": "USP=2 + VAE patch parallel=2 + HSDP + VAE slicing", + "server_type": "vllm-omni", + "server_params": { + "model": "Wan-AI/Wan2.2-I2V-A14B-Diffusers", + "serve_args": { + "usp": 2, + "vae-patch-parallel-size": 2, + "use-hsdp": true, + "vae-use-slicing": true, + "enable-diffusion-pipeline-profiler": true + } + }, + "benchmark_params": [ + { + "name": "832x480_frames81_steps4", + "dataset": "random", + "task": "i2v", + "num-prompts": 10, + "max-concurrency": 1, + "num-input-images": 1, + "seed": 42, + "enable-negative-prompt": true, + "random-request-config": [ + { + "width": 832, + "height": 480, + "num_inference_steps": 4, + "num_frames": 81, + "fps": 16, + "weight": 1 + } + ], + "baseline": { + "throughput_qps": 0.042, + "latency_mean": 21.6, + "peak_memory_mb_mean": 55300 + } + }, + { + "name": "1280x720_frames121_steps4", + "dataset": "random", + "task": "i2v", + "num-prompts": 10, + "max-concurrency": 1, + "num-input-images": 1, + "seed": 42, + "enable-negative-prompt": true, + "random-request-config": [ + { + "width": 1280, + "height": 720, + "num_inference_steps": 4, + "num_frames": 121, + "fps": 16, + "weight": 1 + } + ], + "baseline": { + "throughput_qps": 0.0085, + "latency_mean": 101.6, + "peak_memory_mb_mean": 65200 + } + } + ] + } +] diff --git a/tests/dfx/stability/conftest.py b/tests/dfx/stability/conftest.py index e36c88b9a..30718d4bf 100644 --- a/tests/dfx/stability/conftest.py +++ b/tests/dfx/stability/conftest.py @@ -1,4 +1,14 @@ -"""Stability pytest hooks and fixtures.""" +""" +Stability-specific conftest: when pytest is executed under this directory, +resource monitoring is started before each test and finalized after each test, +so each stability test case gets its own HTML report (one report per case). +No need to wrap pytest with `bash resource_monitor.sh run -- pytest ...`. + +Duration-based benchmark helper functions are hosted in ``helpers.py``, +while this file focuses on pytest fixtures and setup/teardown. +""" + +from __future__ import annotations import subprocess import sys @@ -6,12 +16,66 @@ import threading import pytest +from tests.dfx.conftest import get_benchmark_params_for_server from tests.dfx.stability.helpers import ( finalize_resource_monitor, report_latest_gpu_samples, start_resource_monitor, wait_for_run_dir, ) +from tests.helpers.runtime import OmniServer + +DEFAULT_STABILITY_SERVER_TIMEOUT_ARGS = ["--stage-init-timeout", "600", "--init-timeout", "900"] + +_omni_server_lock = threading.Lock() + + +@pytest.fixture(scope="module") +def omni_server(request: pytest.FixtureRequest): + """Start OmniServer for stability tests, with per-module timeout override.""" + timeout_args = getattr(request.module, "STABILITY_SERVER_TIMEOUT_ARGS", DEFAULT_STABILITY_SERVER_TIMEOUT_ARGS) + with _omni_server_lock: + # Same 5-tuple and CLI composition as ``tests/dfx/perf/scripts/run_benchmark.py`` on main; + # ``serve_args`` from JSON are folded into ``extra_cli_args`` inside + # ``create_unique_server_params``. + test_name, model, deploy_path, stage_overrides, extra_cli_args = request.param + + print(f"Starting OmniServer with test: {test_name}, model: {model}") + server_args = list(timeout_args) + if deploy_path: + server_args = ["--deploy-config", deploy_path] + server_args + if stage_overrides: + server_args = ["--stage-overrides", stage_overrides] + server_args + if extra_cli_args: + server_args = list(extra_cli_args) + server_args + with OmniServer(model, server_args) as server: + server.test_name = test_name + print("OmniServer started successfully") + yield server + print("OmniServer stopping...") + print("OmniServer stopped") + + +@pytest.fixture +def stability_benchmark_params(request: pytest.FixtureRequest, omni_server): + test_name, param_index = request.param + if test_name != omni_server.test_name: + pytest.skip(f"Skipping parameter for {test_name} - current server is {omni_server.test_name}") + + server_to_benchmark_mapping = getattr(request.module, "server_to_benchmark_mapping", None) + if server_to_benchmark_mapping is None: + raise ValueError("server_to_benchmark_mapping must be defined in the test module") + + all_params = get_benchmark_params_for_server(test_name, server_to_benchmark_mapping) + if not all_params: + raise ValueError(f"No benchmark parameters found for test: {test_name}") + if param_index >= len(all_params): + raise ValueError(f"No benchmark parameters found for index {param_index} in test: {test_name}") + + current = param_index + 1 + total = len(all_params) + print(f"\n Running benchmark {current}/{total} for {test_name}") + return {"test_name": test_name, "params": all_params[param_index]} @pytest.fixture(autouse=True) diff --git a/tests/dfx/stability/helpers.py b/tests/dfx/stability/helpers.py index 3956bd213..3a873f69c 100644 --- a/tests/dfx/stability/helpers.py +++ b/tests/dfx/stability/helpers.py @@ -1,17 +1,32 @@ -"""Stability resource monitor helpers.""" +"""Stability helpers for resource monitoring and benchmark execution.""" from __future__ import annotations +import json import os +import random +import re +import shlex import subprocess import sys +import tempfile import threading import time +from collections.abc import Callable from pathlib import Path +from typing import Any + +from tests.dfx.conftest import run_benchmark STABILITY_DIR = Path(__file__).resolve().parent RESOURCE_MONITOR_SCRIPT = STABILITY_DIR / "scripts" / "resource_monitor.sh" REPO_ROOT = STABILITY_DIR.parent.parent.parent +_BUCKET_KEY_PATTERN = re.compile(r"^\(\s*([^,]+)\s*,\s*([^,]+)\s*,\s*([^,]+)\s*\)$") + +RunOneBatchFn = Callable[ + [str, int, str, dict[str, Any], int, float | None, int | None, str, int], + dict[str, Any], +] def start_resource_monitor(): @@ -115,3 +130,375 @@ def finalize_resource_monitor() -> str | None: return None except (FileNotFoundError, OSError, subprocess.TimeoutExpired): return None + + +def _normalize_bench_metrics(raw: dict[str, Any]) -> dict[str, Any]: + completed = int(raw.get("completed", raw.get("completed_requests", 0) or 0)) + failed = int(raw.get("failed", raw.get("failed_requests", 0) or 0)) + duration = float(raw.get("duration", 0.0) or 0.0) + errors = list(raw.get("errors") or []) + if failed and not errors: + errors = [f"{failed} benchmark request(s) failed"] + return {"completed": completed, "failed": failed, "duration": duration, "errors": errors} + + +def _build_base_args(params: dict[str, Any], host: str, port: int) -> list[str]: + exclude = { + "request_rate", + "max_concurrency", + "num_prompts", + "baseline", + "duration_sec", + "num_prompts_per_batch", + } + args = ["--host", host, "--port", str(port)] + for key, value in params.items(): + if key in exclude or value is None: + continue + arg_name = f"--{key.replace('_', '-')}" + if isinstance(value, bool) and value: + args.append(arg_name) + elif isinstance(value, dict): + args.extend([arg_name, json.dumps(value, ensure_ascii=False, separators=(",", ":"))]) + elif not isinstance(value, bool): + args.extend([arg_name, str(value)]) + return args + + +def _build_diffusion_cmd( + host: str, + port: int, + model: str, + params: dict[str, Any], + num_prompts: int, + request_rate: float | None, + max_concurrency: int | None, + output_path: Path, + diffusion_benchmark_script: Path, +) -> list[str]: + skip_keys = { + "request_rate", + "max_concurrency", + "num_prompts", + "baseline", + "duration_sec", + "num_prompts_per_batch", + } + cmd: list[str] = [ + sys.executable, + "-u", + str(diffusion_benchmark_script), + "--host", + host, + "--port", + str(port), + "--model", + model, + "--output-file", + str(output_path), + ] + for key, value in params.items(): + if key in skip_keys or value is None: + continue + flag = f"--{str(key).replace('_', '-')}" + if isinstance(value, bool) and value: + cmd.append(flag) + elif isinstance(value, bool): + continue + elif isinstance(value, (dict, list)): + cmd.extend([flag, json.dumps(value, ensure_ascii=False, separators=(",", ":"))]) + else: + cmd.extend([flag, str(value)]) + + cmd.extend(["--num-prompts", str(num_prompts)]) + if request_rate is not None: + cmd.extend(["--request-rate", str(request_rate)]) + else: + cmd.extend(["--max-concurrency", str(max_concurrency), "--request-rate", "inf"]) + return cmd + + +def _sample_int_from_range_spec(value: Any, rng: random.Random) -> Any: + """Resolve one value that may be scalar or range spec into an int.""" + if isinstance(value, int): + return value + + if isinstance(value, (list, tuple)) and len(value) == 2 and all(isinstance(v, int) for v in value): + low, high = int(value[0]), int(value[1]) + if low > high: + low, high = high, low + return rng.randint(low, high) + + if isinstance(value, dict) and {"min", "max"} <= set(value): + low, high = int(value["min"]), int(value["max"]) + if low > high: + low, high = high, low + return rng.randint(low, high) + + if isinstance(value, str): + raw = value.strip() + if raw.isdigit(): + return int(raw) + if "-" in raw: + parts = [p.strip() for p in raw.split("-", 1)] + if len(parts) == 2 and parts[0].isdigit() and parts[1].isdigit(): + low, high = int(parts[0]), int(parts[1]) + if low > high: + low, high = high, low + return rng.randint(low, high) + + return value + + +def _sample_bucket_key(raw_key: str, rng: random.Random) -> str: + """Sample bucket tuple keys that use range syntax, e.g. ``(128-512, 128-512, 1)``.""" + match = _BUCKET_KEY_PATTERN.match(raw_key.strip()) + if not match: + return raw_key + + sampled_parts: list[int] = [] + for token in match.groups(): + sampled = _sample_int_from_range_spec(token.strip(), rng) + if not isinstance(sampled, int): + return raw_key + sampled_parts.append(sampled) + + # For video buckets (height>0 and num_frames>1), enforce even H/W to avoid + # ffmpeg yuv420p encoding/decoding failures ("Could not open video stream"). + if sampled_parts[0] > 0 and sampled_parts[2] > 1: + sampled_parts[0] = max(2, sampled_parts[0] - (sampled_parts[0] % 2)) + sampled_parts[1] = max(2, sampled_parts[1] - (sampled_parts[1] % 2)) + + return f"({sampled_parts[0]}, {sampled_parts[1]}, {sampled_parts[2]})" + + +def _sample_stability_batch_params(params: dict[str, Any], batch_index: int) -> dict[str, Any]: + """Materialize per-batch random values for configured range fields.""" + sampled = dict(params) + rng = random.Random(time.time_ns() + batch_index) + + for field_name in ( + "random_input_len", + "random_output_len", + "random_mm_base_items_per_request", + "width", + "height", + ): + if field_name in sampled: + sampled[field_name] = _sample_int_from_range_spec(sampled[field_name], rng) + + bucket_config = sampled.get("random_mm_bucket_config") + if isinstance(bucket_config, dict): + sampled_bucket_config: dict[str, float] = {} + for raw_key, probability in bucket_config.items(): + sampled_key = _sample_bucket_key(str(raw_key), rng) + sampled_bucket_config[sampled_key] = sampled_bucket_config.get(sampled_key, 0.0) + float(probability) + sampled["random_mm_bucket_config"] = sampled_bucket_config + + return sampled + + +def _run_one_vllm_bench_batch( + host: str, + port: int, + _model: str, + params: dict[str, Any], + num_prompts: int, + request_rate: float | None, + max_concurrency: int | None, + result_dir: str, + batch_index: int, +) -> dict[str, Any]: + base = _build_base_args(params, host, port) + if request_rate is not None: + args = base + ["--request-rate", str(request_rate), "--num-prompts", str(num_prompts)] + flow = request_rate + else: + args = base + [ + "--max-concurrency", + str(max_concurrency), + "--num-prompts", + str(num_prompts), + "--request-rate", + "inf", + ] + flow = max_concurrency + + # Print the exact per-batch benchmark CLI (randomized params are already materialized). + preview_cmd = ["vllm", "bench", "serve", "--omni", *args] + print(f"\n[Stability][Batch {batch_index}] Benchmark command:") + print(shlex.join(preview_cmd)) + + dataset_name = params.get("dataset_name", "random") + old_benchmark_dir = os.environ.get("BENCHMARK_DIR") + try: + os.environ["BENCHMARK_DIR"] = result_dir + result = run_benchmark( + args=args, + test_name="stability", + flow=flow, + dataset_name=dataset_name, + num_prompt=num_prompts, + random_input_len=params.get("random_input_len"), + random_output_len=params.get("random_output_len"), + ) + return _normalize_bench_metrics(result) + except (FileNotFoundError, OSError) as exc: + return { + "completed": 0, + "failed": 1, + "duration": 0.0, + "errors": [f"Benchmark batch failed: {type(exc).__name__}: {exc}"], + } + finally: + if old_benchmark_dir is not None: + os.environ["BENCHMARK_DIR"] = old_benchmark_dir + elif "BENCHMARK_DIR" in os.environ: + os.environ.pop("BENCHMARK_DIR") + + +def _run_one_diffusion_batch( + host: str, + port: int, + model: str, + params: dict[str, Any], + num_prompts: int, + request_rate: float | None, + max_concurrency: int | None, + _result_dir: str, + _batch_index: int, +) -> dict[str, Any]: + diffusion_benchmark_script = Path(REPO_ROOT / "benchmarks" / "diffusion" / "diffusion_benchmark_serving.py") + with tempfile.NamedTemporaryFile(mode="w", suffix=".json", prefix="stability_diffusion_", delete=False) as tmp: + out_path = Path(tmp.name) + try: + cmd = _build_diffusion_cmd( + host, + port, + model, + params, + num_prompts, + request_rate, + max_concurrency, + out_path, + diffusion_benchmark_script, + ) + proc = subprocess.run( + cmd, + cwd=str(REPO_ROOT), + capture_output=True, + text=True, + ) + if proc.stdout: + print(proc.stdout, end="" if proc.stdout.endswith("\n") else "\n") + if proc.stderr: + print(proc.stderr, end="" if proc.stderr.endswith("\n") else "\n") + if proc.returncode != 0: + return { + "completed": 0, + "failed": 1, + "duration": 0.0, + "errors": [f"diffusion_benchmark_serving.py exited {proc.returncode}"], + } + if not out_path.is_file(): + return { + "completed": 0, + "failed": 1, + "duration": 0.0, + "errors": [f"Missing benchmark output: {out_path}"], + } + with open(out_path, encoding="utf-8") as file: + metrics = json.load(file) + return _normalize_bench_metrics(metrics) + except (FileNotFoundError, OSError, json.JSONDecodeError) as exc: + return { + "completed": 0, + "failed": 1, + "duration": 0.0, + "errors": [f"Diffusion batch failed: {type(exc).__name__}: {exc}"], + } + finally: + out_path.unlink(missing_ok=True) + + +def merge_batch_results(batch_results: list[dict[str, Any]], total_duration_sec: float) -> dict[str, Any]: + if not batch_results: + return {"completed": 0, "failed": 0, "duration": total_duration_sec, "errors": []} + + completed = sum(result.get("completed", 0) for result in batch_results) + failed = sum(result.get("failed", 0) for result in batch_results) + merged: dict[str, Any] = { + "completed": completed, + "failed": failed, + "duration": total_duration_sec, + "errors": [], + } + for result in batch_results: + merged["errors"].extend(result.get("errors") or []) + return merged + + +def print_merged_report(result: dict[str, Any]) -> None: + fmt = "{:<40} {:<10}" + fmt_float = "{:<40} {:<10.2f}" + completed = result.get("completed", 0) + failed = result.get("failed", 0) + duration = float(result.get("duration", 0.0) or 0.0) + print("\n============ Stability Benchmark Summary ============") + print(fmt.format("Successful requests:", completed)) + print(fmt.format("Failed requests:", failed)) + print(fmt_float.format("Total duration (s):", duration)) + print("==================================================\n") + + +def run_stability_benchmark_loop( + host: str, + port: int, + model: str, + duration_sec: int | float, + params: dict[str, Any], + *, + request_rate: float | None, + max_concurrency: int | None, + result_dir: str, + num_prompts_per_batch: int, + run_one_batch: RunOneBatchFn, + result_filename: str | None = None, +) -> dict[str, Any]: + if (request_rate is None) == (max_concurrency is None): + raise ValueError("Exactly one of request_rate or max_concurrency must be specified") + + start_time = time.perf_counter() + batch_results: list[dict[str, Any]] = [] + batch_index = 0 + + while True: + if (time.perf_counter() - start_time) >= duration_sec: + break + sampled_params = _sample_stability_batch_params(params, batch_index) + result = run_one_batch( + host, + port, + model, + sampled_params, + num_prompts_per_batch, + request_rate, + max_concurrency, + result_dir, + batch_index, + ) + batch_results.append(result) + batch_index += 1 + if (time.perf_counter() - start_time) >= duration_sec: + break + + total_duration = time.perf_counter() - start_time + merged = merge_batch_results(batch_results, total_duration) + print_merged_report(merged) + + if result_filename and result_dir: + result_path = Path(result_dir) / result_filename + with open(result_path, "w", encoding="utf-8") as file: + json.dump(merged, file, indent=2, ensure_ascii=False) + + return merged diff --git a/tests/dfx/stability/scripts/test_benchmark_stability.py b/tests/dfx/stability/scripts/test_benchmark_stability.py deleted file mode 100644 index 620241762..000000000 --- a/tests/dfx/stability/scripts/test_benchmark_stability.py +++ /dev/null @@ -1,297 +0,0 @@ -""" -Stability test cases: start OmniServer first, then run benchmark traffic with either -`request-rate` or `max-concurrency` for a fixed duration. No new requests are sent -after the duration is reached, and the test asserts that there are no failed requests. - -The overall flow matches the perf logic: `load_configs`, `modify_stage`, -`create_unique_server_params`, `create_test_parameter_mapping`, -`get_benchmark_params_for_server`, `create_benchmark_indices`, and the -`omni_server` fixture are aligned with perf. Only the benchmark execution -(`run_stability_benchmark`, which is duration-based here) and the test cases differ. - -All test-specific parameters, such as `duration_sec`, `request_rate` / -`max_concurrency`, and `num_prompts_per_batch`, are configured in -`tests/dfx/stability/tests/test.json` and are no longer overridden -through environment variables. -""" - -import json -import os -import threading -import time -from pathlib import Path -from typing import Any - -import pytest - -from tests.dfx.conftest import ( - create_benchmark_indices, - create_test_parameter_mapping, - create_unique_server_params, - get_benchmark_params_for_server, - load_configs, -) -from tests.dfx.perf.scripts.run_benchmark import run_benchmark -from tests.helpers.runtime import OmniServer - -STABILITY_DIR = Path(__file__).resolve().parent.parent -DEPLOY_CONFIGS_DIR = STABILITY_DIR / "deploy" -CONFIG_FILE_PATH = str(STABILITY_DIR / "tests" / "test.json") -DEFAULT_NUM_PROMPTS_PER_BATCH = 20 - - -try: - BENCHMARK_CONFIGS = load_configs(CONFIG_FILE_PATH) -except FileNotFoundError: - BENCHMARK_CONFIGS = [] - -test_params = create_unique_server_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) if BENCHMARK_CONFIGS else [] -server_to_benchmark_mapping = create_test_parameter_mapping(BENCHMARK_CONFIGS) if BENCHMARK_CONFIGS else {} - -_omni_server_lock = threading.Lock() - - -benchmark_indices = create_benchmark_indices(BENCHMARK_CONFIGS, server_to_benchmark_mapping) - - -def _build_base_args(params: dict[str, Any], host: str, port: int) -> list[str]: - exclude = { - "request_rate", - "max_concurrency", - "num_prompts", - "baseline", - "duration_sec", - "num_prompts_per_batch", - } - args = ["--host", host, "--port", str(port)] - for key, value in params.items(): - if key in exclude or value is None: - continue - arg_name = f"--{key.replace('_', '-')}" - if isinstance(value, bool) and value: - args.append(arg_name) - elif isinstance(value, dict): - args.extend([arg_name, json.dumps(value, ensure_ascii=False, separators=(",", ":"))]) - elif not isinstance(value, bool): - args.extend([arg_name, str(value)]) - return args - - -def _run_one_benchmark_batch( - host: str, - port: int, - params: dict[str, Any], - num_prompts: int, - request_rate: float | None, - max_concurrency: int | None, - result_dir: str, - batch_index: int, -) -> dict[str, Any]: - base = _build_base_args(params, host, port) - if request_rate is not None: - args = base + ["--request-rate", str(request_rate), "--num-prompts", str(num_prompts)] - flow = request_rate - else: - args = base + [ - "--max-concurrency", - str(max_concurrency), - "--num-prompts", - str(num_prompts), - "--request-rate", - "inf", - ] - flow = max_concurrency - - dataset_name = params.get("dataset_name", "random") - old_benchmark_dir = os.environ.get("BENCHMARK_DIR") - try: - os.environ["BENCHMARK_DIR"] = result_dir - result = run_benchmark( - args=args, - test_name="stability", - flow=flow, - dataset_name=dataset_name, - num_prompt=num_prompts, - random_input_len=params.get("random_input_len"), - random_output_len=params.get("random_output_len"), - ) - return result - except (FileNotFoundError, OSError) as e: - # Surface batch failure so the stability test does not false-pass when - # run_benchmark fails before writing JSON (e.g. command not found). - return { - "completed": 0, - "failed": 1, - "duration": 0.0, - "errors": [f"Benchmark batch failed: {type(e).__name__}: {e}"], - } - finally: - if old_benchmark_dir is not None: - os.environ["BENCHMARK_DIR"] = old_benchmark_dir - elif "BENCHMARK_DIR" in os.environ: - os.environ.pop("BENCHMARK_DIR") - - -def _merge_batch_results(batch_results: list[dict[str, Any]], total_duration_sec: float) -> dict[str, Any]: - if not batch_results: - return {"completed": 0, "failed": 0, "duration": total_duration_sec, "errors": []} - - completed = sum(r.get("completed", 0) for r in batch_results) - failed = sum(r.get("failed", 0) for r in batch_results) - merged: dict[str, Any] = { - "completed": completed, - "failed": failed, - "duration": total_duration_sec, - "errors": [], - } - for r in batch_results: - merged["errors"].extend(r.get("errors") or []) - return merged - - -def _print_merged_report(result: dict[str, Any]) -> None: - """Print the final summary: successful requests, failed requests, and total duration only.""" - fmt = "{:<40} {:<10}" - fmt_float = "{:<40} {:<10.2f}" - completed = result.get("completed", 0) - failed = result.get("failed", 0) - duration = float(result.get("duration", 0.0) or 0.0) - print("\n============ Stability Benchmark Summary ============") - print(fmt.format("Successful requests:", completed)) - print(fmt.format("Failed requests:", failed)) - print(fmt_float.format("Total duration (s):", duration)) - print("==================================================\n") - - -def run_stability_benchmark( - host: str, - port: int, - duration_sec: int | float, - params: dict[str, Any], - *, - request_rate: float | None = None, - max_concurrency: int | None = None, - result_filename: str | None = None, - result_dir: str = "./", - num_prompts_per_batch: int = DEFAULT_NUM_PROMPTS_PER_BATCH, -) -> dict[str, Any]: - if (request_rate is None) == (max_concurrency is None): - raise ValueError("Exactly one of request_rate or max_concurrency must be specified") - - start_time = time.perf_counter() - batch_results: list[dict[str, Any]] = [] - batch_index = 0 - - while True: - if (time.perf_counter() - start_time) >= duration_sec: - break - result = _run_one_benchmark_batch( - host=host, - port=port, - params=params, - num_prompts=num_prompts_per_batch, - request_rate=request_rate, - max_concurrency=max_concurrency, - result_dir=result_dir, - batch_index=batch_index, - ) - batch_results.append(result) - batch_index += 1 - if (time.perf_counter() - start_time) >= duration_sec: - break - - total_duration = time.perf_counter() - start_time - merged = _merge_batch_results(batch_results, total_duration) - _print_merged_report(merged) - - if result_filename and result_dir: - result_path = Path(result_dir) / result_filename - with open(result_path, "w", encoding="utf-8") as f: - json.dump(merged, f, indent=2, ensure_ascii=False) - - return merged - - -@pytest.fixture(scope="module") -def omni_server(request): - """Start vLLM-Omni server as a subprocess with actual model weights. - Uses session scope so the server starts only once for the entire test session. - Multi-stage initialization can take 10-20+ minutes. - """ - with _omni_server_lock: - test_name, model, stage_config_path, stage_overrides, extra_cli_args = request.param - - print(f"Starting OmniServer with test: {test_name}, model: {model}") - - server_args = ["--stage-init-timeout", "120"] - # --deploy-config and --stage-overrides compose at the CLI (see vllm_omni/entrypoints/utils.py): - # deploy-config sets the base; stage-overrides are applied on top. Both can be set. - if stage_config_path: - server_args = ["--deploy-config", stage_config_path] + server_args - if stage_overrides: - server_args = ["--stage-overrides", stage_overrides] + server_args - if extra_cli_args: - server_args = list(extra_cli_args) + server_args - with OmniServer(model, server_args) as server: - server.test_name = test_name - print("OmniServer started successfully") - yield server - print("OmniServer stopping...") - - print("OmniServer stopped") - - -@pytest.fixture(params=benchmark_indices) -def stability_benchmark_params(request, omni_server): - """Benchmark parameters fixture with proper parametrization (same as perf).""" - test_name, param_index = request.param - - if test_name != omni_server.test_name: - pytest.skip(f"Skipping parameter for {test_name} - current server is {omni_server.test_name}") - - all_params = get_benchmark_params_for_server(test_name, server_to_benchmark_mapping) - - if not all_params: - raise ValueError(f"No benchmark parameters found for test: {test_name}") - - if param_index >= len(all_params): - raise ValueError(f"No benchmark parameters found for index {param_index} in test: {test_name}") - - current = param_index + 1 - total = len(all_params) - print(f"\n Running benchmark {current}/{total} for {test_name}") - - return {"test_name": test_name, "params": all_params[param_index]} - - -@pytest.mark.slow -@pytest.mark.parametrize("omni_server", test_params, indirect=True) -@pytest.mark.parametrize("stability_benchmark_params", benchmark_indices, indirect=True) -def test_benchmark_stability(omni_server, stability_benchmark_params): - """Run the benchmark for a fixed duration using request-rate or max-concurrency and assert zero failed requests.""" - test_name = stability_benchmark_params["test_name"] - params = stability_benchmark_params["params"] - duration_sec = params.get("duration_sec", 300) - num_prompts_per_batch = params.get("num_prompts_per_batch", DEFAULT_NUM_PROMPTS_PER_BATCH) - request_rate = params.get("request_rate") - max_concurrency = params.get("max_concurrency") - - bench_params = { - k: v - for k, v in params.items() - if k not in ("duration_sec", "request_rate", "max_concurrency", "num_prompts_per_batch") - } - - result = run_stability_benchmark( - host=omni_server.host, - port=omni_server.port, - duration_sec=duration_sec, - params=bench_params, - request_rate=request_rate, - max_concurrency=max_concurrency, - result_dir=str(STABILITY_DIR), - num_prompts_per_batch=num_prompts_per_batch, - ) - - assert result.get("failed", 0) == 0, f"[{test_name}] Failed requests detected: {result.get('errors', [])}" - assert result.get("completed", 0) > 0, f"[{test_name}] No requests completed" diff --git a/tests/dfx/stability/scripts/test_stability_qwen3_omni.py b/tests/dfx/stability/scripts/test_stability_qwen3_omni.py new file mode 100644 index 000000000..d1c2af8cf --- /dev/null +++ b/tests/dfx/stability/scripts/test_stability_qwen3_omni.py @@ -0,0 +1,68 @@ +""" +Qwen3-Omni stability: OmniServer + ``vllm bench serve --omni`` for a fixed duration. + +Configuration: ``tests/dfx/stability/tests/test_qwen3_omni.json``. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from tests.dfx.conftest import ( + create_benchmark_indices, + create_test_parameter_mapping, + create_unique_server_params, + load_configs, +) +from tests.dfx.stability.helpers import _run_one_vllm_bench_batch, run_stability_benchmark_loop + +STABILITY_DIR = Path(__file__).resolve().parent.parent +DEPLOY_CONFIGS_DIR = STABILITY_DIR / "deploy" +CONFIG_FILE_PATH = str(STABILITY_DIR / "tests" / "test_qwen3_omni.json") +DEFAULT_NUM_PROMPTS_PER_BATCH = 20 +STABILITY_SERVER_TIMEOUT_ARGS = ["--stage-init-timeout", "600"] + +try: + BENCHMARK_CONFIGS = load_configs(CONFIG_FILE_PATH) +except FileNotFoundError: + BENCHMARK_CONFIGS = [] + +test_params = create_unique_server_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) if BENCHMARK_CONFIGS else [] +server_to_benchmark_mapping = create_test_parameter_mapping(BENCHMARK_CONFIGS) if BENCHMARK_CONFIGS else {} +benchmark_indices = create_benchmark_indices(BENCHMARK_CONFIGS, server_to_benchmark_mapping) + + +@pytest.mark.slow +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +@pytest.mark.parametrize("stability_benchmark_params", benchmark_indices, indirect=True) +def test_stability_qwen3_omni(omni_server, stability_benchmark_params): + test_name = stability_benchmark_params["test_name"] + params = stability_benchmark_params["params"] + duration_sec = params.get("duration_sec", 300) + num_prompts_per_batch = params.get("num_prompts_per_batch", DEFAULT_NUM_PROMPTS_PER_BATCH) + request_rate = params.get("request_rate") + max_concurrency = params.get("max_concurrency") + + bench_params = { + k: v + for k, v in params.items() + if k not in ("duration_sec", "request_rate", "max_concurrency", "num_prompts_per_batch") + } + + result = run_stability_benchmark_loop( + host=omni_server.host, + port=omni_server.port, + model=omni_server.model, + duration_sec=duration_sec, + params=bench_params, + request_rate=request_rate, + max_concurrency=max_concurrency, + result_dir=str(STABILITY_DIR), + num_prompts_per_batch=num_prompts_per_batch, + run_one_batch=_run_one_vllm_bench_batch, + ) + + assert result.get("failed", 0) == 0, f"[{test_name}] Failed requests detected: {result.get('errors', [])}" + assert result.get("completed", 0) > 0, f"[{test_name}] No requests completed" diff --git a/tests/dfx/stability/scripts/test_stability_qwen3_tts.py b/tests/dfx/stability/scripts/test_stability_qwen3_tts.py new file mode 100644 index 000000000..beccd67d9 --- /dev/null +++ b/tests/dfx/stability/scripts/test_stability_qwen3_tts.py @@ -0,0 +1,68 @@ +""" +Qwen3-TTS stability: OmniServer + ``vllm bench serve --omni`` for a fixed duration. + +Configuration: ``tests/dfx/stability/tests/test_qwen3_tts.json``. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from tests.dfx.conftest import ( + create_benchmark_indices, + create_test_parameter_mapping, + create_unique_server_params, + load_configs, +) +from tests.dfx.stability.helpers import _run_one_vllm_bench_batch, run_stability_benchmark_loop + +STABILITY_DIR = Path(__file__).resolve().parent.parent +DEPLOY_CONFIGS_DIR = STABILITY_DIR / "deploy" +CONFIG_FILE_PATH = str(STABILITY_DIR / "tests" / "test_qwen3_tts.json") +DEFAULT_NUM_PROMPTS_PER_BATCH = 20 +STABILITY_SERVER_TIMEOUT_ARGS = ["--stage-init-timeout", "600"] + +try: + BENCHMARK_CONFIGS = load_configs(CONFIG_FILE_PATH) +except FileNotFoundError: + BENCHMARK_CONFIGS = [] + +test_params = create_unique_server_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) if BENCHMARK_CONFIGS else [] +server_to_benchmark_mapping = create_test_parameter_mapping(BENCHMARK_CONFIGS) if BENCHMARK_CONFIGS else {} +benchmark_indices = create_benchmark_indices(BENCHMARK_CONFIGS, server_to_benchmark_mapping) + + +@pytest.mark.slow +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +@pytest.mark.parametrize("stability_benchmark_params", benchmark_indices, indirect=True) +def test_stability_qwen3_tts(omni_server, stability_benchmark_params): + test_name = stability_benchmark_params["test_name"] + params = stability_benchmark_params["params"] + duration_sec = params.get("duration_sec", 300) + num_prompts_per_batch = params.get("num_prompts_per_batch", DEFAULT_NUM_PROMPTS_PER_BATCH) + request_rate = params.get("request_rate") + max_concurrency = params.get("max_concurrency") + + bench_params = { + k: v + for k, v in params.items() + if k not in ("duration_sec", "request_rate", "max_concurrency", "num_prompts_per_batch") + } + + result = run_stability_benchmark_loop( + host=omni_server.host, + port=omni_server.port, + model=omni_server.model, + duration_sec=duration_sec, + params=bench_params, + request_rate=request_rate, + max_concurrency=max_concurrency, + result_dir=str(STABILITY_DIR), + num_prompts_per_batch=num_prompts_per_batch, + run_one_batch=_run_one_vllm_bench_batch, + ) + + assert result.get("failed", 0) == 0, f"[{test_name}] Failed requests detected: {result.get('errors', [])}" + assert result.get("completed", 0) > 0, f"[{test_name}] No requests completed" diff --git a/tests/dfx/stability/scripts/test_stability_qwen_image.py b/tests/dfx/stability/scripts/test_stability_qwen_image.py new file mode 100644 index 000000000..a90e2092f --- /dev/null +++ b/tests/dfx/stability/scripts/test_stability_qwen_image.py @@ -0,0 +1,68 @@ +""" +Qwen-Image stability: OmniServer (diffusion) + ``diffusion_benchmark_serving.py``. + +Configuration: ``tests/dfx/stability/tests/test_qwen_image.json``. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from tests.dfx.conftest import ( + create_benchmark_indices, + create_test_parameter_mapping, + create_unique_server_params, + load_configs, +) +from tests.dfx.stability.helpers import _run_one_diffusion_batch, run_stability_benchmark_loop + +STABILITY_DIR = Path(__file__).resolve().parent.parent +DEPLOY_CONFIGS_DIR = STABILITY_DIR / "deploy" +CONFIG_FILE_PATH = str(STABILITY_DIR / "tests" / "test_qwen_image.json") +DEFAULT_NUM_PROMPTS_PER_BATCH = 20 +STABILITY_SERVER_TIMEOUT_ARGS = ["--stage-init-timeout", "600", "--init-timeout", "900"] + +try: + BENCHMARK_CONFIGS = load_configs(CONFIG_FILE_PATH) +except FileNotFoundError: + BENCHMARK_CONFIGS = [] + +test_params = create_unique_server_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) if BENCHMARK_CONFIGS else [] +server_to_benchmark_mapping = create_test_parameter_mapping(BENCHMARK_CONFIGS) if BENCHMARK_CONFIGS else {} +benchmark_indices = create_benchmark_indices(BENCHMARK_CONFIGS, server_to_benchmark_mapping) + + +@pytest.mark.slow +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +@pytest.mark.parametrize("stability_benchmark_params", benchmark_indices, indirect=True) +def test_stability_qwen_image(omni_server, stability_benchmark_params): + test_name = stability_benchmark_params["test_name"] + params = stability_benchmark_params["params"] + duration_sec = params.get("duration_sec", 300) + num_prompts_per_batch = params.get("num_prompts_per_batch", DEFAULT_NUM_PROMPTS_PER_BATCH) + request_rate = params.get("request_rate") + max_concurrency = params.get("max_concurrency") + + bench_params = { + k: v + for k, v in params.items() + if k not in ("duration_sec", "request_rate", "max_concurrency", "num_prompts_per_batch") + } + + result = run_stability_benchmark_loop( + host=omni_server.host, + port=omni_server.port, + model=omni_server.model, + duration_sec=duration_sec, + params=bench_params, + request_rate=request_rate, + max_concurrency=max_concurrency, + result_dir=str(STABILITY_DIR), + num_prompts_per_batch=num_prompts_per_batch, + run_one_batch=_run_one_diffusion_batch, + ) + + assert result.get("failed", 0) == 0, f"[{test_name}] Failed requests detected: {result.get('errors', [])}" + assert result.get("completed", 0) > 0, f"[{test_name}] No requests completed" diff --git a/tests/dfx/stability/scripts/test_stability_wan22.py b/tests/dfx/stability/scripts/test_stability_wan22.py new file mode 100644 index 000000000..afe9c4d0c --- /dev/null +++ b/tests/dfx/stability/scripts/test_stability_wan22.py @@ -0,0 +1,68 @@ +""" +Wan2.2 T2V stability: OmniServer (diffusion) + ``diffusion_benchmark_serving.py`` / ``v1/videos``. + +Configuration: ``tests/dfx/stability/tests/test_wan22.json``. +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from tests.dfx.conftest import ( + create_benchmark_indices, + create_test_parameter_mapping, + create_unique_server_params, + load_configs, +) +from tests.dfx.stability.helpers import _run_one_diffusion_batch, run_stability_benchmark_loop + +STABILITY_DIR = Path(__file__).resolve().parent.parent +DEPLOY_CONFIGS_DIR = STABILITY_DIR / "deploy" +CONFIG_FILE_PATH = str(STABILITY_DIR / "tests" / "test_wan22.json") +DEFAULT_NUM_PROMPTS_PER_BATCH = 20 +STABILITY_SERVER_TIMEOUT_ARGS = ["--stage-init-timeout", "600", "--init-timeout", "900"] + +try: + BENCHMARK_CONFIGS = load_configs(CONFIG_FILE_PATH) +except FileNotFoundError: + BENCHMARK_CONFIGS = [] + +test_params = create_unique_server_params(BENCHMARK_CONFIGS, DEPLOY_CONFIGS_DIR) if BENCHMARK_CONFIGS else [] +server_to_benchmark_mapping = create_test_parameter_mapping(BENCHMARK_CONFIGS) if BENCHMARK_CONFIGS else {} +benchmark_indices = create_benchmark_indices(BENCHMARK_CONFIGS, server_to_benchmark_mapping) + + +@pytest.mark.slow +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +@pytest.mark.parametrize("stability_benchmark_params", benchmark_indices, indirect=True) +def test_stability_wan22(omni_server, stability_benchmark_params): + test_name = stability_benchmark_params["test_name"] + params = stability_benchmark_params["params"] + duration_sec = params.get("duration_sec", 300) + num_prompts_per_batch = params.get("num_prompts_per_batch", DEFAULT_NUM_PROMPTS_PER_BATCH) + request_rate = params.get("request_rate") + max_concurrency = params.get("max_concurrency") + + bench_params = { + k: v + for k, v in params.items() + if k not in ("duration_sec", "request_rate", "max_concurrency", "num_prompts_per_batch") + } + + result = run_stability_benchmark_loop( + host=omni_server.host, + port=omni_server.port, + model=omni_server.model, + duration_sec=duration_sec, + params=bench_params, + request_rate=request_rate, + max_concurrency=max_concurrency, + result_dir=str(STABILITY_DIR), + num_prompts_per_batch=num_prompts_per_batch, + run_one_batch=_run_one_diffusion_batch, + ) + + assert result.get("failed", 0) == 0, f"[{test_name}] Failed requests detected: {result.get('errors', [])}" + assert result.get("completed", 0) > 0, f"[{test_name}] No requests completed" diff --git a/tests/dfx/stability/tests/test.json b/tests/dfx/stability/tests/test.json deleted file mode 100644 index 255cd5b10..000000000 --- a/tests/dfx/stability/tests/test.json +++ /dev/null @@ -1,77 +0,0 @@ -[ - { - "test_name": "test_qwen3_omni_stability", - "server_params": { - "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", - "stage_overrides": { - "2": { - "max_num_batched_tokens": 1000000 - } - } - }, - "benchmark_params": [ - { - "dataset_name": "random", - "backend": "openai-chat-omni", - "endpoint": "/v1/chat/completions", - "duration_sec": 300, - "request_rate": 0.2, - "num_prompts_per_batch": 10, - "random_input_len": 2500, - "random_output_len": 900, - "ignore_eos": true, - "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration" - }, - { - "dataset_name": "random", - "backend": "openai-chat-omni", - "endpoint": "/v1/chat/completions", - "duration_sec": 300, - "max_concurrency": 2, - "num_prompts_per_batch": 10, - "random_input_len": 2500, - "random_output_len": 900, - "ignore_eos": true, - "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration" - } - ] - }, - { - "test_name": "test_qwen3_omni_stability_async_chunk", - "server_params": { - "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", - "stage_overrides": { - "2": { - "max_num_batched_tokens": 1000000 - } - }, - "extra_cli_args": ["--async-chunk"] - }, - "benchmark_params": [ - { - "dataset_name": "random", - "backend": "openai-chat-omni", - "endpoint": "/v1/chat/completions", - "duration_sec": 300, - "request_rate": 0.2, - "num_prompts_per_batch": 10, - "random_input_len": 2500, - "random_output_len": 900, - "ignore_eos": true, - "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration" - }, - { - "dataset_name": "random", - "backend": "openai-chat-omni", - "endpoint": "/v1/chat/completions", - "duration_sec": 300, - "max_concurrency": 2, - "num_prompts_per_batch": 10, - "random_input_len": 2500, - "random_output_len": 900, - "ignore_eos": true, - "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration" - } - ] - } -] diff --git a/tests/dfx/stability/tests/test_qwen3_omni.json b/tests/dfx/stability/tests/test_qwen3_omni.json new file mode 100644 index 000000000..a16ab805c --- /dev/null +++ b/tests/dfx/stability/tests/test_qwen3_omni.json @@ -0,0 +1,97 @@ +[ + { + "test_name": "test_qwen3_omni_stability", + "server_params": { + "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", + "stage_overrides": { + "2": { + "max_num_batched_tokens": 1000000 + } + } + }, + "benchmark_params": [ + { + "dataset_name": "random-mm", + "backend": "openai-chat-omni", + "endpoint": "/v1/chat/completions", + "duration_sec": 86400, + "request_rate": 0.3, + "num_prompts_per_batch": 10, + "random_input_len": { + "min": 0, + "max": 8000 + }, + "random_output_len": { + "min": 0, + "max": 1000 + }, + "random_range_ratio": 0.0, + "random_mm_base_items_per_request": { + "min": 0, + "max": 6 + }, + "random_mm_num_mm_items_range_ratio": 0.0, + "random_mm_limit_mm_per_prompt": { + "image": 2, + "video": 2, + "audio": 2 + }, + "random_mm_bucket_config": { + "(128-1024, 128-1024, 1)": 0.34, + "(256-1080, 256-1920, 2-16)": 0.33, + "(0, 1-60, 1-3)": 0.33 + }, + "ignore_eos": true, + "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration" + } + ] + }, + { + "test_name": "test_qwen3_omni_stability_async_chunk", + "server_params": { + "model": "Qwen/Qwen3-Omni-30B-A3B-Instruct", + "stage_overrides": { + "2": { + "max_num_batched_tokens": 1000000 + } + }, + "extra_cli_args": ["--async-chunk"] + }, + "benchmark_params": [ + { + "dataset_name": "random-mm", + "backend": "openai-chat-omni", + "endpoint": "/v1/chat/completions", + "duration_sec": 86400, + "max_concurrency": 2, + "num_prompts_per_batch": 10, + "random_input_len": { + "min": 0, + "max": 8000 + }, + "random_output_len": { + "min": 0, + "max": 1000 + }, + "random_range_ratio": 0.0, + "random_mm_base_items_per_request": { + "min": 0, + "max": 6 + }, + "random_mm_num_mm_items_range_ratio": 0.0, + "random_mm_limit_mm_per_prompt": { + "image": 2, + "video": 2, + "audio": 2 + }, + "random_mm_bucket_config": { + "(128-1024, 128-1024, 1)": 0.34, + "(256-1080, 256-1920, 2-16)": 0.33, + "(0, 1-60, 1-3)": 0.33 + }, + "ignore_eos": true, + "percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration" + } + ] + } +] diff --git a/tests/dfx/stability/tests/test_qwen3_tts.json b/tests/dfx/stability/tests/test_qwen3_tts.json new file mode 100644 index 000000000..fbf30d88a --- /dev/null +++ b/tests/dfx/stability/tests/test_qwen3_tts.json @@ -0,0 +1,56 @@ +[ + { + "test_name": "test_qwen3_tts_stability", + "server_params": { + "model": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice" + }, + "benchmark_params": [ + { + "dataset_name": "random", + "backend": "openai-audio-speech", + "endpoint": "/v1/audio/speech", + "duration_sec": 86400, + "request_rate": 0.3, + "num_prompts_per_batch": 10, + "random_input_len": { + "min": 0, + "max": 1000 + }, + "random_output_len": { + "min": 0, + "max": 1000 + }, + "random_range_ratio": 0.0, + "extra_body": { + "voice": "Vivian", + "language": "English" + }, + "ignore_eos": true, + "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration" + }, + { + "dataset_name": "random", + "backend": "openai-audio-speech", + "endpoint": "/v1/audio/speech", + "duration_sec": 86400, + "max_concurrency": 2, + "num_prompts_per_batch": 10, + "random_input_len": { + "min": 0, + "max": 1000 + }, + "random_output_len": { + "min": 0, + "max": 1000 + }, + "random_range_ratio": 0.0, + "extra_body": { + "voice": "Vivian", + "language": "English" + }, + "ignore_eos": true, + "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration" + } + ] + } +] diff --git a/tests/dfx/stability/tests/test_qwen_image.json b/tests/dfx/stability/tests/test_qwen_image.json new file mode 100644 index 000000000..f3dd93f6f --- /dev/null +++ b/tests/dfx/stability/tests/test_qwen_image.json @@ -0,0 +1,28 @@ +[ + { + "test_name": "test_qwen_image_stability", + "server_params": { + "model": "Qwen/Qwen-Image" + }, + "benchmark_params": [ + { + "dataset": "random", + "task": "t2i", + "backend": "vllm-omni", + "duration_sec": 86400, + "max_concurrency": 1, + "num_prompts_per_batch": 10, + "width": { + "min": 512, + "max": 2048 + }, + "height": { + "min": 512, + "max": 2048 + }, + "num_inference_steps": 50, + "enable_negative_prompt": true + } + ] + } +] diff --git a/tests/dfx/stability/tests/test_wan22.json b/tests/dfx/stability/tests/test_wan22.json new file mode 100644 index 000000000..c787ce96a --- /dev/null +++ b/tests/dfx/stability/tests/test_wan22.json @@ -0,0 +1,31 @@ +[ + { + "test_name": "test_wan22_i2v_stability_v1_videos", + "server_params": { + "model": "Wan-AI/Wan2.2-I2V-A14B-Diffusers", + "serve_args": { + "ulysses-degree": 2, + "vae-patch-parallel-size": 2, + "tensor-parallel-size": 1, + "use-hsdp": true, + "vae-use-slicing": true, + "vae-use-tiling": true + } + }, + "benchmark_params": [ + { + "dataset": "random", + "task": "i2v", + "backend": "v1/videos", + "duration_sec": 86400, + "max_concurrency": 1, + "num_prompts_per_batch": 20, + "enable_negative_prompt": true, + "random_request_config": [ + {"width": 832, "height": 480, "num_inference_steps": 4, "num_frames": 81, "fps": 16, "weight": 0.5}, + {"width": 1280, "height": 720, "num_inference_steps": 4, "num_frames": 121, "fps": 16, "weight": 0.5} + ] + } + ] + } +] diff --git a/tests/diffusion/models/glm_image/test_glm_image_sp.py b/tests/diffusion/models/glm_image/test_glm_image_sp.py index 40d1c8730..06a1a116d 100644 --- a/tests/diffusion/models/glm_image/test_glm_image_sp.py +++ b/tests/diffusion/models/glm_image/test_glm_image_sp.py @@ -6,6 +6,8 @@ import pytest from vllm_omni.diffusion.data import DiffusionParallelConfig +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + @pytest.fixture(scope="function", autouse=True) def setup_sp_groups(mocker): @@ -24,9 +26,6 @@ def setup_sp_groups(mocker): yield -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - def test_glm_image_sp_plan_defined(): """Test that _sp_plan is properly defined on GlmImageTransformer2DModel.""" from vllm_omni.diffusion.models.glm_image.glm_image_transformer import ( diff --git a/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py b/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py new file mode 100644 index 000000000..665126df7 --- /dev/null +++ b/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py @@ -0,0 +1,230 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Unit tests for LTX-2.3 pipeline integration. + +These tests verify: +- Pipeline is properly registered in the diffusion registry +- Post-process function is registered +- Cache-DiT enablers are registered +- Pipeline does NOT inherit from LTX2Pipeline +- Vocoder sample rate detection logic +- Re-export module works correctly +""" + +import json +import os +import tempfile + +import pytest + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class TestPipelineIndependence: + """Verify LTX23Pipeline is fully independent from LTX2Pipeline.""" + + def test_ltx23_pipeline_does_not_inherit_from_ltx2(self): + """LTX23Pipeline must NOT inherit from LTX2Pipeline.""" + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline + + assert not issubclass(LTX23Pipeline, LTX2Pipeline), ( + "LTX23Pipeline should be fully independent and not inherit from LTX2Pipeline" + ) + + def test_ltx23_pipeline_is_nn_module(self): + """LTX23Pipeline must be an nn.Module.""" + import torch.nn as nn + + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline + + assert issubclass(LTX23Pipeline, nn.Module) + + def test_ltx23_pipeline_has_progress_bar(self): + """LTX23Pipeline must mix in ProgressBarMixin.""" + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline + from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin + + assert issubclass(LTX23Pipeline, ProgressBarMixin) + + +class TestRegistryIntegration: + """Verify all LTX-2.3 pipeline variants are registered.""" + + def test_pipeline_models_registered(self): + """LTX-2.3 pipeline variants must be in _DIFFUSION_MODELS.""" + from vllm_omni.diffusion.registry import _DIFFUSION_MODELS + + expected = [ + "LTX23Pipeline", + "LTX23ImageToVideoPipeline", + ] + for name in expected: + assert name in _DIFFUSION_MODELS, f"{name} not found in _DIFFUSION_MODELS" + + def test_pipeline_module_paths(self): + """Registry entries must point to the correct modules.""" + from vllm_omni.diffusion.registry import _DIFFUSION_MODELS + + # T2V -> pipeline_ltx2_3 + assert _DIFFUSION_MODELS["LTX23Pipeline"] == ("ltx2", "pipeline_ltx2_3", "LTX23Pipeline") + + # I2V -> pipeline_ltx2_3_image2video + assert _DIFFUSION_MODELS["LTX23ImageToVideoPipeline"] == ( + "ltx2", + "pipeline_ltx2_3_image2video", + "LTX23ImageToVideoPipeline", + ) + + def test_post_process_funcs_registered(self): + """Pipeline variants must map to get_ltx2_post_process_func.""" + from vllm_omni.diffusion.registry import _DIFFUSION_POST_PROCESS_FUNCS + + expected = [ + "LTX23Pipeline", + "LTX23ImageToVideoPipeline", + ] + for name in expected: + assert name in _DIFFUSION_POST_PROCESS_FUNCS, f"{name} not in _DIFFUSION_POST_PROCESS_FUNCS" + assert _DIFFUSION_POST_PROCESS_FUNCS[name] == "get_ltx2_post_process_func" + + def test_cache_dit_enablers_registered(self): + """Pipeline variants must be registered in CUSTOM_DIT_ENABLERS.""" + from vllm_omni.diffusion.cache.cache_dit_backend import CUSTOM_DIT_ENABLERS + + expected = [ + "LTX23Pipeline", + "LTX23ImageToVideoPipeline", + ] + for name in expected: + assert name in CUSTOM_DIT_ENABLERS, f"{name} not in CUSTOM_DIT_ENABLERS" + + +class TestVocoderSampleRateDetection: + """Test _detect_vocoder_output_sample_rate logic.""" + + def test_detects_48khz_from_config(self): + """Should detect output_sampling_rate=48000 from vocoder/config.json.""" + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import _detect_vocoder_output_sample_rate + + with tempfile.TemporaryDirectory() as tmpdir: + vocoder_dir = os.path.join(tmpdir, "vocoder") + os.makedirs(vocoder_dir) + with open(os.path.join(vocoder_dir, "config.json"), "w") as f: + json.dump({"output_sampling_rate": 48000, "input_sampling_rate": 16000}, f) + + result = _detect_vocoder_output_sample_rate(tmpdir) + assert result == 48000 + + def test_returns_none_for_no_output_sr(self): + """Should return None if vocoder config has no output_sampling_rate.""" + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import _detect_vocoder_output_sample_rate + + with tempfile.TemporaryDirectory() as tmpdir: + vocoder_dir = os.path.join(tmpdir, "vocoder") + os.makedirs(vocoder_dir) + with open(os.path.join(vocoder_dir, "config.json"), "w") as f: + json.dump({"sampling_rate": 16000}, f) + + result = _detect_vocoder_output_sample_rate(tmpdir) + assert result is None + + def test_returns_none_for_missing_directory(self): + """Should return None if vocoder directory doesn't exist.""" + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import _detect_vocoder_output_sample_rate + + result = _detect_vocoder_output_sample_rate("/nonexistent/path") + assert result is None + + +class TestPostProcessFunction: + """Test the post-process function factory.""" + + def test_post_process_includes_audio_sample_rate(self): + """Post-process func should include audio_sample_rate when detected.""" + import torch + + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import get_ltx2_post_process_func + + with tempfile.TemporaryDirectory() as tmpdir: + vocoder_dir = os.path.join(tmpdir, "vocoder") + os.makedirs(vocoder_dir) + with open(os.path.join(vocoder_dir, "config.json"), "w") as f: + json.dump({"output_sampling_rate": 48000}, f) + + # Create a minimal od_config mock + class MockConfig: + model = tmpdir + + func = get_ltx2_post_process_func(MockConfig()) + + video = torch.zeros(1, 3, 4, 64, 64) + audio = torch.zeros(1, 1, 48000) + result = func((video, audio)) + + assert isinstance(result, dict) + assert "video" in result + assert "audio" in result + assert result["audio_sample_rate"] == 48000 + + def test_post_process_without_vocoder_config(self): + """Post-process func should work without vocoder config (no audio_sample_rate key).""" + import torch + + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import get_ltx2_post_process_func + + class MockConfig: + model = "/nonexistent/path" + + func = get_ltx2_post_process_func(MockConfig()) + + video = torch.zeros(1, 3, 4, 64, 64) + audio = torch.zeros(1, 1, 16000) + result = func((video, audio)) + + assert isinstance(result, dict) + assert "video" in result + assert "audio" in result + assert "audio_sample_rate" not in result + + +class TestReExportModule: + """Test that pipeline_ltx2_3_image2video.py correctly re-exports.""" + + def test_i2v_classes_importable(self): + """I2V classes must be importable from the re-export module.""" + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline + + assert LTX23ImageToVideoPipeline is not None + + def test_post_process_func_importable(self): + """get_ltx2_post_process_func must be importable from re-export module.""" + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import get_ltx2_post_process_func + + assert callable(get_ltx2_post_process_func) + + def test_i2v_classes_are_same_as_direct_import(self): + """Re-exported classes must be the same objects as direct imports.""" + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23ImageToVideoPipeline as Direct + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import ( + LTX23ImageToVideoPipeline as ReExported, + ) + + assert Direct is ReExported + + +class TestInitExports: + """Test that __init__.py exports all LTX-2.3 classes.""" + + def test_all_ltx23_classes_exported(self): + """All LTX-2.3 pipeline classes must be in the ltx2 package __all__.""" + from vllm_omni.diffusion.models import ltx2 + + expected_classes = [ + "LTX23Pipeline", + "LTX23ImageToVideoPipeline", + ] + for name in expected_classes: + assert hasattr(ltx2, name), f"{name} not exported from ltx2 package" + assert name in ltx2.__all__, f"{name} not in ltx2.__all__" diff --git a/tests/diffusion/models/t5_encoder/test_t5_encoder_prefix.py b/tests/diffusion/models/t5_encoder/test_t5_encoder_prefix.py new file mode 100644 index 000000000..039150f09 --- /dev/null +++ b/tests/diffusion/models/t5_encoder/test_t5_encoder_prefix.py @@ -0,0 +1,164 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for T5EncoderModel prefix handling and weight loading fix.""" + +import pytest +import torch +from transformers import T5Config +from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config + +from vllm_omni.diffusion.models.t5_encoder.t5_encoder import T5EncoderModel + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +_SMALL_T5_CONFIG = dict( + d_model=64, + d_kv=8, + d_ff=128, + num_heads=8, + num_layers=2, + vocab_size=256, + relative_attention_num_buckets=32, + relative_attention_max_distance=128, + is_gated_act=True, + dense_act_fn="gelu_new", + layer_norm_epsilon=1e-6, + feed_forward_proj="gated-gelu", +) + +_T5_MODULE = "vllm_omni.diffusion.models.t5_encoder.t5_encoder" + + +@pytest.fixture +def t5_config() -> T5Config: + return T5Config(**_SMALL_T5_CONFIG) + + +@pytest.fixture(scope="function", autouse=True) +def setup_vllm_config(monkeypatch, mocker): + """Set up VllmConfig and TP=2 mocks for tests.""" + device_config = DeviceConfig(device="cpu") + + monkeypatch.setattr("vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", lambda: 2) + monkeypatch.setattr(f"{_T5_MODULE}.get_tensor_model_parallel_world_size", lambda: 2) + monkeypatch.setattr( + "vllm.model_executor.layers.vocab_parallel_embedding.get_tensor_model_parallel_world_size", + lambda: 2, + ) + + monkeypatch.setattr(f"{_T5_MODULE}.get_tensor_model_parallel_rank", lambda: 0) + monkeypatch.setattr( + "vllm.model_executor.layers.vocab_parallel_embedding.get_tensor_model_parallel_rank", + lambda: 0, + ) + + mock_tp_group = mocker.MagicMock() + mock_tp_group.world_size = 2 + mocker.patch("vllm.distributed.parallel_state.get_tp_group", return_value=mock_tp_group) + + monkeypatch.setattr(f"{_T5_MODULE}.get_act_fn", lambda _: torch.nn.GELU()) + + with set_current_vllm_config(VllmConfig(device_config=device_config)): + yield + + +class TestT5EncoderModelPrefixHandling: + """Test that T5EncoderModel correctly handles prefix attribute.""" + + def test_prefix_stored_in_model(self, t5_config): + """Test that prefix is stored in the model when provided.""" + prefix = "text_encoder" + model = T5EncoderModel(t5_config, prefix=prefix) + assert hasattr(model, "prefix") + assert model.prefix == prefix + + def test_prefix_empty_by_default(self, t5_config): + """Test that prefix defaults to empty string when not provided.""" + model = T5EncoderModel(t5_config) + assert hasattr(model, "prefix") + assert model.prefix == "" + + +class TestT5EncoderModelWeightLoadingWithPrefix: + """Test weight loading with prefix handling.""" + + def test_load_weights_with_prefix(self, t5_config): + """Test that weights without prefix are loaded when model has prefix.""" + config = T5Config(**{**_SMALL_T5_CONFIG, "num_layers": 1}) + model = T5EncoderModel(config, prefix="text_encoder") + + inner_dim = config.num_heads * config.d_kv + + weights = [ + ("encoder.block.0.layer.0.SelfAttention.q.weight", torch.randn(inner_dim, config.d_model)), + ("encoder.block.0.layer.0.SelfAttention.k.weight", torch.randn(inner_dim, config.d_model)), + ("encoder.block.0.layer.0.SelfAttention.v.weight", torch.randn(inner_dim, config.d_model)), + ] + + loaded = model.load_weights(weights) + assert len(loaded) > 0 + + def test_load_weights_embed_tokens_shared_sync(self, t5_config): + """Test that embed_tokens and shared weights are synced.""" + model = T5EncoderModel(t5_config, prefix="text_encoder") + + d_model = t5_config.d_model + vocab_size = t5_config.vocab_size + + embed_weight = torch.randn(vocab_size, d_model) + weights = [ + ("encoder.embed_tokens.weight", embed_weight.clone()), + ] + + model.load_weights(weights) + + shared_param = model.shared.weight + embed_param = model.encoder.embed_tokens.weight + + assert torch.allclose(shared_param, embed_param), ( + "shared and embed_tokens should have the same weights after loading" + ) + + def test_load_weights_shared_without_prefix(self, t5_config): + """Test shared.weight is recognized without relying on dot context.""" + model = T5EncoderModel(t5_config, prefix="text_encoder") + + shared_weight = torch.randn(t5_config.vocab_size, t5_config.d_model) + loaded = model.load_weights([("shared.weight", shared_weight)]) + + assert "shared.weight" in loaded + assert torch.allclose(model.shared.weight, model.encoder.embed_tokens.weight) + + def test_unmatched_weights_are_not_reported_loaded(self, t5_config): + """Test that skipped checkpoint weights are not added to loaded_params.""" + model = T5EncoderModel(t5_config, prefix="text_encoder") + + loaded = model.load_weights( + [ + ( + "text_encoder.encoder.block.0.layer.0.SelfAttention.missing.weight", + torch.randn(t5_config.d_model, t5_config.d_model), + ), + ] + ) + + assert loaded == set() + + +class TestT5EncoderModelWeightLoadingWithoutPrefix: + """Test weight loading without prefix.""" + + def test_load_weights_without_prefix(self, t5_config): + """Test that weights without prefix are loaded correctly.""" + config = T5Config(**{**_SMALL_T5_CONFIG, "num_layers": 1}) + model = T5EncoderModel(config) + + inner_dim = config.num_heads * config.d_kv + + weights = [ + ("encoder.block.0.layer.0.SelfAttention.q.weight", torch.randn(inner_dim, config.d_model)), + ] + + loaded = model.load_weights(weights) + assert len(loaded) > 0 diff --git a/tests/diffusion/offloader/test_module_collector.py b/tests/diffusion/offloader/test_module_collector.py new file mode 100644 index 000000000..ab15ad8df --- /dev/null +++ b/tests/diffusion/offloader/test_module_collector.py @@ -0,0 +1,240 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Unit tests for ModuleDiscovery and SupportsModuleOffload.""" + +from typing import ClassVar + +import pytest +from torch import nn + +from vllm_omni.diffusion.models.interface import SupportsModuleOffload +from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery + +pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model] + +# NOTE: tests for skipped/warned attributes verify the *behavioral* +# outcome (attribute excluded from results) but do not assert on log +# output. vllm's logger sets propagate=False, preventing caplog from +# capturing records. See https://github.com/pytest-dev/pytest/issues/3697 + + +# --------------------------------------------------------------------------- +# Test pipelines +# --------------------------------------------------------------------------- + + +class FallbackPipeline(nn.Module): + """Pipeline with standard attribute names (no protocol).""" + + def __init__(self): + super().__init__() + self.transformer = nn.Linear(10, 10) + self.text_encoder = nn.Linear(10, 10) + self.text_encoder_2 = nn.Linear(10, 10) + self.vae = nn.Linear(10, 10) + + +class NonModuleAttrPipeline(nn.Module): + """Pipeline where an attribute is not an nn.Module (fallback path).""" + + def __init__(self): + super().__init__() + self.transformer = nn.Linear(10, 10) + self.text_encoder = "not_a_module" + self.vae = nn.Linear(10, 10) + + +class DuplicateAttrPipeline(nn.Module): + """Pipeline where two encoder attrs point to the same module.""" + + def __init__(self): + super().__init__() + self.transformer = nn.Linear(10, 10) + encoder = nn.Linear(10, 10) + self.text_encoder = encoder + self.text_encoder_2 = encoder + self.vae = nn.Linear(10, 10) + + +class ProtocolPipeline(nn.Module, SupportsModuleOffload): + """Pipeline with non-standard names, using the protocol.""" + + _dit_modules: ClassVar[list[str]] = ["gen_transformer"] + _encoder_modules: ClassVar[list[str]] = ["mllm", "vision_model"] + _vae_modules: ClassVar[list[str]] = ["gen_vae"] + + def __init__(self): + super().__init__() + self.gen_transformer = nn.Linear(10, 10) + self.mllm = nn.Linear(10, 10) + self.vision_model = nn.Linear(10, 10) + self.gen_vae = nn.Linear(10, 10) + # Standard name present but NOT declared — should be ignored + self.transformer = nn.Linear(10, 10) + + +class MissingAttrPipeline(nn.Module, SupportsModuleOffload): + """Pipeline that declares a non-existent attribute.""" + + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["nonexistent_encoder"] + _vae_modules: ClassVar[list[str]] = ["vae"] + + def __init__(self): + super().__init__() + self.transformer = nn.Linear(10, 10) + self.vae = nn.Linear(10, 10) + + +class MissingIntermediatePipeline(nn.Module, SupportsModuleOffload): + """Pipeline with dotted path referencing non-existent intermediate.""" + + _dit_modules: ClassVar[list[str]] = ["nonexistent.transformer"] + _encoder_modules: ClassVar[list[str]] = [] + _vae_modules: ClassVar[list[str]] = [] + + def __init__(self): + super().__init__() + + +class NestedPipeline(nn.Module, SupportsModuleOffload): + """Pipeline with nested modules accessed via dotted paths.""" + + _dit_modules: ClassVar[list[str]] = ["pipe.transformer"] + _encoder_modules: ClassVar[list[str]] = ["pipe.text_encoder"] + _vae_modules: ClassVar[list[str]] = ["vae"] + + def __init__(self): + super().__init__() + self.pipe = nn.Module() + self.pipe.transformer = nn.Linear(10, 10) + self.pipe.text_encoder = nn.Linear(10, 10) + self.vae = nn.Linear(10, 10) + + +class ResidentPipeline(nn.Module, SupportsModuleOffload): + """Pipeline with resident modules that must stay on GPU.""" + + _dit_modules: ClassVar[list[str]] = ["language_model.model"] + _encoder_modules: ClassVar[list[str]] = [] + _vae_modules: ClassVar[list[str]] = ["vae"] + _resident_modules: ClassVar[list[str]] = [ + "bagel.time_embedder", + "bagel.vae2llm", + ] + + def __init__(self): + super().__init__() + self.language_model = nn.Module() + self.language_model.model = nn.Linear(10, 10) + self.bagel = nn.Module() + self.bagel.time_embedder = nn.Linear(10, 10) + self.bagel.vae2llm = nn.Linear(10, 10) + self.vae = nn.Linear(10, 10) + + +class MultiVaePipeline(nn.Module, SupportsModuleOffload): + """Pipeline with multiple VAEs.""" + + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["text_encoder"] + _vae_modules: ClassVar[list[str]] = ["vae", "audio_vae"] + + def __init__(self): + super().__init__() + self.transformer = nn.Linear(10, 10) + self.text_encoder = nn.Linear(10, 10) + self.vae = nn.Linear(10, 10) + self.audio_vae = nn.Linear(10, 10) + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +class TestFallbackDiscovery: + """Test the fallback attribute scan (no SupportsModuleOffload).""" + + def test_discovers_standard_attrs(self): + pipeline = FallbackPipeline() + result = ModuleDiscovery.discover(pipeline) + + assert not isinstance(pipeline, SupportsModuleOffload) + assert result.dit_names == ["transformer"] + assert result.dits[0] is pipeline.transformer + assert result.encoder_names == ["text_encoder", "text_encoder_2"] + assert result.vaes[0] is pipeline.vae + assert result.resident_modules == [] + + def test_deduplicates_encoders(self): + pipeline = DuplicateAttrPipeline() + result = ModuleDiscovery.discover(pipeline) + + assert len(result.encoders) == 1 + assert result.encoder_names == ["text_encoder"] + + def test_skips_non_module_attr(self): + pipeline = NonModuleAttrPipeline() + result = ModuleDiscovery.discover(pipeline) + + assert len(result.encoders) == 0 + + +class TestProtocolDiscovery: + """Test discovery via SupportsModuleOffload protocol.""" + + def test_discovers_declared_attrs_and_ignores_undeclared(self): + pipeline = ProtocolPipeline() + result = ModuleDiscovery.discover(pipeline) + + assert isinstance(pipeline, SupportsModuleOffload) + assert result.dit_names == ["gen_transformer"] + assert result.encoder_names == ["mllm", "vision_model"] + assert len(result.vaes) == 1 + # self.transformer exists but is NOT in _dit_modules + assert "transformer" not in result.dit_names + # No _resident_modules declared — defaults to empty + assert result.resident_modules == [] + + def test_skips_missing_attr(self): + pipeline = MissingAttrPipeline() + result = ModuleDiscovery.discover(pipeline) + + assert len(result.encoders) == 0 + + def test_skips_missing_intermediate(self): + result = ModuleDiscovery.discover(MissingIntermediatePipeline()) + + assert len(result.dits) == 0 + + def test_dotted_path_resolves_nested_modules(self): + pipeline = NestedPipeline() + result = ModuleDiscovery.discover(pipeline) + + assert result.dit_names == ["pipe.transformer"] + assert result.dits[0] is pipeline.pipe.transformer + assert result.encoder_names == ["pipe.text_encoder"] + assert result.encoders[0] is pipeline.pipe.text_encoder + assert result.vaes[0] is pipeline.vae + + def test_resident_modules(self): + pipeline = ResidentPipeline() + result = ModuleDiscovery.discover(pipeline) + + assert result.resident_names == [ + "bagel.time_embedder", + "bagel.vae2llm", + ] + assert result.resident_modules[0] is pipeline.bagel.time_embedder + assert result.resident_modules[1] is pipeline.bagel.vae2llm + assert result.dits[0] is pipeline.language_model.model + + def test_multiple_vaes(self): + pipeline = MultiVaePipeline() + result = ModuleDiscovery.discover(pipeline) + + assert len(result.vaes) == 2 + assert result.vaes[0] is pipeline.vae + assert result.vaes[1] is pipeline.audio_vae diff --git a/tests/diffusion/quantization/test_component_routing.py b/tests/diffusion/quantization/test_component_routing.py new file mode 100644 index 000000000..c8b3837e2 --- /dev/null +++ b/tests/diffusion/quantization/test_component_routing.py @@ -0,0 +1,408 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for component routing for quantization.""" + +from unittest.mock import MagicMock + +import pytest +import torch +from vllm.model_executor.layers.quantization.base_config import ( + QuantizationConfig, +) +from vllm.model_executor.models.utils import WeightsMapper + +from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import ( + PRE_QUANTIZED_METHODS, +) +from vllm_omni.quantization.component_config import ( + ComponentQuantizationConfig, +) +from vllm_omni.quantization.inc_config import OmniINCConfig + +pytestmark = [pytest.mark.core_model] + + +# --------------------------------------------------------------------------- +# Helpers: lightweight mock quant configs +# --------------------------------------------------------------------------- + + +class _MockQuantConfig(QuantizationConfig): + """Minimal mock that only implements get_name().""" + + def __init__(self, name: str, **attrs): + self._name = name + for k, v in attrs.items(): + setattr(self, k, v) + + def get_name(self) -> str: + return self._name + + def get_quant_method(self, layer, prefix): + return MagicMock() + + @classmethod + def get_supported_act_dtypes(cls): + return [torch.bfloat16, torch.float16] + + def get_min_capability(self): + return 0 + + @classmethod + def from_config(cls, config): + raise NotImplementedError + + def get_config_filenames(self): + return [] + + +def _make_inc_config(block_names="thinker.model.layers,talker.model.layers", extra_config=None): + """Create a real OmniINCConfig with block_name_to_quantize.""" + return OmniINCConfig( + weight_bits=4, + group_size=128, + sym=True, + block_name_to_quantize=block_names, + extra_config=extra_config or {}, + ) + + +THINKER_MAPPER = WeightsMapper( + orig_to_new_prefix={ + "thinker.lm_head.": "language_model.lm_head.", + "thinker.model.": "language_model.model.", + "thinker.": "", + } +) + +TALKER_MAPPER = WeightsMapper( + orig_to_new_prefix={ + "talker.codec_head.": "language_model.lm_head.", + "talker.model.": "language_model.model.", + "talker.thinker_to_talker_proj.": "thinker_to_talker_proj.", + "talker.": "", + } +) + + +# =================================================================== +# 1. OmniINCConfig.apply_vllm_mapper +# =================================================================== + + +class TestApplyVllmMapper: + def test_inc_csv_string_normalized_to_list(self): + """CSV string block_name_to_quantize is split into a list.""" + cfg = _make_inc_config("thinker.model.layers,talker.model.layers") + cfg.apply_vllm_mapper(THINKER_MAPPER) + assert isinstance(cfg.block_name_to_quantize, list) + + def test_thinker_blocks_remapped(self): + """thinker.model.layers -> language_model.model.layers after apply_vllm_mapper.""" + cfg = _make_inc_config("thinker.model.layers,talker.model.layers") + cfg.apply_vllm_mapper(THINKER_MAPPER) + assert any("language_model.model.layers" in b for b in cfg.block_name_to_quantize) + + def test_cross_stage_blocks_kept_unchanged(self): + """Blocks not matching any mapper prefix are kept unchanged (harmless).""" + cfg = _make_inc_config("thinker.model.layers,talker.model.layers") + cfg.apply_vllm_mapper(THINKER_MAPPER) + # talker.model.layers doesn't match any thinker mapper prefix → stays as-is + assert "talker.model.layers" in cfg.block_name_to_quantize + + def test_talker_remap(self): + """talker.model.layers -> language_model.model.layers with talker mapper.""" + cfg = _make_inc_config("thinker.model.layers,talker.model.layers") + cfg.apply_vllm_mapper(TALKER_MAPPER) + assert any("language_model.model.layers" in b for b in cfg.block_name_to_quantize) + # thinker.model.layers doesn't match talker mapper → stays as-is + assert "thinker.model.layers" in cfg.block_name_to_quantize + + def test_extra_config_keys_remapped(self): + """Regex keys in extra_config get their escaped-dot prefixes remapped.""" + extra = { + r".*thinker\.model\.layers\.0\.mlp\.gate.*": {"bits": 16, "data_type": "float"}, + } + cfg = _make_inc_config("thinker.model.layers", extra_config=extra) + cfg.apply_vllm_mapper(THINKER_MAPPER) + # The key should now reference the vLLM runtime path + assert any("language_model" in k for k in cfg.extra_config) + # Original thinker\.model prefix should be replaced + assert not any(r"thinker\.model" in k for k in cfg.extra_config) + + def test_single_block_name(self): + """Only one block name (not CSV) still works.""" + cfg = _make_inc_config("thinker.model.layers") + cfg.apply_vllm_mapper(THINKER_MAPPER) + assert any("language_model.model.layers" in b for b in cfg.block_name_to_quantize) + + def test_already_list_block_names(self): + """block_name_to_quantize already a list (not CSV string) works.""" + cfg = _make_inc_config(["thinker.model.layers", "talker.model.layers"]) + cfg.apply_vllm_mapper(THINKER_MAPPER) + assert isinstance(cfg.block_name_to_quantize, list) + assert any("language_model.model.layers" in b for b in cfg.block_name_to_quantize) + + def test_mutates_in_place(self): + """apply_vllm_mapper mutates the config in place (same as upstream INCConfig).""" + cfg = _make_inc_config("thinker.model.layers") + original_id = id(cfg) + cfg.apply_vllm_mapper(THINKER_MAPPER) + assert id(cfg) == original_id + + # -- Stage prefix tests (runtime prefix = container + internal name) -- + + def test_thinker_block_has_stage_prefix(self): + """Mapped block name must start with 'thinker.' so runtime startswith() works.""" + cfg = _make_inc_config("thinker.model.layers,talker.model.layers") + cfg.apply_vllm_mapper(THINKER_MAPPER) + assert "thinker.language_model.model.layers" in cfg.block_name_to_quantize + + def test_talker_block_has_stage_prefix(self): + """Mapped block name must start with 'talker.' so runtime startswith() works.""" + cfg = _make_inc_config("thinker.model.layers,talker.model.layers") + cfg.apply_vllm_mapper(TALKER_MAPPER) + assert "talker.language_model.model.layers" in cfg.block_name_to_quantize + + def test_thinker_block_matches_runtime_prefix(self): + """Simulates get_layer_config's startswith() check for FusedMoE layers.""" + cfg = _make_inc_config("thinker.model.layers,talker.model.layers") + cfg.apply_vllm_mapper(THINKER_MAPPER) + runtime_prefix = "thinker.language_model.model.layers.0.mlp.experts" + assert any(runtime_prefix.startswith(b) for b in cfg.block_name_to_quantize) + + def test_talker_block_matches_runtime_prefix(self): + """Simulates get_layer_config's startswith() check for talker FusedMoE.""" + cfg = _make_inc_config("thinker.model.layers,talker.model.layers") + cfg.apply_vllm_mapper(TALKER_MAPPER) + runtime_prefix = "talker.language_model.model.layers.0.mlp.experts" + assert any(runtime_prefix.startswith(b) for b in cfg.block_name_to_quantize) + + def test_extra_config_plain_key_has_stage_prefix(self): + """Plain extra_config keys are remapped with stage prefix.""" + extra = { + "talker.model.layers.0.mlp.shared_expert_gate": {"bits": 16}, + } + cfg = _make_inc_config("talker.model.layers", extra_config=extra) + cfg.apply_vllm_mapper(TALKER_MAPPER) + assert "talker.language_model.model.layers.0.mlp.shared_expert_gate" in cfg.extra_config + + def test_extra_config_regex_key_still_works(self): + """Regex extra_config keys use re.search so no stage prefix needed.""" + import re + + extra = { + r".*thinker\.model\.layers\.0\.mlp\.gate.*": {"bits": 16}, + } + cfg = _make_inc_config("thinker.model.layers", extra_config=extra) + cfg.apply_vllm_mapper(THINKER_MAPPER) + runtime_name = "thinker.language_model.model.layers.0.mlp.gate" + matched = any(re.search(k, runtime_name) for k in cfg.extra_config) + assert matched + + +# =================================================================== +# 2. OmniINCConfig upgrade helpers +# =================================================================== + + +class TestOmniINCConfigUpgrade: + def test_maybe_upgrade_none(self): + assert OmniINCConfig.maybe_upgrade(None) is None + + def test_maybe_upgrade_non_inc(self): + """Non-INC configs are passed through unchanged.""" + cfg = _MockQuantConfig("fp8") + assert OmniINCConfig.maybe_upgrade(cfg) is cfg + + def test_maybe_upgrade_already_omni(self): + """Already OmniINCConfig is returned as-is.""" + cfg = _make_inc_config() + assert OmniINCConfig.maybe_upgrade(cfg) is cfg + + def test_maybe_upgrade_vanilla_inc(self): + """Vanilla INCConfig is promoted to OmniINCConfig.""" + from vllm.model_executor.layers.quantization.inc import INCConfig + + vanilla = INCConfig(weight_bits=4, group_size=128, sym=True) + upgraded = OmniINCConfig.maybe_upgrade(vanilla) + assert isinstance(upgraded, OmniINCConfig) + assert upgraded.weight_bits == 4 + assert upgraded.group_size == 128 + + +# =================================================================== +# 2. Three-branch thinker routing (simulated) +# =================================================================== + + +def _simulate_thinker_routing(quant_config): + """Simulate the three-branch routing in thinker __init__. + + Returns (visual_quant_config, language_quant_config, wrapped_vllm_quant). + """ + if isinstance(quant_config, ComponentQuantizationConfig): + visual_quant_config = quant_config.resolve("visual") + language_quant_config = quant_config.resolve("language_model") + return visual_quant_config, language_quant_config, quant_config + elif quant_config is not None: + if quant_config.get_name() in PRE_QUANTIZED_METHODS: + return quant_config, quant_config, quant_config + else: + language_quant_config = quant_config + wrapped = ComponentQuantizationConfig( + component_configs={"language_model": quant_config}, + default_config=None, + ) + return None, language_quant_config, wrapped + else: + return None, None, None + + +class TestThinkerRouting: + def test_none(self): + vis, lang, wrapped = _simulate_thinker_routing(None) + assert vis is None + assert lang is None + assert wrapped is None + + @pytest.mark.parametrize("method", ["modelopt", "modelopt_fp4", "modelopt_mxfp8"]) + def test_pre_quantized_all_components(self, method): + """Pre-quantized methods pass config to all components.""" + cfg = _MockQuantConfig(method) + vis, lang, wrapped = _simulate_thinker_routing(cfg) + assert vis is cfg + assert lang is cfg + assert wrapped is cfg + + def test_fp8_dynamic_language_only(self): + """fp8 dynamic: visual=None, language gets original config.""" + cfg = _MockQuantConfig("fp8") + vis, lang, wrapped = _simulate_thinker_routing(cfg) + assert vis is None + assert lang is cfg + assert isinstance(wrapped, ComponentQuantizationConfig) + assert wrapped.resolve("language_model") is cfg + assert wrapped.resolve("visual") is None + + def test_inc_autoround_language_only(self): + """INC/AutoRound: not in _PRE_QUANTIZED_METHODS -> wrapped like fp8.""" + cfg = _MockQuantConfig("inc") + vis, lang, wrapped = _simulate_thinker_routing(cfg) + assert vis is None + assert lang is cfg + assert isinstance(wrapped, ComponentQuantizationConfig) + + def test_component_config_passthrough(self): + """Explicit ComponentQuantizationConfig is used directly.""" + inner_fp8 = _MockQuantConfig("fp8") + inner_modelopt = _MockQuantConfig("modelopt") + cqc = ComponentQuantizationConfig( + component_configs={ + "visual": inner_modelopt, + "language_model": inner_fp8, + } + ) + vis, lang, wrapped = _simulate_thinker_routing(cqc) + assert vis is inner_modelopt + assert lang is inner_fp8 + assert wrapped is cqc + + +# =================================================================== +# 3. Talker visual routing (init_multi_modal guard) +# =================================================================== + + +def _simulate_talker_visual_routing(quant_config): + """Simulate the talker init_multi_modal visual routing.""" + if quant_config is not None and quant_config.get_name() in PRE_QUANTIZED_METHODS: + return quant_config + return None + + +class TestTalkerVisualRouting: + def test_none(self): + assert _simulate_talker_visual_routing(None) is None + + @pytest.mark.parametrize("method", ["modelopt", "modelopt_fp4", "modelopt_mxfp8"]) + def test_pre_quantized_passes_through(self, method): + """Pre-quantized methods pass quant config to visual.""" + cfg = _MockQuantConfig(method) + assert _simulate_talker_visual_routing(cfg) is cfg + + def test_fp8_blocked(self): + """fp8 dynamic must NOT be passed to visual.""" + cfg = _MockQuantConfig("fp8") + assert _simulate_talker_visual_routing(cfg) is None + + def test_inc_blocked(self): + """INC/AutoRound must NOT be passed to visual (not in _PRE_QUANTIZED_METHODS).""" + cfg = _MockQuantConfig("inc") + assert _simulate_talker_visual_routing(cfg) is None + + +# =================================================================== +# 4. ComponentQuantizationConfig.resolve +# =================================================================== + + +class TestComponentResolve: + def test_longest_prefix_match(self): + a = _MockQuantConfig("a") + b = _MockQuantConfig("b") + cqc = ComponentQuantizationConfig(component_configs={"language_model": a, "language_model.model": b}) + assert cqc.resolve("language_model.model.layers.0") is b + assert cqc.resolve("language_model.lm_head") is a + + def test_no_match_returns_default(self): + a = _MockQuantConfig("a") + default = _MockQuantConfig("default") + cqc = ComponentQuantizationConfig( + component_configs={"language_model": a}, + default_config=default, + ) + assert cqc.resolve("visual") is default + + def test_no_match_no_default_returns_none(self): + a = _MockQuantConfig("a") + cqc = ComponentQuantizationConfig( + component_configs={"language_model": a}, + ) + assert cqc.resolve("visual") is None + + def test_get_name(self): + cqc = ComponentQuantizationConfig(component_configs={}) + assert cqc.get_name() == "component" + + def test_get_quant_method_delegates(self): + """get_quant_method dispatches to the resolved config.""" + inner = _MockQuantConfig("fp8") + cqc = ComponentQuantizationConfig( + component_configs={"language_model": inner}, + ) + layer = MagicMock() + result = cqc.get_quant_method(layer, "language_model.model.layers.0.mlp") + assert result is not None # delegates to inner.get_quant_method + + def test_get_quant_method_returns_none_for_unmatched(self): + """get_quant_method returns None when no config matches.""" + inner = _MockQuantConfig("fp8") + cqc = ComponentQuantizationConfig( + component_configs={"language_model": inner}, + ) + layer = MagicMock() + result = cqc.get_quant_method(layer, "visual.blocks.0.mlp") + assert result is None + + def test_min_capability(self): + a = _MockQuantConfig("a") + a.get_min_capability = lambda: 80 + b = _MockQuantConfig("b") + b.get_min_capability = lambda: 70 + cqc = ComponentQuantizationConfig(component_configs={"x": a, "y": b}) + assert cqc.get_min_capability() == 70 + + def test_min_capability_empty(self): + cqc = ComponentQuantizationConfig(component_configs={}) + assert cqc.get_min_capability() == 0 diff --git a/tests/diffusion/test_diffusers_adapter.py b/tests/diffusion/test_diffusers_adapter.py new file mode 100644 index 000000000..ac2ec2e3f --- /dev/null +++ b/tests/diffusion/test_diffusers_adapter.py @@ -0,0 +1,186 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections import namedtuple +from types import SimpleNamespace + +import pytest +import torch +from diffusers import DiffusionPipeline +from PIL import Image + +from vllm_omni.diffusion.data import ( + DiffusionOutput, + DiffusionParallelConfig, + OmniDiffusionConfig, +) +from vllm_omni.diffusion.models.diffusers_adapter import DiffusersAdapterPipeline +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +def _make_od_config(**overrides) -> OmniDiffusionConfig: + od_config = OmniDiffusionConfig( + model="test/model", + model_class_name="DiffusersAdapterPipeline", + dtype=torch.float16, + diffusion_load_format="diffusers", + diffusers_load_kwargs={}, + diffusers_call_kwargs={}, + output_type="pil", + parallel_config=DiffusionParallelConfig(cfg_parallel_size=1, sequence_parallel_size=1), + cache_backend="none", + ) + for key, value in overrides.items(): + setattr(od_config, key, value) + return od_config + + +def _make_request(**overrides) -> OmniDiffusionRequest: + prompt = overrides.pop("prompt", "a test prompt") + negative_prompt = overrides.pop("negative_prompt", None) + prompt_obj: dict[str, str] = {"prompt": prompt} + if negative_prompt is not None: + prompt_obj["negative_prompt"] = negative_prompt + + defaults = { + "prompts": [prompt_obj], + "sampling_params": OmniDiffusionSamplingParams( + num_inference_steps=20, + guidance_scale=7.5, + height=16, + width=16, + num_frames=1, + num_outputs_per_prompt=1, + seed=42, + output_type="pil", + generator_device="cpu", + ), + } + defaults.update(overrides) + return OmniDiffusionRequest(**defaults) + + +class TestDiffusersAdapterPipeline: + def test_adapter_forward_returns_output(self, mocker): + od_config = _make_od_config() + request = _make_request() + stub_image = Image.new("RGB", (request.sampling_params.width, request.sampling_params.height)) # pyright: ignore[reportArgumentType] + + adapter = DiffusersAdapterPipeline(od_config=od_config) + MockPipelineOutput = namedtuple("MockPipelineOutput", ["image"]) + MockPipeline = type("MockPipeline", (DiffusionPipeline,), {}) + adapter._pipeline = MockPipeline() + + mocker.patch.object( + MockPipeline, + "__call__", + return_value=MockPipelineOutput(image=stub_image), + ) + output = adapter.forward(request) + + assert isinstance(output, DiffusionOutput) + assert isinstance(output.output, MockPipelineOutput) + assert output.output.image is stub_image + + @pytest.mark.parametrize( + "feature_id", + ["cfg_parallel", "ulysses", "ring", "teacache", "cache_dit", "enforce_eager", "quantization"], + ) + def test_adapter_guard_unsupported_feature(self, feature_id): + if feature_id == "cfg_parallel": + od_config = _make_od_config( + parallel_config=DiffusionParallelConfig(cfg_parallel_size=2, sequence_parallel_size=1), + cache_backend="none", + ) + elif feature_id == "ulysses": + od_config = _make_od_config( + parallel_config=DiffusionParallelConfig(cfg_parallel_size=1, ulysses_degree=2), + cache_backend="none", + ) + elif feature_id == "ring": + od_config = _make_od_config( + parallel_config=DiffusionParallelConfig(cfg_parallel_size=1, ring_degree=2), + cache_backend="none", + ) + elif feature_id == "teacache": + od_config = _make_od_config( + parallel_config=DiffusionParallelConfig(cfg_parallel_size=1, sequence_parallel_size=1), + cache_backend="tea_cache", + ) + elif feature_id == "cache_dit": + od_config = _make_od_config( + parallel_config=DiffusionParallelConfig(cfg_parallel_size=1, sequence_parallel_size=1), + cache_backend="cache_dit", + ) + elif feature_id == "enforce_eager": + od_config = _make_od_config(enforce_eager=True) + elif feature_id == "quantization": + od_config = _make_od_config(quantization_config=SimpleNamespace(quant_method="fp8")) + else: + raise ValueError(f"Unknown feature ID: {feature_id}") + + with pytest.raises(NotImplementedError): + DiffusersAdapterPipeline(od_config=od_config) + + def test_adapter_guard_unknown_output_type(self, mocker): + """Test that the adapter wraps an unknown output type as-is. + This is useful when `return_dict=True` and the diffusers pipeline returns an OrderedDict subclass.""" + + adapter = DiffusersAdapterPipeline(od_config=_make_od_config()) + raw_output = {"unexpected": "dict-output"} + + MockPipeline = type("MockPipeline", (DiffusionPipeline,), {}) + adapter._pipeline = MockPipeline() + + mocker.patch.object( + MockPipeline, + "__call__", + return_value=raw_output, + ) + output = adapter.forward(_make_request()) + + assert isinstance(output, DiffusionOutput) + assert output.output == raw_output + + def test_adapter_build_call_kwargs(self): + adapter = DiffusersAdapterPipeline( + od_config=_make_od_config( + diffusers_call_kwargs={ + "guidance_scale": 1.25, + "eta": 0.3, + "output_type": "np", + } + ) + ) + req = _make_request( + prompt="a cat on mars", + negative_prompt="low quality", + sampling_params=OmniDiffusionSamplingParams( + num_inference_steps=9, + guidance_scale=8.0, + height=320, + width=640, + num_frames=8, + num_outputs_per_prompt=2, + seed=123, + output_type="pil", + ), + ) + + kwargs = adapter._build_call_kwargs(req) + + assert kwargs["prompt"] == "a cat on mars" + assert kwargs["negative_prompt"] == "low quality" + assert kwargs["num_inference_steps"] == 9 + assert kwargs["guidance_scale"] == 8.0 + assert kwargs["height"] == 320 + assert kwargs["width"] == 640 + assert kwargs["num_frames"] == 8 + assert kwargs["num_images_per_prompt"] == 2 + assert kwargs["output_type"] == "pil" + assert isinstance(kwargs["generator"], torch.Generator) + assert kwargs["generator"].device.type == "cpu" + assert kwargs["generator"].initial_seed() == 123 diff --git a/tests/distributed/omni_coordinator/test_omni_coord_client_for_stage.py b/tests/distributed/omni_coordinator/test_omni_coord_client_for_stage.py index f095dfd49..0ba19c7ff 100644 --- a/tests/distributed/omni_coordinator/test_omni_coord_client_for_stage.py +++ b/tests/distributed/omni_coordinator/test_omni_coord_client_for_stage.py @@ -17,8 +17,6 @@ from vllm_omni.distributed.omni_coordinator import ( pytestmark = [pytest.mark.core_model, pytest.mark.cpu] -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - def _bind_router() -> tuple[zmq.Context, zmq.Socket, str]: ctx = zmq.Context.instance() diff --git a/tests/distributed/omni_coordinator/test_omni_coordinator.py b/tests/distributed/omni_coordinator/test_omni_coordinator.py index 38a595cc7..eff3d429e 100644 --- a/tests/distributed/omni_coordinator/test_omni_coordinator.py +++ b/tests/distributed/omni_coordinator/test_omni_coordinator.py @@ -41,6 +41,74 @@ def _wait_for_instance_list( return None +def _drain_sub_messages(sub: zmq.Socket, max_seconds: float = 0.4) -> None: + """Drain queued SUB messages for a short window.""" + deadline = time.time() + max_seconds + while time.time() < deadline: + _recv_instance_list(sub, timeout_ms=50) + + +def test_omni_coordinator_pub_coalescing_on_rapid_queue_updates(): + """Rapid updates should be coalesced into fewer PUB messages.""" + router_addr = get_engine_client_zmq_addr( + local_only=False, + host="127.0.0.1", + port=0, + ) + pub_addr = get_engine_client_zmq_addr( + local_only=False, + host="127.0.0.1", + port=0, + ) + coordinator = OmniCoordinator( + router_zmq_addr=router_addr, + pub_zmq_addr=pub_addr, + heartbeat_timeout=1000.0, + ) + + sub_ctx = zmq.Context.instance() + sub = sub_ctx.socket(zmq.SUB) + sub.connect(pub_addr) + sub.setsockopt(zmq.SUBSCRIBE, b"") + + time.sleep(0.3) # PUB/SUB slow-joiner + + client = OmniCoordClientForStage( + router_addr, + "tcp://stage:coalesce", + "tcp://stage:coalesce-out", + 0, + ) + + # Wait for initial registration broadcast and clear any queued messages. + msg = _wait_for_instance_list(sub, expected_count=1) + assert msg is not None + _drain_sub_messages(sub) + + # Burst many queue updates in a short period. + update_count = 80 + for i in range(update_count): + client.update_info(queue_length=i) + + # With publish_min_interval=0.1s, received messages over ~1s should be + # much smaller than update_count (coalescing effect). + window_s = 1.1 + deadline = time.time() + window_s + recv_count = 0 + while time.time() < deadline: + if _recv_instance_list(sub, timeout_ms=100) is not None: + recv_count += 1 + + assert recv_count < update_count // 2, ( + f"expected coalesced PUB traffic, got {recv_count} for {update_count} updates" + ) + + client.close() + coordinator.close() + sub.close(0) + sub_ctx.term() + + def test_omni_coordinator_registration_broadcast(): """Verify that after multiple OmniCoordClientForStage instances register, OmniCoordinator publishes an InstanceList containing all registered instances. diff --git a/tests/e2e/accuracy/qwen3_omni/__init__.py b/tests/e2e/accuracy/qwen3_omni/__init__.py new file mode 100644 index 000000000..79a31c4f1 --- /dev/null +++ b/tests/e2e/accuracy/qwen3_omni/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Qwen3-Omni accuracy benchmarks (Daily-Omni / Seed-TTS ``vllm bench serve --omni``).""" diff --git a/tests/e2e/accuracy/qwen3_omni/qwen3_omni_acc_bench_core.py b/tests/e2e/accuracy/qwen3_omni/qwen3_omni_acc_bench_core.py new file mode 100644 index 000000000..2ce86d504 --- /dev/null +++ b/tests/e2e/accuracy/qwen3_omni/qwen3_omni_acc_bench_core.py @@ -0,0 +1,201 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Shared helpers for Qwen3-Omni Daily-Omni / Seed-TTS ``vllm bench serve --omni`` accuracy runs. + +Local dataset paths are **optional**. When ``VLLM_DAILY_OMNI_QA_JSON`` + ``VLLM_DAILY_OMNI_VIDEO_DIR`` +point to existing files, those are used with inline video. Otherwise the benchmark falls back to +the HuggingFace dataset id (``liarliar/Daily-Omni``); QA loads via ``datasets``, and the first +bench request that needs media downloads ``Videos.tar`` from the Hub when no video dir is set. + +Similarly for Seed-TTS: a local directory wins; otherwise ``--dataset-path`` uses the Hub id +and ``huggingface_hub.snapshot_download`` inside ``resolve_seed_tts_root`` pulls files on demand. + +Use :func:`build_acc_benchmark_cli_argv` to assemble ``argv`` for a live Omni server (host/port/model +and small bench defaults) before ``parse_args`` / ``run_acc_benchmark`` in the accuracy driver. +""" + +from __future__ import annotations + +import json +import os +import shutil +import subprocess +from pathlib import Path +from typing import Any, Protocol + +DEFAULT_DAILY_OMNI_HF_REPO = "liarliar/Daily-Omni" +DEFAULT_SEED_TTS_HF_REPO = "zhaochenyang20/seed-tts-eval" + + +class OmniBenchServerEndpoint(Protocol): + """Anything with ``host`` / ``port`` / ``model`` (e.g. :class:`tests.conftest.OmniServer`).""" + + host: str + port: int + model: str + + +def build_acc_benchmark_cli_argv( + server: OmniBenchServerEndpoint, + *, + skip_seed: bool, + skip_daily: bool, + num_prompts: int | None = None, + max_concurrency: int | None = None, +) -> list[str]: + """Prefix argv for :func:`run_qwen_omni_acc_benchmark.parse_acc_benchmark_args` + :func:`run_acc_benchmark`. + + Wires ``--host`` / ``--port`` / ``--model`` to a running Omni OpenAI server, sets small + ``--num-prompts`` / ``--max-concurrency`` defaults (overridable via ``ACC_BENCH_NUM_PROMPTS`` / + ``ACC_BENCH_MAX_CONCURRENCY``), and when Daily-Omni runs adds ``--daily-omni-repo`` so Hub QA + matches :func:`daily_omni_bench_argv` once ``run_acc_benchmark`` mirrors ``--daily-omni-repo`` into env. + """ + n_prompts = int(os.environ.get("ACC_BENCH_NUM_PROMPTS", "2000")) if num_prompts is None else int(num_prompts) + n_conc = int(os.environ.get("ACC_BENCH_MAX_CONCURRENCY", "10")) if max_concurrency is None else int(max_concurrency) + argv = [ + "--host", + server.host, + "--port", + str(server.port), + "--model", + server.model, + "--num-prompts", + str(n_prompts), + "--max-concurrency", + str(n_conc), + ] + if not skip_daily: + repo = os.environ.get("VLLM_DAILY_OMNI_REPO", DEFAULT_DAILY_OMNI_HF_REPO).strip() or DEFAULT_DAILY_OMNI_HF_REPO + argv.extend(["--daily-omni-repo", repo]) + if skip_seed: + argv.append("--skip-seed-tts") + if skip_daily: + argv.append("--skip-daily-omni") + return argv + + +def daily_omni_bench_argv() -> list[str]: + """CLI args for Daily-Omni (after ``vllm bench serve --omni``).""" + qa = os.environ.get("VLLM_DAILY_OMNI_QA_JSON", "").strip() + vd = os.environ.get("VLLM_DAILY_OMNI_VIDEO_DIR", "").strip() + if qa and vd: + qap = Path(qa).expanduser() + vdp = Path(vd).expanduser() + if qap.is_file() and vdp.is_dir(): + return [ + "--dataset-name", + "daily-omni", + "--daily-omni-qa-json", + str(qap.resolve()), + "--daily-omni-video-dir", + str(vdp.resolve()), + "--daily-omni-inline-local-video", + ] + repo = os.environ.get("VLLM_DAILY_OMNI_REPO", DEFAULT_DAILY_OMNI_HF_REPO).strip() or DEFAULT_DAILY_OMNI_HF_REPO + return [ + "--dataset-name", + "daily-omni", + "--dataset-path", + repo, + ] + + +def seed_tts_bench_argv(*, locale: str = "en") -> list[str]: + """CLI args for Seed-TTS (after ``vllm bench serve --omni``).""" + dp = os.environ.get("VLLM_SEED_TTS_DATASET_PATH", "").strip() + if dp: + p = Path(dp).expanduser() + # Preserve Hugging Face repo ids verbatim. Only canonicalize to an + # absolute path when the value actually exists as a local directory. + dataset_path = str(p.resolve()) if p.exists() and p.is_dir() else dp + else: + dataset_path = ( + os.environ.get("VLLM_SEED_TTS_REPO", DEFAULT_SEED_TTS_HF_REPO).strip() or DEFAULT_SEED_TTS_HF_REPO + ) + out = ["--dataset-name", "seed-tts", "--dataset-path", dataset_path] + root = os.environ.get("SEED_TTS_ROOT", "").strip() + if root: + out.extend(["--seed-tts-root", str(Path(root).expanduser().resolve())]) + out.extend(["--seed-tts-locale", locale]) + return out + + +def find_vllm_cli() -> str: + exe = shutil.which("vllm") + if not exe: + raise FileNotFoundError("Could not find `vllm` on PATH (install vLLM-Omni with CLI entrypoints).") + return exe + + +def run_vllm_bench_subprocess(vllm: str, argv: list[str], *, extra_env: dict[str, str] | None = None) -> None: + env = os.environ.copy() + if extra_env: + env.update(extra_env) + subprocess.run([vllm, *argv], env=env, check=True) + + +def load_benchmark_result(path: Path) -> dict[str, Any]: + with path.open(encoding="utf-8") as f: + return json.load(f) + + +def build_serve_common_argv( + *, + host: str, + port: int, + model: str, + num_prompts: int, + max_concurrency: int, + num_warmups: int, + percentile_metrics: str, + result_dir: Path, + result_filename: str, + ready_check_timeout_sec: int | None = None, +) -> list[str]: + out = [ + "bench", + "serve", + "--omni", + "--host", + host, + "--port", + str(port), + "--model", + model, + "--endpoint", + "/v1/chat/completions", + "--backend", + "openai-chat-omni", + "--request-rate", + "inf", + "--num-prompts", + str(num_prompts), + "--max-concurrency", + str(max_concurrency), + "--no-oversample", + "--num-warmups", + str(num_warmups), + "--percentile-metrics", + percentile_metrics, + "--save-result", + "--result-dir", + str(result_dir), + "--result-filename", + result_filename, + ] + if ready_check_timeout_sec is not None: + out.extend(["--ready-check-timeout-sec", str(int(ready_check_timeout_sec))]) + return out + + +def assert_daily_omni_scored(result: dict[str, Any]) -> None: + acc = result.get("daily_omni_accuracy") + assert acc is not None, "daily_omni_accuracy missing — wrong dataset or benchmark wiring" + assert int(result.get("daily_omni_evaluated_ok", 0) or 0) > 0, "no successful MCQ rows (daily_omni_evaluated_ok==0)" + + +def assert_seed_tts_scored(result: dict[str, Any]) -> None: + err = result.get("seed_tts_eval_setup_error") + assert not err, f"Seed-TTS eval deps/setup failed: {err}" + assert int(result.get("seed_tts_content_evaluated", 0) or 0) > 0, ( + "seed_tts_content_evaluated==0 — enable WER eval and check PCM capture / modalities" + ) diff --git a/tests/e2e/accuracy/qwen3_omni/run_qwen_omni_acc_benchmark.py b/tests/e2e/accuracy/qwen3_omni/run_qwen_omni_acc_benchmark.py new file mode 100644 index 000000000..d30457dcd --- /dev/null +++ b/tests/e2e/accuracy/qwen3_omni/run_qwen_omni_acc_benchmark.py @@ -0,0 +1,428 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +"""Accuracy (and light perf) checks for Qwen3-Omni via ``vllm bench serve --omni``. + +The standalone CLI uses small ``--num-prompts`` / ``--max-concurrency`` defaults suitable for +L4-style smoke runs against an already-running server. The pytest wrappers in +``tests/e2e/accuracy/qwen3_omni/test_qwen3_omni.py`` may still require larger GPUs (currently +H100 / MI325) because they launch the live Omni server inside the test. + +1. **Daily-Omni** — MCQ accuracy fields in the saved JSON (``daily_omni_accuracy``, …); by default the + run **fails** if accuracy is strictly below **0.69** (``--min-daily-omni-accuracy`` / ``ACC_BENCH_MIN_DAILY_OMNI_ACCURACY``). +2. **Seed-TTS** — ``seed-tts-eval``-style metrics when ``--seed-tts-wer-eval`` is used + (WER / SIM / UTMOS keys from :func:`compute_seed_tts_wer_metrics`). + +Prerequisites +------------- +* A running Omni OpenAI-compatible server (same machine or reachable host), e.g.:: + + vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8000 + + On L4 you may need a smaller checkpoint, quantization, or tighter engine flags; this script + only drives the **client** benchmark. + +* ``vllm`` CLI from **vLLM-Omni** (so ``bench serve`` registers ``daily-omni`` / ``seed-tts``). + +* **Daily-Omni** — if local ``qa.json`` + ``Videos/`` are not both provided (CLI or matching env), + the client passes ``--dataset-path`` with a Hub id (default ``liarliar/Daily-Omni``). The **child** + ``vllm bench serve`` process then loads QA via ``datasets.load_dataset`` (needs ``pip install datasets``, + network or HF cache). Without ``--daily-omni-video-dir``, the benchmark **lazily** downloads and + extracts ``Videos.tar`` from the Hub (``huggingface_hub``) on first multimodal request. Override + the dataset repo with ``--daily-omni-repo`` or ``VLLM_DAILY_OMNI_REPO``; override the tar repo + with ``VLLM_DAILY_OMNI_MEDIA_REPO`` if needed. + +* **Seed-TTS** optional extras for WER/SIM/UTMOS:: + + pip install 'vllm-omni[seed-tts-eval]' + +Examples +-------- +Pytest (same checks; needs a running server):: + + pytest -sv tests/e2e/accuracy/qwen3_omni/test_qwen3_omni.py + +Smoke on localhost (server already up):: + + python tests/e2e/accuracy/qwen3_omni/run_qwen_omni_acc_benchmark.py \\ + --model Qwen/Qwen3-Omni-30B-A3B-Instruct \\ + --daily-omni-qa-json ./qa.json \\ + --daily-omni-video-dir ./Videos \\ + --seed-tts-dataset-path ./seed-tts-eval + +Skip one suite, tighten gates:: + + python tests/e2e/accuracy/qwen3_omni/run_qwen_omni_acc_benchmark.py \\ + --skip-daily-omni \\ + --max-seed-tts-mean-wer 0.35 \\ + --min-seed-tts-mean-sim 0.75 +""" + +from __future__ import annotations + +import argparse +import contextlib +import json +import os +import sys +from datetime import datetime +from pathlib import Path +from typing import Any + +from tests.e2e.accuracy.qwen3_omni.qwen3_omni_acc_bench_core import ( + build_serve_common_argv, + daily_omni_bench_argv, + find_vllm_cli, + load_benchmark_result, + run_vllm_bench_subprocess, + seed_tts_bench_argv, +) + +_REPO_ROOT = Path(__file__).resolve().parents[4] + + +def _repo_root() -> Path: + return _REPO_ROOT + + +def _default_result_dir() -> Path: + return Path(__file__).resolve().parent / "results" / "qwen_omni_acc" + + +def _validate_daily_omni(result: dict[str, Any], *, min_accuracy: float | None) -> list[str]: + errs: list[str] = [] + acc = result.get("daily_omni_accuracy") + if acc is None: + errs.append("Missing daily_omni_accuracy (wrong dataset or no gold-evaluated rows).") + return errs + ev = int(result.get("daily_omni_evaluated_ok", 0) or 0) + if ev <= 0: + errs.append("daily_omni_evaluated_ok is 0; no successful MCQ rows to score.") + if min_accuracy is not None and float(acc) + 1e-12 < float(min_accuracy): + errs.append(f"daily_omni_accuracy={acc:.6f} < --min-daily-omni-accuracy={min_accuracy}") + return errs + + +def _validate_seed_tts( + result: dict[str, Any], + *, + max_mean_wer: float | None, + min_mean_sim: float | None, + min_mean_utmos: float | None, +) -> list[str]: + errs: list[str] = [] + setup = result.get("seed_tts_eval_setup_error") + if setup: + errs.append(f"Seed-TTS eval setup failed: {setup}") + return errs + n = int(result.get("seed_tts_content_evaluated", 0) or 0) + if n <= 0: + errs.append("seed_tts_content_evaluated is 0 (enable --seed-tts-wer-eval and check PCM capture).") + mean_wer = result.get("seed_tts_content_error_mean") + if mean_wer is not None and max_mean_wer is not None and float(mean_wer) > float(max_mean_wer) + 1e-12: + errs.append(f"seed_tts_content_error_mean (WER)={mean_wer:.6f} > --max-seed-tts-mean-wer={max_mean_wer}") + sim_m = result.get("seed_tts_sim_mean") + if sim_m is not None and min_mean_sim is not None and float(sim_m) + 1e-12 < float(min_mean_sim): + errs.append(f"seed_tts_sim_mean={sim_m:.6f} < --min-seed-tts-mean-sim={min_mean_sim}") + ut_m = result.get("seed_tts_utmos_mean") + if ut_m is not None and min_mean_utmos is not None and float(ut_m) + 1e-12 < float(min_mean_utmos): + errs.append(f"seed_tts_utmos_mean={ut_m:.6f} < --min-seed-tts-mean-utmos={min_mean_utmos}") + return errs + + +def sync_dataset_env_from_ns(ns: argparse.Namespace) -> None: + """Mirror CLI path flags into env vars read by ``daily_omni_bench_argv`` / ``seed_tts_bench_argv``.""" + repo = getattr(ns, "daily_omni_repo", None) + if repo is not None and str(repo).strip(): + os.environ["VLLM_DAILY_OMNI_REPO"] = str(repo).strip() + if ns.daily_omni_qa_json is not None: + os.environ["VLLM_DAILY_OMNI_QA_JSON"] = str(Path(ns.daily_omni_qa_json).expanduser().resolve()) + if ns.daily_omni_video_dir is not None: + os.environ["VLLM_DAILY_OMNI_VIDEO_DIR"] = str(Path(ns.daily_omni_video_dir).expanduser().resolve()) + if ns.seed_tts_dataset_path is not None: + # ``--seed-tts-dataset-path`` accepts either a local directory or a + # Hugging Face repo id. Only resolve to an absolute filesystem path + # when the value actually exists locally; otherwise preserve the repo + # string verbatim so downstream code can pass it to snapshot_download. + raw = str(ns.seed_tts_dataset_path).strip() + p = Path(raw).expanduser() + os.environ["VLLM_SEED_TTS_DATASET_PATH"] = str(p.resolve()) if p.exists() and p.is_dir() else raw + if ns.seed_tts_root is not None: + os.environ["SEED_TTS_ROOT"] = str(Path(ns.seed_tts_root).expanduser().resolve()) + + +@contextlib.contextmanager +def _preserve_benchmark_dataset_env() -> Any: + """Save/restore dataset-related env vars so benchmark tests don't leak state.""" + keys = ( + "VLLM_DAILY_OMNI_REPO", + "VLLM_DAILY_OMNI_QA_JSON", + "VLLM_DAILY_OMNI_VIDEO_DIR", + "VLLM_SEED_TTS_DATASET_PATH", + "SEED_TTS_ROOT", + ) + original = {k: os.environ.get(k) for k in keys} + try: + yield + finally: + for key, value in original.items(): + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + + +def _build_common_args(ns: argparse.Namespace, *, result_filename: str) -> list[str]: + return build_serve_common_argv( + host=ns.host, + port=ns.port, + model=ns.model, + num_prompts=ns.num_prompts, + max_concurrency=ns.max_concurrency, + num_warmups=ns.num_warmups, + percentile_metrics=ns.percentile_metrics, + result_dir=ns.result_dir, + result_filename=result_filename, + ready_check_timeout_sec=ns.ready_check_timeout_sec, + ) + + +def run_daily_omni(ns: argparse.Namespace, vllm: str) -> Path: + ns.result_dir.mkdir(parents=True, exist_ok=True) + tag = datetime.now().strftime("%Y%m%d-%H%M%S") + result_filename = f"qwen_omni_acc_daily_omni_{tag}.json" + extra = json.loads(ns.daily_extra_body_json) + argv = ( + _build_common_args(ns, result_filename=result_filename) + + daily_omni_bench_argv() + + [ + "--daily-omni-input-mode", + ns.daily_omni_input_mode, + "--extra-body", + json.dumps(extra, ensure_ascii=False, separators=(",", ":")), + ] + ) + if ns.daily_omni_save_eval_items: + argv.append("--daily-omni-save-eval-items") + print("\n$", vllm, *argv, "\n", flush=True) + run_vllm_bench_subprocess(vllm, argv) + out = Path(ns.result_dir) / result_filename + if not out.is_file(): + raise FileNotFoundError(f"Expected result JSON at {out}") + return out + + +def run_seed_tts(ns: argparse.Namespace, vllm: str) -> Path: + ns.result_dir.mkdir(parents=True, exist_ok=True) + tag = datetime.now().strftime("%Y%m%d-%H%M%S") + result_filename = f"qwen_omni_acc_seed_tts_{tag}.json" + extra = json.loads(ns.seed_extra_body_json) + argv = ( + _build_common_args(ns, result_filename=result_filename) + + seed_tts_bench_argv(locale=ns.seed_tts_locale) + + [ + "--seed-tts-wer-eval", + "--extra-body", + json.dumps(extra, ensure_ascii=False, separators=(",", ":")), + ] + ) + if ns.seed_tts_wer_save_items: + argv.append("--seed-tts-wer-save-items") + if ns.seed_tts_file_ref_audio: + argv.append("--seed-tts-file-ref-audio") + extra_env: dict[str, str] = {"SEED_TTS_WER_EVAL": "1"} + if ns.seed_tts_eval_device: + extra_env["SEED_TTS_EVAL_DEVICE"] = ns.seed_tts_eval_device + print("\n$", vllm, *argv, "\n", flush=True) + run_vllm_bench_subprocess(vllm, argv, extra_env=extra_env) + out = Path(ns.result_dir) / result_filename + if not out.is_file(): + raise FileNotFoundError(f"Expected result JSON at {out}") + return out + + +def build_arg_parser() -> argparse.ArgumentParser: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--host", default=os.environ.get("ACC_BENCH_HOST", "127.0.0.1")) + p.add_argument("--port", type=int, default=int(os.environ.get("ACC_BENCH_PORT", "8000"))) + p.add_argument( + "--model", + default=os.environ.get( + "ACC_BENCH_MODEL", + "Qwen/Qwen3-Omni-30B-A3B-Instruct", + ), + help="Model id passed to ``vllm bench serve`` (must match the running server).", + ) + p.add_argument("--num-prompts", type=int, default=int(os.environ.get("ACC_BENCH_NUM_PROMPTS", "2000"))) + p.add_argument("--max-concurrency", type=int, default=int(os.environ.get("ACC_BENCH_MAX_CONCURRENCY", "10"))) + p.add_argument("--num-warmups", type=int, default=int(os.environ.get("ACC_BENCH_NUM_WARMUPS", "0"))) + p.add_argument( + "--percentile-metrics", + default=os.environ.get("ACC_BENCH_PERCENTILE_METRICS", "ttft,tpot,itl,e2el,audio_ttfp,audio_rtf"), + ) + p.add_argument( + "--ready-check-timeout-sec", + type=int, + default=None, + help="If set, forwarded to ``vllm bench serve`` (probe first request until success). " + "Omit to use upstream default (typically skip).", + ) + p.add_argument( + "--result-dir", + type=Path, + default=Path(os.environ.get("ACC_BENCH_RESULT_DIR", str(_default_result_dir()))), + ) + + p.add_argument("--skip-daily-omni", action="store_true") + p.add_argument("--skip-seed-tts", action="store_true") + + p.add_argument( + "--daily-omni-repo", + type=str, + default=None, + help="Hugging Face dataset id for Daily-Omni Hub mode (sets VLLM_DAILY_OMNI_REPO). " + "Ignored when local qa.json + video dir are used.", + ) + p.add_argument( + "--daily-omni-qa-json", + type=Path, + default=None, + help="Optional local qa.json; if omitted with no env, uses Hub liarliar/Daily-Omni.", + ) + p.add_argument( + "--daily-omni-video-dir", + type=Path, + default=None, + help="Optional local Videos root; if omitted, media is fetched lazily from Hub Videos.tar.", + ) + p.add_argument("--daily-omni-input-mode", choices=("all", "visual", "audio"), default="all") + p.add_argument( + "--daily-extra-body-json", + default='{"modalities":["text"]}', + help="JSON merged into each chat request for Daily-Omni (default matches common L4 / text-output runs).", + ) + p.add_argument( + "--daily-omni-save-eval-items", + action="store_true", + help="Sets env via CLI flag so per-item rows are stored in the result JSON.", + ) + p.add_argument( + "--min-daily-omni-accuracy", + type=float, + default=float((os.environ.get("ACC_BENCH_MIN_DAILY_OMNI_ACCURACY") or "0.69").strip() or "0.69"), + help="Fail when daily_omni_accuracy is strictly below this threshold (0–1). " + "Default baseline 0.69; override with env ACC_BENCH_MIN_DAILY_OMNI_ACCURACY or pass 0 to disable the floor.", + ) + + p.add_argument( + "--seed-tts-dataset-path", + type=str, + default=None, + help="Optional local root or Hub id; if omitted, uses zhaochenyang20/seed-tts-eval.", + ) + p.add_argument("--seed-tts-root", type=Path, default=None, help="Optional override for Seed-TTS filesystem root.") + p.add_argument("--seed-tts-locale", choices=("en", "zh"), default="en") + p.add_argument( + "--seed-extra-body-json", + default='{"modalities":["text","audio"]}', + help="JSON for Seed-TTS chat requests (must include audio for synthesis + PCM capture).", + ) + p.add_argument("--seed-tts-wer-save-items", action="store_true") + p.add_argument( + "--seed-tts-file-ref-audio", + action="store_true", + help="Use file:// ref_audio; server must allow local media paths.", + ) + p.add_argument( + "--seed-tts-eval-device", + default=os.environ.get("SEED_TTS_EVAL_DEVICE"), + help="Sets SEED_TTS_EVAL_DEVICE for Whisper / WavLM / UTMOS (e.g. cuda:0).", + ) + p.add_argument( + "--max-seed-tts-mean-wer", + type=float, + default=0.02, + help="If set, fail when seed_tts_content_error_mean is strictly above this value.", + ) + p.add_argument( + "--min-seed-tts-mean-sim", + type=float, + default=None, + help="If set, fail when seed_tts_sim_mean is strictly below this value.", + ) + p.add_argument( + "--min-seed-tts-mean-utmos", + type=float, + default=None, + help="If set, fail when seed_tts_utmos_mean is strictly below this value.", + ) + return p + + +def parse_acc_benchmark_args(argv: list[str] | None = None) -> argparse.Namespace: + """Parse CLI args; when ``argv`` is ``None``, use ``sys.argv[1:]`` (standalone script).""" + if argv is None: + argv = sys.argv[1:] + return build_arg_parser().parse_args(argv) + + +def run_acc_benchmark(ns: argparse.Namespace) -> int: + """Run Daily-Omni and/or Seed-TTS client benches against a running server; return 0 on success.""" + failed: list[str] = [] + + with _preserve_benchmark_dataset_env(): + sync_dataset_env_from_ns(ns) + + vllm = find_vllm_cli() + print(f"Using vLLM CLI: {vllm}", flush=True) + print(f"Repo root (for cwd reference): {_repo_root()}", flush=True) + + if not ns.skip_daily_omni: + path = run_daily_omni(ns, vllm) + print(f"\n[Daily-Omni] result JSON: {path}", flush=True) + data = load_benchmark_result(path) + errs = _validate_daily_omni(data, min_accuracy=ns.min_daily_omni_accuracy) + if errs: + failed.extend([f"[Daily-Omni] {e}" for e in errs]) + else: + print( + f"[Daily-Omni] daily_omni_accuracy={data.get('daily_omni_accuracy')} " + f"evaluated_ok={data.get('daily_omni_evaluated_ok')}", + flush=True, + ) + + if not ns.skip_seed_tts: + path = run_seed_tts(ns, vllm) + print(f"\n[Seed-TTS] result JSON: {path}", flush=True) + data = load_benchmark_result(path) + errs = _validate_seed_tts( + data, + max_mean_wer=ns.max_seed_tts_mean_wer, + min_mean_sim=ns.min_seed_tts_mean_sim, + min_mean_utmos=ns.min_seed_tts_mean_utmos, + ) + if errs: + failed.extend([f"[Seed-TTS] {e}" for e in errs]) + else: + print( + f"[Seed-TTS] mean_wer={data.get('seed_tts_content_error_mean')} " + f"mean_sim={data.get('seed_tts_sim_mean')} mean_utmos={data.get('seed_tts_utmos_mean')} " + f"evaluated={data.get('seed_tts_content_evaluated')}", + flush=True, + ) + + if failed: + print("\nACCURACY CHECK FAILED:", file=sys.stderr) + for line in failed: + print(f" - {line}", file=sys.stderr) + return 1 + + print("\nAll configured accuracy checks passed.", flush=True) + return 0 + + +def main() -> int: + return run_acc_benchmark(parse_acc_benchmark_args()) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/e2e/accuracy/qwen3_omni/test_qwen3_omni.py b/tests/e2e/accuracy/qwen3_omni/test_qwen3_omni.py new file mode 100644 index 000000000..773f7c110 --- /dev/null +++ b/tests/e2e/accuracy/qwen3_omni/test_qwen3_omni.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Qwen3-Omni accuracy benchmarks (Daily-Omni MCQ + Seed-TTS WER) via ``vllm bench serve --omni``. + +Starts a **module-scoped** Omni OpenAI-compatible server (same pattern as ``tests/dfx/perf`` and +``tests/e2e/online_serving/test_qwen3_omni.py``), then runs the client benches against +``omni_server.host`` / ``omni_server.port`` / ``omni_server.model``. + +**Daily-Omni from Hugging Face:** unless ``VLLM_DAILY_OMNI_QA_JSON`` and ``VLLM_DAILY_OMNI_VIDEO_DIR`` +point at a full local tree, the bench uses ``--dataset-path`` (default ``liarliar/Daily-Omni`` via +``VLLM_DAILY_OMNI_REPO`` / ``--daily-omni-repo``). QA loads through ``datasets``; ``Videos.tar`` is +downloaded and extracted under ``HF_HOME`` on demand. The tests patch in +``--daily-omni-inline-local-video`` so multimodal payloads use data URLs (no +``--allowed-local-media-path`` on the server). Use small ``--num-prompts`` defaults suitable for CI +(override with ``ACC_BENCH_NUM_PROMPTS`` / ``ACC_BENCH_MAX_CONCURRENCY``; see +:func:`tests.e2e.accuracy.qwen3_omni.qwen3_omni_acc_bench_core.build_acc_benchmark_cli_argv`). + +This package lives under ``tests/e2e/accuracy/qwen3_omni/``, so pytest still loads +``tests/e2e/accuracy/conftest.py``, which imports ``tests.conftest`` (heavy deps: ``vllm``, ``torch``, …). +A broken or partial install can therefore **fail during collection** before these tests run. + +If ``vllm`` is not on ``PATH``, the tests **skip** instead of erroring. Without +``VLLM_SKIP_ACC_BENCH=1``, a failed bench still yields a **failed** run (non-zero subprocess exit). + +Run:: + + pytest -sv tests/e2e/accuracy/qwen3_omni/test_qwen3_omni.py + +Only the subprocess accuracy marker:: + + pytest -sv tests/e2e/accuracy/qwen3_omni/test_qwen3_omni.py -m qwen3_omni_acc + +Skip when you do not have GPUs, a server, or datasets (CI opt-out):: + + VLLM_SKIP_ACC_BENCH=1 pytest -sv tests/e2e/accuracy/qwen3_omni/test_qwen3_omni.py + +Standalone CLI (expects a server already up; uses ``ACC_BENCH_*`` env defaults):: + + python tests/e2e/accuracy/qwen3_omni/run_qwen_omni_acc_benchmark.py --help +""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +from tests.e2e.accuracy.qwen3_omni import run_qwen_omni_acc_benchmark as _acc_bench +from tests.e2e.accuracy.qwen3_omni.qwen3_omni_acc_bench_core import ( + build_acc_benchmark_cli_argv, + find_vllm_cli, +) +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniServerParams +from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config +from vllm_omni.platforms import current_omni_platform + +_E2E_ROOT = Path(__file__).resolve().parent.parent.parent + +models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"] + +pytestmark = [pytest.mark.full_model, pytest.mark.omni] + +_CI_DEPLOY = get_deploy_config_path("ci/qwen3_omni_moe.yaml") + + +def get_chunk_config(config_path: str | None = None): + """Load the qwen3_omni CI deploy yaml with async_chunk modifications for streaming mode.""" + if config_path is None: + config_path = _CI_DEPLOY + # TODO: remove this workaround once legacy `stage_args` path is deleted. + # The pipeline (qwen3_omni/pipeline.py) already wires + # thinker2talker_async_chunk / talker2code2wav_async_chunk on stage 0/1, + # so only async_chunk needs flipping. Writing nested `engine_args:` into + # the new-schema overlay trips _parse_stage_deploy's legacy branch and + # drops flat fields (load_format, max_num_seqs, ...). + return modify_stage_config(config_path, updates={"async_chunk": True}) + + +if current_omni_platform.is_xpu(): + stage_configs = [_CI_DEPLOY] +else: # CUDA + ROCm MI325 share the same deploy config + stage_configs = [get_chunk_config()] + +test_params = [ + OmniServerParams(model=model, stage_config_path=stage_config) for model in models for stage_config in stage_configs +] + + +def _require_vllm_cli() -> None: + try: + find_vllm_cli() + except FileNotFoundError as exc: + pytest.skip(str(exc)) + + +@pytest.fixture(autouse=True) +def _daily_omni_hub_inline_media(monkeypatch: pytest.MonkeyPatch) -> None: + """Hub / lazy-cache mode uses local files → default ``file://`` needs server allowlist. + + ``run_qwen_omni_acc_benchmark`` binds ``daily_omni_bench_argv`` at import time; patch that copy + so we append ``--daily-omni-inline-local-video`` whenever the core helper did not already set it + (local qa.json + video-dir mode already passes the flag). + """ + orig = _acc_bench.daily_omni_bench_argv + + def _wrapped() -> list[str]: + out = list(orig()) + if "--daily-omni-inline-local-video" not in out: + out.append("--daily-omni-inline-local-video") + return out + + monkeypatch.setattr(_acc_bench, "daily_omni_bench_argv", _wrapped) + monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") + monkeypatch.setenv("VLLM_TEST_CLEAN_GPU_MEMORY", "0") + + +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_qwen3_omni_daily_omni_accuracy_bench(omni_server) -> None: + _require_vllm_cli() + pytest.importorskip("datasets") + pytest.importorskip("huggingface_hub") + ns = _acc_bench.parse_acc_benchmark_args( + build_acc_benchmark_cli_argv(omni_server, skip_seed=True, skip_daily=False) + ) + assert _acc_bench.run_acc_benchmark(ns) == 0 + + +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_qwen3_omni_seed_tts_wer_bench(omni_server) -> None: + _require_vllm_cli() + pytest.importorskip("huggingface_hub") + ns = _acc_bench.parse_acc_benchmark_args( + build_acc_benchmark_cli_argv(omni_server, skip_seed=False, skip_daily=True) + ) + assert _acc_bench.run_acc_benchmark(ns) == 0 diff --git a/tests/e2e/accuracy/test_gebench_h100_smoke.py b/tests/e2e/accuracy/test_gebench_h100_smoke.py index 748919269..2702710e4 100644 --- a/tests/e2e/accuracy/test_gebench_h100_smoke.py +++ b/tests/e2e/accuracy/test_gebench_h100_smoke.py @@ -9,10 +9,10 @@ from benchmarks.accuracy.text_to_image.gbench import main as gbench_main from tests.e2e.accuracy.helpers import infer_model_label, reset_artifact_dir from tests.helpers.mark import hardware_test +pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] + -@pytest.mark.advanced_model @pytest.mark.benchmark -@pytest.mark.diffusion @hardware_test(res={"cuda": "H100"}, num_cards=1) def test_gebench_h100_smoke( gebench_accuracy_servers, diff --git a/tests/e2e/accuracy/test_gedit_bench_h100_smoke.py b/tests/e2e/accuracy/test_gedit_bench_h100_smoke.py index 6227d6368..789f7ec93 100644 --- a/tests/e2e/accuracy/test_gedit_bench_h100_smoke.py +++ b/tests/e2e/accuracy/test_gedit_bench_h100_smoke.py @@ -10,10 +10,10 @@ from benchmarks.accuracy.image_to_image.gedit_bench import main as gedit_main from tests.e2e.accuracy.helpers import infer_model_label, reset_artifact_dir from tests.helpers.mark import hardware_test +pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] + -@pytest.mark.advanced_model @pytest.mark.benchmark -@pytest.mark.diffusion @hardware_test(res={"cuda": "H100"}, num_cards=1) def test_gedit_bench_h100_smoke( gedit_accuracy_servers, diff --git a/tests/e2e/accuracy/test_ltx2_3_video_similarity.py b/tests/e2e/accuracy/test_ltx2_3_video_similarity.py new file mode 100644 index 000000000..dec533d58 --- /dev/null +++ b/tests/e2e/accuracy/test_ltx2_3_video_similarity.py @@ -0,0 +1,410 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +SSIM/PSNR accuracy tests for LTX-2.3. + +1. **Transformer parity** (``test_ltx2_3_transformer_matches_diffusers``): + Swaps our custom transformer into diffusers' ``LTX2Pipeline`` to measure + numerical parity in isolation. Thresholds: SSIM >= 0.95, PSNR >= 28 dB. + Result: SSIM 0.999987 (bit-identical). + +2. **Full pipeline** (``test_ltx2_3_pipeline_matches_diffusers``): + Runs the full vLLM-Omni serving stack (``OmniServer`` -> HTTP API) and + compares per-frame against stock diffusers. Currently skipped because + the OmniServer subprocess creates a different RNG state than in-process + diffusers, producing different initial latents from the same seed. + This is a test infrastructure limitation, not a model accuracy issue. +""" + +from __future__ import annotations + +import gc +import os +import tempfile +from pathlib import Path + +import diffusers +import numpy as np +import pytest +import requests +import torch +from PIL import Image + +from tests.e2e.accuracy.helpers import compute_image_ssim_psnr, model_output_dir +from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniServer + +# Parse diffusers version for compatibility check +_DIFFUSERS_VERSION = tuple(int(x) for x in diffusers.__version__.split(".")[:2] if x.isdigit()) +_DIFFUSERS_038 = _DIFFUSERS_VERSION >= (0, 38) + +MODEL_ID = "dg845/LTX-2.3-Diffusers" +MODEL_ENV_VAR = "VLLM_TEST_LTX23_MODEL" +PROMPT = "A lighthouse on a rocky cliff at sunset, waves crashing below, golden hour lighting" +NEGATIVE_PROMPT = "blurry, low quality, distorted, watermark" +WIDTH = 512 +HEIGHT = 384 +NUM_FRAMES = 25 # ~1 second at 24fps +NUM_INFERENCE_STEPS = 20 +GUIDANCE_SCALE = 4.0 +SEED = 42 + +# Transformer-swap test: near-identical output expected +TRANSFORMER_SSIM_THRESHOLD = 0.95 +TRANSFORMER_PSNR_THRESHOLD = 28.0 + +# Full-pipeline test: allows minor divergence from RNG / pipeline differences +PIPELINE_SSIM_THRESHOLD = 0.94 +PIPELINE_PSNR_THRESHOLD = 28.0 + + +def _model_name() -> str: + return os.environ.get(MODEL_ENV_VAR, MODEL_ID) + + +def _local_files_only(model: str) -> bool: + return Path(model).exists() + + +# --------------------------------------------------------------------------- +# Frame extraction helpers +# --------------------------------------------------------------------------- + + +def _video_to_frames(video_np: np.ndarray) -> list[Image.Image]: + """Convert numpy video to list of PIL Images.""" + while video_np.ndim > 4: + video_np = video_np[0] + if video_np.dtype in (np.float32, np.float64, np.float16): + video_np = np.clip(video_np * 255, 0, 255).astype(np.uint8) + return [Image.fromarray(video_np[t]) for t in range(video_np.shape[0])] + + +def _extract_diffusers_frames(result) -> list[Image.Image]: + """Extract frames from diffusers pipeline output.""" + video = result.frames + if isinstance(video, np.ndarray): + return _video_to_frames(video) + if isinstance(video, list): + if isinstance(video[0], list): + return [img.convert("RGB") for img in video[0]] + if isinstance(video[0], Image.Image): + return [img.convert("RGB") for img in video] + raise ValueError(f"Unexpected output type: {type(video)}") + + +def _extract_mp4_frames(mp4_bytes: bytes) -> list[Image.Image]: + """Extract frames from an MP4 video using ffmpeg.""" + import subprocess + + with tempfile.TemporaryDirectory() as tmpdir: + mp4_path = os.path.join(tmpdir, "video.mp4") + with open(mp4_path, "wb") as f: + f.write(mp4_bytes) + + # Extract video frames as PNG files using ffmpeg + frame_pattern = os.path.join(tmpdir, "frame_%04d.png") + subprocess.run( + ["ffmpeg", "-i", mp4_path, "-vsync", "0", frame_pattern], + capture_output=True, + check=True, + ) + + # Load frames in order + frames = [] + i = 1 + while True: + fpath = os.path.join(tmpdir, f"frame_{i:04d}.png") + if not os.path.exists(fpath): + break + frames.append(Image.open(fpath).convert("RGB").copy()) + i += 1 + return frames + + +# --------------------------------------------------------------------------- +# Comparison helper +# --------------------------------------------------------------------------- + + +def _assert_video_similarity( + *, + model_name: str, + vllm_frames: list[Image.Image], + diffusers_frames: list[Image.Image], + ssim_threshold: float, + psnr_threshold: float, +) -> tuple[float, float]: + """Compare video frames and assert SSIM/PSNR meet thresholds.""" + min_frames = min(len(vllm_frames), len(diffusers_frames)) + assert min_frames > 0, "No frames to compare" + + ssim_scores = [] + psnr_scores = [] + for i in range(min_frames): + ssim_val, psnr_val = compute_image_ssim_psnr( + prediction=vllm_frames[i], + reference=diffusers_frames[i], + ) + ssim_scores.append(ssim_val) + psnr_scores.append(psnr_val) + + avg_ssim = sum(ssim_scores) / len(ssim_scores) + avg_psnr = sum(psnr_scores) / len(psnr_scores) + + print(f"\n{model_name} video similarity ({min_frames} frames):") + print(f" SSIM: avg={avg_ssim:.6f}, min={min(ssim_scores):.6f}, threshold>={ssim_threshold:.6f}") + print(f" PSNR: avg={avg_psnr:.6f} dB, min={min(psnr_scores):.6f} dB, threshold>={psnr_threshold:.6f} dB") + + assert avg_ssim >= ssim_threshold, f"SSIM below threshold: got {avg_ssim:.6f}, expected >= {ssim_threshold:.6f}." + assert avg_psnr >= psnr_threshold, f"PSNR below threshold: got {avg_psnr:.6f}, expected >= {psnr_threshold:.6f}." + return avg_ssim, avg_psnr + + +# --------------------------------------------------------------------------- +# Diffusers baseline (shared by both tests) +# --------------------------------------------------------------------------- + + +def _run_diffusers_baseline(model: str, output_dir: Path) -> list[Image.Image]: + """Generate video using stock diffusers LTX2Pipeline.""" + from diffusers import LTX2Pipeline + + run_pre_test_cleanup(enable_force=True) + pipe = None + try: + pipe = LTX2Pipeline.from_pretrained( + model, torch_dtype=torch.bfloat16, local_files_only=_local_files_only(model) + ).to("cuda") + + generator = torch.Generator(device="cuda").manual_seed(SEED) + result = pipe( + prompt=PROMPT, + negative_prompt=NEGATIVE_PROMPT, + width=WIDTH, + height=HEIGHT, + num_frames=NUM_FRAMES, + num_inference_steps=NUM_INFERENCE_STEPS, + guidance_scale=GUIDANCE_SCALE, + generator=generator, + output_type="np", + ) + frames = _extract_diffusers_frames(result) + for i, f in enumerate(frames): + f.save(output_dir / f"diffusers_frame_{i:04d}.png") + return frames + finally: + del pipe + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + run_post_test_cleanup(enable_force=True) + + +# --------------------------------------------------------------------------- +# Test 1: Transformer-swap parity +# --------------------------------------------------------------------------- + + +def _run_with_custom_transformer(model: str, output_dir: Path) -> list[Image.Image]: + """Run diffusers pipeline with our custom transformer swapped in.""" + from contextlib import nullcontext + + from diffusers import LTX2Pipeline + from vllm.config import VllmConfig, set_current_vllm_config + from vllm.distributed.parallel_state import init_distributed_environment, initialize_model_parallel + + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import create_transformer_from_config, load_transformer_config + + vllm_config = VllmConfig() + ctx = set_current_vllm_config(vllm_config) + ctx.__enter__() + + if not torch.distributed.is_initialized(): + os.environ.setdefault("MASTER_ADDR", "localhost") + os.environ.setdefault("MASTER_PORT", "29503") + os.environ.setdefault("RANK", "0") + os.environ.setdefault("WORLD_SIZE", "1") + init_distributed_environment(world_size=1, rank=0, local_rank=0) + initialize_model_parallel(tensor_model_parallel_size=1) + + local = _local_files_only(model) + pipe = LTX2Pipeline.from_pretrained(model, torch_dtype=torch.bfloat16, local_files_only=local) + + transformer_config = load_transformer_config(model, "transformer", local) + our_transformer = create_transformer_from_config(transformer_config) + + diffusers_state = dict(pipe.transformer.named_parameters()) + + def _weight_iter(): + for name, param in diffusers_state.items(): + yield name, param.data + + our_transformer.load_weights(_weight_iter()) + our_transformer = our_transformer.to(dtype=torch.bfloat16, device="cuda").eval() + + # Compatibility shims for diffusers pipeline + our_transformer.dtype = torch.bfloat16 + if not hasattr(our_transformer, "cache_context"): + our_transformer.cache_context = lambda name: nullcontext() + + del pipe.transformer + pipe.transformer = our_transformer + for name, component in pipe.components.items(): + if name != "transformer" and hasattr(component, "to"): + try: + component.to("cuda") + except Exception: + pass + + generator = torch.Generator(device="cuda").manual_seed(SEED) + result = pipe( + prompt=PROMPT, + negative_prompt=NEGATIVE_PROMPT, + width=WIDTH, + height=HEIGHT, + num_frames=NUM_FRAMES, + num_inference_steps=NUM_INFERENCE_STEPS, + guidance_scale=GUIDANCE_SCALE, + generator=generator, + output_type="np", + ) + frames = _extract_diffusers_frames(result) + for i, f in enumerate(frames): + f.save(output_dir / f"vllm_transformer_frame_{i:04d}.png") + + del pipe, result, our_transformer + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + return frames + + +@pytest.mark.advanced_model +@pytest.mark.benchmark +@pytest.mark.diffusion +@pytest.mark.skipif( + not _DIFFUSERS_038, reason="LTX-2.3 requires diffusers >= 0.38.0 for cross_attn_mod and BWE vocoder" +) +@hardware_test(res={"cuda": "H100"}, num_cards=1) +def test_ltx2_3_transformer_matches_diffusers(accuracy_artifact_root: Path) -> None: + """Transformer-level parity: swap our transformer into diffusers pipeline. + + Isolates transformer numerical accuracy from pipeline-level differences. + Both runs use diffusers' denoising loop, CFG, scheduler, and RNG. + """ + model = _model_name() + output_dir = model_output_dir(accuracy_artifact_root, MODEL_ID) + + diffusers_frames = _run_diffusers_baseline(model=model, output_dir=output_dir) + vllm_frames = _run_with_custom_transformer(model=model, output_dir=output_dir) + + _assert_video_similarity( + model_name=f"{MODEL_ID} (transformer-swap)", + vllm_frames=vllm_frames, + diffusers_frames=diffusers_frames, + ssim_threshold=TRANSFORMER_SSIM_THRESHOLD, + psnr_threshold=TRANSFORMER_PSNR_THRESHOLD, + ) + + +# --------------------------------------------------------------------------- +# Test 2: Full pipeline (OmniServer → HTTP API vs diffusers) +# --------------------------------------------------------------------------- + + +def _run_vllm_omni_serving(model: str, output_dir: Path) -> list[Image.Image]: + """Generate video via the full vLLM-Omni serving stack.""" + server_args = [ + "--model-class-name", + "LTX23Pipeline", + "--stage-init-timeout", + "600", + ] + with OmniServer(model, server_args, use_omni=True) as server: + # Submit generation request + response = requests.post( + f"http://{server.host}:{server.port}/v1/videos", + files={ + "prompt": (None, PROMPT), + "negative_prompt": (None, NEGATIVE_PROMPT), + "model": (None, server.model), + "num_frames": (None, str(NUM_FRAMES)), + "fps": (None, "24"), + "size": (None, f"{WIDTH}x{HEIGHT}"), + "num_inference_steps": (None, str(NUM_INFERENCE_STEPS)), + "guidance_scale": (None, str(GUIDANCE_SCALE)), + "seed": (None, str(SEED)), + }, + timeout=120, + ) + response.raise_for_status() + video_id = response.json()["id"] + + # Poll for completion + import time + + for _ in range(120): + status_resp = requests.get( + f"http://{server.host}:{server.port}/v1/videos/{video_id}", + timeout=30, + ) + status_resp.raise_for_status() + status = status_resp.json()["status"] + if status == "completed": + break + if status in ("error", "failed"): + raise RuntimeError(f"Video generation failed: {status_resp.json()}") + time.sleep(5) + else: + raise TimeoutError(f"Video generation timed out after 600s (id={video_id})") + + # Download video content + content_resp = requests.get( + f"http://{server.host}:{server.port}/v1/videos/{video_id}/content", + timeout=120, + ) + content_resp.raise_for_status() + mp4_bytes = content_resp.content + + # Save MP4 + mp4_path = output_dir / "vllm_omni_pipeline.mp4" + with open(mp4_path, "wb") as f: + f.write(mp4_bytes) + + # Extract frames + frames = _extract_mp4_frames(mp4_bytes) + for i, frame in enumerate(frames): + frame.save(output_dir / f"vllm_pipeline_frame_{i:04d}.png") + return frames + + +@pytest.mark.advanced_model +@pytest.mark.benchmark +@pytest.mark.diffusion +@pytest.mark.skipif( + not _DIFFUSERS_038, reason="LTX-2.3 requires diffusers >= 0.38.0 for cross_attn_mod and BWE vocoder" +) +@hardware_test(res={"cuda": "H100"}, num_cards=1) +def test_ltx2_3_pipeline_matches_diffusers(accuracy_artifact_root: Path) -> None: + """Full-pipeline parity: vLLM-Omni serving stack vs diffusers. + + Runs the complete vLLM-Omni OmniServer (subprocess, HTTP API, video + encoding) and compares per-frame against stock diffusers output. + Follows the Wan2.2 / Qwen Image pattern with seed-based determinism. + """ + model = _model_name() + output_dir = model_output_dir(accuracy_artifact_root, MODEL_ID) + + diffusers_frames = _run_diffusers_baseline(model=model, output_dir=output_dir) + vllm_frames = _run_vllm_omni_serving(model=model, output_dir=output_dir) + + _assert_video_similarity( + model_name=f"{MODEL_ID} (full-pipeline)", + vllm_frames=vllm_frames, + diffusers_frames=diffusers_frames, + ssim_threshold=PIPELINE_SSIM_THRESHOLD, + psnr_threshold=PIPELINE_PSNR_THRESHOLD, + ) diff --git a/tests/e2e/accuracy/test_qwen_image.py b/tests/e2e/accuracy/test_qwen_image.py index 8922d9d10..4b8215d54 100644 --- a/tests/e2e/accuracy/test_qwen_image.py +++ b/tests/e2e/accuracy/test_qwen_image.py @@ -17,6 +17,9 @@ from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup from tests.helpers.mark import hardware_test from tests.helpers.runtime import OmniServer +pytestmark = [pytest.mark.full_model, pytest.mark.diffusion] + + MODEL_ID = "Qwen/Qwen-Image" MODEL_ENV_VAR = "QWEN_IMAGE_MODEL" PROMPT = "A photo of a cat sitting on a laptop keyboard, digital art style." @@ -99,9 +102,7 @@ def _run_diffusers_qwen_image(*, model: str, output_path: Path) -> Image.Image: run_post_test_cleanup(enable_force=True) -@pytest.mark.advanced_model @pytest.mark.benchmark -@pytest.mark.diffusion @hardware_test(res={"cuda": "H100"}, num_cards=1) def test_qwen_image_matches_diffusers(accuracy_artifact_root: Path) -> None: model = _model_name() diff --git a/tests/e2e/accuracy/test_qwen_image_edit.py b/tests/e2e/accuracy/test_qwen_image_edit.py index e17aca6e9..07deecca9 100644 --- a/tests/e2e/accuracy/test_qwen_image_edit.py +++ b/tests/e2e/accuracy/test_qwen_image_edit.py @@ -15,6 +15,9 @@ from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup from tests.helpers.mark import hardware_test from tests.helpers.runtime import OmniServer +pytestmark = [pytest.mark.full_model, pytest.mark.diffusion] + + SINGLE_MODEL = "Qwen/Qwen-Image-Edit" MULTIPLE_MODEL = "Qwen/Qwen-Image-Edit-2509" WIDTH = 512 @@ -169,9 +172,7 @@ def _diffusers_output_multiple_image( ) -@pytest.mark.advanced_model @pytest.mark.benchmark -@pytest.mark.diffusion @hardware_test(res={"cuda": "H100"}, num_cards=1) def test_qwen_image_edit_single_matches_diffusers( accuracy_artifact_root: Path, @@ -196,9 +197,7 @@ def test_qwen_image_edit_single_matches_diffusers( ) -@pytest.mark.advanced_model @pytest.mark.benchmark -@pytest.mark.diffusion @hardware_test(res={"cuda": "H100"}, num_cards=1) @pytest.mark.skip( reason="Skipping as the second image seems to be ignored by the API. Will come back to this later after #2772 is merged." diff --git a/tests/e2e/accuracy/test_qwen_image_layered.py b/tests/e2e/accuracy/test_qwen_image_layered.py index 0ab9cb323..30ad2966f 100644 --- a/tests/e2e/accuracy/test_qwen_image_layered.py +++ b/tests/e2e/accuracy/test_qwen_image_layered.py @@ -17,6 +17,9 @@ from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup from tests.helpers.mark import hardware_test from tests.helpers.runtime import OmniServer +pytestmark = [pytest.mark.full_model, pytest.mark.diffusion] + + MODEL_ID = "Qwen/Qwen-Image-Layered" MODEL_ENV_VAR = "QWEN_IMAGE_LAYERED_MODEL" PROMPT = "decompose into layers" @@ -126,9 +129,7 @@ def _run_diffusers_qwen_image_layered(*, model: str, input_image: Image.Image, o run_post_test_cleanup(enable_force=True) -@pytest.mark.advanced_model @pytest.mark.benchmark -@pytest.mark.diffusion @hardware_test(res={"cuda": "H100"}, num_cards=1) def test_qwen_image_layered_matches_diffusers(accuracy_artifact_root: Path, qwen_bear_image: Image.Image) -> None: model = _model_name() diff --git a/tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py b/tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py index 3aa5da85c..1caef3bff 100644 --- a/tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py +++ b/tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py @@ -50,6 +50,8 @@ from tests.e2e.accuracy.wan22_i2v.wan22_i2v_video_similarity_common import ( from tests.helpers.mark import hardware_test from tests.helpers.runtime import OmniServerParams +pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] + def test_parse_video_metadata_extracts_dimensions_and_fps() -> None: payload = { @@ -537,9 +539,7 @@ def _generate_offline_video(*, image_source: str) -> tuple[Path, Path]: return offline_path, offline_metadata_path -@pytest.mark.advanced_model @pytest.mark.benchmark -@pytest.mark.diffusion @hardware_test(res={"cuda": "H100"}, num_cards=1) def test_wan22_i2v_diffusers_offline_generates_video( wan22_i2v_image_source: str | None, @@ -563,9 +563,7 @@ def test_wan22_i2v_diffusers_offline_generates_video( assert offline_metadata["frame_count"] == NUM_FRAMES -@pytest.mark.advanced_model @pytest.mark.benchmark -@pytest.mark.diffusion @hardware_test(res={"cuda": "H100"}, num_cards=2) @pytest.mark.parametrize("omni_server", SERVER_CASES, indirect=True) def test_wan22_i2v_online_serving_generates_video( @@ -594,9 +592,7 @@ def test_wan22_i2v_online_serving_generates_video( assert online_metadata["frame_count"] == NUM_FRAMES -@pytest.mark.advanced_model @pytest.mark.benchmark -@pytest.mark.diffusion @hardware_test(res={"cuda": "H100"}, num_cards=2) def test_wan22_i2v_serving_matches_diffusers_video_similarity( wan22_i2v_image_source: str | None, diff --git a/tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml b/tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml deleted file mode 100644 index b7768c071..000000000 --- a/tests/e2e/offline_inference/stage_configs/bagel_mooncake_ci.yaml +++ /dev/null @@ -1,79 +0,0 @@ -# stage config for running BAGEL with Mooncake connector for CI e2e tests. -# This config is optimized for single GPU tests with Mooncake inter-stage communication. - -stage_args: - - stage_id: 0 - stage_type: llm - prompt_expand_func: vllm_omni.model_executor.stage_input_processors.bagel.expand_cfg_prompts - runtime: - devices: "0" - engine_args: - model_stage: thinker - max_num_seqs: 1 - model_arch: BagelForConditionalGeneration - worker_type: ar - scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - gpu_memory_utilization: 0.45 - enforce_eager: true - trust_remote_code: true - engine_output_type: text - distributed_executor_backend: mp - enable_prefix_caching: false - max_num_batched_tokens: 32768 - tensor_parallel_size: 1 - load_format: dummy - omni_kv_config: - need_send_cache: true - kv_transfer_criteria: - type: prefill_finished - final_output: true - final_output_type: text - is_comprehension: true - default_sampling_params: - temperature: 0.4 - top_p: 0.9 - top_k: 1 - max_tokens: 2048 - seed: 52 - detokenize: true - repetition_penalty: 1.05 - output_connectors: - to_stage_1: mooncake_connector - - stage_id: 1 - stage_type: diffusion - cfg_kv_collect_func: vllm_omni.model_executor.stage_input_processors.bagel.collect_cfg_kv_caches - runtime: - devices: "0" - engine_args: - model_stage: dit - max_num_seqs: 1 - enforce_eager: true - trust_remote_code: true - distributed_executor_backend: mp - omni_kv_config: - need_recv_cache: true - engine_input_source: [0] - final_output: true - final_output_type: image - is_comprehension: false - default_sampling_params: - seed: 52 - input_connectors: - from_stage_0: mooncake_connector - -# Top-level runtime config with Mooncake connector -runtime: - enabled: true - connectors: - mooncake_connector: - name: MooncakeConnector - extra: - host: "${MOONCAKE_HOST}" - metadata_server: "http://${MOONCAKE_HOST}:${MOONCAKE_HTTP_PORT}/metadata" - master: "${MOONCAKE_HOST}:${MOONCAKE_RPC_PORT}" - segment: 64000000 - localbuf: 64000000 - proto: tcp - edges: - - from: 0 - to: 1 diff --git a/tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml b/tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml deleted file mode 100644 index 504f3c98e..000000000 --- a/tests/e2e/offline_inference/stage_configs/bagel_sharedmemory_ci.yaml +++ /dev/null @@ -1,76 +0,0 @@ -# stage config for running BAGEL with SharedMemory connector for CI e2e tests. -# This config is optimized for single GPU tests with SharedMemory inter-stage communication. - -stage_args: - - stage_id: 0 - stage_type: llm - prompt_expand_func: vllm_omni.model_executor.stage_input_processors.bagel.expand_cfg_prompts - runtime: - devices: "0" - engine_args: - model_stage: thinker - max_num_seqs: 1 - model_arch: OmniBagelForConditionalGeneration - worker_type: ar - scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - gpu_memory_utilization: 0.45 - enforce_eager: true - trust_remote_code: true - engine_output_type: text - distributed_executor_backend: "mp" - enable_prefix_caching: false - max_num_batched_tokens: 32768 - tensor_parallel_size: 1 - load_format: dummy - omni_kv_config: - need_send_cache: true - kv_transfer_criteria: - type: prefill_finished #or special token generated - final_output: true - final_output_type: text - is_comprehension: true - default_sampling_params: - temperature: 0.4 - top_p: 0.9 - top_k: 1 - max_tokens: 2048 - seed: 52 - detokenize: True - repetition_penalty: 1.05 - - - stage_id: 1 - stage_type: diffusion - cfg_kv_collect_func: vllm_omni.model_executor.stage_input_processors.bagel.collect_cfg_kv_caches - runtime: - devices: "0" - engine_args: - model_stage: dit - max_num_seqs: 1 - enforce_eager: true - trust_remote_code: true - distributed_executor_backend: "mp" - omni_kv_config: - need_recv_cache: true - engine_input_source: [0] - - final_output: true - final_output_type: image - is_comprehension: false - default_sampling_params: - seed: 52 - -# Runtime edges -runtime: - enabled: true - # Distributed connectors configuration (optional) - # More connectors will be supported in the future. - connectors: - shared_memory_connector: - name: SharedMemoryConnector - extra: - shm_threshold_bytes: 65536 # 64KB threshold - - - edges: - - from: 0 - to: 1 diff --git a/tests/e2e/offline_inference/test_bagel_img2img.py b/tests/e2e/offline_inference/test_bagel_img2img.py index 66aec80c7..b4de059f2 100644 --- a/tests/e2e/offline_inference/test_bagel_img2img.py +++ b/tests/e2e/offline_inference/test_bagel_img2img.py @@ -15,7 +15,6 @@ Equivalent to running: """ import socket -from pathlib import Path from typing import Any import pytest @@ -24,10 +23,12 @@ from vllm.assets.image import ImageAsset from tests.helpers.mark import hardware_test from tests.helpers.runtime import OmniRunner -from tests.helpers.stage_config import modify_stage_config +from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config from vllm_omni.entrypoints.omni import Omni from vllm_omni.platforms import current_omni_platform +BAGEL_CI_DEPLOY = get_deploy_config_path("ci/bagel.yaml") + # Reference pixel data extracted from the known-good output image # Generated with seed=52, num_inference_steps=15, # prompt='Change the grass color to red', @@ -183,8 +184,8 @@ def _generate_bagel_img2img( return generated_image -def _resolve_stage_config(config_path: str, run_level: str) -> str: - """Resolve stage config based on run level. +def _resolve_deploy_config(config_path: str, run_level: str) -> str: + """Resolve deploy config based on run level. For advanced_model (real weights), strip load_format: dummy so the model falls back to loading real weights from HuggingFace. @@ -193,9 +194,9 @@ def _resolve_stage_config(config_path: str, run_level: str) -> str: return modify_stage_config( config_path, deletes={ - "stage_args": { - 0: ["engine_args.load_format"], - 1: ["engine_args.load_format"], + "stages": { + 0: ["load_format"], + 1: ["load_format"], } }, ) @@ -209,8 +210,7 @@ def _resolve_stage_config(config_path: str, run_level: str) -> str: def test_bagel_img2img_shared_memory_connector(run_level): """Test Bagel img2img with shared memory connector.""" input_image = _load_input_image() - config_path = str(Path(__file__).parent / "stage_configs" / "bagel_sharedmemory_ci.yaml") - config_path = _resolve_stage_config(config_path, run_level) + config_path = _resolve_deploy_config(BAGEL_CI_DEPLOY, run_level) with OmniRunner( "ByteDance-Seed/BAGEL-7B-MoT", stage_configs_path=config_path, diff --git a/tests/e2e/offline_inference/test_bagel_lora.py b/tests/e2e/offline_inference/test_bagel_lora.py index 75f41f9be..785d0c7fb 100644 --- a/tests/e2e/offline_inference/test_bagel_lora.py +++ b/tests/e2e/offline_inference/test_bagel_lora.py @@ -17,14 +17,13 @@ Assertions: import json import os +from pathlib import Path from vllm_omni.inputs.data import OmniSamplingParams from vllm_omni.outputs import OmniRequestOutput os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" -from pathlib import Path - import numpy as np import pytest import torch @@ -33,13 +32,13 @@ from safetensors.torch import save_file from tests.helpers.mark import hardware_test from tests.helpers.runtime import OmniRunner -from tests.helpers.stage_config import modify_stage_config +from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config from vllm_omni.entrypoints.omni import Omni from vllm_omni.lora.request import LoRARequest from vllm_omni.lora.utils import stable_lora_int_id MODEL = "ByteDance-Seed/BAGEL-7B-MoT" -BAGEL_STAGE_CONFIG = str(Path(__file__).parent / "stage_configs" / "bagel_sharedmemory_ci.yaml") +BAGEL_STAGE_CONFIG = get_deploy_config_path("ci/bagel.yaml") DEFAULT_PROMPT = "<|im_start|>A cute cat<|im_end|>" @@ -48,14 +47,14 @@ DEFAULT_PROMPT = "<|im_start|>A cute cat<|im_end|>" # --------------------------------------------------------------------------- -def _resolve_stage_config(config_path: str, run_level: str) -> str: +def _resolve_deploy_config(config_path: str, run_level: str) -> str: if run_level == "advanced_model": return modify_stage_config( config_path, deletes={ - "stage_args": { - 0: ["engine_args.load_format"], - 1: ["engine_args.load_format"], + "stages": { + 0: ["load_format"], + 1: ["load_format"], } }, ) @@ -153,7 +152,7 @@ def _make_file_lora_request(adapter_dir: Path) -> LoRARequest: @hardware_test(res={"cuda": "H100", "rocm": "MI325"}) def test_bagel_lora_scale_and_deactivation(run_level, tmp_path): """Validate LoRA effect, bounded perturbation, and clean deactivation.""" - config_path = _resolve_stage_config(BAGEL_STAGE_CONFIG, run_level) + config_path = _resolve_deploy_config(BAGEL_STAGE_CONFIG, run_level) with OmniRunner(MODEL, stage_configs_path=config_path) as runner: omni = runner.omni lora_request = _make_file_lora_request(tmp_path / "bagel_lora") @@ -188,9 +187,9 @@ def test_bagel_lora_scale_and_deactivation(run_level, tmp_path): f"LoRA scale has no effect: diff_1x={diff_1x:.2f}, diff_2x={diff_2x:.2f}" ) - # (c) Output is not corrupted + # (c) Output is not corrupted (scale=2.0 can produce ~2x the diff of scale=1.0) assert diff_1x < 80, f"LoRA output looks corrupted: diff_1x={diff_1x}" - assert diff_2x < 80, f"LoRA output looks corrupted: diff_2x={diff_2x}" + assert diff_2x < 120, f"LoRA output looks corrupted: diff_2x={diff_2x}" # (d) Deactivation fully restores base model assert diff_restored == 0.0, f"Base model not restored after LoRA deactivation: diff={diff_restored}" diff --git a/tests/e2e/offline_inference/test_bagel_text2img.py b/tests/e2e/offline_inference/test_bagel_text2img.py index 0819f103a..65cd8425c 100644 --- a/tests/e2e/offline_inference/test_bagel_text2img.py +++ b/tests/e2e/offline_inference/test_bagel_text2img.py @@ -21,7 +21,6 @@ import socket import subprocess import tempfile import time -from pathlib import Path from typing import Any import pytest @@ -29,10 +28,13 @@ from PIL import Image from tests.helpers.mark import hardware_test from tests.helpers.runtime import OmniRunner -from tests.helpers.stage_config import modify_stage_config +from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config from vllm_omni.entrypoints.omni import Omni from vllm_omni.platforms import current_omni_platform +BAGEL_CI_DEPLOY = get_deploy_config_path("ci/bagel.yaml") +BAGEL_MOONCAKE_CI_DEPLOY = get_deploy_config_path("ci/bagel_mooncake.yaml") + # Reference pixel data extracted from the known-good output image # Each entry contains (x, y) position and expected (R, G, B) values # "Generated with seed=52, num_inference_steps=15, @@ -172,8 +174,8 @@ def _generate_bagel_image(omni: Omni, prompt: str = DEFAULT_PROMPT) -> Image.Ima return generated_image -def _resolve_stage_config(config_path: str, run_level: str) -> str: - """Resolve stage config based on run level. +def _resolve_deploy_config(config_path: str, run_level: str) -> str: + """Resolve deploy config based on run level. For advanced_model (real weights), strip load_format: dummy so the model falls back to loading real weights from HuggingFace. @@ -182,9 +184,9 @@ def _resolve_stage_config(config_path: str, run_level: str) -> str: return modify_stage_config( config_path, deletes={ - "stage_args": { - 0: ["engine_args.load_format"], - 1: ["engine_args.load_format"], + "stages": { + 0: ["load_format"], + 1: ["load_format"], } }, ) @@ -197,8 +199,7 @@ def _resolve_stage_config(config_path: str, run_level: str) -> str: @hardware_test(res={"cuda": "H100", "rocm": "MI325"}) def test_bagel_text2img_shared_memory_connector(run_level): """Test Bagel text2img with shared memory connector.""" - config_path = str(Path(__file__).parent / "stage_configs" / "bagel_sharedmemory_ci.yaml") - config_path = _resolve_stage_config(config_path, run_level) + config_path = _resolve_deploy_config(BAGEL_CI_DEPLOY, run_level) with OmniRunner( "ByteDance-Seed/BAGEL-7B-MoT", stage_configs_path=config_path, @@ -277,7 +278,7 @@ def _cleanup_mooncake_processes(timeout_secs: int = 5) -> None: def _load_mooncake_config(host: str, rpc_port: int, http_port: int) -> str: - """Load Mooncake config from YAML and substitute placeholders. + """Load Mooncake config from CI overlay and substitute placeholders. Args: host: Mooncake host address. @@ -287,16 +288,13 @@ def _load_mooncake_config(host: str, rpc_port: int, http_port: int) -> str: Returns: Path to the temporary config file with substituted values. """ - config_path = str(Path(__file__).parent / "stage_configs" / "bagel_mooncake_ci.yaml") - with open(config_path) as f: + with open(BAGEL_MOONCAKE_CI_DEPLOY) as f: config_content = f.read() - # Substitute placeholders config_content = config_content.replace("${MOONCAKE_HOST}", host) config_content = config_content.replace("${MOONCAKE_RPC_PORT}", str(rpc_port)) config_content = config_content.replace("${MOONCAKE_HTTP_PORT}", str(http_port)) - # Write to temp file temp_file = tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) temp_file.write(config_content) temp_file.close() @@ -346,7 +344,7 @@ def test_bagel_text2img_mooncake_connector(run_level): http_port=MOONCAKE_HTTP_PORT, ) - temp_config_file = _resolve_stage_config(temp_config_file, run_level) + temp_config_file = _resolve_deploy_config(temp_config_file, run_level) with OmniRunner( "ByteDance-Seed/BAGEL-7B-MoT", stage_configs_path=temp_config_file, diff --git a/tests/e2e/offline_inference/test_bagel_understanding.py b/tests/e2e/offline_inference/test_bagel_understanding.py index c3ed97b42..e342152fc 100644 --- a/tests/e2e/offline_inference/test_bagel_understanding.py +++ b/tests/e2e/offline_inference/test_bagel_understanding.py @@ -21,17 +21,16 @@ Equivalent to running: import os os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" -from pathlib import Path import pytest from vllm.assets.image import ImageAsset from tests.helpers.mark import hardware_test from tests.helpers.runtime import OmniRunner -from tests.helpers.stage_config import modify_stage_config +from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config MODEL_NAME = "ByteDance-Seed/BAGEL-7B-MoT" -STAGE_CONFIG = str(Path(__file__).parent / "stage_configs" / "bagel_sharedmemory_ci.yaml") +STAGE_CONFIG = get_deploy_config_path("ci/bagel.yaml") REFERENCE_TEXT_TEXT2TEXT = "The capital of France is Paris." @@ -44,15 +43,15 @@ REFERENCE_TEXT_IMG2TEXT = ( ) -def _resolve_stage_config(config_path: str, run_level: str) -> str: +def _resolve_deploy_config(config_path: str, run_level: str) -> str: """Strip load_format: dummy for advanced_model (real weights).""" if run_level == "advanced_model": return modify_stage_config( config_path, deletes={ - "stage_args": { - 0: ["engine_args.load_format"], - 1: ["engine_args.load_format"], + "stages": { + 0: ["load_format"], + 1: ["load_format"], } }, ) @@ -74,7 +73,7 @@ def _extract_text(omni_outputs: list) -> str: @hardware_test(res={"cuda": "H100", "rocm": "MI325"}) def test_bagel_text2text(run_level): """Test Bagel text2text produces correct text output.""" - config_path = _resolve_stage_config(STAGE_CONFIG, run_level) + config_path = _resolve_deploy_config(STAGE_CONFIG, run_level) with OmniRunner( MODEL_NAME, stage_configs_path=config_path, @@ -106,7 +105,7 @@ def test_bagel_text2text(run_level): def test_bagel_img2text(run_level): """Test Bagel img2text produces correct text output.""" input_image = ImageAsset("2560px-Gfp-wisconsin-madison-the-nature-boardwalk").pil_image.convert("RGB") - config_path = _resolve_stage_config(STAGE_CONFIG, run_level) + config_path = _resolve_deploy_config(STAGE_CONFIG, run_level) with OmniRunner( MODEL_NAME, stage_configs_path=config_path, diff --git a/tests/e2e/offline_inference/test_flux2_klein_inpaint.py b/tests/e2e/offline_inference/test_flux2_klein.py similarity index 94% rename from tests/e2e/offline_inference/test_flux2_klein_inpaint.py rename to tests/e2e/offline_inference/test_flux2_klein.py index ac7169361..a13767534 100644 --- a/tests/e2e/offline_inference/test_flux2_klein_inpaint.py +++ b/tests/e2e/offline_inference/test_flux2_klein.py @@ -16,6 +16,7 @@ import pytest import torch from PIL import Image, ImageDraw +from vllm_omni.entrypoints.omni import Omni from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput from vllm_omni.platforms import current_omni_platform @@ -24,8 +25,6 @@ REPO_ROOT = Path(__file__).resolve().parents[2] if str(REPO_ROOT) not in sys.path: sys.path.insert(0, str(REPO_ROOT)) -from vllm_omni import Omni - os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1" MODEL = "black-forest-labs/FLUX.2-klein-4B" @@ -66,6 +65,18 @@ def _extract_images_from_output(outputs: list) -> list[Image.Image]: return images +# Regression test for https://github.com/vllm-project/vllm-omni/issues/3097 +@pytest.mark.core_model +@pytest.mark.diffusion +def test_flux2_klein_can_accept_text_inputs(): + model = Omni(model=MODEL) + outputs = model.generate( + "a cup of coffee on the table", + OmniDiffusionSamplingParams(num_inference_steps=2, seed=42), + ) + assert len(outputs[0].images) == 1 + + @pytest.mark.core_model @pytest.mark.diffusion def test_flux2_klein_inpaint_basic(): diff --git a/tests/e2e/offline_inference/test_ming_flash_omni.py b/tests/e2e/offline_inference/test_ming_flash_omni.py index c591e910a..ca0b0fe0d 100644 --- a/tests/e2e/offline_inference/test_ming_flash_omni.py +++ b/tests/e2e/offline_inference/test_ming_flash_omni.py @@ -36,6 +36,21 @@ def build_prompt(user_text: str) -> str: def get_eager_config(): + path = modify_stage_config( + str(Path(__file__).parent.parent / "stage_configs" / "bailingmm_moe_v2_lite_thinker_only_ci.yaml"), + updates={ + "stage_args": { + 0: { + "engine_args.enforce_eager": "true", + }, + }, + }, + ) + return path + + +def get_eager_tts_config(): + """Thinker+talker CI config with enforce_eager on the thinker stage.""" path = modify_stage_config( str(Path(__file__).parent.parent / "stage_configs" / "bailingmm_moe_v2_lite_ci.yaml"), updates={ @@ -49,9 +64,14 @@ def get_eager_config(): return path +# Thinker-only config — used by text-output tests. stage_configs = [get_eager_config()] test_params = [(model, stage_config) for model in models for stage_config in stage_configs] +# Thinker+talker config — used by audio-output tests. +stage_configs_tts = [get_eager_tts_config()] +test_params_tts = [(model, stage_config) for model in models for stage_config in stage_configs_tts] + @pytest.mark.core_model @pytest.mark.omni @@ -140,3 +160,36 @@ def test_mixed_to_text(omni_runner, omni_runner_handler) -> None: request_config = {"prompts": prompt, "images": image, "audios": audio, "modalities": ["text"]} omni_runner_handler.send_request(request_config) + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100"}, num_cards=4) +@pytest.mark.parametrize("omni_runner", test_params_tts, indirect=True) +def test_text_to_audio(omni_runner, omni_runner_handler) -> None: + """ + Test text input with audio output via the thinker+talker pipeline. + Input Modal: text + Output Modal: audio + """ + prompt = build_prompt("请简单介绍一下北京。") + request_config = {"prompts": prompt, "modalities": ["audio"]} + + omni_runner_handler.send_request(request_config) + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100"}, num_cards=4) +@pytest.mark.parametrize("omni_runner", test_params_tts, indirect=True) +def test_image_to_audio(omni_runner, omni_runner_handler) -> None: + """ + Test image + text input with audio output via the thinker+talker pipeline. + Input Modal: image + text + Output Modal: audio + """ + image = generate_synthetic_image(224, 224)["np_array"] + prompt = build_prompt(f"{IMAGE_TOKEN}Describe this image briefly.") + request_config = {"prompts": prompt, "images": image, "modalities": ["audio"]} + + omni_runner_handler.send_request(request_config) diff --git a/tests/e2e/offline_inference/test_omni_sleep_mode.py b/tests/e2e/offline_inference/test_omni_sleep_mode.py index 2ad0b53b0..5a3ae9ab7 100644 --- a/tests/e2e/offline_inference/test_omni_sleep_mode.py +++ b/tests/e2e/offline_inference/test_omni_sleep_mode.py @@ -38,8 +38,8 @@ def get_dynamic_devices(stage_idx, num_stages, tp_size): # Test 1: Diffusion Model (2-Stage BAGEL) @pytest.mark.advanced_model @pytest.mark.omni -@pytest.mark.parametrize("tp_size", [1, 2]) -@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) +@pytest.mark.parametrize("tp_size", [1]) +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=1) @pytest.mark.asyncio async def test_diffusion_model_sleep_tp(tp_size): num_gpus = torch.cuda.device_count() @@ -123,7 +123,7 @@ async def test_multistage_sleep_h100(tp_size): pass acks = await engine.sleep(stage_ids=[0, 1], level=2) - assert len(acks) == 2 * tp_size + assert len(acks) == 2 await engine.wake_up(stage_ids=[0, 1]) async for _ in engine.generate("verify", sampling_params=[SamplingParams(), sp]): diff --git a/tests/e2e/offline_inference/test_quantization_fp8.py b/tests/e2e/offline_inference/test_quantization_fp8.py index 92cf351e3..9801e0ae7 100644 --- a/tests/e2e/offline_inference/test_quantization_fp8.py +++ b/tests/e2e/offline_inference/test_quantization_fp8.py @@ -30,7 +30,6 @@ import os os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" -from pathlib import Path from typing import Any import pytest @@ -38,12 +37,15 @@ import torch from tests.helpers.mark import hardware_test from tests.helpers.runtime import OmniRunner +from tests.helpers.stage_config import get_deploy_config_path from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput from vllm_omni.platforms import current_omni_platform pytestmark = [pytest.mark.core_model, pytest.mark.diffusion] +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + # ─── helpers ────────────────────────────────────────────────────────────────── @@ -102,7 +104,7 @@ def _generate_bagel_image( Returns (generated_image, peak_memory_gib). """ - config_path = str(Path(__file__).parent / "stage_configs" / "bagel_sharedmemory_ci.yaml") + config_path = get_deploy_config_path("ci/bagel.yaml") omni_kwargs: dict[str, Any] = { "model": "ByteDance-Seed/BAGEL-7B-MoT", "stage_configs_path": config_path, diff --git a/tests/e2e/offline_inference/test_qwen2_5_omni_autoround_w4a16.py b/tests/e2e/offline_inference/test_qwen2_5_omni_autoround_w4a16.py new file mode 100644 index 000000000..e59a8c7d7 --- /dev/null +++ b/tests/e2e/offline_inference/test_qwen2_5_omni_autoround_w4a16.py @@ -0,0 +1,170 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""E2E tests for Qwen2.5-Omni AutoRound W4A16 quantized inference. + +These tests cover text, audio, image, video, and mixed-modality inputs +to verify multimodal understanding with quantized weights. + +Requirements: + - CUDA GPUs (4x L4 / 24 GB or equivalent) + - The quantized model checkpoint (Intel/Qwen2.5-Omni-7B-int4-AutoRound) +""" + +import os + +import pytest + +from tests.helpers.mark import hardware_test +from tests.helpers.media import ( + generate_synthetic_audio, + generate_synthetic_image, + generate_synthetic_video, +) +from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config + +QUANTIZED_MODEL = "Intel/Qwen2.5-Omni-7B-int4-AutoRound" +BASELINE_MODEL = "Qwen/Qwen2.5-Omni-7B" + +# Allow overriding via environment for local testing +QUANTIZED_MODEL = os.environ.get("QWEN2_5_OMNI_AUTOROUND_MODEL", QUANTIZED_MODEL) +BASELINE_MODEL = os.environ.get("QWEN2_5_OMNI_BASELINE_MODEL", BASELINE_MODEL) + +_CI_DEPLOY = get_deploy_config_path("ci/qwen2_5_omni.yaml") + + +def _get_stage_config(): + """Build a CI-friendly stage config with eager mode.""" + return modify_stage_config( + _CI_DEPLOY, + updates={ + "stages": { + 0: {"enforce_eager": True}, + 1: {"enforce_eager": True}, + }, + }, + ) + + +stage_config = _get_stage_config() + +# Parametrise: (model, stage_config) +quant_params = [(QUANTIZED_MODEL, stage_config)] + + +# ------------------------------------------------------------------ +# Test: text-only input → text output +# ------------------------------------------------------------------ + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "L4"}, num_cards=4) +@pytest.mark.parametrize("omni_runner", quant_params, indirect=True) +def test_text_to_text(omni_runner, omni_runner_handler): + """Text input → text output with W4A16 quantized Qwen2.5-Omni.""" + request_config = { + "prompts": "What is the capital of China?", + "modalities": ["text"], + } + response = omni_runner_handler.send_request(request_config) + assert response.success, f"Request failed: {response.error_message}" + assert response.text_content and len(response.text_content.strip()) > 0 + + +# ------------------------------------------------------------------ +# Test: audio input → text output +# ------------------------------------------------------------------ + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "L4"}, num_cards=4) +@pytest.mark.parametrize("omni_runner", quant_params, indirect=True) +def test_audio_to_text(omni_runner, omni_runner_handler): + """Audio input → text output with W4A16 quantized Qwen2.5-Omni.""" + audio = generate_synthetic_audio(1, 1, 16000)["np_array"] + if len(audio.shape) == 2: + audio = audio.squeeze() + + request_config = { + "prompts": "What is the content of this audio?", + "audios": (audio, 16000), + "modalities": ["text"], + } + response = omni_runner_handler.send_request(request_config) + assert response.success, f"Request failed: {response.error_message}" + assert response.text_content and len(response.text_content.strip()) > 0 + + +# ------------------------------------------------------------------ +# Test: image input → text output +# ------------------------------------------------------------------ + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "L4"}, num_cards=4) +@pytest.mark.parametrize("omni_runner", quant_params, indirect=True) +def test_image_to_text(omni_runner, omni_runner_handler): + """Image input → text output with W4A16 quantized Qwen2.5-Omni.""" + image = generate_synthetic_image(16, 16)["np_array"] + + request_config = { + "prompts": "Describe what you see in this image.", + "images": image, + "modalities": ["text"], + } + response = omni_runner_handler.send_request(request_config) + assert response.success, f"Request failed: {response.error_message}" + assert response.text_content and len(response.text_content.strip()) > 0 + + +# ------------------------------------------------------------------ +# Test: video input → text output +# ------------------------------------------------------------------ + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "L4"}, num_cards=4) +@pytest.mark.parametrize("omni_runner", quant_params, indirect=True) +def test_video_to_text(omni_runner, omni_runner_handler): + """Video input → text output with W4A16 quantized Qwen2.5-Omni.""" + video = generate_synthetic_video(16, 16, 30)["np_array"] + + request_config = { + "prompts": "Describe the video briefly.", + "videos": video, + "modalities": ["text"], + } + response = omni_runner_handler.send_request(request_config) + assert response.success, f"Request failed: {response.error_message}" + assert response.text_content and len(response.text_content.strip()) > 0 + + +# ------------------------------------------------------------------ +# Test: mixed modality (audio + image + video) → audio output +# ------------------------------------------------------------------ + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "L4"}, num_cards=4) +@pytest.mark.parametrize("omni_runner", quant_params, indirect=True) +def test_mix_to_audio(omni_runner, omni_runner_handler): + """Mixed-modality input → audio output with W4A16 quantized Qwen2.5-Omni.""" + video = generate_synthetic_video(16, 16, 30)["np_array"] + image = generate_synthetic_image(16, 16)["np_array"] + audio = generate_synthetic_audio(1, 1, 16000)["np_array"] + if len(audio.shape) == 2: + audio = audio.squeeze() + + request_config = { + "prompts": "What is recited in the audio? What is in this image? Describe the video briefly.", + "videos": video, + "images": image, + "audios": (audio, 16000), + "modalities": ["audio"], + } + response = omni_runner_handler.send_request(request_config) + assert response.success, f"Request failed: {response.error_message}" diff --git a/tests/e2e/offline_inference/test_qwen3_omni_autoround_w4a16.py b/tests/e2e/offline_inference/test_qwen3_omni_autoround_w4a16.py new file mode 100644 index 000000000..3a3c874b6 --- /dev/null +++ b/tests/e2e/offline_inference/test_qwen3_omni_autoround_w4a16.py @@ -0,0 +1,205 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""E2E tests for Qwen3-Omni AutoRound W4A16 quantized inference. + +These tests cover text, audio, image, video, and mixed-modality inputs +to verify multimodal understanding with quantized weights. + +Requirements: + - CUDA GPUs (2x H100-80G or equivalent) + - The quantized model checkpoint (Intel/Qwen3-Omni-30B-A3B-Instruct-int4-AutoRound) +""" + +import os + +import pytest + +from tests.helpers.mark import hardware_test +from tests.helpers.media import ( + generate_synthetic_audio, + generate_synthetic_image, + generate_synthetic_video, +) +from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config + +QUANTIZED_MODEL = "Intel/Qwen3-Omni-30B-A3B-Instruct-int4-AutoRound" +BASELINE_MODEL = "Qwen/Qwen3-Omni-30B-A3B-Instruct" + +# Allow overriding via environment for local testing +QUANTIZED_MODEL = os.environ.get("QWEN3_OMNI_AUTOROUND_MODEL", QUANTIZED_MODEL) +BASELINE_MODEL = os.environ.get("QWEN3_OMNI_BASELINE_MODEL", BASELINE_MODEL) + +_CI_DEPLOY = get_deploy_config_path("ci/qwen3_omni_moe.yaml") + + +@pytest.fixture(scope="module", autouse=True) +def _qwen3_omni_env(): + """Set env vars required by multi-stage worker spawning. + + Must run before CUDA context init. Reverted after every test module + so that values do not leak into unrelated test files. + """ + with pytest.MonkeyPatch.context() as mp: + mp.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") + mp.setenv("VLLM_TEST_CLEAN_GPU_MEMORY", "0") + yield + + +def _get_stage_config(): + """Build a CI-friendly stage config with eager mode.""" + return modify_stage_config( + _CI_DEPLOY, + updates={ + "stages": { + 0: {"enforce_eager": True}, + 1: {"enforce_eager": True}, + }, + }, + ) + + +stage_config = _get_stage_config() + +# Parametrise: (model, stage_config) +quant_params = [(QUANTIZED_MODEL, stage_config)] + + +# ------------------------------------------------------------------ +# Test: text-only input → text output +# ------------------------------------------------------------------ + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100"}, num_cards=2) +@pytest.mark.parametrize("omni_runner", quant_params, indirect=True) +def test_text_to_text(omni_runner, omni_runner_handler): + """Text input → text output with W4A16 quantized Qwen3-Omni.""" + request_config = { + "prompts": "What is the capital of France?", + "modalities": ["text"], + } + response = omni_runner_handler.send_request(request_config) + assert response.success, f"Request failed: {response.error_message}" + assert response.text_content and len(response.text_content.strip()) > 0 + + +# ------------------------------------------------------------------ +# Test: audio input → text output +# ------------------------------------------------------------------ + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100"}, num_cards=2) +@pytest.mark.parametrize("omni_runner", quant_params, indirect=True) +def test_audio_to_text(omni_runner, omni_runner_handler): + """Audio input → text output with W4A16 quantized Qwen3-Omni.""" + audio = generate_synthetic_audio(1, 1, 16000)["np_array"] + if len(audio.shape) == 2: + audio = audio.squeeze() + + request_config = { + "prompts": "What is the content of this audio?", + "audios": (audio, 16000), + "modalities": ["text"], + } + response = omni_runner_handler.send_request(request_config) + assert response.success, f"Request failed: {response.error_message}" + assert response.text_content and len(response.text_content.strip()) > 0 + + +# ------------------------------------------------------------------ +# Test: image input → text output +# ------------------------------------------------------------------ + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100"}, num_cards=2) +@pytest.mark.parametrize("omni_runner", quant_params, indirect=True) +def test_image_to_text(omni_runner, omni_runner_handler): + """Image input → text output with W4A16 quantized Qwen3-Omni.""" + image = generate_synthetic_image(16, 16)["np_array"] + + request_config = { + "prompts": "Describe what you see in this image.", + "images": image, + "modalities": ["text"], + } + response = omni_runner_handler.send_request(request_config) + assert response.success, f"Request failed: {response.error_message}" + assert response.text_content and len(response.text_content.strip()) > 0 + + +# ------------------------------------------------------------------ +# Test: video input → text output +# ------------------------------------------------------------------ + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100"}, num_cards=2) +@pytest.mark.parametrize("omni_runner", quant_params, indirect=True) +def test_video_to_text(omni_runner, omni_runner_handler): + """Video input → text output with W4A16 quantized Qwen3-Omni.""" + video = generate_synthetic_video(224, 224, 300)["np_array"] + + request_config = { + "prompts": "Describe the video briefly.", + "videos": video, + "modalities": ["text"], + } + response = omni_runner_handler.send_request(request_config) + assert response.success, f"Request failed: {response.error_message}" + assert response.text_content and len(response.text_content.strip()) > 0 + + +# ------------------------------------------------------------------ +# Test: video input → audio output +# ------------------------------------------------------------------ + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100"}, num_cards=2) +@pytest.mark.parametrize("omni_runner", quant_params, indirect=True) +def test_video_to_audio(omni_runner, omni_runner_handler): + """Video input → audio output with W4A16 quantized Qwen3-Omni.""" + video = generate_synthetic_video(224, 224, 300)["np_array"] + + request_config = { + "prompts": "Describe the video briefly.", + "videos": video, + "modalities": ["audio"], + } + response = omni_runner_handler.send_request(request_config) + assert response.success, f"Request failed: {response.error_message}" + + +# ------------------------------------------------------------------ +# Test: mixed modality (audio + image + video) → audio output +# ------------------------------------------------------------------ + + +@pytest.mark.advanced_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100"}, num_cards=2) +@pytest.mark.parametrize("omni_runner", quant_params, indirect=True) +def test_mix_to_audio(omni_runner, omni_runner_handler): + """Mixed-modality input → audio output with W4A16 quantized Qwen3-Omni.""" + video = generate_synthetic_video(224, 224, 300)["np_array"] + image = generate_synthetic_image(16, 16)["np_array"] + audio = generate_synthetic_audio(1, 1, 16000)["np_array"] + if len(audio.shape) == 2: + audio = audio.squeeze() + + request_config = { + "prompts": "What is recited in the audio? What is in this image? Describe the video briefly.", + "videos": video, + "images": image, + "audios": (audio, 16000), + "modalities": ["audio"], + } + response = omni_runner_handler.send_request(request_config) + assert response.success, f"Request failed: {response.error_message}" diff --git a/tests/e2e/offline_inference/test_stable_audio_expansion.py b/tests/e2e/offline_inference/test_stable_audio_expansion.py index 2fc07eb24..a5d3e6d22 100644 --- a/tests/e2e/offline_inference/test_stable_audio_expansion.py +++ b/tests/e2e/offline_inference/test_stable_audio_expansion.py @@ -22,6 +22,8 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput from vllm_omni.platforms import current_omni_platform +pytestmark = [pytest.mark.full_model, pytest.mark.diffusion] + _SAMPLE_RATE = 44100 _CLIP_DURATION_S = 2.0 @@ -72,8 +74,6 @@ def generate_stable_audio_short_clip( return audio -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.cache @hardware_test(res={"cuda": "L4", "xpu": "B60"}) def test_stable_audio_quantization_and_teacache() -> None: diff --git a/tests/e2e/online_serving/test_bagel_expansion.py b/tests/e2e/online_serving/test_bagel_expansion.py index aa289c6d9..21142439b 100644 --- a/tests/e2e/online_serving/test_bagel_expansion.py +++ b/tests/e2e/online_serving/test_bagel_expansion.py @@ -19,6 +19,8 @@ import pytest from tests.helpers.mark import hardware_marks from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data +pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] + PROMPT = "A futuristic city skyline at twilight, cyberpunk style, ultra-detailed, high resolution." NEGATIVE_PROMPT = "low quality, blurry, distorted, deformed, watermark" @@ -112,8 +114,6 @@ def _get_diffusion_feature_cases(model: str): ] -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", _get_diffusion_feature_cases("ByteDance-Seed/BAGEL-7B-MoT"), diff --git a/tests/e2e/online_serving/test_bagel_online.py b/tests/e2e/online_serving/test_bagel_online.py index 745e9a116..a8ec65489 100644 --- a/tests/e2e/online_serving/test_bagel_online.py +++ b/tests/e2e/online_serving/test_bagel_online.py @@ -23,21 +23,19 @@ Equivalent to running: import base64 import os from io import BytesIO -from pathlib import Path import pytest from vllm.assets.image import ImageAsset from tests.helpers.mark import hardware_test from tests.helpers.runtime import OmniServerParams +from tests.helpers.stage_config import get_deploy_config_path os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0" MODEL = "ByteDance-Seed/BAGEL-7B-MoT" -STAGE_CONFIGS_PATH = str( - Path(__file__).parent.parent / "offline_inference" / "stage_configs" / "bagel_sharedmemory_ci.yaml" -) +STAGE_CONFIGS_PATH = get_deploy_config_path("ci/bagel.yaml") TEXT2IMG_PROMPT = "A cute cat" IMG2IMG_PROMPT = "Change the grass color to red" diff --git a/tests/e2e/online_serving/test_diffusers_adapter.py b/tests/e2e/online_serving/test_diffusers_adapter.py new file mode 100644 index 000000000..8b41db13a --- /dev/null +++ b/tests/e2e/online_serving/test_diffusers_adapter.py @@ -0,0 +1,56 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""End-to-end tests for DiffusersAdapterPipeline. + +It tests the full user flow of launching a diffusers-backed model and running inference. +""" + +import pytest +from PIL import Image + +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data + +pytestmark = [pytest.mark.diffusion, pytest.mark.core_model] + + +@pytest.mark.parametrize( + "omni_server", + [ + OmniServerParams( + model="tiny-random/Qwen-Image", + server_args=[ + "--diffusion-load-format", + "diffusers", + "--diffusers-call-kwargs", + '{"height": 512, "width": 0}', # deliberately weird width to be overridden + ], + ), + ], + indirect=True, +) +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_t2i_with_diffusers_adapter( + omni_server: OmniServer, + openai_client: OpenAIClientHandler, +): + messages = dummy_messages_from_mix_data(content_text="a photo of an astronaut riding a horse on mars") + + request_config = { + "model": omni_server.model, + "messages": messages, + "extra_body": { + "width": 512, + "num_inference_steps": 2, + "negative_prompt": "blurry", + "true_cfg_scale": 4.0, + "seed": 42, + }, + } + + response = openai_client.send_diffusion_request(request_config) + image: Image.Image = response[0].images[0] # pyright: ignore[reportOptionalSubscript] + + # Request config has incomplete width/height, so internal assertion in `send_diffusion_request` is incomplete. + assert image.size == (512, 512) diff --git a/tests/e2e/online_serving/test_dynin_omni_expansion.py b/tests/e2e/online_serving/test_dynin_omni_expansion.py index b63fcbd75..da179dbc8 100644 --- a/tests/e2e/online_serving/test_dynin_omni_expansion.py +++ b/tests/e2e/online_serving/test_dynin_omni_expansion.py @@ -5,29 +5,26 @@ Example online tests for Dynin-Omni model. """ import base64 -import gc import os from io import BytesIO from pathlib import Path -import numpy as np import pytest -import soundfile as sf from vllm.assets.image import ImageAsset from tests.helpers.mark import hardware_test -from tests.helpers.media import convert_audio_bytes_to_text from tests.helpers.runtime import OmniServerParams +pytestmark = [pytest.mark.full_model, pytest.mark.omni] + os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0" MODEL = "snu-aidas/Dynin-Omni" STAGE_CONFIG = str(Path(__file__).parent.parent / "stage_configs" / "dynin_omni_ci.yaml") -_WHISPER_SAMPLE_RATE_HZ = 16_000 T2I_PROMPT = "A high quality detailed living room interior photo." -T2S_PROMPT = "Please read this sentence naturally: Hello from Dynin-Omni online serving." +T2S_PROMPT = "Please read this sentence naturally: Hello from online serving." I2I_PROMPT = "Transform this outdoor nature boardwalk scene into a painting style with vivid colors." TEST_PARAMS = [OmniServerParams(model=MODEL, stage_config_path=STAGE_CONFIG, stage_init_timeout=600)] @@ -35,65 +32,6 @@ _STAGE_COUNT = 3 _I2I_STAGE_SAMPLING = {"max_tokens": 1, "temperature": 0.0, "top_p": 1.0, "detokenize": False} -def _prepare_audio_waveform_for_whisper(audio_data: np.ndarray, samplerate: int) -> np.ndarray: - """Normalize decoded audio into a mono 16 kHz float32 waveform for Whisper.""" - if samplerate <= 0: - raise ValueError(f"Invalid audio sample rate: {samplerate}") - - waveform = np.asarray(audio_data, dtype=np.float32) - if waveform.ndim == 0: - raise ValueError("Audio waveform must have at least one dimension") - if waveform.ndim > 1: - waveform = np.mean(waveform, axis=1) - if waveform.size == 0: - raise ValueError("Empty audio waveform") - - if samplerate != _WHISPER_SAMPLE_RATE_HZ: - target_num_samples = max(int(round(waveform.shape[0] * _WHISPER_SAMPLE_RATE_HZ / samplerate)), 1) - source_positions = np.arange(waveform.shape[0], dtype=np.float64) - target_positions = np.linspace( - 0.0, - max(waveform.shape[0] - 1, 0), - num=target_num_samples, - dtype=np.float64, - ) - waveform = np.interp(target_positions, source_positions, waveform).astype(np.float32) - - return np.ascontiguousarray(np.clip(waveform, -1.0, 1.0), dtype=np.float32) - - -def _convert_audio_bytes_to_text_without_ffmpeg(raw_bytes: bytes) -> str: - """Dynin t2s keeps Whisper transcription local to this test module and avoids ffmpeg.""" - import whisper - - data, samplerate = sf.read(BytesIO(raw_bytes), dtype="float32", always_2d=True) - audio_waveform = _prepare_audio_waveform_for_whisper(data, samplerate) - - model = whisper.load_model("small", device="cpu") - try: - transcript = model.transcribe( - audio_waveform, - temperature=0.0, - word_timestamps=True, - condition_on_previous_text=False, - )["text"] - finally: - del model - gc.collect() - - return transcript or "" - - -@pytest.fixture -def dynin_t2s_openai_client(openai_client, monkeypatch): - monkeypatch.setattr( - convert_audio_bytes_to_text, - "convert_audio_bytes_to_text", - _convert_audio_bytes_to_text_without_ffmpeg, - ) - return openai_client - - def _build_t2i_messages(prompt: str) -> list[dict]: return [{"role": "user", "content": [{"type": "text", "text": f"<|t2i|> {prompt}"}]}] @@ -118,8 +56,6 @@ def _build_i2i_messages(prompt: str) -> list[dict]: ] -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}) @pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True) def test_send_i2i_request_001(omni_server, openai_client) -> None: @@ -134,8 +70,6 @@ def test_send_i2i_request_001(omni_server, openai_client) -> None: openai_client.send_diffusion_request(request_config) -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}) @pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True) def test_send_t2i_request_001(omni_server, openai_client) -> None: @@ -147,14 +81,13 @@ def test_send_t2i_request_001(omni_server, openai_client) -> None: openai_client.send_diffusion_request(request_config) -@pytest.mark.core_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}) @pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True) -def test_send_t2s_request_001(omni_server, dynin_t2s_openai_client) -> None: +def test_send_t2s_request_001(omni_server, openai_client) -> None: request_config = { "model": omni_server.model, "messages": _build_t2s_messages(T2S_PROMPT), "modalities": ["audio"], + "audio_ref_text": T2S_PROMPT, } - dynin_t2s_openai_client.send_omni_request(request_config) + openai_client.send_omni_request(request_config) diff --git a/tests/e2e/online_serving/test_flux2_expansion.py b/tests/e2e/online_serving/test_flux2_expansion.py index ce06ad564..9a2b164b3 100644 --- a/tests/e2e/online_serving/test_flux2_expansion.py +++ b/tests/e2e/online_serving/test_flux2_expansion.py @@ -14,6 +14,8 @@ import pytest from tests.helpers.mark import hardware_marks from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler +pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] + FOUR_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"}, num_cards=4) POSITIVE_PROMPT = "A cat sitting on a windowsill" NEGATIVE_PROMPT = "blurry, low quality" @@ -86,8 +88,6 @@ def _get_diffusion_feature_cases(model: str): ] -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", _get_diffusion_feature_cases( diff --git a/tests/e2e/online_serving/test_flux2_klein_inpaint_expansion.py b/tests/e2e/online_serving/test_flux2_klein_inpaint_expansion.py index 0c45bb33f..a3d4f004e 100644 --- a/tests/e2e/online_serving/test_flux2_klein_inpaint_expansion.py +++ b/tests/e2e/online_serving/test_flux2_klein_inpaint_expansion.py @@ -16,6 +16,8 @@ from PIL import Image, ImageDraw from tests.helpers.runtime import OmniServer, OmniServerParams +pytestmark = [pytest.mark.full_model, pytest.mark.diffusion] + MODEL = "black-forest-labs/FLUX.2-klein-4B" _HEIGHT = 512 @@ -66,8 +68,6 @@ def _send_edit_request(host: str, port: int, model: str, image_b64: str, mask_b6 return response.json() -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True) def test_flux2_klein_inpaint_basic(omni_server: OmniServer): input_image_b64 = _image_to_base64_jpeg(Image.new("RGB", (_WIDTH, _HEIGHT), (128, 128, 128))) @@ -92,8 +92,6 @@ def test_flux2_klein_inpaint_basic(omni_server: OmniServer): assert img.size == (_WIDTH, _HEIGHT) -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True) def test_flux2_klein_inpaint_deterministic(omni_server: OmniServer): input_image_b64 = _image_to_base64_jpeg(Image.new("RGB", (_WIDTH, _HEIGHT), (128, 128, 128))) @@ -137,8 +135,6 @@ def test_flux2_klein_inpaint_deterministic(omni_server: OmniServer): ) -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True) def test_flux2_klein_inpaint_multiple_outputs(omni_server: OmniServer): input_image_b64 = _image_to_base64_jpeg(Image.new("RGB", (_WIDTH, _HEIGHT), (128, 128, 128))) diff --git a/tests/e2e/online_serving/test_flux_2_dev_expansion.py b/tests/e2e/online_serving/test_flux_2_dev_expansion.py index addf6f002..953cb448a 100644 --- a/tests/e2e/online_serving/test_flux_2_dev_expansion.py +++ b/tests/e2e/online_serving/test_flux_2_dev_expansion.py @@ -17,6 +17,8 @@ import pytest from tests.helpers.mark import hardware_marks from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data +pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] + MODEL = "black-forest-labs/FLUX.2-dev" PROMPT = "A cinematic mountain landscape at sunrise, dramatic clouds, ultra-detailed, realistic photography." NEGATIVE_PROMPT = "low quality, blurry, distorted, deformed, watermark" @@ -54,8 +56,6 @@ def _get_flux_2_dev_feature_cases(model: str): ] -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", _get_flux_2_dev_feature_cases(MODEL), diff --git a/tests/e2e/online_serving/test_flux_kontext_expansion.py b/tests/e2e/online_serving/test_flux_kontext_expansion.py index 574bb3db4..fd7d7d1b4 100644 --- a/tests/e2e/online_serving/test_flux_kontext_expansion.py +++ b/tests/e2e/online_serving/test_flux_kontext_expansion.py @@ -8,6 +8,8 @@ import pytest from tests.helpers.media import generate_synthetic_image from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data +pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] + EDIT_PROMPT = "Transform this modern, geometrist image into a Vincent van Gogh style impressionist painting." NEGATIVE_PROMPT = "blurry, low quality, modern, geometrist" MODEL = "black-forest-labs/FLUX.1-Kontext-dev" @@ -29,8 +31,6 @@ def _get_diffusion_feature_cases(model: str): ] -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", _get_diffusion_feature_cases(MODEL), @@ -54,8 +54,6 @@ def test_flux_kontext_text_to_image(omni_server: OmniServer, openai_client: Open openai_client.send_diffusion_request(request_config) -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", _get_diffusion_feature_cases(MODEL), @@ -83,8 +81,6 @@ def test_flux_kontext_image_edit(omni_server: OmniServer, openai_client: OpenAIC openai_client.send_diffusion_request(request_config) -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", _get_diffusion_feature_cases(MODEL), @@ -110,8 +106,6 @@ def test_flux_kontext_image_edit_no_negative(omni_server: OmniServer, openai_cli openai_client.send_diffusion_request(request_config) -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", _get_diffusion_feature_cases(MODEL), @@ -135,8 +129,6 @@ def test_flux_kontext_high_resolution(omni_server: OmniServer, openai_client: Op openai_client.send_diffusion_request(request_config) -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", _get_diffusion_feature_cases(MODEL), diff --git a/tests/e2e/online_serving/test_hunyuan_video_15_expansion.py b/tests/e2e/online_serving/test_hunyuan_video_15_expansion.py index 7af4c4344..681d8ff23 100644 --- a/tests/e2e/online_serving/test_hunyuan_video_15_expansion.py +++ b/tests/e2e/online_serving/test_hunyuan_video_15_expansion.py @@ -14,6 +14,8 @@ import pytest from tests.helpers.mark import hardware_marks from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler +pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] + PROMPT = "A cat walking across a sunlit garden, cinematic lighting, slow motion." NEGATIVE_PROMPT = "low quality, blurry, distorted" @@ -62,8 +64,6 @@ def _get_diffusion_feature_cases(model: str): ] -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", _get_diffusion_feature_cases(MODEL), diff --git a/tests/e2e/online_serving/test_longcat_image_edit_expansion.py b/tests/e2e/online_serving/test_longcat_image_edit_expansion.py index 9a96280f2..4be5fe42d 100644 --- a/tests/e2e/online_serving/test_longcat_image_edit_expansion.py +++ b/tests/e2e/online_serving/test_longcat_image_edit_expansion.py @@ -17,6 +17,8 @@ from tests.helpers.mark import hardware_marks from tests.helpers.media import generate_synthetic_image from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data +pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] + EDIT_PROMPT = "Transform this modern image into a cinematic animation style with vibrant colors and soft lighting." NEGATIVE_PROMPT = "blurry, low quality, distorted, oversaturated" SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}) @@ -50,8 +52,6 @@ def _get_diffusion_feature_cases(model: str): ] -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", _get_diffusion_feature_cases("meituan-longcat/LongCat-Image-Edit"), diff --git a/tests/e2e/online_serving/test_longcat_image_expansion.py b/tests/e2e/online_serving/test_longcat_image_expansion.py index 7d12db6a2..b9fd858b0 100644 --- a/tests/e2e/online_serving/test_longcat_image_expansion.py +++ b/tests/e2e/online_serving/test_longcat_image_expansion.py @@ -16,6 +16,8 @@ import pytest from tests.helpers.mark import hardware_marks from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data +pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] + TEXT_TO_IMAGE_PROMPT = ( "A cinematic illustration of a cat typing on a silver laptop, soft window light, highly detailed." ) @@ -51,8 +53,6 @@ def _get_diffusion_feature_cases(model: str): ] -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", _get_diffusion_feature_cases("meituan-longcat/LongCat-Image"), diff --git a/tests/e2e/online_serving/test_nextstep_expansion.py b/tests/e2e/online_serving/test_nextstep_expansion.py index 004a0967b..75290db5d 100644 --- a/tests/e2e/online_serving/test_nextstep_expansion.py +++ b/tests/e2e/online_serving/test_nextstep_expansion.py @@ -14,6 +14,8 @@ from tests.helpers.runtime import ( dummy_messages_from_mix_data, ) +pytestmark = [pytest.mark.full_model, pytest.mark.diffusion] + # L4: 4 GPUs + TP=4; XPU B60: 2 cards (use num_cards={"cuda": 4, "xpu": 4} if needed) FOUR_CARD_MARKS = hardware_marks( res={"cuda": "L4", "xpu": "B60"}, @@ -45,8 +47,6 @@ def _get_diffusion_feature_cases(model: str): ] -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", _get_diffusion_feature_cases(model=os.environ.get("VLLM_TEST_NEXTSTEP_MODEL", _DEFAULT_MODEL)), diff --git a/tests/e2e/online_serving/test_qwen3_omni.py b/tests/e2e/online_serving/test_qwen3_omni.py index 7d1a181d2..c9494d9fe 100644 --- a/tests/e2e/online_serving/test_qwen3_omni.py +++ b/tests/e2e/online_serving/test_qwen3_omni.py @@ -180,14 +180,19 @@ def test_text_to_text_001(omni_server, openai_client) -> None: @pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", prefix_test_params, indirect=True) -@pytest.mark.skip(reason="issue: #2833") def test_thinker_prefix_caching(omni_server, openai_client) -> None: """ Test thinker prefix caching by sending identical requests with an image (i.e., a large shared prefix) and verifying that the second request uses cached tokens - & produces the same output. + & produces the same output with greedy decoding. + + NOTE: The seed for this test is used as a regression test for the issue linked below; + https://github.com/vllm-project/vllm-omni/issues/2833; without passing the sampling + params, this test will fail with the current default stage configs. """ - image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}" + seed = 10 + img_res = generate_synthetic_image(224, 224, seed=seed) + image_data_url = f"data:image/jpeg;base64,{img_res['base64']}" messages = dummy_messages_from_mix_data( system_prompt=get_system_prompt(), image_data_url=image_data_url, @@ -199,16 +204,14 @@ def test_thinker_prefix_caching(omni_server, openai_client) -> None: "messages": messages, "stream": False, "modalities": ["text"], + "sampling_params_list": [{"seed": seed, "temperature": 0, "max_tokens": 16}] * 3, } response_1 = openai_client.send_omni_request(request_config, request_num=1)[0] response_2 = openai_client.send_omni_request(request_config, request_num=1)[0] - assert response_1.success - assert response_2.success - assert response_2.cached_tokens is not None # We should cache the vast majority of the prompt (image + up to last full block), - # and set seed in the CI config, so the second request should give an identical + # and set seed + temperature, so the second request should give an identical # response for the generated input image, even if we use dummy weights - assert response_2.cached_tokens > 0 + assert response_2.cached_tokens is not None and response_2.cached_tokens > 0 assert response_1.text_content == response_2.text_content diff --git a/tests/e2e/online_serving/test_qwen3_omni_expansion.py b/tests/e2e/online_serving/test_qwen3_omni_expansion.py index 6be2ebb2a..2ebf5c7e3 100644 --- a/tests/e2e/online_serving/test_qwen3_omni_expansion.py +++ b/tests/e2e/online_serving/test_qwen3_omni_expansion.py @@ -6,8 +6,6 @@ E2E Online tests for Qwen3-Omni model. import os -os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" - import pytest from tests.helpers.mark import hardware_test @@ -15,6 +13,10 @@ from tests.helpers.media import generate_synthetic_audio, generate_synthetic_ima from tests.helpers.runtime import OmniServerParams, dummy_messages_from_mix_data from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config +pytestmark = [pytest.mark.full_model, pytest.mark.omni] + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + model = "Qwen/Qwen3-Omni-30B-A3B-Instruct" AUDIO_KEY = ["test"] @@ -82,6 +84,7 @@ test_params = [ model=model, stage_config_path=get_async_chunk_config(default_path), use_stage_cli=True, + server_args=["--async-chunk"], ), id="async_chunk", ), @@ -130,8 +133,6 @@ def get_max_batch_size(size_type="few"): return batch_sizes.get(size_type, 5) -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_text_to_audio_001(omni_server, openai_client) -> None: @@ -154,8 +155,6 @@ def test_text_to_audio_001(omni_server, openai_client) -> None: openai_client.send_omni_request(request_config) -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True) def test_text_to_text_audio_001(omni_server, openai_client) -> None: @@ -176,8 +175,6 @@ def test_text_to_text_audio_001(omni_server, openai_client) -> None: openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_text_video_to_text_001(omni_server, openai_client) -> None: @@ -202,8 +199,6 @@ def test_text_video_to_text_001(omni_server, openai_client) -> None: openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True) def test_text_audio_to_text_audio_001(omni_server, openai_client) -> None: @@ -227,8 +222,6 @@ def test_text_audio_to_text_audio_001(omni_server, openai_client) -> None: openai_client.send_omni_request(request_config) -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True) def test_text_audio_to_text_audio_002(omni_server, openai_client) -> None: @@ -254,8 +247,6 @@ def test_text_audio_to_text_audio_002(omni_server, openai_client) -> None: openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True) def test_text_image_to_text_audio_001(omni_server, openai_client) -> None: @@ -280,8 +271,6 @@ def test_text_image_to_text_audio_001(omni_server, openai_client) -> None: openai_client.send_omni_request(request_config) -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True) def test_large_image_to_text_audio_001(omni_server, openai_client) -> None: @@ -310,8 +299,6 @@ def test_large_image_to_text_audio_001(omni_server, openai_client) -> None: openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True) def test_text_video_to_text_audio_001(omni_server, openai_client) -> None: @@ -338,8 +325,6 @@ def test_text_video_to_text_audio_001(omni_server, openai_client) -> None: @pytest.mark.skip(reason="There is a known issue with shape mismatch error.") -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True) def test_mix_to_text_audio_001(omni_server, openai_client) -> None: @@ -369,8 +354,6 @@ def test_mix_to_text_audio_001(omni_server, openai_client) -> None: openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) @pytest.mark.skip(reason="issue: #2827") @@ -398,8 +381,6 @@ def test_audio_in_video_001(omni_server, openai_client) -> None: openai_client.send_omni_request(request_config) -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_audio_in_video_002(omni_server, openai_client) -> None: @@ -427,8 +408,6 @@ def test_audio_in_video_002(omni_server, openai_client) -> None: openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_one_word_prompt_001(omni_server, openai_client) -> None: @@ -463,8 +442,6 @@ def test_one_word_prompt_001(omni_server, openai_client) -> None: print(f"Similarity assertion failed, retrying {attempt + 2}/{_max_retries}: {e!r}") -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_speaker_001(omni_server, openai_client) -> None: @@ -490,8 +467,6 @@ def test_speaker_001(omni_server, openai_client) -> None: openai_client.send_omni_request(request_config) -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) @pytest.mark.skip(reason="Known issue: occasional inaccuracy in voice recognition.") @@ -528,8 +503,6 @@ def test_speaker_002(omni_server, openai_client) -> None: print(f"Gender assertion failed, retrying {attempt + 2}/{_max_retries}: {e!r}") -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_speaker_003(omni_server, openai_client) -> None: @@ -555,8 +528,6 @@ def test_speaker_003(omni_server, openai_client) -> None: openai_client.send_omni_request(request_config) -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_language_001(omni_server, openai_client) -> None: diff --git a/tests/e2e/online_serving/test_qwen3_tts_base_expansion.py b/tests/e2e/online_serving/test_qwen3_tts_base_expansion.py index 527eeaef4..d86f96af0 100644 --- a/tests/e2e/online_serving/test_qwen3_tts_base_expansion.py +++ b/tests/e2e/online_serving/test_qwen3_tts_base_expansion.py @@ -9,15 +9,17 @@ actual model inference, not mocks. import os -os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" -os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0" - import pytest from tests.helpers.mark import hardware_test from tests.helpers.runtime import OmniServerParams from tests.helpers.stage_config import get_deploy_config_path +pytestmark = [pytest.mark.full_model, pytest.mark.omni] + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" +os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0" + MODEL = "Qwen/Qwen3-TTS-12Hz-0.6B-Base" REF_AUDIO_URL = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-TTS-Repo/clone_2.wav" @@ -61,9 +63,6 @@ tts_server_params = [ ] -@pytest.mark.advanced_model -@pytest.mark.core_model -@pytest.mark.omni @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) def test_voice_clone_streaming_001(omni_server, openai_client) -> None: @@ -89,9 +88,6 @@ def test_voice_clone_streaming_001(omni_server, openai_client) -> None: openai_client.send_audio_speech_request(request_config, request_num=get_max_batch_size("few")) -@pytest.mark.advanced_model -@pytest.mark.core_model -@pytest.mark.omni @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) def test_response_format_001(omni_server, openai_client) -> None: diff --git a/tests/e2e/online_serving/test_qwen3_tts_customvoice_expansion.py b/tests/e2e/online_serving/test_qwen3_tts_customvoice_expansion.py index b32303842..d97e70e41 100644 --- a/tests/e2e/online_serving/test_qwen3_tts_customvoice_expansion.py +++ b/tests/e2e/online_serving/test_qwen3_tts_customvoice_expansion.py @@ -9,15 +9,17 @@ actual model inference, not mocks. import os -os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" -os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0" - import pytest from tests.helpers.mark import hardware_test from tests.helpers.runtime import OmniServerParams from tests.helpers.stage_config import get_deploy_config_path +pytestmark = [pytest.mark.full_model, pytest.mark.omni] + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" +os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0" + MODEL = "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice" @@ -59,8 +61,6 @@ tts_server_params = [ ] -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) def test_voice_001(omni_server, openai_client) -> None: @@ -92,8 +92,6 @@ def test_voice_001(omni_server, openai_client) -> None: raise -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) def test_voice_002(omni_server, openai_client) -> None: @@ -117,8 +115,6 @@ def test_voice_002(omni_server, openai_client) -> None: openai_client.send_audio_speech_request(request_config) -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) def test_voice_003(omni_server, openai_client) -> None: @@ -142,8 +138,6 @@ def test_voice_003(omni_server, openai_client) -> None: openai_client.send_audio_speech_request(request_config) -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "L4"}, num_cards=1) @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True) def test_language_001(omni_server, openai_client) -> None: diff --git a/tests/e2e/online_serving/test_qwen_image_edit_expansion.py b/tests/e2e/online_serving/test_qwen_image_edit_expansion.py index c2461977d..bb8bf7dca 100644 --- a/tests/e2e/online_serving/test_qwen_image_edit_expansion.py +++ b/tests/e2e/online_serving/test_qwen_image_edit_expansion.py @@ -11,6 +11,8 @@ from tests.helpers.mark import hardware_marks from tests.helpers.media import generate_synthetic_image from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data +pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] + EDIT_PROMPT = "Transform this modern, geometrist image into a Vincent van Gogh style impressionist painting." MULTI_EDIT_PROMPT = ( "Transform the first image into a Dadaism collage art. " @@ -108,8 +110,6 @@ def _get_diffusion_feature_cases(model: str): ] -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit"), @@ -138,8 +138,6 @@ def test_qwen_image_edit(omni_server: OmniServer, openai_client: OpenAIClientHan openai_client.send_diffusion_request(request_config) -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", _get_diffusion_feature_cases("Qwen/Qwen-Image-Edit-2509"), diff --git a/tests/e2e/online_serving/test_qwen_image_expansion.py b/tests/e2e/online_serving/test_qwen_image_expansion.py index b6f91c13d..0830e4944 100644 --- a/tests/e2e/online_serving/test_qwen_image_expansion.py +++ b/tests/e2e/online_serving/test_qwen_image_expansion.py @@ -15,6 +15,8 @@ import pytest from tests.helpers.mark import hardware_marks from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data +pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] + T2I_PROMPT = "A photo of a cat sitting on a laptop keyboard, digital art style." NEGATIVE_PROMPT = "blurry, low quality" SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}) @@ -117,8 +119,6 @@ def _get_diffusion_feature_cases(model: str): ] -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", _get_diffusion_feature_cases("Qwen/Qwen-Image"), @@ -142,8 +142,6 @@ def test_qwen_image(omni_server: OmniServer, openai_client: OpenAIClientHandler) openai_client.send_diffusion_request(request_config) -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", _get_diffusion_feature_cases("Qwen/Qwen-Image-2512"), diff --git a/tests/e2e/online_serving/test_qwen_image_layered_expansion.py b/tests/e2e/online_serving/test_qwen_image_layered_expansion.py index b958cfc05..4e79beab7 100644 --- a/tests/e2e/online_serving/test_qwen_image_layered_expansion.py +++ b/tests/e2e/online_serving/test_qwen_image_layered_expansion.py @@ -18,6 +18,8 @@ from tests.helpers.mark import hardware_marks from tests.helpers.media import decode_b64_image, generate_synthetic_image from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler, dummy_messages_from_mix_data +pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] + MODEL = "Qwen/Qwen-Image-Layered" EDIT_PROMPT = "Decompose this image into layers." EMPTY_EDIT_PROMPT = "" @@ -71,8 +73,6 @@ FEATURE_CASES = [ ] -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize("omni_server", FEATURE_CASES, indirect=True) def test_feature(omni_server: OmniServer, openai_client: OpenAIClientHandler): """Test feature combinations with Qwen-Image-Layered.""" @@ -149,8 +149,6 @@ def _collect_image_url_items(openai_client: OpenAIClientHandler, request_config: return image_items -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize( "omni_server, expected_layers", LAYERS_GUARD_CASES, @@ -224,8 +222,6 @@ PROMPT_CASES = [ ] -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize("omni_server", PROMPT_CASES, indirect=True) def test_empty_prompt(omni_server: OmniServer, openai_client: OpenAIClientHandler): """Test feature combinations with Qwen-Image-Layered.""" diff --git a/tests/e2e/online_serving/test_sd3_expansion.py b/tests/e2e/online_serving/test_sd3_expansion.py index 37e590ba3..09b50d2e5 100644 --- a/tests/e2e/online_serving/test_sd3_expansion.py +++ b/tests/e2e/online_serving/test_sd3_expansion.py @@ -7,6 +7,8 @@ import pytest from tests.helpers.mark import hardware_marks from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler +pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] + FOUR_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "L4"}, num_cards=4) POSITIVE_PROMPT = "A serene mountain landscape at sunset" NEGATIVE_PROMPT = "blurry, low quality, distorted" @@ -35,8 +37,6 @@ def _get_diffusion_feature_cases(model: str): ] -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", _get_diffusion_feature_cases( diff --git a/tests/e2e/online_serving/test_wan22_expansion.py b/tests/e2e/online_serving/test_wan22_expansion.py index 7e5bc9121..36035f649 100644 --- a/tests/e2e/online_serving/test_wan22_expansion.py +++ b/tests/e2e/online_serving/test_wan22_expansion.py @@ -23,6 +23,8 @@ from tests.helpers.mark import hardware_marks from tests.helpers.media import generate_synthetic_image from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler +pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] + PROMPT = "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." NEGATIVE_PROMPT = "low quality, blurry, distorted face, extra limbs, bad anatomy, watermark, logo, text, ugly, deformed, mutated, jpeg artifacts" SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}) @@ -79,8 +81,6 @@ def _get_wan22_feature_cases(): return cases -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", _get_wan22_feature_cases(), diff --git a/tests/e2e/online_serving/test_wan_2_1_vace_expansion.py b/tests/e2e/online_serving/test_wan_2_1_vace_expansion.py index 1f7a8c072..4f0e96447 100644 --- a/tests/e2e/online_serving/test_wan_2_1_vace_expansion.py +++ b/tests/e2e/online_serving/test_wan_2_1_vace_expansion.py @@ -26,6 +26,8 @@ import pytest from tests.helpers.mark import hardware_marks from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler +pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] + MODEL = "Wan-AI/Wan2.1-VACE-1.3B-diffusers" PROMPT = "A cat walking slowly across a sunlit garden path" @@ -131,8 +133,6 @@ def _get_vace_feature_cases(): ] -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", _get_vace_feature_cases(), diff --git a/tests/e2e/online_serving/test_zimage_expansion.py b/tests/e2e/online_serving/test_zimage_expansion.py index 679233c82..d730db0ca 100644 --- a/tests/e2e/online_serving/test_zimage_expansion.py +++ b/tests/e2e/online_serving/test_zimage_expansion.py @@ -16,6 +16,8 @@ import pytest from tests.helpers.mark import hardware_marks from tests.helpers.runtime import OmniServer, OmniServerParams, OpenAIClientHandler +pytestmark = [pytest.mark.diffusion, pytest.mark.full_model] + MODEL = "Tongyi-MAI/Z-Image-Turbo" PROMPT = "A high-detail studio photo of an orange tabby cat sitting on a laptop keyboard." @@ -101,8 +103,6 @@ def _get_diffusion_feature_cases(): ] -@pytest.mark.advanced_model -@pytest.mark.diffusion @pytest.mark.parametrize( "omni_server", _get_diffusion_feature_cases(), diff --git a/tests/e2e/stage_configs/bailingmm_moe_v2_lite_ci.yaml b/tests/e2e/stage_configs/bailingmm_moe_v2_lite_ci.yaml index fb0c72cc5..53cc73ce0 100644 --- a/tests/e2e/stage_configs/bailingmm_moe_v2_lite_ci.yaml +++ b/tests/e2e/stage_configs/bailingmm_moe_v2_lite_ci.yaml @@ -1,4 +1,8 @@ -# Thinker stage only +# CI stage config for Ming-flash-omni-2.0 thinker+talker pipeline. +# Stage 0: Thinker (multimodal understanding -> text generation) +# Stage 1: Talker (text -> audio waveform via CFM + AudioVAE) +# +# The following config has been verified on 4x H100-80G GPUs stage_args: - stage_id: 0 stage_type: llm @@ -10,10 +14,12 @@ stage_args: model_arch: MingFlashOmniForConditionalGeneration worker_type: ar scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - gpu_memory_utilization: 0.9 + gpu_memory_utilization: 0.74 enforce_eager: false trust_remote_code: true - engine_output_type: latent + # Ming Thinker -> talker bridge reads detokenised text from + # source_output.outputs[0].text (not hidden states). + engine_output_type: text distributed_executor_backend: "mp" enable_prefix_caching: false max_num_batched_tokens: 32768 @@ -22,14 +28,48 @@ stage_args: hf_config_name: llm_config load_format: dummy mm_processor_cache_gb: 0 + compilation_config: + pass_config: + # Disable fused all-reduce to avoid a vllm/flashinfer version mismatch. + fuse_allreduce_rms: false final_output: true final_output_type: text is_comprehension: true default_sampling_params: - temperature: 0.4 + temperature: 0.0 top_p: 0.9 max_tokens: 100 repetition_penalty: 1.05 seed: 42 detokenize: true ignore_eos: false + + - stage_id: 1 + stage_type: llm + runtime: + devices: "3" + max_batch_size: 1 + engine_args: + model_stage: ming_tts + model_arch: MingFlashOmniTalkerForConditionalGeneration + worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker + scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler + gpu_memory_utilization: 0.18 + enforce_eager: true + trust_remote_code: true + engine_output_type: audio + enable_prefix_caching: false + max_num_batched_tokens: 1000000 + tokenizer_subdir: talker/llm + # The HF repo ships BailingMM2Config (thinker-only) at root; + # OmniModelConfig treats that as "stage does not share outer mrope". + hf_config_name: talker_config + load_format: dummy + engine_input_source: [0] + custom_process_input_func: vllm_omni.model_executor.stage_input_processors.ming_flash_omni.thinker2talker + final_output: true + final_output_type: audio + default_sampling_params: + temperature: 0.0 + max_tokens: 1 + seed: 42 diff --git a/tests/e2e/stage_configs/bailingmm_moe_v2_lite_thinker_only_ci.yaml b/tests/e2e/stage_configs/bailingmm_moe_v2_lite_thinker_only_ci.yaml new file mode 100644 index 000000000..fb0c72cc5 --- /dev/null +++ b/tests/e2e/stage_configs/bailingmm_moe_v2_lite_thinker_only_ci.yaml @@ -0,0 +1,35 @@ +# Thinker stage only +stage_args: + - stage_id: 0 + stage_type: llm + runtime: + devices: "0,1,2,3" + max_batch_size: 1 + engine_args: + model_stage: thinker + model_arch: MingFlashOmniForConditionalGeneration + worker_type: ar + scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler + gpu_memory_utilization: 0.9 + enforce_eager: false + trust_remote_code: true + engine_output_type: latent + distributed_executor_backend: "mp" + enable_prefix_caching: false + max_num_batched_tokens: 32768 + max_model_len: 32768 + tensor_parallel_size: 4 + hf_config_name: llm_config + load_format: dummy + mm_processor_cache_gb: 0 + final_output: true + final_output_type: text + is_comprehension: true + default_sampling_params: + temperature: 0.4 + top_p: 0.9 + max_tokens: 100 + repetition_penalty: 1.05 + seed: 42 + detokenize: true + ignore_eos: false diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index 4d69f24c5..0d61f6a67 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -139,6 +139,41 @@ def test_from_cli_args_picks_up_stage_configs_path(): assert args.custom_pipeline_args is None +def test_qwen3_tts_code2wav_injects_max_position_embeddings(monkeypatch): + """Ensure Code2Wav mirrors stage max_model_len into nested HF overrides. + + Qwen3-TTS Code2Wav is a pure decoder stage whose runtime max_model_len can + legitimately exceed the base checkpoint's default text max length. Recent + vLLM validates these values during ModelConfig creation, so we inject + ``talker_config.max_position_embeddings`` before delegating to vLLM. + """ + captured: dict[str, object] = {} + baseline_config = Mock() + + def fake_create_model_config(self): + captured["hf_overrides"] = self.hf_overrides + return baseline_config + + monkeypatch.setattr(EngineArgs, "create_model_config", fake_create_model_config) + monkeypatch.setattr( + OmniModelConfig, + "from_vllm_model_config", + classmethod(lambda cls, model_config, **omni_kwargs: model_config), + ) + + OmniEngineArgs( + model_arch="Qwen3TTSCode2Wav", + max_model_len=65536, + ).create_model_config() + + assert captured["hf_overrides"] == { + "architectures": ["Qwen3TTSCode2Wav"], + "talker_config": { + "max_position_embeddings": 65536, + }, + } + + def test_stage_specific_text_config_override(): """Ensure dependent attributes are updated when using stage-specific config.""" vllm_config = EngineArgs().create_model_config() diff --git a/tests/engine/test_orchestrator.py b/tests/engine/test_orchestrator.py index 0b549f58e..c07762ad5 100644 --- a/tests/engine/test_orchestrator.py +++ b/tests/engine/test_orchestrator.py @@ -10,11 +10,14 @@ from types import SimpleNamespace from typing import Any import janus +import psutil import pytest from vllm.outputs import CompletionOutput, RequestOutput from vllm.sampling_params import SamplingParams +from vllm.v1.engine.core_client import AsyncMPClient from vllm_omni.engine.orchestrator import Orchestrator +from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.outputs import OmniRequestOutput @@ -104,6 +107,42 @@ class FakeOutputProcessor: return None +class _FakeProc: + pid = 1234 + + def __init__(self): + self.terminated = False + self.killed = False + self.join_calls = [] + + def is_alive(self): + return not self.terminated and not self.killed + + def terminate(self): + self.terminated = True + + def kill(self): + self.killed = True + + def join(self, timeout=None): + self.join_calls.append(timeout) + + +class _FakeChildProc: + def __init__(self): + self.terminated = False + self.killed = False + + def is_running(self): + return not self.terminated and not self.killed + + def terminate(self): + self.terminated = True + + def kill(self): + self.killed = True + + def _sampling_params(max_tokens: int = 4) -> SamplingParams: return SamplingParams(max_tokens=max_tokens) @@ -269,6 +308,61 @@ async def _enqueue_abort_request(orchestrator_fixture: OrchestratorFixture, requ ) +def test_stage_engine_core_client_shutdown_cleans_children_if_base_shutdown_fails(monkeypatch): + fake_proc = _FakeProc() + fake_child = _FakeChildProc() + + class FakePsutilProcess: + def __init__(self, pid): + assert pid == fake_proc.pid + + def children(self, recursive=True): + assert recursive + return [fake_child] + + def fail_base_shutdown(self): + raise RuntimeError("base shutdown failed") + + monkeypatch.setattr(psutil, "Process", FakePsutilProcess) + monkeypatch.setattr(psutil, "wait_procs", lambda procs, timeout: (list(procs), [])) + monkeypatch.setattr(AsyncMPClient, "shutdown", fail_base_shutdown) + + client = object.__new__(StageEngineCoreClient) + client._proc = fake_proc + + with pytest.raises(RuntimeError, match="base shutdown failed"): + client.shutdown() + + assert fake_proc.terminated + assert fake_proc.join_calls == [5] + assert fake_child.terminated + + +def test_stage_engine_core_client_shutdown_kills_stubborn_children(monkeypatch): + fake_proc = _FakeProc() + fake_child = _FakeChildProc() + + class FakePsutilProcess: + def __init__(self, pid): + assert pid == fake_proc.pid + + def children(self, recursive=True): + assert recursive + return [fake_child] + + monkeypatch.setattr(psutil, "Process", FakePsutilProcess) + monkeypatch.setattr(psutil, "wait_procs", lambda procs, timeout: ([], list(procs))) + monkeypatch.setattr(AsyncMPClient, "shutdown", lambda self: None) + + client = object.__new__(StageEngineCoreClient) + client._proc = fake_proc + + client.shutdown() + + assert fake_child.terminated + assert fake_child.killed + + @pytest.fixture def orchestrator_factory(): fixtures: list[OrchestratorFixture] = [] diff --git a/tests/engine/test_orchestrator_kv_sender_info.py b/tests/engine/test_orchestrator_kv_sender_info.py index 7e3fe0906..ec3a42e35 100644 --- a/tests/engine/test_orchestrator_kv_sender_info.py +++ b/tests/engine/test_orchestrator_kv_sender_info.py @@ -183,6 +183,52 @@ def test_forward_to_diffusion_uses_engine_input_source_for_kv_sender_info(): } +def test_forward_to_diffusion_returns_terminal_error_for_empty_custom_inputs(): + orchestrator = object.__new__(Orchestrator) + sender_stage = _DummySenderStage({"host": "10.0.0.2", "zmq_port": 50151}) + diffusion_stage = _DummyDiffusionStage(engine_input_source=[0]) + diffusion_stage.custom_process_input_func = lambda *_args, **_kwargs: [] + + class _AsyncQueue: + def __init__(self): + self.items = [] + + async def put(self, item): + self.items.append(item) + + orchestrator.num_stages = 2 + orchestrator.stage_clients = [sender_stage, diffusion_stage] + orchestrator._cfg_tracker = CfgCompanionTracker() + orchestrator.stage_vllm_configs = [None, None] + orchestrator.output_processors = [None, None] + orchestrator.output_async_queue = _AsyncQueue() + orchestrator.request_states = {} + orchestrator._pd_kv_params = {} + + params = OmniDiffusionSamplingParams() + req_state = OrchestratorRequestState( + request_id="req-empty", + prompt={"prompt": "hello"}, + sampling_params_list=[SamplingParams(max_tokens=4), params], + final_stage_id=1, + ) + orchestrator.request_states["req-empty"] = req_state + + output = SimpleNamespace(request_id="req-empty", finished=True) + asyncio.run(Orchestrator._forward_to_next_stage(orchestrator, "req-empty", 0, output, req_state)) + + assert sender_stage.engine_outputs == [output] + assert diffusion_stage.calls == [] + assert len(orchestrator.output_async_queue.items) == 1 + terminal_msg = orchestrator.output_async_queue.items[0] + assert terminal_msg["type"] == "output" + assert terminal_msg["request_id"] == "req-empty" + assert terminal_msg["stage_id"] == 1 + assert terminal_msg["finished"] is True + assert "produced no valid inputs" in terminal_msg["engine_outputs"].error + assert "req-empty" not in orchestrator.request_states + + def test_prewarm_diffusion_attaches_kv_sender_info(): orchestrator = object.__new__(Orchestrator) sender_stage = _DummySenderStage({"host": "10.0.0.3", "zmq_port": 50151}) diff --git a/tests/engine/test_single_stage_mode.py b/tests/engine/test_single_stage_mode.py index 28ccccaa2..8169f586a 100644 --- a/tests/engine/test_single_stage_mode.py +++ b/tests/engine/test_single_stage_mode.py @@ -422,6 +422,39 @@ class TestSingleStageModeDetection: ) assert engine._single_stage_id_filter is None + def test_engine_args_create_only_forwards_explicit_fields(self, mocker: MockerFixture): + from vllm_omni.engine.arg_utils import OmniEngineArgs + + captured: dict[str, Any] = {} + + def fake_resolve(self, model: str, kwargs: dict[str, Any]): + captured.update(kwargs) + return "/fake/path", [_make_stage_cfg(0)] + + mocker.patch.object(AsyncOmniEngine, "_resolve_stage_configs", fake_resolve) + mocker.patch.object(AsyncOmniEngine, "_bootstrap_orchestrator") + mock_thread_cls = mocker.patch("threading.Thread") + mock_future_cls = mocker.patch("concurrent.futures.Future") + mock_future = mocker.Mock() + mock_future.result.return_value = mocker.Mock() + mock_future_cls.return_value = mock_future + mock_thread = mocker.Mock() + mock_thread.is_alive.return_value = False + mock_thread_cls.return_value = mock_thread + + ea = OmniEngineArgs.create(model="ignored", gpu_memory_utilization=0.5) + AsyncOmniEngine(model="fake-model", engine_args=ea) + + assert captured["gpu_memory_utilization"] == 0.5 + assert "model" not in captured + assert "max_num_seqs" not in captured + + def test_bare_engine_args_rejected(self, mocker: MockerFixture): + from vllm_omni.engine.arg_utils import OmniEngineArgs + + with pytest.raises(TypeError, match="OmniEngineArgs.create"): + self._make_engine_no_thread(mocker, engine_args=OmniEngineArgs(model="fake-model")) + def test_master_address_and_port_stored(self, mocker: MockerFixture): engine = self._make_engine_no_thread( mocker, diff --git a/tests/entrypoints/openai_api/conftest_video.py b/tests/entrypoints/openai_api/conftest_video.py new file mode 100644 index 000000000..3d5e0d855 --- /dev/null +++ b/tests/entrypoints/openai_api/conftest_video.py @@ -0,0 +1,31 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shared test helpers for video streaming tests.""" + +from __future__ import annotations + +import io + +import pytest + +np = pytest.importorskip("numpy", reason="numpy required for video stream tests") +PIL = pytest.importorskip("PIL", reason="Pillow required for video stream tests") +from PIL import Image # noqa: E402 + + +def make_jpeg(r: int = 128, g: int = 128, b: int = 128, size: int = 64) -> bytes: + """Create a solid-colour JPEG image.""" + img = Image.new("RGB", (size, size), (r, g, b)) + buf = io.BytesIO() + img.save(buf, format="JPEG", quality=95) + return buf.getvalue() + + +def make_gradient_jpeg(seed: int, size: int = 64) -> bytes: + """Create a random-gradient JPEG that varies based on *seed*.""" + rng = np.random.RandomState(seed) + arr = rng.randint(0, 256, (size, size, 3), dtype=np.uint8) + img = Image.fromarray(arr, "RGB") + buf = io.BytesIO() + img.save(buf, format="JPEG", quality=95) + return buf.getvalue() diff --git a/tests/entrypoints/openai_api/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py index 4cf46f18c..7fcec2a32 100644 --- a/tests/entrypoints/openai_api/test_image_server.py +++ b/tests/entrypoints/openai_api/test_image_server.py @@ -13,16 +13,22 @@ from argparse import Namespace from types import SimpleNamespace import pytest +from fastapi import FastAPI from fastapi.testclient import TestClient from PIL import Image from pytest_mock import MockerFixture from vllm import SamplingParams +from vllm.entrypoints.openai.models.protocol import BaseModelPath +from vllm_omni.entrypoints.async_omni import AsyncOmni +from vllm_omni.entrypoints.openai.api_server import _DiffusionServingModels, router from vllm_omni.entrypoints.openai.image_api_utils import ( encode_image_base64, parse_size, ) +from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.model_executor.stage_input_processors.glm_image import compute_max_tokens pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -569,6 +575,84 @@ def test_multistage_images_async_omni_construction(async_omni_test_client): assert captured[1].guidance_scale == 6.5 +def test_generate_images_async_omni_glm_image_sets_stage0_max_tokens(): + """GLM-Image multistage requests must compute stage-0 max_tokens from size.""" + + class FakeAsyncOmniClass(AsyncOmni): + def __init__(self): + stage_configs = [ + SimpleNamespace(stage_type="llm", is_comprehension=True, model_arch="GlmImageForConditionalGeneration"), + SimpleNamespace(stage_type="diffusion", is_comprehension=False, model_arch="GlmImagePipeline"), + ] + default_sampling_params_list = [ + SamplingParams(temperature=0.1, seed=42), + OmniDiffusionSamplingParams(height=1024, width=1024), + ] + self.engine = SimpleNamespace( + stage_configs=stage_configs, + default_sampling_params_list=default_sampling_params_list, + ) + self.default_sampling_params_list = default_sampling_params_list + self.captured_sampling_params_list = None + self.captured_prompt = None + self._images = [Image.new("RGB", (64, 64), color="green")] + self.od_config = SimpleNamespace(supports_multimodal_inputs=True) + + async def generate(self, prompt, request_id, sampling_params=None, sampling_params_list=None): + self.captured_sampling_params_list = ( + sampling_params_list if sampling_params_list is not None else [sampling_params] + ) + self.captured_prompt = prompt + yield MockGenerationResult([img.copy() for img in self._images]) + + def __class_getitem__(cls, item): + return cls + + def get_diffusion_od_config(self): + return self.od_config + + app = FastAPI() + app.include_router(router) + engine = FakeAsyncOmniClass() + chat_handler = object.__new__(OmniOpenAIServingChat) + chat_handler.engine_client = engine + chat_handler._diffusion_engine = None + app.state.openai_serving_chat = chat_handler + app.state.engine_client = engine + app.state.stage_configs = [ + SimpleNamespace(stage_type="llm", model_arch="GlmImageForConditionalGeneration"), + SimpleNamespace(stage_type="diffusion", model_arch="GlmImagePipeline"), + ] + app.state.openai_serving_models = _DiffusionServingModels( + [BaseModelPath(name="THUDM/GLM-4.5V", model_path="THUDM/GLM-4.5V")] + ) + app.state.args = Namespace( + default_sampling_params='{"1": {"num_inference_steps":4, "guidance_scale":7.5, "generator_device":"cpu"}}', + max_generated_image_size=1048576, + ) + client = TestClient(app) + + response = client.post( + "/v1/images/generations", + json={ + "prompt": "a coral reef", + "n": 1, + "size": "1024x1024", + "seed": 7, + }, + ) + assert response.status_code == 200 + + captured = engine.captured_sampling_params_list + assert captured is not None + assert len(captured) == 2 + assert captured[0].max_tokens == compute_max_tokens(1024, 1024, is_i2i=False) + assert captured[0].extra_args["target_h"] == 1024 + assert captured[0].extra_args["target_w"] == 1024 + assert captured[1].height == 1024 + assert captured[1].width == 1024 + + def test_image_edits_async_omni_stage_configs_only(async_omni_stage_configs_only_client): """Regression: image edits accepts refactored AsyncOmni without stage_list.""" img_bytes = make_test_image_bytes((16, 16)) diff --git a/tests/e2e/online_serving/test_qwen3_omni_realtime_websocket.py b/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py similarity index 91% rename from tests/e2e/online_serving/test_qwen3_omni_realtime_websocket.py rename to tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py index f3b261081..90f8897c5 100644 --- a/tests/e2e/online_serving/test_qwen3_omni_realtime_websocket.py +++ b/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py @@ -23,33 +23,28 @@ from tests.helpers.media import ( generate_synthetic_audio, ) from tests.helpers.runtime import OmniServerParams -from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config +from tests.helpers.stage_config import get_deploy_config_path os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" MODEL = "Qwen/Qwen3-Omni-30B-A3B-Instruct" +# Synthetic input for realtime E2E (``generate_synthetic_audio``); distinct cache file per phrase. +REALTIME_SYNTH_PHRASE_TEXT = "Translate into Chinese: Beijing is the Capital of China" + # The new-schema CI overlay bakes in async_chunk: False and covers CUDA/ROCm/XPU # via its ``platforms:`` section, so one path serves all three. default_stage_config = get_deploy_config_path("ci/qwen3_omni_moe.yaml") - -def _realtime_stage_config_path() -> str: - """CI omni layout without async_chunk; stage 0 thinker max_tokens=10.""" - return modify_stage_config( - default_stage_config, - updates={"stages": {0: {"default_sampling_params.max_tokens": 10}}}, - ) - - realtime_server_params = [ pytest.param( OmniServerParams( model=MODEL, - stage_config_path=_realtime_stage_config_path(), + stage_config_path=default_stage_config, use_stage_cli=True, + server_args=["--no-async-chunk"], ), - id="thinker_max_tokens_10", + id="default", ), ] @@ -169,9 +164,15 @@ class TestQwen3OmniRealtimeWebSocket: """ Short streamed 16 kHz mono PCM16 input; expect streamed PCM16 audio deltas and transcription. Verify Whisper(output audio) aligns with model text (same idea - as multimodal omni e2e). + as multimodal omni e2e). Input speech is synthesized from + ``REALTIME_SYNTH_PHRASE_TEXT``. """ - syn = generate_synthetic_audio(10, 1, sample_rate=16000) + syn = generate_synthetic_audio( + 10, + 1, + sample_rate=16000, + phrase_text=REALTIME_SYNTH_PHRASE_TEXT, + ) wav_bytes = base64.b64decode(syn["base64"]) pcm16 = _pcm16_mono_16k_from_wav_bytes(wav_bytes) diff --git a/tests/entrypoints/openai_api/test_serving_chat_metrics.py b/tests/entrypoints/openai_api/test_serving_chat_metrics.py index d25af6c38..0647c40a3 100644 --- a/tests/entrypoints/openai_api/test_serving_chat_metrics.py +++ b/tests/entrypoints/openai_api/test_serving_chat_metrics.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 """Unit tests for OmniChatCompletionResponse/StreamResponse metrics field.""" +from types import SimpleNamespace + import pytest pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -50,3 +52,36 @@ def test_omni_chat_completion_stream_response_metrics(): ) assert response.modality == "audio" assert response.metrics == {"stage_latency": 0.5} + + +def test_create_image_choice_exposes_diffusion_metrics(): + """Ensure image chat content exposes profiler metrics for clients.""" + from PIL import Image + + from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat + + stage_durations = {"prefill": 0.12, "diffusion": 1.23} + peak_memory_mb = 3210.5 + omni_outputs = SimpleNamespace( + request_output=None, + stage_durations=stage_durations, + peak_memory_mb=peak_memory_mb, + images=[Image.new("RGB", (2, 2), color=(255, 0, 0))], + ) + + choices = OmniOpenAIServingChat._create_image_choice( # type: ignore[misc] + None, + omni_outputs=omni_outputs, + role="assistant", + request=SimpleNamespace(return_token_ids=False), + ) + + assert len(choices) == 1 + content = choices[0].message.content + assert isinstance(content, list) + assert len(content) == 1 + first_item = content[0] + assert first_item["type"] == "image_url" + assert first_item["image_url"]["url"].startswith("data:image/png;base64,") + assert first_item["stage_durations"] == stage_durations + assert first_item["peak_memory_mb"] == peak_memory_mb diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index b388b1860..b78d62d9e 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -18,6 +18,7 @@ from pydantic import ValidationError from pytest_mock import MockerFixture from vllm.entrypoints.openai.engine.protocol import ErrorInfo, ErrorResponse +from vllm_omni.entrypoints.omni_base import OmniEngineDeadError from vllm_omni.entrypoints.openai import api_server as api_server_module from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin from vllm_omni.entrypoints.openai.protocol.audio import ( @@ -1707,6 +1708,77 @@ def test_api_server_create_speech_wraps_error_response_status(mocker: MockerFixt assert response.status_code == 400 +def test_api_server_create_speech_engine_error_response_includes_request_and_stage_id(mocker: MockerFixture): + handler = mocker.MagicMock() + handler.create_speech = mocker.AsyncMock( + side_effect=OmniEngineDeadError( + "engine dead", + error_stage_id=1, + ) + ) + + terminate_mock = mocker.patch.object(api_server_module, "terminate_if_errored") + + app = FastAPI() + app.state.args = SimpleNamespace(log_error_stack=False) + app.state.openai_serving_speech = handler + app.state.engine_client = SimpleNamespace( + engine=SimpleNamespace(is_alive=lambda: False), + errored=True, + ) + app.state.server = SimpleNamespace() + scope = { + "type": "http", + "app": app, + "method": "POST", + "path": "/v1/audio/speech", + "headers": [], + "query_string": b"", + "client": ("127.0.0.1", 12345), + "server": ("testserver", 80), + "scheme": "http", + } + raw_request = Request(scope) + raw_request.state.request_metadata = SimpleNamespace(request_id="speech-req-1") + request = OpenAICreateSpeechRequest(input="Hello") + + response = asyncio.run(api_server_module.create_speech(request, raw_request)) + + assert isinstance(response, JSONResponse) + assert response.status_code == 500 + assert response.body.decode("utf-8") == ( + '{"error":{"message":"engine dead","type":"InternalServerError","param":null,' + '"code":500,"request_id":"speech-req-1","error_stage_id":1}}' + ) + terminate_mock.assert_called_once() + + +def test_omni_engine_error_handler_includes_request_and_stage_id(mocker: MockerFixture): + app = FastAPI() + app.state.args = SimpleNamespace(log_error_stack=False) + app.state.engine_client = SimpleNamespace( + engine=SimpleNamespace(is_alive=lambda: False), + errored=True, + ) + app.state.server = SimpleNamespace() + + terminate_mock = mocker.patch.object(api_server_module, "terminate_if_errored") + api_server_module._register_omni_exception_handlers(app) + + @app.get("/boom") + async def boom(request: Request): + request.state.request_metadata = SimpleNamespace(request_id="speech-req-1") + exc = OmniEngineDeadError("engine dead", error_stage_id=1) + raise exc + + response = TestClient(app).get("/boom") + + assert response.status_code == 500 + assert response.json()["error"]["request_id"] == "speech-req-1" + assert response.json()["error"]["error_stage_id"] == 1 + terminate_mock.assert_called_once() + + class TestWAVHeaderGeneration: """Unit tests for WAV header generation with placeholder values.""" diff --git a/tests/entrypoints/openai_api/test_serving_video_stream.py b/tests/entrypoints/openai_api/test_serving_video_stream.py new file mode 100644 index 000000000..0787eb856 --- /dev/null +++ b/tests/entrypoints/openai_api/test_serving_video_stream.py @@ -0,0 +1,650 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the serving-layer streaming video WebSocket handler.""" + +from __future__ import annotations + +import asyncio +import base64 +import io +import json +import threading +from typing import Any + +import pytest +from PIL import Image + +from vllm_omni.entrypoints.openai import serving_video_stream, video_stream_envs +from vllm_omni.entrypoints.openai.serving_video_stream import ( + OmniStreamingVideoHandler, + StreamingVideoSessionConfig, +) +from vllm_omni.outputs import OmniRequestOutput + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _make_jpeg(r: int = 128, g: int = 128, b: int = 128) -> bytes: + img = Image.new("RGB", (64, 64), (r, g, b)) + buf = io.BytesIO() + img.save(buf, format="JPEG", quality=95) + return buf.getvalue() + + +def _b64(data: bytes) -> str: + return base64.b64encode(data).decode() + + +def _text_result(text: str) -> OmniRequestOutput: + class Output: + pass + + class RequestOutput: + pass + + output = Output() + output.text = text + request_output = RequestOutput() + request_output.outputs = [output] + return OmniRequestOutput(final_output_type="text", request_output=request_output) + + +def _audio_result(audio_data: Any) -> OmniRequestOutput: + class Output: + pass + + class RequestOutput: + pass + + output = Output() + output.multimodal_output = {"audio": audio_data} + request_output = RequestOutput() + request_output.outputs = [output] + return OmniRequestOutput(final_output_type="audio", request_output=request_output) + + +class MockWebSocket: + def __init__(self, messages: list[str] | None = None): + self._messages = list(messages or []) + self._idx = 0 + self.accepted = False + self.sent: list[dict[str, Any]] = [] + + async def accept(self): + self.accepted = True + + async def receive_text(self) -> str: + if self._idx >= len(self._messages): + await asyncio.sleep(999) + msg = self._messages[self._idx] + self._idx += 1 + return msg + + async def send_json(self, data: dict[str, Any]): + self.sent.append(data) + + +class TimedWebSocket: + def __init__(self): + self._q: asyncio.Queue[str] = asyncio.Queue() + self.accepted = False + self.sent: list[dict[str, Any]] = [] + + async def accept(self): + self.accepted = True + + async def receive_text(self) -> str: + return await self._q.get() + + async def send_json(self, data: dict[str, Any]): + self.sent.append(data) + + def put(self, msg: dict[str, Any]): + self._q.put_nowait(json.dumps(msg)) + + def sent_types(self) -> list[str]: + return [m.get("type", "") for m in self.sent] + + +def test_api_server_registers_video_stream_route(): + from vllm_omni.entrypoints.openai.api_server import router + + assert any(getattr(route, "path", None) == "/v1/video/chat/stream" for route in router.routes) + + +@pytest.mark.asyncio +async def test_receive_config_accepts_client_legacy_aliases(): + ws = MockWebSocket( + [ + json.dumps( + { + "type": "session.config", + "model": "test", + "num_sample_frames": 7, + "evs_enabled": False, + "evs_threshold": 0.87, + } + ) + ] + ) + handler = OmniStreamingVideoHandler(chat_service=object()) + + config = await handler._receive_config(ws) + + assert config is not None + assert config.num_frames == 7 + assert config.enable_frame_filter is False + assert config.frame_filter_threshold == 0.87 + + +@pytest.mark.asyncio +async def test_audio_in_video_sets_mm_processor_kwargs(): + captured_requests = [] + + class EmptyEngine: + def generate(self, **_kwargs): + async def _gen(): + if False: + yield None + + return _gen() + + class CapturingHandler(OmniStreamingVideoHandler): + async def _preprocess_to_engine_prompt(self, request): + captured_requests.append(request) + return {"prompt": "x"} + + ws = MockWebSocket() + handler = CapturingHandler(chat_service=object(), engine_client=EmptyEngine()) + config = StreamingVideoSessionConfig(model="test", modalities=["text", "audio"], use_audio_in_video=True) + + await handler._process_query_engine( + ws, + config, + [_b64(_make_jpeg())], + bytearray(b"\x00\x00"), + [], + "what is happening?", + "req-1", + asyncio.Event(), + {}, + ) + + assert captured_requests + assert captured_requests[0].mm_processor_kwargs == {"use_audio_in_video": True} + + +@pytest.mark.asyncio +async def test_audio_in_video_disabled_omits_mm_processor_kwargs(): + captured_requests = [] + + class EmptyEngine: + def generate(self, **_kwargs): + async def _gen(): + if False: + yield None + + return _gen() + + class CapturingHandler(OmniStreamingVideoHandler): + async def _preprocess_to_engine_prompt(self, request): + captured_requests.append(request) + return {"prompt": "x"} + + ws = MockWebSocket() + handler = CapturingHandler(chat_service=object(), engine_client=EmptyEngine()) + config = StreamingVideoSessionConfig(model="test", modalities=["text", "audio"], use_audio_in_video=False) + + await handler._process_query_engine( + ws, + config, + [_b64(_make_jpeg())], + bytearray(b"\x00\x00"), + [], + "what is happening?", + "req-1", + asyncio.Event(), + {}, + ) + + assert captured_requests + assert captured_requests[0].mm_processor_kwargs is None + + +@pytest.mark.asyncio +async def test_query_inline_audio_data_sets_mm_processor_kwargs(): + captured_requests = [] + + class EmptyEngine: + def generate(self, **_kwargs): + async def _gen(): + if False: + yield None + + return _gen() + + class CapturingHandler(OmniStreamingVideoHandler): + async def _preprocess_to_engine_prompt(self, request): + captured_requests.append(request) + return {"prompt": "x"} + + ws = MockWebSocket( + [ + json.dumps({"type": "session.config", "model": "test"}), + json.dumps({"type": "video.frame", "data": _b64(_make_jpeg())}), + json.dumps( + { + "type": "video.query", + "text": "describe", + "audio_data": _b64(b"\x00\x00"), + } + ), + json.dumps({"type": "video.done"}), + ] + ) + handler = CapturingHandler(chat_service=object(), engine_client=EmptyEngine(), idle_timeout=2.0) + + await handler.handle_session(ws) + + assert captured_requests + assert captured_requests[0].mm_processor_kwargs == {"use_audio_in_video": True} + assert "session.done" in [m.get("type") for m in ws.sent] + + +def test_audio_delta_mode_is_read_by_serving_code_at_runtime(monkeypatch): + handler = OmniStreamingVideoHandler(chat_service=object()) + result = _audio_result([object()]) + + monkeypatch.setattr( + OmniStreamingVideoHandler, + "_delta_fast", + classmethod(lambda cls, audio_data, chunks_drained: ("fast-path", chunks_drained)), + ) + monkeypatch.setattr( + OmniStreamingVideoHandler, + "_delta_slow", + classmethod(lambda cls, audio_data, chunks_drained: ("slow-path", chunks_drained)), + ) + + monkeypatch.setenv("VLLM_VIDEO_AUDIO_DELTA_MODE", "fast") + assert handler._extract_audio_delta_b64(result, 0)[0] == "fast-path" + + monkeypatch.setenv("VLLM_VIDEO_AUDIO_DELTA_MODE", "slow") + assert handler._extract_audio_delta_b64(result, 0)[0] == "slow-path" + + +def test_video_stream_envs_strip_and_warn_once_per_invalid_value(monkeypatch): + warnings = [] + + video_stream_envs._warned_invalid_envs.clear() + try: + monkeypatch.setattr( + video_stream_envs.logger, + "warning", + lambda message, *args, **_kwargs: warnings.append((message, args)), + ) + + monkeypatch.setenv("VLLM_VIDEO_ASYNC_CHUNK", " off ") + assert video_stream_envs.VLLM_VIDEO_ASYNC_CHUNK == "off" + assert not warnings + + monkeypatch.setenv("VLLM_VIDEO_ASYNC_CHUNK", "bad") + assert video_stream_envs.VLLM_VIDEO_ASYNC_CHUNK == "on" + assert video_stream_envs.VLLM_VIDEO_ASYNC_CHUNK == "on" + assert len(warnings) == 1 + + monkeypatch.setenv("VLLM_VIDEO_ASYNC_CHUNK", "still_bad") + assert video_stream_envs.VLLM_VIDEO_ASYNC_CHUNK == "on" + assert len(warnings) == 2 + finally: + video_stream_envs._warned_invalid_envs.clear() + + +@pytest.mark.asyncio +async def test_async_chunk_mode_is_read_by_engine_path_at_runtime(monkeypatch): + class TextEngine: + def generate(self, **_kwargs): + async def _gen(): + yield _text_result("hello") + + return _gen() + + class CapturingHandler(OmniStreamingVideoHandler): + async def _preprocess_to_engine_prompt(self, request): + return {"prompt": "x"} + + handler = CapturingHandler(chat_service=object(), engine_client=TextEngine()) + config = StreamingVideoSessionConfig(model="test", modalities=["text"]) + + monkeypatch.setenv("VLLM_VIDEO_ASYNC_CHUNK", "on") + ws_on = MockWebSocket() + await handler._process_query_engine( + ws_on, + config, + [_b64(_make_jpeg())], + bytearray(), + [], + "describe", + "req-on", + asyncio.Event(), + {}, + ) + assert {"type": "response.text.delta", "delta": "hello"} in ws_on.sent + + monkeypatch.setenv("VLLM_VIDEO_ASYNC_CHUNK", "off") + ws_off = MockWebSocket() + await handler._process_query_engine( + ws_off, + config, + [_b64(_make_jpeg())], + bytearray(), + [], + "describe", + "req-off", + asyncio.Event(), + {}, + ) + assert {"type": "response.text.done", "text": "hello"} in ws_off.sent + assert not any(m.get("type") == "response.text.delta" for m in ws_off.sent) + + +@pytest.mark.asyncio +async def test_query_without_engine_client_sends_error(): + ws = MockWebSocket() + handler = OmniStreamingVideoHandler(chat_service=object(), engine_client=None) + + await handler._process_query( + ws, + StreamingVideoSessionConfig(model="test"), + [], + bytearray(), + [], + "describe", + "req-1", + asyncio.Event(), + {}, + ) + + assert {"type": "error", "message": "Streaming video requires an engine client"} in ws.sent + + +@pytest.mark.asyncio +async def test_new_query_cancels_in_flight_query(): + query_started = asyncio.Event() + query_cancelled = asyncio.Event() + calls = 0 + + class BlockingHandler(OmniStreamingVideoHandler): + async def _process_query(self, *args, **kwargs): + nonlocal calls + calls += 1 + if calls > 1: + return + query_started.set() + try: + await asyncio.sleep(999) + except asyncio.CancelledError: + query_cancelled.set() + raise + + ws = TimedWebSocket() + handler = BlockingHandler(chat_service=object(), idle_timeout=5.0) + task = asyncio.create_task(handler.handle_session(ws)) + + ws.put({"type": "session.config", "model": "test"}) + await asyncio.sleep(0) + ws.put({"type": "video.frame", "data": _b64(_make_jpeg())}) + await asyncio.sleep(0) + ws.put({"type": "video.query", "text": "describe"}) + await asyncio.wait_for(query_started.wait(), timeout=2.0) + + ws.put({"type": "video.query", "text": "interrupt"}) + await asyncio.wait_for(query_cancelled.wait(), timeout=2.0) + ws.put({"type": "video.done"}) + + await asyncio.wait_for(task, timeout=2.0) + assert "session.done" in ws.sent_types() + + +@pytest.mark.asyncio +async def test_video_done_waits_for_in_flight_query(): + query_started = asyncio.Event() + allow_finish = asyncio.Event() + query_finished = asyncio.Event() + + class BlockingHandler(OmniStreamingVideoHandler): + async def _process_query(self, *args, **kwargs): + query_started.set() + await allow_finish.wait() + query_finished.set() + + ws = TimedWebSocket() + handler = BlockingHandler(chat_service=object(), idle_timeout=5.0) + task = asyncio.create_task(handler.handle_session(ws)) + + ws.put({"type": "session.config", "model": "test"}) + await asyncio.sleep(0) + ws.put({"type": "video.frame", "data": _b64(_make_jpeg())}) + await asyncio.sleep(0) + ws.put({"type": "video.query", "text": "describe"}) + await asyncio.wait_for(query_started.wait(), timeout=2.0) + + ws.put({"type": "video.done"}) + await asyncio.sleep(0.05) + assert not task.done() + assert not query_finished.is_set() + + allow_finish.set() + await asyncio.wait_for(task, timeout=2.0) + + assert query_finished.is_set() + assert "session.done" in ws.sent_types() + + +@pytest.mark.asyncio +async def test_frame_prewarm_does_not_block_following_query(monkeypatch): + decode_started = threading.Event() + release_decode = threading.Event() + query_started = asyncio.Event() + + def blocked_decode(raw_bytes: bytes): + decode_started.set() + release_decode.wait(timeout=2.0) + return Image.open(io.BytesIO(raw_bytes)).convert("RGB") + + class BlockingHandler(OmniStreamingVideoHandler): + async def _process_query(self, *args, **kwargs): + query_started.set() + + monkeypatch.setattr(serving_video_stream, "_decode_frame_bytes", blocked_decode) + + ws = TimedWebSocket() + handler = BlockingHandler(chat_service=object(), idle_timeout=5.0) + task = asyncio.create_task(handler.handle_session(ws)) + + ws.put({"type": "session.config", "model": "test"}) + await asyncio.sleep(0) + ws.put({"type": "video.frame", "data": _b64(_make_jpeg())}) + + for _ in range(100): + if decode_started.is_set(): + break + await asyncio.sleep(0.01) + assert decode_started.is_set() + + ws.put({"type": "video.query", "text": "describe"}) + await asyncio.wait_for(query_started.wait(), timeout=2.0) + + release_decode.set() + ws.put({"type": "video.done"}) + await asyncio.wait_for(task, timeout=2.0) + assert "session.done" in ws.sent_types() + + +@pytest.mark.asyncio +async def test_client_cannot_send_internal_frame_decode_failed_message(): + captured_frames: list[list[str]] = [] + frame = _b64(_make_jpeg()) + + class CapturingHandler(OmniStreamingVideoHandler): + async def _process_query( + self, + websocket, + config, + frame_buffer, + audio_buffer, + message_history, + query_text, + request_id, + interrupt_event, + prewarmed_frames, + ): + captured_frames.append(list(frame_buffer)) + + ws = TimedWebSocket() + handler = CapturingHandler(chat_service=object(), idle_timeout=5.0) + task = asyncio.create_task(handler.handle_session(ws)) + + ws.put({"type": "session.config", "model": "test"}) + await asyncio.sleep(0) + ws.put({"type": "video.frame", "data": frame}) + await asyncio.sleep(0) + ws.put({"type": "_internal.frame_decode_failed", "b64": frame}) + await asyncio.sleep(0) + ws.put({"type": "video.query", "text": "describe"}) + await asyncio.sleep(0) + ws.put({"type": "video.done"}) + await asyncio.wait_for(task, timeout=2.0) + + assert {"type": "error", "message": "Unknown type: _internal.frame_decode_failed"} in ws.sent + assert captured_frames == [[frame]] + + +@pytest.mark.asyncio +async def test_failed_frame_prewarm_removes_frame_before_query(): + ws = TimedWebSocket() + handler = OmniStreamingVideoHandler(chat_service=object(), idle_timeout=5.0) + task = asyncio.create_task(handler.handle_session(ws)) + + ws.put({"type": "session.config", "model": "test", "enable_frame_filter": False}) + await asyncio.sleep(0) + ws.put({"type": "video.frame", "data": _b64(b"not-a-jpeg")}) + + for _ in range(100): + if any(m.get("message") == "Frame decode failed" for m in ws.sent): + break + await asyncio.sleep(0.01) + + assert {"type": "error", "message": "Frame decode failed"} in ws.sent + + ws.put({"type": "video.query", "text": "describe"}) + await asyncio.sleep(0) + ws.put({"type": "video.done"}) + await asyncio.wait_for(task, timeout=2.0) + + assert {"type": "error", "message": "No frames buffered"} in ws.sent + + +@pytest.mark.asyncio +async def test_frame_filter_error_sends_invalid_image(monkeypatch): + def fail_should_retain(self, frame_jpeg): + raise ValueError("decode failed") + + monkeypatch.setattr(serving_video_stream.FrameSimilarityFilter, "should_retain", fail_should_retain) + + ws = TimedWebSocket() + handler = OmniStreamingVideoHandler(chat_service=object(), idle_timeout=5.0) + task = asyncio.create_task(handler.handle_session(ws)) + + ws.put({"type": "session.config", "model": "test"}) + await asyncio.sleep(0) + ws.put({"type": "video.frame", "data": _b64(_make_jpeg())}) + await asyncio.sleep(0) + ws.put({"type": "video.done"}) + await asyncio.wait_for(task, timeout=2.0) + + assert {"type": "error", "message": "Invalid image data"} in ws.sent + assert "session.done" in ws.sent_types() + + +@pytest.mark.asyncio +async def test_audio_buffer_overflow_clears_buffer_before_query(monkeypatch): + captured_audio_lengths: list[int] = [] + + class EmptyEngine: + def generate(self, **_kwargs): + async def _gen(): + if False: + yield None + + return _gen() + + class CapturingHandler(OmniStreamingVideoHandler): + async def _process_query_engine( + self, + websocket, + config, + frame_buffer, + audio_buffer, + message_history, + query_text, + request_id, + interrupt_event, + prewarmed_frames, + ): + captured_audio_lengths.append(len(audio_buffer)) + + monkeypatch.setattr(serving_video_stream, "_MAX_AUDIO_BUFFER_BYTES", 4) + + ws = TimedWebSocket() + handler = CapturingHandler(chat_service=object(), engine_client=EmptyEngine(), idle_timeout=5.0) + task = asyncio.create_task(handler.handle_session(ws)) + + ws.put({"type": "session.config", "model": "test"}) + await asyncio.sleep(0) + ws.put({"type": "audio.chunk", "data": _b64(b"1234")}) + await asyncio.sleep(0) + ws.put({"type": "audio.chunk", "data": _b64(b"5")}) + await asyncio.sleep(0) + ws.put({"type": "video.frame", "data": _b64(_make_jpeg())}) + await asyncio.sleep(0) + ws.put({"type": "video.query", "text": "describe"}) + await asyncio.sleep(0) + ws.put({"type": "video.done"}) + await asyncio.wait_for(task, timeout=2.0) + + assert {"type": "error", "message": "Audio buffer overflow"} in ws.sent + assert captured_audio_lengths == [0] + + +def test_build_messages_keeps_recent_history_text_only(): + handler = OmniStreamingVideoHandler(chat_service=object()) + old_frame = _b64(_make_jpeg(1, 2, 3)) + current_frame = _b64(_make_jpeg(4, 5, 6)) + history = [ + {"role": "user", "content": [{"type": "text", "text": "old question"}]}, + {"role": "assistant", "content": "old answer"}, + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{old_frame}"}}, + {"type": "input_audio", "input_audio": {"data": "ignored", "format": "wav"}}, + {"type": "text", "text": "recent question"}, + ], + }, + {"role": "assistant", "content": "recent answer"}, + ] + + messages, user_message = handler._build_messages( + StreamingVideoSessionConfig(model="test", num_frames=1), + [current_frame], + bytearray(), + history, + "current question", + {}, + ) + + assert messages[0] == {"role": "user", "content": "recent question"} + assert messages[1] == {"role": "assistant", "content": "recent answer"} + assert messages[2] == user_message + assert user_message["content"][-1] == {"type": "text", "text": "current question"} diff --git a/tests/entrypoints/openai_api/test_video_frame_filter.py b/tests/entrypoints/openai_api/test_video_frame_filter.py new file mode 100644 index 000000000..a734aec75 --- /dev/null +++ b/tests/entrypoints/openai_api/test_video_frame_filter.py @@ -0,0 +1,89 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for FrameSimilarityFilter (Phase 2 EVS).""" + +from __future__ import annotations + +import pytest + +from tests.entrypoints.openai_api.conftest_video import ( + make_gradient_jpeg, + make_jpeg, +) +from vllm_omni.entrypoints.openai.video_frame_filter import FrameSimilarityFilter + + +class TestFrameSimilarityFilter: + def test_first_frame_always_retained(self): + f = FrameSimilarityFilter(threshold=0.99) + assert f.should_retain(make_jpeg()) is True + + def test_identical_frames_dropped(self): + f = FrameSimilarityFilter(threshold=0.90) + frame = make_jpeg(100, 100, 100) + assert f.should_retain(frame) is True + assert f.should_retain(frame) is False + + def test_very_different_frames_retained(self): + f = FrameSimilarityFilter(threshold=0.95) + assert f.should_retain(make_jpeg(255, 255, 255)) is True + assert f.should_retain(make_jpeg(0, 0, 0)) is True + + def test_low_threshold_keeps_slightly_different(self): + f = FrameSimilarityFilter(threshold=0.50) + assert f.should_retain(make_jpeg(100, 100, 100)) is True + assert f.should_retain(make_jpeg(110, 110, 110)) is True + + def test_random_frames_all_retained(self): + f = FrameSimilarityFilter(threshold=0.95) + for i in range(5): + assert f.should_retain(make_gradient_jpeg(seed=i)) is True + + def test_reset_clears_state(self): + f = FrameSimilarityFilter(threshold=0.90) + frame = make_jpeg(128, 128, 128) + assert f.should_retain(frame) is True + assert f.should_retain(frame) is False + f.reset() + assert f.should_retain(frame) is True + + def test_stats_counting(self): + f = FrameSimilarityFilter(threshold=0.90) + frame = make_jpeg(50, 50, 50) + f.should_retain(frame) # retained + f.should_retain(frame) # dropped + f.should_retain(frame) # dropped + stats = f.stats + assert stats["retained_count"] == 1 + assert stats["dropped_count"] == 2 + assert stats["total_count"] == 3 + assert abs(stats["drop_rate"] - 2.0 / 3.0) < 1e-6 + + def test_stats_empty(self): + stats = FrameSimilarityFilter().stats + assert stats["total_count"] == 0 + assert stats["drop_rate"] == 0.0 + + def test_stats_reset(self): + f = FrameSimilarityFilter(threshold=0.90) + f.should_retain(make_jpeg()) + f.reset() + assert f.stats["total_count"] == 0 + + def test_invalid_threshold(self): + with pytest.raises(ValueError, match="threshold"): + FrameSimilarityFilter(threshold=1.5) + with pytest.raises(ValueError, match="threshold"): + FrameSimilarityFilter(threshold=-0.1) + + def test_invalid_thumbnail_size(self): + with pytest.raises(ValueError, match="thumbnail_size"): + FrameSimilarityFilter(thumbnail_size=0) + + def test_different_image_sizes_same_colour(self): + """Filter should handle frames of varying resolutions.""" + f = FrameSimilarityFilter(threshold=0.90) + small = make_jpeg(100, 100, 100, size=32) + large = make_jpeg(100, 100, 100, size=256) + assert f.should_retain(small) is True + assert f.should_retain(large) is False diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py index 1cc561665..a29f4493c 100644 --- a/tests/entrypoints/openai_api/test_video_server.py +++ b/tests/entrypoints/openai_api/test_video_server.py @@ -29,6 +29,7 @@ from vllm_omni.entrypoints.openai.protocol.videos import ( from vllm_omni.entrypoints.openai.serving_video import OmniOpenAIServingVideo from vllm_omni.entrypoints.openai.storage import LocalStorageManager from vllm_omni.entrypoints.openai.stores import AsyncDictStore, TaskRegistry +from vllm_omni.inputs.data import OmniDiffusionSamplingParams pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -60,6 +61,7 @@ class MockVideoResult: class FakeAsyncOmni: def __init__(self): self.stage_configs = [SimpleNamespace(stage_type="diffusion")] + self.default_sampling_params_list = [OmniDiffusionSamplingParams()] self.captured_prompt = None self.captured_sampling_params_list = None @@ -442,6 +444,87 @@ def test_frame_interpolation_params_pass_to_diffusion_sampling_params(test_clien assert captured.frame_interpolation_model_path == "local-rife" +def test_default_sampling_params_apply_to_video_requests(test_client, mocker: MockerFixture): + mocker.patch( + "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes", + return_value=b"fake-video", + ) + engine = test_client.app.state.openai_serving_video._engine_client + engine.default_sampling_params_list = [ + OmniDiffusionSamplingParams( + num_inference_steps=4, + guidance_scale=7.5, + generator_device="cpu", + enable_frame_interpolation=True, + frame_interpolation_exp=2, + frame_interpolation_scale=0.5, + frame_interpolation_model_path="default-rife", + ) + ] + + response = test_client.post( + "/v1/videos", + data={ + "prompt": "default param pass-through", + }, + ) + + assert response.status_code == 200 + video_id = response.json()["id"] + _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value) + + captured = engine.captured_sampling_params_list[0] + assert captured.num_inference_steps == 4 + assert captured.guidance_scale == 7.5 + assert captured.generator_device == "cpu" + assert captured.enable_frame_interpolation is True + assert captured.frame_interpolation_exp == 2 + assert captured.frame_interpolation_scale == 0.5 + assert captured.frame_interpolation_model_path == "default-rife" + + +def test_request_params_override_default_video_sampling_params(test_client, mocker: MockerFixture): + mocker.patch( + "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes", + return_value=b"fake-video", + ) + engine = test_client.app.state.openai_serving_video._engine_client + engine.default_sampling_params_list = [ + OmniDiffusionSamplingParams( + num_inference_steps=4, + guidance_scale=7.5, + enable_frame_interpolation=True, + frame_interpolation_exp=2, + frame_interpolation_scale=0.5, + frame_interpolation_model_path="default-rife", + ) + ] + + response = test_client.post( + "/v1/videos", + data={ + "prompt": "explicit override", + "num_inference_steps": "8", + "enable_frame_interpolation": "false", + "frame_interpolation_exp": "1", + "frame_interpolation_scale": "1.0", + "frame_interpolation_model_path": "custom-rife", + }, + ) + + assert response.status_code == 200 + video_id = response.json()["id"] + _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value) + + captured = engine.captured_sampling_params_list[0] + assert captured.num_inference_steps == 8 + assert captured.guidance_scale == 7.5 + assert captured.enable_frame_interpolation is False + assert captured.frame_interpolation_exp == 1 + assert captured.frame_interpolation_scale == 1.0 + assert captured.frame_interpolation_model_path == "custom-rife" + + def test_worker_fps_multiplier_is_applied_to_async_encoding(test_client, mocker: MockerFixture): fps_values = [] engine = test_client.app.state.openai_serving_video._engine_client @@ -1149,6 +1232,39 @@ def test_sync_frame_interpolation_params_pass_to_sampling_params(test_client, mo assert kwargs["fps"] == 8 +def test_sync_default_sampling_params_apply_to_video_requests(test_client, mocker: MockerFixture): + _mock_encode_video_bytes(mocker) + engine = test_client.app.state.openai_serving_video._engine_client + engine.default_sampling_params_list = [ + OmniDiffusionSamplingParams( + num_inference_steps=4, + guidance_scale=7.5, + enable_frame_interpolation=True, + frame_interpolation_exp=2, + frame_interpolation_scale=0.5, + frame_interpolation_model_path="default-rife", + ) + ] + + response = test_client.post( + "/v1/videos/sync", + data={ + "prompt": "sync default param pass-through", + "fps": "8", + }, + ) + + assert response.status_code == 200 + engine = test_client.app.state.openai_serving_video._engine_client + captured = engine.captured_sampling_params_list[0] + assert captured.num_inference_steps == 4 + assert captured.guidance_scale == 7.5 + assert captured.enable_frame_interpolation is True + assert captured.frame_interpolation_exp == 2 + assert captured.frame_interpolation_scale == 0.5 + assert captured.frame_interpolation_model_path == "default-rife" + + def test_worker_fps_multiplier_is_applied_to_sync_encoding(test_client, mocker: MockerFixture): engine = test_client.app.state.openai_serving_video._engine_client fps_values = [] diff --git a/tests/entrypoints/openai_api/test_video_stream_handler.py b/tests/entrypoints/openai_api/test_video_stream_handler.py new file mode 100644 index 000000000..bee5e0de3 --- /dev/null +++ b/tests/entrypoints/openai_api/test_video_stream_handler.py @@ -0,0 +1,593 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Runtime behavior tests for VideoStreamHandler. + +Every test creates a mock WebSocket, drives handle_session through a +specific code path, and asserts on the JSON messages actually sent back. +No inspect.getsource() tricks — these test real async execution. +""" + +from __future__ import annotations + +import asyncio +import base64 +import io +import json +from unittest.mock import AsyncMock + +import pytest +from PIL import Image + +from vllm_omni.entrypoints.openai.video_stream_session import ( + VideoStreamHandler, +) + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_jpeg(r: int = 128, g: int = 128, b: int = 128) -> bytes: + img = Image.new("RGB", (64, 64), (r, g, b)) + buf = io.BytesIO() + img.save(buf, format="JPEG", quality=95) + return buf.getvalue() + + +def _b64(data: bytes) -> str: + return base64.b64encode(data).decode() + + +def _config_msg(**overrides) -> str: + msg = {"type": "session.config", "model": "test", "evs_enabled": False} + msg.update(overrides) + return json.dumps(msg) + + +def _frame_msg(jpeg: bytes | None = None) -> str: + if jpeg is None: + jpeg = _make_jpeg() + return json.dumps({"type": "video.frame", "data": _b64(jpeg)}) + + +def _query_msg(text: str = "What do you see?") -> str: + return json.dumps({"type": "video.query", "text": text}) + + +def _done_msg() -> str: + return json.dumps({"type": "video.done"}) + + +def _audio_msg(pcm: bytes = b"\x00" * 320) -> str: + return json.dumps({"type": "audio.chunk", "data": _b64(pcm)}) + + +def _make_sse_line(content: str) -> str: + payload = {"choices": [{"delta": {"content": content}}]} + return f"data: {json.dumps(payload)}\n\n" + + +class MockWebSocket: + """Async-compatible mock WebSocket that feeds messages from a list.""" + + def __init__(self, messages: list[str]): + self._messages = list(messages) + self._idx = 0 + self._sent: list[dict | bytes] = [] + self._accepted = False + + async def accept(self): + self._accepted = True + + async def receive_text(self) -> str: + if self._idx >= len(self._messages): + # Simulate connection close by hanging forever (will be timed out) + await asyncio.sleep(999) + msg = self._messages[self._idx] + self._idx += 1 + return msg + + async def send_json(self, data: dict): + self._sent.append(data) + + async def send_bytes(self, data: bytes): + self._sent.append(data) + + @property + def sent_messages(self) -> list[dict]: + return [m for m in self._sent if isinstance(m, dict)] + + def sent_types(self) -> list[str]: + return [m.get("type", "") for m in self.sent_messages] + + +class MockChatHandler: + """Mock OmniOpenAIServingChat for testing _handle_query paths.""" + + def __init__(self, response=None): + self._response = response + + async def create_chat_completion(self, request, raw_request=None): + if self._response is not None: + return self._response + + # Default: return a simple streaming generator + async def _gen(): + yield _make_sse_line("Hello") + yield _make_sse_line("World") + yield "data: [DONE]\n\n" + + return _gen() + + +# --------------------------------------------------------------------------- +# _receive_config path tests +# --------------------------------------------------------------------------- + + +class TestReceiveConfig: + @pytest.mark.asyncio + async def test_config_timeout(self): + """No message within config_timeout → error sent, session ends.""" + ws = MockWebSocket([]) # no messages — will hang + handler = VideoStreamHandler(chat_handler=MockChatHandler(), config_timeout=0.05) + await handler.handle_session(ws) + + assert ws._accepted + errors = [m for m in ws.sent_messages if m["type"] == "error"] + assert len(errors) == 1 + assert "Timeout" in errors[0]["message"] + + @pytest.mark.asyncio + async def test_config_invalid_json(self): + """Non-JSON config message → error, session ends.""" + ws = MockWebSocket(["not json"]) + handler = VideoStreamHandler(chat_handler=MockChatHandler(), config_timeout=1.0) + await handler.handle_session(ws) + + errors = [m for m in ws.sent_messages if m["type"] == "error"] + assert any("Invalid JSON" in e["message"] for e in errors) + + @pytest.mark.asyncio + async def test_config_wrong_type(self): + """Message with wrong type → error, session ends.""" + ws = MockWebSocket([json.dumps({"type": "wrong"})]) + handler = VideoStreamHandler(chat_handler=MockChatHandler(), config_timeout=1.0) + await handler.handle_session(ws) + + errors = [m for m in ws.sent_messages if m["type"] == "error"] + assert any("Expected session.config" in e["message"] for e in errors) + + @pytest.mark.asyncio + async def test_config_invalid_field_type(self): + """Config with wrong field type → error, session ends.""" + ws = MockWebSocket( + [ + json.dumps( + { + "type": "session.config", + "model": "test", + "max_frames": "potato", + } + ) + ] + ) + handler = VideoStreamHandler(chat_handler=MockChatHandler(), config_timeout=1.0) + await handler.handle_session(ws) + + errors = [m for m in ws.sent_messages if m["type"] == "error"] + assert any("max_frames" in e["message"] for e in errors) + + +# --------------------------------------------------------------------------- +# Normal session flow +# --------------------------------------------------------------------------- + + +class TestNormalFlow: + @pytest.mark.asyncio + async def test_frame_query_done(self): + """Happy path: config → frame → query → done → session.done.""" + ws = MockWebSocket( + [ + _config_msg(), + _frame_msg(), + _query_msg("Describe."), + _done_msg(), + ] + ) + handler = VideoStreamHandler(chat_handler=MockChatHandler(), idle_timeout=2.0) + await handler.handle_session(ws) + + types = ws.sent_types() + assert "response.start" in types + assert "response.text.delta" in types + assert "response.text.done" in types + assert "session.done" in types + + # Check text content was streamed + deltas = [m["delta"] for m in ws.sent_messages if m.get("type") == "response.text.delta"] + assert "Hello" in deltas + assert "World" in deltas + + # response.text.done has full text + done_msg = next(m for m in ws.sent_messages if m.get("type") == "response.text.done") + assert done_msg["text"] == "HelloWorld" + + @pytest.mark.asyncio + async def test_multiple_frames_before_query(self): + """Multiple frames accumulate, query uses sampled frames.""" + frames = [_frame_msg(_make_jpeg(r=i * 50)) for i in range(5)] + ws = MockWebSocket( + [ + _config_msg(num_sample_frames=3), + *frames, + _query_msg(), + _done_msg(), + ] + ) + handler = VideoStreamHandler(chat_handler=MockChatHandler(), idle_timeout=2.0) + await handler.handle_session(ws) + assert "session.done" in ws.sent_types() + + @pytest.mark.asyncio + async def test_evs_stats_in_done(self): + """EVS stats sent before session.done when evs_enabled=True.""" + frame = _frame_msg(_make_jpeg(100, 100, 100)) + ws = MockWebSocket( + [ + _config_msg(evs_enabled=True, evs_threshold=0.90), + frame, + frame, # second should be dropped + _done_msg(), + ] + ) + handler = VideoStreamHandler(chat_handler=MockChatHandler(), idle_timeout=2.0) + await handler.handle_session(ws) + + evs_msgs = [m for m in ws.sent_messages if m.get("type") == "response.evs_stats"] + assert len(evs_msgs) == 1 + assert evs_msgs[0]["retained_count"] == 1 + assert evs_msgs[0]["dropped_count"] == 1 + + +# --------------------------------------------------------------------------- +# _handle_query error paths +# --------------------------------------------------------------------------- + + +class TestHandleQueryErrors: + @pytest.mark.asyncio + async def test_query_with_no_frames(self): + """Query before any frames → error, session continues.""" + ws = MockWebSocket( + [ + _config_msg(), + _query_msg("Hello?"), + _done_msg(), + ] + ) + handler = VideoStreamHandler(chat_handler=MockChatHandler(), idle_timeout=2.0) + await handler.handle_session(ws) + + errors = [m for m in ws.sent_messages if m["type"] == "error"] + assert any("No frames" in e["message"] for e in errors) + # Session still ends cleanly + assert "session.done" in ws.sent_types() + + @pytest.mark.asyncio + async def test_query_returns_error_response(self): + """create_chat_completion returns ErrorResponse → error forwarded.""" + from vllm.entrypoints.openai.engine.protocol import ErrorResponse + + err = ErrorResponse( + message="model overloaded", + type="server_error", + code=503, + ) + + chat = AsyncMock() + chat.create_chat_completion = AsyncMock(return_value=err) + + ws = MockWebSocket( + [ + _config_msg(), + _frame_msg(), + _query_msg(), + _done_msg(), + ] + ) + handler = VideoStreamHandler(chat_handler=chat, idle_timeout=2.0) + await handler.handle_session(ws) + + # Should see response.start, then error, then response.text.done + assert "response.start" in ws.sent_types() + assert "response.text.done" in ws.sent_types() + errors = [m for m in ws.sent_messages if m["type"] == "error"] + assert any("model overloaded" in e["message"] for e in errors) + + @pytest.mark.asyncio + async def test_query_generator_raises(self): + """Exception during streaming → error sent, session continues.""" + + async def _exploding_gen(): + yield _make_sse_line("partial") + raise RuntimeError("CUDA OOM") + + chat = MockChatHandler() + chat._response = None # override below + + class BoomChat: + async def create_chat_completion(self, req, raw_request=None): + return _exploding_gen() + + ws = MockWebSocket( + [ + _config_msg(), + _frame_msg(), + _query_msg(), + _done_msg(), + ] + ) + handler = VideoStreamHandler(chat_handler=BoomChat(), idle_timeout=2.0) + await handler.handle_session(ws) + + errors = [m for m in ws.sent_messages if m["type"] == "error"] + assert any("Query processing failed" in e["message"] for e in errors) + # Crucially: no raw exception message leaked + assert not any("CUDA OOM" in e["message"] for e in errors) + # Protocol: response.text.done always sent after response.start + assert "response.text.done" in ws.sent_types() + # Session still ends cleanly + assert "session.done" in ws.sent_types() + + @pytest.mark.asyncio + async def test_empty_query_text_rejected(self): + """video.query with empty text → error, session continues.""" + ws = MockWebSocket( + [ + _config_msg(), + _frame_msg(), + json.dumps({"type": "video.query", "text": ""}), + _done_msg(), + ] + ) + handler = VideoStreamHandler(chat_handler=MockChatHandler(), idle_timeout=2.0) + await handler.handle_session(ws) + + errors = [m for m in ws.sent_messages if m["type"] == "error"] + assert any("non-empty" in e["message"] for e in errors) + assert "session.done" in ws.sent_types() + + +# --------------------------------------------------------------------------- +# Timeout paths +# --------------------------------------------------------------------------- + + +class TestTimeouts: + @pytest.mark.asyncio + async def test_idle_timeout(self): + """No message after config within idle_timeout → error, session ends.""" + ws = MockWebSocket([_config_msg()]) # config then nothing + handler = VideoStreamHandler(chat_handler=MockChatHandler(), idle_timeout=0.05) + await handler.handle_session(ws) + + errors = [m for m in ws.sent_messages if m["type"] == "error"] + assert any("Idle timeout" in e["message"] for e in errors) + + +# --------------------------------------------------------------------------- +# Reader error paths +# --------------------------------------------------------------------------- + + +class TestReaderErrors: + @pytest.mark.asyncio + async def test_invalid_json_mid_session(self): + """Invalid JSON mid-session → error sent, session continues.""" + ws = MockWebSocket( + [ + _config_msg(), + "{{bad json", + _done_msg(), + ] + ) + handler = VideoStreamHandler(chat_handler=MockChatHandler(), idle_timeout=2.0) + await handler.handle_session(ws) + + errors = [m for m in ws.sent_messages if m["type"] == "error"] + assert any("Invalid JSON" in e["message"] for e in errors) + assert "session.done" in ws.sent_types() + + @pytest.mark.asyncio + async def test_non_dict_message(self): + """Non-object JSON → error sent, continues.""" + ws = MockWebSocket( + [ + _config_msg(), + json.dumps([1, 2, 3]), + _done_msg(), + ] + ) + handler = VideoStreamHandler(chat_handler=MockChatHandler(), idle_timeout=2.0) + await handler.handle_session(ws) + + errors = [m for m in ws.sent_messages if m["type"] == "error"] + assert any("JSON objects" in e["message"] for e in errors) + + @pytest.mark.asyncio + async def test_unknown_message_type(self): + """Unknown type → error, session continues.""" + ws = MockWebSocket( + [ + _config_msg(), + json.dumps({"type": "teleport.now"}), + _done_msg(), + ] + ) + handler = VideoStreamHandler(chat_handler=MockChatHandler(), idle_timeout=2.0) + await handler.handle_session(ws) + + errors = [m for m in ws.sent_messages if m["type"] == "error"] + assert any("Unknown message type" in e["message"] for e in errors) + assert "session.done" in ws.sent_types() + + @pytest.mark.asyncio + async def test_invalid_base64_frame(self): + """Invalid base64 in video.frame → error, continues.""" + ws = MockWebSocket( + [ + _config_msg(), + json.dumps({"type": "video.frame", "data": "!!!not-base64!!!"}), + _done_msg(), + ] + ) + handler = VideoStreamHandler(chat_handler=MockChatHandler(), idle_timeout=2.0) + await handler.handle_session(ws) + + errors = [m for m in ws.sent_messages if m["type"] == "error"] + assert any("Invalid base64" in e["message"] for e in errors) + assert "session.done" in ws.sent_types() + + +# --------------------------------------------------------------------------- +# Concurrency: frames arrive during query processing +# --------------------------------------------------------------------------- + + +class TestConcurrency: + @pytest.mark.asyncio + async def test_frames_during_query(self): + """Frames sent while a query is processing must be buffered.""" + query_started = asyncio.Event() + query_can_finish = asyncio.Event() + captured_frame_count: list[int] = [] + + class SlowChat: + async def create_chat_completion(self, req, raw_request=None): + query_started.set() + await query_can_finish.wait() + + async def _gen(): + yield _make_sse_line("answer") + yield "data: [DONE]\n\n" + + return _gen() + + class InspectingSlowChat(SlowChat): + """On the second query, record how many frames are in the request.""" + + _call_count = 0 + + async def create_chat_completion(self, req, raw_request=None): + self._call_count += 1 + if self._call_count == 2: + # Count image_url blocks in the request to verify buffering + for msg in req.messages: + if isinstance(msg.get("content", msg), list): + content = msg.get("content", msg) + else: + content = getattr(msg, "content", []) + if isinstance(content, list): + n = sum( + 1 + for p in content + if (isinstance(p, dict) and p.get("type") == "image_url") + or (hasattr(p, "type") and p.type == "image_url") + ) + captured_frame_count.append(n) + return await super().create_chat_completion(req, raw_request) + + # Custom WebSocket that can inject frames after query starts + class TimedWebSocket(MockWebSocket): + def __init__(self): + self._q: asyncio.Queue[str] = asyncio.Queue() + self._sent: list[dict | bytes] = [] + self._accepted = False + + async def accept(self): + self._accepted = True + + async def receive_text(self) -> str: + return await self._q.get() + + def put(self, msg: str): + self._q.put_nowait(msg) + + ws = TimedWebSocket() + chat = InspectingSlowChat() + handler = VideoStreamHandler(chat_handler=chat, idle_timeout=5.0) + + session_task = asyncio.create_task(handler.handle_session(ws)) + + # 1. Config + ws.put(_config_msg()) + await asyncio.sleep(0.01) + + # 2. Initial frame + query + ws.put(_frame_msg(_make_jpeg(10, 10, 10))) + await asyncio.sleep(0.01) + ws.put(_query_msg("first")) + + # 3. Wait for query to start processing + await asyncio.wait_for(query_started.wait(), timeout=2.0) + + # 4. Send more frames WHILE query is running + ws.put(_frame_msg(_make_jpeg(200, 200, 200))) + ws.put(_frame_msg(_make_jpeg(50, 50, 50))) + await asyncio.sleep(0.05) + + # 5. Let query finish + query_can_finish.set() + await asyncio.sleep(0.1) + + # 6. Second query — should see all 3 frames (1 initial + 2 during) + query_started.clear() + query_can_finish.clear() + ws.put(_query_msg("second")) + await asyncio.wait_for(query_started.wait(), timeout=2.0) + query_can_finish.set() + await asyncio.sleep(0.1) + + # 7. Done + ws.put(_done_msg()) + + await asyncio.wait_for(session_task, timeout=5.0) + + assert "session.done" in ws.sent_types() + # Verify: the 2 frames sent during query 1 were buffered and + # included in query 2 (total 3 frames). + assert captured_frame_count == [3] + + +# --------------------------------------------------------------------------- +# SSE parsing (kept from original, these are valid unit tests) +# --------------------------------------------------------------------------- + + +class TestParseSSEDeltas: + def test_single_delta(self): + assert VideoStreamHandler._parse_sse_deltas(_make_sse_line("hello")) == ["hello"] + + def test_multiple_deltas(self): + chunk = _make_sse_line("a") + _make_sse_line("b") + assert VideoStreamHandler._parse_sse_deltas(chunk) == ["a", "b"] + + def test_done_skipped(self): + chunk = _make_sse_line("x") + "data: [DONE]\n\n" + assert VideoStreamHandler._parse_sse_deltas(chunk) == ["x"] + + def test_empty_content_skipped(self): + p = json.dumps({"choices": [{"delta": {"content": ""}}]}) + assert VideoStreamHandler._parse_sse_deltas(f"data: {p}\n\n") == [] + + def test_malformed_json_skipped(self): + chunk = "data: {bad}\n\n" + _make_sse_line("ok") + assert VideoStreamHandler._parse_sse_deltas(chunk) == ["ok"] + + def test_empty_string(self): + assert VideoStreamHandler._parse_sse_deltas("") == [] + + def test_unicode(self): + assert VideoStreamHandler._parse_sse_deltas(_make_sse_line("你好🌍")) == ["你好🌍"] diff --git a/tests/entrypoints/openai_api/test_video_stream_session.py b/tests/entrypoints/openai_api/test_video_stream_session.py new file mode 100644 index 000000000..b0a12b1e4 --- /dev/null +++ b/tests/entrypoints/openai_api/test_video_stream_session.py @@ -0,0 +1,289 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for VideoStreamSession (Phase 2 + 3).""" + +from __future__ import annotations + +import pytest + +from tests.entrypoints.openai_api.conftest_video import ( + make_gradient_jpeg, + make_jpeg, +) +from vllm_omni.entrypoints.openai.video_stream_session import ( + VideoStreamConfig, + VideoStreamSession, +) + +# --------------------------------------------------------------------------- +# VideoStreamConfig +# --------------------------------------------------------------------------- + + +class TestVideoStreamConfig: + def test_defaults(self): + cfg = VideoStreamConfig() + assert cfg.max_frames == 64 + assert cfg.num_sample_frames == 16 + assert cfg.evs_enabled is True + assert cfg.evs_threshold == 0.95 + + def test_from_dict(self): + cfg = VideoStreamConfig.from_dict( + { + "model": "test-model", + "max_frames": 32, + "evs_threshold": 0.90, + "unknown_field": "ignored", + } + ) + assert cfg.model == "test-model" + assert cfg.max_frames == 32 + assert cfg.evs_threshold == 0.90 + + def test_from_dict_empty(self): + cfg = VideoStreamConfig.from_dict({}) + assert cfg.model == "" + assert cfg.max_frames == 64 + + def test_from_dict_invalid_type(self): + with pytest.raises(TypeError, match="max_frames.*expected int.*got str"): + VideoStreamConfig.from_dict({"max_frames": "potato"}) + + def test_from_dict_invalid_bool(self): + with pytest.raises(TypeError, match="evs_enabled.*expected bool"): + VideoStreamConfig.from_dict({"evs_enabled": "yes"}) + + def test_from_dict_invalid_modalities(self): + with pytest.raises(TypeError, match="modalities.*expected list"): + VideoStreamConfig.from_dict({"modalities": 42}) + + def test_from_dict_evs_threshold_int_accepted(self): + """JSON doesn't distinguish int/float — int 1 is a valid threshold.""" + cfg = VideoStreamConfig.from_dict({"evs_threshold": 1}) + assert cfg.evs_threshold == 1 + + +# --------------------------------------------------------------------------- +# Frame buffer & sliding window (uses deque) +# --------------------------------------------------------------------------- + + +class TestFrameBuffer: + def test_add_frame_basic(self): + cfg = VideoStreamConfig(evs_enabled=False, max_frames=10) + session = VideoStreamSession(cfg) + assert session.add_frame(make_jpeg()) is True + assert session.frame_count == 1 + + def test_add_frame_too_large(self): + cfg = VideoStreamConfig(evs_enabled=False, max_frames=10) + session = VideoStreamSession(cfg) + huge = b"\xff" * (10 * 1024 * 1024 + 1) # just over 10 MB + with pytest.raises(ValueError, match="Frame too large"): + session.add_frame(huge) + + def test_sliding_window(self): + cfg = VideoStreamConfig(evs_enabled=False, max_frames=3) + session = VideoStreamSession(cfg) + for i in range(5): + session.add_frame(make_jpeg(r=i * 50)) + assert session.frame_count == 3 + + def test_sliding_window_keeps_newest(self): + cfg = VideoStreamConfig(evs_enabled=False, max_frames=2) + session = VideoStreamSession(cfg) + f1 = make_jpeg(10, 10, 10) + f2 = make_jpeg(20, 20, 20) + f3 = make_jpeg(30, 30, 30) + session.add_frame(f1) + session.add_frame(f2) + session.add_frame(f3) + sampled = session.sample_frames() + assert len(sampled) == 2 + assert sampled[0] == f2 + assert sampled[1] == f3 + + +# --------------------------------------------------------------------------- +# EVS integration +# --------------------------------------------------------------------------- + + +class TestEVSIntegration: + def test_evs_drops_identical_frames(self): + cfg = VideoStreamConfig(evs_enabled=True, evs_threshold=0.90) + session = VideoStreamSession(cfg) + frame = make_jpeg(100, 100, 100) + assert session.add_frame(frame) is True + assert session.add_frame(frame) is False + assert session.frame_count == 1 + + def test_evs_keeps_different_frames(self): + cfg = VideoStreamConfig(evs_enabled=True, evs_threshold=0.95) + session = VideoStreamSession(cfg) + for i in range(5): + assert session.add_frame(make_gradient_jpeg(seed=i)) is True + assert session.frame_count == 5 + + def test_evs_disabled(self): + cfg = VideoStreamConfig(evs_enabled=False) + session = VideoStreamSession(cfg) + frame = make_jpeg() + assert session.add_frame(frame) is True + assert session.add_frame(frame) is True + assert session.frame_count == 2 + + def test_evs_stats(self): + cfg = VideoStreamConfig(evs_enabled=True, evs_threshold=0.90) + session = VideoStreamSession(cfg) + frame = make_jpeg() + session.add_frame(frame) + session.add_frame(frame) + stats = session.evs_stats + assert stats is not None + assert stats["retained_count"] == 1 + assert stats["dropped_count"] == 1 + + def test_evs_stats_none_when_disabled(self): + cfg = VideoStreamConfig(evs_enabled=False) + session = VideoStreamSession(cfg) + assert session.evs_stats is None + + +# --------------------------------------------------------------------------- +# Uniform sampling +# --------------------------------------------------------------------------- + + +class TestSampling: + def test_sample_exact(self): + cfg = VideoStreamConfig(evs_enabled=False, num_sample_frames=3, max_frames=10) + session = VideoStreamSession(cfg) + frames = [make_gradient_jpeg(i) for i in range(3)] + for f in frames: + session.add_frame(f) + sampled = session.sample_frames() + assert len(sampled) == 3 + assert sampled == frames + + def test_sample_fewer_than_requested(self): + cfg = VideoStreamConfig(evs_enabled=False, num_sample_frames=10, max_frames=64) + session = VideoStreamSession(cfg) + for i in range(3): + session.add_frame(make_gradient_jpeg(i)) + assert len(session.sample_frames()) == 3 + + def test_sample_uniform(self): + cfg = VideoStreamConfig(evs_enabled=False, num_sample_frames=4, max_frames=64) + session = VideoStreamSession(cfg) + for i in range(10): + session.add_frame(make_gradient_jpeg(i)) + sampled = session.sample_frames() + assert len(sampled) == 4 + expected = [session._frames[i] for i in [0, 3, 6, 9]] + assert sampled == expected + + def test_sample_empty(self): + cfg = VideoStreamConfig(evs_enabled=False) + session = VideoStreamSession(cfg) + assert session.sample_frames() == [] + + def test_sample_single_frame_from_multi_frame_buffer(self): + cfg = VideoStreamConfig(evs_enabled=False, num_sample_frames=1, max_frames=64) + session = VideoStreamSession(cfg) + first = make_gradient_jpeg(0) + second = make_gradient_jpeg(1) + session.add_frame(first) + session.add_frame(second) + + assert session.sample_frames() == [second] + + +# --------------------------------------------------------------------------- +# Audio buffer (Phase 3) +# --------------------------------------------------------------------------- + + +class TestAudioBuffer: + def test_add_audio_chunk(self): + session = VideoStreamSession(VideoStreamConfig()) + assert session.has_audio is False + session.add_audio_chunk(b"\x00" * 100) + assert session.has_audio is True + + def test_clear_audio(self): + session = VideoStreamSession(VideoStreamConfig()) + session.add_audio_chunk(b"\x00" * 100) + session.clear_audio() + assert session.has_audio is False + + +# --------------------------------------------------------------------------- +# build_chat_request +# --------------------------------------------------------------------------- + + +class TestBuildChatRequest: + def test_video_only_request(self): + cfg = VideoStreamConfig(model="test-model", evs_enabled=False, num_sample_frames=4) + session = VideoStreamSession(cfg) + for i in range(4): + session.add_frame(make_gradient_jpeg(i)) + + request = session.build_chat_request("Describe this scene.") + assert request.model == "test-model" + assert request.stream is True + + content = request.messages[0]["content"] + assert len(content) == 5 # 4 image_url + 1 text + image_parts = [p for p in content if p["type"] == "image_url"] + text_parts = [p for p in content if p["type"] == "text"] + assert len(image_parts) == 4 + assert len(text_parts) == 1 + assert text_parts[0]["text"] == "Describe this scene." + + mm_kw = getattr(request, "mm_processor_kwargs", None) + assert mm_kw is None or not mm_kw.get("use_audio_in_video", False) + + def test_video_plus_audio_request(self): + cfg = VideoStreamConfig(model="test-model", evs_enabled=False, num_sample_frames=2) + session = VideoStreamSession(cfg) + session.add_frame(make_gradient_jpeg(0)) + session.add_frame(make_gradient_jpeg(1)) + session.add_audio_chunk(b"\x00" * 3200) + + request = session.build_chat_request("What is being said?") + + content = request.messages[0]["content"] + assert len(content) == 4 # 2 image_url + 1 audio_url + 1 text + audio_parts = [p for p in content if p["type"] == "audio_url"] + assert len(audio_parts) == 1 + # RFC 3551: audio/L16 for linear 16-bit PCM + assert audio_parts[0]["audio_url"]["url"].startswith("data:audio/L16;rate=16000;base64,") + + mm_kw = getattr(request, "mm_processor_kwargs", None) or {} + assert mm_kw.get("use_audio_in_video") is True + + def test_image_url_is_valid_base64(self): + cfg = VideoStreamConfig(evs_enabled=False, num_sample_frames=1) + session = VideoStreamSession(cfg) + session.add_frame(make_jpeg(200, 100, 50)) + request = session.build_chat_request("test") + content = request.messages[0]["content"] + img_url = content[0]["image_url"]["url"] + assert img_url.startswith("data:image/jpeg;base64,") + import base64 + + b64_data = img_url.split(",", 1)[1] + decoded = base64.b64decode(b64_data) + assert len(decoded) > 0 + + def test_clear_audio_after_query(self): + session = VideoStreamSession(VideoStreamConfig(evs_enabled=False)) + session.add_frame(make_jpeg()) + session.add_audio_chunk(b"\x00" * 100) + session.build_chat_request("test") + session.clear_audio() + assert session.has_audio is False + assert session.frame_count == 1 diff --git a/tests/entrypoints/test_omni_entrypoints.py b/tests/entrypoints/test_omni_entrypoints.py index 289cf2673..a96b4dd1d 100644 --- a/tests/entrypoints/test_omni_entrypoints.py +++ b/tests/entrypoints/test_omni_entrypoints.py @@ -1,5 +1,6 @@ from __future__ import annotations +import argparse import queue from collections.abc import Callable from types import SimpleNamespace @@ -14,6 +15,7 @@ from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError from vllm_omni.entrypoints.async_omni import AsyncOmni from vllm_omni.entrypoints.omni import Omni +from vllm_omni.entrypoints.omni_base import OmniEngineDeadError pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -165,6 +167,31 @@ def _patch_engine(monkeypatch: pytest.MonkeyPatch, engine: FakeAsyncOmniEngine) monkeypatch.setattr("vllm_omni.entrypoints.omni_base.omni_snapshot_download", lambda model: model) +def test_from_cli_args_only_nulls_untyped_override_fields(monkeypatch: pytest.MonkeyPatch): + from vllm_omni.entrypoints.omni import Omni + + captured: dict[str, Any] = {} + + def fake_engine(*args: Any, **kwargs: Any) -> FakeAsyncOmniEngine: + captured.update(kwargs) + return FakeAsyncOmniEngine() + + monkeypatch.setattr("vllm_omni.entrypoints.omni_base.AsyncOmniEngine", fake_engine) + monkeypatch.setattr("vllm_omni.entrypoints.omni_base.omni_snapshot_download", lambda model: model) + monkeypatch.setattr("sys.argv", ["prog"]) + + parser = argparse.ArgumentParser() + parser.add_argument("--gpu-memory-utilization", type=float, default=0.9) + parser.add_argument("--hsdp-shard-size", type=int, default=-1) + args = parser.parse_args([]) + args.model = "fake-model" + + Omni.from_cli_args(args, parser=parser) + + assert captured["gpu_memory_utilization"] is None + assert captured["hsdp_shard_size"] == -1 + + def _make_base(): from vllm_omni.entrypoints.omni_base import OmniBase @@ -335,6 +362,18 @@ def _enqueue_error_message(engine: FakeAsyncOmniEngine, msg: dict[str, Any]) -> ) +def _enqueue_fatal_error_message(engine: FakeAsyncOmniEngine, msg: dict[str, Any]) -> None: + engine.output_q.put_nowait( + { + "type": "error", + "fatal": True, + "request_id": msg["request_id"], + "stage_id": 2, + "error": "engine dead", + } + ) + + @pytest.mark.asyncio async def test_get_supported_tasks_returns_engine_supported_tasks(): omni = object.__new__(AsyncOmni) @@ -546,18 +585,22 @@ async def test_async_omni_abort_forwards_to_engine(monkeypatch: pytest.MonkeyPat @pytest.mark.asyncio -async def test_async_omni_propagates_engine_error(monkeypatch: pytest.MonkeyPatch): - engine = FakeAsyncOmniEngine(stage_metadata=THREE_STAGE_META, on_add_request=_enqueue_error_message) +async def test_async_omni_propagates_fatal_error_context(monkeypatch: pytest.MonkeyPatch): + engine = FakeAsyncOmniEngine(stage_metadata=THREE_STAGE_META, on_add_request=_enqueue_fatal_error_message) _patch_engine(monkeypatch, engine) app = AsyncOmni("dummy-model") try: - with pytest.raises(RuntimeError, match="engine boom"): + with pytest.raises(EngineDeadError, match="engine dead") as exc_info: async for _ in app.generate(prompt="hello", request_id="req-1"): pass finally: app.shutdown() + assert isinstance(exc_info.value, OmniEngineDeadError) + assert str(exc_info.value) == "engine dead" + assert getattr(exc_info.value, "error_stage_id") == 2 + def test_omni_generate_py_generator_yields_final_outputs_for_each_request(monkeypatch: pytest.MonkeyPatch): sampling_params = [SamplingParams(max_tokens=8) for _ in range(3)] diff --git a/tests/examples/offline_inference/test_text_to_image.py b/tests/examples/offline_inference/test_text_to_image.py index f24506587..041c32dc4 100644 --- a/tests/examples/offline_inference/test_text_to_image.py +++ b/tests/examples/offline_inference/test_text_to_image.py @@ -11,7 +11,8 @@ from tests.examples.helpers import EXAMPLES, ExampleRunner, ReadmeSnippet from tests.helpers.assertions import assert_image_valid from tests.helpers.mark import hardware_marks -pytestmark = [pytest.mark.advanced_model, pytest.mark.example, *hardware_marks(res={"cuda": "H100"})] +pytestmark = [pytest.mark.full_model, pytest.mark.example, *hardware_marks(res={"cuda": "H100"})] + T2I_SCRIPT = EXAMPLES / "offline_inference" / "text_to_image" / "text_to_image.py" README_PATH = T2I_SCRIPT.with_name("README.md") diff --git a/tests/examples/online_serving/test_qwen2_5_omni.py b/tests/examples/online_serving/test_qwen2_5_omni.py index b8acfef84..b3e49b8d9 100644 --- a/tests/examples/online_serving/test_qwen2_5_omni.py +++ b/tests/examples/online_serving/test_qwen2_5_omni.py @@ -4,9 +4,6 @@ See examples/online_serving/qwen2_5_omni/README.md """ import os - -os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" - from pathlib import Path import pytest @@ -21,7 +18,9 @@ from tests.helpers.media import convert_audio_file_to_text, cosine_similarity_te from tests.helpers.runtime import OmniServerParams from tests.helpers.stage_config import get_deploy_config_path -pytestmark = [pytest.mark.advanced_model, pytest.mark.example] +pytestmark = [pytest.mark.full_model, pytest.mark.example, pytest.mark.omni] + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" models = ["Qwen/Qwen2.5-Omni-7B"] @@ -40,8 +39,6 @@ test_params = [ common_args = ["python", os.path.join(example_dir, "openai_chat_completion_client_for_multimodal_generation.py")] -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_send_multimodal_request_001(omni_server) -> None: @@ -77,8 +74,6 @@ def test_send_multimodal_request_001(omni_server) -> None: # TODO: Verify the E2E latency after confirmation baseline. -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_send_multimodal_request_002(omni_server) -> None: @@ -114,8 +109,6 @@ def test_send_multimodal_request_002(omni_server) -> None: # TODO: Verify the E2E latency after confirmation baseline. -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_send_multimodal_request_003(omni_server) -> None: @@ -141,8 +134,6 @@ def test_send_multimodal_request_003(omni_server) -> None: # TODO: Verify the E2E latency after confirmation baseline. -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_modality_control_001(omni_server) -> None: @@ -171,8 +162,6 @@ def test_modality_control_001(omni_server) -> None: # TODO: Verify the E2E latency after confirmation baseline. -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_modality_control_002(omni_server) -> None: @@ -200,8 +189,6 @@ def test_modality_control_002(omni_server) -> None: # TODO: Verify the E2E latency after confirmation baseline. -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_modality_control_003(omni_server) -> None: @@ -238,8 +225,6 @@ def test_modality_control_003(omni_server) -> None: # TODO: Verify the E2E latency after confirmation baseline. -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards={"cuda": 4, "rocm": 2}) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_stream_001(omni_server) -> None: diff --git a/tests/examples/online_serving/test_qwen3_omni.py b/tests/examples/online_serving/test_qwen3_omni.py index e2c3b3384..e52a2bf5a 100644 --- a/tests/examples/online_serving/test_qwen3_omni.py +++ b/tests/examples/online_serving/test_qwen3_omni.py @@ -4,9 +4,6 @@ See examples/online_serving/qwen3_omni/README.md """ import os - -os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" - from pathlib import Path import pytest @@ -21,7 +18,9 @@ from tests.helpers.media import convert_audio_file_to_text, cosine_similarity_te from tests.helpers.runtime import OmniServerParams from tests.helpers.stage_config import get_deploy_config_path -pytestmark = [pytest.mark.advanced_model, pytest.mark.example] +pytestmark = [pytest.mark.full_model, pytest.mark.example, pytest.mark.omni] + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"] @@ -39,8 +38,6 @@ test_params = [ common_args = ["python", os.path.join(example_dir, "openai_chat_completion_client_for_multimodal_generation.py")] -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_send_multimodal_request_001(omni_server) -> None: @@ -69,8 +66,6 @@ def test_send_multimodal_request_001(omni_server) -> None: # TODO: Verify the E2E latency after confirmation baseline. -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_send_multimodal_request_002(omni_server) -> None: @@ -102,8 +97,6 @@ def test_send_multimodal_request_002(omni_server) -> None: # TODO: Verify the E2E latency after confirmation baseline. -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_send_multimodal_request_003(omni_server) -> None: @@ -119,8 +112,6 @@ def test_send_multimodal_request_003(omni_server) -> None: # TODO: Verify the E2E latency after confirmation baseline. -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_modality_control_001(omni_server) -> None: @@ -143,8 +134,6 @@ def test_modality_control_001(omni_server) -> None: # TODO: Verify the E2E latency after confirmation baseline. -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_modality_control_002(omni_server) -> None: @@ -167,8 +156,6 @@ def test_modality_control_002(omni_server) -> None: # TODO: Verify the E2E latency after confirmation baseline. -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_modality_control_003(omni_server) -> None: @@ -199,8 +186,6 @@ def test_modality_control_003(omni_server) -> None: # TODO: Verify the E2E latency after confirmation baseline. -@pytest.mark.advanced_model -@pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_stream_001(omni_server) -> None: diff --git a/tests/examples/online_serving/test_text_to_image.py b/tests/examples/online_serving/test_text_to_image.py index 6f89cb549..ee0a1fedb 100644 --- a/tests/examples/online_serving/test_text_to_image.py +++ b/tests/examples/online_serving/test_text_to_image.py @@ -18,7 +18,7 @@ from tests.helpers.assertions import assert_image_valid from tests.helpers.mark import hardware_marks from tests.helpers.runtime import OmniServer, OmniServerParams -pytestmark = [pytest.mark.advanced_model, pytest.mark.example, *hardware_marks(res={"cuda": "H100"})] +pytestmark = [pytest.mark.full_model, pytest.mark.example, *hardware_marks(res={"cuda": "H100"})] T2I_ONLINE_CLIENT = EXAMPLES / "online_serving" / "text_to_image" / "openai_chat_client.py" EXAMPLE_OUTPUT_SUBFOLDER = "example_online_t2i" diff --git a/tests/helpers/assertions.py b/tests/helpers/assertions.py index 533806477..604b76b62 100644 --- a/tests/helpers/assertions.py +++ b/tests/helpers/assertions.py @@ -11,7 +11,9 @@ import numpy as np import soundfile as sf from PIL import Image -from tests.helpers.media import cosine_similarity_text +from tests.helpers.media import ( + cosine_similarity_text, +) _GENDER_PIPELINE = None _GENDER_PIPELINE_LOCK = threading.Lock() @@ -56,7 +58,7 @@ def assert_image_diffusion_response( f"Expected {num_outputs_per_prompt} images, got {len(response.images)}" ) - if run_level == "advanced_model": + if run_level in {"advanced_model", "full_model"}: width = extra_body.get("width") height = extra_body.get("height") @@ -407,7 +409,8 @@ def assert_omni_response(response: Any, request_config: dict[str, Any], run_leve modalities = request_config.get("modalities", ["text", "audio"]) - if run_level == "advanced_model": + if run_level in {"advanced_model", "full_model"}: + # Verify output success if "audio" in modalities: assert response.audio_content is not None, "No audio output is generated" print(f"audio content is: {response.audio_content}") @@ -417,12 +420,11 @@ def assert_omni_response(response: Any, request_config: dict[str, Any], run_leve response.audio_bytes, speaker, ) - if "text" in modalities: assert response.text_content is not None, "No text output is generated" print(f"text content is: {response.text_content}") - # Verify image description + # Verify keywords in output word_types = ["text", "image", "audio", "video"] keywords_dict = request_config.get("key_words", {}) for word_type in word_types: @@ -441,11 +443,25 @@ def assert_omni_response(response: Any, request_config: dict[str, Any], run_leve ) # Verify similarity (Whisper transcript vs streamed/detokenized text) - if "text" in modalities and "audio" in modalities: - assert response.similarity is not None and response.similarity > 0.9, ( - "The audio content is not same as the text" - ) - print(f"similarity is: {response.similarity}") + if "audio" in modalities: + audio_ref_text = request_config.get("audio_ref_text") + if "text" in modalities: + transcript = (response.audio_content or "").strip() + text_output = (response.text_content or "").strip() + similarity = cosine_similarity_text( + transcript.lower(), + text_output.lower(), + ) + assert similarity > 0.9, "The audio content is not same as the text" + print(f"similarity is: {similarity}") + if audio_ref_text: + audio_similarity = cosine_similarity_text( + response.audio_content.lower(), + str(audio_ref_text).lower(), + ) + assert audio_similarity > 0.9, ( + f"The audio content does not match reference text: similarity={audio_similarity:.3f}" + ) def assert_audio_speech_response(response: Any, request_config: dict[str, Any], run_level: str) -> None: @@ -464,7 +480,7 @@ def assert_audio_speech_response(response: Any, request_config: dict[str, Any], elif req_fmt == "wav" and response.audio_format: assert req_fmt in response.audio_format - if run_level == "advanced_model" and req_fmt != "pcm": + if run_level in {"advanced_model", "full_model"} and req_fmt != "pcm": expected_text = request_config.get("input") if expected_text: transcript = (response.audio_content or "").strip() diff --git a/tests/helpers/fixtures/run_args.py b/tests/helpers/fixtures/run_args.py index b18a64b98..975584d20 100644 --- a/tests/helpers/fixtures/run_args.py +++ b/tests/helpers/fixtures/run_args.py @@ -6,8 +6,8 @@ def pytest_addoption(parser): "--run-level", action="store", default="core_model", - choices=["core_model", "advanced_model"], - help="Test level to run: L2, L3", + choices=["core_model", "advanced_model", "full_model"], + help="Test level to run: L2, L3, L4", ) diff --git a/tests/helpers/fixtures/runtime.py b/tests/helpers/fixtures/runtime.py index 77a468f2f..4cae13cd6 100644 --- a/tests/helpers/fixtures/runtime.py +++ b/tests/helpers/fixtures/runtime.py @@ -36,7 +36,7 @@ def omni_server(request: pytest.FixtureRequest, run_level: str, model_prefix: st model = model_prefix + params.model port = params.port stage_config_path = params.stage_config_path - if run_level == "advanced_model" and stage_config_path is not None: + if run_level in {"advanced_model", "full_model"} and stage_config_path is not None: with open(stage_config_path, encoding="utf-8") as f: cfg = yaml.safe_load(f) or {} # Strip ``load_format: dummy`` (CI overlay default) so advanced_model diff --git a/tests/helpers/media.py b/tests/helpers/media.py index 4463acbbb..c0fb97171 100644 --- a/tests/helpers/media.py +++ b/tests/helpers/media.py @@ -3,6 +3,7 @@ import base64 import concurrent.futures import gc +import hashlib import io import logging import math @@ -66,6 +67,7 @@ def generate_synthetic_audio( num_channels: int, sample_rate: int = 48000, *, + phrase_text: str = "test", force_regenerate: bool = False, cache_dir: Path | str | None = None, ) -> dict[str, Any]: @@ -74,12 +76,16 @@ def generate_synthetic_audio( Caches the WAV under ``cache_dir`` when given, else under the default temp subdirectory. Reuses the file when the same - ``duration`` / ``num_channels`` / ``sample_rate`` are requested unless - ``force_regenerate`` is true. + ``duration`` / ``num_channels`` / ``sample_rate`` / ``phrase_text`` are + requested unless ``force_regenerate`` is true. + + The cache filename includes a SHA-256 digest of ``phrase_text`` so different + phrases never share a WAV cache entry. """ root = _resolve_synthetic_media_cache_dir(cache_dir) root.mkdir(parents=True, exist_ok=True) - cache_path = root / f"synth_audio_d{duration}_ch{num_channels}_sr{sample_rate}.wav" + phrase_key = hashlib.sha256(phrase_text.encode("utf-8")).hexdigest() + cache_path = root / f"synth_audio_d{duration}_ch{num_channels}_sr{sample_rate}_pt{phrase_key}.wav" if not force_regenerate and cache_path.is_file(): data, _sr = sf.read(str(cache_path), dtype="float32", always_2d=True) @@ -204,7 +210,6 @@ def generate_synthetic_audio( enhanced = enhanced / peak * 0.95 return enhanced.astype(np.float32) - phrase_text = "test" num_samples = int(sample_rate * max(1, duration)) audio_data = np.zeros((num_samples, num_channels), dtype=np.float32) @@ -442,12 +447,16 @@ def generate_synthetic_image( *, force_regenerate: bool = False, cache_dir: Path | str | None = None, + seed: int | None = None, ) -> dict[str, Any]: """ Random colored squares on white background. Caches JPEG by ``width`` / ``height`` unless ``force_regenerate`` is true. Cache root: ``cache_dir`` if given, else the default temp subdirectory. """ + if seed is not None: + random.seed(seed) + root = _resolve_synthetic_media_cache_dir(cache_dir) root.mkdir(parents=True, exist_ok=True) cache_path = root / f"synth_image_w{width}_h{height}.jpg" diff --git a/tests/helpers/runtime.py b/tests/helpers/runtime.py index 9520aff0c..0cf0f9e48 100644 --- a/tests/helpers/runtime.py +++ b/tests/helpers/runtime.py @@ -35,7 +35,6 @@ from tests.helpers.env import run_forced_gpu_cleanup_round from tests.helpers.media import ( _merge_base64_audio_to_segment, convert_audio_bytes_to_text, - cosine_similarity_text, decode_b64_image, ) from vllm_omni.config.stage_config import resolve_deploy_yaml @@ -500,7 +499,6 @@ class OmniResponse: audio_content: str | None = None audio_format: str | None = None audio_bytes: bytes | None = None - similarity: float | None = None e2e_latency: float | None = None success: bool = False error_message: str | None = None @@ -542,19 +540,15 @@ class OpenAIClientHandler: text_content += content result.e2e_latency = time.perf_counter() - start_time audio_content = None - similarity = None if audio_data: merged_seg = _merge_base64_audio_to_segment(audio_data) wav_buf = BytesIO() merged_seg.export(wav_buf, format="wav") result.audio_bytes = wav_buf.getvalue() audio_content = convert_audio_bytes_to_text(result.audio_bytes) - if audio_content and text_content: - similarity = cosine_similarity_text(audio_content.lower(), text_content.lower()) result.text_content = text_content result.audio_data = audio_data result.audio_content = audio_content - result.similarity = similarity result.success = True except Exception as e: result.error_message = f"Stream processing error: {str(e)}" @@ -578,15 +572,11 @@ class OpenAIClientHandler: result.cached_tokens = details.cached_tokens result.e2e_latency = time.perf_counter() - start_time audio_content = None - similarity = None if audio_data: result.audio_bytes = base64.b64decode(audio_data) audio_content = convert_audio_bytes_to_text(result.audio_bytes) - if audio_content and text_content: - similarity = cosine_similarity_text(audio_content.lower(), text_content.lower()) result.text_content = text_content result.audio_content = audio_content - result.similarity = similarity result.success = True except Exception as e: result.error_message = f"Non-stream processing error: {str(e)}" @@ -630,6 +620,9 @@ class OpenAIClientHandler: mm = dict(extra_body.get("mm_processor_kwargs") or {}) mm["use_audio_in_video"] = True extra_body["mm_processor_kwargs"] = mm + if "sampling_params_list" in request_config: + extra_body["sampling_params_list"] = request_config["sampling_params_list"] + create_kwargs: dict[str, Any] = { "model": request_config.get("model"), "messages": request_config.get("messages"), diff --git a/tests/helpers/stage_config.py b/tests/helpers/stage_config.py index 29a80372e..2bb017b81 100644 --- a/tests/helpers/stage_config.py +++ b/tests/helpers/stage_config.py @@ -411,6 +411,79 @@ _CI_OVERLAYS: dict[str, dict[str, Any]] = { }, }, }, + "bagel": { + "base_config": "bagel.yaml", + "stages": [ + { + "stage_id": 0, + "max_num_seqs": 3, + "gpu_memory_utilization": 0.45, + "load_format": "dummy", + }, + { + "stage_id": 1, + "max_num_seqs": 1, + "load_format": "dummy", + }, + ], + }, + "bagel_think": { + "base_config": "bagel_think.yaml", + "stages": [ + { + "stage_id": 0, + "max_num_seqs": 3, + "gpu_memory_utilization": 0.45, + "load_format": "dummy", + }, + { + "stage_id": 1, + "max_num_seqs": 1, + "load_format": "dummy", + }, + ], + }, + "bagel_single_stage": { + "base_config": "bagel_single_stage.yaml", + "stages": [ + { + "stage_id": 0, + "max_num_seqs": 1, + "load_format": "dummy", + }, + ], + }, + "bagel_mooncake": { + "base_config": "bagel.yaml", + "stages": [ + { + "stage_id": 0, + "max_num_seqs": 1, + "gpu_memory_utilization": 0.45, + "load_format": "dummy", + "output_connectors": {"to_stage_1": "mooncake_connector"}, + }, + { + "stage_id": 1, + "max_num_seqs": 1, + "load_format": "dummy", + "input_connectors": {"from_stage_0": "mooncake_connector"}, + }, + ], + "connectors": { + "mooncake_connector": { + "name": "MooncakeConnector", + "extra": { + "host": "${MOONCAKE_HOST}", + "metadata_server": "http://${MOONCAKE_HOST}:${MOONCAKE_HTTP_PORT}/metadata", + "master": "${MOONCAKE_HOST}:${MOONCAKE_RPC_PORT}", + "segment": 64000000, + "localbuf": 64000000, + "proto": "tcp", + }, + }, + }, + }, # Single-stage thinker-only topology for the abort test. "qwen2_5_omni_thinker_only": { "async_chunk": False, diff --git a/tests/model_executor/models/ming_flash_omni/test_talker_cfm.py b/tests/model_executor/models/ming_flash_omni/test_talker_cfm.py new file mode 100644 index 000000000..419ce00da --- /dev/null +++ b/tests/model_executor/models/ming_flash_omni/test_talker_cfm.py @@ -0,0 +1,146 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from vllm_omni.model_executor.models.ming_flash_omni.talker_module import ( + CFM, + Aggregator, + CFMGraphExecutor, + CFMGraphExecutorPool, + DiT, +) + +torch = pytest.importorskip("torch") +pytest.importorskip("x_transformers") + +pytestmark = [ + pytest.mark.core_model, + pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA for graph capture"), +] + +_LATENT_DIM = 8 +_PATCH_SIZE = 4 +_HIS_PATCH_SIZE = 8 +_LLM_HIDDEN = 16 +_DIT_HIDDEN = 32 +_AGG_HIDDEN = 32 +_NUM_HEADS = 4 +_DEPTH = 2 +_STEPS = 5 +_DTYPE = torch.float32 + + +def _warmup_pipeline(cfm: CFM, aggregator: Aggregator, stop_head: torch.nn.Linear, device: torch.device) -> None: + llm_cond = torch.randn(1, 1, _LLM_HIDDEN, device=device, dtype=_DTYPE) + lat_cond = torch.randn(1, _HIS_PATCH_SIZE, _LATENT_DIM, device=device, dtype=_DTYPE) + y0 = torch.randn(1, _PATCH_SIZE, _LATENT_DIM, device=device, dtype=_DTYPE) + t = torch.linspace(0.0, 1.0, _STEPS + 1, device=device, dtype=_DTYPE) + sde_args = torch.tensor([2.0, 0.25, 0.0], device=device, dtype=_DTYPE) + sde_rnd = torch.randn(_STEPS, 1, _PATCH_SIZE, _LATENT_DIM, device=device, dtype=_DTYPE) + + with torch.no_grad(): + gen_lat = cfm.sample(llm_cond, lat_cond, y0, t, sde_args, sde_rnd) + aggregator(gen_lat) + stop_head(llm_cond[:, -1, :]).softmax(dim=-1) + torch.cuda.synchronize(device) + + +def _build_pipeline(): + device = torch.device("cuda") + dit = ( + DiT( + in_channels=_LATENT_DIM, + hidden_size=_DIT_HIDDEN, + depth=_DEPTH, + num_heads=_NUM_HEADS, + mlp_ratio=2.0, + llm_cond_dim=_LLM_HIDDEN, + ) + .to(device=device, dtype=_DTYPE) + .eval() + ) + cfm = CFM(dit, steps=_STEPS, sway_sampling_coef=-1.0).to(device=device, dtype=_DTYPE).eval() + aggregator = ( + Aggregator( + in_channels=_LATENT_DIM, + hidden_size=_AGG_HIDDEN, + depth=_DEPTH, + num_heads=_NUM_HEADS, + mlp_ratio=2.0, + llm_input_dim=_LLM_HIDDEN, + ) + .to(device=device, dtype=_DTYPE) + .eval() + ) + stop_head = torch.nn.Linear(_LLM_HIDDEN, 2).to(device=device, dtype=_DTYPE).eval() + + config = SimpleNamespace(steps=_STEPS, patch_size=_PATCH_SIZE) + _warmup_pipeline(cfm, aggregator, stop_head, device) + return config, cfm, aggregator, stop_head, device + + +class TestCFMGraphExecutor: + """Capture once, replay twice: outputs must stay consistently-shaped.""" + + def test_execute_shapes_and_replay(self) -> None: + config, cfm, aggregator, stop_head, device = _build_pipeline() + executor = CFMGraphExecutor(config, cfm, aggregator, stop_head) + + bsz = 1 + input_tensor = torch.randn(bsz, 1, _LLM_HIDDEN, device=device, dtype=_DTYPE) + his_lat = torch.randn(bsz, _HIS_PATCH_SIZE, _LATENT_DIM, device=device, dtype=_DTYPE) + + gen_lat, inputs_embeds, stop_out = executor.execute(input_tensor, his_lat) + torch.cuda.synchronize() + + assert gen_lat.shape == (bsz, _PATCH_SIZE, _LATENT_DIM) + assert inputs_embeds.shape == (bsz, 1, _LLM_HIDDEN) + assert stop_out.shape == (bsz, 2) + assert torch.isfinite(gen_lat).all() + assert torch.isfinite(inputs_embeds).all() + # stop_head output is softmax-normalized across the last dim. + assert torch.allclose(stop_out.sum(dim=-1), torch.ones(bsz, device=device, dtype=_DTYPE), atol=1e-4) + + # Replay the captured graph with fresh inputs — shapes must match. + new_input = torch.randn_like(input_tensor) + new_his = torch.randn_like(his_lat) + gen_lat2, inputs_embeds2, stop_out2 = executor.execute(new_input, new_his) + torch.cuda.synchronize() + assert gen_lat2.shape == gen_lat.shape + assert inputs_embeds2.shape == inputs_embeds.shape + assert stop_out2.shape == stop_out.shape + assert executor.initialized is True + + def test_execute_is_noninplace_on_inputs(self) -> None: + config, cfm, aggregator, stop_head, device = _build_pipeline() + executor = CFMGraphExecutor(config, cfm, aggregator, stop_head) + + input_tensor = torch.randn(1, 1, _LLM_HIDDEN, device=device, dtype=_DTYPE) + his_lat = torch.randn(1, _HIS_PATCH_SIZE, _LATENT_DIM, device=device, dtype=_DTYPE) + snapshot_input = input_tensor.clone() + snapshot_his = his_lat.clone() + + executor.execute(input_tensor, his_lat) + torch.cuda.synchronize() + assert torch.equal(input_tensor, snapshot_input) + assert torch.equal(his_lat, snapshot_his) + + +class TestCFMGraphExecutorPool: + def test_pool_acquires_and_releases(self) -> None: + config, cfm, aggregator, stop_head, device = _build_pipeline() + pool = CFMGraphExecutorPool(config, cfm, aggregator, stop_head, pool_size=2) + + input_tensor = torch.randn(1, 1, _LLM_HIDDEN, device=device, dtype=_DTYPE) + his_lat = torch.randn(1, _HIS_PATCH_SIZE, _LATENT_DIM, device=device, dtype=_DTYPE) + + gen_lat, inputs_embeds, stop_out = pool.execute(input_tensor, his_lat) + torch.cuda.synchronize() + assert gen_lat.shape == (1, _PATCH_SIZE, _LATENT_DIM) + assert inputs_embeds.shape == (1, 1, _LLM_HIDDEN) + assert stop_out.shape == (1, 2) + assert pool.pool.qsize() == 2 diff --git a/tests/model_executor/models/ming_flash_omni/test_talker_modules.py b/tests/model_executor/models/ming_flash_omni/test_talker_modules.py new file mode 100644 index 000000000..4cbbc887a --- /dev/null +++ b/tests/model_executor/models/ming_flash_omni/test_talker_modules.py @@ -0,0 +1,162 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from __future__ import annotations + +import pytest + +from vllm_omni.model_executor.models.ming_flash_omni.talker_module import CFM, Aggregator, DiT + +torch = pytest.importorskip("torch") +pytest.importorskip("x_transformers") + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +_LATENT_DIM = 8 +_PATCH_SIZE = 4 +_HIS_PATCH_SIZE = 8 +_LLM_HIDDEN = 16 +_DIT_HIDDEN = 32 +_AGG_HIDDEN = 32 +_NUM_HEADS = 4 +_DEPTH = 2 +_STEPS = 5 + + +def _make_dit() -> DiT: + return DiT( + in_channels=_LATENT_DIM, + hidden_size=_DIT_HIDDEN, + depth=_DEPTH, + num_heads=_NUM_HEADS, + mlp_ratio=2.0, + llm_cond_dim=_LLM_HIDDEN, + ) + + +def _make_aggregator() -> Aggregator: + return Aggregator( + in_channels=_LATENT_DIM, + hidden_size=_AGG_HIDDEN, + depth=_DEPTH, + num_heads=_NUM_HEADS, + mlp_ratio=2.0, + llm_input_dim=_LLM_HIDDEN, + ) + + +class TestDiTDummyForward: + """DiT with dummy weights runs forward + CFG-doubled forward.""" + + def test_forward_shape(self) -> None: + dit = _make_dit().eval() + bsz = 2 + x = torch.randn(bsz, _PATCH_SIZE, _LATENT_DIM) + t = torch.zeros(bsz) + c = torch.randn(bsz, 1, _LLM_HIDDEN) + latent_history = torch.randn(bsz, _HIS_PATCH_SIZE, _LATENT_DIM) + + with torch.no_grad(): + out = dit(x, t, c, latent_history) + + # Output preserves the concatenated (history + time/cond prefix + x) + # token axis: history + 1 (time+cond) + patch. + assert out.shape == (bsz, _HIS_PATCH_SIZE + 1 + _PATCH_SIZE, _LATENT_DIM) + + def test_forward_with_cfg_trims_to_patch(self) -> None: + dit = _make_dit().eval() + bsz = 1 + x = torch.randn(bsz, _PATCH_SIZE, _LATENT_DIM) + t = torch.zeros(()) + c = torch.randn(bsz, 1, _LLM_HIDDEN) + latent_history = torch.randn(bsz, _HIS_PATCH_SIZE, _LATENT_DIM) + + with torch.no_grad(): + out = dit.forward_with_cfg(x, t, c, latent_history) + + # CFG doubles the batch and trims the output to the patch window. + assert out.shape == (2 * bsz, _PATCH_SIZE, _LATENT_DIM) + + +class TestAggregatorDummyForward: + """Aggregator with dummy weights maps latent patch -> LLM hidden.""" + + def test_forward_shape(self) -> None: + agg = _make_aggregator().eval() + bsz = 3 + gen_lat = torch.randn(bsz, _PATCH_SIZE, _LATENT_DIM) + + with torch.no_grad(): + out = agg(gen_lat) + + assert out.shape == (bsz, 1, _LLM_HIDDEN) + + def test_forward_is_finite(self) -> None: + agg = _make_aggregator().eval() + gen_lat = torch.randn(1, _PATCH_SIZE, _LATENT_DIM) + with torch.no_grad(): + out = agg(gen_lat) + assert torch.isfinite(out).all() + + +class TestCFMSampleDummy: + """CFM.sample drives DiT.forward_with_cfg through the integration loop.""" + + def test_sample_shape_and_finite(self) -> None: + cfm = CFM(_make_dit(), steps=_STEPS, sway_sampling_coef=-1.0).eval() + bsz = 1 + llm_cond = torch.randn(bsz, 1, _LLM_HIDDEN) + lat_cond = torch.randn(bsz, _HIS_PATCH_SIZE, _LATENT_DIM) + y0 = torch.randn(bsz, _PATCH_SIZE, _LATENT_DIM) + # Grid used by the talker; must span [0, 1] inclusive. + t = torch.linspace(0.0, 1.0, _STEPS + 1) + sde_args = torch.tensor([2.0, 0.0, 0.0]) # cfg=2.0, sigma=0, temp=0 + sde_rnd = torch.zeros(_STEPS, bsz, _PATCH_SIZE, _LATENT_DIM) + + with torch.no_grad(): + out = cfm.sample(llm_cond, lat_cond, y0, t, sde_args, sde_rnd) + + assert out.shape == y0.shape + assert torch.isfinite(out).all() + + def test_sample_zero_cfg_reduces_to_unguided(self) -> None: + """With cfg=0 the guidance term drops, but output shape is still valid.""" + cfm = CFM(_make_dit(), steps=_STEPS, sway_sampling_coef=None).eval() + bsz = 2 + llm_cond = torch.randn(bsz, 1, _LLM_HIDDEN) + lat_cond = torch.randn(bsz, _HIS_PATCH_SIZE, _LATENT_DIM) + y0 = torch.zeros(bsz, _PATCH_SIZE, _LATENT_DIM) + t = torch.linspace(0.0, 1.0, _STEPS + 1) + sde_args = torch.tensor([0.0, 0.0, 0.0]) + sde_rnd = torch.zeros(_STEPS, bsz, _PATCH_SIZE, _LATENT_DIM) + + with torch.no_grad(): + out = cfm.sample(llm_cond, lat_cond, y0, t, sde_args, sde_rnd) + + assert out.shape == (bsz, _PATCH_SIZE, _LATENT_DIM) + assert torch.isfinite(out).all() + + +class TestTalkerPipelineDummyWiring: + """End-to-end wiring of DiT -> CFM.sample -> Aggregator with dummy weights.""" + + def test_cfm_then_aggregator(self) -> None: + dit = _make_dit().eval() + cfm = CFM(dit, steps=_STEPS, sway_sampling_coef=-1.0).eval() + agg = _make_aggregator().eval() + + bsz = 1 + llm_cond = torch.randn(bsz, 1, _LLM_HIDDEN) + lat_cond = torch.randn(bsz, _HIS_PATCH_SIZE, _LATENT_DIM) + y0 = torch.randn(bsz, _PATCH_SIZE, _LATENT_DIM) + t = torch.linspace(0.0, 1.0, _STEPS + 1) + sde_args = torch.tensor([2.0, 0.0, 0.0]) + sde_rnd = torch.zeros(_STEPS, bsz, _PATCH_SIZE, _LATENT_DIM) + + with torch.no_grad(): + gen_lat = cfm.sample(llm_cond, lat_cond, y0, t, sde_args, sde_rnd) + agg_out = agg(gen_lat) + + assert gen_lat.shape == (bsz, _PATCH_SIZE, _LATENT_DIM) + assert agg_out.shape == (bsz, 1, _LLM_HIDDEN) + assert torch.isfinite(agg_out).all() diff --git a/tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py b/tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py index 847adae06..c7b023361 100644 --- a/tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py +++ b/tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py @@ -137,7 +137,7 @@ class SyntheticModel(nn.Module): def compute_mm_logits( self, hidden_states: torch.Tensor, - mm_sampling_tensors=None, + cfg_alpha: torch.Tensor, ): """Eager fallback path: replicate what the wrapper does.""" at = self.acoustic_transformer @@ -216,6 +216,10 @@ def _random_hidden(batch_size, device=DEVICE, dtype=torch.bfloat16): return torch.randn(batch_size, HIDDEN_DIM, device=device, dtype=dtype) +def _cfg_alpha(batch_size, value=1.2, device=DEVICE): + return torch.full((batch_size,), value, device=device, dtype=torch.float32) + + def _unpack_audio_codes(result): """Unpack (fake_eos, {"audio": [list of tensors]}) into (fake_eos, audio_codes).""" fake_eos, mm_tokens = result @@ -235,7 +239,7 @@ def test_exact_size_output_format(model, wrapper, batch_size): """Graph path returns correctly shaped and bounded outputs.""" hidden = _random_hidden(batch_size) with torch.no_grad(): - graph_eos, graph_codes = _unpack_audio_codes(wrapper(hidden)) + graph_eos, graph_codes = _unpack_audio_codes(wrapper(hidden, cfg_alpha=_cfg_alpha(batch_size))) assert graph_eos.shape == (batch_size,) assert graph_codes.shape == (batch_size, 1 + N_ACOUSTIC_CODEBOOK) # fake_eos should be 0.0 or 1.0 @@ -248,11 +252,12 @@ def test_exact_size_output_format(model, wrapper, batch_size): def test_exact_size_deterministic(model, wrapper, batch_size): """Same input + same RNG state produces identical CUDA graph output.""" hidden = _random_hidden(batch_size) + cfg_alpha = _cfg_alpha(batch_size) with torch.no_grad(): torch.manual_seed(42) - eos1, codes1 = _unpack_audio_codes(wrapper(hidden)) + eos1, codes1 = _unpack_audio_codes(wrapper(hidden, cfg_alpha=cfg_alpha)) torch.manual_seed(42) - eos2, codes2 = _unpack_audio_codes(wrapper(hidden)) + eos2, codes2 = _unpack_audio_codes(wrapper(hidden, cfg_alpha=cfg_alpha)) torch.testing.assert_close(eos1, eos2, atol=0, rtol=0) torch.testing.assert_close(codes1, codes2, atol=0, rtol=0) @@ -267,7 +272,7 @@ def test_padded_output_shape(model, wrapper, batch_size): """Padded decode must return output trimmed to actual batch size.""" hidden = _random_hidden(batch_size) with torch.no_grad(): - graph_eos, graph_codes = _unpack_audio_codes(wrapper(hidden)) + graph_eos, graph_codes = _unpack_audio_codes(wrapper(hidden, cfg_alpha=_cfg_alpha(batch_size))) assert graph_eos.shape == (batch_size,) assert graph_codes.shape == (batch_size, 1 + N_ACOUSTIC_CODEBOOK) @@ -277,7 +282,7 @@ def test_padded_output_bounded(model, wrapper, batch_size): """Padded output audio codes should be non-negative integers.""" hidden = _random_hidden(batch_size) with torch.no_grad(): - graph_eos, graph_codes = _unpack_audio_codes(wrapper(hidden)) + graph_eos, graph_codes = _unpack_audio_codes(wrapper(hidden, cfg_alpha=_cfg_alpha(batch_size))) # fake_eos should be 0.0 or 1.0 assert torch.all((graph_eos == 0.0) | (graph_eos == 1.0)) # Audio codes should be non-negative @@ -293,11 +298,12 @@ def test_padded_output_bounded(model, wrapper, batch_size): def test_fallback_eager_exact_match(model, wrapper, batch_size): """Cudagraph fallback to eager. Two eager runs should produce identical results.""" hidden = _random_hidden(batch_size) + alpha = _cfg_alpha(batch_size) with torch.no_grad(): torch.manual_seed(100) - eager_eos, eager_codes = _unpack_audio_codes(model.compute_mm_logits(hidden)) + eager_eos, eager_codes = _unpack_audio_codes(model.compute_mm_logits(hidden, cfg_alpha=alpha)) torch.manual_seed(100) - graph_eos, graph_codes = _unpack_audio_codes(wrapper(hidden)) + graph_eos, graph_codes = _unpack_audio_codes(wrapper(hidden, cfg_alpha=alpha)) torch.testing.assert_close(graph_eos, eager_eos, atol=0, rtol=0) torch.testing.assert_close(graph_codes, eager_codes, atol=0, rtol=0) @@ -310,12 +316,13 @@ def test_fallback_eager_exact_match(model, wrapper, batch_size): def test_disabled_wrapper_matches_eager(model, wrapper): """Cudagraph fallback to eager. Two eager runs should produce identical results.""" hidden = _random_hidden(4) + alpha = _cfg_alpha(4) wrapper.enabled = False with torch.no_grad(): torch.manual_seed(200) - eager_eos, eager_codes = _unpack_audio_codes(model.compute_mm_logits(hidden)) + eager_eos, eager_codes = _unpack_audio_codes(model.compute_mm_logits(hidden, cfg_alpha=alpha)) torch.manual_seed(200) - graph_eos, graph_codes = _unpack_audio_codes(wrapper(hidden)) + graph_eos, graph_codes = _unpack_audio_codes(wrapper(hidden, cfg_alpha=alpha)) wrapper.enabled = True torch.testing.assert_close(graph_eos, eager_eos, atol=0, rtol=0) torch.testing.assert_close(graph_codes, eager_codes, atol=0, rtol=0) @@ -329,10 +336,11 @@ def test_disabled_wrapper_matches_eager(model, wrapper): def test_deterministic_across_calls(model, wrapper): """Same input + same RNG state. Two cudagraph runs should produce identical results.""" hidden = _random_hidden(4) + alpha = _cfg_alpha(4) with torch.no_grad(): torch.manual_seed(300) - eos1, codes1 = _unpack_audio_codes(wrapper(hidden)) + eos1, codes1 = _unpack_audio_codes(wrapper(hidden, cfg_alpha=alpha)) torch.manual_seed(300) - eos2, codes2 = _unpack_audio_codes(wrapper(hidden)) + eos2, codes2 = _unpack_audio_codes(wrapper(hidden, cfg_alpha=alpha)) torch.testing.assert_close(eos1, eos2, atol=0, rtol=0) torch.testing.assert_close(codes1, codes2, atol=0, rtol=0) diff --git a/tests/test_arg_utils.py b/tests/test_arg_utils.py index dab5ed687..ae640b2d8 100644 --- a/tests/test_arg_utils.py +++ b/tests/test_arg_utils.py @@ -351,3 +351,138 @@ def test_ambiguous_field_non_strict_routes_to_orchestrator(caplog): assert orch.deploy_config == "x" assert "deploy_config" not in engine assert any("both OrchestratorArgs" in r.message for r in caplog.records) + + +# Sentinel-default precedence invariants (#3035) + + +def _build_full_serve_parser(): + from vllm.utils.argparse_utils import FlexibleArgumentParser + + try: + from vllm.entrypoints.openai.cli_args import make_arg_parser + except ImportError: + pytest.skip("vllm parser not importable") + return make_arg_parser(FlexibleArgumentParser()) + + +def test_nullify_stage_engine_defaults_resets_inherited_defaults(): + import argparse + + from vllm_omni.engine.arg_utils import ( + deploy_override_field_names, + nullify_stage_engine_defaults, + ) + + parser = _build_full_serve_parser() + nullify_stage_engine_defaults(parser) + + override_dests = deploy_override_field_names() + offenders = [ + (a.dest, a.default) + for a in parser._actions + if a.dest not in ("help", "version") + and a.option_strings + and a.dest in override_dests + and a.default is not None + and a.default is not argparse.SUPPRESS + ] + assert not offenders, f"Stage flags with non-None defaults after nullify: {offenders}" + + +def test_non_override_flags_keep_real_defaults_after_nullify(): + import argparse + + from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults + + parser = argparse.ArgumentParser() + parser.add_argument("--hsdp-shard-size", type=int, default=-1, help="HSDP shard size.") + parser.add_argument("--max-num-seqs", type=int, default=64, help="Max num seqs.") + nullify_stage_engine_defaults(parser) + + hsdp = next(a for a in parser._actions if a.dest == "hsdp_shard_size") + max_num_seqs = next(a for a in parser._actions if a.dest == "max_num_seqs") + assert hsdp.default == -1 + assert max_num_seqs.default is None + + +def test_help_text_preserves_default_after_nullify(): + # Real defaults must stay visible in --help even though parser stores None. + import argparse + + from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults + + parser = argparse.ArgumentParser() + parser.add_argument("--max-num-seqs", type=int, default=42, help="Example knob.") + nullify_stage_engine_defaults(parser) + + action = next(a for a in parser._actions if a.dest == "max_num_seqs") + assert action.default is None + assert "(default: 42)" in action.help + + +_OMNIENGINEARGS_USER_INPUT_FIELDS = frozenset( + { + "model_stage", + "model_arch", + "engine_output_type", + "hf_config_name", + "custom_process_next_stage_input_func", + "subtalker_sampling_params", + "async_chunk", + "omni_kv_config", + "quantization_config", + "worker_type", + "task_type", + "worker_cls", + "enable_sleep_mode", + "omni_master_address", + "omni_master_port", + "stage_configs_path", + "output_modalities", + "log_stats", + "custom_pipeline_args", + } +) + + +def test_omniengineargs_user_input_fields_default_to_none(): + try: + from vllm_omni.engine.arg_utils import OmniEngineArgs + except Exception as exc: + pytest.skip(f"OmniEngineArgs not importable: {exc}") + + offenders = [ + (f.name, f.default) + for f in fields(OmniEngineArgs) + if f.name in _OMNIENGINEARGS_USER_INPUT_FIELDS + and f.default is not dataclasses.MISSING + and f.default is not None + ] + assert not offenders, f"User-input fields with non-None defaults: {offenders}" + + +def test_omniengineargs_create_tracks_explicit_fields(): + try: + from vllm_omni.engine.arg_utils import OmniEngineArgs + except Exception as exc: + pytest.skip(f"OmniEngineArgs not importable: {exc}") + + ea = OmniEngineArgs.create(model="x", gpu_memory_utilization=0.5) + assert ea._explicit_fields == frozenset({"model", "gpu_memory_utilization"}) + assert ea.explicit_kwargs() == {"model": "x", "gpu_memory_utilization": 0.5} + + +def test_omniengineargs_bare_constructor_has_no_explicit_tracking(): + try: + from vllm_omni.engine.arg_utils import OmniEngineArgs + except Exception as exc: + pytest.skip(f"OmniEngineArgs not importable: {exc}") + + ea = OmniEngineArgs(model="x") + assert not hasattr(ea, "_explicit_fields") + assert "model" in ea.explicit_kwargs() + + +# dataclasses already imported via ``from dataclasses import dataclass, fields`` +import dataclasses # noqa: E402 -- needed for MISSING sentinel above diff --git a/tests/test_config_factory.py b/tests/test_config_factory.py index 6cf8dcd40..3148d8dc6 100644 --- a/tests/test_config_factory.py +++ b/tests/test_config_factory.py @@ -4,6 +4,7 @@ Unit tests for StageConfigFactory and related classes. """ +import warnings from dataclasses import dataclass from pathlib import Path @@ -375,6 +376,31 @@ class TestStageResolutionHelpers: } assert overridden == ["tensor_parallel_size"] + def test_strip_parent_engine_args_keeps_allowed_media_access_controls(self): + from dataclasses import fields as dc_fields + + from vllm.engine.arg_utils import EngineArgs + + parent_fields = {f.name: f for f in dc_fields(EngineArgs)} + filtered, overridden = strip_parent_engine_args( + { + "model": "some/model", + "stage_configs_path": "/tmp/stages.yaml", + "allowed_local_media_path": "/data/qwentts", + "allowed_media_domains": ["example.com"], + }, + parent_fields=parent_fields, + keep_keys={"allowed_local_media_path", "allowed_media_domains"}, + strip_keys={"stage_configs_path"}, + no_warn_keys={"model"}, + ) + + assert filtered == { + "allowed_local_media_path": "/data/qwentts", + "allowed_media_domains": ["example.com"], + } + assert overridden == [] + class TestPipelineYamlParsing: """Tests for pipeline YAML file parsing (@ZJY0516).""" @@ -1171,99 +1197,89 @@ class TestCLIOverrideFlow: assert s.runtime_overrides["enforce_eager"] is True -class TestCLIExplicitPrecedence: - """Verify YAML > argparse defaults; explicit CLI args > YAML.""" +class TestSentinelDefaultPrecedence: + """Caller-typed (non-None) values win over YAML; None values fall through + to YAML / dataclass defaults (#3035).""" - def _stages(self, cli_overrides, cli_explicit_keys): + def _stages(self, cli_overrides): import vllm_omni.model_executor.models.qwen3_omni.pipeline # noqa: F401 return StageConfigFactory._create_from_registry( "qwen3_omni_moe", cli_overrides=cli_overrides, - cli_explicit_keys=cli_explicit_keys, ) - def test_explicit_cli_overrides_yaml(self): - """User-typed --max-num-seqs wins over the deploy YAML value.""" - stages = self._stages( - cli_overrides={"max_num_seqs": 999}, - cli_explicit_keys={"max_num_seqs"}, - ) - # Stage 2 yaml has max_num_seqs=1; explicit CLI must beat it. + def test_typed_kwarg_overrides_yaml(self): + stages = self._stages({"max_num_seqs": 999}) assert stages[2].runtime_overrides.get("max_num_seqs") == 999 - def test_default_cli_does_not_override_yaml(self): - """Argparse defaults must NOT clobber values that are present in YAML.""" - stages = self._stages( - cli_overrides={"max_num_seqs": 256}, - cli_explicit_keys=set(), # user typed nothing - ) - # Stage 2's YAML value (1) should win because the user didn't type --max-num-seqs. - assert stages[2].runtime_overrides.get("max_num_seqs") != 256 + def test_none_value_skipped_yaml_wins(self): + stages = self._stages({"max_num_seqs": None}) + assert stages[2].runtime_overrides.get("max_num_seqs") is None + assert stages[2].yaml_engine_args.get("max_num_seqs") == 1 - def test_default_cli_fills_missing_yaml_field(self): - """Argparse defaults still fill fields the YAML doesn't set.""" - stages = self._stages( - cli_overrides={"some_unrelated_knob": "fallback"}, - cli_explicit_keys=set(), - ) - # Field absent from YAML → CLI default flows through as a fallback. - assert stages[0].runtime_overrides.get("some_unrelated_knob") == "fallback" + def test_empty_kwargs_yaml_only(self): + stages = self._stages({}) + for stage in stages: + assert stage.runtime_overrides == {} - def test_per_stage_overrides_always_explicit(self): - """``stage__*`` keys are always treated as explicit.""" - stages = self._stages( - cli_overrides={"stage_0_gpu_memory_utilization": 0.42}, - cli_explicit_keys=set(), # not in the explicit set, but per-stage - ) + def test_typed_kwarg_equal_to_dataclass_default_still_overrides(self): + # Caller intent honored regardless of value coincidence (no heuristic). + stages = self._stages({"gpu_memory_utilization": 0.9}) + assert stages[2].runtime_overrides.get("gpu_memory_utilization") == 0.9 + + def test_per_stage_kwarg_routed_to_correct_stage(self): + stages = self._stages({"stage_0_gpu_memory_utilization": 0.42}) assert stages[0].runtime_overrides.get("gpu_memory_utilization") == 0.42 + assert stages[2].runtime_overrides.get("gpu_memory_utilization") is None - def test_none_explicit_set_treats_all_as_explicit(self): - """Programmatic Omni() callers (cli_explicit_keys=None) keep current behavior.""" - stages = self._stages( - cli_overrides={"max_num_seqs": 999}, - cli_explicit_keys=None, - ) - assert stages[2].runtime_overrides.get("max_num_seqs") == 999 - - def test_explicit_async_chunk_false_overrides_yaml(self): - """``--no-async-chunk`` flips the deploy-level async_chunk to False even - when the YAML sets it to True. Verifies that the per-stage - ``async_chunk: True`` injection in ``merge_pipeline_deploy`` is skipped - and that ``async_chunk`` does not leak through ``_merge_cli_overrides``. - """ - stages = self._stages( - cli_overrides={"async_chunk": False}, - cli_explicit_keys={"async_chunk"}, - ) - # qwen3_omni_moe.yaml has `async_chunk: true`, so by default every - # stage's engine_args would carry it. With the explicit override, it - # must NOT show up. + def test_async_chunk_false_overrides_yaml_true(self): + stages = self._stages({"async_chunk": False}) for stage in stages: assert stage.yaml_engine_args.get("async_chunk") is not True - assert stage.runtime_overrides.get("async_chunk") is None - def test_default_async_chunk_leaves_yaml_alone(self): - """An unset ``--async-chunk`` (default None) must leave the YAML's True - in force on every stage.""" - stages = self._stages( - cli_overrides={"async_chunk": None}, - cli_explicit_keys=set(), - ) - # qwen3_omni_moe.yaml: `async_chunk: true` → injected on every stage. + def test_async_chunk_none_keeps_yaml_true(self): + stages = self._stages({"async_chunk": None}) for stage in stages: assert stage.yaml_engine_args.get("async_chunk") is True - def test_explicit_enable_prefix_caching_overrides_yaml(self): - """``--enable-prefix-caching`` (global) flips every stage's - ``enable_prefix_caching`` to True regardless of the YAML default.""" - stages = self._stages( - cli_overrides={"enable_prefix_caching": True}, - cli_explicit_keys={"enable_prefix_caching"}, - ) + def test_enable_prefix_caching_typed_overrides_yaml(self): + stages = self._stages({"enable_prefix_caching": True}) for stage in stages: assert stage.runtime_overrides.get("enable_prefix_caching") is True + def test_omni_with_vars_args_anti_pattern_is_safe(self): + # Omni(**vars(args)) with mostly-None namespace must not clobber YAML. + simulated_vars_args = { + "gpu_memory_utilization": None, + "max_num_seqs": None, + "async_chunk": None, + "enable_prefix_caching": None, + "dtype": None, + } + stages = self._stages(simulated_vars_args) + for stage in stages: + assert stage.runtime_overrides == {} + + def test_create_from_registry_no_cli_explicit_keys_param(self): + import inspect + + sig = inspect.signature(StageConfigFactory._create_from_registry) + named = [p for p in sig.parameters.values() if p.kind != p.VAR_KEYWORD] + assert "cli_explicit_keys" not in {p.name for p in named} + + def test_cli_explicit_keys_kwarg_emits_deprecation(self): + import vllm_omni.model_executor.models.qwen3_omni.pipeline # noqa: F401 + + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + StageConfigFactory._create_from_registry( + "qwen3_omni_moe", + cli_overrides={}, + cli_explicit_keys={"max_num_seqs"}, + ) + assert any(issubclass(x.category, DeprecationWarning) for x in w) + def test_async_chunk_dispatches_processors(self): """A single ``qwen3_tts`` pipeline picks per-chunk vs end-to-end processors based on ``deploy.async_chunk``, without needing a diff --git a/vllm_omni/benchmarks/data_modules/daily_omni_dataset.py b/vllm_omni/benchmarks/data_modules/daily_omni_dataset.py index 01b86d0fd..65918414f 100644 --- a/vllm_omni/benchmarks/data_modules/daily_omni_dataset.py +++ b/vllm_omni/benchmarks/data_modules/daily_omni_dataset.py @@ -9,7 +9,9 @@ Supports loading QA metadata from: - Local JSON file (``qa_json_path``): recommended for offline/air-gapped environments - HuggingFace datasets (``dataset_path``): legacy online mode -The videos must be separately downloaded and extracted from Videos.tar. +Video/audio files normally come from extracted ``Videos.tar``. When ``--daily-omni-video-dir`` +is not set, the first request that needs on-disk media downloads that archive from the Hugging Face +dataset repo (``huggingface_hub``) and caches it under ``HF_HOME``. Why ``BenchmarkDataset`` instead of ``HuggingFaceDataset``? vLLM's ``HuggingFaceDataset`` is a thin wrapper whose ``__init__`` always ends by calling @@ -52,6 +54,9 @@ Usage: import base64 import json import logging +import os +import shutil +import tarfile from dataclasses import dataclass from pathlib import Path from typing import Any, Literal @@ -73,6 +78,103 @@ except ImportError: logger = logging.getLogger(__name__) +def _daily_omni_hf_cache_root() -> Path: + return Path(os.environ.get("HF_HOME", Path.home() / ".cache" / "huggingface")).expanduser().resolve() + + +def _daily_omni_tar_fingerprint(tar_path: Path) -> str: + st = tar_path.stat() + return f"v1:{st.st_size}:{int(st.st_mtime_ns)}" + + +def _daily_omni_find_videos_root_in_extract(tmp: Path) -> Path: + """Return directory whose children are ``video_id`` folders with ``*_video.mp4``.""" + videos = tmp / "Videos" + if videos.is_dir(): + return videos + for child in sorted(tmp.iterdir()): + if child.is_dir() and not child.name.startswith("."): + probe = child / f"{child.name}_video.mp4" + if probe.is_file(): + return tmp + raise RuntimeError( + f"Unrecognized layout after extracting Daily-Omni Videos.tar under {tmp} " + "(expected top-level 'Videos/' or per-video_id subdirs)." + ) + + +def ensure_daily_omni_hub_videos_dir(repo_id: str) -> Path: + """Download ``Videos.tar`` from the Hugging Face dataset repo and return the ``Videos`` root. + + The returned path matches ``--daily-omni-video-dir`` (directory containing ``{{video_id}}/``). + + Cached under ``HF_HOME`` / ``vllm_omni/daily_omni_media/``. Reuses extraction when the + tarball fingerprint matches. + + Raises: + ImportError: if ``huggingface_hub`` is not installed. + FileNotFoundError / RuntimeError: if the archive is missing or malformed. + """ + rid = (repo_id or "").strip() + if not rid: + raise ValueError("repo_id is required to download Daily-Omni Videos.tar") + + try: + from huggingface_hub import hf_hub_download + except ImportError as e: + raise ImportError( + "Daily-Omni Hub media download requires huggingface_hub. " + "Install it (e.g. with vLLM) or provide --daily-omni-video-dir with a local extract." + ) from e + + safe = rid.replace("/", "__").replace("\\", "_") + staging_root = _daily_omni_hf_cache_root() / "vllm_omni" / "daily_omni_media" / safe + videos_dir = staging_root / "Videos" + marker = staging_root / ".videos_extracted" + + tar_path: Path | None = None + for fname in ("Videos.tar", "videos.tar"): + try: + tar_path = Path(hf_hub_download(repo_id=rid, filename=fname, repo_type="dataset")) + break + except Exception: + continue + if tar_path is None or not tar_path.is_file(): + raise FileNotFoundError( + f"Could not download Videos.tar from Hugging Face dataset {rid!r} (tried Videos.tar / videos.tar)." + ) + + fp = _daily_omni_tar_fingerprint(tar_path) + if marker.is_file() and videos_dir.is_dir(): + try: + if marker.read_text(encoding="utf-8").strip() == fp: + next(videos_dir.iterdir()) + logger.info("Reusing cached Daily-Omni Videos extract at %s", videos_dir) + return videos_dir + except (OSError, StopIteration): + shutil.rmtree(videos_dir, ignore_errors=True) + marker.unlink(missing_ok=True) + + staging_root.mkdir(parents=True, exist_ok=True) + work = staging_root / "_extract_work" + shutil.rmtree(work, ignore_errors=True) + work.mkdir(parents=True) + try: + logger.info("Extracting Daily-Omni Videos.tar from %s (repo=%s)", tar_path, rid) + with tarfile.open(tar_path, "r:*") as tf: + tf.extractall(path=work, filter="data") + found = _daily_omni_find_videos_root_in_extract(work) + if videos_dir.exists(): + shutil.rmtree(videos_dir, ignore_errors=True) + shutil.move(str(found), str(videos_dir)) + finally: + shutil.rmtree(work, ignore_errors=True) + + marker.write_text(fp, encoding="utf-8") + logger.info("Daily-Omni Hub media ready at %s", videos_dir) + return videos_dir + + class _ListDatasetIterator: """Simple iterator wrapper around a list to mimic HuggingFace streaming dataset behavior.""" @@ -142,7 +244,10 @@ class DailyOmniDataset(BenchmarkDataset): - Local JSON file (``qa_json_path``): recommended for offline/air-gapped environments - HuggingFace datasets (``dataset_path``): legacy online mode - The videos must be separately downloaded and extracted from Videos.tar. + Video/audio files normally come from extracted ``Videos.tar``. When ``video_dir`` is not set, + the first sample that needs on-disk media downloads that archive from the Hugging Face dataset + repo (env ``VLLM_DAILY_OMNI_MEDIA_REPO`` overrides the repo id; else ``dataset_path`` or + :data:`DEFAULT_HF_DATASET_ID`). Args: qa_json_path: Path to local qa.json file (offline mode, preferred). When provided, @@ -151,7 +256,8 @@ class DailyOmniDataset(BenchmarkDataset): ``qa_json_path`` is not provided (legacy online mode). dataset_split: Dataset split to use (default: "train"). Used only in online mode. random_seed: Random seed for shuffling - video_dir: Directory containing extracted video files (default: None) + video_dir: Directory containing extracted video files (default: None; may be filled lazily + from Hub — see above). input_mode: Which modalities to send, matching upstream Daily-Omni ``--input_mode``: ``all`` — video + WAV (default; official audio-visual protocol); ``visual`` — video only; @@ -308,13 +414,14 @@ class DailyOmniDataset(BenchmarkDataset): "Install with: pip install datasets, or use local JSON mode instead." ) - ds = load_dataset( - self.dataset_path, - name=self.dataset_subset, - split=self.dataset_split, - streaming=self._hf_streaming, - trust_remote_code=self.trust_remote_code, - ) + load_kw: dict[str, Any] = { + "split": self.dataset_split, + "streaming": self._hf_streaming, + "trust_remote_code": self.trust_remote_code, + } + if self.dataset_subset is not None: + load_kw["name"] = self.dataset_subset + ds = load_dataset(self.dataset_path, **load_kw) if not getattr(self, "disable_shuffle", False): ds = ds.shuffle(seed=self.random_seed) self.data = ds @@ -569,6 +676,17 @@ class DailyOmniDataset(BenchmarkDataset): "audio_url": {"url": path.as_uri()}, } + def _lazy_ensure_hub_media_dir(self) -> None: + """If ``video_dir`` was not configured, download and extract ``Videos.tar`` once from HF.""" + if self.video_dir is not None: + return + repo = os.environ.get("VLLM_DAILY_OMNI_MEDIA_REPO", "").strip() + if not repo: + repo = (self.dataset_path or "").strip() + if not repo: + repo = self.DEFAULT_HF_DATASET_ID + self.video_dir = ensure_daily_omni_hub_videos_dir(repo) + def _get_video_content( self, video_id: str, @@ -589,6 +707,8 @@ class DailyOmniDataset(BenchmarkDataset): url = f"https://{url.lstrip('/')}" return {"type": "video_url", "video_url": {"url": url}} + self._lazy_ensure_hub_media_dir() + if self.video_dir and video_id: video_path = self._resolve_local_video_path(video_id) if video_path is not None: @@ -615,7 +735,13 @@ class DailyOmniDataset(BenchmarkDataset): return None def _get_audio_content(self, video_id: str) -> dict[str, Any] | None: - """Resolve extracted WAV for OpenAI-style ``audio_url`` (local files only).""" + """Resolve extracted WAV for OpenAI-style ``audio_url`` (local files under ``video_dir``). + + Uses the same tree as video (``{video_id}/{video_id}_audio.wav``), including after lazy + Hub ``Videos.tar`` extraction when ``video_dir`` was unset. + """ + self._lazy_ensure_hub_media_dir() + if not self.video_dir or not video_id: logger.warning( "Daily-Omni input_mode %r requires --daily-omni-video-dir with %s", diff --git a/vllm_omni/benchmarks/data_modules/daily_omni_eval.py b/vllm_omni/benchmarks/data_modules/daily_omni_eval.py index ecc9edc84..f191cf2fe 100644 --- a/vllm_omni/benchmarks/data_modules/daily_omni_eval.py +++ b/vllm_omni/benchmarks/data_modules/daily_omni_eval.py @@ -5,7 +5,8 @@ Compares model ``generated_text`` to dataset ``Answer`` (A/B/C/D). **Alignment with open-source** (`Lliar-liar/Daily-Omni` ``test_model/.../testmodel.py``): - Answer extraction defaults to the same rules as ``extract_choice_letter`` (strip after an - ``assistant`` marker, then leading ``A``–``D``, else first ``\\b[A-D]\\b``). Set env + ``assistant`` marker, then leading ``A``–``D``, else ``\\b[A-D]\\b``, else a CJK-safe + non-letter-boundary pass). Set env ``DAILY_OMNI_EXTRACT_MODE=relaxed`` to use the older vLLM-Omni heuristics (last ``answer:``, tail scan, etc.). - Overall accuracy comparable to the official script uses **successful HTTP responses only** as @@ -52,6 +53,16 @@ def extract_choice_letter_official(text: str | None) -> str | None: fallback = re.search(r"\b([A-D])\b", candidate.upper()) if fallback: return fallback.group(1) + # ``\b`` is ASCII/Latin-word-centric; CJK (e.g. "选B", "答案：B") has no boundary before B. + loose = list( + re.finditer( + r"(?:[^A-Za-z]|^)([A-D])(?:[^A-Za-z]|$)", + candidate, + flags=re.IGNORECASE, + ) + ) + if loose: + return loose[-1].group(1).upper() return None diff --git a/vllm_omni/benchmarks/data_modules/seed_tts_dataset.py b/vllm_omni/benchmarks/data_modules/seed_tts_dataset.py index ca6de4cb2..4f072e1a2 100644 --- a/vllm_omni/benchmarks/data_modules/seed_tts_dataset.py +++ b/vllm_omni/benchmarks/data_modules/seed_tts_dataset.py @@ -94,7 +94,7 @@ def _load_meta_rows(meta_file: Path) -> list[_SeedTTSRow]: return rows -def resolve_seed_tts_root(dataset_path: str | None, *, explicit_root: str | None) -> Path: +def resolve_seed_tts_root(dataset_path: str | None, *, explicit_root: str | None, locale: str = "en") -> Path: """Return directory containing ``{locale}/meta.lst`` and ``{locale}/prompt-wavs/``.""" if explicit_root: root = Path(explicit_root).expanduser().resolve() @@ -117,7 +117,14 @@ def resolve_seed_tts_root(dataset_path: str | None, *, explicit_root: str | None "Install huggingface_hub to download Seed-TTS from the Hub, or clone the dataset " "locally and pass --dataset-path / --seed-tts-root to that directory." ) from e - cache = snapshot_download(repo_id=repo_id, repo_type="dataset") + # Download only the requested locale subtree instead of the whole dataset + # repo. This avoids large, flaky nightly downloads when we only need e.g. + # ``en/meta.lst`` + ``en/prompt-wavs/**``. + cache = snapshot_download( + repo_id=repo_id, + repo_type="dataset", + allow_patterns=[f"{locale}/**"], + ) return Path(cache).resolve() @@ -172,7 +179,7 @@ class SeedTTSDataset(BenchmarkDataset): disable_shuffle=disable_shuffle, **kwargs, ) - self._root = resolve_seed_tts_root(self.dataset_path, explicit_root=self._explicit_root) + self._root = resolve_seed_tts_root(self.dataset_path, explicit_root=self._explicit_root, locale=self.locale) self._rows: list[_SeedTTSRow] = [] self.load_data() diff --git a/vllm_omni/benchmarks/data_modules/seed_tts_eval.py b/vllm_omni/benchmarks/data_modules/seed_tts_eval.py index d5f1b6470..d8c37af13 100644 --- a/vllm_omni/benchmarks/data_modules/seed_tts_eval.py +++ b/vllm_omni/benchmarks/data_modules/seed_tts_eval.py @@ -1,7 +1,7 @@ """Seed-TTS WER aligned with Bytedance ``seed-tts-eval`` / ``run_wer.py``. Matches the published protocol (see Hugging Face dataset card and -https://github.com/BytedanceSpeech/seed-tts-eval): +https://github.com/zhaochenyang20/seed-tts-eval): - **EN**: ``openai/whisper-large-v3`` via ``transformers``, audio resampled to **16 kHz** (same as ``run_wer.py``). @@ -489,7 +489,7 @@ def compute_seed_tts_wer_metrics( sim_failed = 0 sim_skipped_no_ref = 0 utmos_failed = 0 - utmos_on = _eval_submetric_enabled("SEED_TTS_UTMOS_EVAL", default=True) + utmos_on = _eval_submetric_enabled("SEED_TTS_UTMOS_EVAL", default=False) for req, out in zip(input_requests, outputs, strict=True): assert isinstance(req, SeedTTSSampleRequest) @@ -622,7 +622,7 @@ def compute_seed_tts_wer_metrics( errs.append(wer) sim_v: float | None = None - if _eval_submetric_enabled("SEED_TTS_SIM_EVAL", default=True): + if _eval_submetric_enabled("SEED_TTS_SIM_EVAL", default=False): ref_path = getattr(req, "seed_tts_ref_wav_path", "") or "" if ref_path and os.path.isfile(ref_path): try: diff --git a/vllm_omni/config/model.py b/vllm_omni/config/model.py index 96a34a8d7..72731eda8 100644 --- a/vllm_omni/config/model.py +++ b/vllm_omni/config/model.py @@ -5,7 +5,10 @@ from pydantic import ConfigDict, TypeAdapter from vllm.config import ModelConfig from vllm.config.utils import config from vllm.logger import init_logger -from vllm.transformers_utils.config import get_hf_text_config +from vllm.transformers_utils.config import ( + get_hf_text_config, + thinker_uses_mrope, +) from vllm.transformers_utils.model_arch_config_convertor import ( ModelArchConfigConvertorBase, ) @@ -49,6 +52,25 @@ class OmniModelArchConfigConvertor(ModelArchConfigConvertorBase): if quant_cfg is not None: return quant_cfg + # Fall back to top-level quantization_config + top_quant = super().get_quantization_config() + if top_quant is not None: + block_names = top_quant.get("block_name_to_quantize") + if block_names is not None: + # NOTE: This assumes stage_config_name follows the HF + # ``_config`` convention (e.g. thinker_config → + # prefix "thinker."). removesuffix is a no-op when + # the suffix doesn't match, so a non-standard name + # would just use itself as prefix — safe but worth + # verifying if new stage names are introduced. + hf_prefix = self.stage_config_name.removesuffix("_config") + "." + if isinstance(block_names, str): + block_names = [b.strip() for b in block_names.split(",")] + if isinstance(block_names, list) and not any(b.startswith(hf_prefix) for b in block_names): + # This stage is not listed → no quantization. + return None + return top_quant + # For non-thinker stages (talker, code2wav) whose text_config # has no quantization_config, return None so quantization is # not applied to stages that were not quantized. @@ -114,6 +136,7 @@ class OmniModelConfig(ModelConfig): codec_frame_rate_hz: float | None = None task_type: str | None = None enable_sleep_mode: bool = False + has_sampling_extra_args: bool = False @property def registry(self): @@ -125,6 +148,18 @@ class OmniModelConfig(ModelConfig): return [self.model_arch] return super().architectures + @property + def uses_mrope(self) -> bool: + if self.hf_config_name is not None: + # talker_config/thinker_config/etc + stage_config = getattr(self.hf_config, self.hf_config_name, None) + if stage_config is None: + # Check the named sub-config's text_config directly. + # Handles mrope resolution of stage-specific cls + # (e.g., talker runs as a standalone cls) + return thinker_uses_mrope(self.hf_config) + return super().uses_mrope + @property def embedding_size(self): if self.hf_config_name is not None: diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index c06cc90e4..75665f661 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/vllm_omni/config/pipeline_registry.py @@ -43,6 +43,22 @@ _OMNI_PIPELINES: dict[str, tuple[str, str]] = { "vllm_omni.model_executor.models.qwen3_tts.pipeline", "QWEN3_TTS_PIPELINE", ), + "bagel": ( + "vllm_omni.model_executor.models.bagel.pipeline", + "BAGEL_PIPELINE", + ), + "bagel_think": ( + "vllm_omni.model_executor.models.bagel.pipeline", + "BAGEL_THINK_PIPELINE", + ), + "bagel_single_stage": ( + "vllm_omni.model_executor.models.bagel.pipeline", + "BAGEL_SINGLE_STAGE_PIPELINE", + ), + "glm_image": ( + "vllm_omni.model_executor.models.glm_image.pipeline", + "GLM_IMAGE_PIPELINE", + ), "voxcpm2": ( "vllm_omni.model_executor.models.voxcpm2.pipeline", "VOXCPM2_PIPELINE", diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 9839a5389..e133c2aef 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -28,6 +28,15 @@ def get_pipeline_path(model_dir: str, filename: str) -> Path: logger = init_logger(__name__) +def _warn_deprecated_kwargs(kwargs: dict[str, Any]) -> None: + if "cli_explicit_keys" in kwargs: + warnings.warn( + "cli_explicit_keys= is deprecated and ignored. Remove the kwarg.", + DeprecationWarning, + stacklevel=3, + ) + + _STAGE_OVERRIDE_PATTERN = re.compile(r"^stage_(\d+)_(.+)$") @@ -169,6 +178,12 @@ class StagePipelineConfig: prompt_expand_func: str | None = None cfg_kv_collect_func: str | None = None omni_kv_config: dict[str, Any] | None = None + # Model subdirectory indirections: for multi-component HF repos where the + # stage's config/tokenizer lives in a subdirectory (e.g. GLM-Image's AR + # config is in ``vision_language_encoder/``). Consumed at stage-init time + # by ``stage_init_utils._resolve_model_tokenizer_paths``. + model_subdir: str | None = None + tokenizer_subdir: str | None = None extras: dict[str, Any] = field(default_factory=dict) @@ -185,6 +200,11 @@ class PipelineConfig: # matches ``hf_config.architectures[*]`` against this tuple to route # to the correct pipeline. Leave empty for models with unique model_type. hf_architectures: tuple[str, ...] = () + # Diffusers pipeline class name: for models that ship a ``model_index.json`` + # (no root ``config.json``), the ``_class_name`` field is matched against + # this value to auto-detect the pipeline. Only needed for diffusers-style + # multi-component repos (e.g. GLM-Image). ``None`` = not a diffusers model. + diffusers_class_name: str | None = None def get_stage(self, stage_id: int) -> StagePipelineConfig | None: """Look up a stage by its ID.""" @@ -702,6 +722,13 @@ def _build_engine_args( engine_args["engine_output_type"] = ps.engine_output_type if next_stage_proc: engine_args["custom_process_next_stage_input_func"] = next_stage_proc + # Subdirectory indirections from StagePipelineConfig (structural, not + # deployment knobs). Deploy YAML ``engine_extras`` can still override + # these per-stage if needed. + if ps.model_subdir: + engine_args["model_subdir"] = ps.model_subdir + if ps.tokenizer_subdir: + engine_args["tokenizer_subdir"] = ps.tokenizer_subdir # Pipeline-wide top-level DeployConfig settings, applied to every stage. for name in _PIPELINE_WIDE_ENGINE_FIELDS: @@ -716,8 +743,11 @@ def _build_engine_args( continue engine_args[k] = v engine_args.update(ds.engine_extras) - if deploy.async_chunk: - engine_args["async_chunk"] = True + # Materialize the resolved pipeline-wide async_chunk value into every + # stage so explicit False overrides do not get lost downstream. + engine_args["async_chunk"] = bool(deploy.async_chunk) + if ps.omni_kv_config: + engine_args["omni_kv_config"] = dict(ps.omni_kv_config) return engine_args @@ -737,6 +767,10 @@ def _build_extras( extras["output_connectors"] = dict(ds.output_connectors) if ds is not None and ds.input_connectors: extras["input_connectors"] = dict(ds.input_connectors) + if ps.prompt_expand_func: + extras["prompt_expand_func"] = ps.prompt_expand_func + if ps.cfg_kv_collect_func: + extras["cfg_kv_collect_func"] = ps.cfg_kv_collect_func if ps.extras: extras.update(ps.extras) return extras @@ -978,17 +1012,14 @@ class StageConfigFactory: model: str, cli_overrides: dict[str, Any] | None = None, deploy_config_path: str | None = None, - cli_explicit_keys: set[str] | None = None, + **deprecated_kwargs: Any, ) -> list[StageConfig] | None: """Load pipeline + deploy config, merge with CLI overrides. Checks _PIPELINE_REGISTRY first (new path), falls back to legacy YAML. - - ``cli_explicit_keys`` is the set of CLI keys the user actually typed - (captured at the parser layer in ``vllm serve``). When ``None`` — - which is the case for programmatic ``Omni()`` callers — every kwarg - in ``cli_overrides`` is treated as explicit. """ + _warn_deprecated_kwargs(deprecated_kwargs) + if cli_overrides is None: cli_overrides = {} @@ -997,7 +1028,7 @@ class StageConfigFactory: # --- New path: check pipeline registry by model_type first --- model_type, hf_config = cls._auto_detect_model_type(model, trust_remote_code=trust_remote_code) if model_type and model_type in _PIPELINE_REGISTRY: - return cls._create_from_registry(model_type, cli_overrides, deploy_config_path, cli_explicit_keys) + return cls._create_from_registry(model_type, cli_overrides, deploy_config_path) # --- HF architecture fallback: some models report a generic # model_type that collides with another model. Match by the @@ -1007,9 +1038,7 @@ class StageConfigFactory: if hf_archs: for registered in _PIPELINE_REGISTRY.values(): if hf_archs.intersection(registered.hf_architectures): - return cls._create_from_registry( - registered.model_type, cli_overrides, deploy_config_path, cli_explicit_keys - ) + return cls._create_from_registry(registered.model_type, cli_overrides, deploy_config_path) # --- Legacy path: load from pipeline YAML --- pipeline = cls._load_pipeline(model, trust_remote_code=trust_remote_code) @@ -1021,14 +1050,14 @@ class StageConfigFactory: if errors: logger.warning(f"Pipeline validation warnings for {model}: {errors}") - # Inject pipeline-wide async_chunk into ALL stages' engine_args. - # The legacy loader (load_stage_configs_from_yaml) sets async_chunk - # on every stage so that build_engine_args_dict() can inject the - # stage_connector_spec. AsyncOmniEngine.__init__ also reads it - # from stage_configs[0].engine_args.async_chunk. - if pipeline.async_chunk: - for stage in pipeline.stages: - stage.yaml_engine_args.setdefault("async_chunk", True) + # Materialize the resolved pipeline-wide async_chunk value into every + # stage so build_engine_args_dict() can inject the stage connector + # spec and explicit False overrides are preserved. + resolved_async_chunk = cli_overrides.get("async_chunk") + if resolved_async_chunk is None: + resolved_async_chunk = bool(pipeline.async_chunk) + for stage in pipeline.stages: + stage.yaml_engine_args["async_chunk"] = bool(resolved_async_chunk) # Apply CLI overrides result: list[StageConfig] = [] @@ -1045,20 +1074,15 @@ class StageConfigFactory: model_type: str, cli_overrides: dict[str, Any], deploy_config_path: str | None = None, - cli_explicit_keys: set[str] | None = None, + **deprecated_kwargs: Any, ) -> list[StageConfig]: """Create StageConfigs from pipeline registry + deploy YAML. - Precedence (high → low): - explicit CLI args > deploy YAML > parser default CLI values - - ``cli_explicit_keys`` carries the set of long-option attribute names - the user actually typed (captured in ``OmniServeCommand.cmd``). Any - kwarg whose key is not in that set is treated as a parser default - and is only used to fill fields YAML doesn't already cover. When the - set is ``None`` (programmatic ``Omni()`` callers, which have no - argparse layer), every kwarg is treated as explicit. + Precedence: caller-typed (non-None) value > deploy YAML > + StageDeployConfig dataclass default. """ + _warn_deprecated_kwargs(deprecated_kwargs) + # Resolve deploy config path if deploy_config_path is None: deploy_path = _DEPLOY_DIR / f"{model_type}.yaml" @@ -1075,7 +1099,7 @@ class StageConfigFactory: deploy_cfg = load_deploy_config(deploy_path) cli_async_chunk = cli_overrides.get("async_chunk") - if cli_async_chunk is not None and (cli_explicit_keys is None or "async_chunk" in cli_explicit_keys): + if cli_async_chunk is not None: deploy_cfg.async_chunk = bool(cli_async_chunk) pipeline_key = deploy_cfg.pipeline or model_type @@ -1089,25 +1113,10 @@ class StageConfigFactory: stages = merge_pipeline_deploy(pipeline_cfg, deploy_cfg, cli_overrides) - # Precedence: explicit CLI > yaml > parser-default CLI. - # Per-stage (``stage_N_*``) keys are always treated as explicit. - explicit_overrides: dict[str, Any] = {} - default_overrides: dict[str, Any] = {} - for key, value in cli_overrides.items(): - if value is None: - continue - is_per_stage = bool(re.match(r"stage_\d+_", key)) - is_explicit = cli_explicit_keys is None or key in cli_explicit_keys or is_per_stage - if is_explicit: - explicit_overrides[key] = value - else: - default_overrides[key] = value + explicit_overrides = {k: v for k, v in cli_overrides.items() if v is not None} for stage in stages: - yaml_keys = set(stage.yaml_engine_args) - fallback = {k: v for k, v in default_overrides.items() if k not in yaml_keys} - merged = {**fallback, **explicit_overrides} - stage.runtime_overrides = cls._merge_cli_overrides(stage, merged) + stage.runtime_overrides = cls._merge_cli_overrides(stage, explicit_overrides) return stages @@ -1378,6 +1387,27 @@ class StageConfigFactory: except Exception as e: logger.debug(f"Failed to auto-detect model type for {model}: {e}") + # Fallback for diffusers-style models: check model_index.json. + # Some models (e.g. GLM-Image) have no root config.json but ship a + # model_index.json with _class_name that maps to a pipeline key via + # PipelineConfig.diffusers_class_name. + try: + from vllm.transformers_utils.config import get_hf_file_to_dict + + model_index = get_hf_file_to_dict("model_index.json", model, revision=None) + if model_index and "_class_name" in model_index: + class_name = model_index["_class_name"] + for pipeline_cfg in _PIPELINE_REGISTRY.values(): + if pipeline_cfg.diffusers_class_name == class_name: + logger.info( + "Detected pipeline %r from model_index.json (_class_name=%r)", + pipeline_cfg.model_type, + class_name, + ) + return pipeline_cfg.model_type, None + except Exception: + pass + # Final fallback: some models (e.g. CosyVoice3) ship an empty # config.json and rely on naming conventions. Match the model path # basename against registered pipeline keys — longest match wins diff --git a/vllm_omni/deploy/bagel.yaml b/vllm_omni/deploy/bagel.yaml new file mode 100644 index 000000000..9d2f1f8ff --- /dev/null +++ b/vllm_omni/deploy/bagel.yaml @@ -0,0 +1,48 @@ +# BAGEL-7B-MoT deploy: CUDA defaults, verified on NVIDIA A100 (80GB). +# +# Stage 0 (Thinker) and Stage 1 (DiT) share a single GPU by default. +# For dual-GPU setups, set stage 1 devices to "1". +# +# Fields omitted from a stage fall back to StageDeployConfig dataclass +# defaults (see vllm_omni/config/stage_config.py). + +async_chunk: false + +stages: + - stage_id: 0 + max_num_seqs: 3 + gpu_memory_utilization: 0.45 + devices: "0" + default_sampling_params: + temperature: 0.4 + top_p: 0.9 + top_k: 1 + max_tokens: 2048 + seed: 52 + detokenize: true + repetition_penalty: 1.05 + + - stage_id: 1 + max_num_seqs: 1 + enforce_eager: true + devices: "0" + input_connectors: + from_stage_0: shared_memory_connector + default_sampling_params: + seed: 52 + +connectors: + shared_memory_connector: + name: SharedMemoryConnector + extra: + shm_threshold_bytes: 65536 + + rdma_connector: + name: MooncakeTransferEngineConnector + extra: + host: "auto" + zmq_port: 50051 + protocol: "rdma" + device_name: "" + memory_pool_size: 4294967296 + memory_pool_device: "cpu" diff --git a/vllm_omni/deploy/bagel_single_stage.yaml b/vllm_omni/deploy/bagel_single_stage.yaml new file mode 100644 index 000000000..8470124ec --- /dev/null +++ b/vllm_omni/deploy/bagel_single_stage.yaml @@ -0,0 +1,22 @@ +# BAGEL-7B-MoT single-stage deploy: all modalities handled by the DiT stage. +# +# The DiT stage contains a full LLM (Qwen2-MoT), ViT, VAE, and tokenizer, +# so it supports text2img, img2img, img2text, text2text, and think mode +# without a separate Thinker (AR) stage. +# +# Select this topology via: +# vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni \ +# --deploy-config vllm_omni/deploy/bagel_single_stage.yaml +# +# Or programmatically: +# Omni(model="...", deploy_config_path="vllm_omni/deploy/bagel_single_stage.yaml") + +pipeline: bagel_single_stage +async_chunk: false + +stages: + - stage_id: 0 + max_num_seqs: 1 + devices: "0" + default_sampling_params: + seed: 52 diff --git a/vllm_omni/deploy/bagel_think.yaml b/vllm_omni/deploy/bagel_think.yaml new file mode 100644 index 000000000..d7adf7f66 --- /dev/null +++ b/vllm_omni/deploy/bagel_think.yaml @@ -0,0 +1,16 @@ +# BAGEL-7B-MoT think-mode deploy. +# +# Inherits all settings from bagel.yaml; only overrides the pipeline to +# bagel_think which uses expand_cfg_prompts_think and omits +# kv_transfer_criteria so the Thinker decodes tokens before +# transferring KV to DiT. +# +# Select this topology via: +# python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT --think +# +# Or explicitly: +# vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni \ +# --deploy-config vllm_omni/deploy/bagel_think.yaml + +base_config: bagel.yaml +pipeline: bagel_think diff --git a/vllm_omni/deploy/glm_image.yaml b/vllm_omni/deploy/glm_image.yaml new file mode 100644 index 000000000..47cd02aa8 --- /dev/null +++ b/vllm_omni/deploy/glm_image.yaml @@ -0,0 +1,38 @@ +# GLM-Image deploy: AR (stage 0) + Diffusion (stage 1). +# Topology declared in vllm_omni/model_executor/models/glm_image/pipeline.py. +# +# Fields omitted from a stage fall back to StageDeployConfig dataclass +# defaults (see vllm_omni/config/stage_config.py). Pipeline-wide settings +# (trust_remote_code, distributed_executor_backend, etc.) use DeployConfig +# defaults unless overridden here. +async_chunk: false + +stages: + # Stage 0: AR Model (GlmImageForConditionalGeneration) + # Generates prior_token_ids for conditioning the diffusion process. + - stage_id: 0 + max_num_seqs: 1 + gpu_memory_utilization: 0.6 + enforce_eager: false + max_num_batched_tokens: 32768 + devices: "0" + default_sampling_params: + temperature: 0.9 + top_p: 0.75 + top_k: 16512 + stop_token_ids: [16385] + seed: 42 + detokenize: false + + # Stage 1: Diffusion (DiT + VAE) + # Receives prior_token_ids from AR, performs denoising + VAE decode. + - stage_id: 1 + max_num_seqs: 1 + enforce_eager: true + devices: "1" + default_sampling_params: + seed: 42 + num_inference_steps: 50 + guidance_scale: 1.5 + height: 1024 + width: 1024 diff --git a/vllm_omni/deploy/qwen3_omni_moe.yaml b/vllm_omni/deploy/qwen3_omni_moe.yaml index fb8b61621..39baed6bd 100644 --- a/vllm_omni/deploy/qwen3_omni_moe.yaml +++ b/vllm_omni/deploy/qwen3_omni_moe.yaml @@ -46,7 +46,6 @@ stages: - stage_id: 2 gpu_memory_utilization: 0.1 - max_num_seqs: 1 enforce_eager: true async_scheduling: false max_num_batched_tokens: 51200 @@ -67,10 +66,11 @@ platforms: - stage_id: 0 gpu_memory_utilization: 0.6 tensor_parallel_size: 2 + max_num_batched_tokens: 8192 devices: "0,1" - stage_id: 1 gpu_memory_utilization: 0.6 - enforce_eager: true + max_num_batched_tokens: 8192 devices: "2" - stage_id: 2 gpu_memory_utilization: 0.3 diff --git a/vllm_omni/deploy/qwen3_tts.yaml b/vllm_omni/deploy/qwen3_tts.yaml index 5948007a9..522ea7c58 100644 --- a/vllm_omni/deploy/qwen3_tts.yaml +++ b/vllm_omni/deploy/qwen3_tts.yaml @@ -71,16 +71,3 @@ stages: max_tokens: 65536 seed: 42 repetition_penalty: 1.0 - -platforms: - npu: - stages: - # NPU does not yet support async-scheduling for TTS, and the - # talker fits at max_num_seqs=1 only. - - stage_id: 0 - max_num_seqs: 1 - enforce_eager: true - async_scheduling: false - - stage_id: 1 - gpu_memory_utilization: 0.2 - async_scheduling: false diff --git a/vllm_omni/deploy/voxtral_tts.yaml b/vllm_omni/deploy/voxtral_tts.yaml index b0899a399..87d999c67 100644 --- a/vllm_omni/deploy/voxtral_tts.yaml +++ b/vllm_omni/deploy/voxtral_tts.yaml @@ -36,6 +36,8 @@ stages: max_tokens: 2048 seed: 42 repetition_penalty: 1.1 + extra_args: + cfg_alpha: 1.2 tokenizer_mode: mistral config_format: mistral load_format: mistral diff --git a/vllm_omni/diffusion/attention/backends/flash_attn.py b/vllm_omni/diffusion/attention/backends/flash_attn.py index b6ab3a57a..d38ea4f6e 100644 --- a/vllm_omni/diffusion/attention/backends/flash_attn.py +++ b/vllm_omni/diffusion/attention/backends/flash_attn.py @@ -96,7 +96,7 @@ class FlashAttentionImpl(AttentionImpl): value: torch.Tensor, attn_metadata: AttentionMetadata = None, ) -> torch.Tensor: - """CUDA/ROCm flash attention implementation.""" + """CUDA/ROCm/MUSA flash attention implementation.""" from vllm_omni.diffusion.attention.backends.utils.fa import ( HAS_FLASH_ATTN, flash_attn_func, @@ -209,13 +209,3 @@ class FlashAttentionImpl(AttentionImpl): layout="BNSD", ) return output - - def forward_musa( - self, - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - attn_metadata: AttentionMetadata = None, - ) -> torch.Tensor: - # XXX (MUSA): MUSA uses the same implementation as XPU (mate only provides flash_attn_varlen_func) - return self.forward_xpu(query, key, value, attn_metadata) diff --git a/vllm_omni/diffusion/attention/backends/utils/fa.py b/vllm_omni/diffusion/attention/backends/utils/fa.py index fe6051f8b..18886871a 100644 --- a/vllm_omni/diffusion/attention/backends/utils/fa.py +++ b/vllm_omni/diffusion/attention/backends/utils/fa.py @@ -41,7 +41,7 @@ elif current_omni_platform.is_xpu(): pass elif current_omni_platform.is_musa(): try: - from mate import flash_attn_varlen_func # noqa: F401 + from flash_attn_interface import flash_attn_func, flash_attn_varlen_func # noqa: F401 except (ImportError, ModuleNotFoundError): pass else: @@ -82,8 +82,8 @@ HAS_FLASH_ATTN = flash_attn_func is not None or flash_attn_varlen_func is not No @lru_cache(maxsize=1) def is_mate_available() -> bool: - """Check if MATE (MUSA Flash Attention) is available.""" - return current_omni_platform.is_musa() and flash_attn_varlen_func is not None + """Check if MATE (MUSA AI Tensor Engine) is available.""" + return current_omni_platform.is_musa() and flash_attn_func is not None or flash_attn_varlen_func is not None def _index_first_axis(tensor, indices): diff --git a/vllm_omni/diffusion/cache/cache_dit_backend.py b/vllm_omni/diffusion/cache/cache_dit_backend.py index d5397dd16..1b1800dac 100644 --- a/vllm_omni/diffusion/cache/cache_dit_backend.py +++ b/vllm_omni/diffusion/cache/cache_dit_backend.py @@ -1337,6 +1337,8 @@ CUSTOM_DIT_ENABLERS.update( "LTX2ImageToVideoPipeline": enable_cache_for_ltx2, "LTX2TwoStagesPipeline": enable_cache_for_ltx2, "LTX2ImageToVideoTwoStagesPipeline": enable_cache_for_ltx2, + "LTX23Pipeline": enable_cache_for_ltx2, + "LTX23ImageToVideoPipeline": enable_cache_for_ltx2, "BagelPipeline": enable_cache_for_bagel, "GlmImagePipeline": enable_cache_for_glm_image, "Flux2Pipeline": enable_cache_for_flux2, diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 012c0130c..40c943a41 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -452,7 +452,14 @@ class OmniDiffusionConfig: custom_pipeline_args: dict[str, Any] | None = None # Diffusion model loading format - diffusion_load_format: str = "default" # "default", "custom_pipeline", "dummy" + # "default", "custom_pipeline", "dummy", "diffusers" (HF diffusers adapter) + diffusion_load_format: str = "default" + + # Diffusers adapter kwargs + # kwargs forwarded to DiffusionPipeline.from_pretrained() + diffusers_load_kwargs: dict[str, Any] = field(default_factory=dict) + # kwargs forwarded to pipeline.__call__() + diffusers_call_kwargs: dict[str, Any] = field(default_factory=dict) # http server endpoint config, would be ignored in local mode host: str | None = None @@ -648,6 +655,12 @@ class OmniDiffusionConfig: elif self.max_cpu_loras < 1: raise ValueError("max_cpu_loras must be >= 1 for diffusion LoRA") + if self.diffusion_load_format != "diffusers" and (self.diffusers_load_kwargs or self.diffusers_call_kwargs): + raise ValueError( + "diffusers_load_kwargs and diffusers_call_kwargs are only " + "valid together with diffusion_load_format=diffusers" + ) + def set_tf_model_config(self, tf_config: "TransformerConfig") -> None: """Assign `tf_model_config` and propagate quantization if detected. @@ -686,6 +699,10 @@ class OmniDiffusionConfig: """ from vllm.transformers_utils.config import get_hf_file_to_dict + # Default model_class_name for diffusers adapter + if self.model_class_name is None and self.diffusion_load_format == "diffusers": + self.model_class_name = "DiffusersAdapterPipeline" + try: config_dict = get_hf_file_to_dict("model_index.json", self.model) if config_dict is not None: @@ -693,32 +710,42 @@ class OmniDiffusionConfig: self.model_class_name = config_dict.get("_class_name", None) self.update_multimodal_support() - tf_config_dict = get_hf_file_to_dict("transformer/config.json", self.model) - self.tf_model_config = TransformerConfig.from_dict(tf_config_dict) + # Skip transformer config loading for diffusers adapter + # (non-DiT models don't have a separate transformer folder/config) + if self.diffusion_load_format == "diffusers": + self.tf_model_config = TransformerConfig() + else: + tf_config_dict = get_hf_file_to_dict("transformer/config.json", self.model) + self.tf_model_config = TransformerConfig.from_dict(tf_config_dict) else: raise FileNotFoundError("model_index.json not found") except (AttributeError, OSError, ValueError, FileNotFoundError): - cfg = get_hf_file_to_dict("config.json", self.model) - if cfg is None: - raise ValueError(f"Could not find config.json or model_index.json for model {self.model}") - - self.tf_model_config = TransformerConfig.from_dict(cfg) - model_type = cfg.get("model_type") - architectures = cfg.get("architectures") or [] - - if model_type == "bagel" or "BagelForConditionalGeneration" in architectures: - self.model_class_name = "BagelPipeline" + # Skip transformer config loading for diffusers adapter + # (non-DiT models don't have a separate transformer folder/config) + if self.diffusion_load_format == "diffusers": self.tf_model_config = TransformerConfig() - self.update_multimodal_support() - elif model_type == "nextstep": - if self.model_class_name is None: - self.model_class_name = "NextStep11Pipeline" - self.tf_model_config = TransformerConfig() - self.update_multimodal_support() - elif architectures and len(architectures) == 1: - self.model_class_name = architectures[0] else: - raise + cfg = get_hf_file_to_dict("config.json", self.model) + if cfg is None: + raise ValueError(f"Could not find config.json or model_index.json for model {self.model}") + + self.tf_model_config = TransformerConfig.from_dict(cfg) + model_type = cfg.get("model_type") + architectures = cfg.get("architectures") or [] + + if model_type == "bagel" or "BagelForConditionalGeneration" in architectures: + self.model_class_name = "BagelPipeline" + self.tf_model_config = TransformerConfig() + self.update_multimodal_support() + elif model_type == "nextstep": + if self.model_class_name is None: + self.model_class_name = "NextStep11Pipeline" + self.tf_model_config = TransformerConfig() + self.update_multimodal_support() + elif architectures and len(architectures) == 1: + self.model_class_name = architectures[0] + else: + raise @classmethod def from_kwargs(cls, **kwargs: Any) -> "OmniDiffusionConfig": @@ -732,7 +759,7 @@ class OmniDiffusionConfig: # Backwards-compatibility: map "quantization" to "quantization_config" # so callers using the old field name still work. - if "quantization" in kwargs and "quantization_config" not in kwargs: + if "quantization" in kwargs and kwargs.get("quantization_config") is None: kwargs["quantization_config"] = kwargs.pop("quantization") else: kwargs.pop("quantization", None) @@ -743,6 +770,12 @@ class OmniDiffusionConfig: cache_backend = os.environ.get("DIFFUSION_CACHE_BACKEND") or os.environ.get("DIFFUSION_CACHE_ADAPTER") kwargs["cache_backend"] = cache_backend.lower() if cache_backend else "none" + # Falsy-value check for not-None fields (convert potential None values in YAML config to empty containers) + if "diffusers_load_kwargs" in kwargs and kwargs["diffusers_load_kwargs"] is None: + kwargs["diffusers_load_kwargs"] = {} + if "diffusers_call_kwargs" in kwargs and kwargs["diffusers_call_kwargs"] is None: + kwargs["diffusers_call_kwargs"] = {} + # Filter kwargs to only include valid fields valid_fields = {f.name for f in fields(cls)} filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_fields} diff --git a/vllm_omni/diffusion/layers/norm.py b/vllm_omni/diffusion/layers/norm.py index 6096ad7c3..f397c1a85 100644 --- a/vllm_omni/diffusion/layers/norm.py +++ b/vllm_omni/diffusion/layers/norm.py @@ -108,3 +108,81 @@ class RMSNorm(CustomOp): out = x * torch.rsqrt(variance + self.variance_epsilon) out = self.weight.to(torch.float32) * out return out.to(input_dtype) + + +class RMSNormVAE(CustomOp): + """Root Mean Square Layer Normalization for Channel-First or Last""" + + def __init__( + self, + dim: int, + channel_first: bool = True, + images: bool = True, + bias: bool = False, + epsilon: float = 1e-6, + ) -> None: + super().__init__() + broadcastable_dims = (1, 1, 1) if not images else (1, 1) + shape = (dim, *broadcastable_dims) if channel_first else (dim,) + + self.channel_first = channel_first + self.scale = dim**0.5 + self.gamma = nn.Parameter(torch.ones(shape)) + self.bias = nn.Parameter(torch.zeros(shape)) if bias else None + self.epsilon = epsilon + + self.gamma_rmsnorm = None + + def forward_cuda( + self, + x: torch.Tensor, + ) -> torch.Tensor: + return self.forward_native(x) + + def forward_hip( + self, + x: torch.Tensor, + ) -> torch.Tensor: + return self.forward_native(x) + + def forward_npu( + self, + x: torch.Tensor, + ) -> torch.Tensor: + import torch_npu + + if self.gamma_rmsnorm is None: + self.gamma_rmsnorm = self.gamma.reshape(-1) + + if self.channel_first: + x = x.transpose(1, -1) + out = torch_npu.npu_rms_norm(x, self.gamma_rmsnorm, epsilon=self.epsilon)[0].transpose(1, -1) + else: + out = torch_npu.npu_rms_norm(x, self.gamma_rmsnorm, epsilon=self.epsilon)[0] + + if self.bias is not None: + out = out + self.bias + return out + + def forward_xpu( + self, + x: torch.Tensor, + ) -> torch.Tensor: + return self.forward_native(x) + + def forward_native( + self, + x: torch.Tensor, + ) -> torch.Tensor: + out = ( + F.normalize( + x, + dim=(1 if self.channel_first else -1), + eps=self.epsilon, + ) + * self.scale + * self.gamma + ) + if self.bias is not None: + out = out + self.bias + return out diff --git a/vllm_omni/diffusion/layers/rope.py b/vllm_omni/diffusion/layers/rope.py index 61ddb4d84..127e1b2cb 100644 --- a/vllm_omni/diffusion/layers/rope.py +++ b/vllm_omni/diffusion/layers/rope.py @@ -72,18 +72,18 @@ class RotaryEmbedding(CustomOp): of 1st half and 2nd half (GPT-NeoX style). """ - def __init__( - self, - is_neox_style: bool = False, - ) -> None: + def __init__(self, is_neox_style: bool = False) -> None: super().__init__() self.is_neox_style = is_neox_style self.interleaved = not is_neox_style self.apply_rotary_emb_flash_attn = None + self.has_mindie = False if find_spec("flash_attn") is not None: from flash_attn.ops.triton.rotary import apply_rotary self.apply_rotary_emb_flash_attn = apply_rotary + if find_spec("mindiesd") is not None: + self.has_mindie = True def forward_cuda( self, @@ -132,7 +132,7 @@ class RotaryEmbedding(CustomOp): cos: torch.Tensor, sin: torch.Tensor, ) -> torch.Tensor: - if find_spec("mindiesd"): + if self.has_mindie: return apply_rotary_emb_mindiesd(x, cos, sin, self.interleaved) else: return self.forward_native(x, cos, sin) @@ -167,6 +167,56 @@ class RotaryEmbedding(CustomOp): ) +class RotaryEmbeddingWan(RotaryEmbedding): + """ + rotary positional embedding for Wan. + interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead + of 1st half and 2nd half (GPT-NeoX style). + """ + + def __init__(self, is_neox_style: bool = False, half_head_dim: bool = False) -> None: + super().__init__(is_neox_style=is_neox_style) + self.half_head_dim = half_head_dim + + def forward_cuda( + self, + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + ) -> torch.Tensor: + return self.forward_native(x, cos, sin) + + def forward_npu( + self, + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + ) -> torch.Tensor: + if self.has_mindie: + if cos.dim() > 2: + cos = cos.reshape(-1, cos.shape[-1]) + sin = sin.reshape(-1, sin.shape[-1]) + return apply_rotary_emb_mindiesd(x, cos, sin, self.interleaved, self.half_head_dim) + else: + return self.forward_native(x, cos, sin) + + def forward_native( + self, + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + ) -> torch.Tensor: + x1, x2 = x.unflatten(-1, (-1, 2)).unbind(-1) + rotated = torch.stack( + ( + x1 * cos - x2 * sin, + x1 * sin + x2 * cos, + ), + dim=-1, + ) + return rotated.flatten(-2, -1).to(x.dtype) + + def apply_rope_to_qk( rope: RotaryEmbedding, query: torch.Tensor, diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index 146afb26f..91f3574b1 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -32,6 +32,7 @@ from vllm.utils.torch_utils import set_default_torch_dtype from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.distributed.hsdp import HSDPInferenceConfig from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter +from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import DiffusersAdapterPipeline from vllm_omni.diffusion.registry import initialize_model if TYPE_CHECKING: @@ -257,11 +258,14 @@ class DiffusersPipelineLoader: self, od_config: OmniDiffusionConfig, load_device: str, - load_format: str = "default", + load_format: str | None = "default", custom_pipeline_name: str | None = None, device: torch.device | None = None, ) -> nn.Module: """Load a model with the given configurations.""" + if load_format is None: + load_format = "default" + # CPU offload + FP8: load weights on device for FP8 quantization if load_device == "cpu" and od_config.quantization_config is not None: load_device = device.type @@ -277,11 +281,21 @@ class DiffusersPipelineLoader: with target_device: if load_format == "default": model = initialize_model(od_config) + elif load_format == "diffusers": + model = DiffusersAdapterPipeline(od_config=od_config, device=target_device) elif load_format == "custom_pipeline": model_cls = resolve_obj_by_qualname(custom_pipeline_name) model = model_cls(od_config=od_config) + else: + # 'dummy' format should not call this function at all + raise ValueError(f"Unknown load_format: {load_format}") logger.debug("Loading weights on %s ...", load_device) - if self._is_gguf_quantization(od_config): + if load_format == "diffusers": + # DiffusersAdapterPipeline.load_weights() calls + # DiffusionPipeline.from_pretrained() internally — it does + # NOT use our native (customized) pipeline classes. + cast(DiffusersAdapterPipeline, model).load_weights() + elif self._is_gguf_quantization(od_config): self._load_weights_with_gguf(model, od_config) else: # Quantization does not happen in `load_weights` but after it diff --git a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py index 90baf5f67..0290e9502 100644 --- a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py +++ b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py @@ -12,6 +12,7 @@ from collections.abc import Iterable from copy import deepcopy from dataclasses import dataclass from math import isqrt +from typing import ClassVar import numpy as np import torch @@ -26,6 +27,7 @@ from vllm.transformers_utils.configs.bagel import BagelConfig from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader +from vllm_omni.diffusion.models.interface import SupportsModuleOffload from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific @@ -148,12 +150,25 @@ class SiglipNaViTWrapper(nn.Module): return outputs.last_hidden_state.squeeze(0) -class BagelPipeline(nn.Module, DiffusionPipelineProfilerMixin): +class BagelPipeline(nn.Module, SupportsModuleOffload, DiffusionPipelineProfilerMixin): """Bagel generation pipeline (MoT) packaged for vllm-omni diffusion engine. This pipeline is self-contained and uses the ported Bagel core files. """ + _dit_modules: ClassVar[list[str]] = ["language_model.model"] + _encoder_modules: ClassVar[list[str]] = [] + _vae_modules: ClassVar[list[str]] = ["vae"] + _resident_modules: ClassVar[list[str]] = [ + "bagel.time_embedder", + "bagel.vae2llm", + "bagel.llm2vae", + "bagel.latent_pos_embed", + "bagel.vit_model", + "bagel.connector", + "bagel.vit_pos_embed", + ] + def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): super().__init__() self.od_config = od_config diff --git a/vllm_omni/diffusion/models/diffusers_adapter/__init__.py b/vllm_omni/diffusion/models/diffusers_adapter/__init__.py new file mode 100644 index 000000000..c8dd51c8e --- /dev/null +++ b/vllm_omni/diffusion/models/diffusers_adapter/__init__.py @@ -0,0 +1,12 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Diffusers backend adapter for vLLM-Omni.""" + +from vllm_omni.diffusion.models.diffusers_adapter.pipeline_diffusers_adapter import ( + DiffusersAdapterPipeline, +) + +__all__ = [ + "DiffusersAdapterPipeline", +] diff --git a/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py b/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py new file mode 100644 index 000000000..8a1fdfc08 --- /dev/null +++ b/vllm_omni/diffusion/models/diffusers_adapter/pipeline_diffusers_adapter.py @@ -0,0 +1,358 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Diffusers backend adapter for vLLM-Omni. + +Provides a black-box wrapper around any 🤗 Diffusers pipeline, enabling +vLLM-Omni to directly serve Diffusers models with near-zero per-model code. + +The adapter delegates full pipeline execution to diffusers' ``__call__()``. +It does NOT support: +- CFG parallel (diffusers handles CFG via guidance_scale internally) +- Sequence parallel (requires model-specific attention surgery) +- TeaCache / Cache-DiT (requires hooking into transformer blocks) +- Step-wise execution (continuous batching) +""" + +import logging +import os +from typing import Any + +import torch +from diffusers.pipelines.pipeline_utils import DiffusionPipeline +from torch import nn + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.inputs.data import OmniPromptType +from vllm_omni.platforms import current_omni_platform + +logger = logging.getLogger(__name__) + + +class DiffusersAdapterPipeline(nn.Module, DiffusionPipelineProfilerMixin): + """Black-box adapter that delegates full pipeline execution to a diffusers pipeline. + + Usage:: + + adapter = DiffusersAdapterPipeline(od_config=od_config) + adapter.load_weights() # calls DiffusionPipeline.from_pretrained() + output = adapter.forward(req) + + Step-wise execution is explicitly rejected — diffusers encapsulates the + full denoising loop internally. Use native pipelines for continuous + batching mode. + """ + + supports_step_execution: bool = False + + def __init__(self, *, od_config: OmniDiffusionConfig, device: torch.device | None = None): + super().__init__() + self._pipeline: DiffusionPipeline + self.od_config = od_config + self.device = device + self._capabilities: dict[str, Any] = {} + self._raise_unsupported_features() + + self.setup_diffusion_pipeline_profiler( + enable_diffusion_pipeline_profiler=od_config.enable_diffusion_pipeline_profiler, + profiler_targets=["forward"], + ) + if od_config.enable_diffusion_pipeline_profiler: + logger.info("Profiling enabled for DiffusersAdapterPipeline. Only 'forward' is supported.") + + # ------------------------------------------------------------------ + # Weight loading + # ------------------------------------------------------------------ + + def load_weights(self) -> None: + """Load the diffusers pipeline via ``DiffusionPipeline.from_pretrained()``.""" + + model_id = self.od_config.model + dtype = self.od_config.dtype + + load_kwargs = { + "torch_dtype": dtype, + **self.od_config.diffusers_load_kwargs, + } + logger.debug(f"Loading diffusers pipeline with kwargs: {load_kwargs}") + + self._pipeline = DiffusionPipeline.from_pretrained( + model_id, + **load_kwargs, + ).to(self.device) + + # CPU offloading + if self.od_config.enable_layerwise_offload: + self._pipeline.enable_sequential_cpu_offload() + elif self.od_config.enable_cpu_offload: + self._pipeline.enable_model_cpu_offload() + + # VAE slicing and tiling: try-catch because not all models have VAE + if self.od_config.vae_use_slicing: + try: + self._pipeline.enable_vae_slicing() + except Exception as e: + logger.warning( + f"Failed to enable VAE slicing for diffusers pipeline {self._pipeline.__class__.__name__}: {e}" + ) + if self.od_config.vae_use_tiling: + try: + self._pipeline.enable_vae_tiling() + except Exception as e: + logger.warning( + f"Failed to enable VAE tiling for diffusers pipeline {self._pipeline.__class__.__name__}: {e}" + ) + + # Attention backend + self._set_attention_backend() + + # ------------------------------------------------------------------ + # Step-wise execution — explicitly rejected + # ------------------------------------------------------------------ + + def prepare_encode(self, **_: Any) -> Any: + raise NotImplementedError( + "Step-wise execution is not yet supported with the diffusers backend. " + "Use a native pipeline for continuous batching mode." + ) + + def denoise_step(self, **_: Any) -> torch.Tensor | None: + raise NotImplementedError( + "Step-wise execution is not yet supported with the diffusers backend. " + "Use a native pipeline for continuous batching mode." + ) + + def step_scheduler(self, **_: Any) -> None: + raise NotImplementedError( + "Step-wise execution is not yet supported with the diffusers backend. " + "Use a native pipeline for continuous batching mode." + ) + + def post_decode(self, **_: Any) -> Any: + raise NotImplementedError( + "Step-wise execution is not yet supported with the diffusers backend. " + "Use a native pipeline for continuous batching mode." + ) + + # ------------------------------------------------------------------ + # Forward pass + # ------------------------------------------------------------------ + + def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput: + """Full delegation to diffusers ``pipeline.__call__()``.""" + + kwargs = self._build_call_kwargs(req) + logger.debug(f"Calling diffusers pipeline with kwargs: {kwargs}") + + with torch.inference_mode(): + output = self._pipeline(**kwargs) # pyright: ignore[reportCallIssue] + + return self._wrap_output(output) + + # ------------------------------------------------------------------ + # Validation guards + # ------------------------------------------------------------------ + + def _raise_unsupported_features(self) -> None: + """Raise an error for incompatible feature switches.""" + pc = self.od_config.parallel_config + if pc.cfg_parallel_size > 1: + raise NotImplementedError( + "CFG parallel is not supported with the diffusers backend. " + "Diffusers handles CFG internally via guidance_scale." + ) + if pc.sequence_parallel_size is not None and pc.sequence_parallel_size > 1: + raise NotImplementedError( + "Sequence parallel is not supported with the diffusers backend. " + "It requires model-specific attention surgery." + ) + if self.od_config.cache_backend not in ("none", None): + raise NotImplementedError( + f"Cache backend '{self.od_config.cache_backend}' is not supported " + "with the diffusers backend. TeaCache/Cache-DiT require hooking " + "into individual transformer blocks." + ) + if self.od_config.enforce_eager: + raise NotImplementedError( + "Eager execution is not supported with the diffusers backend. " + "Use a native pipeline for continuous batching mode." + ) + if self.od_config.quantization_config is not None: + raise NotImplementedError( + "Quantization is not supported with the diffusers backend. Use a native pipeline for quantization." + ) + + # ------------------------------------------------------------------ + # Wrap settings, inputs, and outputs + # ------------------------------------------------------------------ + + def _set_attention_backend(self) -> None: + """Set the attention backend. + + Roughly follow the logic in vllm_omni/diffusion/attention/backends/utils/fa.py, + But also consider the available attention backends in diffusers. + (See: https://huggingface.co/docs/diffusers/optimization/attention_backends) + """ + if not hasattr(self._pipeline, "transformer"): + logging.info("No transformer found in diffusers pipeline. Skipping attention backend setting.") + return + + attention_backend_config = self.od_config.attention_backend or os.environ.get("DIFFUSION_ATTENTION_BACKEND") + attention_backend_attempts: list[str] = [] + match attention_backend_config: + case "FLASH_ATTN" | None: + if current_omni_platform.is_rocm(): + attention_backend_attempts.append("aiter") + elif current_omni_platform.is_xpu(): + attention_backend_attempts.append("_native_xla") + elif current_omni_platform.is_musa(): + logger.warning( + "Unknown diffusers attention backend option for MUSA platform. Falling back to SDPA." + ) + attention_backend_attempts.append("native") + else: + attention_backend_attempts.extend( + [ + "_flash_3_hub", + "_flash_3_varlen_hub", + "_flash_3", + "_flash_varlen_3", + "flash_hub", + "flash_varlen_hub", + "flash", + "flash_varlen", + "_native_flash", + ] + ) + case "SAGE_ATTN": + attention_backend_attempts.extend(["sage_hub", "sage", "sage", "sage_varlen"]) + case "ASCEND": + attention_backend_attempts.append("_native_npu") + case "TORCH_SDPA": + attention_backend_attempts.append("native") + case _: + logger.warning(f"Invalid attention backend: {attention_backend_config}. Falling back to SDPA.") + attention_backend_attempts.append("native") + + attempt_errors: list[str] = [] + set_backend: str | None = None + for backend in attention_backend_attempts: + try: + self._pipeline.transformer.set_attention_backend(backend) + set_backend = backend + break + except Exception as e: + attempt_errors.append(str(e)) + + # If all attempts fail, fallback to SDPA and warn the user about the failures + if len(attempt_errors) == len(attention_backend_attempts): + self._pipeline.transformer.set_attention_backend("native") + logger.warning( + f"Failed to set attention backend '{attention_backend_config}' for " + f"diffusers pipeline {self._pipeline.__class__.__name__}. " + "Falling back to SDPA. " + f"The following attempts were made: {dict(zip(attention_backend_attempts, attempt_errors))}" + ) + return + + # If some attempts fail, only warn the user about the failures + logger.info( + f"Set diffusers attention backend to '{set_backend}', adapted from " + f"user config value '{attention_backend_config}'." + ) + if len(attempt_errors) > 0: + logger.warning( + f"The following failed attempts were made before choosing this diffusers backend: " + f"{dict(zip(attention_backend_attempts, attempt_errors))}" + ) + + def _build_call_kwargs(self, req: OmniDiffusionRequest) -> dict[str, Any]: + """Translate ``OmniDiffusionRequest`` into diffusers ``__call__`` kwargs.""" + sampling = req.sampling_params + prompt, neg_prompt = self._extract_prompt(req.prompts) + + # Merge user-provided call kwargs from stage/CLI defaults. + # Request-time parameters take precedence over stage-config defaults + call_kwargs = self.od_config.diffusers_call_kwargs + kwargs: dict[str, Any] = { + **call_kwargs, + "prompt": prompt, + "num_inference_steps": sampling.num_inference_steps, + "guidance_scale": sampling.guidance_scale, + "output_type": sampling.output_type or self.od_config.output_type, + } + + if sampling.height is not None: + kwargs["height"] = sampling.height + if sampling.width is not None: + kwargs["width"] = sampling.width + if sampling.num_frames is not None and sampling.num_frames > 1: + kwargs["num_frames"] = sampling.num_frames + if sampling.num_outputs_per_prompt is not None and sampling.num_outputs_per_prompt > 1: + kwargs["num_images_per_prompt"] = sampling.num_outputs_per_prompt + + if neg_prompt is not None: + kwargs["negative_prompt"] = neg_prompt + + if sampling.generator is not None: + kwargs["generator"] = sampling.generator + elif sampling.seed is not None: + kwargs["generator"] = torch.Generator(device=sampling.generator_device).manual_seed(sampling.seed) + else: + kwargs["generator"] = torch.Generator(device=sampling.generator_device) + + if sampling.latents is not None: + kwargs["latents"] = sampling.latents + + return kwargs + + @staticmethod + def _extract_prompt(prompt_obj: list[OmniPromptType]) -> tuple[str | list[str], str | list[str] | None]: + """Extract the text prompts and negative prompts from a list of prompt objects.""" + if len(prompt_obj) == 1: + if isinstance(prompt_obj[0], str): + return prompt_obj[0], None + else: + return prompt_obj[0].get("prompt", ""), prompt_obj[0].get("negative_prompt", None) + + prompts = [] + negative_prompts: list[str] | None = [] + for prompt in prompt_obj: + if isinstance(prompt, str): + prompts.append(prompt) + else: + prompts.append(prompt.get("prompt", "")) + negative_prompts.append(prompt.get("negative_prompt", "")) + if all(not np for np in negative_prompts): + negative_prompts = None + return prompts, negative_prompts + + @staticmethod + def _extract_negative_prompt(prompt_obj: Any) -> str | None: + """Extract the negative prompt from a prompt object, if present.""" + if isinstance(prompt_obj, dict): + return prompt_obj.get("negative_prompt") + return getattr(prompt_obj, "negative_prompt", None) + + def _wrap_output(self, output: Any) -> DiffusionOutput: + """Convert diffusers pipeline output to ``DiffusionOutput``. + + Diffusers output types: + - ``ImagePipelineOutput(images=...)`` — text2img, img2img + - ``VideoPipelineOutput(frames=...)`` — text2vid, img2vid + """ + from vllm_omni.diffusion.data import DiffusionOutput + + if hasattr(output, "images"): + # Preserve diffusers image format (`output_type`) + return DiffusionOutput(output=output.images) + + if hasattr(output, "frames"): + # Preserve diffusers video format (`output_type`) + return DiffusionOutput(output=output.frames) + + if hasattr(output, "audios"): + return DiffusionOutput(output=output.audios) + + return DiffusionOutput(output=output) diff --git a/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py b/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py index 021d28c2a..2da0038bb 100644 --- a/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py +++ b/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py @@ -829,14 +829,21 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi """Taking only the first image for now.""", ) first_prompt = req.prompts[0] - prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "") + if isinstance(first_prompt, str): + multi_modal_data = {} + prompt = first_prompt + raw_image = None + mask_image = None + reference_image = None + else: + multi_modal_data = first_prompt.get("multi_modal_data", {}) + prompt = first_prompt.get("prompt") or "" + raw_image = multi_modal_data.get("image") + mask_image = multi_modal_data.get("mask_image") + reference_image = multi_modal_data.get("reference_image") - if ( - raw_image := None - if isinstance(first_prompt, str) - else first_prompt.get("multi_modal_data", {}).get("image") - ) is None: - pass # use image from param list + if raw_image is None: + image = None elif isinstance(raw_image, list): image = [PIL.Image.open(im) if isinstance(im, str) else cast(PIL.Image.Image, im) for im in raw_image] else: @@ -956,11 +963,6 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi height = height or self.default_sample_size * self.vae_scale_factor width = width or self.default_sample_size * self.vae_scale_factor - # Get mask_image and reference_image - multi_modal_data = req.prompts[0].get("multi_modal_data", {}) if req.prompts else {} - mask_image = multi_modal_data.get("mask_image") - reference_image = multi_modal_data.get("reference_image") - if mask_image is not None and (image is None or (isinstance(image, list) and len(image) == 0)): raise ValueError("image must be provided when using mask_image for inpainting") diff --git a/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py b/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py index 7ff42a5f0..ddb32aa20 100644 --- a/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py +++ b/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py @@ -4,7 +4,7 @@ import math from collections.abc import Iterable from enum import Enum -from typing import Any +from typing import TYPE_CHECKING, Any import torch import torch.nn as nn @@ -19,6 +19,9 @@ from vllm.model_executor.layers.linear import ( ) from vllm.model_executor.model_loader.weight_utils import default_weight_loader +if TYPE_CHECKING: + from vllm.model_executor.layers.quantization.base_config import QuantizationConfig + from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.cache.base import CachedTransformer @@ -465,6 +468,7 @@ class GlmImageAttention(nn.Module): parallel_config: DiffusionParallelConfig | None = None, out_bias: bool = True, eps: float = 1e-5, + quant_config: "QuantizationConfig | None" = None, ): super().__init__() self.dim = dim @@ -480,6 +484,7 @@ class GlmImageAttention(nn.Module): total_num_kv_heads=num_heads, bias=True, return_bias=False, + quant_config=quant_config, ) # QK normalization (LayerNorm, not RMSNorm for GLM-Image) @@ -495,6 +500,7 @@ class GlmImageAttention(nn.Module): bias=out_bias, input_is_parallel=True, return_bias=False, + quant_config=quant_config, ), nn.Dropout(0.0), ] @@ -656,6 +662,7 @@ class ColumnParallelGELU(nn.Module): *, approximate: str = "none", bias: bool = True, + quant_config: "QuantizationConfig | None" = None, ): super().__init__() self.proj = ColumnParallelLinear( @@ -664,6 +671,7 @@ class ColumnParallelGELU(nn.Module): bias=bias, gather_output=False, return_bias=False, + quant_config=quant_config, ) self.approximate = approximate @@ -679,6 +687,7 @@ class ColumnParallelSiLU(nn.Module): dim_out: int, *, bias: bool = True, + quant_config: "QuantizationConfig | None" = None, ): super().__init__() self.proj = ColumnParallelLinear( @@ -687,6 +696,7 @@ class ColumnParallelSiLU(nn.Module): bias=bias, gather_output=False, return_bias=False, + quant_config=quant_config, ) def forward(self, x: torch.Tensor) -> torch.Tensor: @@ -703,6 +713,7 @@ class GlmImageFeedForward(nn.Module): inner_dim: int | None = None, bias: bool = True, activation_fn: str = "gelu", + quant_config: "QuantizationConfig | None" = None, ): super().__init__() inner_dim = inner_dim or int(dim * mult) @@ -710,7 +721,7 @@ class GlmImageFeedForward(nn.Module): if activation_fn == "linear-silu": layers: list[nn.Module] = [ - ColumnParallelSiLU(dim, inner_dim, bias=bias), + ColumnParallelSiLU(dim, inner_dim, bias=bias, quant_config=quant_config), nn.Identity(), RowParallelLinear( inner_dim, @@ -718,12 +729,13 @@ class GlmImageFeedForward(nn.Module): bias=bias, input_is_parallel=True, return_bias=False, + quant_config=quant_config, ), ] else: approximate = "tanh" if activation_fn == "gelu-approximate" else "none" layers = [ - ColumnParallelGELU(dim, inner_dim, approximate=approximate, bias=bias), + ColumnParallelGELU(dim, inner_dim, approximate=approximate, bias=bias, quant_config=quant_config), nn.Identity(), RowParallelLinear( inner_dim, @@ -731,6 +743,7 @@ class GlmImageFeedForward(nn.Module): bias=bias, input_is_parallel=True, return_bias=False, + quant_config=quant_config, ), ] @@ -752,6 +765,7 @@ class GlmImageTransformerBlock(nn.Module): attention_head_dim: int = 40, time_embed_dim: int = 512, ffn_hidden_dim: int | None = None, + quant_config: "QuantizationConfig | None" = None, parallel_config: DiffusionParallelConfig | None = None, ) -> None: super().__init__() @@ -762,13 +776,20 @@ class GlmImageTransformerBlock(nn.Module): dim=dim, num_heads=num_attention_heads, head_dim=attention_head_dim, + quant_config=quant_config, parallel_config=parallel_config, ) # 2. Feedforward self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5) self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5) - self.ff = GlmImageFeedForward(dim=dim, dim_out=dim, inner_dim=ffn_hidden_dim, activation_fn="gelu-approximate") + self.ff = GlmImageFeedForward( + dim=dim, + dim_out=dim, + inner_dim=ffn_hidden_dim, + activation_fn="gelu-approximate", + quant_config=quant_config, + ) def forward( self, @@ -879,6 +900,7 @@ class GlmImageTransformer2DModel(CachedTransformer): def __init__( self, od_config: OmniDiffusionConfig, + quant_config: "QuantizationConfig | None" = None, ): super().__init__() @@ -932,11 +954,19 @@ class GlmImageTransformer2DModel(CachedTransformer): # 2. Patch & Text-timestep embedding self.image_projector = GlmImageImageProjector(in_channels, inner_dim, patch_size) self.glyph_projector = GlmImageFeedForward( - dim=text_embed_dim, dim_out=inner_dim, inner_dim=inner_dim, activation_fn="gelu" + dim=text_embed_dim, + dim_out=inner_dim, + inner_dim=inner_dim, + activation_fn="gelu", + quant_config=quant_config, ) self.prior_token_embedding = nn.Embedding(prior_vq_quantizer_codebook_size, inner_dim) self.prior_projector = GlmImageFeedForward( - dim=inner_dim, dim_out=inner_dim, inner_dim=inner_dim, activation_fn="linear-silu" + dim=inner_dim, + dim_out=inner_dim, + inner_dim=inner_dim, + activation_fn="linear-silu", + quant_config=quant_config, ) # Prepare module for SP (encapsulates patch embedding and RoPE for _sp_plan) @@ -958,6 +988,7 @@ class GlmImageTransformer2DModel(CachedTransformer): attention_head_dim, time_embed_dim, ffn_hidden_dim=ffn_hidden_dim, + quant_config=quant_config, parallel_config=self.parallel_config, ) for _ in range(num_layers) diff --git a/vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py b/vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py index 038636499..97cba18c2 100644 --- a/vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py +++ b/vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py @@ -301,7 +301,10 @@ class GlmImagePipeline(nn.Module, DiffusionPipelineProfilerMixin): # Load transformer (DiT) logger.info("Loading GlmImageTransformer2DModel (DiT)...") - self.transformer = GlmImageTransformer2DModel(od_config=od_config) + self.transformer = GlmImageTransformer2DModel( + od_config=od_config, + quant_config=od_config.quantization_config, + ) # Weight sources for DiT loading self.weights_sources = [ diff --git a/vllm_omni/diffusion/models/interface.py b/vllm_omni/diffusion/models/interface.py index ef906472b..00d54420d 100644 --- a/vllm_omni/diffusion/models/interface.py +++ b/vllm_omni/diffusion/models/interface.py @@ -58,6 +58,27 @@ class SupportsStepExecution(Protocol): """Decode output after denoise loop.""" +@runtime_checkable +class SupportsModuleOffload(Protocol): + """Declares which submodules participate in CPU offload. + + All attribute names support dotted paths for nested submodules + (e.g. ``"pipe.transformer"``). + + Attributes: + _dit_modules: Denoising submodules (on GPU during diffusion). + _encoder_modules: Encoder submodules (offloaded during diffusion). + _vae_modules: VAE(s) (always on GPU). + _resident_modules: Extra modules pinned on GPU during layerwise + offloading. Optional, defaults to ``[]``. + """ + + _dit_modules: ClassVar[list[str]] + _encoder_modules: ClassVar[list[str]] + _vae_modules: ClassVar[list[str]] + _resident_modules: ClassVar[list[str]] = [] + + def supports_step_execution(pipeline: object) -> bool: """Return whether `pipeline` implements :class:`SupportsStepExecution`.""" diff --git a/vllm_omni/diffusion/models/ltx2/__init__.py b/vllm_omni/diffusion/models/ltx2/__init__.py index 2a78b61ba..bf57d3f9e 100644 --- a/vllm_omni/diffusion/models/ltx2/__init__.py +++ b/vllm_omni/diffusion/models/ltx2/__init__.py @@ -10,6 +10,10 @@ from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import ( get_ltx2_post_process_func, load_transformer_config, ) +from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import ( + LTX23ImageToVideoPipeline, + LTX23Pipeline, +) from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_image2video import ( LTX2I2VDMD2Pipeline, LTX2ImageToVideoPipeline, @@ -25,6 +29,8 @@ __all__ = [ "LTX2LatentUpsamplePipeline", "LTX2TwoStagesPipeline", "LTX2ImageToVideoTwoStagesPipeline", + "LTX23Pipeline", + "LTX23ImageToVideoPipeline", "get_ltx2_post_process_func", "load_transformer_config", "create_transformer_from_config", diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py b/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py index 0ae71aa2d..cd8eee9e0 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py @@ -436,6 +436,10 @@ class LTX2AudioVideoAttnProcessor: sequence_length=sequence_length, ) + # Compute gate logits from original hidden_states (before attention) + if attn.to_gate_logits is not None: + gate_logits = attn.to_gate_logits(hidden_states) + if is_self_attention: encoder_hidden_states = hidden_states @@ -471,6 +475,14 @@ class LTX2AudioVideoAttnProcessor: hidden_states = hidden_states.flatten(2, 3) hidden_states = hidden_states.to(query.dtype) + # LTX-2.3: per-head gated attention + if attn.to_gate_logits is not None: + hidden_states = hidden_states.unflatten(2, (attn.heads, -1)) # [B, T, H, D] + # 2.0 * sigmoid so zero-init gates produce 1.0 (identity) + gates = 2.0 * torch.sigmoid(gate_logits) # [B, T, H] + hidden_states = hidden_states * gates.unsqueeze(-1) + hidden_states = hidden_states.flatten(2, 3) + hidden_states = attn.to_out[0](hidden_states) if isinstance(hidden_states, tuple): hidden_states = hidden_states[0] @@ -501,11 +513,16 @@ class LTX2Attention(torch.nn.Module): norm_eps: float = 1e-6, norm_elementwise_affine: bool = True, rope_type: str = "interleaved", + apply_gated_attention: bool = False, processor=None, ): super().__init__() - if qk_norm != "rms_norm_across_heads": - raise NotImplementedError("Only 'rms_norm_across_heads' is supported as a valid value for `qk_norm`.") + # LTX-2 uses "rms_norm_across_heads", LTX-2.3 uses "rms_norm" -- both + # map to the same RMSNorm implementation applied across Q/K heads. + if qk_norm not in ("rms_norm_across_heads", "rms_norm"): + raise NotImplementedError( + f"Only 'rms_norm_across_heads' and 'rms_norm' are supported for `qk_norm`, got {qk_norm!r}." + ) kv_heads = heads if kv_heads is None else kv_heads @@ -563,18 +580,34 @@ class LTX2Attention(torch.nn.Module): self.heads = self.query_num_heads tp_size = get_tensor_model_parallel_world_size() - self.norm_q = TensorParallelRMSNorm( - dim_head * self.query_num_heads, - eps=norm_eps, - elementwise_affine=norm_elementwise_affine, - tp_size=tp_size, - ) - self.norm_k = TensorParallelRMSNorm( - dim_head * self.kv_num_heads, - eps=norm_eps, - elementwise_affine=norm_elementwise_affine, - tp_size=tp_size, - ) + # At TP > 1 with rms_norm_across_heads, use TensorParallelRMSNorm + # which all-reduces squared sums to match global RMS statistics. + # At TP=1, use torch.nn.RMSNorm which is numerically identical to + # the diffusers reference (verified via hook-based comparison). + if tp_size > 1 and qk_norm == "rms_norm_across_heads": + self.norm_q = TensorParallelRMSNorm( + dim_head * self.query_num_heads, + eps=norm_eps, + elementwise_affine=norm_elementwise_affine, + tp_size=tp_size, + ) + self.norm_k = TensorParallelRMSNorm( + dim_head * self.kv_num_heads, + eps=norm_eps, + elementwise_affine=norm_elementwise_affine, + tp_size=tp_size, + ) + else: + self.norm_q = torch.nn.RMSNorm( + dim_head * self.query_num_heads, + eps=norm_eps, + elementwise_affine=norm_elementwise_affine, + ) + self.norm_k = torch.nn.RMSNorm( + dim_head * self.kv_num_heads, + eps=norm_eps, + elementwise_affine=norm_elementwise_affine, + ) self.to_out = torch.nn.ModuleList( [ @@ -596,6 +629,12 @@ class LTX2Attention(torch.nn.Module): causal=False, ) + # LTX-2.3: per-head gated attention + if apply_gated_attention: + self.to_gate_logits = nn.Linear(query_dim, heads, bias=True) + else: + self.to_gate_logits = None + if processor is None: processor = self._default_processor_cls() self.set_processor(processor) @@ -697,6 +736,10 @@ class LTX2VideoTransformerBlock(nn.Module): audio_num_attention_heads: int, audio_attention_head_dim, audio_cross_attention_dim: int, + video_gated_attn: bool = False, + video_cross_attn_adaln: bool = False, + audio_gated_attn: bool = False, + audio_cross_attn_adaln: bool = False, qk_norm: str = "rms_norm_across_heads", activation_fn: str = "gelu-approximate", attention_bias: bool = True, @@ -704,8 +747,12 @@ class LTX2VideoTransformerBlock(nn.Module): eps: float = 1e-6, elementwise_affine: bool = False, rope_type: str = "interleaved", + perturbed_attn: bool = False, ): super().__init__() + self.video_cross_attn_adaln = video_cross_attn_adaln + self.audio_cross_attn_adaln = audio_cross_attn_adaln + self.perturbed_attn = perturbed_attn # 1. Self-Attention (video and audio) self.norm1 = _make_rms_norm(dim, eps=eps, elementwise_affine=elementwise_affine) @@ -719,6 +766,7 @@ class LTX2VideoTransformerBlock(nn.Module): out_bias=attention_out_bias, qk_norm=qk_norm, rope_type=rope_type, + apply_gated_attention=video_gated_attn, ) self.audio_norm1 = _make_rms_norm(audio_dim, eps=eps, elementwise_affine=elementwise_affine) @@ -732,6 +780,7 @@ class LTX2VideoTransformerBlock(nn.Module): out_bias=attention_out_bias, qk_norm=qk_norm, rope_type=rope_type, + apply_gated_attention=audio_gated_attn, ) # 2. Prompt Cross-Attention @@ -746,6 +795,7 @@ class LTX2VideoTransformerBlock(nn.Module): out_bias=attention_out_bias, qk_norm=qk_norm, rope_type=rope_type, + apply_gated_attention=video_gated_attn, ) self.audio_norm2 = _make_rms_norm(audio_dim, eps=eps, elementwise_affine=elementwise_affine) @@ -759,10 +809,10 @@ class LTX2VideoTransformerBlock(nn.Module): out_bias=attention_out_bias, qk_norm=qk_norm, rope_type=rope_type, + apply_gated_attention=audio_gated_attn, ) # 3. Audio-to-Video (a2v) and Video-to-Audio (v2a) Cross-Attention - # Audio-to-Video (a2v) Attention --> Q: Video; K,V: Audio self.audio_to_video_norm = _make_rms_norm(dim, eps=eps, elementwise_affine=elementwise_affine) self.audio_to_video_attn = LTX2Attention( query_dim=dim, @@ -774,9 +824,9 @@ class LTX2VideoTransformerBlock(nn.Module): out_bias=attention_out_bias, qk_norm=qk_norm, rope_type=rope_type, + apply_gated_attention=video_gated_attn, ) - # Video-to-Audio (v2a) Attention --> Q: Audio; K,V: Video self.video_to_audio_norm = _make_rms_norm(audio_dim, eps=eps, elementwise_affine=elementwise_affine) self.video_to_audio_attn = LTX2Attention( query_dim=audio_dim, @@ -788,6 +838,7 @@ class LTX2VideoTransformerBlock(nn.Module): out_bias=attention_out_bias, qk_norm=qk_norm, rope_type=rope_type, + apply_gated_attention=audio_gated_attn, ) # 4. Feedforward layers @@ -798,14 +849,33 @@ class LTX2VideoTransformerBlock(nn.Module): self.audio_ff = LTX2FeedForward(audio_dim, activation_fn=activation_fn) # 5. Per-Layer Modulation Parameters - # Self-Attention / Feedforward AdaLayerNorm-Zero mod params - self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5) - self.audio_scale_shift_table = nn.Parameter(torch.randn(6, audio_dim) / audio_dim**0.5) + # LTX-2.3 with cross_attn_adaln uses 9 params (extra 3 for cross-attn modulation); + # LTX-2 uses 6. + video_mod_param_num = 9 if self.video_cross_attn_adaln else 6 + audio_mod_param_num = 9 if self.audio_cross_attn_adaln else 6 + self.scale_shift_table = nn.Parameter(torch.randn(video_mod_param_num, dim) / dim**0.5) + self.audio_scale_shift_table = nn.Parameter(torch.randn(audio_mod_param_num, audio_dim) / audio_dim**0.5) + + # Prompt cross-attn additional modulation params (LTX-2.3) + self.cross_attn_adaln = video_cross_attn_adaln or audio_cross_attn_adaln + if self.cross_attn_adaln: + self.prompt_scale_shift_table = nn.Parameter(torch.randn(2, dim)) + self.audio_prompt_scale_shift_table = nn.Parameter(torch.randn(2, audio_dim)) # Per-layer a2v, v2a Cross-Attention mod params self.video_a2v_cross_attn_scale_shift_table = nn.Parameter(torch.randn(5, dim)) self.audio_a2v_cross_attn_scale_shift_table = nn.Parameter(torch.randn(5, audio_dim)) + @staticmethod + def get_mod_params( + scale_shift_table: torch.Tensor, temb: torch.Tensor, batch_size: int + ) -> tuple[torch.Tensor, ...]: + num_ada_params = scale_shift_table.shape[0] + ada_values = scale_shift_table[None, None].to(temb.device) + temb.reshape( + batch_size, temb.shape[1], num_ada_params, -1 + ) + return ada_values.unbind(dim=2) + def forward( self, hidden_states: torch.Tensor, @@ -818,143 +888,170 @@ class LTX2VideoTransformerBlock(nn.Module): temb_ca_audio_scale_shift: torch.Tensor, temb_ca_gate: torch.Tensor, temb_ca_audio_gate: torch.Tensor, + temb_prompt: torch.Tensor | None = None, + temb_prompt_audio: torch.Tensor | None = None, video_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, audio_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, ca_video_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, ca_audio_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, encoder_attention_mask: torch.Tensor | None = None, audio_encoder_attention_mask: torch.Tensor | None = None, + self_attention_mask: torch.Tensor | None = None, + audio_self_attention_mask: torch.Tensor | None = None, a2v_cross_attention_mask: torch.Tensor | None = None, v2a_cross_attention_mask: torch.Tensor | None = None, + use_a2v_cross_attention: bool = True, + use_v2a_cross_attention: bool = True, + perturbation_mask: torch.Tensor | None = None, + all_perturbed: bool | None = None, ) -> torch.Tensor: batch_size = hidden_states.size(0) # 1. Video and Audio Self-Attention - norm_hidden_states = self.norm1(hidden_states) + # 1.1. Video Self-Attention + video_ada_params = self.get_mod_params(self.scale_shift_table, temb, batch_size) + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = video_ada_params[:6] + if self.video_cross_attn_adaln: + shift_text_q, scale_text_q, gate_text_q = video_ada_params[6:9] - num_ada_params = self.scale_shift_table.shape[0] - ada_values = self.scale_shift_table[None, None].to(temb.device) + temb.reshape( - batch_size, temb.size(1), num_ada_params, -1 - ) - shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ada_values.unbind(dim=2) + norm_hidden_states = self.norm1(hidden_states) norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa attn_hidden_states = self.attn1( hidden_states=norm_hidden_states, encoder_hidden_states=None, query_rotary_emb=video_rotary_emb, + attention_mask=self_attention_mask, ) hidden_states = hidden_states + attn_hidden_states * gate_msa - norm_audio_hidden_states = self.audio_norm1(audio_hidden_states) - - num_audio_ada_params = self.audio_scale_shift_table.shape[0] - audio_ada_values = self.audio_scale_shift_table[None, None].to(temb_audio.device) + temb_audio.reshape( - batch_size, temb_audio.size(1), num_audio_ada_params, -1 - ) + # 1.2. Audio Self-Attention + audio_ada_params = self.get_mod_params(self.audio_scale_shift_table, temb_audio, batch_size) audio_shift_msa, audio_scale_msa, audio_gate_msa, audio_shift_mlp, audio_scale_mlp, audio_gate_mlp = ( - audio_ada_values.unbind(dim=2) + audio_ada_params[:6] ) + if self.audio_cross_attn_adaln: + audio_shift_text_q, audio_scale_text_q, audio_gate_text_q = audio_ada_params[6:9] + + norm_audio_hidden_states = self.audio_norm1(audio_hidden_states) norm_audio_hidden_states = norm_audio_hidden_states * (1 + audio_scale_msa) + audio_shift_msa attn_audio_hidden_states = self.audio_attn1( hidden_states=norm_audio_hidden_states, encoder_hidden_states=None, query_rotary_emb=audio_rotary_emb, + attention_mask=audio_self_attention_mask, ) audio_hidden_states = audio_hidden_states + attn_audio_hidden_states * audio_gate_msa - # 2. Video and Audio Cross-Attention with the text embeddings + # 2. Video and Audio Cross-Attention with text embeddings (Q: Video/Audio; K,V: Text) + # LTX-2.3: compute prompt modulation params for K/V + if self.cross_attn_adaln and temb_prompt is not None: + video_prompt_ada_params = self.get_mod_params(self.prompt_scale_shift_table, temb_prompt, batch_size) + shift_text_kv, scale_text_kv = video_prompt_ada_params + + audio_prompt_ada_params = self.get_mod_params( + self.audio_prompt_scale_shift_table, temb_prompt_audio, batch_size + ) + audio_shift_text_kv, audio_scale_text_kv = audio_prompt_ada_params + + # 2.1. Video-Text Cross-Attention norm_hidden_states = self.norm2(hidden_states) + if self.video_cross_attn_adaln: + norm_hidden_states = norm_hidden_states * (1 + scale_text_q) + shift_text_q + if self.cross_attn_adaln and temb_prompt is not None: + encoder_hidden_states = encoder_hidden_states * (1 + scale_text_kv) + shift_text_kv + attn_hidden_states = self.attn2( norm_hidden_states, encoder_hidden_states=encoder_hidden_states, query_rotary_emb=None, attention_mask=encoder_attention_mask, ) + if self.video_cross_attn_adaln: + attn_hidden_states = attn_hidden_states * gate_text_q hidden_states = hidden_states + attn_hidden_states + # 2.2. Audio-Text Cross-Attention norm_audio_hidden_states = self.audio_norm2(audio_hidden_states) + if self.audio_cross_attn_adaln: + norm_audio_hidden_states = norm_audio_hidden_states * (1 + audio_scale_text_q) + audio_shift_text_q + if self.cross_attn_adaln and temb_prompt is not None: + audio_encoder_hidden_states = audio_encoder_hidden_states * (1 + audio_scale_text_kv) + audio_shift_text_kv + attn_audio_hidden_states = self.audio_attn2( norm_audio_hidden_states, encoder_hidden_states=audio_encoder_hidden_states, query_rotary_emb=None, attention_mask=audio_encoder_attention_mask, ) + if self.audio_cross_attn_adaln: + attn_audio_hidden_states = attn_audio_hidden_states * audio_gate_text_q audio_hidden_states = audio_hidden_states + attn_audio_hidden_states # 3. Audio-to-Video (a2v) and Video-to-Audio (v2a) Cross-Attention - norm_hidden_states = self.audio_to_video_norm(hidden_states) - norm_audio_hidden_states = self.video_to_audio_norm(audio_hidden_states) + if use_a2v_cross_attention or use_v2a_cross_attention: + norm_hidden_states = self.audio_to_video_norm(hidden_states) + norm_audio_hidden_states = self.video_to_audio_norm(audio_hidden_states) - # Combine global and per-layer cross attention modulation parameters - # Video - video_per_layer_ca_scale_shift = self.video_a2v_cross_attn_scale_shift_table[:4, :] - video_per_layer_ca_gate = self.video_a2v_cross_attn_scale_shift_table[4:, :] + # Combine global and per-layer cross attention modulation parameters + # Video + video_per_layer_ca_scale_shift = self.video_a2v_cross_attn_scale_shift_table[:4, :] + video_per_layer_ca_gate = self.video_a2v_cross_attn_scale_shift_table[4:, :] - video_ca_scale_shift_table = ( - video_per_layer_ca_scale_shift[:, :, ...].to(temb_ca_scale_shift.dtype) - + temb_ca_scale_shift.reshape(batch_size, temb_ca_scale_shift.shape[1], 4, -1) - ).unbind(dim=2) - video_ca_gate = ( - video_per_layer_ca_gate[:, :, ...].to(temb_ca_gate.dtype) - + temb_ca_gate.reshape(batch_size, temb_ca_gate.shape[1], 1, -1) - ).unbind(dim=2) + video_ca_ada_params = self.get_mod_params(video_per_layer_ca_scale_shift, temb_ca_scale_shift, batch_size) + video_ca_gate_param = self.get_mod_params(video_per_layer_ca_gate, temb_ca_gate, batch_size) - video_a2v_ca_scale, video_a2v_ca_shift, video_v2a_ca_scale, video_v2a_ca_shift = video_ca_scale_shift_table - a2v_gate = video_ca_gate[0].squeeze(2) + video_a2v_ca_scale, video_a2v_ca_shift, video_v2a_ca_scale, video_v2a_ca_shift = video_ca_ada_params + a2v_gate = video_ca_gate_param[0].squeeze(2) - # Audio - audio_per_layer_ca_scale_shift = self.audio_a2v_cross_attn_scale_shift_table[:4, :] - audio_per_layer_ca_gate = self.audio_a2v_cross_attn_scale_shift_table[4:, :] + # Audio + audio_per_layer_ca_scale_shift = self.audio_a2v_cross_attn_scale_shift_table[:4, :] + audio_per_layer_ca_gate = self.audio_a2v_cross_attn_scale_shift_table[4:, :] - audio_ca_scale_shift_table = ( - audio_per_layer_ca_scale_shift[:, :, ...].to(temb_ca_audio_scale_shift.dtype) - + temb_ca_audio_scale_shift.reshape(batch_size, temb_ca_audio_scale_shift.shape[1], 4, -1) - ).unbind(dim=2) - audio_ca_gate = ( - audio_per_layer_ca_gate[:, :, ...].to(temb_ca_audio_gate.dtype) - + temb_ca_audio_gate.reshape(batch_size, temb_ca_audio_gate.shape[1], 1, -1) - ).unbind(dim=2) + audio_ca_ada_params = self.get_mod_params( + audio_per_layer_ca_scale_shift, temb_ca_audio_scale_shift, batch_size + ) + audio_ca_gate_param = self.get_mod_params(audio_per_layer_ca_gate, temb_ca_audio_gate, batch_size) - audio_a2v_ca_scale, audio_a2v_ca_shift, audio_v2a_ca_scale, audio_v2a_ca_shift = audio_ca_scale_shift_table - v2a_gate = audio_ca_gate[0].squeeze(2) + audio_a2v_ca_scale, audio_a2v_ca_shift, audio_v2a_ca_scale, audio_v2a_ca_shift = audio_ca_ada_params + v2a_gate = audio_ca_gate_param[0].squeeze(2) - # Audio-to-Video Cross Attention: Q: Video; K,V: Audio - mod_norm_hidden_states = norm_hidden_states * (1 + video_a2v_ca_scale.squeeze(2)) + video_a2v_ca_shift.squeeze( - 2 - ) - mod_norm_audio_hidden_states = norm_audio_hidden_states * ( - 1 + audio_a2v_ca_scale.squeeze(2) - ) + audio_a2v_ca_shift.squeeze(2) + # 3.2. Audio-to-Video Cross Attention: Q: Video; K,V: Audio + if use_a2v_cross_attention: + mod_norm_hidden_states = norm_hidden_states * ( + 1 + video_a2v_ca_scale.squeeze(2) + ) + video_a2v_ca_shift.squeeze(2) + mod_norm_audio_hidden_states = norm_audio_hidden_states * ( + 1 + audio_a2v_ca_scale.squeeze(2) + ) + audio_a2v_ca_shift.squeeze(2) - a2v_attn_hidden_states = self.audio_to_video_attn( - mod_norm_hidden_states, - encoder_hidden_states=mod_norm_audio_hidden_states, - query_rotary_emb=ca_video_rotary_emb, - key_rotary_emb=ca_audio_rotary_emb, - attention_mask=a2v_cross_attention_mask, - ) + a2v_attn_hidden_states = self.audio_to_video_attn( + mod_norm_hidden_states, + encoder_hidden_states=mod_norm_audio_hidden_states, + query_rotary_emb=ca_video_rotary_emb, + key_rotary_emb=ca_audio_rotary_emb, + attention_mask=a2v_cross_attention_mask, + ) + hidden_states = hidden_states + a2v_gate * a2v_attn_hidden_states - hidden_states = hidden_states + a2v_gate * a2v_attn_hidden_states + # 3.3. Video-to-Audio Cross Attention: Q: Audio; K,V: Video + if use_v2a_cross_attention: + mod_norm_hidden_states = norm_hidden_states * ( + 1 + video_v2a_ca_scale.squeeze(2) + ) + video_v2a_ca_shift.squeeze(2) + mod_norm_audio_hidden_states = norm_audio_hidden_states * ( + 1 + audio_v2a_ca_scale.squeeze(2) + ) + audio_v2a_ca_shift.squeeze(2) - # Video-to-Audio Cross Attention: Q: Audio; K,V: Video - mod_norm_hidden_states = norm_hidden_states * (1 + video_v2a_ca_scale.squeeze(2)) + video_v2a_ca_shift.squeeze( - 2 - ) - mod_norm_audio_hidden_states = norm_audio_hidden_states * ( - 1 + audio_v2a_ca_scale.squeeze(2) - ) + audio_v2a_ca_shift.squeeze(2) - - v2a_attn_hidden_states = self.video_to_audio_attn( - mod_norm_audio_hidden_states, - encoder_hidden_states=mod_norm_hidden_states, - query_rotary_emb=ca_audio_rotary_emb, - key_rotary_emb=ca_video_rotary_emb, - attention_mask=v2a_cross_attention_mask, - ) - - audio_hidden_states = audio_hidden_states + v2a_gate * v2a_attn_hidden_states + v2a_attn_hidden_states = self.video_to_audio_attn( + mod_norm_audio_hidden_states, + encoder_hidden_states=mod_norm_hidden_states, + query_rotary_emb=ca_audio_rotary_emb, + key_rotary_emb=ca_video_rotary_emb, + attention_mask=v2a_cross_attention_mask, + ) + audio_hidden_states = audio_hidden_states + v2a_gate * v2a_attn_hidden_states # 4. Feedforward norm_hidden_states = self.norm3(hidden_states) * (1 + scale_mlp) + shift_mlp @@ -1350,8 +1447,15 @@ class LTX2VideoTransformer3DModel(nn.Module): timestep_scale_multiplier: int = 1000, cross_attn_timestep_scale_multiplier: int = 1000, rope_type: str = "interleaved", + use_prompt_embeddings: bool = True, + perturbed_attn: bool = False, + gated_attn: bool = False, + cross_attn_mod: bool = False, + audio_gated_attn: bool = False, + audio_cross_attn_mod: bool = False, ) -> None: super().__init__() + self.perturbed_attn = perturbed_attn out_channels = out_channels or in_channels audio_out_channels = audio_out_channels or audio_in_channels @@ -1401,18 +1505,33 @@ class LTX2VideoTransformer3DModel(nn.Module): self.audio_proj_in = nn.Linear(audio_in_channels, audio_inner_dim) # 2. Prompt embeddings - self.caption_projection = PixArtAlphaTextProjection(in_features=caption_channels, hidden_size=inner_dim) - self.audio_caption_projection = PixArtAlphaTextProjection( - in_features=caption_channels, hidden_size=audio_inner_dim - ) + # LTX-2 (use_prompt_embeddings=True): caption projection in the transformer + # LTX-2.3 (use_prompt_embeddings=False): caption projection in the connectors + if use_prompt_embeddings: + self.caption_projection = PixArtAlphaTextProjection(in_features=caption_channels, hidden_size=inner_dim) + self.audio_caption_projection = PixArtAlphaTextProjection( + in_features=caption_channels, hidden_size=audio_inner_dim + ) # 3. Timestep Modulation Params and Embedding - # 3.1. Global Timestep Modulation Parameters (except for cross-attention) and timestep + size embedding - # time_embed and audio_time_embed calculate both the timestep embedding and (global) modulation parameters - self.time_embed = LTX2AdaLayerNormSingle(inner_dim, num_mod_params=6, use_additional_conditions=False) - self.audio_time_embed = LTX2AdaLayerNormSingle( - audio_inner_dim, num_mod_params=6, use_additional_conditions=False + # 3.1. Global Timestep Modulation Parameters + # LTX-2.3 with cross_attn_mod uses 9 mod params (extra 3 for cross-attn); LTX-2 uses 6. + video_num_mod_params = 9 if cross_attn_mod else 6 + audio_num_mod_params = 9 if audio_cross_attn_mod else 6 + self.time_embed = LTX2AdaLayerNormSingle( + inner_dim, num_mod_params=video_num_mod_params, use_additional_conditions=False ) + self.audio_time_embed = LTX2AdaLayerNormSingle( + audio_inner_dim, num_mod_params=audio_num_mod_params, use_additional_conditions=False + ) + + # 3.3. LTX-2.3: Prompt modulation from sigma + self.prompt_modulation = cross_attn_mod or audio_cross_attn_mod + if self.prompt_modulation: + self.prompt_adaln = LTX2AdaLayerNormSingle(inner_dim, num_mod_params=2, use_additional_conditions=False) + self.audio_prompt_adaln = LTX2AdaLayerNormSingle( + audio_inner_dim, num_mod_params=2, use_additional_conditions=False + ) # 3.2. Global Cross Attention Modulation Parameters # Used in the audio-to-video and video-to-audio cross attention layers as a global set of modulation params, @@ -1516,6 +1635,10 @@ class LTX2VideoTransformer3DModel(nn.Module): audio_num_attention_heads=audio_num_attention_heads, audio_attention_head_dim=audio_attention_head_dim, audio_cross_attention_dim=audio_cross_attention_dim, + video_gated_attn=gated_attn, + video_cross_attn_adaln=cross_attn_mod, + audio_gated_attn=audio_gated_attn, + audio_cross_attn_adaln=audio_cross_attn_mod, qk_norm=qk_norm, activation_fn=activation_fn, attention_bias=attention_bias, @@ -1523,6 +1646,7 @@ class LTX2VideoTransformer3DModel(nn.Module): eps=norm_eps, elementwise_affine=norm_elementwise_affine, rope_type=rope_type, + perturbed_attn=perturbed_attn, ) for _ in range(num_layers) ] @@ -1558,6 +1682,8 @@ class LTX2VideoTransformer3DModel(nn.Module): audio_encoder_hidden_states: torch.Tensor, timestep: torch.LongTensor, audio_timestep: torch.LongTensor | None = None, + sigma: torch.Tensor | None = None, + audio_sigma: torch.Tensor | None = None, encoder_attention_mask: torch.Tensor | None = None, audio_encoder_attention_mask: torch.Tensor | None = None, num_frames: int | None = None, @@ -1569,6 +1695,7 @@ class LTX2VideoTransformer3DModel(nn.Module): audio_coords: torch.Tensor | None = None, attention_kwargs: dict[str, Any] | None = None, return_dict: bool = True, + **kwargs, # Accept extra diffusers kwargs (isolate_modalities, perturbation_mask, etc.) ) -> torch.Tensor: """ Forward pass for LTX-2.0 audiovisual video transformer. @@ -1711,54 +1838,55 @@ class LTX2VideoTransformer3DModel(nn.Module): ) audio_cross_attn_v2a_gate = audio_cross_attn_v2a_gate.view(batch_size, -1, audio_cross_attn_v2a_gate.shape[-1]) - # 4. Prepare prompt embeddings - encoder_hidden_states = self.caption_projection(encoder_hidden_states) - encoder_hidden_states = encoder_hidden_states.view(batch_size, -1, hidden_states.size(-1)) + # 3.3. LTX-2.3: Compute prompt modulation from sigma + audio_sigma = audio_sigma if audio_sigma is not None else sigma + if self.prompt_modulation and sigma is not None: + temb_prompt, _ = self.prompt_adaln(sigma.flatten(), batch_size=batch_size, hidden_dtype=hidden_states.dtype) + temb_prompt_audio, _ = self.audio_prompt_adaln( + audio_sigma.flatten(), batch_size=batch_size, hidden_dtype=audio_hidden_states.dtype + ) + temb_prompt = temb_prompt.view(batch_size, -1, temb_prompt.size(-1)) + temb_prompt_audio = temb_prompt_audio.view(batch_size, -1, temb_prompt_audio.size(-1)) + else: + temb_prompt = temb_prompt_audio = None - audio_encoder_hidden_states = self.audio_caption_projection(audio_encoder_hidden_states) - audio_encoder_hidden_states = audio_encoder_hidden_states.view(batch_size, -1, audio_hidden_states.size(-1)) + # 4. Prepare prompt embeddings + # LTX-2: caption projection is in the transformer + # LTX-2.3: caption projection is in the connectors (already applied) + if hasattr(self, "caption_projection"): + encoder_hidden_states = self.caption_projection(encoder_hidden_states) + encoder_hidden_states = encoder_hidden_states.view(batch_size, -1, hidden_states.size(-1)) + + if hasattr(self, "audio_caption_projection"): + audio_encoder_hidden_states = self.audio_caption_projection(audio_encoder_hidden_states) + audio_encoder_hidden_states = audio_encoder_hidden_states.view(batch_size, -1, audio_hidden_states.size(-1)) # 5. Run transformer blocks for block in self.transformer_blocks: + block_kwargs = { + "hidden_states": hidden_states, + "audio_hidden_states": audio_hidden_states, + "encoder_hidden_states": encoder_hidden_states, + "audio_encoder_hidden_states": audio_encoder_hidden_states, + "temb": temb, + "temb_audio": temb_audio, + "temb_ca_scale_shift": video_cross_attn_scale_shift, + "temb_ca_audio_scale_shift": audio_cross_attn_scale_shift, + "temb_ca_gate": video_cross_attn_a2v_gate, + "temb_ca_audio_gate": audio_cross_attn_v2a_gate, + "temb_prompt": temb_prompt, + "temb_prompt_audio": temb_prompt_audio, + "video_rotary_emb": video_rotary_emb, + "audio_rotary_emb": audio_rotary_emb, + "ca_video_rotary_emb": video_cross_attn_rotary_emb, + "ca_audio_rotary_emb": audio_cross_attn_rotary_emb, + "encoder_attention_mask": encoder_attention_mask, + "audio_encoder_attention_mask": audio_encoder_attention_mask, + } if torch.is_grad_enabled() and self.gradient_checkpointing: - hidden_states, audio_hidden_states = self._gradient_checkpointing_func( - block, - hidden_states, - audio_hidden_states, - encoder_hidden_states, - audio_encoder_hidden_states, - temb, - temb_audio, - video_cross_attn_scale_shift, - audio_cross_attn_scale_shift, - video_cross_attn_a2v_gate, - audio_cross_attn_v2a_gate, - video_rotary_emb, - audio_rotary_emb, - video_cross_attn_rotary_emb, - audio_cross_attn_rotary_emb, - encoder_attention_mask, - audio_encoder_attention_mask, - ) + hidden_states, audio_hidden_states = self._gradient_checkpointing_func(block, **block_kwargs) else: - hidden_states, audio_hidden_states = block( - hidden_states, - audio_hidden_states, - encoder_hidden_states, - audio_encoder_hidden_states, - temb, - temb_audio, - video_cross_attn_scale_shift, - audio_cross_attn_scale_shift, - video_cross_attn_a2v_gate, - audio_cross_attn_v2a_gate, - video_rotary_emb, - audio_rotary_emb, - video_cross_attn_rotary_emb, - audio_cross_attn_rotary_emb, - encoder_attention_mask, - audio_encoder_attention_mask, - ) + hidden_states, audio_hidden_states = block(**block_kwargs) # 6. Output layers (including unpatchification) scale_shift_values = self.scale_shift_table[None, None] + embedded_timestep[:, :, None] @@ -1826,6 +1954,14 @@ class LTX2VideoTransformer3DModel(nn.Module): weight_loader(param, loaded_weight, shard_id) break else: + if name not in params_dict: + logger.warning( + "Skipping transformer weight %s -- not found in model " + "parameters. This may indicate an incomplete " + "implementation or checkpoint mismatch.", + name, + ) + continue param = params_dict[name] weight_loader = getattr(param, "weight_loader", None) if weight_loader is not None: diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py index f06ffab16..4f62d72c9 100644 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py +++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py @@ -9,7 +9,7 @@ import json import os from collections.abc import Iterable from contextlib import nullcontext -from typing import Any +from typing import Any, ClassVar import numpy as np import torch @@ -34,6 +34,7 @@ from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin +from vllm_omni.diffusion.models.interface import SupportsModuleOffload from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.lora.request import LoRARequest @@ -1152,9 +1153,13 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin): return loader.load_weights(weights) -class LTX2TwoStagesPipeline(nn.Module): +class LTX2TwoStagesPipeline(nn.Module, SupportsModuleOffload): """LTX2TwoStagesPipeline is for two stages image to video generation""" + _dit_modules: ClassVar[list[str]] = ["pipe.transformer"] + _encoder_modules: ClassVar[list[str]] = ["pipe.text_encoder"] + _vae_modules: ClassVar[list[str]] = ["pipe.vae", "pipe.audio_vae"] + def __init__( self, *, diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py new file mode 100644 index 000000000..dd0a57179 --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py @@ -0,0 +1,998 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +Fully independent LTX-2.3 pipeline for vLLM-Omni. + +This pipeline does NOT inherit from LTX2Pipeline because: +- LTX-2.3 uses a different text encoding strategy (flatten ALL 49 hidden states + vs. LTX-2's _pack_text_embeds with per-layer normalization and pooling) +- LTX-2.3 connectors expect the padding_side API (not additive_mask) +- LTX-2.3 uses a BWE vocoder outputting 48kHz audio (not 16kHz) +- LTX-2.3 transformer requires the sigma parameter for prompt modulation +- CPU offloading is required for the 22B transformer (~44GB VRAM) +""" + +from __future__ import annotations + +import copy +import json +import os +from collections.abc import Iterable +from contextlib import nullcontext +from typing import Any + +import numpy as np +import torch +from diffusers import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video, FlowMatchEulerDiscreteScheduler +from diffusers.pipelines.ltx2 import LTX2TextConnectors +from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder +from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import retrieve_timesteps +from diffusers.utils.torch_utils import randn_tensor +from diffusers.video_processor import VideoProcessor +from huggingface_hub import hf_hub_download +from torch import nn +from transformers import AutoTokenizer, Gemma3ForConditionalGeneration +from vllm.logger import init_logger +from vllm.model_executor.models.utils import AutoWeightsLoader + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.utils import get_local_device +from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader +from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin +from vllm_omni.diffusion.request import OmniDiffusionRequest + +from .pipeline_ltx2 import ( + _get_prompt_field, + calculate_shift, + create_transformer_from_config, + load_transformer_config, +) + +logger = init_logger(__name__) + +# Try to import LTX2VocoderWithBWE (diffusers >= 0.38.0) +try: + from diffusers.pipelines.ltx2.vocoder import LTX2VocoderWithBWE +except ImportError: + LTX2VocoderWithBWE = None + + +def _detect_vocoder_output_sample_rate(model: str) -> int | None: + """Detect the vocoder output sample rate from vocoder/config.json. + + This runs at factory time (engine process) so the rate is captured in + the post-process closure and doesn't need cross-process communication. + + Returns: + Output sample rate (e.g. 48000 for LTX-2.3 BWE vocoder) or None. + """ + vocoder_config_path = os.path.join(model, "vocoder", "config.json") + if not os.path.exists(vocoder_config_path): + try: + vocoder_config_path = hf_hub_download(model, "vocoder/config.json") + except Exception: + return None + try: + with open(vocoder_config_path) as f: + cfg = json.load(f) + return cfg.get("output_sampling_rate") + except Exception: + return None + + +def get_ltx2_post_process_func(od_config: OmniDiffusionConfig): + """Factory for the LTX-2.3 post-process function. + + Detects the vocoder output sample rate at factory time and captures it + in the closure so that the audio_sample_rate flows through + DiffusionEngine -> OmniRequestOutput -> serving_video. + """ + output_sr = _detect_vocoder_output_sample_rate(od_config.model) + + def post_process_func(output: tuple[torch.Tensor, torch.Tensor] | torch.Tensor): + if isinstance(output, tuple) and len(output) == 2: + video, audio = output + if isinstance(audio, torch.Tensor): + audio = audio.detach().cpu() + result: dict[str, Any] = {"video": video, "audio": audio} + if output_sr is not None: + result["audio_sample_rate"] = output_sr + return result + return output + + return post_process_func + + +class LTX23Pipeline(nn.Module, ProgressBarMixin): + """Fully independent LTX-2.3 pipeline. + + Key differences from LTX2Pipeline: + - Text encoding: uses ALL 49 hidden states from Gemma-3-12B, flattened + - Connectors: uses padding_side API (not additive_mask) + - Vocoder: uses LTX2VocoderWithBWE (48kHz output) + - Transformer: passes sigma for prompt_adaln + - CPU offloading: text encoder, connectors, VAE, vocoder stay on CPU + """ + + def __init__( + self, + *, + od_config: OmniDiffusionConfig, + prefix: str = "", + ): + super().__init__() + self.od_config = od_config + self.device = get_local_device() + dtype = getattr(od_config, "dtype", torch.bfloat16) + model = od_config.model + local_files_only = os.path.exists(model) + + # Weight sources for transformer (loaded via AutoWeightsLoader) + self.weights_sources = [ + DiffusersPipelineLoader.ComponentSource( + model_or_path=od_config.model, + subfolder="transformer", + revision=None, + prefix="transformer.", + fall_back_to_pt=True, + ), + ] + + # --- Tokenizer (lightweight, stays wherever) --- + self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) + + # --- Text encoder: load on CPU, move to GPU only during encoding --- + with torch.device("cpu"): + self.text_encoder = Gemma3ForConditionalGeneration.from_pretrained( + model, subfolder="text_encoder", torch_dtype=dtype, local_files_only=local_files_only + ) + + # --- Connectors: CPU (LTX-2.3 connectors include caption projection) --- + self.connectors = LTX2TextConnectors.from_pretrained( + model, subfolder="connectors", torch_dtype=dtype, local_files_only=local_files_only + ) + + # --- VAE, Audio VAE: CPU --- + self.vae = AutoencoderKLLTX2Video.from_pretrained( + model, subfolder="vae", torch_dtype=dtype, local_files_only=local_files_only + ) + self.audio_vae = AutoencoderKLLTX2Audio.from_pretrained( + model, subfolder="audio_vae", torch_dtype=dtype, local_files_only=local_files_only + ) + + # --- Vocoder: prefer BWE vocoder (48kHz) for LTX-2.3 --- + vocoder_cls = LTX2VocoderWithBWE or LTX2Vocoder + try: + self.vocoder = vocoder_cls.from_pretrained( + model, subfolder="vocoder", torch_dtype=dtype, local_files_only=local_files_only + ) + except (TypeError, OSError, ValueError): + self.vocoder = LTX2Vocoder.from_pretrained( + model, subfolder="vocoder", torch_dtype=dtype, local_files_only=local_files_only + ) + + # --- Transformer: created empty, weights loaded via AutoWeightsLoader --- + transformer_config = load_transformer_config(model, "transformer", local_files_only) + self.transformer = create_transformer_from_config(transformer_config) + + # --- Scheduler --- + self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( + model, subfolder="scheduler", local_files_only=local_files_only + ) + + # --- Derived compression ratios --- + self.vae_spatial_compression_ratio = self.vae.spatial_compression_ratio if self.vae is not None else 32 + self.vae_temporal_compression_ratio = self.vae.temporal_compression_ratio if self.vae is not None else 8 + self.audio_vae_mel_compression_ratio = self.audio_vae.mel_compression_ratio if self.audio_vae is not None else 4 + self.audio_vae_temporal_compression_ratio = ( + self.audio_vae.temporal_compression_ratio if self.audio_vae is not None else 4 + ) + self.transformer_spatial_patch_size = self.transformer.config.patch_size if self.transformer is not None else 1 + self.transformer_temporal_patch_size = ( + self.transformer.config.patch_size_t if self.transformer is not None else 1 + ) + self.audio_sampling_rate = self.audio_vae.config.sample_rate if self.audio_vae is not None else 16000 + self.audio_hop_length = self.audio_vae.config.mel_hop_length if self.audio_vae is not None else 160 + + self.video_processor = VideoProcessor(vae_scale_factor=self.vae_spatial_compression_ratio) + + # Tokenizer max length + tokenizer_max_length = 1024 + if self.tokenizer is not None: + tokenizer_max_length = self.tokenizer.model_max_length + if tokenizer_max_length is None or tokenizer_max_length > 100000: + encoder_config = getattr(self.text_encoder, "config", None) + config_max_len = getattr(encoder_config, "max_position_embeddings", None) + if config_max_len is None: + config_max_len = getattr(encoder_config, "max_seq_len", None) + tokenizer_max_length = config_max_len or 1024 + self.tokenizer_max_length = int(tokenizer_max_length) + + # Pipeline state + self._guidance_scale = None + self._attention_kwargs = None + self._interrupt = False + self._num_timesteps = None + self._current_timestep = None + + # ------------------------------------------------------------------ + # Text Encoding (LTX-2.3 specific) + # ------------------------------------------------------------------ + + def _get_gemma_prompt_embeds( + self, + prompt: str | list[str], + num_videos_per_prompt: int = 1, + max_sequence_length: int = 1024, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ): + """Encode prompts using Gemma-3-12B, returning ALL 49 hidden states flattened. + + LTX-2.3 differs from LTX-2 in text encoding: + - LTX-2: uses _pack_text_embeds (layer selection + pooling) + - LTX-2.3: stacks ALL 49 hidden states and flattens to [B, seq, 188160] + The connectors unflatten, apply per_token_rms_norm, and project internally. + """ + device = device or self.device + dtype = dtype or self.text_encoder.dtype + + prompt = [prompt] if isinstance(prompt, str) else prompt + batch_size = len(prompt) + + if self.tokenizer is not None: + self.tokenizer.padding_side = "left" + if self.tokenizer.pad_token is None: + self.tokenizer.pad_token = self.tokenizer.eos_token + + prompt = [p.strip() for p in prompt] + text_inputs = self.tokenizer( + prompt, + padding="max_length", + max_length=max_sequence_length, + truncation=True, + add_special_tokens=True, + return_tensors="pt", + ) + text_input_ids = text_inputs.input_ids.to(device) + prompt_attention_mask = text_inputs.attention_mask.to(device) + + # Move text encoder to GPU for encoding + self.text_encoder.to(device) + text_encoder_outputs = self.text_encoder( + input_ids=text_input_ids, + attention_mask=prompt_attention_mask, + output_hidden_states=True, + ) + # Move text encoder back to CPU immediately + self.text_encoder.to("cpu") + torch.cuda.empty_cache() + + hidden_states = text_encoder_outputs.hidden_states + + # LTX-2.3: Stack ALL 49 hidden states and flatten + # [49 x (B, seq, 3840)] -> [B, seq, 3840, 49] -> [B, seq, 188160] + prompt_embeds = torch.stack(hidden_states, dim=-1).flatten(2, 3).to(dtype=dtype) + + _, seq_len, _ = prompt_embeds.shape + prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1) + prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1) + + prompt_attention_mask = prompt_attention_mask.view(batch_size, -1) + prompt_attention_mask = prompt_attention_mask.repeat(num_videos_per_prompt, 1) + + return prompt_embeds, prompt_attention_mask + + def encode_prompt( + self, + prompt: str | list[str], + negative_prompt: str | list[str] | None = None, + do_classifier_free_guidance: bool = True, + num_videos_per_prompt: int = 1, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + max_sequence_length: int = 1024, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ): + device = device or self.device + + prompt = [prompt] if isinstance(prompt, str) else prompt + if prompt is not None: + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + if prompt_embeds is None: + prompt_embeds, prompt_attention_mask = self._get_gemma_prompt_embeds( + prompt=prompt, + num_videos_per_prompt=num_videos_per_prompt, + max_sequence_length=max_sequence_length, + device=device, + dtype=dtype, + ) + + if do_classifier_free_guidance and negative_prompt_embeds is None: + negative_prompt = negative_prompt or "" + negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt + + if prompt is not None and type(prompt) is not type(negative_prompt): + raise TypeError( + f"`negative_prompt` should be the same type as `prompt`, but got {type(negative_prompt)} !=" + f" {type(prompt)}." + ) + if isinstance(negative_prompt, list) and batch_size != len(negative_prompt): + raise ValueError( + f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" + f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" + " the batch size of `prompt`." + ) + + negative_prompt_embeds, negative_prompt_attention_mask = self._get_gemma_prompt_embeds( + prompt=negative_prompt, + num_videos_per_prompt=num_videos_per_prompt, + max_sequence_length=max_sequence_length, + device=device, + dtype=dtype, + ) + + return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask + + # ------------------------------------------------------------------ + # Latent utilities (shared with LTX2Pipeline) + # ------------------------------------------------------------------ + + @staticmethod + def _pack_latents(latents: torch.Tensor, patch_size: int = 1, patch_size_t: int = 1) -> torch.Tensor: + batch_size, num_channels, num_frames, height, width = latents.shape + post_patch_num_frames = num_frames // patch_size_t + post_patch_height = height // patch_size + post_patch_width = width // patch_size + latents = latents.reshape( + batch_size, + -1, + post_patch_num_frames, + patch_size_t, + post_patch_height, + patch_size, + post_patch_width, + patch_size, + ) + latents = latents.permute(0, 2, 4, 6, 1, 3, 5, 7).flatten(4, 7).flatten(1, 3) + return latents + + @staticmethod + def _unpack_latents( + latents: torch.Tensor, + num_frames: int, + height: int, + width: int, + patch_size: int = 1, + patch_size_t: int = 1, + ) -> torch.Tensor: + batch_size = latents.size(0) + latents = latents.reshape(batch_size, num_frames, height, width, -1, patch_size_t, patch_size, patch_size) + latents = latents.permute(0, 4, 1, 5, 2, 6, 3, 7).flatten(6, 7).flatten(4, 5).flatten(2, 3) + return latents + + @staticmethod + def _normalize_latents( + latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, scaling_factor: float = 1.0 + ) -> torch.Tensor: + latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) + latents_std = latents_std.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) + latents = (latents - latents_mean) * scaling_factor / latents_std + return latents + + @staticmethod + def _normalize_audio_latents(latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor): + latents_mean = latents_mean.to(latents.device, latents.dtype) + latents_std = latents_std.to(latents.device, latents.dtype) + return (latents - latents_mean) / latents_std + + @staticmethod + def _denormalize_latents( + latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, scaling_factor: float = 1.0 + ) -> torch.Tensor: + latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) + latents_std = latents_std.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) + latents = latents * latents_std / scaling_factor + latents_mean + return latents + + @staticmethod + def _denormalize_audio_latents(latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor): + latents_mean = latents_mean.to(latents.device, latents.dtype) + latents_std = latents_std.to(latents.device, latents.dtype) + return (latents * latents_std) + latents_mean + + @staticmethod + def _pack_audio_latents( + latents: torch.Tensor, patch_size: int | None = None, patch_size_t: int | None = None + ) -> torch.Tensor: + if patch_size is not None and patch_size_t is not None: + batch_size, num_channels, latent_length, latent_mel_bins = latents.shape + post_patch_latent_length = latent_length / patch_size_t + post_patch_mel_bins = latent_mel_bins / patch_size + latents = latents.reshape( + batch_size, -1, post_patch_latent_length, patch_size_t, post_patch_mel_bins, patch_size + ) + latents = latents.permute(0, 2, 4, 1, 3, 5).flatten(3, 5).flatten(1, 2) + else: + latents = latents.transpose(1, 2).flatten(2, 3) + return latents + + @staticmethod + def _unpack_audio_latents( + latents: torch.Tensor, + latent_length: int, + num_mel_bins: int, + patch_size: int | None = None, + patch_size_t: int | None = None, + ) -> torch.Tensor: + if patch_size is not None and patch_size_t is not None: + batch_size = latents.size(0) + latents = latents.reshape(batch_size, latent_length, num_mel_bins, -1, patch_size_t, patch_size) + latents = latents.permute(0, 3, 1, 4, 2, 5).flatten(4, 5).flatten(2, 3) + else: + latents = latents.unflatten(2, (-1, num_mel_bins)).transpose(1, 2) + return latents + + @staticmethod + def _unpad_audio_latents(latents: torch.Tensor, num_frames: int) -> torch.Tensor: + return latents[:, :num_frames] + + # ------------------------------------------------------------------ + # Latent preparation + # ------------------------------------------------------------------ + + def prepare_latents( + self, + batch_size: int = 1, + num_channels_latents: int = 128, + height: int = 512, + width: int = 768, + num_frames: int = 121, + noise_scale: float = 0.0, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | None = None, + latents: torch.Tensor | None = None, + ) -> torch.Tensor: + if latents is not None: + if latents.ndim == 5: + latents = self._normalize_latents( + latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor + ) + latents = self._pack_latents( + latents, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size + ) + if latents.ndim != 3: + raise ValueError(f"Provided `latents` has shape {latents.shape}, expected [batch, seq, features].") + noise = randn_tensor(latents.shape, generator=generator, device=latents.device, dtype=latents.dtype) + latents = noise_scale * noise + (1 - noise_scale) * latents + return latents.to(device=device, dtype=dtype) + + height = height // self.vae_spatial_compression_ratio + width = width // self.vae_spatial_compression_ratio + num_frames = (num_frames - 1) // self.vae_temporal_compression_ratio + 1 + shape = (batch_size, num_channels_latents, num_frames, height, width) + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + latents = self._pack_latents(latents, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size) + return latents + + def prepare_audio_latents( + self, + batch_size: int = 1, + num_channels_latents: int = 8, + audio_latent_length: int = 1, + num_mel_bins: int = 64, + noise_scale: float = 0.0, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, int, int]: + original_latent_length = audio_latent_length + padded_latent_length = original_latent_length + latent_mel_bins = num_mel_bins // self.audio_vae_mel_compression_ratio + + if latents is not None: + if latents.ndim == 4: + latents = self._pack_audio_latents(latents) + if latents.ndim != 3: + raise ValueError(f"Provided `latents` has shape {latents.shape}, expected [batch, seq, features].") + latents = self._normalize_audio_latents(latents, self.audio_vae.latents_mean, self.audio_vae.latents_std) + noise = randn_tensor(latents.shape, generator=generator, device=latents.device, dtype=latents.dtype) + latents = noise_scale * noise + (1 - noise_scale) * latents + return latents.to(device=device, dtype=dtype), original_latent_length, padded_latent_length + + shape = (batch_size, num_channels_latents, padded_latent_length, latent_mel_bins) + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + latents = self._pack_audio_latents(latents) + return latents, original_latent_length, padded_latent_length + + # ------------------------------------------------------------------ + # Properties + # ------------------------------------------------------------------ + + @property + def guidance_scale(self): + return self._guidance_scale + + @property + def do_classifier_free_guidance(self): + return self._guidance_scale is not None and self._guidance_scale > 1.0 + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def current_timestep(self): + return self._current_timestep + + @property + def interrupt(self): + return self._interrupt + + # ------------------------------------------------------------------ + # Input validation + # ------------------------------------------------------------------ + + def check_inputs( + self, + prompt, + height, + width, + prompt_embeds=None, + negative_prompt_embeds=None, + prompt_attention_mask=None, + negative_prompt_attention_mask=None, + ): + if height % 32 != 0 or width % 32 != 0: + raise ValueError(f"`height` and `width` must be divisible by 32 but are {height} and {width}.") + if prompt is not None and prompt_embeds is not None: + raise ValueError("Cannot forward both `prompt` and `prompt_embeds`.") + elif prompt is None and prompt_embeds is None: + raise ValueError("Provide either `prompt` or `prompt_embeds`.") + elif prompt is not None and not isinstance(prompt, (str, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + + if prompt_embeds is not None and prompt_attention_mask is None: + raise ValueError("Must provide `prompt_attention_mask` when specifying `prompt_embeds`.") + + if negative_prompt_embeds is not None and negative_prompt_attention_mask is None: + raise ValueError("Must provide `negative_prompt_attention_mask` when specifying `negative_prompt_embeds`.") + + if prompt_embeds is not None and negative_prompt_embeds is not None: + if prompt_embeds.shape != negative_prompt_embeds.shape: + raise ValueError( + "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but" + f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`" + f" {negative_prompt_embeds.shape}." + ) + if prompt_attention_mask.shape != negative_prompt_attention_mask.shape: + raise ValueError( + "`prompt_attention_mask` and `negative_prompt_attention_mask` must have the same shape when " + "passed directly, but got: `prompt_attention_mask` " + f"{prompt_attention_mask.shape} != `negative_prompt_attention_mask` " + f"{negative_prompt_attention_mask.shape}." + ) + + # ------------------------------------------------------------------ + # Cache context + # ------------------------------------------------------------------ + + def _transformer_cache_context(self, context_name: str): + cache_context = getattr(self.transformer, "cache_context", None) + if callable(cache_context): + return cache_context(context_name) + return nullcontext() + + # ------------------------------------------------------------------ + # Forward pass + # ------------------------------------------------------------------ + + @torch.no_grad() + def forward( + self, + req: OmniDiffusionRequest, + prompt: str | list[str] | None = None, + negative_prompt: str | list[str] | None = None, + height: int | None = None, + width: int | None = None, + num_frames: int | None = None, + frame_rate: float | None = None, + num_inference_steps: int | None = None, + sigmas: list[float] | None = None, + timesteps: list[int] | None = None, + guidance_scale: float = 4.0, + noise_scale: float = 0.0, + num_videos_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + latents: torch.Tensor | None = None, + audio_latents: torch.Tensor | None = None, + prompt_embeds: torch.Tensor | None = None, + negative_prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + negative_prompt_attention_mask: torch.Tensor | None = None, + decode_timestep: float | list[float] = 0.0, + decode_noise_scale: float | list[float] | None = None, + output_type: str = "np", + return_dict: bool = True, + attention_kwargs: dict[str, Any] | None = None, + max_sequence_length: int | None = None, + ) -> DiffusionOutput: + # ---- Extract from request ---- + prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt + if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts): + negative_prompt = None + elif req.prompts: + negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts] + + height = req.sampling_params.height or height or 512 + width = req.sampling_params.width or width or 768 + num_frames = req.sampling_params.num_frames or num_frames or 121 + frame_rate = req.sampling_params.resolved_frame_rate or frame_rate or 24.0 + num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps or 40 + # Enforce minimum of 2 timesteps for flow matching scheduler + if timesteps is None: + num_inference_steps = max(int(num_inference_steps), 2) + elif len(timesteps) < 2: + raise ValueError("`timesteps` must contain at least 2 values for FlowMatchEulerDiscreteScheduler.") + num_videos_per_prompt = ( + req.sampling_params.num_outputs_per_prompt + if req.sampling_params.num_outputs_per_prompt > 0 + else num_videos_per_prompt or 1 + ) + max_sequence_length = ( + req.sampling_params.max_sequence_length or max_sequence_length or self.tokenizer_max_length + ) + + if req.sampling_params.guidance_scale_provided: + guidance_scale = req.sampling_params.guidance_scale + + if generator is None: + generator = req.sampling_params.generator + if generator is None and req.sampling_params.seed is not None: + generator = torch.Generator(device=self.device).manual_seed(req.sampling_params.seed) + + latents = req.sampling_params.latents if req.sampling_params.latents is not None else latents + audio_latents = ( + req.sampling_params.audio_latents + if req.sampling_params.audio_latents is not None + else req.sampling_params.extra_args.get("audio_latents", audio_latents) + ) + + # Override with pre-computed embeddings if provided in request + req_prompt_embeds = [_get_prompt_field(p, "prompt_embeds") for p in req.prompts] + if any(p is not None for p in req_prompt_embeds): + prompt_embeds = torch.stack(req_prompt_embeds) + + req_negative_prompt_embeds = [_get_prompt_field(p, "negative_prompt_embeds") for p in req.prompts] + if any(p is not None for p in req_negative_prompt_embeds): + negative_prompt_embeds = torch.stack(req_negative_prompt_embeds) + + req_prompt_attention_masks = [ + _get_prompt_field(p, "prompt_attention_mask") or _get_prompt_field(p, "attention_mask") for p in req.prompts + ] + if any(m is not None for m in req_prompt_attention_masks): + prompt_attention_mask = torch.stack(req_prompt_attention_masks) + + req_negative_attention_masks = [ + _get_prompt_field(p, "negative_prompt_attention_mask") or _get_prompt_field(p, "negative_attention_mask") + for p in req.prompts + ] + if any(m is not None for m in req_negative_attention_masks): + negative_prompt_attention_mask = torch.stack(req_negative_attention_masks) + + if req.sampling_params.decode_timestep is not None: + decode_timestep = req.sampling_params.decode_timestep + if req.sampling_params.decode_noise_scale is not None: + decode_noise_scale = req.sampling_params.decode_noise_scale + if req.sampling_params.output_type is not None: + output_type = req.sampling_params.output_type + + self.check_inputs( + prompt=prompt, + height=height, + width=width, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + prompt_attention_mask=prompt_attention_mask, + negative_prompt_attention_mask=negative_prompt_attention_mask, + ) + + self._guidance_scale = guidance_scale + self._attention_kwargs = attention_kwargs + self._interrupt = False + self._current_timestep = None + + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + + device = self.device + + # ---- Encode prompts ---- + ( + prompt_embeds, + prompt_attention_mask, + negative_prompt_embeds, + negative_prompt_attention_mask, + ) = self.encode_prompt( + prompt=prompt, + negative_prompt=negative_prompt, + do_classifier_free_guidance=self.do_classifier_free_guidance, + num_videos_per_prompt=num_videos_per_prompt, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + prompt_attention_mask=prompt_attention_mask, + negative_prompt_attention_mask=negative_prompt_attention_mask, + max_sequence_length=max_sequence_length, + device=device, + ) + + # ---- Connectors (LTX-2.3: padding_side API) ---- + # Concatenate negative + positive embeddings BEFORE connector call, + # matching diffusers which calls connectors once with batch=2. + # This ensures batch-dependent operations produce identical results. + if self.do_classifier_free_guidance: + prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) + prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0) + + self.connectors.to(device) + tokenizer_padding_side = getattr(self.tokenizer, "padding_side", "left") + connector_prompt_embeds, connector_audio_prompt_embeds, connector_attention_mask = self.connectors( + prompt_embeds, prompt_attention_mask, padding_side=tokenizer_padding_side + ) + self.connectors.to("cpu") + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + # ---- Prepare latents ---- + latent_num_frames = (num_frames - 1) // self.vae_temporal_compression_ratio + 1 + latent_height = height // self.vae_spatial_compression_ratio + latent_width = width // self.vae_spatial_compression_ratio + if latents is not None and latents.ndim == 5: + _, _, latent_num_frames, latent_height, latent_width = latents.shape + + num_channels_latents = self.transformer.config.in_channels + latents = self.prepare_latents( + batch_size * num_videos_per_prompt, + num_channels_latents, + height, + width, + num_frames, + noise_scale, + torch.float32, + device, + generator, + latents, + ) + + duration_s = num_frames / frame_rate + audio_latents_per_second = ( + self.audio_sampling_rate / self.audio_hop_length / float(self.audio_vae_temporal_compression_ratio) + ) + audio_num_frames = round(duration_s * audio_latents_per_second) + if audio_latents is not None and audio_latents.ndim == 4: + _, _, audio_num_frames, _ = audio_latents.shape + + num_mel_bins = self.audio_vae.config.mel_bins if self.audio_vae is not None else 64 + latent_mel_bins = num_mel_bins // self.audio_vae_mel_compression_ratio + num_channels_latents_audio = self.audio_vae.config.latent_channels if self.audio_vae is not None else 8 + audio_latents, original_audio_num_frames, padded_audio_num_frames = self.prepare_audio_latents( + batch_size * num_videos_per_prompt, + num_channels_latents=num_channels_latents_audio, + audio_latent_length=audio_num_frames, + num_mel_bins=num_mel_bins, + noise_scale=noise_scale, + dtype=torch.float32, + device=device, + generator=generator, + latents=audio_latents, + ) + + # ---- Scheduler setup ---- + sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas + # Use max_image_seq_len (not actual video_sequence_length) for mu calculation, + # matching diffusers' LTX2Pipeline which hardcodes this value. + mu = calculate_shift( + self.scheduler.config.get("max_image_seq_len", 4096), + self.scheduler.config.get("base_image_seq_len", 1024), + self.scheduler.config.get("max_image_seq_len", 4096), + self.scheduler.config.get("base_shift", 0.95), + self.scheduler.config.get("max_shift", 2.05), + ) + audio_scheduler = copy.deepcopy(self.scheduler) + _ = retrieve_timesteps(audio_scheduler, num_inference_steps, device, timesteps, sigmas=sigmas, mu=mu) + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, + num_inference_steps, + device, + timesteps, + sigmas=sigmas, + mu=mu, + ) + self._num_timesteps = len(timesteps) + + # ---- RoPE coordinates ---- + video_coords = self.transformer.rope.prepare_video_coords( + latents.shape[0], + latent_num_frames, + latent_height, + latent_width, + latents.device, + fps=frame_rate, + ) + audio_coords = self.transformer.audio_rope.prepare_audio_coords( + audio_latents.shape[0], + padded_audio_num_frames, + audio_latents.device, + ) + + # ---- CFG: duplicate coords for batch=2 ---- + # Connector outputs are already batch=2 (neg+pos concatenated before connector call) + if self.do_classifier_free_guidance: + video_coords = video_coords.repeat((2,) + (1,) * (video_coords.ndim - 1)) + audio_coords = audio_coords.repeat((2,) + (1,) * (audio_coords.ndim - 1)) + + # ---- Denoising loop ---- + # Uses x0-space CFG (delta formulation) matching diffusers' LTX2Pipeline. + # The velocity predictions are converted to x0, guidance is applied in x0 + # space, then converted back to velocity for the scheduler step. + with self.progress_bar(total=len(timesteps)) as pbar: + for i, t in enumerate(timesteps): + if self.interrupt: + continue + + self._current_timestep = t + + # Duplicate latents for CFG (uncond + cond) + latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents + latent_model_input = latent_model_input.to(connector_prompt_embeds.dtype) + audio_latent_model_input = ( + torch.cat([audio_latents] * 2) if self.do_classifier_free_guidance else audio_latents + ) + audio_latent_model_input = audio_latent_model_input.to(connector_prompt_embeds.dtype) + ts = t.expand(latent_model_input.shape[0]) + + with self._transformer_cache_context("cond_uncond"): + noise_pred_video, noise_pred_audio = self.transformer( + hidden_states=latent_model_input, + audio_hidden_states=audio_latent_model_input, + encoder_hidden_states=connector_prompt_embeds, + audio_encoder_hidden_states=connector_audio_prompt_embeds, + timestep=ts, + sigma=ts, # LTX-2.3: sigma for prompt_adaln + encoder_attention_mask=connector_attention_mask, + audio_encoder_attention_mask=connector_attention_mask, + num_frames=latent_num_frames, + height=latent_height, + width=latent_width, + fps=frame_rate, + audio_num_frames=padded_audio_num_frames, + video_coords=video_coords, + audio_coords=audio_coords, + attention_kwargs=attention_kwargs, + return_dict=False, + ) + + noise_pred_video = noise_pred_video.float() + noise_pred_audio = noise_pred_audio.float() + + # CFG in x0-space (delta formulation matching diffusers) + if self.do_classifier_free_guidance: + noise_pred_video_uncond, noise_pred_video_cond = noise_pred_video.chunk(2) + # Convert velocity to x0: x0 = sample - velocity * sigma + x0_video_cond = latents - noise_pred_video_cond * self.scheduler.sigmas[i] + x0_video_uncond = latents - noise_pred_video_uncond * self.scheduler.sigmas[i] + video_cfg_delta = (guidance_scale - 1) * (x0_video_cond - x0_video_uncond) + x0_video_guided = x0_video_cond + video_cfg_delta + + noise_pred_audio_uncond, noise_pred_audio_cond = noise_pred_audio.chunk(2) + x0_audio_cond = audio_latents - noise_pred_audio_cond * audio_scheduler.sigmas[i] + x0_audio_uncond = audio_latents - noise_pred_audio_uncond * audio_scheduler.sigmas[i] + audio_cfg_delta = (guidance_scale - 1) * (x0_audio_cond - x0_audio_uncond) + x0_audio_guided = x0_audio_cond + audio_cfg_delta + + # Convert x0 back to velocity: v = (sample - x0) / sigma + noise_pred_video = (latents - x0_video_guided) / self.scheduler.sigmas[i] + noise_pred_audio = (audio_latents - x0_audio_guided) / audio_scheduler.sigmas[i] + + latents = self.scheduler.step(noise_pred_video, t, latents, return_dict=False)[0] + audio_latents = audio_scheduler.step(noise_pred_audio, t, audio_latents, return_dict=False)[0] + + pbar.update() + + # ---- Unpack and denormalize ---- + latents = self._unpack_latents( + latents, + latent_num_frames, + latent_height, + latent_width, + self.transformer_spatial_patch_size, + self.transformer_temporal_patch_size, + ) + latents = self._denormalize_latents( + latents, + self.vae.latents_mean, + self.vae.latents_std, + self.vae.config.scaling_factor, + ) + + audio_latents = self._unpad_audio_latents(audio_latents, original_audio_num_frames) + audio_latents = self._denormalize_audio_latents( + audio_latents, + self.audio_vae.latents_mean, + self.audio_vae.latents_std, + ) + audio_latents = self._unpack_audio_latents( + audio_latents, + original_audio_num_frames, + num_mel_bins=latent_mel_bins, + ) + + # ---- Decode ---- + if output_type == "latent": + video = latents + audio = audio_latents + else: + latents = latents.to(connector_prompt_embeds.dtype) + + if not self.vae.config.timestep_conditioning: + timestep_decode = None + else: + noise = randn_tensor(latents.shape, generator=generator, device=device, dtype=latents.dtype) + if not isinstance(decode_timestep, list): + decode_timestep = [decode_timestep] * batch_size + if decode_noise_scale is None: + decode_noise_scale = decode_timestep + elif not isinstance(decode_noise_scale, list): + decode_noise_scale = [decode_noise_scale] * batch_size + timestep_decode = torch.tensor(decode_timestep, device=device, dtype=latents.dtype) + decode_noise_scale_t = torch.tensor(decode_noise_scale, device=device, dtype=latents.dtype)[ + :, None, None, None, None + ] + latents = (1 - decode_noise_scale_t) * latents + decode_noise_scale_t * noise + + # Move VAE, audio_vae, vocoder to GPU for decoding + self.vae.to(device) + latents = latents.to(self.vae.dtype) + video = self.vae.decode(latents, timestep_decode, return_dict=False)[0] + video = self.video_processor.postprocess_video(video, output_type=output_type) + self.vae.to("cpu") + + self.audio_vae.to(device) + audio_latents = audio_latents.to(self.audio_vae.dtype) + generated_mel_spectrograms = self.audio_vae.decode(audio_latents, return_dict=False)[0] + self.audio_vae.to("cpu") + + self.vocoder.to(device) + audio = self.vocoder(generated_mel_spectrograms) + self.vocoder.to("cpu") + torch.cuda.empty_cache() + + return DiffusionOutput(output=(video, audio)) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader(self) + return loader.load_weights(weights) + + +class LTX23ImageToVideoPipeline(nn.Module): + """LTX-2.3 image-to-video pipeline placeholder.""" + + def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): + super().__init__() + raise NotImplementedError( + "LTX23ImageToVideoPipeline is not yet implemented. " + "Use LTX23Pipeline for single-stage text-to-video generation." + ) diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3_image2video.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3_image2video.py new file mode 100644 index 000000000..d30ef5464 --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3_image2video.py @@ -0,0 +1,18 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Re-exports for LTX-2.3 I2V pipeline variants. + +The registry loads pipeline classes by (mod_folder, mod_relname, cls_name). +This module exposes the I2V class names so the registry can find them. +""" + +from .pipeline_ltx2_3 import ( + LTX23ImageToVideoPipeline, + get_ltx2_post_process_func, # noqa: F401 - loaded by registry via getattr +) + +__all__ = [ + "LTX23ImageToVideoPipeline", + "get_ltx2_post_process_func", +] diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py index 50a71a54b..4cc65f749 100644 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py +++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_image2video.py @@ -6,7 +6,7 @@ from __future__ import annotations import copy import os from collections.abc import Iterable -from typing import Any +from typing import Any, ClassVar import numpy as np import PIL.Image @@ -26,6 +26,7 @@ from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin +from vllm_omni.diffusion.models.interface import SupportsModuleOffload from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.lora.request import LoRARequest @@ -733,11 +734,15 @@ class LTX2ImageToVideoPipeline(LTX2Pipeline): return DiffusionOutput(output=(video, audio)) -class LTX2ImageToVideoTwoStagesPipeline(nn.Module): +class LTX2ImageToVideoTwoStagesPipeline(nn.Module, SupportsModuleOffload): """LTXImageToVideoTwoStagesPipeline is for two stages image to video generation""" support_image_input = True + _dit_modules: ClassVar[list[str]] = ["pipe.transformer"] + _encoder_modules: ClassVar[list[str]] = ["pipe.text_encoder"] + _vae_modules: ClassVar[list[str]] = ["pipe.vae", "pipe.audio_vae"] + def __init__( self, *, diff --git a/vllm_omni/diffusion/models/omnigen2/pipeline_omnigen2.py b/vllm_omni/diffusion/models/omnigen2/pipeline_omnigen2.py index 04720c932..46d634bfd 100644 --- a/vllm_omni/diffusion/models/omnigen2/pipeline_omnigen2.py +++ b/vllm_omni/diffusion/models/omnigen2/pipeline_omnigen2.py @@ -8,7 +8,7 @@ import os import warnings from collections.abc import Iterable from dataclasses import dataclass -from typing import Any +from typing import Any, ClassVar import numpy as np import PIL.Image @@ -32,6 +32,7 @@ from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader +from vllm_omni.diffusion.models.interface import SupportsModuleOffload from vllm_omni.diffusion.models.omnigen2.omnigen2_transformer import ( OmniGen2RotaryPosEmbed, OmniGen2Transformer2DModel, @@ -620,7 +621,7 @@ def retrieve_timesteps( return timesteps, num_inference_steps -class OmniGen2Pipeline(CFGParallelMixin, nn.Module): +class OmniGen2Pipeline(CFGParallelMixin, nn.Module, SupportsModuleOffload): """ Pipeline for text-to-image generation using OmniGen2. @@ -634,6 +635,10 @@ class OmniGen2Pipeline(CFGParallelMixin, nn.Module): od_config (OmniDiffusionConfig): The OmniDiffusion configuration. """ + _dit_modules: ClassVar[list[str]] = ["transformer"] + _encoder_modules: ClassVar[list[str]] = ["mllm"] + _vae_modules: ClassVar[list[str]] = ["vae"] + def __init__( self, *, diff --git a/vllm_omni/diffusion/models/t5_encoder/t5_encoder.py b/vllm_omni/diffusion/models/t5_encoder/t5_encoder.py index 7b0e842d0..0b81a71f1 100644 --- a/vllm_omni/diffusion/models/t5_encoder/t5_encoder.py +++ b/vllm_omni/diffusion/models/t5_encoder/t5_encoder.py @@ -328,6 +328,7 @@ class T5EncoderModel(nn.Module): def __init__(self, config: T5Config, prefix: str = ""): super().__init__() self.config = config + self.prefix = prefix self.shared = VocabParallelEmbedding(config.vocab_size, config.d_model) self.encoder = T5Stack(config, self.shared, prefix=f"{prefix}.encoder") @@ -359,29 +360,62 @@ class T5EncoderModel(nn.Module): ("wi", "wi_1", 1), ] + model_prefix = self.prefix + params_dict = dict(self.named_parameters()) loaded_params: set[str] = set() for name, loaded_weight in weights: original_name = name + + if model_prefix and name.startswith(model_prefix + "."): + name = name[len(model_prefix) + 1 :] + lookup_name = name + matched = False for param_name, weight_name, shard_id in stacked_params_mapping: if f".{weight_name}." not in name: continue - lookup_name = name.replace(weight_name, param_name) - if lookup_name not in params_dict: - continue - param = params_dict[lookup_name] - weight_loader = param.weight_loader - weight_loader(param, loaded_weight, shard_id) - break - else: - if name not in params_dict: - continue - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) + lookup_name = name.replace(f".{weight_name}.", f".{param_name}.", 1) + if lookup_name in params_dict: + param = params_dict[lookup_name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + matched = True + break + + if not matched: + if name in params_dict: + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + matched = True + + is_embed = "encoder.embed_tokens" in lookup_name + is_shared = lookup_name.startswith("shared.") or ".shared." in lookup_name + target_name = None + + if is_embed or is_shared: + if is_embed: + target_name = lookup_name.replace("encoder.embed_tokens", "shared") + else: + target_name = lookup_name.replace("shared.", "encoder.embed_tokens.", 1) + + if not matched and target_name in params_dict: + weight_loader = getattr(params_dict[target_name], "weight_loader", default_weight_loader) + weight_loader(params_dict[target_name], loaded_weight) + loaded_params.add(target_name) + matched = True + + if not matched: + continue + + if target_name is not None and target_name in params_dict and target_name not in loaded_params: + if target_name != lookup_name: + weight_loader = getattr(params_dict[target_name], "weight_loader", default_weight_loader) + weight_loader(params_dict[target_name], loaded_weight) + loaded_params.add(target_name) loaded_params.add(original_name) loaded_params.add(lookup_name) diff --git a/vllm_omni/diffusion/models/wan2_2/__init__.py b/vllm_omni/diffusion/models/wan2_2/__init__.py index 97808df29..ecbfdd472 100644 --- a/vllm_omni/diffusion/models/wan2_2/__init__.py +++ b/vllm_omni/diffusion/models/wan2_2/__init__.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from .patch_diffusers import patch_wan_rms_norm from .pipeline_wan2_2 import ( Wan22Pipeline, WanT2VDMD2Pipeline, @@ -51,3 +52,5 @@ __all__ = [ "VaceWanTransformerBlock", "WanVACETransformer3DModel", ] + +patch_wan_rms_norm() diff --git a/vllm_omni/diffusion/models/wan2_2/patch_diffusers.py b/vllm_omni/diffusion/models/wan2_2/patch_diffusers.py new file mode 100644 index 000000000..9b590fc5e --- /dev/null +++ b/vllm_omni/diffusion/models/wan2_2/patch_diffusers.py @@ -0,0 +1,14 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import sys + +from vllm_omni.diffusion.layers.norm import RMSNormVAE + + +def patch_wan_rms_norm(): + """Patch diffusers Wan RMSNorm implementation with RMSNormVAE.""" + + for module_name, module in sys.modules.items(): + if hasattr(module, "WanRMS_norm"): + setattr(module, "WanRMS_norm", RMSNormVAE) diff --git a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py index d4d81b78e..98e10b408 100644 --- a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py +++ b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py @@ -30,40 +30,12 @@ from vllm_omni.diffusion.distributed.sp_plan import ( from vllm_omni.diffusion.forward_context import get_forward_context from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNorm from vllm_omni.diffusion.layers.norm import LayerNorm, RMSNorm +from vllm_omni.diffusion.layers.rope import RotaryEmbeddingWan from vllm_omni.platforms import current_omni_platform logger = init_logger(__name__) -def apply_rotary_emb_wan( - hidden_states: torch.Tensor, - freqs_cos: torch.Tensor, - freqs_sin: torch.Tensor, -) -> torch.Tensor: - """ - Apply rotary embeddings to input tensors using the given frequency tensors. - - Args: - hidden_states: Input tensor of shape [B, S, H, D] - freqs_cos: Cosine frequencies - freqs_sin: Sine frequencies - - Returns: - Tensor with rotary embeddings applied - """ - x1, x2 = hidden_states.unflatten(-1, (-1, 2)).unbind(-1) - cos = freqs_cos[..., 0::2] - sin = freqs_sin[..., 1::2] - rotated = torch.stack( - ( - x1 * cos - x2 * sin, - x1 * sin + x2 * cos, - ), - dim=-1, - ) - return rotated.flatten(-2, -1).to(hidden_states.dtype) - - class DistributedRMSNorm(nn.Module): """ RMSNorm that computes global RMS across tensor parallel ranks. @@ -427,9 +399,10 @@ class WanSelfAttention(nn.Module): # Apply rotary embeddings if rotary_emb is not None: + self.rotary_embedding = RotaryEmbeddingWan(is_neox_style=False, half_head_dim=True) freqs_cos, freqs_sin = rotary_emb - query = apply_rotary_emb_wan(query, freqs_cos, freqs_sin) - key = apply_rotary_emb_wan(key, freqs_cos, freqs_sin) + query = self.rotary_embedding(query, freqs_cos, freqs_sin) + key = self.rotary_embedding(key, freqs_cos, freqs_sin) # Create attention metadata if mask is provided attn_metadata = None @@ -871,6 +844,10 @@ class WanTransformer3DModel(nn.Module): self.timestep_proj_prepare = TimestepProjPrepare() self.output_scale_shift_prepare = OutputScaleShiftPrepare(inner_dim) + # ROPE helper + self._cached_rope_emb = None + self._cached_rope_resolution = None + @property def dtype(self) -> torch.dtype: """Return the dtype of the model parameters.""" @@ -892,7 +869,14 @@ class WanTransformer3DModel(nn.Module): post_patch_width = width // p_w # Compute RoPE embeddings (sharded by _sp_plan via split_output=True) - rotary_emb = self.rope(hidden_states) + current_rope_resolution = (post_patch_num_frames, post_patch_height, post_patch_width) + if self._cached_rope_resolution == current_rope_resolution and self._cached_rope_emb is not None: + rotary_emb = self._cached_rope_emb + else: + freqs_cos, freqs_sin = self.rope(hidden_states) + rotary_emb = (freqs_cos[..., 0::2].to(hidden_states.dtype), freqs_sin[..., 1::2].to(hidden_states.dtype)) + self._hidden_states_shape = hidden_states.shape + self._cached_rope_emb = rotary_emb # Patch embedding and flatten to sequence # (hidden_states is sharded at blocks.0 input by _sp_plan) diff --git a/vllm_omni/diffusion/models/wan2_2/wan2_2_vace_transformer.py b/vllm_omni/diffusion/models/wan2_2/wan2_2_vace_transformer.py index c48938e1b..5060f1904 100644 --- a/vllm_omni/diffusion/models/wan2_2/wan2_2_vace_transformer.py +++ b/vllm_omni/diffusion/models/wan2_2/wan2_2_vace_transformer.py @@ -123,6 +123,10 @@ class WanVACETransformer3DModel(WanTransformer3DModel): ] ) + # ROPE helper + self._cached_rope_emb = None + self._cached_rope_resolution = None + def embed_vace_context( self, vace_context: torch.Tensor, @@ -165,7 +169,14 @@ class WanVACETransformer3DModel(WanTransformer3DModel): post_patch_width = width // p_w # Compute RoPE embeddings (sharded by _sp_plan via split_output=True) - rotary_emb = self.rope(hidden_states) + current_rope_resolution = (post_patch_num_frames, post_patch_height, post_patch_width) + if self._cached_rope_resolution == current_rope_resolution and self._cached_rope_emb is not None: + rotary_emb = self._cached_rope_emb + else: + freqs_cos, freqs_sin = self.rope(hidden_states) + rotary_emb = (freqs_cos[..., 0::2].to(hidden_states.dtype), freqs_sin[..., 1::2].to(hidden_states.dtype)) + self._hidden_states_shape = hidden_states.shape + self._cached_rope_emb = rotary_emb # Patch embedding and flatten to sequence hidden_states = self.patch_embedding(hidden_states) diff --git a/vllm_omni/diffusion/offloader/layerwise_backend.py b/vllm_omni/diffusion/offloader/layerwise_backend.py index d1216e2f2..9979d01b1 100644 --- a/vllm_omni/diffusion/offloader/layerwise_backend.py +++ b/vllm_omni/diffusion/offloader/layerwise_backend.py @@ -298,13 +298,20 @@ class LayerWiseOffloadBackend(OffloadBackend): for enc in modules.encoders: enc.to(self.device) - # Move VAE to GPU if available - if modules.vae is not None: + # Move VAE(s) to GPU if available + for vae in modules.vaes: try: - modules.vae.to(self.device, non_blocking=True) + vae.to(self.device, non_blocking=True) except Exception as exc: logger.debug("Failed to move VAE to GPU: %s", exc) + # Move resident modules to GPU (small modules needed every forward) + for name, module in zip(modules.resident_names, modules.resident_modules): + try: + module.to(self.device) + except Exception as exc: + logger.debug("Failed to move resident module %s to GPU: %s", name, exc) + logger.info("Applying layer-wise offloading on %s", modules.dit_names) # Apply block-wise offloading hook for each of the blocks in DiT model(s) diff --git a/vllm_omni/diffusion/offloader/module_collector.py b/vllm_omni/diffusion/offloader/module_collector.py index a09a33700..dfd81e98b 100644 --- a/vllm_omni/diffusion/offloader/module_collector.py +++ b/vllm_omni/diffusion/offloader/module_collector.py @@ -1,11 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from dataclasses import dataclass +from dataclasses import dataclass, field +from operator import attrgetter from torch import nn from vllm.logger import init_logger +from vllm_omni.diffusion.models.interface import SupportsModuleOffload + logger = init_logger(__name__) @@ -15,15 +18,81 @@ class PipelineModules: dit_names: list[str] encoders: list[nn.Module] encoder_names: list[str] - vae: nn.Module | None = None + vaes: list[nn.Module] + resident_modules: list[nn.Module] = field(default_factory=list) + resident_names: list[str] = field(default_factory=list) class ModuleDiscovery: - """Discovers pipeline components for offloading""" + """Discovers pipeline components for offloading. - DIT_ATTRS = ["transformer", "transformer_2", "dit", "sr_dit", "language_model", "transformer_blocks", "model"] - ENCODER_ATTRS = ["text_encoder", "text_encoder_2", "text_encoder_3", "image_encoder"] - VAE_ATTRS = ["vae", "audio_vae"] + If the pipeline implements :class:`SupportsModuleOffload`, + its ``_dit_modules``, ``_encoder_modules``, and ``_vae_modules`` + class variables are used directly. Otherwise, falls back to + scanning well-known attribute names. + """ + + # Fallback attribute names for pipelines that do not implement + # SupportsModuleOffload. + _FALLBACK_DIT_ATTRS = [ + "transformer", + "transformer_2", + "dit", + "sr_dit", + "language_model", + "transformer_blocks", + "model", + ] + _FALLBACK_ENCODER_ATTRS = [ + "text_encoder", + "text_encoder_2", + "text_encoder_3", + "image_encoder", + ] + _FALLBACK_VAE_ATTRS = [ + "vae", + "audio_vae", + ] + + @staticmethod + def _collect_modules( + pipeline: nn.Module, + attr_names: list[str], + *, + warn_missing: bool = False, + ) -> tuple[list[nn.Module], list[str]]: + """Resolve attribute names to (module, name) pairs, skipping missing. + + Supports dotted paths via :func:`operator.attrgetter`. + Warns on missing attributes when *warn_missing* is True. + """ + modules: list[nn.Module] = [] + names: list[str] = [] + seen: set[int] = set() + for attr in attr_names: + try: + module = attrgetter(attr)(pipeline) + except AttributeError: + module = None + if module is None: + if warn_missing: + logger.warning( + "Pipeline declares '%s' as offloadable but the attribute does not exist or is None", + attr, + ) + continue + if not isinstance(module, nn.Module): + logger.warning( + "Expected '%s' to be nn.Module, got %r", + attr, + type(module), + ) + continue + if id(module) not in seen: + seen.add(id(module)) + modules.append(module) + names.append(attr) + return modules, names @staticmethod def discover(pipeline: nn.Module) -> PipelineModules: @@ -35,46 +104,29 @@ class ModuleDiscovery: Returns: PipelineModules with lists of discovered modules and names """ - # Collect DiT/transformer modules - dit_modules: list[nn.Module] = [] - dit_names: list[str] = [] - for attr in ModuleDiscovery.DIT_ATTRS: - if not hasattr(pipeline, attr): - continue - module_obj = getattr(pipeline, attr) - if module_obj is None: - continue + declared = isinstance(pipeline, SupportsModuleOffload) + if declared: + dit_attrs = pipeline._dit_modules + enc_attrs = pipeline._encoder_modules + vae_attrs = pipeline._vae_modules + res_attrs = pipeline._resident_modules + else: + dit_attrs = ModuleDiscovery._FALLBACK_DIT_ATTRS + enc_attrs = ModuleDiscovery._FALLBACK_ENCODER_ATTRS + vae_attrs = ModuleDiscovery._FALLBACK_VAE_ATTRS + res_attrs = [] - if not isinstance(module_obj, nn.Module): - logger.warning(f"Expected {attr} to be nn.Module, got {type(module_obj)!r}") - continue - - if module_obj in dit_modules: - continue - - dit_modules.append(module_obj) - dit_names.append(attr) - - # Collect all encoders - encoders: list[nn.Module] = [] - encoder_names: list[str] = [] - for attr in ModuleDiscovery.ENCODER_ATTRS: - if hasattr(pipeline, attr) and getattr(pipeline, attr) is not None: - encoders.append(getattr(pipeline, attr)) - encoder_names.append(attr) - - # Collect VAE - vae = None - for attr in ModuleDiscovery.VAE_ATTRS: - module = getattr(pipeline, attr, None) - if module is not None: - vae = module - break + dit_modules, dit_names = ModuleDiscovery._collect_modules(pipeline, dit_attrs, warn_missing=declared) + encoders, encoder_names = ModuleDiscovery._collect_modules(pipeline, enc_attrs, warn_missing=declared) + vaes, _ = ModuleDiscovery._collect_modules(pipeline, vae_attrs, warn_missing=declared) + residents, resident_names = ModuleDiscovery._collect_modules(pipeline, res_attrs, warn_missing=declared) return PipelineModules( dits=dit_modules, dit_names=dit_names, encoders=encoders, encoder_names=encoder_names, - vae=vae, + vaes=vaes, + resident_modules=residents, + resident_names=resident_names, ) diff --git a/vllm_omni/diffusion/offloader/sequential_backend.py b/vllm_omni/diffusion/offloader/sequential_backend.py index 46f48e99c..06454ad5c 100644 --- a/vllm_omni/diffusion/offloader/sequential_backend.py +++ b/vllm_omni/diffusion/offloader/sequential_backend.py @@ -210,10 +210,10 @@ class ModelLevelOffloadBackend(OffloadBackend): for enc in modules.encoders: enc.to(self.device) - # Move VAE to GPU if available - if modules.vae is not None: + # Move VAE(s) to GPU if available + for vae in modules.vaes: try: - modules.vae.to(self.device, non_blocking=True) + vae.to(self.device, non_blocking=True) except Exception as exc: logger.debug("Failed to move VAE to GPU: %s", exc) diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 4001109cc..ed7a03c2b 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -93,6 +93,16 @@ _DIFFUSION_MODELS = { "pipeline_ltx2_image2video", "LTX2I2VDMD2Pipeline", ), + "LTX23Pipeline": ( + "ltx2", + "pipeline_ltx2_3", + "LTX23Pipeline", + ), + "LTX23ImageToVideoPipeline": ( + "ltx2", + "pipeline_ltx2_3_image2video", + "LTX23ImageToVideoPipeline", + ), "StableAudioPipeline": ( "stable_audio", "pipeline_stable_audio", @@ -208,6 +218,11 @@ _DIFFUSION_MODELS = { "pipeline_omnivoice", "OmniVoicePipeline", ), + "DiffusersAdapterPipeline": ( + "diffusers_adapter", + "pipeline_diffusers_adapter", + "DiffusersAdapterPipeline", + ), } @@ -379,6 +394,8 @@ _DIFFUSION_POST_PROCESS_FUNCS = { "LTX2ImageToVideoTwoStagesPipeline": "get_ltx2_post_process_func", "LTX2T2VDMD2Pipeline": "get_ltx2_post_process_func", "LTX2I2VDMD2Pipeline": "get_ltx2_post_process_func", + "LTX23Pipeline": "get_ltx2_post_process_func", + "LTX23ImageToVideoPipeline": "get_ltx2_post_process_func", "StableAudioPipeline": "get_stable_audio_post_process_func", "WanImageToVideoPipeline": "get_wan22_i2v_post_process_func", "WanT2VDMD2Pipeline": "get_wan22_post_process_func", diff --git a/vllm_omni/distributed/omni_coordinator/omni_coordinator.py b/vllm_omni/distributed/omni_coordinator/omni_coordinator.py index 7ff608f2f..2c7c8fbb9 100644 --- a/vllm_omni/distributed/omni_coordinator/omni_coordinator.py +++ b/vllm_omni/distributed/omni_coordinator/omni_coordinator.py @@ -61,6 +61,7 @@ class OmniCoordinator: self._publish_min_interval: float = 0.1 # seconds self._pending_broadcast: bool = False + self._pending_lock = threading.Lock() self._running = True self._closed = False @@ -84,13 +85,13 @@ class OmniCoordinator: """Add a new instance based on an incoming event.""" with self._lock: self._add_new_instance_locked(event) - self.publish_instance_list_update() + self._schedule_broadcast() def update_instance_info(self, event: InstanceEvent) -> None: """Update an existing instance based on an incoming event.""" with self._lock: self._update_instance_info_locked(event) - self.publish_instance_list_update() + self._schedule_broadcast() def remove_instance(self, event: InstanceEvent) -> None: """Mark an instance as removed / down based on an incoming event. @@ -101,10 +102,15 @@ class OmniCoordinator: """ with self._lock: self._remove_instance_locked(event) - self.publish_instance_list_update() + self._schedule_broadcast() - def publish_instance_list_update(self) -> None: - """Publish the current active instance list to all subscribers.""" + def publish_instance_list_update(self) -> bool: + """Publish the current active instance list to all subscribers. + + Returns: + True if the PUB send succeeded, False if it was dropped (e.g. + socket not ready when using ``zmq.NOBLOCK``). + """ active_list = self.get_active_instances() payload = asdict(active_list) data = json.dumps(payload).encode("utf-8") @@ -113,20 +119,18 @@ class OmniCoordinator: try: # PUB socket is best-effort; drop update if not ready. self._pub.send(data, flags=zmq.NOBLOCK) + return True except (zmq.Again, zmq.ZMQError): # Silently ignore send failures; next update will catch up. - return + return False - def _schedule_broadcast(self, force: bool) -> None: - """Schedule a broadcast, optionally bypassing throttling. + def _schedule_broadcast(self) -> None: + """Request a broadcast to be flushed by the periodic loop. - When ``force`` is True, publish immediately. Otherwise, mark a pending - broadcast that will be flushed by the periodic loop at most once per - ``_publish_min_interval``. + All broadcast requests are coalesced via ``_pending_broadcast`` and + flushed at most once per ``_publish_min_interval``. """ - if force: - self.publish_instance_list_update() - else: + with self._pending_lock: self._pending_broadcast = True def _mark_instance_error_locked(self, info: InstanceInfo) -> None: @@ -152,8 +156,8 @@ class OmniCoordinator: for input_addr in to_delete: del self._instances[input_addr] if timed_out: - # Instance liveness changed; force immediate broadcast. - self._schedule_broadcast(force=True) + # Instance liveness changed; request broadcast. + self._schedule_broadcast() def close(self) -> None: """Shut down background threads and close all ZMQ sockets.""" @@ -228,9 +232,9 @@ class OmniCoordinator: def _periodic_loop(self) -> None: """Periodic loop to check heartbeat timeouts and flush broadcasts. - Heartbeat timeouts are checked on their original cadence, while - queue_length / non-liveness updates are coalesced and flushed at - most once per ``_publish_min_interval``. + Heartbeat timeouts are checked on their original cadence, while all + broadcast requests are coalesced and flushed at most once per + ``_publish_min_interval``. """ heartbeat_interval = max(1.0, min(self._heartbeat_timeout / 2.0, 5.0)) loop_interval = self._publish_min_interval @@ -243,9 +247,18 @@ class OmniCoordinator: self._check_heartbeat_timeouts() last_heartbeat_check = now - if self._pending_broadcast: - self.publish_instance_list_update() - self._pending_broadcast = False + with self._pending_lock: + has_pending_broadcast = self._pending_broadcast + + if not has_pending_broadcast: + if self._stop_event.wait(timeout=loop_interval): + break + continue + + # Publish outside lock. Clear pending only on success. + if self.publish_instance_list_update(): + with self._pending_lock: + self._pending_broadcast = False if self._stop_event.wait(timeout=loop_interval): break @@ -267,27 +280,23 @@ class OmniCoordinator: info.status = StageStatus.UP promote = True if promote: - self._schedule_broadcast(force=True) + self._schedule_broadcast() return # Check-and-act under single lock to avoid TOCTOU race (duplicate # registration when concurrent events arrive for the same instance). with self._lock: - force_broadcast = False if input_addr not in self._instances: self._add_new_instance_locked(event) - force_broadcast = True else: if event.status == StageStatus.DOWN: self._remove_instance_locked(event) - force_broadcast = True else: self._update_instance_info_locked(event) - # New instances / DOWN events are broadcast immediately; other - # updates (e.g. queue_length changes) are throttled via the - # periodic loop. - self._schedule_broadcast(force=force_broadcast) + # Any non-heartbeat state change that affects the active list + # is coalesced and flushed via the periodic loop. + self._schedule_broadcast() except (KeyError, ValueError, TypeError) as e: logger.warning("Dropping malformed event: %s", e) diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index c93bd32c2..a6fff97e7 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -170,6 +170,7 @@ class OmniEngineArgs(EngineArgs): output_modalities: list[str] | None = None log_stats: bool = False custom_pipeline_args: dict[str, Any] | None = None + has_sampling_extra_args: bool = False def __post_init__(self) -> None: if self.worker_cls is None: @@ -184,8 +185,26 @@ class OmniEngineArgs(EngineArgs): def from_cli_args(cls, args: argparse.Namespace) -> "OmniEngineArgs": attrs = [attr.name for attr in dataclasses.fields(cls)] engine_args = cls(**{attr: getattr(args, attr) for attr in attrs if hasattr(args, attr)}) + engine_args._explicit_fields = frozenset( + attr for attr in attrs if hasattr(args, attr) and getattr(args, attr) is not None + ) return engine_args + @classmethod + def create(cls, **explicit: Any) -> "OmniEngineArgs": + """Tracks caller-set fields for ``Omni(..., engine_args=ea)``.""" + ea = cls(**explicit) + ea._explicit_fields = frozenset(explicit.keys()) + return ea + + def explicit_kwargs(self) -> dict[str, Any]: + explicit = getattr(self, "_explicit_fields", None) + if explicit is None: + return { + f.name: getattr(self, f.name) for f in dataclasses.fields(self) if getattr(self, f.name) is not None + } + return {k: getattr(self, k) for k in explicit} + def _ensure_omni_models_registered(self): if hasattr(self, "_omni_models_registered"): return True @@ -246,6 +265,13 @@ class OmniEngineArgs(EngineArgs): if model_type is not None: self.hf_overrides.setdefault("model_type", model_type) + # Stage wrappers (e.g. Code2Wav) may need max_model_len larger + # than the base checkpoint's text max_position_embeddings. + if self.model_arch == "Qwen3TTSCode2Wav" and self.max_model_len is not None: + self.hf_overrides.setdefault("talker_config", {}).setdefault( + "max_position_embeddings", int(self.max_model_len) + ) + # For models whose HF config.json is empty or lacks model_type # (e.g. CosyVoice3), AutoConfig.from_pretrained fails because it # cannot determine which config class to use from the empty dict. @@ -319,6 +345,7 @@ class OmniEngineArgs(EngineArgs): subtalker_sampling_params=self.subtalker_sampling_params, omni_kv_config=self.omni_kv_config, task_type=self.task_type, + has_sampling_extra_args=self.has_sampling_extra_args, ) return omni_config @@ -433,12 +460,50 @@ SHARED_FIELDS: frozenset[str] = frozenset( } ) +_DEPLOY_ENGINE_ARG_OVERRIDE_FIELDS: frozenset[str] = frozenset( + { + # Capacity / scheduling. + "async_scheduling", + "max_model_len", + "max_num_batched_tokens", + "max_num_seqs", + # Memory / parallelism. + "data_parallel_size", + "gpu_memory_utilization", + "pipeline_parallel_size", + "tensor_parallel_size", + # Execution / loading. + "enforce_eager", + "distributed_executor_backend", + "dtype", + "quantization", + "trust_remote_code", + # Caching / chunking. + "async_chunk", + "enable_prefix_caching", + "enable_chunked_prefill", + # Model-specific engine extras. + "subtalker_sampling_params", + } +) + +_DEPLOY_RUNTIME_OVERRIDE_FIELDS: frozenset[str] = frozenset( + { + "devices", + } +) + def orchestrator_field_names() -> frozenset[str]: """Return the names of every field on OrchestratorArgs.""" return frozenset(f.name for f in fields(OrchestratorArgs)) +def deploy_override_field_names() -> frozenset[str]: + """Return kwargs whose parser defaults must not override deploy YAML.""" + return _DEPLOY_ENGINE_ARG_OVERRIDE_FIELDS | _DEPLOY_RUNTIME_OVERRIDE_FIELDS + + def internal_blacklist_keys() -> frozenset[str]: """Return the set of CLI keys that must never be forwarded as per-stage engine overrides. @@ -586,3 +651,23 @@ def orchestrator_args_from_argparse(args: Any) -> OrchestratorArgs: if value is not None or f.default is None: kwargs[f.name] = value return OrchestratorArgs(**kwargs) + + +def nullify_stage_engine_defaults(parser: argparse.ArgumentParser) -> None: + """Reset stage-level engine flag defaults to ``None``; preserve real + default in help text. Only deploy-YAML override fields are touched. + Idempotent.""" + override_dests = deploy_override_field_names() + + for action in parser._actions: + if action.dest in ("help", "version") or not action.option_strings: + continue + if action.dest not in override_dests: + continue + if action.default is None or action.default is argparse.SUPPRESS: + continue + if action.help and "(default:" not in action.help and "%(default)" not in action.help: + action.help = f"{action.help} (default: {action.default})" + action.default = None + + parser._omni_nullified = True # type: ignore[attr-defined] diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index a26102d05..6ddaf993e 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -105,7 +105,13 @@ _STARTUP_POLL_INTERVAL_S = 1.0 # Fields that must survive the "equal to default → strip" filter because # diffusion stages need them even when equal to vllm's default value # (e.g. colocate worker setup relies on worker_extension_cls being forwarded). -_PARENT_ARGS_KEEP: frozenset[str] = frozenset({"worker_extension_cls"}) +_PARENT_ARGS_KEEP: frozenset[str] = frozenset( + { + "worker_extension_cls", + "allowed_local_media_path", + "allowed_media_domains", + } +) # Omni orchestrator-level fields consumed by ``_resolve_stage_configs`` that # must never leak into per-stage EngineArgs (``stage_configs_path`` would @@ -273,13 +279,16 @@ class AsyncOmniEngine: logger.info(f"[AsyncOmniEngine] Initializing with model {model}") - # Merge typed engine_args fields into kwargs; explicit kwargs take priority. + # Merge tracked engine_args fields into kwargs; explicit kwargs take priority. if engine_args is not None: - ea_dict = { - f.name: getattr(engine_args, f.name) - for f in dataclasses.fields(engine_args) - if not f.name.startswith("_") - } + if not hasattr(engine_args, "_explicit_fields"): + raise TypeError( + "engine_args=OmniEngineArgs(...) is ambiguous under " + "sentinel-default precedence. Use " + "OmniEngineArgs.create(**explicit) or pass explicit kwargs " + "directly." + ) + ea_dict = engine_args.explicit_kwargs() # Remove model since it is passed as a positional arg already. ea_dict.pop("model", None) kwargs = {**ea_dict, **kwargs} @@ -1232,6 +1241,8 @@ class AsyncOmniEngine: if not isinstance(default_sampling_params, dict): default_sampling_params = None stage_default_sampling_params = default_sampling_params.get("0", {}) if default_sampling_params else {} + if normalized_kwargs.get("dtype") is None: + normalized_kwargs["dtype"] = "auto" # TODO: hack, convert dtype to string to avoid non-premitive omegaconf create error. if "dtype" in normalized_kwargs and not isinstance(normalized_kwargs["dtype"], str): @@ -1319,6 +1330,12 @@ class AsyncOmniEngine: if "dtype" in normalized_kwargs: stage_engine_args["dtype"] = normalized_kwargs["dtype"] + # New split fields for diffusers adapter kwargs. + if kwargs.get("diffusers_load_kwargs") is not None: + stage_engine_args["diffusers_load_kwargs"] = kwargs["diffusers_load_kwargs"] + if kwargs.get("diffusers_call_kwargs") is not None: + stage_engine_args["diffusers_call_kwargs"] = kwargs["diffusers_call_kwargs"] + default_stage_cfg = [ { "stage_id": 0, @@ -1376,10 +1393,7 @@ class AsyncOmniEngine: stage_configs_path = kwargs.get("stage_configs_path", None) deploy_config_path = kwargs.pop("deploy_config", None) stage_overrides_json = kwargs.pop("stage_overrides", None) - # Set of CLI keys the user actually typed; ``None`` means we have no - # parser-level info (e.g. programmatic Omni() call) and the lower - # layers should treat all kwargs as explicit. - cli_explicit_keys = kwargs.pop("_cli_explicit_keys", None) + kwargs.pop("_cli_explicit_keys", None) explicit_stage_configs = kwargs.pop("stage_configs", None) if explicit_stage_configs is not None: logger.warning( @@ -1412,7 +1426,6 @@ class AsyncOmniEngine: default_stage_cfg_factory=lambda: self._create_default_diffusion_stage_cfg(kwargs), deploy_config_path=deploy_config_path, stage_overrides=stage_overrides, - cli_explicit_keys=cli_explicit_keys, ) # Inject diffusion LoRA-related knobs from kwargs if not present in the stage config. @@ -1438,7 +1451,10 @@ class AsyncOmniEngine: if lora_scale is not None: if not hasattr(cfg.engine_args, "lora_scale") or cfg.engine_args.lora_scale is None: cfg.engine_args.lora_scale = lora_scale + # Prefer explicit quantization_config; fallback to legacy --quantization. quantization_config = kwargs.get("quantization_config") + if quantization_config is None: + quantization_config = kwargs.get("quantization") if quantization_config is not None: if ( not hasattr(cfg.engine_args, "quantization_config") diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index ca60c7531..a7edf4339 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -33,6 +33,7 @@ from vllm_omni.engine.serialization import serialize_additional_information from vllm_omni.metrics.stats import StageRequestStats as StageRequestMetrics from vllm_omni.metrics.stats import StageStats from vllm_omni.metrics.utils import count_tokens_from_outputs +from vllm_omni.outputs import OmniRequestOutput logger = init_logger(__name__) @@ -181,6 +182,7 @@ class Orchestrator: self._shutdown_event = asyncio.Event() self._stages_shutdown = False self._fatal_error: str | None = None + self._fatal_error_stage_id: int | None = None async def run(self) -> None: """Main entry point for the Orchestrator event loop.""" @@ -331,6 +333,7 @@ class Orchestrator: e, ) self._fatal_error = str(e) + self._fatal_error_stage_id = stage_id for req_id, req_state in list(self.request_states.items()): if stage_id in req_state.stage_submit_ts: await self.output_async_queue.put( @@ -339,6 +342,7 @@ class Orchestrator: "error": str(e), "fatal": True, "request_id": req_id, + "stage_id": stage_id, } ) self.request_states.pop(req_id, None) @@ -686,6 +690,32 @@ class Orchestrator: next_stage_id, ) if isinstance(diffusion_prompt, list): + if not diffusion_prompt: + error_output = OmniRequestOutput.from_error( + req_id, + f"Stage-{stage_id} produced no valid inputs for diffusion stage-{next_stage_id}", + ) + logger.warning( + "[Orchestrator] req=%s stage=%d produced empty diffusion inputs for stage=%d; " + "routing terminal error output", + req_id, + stage_id, + next_stage_id, + ) + await self.output_async_queue.put( + { + "type": "output", + "request_id": req_id, + "stage_id": next_stage_id, + "engine_outputs": error_output, + "metrics": None, + "finished": True, + } + ) + self._pd_kv_params.pop(req_id, None) + self._cfg_tracker.cleanup_parent(req_id) + self.request_states.pop(req_id, None) + return diffusion_prompt = diffusion_prompt[0] else: diffusion_prompt = req_state.prompt @@ -1135,6 +1165,7 @@ class Orchestrator: "error": self._fatal_error, "fatal": True, "request_id": req_id, + "stage_id": self._fatal_error_stage_id, } ) notified.add(req_id) @@ -1150,6 +1181,7 @@ class Orchestrator: "error": self._fatal_error, "fatal": True, "request_id": req_id, + "stage_id": self._fatal_error_stage_id, } ) self.request_states.pop(req_id, None) diff --git a/vllm_omni/engine/stage_engine_core_client.py b/vllm_omni/engine/stage_engine_core_client.py index f6683c1e0..844be7194 100644 --- a/vllm_omni/engine/stage_engine_core_client.py +++ b/vllm_omni/engine/stage_engine_core_client.py @@ -13,6 +13,7 @@ import weakref from typing import TYPE_CHECKING, Any from urllib.parse import urlparse +import psutil from vllm.logger import init_logger from vllm.v1.engine import EngineCoreRequest from vllm.v1.engine.core_client import AsyncMPClient, DPLBAsyncMPClient @@ -30,6 +31,8 @@ if TYPE_CHECKING: logger = init_logger(__name__) +SHUTDOWN_TIMEOUT_S = 5 + class StageEngineCoreClientBase: """Shared stage-aware behavior for async EngineCore clients. @@ -420,13 +423,37 @@ class StageEngineCoreClientBase: def shutdown(self) -> None: """Shutdown managed resources and any externally spawned subprocess.""" - super().shutdown() - if self._proc is not None and self._proc.is_alive(): - self._proc.terminate() - self._proc.join(timeout=5) - if self._proc.is_alive(): - self._proc.kill() - self._proc = None + child_procs: list[psutil.Process] = [] + if self._proc is not None and self._proc.pid is not None: + try: + child_procs = psutil.Process(self._proc.pid).children(recursive=True) + except psutil.Error: + child_procs = [] + + try: + super().shutdown() + finally: + if self._proc is not None and self._proc.is_alive(): + self._proc.terminate() + self._proc.join(timeout=SHUTDOWN_TIMEOUT_S) + if self._proc.is_alive(): + self._proc.kill() + self._proc.join(timeout=SHUTDOWN_TIMEOUT_S) + + alive_children = [proc for proc in child_procs if proc.is_running()] + for proc in alive_children: + try: + proc.terminate() + except psutil.Error: + pass + _, still_alive = psutil.wait_procs(alive_children, timeout=SHUTDOWN_TIMEOUT_S) + for proc in still_alive: + try: + proc.kill() + except psutil.Error: + pass + # The process handle is no longer reliable after best-effort cleanup. + self._proc = None class StageEngineCoreClient(StageEngineCoreClientBase, AsyncMPClient): diff --git a/vllm_omni/engine/stage_engine_core_proc.py b/vllm_omni/engine/stage_engine_core_proc.py index 689378a79..2ab8b37dd 100644 --- a/vllm_omni/engine/stage_engine_core_proc.py +++ b/vllm_omni/engine/stage_engine_core_proc.py @@ -23,10 +23,12 @@ from vllm.utils.system_utils import ( get_mp_context, set_process_title, ) -from vllm.v1.engine.core import EngineCoreProc +from vllm.v1.engine import EngineCoreRequestType +from vllm.v1.engine.core import EngineCoreProc, EngineShutdownState from vllm.v1.engine.utils import ( EngineHandshakeMetadata, EngineZmqAddresses, + SignalCallback, get_engine_zmq_addresses, ) from vllm.v1.utils import shutdown @@ -54,18 +56,9 @@ class StageEngineCoreProc(EngineCoreProc): **kwargs: Any, ) -> None: """Launch StageEngineCoreProc busy loop in background process.""" - shutdown_requested = False + signal_callback: SignalCallback | None = None maybe_register_config_serialize_by_value() - def signal_handler(signum: int, frame: Any) -> None: - nonlocal shutdown_requested - if not shutdown_requested: - shutdown_requested = True - raise SystemExit() - - signal.signal(signal.SIGTERM, signal_handler) - signal.signal(signal.SIGINT, signal_handler) - engine_core: StageEngineCoreProc | None = None try: vllm_config: VllmConfig = kwargs["vllm_config"] @@ -88,6 +81,19 @@ class StageEngineCoreProc(EngineCoreProc): engine_index=dp_rank, **kwargs, ) + + def wakeup_engine() -> None: + engine_core.input_queue.put_nowait((EngineCoreRequestType.WAKEUP, None)) + + signal_callback = SignalCallback(wakeup_engine) + + def signal_handler(signum: int, frame: Any) -> None: + engine_core.shutdown_state = EngineShutdownState.REQUESTED + signal_callback.trigger() + + signal.signal(signal.SIGTERM, signal_handler) + signal.signal(signal.SIGINT, signal_handler) + engine_core.run_busy_loop() except SystemExit: @@ -101,6 +107,10 @@ class StageEngineCoreProc(EngineCoreProc): engine_core._send_engine_dead() raise finally: + signal.signal(signal.SIGTERM, signal.SIG_DFL) + signal.signal(signal.SIGINT, signal.SIG_DFL) + if signal_callback is not None: + signal_callback.stop() if engine_core is not None: engine_core.shutdown() diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index cd6cd6a69..9fdcd6216 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -14,7 +14,7 @@ import multiprocessing as mp import os import time from collections.abc import Callable, Sequence -from dataclasses import dataclass +from dataclasses import dataclass, replace from typing import Any, Literal from vllm.logger import init_logger @@ -29,6 +29,7 @@ from vllm_omni.entrypoints.stage_utils import _to_dict, set_stage_devices from vllm_omni.entrypoints.utils import filter_dataclass_kwargs, resolve_model_config_path from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniSamplingParams from vllm_omni.platforms import current_omni_platform +from vllm_omni.quantization.inc_config import OmniINCConfig logger = init_logger(__name__) @@ -411,6 +412,10 @@ def build_engine_args_dict( if stage_type != "diffusion": resolve_worker_cls(engine_args_dict) + # Check whether the stage's default_sampling_params defines extra_args. + default_sp = _to_dict(getattr(stage_config, "default_sampling_params", {})) + engine_args_dict["has_sampling_extra_args"] = bool(default_sp.get("extra_args")) + return engine_args_dict @@ -455,6 +460,11 @@ def build_vllm_config( ) executor_class = Executor.get_class(vllm_config) + # Upgrade vanilla INCConfig to OmniINCConfig for multi-stage models. + upgraded = OmniINCConfig.maybe_upgrade(vllm_config.quant_config) + if upgraded is not vllm_config.quant_config: + vllm_config = replace(vllm_config, quant_config=upgraded) + return vllm_config, executor_class @@ -591,12 +601,22 @@ def release_device_locks(lock_fds: list[int]) -> None: def load_omni_transfer_config_for_model(model: str, config_path: str | None) -> Any: - """Load omni transfer config from an explicit path or resolved model config.""" + """Load omni transfer config from an explicit path or resolved model config. + + Resolves ``base_config`` inheritance (CI overlay → base deploy YAML) so + that connectors defined in the base config are visible to the transfer + config parser. + """ from vllm_omni.distributed.omni_connectors import load_omni_transfer_config try: resolved_config_path = config_path or resolve_model_config_path(model) - return load_omni_transfer_config(resolved_config_path) + if resolved_config_path is None: + return None + from vllm_omni.config.stage_config import resolve_deploy_yaml + + resolved_dict = resolve_deploy_yaml(resolved_config_path) + return load_omni_transfer_config(config_dict=resolved_dict) except Exception as e: logger.warning("[stage_init] Failed to load transfer config: %s", e) return None diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 67425ea2e..119232ce7 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -27,7 +27,10 @@ from vllm.v1.engine.exceptions import EngineDeadError from vllm_omni.diffusion.data import OmniACK, OmniSleepTask, OmniWakeTask from vllm_omni.entrypoints.client_request_state import ClientRequestState -from vllm_omni.entrypoints.omni_base import OmniBase +from vllm_omni.entrypoints.omni_base import ( + OmniBase, + OmniEngineDeadError, +) from vllm_omni.metrics.stats import OrchestratorAggregator as OrchestratorMetrics from vllm_omni.outputs import OmniRequestOutput from vllm_omni.platforms import current_omni_platform @@ -494,7 +497,10 @@ class AsyncOmni(EngineClient, OmniBase): stage_id = result.get("stage_id", 0) if result.get("type") == "error" and result.get("fatal"): - raise EngineDeadError(result.get("error", "")) + raise OmniEngineDeadError( + result.get("error", ""), + error_stage_id=result.get("stage_id"), + ) # Check for errors if "error" in result: @@ -580,6 +586,18 @@ class AsyncOmni(EngineClient, OmniBase): except asyncio.CancelledError: raise + except OmniEngineDeadError as e: + logger.error("[AsyncOmni] Engine dead: %s", e) + for req_state in list(self.request_states.values()): + error_msg = { + "type": "error", + "error": str(e), + "fatal": True, + "request_id": req_state.request_id, + } + if e.error_stage_id is not None: + error_msg["stage_id"] = e.error_stage_id + await req_state.queue.put(error_msg) except EngineDeadError as e: logger.error("[AsyncOmni] Engine dead: %s", e) for req_state in list(self.request_states.values()): @@ -884,7 +902,7 @@ class AsyncOmni(EngineClient, OmniBase): @property def dead_error(self) -> BaseException: """EngineClient abstract property implementation.""" - return EngineDeadError() + return OmniEngineDeadError() # ==================== EngineClient Interface ==================== diff --git a/vllm_omni/entrypoints/cli/benchmark/serve.py b/vllm_omni/entrypoints/cli/benchmark/serve.py index d281432e5..44d38303f 100644 --- a/vllm_omni/entrypoints/cli/benchmark/serve.py +++ b/vllm_omni/entrypoints/cli/benchmark/serve.py @@ -33,6 +33,8 @@ def add_daily_omni_cli_args(parser: argparse.ArgumentParser) -> None: default=None, help="Root directory of extracted Daily-Omni videos (contents of Videos.tar: " "each video_id in its own subdir with {video_id}_video.mp4). " + "If omitted, Videos.tar is downloaded from the Hugging Face dataset repo on first multimodal " + "request. " "When using file URLs, you MUST start the vLLM server with " "--allowed-local-media-path set to this same directory (or a parent), " "otherwise requests fail with 'Cannot load local files without " @@ -116,7 +118,7 @@ def add_seed_tts_cli_args(parser: argparse.ArgumentParser) -> None: default=False, help="Keep synthesized audio as 24 kHz mono PCM for WER (works with " "--backend openai-audio-speech or openai-chat-omni). Scoring follows " - "BytedanceSpeech/seed-tts-eval (Whisper-large-v3 / Paraformer-zh + jiwer). " + "zhaochenyang20/seed-tts-eval (Whisper-large-v3 / Paraformer-zh + jiwer). " "Sets SEED_TTS_WER_EVAL=1. Install: pip install 'vllm-omni[seed-tts-eval]'. " "Optional: SEED_TTS_EVAL_DEVICE, SEED_TTS_HF_WHISPER_MODEL.", ) diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index e355ee679..a6fd39741 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -9,7 +9,6 @@ import argparse import json import os import signal -import sys from types import FrameType from typing import Any @@ -20,9 +19,9 @@ from vllm.entrypoints.utils import VLLM_SUBCMD_PARSER_EPILOG from vllm.logger import init_logger from vllm.utils.argparse_utils import FlexibleArgumentParser +from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults from vllm_omni.entrypoints.cli.logo import log_logo from vllm_omni.entrypoints.openai.api_server import omni_run_server -from vllm_omni.entrypoints.utils import detect_explicit_cli_keys logger = init_logger(__name__) @@ -95,10 +94,6 @@ class OmniServeCommand(CLISubcommand): if hasattr(args, "model_tag") and args.model_tag is not None: args.model = args.model_tag - # Stash the set of long-option keys the user actually typed so the - # stage-config factory can give YAML precedence over argparse defaults. - args._cli_explicit_keys = detect_explicit_cli_keys(sys.argv[1:], OmniServeCommand._parser) - if args.headless: run_headless(args) else: @@ -267,6 +262,40 @@ class OmniServeCommand(CLISubcommand): default=None, help="Override the diffusion pipeline class name (e.g. LTX2ImageToVideoPipeline).", ) + omni_config_group.add_argument( + "--diffusion-load-format", + dest="diffusion_load_format", + type=str, + default=None, + choices=["default", "custom_pipeline", "dummy", "diffusers"], + help=( + "How to load the diffusion pipeline: native/registry (default), " + "custom_pipeline, dummy, or diffusers for the HF diffusers adapter." + ), + ) + omni_config_group.add_argument( + "--diffusers-load-kwargs", + dest="diffusers_load_kwargs", + type=json.loads, + default="{}", + help=( + "JSON object passed to DiffusionPipeline.from_pretrained()." + "It overrides corresponding parameters in the standard vLLM-Omni interface." + '(e.g. \'{"use_safetensors": true, "variant": "fp16"}\').' + ), + ) + omni_config_group.add_argument( + "--diffusers-call-kwargs", + dest="diffusers_call_kwargs", + type=json.loads, + default="{}", + help=( + "JSON object passed to pipeline.__call__(). " + "Useful for model-specific sampling parameters not covered by the vLLM-Omni interface." + "During request time, it is overridden by corresponding parameters in the vLLM-Omni interface." + '(e.g. \'{"num_inference_steps": 30, "guidance_scale": 7.5}\').' + ), + ) omni_config_group.add_argument( "--usp", "--ulysses-degree", @@ -451,6 +480,8 @@ class OmniServeCommand(CLISubcommand): # Stash via type(self) so the docs hook (which execs this function in a # sandboxed globals dict via ``DummySelf``) doesn't fail on a NameError. type(self)._parser = serve_parser + + nullify_stage_engine_defaults(serve_parser) return serve_parser @@ -506,15 +537,11 @@ def run_headless(args: argparse.Namespace) -> None: raise ValueError("headless mode requires worker_backend=multi_process") args_dict = vars(args).copy() - # Preserve the explicit-keys set captured at parse time so per-stage yaml - # values (e.g. stage 1's ``gpu_memory_utilization: 0.5``) are not - # overwritten by argparse defaults for flags the user didn't type. - cli_explicit_keys = args_dict.pop("_cli_explicit_keys", None) + args_dict.pop("_cli_explicit_keys", None) config_path, stage_configs = load_and_resolve_stage_configs( model, args_dict.get("stage_configs_path"), args_dict, - cli_explicit_keys=cli_explicit_keys, ) # Locate the stage config that matches stage_id. diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index 3180c9c80..8ffa39e16 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/vllm_omni/entrypoints/omni_base.py @@ -17,7 +17,7 @@ from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError from vllm_omni.engine.async_omni_engine import AsyncOmniEngine from vllm_omni.entrypoints.client_request_state import ClientRequestState from vllm_omni.entrypoints.pd_utils import PDDisaggregationMixin -from vllm_omni.entrypoints.utils import detect_explicit_cli_keys, get_final_stage_id_for_e2e +from vllm_omni.entrypoints.utils import get_final_stage_id_for_e2e from vllm_omni.metrics.stats import OrchestratorAggregator as OrchestratorMetrics from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific from vllm_omni.outputs import OmniRequestOutput @@ -28,6 +28,23 @@ if TYPE_CHECKING: logger = init_logger(__name__) +class OmniEngineDeadError(EngineDeadError): + _DEFAULT_MESSAGE = EngineDeadError().args[0] + error_stage_id: int | None + + def __init__( + self, + message: str | None = None, + *, + error_stage_id: int | None = None, + suppress_context: bool = False, + ) -> None: + resolved_message = message or self._DEFAULT_MESSAGE + Exception.__init__(self, resolved_message) + self.__suppress_context__ = suppress_context + self.error_stage_id = error_stage_id + + def _weak_shutdown_engine(engine: AsyncOmniEngine) -> None: """Best-effort engine cleanup for GC finalization.""" try: @@ -79,38 +96,23 @@ class OmniBase(PDDisaggregationMixin): parser: argparse.ArgumentParser | None = None, **overrides: Any, ) -> OmniBase: - """Construct an ``Omni`` / ``AsyncOmni`` from an ``argparse.Namespace``. + """Build from argparse. If ``parser`` is passed and not yet nullified, + un-typed engine fields are reset to ``None``.""" + kwargs: dict[str, Any] = {k: v for k, v in vars(args).items() if not k.startswith("_")} - Mirrors the ``EngineArgs.from_cli_args`` pattern used upstream and in - ``OmniEngineArgs.from_cli_args``. This is the recommended entry point - for any argparse-based caller (offline scripts, tests, CI): it - expands ``vars(args)`` into kwargs and automatically captures which - flags the user typed on the command line so that argparse defaults - do not silently override deploy YAML values. + if parser is not None and not getattr(parser, "_omni_nullified", False): + from vllm_omni.engine.arg_utils import ( + deploy_override_field_names, + ) + from vllm_omni.entrypoints.utils import detect_explicit_cli_keys - Passing ``parser`` is strongly recommended: without it, flag-to-dest - resolution falls back to a name-based heuristic that misidentifies - flags with ``dest=`` overrides, alias flags, and ``--disable-X`` / - ``store_false`` pairs. See :func:`detect_explicit_cli_keys`. + explicit = detect_explicit_cli_keys(sys.argv[1:], parser) or set() + override_dests = deploy_override_field_names() + for key in list(kwargs): + if key in override_dests and key not in explicit: + kwargs[key] = None - Args: - args: Parsed argparse namespace from ``parser.parse_args()``. - parser: The argparse parser used to produce ``args``. When - provided, each user-typed flag is resolved to its real - ``dest`` via the parser's action table. - **overrides: Extra keyword arguments that take precedence over - attributes on ``args``. - - Example:: - - parser = FlexibleArgumentParser() - parser.add_argument("--model", required=True) - args = parser.parse_args() - omni = Omni.from_cli_args(args, parser=parser) # preferred - omni = Omni.from_cli_args(args, parser=parser, model="other") - """ - kwargs: dict[str, Any] = {**vars(args), **overrides} - kwargs["_cli_explicit_keys"] = detect_explicit_cli_keys(sys.argv[1:], parser) + kwargs.update(overrides) return cls(**kwargs) def __init__( @@ -119,6 +121,7 @@ class OmniBase(PDDisaggregationMixin): **kwargs: Any, ) -> None: engine_args: OmniEngineArgs | None = kwargs.pop("engine_args", None) + stage_init_timeout = kwargs.pop("stage_init_timeout", 300) init_timeout = kwargs.pop("init_timeout", 600) log_stats = kwargs.pop("log_stats", False) @@ -296,8 +299,12 @@ class OmniBase(PDDisaggregationMixin): if msg_type == "error": error_text = msg.get("error", "Orchestrator returned an error message") + stage_id = msg.get("stage_id") if msg.get("fatal"): - raise EngineDeadError(error_text) + raise OmniEngineDeadError( + error_text, + error_stage_id=stage_id, + ) raise RuntimeError(error_text) if msg_type != "output": @@ -352,7 +359,10 @@ class OmniBase(PDDisaggregationMixin): ) # NOTE: O(n_stages) check for every error. if self.errored: - raise EngineDeadError(error_text) + raise OmniEngineDeadError( + error_text, + error_stage_id=stage_id, + ) raise EngineGenerateError(error_text) def _process_single_result( diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 41ccbcf1d..6345f0006 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -117,6 +117,7 @@ from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech from vllm_omni.entrypoints.openai.serving_speech_stream import OmniStreamingSpeechHandler from vllm_omni.entrypoints.openai.serving_video import OmniOpenAIServingVideo, ReferenceImage +from vllm_omni.entrypoints.openai.serving_video_stream import OmniStreamingVideoHandler from vllm_omni.entrypoints.openai.storage import STORAGE_MANAGER from vllm_omni.entrypoints.openai.stores import VIDEO_STORE, VIDEO_TASKS from vllm_omni.entrypoints.openai.utils import get_stage_type, parse_lora_request @@ -220,33 +221,64 @@ def _register_omni_exception_handlers(app) -> None: req: Request, exc: EngineDeadError | EngineGenerateError, ): - request_id = req.state.request_metadata.request_id if hasattr(req.state, "request_metadata") else None + request_id = _get_request_id_from_request(req) if req.app.state.args.log_error_stack: logger.exception("Engine Exception caught. Request id: %s", request_id) - engine = req.app.state.engine_client - if isinstance(exc, EngineDeadError): - # Log Omni-specific diagnostic information for dead engines. - orchestrator_alive = engine.engine.is_alive() if hasattr(engine, "engine") else "N/A" - logger.error( - "EngineDeadError: orchestrator_alive=%s, errored=%s, request_id=%s", - orchestrator_alive, - engine.errored, - request_id, - ) - - terminate_if_errored( - server=req.app.state.server, - engine=engine, - ) - err = create_error_response(exc) - return JSONResponse(err.model_dump(), status_code=err.error.code) + return _create_engine_error_json_response(req, exc) app.exception_handler(EngineGenerateError)(omni_engine_error_handler) app.exception_handler(EngineDeadError)(omni_engine_error_handler) +def _get_request_id_from_request(req: Request) -> str | None: + return req.state.request_metadata.request_id if hasattr(req.state, "request_metadata") else None + + +def _build_engine_error_payload( + exc: EngineDeadError | EngineGenerateError, + *, + request_id: str | None, +) -> tuple[dict[str, Any], int]: + err = create_error_response(exc) + payload = err.model_dump() + error_body = payload.get("error", {}) + + error_body["request_id"] = request_id + error_body["error_stage_id"] = getattr(exc, "error_stage_id", None) + + return payload, err.error.code + + +def _create_engine_error_json_response( + req: Request, + exc: EngineDeadError | EngineGenerateError, +) -> JSONResponse: + request_id = _get_request_id_from_request(req) + error_stage_id = getattr(exc, "error_stage_id", None) + engine = req.app.state.engine_client + + if isinstance(exc, EngineDeadError): + # Log Omni-specific diagnostic information for dead engines. + orchestrator_alive = engine.engine.is_alive() if hasattr(engine, "engine") else "N/A" + logger.error( + "EngineDeadError: orchestrator_alive=%s, errored=%s, request_id=%s, error_stage_id=%s", + orchestrator_alive, + engine.errored, + request_id, + error_stage_id, + ) + + terminate_if_errored( + server=req.app.state.server, + engine=engine, + ) + + payload, status_code = _build_engine_error_payload(exc, request_id=request_id) + return JSONResponse(content=payload, status_code=status_code) + + class _DiffusionServingModels: """Minimal OpenAIServingModels implementation for diffusion-only servers. @@ -575,6 +607,7 @@ async def omni_init_app_state( stage_configs=diffusion_stage_configs, ) state.openai_streaming_speech = None + state.openai_streaming_video = None state.enable_server_load_tracking = getattr(args, "enable_server_load_tracking", False) state.server_load_metrics = 0 @@ -876,6 +909,14 @@ async def omni_init_app_state( state.openai_streaming_speech = OmniStreamingSpeechHandler( speech_service=state.openai_serving_speech, ) + state.openai_streaming_video = ( + OmniStreamingVideoHandler( + chat_service=state.openai_serving_chat, + engine_client=engine_client, + ) + if state.openai_serving_chat is not None + else None + ) state.openai_serving_realtime = OpenAIServingRealtime( engine_client=engine_client, models=state.openai_serving_models, @@ -930,8 +971,8 @@ async def create_chat_completion(request: ChatCompletionRequest, raw_request: Re return base_server.create_error_response(message="The model does not support Chat Completions API") try: generator = await handler.create_chat_completion(request, raw_request) - except (EngineGenerateError, EngineDeadError): - raise # Propagate to the global Omni exception handler + except (EngineGenerateError, EngineDeadError) as exc: + return _create_engine_error_json_response(raw_request, exc) except Exception as e: logger.exception("Chat completion failed: %s", e) raise HTTPException(status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=str(e)) from e @@ -1027,8 +1068,8 @@ async def create_speech(request: OpenAICreateSpeechRequest, raw_request: Request status_code=result.error.code if result.error else 400, ) return result - except (EngineGenerateError, EngineDeadError): - raise # Propagate to the global Omni exception handler + except (EngineGenerateError, EngineDeadError) as exc: + return _create_engine_error_json_response(raw_request, exc) except Exception as e: raise HTTPException(status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=str(e)) from e @@ -1063,8 +1104,8 @@ async def create_speech_batch(request: BatchSpeechRequest, raw_request: Request) status_code=result.error.code if result.error else 400, ) return JSONResponse(content=result.model_dump()) - except (EngineGenerateError, EngineDeadError): - raise # Propagate to the global Omni exception handler + except (EngineGenerateError, EngineDeadError) as exc: + return _create_engine_error_json_response(raw_request, exc) except ValueError as e: raise HTTPException(status_code=HTTPStatus.BAD_REQUEST.value, detail=str(e)) from e except Exception as e: @@ -1261,6 +1302,23 @@ async def streaming_speech(websocket: WebSocket): await handler.handle_session(websocket) +@router.websocket("/v1/video/chat/stream") +async def streaming_video_chat(websocket: WebSocket): + """WebSocket endpoint for streaming video input chat.""" + handler = getattr(websocket.app.state, "openai_streaming_video", None) + if handler is None: + await websocket.accept() + await websocket.send_json( + { + "type": "error", + "message": "Streaming video chat is not available", + } + ) + await websocket.close() + return + await handler.handle_session(websocket) + + @router.websocket("/v1/realtime") async def realtime_websocket(websocket: WebSocket): """WebSocket endpoint for OpenAI-style realtime interactions.""" @@ -1522,8 +1580,8 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request) data=image_data, ) - except (EngineGenerateError, EngineDeadError): - raise # Propagate to the global Omni exception handler + except (EngineGenerateError, EngineDeadError) as exc: + return _create_engine_error_json_response(raw_request, exc) except HTTPException: raise except ValueError as e: @@ -1759,8 +1817,8 @@ async def edit_images( size=size_str, ) - except (EngineGenerateError, EngineDeadError): - raise # Propagate to the global Omni exception handler + except (EngineGenerateError, EngineDeadError) as exc: + return _create_engine_error_json_response(raw_request, exc) except HTTPException: raise except ValueError as e: @@ -2330,9 +2388,9 @@ async def _parse_video_form( true_cfg_scale: float | None = Form(default=None), seed: int | None = Form(default=None), negative_prompt: str | None = Form(default=None), - enable_frame_interpolation: bool = Form(default=False), - frame_interpolation_exp: int = Form(default=1, ge=1), - frame_interpolation_scale: float = Form(default=1.0, gt=0.0), + enable_frame_interpolation: bool | None = Form(default=None), + frame_interpolation_exp: int | None = Form(default=None, ge=1), + frame_interpolation_scale: float | None = Form(default=None, gt=0.0), frame_interpolation_model_path: str | None = Form(default=None), lora: str | None = Form(default=None), extra_params: str | None = Form(default=None), @@ -2481,8 +2539,8 @@ async def create_video_sync( status_code=HTTPStatus.GATEWAY_TIMEOUT.value, detail=f"Video generation timed out after {VIDEO_SYNC_TIMEOUT_S}s.", ) - except (EngineGenerateError, EngineDeadError): - raise # Propagate to the global Omni exception handler + except (EngineGenerateError, EngineDeadError) as exc: + return _create_engine_error_json_response(raw_request, exc) except HTTPException: raise except Exception as exc: @@ -2748,6 +2806,8 @@ if __name__ == "__main__": from vllm.entrypoints.openai.cli_args import make_arg_parser + from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults + parser = argparse.ArgumentParser(description="vLLM-Omni OpenAI-Compatible REST API server") parser = make_arg_parser(parser) registered_flags = set() @@ -2759,6 +2819,7 @@ if __name__ == "__main__": parser.add_argument( "--enable-sleep-mode", action="store_true", default=False, help="Enable GPU memory pool for sleep mode." ) + nullify_stage_engine_defaults(parser) args = parser.parse_args() if not hasattr(args, "model_tag"): setattr(args, "model_tag", args.model) diff --git a/vllm_omni/entrypoints/openai/protocol/audio.py b/vllm_omni/entrypoints/openai/protocol/audio.py index 8468efd86..59b5777a8 100644 --- a/vllm_omni/entrypoints/openai/protocol/audio.py +++ b/vllm_omni/entrypoints/openai/protocol/audio.py @@ -1,5 +1,5 @@ import math -from typing import Literal +from typing import Any, Literal import numpy as np from pydantic import AliasChoices, BaseModel, Field, field_validator, model_validator @@ -74,6 +74,10 @@ class OpenAICreateSpeechRequest(BaseModel): ge=0, description="Per-request initial chunk size override. If null, computed dynamically based on server load.", ) + extra_params: dict[str, Any] | None = Field( + default=None, + description=("Optional model-specific parameters passed directly to the model's extra_args."), + ) @field_validator("stream_format") @classmethod diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 06b739e3b..0e4919f82 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -559,6 +559,16 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): engine_prompt["additional_information"] = {} engine_prompt["additional_information"]["language"] = [language.strip()] + # Style instruction — used by Ming-flash-omni instruct TTS path + # (ming_task="instruct"). For the omni speech path the thinker2talker + # bridge drops this field to match upstream omni_audio_generation + # which hardcodes instruction=None. + instructions = getattr(request, "instructions", None) + if instructions is not None and isinstance(instructions, str) and instructions.strip(): + if "additional_information" not in engine_prompt or engine_prompt["additional_information"] is None: + engine_prompt["additional_information"] = {} + engine_prompt["additional_information"]["instruction"] = instructions.strip() + return conversation, [engine_prompt] async def _inject_audio_from_video_urls( @@ -1636,7 +1646,13 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): logger.warning(f"Unsupported final output type: {omni_outputs.final_output_type}") continue if omni_outputs.metrics: - response_metrics = omni_outputs.metrics + response_metrics = dict(omni_outputs.metrics) + if omni_outputs.final_output_type == "image": + # Expose diffusion profiler metrics on the top-level response for benchmarks / clients. + if response_metrics is None: + response_metrics = {} + response_metrics.setdefault("stage_durations", omni_outputs.stage_durations or {}) + response_metrics.setdefault("peak_memory_mb", float(omni_outputs.peak_memory_mb or 0.0)) choices.extend(choices_data) response = OmniChatCompletionResponse( @@ -1930,7 +1946,8 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): final_res = omni_outputs.request_output # OMNI: Access multimodal_output from CompletionOutput (outputs[0]), not from RequestOutput # Reference: examples/offline_inference/qwen3_omni/end2end.py line 421 - audio_data = final_res.outputs[0].multimodal_output.get("audio") + mm_output = final_res.outputs[0].multimodal_output + audio_data = mm_output.get("audio") if isinstance(audio_data, list): if stream: audio_tensor = audio_data[-1] @@ -1944,9 +1961,20 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): if audio_tensor.ndim > 1: audio_tensor = audio_tensor.flatten() + # Prefer the talker-reported sample rate when present. Qwen3-Omni + # omits "sr" and runs at 24kHz; Ming-flash-omni surfaces a 44.1kHz + # AudioVAE rate via multimodal_output["sr"]. + sr_raw = mm_output.get("sr") + if sr_raw is None: + sample_rate = 24000 + elif hasattr(sr_raw, "item"): + sample_rate = int(sr_raw.item()) + else: + sample_rate = int(sr_raw) + audio_obj = CreateAudio( audio_tensor=audio_tensor, - sample_rate=24000, + sample_rate=sample_rate, response_format="wav", speed=1.0, stream_format="audio", @@ -2164,6 +2192,24 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): comprehension_idx = idx break + is_glm_image_multistage = False + for stage_cfg in stage_configs: + if isinstance(stage_cfg, dict): + engine_args = stage_cfg.get("engine_args") + model_arch = stage_cfg.get("model_arch") + else: + engine_args = getattr(stage_cfg, "engine_args", None) + model_arch = getattr(stage_cfg, "model_arch", None) + + if model_arch is None and engine_args is not None: + if isinstance(engine_args, dict): + model_arch = engine_args.get("model_arch") + else: + model_arch = getattr(engine_args, "model_arch", None) + if model_arch in {"GlmImageForConditionalGeneration", "GlmImagePipeline"}: + is_glm_image_multistage = True + break + sampling_params_list: list[Any] = [] for idx, stage_cfg in enumerate(stage_configs): stage_type = get_stage_type(stage_cfg) @@ -2212,6 +2258,25 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): default_stage_params.lora_scale = lora_scale except Exception as e: # pragma: no cover - safeguard logger.warning("Failed to parse LoRA request: %s", e) + elif ( + is_glm_image_multistage + and comprehension_idx is not None + and idx == comprehension_idx + and isinstance(default_stage_params, SamplingParams) + and height is not None + and width is not None + ): + from vllm_omni.model_executor.stage_input_processors.glm_image import compute_max_tokens + + default_stage_params.max_tokens = compute_max_tokens( + int(height), + int(width), + is_i2i=bool(reference_images), + ) + extra_args = dict(getattr(default_stage_params, "extra_args", {}) or {}) + extra_args["target_h"] = int(height) + extra_args["target_w"] = int(width) + default_stage_params.extra_args = extra_args sampling_params_list.append(default_stage_params) diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 7161aa318..5bc84b18a 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -8,13 +8,14 @@ import re import struct import time from concurrent.futures import ThreadPoolExecutor +from http import HTTPStatus from pathlib import Path from typing import Any import numpy as np import soundfile as sf import torch -from fastapi import Request, UploadFile +from fastapi import HTTPException, Request, UploadFile from fastapi.responses import Response, StreamingResponse from transformers.utils.hub import cached_file from vllm.entrypoints.launcher import terminate_if_errored @@ -44,6 +45,12 @@ from vllm_omni.model_executor.models.fish_speech.prompt_utils import ( estimate_fish_voice_clone_prompt_len_from_normalized, normalize_fish_voice_clone_texts, ) +from vllm_omni.model_executor.models.ming_flash_omni.prompt_utils import ( + DEFAULT_PROMPT as MING_DEFAULT_PROMPT, +) +from vllm_omni.model_executor.models.ming_flash_omni.prompt_utils import ( + create_instruction as ming_create_instruction, +) from vllm_omni.outputs import OmniRequestOutput logger = init_logger(__name__) @@ -57,6 +64,7 @@ _OMNIVOICE_TTS_MODEL_STAGES = {"omnivoice_generator"} _VOXCPM_TTS_MODEL_STAGES = {"latent_generator", "vae"} _VOXCPM2_TTS_MODEL_STAGES = {"latent_generator"} _MOSS_TTS_MODEL_STAGES = {"moss_tts_nano"} +_MING_TTS_MODEL_STAGES = {"ming_tts"} _TTS_MODEL_STAGES: set[str] = ( _VOXTRAL_TTS_MODEL_STAGES | _QWEN3_TTS_MODEL_STAGES @@ -66,6 +74,7 @@ _TTS_MODEL_STAGES: set[str] = ( | _VOXCPM_TTS_MODEL_STAGES | _VOXCPM2_TTS_MODEL_STAGES | _MOSS_TTS_MODEL_STAGES + | _MING_TTS_MODEL_STAGES ) _TTS_LANGUAGES: set[str] = { "Auto", @@ -318,6 +327,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): for stage in self.engine_client.stage_configs ) return "voxcpm" if has_vae_stage or model_stage == "vae" else "voxcpm2" + if model_stage in _MING_TTS_MODEL_STAGES: + return "ming_flash_omni_tts" return None def _compute_max_instructions_length(self) -> int: @@ -341,6 +352,11 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): def _load_supported_speakers(self) -> set[str]: """Load supported speakers (case-insensitive) from the model configuration.""" + if self._tts_model_type == "ming_flash_omni_tts": + # Ming-flash-omni drives speaker selection via the caption JSON + # (audio_sequence[0]["说话人"]) rather than a spk_id table, so there + # is no static speaker list to surface here. + return set() try: if self._tts_model_type == "voxcpm": return set() @@ -844,6 +860,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): return None # VoxCPM2 accepts any text input if self._tts_model_type == "moss_tts_nano": return self._validate_moss_tts_request(request) + if self._tts_model_type == "ming_flash_omni_tts": + return self._validate_ming_tts_request(request) return self._validate_qwen_tts_request(request) def _voxcpm2_encode(self, text: str) -> list[int]: @@ -865,6 +883,41 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): ids = self._voxcpm2_tokenizer.encode(text, add_special_tokens=True) return split_multichar_chinese(ids, self._voxcpm2_split_map) + def _validate_ming_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None: + """Validate Ming-flash-omni standalone-talker request parameters.""" + if not request.input or not request.input.strip(): + return "Input text cannot be empty" + if request.instructions is not None: + if not isinstance(request.instructions, str): + return "instructions must be a string" + if len(request.instructions) > self._max_instructions_length: + return f"instructions exceeds max length {self._max_instructions_length}" + + if request.task_type is not None: + return "'task_type' is not supported for Ming-flash-omni TTS" + if request.language is not None: + return "'language' is not supported for Ming-flash-omni TTS (language is inferred from input text)" + if request.x_vector_only_mode is not None: + return "'x_vector_only_mode' is not supported for Ming-flash-omni TTS" + if request.initial_codec_chunk_frames is not None: + return "'initial_codec_chunk_frames' is not supported for Ming-flash-omni TTS" + + # Per-request voice cloning from raw audio is not yet wired up: Ming + # extracts spk_emb / prompt_wav_lat / prompt_wav_emb model-side via + # register_prompt_wav() at engine init. For ad-hoc cloning, callers + # should pre-compute speaker_embedding and pass it directly. + if request.ref_audio is not None: + return ( + "'ref_audio' is not yet supported for Ming-flash-omni TTS; " + "use a preset 'voice' or 'speaker_embedding' instead" + ) + if request.ref_text is not None: + return "'ref_text' is not yet supported for Ming-flash-omni TTS" + + if request.max_new_tokens is not None and request.max_new_tokens <= 0: + return "'max_new_tokens' must be a positive integer" + return None + def _validate_ref_audio_format(self, ref_audio: str) -> str | None: """Validate ref_audio is a supported URI format. Returns error or None.""" if not ( @@ -1563,6 +1616,52 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): }, } + # ---- Ming-flash-omni standalone-talker (TTS) helpers ---- + + def _build_ming_prompt(self, request: OpenAICreateSpeechRequest) -> dict[str, Any]: + # request.instructions accepts two forms: + # 1. Plain text: mapped to the caption's 风格 (style) field + # 2. JSON object: parsed and splatted into the caption. Unlocks + # Unknown keys are dropped by `ming_create_instruction`. + caption_fields: dict[str, Any] = {} + if request.instructions: + stripped = request.instructions.strip() + if stripped.startswith("{"): + try: + parsed = json.loads(stripped) + except json.JSONDecodeError: + parsed = None + if isinstance(parsed, dict): + caption_fields.update(parsed) + else: + caption_fields["风格"] = request.instructions + else: + caption_fields["风格"] = request.instructions + + has_spk_emb = request.speaker_embedding is not None + + # TTS path applies ming task type `instruct`. + # voice_name enables talker-side voice preset resolution (e.g. "DB30"). + additional_information: dict[str, Any] = { + "ming_task": "instruct", + "prompt": MING_DEFAULT_PROMPT, + "text": request.input, + "instruction": ming_create_instruction(caption_fields), + "voice_name": request.voice or None, + "use_zero_spk_emb": not has_spk_emb, + "max_decode_steps": request.max_new_tokens or _TTS_MAX_NEW_TOKENS_MAX, + "cfg": 2.0, + "sigma": 0.25, + "temperature": 0.0, + } + if has_spk_emb: + # Passed as plain float list + additional_information["spk_emb"] = list(request.speaker_embedding) + return { + "prompt_token_ids": [0], + "additional_information": additional_information, + } + # ---- Common speech generation helpers ---- async def _prepare_speech_generation( @@ -1628,6 +1727,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): elif self._tts_model_type == "moss_tts_nano": tts_params = await self._build_moss_tts_params(request) prompt = {"prompt_token_ids": [1], "additional_information": tts_params} + elif self._tts_model_type == "ming_flash_omni_tts": + prompt = self._build_ming_prompt(request) + tts_params = {} else: tts_params = self._build_tts_params(request) # Resolve ref_audio (explicit or auto-set for uploaded voices) @@ -1675,6 +1777,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): model_type = "voxcpm" elif self._tts_model_type == "voxcpm2": model_type = "voxcpm2" + elif self._tts_model_type == "ming_flash_omni_tts": + model_type = "ming_flash_omni_tts" elif self._is_tts: model_type = tts_params.get("task_type", ["unknown"])[0] else: @@ -1711,6 +1815,21 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): max_tokens, ) + # Apply model-specific extra parameters + if request.extra_params is not None and sampling_params_list: + if not isinstance(request.extra_params, dict): + raise HTTPException( + status_code=HTTPStatus.BAD_REQUEST.value, + detail="extra_params must be a JSON object/dict.", + ) + import copy + + sampling_params_list = copy.deepcopy(sampling_params_list) + if sampling_params_list[0].extra_args is None: + sampling_params_list[0].extra_args = {} + sampling_params_list[0].extra_args.update(request.extra_params) + logger.info("Applied extra_params: %s", request.extra_params) + # Fish defaults come from stage_configs YAML. Only override when the caller # explicitly requests a different generation length. if self._is_fish_speech and request.max_new_tokens is not None and sampling_params_list: diff --git a/vllm_omni/entrypoints/openai/serving_video.py b/vllm_omni/entrypoints/openai/serving_video.py index 741295c7c..a4be330eb 100644 --- a/vllm_omni/entrypoints/openai/serving_video.py +++ b/vllm_omni/entrypoints/openai/serving_video.py @@ -3,6 +3,7 @@ from __future__ import annotations +import copy import time from dataclasses import dataclass from http import HTTPStatus @@ -95,7 +96,7 @@ class OmniOpenAIServingVideo: if request.negative_prompt is not None: prompt["negative_prompt"] = request.negative_prompt - gen_params = OmniDiffusionSamplingParams() + gen_params = self._resolve_default_sampling_params() input_image = None if reference_image is None else reference_image.data vp = request.resolve_video_params() @@ -113,22 +114,27 @@ class OmniOpenAIServingVideo: if vp.fps is not None: gen_params.fps = vp.fps gen_params.frame_rate = float(vp.fps) - gen_params.enable_frame_interpolation = request.enable_frame_interpolation - gen_params.frame_interpolation_exp = request.frame_interpolation_exp - gen_params.frame_interpolation_scale = request.frame_interpolation_scale - gen_params.frame_interpolation_model_path = request.frame_interpolation_model_path + provided_fields = request.model_fields_set + if "enable_frame_interpolation" in provided_fields: + gen_params.enable_frame_interpolation = request.enable_frame_interpolation + if "frame_interpolation_exp" in provided_fields: + gen_params.frame_interpolation_exp = request.frame_interpolation_exp + if "frame_interpolation_scale" in provided_fields: + gen_params.frame_interpolation_scale = request.frame_interpolation_scale + if "frame_interpolation_model_path" in provided_fields: + gen_params.frame_interpolation_model_path = request.frame_interpolation_model_path - if request.num_inference_steps is not None: + if "num_inference_steps" in provided_fields and request.num_inference_steps is not None: gen_params.num_inference_steps = request.num_inference_steps - if request.guidance_scale is not None: + if "guidance_scale" in provided_fields and request.guidance_scale is not None: gen_params.guidance_scale = request.guidance_scale - if request.guidance_scale_2 is not None: + if "guidance_scale_2" in provided_fields and request.guidance_scale_2 is not None: gen_params.guidance_scale_2 = request.guidance_scale_2 - if request.true_cfg_scale is not None: + if "true_cfg_scale" in provided_fields and request.true_cfg_scale is not None: gen_params.true_cfg_scale = request.true_cfg_scale - if request.seed is not None: + if "seed" in provided_fields and request.seed is not None: gen_params.seed = request.seed - if request.boundary_ratio is not None: + if "boundary_ratio" in provided_fields and request.boundary_ratio is not None: gen_params.boundary_ratio = request.boundary_ratio logger.info( @@ -136,7 +142,7 @@ class OmniOpenAIServingVideo: request.boundary_ratio, gen_params.boundary_ratio, ) - if request.flow_shift is not None: + if "flow_shift" in provided_fields and request.flow_shift is not None: gen_params.extra_args["flow_shift"] = request.flow_shift # Apply model-specific extra parameters @@ -263,6 +269,17 @@ class OmniOpenAIServingVideo: return int(multiplier) return 1 + def _resolve_default_sampling_params(self) -> OmniDiffusionSamplingParams: + default_sampling_params_list = getattr(self._engine_client, "default_sampling_params_list", None) + if default_sampling_params_list: + for params in default_sampling_params_list: + if isinstance(params, OmniDiffusionSamplingParams): + # Requests mutate sampling params in-place, including + # nested dict fields like extra_args. Deep-copy the stage + # defaults so one request cannot leak state into another. + return copy.deepcopy(params) + return OmniDiffusionSamplingParams() + @staticmethod def _apply_lora(lora_body: Any, gen_params: OmniDiffusionSamplingParams) -> None: try: diff --git a/vllm_omni/entrypoints/openai/serving_video_stream.py b/vllm_omni/entrypoints/openai/serving_video_stream.py new file mode 100644 index 000000000..a76b241c5 --- /dev/null +++ b/vllm_omni/entrypoints/openai/serving_video_stream.py @@ -0,0 +1,971 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""WebSocket handler for streaming video input understanding. + +Accepts video frames incrementally via WebSocket, buffers them, and +generates text + optional audio responses using the existing Qwen3-Omni +multi-stage pipeline (thinker -> talker -> code2wav). + +Protocol: + Client -> Server: + {"type": "session.config", ...} # Session config (sent once) + {"type": "video.frame", "data": "..."} # base64 JPEG/PNG frame + {"type": "audio.chunk", "data": "..."} # base64 PCM16 16kHz mono + {"type": "video.query", "text": "..."} # Submit query about buffered frames + {"type": "video.done"} # End of session + + Server -> Client: + {"type": "response.start"} + {"type": "response.text.delta", "delta": "..."} + {"type": "response.text.done", "text": "..."} + {"type": "response.audio.delta", "data": "...", "format": "wav"} + {"type": "response.audio.done"} + {"type": "session.done"} + {"type": "error", "message": "..."} +""" + +import asyncio +import base64 +import hashlib +import io +import json +import time as _time +import uuid +import wave +from typing import Any + +import torch +from fastapi import WebSocket, WebSocketDisconnect +from PIL import Image +from pydantic import BaseModel, Field, ValidationError +from vllm.logger import init_logger + +from vllm_omni.entrypoints.openai import video_stream_envs +from vllm_omni.entrypoints.openai.video_frame_filter import FrameSimilarityFilter +from vllm_omni.entrypoints.openai.video_stream_context import ( + text_only_message, +) +from vllm_omni.outputs import OmniRequestOutput + +logger = init_logger(__name__) + +_DEFAULT_IDLE_TIMEOUT = 60.0 +_DEFAULT_CONFIG_TIMEOUT = 10.0 +_MAX_FRAME_SIZE = 10 * 1024 * 1024 # 10MB per frame +_MAX_BUFFER_FRAMES = 64 +_MAX_AUDIO_BUFFER_BYTES = 4 * 1024 * 1024 +_MAX_MSG_QUEUE = 200 +_CODEC_FRAME_SAMPLES = 1920 # CausalConv leading-edge artifact length +_BAD_FRAME = object() + + +def _decode_frame_bytes(raw_bytes: bytes) -> Any: + return Image.open(io.BytesIO(raw_bytes)).convert("RGB") + + +class StreamingVideoSessionConfig(BaseModel): + """Configuration sent as the first WebSocket message.""" + + model: str | None = None + modalities: list[str] = Field( + default_factory=lambda: ["text", "audio"], + description="Output modalities: 'text', 'audio', or both.", + ) + num_frames: int = Field( + default=4, + ge=1, + le=128, + description="Max frames to sample from buffer for the model.", + ) + max_frames: int = Field( + default=50, + ge=1, + le=256, + description="Max frames to keep in the buffer.", + ) + system_prompt: str | None = Field( + default=None, + description="Custom system prompt.", + ) + use_audio_in_video: bool = Field( + default=True, + description="Interleave audio chunks with video frames when audio input is present.", + ) + sampling_params_list: list[dict[str, Any]] | None = Field( + default=None, + description="Per-stage sampling params [thinker, talker, code2wav].", + ) + enable_frame_filter: bool = Field( + default=True, + description="EVS pixel-similarity pre-filter to drop near-duplicate frames.", + ) + frame_filter_threshold: float = Field( + default=0.95, + ge=0.0, + le=1.0, + description="EVS similarity threshold (higher = keep more frames).", + ) + + +class OmniStreamingVideoHandler: + """Handles WebSocket sessions for streaming video input. + + Supports: + - Concurrent frame reception during query processing (reader/processor split) + - PCM audio input (``audio.chunk``) + - Async-chunk incremental audio output via ``engine_client.generate()`` + - Multi-turn conversation history + - Soft interrupt (new query cancels current generation) + """ + + def __init__( + self, + chat_service: Any, + idle_timeout: float = _DEFAULT_IDLE_TIMEOUT, + config_timeout: float = _DEFAULT_CONFIG_TIMEOUT, + engine_client: Any | None = None, + ) -> None: + self._chat_service = chat_service + self._idle_timeout = idle_timeout + self._config_timeout = config_timeout + self._engine_client = engine_client + + async def handle_session(self, websocket: WebSocket) -> None: + """Main session loop for a single WebSocket connection.""" + await websocket.accept() + + try: + config = await self._receive_config(websocket) + if config is None: + return + + frame_buffer: list[str] = [] # base64-encoded JPEG frames + # Per-frame PIL cache + uuid for mm_hash reuse. Aligned with frame_buffer by index. + frame_pil_cache: dict[str, tuple[Any, str] | object] = {} # b64 -> (PIL.Image, uuid) or _BAD_FRAME + frame_filter = ( + FrameSimilarityFilter(threshold=config.frame_filter_threshold) if config.enable_frame_filter else None + ) + audio_buffer = bytearray() # raw PCM16 16kHz mono + message_history: list[dict[str, Any]] = [] + active_request_id: str | None = None + prev_request_id: str | None = None # abort target iff prev was interrupted + prev_was_interrupted: bool = False + interrupt_event = asyncio.Event() + prewarm_tasks: set[asyncio.Task[Any]] = set() + query_task: asyncio.Task[Any] | None = None + + msg_queue: asyncio.Queue[dict[str, Any] | None] = asyncio.Queue(maxsize=_MAX_MSG_QUEUE) + + async def _reader() -> None: + """Receive WebSocket messages and enqueue them.""" + try: + while True: + try: + raw = await asyncio.wait_for( + websocket.receive_text(), + timeout=self._idle_timeout, + ) + except asyncio.TimeoutError: + await self._send_error(websocket, "Idle timeout") + await msg_queue.put(None) + return + + try: + msg = json.loads(raw) + except json.JSONDecodeError: + await self._send_error(websocket, "Invalid JSON") + continue + + if not isinstance(msg, dict): + await self._send_error(websocket, "Messages must be JSON objects") + continue + + msg_type = str(msg.get("type", "")) + if msg_type.startswith("_internal."): + await self._send_error(websocket, f"Unknown type: {msg_type}") + continue + + await msg_queue.put(msg) + if msg.get("type") == "video.done": + return + except WebSocketDisconnect: + await msg_queue.put(None) + except Exception: + await msg_queue.put(None) + raise + + async def _cancel_active_query(*, abort_now: bool = False) -> None: + """Signal soft interrupt for the active query.""" + nonlocal active_request_id, prev_was_interrupted, query_task + if active_request_id is not None: + interrupt_event.set() + prev_was_interrupted = True + logger.info("Interrupt signaled for %s", active_request_id) + if abort_now and self._engine_client: + try: + await self._engine_client.abort(active_request_id) + except Exception: + logger.debug("Abort failed for %s", active_request_id, exc_info=True) + if query_task is not None and not query_task.done(): + query_task.cancel() + await asyncio.gather(query_task, return_exceptions=True) + query_task = None + + async def _processor() -> None: + """Process enqueued messages.""" + nonlocal active_request_id, prev_request_id, prev_was_interrupted, query_task + + while True: + msg = await msg_queue.get() + if msg is None: + await _cancel_active_query(abort_now=True) + return + + msg_type = msg.get("type") + + if msg_type == "_internal.frame_decode_failed": + frame_data = msg.get("b64", "") + removed = frame_data in frame_buffer + if removed: + frame_buffer[:] = [f for f in frame_buffer if f != frame_data] + if frame_pil_cache.get(frame_data) is _BAD_FRAME: + frame_pil_cache.pop(frame_data, None) + if removed: + await self._send_error(websocket, "Frame decode failed") + + elif msg_type == "video.frame": + frame_data = msg.get("data", "") + if not frame_data: + continue + if len(frame_data) > _MAX_FRAME_SIZE: + await self._send_error(websocket, "Frame too large") + continue + try: + raw_bytes = base64.b64decode(frame_data, validate=True) + except Exception: + await self._send_error(websocket, "Invalid image data") + continue + if frame_filter is not None: + try: + if not frame_filter.should_retain(raw_bytes): + continue + except Exception: + await self._send_error(websocket, "Invalid image data") + continue + max_buf = config.max_frames + if len(frame_buffer) >= max_buf: + dropped = frame_buffer.pop(0) + frame_pil_cache.pop(dropped, None) + frame_buffer.append(frame_data) + # Prewarm: decode PIL off the event loop so query-time chat_template + # can skip base64+Image.open. uuid=md5 lets mm_cache dedupe identical frames. + if frame_data not in frame_pil_cache: + mm_uuid = hashlib.md5(raw_bytes, usedforsecurity=False).hexdigest() + + async def _prewarm(b64: str, b: bytes, u: str) -> None: + try: + pil = await asyncio.to_thread(_decode_frame_bytes, b) + frame_pil_cache[b64] = (pil, u) + except Exception: + frame_pil_cache[b64] = _BAD_FRAME + logger.warning("prewarm decode failed for frame (len=%d)", len(b)) + try: + msg_queue.put_nowait({"type": "_internal.frame_decode_failed", "b64": b64}) + except asyncio.QueueFull: + logger.warning( + "frame decode failure event dropped because message queue is full" + ) + + task = asyncio.create_task(_prewarm(frame_data, raw_bytes, mm_uuid)) + prewarm_tasks.add(task) + task.add_done_callback(prewarm_tasks.discard) + + elif msg_type == "audio.chunk": + data_b64 = msg.get("data", "") + try: + pcm_bytes = base64.b64decode(data_b64) + except Exception: + continue + if len(audio_buffer) + len(pcm_bytes) > _MAX_AUDIO_BUFFER_BYTES: + await self._send_error(websocket, "Audio buffer overflow") + audio_buffer.clear() + continue + audio_buffer.extend(pcm_bytes) + + elif msg_type == "video.query": + await _cancel_active_query() + + query_text = msg.get("text", "") + audio_data_b64 = msg.get("audio_data") + if audio_data_b64: + try: + decoded = base64.b64decode(audio_data_b64) + if len(audio_buffer) + len(decoded) <= _MAX_AUDIO_BUFFER_BYTES: + audio_buffer.extend(decoded) + else: + await self._send_error(websocket, "Audio buffer overflow") + audio_buffer.clear() + except Exception: + pass + + if not frame_buffer: + await self._send_error(websocket, "No frames buffered") + continue + + # Abort only if the previous turn was interrupted mid-flight. + # A naturally-finished request is already released by the scheduler; + # aborting it again can race with stage-1/2 tear-down and has been + # observed to crash flash_attn with a mixed prefill+decode batch + # (scheduler_metadata shape mismatch) under longer sessions. + if prev_was_interrupted and prev_request_id and self._engine_client: + try: + await self._engine_client.abort(prev_request_id) + except Exception: + pass + await asyncio.sleep(0.1) + prev_was_interrupted = False + + request_id = f"video-{uuid.uuid4().hex[:12]}" + active_request_id = request_id + interrupt_event.clear() + query_frames = list(frame_buffer) + query_audio_buffer = bytearray(audio_buffer) + audio_buffer.clear() + query_prewarmed_frames = dict(frame_pil_cache) + + async def _run_query() -> None: + nonlocal active_request_id, prev_request_id + try: + await self._process_query( + websocket, + config, + query_frames, + query_audio_buffer, + message_history, + query_text, + request_id, + interrupt_event, + query_prewarmed_frames, + ) + finally: + if active_request_id == request_id: + prev_request_id = request_id + active_request_id = None + + query_task = asyncio.create_task(_run_query()) + + elif msg_type == "video.done": + if query_task is not None and not query_task.done(): + await asyncio.gather(query_task, return_exceptions=True) + query_task = None + await websocket.send_json({"type": "session.done"}) + return + + elif msg_type == "ping": + try: + await websocket.send_json({"type": "pong"}) + except Exception: + pass + + else: + await self._send_error(websocket, f"Unknown type: {msg_type}") + + reader_task = asyncio.create_task(_reader()) + try: + await _processor() + finally: + reader_task.cancel() + try: + await reader_task + except (asyncio.CancelledError, Exception): + pass + for t in list(prewarm_tasks): + t.cancel() + if prewarm_tasks: + await asyncio.gather(*prewarm_tasks, return_exceptions=True) + if query_task is not None and not query_task.done(): + await _cancel_active_query(abort_now=True) + + except WebSocketDisconnect: + logger.info("Streaming video: client disconnected") + except Exception as e: + logger.exception("Streaming video session error: %s", e) + try: + await self._send_error(websocket, f"Internal error: {e}") + except Exception: + pass + + async def _receive_config(self, websocket: WebSocket) -> StreamingVideoSessionConfig | None: + """Wait for and validate the session.config message.""" + try: + raw = await asyncio.wait_for( + websocket.receive_text(), + timeout=self._config_timeout, + ) + except asyncio.TimeoutError: + await self._send_error(websocket, "Timeout waiting for session.config") + return None + + try: + msg = json.loads(raw) + except json.JSONDecodeError: + await self._send_error(websocket, "Invalid JSON in session.config") + return None + + if not isinstance(msg, dict) or msg.get("type") != "session.config": + await self._send_error( + websocket, + f"Expected session.config, got: {msg.get('type') if isinstance(msg, dict) else type(msg).__name__}", + ) + return None + + config_data = {k: v for k, v in msg.items() if k != "type"} + alias_map = { + "num_sample_frames": "num_frames", + "evs_enabled": "enable_frame_filter", + "evs_threshold": "frame_filter_threshold", + } + for old_key, new_key in alias_map.items(): + if old_key in config_data and new_key not in config_data: + config_data[new_key] = config_data[old_key] + + try: + config = StreamingVideoSessionConfig(**config_data) + except ValidationError as e: + await self._send_error(websocket, f"Invalid session config: {e}") + return None + + return config + + async def _process_query( + self, + websocket: WebSocket, + config: StreamingVideoSessionConfig, + frame_buffer: list[str], + audio_buffer: bytearray, + message_history: list[dict[str, Any]], + query_text: str, + request_id: str, + interrupt_event: asyncio.Event, + prewarmed_frames: dict[str, tuple[Any, str]], + ) -> None: + """Build prompt, run inference, stream text + audio response.""" + + if self._engine_client is None: + await self._send_error(websocket, "Streaming video requires an engine client") + return + + await self._process_query_engine( + websocket, + config, + frame_buffer, + audio_buffer, + message_history, + query_text, + request_id, + interrupt_event, + prewarmed_frames, + ) + + # ------------------------------------------------------------------ + # Engine-client path (async_chunk audio streaming) + # ------------------------------------------------------------------ + + async def _process_query_engine( + self, + websocket: WebSocket, + config: StreamingVideoSessionConfig, + frame_buffer: list[str], + audio_buffer: bytearray, + message_history: list[dict[str, Any]], + query_text: str, + request_id: str, + interrupt_event: asyncio.Event, + prewarmed_frames: dict[str, tuple[Any, str]], + ) -> None: + """Direct engine_client.generate() path for async_chunk audio.""" + from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionRequest, + ) + + messages, user_message = self._build_messages( + config, + frame_buffer, + audio_buffer, + message_history, + query_text, + prewarmed_frames, + ) + + request_kwargs: dict[str, Any] = { + "model": config.model or "default", + "messages": messages, + "stream": True, + "modalities": config.modalities, + "add_generation_prompt": True, + "continue_final_message": False, + "add_special_tokens": False, + } + if config.use_audio_in_video and len(audio_buffer) > 0: + request_kwargs["mm_processor_kwargs"] = { + "use_audio_in_video": True, + } + if config.sampling_params_list: + request_kwargs["sampling_params_list"] = config.sampling_params_list + + try: + chat_request = ChatCompletionRequest(**request_kwargs) + except Exception as e: + await self._send_error(websocket, f"Failed to build request: {e}") + return + + try: + engine_prompt = await self._preprocess_to_engine_prompt(chat_request) + except Exception as e: + await self._send_error(websocket, f"Preprocess failed: {e}") + return + + await websocket.send_json({"type": "response.start"}) + text_parts: list[str] = [] + text_done_sent = False + audio_chunk_count = 0 + # Number of per-step tensors in OmniRequestOutput.audio_data already + # drained. Used by the fast path to skip already-emitted history. + audio_chunks_drained = 0 + previous_text = "" + interrupted = False + t_start = _time.monotonic() + t_first_text = None + t_first_audio = None + + # Wire-level async-chunk switch. "off" means + # buffer all deltas server-side and flush once at the end; the engine + # pipeline still overlaps internally. + async_chunk_mode = video_stream_envs.VLLM_VIDEO_ASYNC_CHUNK + streaming = async_chunk_mode == "on" + audio_tail_tensors: list[Any] = [] + + try: + result_gen = self._engine_client.generate( + prompt=engine_prompt, + request_id=request_id, + output_modalities=config.modalities, + ) + + async for output in result_gen: + # Soft interrupt: drain without sending + if interrupt_event.is_set(): + if not interrupted: + logger.info("Generation interrupted — draining") + interrupted = True + continue + + if not isinstance(output, OmniRequestOutput): + continue + + out_type = getattr(output, "final_output_type", "text") + + if out_type == "audio": + if streaming and not text_done_sent: + full_text = "".join(text_parts) + await websocket.send_json({"type": "response.text.done", "text": full_text}) + text_done_sent = True + + if t_first_audio is None: + t_first_audio = _time.monotonic() + audio_chunk_count += 1 + if streaming: + b64, audio_chunks_drained = self._extract_audio_delta_b64( + output, + audio_chunks_drained, + ) + if b64: + await websocket.send_json( + { + "type": "response.audio.delta", + "data": b64, + "format": "wav", + } + ) + else: + audio_data = self._get_audio_data(output) + if audio_data is not None: + if isinstance(audio_data, list): + audio_tail_tensors = list(audio_data) + else: + audio_tail_tensors = [audio_data] + else: + delta_text, previous_text = self._extract_text_delta( + output, + previous_text, + ) + if delta_text: + if t_first_text is None: + t_first_text = _time.monotonic() + text_parts.append(delta_text) + if streaming: + await websocket.send_json({"type": "response.text.delta", "delta": delta_text}) + + if not text_done_sent: + full_text = "".join(text_parts) + await websocket.send_json({"type": "response.text.done", "text": full_text}) + text_done_sent = True + + if not streaming and audio_tail_tensors: + try: + coalesced = ( + audio_tail_tensors[0] if len(audio_tail_tensors) == 1 else torch.cat(audio_tail_tensors, dim=-1) + ) + tail_np = self._tensor_to_1d_np(coalesced) + b64, _ = self._encode_tail( + tail_np, + 0, + new_drained=len(audio_tail_tensors), + is_first=True, + ) + if b64: + await websocket.send_json( + { + "type": "response.audio.delta", + "data": b64, + "format": "wav", + } + ) + except Exception: + logger.exception("Failed to coalesce off-path audio") + + if audio_chunk_count > 0: + await websocket.send_json({"type": "response.audio.done"}) + + response_text = "".join(text_parts) + message_history.append(user_message) + message_history.append({"role": "assistant", "content": response_text}) + + t_end = _time.monotonic() + logger.info( + "[TIMING] mode=%s total=%.2fs first_text=%.2fs first_audio=%.2fs audio_chunks=%d", + async_chunk_mode, + t_end - t_start, + (t_first_text - t_start) if t_first_text else -1, + (t_first_audio - t_start) if t_first_audio else -1, + audio_chunk_count, + ) + + except Exception: + logger.exception("Engine query failed") + await self._send_error(websocket, "Query processing failed") + + if not text_done_sent: + full_text = "".join(text_parts) + await websocket.send_json({"type": "response.text.done", "text": full_text}) + + # ------------------------------------------------------------------ + # Message building + # ------------------------------------------------------------------ + + def _build_messages( + self, + config: StreamingVideoSessionConfig, + frame_buffer: list[str], + audio_buffer: bytearray, + message_history: list[dict[str, Any]], + query_text: str, + prewarmed_frames: dict[str, tuple[Any, str]], + ) -> tuple[list[dict[str, Any]], dict[str, Any]]: + """Build OpenAI-style messages list and the current user message. + + Returns (messages, user_message). + """ + # Stride sampling (index 0 anchor, last slot = newest). Covers full buffer + stable mm_hash. + n_buf = len(frame_buffer) + if n_buf <= config.num_frames: + frames = list(frame_buffer) + else: + stride = max(1, n_buf // config.num_frames) + idx = [i * stride for i in range(config.num_frames - 1)] + [n_buf - 1] + frames = [frame_buffer[i] for i in idx] + + # Prefer prewarmed PIL + uuid so mm_cache can dedupe by hash. + prewarmed = prewarmed_frames or {} + user_content: list[dict] = [] + for frame_b64 in frames: + cached = prewarmed.get(frame_b64) + if cached is _BAD_FRAME: + continue + if cached is not None: + pil, pil_uuid = cached + user_content.append( + { + "type": "image_pil", + "image_pil": pil, + "uuid": pil_uuid, + } + ) + else: + user_content.append( + { + "type": "image_url", + "image_url": {"url": f"data:image/jpeg;base64,{frame_b64}"}, + } + ) + + if len(audio_buffer) > 0: + wav_b64 = self._pcm_to_wav_b64(bytes(audio_buffer)) + user_content.append( + { + "type": "input_audio", + "input_audio": { + "data": wav_b64, + "format": "wav", + }, + } + ) + + if query_text: + user_content.append({"type": "text", "text": query_text}) + + user_message: dict[str, Any] = {"role": "user", "content": user_content} + + messages: list[dict[str, Any]] = [] + if config.system_prompt: + messages.append({"role": "system", "content": config.system_prompt}) + + # Add text-only history (strip images/audio from old turns). + # Keep only the last turn (2 messages) to keep prompt short + # enough for single-step mm_encoder scheduling. When prompt + # exceeds ~50 tokens, the V1 scheduler splits mm_encoder and + # text prefill, causing incomplete thinker embeddings and + # garbled audio. + recent_history = message_history[-2:] if len(message_history) > 2 else message_history + for hist_msg in recent_history: + messages.append(self._text_only_message(hist_msg)) + + messages.append(user_message) + + return messages, user_message + + # ------------------------------------------------------------------ + # Audio helpers + # ------------------------------------------------------------------ + + @staticmethod + def _pcm_to_wav_b64(pcm_data: bytes, sample_rate: int = 16000) -> str: + """Wrap raw PCM16 mono in a WAV container and return base64.""" + buf = io.BytesIO() + with wave.open(buf, "wb") as wf: + wf.setnchannels(1) + wf.setsampwidth(2) + wf.setframerate(sample_rate) + wf.writeframes(pcm_data) + return base64.b64encode(buf.getvalue()).decode() + + @classmethod + def _extract_audio_delta_b64( + cls, + result: OmniRequestOutput, + chunks_drained: int, + ) -> tuple[str | None, int]: + """Return (base64 WAV of new samples, updated chunks_drained). + + `chunks_drained` is the number of per-step tensors in + ``audio_data`` that have already been emitted. Each engine step appends + one tensor, so new samples are ``audio_data[chunks_drained:]`` — no + matter how many steps accumulated between reads (handles backpressure + cleanly, unlike a simple ``audio_data[-1]``). + + Two paths, selected at runtime by ``VLLM_VIDEO_AUDIO_DELTA_MODE``: + * fast — only D2H the new tail. Per-call cost ∝ new chunks. + * slow — full cat + D2H each call. Per-call cost ∝ total history. + Retained for A/B; remove once downstream callers confirm. + """ + audio_data = cls._get_audio_data(result) + if audio_data is None: + return None, chunks_drained + + if video_stream_envs.VLLM_VIDEO_AUDIO_DELTA_MODE == "slow": + return cls._delta_slow(audio_data, chunks_drained) + return cls._delta_fast(audio_data, chunks_drained) + + @staticmethod + def _get_audio_data(result: OmniRequestOutput): + """Navigate OmniRequestOutput → multimodal_output['audio']. None on miss.""" + request_output = getattr(result, "request_output", None) + if request_output is None: + return None + outputs = getattr(request_output, "outputs", None) + if not isinstance(outputs, list) or not outputs: + return None + mm_output = getattr(outputs[0], "multimodal_output", None) + if not isinstance(mm_output, dict): + return None + return mm_output.get("audio") + + @classmethod + def _delta_fast( + cls, + audio_data, + chunks_drained: int, + ) -> tuple[str | None, int]: + """Emit only tensors appended since the last call.""" + # Single tensor: output_processor hands us one tensor before it becomes a + # list (see output_processor.py:89). Treat it as chunk #0. + if not isinstance(audio_data, list): + if chunks_drained >= 1: + return None, chunks_drained + tail_np = cls._tensor_to_1d_np(audio_data) + return cls._encode_tail(tail_np, chunks_drained, new_drained=1, is_first=True) + + n = len(audio_data) + if n <= chunks_drained: + return None, chunks_drained + + new_chunks = audio_data[chunks_drained:] + tail = new_chunks[0] if len(new_chunks) == 1 else torch.cat(new_chunks, dim=-1) + tail_np = cls._tensor_to_1d_np(tail) + return cls._encode_tail(tail_np, chunks_drained, new_drained=n, is_first=(chunks_drained == 0)) + + @classmethod + def _delta_slow( + cls, + audio_data, + chunks_drained: int, + ) -> tuple[str | None, int]: + """Pre-fix behaviour: concat everything each call and slice on CPU.""" + if isinstance(audio_data, list): + if not audio_data: + return None, chunks_drained + audio_tensor = torch.cat(audio_data, dim=-1) + new_drained = len(audio_data) + else: + audio_tensor = audio_data + new_drained = 1 + + full_np = cls._tensor_to_1d_np(audio_tensor) + if full_np is None: + return None, chunks_drained + # chunks_drained doesn't map directly to sample offset without tracking + # per-chunk lengths, so we re-derive: replay the tail that corresponds + # to chunks appended since last call by slicing off the part produced + # by the already-drained prefix. For slow path this is intentionally + # wasteful — the point is to reproduce the pre-fix hot loop. + if chunks_drained == 0: + tail_np = full_np + else: + # Recover prefix length by re-concatenating the already-drained + # prefix tensors (cost intentionally identical to the baseline + # implementation this was lifted from). + if isinstance(audio_data, list) and chunks_drained < len(audio_data): + prefix_len = sum(int(t.shape[-1]) for t in audio_data[:chunks_drained]) + tail_np = full_np[prefix_len:] + else: + tail_np = full_np[0:0] + return cls._encode_tail(tail_np, chunks_drained, new_drained=new_drained, is_first=(chunks_drained == 0)) + + @classmethod + def _encode_tail( + cls, + tail_np, + old_drained: int, + *, + new_drained: int, + is_first: bool, + ) -> tuple[str | None, int]: + """Strip the CausalConv leading artifact on first emit, then b64-encode.""" + if tail_np is None or len(tail_np) == 0: + return None, new_drained + if is_first and len(tail_np) > _CODEC_FRAME_SAMPLES * 2: + tail_np = tail_np[_CODEC_FRAME_SAMPLES:] + if len(tail_np) == 0: + return None, new_drained + try: + return cls._encode_audio_wav_b64(tail_np), new_drained + except Exception: + logger.exception("Failed to encode audio delta WAV") + return None, old_drained + + @staticmethod + def _tensor_to_1d_np(t): + """Tensor → flat float32 numpy on CPU. None on failure.""" + if t is None or not hasattr(t, "float"): + return None + arr = t.float().detach().cpu().numpy() + if arr.ndim > 1: + arr = arr.flatten() + return arr + + @staticmethod + def _encode_audio_wav_b64(audio_np) -> str: + """Encode numpy float32 audio to base64 WAV (24kHz).""" + from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin + from vllm_omni.entrypoints.openai.protocol.audio import CreateAudio + + audio_obj = CreateAudio( + audio_tensor=audio_np, + sample_rate=24000, + response_format="wav", + speed=1.0, + stream_format="audio", + base64_encode=True, + ) + mixin = AudioMixin() + resp = mixin.create_audio(audio_obj) + return resp.audio_data + + @staticmethod + def _extract_text_delta( + result: OmniRequestOutput, + previous_text: str, + ) -> tuple[str, str]: + """Extract incremental text delta from OmniRequestOutput.""" + if result.final_output_type != "text": + return "", previous_text + + request_output = getattr(result, "request_output", None) + if request_output is None: + return "", previous_text + + outputs = getattr(request_output, "outputs", None) + if not isinstance(outputs, list) or not outputs: + return "", previous_text + + text = getattr(outputs[0], "text", None) + if not isinstance(text, str) or not text: + return "", previous_text + + if text.startswith(previous_text): + return text[len(previous_text) :], text + return text, text + + # ------------------------------------------------------------------ + # Preprocessing + # ------------------------------------------------------------------ + + async def _preprocess_to_engine_prompt(self, request) -> Any: + """Use the chat handler's preprocessing to build an engine prompt.""" + handler = self._chat_service + renderer = handler.renderer + + _conversation, engine_prompts = await handler._preprocess_chat( + request, + request.messages, + default_template=getattr(request, "chat_template", None) or handler.chat_template, + default_template_content_format=handler.chat_template_content_format, + renderer=renderer, + add_generation_prompt=request.add_generation_prompt, + continue_final_message=request.continue_final_message, + add_special_tokens=request.add_special_tokens, + ) + return engine_prompts[0] + + # ------------------------------------------------------------------ + # Utilities + # ------------------------------------------------------------------ + + _text_only_message = staticmethod(text_only_message) + + async def _send_error(self, websocket: WebSocket, message: str) -> None: + """Send an error message to the client.""" + try: + await websocket.send_json({"type": "error", "message": message}) + except Exception: + pass diff --git a/vllm_omni/entrypoints/openai/video_frame_filter.py b/vllm_omni/entrypoints/openai/video_frame_filter.py new file mode 100644 index 000000000..4ea427835 --- /dev/null +++ b/vllm_omni/entrypoints/openai/video_frame_filter.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""EVS (Efficient Video Sampling) frame pre-filter for streaming video input. + +Lightweight pixel-level similarity filter that runs before frames reach the +vision encoder. For static or slow-moving scenes this can reduce the number +of frames by 2-5x, proportionally cutting encoder compute and KV-cache usage. + +Usage: + filter = FrameSimilarityFilter(threshold=0.95) + for jpeg_bytes in incoming_frames: + if filter.should_retain(jpeg_bytes): + buffer.append(jpeg_bytes) +""" + +from __future__ import annotations + +import io + +import numpy as np +from PIL import Image + +_DEFAULT_THRESHOLD = 0.95 +_DEFAULT_THUMBNAIL_SIZE = 64 + + +class FrameSimilarityFilter: + """Drop near-duplicate frames based on pixel-level similarity. + + Each incoming JPEG frame is down-scaled to a small thumbnail and compared + against the last *retained* frame. If the normalised similarity exceeds + ``threshold`` the frame is considered redundant and dropped. + + Args: + threshold: Similarity score in [0, 1] above which a frame is dropped. + Higher values keep more frames (less aggressive filtering). + Default 0.95 works well for typical webcam / surveillance feeds. + thumbnail_size: Edge length of the square thumbnail used for + comparison. Larger values are more accurate but slower. + """ + + def __init__( + self, + threshold: float = _DEFAULT_THRESHOLD, + thumbnail_size: int = _DEFAULT_THUMBNAIL_SIZE, + ) -> None: + if not 0.0 <= threshold <= 1.0: + raise ValueError(f"threshold must be in [0, 1], got {threshold}") + if thumbnail_size < 1: + raise ValueError(f"thumbnail_size must be >= 1, got {thumbnail_size}") + + self._threshold = threshold + self._thumbnail_size = thumbnail_size + self._last_retained: np.ndarray | None = None + self._retained_count = 0 + self._dropped_count = 0 + + # ------------------------------------------------------------------ + # Public API + # ------------------------------------------------------------------ + + def should_retain(self, frame_jpeg: bytes) -> bool: + """Return ``True`` if *frame_jpeg* is sufficiently different from the + last retained frame and should be kept in the buffer.""" + current = self._decode_and_resize(frame_jpeg) + + if self._last_retained is None: + self._last_retained = current + self._retained_count += 1 + return True + + similarity = self._compute_similarity(self._last_retained, current) + if similarity >= self._threshold: + self._dropped_count += 1 + return False + + self._last_retained = current + self._retained_count += 1 + return True + + def reset(self) -> None: + """Clear internal state so the next frame is always retained.""" + self._last_retained = None + self._retained_count = 0 + self._dropped_count = 0 + + @property + def stats(self) -> dict[str, int | float]: + """Return filtering statistics.""" + total = self._retained_count + self._dropped_count + return { + "retained_count": self._retained_count, + "dropped_count": self._dropped_count, + "total_count": total, + "drop_rate": self._dropped_count / total if total > 0 else 0.0, + } + + # ------------------------------------------------------------------ + # Internals + # ------------------------------------------------------------------ + + @staticmethod + def _compute_similarity(a: np.ndarray, b: np.ndarray) -> float: + """Normalised pixel similarity in [0, 1]. 1 = identical.""" + mse = float(np.mean((a.astype(np.float32) - b.astype(np.float32)) ** 2)) + return 1.0 - mse / (255.0 * 255.0) + + def _decode_and_resize(self, jpeg_bytes: bytes) -> np.ndarray: + """Decode JPEG and resize to a small square thumbnail for fast + comparison.""" + img = Image.open(io.BytesIO(jpeg_bytes)) + img = img.resize( + (self._thumbnail_size, self._thumbnail_size), + Image.Resampling.BILINEAR, + ).convert("RGB") + return np.asarray(img, dtype=np.uint8) diff --git a/vllm_omni/entrypoints/openai/video_stream_context.py b/vllm_omni/entrypoints/openai/video_stream_context.py new file mode 100644 index 000000000..e72dfcd41 --- /dev/null +++ b/vllm_omni/entrypoints/openai/video_stream_context.py @@ -0,0 +1,28 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Small helpers shared by streaming video handlers.""" + +from __future__ import annotations + +from typing import Any + + +def text_only_message(message: dict[str, Any]) -> dict[str, Any]: + """Return a history message with multimodal content stripped out.""" + role = message.get("role", "user") + content = message.get("content", "") + + if isinstance(content, str): + return {"role": role, "content": content} + + if not isinstance(content, list): + return {"role": role, "content": ""} + + text_parts: list[str] = [] + for part in content: + if isinstance(part, dict) and part.get("type") == "text": + text = part.get("text") + if isinstance(text, str): + text_parts.append(text) + + return {"role": role, "content": "".join(text_parts)} diff --git a/vllm_omni/entrypoints/openai/video_stream_envs.py b/vllm_omni/entrypoints/openai/video_stream_envs.py new file mode 100644 index 000000000..f8571e7f9 --- /dev/null +++ b/vllm_omni/entrypoints/openai/video_stream_envs.py @@ -0,0 +1,62 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Environment variables for the streaming video OpenAI entrypoint.""" + +import logging +import os +from collections.abc import Callable +from typing import TYPE_CHECKING, Literal + +if TYPE_CHECKING: + VLLM_VIDEO_AUDIO_DELTA_MODE: Literal["fast", "slow"] = "fast" + VLLM_VIDEO_ASYNC_CHUNK: Literal["on", "off"] = "on" + +logger = logging.getLogger(__name__) +_warned_invalid_envs: set[tuple[str, str]] = set() +_VIDEO_AUDIO_DELTA_MODE = "VLLM_VIDEO_AUDIO_DELTA_MODE" +_VIDEO_ASYNC_CHUNK = "VLLM_VIDEO_ASYNC_CHUNK" + + +def _choice_env( + name: str, + default: str, + allowed: tuple[str, ...], +) -> str: + value = os.getenv(name, default).strip().lower() + if value in allowed: + return value + warning_key = (name, value) + if warning_key not in _warned_invalid_envs: + logger.warning("%s=%s not recognized; falling back to %r", name, value, default) + _warned_invalid_envs.add(warning_key) + return default + + +environment_variables: dict[str, Callable[[], str]] = { + _VIDEO_AUDIO_DELTA_MODE: lambda: _choice_env( + _VIDEO_AUDIO_DELTA_MODE, + "fast", + ("fast", "slow"), + ), + _VIDEO_ASYNC_CHUNK: lambda: _choice_env( + _VIDEO_ASYNC_CHUNK, + "on", + ("on", "off"), + ), +} + + +def __getattr__(name: str) -> str: + if name in environment_variables: + return environment_variables[name]() + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +def __dir__() -> list[str]: + return list(environment_variables.keys()) + + +def is_set(name: str) -> bool: + if name in environment_variables: + return name in os.environ + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/vllm_omni/entrypoints/openai/video_stream_session.py b/vllm_omni/entrypoints/openai/video_stream_session.py new file mode 100644 index 000000000..ead3b1ec5 --- /dev/null +++ b/vllm_omni/entrypoints/openai/video_stream_session.py @@ -0,0 +1,439 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Video streaming session manager and WebSocket handler. + +Provides ``VideoStreamConfig``, ``VideoStreamSession`` (frame/audio buffer +with EVS filtering), and ``VideoStreamHandler`` (WebSocket session loop). +Builds standard ``ChatCompletionRequest`` so ``OmniOpenAIServingChat`` is +reused with zero changes. +""" + +from __future__ import annotations + +import asyncio +import base64 +import json +from collections import deque +from typing import Any + +from fastapi import WebSocket, WebSocketDisconnect +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionRequest, +) +from vllm.entrypoints.openai.engine.protocol import ErrorResponse +from vllm.logger import init_logger + +from vllm_omni.entrypoints.openai.video_frame_filter import FrameSimilarityFilter + +logger = init_logger(__name__) + +_DEFAULT_IDLE_TIMEOUT = 60.0 +_DEFAULT_CONFIG_TIMEOUT = 10.0 +_MAX_FRAME_BYTES = 10 * 1024 * 1024 # 10 MB per frame + + +class VideoStreamConfig: + """Per-session configuration sent by the client in ``session.config``.""" + + __slots__ = ( + "model", + "modalities", + "max_frames", + "num_sample_frames", + "evs_enabled", + "evs_threshold", + ) + + def __init__( + self, + model: str = "", + modalities: list[str] | None = None, + max_frames: int = 64, + num_sample_frames: int = 16, + evs_enabled: bool = True, + evs_threshold: float = 0.95, + ) -> None: + self.model = model + self.modalities = modalities if modalities is not None else ["text"] + self.max_frames = max_frames + self.num_sample_frames = num_sample_frames + self.evs_enabled = evs_enabled + self.evs_threshold = evs_threshold + + # JSON doesn't distinguish int/float, so float fields accept both. + _FIELD_TYPES: dict[str, tuple[type, ...]] = { + "model": (str,), + "modalities": (list,), + "max_frames": (int,), + "num_sample_frames": (int,), + "evs_enabled": (bool,), + "evs_threshold": (int, float), + } + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> VideoStreamConfig: + known = set(cls.__slots__) + filtered = {k: v for k, v in data.items() if k in known} + + for field, accepted in cls._FIELD_TYPES.items(): + if field in filtered and not isinstance(filtered[field], accepted): + names = "/".join(t.__name__ for t in accepted) + raise TypeError(f"Invalid type for '{field}': expected {names}, got {type(filtered[field]).__name__}") + + return cls(**filtered) + + +class VideoStreamSession: + """Manages frame and audio buffers for a single streaming session. + + Frames are optionally filtered through ``FrameSimilarityFilter`` (EVS) + before being stored in a fixed-size ring buffer. ``build_chat_request`` + merges buffers into a ``ChatCompletionRequest`` with ``image_url`` (+ + ``audio_url``) content blocks. + """ + + def __init__(self, config: VideoStreamConfig) -> None: + self._config = config + self._frame_filter: FrameSimilarityFilter | None = ( + FrameSimilarityFilter(threshold=config.evs_threshold) if config.evs_enabled else None + ) + self._frames: deque[bytes] = deque(maxlen=config.max_frames) + self._audio_chunks: list[bytes] = [] + + def add_frame(self, jpeg_bytes: bytes) -> bool: + """Add a JPEG frame after EVS filtering. Returns ``True`` if kept.""" + if len(jpeg_bytes) > _MAX_FRAME_BYTES: + raise ValueError(f"Frame too large: {len(jpeg_bytes)} bytes (limit {_MAX_FRAME_BYTES})") + + if self._frame_filter and not self._frame_filter.should_retain(jpeg_bytes): + return False + + self._frames.append(jpeg_bytes) + return True + + def sample_frames(self) -> list[bytes]: + """Return up to ``num_sample_frames`` uniformly sampled frames.""" + n = len(self._frames) + k = min(n, self._config.num_sample_frames) + if k == 0: + return [] + if k == 1: + return [self._frames[-1]] + if k >= n: + return list(self._frames) + indices = [int(i * (n - 1) / (k - 1)) for i in range(k)] + return [self._frames[i] for i in indices] + + @property + def frame_count(self) -> int: + return len(self._frames) + + def add_audio_chunk(self, pcm_bytes: bytes) -> None: + """Append raw PCM 16 kHz audio bytes.""" + self._audio_chunks.append(pcm_bytes) + + def clear_audio(self) -> None: + """Clear audio buffer after a query is submitted.""" + self._audio_chunks.clear() + + @property + def has_audio(self) -> bool: + return bool(self._audio_chunks) + + def build_chat_request(self, query_text: str) -> ChatCompletionRequest: + """Build a ``ChatCompletionRequest`` from the current buffers.""" + sampled = self.sample_frames() + content_parts: list[dict[str, Any]] = [] + + for frame in sampled: + frame_b64 = base64.b64encode(frame).decode() + content_parts.append( + { + "type": "image_url", + "image_url": { + "url": f"data:image/jpeg;base64,{frame_b64}", + }, + } + ) + + if self.has_audio: + combined_pcm = b"".join(self._audio_chunks) + audio_b64 = base64.b64encode(combined_pcm).decode() + content_parts.append( + { + "type": "audio_url", + "audio_url": { + "url": f"data:audio/L16;rate=16000;base64,{audio_b64}", + }, + } + ) + + content_parts.append({"type": "text", "text": query_text}) + + messages: list[dict[str, Any]] = [ + {"role": "user", "content": content_parts}, + ] + + request_dict: dict[str, Any] = { + "model": self._config.model, + "messages": messages, + "stream": True, + } + + if self.has_audio: + request_dict["mm_processor_kwargs"] = { + "use_audio_in_video": True, + } + + return ChatCompletionRequest(**request_dict) + + @property + def evs_stats(self) -> dict[str, int | float] | None: + return self._frame_filter.stats if self._frame_filter else None + + +class VideoStreamHandler: + """Drives a ``VideoStreamSession`` over a FastAPI ``WebSocket``. + + Instantiate once at server startup and call ``handle_session`` for each + incoming connection. Uses two concurrent tasks (_reader + _processor) + so frames keep arriving while a query is being processed. + """ + + def __init__( + self, + chat_handler: Any, + idle_timeout: float = _DEFAULT_IDLE_TIMEOUT, + config_timeout: float = _DEFAULT_CONFIG_TIMEOUT, + ) -> None: + self._chat_handler = chat_handler + self._idle_timeout = idle_timeout + self._config_timeout = config_timeout + + async def handle_session(self, websocket: WebSocket) -> None: + """Main session loop for one WebSocket connection.""" + await websocket.accept() + + try: + config = await self._receive_config(websocket) + if config is None: + return + + session = VideoStreamSession(config) + logger.info( + "Video stream session started: model=%s modalities=%s max_frames=%d evs=%s", + config.model, + config.modalities, + config.max_frames, + config.evs_enabled, + ) + + msg_queue: asyncio.Queue[dict[str, Any] | None] = asyncio.Queue() + + async def _reader() -> None: + try: + while True: + try: + raw = await asyncio.wait_for( + websocket.receive_text(), + timeout=self._idle_timeout, + ) + except asyncio.TimeoutError: + await self._send_error(websocket, "Idle timeout: no message received") + await msg_queue.put(None) + return + + try: + msg = json.loads(raw) + except json.JSONDecodeError: + await self._send_error(websocket, "Invalid JSON") + continue + + if not isinstance(msg, dict): + await self._send_error(websocket, "Messages must be JSON objects") + continue + + await msg_queue.put(msg) + + if msg.get("type") == "video.done": + return + except WebSocketDisconnect: + await msg_queue.put(None) + except Exception: + await msg_queue.put(None) + raise + + async def _processor() -> None: + while True: + msg = await msg_queue.get() + if msg is None: + return + + msg_type = msg.get("type") + + if msg_type == "video.frame": + data_b64 = msg.get("data", "") + try: + jpeg_bytes = base64.b64decode(data_b64) + except Exception: + await self._send_error(websocket, "Invalid base64 in video.frame") + continue + try: + session.add_frame(jpeg_bytes) + except ValueError as exc: + await self._send_error(websocket, str(exc)) + continue + + elif msg_type == "audio.chunk": + data_b64 = msg.get("data", "") + try: + pcm_bytes = base64.b64decode(data_b64) + except Exception: + await self._send_error(websocket, "Invalid base64 in audio.chunk") + continue + session.add_audio_chunk(pcm_bytes) + + elif msg_type == "video.query": + query_text = msg.get("text", "") + if not query_text: + await self._send_error( + websocket, + "video.query requires a non-empty 'text' field", + ) + continue + await self._handle_query(websocket, session, query_text) + + elif msg_type == "video.done": + evs = session.evs_stats + if evs is not None: + await websocket.send_json({"type": "response.evs_stats", **evs}) + await websocket.send_json({"type": "session.done"}) + return + + else: + await self._send_error(websocket, f"Unknown message type: {msg_type}") + + reader_task = asyncio.create_task(_reader()) + try: + await _processor() + finally: + reader_task.cancel() + try: + await reader_task + except asyncio.CancelledError: + pass + except Exception: + logger.debug( + "Reader task raised an exception", + exc_info=True, + ) + + except WebSocketDisconnect: + logger.info("Video stream: client disconnected") + except Exception: + logger.exception("Video stream session error") + try: + await self._send_error(websocket, "Internal server error") + except Exception: + pass + + async def _receive_config(self, websocket: WebSocket) -> VideoStreamConfig | None: + try: + raw = await asyncio.wait_for( + websocket.receive_text(), + timeout=self._config_timeout, + ) + except asyncio.TimeoutError: + await self._send_error(websocket, "Timeout waiting for session.config") + return None + + try: + msg = json.loads(raw) + except json.JSONDecodeError: + await self._send_error(websocket, "Invalid JSON in session.config") + return None + + if not isinstance(msg, dict) or msg.get("type") != "session.config": + await self._send_error( + websocket, + f"Expected session.config, got: {msg.get('type') if isinstance(msg, dict) else type(msg).__name__}", + ) + return None + + try: + return VideoStreamConfig.from_dict(msg) + except TypeError as exc: + await self._send_error(websocket, str(exc)) + return None + + async def _handle_query( + self, + websocket: WebSocket, + session: VideoStreamSession, + query_text: str, + ) -> None: + """Build a ChatCompletionRequest, run it, and stream the response.""" + if session.frame_count == 0: + await self._send_error( + websocket, + "No frames available. Send video.frame before video.query.", + ) + return + + request = session.build_chat_request(query_text) + await websocket.send_json({"type": "response.start"}) + + # After response.start, the protocol contract requires us to always + # send response.text.done — even on error / exception paths. + text_parts: list[str] = [] + try: + # raw_request=None: serving_chat.py guards with `if raw_request:` + generator = await self._chat_handler.create_chat_completion(request, raw_request=None) + + if isinstance(generator, ErrorResponse): + error_msg = generator.error.message if generator.error else "Unknown error" + await self._send_error(websocket, error_msg) + else: + async for chunk in generator: + if isinstance(chunk, str): + for delta in self._parse_sse_deltas(chunk): + text_parts.append(delta) + await websocket.send_json( + { + "type": "response.text.delta", + "delta": delta, + } + ) + + except Exception: + logger.exception("Query failed") + await self._send_error(websocket, "Query processing failed") + + full_text = "".join(text_parts) + await websocket.send_json({"type": "response.text.done", "text": full_text}) + + session.clear_audio() + + @staticmethod + def _parse_sse_deltas(chunk: str) -> list[str]: + """Extract content deltas from SSE-formatted chunk.""" + deltas: list[str] = [] + for line in chunk.split("\n"): + line = line.strip() + if not line or not line.startswith("data: ") or line == "data: [DONE]": + continue + try: + data = json.loads(line[6:]) + delta = data.get("choices", [{}])[0].get("delta", {}).get("content", "") + if delta: + deltas.append(delta) + except (json.JSONDecodeError, IndexError, AttributeError): + pass + return deltas + + @staticmethod + async def _send_error(websocket: WebSocket, message: str) -> None: + try: + await websocket.send_json({"type": "error", "message": message}) + except Exception: + pass diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py index 04c0ffe47..5541fde30 100644 --- a/vllm_omni/entrypoints/utils.py +++ b/vllm_omni/entrypoints/utils.py @@ -21,6 +21,18 @@ PROJECT_ROOT = Path(__file__).parent.parent.parent logger = init_logger(__name__) + +def _warn_deprecated_explicit_keys(kwargs: dict[str, Any]) -> None: + if "cli_explicit_keys" in kwargs: + import warnings + + warnings.warn( + "cli_explicit_keys= is deprecated and ignored. Remove the kwarg.", + DeprecationWarning, + stacklevel=3, + ) + + _DIFFUSERS_CLASS_TO_CONFIG: dict[str, str] = { "GlmImagePipeline": "glm_image", } @@ -331,6 +343,10 @@ def resolve_model_config_path(model: str) -> str: if os.path.exists(complete_config_path): return str(complete_config_path) + deploy_config_path = PROJECT_ROOT / "vllm_omni" / "deploy" / model_type_str + if os.path.exists(deploy_config_path): + return str(deploy_config_path) + stage_config_file = f"vllm_omni/model_executor/stage_configs/{normalized_model_type}.yaml" stage_config_path = PROJECT_ROOT / stage_config_file if not os.path.exists(stage_config_path): @@ -343,7 +359,7 @@ def load_stage_configs_from_model( base_engine_args: dict | None = None, deploy_config_path: str | None = None, stage_overrides: dict[str, dict[str, Any]] | None = None, - cli_explicit_keys: set[str] | None = None, + **deprecated_kwargs: Any, ) -> list: """Load stage configurations from model's default config file. @@ -358,33 +374,25 @@ def load_stage_configs_from_model( base_engine_args: Base engine args to merge as CLI overrides. deploy_config_path: Optional explicit deploy config path. stage_overrides: Per-stage overrides from --stage-overrides. - cli_explicit_keys: Set of CLI keys the user actually typed. When - provided, only these keys override deploy YAML; argparse defaults - stay subordinate to YAML. ``None`` means treat every kwarg as - explicit (programmatic ``Omni()`` calls). Returns: List of stage configuration dictionaries """ + _warn_deprecated_explicit_keys(deprecated_kwargs) + if base_engine_args is None: base_engine_args = {} cli_overrides = _convert_dataclasses_to_dict(dict(base_engine_args)) - # Per-stage JSON overrides are always explicit (the user typed --stage-overrides). - explicit = set(cli_explicit_keys) if cli_explicit_keys is not None else None if stage_overrides: for stage_id_str, overrides in stage_overrides.items(): for key, val in overrides.items(): - stage_key = f"stage_{stage_id_str}_{key}" - cli_overrides[stage_key] = val - if explicit is not None: - explicit.add(stage_key) + cli_overrides[f"stage_{stage_id_str}_{key}"] = val stages = StageConfigFactory.create_from_model( model, cli_overrides=cli_overrides, deploy_config_path=deploy_config_path, - cli_explicit_keys=explicit, ) if stages is not None: # Convert StageConfig objects to OmegaConf for backward compat @@ -545,7 +553,7 @@ def load_and_resolve_stage_configs( default_stage_cfg_factory: Any = None, deploy_config_path: str | None = None, stage_overrides: dict[str, dict[str, Any]] | None = None, - cli_explicit_keys: set[str] | None = None, + **deprecated_kwargs: Any, ) -> tuple[str, list]: """Load stage configurations from model or YAML file with fallback to defaults. @@ -589,6 +597,8 @@ def load_and_resolve_stage_configs( stage_configs_path, ) + _warn_deprecated_explicit_keys(deprecated_kwargs) + if deploy_config_path is not None: config_path = deploy_config_path stage_configs = load_stage_configs_from_model( @@ -596,7 +606,6 @@ def load_and_resolve_stage_configs( base_engine_args=kwargs, deploy_config_path=deploy_config_path, stage_overrides=stage_overrides, - cli_explicit_keys=cli_explicit_keys, ) if not stage_configs: if default_stage_cfg_factory is not None: @@ -610,7 +619,6 @@ def load_and_resolve_stage_configs( model, base_engine_args=kwargs, stage_overrides=stage_overrides, - cli_explicit_keys=cli_explicit_keys, ) if not stage_configs: if default_stage_cfg_factory is not None: diff --git a/vllm_omni/model_executor/models/bagel/bagel.py b/vllm_omni/model_executor/models/bagel/bagel.py index cbb775680..3714836f2 100644 --- a/vllm_omni/model_executor/models/bagel/bagel.py +++ b/vllm_omni/model_executor/models/bagel/bagel.py @@ -203,6 +203,13 @@ class OmniBagelDataParser(MultiModalDataParser): class OmniBagelMultiModalProcessor(BaseMultiModalProcessor[OmniBagelProcessingInfo]): IMG2IMG_PLACEHOLDER = "<|fim_middle|>" + @staticmethod + def _mm_kwargs_for_bagel_img2img_hf(mm_kwargs: Mapping[str, object]) -> dict[str, object]: + # OpenAI / GLM-style serving may pass target_h/target_w for output grid sizing. + # BagelProcessor does not accept these in img2img mode; strip here so callers + # (e.g. serving_chat) can stay model-agnostic. + return {k: v for k, v in mm_kwargs.items() if k not in ("target_h", "target_w")} + def _cached_apply_hf_processor(self, inputs, timing_ctx): # img2img: prompt text must be modified based on mm data presence, # so text and mm data cannot be tokenized separately — bypass cache. @@ -248,7 +255,7 @@ class OmniBagelMultiModalProcessor(BaseMultiModalProcessor[OmniBagelProcessingIn if "images" in img2img_data: del img2img_data["images"] img2img_data["images"] = img2img_data.pop("pixel_values_img2img") - kwargs_img2img = dict(mm_kwargs) + kwargs_img2img = self._mm_kwargs_for_bagel_img2img_hf(mm_kwargs) kwargs_img2img["is_img2img"] = True out_img2img = super()._call_hf_processor(prompt, img2img_data, kwargs_img2img, tok_kwargs) if "pixel_values" in out_img2img: @@ -262,7 +269,7 @@ class OmniBagelMultiModalProcessor(BaseMultiModalProcessor[OmniBagelProcessingIn elif has_img2img: mm_data = dict(mm_data) mm_data["images"] = mm_data.pop("pixel_values_img2img") - mm_kwargs = dict(mm_kwargs) + mm_kwargs = self._mm_kwargs_for_bagel_img2img_hf(mm_kwargs) mm_kwargs["is_img2img"] = True outputs = super()._call_hf_processor(prompt, mm_data, mm_kwargs, tok_kwargs) if "pixel_values" in outputs: @@ -581,6 +588,9 @@ class OmniBagelForConditionalGeneration(BagelForConditionalGeneration): existing = self._ropes_metadata.get(rid) if existing and "image_shape" in existing and "image_shape" not in meta: continue + ropes = meta.get("ropes") + if ropes: + meta["ropes"] = [int(r.item()) if isinstance(r, torch.Tensor) else r for r in ropes] self._ropes_metadata[rid] = meta def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict: @@ -724,7 +734,7 @@ class OmniBagelForConditionalGeneration(BagelForConditionalGeneration): positions = self._adjust_positions_for_img2img(positions, input_ids) use_mot = True else: - rope = int(positions[seq_len - 1].item()) + 1 + rope = positions[seq_len - 1] + 1 self._ropes_pending.append({"ropes": [rope]}) if use_mot: diff --git a/vllm_omni/model_executor/models/bagel/pipeline.py b/vllm_omni/model_executor/models/bagel/pipeline.py new file mode 100644 index 000000000..c68a531c2 --- /dev/null +++ b/vllm_omni/model_executor/models/bagel/pipeline.py @@ -0,0 +1,111 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""BAGEL-7B-MoT pipeline topologies (frozen). + +Two-stage (default): + Stage 0: Thinker — multimodal understanding + text generation (AR) + Stage 1: DiT — diffusion image generation + +Two-stage think: + Same as two-stage but the Thinker decodes ... tokens before + KV transfer. Uses expand_cfg_prompts_think (companion max_tokens=1) and + omits kv_transfer_criteria so transfer happens after EOS, not after prefill. + +Single-stage: + Stage 0: DiT — self-contained diffusion stage that handles all modalities + (text2img, img2img, img2text, text2text, think) internally via its + own LLM, ViT, VAE, and tokenizer. +""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +_PROC = "vllm_omni.model_executor.stage_input_processors.bagel" + +BAGEL_PIPELINE = PipelineConfig( + model_type="bagel", + model_arch="OmniBagelForConditionalGeneration", + hf_architectures=("BagelForConditionalGeneration",), + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="thinker", + execution_type=StageExecutionType.LLM_AR, + input_sources=(), + final_output=True, + final_output_type="text", + owns_tokenizer=True, + requires_multimodal_data=True, + model_arch="OmniBagelForConditionalGeneration", + engine_output_type="text", + prompt_expand_func=f"{_PROC}.expand_cfg_prompts", + omni_kv_config={ + "need_send_cache": True, + "kv_transfer_criteria": {"type": "prefill_finished"}, + }, + sampling_constraints={"detokenize": True}, + ), + StagePipelineConfig( + stage_id=1, + model_stage="dit", + execution_type=StageExecutionType.DIFFUSION, + input_sources=(0,), + final_output=True, + final_output_type="image", + cfg_kv_collect_func=f"{_PROC}.collect_cfg_kv_caches", + omni_kv_config={"need_recv_cache": True}, + ), + ), +) + +BAGEL_THINK_PIPELINE = PipelineConfig( + model_type="bagel_think", + model_arch="OmniBagelForConditionalGeneration", + hf_architectures=(), + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="thinker", + execution_type=StageExecutionType.LLM_AR, + input_sources=(), + final_output=True, + final_output_type="text", + owns_tokenizer=True, + requires_multimodal_data=True, + model_arch="OmniBagelForConditionalGeneration", + engine_output_type="text", + prompt_expand_func=f"{_PROC}.expand_cfg_prompts_think", + omni_kv_config={"need_send_cache": True}, + sampling_constraints={"detokenize": True}, + ), + StagePipelineConfig( + stage_id=1, + model_stage="dit", + execution_type=StageExecutionType.DIFFUSION, + input_sources=(0,), + final_output=True, + final_output_type="image", + cfg_kv_collect_func=f"{_PROC}.collect_cfg_kv_caches", + omni_kv_config={"need_recv_cache": True}, + ), + ), +) + +BAGEL_SINGLE_STAGE_PIPELINE = PipelineConfig( + model_type="bagel_single_stage", + model_arch="BagelForConditionalGeneration", + hf_architectures=(), + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="dit", + execution_type=StageExecutionType.DIFFUSION, + input_sources=(), + final_output=True, + final_output_type="image", + ), + ), +) diff --git a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py index 3a904442f..ecfeec381 100644 --- a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py +++ b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py @@ -122,6 +122,9 @@ class CodePredictorAttention(nn.Module): super().__init__() self.num_heads = config.num_attention_heads self.num_kv_heads = config.num_key_value_heads + assert self.num_heads % self.num_kv_heads == 0 + self.is_gqa = self.num_kv_heads != self.num_heads + self.num_queries_per_kv = self.num_heads // self.num_kv_heads self.head_dim = getattr( config, "head_dim", @@ -129,7 +132,7 @@ class CodePredictorAttention(nn.Module): ) self.hidden_size = config.hidden_size self.scaling = self.head_dim**-0.5 - self._use_gqa = self.num_kv_heads != self.num_heads + self.max_seq = int(config.num_code_groups) + 1 # Separate q/k/v projections matching HF (no fused packing) bias = getattr(config, "attention_bias", False) @@ -140,6 +143,73 @@ class CodePredictorAttention(nn.Module): self.q_norm = _RMSNorm(self.head_dim, eps=config.rms_norm_eps) self.k_norm = _RMSNorm(self.head_dim, eps=config.rms_norm_eps) + if current_omni_platform.is_npu(): + if self.max_seq > 2048: + raise ValueError( + "Qwen3-TTS code predictor NPU fusion attention uses a fixed 2048x2048 " + f"causal mask, but max_seq={self.max_seq} exceeds the mask size." + ) + # Ascend SDPA is_causal migration example uses a fixed 2048x2048 + # compressed causal mask with sparse_mode=2. + fusion_mask = torch.triu( + torch.ones(2048, 2048, dtype=torch.bool), + diagonal=1, + ) + self.register_buffer("_fusion_causal_mask", fusion_mask, persistent=False) + + def _forward_npu_attention( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + bsz: int, + seq_len: int, + ) -> torch.Tensor: + import torch_npu + + q_f, k_f, v_f = q, k, v + if self.is_gqa: + k_f = ( + k[:, :, None, :, :] + .expand(bsz, self.num_kv_heads, self.num_queries_per_kv, seq_len, self.head_dim) + .reshape(bsz, self.num_heads, seq_len, self.head_dim) + ) + v_f = ( + v[:, :, None, :, :] + .expand(bsz, self.num_kv_heads, self.num_queries_per_kv, seq_len, self.head_dim) + .reshape(bsz, self.num_heads, seq_len, self.head_dim) + ) + + mask = self._fusion_causal_mask + mask = mask.contiguous() + q_f = q_f.contiguous() + k_f = k_f.contiguous() + v_f = v_f.contiguous() + return torch_npu.npu_fusion_attention( + q_f, + k_f, + v_f, + self.num_heads, + "BNSD", + pse=None, + padding_mask=None, + atten_mask=mask, + scale=float(self.scaling), + keep_prob=1.0, + # Keep torch_npu's API spelling. + pre_tockens=2147483647, + next_tockens=2147483647, + inner_precise=0, + prefix=None, + actual_seq_qlen=None, + actual_seq_kvlen=None, + # Ascend SDPA is_causal migration example uses sparse_mode=2. + sparse_mode=2, + gen_mask_parallel=True, + # Keep sync=True for the NPU fused attention path. + sync=True, + )[0] + def forward( self, hidden_states: torch.Tensor, @@ -160,14 +230,17 @@ class CodePredictorAttention(nn.Module): q = (q * cos) + (_rotate_half(q) * sin) k = (k * cos) + (_rotate_half(k) * sin) - attn_out = F.scaled_dot_product_attention( - q, - k, - v, - scale=self.scaling, - is_causal=True, - enable_gqa=self._use_gqa, - ) + if not current_omni_platform.is_npu(): + attn_out = F.scaled_dot_product_attention( + q, + k, + v, + scale=self.scaling, + is_causal=True, + enable_gqa=self.is_gqa, + ) + else: + attn_out = self._forward_npu_attention(q, k, v, bsz, seq_len) attn_out = attn_out.transpose(1, 2).reshape(bsz, seq_len, -1) return self.o_proj(attn_out) @@ -384,7 +457,7 @@ class CodePredictorWrapper(nn.Module): self._bucket_pos_ids: dict[int, torch.Tensor] = {} self._lm_heads_list: list[nn.Module] | None = None self._codec_embeds_list: list[nn.Module] | None = None - self._cuda_graphs: dict[int, tuple[torch.cuda.CUDAGraph, torch.Tensor]] = {} + self._device_graphs: dict[int, tuple] = {} # (graph, static_output) per bucket def get_input_embeddings(self) -> nn.ModuleList: return self.model.get_input_embeddings() @@ -412,7 +485,7 @@ class CodePredictorWrapper(nn.Module): self._proj_buf = torch.zeros(bsz, max_seq, self._cp_hidden, dtype=dtype, device=device) def _setup_compile(self) -> None: - """Lazily set up torch.compile with optional CUDA graph capture.""" + """Lazily set up torch.compile with optional device graph capture.""" if self._compiled_model_fwd is not None: return @@ -424,8 +497,16 @@ class CodePredictorWrapper(nn.Module): self._codec_embeds_list = list(self.model.codec_embedding) if not current_omni_platform.supports_torch_inductor(): - logger.warning_once("code_predictor: torch.compile disabled") + # NPU or other platforms without Inductor support self._compiled_model_fwd = self.model.forward + + if current_omni_platform.is_npu() and self._wrapper_config.use_cuda_graphs: + # For NPU, use eager + NPU graphs (no torch.compile) + self._warmup_buckets() + self._capture_npu_graphs() + logger.info("code_predictor: eager mode + NPU graphs") + else: + logger.warning_once("code_predictor: torch.compile disabled") return # torch.compile fuses RMSNorm/RoPE in ways that lose float32 @@ -491,10 +572,28 @@ class CodePredictorWrapper(nn.Module): with torch.cuda.graph(g, pool=pool): static_output = self._compiled_model_fwd(static_input, pos_ids) - self._cuda_graphs[bsz] = (g, static_output) + self._device_graphs[bsz] = (g, static_output) logger.info("code_predictor: captured CUDA graphs for buckets %s", self._bucket_sizes) + def _capture_npu_graphs(self) -> None: + """Capture an NPU graph per bucket using torch_npu's NPUGraph.""" + max_seq = self._num_groups + 1 + proj_buf = self._proj_buf + pool = torch.npu.graph_pool_handle() + + for bsz in self._bucket_sizes: + static_input = proj_buf[:bsz, :max_seq, :] + pos_ids = self._bucket_pos_ids[bsz] + + g = torch.npu.NPUGraph() + with torch.npu.graph(g, pool=pool): + static_output = self._compiled_model_fwd(static_input, pos_ids) + + self._device_graphs[bsz] = (g, static_output) + + logger.info("code_predictor: captured NPU graphs for buckets %s", self._bucket_sizes) + # ------------------------------------------------------------------ # Forward -- re-prefill + inline sampling # ------------------------------------------------------------------ @@ -544,8 +643,8 @@ class CodePredictorWrapper(nn.Module): torch.arange(max_seq, device=device, dtype=torch.long).unsqueeze(0).expand(padded_bsz, -1).contiguous() ) - # Use captured CUDA graph if available, otherwise call compiled fn. - cuda_graph_entry = self._cuda_graphs.get(padded_bsz) + # Use captured device graph if available, otherwise call compiled fn. + device_graph_entry = self._device_graphs.get(padded_bsz) # Prepare sampling parameters stored_mode = self._wrapper_config.sampling_mode == "stored" @@ -570,10 +669,10 @@ class CodePredictorWrapper(nn.Module): # Autoregressive loop: predict layers 1..G-1 for step in range(1, num_groups): - # Run transformer (CUDA graph replay or compiled forward) - if cuda_graph_entry is not None: - cuda_graph_entry[0].replay() - hidden_out = cuda_graph_entry[1] + # Run transformer (device graph replay or compiled forward) + if device_graph_entry is not None: + device_graph_entry[0].replay() + hidden_out = device_graph_entry[1] else: hidden_out = model_fwd(proj_buf[:padded_bsz, :max_seq, :], full_pos_ids) diff --git a/vllm_omni/model_executor/models/glm_image/__init__.py b/vllm_omni/model_executor/models/glm_image/__init__.py index d37044c09..6153f07c0 100644 --- a/vllm_omni/model_executor/models/glm_image/__init__.py +++ b/vllm_omni/model_executor/models/glm_image/__init__.py @@ -1,3 +1,16 @@ -from .glm_image_ar import GlmImageForConditionalGeneration +def __getattr__(name: str): + """Lazy import to avoid importing transformers.models.glm_image at module init. + + The AR model depends on ``transformers.models.glm_image`` which is only + available when the model weights are present. Importing it eagerly in + ``__init__.py`` breaks the pipeline registry lookup for environments that + don't have this custom transformers extension installed. + """ + if name == "GlmImageForConditionalGeneration": + from .glm_image_ar import GlmImageForConditionalGeneration + + return GlmImageForConditionalGeneration + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + __all__ = ["GlmImageForConditionalGeneration"] diff --git a/vllm_omni/model_executor/models/glm_image/pipeline.py b/vllm_omni/model_executor/models/glm_image/pipeline.py new file mode 100644 index 000000000..bd041292c --- /dev/null +++ b/vllm_omni/model_executor/models/glm_image/pipeline.py @@ -0,0 +1,47 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""GLM-Image pipeline topologies (frozen). +Two-stage (default): + Stage 0: AR — multimodal understanding + token_ids generation + Stage 1: DiT — diffusion image generation +""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +GLM_IMAGE_PIPELINE = PipelineConfig( + model_type="glm_image", + model_arch="GlmImageForConditionalGeneration", + hf_architectures=("GlmImageForConditionalGeneration",), + diffusers_class_name="GlmImagePipeline", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="ar", + execution_type=StageExecutionType.LLM_AR, + requires_multimodal_data=True, + input_sources=(), + final_output=False, + owns_tokenizer=True, + model_arch="GlmImageForConditionalGeneration", + engine_output_type="token_ids", + model_subdir="vision_language_encoder", + tokenizer_subdir="processor", + ), + StagePipelineConfig( + stage_id=1, + model_stage="dit", + execution_type=StageExecutionType.DIFFUSION, + input_sources=(0,), + requires_multimodal_data=True, + final_output=True, + final_output_type="image", + model_arch="GlmImagePipeline", + custom_process_input_func="vllm_omni.model_executor.stage_input_processors.glm_image.ar2diffusion", + omni_kv_config={"need_recv_cache": False}, + ), + ), +) diff --git a/vllm_omni/model_executor/models/mimo_audio/pipeline.py b/vllm_omni/model_executor/models/mimo_audio/pipeline.py index 33e78a172..1cc15af50 100644 --- a/vllm_omni/model_executor/models/mimo_audio/pipeline.py +++ b/vllm_omni/model_executor/models/mimo_audio/pipeline.py @@ -25,7 +25,7 @@ MIMO_AUDIO_PIPELINE = PipelineConfig( # ``model_executor/models/registry.py``; it resolves to the internal # class ``MiMoAudioForConditionalGeneration`` in ``mimo_audio.py``. model_arch="MiMoAudioModel", - hf_architectures=("MiMoAudioModel",), + hf_architectures=("MiMoAudioModel", "MiMoV2ASRForCausalLM"), stages=( StagePipelineConfig( stage_id=0, diff --git a/vllm_omni/model_executor/models/ming_flash_omni/__init__.py b/vllm_omni/model_executor/models/ming_flash_omni/__init__.py index d7fa44fd7..4cd086e64 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/__init__.py +++ b/vllm_omni/model_executor/models/ming_flash_omni/__init__.py @@ -2,6 +2,7 @@ # Copyright 2025 The vLLM-Omni team. from .ming_flash_omni import MingFlashOmniForConditionalGeneration +from .ming_flash_omni_talker import MingFlashOmniTalkerForConditionalGeneration from .ming_flash_omni_thinker import ( MingFlashOmniThinkerDummyInputsBuilder, MingFlashOmniThinkerForConditionalGeneration, @@ -11,6 +12,7 @@ from .ming_flash_omni_thinker import ( __all__ = [ "MingFlashOmniForConditionalGeneration", + "MingFlashOmniTalkerForConditionalGeneration", "MingFlashOmniThinkerForConditionalGeneration", "MingFlashOmniThinkerProcessingInfo", "MingFlashOmniThinkerMultiModalProcessor", diff --git a/vllm_omni/model_executor/models/ming_flash_omni/audio_vae.py b/vllm_omni/model_executor/models/ming_flash_omni/audio_vae.py new file mode 100644 index 000000000..9d5c266b4 --- /dev/null +++ b/vllm_omni/model_executor/models/ming_flash_omni/audio_vae.py @@ -0,0 +1,390 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2025 The vLLM-Omni team. +# Copyright (c) Ant Group. All rights reserved. +# Adapted from: +# https://github.com/inclusionAI/Ming/tree/e58533db227031990c5a6864dcf5f08fb53ed0d2/AudioVAE + +from __future__ import annotations + +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers import PretrainedConfig, PreTrainedModel, Qwen2Config, Qwen2Model +from vllm.logger import init_logger + +logger = init_logger(__name__) +try: + import flash_attn # noqa: F401 +except (ImportError, ModuleNotFoundError): + flash_attn = None + logger.warning( + "flash_attn is not available, the model may not yield the " + "exactly same result as the transformers implementation " + "in the audio tower part." + ) + + +class AudioVAEConfig(PretrainedConfig): + model_type = "audio_vae" + + def __init__( + self, + sample_rate: int = 44100, + enc_kwargs: dict | None = None, + dec_kwargs: dict | None = None, + init_method: str = "kaiming", + patch_size: int = 4, + **kwargs, + ): + self.sample_rate = sample_rate + self.enc_kwargs = enc_kwargs or {} + self.dec_kwargs = dec_kwargs or {} + self.init_method = init_method + self.patch_size = patch_size + super().__init__(**kwargs) + + +class ISTFT(nn.Module): + def __init__(self, n_fft: int, hop_length: int, win_length: int, padding: str = "same"): + super().__init__() + if padding not in ["center", "same"]: + raise ValueError("Padding must be 'center' or 'same'.") + self.padding = padding + self.n_fft = n_fft + self.hop_length = hop_length + self.win_length = win_length + window = torch.hann_window(win_length) + self.register_buffer("window", window) + self.buffer_len = self.win_length - self.hop_length + + def _buffer_process(self, x, buffer, pad, last_chunk=False, streaming=False): + if streaming: + if buffer is None: + x = x[:, pad:] + if buffer is not None: + x[:, : self.buffer_len] += buffer + buffer = x[:, -self.buffer_len :] + if not last_chunk: + x = x[:, : -self.buffer_len] + else: + x = x[:, :-pad] + else: + x = x[:, pad:-pad] + return x, buffer + + def forward(self, spec, audio_buffer=None, window_buffer=None, streaming=False, last_chunk=False): + if self.padding == "center": + return torch.istft(spec, self.n_fft, self.hop_length, self.win_length, self.window, center=True) + elif self.padding == "same": + pad = (self.win_length - self.hop_length) // 2 + else: + raise ValueError("Padding must be 'center' or 'same'.") + + B, N, T = spec.shape + ifft = torch.fft.irfft(spec, self.n_fft, dim=1, norm="backward") + ifft = ifft * self.window[None, :, None] + + output_size = (T - 1) * self.hop_length + self.win_length + y = torch.nn.functional.fold( + ifft, + output_size=(1, output_size), + kernel_size=(1, self.win_length), + stride=(1, self.hop_length), + )[:, 0, 0, :] + + y, audio_buffer = self._buffer_process(y, audio_buffer, pad, last_chunk=last_chunk, streaming=streaming) + + window_sq = self.window.square().expand(1, T, -1).transpose(1, 2) + window_envelope = ( + torch.nn.functional.fold( + window_sq, + output_size=(1, output_size), + kernel_size=(1, self.win_length), + stride=(1, self.hop_length), + ) + .squeeze(0) + .squeeze(0) + ) + + window_envelope, window_buffer = self._buffer_process( + window_envelope, window_buffer, pad, last_chunk=last_chunk, streaming=streaming + ) + window_envelope = window_envelope.squeeze() + + assert (window_envelope > 1e-11).all() + y = y / window_envelope + + return y, audio_buffer, window_buffer + + +class ISTFTHead(nn.Module): + def __init__(self, dim: int, n_fft: int, hop_length: int, padding: str = "same"): + super().__init__() + out_dim = n_fft + 2 + self.out = nn.Linear(dim, out_dim) + self.istft = ISTFT(n_fft=n_fft, hop_length=hop_length, win_length=n_fft, padding=padding) + + def forward(self, x, audio_buffer=None, window_buffer=None, streaming=False, last_chunk=False): + x_pred = self.out(x) + x_pred = x_pred.transpose(1, 2) + mag, p = x_pred.chunk(2, dim=1) + mag = torch.exp(mag) + mag = torch.clip(mag, max=1e2) + x = torch.cos(p) + y = torch.sin(p) + S = mag * (x + 1j * y) + audio, audio_buffer, window_buffer = self.istft( + S, audio_buffer=audio_buffer, window_buffer=window_buffer, streaming=streaming, last_chunk=last_chunk + ) + return audio.unsqueeze(1), x_pred, audio_buffer, window_buffer + + +class StreamingLinearUpsample(nn.Module): + def __init__(self, scale_factor=4): + super().__init__() + self.scale_factor = scale_factor + self.upsampler = nn.Upsample(scale_factor=scale_factor, mode="linear", align_corners=False) + + def forward(self, x, state=None, is_last=False): + if state is None: + state = {"prev_chunk": None, "history_last": None, "is_first": True} + + if x is None and not is_last: + return None, state + + if state["is_first"] and is_last: + out = self.upsampler(x.transpose(1, 2)).transpose(1, 2) + return out, None + + output_chunks = [] + + if state["is_first"]: + state["prev_chunk"] = x + state["is_first"] = False + if not is_last: + return None, state + + if state["prev_chunk"] is not None: + p = state["prev_chunk"].transpose(1, 2) + + if state["history_last"] is None: + lookahead = x[:, :1, :].transpose(1, 2) + inp = torch.cat([p, lookahead], dim=2) + up = self.upsampler(inp) + out_prev = up[:, :, : p.size(2) * self.scale_factor] + else: + lookahead = x[:, :1, :].transpose(1, 2) + inp = torch.cat([state["history_last"], p, lookahead], dim=2) + up = self.upsampler(inp) + start = self.scale_factor + end = start + p.size(2) * self.scale_factor + out_prev = up[:, :, start:end] + + output_chunks.append(out_prev.transpose(1, 2)) + state["history_last"] = p[:, :, -1:] + state["prev_chunk"] = x + + if is_last: + p = state["prev_chunk"].transpose(1, 2) + inp = torch.cat([state["history_last"], p], dim=2) + up = self.upsampler(inp) + out_last = up[:, :, self.scale_factor :] + output_chunks.append(out_last.transpose(1, 2)) + state = None + + final_out = torch.cat(output_chunks, dim=1) if output_chunks else None + return final_out, state + + +class Decoder(nn.Module): + def __init__(self, decoder_args, output_dim=320, latent_dim=64, patch_size=-1): + super().__init__() + config = Qwen2Config.from_dict(config_dict=decoder_args) + if flash_attn is None: + config._attn_implementation = "sdpa" + self.decoder = Qwen2Model(config) + self.output_dim = output_dim + self.latent_dim = latent_dim + self.fc1 = nn.Linear(latent_dim, config.hidden_size) + self.hop_length = output_dim + self.head = ISTFTHead( + dim=config.hidden_size, n_fft=self.hop_length * 4, hop_length=self.hop_length, padding="same" + ) + self.patch_size = patch_size + if self.patch_size != -1: + self.upsampling = StreamingLinearUpsample(scale_factor=patch_size) + + def low_level_reconstruct(self, x, past_key_values=None, use_cache=False, stream_state=None, last_chunk=False): + upsample_state, audio_buffer, window_buffer = stream_state + bsz, device, dtype = x.size(0), x.device, x.dtype + x = self.fc1(x) + if self.patch_size != -1: + if use_cache: + x, upsample_state = self.upsampling(x, state=upsample_state, is_last=last_chunk) + if x is None: + stream_state = (upsample_state, audio_buffer, window_buffer) + return torch.empty(bsz, 1, 0, device=device, dtype=dtype), stream_state, past_key_values + else: + x = self.upsampling.upsampler(x.transpose(1, 2)).transpose(1, 2) + + hidden_states_list = [] + + if use_cache and getattr(self.decoder.config, "sliding_window", None) is not None: + sw_size = self.decoder.config.sliding_window + target_len = sw_size - 1 + if past_key_values is None: + past_len = 0 + elif hasattr(past_key_values, "get_seq_length"): + past_len = past_key_values.get_seq_length() + elif isinstance(past_key_values, tuple) and len(past_key_values) > 0: + past_len = past_key_values[0][0].shape[-2] + else: + past_len = 0 + + curr_len = x.shape[1] + + if past_len < target_len and (past_len + curr_len) >= sw_size: + fill_len = target_len - past_len + x_fill = x[:, :fill_len, :] + outputs = self.decoder(inputs_embeds=x_fill, past_key_values=past_key_values, use_cache=use_cache) + hidden_states_list.append(outputs.last_hidden_state) + past_key_values = outputs.past_key_values + x = x[:, fill_len:, :] + + outputs = self.decoder(inputs_embeds=x, past_key_values=past_key_values, use_cache=use_cache) + hidden_states_list.append(outputs.last_hidden_state) + past_key_values = outputs.past_key_values + + if len(hidden_states_list) > 1: + full_hidden_state = torch.cat(hidden_states_list, dim=1) + else: + full_hidden_state = hidden_states_list[0] + + x_out, _, audio_buffer, window_buffer = self.head( + full_hidden_state, + streaming=use_cache, + audio_buffer=audio_buffer, + window_buffer=window_buffer, + last_chunk=last_chunk, + ) + + stream_state = (upsample_state, audio_buffer, window_buffer) + return x_out, stream_state, past_key_values + + +class Encoder(nn.Module): + def __init__(self, encoder_args, input_dim=320, hop_size=320, latent_dim=64, patch_size=-1): + super().__init__() + config = Qwen2Config.from_dict(config_dict=encoder_args) + if flash_attn is None: + config._attn_implementation = "sdpa" + self.encoder = Qwen2Model(config) + self.input_dim = input_dim + self.hop_size = hop_size + self.latent_dim = latent_dim + self.fc1 = nn.Linear(input_dim, config.hidden_size, bias=False) + self.fc2 = nn.Linear(config.hidden_size, config.hidden_size) + self.fc3 = nn.Linear(config.hidden_size, latent_dim * 2) + self.norm = nn.LayerNorm(config.hidden_size) + self.patch_size = patch_size + if patch_size != -1: + config.num_hidden_layers = 4 + self.aggregator = Qwen2Model(config) + self.cls_embed = nn.Parameter(torch.rand(1, 1, config.hidden_size)) + self.cls_embed.data.normal_(0, 0.02) + + def get_frames(self, x): + num_frames_total = (x.size(-1) + self.hop_size - 1) // self.hop_size + expected_len = (num_frames_total - 1) * self.hop_size + self.input_dim + padding_needed = expected_len - x.size(-1) + waveform = F.pad(x, (0, padding_needed), value=0.0) + frames = waveform.unfold(dimension=-1, size=self.input_dim, step=self.hop_size) + return frames + + def pad_patch_insert_cls(self, x): + bsz, _, dim = x.size() + num_frame = x.size(1) + r = num_frame % self.patch_size + pad_num = self.patch_size - r if r else 0 + x = F.pad(x, (0, 0, 0, pad_num), value=0.0) + x = x.reshape(-1, self.patch_size, dim) + x = torch.cat((x, self.cls_embed.expand(x.size(0), -1, -1)), dim=1) + x = x.reshape(bsz, -1, dim) + return x + + def forward(self, waveform): + x = self.get_frames(waveform) + x = self.fc1(x) + x = self.fc2(x) + x = self.encoder(inputs_embeds=x) + x = x.last_hidden_state + + if self.patch_size != -1: + x = self.pad_patch_insert_cls(x) + x = self.aggregator(inputs_embeds=x) + x = x.last_hidden_state + bsz, _, dim = x.size() + x = x.reshape(-1, self.patch_size + 1, dim) + x = x[:, -1:, :].reshape(bsz, -1, dim) + + x = self.fc3(x) + return x, waveform.unsqueeze(1) + + +class AudioVAE(PreTrainedModel): + config_class = AudioVAEConfig + + def __init__(self, config: AudioVAEConfig): + super().__init__(config) + self.encoder = Encoder( + encoder_args=config.enc_kwargs["backbone"], + input_dim=config.enc_kwargs["input_dim"], + hop_size=config.enc_kwargs.get("hop_size", 320), + latent_dim=config.enc_kwargs["latent_dim"], + patch_size=config.patch_size, + ) + self.decoder = Decoder( + decoder_args=config.dec_kwargs["backbone"], + output_dim=config.dec_kwargs["output_dim"], + latent_dim=config.dec_kwargs["latent_dim"], + patch_size=config.patch_size, + ) + self.post_init() + + def _init_weights(self, module): + std = 0.02 + if isinstance(module, nn.Linear): + if self.config.init_method == "kaiming": + nn.init.kaiming_normal_(module.weight, mode="fan_in", nonlinearity="relu") + else: + module.weight.data.normal_(mean=0.0, std=std) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.Embedding): + module.weight.data.normal_(mean=0.0, std=std) + if module.padding_idx is not None: + module.weight.data[module.padding_idx].zero_() + + def encode_latent(self, waveform, waveform_length): + from diffusers.models.autoencoders.autoencoder_oobleck import OobleckDiagonalGaussianDistribution + + frame_num = torch.ceil(waveform_length / self.config.enc_kwargs["input_dim"]).to(torch.int32) + if self.config.patch_size != -1: + frame_num = torch.ceil(frame_num / self.config.patch_size) + h, y = self.encoder(waveform) + h = h.transpose(1, 2) + + posterior = OobleckDiagonalGaussianDistribution(h) + latent = posterior.sample() + latent = latent.transpose(1, 2) + return latent, frame_num + + def decode(self, latent, past_key_values=None, use_cache=False, stream_state=(None, None, None), last_chunk=False): + waveform, stream_state, past_key_values = self.decoder.low_level_reconstruct( + latent, + past_key_values=past_key_values, + use_cache=use_cache, + stream_state=stream_state, + last_chunk=last_chunk, + ) + return waveform, stream_state, past_key_values diff --git a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni.py b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni.py index 87728890b..462c20438 100644 --- a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni.py +++ b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni.py @@ -16,7 +16,23 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Ming-flash-omni-2.0 unified model (thinker + imagegen + talker).""" +"""Ming-flash-omni-2.0 thinker / image-gen wrapper. + +This class is the multimodal-registered entry point for Ming stages that +share the thinker's backbone: comprehension / text generation (`thinker`) +and diffusion conditioning for image generation (`imagegen`, not yet +implemented). + +The talker deliberately lives elsewhere. Upstream Ming hands text (not +hidden states) from the thinker to the talker, and the talker then +tokenises that string with its own Qwen2 tokenizer and runs an entirely +self-contained LLM + CFM + AudioVAE pipeline. Because it has no +multimodal inputs, it belongs in the non-MM-registered +`MingFlashOmniTalkerForConditionalGeneration` — routing it through +this wrapper would force it through vLLM's multimodal preprocess path +and trigger a hidden-size mismatch between the outer Ming config +(4096, thinker's LLM) and the talker's Qwen2 backbone (896). +""" from collections.abc import Iterable @@ -40,7 +56,10 @@ from vllm.sequence import IntermediateTensors from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights -from vllm_omni.transformers_utils.configs.ming_flash_omni import BailingMM2Config, MingFlashOmniConfig +from vllm_omni.transformers_utils.configs.ming_flash_omni import ( + BailingMM2Config, + MingFlashOmniConfig, +) from .ming_flash_omni_thinker import ( MingFlashOmniThinkerDummyInputsBuilder, @@ -63,7 +82,7 @@ class MingFlashOmniForConditionalGeneration( SupportsMRoPE, CustomProcessMixin, ): - """Unified Ming-flash-omni-2.0 model combining thinker, imagegen, and talker.""" + """Ming-flash-omni-2.0 thinker + image-gen wrapper.""" supports_multimodal = True requires_raw_input_tokens: bool = True @@ -75,19 +94,26 @@ class MingFlashOmniForConditionalGeneration( self.has_postprocess = False config = vllm_config.model_config.hf_config - - self.vllm_config = vllm_config self.config = config - - if isinstance(config, MingFlashOmniConfig): - thinker_config = config.thinker_config - else: - thinker_config = config - - self.thinker_config: BailingMM2Config = thinker_config self.model_stage = vllm_config.model_config.model_stage + if self.model_stage == "talker": + raise ValueError( + "MingFlashOmniForConditionalGeneration does not support " + "model_stage='talker'. Use " + "model_arch='MingFlashOmniTalkerForConditionalGeneration' " + "directly — the talker has a self-contained LLM that " + "tokenises text itself and does not need the multimodal " + "preprocess path. See stage_configs/ming_flash_omni.yaml " + "stage 1 and stage_configs/ming_flash_omni_tts.yaml." + ) + if self.model_stage == "thinker": + if isinstance(config, MingFlashOmniConfig): + thinker_config: BailingMM2Config = config.thinker_config + else: + thinker_config: BailingMM2Config = config + thinker_vllm_config = vllm_config.with_hf_config( thinker_config, architectures=["MingFlashOmniThinkerForConditionalGeneration"] ) @@ -97,8 +123,7 @@ class MingFlashOmniForConditionalGeneration( architectures=["MingFlashOmniThinkerForConditionalGeneration"], ) self.model = self.thinker - self.imagegen = None - self.talker = None + self.make_empty_intermediate_tensors = self.thinker.make_empty_intermediate_tensors elif self.model_stage == "imagegen": # TODO: Implement image generator stage @@ -106,22 +131,12 @@ class MingFlashOmniForConditionalGeneration( "Image generation stage is not yet implemented. Please use model_stage='thinker' for now." ) - elif self.model_stage == "talker": - # TODO: Implement talker (TTS) stage - raise NotImplementedError( - "Talker (TTS) stage is not yet implemented. Please use model_stage='thinker' for now." - ) - else: raise ValueError( - f"Invalid model_stage: {self.model_stage}. Must be one of: 'thinker', 'imagegen', 'talker'" + f"Invalid model_stage: {self.model_stage!r}. Must be one of: 'thinker', 'imagegen'. " + f"For the talker stage, use MingFlashOmniTalkerForConditionalGeneration directly." ) - # Set up intermediate tensors - self.make_empty_intermediate_tensors = ( - self.thinker.make_empty_intermediate_tensors if self.model_stage == "thinker" else lambda: None - ) - def forward( self, input_ids: torch.Tensor, @@ -154,7 +169,7 @@ class MingFlashOmniForConditionalGeneration( ): if hasattr(self.model, "sample"): return self.model.sample(logits, sampling_metadata) - raise NotImplementedError("sample method not available on current stage") + return None def get_mrope_input_positions(self, *args, **kwargs): if hasattr(self.model, "get_mrope_input_positions"): @@ -162,36 +177,9 @@ class MingFlashOmniForConditionalGeneration( raise NotImplementedError("get_mrope_input_positions not available on current stage") def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loaded_weights = set() - thinker_weights = [] - imagegen_weights = [] - talker_weights = [] - - for name, value in weights: - if name.startswith("thinker."): - thinker_weights.append((name, value)) - elif name.startswith("imagegen."): - imagegen_weights.append((name, value)) - elif name.startswith("talker."): - talker_weights.append((name, value)) - else: - # Weights without prefix go to thinker by default - thinker_weights.append((name, value)) - - if self.model_stage == "thinker" and thinker_weights: - # Remove "thinker." prefix before loading - thinker_weights_stripped = [ - (name.replace("thinker.", "", 1) if name.startswith("thinker.") else name, value) - for name, value in thinker_weights - ] - thinker_loaded = self.thinker.load_weights(thinker_weights_stripped) - thinker_loaded = add_prefix_to_loaded_weights(thinker_loaded, "thinker") - loaded_weights.update(thinker_loaded) - - # TODO: Load imagegen weights when implemented - # TODO: Load talker weights when implemented - - return loaded_weights + stripped = ((name.removeprefix("thinker."), value) for name, value in weights) + thinker_loaded = self.thinker.load_weights(stripped) + return add_prefix_to_loaded_weights(thinker_loaded, "thinker") def get_mm_mapping(self) -> MultiModelKeys: return MultiModelKeys.from_string_field( @@ -202,9 +190,7 @@ class MingFlashOmniForConditionalGeneration( @property def sampler(self): - if hasattr(self.model, "sampler"): - return self.model.sampler - return None + return getattr(self.model, "sampler", None) def embed_input_ids( self, diff --git a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py new file mode 100644 index 000000000..08ed9e854 --- /dev/null +++ b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_talker.py @@ -0,0 +1,586 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2025 The vLLM-Omni team. +# Copyright (c) Ant Group. All rights reserved. +# Adapted from: +# https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/modeling_bailing_talker.py +"""Ming-flash-omni-2.0 talker (TTS) stage model.""" + +from __future__ import annotations + +import glob as glob_module +import os +from collections.abc import Iterable +from dataclasses import dataclass +from functools import cached_property +from typing import Any + +import torch +import torch.nn as nn +from safetensors.torch import load_file +from transformers import AutoTokenizer, Qwen2Config, Qwen2Model +from transformers.utils.hub import cached_file +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.model_executor.models.utils import AutoWeightsLoader +from vllm.sequence import IntermediateTensors + +from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin +from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific +from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.transformers_utils.configs.ming_flash_omni import MingFlashOmniTalkerConfig + +from .audio_vae import AudioVAE, AudioVAEConfig +from .prompt_utils import DEFAULT_PROMPT as MING_DEFAULT_PROMPT +from .talker_module import CFM, Aggregator, DiT, MingAudioGenerator, build_tts_input +from .text_processing import segment_and_normalize +from .voice_presets import VoicePresetRegistry + +logger = init_logger(__name__) + + +@dataclass(slots=True) +class _GenerationParams: + """Resolved sampling / decoding parameters for one forward call.""" + + prompt: str + instruction: str | None + cfg: float + sigma: float + temperature: float + max_steps: int + use_zero_spk_emb: bool + max_text_length: int + use_static_cache: bool + stream_decode: bool + + +@dataclass(slots=True) +class _VoiceContext: + """Voice cloning inputs resolved from request info + presets.""" + + spk_emb: Any # list[Tensor] | Tensor | list[float] | None + prompt_text: str | None + prompt_wav_lat: torch.Tensor | None + prompt_wav_emb: torch.Tensor | None + already_projected: bool + + +class MingFlashOmniTalkerForConditionalGeneration(nn.Module, CustomProcessMixin): + """Ming-flash-omni-2.0 talker stage: text -> audio waveform. + + Uses Qwen2 LLM + CFM (Conditional Flow Matching with DiT) + Aggregator + in an autoregressive loop to produce continuous audio latents, then + AudioVAE decodes latents to waveforms. + """ + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.have_multimodal_outputs = True + self.has_preprocess = False + self.has_postprocess = False + + self.vllm_config = vllm_config + root_config = vllm_config.model_config.hf_config + + model_path = vllm_config.model_config.model + self._model_path = model_path + self.talker_dir = ( + os.path.join(model_path, "talker") if os.path.isdir(os.path.join(model_path, "talker")) else model_path + ) + + # When used standalone (model_arch=MingFlashOmniTalkerForConditionalGeneration), + # the root hf_config may be BailingMM2Config (thinker-only) due to model file structure + # Resolve talker config from talker/config.json in that case. + config = ( + root_config + if isinstance(root_config, MingFlashOmniTalkerConfig) + else self._resolve_talker_config(root_config, self.talker_dir, model_path) + ) + self.config = config + + self._standalone = prefix in ("", "talker") + if self._standalone: + self.allow_patterns_overrides = ["talker/model*.safetensors"] + self.fall_back_to_pt_during_load = False + + # LLM + llm_config = self._resolve_llm_config(config, self.talker_dir, model_path) + llm_config._attn_implementation = "sdpa" + self.llm_config = llm_config + self.hidden_size = llm_config.hidden_size + self.latent_dim = config.latent_dim + self.patch_size = config.patch_size + self.his_patch_size = config.history_patch_size + self.cfg_strength = config.cfg_strength + + self.model = Qwen2Model(llm_config) + self.cfm = CFM( + DiT(llm_input_dim=self.hidden_size, **config.flowmodel), + steps=config.steps, + ) + self.aggregator = Aggregator(llm_input_dim=self.hidden_size, **config.aggregator) + self.stop_head = nn.Linear(self.hidden_size, 2, bias=True) + # CAMPPlus 192-dim -> hidden + self.spk_head = nn.Linear(192, self.hidden_size, bias=True) + + # AudioVAE + self.audio_vae, self._vae_weight_source = self._init_audio_vae(config, self.talker_dir, model_path) + + self._use_cuda_graphs = not vllm_config.model_config.enforce_eager + + self.audio_generator = MingAudioGenerator( + config=self.config, + llm_config=self.llm_config, + model=self.model, + cfm=self.cfm, + aggregator=self.aggregator, + stop_head=self.stop_head, + audio_vae=self.audio_vae, + patch_size=self.patch_size, + his_patch_size=self.his_patch_size, + latent_dim=self.latent_dim, + cfg_strength=self.cfg_strength, + use_cuda_graphs=self._use_cuda_graphs, + ) + self.voice_presets = VoicePresetRegistry( + talker_dir=self.talker_dir, + model_path=self._model_path, + download_dir=vllm_config.load_config.download_dir, + audio_vae=self.audio_vae, + aggregator=self.aggregator, + spk_head=self.spk_head, + patch_size=self.patch_size, + ) + + @property + def device(self) -> torch.device: + return next(self.model.parameters()).device + + @property + def dtype(self) -> torch.dtype: + return next(self.model.parameters()).dtype + + @cached_property + def tokenizer(self): + # Lazy Qwen2 tokenizer resolution: + # 1. Try local dirs first (talker/llm, talker, and then model root). + # 2. HF repo-id fallback: talker/llm is the canonical tokenizer location. + candidates = (os.path.join(self.talker_dir, "llm"), self.talker_dir, self._model_path) + for path in candidates: + if os.path.isdir(path): + try: + logger.debug("Resolving talker tokenizer from local dir %s", path) + return AutoTokenizer.from_pretrained(path, trust_remote_code=True) + except Exception: + continue + for subfolder in ("talker/llm", "llm"): + try: + logger.debug("Resolving talker tokenizer from HF subfolder %s", subfolder) + return AutoTokenizer.from_pretrained(self._model_path, subfolder=subfolder, trust_remote_code=True) + except Exception: + continue + logger.debug("Falling back to raw model_path tokenizer resolution") + return AutoTokenizer.from_pretrained(self._model_path, trust_remote_code=True) + + @staticmethod + def _resolve_talker_config(config, talker_dir: str, model_path: str) -> MingFlashOmniTalkerConfig: + """Resolve MingFlashOmniTalkerConfig when the root config is not one. + + This happens in standalone TTS mode where hf_config is BailingMM2Config. + """ + # If the root config wraps a talker_config, use it + talker_config = getattr(config, "talker_config", None) + if isinstance(talker_config, MingFlashOmniTalkerConfig): + return talker_config + + # Try loading from talker/config.json + if os.path.isdir(talker_dir): + try: + resolved = MingFlashOmniTalkerConfig.from_pretrained(talker_dir) + logger.info("Resolved talker config from %s", talker_dir) + return resolved + except Exception: + pass + + try: + resolved = MingFlashOmniTalkerConfig.from_pretrained(model_path, subfolder="talker", trust_remote_code=True) + logger.info("Resolved talker config from %s/talker (HF hub)", model_path) + return resolved + except Exception as e: + raise ValueError( + f"Cannot resolve MingFlashOmniTalkerConfig. The root config " + f"is {type(config).__name__}, and talker/config.json was not " + f"found at {talker_dir} or via HF hub: {e}" + ) from e + + @staticmethod + def _resolve_llm_config(config: MingFlashOmniTalkerConfig, talker_dir: str, model_path: str) -> Qwen2Config: + """Resolve the Qwen2 LLM config for the talker backbone.""" + + if config.llm_config is not None: + return Qwen2Config(**config.llm_config) if isinstance(config.llm_config, dict) else config.llm_config + + # Try local talker/llm directory + llm_dir = os.path.join(talker_dir, "llm") + if os.path.isdir(llm_dir): + return Qwen2Config.from_pretrained(llm_dir) + + # HF hub fallback + for subfolder in ("talker/llm", "llm"): + try: + return Qwen2Config.from_pretrained(model_path, subfolder=subfolder, trust_remote_code=True) + except Exception: + continue + + raise ValueError( + f"Cannot find talker LLM config at {llm_dir}. " + "Either provide llm_config in MingFlashOmniTalkerConfig or " + "ensure the model path contains talker/llm/config.json." + ) + + @staticmethod + def _init_audio_vae( + config: MingFlashOmniTalkerConfig, talker_dir: str, model_path: str + ) -> tuple[AudioVAE | None, str | tuple[str, str] | None]: + """Initialize AudioVAE and return (vae, weight_source). + + weight_source is either a local directory path (str) or an + (repo_id, subfolder) tuple for HF hub downloads, or None. + """ + vae_path = config.audio_vae_path or os.path.join(talker_dir, "vae") + + # Try local directory first + if os.path.isdir(vae_path): + try: + vae_config = AudioVAEConfig.from_pretrained(vae_path) + vae = AudioVAE(vae_config) + logger.info("Initialized AudioVAE from %s (sr=%d)", vae_path, vae_config.sample_rate) + return vae, vae_path + except Exception as e: + logger.warning("Failed to initialize AudioVAE from %s: %s", vae_path, e) + return None, None + + # HF hub fallback + for subfolder in ("talker/vae", "vae"): + try: + vae_config = AudioVAEConfig.from_pretrained(model_path, subfolder=subfolder, trust_remote_code=True) + vae = AudioVAE(vae_config) + logger.info(f"Initialized AudioVAE from {model_path}/{subfolder}") + return vae, (model_path, subfolder) + except Exception: + continue + + logger.info("AudioVAE not found at %s; waveform decoding unavailable", vae_path) + return None, None + + def compute_logits(self, hidden_states: torch.Tensor, sampling_metadata=None) -> torch.Tensor | None: + return None + + def sample(self, logits: torch.Tensor, sampling_metadata): + return None + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings=None, + is_multimodal=None, + ) -> torch.Tensor: + return self.model.get_input_embeddings()(input_ids) + + def make_empty_intermediate_tensors( + self, batch_size: int, dtype: torch.dtype, device: torch.device + ) -> IntermediateTensors | None: + return None + + def get_dummy_runtime_additional_information(self, num_reqs: int) -> list[dict[str, object]]: + info: dict[str, object] = {"text": "dummy", "use_zero_spk_emb": True, "max_steps": 1} + return [info for _ in range(num_reqs)] + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + runtime_additional_information: list[dict] | None = None, + **kwargs, + ) -> OmniOutput: + """Run TTS generation and return audio output. + + The full autoregressive generation loop is executed inside this method. + """ + additional_info = self._extract_additional_info(runtime_additional_information) + params = self._resolve_generation_params(additional_info) + voice = self._resolve_voice(additional_info) + + latents = self._generate_latents( + input_ids=input_ids, + inputs_embeds=inputs_embeds, + text=additional_info.get("text", ""), + params=params, + voice=voice, + ) + return self._decode_to_output(latents, stream_decode=params.stream_decode) + + @staticmethod + def _extract_additional_info( + runtime_additional_information: list[dict] | None, + ) -> dict[str, Any]: + if runtime_additional_information and len(runtime_additional_information) > 0: + return runtime_additional_information[0] or {} + return {} + + def _resolve_generation_params(self, additional_info: dict[str, Any]) -> _GenerationParams: + # "omni" : thinker -> talker hand-off with hardcoded defaults + # "instruct": standalone TTS with caller-supplied sampling knobs + ming_task = additional_info.get("ming_task", "instruct") + + if ming_task == "omni": + prompt = MING_DEFAULT_PROMPT + instruction = None + use_zero_spk_emb = additional_info.get("spk_emb") is None + cfg = 2.0 + sigma = 0.25 + temperature = 0.0 + max_steps = 200 + else: + prompt = additional_info.get("prompt", MING_DEFAULT_PROMPT) + instruction = additional_info.get("instruction", None) + use_zero_spk_emb = additional_info.get("use_zero_spk_emb", False) + cfg = additional_info.get("cfg", self.cfg_strength) + sigma = additional_info.get("sigma", 0.25) + temperature = additional_info.get("temperature", 0.0) + max_steps = int(additional_info.get("max_steps", additional_info.get("max_decode_steps", 200))) + + return _GenerationParams( + prompt=prompt, + instruction=instruction, + cfg=cfg, + sigma=sigma, + temperature=temperature, + max_steps=max_steps, + use_zero_spk_emb=use_zero_spk_emb, + max_text_length=int(additional_info.get("max_text_length", 50)), + use_static_cache=bool(additional_info.get("use_static_cache", True)), + stream_decode=bool(additional_info.get("stream_decode", True)), + ) + + def _resolve_voice(self, additional_info: dict[str, Any]) -> _VoiceContext: + spk_emb = additional_info.get("spk_emb", None) + prompt_text = additional_info.get("prompt_text", None) + prompt_wav_lat = additional_info.get("prompt_wav_lat", None) + prompt_wav_emb = additional_info.get("prompt_wav_emb", None) + already_projected = False + + voice_name = additional_info.get("voice_name", None) + if voice_name and spk_emb is None and voice_name in self.voice_presets: + preset = self.voice_presets.get(voice_name) or {} + prompt_wav_lat = preset.get("prompt_wav_lat") + prompt_wav_emb = preset.get("prompt_wav_emb") + spk_emb = preset.get("spk_emb") + already_projected = True + if prompt_text is None: + prompt_text = preset.get("prompt_text") + + return _VoiceContext( + spk_emb=spk_emb, + prompt_text=prompt_text, + prompt_wav_lat=prompt_wav_lat, + prompt_wav_emb=prompt_wav_emb, + already_projected=already_projected, + ) + + def _project_spk_emb( + self, spk_emb: Any, already_projected: bool, use_zero_spk_emb: bool + ) -> list[torch.Tensor] | None: + if spk_emb is None: + if use_zero_spk_emb: + return [torch.zeros(1, self.hidden_size, device=self.device, dtype=self.dtype)] + return None + + if already_projected: + return spk_emb if isinstance(spk_emb, list) else [spk_emb] + + if isinstance(spk_emb, torch.Tensor): + tensors = [spk_emb] + elif isinstance(spk_emb, list) and spk_emb and isinstance(spk_emb[0], (int, float)): + tensors = [torch.tensor(spk_emb, dtype=self.dtype).unsqueeze(0)] + elif isinstance(spk_emb, list): + tensors = spk_emb + else: + tensors = [spk_emb] + return [self.spk_head(t.to(device=self.device, dtype=self.dtype)) for t in tensors] + + def _generate_latents( + self, + *, + input_ids: torch.Tensor, + inputs_embeds: torch.Tensor | None, + text: str, + params: _GenerationParams, + voice: _VoiceContext, + ) -> list[torch.Tensor]: + generator = self.audio_generator + + if inputs_embeds is not None: + # Caller pre-built embeddings — run a single AR pass. + return generator.generate_latents( + inputs_embeds=inputs_embeds, + prompt_wav_lat=voice.prompt_wav_lat, + max_steps=params.max_steps, + cfg=params.cfg, + sigma=params.sigma, + temperature=params.temperature, + use_static_cache=params.use_static_cache, + ) + + spk_emb = self._project_spk_emb(voice.spk_emb, voice.already_projected, params.use_zero_spk_emb) + text_segments = segment_and_normalize(text, max_length=params.max_text_length) if text else [] + + if not text_segments: + # vLLM passes 1D input_ids; Qwen2Model expects (batch, seq). + inputs_embeds = self.model.get_input_embeddings()(input_ids.to(self.device)).unsqueeze(0) + return generator.generate_latents( + inputs_embeds=inputs_embeds, + prompt_wav_lat=voice.prompt_wav_lat, + max_steps=params.max_steps, + cfg=params.cfg, + sigma=params.sigma, + temperature=params.temperature, + use_static_cache=params.use_static_cache, + ) + + all_latents: list[torch.Tensor] = [] + for segment in text_segments: + seg_embeds, _ = build_tts_input( + tokenizer=self.tokenizer, + embed_tokens=self.model.get_input_embeddings(), + device=self.device, + dtype=torch.bfloat16, + text=segment, + prompt=params.prompt, + spk_emb=spk_emb, + instruction=params.instruction, + prompt_text=voice.prompt_text, + prompt_wav_emb=voice.prompt_wav_emb, + ) + effective_max_steps = generator.duration_capped_steps(len(segment), params.max_steps) + all_latents.extend( + generator.generate_latents( + inputs_embeds=seg_embeds, + prompt_wav_lat=voice.prompt_wav_lat, + max_steps=effective_max_steps, + cfg=params.cfg, + sigma=params.sigma, + temperature=params.temperature, + use_static_cache=params.use_static_cache, + ) + ) + return all_latents + + def _decode_to_output(self, latents: list[torch.Tensor], *, stream_decode: bool) -> OmniOutput: + multimodal_outputs: dict[str, Any] = {} + if latents and self.audio_vae is not None: + waveform = self.audio_generator.decode_to_waveform(latents, stream_decode=stream_decode) + if not stream_decode: + waveform = self.audio_generator.trim_trailing_silence(waveform) + multimodal_outputs["audio"] = waveform.detach().float().cpu() + multimodal_outputs["sr"] = torch.tensor(self.audio_vae.config.sample_rate) + elif latents: + all_lat = torch.cat(latents, dim=1) + multimodal_outputs["audio_latents"] = all_lat.detach().float().cpu() + + return OmniOutput(text_hidden_states=None, multimodal_outputs=multimodal_outputs) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load weights for all talker components. + + The talker's HF checkpoint (talker/model.safetensors) stores + weights with prefixes matching this module's submodule names directly. + And AudioVAE weights live in a separate file under talker/vae/ + """ + # Standalone: bypass the default loader's iterator (torch.load on + # .safetensors crashes) and read talker/model*.safetensors directly. + if self._standalone: + weights = self._iter_talker_safetensors() + + loader = AutoWeightsLoader( + self, + skip_prefixes=["audio_vae."], # loaded separately + skip_substrs=["rotary_embed.inv_freq"], # non-persistent buffer + ) + loaded = loader.load_weights(weights) + logger.info("Loaded %d talker weights from checkpoint", len(loaded)) + + if self.audio_vae is not None and self._vae_weight_source is not None: + loaded.update(self._load_vae_weights()) + + # Register voice presets after all weights (incl. VAE) are loaded. + try: + self.voice_presets.load_presets_from_manifest(device=self.device, dtype=self.dtype) + except Exception as e: # pragma: no cover — best-effort + logger.warning("Voice preset loading failed (non-fatal): %s", e) + + return loaded + + def _iter_talker_safetensors(self) -> Iterable[tuple[str, torch.Tensor]]: + """Yield (name, tensor) pairs from talker/model*.safetensors.""" + model_path = self._model_path + # Try local path first + for candidate in (os.path.join(model_path, "talker"), model_path): + sf_files = sorted(glob_module.glob(os.path.join(candidate, "model*.safetensors"))) + if sf_files: + for sf_path in sf_files: + yield from load_file(sf_path, device="cpu").items() + return + + # HF hub fallback: download only the talker checkpoint files + model_root = download_weights_from_hf_specific( + model_path, + self.vllm_config.load_config.download_dir, + allow_patterns=["talker/model*.safetensors"], + ) + talker_dir = os.path.join(model_root, "talker") + sf_files = sorted(glob_module.glob(os.path.join(talker_dir, "model*.safetensors"))) + if not sf_files: + raise RuntimeError(f"No talker safetensors found under {model_root}. Expected talker/model*.safetensors.") + for sf_path in sf_files: + yield from load_file(sf_path, device="cpu").items() + + def _load_vae_weights(self) -> set[str]: + """Load AudioVAE weights from talker/vae/model.safetensors.""" + if self.audio_vae is None or self._vae_weight_source is None: + return set() + + # Resolve safetensors file paths from the weight source + safetensors_files: list[str] = [] + source = self._vae_weight_source + if isinstance(source, str): + # Local directory path + safetensors_files = sorted(glob_module.glob(os.path.join(source, "*.safetensors"))) + elif isinstance(source, tuple): + # (repo_id, subfolder) for HF hub + repo_id, subfolder = source + for filename in ("model.safetensors", "diffusion_pytorch_model.safetensors"): + try: + cached = cached_file(repo_id, filename, subfolder=subfolder) + except Exception: + cached = None + if cached is not None: + safetensors_files.append(cached) + break + + if not safetensors_files: + logger.warning("No AudioVAE safetensors files found for source=%s", source) + return set() + + vae_state_keys = set(self.audio_vae.state_dict().keys()) + vae_loader = AutoWeightsLoader(self.audio_vae) + loaded: set[str] = set() + for sf_path in safetensors_files: + file_weights = load_file(sf_path, device="cpu") + matched = ((name, tensor) for name, tensor in file_weights.items() if name in vae_state_keys) + loaded.update(f"audio_vae.{name}" for name in vae_loader.load_weights(matched)) + + logger.info("Loaded %d AudioVAE weights from %s", len(loaded), source) + return loaded diff --git a/vllm_omni/model_executor/models/ming_flash_omni/prompt_utils.py b/vllm_omni/model_executor/models/ming_flash_omni/prompt_utils.py new file mode 100644 index 000000000..4271114bc --- /dev/null +++ b/vllm_omni/model_executor/models/ming_flash_omni/prompt_utils.py @@ -0,0 +1,50 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2025 The vLLM-Omni team. +# Copyright (c) Ant Group. All rights reserved. +# Adapted from Ming repo's usage cookbook: +# https://github.com/inclusionAI/Ming/blob/3954fcb880ff5e61ff128bcf7f1ec344d46a6fe3/cookbook.ipynb +"""Shared prompt-building helpers for Ming-flash-omni standalone talker.""" + +import copy +import json +from typing import Any + +DEFAULT_PROMPT = "Please generate speech based on the following description.\n" + +BASE_CAPTION_TEMPLATE: dict[str, Any] = { + "audio_sequence": [ + { + "序号": 1, + "说话人": "speaker_1", + "方言": None, + "风格": None, + "语速": None, + "基频": None, + "音量": None, + "情感": None, + "BGM": { + "Genre": None, + "Mood": None, + "Instrument": None, + "Theme": None, + "ENV": None, + "SNR": None, + }, + "IP": None, + } + ] +} + + +def create_instruction(user_input: dict[str, Any]) -> str: + """Return a JSON caption string for ``audio_sequence[0]``. + + Only keys already present on the base template are merged in; unknown + keys are silently ignored to keep the output schema stable. + """ + caption = copy.deepcopy(BASE_CAPTION_TEMPLATE) + item = caption["audio_sequence"][0] + for key, value in user_input.items(): + if key in item: + item[key] = value + return json.dumps(caption, ensure_ascii=False) diff --git a/vllm_omni/model_executor/models/ming_flash_omni/spk_embedding.py b/vllm_omni/model_executor/models/ming_flash_omni/spk_embedding.py new file mode 100644 index 000000000..68dbfe650 --- /dev/null +++ b/vllm_omni/model_executor/models/ming_flash_omni/spk_embedding.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2025 The vLLM-Omni team. +# Copyright (c) Ant Group. All rights reserved. +# Ported from: +# https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/modeling_bailing_talker.py + +from __future__ import annotations + +import torch + + +class SpkembExtractor: + """CAMPPlus ONNX-based speaker embedding extractor (runs on CPU).""" + + def __init__(self, campplus_model: str, target_sr: int = 16000): + import onnxruntime + import torchaudio.compliance.kaldi as kaldi + + self.kaldi = kaldi + option = onnxruntime.SessionOptions() + option.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL + option.intra_op_num_threads = 2 + self.campplus_session = onnxruntime.InferenceSession( + campplus_model, sess_options=option, providers=["CPUExecutionProvider"] + ) + self.target_sr = target_sr + + def _extract_spk_embedding(self, speech): + feat = self.kaldi.fbank(speech, num_mel_bins=80, dither=0, sample_frequency=16000) + feat = feat - feat.mean(dim=0, keepdim=True) + embedding = ( + self.campplus_session.run( + None, + {self.campplus_session.get_inputs()[0].name: feat.unsqueeze(dim=0).cpu().numpy()}, + )[0] + .flatten() + .tolist() + ) + embedding = torch.tensor([embedding]) + return embedding + + def __call__(self, waveform, **kwargs) -> torch.Tensor | None: + spk_emb = self._extract_spk_embedding(waveform) + return spk_emb diff --git a/vllm_omni/model_executor/models/ming_flash_omni/talker_module.py b/vllm_omni/model_executor/models/ming_flash_omni/talker_module.py new file mode 100644 index 000000000..80acbaad0 --- /dev/null +++ b/vllm_omni/model_executor/models/ming_flash_omni/talker_module.py @@ -0,0 +1,1151 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2025 The vLLM-Omni team. +# Copyright (c) Ant Group. All rights reserved. +# Ported from: +# https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/talker_module/dit.py +# +# Ported from: +# https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/talker_module/modules.py +# Ported from: +# https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/talker_module/cfm.py + +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. + +# Partial of the following source code +# is licensed under the license found in the +# LICENSE file in the root directory of this source tree. +# -------------------------------------------------------- +# References: +# GLIDE: https://github.com/openai/glide-text2im +# MAE: https://github.com/facebookresearch/mae/blob/main/models_mae.py +# -------------------------------------------------------- +import logging +import math +from functools import cached_property +from queue import Queue +from threading import Lock + +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers import PreTrainedTokenizerBase, Qwen2Config, Qwen2Model, StaticCache +from vllm.logger import init_logger +from x_transformers.x_transformers import RotaryEmbedding, apply_rotary_pos_emb + +from .audio_vae import AudioVAE + +logger = init_logger(__name__) + + +######################################################################## +# DiT Modules +# Ported from: +# https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/talker_module/modules.py +# Ported from: +# https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/talker_module/dit.py +######################################################################## + + +class RMSNorm(nn.Module): + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.weight.dtype in [torch.float16, torch.bfloat16]: + x = x.to(self.weight.dtype) + x = F.rms_norm(x, normalized_shape=(x.shape[-1],), weight=self.weight, eps=self.eps) + return x + + +class FeedForward(nn.Module): + def __init__( + self, dim: int, dim_out: int | None = None, mult: float = 4, dropout: float = 0.0, approximate: str = "none" + ): + super().__init__() + inner_dim = int(dim * mult) + dim_out = dim_out if dim_out is not None else dim + + activation = nn.GELU(approximate=approximate) + project_in = nn.Sequential(nn.Linear(dim, inner_dim), activation) + self.ff = nn.Sequential(project_in, nn.Dropout(dropout), nn.Linear(inner_dim, dim_out)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.ff(x) + + +class Attention(nn.Module): + def __init__( + self, + dim: int, + heads: int = 8, + dim_head: int = 64, + dropout: float = 0.0, + qk_norm: str | None = None, + pe_attn_head: int | None = None, + attn_mask_enabled: bool = True, + ): + super().__init__() + self.dim = dim + self.heads = heads + self.inner_dim = dim_head * heads + self.dropout = dropout + + self.to_q = nn.Linear(dim, self.inner_dim) + self.to_k = nn.Linear(dim, self.inner_dim) + self.to_v = nn.Linear(dim, self.inner_dim) + if qk_norm is None: + self.q_norm = None + self.k_norm = None + elif qk_norm == "rms_norm": + self.q_norm = RMSNorm(dim_head) + self.k_norm = RMSNorm(dim_head) + else: + raise ValueError(f"Unimplemented qk_norm: {qk_norm}") + + self.to_out = nn.ModuleList([]) + self.to_out.append(nn.Linear(self.inner_dim, dim)) + self.to_out.append(nn.Dropout(dropout)) + + self.pe_attn_head = pe_attn_head + self.attn_mask_enabled = attn_mask_enabled + + def forward( + self, + x: torch.Tensor, + mask: torch.Tensor | None = None, + rope: tuple[torch.Tensor, torch.Tensor | None] | None = None, + ) -> torch.Tensor: + batch_size = x.shape[0] + + query = self.to_q(x) + key = self.to_k(x) + value = self.to_v(x) + + inner_dim = key.shape[-1] + head_dim = inner_dim // self.heads + query = query.view(batch_size, -1, self.heads, head_dim).transpose(1, 2) + key = key.view(batch_size, -1, self.heads, head_dim).transpose(1, 2) + value = value.view(batch_size, -1, self.heads, head_dim).transpose(1, 2) + + if self.q_norm is not None: + query = self.q_norm(query) + if self.k_norm is not None: + key = self.k_norm(key) + + if rope is not None: + freqs, xpos_scale = rope + q_xpos_scale, k_xpos_scale = (xpos_scale, xpos_scale**-1.0) if xpos_scale is not None else (1.0, 1.0) + + if self.pe_attn_head is not None: + on = self.pe_attn_head + query[:, :on, :, :] = apply_rotary_pos_emb(query[:, :on, :, :], freqs, q_xpos_scale) + key[:, :on, :, :] = apply_rotary_pos_emb(key[:, :on, :, :], freqs, k_xpos_scale) + else: + query = apply_rotary_pos_emb(query, freqs, q_xpos_scale) + key = apply_rotary_pos_emb(key, freqs, k_xpos_scale) + + if self.attn_mask_enabled and mask is not None: + valid_sample_indices = mask.any(dim=1) + final_output = torch.zeros_like(query).to(query.device) + + attn_mask = mask[valid_sample_indices] + query = query[valid_sample_indices] + key = key[valid_sample_indices] + value = value[valid_sample_indices] + attn_mask = attn_mask.unsqueeze(1).unsqueeze(1) + attn_mask = attn_mask.expand(valid_sample_indices.sum().item(), self.heads, query.shape[-2], key.shape[-2]) + else: + attn_mask = None + + x = F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=0.0, is_causal=False) + if self.attn_mask_enabled and mask is not None: + final_output[valid_sample_indices] = x + x = final_output + + x = x.transpose(1, 2).reshape(batch_size, -1, self.heads * head_dim) + x = x.to(query.dtype) + + x = self.to_out[0](x) + x = self.to_out[1](x) + + if mask is not None: + mask = mask.unsqueeze(-1) + x = x.masked_fill(~mask, 0.0) + + return x + + +class DiTBlock(nn.Module): + """A DiT block with pre-norm and residual connections.""" + + def __init__( + self, + hidden_size: int, + num_heads: int, + mlp_ratio: float = 4.0, + dropout: float = 0.1, + qk_norm: str | None = None, + pe_attn_head: int | None = None, + attn_mask_enabled: bool = True, + **kwargs, + ): + super().__init__() + self.norm1 = RMSNorm(hidden_size) + self.attn = Attention( + dim=hidden_size, + heads=num_heads, + dim_head=hidden_size // num_heads, + dropout=dropout, + qk_norm=qk_norm, + pe_attn_head=pe_attn_head, + attn_mask_enabled=attn_mask_enabled, + ) + self.norm2 = RMSNorm(hidden_size) + self.mlp = FeedForward(dim=hidden_size, mult=mlp_ratio, dropout=dropout, approximate="tanh") + + def forward( + self, + x: torch.Tensor, + mask: torch.Tensor | None, + rope: tuple[torch.Tensor, torch.Tensor | None] | None, + ) -> torch.Tensor: + x = x + self.attn(self.norm1(x), mask=mask, rope=rope) + x = x + self.mlp(self.norm2(x)) + return x + + +class FinalLayer(nn.Module): + """The final layer of DiT.""" + + def __init__(self, hidden_size: int, out_channels: int): + super().__init__() + self.norm_final = RMSNorm(hidden_size) + self.linear = nn.Linear(hidden_size, out_channels, bias=True) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.norm_final(x) + x = self.linear(x) + return x + + +class SinusPositionEmbedding(nn.Module): + def __init__(self, dim: int): + super().__init__() + self.dim = dim + + def forward(self, x: torch.Tensor, scale: float = 1000) -> torch.Tensor: + device = x.device + half_dim = self.dim // 2 + emb = math.log(10000) / (half_dim - 1) + emb = torch.exp(torch.arange(half_dim, device=device).float() * -emb) + emb = scale * x.unsqueeze(1) * emb.unsqueeze(0) + emb = torch.cat((emb.sin(), emb.cos()), dim=-1) + return emb + + +class TimestepEmbedder(nn.Module): + def __init__(self, dim: int, freq_embed_dim: int = 256): + super().__init__() + self.time_embed = SinusPositionEmbedding(freq_embed_dim) + self.time_mlp = nn.Sequential(nn.Linear(freq_embed_dim, dim), nn.SiLU(), nn.Linear(dim, dim)) + + def forward(self, timestep: torch.Tensor) -> torch.Tensor: + time_hidden = self.time_embed(timestep) + time_hidden = time_hidden.to(timestep.dtype) + time = self.time_mlp(time_hidden) + return time + + +class CondEmbedder(nn.Module): + """Embeds LLM hidden states with optional CFG dropout.""" + + def __init__(self, input_feature_size: int, hidden_size: int): + super().__init__() + self.cond_embedder = nn.Linear(input_feature_size, hidden_size) + + def forward(self, llm_cond: torch.Tensor) -> torch.Tensor: + return self.cond_embedder(llm_cond) + + +class DiT(nn.Module): + """Diffusion model with a Transformer backbone for audio latent generation.""" + + def __init__( + self, + in_channels: int = 64, + hidden_size: int = 1024, + depth: int = 28, + num_heads: int = 16, + mlp_ratio: float = 4.0, + llm_cond_dim: int = 896, + **kwargs, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = in_channels + self.num_heads = num_heads + + self.t_embedder = TimestepEmbedder(hidden_size) + self.x_embedder = nn.Linear(in_channels, hidden_size) + self.c_embedder = CondEmbedder(llm_cond_dim, hidden_size) + if "spk_dim" in kwargs: + self.spk_embedder = nn.Linear(kwargs["spk_dim"], hidden_size) + else: + self.spk_embedder = None + self.hidden_size = hidden_size + + self.rotary_embed = RotaryEmbedding(hidden_size // num_heads) + + self.blocks = nn.ModuleList( + [DiTBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio, **kwargs) for _ in range(depth)] + ) + self.final_layer = FinalLayer(hidden_size, self.out_channels) + + def forward( + self, + x: torch.Tensor, + t: torch.Tensor, + c: torch.Tensor, + latent_history: torch.Tensor, + spk_emb: torch.Tensor | None = None, + ) -> torch.Tensor: + x = torch.cat([latent_history, x], dim=1) + x = self.x_embedder(x) + t = self.t_embedder(t).unsqueeze(1) + c = self.c_embedder(c) + y = t + c + if spk_emb is None: + assert self.spk_embedder is None + x = torch.cat([y, x], dim=1) + else: + x = torch.cat([self.spk_embedder(spk_emb), y, x], dim=1) + rope = self.rotary_embed.forward_from_seq_len(x.shape[1]) + + for block in self.blocks: + x = block(x, None, rope) + x = self.final_layer(x) + return x + + def forward_with_cfg( + self, + x: torch.Tensor, + t: torch.Tensor, + c: torch.Tensor, + latent_history: torch.Tensor, + spk_emb: torch.Tensor | None = None, + ) -> torch.Tensor: + """Forward with classifier-free guidance (doubles batch for CFG).""" + x = torch.cat([x, x], dim=0) + latent_history = torch.cat([latent_history, latent_history], dim=0) + fake_latent = torch.zeros_like(c) + c = torch.cat([c, fake_latent], dim=0) + if t.ndim == 0: + t = t.repeat(x.shape[0]) + if spk_emb is not None: + spk_emb = torch.cat([spk_emb, spk_emb], dim=0) + model_out = self.forward(x, t, c, latent_history, spk_emb) + return model_out[:, -x.shape[1] :, :] + + +######################################################################################### +# CFM +# Ported from: +# https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/talker_module/cfm.py +######################################################################################### + + +def get_epss_timesteps(n, device, dtype): + dt = 1 / 32 + predefined_timesteps = { + 5: [0, 2, 4, 8, 16, 32], + 6: [0, 2, 4, 6, 8, 16, 32], + 7: [0, 2, 4, 6, 8, 16, 24, 32], + 10: [0, 2, 4, 6, 8, 12, 16, 20, 24, 28, 32], + 12: [0, 2, 4, 6, 8, 10, 12, 14, 16, 20, 24, 28, 32], + 16: [0, 1, 2, 3, 4, 5, 6, 7, 8, 10, 12, 14, 16, 20, 24, 28, 32], + } + t = predefined_timesteps.get(n, []) + if not t: + return torch.linspace(0, 1, n + 1, device=device, dtype=dtype) + return dt * torch.tensor(t, device=device, dtype=dtype) + + +class CFM(nn.Module): + """Conditional Flow Matching module for audio latent generation.""" + + def __init__(self, model: nn.Module, steps: int = 10, sway_sampling_coef: float | None = -1.0): + """ + Args: + model: DiT used for the velocity prediction. + steps: number of integration steps per sample call. + sway_sampling_coef: coefficient used to skew the integration + grid towards low-noise timesteps. Defaults to -1.0 which + packs more steps near t=0, where prediction error is highest. + Set to `None` to use the linear grid as-is. + """ + super().__init__() + self.model = model + self.steps = steps + self.sway_sampling_coef = sway_sampling_coef + + @torch.no_grad() + def sample( + self, + llm_cond: torch.Tensor, + lat_cond: torch.Tensor, + y0: torch.Tensor, + t: torch.Tensor, + sde_args: torch.Tensor, + sde_rnd: torch.Tensor, + ): + """Sample audio latent via ODE/SDE integration with CFG. + + Args: + llm_cond: LLM hidden state (B, 1, hidden_size) + lat_cond: latent history (B, his_patch_size, latent_dim) + y0: initial noise (B, patch_size, latent_dim) + t: timesteps from get_epss_timesteps + sde_args: [cfg_strength, sigma, temperature] + sde_rnd: random noise for SDE steps (steps, B, patch_size, latent_dim) + """ + + def fn(fn_t, x): + pred_cfg = self.model.forward_with_cfg(x, fn_t, llm_cond, lat_cond, None) + pred, null_pred = torch.chunk(pred_cfg, 2, dim=0) + return pred + (pred - null_pred) * sde_args[0] + + if self.sway_sampling_coef is not None: + t = t + self.sway_sampling_coef * (torch.cos(torch.pi / 2 * t) - 1 + t) + + for step in range(self.steps): + dt = t[step + 1] - t[step] + y0 = y0 + fn(t[step], y0) * dt + y0 = y0 + sde_args[1] * (sde_args[2] ** 0.5) * (dt.abs() ** 0.5) * sde_rnd[step] + + return y0 + + +class CFMGraphExecutor: + """CUDA graph-accelerated executor for CFM + Aggregator + StopHead pipeline.""" + + def __init__(self, config, cfm, aggregator, stop_head: nn.Linear): + self.config = config + self.cfm = cfm + self.aggregator = aggregator + self.stop_head = stop_head + self.initialized = False + + self.last_hidden_state_placeholder = None + self.his_lat_placeholder = None + self.randn_like_placeholder = None + self.t_placeholder = None + self.sde_args_placeholder = None + self.sde_rnd_placeholder = None + self.gen_lat_placeholder = None + self.inputs_embeds_placeholder = None + self.stop_out_placeholder = None + self.graph = None + + def execute( + self, + input_tensor: torch.Tensor, + his_lat: torch.Tensor, + cfg_strength: float = 2.0, + sigma: float = 0.25, + temperature: float = 0.0, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + bat_size, his_patch_size, z_dim = his_lat.shape + randn_tensor = torch.randn( + (bat_size, self.config.patch_size, z_dim), device=input_tensor.device, dtype=input_tensor.dtype + ) + t = get_epss_timesteps(self.config.steps, device=input_tensor.device, dtype=input_tensor.dtype) + sde_rnd = torch.randn( + (self.config.steps, *randn_tensor.shape), device=input_tensor.device, dtype=input_tensor.dtype + ) + + if not self.initialized: + self._initialize_graph(input_tensor, his_lat, randn_tensor, sde_rnd) + + self.last_hidden_state_placeholder.copy_(input_tensor) + self.his_lat_placeholder.copy_(his_lat) + self.randn_like_placeholder.copy_(randn_tensor) + self.t_placeholder.copy_(t) + self.sde_args_placeholder[0] = cfg_strength + self.sde_args_placeholder[1] = sigma + self.sde_args_placeholder[2] = temperature + self.sde_rnd_placeholder.copy_(sde_rnd) + + self.graph.replay() + + gen_lat = torch.empty_like(self.gen_lat_placeholder) + gen_lat.copy_(self.gen_lat_placeholder) + + inputs_embeds = torch.empty_like(self.inputs_embeds_placeholder) + inputs_embeds.copy_(self.inputs_embeds_placeholder) + + stop_out = torch.empty_like(self.stop_out_placeholder) + stop_out.copy_(self.stop_out_placeholder) + + return gen_lat, inputs_embeds, stop_out + + def _initialize_graph( + self, + input_tensor: torch.Tensor, + his_lat: torch.Tensor, + randn_tensor: torch.Tensor, + sde_rnd: torch.Tensor, + ) -> None: + self.last_hidden_state_placeholder = torch.empty_like(input_tensor) + self.his_lat_placeholder = torch.empty_like(his_lat) + self.randn_like_placeholder = torch.empty_like(randn_tensor) + self.t_placeholder = get_epss_timesteps(self.config.steps, device=input_tensor.device, dtype=input_tensor.dtype) + self.sde_args_placeholder = torch.empty(3, device=input_tensor.device, dtype=input_tensor.dtype) + self.sde_rnd_placeholder = torch.empty_like(sde_rnd) + + self.graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(self.graph): + self.gen_lat_placeholder = self.cfm.sample( + self.last_hidden_state_placeholder, + self.his_lat_placeholder, + self.randn_like_placeholder, + self.t_placeholder, + self.sde_args_placeholder, + self.sde_rnd_placeholder, + ) + self.inputs_embeds_placeholder = self.aggregator(self.gen_lat_placeholder) + self.stop_out_placeholder = self.stop_head(self.last_hidden_state_placeholder[:, -1, :]).softmax(dim=-1) + + self.initialized = True + + +class CFMGraphExecutorPool: + """Thread-safe pool of CFMGraphExecutors for concurrent inference.""" + + def __init__(self, config, cfm, aggregator, stop_head: nn.Linear, pool_size: int = 1): + self.config = config + self.cfm = cfm + self.aggregator = aggregator + self.stop_head = stop_head + self.pool_size = pool_size + self.pool = Queue(maxsize=pool_size) + self.lock = Lock() + + for _ in range(pool_size): + executor = CFMGraphExecutor(config, cfm, aggregator, stop_head) + self.pool.put(executor) + + def acquire(self) -> CFMGraphExecutor: + return self.pool.get() + + def release(self, executor: CFMGraphExecutor) -> None: + self.pool.put(executor) + + def execute( + self, + input_tensor: torch.Tensor, + his_lat: torch.Tensor, + cfg_strength: float = 2.0, + sigma: float = 0.25, + temperature: float = 0.0, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + executor = self.acquire() + try: + return executor.execute( + input_tensor, his_lat, cfg_strength=cfg_strength, sigma=sigma, temperature=temperature + ) + finally: + self.release(executor) + + +######################################################################## +# Audio Postprocess +# Adapted from: +# https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/modeling_bailing_talker.py +######################################################################## + + +@torch.no_grad() +def resample(waveform: torch.Tensor, orig_sr: int, target_sr: int) -> torch.Tensor: + """Resample a waveform via linear interpolation (no torchaudio dep). + + Args: + waveform: Tensor shaped ``(..., num_samples)``. + orig_sr: Source sample rate (Hz); must be > 0. + target_sr: Target sample rate (Hz); must be > 0. + + Raises: + ValueError: If sample rates are non-positive, the waveform is empty, + or the resampled length would round to zero. + """ + if orig_sr <= 0: + raise ValueError(f"orig_sr must be positive, got {orig_sr}") + if target_sr <= 0: + raise ValueError(f"target_sr must be positive, got {target_sr}") + if waveform.numel() == 0 or waveform.shape[-1] == 0: + raise ValueError("waveform must contain at least one sample") + if orig_sr == target_sr: + return waveform + + ratio = target_sr / orig_sr + new_len = int(waveform.shape[-1] * ratio) + if new_len <= 0: + raise ValueError( + f"resampled waveform would be empty for input length {waveform.shape[-1]}, " + f"orig_sr={orig_sr}, target_sr={target_sr}" + ) + return torch.nn.functional.interpolate( + waveform.unsqueeze(0), + size=new_len, + mode="linear", + align_corners=False, + ).squeeze(0) + + +def trim_trailing_silence( + waveform: torch.Tensor, + sample_rate: int, + sil_th: float = 1e-3, + tail_silence_s: float = 0.3, +) -> torch.Tensor: + """Trim low-energy tail while keeping a short trailing silence. + + Works on 2-D ``(channels, samples)`` or 3-D ``(batch, channels, samples)`` + tensors. Any other shape is returned unchanged. + """ + if waveform.numel() == 0: + return waveform + + original_dim = waveform.dim() + if original_dim == 3: + speech = waveform[:, 0, :] + elif original_dim == 2: + speech = waveform + else: + return waveform + + frame_step = int(sample_rate * 0.1) + frame_size = int(sample_rate * 0.1) + if speech.shape[-1] < frame_size: + keep = min(speech.shape[-1], int(tail_silence_s * sample_rate)) + trimmed = speech[..., :keep] + else: + num_frame = (speech.shape[-1] - frame_size) // frame_step + 1 + cur_len = (num_frame - 1) * frame_step + frame_size + speech = speech[..., :cur_len] + spe_frames = speech.unfold(-1, frame_size, frame_step) + scores = spe_frames.abs().mean(dim=-1) + scores = scores.mean(dim=list(range(scores.dim() - 1))) + idx = scores.shape[0] - 1 + while idx >= 0 and scores[idx] <= sil_th: + idx -= 1 + if idx < 0: + keep = min(speech.shape[-1], int(tail_silence_s * sample_rate)) + trimmed = speech[..., :keep] + else: + non_sil_len = idx * frame_step + frame_size + int(tail_silence_s * sample_rate) + non_sil_len = min(non_sil_len, speech.shape[-1]) + trimmed = speech[..., :non_sil_len] + + if original_dim == 3: + return trimmed.unsqueeze(1) + return trimmed + + +def silence_holder( + speech: torch.Tensor, + sample_rate: int, + sil_cache: dict | None = None, + last_chunk: bool = True, + sil_th: float = 1e-3, + last_sil: float = 0.3, +) -> tuple[torch.Tensor, dict]: + """Ming-style streaming silence holder. + + Used during streaming VAE decode to defer emission of silent regions + until a non-silent frame arrives (or the stream ends). ``sil_cache`` + is carried across chunks and updated in place. + """ + if speech.numel() == 0: + return speech, sil_cache or {"holder": [], "buffer": []} + + frame_step = int(sample_rate * 0.1) + frame_size = int(sample_rate * 0.1) + if sil_cache is None: + sil_cache = {"holder": [], "buffer": []} + + if sil_cache["buffer"]: + speech = torch.cat([*sil_cache["buffer"], speech], dim=-1) + sil_cache["buffer"] = [] + + if speech.shape[-1] < frame_size: + sil_cache["buffer"].append(speech) + if last_chunk: + speech = torch.cat(sil_cache["holder"] + sil_cache["buffer"], dim=-1) + return speech[..., : int(last_sil * sample_rate)], sil_cache + return torch.zeros((*speech.shape[:-1], 0), device=speech.device, dtype=speech.dtype), sil_cache + + num_frame = (speech.shape[-1] - frame_size) // frame_step + 1 + cur_len = (num_frame - 1) * frame_step + frame_size + if speech.shape[-1] > cur_len: + sil_cache["buffer"].append(speech[..., cur_len:]) + speech = speech[..., :cur_len] + + spe_frames = speech.unfold(-1, frame_size, frame_step) + scores = spe_frames.abs().mean(dim=-1) + scores = scores.mean(dim=list(range(scores.dim() - 1))) + idx = scores.shape[0] - 1 + while idx >= 0 and scores[idx] <= sil_th: + idx -= 1 + + if idx < 0: + sil_cache["holder"].append(speech) + if last_chunk: + speech = torch.cat(sil_cache["holder"] + sil_cache["buffer"], dim=-1) + return speech[..., : int(last_sil * sample_rate)], sil_cache + return torch.zeros((*speech.shape[:-1], 0), device=speech.device, dtype=speech.dtype), sil_cache + + non_sil_len = idx * frame_step + frame_size + if last_chunk: + non_sil_len += int(last_sil * sample_rate) + non_sil_len = min(non_sil_len, speech.shape[-1]) + speech_out = torch.cat([*sil_cache["holder"], speech[..., :non_sil_len]], dim=-1) + sil_cache["holder"] = [] + if non_sil_len < speech.shape[-1]: + sil_cache["holder"].append(speech[..., non_sil_len:]) + return speech_out, sil_cache + + +######################################################################## +# Audio Postprocess +# Ported from: +# https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/talker_module/aggregator.py +######################################################################## + + +class Aggregator(nn.Module): + """Maps generated audio latent patches back to LLM embedding space.""" + + def __init__( + self, + in_channels: int = 64, + hidden_size: int = 1152, + depth: int = 28, + num_heads: int = 16, + mlp_ratio: float = 4.0, + llm_input_dim: int = 896, + **kwargs, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = in_channels + self.num_heads = num_heads + + self.word_embedder = nn.Embedding(1, hidden_size) + self.x_embedder = nn.Linear(in_channels, hidden_size) + self.hidden_size = hidden_size + + self.rotary_embed = RotaryEmbedding(hidden_size // num_heads) + + self.blocks = nn.ModuleList( + [DiTBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio, **kwargs) for _ in range(depth)] + ) + self.final_layer = FinalLayer(hidden_size, llm_input_dim) + + def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor: + x = self.x_embedder(x) + cls_embed = self.word_embedder(torch.zeros((x.shape[0], 1), dtype=torch.long, device=x.device)) + x = torch.cat([cls_embed, x], dim=1) + + rope = self.rotary_embed.forward_from_seq_len(x.shape[1]) + if mask is not None: + mask_pad = mask.clone().detach()[:, :1] + mask = torch.cat([mask_pad, mask], dim=-1) + for block in self.blocks: + x = block(x, mask, rope) + x = self.final_layer(x) + x = x[:, :1, :] + return x + + +######################################################################## +# Prompt Builder +# Adapted from: +# https://github.com/inclusionAI/Ming/blob/e58533db227031990c5a6864dcf5f08fb53ed0d2/modeling_bailing_talker.py +######################################################################## + +_MUSIC_TAGS = ("Genre: ", "Mood: ", "Instrument: ", "Theme: ", "Duration: ") + + +def _looks_like_music_prompt(text: str) -> bool: + return all(tag in text for tag in _MUSIC_TAGS) + + +def build_tts_input( + *, + tokenizer: PreTrainedTokenizerBase, + embed_tokens: torch.nn.Module, + device: torch.device, + dtype: torch.dtype, + text: str, + prompt: str, + spk_emb: list[torch.Tensor] | None = None, + instruction: str | None = None, + prompt_text: str | None = None, + prompt_wav_emb: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Build (inputs_embeds, input_ids) for one TTS segment. + + Args: + tokenizer: HF tokenizer + embed_tokens: The LLM's input-embedding module + device: Device to place the returned tensors on. + dtype: dtype for the returned `inputs_embeds`. + text: Text to synthesize. + prompt: System-level instruction prompt prepended to the user turn. + spk_emb: Optional list of speaker embeddings already projected into + LLM hidden dim; each is injected at a `<|vision_start|>` slot. + instruction: Optional free-form instruction + prompt_text: Reference text for zero-shot voice cloning. + prompt_wav_emb: Reference-wav embeddings to inject. + """ + spk_emb_prompt: list[int] = [] + if spk_emb is not None: + for i in range(len(spk_emb)): + spk_emb_prompt.extend( + tokenizer.encode(f" speaker_{i + 1}:") + + tokenizer.encode("<|vision_start|>") + + tokenizer.encode("<|vision_pad|>") + + tokenizer.encode("<|vision_end|>\n") + ) + + instruction_prompt: list[int] = [] + if instruction is not None: + instruction_prompt = tokenizer.encode(instruction) + tokenizer.encode("<|im_end|>") + + prompt_text_token: list[int] = [] + prompt_latent_token: list[int] = [] + if prompt_wav_emb is not None and prompt_text is not None: + prompt_text_token = tokenizer.encode(prompt_text) + prompt_latent_token = tokenizer.encode("") * prompt_wav_emb.size(1) + + prompt2 = [] if _looks_like_music_prompt(text) else tokenizer.encode(" Text input:\n") + + input_part = ( + tokenizer.encode("<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n") + + tokenizer.encode("<|im_start|>user\n") + + tokenizer.encode(prompt) + + spk_emb_prompt + + prompt2 + + prompt_text_token + + tokenizer.encode(text) + + tokenizer.encode("<|im_end|>\n") + + tokenizer.encode("<|im_start|>assistant\n") + + instruction_prompt + + tokenizer.encode("