Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
380127049a | ||
|
|
743953b821 | ||
|
|
86166fd988 |
No files matched your search
@@ -107,11 +107,6 @@ steps:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py -m core_model
|
||||
- >-
|
||||
timeout 20m pytest -s -v
|
||||
tests/diffusion/models/ltx2/test_ltx2_transformer_ulysses.py
|
||||
-m 'core_model and rocm and cards_2'
|
||||
--run-level "core_model"
|
||||
|
||||
- label: "Diffusion GPU Worker Test"
|
||||
agent_pool: mi300_2
|
||||
@@ -127,13 +122,7 @@ steps:
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
# Preserve legacy one-GPU tests without cards_1 while excluding every
|
||||
# registered multi-GPU marker from this one-GPU job.
|
||||
- >-
|
||||
timeout 15m pytest -sv tests/diffusion/models/
|
||||
-m 'core_model and cuda and not
|
||||
(cards_2 or cards_3 or cards_4 or cards_5 or cards_6 or cards_7 or cards_8)'
|
||||
--run-level "core_model"
|
||||
- timeout 15m pytest -sv tests/diffusion/models/ -m 'core_model and cuda' --run-level "core_model"
|
||||
|
||||
- label: "Diffusion · Wan22 Test"
|
||||
agent_pool: mi300_1
|
||||
|
||||
@@ -79,8 +79,6 @@ steps:
|
||||
steps:
|
||||
- label: "TTS · Qwen3-TTS CustomVoice Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/features/prefill_phase/
|
||||
- vllm_omni/worker/gpu_model_runner.py
|
||||
- tests/helpers/media.py
|
||||
- tests/e2e/online_serving/test_qwen3_tts_customvoice.py
|
||||
- tests/e2e/offline_inference/test_qwen3_tts_customvoice.py
|
||||
@@ -98,7 +96,7 @@ steps:
|
||||
timeout 40m bash -c "
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_ALLOW_LONG_MAX_MODEL_LEN=1
|
||||
pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py tests/e2e/offline_inference/test_qwen3_tts_customvoice.py tests/e2e/features/prefill_phase/test_chunked_prefill_tail.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py tests/e2e/offline_inference/test_qwen3_tts_customvoice.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
"
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
|
||||
@@ -76,7 +76,7 @@ steps:
|
||||
|
||||
- label: "Diffusion · Other Test"
|
||||
commands:
|
||||
- timeout 15m pytest -sv tests/diffusion/*.py tests/diffusion/ar_diffusion -m 'core_model and cuda' --run-level "core_model"
|
||||
- timeout 15m pytest -sv tests/diffusion/*.py -m 'core_model and cuda' --run-level "core_model"
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
- group: ":card_index_dividers: Tiny Model Tests [Single GPU]"
|
||||
@@ -93,7 +93,7 @@ steps:
|
||||
commands:
|
||||
- |
|
||||
timeout 15m bash -c "
|
||||
pytest -sv tests/engine/ tests/model_executor/ tests/worker/test_batched_omni_output.py -m 'core_model and cuda' --run-level 'core_model'
|
||||
pytest -sv tests/engine/ tests/model_executor/ -m 'core_model and cuda' --run-level 'core_model'
|
||||
"
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
|
||||
@@ -19,8 +19,6 @@ steps:
|
||||
pytest -s -v
|
||||
tests/diffusion/layers/test_activation.py
|
||||
tests/diffusion/layers/test_fused_qk_norm_rope_npu.py
|
||||
tests/diffusion/models/minimax_h3/test_minimax_h3_qwen3vl_swiglu.py
|
||||
tests/diffusion/models/minimax_h3/test_minimax_h3_modulation.py
|
||||
-m 'core_model and npu and A3'
|
||||
--run-level 'core_model'
|
||||
|
||||
|
||||
@@ -2,8 +2,6 @@
|
||||
|
||||
You may find information about contributing to vLLM-Omni on [Contributing](https://vllm-omni.readthedocs.io/en/latest/contributing/).
|
||||
|
||||
AI-assisted contributions must follow the [AI Usage Policy](docs/contributing/README.md#ai-usage-policy), including disclosure, human review, and validation requirements.
|
||||
|
||||
Local `pre-commit` gates (SPDX, forbidden imports, `torch.cuda`, shellcheck,
|
||||
mypy, test marks, TTS adapter ratchet, markdownlint, Buildkite schema, and
|
||||
how to extend allowlists) are documented in the
|
||||
|
||||
@@ -17,18 +17,18 @@ The main entrypoint is:
|
||||
vllm serve Qwen/Qwen-Image --omni --port 8099
|
||||
```
|
||||
|
||||
1. Run a minimal benchmark:
|
||||
2. Run a minimal benchmark:
|
||||
|
||||
```bash
|
||||
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
--base-url http://localhost:8099 \
|
||||
--model Qwen/Qwen-Image \
|
||||
--task t2i \
|
||||
--dataset vbench \
|
||||
--num-prompts 5
|
||||
--base-url http://localhost:8099 \
|
||||
--model Qwen/Qwen-Image \
|
||||
--task t2i \
|
||||
--dataset vbench \
|
||||
--num-prompts 5
|
||||
```
|
||||
|
||||
### Notes
|
||||
**Notes**
|
||||
|
||||
- By default, image tasks talk to `http://<host>:<port>/v1/chat/completions`; video tasks talk to `/v1/videos`.
|
||||
- If you run the server on another host or port, pass `--base-url` accordingly.
|
||||
@@ -52,14 +52,14 @@ Example (`t2v`):
|
||||
|
||||
```bash
|
||||
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
--base-url http://localhost:8099 \
|
||||
--model Wan-AI/Wan2.2-T2V-A14B-Diffusers \
|
||||
--task t2v \
|
||||
--dataset vbench \
|
||||
--num-prompts 50 \
|
||||
--width 640 --height 480 \
|
||||
--num-frames 81 --fps 16 \
|
||||
--num-inference-steps 40
|
||||
--base-url http://localhost:8099 \
|
||||
--model Wan-AI/Wan2.2-T2V-A14B-Diffusers \
|
||||
--task t2v \
|
||||
--dataset vbench \
|
||||
--num-prompts 50 \
|
||||
--width 640 --height 480 \
|
||||
--num-frames 81 --fps 16 \
|
||||
--num-inference-steps 40
|
||||
```
|
||||
|
||||
Note: `vbench` can also be used for other tasks such as `t2i` / `i2v` (and `i2i`). For `t2i`, the loader reuses VBench t2v text prompts; for `i2v` / `i2i`, it loads the VBench i2v dataset (with image paths).
|
||||
@@ -140,21 +140,6 @@ Traffic / concurrency flags:
|
||||
- `--request-rate`: Target request rate (requests/second). If set to `inf`, the script sends all requests immediately.
|
||||
- `--max-concurrency`: Max number of in-flight requests (default: `1`). This can hard-cap the achieved QPS: if it is too small, requests will queue behind the semaphore, and both achieved throughput and observed SLO attainment can be skewed.
|
||||
|
||||
### Video job timeout and failures
|
||||
|
||||
`--video-job-timeout` controls how long each `/v1/videos` job may be polled,
|
||||
including time queued on the server (default: `900` seconds). For long video
|
||||
jobs at high concurrency, use a larger budget, for example
|
||||
`--video-job-timeout 1800`. This also applies to warmup jobs. The polling budget
|
||||
starts after job creation; it excludes time waiting for the client concurrency
|
||||
semaphore. Individual HTTP calls still use the aiohttp session timeout.
|
||||
|
||||
An expired job is counted as failed and deleted, which cancels queued or running
|
||||
server generation. The benchmark continues processing the remaining requests.
|
||||
The progress bar counts both successful and failed requests. The final report
|
||||
shows the failure count and the first 10 errors; `--output-file` saves all errors
|
||||
with their request IDs in `request_errors`.
|
||||
|
||||
### Batched warmup note
|
||||
|
||||
For batched serving runs, warm the same in-flight shape you plan to measure.
|
||||
|
||||
@@ -1,6 +1,3 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import json
|
||||
@@ -46,7 +43,6 @@ class RequestFuncInput:
|
||||
video_paths: list[str] | None = None
|
||||
request_id: str = field(default_factory=lambda: str(uuid.uuid4()))
|
||||
default_bot_task: str | None = DEFAULT_EDITS_BOT_TASK
|
||||
video_job_timeout: float = 900.0
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -335,64 +331,64 @@ async def async_request_v1_videos(
|
||||
output = RequestFuncOutput()
|
||||
output.start_time = time.perf_counter()
|
||||
|
||||
files = dict(input.extra_body)
|
||||
if input.prompt:
|
||||
files.setdefault("prompt", input.prompt)
|
||||
if input.width and input.height:
|
||||
files.setdefault("height", input.height)
|
||||
files.setdefault("width", input.width)
|
||||
if input.num_frames:
|
||||
files.setdefault("num_frames", input.num_frames)
|
||||
if input.num_inference_steps:
|
||||
files.setdefault("num_inference_steps", input.num_inference_steps)
|
||||
if input.seed is not None:
|
||||
files.setdefault("seed", input.seed)
|
||||
if input.fps:
|
||||
files.setdefault("fps", input.fps)
|
||||
|
||||
form = aiohttp.FormData()
|
||||
for k, v in files.items():
|
||||
form.add_field(k, str(v))
|
||||
|
||||
image_file = None
|
||||
job_id = None
|
||||
try:
|
||||
files = dict(input.extra_body)
|
||||
if input.prompt:
|
||||
files.setdefault("prompt", input.prompt)
|
||||
if input.width and input.height:
|
||||
files.setdefault("height", input.height)
|
||||
files.setdefault("width", input.width)
|
||||
if input.num_frames:
|
||||
files.setdefault("num_frames", input.num_frames)
|
||||
if input.num_inference_steps:
|
||||
files.setdefault("num_inference_steps", input.num_inference_steps)
|
||||
if input.seed is not None:
|
||||
files.setdefault("seed", input.seed)
|
||||
if input.fps:
|
||||
files.setdefault("fps", input.fps)
|
||||
if input.image_paths and input.video_paths:
|
||||
output.error = "Only one of image_paths or video_paths can be provided"
|
||||
output.success = False
|
||||
return output
|
||||
|
||||
form = aiohttp.FormData()
|
||||
for k, v in files.items():
|
||||
form.add_field(k, str(v))
|
||||
|
||||
if input.image_paths and input.video_paths:
|
||||
output.error = "Only one of image_paths or video_paths can be provided"
|
||||
output.success = False
|
||||
return output
|
||||
|
||||
if input.image_paths and len(input.image_paths) > 0:
|
||||
image_path = input.image_paths[0]
|
||||
image_file = open(image_path, "rb")
|
||||
if input.image_paths and len(input.image_paths) > 0:
|
||||
image_path = input.image_paths[0]
|
||||
image_file = open(image_path, "rb")
|
||||
form.add_field(
|
||||
"input_reference",
|
||||
image_file,
|
||||
filename=os.path.basename(image_path),
|
||||
content_type=_guess_mime_type(image_path),
|
||||
)
|
||||
elif input.video_paths and len(input.video_paths) > 0:
|
||||
video_path = input.video_paths[0]
|
||||
if _is_minimax_h3_mixed_reference(input.model):
|
||||
with open(video_path, "rb") as reference_file:
|
||||
video_b64 = base64.b64encode(reference_file.read()).decode("ascii")
|
||||
form.add_field(
|
||||
"video_reference",
|
||||
json.dumps([{"video_url": f"data:video/mp4;base64,{video_b64}"}]),
|
||||
)
|
||||
else:
|
||||
image_file = open(video_path, "rb")
|
||||
form.add_field(
|
||||
"input_reference",
|
||||
image_file,
|
||||
filename=os.path.basename(image_path),
|
||||
content_type=_guess_mime_type(image_path),
|
||||
filename=os.path.basename(video_path),
|
||||
content_type=_guess_mime_type(video_path),
|
||||
)
|
||||
elif input.video_paths and len(input.video_paths) > 0:
|
||||
video_path = input.video_paths[0]
|
||||
if _is_minimax_h3_mixed_reference(input.model):
|
||||
with open(video_path, "rb") as reference_file:
|
||||
video_b64 = base64.b64encode(reference_file.read()).decode("ascii")
|
||||
form.add_field(
|
||||
"video_reference",
|
||||
json.dumps([{"video_url": f"data:video/mp4;base64,{video_b64}"}]),
|
||||
)
|
||||
else:
|
||||
image_file = open(video_path, "rb")
|
||||
form.add_field(
|
||||
"input_reference",
|
||||
image_file,
|
||||
filename=os.path.basename(video_path),
|
||||
content_type=_guess_mime_type(video_path),
|
||||
)
|
||||
|
||||
job_status = None
|
||||
poll_json = {}
|
||||
resp_json = {}
|
||||
job_id = None
|
||||
job_status = None
|
||||
poll_json = {}
|
||||
resp_json = {}
|
||||
|
||||
try:
|
||||
# invoke a post request (POST /v1/videos)
|
||||
async with session.post(input.api_url, data=form) as response:
|
||||
if response.status == 200:
|
||||
@@ -410,7 +406,7 @@ async def async_request_v1_videos(
|
||||
|
||||
# invoke a poll request (GET /v1/videos/{video_id})
|
||||
poll_interval = 2.0 # Unit(s)
|
||||
timeout_seconds = input.video_job_timeout
|
||||
timeout_seconds = 600.0
|
||||
deadline = time.perf_counter() + timeout_seconds
|
||||
job_url = f"{input.api_url}/{job_id}"
|
||||
|
||||
@@ -426,11 +422,8 @@ async def async_request_v1_videos(
|
||||
poll_json = await poll_response.json()
|
||||
job_status = poll_json.get("status")
|
||||
|
||||
if job_status not in {"completed", "failed"} and time.perf_counter() >= deadline:
|
||||
output.error = (
|
||||
f"Timed out after {timeout_seconds:g}s waiting for video job {job_id} to complete. "
|
||||
"Increase --video-job-timeout for long-running or queued video jobs."
|
||||
)
|
||||
if time.perf_counter() >= deadline:
|
||||
output.error = f"Timed out waiting for video job {job_id} to complete."
|
||||
output.success = False
|
||||
return output
|
||||
|
||||
@@ -459,7 +452,7 @@ async def async_request_v1_videos(
|
||||
elif "peak_memory_mb" in resp_json:
|
||||
output.peak_memory_mb = resp_json["peak_memory_mb"]
|
||||
except Exception as e:
|
||||
output.error = f"{type(e).__name__}: {e}"
|
||||
output.error = str(e)
|
||||
output.success = False
|
||||
finally:
|
||||
if image_file is not None:
|
||||
@@ -472,13 +465,13 @@ async def async_request_v1_videos(
|
||||
except Exception as e:
|
||||
print(f"Failed to clean up video job {job_id}: {e}")
|
||||
|
||||
output.latency = time.perf_counter() - output.start_time
|
||||
if pbar is not None:
|
||||
pbar.update(1)
|
||||
output.latency = time.perf_counter() - output.start_time
|
||||
|
||||
if output.success and input.slo_ms is not None:
|
||||
output.slo_achieved = (output.latency * 1000.0) <= float(input.slo_ms)
|
||||
|
||||
if pbar:
|
||||
pbar.update(1)
|
||||
return output
|
||||
|
||||
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
# adapted from fastvideo
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
Benchmark online serving for diffusion models (Image/Video Generation).
|
||||
@@ -87,7 +87,6 @@ import base64
|
||||
import glob
|
||||
import json
|
||||
import logging
|
||||
import math
|
||||
import os
|
||||
import random
|
||||
import tempfile
|
||||
@@ -1117,11 +1116,6 @@ def calculate_metrics(
|
||||
"duration": total_duration,
|
||||
"completed_requests": num_success,
|
||||
"failed_requests": len(error_outputs),
|
||||
"request_errors": [
|
||||
{"request_id": req.request_id, "error": out.error}
|
||||
for req, out in zip(requests_list, outputs)
|
||||
if not out.success
|
||||
],
|
||||
"throughput_qps": num_success / total_duration if total_duration > 0 else 0,
|
||||
"latency_mean": np.mean(latencies) if latencies else 0,
|
||||
"latency_median": np.median(latencies) if latencies else 0,
|
||||
@@ -1267,9 +1261,6 @@ def _default_endpoint_for_task(task: str) -> str:
|
||||
|
||||
|
||||
async def benchmark(args):
|
||||
if not math.isfinite(args.video_job_timeout) or args.video_job_timeout <= 0:
|
||||
raise ValueError("--video-job-timeout must be a positive finite number of seconds.")
|
||||
|
||||
# Construct base_url if not provided
|
||||
if args.base_url is None:
|
||||
args.base_url = f"http://{args.host}:{args.port}"
|
||||
@@ -1321,8 +1312,6 @@ async def benchmark(args):
|
||||
print("Loading requests...")
|
||||
requests_list = dataset.get_requests()
|
||||
print(f"Prepared {len(requests_list)} requests from {args.dataset} dataset.")
|
||||
for req in requests_list:
|
||||
req.video_job_timeout = args.video_job_timeout
|
||||
|
||||
if args.return_stage_metrics and args.endpoint in _STAGE_METRICS_ENDPOINTS:
|
||||
for req in requests_list:
|
||||
@@ -1387,8 +1376,6 @@ async def benchmark(args):
|
||||
metrics["model"] = args.model
|
||||
metrics["dataset"] = args.dataset
|
||||
metrics["task"] = args.task
|
||||
if args.endpoint == "/v1/videos":
|
||||
metrics["video_job_timeout"] = args.video_job_timeout
|
||||
if args.endpoint == "/v1/images/edits":
|
||||
metrics["bot_task"] = args.bot_task
|
||||
|
||||
@@ -1411,11 +1398,6 @@ async def benchmark(args):
|
||||
)
|
||||
)
|
||||
print("{:<40} {}/{:<15}".format("Successful requests:", metrics["completed_requests"], len(requests_list)))
|
||||
print("{:<40} {:<15}".format("Failed requests:", metrics["failed_requests"]))
|
||||
for failure in metrics["request_errors"][:10]:
|
||||
print(f"Request {failure['request_id']} failed: {failure['error']}")
|
||||
if metrics["failed_requests"] > 10:
|
||||
print("Showing first 10 failures; use --output-file to save all request errors.")
|
||||
|
||||
# Section 3: Performance Metrics
|
||||
print(f"{'-' * 50}")
|
||||
@@ -1499,13 +1481,6 @@ if __name__ == "__main__":
|
||||
help="Path to local dataset file (optional).",
|
||||
)
|
||||
parser.add_argument("--num-prompts", type=int, default=10, help="Number of prompts to benchmark.")
|
||||
parser.add_argument(
|
||||
"--video-job-timeout",
|
||||
type=float,
|
||||
default=900.0,
|
||||
help="Maximum polling time in seconds per video job, including server queue time (default: 900). "
|
||||
"Increase this for long-running video generation or high concurrency.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-concurrency",
|
||||
type=int,
|
||||
|
||||
@@ -0,0 +1,57 @@
|
||||
# Host weight dependency memory diagnostic
|
||||
|
||||
`safetensors_retention.py` isolates repeated CPU `get_tensor()` calls from
|
||||
access to one cached tensor view. It imports neither vLLM nor HWR and uses a
|
||||
synthetic float32 payload controlled by `--tensor-elements` (default: 64).
|
||||
Tensor size can affect retention: a negative result for 16 elements does not
|
||||
rule out growth for 64 elements. The selected size is recorded in the JSON
|
||||
arguments and used for both file creation and the correctness check.
|
||||
HWR itself calls `get_tensor()` when
|
||||
acquiring a lease, then exposes cached views; this probe does not measure HWR
|
||||
requests or repeated lease acquisition.
|
||||
|
||||
Run one mode per fresh process in the same Python environment and CPU affinity:
|
||||
|
||||
```bash
|
||||
git rev-parse HEAD > revision.txt
|
||||
git status --short > working-tree.txt
|
||||
git diff > working-tree.patch
|
||||
# Select an allowed CPU from your process affinity; 0 is only an example.
|
||||
CUDA_VISIBLE_DEVICES='' taskset -c 0 timeout 120s python \
|
||||
benchmarks/host_weight_runtime/safetensors_retention.py \
|
||||
--mode get_tensor --iterations 10000 --sample-every 5000 > feasibility.json
|
||||
```
|
||||
|
||||
For a size comparison, run 16 and 64 elements twice each in fresh processes,
|
||||
with cached reuse as a control at both sizes:
|
||||
|
||||
```bash
|
||||
for repetition in 1 2; do
|
||||
for elements in 16 64; do
|
||||
for mode in get_tensor reuse; do
|
||||
CUDA_VISIBLE_DEVICES='' taskset -c 0 timeout 120s python \
|
||||
benchmarks/host_weight_runtime/safetensors_retention.py \
|
||||
--mode "$mode" --tensor-elements "$elements" \
|
||||
> "$mode-$elements-$repetition.json"
|
||||
done
|
||||
done
|
||||
done
|
||||
```
|
||||
|
||||
The diagnostic sets Torch to one CPU thread, warms up before the baseline,
|
||||
and reports preparation separately from timed loop work. Checkpoints run GC
|
||||
and report Linux private/anonymous memory, RSS, and open descriptors. It checks
|
||||
that transient tensor objects and the final cached view are released. Linux
|
||||
`/proc/self/smaps_rollup` must be readable; invalid arguments and missing proc
|
||||
support fail explicitly. Keep raw JSON and the repository snapshot with results.
|
||||
|
||||
Compare the final loop sample with the warmed baseline, then inspect the closed
|
||||
sample. Private-memory growth after Python tensor objects disappear is evidence
|
||||
of retention in the dependency/allocator path, not proof of a leak's root cause.
|
||||
RSS can include file-backed pages; it is not a private-memory metric. Samples
|
||||
also include small diagnostic bookkeeping allocations. Do not assert a portable
|
||||
memory threshold or turn this synthetic loop into a per-request service estimate.
|
||||
|
||||
Do not drop shared caches, trim allocators, or add serving-time GC to make the
|
||||
numbers look smaller. A dependency replacement or version change requires its
|
||||
own reproducible comparison and correctness/lifecycle validation.
|
||||
@@ -0,0 +1,120 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Linux CPU diagnostic: repeated safetensors materialization versus view reuse."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import gc
|
||||
import json
|
||||
import os
|
||||
import platform
|
||||
import tempfile
|
||||
import time
|
||||
import weakref
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def positive_int(value: str) -> int:
|
||||
parsed = int(value)
|
||||
if parsed <= 0:
|
||||
raise argparse.ArgumentTypeError("must be positive")
|
||||
return parsed
|
||||
|
||||
|
||||
def sample(phase: str, iterations: int, loop_seconds: float) -> dict[str, object]:
|
||||
gc.collect()
|
||||
counters = {}
|
||||
for line in Path("/proc/self/smaps_rollup").read_text().splitlines():
|
||||
fields = line.split()
|
||||
if len(fields) == 3 and fields[2] == "kB":
|
||||
counters[fields[0].removesuffix(":")] = int(fields[1]) * 1024
|
||||
return {
|
||||
"phase": phase,
|
||||
"iterations": iterations,
|
||||
"loop_seconds": loop_seconds,
|
||||
"rss_bytes": counters["Rss"],
|
||||
"private_bytes": counters["Private_Clean"] + counters["Private_Dirty"],
|
||||
"anonymous_bytes": counters["Anonymous"],
|
||||
"fd_count": len(tuple(Path("/proc/self/fd").iterdir())),
|
||||
}
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--mode", choices=("get_tensor", "reuse"), required=True)
|
||||
parser.add_argument(
|
||||
"--tensor-elements", type=positive_int, default=64, help="Number of float32 elements in the synthetic tensor"
|
||||
)
|
||||
parser.add_argument("--iterations", type=positive_int, default=1_000_000)
|
||||
parser.add_argument("--sample-every", type=positive_int, default=100_000)
|
||||
parser.add_argument("--warmup", type=positive_int, default=1000)
|
||||
args = parser.parse_args()
|
||||
if not Path("/proc/self/smaps_rollup").is_file():
|
||||
parser.error("requires Linux /proc/self/smaps_rollup")
|
||||
|
||||
preparation_start = time.monotonic()
|
||||
import safetensors
|
||||
import torch
|
||||
from safetensors import safe_open
|
||||
from safetensors.torch import save_file
|
||||
|
||||
torch.set_num_threads(1)
|
||||
torch.set_num_interop_threads(1)
|
||||
samples = [sample("before_file", 0, 0.0)]
|
||||
with tempfile.TemporaryDirectory(prefix="safetensors-retention-") as directory:
|
||||
path = Path(directory) / "synthetic.safetensors"
|
||||
save_file({"weight": torch.arange(args.tensor_elements, dtype=torch.float32)}, str(path))
|
||||
with safe_open(path, framework="pt", device="cpu") as reader:
|
||||
cached = reader.get_tensor("weight")
|
||||
assert torch.equal(cached, torch.arange(args.tensor_elements, dtype=torch.float32))
|
||||
for _ in range(args.warmup):
|
||||
tensor = reader.get_tensor("weight") if args.mode == "get_tensor" else cached
|
||||
del tensor
|
||||
samples.append(sample("baseline", 0, 0.0))
|
||||
preparation_seconds = time.monotonic() - preparation_start
|
||||
completed = 0
|
||||
loop_seconds = 0.0
|
||||
while completed < args.iterations:
|
||||
count = min(args.sample_every, args.iterations - completed)
|
||||
start = time.monotonic()
|
||||
for _ in range(count):
|
||||
tensor = reader.get_tensor("weight") if args.mode == "get_tensor" else cached
|
||||
last_tensor = weakref.ref(tensor)
|
||||
del tensor
|
||||
loop_seconds += time.monotonic() - start
|
||||
completed += count
|
||||
snapshot = sample("loop", completed, loop_seconds)
|
||||
snapshot["last_tensor_alive"] = last_tensor() is not None
|
||||
if args.mode == "get_tensor":
|
||||
assert last_tensor() is None, "diagnostic retained a transient tensor"
|
||||
samples.append(snapshot)
|
||||
del cached
|
||||
del reader
|
||||
samples.append(sample("closed", completed, loop_seconds))
|
||||
assert last_tensor() is None, "last tensor survived reader/view teardown"
|
||||
|
||||
print(
|
||||
json.dumps(
|
||||
{
|
||||
"schema_version": 1,
|
||||
"arguments": vars(args),
|
||||
"python": platform.python_version(),
|
||||
"torch": torch.__version__,
|
||||
"safetensors": safetensors.__version__,
|
||||
"kernel": platform.release(),
|
||||
"hostname": platform.node(),
|
||||
"pid": os.getpid(),
|
||||
"cpu_affinity": sorted(os.sched_getaffinity(0)),
|
||||
"torch_threads": torch.get_num_threads(),
|
||||
"pythonmalloc": os.environ.get("PYTHONMALLOC", "default"),
|
||||
"preparation_seconds": preparation_seconds,
|
||||
"samples": samples,
|
||||
},
|
||||
indent=2,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -85,25 +85,6 @@ vllm bench serve --omni \
|
||||
--save-result --result-dir ./results
|
||||
```
|
||||
|
||||
#### Set the sample rate for models that are not 24 kHz
|
||||
|
||||
Streamed PCM carries no header, so the harness derives `audio_duration` from the
|
||||
byte count and a **default 24000 Hz** assumption
|
||||
(`vllm_omni/metrics/definitions.py:DEFAULT_AUDIO_SAMPLE_RATE`). For a model that
|
||||
emits another rate, export the real one or every duration-derived metric is
|
||||
wrong by `real_sr / 24000`:
|
||||
|
||||
```bash
|
||||
# Audio8 TTS, Fish Speech S2 Pro, Ming-omni-tts, ... all decode at 44.1 kHz
|
||||
export VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE=44100
|
||||
```
|
||||
|
||||
Without it, a 44.1 kHz model reports `audio_duration` 1.84x too long, so
|
||||
`audio_rtf` looks 1.84x *better* than reality and `audio_underrun` /
|
||||
`Streaming continuity OK rate` are measured against inflated chunk budgets.
|
||||
A quick sanity check: compare `Mean AUDIO_DURATION (s)` against the actual
|
||||
length of a saved WAV for the same prompt.
|
||||
|
||||
#### Streaming continuity (`audio_underrun`)
|
||||
|
||||
`audio_underrun` (seconds) is the per-request worst-case buffer deficit
|
||||
|
||||
@@ -17,17 +17,6 @@ models:
|
||||
extra_params:
|
||||
lang: en
|
||||
text_normalization: true
|
||||
# Audio8 TTS Preview: DualAR, 44.1 kHz. trust_remote_code stays false.
|
||||
# Bench at 44.1 kHz: export VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE=44100, otherwise
|
||||
# the streamed-PCM duration defaults to 24 kHz and RTF looks 1.84x too good.
|
||||
Audio8/Audio8-TTS-Preview-0.6b:
|
||||
supported_tasks: [default_voice, voice_clone]
|
||||
backend: openai-audio-speech
|
||||
endpoint: /v1/audio/speech
|
||||
trust_remote_code: false
|
||||
task_extra_body:
|
||||
default_voice: {}
|
||||
voice_clone: {}
|
||||
|
||||
# -CustomVoice checkpoints lack speaker_encoder weights, so voice_clone is
|
||||
# NOT supported (an attempt raises ValueError from _extract_speaker_embedding
|
||||
|
||||
@@ -127,7 +127,6 @@ nav:
|
||||
- Async Chunk: design/feature/async_chunk.md
|
||||
- Async Diffusion Output: design/feature/async_diffusion_output.md
|
||||
- Async Omni Output Materialization: design/feature/omni_async_output_materialization.md
|
||||
- Prefill/Decode Phase Contract: design/feature/preprocess_phase_contract.md
|
||||
- Automatic Prefix Caching: design/feature/prefix_caching.md
|
||||
- Realtime AR-Diffusion: design/feature/realtime_ar_diffusion.md
|
||||
- Communication:
|
||||
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 272 KiB After Width: | Height: | Size: 274 KiB |
@@ -436,76 +436,6 @@ This should be seen as an edge case, and if this behavior can be avoided by sett
|
||||
|
||||
</details>
|
||||
|
||||
### Image / video generation
|
||||
|
||||
For `/v1/images/generations`, `/v1/images/edits`, and `/v1/videos`, set `--endpoint` to that path and omit `--backend`. The endpoint path is registered as the bench request adapter, so a separate named backend is not required.
|
||||
|
||||
Example (`/v1/images/generations`):
|
||||
|
||||
```bash
|
||||
vllm bench serve --omni \
|
||||
--endpoint /v1/images/generations \
|
||||
--dataset-name random \
|
||||
--model ~/models/Qwen/Qwen-Image \
|
||||
--tokenizer ~/models/Qwen/Qwen-Image/tokenizer \
|
||||
--max-concurrency 1 \
|
||||
--num-warmups 2 \
|
||||
--num-prompts 10 \
|
||||
--random-input-len 64 \
|
||||
--random-output-len 1 \
|
||||
--ignore-eos \
|
||||
--percentile-metrics e2el \
|
||||
--extra-body '{
|
||||
"num_inference_steps": 20,
|
||||
"seed": 42,
|
||||
"true_cfg_scale": 4.0
|
||||
}'
|
||||
```
|
||||
|
||||
If successful, following output is like:
|
||||
|
||||
```text
|
||||
============ Serving Benchmark Result ============
|
||||
Successful requests: 3
|
||||
Failed requests: 0
|
||||
Maximum request concurrency: 1
|
||||
Benchmark duration (s): 8.42
|
||||
Request throughput (req/s): 0.36
|
||||
Peak concurrent requests: 2.00
|
||||
-------------------Peak Memory--------------------
|
||||
Mean PEAK_MEMORY_MB (MB): 58832.00
|
||||
Median PEAK_MEMORY_MB (MB): 58832.00
|
||||
P99 PEAK_MEMORY_MB (MB): 58832.00
|
||||
----------------End-to-end Latency----------------
|
||||
Mean E2EL (ms): 2798.16
|
||||
Median E2EL (ms): 2799.83
|
||||
P99 E2EL (ms): 2800.03
|
||||
================== Image Result ==================
|
||||
Total images generated: 3
|
||||
Image throughput (img/s): 0.36
|
||||
Average pixels per image: 1048576.00
|
||||
Mean denoise step latency (ms): 136.59
|
||||
---------------- Image Generation ----------------
|
||||
Mean IMAGE_GENERATION (ms): 2731.84
|
||||
Median IMAGE_GENERATION (ms): 2732.02
|
||||
P99 IMAGE_GENERATION (ms): 2734.65
|
||||
==================================================
|
||||
```
|
||||
|
||||
Use the same pattern with `--endpoint /v1/images/edits` or `--endpoint /v1/videos` (and model-specific `--extra-body` as needed). Pure image/video runs omit the Text Result section when there is no generated text.
|
||||
|
||||
`/v1/images/edits` defaults to **non-streaming JSON** (`stream=false`) so single-stage edit models are not rejected by the server. For multi-stage pipelines that need SSE (AR TTFT / image chunks), pass `"stream": true` in `--extra-body`.
|
||||
|
||||
`/v1/videos` is an async job API: the client creates a job, then polls until `completed`/`failed`. Measured **e2el** therefore includes client poll sleep and any overshoot after the job actually finishes. Tune polling via `--extra-body`:
|
||||
|
||||
- `poll_interval_s` (default `2.0`): sleep between status polls
|
||||
- `poll_timeout_s` (default `21600`, i.e. 6 hours): give up waiting for the job
|
||||
|
||||
**VIDEO_RTF** prefers server-reported generation time when available:
|
||||
|
||||
- `video_rtf = video_generation_time_ms / 1000 / video_duration`, where `video_generation_time_ms` comes from response `stage_durations` (or `inference_time_s` as fallback)
|
||||
- If generation time is missing, falls back to `e2el / video_duration` (this fallback *does* include poll overhead)
|
||||
|
||||
### Multi-Stage Benchmark
|
||||
|
||||
<details class="admonition abstract" markdown="1">
|
||||
|
||||
@@ -262,28 +262,6 @@ or serving, explain the required L2/L3 environment and provide the exact command
|
||||
|
||||
Repository skills guide the agent, but they do not replace contributor judgment, the accepted issue or RFC, required test evidence, or maintainer review. Review all generated changes before submitting them.
|
||||
|
||||
## AI Usage Policy
|
||||
|
||||
AI tools are welcome in vLLM-Omni contributions. Contributors, including maintainers, remain responsible for everything they submit. These requirements apply to pull requests, issues, discussions, and review comments:
|
||||
|
||||
- **Disclose AI assistance.** Name the tools used and describe their role in the submission, including assistance with code, tests, documentation, or writing.
|
||||
- **Understand your changes.** Review every generated change and be able to explain its behavior, design choices, and interaction with vLLM-Omni without relying on AI to answer for you.
|
||||
- **Validate before submitting.** Run relevant checks and report actual results. For model, kernel, or performance changes, include the hardware, model, configuration, and commands needed to reproduce accuracy or benchmark evidence. Clearly identify checks you could not run; generated claims are not test results.
|
||||
- **Keep communication useful.** Verify technical claims, search for existing reports, and edit AI-assisted text for accuracy and brevity. Submit only content you have personally reviewed, and take responsibility for responding to review feedback.
|
||||
- **Label generated media.** Model-generated images, video, and audio are valid examples or test artifacts for this multimodal project. Identify them as generated output and provide the model, inputs, and reproduction settings. Do not present them as ground-truth references or independent evidence of correctness.
|
||||
|
||||
For example, a PR description could include:
|
||||
|
||||
```text
|
||||
AI assistance: Used Codex to draft the regression test and PR description.
|
||||
I reviewed the changes and verified the test reproduces the reported failure.
|
||||
Validation: <commands, results, hardware, and any checks not run>.
|
||||
```
|
||||
|
||||
Maintainer review time is limited. Contributions that are unverified, misleading, or beyond the author's understanding may be closed. Repeated violations may lead to contribution restrictions. New contributors are welcome to ask for help and submit small, well-understood changes.
|
||||
|
||||
This policy is adapted from [Ghostty's AI Usage Policy](https://github.com/ghostty-org/ghostty/blob/main/AI_POLICY.md) for vLLM-Omni's contribution and multimodal validation workflows. The [coding-agent workflow](#using-repository-skills-with-coding-agents) provides practical guidance for applying it.
|
||||
|
||||
## Issues
|
||||
|
||||
If you encounter a bug or have a feature request, please search existing issues first to see if it has already been reported. If not, please file a new issue, providing as much relevant information as possible.
|
||||
|
||||
@@ -779,31 +779,6 @@ turning them on. For Qwen-Image-style serving examples, document
|
||||
`--step-execution` as the feature gate and `--max-num-seqs N` as the
|
||||
companion batching knob.
|
||||
|
||||
### Streaming Output and Mid-Stream Interactions
|
||||
|
||||
To add streaming generation with optional mid-stream prompt updates to a new diffusion pipeline (often a video generation pipeline):
|
||||
|
||||
0. **Model should internally support** chunked video emission and bounded VAE decode of a chunk.
|
||||
1. **Implement stepwise execution** so interaction RPCs can be consumed at chunk
|
||||
boundaries (see [Step Execution](#step-execution) above).
|
||||
2. **Mix in** ``InteractionMixin`` on the pipeline class.
|
||||
3. **Implement** ``peek_chunk_media(StepRequestState) -> ChunkMediaSpec`` to offer information
|
||||
about the next chunk to generate.
|
||||
This is called whenever one supported interaction modality needs frame-level calculation,
|
||||
e.g., per-frame camera trajectory. If only chunk-level calculation is involved, peeking is skipped.
|
||||
4. **Optionally implement** ``prepare_next_chunk(StepRequestState)`` if the
|
||||
pipeline must rebuild something (e.g., per-chunk state) after an interaction is applied.
|
||||
This function is called at chunk boundary.
|
||||
5. **Register** supported modalities and handlers in ``vllm_omni/diffusion/interaction/registry.py``
|
||||
(keyed by ``od_config.model_class_name`` → modality key as in ``req["multi_modal_data"]``).
|
||||
|
||||
Serve with ``--diffusion-streaming-output`` (auto-enables step execution). See
|
||||
[Diffusion Execution Modes](../../user_guide/diffusion/execution_modes.md#streaming-output)
|
||||
and the example script at `examples/online_serving/streaming_video_generation/`.
|
||||
|
||||
**Interaction payload contract:** Although ``OmniInteractionPrompt`` type annocation requires ``event_id``,
|
||||
the WebSocket API payload may omit ``event_id`` and let the API layer assigns one.
|
||||
|
||||
### Cache Acceleration
|
||||
|
||||
#### TeaCache
|
||||
|
||||
@@ -222,16 +222,6 @@ __all__ = ["Qwen3OmniMoeForConditionalGeneration"]
|
||||
|
||||
## Key Components
|
||||
|
||||
### Preprocess phase metadata
|
||||
|
||||
Models with `has_preprocess=True`, including out-of-tree plugins, can rely on
|
||||
the runner's stable `_omni_is_prefill`, `_omni_prompt_len`, and
|
||||
`_omni_num_computed_tokens` fields. Use request progress to distinguish prompt
|
||||
prefill from decode: a one-token span can still be a prefill tail. See the
|
||||
[runner-to-model phase contract](../../design/feature/preprocess_phase_contract.md)
|
||||
for types, batched-hook semantics, MTP routing guarantees, and the correctness
|
||||
workaround for vLLM-Omni 0.20.x and earlier.
|
||||
|
||||
### 1. Model Interfaces
|
||||
|
||||
Your model should implement the appropriate interfaces:
|
||||
|
||||
@@ -22,15 +22,11 @@ A pipeline opts in by implementing `SupportsARDiffusionPipeline`. Its immutable
|
||||
- named fixed-length cross-attention KV, model-specific scratch geometry, and
|
||||
the pipeline's requested session-capacity upper bound.
|
||||
|
||||
The runner binds an `ARDiffusionKVState` only around one runner invocation
|
||||
(`execute_model` or `execute_stepwise`). The pipeline may use it during that
|
||||
context but must not retain it. Scratch and committed pages stay on the
|
||||
runner-owned session object, so unbinding does not drop an in-flight stepwise
|
||||
chunk. Tick-mode sessions survive subsequent requests with the same
|
||||
`session_id`. On the stepwise path the session id is the request id and is
|
||||
released when that request finishes or fails. Request reset, explicit close,
|
||||
LRU eviction, and forward exceptions all release runner-owned KV and call the
|
||||
pipeline's reset/close lifecycle hook. Cross-attention KV is allocated lazily
|
||||
The runner binds an `ARDiffusionKVState` only around one `forward()` call. The
|
||||
pipeline may use it during that context but must not retain it. The session
|
||||
survives subsequent requests with the same `session_id`. Request reset, explicit
|
||||
close, LRU eviction, and forward exceptions all release runner-owned KV and call
|
||||
the pipeline's reset/close lifecycle hook. Cross-attention KV is allocated lazily
|
||||
per session and logical KV branch, and is released on the same path. A pipeline
|
||||
publishes one named cache transactionally by yielding exactly one `(k, v)` pair
|
||||
per layer:
|
||||
@@ -52,12 +48,11 @@ successful forward. The runner also exposes `reset_session()` and
|
||||
|
||||
## Current execution limits
|
||||
|
||||
The experimental runner supports `max_num_seqs=1`. Request-batch execution is
|
||||
still rejected. Step execution is allowed only when the loaded pipeline
|
||||
implements `SupportsStepExecution`; other AR-Diffusion pipelines still fail at
|
||||
load if `step_execution=True`. Capable pipelines bind KV around each
|
||||
`execute_stepwise` invocation instead of inheriting the generic runner entry
|
||||
point unwrapped.
|
||||
The experimental runner supports request-mode execution with
|
||||
`max_num_seqs=1`. It rejects step execution and request-batch execution at load
|
||||
time, and its inherited batch/step entry points also fail explicitly. Supporting
|
||||
those paths requires a batch-aware state-binding contract rather than silently
|
||||
bypassing `bind_ar_diffusion_state()`.
|
||||
|
||||
The implementation may retain multiple sessions per model instance. It selects
|
||||
the largest feasible capacity up to the pipeline's `session_capacity`, based on
|
||||
|
||||
@@ -323,24 +323,18 @@ inputs, fields, and ownership rules are described below.
|
||||
```mermaid
|
||||
flowchart TB
|
||||
layer1["Layer 1 · Authoring inputs<br/>PipelineConfig + DeployConfig"]
|
||||
layer2["Layer 2 · Production resolution boundary<br/>resolve_omni_config()<br/>(StageConfigFactory + VllmOmniConfig internally)"]
|
||||
layer3["Layer 3 · Startup hand-off<br/>OmniConfigResolution<br/>PipelineConfig + temporary OmegaConf stage bridge"]
|
||||
layer2["Layer 2 · Resolve once<br/>StageConfigFactory.create_from_model()<br/>VllmOmniConfig.from_pipeline_config()"]
|
||||
layer3["Layer 3 · Transport-safe control plane<br/>VllmOmniConfig"]
|
||||
layer4["Layer 4 · Runtime launch planning<br/>StageRuntime"]
|
||||
layer5["Layer 5 · Engine materialization<br/>VllmConfig / OmniDiffusionConfig"]
|
||||
|
||||
layer1 --> layer2 --> layer3 --> layer4 --> layer5
|
||||
```
|
||||
|
||||
The production resolution boundary is `resolve_omni_config()` in
|
||||
The single resolution boundary is implemented by
|
||||
`StageConfigFactory.create_from_model()` and
|
||||
`VllmOmniConfig.from_pipeline_config()` in
|
||||
[`vllm_omni/config`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/config).
|
||||
It delegates typed construction to `StageConfigFactory.create_from_model()` and
|
||||
`VllmOmniConfig.from_pipeline_config()`, then returns an
|
||||
`OmniConfigResolution` consumed by both `AsyncOmniEngine` and headless startup.
|
||||
Until `StageRuntime` consumes typed stage configs directly, this envelope carries
|
||||
the effective `PipelineConfig` alongside OmegaConf-compatible `stage_configs` as
|
||||
a temporary runtime bridge; both views describe the same resolved topology,
|
||||
including injected stages.
|
||||
|
||||
The legacy `stage_args` YAML path has been removed. Model topology now resolves
|
||||
through `PipelineConfig`, with runtime overrides supplied by `DeployConfig`.
|
||||
|
||||
@@ -360,12 +354,9 @@ The important ownership rules are:
|
||||
deploy-time defaults; it does not redefine the model graph.
|
||||
3. CLI and Python overrides are applied at the resolution boundary, with
|
||||
per-stage overrides taking precedence over global values where supported.
|
||||
4. `OmniConfigResolution` is the sole startup hand-off. Its `pipeline_config`
|
||||
and temporary `stage_configs` compatibility view must describe the same
|
||||
topology.
|
||||
5. `StageRuntime` owns launch planning and replica lifecycle. `ReplicaInitPlan`
|
||||
4. `StageRuntime` owns launch planning and replica lifecycle. `ReplicaInitPlan`
|
||||
is runtime-private state, not a user configuration object.
|
||||
6. `VllmConfig` and the enriched `OmniDiffusionConfig` are materialized in the
|
||||
5. `VllmConfig` and the enriched `OmniDiffusionConfig` are materialized in the
|
||||
process that owns the corresponding engine.
|
||||
|
||||
## Main features
|
||||
|
||||
@@ -42,6 +42,11 @@ V1 does not include:
|
||||
|
||||
## Motivation and use cases
|
||||
|
||||
For investigating CPU memory retained by dependency tensor materialization,
|
||||
see the [standalone safetensors diagnostic](https://github.com/vllm-project/vllm-omni/blob/main/benchmarks/host_weight_runtime/README.md).
|
||||
It distinguishes repeated dependency calls from reuse of cached views and does
|
||||
not establish per-request HWR leakage.
|
||||
|
||||
Model loading can create the same final host representation repeatedly. This
|
||||
is especially expensive when loading performs checkpoint decoding, tensor
|
||||
renaming, TP slicing, quantization, packing, or scale construction before GPU
|
||||
|
||||
@@ -53,11 +53,6 @@ must use distinct hook names and remove only hooks they own.
|
||||
|
||||
## Discovery and topology
|
||||
|
||||
`resolve_offload_plan(pipeline, config)` is the single topology entry point.
|
||||
Every backend calls it once when `enable()` starts and then consumes the
|
||||
resolved components, block stacks, resident heads, and staged-residency flags
|
||||
it returns. Backends do not re-read model declarations.
|
||||
|
||||
Pipeline component discovery prefers `SupportsComponentDiscovery` declarations:
|
||||
|
||||
- `_dit_modules`;
|
||||
@@ -66,26 +61,16 @@ Pipeline component discovery prefers `SupportsComponentDiscovery` declarations:
|
||||
- `_resident_modules`.
|
||||
|
||||
Dotted paths are supported. Legacy pipelines may use the fallback scan of
|
||||
well-known attribute names; it now warns once per pipeline class, and new
|
||||
integrations must declare components explicitly.
|
||||
well-known attribute names, but new integrations should declare components
|
||||
explicitly.
|
||||
|
||||
Block topology comes from the pipeline `OffloadPlan` first: `block_attrs` maps
|
||||
each DiT path to its ordered block containers, `encoder_block_attrs` declares
|
||||
streamable encoder stacks, `on_demand_component_paths` marks pipeline-managed
|
||||
residency, `resident_dit_paths` marks the DiTs that may hold resident layers,
|
||||
and `encoder_dlo_weight_replication` marks the encoders whose loader-produced
|
||||
weights are safe for AllGather. DiTs absent from the plan fall back to
|
||||
`_layerwise_offload_blocks_attrs` (including the deprecated singular-name
|
||||
compatibility path).
|
||||
|
||||
The resolver is pure: it moves no tensor, installs no hook, writes no module
|
||||
attribute, and reads no process group — multi-rank facts come from
|
||||
`OffloadConfig`. It owns topology validation, so a configuration the model
|
||||
cannot serve fails before the first component is placed. Explicit component
|
||||
selection turns every declaration problem into an error, while the
|
||||
compatibility topology (no component selector) keeps the historical
|
||||
warn-and-skip behavior. Metadata still describes structure only; the backend
|
||||
remains responsible for transfer, synchronization, and storage ownership.
|
||||
Both layerwise backends first consume pipeline `OffloadPlan` metadata. A
|
||||
plan's `block_attrs` maps each DiT path to its ordered block containers, while
|
||||
`encoder_block_attrs` declares streamable encoder stacks. DiTs absent from the
|
||||
plan fall back to `_layerwise_offload_blocks_attrs` (including the deprecated
|
||||
singular-name compatibility path). Discovery metadata describes structure
|
||||
only; the backend remains responsible for transfer, synchronization, and
|
||||
storage ownership.
|
||||
|
||||
## Cross-strategy invariants
|
||||
|
||||
|
||||
@@ -1,172 +0,0 @@
|
||||
# Runner-to-model prefill/decode phase contract
|
||||
|
||||
## Overview
|
||||
|
||||
Autoregressive models with a custom `preprocess` hook need to distinguish
|
||||
prompt processing from decoding. A one-token span can be either a complete
|
||||
one-token prompt, the last chunk of a longer prompt, or a decode step. Using
|
||||
span length as the phase signal can replace a prompt embedding with a generated
|
||||
audio embedding and advance teacher-forced text too early without an exception.
|
||||
|
||||
[Issue #4382](https://github.com/vllm-project/vllm-omni/issues/4382) requests a
|
||||
stable contract for the metadata and routing guards introduced by
|
||||
[PR #3662](https://github.com/vllm-project/vllm-omni/pull/3662).
|
||||
The shared runner already implements those guards. This design documents the
|
||||
existing interface and adds runner-level regression coverage rather than
|
||||
introducing a second phase detector or changing scheduling policy.
|
||||
|
||||
## Scope and objectives
|
||||
|
||||
- Make the three existing phase keys safe to depend on in out-of-tree models.
|
||||
- Preserve prompt embeddings and model state on a one-token prefill tail.
|
||||
- Keep normal and batched decode routing consistent in mixed batches.
|
||||
- Explain the correctness workaround required on older releases.
|
||||
|
||||
This contract does not enable chunked prefill for models whose other state
|
||||
handling does not support it. It does not change the scheduler, introduce new
|
||||
kwargs, or standardize other `_omni_*` fields. The implementation adds
|
||||
documentation and regression coverage; no hot-path allocation is added.
|
||||
|
||||
## Design
|
||||
|
||||
### Ownership and data flow
|
||||
|
||||
```text
|
||||
Scheduler / current input batch
|
||||
prompt_token_ids + num_computed_tokens_cpu (before this step)
|
||||
|
|
||||
OmniGPUModelRunner._preprocess
|
||||
|
|
||||
per-request phase metadata
|
||||
/ \
|
||||
prefill / other spans single-token decode
|
||||
preprocess preprocess or preprocess_decode_batch
|
||||
| |
|
||||
prompt embedding runner-managed talker_mtp
|
||||
\ /
|
||||
model forward
|
||||
```
|
||||
|
||||
The producer is `OmniGPUModelRunner._preprocess` in
|
||||
`vllm_omni/worker/gpu_model_runner.py`. `GPUARModelRunner` and
|
||||
`GPUGenerationModelRunner` inherit it. The current `OmniNPUModelRunner` also
|
||||
inherits this method; device-specific execution still needs its own hardware
|
||||
validation. A plugin runner that overrides preprocessing is responsible for
|
||||
preserving this interface.
|
||||
|
||||
### Stable model-facing fields
|
||||
|
||||
The following names, types, and meanings are a supported interface for in-tree
|
||||
and out-of-tree model hooks despite the leading underscore. Models should treat
|
||||
them as read-only runner-owned values. Changes to their meanings or removal
|
||||
require an explicit compatibility/migration plan.
|
||||
|
||||
| Key | Python type | Meaning |
|
||||
| --- | --- | --- |
|
||||
| `_omni_prompt_len` | `int` | Length of the current request state's `prompt_token_ids`. This is the stage's prompt length, not the current chunk size or raw text length. |
|
||||
| `_omni_num_computed_tokens` | `int` | The row's `input_batch.num_computed_tokens_cpu` value **before** processing the scheduled span. It includes already-computed/cached progress, not just tokens computed in the current call. |
|
||||
| `_omni_is_prefill` | `bool` | Exactly `_omni_num_computed_tokens < _omni_prompt_len`. Equality and greater progress mean decode. |
|
||||
|
||||
The runner refreshes the fields after current request state has been gathered,
|
||||
immediately before selecting the per-row hook. Existing values in intermediate
|
||||
model state cannot override the current phase. Row indices follow the current
|
||||
input batch (including any reordering), not request arrival order. A resumed
|
||||
or streaming request is classified from its **current** prompt/progress values;
|
||||
the phase is not latched for the lifetime of the request.
|
||||
|
||||
The normal hook receives these fields as keyword arguments:
|
||||
|
||||
```python
|
||||
def preprocess(self, input_ids, input_embeds, **info):
|
||||
if info["_omni_is_prefill"]:
|
||||
offset = info["_omni_num_computed_tokens"]
|
||||
# Select the scheduled slice of the model's prompt embeddings.
|
||||
...
|
||||
else:
|
||||
# Advance decode-only text/audio state.
|
||||
...
|
||||
```
|
||||
|
||||
For `preprocess_decode_batch(input_ids=..., req_infos=...)`, each `req_infos`
|
||||
entry contains the same fields, aligned with its corresponding input row.
|
||||
Every such row has `_omni_is_prefill == False`. The optional earlier
|
||||
`preprocess_batch` hook is a separate whole-batch preparation hook; these freshly
|
||||
computed per-row fields are not promised at that earlier boundary.
|
||||
|
||||
### Routing guarantee
|
||||
|
||||
With `has_preprocess=True`, the batched decode hook is eligible only when it is
|
||||
callable, the runner has `talker_mtp`, the scheduled span is one token, and
|
||||
`_omni_is_prefill` is false. Otherwise the normal `preprocess` hook runs.
|
||||
|
||||
Runner-managed MTP is likewise restricted to single-token decode rows. A prefill
|
||||
row, including a one-token tail, is never added to that MTP batch: the runner
|
||||
does not consume its `mtp_inputs`, generate codes for it, or overwrite its
|
||||
preprocessed embedding with an MTP embedding. Decode rows in the same step
|
||||
continue to receive their MTP outputs at their own token offsets. Models may
|
||||
still implement their own model-specific work inside `preprocess`.
|
||||
|
||||
| Prompt length | Computed before step | Scheduled span | Phase | Runner MTP eligible |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| 5 | 0 | 4 | Prefill | No |
|
||||
| 5 | 4 | 1 | Prefill tail | No |
|
||||
| 5 | 5 | 1 | Decode | Yes |
|
||||
| 5 | 6 | 1 | Decode | Yes |
|
||||
| 1 | 0 | 1 | Prefill | No |
|
||||
| 5 | 5 | 2 | Decode | No (not a single-token span) |
|
||||
|
||||
### Compatibility
|
||||
|
||||
!!! warning "vLLM-Omni 0.20.x and earlier"
|
||||
AR talker models with custom preprocessing that infer phase from span length
|
||||
must use `enable_chunked_prefill: false` for correctness on these releases.
|
||||
They lack the phase metadata and the runner's prefill-tail MTP protection.
|
||||
A model-only heuristic cannot prevent the older runner from overwriting a
|
||||
one-token prompt embedding. Disabling chunking avoids split tails; it does
|
||||
not make span length a reliable phase detector for every possible prompt.
|
||||
|
||||
The mechanism landed in #3662 (v0.21.0rc2 / v0.22.0 and later); this document
|
||||
formalizes that existing behavior. Prefer upgrading the matching vLLM and
|
||||
vLLM-Omni release pair and consuming the explicit phase keys. Do not silently
|
||||
fall back to `input_ids.numel() > 1` when the model requires chunking correctness.
|
||||
For a custom or backported runtime, verify both metadata and MTP guards.
|
||||
|
||||
## Correctness and testing plan
|
||||
|
||||
Runner tests call the production `_preprocess` and `_talker_mtp_forward` methods
|
||||
with CPU buffers and a deterministic MTP model. They inspect kwargs at the model
|
||||
boundary, resulting embeddings, and per-request generated codes. Coverage must
|
||||
include a multi-step prompt split with a one-token tail followed by decode,
|
||||
mixed prefill/decode batches in different orders, normal/batched decode hooks,
|
||||
adjacent decode requests in the same hook call with exact per-request codes,
|
||||
single-token prompts, cached progress, and decode spans longer than one token.
|
||||
No test should pass merely because MTP was disabled for every row.
|
||||
|
||||
Real-weight validation uses Qwen3-TTS CustomVoice through the full talker and
|
||||
code2wav pipeline. Record the actual scheduled prompt progress, assert that at
|
||||
least one one-token prefill tail was observed, check its hook metadata and MTP
|
||||
exclusion, and require finite non-empty output audio. Compare to an unchunked
|
||||
run with the same real input, greedy sampling, and `VLLM_BATCH_INVARIANT=1`.
|
||||
Greedy sampling alone does not guarantee identical codes across different
|
||||
prefill shapes: floating-point reduction differences can change an argmax and
|
||||
diverge in subsequent autoregressive steps. Invariant kernels keep this exact
|
||||
comparison focused on the phase contract. The two engines must be initialized
|
||||
separately because their chunking settings and token budgets differ. GPU tests belong under
|
||||
`tests/e2e/features/`; existing model tests remain useful adjacent regression
|
||||
coverage. Functional audio checks do not establish throughput or all-model
|
||||
semantic parity; this change makes no performance claim.
|
||||
|
||||
## Open questions and discussions
|
||||
|
||||
There are no new field names or execution modes to choose. Future runner
|
||||
implementations, including Model Runner V2, should preserve these semantics or
|
||||
provide an explicit migration path. Hardware validation beyond the tested
|
||||
platform remains separate from the shared Python routing tests.
|
||||
|
||||
## References
|
||||
|
||||
- [Issue #4382](https://github.com/vllm-project/vllm-omni/issues/4382)
|
||||
- [Original metadata and guards, #3662](https://github.com/vllm-project/vllm-omni/pull/3662)
|
||||
- [Qwen2.5-Omni consumer fix, #5019](https://github.com/vllm-project/vllm-omni/pull/5019)
|
||||
- [Adding an Omni model](../../contributing/model/adding_omni_model.md)
|
||||
- [Test writing guide](../../contributing/ci/test_writing_guide.md)
|
||||
@@ -66,8 +66,7 @@ The checkpoint keeps these verified contracts:
|
||||
- scheduler data-plane append over a resumable request;
|
||||
- stale epoch/turn/response fencing;
|
||||
- `/v1/duplex` and OpenAI Realtime projection;
|
||||
- optional OpenAI `server_vad` turn detection backed by a process-shared Silero
|
||||
model and per-session VAD state;
|
||||
- optional OpenAI `server_vad` turn detection backed by per-session Silero VAD;
|
||||
- existing JoyVL behavior.
|
||||
|
||||
The checkpoint does not claim:
|
||||
@@ -653,7 +652,7 @@ EOS decisions advance the model conversation.
|
||||
|
||||
Explicit Realtime clients may still use `input_audio_buffer.commit` to create a
|
||||
conversation item. The translator validates that wire input before producing a
|
||||
commit carrying `item_id`. Native auto-response may already have
|
||||
commit carrying `realtime_item_id`. Native auto-response may already have
|
||||
streamed those PCM samples into the runtime, so the runner accepts a validated
|
||||
commit even when no runtime-side chunk remains. A truly empty explicit buffer
|
||||
continues to return `input_audio_buffer_empty`.
|
||||
@@ -678,7 +677,8 @@ multi-turn runtime evidence.
|
||||
|
||||
During auto-response overlap, `preserve_realtime_input` distinguishes "do not
|
||||
append this silent chunk to the native buffer" from "clear the open Realtime
|
||||
item". Silent overlap no longer discards earlier user PCM.
|
||||
item". Silent overlap no longer discards earlier user PCM. A separate explicit
|
||||
`server_vad` policy is documented in the Realtime contract above.
|
||||
|
||||
The first chunk of one overlapping input item also reserves its target model
|
||||
turn. A later Realtime commit uses that reserved identity even if response EOS
|
||||
|
||||
@@ -24,7 +24,6 @@ implementation contract; it is not, by itself, a general support claim.
|
||||
- [Async Chunk](feature/async_chunk.md)
|
||||
- [Async Diffusion Output](feature/async_diffusion_output.md)
|
||||
- [Async Omni Output Materialization](feature/omni_async_output_materialization.md)
|
||||
- [Runner-to-model Prefill/Decode Phase Contract](feature/preprocess_phase_contract.md)
|
||||
- [Automatic Prefix Caching in Omni Models](feature/prefix_caching.md)
|
||||
- [Realtime AR-Diffusion Sessions](feature/realtime_ar_diffusion.md)
|
||||
|
||||
|
||||
@@ -8,44 +8,22 @@ owners:
|
||||
- "@SamitHuang"
|
||||
- "@fhfuih"
|
||||
primary_code_paths:
|
||||
- vllm_omni/diffusion/diffusion_engine.py
|
||||
- vllm_omni/diffusion/request.py
|
||||
- vllm_omni/diffusion/data.py
|
||||
- vllm_omni/diffusion/sched/**
|
||||
- vllm_omni/diffusion/executor/**
|
||||
- vllm_omni/diffusion/worker/**
|
||||
related_code_paths:
|
||||
- vllm_omni/diffusion/stage_diffusion_client.py
|
||||
- vllm_omni/diffusion/stage_diffusion_proc.py
|
||||
- vllm_omni/diffusion/inline_stage_diffusion_client.py
|
||||
- vllm_omni/diffusion/ipc.py
|
||||
- vllm_omni/diffusion/hooks/**
|
||||
- vllm_omni/diffusion/postprocess/**
|
||||
depends_on:
|
||||
- ../engine_orchestration.md
|
||||
- ../stage_runtime.md
|
||||
- ../input_output_modality_contracts.md
|
||||
- diffusion_model_integration.md
|
||||
validation_paths:
|
||||
- tests/diffusion/test_diffusion_engine.py
|
||||
- tests/diffusion/test_diffusion_engine_cleanup.py
|
||||
- tests/diffusion/test_diffusion_engine_dummy_run.py
|
||||
- tests/diffusion/test_diffusion_engine_rpc_routing.py
|
||||
- tests/diffusion/test_diffusion_scheduler.py
|
||||
- tests/diffusion/test_diffusion_model_runner.py
|
||||
- tests/diffusion/test_diffusion_worker.py
|
||||
- tests/diffusion/test_multiproc_engine_concurrency.py
|
||||
- tests/diffusion/test_uniproc_executor.py
|
||||
- tests/diffusion/test_result_pump.py
|
||||
- tests/diffusion/test_async_output_worker.py
|
||||
- tests/diffusion/test_diffusion_ipc.py
|
||||
- tests/diffusion/test_diffusion_streaming_output.py
|
||||
- tests/diffusion/test_stage_diffusion_proc.py
|
||||
- tests/diffusion/test_inline_stage_diffusion_client.py
|
||||
- tests/diffusion/diffusion_backend/**
|
||||
- tests/diffusion/hooks/**
|
||||
upstream_refs:
|
||||
- diffusers.DiffusionPipeline
|
||||
last_reviewed: 2026-09-02
|
||||
last_verified_commit: e51fe6ec1b9a9a0e14bb1fdb296d61b6593b93c6
|
||||
last_reviewed: 2026-07-16
|
||||
---
|
||||
|
||||
# Diffusion runtime
|
||||
@@ -53,380 +31,6 @@ last_verified_commit: e51fe6ec1b9a9a0e14bb1fdb296d61b6593b93c6
|
||||
The diffusion runtime owns request admission, scheduling, execution, progress,
|
||||
output, cancellation, and cleanup inside a diffusion stage.
|
||||
|
||||
It is easiest to think of the runtime as a small control loop:
|
||||
|
||||
1. the **scheduler** chooses ready requests;
|
||||
2. the **executor** sends that choice to the worker layer;
|
||||
3. each **worker** manages its device and delegates model work to a **runner**;
|
||||
4. the runner calls the model pipeline and returns per-request results; and
|
||||
5. the engine feeds those results back to the scheduler and output streams.
|
||||
|
||||
!!! note
|
||||
|
||||
This page describes the runtime plumbing. Model loading and pipeline
|
||||
contracts belong to [Diffusion model integration](diffusion_model_integration.md).
|
||||
Batch compatibility and step-state details belong to
|
||||
[Diffusion continuous batching](continuous_batching.md).
|
||||
|
||||
## Goals and non-goals
|
||||
|
||||
The runtime aims to keep request policy separate from device execution. It
|
||||
also keeps request identity stable from admission through the final output,
|
||||
including cancellation and failures.
|
||||
|
||||
This module does **not** choose the next Omni stage, place stage replicas, or
|
||||
define a model's denoising algorithm:
|
||||
|
||||
- cross-stage routing belongs to [Engine orchestration](../engine_orchestration.md);
|
||||
- process placement and replica lifecycle belong to
|
||||
[Stage runtime](../stage_runtime.md); and
|
||||
- model-specific work belongs to
|
||||
[Diffusion model integration](diffusion_model_integration.md).
|
||||
|
||||
## Runtime at a glance
|
||||
|
||||
The stage owns `DiffusionEngine` and a `DiffusionExecutor` backend. The inline
|
||||
stage client keeps that stack in the caller process; a process-backed stage
|
||||
runs it inside `StageDiffusionProc`.
|
||||
|
||||
The executor then chooses how the worker runs:
|
||||
|
||||
- **`uni`** (default when `num_gpus == 1`): `UniProcDiffusionExecutor` builds
|
||||
one in-process worker. There is no worker subprocess and no shared-memory
|
||||
RPC.
|
||||
- **`mp`** (default when `num_gpus > 1`, or when set explicitly):
|
||||
`MultiprocDiffusionExecutor` starts one `WorkerProc` per device.
|
||||
|
||||
Either way, each worker owns its device, distributed state,
|
||||
`DiffusionWorker`, `DiffusionModelRunner`, and pipeline instance.
|
||||
|
||||
```mermaid
|
||||
flowchart TB
|
||||
client["Stage client"]
|
||||
|
||||
subgraph stage["Diffusion stage process"]
|
||||
stageProc["StageDiffusionProc<br/>stage transport and lifecycle"]
|
||||
engine["DiffusionEngine<br/>admission and control loop"]
|
||||
scheduler["RequestScheduler or StepScheduler<br/>request state and policy"]
|
||||
executor["DiffusionExecutor<br/>uni or mp backend"]
|
||||
streams["Per-request output streams"]
|
||||
|
||||
stageProc --> engine
|
||||
engine --> scheduler
|
||||
scheduler -->|"DiffusionSchedulerOutput"| engine
|
||||
engine --> executor
|
||||
engine --> streams
|
||||
end
|
||||
|
||||
subgraph workers["Worker layer"]
|
||||
workerProc["WorkerProc<br/>mp IPC loop only"]
|
||||
worker["DiffusionWorker<br/>device and distributed setup"]
|
||||
runner["DiffusionModelRunner<br/>model state and execution"]
|
||||
pipeline["Diffusion pipeline<br/>model-specific computation"]
|
||||
|
||||
workerProc --> worker
|
||||
worker --> runner --> pipeline
|
||||
end
|
||||
|
||||
client -->|"process-backed: ZMQ"| stageProc
|
||||
client -.->|"inline: direct call"| engine
|
||||
executor -->|"mp: broadcast / result queues"| workerProc
|
||||
executor -.->|"uni: in-process call"| worker
|
||||
streams -->|"process-backed"| stageProc
|
||||
stageProc -->|"ZMQ results"| client
|
||||
streams -.->|"inline"| client
|
||||
```
|
||||
|
||||
!!! info "Stage process vs worker process"
|
||||
|
||||
A deployed stage may already run in a `StageDiffusionProc` child process.
|
||||
That is separate from the optional `mp` worker processes. With `uni`, the
|
||||
worker lives in the same process as the engine.
|
||||
|
||||
## Component responsibilities
|
||||
|
||||
| Component | Owns | Does not own |
|
||||
| --- | --- | --- |
|
||||
| `DiffusionEngine` | Admission, the busy loop, scheduler/executor coordination, RPC ordering, output streams, warmup, abort routing | Batch policy details, device setup, model computation |
|
||||
| `BaseScheduler` | Request states, waiting/running sets, capacity, compatibility, terminal transitions | IPC, worker calls, output formatting |
|
||||
| `RequestScheduler` | Complete-request waves and optional admission delay for request-level batching | Denoising-step progress |
|
||||
| `StepScheduler` | Per-request denoising progress and step-wise completion | Pipeline tensors and model state |
|
||||
| `DiffusionExecutor` | Execution backend contract, worker RPC, health, shutdown | Admission and request-state transitions |
|
||||
| `UniProcDiffusionExecutor` | Single in-process worker for `num_gpus == 1`; no IPC or async output pump | Multi-GPU execution |
|
||||
| `MultiprocDiffusionExecutor` | Worker processes, shared-memory message queues, result dispatch, worker monitoring | In-process single-GPU path |
|
||||
| `WorkerProc` | One `mp` worker process's IPC loop and reply rules | Scheduling policy; unused by `uni` |
|
||||
| `DiffusionWorker` | Device/distributed setup, LoRA activation, profiling, sleep/wake, runner delegation | Request admission |
|
||||
| `DiffusionModelRunner` | Pipeline loading, request-local model state, cache/compile setup, request or step execution | Queueing and cross-stage routing |
|
||||
|
||||
The boundary between policy and execution is
|
||||
`DiffusionSchedulerOutput`. It contains newly admitted request payloads,
|
||||
IDs for requests whose runner state is already cached, finished IDs,
|
||||
optional KV-prefetch work, and per-request `diffusion_kv_metadata` on
|
||||
`NewRequestData` envelopes when paged Diffusion KV is enabled. Executors and
|
||||
workers consume this output; they do not decide what should run next.
|
||||
|
||||
`BaseScheduler.initialize()` sets `max_num_running_reqs` from
|
||||
`od_config.max_num_seqs` (default 1). The engine may override that capacity
|
||||
for distributed layerwise offload with AllGather DP concurrency.
|
||||
|
||||
## One scheduling tick
|
||||
|
||||
Async requests share one engine busy loop. Worker calls and control RPCs pass
|
||||
through that loop, so they cannot race on the executor transport.
|
||||
|
||||
```mermaid
|
||||
sequenceDiagram
|
||||
participant Caller
|
||||
participant Engine as DiffusionEngine
|
||||
participant Scheduler
|
||||
participant Executor
|
||||
participant Worker
|
||||
participant Runner
|
||||
|
||||
Caller->>Engine: add request
|
||||
Engine->>Scheduler: add_request(request)
|
||||
Engine->>Scheduler: schedule()
|
||||
Scheduler-->>Engine: DiffusionSchedulerOutput
|
||||
Engine->>Executor: execute_batch() or execute_step()
|
||||
Executor->>Worker: run worker method
|
||||
Worker->>Runner: execute model or one step
|
||||
Runner-->>Worker: RunnerOutput(s)
|
||||
Worker-->>Executor: result
|
||||
Executor-->>Engine: BaseRunnerOutput
|
||||
Engine->>Scheduler: update_from_output(...)
|
||||
Scheduler-->>Engine: finished request IDs
|
||||
Engine-->>Caller: chunk or final output
|
||||
```
|
||||
|
||||
The engine catches request execution errors and turns them into per-request
|
||||
error outputs. A dead worker group is different: the executor marks itself
|
||||
failed, health checks raise `EngineDeadError`, and the owning stage client
|
||||
handles the stage-level failure.
|
||||
|
||||
## Execution modes
|
||||
|
||||
The engine resolves one mode at startup and binds the matching scheduler and
|
||||
executor call.
|
||||
|
||||
| Mode | Scheduler | Executor call | Runner path |
|
||||
| --- | --- | --- | --- |
|
||||
| Request batch | `RequestScheduler` | `execute_batch()` | `execute_model()` for one request, or `execute_model_batch()` for a fused batch |
|
||||
| Step batch | `StepScheduler` | `execute_step()` | `execute_stepwise()` |
|
||||
|
||||
**Request mode** runs a complete pipeline forward for each scheduled wave. A
|
||||
single-request wave is the conservative path. A multi-request wave usually
|
||||
uses fused `execute_model_batch()` and requires the pipeline to declare
|
||||
request-batch support.
|
||||
|
||||
**Distributed layerwise offload with AllGather (DLO DP concurrency)** is a
|
||||
separate multiproc dispatch path. It activates when
|
||||
`data_parallel_size > 1`, `enable_distributed_layerwise_offload` is set, and
|
||||
`dlo_use_allgather` is true. The engine then sets `dp_concurrent = True` and
|
||||
raises `scheduler.max_num_running_reqs` to `dp_size`, overriding the
|
||||
`max_num_seqs` value from `initialize()`. The scheduler still emits a
|
||||
multi-request `DiffusionSchedulerOutput`; the multiproc executor routes that
|
||||
wave through `execute_request()` instead of fused `execute_model_batch()`.
|
||||
|
||||
Optional admission coalescing can fill a `dp_size` wave: when
|
||||
`request_batch_max_wait_ms > 0`, `RequestScheduler` may wait briefly before
|
||||
the first schedule of a wave (under `dp_concurrent`, the stable window is
|
||||
`min(0.3s, wait/2)`). With the default `request_batch_max_wait_ms == 0`,
|
||||
there is no wait.
|
||||
|
||||
Before dispatch, `MultiprocDiffusionExecutor.execute_request()` rejects waves
|
||||
whose requests differ in sampling-parameter compatibility or `extra_args`
|
||||
(AllGather requires every DP rank to follow the same forward schedule). Each
|
||||
worker rank then picks one envelope from the wave:
|
||||
`req[dp_rank % len(req)]`, so every rank enters the same collective while
|
||||
computing different requests.
|
||||
|
||||
**Step mode** keeps `StepRequestState` in the runner. New requests run
|
||||
`prepare_encode()` once; every tick runs `denoise_step()` and
|
||||
`step_scheduler()`; completed requests run `post_decode()`. The scheduler keeps
|
||||
only lifecycle and progress metadata—it does not hold model tensors.
|
||||
|
||||
**Step mode is also the path for streaming diffusion output**. Enabling
|
||||
`streaming_output` turns on step execution if needed. Chunk-capable pipelines
|
||||
can emit intermediate outputs through the same request stream; final-only step
|
||||
pipelines still emit only the final result. If the pipeline does not implement
|
||||
step execution, initialization fails.
|
||||
|
||||
!!! tip
|
||||
|
||||
For user-facing flags and examples, see
|
||||
[Diffusion execution modes](../../../user_guide/diffusion/execution_modes.md).
|
||||
That guide explains how to select a mode; this page explains what happens
|
||||
after the choice is made.
|
||||
|
||||
## Request state and identity
|
||||
|
||||
The scheduler is the source of truth for an admitted request. Its normal state
|
||||
flow is:
|
||||
|
||||
```mermaid
|
||||
stateDiagram-v2
|
||||
[*] --> WAITING: add_request
|
||||
WAITING --> RUNNING: schedule
|
||||
WAITING --> FINISHED_ABORTED: abort
|
||||
RUNNING --> FINISHED_ABORTED: abort
|
||||
RUNNING --> FINISHED_COMPLETED: successful output
|
||||
RUNNING --> FINISHED_ERROR: failed or missing output
|
||||
FINISHED_COMPLETED --> [*]: deliver and remove state
|
||||
FINISHED_ABORTED --> [*]: deliver and remove state
|
||||
FINISHED_ERROR --> [*]: deliver and remove state
|
||||
```
|
||||
|
||||
`PREEMPTED` and `BaseScheduler.preempt_request()` exist on the scheduler API,
|
||||
but the engine loop does not call them today—only scheduler unit tests do.
|
||||
|
||||
`request_id` ties together scheduler state, runner state, executor results, and
|
||||
the output queue. Batch positions are temporary and must never replace request
|
||||
identity when results are mapped back.
|
||||
|
||||
Only compatible requests enter the same running wave. The scheduler compares a
|
||||
mode-specific sampling-parameter key and stops at the first incompatible
|
||||
waiting request. This is intentionally conservative and can cause
|
||||
head-of-line blocking.
|
||||
|
||||
## Executor and IPC
|
||||
|
||||
`DiffusionExecutor` is the backend interface. Built-in backends are selected by
|
||||
`distributed_executor_backend`:
|
||||
|
||||
| Backend | When chosen | Worker layout |
|
||||
| --- | --- | --- |
|
||||
| `uni` | Default for `num_gpus == 1`, or set explicitly | One in-process worker |
|
||||
| `mp` | Default for `num_gpus > 1`, or set explicitly | One worker process per device |
|
||||
|
||||
Ray and external-launcher diffusion backends are not implemented. A custom
|
||||
`DiffusionExecutor` subclass or import path is also accepted.
|
||||
|
||||
### Uniproc backend
|
||||
|
||||
`UniProcDiffusionExecutor` constructs `WorkerWrapperBase` in the engine process
|
||||
and calls worker methods directly. It avoids a second model load, MessageQueue
|
||||
rings, ZMQ IPC sockets, and `/dev/shm` tensor packing. RPC timeouts are
|
||||
accepted for interface parity but are not enforced: a hung worker blocks the
|
||||
calling thread. Sticky accelerator faults mark the executor dead through
|
||||
failure callbacks; ordinary request errors stay per-request.
|
||||
|
||||
### Multiproc backend
|
||||
|
||||
`MultiprocDiffusionExecutor`:
|
||||
|
||||
1. creates a broadcast queue shared by all workers;
|
||||
2. spawns one worker process per configured device;
|
||||
3. waits until every worker reports ready;
|
||||
4. sends generation and control calls as RPC messages;
|
||||
5. applies rank-aware reply rules so only expected ranks respond; and
|
||||
6. monitors worker process sentinels and fails the executor if a worker dies.
|
||||
|
||||
In request mode on this backend, each async output produces two correlated
|
||||
messages keyed by `async_output_id`:
|
||||
|
||||
- `COMPUTE_DONE` — the forward finished and the device can start the next
|
||||
request.
|
||||
- `OUTPUT_READY` — background D2H/SHM packing finished and the final output
|
||||
is ready.
|
||||
|
||||
The worker queues background packing work before it enqueues `COMPUTE_DONE`,
|
||||
but either message may reach the executor first. The result pump and
|
||||
`execute_batch()` must handle both orderings. Step mode keeps the synchronous
|
||||
result path and does not start those pumps.
|
||||
|
||||
See [Async diffusion output](../../feature/async_diffusion_output.md) for the
|
||||
multiproc request-mode timeline.
|
||||
|
||||
## Worker and runner boundary
|
||||
|
||||
On the `mp` path, `WorkerProc` receives messages and calls methods through
|
||||
`WorkerWrapperBase`. On the `uni` path, the executor owns that wrapper
|
||||
directly. In both cases, `DiffusionWorker` handles the parts tied to a device:
|
||||
distributed initialization, model-runner construction, LoRA activation,
|
||||
profiling, and memory sleep/wake. It delegates actual model work to
|
||||
`DiffusionModelRunner`.
|
||||
|
||||
The runner owns long-lived model-side state:
|
||||
|
||||
- the loaded pipeline and compilation setup;
|
||||
- cache and offload integration;
|
||||
- random-generator setup;
|
||||
- request batches for complete forwards; and
|
||||
- `StepRequestState` plus `InputBatch` for step execution.
|
||||
|
||||
This split keeps infrastructure out of model pipelines. A pipeline receives a
|
||||
request batch or step state and performs model work; it should not inspect
|
||||
engine queues or change scheduler state.
|
||||
|
||||
## Lifecycle and cleanup
|
||||
|
||||
### Startup
|
||||
|
||||
`DiffusionEngine.make_engine()` resolves the engine class, constructs the
|
||||
executor and scheduler, then normally runs a small dummy request to warm up the
|
||||
model.
|
||||
|
||||
Warmup is skipped only when distributed layerwise offload with AllGather
|
||||
is enabled (via `enable_distributed_layerwise_offload` and `dlo_use_allgather`)
|
||||
and `max(data_parallel_size, sequence_parallel_size) > 1`, because
|
||||
the dummy run sends one request while AllGather needs every shard rank to
|
||||
enter the same collective. If initialization or warmup fails, the runtime
|
||||
closes scheduler state and worker resources before returning the error.
|
||||
|
||||
### Cancellation
|
||||
|
||||
`abort()` places request IDs on an abort queue. The busy loop marks those
|
||||
requests `FINISHED_ABORTED`; finalization removes scheduler state and emits an
|
||||
aborted output when a consumer still exists. The runner removes cached
|
||||
step-mode state when it receives a scheduler output that reports the finished
|
||||
request ID.
|
||||
|
||||
On the multiproc request-mode async path, abort or consumer drop while output
|
||||
is still materializing should release the associated async-output bookkeeping
|
||||
so late results are not retained after the request ends
|
||||
*(pending #6253/#6439/#6580)*. Today, an unconsumed `OUTPUT_READY` is cached
|
||||
by design until `wait_output_ready()` or executor teardown. Dropping an
|
||||
output-stream consumer removes that consumer's queue. It does not take
|
||||
ownership of scheduler cleanup.
|
||||
|
||||
### Shutdown and worker failure
|
||||
|
||||
`DiffusionEngine.close()` is idempotent. It stops the busy loop, wakes pending
|
||||
streams with an error, closes scheduler state, and shuts down the executor.
|
||||
|
||||
- **`mp`**: the executor asks workers to stop, waits for them, then terminates
|
||||
workers that miss the grace period. Shutdown marks unfinished RPC and
|
||||
async-output futures as failed, then clears them. Completed async outputs
|
||||
cached for later `wait_output_ready()` calls are dropped with the executor
|
||||
object.
|
||||
- **`uni`**: the executor shuts down the in-process worker, drops the final
|
||||
model reference, and empties the accelerator cache so a later engine can reuse
|
||||
the device.
|
||||
|
||||
!!! warning
|
||||
|
||||
Do not continue using a worker group after a fatal collective timeout or
|
||||
unexpected worker exit. Distributed state may be incomplete. The
|
||||
multiprocess executor deliberately fails closed so the stage can restart.
|
||||
The uniproc executor fails closed when the accelerator context itself is
|
||||
poisoned.
|
||||
|
||||
## Extension boundaries
|
||||
|
||||
- A custom engine may set `default_diffusion_model_runner_cls`; an explicit
|
||||
`diffusion_model_runner_cls` configuration still wins.
|
||||
- Tests and custom engine integrations may inject a `BaseScheduler` subclass.
|
||||
`SchedulerInterface` remains only as a deprecated compatibility name.
|
||||
- `distributed_executor_backend` may be `"uni"`, `"mp"`, a custom `DiffusionExecutor`
|
||||
subclass, or an import path. A backend must preserve scheduler output,
|
||||
request identity, health, and cleanup contracts. Use `"mp"` when you need
|
||||
process isolation or multi-GPU; `"uni"` is the single-GPU default.
|
||||
- Worker extensions go through `WorkerWrapperBase`. They add worker methods;
|
||||
they do not become a second scheduler or request lifecycle.
|
||||
|
||||
These are advanced Python integration points, not stable end-user CLI
|
||||
customization surfaces.
|
||||
|
||||
## Candidate invariants
|
||||
|
||||
### DIFF-RUNTIME-INV-001: One lifecycle owner
|
||||
@@ -437,50 +41,19 @@ until completion, cancellation, or failure.
|
||||
### DIFF-RUNTIME-INV-002: Execution follows scheduler output
|
||||
|
||||
**Rule:** Executors and workers MUST execute scheduler decisions without
|
||||
admitting, reordering, or forwarding requests independently. Distributed
|
||||
layerwise offload with AllGather is an exception at dispatch time: the
|
||||
scheduler may schedule a multi-request wave, but each DP rank executes one
|
||||
envelope from that wave so every rank enters the same collective schedule.
|
||||
admitting, reordering, or forwarding requests independently.
|
||||
|
||||
### DIFF-RUNTIME-INV-003: Terminal cleanup is complete
|
||||
|
||||
**Rule:** Every terminal path MUST release request state, temporary tensors,
|
||||
hooks, and runtime-owned resources. Async-output bookkeeping on abort or
|
||||
consumer drop before consumption is also required *(pending #6253/#6439/#6580)*.
|
||||
hooks, and runtime-owned resources.
|
||||
|
||||
### DIFF-RUNTIME-INV-004: Optional features use runtime hooks
|
||||
|
||||
**Rule:** Cache, profiling, offload, and parallel features SHOULD integrate at
|
||||
defined hooks instead of creating another request lifecycle.
|
||||
|
||||
### DIFF-RUNTIME-INV-005: Results keep request identity
|
||||
|
||||
**Rule:** Batched and asynchronous results MUST be mapped by stable request ID,
|
||||
not by assumed completion order.
|
||||
|
||||
## Safe-change guide
|
||||
|
||||
Test the smallest affected slice, then cover its neighboring boundary:
|
||||
|
||||
| Change area | Minimum evidence |
|
||||
| --- | --- |
|
||||
| Admission or state transitions | `test_diffusion_scheduler.py`, including duplicate IDs, compatibility, abort, and missing output |
|
||||
| Engine loop or output delivery | `test_diffusion_engine.py`, `test_diffusion_engine_cleanup.py`, `test_diffusion_engine_rpc_routing.py` |
|
||||
| Executor or IPC | `test_multiproc_engine_concurrency.py`, `test_uniproc_executor.py`, `test_result_pump.py`, `test_diffusion_ipc.py`, `test_async_output_worker.py` |
|
||||
| Worker or runner | `test_diffusion_worker.py`, `test_diffusion_model_runner.py` |
|
||||
| Stage boundary | `test_stage_diffusion_proc.py`, `test_inline_stage_diffusion_client.py` |
|
||||
| Warmup or streaming | `test_diffusion_engine_dummy_run.py`, `test_diffusion_streaming_output.py` |
|
||||
|
||||
Always exercise success, cancellation, per-request failure, fatal worker
|
||||
failure, repeated shutdown, one request, and multiple compatible requests.
|
||||
For step-mode changes, also test partial progress and cleanup of runner state.
|
||||
|
||||
## Related documents
|
||||
|
||||
- [Diffusion module overview](index.md)
|
||||
- [Diffusion model integration](diffusion_model_integration.md)
|
||||
- [Diffusion continuous batching](continuous_batching.md)
|
||||
- [Diffusion execution modes](../../../user_guide/diffusion/execution_modes.md)
|
||||
- [Async diffusion output](../../feature/async_diffusion_output.md)
|
||||
- [Engine orchestration](../engine_orchestration.md)
|
||||
- [Stage runtime](../stage_runtime.md)
|
||||
Test success, cancellation, failure, shutdown, multiple requests, and resource
|
||||
cleanup.
|
||||
@@ -124,7 +124,6 @@ def extract_omni_serve_subparser_init():
|
||||
"make_arg_parser": lambda parser: parser, # no-op for doc
|
||||
"_ensure_vllm_platform": lambda: None, # no-op for doc
|
||||
"_nonneg_finite_float": float,
|
||||
"_parse_stage_overrides": __import__("json").loads,
|
||||
"VLLM_SUBCMD_PARSER_EPILOG": "",
|
||||
"logger": logger,
|
||||
"DummySubparsers": DummySubparsers,
|
||||
|
||||
@@ -64,7 +64,6 @@ th {
|
||||
| `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `LongCatVideoAvatarPipeline` | LongCat-Video-Avatar-1.5 A2V/AI2V (native single-speaker and multi-speaker AI2V/AVC) | `meituan-longcat/LongCat-Video-Avatar-1.5` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/meituan-longcat/LongCat-Video-Avatar-1.5.md) |
|
||||
| `BooguImagePipeline` | Boogu-Image | `Boogu/Boogu-Image-0.1-Base`, `Boogu/Boogu-Image-0.1-Edit`, `Boogu/Boogu-Image-0.1-Edit-Turbo` | ✅︎ | | | ✅︎ | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Boogu/Boogu-Image.md) |
|
||||
| `BooguImageTurboPipeline` | Boogu-Image Turbo distilled T2I | `Boogu/Boogu-Image-0.1-Turbo` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/Boogu/Boogu-Image.md) |
|
||||
| `StableDiffusionXLPipeline` | Stable-Diffusion-XL | `stabilityai/stable-diffusion-xl-base-1.0` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | ✅︎ | ✅︎ | | ✅︎ | — |
|
||||
| `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ | — |
|
||||
@@ -97,7 +96,6 @@ th {
|
||||
| `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` | ✅︎ | ✅︎ | | ✅︎ | — |
|
||||
| `Audio8TTSSlowARForConditionalGeneration` | Audio8 TTS Preview | `Audio8/Audio8-TTS-Preview-0.6b` | ✅︎ | | | | — |
|
||||
| `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only), SenseNova-U1.5 | `SenseNova/SenseNova-U1-8B-MoT`, `sensenova/SenseNova-U1.5-8B-MoT` | ✅︎ | | | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/SenseNova/SenseNova-U1.5.md) |
|
||||
| `LancePipeline` | Lance | `bytedance-research/Lance` | ✅︎ | | | | — |
|
||||
| `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | | — |
|
||||
|
||||
@@ -64,7 +64,7 @@ connect
|
||||
<- msgpack(policy_server_config)
|
||||
|
||||
infer
|
||||
-> msgpack({"endpoint": "infer", "session_id": "...", "seed": <optional int>, ...observation})
|
||||
-> msgpack({"endpoint": "infer", "session_id": "...", ...observation})
|
||||
<- msgpack(ndarray | dict[str, ndarray])
|
||||
|
||||
reset
|
||||
@@ -92,12 +92,6 @@ frames are not observation messages.
|
||||
- The policy pipeline owns observation transforms and persistent model state,
|
||||
normally keyed by `session_id`; the API layer forwards the raw observation
|
||||
dictionary.
|
||||
- An optional integer `seed` in an inference message becomes the engine
|
||||
request's `sampling_params.seed`. A policy whose sampling consumes the
|
||||
request generator (GR00T-N1.7) returns the same action chunk for the same
|
||||
observation and seed. If omitted, the engine assigns a random per-request
|
||||
seed, so unseeded inferences are independent of each other. The API layer
|
||||
consumes `seed`; it is not forwarded inside the observation.
|
||||
|
||||
## Limits and Errors
|
||||
|
||||
|
||||
@@ -1,17 +1,17 @@
|
||||
# Realtime Duplex API
|
||||
|
||||
vLLM-Omni serves full-duplex speech models over a WebSocket endpoint,
|
||||
`/v1/realtime`, that speaks the OpenAI Realtime vocabulary plus a
|
||||
`/v1/realtime?duplex=1`, that speaks the OpenAI Realtime vocabulary plus a
|
||||
set of duplex extensions: the model decides when to listen and when to speak,
|
||||
the user can talk over the assistant, playback progress is reported back so
|
||||
history stays honest, and a dropped connection can resume the same session.
|
||||
This page covers how to run a duplex deployment, how to drive it from Python
|
||||
with `vllm_omni.clients.duplex.DuplexClient`, and the complete wire contract.
|
||||
|
||||
The handler is enabled by either a model-native `session_mode: duplex`
|
||||
deployment or an explicit `duplex_session` configuration for turn-based
|
||||
Server VAD. The runtime architecture is described in
|
||||
[Full-Duplex Runtime (MiniCPM-o 4.5)](../design/fullduplex.md).
|
||||
The endpoint is served for models that ship a duplex plugin — currently
|
||||
MiniCPM-o 4.5, PersonaPlex, and Nemotron VoiceChat — and only for deploy
|
||||
configurations that declare `session_mode: duplex`. The runtime architecture
|
||||
is described in [Full-Duplex Runtime (MiniCPM-o 4.5)](../design/fullduplex.md).
|
||||
|
||||
## Quick Start
|
||||
|
||||
@@ -27,7 +27,7 @@ vllm-omni serve openbmb/MiniCPM-o-4_5 \
|
||||
|
||||
`vllm_omni/deploy/minicpmo_4_5.yaml` declares `session_mode: duplex` and
|
||||
`duplex_session.max_sessions: 4`, which mounts the WebSocket routes
|
||||
`ws://<host>:8099/v1/realtime` (this page) and `ws://<host>:8099/v1/duplex`
|
||||
`ws://<host>:8099/v1/realtime?duplex=1` (this page) and `ws://<host>:8099/v1/duplex`
|
||||
(the native dialect for raw-protocol clients) next to the usual HTTP API.
|
||||
PersonaPlex and Nemotron VoiceChat use `vllm_omni/deploy/personaplex.yaml` and
|
||||
`vllm_omni/deploy/nemotron_labs_voicechat_duplex.yaml` in the same way.
|
||||
@@ -240,7 +240,7 @@ on the same helpers.
|
||||
|
||||
## API Reference
|
||||
|
||||
This section is the complete wire contract of `/v1/realtime`: how
|
||||
This section is the complete wire contract of `/v1/realtime?duplex=1`: how
|
||||
the dialect relates to the OpenAI Realtime protocol, which surfaces each
|
||||
model's capability flags gate, a catalogue of every client-to-server and
|
||||
server-to-client event, and a JSON example for each. `DuplexClient` speaks
|
||||
@@ -253,11 +253,10 @@ its protocol tests.
|
||||
|
||||
### Endpoint and transport
|
||||
|
||||
The endpoint is `ws(s)://<host>/v1/realtime`. Bare connections select the
|
||||
session handler when the deployment declares either `session_mode: duplex`
|
||||
or `duplex_session`. `?duplex=1` remains a compatibility alias, while
|
||||
`?duplex=0` selects the legacy Realtime handler. Optional query parameters
|
||||
are `model`, `session_id`,
|
||||
The endpoint is `ws(s)://<host>/v1/realtime?duplex=1`; the `duplex` query
|
||||
parameter (`1`, `true`, or `on`) selects the duplex session handler, and the
|
||||
route is mounted only for deployments whose deploy configuration declares
|
||||
`session_mode: duplex`. Optional query parameters are `model`, `session_id`,
|
||||
`autostart` (`0` means resume-only), `resume`, and `native_duplex`
|
||||
(`minicpmo45_native_duplex` is accepted as a deprecated alias and folded
|
||||
into the canonical name). Every message is one JSON object per WebSocket
|
||||
@@ -343,14 +342,9 @@ Semantic divergences hidden behind shared names:
|
||||
before any commit; a commit may end in `response.listen` and no
|
||||
response, and the model may open a response with no commit at all.
|
||||
OpenAI: commit ⇒ item, `response.create` ⇒ exactly one response.
|
||||
- With `turn_detection={"type":"server_vad"}`, omitted `interrupt_response`
|
||||
defaults to `false` on turn-based models such as Qwen3-Omni, enabling
|
||||
endpointing and automatic responses without interruption. Explicit `true`
|
||||
is rejected there. Model-native duplex keeps the `true` default and rejects
|
||||
`false`; its VAD implies `overlap_policy=barge_in_on_speech`.
|
||||
This capability-dependent default differs from OpenAI's `true` default;
|
||||
the effective configuration is returned in `session.created` / `session.updated`.
|
||||
`semantic_vad` is unsupported.
|
||||
- `turn_detection.interrupt_response=false` is rejected;
|
||||
`create_response` is ignored; `semantic_vad` is unsupported; the VAD
|
||||
runs per session (Silero) and implies `overlap_policy=barge_in_on_speech`.
|
||||
- `rate_limits.updated` is always an empty list (compatibility only).
|
||||
- Cancellation never reuses a `response_id`; truncation is driven by
|
||||
`playback.ack` as well as `conversation.item.truncate`.
|
||||
@@ -1207,10 +1201,9 @@ Realtime projector renames and fans out these events before they reach a
|
||||
barge-in, or audio truncation; Nemotron VoiceChat does not support barge-in
|
||||
or audio truncation; camera frames are consumed only by MiniCPM-o 4.5; tool
|
||||
calls are produced only by Nemotron VoiceChat.
|
||||
- `turn_detection` supports `server_vad` or `null`, not `semantic_vad`.
|
||||
Turn-based models support endpointing without interruption, including
|
||||
`create_response=false` for explicit response creation. Model-native duplex
|
||||
requires `interrupt_response=true` and `create_response=true`.
|
||||
- `turn_detection` supports only `server_vad` with `interrupt_response=true`;
|
||||
`semantic_vad`, `interrupt_response=false`, and `create_response` are not
|
||||
supported. Server VAD requires a model-native session.
|
||||
- `input_audio_transcription` and `input_audio_noise_reduction` are accepted
|
||||
and echoed but no separate transcription or noise-reduction stage runs;
|
||||
transcripts come from the model.
|
||||
|
||||
+12
-31
@@ -279,15 +279,10 @@ curl -X POST http://localhost:8091/v1/audio/voices \
|
||||
|
||||
## Streaming Text Input (WebSocket)
|
||||
|
||||
The `/v1/audio/speech/stream` WebSocket endpoint accepts text incrementally.
|
||||
By default (`split_granularity=none`) it buffers until `input.done` and
|
||||
synthesizes that flush as **one** TTS request, which keeps long-form timbre
|
||||
stable. Set `split_granularity` to `sentence` or `clause` to emit a request
|
||||
at each detected boundary (lower time-to-first-audio for STT/LLM pipelines).
|
||||
The `/v1/audio/speech/stream` WebSocket endpoint accepts text incrementally and generates audio per sentence as boundaries are detected.
|
||||
|
||||
> Note: `stream_audio` only changes how **audio bytes** are framed (one WAV/PCM
|
||||
> payload vs chunked PCM). Text segmentation is controlled separately by
|
||||
> `split_granularity`.
|
||||
> Note: text input is always streamed incrementally. Audio output remains sentence-scoped:
|
||||
> use `stream_audio=false` for one binary frame per sentence, or `stream_audio=true` for one or more PCM chunks per sentence.
|
||||
|
||||
### WebSocket Protocol
|
||||
|
||||
@@ -319,16 +314,14 @@ upstream LLM) pays the WebSocket handshake once instead of once per utterance.
|
||||
|
||||
- The session config is sticky. Send `input.text` again straight after
|
||||
`session.done` to reuse it, or send another `session.config` first to change
|
||||
voice, format, or reference audio. A `session.config` sent in the middle of
|
||||
an utterance is rejected so no pending input is silently dropped and so a
|
||||
split utterance cannot end up half in one voice and half in another.
|
||||
- An utterance is the flush unit, not a linguistic one: it is one `input.done`
|
||||
cycle. `utterance_index` counts those flushes across the connection, so it
|
||||
tells you which `input.done` a frame belongs to. `sentence_index` counts the
|
||||
TTS requests inside one flush and so pairs with `total_sentences`: with the
|
||||
default `split_granularity=none` that is always `sentence_index: 0` of
|
||||
`total_sentences: 1` (or `0` for an empty buffer), while `sentence` or
|
||||
`clause` counts the linguistic units actually synthesized.
|
||||
voice, format, or reference audio. A `session.config` sent while text is
|
||||
still buffered is rejected so no pending input is silently dropped.
|
||||
- An utterance is the flush unit, not a linguistic one: it is whatever text was
|
||||
buffered when `input.done` arrived, of any length, synthesized as one request.
|
||||
`utterance_index` counts those flushes across the connection, so it tells you
|
||||
which `input.done` a frame belongs to. `sentence_index` counts within one
|
||||
flush and so pairs with `total_sentences`, which means every utterance reports
|
||||
`sentence_index: 0` of `total_sentences: 1` (or `0` for an empty buffer).
|
||||
- End the connection with `session.close`, or by closing the socket. An idle
|
||||
connection is still closed after the server's idle timeout, which now also
|
||||
applies to the gap between utterances.
|
||||
@@ -339,15 +332,7 @@ All REST API parameters are supported, plus:
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
| ----------- | ------ | --------- | ------------- |
|
||||
| `stream_audio` | bool | false | Stream one or more PCM chunks for each TTS request over WebSocket |
|
||||
| `split_granularity` | string | `"none"` | `"none"`: one request per `input.done`. `"sentence"`: split on `.!?` plus CJK `。!?…`, Indic danda `।॥`, and Arabic `؟`. `"clause"`: also split on `,;,;،؛`. |
|
||||
| `seed` | integer | null | Forwarded to the speech engine for this session |
|
||||
|
||||
ASCII punctuation only ends a unit when whitespace or `input.done` follows it,
|
||||
and decimals (`3.14`), thousands separators (`1,000`), abbreviations (`Dr.`,
|
||||
`e.g.`) and initials (`J. R.`) are not treated as boundaries. A punctuation run
|
||||
and any closing quote or bracket stay with the unit they close, so `Wait...`
|
||||
and `He said "Hello."` are one request each.
|
||||
| `stream_audio` | bool | false | Stream one or more PCM chunks for the buffered input over WebSocket |
|
||||
|
||||
```bash
|
||||
DELETE /v1/audio/voices/{name}
|
||||
@@ -767,10 +752,6 @@ Fish Speech uses `ref_audio` and `ref_text` for voice cloning (no `task_type` ne
|
||||
| ------- | ------------- |
|
||||
| `k2-fsa/OmniVoice` | Pure-diffusion TTS. Supports voice cloning via `ref_audio` (with optional `ref_text`); no built-in voice presets. |
|
||||
|
||||
OmniVoice uses packed variable-length attention for batched generator execution. The attention operator accepts FP16 and BF16 inputs, so its query,
|
||||
key, and value tensors are evaluated in BF16 even when the stage is configured with `dtype: float32`; the attention output is converted back to the model's
|
||||
hidden-state dtype before the output projection. Consequently, a float32 stage configuration does not imply FP32 attention arithmetic.
|
||||
|
||||
### VoxCPM2
|
||||
|
||||
| Model | Description |
|
||||
|
||||
@@ -68,8 +68,6 @@ curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o output.mp4
|
||||
| Parameter | Type | Default | Description |
|
||||
| ----------- | ------ | --------- | ------------- |
|
||||
| `input_reference` | file | null | Uploaded reference image or video for image-to-video/video-to-video requests |
|
||||
| `control_reference` | file | null | Optional uploaded image/video control, up to 512 MiB, for models that declare control-upload support |
|
||||
| `control_type` | string | null | Model control name associated with `control_reference`; currently Cosmos3 supports `edge`, `blur`, `depth`, `seg`, and `wsm` |
|
||||
| `image_reference` | string | null | JSON-encoded reference image payload; do not combine with `input_reference` or `video_reference` |
|
||||
| `video_reference` | string | null | JSON-encoded reference video payload; do not combine with `input_reference` or `image_reference` |
|
||||
| `audio_reference` | string | null | JSON-encoded audio reference for speech-to-video: `{"audio_url": "..."}` — supports HTTP(s) URLs or base64 data URLs |
|
||||
@@ -179,23 +177,6 @@ corresponding addition. `negative_metadata_mode` accepts `same`, `inverse`, or
|
||||
`none` and defaults to `same` for transfer. See the Cosmos3 recipe for complete
|
||||
examples.
|
||||
|
||||
A client that cannot place the control on the server filesystem can upload one
|
||||
control with `control_reference` and identify it with `control_type`. The API
|
||||
streams the upload to request-scoped storage, supplies its path to the model,
|
||||
and removes it after synchronous or asynchronous generation completes. Other
|
||||
options for the selected control can remain in `extra_params`; do not also set
|
||||
`control` or `control_path` there. Uploads larger than 512 MiB are rejected.
|
||||
|
||||
```bash
|
||||
curl -s http://localhost:8091/v1/videos/sync \
|
||||
-F "prompt=Preserve the scene while following the world-state control" \
|
||||
-F "input_reference=@input.mp4;type=video/mp4" \
|
||||
-F "control_reference=@wsm.mp4;type=video/mp4" \
|
||||
-F "control_type=wsm" \
|
||||
-F 'extra_params={"wsm":{"control_weight":1.0}}' \
|
||||
-o output.mp4
|
||||
```
|
||||
|
||||
HTTP redirects for `image_reference.image_url` follow vLLM's
|
||||
`VLLM_MEDIA_URL_ALLOW_REDIRECTS` setting. Before starting the server, set it to
|
||||
`1` (the default) to allow redirects or `0` to reject them. A redirect target
|
||||
|
||||
@@ -8,7 +8,7 @@ are exposed.
|
||||
## Choose a Mode
|
||||
|
||||
| Goal | CLI configuration |
|
||||
| --- | --- |
|
||||
|---|---|
|
||||
| Serial request execution | `--max-num-seqs 1` |
|
||||
| Fused request-level batching | `--max-num-seqs N` |
|
||||
| Single-request step execution | `--step-execution --max-num-seqs 1` |
|
||||
@@ -177,10 +177,8 @@ wait
|
||||
The scheduler may batch these requests when their sampling parameters are
|
||||
compatible. See the
|
||||
[Image Generation API](../../serving/image_generation_api.md) for response
|
||||
formats and additional client examples.
|
||||
|
||||
Streaming models can use model-specific streaming endpoints documented at
|
||||
[Streaming Video Output API](../../serving/streaming_video_output_api.md).
|
||||
formats and additional client examples. Streaming models can use
|
||||
model-specific streaming endpoints documented by their serving guide.
|
||||
|
||||
## Python API
|
||||
|
||||
@@ -243,7 +241,7 @@ For step execution, set `step_execution: true` and remove
|
||||
## CLI Reference
|
||||
|
||||
| Flag | Default | Effect |
|
||||
| --- | ---: | --- |
|
||||
|---|---:|---|
|
||||
| `--step-execution` | disabled | Select step-wise scheduling |
|
||||
| `--max-num-seqs` | `1` for diffusion stages | Set request- or step-scheduler capacity |
|
||||
| `--request-batch-max-wait-ms` | `0` | Wait for burst coalescing in request mode |
|
||||
|
||||
@@ -249,14 +249,8 @@ MiniMax-H3 supports two few-step mechanisms that must not be conflated:
|
||||
- **Checkpoint-pinned schedule**: a merged release writes `base_schedule` into
|
||||
`model_index.json`. Requests must pass `num_inference_steps` as the interval
|
||||
count (for example `4` for `[1.0, 0.7, 0.4, 0.15, 0.0]`).
|
||||
- **Runtime Turbo LoRA**: each LightX2V Turbo artifact carries its own contract,
|
||||
read from its filename. A four-step artifact requests five sigma points and an
|
||||
eight-step one requests nine; the 768p retrains enforce `flow_shift=6` and the
|
||||
544p artifacts `flow_shift=12`. `audio_flow_shift=3` across the family. A
|
||||
request that does not match the loaded artifact is rejected by name. Alpha
|
||||
comes from the artifact's metadata, or 8 when it declares none. `ref2v`
|
||||
artifacts require `--task-type ref2va`; `fl2v` artifacts serve T2VA and FL2VA.
|
||||
The ComfyUI exports use a fused-QKV layout and are not supported.
|
||||
- **Runtime Turbo LoRA**: LightX2V Turbo artifacts request five sigma points and
|
||||
enforce `flow_shift=6`, `audio_flow_shift=3`.
|
||||
- **Runtime native LoRA**: FlashGen-style artifacts declare
|
||||
`key_format=minimax-h3-native` and embed `base_schedule` in safetensors
|
||||
metadata. When active, the adapter schedule overrides the base checkpoint and
|
||||
|
||||
@@ -162,9 +162,9 @@ The following tables show which models support each feature:
|
||||
| **Wan2.2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (encode/decode) | ❌ | ❌ |
|
||||
| **Wan2.2-S2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ❌ | ❌ |
|
||||
| **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **LTX-2** | ❌ | ✅ | ✅ (Ulysses only) | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **LTX-2.3** | ❌ | ✅ | ✅ (Ulysses only) | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **LTX-2.5** | ❌ | ❓ (one-stage) | ✅* (Ulysses only) | ✅ (Full only) | ❓ | ❌ | ✅ | ✅ | ✅ (decode) | ❓ (FP8) | ❌ |
|
||||
| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **LTX-2.3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **LTX-2.5** | ❌ | ❓ (one-stage) | ❓ (Ulysses only) | ✅ (Full only) | ❓ | ❌ | ✅ | ✅ | ✅ (decode) | ❓ (FP8) | ❌ |
|
||||
| **SANA-Video-2B T2V I2V** | ❌ | ❌ | ✅ (`--usp`, frame-sharded) | ✅ | ✅ (TP=2 only) | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅* |
|
||||
| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ |
|
||||
|
||||
@@ -126,7 +126,11 @@ python examples/offline_inference/text_to_speech/cosyvoice3/end2end.py \
|
||||
## Fish Speech S2 Pro
|
||||
|
||||
4B dual-AR text-to-speech model from FishAudio with the DAC codec at 44.1 kHz.
|
||||
No extra packages are required; the DAC codec is vendored in vLLM-Omni.
|
||||
|
||||
### Prerequisites
|
||||
```bash
|
||||
pip install fish-speech
|
||||
```
|
||||
|
||||
### Quick start
|
||||
```bash
|
||||
|
||||
@@ -155,8 +155,13 @@ The client supports `--api-base`, `--model`, `--text`, `--ref-audio`, `--ref-tex
|
||||
|
||||
## Fish Speech S2 Pro
|
||||
|
||||
4B dual-AR TTS at 44.1 kHz. Server uses the DAC codec, which is vendored in
|
||||
vLLM-Omni (no extra packages required).
|
||||
4B dual-AR TTS at 44.1 kHz. Server uses the DAC codec.
|
||||
|
||||
### Prerequisites
|
||||
|
||||
```bash
|
||||
pip install fish-speech
|
||||
```
|
||||
|
||||
### Launch
|
||||
|
||||
|
||||
@@ -30,6 +30,7 @@ from typing import Any
|
||||
import numpy as np
|
||||
|
||||
_MODEL = "robbyant/lingbot-world-v2-14b-causal-fast-diffusers"
|
||||
_NUM_INFERENCE_STEPS = 4
|
||||
_MAX_SEQUENCE_LENGTH = 512
|
||||
_MAX_PIXEL_AREA = 480 * 832
|
||||
_MAX_RAW_FRAMES = 117
|
||||
@@ -208,6 +209,7 @@ def build_request(
|
||||
"height": args.height,
|
||||
"width": args.width,
|
||||
"num_frames": args.num_frames,
|
||||
"num_inference_steps": _NUM_INFERENCE_STEPS,
|
||||
"max_sequence_length": _MAX_SEQUENCE_LENGTH,
|
||||
"seed": args.seed,
|
||||
"fps": args.fps,
|
||||
|
||||
@@ -155,6 +155,7 @@ async def run(argv: Sequence[str] | None = None) -> Path:
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
num_frames=9,
|
||||
num_inference_steps=4,
|
||||
max_sequence_length=512,
|
||||
seed=args.seed,
|
||||
output_type="latent",
|
||||
|
||||
@@ -5,12 +5,10 @@ This example shows how to use vLLM-Omni for running offline inference
|
||||
with the correct prompt format on Qwen2.5-Omni
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import functools
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
from typing import Any, NamedTuple
|
||||
from typing import NamedTuple
|
||||
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
@@ -28,20 +26,6 @@ from vllm_omni.utils.tracking_parser import TrackingArgumentParser
|
||||
SEED = 42
|
||||
|
||||
|
||||
def parse_json_object(value: str, flag_name: str = "argument") -> dict[str, Any]:
|
||||
"""Parse a CLI value as a JSON object, attributing errors to ``flag_name``."""
|
||||
try:
|
||||
config = json.loads(value)
|
||||
except json.JSONDecodeError as e:
|
||||
raise argparse.ArgumentTypeError(f"{flag_name} must be valid JSON: {e}") from e
|
||||
if not isinstance(config, dict):
|
||||
raise argparse.ArgumentTypeError(f"{flag_name} must be a JSON object")
|
||||
return config
|
||||
|
||||
|
||||
parse_profiler_config = functools.partial(parse_json_object, flag_name="--profiler-config")
|
||||
|
||||
|
||||
class QueryResult(NamedTuple):
|
||||
inputs: dict
|
||||
limit_mm_per_prompt: dict[str, int]
|
||||
@@ -394,10 +378,11 @@ def main(args):
|
||||
for i, prompt in enumerate(prompts):
|
||||
prompt["modalities"] = output_modalities
|
||||
|
||||
profiler_enabled = args.profiler_config is not None
|
||||
if profiler_enabled:
|
||||
print("[Profiler] Starting profiling...")
|
||||
profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
|
||||
if profiler_enabled and hasattr(omni, "start_profile"):
|
||||
omni.start_profile(stages=[0])
|
||||
elif profiler_enabled:
|
||||
print("[Warn] VLLM_TORCH_PROFILER_DIR is set, but current engine does not support profiler controls.")
|
||||
omni_generator = omni.generate(prompts, sampling_params_list, py_generator=args.py_generator)
|
||||
|
||||
# Determine output directory: prefer --output-dir; fallback to --output-wav
|
||||
@@ -433,7 +418,7 @@ def main(args):
|
||||
print(f"Request ID: {request_id}, Saved audio to {output_wav}")
|
||||
|
||||
processed_count += 1
|
||||
if profiler_enabled and processed_count >= total_requests:
|
||||
if profiler_enabled and hasattr(omni, "stop_profile") and processed_count >= total_requests:
|
||||
print(f"[Info] Processed {processed_count}/{total_requests}. Stopping profiler inside active loop...")
|
||||
# Stop the profiler while workers are still alive
|
||||
omni.stop_profile()
|
||||
@@ -491,12 +476,6 @@ def parse_args():
|
||||
default=300,
|
||||
help="Timeout for initializing stages in seconds (default: 300)",
|
||||
)
|
||||
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(
|
||||
"--shm-threshold-bytes",
|
||||
type=int,
|
||||
|
||||
@@ -14,7 +14,6 @@ list of supported architectures across all modalities, see
|
||||
|
||||
| Model | HuggingFace repo | Stages | Voice cloning | Streaming | Special modes | Sample rate |
|
||||
|---|---|---|---|---|---|---|
|
||||
| Audio8 TTS Preview | `Audio8/Audio8-TTS-Preview-0.6b` | dual-AR | ✓ | ✓ | 11 languages | 44.1 kHz |
|
||||
| CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ | — | 24 kHz |
|
||||
| Fish Speech S2 Pro | `fishaudio/s2-pro` | dual-AR | ✓ | ✓ | — | 44.1 kHz |
|
||||
| Gepard-1.0 | `nineninesix/gepard-1.0` | single (native AR + NanoCodec) | — (zero-shot; cloning WIP) | — (serving WIP) | zero-shot | 22.05 kHz |
|
||||
@@ -45,57 +44,6 @@ python examples/offline_inference/text_to_speech/<model>/end2end.py \
|
||||
|
||||
---
|
||||
|
||||
## Audio8 TTS Preview
|
||||
|
||||
0.6B DualAR text-to-speech model (Audio8) with its own 44.1 kHz neural audio
|
||||
codec: a slow AR transformer predicts one semantic token per codec frame, and a
|
||||
4-layer fast AR predicts that frame's remaining 9 codebooks.
|
||||
|
||||
### Prerequisites
|
||||
No extra packages: the codec architecture ships in tree and its weights
|
||||
(`codec.pth`) come with the checkpoint.
|
||||
|
||||
### Quick start
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/audio8_tts/end2end.py \
|
||||
--text "Welcome to Audio8 TTS, a compact multilingual text to speech model."
|
||||
```
|
||||
|
||||
### Voice cloning
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/audio8_tts/end2end.py \
|
||||
--text "Welcome to Audio8 TTS." \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--ref-text "The exact transcript of the reference recording."
|
||||
```
|
||||
The reference transcript must match what is actually spoken in the clip;
|
||||
a mismatch degrades both stability and speaker similarity.
|
||||
|
||||
### Streaming
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/audio8_tts/end2end.py \
|
||||
--text "Welcome to Audio8 TTS." --streaming
|
||||
```
|
||||
|
||||
### Notes
|
||||
- Output: 44.1 kHz mono WAV; the codec emits ~21.5 frames/s (2048 samples per frame).
|
||||
- Context is 2048 packed text+audio positions, so `max_model_len` is 2048.
|
||||
- Greedy decoding (`temperature: 0`) is degenerate for this checkpoint — it
|
||||
ends the utterance almost immediately. Keep the deploy defaults
|
||||
(`temperature 0.7`, `top_k 50`, `top_p 0.9`), which also match
|
||||
`generation_config.json`.
|
||||
- Repetition-Aware Sampling is applied to the semantic token: a repeat inside
|
||||
the last 10 frames is re-drawn from a flatter distribution instead of being
|
||||
masked out.
|
||||
- `--num-prompts 4` smoke-tests concurrency (stage 0 runs `max_num_seqs: 4`).
|
||||
|
||||
### Measured behaviour (1x H20, shared GPU)
|
||||
`--streaming` on a 58-character English sentence: TTFA ~0.74 s (4-frame first
|
||||
chunk), 7 chunks, 1.55 s wall for 3.16 s of audio (RTF ~0.49). Treat these as
|
||||
smoke-test magnitudes, not a benchmark.
|
||||
|
||||
---
|
||||
|
||||
## CosyVoice3
|
||||
|
||||
2-stage TTS pipeline (`talker` + `code2wav`) at 24 kHz.
|
||||
@@ -182,7 +130,11 @@ Text → [Stage 0: AR] → Speech Tokens → [Stage 1: DiT + HiFT] → Audio (24
|
||||
## Fish Speech S2 Pro
|
||||
|
||||
4B dual-AR text-to-speech model from FishAudio with the DAC codec at 44.1 kHz.
|
||||
No extra packages are required; the DAC codec is vendored in vLLM-Omni.
|
||||
|
||||
### Prerequisites
|
||||
```bash
|
||||
pip install fish-speech
|
||||
```
|
||||
|
||||
### Quick start
|
||||
```bash
|
||||
|
||||
@@ -1,237 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""Offline inference demo for Audio8 TTS Preview 0.6B via vLLM Omni.
|
||||
|
||||
Usage:
|
||||
# Text-only synthesis
|
||||
python end2end.py --text "Welcome to Audio8 TTS."
|
||||
|
||||
# Zero-shot voice cloning (the transcript must match the reference audio)
|
||||
python end2end.py --text "Welcome to Audio8 TTS." \
|
||||
--ref-audio reference.wav --ref-text "The exact transcript of the reference recording."
|
||||
|
||||
# Streaming (per-chunk TTFA / inter-chunk timings)
|
||||
python end2end.py --text "Welcome to Audio8 TTS." --streaming
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import time
|
||||
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
import torch
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
from vllm_omni import AsyncOmni, Omni
|
||||
from vllm_omni.model_executor.models.audio8_tts.codec_utils import estimate_reference_code_frames
|
||||
from vllm_omni.model_executor.models.audio8_tts.prompt_utils import (
|
||||
build_text_only_prompt_ids,
|
||||
estimate_voice_clone_prompt_len,
|
||||
normalize_text,
|
||||
)
|
||||
from vllm_omni.utils.tracking_parser import TrackingArgumentParser
|
||||
|
||||
DEFAULT_MODEL = "Audio8/Audio8-TTS-Preview-0.6b"
|
||||
|
||||
|
||||
def build_prompt(
|
||||
text: str,
|
||||
ref_audio_path: str | None = None,
|
||||
ref_text: str | None = None,
|
||||
model_name: str = DEFAULT_MODEL,
|
||||
) -> dict:
|
||||
"""Build an engine prompt, using the same protocol as online serving."""
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_name)
|
||||
if ref_audio_path is None and ref_text is None:
|
||||
prompt_ids, normalized_text = build_text_only_prompt_ids(tokenizer, text)
|
||||
return {
|
||||
"prompt_token_ids": prompt_ids,
|
||||
"additional_information": {"text": [normalized_text]},
|
||||
}
|
||||
|
||||
if not ref_audio_path or not ref_text:
|
||||
raise ValueError("Audio8 TTS voice cloning requires both --ref-audio and --ref-text")
|
||||
|
||||
normalized_text = normalize_text(text)
|
||||
normalized_ref_text = normalize_text(ref_text, add_default_speaker=True)
|
||||
ref_audio_wav, ref_audio_sr = sf.read(ref_audio_path, dtype="float32", always_2d=False)
|
||||
if ref_audio_wav.ndim > 1:
|
||||
ref_audio_wav = ref_audio_wav.mean(axis=-1)
|
||||
ref_frames = estimate_reference_code_frames(int(ref_audio_wav.shape[0]), int(ref_audio_sr))
|
||||
placeholder_len = estimate_voice_clone_prompt_len(tokenizer, normalized_text, normalized_ref_text, ref_frames)
|
||||
|
||||
# The clone prompt embeds the reference audio's own codec codes, so the
|
||||
# model builds it in preprocess(); we only reserve the exact length here.
|
||||
return {
|
||||
"prompt_token_ids": [1] * placeholder_len,
|
||||
"additional_information": {
|
||||
"text": normalized_text,
|
||||
"ref_text": normalized_ref_text,
|
||||
"ref_audio_wav": torch.from_numpy(np.asarray(ref_audio_wav, dtype=np.float32)),
|
||||
"ref_audio_sr": int(ref_audio_sr),
|
||||
"audio8_structured_voice_clone": True,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def extract_audio(multimodal_output: dict) -> tuple[torch.Tensor, int]:
|
||||
"""Pull the waveform and sample rate out of a multimodal output.
|
||||
|
||||
The output processor concatenates the per-step delta tensors under
|
||||
``model_outputs``; ``audio`` is the older alias. Never use ``a or b`` on
|
||||
these values -- truthiness on a tensor raises.
|
||||
"""
|
||||
audio = multimodal_output.get("model_outputs")
|
||||
if audio is None:
|
||||
audio = multimodal_output.get("audio")
|
||||
if audio is None:
|
||||
raise ValueError(f"No audio in multimodal_output: {sorted(multimodal_output)}")
|
||||
if isinstance(audio, list):
|
||||
audio = torch.cat([torch.as_tensor(chunk).reshape(-1) for chunk in audio], dim=0)
|
||||
audio = torch.as_tensor(audio).reshape(-1)
|
||||
|
||||
sr_raw = multimodal_output.get("sr")
|
||||
if isinstance(sr_raw, list):
|
||||
sr_raw = sr_raw[-1] if sr_raw else None
|
||||
if sr_raw is None:
|
||||
raise ValueError("Missing sample rate in multimodal_output")
|
||||
sample_rate = int(sr_raw.item()) if hasattr(sr_raw, "item") else int(sr_raw)
|
||||
return audio, sample_rate
|
||||
|
||||
|
||||
def _save_wav(output_dir: str, request_id: str, audio: torch.Tensor, sample_rate: int) -> None:
|
||||
out_wav = os.path.join(output_dir, f"output_{request_id}.wav")
|
||||
sf.write(out_wav, audio.float().cpu().numpy(), samplerate=sample_rate, format="WAV")
|
||||
duration = audio.numel() / sample_rate
|
||||
# print(), not logging: vLLM reconfigures logging on engine start and
|
||||
# disables loggers created before that, which would swallow these lines.
|
||||
print(f"Request {request_id}: saved {out_wav} (sr={sample_rate}, {duration:.2f}s)")
|
||||
|
||||
|
||||
def _omni_kwargs(args, model_name: str) -> dict:
|
||||
kwargs = {
|
||||
"model": model_name,
|
||||
"log_stats": args.log_stats,
|
||||
"stage_init_timeout": args.stage_init_timeout,
|
||||
}
|
||||
if args.deploy_config:
|
||||
kwargs["deploy_config"] = args.deploy_config
|
||||
return kwargs
|
||||
|
||||
|
||||
def main(args) -> None:
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
model_name = args.model or DEFAULT_MODEL
|
||||
prompts = [
|
||||
build_prompt(text, args.ref_audio, args.ref_text, model_name) for text in ([args.text] * args.num_prompts)
|
||||
]
|
||||
|
||||
omni = Omni(**_omni_kwargs(args, model_name))
|
||||
t_start = time.perf_counter()
|
||||
total_audio_seconds = 0.0
|
||||
for stage_outputs in omni.generate(prompts):
|
||||
request_output = stage_outputs
|
||||
if request_output is None or not request_output.outputs:
|
||||
continue
|
||||
audio, sample_rate = extract_audio(request_output.outputs[0].multimodal_output)
|
||||
assert audio.numel() > 0, "Audio8 TTS produced an empty waveform"
|
||||
total_audio_seconds += audio.numel() / sample_rate
|
||||
_save_wav(args.output_dir, request_output.request_id, audio, sample_rate)
|
||||
elapsed = time.perf_counter() - t_start
|
||||
rtf = elapsed / max(total_audio_seconds, 1e-6)
|
||||
print(f"Total {elapsed * 1000:.1f} ms for {total_audio_seconds:.2f} s of audio (RTF={rtf:.3f})")
|
||||
|
||||
|
||||
async def main_streaming(args) -> None:
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
model_name = args.model or DEFAULT_MODEL
|
||||
prompt = build_prompt(args.text, args.ref_audio, args.ref_text, model_name)
|
||||
|
||||
omni = AsyncOmni(**_omni_kwargs(args, model_name))
|
||||
request_id = "0"
|
||||
chunks: list[torch.Tensor] = []
|
||||
sample_rate = None
|
||||
t_start = time.perf_counter()
|
||||
t_prev = t_start
|
||||
chunk_index = 0
|
||||
|
||||
async for stage_output in omni.generate(prompt, request_id=request_id):
|
||||
multimodal_output = stage_output.outputs[0].multimodal_output
|
||||
if not multimodal_output:
|
||||
continue
|
||||
audio = multimodal_output.get("model_outputs")
|
||||
if audio is None:
|
||||
audio = multimodal_output.get("audio")
|
||||
if audio is not None:
|
||||
if isinstance(audio, list):
|
||||
chunks.extend(torch.as_tensor(chunk).reshape(-1) for chunk in audio)
|
||||
else:
|
||||
chunks.append(torch.as_tensor(audio).reshape(-1))
|
||||
sr_raw = multimodal_output.get("sr")
|
||||
if isinstance(sr_raw, list):
|
||||
sr_raw = sr_raw[-1] if sr_raw else None
|
||||
if sr_raw is not None:
|
||||
sample_rate = int(sr_raw.item()) if hasattr(sr_raw, "item") else int(sr_raw)
|
||||
|
||||
if stage_output.finished:
|
||||
break
|
||||
now = time.perf_counter()
|
||||
if chunk_index == 0:
|
||||
print(f"chunk 0: TTFA={(now - t_start) * 1000:.1f} ms")
|
||||
else:
|
||||
print(f"chunk {chunk_index}: inter_chunk={(now - t_prev) * 1000:.1f} ms")
|
||||
t_prev = now
|
||||
chunk_index += 1
|
||||
|
||||
if not chunks or sample_rate is None:
|
||||
raise RuntimeError("Audio8 TTS streaming produced no audio")
|
||||
_save_wav(args.output_dir, request_id, torch.cat(chunks, dim=0), sample_rate)
|
||||
print(f"Streaming done: total={(time.perf_counter() - t_start) * 1000:.1f} ms chunks={chunk_index}")
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = TrackingArgumentParser(description="Audio8 TTS Preview offline inference with vLLM Omni")
|
||||
parser.add_argument("--model", type=str, default=DEFAULT_MODEL, help="Model path or HuggingFace repo ID.")
|
||||
parser.add_argument(
|
||||
"--text",
|
||||
type=str,
|
||||
default="Welcome to Audio8 TTS, a compact multilingual text to speech model.",
|
||||
help="Text to synthesize.",
|
||||
)
|
||||
parser.add_argument("--ref-audio", type=str, default=None, help="Reference audio path for voice cloning.")
|
||||
parser.add_argument(
|
||||
"--ref-text",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Transcript of --ref-audio; required for voice cloning and must match the recording.",
|
||||
)
|
||||
parser.add_argument("--output-dir", type=str, default="output_audio", help="Directory for output WAV files.")
|
||||
parser.add_argument(
|
||||
"--num-prompts",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Duplicate the prompt N times to smoke-test concurrent requests.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--deploy-config",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Override the deploy config path; defaults to vllm_omni/deploy/audio8_tts.yaml via model_type.",
|
||||
)
|
||||
parser.add_argument("--streaming", action="store_true", help="Use AsyncOmni and report per-chunk latency.")
|
||||
parser.add_argument("--log-stats", action="store_true", help="Enable engine stats logging.")
|
||||
parser.add_argument("--stage-init-timeout", type=int, default=600, help="Per-stage init timeout in seconds.")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parsed_args = parse_args()
|
||||
if parsed_args.streaming:
|
||||
asyncio.run(main_streaming(parsed_args))
|
||||
else:
|
||||
main(parsed_args)
|
||||
@@ -1,12 +1,9 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import argparse
|
||||
import functools
|
||||
import json
|
||||
import os
|
||||
import urllib.request
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
@@ -32,20 +29,6 @@ def _default_ref_audio() -> str:
|
||||
return str(dest)
|
||||
|
||||
|
||||
def parse_json_object(value: str, flag_name: str = "argument") -> dict[str, Any]:
|
||||
"""Parse a CLI value as a JSON object, attributing errors to ``flag_name``."""
|
||||
try:
|
||||
config = json.loads(value)
|
||||
except json.JSONDecodeError as e:
|
||||
raise argparse.ArgumentTypeError(f"{flag_name} must be valid JSON: {e}") from e
|
||||
if not isinstance(config, dict):
|
||||
raise argparse.ArgumentTypeError(f"{flag_name} must be a JSON object")
|
||||
return config
|
||||
|
||||
|
||||
parse_profiler_config = functools.partial(parse_json_object, flag_name="--profiler-config")
|
||||
|
||||
|
||||
def run_e2e():
|
||||
parser = argparse.ArgumentParser()
|
||||
# ""FunAudioLLM/Fun-CosyVoice3-0.5B-2512
|
||||
@@ -81,12 +64,6 @@ def run_e2e():
|
||||
required=True,
|
||||
help="Path to tokenizer directory (e.g., <model_path>/CosyVoice-BlankEN).",
|
||||
)
|
||||
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"}\'.',
|
||||
)
|
||||
args = parser.parse_args()
|
||||
# Ensure tokenizer directory exists
|
||||
if not os.path.exists(args.tokenizer):
|
||||
@@ -102,7 +79,6 @@ def run_e2e():
|
||||
deploy_config=args.deploy_config,
|
||||
tokenizer=args.tokenizer,
|
||||
log_stats=True,
|
||||
profiler_config=args.profiler_config,
|
||||
)
|
||||
|
||||
sampling_cfg = {"top_p": 0.8, "top_k": 25, "eos_token_id": 6561 + 1}
|
||||
@@ -173,8 +149,8 @@ def run_e2e():
|
||||
|
||||
sampling_params_list = [gpt_sampling, s2mel_sampling]
|
||||
|
||||
profiler_enabled = args.profiler_config is not None
|
||||
if profiler_enabled:
|
||||
# Start profiling (requires VLLM_TORCH_PROFILER_DIR env var)
|
||||
if os.environ.get("VLLM_TORCH_PROFILER_DIR"):
|
||||
print("Starting profiler...")
|
||||
omni.start_profile()
|
||||
|
||||
@@ -182,7 +158,7 @@ def run_e2e():
|
||||
outputs = list(omni.generate(prompts, sampling_params_list=sampling_params_list[:2]))
|
||||
|
||||
# Stop profiling and get results
|
||||
if profiler_enabled:
|
||||
if os.environ.get("VLLM_TORCH_PROFILER_DIR"):
|
||||
print("Stopping profiler...")
|
||||
profile_results = omni.stop_profile()
|
||||
print(f"Profile traces saved to: {profile_results}")
|
||||
|
||||
@@ -264,7 +264,7 @@ python examples/online_serving/openai_chat_completion_client_for_multimodal_gene
|
||||
|
||||
#### Realtime WebSocket client (`openai_realtime_client.py`)
|
||||
|
||||
[`openai_realtime_client.py`](./openai_realtime_client.py) connects to **`ws://<host>:<port>/v1/realtime`**, streams a local WAV as **PCM16 mono @ 16 kHz** in fixed-size chunks, and receives **`response.audio.*`**, **`response.output_text.*`**, and transcript events. By default it uses explicit `input_audio_buffer.commit`; pass **`--server-vad`** to let the server detect speech endpoints and commit the turn automatically. The client concatenates audio deltas and writes **`--output-wav`** (model output is typically **24 kHz**). Optional **`--delta-dump-dir`** saves each delta as `delta_000001.wav`, … for debugging.
|
||||
[`openai_realtime_client.py`](./openai_realtime_client.py) connects to **`ws://<host>:<port>/v1/realtime`**, streams a local WAV as **PCM16 mono @ 16 kHz** in fixed-size chunks (OpenAI-style `input_audio_buffer.append` / `commit`), and receives **`response.audio.delta`** (incremental PCM for the reply) plus **`transcription.*`** events. By default it concatenates audio deltas and writes **`--output-wav`** (model output is typically **24 kHz**). Optional **`--delta-dump-dir`** saves each delta as `delta_000001.wav`, … for debugging.
|
||||
|
||||
Streaming input works well for translation-style use cases; if the Thinker runs while input is still incomplete, consider limiting **`max_tokens`** in your session / server defaults to avoid over-generation.
|
||||
|
||||
@@ -282,7 +282,6 @@ python openai_realtime_client.py \
|
||||
--model Qwen/Qwen3-Omni-30B-A3B-Instruct \
|
||||
--input-wav /path/to/input_16k_mono.wav \
|
||||
--output-wav realtime_output.wav \
|
||||
--server-vad \
|
||||
--delta-dump-dir ./rt_delta_wavs
|
||||
```
|
||||
|
||||
@@ -298,32 +297,15 @@ python openai_realtime_client.py \
|
||||
| `--chunk-ms` | `200` | Size of each uploaded audio chunk (milliseconds of audio) |
|
||||
| `--send-delay-ms` | `0` | Delay between chunk sends (simulate realtime upload) |
|
||||
| `--delta-dump-dir` | *(optional)* | Directory to write per-`response.audio.delta` WAV files |
|
||||
| `--server-vad` | disabled | Enable server-side Silero VAD instead of sending explicit commits |
|
||||
| `--num-requests` | `1` | Number of sequential sessions (see `--concurrency`) |
|
||||
| `--concurrency` | `1` | Max concurrent WebSocket sessions when `--num-requests` > 1 |
|
||||
|
||||
Server VAD uses the pinned Silero v6.2 ONNX artifact from `istupakov/silero-vad-onnx` and never downloads model files while processing a session. Make the artifact available in the Hugging Face cache before startup, or configure a local artifact in the deployment YAML:
|
||||
|
||||
```yaml
|
||||
duplex_session:
|
||||
server_vad_model_path: /models/silero_vad.onnx
|
||||
```
|
||||
|
||||
Start from the bundled Qwen3-Omni deployment YAML, add the fields above at the top level, and serve that complete configuration:
|
||||
Ensure the server is running, for example:
|
||||
|
||||
```bash
|
||||
cp vllm_omni/deploy/qwen3_omni_moe.yaml /path/to/qwen3_omni_server_vad.yaml
|
||||
# Edit /path/to/qwen3_omni_server_vad.yaml and add duplex_session.
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct \
|
||||
--omni \
|
||||
--port 8091 \
|
||||
--deploy-config /path/to/qwen3_omni_server_vad.yaml
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
|
||||
```
|
||||
|
||||
Keep Qwen's default `session_mode: turn`; `duplex_session` enables the Realtime handler without changing scheduler
|
||||
semantics. Bare `/v1/realtime` selects that handler. The client uses `?duplex=0` only for the existing non-Server-VAD
|
||||
wire flow; `?duplex=1` remains a supported compatibility alias.
|
||||
|
||||
The Python client supports the following command-line arguments:
|
||||
|
||||
- `--query-type` (or `-q`): Query type (default: `use_video`). Options: `text`, `use_audio`, `use_image`, `use_video`
|
||||
|
||||
@@ -1,12 +1,9 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""Realtime client for vLLM-Omni /v1/realtime (audio + text events).
|
||||
|
||||
This client:
|
||||
1) Reads a local WAV file (must be mono, 16-bit PCM, 16kHz),
|
||||
2) Streams PCM16 chunks to /v1/realtime with OpenAI-style events,
|
||||
3) Receives response.audio.*, response.output_text.*, and transcript events,
|
||||
3) Receives response.audio.* and transcription.* events,
|
||||
4) Saves synthesized audio to an output WAV file and optional text file.
|
||||
|
||||
By default each ``response.audio.delta`` is treated as an **incremental PCM**
|
||||
@@ -21,7 +18,6 @@ Usage:
|
||||
--model Qwen/Qwen3-Omni-30B-A3B-Instruct \
|
||||
--input-wav input_16k_mono.wav \
|
||||
--output-wav realtime_output.wav \
|
||||
--server-vad \
|
||||
--delta-dump-dir ./rt_delta_wavs
|
||||
|
||||
Dependencies:
|
||||
@@ -36,7 +32,6 @@ import base64
|
||||
import json
|
||||
import wave
|
||||
from pathlib import Path
|
||||
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
|
||||
|
||||
try:
|
||||
import websockets
|
||||
@@ -75,14 +70,6 @@ def _write_wav_pcm16(path: Path, pcm16_bytes: bytes, sample_rate_hz: int) -> Non
|
||||
wf.writeframes(pcm16_bytes)
|
||||
|
||||
|
||||
def _with_realtime_route(url: str, *, server_vad: bool) -> str:
|
||||
parts = urlsplit(url)
|
||||
query = [(key, value) for key, value in parse_qsl(parts.query, keep_blank_values=True) if key != "duplex"]
|
||||
if not server_vad:
|
||||
query.append(("duplex", "0"))
|
||||
return urlunsplit(parts._replace(query=urlencode(query)))
|
||||
|
||||
|
||||
async def run_client(
|
||||
url: str,
|
||||
model: str,
|
||||
@@ -92,7 +79,6 @@ async def run_client(
|
||||
chunk_ms: int,
|
||||
send_delay_ms: int,
|
||||
delta_dump_dir: Path | None,
|
||||
server_vad: bool,
|
||||
request_idx: int = 1,
|
||||
total_requests: int = 1,
|
||||
) -> None:
|
||||
@@ -106,44 +92,23 @@ async def run_client(
|
||||
delta_index = 0
|
||||
text_chunks: list[str] = []
|
||||
final_text: str = ""
|
||||
final_audio_transcript: str = ""
|
||||
|
||||
if delta_dump_dir is not None:
|
||||
delta_dump_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Keep the CLI mode authoritative even if the input URL already contains
|
||||
# an explicit duplex query parameter.
|
||||
url = _with_realtime_route(url, server_vad=server_vad)
|
||||
async with websockets.connect(url, max_size=64 * 1024 * 1024) as ws:
|
||||
# 1) Validate model.
|
||||
if server_vad:
|
||||
session_update = {
|
||||
"type": "session.update",
|
||||
"session": {
|
||||
await ws.send(
|
||||
json.dumps(
|
||||
{
|
||||
"type": "session.update",
|
||||
"model": model,
|
||||
"audio": {
|
||||
"input": {
|
||||
"format": {"type": "audio/pcm", "rate": 16_000},
|
||||
"turn_detection": {
|
||||
"type": "server_vad",
|
||||
"threshold": 0.5,
|
||||
"prefix_padding_ms": 300,
|
||||
"silence_duration_ms": 500,
|
||||
"create_response": True,
|
||||
"interrupt_response": False,
|
||||
},
|
||||
}
|
||||
},
|
||||
},
|
||||
}
|
||||
else:
|
||||
# Preserve the existing non-duplex wire format.
|
||||
session_update = {"type": "session.update", "model": model}
|
||||
await ws.send(json.dumps(session_update))
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
# 2) Preserve the legacy explicit-commit flow unless Server VAD owns endpointing.
|
||||
if not server_vad:
|
||||
await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": False}))
|
||||
# 2) Start generation once (non-final commit).
|
||||
await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": False}))
|
||||
|
||||
# 3) Stream audio chunks.
|
||||
for i in range(0, len(pcm16), chunk_bytes):
|
||||
@@ -159,20 +124,8 @@ async def run_client(
|
||||
if send_delay_ms > 0:
|
||||
await asyncio.sleep(send_delay_ms / 1000.0)
|
||||
|
||||
# 4) Finalize the turn explicitly, or append silence for Server VAD endpointing.
|
||||
if server_vad:
|
||||
trailing_silence = bytes(16_000 * 2)
|
||||
for i in range(0, len(trailing_silence), chunk_bytes):
|
||||
await ws.send(
|
||||
json.dumps(
|
||||
{
|
||||
"type": "input_audio_buffer.append",
|
||||
"audio": base64.b64encode(trailing_silence[i : i + chunk_bytes]).decode("utf-8"),
|
||||
}
|
||||
)
|
||||
)
|
||||
else:
|
||||
await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": True}))
|
||||
# 4) Final commit closes input stream.
|
||||
await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": True}))
|
||||
|
||||
# 5) Receive server events until audio done.
|
||||
while True:
|
||||
@@ -191,7 +144,7 @@ async def run_client(
|
||||
sr = event.get("sample_rate_hz")
|
||||
if isinstance(sr, int) and sr > 0:
|
||||
output_sample_rate = sr
|
||||
audio_b64 = event.get("delta" if server_vad else "audio", "")
|
||||
audio_b64 = event.get("audio", "")
|
||||
if audio_b64:
|
||||
pcm_delta = base64.b64decode(audio_b64)
|
||||
incremental_pcm_parts.append(pcm_delta)
|
||||
@@ -205,14 +158,14 @@ async def run_client(
|
||||
)
|
||||
continue
|
||||
|
||||
if event_type in {"response.output_text.delta", "transcription.delta"}:
|
||||
if event_type == "transcription.delta":
|
||||
delta = event.get("delta", "")
|
||||
if delta:
|
||||
text_chunks.append(delta)
|
||||
print(delta, end="", flush=True)
|
||||
continue
|
||||
|
||||
if event_type in {"response.output_text.done", "transcription.done"}:
|
||||
if event_type == "transcription.done":
|
||||
final_text = event.get("text", "") or "".join(text_chunks)
|
||||
usage = event.get("usage")
|
||||
final_text_with_tag = f"Final transcription: {final_text}"
|
||||
@@ -223,16 +176,7 @@ async def run_client(
|
||||
print(f"{log_prefix}text usage: {usage}")
|
||||
continue
|
||||
|
||||
if event_type == "response.audio_transcript.done":
|
||||
final_audio_transcript = event.get("transcript", "")
|
||||
continue
|
||||
|
||||
if event_type == "response.audio.done":
|
||||
if not server_vad:
|
||||
break
|
||||
continue
|
||||
|
||||
if event_type == "response.done" and server_vad:
|
||||
break
|
||||
|
||||
if event_type == "error":
|
||||
@@ -247,7 +191,7 @@ async def run_client(
|
||||
print(f"{log_prefix}Saved realtime audio to: {output_wav} (incremental chunks joined)")
|
||||
|
||||
if output_text is not None:
|
||||
text_to_save = final_text or "".join(text_chunks) or final_audio_transcript
|
||||
text_to_save = final_text if final_text else "".join(text_chunks)
|
||||
output_text.parent.mkdir(parents=True, exist_ok=True)
|
||||
output_text.write_text(text_to_save, encoding="utf-8")
|
||||
print(f"{log_prefix}Saved realtime text to: {output_text}")
|
||||
@@ -269,7 +213,6 @@ async def run_clients_concurrent(
|
||||
chunk_ms: int,
|
||||
send_delay_ms: int,
|
||||
delta_dump_dir: Path | None,
|
||||
server_vad: bool,
|
||||
num_requests: int,
|
||||
concurrency: int,
|
||||
) -> None:
|
||||
@@ -292,7 +235,6 @@ async def run_clients_concurrent(
|
||||
chunk_ms=chunk_ms,
|
||||
send_delay_ms=send_delay_ms,
|
||||
delta_dump_dir=per_delta_dir,
|
||||
server_vad=server_vad,
|
||||
request_idx=index,
|
||||
total_requests=num_requests,
|
||||
)
|
||||
@@ -341,11 +283,6 @@ def main() -> None:
|
||||
default=None,
|
||||
help="If set, each response.audio.delta is saved as delta_NNNNNN.wav under this directory",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--server-vad",
|
||||
action="store_true",
|
||||
help="Let the server detect speech endpoints instead of sending input_audio_buffer.commit",
|
||||
)
|
||||
parser.add_argument("--num-requests", type=int, default=1, help="Total number of requests to send")
|
||||
parser.add_argument(
|
||||
"--concurrency",
|
||||
@@ -372,7 +309,6 @@ def main() -> None:
|
||||
chunk_ms=args.chunk_ms,
|
||||
send_delay_ms=args.send_delay_ms,
|
||||
delta_dump_dir=args.delta_dump_dir,
|
||||
server_vad=args.server_vad,
|
||||
)
|
||||
)
|
||||
else:
|
||||
@@ -386,7 +322,6 @@ def main() -> None:
|
||||
chunk_ms=args.chunk_ms,
|
||||
send_delay_ms=args.send_delay_ms,
|
||||
delta_dump_dir=args.delta_dump_dir,
|
||||
server_vad=args.server_vad,
|
||||
num_requests=args.num_requests,
|
||||
concurrency=concurrency,
|
||||
)
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
# Streaming Video Generation
|
||||
|
||||
This example uses the custom WebSocket endpoint `WS /v1/realtime/video` to receive a video byte stream as chunks are produced.
|
||||
It covers text-to-video, and image-conditioned models through `--image-reference`.
|
||||
It covers text-only video generation. Image/reference input is intentionally not included for now.
|
||||
|
||||
## Start The Server
|
||||
|
||||
@@ -20,7 +20,7 @@ The `--diffusion-streaming-output` CLI flag is forwarded as `streaming_output=Tr
|
||||
|
||||
| Direction | Message | Format | Description |
|
||||
| --- | --- | --- | --- |
|
||||
| Client to server | `session.start` | JSON text: `{"type":"session.start","model":"...","prompt":"...","format":"m4s"}` | Starts generation. `format` is optional and accepts `m4s` (default). Sampling fields such as `width`, `height`, `fps`, `num_frames`, and `extra_params` may be included. Image-conditioned models take a first frame as `image_reference`: `{"image_url":"https://..."}` or a `data:` URL. |
|
||||
| Client to server | `session.start` | JSON text: `{"type":"session.start","model":"...","prompt":"...","format":"m4s"}` | Starts generation. `format` is optional and accepts `m4s` (default). Sampling fields such as `width`, `height`, `fps`, `num_frames`, and `extra_params` may be included. |
|
||||
| Client to server | `session.interaction` | JSON text: `{"type":"session.interaction","interaction":{"event_id":"xxx","event":{"prompt":"..."},"transition_chunks":3}}` | Updates the active prompt midway through generation. `event_id` is optional and `transition_chunks` defaults to the model setting. |
|
||||
| Server to client | `video.start` | JSON text: `{"type":"video.start","request_id":"...","format":"m4s","config":{...}}` | Confirms the session and mirrors the accepted `format`. |
|
||||
| Server to client | `video.chunk_metadata` | JSON text: `{"type":"video.chunk_metadata","request_id":"...","kind":"media","transport_chunk_index":0,"generation_chunk_index":0,"num_frames":9,"byte_length":1234,"started_event_ids":[],"active_event_ids":[],"completed_event_ids":[]}` | Precedes each binary frame and describes the immediately following payload. |
|
||||
@@ -62,7 +62,7 @@ python streaming_video_client.py \
|
||||
The client sends one `session.start` message, prints each received binary video chunk with its byte size and elapsed time, and saves the received bytes to `--output` after `session.done`.
|
||||
The client remuxes the gathered stream to a regular progressive MP4 file so that local playback knows the video duration.
|
||||
|
||||
Schedule midway prompt updates with `--prompt-updates`. None are sent unless you ask for them, and only pipelines that implement midway prompt updates accept them (LingBot-World, for one, does not, and rejects the whole session). Each entry uses `"at"` as seconds on the client clock after the server sends `video.start`:
|
||||
Schedule midway prompt updates with `--prompt-updates`. Each entry uses `"at"` as seconds on the client clock after the server sends `video.start`:
|
||||
|
||||
```bash
|
||||
python streaming_video_client.py \
|
||||
@@ -73,19 +73,6 @@ python streaming_video_client.py \
|
||||
]'
|
||||
```
|
||||
|
||||
The Helios demo continues the default prompt into an underwater storm:
|
||||
|
||||
```bash
|
||||
python streaming_video_client.py \
|
||||
--model BestWishYsh/Helios-Distilled \
|
||||
--prompt-updates '[
|
||||
{"at": 4,
|
||||
"prompt": "An underwater tornado appears and affects the ocean floor in a dramatic and chaotic scene. The water is murky, swirling violently, carrying debris and marine life into the vortex. The tropical fish on the scene all swim in panic, trying to avoid the powerful currents. The camera remains stationary, capturing the intensity of the underwater tornado as it disrupts the serene ocean floor. Close-up shot emphasizing the turbulent motion and destruction."},
|
||||
{"at": 11,
|
||||
"prompt": "The swirling underwater vortex now seizes a heavy, encrusted treasure chest, its lid flapping open as it is smashed onto the ocean floor. Gold coins and silver trinkets spill out, glittering briefly in the murky water before being swept instantly into the violent funnel. The heavy wooden box tumbles end over end, colliding with floating rocks and adding to the debris field. Swirling sediment and bubbles surround the spilling fortune, highlighting the chaotic power of the storm as it ravages the seabed. Close-up shot emphasizing the turbulent motion and destruction."}
|
||||
]'
|
||||
```
|
||||
|
||||
## Run The Gradio Demo
|
||||
|
||||
```bash
|
||||
@@ -119,31 +106,3 @@ Disable it in the CLI example with `--no-helios-distilled-preset`, or override/e
|
||||
python streaming_video_client.py \
|
||||
--extra-params '{"pyramid_num_inference_steps_list":[2, 2, 2]}'
|
||||
```
|
||||
|
||||
### LingBot-World 2.0
|
||||
|
||||
`robbyant/lingbot-world-v2-14b-causal-fast-diffusers` is image-conditioned and runs on the
|
||||
AR-Diffusion engine, so it needs the stepwise deploy config and a first frame. See
|
||||
[the recipe](../../../recipes/Robbyant/LingBot-World-2.0.md) for the constraints on
|
||||
`width`/`height` and `num_frames`.
|
||||
|
||||
```bash
|
||||
vllm serve robbyant/lingbot-world-v2-14b-causal-fast-diffusers \
|
||||
--omni \
|
||||
--deploy-config vllm_omni/deploy/lingbot_world_v2_stepwise.yaml \
|
||||
--port 8000
|
||||
```
|
||||
|
||||
```bash
|
||||
python streaming_video_client.py \
|
||||
--model robbyant/lingbot-world-v2-14b-causal-fast-diffusers \
|
||||
--prompt "The camera moves slowly forward through the scene." \
|
||||
--image-reference /path/to/first_frame.png \
|
||||
--width 832 --height 480 --num-frames 33 --fps 16 --seed 42 \
|
||||
--extra-params '{"camera_action_script":[[["w"],["w"],["w"]],[["a"],[],[]],[[],[],[]]]}' \
|
||||
--output lingbot_world_v2_stream.mp4
|
||||
```
|
||||
|
||||
`camera_action_script` carries one three-latent-frame WASD action list per generated
|
||||
chunk. Mid-session `session.interaction` only updates the prompt, so a LingBot rollout
|
||||
follows the camera script it started with.
|
||||
@@ -1,12 +1,11 @@
|
||||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""WebSocket client for `/v1/realtime/video`.
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Text-only WebSocket client for `/v1/realtime/video`.
|
||||
|
||||
The client prints a line for every binary video chunk it receives and writes
|
||||
the received bytes to disk when the session finishes or is interrupted.
|
||||
``--image-reference`` sends a first frame for image-conditioned models such as
|
||||
LingBot-World 2.0; text-only models need no image.
|
||||
Image/reference input is intentionally not supported in this example yet.
|
||||
|
||||
Requirements:
|
||||
pip install av websockets
|
||||
@@ -14,11 +13,9 @@ Requirements:
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import base64
|
||||
import contextlib
|
||||
import io
|
||||
import json
|
||||
import mimetypes
|
||||
import os
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
@@ -56,18 +53,6 @@ class ScheduledPromptUpdate:
|
||||
transition_chunks: int
|
||||
|
||||
|
||||
def _image_reference(value: str) -> dict[str, Any]:
|
||||
"""Accept an http(s) URL as-is, or inline a local image as a data URL."""
|
||||
if value.startswith(("http://", "https://", "data:")):
|
||||
return {"image_url": value}
|
||||
path = Path(value)
|
||||
if not path.is_file():
|
||||
raise argparse.ArgumentTypeError(f"--image-reference is neither a URL nor a readable file: {value}")
|
||||
media_type = mimetypes.guess_type(path.name)[0] or "image/png"
|
||||
encoded = base64.b64encode(path.read_bytes()).decode()
|
||||
return {"image_url": f"data:{media_type};base64,{encoded}"}
|
||||
|
||||
|
||||
def _parse_extra_params(value: str) -> dict[str, Any]:
|
||||
try:
|
||||
parsed = json.loads(value)
|
||||
@@ -131,7 +116,6 @@ def build_session_start(args: argparse.Namespace) -> dict[str, Any]:
|
||||
"prompt": args.prompt,
|
||||
}
|
||||
|
||||
_maybe_set(payload, "image_reference", args.image_reference)
|
||||
_maybe_set(payload, "negative_prompt", args.negative_prompt)
|
||||
_maybe_set(payload, "width", args.width)
|
||||
_maybe_set(payload, "height", args.height)
|
||||
@@ -378,12 +362,6 @@ def parse_args() -> argparse.Namespace:
|
||||
parser.add_argument("--fps", type=int, default=16, help="Frames per second")
|
||||
parser.add_argument("--num-frames", type=int, default=429, help="Number of generated frames")
|
||||
|
||||
parser.add_argument(
|
||||
"--image-reference",
|
||||
type=_image_reference,
|
||||
default=None,
|
||||
help="First frame for image-conditioned models: a local image path or an http(s) URL.",
|
||||
)
|
||||
parser.add_argument("--negative-prompt", default=None, help="Negative prompt")
|
||||
parser.add_argument("--num-inference-steps", type=int, default=None, help="Number of diffusion steps")
|
||||
parser.add_argument("--guidance-scale", type=float, default=1.0, help="Classifier-free guidance scale")
|
||||
@@ -405,16 +383,36 @@ def parse_args() -> argparse.Namespace:
|
||||
default=None,
|
||||
help="JSON object merged into extra_params; overrides preset keys on conflict.",
|
||||
)
|
||||
# Midway prompt updates are pipeline-specific: a model that does not implement
|
||||
# them (LingBot-World, for one) rejects the whole session on the first one, so
|
||||
# nothing is scheduled unless the caller asks for it.
|
||||
parser.add_argument(
|
||||
"--prompt-updates",
|
||||
type=_parse_prompt_updates,
|
||||
default=[],
|
||||
default=json.dumps(
|
||||
[
|
||||
{
|
||||
"at": 4,
|
||||
"prompt": (
|
||||
"An underwater tornado appears and affects the ocean floor in a dramatic and chaotic scene. "
|
||||
"The water is murky, swirling violently, carrying debris and marine life into the vortex. "
|
||||
"The tropical fish on the scene all swim in panic, trying to avoid the powerful currents. "
|
||||
"The camera remains stationary, capturing the intensity of the underwater tornado as it disrupts the serene ocean floor. "
|
||||
"Close-up shot emphasizing the turbulent motion and destruction."
|
||||
),
|
||||
},
|
||||
{
|
||||
"at": 11,
|
||||
"prompt": (
|
||||
"The swirling underwater vortex now seizes a heavy, encrusted treasure chest, its lid flapping open as it is smashed onto the ocean floor. "
|
||||
"Gold coins and silver trinkets spill out, glittering briefly in the murky water before being swept instantly into the violent funnel. "
|
||||
"The heavy wooden box tumbles end over end, colliding with floating rocks and adding to the debris field. "
|
||||
"Swirling sediment and bubbles surround the spilling fortune, highlighting the chaotic power of the storm as it ravages the seabed. "
|
||||
"Close-up shot emphasizing the turbulent motion and destruction."
|
||||
),
|
||||
# "transition_chunks": 3,
|
||||
},
|
||||
]
|
||||
),
|
||||
help=(
|
||||
"JSON array of scheduled prompt updates; none are sent by default. "
|
||||
"Each object requires "
|
||||
"JSON array of scheduled prompt updates. Each object requires "
|
||||
"'at' (seconds after video.start) and 'prompt'. Optional "
|
||||
f"'transition_chunks' defaults to {DEFAULT_TRANSITION_CHUNKS}."
|
||||
),
|
||||
|
||||
@@ -14,7 +14,6 @@ For the full list of supported architectures across all modalities, see
|
||||
|
||||
| Model | HuggingFace repo | Voice cloning | Streaming | Voice presets / upload | Gradio demo |
|
||||
| --- | --- | --- | --- | --- | --- |
|
||||
| Audio8 TTS Preview | `Audio8/Audio8-TTS-Preview-0.6b` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | uploaded audio voice only; no presets | ✓ |
|
||||
| Fish Speech S2 Pro | `fishaudio/s2-pro` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | ✓ |
|
||||
| GLM-TTS | `zai-org/GLM-TTS` | ✓ (`ref_audio`+`ref_text`, required) | ✓ (PCM stream) | — | ✓ |
|
||||
| IndexTTS-2 | `IndexTeam/IndexTTS-2` | ✓ (`ref_audio` or uploaded `voice`) | `stream=true` response, non-chunk | uploaded audio voice only; no presets | — |
|
||||
@@ -254,152 +253,16 @@ python examples/online_serving/text_to_speech/indextts2/speech_client.py \
|
||||
|
||||
---
|
||||
|
||||
## Audio8 TTS Preview
|
||||
## Fish Speech S2 Pro
|
||||
|
||||
0.6B DualAR TTS at 44.1 kHz, 11 languages, zero-shot voice cloning.
|
||||
4B dual-AR TTS at 44.1 kHz. Server uses the DAC codec.
|
||||
|
||||
### Prerequisites
|
||||
|
||||
None beyond the base install: the neural audio codec is implemented in tree and
|
||||
its weights (`codec.pth`) ship with the checkpoint.
|
||||
|
||||
### Launch
|
||||
|
||||
```bash
|
||||
vllm serve Audio8/Audio8-TTS-Preview-0.6b --omni --port 8092
|
||||
# or:
|
||||
./audio8_tts/run_server.sh
|
||||
pip install fish-speech
|
||||
```
|
||||
|
||||
The deploy config auto-loads from `vllm_omni/deploy/audio8_tts.yaml` (HF
|
||||
`model_type` is `arktts`). Do **not** pass `--trust-remote-code`: vllm-omni
|
||||
registers its own `arktts` config, and transformers would otherwise prefer the
|
||||
checkpoint's remote code and bypass it.
|
||||
|
||||
### Text-only synthesis
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8092/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "Welcome to Audio8 TTS.",
|
||||
"response_format": "wav"
|
||||
}' --output output.wav
|
||||
```
|
||||
|
||||
### Voice cloning
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8092/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "Welcome to Audio8 TTS.",
|
||||
"ref_audio": "https://example.com/reference.wav",
|
||||
"ref_text": "The exact transcript of the reference recording."
|
||||
}' --output cloned.wav
|
||||
```
|
||||
|
||||
`ref_audio` also accepts a `data:audio/wav;base64,...` URL. `ref_text` is
|
||||
mandatory whenever `ref_audio` is present and must match the recording.
|
||||
Uploading a voice via `POST /v1/audio/voices` lets the server reuse the encoded
|
||||
reference codes across requests (`voice: "<name>"`).
|
||||
|
||||
### Streaming
|
||||
|
||||
```bash
|
||||
python audio8_tts/speech_client.py --text "Welcome to Audio8 TTS." --stream --output out.pcm
|
||||
ffplay -f s16le -ar 44100 -ac 1 out.pcm
|
||||
```
|
||||
|
||||
### Gradio demo
|
||||
|
||||
```bash
|
||||
./audio8_tts/run_gradio_demo.sh # server + demo
|
||||
python audio8_tts/gradio_demo.py --api-base http://localhost:8092 # demo only
|
||||
```
|
||||
|
||||
### Benchmark
|
||||
|
||||
```bash
|
||||
export VLLM_OMNI_BENCH_AUDIO_SAMPLE_RATE=44100 # required: harness defaults to 24 kHz
|
||||
vllm bench serve --omni --host 127.0.0.1 --port 8092 \
|
||||
--model Audio8/Audio8-TTS-Preview-0.6b \
|
||||
--backend openai-audio-speech --endpoint /v1/audio/speech \
|
||||
--dataset-name seed-tts-text --dataset-path benchmarks/build_dataset/seed_tts_smoke \
|
||||
--seed-tts-locale en --num-prompts 20 --num-warmups 2 \
|
||||
--max-concurrency 1 --request-rate inf \
|
||||
--percentile-metrics ttft,e2el,audio_rtf,audio_ttfp,audio_duration,audio_underrun
|
||||
```
|
||||
|
||||
Measured on 1x H20 **shared with an unrelated job holding ~80% SM**, so treat
|
||||
these as a lower bound. 20 prompts from `seed_tts_smoke/en`, mean audio 4.5 s,
|
||||
deploy defaults:
|
||||
|
||||
| | HF reference (batch=1) | vllm-omni c=1 | c=4 | c=8 |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| RTF (mean) | 1.008 | **0.19** | 0.30 | 0.54 |
|
||||
| Throughput (req/s) | 0.23 | 1.18 | 2.84 | **2.93** |
|
||||
| E2E latency (mean, ms) | 4375 | 850 | 1349 | 2361 |
|
||||
| Time to first packet (mean, ms) | n/a (no streaming) | 63 | 105 | 1124 |
|
||||
| Underrun p99 / continuity OK | n/a | 0 s / 100% | 0 s / 100% | 0 s / 100% |
|
||||
|
||||
~5.3x lower per-request RTF than the upstream reference implementation at
|
||||
concurrency 1, ~12.8x aggregate throughput at concurrency 8.
|
||||
|
||||
#### Stage-0 `max_num_seqs`: first-packet vs throughput
|
||||
|
||||
This subsection is operator tuning guidance, not a claim this PR optimizes
|
||||
anything: the shipped deploy default is unchanged (`max_num_seqs: 4`). The
|
||||
numbers below are a single run on the same contended H20 as above (one unrelated
|
||||
job holding ~80% SM, no repeats), so read them as directional, not as measured
|
||||
speedups. Same setup, varying stage 0's `max_num_seqs` (stage 1 stays at 1):
|
||||
|
||||
| stage-0 `max_num_seqs` | client concurrency | req/s | RTF | TTFP mean / p99 (ms) | underrun p99 (s) |
|
||||
| --- | --- | --- | --- | --- | --- |
|
||||
| 4 (default) | 1 | 1.18 | 0.19 | 63 / 67 | 0.00 |
|
||||
| 4 | 4 | 2.84 | 0.30 | 105 / 148 | 0.00 |
|
||||
| 4 | 8 | 2.93 | 0.54 | 1124 / 1642 | 0.00 |
|
||||
| 8 | 1 | 1.16 | 0.19 | 64 / 70 | 0.00 |
|
||||
| 8 | 4 | 2.88 | 0.30 | 107 / 148 | 0.00 |
|
||||
| 8 | 8 | **3.72** | 0.40 | **247** / 733 | 0.22 |
|
||||
| 8 | 16 | 3.91 | 0.65 | 1366 / 2530 | 0.29 |
|
||||
|
||||
Reading of this:
|
||||
|
||||
- Below the limit (`concurrency <= max_num_seqs`) the setting is irrelevant, as
|
||||
expected — c=1 and c=4 are identical for both.
|
||||
- At c=8, raising the limit to 8 admits every request instead of queueing half
|
||||
of them: in this single run, roughly +27% throughput and ~4.5x lower
|
||||
first-packet latency.
|
||||
- Throughput saturates around 3.7-3.9 req/s (~17 s of audio per second). Past
|
||||
c=8 you only buy queueing delay.
|
||||
- The cost of admitting more streams is a 0.22-0.29 s worst-case buffer deficit
|
||||
in 1 of 20 requests (`Streaming continuity OK rate` 95%), because stage 1 runs
|
||||
`max_num_seqs: 1` and the codec window round-robins across streams. Client-side
|
||||
prebuffering (the bundled Gradio demo holds 2 chunks) absorbs that; raising
|
||||
stage 1's `max_num_seqs` is the untested knob if you need both.
|
||||
|
||||
The shipped default stays at 4 because every measured point there had zero
|
||||
underrun. Raise it to 8 if first-packet latency under load matters more than a
|
||||
rare sub-300 ms gap.
|
||||
|
||||
### Notes
|
||||
|
||||
- Output: 44.1 kHz mono; ~21.5 codec frames per second.
|
||||
- No built-in speaker presets. Omit `voice` for a random timbre, or clone one.
|
||||
- `max_new_tokens` caps generated codec frames (1 frame ~= 46 ms).
|
||||
- Context is 2048 packed text+audio positions (`max_model_len: 2048`).
|
||||
- Do not send `temperature: 0`: greedy decoding is degenerate for this
|
||||
checkpoint and truncates the utterance. The deploy defaults
|
||||
(0.7 / top_k 50 / top_p 0.9) mirror `generation_config.json`.
|
||||
|
||||
---
|
||||
|
||||
## Fish Speech S2 Pro
|
||||
|
||||
4B dual-AR TTS at 44.1 kHz. Server uses the DAC codec, which is vendored in
|
||||
vLLM-Omni (no extra packages required).
|
||||
|
||||
### Kvcache attention fast path
|
||||
|
||||
Fish Speech S2 Pro uses a Triton decode-only kvcache attention fast path by
|
||||
@@ -562,46 +425,6 @@ python examples/online_serving/text_to_speech/ming_flash_omni_tts/speech_client.
|
||||
|
||||
---
|
||||
|
||||
## MOSS-TTS Local Transformer v1.5
|
||||
|
||||
For a single H200, the optional
|
||||
[`moss_tts_local_h200.yaml`](../../../vllm_omni/deploy/moss_tts_local_h200.yaml)
|
||||
deployment places both the talker and codec on logical GPU 0:
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=0 VLLM_OMNI_EVENT_DRIVEN_ORCH=1 \
|
||||
vllm serve OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 \
|
||||
--omni --trust-remote-code \
|
||||
--deploy-config vllm_omni/deploy/moss_tts_local_h200.yaml \
|
||||
--disable-log-stats
|
||||
```
|
||||
|
||||
Run from the repository root and select an available physical GPU with
|
||||
`CUDA_VISIBLE_DEVICES`. The command enables event-driven orchestration and
|
||||
disables detailed per-request statistics logging to reduce CPU overhead at
|
||||
high concurrency. Remove `--disable-log-stats` when those statistics are needed;
|
||||
keep these settings identical when comparing performance.
|
||||
This preset configures 256 request slots per stage,
|
||||
a 32 GiB talker KV cache, talker CUDA Graph buckets through 512 scheduled tokens,
|
||||
and codec batch buckets through 256. It requires more than 80 GiB of GPU memory;
|
||||
request capacity also depends on input and generated lengths. The default
|
||||
[`moss_tts_local.yaml`](../../../vllm_omni/deploy/moss_tts_local.yaml) remains
|
||||
available for smaller deployments.
|
||||
|
||||
The preset enables the optional codec backend with
|
||||
`hf_overrides.codec_attention_backend: triton` on stage 1. It preserves the
|
||||
streaming ring-cache mask and uses BF16 attention with 64-dimensional heads;
|
||||
other attention shapes use PyTorch SDPA. Set the backend to `sdpa` to use the
|
||||
default attention implementation. Codec terminal tails share execution only
|
||||
when they already map to the same padded CUDA Graph; returned audio retains
|
||||
each request's actual length.
|
||||
|
||||
Voice cloning requests use `ref_audio` and `ref_text`. For streaming output,
|
||||
set `stream: true`, `stream_format: "audio"`, and `response_format: "pcm"`;
|
||||
the native PCM format is 48 kHz, stereo, signed 16-bit little-endian. Compute
|
||||
audio throughput as `PCM bytes / (48000 * 2 * 2) / elapsed seconds`, and compare
|
||||
the same requests, concurrency, and reference-cache state.
|
||||
|
||||
## MOSS-TTS-Nano
|
||||
|
||||
Single-stage 0.1B AR LM + MOSS-Audio-Tokenizer-Nano codec at 48 kHz mono. Every request must include `ref_audio`; there are no built-in speaker presets.
|
||||
@@ -845,21 +668,13 @@ Raw PCM streaming requires `stream_format="audio"`, `response_format="pcm"`, and
|
||||
|
||||
### Streaming WebSocket
|
||||
|
||||
The `/v1/audio/speech/stream` endpoint accepts text incrementally and, by default, synthesizes the buffered text as one continuous request on `input.done`:
|
||||
The `/v1/audio/speech/stream` endpoint accepts text incrementally and synthesizes the buffered text as one continuous request on `input.done`:
|
||||
|
||||
```bash
|
||||
python qwen3_tts/streaming_speech_client.py --text "Hello world. How are you? I am fine."
|
||||
python qwen3_tts/streaming_speech_client.py --text "..." --simulate-stt --stt-delay 0.1
|
||||
```
|
||||
|
||||
For per-sentence audio (including Indic danda `।`) before `input.done`:
|
||||
|
||||
```bash
|
||||
python qwen3_tts/streaming_speech_client.py \
|
||||
--text "नमस्ते। कैसे हो?" \
|
||||
--split-granularity sentence
|
||||
```
|
||||
|
||||
`input.done` flushes without closing, so repeating `--text` synthesizes several utterances over one connection:
|
||||
|
||||
```bash
|
||||
|
||||
@@ -1,258 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""Gradio demo for Audio8 TTS Preview via the /v1/audio/speech API.
|
||||
|
||||
Supports:
|
||||
- text-to-speech synthesis in 11 languages
|
||||
- zero-shot voice cloning from reference audio (upload, microphone, or URL)
|
||||
- streaming (progressive PCM chunks) and non-streaming modes
|
||||
|
||||
Usage:
|
||||
# Start the server first (see run_server.sh), then:
|
||||
python gradio_demo.py --api-base http://localhost:8092
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import io
|
||||
|
||||
try:
|
||||
import gradio as gr
|
||||
except ImportError:
|
||||
raise ImportError("gradio is required to run this demo. Install it with: pip install 'vllm-omni[demo]'") from None
|
||||
import httpx
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
|
||||
#: Audio8 TTS decodes at 44.1 kHz; raw PCM must be interpreted at that rate.
|
||||
PCM_SAMPLE_RATE = 44100
|
||||
DEFAULT_API_BASE = "http://localhost:8092"
|
||||
|
||||
|
||||
def encode_audio_to_base64(audio_data: tuple) -> str:
|
||||
"""Encode a Gradio ``(sample_rate, numpy_array)`` input as a data URL."""
|
||||
sample_rate, audio_np = audio_data
|
||||
if audio_np.dtype != np.int16:
|
||||
if audio_np.dtype in (np.float32, np.float64):
|
||||
audio_np = (np.clip(audio_np, -1.0, 1.0) * 32767).astype(np.int16)
|
||||
else:
|
||||
audio_np = audio_np.astype(np.int16)
|
||||
buf = io.BytesIO()
|
||||
sf.write(buf, audio_np, sample_rate, format="WAV")
|
||||
return f"data:audio/wav;base64,{base64.b64encode(buf.getvalue()).decode('utf-8')}"
|
||||
|
||||
|
||||
def build_payload(
|
||||
text: str,
|
||||
ref_audio: tuple | None,
|
||||
ref_audio_url: str,
|
||||
ref_text: str,
|
||||
response_format: str = "wav",
|
||||
stream: bool = False,
|
||||
) -> dict:
|
||||
"""Build the ``/v1/audio/speech`` payload.
|
||||
|
||||
Audio8 TTS has no built-in speakers: omit ``voice`` for a random timbre, or
|
||||
supply ``ref_audio`` + ``ref_text`` to clone a voice.
|
||||
"""
|
||||
if not text or not text.strip():
|
||||
raise gr.Error("Please enter text to synthesize.")
|
||||
|
||||
payload: dict = {
|
||||
"input": text.strip(),
|
||||
"response_format": "pcm" if stream else response_format,
|
||||
"stream": stream,
|
||||
}
|
||||
if stream:
|
||||
payload["stream_format"] = "audio"
|
||||
|
||||
ref_url = ref_audio_url.strip() if ref_audio_url else ""
|
||||
if ref_audio is not None:
|
||||
payload["ref_audio"] = encode_audio_to_base64(ref_audio)
|
||||
elif ref_url:
|
||||
payload["ref_audio"] = ref_url
|
||||
|
||||
if "ref_audio" in payload:
|
||||
if not ref_text or not ref_text.strip():
|
||||
raise gr.Error("Voice cloning requires the transcript of the reference audio.")
|
||||
payload["ref_text"] = ref_text.strip()
|
||||
|
||||
return payload
|
||||
|
||||
|
||||
def generate_speech(api_base: str, text: str, ref_audio, ref_audio_url, ref_text, response_format):
|
||||
"""Non-streaming request; returns the whole waveform."""
|
||||
payload = build_payload(text, ref_audio, ref_audio_url, ref_text, response_format, stream=False)
|
||||
try:
|
||||
with httpx.Client(timeout=300.0) as client:
|
||||
resp = client.post(
|
||||
f"{api_base}/v1/audio/speech",
|
||||
json=payload,
|
||||
headers={"Content-Type": "application/json", "Authorization": "Bearer EMPTY"},
|
||||
)
|
||||
except httpx.TimeoutException:
|
||||
raise gr.Error("Request timed out. The server may be busy.") from None
|
||||
except httpx.ConnectError:
|
||||
raise gr.Error(f"Cannot connect to server at {api_base}. Is the server running?") from None
|
||||
|
||||
if resp.status_code != 200:
|
||||
raise gr.Error(f"Server error ({resp.status_code}): {resp.text}")
|
||||
if "application/json" in resp.headers.get("content-type", ""):
|
||||
raise gr.Error(f"Server error: {resp.text}")
|
||||
|
||||
try:
|
||||
if response_format == "pcm":
|
||||
audio_np = np.frombuffer(resp.content, dtype=np.int16).astype(np.float32) / 32767.0
|
||||
return (PCM_SAMPLE_RATE, audio_np)
|
||||
audio_np, sample_rate = sf.read(io.BytesIO(resp.content))
|
||||
if audio_np.ndim > 1:
|
||||
audio_np = audio_np[:, 0]
|
||||
return (sample_rate, audio_np.astype(np.float32))
|
||||
except Exception as e:
|
||||
raise gr.Error(f"Failed to decode audio: {e}") from None
|
||||
|
||||
|
||||
def generate_speech_stream(api_base, text, ref_audio, ref_audio_url, ref_text, response_format, chunk_seconds=0.5):
|
||||
"""Streaming request; yields PCM chunks for progressive playback."""
|
||||
payload = build_payload(text, ref_audio, ref_audio_url, ref_text, response_format, stream=True)
|
||||
min_chunk_samples = int(PCM_SAMPLE_RATE * chunk_seconds)
|
||||
prebuffer_chunks = 2
|
||||
pending: list[np.ndarray] = []
|
||||
pending_len = 0
|
||||
chunks_yielded = 0
|
||||
prebuffer: list[np.ndarray] = []
|
||||
leftover = b""
|
||||
|
||||
try:
|
||||
with (
|
||||
httpx.Client(timeout=300.0) as client,
|
||||
client.stream(
|
||||
"POST",
|
||||
f"{api_base}/v1/audio/speech",
|
||||
json=payload,
|
||||
headers={"Content-Type": "application/json", "Authorization": "Bearer EMPTY"},
|
||||
) as resp,
|
||||
):
|
||||
if resp.status_code != 200:
|
||||
resp.read()
|
||||
raise gr.Error(f"Server error ({resp.status_code}): {resp.text}")
|
||||
for chunk in resp.iter_bytes():
|
||||
if not chunk:
|
||||
continue
|
||||
raw = leftover + chunk
|
||||
# s16le: never split a sample across chunk boundaries.
|
||||
usable = len(raw) - (len(raw) % 2)
|
||||
leftover = raw[usable:]
|
||||
if usable == 0:
|
||||
continue
|
||||
samples = np.frombuffer(raw[:usable], dtype=np.int16).copy().astype(np.float32) / 32767.0
|
||||
pending.append(samples)
|
||||
pending_len += len(samples)
|
||||
if pending_len < min_chunk_samples:
|
||||
continue
|
||||
audio_chunk = np.concatenate(pending)
|
||||
pending.clear()
|
||||
pending_len = 0
|
||||
if chunks_yielded < prebuffer_chunks:
|
||||
prebuffer.append(audio_chunk)
|
||||
chunks_yielded += 1
|
||||
if chunks_yielded == prebuffer_chunks:
|
||||
yield (PCM_SAMPLE_RATE, np.concatenate(prebuffer))
|
||||
prebuffer.clear()
|
||||
else:
|
||||
yield (PCM_SAMPLE_RATE, audio_chunk)
|
||||
remaining = prebuffer + pending
|
||||
if remaining:
|
||||
yield (PCM_SAMPLE_RATE, np.concatenate(remaining))
|
||||
except httpx.TimeoutException:
|
||||
raise gr.Error("Request timed out. The server may be busy.") from None
|
||||
except httpx.ConnectError:
|
||||
raise gr.Error(f"Cannot connect to server at {api_base}. Is the server running?") from None
|
||||
|
||||
|
||||
def on_stream_change(stream: bool):
|
||||
"""Streaming only makes sense with raw PCM."""
|
||||
if stream:
|
||||
return gr.update(value="pcm", interactive=False)
|
||||
return gr.update(value="wav", interactive=True)
|
||||
|
||||
|
||||
def build_interface(api_base: str, stream_chunk_seconds: float = 0.5):
|
||||
with gr.Blocks(title="Audio8 TTS Preview Demo") as demo:
|
||||
gr.Markdown("# Audio8 TTS Preview 0.6B - Text to Speech")
|
||||
gr.Markdown(f"**Server:** `{api_base}` | **Model:** Audio8/Audio8-TTS-Preview-0.6b | **Output:** 44.1 kHz")
|
||||
|
||||
with gr.Row():
|
||||
with gr.Column(scale=3):
|
||||
text_input = gr.Textbox(label="Text to Synthesize", placeholder="Enter text here...", lines=4)
|
||||
|
||||
with gr.Accordion("Voice Cloning (optional)", open=False):
|
||||
gr.Markdown(
|
||||
"Upload or link a short reference recording (5-30 s) and paste its **exact** transcript. "
|
||||
"A transcript that does not match the audio degrades both stability and speaker similarity."
|
||||
)
|
||||
ref_audio = gr.Audio(label="Reference Audio", type="numpy", sources=["upload", "microphone"])
|
||||
ref_audio_url = gr.Textbox(
|
||||
label="Reference Audio URL (alternative to upload)",
|
||||
placeholder="https://example.com/reference.wav",
|
||||
lines=1,
|
||||
)
|
||||
ref_text = gr.Textbox(
|
||||
label="Reference Audio Transcript (required for cloning)",
|
||||
placeholder="Exact transcript of the reference audio...",
|
||||
lines=2,
|
||||
)
|
||||
|
||||
with gr.Row():
|
||||
response_format = gr.Dropdown(
|
||||
choices=["wav", "mp3", "flac", "pcm"], value="wav", label="Audio Format", scale=1
|
||||
)
|
||||
stream_checkbox = gr.Checkbox(
|
||||
label="Stream output", value=False, info="Progressive PCM streaming", scale=1
|
||||
)
|
||||
|
||||
generate_btn = gr.Button("Generate Speech", variant="primary", size="lg")
|
||||
|
||||
with gr.Column(scale=2):
|
||||
audio_output = gr.Audio(label="Generated Audio", interactive=False, streaming=True, autoplay=True)
|
||||
gr.Markdown(
|
||||
"### About\n"
|
||||
"- **Audio8 TTS Preview** — 0.6B DualAR model (slow AR + fast AR codebook decoder)\n"
|
||||
"- **Languages**: Cantonese, Chinese, Dutch, English, French, German, "
|
||||
"Italian, Japanese, Korean, Polish, Spanish\n"
|
||||
"- **Voice cloning**: zero-shot from a reference clip + transcript\n"
|
||||
"- **44.1 kHz** output, ~21.5 codec frames per second"
|
||||
)
|
||||
|
||||
stream_checkbox.change(fn=on_stream_change, inputs=[stream_checkbox], outputs=[response_format])
|
||||
all_inputs = [text_input, ref_audio, ref_audio_url, ref_text, response_format]
|
||||
|
||||
def dispatch(stream_enabled, *args):
|
||||
if stream_enabled:
|
||||
yield from generate_speech_stream(api_base, *args, chunk_seconds=stream_chunk_seconds)
|
||||
else:
|
||||
yield generate_speech(api_base, *args)
|
||||
|
||||
generate_btn.click(fn=dispatch, inputs=[stream_checkbox] + all_inputs, outputs=[audio_output])
|
||||
demo.queue()
|
||||
return demo
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Gradio demo for Audio8 TTS Preview")
|
||||
parser.add_argument("--api-base", default=DEFAULT_API_BASE, help=f"API base URL (default: {DEFAULT_API_BASE})")
|
||||
parser.add_argument("--host", default="0.0.0.0", help="Gradio host")
|
||||
parser.add_argument("--port", type=int, default=7861, help="Gradio port")
|
||||
parser.add_argument("--share", action="store_true", help="Create a public share link")
|
||||
parser.add_argument("--stream-chunk-seconds", type=float, default=0.5, help="Seconds of audio per yielded chunk")
|
||||
args = parser.parse_args()
|
||||
|
||||
print(f"Connecting to vLLM server at: {args.api_base}")
|
||||
build_interface(args.api_base, stream_chunk_seconds=args.stream_chunk_seconds).launch(
|
||||
server_name=args.host, server_port=args.port, share=args.share
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,44 +0,0 @@
|
||||
#!/bin/bash
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
# Launch the Audio8 TTS Preview server + Gradio demo together.
|
||||
#
|
||||
# Usage:
|
||||
# ./run_gradio_demo.sh
|
||||
# CUDA_VISIBLE_DEVICES=0 PORT=8092 GRADIO_PORT=7861 ./run_gradio_demo.sh
|
||||
|
||||
set -e
|
||||
|
||||
MODEL="${MODEL:-Audio8/Audio8-TTS-Preview-0.6b}"
|
||||
PORT="${PORT:-8092}"
|
||||
GRADIO_PORT="${GRADIO_PORT:-7861}"
|
||||
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
|
||||
|
||||
echo "Starting Audio8 TTS Preview server (port $PORT)..."
|
||||
vllm serve "$MODEL" \
|
||||
--omni \
|
||||
--host 0.0.0.0 \
|
||||
--port "$PORT" &
|
||||
SERVER_PID=$!
|
||||
|
||||
cleanup() {
|
||||
echo "Stopping server (PID $SERVER_PID)..."
|
||||
kill $SERVER_PID 2>/dev/null
|
||||
wait $SERVER_PID 2>/dev/null
|
||||
}
|
||||
trap cleanup EXIT
|
||||
|
||||
echo "Waiting for server to start..."
|
||||
for _ in $(seq 1 120); do
|
||||
if curl -s "http://localhost:$PORT/health" > /dev/null 2>&1; then
|
||||
echo "Server ready."
|
||||
break
|
||||
fi
|
||||
sleep 2
|
||||
done
|
||||
|
||||
echo "Starting Gradio demo (port $GRADIO_PORT)..."
|
||||
python "$SCRIPT_DIR/gradio_demo.py" \
|
||||
--api-base "http://localhost:$PORT" \
|
||||
--port "$GRADIO_PORT"
|
||||
@@ -1,24 +0,0 @@
|
||||
#!/bin/bash
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
# Launch a vLLM-Omni server for Audio8 TTS Preview 0.6B.
|
||||
#
|
||||
# Usage:
|
||||
# ./run_server.sh
|
||||
# CUDA_VISIBLE_DEVICES=0 PORT=8092 ./run_server.sh
|
||||
|
||||
set -e
|
||||
|
||||
MODEL="${MODEL:-Audio8/Audio8-TTS-Preview-0.6b}"
|
||||
PORT="${PORT:-8092}"
|
||||
|
||||
echo "Starting Audio8 TTS Preview server with model: $MODEL"
|
||||
|
||||
# --trust-remote-code is deliberately NOT passed: vllm-omni registers its own
|
||||
# `arktts` config, and transformers would otherwise prefer the checkpoint's
|
||||
# remote code and bypass it.
|
||||
vllm serve "$MODEL" \
|
||||
--omni \
|
||||
--host 0.0.0.0 \
|
||||
--port "$PORT"
|
||||
@@ -1,131 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""Client for Audio8 TTS Preview via the /v1/audio/speech endpoint.
|
||||
|
||||
Examples:
|
||||
# Basic TTS
|
||||
python speech_client.py --text "Welcome to Audio8 TTS."
|
||||
|
||||
# Zero-shot voice cloning (the transcript must match the reference audio)
|
||||
python speech_client.py --text "Welcome to Audio8 TTS." \
|
||||
--ref-audio ref.wav --ref-text "The exact transcript of the reference recording."
|
||||
|
||||
# Streaming PCM output
|
||||
python speech_client.py --text "Welcome to Audio8 TTS." --stream --output output.pcm
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import os
|
||||
|
||||
import httpx
|
||||
|
||||
DEFAULT_API_BASE = "http://localhost:8092"
|
||||
DEFAULT_API_KEY = "EMPTY"
|
||||
DEFAULT_MODEL = "Audio8/Audio8-TTS-Preview-0.6b"
|
||||
#: The codec runs at 44.1 kHz, so raw PCM must be played back at that rate.
|
||||
PCM_SAMPLE_RATE = 44100
|
||||
|
||||
|
||||
def encode_audio_to_base64(audio_path: str) -> str:
|
||||
"""Encode a local audio file as a base64 data URL."""
|
||||
if not os.path.exists(audio_path):
|
||||
raise FileNotFoundError(f"Audio file not found: {audio_path}")
|
||||
ext = audio_path.lower().rsplit(".", 1)[-1]
|
||||
mime_type = {"wav": "audio/wav", "mp3": "audio/mpeg", "flac": "audio/flac", "ogg": "audio/ogg"}.get(
|
||||
ext, "audio/wav"
|
||||
)
|
||||
with open(audio_path, "rb") as handle:
|
||||
audio_b64 = base64.b64encode(handle.read()).decode("utf-8")
|
||||
return f"data:{mime_type};base64,{audio_b64}"
|
||||
|
||||
|
||||
def build_payload(args) -> dict:
|
||||
payload = {
|
||||
"model": args.model,
|
||||
"input": args.text,
|
||||
"response_format": args.response_format,
|
||||
}
|
||||
if args.ref_audio:
|
||||
if args.ref_audio.startswith(("http://", "https://")):
|
||||
payload["ref_audio"] = args.ref_audio
|
||||
else:
|
||||
payload["ref_audio"] = encode_audio_to_base64(args.ref_audio)
|
||||
if args.ref_text:
|
||||
payload["ref_text"] = args.ref_text
|
||||
if args.max_new_tokens is not None:
|
||||
payload["max_new_tokens"] = args.max_new_tokens
|
||||
if args.stream:
|
||||
payload["stream"] = True
|
||||
payload["stream_format"] = "audio"
|
||||
payload["response_format"] = "pcm"
|
||||
return payload
|
||||
|
||||
|
||||
def run_tts(args) -> None:
|
||||
payload = build_payload(args)
|
||||
api_url = f"{args.api_base}/v1/audio/speech"
|
||||
headers = {"Content-Type": "application/json", "Authorization": f"Bearer {args.api_key}"}
|
||||
|
||||
print(f"Model: {args.model}")
|
||||
print(f"Text: {args.text}")
|
||||
if args.ref_audio:
|
||||
print(f"Voice cloning: ref_audio={args.ref_audio!r} ref_text={args.ref_text!r}")
|
||||
|
||||
if args.stream:
|
||||
output_path = args.output or "output.pcm"
|
||||
with (
|
||||
httpx.Client(timeout=300.0) as client,
|
||||
client.stream("POST", api_url, json=payload, headers=headers) as resp,
|
||||
):
|
||||
if resp.status_code != 200:
|
||||
print(f"Error {resp.status_code}: {resp.read().decode()}")
|
||||
return
|
||||
total_bytes = 0
|
||||
with open(output_path, "wb") as handle:
|
||||
for chunk in resp.iter_bytes():
|
||||
handle.write(chunk)
|
||||
total_bytes += len(chunk)
|
||||
seconds = total_bytes / 2 / PCM_SAMPLE_RATE
|
||||
print(f"Streamed {total_bytes} bytes (~{seconds:.2f}s of s16le @ {PCM_SAMPLE_RATE} Hz) to {output_path}")
|
||||
print(f"Play with: ffplay -f s16le -ar {PCM_SAMPLE_RATE} -ac 1 {output_path}")
|
||||
return
|
||||
|
||||
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}: {response.text}")
|
||||
return
|
||||
if response.headers.get("content-type", "").startswith("application/json"):
|
||||
print(f"Error: {response.text}")
|
||||
return
|
||||
|
||||
output_path = args.output or f"output.{args.response_format}"
|
||||
with open(output_path, "wb") as handle:
|
||||
handle.write(response.content)
|
||||
print(f"Audio saved to: {output_path} ({len(response.content)} bytes)")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Audio8 TTS Preview 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")
|
||||
parser.add_argument("--text", required=True, help="Text to synthesize")
|
||||
parser.add_argument("--ref-audio", default=None, help="Reference audio for voice cloning (path or URL)")
|
||||
parser.add_argument("--ref-text", default=None, help="Transcript of the reference audio")
|
||||
parser.add_argument("--max-new-tokens", type=int, default=None, help="Cap the number of generated codec frames")
|
||||
parser.add_argument("--stream", action="store_true", help="Stream PCM instead of returning a whole file")
|
||||
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")
|
||||
run_tts(parser.parse_args())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -24,11 +24,6 @@ Usage:
|
||||
--text "Hello world. How are you? I am fine." \
|
||||
--simulate-stt --stt-delay 0.1
|
||||
|
||||
# Opt into per-sentence synthesis (Indic danda, CJK, Latin)
|
||||
python streaming_speech_client.py \
|
||||
--text "नमस्ते। कैसे हो?" \
|
||||
--split-granularity sentence
|
||||
|
||||
# Receive JSON sidecar chunks with word-level timestamps
|
||||
python streaming_speech_client.py \
|
||||
--text "Hello world. How are you?" \
|
||||
@@ -78,9 +73,8 @@ except ImportError:
|
||||
def frame_basename(msg: dict) -> str:
|
||||
"""Name a file after the utterance and sentence a frame belongs to.
|
||||
|
||||
Every utterance uses `sentence_index` for units inside that flush.
|
||||
With default `split_granularity=none` that index stays 0; sentence mode
|
||||
uses 0, 1, ... so files from one connection do not overwrite each other.
|
||||
Every utterance is sentence 0, so the utterance index is what keeps the
|
||||
files of one connection from overwriting each other.
|
||||
"""
|
||||
return f"utterance_{msg['utterance_index']:03d}_sentence_{msg['sentence_index']:03d}"
|
||||
|
||||
@@ -374,13 +368,6 @@ def main():
|
||||
)
|
||||
parser.add_argument("--speed", type=float, default=1.0, help="Playback speed (0.25-4.0)")
|
||||
parser.add_argument("--max-new-tokens", type=int, default=None, help="Max tokens")
|
||||
parser.add_argument("--seed", type=int, default=None, help="Sampling seed forwarded to the server")
|
||||
parser.add_argument(
|
||||
"--split-granularity",
|
||||
default=None,
|
||||
choices=["none", "sentence", "clause"],
|
||||
help="Text split mode (default: server none = one request per input.done)",
|
||||
)
|
||||
|
||||
# Base task options
|
||||
parser.add_argument("--ref-audio", default=None, help="Reference audio")
|
||||
@@ -422,8 +409,6 @@ def main():
|
||||
"max_new_tokens",
|
||||
"ref_audio",
|
||||
"ref_text",
|
||||
"seed",
|
||||
"split_granularity",
|
||||
]:
|
||||
val = getattr(args, key.replace("-", "_"), None)
|
||||
if val is not None:
|
||||
|
||||
+5
-7
@@ -53,9 +53,7 @@ dev = [
|
||||
"soundfile==0.14.0",
|
||||
"imageio[ffmpeg]==2.37.4",
|
||||
"opencv-python==5.0.0.93",
|
||||
# 0.3.12 includes acknowledged TCP writes (Mooncake #2850): completion
|
||||
# must mean the receive buffer is ready before zero-copy consumers read it.
|
||||
"mooncake-transfer-engine-cuda13==0.3.12.post1",
|
||||
"mooncake-transfer-engine-cuda13==0.3.11.post1",
|
||||
"av==18.0.0", # for ComfyUI tests
|
||||
"openpyxl==3.1.5", # for nightly CI
|
||||
"pyttsx3==2.99",
|
||||
@@ -77,6 +75,7 @@ dev = [
|
||||
# Provides top-level cosyvoice2; do not co-install stepaudio2-minicpmo,
|
||||
# whose top-level stepaudio2 package conflicts with this distribution.
|
||||
"step-audio2==1.0.0", # Step-Audio2 / MiniCPM-o 4.5 Token2Wav (flashcosyvoice / hyperpyyaml)
|
||||
"silero-vad==6.1", # Realtime server-VAD duplex E2E
|
||||
]
|
||||
|
||||
demo = [
|
||||
@@ -86,6 +85,9 @@ demo = [
|
||||
"cn2an>=0.5.22", # Chinese number normalization for TTS (e.g. "1次" -> "一次")
|
||||
]
|
||||
|
||||
server-vad = [
|
||||
"silero-vad==6.1",
|
||||
]
|
||||
|
||||
forced-aligner = [
|
||||
"qwen-asr>=0.0.6", # faithful Japanese/Korean word segmentation for the TTS forced aligner
|
||||
@@ -374,7 +376,3 @@ ba = "ba"
|
||||
SER = "SER"
|
||||
# NumPy ndarray.flags.writeable (canonical flag name, not "writable").
|
||||
writeable = "writeable"
|
||||
# "convnet" abbreviates convolutional network. It appears in the vendored Fish
|
||||
# Speech DAC codec identifiers (convnet_type, transconvnet_type) and is not a
|
||||
# typo of "convent".
|
||||
convnet = "convnet"
|
||||
@@ -1,13 +1,12 @@
|
||||
# Boogu-Image
|
||||
|
||||
> Text-to-image and image-editing online serving
|
||||
(Boogu-Image-0.1-Base / -Turbo / -Edit / -Edit-Turbo)
|
||||
(Boogu-Image-0.1-Base / -Edit / -Edit-Turbo)
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: Boogu
|
||||
- Model: `Boogu/Boogu-Image-0.1-Base` (text-to-image),
|
||||
`Boogu/Boogu-Image-0.1-Turbo` (four-step text-to-image),
|
||||
`Boogu/Boogu-Image-0.1-Edit` (image editing), and
|
||||
`Boogu/Boogu-Image-0.1-Edit-Turbo` (four-step image editing)
|
||||
- Task: Text-to-image generation and text-guided image editing (TI2I)
|
||||
@@ -17,9 +16,8 @@
|
||||
## When to use this recipe
|
||||
|
||||
Use this recipe when you want a known-good starting point for serving
|
||||
`Boogu/Boogu-Image-0.1-Base`, `Boogu/Boogu-Image-0.1-Turbo`,
|
||||
`Boogu/Boogu-Image-0.1-Edit`, or `Boogu/Boogu-Image-0.1-Edit-Turbo` with
|
||||
vLLM-Omni's native pipeline (no
|
||||
`Boogu/Boogu-Image-0.1-Base`, `Boogu/Boogu-Image-0.1-Edit`, or
|
||||
`Boogu/Boogu-Image-0.1-Edit-Turbo` with vLLM-Omni's native pipeline (no
|
||||
`--diffusion-load-format diffusers`, and the upstream `boogu` package is not
|
||||
required).
|
||||
|
||||
@@ -27,14 +25,12 @@ Boogu-Image-0.1 is an Apache-2.0 unified image generation and editing model
|
||||
family. The Base text-to-image checkpoint pairs a Qwen3-VL multimodal encoder
|
||||
with a Diffusion Transformer (DiT) and a flow-match Euler scheduler with
|
||||
time-shift. It handles photorealistic generation and Chinese/English text
|
||||
rendering. The Edit checkpoint uses the native `BooguImagePipeline`, while the
|
||||
distilled Turbo and Edit-Turbo checkpoints run the four-step DMD path through
|
||||
`BooguImageTurboPipeline`.
|
||||
rendering. The Edit and Edit-Turbo checkpoints use the same native
|
||||
`BooguImagePipeline` for image editing.
|
||||
|
||||
## References
|
||||
|
||||
- Upstream model card: <https://huggingface.co/Boogu/Boogu-Image-0.1-Base>
|
||||
- Turbo model card: <https://huggingface.co/Boogu/Boogu-Image-0.1-Turbo>
|
||||
- Edit model card: <https://huggingface.co/Boogu/Boogu-Image-0.1-Edit>
|
||||
- Edit-Turbo 1K hotfix: <https://huggingface.co/Boogu/Boogu-Image-0.1-Edit-Turbo/tree/hotfix-1k-20260708>
|
||||
- Project page: <https://boogu.org>
|
||||
@@ -190,19 +186,6 @@ ID passed to `vllm serve`. For example, use
|
||||
corresponding Edit model.
|
||||
|
||||
Base text-to-image:
|
||||
## Fast text-to-image (Boogu-Image-0.1-Turbo)
|
||||
|
||||
Turbo is the distilled text-to-image checkpoint. Its `model_index.json`
|
||||
declares `BooguImageTurboPipeline`, which runs the four-step DMD path instead
|
||||
of the scheduler-driven path used by Base.
|
||||
|
||||
### Command
|
||||
|
||||
```bash
|
||||
vllm serve Boogu/Boogu-Image-0.1-Turbo --omni --port 8091
|
||||
```
|
||||
|
||||
### Verification
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/images/generations \
|
||||
@@ -247,23 +230,6 @@ curl -X POST http://localhost:8091/v1/images/edits \
|
||||
checkpoint-declared FP8 configuration. With the regular checkpoints used for
|
||||
online FP8, the MLLM remains in BF16. The VAE and scheduler are outside this
|
||||
quantization routing.
|
||||
"model": "Boogu/Boogu-Image-0.1-Turbo",
|
||||
"prompt": "A mountain lake at sunset, photorealistic, cinematic lighting",
|
||||
"size": "1024x1024",
|
||||
"seed": 42
|
||||
}' | jq -r '.data[0].b64_json' | base64 -d > output-turbo.png
|
||||
```
|
||||
|
||||
### Notes
|
||||
|
||||
- **Required settings:** `guidance_scale=1.0` and `guidance_scale_2=1.0`; any
|
||||
other value is rejected. `num_inference_steps` defaults to 4. All three are
|
||||
pipeline defaults, so the request above omits them.
|
||||
- **Resolution:** upstream constrains Turbo to 1K, so `1024x1024` is the
|
||||
supported working resolution. The shared pipeline still clamps at 2K; larger
|
||||
sizes are not validated for this checkpoint.
|
||||
- **Known limitations:** the same single-GPU limitations as Base apply; CPU
|
||||
offload, Cache-DiT, and multi-GPU parallelism are not yet validated.
|
||||
|
||||
## Image editing (Boogu-Image-0.1-Edit)
|
||||
|
||||
@@ -388,18 +354,16 @@ curl -s http://localhost:8091/v1/chat/completions \
|
||||
## Fast image editing (Boogu-Image-0.1-Edit-Turbo)
|
||||
|
||||
Edit-Turbo is the distilled image-editing checkpoint. Its `model_index.json`
|
||||
declares `BooguImagePipeline`, so `--model-class-name BooguImageTurboPipeline`
|
||||
is required to reach the four-step DMD path; without it the checkpoint is
|
||||
served on the dense Edit path. The upstream 1K hotfix is recommended for
|
||||
stable results and is pinned below to avoid silently loading the older
|
||||
checkpoint from the repository's default revision.
|
||||
declares `BooguImagePipeline`, so it uses the same native TI2I path as Edit.
|
||||
The upstream 1K hotfix is recommended for stable results and is pinned below
|
||||
to avoid silently loading the older checkpoint from the repository's default
|
||||
revision.
|
||||
|
||||
### Command
|
||||
|
||||
```bash
|
||||
vllm serve Boogu/Boogu-Image-0.1-Edit-Turbo \
|
||||
--omni \
|
||||
--model-class-name BooguImageTurboPipeline \
|
||||
--revision hotfix-1k-20260708 \
|
||||
--port 8091
|
||||
```
|
||||
@@ -443,10 +407,11 @@ curl -s http://localhost:8091/v1/chat/completions \
|
||||
|
||||
- **Pinned revision:** use `hotfix-1k-20260708`; upstream recommends the 1K
|
||||
hotfix over the 1.5K variant for more stable results.
|
||||
- **Required settings:** `guidance_scale=1.0` and `guidance_scale_2=1.0`; the
|
||||
Turbo pipeline rejects any other value, so the `guidance_scale_2=2.0` shown
|
||||
for the regular Edit checkpoint does not apply here. `num_inference_steps`
|
||||
defaults to 4.
|
||||
- **Recommended settings:** `num_inference_steps=4` and
|
||||
`guidance_scale=1.0`. Edit-Turbo is guidance-distilled, meaning the guidance
|
||||
behavior is already baked into the distilled checkpoint. A scale of `1.0`
|
||||
avoids applying additional classifier-free guidance; unlike the regular Edit
|
||||
checkpoint, it should not be increased to `5.0`.
|
||||
- **Reference images:** the same single-reference and `align_res` behavior as
|
||||
the regular Edit checkpoint applies.
|
||||
- **Known limitations:** the same single-GPU limitations as Base and Edit apply;
|
||||
|
||||
@@ -219,51 +219,3 @@ the audio-track caveat.
|
||||
5 chunks, `key_frames=0` = summarize all chunk frames) are aligned to the JoyVL
|
||||
reference adapter so per-tick behavior matches the released model; the framework only
|
||||
supplies the serving structure.
|
||||
|
||||
## Native multi-stage pipeline (experimental)
|
||||
|
||||
An experimental native speech path is available alongside the existing Day-0 serving path:
|
||||
|
||||
```text
|
||||
one request with image/video and text
|
||||
-> JoyAI complete action text
|
||||
-> Qwen3-TTS Talker
|
||||
-> Code2Wav
|
||||
-> 24 kHz audio
|
||||
```
|
||||
|
||||
It runs behind one `vllm serve --omni` endpoint:
|
||||
|
||||
```bash
|
||||
vllm serve jdopensource/JoyAI-VL-Interaction-Preview --omni \
|
||||
--deploy-config vllm_omni/deploy/joyai_vl_interaction.yaml \
|
||||
--port 8091
|
||||
```
|
||||
|
||||
With `modalities: ["text", "audio"]`, its behavior is:
|
||||
|
||||
| JoyAI action | Native output |
|
||||
|---|---|
|
||||
| `</response> <text>` | action text + speech for `<text>` |
|
||||
| `</silence>` | action text + empty audio output (zero samples) |
|
||||
| `</response> <note> </delegation> <question>` | action text + speech for `<note>` only |
|
||||
|
||||
Native speech supports the Qwen3-TTS `CustomVoice` task. The request-level `voice` and
|
||||
`language` fields select the speaker and language; when omitted, they default to
|
||||
`Vivian` and `Auto`. Values supplied as `tts_speaker` and `tts_language` inside
|
||||
`additional_information` take priority.
|
||||
|
||||
This experimental native multi-stage pipeline is request-scoped, stateless, and
|
||||
all-sync. It currently does not include:
|
||||
|
||||
- continuous frame sessions, persistent standing instructions, or session memory;
|
||||
- audio input / ASR, full-duplex interaction, or barge-in;
|
||||
- background Agent execution for delegated questions;
|
||||
- async-chunk support from Talker to Code2Wav;
|
||||
- integration with the existing WebUI.
|
||||
|
||||
For actions with spoken text, the Talker starts only after the complete JoyAI action is
|
||||
available, and Code2Wav starts only after the Talker completes. Silence actions skip both
|
||||
TTS stages and return an empty audio output with zero samples. All three stages are placed
|
||||
on GPU 0 by the provided
|
||||
[`deploy config`](../../vllm_omni/deploy/joyai_vl_interaction.yaml).
|
||||
@@ -110,29 +110,6 @@ vllm serve Lightricks/LTX-2.5-Diffusers \
|
||||
|
||||
Use the T2V or I2V request below after the server is ready.
|
||||
|
||||
### Multi-GPU low-latency path
|
||||
|
||||
For low-latency 1920x1088 distilled two-stage serving, match the Ulysses and
|
||||
VAE patch-parallel degrees to the GPU count. For example, on four GPUs:
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=0,1,2,3 \
|
||||
vllm serve Lightricks/LTX-2.5-Diffusers \
|
||||
--omni \
|
||||
--model-class-name LTX2DistilledTwoStagePipeline \
|
||||
--usp 4 \
|
||||
--vae-patch-parallel-size 4 \
|
||||
--vae-use-tiling \
|
||||
--host 0.0.0.0 \
|
||||
--port 8000 \
|
||||
--stage-init-timeout 900
|
||||
```
|
||||
|
||||
Strict Ulysses only: `((F - 1) / 8 + 1) * (H / 32) * (W / 32)` and TP-local
|
||||
SP attention head counts must be divisible by `ulysses_degree`. Use each
|
||||
phase's `H, W` (half/full resolution for two-stage); non-divisible shapes fail.
|
||||
Each GPU must still fit the resident weights. Warm up with the production shape.
|
||||
|
||||
## Offline inference
|
||||
|
||||
Choose values from the pipeline table. For example, the distilled two-stage
|
||||
|
||||
@@ -101,12 +101,6 @@ sigmas and input latents.
|
||||
Ordinary two-stage uses layer-fused LoRA for an unquantized BF16 Transformer
|
||||
and automatically switches to dynamic LoRA when quantization is enabled.
|
||||
|
||||
## Sequence Parallelism
|
||||
|
||||
Strict Ulysses only: `((F - 1) / 8 + 1) * (H / 32) * (W / 32)` and TP-local
|
||||
SP attention head counts must be divisible by `ulysses_degree`. Use each
|
||||
phase's `H, W` (half/full resolution for two-stage); non-divisible shapes fail.
|
||||
|
||||
## Serving
|
||||
|
||||
Start either one-stage checkpoint:
|
||||
@@ -341,8 +335,9 @@ bundled offline CLI do not currently expose `sigmas`.
|
||||
`--max-num-seqs 1`.
|
||||
- `--cfg-parallel-size` shards the complete LTX guidance plan, including STG,
|
||||
modality guidance, and rescale-compatible prediction gathering.
|
||||
- Ulysses masks padded audio latents; video sequences must satisfy the
|
||||
[sequence-parallel divisibility constraints](#sequence-parallelism).
|
||||
- Sequence parallelism may pad audio latents. Pure Ulysses masks the padding;
|
||||
Ring cannot, so audio length must be SP-divisible. Use `ring_degree=1` or a
|
||||
divisible request shape.
|
||||
- Cache-DiT is one-stage only; multi-stage configurations are rejected.
|
||||
|
||||
## Operational Notes
|
||||
|
||||
@@ -68,24 +68,16 @@ The `/v1/videos` request schema and `extra_params.task` values (`t2va`,
|
||||
|
||||
## Turbo LoRA
|
||||
|
||||
The deploy config below carries the four-step 768p contract -- five sigma
|
||||
points, `flow_shift=6`, `audio_flow_shift=3` -- so requests that omit sampling
|
||||
controls do not inherit the 50-step base schedule:
|
||||
MiniMax-H3 Turbo uses five sigma points, `flow_shift=6`, and
|
||||
`audio_flow_shift=3`. Start Turbo deployments with the dedicated defaults so
|
||||
requests that omit sampling controls do not inherit the 50-step base schedule:
|
||||
|
||||
```bash
|
||||
vllm-omni serve MiniMaxAI/MiniMax-H3 \
|
||||
--omni \
|
||||
--lora-path /path/to/minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors \
|
||||
--lora-path /path/to/MiniMax-H3-Turbo \
|
||||
--deploy-config vllm_omni/deploy/minimax_h3_disaggregated_turbo.yaml
|
||||
```
|
||||
|
||||
`--lora-path` must name one artifact, or a directory holding exactly one; the
|
||||
Turbo repository holds several and a directory of them is rejected as
|
||||
ambiguous. Serving a different artifact means overriding those defaults with
|
||||
that artifact's own contract -- see the
|
||||
[Turbo LoRA table](MiniMax-H3.md#turbo-lora) for every published row.
|
||||
|
||||
The base deployment intentionally retains 50 inference steps for non-LoRA
|
||||
quality. This disaggregated topology is validated on the `fl2v` artifacts,
|
||||
which serve T2VA and FL2VA; the `ref2v` artifacts need a `--task-type ref2va`
|
||||
server.
|
||||
quality. Turbo LoRA supports T2VA and FL2VA requests, not Ref2VA.
|
||||
@@ -869,51 +869,13 @@ balanced switch order.
|
||||
|
||||
### Turbo LoRA
|
||||
|
||||
The eight Diffusers-layout LightX2V Turbo artifacts are supported. The
|
||||
filename records the contract, so the server reads the step count, task family
|
||||
and flow shift from it and validates each request against the artifact that is
|
||||
loaded. It does not rewrite request or deploy-config sampling values: the
|
||||
request must carry that artifact's own settings, listed here, or it is
|
||||
rejected.
|
||||
Only the native Diffusers 4-step FL2VA/T2VA v1.0 artifact is supported:
|
||||
|
||||
| Artifact | Task | Forwards | `num_inference_steps` | `flow_shift` | declared `alpha` |
|
||||
| --- | --- | ---: | ---: | ---: | ---: |
|
||||
| `minimax_h3_fl2v_turbo_4step_v0.1.safetensors` | T2VA / FL2VA | 4 | 5 | 12 | none -> 8 |
|
||||
| `minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors` | T2VA / FL2VA | 4 | 5 | 6 | 128 |
|
||||
| `minimax_h3_fl2v_turbo_4step_v1.1_768p_bf16.safetensors` | T2VA / FL2VA | 4 | 5 | 6 | 128 |
|
||||
| `minimax_h3_fl2v_turbo_4step_v1.2_768p_bf16.safetensors` | T2VA / FL2VA | 4 | 5 | 6 | 8 |
|
||||
| `minimax_h3_fl2v_turbo_8step_v1.0_bf16.safetensors` | T2VA / FL2VA | 8 | 9 | 12 | 8 |
|
||||
| `minimax_h3_fl2v_turbo_8step_v1.0_768p_bf16.safetensors` | T2VA / FL2VA | 8 | 9 | 6 | 8 |
|
||||
| `minimax_h3_ref2v_turbo_4step_v0.1_bf16.safetensors` | Ref2VA | 4 | 5 | 12 | 8 |
|
||||
| `minimax_h3_ref2v_turbo_8step_v1.0_768p_bf16.safetensors` | Ref2VA | 8 | 9 | 6 | 8 |
|
||||
```text
|
||||
lightx2v/Minimax-h3-Turbo/minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors
|
||||
```
|
||||
|
||||
`audio_flow_shift` is `3.0` across the family. Each row is the complete
|
||||
published filename; use it verbatim as `TURBO_FILE` below.
|
||||
|
||||
Alpha needs no manual compensation: the server reads it from the artifact's
|
||||
metadata, falling back to 8 with a warning for
|
||||
`minimax_h3_fl2v_turbo_4step_v0.1`, the one artifact that declares none. The
|
||||
request-level `scale` is a further multiplier on top of it.
|
||||
|
||||
> [!NOTE]
|
||||
> Every artifact is rank 128 and the delta is applied at `scale * alpha /
|
||||
> rank`, so the `alpha=8` rows drive at 1/16 the strength of the `alpha=128`
|
||||
> rows. 8 is the default of LightX2V's reference script, which never reads the
|
||||
> metadata; its documented `v0.1` command does not override that default.
|
||||
|
||||
Every artifact except `minimax_h3_fl2v_turbo_4step_v0.1` also ships a
|
||||
`_comfyui_` export of the same weights. Those fuse Q/K/V into one projection
|
||||
and are **not** supported; downloading one is refused by name. Take the
|
||||
Diffusers file. The filename is the contract, so do not rename an artifact
|
||||
either -- a renamed file is rejected rather than served on a guess.
|
||||
|
||||
FL2VA artifacts serve `t2va` and `fl2va` on any FL2VA or combined server.
|
||||
Ref2VA artifacts require
|
||||
`--task-type ref2va`: a combined server serves `ref2va` from a second DiT that
|
||||
the adapter cannot bind to, so loading one there is refused rather than silently
|
||||
running an undistilled model on the few-step schedule.
|
||||
|
||||
Download the artifact you want:
|
||||
Download only that file:
|
||||
|
||||
```bash
|
||||
export TURBO_DIR=/path/to/minimax-h3-turbo
|
||||
@@ -922,19 +884,10 @@ hf download lightx2v/Minimax-h3-Turbo "${TURBO_FILE}" --local-dir "${TURBO_DIR}"
|
||||
export TURBO_LORA="${TURBO_DIR}/${TURBO_FILE}"
|
||||
```
|
||||
|
||||
`--lora-path` accepts one artifact, or a directory holding exactly one.
|
||||
|
||||
> [!IMPORTANT]
|
||||
> This changes earlier behaviour. `--lora-path /path/to/minimax-h3-turbo`
|
||||
> pointing at a full clone of the Turbo repository used to select the v1.0 768p
|
||||
> file implicitly; a directory holding several recognized artifacts is now
|
||||
> rejected as ambiguous. Name the artifact instead:
|
||||
> `--lora-path /path/to/minimax-h3-turbo/minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors`.
|
||||
|
||||
Start from a non-offloaded or DLO FL2VA server command and add
|
||||
`--task-type fl2va --lora-backend peft --lora-path "${TURBO_LORA}"`.
|
||||
`--lora-path` preloads the adapter; each request still activates it and must
|
||||
carry that artifact's sampling settings:
|
||||
`--lora-path` preloads the adapter; each request still activates it and uses
|
||||
the published sampling settings:
|
||||
|
||||
```bash
|
||||
-F 'num_inference_steps=5' \
|
||||
@@ -943,25 +896,14 @@ carry that artifact's sampling settings:
|
||||
-F "lora={\"name\":\"h3-turbo-v1.0\",\"path\":\"${TURBO_LORA}\",\"scale\":1.0}"
|
||||
```
|
||||
|
||||
Switching to another FL2VA artifact means repointing `TURBO_FILE`, which moves
|
||||
both `--lora-path` and the request's `lora.path`, and carrying that row's
|
||||
`num_inference_steps` and `flow_shift`: `9` and `6` for `8step_v1.0_768p`, `9`
|
||||
and `12` for the 544p `8step_v1.0`. A request that does not match the loaded
|
||||
artifact is rejected, so a mismatch cannot silently degrade output.
|
||||
|
||||
The two `ref2v` rows are not served by this FL2VA command. Start a
|
||||
`--task-type ref2va` server, take a request from
|
||||
[Ref2VA](#3-ref2va-image-only-imageaudio-or-mixed-references) and override
|
||||
`num_inference_steps` and `flow_shift` with that row's values; those examples
|
||||
already send `audio_flow_shift=3.0`.
|
||||
|
||||
For FL2VA, change `task` and add `input_reference` as shown above. This
|
||||
integration is dynamic-only and does not support prefusion or LoRA composition. DLO is
|
||||
For FL2VA, change `task` and add `input_reference` as shown above. The 8-step,
|
||||
ComfyUI, Ref2VA, and v1.1 artifacts are not supported. This integration is
|
||||
dynamic-only and does not support prefusion or LoRA composition. DLO is
|
||||
supported by keeping the request-switchable LoRA A/B buffers resident on the
|
||||
accelerator while DLO streams only the base blocks; budget for this additional
|
||||
fixed HBM usage. Model-level and standard layerwise offload remain unsupported.
|
||||
The requested sigma points always number one more than the artifact's denoiser
|
||||
evaluations.
|
||||
The five requested sigma points produce the four denoiser evaluations expected
|
||||
by the Turbo artifact.
|
||||
|
||||
### FlashGen native LoRA
|
||||
|
||||
|
||||
@@ -16,7 +16,7 @@ Use this recipe for 24 kHz multilingual TTS with voice cloning (20 languages inc
|
||||
Chinese and English). Choose a variant based on your latency and quality requirements:
|
||||
|
||||
| Model | Params | Use case |
|
||||
| --- | --- | --- |
|
||||
|---|---|---|
|
||||
| MOSS-TTS | 8B | General TTS, highest quality |
|
||||
| MOSS-TTS-v1.5 | 8B | General TTS upgrade of 1.0: 31 languages, steadier cloning, `[pause Xs]` markers (set `language` for best results); same `MossTTSDelay` API |
|
||||
| MOSS-TTS-Realtime | 1.7B | Lowest latency (TTFB ~180 ms), streaming-first |
|
||||
@@ -24,17 +24,9 @@ Chinese and English). Choose a variant based on your latency and quality require
|
||||
| MOSS-SoundEffect | 8B | Sound effect synthesis from text description |
|
||||
| MOSS-VoiceGenerator | 1.7B | Zero-shot voice design |
|
||||
|
||||
The variants above share the same codec (`OpenMOSS-Team/MOSS-Audio-Tokenizer`, ~7 GB) and
|
||||
All variants share the same codec (`OpenMOSS-Team/MOSS-Audio-Tokenizer`, ~7 GB) and
|
||||
output 24 kHz mono audio.
|
||||
|
||||
MOSS-TTS-Local-Transformer-v1.5 uses MOSS-Audio-Tokenizer-v2 and outputs 48 kHz
|
||||
stereo audio. For Local voice cloning through `/v1/audio/speech`, provide an
|
||||
accurate `ref_text` transcript alongside `ref_audio`. The adapter uses the
|
||||
reference transcript and audio as a continuation prefix, then generates only
|
||||
the requested `input` speech. Without a nonblank `ref_text`, Local uses
|
||||
audio-reference generation. Reference transcripts must match the reference
|
||||
audio; they are not style instructions.
|
||||
|
||||
## References
|
||||
|
||||
- Offline inference example: [`examples/offline_inference/text_to_speech/moss_tts/`](../../examples/offline_inference/text_to_speech/moss_tts/)
|
||||
|
||||
@@ -85,7 +85,6 @@ recipes/
|
||||
| [`Tencent/Covo-Audio-Chat.md`](./Tencent/Covo-Audio-Chat.md) | Online serving for audio chat | 1x A100 80GB |
|
||||
| [`Tencent/HunyuanImage-3.0-Instruct.md`](./Tencent/HunyuanImage-3.0-Instruct.md) | DiT-only text-to-image serving and benchmark, including ModelOpt mixed FP8/NVFP4 | 4x H100/H800 80GB / 2x B200 |
|
||||
| [`Wan-AI/Wan2.2-T2V.md`](./Wan-AI/Wan2.2-T2V.md) | Text-to-video serving (Wan2.2 14B) with Skip-Softmax sparse attention | 1x datacenter Blackwell (B200/B300) |
|
||||
| [`Wan-AI/Wan2.2-TI2V-5B-A100.md`](./Wan-AI/Wan2.2-TI2V-5B-A100.md) | Unified 720p text-to-video and image-to-video serving (Wan2.2 5B) | 1x A100-SXM4-80GB |
|
||||
| [`Wan-AI/Wan2.2-I2V.md`](./Wan-AI/Wan2.2-I2V.md) | Image-to-video serving (Wan2.2 14B) | 8x Ascend NPU (A2/A3) |
|
||||
| [`Wan-AI/Wan2.2-S2V.md`](./Wan-AI/Wan2.2-S2V.md) | Speech-to-video serving (Wan2.2 14B) | 2x A100/H100 80GB |
|
||||
| [`Wan-AI/Wan2.1-VACE.md`](./Wan-AI/Wan2.1-VACE.md) | Unified T2V, I2V, V2LF, FLF2V, inpaint, and R2V | 1x RTX 5090 (1.3B) / 1x L40S 48GB with layerwise offload (14B) |
|
||||
|
||||
@@ -7,31 +7,13 @@
|
||||
- Vendor: Robbyant
|
||||
- Model: `robbyant/lingbot-world-v2-14b-causal-fast-diffusers`
|
||||
- Task: image-conditioned interactive world generation
|
||||
- Modes: offline trajectory replay, stepwise one-request streaming, and in-process realtime AR-Diffusion ticks (deprecated)
|
||||
- Modes: offline trajectory replay and in-process realtime AR-Diffusion ticks
|
||||
- Hardware validated: NVIDIA H200 and B200
|
||||
- Maintainer: Community
|
||||
|
||||
The checkpoint is separately licensed under CC BY-NC-SA and restricted to
|
||||
non-commercial use. The vLLM-Omni integration code remains Apache-2.0.
|
||||
|
||||
## Which path to use
|
||||
|
||||
There are three ways to drive this model, each documented in its own section
|
||||
below.
|
||||
|
||||
- **Realtime stepwise** — see [Streaming video serving](#streaming-video-serving).
|
||||
The suggested path: one `WS /v1/realtime/video` session produces the whole
|
||||
rollout, one video chunk per AR block. It cannot take realtime camera
|
||||
interaction yet; that is being added, and this path is the one that will
|
||||
remain.
|
||||
- **Realtime tick** *(deprecated)* — see
|
||||
[Realtime in-process generation](#realtime-in-process-generation-deprecated).
|
||||
The only path that takes realtime camera interaction today, so it stays until
|
||||
the stepwise path can, but it takes no new work and will be removed.
|
||||
- **Offline** — see [Offline generation](#offline-generation). Replays a fixed
|
||||
pose/intrinsics trajectory in one request and writes an MP4. Use it when the
|
||||
camera path is known up front and streaming is not needed.
|
||||
|
||||
## Offline generation
|
||||
|
||||
The offline path consumes one source image and an action directory containing
|
||||
@@ -52,12 +34,7 @@ counts must be `9 + 12k`, up to 117 frames. Camera files may contain a
|
||||
longer bounded source trajectory (the official example contains 269 frames);
|
||||
the runtime consumes only the prefix needed by the request.
|
||||
|
||||
## Realtime in-process generation (deprecated)
|
||||
|
||||
Use [Streaming video serving](#streaming-video-serving) instead unless you need
|
||||
mid-session camera interaction, which only this path has. It still works and
|
||||
emits no runtime warning; removal is tracked as B4 of the LingBot World 2.0
|
||||
roadmap ([#6672](https://github.com/vllm-project/vllm-omni/issues/6672)).
|
||||
## Realtime in-process generation
|
||||
|
||||
The realtime example keeps AR-Diffusion KV and model-owned state across
|
||||
requests. Each JSONL line describes the prompt and/or three latent-frame camera
|
||||
@@ -84,81 +61,12 @@ exercises the same `ARDiffusionSessionManager -> ARDiffusionOmniTickConsumer
|
||||
-> AsyncOmni -> ARDiffusionEngine` path used by a future HTTP or WebSocket
|
||||
transport.
|
||||
|
||||
Realtime serving now goes through the generic `WS /v1/realtime/video`
|
||||
transport, which exposes no LingBot-specific event fields; the tick loop above
|
||||
stays available only for the older control plane.
|
||||
|
||||
## Streaming video serving
|
||||
|
||||
The stepwise path keeps AR-Diffusion paged KV but issues **one** request for
|
||||
the whole rollout: `prepare_encode` runs once, then every AR block is four DMD
|
||||
steps followed by one streamed chunk. Identity metadata uses
|
||||
`session_id = request_id` with contiguous `chunk_index` values from zero.
|
||||
Camera motion for this landing is request-scoped: one three-frame action list
|
||||
per chunk, fixed when the rollout starts.
|
||||
|
||||
Serve it with the AR-Diffusion deploy config, which selects the AR-Diffusion
|
||||
engine and enables streamed step execution:
|
||||
|
||||
```bash
|
||||
vllm serve robbyant/lingbot-world-v2-14b-causal-fast-diffusers \
|
||||
--omni \
|
||||
--deploy-config vllm_omni/deploy/lingbot_world_v2_stepwise.yaml \
|
||||
--port 8000
|
||||
```
|
||||
|
||||
`--deploy-config` is required: the `lingbot_world` pipeline deliberately
|
||||
registers no default deploy config, so offline replay and the tick example
|
||||
keep their request-mode topology when no deploy config is given, and the
|
||||
stepwise serving topology is only ever an explicit choice.
|
||||
|
||||
Clients then use the generic WebSocket protocol documented in
|
||||
[`docs/serving/video_stream_api.md`](../../docs/serving/video_stream_api.md):
|
||||
`session.start` begins one rollout and each AR block arrives as a binary video
|
||||
chunk. The model is image-conditioned, so the first frame is required and rides
|
||||
on `image_reference` as an `http(s)` or `data:` URL; per-chunk camera actions
|
||||
ride on `extra_params`:
|
||||
|
||||
```json
|
||||
{"type": "session.start", "model": "robbyant/lingbot-world-v2-14b-causal-fast-diffusers",
|
||||
"prompt": "The camera moves slowly forward through the scene.",
|
||||
"image_reference": {"image_url": "data:image/png;base64,<first frame>"},
|
||||
"width": 832, "height": 480, "num_frames": 33,
|
||||
"extra_params": {"flow_shift": 5.0,
|
||||
"camera_action_script": [[["w"], ["w"], ["w"]], [["a"], [], []], [[], [], []]]}}
|
||||
```
|
||||
|
||||
`camera_action_script` must hold exactly one action list per generated
|
||||
chunk, and a request generates `((num_frames - 1) // 4 + 1) // 3` chunks:
|
||||
three for `num_frames: 33`, seven for `num_frames: 81`.
|
||||
|
||||
The bundled client speaks this protocol:
|
||||
|
||||
```bash
|
||||
python examples/online_serving/streaming_video_generation/streaming_video_client.py \
|
||||
--model robbyant/lingbot-world-v2-14b-causal-fast-diffusers \
|
||||
--prompt "The camera moves slowly forward through the scene." \
|
||||
--image-reference /path/to/first_frame.png \
|
||||
--width 832 --height 480 --num-frames 33 --fps 16 --seed 42 \
|
||||
--extra-params '{"camera_action_script":[[["w"],["w"],["w"]],[["a"],[],[]],[[],[],[]]]}' \
|
||||
--output lingbot_world_v2_stream.mp4
|
||||
```
|
||||
|
||||
A served request may instead point at a pose/intrinsics trajectory with
|
||||
`extra_params.action_path`, which is resolved inside the trusted root set by
|
||||
`model_config.lingbot_action_root` or `VLLM_OMNI_LINGBOT_ACTION_ROOT`; a server
|
||||
started without that root configured accepts only `camera_action_script`.
|
||||
|
||||
Requested `width`/`height` must match `ar_diffusion_width`/`ar_diffusion_height`
|
||||
in the deploy config, because the AR cache geometry is fixed at load time.
|
||||
Blocks are decoded independently, so seams between chunks are possible; a
|
||||
session-owned streaming decoder is tracked separately.
|
||||
|
||||
Mid-session `session.interaction` for camera control is not wired yet, so a
|
||||
served rollout follows the script it started with.
|
||||
|
||||
The deprecated tick example remains available for the older
|
||||
one-block-per-`generate()` control plane.
|
||||
This PR intentionally does not define a public realtime HTTP/WebSocket schema.
|
||||
An online serving client should be added together with the public transport
|
||||
API rather than exposing LingBot-specific event fields from the generic model
|
||||
runtime. The structured camera-interaction frontend and transport work is
|
||||
tracked separately in
|
||||
[vllm-project/vllm-omni#5527](https://github.com/vllm-project/vllm-omni/pull/5527).
|
||||
|
||||
## Realtime identity and controls
|
||||
|
||||
@@ -196,16 +104,10 @@ tested commit.
|
||||
## Current limitations
|
||||
|
||||
- Only the 14B causal-fast checkpoint is supported.
|
||||
- The tick control plane is internal; the public transport is the stepwise
|
||||
`WS /v1/realtime/video` path, which cannot take mid-session camera
|
||||
interaction yet.
|
||||
- Stepwise serving requires an explicit
|
||||
`--deploy-config vllm_omni/deploy/lingbot_world_v2_stepwise.yaml`; there is
|
||||
no default deploy config for this model.
|
||||
- The realtime control plane is internal; there is no public server transport yet.
|
||||
- AR-Diffusion stages currently require one replica because session-affine
|
||||
routing across replicas is not implemented.
|
||||
- Tick mode generates one AR block per request. Stepwise mode generates many
|
||||
AR blocks in one request. `max_num_seqs` must be one in both cases.
|
||||
- One AR block is generated per request and `max_num_seqs` must be one.
|
||||
- Stateful streaming VAE decode is not implemented; the realtime example emits
|
||||
latent chunks.
|
||||
- SP/USP, pipeline/CFG parallelism, HSDP, VAE parallelism, quantization,
|
||||
@@ -217,7 +119,4 @@ tested commit.
|
||||
- Checkpoint: <https://huggingface.co/robbyant/lingbot-world-v2-14b-causal-fast-diffusers>
|
||||
- Official implementation: <https://github.com/robbyant/lingbot-world-v2>
|
||||
- Offline example: [`examples/offline_inference/diffusion/lingbot_world_v2.py`](../../examples/offline_inference/diffusion/lingbot_world_v2.py)
|
||||
- Streaming serving client: [`examples/online_serving/streaming_video_generation/`](../../examples/online_serving/streaming_video_generation/README.md)
|
||||
- Streaming serving end-to-end test: [`tests/e2e/online_serving/test_lingbot_world_v2_stepwise.py`](../../tests/e2e/online_serving/test_lingbot_world_v2_stepwise.py)
|
||||
- Stepwise latent-level end-to-end test: [`tests/e2e/offline_inference/test_lingbot_world_v2_stepwise.py`](../../tests/e2e/offline_inference/test_lingbot_world_v2_stepwise.py)
|
||||
- Realtime design: [`docs/design/feature/realtime_ar_diffusion.md`](../../docs/design/feature/realtime_ar_diffusion.md)
|
||||
@@ -1,184 +0,0 @@
|
||||
# Wan2.2 TI2V 5B — A100 80GB
|
||||
|
||||
> Single-GPU, native 720p text-to-video and image-to-video serving
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: Wan-AI
|
||||
- Model: `Wan-AI/Wan2.2-TI2V-5B-Diffusers`
|
||||
- Task: Unified text-to-video (T2V) and image-to-video (I2V) generation
|
||||
- Mode: Online serving with the OpenAI-compatible Videos API
|
||||
- Hardware: 1x NVIDIA A100-SXM4-80GB
|
||||
- Maintainer: Community
|
||||
|
||||
## When to use this recipe
|
||||
|
||||
Use this recipe to serve both T2V and I2V from one dense 5B checkpoint on one
|
||||
A100 80GB. It qualifies the model's native 720p profile without CPU offload or
|
||||
multi-GPU parallelism. The synchronous endpoint is used below because it
|
||||
returns the MP4 and request timing in one response; the asynchronous Videos API
|
||||
uses the same server and sampling parameters.
|
||||
|
||||
## Supported model contract
|
||||
|
||||
| Task | Required input | Qualified output | Endpoint |
|
||||
| --- | --- | --- | --- |
|
||||
| T2V | Text prompt; omit `input_reference` | 1280x704, 121 frames, 24 fps | `POST /v1/videos/sync` |
|
||||
| I2V | Text prompt plus one image uploaded as `input_reference` | 1280x704, 121 frames, 24 fps | `POST /v1/videos/sync` |
|
||||
|
||||
The upstream checkpoint supports both tasks at 720p and 24 fps. For native
|
||||
landscape 720p, use `1280x704`; portrait generation uses `704x1280`. Wan video
|
||||
lengths follow the temporal VAE contract (`4k+1` frames). This recipe validates
|
||||
one output per request with 50 denoising steps and seed 42.
|
||||
|
||||
## References
|
||||
|
||||
- Upstream model card:
|
||||
<https://huggingface.co/Wan-AI/Wan2.2-TI2V-5B-Diffusers>
|
||||
- Online T2V example:
|
||||
[`examples/online_serving/text_to_video`](../../examples/online_serving/text_to_video)
|
||||
- Online I2V example:
|
||||
[`examples/online_serving/image_to_video`](../../examples/online_serving/image_to_video)
|
||||
- Offline I2V example:
|
||||
[`examples/offline_inference/image_to_video`](../../examples/offline_inference/image_to_video)
|
||||
- Community recipe tracker:
|
||||
[vllm-project/vllm-omni#2645](https://github.com/vllm-project/vllm-omni/issues/2645)
|
||||
|
||||
## Hardware
|
||||
|
||||
- Accelerator model and per-device memory: NVIDIA A100-SXM4-80GB, 81,920 MiB
|
||||
- Number of devices: 1
|
||||
- Device interconnect: not applicable to this single-GPU profile
|
||||
- Host memory: 1 TiB installed on the qualification host; CPU offload was not used
|
||||
- Qualification scope: BF16, batch size 1, native FlashAttention, eager mode,
|
||||
VAE tiling, T2V and I2V at 1280x704 / 121 frames / 24 fps
|
||||
|
||||
## Software environment
|
||||
|
||||
- OS: Ubuntu 24.04.3 LTS, Linux 6.8.0
|
||||
- Python: 3.12.3
|
||||
- NVIDIA driver / host toolkit: 550.127.05 / CUDA 12.8
|
||||
- PyTorch: 2.13.0+cu129
|
||||
- Transformers / Diffusers: 5.14.1 / 0.40.0
|
||||
- vLLM: 0.28.0+cu129
|
||||
- vLLM-Omni: commit `ff6e906a25de1ca3122c0d0d5250fde5a0ddbc7b`
|
||||
|
||||
The default vLLM 0.28.0 wheel targets CUDA 13.0. Driver 550 cannot run that
|
||||
variant, so this qualification used the official CUDA 12.9 wheel:
|
||||
|
||||
```bash
|
||||
uv venv --python 3.12 --seed
|
||||
source .venv/bin/activate
|
||||
uv pip install \
|
||||
'https://github.com/vllm-project/vllm/releases/download/v0.28.0/vllm-0.28.0%2Bcu129-cp38-abi3-manylinux_2_28_x86_64.whl' \
|
||||
--torch-backend=auto
|
||||
uv pip install -e .
|
||||
```
|
||||
|
||||
## Command
|
||||
|
||||
Start the server from the repository root:
|
||||
|
||||
```bash
|
||||
vllm serve Wan-AI/Wan2.2-TI2V-5B-Diffusers \
|
||||
--omni \
|
||||
--port 8092 \
|
||||
--vae-use-tiling \
|
||||
--enforce-eager
|
||||
```
|
||||
|
||||
In another terminal, download the public reference image:
|
||||
|
||||
```bash
|
||||
wget -O cherry_blossom.jpg \
|
||||
https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg
|
||||
```
|
||||
|
||||
Generate a text-only video:
|
||||
|
||||
```bash
|
||||
curl --fail-with-body -D t2v-headers.txt \
|
||||
-X POST http://127.0.0.1:8092/v1/videos/sync \
|
||||
-F "prompt=Two anthropomorphic cats in comfortable boxing gear and bright gloves fight intensely on a spotlighted stage, cinematic lighting, smooth motion." \
|
||||
-F "negative_prompt=blurry, low quality, static, distorted, artifacts, watermark, text" \
|
||||
-F "width=1280" -F "height=704" \
|
||||
-F "num_frames=121" -F "fps=24" \
|
||||
-F "num_inference_steps=50" -F "guidance_scale=5.0" \
|
||||
-F "flow_shift=5.0" -F "seed=42" \
|
||||
--output wan22-ti2v-5b-t2v.mp4
|
||||
```
|
||||
|
||||
Generate an image-conditioned video from the same server:
|
||||
|
||||
```bash
|
||||
curl --fail-with-body -D i2v-headers.txt \
|
||||
-X POST http://127.0.0.1:8092/v1/videos/sync \
|
||||
-F "input_reference=@cherry_blossom.jpg" \
|
||||
-F "prompt=Cherry blossoms sway gently in the breeze while petals drift across the path, cinematic, smooth natural motion." \
|
||||
-F "negative_prompt=blurry, low quality, static, distorted, artifacts, watermark, text" \
|
||||
-F "width=1280" -F "height=704" \
|
||||
-F "num_frames=121" -F "fps=24" \
|
||||
-F "num_inference_steps=50" -F "guidance_scale=5.0" \
|
||||
-F "flow_shift=5.0" -F "seed=42" \
|
||||
--output wan22-ti2v-5b-i2v.mp4
|
||||
```
|
||||
|
||||
## Verification
|
||||
|
||||
Check both MP4 files with `ffprobe`:
|
||||
|
||||
```bash
|
||||
for video in wan22-ti2v-5b-t2v.mp4 wan22-ti2v-5b-i2v.mp4; do
|
||||
ffprobe -v error -select_streams v:0 -count_frames \
|
||||
-show_entries stream=codec_name,width,height,r_frame_rate,nb_read_frames,duration \
|
||||
-of default=noprint_wrappers=1 "$video"
|
||||
done
|
||||
```
|
||||
|
||||
Each file reported:
|
||||
|
||||
```text
|
||||
codec_name=h264
|
||||
width=1280
|
||||
height=704
|
||||
r_frame_rate=24/1
|
||||
duration=5.041667
|
||||
nb_read_frames=121
|
||||
```
|
||||
|
||||
The synchronous response headers and an independent 200 ms `nvidia-smi`
|
||||
sampler produced these steady-state results after server warmup:
|
||||
|
||||
| Task | HTTP | `X-Inference-Time-S` | Server peak | Observed device peak | MP4 size |
|
||||
| --- | ---: | ---: | ---: | ---: | ---: |
|
||||
| T2V | 200 | 288.267 s | 26,546 MB | 27,533 MiB | 7,321,179 bytes |
|
||||
| I2V | 200 | 299.439 s | 29,766 MB | 30,333 MiB | 15,667,146 bytes |
|
||||
|
||||
Both outputs were decoded in full and visually inspected. The T2V output
|
||||
contained the prompted boxing cats. The I2V first frame preserved the uploaded
|
||||
cherry-blossom scene (5.81 mean absolute pixel error after resize and H.264
|
||||
encoding), and later frames showed motion.
|
||||
|
||||
## Notes
|
||||
|
||||
- Memory usage: model loading used 21.2239 GiB. The measured I2V device peak
|
||||
was 30,333 MiB, leaving more than 50 GiB free on this A100.
|
||||
- `--vae-use-tiling` is retained in the known-good command even though this
|
||||
profile has ample headroom; it bounds decode memory at 720p.
|
||||
- `--enforce-eager` removes graph-compilation variability from the measured
|
||||
profile. Remove it only after separately warming and measuring compiled mode.
|
||||
- `flow_shift=5.0` is the 720p setting. The shared 480p examples use 12.0.
|
||||
- The generated MP4 contains video only; this checkpoint does not generate audio.
|
||||
- The first request includes neither model-load time nor the server's startup
|
||||
warmup. Server initialization took approximately 26 seconds with a warm cache.
|
||||
|
||||
## Supported features
|
||||
|
||||
| Feature | Status for this profile | Shared guide |
|
||||
| --- | --- | --- |
|
||||
| OpenAI-compatible Videos API | Qualified for synchronous T2V and I2V | [Videos API](../../docs/serving/videos_api.md) |
|
||||
| VAE tiling | Enabled and qualified on one A100 | [VAE parallelism and tiling](../../docs/user_guide/diffusion/parallelism/vae_parallelism.md) |
|
||||
| Cache-DiT / TaylorSeer | Available but not enabled in the measurements above | [Cache-DiT](../../docs/user_guide/diffusion/cache_acceleration/cache_dit.md) |
|
||||
| TeaCache | Available but not qualified in this recipe | [TeaCache](../../docs/user_guide/diffusion/cache_acceleration/teacache.md) |
|
||||
| CPU or layerwise offload | Not needed for the qualified 80GB profile | [CPU offload](../../docs/user_guide/diffusion/cpu_offload.md) |
|
||||
| TP / SP / CFG parallelism | Not applicable to this single-GPU profile | [Parallelism overview](../../docs/user_guide/diffusion/parallelism/overview.md) |
|
||||
@@ -207,18 +207,6 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync \
|
||||
-F 'extra_params={"depth":{"control_path":"/path/to/depth_control.mp4"},"max_frames":121,"resolution":"720","num_video_frames_per_chunk":121}' \
|
||||
-o cosmos3_transfer_depth.mp4
|
||||
|
||||
# The same transfer control can be uploaded by a remote client instead of
|
||||
# being placed on the server filesystem. Uploads are limited to 512 MiB.
|
||||
curl -sS -X POST http://localhost:8000/v1/videos/sync \
|
||||
-H "Accept: video/mp4" \
|
||||
-F "model=nvidia/Cosmos3-Nano" \
|
||||
-F "prompt=Generate a realistic scene following the provided world-state control." \
|
||||
-F "input_reference=@/path/to/input.mp4;type=video/mp4" \
|
||||
-F "control_reference=@/path/to/wsm.mp4;type=video/mp4" \
|
||||
-F "control_type=wsm" \
|
||||
-F 'extra_params={"wsm":{"control_weight":1.0},"max_frames":121,"resolution":"720","num_video_frames_per_chunk":121}' \
|
||||
-o cosmos3_transfer_wsm.mp4
|
||||
|
||||
# Text-to-video-with-sound
|
||||
curl -sS -X POST http://localhost:8000/v1/videos/sync \
|
||||
-H "Accept: video/mp4" \
|
||||
|
||||
@@ -36,13 +36,19 @@ Other hardware and configurations are welcome as community validation lands.
|
||||
|
||||
- OS: Linux
|
||||
- Python: 3.10+
|
||||
- Required package: `fish-speech` (for DAC codec)
|
||||
- CUDA 12.8
|
||||
- vLLM version: 0.19.0
|
||||
- vLLM-Omni version or commit: c93359bb354a6aa5c14d062430cb85b2c4db251e
|
||||
|
||||
No extra packages are required: the DAC codec architecture is vendored in
|
||||
vLLM-Omni (`vllm_omni/model_executor/models/fish_speech/dac_modules/`), so the
|
||||
model runs out of the box, including in the official Docker image.
|
||||
```bash
|
||||
# Install PortAudio dependency if missing (required by pyaudio which is a dependency of fish-speech)
|
||||
# For Ubuntu/Debian:
|
||||
if ! dpkg -l | grep -q libportaudio2; then
|
||||
sudo apt-get update && sudo apt-get install -y libportaudio2 portaudio19-dev
|
||||
fi
|
||||
pip install fish-speech
|
||||
```
|
||||
|
||||
#### Kvcache attention fast path
|
||||
|
||||
@@ -123,6 +129,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
|
||||
- OS: Linux
|
||||
- Python: 3.10+
|
||||
- Required package: `fish-speech` (for DAC codec)
|
||||
- CUDA 12.8
|
||||
- vLLM version: 0.19.0+
|
||||
|
||||
|
||||
@@ -457,136 +457,5 @@ def test_text_benchmark_still_reports_ttft(capsys):
|
||||
assert "Time to First Token" in out, "text benchmarks must keep TTFT"
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Suppress text metrics for pure image / video endpoints
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def _make_image_output(prompt_len: int) -> MixRequestFuncOutput:
|
||||
output = MixRequestFuncOutput()
|
||||
output.success = True
|
||||
output.prompt_len = prompt_len
|
||||
output.output_tokens = 0
|
||||
output.generated_text = ""
|
||||
output.ttft = 0.0
|
||||
output.text_latency = 0.0
|
||||
output.latency = 2.5
|
||||
output.start_time = 0.0
|
||||
output.itl = []
|
||||
output.image_count = 1
|
||||
output.image_generation_time_ms = 2000.0
|
||||
output.image_pixels = 1024 * 1024
|
||||
output.input_audio_duration = 0.0
|
||||
output.error = ""
|
||||
return output
|
||||
|
||||
|
||||
def _make_video_output(prompt_len: int) -> MixRequestFuncOutput:
|
||||
output = MixRequestFuncOutput()
|
||||
output.success = True
|
||||
output.prompt_len = prompt_len
|
||||
output.output_tokens = 0
|
||||
output.generated_text = ""
|
||||
output.ttft = 0.0
|
||||
output.text_latency = 0.0
|
||||
output.latency = 5.0
|
||||
output.start_time = 0.0
|
||||
output.itl = []
|
||||
output.video_duration = 2.0
|
||||
output.video_frames = 48
|
||||
output.video_generation_time_ms = 4000.0
|
||||
output.video_rtf = 2.0
|
||||
output.input_audio_duration = 0.0
|
||||
output.error = ""
|
||||
return output
|
||||
|
||||
|
||||
_MEDIA_PERCENTILE_METRICS = ["e2el", "ttft"]
|
||||
|
||||
|
||||
def test_pure_image_benchmark_omits_text_result(capsys):
|
||||
"""Pure image generation must not print Text Result or invent peak tok/s."""
|
||||
outputs = [_make_image_output(64), _make_image_output(64)]
|
||||
|
||||
metrics, actual_output_lens = calculate_metrics(
|
||||
input_requests=[],
|
||||
outputs=outputs,
|
||||
dur_s=5.0,
|
||||
tokenizer=None,
|
||||
selected_percentiles=[99.0],
|
||||
goodput_config_dict={},
|
||||
task_type=TaskType.GENERATION,
|
||||
selected_percentile_metrics=_MEDIA_PERCENTILE_METRICS,
|
||||
max_concurrency=None,
|
||||
request_rate=float("inf"),
|
||||
benchmark_duration=5.0,
|
||||
)
|
||||
|
||||
out = capsys.readouterr().out
|
||||
assert actual_output_lens == [0, 0]
|
||||
assert metrics.total_output == 0
|
||||
assert metrics.max_output_tokens_per_s == 0.0
|
||||
assert " Text Result " not in out
|
||||
assert "Peak output token throughput" not in out
|
||||
assert "Time to First Token" not in out
|
||||
assert " Image Result " in out
|
||||
|
||||
|
||||
def test_pure_video_benchmark_omits_text_result(capsys):
|
||||
"""Pure video generation must not print Text Result or invent peak tok/s."""
|
||||
outputs = [_make_video_output(63), _make_video_output(63)]
|
||||
|
||||
metrics, actual_output_lens = calculate_metrics(
|
||||
input_requests=[],
|
||||
outputs=outputs,
|
||||
dur_s=10.0,
|
||||
tokenizer=None,
|
||||
selected_percentiles=[99.0],
|
||||
goodput_config_dict={},
|
||||
task_type=TaskType.GENERATION,
|
||||
selected_percentile_metrics=_MEDIA_PERCENTILE_METRICS,
|
||||
max_concurrency=None,
|
||||
request_rate=float("inf"),
|
||||
benchmark_duration=10.0,
|
||||
)
|
||||
|
||||
out = capsys.readouterr().out
|
||||
assert actual_output_lens == [0, 0]
|
||||
assert metrics.total_output == 0
|
||||
assert metrics.max_output_tokens_per_s == 0.0
|
||||
assert " Text Result " not in out
|
||||
assert "Peak output token throughput" not in out
|
||||
assert "Time to First Token" not in out
|
||||
assert " Video Result " in out
|
||||
|
||||
|
||||
def test_image_with_generated_text_still_reports_text_result(capsys):
|
||||
"""Image edits / AR text alongside images must keep Text Result."""
|
||||
output = _make_image_output(64)
|
||||
output.output_tokens = 8
|
||||
output.generated_text = "revised caption"
|
||||
output.ttft = 0.05
|
||||
output.itl = [0.01] * 7
|
||||
|
||||
calculate_metrics(
|
||||
input_requests=[],
|
||||
outputs=[output],
|
||||
dur_s=3.0,
|
||||
tokenizer=_EmptyAwareTokenizer(),
|
||||
selected_percentiles=[99.0],
|
||||
goodput_config_dict={},
|
||||
task_type=TaskType.GENERATION,
|
||||
selected_percentile_metrics=_MEDIA_PERCENTILE_METRICS,
|
||||
max_concurrency=None,
|
||||
request_rate=float("inf"),
|
||||
benchmark_duration=3.0,
|
||||
)
|
||||
|
||||
out = capsys.readouterr().out
|
||||
assert " Text Result " in out
|
||||
assert "Time to First Token" in out
|
||||
assert " Image Result " in out
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v", "-s"])
|
||||
@@ -22,16 +22,9 @@ from vllm_omni.benchmarks.data_modules.seed_tts_dataset import (
|
||||
)
|
||||
from vllm_omni.benchmarks.patch.patch import (
|
||||
MixRequestFuncOutput,
|
||||
_add_video_extra_body_to_form,
|
||||
_add_video_reference_to_form,
|
||||
_apply_stage0_token_timings,
|
||||
_apply_video_metrics_from_payload,
|
||||
_attach_seed_tts_to_request_func_input,
|
||||
_build_benchmark_session,
|
||||
_omni_request_timeout_s,
|
||||
async_request_openai_chat_omni_completions,
|
||||
async_request_openai_image_edits_omni,
|
||||
async_request_openai_image_generations_omni,
|
||||
async_request_openai_realtime_duplex,
|
||||
should_request_stage_metrics,
|
||||
)
|
||||
@@ -1203,337 +1196,5 @@ async def test_prompt_len_assigned_from_usage(mocker: MockerFixture):
|
||||
)
|
||||
|
||||
|
||||
class TestOmniRequestTimeout:
|
||||
"""``--omni-request-timeout-s`` precedence: explicit value > 900 s default."""
|
||||
|
||||
_OVERRIDE = "vllm_omni.benchmarks.patch.patch._REQUEST_TIMEOUT_OVERRIDE_S"
|
||||
|
||||
def test_default_timeout_is_900s_when_not_set(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(self._OVERRIDE, None)
|
||||
assert _omni_request_timeout_s() == 900.0
|
||||
|
||||
def test_explicit_value_wins_over_default(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(self._OVERRIDE, 123.5)
|
||||
assert _omni_request_timeout_s() == 123.5
|
||||
|
||||
def test_non_positive_restores_legacy_6h_cap(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(self._OVERRIDE, 0)
|
||||
assert _omni_request_timeout_s() == 6 * 60 * 60.0
|
||||
|
||||
async def test_benchmark_session_uses_configured_timeout(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(self._OVERRIDE, 42.0)
|
||||
session = _build_benchmark_session(max_concurrency=8, ssl_setting=False)
|
||||
try:
|
||||
assert session.timeout.total == 42.0
|
||||
assert session.connector.limit == 8
|
||||
assert session.connector.limit_per_host == 8
|
||||
finally:
|
||||
await session.close()
|
||||
|
||||
async def test_hung_server_request_times_out_as_failed(self, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""A server that accepts the request but never responds must surface as ``failed``."""
|
||||
|
||||
async def handler(reader, writer):
|
||||
# Drain the request but never respond — the original hang failure mode.
|
||||
try:
|
||||
while await reader.read(4096):
|
||||
pass
|
||||
except ConnectionResetError:
|
||||
pass
|
||||
finally:
|
||||
# Close our side of the socket: the client aborts after the
|
||||
# timeout and only sends FIN, so without writer.close() the
|
||||
# half-closed connection keeps Server.wait_closed() (which on
|
||||
# Python 3.12+ waits for every accepted connection to drop)
|
||||
# blocked forever.
|
||||
writer.close()
|
||||
try:
|
||||
await writer.wait_closed()
|
||||
except (ConnectionResetError, BrokenPipeError):
|
||||
pass
|
||||
|
||||
monkeypatch.setattr(self._OVERRIDE, 1.0)
|
||||
server = await asyncio.start_server(handler, "127.0.0.1", 0)
|
||||
port = server.sockets[0].getsockname()[1]
|
||||
session = _build_benchmark_session(max_concurrency=1, ssl_setting=False)
|
||||
try:
|
||||
request_input = _seed_tts_request_func_input()
|
||||
request_input.api_url = f"http://127.0.0.1:{port}/v1/chat/completions"
|
||||
output = await async_request_openai_chat_omni_completions(request_input, session)
|
||||
assert output.success is False
|
||||
assert output.error
|
||||
finally:
|
||||
await session.close()
|
||||
server.close()
|
||||
await server.wait_closed()
|
||||
|
||||
|
||||
def test_video_rtf_prefers_generation_time_over_poll_latency():
|
||||
"""RTF must not grow with client poll_interval overshoot in output.latency."""
|
||||
payload = {
|
||||
"duration_s": 2.0,
|
||||
"num_frames": 48,
|
||||
"fps": 24.0,
|
||||
"stage_durations": {"stage_0_gen_ms": 4000.0},
|
||||
}
|
||||
request_body: dict[str, object] = {}
|
||||
|
||||
short_poll = MixRequestFuncOutput()
|
||||
short_poll.latency = 4.2 # ~generation + small poll overshoot
|
||||
_apply_video_metrics_from_payload(short_poll, payload, request_body)
|
||||
|
||||
long_poll = MixRequestFuncOutput()
|
||||
long_poll.latency = 8.0 # same job, larger poll_interval_s overshoot
|
||||
_apply_video_metrics_from_payload(long_poll, payload, request_body)
|
||||
|
||||
assert short_poll.video_generation_time_ms == pytest.approx(4000.0)
|
||||
assert long_poll.video_generation_time_ms == pytest.approx(4000.0)
|
||||
assert short_poll.video_rtf == pytest.approx(2.0)
|
||||
assert long_poll.video_rtf == pytest.approx(2.0)
|
||||
assert short_poll.video_rtf == long_poll.video_rtf
|
||||
|
||||
|
||||
def test_video_rtf_falls_back_to_e2e_latency_without_generation_time():
|
||||
output = MixRequestFuncOutput()
|
||||
output.latency = 6.0
|
||||
_apply_video_metrics_from_payload(
|
||||
output,
|
||||
{"duration_s": 2.0, "num_frames": 48, "fps": 24.0},
|
||||
{},
|
||||
)
|
||||
assert output.video_generation_time_ms == 0.0
|
||||
assert output.video_rtf == pytest.approx(3.0)
|
||||
|
||||
|
||||
# 1x1 PNG used as a valid b64_json image payload in edit-client tests.
|
||||
_MIN_PNG_B64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg=="
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_image_edits_defaults_to_non_streaming_json(mocker: MockerFixture) -> None:
|
||||
"""Single-stage servers reject stream=true; default path must send stream=false + JSON."""
|
||||
|
||||
class MockJsonResponse:
|
||||
status = 200
|
||||
|
||||
def __init__(self):
|
||||
self._payload = {
|
||||
"created": 1,
|
||||
"data": [
|
||||
{
|
||||
"b64_json": _MIN_PNG_B64,
|
||||
"stage_durations": {"stage_0_gen_ms": 12.5},
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
async def json(self):
|
||||
return self._payload
|
||||
|
||||
async def text(self):
|
||||
return json.dumps(self._payload)
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *_args):
|
||||
return None
|
||||
|
||||
captured_stream: list[str] = []
|
||||
real_add_field = None
|
||||
|
||||
def tracking_add_field(self, name, value=None, **kwargs):
|
||||
if name == "stream":
|
||||
captured_stream.append(str(value))
|
||||
return real_add_field(self, name, value, **kwargs)
|
||||
|
||||
import aiohttp
|
||||
|
||||
real_add_field = aiohttp.FormData.add_field
|
||||
mocker.patch.object(aiohttp.FormData, "add_field", tracking_add_field)
|
||||
|
||||
mock_session = mocker.AsyncMock()
|
||||
mock_session.post = mocker.MagicMock(return_value=MockJsonResponse())
|
||||
|
||||
request = RequestFuncInput(
|
||||
model="single-stage-edit",
|
||||
model_name="single-stage-edit",
|
||||
prompt="make it sunny",
|
||||
api_url="http://test.com/v1/images/edits",
|
||||
prompt_len=4,
|
||||
output_len=1,
|
||||
multi_modal_content=[{"type": "image_url", "image_url": {"url": f"data:image/png;base64,{_MIN_PNG_B64}"}}],
|
||||
extra_body={"num_inference_steps": 4},
|
||||
)
|
||||
output = await async_request_openai_image_edits_omni(request, mock_session, pbar=None)
|
||||
|
||||
assert captured_stream == ["false"]
|
||||
assert output.success is True
|
||||
assert not output.error
|
||||
assert output.image_count == 1
|
||||
assert output.image_generation_time_ms == pytest.approx(12.5)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_image_edits_stream_true_uses_sse_path(mocker: MockerFixture) -> None:
|
||||
"""Explicit stream=true keeps the multi-stage SSE client path."""
|
||||
sse_chunk = (
|
||||
b'data: {"type":"image","data":[{"b64_json":"' + _MIN_PNG_B64.encode() + b'"}]}\n\n' + b"data: [DONE]\n\n"
|
||||
)
|
||||
mock_response = MockResponse(200, [sse_chunk])
|
||||
captured_stream: list[str] = []
|
||||
|
||||
import aiohttp
|
||||
|
||||
real_add_field = aiohttp.FormData.add_field
|
||||
|
||||
def tracking_add_field(self, name, value=None, **kwargs):
|
||||
if name == "stream":
|
||||
captured_stream.append(str(value))
|
||||
return real_add_field(self, name, value, **kwargs)
|
||||
|
||||
mocker.patch.object(aiohttp.FormData, "add_field", tracking_add_field)
|
||||
|
||||
mock_session = mocker.AsyncMock()
|
||||
mock_session.post = mocker.MagicMock(return_value=mock_response)
|
||||
|
||||
request = RequestFuncInput(
|
||||
model="multi-stage-edit",
|
||||
model_name="multi-stage-edit",
|
||||
prompt="edit",
|
||||
api_url="http://test.com/v1/images/edits",
|
||||
prompt_len=2,
|
||||
output_len=1,
|
||||
multi_modal_content=[{"type": "image_url", "image_url": {"url": f"data:image/png;base64,{_MIN_PNG_B64}"}}],
|
||||
extra_body={"stream": True},
|
||||
)
|
||||
output = await async_request_openai_image_edits_omni(request, mock_session, pbar=None)
|
||||
|
||||
assert captured_stream == ["true"]
|
||||
assert output.success is True
|
||||
assert output.image_count == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_image_generations_e2el_includes_json_body_consume(mocker: MockerFixture) -> None:
|
||||
"""E2EL must include body transfer/decode, not stop at HTTP headers."""
|
||||
|
||||
class SlowJsonResponse:
|
||||
status = 200
|
||||
|
||||
def __init__(self):
|
||||
self._payload = {
|
||||
"created": 1,
|
||||
"data": [{"b64_json": _MIN_PNG_B64, "stage_durations": {"stage_0_gen_ms": 1.0}}],
|
||||
}
|
||||
|
||||
async def json(self):
|
||||
await asyncio.sleep(0.05)
|
||||
return self._payload
|
||||
|
||||
async def text(self):
|
||||
return json.dumps(self._payload)
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *_args):
|
||||
return None
|
||||
|
||||
mock_session = mocker.AsyncMock()
|
||||
mock_session.post = mocker.MagicMock(return_value=SlowJsonResponse())
|
||||
request = RequestFuncInput(
|
||||
model="img-gen",
|
||||
model_name="img-gen",
|
||||
prompt="a cat",
|
||||
api_url="http://test.com/v1/images/generations",
|
||||
prompt_len=2,
|
||||
output_len=1,
|
||||
extra_body={},
|
||||
)
|
||||
output = await async_request_openai_image_generations_omni(request, mock_session, pbar=None)
|
||||
assert output.success is True
|
||||
assert output.latency >= 0.05
|
||||
|
||||
|
||||
def test_video_local_image_reference_not_forwarded_as_raw_extra_field(tmp_path, mocker: MockerFixture) -> None:
|
||||
"""Local image_reference must only be uploaded via dedicated serializer."""
|
||||
import aiohttp
|
||||
|
||||
ref_path = tmp_path / "ref.png"
|
||||
ref_path.write_bytes(base64.b64decode(_MIN_PNG_B64))
|
||||
|
||||
captured: list[tuple[str, object]] = []
|
||||
real_add_field = aiohttp.FormData.add_field
|
||||
|
||||
def tracking_add_field(self, name, value=None, **kwargs):
|
||||
captured.append((str(name), value))
|
||||
return real_add_field(self, name, value, **kwargs)
|
||||
|
||||
mocker.patch.object(aiohttp.FormData, "add_field", tracking_add_field)
|
||||
|
||||
form = aiohttp.FormData()
|
||||
extra_body = {
|
||||
"num_inference_steps": 2,
|
||||
"image_reference": str(ref_path),
|
||||
}
|
||||
request_body = {"model": "vid", "prompt": "p", **extra_body}
|
||||
_add_video_extra_body_to_form(form, extra_body, request_body)
|
||||
assert _add_video_reference_to_form(form, extra_body["image_reference"]) is True
|
||||
|
||||
field_names = [name for name, _ in captured]
|
||||
assert "image_reference" not in field_names
|
||||
assert field_names.count("input_reference") == 1
|
||||
# Dedicated uploader sends file bytes, not the raw local path string.
|
||||
uploaded = next(value for name, value in captured if name == "input_reference")
|
||||
assert uploaded == base64.b64decode(_MIN_PNG_B64)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"reference",
|
||||
[
|
||||
{"image_url": "https://example.com/ref.png"},
|
||||
{"file_id": "file-abc"},
|
||||
[{"image_url": "https://example.com/a.png"}, {"file_id": "file-xyz"}],
|
||||
],
|
||||
)
|
||||
def test_video_structured_image_reference_serialized_to_form(reference: object, mocker: MockerFixture) -> None:
|
||||
"""Object-form image_reference must be JSON-serialized, not silently dropped."""
|
||||
import aiohttp
|
||||
|
||||
captured: list[tuple[str, object]] = []
|
||||
real_add_field = aiohttp.FormData.add_field
|
||||
|
||||
def tracking_add_field(self, name, value=None, **kwargs):
|
||||
captured.append((str(name), value))
|
||||
return real_add_field(self, name, value, **kwargs)
|
||||
|
||||
mocker.patch.object(aiohttp.FormData, "add_field", tracking_add_field)
|
||||
|
||||
form = aiohttp.FormData()
|
||||
extra_body = {"image_reference": reference}
|
||||
request_body = {"model": "vid", "prompt": "p", **extra_body}
|
||||
_add_video_extra_body_to_form(form, extra_body, request_body)
|
||||
assert _add_video_reference_to_form(form, reference) is True
|
||||
|
||||
field_names = [name for name, _ in captured]
|
||||
# Reserved key must not be double-forwarded as a generic extra field;
|
||||
# exactly one dedicated image_reference field.
|
||||
assert field_names.count("image_reference") == 1
|
||||
assert "input_reference" not in field_names
|
||||
payload = next(value for name, value in captured if name == "image_reference")
|
||||
assert json.loads(payload) == reference
|
||||
|
||||
|
||||
def test_video_unsupported_image_reference_raises() -> None:
|
||||
import aiohttp
|
||||
|
||||
form = aiohttp.FormData()
|
||||
with pytest.raises(ValueError, match="Unsupported image_reference"):
|
||||
_add_video_reference_to_form(form, {"not_a_supported_key": "x"})
|
||||
with pytest.raises(ValueError, match="Unsupported image_reference"):
|
||||
_add_video_reference_to_form(form, "/tmp/does-not-exist-ref.png")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v", "-s"])
|
||||
@@ -1,11 +1,9 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
# ruff: noqa: E402, I001
|
||||
import argparse
|
||||
import math
|
||||
import os
|
||||
import sys
|
||||
import types
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
@@ -108,9 +106,11 @@ def test_resolve_seed_tts_root_downloads_only_requested_locale(monkeypatch, tmp_
|
||||
captured["allow_patterns"] = allow_patterns
|
||||
return str(downloaded_root)
|
||||
|
||||
from vllm_omni.transformers_utils import repo_utils
|
||||
|
||||
monkeypatch.setattr(repo_utils.hf_api(), "snapshot_download", fake_snapshot_download)
|
||||
monkeypatch.setitem(
|
||||
sys.modules,
|
||||
"huggingface_hub",
|
||||
types.SimpleNamespace(snapshot_download=fake_snapshot_download),
|
||||
)
|
||||
|
||||
resolved = resolve_seed_tts_root(
|
||||
"zhaochenyang20/seed-tts-eval",
|
||||
|
||||
@@ -26,7 +26,6 @@ _MODULE_NAME = "vllm_omni.benchmarks.data_modules.seed_tts_dataset"
|
||||
|
||||
if _MODULE_NAME not in sys.modules:
|
||||
_spec = importlib.util.spec_from_file_location(_MODULE_NAME, _MODULE_PATH)
|
||||
assert _spec is not None and _spec.loader is not None
|
||||
_mod = importlib.util.module_from_spec(_spec)
|
||||
sys.modules[_MODULE_NAME] = _mod
|
||||
_spec.loader.exec_module(_mod)
|
||||
@@ -238,8 +237,7 @@ def test_seed_tts_design_dataset_rejects_missing_description(seed_tts_design_roo
|
||||
assert req.seed_tts_utterance_id == "ok"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("audio_seconds", [0.1, 61])
|
||||
def test_seed_tts_whisper_transcribe_passes_attention_mask(monkeypatch, audio_seconds):
|
||||
def test_seed_tts_whisper_transcribe_passes_attention_mask(monkeypatch):
|
||||
from vllm_omni.benchmarks.data_modules import seed_tts_eval
|
||||
|
||||
calls = {}
|
||||
@@ -254,9 +252,8 @@ def test_seed_tts_whisper_transcribe_passes_attention_mask(monkeypatch, audio_se
|
||||
return self
|
||||
|
||||
class FakeProcessor:
|
||||
def __call__(self, wav, *, sampling_rate, return_tensors, return_attention_mask=False, **kwargs):
|
||||
def __call__(self, wav, *, sampling_rate, return_tensors, return_attention_mask=False):
|
||||
calls["return_attention_mask"] = return_attention_mask
|
||||
calls["processor_kwargs"] = kwargs
|
||||
assert sampling_rate == 16000
|
||||
assert return_tensors == "pt"
|
||||
assert len(wav) > 0
|
||||
@@ -265,10 +262,9 @@ def test_seed_tts_whisper_transcribe_passes_attention_mask(monkeypatch, audio_se
|
||||
attention_mask=FakeTensor("mask") if return_attention_mask else None,
|
||||
)
|
||||
|
||||
def get_decoder_prompt_ids(self, *, language, task, no_timestamps=True):
|
||||
def get_decoder_prompt_ids(self, *, language, task):
|
||||
assert language == "english"
|
||||
assert task == "transcribe"
|
||||
calls["no_timestamps"] = no_timestamps
|
||||
return [(1, 2)]
|
||||
|
||||
def batch_decode(self, predicted_ids, *, skip_special_tokens):
|
||||
@@ -287,18 +283,10 @@ def test_seed_tts_whisper_transcribe_passes_attention_mask(monkeypatch, audio_se
|
||||
monkeypatch.setattr(seed_tts_eval, "_en_model", FakeModel())
|
||||
monkeypatch.setattr(seed_tts_eval, "_device", "cuda:1")
|
||||
|
||||
text = seed_tts_eval._transcribe_en_f32_16k(np.ones(int(audio_seconds * 16000), dtype=np.float32))
|
||||
text = seed_tts_eval._transcribe_en_f32_16k(np.ones(1600, dtype=np.float32))
|
||||
|
||||
assert text == "hello"
|
||||
assert calls["return_attention_mask"] is True
|
||||
assert calls["input_features"].device == "cuda:1"
|
||||
assert calls["generate_kwargs"]["attention_mask"].device == "cuda:1"
|
||||
assert calls["generate_kwargs"]["forced_decoder_ids"] == [(1, 2)]
|
||||
if audio_seconds > 30:
|
||||
assert calls["processor_kwargs"] == {"truncation": False, "padding": "longest"}
|
||||
assert calls["generate_kwargs"]["return_timestamps"] is True
|
||||
assert calls["no_timestamps"] is False
|
||||
else:
|
||||
assert calls["processor_kwargs"] == {}
|
||||
assert "return_timestamps" not in calls["generate_kwargs"]
|
||||
assert calls["no_timestamps"] is True
|
||||
@@ -167,26 +167,6 @@ def test_preprocess_serve_args_applies_safe_omniinteract_prompt_default(
|
||||
{"bot_task": "recaption"},
|
||||
{"backend", "print_stage", "bot_task"},
|
||||
),
|
||||
(
|
||||
[
|
||||
"--endpoint",
|
||||
"/v1/images/edits",
|
||||
"--image-edits-bot-task",
|
||||
"think",
|
||||
],
|
||||
{"bot_task": "think"},
|
||||
{"endpoint", "bot_task"},
|
||||
),
|
||||
(
|
||||
[
|
||||
"--backend",
|
||||
"/v1/images/edits",
|
||||
"--image-edits-bot-task",
|
||||
"recaption",
|
||||
],
|
||||
{"bot_task": "recaption"},
|
||||
{"backend", "bot_task"},
|
||||
),
|
||||
(
|
||||
["--extra-body", '{"bot_task":"vanilla"}'],
|
||||
{"bot_task": "vanilla"},
|
||||
@@ -202,7 +182,6 @@ def test_omni_args_parse_and_preprocess(
|
||||
parser = TrackingArgumentParser()
|
||||
parser.add_argument("--extra-body", type=json.loads, default=None)
|
||||
parser.add_argument("--backend", default="openai-chat-omni")
|
||||
parser.add_argument("--endpoint", default="/v1/chat/completions")
|
||||
add_omni_args(parser)
|
||||
|
||||
args = parser.parse_args(argv)
|
||||
@@ -379,17 +358,3 @@ def test_bench_serve_cli_mocks_http_request(tmp_path: Path):
|
||||
)
|
||||
assert bench_requests
|
||||
assert all(url == expected_url for url in bench_requests), f"Unexpected target URLs: {bench_requests}"
|
||||
|
||||
|
||||
def test_omni_request_timeout_s_flag_defaults_and_parses() -> None:
|
||||
parser = TrackingArgumentParser()
|
||||
add_omni_args(parser)
|
||||
|
||||
args = parser.parse_args([])
|
||||
assert args.omni_request_timeout_s is None
|
||||
|
||||
args = parser.parse_args(["--omni-request-timeout-s", "0"])
|
||||
assert args.omni_request_timeout_s == 0.0
|
||||
|
||||
args = parser.parse_args(["--omni-request-timeout-s", "3600"])
|
||||
assert args.omni_request_timeout_s == 3600.0
|
||||
@@ -6,6 +6,7 @@ Unit tests for StageConfigFactory and related classes.
|
||||
|
||||
import importlib
|
||||
import warnings
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
@@ -36,7 +37,6 @@ from vllm_omni.config.stage_config import (
|
||||
normalize_pipeline_cli_overrides,
|
||||
pipeline_cfg_resolver,
|
||||
)
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
from vllm_omni.engine.arg_utils import SHARED_FIELDS, internal_blacklist_keys
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
@@ -451,8 +451,13 @@ class TestStageConfigFactory:
|
||||
|
||||
def test_default_diffusion_with_parallel_config(self):
|
||||
"""Test diffusion config calculates devices from parallel_config."""
|
||||
|
||||
@dataclass
|
||||
class MockParallelConfig:
|
||||
world_size: int = 4
|
||||
|
||||
kwargs = {
|
||||
"parallel_config": DiffusionParallelConfig(tensor_parallel_size=4),
|
||||
"parallel_config": MockParallelConfig(),
|
||||
"cache_backend": "tea_cache",
|
||||
}
|
||||
configs = StageConfigFactory.create_default_diffusion(kwargs)
|
||||
@@ -932,44 +937,6 @@ class TestPipelineRegistration:
|
||||
assert isinstance(omni_config, VllmOmniConfig)
|
||||
assert omni_config.stage_by_id(0).diffusion_config.model == "fake/model"
|
||||
|
||||
def test_cosyvoice3_deploy_resolves_hash_snapshot_without_hf_metadata(self, tmp_path):
|
||||
snapshot = tmp_path / "models--FunAudioLLM--Fun-CosyVoice3-0.5B-2512" / "snapshots" / ("a" * 40)
|
||||
snapshot.mkdir(parents=True)
|
||||
(snapshot / "config.json").write_text("{}\n", encoding="utf-8")
|
||||
model = str(snapshot)
|
||||
deploy_path = get_deploy_config_path("cosyvoice3.yaml")
|
||||
|
||||
assert StageConfigFactory.try_infer_model_type(model, trust_remote_code=False) is None
|
||||
assert StageConfigFactory.get_pipeline_config(model, trust_remote_code=False) is None
|
||||
|
||||
omni_config = StageConfigFactory.create_from_model(
|
||||
model,
|
||||
trust_remote_code=False,
|
||||
cli_overrides={},
|
||||
deploy_config_path=deploy_path,
|
||||
)
|
||||
legacy_configs, _ = StageConfigFactory.create_legacy_stage_configs_from_model(
|
||||
model,
|
||||
trust_remote_code=False,
|
||||
cli_overrides={},
|
||||
deploy_config_path=deploy_path,
|
||||
)
|
||||
|
||||
assert omni_config is not None
|
||||
assert omni_config.pipeline_config.model_type == "cosyvoice3"
|
||||
assert [
|
||||
(stage.stage_id, stage.stage_type, stage.worker_type, stage.model_stage)
|
||||
for stage in omni_config.stage_configs
|
||||
] == [
|
||||
(0, StageType.LLM, "ar", "cosyvoice3_talker"),
|
||||
(1, StageType.LLM, "generation", "cosyvoice3_code2wav"),
|
||||
]
|
||||
assert all(stage.model_config.model == model for stage in omni_config.stage_configs)
|
||||
assert all(stage.stage_type is not StageType.DIFFUSION for stage in omni_config.stage_configs)
|
||||
|
||||
assert legacy_configs is not None
|
||||
assert [stage.stage_type for stage in legacy_configs] == [StageType.LLM, StageType.LLM]
|
||||
|
||||
def test_pipeline_registration(self, clean_pipeline_registry):
|
||||
"""Ensure that we can register and create a custom pipeline config."""
|
||||
new_model_type = "new_model_type"
|
||||
@@ -1508,11 +1475,6 @@ stages:
|
||||
deploy = load_deploy_config(Path(get_deploy_config_path("minimax_h3_disaggregated.yaml")))
|
||||
stages = merge_pipeline_deploy(pipeline, deploy)
|
||||
|
||||
assert stages[0].yaml_engine_args["model_arch"] == "MiniMaxH3TextEncoder"
|
||||
assert stages[1].yaml_engine_args["model_arch"] == "MiniMaxH3Pipeline"
|
||||
assert stages[0].yaml_runtime["num_replicas"] == 1
|
||||
assert stages[1].yaml_runtime["num_replicas"] == 1
|
||||
assert stages[1].yaml_engine_args["model_loaded"] == {"text_encoder": False}
|
||||
assert stages[0].yaml_engine_args["max_num_seqs"] == 1
|
||||
assert stages[0].yaml_engine_args["model_path_resolver"].endswith(".resolve_minimax_h3_model_root")
|
||||
assert stages[1].yaml_engine_args["model_path_resolver"].endswith(".resolve_minimax_h3_diffusion_model_path")
|
||||
@@ -1635,7 +1597,7 @@ stages:
|
||||
assert stage.max_model_len == 1024
|
||||
assert stage.max_num_batched_tokens == 1024
|
||||
assert stage.max_num_seqs == 1
|
||||
assert stage.gpu_memory_utilization == 0.8
|
||||
assert stage.gpu_memory_utilization == 0.7
|
||||
assert stage.skip_mm_profiling is True
|
||||
assert stage.enforce_eager is True
|
||||
assert stage.async_scheduling is False
|
||||
|
||||
@@ -167,7 +167,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
|
||||
EnvironmentVariableCategory.PUBLIC_OMNI: 24,
|
||||
EnvironmentVariableCategory.INHERITED_VLLM: 20,
|
||||
EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27,
|
||||
EnvironmentVariableCategory.MODEL_SPECIFIC: 61,
|
||||
EnvironmentVariableCategory.MODEL_SPECIFIC: 62,
|
||||
EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 20,
|
||||
EnvironmentVariableCategory.INTERNAL: 2,
|
||||
}
|
||||
@@ -179,7 +179,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
|
||||
)
|
||||
assert {disposition: disposition_counts[disposition] for disposition in ModelEnvironmentVariableDisposition} == {
|
||||
ModelEnvironmentVariableDisposition.PROMOTE: 37,
|
||||
ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 6,
|
||||
ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 7,
|
||||
ModelEnvironmentVariableDisposition.EXTERNAL: 0,
|
||||
ModelEnvironmentVariableDisposition.INTERNALIZE: 13,
|
||||
ModelEnvironmentVariableDisposition.DEPRECATE_REMOVE: 5,
|
||||
|
||||
@@ -1,52 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""The LingBot stepwise deploy config must keep serving on the AR-Diffusion path.
|
||||
|
||||
`WS /v1/realtime/video` only reaches the stepwise code when the stage selects the
|
||||
AR-Diffusion engine and asks for streamed step execution, so those fields are
|
||||
pinned here rather than left to drift.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.helpers.stage_config import get_deploy_config_path
|
||||
from vllm_omni.config.stage_config import load_deploy_config
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
_DEPLOY = "lingbot_world_v2_stepwise.yaml"
|
||||
|
||||
|
||||
def _stage():
|
||||
deploy = load_deploy_config(get_deploy_config_path(_DEPLOY))
|
||||
assert len(deploy.stages) == 1, "stepwise serving is a single diffusion stage"
|
||||
return deploy.stages[0]
|
||||
|
||||
|
||||
def test_stage_selects_the_ar_diffusion_engine_and_pipeline() -> None:
|
||||
stage = _stage()
|
||||
assert stage.model_class_name == "LingBotWorldCausalDMDPipeline"
|
||||
# engine_backend/model_config reach the engine through the passthrough, the
|
||||
# same way the DreamZero AR-Diffusion deploy config does.
|
||||
assert stage.engine_extras["engine_backend"].endswith("ar_diffusion.engine.ARDiffusionEngine")
|
||||
|
||||
|
||||
def test_stage_requests_streamed_step_execution_on_one_sequence() -> None:
|
||||
stage = _stage()
|
||||
assert stage.step_execution is True
|
||||
assert stage.engine_extras["streaming_output"] is True
|
||||
# AR-Diffusion binds one session per runner invocation; batching requests
|
||||
# would break that contract, so the runner rejects max_num_seqs > 1.
|
||||
assert stage.max_num_seqs == 1
|
||||
|
||||
|
||||
def test_stage_declares_the_fixed_ar_cache_geometry() -> None:
|
||||
stage = _stage()
|
||||
model_config = stage.engine_extras["model_config"]
|
||||
# The KV cache geometry is fixed at load time, so a served request must ask
|
||||
# for exactly this resolution.
|
||||
assert model_config["ar_diffusion_height"] == 480
|
||||
assert model_config["ar_diffusion_width"] == 832
|
||||
assert model_config["ar_diffusion_kv_config"]["gpu_memory_fraction"] == 0.6
|
||||
@@ -261,57 +261,6 @@ def test_from_pipeline_config_rejects_unowned_deploy_engine_extras(engine_extras
|
||||
VllmOmniConfig.from_pipeline_config(pipeline, user_deploy_config=deploy)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stage_id", [0, 2], ids=["ar", "generation"])
|
||||
def test_llm_additional_config_roundtrip_and_isolation(stage_id, monkeypatch):
|
||||
from vllm_omni.engine import stage_init_utils
|
||||
from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config
|
||||
|
||||
# Worker discovery requires hardware support; this test covers config transport.
|
||||
monkeypatch.setattr(stage_init_utils, "resolve_worker_cls", lambda engine_args: None)
|
||||
additional_config = {"backend_options": {"enabled": True}}
|
||||
pipeline = _resolve_pipeline_or_skip("minicpmo_4_5")
|
||||
deploy = DeployConfig(
|
||||
stages=[StageDeployConfig(stage_id=i, engine_extras={"additional_config": additional_config}) for i in (0, 2)]
|
||||
)
|
||||
config = VllmOmniConfig.from_pipeline_config(pipeline, user_deploy_config=deploy)
|
||||
stage = config.stage_by_id(stage_id)
|
||||
assert stage.runtime_config.additional_config == additional_config
|
||||
engine_args = build_engine_args_dict_from_omni_stage_config(stage, model="test-model")
|
||||
assert engine_args["additional_config"] == additional_config
|
||||
|
||||
engine_args["additional_config"]["backend_options"]["enabled"] = False
|
||||
assert stage.runtime_config.additional_config["backend_options"]["enabled"] is True
|
||||
stage.runtime_config.additional_config["backend_options"]["enabled"] = False
|
||||
assert config.stage_by_id(2 if stage_id == 0 else 0).runtime_config.additional_config == additional_config
|
||||
assert additional_config["backend_options"]["enabled"] is True
|
||||
|
||||
|
||||
@pytest.mark.parametrize("deploy_name", ["minicpmo_4_5", "minicpmo_4_5_2gpu", "minicpmo_4_5_3gpu"])
|
||||
def test_minicpmo_npu_additional_config_reaches_engine_args(monkeypatch, deploy_name):
|
||||
from vllm_omni.engine import stage_init_utils
|
||||
from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
monkeypatch.setattr(current_omni_platform, "device_name", "npu")
|
||||
monkeypatch.setattr(stage_init_utils, "resolve_worker_cls", lambda engine_args: None)
|
||||
stage = _from_pipeline_key("minicpmo_4_5", deploy_config_path=deploy_name).stage_by_id(2)
|
||||
expected = {"code2wav_enable_npu_graph": True, "code2wav_max_npu_graphs": 32}
|
||||
assert stage.runtime_config.additional_config == expected
|
||||
engine_args = build_engine_args_dict_from_omni_stage_config(stage, model="test-model")
|
||||
assert engine_args["additional_config"] == expected
|
||||
|
||||
|
||||
def test_diffusion_additional_config_keeps_diffusion_owner():
|
||||
from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config
|
||||
|
||||
additional_config = {"torchair_graph_config": {"enabled": True}}
|
||||
stage = _from_pipeline_key("dreamzero", cli_overrides={"additional_config": additional_config}).stage_by_id(0)
|
||||
assert stage.runtime_config.additional_config is None
|
||||
assert stage.diffusion_config.additional_config == additional_config
|
||||
engine_args = build_engine_args_dict_from_omni_stage_config(stage, model="test-model")
|
||||
assert engine_args["additional_config"] == additional_config
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("engine_extras", "cli_overrides"),
|
||||
[
|
||||
@@ -583,17 +532,6 @@ def test_from_pipeline_config_dispatches_async_chunk_processors_without_mutating
|
||||
assert pipeline.get_stage(1).custom_process_input_func is None
|
||||
|
||||
|
||||
def test_joyai_code2wav_waits_for_full_payload():
|
||||
config = _from_pipeline_key("joyai_vl_interaction")
|
||||
talker = config.stage_by_id(1)
|
||||
code2wav = config.stage_by_id(2)
|
||||
|
||||
assert talker.custom_process_next_stage_input_func.endswith("talker2code2wav_full_payload")
|
||||
assert code2wav.custom_process_input_func.endswith("talker2code2wav_token_only")
|
||||
assert code2wav.connector_config.async_chunk is False
|
||||
assert code2wav.model_config.requires_full_payload_input is True
|
||||
|
||||
|
||||
def test_vllm_omni_stage_config_public_fields_use_typed_stage_realizations():
|
||||
assert not hasattr(BaseVllmOmniStageConfig, "from_stage_config")
|
||||
assert not hasattr(BaseVllmOmniStageConfig, "to_legacy_stage_config")
|
||||
@@ -621,7 +559,6 @@ def test_vllm_omni_stage_config_public_fields_use_typed_stage_realizations():
|
||||
|
||||
def test_runtime_config_fields_match_structured_runtime_scope():
|
||||
assert {f.name for f in fields(OmniStageRuntimeConfig)} == {
|
||||
"additional_config",
|
||||
"distributed_executor_backend",
|
||||
"worker_cls",
|
||||
"devices",
|
||||
@@ -1316,70 +1253,6 @@ def test_diffusion_config_from_kwargs_reuses_legacy_normalization(monkeypatch):
|
||||
assert cfg.diffusers_call_kwargs == {}
|
||||
|
||||
|
||||
def test_diffusion_config_none_values_preserve_dataclass_defaults():
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
|
||||
normalized = OmniDiffusionConfig.normalize_init_kwargs(
|
||||
{
|
||||
"lora_scale": None,
|
||||
"enable_sleep_mode": None,
|
||||
"diffusers_load_kwargs": None,
|
||||
}
|
||||
)
|
||||
|
||||
assert "lora_scale" not in normalized
|
||||
assert "enable_sleep_mode" not in normalized
|
||||
assert normalized["diffusers_load_kwargs"] == {}
|
||||
|
||||
config = OmniDiffusionConfig.from_kwargs(
|
||||
lora_scale=None,
|
||||
enable_sleep_mode=None,
|
||||
)
|
||||
assert config.lora_scale == 1.0
|
||||
assert config.enable_sleep_mode is False
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("canonical_key", "alias_key", "canonical_value", "alias_value"),
|
||||
[
|
||||
("lora_scale", "static_lora_scale", 0.75, 0.25),
|
||||
(
|
||||
"quantization_config",
|
||||
"diffusion_quantization_config",
|
||||
{"method": "canonical"},
|
||||
{"method": "diffusion-alias"},
|
||||
),
|
||||
(
|
||||
"quantization_config",
|
||||
"quantization",
|
||||
{"method": "canonical"},
|
||||
"legacy-alias",
|
||||
),
|
||||
("diffusion_kv_cache_dtype", "kv_cache_dtype", "fp8", "fp16"),
|
||||
("diffusion_kv_cache_skip_steps", "kv_cache_skip_steps", "0-1", "2-3"),
|
||||
("diffusion_kv_cache_skip_layers", "kv_cache_skip_layers", "1-2", "3-4"),
|
||||
("streaming_output", "diffusion_streaming_output", False, True),
|
||||
],
|
||||
)
|
||||
def test_diffusion_alias_conflicts_prefer_canonical_key(
|
||||
canonical_key,
|
||||
alias_key,
|
||||
canonical_value,
|
||||
alias_value,
|
||||
):
|
||||
from vllm_omni.diffusion.data import normalize_omni_diffusion_kwargs
|
||||
|
||||
normalized = normalize_omni_diffusion_kwargs(
|
||||
{
|
||||
canonical_key: canonical_value,
|
||||
alias_key: alias_value,
|
||||
}
|
||||
)
|
||||
|
||||
assert normalized[canonical_key] == canonical_value
|
||||
assert alias_key not in normalized
|
||||
|
||||
|
||||
def test_from_pipeline_config_normalizes_diffusion_config_aliases_from_engine_args(tmp_path, monkeypatch):
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
|
||||
@@ -19,9 +19,6 @@ from vllm.v1.core.sched.output import SchedulerOutput
|
||||
from vllm.v1.outputs import ModelRunnerOutput
|
||||
from vllm.v1.request import Request, RequestStatus, StreamingUpdate
|
||||
from vllm_omni.core.sched.omni_ar_scheduler import OmniARAsyncScheduler, OmniARScheduler
|
||||
from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import (
|
||||
OmniChunkTransferAdapter,
|
||||
)
|
||||
|
||||
# isort: on
|
||||
|
||||
@@ -65,43 +62,6 @@ def _make_update(prompt_token_ids: list[int] | None = None) -> StreamingUpdate:
|
||||
)
|
||||
|
||||
|
||||
def _make_talker_adapter(
|
||||
*,
|
||||
max_model_len: int = 100,
|
||||
recompute_on_capacity: bool = True,
|
||||
) -> OmniChunkTransferAdapter:
|
||||
adapter = OmniChunkTransferAdapter.__new__(OmniChunkTransferAdapter)
|
||||
adapter.receives_chunks = False
|
||||
adapter._max_model_len = max_model_len
|
||||
adapter._streaming_prompt_previous_chunks = 1
|
||||
adapter._streaming_prompt_recompute_on_capacity = recompute_on_capacity
|
||||
adapter._streaming_condition_lengths = {}
|
||||
adapter._streaming_condition_seqs = {}
|
||||
adapter.segment_finished_requests = set()
|
||||
adapter.requests_num_chunks_sent = {}
|
||||
return adapter
|
||||
|
||||
|
||||
def _make_talker_update(
|
||||
condition_len: int,
|
||||
*,
|
||||
reserve: int = 10,
|
||||
condition_seq: int | None,
|
||||
) -> StreamingUpdate:
|
||||
update = _make_update([0] * condition_len)
|
||||
update.model_intermediate_buffer = {
|
||||
"native_duplex": True,
|
||||
"ids": {"prompt": [1]},
|
||||
"meta": {
|
||||
"next_stage_prompt_len": condition_len,
|
||||
"next_stage_generation_tokens": reserve,
|
||||
},
|
||||
}
|
||||
if condition_seq is not None:
|
||||
update.model_intermediate_buffer["meta"]["streaming_condition_seq"] = condition_seq
|
||||
return update
|
||||
|
||||
|
||||
def _run_resumable_segment_stop(
|
||||
session: Request,
|
||||
*,
|
||||
@@ -578,288 +538,6 @@ def test_explicit_model_intermediate_prompt_replacement_releases_cache_and_water
|
||||
sched.encoder_cache_manager.free.assert_called_once_with(session)
|
||||
|
||||
|
||||
def test_talker_capacity_exact_fit_extends_from_declared_length_without_ids_prompt() -> None:
|
||||
sched = _make_scheduler(stage_id=1)
|
||||
adapter = _make_talker_adapter(max_model_len=100)
|
||||
sched.chunk_transfer_adapter = adapter
|
||||
session = _make_request()
|
||||
session.external_req_id = "external-capacity-exact-fit"
|
||||
session.prompt_token_ids = [0] * 70
|
||||
session._all_token_ids.clear()
|
||||
session._all_token_ids.extend(session.prompt_token_ids)
|
||||
session._output_token_ids.clear()
|
||||
session.append_output_token_ids([7, 8, 9, 10])
|
||||
session.num_prompt_tokens = 70
|
||||
session.num_computed_tokens = 74
|
||||
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
|
||||
sched.num_waiting_for_streaming_input = 1
|
||||
adapter._streaming_condition_lengths[session.request_id] = 20
|
||||
adapter._streaming_condition_seqs[session.request_id] = 0
|
||||
update = _make_talker_update(16, condition_seq=1)
|
||||
update.model_intermediate_buffer["ids"] = {"tts": [1, 2, 3]}
|
||||
|
||||
sched._update_request_as_session(session, update)
|
||||
|
||||
assert session.num_computed_tokens == 74
|
||||
assert session.num_prompt_tokens == 90
|
||||
assert list(session._all_token_ids[-20:]) == [7, 8, 9, 10] + [0] * 16
|
||||
assert update.model_intermediate_buffer["meta"] == {
|
||||
"next_stage_prompt_len": 16,
|
||||
"next_stage_generation_tokens": 10,
|
||||
"streaming_condition_seq": 1,
|
||||
"streaming_prompt_recompute": False,
|
||||
}
|
||||
assert adapter._streaming_condition_lengths[session.request_id] == 16
|
||||
assert adapter._streaming_condition_seqs[session.request_id] == 1
|
||||
assert session._omni_segment_generation == 1
|
||||
assert session.status == RequestStatus.WAITING
|
||||
assert sched.num_waiting_for_streaming_input == 0
|
||||
sched._free_request_blocks.assert_not_called()
|
||||
sched.encoder_cache_manager.free.assert_not_called()
|
||||
|
||||
|
||||
def test_talker_first_update_seeds_sender_condition_tracking() -> None:
|
||||
sched = _make_scheduler(stage_id=1)
|
||||
adapter = _make_talker_adapter(max_model_len=100)
|
||||
sched.chunk_transfer_adapter = adapter
|
||||
session = _make_request()
|
||||
session.external_req_id = "external-capacity-first-update"
|
||||
session.prompt_token_ids = [0] * 20
|
||||
session._all_token_ids.clear()
|
||||
session._all_token_ids.extend(session.prompt_token_ids)
|
||||
session._output_token_ids.clear()
|
||||
session.append_output_token_ids([7, 8, 9])
|
||||
session.num_prompt_tokens = 20
|
||||
session.num_computed_tokens = 23
|
||||
session.model_intermediate_buffer = {
|
||||
"native_duplex": True,
|
||||
"meta": {
|
||||
"next_stage_prompt_len": 20,
|
||||
"next_stage_generation_tokens": 10,
|
||||
"streaming_condition_seq": 0,
|
||||
},
|
||||
}
|
||||
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
|
||||
sched.num_waiting_for_streaming_input = 1
|
||||
update = _make_talker_update(16, condition_seq=1)
|
||||
|
||||
sched._update_request_as_session(session, update)
|
||||
|
||||
assert update.model_intermediate_buffer["meta"]["streaming_condition_seq"] == 1
|
||||
assert update.model_intermediate_buffer["meta"]["streaming_prompt_recompute"] is False
|
||||
assert adapter._streaming_condition_lengths[session.request_id] == 16
|
||||
assert adapter._streaming_condition_seqs[session.request_id] == 1
|
||||
assert session.num_computed_tokens == 23
|
||||
assert session.num_prompt_tokens == 39
|
||||
sched._free_request_blocks.assert_not_called()
|
||||
sched.encoder_cache_manager.free.assert_not_called()
|
||||
|
||||
|
||||
def test_talker_first_sliding_update_recomputes_from_seeded_condition() -> None:
|
||||
sched = _make_scheduler(stage_id=1)
|
||||
adapter = _make_talker_adapter(recompute_on_capacity=False)
|
||||
sched.chunk_transfer_adapter = adapter
|
||||
session = _make_request()
|
||||
session.external_req_id = "external-sliding-first-update"
|
||||
session.prompt_token_ids = [0] * 20
|
||||
session._all_token_ids.clear()
|
||||
session._all_token_ids.extend(session.prompt_token_ids)
|
||||
session._output_token_ids.clear()
|
||||
previous_codes = [7, 8, 9]
|
||||
session.append_output_token_ids(previous_codes)
|
||||
session.num_prompt_tokens = 20
|
||||
session.num_computed_tokens = 23
|
||||
session.model_intermediate_buffer = {
|
||||
"native_duplex": True,
|
||||
"meta": {
|
||||
"next_stage_prompt_len": 20,
|
||||
"next_stage_generation_tokens": 10,
|
||||
"streaming_condition_seq": 0,
|
||||
},
|
||||
}
|
||||
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
|
||||
sched.num_waiting_for_streaming_input = 1
|
||||
update = _make_talker_update(16, condition_seq=1)
|
||||
|
||||
sched._update_request_as_session(session, update)
|
||||
|
||||
assert session.num_computed_tokens == 0
|
||||
assert session.num_prompt_tokens == 39
|
||||
assert update.model_intermediate_buffer["ids"]["streaming_prompt_previous_codes"] == previous_codes
|
||||
assert update.model_intermediate_buffer["meta"]["streaming_prompt_recompute"] is True
|
||||
assert adapter._streaming_condition_lengths[session.request_id] == 16
|
||||
assert adapter._streaming_condition_seqs[session.request_id] == 1
|
||||
sched._free_request_blocks.assert_called_once_with(session)
|
||||
sched.encoder_cache_manager.free.assert_called_once_with(session)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("condition_seq", "error"),
|
||||
[
|
||||
(None, "missing streaming_condition_seq"),
|
||||
(0, "expected=1, received=0"),
|
||||
(2, "expected=1, received=2"),
|
||||
],
|
||||
)
|
||||
def test_talker_invalid_condition_sequence_does_not_advance_tracking(mocker, condition_seq, error) -> None:
|
||||
sched = _make_scheduler(stage_id=1)
|
||||
adapter = _make_talker_adapter()
|
||||
adapter.record_receive_failure = mocker.MagicMock()
|
||||
adapter._streaming_condition_lengths["req-ar-streaming-test"] = 20
|
||||
adapter._streaming_condition_seqs["req-ar-streaming-test"] = 0
|
||||
sched.chunk_transfer_adapter = adapter
|
||||
session = _make_request()
|
||||
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
|
||||
sched.num_waiting_for_streaming_input = 1
|
||||
update = _make_talker_update(16, condition_seq=condition_seq)
|
||||
|
||||
sched._update_request_as_session(session, update)
|
||||
|
||||
failure = adapter.record_receive_failure.call_args.args
|
||||
assert failure[0] == session.request_id
|
||||
assert error in failure[1]
|
||||
assert adapter._streaming_condition_lengths[session.request_id] == 20
|
||||
assert adapter._streaming_condition_seqs[session.request_id] == 0
|
||||
assert session.status == RequestStatus.WAITING_FOR_STREAMING_REQ
|
||||
|
||||
|
||||
def test_talker_first_update_that_cannot_fit_window_is_request_local(mocker) -> None:
|
||||
sched = _make_scheduler(stage_id=1)
|
||||
adapter = _make_talker_adapter()
|
||||
adapter.record_receive_failure = mocker.MagicMock()
|
||||
sched.chunk_transfer_adapter = adapter
|
||||
session = _make_request()
|
||||
session.prompt_token_ids = [0] * 20
|
||||
session._all_token_ids.clear()
|
||||
session._all_token_ids.extend(session.prompt_token_ids)
|
||||
session._output_token_ids.clear()
|
||||
session.append_output_token_ids(list(range(9)))
|
||||
session.num_prompt_tokens = 20
|
||||
session.num_computed_tokens = 29
|
||||
session.model_intermediate_buffer = {
|
||||
"native_duplex": True,
|
||||
"meta": {
|
||||
"next_stage_prompt_len": 20,
|
||||
"next_stage_generation_tokens": 10,
|
||||
"streaming_condition_seq": 0,
|
||||
},
|
||||
}
|
||||
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
|
||||
sched.num_waiting_for_streaming_input = 1
|
||||
update = _make_talker_update(62, condition_seq=1)
|
||||
|
||||
sched._update_request_as_session(session, update)
|
||||
|
||||
failure = adapter.record_receive_failure.call_args.args
|
||||
assert failure[0] == session.request_id
|
||||
assert "sliding streaming prompt plus generation reserve exceeds max_model_len" in failure[1]
|
||||
assert adapter._streaming_condition_lengths == {}
|
||||
assert adapter._streaming_condition_seqs == {}
|
||||
assert session.status == RequestStatus.WAITING_FOR_STREAMING_REQ
|
||||
|
||||
|
||||
def test_talker_capacity_overflow_recomputes_then_resumes_accumulation() -> None:
|
||||
sched = _make_scheduler(stage_id=1)
|
||||
adapter = _make_talker_adapter(max_model_len=100)
|
||||
sched.chunk_transfer_adapter = adapter
|
||||
session = _make_request()
|
||||
session.external_req_id = "external-capacity-rollover"
|
||||
session.prompt_token_ids = [0] * 70
|
||||
session._all_token_ids.clear()
|
||||
session._all_token_ids.extend(session.prompt_token_ids)
|
||||
session._output_token_ids.clear()
|
||||
previous_codes = [7, 8, 9, 10]
|
||||
session.append_output_token_ids(previous_codes)
|
||||
session.num_prompt_tokens = 70
|
||||
session.num_computed_tokens = 74
|
||||
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
|
||||
sched.num_waiting_for_streaming_input = 1
|
||||
adapter._streaming_condition_lengths[session.request_id] = 20
|
||||
adapter._streaming_condition_seqs[session.request_id] = 0
|
||||
adapter.segment_finished_requests.add(session.request_id)
|
||||
adapter.requests_num_chunks_sent[session.external_req_id] = 74
|
||||
overflow_update = _make_talker_update(17, condition_seq=1)
|
||||
|
||||
sched._update_request_as_session(session, overflow_update)
|
||||
|
||||
assert session.num_computed_tokens == 0
|
||||
assert session.num_prompt_tokens == 41
|
||||
assert session.prompt_token_ids == [0] * 41
|
||||
assert overflow_update.model_intermediate_buffer["ids"]["streaming_prompt_previous_codes"] == previous_codes
|
||||
assert overflow_update.model_intermediate_buffer["meta"]["streaming_condition_seq"] == 1
|
||||
assert overflow_update.model_intermediate_buffer["meta"]["streaming_prompt_recompute"] is True
|
||||
assert adapter.segment_finished_requests == set()
|
||||
assert adapter.requests_num_chunks_sent == {}
|
||||
assert session._omni_segment_generation == 1
|
||||
sched._free_request_blocks.assert_called_once_with(session)
|
||||
sched.encoder_cache_manager.free.assert_called_once_with(session)
|
||||
|
||||
# Once the replacement has been recomputed, later conditions append to
|
||||
# the new full-attention prefix until capacity is approached again.
|
||||
next_codes = [21, 22, 23, 24, 25, 26]
|
||||
session.append_output_token_ids(next_codes)
|
||||
session.num_computed_tokens = 47
|
||||
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
|
||||
sched.num_waiting_for_streaming_input = 1
|
||||
append_update = _make_talker_update(20, condition_seq=2)
|
||||
|
||||
sched._update_request_as_session(session, append_update)
|
||||
|
||||
assert session.num_computed_tokens == 47
|
||||
assert session.num_prompt_tokens == 67
|
||||
assert list(session._all_token_ids[-26:]) == next_codes + [0] * 20
|
||||
assert append_update.model_intermediate_buffer["meta"]["streaming_condition_seq"] == 2
|
||||
assert append_update.model_intermediate_buffer["meta"]["streaming_prompt_recompute"] is False
|
||||
assert session._omni_segment_generation == 2
|
||||
assert sched._free_request_blocks.call_count == 1
|
||||
assert sched.encoder_cache_manager.free.call_count == 1
|
||||
|
||||
# A later overflow replaces the accumulated prefix again rather than
|
||||
# reverting to recompute-on-every-condition behavior.
|
||||
second_rollover_codes = [31, 32, 33, 34, 35, 36]
|
||||
session.num_computed_tokens = session.num_prompt_tokens
|
||||
session.append_output_token_ids(second_rollover_codes)
|
||||
session.num_computed_tokens += len(second_rollover_codes)
|
||||
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
|
||||
sched.num_waiting_for_streaming_input = 1
|
||||
second_overflow_update = _make_talker_update(18, condition_seq=3)
|
||||
|
||||
sched._update_request_as_session(session, second_overflow_update)
|
||||
|
||||
assert session.num_computed_tokens == 0
|
||||
assert session.num_prompt_tokens == 44
|
||||
assert second_overflow_update.model_intermediate_buffer["ids"]["streaming_prompt_previous_codes"] == (
|
||||
second_rollover_codes
|
||||
)
|
||||
assert second_overflow_update.model_intermediate_buffer["meta"]["streaming_condition_seq"] == 3
|
||||
assert second_overflow_update.model_intermediate_buffer["meta"]["streaming_prompt_recompute"] is True
|
||||
assert session._omni_segment_generation == 3
|
||||
assert sched._free_request_blocks.call_count == 2
|
||||
assert sched.encoder_cache_manager.free.call_count == 2
|
||||
|
||||
|
||||
def test_talker_invalid_capacity_update_is_reported_per_request(mocker) -> None:
|
||||
sched = _make_scheduler(stage_id=1)
|
||||
adapter = _make_talker_adapter(max_model_len=100)
|
||||
adapter.record_receive_failure = mocker.MagicMock()
|
||||
sched.chunk_transfer_adapter = adapter
|
||||
session = _make_request()
|
||||
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
|
||||
sched.num_waiting_for_streaming_input = 1
|
||||
update = _make_talker_update(91, condition_seq=0)
|
||||
|
||||
sched._update_request_as_session(session, update)
|
||||
|
||||
adapter.record_receive_failure.assert_called_once_with(
|
||||
session.request_id,
|
||||
"fresh streaming prompt plus generation reserve exceeds max_model_len: prompt=91, reserve=10, limit=100",
|
||||
)
|
||||
assert session.status == RequestStatus.WAITING_FOR_STREAMING_REQ
|
||||
assert sched.num_waiting_for_streaming_input == 1
|
||||
sched._free_request_blocks.assert_not_called()
|
||||
sched.encoder_cache_manager.free.assert_not_called()
|
||||
|
||||
|
||||
def test_ready_async_chunk_prompt_replacement_releases_stale_kv_once() -> None:
|
||||
sched = _make_scheduler(stage_id=1)
|
||||
session = _make_request()
|
||||
|
||||
@@ -1,6 +1,3 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""
|
||||
Performance benchmark CI runner for diffusion models.
|
||||
|
||||
@@ -359,8 +356,6 @@ def _resolve_offline_model(model: str) -> str:
|
||||
"""
|
||||
import huggingface_hub
|
||||
|
||||
from vllm_omni.transformers_utils.repo_utils import hf_api
|
||||
|
||||
if not model or os.path.isdir(model):
|
||||
return model
|
||||
|
||||
@@ -385,7 +380,9 @@ def _resolve_offline_model(model: str) -> str:
|
||||
if len(parts) >= 3:
|
||||
repo_id = "/".join(parts[:2])
|
||||
subfolder = "/".join(parts[2:])
|
||||
snapshot_root = hf_api().snapshot_download(
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
snapshot_root = snapshot_download(
|
||||
repo_id,
|
||||
allow_patterns=[f"{subfolder}/**"],
|
||||
local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE,
|
||||
@@ -394,7 +391,9 @@ def _resolve_offline_model(model: str) -> str:
|
||||
|
||||
if not huggingface_hub.constants.HF_HUB_OFFLINE:
|
||||
return model
|
||||
return hf_api().snapshot_download(model, local_files_only=True)
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
return snapshot_download(model, local_files_only=True)
|
||||
|
||||
|
||||
class DiffusionServer:
|
||||
|
||||
@@ -1,5 +1,4 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""CPU contracts for capability-driven AR-Diffusion sessions."""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -11,7 +10,6 @@ from types import SimpleNamespace
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.sched.interface import CachedRequestData, DiffusionSchedulerOutput
|
||||
from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
|
||||
from vllm_omni.experimental.ar_diffusion.capability import (
|
||||
ARDiffusionCrossAttentionKVSpec,
|
||||
@@ -29,22 +27,6 @@ NEG = "negative"
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def platform_synchronize_calls(monkeypatch) -> list[None]:
|
||||
"""Stub the platform barrier the runner takes at every chunk boundary.
|
||||
|
||||
These tests run on a CPU host, where the resolved platform is
|
||||
``UnspecifiedOmniPlatform`` and ``synchronize()`` raises. Tests that need
|
||||
a failing barrier override this with their own fake.
|
||||
"""
|
||||
calls: list[None] = []
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.experimental.ar_diffusion.runner.current_omni_platform",
|
||||
SimpleNamespace(synchronize=lambda: calls.append(None)),
|
||||
)
|
||||
return calls
|
||||
|
||||
|
||||
def lingbot_like_spec(*, capacity: int = 2) -> ARDiffusionKVCacheSpec:
|
||||
"""Single-kv_branch causal DMD: 3 latent frames/block + sink/window + text KV."""
|
||||
return ARDiffusionKVCacheSpec(
|
||||
@@ -132,43 +114,6 @@ class BatchCapablePipeline(CapablePipeline):
|
||||
supports_request_batch = True
|
||||
|
||||
|
||||
class StepCapablePipeline(CapablePipeline):
|
||||
supports_step_execution = True
|
||||
|
||||
def prepare_encode(self, state, **kwargs):
|
||||
del kwargs
|
||||
return state
|
||||
|
||||
def denoise_step(self, input_batch, *, states=None, **kwargs):
|
||||
del input_batch, states, kwargs
|
||||
return None
|
||||
|
||||
def step_scheduler(self, state, noise_pred, **kwargs):
|
||||
del state, noise_pred, kwargs
|
||||
|
||||
def post_decode(self, state, **kwargs):
|
||||
del state, kwargs
|
||||
from vllm_omni.diffusion.data import DiffusionOutput
|
||||
|
||||
return DiffusionOutput()
|
||||
|
||||
|
||||
def scheduler_output(
|
||||
request_ids: tuple[str, ...] = ("req-1",),
|
||||
*,
|
||||
finished: tuple[str, ...] = (),
|
||||
) -> DiffusionSchedulerOutput:
|
||||
"""Build the real scheduler payload; stepwise requests arrive as cached rows."""
|
||||
return DiffusionSchedulerOutput(
|
||||
step_id=0,
|
||||
scheduled_new_reqs=[],
|
||||
scheduled_cached_reqs=CachedRequestData(request_ids=list(request_ids)),
|
||||
finished_req_ids=set(finished),
|
||||
num_running_reqs=len(request_ids),
|
||||
num_waiting_reqs=0,
|
||||
)
|
||||
|
||||
|
||||
def make_runner(
|
||||
pipeline: object,
|
||||
*,
|
||||
@@ -195,7 +140,6 @@ def make_runner(
|
||||
runner._sessions = OrderedDict()
|
||||
runner._session_capacity = 0
|
||||
runner._perf_e2e_times = []
|
||||
runner._stepwise_chunk_started = {}
|
||||
runner._preallocate_kv_cache(available_bytes=available_bytes)
|
||||
return runner
|
||||
|
||||
@@ -451,16 +395,12 @@ def test_synchronize_exception_uses_forward_cleanup_path(monkeypatch):
|
||||
ctx.ensure_video_slots(torch.device("cpu"))
|
||||
return object()
|
||||
|
||||
def synchronize_boom():
|
||||
def synchronize_boom(device):
|
||||
raise RuntimeError("asynchronous kernel failed")
|
||||
|
||||
monkeypatch.setattr(DiffusionModelRunner, "execute_model", return_after_allocation)
|
||||
# The runner defers the device barrier to the platform, so the failure is
|
||||
# injected there rather than on a device-specific torch entry point.
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.experimental.ar_diffusion.runner.current_omni_platform",
|
||||
SimpleNamespace(synchronize=synchronize_boom),
|
||||
)
|
||||
monkeypatch.setattr(torch.accelerator, "synchronize", synchronize_boom)
|
||||
runner.device = torch.device("cuda")
|
||||
request = SimpleNamespace(
|
||||
request_id="broken-request",
|
||||
sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}),
|
||||
@@ -484,11 +424,6 @@ def test_ar_runner_rejects_step_and_request_batch_modes():
|
||||
make_runner(BatchCapablePipeline(lingbot_like_spec()))
|
||||
|
||||
|
||||
def test_ar_runner_allows_step_execution_when_pipeline_implements_the_contract():
|
||||
runner = make_runner(StepCapablePipeline(lingbot_like_spec()), step_execution=True)
|
||||
assert runner.kv_cache is not None
|
||||
|
||||
|
||||
def test_ar_runner_defensively_rejects_inherited_batch_and_step_entrypoints():
|
||||
runner = object.__new__(ARDiffusionModelRunner)
|
||||
with pytest.raises(RuntimeError, match="request-batch execution"):
|
||||
@@ -497,144 +432,6 @@ def test_ar_runner_defensively_rejects_inherited_batch_and_step_entrypoints():
|
||||
runner.execute_stepwise(None)
|
||||
|
||||
|
||||
def test_execute_stepwise_binds_request_id_session_and_unbinds_after_call(monkeypatch):
|
||||
from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput
|
||||
|
||||
pipeline = StepCapablePipeline(lingbot_like_spec())
|
||||
runner = make_runner(pipeline, step_execution=True)
|
||||
bound_during = []
|
||||
|
||||
def fake_execute(self, scheduler_output):
|
||||
del self, scheduler_output
|
||||
bound_during.append(pipeline.bound_state is not None)
|
||||
return BatchRunnerOutput.from_list([RunnerOutput(request_id="req-1", finished=False)])
|
||||
|
||||
monkeypatch.setattr(DiffusionModelRunner, "execute_stepwise", fake_execute)
|
||||
output = runner.execute_stepwise(scheduler_output())
|
||||
|
||||
assert bound_during == [True]
|
||||
assert pipeline.bound_state is None
|
||||
assert pipeline.binds == ["req-1"]
|
||||
assert "req-1" in runner._sessions
|
||||
assert output.get_request_output("req-1") is not None
|
||||
assert not pipeline.closes
|
||||
|
||||
|
||||
def test_execute_stepwise_closes_session_when_request_finishes(monkeypatch):
|
||||
from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput
|
||||
|
||||
pipeline = StepCapablePipeline(lingbot_like_spec())
|
||||
runner = make_runner(pipeline, step_execution=True)
|
||||
|
||||
def fake_execute(self, scheduler_output):
|
||||
del self, scheduler_output
|
||||
return BatchRunnerOutput.from_list([RunnerOutput(request_id="req-1", finished=True)])
|
||||
|
||||
monkeypatch.setattr(DiffusionModelRunner, "execute_stepwise", fake_execute)
|
||||
runner.execute_stepwise(scheduler_output())
|
||||
|
||||
assert pipeline.bound_state is None
|
||||
assert pipeline.closes == ["req-1"]
|
||||
assert "req-1" not in runner._sessions
|
||||
|
||||
|
||||
def test_execute_stepwise_exception_releases_kv_and_does_not_resume_pages(monkeypatch):
|
||||
pipeline = StepCapablePipeline(lingbot_like_spec())
|
||||
runner = make_runner(pipeline, step_execution=True)
|
||||
first = runner._get_or_create_session("req-1")
|
||||
|
||||
def boom(self, scheduler_output):
|
||||
del self, scheduler_output
|
||||
raise RuntimeError("step failed")
|
||||
|
||||
monkeypatch.setattr(DiffusionModelRunner, "execute_stepwise", boom)
|
||||
with pytest.raises(RuntimeError, match="step failed"):
|
||||
runner.execute_stepwise(scheduler_output())
|
||||
|
||||
assert pipeline.bound_state is None
|
||||
assert pipeline.closes == ["req-1"]
|
||||
assert "req-1" not in runner._sessions
|
||||
second = runner._get_or_create_session("req-1")
|
||||
assert second is not first
|
||||
|
||||
|
||||
def test_execute_stepwise_closes_session_for_scheduler_aborted_request(monkeypatch):
|
||||
from vllm_omni.diffusion.worker.utils import BatchRunnerOutput
|
||||
|
||||
pipeline = StepCapablePipeline(lingbot_like_spec())
|
||||
runner = make_runner(pipeline, step_execution=True)
|
||||
runner._get_or_create_session("req-1")
|
||||
|
||||
def fake_execute(self, scheduler_output):
|
||||
del self, scheduler_output
|
||||
return BatchRunnerOutput.from_list([])
|
||||
|
||||
monkeypatch.setattr(DiffusionModelRunner, "execute_stepwise", fake_execute)
|
||||
# An aborted request is never scheduled again and never reports finished,
|
||||
# so the retire path is the only thing that can free its KV.
|
||||
runner.execute_stepwise(scheduler_output(request_ids=(), finished=("req-1",)))
|
||||
|
||||
assert "req-1" not in runner._sessions
|
||||
assert pipeline.closes == ["req-1"]
|
||||
|
||||
|
||||
def test_execute_stepwise_times_once_per_chunk_not_once_per_step(monkeypatch, platform_synchronize_calls):
|
||||
from vllm_omni.diffusion.data import DiffusionOutput
|
||||
from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput
|
||||
|
||||
pipeline = StepCapablePipeline(lingbot_like_spec())
|
||||
runner = make_runner(pipeline, step_execution=True)
|
||||
results = [None, None, None, DiffusionOutput()]
|
||||
|
||||
def fake_execute(self, scheduler_output):
|
||||
del self, scheduler_output
|
||||
return BatchRunnerOutput.from_list([RunnerOutput(request_id="req-1", finished=False, result=results.pop(0))])
|
||||
|
||||
monkeypatch.setattr(DiffusionModelRunner, "execute_stepwise", fake_execute)
|
||||
for _ in range(4):
|
||||
runner.execute_stepwise(scheduler_output())
|
||||
|
||||
# Four denoise steps, one emitted AR block: one timing sample, matching the
|
||||
# one-sample-per-block meaning request mode already has.
|
||||
assert len(runner._perf_e2e_times) == 1
|
||||
assert len(platform_synchronize_calls) == 1
|
||||
assert "req-1" not in runner._stepwise_chunk_started
|
||||
|
||||
|
||||
def test_execute_stepwise_synchronize_exception_fails_closed(monkeypatch):
|
||||
"""The chunk barrier is where an asynchronous device error surfaces. It must
|
||||
clean up like an in-step failure instead of leaving the session, the cached
|
||||
step state and the timing entry to the scheduler's next cycle."""
|
||||
from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput
|
||||
|
||||
pipeline = StepCapablePipeline(lingbot_like_spec())
|
||||
runner = make_runner(pipeline, step_execution=True)
|
||||
runner.state_cache = {"req-1": object()}
|
||||
|
||||
def fake_execute(self, scheduler_output):
|
||||
del self, scheduler_output
|
||||
return BatchRunnerOutput.from_list([RunnerOutput(request_id="req-1", finished=False, result=object())])
|
||||
|
||||
def synchronize_boom():
|
||||
raise RuntimeError("asynchronous kernel failed")
|
||||
|
||||
monkeypatch.setattr(DiffusionModelRunner, "execute_stepwise", fake_execute)
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.experimental.ar_diffusion.runner.current_omni_platform",
|
||||
SimpleNamespace(synchronize=synchronize_boom),
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="asynchronous kernel failed"):
|
||||
runner.execute_stepwise(scheduler_output())
|
||||
|
||||
assert pipeline.bound_state is None
|
||||
assert pipeline.closes == ["req-1"]
|
||||
assert "req-1" not in runner._sessions
|
||||
assert "req-1" not in runner.state_cache
|
||||
assert "req-1" not in runner._stepwise_chunk_started
|
||||
assert not runner._perf_e2e_times
|
||||
|
||||
|
||||
def test_model_specific_warmup_provider_is_consumed(monkeypatch):
|
||||
requests = [object(), object()]
|
||||
pipeline = WarmupPipeline(lingbot_like_spec(), requests)
|
||||
|
||||
@@ -25,9 +25,12 @@ from vllm_omni.experimental.ar_diffusion.kv_cache import (
|
||||
)
|
||||
from vllm_omni.experimental.ar_diffusion.kv_cache.paged import chunk_window_skipped_tokens
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
|
||||
BLOCK = 16
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def make_spec(*, chunk_size=BLOCK, window_chunks=2, sink_chunks=0, reset_at_boundary=False):
|
||||
@@ -152,155 +155,14 @@ def test_paged_pool_layout_exposes_flat_slot_views():
|
||||
dtype=torch.float32,
|
||||
device=torch.device("cpu"),
|
||||
)
|
||||
key_cache, value_cache = kv_pools[0]
|
||||
assert key_cache.shape == value_cache.shape == (4, BLOCK, 4, 64)
|
||||
assert kv_pools[0].shape == (2, 4, BLOCK, 4, 64)
|
||||
assert k_pools[0].shape == (4 * BLOCK, 4, 64)
|
||||
assert v_pools[0].shape == (4 * BLOCK, 4, 64)
|
||||
|
||||
# The flat views must alias their block-shaped cache, so a slot write is
|
||||
# visible through the layout the attention kernel reads.
|
||||
k_pools[0][BLOCK + 3].fill_(7)
|
||||
v_pools[0][2 * BLOCK + 5].fill_(11)
|
||||
assert torch.equal(key_cache[1, 3], k_pools[0][BLOCK + 3])
|
||||
assert torch.equal(value_cache[2, 5], v_pools[0][2 * BLOCK + 5])
|
||||
|
||||
|
||||
def test_key_and_value_caches_do_not_share_storage():
|
||||
"""K and V must be separate allocations, not halves of one tensor.
|
||||
|
||||
The paged-write custom op declares both as mutated. When they alias one
|
||||
storage, inductor's reinplace pass can re-inplace only the first of the
|
||||
two, and auto_functionalized_v2's clone of the entire second pool survives
|
||||
into the compiled graph -- one full pool copied per compiled region per
|
||||
denoising step.
|
||||
"""
|
||||
kv_pools, k_pools, v_pools = allocate_kv_pool_with_views(
|
||||
num_blocks=4,
|
||||
block_size=BLOCK,
|
||||
num_layers=2,
|
||||
num_kv_heads=4,
|
||||
head_dim=64,
|
||||
dtype=torch.float32,
|
||||
device=torch.device("cpu"),
|
||||
)
|
||||
seen = set()
|
||||
for layer, (key_cache, value_cache) in enumerate(kv_pools):
|
||||
k_storage = key_cache.untyped_storage().data_ptr()
|
||||
v_storage = value_cache.untyped_storage().data_ptr()
|
||||
assert k_storage != v_storage, f"layer {layer}: K and V share one allocation"
|
||||
assert k_pools[layer].untyped_storage().data_ptr() == k_storage
|
||||
assert v_pools[layer].untyped_storage().data_ptr() == v_storage
|
||||
seen.update((k_storage, v_storage))
|
||||
# Every layer's K and V are distinct allocations as well.
|
||||
assert len(seen) == 2 * len(kv_pools)
|
||||
|
||||
# Negative control: writing V must not disturb K.
|
||||
k_pools[0].fill_(0)
|
||||
v_pools[0].fill_(5)
|
||||
assert torch.equal(k_pools[0], torch.zeros_like(k_pools[0]))
|
||||
|
||||
|
||||
def _pool_sized_allocations(code: str, pool_numel: int) -> list[int]:
|
||||
"""Sizes of generated allocations at least as large as one pool."""
|
||||
import re
|
||||
|
||||
sizes = []
|
||||
for shape in re.findall(r"empty_strided_\w+\(\(([\d, ]*?)\)", code):
|
||||
dims = [int(part) for part in shape.replace(" ", "").strip(",").split(",") if part]
|
||||
numel = 1
|
||||
for dim_size in dims:
|
||||
numel *= dim_size
|
||||
if dims and numel >= pool_numel:
|
||||
sizes.append(numel)
|
||||
return sizes
|
||||
|
||||
|
||||
def _generated_code_for_pools(key_cache, value_cache):
|
||||
"""Compile a paged-write-shaped op over these pools, return the generated code."""
|
||||
from torch._inductor import config as inductor_config
|
||||
from torch._inductor.utils import run_and_get_code
|
||||
|
||||
class _Holder(torch.nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.register_buffer("key_cache", key_cache)
|
||||
self.register_buffer("value_cache", value_cache)
|
||||
|
||||
def forward(self, value):
|
||||
return torch.ops.test_kv_pool_clone.write(self.key_cache, self.value_cache, value)
|
||||
|
||||
torch._dynamo.reset()
|
||||
# Inductor caches compiled code across ``torch._dynamo.reset()``, so the
|
||||
# second call in a process would otherwise return the first one's code and
|
||||
# both branches of this test would describe the same graph.
|
||||
with inductor_config.patch(force_disable_caches=True):
|
||||
_, codes = run_and_get_code(torch.compile(_Holder(), backend="inductor", fullgraph=True), torch.ones(1))
|
||||
return "\n".join(codes)
|
||||
|
||||
|
||||
def test_compiled_paged_write_does_not_clone_the_pool():
|
||||
"""The compiled graph must write the pools in place, not through a clone.
|
||||
|
||||
``test_key_and_value_caches_do_not_share_storage`` pins the mechanism this
|
||||
change introduces -- two allocations rather than one. This pins the effect
|
||||
that mechanism buys, which is the reason the change exists: with K and V as
|
||||
halves of one tensor, inductor's reinplace pass re-inplaces only the first,
|
||||
and ``auto_functionalized_v2``'s clone of the pool survives into the
|
||||
generated code as a pool-sized allocation plus a copy, once per compiled
|
||||
region per denoising step.
|
||||
|
||||
The shared layout is compiled here as a positive control. Without it, a
|
||||
codegen change that stops emitting ``empty_strided_*`` would leave the
|
||||
detector matching nothing and this test passing for the wrong reason.
|
||||
"""
|
||||
if "write" not in dir(getattr(torch.ops, "test_kv_pool_clone", object())):
|
||||
|
||||
@torch.library.custom_op("test_kv_pool_clone::write", mutates_args=("key_pool", "value_pool"))
|
||||
def _write(key_pool: torch.Tensor, value_pool: torch.Tensor, value: torch.Tensor) -> torch.Tensor:
|
||||
key_pool[0, 0, 0, 0] = value[0]
|
||||
value_pool[0, 0, 0, 0] = value[0]
|
||||
return value * 2
|
||||
|
||||
@_write.register_fake
|
||||
def _(key_pool, value_pool, value):
|
||||
return torch.empty_like(value)
|
||||
|
||||
# The stand-in models the real op's contract; that declaration is what
|
||||
# drives the reinplace pass, so a drift here would silently test nothing.
|
||||
real_schema = torch.ops.vllm_omni.ar_diffusion_paged_write_attn.default._schema
|
||||
real_mutated = [arg.name for arg in real_schema.arguments if arg.alias_info and arg.alias_info.is_write]
|
||||
assert real_mutated == ["key_pool", "value_pool"], (
|
||||
f"the real paged-write op now mutates {real_mutated}; this test still models two pools"
|
||||
)
|
||||
|
||||
num_blocks, heads, dim = 8, 4, 16
|
||||
cache_shape = (num_blocks, BLOCK, heads, dim)
|
||||
pool_numel = num_blocks * BLOCK * heads * dim
|
||||
|
||||
# Positive control: one allocation, K and V aliasing it.
|
||||
shared = torch.empty(2, *cache_shape)
|
||||
shared_allocations = _pool_sized_allocations(_generated_code_for_pools(shared[0], shared[1]), pool_numel)
|
||||
assert shared_allocations, (
|
||||
"the shared layout produced no pool-sized allocation, so this test can no "
|
||||
"longer tell the two layouts apart -- inductor's codegen or the detector changed"
|
||||
)
|
||||
|
||||
# The allocator under test.
|
||||
kv_pools, _, _ = allocate_kv_pool_with_views(
|
||||
num_blocks=num_blocks,
|
||||
block_size=BLOCK,
|
||||
num_layers=1,
|
||||
num_kv_heads=heads,
|
||||
head_dim=dim,
|
||||
dtype=torch.float32,
|
||||
device=torch.device("cpu"),
|
||||
)
|
||||
key_cache, value_cache = kv_pools[0]
|
||||
separate_allocations = _pool_sized_allocations(_generated_code_for_pools(key_cache, value_cache), pool_numel)
|
||||
assert not separate_allocations, (
|
||||
"the compiled graph allocates a pool-sized buffer, so the pool is being cloned "
|
||||
f"rather than written in place: {separate_allocations} (control saw {shared_allocations})"
|
||||
)
|
||||
assert torch.equal(kv_pools[0][0, 1, 3], k_pools[0][BLOCK + 3])
|
||||
assert torch.equal(kv_pools[0][1, 2, 5], v_pools[0][2 * BLOCK + 5])
|
||||
|
||||
|
||||
def test_build_manager_allocate_free_roundtrip():
|
||||
|
||||
@@ -1,193 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""CUDA-only regression test for the AR-Diffusion paged KV write op.
|
||||
|
||||
Split out of ``test_kv_cache.py``: that module is swept into CI's
|
||||
``core_model and cpu`` job via a module-level ``pytest.mark.cpu``, and
|
||||
``hardware_test(..., num_cards=1)`` adds no ``skipif`` guard for a missing
|
||||
GPU, so this test was being collected (and would fail to import a CUDA
|
||||
device) on CPU-only runners. Keeping it in its own, non-``cpu`` module means
|
||||
it is only ever collected by a job that explicitly selects ``cuda`` for this
|
||||
path -- none does yet, so for now this is intentionally not wired into any
|
||||
CI job.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
from vllm_omni.experimental.ar_diffusion.kv_cache import allocate_kv_pool_with_views
|
||||
|
||||
BLOCK = 16
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_real_paged_write_op_does_not_clone_the_pool_on_cuda():
|
||||
"""Same assertion as the CPU test, against the real op's schema and CUDA codegen.
|
||||
|
||||
The CPU test compiles a stand-in with the same ``mutates_args`` contract,
|
||||
because the real op needs a device and an attention kernel. That leaves
|
||||
one thing unverified: whether Triton codegen reinplaces the same way,
|
||||
which is the path the reported regression was observed on.
|
||||
|
||||
This does not call the real op directly: ``ar_diffusion_paged_write_attn``
|
||||
runs FlashAttention's paged ``flash_attn_varlen_func`` internally, and
|
||||
compiling that under ``fullgraph=True`` can fail for reasons that have
|
||||
nothing to do with cloning (a missing kernel for this toy shape, a graph
|
||||
break, or an autotune illegal-memory-access on the control pool). Instead
|
||||
this registers a second stand-in that copies the real op's full argument
|
||||
list and ``mutates_args`` declaration -- so reinplace sees the identical
|
||||
contract the real op presents -- but skips straight to the pool writes,
|
||||
the same way the real op's own implementation does before it calls
|
||||
attention. That isolates "does Inductor clone this schema on CUDA" from
|
||||
"does FlashAttention compile", which is a separate, already-covered
|
||||
concern.
|
||||
|
||||
Note the pools passed here are the flat views, since that is what the op
|
||||
takes -- it indexes them by slot and unflattens internally.
|
||||
"""
|
||||
import re
|
||||
|
||||
from torch._inductor import config as inductor_config
|
||||
from torch._inductor.utils import run_and_get_code
|
||||
|
||||
device = torch.device("cuda")
|
||||
num_blocks, heads, dim, num_tokens = 8, 4, 64, 8
|
||||
pool_numel = num_blocks * BLOCK * heads * dim
|
||||
# The size a genuine clone would produce: K and V pool-sized allocations,
|
||||
# coalesced into one buffer the way `auto_functionalized_v2` was cloning
|
||||
# them before this PR's fix. Matching this exact size (rather than "any
|
||||
# allocation at least as large as one pool") avoids false positives from
|
||||
# unrelated workspace/scratch buffers Inductor may emit in that range.
|
||||
clone_numel = 2 * pool_numel
|
||||
|
||||
if "write" not in dir(getattr(torch.ops, "test_kv_pool_clone_schema", object())):
|
||||
|
||||
@torch.library.custom_op("test_kv_pool_clone_schema::write", mutates_args=("key_pool", "value_pool"))
|
||||
def _write(
|
||||
query: torch.Tensor,
|
||||
k_curr: torch.Tensor,
|
||||
v_curr: torch.Tensor,
|
||||
k_act: torch.Tensor | None,
|
||||
v_act: torch.Tensor | None,
|
||||
key_pool: torch.Tensor,
|
||||
value_pool: torch.Tensor,
|
||||
block_size: int,
|
||||
video_slots: torch.Tensor,
|
||||
action_slots: torch.Tensor,
|
||||
block_table: torch.Tensor,
|
||||
query_start_loc: torch.Tensor,
|
||||
seq_lens: torch.Tensor,
|
||||
max_query_len: int,
|
||||
max_seq_len: int,
|
||||
softmax_scale: float,
|
||||
) -> torch.Tensor:
|
||||
# Same pool-write side effect as `_paged_write_attn_impl`, minus
|
||||
# the `ar_diffusion_paged_attention` call -- that is the part
|
||||
# this test is not trying to exercise.
|
||||
key_pool[video_slots] = k_curr.to(key_pool.dtype)
|
||||
value_pool[video_slots] = v_curr.to(value_pool.dtype)
|
||||
if k_act is not None and v_act is not None and k_act.shape[0] > 0:
|
||||
key_pool[action_slots] = k_act.to(key_pool.dtype)
|
||||
value_pool[action_slots] = v_act.to(value_pool.dtype)
|
||||
return query * 2
|
||||
|
||||
@_write.register_fake
|
||||
def _(
|
||||
query,
|
||||
k_curr,
|
||||
v_curr,
|
||||
k_act,
|
||||
v_act,
|
||||
key_pool,
|
||||
value_pool,
|
||||
block_size,
|
||||
video_slots,
|
||||
action_slots,
|
||||
block_table,
|
||||
query_start_loc,
|
||||
seq_lens,
|
||||
max_query_len,
|
||||
max_seq_len,
|
||||
softmax_scale,
|
||||
):
|
||||
return torch.empty_like(query)
|
||||
|
||||
class _Holder(torch.nn.Module):
|
||||
def __init__(self, key_flat, value_flat):
|
||||
super().__init__()
|
||||
self.register_buffer("key_flat", key_flat)
|
||||
self.register_buffer("value_flat", value_flat)
|
||||
|
||||
def forward(self, query, k_curr, v_curr, video_slots, action_slots, block_table, query_start_loc, seq_lens):
|
||||
return torch.ops.test_kv_pool_clone_schema.write(
|
||||
query,
|
||||
k_curr,
|
||||
v_curr,
|
||||
None,
|
||||
None,
|
||||
self.key_flat,
|
||||
self.value_flat,
|
||||
BLOCK,
|
||||
video_slots,
|
||||
action_slots,
|
||||
block_table,
|
||||
query_start_loc,
|
||||
seq_lens,
|
||||
num_tokens,
|
||||
num_tokens,
|
||||
dim**-0.5,
|
||||
)
|
||||
|
||||
def compiled_pool_allocations(key_flat, value_flat) -> list[int]:
|
||||
args = (
|
||||
torch.randn(num_tokens, heads, dim, dtype=torch.bfloat16, device=device),
|
||||
torch.randn(num_tokens, heads, dim, dtype=torch.bfloat16, device=device),
|
||||
torch.randn(num_tokens, heads, dim, dtype=torch.bfloat16, device=device),
|
||||
torch.arange(num_tokens, device=device),
|
||||
torch.empty(0, dtype=torch.long, device=device),
|
||||
torch.arange(num_blocks, device=device, dtype=torch.int32).unsqueeze(0),
|
||||
torch.tensor([0, num_tokens], device=device, dtype=torch.int32),
|
||||
torch.tensor([num_tokens], device=device, dtype=torch.int32),
|
||||
)
|
||||
torch._dynamo.reset()
|
||||
# See the CPU test: inductor's cache outlives torch._dynamo.reset().
|
||||
with inductor_config.patch(force_disable_caches=True):
|
||||
_, codes = run_and_get_code(
|
||||
torch.compile(_Holder(key_flat, value_flat), backend="inductor", fullgraph=True), *args
|
||||
)
|
||||
sizes = []
|
||||
for shape in re.findall(r"empty_strided_cuda\(\(([\d, ]*?)\)", "\n".join(codes)):
|
||||
dims = [int(part) for part in shape.replace(" ", "").strip(",").split(",") if part]
|
||||
numel = 1
|
||||
for dim_size in dims:
|
||||
numel *= dim_size
|
||||
if dims and numel >= pool_numel:
|
||||
sizes.append(numel)
|
||||
return sizes
|
||||
|
||||
# Positive control: one allocation, K and V as its two halves. This must
|
||||
# produce an allocation of exactly `clone_numel` -- that is the clone
|
||||
# signature the assertion below checks for, not "no allocation this big".
|
||||
shared = torch.empty(2, num_blocks * BLOCK, heads, dim, dtype=torch.bfloat16, device=device)
|
||||
control = compiled_pool_allocations(shared[0], shared[1])
|
||||
assert clone_numel in control, (
|
||||
f"the shared layout produced no {clone_numel}-element allocation ({control}), so this test can no "
|
||||
"longer tell the two layouts apart -- inductor's codegen or the detector changed"
|
||||
)
|
||||
|
||||
_, k_pools, v_pools = allocate_kv_pool_with_views(
|
||||
num_blocks=num_blocks,
|
||||
block_size=BLOCK,
|
||||
num_layers=1,
|
||||
num_kv_heads=heads,
|
||||
head_dim=dim,
|
||||
dtype=torch.bfloat16,
|
||||
device=device,
|
||||
)
|
||||
separate = compiled_pool_allocations(k_pools[0], v_pools[0])
|
||||
assert clone_numel not in separate, (
|
||||
f"the compiled graph allocates a {clone_numel}-element buffer, so the pool is being cloned "
|
||||
f"rather than written in place: {separate} (control saw {control})"
|
||||
)
|
||||
@@ -1,174 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
import weakref
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from diffusers.models.autoencoders import AutoencoderKLWan
|
||||
|
||||
from vllm_omni.diffusion.distributed.autoencoders import wan_spatial_shard
|
||||
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import (
|
||||
DistributedAutoencoderKLWan,
|
||||
OmniAutoencoderKLWan,
|
||||
)
|
||||
from vllm_omni.diffusion.models.interface import supports_chunked_vae_decode
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def vae():
|
||||
# Real causal convolutions and caches, with small random weights for CPU CI.
|
||||
return OmniAutoencoderKLWan(
|
||||
base_dim=4,
|
||||
z_dim=2,
|
||||
dim_mult=[1, 2],
|
||||
num_res_blocks=1,
|
||||
temperal_downsample=[True],
|
||||
scale_factor_spatial=2,
|
||||
scale_factor_temporal=2,
|
||||
).eval()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("patch_size", [None, 2])
|
||||
@pytest.mark.parametrize("return_dict", [False, True])
|
||||
@torch.no_grad()
|
||||
def test_chunked_decode_matches_diffusers(vae, patch_size, return_dict):
|
||||
if patch_size is not None:
|
||||
vae = OmniAutoencoderKLWan.from_config(vae.config, patch_size=patch_size, out_channels=12).eval()
|
||||
z = torch.randn(1, 2, 3, 2, 2)
|
||||
expected = AutoencoderKLWan._decode(vae, z).sample
|
||||
actual = vae.decode(z, return_dict=return_dict)
|
||||
torch.testing.assert_close(actual.sample if return_dict else actual[0], expected, rtol=0, atol=0)
|
||||
chunks: list[torch.Tensor] = []
|
||||
assert supports_chunked_vae_decode(vae)
|
||||
assert not supports_chunked_vae_decode(object())
|
||||
assert vae.decode_with_chunks(z, on_chunk=chunks.append) is None
|
||||
assert [chunk.shape[2] for chunk in chunks] == [1, 2, 2]
|
||||
torch.testing.assert_close(torch.cat(chunks, dim=2), expected, rtol=0, atol=0)
|
||||
assert all(entry is None for entry in vae._feat_map)
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def test_slicing_and_forward_hooks(vae):
|
||||
z = torch.randn(2, 2, 2, 2, 2)
|
||||
vae.enable_slicing()
|
||||
batches = []
|
||||
handle = vae.post_quant_conv.register_forward_pre_hook(lambda module, args: batches.append(args[0].shape[0]))
|
||||
vae._hf_hook = SimpleNamespace(pre_forward=Mock())
|
||||
try:
|
||||
assert vae.decode(z).sample.shape[0] == 2
|
||||
assert batches == [1, 1]
|
||||
vae._hf_hook.pre_forward.assert_called_once_with(vae)
|
||||
batches.clear()
|
||||
with pytest.raises(ValueError, match="use_slicing"):
|
||||
vae.decode_with_chunks(z, on_chunk=Mock())
|
||||
assert batches == []
|
||||
vae._hf_hook.pre_forward.reset_mock()
|
||||
vae.decode_with_chunks(z[:1], on_chunk=Mock())
|
||||
assert batches == [1]
|
||||
vae._hf_hook.pre_forward.assert_called_once_with(vae)
|
||||
finally:
|
||||
handle.remove()
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def test_callback_failure_drains_and_clears_cache(vae):
|
||||
z = torch.randn(1, 2, 4, 2, 2)
|
||||
decoded = Mock(return_value=None)
|
||||
handle = vae.decoder.register_forward_hook(lambda *args: decoded())
|
||||
error = RuntimeError("consumer failed")
|
||||
callback = Mock(side_effect=error)
|
||||
try:
|
||||
with pytest.raises(RuntimeError, match="consumer failed") as exc:
|
||||
vae.decode_with_chunks(z, on_chunk=callback)
|
||||
assert exc.value is error
|
||||
assert callback.call_count == 1
|
||||
assert decoded.call_count == z.shape[2]
|
||||
assert all(entry is None for entry in vae._feat_map)
|
||||
finally:
|
||||
handle.remove()
|
||||
# The same instance can decode again after the failure.
|
||||
chunks: list[torch.Tensor] = []
|
||||
vae.decode_with_chunks(z, on_chunk=chunks.append)
|
||||
torch.testing.assert_close(torch.cat(chunks, 2), vae.decode(z).sample, rtol=0, atol=0)
|
||||
|
||||
|
||||
def test_chunked_decode_rejects_spatial_tiling(vae):
|
||||
vae.enable_tiling(tile_sample_min_height=2, tile_sample_min_width=2)
|
||||
with pytest.raises(ValueError, match="spatial tiling"):
|
||||
vae.decode_with_chunks(torch.zeros(1, 2, 2, 2, 2), on_chunk=Mock())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("rank", [0, 1])
|
||||
@pytest.mark.parametrize("fail_callback", [False, True])
|
||||
def test_spatial_callback_ownership_and_bounded_retention(monkeypatch, rank, fail_callback):
|
||||
# Isolate the temporal loop from halo kernels; weakrefs detect retention of
|
||||
# decoded GPU-sized tensors even after the consumer has failed.
|
||||
monkeypatch.setattr(wan_spatial_shard, "install_wan_spatial_shard_decode", lambda *args, **kwargs: None)
|
||||
monkeypatch.setattr(wan_spatial_shard, "_rank_world", lambda group: (rank, 2))
|
||||
references: list[weakref.ReferenceType[torch.Tensor]] = []
|
||||
|
||||
def decode(x, **kwargs):
|
||||
assert sum(ref() is not None for ref in references) <= 1
|
||||
chunk = x.clone()
|
||||
references.append(weakref.ref(chunk))
|
||||
return chunk
|
||||
|
||||
vae = SimpleNamespace(
|
||||
config=SimpleNamespace(patch_size=None),
|
||||
clear_cache=Mock(),
|
||||
post_quant_conv=lambda z: z,
|
||||
decoder=decode,
|
||||
_feat_map=[],
|
||||
)
|
||||
callback = Mock(side_effect=RuntimeError("consumer failed") if fail_callback else None)
|
||||
z = torch.zeros(1, 3, 6, 2, 2)
|
||||
if rank == 0 and fail_callback:
|
||||
with pytest.raises(RuntimeError, match="consumer failed"):
|
||||
wan_spatial_shard.spatial_shard_decode(vae, z, group=None, on_chunk=callback)
|
||||
else:
|
||||
assert wan_spatial_shard.spatial_shard_decode(vae, z, group=None, on_chunk=callback) is None
|
||||
assert len(references) == z.shape[2]
|
||||
assert callback.call_count == (0 if rank else 1 if fail_callback else z.shape[2])
|
||||
assert vae.clear_cache.call_count == 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize("rank", [0, 1])
|
||||
@torch.no_grad()
|
||||
def test_distributed_small_latent_callback_ownership(vae, monkeypatch, rank):
|
||||
vae.is_distributed_enabled = lambda: True
|
||||
vae.distributed_executor = SimpleNamespace(group=None)
|
||||
monkeypatch.setattr(torch.distributed, "get_rank", lambda **kwargs: rank)
|
||||
callback = Mock()
|
||||
vae.decode_with_chunks(torch.zeros(1, 2, 2, 2, 2), on_chunk=callback)
|
||||
assert callback.call_count == (2 if rank == 0 else 0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", ["spatial_shard_height", "spatial_shard_width", "tile"])
|
||||
def test_distributed_chunked_routing(monkeypatch, mode):
|
||||
vae = DistributedAutoencoderKLWan.__new__(DistributedAutoencoderKLWan)
|
||||
torch.nn.Module.__init__(vae)
|
||||
vae.use_tiling = True
|
||||
vae.use_slicing = False
|
||||
vae.tile_sample_min_height = vae.tile_sample_min_width = 2
|
||||
vae.spatial_compression_ratio = 2
|
||||
vae.distributed_executor = SimpleNamespace(group=None, parallel_mode=mode, parallel_size=2)
|
||||
monkeypatch.setattr(vae, "is_distributed_enabled", lambda: True)
|
||||
monkeypatch.setattr(torch.distributed, "get_world_size", lambda **kwargs: 2)
|
||||
decode = Mock()
|
||||
monkeypatch.setattr(wan_spatial_shard, "spatial_shard_decode", decode)
|
||||
callback = Mock()
|
||||
z = torch.zeros(1, 2, 2, 2, 2)
|
||||
if mode == "tile":
|
||||
with pytest.raises(ValueError, match="require spatial-shard"):
|
||||
vae.decode_with_chunks(z, on_chunk=callback)
|
||||
decode.assert_not_called()
|
||||
else:
|
||||
assert vae.decode_with_chunks(z, on_chunk=callback) is None
|
||||
decode.assert_called_once_with(
|
||||
vae, z, group=None, return_dict=True, split_dim=mode.removeprefix("spatial_shard_"), on_chunk=callback
|
||||
)
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Shared test helpers for diffusion LoRA tests."""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -26,39 +26,15 @@ class DummyBaseLayerWithLoRA(torch.nn.Module):
|
||||
tuple[list[torch.Tensor | None] | torch.Tensor, list[torch.Tensor | None] | torch.Tensor]
|
||||
] = []
|
||||
self.reset_calls: int = 0
|
||||
self.suspend_calls: int = 0
|
||||
self.resume_calls: int = 0
|
||||
self.active_slices: tuple[bool, ...] = ()
|
||||
self.suspended_slices: tuple[bool, ...] | None = None
|
||||
self.create_calls: int = 0
|
||||
|
||||
def set_lora(self, index: int, lora_a, lora_b):
|
||||
assert index == 0
|
||||
self.set_calls.append((lora_a, lora_b))
|
||||
if isinstance(lora_b, list):
|
||||
self.active_slices = tuple(b is not None for b in lora_b)
|
||||
else:
|
||||
self.active_slices = (True,)
|
||||
self.suspended_slices = None
|
||||
|
||||
def reset_lora(self, index: int):
|
||||
assert index == 0
|
||||
self.reset_calls += 1
|
||||
self.active_slices = ()
|
||||
self.suspended_slices = None
|
||||
|
||||
def suspend_lora(self) -> None:
|
||||
if self.suspended_slices is not None:
|
||||
return
|
||||
self.suspended_slices = self.active_slices
|
||||
self.active_slices = (False,) * len(self.active_slices)
|
||||
self.suspend_calls += 1
|
||||
|
||||
def resume_lora(self) -> None:
|
||||
if self.suspended_slices is not None:
|
||||
self.active_slices = self.suspended_slices
|
||||
self.suspended_slices = None
|
||||
self.resume_calls += 1
|
||||
|
||||
def create_lora_weights(self, max_loras, lora_config, model_config):
|
||||
self.create_calls += 1
|
||||
|
||||
@@ -30,38 +30,14 @@ class _DummyLoRALayer:
|
||||
tuple[list[torch.Tensor | None] | torch.Tensor, list[torch.Tensor | None] | torch.Tensor]
|
||||
] = []
|
||||
self.reset_calls: int = 0
|
||||
self.suspend_calls: int = 0
|
||||
self.resume_calls: int = 0
|
||||
self.active_slices: tuple[bool, ...] = ()
|
||||
self.suspended_slices: tuple[bool, ...] | None = None
|
||||
|
||||
def set_lora(self, index: int, lora_a, lora_b):
|
||||
assert index == 0
|
||||
self.set_calls.append((lora_a, lora_b))
|
||||
if isinstance(lora_b, list):
|
||||
self.active_slices = tuple(b is not None for b in lora_b)
|
||||
else:
|
||||
self.active_slices = (True,)
|
||||
self.suspended_slices = None
|
||||
|
||||
def reset_lora(self, index: int):
|
||||
assert index == 0
|
||||
self.reset_calls += 1
|
||||
self.active_slices = ()
|
||||
self.suspended_slices = None
|
||||
|
||||
def suspend_lora(self) -> None:
|
||||
if self.suspended_slices is not None:
|
||||
return
|
||||
self.suspended_slices = self.active_slices
|
||||
self.active_slices = (False,) * len(self.active_slices)
|
||||
self.suspend_calls += 1
|
||||
|
||||
def resume_lora(self) -> None:
|
||||
if self.suspended_slices is not None:
|
||||
self.active_slices = self.suspended_slices
|
||||
self.suspended_slices = None
|
||||
self.resume_calls += 1
|
||||
|
||||
|
||||
# Aliases for backward compatibility within this file
|
||||
@@ -747,104 +723,3 @@ def test_lora_manager_discovers_unet_component(monkeypatch):
|
||||
assert "unet.down_block.proj" in manager._lora_modules
|
||||
# Verify the module was actually replaced in the tree (not just recorded)
|
||||
assert isinstance(pipeline.unet.down_block.proj, _DummyBaseLayerWithLoRA)
|
||||
|
||||
|
||||
def _suspend_harness(monkeypatch, adapter_id: int = 7, rank: int = 2):
|
||||
"""Manager wired to a single dummy layer, ready to activate `adapter_id`."""
|
||||
import vllm_omni.diffusion.lora.manager as manager_mod
|
||||
|
||||
monkeypatch.setattr(manager_mod, "BaseLayerWithLoRA", _DummyBaseLayerWithLoRA)
|
||||
|
||||
lora_model = _DummyLM(rank=rank)
|
||||
manager = DiffusionLoRAManager(
|
||||
pipeline=_DummyPipeline(),
|
||||
device=torch.device("cpu"),
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
monkeypatch.setattr(manager, "_load_adapter", lambda _req: (lora_model, type("PH", (), {"r": rank})()))
|
||||
manager._registered_adapters = {adapter_id: lora_model}
|
||||
manager._lora_modules = {"transformer.foo": lora_model.transformer.foo}
|
||||
return manager, lora_model.transformer.foo
|
||||
|
||||
|
||||
def test_reactivating_same_adapter_resumes_without_rebinding(monkeypatch):
|
||||
"""activate -> deactivate -> activate at the same scale must re-arm the
|
||||
saved mask instead of re-uploading every layer."""
|
||||
adapter_id = 7
|
||||
req = _dummy_lora_request(adapter_id)
|
||||
manager, layer = _suspend_harness(monkeypatch, adapter_id)
|
||||
|
||||
manager.set_active_adapter(req, lora_scale=0.5)
|
||||
assert len(layer.set_calls) == 1
|
||||
assert layer.active_slices == (True,)
|
||||
|
||||
manager.set_active_adapter(None)
|
||||
assert manager._suspended_adapter_id == adapter_id
|
||||
assert layer.suspend_calls == 1
|
||||
assert layer.active_slices == (False,)
|
||||
# Suspending must not tear the upload down.
|
||||
assert layer.reset_calls == 0
|
||||
|
||||
manager.set_active_adapter(req, lora_scale=0.5)
|
||||
assert layer.resume_calls == 1
|
||||
assert len(layer.set_calls) == 1, "resume must not rebind"
|
||||
assert layer.active_slices == (True,)
|
||||
assert manager._active_adapter_id == adapter_id
|
||||
assert manager._suspended_adapter_id is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("second_scale", [0.5, 0.25])
|
||||
def test_scale_or_adapter_change_forces_rebind(monkeypatch, second_scale):
|
||||
"""A different scale, or a different adapter, must take the full bind
|
||||
path rather than re-arming a stale mask."""
|
||||
adapter_id = 7
|
||||
manager, layer = _suspend_harness(monkeypatch, adapter_id)
|
||||
|
||||
manager.set_active_adapter(_dummy_lora_request(adapter_id), lora_scale=0.5)
|
||||
manager.set_active_adapter(None)
|
||||
assert manager._suspended_adapter_id == adapter_id
|
||||
|
||||
# Same id at a new scale, or a different id entirely.
|
||||
other_id = adapter_id if second_scale != 0.5 else adapter_id + 1
|
||||
if other_id != adapter_id:
|
||||
manager._registered_adapters[other_id] = manager._registered_adapters[adapter_id]
|
||||
|
||||
manager.set_active_adapter(_dummy_lora_request(other_id), lora_scale=second_scale)
|
||||
assert len(layer.set_calls) == 2, "must rebind"
|
||||
assert manager._suspended_adapter_id is None
|
||||
assert manager._active_adapter_id == other_id
|
||||
|
||||
|
||||
def test_partial_packed_mask_survives_suspend_resume(monkeypatch):
|
||||
"""A packed layer where only some slices carry LoRA must come back with
|
||||
the same mask, not an all-True one."""
|
||||
adapter_id = 7
|
||||
manager, layer = _suspend_harness(monkeypatch, adapter_id)
|
||||
|
||||
# Emulate a packed bind that left slice 1 empty.
|
||||
layer.set_lora(0, [torch.ones(2, 2), None], [torch.ones(2, 2), None])
|
||||
manager._active_adapter_id = adapter_id
|
||||
manager._update_adapter_scale(adapter_id, 0.5)
|
||||
assert layer.active_slices == (True, False)
|
||||
|
||||
manager.set_active_adapter(None)
|
||||
assert layer.active_slices == (False, False)
|
||||
|
||||
manager.set_active_adapter(_dummy_lora_request(adapter_id), lora_scale=0.5)
|
||||
assert layer.active_slices == (True, False), "mask must survive unchanged"
|
||||
|
||||
|
||||
def test_removing_suspended_adapter_drops_the_upload(monkeypatch):
|
||||
"""A removed adapter can never be resumed, so its weights must not stay
|
||||
in the stacked buffers."""
|
||||
adapter_id = 7
|
||||
manager, layer = _suspend_harness(monkeypatch, adapter_id)
|
||||
|
||||
manager.set_active_adapter(_dummy_lora_request(adapter_id), lora_scale=0.5)
|
||||
manager.set_active_adapter(None)
|
||||
assert manager._suspended_adapter_id == adapter_id
|
||||
|
||||
manager.remove_adapter(adapter_id)
|
||||
assert manager._suspended_adapter_id is None
|
||||
assert layer.reset_calls == 1, "the upload must be torn down"
|
||||
assert layer.suspended_slices is None
|
||||
@@ -13,6 +13,7 @@ from unittest.mock import MagicMock
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from huggingface_hub import snapshot_download
|
||||
from safetensors.torch import save_file
|
||||
from vllm.config.load import LoadConfig
|
||||
|
||||
@@ -31,7 +32,6 @@ from vllm_omni.diffusion.models.helios import HeliosPipeline
|
||||
from vllm_omni.diffusion.models.host_weight_contract import FinalLayoutModelContract
|
||||
from vllm_omni.diffusion.registry import initialize_model
|
||||
from vllm_omni.quantization.component_config import ComponentQuantizationConfig
|
||||
from vllm_omni.transformers_utils.repo_utils import hf_api
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
@@ -41,7 +41,7 @@ model_path = "hf-internal-testing/tiny-helios-modular-pipe"
|
||||
@pytest.fixture(scope="module")
|
||||
def prefetch_helios_model():
|
||||
"""Downloads the tiny helios model prior to running a test."""
|
||||
hf_api().snapshot_download(model_path)
|
||||
snapshot_download(model_path)
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import contextlib
|
||||
|
||||
@@ -14,10 +14,10 @@ pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
def test_prefetch_subfolders_propagates_revision(monkeypatch):
|
||||
calls = []
|
||||
|
||||
def fake_snapshot_download(self, **kwargs):
|
||||
def fake_snapshot_download(**kwargs):
|
||||
calls.append(kwargs)
|
||||
|
||||
monkeypatch.setattr(huggingface_hub.HfApi, "snapshot_download", fake_snapshot_download)
|
||||
monkeypatch.setattr(huggingface_hub, "snapshot_download", fake_snapshot_download)
|
||||
monkeypatch.setattr(hub_prefetch, "_repo_prefetch_lock", lambda _model: contextlib.nullcontext())
|
||||
|
||||
hub_prefetch.prefetch_subfolders(
|
||||
|
||||
@@ -1,30 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.models.utils import make_attention_mask
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("seq_lengths", [[16], [16, 16]])
|
||||
def test_attention_mask_is_skipped_for_dense_batch(seq_lengths):
|
||||
hidden_states = torch.empty(len(seq_lengths), 16, 64)
|
||||
|
||||
assert make_attention_mask(hidden_states, seq_lengths) is None
|
||||
|
||||
|
||||
def test_attention_mask_preserves_variable_length_padding():
|
||||
hidden_states = torch.empty(2, 16, 64)
|
||||
|
||||
attention_mask = make_attention_mask(hidden_states, [16, 12])
|
||||
|
||||
expected = torch.tensor(
|
||||
[
|
||||
[True] * 16,
|
||||
[True] * 12 + [False] * 4,
|
||||
]
|
||||
)
|
||||
torch.testing.assert_close(attention_mask, expected)
|
||||
@@ -1,96 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""L1 tests for checkpoint-to-pipeline routing of Boogu-Image Turbo."""
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
|
||||
def test_turbo_registry_entry_loads_the_turbo_class():
|
||||
from vllm_omni.diffusion.models.boogu_image import BooguImageTurboPipeline
|
||||
from vllm_omni.diffusion.registry import DiffusionModelRegistry
|
||||
|
||||
assert DiffusionModelRegistry._try_load_model_cls("BooguImageTurboPipeline") is BooguImageTurboPipeline
|
||||
|
||||
|
||||
def test_turbo_entry_disables_request_batching():
|
||||
from vllm_omni.diffusion.models.boogu_image import BooguImagePipeline, BooguImageTurboPipeline
|
||||
|
||||
assert BooguImagePipeline.supports_request_batch
|
||||
assert not BooguImageTurboPipeline.supports_request_batch
|
||||
|
||||
|
||||
def test_turbo_process_funcs_resolve_to_the_shared_boogu_helpers():
|
||||
from vllm_omni.diffusion.models import boogu_image
|
||||
from vllm_omni.diffusion.registry import _DIFFUSION_POST_PROCESS_FUNCS, _DIFFUSION_PRE_PROCESS_FUNCS
|
||||
|
||||
module = boogu_image.pipeline_boogu_image
|
||||
for funcs in (_DIFFUSION_PRE_PROCESS_FUNCS, _DIFFUSION_POST_PROCESS_FUNCS):
|
||||
name = funcs["BooguImageTurboPipeline"]
|
||||
assert name == funcs["BooguImagePipeline"]
|
||||
assert callable(getattr(module, name))
|
||||
|
||||
|
||||
def test_turbo_inherits_the_base_multimodal_metadata():
|
||||
from vllm_omni.diffusion.model_metadata import get_diffusion_model_metadata
|
||||
|
||||
metadata = get_diffusion_model_metadata("BooguImageTurboPipeline")
|
||||
|
||||
assert metadata == get_diffusion_model_metadata("BooguImagePipeline")
|
||||
assert metadata.supports_multimodal_inputs
|
||||
assert metadata.max_multimodal_image_inputs == 1
|
||||
|
||||
|
||||
def test_turbo_model_index_resolves_and_enriches(tmp_path):
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig, resolve_model_class_name
|
||||
|
||||
(tmp_path / "model_index.json").write_text(
|
||||
json.dumps({"_class_name": "BooguImageTurboPipeline"}),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
assert resolve_model_class_name(str(tmp_path)) == "BooguImageTurboPipeline"
|
||||
|
||||
config = OmniDiffusionConfig(model=str(tmp_path))
|
||||
config.enrich_config()
|
||||
|
||||
assert config.model_class_name == "BooguImageTurboPipeline"
|
||||
assert config.supports_multimodal_inputs
|
||||
assert config.max_multimodal_image_inputs == 1
|
||||
|
||||
|
||||
def test_edit_turbo_override_selects_the_turbo_class(tmp_path):
|
||||
"""Edit-Turbo publishes ``BooguImagePipeline``; ``--model-class-name`` must win."""
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.registry import DiffusionModelRegistry
|
||||
|
||||
(tmp_path / "model_index.json").write_text(
|
||||
json.dumps({"_class_name": "BooguImagePipeline"}),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
config = OmniDiffusionConfig(model=str(tmp_path), model_class_name="BooguImageTurboPipeline")
|
||||
config.enrich_config()
|
||||
|
||||
assert config.model_class_name == "BooguImageTurboPipeline"
|
||||
assert DiffusionModelRegistry._try_load_model_cls(config.model_class_name)._is_turbo
|
||||
|
||||
|
||||
def test_edit_turbo_without_override_keeps_the_base_class(tmp_path):
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.registry import DiffusionModelRegistry
|
||||
|
||||
(tmp_path / "model_index.json").write_text(
|
||||
json.dumps({"_class_name": "BooguImagePipeline"}),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
config = OmniDiffusionConfig(model=str(tmp_path))
|
||||
config.enrich_config()
|
||||
|
||||
assert config.model_class_name == "BooguImagePipeline"
|
||||
assert not DiffusionModelRegistry._try_load_model_cls(config.model_class_name)._is_turbo
|
||||
@@ -1180,21 +1180,6 @@ def test_transfer_fps_matches_resolved_frame_rate_precedence() -> None:
|
||||
assert cfg.fps == sp.resolved_frame_rate == 12.0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("use_path", [False, True], ids=["image", "path"])
|
||||
def test_transfer_pil_conversion_returns_writable_array(tmp_path, use_path: bool) -> None:
|
||||
from vllm_omni.diffusion.models.cosmos3 import transfer
|
||||
|
||||
image = Image.new("RGB", (5, 4), "red")
|
||||
value = image
|
||||
if use_path:
|
||||
value = tmp_path / "control.png"
|
||||
image.save(value)
|
||||
|
||||
array = transfer._pil_to_uint8_rgb(value)
|
||||
|
||||
assert array.flags.writeable
|
||||
|
||||
|
||||
def test_transfer_vae_executor_requires_distributed_vae() -> None:
|
||||
from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniDiffusersPipeline
|
||||
|
||||
@@ -2401,42 +2386,6 @@ def test_forward_transfer_runs_multichunk_overlap_path(
|
||||
torch.testing.assert_close(captured["targets"][1][:, :, 0], torch.full((1, 3, 16, 16), -0.2))
|
||||
|
||||
|
||||
def test_forward_transfer_non_output_rank_uses_canonical_envelope(
|
||||
make_cosmos3_pipeline,
|
||||
sequential_cfg_parallel,
|
||||
) -> None:
|
||||
pipeline = make_cosmos3_pipeline()
|
||||
pipeline.vae.distributed_executor = SimpleNamespace(rank=1)
|
||||
pipeline.vae.is_distributed_enabled = lambda: True
|
||||
pipeline._transfer_bucket_size = lambda sp, source_hw: (16, 16, "1,1")
|
||||
pipeline._format_and_tokenize_prompts = lambda *args, **kwargs: (_ids(2), _mask(), _ids(1), _mask())
|
||||
pipeline._set_flow_shift = lambda *_args, **_kwargs: None
|
||||
decoded = torch.zeros(1, 3, 1, 16, 16)
|
||||
pipeline._decode_latents = lambda latents: decoded
|
||||
|
||||
request = SimpleNamespace(
|
||||
prompts=[{"prompt": "transfer", "modalities": ["video"]}],
|
||||
sampling_params=make_sampling_params(
|
||||
height=16,
|
||||
width=16,
|
||||
num_inference_steps=1,
|
||||
guidance_scale=1.0,
|
||||
extra_args={
|
||||
"edge": {"control": torch.zeros(3, 1, 16, 16, dtype=torch.uint8)},
|
||||
"max_frames": 1,
|
||||
"num_video_frames_per_chunk": 1,
|
||||
},
|
||||
),
|
||||
)
|
||||
|
||||
output = pipeline.forward(request)
|
||||
|
||||
assert set(output.output) == {"payload", "metadata"}
|
||||
assert set(output.output["payload"]) == {"video"}
|
||||
torch.testing.assert_close(output.output["payload"]["video"], decoded)
|
||||
assert output.output["metadata"] == {"video": {"fps": 24.0}}
|
||||
|
||||
|
||||
def test_diffuse_keeps_paired_cfg_when_cache_dit_active(make_cosmos3_pipeline) -> None:
|
||||
"""With cache-dit active the uncond pass is kept even outside the guidance
|
||||
interval (so cache-dit's has_separate_cfg parity stays in phase), and the
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -81,8 +81,9 @@ def test_from_config_loads_local_diffusers_component(tmp_path, monkeypatch: pyte
|
||||
|
||||
|
||||
def test_from_config_downloads_component_from_hf_repo(tmp_path, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
import huggingface_hub
|
||||
|
||||
from vllm_omni.diffusion.models.cosmos3 import sound_tokenizer
|
||||
from vllm_omni.transformers_utils import repo_utils
|
||||
|
||||
cache_dir = tmp_path / "hf"
|
||||
_write_component(cache_dir, checkpoint_name=DIFFUSERS_SOUND_TOKENIZER_CHECKPOINT_NAME)
|
||||
@@ -94,7 +95,7 @@ def test_from_config_downloads_component_from_hf_repo(tmp_path, monkeypatch: pyt
|
||||
calls.append((repo_id, revision, allow_patterns))
|
||||
return str(cache_dir)
|
||||
|
||||
monkeypatch.setattr(repo_utils.hf_api(), "snapshot_download", fake_snapshot_download)
|
||||
monkeypatch.setattr(huggingface_hub, "snapshot_download", fake_snapshot_download)
|
||||
|
||||
sound_tokenizer.Cosmos3SoundTokenizer.from_config(
|
||||
SimpleNamespace(
|
||||
|
||||
@@ -383,117 +383,6 @@ def test_transformer_sharding_offload_and_patch_round_trip_contracts() -> None:
|
||||
)
|
||||
|
||||
|
||||
def test_gen_sp_plan_auto_pads_hidden_states_and_rope_together() -> None:
|
||||
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer
|
||||
|
||||
prepare_plan = Cosmos3VFMTransformer._sp_plan["gen_sp_prepare"]
|
||||
assert set(prepare_plan) == {0, 1, 2}
|
||||
assert all(spec.auto_pad for spec in prepare_plan.values())
|
||||
assert [prepare_plan[index].split_dim for index in range(3)] == [1, 1, 1]
|
||||
|
||||
|
||||
def test_gen_sp_prepare_resets_branch_local_padding_metadata() -> None:
|
||||
from vllm_omni.diffusion.forward_context import ForwardContext, override_forward_context
|
||||
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3GenSPPrepare
|
||||
|
||||
ctx = ForwardContext(sp_padding_size=1, sp_original_seq_len=12121)
|
||||
hidden = torch.zeros(1, 4, 8)
|
||||
rope = torch.zeros(1, 4, 1, 2)
|
||||
with override_forward_context(ctx):
|
||||
output = Cosmos3GenSPPrepare()(hidden, rope, rope)
|
||||
|
||||
assert output[0] is hidden
|
||||
assert output[1] is rope
|
||||
assert output[2] is rope
|
||||
assert ctx.sp_padding_size == 0
|
||||
assert ctx.sp_original_seq_len is None
|
||||
|
||||
|
||||
def test_sp_attention_masks_padded_gen_and_joint_und_keys() -> None:
|
||||
from vllm_omni.diffusion.forward_context import ForwardContext, override_forward_context
|
||||
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3CrossAttention
|
||||
|
||||
class RecordingAttention(nn.Module):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.metadata: Any | None = None
|
||||
|
||||
def forward(self, query, key, value, metadata):
|
||||
del key, value
|
||||
self.metadata = metadata
|
||||
return query
|
||||
|
||||
module = object.__new__(Cosmos3CrossAttention)
|
||||
nn.Module.__init__(module)
|
||||
module.num_heads_local = 2
|
||||
module.head_dim = 4
|
||||
recording_attention = RecordingAttention()
|
||||
module.attn = recording_attention
|
||||
|
||||
q = torch.zeros(1, 2, 2, 4)
|
||||
k = torch.zeros_like(q)
|
||||
v = torch.zeros_like(q)
|
||||
k_und = torch.zeros(1, 3, 2, 4)
|
||||
v_und = torch.zeros_like(k_und)
|
||||
ctx = ForwardContext(sp_padding_size=1, sp_original_seq_len=3)
|
||||
with override_forward_context(ctx):
|
||||
output = module._forward_sp(q, k, v, k_und, v_und)
|
||||
|
||||
assert output.shape == (1, 2, 8)
|
||||
assert recording_attention.metadata is not None
|
||||
assert recording_attention.metadata.attn_mask.tolist() == [[True, True, True, False]]
|
||||
assert recording_attention.metadata.joint_attn_mask.tolist() == [[True, True, True]]
|
||||
|
||||
|
||||
def test_ulysses_2d_mask_uses_key_length_with_kv_only_joint_context(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
|
||||
from vllm_omni.diffusion.attention.parallel import ulysses
|
||||
|
||||
def fake_all_to_all(
|
||||
_group: object,
|
||||
tensor: torch.Tensor,
|
||||
_scatter_idx: int,
|
||||
_gather_idx: int,
|
||||
_use_sync: bool,
|
||||
) -> torch.Tensor:
|
||||
return torch.cat([tensor, tensor], dim=1)[:, :, :1]
|
||||
|
||||
monkeypatch.setattr(ulysses.SeqAllToAll4D, "apply", staticmethod(fake_all_to_all))
|
||||
strategy = ulysses.UlyssesParallelAttention(
|
||||
sp_group=SimpleNamespace(
|
||||
ulysses_group=object(),
|
||||
ulysses_world_size=2,
|
||||
ulysses_rank=0,
|
||||
ring_world_size=1,
|
||||
),
|
||||
scatter_idx=2,
|
||||
gather_idx=1,
|
||||
use_sync=False,
|
||||
)
|
||||
query = torch.zeros(1, 2, 2, 4)
|
||||
key = torch.zeros_like(query)
|
||||
value = torch.zeros_like(query)
|
||||
joint_query = torch.zeros(1, 0, 2, 4)
|
||||
joint_key = torch.zeros(1, 3, 2, 4)
|
||||
metadata = AttentionMetadata(
|
||||
attn_mask=torch.tensor([[True, True, True, False]]),
|
||||
joint_attn_mask=torch.ones(1, 3, dtype=torch.bool),
|
||||
joint_query=joint_query,
|
||||
joint_key=joint_key,
|
||||
joint_value=torch.zeros_like(joint_key),
|
||||
joint_strategy="front",
|
||||
)
|
||||
|
||||
query_out, key_out, _, metadata_out, _ = strategy.pre_attention(query, key, value, metadata)
|
||||
|
||||
assert query_out.shape[1] == 4
|
||||
assert key_out.shape[1] == 7
|
||||
assert metadata_out is not None
|
||||
assert metadata_out.attn_mask.tolist() == [[True, True, True, True, True, True, False]]
|
||||
|
||||
|
||||
def test_forward_returns_video_prediction(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
|
||||
|
||||
@@ -832,7 +721,7 @@ def test_forward_returns_video_plus_optional_modality_predictions(
|
||||
assert [tuple(tensor.shape) for tensor in output] == expected_shapes
|
||||
|
||||
|
||||
def test_forward_with_sound_odd_sequence_uses_auto_padding_contract(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
def test_forward_with_sound_ulysses_error_mentions_combined_sequence(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
import vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 as cosmos3_module
|
||||
|
||||
model = cosmos3_module.Cosmos3VFMTransformer(
|
||||
@@ -843,17 +732,16 @@ def test_forward_with_sound_odd_sequence_uses_auto_padding_contract(monkeypatch:
|
||||
)
|
||||
monkeypatch.setattr(cosmos3_module, "_get_ulysses_state", lambda: (2, 0, None))
|
||||
|
||||
output = model(
|
||||
hidden_states=torch.zeros(1, 2, 1, 1, 2),
|
||||
timestep=torch.tensor([1.0]),
|
||||
text_ids=torch.tensor([[1, 2]], dtype=torch.long),
|
||||
text_mask=torch.ones(1, 2, dtype=torch.long),
|
||||
video_shape=(1, 1, 2),
|
||||
fps=24.0,
|
||||
sound_latents=torch.zeros(1, 3, 1),
|
||||
)
|
||||
|
||||
assert [tuple(tensor.shape) for tensor in output] == [(1, 2, 1, 1, 2), (1, 3, 1)]
|
||||
with pytest.raises(ValueError, match=r"GEN sequence length \(3 = video tokens 2 \+ sound tokens 1\)"):
|
||||
model(
|
||||
hidden_states=torch.zeros(1, 2, 1, 1, 2),
|
||||
timestep=torch.tensor([1.0]),
|
||||
text_ids=torch.tensor([[1, 2]], dtype=torch.long),
|
||||
text_mask=torch.ones(1, 2, dtype=torch.long),
|
||||
video_shape=(1, 1, 2),
|
||||
fps=24.0,
|
||||
sound_latents=torch.zeros(1, 3, 1),
|
||||
)
|
||||
|
||||
|
||||
def test_sound_latent_frames_padded_for_sequence_parallel(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
|
||||
@@ -1,177 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""The request's seed must reach GR00T's flow-matching noise.
|
||||
|
||||
The runner builds ``sampling_params.generator`` from ``sampling_params.seed``, the pipeline hands
|
||||
that generator to the policy, and the action head draws its initial noise from it, so the same
|
||||
seed reproduces the same action chunk and different seeds give independent ones. These tests pin
|
||||
both ends of that plumbing.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from torch import nn
|
||||
from transformers.feature_extraction_utils import BatchFeature
|
||||
|
||||
from vllm_omni.diffusion.models.gr00t import pipeline_gr00t
|
||||
from vllm_omni.diffusion.models.gr00t.configs.gr00t_n1d7 import Gr00tN1d7Config
|
||||
from vllm_omni.diffusion.models.gr00t.modeling.gr00t_n1d7 import Gr00tN1d7ActionHead
|
||||
from vllm_omni.diffusion.models.gr00t.pipeline_gr00t import Gr00tN1d7Pipeline
|
||||
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
ROBOT_OBS = {
|
||||
"images": {"cam": np.zeros((1, 1, 8, 8, 3), dtype=np.uint8)},
|
||||
"state": {"joint": np.zeros((1, 1, 2), dtype=np.float32)},
|
||||
"prompt": "pick the cube",
|
||||
}
|
||||
|
||||
|
||||
class _OptionsRecordingPolicy:
|
||||
"""Stand-in for ``Gr00tPolicy``; the pipeline only needs ``language_key`` and ``get_action``."""
|
||||
|
||||
def __init__(self, *, model_path: str, embodiment_tag: str, device: str, strict: bool) -> None:
|
||||
del model_path, embodiment_tag, device, strict # the pipeline's constructor contract; unused here
|
||||
self.language_key = "annotation.language.language_instruction"
|
||||
self.seen_options: dict[str, object] | None = None
|
||||
|
||||
def get_action(self, obs: dict[str, object], options: dict[str, object] | None = None):
|
||||
self.seen_options = options
|
||||
return {"arm": np.zeros((1, 2, 2), dtype=np.float32)}, {}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def pipeline(monkeypatch) -> Gr00tN1d7Pipeline:
|
||||
monkeypatch.setattr(pipeline_gr00t, "Gr00tPolicy", _OptionsRecordingPolicy)
|
||||
od_config = SimpleNamespace(
|
||||
model="nvidia/GR00T-N1.7-3B",
|
||||
model_config={"embodiment_tag": "LIBERO_PANDA", "strict": False},
|
||||
custom_pipeline_args={},
|
||||
)
|
||||
return Gr00tN1d7Pipeline(od_config=od_config)
|
||||
|
||||
|
||||
def _request(**sampling_kwargs) -> OmniDiffusionRequest:
|
||||
return OmniDiffusionRequest(
|
||||
prompt="pick",
|
||||
request_id="req",
|
||||
sampling_params=OmniDiffusionSamplingParams(extra_args={"robot_obs": ROBOT_OBS}, **sampling_kwargs),
|
||||
)
|
||||
|
||||
|
||||
def test_runner_generator_is_passed_through_untouched(pipeline):
|
||||
"""In serving the runner has already built ``sampling_params.generator``; the head must get that object."""
|
||||
generator = torch.Generator().manual_seed(99)
|
||||
|
||||
assert pipeline.forward(_request(generator=generator)).error is None
|
||||
|
||||
assert pipeline.policy.seen_options["generator"] is generator
|
||||
|
||||
|
||||
def test_seed_without_generator_builds_one_on_the_pipeline_device(pipeline):
|
||||
"""Callers that bypass the runner (direct use, tests) still get the seed honoured."""
|
||||
pipeline.forward(_request(seed=1234))
|
||||
|
||||
generator = pipeline.policy.seen_options["generator"]
|
||||
assert isinstance(generator, torch.Generator)
|
||||
assert generator.initial_seed() == 1234
|
||||
assert generator.device.type == pipeline.device
|
||||
|
||||
|
||||
def test_auto_assigned_seed_reaches_the_head(pipeline):
|
||||
"""``OmniDiffusionRequest`` seeds every request, so the head is never left on the global RNG."""
|
||||
req = _request()
|
||||
assert req.sampling_params.seed is not None
|
||||
|
||||
pipeline.forward(req)
|
||||
|
||||
assert pipeline.policy.seen_options["generator"].initial_seed() == req.sampling_params.seed
|
||||
|
||||
|
||||
def _captured_warnings(monkeypatch) -> list[tuple]:
|
||||
"""vLLM loggers do not propagate to the root logger, so ``caplog`` cannot see them; record calls instead."""
|
||||
captured: list[tuple] = []
|
||||
monkeypatch.setattr(pipeline_gr00t.logger, "warning", lambda *args, **kwargs: captured.append(args))
|
||||
return captured
|
||||
|
||||
|
||||
def test_unexpected_generator_form_warns_and_falls_back_to_the_seed(pipeline, monkeypatch):
|
||||
"""GR00T draws one noise tensor per request, so a list of several generators cannot be honoured."""
|
||||
warnings = _captured_warnings(monkeypatch)
|
||||
generators = [torch.Generator().manual_seed(1), torch.Generator().manual_seed(2)]
|
||||
|
||||
pipeline.forward(_request(generator=generators, seed=7))
|
||||
|
||||
assert len(warnings) == 1
|
||||
assert "torch.Generator" in warnings[0][0]
|
||||
assert pipeline.policy.seen_options["generator"].initial_seed() == 7
|
||||
|
||||
|
||||
def test_unexpected_generator_form_without_a_seed_warns_and_leaves_the_global_rng(pipeline, monkeypatch):
|
||||
warnings = _captured_warnings(monkeypatch)
|
||||
generators = [torch.Generator().manual_seed(1), torch.Generator().manual_seed(2)]
|
||||
|
||||
pipeline.forward(_request(generator=generators))
|
||||
|
||||
assert len(warnings) == 1
|
||||
assert pipeline.policy.seen_options is None
|
||||
|
||||
|
||||
def _noise_only_head(action_horizon: int, action_dim: int) -> Gr00tN1d7ActionHead:
|
||||
"""An action head whose single denoising step is the identity, so ``action_pred`` is the initial noise."""
|
||||
config = Gr00tN1d7Config(
|
||||
max_action_dim=action_dim,
|
||||
action_horizon=action_horizon,
|
||||
num_inference_timesteps=1,
|
||||
add_pos_embed=False,
|
||||
use_alternate_vl_dit=False,
|
||||
)
|
||||
head = Gr00tN1d7ActionHead.__new__(Gr00tN1d7ActionHead)
|
||||
nn.Module.__init__(head)
|
||||
head.config = config
|
||||
head.action_dim = config.max_action_dim
|
||||
head.action_horizon = config.action_horizon
|
||||
head.num_inference_timesteps = config.num_inference_timesteps
|
||||
head.num_timestep_buckets = config.num_timestep_buckets
|
||||
head.action_encoder = lambda actions, timesteps, embodiment_id: actions
|
||||
head.model = lambda hidden_states, encoder_hidden_states, timestep: hidden_states
|
||||
head.action_decoder = lambda hidden_states, embodiment_id: torch.zeros_like(hidden_states)
|
||||
return head
|
||||
|
||||
|
||||
def _predict(head: Gr00tN1d7ActionHead, options) -> torch.Tensor:
|
||||
batch, seq, dim = 2, 3, head.action_dim
|
||||
return head.get_action_with_features(
|
||||
backbone_features=torch.zeros(batch, seq, dim),
|
||||
state_features=torch.zeros(batch, 1, dim),
|
||||
embodiment_id=torch.zeros(batch, dtype=torch.long),
|
||||
backbone_output=BatchFeature(data={}),
|
||||
action_input=BatchFeature(data={}),
|
||||
options=options,
|
||||
)["action_pred"]
|
||||
|
||||
|
||||
def test_head_draws_its_noise_from_the_request_generator():
|
||||
head = _noise_only_head(action_horizon=4, action_dim=3)
|
||||
expected = torch.randn(2, 4, 3, generator=torch.Generator().manual_seed(5))
|
||||
|
||||
assert torch.equal(_predict(head, {"generator": torch.Generator().manual_seed(5)}), expected)
|
||||
assert not torch.equal(_predict(head, {"generator": torch.Generator().manual_seed(6)}), expected)
|
||||
|
||||
|
||||
def test_head_without_generator_stays_on_the_global_rng():
|
||||
"""Upstream Isaac-GR00T parity: no generator means a bare ``torch.randn`` on the global RNG."""
|
||||
head = _noise_only_head(action_horizon=4, action_dim=3)
|
||||
|
||||
torch.manual_seed(5)
|
||||
expected = torch.randn(2, 4, 3)
|
||||
torch.manual_seed(5)
|
||||
|
||||
assert torch.equal(_predict(head, None), expected)
|
||||
@@ -1,8 +1,4 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
@@ -16,12 +12,12 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class FakeGr00tPolicy:
|
||||
instances: list["FakeGr00tPolicy"] = []
|
||||
instances = []
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
self.kwargs = kwargs
|
||||
self.reset_calls = 0
|
||||
self.seen_obs: dict[str, Any] | None = None
|
||||
self.seen_obs = None
|
||||
self.embodiment_tag = SimpleNamespace(value="fake_embodiment")
|
||||
self.language_key = "annotation.language.language_instruction"
|
||||
self.modality_configs = {
|
||||
@@ -44,7 +40,7 @@ class FakeGr00tPolicy:
|
||||
)
|
||||
FakeGr00tPolicy.instances.append(self)
|
||||
|
||||
def get_action(self, obs, options=None):
|
||||
def get_action(self, obs):
|
||||
self.seen_obs = obs
|
||||
return {
|
||||
"arm": np.array([[[1.0, 2.0]]], dtype=np.float64),
|
||||
@@ -111,7 +107,6 @@ def test_forward_returns_dict_actions_in_output():
|
||||
assert actions["arm"].dtype == np.float32
|
||||
np.testing.assert_allclose(actions["arm"], np.array([[[1.0, 2.0]]], dtype=np.float32))
|
||||
policy = FakeGr00tPolicy.instances[0]
|
||||
assert policy.seen_obs is not None
|
||||
assert "video" in policy.seen_obs
|
||||
assert policy.seen_obs["language"] == {"annotation.language.language_instruction": [["pick the cube"]]}
|
||||
assert "images" not in policy.seen_obs
|
||||
|
||||
@@ -1,292 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""GR00T-N1.7 eval image transform: chosen from ``processor_config.json`` and bit-exact with Isaac-GR00T.
|
||||
|
||||
``GOLDEN`` holds sha256 digests of the ``(C, H, W)`` uint8 output of Isaac-GR00T's own eval transform
|
||||
(``build_image_transformations_albumentations(...)[1]`` applied through ``apply_with_replay``) for the
|
||||
synthetic frames built below, at Isaac-GR00T 51d4c89 / albumentations 1.4.18 / opencv 5.0.0.93. The frame
|
||||
generators are integer-only so they are identical on every platform.
|
||||
|
||||
Regenerating a digest needs Isaac-GR00T installed, whose transformers pin conflicts with this repo's, so it
|
||||
cannot run here: in an Isaac-GR00T environment push each ``synthetic_frames()`` frame through that transform
|
||||
with ``letter_box_transform`` False and True and hash the output with ``digest()``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
import torchvision.transforms.v2 as transforms
|
||||
|
||||
from vllm_omni.diffusion.models.gr00t.modeling import processing_gr00t_n1d7 as processing
|
||||
from vllm_omni.diffusion.models.gr00t.modeling.processing_gr00t_n1d7 import (
|
||||
AlbumentationsEvalImageTransform,
|
||||
Gr00tN1d7Processor,
|
||||
LetterBoxTransform,
|
||||
_build_eval_image_transform,
|
||||
_fractional_center_crop,
|
||||
_smallest_edge_resize,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
# nvidia/GR00T-N1.7-3B and nvidia/GR00T-N1.7-DROID both ship exactly these processor_kwargs.
|
||||
TARGET, CROP, EDGE, FRACTION = [256, 256], [230, 230], 256, 0.95
|
||||
MODALITY_CONFIGS = {
|
||||
"new_embodiment": {
|
||||
"video": {"delta_indices": [0], "modality_keys": ["cam"]},
|
||||
"state": {"delta_indices": [0], "modality_keys": ["joint"]},
|
||||
"action": {"delta_indices": [0, 1], "modality_keys": ["joint"]},
|
||||
}
|
||||
}
|
||||
|
||||
# Isaac-GR00T eval transform output for synthetic_frames(); key = "<frame>|letterbox=<flag>".
|
||||
GOLDEN: dict[str, tuple[tuple[int, int, int], str]] = {
|
||||
"gradient_256x256|letterbox=False": (
|
||||
(3, 256, 256),
|
||||
"4e90a1a3c49539b3648aad752bba59af452bd7231e5712e871ab941cbea1d715",
|
||||
),
|
||||
"checker_256x256|letterbox=False": (
|
||||
(3, 256, 256),
|
||||
"d2225b94efea5ab3f1c3e6dbae52b7afa3e0f7172d5de24439a2e729ab56080c",
|
||||
),
|
||||
"noise_256x256|letterbox=False": (
|
||||
(3, 256, 256),
|
||||
"a427ae3b84ebcc6e87470c79c0e45cf684770d6e31158a49f8736685e0da72c6",
|
||||
),
|
||||
"noise_512x768|letterbox=False": (
|
||||
(3, 256, 383),
|
||||
"45bc3a309f3d06c3f31051911bc8a544ff3f067b023b999dd9c2b7e6029ecc93",
|
||||
),
|
||||
"gradient_256x256|letterbox=True": (
|
||||
(3, 256, 256),
|
||||
"4e90a1a3c49539b3648aad752bba59af452bd7231e5712e871ab941cbea1d715",
|
||||
),
|
||||
"checker_256x256|letterbox=True": (
|
||||
(3, 256, 256),
|
||||
"d2225b94efea5ab3f1c3e6dbae52b7afa3e0f7172d5de24439a2e729ab56080c",
|
||||
),
|
||||
"noise_256x256|letterbox=True": ((3, 256, 256), "a427ae3b84ebcc6e87470c79c0e45cf684770d6e31158a49f8736685e0da72c6"),
|
||||
"noise_512x768|letterbox=True": ((3, 256, 256), "52a3db185e6ee4f00617b87711dd198b8f2e2ad53865713d34e56225a4ec61a2"),
|
||||
}
|
||||
|
||||
|
||||
def _hash_noise(height: int, width: int, seed: int) -> np.ndarray:
|
||||
"""Integer-only pseudo-noise (murmur-style mixing on uint32); no RNG or float, so identical everywhere."""
|
||||
ys, xs = np.mgrid[0:height, 0:width].astype(np.uint32)
|
||||
channels = []
|
||||
for channel in range(3):
|
||||
v = xs * np.uint32(73856093) ^ ys * np.uint32(19349663) ^ np.uint32((seed * 3 + channel) * 83492791)
|
||||
v ^= v >> np.uint32(13)
|
||||
v *= np.uint32(0x5BD1E995)
|
||||
v ^= v >> np.uint32(15)
|
||||
channels.append((v & np.uint32(0xFF)).astype(np.uint8))
|
||||
return np.stack(channels, axis=-1)
|
||||
|
||||
|
||||
def synthetic_frames() -> dict[str, np.ndarray]:
|
||||
"""Deterministic HWC uint8 RGB frames: smooth, hard-edged, noisy, and a non-square noisy one.
|
||||
|
||||
The non-square frame is 512x768 so both INTER_AREA downscales (2x without letterbox, 3x with it)
|
||||
are integer factors, which OpenCV computes with exact integer sums: the goldens do not depend on
|
||||
the CPU architecture. The 243->256 upscale is a fixed-point bilinear on uint8, also exact.
|
||||
"""
|
||||
ys, xs = np.mgrid[0:256, 0:256].astype(np.int64)
|
||||
frames = {
|
||||
"gradient_256x256": np.stack([xs % 256, ys % 256, (xs + ys) // 2 % 256], axis=-1).astype(np.uint8),
|
||||
"checker_256x256": np.stack(
|
||||
[((ys // 16 + xs // 16) % 2) * 255, ((ys // 8 + xs // 8) % 2) * 255, ((ys // 32) % 2) * 255],
|
||||
axis=-1,
|
||||
).astype(np.uint8),
|
||||
"noise_256x256": _hash_noise(256, 256, seed=1),
|
||||
"noise_512x768": _hash_noise(512, 768, seed=2),
|
||||
}
|
||||
return frames
|
||||
|
||||
|
||||
def digest(tensor: torch.Tensor) -> str:
|
||||
return hashlib.sha256(np.ascontiguousarray(tensor.numpy()).tobytes()).hexdigest()
|
||||
|
||||
|
||||
class _FakeTokenizer:
|
||||
padding_side = "right"
|
||||
|
||||
|
||||
class _FakeVlmProcessor:
|
||||
"""The two attributes ``Gr00tN1d7Processor`` and its collator touch on the Qwen3-VL processor."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.tokenizer = _FakeTokenizer()
|
||||
|
||||
def apply_chat_template(self, *_args: object, **_kwargs: object) -> str:
|
||||
return "<prompt>"
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def fake_vlm_processor(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""The real Qwen3-VL processor is a gated download and irrelevant to the image pipeline."""
|
||||
monkeypatch.setattr(processing, "build_processor", lambda _model_name, _kwargs: _FakeVlmProcessor())
|
||||
|
||||
|
||||
def _processor(**kwargs) -> Gr00tN1d7Processor:
|
||||
return Gr00tN1d7Processor(
|
||||
modality_configs=MODALITY_CONFIGS, image_crop_size=CROP, image_target_size=TARGET, **kwargs
|
||||
)
|
||||
|
||||
|
||||
def _has_letterbox(compose: transforms.Compose) -> bool:
|
||||
return any(isinstance(step, LetterBoxTransform) for step in compose.transforms)
|
||||
|
||||
|
||||
def test_processor_builds_the_pipeline_named_by_processor_config():
|
||||
albumentations = _processor(use_albumentations=True, shortest_image_edge=EDGE, crop_fraction=FRACTION)
|
||||
transform = albumentations.eval_image_transform
|
||||
assert isinstance(transform, AlbumentationsEvalImageTransform)
|
||||
assert (transform.max_size, transform.fraction, transform.letter_box_transform) == (EDGE, FRACTION, False)
|
||||
assert albumentations.use_albumentations is True
|
||||
|
||||
torchvision = _processor(use_albumentations=False)
|
||||
assert isinstance(torchvision.eval_image_transform, transforms.Compose)
|
||||
assert not _has_letterbox(torchvision.eval_image_transform)
|
||||
|
||||
letterboxed = _processor(use_albumentations=False, letter_box_transform=True)
|
||||
assert _has_letterbox(letterboxed.eval_image_transform)
|
||||
assert _processor(use_albumentations=True, letter_box_transform=True).eval_image_transform.letter_box_transform
|
||||
|
||||
|
||||
def test_legacy_config_without_the_flags_keeps_the_torchvision_recipe():
|
||||
"""Checkpoints predating the flags behave like Isaac-GR00T's defaults; training-only keys stay ignored."""
|
||||
processor = _processor(state_dropout_prob=0.2, extra_augmentation_config=None)
|
||||
assert processor.use_albumentations is False
|
||||
assert processor.letter_box_transform is False
|
||||
assert isinstance(processor.eval_image_transform, transforms.Compose)
|
||||
assert not _has_letterbox(processor.eval_image_transform)
|
||||
|
||||
|
||||
def test_fraction_and_max_size_fall_back_like_isaac_gr00t():
|
||||
transform = _build_eval_image_transform(
|
||||
TARGET, CROP, shortest_image_edge=None, crop_fraction=None, use_albumentations=True
|
||||
)
|
||||
assert isinstance(transform, AlbumentationsEvalImageTransform)
|
||||
assert transform.fraction == CROP[0] / TARGET[0]
|
||||
assert transform.max_size == TARGET[0]
|
||||
with pytest.raises(ValueError, match="crop_fraction is None"):
|
||||
_build_eval_image_transform(TARGET, None, crop_fraction=None, use_albumentations=True)
|
||||
|
||||
|
||||
def test_geometry_follows_isaac_rounding_rules():
|
||||
frame = _hash_noise(540, 640, seed=3)
|
||||
resized = _smallest_edge_resize(frame, EDGE)
|
||||
assert resized.shape == (256, 303, 3) # round() per dimension, shortest edge -> 256
|
||||
cropped = _fractional_center_crop(resized, FRACTION)
|
||||
assert cropped.shape == (243, 287, 3) # int() truncation: 243.2 -> 243, 287.85 -> 287
|
||||
assert np.array_equal(cropped, resized[6:249, 8:295]) # (256-243)//2, (303-287)//2
|
||||
square = frame[:256, :256]
|
||||
assert _smallest_edge_resize(square, EDGE) is square # scale == 1.0 is a no-op, not a resample
|
||||
|
||||
|
||||
@pytest.mark.parametrize("case", sorted(GOLDEN))
|
||||
def test_albumentations_recipe_is_bit_exact_with_isaac_gr00t(case: str):
|
||||
name, flag = case.split("|letterbox=")
|
||||
transform = _build_eval_image_transform(
|
||||
TARGET,
|
||||
CROP,
|
||||
shortest_image_edge=EDGE,
|
||||
crop_fraction=FRACTION,
|
||||
use_albumentations=True,
|
||||
letter_box_transform=(flag == "True"),
|
||||
)
|
||||
out = transform(synthetic_frames()[name])
|
||||
expected_shape, expected_digest = GOLDEN[case]
|
||||
assert out.dtype == torch.uint8
|
||||
assert tuple(out.shape) == expected_shape
|
||||
assert digest(out) == expected_digest
|
||||
|
||||
|
||||
def test_letterbox_is_gated_on_both_recipes():
|
||||
frame = _hash_noise(540, 640, seed=3)
|
||||
|
||||
def albumentations(letterbox: bool) -> torch.Tensor:
|
||||
return _build_eval_image_transform(
|
||||
TARGET,
|
||||
CROP,
|
||||
shortest_image_edge=EDGE,
|
||||
crop_fraction=FRACTION,
|
||||
use_albumentations=True,
|
||||
letter_box_transform=letterbox,
|
||||
)(frame)
|
||||
|
||||
padded, unpadded = albumentations(True), albumentations(False)
|
||||
assert tuple(padded.shape) == (3, 256, 256) # 540x640 -> 640x640 -> 256x256 -> 243x243 -> 256x256
|
||||
assert tuple(unpadded.shape) == (3, 256, 302) # aspect preserved: 256x303 -> 243x287 -> 256x302
|
||||
|
||||
tv_padded = _build_eval_image_transform(TARGET, CROP, use_albumentations=False, letter_box_transform=True)(frame)
|
||||
tv_unpadded = _build_eval_image_transform(TARGET, CROP, use_albumentations=False, letter_box_transform=False)(frame)
|
||||
assert tuple(tv_padded.shape) == tuple(tv_unpadded.shape) == (3, 256, 256)
|
||||
assert not torch.equal(tv_padded, tv_unpadded)
|
||||
assert (tv_padded[:, :2, :] == 0).all() and not (tv_unpadded[:, :2, :] == 0).all() # black bars only when gated
|
||||
|
||||
|
||||
def test_albumentations_recipe_rejects_frames_outside_the_contract():
|
||||
transform = _build_eval_image_transform(TARGET, CROP, use_albumentations=True)
|
||||
with pytest.raises(ValueError, match="HWC uint8"):
|
||||
transform(np.zeros((256, 256, 3), dtype=np.float32))
|
||||
with pytest.raises(ValueError, match="HWC uint8"):
|
||||
transform(np.zeros((3, 256, 256), dtype=np.uint8))
|
||||
with pytest.raises(ValueError, match="HWC uint8"):
|
||||
transform(np.zeros((256, 256), dtype=np.uint8))
|
||||
|
||||
|
||||
def test_get_vlm_inputs_runs_the_selected_transform_per_frame():
|
||||
processor = _processor(use_albumentations=True, shortest_image_edge=EDGE, crop_fraction=FRACTION)
|
||||
frames = np.stack([_hash_noise(540, 640, seed=3)] * 2) # (T, H, W, C) like policy.py hands over
|
||||
vlm_inputs = processor._get_vlm_inputs(
|
||||
image_keys=["cam"],
|
||||
images={"cam": frames},
|
||||
masks=None,
|
||||
image_transform=processor.eval_image_transform,
|
||||
language="x",
|
||||
)
|
||||
images = vlm_inputs["vlm_content"]["images"]
|
||||
assert len(images) == 2
|
||||
assert all(img.size == (302, 256) for img in images) # PIL reports (W, H); no letterbox for this checkpoint
|
||||
|
||||
|
||||
def test_flags_round_trip_through_save_pretrained_and_from_pretrained(tmp_path):
|
||||
saved = _processor(
|
||||
use_albumentations=True, letter_box_transform=True, shortest_image_edge=EDGE, crop_fraction=FRACTION
|
||||
)
|
||||
saved.save_pretrained(tmp_path)
|
||||
kwargs = json.loads((tmp_path / "processor_config.json").read_text())["processor_kwargs"]
|
||||
assert kwargs["use_albumentations"] is True
|
||||
assert kwargs["letter_box_transform"] is True
|
||||
|
||||
loaded = Gr00tN1d7Processor.from_pretrained(tmp_path)
|
||||
assert loaded.use_albumentations is True
|
||||
assert loaded.letter_box_transform is True
|
||||
assert isinstance(loaded.eval_image_transform, AlbumentationsEvalImageTransform)
|
||||
assert loaded.eval_image_transform.letter_box_transform is True
|
||||
|
||||
|
||||
def test_from_pretrained_defaults_missing_flags_to_isaac_gr00t_defaults(tmp_path):
|
||||
(tmp_path / "processor_config.json").write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"processor_class": "Gr00tN1d7Processor",
|
||||
"processor_kwargs": {
|
||||
"modality_configs": MODALITY_CONFIGS,
|
||||
"image_crop_size": CROP,
|
||||
"image_target_size": TARGET,
|
||||
},
|
||||
}
|
||||
)
|
||||
)
|
||||
(tmp_path / "statistics.json").write_text("{}")
|
||||
loaded = Gr00tN1d7Processor.from_pretrained(tmp_path)
|
||||
assert loaded.use_albumentations is False
|
||||
assert loaded.letter_box_transform is False
|
||||
assert isinstance(loaded.eval_image_transform, transforms.Compose)
|
||||
@@ -1,5 +1,4 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -14,7 +13,6 @@ from vllm_omni.diffusion.models.lingbot_world.actions import (
|
||||
LingBotCameraControlReducer,
|
||||
integrate_lingbot_camera_actions,
|
||||
parse_lingbot_camera_action_frames,
|
||||
parse_lingbot_camera_action_script,
|
||||
)
|
||||
from vllm_omni.experimental.ar_diffusion.session import (
|
||||
ARDiffusionSession,
|
||||
@@ -294,8 +292,3 @@ def test_session_reducer_keeps_prompt_and_actions_atomic_and_resets() -> None:
|
||||
assert consumer.ticks[1].controls[0].data["frames"] == ((), (), ())
|
||||
|
||||
asyncio.run(exercise())
|
||||
|
||||
|
||||
def test_parse_camera_action_script_rejects_wrong_chunk_width() -> None:
|
||||
with pytest.raises(ValueError, match="exactly 3 per-latent-frame"):
|
||||
parse_lingbot_camera_action_script([[["w"], ["w"]]], frames_per_chunk=3)
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -15,10 +15,8 @@ from diffusers.utils.torch_utils import randn_tensor as _diffusers_randn_tensor
|
||||
from PIL import Image
|
||||
from torch import nn
|
||||
|
||||
import vllm_omni.diffusion.models.lingbot_world.dmd_block as lingbot_dmd_block
|
||||
import vllm_omni.diffusion.models.lingbot_world.pipeline as lingbot_pipeline
|
||||
from tests.diffusion.models.wan2_2.conftest import noop_progress_bar
|
||||
from vllm_omni.diffusion.models.interface import SupportsStepExecution, supports_step_execution
|
||||
from vllm_omni.diffusion.models.lingbot_world.actions import (
|
||||
integrate_lingbot_camera_actions,
|
||||
)
|
||||
@@ -26,8 +24,6 @@ from vllm_omni.diffusion.models.lingbot_world.camera import CameraTrajectory as
|
||||
from vllm_omni.diffusion.models.lingbot_world.camera import (
|
||||
build_plucker_embedding as _real_build_plucker_embedding,
|
||||
)
|
||||
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
||||
from vllm_omni.diffusion.worker.utils import StepRequestState
|
||||
from vllm_omni.experimental.ar_diffusion.tick_protocol import (
|
||||
ARDiffusionControlInput,
|
||||
ARDiffusionTickRequest,
|
||||
@@ -257,8 +253,6 @@ class _SamplingParams:
|
||||
if extra_args is not None:
|
||||
self.extra_args.update(extra_args)
|
||||
self.latents = None
|
||||
self.guidance_scale = None
|
||||
self.guidance_scale_2 = None
|
||||
|
||||
|
||||
class _RequestBatch:
|
||||
@@ -351,6 +345,7 @@ def _stub_pipeline_dependencies(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"FlowUniPCMultistepScheduler": _FakeScheduler,
|
||||
"CausalLingBotWorldTransformer3DModel": _FakeTransformerFactory,
|
||||
"get_local_device": lambda: torch.device("cpu"),
|
||||
"set_forward_context_denoise_step_idx": lambda index: None,
|
||||
"prefetch_subfolders": prefetch_subfolders,
|
||||
"from_pretrained_with_prefetch": from_pretrained_with_prefetch,
|
||||
"load_transformer_config": load_transformer_config,
|
||||
@@ -363,9 +358,6 @@ def _stub_pipeline_dependencies(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
}
|
||||
for name, value in replacements.items():
|
||||
monkeypatch.setattr(lingbot_pipeline, name, value)
|
||||
# The DMD block math lives in its own module; stub the symbols it imports.
|
||||
monkeypatch.setattr(lingbot_dmd_block, "set_forward_context_denoise_step_idx", lambda index: None)
|
||||
monkeypatch.setattr(lingbot_dmd_block, "randn_tensor", _diffusers_randn_tensor)
|
||||
monkeypatch.setattr(lingbot_pipeline, "_loader_state", loader_state, raising=False)
|
||||
|
||||
|
||||
@@ -782,7 +774,6 @@ def test_denoise_state_stays_fp32_while_transformer_inputs_use_model_dtype() ->
|
||||
return torch.zeros(shape, device=device, dtype=dtype)
|
||||
|
||||
module.randn_tensor = randn
|
||||
lingbot_dmd_block.randn_tensor = randn
|
||||
result = pipeline(_request())
|
||||
|
||||
assert requested_noise_dtypes == [torch.float32] * 4
|
||||
@@ -1148,7 +1139,7 @@ def test_first_frame_condition_and_camera_fold_match_transformer_contract() -> N
|
||||
module = _load_pipeline_module()
|
||||
transformer = _RecordingTransformer()
|
||||
pipeline = _pipeline(module, transformer=transformer)
|
||||
module.randn_tensor = lingbot_dmd_block.randn_tensor = lambda shape, **kwargs: torch.full(
|
||||
module.randn_tensor = lambda shape, **kwargs: torch.full(
|
||||
shape,
|
||||
-99.0,
|
||||
device=kwargs["device"],
|
||||
@@ -1754,288 +1745,6 @@ def test_request_cache_becomes_unreachable_after_transformer_error() -> None:
|
||||
assert not hasattr(pipeline, "transformer_cache")
|
||||
|
||||
|
||||
class _FakeARState:
|
||||
def __init__(self, session_id: str = "req-1") -> None:
|
||||
self.session_id = session_id
|
||||
self.commits: list[str] = []
|
||||
|
||||
def get_kv_caches(self, branch, *, seq_len, commit_current):
|
||||
del branch, seq_len, commit_current
|
||||
return [SimpleNamespace()]
|
||||
|
||||
def commit_paged_context(self, branch):
|
||||
self.commits.append(branch)
|
||||
|
||||
def clear_cross_attention(self):
|
||||
return None
|
||||
|
||||
def is_cross_attention_populated(self, branch, name):
|
||||
del branch, name
|
||||
return True
|
||||
|
||||
def get_cross_attention_kv(self, branch, name):
|
||||
del branch, name
|
||||
zeros = torch.zeros(1, 1, 2, 4)
|
||||
return [{"k": zeros, "v": zeros} for _ in range(2)]
|
||||
|
||||
|
||||
def _empty_action_script(num_chunks: int):
|
||||
return tuple(((), (), ()) for _ in range(num_chunks))
|
||||
|
||||
|
||||
def _stepwise_state(
|
||||
*,
|
||||
request_id: str = "req-1",
|
||||
num_frames: int = 21,
|
||||
script=None,
|
||||
seed: int = 17,
|
||||
):
|
||||
num_chunks = ((num_frames - 1) // 4 + 1) // 3
|
||||
sampling = _SamplingParams(num_frames=num_frames, seed=seed)
|
||||
sampling.extra_args["_lingbot_camera_trajectory"] = None
|
||||
sampling.extra_args["_lingbot_camera_actions"] = None
|
||||
sampling.extra_args["_lingbot_camera_action_script"] = (
|
||||
script if script is not None else _empty_action_script(num_chunks)
|
||||
)
|
||||
return StepRequestState(
|
||||
request_id=request_id,
|
||||
sampling=sampling,
|
||||
prompt=_prompt(),
|
||||
)
|
||||
|
||||
|
||||
def _run_stepwise(pipeline, state):
|
||||
outputs = []
|
||||
pipeline.prepare_encode(state)
|
||||
while not state.request_denoise_completed:
|
||||
noise = pipeline.denoise_step(None, states=[state])
|
||||
pipeline.step_scheduler(state, noise)
|
||||
if state.chunk_denoise_completed:
|
||||
outputs.append(pipeline.post_decode(state))
|
||||
return outputs
|
||||
|
||||
|
||||
def test_pipeline_declares_step_execution_support() -> None:
|
||||
module = _load_pipeline_module()
|
||||
pipeline = _pipeline(module)
|
||||
|
||||
assert pipeline.supports_step_execution is True
|
||||
assert isinstance(pipeline, SupportsStepExecution)
|
||||
assert supports_step_execution(pipeline) is True
|
||||
|
||||
|
||||
def test_preprocess_materializes_camera_action_script_without_action_path() -> None:
|
||||
module = _load_pipeline_module()
|
||||
sampling = _SamplingParams(include_action=False)
|
||||
sampling.extra_args["camera_action_script"] = [[["w"], ["w"], ["w"]], [["a"], [], []]]
|
||||
request = OmniDiffusionRequest(prompt=_prompt(), sampling_params=sampling, request_id="req-1")
|
||||
|
||||
result = module.get_lingbot_world_pre_process_func(_od_config())(request)
|
||||
|
||||
assert result.sampling_params.extra_args["_lingbot_camera_trajectory"] is None
|
||||
assert result.sampling_params.extra_args["_lingbot_camera_actions"] is None
|
||||
assert result.sampling_params.extra_args["_lingbot_camera_action_script"] == (
|
||||
(("w",), ("w",), ("w",)),
|
||||
(("a",), (), ()),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"extra_args",
|
||||
[
|
||||
pytest.param({}, id="request_mode"),
|
||||
pytest.param(_tick_extra_args(chunk_index=0), id="tick_mode"),
|
||||
],
|
||||
)
|
||||
def test_forward_rejects_a_stepwise_camera_action_script(extra_args) -> None:
|
||||
"""A script only steers step execution, so request mode must not drop it silently."""
|
||||
module = _load_pipeline_module()
|
||||
pipeline = _pipeline(module)
|
||||
sampling = _SamplingParams(
|
||||
include_action=False,
|
||||
extra_args={
|
||||
**extra_args,
|
||||
"_lingbot_camera_trajectory": None,
|
||||
"_lingbot_camera_action_script": _empty_action_script(1),
|
||||
},
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="camera_action_script is read only by LingBot step execution"):
|
||||
pipeline(_request(sampling=sampling))
|
||||
|
||||
|
||||
def test_stepwise_progress_metadata_and_commit_trace() -> None:
|
||||
module = _load_pipeline_module()
|
||||
transformer = _RecordingTransformer()
|
||||
pipeline = _pipeline(module, transformer=transformer)
|
||||
pipeline._ar_height = 16
|
||||
pipeline._ar_width = 16
|
||||
state = _stepwise_state(num_frames=21)
|
||||
fake = _FakeARState(state.request_id)
|
||||
|
||||
with pipeline.bind_ar_diffusion_state(state.request_id, fake):
|
||||
outputs = _run_stepwise(pipeline, state)
|
||||
|
||||
assert pipeline._ar_diffusion_kv_state is None
|
||||
assert len(outputs) == 2
|
||||
assert [output.chunk_index for output in outputs] == [0, 1]
|
||||
assert all(output.total_chunks == 2 for output in outputs)
|
||||
assert outputs[-1].finished is True
|
||||
assert [call["update_cache"] for call in transformer.calls] == [False, False, False, False, True] * 2
|
||||
assert [call["start_frame"] for call in transformer.calls] == [0] * 5 + [3] * 5
|
||||
assert fake.commits == ["main", "main"]
|
||||
for chunk_index, output in enumerate(outputs):
|
||||
metadata = output.output["metadata"]["ar_diffusion"]
|
||||
assert metadata == {
|
||||
"session_id": "req-1",
|
||||
"request_id": "req-1",
|
||||
"chunk_index": chunk_index,
|
||||
"applied_event_ids": [],
|
||||
}
|
||||
assert output.output["payload"]["latents"].shape == (1, 16, 3, 2, 2)
|
||||
|
||||
|
||||
def test_stepwise_matches_tick_transformer_trace_and_latents() -> None:
|
||||
module = _load_pipeline_module()
|
||||
transformer = _RecordingTransformer()
|
||||
pipeline = _pipeline(module, transformer=transformer)
|
||||
pipeline._ar_height = 16
|
||||
pipeline._ar_width = 16
|
||||
script = ((("w",), ("w",), ("w",)), (("d",), ("d",), ("d",)))
|
||||
fake = _FakeARState("world-1")
|
||||
|
||||
with pipeline.bind_ar_diffusion_state("world-1", fake):
|
||||
tick_latents = []
|
||||
for chunk_index, frames in enumerate(script):
|
||||
tick = ARDiffusionTickRequest(
|
||||
session_id="world-1",
|
||||
request_id="legacy-lingbot-request",
|
||||
chunk_index=chunk_index,
|
||||
controls=(
|
||||
ARDiffusionControlInput(
|
||||
track="camera",
|
||||
schema="lingbot.camera_actions.v1",
|
||||
data={"mode": "frames", "frames": [list(frame) for frame in frames]},
|
||||
),
|
||||
),
|
||||
)
|
||||
sampling = _SamplingParams(extra_args=tick.to_extra_args(), seed=17)
|
||||
sampling.extra_args["_lingbot_camera_trajectory"] = None
|
||||
sampling.extra_args["_lingbot_camera_actions"] = frames
|
||||
result = pipeline(_request(sampling=sampling))
|
||||
tick_latents.append(result.output["payload"]["latents"].clone())
|
||||
tick_calls = list(transformer.calls)
|
||||
transformer.calls.clear()
|
||||
pipeline.close_ar_diffusion_session("world-1")
|
||||
|
||||
stepwise_state = _stepwise_state(request_id="req-1", num_frames=21, script=script, seed=17)
|
||||
stepwise_fake = _FakeARState(stepwise_state.request_id)
|
||||
with pipeline.bind_ar_diffusion_state(stepwise_state.request_id, stepwise_fake):
|
||||
stepwise_outputs = _run_stepwise(pipeline, stepwise_state)
|
||||
|
||||
assert [call["update_cache"] for call in transformer.calls] == [call["update_cache"] for call in tick_calls]
|
||||
assert [call["start_frame"] for call in transformer.calls] == [call["start_frame"] for call in tick_calls]
|
||||
for tick_latent, output in zip(tick_latents, stepwise_outputs, strict=True):
|
||||
torch.testing.assert_close(output.output["payload"]["latents"], tick_latent)
|
||||
|
||||
|
||||
def test_stepwise_trajectory_camera_matches_request_mode_under_non_uniform_speed(monkeypatch) -> None:
|
||||
"""A trajectory that slows down between blocks must condition every stepwise
|
||||
chunk exactly as request mode does. Request mode embeds the whole trajectory
|
||||
once; embedding per chunk would re-normalize each block's framewise
|
||||
translations by its own largest step and present a slow block as full-speed
|
||||
motion."""
|
||||
from vllm_omni.diffusion.models.lingbot_world import camera as camera_module
|
||||
|
||||
module = _load_pipeline_module()
|
||||
# The fixture stubs the ray embedding with a pose-blind ramp; parity needs
|
||||
# the real geometry so a scale mismatch actually shows up.
|
||||
monkeypatch.setattr(module, "build_plucker_embedding", camera_module.build_plucker_embedding)
|
||||
transformer = _RecordingTransformer()
|
||||
pipeline = _pipeline(module, transformer=transformer)
|
||||
pipeline._ar_height = 16
|
||||
pipeline._ar_width = 16
|
||||
|
||||
# Six latent frames -> two 3-frame blocks. The first block moves at full
|
||||
# speed and everything after it at a tenth of that, so a per-chunk
|
||||
# normalization would rescale the second block by 10x.
|
||||
steps = torch.tensor([1.0] * 3 + [0.1] * 18)
|
||||
poses = torch.eye(4).repeat(21, 1, 1)
|
||||
poses[:, 2, 3] = torch.cumsum(steps, dim=0) - steps[0]
|
||||
trajectory = _CameraTrajectory(
|
||||
poses=poses,
|
||||
intrinsics=torch.tensor([[100.0, 100.0, 8.0, 8.0]]).repeat(21, 1),
|
||||
)
|
||||
|
||||
sampling = _SamplingParams(num_frames=21)
|
||||
sampling.extra_args["_lingbot_camera_trajectory"] = trajectory
|
||||
pipeline(_request(sampling=sampling))
|
||||
request_cameras = [call["camera_hidden_states"] for call in transformer.calls]
|
||||
transformer.calls.clear()
|
||||
|
||||
state = _stepwise_state(num_frames=21)
|
||||
state.sampling.extra_args["_lingbot_camera_trajectory"] = trajectory
|
||||
state.sampling.extra_args.pop("_lingbot_camera_action_script")
|
||||
with pipeline.bind_ar_diffusion_state(state.request_id, _FakeARState(state.request_id)):
|
||||
_run_stepwise(pipeline, state)
|
||||
stepwise_cameras = [call["camera_hidden_states"] for call in transformer.calls]
|
||||
|
||||
# Two blocks x (four probes + one commit), in the same order on both paths.
|
||||
assert len(request_cameras) == len(stepwise_cameras) == 10
|
||||
for stepwise, request in zip(stepwise_cameras, request_cameras, strict=True):
|
||||
torch.testing.assert_close(stepwise, request)
|
||||
|
||||
|
||||
def test_stepwise_emits_latents_without_touching_the_vae() -> None:
|
||||
"""Latent mode must stay latent: streaming decode is opt-in per request."""
|
||||
module = _load_pipeline_module()
|
||||
pipeline = _pipeline(module, transformer=_RecordingTransformer())
|
||||
pipeline._ar_height = 16
|
||||
pipeline._ar_width = 16
|
||||
state = _stepwise_state(num_frames=21)
|
||||
fake = _FakeARState(state.request_id)
|
||||
|
||||
with pipeline.bind_ar_diffusion_state(state.request_id, fake):
|
||||
outputs = _run_stepwise(pipeline, state)
|
||||
|
||||
assert [set(output.output["payload"]) for output in outputs] == [{"latents"}, {"latents"}]
|
||||
assert pipeline.vae.decode_inputs == []
|
||||
|
||||
|
||||
def test_stepwise_decodes_each_chunk_for_streaming_consumers(monkeypatch) -> None:
|
||||
"""A non-latent request must stream pixels under the primary "video" key."""
|
||||
import diffusers.video_processor as video_processor_module
|
||||
|
||||
module = _load_pipeline_module()
|
||||
processed = []
|
||||
|
||||
class VideoProcessor:
|
||||
def __init__(self, *, vae_scale_factor):
|
||||
assert vae_scale_factor == 8
|
||||
|
||||
def postprocess_video(self, video, *, output_type):
|
||||
processed.append((tuple(video.shape), output_type))
|
||||
return f"frames-{len(processed)}"
|
||||
|
||||
monkeypatch.setattr(video_processor_module, "VideoProcessor", VideoProcessor)
|
||||
pipeline = _pipeline(module, transformer=_RecordingTransformer())
|
||||
pipeline._ar_height = 16
|
||||
pipeline._ar_width = 16
|
||||
state = _stepwise_state(num_frames=21)
|
||||
state.sampling.output_type = "np"
|
||||
fake = _FakeARState(state.request_id)
|
||||
|
||||
with pipeline.bind_ar_diffusion_state(state.request_id, fake):
|
||||
outputs = _run_stepwise(pipeline, state)
|
||||
|
||||
# One decode per AR block, each seeing only that block's three latent frames.
|
||||
assert [tuple(latents.shape) for latents in pipeline.vae.decode_inputs] == [(1, 16, 3, 2, 2)] * 2
|
||||
assert [output_type for _, output_type in processed] == ["np", "np"]
|
||||
assert [output.output["payload"] for output in outputs] == [{"video": "frames-1"}, {"video": "frames-2"}]
|
||||
# Identity metadata survives the switch to pixel output.
|
||||
assert [output.output["metadata"]["ar_diffusion"]["chunk_index"] for output in outputs] == [0, 1]
|
||||
|
||||
|
||||
def test_registry_and_model_exports_resolve_official_pipeline_class_name() -> None:
|
||||
module = _load_pipeline_module()
|
||||
resolved, entry, cache_acceleration_disabled, preprocess_name, preprocess = _resolve_pipeline_through_real_registry(
|
||||
@@ -2046,7 +1755,7 @@ def test_registry_and_model_exports_resolve_official_pipeline_class_name() -> No
|
||||
assert resolved is module.LingBotWorldCausalDMDPipeline
|
||||
assert cache_acceleration_disabled
|
||||
assert preprocess_name == "get_lingbot_world_pre_process_func"
|
||||
request = OmniDiffusionRequest(prompt=_prompt(), sampling_params=_SamplingParams(), request_id="req-1")
|
||||
request = SimpleNamespace(prompt=_prompt(), sampling_params=_SamplingParams())
|
||||
assert preprocess(request) is request
|
||||
assert isinstance(request.sampling_params.extra_args["_lingbot_camera_trajectory"], _CameraTrajectory)
|
||||
assert module.LingBotWorldCausalDMDPipeline.__name__ == "LingBotWorldCausalDMDPipeline"
|
||||
|
||||
@@ -5,7 +5,6 @@
|
||||
|
||||
import json
|
||||
from dataclasses import replace
|
||||
from threading import Lock
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
|
||||
@@ -501,32 +500,6 @@ def test_ltx25_checkpoint_selects_distilled_two_stage_profile(tmp_path, monkeypa
|
||||
assert pipe.pipeline_recipe is LTX25_DISTILLED_TWO_STAGE_RECIPE
|
||||
|
||||
|
||||
def test_ltx_runtime_profiles_phase_methods_separately(monkeypatch):
|
||||
pipe = object.__new__(LTX2DistilledTwoStagePipeline)
|
||||
pipe.enable_diffusion_pipeline_profiler = True
|
||||
pipe._profiler_lock = Lock()
|
||||
pipe._stage_durations = {}
|
||||
monkeypatch.setattr(ltx2_runtime.current_omni_platform, "is_available", lambda: False)
|
||||
|
||||
with pipe._profile_phase_method("run_phase", "generate_lowres"):
|
||||
pass
|
||||
with pipe._profile_phase_method("run_phase", "refine"):
|
||||
pass
|
||||
with pipe._profile_phase_method("decode_phase"):
|
||||
pass
|
||||
output = DiffusionOutput()
|
||||
assert pipe._attach_profiled_stage_durations(output) is output
|
||||
|
||||
expected_keys = {
|
||||
"LTX2DistilledTwoStagePipeline.generate_lowres.run_phase",
|
||||
"LTX2DistilledTwoStagePipeline.refine.run_phase",
|
||||
"LTX2DistilledTwoStagePipeline.decode_phase",
|
||||
}
|
||||
assert set(pipe.stage_durations) == expected_keys
|
||||
assert set(output.stage_durations) == expected_keys
|
||||
assert all(duration >= 0.0 for duration in pipe.stage_durations.values())
|
||||
|
||||
|
||||
def test_ltx25_full_config_enrichment_uses_selected_transformer_subfolder(tmp_path):
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
|
||||
|
||||
@@ -103,7 +103,7 @@ def test_prepare_audio_latents_samples_directly_in_packed_token_space():
|
||||
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
|
||||
|
||||
|
||||
def test_prepare_audio_latents_pads_short_audio_for_sp():
|
||||
def test_prepare_audio_latents_pads_generated_dummy_length_for_sp():
|
||||
pipeline = _make_pipeline(LTX2Pipeline, sequence_parallel_size=2)
|
||||
|
||||
latents, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents(
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""Unit tests for LTX guidance and forward-parallel behavior."""
|
||||
|
||||
@@ -85,14 +85,14 @@ class TestCFGParallelHelpers:
|
||||
|
||||
LTX_GUIDANCE_EXECUTOR.validate_guidance_world_size(plan, guidance_world_size)
|
||||
|
||||
def test_guidance_parallel_preserves_full_guidance_plan(self, monkeypatch):
|
||||
def test_guidance_parallel_dummy_warms_full_guidance_plan(self, monkeypatch):
|
||||
from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
|
||||
from vllm_omni.diffusion.models.ltx2.ltx2_recipes import LTX23_ONE_STAGE_RECIPE
|
||||
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
|
||||
|
||||
monkeypatch.setattr(ltx2_runtime, "get_guidance_parallel_world_size", lambda: 2)
|
||||
pipe = object.__new__(LTX2Pipeline)
|
||||
req = SimpleNamespace()
|
||||
req = SimpleNamespace(is_dummy_run=lambda: True)
|
||||
request_inputs = SimpleNamespace(guidance=LTX23_ONE_STAGE_RECIPE.request_guidance)
|
||||
|
||||
guidance_parallel_ready = pipe._setup_forward_runtime(req, request_inputs, attention_kwargs=None)
|
||||
|
||||
@@ -10,7 +10,6 @@ from dataclasses import replace
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
|
||||
import huggingface_hub
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
@@ -176,19 +175,7 @@ def test_ltx_device_uint8_output_rejects_invalid_contract(video, error):
|
||||
_prepare_ltx2_video_output(video, do_normalize=True)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"parallel_config",
|
||||
[
|
||||
None,
|
||||
SimpleNamespace(ulysses_degree=4),
|
||||
SimpleNamespace(tensor_parallel_size=2),
|
||||
SimpleNamespace(use_hsdp=True, hsdp_shard_size=2),
|
||||
SimpleNamespace(ulysses_degree=2, tensor_parallel_size=2),
|
||||
SimpleNamespace(ulysses_degree=2, use_hsdp=True, hsdp_shard_size=2),
|
||||
],
|
||||
ids=["single", "ulysses", "tp", "hsdp", "tp-ulysses", "hsdp-ulysses"],
|
||||
)
|
||||
def test_ltx_supported_parallel_configs_keep_output_profiling(tmp_path, monkeypatch, parallel_config):
|
||||
def test_ltx_profiler_covers_video_output_transport(tmp_path, monkeypatch):
|
||||
from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
|
||||
|
||||
(tmp_path / "model_index.json").write_text(json.dumps({"vocoder": ["ltx2", "LTX2Vocoder"]}))
|
||||
@@ -206,13 +193,7 @@ def test_ltx_supported_parallel_configs_keep_output_profiling(tmp_path, monkeypa
|
||||
monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components)
|
||||
monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", capture_profiler_setup)
|
||||
|
||||
LTX2Pipeline(
|
||||
od_config=SimpleNamespace(
|
||||
model=str(tmp_path),
|
||||
parallel_config=parallel_config,
|
||||
enable_diffusion_pipeline_profiler=False,
|
||||
)
|
||||
)
|
||||
LTX2Pipeline(od_config=SimpleNamespace(model=str(tmp_path), enable_diffusion_pipeline_profiler=False))
|
||||
|
||||
assert "video_processor.postprocess_video" in profiler_setup["profiler_targets"]
|
||||
assert "_prepare_video_output_for_transport" in profiler_setup["profiler_targets"]
|
||||
@@ -303,11 +284,11 @@ def test_ltx_artifact_uses_source_revision_and_hub_fallback(
|
||||
filename = "ltx-sidecar.safetensors"
|
||||
calls = []
|
||||
|
||||
def fake_download(self, **kwargs):
|
||||
def fake_download(**kwargs):
|
||||
calls.append(kwargs)
|
||||
return "/cache/ltx-sidecar.safetensors"
|
||||
|
||||
monkeypatch.setattr(huggingface_hub.HfApi, "hf_hub_download", fake_download)
|
||||
monkeypatch.setattr(ltx2_components, "hf_hub_download", fake_download)
|
||||
|
||||
assert (
|
||||
resolve_ltx_artifact(
|
||||
@@ -332,9 +313,7 @@ def test_ltx_artifact_prefers_model_root(tmp_path, monkeypatch):
|
||||
filename = "ltx-sidecar.safetensors"
|
||||
expected = tmp_path / filename
|
||||
expected.write_bytes(b"sidecar")
|
||||
monkeypatch.setattr(
|
||||
huggingface_hub.HfApi, "hf_hub_download", lambda *_args, **_kwargs: pytest.fail("unexpected Hub lookup")
|
||||
)
|
||||
monkeypatch.setattr(ltx2_components, "hf_hub_download", lambda **_kwargs: pytest.fail("unexpected Hub lookup"))
|
||||
|
||||
assert resolve_ltx_artifact(
|
||||
str(tmp_path),
|
||||
@@ -348,11 +327,11 @@ def test_ltx_artifact_prefers_model_root(tmp_path, monkeypatch):
|
||||
def test_ltx_artifact_local_model_missing_sidecar_falls_back_to_hub(tmp_path, monkeypatch):
|
||||
calls = []
|
||||
|
||||
def fake_download(self, **kwargs):
|
||||
def fake_download(**kwargs):
|
||||
calls.append(kwargs)
|
||||
return "/cache/ltx-sidecar.safetensors"
|
||||
|
||||
monkeypatch.setattr(huggingface_hub.HfApi, "hf_hub_download", fake_download)
|
||||
monkeypatch.setattr(ltx2_components, "hf_hub_download", fake_download)
|
||||
|
||||
assert (
|
||||
resolve_ltx_artifact(
|
||||
@@ -404,42 +383,14 @@ def test_ltx_checkpoint_validation_rejects_mismatched_scheduler_metadata(expecte
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"pipeline_cls",
|
||||
[
|
||||
LTX2Pipeline,
|
||||
LTX2DistilledOneStagePipeline,
|
||||
LTX2TwoStagePipeline,
|
||||
LTX2DistilledTwoStagePipeline,
|
||||
LTX2T2VDMD2Pipeline,
|
||||
LTX2I2VDMD2Pipeline,
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
("parallel_config", "error"),
|
||||
[
|
||||
(SimpleNamespace(ulysses_mode="advanced_uaa"), "does not support ulysses_mode='advanced_uaa'"),
|
||||
(SimpleNamespace(ring_degree=2), "pure Ulysses sequence parallelism only"),
|
||||
(SimpleNamespace(ulysses_degree=2, ring_degree=2), "pure Ulysses sequence parallelism only"),
|
||||
(SimpleNamespace(allgather_degree=2), "pure Ulysses sequence parallelism only"),
|
||||
],
|
||||
ids=["advanced-uaa", "ring", "hybrid", "allgather-kv"],
|
||||
)
|
||||
def test_ltx_rejects_unsupported_sp_before_component_initialization(pipeline_cls, parallel_config, error, monkeypatch):
|
||||
from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
|
||||
|
||||
def unexpected_model_access(*_args, **_kwargs):
|
||||
pytest.fail("Unsupported SP must be rejected before accessing model files or initializing components")
|
||||
|
||||
monkeypatch.setattr(ltx2_runtime, "detect_ltx_model_version", unexpected_model_access)
|
||||
monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", unexpected_model_access)
|
||||
def test_ltx_rejects_advanced_uaa_before_component_initialization():
|
||||
od_config = SimpleNamespace(
|
||||
model="unused",
|
||||
parallel_config=parallel_config,
|
||||
parallel_config=SimpleNamespace(ulysses_mode="advanced_uaa"),
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match=error):
|
||||
pipeline_cls(od_config=od_config)
|
||||
with pytest.raises(ValueError, match="does not support ulysses_mode='advanced_uaa'"):
|
||||
LTX2Pipeline(od_config=od_config)
|
||||
|
||||
|
||||
def test_ltx23_checkpoint_selects_version_specific_one_stage_profile(tmp_path, monkeypatch):
|
||||
@@ -1458,19 +1409,39 @@ def test_ltx_distilled_forward_rejects_fixed_recipe_overrides(direct_kwargs, sam
|
||||
pipeline.forward(req, **direct_kwargs)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"pipeline_cls",
|
||||
[
|
||||
LTX2Pipeline,
|
||||
LTX2DistilledOneStagePipeline,
|
||||
LTX2TwoStagePipeline,
|
||||
LTX2DistilledTwoStagePipeline,
|
||||
LTX2T2VDMD2Pipeline,
|
||||
LTX2I2VDMD2Pipeline,
|
||||
],
|
||||
)
|
||||
def test_ltx_pipelines_disable_startup_dummy_run(pipeline_cls):
|
||||
assert pipeline_cls.dummy_run_num_frames == 0
|
||||
def test_ltx_distilled_dummy_run_uses_fixed_recipe_values():
|
||||
from vllm_omni.diffusion.request import DUMMY_DIFFUSION_REQUEST_ID, OmniDiffusionRequest
|
||||
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
req = DiffusionRequestBatch(
|
||||
[
|
||||
OmniDiffusionRequest(
|
||||
prompt="dummy run",
|
||||
sampling_params=OmniDiffusionSamplingParams(
|
||||
height=512,
|
||||
width=512,
|
||||
num_frames=LTX2DistilledPipeline.dummy_run_num_frames,
|
||||
num_inference_steps=1,
|
||||
guidance_scale=0.0,
|
||||
),
|
||||
request_id=DUMMY_DIFFUSION_REQUEST_ID,
|
||||
)
|
||||
]
|
||||
)
|
||||
pipeline = _make_ltx_request_pipe(LTX2DistilledPipeline)
|
||||
seen = {}
|
||||
|
||||
def fake_run_recipe(_req, request_inputs, **_kwargs):
|
||||
seen["request_inputs"] = request_inputs
|
||||
return "dummy-output"
|
||||
|
||||
object.__setattr__(pipeline, "_run_recipe", fake_run_recipe)
|
||||
|
||||
assert pipeline.forward(req) == "dummy-output"
|
||||
assert seen["request_inputs"].num_frames == 1
|
||||
assert seen["request_inputs"].num_inference_steps == 8
|
||||
assert seen["request_inputs"].guidance == LTXGuidanceSpec.positive_only()
|
||||
|
||||
|
||||
def test_ltx_custom_sigmas_bypass_scheduler_shifting():
|
||||
|
||||
@@ -1,100 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import timedelta
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
import torch.nn.functional as F
|
||||
|
||||
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
|
||||
from vllm_omni.diffusion.models.ltx2.ltx2_sequence_parallel import (
|
||||
LTX2VideoToAudioParallelAttention,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _run_video_to_audio_parity(rank: int, world_size: int, master_port: int) -> None:
|
||||
dist.init_process_group(
|
||||
"gloo",
|
||||
init_method=f"tcp://127.0.0.1:{master_port}",
|
||||
rank=rank,
|
||||
world_size=world_size,
|
||||
timeout=timedelta(seconds=60),
|
||||
)
|
||||
try:
|
||||
# Non-contiguous SP subgroups match TP-fastest rank ordering. Different
|
||||
# seeds across groups catch collectives that accidentally use WORLD.
|
||||
group_count = world_size // 2
|
||||
ulysses_group = dist.group.WORLD
|
||||
if group_count > 1:
|
||||
for offset in range(group_count):
|
||||
ranks = list(range(offset, world_size, group_count))
|
||||
group = dist.new_group(ranks)
|
||||
if rank in ranks:
|
||||
ulysses_group = group
|
||||
ulysses_rank = dist.get_rank(ulysses_group)
|
||||
sp_group = SimpleNamespace(
|
||||
ring_world_size=1,
|
||||
ulysses_group=ulysses_group,
|
||||
ulysses_world_size=2,
|
||||
ulysses_rank=ulysses_rank,
|
||||
)
|
||||
strategy = LTX2VideoToAudioParallelAttention(sp_group)
|
||||
|
||||
for batch_size, masked in [(1, False), (2, True)]:
|
||||
torch.manual_seed(17 + rank % group_count)
|
||||
query = torch.randn(batch_size, 3, 4, 8)
|
||||
global_key = torch.randn(batch_size, 6, 4, 8)
|
||||
global_value = torch.randn(batch_size, 6, 4, 8)
|
||||
key = global_key.chunk(2, dim=1)[ulysses_rank].contiguous()
|
||||
value = global_value.chunk(2, dim=1)[ulysses_rank].contiguous()
|
||||
mask = None
|
||||
if masked:
|
||||
mask = torch.ones(batch_size, 1, 1, 6, dtype=torch.bool)
|
||||
mask[0, ..., -1] = False
|
||||
mask[1, ..., -2:] = False
|
||||
metadata = AttentionMetadata(attn_mask=mask) if masked else None
|
||||
|
||||
local_query, local_key, local_value, local_metadata, ctx = strategy.pre_attention(
|
||||
query, key, value, metadata
|
||||
)
|
||||
assert local_metadata is metadata
|
||||
local_output = F.scaled_dot_product_attention(
|
||||
local_query.transpose(1, 2),
|
||||
local_key.transpose(1, 2),
|
||||
local_value.transpose(1, 2),
|
||||
attn_mask=local_metadata.attn_mask if local_metadata is not None else None,
|
||||
).transpose(1, 2)
|
||||
actual = strategy.post_attention(local_output, ctx)
|
||||
expected = F.scaled_dot_product_attention(
|
||||
query.transpose(1, 2),
|
||||
global_key.transpose(1, 2),
|
||||
global_value.transpose(1, 2),
|
||||
attn_mask=mask,
|
||||
).transpose(1, 2)
|
||||
|
||||
torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-6)
|
||||
finally:
|
||||
dist.destroy_process_group()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("world_size", [2, 4], ids=["world", "subgroups"])
|
||||
def test_ltx_video_to_audio_sp_matches_replicated_attention(unused_tcp_port, world_size):
|
||||
torch.multiprocessing.spawn(
|
||||
_run_video_to_audio_parity,
|
||||
args=(world_size, unused_tcp_port),
|
||||
nprocs=world_size,
|
||||
)
|
||||
|
||||
|
||||
def test_ltx_video_to_audio_sp_rejects_ring_parallelism():
|
||||
sp_group = SimpleNamespace(ring_world_size=2)
|
||||
|
||||
with pytest.raises(NotImplementedError, match="ring_degree must be 1"):
|
||||
LTX2VideoToAudioParallelAttention(sp_group)
|
||||
Loaded 100 of 442 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user