Merge main and preserve event-driven defaults in shared engine bases
Signed-off-by: Sy03 <1370724210@qq.com>
This commit is contained in:
commit
d2bb944863
371 files changed
+43235
-38379
No files matched your search
@@ -205,6 +205,71 @@ steps:
|
||||
pytest -s -v tests/engine/test_async_omni_engine_abort.py
|
||||
'
|
||||
|
||||
- label: "ROCm · Entrypoints GPU Coverage (R2-02)"
|
||||
agent_pool: mi300_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: NonBlocking
|
||||
timeout_in_minutes: 30
|
||||
env:
|
||||
VLLM_WORKER_MULTIPROC_METHOD: spawn
|
||||
artifact_paths:
|
||||
- "artifacts/rocm-r2-02/**/*"
|
||||
commands:
|
||||
# Native AMD tests run from /app/vllm-omni, while Buildkite uploads
|
||||
# artifacts relative to its checkout directory.
|
||||
- export R2_02_ARTIFACT_DIR="$$BUILDKITE_BUILD_CHECKOUT_PATH/artifacts/rocm-r2-02"
|
||||
- mkdir -p "$$R2_02_ARTIFACT_DIR"
|
||||
- >-
|
||||
ps -eo pid=,ppid=,lstart=,stat=,comm= >
|
||||
"$$R2_02_ARTIFACT_DIR/processes-before.txt"
|
||||
# Verbose pytest output and JUnit retain selected node IDs without a
|
||||
# second collection-only pass over the same GPU-backed test tree.
|
||||
# The inner timeout leaves five minutes for teardown verification and
|
||||
# artifact generation before the Buildkite step timeout.
|
||||
- >-
|
||||
timeout --signal=TERM --kill-after=1m 25m
|
||||
pytest -s -v -ra
|
||||
tests/entrypoints/
|
||||
-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"
|
||||
--durations=0
|
||||
--junitxml="$$R2_02_ARTIFACT_DIR/pytest.xml"
|
||||
2>&1 | tee "$$R2_02_ARTIFACT_DIR/pytest.log"
|
||||
- >-
|
||||
tail -n 100 "$$R2_02_ARTIFACT_DIR/pytest.log" >
|
||||
"$$R2_02_ARTIFACT_DIR/pytest-summary.txt"
|
||||
- |
|
||||
sleep 5
|
||||
ps -eo pid=,ppid=,lstart=,stat=,comm= > "$$R2_02_ARTIFACT_DIR/processes-after.txt"
|
||||
# Compare PID plus process start time so PID reuse cannot hide a leak.
|
||||
# Double dollars preserve awk field references through Buildkite interpolation.
|
||||
awk '
|
||||
NR == FNR {
|
||||
if ($$9 != "ps") {
|
||||
baseline[$$1 "|" $$3 "|" $$4 "|" $$5 "|" $$6 "|" $$7] = 1
|
||||
}
|
||||
next
|
||||
}
|
||||
$$9 != "ps" {
|
||||
identity = $$1 "|" $$3 "|" $$4 "|" $$5 "|" $$6 "|" $$7
|
||||
if (!(identity in baseline)) {
|
||||
print
|
||||
leaked = 1
|
||||
}
|
||||
}
|
||||
END {
|
||||
if (leaked) {
|
||||
print "status=FAIL leaked_processes_detected=1"
|
||||
exit 1
|
||||
}
|
||||
print "status=PASS leaked_processes=0"
|
||||
}
|
||||
' "$$R2_02_ARTIFACT_DIR/processes-before.txt" \
|
||||
"$$R2_02_ARTIFACT_DIR/processes-after.txt" \
|
||||
2>&1 | tee "$$R2_02_ARTIFACT_DIR/process-cleanup.txt"
|
||||
|
||||
# - label: "Omni Model Test Qwen3-Omni"
|
||||
# agent_pool: mi300_2
|
||||
# depends_on: amd-build
|
||||
|
||||
@@ -38,9 +38,13 @@ x-minicpmo-4-5: &minicpmo_4_5
|
||||
x-minicpmo-4-5-duplex: &minicpmo_4_5_duplex
|
||||
- vllm_omni/deploy/minicpmo_4_5_duplex.yaml
|
||||
- vllm_omni/engine/duplex/
|
||||
- vllm_omni/engine/duplex_omni_engine.py
|
||||
- vllm_omni/engine/duplex_orchestrator.py
|
||||
- vllm_omni/engine/omni_engine_base.py
|
||||
- vllm_omni/entrypoints/duplex/
|
||||
- vllm_omni/entrypoints/duplex_omni.py
|
||||
- vllm_omni/entrypoints/async_omni_base.py
|
||||
- vllm_omni/clients/
|
||||
- vllm_omni/entrypoints/duplex_request_client.py
|
||||
- vllm_omni/outputs/duplex.py
|
||||
- vllm_omni/model_executor/duplex_sampling.py
|
||||
- vllm_omni/experimental/fullduplex/
|
||||
@@ -249,7 +253,11 @@ source_file_dependencies:
|
||||
- tests/assets/minicpmo_4_5/
|
||||
- tests/clients/
|
||||
- tests/engine/duplex/
|
||||
- tests/engine/test_duplex_omni_engine.py
|
||||
- tests/engine/test_duplex_orchestrator.py
|
||||
- tests/entrypoints/duplex/
|
||||
- tests/entrypoints/test_duplex_omni.py
|
||||
- tests/entrypoints/openai_api/test_duplex_api_server.py
|
||||
- tests/model_executor/models/minicpmo_4_5/duplex/
|
||||
|
||||
omni_minicpmo_4_5_duplex_perf:
|
||||
|
||||
@@ -66,7 +66,9 @@ steps:
|
||||
# 0.26 startup regression is understood and fixed.
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and H100" --run-level "advanced_model"
|
||||
# Turn-based server VAD over /v1/realtime is deferred to RFC #7181 PR 5
|
||||
# (see the P0 thread on #7413); the test stays in the tree and runs again there.
|
||||
- pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and H100" --run-level "advanced_model" --deselect tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py::TestQwen3OmniRealtimeWebSocket::test_server_vad_multi_turn_without_client_commit
|
||||
mirror_hardwares: h100_2
|
||||
|
||||
- label: "Entrypoint Test with L4"
|
||||
|
||||
@@ -191,13 +191,14 @@ steps:
|
||||
# from downgrading transformers; see
|
||||
# examples/offline_inference/text_to_speech/README.md.
|
||||
- label: "TTS · Gepard-1.0 · L4"
|
||||
timeout_in_minutes: 60
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
- printf '%s\n' "transformers>=5.5.3" "huggingface-hub>=1.0" > /tmp/gepard-overrides.txt
|
||||
- uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" "nemo_toolkit[tts]==2.7.3" --overrides /tmp/gepard-overrides.txt
|
||||
- >-
|
||||
pytest -s -v tests/e2e/offline_inference/test_gepard_tts.py
|
||||
tests/e2e/online_serving/test_gepard_tts.py
|
||||
-m "tts and cuda"
|
||||
--run-level "full_model"
|
||||
mirror_hardwares: l4_1
|
||||
|
||||
@@ -139,6 +139,89 @@ You can configure per-stage sampling parameters for multi-stage models.
|
||||
>
|
||||
> Do not use `frame` and `references` together. Task routing is automatic from which inputs you connect.
|
||||
|
||||
For MiniMax-H3 Ref2VA, **Video References** accepts up to 9 images (`image_1`–`image_9`),
|
||||
3 videos (`video_1`–`video_3`), and 3 audio clips (`audio_1`–`audio_3`), with at most
|
||||
12 connected inputs in total. Any mixture containing at least one image or video
|
||||
is supported; empty and audio-only references are rejected. For example, you can
|
||||
combine 6 images, 3 videos, and 3 audio clips in one request.
|
||||
|
||||
Within each media type, references follow slot-number order, skipping unconnected
|
||||
slots. For example, connecting `image_2` and `image_9` sends `image_2` as the first
|
||||
image and `image_9` as the second. Each image slot uses the first image in its batch.
|
||||
Existing connections to `image_1`, `image_2`, `audio_1`, `audio_2`, `video_1`, and
|
||||
`video_2` remain valid in saved workflows.
|
||||
|
||||
#### MiniMax-H3 Reference to Video
|
||||
|
||||
Load [MiniMax-H3 Reference to Video](example_workflows/vLLM-Omni%20MiniMax-H3%20Reference%20to%20Video.json)
|
||||
from **Templates → ComfyUI-vLLM-Omni**, or drag the JSON onto the canvas.
|
||||
It adapts the [official ComfyUI R2V workflow](https://github.com/Comfy-Org/workflow_templates/blob/main/templates/video_minimax_h3_r2v.json)
|
||||
to remote vLLM-Omni execution. Model loading, sampling, and VAE decoding run on
|
||||
the server; ComfyUI loads references and saves the returned video.
|
||||
|
||||
Use an extension version containing the audio-preserving output fix from
|
||||
[#7456](https://github.com/vllm-project/vllm-omni/pull/7456). Older versions can
|
||||
drop the server's audio while decoding the response.
|
||||
|
||||
1. Start a Ref2VA-capable service using the [MiniMax-H3 recipe](../../recipes/MiniMaxAI/MiniMax-H3.md).
|
||||
Set **Generate Video** to its `/v1` URL and served model name.
|
||||
2. Select your own image in **Load Image**. The template's filenames are placeholders
|
||||
relative to ComfyUI's input directory; no sample assets or model weights are bundled.
|
||||
3. Connect **Load Video** and **Load Audio** when needed. Duplicate the loaders to
|
||||
fill more reference slots, or disconnect the image when using video-only input.
|
||||
Keep the references connected to **Generate Video** and leave `frame` disconnected.
|
||||
4. Match prompt tags such as `<Picture 1>`, `<Video 1>`, and `<Audio 1>` to the
|
||||
connected references, counting each media type separately and skipping empty slots.
|
||||
5. Run the workflow. **Save Video** writes an MP4 under `output/video/` and preserves
|
||||
the generated audio when the audio-output prerequisite is installed.
|
||||
|
||||
The default is **1344×768, 24 FPS, 124 frames** (about 5.17 seconds), 50 sampling
|
||||
points, video flow shift 12, audio flow shift 3, and seed 42. The duration widget
|
||||
is set to 5.167 seconds, which converts to 124 frames at 24 FPS. Other H3 canvas presets
|
||||
are 1024×768, 768×768, 768×1024, and 768×1344. Keep frame counts at `17k+5` within
|
||||
the 4–15 second output range; examples are 107, 124, 209, and 345 frames.
|
||||
Reference video/audio clips must each be 2–15 seconds, with at most 15 seconds of
|
||||
reference video and at most 15 seconds of standalone reference audio per request.
|
||||
|
||||
For **Ref2VA Turbo**, use the same template. Configure a Ref2VA-only server with
|
||||
`--lora-backend peft --lora-path "$TURBO_LORA" --task-type ref2va`, following the
|
||||
recipe's hardware setup. For `minimax_h3_ref2v_turbo_4step_v0.1_bf16.safetensors`,
|
||||
set **LoRA** `local_path` to that file's path on the server and connect it to
|
||||
**Generate Video**. Use `num_inference_steps=5`, `flow_shift=12`,
|
||||
`audio_flow_shift=3`, and LoRA `scale=1`. The API counts sigma points, so five
|
||||
points produce four denoiser evaluations. Use the Diffusers artifact, not its
|
||||
`_comfyui_` export. To return to base mode, disconnect LoRA and restore 50 steps.
|
||||
The 8-step v1.0 Ref2VA adapter instead requires nine sampling points and video
|
||||
flow shift 6. FastH3 is a separate startup-fused T2VA deployment.
|
||||
|
||||
Validate the template's wiring and defaults locally with:
|
||||
|
||||
```bash
|
||||
python -m pytest tests/e2e/features/comfyui/test_h3_reference_workflow.py -q
|
||||
```
|
||||
|
||||
For a real-model check, run the workflow against H3 and inspect the saved file:
|
||||
|
||||
```bash
|
||||
ffprobe -v error -show_entries stream=codec_type,width,height,r_frame_rate,sample_rate,channels \
|
||||
-show_entries format=duration -of json "$OUTPUT_MP4"
|
||||
```
|
||||
|
||||
Check for 24 FPS video, stereo audio, the requested canvas, and aligned duration;
|
||||
also play the result to assess reference conditioning and audio/video synchronization.
|
||||
Record the server commit, model/adapter, input assets, prompt, settings, and output
|
||||
alongside the result. Local schema or mocked-server checks do not establish H3
|
||||
generation quality or replace this real-model validation.
|
||||
|
||||
### MiniMax H3 Text to Video
|
||||
|
||||
Import [the H3 text-to-video template](example_workflows/MiniMax_H3_Text_to_Video.json)
|
||||
for remote video generation with audio. It includes connected base presets and
|
||||
optional Turbo sampling, H3 parameters, and Remote LoRA nodes. See the
|
||||
[workflow guide](docs/minimax-h3-t2v.md) for server setup, artifact-specific Turbo
|
||||
settings, and saved-video/audio validation. Set Remote LoRA’s `local_path` to
|
||||
the downloaded Turbo artifact on your server before enabling Turbo.
|
||||
|
||||
#### H3 video upscale (WF-07)
|
||||
|
||||
The **vLLM-Omni MiniMax H3 Video Upscale** template generates video remotely, upscales it with SeedVR2, and saves the original and upscaled videos with the generated audio and FPS. See [workflow setup](docs/wf07-h3-upscale.md).
|
||||
|
||||
@@ -10,6 +10,9 @@ from .utils.api_client import VLLMOmniClient
|
||||
from .utils.logger import get_logger
|
||||
from .utils.models import lookup_model_spec
|
||||
from .utils.types import (
|
||||
MAX_REFERENCE_AUDIOS,
|
||||
MAX_REFERENCE_IMAGES,
|
||||
MAX_REFERENCE_VIDEOS,
|
||||
AudioFormat,
|
||||
AutoregressionSamplingParams,
|
||||
DiffusionSamplingParams,
|
||||
@@ -918,6 +921,10 @@ class VLLMOmniVideoReferences:
|
||||
"audio_2": ("AUDIO",),
|
||||
"video_1": ("VIDEO",),
|
||||
"video_2": ("VIDEO",),
|
||||
# Append ports to preserve connections in saved workflows.
|
||||
**{f"image_{i}": ("IMAGE",) for i in range(3, MAX_REFERENCE_IMAGES + 1)},
|
||||
**{f"audio_{i}": ("AUDIO",) for i in range(3, MAX_REFERENCE_AUDIOS + 1)},
|
||||
**{f"video_{i}": ("VIDEO",) for i in range(3, MAX_REFERENCE_VIDEOS + 1)},
|
||||
},
|
||||
}
|
||||
|
||||
@@ -934,21 +941,26 @@ class VLLMOmniVideoReferences:
|
||||
audio_2: AudioInput | None = None,
|
||||
video_1: VideoInput | None = None,
|
||||
video_2: VideoInput | None = None,
|
||||
image_3: torch.Tensor | None = None,
|
||||
image_4: torch.Tensor | None = None,
|
||||
image_5: torch.Tensor | None = None,
|
||||
image_6: torch.Tensor | None = None,
|
||||
image_7: torch.Tensor | None = None,
|
||||
image_8: torch.Tensor | None = None,
|
||||
image_9: torch.Tensor | None = None,
|
||||
audio_3: AudioInput | None = None,
|
||||
video_3: VideoInput | None = None,
|
||||
**kwargs,
|
||||
):
|
||||
if kwargs:
|
||||
logger.info("Uncaught kwargs: %s", kwargs)
|
||||
refs = VideoReferences()
|
||||
if image_1 is not None:
|
||||
refs["image_1"] = image_1
|
||||
if image_2 is not None:
|
||||
refs["image_2"] = image_2
|
||||
if audio_1 is not None:
|
||||
refs["audio_1"] = audio_1
|
||||
if audio_2 is not None:
|
||||
refs["audio_2"] = audio_2
|
||||
if video_1 is not None:
|
||||
refs["video_1"] = video_1
|
||||
if video_2 is not None:
|
||||
refs["video_2"] = video_2
|
||||
for kind, values in (
|
||||
("image", (image_1, image_2, image_3, image_4, image_5, image_6, image_7, image_8, image_9)),
|
||||
("video", (video_1, video_2, video_3)),
|
||||
("audio", (audio_1, audio_2, audio_3)),
|
||||
):
|
||||
for index, value in enumerate(values, start=1):
|
||||
if value is not None:
|
||||
refs[f"{kind}_{index}"] = value
|
||||
return (refs,)
|
||||
@@ -20,6 +20,7 @@ from comfy_api.input import AudioInput, VideoInput
|
||||
|
||||
from .format import (
|
||||
audio_to_base64,
|
||||
audio_to_bytes,
|
||||
base64_to_audio,
|
||||
base64_to_image_tensor,
|
||||
bytes_to_audio,
|
||||
@@ -31,7 +32,13 @@ from .format import (
|
||||
)
|
||||
from .logger import get_logger, pretty_printer
|
||||
from .models import lookup_model_spec
|
||||
from .types import AudioFormat
|
||||
from .types import (
|
||||
MAX_REFERENCE_AUDIOS,
|
||||
MAX_REFERENCE_IMAGES,
|
||||
MAX_REFERENCE_VIDEOS,
|
||||
MAX_TOTAL_REFERENCES,
|
||||
AudioFormat,
|
||||
)
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
@@ -304,38 +311,44 @@ class VLLMOmniClient:
|
||||
|
||||
# === multimodal inputs (first-last-frames, references, etc.) ===
|
||||
input_reference_image: torch.Tensor | None = None
|
||||
audio_reference: AudioInput | None = None
|
||||
reference_videos: list[VideoInput] = []
|
||||
video_task: str | None = None
|
||||
|
||||
if frame is not None:
|
||||
input_reference_image = frame
|
||||
video_task = "fl2va"
|
||||
elif references is not None:
|
||||
images = [references[k] for k in ("image_1", "image_2") if k in references and references[k] is not None]
|
||||
audios = [references[k] for k in ("audio_1", "audio_2") if k in references and references[k] is not None]
|
||||
videos = [references[k] for k in ("video_1", "video_2") if k in references and references[k] is not None]
|
||||
|
||||
if not images and not audios and not videos:
|
||||
raise ValueError("references is empty; connect at least one image/audio/video.")
|
||||
|
||||
if videos:
|
||||
if images or audios:
|
||||
raise ValueError(
|
||||
"Video references cannot be combined with image or audio references. "
|
||||
"Connect only video_1/video_2 for multi-video Ref2VA."
|
||||
)
|
||||
reference_videos = videos
|
||||
video_task = "ref2va"
|
||||
elif len(images) == 1 and len(audios) == 1:
|
||||
input_reference_image = images[0]
|
||||
audio_reference = audios[0]
|
||||
video_task = "ref2va"
|
||||
else:
|
||||
reference_formats = (
|
||||
("image", MAX_REFERENCE_IMAGES, "png", "image/png", image_tensor_to_png_bytes),
|
||||
("video", MAX_REFERENCE_VIDEOS, "mp4", "video/mp4", video_to_bytes),
|
||||
("audio", MAX_REFERENCE_AUDIOS, "mp3", "audio/mpeg", audio_to_bytes),
|
||||
)
|
||||
supported_inputs = {f"{kind}_{i}" for kind, limit, *_ in reference_formats for i in range(1, limit + 1)}
|
||||
connected = {name: value for name, value in references.items() if value is not None}
|
||||
unsupported = connected.keys() - supported_inputs
|
||||
if unsupported:
|
||||
raise ValueError(f"Unsupported reference input(s): {', '.join(sorted(unsupported))}.")
|
||||
if not any(name.startswith(("image_", "video_")) for name in connected):
|
||||
raise ValueError(
|
||||
"Invalid references combination. Supported modes: "
|
||||
"(1) one or more videos only, or (2) exactly one image and one audio."
|
||||
"references requires at least one image or video; audio-only inputs are not supported."
|
||||
)
|
||||
if len(connected) > MAX_TOTAL_REFERENCES:
|
||||
raise ValueError(
|
||||
f"references supports at most {MAX_TOTAL_REFERENCES} inputs in total "
|
||||
f"(up to {MAX_REFERENCE_IMAGES} images, {MAX_REFERENCE_VIDEOS} videos, "
|
||||
f"and {MAX_REFERENCE_AUDIOS} audios)."
|
||||
)
|
||||
for kind, limit, extension, content_type, encode in reference_formats:
|
||||
for index in range(1, limit + 1):
|
||||
name = f"{kind}_{index}"
|
||||
if name in connected:
|
||||
filename = f"{name}.{extension}"
|
||||
form.add_field(
|
||||
"input_references",
|
||||
encode(connected[name], filename),
|
||||
filename=filename,
|
||||
content_type=content_type,
|
||||
)
|
||||
video_task = "ref2va"
|
||||
else:
|
||||
video_task = "t2va"
|
||||
|
||||
@@ -348,21 +361,6 @@ class VLLMOmniClient:
|
||||
content_type="image/png",
|
||||
)
|
||||
|
||||
if audio_reference is not None:
|
||||
form.add_field(
|
||||
"audio_reference",
|
||||
json.dumps({"audio_url": audio_to_base64(audio_reference)}, ensure_ascii=False),
|
||||
)
|
||||
|
||||
for idx, video in enumerate(reference_videos, start=1):
|
||||
video_filename = f"reference_{idx}.mp4"
|
||||
form.add_field(
|
||||
"input_references",
|
||||
video_to_bytes(video, video_filename),
|
||||
filename=video_filename,
|
||||
content_type="video/mp4",
|
||||
)
|
||||
|
||||
# === model specific params. Either use a specialized builder, or add flattened fields as-is ===
|
||||
if model_params is not None:
|
||||
model_params = dict(model_params)
|
||||
|
||||
@@ -30,6 +30,12 @@ class MiniMaxH3ModelSpecificParams(dict):
|
||||
pass
|
||||
|
||||
|
||||
MAX_REFERENCE_IMAGES = 9
|
||||
MAX_REFERENCE_VIDEOS = 3
|
||||
MAX_REFERENCE_AUDIOS = 3
|
||||
MAX_TOTAL_REFERENCES = 12
|
||||
|
||||
|
||||
class VideoReferences(dict):
|
||||
pass
|
||||
|
||||
|
||||
@@ -0,0 +1,94 @@
|
||||
# MiniMax H3 Text to Video (WF-01)
|
||||
|
||||
Import [MiniMax_H3_Text_to_Video.json](../example_workflows/MiniMax_H3_Text_to_Video.json)
|
||||
into ComfyUI, or select it under **Templates → ComfyUI-vLLM-Omni**.
|
||||
The template connects existing nodes to a remote H3 service and saves the
|
||||
returned video with audio. ComfyUI does not load H3 weights.
|
||||
|
||||
## Setup
|
||||
|
||||
Install the [ComfyUI extension](../README.md#installation) and follow the
|
||||
[H3 deployment recipe](../../../recipes/MiniMaxAI/MiniMax-H3.md) for your hardware.
|
||||
For a two-GPU FL2VA server with sufficient memory:
|
||||
|
||||
```bash
|
||||
export MODEL=/path/to/MiniMax-H3/FL2VA
|
||||
CUDA_VISIBLE_DEVICES=0,1 VLLM_WORKER_MULTIPROC_METHOD=spawn \
|
||||
VLLM_OMNI_VIDEO_SYNC_TIMEOUT=14400 \
|
||||
vllm serve "$MODEL" --omni --host 127.0.0.1 --port 8000 \
|
||||
--served-model-name MiniMaxAI/MiniMax-H3 --trust-remote-code \
|
||||
--task-type fl2va --num-gpus 2 --tensor-parallel-size 2 \
|
||||
--text-encoder-tp-size 2 --enforce-eager
|
||||
```
|
||||
|
||||
Set Generate Video's **url** to the address reachable from ComfyUI, including
|
||||
`/v1`, and match its **model** to the served name. The defaults are
|
||||
`http://localhost:8000/v1` and `MiniMaxAI/MiniMax-H3`.
|
||||
|
||||
## Base and Turbo
|
||||
|
||||
Base presets are connected by default. Write the scene and its sound in the
|
||||
same **prompt**; leave **frame**, **references**, and **fast_h3** disconnected.
|
||||
The template uses 1344×768, 24 FPS, seed 1101 and **duration** 5.167 seconds
|
||||
(124 frames, satisfying H3's `17k+5` constraint). The existing client infers
|
||||
`16:9` from these dimensions. For portrait output, use 768×1344.
|
||||
|
||||
| Setting | Base | Turbo v1.0 768p |
|
||||
| --- | --- | --- |
|
||||
| Inference steps | 50 | 5 (4 forwards) |
|
||||
| Video flow shift | 12 | 6 |
|
||||
| Audio flow shift | 3 | 3 |
|
||||
| Remote LoRA | Disconnected | Connected, scale 1 |
|
||||
|
||||
For Turbo, use the **Diffusers-layout** artifact from
|
||||
[lightx2v/Minimax-h3-Turbo](https://huggingface.co/lightx2v/Minimax-h3-Turbo):
|
||||
|
||||
```text
|
||||
minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors
|
||||
```
|
||||
|
||||
Add `--lora-backend peft --lora-path "$TURBO_LORA"` to the server command,
|
||||
where `TURBO_LORA` is the full server-side path to that file. Enter the same
|
||||
path in Remote LoRA's **local_path**; retain name `h3-turbo-v1.0-768p`, scale
|
||||
`1.0`, and **int_id** `0` so the existing API derives the ID from the path.
|
||||
The template leaves the path blank for users to configure. The `_comfyui_`
|
||||
export is not supported by this backend.
|
||||
|
||||
Connect **Turbo sampling**, **Turbo H3 params**, and **Remote LoRA** to Generate
|
||||
Video's `sampling_params`, `model_params`, and `lora` inputs, replacing the
|
||||
Base links. To return to Base, restore both Base links and disconnect LoRA.
|
||||
Other Turbo artifacts may need different settings; see the
|
||||
[artifact table](../../../recipes/MiniMaxAI/MiniMax-H3.md#turbo-lora).
|
||||
|
||||
If importing an older template with a frame-count widget, reimport this JSON
|
||||
or set **duration** to **5.167**; do not interpret 124 frames as 124 seconds.
|
||||
|
||||
## Validation
|
||||
|
||||
Queue the configured workflow and inspect the MP4 saved under
|
||||
`ComfyUI/output/video/MiniMax_H3_WF01_*.mp4`:
|
||||
|
||||
```bash
|
||||
ffprobe -v error -show_streams -show_format -of json "$OUTPUT"
|
||||
ffmpeg -v error -i "$OUTPUT" -f null -
|
||||
```
|
||||
|
||||
Check the requested dimensions, frame count, 24 FPS and both video/audio
|
||||
streams, then watch and listen to the saved file. The current upstream
|
||||
extension includes the shared audio-preserving decoder fix.
|
||||
|
||||
On September 16, the narrowed workflow was validated through ComfyUI using
|
||||
existing upstream Python code at `4a8ac297`, two H20-3e GPUs, vLLM 0.29.0,
|
||||
PyTorch 2.13.0+cu129 and ComfyUI `1d48d9cf`. The run used the template's
|
||||
forest-stream prompt, explicit server-side Turbo LoRA path, and settings
|
||||
above, with `TORCH_SDPA`, `--usp 1 --ring 1`, and tiled VAE parallelism on two
|
||||
GPUs. ComfyUI execution took 131.55 seconds. SaveVideo produced 124 H.264
|
||||
frames at 1344×768 and 24 FPS, plus 32 kHz stereo AAC audio. Video/audio
|
||||
durations were 5.166667/5.167000 seconds; full decoding passed and audio was
|
||||
nonzero. This is model-generated output, not a ground-truth reference or a
|
||||
subjective audio-quality assessment. Base graph validation passed with Turbo
|
||||
nodes disconnected; full 50-step Base quality was not evaluated.
|
||||
|
||||
Based on the official
|
||||
[ComfyUI H3 T2V template](https://github.com/Comfy-Org/workflow_templates/blob/main/templates/video_minimax_h3_t2v.json),
|
||||
using remote inference instead of local model loaders.
|
||||
@@ -0,0 +1,362 @@
|
||||
{
|
||||
"last_node_id": 8,
|
||||
"last_link_id": 3,
|
||||
"nodes": [
|
||||
{
|
||||
"id": 1,
|
||||
"type": "VLLMOmniDiffusionSampling",
|
||||
"pos": [
|
||||
20,
|
||||
20
|
||||
],
|
||||
"size": [
|
||||
310,
|
||||
280
|
||||
],
|
||||
"flags": {},
|
||||
"order": 0,
|
||||
"mode": 0,
|
||||
"inputs": [],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "diffusion sampling params",
|
||||
"type": "SAMPLING_PARAMS",
|
||||
"links": [
|
||||
1
|
||||
],
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "VLLMOmniDiffusionSampling"
|
||||
},
|
||||
"widgets_values": [
|
||||
1,
|
||||
50,
|
||||
1.0,
|
||||
1.0,
|
||||
false,
|
||||
false,
|
||||
1101,
|
||||
"fixed"
|
||||
],
|
||||
"title": "Base sampling"
|
||||
},
|
||||
{
|
||||
"id": 2,
|
||||
"type": "VLLMOmniMiniMaxH3Params",
|
||||
"pos": [
|
||||
20,
|
||||
340
|
||||
],
|
||||
"size": [
|
||||
310,
|
||||
140
|
||||
],
|
||||
"flags": {},
|
||||
"order": 1,
|
||||
"mode": 0,
|
||||
"inputs": [],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "MiniMax-H3 video params",
|
||||
"type": "VIDEO_PARAMS",
|
||||
"links": [
|
||||
2
|
||||
],
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "VLLMOmniMiniMaxH3Params"
|
||||
},
|
||||
"widgets_values": [
|
||||
3.0,
|
||||
12.0
|
||||
],
|
||||
"title": "Base H3 params"
|
||||
},
|
||||
{
|
||||
"id": 3,
|
||||
"type": "VLLMOmniGenerateVideo",
|
||||
"pos": [
|
||||
430,
|
||||
20
|
||||
],
|
||||
"size": [
|
||||
430,
|
||||
440
|
||||
],
|
||||
"flags": {},
|
||||
"order": 2,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "frame",
|
||||
"type": "IMAGE",
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "references",
|
||||
"type": "VIDEO_REFERENCES",
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "sampling_params",
|
||||
"type": "SAMPLING_PARAMS",
|
||||
"link": 1
|
||||
},
|
||||
{
|
||||
"name": "lora",
|
||||
"type": "REMOTE_LORA",
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "model_params",
|
||||
"type": "VIDEO_PARAMS",
|
||||
"link": 2
|
||||
},
|
||||
{
|
||||
"name": "fast_h3",
|
||||
"type": "FASTH3_DEPLOYMENT",
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "duration",
|
||||
"type": "FLOAT",
|
||||
"widget": {
|
||||
"name": "duration"
|
||||
},
|
||||
"link": null
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "video",
|
||||
"type": "VIDEO",
|
||||
"links": [
|
||||
3
|
||||
],
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "VLLMOmniGenerateVideo"
|
||||
},
|
||||
"widgets_values": [
|
||||
"http://localhost:8000/v1",
|
||||
"MiniMaxAI/MiniMax-H3",
|
||||
"A small stream flows through a sunlit forest. The camera slowly moves toward the water. Birds chirp in the trees and water gurgles over smooth stones. Natural continuous sound, no music.",
|
||||
"",
|
||||
1344,
|
||||
768,
|
||||
24,
|
||||
5.167
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 4,
|
||||
"type": "SaveVideo",
|
||||
"pos": [
|
||||
940,
|
||||
20
|
||||
],
|
||||
"size": [
|
||||
440,
|
||||
440
|
||||
],
|
||||
"flags": {},
|
||||
"order": 3,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "video",
|
||||
"type": "VIDEO",
|
||||
"link": 3
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "video",
|
||||
"type": "VIDEO",
|
||||
"links": null,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "SaveVideo"
|
||||
},
|
||||
"widgets_values": [
|
||||
"video/MiniMax_H3_WF01",
|
||||
"mp4",
|
||||
"auto"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 5,
|
||||
"type": "VLLMOmniRemoteLoRA",
|
||||
"pos": [
|
||||
430,
|
||||
540
|
||||
],
|
||||
"size": [
|
||||
430,
|
||||
160
|
||||
],
|
||||
"flags": {},
|
||||
"order": 4,
|
||||
"mode": 0,
|
||||
"inputs": [],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "lora",
|
||||
"type": "REMOTE_LORA",
|
||||
"links": null,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "VLLMOmniRemoteLoRA"
|
||||
},
|
||||
"widgets_values": [
|
||||
"",
|
||||
"h3-turbo-v1.0-768p",
|
||||
1.0,
|
||||
0
|
||||
],
|
||||
"title": "Optional Turbo: server LoRA path"
|
||||
},
|
||||
{
|
||||
"id": 6,
|
||||
"type": "VLLMOmniDiffusionSampling",
|
||||
"pos": [
|
||||
20,
|
||||
540
|
||||
],
|
||||
"size": [
|
||||
310,
|
||||
280
|
||||
],
|
||||
"flags": {},
|
||||
"order": 5,
|
||||
"mode": 0,
|
||||
"inputs": [],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "diffusion sampling params",
|
||||
"type": "SAMPLING_PARAMS",
|
||||
"links": null,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "VLLMOmniDiffusionSampling"
|
||||
},
|
||||
"widgets_values": [
|
||||
1,
|
||||
5,
|
||||
1.0,
|
||||
1.0,
|
||||
false,
|
||||
false,
|
||||
1101,
|
||||
"fixed"
|
||||
],
|
||||
"title": "Turbo sampling (4 forwards)"
|
||||
},
|
||||
{
|
||||
"id": 7,
|
||||
"type": "VLLMOmniMiniMaxH3Params",
|
||||
"pos": [
|
||||
430,
|
||||
740
|
||||
],
|
||||
"size": [
|
||||
310,
|
||||
140
|
||||
],
|
||||
"flags": {},
|
||||
"order": 6,
|
||||
"mode": 0,
|
||||
"inputs": [],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "MiniMax-H3 video params",
|
||||
"type": "VIDEO_PARAMS",
|
||||
"links": null,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "VLLMOmniMiniMaxH3Params"
|
||||
},
|
||||
"widgets_values": [
|
||||
3.0,
|
||||
6.0
|
||||
],
|
||||
"title": "Turbo H3 params"
|
||||
},
|
||||
{
|
||||
"id": 8,
|
||||
"type": "Note",
|
||||
"pos": [
|
||||
940,
|
||||
520
|
||||
],
|
||||
"size": [
|
||||
440,
|
||||
440
|
||||
],
|
||||
"flags": {},
|
||||
"order": 7,
|
||||
"mode": 0,
|
||||
"inputs": [],
|
||||
"outputs": [],
|
||||
"properties": {
|
||||
"Node name for S&R": "Note"
|
||||
},
|
||||
"widgets_values": [
|
||||
"MiniMax H3 Text to Video (WF-01)\n\nBASE: connected presets, 50 steps / video flow 12 / audio flow 3.\n\nTURBO: set Remote LoRA local_path to the SERVER path of minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors from lightx2v/Minimax-h3-Turbo. Keep name h3-turbo-v1.0-768p, scale 1.0 and int_id 0. Connect Turbo sampling, Turbo H3 params, and LoRA to Generate Video, replacing the two Base links. Use 5 steps (4 forwards), video flow 6 and audio flow 3. Do not use the ComfyUI-layout artifact.\n\nKeep frame, references and fast_h3 disconnected. 1344x768, 24 FPS, duration 5.167 seconds = 124 frames (17*7+5). The client selects the named 16:9 preset from the dimensions. Save Video preserves video and audio.\n\nSet URL and served model name for your server. See docs/minimax-h3-t2v.md for setup and validation. Source: Comfy-Org/workflow_templates/templates/video_minimax_h3_t2v.json."
|
||||
],
|
||||
"title": "Setup / Base and Turbo"
|
||||
}
|
||||
],
|
||||
"links": [
|
||||
[
|
||||
1,
|
||||
1,
|
||||
0,
|
||||
3,
|
||||
2,
|
||||
"SAMPLING_PARAMS"
|
||||
],
|
||||
[
|
||||
2,
|
||||
2,
|
||||
0,
|
||||
3,
|
||||
4,
|
||||
"VIDEO_PARAMS"
|
||||
],
|
||||
[
|
||||
3,
|
||||
3,
|
||||
0,
|
||||
4,
|
||||
0,
|
||||
"VIDEO"
|
||||
]
|
||||
],
|
||||
"groups": [],
|
||||
"config": {},
|
||||
"extra": {
|
||||
"ds": {
|
||||
"scale": 0.8,
|
||||
"offset": [
|
||||
50,
|
||||
50
|
||||
]
|
||||
}
|
||||
},
|
||||
"version": 0.4
|
||||
}
|
||||
+786
@@ -0,0 +1,786 @@
|
||||
{
|
||||
"id": "e7a4e606-6510-4c4e-b127-9a8b6bf3ea71",
|
||||
"revision": 0,
|
||||
"last_node_id": 11,
|
||||
"last_link_id": 5,
|
||||
"nodes": [
|
||||
{
|
||||
"id": 1,
|
||||
"type": "LoadImage",
|
||||
"pos": [
|
||||
40,
|
||||
100
|
||||
],
|
||||
"size": [
|
||||
320,
|
||||
300
|
||||
],
|
||||
"flags": {},
|
||||
"order": 0,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"localized_name": "image",
|
||||
"name": "image",
|
||||
"type": "COMBO",
|
||||
"widget": {
|
||||
"name": "image"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "choose file to upload",
|
||||
"name": "upload",
|
||||
"type": "IMAGEUPLOAD",
|
||||
"widget": {
|
||||
"name": "upload"
|
||||
},
|
||||
"link": null
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"localized_name": "IMAGE",
|
||||
"name": "IMAGE",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
1
|
||||
]
|
||||
},
|
||||
{
|
||||
"localized_name": "MASK",
|
||||
"name": "MASK",
|
||||
"type": "MASK",
|
||||
"links": null
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "LoadImage"
|
||||
},
|
||||
"widgets_values": [
|
||||
"h3_reference.png",
|
||||
"image"
|
||||
],
|
||||
"title": "Reference image \u2014 select a local file"
|
||||
},
|
||||
{
|
||||
"type": "LoadVideo",
|
||||
"size": [
|
||||
320,
|
||||
200
|
||||
],
|
||||
"flags": {},
|
||||
"inputs": [],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "VIDEO",
|
||||
"type": "VIDEO",
|
||||
"links": null
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "LoadVideo"
|
||||
},
|
||||
"widgets_values": [
|
||||
"h3_reference.mp4"
|
||||
],
|
||||
"id": 2,
|
||||
"pos": [
|
||||
40,
|
||||
470
|
||||
],
|
||||
"order": 1,
|
||||
"mode": 0,
|
||||
"title": "Optional reference video"
|
||||
},
|
||||
{
|
||||
"id": 3,
|
||||
"type": "LoadAudio",
|
||||
"pos": [
|
||||
40,
|
||||
1100
|
||||
],
|
||||
"size": [
|
||||
320,
|
||||
160
|
||||
],
|
||||
"flags": {},
|
||||
"order": 2,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"localized_name": "audio",
|
||||
"name": "audio",
|
||||
"type": "COMBO",
|
||||
"widget": {
|
||||
"name": "audio"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "audioUI",
|
||||
"name": "audioUI",
|
||||
"type": "AUDIO_UI",
|
||||
"widget": {
|
||||
"name": "audioUI"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "choose file to upload",
|
||||
"name": "upload",
|
||||
"type": "AUDIOUPLOAD",
|
||||
"widget": {
|
||||
"name": "upload"
|
||||
},
|
||||
"link": null
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"localized_name": "AUDIO",
|
||||
"name": "AUDIO",
|
||||
"type": "AUDIO",
|
||||
"links": null
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "LoadAudio"
|
||||
},
|
||||
"widgets_values": [
|
||||
"h3_reference.wav",
|
||||
null,
|
||||
""
|
||||
],
|
||||
"title": "Optional reference audio"
|
||||
},
|
||||
{
|
||||
"id": 4,
|
||||
"type": "VLLMOmniVideoReferences",
|
||||
"pos": [
|
||||
420,
|
||||
100
|
||||
],
|
||||
"size": [
|
||||
300,
|
||||
365
|
||||
],
|
||||
"flags": {},
|
||||
"order": 3,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "image_1",
|
||||
"type": "IMAGE",
|
||||
"shape": 7,
|
||||
"link": 1
|
||||
},
|
||||
{
|
||||
"name": "image_2",
|
||||
"type": "IMAGE",
|
||||
"shape": 7,
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "audio_1",
|
||||
"type": "AUDIO",
|
||||
"shape": 7,
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "audio_2",
|
||||
"type": "AUDIO",
|
||||
"shape": 7,
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "video_1",
|
||||
"type": "VIDEO",
|
||||
"shape": 7,
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "video_2",
|
||||
"type": "VIDEO",
|
||||
"shape": 7,
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "image_3",
|
||||
"type": "IMAGE",
|
||||
"shape": 7,
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "image_4",
|
||||
"type": "IMAGE",
|
||||
"shape": 7,
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "image_5",
|
||||
"type": "IMAGE",
|
||||
"shape": 7,
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "image_6",
|
||||
"type": "IMAGE",
|
||||
"shape": 7,
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "image_7",
|
||||
"type": "IMAGE",
|
||||
"shape": 7,
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "image_8",
|
||||
"type": "IMAGE",
|
||||
"shape": 7,
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "image_9",
|
||||
"type": "IMAGE",
|
||||
"shape": 7,
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "audio_3",
|
||||
"type": "AUDIO",
|
||||
"shape": 7,
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"name": "video_3",
|
||||
"type": "VIDEO",
|
||||
"shape": 7,
|
||||
"link": null
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"localized_name": "references",
|
||||
"name": "references",
|
||||
"type": "VIDEO_REFERENCES",
|
||||
"links": [
|
||||
2
|
||||
]
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "VLLMOmniVideoReferences"
|
||||
},
|
||||
"widgets_values": [],
|
||||
"title": "References \u2014 up to 12 total"
|
||||
},
|
||||
{
|
||||
"id": 5,
|
||||
"type": "VLLMOmniDiffusionSampling",
|
||||
"pos": [
|
||||
420,
|
||||
530
|
||||
],
|
||||
"size": [
|
||||
310,
|
||||
275
|
||||
],
|
||||
"flags": {},
|
||||
"order": 4,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"localized_name": "n",
|
||||
"name": "n",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "n"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "num_inference_steps",
|
||||
"name": "num_inference_steps",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "num_inference_steps"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "guidance_scale",
|
||||
"name": "guidance_scale",
|
||||
"type": "FLOAT",
|
||||
"widget": {
|
||||
"name": "guidance_scale"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "true_cfg_scale",
|
||||
"name": "true_cfg_scale",
|
||||
"type": "FLOAT",
|
||||
"widget": {
|
||||
"name": "true_cfg_scale"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "vae_use_slicing",
|
||||
"name": "vae_use_slicing",
|
||||
"type": "BOOLEAN",
|
||||
"widget": {
|
||||
"name": "vae_use_slicing"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "vae_use_tiling",
|
||||
"name": "vae_use_tiling",
|
||||
"type": "BOOLEAN",
|
||||
"widget": {
|
||||
"name": "vae_use_tiling"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "seed",
|
||||
"name": "seed",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "seed"
|
||||
},
|
||||
"link": null
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"localized_name": "diffusion sampling params",
|
||||
"name": "diffusion sampling params",
|
||||
"type": "SAMPLING_PARAMS",
|
||||
"links": [
|
||||
3
|
||||
]
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "VLLMOmniDiffusionSampling"
|
||||
},
|
||||
"widgets_values": [
|
||||
1,
|
||||
50,
|
||||
1.0,
|
||||
1.0,
|
||||
false,
|
||||
false,
|
||||
42,
|
||||
"fixed"
|
||||
],
|
||||
"title": "Sampling \u2014 base 50 / Turbo 5"
|
||||
},
|
||||
{
|
||||
"id": 6,
|
||||
"type": "VLLMOmniMiniMaxH3Params",
|
||||
"pos": [
|
||||
790,
|
||||
750
|
||||
],
|
||||
"size": [
|
||||
310,
|
||||
105
|
||||
],
|
||||
"flags": {},
|
||||
"order": 5,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"localized_name": "audio_flow_shift",
|
||||
"name": "audio_flow_shift",
|
||||
"type": "FLOAT",
|
||||
"widget": {
|
||||
"name": "audio_flow_shift"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "flow_shift",
|
||||
"name": "flow_shift",
|
||||
"type": "FLOAT",
|
||||
"widget": {
|
||||
"name": "flow_shift"
|
||||
},
|
||||
"link": null
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"localized_name": "MiniMax-H3 video params",
|
||||
"name": "MiniMax-H3 video params",
|
||||
"type": "VIDEO_PARAMS",
|
||||
"links": [
|
||||
4
|
||||
]
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "VLLMOmniMiniMaxH3Params"
|
||||
},
|
||||
"widgets_values": [
|
||||
3.0,
|
||||
12.0
|
||||
],
|
||||
"title": "H3 video/audio flow shifts"
|
||||
},
|
||||
{
|
||||
"type": "VLLMOmniRemoteLoRA",
|
||||
"size": [
|
||||
310,
|
||||
180
|
||||
],
|
||||
"flags": {},
|
||||
"inputs": [],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "lora",
|
||||
"type": "REMOTE_LORA",
|
||||
"links": null
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "VLLMOmniRemoteLoRA"
|
||||
},
|
||||
"widgets_values": [
|
||||
"",
|
||||
"h3-ref2va-turbo",
|
||||
1.0,
|
||||
0
|
||||
],
|
||||
"id": 7,
|
||||
"pos": [
|
||||
420,
|
||||
940
|
||||
],
|
||||
"order": 6,
|
||||
"mode": 0,
|
||||
"title": "Optional Turbo LoRA \u2014 connect to enable"
|
||||
},
|
||||
{
|
||||
"id": 8,
|
||||
"type": "VLLMOmniGenerateVideo",
|
||||
"pos": [
|
||||
800,
|
||||
100
|
||||
],
|
||||
"size": [
|
||||
470,
|
||||
560
|
||||
],
|
||||
"flags": {},
|
||||
"order": 7,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"localized_name": "frame",
|
||||
"name": "frame",
|
||||
"shape": 7,
|
||||
"type": "IMAGE",
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "references",
|
||||
"name": "references",
|
||||
"shape": 7,
|
||||
"type": "VIDEO_REFERENCES",
|
||||
"link": 2
|
||||
},
|
||||
{
|
||||
"localized_name": "sampling_params",
|
||||
"name": "sampling_params",
|
||||
"shape": 7,
|
||||
"type": "SAMPLING_PARAMS",
|
||||
"link": 3
|
||||
},
|
||||
{
|
||||
"localized_name": "lora",
|
||||
"name": "lora",
|
||||
"shape": 7,
|
||||
"type": "REMOTE_LORA",
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "model_params",
|
||||
"name": "model_params",
|
||||
"shape": 7,
|
||||
"type": "VIDEO_PARAMS",
|
||||
"link": 4
|
||||
},
|
||||
{
|
||||
"localized_name": "url",
|
||||
"name": "url",
|
||||
"type": "STRING",
|
||||
"widget": {
|
||||
"name": "url"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "model",
|
||||
"name": "model",
|
||||
"type": "STRING",
|
||||
"widget": {
|
||||
"name": "model"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "prompt",
|
||||
"name": "prompt",
|
||||
"type": "STRING",
|
||||
"widget": {
|
||||
"name": "prompt"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "negative_prompt",
|
||||
"name": "negative_prompt",
|
||||
"type": "STRING",
|
||||
"widget": {
|
||||
"name": "negative_prompt"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "width",
|
||||
"name": "width",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "width"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "height",
|
||||
"name": "height",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "height"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "fps",
|
||||
"name": "fps",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "fps"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "duration",
|
||||
"name": "duration",
|
||||
"type": "FLOAT",
|
||||
"widget": {
|
||||
"name": "duration"
|
||||
},
|
||||
"link": null
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"localized_name": "video",
|
||||
"name": "video",
|
||||
"type": "VIDEO",
|
||||
"links": [
|
||||
5
|
||||
]
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "VLLMOmniGenerateVideo"
|
||||
},
|
||||
"widgets_values": [
|
||||
"http://localhost:8000/v1",
|
||||
"MiniMaxAI/MiniMax-H3",
|
||||
"Use the subject and visual style from <Picture 1>. Create a cinematic shot with a gentle camera move and natural ambient sound.",
|
||||
"",
|
||||
1344,
|
||||
768,
|
||||
24,
|
||||
5.167
|
||||
],
|
||||
"title": "MiniMax-H3 Reference to Video"
|
||||
},
|
||||
{
|
||||
"id": 9,
|
||||
"type": "SaveVideo",
|
||||
"pos": [
|
||||
1360,
|
||||
100
|
||||
],
|
||||
"size": [
|
||||
410,
|
||||
420
|
||||
],
|
||||
"flags": {},
|
||||
"order": 8,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"localized_name": "video",
|
||||
"name": "video",
|
||||
"type": "VIDEO",
|
||||
"link": 5
|
||||
},
|
||||
{
|
||||
"localized_name": "filename_prefix",
|
||||
"name": "filename_prefix",
|
||||
"type": "STRING",
|
||||
"widget": {
|
||||
"name": "filename_prefix"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "format",
|
||||
"name": "format",
|
||||
"type": "COMBO",
|
||||
"widget": {
|
||||
"name": "format"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "codec",
|
||||
"name": "codec",
|
||||
"type": "COMFY_DYNAMICCOMBO_V3",
|
||||
"widget": {
|
||||
"name": "codec"
|
||||
},
|
||||
"link": null
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"localized_name": "video",
|
||||
"name": "video",
|
||||
"type": "VIDEO",
|
||||
"links": null
|
||||
}
|
||||
],
|
||||
"properties": {},
|
||||
"widgets_values": [
|
||||
"video/MiniMax-H3-Ref2VA",
|
||||
"mp4",
|
||||
"h264"
|
||||
],
|
||||
"title": "Save generated video with audio"
|
||||
},
|
||||
{
|
||||
"id": 10,
|
||||
"type": "MarkdownNote",
|
||||
"pos": [
|
||||
1360,
|
||||
590
|
||||
],
|
||||
"size": [
|
||||
440,
|
||||
595
|
||||
],
|
||||
"flags": {},
|
||||
"order": 9,
|
||||
"mode": 0,
|
||||
"inputs": [],
|
||||
"outputs": [],
|
||||
"title": "Reference inputs and output settings",
|
||||
"properties": {},
|
||||
"widgets_values": [
|
||||
"## Reference to Video\nSelect an image, set the server URL and served model name, then run.\n\nConnect the optional video/audio loaders to Video References, or duplicate loaders for more inputs: up to 9 images, 3 videos, 3 audio clips, 12 total. At least one image or video is required.\n\nSlots are numbered separately per media type. Empty slots are skipped. Update prompt tags such as <Picture 1>, <Video 1>, <Audio 1> to match the connected references. Videos keep their own soundtrack in the reference upload.\n\n**Output:** 24 FPS; 124 frames (~5.17 s). Valid frame counts are 107, 124, 141, 158, 175, 192, 209, 226, 243, 260, 277, 294, 311, 328, 345.\n\n**Width x height:** 1344x768 (16:9), 1024x768 (4:3), 768x768 (1:1), 768x1024 (3:4), 768x1344 (9:16).\n\nEach reference video/audio clip must be 2\u201315 s. Total reference-video duration <=15 s; total standalone-audio duration <=15 s.\n\nRequires the expanded Video References node and the audio-preserving video output fix. See the extension README for setup and validation status."
|
||||
],
|
||||
"color": "#432",
|
||||
"bgcolor": "#000"
|
||||
},
|
||||
{
|
||||
"id": 11,
|
||||
"type": "MarkdownNote",
|
||||
"pos": [
|
||||
790,
|
||||
940
|
||||
],
|
||||
"size": [
|
||||
480,
|
||||
455
|
||||
],
|
||||
"flags": {},
|
||||
"order": 10,
|
||||
"mode": 0,
|
||||
"inputs": [],
|
||||
"outputs": [],
|
||||
"title": "Optional Turbo \u2014 same workflow",
|
||||
"properties": {},
|
||||
"widgets_values": [
|
||||
"## Turbo Ref2VA\nBase mode: leave LoRA disconnected; use 50 steps, video flow shift 12, audio flow shift 3.\n\nFor the 4-step v0.1 Ref2VA Turbo adapter:\n1. Start a Ref2VA-only server with the Diffusers adapter loaded (see README).\n2. Set LoRA local_path to its path on the server, then connect LoRA to Generate Video.\n3. Set Sampling num_inference_steps to **5**; keep shifts at **12 / 3**, scale **1**.\n\nvLLM-Omni counts sigma points: 5 points give 4 denoiser evaluations. Use minimax_h3_ref2v_turbo_4step_v0.1_bf16.safetensors, not the _comfyui_ export.\n\nDisconnect LoRA and restore 50 steps to return to base mode. FastH3 is a separate T2VA deployment, not this option.\n\nAdapted from the official ComfyUI MiniMax-H3 R2V workflow: https://github.com/Comfy-Org/workflow_templates/blob/main/templates/video_minimax_h3_r2v.json"
|
||||
],
|
||||
"color": "#432",
|
||||
"bgcolor": "#000"
|
||||
}
|
||||
],
|
||||
"links": [
|
||||
[
|
||||
1,
|
||||
1,
|
||||
0,
|
||||
4,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
2,
|
||||
4,
|
||||
0,
|
||||
8,
|
||||
1,
|
||||
"VIDEO_REFERENCES"
|
||||
],
|
||||
[
|
||||
3,
|
||||
5,
|
||||
0,
|
||||
8,
|
||||
2,
|
||||
"SAMPLING_PARAMS"
|
||||
],
|
||||
[
|
||||
4,
|
||||
6,
|
||||
0,
|
||||
8,
|
||||
4,
|
||||
"VIDEO_PARAMS"
|
||||
],
|
||||
[
|
||||
5,
|
||||
8,
|
||||
0,
|
||||
9,
|
||||
0,
|
||||
"VIDEO"
|
||||
]
|
||||
],
|
||||
"groups": [],
|
||||
"config": {},
|
||||
"extra": {
|
||||
"workflowRendererVersion": "LG",
|
||||
"ds": {
|
||||
"scale": 0.85,
|
||||
"offset": [
|
||||
100,
|
||||
100
|
||||
]
|
||||
}
|
||||
},
|
||||
"version": 0.4
|
||||
}
|
||||
@@ -0,0 +1,333 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Benchmark the Qwen-Image-Edit select01 modulation hot path.
|
||||
|
||||
This script intentionally starts with the current PyTorch implementation. It is
|
||||
used as the before/after baseline for SGLang PR #20395/#21318 style fusion:
|
||||
|
||||
* layernorm + scale/shift + binary select01 gate
|
||||
* residual + gated update + layernorm + scale/shift + binary select01 gate
|
||||
|
||||
Run on a CUDA machine for actionable numbers, for example:
|
||||
|
||||
python benchmarks/kernels/bench_qwen_select01_modulation.py \
|
||||
--device cuda --dtype bf16 --seq-lens 4096 8192 16384 \
|
||||
--output-json /tmp/qwen_select01_baseline.json
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import statistics
|
||||
import sys
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
for candidate in Path(__file__).resolve().parents:
|
||||
if (candidate / "vllm_omni").is_dir():
|
||||
if str(candidate) not in sys.path:
|
||||
sys.path.insert(0, str(candidate))
|
||||
break
|
||||
|
||||
from vllm_omni.diffusion.layers.qwen_select01_modulation import ( # noqa: E402
|
||||
fused_layernorm_select01,
|
||||
fused_residual_layernorm_select01,
|
||||
)
|
||||
|
||||
|
||||
def _dtype(name: str) -> torch.dtype:
|
||||
return {
|
||||
"fp32": torch.float32,
|
||||
"float32": torch.float32,
|
||||
"fp16": torch.float16,
|
||||
"float16": torch.float16,
|
||||
"bf16": torch.bfloat16,
|
||||
"bfloat16": torch.bfloat16,
|
||||
}[name]
|
||||
|
||||
|
||||
def _select01_modulation(
|
||||
mod_params: torch.Tensor,
|
||||
index: torch.Tensor | None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
shift, scale, gate = mod_params.chunk(3, dim=-1)
|
||||
|
||||
if index is None:
|
||||
return scale.unsqueeze(1), shift.unsqueeze(1), gate.unsqueeze(1)
|
||||
|
||||
actual_batch = shift.size(0) // 2
|
||||
shift_0, shift_1 = shift[:actual_batch], shift[actual_batch:]
|
||||
scale_0, scale_1 = scale[:actual_batch], scale[actual_batch:]
|
||||
gate_0, gate_1 = gate[:actual_batch], gate[actual_batch:]
|
||||
|
||||
idx = index.unsqueeze(-1) == 0
|
||||
shift = torch.where(idx, shift_0.unsqueeze(1), shift_1.unsqueeze(1))
|
||||
scale = torch.where(idx, scale_0.unsqueeze(1), scale_1.unsqueeze(1))
|
||||
gate = torch.where(idx, gate_0.unsqueeze(1), gate_1.unsqueeze(1))
|
||||
return scale, shift, gate
|
||||
|
||||
|
||||
def _native_layernorm_select01(
|
||||
x: torch.Tensor,
|
||||
mod_params: torch.Tensor,
|
||||
index: torch.Tensor,
|
||||
eps: float,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
scale, shift, gate = _select01_modulation(mod_params, index)
|
||||
out = F.layer_norm(x.float(), (x.shape[-1],), eps=eps).to(x.dtype)
|
||||
return out * (1 + scale) + shift, gate
|
||||
|
||||
|
||||
def _native_residual_layernorm_select01(
|
||||
x: torch.Tensor,
|
||||
residual: torch.Tensor,
|
||||
residual_gate: torch.Tensor,
|
||||
mod_params: torch.Tensor,
|
||||
index: torch.Tensor,
|
||||
eps: float,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
scale, shift, gate = _select01_modulation(mod_params, index)
|
||||
residual_out = residual + residual_gate * x
|
||||
out = F.layer_norm(residual_out.float(), (residual_out.shape[-1],), eps=eps).to(residual_out.dtype)
|
||||
return out * (1 + scale) + shift, residual_out, gate
|
||||
|
||||
|
||||
def _time_cuda(fn: Callable[[], object], warmup: int, iters: int) -> list[float]:
|
||||
for _ in range(warmup):
|
||||
fn()
|
||||
torch.accelerator.synchronize()
|
||||
|
||||
times: list[float] = []
|
||||
for _ in range(iters):
|
||||
start = torch.cuda.Event(enable_timing=True)
|
||||
end = torch.cuda.Event(enable_timing=True)
|
||||
start.record()
|
||||
fn()
|
||||
end.record()
|
||||
torch.accelerator.synchronize()
|
||||
times.append(float(start.elapsed_time(end)))
|
||||
return times
|
||||
|
||||
|
||||
def _reset_peak_memory_stats(device: torch.device) -> None:
|
||||
reset_peak_memory_stats = getattr(torch.accelerator, "reset_peak_memory_stats", None)
|
||||
if reset_peak_memory_stats is None:
|
||||
reset_peak_memory_stats = getattr(getattr(torch, device.type), "reset_peak_memory_stats")
|
||||
reset_peak_memory_stats(device)
|
||||
|
||||
|
||||
def _max_memory_allocated(device: torch.device) -> int:
|
||||
max_memory_allocated = getattr(torch.accelerator, "max_memory_allocated", None)
|
||||
if max_memory_allocated is None:
|
||||
max_memory_allocated = getattr(getattr(torch, device.type), "max_memory_allocated")
|
||||
return int(max_memory_allocated(device))
|
||||
|
||||
|
||||
def _time_cpu(fn: Callable[[], object], warmup: int, iters: int) -> list[float]:
|
||||
for _ in range(warmup):
|
||||
fn()
|
||||
|
||||
times: list[float] = []
|
||||
for _ in range(iters):
|
||||
t0 = time.perf_counter()
|
||||
fn()
|
||||
times.append((time.perf_counter() - t0) * 1000.0)
|
||||
return times
|
||||
|
||||
|
||||
def _summarize(times: list[float]) -> dict[str, float]:
|
||||
ordered = sorted(times)
|
||||
return {
|
||||
"median_ms": statistics.median(ordered),
|
||||
"min_ms": ordered[0],
|
||||
"p20_ms": ordered[max(0, int(len(ordered) * 0.20) - 1)],
|
||||
"p80_ms": ordered[min(len(ordered) - 1, int(len(ordered) * 0.80))],
|
||||
}
|
||||
|
||||
|
||||
def _make_inputs(
|
||||
batch: int,
|
||||
seq_len: int,
|
||||
hidden_size: int,
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
) -> dict[str, torch.Tensor]:
|
||||
torch.manual_seed(0)
|
||||
return {
|
||||
"x": torch.randn(batch, seq_len, hidden_size, device=device, dtype=dtype),
|
||||
"residual": torch.randn(batch, seq_len, hidden_size, device=device, dtype=dtype),
|
||||
"residual_gate": torch.randn(batch, 1, hidden_size, device=device, dtype=dtype),
|
||||
"mod_params": torch.randn(batch * 2, hidden_size * 3, device=device, dtype=dtype),
|
||||
"index": torch.randint(0, 2, (batch, seq_len), device=device, dtype=torch.int64),
|
||||
}
|
||||
|
||||
|
||||
def _bench_one(
|
||||
name: str,
|
||||
fn: Callable[[], object],
|
||||
device: torch.device,
|
||||
warmup: int,
|
||||
iters: int,
|
||||
) -> dict[str, object]:
|
||||
if device.type == "cuda":
|
||||
_reset_peak_memory_stats(device)
|
||||
times = _time_cuda(fn, warmup, iters)
|
||||
peak_memory_mb = _max_memory_allocated(device) / 1024**2
|
||||
else:
|
||||
times = _time_cpu(fn, warmup, iters)
|
||||
peak_memory_mb = None
|
||||
|
||||
return {
|
||||
"name": name,
|
||||
**_summarize(times),
|
||||
"peak_memory_mb": peak_memory_mb,
|
||||
}
|
||||
|
||||
|
||||
def _run(args: argparse.Namespace) -> dict[str, object]:
|
||||
device = torch.device(args.device)
|
||||
if device.type == "cuda" and not torch.cuda.is_available():
|
||||
raise SystemExit("CUDA is not available on this machine.")
|
||||
|
||||
dtype = _dtype(args.dtype)
|
||||
rows: list[dict[str, object]] = []
|
||||
|
||||
for seq_len in args.seq_lens:
|
||||
tensors = _make_inputs(args.batch_size, seq_len, args.hidden_size, dtype, device)
|
||||
impls = ["native"] if args.impl == "native" else ["fused"] if args.impl == "fused" else ["native", "fused"]
|
||||
|
||||
def first_norm() -> object:
|
||||
return _native_layernorm_select01(
|
||||
tensors["x"],
|
||||
tensors["mod_params"],
|
||||
tensors["index"],
|
||||
args.eps,
|
||||
)
|
||||
|
||||
def residual_norm() -> object:
|
||||
return _native_residual_layernorm_select01(
|
||||
tensors["x"],
|
||||
tensors["residual"],
|
||||
tensors["residual_gate"],
|
||||
tensors["mod_params"],
|
||||
tensors["index"],
|
||||
args.eps,
|
||||
)
|
||||
|
||||
callables: list[tuple[str, Callable[[], object]]] = []
|
||||
if "native" in impls:
|
||||
callables.extend(
|
||||
[
|
||||
("native_layernorm_select01", first_norm),
|
||||
("native_residual_layernorm_select01", residual_norm),
|
||||
]
|
||||
)
|
||||
if "fused" in impls:
|
||||
|
||||
def fused_first_norm() -> object:
|
||||
return fused_layernorm_select01(
|
||||
tensors["x"],
|
||||
tensors["mod_params"],
|
||||
tensors["index"],
|
||||
args.eps,
|
||||
)
|
||||
|
||||
def fused_residual_norm() -> object:
|
||||
return fused_residual_layernorm_select01(
|
||||
tensors["x"],
|
||||
tensors["residual"],
|
||||
tensors["residual_gate"].expand_as(tensors["residual"]),
|
||||
tensors["mod_params"],
|
||||
tensors["index"],
|
||||
args.eps,
|
||||
)
|
||||
|
||||
callables.extend(
|
||||
[
|
||||
("fused_layernorm_select01", fused_first_norm),
|
||||
("fused_residual_layernorm_select01", fused_residual_norm),
|
||||
]
|
||||
)
|
||||
|
||||
if args.check and "fused" in impls:
|
||||
ref_norm, ref_gate = first_norm()
|
||||
actual_norm, actual_gate = fused_layernorm_select01(
|
||||
tensors["x"],
|
||||
tensors["mod_params"],
|
||||
tensors["index"],
|
||||
args.eps,
|
||||
)
|
||||
torch.testing.assert_close(actual_norm, ref_norm, atol=args.check_atol, rtol=args.check_rtol)
|
||||
torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0)
|
||||
ref_norm, ref_residual, ref_gate = residual_norm()
|
||||
actual_norm, actual_residual, actual_gate = fused_residual_layernorm_select01(
|
||||
tensors["x"],
|
||||
tensors["residual"],
|
||||
tensors["residual_gate"].expand_as(tensors["residual"]),
|
||||
tensors["mod_params"],
|
||||
tensors["index"],
|
||||
args.eps,
|
||||
)
|
||||
torch.testing.assert_close(actual_norm, ref_norm, atol=args.check_atol, rtol=args.check_rtol)
|
||||
torch.testing.assert_close(actual_residual, ref_residual, atol=args.check_atol, rtol=args.check_rtol)
|
||||
torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0)
|
||||
|
||||
for name, fn in callables:
|
||||
row = _bench_one(name, fn, device, args.warmup, args.iters)
|
||||
row.update(
|
||||
{
|
||||
"batch_size": args.batch_size,
|
||||
"seq_len": seq_len,
|
||||
"hidden_size": args.hidden_size,
|
||||
"dtype": args.dtype,
|
||||
"device": str(device),
|
||||
}
|
||||
)
|
||||
rows.append(row)
|
||||
print(
|
||||
f"{name:<38} B={args.batch_size:<2} L={seq_len:<6} C={args.hidden_size:<5} "
|
||||
f"median={row['median_ms']:.4f} ms p20={row['p20_ms']:.4f} p80={row['p80_ms']:.4f}"
|
||||
)
|
||||
|
||||
result = {
|
||||
"benchmark": "qwen_select01_modulation",
|
||||
"torch": torch.__version__,
|
||||
"cuda_available": torch.cuda.is_available(),
|
||||
"gpu": torch.cuda.get_device_name(device) if device.type == "cuda" else None,
|
||||
"rows": rows,
|
||||
}
|
||||
|
||||
if args.output_json:
|
||||
output_path = Path(args.output_json)
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
output_path.write_text(json.dumps(result, indent=2), encoding="utf-8")
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def _parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
default_device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
parser.add_argument("--device", default=default_device)
|
||||
parser.add_argument("--dtype", default="bf16", choices=["fp32", "float32", "fp16", "float16", "bf16", "bfloat16"])
|
||||
parser.add_argument("--batch-size", type=int, default=1)
|
||||
parser.add_argument("--seq-lens", type=int, nargs="+", default=[4096, 8192, 16384])
|
||||
parser.add_argument("--hidden-size", type=int, default=3072)
|
||||
parser.add_argument("--eps", type=float, default=1e-6)
|
||||
parser.add_argument("--warmup", type=int, default=20)
|
||||
parser.add_argument("--iters", type=int, default=100)
|
||||
parser.add_argument("--impl", choices=["native", "fused", "all"], default="native")
|
||||
parser.add_argument("--check", action="store_true")
|
||||
parser.add_argument("--check-atol", type=float, default=5e-2)
|
||||
parser.add_argument("--check-rtol", type=float, default=5e-2)
|
||||
parser.add_argument("--output-json")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
_run(_parse_args())
|
||||
@@ -44,6 +44,7 @@ RUN python3 -m pip install -U pip && \
|
||||
"onnx>=1.16.0" \
|
||||
"librosa>=0.10.0" \
|
||||
"datasets>=3.0.0" \
|
||||
"torchmetrics==1.9.0" \
|
||||
--index-url http://cache-service.nginx-pypi-cache.svc.cluster.local/pypi/simple \
|
||||
--trusted-host cache-service.nginx-pypi-cache.svc.cluster.local \
|
||||
--default-timeout=60
|
||||
|
||||
+6
-1
@@ -5,6 +5,9 @@
|
||||
Main entry points for vLLM-Omni inference and serving.
|
||||
|
||||
- [vllm_omni.entrypoints.async_omni.AsyncOmni][]
|
||||
- [vllm_omni.entrypoints.async_omni_base.AsyncOmniBase][]
|
||||
- [vllm_omni.entrypoints.duplex_omni.DuplexOmni][]
|
||||
- [vllm_omni.entrypoints.duplex_omni.DuplexSessionHandle][]
|
||||
- [vllm_omni.engine.cfg_companion_tracker.CfgCompanionTracker][]
|
||||
- [vllm_omni.entrypoints.cli.benchmark.base.OmniBenchmarkSubcommandBase][]
|
||||
- [vllm_omni.entrypoints.cli.benchmark.main.OmniBenchmarkSubcommand][]
|
||||
@@ -17,12 +20,13 @@ Main entry points for vLLM-Omni inference and serving.
|
||||
|
||||
## Clients
|
||||
|
||||
Client libraries that talk to a running vLLM-Omni server (the full-duplex Realtime client and its per-model session presets).
|
||||
Client libraries for full-duplex sessions: the Realtime WebSocket client, the in-process client over `DuplexOmni`, and the per-model session presets.
|
||||
|
||||
- [vllm_omni.clients.duplex.AudioDelta][]
|
||||
- [vllm_omni.clients.duplex.AudioFormat][]
|
||||
- [vllm_omni.clients.duplex.ConnectionResumed][]
|
||||
- [vllm_omni.clients.duplex.DuplexClient][]
|
||||
- [vllm_omni.clients.duplex.DuplexClientBase][]
|
||||
- [vllm_omni.clients.duplex.DuplexClientError][]
|
||||
- [vllm_omni.clients.duplex.DuplexConnectionError][]
|
||||
- [vllm_omni.clients.duplex.DuplexEvent][]
|
||||
@@ -44,6 +48,7 @@ Client libraries that talk to a running vLLM-Omni server (the full-duplex Realti
|
||||
- [vllm_omni.clients.duplex.TextDelta][]
|
||||
- [vllm_omni.clients.duplex.TranscriptDelta][]
|
||||
- [vllm_omni.clients.duplex.WebSocketTransport][]
|
||||
- [vllm_omni.clients.inline_duplex.InlineDuplexClient][]
|
||||
|
||||
## Inputs
|
||||
|
||||
|
||||
@@ -82,6 +82,7 @@ depends on the installed kernels and model path.
|
||||
| --- | --- | --- | --- | --- |
|
||||
| `SPEAKER_SAMPLES_DIR` | Filesystem path; default `~/.cache/vllm-omni/speakers` | Speech server; read when speaker storage initializes | Environment-only setting. The directory is created; filesystem errors propagate. | Stable |
|
||||
| `SPEAKER_MAX_UPLOADED` | Integer; default `1000` | Speech server; read when speaker storage initializes | Environment-only setting. A non-integer logs a warning and uses `1000`; range is not otherwise validated. | Stable |
|
||||
| `VLLM_OMNI_ABORT_TIMEOUT` | Float seconds; default `2` | Engine abort wait for Videos DELETE and `generate()` cancel/error cleanup; read when `async_omni` imports | Environment-only setting. A non-float raises `ValueError` during import. | Experimental |
|
||||
| `VLLM_OMNI_ASYNC_OUTPUT_TIMEOUT` | Float seconds; default `600` | Diffusion engine async-output wait in `step_streaming`; resolved per call on the request path, not at import | Environment-only setting. A non-float or `<=0` value warns once and uses the default. | Experimental |
|
||||
| `VLLM_OMNI_EVENT_DRIVEN_ORCH` | `1`, `true`, `yes` or `on` enables; default `0` (off), except Qwen3-TTS defaults to on | Orchestration loop and the serving-side final-output drain; read once when the `Orchestrator` is constructed | Environment-only setting. An explicit value always wins. Values are stripped and case-normalized; any unrecognized value leaves the legacy poll loop selected. | Experimental |
|
||||
| `VLLM_OMNI_INPUT_WAIT_TIMEOUT_S` | Float seconds; default `600`; `<=0` disables | Full-payload input coordinator, not async-chunk transfer; read when the scheduler module imports in each worker | Environment-only setting. A non-float logs a warning and uses `600`. | Stable operational control |
|
||||
|
||||
@@ -17,6 +17,7 @@ Common `PipelineConfig` fields include:
|
||||
| ------- | ------------- |
|
||||
| `model_type` | Pipeline identifier used during model and config resolution. |
|
||||
| `default_deploy_config_name` | Bundled deploy YAML loaded when the user does not pass `deploy_config`. |
|
||||
| `duplex_plugin` | Dotted path of the model's `DuplexModelPlugin`. Set only for full-duplex models; it makes `vllm-omni serve` run the model through `DuplexOmni` (a server whose every surface runs on a duplex session) and the engine host a `DuplexOrchestrator` with the plugin loaded. |
|
||||
| `model_arch` | Default Hugging Face architecture for the pipeline. |
|
||||
| `hf_architectures` | Architecture names used to identify checkpoints whose `model_type` is shared. |
|
||||
| `hf_config_predicate` | Optional predicate used to select between pipelines with otherwise identical HF metadata. |
|
||||
@@ -52,7 +53,7 @@ The new deploy schema lives under `vllm_omni/deploy/` and is paired with a froze
|
||||
| ------- | ------ | ---------- | --------- | ------------- |
|
||||
| `base_config` | str (path) | optional | — | Overlay parent (relative or absolute). `stages:` / `platforms:` deep-merged by stage_id; other scalars overlay-wins. Intended for user-authored overlays; prod yamls stay flat. |
|
||||
| `async_chunk` | bool | optional | `true` | Enable chunked streaming between stages. Pin to `false` if the pipeline runs end-to-end. |
|
||||
| `session_mode` | str | optional | `"turn"` | Session behavior. MiniCPM-o 4.5 deploy YAMLs set `"duplex"` so `/v1/realtime` and chat share the same profile. |
|
||||
| `session_mode` | str | optional | `"turn"` | Session behavior. Duplex models (pipelines with a `duplex_plugin`, e.g. MiniCPM-o 4.5) set `"duplex"`; `vllm-omni serve` then serves it over `/v1/realtime?duplex=1` and `/v1/chat/completions`, both on duplex sessions. |
|
||||
| `active_stream_window` | int | optional | `0` | Number of active downstream stream slots; `0` preserves all-stream cycling. |
|
||||
| `duplex_session` | dict | optional | runtime defaults | Full-duplex session lifecycle, buffering, replay, and capacity limits. |
|
||||
| `connectors` | dict | optional | `null` | Named connector specs (`{name, extra}`). Referenced by each stage's `input_connectors` / `output_connectors`. See [Connector schema](#connector-schema). |
|
||||
|
||||
@@ -131,13 +131,13 @@ for the validated batching, attention, and quantization combinations.
|
||||
|
||||
### MiniMax-H3: conditioning, denoising, and decode
|
||||
|
||||
MiniMax-H3 illustrates logical component and phase boundaries between text
|
||||
encoding, task-specific denoising, and media decoding:
|
||||
`Text Encoder → DiT → VAE Decoder`. The current registered topology co-locates
|
||||
these components in one diffusion stage. These boundaries are useful for
|
||||
profiling and offload; independent placement would require a future
|
||||
multi-stage `PipelineConfig`. The stage mapping follows the implementation and
|
||||
serving recipe; the [release overview](https://minimaxi.com/blog/minimax-h3)
|
||||
MiniMax-H3 combines text and reference-media encoding, task-specific denoising,
|
||||
and media decoding. Standard serving co-locates these components in one diffusion
|
||||
stage. The opt-in
|
||||
[disaggregated topology](https://github.com/vllm-project/vllm-omni/blob/main/recipes/MiniMaxAI/MiniMax-H3-Disaggregated.md)
|
||||
runs Qwen and the video/audio VAE encoders in Stage 0, then sends their conditioning
|
||||
to a Stage 1 containing the DiTs and VAE decoders. The stage mapping follows the
|
||||
implementation and serving recipe; the [release overview](https://minimaxi.com/blog/minimax-h3)
|
||||
notes that a detailed technical report is forthcoming.
|
||||
|
||||
#### Model architecture by component
|
||||
@@ -149,16 +149,15 @@ notes that a detailed technical report is forthcoming.
|
||||
| **VAE Decoder** | Video VAE and audio VAE decode the generated latents and assemble the synchronized output | Decoded video frames and stereo waveform, with FPS and audio sample rate metadata |
|
||||
|
||||
Both DiT partitions remain behind one task-selection point; `task` chooses
|
||||
which partition runs. In the current implementation, the three logical
|
||||
components execute inside one diffusion stage. This is different from
|
||||
Qwen3-Omni's three independently scheduled stages, even though both models
|
||||
produce synchronized audio and video-like outputs.
|
||||
which partition runs.
|
||||
|
||||
#### Representative component execution
|
||||
|
||||
The following table describes the standard single-stage deployment.
|
||||
|
||||
| Logical component | Batching | Attention and execution | Parallelism | Quantization |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| **Text Encoder** | The current co-located pipeline follows the diffusion request scheduler; independent prompt/reference batching would require a separate multi-stage topology | Qwen3-VL attention and multimodal preprocessing | `--text-encoder-tp-size N` shards the encoder across the first `N` DiT ranks; independent placement is not part of the current topology | BF16/FP32 baseline; the H3 DiT FP8 path does not quantize the text encoder |
|
||||
| **Text Encoder** | Follows the diffusion request scheduler | Qwen3-VL attention and multimodal preprocessing | `--text-encoder-tp-size N` shards the encoder across the first `N` DiT ranks | BF16/FP32 baseline; the H3 DiT FP8 path does not quantize the text encoder |
|
||||
| **DiT** | The current H3 implementation executes one generation request per diffusion batch; use concurrency for service-level throughput | cuDNN attention for the validated two-GPU consumer profile; TRTLLM attention or FlashAttention-4 on supported Blackwell profiles | TP2 plus text-encoder TP and Ulysses/VAE parallel groups on multi-GPU profiles; DLO/CPU offload trade memory for transfer time | Online FP8 applies to eligible DiT linears and is incompatible with layerwise offload |
|
||||
| **VAE Decoder** | Decode follows each generation request; tile/patch work can be distributed even when denoising is not batched | VAE decode kernels rather than DiT attention | VAE patch parallelism and native tiled decode within the diffusion stage | BF16/FP32 baseline; the documented H3 FP8 path leaves both VAEs unchanged |
|
||||
|
||||
|
||||
@@ -76,6 +76,34 @@ preserve common `self` or `cross` policies. A model also decides whether its
|
||||
path is compatible with automatic TRTLLM selection, because only the model
|
||||
knows whether masking and packing satisfy the kernel contract.
|
||||
|
||||
### Model-owned implementation specializations
|
||||
|
||||
Use `Attention(..., impl_overrides={backend_name: implementation_class})`
|
||||
when a model needs different computation within an existing backend contract.
|
||||
The model defines the implementation in its own directory and passes it when
|
||||
constructing the layer:
|
||||
|
||||
```python
|
||||
# MyModelVSAImpl is a model-owned subclass of FastVideoVSAImpl.
|
||||
self.attention = Attention(
|
||||
num_heads=num_heads,
|
||||
head_size=head_size,
|
||||
softmax_scale=head_size**-0.5,
|
||||
causal=False,
|
||||
role="self",
|
||||
impl_overrides={"FASTVIDEO_VSA": MyModelVSAImpl},
|
||||
)
|
||||
```
|
||||
|
||||
The override is applied after role/platform selection, keyed by the selected
|
||||
backend's `get_name()`. It must subclass the **selected implementation**;
|
||||
incompatible platform implementations raise `TypeError`. Unselected overrides
|
||||
have no effect. Backend capabilities, constructor options, and shared
|
||||
parallel dispatch remain unchanged and must be honored by the specialization.
|
||||
|
||||
Use a new backend for a new provider/capability contract. `custom_attention`
|
||||
instead owns communication and requires `skip_sequence_parallel=True`.
|
||||
|
||||
## Adding or changing a backend
|
||||
|
||||
An implementation change is complete only when it:
|
||||
@@ -91,8 +119,9 @@ An implementation change is complete only when it:
|
||||
|
||||
## FastVideo VSA model metadata
|
||||
|
||||
FastVideo VSA support is scoped to `FastVideo/FastWan2.2-TI2V-5B-Diffusers`, whose text-to-video and image-to-video modes both use `Wan22Pipeline`. It operates on a flattened DiT sequence but partitions tokens in
|
||||
the original latent video grid. Wan integrations therefore attach the
|
||||
FastWan's text-to-video and image-to-video modes both use `Wan22Pipeline`.
|
||||
It operates on a flattened DiT sequence but partitions tokens in the original
|
||||
latent video grid. Wan integrations therefore attach the
|
||||
post-patch `(T, H, W)` shape as `vsa_dit_seq_shape` attention metadata. The
|
||||
backend validates that `T * H * W` equals the sequence length, derives the
|
||||
runtime block count, and selects the configured top-k key/value blocks for
|
||||
@@ -104,3 +133,8 @@ the checkpoint contains those weights. Checkpoints without the projection do
|
||||
not allocate or execute it. When top-k selects every block, native checkpoints
|
||||
route to SDPA, while FastVideo DMD checkpoints preserve the VSA all-block path
|
||||
to retain checkpoint semantics.
|
||||
|
||||
MiniMax-H3 supplies `MiniMaxH3VSAImpl` through `impl_overrides`.
|
||||
`models/minimax_h3/attention/fastvideo_h3.py` owns its prefix layout, video tiling, and
|
||||
learned compression gate. Pooling, block-map construction, and tile64 provider
|
||||
calls are shared through `attention/ops/block_sparse.py`.
|
||||
+359
-1056
File diff suppressed because it is too large.
Load diff
@@ -16,6 +16,7 @@ required_reviewers:
|
||||
- "@yenuo26"
|
||||
- "@NickCao"
|
||||
primary_code_paths:
|
||||
- vllm_omni/engine/omni_engine_base.py
|
||||
- vllm_omni/engine/async_omni_engine.py
|
||||
- vllm_omni/engine/async_engine_utils.py
|
||||
- vllm_omni/engine/orchestrator.py
|
||||
@@ -28,8 +29,12 @@ related_code_paths:
|
||||
- vllm_omni/engine/stage_pool.py
|
||||
- vllm_omni/outputs/output_processor.py
|
||||
- vllm_omni/distributed/omni_coordinator/**
|
||||
# Duplex kernel design is owned by ../fullduplex.md; listed here because the
|
||||
# orchestrator hosts the duplex control plane.
|
||||
# The duplex siblings (DuplexOmniEngine / DuplexOrchestrator) and the
|
||||
# engine-resident session framework are owned by ../fullduplex.md; listed
|
||||
# here because they fill the template seams of OmniEngineBase and
|
||||
# OrchestratorBase.
|
||||
- vllm_omni/engine/duplex_omni_engine.py
|
||||
- vllm_omni/engine/duplex_orchestrator.py
|
||||
- vllm_omni/engine/duplex/**
|
||||
depends_on:
|
||||
- input_output_modality_contracts.md
|
||||
@@ -48,6 +53,8 @@ validation_paths:
|
||||
- tests/engine/test_cfg_companion_lifecycle.py
|
||||
- tests/engine/test_rpc_result_router.py
|
||||
- tests/engine/duplex/**
|
||||
- tests/engine/test_duplex_orchestrator.py
|
||||
- tests/engine/test_duplex_omni_engine.py
|
||||
- tests/engine/test_duplex_import_boundary.py
|
||||
upstream_refs:
|
||||
- vllm.v1.engine.EngineCoreRequest
|
||||
|
||||
@@ -19,7 +19,9 @@ required_reviewers:
|
||||
primary_code_paths:
|
||||
- vllm_omni/entrypoints/omni.py
|
||||
- vllm_omni/entrypoints/async_omni.py
|
||||
- vllm_omni/entrypoints/async_omni_base.py
|
||||
- vllm_omni/entrypoints/omni_base.py
|
||||
- vllm_omni/entrypoints/duplex_omni.py
|
||||
- vllm_omni/entrypoints/cli/**
|
||||
- vllm_omni/entrypoints/openai/**
|
||||
- vllm_omni/entrypoints/openpi/**
|
||||
@@ -27,14 +29,13 @@ primary_code_paths:
|
||||
- vllm_omni/entrypoints/stage_utils.py
|
||||
- vllm_omni/entrypoints/utils.py
|
||||
- vllm_omni/entrypoints/duplex/**
|
||||
- vllm_omni/entrypoints/duplex_request_client.py
|
||||
- vllm_omni/clients/**
|
||||
primary_path_exceptions:
|
||||
- path: vllm_omni/entrypoints/openai/errors.py
|
||||
owner: error_contracts.md
|
||||
- path: vllm_omni/entrypoints/duplex/**
|
||||
owner: ../fullduplex.md
|
||||
- path: vllm_omni/entrypoints/duplex_request_client.py
|
||||
- path: vllm_omni/entrypoints/duplex_omni.py
|
||||
owner: ../fullduplex.md
|
||||
- path: vllm_omni/clients/**
|
||||
owner: ../fullduplex.md
|
||||
@@ -42,7 +43,9 @@ related_code_paths:
|
||||
- vllm_omni/errors.py
|
||||
- vllm_omni/inputs/**
|
||||
- vllm_omni/outputs/**
|
||||
- vllm_omni/engine/omni_engine_base.py
|
||||
- vllm_omni/engine/async_omni_engine.py
|
||||
- vllm_omni/engine/duplex_omni_engine.py
|
||||
- vllm_omni/config/**
|
||||
- vllm_omni/deploy/**
|
||||
depends_on:
|
||||
@@ -53,7 +56,9 @@ validation_paths:
|
||||
- tests/entrypoints/test_omni_entrypoints.py
|
||||
- tests/entrypoints/test_async_omni.py
|
||||
- tests/entrypoints/test_async_omni_pause_sleep_routing.py
|
||||
- tests/entrypoints/test_async_omni_duplex.py
|
||||
- tests/entrypoints/test_duplex_omni.py
|
||||
- tests/entrypoints/duplex/**
|
||||
- tests/entrypoints/openai_api/test_duplex_api_server.py
|
||||
- tests/entrypoints/test_serve.py
|
||||
- tests/entrypoints/test_stream_finish_reason.py
|
||||
- tests/entrypoints/openai/**
|
||||
|
||||
@@ -81,6 +81,18 @@ To build a simple workflow yourself,
|
||||
- For some multi-stage models like BAGEL, [only one stage's sampling parameters are exposed and tunable via vLLM-Omni's online serving API](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/examples/online_serving/bagel/). Thus, these models are treated as single-stage ones. Please check the vLLM-Omni documentation on how to correctly set each model's sampling parameters.
|
||||
- For multi-stage models where all stages are either autoregression or diffusion, you can also connect only a single Sampling Params node, indicating that this set of sampling parameters will be used for all stages.
|
||||
|
||||
## MiniMax H3 text-to-video workflow
|
||||
|
||||
The **MiniMax H3 Text to Video** template uses the existing remote Generate Video,
|
||||
H3 Params, sampling, and Remote LoRA nodes. It includes Base settings and an
|
||||
optional Turbo preset, with native resolution, 24 FPS, and `17k+5` frame counts.
|
||||
ComfyUI receives the generated video and audio from the server; it does not load
|
||||
H3 weights locally. Configure the server-side LoRA path before enabling Turbo.
|
||||
|
||||
See the [H3 workflow guide](https://github.com/vllm-project/vllm-omni/blob/main/apps/ComfyUI-vLLM-Omni/docs/minimax-h3-t2v.md)
|
||||
for server setup, template import, Turbo configuration, recorded validation,
|
||||
and saved-video/audio checks.
|
||||
|
||||
## Examples & Screenshots
|
||||
|
||||
Please read the [ComfyUI integration's README](https://github.com/vllm-project/vllm-omni/tree/main/apps/ComfyUI-vLLM-Omni) for more info.
|
||||
@@ -25,8 +25,8 @@ CATEGORIES = {
|
||||
"clients": {
|
||||
"name": "Clients",
|
||||
"description": (
|
||||
"Client libraries that talk to a running vLLM-Omni server "
|
||||
"(the full-duplex Realtime client and its per-model session presets)."
|
||||
"Client libraries for full-duplex sessions: the Realtime WebSocket client, "
|
||||
"the in-process client over `DuplexOmni`, and the per-model session presets."
|
||||
),
|
||||
},
|
||||
"inputs": {
|
||||
|
||||
@@ -107,7 +107,7 @@ th {
|
||||
| `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | | — |
|
||||
| `CovoAudioForConditionalGeneration` | Covo-Audio-Chat | `tencent/Covo-Audio-Chat` | ✅︎ | | | | — |
|
||||
| `DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | ✅︎ | — |
|
||||
| `MiniCPMO45OmniForConditionalGeneration` | MiniCPM-o 4.5 | `openbmb/MiniCPM-o-4_5` | ✅︎ | | ✅︎ | | — |
|
||||
| `MiniCPMO45OmniForConditionalGeneration` | MiniCPM-o 4.5 | `openbmb/MiniCPM-o-4_5` | ✅︎ | | ✅︎ | | [Repository](https://github.com/vllm-project/vllm-omni/blob/main/recipes/OpenBMB/MiniCPM-o-4_5.md) |
|
||||
| `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | — |
|
||||
| `GepardTalkerForConditionalGeneration` | Gepard-1.0 | `nineninesix/gepard-1.0` | ✅︎ | | | ✅︎ | — |
|
||||
| `HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | | — |
|
||||
|
||||
+12
-5
@@ -145,15 +145,22 @@ interchangeably.
|
||||
| Incremental text input for speech synthesis | `WS /v1/audio/speech/stream` | Send text events and receive audio | [Streaming Text to Speech](speech_api.md#streaming-text-input-websocket) |
|
||||
| Live video understanding | `WS /v1/video/chat/stream` | Send video frames and receive text/audio | [Streaming Video Input](video_stream_api.md) |
|
||||
| Turn-based realtime audio | `WS /v1/realtime` | Stream one audio input and receive transcript/audio events | [Realtime Audio](realtime_api.md) |
|
||||
| Continuous speech-to-speech interaction | `WS /v1/realtime?duplex=1` or `WS /v1/duplex` | Listen and speak concurrently with session control | [Full Duplex](full_duplex_api.md) |
|
||||
| Continuous speech-to-speech interaction | `WS /v1/realtime?duplex=1` (alias `WS /v1/duplex`) | Listen and speak concurrently with session control | [Full Duplex](full_duplex_api.md) |
|
||||
| Generated video chunks | `WS /v1/realtime/video` | Start a diffusion request and receive fragmented MP4 | [Streaming Video Output](streaming_video_output_api.md) |
|
||||
| Robot policy inference | `WS /v1/realtime/robot/openpi` | Send MessagePack observations and receive action arrays | [OpenPI Robot Policy](openpi_api.md) |
|
||||
|
||||
All six routes are model- or configuration-dependent. In particular,
|
||||
`/v1/realtime` is not full duplex unless the client sets `duplex=1` and the
|
||||
deployment explicitly enables duplex sessions. Clients should also verify the
|
||||
duplex capability payload because the query-parameter form falls back to the
|
||||
ordinary realtime handler when duplex is unavailable.
|
||||
`/v1/realtime` is full duplex when the model is a duplex model (its pipeline
|
||||
declares a `duplex_plugin` and the deploy configuration sets
|
||||
`session_mode: duplex`); a stock Realtime client needs no vendor query
|
||||
parameter, and `duplex=0` is the explicit opt-out. Such a server serves the
|
||||
websocket route plus `POST /v1/chat/completions`, and no other turn-based
|
||||
HTTP route. On a duplex server that chat route is not the turn-based path:
|
||||
each request runs on a short-lived duplex session, so it holds one of
|
||||
`duplex_session.max_sessions` for its lifetime and answers at the model's
|
||||
real-time pace -- see [Full Duplex](full_duplex_api.md). Clients should
|
||||
also verify the duplex capability payload because the query-parameter form
|
||||
falls back to the ordinary realtime handler when duplex is unavailable.
|
||||
|
||||
## Related Endpoints
|
||||
|
||||
|
||||
@@ -15,40 +15,52 @@ capability gates, see the [Realtime Duplex API guide](realtime_duplex_api.md).
|
||||
|
||||
| Endpoint | Protocol | Recommended use |
|
||||
| --- | --- | --- |
|
||||
| `WS /v1/realtime?duplex=1` | OpenAI Realtime-style event projection | Applications and browser clients |
|
||||
| `WS /v1/duplex` | Native vLLM-Omni duplex events | Runtime integration and low-level testing |
|
||||
|
||||
Both endpoints use the same duplex engine and require the same model-side
|
||||
adapter. Prefer the Realtime projection unless the native lifecycle events are
|
||||
specifically required.
|
||||
| `WS /v1/realtime?duplex=1` | OpenAI Realtime-style events (the normative contract) | Applications and browser clients |
|
||||
| `WS /v1/duplex` | Alias of `/v1/realtime?duplex=1` (same protocol, same handler) | Clients that prefer a dedicated path |
|
||||
| Python: `DuplexOmni` / `InlineDuplexClient` | Typed commands and events in-process | Embedding the model without a server |
|
||||
|
||||
## Enable Full Duplex
|
||||
|
||||
The route becomes usable only when the deployment configuration explicitly
|
||||
sets:
|
||||
A model is served full duplex when its registered pipeline declares a
|
||||
`duplex_plugin` (the model's `DuplexModelPlugin`). That declaration alone
|
||||
decides it: `vllm-omni serve` constructs `DuplexOmni` instead of `AsyncOmni`,
|
||||
and every surface the server exposes is backed by a duplex session. It serves
|
||||
`/v1/realtime?duplex=1` (and its alias `/v1/duplex`),
|
||||
`POST /v1/chat/completions`, `/v1/models` and `/health`; every other
|
||||
turn-based HTTP route (speech, batch, embeddings, video, ...) answers "not
|
||||
available". Turn-based use of such a model stays available offline through the
|
||||
Python API (`Omni` / `AsyncOmni`).
|
||||
|
||||
`/v1/chat/completions` is the ordinary chat service running on the duplex
|
||||
engine: a request is a turn-based generation on the same stages, served
|
||||
alongside the live websocket sessions, with the request options the chat
|
||||
service supports. It is wired only when the model's plugin declares
|
||||
`DuplexCapabilities.supports_chat_completions`; on any other duplex model the
|
||||
route reports "not available". A model that should not serve the route at all
|
||||
lists it in the deploy configuration's `endpoint_restrictions`.
|
||||
|
||||
The deploy configuration of such a model must agree:
|
||||
|
||||
```yaml
|
||||
session_mode: duplex
|
||||
```
|
||||
|
||||
The selected model pipeline must also provide a duplex serving adapter.
|
||||
Model-native deployments configure the corresponding engine runtime extension
|
||||
and control plane as part of their registered pipeline.
|
||||
A duplex model started with a deploy configuration that does not set it fails
|
||||
at startup rather than falling back to turn-based serving.
|
||||
|
||||
!!! warning
|
||||
|
||||
`WS /v1/duplex` fails with `Duplex API is not available` when duplex is not
|
||||
enabled. By contrast, `/v1/realtime?duplex=1` falls back to the ordinary
|
||||
turn-based realtime handler when no duplex handler exists. Confirm that
|
||||
On a deployment that is *not* duplex, `WS /v1/duplex` fails with
|
||||
`Duplex API is not available` and `/v1/realtime?duplex=1` falls back to the
|
||||
ordinary turn-based realtime handler. Confirm that
|
||||
`session.created.session.capabilities` is present before treating the
|
||||
connection as full duplex.
|
||||
|
||||
The current unified-runtime integrations are:
|
||||
|
||||
- MiniCPM-o 4.5, using `vllm_omni/deploy/minicpmo_4_5.yaml`
|
||||
(`session_mode: duplex`);
|
||||
- PersonaPlex, whose default `vllm_omni/deploy/personaplex.yaml` enables duplex;
|
||||
- Nemotron VoiceChat, via its registered duplex plugin package.
|
||||
**MiniCPM-o 4.5** (`vllm_omni/deploy/minicpmo_4_5.yaml`) is the only model
|
||||
served over this endpoint today. PersonaPlex and Nemotron VoiceChat still carry
|
||||
their pre-framework duplex code: their pipelines declare no `duplex_plugin`, so
|
||||
they run turn-based until the follow-up PRs port them to the plugin contract
|
||||
(RFC [vllm-omni#7181](https://github.com/vllm-project/vllm-omni/issues/7181)).
|
||||
|
||||
JoyVL is a separate HTTP interaction orchestrator and does not use these
|
||||
WebSocket endpoints. See [Standalone Experimental Servers](standalone_servers.md).
|
||||
@@ -101,22 +113,23 @@ deferred, treated as a short acknowledgement, or used to interrupt output.
|
||||
The `session.created` payload includes capability fields such as
|
||||
`supports_barge_in`, `supports_playback_ack`, `supports_multi_session`,
|
||||
`supports_session_resume`, and `chunk_period_ms`. Treat this payload as the
|
||||
runtime contract.
|
||||
runtime contract and branch on the flags, never on the model name: a model
|
||||
that supports native overlapping speech may still advertise
|
||||
`supports_barge_in=false` when destructive output interruption and model-state
|
||||
rewind are not validated for it. Capacity and session-resume behavior also
|
||||
depend on the selected deployment configuration. The per-model table lives in
|
||||
the [Realtime Duplex API guide](realtime_duplex_api.md#capability-negotiation-by-model).
|
||||
|
||||
For example, PersonaPlex supports native overlapping speech but currently
|
||||
advertises `supports_barge_in=false`; destructive output interruption and
|
||||
model-state rewind have not been validated for that integration. Capacity and
|
||||
session-resume behavior also depend on the selected deployment configuration.
|
||||
## Python API
|
||||
|
||||
## Native Protocol
|
||||
`vllm_omni.entrypoints.duplex_omni.DuplexOmni` runs the same engine-resident
|
||||
sessions in-process: `open_session()` returns a `DuplexSessionHandle` whose
|
||||
`submit()` takes typed `DuplexCommand` objects and whose `events()` yields
|
||||
typed `DuplexEvent` objects (each with a `to_realtime()` wire rendering).
|
||||
`vllm_omni.clients.inline_duplex.InlineDuplexClient` exposes that handle
|
||||
behind the `DuplexClient` API. See the
|
||||
[Realtime Duplex API guide](realtime_duplex_api.md#using-the-python-api).
|
||||
|
||||
`WS /v1/duplex` exposes lower-level lifecycle names including
|
||||
`session.create`, `input_audio_buffer.append`, `turn.signal`, `playback.ack`,
|
||||
and `session.close`. It returns native session, input, response, overlap, and
|
||||
error events. This protocol is experimental and may evolve with the runtime;
|
||||
applications should use the provided Realtime client where possible.
|
||||
|
||||
See the [MiniCPM-o example](https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/minicpmo),
|
||||
[PersonaPlex example](https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/personaplex),
|
||||
See the [MiniCPM-o example](https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/minicpmo)
|
||||
and the [full-duplex runtime design](../design/fullduplex.md)
|
||||
for model-specific validation and architecture details.
|
||||
+214
-135
@@ -1,17 +1,18 @@
|
||||
# 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 only — and just for deploy configurations that declare
|
||||
`session_mode: duplex`. PersonaPlex and Nemotron VoiceChat are ported to the
|
||||
plugin contract in follow-up PRs and are not served over this endpoint yet. The runtime architecture
|
||||
is described in [Full-Duplex Runtime (MiniCPM-o 4.5)](../design/fullduplex.md).
|
||||
|
||||
## Quick Start
|
||||
|
||||
@@ -26,11 +27,38 @@ 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`
|
||||
(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.
|
||||
`duplex_session.max_sessions: 4`. Because the MiniCPM-o 4.5 pipeline declares a
|
||||
`duplex_plugin`, `vllm-omni serve` runs it through `DuplexOmni`: the server
|
||||
mounts `ws://<host>:8099/v1/realtime?duplex=1` (this page; `ws://<host>:8099/v1/duplex`
|
||||
is an alias of the same route), `POST /v1/chat/completions`, `/v1/models` and
|
||||
`/health`, and nothing else. The remaining turn-based HTTP routes (speech,
|
||||
batch, embeddings, video, ...) are not served by a duplex model.
|
||||
|
||||
`/v1/chat/completions` is not a turn-based path on a duplex server: each
|
||||
request opens a short-lived duplex session and reads the answer off it. Audio
|
||||
content becomes an ordinary committed turn; text is seeded as the session's
|
||||
opening turn, because a model-native model takes a turn only when it hears
|
||||
speech. A model that cannot be seeded that way declares so, and a text request
|
||||
to it is refused with 400. Two consequences are worth sizing for:
|
||||
|
||||
- **Every request holds an admission slot** for its lifetime, so
|
||||
`duplex_session.max_sessions` caps HTTP concurrency as well as websocket
|
||||
sessions. Beyond it a request is refused with HTTP 503.
|
||||
- **Latency follows the model's own clock.** A model-native session generates
|
||||
per audio unit, so even a short answer costs what that many units cost --
|
||||
noticeably more than the same model would take turn-based.
|
||||
|
||||
`n > 1`, `logprobs`, `tools`, and image or video content parts have no
|
||||
representation in one duplex turn and are refused with HTTP 400 -- a Realtime
|
||||
conversation item carries text and audio only. `temperature`,
|
||||
`max_completion_tokens`, `modalities` and
|
||||
`chat_template_kwargs.use_tts_template` apply; the other
|
||||
`chat_template_kwargs` keys are logged as ignored. `extra_body` is passed to
|
||||
the session unchanged, exactly as a websocket client's session `extra_body`
|
||||
is, which is how a request asking for `modalities: ["text", "audio"]` supplies
|
||||
the `ref_audio` MiniCPM-o requires for audio output. A model that should not
|
||||
serve the route lists it in the deploy config's `endpoint_restrictions`,
|
||||
exactly as a turn-based model does.
|
||||
|
||||
### Run the Example Client
|
||||
|
||||
@@ -81,25 +109,24 @@ ships a preset:
|
||||
|
||||
| Preset | Input / output audio | What it sets |
|
||||
| --- | --- | --- |
|
||||
| `vllm_omni.clients.minicpmo_4_5.create_duplex_session_config(ref_audio=...)` | 16 kHz `pcm16` / 24 kHz `pcm16` | `extra_body.native_duplex=True` (opt in to the model-native lane), `force_listen_count=0`, `overlap_policy="listen_only"`, `playback_commit_policy="ack_only"`; `ref_audio` is the assistant voice clip |
|
||||
| `vllm_omni.clients.minicpmo_4_5.create_duplex_session_config(ref_audio=...)` | 16 kHz `pcm16` / 24 kHz `pcm16` | `force_listen_count=0`, `overlap_policy="listen_only"`, `playback_commit_policy="ack_only"`; `ref_audio` is the assistant voice clip |
|
||||
| `vllm_omni.clients.personaplex.create_duplex_session_config(voice="NATF2.pt", persona="")` | 24 kHz `pcm_f32le` / 24 kHz `pcm16` | bundled `.pt` voice prompt and the persona as `instructions` |
|
||||
| `SessionConfig(...)` | 16 kHz `pcm16` / 24 kHz `pcm16` | model-neutral defaults; pass `extra_body`, `turn_detection`, `overlap_policy`, `playback_commit_policy`, `instructions`, `voice`, `temperature` yourself |
|
||||
|
||||
Keyword overrides on a preset replace the corresponding `SessionConfig`
|
||||
field, e.g. `create_duplex_session_config(ref_audio=..., temperature=0.6)`.
|
||||
|
||||
Constructor options: `session_id` (names the session the client creates; it
|
||||
does not resume or take over an existing session — entering the client always
|
||||
performs the `session.update` handshake), `reconnect` (a `ReconnectPolicy`,
|
||||
default five attempts with 0.25–4 s jittered backoff; `None` disables
|
||||
auto-resume), `heartbeat_interval_s` (default 30 s; `None` disables the lease
|
||||
heartbeat), `handshake_timeout_s`, and `connect` (a custom transport factory,
|
||||
used by tests and benchmarks). Resume is supported only as automatic reconnect
|
||||
within the same `DuplexClient` instance (see
|
||||
[Reconnect and resume](#reconnect-and-resume)); re-attaching from a new client
|
||||
or process requires the wire-level `session.resume` handshake
|
||||
(`resume_token`, `incarnation`, `last_received_server_event_seq`), which this
|
||||
client does not expose.
|
||||
Constructor options: `reconnect` (a `ReconnectPolicy`, default five
|
||||
attempts with 0.25–4 s jittered backoff; `None` disables auto-resume),
|
||||
`heartbeat_interval_s` (default 30 s; `None` disables the lease heartbeat),
|
||||
`handshake_timeout_s`, and `connect` (a custom transport factory, used by
|
||||
tests and benchmarks). The session id is allocated by the server and becomes
|
||||
`client.session_id` once `session.created` arrives; a client cannot choose
|
||||
it. Resume is supported only as automatic reconnect within the same
|
||||
`DuplexClient` instance (see [Reconnect and resume](#reconnect-and-resume));
|
||||
re-attaching from a new client or process requires the wire-level
|
||||
`session.resume` handshake (`resume_token`,
|
||||
`last_received_server_event_seq`), which this client does not expose.
|
||||
|
||||
### Stream audio and commit a turn
|
||||
|
||||
@@ -195,7 +222,7 @@ the stamped `event_id`.
|
||||
On a transport drop the client reconnects with the `ReconnectPolicy` backoff
|
||||
and sends `session.resume` from the last acknowledged `server_event_seq`.
|
||||
Replayed events are deduplicated, consumers see one `ConnectionResumed`
|
||||
event, and `client.incarnation` / `client.resume_token` are refreshed. The
|
||||
event, and `client.resume_token` is refreshed. The
|
||||
server keeps a detached session alive for `disconnect_grace_s` (30 s by
|
||||
default); after that a resume fails with `session_resume_expired`. Resume is
|
||||
gated by `capabilities.supports_session_resume` (MiniCPM-o 4.5 and Nemotron
|
||||
@@ -238,9 +265,67 @@ definitions out of the dependency-free client. `audio_bytes(response_id)`, `coun
|
||||
`openai-realtime-duplex` benchmark backend (`vllm-omni bench serve`) is built
|
||||
on the same helpers.
|
||||
|
||||
## Using the Python API
|
||||
|
||||
`vllm_omni.entrypoints.duplex_omni.DuplexOmni` is the in-process Python API
|
||||
of a duplex model: sessions run inside the engine, and the class only opens,
|
||||
resumes and closes them and pipes typed commands in and typed events out
|
||||
through a `DuplexSessionHandle`.
|
||||
|
||||
```python
|
||||
from vllm_omni.engine.duplex.commands import AppendAudio, Commit
|
||||
from vllm_omni.engine.duplex.events import AudioDelta, SessionClosed
|
||||
from vllm_omni.entrypoints.duplex_omni import DuplexOmni
|
||||
|
||||
omni = DuplexOmni(model="openbmb/MiniCPM-o-4_5", trust_remote_code=True)
|
||||
async with await omni.open_session({"modalities": ["audio", "text"], "ref_audio": ref_audio_data_url}) as session:
|
||||
async def consume():
|
||||
async for event in session.events(): # typed DuplexEvent objects
|
||||
if isinstance(event, AudioDelta):
|
||||
play(event.audio)
|
||||
elif isinstance(event, ResponseDone):
|
||||
break # one turn is enough here
|
||||
task = asyncio.create_task(consume())
|
||||
await session.submit(AppendAudio(audio=pcm16_bytes, format="pcm16", sample_rate_hz=16000))
|
||||
await session.submit(Commit())
|
||||
await task
|
||||
# leaving the block closes the session (the last event is SessionClosed)
|
||||
```
|
||||
|
||||
`open_session` takes a `DuplexSessionConfig` or the same Realtime `session`
|
||||
object the WebSocket handshake takes; the session id is always allocated by
|
||||
`DuplexOmni`. `DuplexSessionHandle` offers `submit(command)` plus one
|
||||
convenience wrapper per command (`append_audio`, `commit`, `create_response`,
|
||||
`cancel_response`, `barge_in`, `ack_playback`, `update`, `create_item`,
|
||||
`delete_item`, `truncate_item`, `heartbeat`, ...). `events()` is a
|
||||
single-consumer async iterator that ends after `session.closed` /
|
||||
`session.expired`; rejected commands come back as `ErrorEvent`s on it.
|
||||
Every event renders the wire JSON with `to_realtime()`, so anything written
|
||||
against the WebSocket protocol works unchanged on the typed stream.
|
||||
|
||||
`vllm_omni.clients.inline_duplex.InlineDuplexClient` wraps a `DuplexOmni`
|
||||
behind the `DuplexClient` API, so the same application code runs in-process
|
||||
(no server) or over the network:
|
||||
|
||||
```python
|
||||
from vllm_omni.clients.inline_duplex import InlineDuplexClient
|
||||
from vllm_omni.clients.minicpmo_4_5 import create_duplex_session_config
|
||||
|
||||
config = create_duplex_session_config(ref_audio=audio_data_url(Path("reference_voice.wav")))
|
||||
async with InlineDuplexClient(omni, model="openbmb/MiniCPM-o-4_5", config=config) as client:
|
||||
await client.stream_pcm(question)
|
||||
await client.commit()
|
||||
async for response in client.responses():
|
||||
...
|
||||
```
|
||||
|
||||
`examples/online_serving/barge_in_client.py --inline` drives the barge-in
|
||||
scenario this way. The inline client has no transport, so it has no
|
||||
reconnect, resume, heartbeat or event-acknowledgement machinery.
|
||||
|
||||
## 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,21 +338,26 @@ 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`,
|
||||
`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
|
||||
The endpoint is `ws(s)://<host>/v1/realtime`. On a deployment whose pipeline
|
||||
declares a `duplex_plugin` and whose deploy configuration declares
|
||||
`session_mode: duplex`, a bare connection is already a duplex session, so a
|
||||
stock Realtime client needs no vendor query parameter. `?duplex=1` (`1`,
|
||||
`true` or `on`) is the spelling this page uses and stays valid;
|
||||
`?duplex=0` (`0`, `false`, `off`) explicitly opts out and selects the
|
||||
turn-based Realtime handler, which a duplex server does not mount --
|
||||
that connection is answered with `Realtime API is not available`.
|
||||
`ws(s)://<host>/v1/duplex` is an alias that always selects the duplex
|
||||
handler. Optional query parameters are `model`, `autostart` (`0` means
|
||||
resume-only) and `resume`. Every message is one JSON object per WebSocket
|
||||
text frame, discriminated by `type`. Inbound events are applied in arrival
|
||||
order through one per-session mailbox; outbound events preserve that order
|
||||
and carry a monotonically increasing `server_event_seq`, which is the replay
|
||||
cursor for `session.resume`. Client events may carry an `event_id`, echoed in
|
||||
`error.error.event_id` when the event is rejected.
|
||||
|
||||
Session identity is `(session_id, incarnation)`; `attachment_generation`
|
||||
Session identity is the server-allocated `session_id` (`duplex-<uuid>`,
|
||||
announced in `session.created`; a `session_id` the client puts in the session
|
||||
object is ignored, and ids are never reused); `attachment_generation`
|
||||
names the current socket; `epoch` is the interruption fence that advances on
|
||||
every cancel, clear, or barge-in; `turn_id` counts committed model turns;
|
||||
`response_id` (item id `item_<response_id>`) names one assistant utterance.
|
||||
@@ -281,12 +371,13 @@ with no OpenAI counterpart.
|
||||
|
||||
The event vocabulary is uniform, but several surfaces are gated by the
|
||||
`capabilities` object the server returns in `session.created`; a client must
|
||||
branch on those flags rather than on the model name. The current model
|
||||
plugins advertise:
|
||||
branch on those flags rather than on the model name. MiniCPM-o 4.5 is the
|
||||
only model on the plugin contract today; the other two columns record what
|
||||
their integrations advertise once the follow-up PRs port them:
|
||||
|
||||
| Capability | MiniCPM-o 4.5 | PersonaPlex | Nemotron VoiceChat | Gated surface |
|
||||
| --- | --- | --- | --- | --- |
|
||||
| `implementation_level` | `model_native_duplex` (when `extra_body.native_duplex=true`) | `model_native_duplex` | `model_native_duplex` | model-owned `response.listen` / `response.speak`; the chat-fallback lane otherwise |
|
||||
| `implementation_level` | `model_native_duplex` | `model_native_duplex` | `model_native_duplex` | model-owned `response.listen` / `response.speak` (constant: every duplex model is model-native) |
|
||||
| `chunk_period_ms` | 1000 | 80 | 80 | the model unit that `response.listen` decisions and camera frames align to |
|
||||
| `supports_session_resume` | yes | no | yes | `session.resume`, `session.resumed`, `session.replaced`, replay after a transport drop |
|
||||
| `supports_barge_in` | yes | no | no | `barge_in`, `turn.signal{event:"barge_in"}`, `overlap_policy=barge_in_on_speech`, `turn_detection.server_vad` |
|
||||
@@ -316,14 +407,8 @@ work unmodified. The overlap falls into the three tiers used by the event catalo
|
||||
| Object shapes | `realtime.session`, `realtime.item`, `realtime.response`; item content parts `input_text` / `input_audio` / `output_text` / `output_audio`; `function_call` and `function_call_output` items; `previous_item_id` chaining; `output_index` / `content_index` addressing |
|
||||
| Sequencing | `response.created → output_item.added → content_part.added → deltas → *.done → content_part.done → output_item.done → response.done → rate_limits.updated`; `speech_started → speech_stopped → committed` |
|
||||
|
||||
Audio/transcript output uses only the current (non-beta) OpenAI event names
|
||||
(`response.output_audio.delta`, not the deprecated `response.audio.delta`),
|
||||
matching `openai.types.realtime` in the official `openai` Python SDK and what
|
||||
clients built against it (e.g. `livekit-plugins-openai`'s non-Azure code path)
|
||||
expect; a client still on the old beta SDK types would not recognize these
|
||||
events. Session fields and `conversation.item.*` events, in contrast, are
|
||||
deliberately dual: the server emits and accepts **both** the beta and the GA
|
||||
spellings at once (`output_modalities` alongside `modalities`,
|
||||
The server deliberately emits and accepts **both** the beta and the GA
|
||||
spellings at once (`response.output_audio.delta` alongside `output_modalities`,
|
||||
`conversation.item.added` alongside `conversation.item.created`) so either
|
||||
generation of OpenAI client parses the stream.
|
||||
|
||||
@@ -333,7 +418,7 @@ A stock client ignores the extra keys; the extensions are additive.
|
||||
|
||||
| Message | Extension |
|
||||
| --- | --- |
|
||||
| `session.created` / `session.updated` | top-level `incarnation`, `attachment_generation`, `resume_token`; inside `session`: `state`, `turn_state`, `epoch`, `turn_id`, `active_request_id`, `active_response_id`, `active_response_turn_id`, `overlap_policy`, `overlap_*_ms/rms`, `playback_commit_policy`, `playback`, `capabilities`, `ref_audio`, `extra_body`, `idle_timeout_s`, `response_format` |
|
||||
| `session.created` / `session.updated` | top-level `attachment_generation`, `resume_token`; inside `session`: `state`, `turn_state`, `epoch`, `turn_id`, `active_request_id`, `active_response_id`, `active_response_turn_id`, `overlap_policy`, `overlap_*_ms/rms`, `playback_commit_policy`, `playback`, `capabilities`, `ref_audio`, `extra_body`, `idle_timeout_s`, `response_format` |
|
||||
| `response.created` / `response.done` / `response.listen` | `response_id` at top level; the raw duplex event under `response.metadata` (`duplex_event`); `status_details.reason` uses vLLM reasons (`barge_in`, `client_cancelled`, `new_response`, …) |
|
||||
| `response.output_audio.delta` | `format`, `sample_rate_hz`, `metadata{session_id, epoch, model_speak, end_of_turn, audio_duration_ms, audio_text_marks, playback}` |
|
||||
| `response.speak` (inserted before the first delta) | not an OpenAI event, but rides the OpenAI response envelope (`response_id`, `item_id`, `output_index`, `content_index`) |
|
||||
@@ -349,14 +434,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`.
|
||||
@@ -365,11 +445,11 @@ Semantic divergences hidden behind shared names:
|
||||
|
||||
| Area | Messages / fields |
|
||||
| --- | --- |
|
||||
| Session lifetime | `session.heartbeat` / `session.heartbeat_ack`, `session.event_ack`, `session.close` / `session.closed`, `session.resume` / `session.resumed` / `session.replaced` / `session.expired` / `session.resync_required`; `server_event_seq`, `resume_token`, `incarnation`, `attachment_generation`; query params `?duplex=1`, `autostart`, `resume`, `native_duplex` |
|
||||
| Session lifetime | `session.heartbeat` / `session.heartbeat_ack`, `session.event_ack`, `session.close` / `session.closed`, `session.resume` / `session.resumed` / `session.replaced` / `session.expired` / `session.resync_required`; `server_event_seq`, `resume_token`, `attachment_generation`; query params `?duplex=1`, `autostart`, `resume` |
|
||||
| Turn-taking | `response.speak`, `response.listen`, `overlap.decision`, `overlap_policy`, `barge_in`, `turn.signal`, `input.cancel`, `input.text.append`, `epoch` / `turn_id` fencing, `force_listen` |
|
||||
| Playback truth | `playback.ack` / `playback.acknowledged`, `playback_commit_policy` (`ack_only` \| `commit_all_on_done`), `playback{generated_ms, sent_ms, played_ms, committed_ms}`, `history_committed`, `audio_text_marks` |
|
||||
| Model negotiation | `capabilities{implementation_level, supports_input_append, supports_barge_in, supports_session_resume, chunk_period_ms, input_modes, …}`, `ref_audio`, `extra_body` (`auto_response`, `native_duplex` — deprecated alias `minicpmo45_native_duplex` is folded into it, `force_listen_count`, `duplex_initial_user_text`) |
|
||||
| Diagnostics | `runtime.control`, `duplex.function_call.done`, the full error-code vocabulary (`REALTIME_ERROR_TYPES_BY_CODE` in `vllm_omni/entrypoints/duplex/realtime_state.py`) |
|
||||
| Model negotiation | `capabilities{implementation_level, supports_input_append, supports_barge_in, supports_session_resume, chunk_period_ms, input_modes, …}`, `ref_audio`, `extra_body` (`auto_response`, `force_listen_count`, `duplex_initial_user_text`) |
|
||||
| Diagnostics | `duplex.function_call.done`, the full error-code vocabulary (`REALTIME_ERROR_TYPES_BY_CODE` in `vllm_omni/engine/duplex/events.py`) |
|
||||
|
||||
#### Consequences for clients
|
||||
|
||||
@@ -433,11 +513,15 @@ remains an approximation, not word-level audio alignment.
|
||||
### Event catalogue
|
||||
|
||||
Every event type the Realtime route accepts or emits, with a worked example in the next section: 21
|
||||
client→server and 42 server→client. Aliases (`push_chunk`,
|
||||
`input.audio.append`, `input_text.append`, `push_text`, `signal_turn`,
|
||||
`audio.playback_ack`, `close_session`, `close`, `session_close`) share the
|
||||
payload of their canonical event and are not listed separately; the alias
|
||||
table lives in the normative contract section of [Full-Duplex Runtime (MiniCPM-o 4.5)](../design/fullduplex.md).
|
||||
client→server and 42 server→client. Aliases (`input_text.append`,
|
||||
`push_text`, `signal_turn`, `audio.playback_ack`, `close_session`, `close`)
|
||||
share the payload of their canonical event and are not listed
|
||||
separately; the alias table is `translate_realtime_command` in
|
||||
`vllm_omni/engine/duplex/realtime_commands.py`. The pre-Realtime WAV-append
|
||||
aliases `push_chunk` and `input.audio.append` are **not** accepted on this
|
||||
route: send `input_audio_buffer.append` with one of the supported input
|
||||
formats (`pcm16`, `pcm_s16le`, `s16le`, `pcm_f32le`, `g711_ulaw`,
|
||||
`g711_alaw`).
|
||||
|
||||
#### Client to server
|
||||
|
||||
@@ -447,18 +531,18 @@ table lives in the normative contract section of [Full-Duplex Runtime (MiniCPM-o
|
||||
| `session.heartbeat` | 3 | Keepalive; refreshes the engine session lease. |
|
||||
| `session.event_ack` | 3 | Acknowledges received server events by `server_event_seq`; trims the resume replay journal. |
|
||||
| `session.close` | 3 | Graceful close; answered with `session.closed`. |
|
||||
| `session.resume` | 3 | Re-attach a live session on a new socket (`resume_token`, `incarnation`, `last_received_server_event_seq`). |
|
||||
| `session.resume` | 3 | Re-attach a live session on a new socket (`session_id`, `resume_token`, `last_received_server_event_seq`). |
|
||||
| `input_audio_buffer.append` | 2 | Append one audio chunk. OpenAI name; extended with `is_speech`, `video_frames`, `duration_ms`, `audio_end_ms`, per-event `format`/`sample_rate_hz`. |
|
||||
| `input_audio_buffer.commit` | 2 | Seal the buffered utterance into a user item. OpenAI name; extended with `final`, `response_create`, `is_speech:false`. |
|
||||
| `input.commit` | 3 | Alias of `input_audio_buffer.commit` accepted by the runner. |
|
||||
| `input_audio_buffer.clear` | 1 | Drop un-committed input audio. |
|
||||
| `input.cancel` | 3 | Cancel pending input; advances `epoch`. |
|
||||
| `input.text.append` | 3 | Append text to the open input item (chat-fallback lane only). |
|
||||
| `input.text.append` | 3 | Append text input; model-native sessions reject it with `native_text_append_unsupported`. |
|
||||
| `conversation.item.create` | 1 | Add a completed user item (text, audio, or `function_call_output`) to history. |
|
||||
| `conversation.item.retrieve` | 1 | Fetch a stored item by id. |
|
||||
| `conversation.item.delete` | 1 | Delete a stored item by id. |
|
||||
| `conversation.item.truncate` | 1 | Truncate an assistant item's audio/transcript at `audio_end_ms`. |
|
||||
| `response.create` | 1 | Explicitly request a response with per-response overrides. |
|
||||
| `response.create` | 1 | Explicitly request a response with per-response overrides. Requires committed audio input; otherwise `response_create_without_input`. |
|
||||
| `response.cancel` | 1 | Cancel the active (or named) response; advances `epoch`. |
|
||||
| `output_audio_buffer.clear` | 1 | Discard queued output audio; advances `epoch` (OpenAI: WebRTC-only, here also WebSocket). |
|
||||
| `barge_in` | 3 | Explicit hard interrupt; requires `capabilities.supports_barge_in`. |
|
||||
@@ -469,10 +553,10 @@ table lives in the normative contract section of [Full-Duplex Runtime (MiniCPM-o
|
||||
|
||||
| Event type | Tier | Description |
|
||||
| --- | --- | --- |
|
||||
| `session.created` | 2 | Session opened. OpenAI name; adds `incarnation`, `attachment_generation`, `resume_token` and vLLM-Omni keys inside `session` (`epoch`, `turn_id`, `playback`, `capabilities`, ...). |
|
||||
| `session.created` | 2 | Session opened. OpenAI name; adds `attachment_generation`, `resume_token` and vLLM-Omni keys inside `session` (`id` is the server-allocated session id, `epoch`, `turn_id`, `playback`, `capabilities`, ...). |
|
||||
| `session.updated` | 2 | Echo of the effective session config after every `session.update` (same extended `session` object). |
|
||||
| `session.heartbeat_ack` | 3 | Reply to `session.heartbeat`. |
|
||||
| `session.closed` | 3 | Last event on the socket; `reason` ∈ `session_close`, `timeout`, `disconnect`, ... |
|
||||
| `session.closed` | 3 | Last event on the socket, emitted once the engine released the session; `reason` ∈ `client_close`, `disconnect`, `timeout`, `transport_error`, `shutdown`, ... |
|
||||
| `session.resumed` | 3 | Resume accepted; carries the new `attachment_generation` and rotated `resume_token`; journaled events are replayed after it. |
|
||||
| `session.replaced` | 3 | Sent to the superseded socket when another socket resumes the session. |
|
||||
| `session.expired` | 3 | Engine lease reaped (`disconnect_grace_expired`, `idle_ttl_expired`); socket closes after it. |
|
||||
@@ -508,8 +592,8 @@ table lives in the normative contract section of [Full-Duplex Runtime (MiniCPM-o
|
||||
| `duplex.function_call.done` | 3 | Raw fallback when a model tool call is malformed. |
|
||||
| `output_audio_buffer.cleared` | 1 | Reply to `output_audio_buffer.clear`. |
|
||||
| `playback.acknowledged` | 3 | Reply to `playback.ack`; returns the playback ledger and `history_committed`. |
|
||||
| `overlap.decision` | 3 | How an append that overlapped an active response was handled (`drop` \| `listen` \| `barge_in`). |
|
||||
| `runtime.control` | 3 | Diagnostic: an engine control signal was partially unsupported. |
|
||||
| `overlap.decision` | 3 | How an append that overlapped an active response was handled (`drop` \| `listen` \| `barge_in`); details under `metadata`. |
|
||||
| `turn.event` | 3 | A turn-taking signal the session accepted (`event`, resulting `turn_state`). |
|
||||
| `error` | 2 | Error envelope. OpenAI shape and classes; vLLM-Omni error codes. |
|
||||
|
||||
### Message examples
|
||||
@@ -528,7 +612,6 @@ additionally carries `"server_event_seq": <int>` (omitted below for brevity).
|
||||
"event_id": "evt_0001",
|
||||
"session": {
|
||||
"model": "openbmb/MiniCPM-o-4_5",
|
||||
"session_id": "sess_demo_01",
|
||||
"modalities": ["audio", "text"],
|
||||
"instructions": "You are a concise voice assistant.",
|
||||
"input_audio_format": "pcm16",
|
||||
@@ -541,7 +624,6 @@ additionally carries `"server_event_seq": <int>` (omitted below for brevity).
|
||||
"idle_timeout_s": 300,
|
||||
"extra_body": {
|
||||
"auto_response": true,
|
||||
"native_duplex": true,
|
||||
"force_listen_count": 0
|
||||
}
|
||||
}
|
||||
@@ -553,13 +635,12 @@ additionally carries `"server_event_seq": <int>` (omitted below for brevity).
|
||||
```json
|
||||
{
|
||||
"type": "session.created",
|
||||
"incarnation": 0,
|
||||
"attachment_generation": 1,
|
||||
"resume_token": "rt_3f9c…",
|
||||
"session": {
|
||||
"object": "realtime.session",
|
||||
"type": "realtime",
|
||||
"id": "sess_demo_01",
|
||||
"id": "duplex-5f2a…",
|
||||
"model": "openbmb/MiniCPM-o-4_5",
|
||||
"state": "open",
|
||||
"turn_state": "idle",
|
||||
@@ -612,7 +693,7 @@ additionally carries `"server_event_seq": <int>` (omitted below for brevity).
|
||||
`session.created` and every later `session.update`)
|
||||
|
||||
```json
|
||||
{"type": "session.updated", "session": {"object": "realtime.session", "id": "sess_demo_01", "…": "…"}}
|
||||
{"type": "session.updated", "session": {"object": "realtime.session", "id": "duplex-5f2a…", "…": "…"}}
|
||||
```
|
||||
|
||||
`C→S session.heartbeat` / `S→C session.heartbeat_ack`
|
||||
@@ -622,7 +703,7 @@ additionally carries `"server_event_seq": <int>` (omitted below for brevity).
|
||||
```
|
||||
|
||||
```json
|
||||
{"type": "session.heartbeat_ack", "session_id": "sess_demo_01"}
|
||||
{"type": "session.heartbeat_ack", "session_id": "duplex-5f2a…"}
|
||||
```
|
||||
|
||||
`C→S session.event_ack`
|
||||
@@ -638,7 +719,7 @@ additionally carries `"server_event_seq": <int>` (omitted below for brevity).
|
||||
```
|
||||
|
||||
```json
|
||||
{"type": "session.closed", "session_id": "sess_demo_01", "reason": "session_close"}
|
||||
{"type": "session.closed", "session_id": "duplex-5f2a…", "epoch": 3, "reason": "client_close", "event": {"type": "session.closed", "reason": "client_close"}}
|
||||
```
|
||||
|
||||
`C→S session.resume` (new socket, `?duplex=1&autostart=0`)
|
||||
@@ -646,8 +727,7 @@ additionally carries `"server_event_seq": <int>` (omitted below for brevity).
|
||||
```json
|
||||
{
|
||||
"type": "session.resume",
|
||||
"session_id": "sess_demo_01",
|
||||
"incarnation": 0,
|
||||
"session_id": "duplex-5f2a…",
|
||||
"resume_token": "rt_3f9c…",
|
||||
"last_received_server_event_seq": 40
|
||||
}
|
||||
@@ -659,8 +739,8 @@ events with `server_event_seq > 40` are replayed in order)
|
||||
```json
|
||||
{
|
||||
"type": "session.resumed",
|
||||
"session_id": "sess_demo_01",
|
||||
"incarnation": 0,
|
||||
"session_id": "duplex-5f2a…",
|
||||
"session": {"id": "duplex-5f2a…", "state": "open", "epoch": 3, "turn_id": 2, "...": "the engine's current session object"},
|
||||
"attachment_generation": 2,
|
||||
"resume_token": "rt_a71e…"
|
||||
}
|
||||
@@ -669,19 +749,19 @@ events with `server_event_seq > 40` are replayed in order)
|
||||
`S→C session.replaced` (sent to the superseded socket, which is then closed)
|
||||
|
||||
```json
|
||||
{"type": "session.replaced", "session_id": "sess_demo_01", "attachment_generation": 1}
|
||||
{"type": "session.replaced", "session_id": "duplex-5f2a…", "attachment_generation": 1}
|
||||
```
|
||||
|
||||
`S→C session.expired`
|
||||
|
||||
```json
|
||||
{"type": "session.expired", "session_id": "sess_demo_01", "incarnation": 0, "reason": "disconnect_grace_expired"}
|
||||
{"type": "session.expired", "session_id": "duplex-5f2a…", "reason": "disconnect_grace_expired"}
|
||||
```
|
||||
|
||||
`S→C session.resync_required`
|
||||
|
||||
```json
|
||||
{"type": "session.resync_required", "session_id": "sess_demo_01", "reason": "journal_gap"}
|
||||
{"type": "session.resync_required", "session_id": "duplex-5f2a…", "reason": "journal_gap"}
|
||||
```
|
||||
|
||||
#### Input
|
||||
@@ -772,7 +852,7 @@ silence/noise declaration (no runtime append; answered with `response.listen`)
|
||||
"item_id": "item_9b2d…",
|
||||
"event": {
|
||||
"type": "input.committed",
|
||||
"session_id": "sess_demo_01",
|
||||
"session_id": "duplex-5f2a…",
|
||||
"turn_id": 1,
|
||||
"epoch": 0,
|
||||
"history_len": 1,
|
||||
@@ -802,7 +882,7 @@ deferred commit during an active response (`event.response_create_deferred`)
|
||||
"type": "input_audio_buffer.committed",
|
||||
"previous_item_id": "item_9b2d…",
|
||||
"item_id": "item_c044…",
|
||||
"event": {"type": "input.committed", "session_id": "sess_demo_01", "turn_id": 2, "epoch": 0, "history_len": 3, "response_create_deferred": true, "message": {"…": "…"}}
|
||||
"event": {"type": "input.committed", "session_id": "duplex-5f2a…", "turn_id": 2, "epoch": 0, "history_len": 3, "response_create_deferred": true, "message": {"…": "…"}}
|
||||
}
|
||||
```
|
||||
|
||||
@@ -911,7 +991,7 @@ deferred commit during an active response (`event.response_create_deferred`)
|
||||
"status_details": null,
|
||||
"output": [],
|
||||
"modalities": ["audio", "text"],
|
||||
"metadata": {"duplex_event": {"type": "response.created", "session_id": "sess_demo_01", "response_id": "resp_01", "epoch": 0, "turn_id": 1}}
|
||||
"metadata": {"duplex_event": {"type": "response.created", "session_id": "duplex-5f2a…", "response_id": "resp_01", "epoch": 0, "turn_id": 1}}
|
||||
}
|
||||
}
|
||||
```
|
||||
@@ -947,7 +1027,7 @@ deferred commit during an active response (`event.response_create_deferred`)
|
||||
"item_id": "item_resp_01",
|
||||
"output_index": 0,
|
||||
"content_index": 0,
|
||||
"metadata": {"session_id": "sess_demo_01", "epoch": 0, "model_speak": true}
|
||||
"metadata": {"session_id": "duplex-5f2a…", "epoch": 0, "model_speak": true}
|
||||
}
|
||||
```
|
||||
|
||||
@@ -964,7 +1044,7 @@ deferred commit during an active response (`event.response_create_deferred`)
|
||||
"format": "pcm16",
|
||||
"sample_rate_hz": 24000,
|
||||
"metadata": {
|
||||
"session_id": "sess_demo_01",
|
||||
"session_id": "duplex-5f2a…",
|
||||
"epoch": 0,
|
||||
"model_speak": true,
|
||||
"end_of_turn": false,
|
||||
@@ -1029,7 +1109,7 @@ deferred commit during an active response (`event.response_create_deferred`)
|
||||
"status": "completed",
|
||||
"status_details": null,
|
||||
"output": [{"id": "item_resp_01", "object": "realtime.item", "type": "message", "role": "assistant", "status": "completed", "content": [{"type": "output_audio", "transcript": "It is sunny, 24 degrees."}]}],
|
||||
"metadata": {"type": "response.done", "session_id": "sess_demo_01", "response_id": "resp_01", "epoch": 0, "committed": true, "playback": {"generated_ms": 2400, "sent_ms": 2400, "played_ms": 2400, "committed_ms": 2400}}
|
||||
"metadata": {"type": "response.done", "session_id": "duplex-5f2a…", "response_id": "resp_01", "epoch": 0, "committed": true, "playback": {"generated_ms": 2400, "sent_ms": 2400, "played_ms": 2400, "committed_ms": 2400}}
|
||||
}
|
||||
}
|
||||
```
|
||||
@@ -1055,9 +1135,9 @@ deferred commit during an active response (`event.response_create_deferred`)
|
||||
```json
|
||||
{
|
||||
"type": "response.listen",
|
||||
"session_id": "sess_demo_01",
|
||||
"session_id": "duplex-5f2a…",
|
||||
"epoch": 0,
|
||||
"response": {"object": "realtime.response", "status": "listening", "metadata": {"type": "response.listen", "session_id": "sess_demo_01", "epoch": 0, "reason": "model_listen"}}
|
||||
"response": {"object": "realtime.response", "status": "listening", "metadata": {"type": "response.listen", "session_id": "duplex-5f2a…", "epoch": 0, "reason": "model_listen"}}
|
||||
}
|
||||
```
|
||||
|
||||
@@ -1104,7 +1184,7 @@ function-call response (Nemotron VoiceChat tools)
|
||||
raw fallback when a model tool call is malformed (`call_id`/`name` missing)
|
||||
|
||||
```json
|
||||
{"type": "duplex.function_call.done", "event": {"type": "function_call.done", "session_id": "sess_demo_01", "epoch": 0, "call_id": null, "name": null, "arguments": "{"}}
|
||||
{"type": "duplex.function_call.done", "event": {"type": "function_call.done", "session_id": "duplex-5f2a…", "epoch": 0, "call_id": null, "name": null, "arguments": "{"}}
|
||||
```
|
||||
|
||||
#### Playback
|
||||
@@ -1120,7 +1200,7 @@ raw fallback when a model tool call is malformed (`call_id`/`name` missing)
|
||||
"type": "playback.acknowledged",
|
||||
"event": {
|
||||
"type": "playback.acknowledged",
|
||||
"session_id": "sess_demo_01",
|
||||
"session_id": "duplex-5f2a…",
|
||||
"epoch": 0,
|
||||
"item_id": "item_resp_01",
|
||||
"played_ms": 1000,
|
||||
@@ -1145,20 +1225,17 @@ truncating ack after an interruption (what was actually heard)
|
||||
```json
|
||||
{
|
||||
"type": "overlap.decision",
|
||||
"session_id": "sess_demo_01",
|
||||
"session_id": "duplex-5f2a…",
|
||||
"epoch": 0,
|
||||
"policy": "listen_only",
|
||||
"action": "listen",
|
||||
"reason": "long_overlap_speech",
|
||||
"overlap_ms": 1400,
|
||||
"buffer_audio": true,
|
||||
"defer_runtime_append": true,
|
||||
"force_listen": true
|
||||
"metadata": {"overlap_ms": 1400, "buffer_audio": true, "defer_runtime_append": true, "force_listen": true}
|
||||
}
|
||||
```
|
||||
|
||||
```json
|
||||
{"type": "overlap.decision", "session_id": "sess_demo_01", "epoch": 0, "policy": "barge_in_on_speech", "action": "barge_in", "reason": "server_vad_utterance_active"}
|
||||
{"type": "overlap.decision", "session_id": "duplex-5f2a…", "epoch": 0, "policy": "barge_in_on_speech", "action": "barge_in", "reason": "server_vad_utterance_active"}
|
||||
```
|
||||
|
||||
server-VAD session (`session.update` fragment) and the annotated append the
|
||||
@@ -1172,12 +1249,6 @@ translator forwards internally (shown for reference; clients never see it)
|
||||
{"type": "input_audio_buffer.append", "audio": "…", "format": "pcm16", "sample_rate_hz": 16000, "is_speech": true, "force_listen": true, "vad": {"backend": "silero", "is_speech": true, "speech_active": true, "speech_started": true, "speech_stopped": false, "speech_probability": 0.93}}
|
||||
```
|
||||
|
||||
`S→C runtime.control` (only when an engine signal was partially unsupported)
|
||||
|
||||
```json
|
||||
{"type": "runtime.control", "session_id": "sess_demo_01", "epoch": 1, "result": {"…": "redacted engine control result"}}
|
||||
```
|
||||
|
||||
#### Errors
|
||||
|
||||
`S→C error` — Realtime shape
|
||||
@@ -1203,71 +1274,79 @@ translator forwards internally (shown for reference; clients never see it)
|
||||
```
|
||||
|
||||
```json
|
||||
{"type": "error", "error": {"type": "rate_limit_error", "code": "duplex_session_busy", "message": "duplex session capacity reached (max_sessions=4)"}}
|
||||
{"type": "error", "error": {"type": "rate_limit_error", "code": "resource_exhausted", "message": "duplex_session_capacity_exhausted: limit=4"}}
|
||||
```
|
||||
|
||||
Handshake-stage errors (before the Realtime projector is bound) use the
|
||||
flat native shape:
|
||||
|
||||
```json
|
||||
{"type": "error", "error": "Unknown or expired duplex session: sess_demo_01", "code": "session_resume_expired"}
|
||||
{"type": "error", "error": "Unknown or expired duplex session: duplex-5f2a…", "code": "session_resume_expired"}
|
||||
```
|
||||
|
||||
#### Coverage
|
||||
|
||||
The examples above cover every event type the Realtime route emits or
|
||||
accepts (verified by diffing the `type` literals in
|
||||
`vllm_omni/entrypoints/duplex/` against this document). Not shown, because
|
||||
they never reach a `/v1/realtime?duplex=1` client unprojected: the native
|
||||
`/v1/duplex` dialect (`session.create` / `open_session` / `session.config`,
|
||||
`input.committed`, `input.cancelled`, `audio.cancelled`,
|
||||
`response.output_audio.delta`, `response.text.delta`, `response.message`,
|
||||
`function_call.done`) — their Realtime projections are the examples above
|
||||
(see the name map at the end of this section) — and the runner-internal markers `__timeout__`,
|
||||
`__disconnect__`, `__replaced_attachment__`. Input aliases (`push_chunk`,
|
||||
`input.audio.append`, `input_text.append`, `push_text`, `signal_turn`,
|
||||
`audio.playback_ack`, `close_session`, `close`, `session_close`) share the
|
||||
payload of their canonical event (see the alias table in [Full-Duplex Runtime (MiniCPM-o 4.5)](../design/fullduplex.md)).
|
||||
`vllm_omni/engine/duplex/events.py` and `vllm_omni/engine/duplex/commands.py`
|
||||
against this document). Not shown, because they never reach a client
|
||||
unprojected: the session-internal events the runner produces
|
||||
(`input.committed`, `input.cancelled`, `audio.cancelled`,
|
||||
`response.output_audio.delta`, `response.text.delta`, `function_call.done`)
|
||||
— their Realtime projections are the examples above (see the name map at the
|
||||
end of this section). Every public event is a typed `DuplexEvent` whose
|
||||
`to_realtime()` is the wire object; every client event is parsed into a
|
||||
typed `DuplexCommand` by `command_from_realtime()`.
|
||||
|
||||
#### Native to Realtime name map
|
||||
#### Internal to Realtime name map
|
||||
|
||||
The native `/v1/duplex` dialect is what the session runner produces; the
|
||||
Realtime projector renames and fans out these events before they reach a
|
||||
`/v1/realtime?duplex=1` client.
|
||||
The session runner emits these session-internal events; the Realtime
|
||||
projector (`vllm_omni/engine/duplex/realtime_events.py`) renames and fans
|
||||
them out into typed events before they reach a client.
|
||||
|
||||
| Native / internal event | Realtime projection |
|
||||
| Internal event | Realtime projection |
|
||||
| --- | --- |
|
||||
| `session.create` (`open_session`, `session.config`) | first `session.update` |
|
||||
| `input.committed` | `conversation.item.added` / `conversation.item.created`, `input_audio_buffer.committed`, `conversation.item.input_audio_transcription.completed`, `conversation.item.done` |
|
||||
| `input.cancelled` | `input_audio_buffer.cleared` |
|
||||
| `response.output_audio.delta` (with `audio_transcript`) | optional `response.speak`, then `response.output_audio.delta` and `response.output_audio_transcript.delta` |
|
||||
| `response.output_audio.done` / `response.output_text.done` | `response.output_audio.done` / `response.output_text.done` |
|
||||
| `response.text.delta` | `response.output_text.delta` |
|
||||
| `response.message` (chat-fallback raw chunk) | passed through |
|
||||
| `audio.cancelled` | optional `output_audio_buffer.cleared`, then the cancelled terminal events and `response.done` with `status: "cancelled"` |
|
||||
| `function_call.done` | `response.created`, `response.output_item.added` (`function_call` item), `response.function_call_arguments.delta` / `.done`, `response.output_item.done`, `response.done` |
|
||||
| `playback.acknowledged` | `playback.acknowledged` (wrapped verbatim under `event`) |
|
||||
| `runtime.control`, `session.resumed`, `session.replaced`, `session.expired`, `session.resync_required`, `session.heartbeat_ack` | passed through unchanged |
|
||||
| `session.resumed`, `session.replaced`, `session.expired`, `session.resync_required`, `session.heartbeat_ack` | passed through unchanged |
|
||||
|
||||
## Known Limitations
|
||||
|
||||
- Only MiniCPM-o 4.5 is served over this endpoint today; PersonaPlex and
|
||||
Nemotron VoiceChat arrive with the follow-up PRs that port them to the
|
||||
plugin contract.
|
||||
- Several surfaces are capability-gated per model (see *Capability
|
||||
negotiation by model* above): PersonaPlex does not support session resume,
|
||||
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`.
|
||||
- A response is generated from committed **audio**, so the OpenAI text-prompt
|
||||
shape does not drive one: `conversation.item.create` with an `input_text`
|
||||
part adds the item to history, but a following `response.create` with no
|
||||
committed audio is rejected with `response_create_without_input`, and
|
||||
`input.text.append` is rejected with `native_text_append_unsupported`.
|
||||
Text-to-speech has its own shape here — put the text in
|
||||
`extra_body.duplex_initial_user_text` on `session.update`, then stream audio
|
||||
units (silence is enough). The seeded turn is what the model answers.
|
||||
- `turn_detection` supports only `server_vad` with `interrupt_response=true`;
|
||||
`semantic_vad`, `interrupt_response=false`, and `create_response` are not
|
||||
supported.
|
||||
- `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.
|
||||
- `rate_limits.updated` is emitted for compatibility only and always carries
|
||||
an empty list.
|
||||
- Session capacity is bounded by `duplex_session.max_sessions` in the deploy
|
||||
configuration; admission beyond it fails with `duplex_session_busy` or
|
||||
`resource_exhausted`.
|
||||
- The native `/v1/duplex` route speaks the internal dialect and does not
|
||||
support `session.resume`; use `/v1/realtime?duplex=1` for the contract on
|
||||
this page.
|
||||
configuration; admission beyond it fails with `resource_exhausted`
|
||||
(`rate_limit_error`, retryable).
|
||||
- A duplex deployment serves the duplex session surfaces only: the websocket
|
||||
route and `/v1/chat/completions`, the latter on a session per request (see
|
||||
[Start the Server](#start-the-server)). The other turn-based HTTP routes
|
||||
report "not available" on a MiniCPM-o 4.5 server; turn-based use of the
|
||||
model stays available offline through `Omni` / `AsyncOmni`.
|
||||
@@ -6,6 +6,7 @@ vLLM-Omni provides an OpenAI-compatible API for text-to-speech (TTS) generation.
|
||||
- **Fish Speech S2 Pro** (`fishaudio/s2-pro`) -- Dual-AR TTS with DAC codec. Supports text-to-speech and voice cloning via reference audio. Output: 44.1 kHz.
|
||||
- **Voxtral TTS** (`mistralai/Voxtral-4B-TTS-2603`) -- AR + FlowMatching TTS with preset voices. Output: 24 kHz.
|
||||
- **CosyVoice3** (`FunAudioLLM/Fun-CosyVoice3-0.5B-2512`) -- 2-stage talker + flow-matching code2wav. Voice cloning via `ref_audio` + `ref_text` (no presets). Output: 24 kHz.
|
||||
- **Gepard-1.0** (`nineninesix/gepard-1.0`) -- single-stage native-AR TTS with a 22.05 kHz NanoCodec. Zero-shot default voice only.
|
||||
|
||||
See the [Supported Models](#supported-models) section below for the full list, including OmniVoice, VoxCPM2, MOSS-TTS-Nano, and Breeze-TTS-2.
|
||||
|
||||
@@ -40,11 +41,16 @@ vllm serve mistralai/Voxtral-4B-TTS-2603 --omni --port 8091
|
||||
# CosyVoice3 (voice cloning only — supply ref_audio + ref_text per request)
|
||||
vllm serve FunAudioLLM/Fun-CosyVoice3-0.5B-2512 \
|
||||
--omni --port 8091 --trust-remote-code
|
||||
|
||||
# Gepard-1.0 (zero-shot default voice; packaged gepard.yaml)
|
||||
vllm-omni serve nineninesix/gepard-1.0 --omni --port 8091 --trust-remote-code \
|
||||
--stage-init-timeout 900 \
|
||||
--deploy-config vllm_omni/deploy/gepard.yaml
|
||||
```
|
||||
|
||||
### Generate Speech
|
||||
|
||||
**Using curl:**
|
||||
**Qwen3-TTS CustomVoice, using curl:**
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
@@ -56,7 +62,24 @@ curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
}' --output output.wav
|
||||
```
|
||||
|
||||
**Using Python:**
|
||||
**Gepard-1.0, using curl:**
|
||||
|
||||
Gepard accepts `input` (required), `voice` (`"default"`), `response_format`,
|
||||
`stream` / `stream_format`, `max_new_tokens`, and `seed`. Unsupported fields
|
||||
include `speed`, `language`, `instructions`, `task_type`, `ref_audio`,
|
||||
`ref_text`, `extra_params`, and `word_timestamps`.
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "Hello, this is Gepard speaking.",
|
||||
"voice": "default",
|
||||
"seed": 7
|
||||
}' --output output.wav
|
||||
```
|
||||
|
||||
**Qwen3-TTS CustomVoice, using Python:**
|
||||
|
||||
```python
|
||||
import httpx
|
||||
@@ -761,6 +784,28 @@ Fish Speech uses `ref_audio` and `ref_text` for voice cloning (no `task_type` ne
|
||||
| ------- | ------------- |
|
||||
| `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | Voice cloning from `ref_audio` + `ref_text`. No built-in voice presets — upload a voice or pass `ref_audio`/`ref_text` per request. |
|
||||
|
||||
### Gepard-1.0
|
||||
|
||||
| Model | Description |
|
||||
| ----- | ----------- |
|
||||
| `nineninesix/gepard-1.0` | Zero-shot native-AR TTS. 22.05 kHz mono. `voice` must be omitted or `"default"`. |
|
||||
|
||||
Gepard request fields:
|
||||
|
||||
| Field | Behavior |
|
||||
| ----- | -------- |
|
||||
| `input` | Required. Empty/whitespace-only returns 400. |
|
||||
| `voice` | Omitted or `"default"` only. Other values 400. |
|
||||
| `response_format` | `wav` default. Non-streaming: `wav`/`pcm`/`flac`/`mp3`. `opus` returns 400 (22.05 kHz is not an Opus sample rate). Streaming: `pcm`/`wav` only. |
|
||||
| `stream` / `stream_format` | SSE (`speech.audio.*`) and raw `audio` byte streaming. |
|
||||
| `speed` | Must be `1.0`. Gepard has no native speed control. |
|
||||
| `max_new_tokens` | Frame budget (1 token = 1 frame = 1024 samples ≈ 46.4 ms). Default 1000 from deploy YAML; adapter bounds 1..4096. |
|
||||
| `seed` | Optional. Reaches the in-model 32-head sampler. The packaged YAML currently pins `seed: 42`, so serving is deterministic by default until that pin is removed. |
|
||||
| `extra_params` | Any key returns 400, including `temperature`/`top_p`/`top_k`. |
|
||||
| Cloning / style fields | `ref_audio`, `ref_text`, `speaker_embedding`, `instructions`, `language`, `task_type`, `word_timestamps`, and similar declared-but-unsupported fields return 400. |
|
||||
|
||||
See the [Gepard section of the online TTS hub](../user_guide/examples/online_serving/text_to_speech.md#gepard-10) for launch commands. Native-AR recompute preemption is a known limitation under concurrency.
|
||||
|
||||
### OmniVoice
|
||||
|
||||
| Model | Description |
|
||||
|
||||
@@ -50,6 +50,14 @@ WebSocket /v1/video/chat/stream
|
||||
| Server -> Client | `session.done` | none | Session closed. |
|
||||
| Server -> Client | `error` | `message` | Recoverable protocol or generation error. |
|
||||
|
||||
### Frame Consumption Reporting
|
||||
|
||||
When buffered frames include client-supplied `frame_id` values, the server emits `video.frames.consumed` after the first engine output. Its `frame_ids`, `frames`, and `latest_pts_ms` describe the images included in that query's prompt.
|
||||
|
||||
Queries stride-sample the buffered frames, keeping the last frame, then exclude frames already known to have failed decoding. Excluded frames are not replaced with other buffered frames. Frames whose background decoding has not finished remain eligible through the image URL path.
|
||||
|
||||
**Bugfix compatibility note:** `video.frames.consumed` now excludes known decode failures that were already excluded from the prompt. Older versions could report those frames and their timestamps as consumed. The event name and fields are unchanged; an empty selection reports empty lists and `latest_pts_ms: null`.
|
||||
|
||||
### `session.config` Fields
|
||||
|
||||
| Field | Type | Default | Description |
|
||||
|
||||
@@ -109,6 +109,16 @@ curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o output.mp4
|
||||
The final content is available from `/v1/videos/{video_id}/content` after the
|
||||
job status becomes `completed`.
|
||||
|
||||
`queued` means the request is still waiting for diffusion scheduler admission.
|
||||
The status changes to `in_progress` when the scheduler first selects the
|
||||
request for execution.
|
||||
|
||||
`DELETE /v1/videos/{video_id}` issues a bounded engine abort
|
||||
(`VLLM_OMNI_ABORT_TIMEOUT`, default 2s), then cancels the frontend
|
||||
task. Cancellation cleanup is also bounded and best-effort: it confirms
|
||||
the abort was queued, and the current request batch may still drain.
|
||||
The job is then re-read so a completed save is not orphaned.
|
||||
|
||||
### Synchronous Response
|
||||
|
||||
`POST /v1/videos/sync` blocks until generation finishes and returns raw video
|
||||
|
||||
@@ -1,24 +1,33 @@
|
||||
# FastVideo VSA
|
||||
|
||||
FastVideo Variable Sparse Attention (VSA) accelerates the `FastVideo/FastWan2.2-TI2V-5B-Diffusers` model's
|
||||
self-attention by partitioning the post-patch latent grid into spatiotemporal
|
||||
FastVideo Variable Sparse Attention (VSA) partitions the post-patch latent grid into spatiotemporal
|
||||
blocks. For every query block, VSA scores the key/value blocks and computes
|
||||
attention only against the selected top-k blocks.
|
||||
|
||||
The supported checkpoint provides both text-to-video and image-to-video modes through `Wan22Pipeline`; the separate Wan I2V-14B, S2V, and VACE pipelines are outside this backend's supported scope.
|
||||
## Supported models
|
||||
|
||||
MiniMax-H3 served with a FastH3 VSA adapter uses a second, H3-specific route
|
||||
through the same backend. See the
|
||||
[MiniMax-H3 recipe](https://recipes.vllm.ai/MiniMaxAI/MiniMax-H3) for its geometry,
|
||||
supported topologies, and commands.
|
||||
| Model / checkpoint | Required adapter | Tasks | Sequence parallelism |
|
||||
| --- | --- | --- | --- |
|
||||
| `FastVideo/FastWan2.2-TI2V-5B-Diffusers` | None | T2V, I2V through `Wan22Pipeline` | Disabled |
|
||||
| `MiniMaxAI/MiniMax-H3` | [FastH3 VSA (recipe)](https://github.com/vllm-project/vllm-omni/blob/main/recipes/MiniMaxAI/MiniMax-H3.md#fasth3-adapter) | T2VA | Disabled or pure Ulysses |
|
||||
| Wan I2V-14B, S2V, VACE | — | Unsupported | — |
|
||||
|
||||
VSA is a CUDA-only, explicitly selected backend. It requires the
|
||||
`fastvideo-kernel` package and currently supports non-causal self-attention
|
||||
with equal query and key/value sequence lengths. Unsupported shapes, masks,
|
||||
dtypes, sequence-parallel execution, or kernel failures fall back to
|
||||
`TORCH_SDPA` and emit a warning with the reason. On the Wan route, an active
|
||||
sequence-parallel context is one of those fallbacks; the H3 route supports pure
|
||||
Ulysses and rejects ring or all-gather sequence parallelism at startup.
|
||||
## Installation
|
||||
|
||||
Install vLLM-Omni from this checkout with the `vsa` extra:
|
||||
|
||||
```bash
|
||||
uv pip install -e '.[vsa]'
|
||||
```
|
||||
|
||||
The extra installs the tested kernel dependency automatically. Prebuilt kernels require
|
||||
Linux, Python 3.12, and glibc 2.34 or newer (x86-64 or aarch64). The full
|
||||
FastVideo framework and provider environment variables are not required.
|
||||
|
||||
`FASTVIDEO_VSA` selects the attention algorithm. On SM120, the H3 integration
|
||||
currently executes FastVideo's 64-token Triton block-sparse kernel; it does not
|
||||
dispatch to FlashInfer. The `vsa` extra installs `fastvideo-kernel==0.3.4` for
|
||||
this execution path.
|
||||
|
||||
## Enable the backend
|
||||
|
||||
@@ -27,24 +36,18 @@ flag. Use `--fastvideo-vsa-topk` to set the number of key/value blocks retained
|
||||
for every query block:
|
||||
|
||||
```bash
|
||||
vllm-omni serve <model> \
|
||||
vllm serve <model> --omni \
|
||||
--diffusion-attention-backend FASTVIDEO_VSA \
|
||||
--fastvideo-vsa-topk 64
|
||||
```
|
||||
|
||||
The backwards-compatible environment variable selects the backend with the
|
||||
default `topk=64`:
|
||||
<details markdown="1">
|
||||
<summary>Alternative configuration formats</summary>
|
||||
|
||||
Equivalent structured configuration:
|
||||
|
||||
```bash
|
||||
export DIFFUSION_ATTENTION_BACKEND=FASTVIDEO_VSA
|
||||
vllm-omni serve <model>
|
||||
```
|
||||
|
||||
To tune top-k, pass the CLI backend and top-k flags together as shown above,
|
||||
or use the equivalent structured configuration:
|
||||
|
||||
```bash
|
||||
vllm-omni serve <model> \
|
||||
vllm serve <model> --omni \
|
||||
--diffusion-attention-config \
|
||||
'{"default":{"backend":"FASTVIDEO_VSA","fastvideo_vsa_topk":64}}'
|
||||
```
|
||||
@@ -73,8 +76,15 @@ stages:
|
||||
fastvideo_vsa_topk: 64
|
||||
```
|
||||
|
||||
</details>
|
||||
|
||||
## Choose top-k
|
||||
|
||||
H3 uses 64-token video blocks and keeps
|
||||
all prefix blocks; see the [FastH3 VSA recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/MiniMaxAI/MiniMax-H3.md#fasth3-adapter).
|
||||
|
||||
### Wan top-k behavior
|
||||
|
||||
At runtime the backend logs the sequence shape and derived block count:
|
||||
|
||||
```text
|
||||
@@ -117,8 +127,16 @@ selecting VSA does not turn a native checkpoint into a distilled model.
|
||||
|
||||
## Verify routing and fallback
|
||||
|
||||
Check the startup and first-forward logs instead of assuming that selecting
|
||||
the backend guarantees sparse execution:
|
||||
VSA is a CUDA-only, explicitly selected backend. It requires the
|
||||
`fastvideo-kernel` package and currently supports non-causal self-attention
|
||||
with equal query and key/value sequence lengths. Unsupported shapes, masks,
|
||||
dtypes, sequence-parallel execution, or kernel failures fall back to
|
||||
`TORCH_SDPA` and emit a warning with the reason. On the Wan route, an active
|
||||
sequence-parallel context is one of those fallbacks; the H3 route supports pure
|
||||
Ulysses and rejects ring or all-gather sequence parallelism at startup.
|
||||
H3 accelerator faults propagate instead of attempting dense recovery.
|
||||
|
||||
On the Wan route, check the startup and first-forward logs:
|
||||
|
||||
- `route=VSA` means top-k block selection is active.
|
||||
- `route=VSA_ALL_BLOCKS` means the FastVideo DMD checkpoint retained all
|
||||
@@ -126,6 +144,10 @@ the backend guarantees sparse execution:
|
||||
- `route=SDPA` or `FASTVIDEO_VSA falling back to SDPA: ...` means dense SDPA
|
||||
executed; the warning includes the reason.
|
||||
|
||||
For H3, check `FastH3 adapter active` at startup and
|
||||
`FASTVIDEO_VSA H3 routing` during DiT execution, as described in the
|
||||
[model recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/MiniMaxAI/MiniMax-H3.md#fasth3-adapter).
|
||||
|
||||
The Wan route requires CUDA tensors in FP16 or BF16, 256-token blocks,
|
||||
standard `head_size**-0.5` scaling, equal Q/K/V head counts, no attention mask,
|
||||
and no active sequence-parallel context. The MiniMax-H3 route uses 64-token
|
||||
|
||||
@@ -18,6 +18,7 @@ For the full list of supported architectures across all modalities, see
|
||||
| --- | --- | --- | --- | --- | --- |
|
||||
| CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | — |
|
||||
| Fish Speech S2 Pro | `fishaudio/s2-pro` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | ✓ |
|
||||
| Gepard-1.0 | `nineninesix/gepard-1.0` | — (zero-shot default voice) | ✓ (PCM / WAV stream) | `default` only | — |
|
||||
| higgs-audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` | ✓ (`ref_audio`+`ref_text`) | ✓ (codec_streaming) | — | — |
|
||||
| GLM-TTS | `zai-org/GLM-TTS` | ✓ (`ref_audio`+`ref_text`, required) | ✓ (PCM stream) | — | ✓ |
|
||||
| OmniVoice | `k2-fsa/OmniVoice` | (offline only) | — | — | — |
|
||||
@@ -92,7 +93,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
}' --no-buffer | play -t raw -r 24000 -e signed -b 16 -c 1 -
|
||||
```
|
||||
|
||||
Adjust the player's sample rate to match the model (44.1 kHz for Fish Speech, 48 kHz for VoxCPM2, 24 kHz for the others).
|
||||
Adjust the player's sample rate to match the model (44.1 kHz for Fish Speech, 48 kHz for VoxCPM2, 22.05 kHz for Gepard, 24 kHz for the others).
|
||||
|
||||
For full request-shape documentation (all parameters, response formats, error codes), see the [Speech API reference](https://github.com/vllm-project/vllm-omni/tree/main/docs/serving/speech_api.md).
|
||||
|
||||
@@ -212,6 +213,61 @@ python fish_speech/gradio_demo.py --api-base http://localhost:8091 # if server
|
||||
|
||||
---
|
||||
|
||||
## Gepard-1.0
|
||||
|
||||
Single-stage native AR TTS at 22.05 kHz mono. Zero-shot only: omit `voice` or pass `"default"`. Voice cloning from reference audio is not available yet.
|
||||
|
||||
### Prerequisites
|
||||
|
||||
Same NeMo NanoCodec install as the [offline Gepard section](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_speech/README.md#gepard-10). On a host whose CUDA toolkit cannot build kernels — no `nvcc`/`ninja`, or a consumer Blackwell (`sm_120`) card — also `export VLLM_USE_FLASHINFER_SAMPLER=0` before launch. That env is not a deploy-YAML field.
|
||||
|
||||
### Launch
|
||||
|
||||
```bash
|
||||
vllm-omni serve nineninesix/gepard-1.0 --omni --port 8091 --trust-remote-code \
|
||||
--stage-init-timeout 900 \
|
||||
--deploy-config vllm_omni/deploy/gepard.yaml
|
||||
# or:
|
||||
./gepard/run_server.sh
|
||||
```
|
||||
|
||||
`--stage-init-timeout 900` matches the online e2e fixture and the offline example. Serve defaults to 300s, which is often too short for a cold download of the talker plus NanoCodec.
|
||||
|
||||
The packaged `vllm_omni/deploy/gepard.yaml` must be passed with `--deploy-config`. The checkpoint self-identifies as `qwen3_5_text`, so omitting the YAML launches a diffusion fallback instead of the Gepard pipeline. The YAML sets `async_chunk: false`, `max_num_seqs: 4`, and currently pins `seed: 42`, so serving is deterministic by default until that YAML seed is removed. Pass an explicit per-request `seed` in tests and clients rather than depending on either default.
|
||||
|
||||
### Sending requests
|
||||
|
||||
```bash
|
||||
python examples/online_serving/text_to_speech/gepard/speech_client.py \
|
||||
--text "Hello, this is Gepard speaking."
|
||||
|
||||
python examples/online_serving/text_to_speech/gepard/speech_client.py \
|
||||
--text "Hello, this is Gepard speaking." --seed 7 --stream --output output.pcm
|
||||
```
|
||||
|
||||
Gepard accepts `input` (required), `voice` (`"default"`), `response_format`, `stream` / `stream_format`, `max_new_tokens`, and `seed`. Unsupported fields include `speed`, `language`, `instructions`, `task_type`, `ref_audio`, `ref_text`, `extra_params`, and `word_timestamps`.
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"input": "Hello, this is Gepard speaking.",
|
||||
"voice": "default",
|
||||
"seed": 7
|
||||
}' --output output.wav
|
||||
```
|
||||
|
||||
### Notes
|
||||
|
||||
- Output: 22.05 kHz mono. `max_new_tokens` is a **frame** budget (1 token = 1 frame = 1024 samples ≈ 46.4 ms at 21.5 fps; adapter bounds 1..4096).
|
||||
- Supported request fields: `input` (required), `voice` (`default` only), `response_format` (`wav` default; `wav/pcm/flac/mp3` non-streaming; `opus` 400 because 22.05 kHz is not an Opus sample rate; streaming `pcm/wav` only), `stream` / `stream_format`, `max_new_tokens`, `seed`.
|
||||
- Unsupported: `speed`, `extra_params` (including `temperature`/`top_p`/`top_k`), `ref_audio`, `ref_text`, `speaker_embedding`, `task_type`, `instructions`, `language`, and `word_timestamps`.
|
||||
- Concurrent requests at `max_num_seqs: 4` are supported. Native-AR recompute preemption is a known limitation of this architecture (a request that is preempted mid-generation can resume incorrectly); keep concurrency at or below `max_num_seqs` and treat preemption as out of scope until the platform fix lands.
|
||||
- Optional comparison against the upstream Gepard reference server needs Blackwell/Hopper + CUDA 13 + Postgres and is not part of CI.
|
||||
- For the request-field table, see the [Speech API Gepard section](https://github.com/vllm-project/vllm-omni/blob/main/docs/serving/speech_api.md#gepard-10).
|
||||
|
||||
---
|
||||
|
||||
## higgs-audio v2
|
||||
|
||||
2-stage TTS at 24 kHz: a vLLM-native Llama-3.2-3B talker with a DualFFN audio expert (Stage 0) feeding a HiggsAudio codec decoder (Stage 1) that streams chunks back to the client.
|
||||
|
||||
@@ -19,7 +19,7 @@ list of supported architectures across all modalities, see
|
||||
| Breeze-TTS-2 | `BreezeBlue/Breeze-TTS-2` | 2 (talker + codec) | ✓ | ✓ (async chunk) | voice design / voice direction (`--instruction`) | 24 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 |
|
||||
| Gepard-1.0 | `nineninesix/gepard-1.0` | single (native AR + NanoCodec) | — (zero-shot; cloning later) | ✓ (online) | zero-shot | 22.05 kHz |
|
||||
| GLM-TTS | `zai-org/GLM-TTS` | 2 (AR + DiT) | ✓ (required) | ✓ | — | 24 kHz |
|
||||
| Ming-omni-tts | `inclusionAI/Ming-omni-tts-0.5B` | 2 (AR + audio VAE) | ✓ | ✓ | style / IP / dialect / TTA / podcast | 44.1 kHz |
|
||||
| Ming-flash-omni-TTS | `Jonathan1909/Ming-flash-omni-2.0` | single (talker only) | — (caption-controlled) | — | style / IP / basic captions | 44.1 kHz |
|
||||
@@ -167,7 +167,7 @@ python examples/offline_inference/text_to_speech/glm_tts/end2end.py \
|
||||
```
|
||||
|
||||
### Architecture
|
||||
```
|
||||
```text
|
||||
Text → [Stage 0: AR] → Speech Tokens → [Stage 1: DiT + HiFT] → Audio (24 kHz)
|
||||
(Llama-based) (32k vocab) (Flow Matching)
|
||||
```
|
||||
@@ -491,8 +491,8 @@ Not yet — PR1 is zero-shot only (the learned `null_prefix` default voice). Ref
|
||||
### Notes
|
||||
- Output: 22.05 kHz mono WAV.
|
||||
- First run downloads two checkpoints, not one: the model itself, and the NanoCodec decoder it names in `codec_id` (`nvidia/nemo-nano-codec-22khz-1.89kbps-21.5fps` by default). Both are fetched automatically; expect the first startup to be correspondingly slower.
|
||||
- Offline only for now; the `/v1/audio/speech` adapter is a follow-up PR.
|
||||
- Deploy config: `vllm_omni/deploy/gepard.yaml` (the example's default; copy it and pass `--deploy-config` to change it).
|
||||
- Offline and online: `examples/offline_inference/text_to_speech/gepard/end2end.py` and `/v1/audio/speech` (see the [online Gepard section](../../online_serving/text_to_speech/README.md#gepard-10)).
|
||||
- Deploy config: `vllm_omni/deploy/gepard.yaml` (the example's default; copy it and pass `--deploy-config` to change it). Online serving **must** pass `--deploy-config vllm_omni/deploy/gepard.yaml` because the checkpoint self-identifies as `qwen3_5_text`.
|
||||
- Generation length and reproducibility are stage settings, not script flags. One output token is one audio frame, so `max_tokens` in the YAML is the frame budget; `seed` makes the in-model 32-head sampling reproducible. The script deliberately passes no `SamplingParams`: one supplied by a caller replaces the stage defaults wholesale rather than merging, which would drop the pipeline's stop token and run every request to `max_tokens`.
|
||||
- Text length: short texts are repeated internally to match the training layout (the checkpoint's `text_repetition` block); the upper bound is the stage's `max_model_len`, enforced by the engine. Empty text is rejected rather than voiced.
|
||||
- `VLLM_GEPARD_GREEDY=1` swaps the 32-head Gumbel-max sampling for argmax, for reproducible comparisons that do not depend on a seed.
|
||||
|
||||
@@ -14,7 +14,10 @@ Scenario (see ``barge_in_client_flow.md`` for the diagram)
|
||||
|
||||
The demo is model-neutral: it drives any duplex model through the public
|
||||
:class:`vllm_omni.clients.duplex.DuplexClient` over
|
||||
``/v1/realtime?duplex=1``. Model-specific session shape comes from the
|
||||
``/v1/realtime?duplex=1``, or — with ``--inline`` — through
|
||||
:class:`vllm_omni.clients.inline_duplex.InlineDuplexClient` on an in-process
|
||||
:class:`vllm_omni.entrypoints.duplex_omni.DuplexOmni` (no server needed; the
|
||||
model loads in this process). Model-specific session shape comes from the
|
||||
per-model presets under ``vllm_omni.clients.<model>`` selected with
|
||||
``--preset`` (default ``minicpmo_4_5``, whose model decides listen/speak on
|
||||
its own — no VAD).
|
||||
@@ -31,6 +34,15 @@ then:
|
||||
--interrupt-wav follow_up_16k.wav \
|
||||
--output-dir ./duplex_out
|
||||
|
||||
or, without a server, in-process (``--deploy-config`` is optional and defaults
|
||||
to the model's bundled duplex profile):
|
||||
|
||||
python barge_in_client.py --inline \
|
||||
--model /path/to/MiniCPM-o-4_5 \
|
||||
--ref-audio /path/to/reference_voice.wav \
|
||||
--question-wav question_16k.wav \
|
||||
--interrupt-wav follow_up_16k.wav
|
||||
|
||||
Input WAVs must be mono 16 kHz PCM16; when a preset's session takes a
|
||||
different capture format (the PersonaPlex preset streams 24 kHz float32),
|
||||
the audio is converted before streaming. ``--ref-audio`` is required by the
|
||||
@@ -55,6 +67,7 @@ if str(REPO_ROOT) not in sys.path:
|
||||
from vllm_omni.clients.duplex import ( # noqa: E402
|
||||
AudioFormat,
|
||||
DuplexClient,
|
||||
DuplexClientBase,
|
||||
EventCollector,
|
||||
SessionConfig,
|
||||
acknowledge_collected_playback,
|
||||
@@ -139,6 +152,39 @@ def _fold_responses(collector: EventCollector) -> dict[str, dict[str, object]]:
|
||||
return out
|
||||
|
||||
|
||||
def _make_client(args: argparse.Namespace, config: SessionConfig) -> DuplexClientBase:
|
||||
"""Websocket client by default; ``--inline`` drives an in-process DuplexOmni."""
|
||||
if not args.inline:
|
||||
return DuplexClient(
|
||||
args.url,
|
||||
model=args.model,
|
||||
config=config,
|
||||
reconnect=None,
|
||||
heartbeat_interval_s=None,
|
||||
handshake_timeout_s=args.timeout_s,
|
||||
)
|
||||
# Imported lazily: the in-process path loads the model (and its engine
|
||||
# dependencies) in this process, which the websocket path never needs.
|
||||
from vllm_omni.clients.inline_duplex import InlineDuplexClient
|
||||
from vllm_omni.entrypoints.duplex_omni import DuplexOmni
|
||||
|
||||
# Models shipped with custom HF code (MiniCPM-o) need trust_remote_code to
|
||||
# resolve their pipeline; harmless for the others.
|
||||
omni_kwargs: dict[str, object] = {"model": args.model, "trust_remote_code": True}
|
||||
if args.deploy_config:
|
||||
omni_kwargs["deploy_config"] = args.deploy_config
|
||||
omni = DuplexOmni(**omni_kwargs)
|
||||
client = InlineDuplexClient(
|
||||
omni,
|
||||
model=args.model,
|
||||
config=config,
|
||||
handshake_timeout_s=args.timeout_s,
|
||||
)
|
||||
# The client owns no engine; stop the in-process one when the run ends.
|
||||
client.owned_omni = omni # type: ignore[attr-defined]
|
||||
return client
|
||||
|
||||
|
||||
async def run(args: argparse.Namespace) -> int:
|
||||
config = _session_config(args.preset, ref_audio_path=args.ref_audio)
|
||||
question = _convert_to_session_format(read_pcm16_wav(Path(args.question_wav)), config.input_audio)
|
||||
@@ -148,105 +194,117 @@ async def run(args: argparse.Namespace) -> int:
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
collector = EventCollector()
|
||||
client = DuplexClient(
|
||||
args.url,
|
||||
model=args.model,
|
||||
config=config,
|
||||
reconnect=None,
|
||||
heartbeat_interval_s=None,
|
||||
handshake_timeout_s=args.timeout_s,
|
||||
)
|
||||
client = _make_client(args, config)
|
||||
# GREET: open the voice session (handshake happens on enter).
|
||||
async with client:
|
||||
consume_task = asyncio.create_task(collector.consume(client))
|
||||
try:
|
||||
async with client:
|
||||
consume_task = asyncio.create_task(collector.consume(client))
|
||||
|
||||
# ASK: stream the question at realtime pace, then commit the turn.
|
||||
# The model decides on its own to answer (no VAD).
|
||||
print("[1] streaming question ...", file=sys.stderr)
|
||||
await client.stream_pcm(question, chunk_ms=chunk_ms)
|
||||
await client.commit(create_response=False)
|
||||
await wait_for_condition(
|
||||
lambda: collector.count("response.created") > 0,
|
||||
timeout_s=args.timeout_s,
|
||||
label="first response.created",
|
||||
)
|
||||
print("[1] assistant is answering", file=sys.stderr)
|
||||
# ASK: stream the question at realtime pace, then commit the turn.
|
||||
# The model decides on its own to answer (no VAD).
|
||||
print("[1] streaming question ...", file=sys.stderr)
|
||||
await client.stream_pcm(question, chunk_ms=chunk_ms)
|
||||
await client.commit(create_response=False)
|
||||
await wait_for_condition(
|
||||
lambda: collector.count("response.created") > 0,
|
||||
timeout_s=args.timeout_s,
|
||||
label="first response.created",
|
||||
)
|
||||
print("[1] assistant is answering", file=sys.stderr)
|
||||
|
||||
# Let the answer play for a moment, as a listening user would.
|
||||
await asyncio.sleep(args.listen_s)
|
||||
# Let the answer play for a moment, as a listening user would.
|
||||
await asyncio.sleep(args.listen_s)
|
||||
|
||||
# INTERRUPT: talk over the assistant while its audio still streams.
|
||||
# The serving overlap policy decides: enough overlapped speech barges
|
||||
# in (the first response is cancelled) or a short remark is deferred
|
||||
# until the model yields the turn.
|
||||
print("[2] interrupting mid-answer ...", file=sys.stderr)
|
||||
ids_before_interrupt = set(collector.response_ids)
|
||||
await client.stream_pcm(interrupt, chunk_ms=chunk_ms)
|
||||
await client.commit(create_response=False)
|
||||
# INTERRUPT: talk over the assistant while its audio still streams.
|
||||
# The serving overlap policy decides: enough overlapped speech barges
|
||||
# in (the first response is cancelled) or a short remark is deferred
|
||||
# until the model yields the turn.
|
||||
print("[2] interrupting mid-answer ...", file=sys.stderr)
|
||||
ids_before_interrupt = set(collector.response_ids)
|
||||
await client.stream_pcm(interrupt, chunk_ms=chunk_ms)
|
||||
await client.commit(create_response=False)
|
||||
|
||||
# ANSWER AGAIN: wait for a response created *after* the interruption
|
||||
# to finish. The first answer may still be streaming here; matching on
|
||||
# any response.done would let its completion satisfy the wait and cut
|
||||
# off the actual follow-up answer at close.
|
||||
def follow_up_answer_done() -> bool:
|
||||
return any(
|
||||
event.get("type") == "response.done"
|
||||
and collector.response_id(event) is not None
|
||||
and collector.response_id(event) not in ids_before_interrupt
|
||||
for event in collector.events
|
||||
# ANSWER AGAIN: wait for a response created *after* the interruption
|
||||
# to finish. The first answer may still be streaming here; matching on
|
||||
# any response.done would let its completion satisfy the wait and cut
|
||||
# off the actual follow-up answer at close.
|
||||
def follow_up_answer_done() -> bool:
|
||||
return any(
|
||||
event.get("type") == "response.done"
|
||||
and collector.response_id(event) is not None
|
||||
and collector.response_id(event) not in ids_before_interrupt
|
||||
for event in collector.events
|
||||
)
|
||||
|
||||
await wait_for_condition(
|
||||
follow_up_answer_done,
|
||||
timeout_s=args.timeout_s,
|
||||
label="response.done for the response answering the interruption",
|
||||
)
|
||||
|
||||
await wait_for_condition(
|
||||
follow_up_answer_done,
|
||||
timeout_s=args.timeout_s,
|
||||
label="response.done for the response answering the interruption",
|
||||
)
|
||||
# HANG UP: report playback, close the session cleanly.
|
||||
await acknowledge_collected_playback(client, collector)
|
||||
await client.close(timeout_s=args.timeout_s)
|
||||
try:
|
||||
await asyncio.wait_for(consume_task, timeout=5.0)
|
||||
except (asyncio.TimeoutError, asyncio.CancelledError):
|
||||
consume_task.cancel()
|
||||
|
||||
# HANG UP: report playback, close the session cleanly.
|
||||
await acknowledge_collected_playback(client, collector)
|
||||
await client.close(timeout_s=args.timeout_s)
|
||||
try:
|
||||
await asyncio.wait_for(consume_task, timeout=5.0)
|
||||
except (asyncio.TimeoutError, asyncio.CancelledError):
|
||||
consume_task.cancel()
|
||||
|
||||
# Report: one WAV per response; statuses show what the overlap did.
|
||||
responses = _fold_responses(collector)
|
||||
summary: list[dict[str, object]] = []
|
||||
for index, (response_id, item) in enumerate(responses.items(), start=1):
|
||||
audio = item["audio"]
|
||||
assert isinstance(audio, bytes)
|
||||
sample_rate_hz = int(item["sample_rate_hz"]) # type: ignore[arg-type]
|
||||
wav_path = output_dir / f"response_{index}_{item['status']}.wav"
|
||||
if audio:
|
||||
write_pcm16_wav(wav_path, audio, sample_rate_hz=sample_rate_hz)
|
||||
summary.append(
|
||||
{
|
||||
"response_id": response_id,
|
||||
"status": item["status"],
|
||||
"audio_s": round(len(audio) / (sample_rate_hz * 2), 2),
|
||||
"text": item["text"],
|
||||
"wav": wav_path.name if audio else None,
|
||||
}
|
||||
# Report: one WAV per response; statuses show what the overlap did.
|
||||
responses = _fold_responses(collector)
|
||||
summary: list[dict[str, object]] = []
|
||||
for index, (response_id, item) in enumerate(responses.items(), start=1):
|
||||
audio = item["audio"]
|
||||
assert isinstance(audio, bytes)
|
||||
sample_rate_hz = int(item["sample_rate_hz"]) # type: ignore[arg-type]
|
||||
wav_path = output_dir / f"response_{index}_{item['status']}.wav"
|
||||
if audio:
|
||||
write_pcm16_wav(wav_path, audio, sample_rate_hz=sample_rate_hz)
|
||||
summary.append(
|
||||
{
|
||||
"response_id": response_id,
|
||||
"status": item["status"],
|
||||
"audio_s": round(len(audio) / (sample_rate_hz * 2), 2),
|
||||
"text": item["text"],
|
||||
"wav": wav_path.name if audio else None,
|
||||
}
|
||||
)
|
||||
(output_dir / "summary.json").write_text(json.dumps(summary, indent=2, ensure_ascii=False))
|
||||
(output_dir / "events.jsonl").write_text(
|
||||
"".join(json.dumps(event, ensure_ascii=False) + "\n" for event in collector.events),
|
||||
encoding="utf-8",
|
||||
)
|
||||
(output_dir / "summary.json").write_text(json.dumps(summary, indent=2, ensure_ascii=False))
|
||||
(output_dir / "events.jsonl").write_text(
|
||||
"".join(json.dumps(event, ensure_ascii=False) + "\n" for event in collector.events),
|
||||
encoding="utf-8",
|
||||
)
|
||||
print(json.dumps(summary, indent=2, ensure_ascii=False))
|
||||
cancelled = sum(1 for item in summary if item["status"] == "cancelled")
|
||||
print(
|
||||
f"\n{len(summary)} responses ({cancelled} interrupted); outputs in {output_dir}",
|
||||
file=sys.stderr,
|
||||
)
|
||||
print(json.dumps(summary, indent=2, ensure_ascii=False))
|
||||
cancelled = sum(1 for item in summary if item["status"] == "cancelled")
|
||||
print(
|
||||
f"\n{len(summary)} responses ({cancelled} interrupted); outputs in {output_dir}",
|
||||
file=sys.stderr,
|
||||
)
|
||||
finally:
|
||||
owned_omni = getattr(client, "owned_omni", None)
|
||||
if owned_omni is not None:
|
||||
owned_omni.shutdown()
|
||||
return 0
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
parser.add_argument("--url", default="ws://127.0.0.1:8099/v1/realtime")
|
||||
parser.add_argument("--model", default="openbmb/MiniCPM-o-4_5")
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
default="openbmb/MiniCPM-o-4_5",
|
||||
help="served model name, or the model path with --inline",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--inline",
|
||||
action="store_true",
|
||||
help="run the model in this process through DuplexOmni + InlineDuplexClient instead of connecting to --url",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--deploy-config",
|
||||
default=None,
|
||||
help="deploy YAML for --inline (default: the model's bundled profile)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--preset",
|
||||
choices=["minicpmo_4_5", "personaplex", "none"],
|
||||
|
||||
@@ -1,8 +1,10 @@
|
||||
# `barge_in_client.py` — what happens, at a glance
|
||||
|
||||
The demo plays one everyday conversation pattern against the full-duplex
|
||||
server: ask something, get talked-over mid-answer, and see the assistant
|
||||
cope. High-level flow below — for the exact wire events, read the
|
||||
The demo plays one everyday conversation pattern against a full-duplex
|
||||
model: ask something, get talked-over mid-answer, and see the assistant
|
||||
cope. It runs either against a server over `/v1/realtime?duplex=1` or, with
|
||||
`--inline`, against an in-process `DuplexOmni` — the conversation below is the
|
||||
same in both modes. High-level flow — for the exact wire events, read the
|
||||
docstring in `barge_in_client.py`.
|
||||
|
||||
```mermaid
|
||||
@@ -48,7 +50,8 @@ answer is saved as `response_1_cancelled.wav`, a completed one as
|
||||
duration per answer.
|
||||
|
||||
Companion: the runtime architecture lives in
|
||||
[`docs/design/fullduplex.md`](../../docs/design/fullduplex.md); the client
|
||||
API is `vllm_omni.clients.duplex.DuplexClient`, documented with the wire
|
||||
protocol in
|
||||
[`docs/design/fullduplex.md`](../../docs/design/fullduplex.md); the client API is
|
||||
`vllm_omni.clients.duplex.DuplexClient` over the websocket and
|
||||
`vllm_omni.clients.inline_duplex.InlineDuplexClient` in-process — both are
|
||||
documented with the wire protocol in
|
||||
[`docs/serving/realtime_duplex_api.md`](../../docs/serving/realtime_duplex_api.md).
|
||||
@@ -5,10 +5,10 @@
|
||||
# Server side for barge_in_client.py: start MiniCPM-o 4.5 in native
|
||||
# full-duplex mode.
|
||||
#
|
||||
# The duplex deploy profile registers the WS /v1/duplex endpoint (with
|
||||
# /v1/realtime?duplex=1 projecting onto the same handler), enables the
|
||||
# engine duplex control plane, and bounds the Thinker/Talker/Code2Wav
|
||||
# stages to four live duplex sessions on one GPU. See
|
||||
# MiniCPM-o 4.5 declares a duplex plugin, so vllm-omni serve runs it through
|
||||
# DuplexOmni: the server is duplex-only (WS /v1/realtime?duplex=1, alias
|
||||
# /v1/duplex, plus /v1/models and /health) and bounds the
|
||||
# Thinker/Talker/Code2Wav stages to four live duplex sessions on one GPU. See
|
||||
# vllm_omni/deploy/minicpmo_4_5.yaml (session_mode: duplex) for the session
|
||||
# limits; every other duplex_session knob keeps its runtime default
|
||||
# (idle TTL 300s, disconnect grace 30s, 16MiB pending input per session).
|
||||
@@ -22,4 +22,4 @@ exec vllm-omni serve "$MODEL" \
|
||||
--deploy-config vllm_omni/deploy/minicpmo_4_5.yaml \
|
||||
--trust-remote-code \
|
||||
--host 0.0.0.0 --port "$PORT"
|
||||
# duplex endpoints: ws://<host>:$PORT/v1/duplex and /v1/realtime?duplex=1
|
||||
# duplex endpoint: ws://<host>:$PORT/v1/realtime?duplex=1 (alias /v1/duplex)
|
||||
@@ -445,7 +445,6 @@ async def run_demo(args: argparse.Namespace) -> dict[str, object]:
|
||||
ref_audio=_ref_audio_data_url(args.ref_audio),
|
||||
temperature=args.temperature,
|
||||
),
|
||||
session_id=args.session_id,
|
||||
reconnect=None,
|
||||
heartbeat_interval_s=None,
|
||||
handshake_timeout_s=args.timeout_s,
|
||||
@@ -638,7 +637,6 @@ def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--url", default="ws://localhost:8099/v1/realtime?duplex=1")
|
||||
parser.add_argument("--model", default="openbmb/MiniCPM-o-4_5")
|
||||
parser.add_argument("--session-id")
|
||||
parser.add_argument(
|
||||
"--input-wav",
|
||||
help="Mono 16 kHz PCM16 clip to stream. Optional when --input-video carries the audio track.",
|
||||
|
||||
@@ -86,7 +86,6 @@
|
||||
url.protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
|
||||
url.searchParams.set('duplex', '1');
|
||||
url.searchParams.set('model', config.model || 'openbmb/MiniCPM-o-4_5');
|
||||
url.searchParams.set('native_duplex', '1');
|
||||
url.searchParams.set('autostart', '0');
|
||||
return url.toString();
|
||||
}
|
||||
@@ -458,7 +457,6 @@
|
||||
settled = true;
|
||||
const extraBody = {
|
||||
auto_response: true,
|
||||
native_duplex: true,
|
||||
};
|
||||
const session = {
|
||||
modalities: ['audio', 'text'],
|
||||
|
||||
@@ -17,6 +17,7 @@ For the full list of supported architectures across all modalities, see
|
||||
| Breeze-TTS-2 | `BreezeBlue/Breeze-TTS-2` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | speaker tags (`S0`..`S9`, default `S0`) | — |
|
||||
| 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) | — | ✓ |
|
||||
| Gepard-1.0 | `nineninesix/gepard-1.0` | — (zero-shot default voice) | ✓ (PCM / WAV stream) | `default` only | — |
|
||||
| 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 | — |
|
||||
| IndexTTS-2.5 | native `checkpoints/` bundle | ✓ (`ref_audio` or uploaded `voice`) | `stream=true` response, non-chunk | uploaded audio voice only; no presets | — |
|
||||
@@ -96,12 +97,52 @@ curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
}' --no-buffer | play -t raw -r 24000 -e signed -b 16 -c 1 -
|
||||
```
|
||||
|
||||
Adjust the player's sample rate to match the model (44.1 kHz for Fish Speech, 48 kHz for VoxCPM2, 22.05 kHz for IndexTTS-2, and 24 kHz for many others).
|
||||
Adjust the player's sample rate to match the model (44.1 kHz for Fish Speech, 48 kHz for VoxCPM2, 22.05 kHz for Gepard and IndexTTS-2, and 24 kHz for many others).
|
||||
|
||||
For full request-shape documentation (all parameters, response formats, error codes), see the [Speech API reference](../../../docs/serving/speech_api.md).
|
||||
|
||||
---
|
||||
|
||||
## Gepard-1.0
|
||||
|
||||
Single-stage native AR TTS at 22.05 kHz mono. Zero-shot only: omit `voice` or pass `"default"`. Voice cloning from reference audio is not available yet.
|
||||
|
||||
### Prerequisites
|
||||
|
||||
Same NeMo NanoCodec install as the [offline Gepard section](../../offline_inference/text_to_speech/README.md#gepard-10). On a host whose CUDA toolkit cannot build kernels — no `nvcc`/`ninja`, or a consumer Blackwell (`sm_120`) card — also `export VLLM_USE_FLASHINFER_SAMPLER=0` before launch (same note as offline). That env is not a deploy-YAML field.
|
||||
|
||||
### Launch
|
||||
|
||||
```bash
|
||||
vllm-omni serve nineninesix/gepard-1.0 --omni --port 8091 --trust-remote-code \
|
||||
--stage-init-timeout 900 \
|
||||
--deploy-config vllm_omni/deploy/gepard.yaml
|
||||
# or:
|
||||
./gepard/run_server.sh
|
||||
```
|
||||
|
||||
`--stage-init-timeout 900` matches the online e2e fixture and the offline example. Serve defaults to 300s, which is often too short for a cold download of the talker plus NanoCodec.
|
||||
|
||||
The packaged `vllm_omni/deploy/gepard.yaml` must be passed with `--deploy-config`. The checkpoint self-identifies as `qwen3_5_text`, so omitting the YAML launches a diffusion fallback instead of the Gepard pipeline. The YAML sets `async_chunk: false`, `max_num_seqs: 4`, and currently pins `seed: 42`, so serving is deterministic by default until that YAML seed is removed. Pass an explicit per-request `seed` in tests and clients rather than depending on either default.
|
||||
|
||||
### Sending requests
|
||||
|
||||
```bash
|
||||
python examples/online_serving/text_to_speech/gepard/speech_client.py \
|
||||
--text "Hello, this is Gepard speaking."
|
||||
|
||||
python examples/online_serving/text_to_speech/gepard/speech_client.py \
|
||||
--text "Hello, this is Gepard speaking." --seed 7 --stream --output output.pcm
|
||||
```
|
||||
|
||||
### Notes
|
||||
|
||||
- Output: 22.05 kHz mono. `max_new_tokens` is a **frame** budget (1 token = 1 frame = 1024 samples ≈ 46.4 ms at 21.5 fps; adapter bounds 1..4096).
|
||||
- Supported request fields: `input` (required), `voice` (`default` only), `response_format` (`wav` default; `wav/pcm/flac/mp3` non-streaming; `opus` 400 because 22.05 kHz is not an Opus sample rate; streaming `pcm/wav` only), `stream` / `stream_format`, `max_new_tokens`, `seed`.
|
||||
- Unsupported: `speed`, `extra_params` (including `temperature`/`top_p`/`top_k`), `ref_audio`, `ref_text`, `speaker_embedding`, `task_type`, `instructions`, `language`, and `word_timestamps`.
|
||||
- Concurrent requests at `max_num_seqs: 4` are supported. Native-AR recompute preemption is a known limitation of this architecture (a request that is preempted mid-generation can resume incorrectly); keep concurrency at or below `max_num_seqs` and treat preemption as out of scope until the platform fix lands.
|
||||
- Optional comparison against the upstream Gepard reference server needs Blackwell/Hopper + CUDA 13 + Postgres and is not part of CI.
|
||||
|
||||
## dots.tts
|
||||
|
||||
Single-stage, text-only TTS at 48 kHz. The current integration supports
|
||||
|
||||
@@ -0,0 +1,33 @@
|
||||
#!/bin/bash
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
# Launch vLLM-Omni server for Gepard-1.0 TTS (zero-shot default voice).
|
||||
#
|
||||
# Usage:
|
||||
# ./run_server.sh
|
||||
# CUDA_VISIBLE_DEVICES=0 ./run_server.sh
|
||||
#
|
||||
# Packaged deploy config: vllm_omni/deploy/gepard.yaml (async_chunk=false,
|
||||
# 22.05 kHz mono, default seed 42 until the YAML seed is removed).
|
||||
# Cold start downloads the talker and NeMo NanoCodec; the serve default of
|
||||
# --stage-init-timeout 300 is too tight (offline end2end.py uses 900).
|
||||
#
|
||||
# On a host whose CUDA toolkit cannot JIT FlashInfer (no nvcc/ninja, or a
|
||||
# consumer Blackwell sm_120 card), also:
|
||||
# export VLLM_USE_FLASHINFER_SAMPLER=0
|
||||
|
||||
set -e
|
||||
|
||||
MODEL="${MODEL:-nineninesix/gepard-1.0}"
|
||||
PORT="${PORT:-8091}"
|
||||
|
||||
echo "Starting Gepard-1.0 server with model: $MODEL"
|
||||
|
||||
vllm-omni serve "$MODEL" \
|
||||
--host 0.0.0.0 \
|
||||
--port "$PORT" \
|
||||
--trust-remote-code \
|
||||
--omni \
|
||||
--stage-init-timeout 900 \
|
||||
--deploy-config vllm_omni/deploy/gepard.yaml
|
||||
@@ -0,0 +1,116 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""Client for Gepard-1.0 TTS via /v1/audio/speech.
|
||||
|
||||
Gepard is zero-shot: omit ``voice`` or pass ``"default"``. Output is 22.05 kHz
|
||||
mono. ``seed`` is optional; the packaged deploy YAML currently pins ``seed: 42``
|
||||
until that default is removed, so requests without an explicit seed are still
|
||||
deterministic.
|
||||
|
||||
Examples:
|
||||
python speech_client.py --text "Hello, this is Gepard speaking."
|
||||
python speech_client.py --text "Hello, this is Gepard speaking." --seed 7
|
||||
python speech_client.py --text "Hello, this is Gepard speaking." --stream --output output.pcm
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
|
||||
import httpx
|
||||
|
||||
DEFAULT_API_BASE = "http://localhost:8091"
|
||||
DEFAULT_API_KEY = "EMPTY"
|
||||
DEFAULT_MODEL = "nineninesix/gepard-1.0"
|
||||
|
||||
|
||||
def run_tts(args) -> None:
|
||||
payload = {
|
||||
"model": args.model,
|
||||
"input": args.text,
|
||||
"voice": args.voice,
|
||||
"response_format": args.response_format,
|
||||
}
|
||||
if args.seed is not None:
|
||||
payload["seed"] = args.seed
|
||||
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"
|
||||
if args.response_format not in ("pcm", "wav"):
|
||||
print(f"Note: streaming requires pcm/wav; overriding --response-format {args.response_format} to 'pcm'")
|
||||
payload["response_format"] = "pcm"
|
||||
|
||||
print(f"Model: {args.model}")
|
||||
print(f"Text: {args.text}")
|
||||
print(f"Voice: {args.voice}")
|
||||
if args.seed is not None:
|
||||
print(f"Seed: {args.seed}")
|
||||
print("Generating audio...")
|
||||
|
||||
api_url = f"{args.api_base}/v1/audio/speech"
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {args.api_key}",
|
||||
}
|
||||
|
||||
if args.stream:
|
||||
output_path = args.output or ("output.wav" if payload["response_format"] == "wav" else "output.pcm")
|
||||
with httpx.Client(timeout=300.0) as client:
|
||||
with client.stream("POST", api_url, json=payload, headers=headers) as resp:
|
||||
if resp.status_code != 200:
|
||||
print(f"Error: {resp.status_code}")
|
||||
print(resp.read().decode())
|
||||
return
|
||||
total_bytes = 0
|
||||
with open(output_path, "wb") as f:
|
||||
for chunk in resp.iter_bytes():
|
||||
f.write(chunk)
|
||||
total_bytes += len(chunk)
|
||||
print(f"Streamed {total_bytes} bytes to: {output_path}")
|
||||
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}")
|
||||
print(response.text)
|
||||
return
|
||||
|
||||
output_path = args.output or "output.wav"
|
||||
with open(output_path, "wb") as f:
|
||||
f.write(response.content)
|
||||
print(f"Audio saved to: {output_path}")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="Gepard-1.0 TTS client")
|
||||
parser.add_argument("--api-base", default=DEFAULT_API_BASE, help="API base URL")
|
||||
parser.add_argument("--api-key", default=DEFAULT_API_KEY, help="API key")
|
||||
parser.add_argument("--model", "-m", default=DEFAULT_MODEL, help="Model name")
|
||||
parser.add_argument("--text", required=True, help="Text to synthesize")
|
||||
parser.add_argument("--voice", default="default", help="Voice name (only 'default' is supported)")
|
||||
parser.add_argument("--seed", type=int, default=None, help="Sampling seed")
|
||||
parser.add_argument(
|
||||
"--max-new-tokens",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Frame budget (1 token = 1 frame ≈ 46.4 ms at 21.5 fps)",
|
||||
)
|
||||
parser.add_argument("--stream", action="store_true", help="Enable streaming (PCM output)")
|
||||
parser.add_argument(
|
||||
"--response-format",
|
||||
default="wav",
|
||||
choices=["wav", "mp3", "flac", "pcm"],
|
||||
help="Audio format (default: wav). Streaming is pcm/wav only. Opus is not supported at 22.05 kHz.",
|
||||
)
|
||||
parser.add_argument("--output", "-o", default=None, help="Output file path")
|
||||
args = parser.parse_args()
|
||||
run_tts(args)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,3 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""Batch speech client for Qwen3-TTS via /v1/audio/speech/batch endpoint.
|
||||
|
||||
This script demonstrates how to synthesize multiple texts in a single request.
|
||||
@@ -7,7 +10,7 @@ the reference for each item.
|
||||
|
||||
Start the server (with batch-optimized stage settings for best throughput):
|
||||
|
||||
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
|
||||
vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base \
|
||||
--omni \
|
||||
--trust-remote-code \
|
||||
--stage-overrides '{"0":{"max_num_seqs":4,"gpu_memory_utilization":0.2},
|
||||
|
||||
@@ -1,3 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""Gradio demo for Qwen3-TTS with gapless streaming audio playback.
|
||||
|
||||
Uses a custom AudioWorklet-based player for gap-free streaming,
|
||||
@@ -351,7 +354,7 @@ def generate_speech(
|
||||
raise gr.Error(f"Failed to decode audio: {e}")
|
||||
|
||||
|
||||
def create_app(api_base: str):
|
||||
def create_app(api_base: str, task_type: str = "CustomVoice"):
|
||||
"""Create the FastAPI app with streaming proxy + Gradio UI."""
|
||||
fastapi_app = FastAPI()
|
||||
|
||||
@@ -463,10 +466,11 @@ def create_app(api_base: str):
|
||||
lines=4,
|
||||
)
|
||||
with gr.Row():
|
||||
task_type = gr.Radio(
|
||||
task_type_input = gr.Radio(
|
||||
choices=TASK_TYPES,
|
||||
value="CustomVoice",
|
||||
value=task_type,
|
||||
label="Task Type",
|
||||
interactive=False,
|
||||
scale=2,
|
||||
)
|
||||
language = gr.Dropdown(
|
||||
@@ -479,17 +483,21 @@ def create_app(api_base: str):
|
||||
choices=voices,
|
||||
value=voices[0] if voices else None,
|
||||
label="speaker",
|
||||
visible=True,
|
||||
visible=task_type == "CustomVoice",
|
||||
allow_custom_value=True,
|
||||
)
|
||||
instructions = gr.Textbox(
|
||||
label="Instructions",
|
||||
placeholder="e.g., Speak with excitement / A warm, friendly female voice",
|
||||
lines=2,
|
||||
visible=True,
|
||||
info="Optional style/emotion instructions",
|
||||
visible=task_type != "Base",
|
||||
info=(
|
||||
"Required: describe the voice style"
|
||||
if task_type == "VoiceDesign"
|
||||
else "Optional style/emotion instructions"
|
||||
),
|
||||
)
|
||||
with gr.Column(visible=False) as ref_group:
|
||||
with gr.Column(visible=task_type == "Base") as ref_group:
|
||||
ref_audio = gr.Audio(
|
||||
label="Reference Audio (upload for voice cloning)",
|
||||
type="numpy",
|
||||
@@ -563,7 +571,7 @@ def create_app(api_base: str):
|
||||
autoplay=True,
|
||||
visible=False,
|
||||
)
|
||||
with gr.Column(visible=True) as examples_cv:
|
||||
with gr.Column(visible=task_type == "CustomVoice") as examples_cv:
|
||||
gr.Examples(
|
||||
examples=[
|
||||
[
|
||||
@@ -588,7 +596,7 @@ def create_app(api_base: str):
|
||||
inputs=[text_input, voice, language, instructions],
|
||||
label="examples",
|
||||
)
|
||||
with gr.Column(visible=False) as examples_vd:
|
||||
with gr.Column(visible=task_type == "VoiceDesign") as examples_vd:
|
||||
gr.Examples(
|
||||
examples=[
|
||||
[
|
||||
@@ -605,7 +613,7 @@ def create_app(api_base: str):
|
||||
inputs=[text_input, language, instructions],
|
||||
label="examples",
|
||||
)
|
||||
with gr.Column(visible=False) as examples_base:
|
||||
with gr.Column(visible=task_type == "Base") as examples_base:
|
||||
gr.Examples(
|
||||
examples=[
|
||||
[
|
||||
@@ -653,9 +661,9 @@ def create_app(api_base: str):
|
||||
gr.update(visible=is_base), # examples_base
|
||||
)
|
||||
|
||||
task_type.change(
|
||||
task_type_input.change(
|
||||
fn=on_task_change,
|
||||
inputs=[task_type],
|
||||
inputs=[task_type_input],
|
||||
outputs=[
|
||||
voice,
|
||||
instructions,
|
||||
@@ -703,7 +711,7 @@ def create_app(api_base: str):
|
||||
|
||||
all_inputs = [
|
||||
text_input,
|
||||
task_type,
|
||||
task_type_input,
|
||||
voice,
|
||||
language,
|
||||
instructions,
|
||||
@@ -788,6 +796,9 @@ def parse_args():
|
||||
description="Gradio demo for Qwen3-TTS with gapless AudioWorklet streaming.",
|
||||
)
|
||||
add_common_args(parser)
|
||||
parser.add_argument(
|
||||
"--task-type", choices=TASK_TYPES, default="CustomVoice", help="Task matching the served checkpoint"
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
@@ -796,7 +807,7 @@ def main():
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
print(f"Connecting to vLLM server at: {args.api_base}")
|
||||
|
||||
app = create_app(args.api_base)
|
||||
app = create_app(args.api_base, task_type=args.task_type)
|
||||
|
||||
import uvicorn
|
||||
|
||||
|
||||
@@ -8,25 +8,28 @@ to generate audio from text using Qwen3-TTS models.
|
||||
|
||||
Examples:
|
||||
# CustomVoice task (predefined speaker)
|
||||
python openai_speech_client.py --text "Hello, how are you?" --voice vivian
|
||||
python openai_speech_client.py --text "Hello, how are you?" --speaker vivian
|
||||
|
||||
# CustomVoice with emotion instruction
|
||||
python openai_speech_client.py --text "I'm so happy!" --voice vivian \
|
||||
python openai_speech_client.py --text "I'm so happy!" --speaker vivian \
|
||||
--instructions "Speak with excitement"
|
||||
|
||||
# VoiceDesign task (voice from description)
|
||||
python openai_speech_client.py --text "Hello world" \
|
||||
--model Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign \
|
||||
--task-type VoiceDesign \
|
||||
--instructions "A warm, friendly female voice"
|
||||
|
||||
# Base task (voice cloning)
|
||||
python openai_speech_client.py --text "Hello world" \
|
||||
--model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
|
||||
--task-type Base \
|
||||
--ref-audio "https://example.com/reference.wav" \
|
||||
--ref-text "This is the reference transcript"
|
||||
|
||||
# Base task with pre-computed speaker embedding
|
||||
python openai_speech_client.py --text "Hello world" \
|
||||
--model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
|
||||
--task-type Base \
|
||||
--speaker-embedding embedding.json
|
||||
"""
|
||||
|
||||
@@ -1,4 +1,7 @@
|
||||
#!/bin/bash
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
# Launch both vLLM server and Gradio demo for Qwen3-TTS
|
||||
#
|
||||
# Usage:
|
||||
@@ -179,7 +182,7 @@ fi
|
||||
echo ""
|
||||
echo "Starting Gradio demo..."
|
||||
cd "$SCRIPT_DIR"
|
||||
GRADIO_CMD=("python" "gradio_demo.py" "--api-base" "$API_BASE" "--host" "$GRADIO_IP" "--port" "$GRADIO_PORT")
|
||||
GRADIO_CMD=("python" "gradio_demo.py" "--api-base" "$API_BASE" "--host" "$GRADIO_IP" "--port" "$GRADIO_PORT" "--task-type" "$TASK_TYPE")
|
||||
if [ "$GRADIO_SHARE" = true ]; then
|
||||
GRADIO_CMD+=("--share")
|
||||
fi
|
||||
|
||||
@@ -123,6 +123,11 @@ fa4 = [
|
||||
"flash-attn-4[cu13]==4.0.0b18",
|
||||
]
|
||||
|
||||
# FastVideo VSA kernels for the supported Linux/Python 3.12 wheel environment.
|
||||
vsa = [
|
||||
"fastvideo-kernel==0.3.4",
|
||||
]
|
||||
|
||||
docs = [
|
||||
"mkdocs>=1.5.0",
|
||||
"mkdocs-api-autonav",
|
||||
|
||||
@@ -1,9 +1,10 @@
|
||||
# MiniMax H3 disaggregated text encoder
|
||||
# MiniMax H3 disaggregated encoder stage
|
||||
|
||||
This opt-in topology runs the Qwen3-VL text encoder as a vLLM stage and sends
|
||||
its hidden states and token-role metadata to an encoder-free diffusion stage.
|
||||
The standard MiniMax H3 recipes remain single-stage and continue to load the
|
||||
text encoder inside the diffusion pipeline.
|
||||
This opt-in topology runs the Qwen3-VL text encoder, video VAE encoder, and
|
||||
audio WVAE encoder together in a vLLM stage. It sends their conditioning
|
||||
payload to an encoder-free diffusion stage. The standard MiniMax H3 recipes
|
||||
remain single-stage and continue to run all components inside the diffusion
|
||||
pipeline.
|
||||
|
||||
## Start the server
|
||||
|
||||
@@ -23,6 +24,14 @@ deployment override for the available hardware. Diffusion quantization,
|
||||
layerwise offload, distributed layerwise offload, VAE parallelism, and USP
|
||||
settings use the same stage 1 options documented in [MiniMax-H3.md](MiniMax-H3.md).
|
||||
|
||||
Stage 0's `tensor_parallel_size` defines one shared rank group for all three
|
||||
encoder roles; it does not choose each role's execution strategy. The required
|
||||
`hf_overrides.minimax_h3_encoder_components` mapping declares those strategies
|
||||
explicitly: `text_encoder` uses tensor parallelism, `video_vae` uses patch
|
||||
parallelism, and `audio_vae` runs on the group leader. All three role entries
|
||||
are mandatory. They share the stage's `max_num_seqs` scheduler and batch limit;
|
||||
the configuration does not create independent role schedulers or world sizes.
|
||||
|
||||
For example, this five-GPU topology assigns one GPU to the encoder and four
|
||||
to the diffusion stage. `--stage-overrides` keeps placement and parallelism
|
||||
scoped to the owning stage rather than broadcasting an override to both:
|
||||
@@ -57,8 +66,12 @@ The Stage 1 VAE patch-parallel options remain independent of offload and
|
||||
quantization. See [MiniMax-H3.md](MiniMax-H3.md) for memory requirements and
|
||||
hardware-qualified profiles before combining these options.
|
||||
|
||||
Stage 1 sets `model_loaded.text_encoder: false`; it must not load or download
|
||||
text-encoder weights. This H3 topology explicitly keeps its single-replica
|
||||
Stage 1 sets `model_loaded.text_encoder: false` and
|
||||
`model_loaded.vae_encoder: false`. It loads only the VAE decoders and skips
|
||||
tokenizer, processor, and text-encoder downloads. It requires complete text and
|
||||
media conditioning from Stage 0, with no fallback to local encoding.
|
||||
|
||||
This H3 topology explicitly keeps its single-replica
|
||||
diffusion stage inline, avoiding serialization of decoded video through a
|
||||
subprocess. It expects both stages to run in one deployment; cross-node payload
|
||||
transport is outside this configuration.
|
||||
|
||||
@@ -23,6 +23,12 @@ One vLLM-Omni diffusion stage can load both DiTs while instantiating the
|
||||
tokenizer, processor, Qwen3-VL text encoder, video VAE, and audio VAE only
|
||||
once. Requests select the DiT with `extra_params.task`.
|
||||
|
||||
This single-stage recipe uses the default `model_loaded.text_encoder: true` and
|
||||
`model_loaded.vae_encoder: true`. Precomputed text conditioning can replace the
|
||||
local Qwen call. Setting `model_loaded.text_encoder: false` while leaving
|
||||
`model_loaded.vae_encoder: true` requires precomputed text conditioning;
|
||||
reference media is still encoded locally.
|
||||
|
||||
The generated MP4 contains H.264 video and synchronized stereo audio.
|
||||
|
||||
## Prerequisites
|
||||
@@ -35,7 +41,7 @@ hf auth login
|
||||
export MODEL=MiniMaxAI/MiniMax-H3
|
||||
```
|
||||
|
||||
The vLLM-Omni pipeline downloads `FL2VA/**`, `Ref2VA/model_index.json`, and
|
||||
By default, the pipeline downloads `FL2VA/**`, `Ref2VA/model_index.json`, and
|
||||
`Ref2VA/transformer/**`. It does not download or load the diffusers-format
|
||||
`transformer`, `transformer_ref`, or `vae` weights at the repository root, nor
|
||||
duplicate Ref2VA copies of shared components.
|
||||
@@ -1230,11 +1236,9 @@ vllm serve "${MODEL_ROOT}/FL2VA" \
|
||||
mode does not support `cache_backend`.
|
||||
- The first regional-compile request is a warmup and should not be included in
|
||||
steady-state performance measurements.
|
||||
- The serving path accepts fewer references than the model supports. H3 documents up
|
||||
to 9 images, 3 video clips, and 3 audio clips (12 files) per Omni Reference
|
||||
request; the current vLLM-Omni path takes exactly one image plus one audio
|
||||
reference, or one or more videos with no separate `audio_reference` (it uses the
|
||||
source soundtracks).
|
||||
- Ref2VA accepts up to 9 images, 3 video clips, and 3 audio clips, with at most
|
||||
12 references in total. At least one image or video is required; audio-only
|
||||
requests are rejected.
|
||||
- The 768 px short-edge mode is available for T2VA and FL2VA; 1344x768 is the
|
||||
documented 16:9 request shape.
|
||||
- `--cfg-parallel-size > 1` is rejected by design (CFG-distilled, no negative branch).
|
||||
@@ -1247,9 +1251,6 @@ vllm serve "${MODEL_ROOT}/FL2VA" \
|
||||
in host memory on every rank during startup before retaining only each rank's
|
||||
shard. Size startup host memory for that transient peak.
|
||||
- TeaCache and Cache-DiT cannot be enabled on the same server.
|
||||
- Image+audio Ref2VA accepts exactly one image and one audio reference.
|
||||
- Video Ref2VA accepts one or more video files, but not an additional standalone
|
||||
audio reference.
|
||||
- Pure Ulysses still replicates the full DiT on every rank, so smaller-memory GPUs
|
||||
cannot use `--usp N --tp 1` as a resident capacity path. Use DiT tensor parallelism
|
||||
or model-level CPU offload; text-encoder TP alone is not sufficient.
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
# MiniCPM-o 4.5
|
||||
|
||||
> Online serving and offline inference for omni multimodal chat
|
||||
> Online full-duplex and chat-completions serving, and offline inference, for omni multimodal chat
|
||||
> (text / image / audio / video → text + 24 kHz speech)
|
||||
|
||||
## Summary
|
||||
@@ -9,8 +9,12 @@
|
||||
- Model: [`openbmb/MiniCPM-o-4_5`](https://huggingface.co/openbmb/MiniCPM-o-4_5)
|
||||
- Task: Omni multimodal chat — accepts text / image / audio / video input;
|
||||
emits text and 24 kHz mono speech in the same response
|
||||
- Mode: Online serving via the OpenAI-compatible `/v1/chat/completions`
|
||||
API (plus Gradio demo), and offline inference via `Omni.generate`
|
||||
- Mode: Online full-duplex serving over the OpenAI Realtime protocol
|
||||
(`/v1/realtime?duplex=1`, alias `/v1/duplex`), the OpenAI-compatible
|
||||
`/v1/chat/completions` API (plus Gradio demo) on the same server, and
|
||||
offline turn-based inference via `Omni.generate`. Besides those the server
|
||||
exposes only `/v1/models` and `/health`; the other turn-based HTTP routes
|
||||
are not available on a duplex deployment.
|
||||
- Maintainer: [`@tc-mb`](https://github.com/tc-mb) (MiniCPM-V / MiniCPM-o team)
|
||||
|
||||
## When to use this recipe
|
||||
@@ -18,9 +22,10 @@
|
||||
Use this recipe as a known-good starting point for serving
|
||||
`openbmb/MiniCPM-o-4_5` on vLLM-Omni. MiniCPM-o 4.5 is the omni member
|
||||
of the MiniCPM-o family — it runs a multimodal thinker, a streaming
|
||||
MiniCPMTTS codec talker, and a separate batched Code2Wav stage so a single
|
||||
`/v1/chat/completions` call can return text and 24 kHz speech in one
|
||||
shot. The recommended batching deploy isolates the Thinker on GPU 0 and
|
||||
MiniCPMTTS codec talker, and a separate batched Code2Wav stage, so a live
|
||||
duplex session streams text and 24 kHz speech while it keeps listening, and a
|
||||
single `/v1/chat/completions` call returns both in one shot. The
|
||||
recommended batching deploy isolates the Thinker on GPU 0 and
|
||||
co-locates Talker and Code2Wav on GPU 1; 1-GPU, 3-GPU, and 8x4090 layouts are
|
||||
also provided.
|
||||
|
||||
@@ -36,8 +41,11 @@ also provided.
|
||||
[`vllm_omni/deploy/minicpmo_4_5_3gpu.yaml`](../../vllm_omni/deploy/minicpmo_4_5_3gpu.yaml)
|
||||
- 8x RTX 4090 layout:
|
||||
[`vllm_omni/deploy/minicpmo_4_5_8x4090.yaml`](../../vllm_omni/deploy/minicpmo_4_5_8x4090.yaml)
|
||||
- Online example + Gradio demo:
|
||||
- Online examples (Realtime CLI demo, browser client, barge-in client,
|
||||
chat-completions client, curl script, Gradio demo):
|
||||
[`examples/online_serving/minicpmo/`](../../examples/online_serving/minicpmo/)
|
||||
- Duplex API and Python client docs:
|
||||
[`docs/serving/realtime_duplex_api.md`](../../docs/serving/realtime_duplex_api.md)
|
||||
- Offline end-to-end example:
|
||||
[`examples/offline_inference/minicpmo/`](../../examples/offline_inference/minicpmo/)
|
||||
- Pipeline / talker source:
|
||||
@@ -179,6 +187,13 @@ as one aggregated chunk, so the table reports Stage 0 metrics.
|
||||
|
||||
#### Verification
|
||||
|
||||
**Health and model listing**:
|
||||
|
||||
```bash
|
||||
curl http://localhost:8099/health
|
||||
curl http://localhost:8099/v1/models
|
||||
```
|
||||
|
||||
**Quick smoke test (text-only output)**:
|
||||
|
||||
```bash
|
||||
@@ -191,8 +206,31 @@ curl http://localhost:8099/v1/chat/completions \
|
||||
}'
|
||||
```
|
||||
|
||||
**Text + speech in one response** (the headline 4.5 feature). The model
|
||||
bridge conditions the Talker from the generated assistant span, so the
|
||||
**One duplex turn over the Realtime WebSocket** (text + speech stream back
|
||||
while the session keeps listening):
|
||||
|
||||
```bash
|
||||
python examples/online_serving/minicpmo/realtime_duplex_demo.py \
|
||||
--url ws://localhost:8099/v1/realtime?duplex=1 \
|
||||
--model openbmb/MiniCPM-o-4_5 \
|
||||
--input-wav /path/to/input_16k_mono_pcm16.wav \
|
||||
--ref-audio /path/to/MiniCPM-o-Demo/assets/ref_audio/ref_minicpm_signature.wav \
|
||||
--output-dir /tmp/minicpmo_realtime_duplex_demo
|
||||
```
|
||||
|
||||
The demo prints transcript deltas as they arrive and saves the streamed
|
||||
24 kHz audio to WAV files. The same session shape is available from Python
|
||||
through `vllm_omni.clients.duplex.DuplexClient` (over the server) or
|
||||
`vllm_omni.clients.inline_duplex.InlineDuplexClient` (in-process, no server;
|
||||
`examples/online_serving/barge_in_client.py --inline`).
|
||||
|
||||
**Barge-in**: `examples/online_serving/barge_in_client.py` plays a question,
|
||||
interrupts the answer with a second WAV and checks that the first response is
|
||||
cancelled and the second one is answered.
|
||||
|
||||
**Text + speech in one response over `/v1/chat/completions`** (the headline
|
||||
4.5 feature, served by the ordinary chat service on the duplex engine). The
|
||||
model bridge conditions the Talker from the generated assistant span, so the
|
||||
generic serving layer does not inject MiniCPM-specific template defaults.
|
||||
`use_tts_template=true` remains supported when explicitly requested:
|
||||
|
||||
@@ -243,6 +281,12 @@ python examples/online_serving/minicpmo/gradio_demo.py \
|
||||
Open `http://<host>:7862` and try a text prompt with the **"Generate
|
||||
speech output (TTS)"** checkbox on / off.
|
||||
|
||||
**Turn-based generation** is also available offline
|
||||
(`examples/offline_inference/minicpmo/`, `Omni.generate`).
|
||||
|
||||
The Daily-Omni numbers above were measured on the `/v1/chat/completions`
|
||||
route.
|
||||
|
||||
#### Notes
|
||||
|
||||
- Memory budget: Thinker, Talker, and Code2Wav reserve 0.55, 0.15, and
|
||||
@@ -326,9 +370,13 @@ vllm serve openbmb/MiniCPM-o-4_5 --omni \
|
||||
and defaults to deterministic seed 42. Stage-1 deploy sampling parameters
|
||||
control only vLLM's binary continue/stop token.
|
||||
|
||||
- **Output audio**: 24 kHz mono WAV inside the OpenAI-style
|
||||
`message.audio.data` (base64). The Gradio demo's WAV player decodes
|
||||
this automatically.
|
||||
- **Output audio**: 24 kHz mono. On `/v1/chat/completions` it is base64 WAV
|
||||
inside the OpenAI-style `message.audio.data` (the Gradio demo's WAV player
|
||||
decodes this automatically); on a duplex session it streams as base64
|
||||
`response.audio.delta` events; offline `Omni.generate` returns WAV.
|
||||
- **Response choices**: on `/v1/chat/completions` text and audio are separate
|
||||
choices. SDK clients should select the choice whose `message.audio.data` is
|
||||
populated rather than assuming `choices[0]` contains audio.
|
||||
|
||||
- **Routing**: MiniCPM-o 4.5 and 2.6 both ship `architectures=
|
||||
["MiniCPMO"]` in HF config; routing is disambiguated by
|
||||
@@ -340,6 +388,3 @@ vllm serve openbmb/MiniCPM-o-4_5 --omni \
|
||||
- **Async chunking**: enabled in all deploy configs. Talker sends
|
||||
25-code chunks with three-code left context to Code2Wav through
|
||||
`SharedMemoryConnector`; terminal chunks flush held lookahead state.
|
||||
- **Response choices**: text and audio are separate choices. SDK clients
|
||||
should select the choice whose `message.audio.data` is populated rather
|
||||
than assuming `choices[0]` contains audio.
|
||||
+1
-1
@@ -71,7 +71,7 @@ recipes/
|
||||
| [`cosmos3/Cosmos3-Nano.md`](./cosmos3/Cosmos3-Nano.md) | Text-to-image, text-to-video, image-to-video, video-to-video generation, text to video with sound, action policy | 1x H200 141GB / B300 |
|
||||
| [`cosmos3/Cosmos3-Edge.md`](./cosmos3/Cosmos3-Edge.md) | T2I / T2V / I2V + action (Physical-AI) generation on the Nemotron-based Edge transformer | 1x GPU (~8 GiB) |
|
||||
| [`cosmos3/Cosmos3-Super.md`](./cosmos3/Cosmos3-Super.md) | 64B T2I / T2V / I2V / V2V generation (+ optional audio) / Action policy | 8x H200/H100/A100 / 2x H200 / B300 |
|
||||
| [`OpenBMB/MiniCPM-o-4_5.md`](./OpenBMB/MiniCPM-o-4_5.md) | Online serving for omni multimodal chat (text / image / audio / video → text + 24 kHz speech) | 2x A100/H100 80GB / 3x mid-tier GPU / 8x RTX 4090 24GB |
|
||||
| [`OpenBMB/MiniCPM-o-4_5.md`](./OpenBMB/MiniCPM-o-4_5.md) | Online full-duplex serving and offline inference for omni multimodal chat (text / image / audio / video → text + 24 kHz speech) | 2x A100/H100 80GB / 3x mid-tier GPU / 8x RTX 4090 24GB |
|
||||
| [`OpenBMB/VoxCPM2.md`](./OpenBMB/VoxCPM2.md) | Online + offline TTS with native AR pipeline (48 kHz, 30+ languages) | 1x RTX 4090 24GB |
|
||||
| [`OpenMOSS/MOSS-TTS.md`](./OpenMOSS/MOSS-TTS.md) | Online + offline multilingual TTS (MOSS-TTS family, 8B) | 1x H100 80GB |
|
||||
| [`OpenMOSS/MOSS-TTS-Nano-NPU.md`](./OpenMOSS/MOSS-TTS-Nano-NPU.md) | Online multilingual voice-cloned TTS | 1x Ascend 910B3 64GB |
|
||||
|
||||
@@ -32,13 +32,13 @@ mode is selected per request:
|
||||
T2V/I2V `/v1/videos/sync` request to also generate synchronized audio, muxed into
|
||||
the mp4 as AAC 48 kHz stereo. See the official model card's "Video + Audio" examples.
|
||||
- **Action** — pass `extra_params={"action_mode": ...}` to drive Physical-AI tasks:
|
||||
- `forward_dynamics` — given a first frame or video **and** an action trajectory,
|
||||
- `forward_dynamics` — given a first frame or video **and** an action trajectory,
|
||||
roll out the resulting video. Synchronous: `POST /v1/videos/sync`.
|
||||
- `policy` — given a first frame or video and a language instruction,
|
||||
- `policy` — given a first frame or video and a language instruction,
|
||||
**predict** the action trajectory (and a rollout video). Use the async
|
||||
`POST /v1/videos` endpoint and read the predicted action from the top-level
|
||||
`action` field.
|
||||
- `inverse_dynamics` — given a video, **recover** the action trajectory. Use
|
||||
- `inverse_dynamics` — given a video, **recover** the action trajectory. Use
|
||||
the async `POST /v1/videos` endpoint and read the recovered action from
|
||||
the top-level `action` field
|
||||
(`{data, shape, dtype, raw_action_dim, domain_id}`).
|
||||
@@ -106,6 +106,25 @@ vllm serve nvidia/Cosmos3-Nano \
|
||||
--init-timeout 1800
|
||||
```
|
||||
|
||||
For Cosmos3, SeaCache is the recommended default choice when opting into
|
||||
diffusion caching. Caching remains opt-in; add `--cache-backend sea_cache` to
|
||||
the command above. Its default threshold and maximum cached-step streak are
|
||||
tuned for Cosmos3, so no `--cache-config` is required.
|
||||
|
||||
Override individual defaults with a JSON cache configuration; for example:
|
||||
|
||||
```bash
|
||||
vllm serve nvidia/Cosmos3-Nano \
|
||||
--omni \
|
||||
--cache-backend sea_cache \
|
||||
--cache-config '{"sea_threshold":0.2,"sea_max_consecutive_cached":3}'
|
||||
```
|
||||
|
||||
Lower `sea_threshold` values and smaller `sea_max_consecutive_cached` caps are
|
||||
more conservative. Higher values allow more cached steps and may improve
|
||||
speed, but can increase quality loss. Setting `sea_max_consecutive_cached` to
|
||||
`0` removes the streak cap.
|
||||
|
||||
To run **without** guardrails (you are responsible for license compliance),
|
||||
add `--no-guardrails` (no token/`cosmos-guardrail` needed). For extra GPUs use
|
||||
`--ulysses-degree N` (context parallel) or `--tensor-parallel-size N`;
|
||||
@@ -382,12 +401,12 @@ OpenPI route uses the normal API authentication middleware. RoboLab's
|
||||
|
||||
#### Notes
|
||||
|
||||
- **Measured latency (1x B300, bf16, guardrails off):**
|
||||
- T2I 1024² — 10 / 25 / 50 steps → ~0.4 / 0.7 / **1.3 s**
|
||||
- T2V 1280×720 @ 35 steps — 25 / 49 / 93 / **189** frames → ~7 / 15 / 33 / **~93 s**
|
||||
- I2V 1280×720, 189 frames @ 35 steps → ~**99 s**
|
||||
- Action 640×480 @ 30 steps — forward-dynamics 61f ~**4 s**, policy 17f ~**1–3 s**.
|
||||
- Guardrails-on overhead: ~8% on T2I, negligible on video.
|
||||
- **Measured latency (1x B300, bf16, guardrails off, diffusion cache off):**
|
||||
- T2I 1024² — 10 / 25 / 50 steps → ~0.4 / 0.7 / **1.3 s**
|
||||
- T2V 1280×720 @ 35 steps — 25 / 49 / 93 / **189** frames → ~7 / 15 / 33 / **~93 s**
|
||||
- I2V 1280×720, 189 frames @ 35 steps → ~**99 s**
|
||||
- Action 640×480 @ 30 steps — forward-dynamics 61f ~**4 s**, policy 17f ~**1–3 s**.
|
||||
- Guardrails-on overhead: ~8% on T2I, negligible on video.
|
||||
- **Memory:** transformer ~17 GiB (bf16); peak ~46 GiB for 720p video on 1 GPU;
|
||||
full repo (transformer + Wan VAE + Qwen3-VL vision encoder + audio tokenizer)
|
||||
~33 GB on disk.
|
||||
@@ -443,12 +462,12 @@ OpenPI route uses the normal API authentication middleware. RoboLab's
|
||||
`conditioning_fps`, `action_chunk_size`, `raw_action_dim`, `deterministic_seed`,
|
||||
and `session_id`.
|
||||
- **Known limitations:**
|
||||
- Guardrails-on requires `cosmos-guardrail` **and** access to the gated
|
||||
- Guardrails-on requires `cosmos-guardrail` **and** access to the gated
|
||||
`nvidia/Cosmos-1.0-Guardrail` repo (accept license + `HF_TOKEN`); otherwise
|
||||
the server fails at pipeline build with a gated-repo / safety-checker error.
|
||||
- A guardrail-blocked prompt currently returns HTTP 500
|
||||
- A guardrail-blocked prompt currently returns HTTP 500
|
||||
(`"Guardrail blocked prompt"`).
|
||||
- Action `forward_dynamics`, `policy`, and `inverse_dynamics` are supported
|
||||
- Action `forward_dynamics`, `policy`, and `inverse_dynamics` are supported
|
||||
online. Use async `POST /v1/videos` when you need the predicted/recovered
|
||||
action payload under the top-level `action` field; sync `/v1/videos/sync`
|
||||
returns raw MP4 bytes and does not expose action metadata in the response body.
|
||||
@@ -579,14 +598,14 @@ so an md5 comparison across two runs also works as a smoke check.
|
||||
a ~390 s build of the VAE decode path *for that shape*, so warming 1024² images
|
||||
does nothing for 189-frame video. At 720p / 189 frames / 35 steps this is
|
||||
**540 s cold vs 161 s warm, with byte-identical output**.
|
||||
- **Measured on 1x MI350X (bf16, guardrails off, warm):** T2I 1024² @ 50 steps
|
||||
- **Measured on 1x MI350X (bf16, guardrails off, diffusion cache off, warm):** T2I 1024² @ 50 steps
|
||||
**~2.7 s**; T2V 1280×720 / 189 frames @ 35 steps **~161 s**, of which ~92% is
|
||||
the DiT denoise loop and ~4% VAE decode, so optimization effort belongs in the
|
||||
denoise loop. The optional flags act on different terms of the memory bill and
|
||||
are therefore complementary rather than redundant:
|
||||
|
||||
| Flag | Latency | Peak reserved | Peak allocated | Acts on |
|
||||
|---|---|---|---|---|
|
||||
| --- | --- | --- | --- | --- |
|
||||
| *(none)* | 161 s | 120 GiB | 95 GiB | — |
|
||||
| `--vae-use-tiling` | 183 s (+13.5%) | **38 GiB** (−68%) | **36 GiB** (−62%) | decode activations |
|
||||
| `--enable-layerwise-offload` | 162 s (+0.8%) | 84 GiB (−30%) | 69 GiB (−27%) | weights |
|
||||
@@ -768,11 +787,11 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync \
|
||||
|
||||
#### Notes
|
||||
|
||||
- **Measured latency (1x Ascend 910B / 910C, bf16, guardrails off):**
|
||||
- T2I 1024² — 10 steps → ~8 s
|
||||
- T2V 1280×720 @ 20 steps — 49 frames → ~55 s
|
||||
- I2V 1280×720 @ 10 steps — 25 frames → ~25 s
|
||||
- V2V 480×320 @ 10 steps — 17 frames → ~12 s
|
||||
- **Measured latency (1x Ascend 910B / 910C, bf16, guardrails off, diffusion cache off):**
|
||||
- T2I 1024² — 10 steps → ~8 s
|
||||
- T2V 1280×720 @ 20 steps — 49 frames → ~55 s
|
||||
- I2V 1280×720 @ 10 steps — 25 frames → ~25 s
|
||||
- V2V 480×320 @ 10 steps — 17 frames → ~12 s
|
||||
- **Memory:** transformer ~17 GiB (bf16); peak ~46 GiB for 720p video on 1 NPU;
|
||||
full repo (transformer + Wan VAE + Qwen3-VL vision encoder + audio tokenizer)
|
||||
~33 GB on disk.
|
||||
@@ -785,9 +804,9 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync \
|
||||
(for model loading), `--tensor-parallel-size 8` for multi-NPU, and
|
||||
`--model-class-name Cosmos3OmniDiffusersPipeline` to force the pipeline class.
|
||||
- **Known limitations:**
|
||||
- Transfer V2V with `extra_params` (`edge`/`blur`/`depth`/`seg`/`wsm`) hits a
|
||||
- Transfer V2V with `extra_params` (`edge`/`blur`/`depth`/`seg`/`wsm`) hits a
|
||||
resolution-parsing bug; basic V2V without transfer hints works.
|
||||
- FP8 online quantization and layerwise offload are not supported on NPU.
|
||||
- FP8 online quantization and layerwise offload are not supported on NPU.
|
||||
|
||||
### 1x Ascend 910B / 910C (Atlas A2 / A3) — Offline generation
|
||||
|
||||
|
||||
@@ -62,6 +62,21 @@ disable. `--enable-layerwise-offload` reduces VRAM on smaller GPUs;
|
||||
generation from ~83 GB to ~55 GB per GPU (2-GPU) with BF16-level quality (T2V
|
||||
composition can shift at the same seed).
|
||||
|
||||
SeaCache is the recommended default choice when opting into diffusion caching,
|
||||
which remains opt-in. Add `--cache-backend sea_cache` to either command above;
|
||||
its defaults are tuned for Cosmos3.
|
||||
|
||||
Override individual defaults with `--cache-config`; for example:
|
||||
|
||||
```bash
|
||||
--cache-config '{"sea_threshold":0.2,"sea_max_consecutive_cached":3}'
|
||||
```
|
||||
|
||||
Lower `sea_threshold` values and smaller `sea_max_consecutive_cached` caps are
|
||||
more conservative. Higher values allow more cached steps and may improve
|
||||
speed, but can increase quality loss. Setting `sea_max_consecutive_cached` to
|
||||
`0` removes the streak cap.
|
||||
|
||||
#### Verification
|
||||
|
||||
Requests are identical to Nano (see [`Cosmos3-Nano.md`](./Cosmos3-Nano.md) for full
|
||||
@@ -109,12 +124,12 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \
|
||||
|
||||
#### Notes
|
||||
|
||||
- **Measured (2x B300, bf16, guardrails off, official 2-GPU config above):**
|
||||
- T2I 1024², 50 steps → **~6 s**
|
||||
- T2V 1280×720, 189 frames, 35 steps → **~197 s**
|
||||
- I2V 1280×720, 189 frames, 35 steps → **~200 s**
|
||||
- T2V + sound (189 frames, 35 steps) → **~198 s**, output muxes **AAC 48 kHz stereo**
|
||||
- (NVIDIA's reference: 8×H200 @ 50 steps ≈ 55 s/video; 2×H200 @ 35 steps ≈ 3 min/video.)
|
||||
- **Measured (2x B300, bf16, guardrails off, diffusion cache off, official 2-GPU config above):**
|
||||
- T2I 1024², 50 steps → **~6 s**
|
||||
- T2V 1280×720, 189 frames, 35 steps → **~197 s**
|
||||
- I2V 1280×720, 189 frames, 35 steps → **~200 s**
|
||||
- T2V + sound (189 frames, 35 steps) → **~198 s**, output muxes **AAC 48 kHz stereo**
|
||||
- (NVIDIA's reference: 8×H200 @ 50 steps ≈ 55 s/video; 2×H200 @ 35 steps ≈ 3 min/video.)
|
||||
- **Memory:** ~61.5 GiB per GPU when sharded across 2 GPUs (HSDP shard 2); repo ~135 GB on disk.
|
||||
- Same generation defaults, supported sizes, V2V reference-video controls
|
||||
(`condition_frame_indexes_vision`, `condition_video_keep`), and
|
||||
|
||||
@@ -178,7 +178,7 @@ mistralai/Voxtral-4B-TTS-2603 checkpoint. running the
|
||||
voice-cloning e2e test ([`Voice cloning (capability gated upstream)`](../../examples/offline_inference/text_to_speech/README.md#voice-cloning-capability-gated-upstream)
|
||||
) or adding a `ref_audio` request parameter will fail
|
||||
because the public checkpoint does not provide the encoder weights needed to
|
||||
convert reference audio into conditioning features. The expected failure is: `RuntimeError: encode_waveforms requires encoder weights which are not available in the open-source checkpoint.`
|
||||
convert reference audio into conditioning features. The expected failure is a 400 BadRequest `Voice cloning with 'ref_audio' requires the full Voxtral checkpoint with encoder weights. The open-source variant only supports preset voices via the 'voice' parameter.`.
|
||||
|
||||
### 1x NVIDIA RTX A6000 48GB
|
||||
|
||||
|
||||
@@ -64,6 +64,177 @@ def test_hugging_face_subset_names_are_splits(monkeypatch):
|
||||
assert sample.task_type == "correction"
|
||||
|
||||
|
||||
def test_local_hf_dataset_directory_is_routed_to_hf_loader(tmp_path, monkeypatch):
|
||||
# An existing local Hugging Face dataset layout (directory with a
|
||||
# ``data/*.parquet`` file) must be loaded through ``datasets.load_dataset``
|
||||
# instead of being parsed as a JSON/JSONL manifest.
|
||||
hf_dir = tmp_path / "Omni-DuplexEval"
|
||||
data_dir = hf_dir / "data"
|
||||
data_dir.mkdir(parents=True)
|
||||
(data_dir / "train-00000-of-00001.parquet").write_bytes(b"PAR1")
|
||||
calls = []
|
||||
|
||||
def fake_load_dataset(name, **kwargs):
|
||||
calls.append((name, kwargs))
|
||||
return [{"id": "pr", "split": "PR_correction", "question_text": "Correct this."}]
|
||||
|
||||
monkeypatch.setitem(sys.modules, "datasets", SimpleNamespace(load_dataset=fake_load_dataset))
|
||||
samples = load_samples(hf_dir)
|
||||
assert calls == [(str(hf_dir), {})]
|
||||
assert len(samples) == 1
|
||||
assert samples[0].task_type == "correction"
|
||||
|
||||
|
||||
def test_local_parquet_file_is_routed_to_hf_loader(tmp_path, monkeypatch):
|
||||
parquet = tmp_path / "samples.parquet"
|
||||
parquet.write_bytes(b"PAR1")
|
||||
calls = []
|
||||
|
||||
def fake_load_dataset(name, **kwargs):
|
||||
calls.append((name, kwargs))
|
||||
return [{"id": "pr", "split": "PR_correction", "question_text": "Correct this."}]
|
||||
|
||||
monkeypatch.setitem(sys.modules, "datasets", SimpleNamespace(load_dataset=fake_load_dataset))
|
||||
samples = load_samples(str(parquet))
|
||||
assert calls == [("parquet", {"data_files": str(parquet)})]
|
||||
assert len(samples) == 1
|
||||
assert samples[0].family == "pr"
|
||||
|
||||
|
||||
def test_hf_collapse_to_train_keeps_row_split_identity(tmp_path, monkeypatch):
|
||||
# datasets.load_dataset collapses a config-less local mirror (a bare data/
|
||||
# directory) into a single generic ``train`` split. Rows that already carry
|
||||
# their own split/subset identity must keep it so family inference is not
|
||||
# poisoned by the generic name (regression for the direct-read path).
|
||||
hf_dir = tmp_path / "Omni-DuplexEval"
|
||||
data_dir = hf_dir / "data"
|
||||
data_dir.mkdir(parents=True)
|
||||
(data_dir / "train-00000-of-00021.parquet").write_bytes(b"PAR1")
|
||||
|
||||
def fake_load_dataset(name, **kwargs):
|
||||
return {
|
||||
"train": [
|
||||
{"id": "pr", "split": "PR_correction", "question_text": "Correct this."},
|
||||
{"id": "pr-subset", "subset": "PR_event_reminder", "question_text": "Remind me."},
|
||||
{"id": "rtd", "split": "RTD_OCR", "question_text": "Read this."},
|
||||
]
|
||||
}
|
||||
|
||||
monkeypatch.setitem(sys.modules, "datasets", SimpleNamespace(load_dataset=fake_load_dataset))
|
||||
samples = load_samples(hf_dir)
|
||||
by_id = {sample.id: sample for sample in samples}
|
||||
assert by_id["pr"].split == "PR_correction"
|
||||
assert by_id["pr"].family == "pr"
|
||||
assert by_id["pr"].task_type == "correction"
|
||||
assert by_id["pr-subset"].split == "PR_event_reminder"
|
||||
assert by_id["pr-subset"].family == "pr"
|
||||
assert by_id["pr-subset"].task_type == "proactive_reminder"
|
||||
assert by_id["rtd"].split == "RTD_OCR"
|
||||
assert by_id["rtd"].family == "rtd"
|
||||
|
||||
|
||||
def test_hf_config_layout_stamps_split_name_when_row_has_none(tmp_path, monkeypatch):
|
||||
# With a proper config-per-split mirror the loader still stamps the Hugging
|
||||
# Face split name onto rows that do not identify themselves (unchanged).
|
||||
hf_dir = tmp_path / "Omni-DuplexEval"
|
||||
data_dir = hf_dir / "data"
|
||||
data_dir.mkdir(parents=True)
|
||||
(data_dir / "RTD_OCR-00000-of-00001.parquet").write_bytes(b"PAR1")
|
||||
(data_dir / "PR_correction-00000-of-00001.parquet").write_bytes(b"PAR1")
|
||||
|
||||
def fake_load_dataset(name, **kwargs):
|
||||
return {
|
||||
"RTD_OCR": [{"id": "rtd", "question_text": "Read this."}],
|
||||
"PR_correction": [{"id": "pr", "question_text": "Correct this."}],
|
||||
}
|
||||
|
||||
monkeypatch.setitem(sys.modules, "datasets", SimpleNamespace(load_dataset=fake_load_dataset))
|
||||
samples = load_samples(hf_dir)
|
||||
by_id = {sample.id: sample for sample in samples}
|
||||
assert by_id["rtd"].split == "RTD_OCR"
|
||||
assert by_id["rtd"].family == "rtd"
|
||||
assert by_id["pr"].split == "PR_correction"
|
||||
assert by_id["pr"].family == "pr"
|
||||
assert by_id["pr"].task_type == "correction"
|
||||
|
||||
|
||||
def test_hf_collapse_to_train_without_row_identity_raises_actionable_error(tmp_path, monkeypatch):
|
||||
# A row that carries no split/family/task_type and lands in a generic
|
||||
# collapsed ``train`` split cannot be routed; the loader must raise a clear,
|
||||
# actionable ValueError instead of the bare family_for_split message.
|
||||
hf_dir = tmp_path / "Omni-DuplexEval"
|
||||
data_dir = hf_dir / "data"
|
||||
data_dir.mkdir(parents=True)
|
||||
(data_dir / "train-00000-of-00021.parquet").write_bytes(b"PAR1")
|
||||
|
||||
def fake_load_dataset(name, **kwargs):
|
||||
return {"train": [{"id": "x", "question_text": "What is this?"}]}
|
||||
|
||||
monkeypatch.setitem(sys.modules, "datasets", SimpleNamespace(load_dataset=fake_load_dataset))
|
||||
with pytest.raises(ValueError, match="Hugging Face splits"):
|
||||
load_samples(hf_dir)
|
||||
|
||||
|
||||
def test_non_manifest_file_raises_clear_error(tmp_path):
|
||||
# A plain non-JSON file (e.g. a README) must surface a clear ValueError
|
||||
# listing the accepted inputs instead of a raw JSONDecodeError.
|
||||
bad = tmp_path / "README.md"
|
||||
bad.write_text("# readme\nnot a manifest", encoding="utf-8")
|
||||
with pytest.raises(ValueError, match="JSON/JSONL manifest"):
|
||||
load_samples(str(bad))
|
||||
|
||||
|
||||
def _write_parquet(path, rows):
|
||||
pyarrow = pytest.importorskip("pyarrow")
|
||||
pyarrow.parquet.write_table(pyarrow.Table.from_pylist(rows), str(path))
|
||||
|
||||
|
||||
def test_local_parquet_split_filters_rows_by_identity(tmp_path):
|
||||
# Regression: a single ``.parquet`` file is exposed by ``load_dataset`` as a
|
||||
# generic ``train`` physical split, so ``--split RTD_OCR`` used to raise
|
||||
# ``Unknown split "RTD_OCR"``. The requested split must instead filter the
|
||||
# rows by their preserved identity.
|
||||
pytest.importorskip("datasets")
|
||||
parquet = tmp_path / "samples.parquet"
|
||||
_write_parquet(
|
||||
parquet,
|
||||
[
|
||||
{"id": "rtd", "split": "RTD_OCR", "question_text": "Read this."},
|
||||
{"id": "pr", "split": "PR_correction", "question_text": "Correct this."},
|
||||
],
|
||||
)
|
||||
assert {sample.id for sample in load_samples(str(parquet))} == {"rtd", "pr"}
|
||||
rtd = load_samples(str(parquet), split="RTD_OCR")
|
||||
assert [sample.id for sample in rtd] == ["rtd"]
|
||||
assert rtd[0].family == "rtd"
|
||||
pr = load_samples(str(parquet), split="PR_correction")
|
||||
assert [sample.id for sample in pr] == ["pr"]
|
||||
assert pr[0].family == "pr"
|
||||
assert pr[0].task_type == "correction"
|
||||
|
||||
|
||||
def test_local_directory_named_splits_resolve_without_row_identity(tmp_path):
|
||||
# Regression: a single-configuration mirror whose data files are named after
|
||||
# the benchmark splits must resolve each split for real (no fake loader) and
|
||||
# without any per-row identity.
|
||||
pytest.importorskip("datasets")
|
||||
hf_dir = tmp_path / "Omni-DuplexEval"
|
||||
data_dir = hf_dir / "data"
|
||||
data_dir.mkdir(parents=True)
|
||||
_write_parquet(data_dir / "RTD_OCR-00000-of-00001.parquet", [{"id": "rtd", "question_text": "Read this."}])
|
||||
_write_parquet(
|
||||
data_dir / "PR_correction-00000-of-00001.parquet",
|
||||
[{"id": "pr", "question_text": "Correct this."}],
|
||||
)
|
||||
by_id = {sample.id: sample for sample in load_samples(hf_dir)}
|
||||
assert by_id["rtd"].split == "RTD_OCR"
|
||||
assert by_id["rtd"].family == "rtd"
|
||||
assert by_id["pr"].split == "PR_correction"
|
||||
assert by_id["pr"].family == "pr"
|
||||
assert by_id["pr"].task_type == "correction"
|
||||
assert [sample.id for sample in load_samples(hf_dir, split="RTD_OCR")] == ["rtd"]
|
||||
|
||||
|
||||
def test_response_aliases_and_clock_guard():
|
||||
assert split_text("One. Two!") == ["One.", "Two!"]
|
||||
assert normalize_response_items({"chunks": [{"text": "x", "current_time": 800}]}) == [
|
||||
@@ -141,3 +312,113 @@ def test_frame_sample_content_passes_frames(tmp_path, monkeypatch):
|
||||
judge_video_mode="frame-sample",
|
||||
)
|
||||
assert score["content"]["frame_count"] == 2
|
||||
|
||||
|
||||
def test_local_directory_collapsed_train_with_row_identity_and_split_filter(tmp_path):
|
||||
# Reviewer-pointed gap: a collapsed local mirror (data/train-*.parquet)
|
||||
# whose rows carry their own split identity must correctly filter by
|
||||
# --split RTD_OCR.
|
||||
pytest.importorskip("datasets")
|
||||
hf_dir = tmp_path / "Omni-DuplexEval"
|
||||
data_dir = hf_dir / "data"
|
||||
data_dir.mkdir(parents=True)
|
||||
_write_parquet(
|
||||
data_dir / "train-00000-of-00001.parquet",
|
||||
[
|
||||
{"id": "rtd", "split": "RTD_OCR", "question_text": "Read this."},
|
||||
{"id": "pr", "split": "PR_correction", "question_text": "Correct this."},
|
||||
],
|
||||
)
|
||||
all_samples = load_samples(hf_dir)
|
||||
assert {s.id for s in all_samples} == {"rtd", "pr"}
|
||||
rtd = load_samples(hf_dir, split="RTD_OCR")
|
||||
assert [s.id for s in rtd] == ["rtd"]
|
||||
assert rtd[0].family == "rtd"
|
||||
|
||||
|
||||
def test_local_directory_mistyped_split_raises_clear_error(tmp_path):
|
||||
# A named-shard mirror with --split RTD_TYPO must raise ValueError
|
||||
# listing the observed splits instead of silently returning nothing.
|
||||
pytest.importorskip("datasets")
|
||||
hf_dir = tmp_path / "Omni-DuplexEval"
|
||||
data_dir = hf_dir / "data"
|
||||
data_dir.mkdir(parents=True)
|
||||
_write_parquet(
|
||||
data_dir / "RTD_OCR-00000-of-00001.parquet",
|
||||
[{"id": "rtd", "question_text": "Read this."}],
|
||||
)
|
||||
_write_parquet(
|
||||
data_dir / "PR_correction-00000-of-00001.parquet",
|
||||
[{"id": "pr", "question_text": "Correct this."}],
|
||||
)
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
load_samples(hf_dir, split="RTD_TYPO")
|
||||
message = str(exc_info.value)
|
||||
assert "matched zero" in message
|
||||
assert "RTD_OCR" in message or "PR_correction" in message
|
||||
|
||||
|
||||
def test_manifest_without_split_identity_split_override_still_works(tmp_path):
|
||||
# Manifest rows without split/subset/config identity must keep the
|
||||
# base override behaviour: --split RTD_OCR stamps the split so rows
|
||||
# are kept instead of rejected (regression test for P2).
|
||||
manifest = tmp_path / "samples.json"
|
||||
manifest.write_text(
|
||||
json.dumps([{"id": "a", "video": "clip.mp4"}, {"id": "b", "video": "clip2.mp4"}]),
|
||||
encoding="utf-8",
|
||||
)
|
||||
samples = load_samples(manifest, split="RTD_OCR")
|
||||
assert len(samples) == 2
|
||||
assert all(s.split == "RTD_OCR" for s in samples)
|
||||
assert all(s.family == "rtd" for s in samples)
|
||||
|
||||
|
||||
def test_iterable_without_split_identity_split_override_still_works():
|
||||
# Iterable rows without split/subset/config identity must keep the
|
||||
# base override behaviour (regression test for F5/P2).
|
||||
data = [{"id": "a", "video": "clip.mp4"}, {"id": "b", "video": "clip2.mp4"}]
|
||||
samples = load_samples(data, split="RTD_OCR")
|
||||
assert len(samples) == 2
|
||||
assert all(s.split == "RTD_OCR" for s in samples)
|
||||
assert all(s.family == "rtd" for s in samples)
|
||||
|
||||
|
||||
def test_audio_feature_decode_false_preserves_bytes_key(tmp_path):
|
||||
# When a parquet file carries an Audio feature column, the loader
|
||||
# must recast it with decode=False so rows contain the ``bytes`` key
|
||||
# instead of ``array``/``sampling_rate``.
|
||||
pytest.importorskip("datasets")
|
||||
pytest.importorskip("pyarrow")
|
||||
import datasets
|
||||
|
||||
audio_bytes = b"\x00\x01\x02"
|
||||
ds = datasets.Dataset.from_dict(
|
||||
{
|
||||
"id": ["a"],
|
||||
"question_text": ["Hello"],
|
||||
"task_type": ["pr_correction"],
|
||||
"audio": [{"bytes": audio_bytes, "path": None}],
|
||||
},
|
||||
features=datasets.Features(
|
||||
{
|
||||
"id": datasets.Value("string"),
|
||||
"question_text": datasets.Value("string"),
|
||||
"task_type": datasets.Value("string"),
|
||||
"audio": datasets.Audio(sampling_rate=16000),
|
||||
}
|
||||
),
|
||||
)
|
||||
parquet = tmp_path / "audio_sample.parquet"
|
||||
ds.to_parquet(str(parquet))
|
||||
|
||||
samples = load_samples(str(parquet))
|
||||
assert len(samples) == 1
|
||||
raw = samples[0].raw
|
||||
assert raw is not None
|
||||
audio_val = raw.get("audio")
|
||||
assert audio_val is not None
|
||||
# The Audio feature was decoded with decode=False, so the struct
|
||||
# keeps the ``bytes`` key that materialize_media needs.
|
||||
assert isinstance(audio_val, dict) and "bytes" in audio_val
|
||||
assert audio_val["bytes"] == audio_bytes
|
||||
assert "path" in audio_val
|
||||
@@ -0,0 +1,67 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.benchmarks.duplex.omni_duplex_eval_media import materialize_media
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_dict_with_bytes_materializes_a_real_file(tmp_path):
|
||||
# A Hugging Face Audio/Video feature value {"bytes": ..., "path": ...} used
|
||||
# to short-circuit on ``path`` (the artifact name inside the dataset, which
|
||||
# does not exist on disk), so downstream ffmpeg saw a missing file. Inline
|
||||
# bytes must win and be written to a real file whose content equals bytes.
|
||||
payload = b"\x00\xff\x01RIFFfake-content"
|
||||
resolved = materialize_media({"bytes": payload, "path": "question_audio.wav"}, tmp_path, "s1_question", ".wav")
|
||||
assert resolved.is_file()
|
||||
assert resolved.read_bytes() == payload
|
||||
assert resolved.parent == tmp_path
|
||||
assert resolved.name == "s1_question.wav"
|
||||
|
||||
|
||||
def test_dict_extension_comes_from_path_safely(tmp_path):
|
||||
payload = b"flac-ish-bytes"
|
||||
resolved = materialize_media({"bytes": payload, "path": "audio/clip.FLAC"}, tmp_path, "s1", ".wav")
|
||||
assert resolved.name == "s1.flac"
|
||||
assert resolved.read_bytes() == payload
|
||||
# A path with no usable extension falls back to the caller's suffix.
|
||||
resolved = materialize_media({"bytes": payload, "path": "audio/clip"}, tmp_path, "s2", ".wav")
|
||||
assert resolved.name == "s2.wav"
|
||||
assert resolved.read_bytes() == payload
|
||||
|
||||
|
||||
def test_materialization_is_idempotent_and_updated(tmp_path):
|
||||
payload = b"abc"
|
||||
first = materialize_media({"bytes": payload, "path": "q.wav"}, tmp_path, "s", ".wav")
|
||||
second = materialize_media({"bytes": payload, "path": "q.wav"}, tmp_path, "s", ".wav")
|
||||
assert first == second
|
||||
assert first.read_bytes() == payload
|
||||
replacement = b"xyz"
|
||||
third = materialize_media({"bytes": replacement, "path": "q.wav"}, tmp_path, "s", ".wav")
|
||||
assert third == first
|
||||
assert third.read_bytes() == replacement
|
||||
|
||||
|
||||
def test_str_path_and_path_only_mapping_returned_as_is(tmp_path):
|
||||
existing = tmp_path / "clip.mp4"
|
||||
existing.write_bytes(b"video")
|
||||
assert materialize_media(str(existing), tmp_path, "s", ".mp4") == existing
|
||||
assert materialize_media({"path": str(existing)}, tmp_path, "s", ".mp4") == existing
|
||||
|
||||
|
||||
def test_raw_bytes_value_is_written(tmp_path):
|
||||
payload = b"raw-media"
|
||||
resolved = materialize_media(payload, tmp_path, "r", ".bin")
|
||||
assert resolved.name == "r.bin"
|
||||
assert resolved.read_bytes() == payload
|
||||
|
||||
|
||||
def test_unresolvable_value_raises_clear_error(tmp_path):
|
||||
with pytest.raises(ValueError, match="cannot materialize"):
|
||||
materialize_media(None, tmp_path, "s", ".wav")
|
||||
with pytest.raises(ValueError, match="cannot materialize"):
|
||||
materialize_media({"path": None}, tmp_path, "s", ".wav")
|
||||
@@ -135,6 +135,7 @@ class MockResponse:
|
||||
async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch):
|
||||
class FakeRealtimeClient:
|
||||
last_instance = None
|
||||
instances: list = []
|
||||
|
||||
def __init__(self, url):
|
||||
assert url == "ws://localhost:8000/v1/realtime?duplex=1"
|
||||
@@ -143,6 +144,9 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch)
|
||||
self.sent = []
|
||||
self.response_count = 0
|
||||
self.ack_count = 0
|
||||
self.silence_seconds = []
|
||||
self.closed = 0
|
||||
type(self).instances.append(self)
|
||||
type(self).last_instance = self
|
||||
|
||||
async def __aenter__(self):
|
||||
@@ -155,10 +159,22 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch)
|
||||
assert model == "openbmb/MiniCPM-o-4_5"
|
||||
self.configure_kwargs = kwargs
|
||||
|
||||
async def stream_silence(self, *, seconds, chunk_ms=200, until=None):
|
||||
# A model-native session speaks off its seeded context once audio
|
||||
# units arrive; the silence itself carries no content. The probe
|
||||
# stops the silence as soon as the turn settles.
|
||||
self.silence_seconds.append(seconds)
|
||||
self._emit_response()
|
||||
assert until is not None and until()
|
||||
return 1.0
|
||||
|
||||
async def send(self, event):
|
||||
self.sent.append(event)
|
||||
if event["type"] != "response.create":
|
||||
return
|
||||
self._emit_response()
|
||||
|
||||
def _emit_response(self):
|
||||
self.response_count += 1
|
||||
response_id = f"resp-{self.response_count}"
|
||||
now = time.monotonic()
|
||||
@@ -210,7 +226,7 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch)
|
||||
self.ack_count += 1
|
||||
|
||||
async def close_session(self, **_kwargs):
|
||||
return None
|
||||
self.closed += 1
|
||||
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.benchmarks.patch.patch._RealtimeTTSProbe",
|
||||
@@ -239,30 +255,27 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch)
|
||||
session=None,
|
||||
)
|
||||
|
||||
client = FakeRealtimeClient.last_instance
|
||||
assert client.configure_kwargs["native_duplex"] is False
|
||||
assert client.configure_kwargs["extra_body"] == {
|
||||
"ref_audio": "data:audio/wav;base64,AAAA",
|
||||
"return_stage_metrics": True,
|
||||
}
|
||||
assert client.sent == [
|
||||
event
|
||||
for index in range(4)
|
||||
for event in (
|
||||
{
|
||||
"type": "conversation.item.create",
|
||||
"item": {
|
||||
"type": "message",
|
||||
"role": "user",
|
||||
"content": [{"type": "input_text", "text": f"text {index}"}],
|
||||
},
|
||||
},
|
||||
{"type": "response.create"},
|
||||
)
|
||||
]
|
||||
assert client.ack_count == 4
|
||||
# One session per utterance: a model-native duplex session takes its text
|
||||
# once, in the session context, so it cannot be re-seeded for a second one.
|
||||
sessions = FakeRealtimeClient.instances
|
||||
assert len(sessions) == 4
|
||||
for index, client in enumerate(sessions):
|
||||
assert "native_duplex" not in client.configure_kwargs
|
||||
assert client.configure_kwargs["extra_body"] == {
|
||||
"ref_audio": "data:audio/wav;base64,AAAA",
|
||||
"return_stage_metrics": True,
|
||||
"duplex_initial_user_text": f"text {index}",
|
||||
"force_listen_count": 0,
|
||||
}
|
||||
# The target text rides the session context, never a conversation item:
|
||||
# a text-only response.create is rejected by a model-native session.
|
||||
assert client.sent == []
|
||||
assert client.silence_seconds and all(seconds > 0 for seconds in client.silence_seconds)
|
||||
assert client.ack_count == 1
|
||||
assert client.closed == 1
|
||||
|
||||
assert output.success is True
|
||||
assert output.generated_text == "turn 1 turn 2 turn 3 turn 4"
|
||||
assert output.generated_text == "turn 1 turn 1 turn 1 turn 1"
|
||||
assert output.audio_duration == pytest.approx(0.4)
|
||||
assert output.ttft > 0
|
||||
assert output.audio_ttfp > output.ttft
|
||||
@@ -281,11 +294,12 @@ async def test_seed_tts_realtime_duplex_exports_per_request_metrics(monkeypatch)
|
||||
"session_id": session_id,
|
||||
"request_index": index,
|
||||
"utterance_id": f"utt-{index}",
|
||||
"response_id": f"resp-{index + 1}",
|
||||
# Response ids are per session, and each utterance now gets its own.
|
||||
"response_id": "resp-1",
|
||||
"source": "client_monotonic_receive",
|
||||
"measurement_origin": {
|
||||
"ttft": "conversation.item.create client send to first non-empty text delta",
|
||||
"ttfp": "conversation.item.create client send to first audio packet",
|
||||
"ttft": "first silence append client send to first non-empty text delta",
|
||||
"ttfp": "first silence append client send to first audio packet",
|
||||
"rtf": "request-start-to-last-audio receive time divided by emitted audio duration",
|
||||
"tpot": "Stage-0 engine mean time per output token",
|
||||
},
|
||||
|
||||
@@ -467,6 +467,27 @@ def test_tolerated_playback_ack_rejection_is_a_warning_not_a_failure():
|
||||
oi._raise_if_session_terminated(collector, 0, warnings=warnings)
|
||||
|
||||
|
||||
def test_our_own_close_is_expected_even_though_the_server_stamps_a_reason():
|
||||
"""``session.close`` is answered with ``session.closed`` carrying ``client_close``.
|
||||
|
||||
The guard exists to catch a session that ended for a reason we did not ask
|
||||
for. A close we requested is not that, whether or not the server names it.
|
||||
"""
|
||||
collector = _collector(({"type": "session.closed", "reason": "client_close"}, 1.0))
|
||||
oi._raise_if_session_terminated(collector, 0, explicit_close_from=0)
|
||||
|
||||
collector = _collector(({"type": "session.closed", "event": {"reason": "client_close"}}, 1.0))
|
||||
oi._raise_if_session_terminated(collector, 0, explicit_close_from=0)
|
||||
|
||||
collector = _collector(({"type": "session.closed", "reason": "disconnect"}, 1.0))
|
||||
with pytest.raises(RuntimeError, match="Unexpected session.closed: disconnect"):
|
||||
oi._raise_if_session_terminated(collector, 0, explicit_close_from=0)
|
||||
|
||||
collector = _collector(({"type": "session.expired", "reason": "timeout"}, 1.0))
|
||||
with pytest.raises(RuntimeError, match="session.expired: timeout"):
|
||||
oi._raise_if_session_terminated(collector, 0, explicit_close_from=0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("event", "match"),
|
||||
[
|
||||
@@ -561,8 +582,8 @@ class _RealtimeClient(oi._RealtimeSession):
|
||||
|
||||
instances: list[_RealtimeClient] = []
|
||||
|
||||
def __init__(self, config: oi.OmniInteractBenchmarkConfig, session_id: str, reference_audio: str):
|
||||
super().__init__(config, session_id, reference_audio)
|
||||
def __init__(self, config: oi.OmniInteractBenchmarkConfig, reference_audio: str):
|
||||
super().__init__(config, reference_audio)
|
||||
self.acks: list[tuple[str, int]] = []
|
||||
self.instances.append(self)
|
||||
|
||||
@@ -621,7 +642,7 @@ async def test_public_runner_executes_one_prepared_session(tmp_path: Path, monke
|
||||
assert "autostart=0" in _RealtimeClient.instances[-1].url
|
||||
session_config = _RealtimeClient.instances[-1].session_config
|
||||
assert session_config.extra_body["custom"] == "value"
|
||||
assert session_config.extra_body["native_duplex"] is True
|
||||
assert "native_duplex" not in session_config.extra_body
|
||||
assert session_config.ref_audio == "data:audio/wav;base64,ref"
|
||||
acks = _RealtimeClient.instances[-1].acks
|
||||
# Cumulative acks for the one response: an optional 0 ms checkpoint the
|
||||
|
||||
@@ -0,0 +1,179 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""The Seed-TTS Realtime probe queues for a duplex session slot instead of failing."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.benchmarks.patch import patch as bench_patch
|
||||
from vllm_omni.clients import duplex as duplex_client
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class _AdmissionGate:
|
||||
"""Stands in for ``DuplexClient``: refuses admission ``refusals`` times, then admits."""
|
||||
|
||||
instances: list[_AdmissionGate] = []
|
||||
refusals = 0
|
||||
refusal_code = "resource_exhausted"
|
||||
|
||||
def __init__(self, url: str, **kwargs: object) -> None:
|
||||
del url, kwargs
|
||||
self.entered = False
|
||||
type(self).instances.append(self)
|
||||
|
||||
async def __aenter__(self) -> _AdmissionGate:
|
||||
cls = type(self)
|
||||
if cls.refusals > 0:
|
||||
cls.refusals -= 1
|
||||
raise duplex_client.DuplexProtocolError("duplex_session_capacity_exhausted: limit=4", code=cls.refusal_code)
|
||||
self.entered = True
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb) -> None:
|
||||
del exc_type, exc, tb
|
||||
|
||||
async def events(self): # the collector's consume loop ends at once
|
||||
return
|
||||
yield # pragma: no cover
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def gate(monkeypatch: pytest.MonkeyPatch) -> type[_AdmissionGate]:
|
||||
_AdmissionGate.instances = []
|
||||
_AdmissionGate.refusals = 0
|
||||
_AdmissionGate.refusal_code = "resource_exhausted"
|
||||
monkeypatch.setattr(duplex_client, "DuplexClient", _AdmissionGate)
|
||||
return _AdmissionGate
|
||||
|
||||
|
||||
async def _configure(probe: bench_patch._RealtimeTTSProbe) -> None:
|
||||
try:
|
||||
await probe.configure("test-model", auto_response=True)
|
||||
finally:
|
||||
await probe.__aexit__(None, None, None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_configure_waits_for_a_free_session_slot(gate) -> None:
|
||||
"""Two refusals, then a slot: the request is served, not counted as failed."""
|
||||
gate.refusals = 2
|
||||
probe = bench_patch._RealtimeTTSProbe("ws://test/v1/realtime?duplex=1")
|
||||
await _configure(probe)
|
||||
assert len(gate.instances) == 3
|
||||
assert gate.instances[-1].entered is True
|
||||
assert probe._client is gate.instances[-1]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_configure_raises_other_rejections_at_once(gate) -> None:
|
||||
gate.refusals = 1
|
||||
gate.refusal_code = "unsupported_turn_detection"
|
||||
probe = bench_patch._RealtimeTTSProbe("ws://test/v1/realtime?duplex=1")
|
||||
with pytest.raises(duplex_client.DuplexProtocolError):
|
||||
await _configure(probe)
|
||||
assert len(gate.instances) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_configure_gives_up_when_the_slot_wait_budget_is_spent(gate, monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(bench_patch._RealtimeTTSProbe, "_SESSION_SLOT_WAIT_S", 0.0)
|
||||
gate.refusals = 5
|
||||
probe = bench_patch._RealtimeTTSProbe("ws://test/v1/realtime?duplex=1")
|
||||
with pytest.raises(duplex_client.DuplexProtocolError) as excinfo:
|
||||
await _configure(probe)
|
||||
assert excinfo.value.code == "resource_exhausted"
|
||||
assert len(gate.instances) == 1
|
||||
|
||||
|
||||
class _SilenceSink:
|
||||
"""Stands in for a configured ``DuplexClient``: records every appended chunk."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.config = duplex_client.SessionConfig()
|
||||
self.chunks: list[bytes] = []
|
||||
|
||||
async def append_audio(self, pcm: bytes, *, is_speech: bool | None = None, video_frames=None) -> None:
|
||||
del video_frames
|
||||
assert is_speech is False
|
||||
self.chunks.append(pcm)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_silence_stops_at_the_turn_end(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Silence is fed until the turn settles, not for the whole cap."""
|
||||
probe = bench_patch._RealtimeTTSProbe("ws://test/v1/realtime?duplex=1")
|
||||
sink = _SilenceSink()
|
||||
probe._client = sink
|
||||
|
||||
async def no_wait(_: float) -> None:
|
||||
return
|
||||
|
||||
monkeypatch.setattr(bench_patch.asyncio, "sleep", no_wait)
|
||||
streamed = await probe.stream_silence(seconds=12.0, chunk_ms=200, until=lambda: len(sink.chunks) >= 3)
|
||||
assert len(sink.chunks) == 3
|
||||
assert streamed == pytest.approx(0.6)
|
||||
assert all(len(chunk) == sink.config.input_audio.byte_count(200) for chunk in sink.chunks)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_silence_runs_to_the_cap_when_the_turn_never_settles(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
probe = bench_patch._RealtimeTTSProbe("ws://test/v1/realtime?duplex=1")
|
||||
sink = _SilenceSink()
|
||||
probe._client = sink
|
||||
|
||||
async def no_wait(_: float) -> None:
|
||||
return
|
||||
|
||||
monkeypatch.setattr(bench_patch.asyncio, "sleep", no_wait)
|
||||
streamed = await probe.stream_silence(seconds=1.0, chunk_ms=200, until=lambda: False)
|
||||
assert len(sink.chunks) == 5
|
||||
assert streamed == pytest.approx(1.0)
|
||||
|
||||
|
||||
class _TurnEvents:
|
||||
def __init__(self, response_ids: list[str], audio: dict[str, bytes], text: dict[str, str] | None = None) -> None:
|
||||
self.response_ids = response_ids
|
||||
self._audio = audio
|
||||
self._text = text or {}
|
||||
|
||||
def audio_bytes(self, response_id: str) -> bytes:
|
||||
return self._audio.get(response_id, b"")
|
||||
|
||||
def response_text(self, response_id: str) -> str:
|
||||
return self._text.get(response_id, "")
|
||||
|
||||
|
||||
def test_turn_response_is_the_first_one_with_audio_even_if_the_model_spoke_again(caplog) -> None:
|
||||
events = _TurnEvents(["r0", "r1", "r2"], {"r1": b"\x01\x00", "r2": b"\x02\x00"})
|
||||
with caplog.at_level("WARNING", logger=bench_patch.logger.name):
|
||||
assert bench_patch._seed_tts_turn_response_id(events, 0, 0) == "r1"
|
||||
assert "spoke again" in caplog.text
|
||||
|
||||
|
||||
def test_turn_response_ignores_responses_from_before_the_turn() -> None:
|
||||
events = _TurnEvents(["r0", "r1"], {"r0": b"\x01\x00", "r1": b"\x02\x00"})
|
||||
assert bench_patch._seed_tts_turn_response_id(events, 1, 0) == "r1"
|
||||
|
||||
|
||||
def test_turn_without_audio_is_an_error() -> None:
|
||||
events = _TurnEvents(["r0"], {})
|
||||
with pytest.raises(RuntimeError, match="no audio response"):
|
||||
bench_patch._seed_tts_turn_response_id(events, 0, 0)
|
||||
|
||||
|
||||
def test_stall_report_says_the_model_never_answered() -> None:
|
||||
events = _TurnEvents(["r0"], {"r0": b"\x01\x00"})
|
||||
report = bench_patch._seed_tts_turn_stall_report(events, 1, 3, 30.0)
|
||||
assert "turn 3 never started a response" in report
|
||||
assert "30.0s of silence" in report
|
||||
|
||||
|
||||
def test_stall_report_lists_a_response_that_never_finished() -> None:
|
||||
events = _TurnEvents(["r1"], {"r1": b"\x01\x00\x02\x00"}, {"r1": "hello"})
|
||||
report = bench_patch._seed_tts_turn_stall_report(events, 0, 0, 4.2)
|
||||
assert "started 1 response(s) after 4.2s of silence" in report
|
||||
assert "r1 (4 audio bytes, text 'hello')" in report
|
||||
+425
-114
@@ -1,6 +1,12 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Unit tests for the public duplex client (fake transport, no server)."""
|
||||
"""Unit tests for the public duplex clients (fake transports, no server).
|
||||
|
||||
``DuplexClient`` runs over a fake websocket; ``InlineDuplexClient`` over a
|
||||
fake ``DuplexOmni`` whose handle records typed commands and replays typed
|
||||
events. The demux / handshake tests that only need "feed a server event" are
|
||||
parameterized over both clients.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -8,6 +14,8 @@ import asyncio
|
||||
import base64
|
||||
import json
|
||||
import wave
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
from urllib.parse import parse_qs, urlsplit
|
||||
|
||||
import pytest
|
||||
@@ -16,8 +24,11 @@ from vllm_omni.clients.duplex import (
|
||||
AudioFormat,
|
||||
ConnectionResumed,
|
||||
DuplexClient,
|
||||
DuplexClientBase,
|
||||
DuplexConnectionError,
|
||||
DuplexProtocolError,
|
||||
DuplexSessionClosedError,
|
||||
ErrorEvent,
|
||||
EventCollector,
|
||||
ReconnectPolicy,
|
||||
SessionConfig,
|
||||
@@ -31,17 +42,26 @@ from vllm_omni.clients.duplex import (
|
||||
summarize_session_request_metrics,
|
||||
write_pcm16_wav,
|
||||
)
|
||||
from vllm_omni.clients.inline_duplex import InlineDuplexClient
|
||||
from vllm_omni.engine.duplex import commands as duplex_commands
|
||||
from vllm_omni.engine.duplex import events as duplex_events
|
||||
from vllm_omni.engine.duplex.messages import DuplexSessionError
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
SESSION_ID = "duplex-0123456789abcdef"
|
||||
# session.created is not journaled by the server, so it carries no server_event_seq.
|
||||
SESSION_CREATED = {
|
||||
"type": "session.created",
|
||||
"session": {"session_id": "sess-1"},
|
||||
"incarnation": 0,
|
||||
"session": {"id": SESSION_ID, "model": "test-model"},
|
||||
"attachment_generation": 1,
|
||||
"resume_token": "tok-1",
|
||||
"server_event_seq": 1,
|
||||
}
|
||||
SESSION_CLOSED = {"type": "session.closed", "session_id": "sess-1", "server_event_seq": 99}
|
||||
SESSION_CLOSED = {"type": "session.closed", "session_id": SESSION_ID, "reason": "client_close", "server_event_seq": 99}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fake websocket transport
|
||||
|
||||
|
||||
class FakeSocket:
|
||||
@@ -88,14 +108,139 @@ def make_client(*sockets: FakeSocket, **kwargs) -> tuple[DuplexClient, list[str]
|
||||
return client, calls
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fake DuplexOmni for the inline client
|
||||
|
||||
|
||||
class _WireEvent:
|
||||
"""A typed-event stand-in that renders a fixed wire payload."""
|
||||
|
||||
def __init__(self, raw: dict[str, object]) -> None:
|
||||
self.raw = raw
|
||||
|
||||
@property
|
||||
def is_terminal(self) -> bool:
|
||||
return self.raw.get("type") in {"session.closed", "session.expired"}
|
||||
|
||||
def to_realtime(self) -> dict[str, object]:
|
||||
return dict(self.raw)
|
||||
|
||||
|
||||
class FakeInlineHandle:
|
||||
def __init__(self, session_id: str) -> None:
|
||||
self.session_id = session_id
|
||||
self.closed = False
|
||||
self.close_calls = 0
|
||||
self.commands: list[duplex_commands.DuplexCommand] = []
|
||||
self.consumer_active = False
|
||||
self._outbox: asyncio.Queue = asyncio.Queue()
|
||||
|
||||
def deliver(self, event: Any) -> None:
|
||||
self._outbox.put_nowait(event)
|
||||
|
||||
def feed(self, raw: dict[str, object]) -> None:
|
||||
self.deliver(_WireEvent(raw))
|
||||
|
||||
async def submit(self, command: duplex_commands.DuplexCommand) -> None:
|
||||
if self.closed:
|
||||
raise DuplexSessionError("closed", code="session_closed", session_id=self.session_id)
|
||||
self.commands.append(command)
|
||||
|
||||
async def events(self):
|
||||
self.consumer_active = True
|
||||
try:
|
||||
while True:
|
||||
event = await self._outbox.get()
|
||||
yield event
|
||||
if event.is_terminal:
|
||||
return
|
||||
finally:
|
||||
self.consumer_active = False
|
||||
|
||||
async def close(self, *, reason: str = "client_close", timeout: float | None = None) -> None:
|
||||
self.close_calls += 1
|
||||
if self.closed:
|
||||
return
|
||||
self.closed = True
|
||||
self.deliver(duplex_events.SessionClosed(session_id=self.session_id, reason=reason))
|
||||
|
||||
|
||||
class FakeInlineOmni:
|
||||
def __init__(self) -> None:
|
||||
self.opened: list[dict[str, object]] = []
|
||||
self.handles: list[FakeInlineHandle] = []
|
||||
|
||||
async def open_session(self, config: dict[str, object]) -> FakeInlineHandle:
|
||||
self.opened.append(dict(config))
|
||||
handle = FakeInlineHandle(SESSION_ID)
|
||||
self.handles.append(handle)
|
||||
handle.deliver(
|
||||
duplex_events.SessionCreated(
|
||||
session_id=SESSION_ID,
|
||||
session={"id": SESSION_ID, "model": config.get("model")},
|
||||
)
|
||||
)
|
||||
return handle
|
||||
|
||||
|
||||
def make_inline_client(**kwargs) -> tuple[InlineDuplexClient, FakeInlineOmni]:
|
||||
omni = FakeInlineOmni()
|
||||
kwargs.setdefault("handshake_timeout_s", 5.0)
|
||||
return InlineDuplexClient(omni, model="test-model", **kwargs), omni
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# One rig over both clients: ``feed`` a server event, drive ``client``
|
||||
|
||||
|
||||
class Rig:
|
||||
def __init__(self, client: DuplexClientBase, feed: Callable[[dict[str, object]], None]) -> None:
|
||||
self.client = client
|
||||
self.feed = feed
|
||||
|
||||
|
||||
def _websocket_rig(*, handshake: dict[str, object] | None = SESSION_CREATED) -> Rig:
|
||||
sock = FakeSocket()
|
||||
if handshake is not None:
|
||||
sock.feed(handshake)
|
||||
client, _ = make_client(sock)
|
||||
return Rig(client, sock.feed)
|
||||
|
||||
|
||||
def _inline_rig(*, handshake: dict[str, object] | None = SESSION_CREATED) -> Rig:
|
||||
client, omni = make_inline_client()
|
||||
original_open = omni.open_session
|
||||
|
||||
async def open_session(config: dict[str, object]) -> FakeInlineHandle:
|
||||
handle = await original_open(config)
|
||||
if handshake is not SESSION_CREATED:
|
||||
# Replace the default session.created with the requested first event.
|
||||
handle._outbox = asyncio.Queue()
|
||||
if handshake is not None:
|
||||
handle.feed(handshake)
|
||||
return handle
|
||||
|
||||
omni.open_session = open_session # type: ignore[method-assign]
|
||||
|
||||
def feed(raw: dict[str, object]) -> None:
|
||||
omni.handles[-1].feed(raw)
|
||||
|
||||
return Rig(client, feed)
|
||||
|
||||
|
||||
@pytest.fixture(params=["websocket", "inline"])
|
||||
def make_rig(request: pytest.FixtureRequest) -> Callable[..., Rig]:
|
||||
return _websocket_rig if request.param == "websocket" else _inline_rig
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# SessionConfig
|
||||
|
||||
|
||||
def test_session_config_payload_defaults():
|
||||
payload = SessionConfig().to_session_payload(model="m", session_id="s")
|
||||
payload = SessionConfig().to_session_payload(model="m")
|
||||
assert payload["model"] == "m"
|
||||
assert payload["session_id"] == "s"
|
||||
assert "session_id" not in payload
|
||||
assert payload["modalities"] == ["audio", "text"]
|
||||
assert payload["input_audio_format"] == "pcm16"
|
||||
assert payload["output_audio_format"] == "pcm16"
|
||||
@@ -137,19 +282,59 @@ def test_audio_format_math():
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Handshake and lifecycle
|
||||
# Handshake and lifecycle (both clients)
|
||||
|
||||
|
||||
async def test_handshake_adopts_session_and_acks():
|
||||
@pytest.mark.asyncio
|
||||
async def test_handshake_adopts_server_allocated_session_id(make_rig):
|
||||
rig = make_rig()
|
||||
async with rig.client as client:
|
||||
assert client.session_id == SESSION_ID
|
||||
assert client.session_info == {"id": SESSION_ID, "model": "test-model"}
|
||||
rig.feed(SESSION_CLOSED)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handshake_error_raises_protocol_error(make_rig):
|
||||
rig = make_rig(handshake={"type": "error", "error": {"code": "unsupported_audio_format", "message": "bad"}})
|
||||
with pytest.raises(DuplexProtocolError) as excinfo:
|
||||
async with rig.client:
|
||||
pass
|
||||
assert excinfo.value.code == "unsupported_audio_format"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_session_expired_raises_from_event_stream(make_rig):
|
||||
rig = make_rig()
|
||||
async with rig.client as client:
|
||||
rig.feed(
|
||||
{"type": "session.expired", "session_id": SESSION_ID, "reason": "lease_expired", "server_event_seq": 2}
|
||||
)
|
||||
with pytest.raises(DuplexSessionClosedError) as excinfo:
|
||||
async for _ in client.events():
|
||||
pass
|
||||
assert "lease_expired" in excinfo.value.reason
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Websocket transport specifics
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_websocket_handshake_sends_session_update_and_acks_sequenced_events():
|
||||
sock = FakeSocket()
|
||||
sock.feed(SESSION_CREATED)
|
||||
client, calls = make_client(sock)
|
||||
async with client:
|
||||
assert client.session_id == "sess-1"
|
||||
assert client.session_id == SESSION_ID
|
||||
assert client.resume_token == "tok-1"
|
||||
assert calls == ["ws://test-host:8099/v1/realtime?duplex=1&model=test-model&autostart=0"]
|
||||
assert sock.sent[0]["type"] == "session.update"
|
||||
assert sock.sent[0]["session"]["model"] == "test-model"
|
||||
assert "session_id" not in sock.sent[0]["session"]
|
||||
# session.created is unjournaled (no seq): nothing to ack yet.
|
||||
assert _acks(sock) == []
|
||||
sock.feed({"type": "response.listen", "server_event_seq": 1})
|
||||
await _drain(lambda: {"type": "session.event_ack", "server_event_seq": 1} in _acks(sock))
|
||||
close_task = asyncio.create_task(client.close())
|
||||
await _drain(lambda: "session.close" in sock.sent_types())
|
||||
@@ -158,28 +343,7 @@ async def test_handshake_adopts_session_and_acks():
|
||||
assert "session.close" in sock.sent_types()
|
||||
|
||||
|
||||
async def test_handshake_error_raises_protocol_error():
|
||||
sock = FakeSocket()
|
||||
sock.feed({"type": "error", "error": {"code": "unsupported_audio_format", "message": "bad"}})
|
||||
client, _ = make_client(sock)
|
||||
with pytest.raises(DuplexProtocolError) as excinfo:
|
||||
async with client:
|
||||
pass
|
||||
assert excinfo.value.code == "unsupported_audio_format"
|
||||
|
||||
|
||||
async def test_session_expired_raises_from_event_stream():
|
||||
sock = FakeSocket()
|
||||
sock.feed(SESSION_CREATED)
|
||||
client, _ = make_client(sock)
|
||||
async with client:
|
||||
sock.feed({"type": "session.expired", "reason": "lease_expired", "server_event_seq": 2})
|
||||
with pytest.raises(DuplexSessionClosedError) as excinfo:
|
||||
async for _ in client.events():
|
||||
pass
|
||||
assert "lease_expired" in excinfo.value.reason
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resync_required_drops_resume_credential():
|
||||
sock = FakeSocket()
|
||||
sock.feed(SESSION_CREATED)
|
||||
@@ -187,7 +351,7 @@ async def test_resync_required_drops_resume_credential():
|
||||
async with client:
|
||||
assert client.resume_token == "tok-1"
|
||||
stream = client.events()
|
||||
sock.feed({"type": "session.resync_required", "session_id": "sess-1", "server_event_seq": 2})
|
||||
sock.feed({"type": "session.resync_required", "session_id": SESSION_ID, "server_event_seq": 2})
|
||||
async for event in stream:
|
||||
if event.type == "session.resync_required":
|
||||
break
|
||||
@@ -202,9 +366,10 @@ async def test_resync_required_drops_resume_credential():
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Input events
|
||||
# Input events (websocket wire shape)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_append_audio_tracks_cumulative_end_ms():
|
||||
sock = FakeSocket()
|
||||
sock.feed(SESSION_CREATED)
|
||||
@@ -224,6 +389,7 @@ async def test_append_audio_tracks_cumulative_end_ms():
|
||||
sock.feed(SESSION_CLOSED)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_append_audio_strips_video_frame_data_url_prefix():
|
||||
sock = FakeSocket()
|
||||
sock.feed(SESSION_CREATED)
|
||||
@@ -238,6 +404,7 @@ async def test_append_audio_strips_video_frame_data_url_prefix():
|
||||
sock.feed(SESSION_CLOSED)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_pcm_chunking():
|
||||
sock = FakeSocket()
|
||||
sock.feed(SESSION_CREATED)
|
||||
@@ -249,6 +416,7 @@ async def test_stream_pcm_chunking():
|
||||
sock.feed(SESSION_CLOSED)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_interruption_primitives_send_documented_events():
|
||||
sock = FakeSocket()
|
||||
sock.feed(SESSION_CREATED)
|
||||
@@ -262,22 +430,21 @@ async def test_interruption_primitives_send_documented_events():
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Response demultiplexing
|
||||
# Response demultiplexing (both clients)
|
||||
|
||||
|
||||
def _b64(data: bytes) -> str:
|
||||
return base64.b64encode(data).decode("ascii")
|
||||
|
||||
|
||||
async def test_response_handle_flow():
|
||||
sock = FakeSocket()
|
||||
sock.feed(SESSION_CREATED)
|
||||
@pytest.mark.asyncio
|
||||
async def test_response_handle_flow(make_rig):
|
||||
rig = make_rig()
|
||||
chunk = b"\x00\x01" * 2400 # 100 ms of pcm16 @ 24 kHz
|
||||
client, _ = make_client(sock)
|
||||
async with client:
|
||||
sock.feed({"type": "response.created", "response": {"id": "resp-1"}, "server_event_seq": 2})
|
||||
sock.feed({"type": "response.speak", "response_id": "resp-1", "server_event_seq": 3})
|
||||
sock.feed(
|
||||
async with rig.client as client:
|
||||
rig.feed({"type": "response.created", "response": {"id": "resp-1"}, "server_event_seq": 2})
|
||||
rig.feed({"type": "response.speak", "response_id": "resp-1", "server_event_seq": 3})
|
||||
rig.feed(
|
||||
{
|
||||
"type": "response.output_audio.delta",
|
||||
"response_id": "resp-1",
|
||||
@@ -286,7 +453,7 @@ async def test_response_handle_flow():
|
||||
"server_event_seq": 4,
|
||||
}
|
||||
)
|
||||
sock.feed(
|
||||
rig.feed(
|
||||
{
|
||||
"type": "response.output_audio_transcript.delta",
|
||||
"response_id": "resp-1",
|
||||
@@ -294,7 +461,7 @@ async def test_response_handle_flow():
|
||||
"server_event_seq": 5,
|
||||
}
|
||||
)
|
||||
sock.feed({"type": "response.done", "response_id": "resp-1", "server_event_seq": 6})
|
||||
rig.feed({"type": "response.done", "response_id": "resp-1", "server_event_seq": 6})
|
||||
|
||||
async for response in client.responses():
|
||||
chunks = [piece async for piece in response.audio()]
|
||||
@@ -304,49 +471,46 @@ async def test_response_handle_flow():
|
||||
assert response.played_ms == pytest.approx(100.0)
|
||||
assert response.finished
|
||||
break
|
||||
sock.feed(SESSION_CLOSED)
|
||||
rig.feed(SESSION_CLOSED)
|
||||
|
||||
|
||||
async def test_listen_decision_yields_finished_silent_handle():
|
||||
sock = FakeSocket()
|
||||
sock.feed(SESSION_CREATED)
|
||||
client, _ = make_client(sock)
|
||||
async with client:
|
||||
sock.feed({"type": "response.listen", "server_event_seq": 2})
|
||||
@pytest.mark.asyncio
|
||||
async def test_listen_decision_yields_finished_silent_handle(make_rig):
|
||||
rig = make_rig()
|
||||
async with rig.client as client:
|
||||
rig.feed({"type": "response.listen", "server_event_seq": 2})
|
||||
async for response in client.responses():
|
||||
assert response.decision == "listen"
|
||||
assert response.finished
|
||||
assert [piece async for piece in response.audio()] == []
|
||||
break
|
||||
sock.feed(SESSION_CLOSED)
|
||||
rig.feed(SESSION_CLOSED)
|
||||
|
||||
|
||||
async def test_listen_terminates_active_response():
|
||||
sock = FakeSocket()
|
||||
sock.feed(SESSION_CREATED)
|
||||
client, _ = make_client(sock)
|
||||
async with client:
|
||||
sock.feed({"type": "response.created", "response": {"id": "resp-1"}, "server_event_seq": 2})
|
||||
sock.feed({"type": "response.listen", "response_id": "resp-1", "server_event_seq": 3})
|
||||
@pytest.mark.asyncio
|
||||
async def test_listen_terminates_active_response(make_rig):
|
||||
rig = make_rig()
|
||||
async with rig.client as client:
|
||||
rig.feed({"type": "response.created", "response": {"id": "resp-1"}, "server_event_seq": 2})
|
||||
rig.feed({"type": "response.listen", "response_id": "resp-1", "server_event_seq": 3})
|
||||
async for response in client.responses():
|
||||
await response.wait(timeout_s=5.0)
|
||||
assert response.decision == "listen"
|
||||
break
|
||||
sock.feed(SESSION_CLOSED)
|
||||
rig.feed(SESSION_CLOSED)
|
||||
|
||||
|
||||
async def test_id_less_listen_never_closes_an_in_flight_response():
|
||||
@pytest.mark.asyncio
|
||||
async def test_id_less_listen_never_closes_an_in_flight_response(make_rig):
|
||||
# An id-less listen is a standalone decision beat (e.g. a silence-skip
|
||||
# while a response streams); it must surface as its own finished handle
|
||||
# and leave the in-flight response to its real terminal.
|
||||
sock = FakeSocket()
|
||||
sock.feed(SESSION_CREATED)
|
||||
client, _ = make_client(sock)
|
||||
async with client:
|
||||
sock.feed({"type": "response.created", "response": {"id": "resp-1"}, "server_event_seq": 2})
|
||||
sock.feed({"type": "response.speak", "response_id": "resp-1", "server_event_seq": 3})
|
||||
sock.feed({"type": "response.listen", "server_event_seq": 4})
|
||||
sock.feed({"type": "response.done", "response_id": "resp-1", "server_event_seq": 5})
|
||||
rig = make_rig()
|
||||
async with rig.client as client:
|
||||
rig.feed({"type": "response.created", "response": {"id": "resp-1"}, "server_event_seq": 2})
|
||||
rig.feed({"type": "response.speak", "response_id": "resp-1", "server_event_seq": 3})
|
||||
rig.feed({"type": "response.listen", "server_event_seq": 4})
|
||||
rig.feed({"type": "response.done", "response_id": "resp-1", "server_event_seq": 5})
|
||||
seen: list[tuple[str | None, str | None, bool]] = []
|
||||
async for response in client.responses():
|
||||
await response.wait(timeout_s=5.0)
|
||||
@@ -355,20 +519,19 @@ async def test_id_less_listen_never_closes_an_in_flight_response():
|
||||
break
|
||||
assert ("resp-1", "speak", True) in seen
|
||||
assert (None, "listen", True) in seen
|
||||
sock.feed(SESSION_CLOSED)
|
||||
rig.feed(SESSION_CLOSED)
|
||||
|
||||
|
||||
async def test_terminal_listen_keeps_spoken_decision():
|
||||
@pytest.mark.asyncio
|
||||
async def test_terminal_listen_keeps_spoken_decision(make_rig):
|
||||
# A terminal listen after the response spoke is the model yielding the
|
||||
# turn; the handle must stay decision="speak" with its audio intact.
|
||||
sock = FakeSocket()
|
||||
sock.feed(SESSION_CREATED)
|
||||
rig = make_rig()
|
||||
chunk = b"\x00\x01" * 2400
|
||||
client, _ = make_client(sock)
|
||||
async with client:
|
||||
sock.feed({"type": "response.created", "response": {"id": "resp-1"}, "server_event_seq": 2})
|
||||
sock.feed({"type": "response.speak", "response_id": "resp-1", "server_event_seq": 3})
|
||||
sock.feed(
|
||||
async with rig.client as client:
|
||||
rig.feed({"type": "response.created", "response": {"id": "resp-1"}, "server_event_seq": 2})
|
||||
rig.feed({"type": "response.speak", "response_id": "resp-1", "server_event_seq": 3})
|
||||
rig.feed(
|
||||
{
|
||||
"type": "response.output_audio.delta",
|
||||
"response_id": "resp-1",
|
||||
@@ -377,24 +540,23 @@ async def test_terminal_listen_keeps_spoken_decision():
|
||||
"server_event_seq": 4,
|
||||
}
|
||||
)
|
||||
sock.feed({"type": "response.listen", "response_id": "resp-1", "server_event_seq": 5})
|
||||
rig.feed({"type": "response.listen", "response_id": "resp-1", "server_event_seq": 5})
|
||||
async for response in client.responses():
|
||||
chunks = [piece async for piece in response.audio()]
|
||||
assert chunks == [chunk]
|
||||
assert response.decision == "speak"
|
||||
assert response.finished
|
||||
break
|
||||
sock.feed(SESSION_CLOSED)
|
||||
rig.feed(SESSION_CLOSED)
|
||||
|
||||
|
||||
async def test_error_event_surfaces_on_response_path():
|
||||
@pytest.mark.asyncio
|
||||
async def test_error_event_surfaces_on_response_path(make_rig):
|
||||
# A rejected send produces no response; a consumer waiting in responses()
|
||||
# must see the rejection instead of waiting forever.
|
||||
sock = FakeSocket()
|
||||
sock.feed(SESSION_CREATED)
|
||||
client, _ = make_client(sock)
|
||||
async with client:
|
||||
sock.feed(
|
||||
rig = make_rig()
|
||||
async with rig.client as client:
|
||||
rig.feed(
|
||||
{
|
||||
"type": "error",
|
||||
"error": {"code": "input_audio_buffer_empty", "message": "empty commit"},
|
||||
@@ -405,9 +567,37 @@ async def test_error_event_surfaces_on_response_path():
|
||||
async for _ in client.responses():
|
||||
pass
|
||||
assert excinfo.value.code == "input_audio_buffer_empty"
|
||||
rig.feed(SESSION_CLOSED)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_close_after_error_reports_the_server_reason_to_a_sender():
|
||||
# A client streaming input sits in send(), not in responses(), so the
|
||||
# error text has to survive into the close marker or the session looks
|
||||
# like it ended for no reason.
|
||||
sock = FakeSocket()
|
||||
sock.feed(SESSION_CREATED)
|
||||
client, _ = make_client(sock)
|
||||
async with client:
|
||||
sock.feed(
|
||||
{
|
||||
"type": "error",
|
||||
"error": {
|
||||
"code": "runtime_data_plane_stream_failed",
|
||||
"message": "context_length_exceeded: streaming session prompt would grow to 41207 tokens",
|
||||
},
|
||||
"server_event_seq": 2,
|
||||
}
|
||||
)
|
||||
sock.feed(SESSION_CLOSED)
|
||||
await _drain(lambda: client._closed.is_set())
|
||||
with pytest.raises(DuplexSessionClosedError) as excinfo:
|
||||
await client.send({"type": "input_audio_buffer.append", "audio": ""})
|
||||
assert "runtime_data_plane_stream_failed" in str(excinfo.value)
|
||||
assert "context_length_exceeded" in str(excinfo.value)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_slow_audio_consumer_drops_oldest_instead_of_stalling():
|
||||
from vllm_omni.clients.duplex import AudioDelta, ResponseHandle
|
||||
|
||||
@@ -425,24 +615,24 @@ async def test_slow_audio_consumer_drops_oldest_instead_of_stalling():
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Resume
|
||||
# Resume (websocket only)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resume_after_transport_drop():
|
||||
first = FakeSocket()
|
||||
first.feed(SESSION_CREATED)
|
||||
second = FakeSocket()
|
||||
# Mirror the real resume activation payload
|
||||
# (entrypoints/duplex/serving.py, activation_payload_factory): it carries
|
||||
# no nested "session" object, only the identity fields.
|
||||
# the identity fields plus the (empty) public session object.
|
||||
second.feed(
|
||||
{
|
||||
"type": "session.resumed",
|
||||
"session_id": "sess-1",
|
||||
"incarnation": 0,
|
||||
"attachment_generation": 1,
|
||||
"session_id": SESSION_ID,
|
||||
"session": {},
|
||||
"attachment_generation": 2,
|
||||
"resume_token": "tok-2",
|
||||
"server_event_seq": 5,
|
||||
}
|
||||
)
|
||||
client, calls = make_client(
|
||||
@@ -453,34 +643,39 @@ async def test_resume_after_transport_drop():
|
||||
async with client:
|
||||
stream = client.events()
|
||||
|
||||
async def take_two():
|
||||
async def collect_until_resumed():
|
||||
seen = []
|
||||
async for event in stream:
|
||||
seen.append(event)
|
||||
if len(seen) == 2:
|
||||
if isinstance(event, SessionResumed):
|
||||
return seen
|
||||
|
||||
consumer = asyncio.create_task(take_two())
|
||||
consumer = asyncio.create_task(collect_until_resumed())
|
||||
await asyncio.sleep(0) # let the consumer subscribe before the drop
|
||||
await asyncio.sleep(0)
|
||||
first.feed({"type": "response.listen", "server_event_seq": 7})
|
||||
await _drain(lambda: {"type": "session.event_ack", "server_event_seq": 7} in _acks(first))
|
||||
first.feed(RuntimeError("transport dropped"))
|
||||
seen = await asyncio.wait_for(consumer, timeout=5.0)
|
||||
assert isinstance(seen[0], ConnectionResumed)
|
||||
assert isinstance(seen[1], SessionResumed)
|
||||
assert [type(event) for event in seen[-2:]] == [ConnectionResumed, SessionResumed]
|
||||
assert client.resume_token == "tok-2"
|
||||
# The superseded transport is closed, not leaked.
|
||||
assert first.closed
|
||||
# The activation payload has no session object; the info captured at
|
||||
# the original handshake must survive the resume.
|
||||
assert client.session_info == {"session_id": "sess-1"}
|
||||
assert client.session_info == {"id": SESSION_ID, "model": "test-model"}
|
||||
assert len(calls) == 2
|
||||
resume = second.sent[0]
|
||||
assert resume["type"] == "session.resume"
|
||||
assert resume["session_id"] == "sess-1"
|
||||
assert resume["incarnation"] == 0
|
||||
assert resume["resume_token"] == "tok-1"
|
||||
assert resume["last_received_server_event_seq"] == 1
|
||||
assert resume == {
|
||||
"type": "session.resume",
|
||||
"session_id": SESSION_ID,
|
||||
"resume_token": "tok-1",
|
||||
"last_received_server_event_seq": 7,
|
||||
}
|
||||
second.feed(SESSION_CLOSED)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_reconnect_policy_surfaces_closed():
|
||||
sock = FakeSocket()
|
||||
sock.feed(SESSION_CREATED)
|
||||
@@ -493,6 +688,7 @@ async def test_no_reconnect_policy_surfaces_closed():
|
||||
pass
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_heartbeat_survives_transient_send_failure():
|
||||
sock = FakeSocket()
|
||||
sock.feed(SESSION_CREATED)
|
||||
@@ -514,6 +710,7 @@ async def test_heartbeat_survives_transient_send_failure():
|
||||
sock.feed(SESSION_CLOSED)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_aexit_sends_session_close_on_error_path():
|
||||
sock = FakeSocket()
|
||||
sock.feed(SESSION_CREATED)
|
||||
@@ -534,6 +731,7 @@ def test_target_url_forces_autostart_off():
|
||||
assert query["autostart"] == ["0"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fatal_resume_error_gives_up():
|
||||
first = FakeSocket()
|
||||
first.feed(SESSION_CREATED)
|
||||
@@ -553,6 +751,121 @@ async def test_fatal_resume_error_gives_up():
|
||||
assert len(calls) == 2 # no retry after a fatal resume error
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Inline client over a fake DuplexOmni
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_inline_handshake_opens_session_with_payload_and_adopts_id():
|
||||
client, omni = make_inline_client(config=SessionConfig(ref_audio="data:audio/wav;base64,AAA="))
|
||||
async with client:
|
||||
assert omni.opened == [client.config.to_session_payload(model="test-model")]
|
||||
assert omni.opened[0]["ref_audio"] == "data:audio/wav;base64,AAA="
|
||||
assert client.session_id == SESSION_ID
|
||||
handle = omni.handles[0]
|
||||
# The session was opened with the full session object already; the
|
||||
# base handshake's session.update must not be re-submitted as a patch
|
||||
# (a ref_audio patch would be rejected by the runner).
|
||||
assert handle.commands == []
|
||||
await client.send({"type": "session.update", "session": {"instructions": "later"}})
|
||||
assert [type(command) for command in handle.commands] == [duplex_commands.UpdateSession]
|
||||
assert handle.commands[0].patch == {"instructions": "later"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_inline_client_submits_typed_commands_in_order():
|
||||
client, omni = make_inline_client()
|
||||
async with client:
|
||||
handle = omni.handles[0]
|
||||
pcm = b"\x01\x00" * 1600
|
||||
await client.append_audio(pcm, is_speech=True)
|
||||
await client.commit(create_response=True)
|
||||
await client.cancel_response("resp-1")
|
||||
await client.ack_playback(120.9, response_id="resp-1")
|
||||
assert [type(command) for command in handle.commands] == [
|
||||
duplex_commands.AppendAudio,
|
||||
duplex_commands.Commit,
|
||||
duplex_commands.CancelResponse,
|
||||
duplex_commands.AckPlayback,
|
||||
]
|
||||
append, commit, cancel, ack = handle.commands
|
||||
assert append.is_speech is True
|
||||
assert append.audio_end_ms == 100
|
||||
assert append.event_id
|
||||
assert commit.final is True and commit.create_response is True
|
||||
assert cancel.response_id == "resp-1"
|
||||
assert ack.played_ms == 120 and ack.response_id == "resp-1"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_inline_session_close_maps_to_handle_close():
|
||||
client, omni = make_inline_client()
|
||||
async with client:
|
||||
pass
|
||||
handle = omni.handles[0]
|
||||
assert handle.closed
|
||||
assert handle.close_calls >= 1
|
||||
assert not any(isinstance(command, duplex_commands.CloseSession) for command in handle.commands)
|
||||
# The typed session.closed reached the client as the expected end.
|
||||
with pytest.raises(DuplexSessionClosedError):
|
||||
await client.send({"type": "session.heartbeat"})
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_inline_client_forwards_heartbeats_and_drops_event_acks():
|
||||
client, omni = make_inline_client()
|
||||
async with client:
|
||||
handle = omni.handles[0]
|
||||
await client.send({"type": "session.heartbeat"})
|
||||
await client.send({"type": "session.event_ack", "server_event_seq": 3})
|
||||
# A heartbeat is a lease touch the engine must see; the ack is a journal concern.
|
||||
assert [type(command).__name__ for command in handle.commands] == ["Heartbeat"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_inline_malformed_payload_surfaces_as_error_event():
|
||||
client, omni = make_inline_client()
|
||||
async with client:
|
||||
handle = omni.handles[0]
|
||||
# Subscribe before sending: the translation error is dispatched inline.
|
||||
first_event = asyncio.create_task(client.wait_for("error", timeout_s=5.0))
|
||||
await asyncio.sleep(0)
|
||||
event_id = await client.send({"type": "playback.ack"}) # played_ms missing
|
||||
event = await first_event
|
||||
assert isinstance(event, ErrorEvent)
|
||||
assert event.code == "bad_event"
|
||||
assert event.related_event_id == event_id
|
||||
assert handle.commands == []
|
||||
with pytest.raises(DuplexProtocolError) as excinfo:
|
||||
async for _ in client.responses():
|
||||
pass
|
||||
assert excinfo.value.code == "bad_event"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_inline_teardown_cancels_pump_and_closes_handle():
|
||||
client, omni = make_inline_client()
|
||||
async with client:
|
||||
handle = omni.handles[0]
|
||||
await _drain(lambda: handle.consumer_active)
|
||||
pump = client._pump_task
|
||||
assert pump is not None and not pump.done()
|
||||
assert pump.done()
|
||||
assert not handle.consumer_active
|
||||
assert handle.closed
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_inline_send_after_handle_closed_raises_connection_error():
|
||||
client, omni = make_inline_client()
|
||||
async with client:
|
||||
handle = omni.handles[0]
|
||||
handle.closed = True # the engine ended the session without an event yet
|
||||
with pytest.raises(DuplexConnectionError, match="not open"):
|
||||
await client.send({"type": "input_audio_buffer.clear"})
|
||||
handle.closed = False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Collector
|
||||
|
||||
@@ -600,8 +913,7 @@ def test_build_realtime_url_with_model_extra_query():
|
||||
url = build_realtime_url(
|
||||
"ws://localhost:8099/v1/realtime?custom=1",
|
||||
"openbmb/MiniCPM-o-4_5",
|
||||
session_id="session-a",
|
||||
extra_query={"native_duplex": "1"},
|
||||
extra_query={"probe": "1"},
|
||||
)
|
||||
|
||||
query = parse_qs(urlsplit(url).query)
|
||||
@@ -609,8 +921,7 @@ def test_build_realtime_url_with_model_extra_query():
|
||||
"custom": ["1"],
|
||||
"duplex": ["1"],
|
||||
"model": ["openbmb/MiniCPM-o-4_5"],
|
||||
"native_duplex": ["1"],
|
||||
"session_id": ["session-a"],
|
||||
"probe": ["1"],
|
||||
}
|
||||
|
||||
|
||||
@@ -619,12 +930,11 @@ def test_build_realtime_url_resume_only_when_autostart_disabled():
|
||||
"ws://localhost:8099/v1/realtime?duplex=1",
|
||||
"openbmb/MiniCPM-o-4_5",
|
||||
autostart=False,
|
||||
extra_query={"native_duplex": "1"},
|
||||
)
|
||||
|
||||
query = parse_qs(urlsplit(url).query)
|
||||
assert query["autostart"] == ["0"]
|
||||
assert query["native_duplex"] == ["1"]
|
||||
assert "session_id" not in query
|
||||
|
||||
|
||||
def test_event_collector_partitions_audio_by_response():
|
||||
@@ -989,6 +1299,7 @@ def test_duplex_unit_boundary_and_residual_math():
|
||||
assert chunk_period_ms([{"type": "session.created", "session": {}}]) == 1000
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_pcm_sends_each_units_composite_beside_its_base_frame():
|
||||
sock = FakeSocket()
|
||||
sock.feed(SESSION_CREATED)
|
||||
@@ -1011,16 +1322,16 @@ async def test_stream_pcm_sends_each_units_composite_beside_its_base_frame():
|
||||
sock.feed(SESSION_CLOSED)
|
||||
|
||||
|
||||
def test_build_realtime_url_native_duplex_flag_and_http_scheme():
|
||||
url = build_realtime_url("http://localhost:8099/v1/realtime", None, native_duplex=None)
|
||||
def test_build_realtime_url_rewrites_http_scheme_and_rejects_others():
|
||||
url = build_realtime_url("http://localhost:8099/v1/realtime", None)
|
||||
parts = urlsplit(url)
|
||||
assert parts.scheme == "ws"
|
||||
assert parse_qs(parts.query) == {"duplex": ["1"]}
|
||||
|
||||
url = build_realtime_url("https://host/v1/realtime", "m", native_duplex=False)
|
||||
url = build_realtime_url("https://host/v1/realtime", "m", autostart=True)
|
||||
parts = urlsplit(url)
|
||||
assert parts.scheme == "wss"
|
||||
assert parse_qs(parts.query)["native_duplex"] == ["0"]
|
||||
assert parse_qs(parts.query) == {"duplex": ["1"], "model": ["m"], "autostart": ["1"]}
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
build_realtime_url("ftp://host/v1/realtime", "m")
|
||||
|
||||
@@ -18,13 +18,13 @@ from vllm_omni.clients.personaplex import (
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_minicpmo45_session_config_matches_native_duplex_deployment():
|
||||
def test_minicpmo45_session_config_matches_duplex_deployment():
|
||||
config = create_minicpmo45_session_config(ref_audio="data:audio/wav;base64,AAA=", temperature=0.0)
|
||||
payload = config.to_session_payload(model="openbmb/MiniCPM-o-4_5")
|
||||
assert payload["ref_audio"] == "data:audio/wav;base64,AAA="
|
||||
assert payload["overlap_policy"] == "listen_only"
|
||||
assert payload["playback_commit_policy"] == "ack_only"
|
||||
assert payload["extra_body"]["native_duplex"] is True
|
||||
assert "native_duplex" not in payload["extra_body"]
|
||||
assert payload["extra_body"]["auto_response"] is True
|
||||
assert payload["temperature"] == 0.0
|
||||
|
||||
@@ -41,8 +41,9 @@ def test_personaplex_session_config_matches_deployment():
|
||||
@pytest.mark.parametrize(
|
||||
"duplex_package",
|
||||
[
|
||||
# PersonaPlex and Nemotron VoiceChat rejoin this list with the follow-up
|
||||
# PRs that port them to the duplex plugin framework.
|
||||
"vllm_omni.model_executor.models.minicpmo_4_5.duplex",
|
||||
"vllm_omni.model_executor.models.personaplex.duplex",
|
||||
],
|
||||
)
|
||||
def test_model_duplex_packages_stay_clear_of_client_library(duplex_package: str):
|
||||
|
||||
@@ -1328,7 +1328,11 @@ class TestDeployConfigLoading:
|
||||
stages = merge_pipeline_deploy(pipeline, deploy)
|
||||
|
||||
assert deploy.session_mode == "duplex"
|
||||
assert deploy.active_stream_window == max_sessions
|
||||
# ``0`` disables the limiter: the configured session/stage capacity is
|
||||
# what bounds the pipeline. Any positive window serializes concurrent
|
||||
# audio first-packet generation (measured: TTFP -23% at 0 vs 4 on
|
||||
# L20X, 128 requests / concurrency 8).
|
||||
assert deploy.active_stream_window == 0
|
||||
assert deploy.duplex_session.max_sessions == max_sessions
|
||||
assert [stage.session_mode for stage in stages] == ["duplex", "duplex", "duplex"]
|
||||
assert [stage.to_omegaconf().session_mode for stage in stages] == ["duplex", "duplex", "duplex"]
|
||||
@@ -1558,7 +1562,10 @@ stages:
|
||||
monkeypatch.setattr(stage_init_utils, "resolve_worker_cls", lambda _engine_args: None)
|
||||
|
||||
deploy = load_deploy_config(get_deploy_config_path("minimax_h3_disaggregated.yaml"))
|
||||
assert deploy.stages[1].engine_extras["model_loaded"] == {"text_encoder": False}
|
||||
assert deploy.stages[1].engine_extras["model_loaded"] == {
|
||||
"text_encoder": False,
|
||||
"vae_encoder": False,
|
||||
}
|
||||
stages = merge_pipeline_deploy(OMNI_PIPELINES["minimax_h3_disaggregated"], deploy)
|
||||
resolved = [stage_init_utils.build_engine_args_dict(stage.to_omegaconf(), str(model_root)) for stage in stages]
|
||||
|
||||
@@ -1583,12 +1590,20 @@ 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[0].yaml_engine_args["model_arch"] == "MiniMaxH3Encoder"
|
||||
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[1].yaml_engine_args["model_loaded"] == {
|
||||
"text_encoder": False,
|
||||
"vae_encoder": False,
|
||||
}
|
||||
assert stages[0].yaml_engine_args["max_num_seqs"] == 1
|
||||
assert stages[0].yaml_engine_args["hf_overrides"]["minimax_h3_encoder_components"] == {
|
||||
"text_encoder": {"parallel_mode": "tp"},
|
||||
"video_vae": {"parallel_mode": "patch"},
|
||||
"audio_vae": {"parallel_mode": "leader"},
|
||||
}
|
||||
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")
|
||||
parallel = stages[1].yaml_engine_args["parallel_config"]
|
||||
@@ -1600,6 +1615,9 @@ stages:
|
||||
|
||||
turbo = load_deploy_config(Path(get_deploy_config_path("minimax_h3_disaggregated_turbo.yaml")))
|
||||
turbo_stages = merge_pipeline_deploy(pipeline, turbo)
|
||||
assert turbo_stages[0].yaml_engine_args["model_arch"] == stages[0].yaml_engine_args["model_arch"]
|
||||
assert turbo_stages[1].yaml_engine_args["model_loaded"] == stages[1].yaml_engine_args["model_loaded"]
|
||||
assert turbo_stages[0].yaml_engine_args["hf_overrides"] == stages[0].yaml_engine_args["hf_overrides"]
|
||||
turbo_sampling = turbo_stages[1].yaml_extras["default_sampling_params"]
|
||||
assert turbo_sampling["num_inference_steps"] == 5
|
||||
assert turbo_sampling["extra_args"] == {"flow_shift": 6.0, "audio_flow_shift": 3.0}
|
||||
@@ -2587,6 +2605,26 @@ class TestPlatformOverrides:
|
||||
assert rocm.stages[0].enforce_eager is None
|
||||
assert rocm.stages[1].enforce_eager is True
|
||||
|
||||
@pytest.mark.parametrize("deploy_name", ["qwen3_tts.yaml", "qwen3_tts_high_concurrency.yaml"])
|
||||
@pytest.mark.parametrize("platform", ["cuda", "npu", "rocm"])
|
||||
def test_qwen3_tts_default_code2wav_dtype_is_bf16(self, deploy_name, platform):
|
||||
deploy = load_deploy_config(Path(get_deploy_config_path(deploy_name)))
|
||||
deploy = _apply_platform_overrides(deploy, platform=platform)
|
||||
stages = merge_pipeline_deploy(resolve_pipeline_config("qwen3_tts"), deploy)
|
||||
|
||||
assert stages[1].yaml_engine_args["dtype"] == "bfloat16"
|
||||
|
||||
@pytest.mark.parametrize("deploy_name", ["qwen3_tts.yaml", "qwen3_tts_high_concurrency.yaml"])
|
||||
@pytest.mark.parametrize("platform", ["cuda", "npu"])
|
||||
@pytest.mark.parametrize("dtype", ["float32", "bfloat16", "float16"])
|
||||
def test_qwen3_tts_propagates_explicit_code2wav_dtype(self, deploy_name, platform, dtype):
|
||||
deploy = load_deploy_config(Path(get_deploy_config_path(deploy_name)))
|
||||
deploy.stages[1].engine_extras["dtype"] = dtype
|
||||
deploy = _apply_platform_overrides(deploy, platform=platform)
|
||||
stages = merge_pipeline_deploy(resolve_pipeline_config("qwen3_tts"), deploy)
|
||||
|
||||
assert stages[1].yaml_engine_args["dtype"] == dtype
|
||||
|
||||
def test_higgs_audio_v3_rocm_uses_triton_attention(self):
|
||||
deploy_path = Path(get_deploy_config_path("higgs_multimodal_qwen3.yaml"))
|
||||
|
||||
|
||||
@@ -164,7 +164,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
|
||||
"""Make an inventory expansion an explicit review decision."""
|
||||
category_counts = Counter(item.category for item in ENVIRONMENT_VARIABLE_INVENTORY.values())
|
||||
assert category_counts == {
|
||||
EnvironmentVariableCategory.PUBLIC_OMNI: 27,
|
||||
EnvironmentVariableCategory.PUBLIC_OMNI: 28,
|
||||
EnvironmentVariableCategory.INHERITED_VLLM: 20,
|
||||
EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27,
|
||||
EnvironmentVariableCategory.MODEL_SPECIFIC: 64,
|
||||
|
||||
@@ -149,6 +149,9 @@ def _make_scheduler_stub(requests: list[_Request]) -> SimpleNamespace:
|
||||
for name in _MIXIN_UPDATE_HELPERS:
|
||||
setattr(scheduler, name, MethodType(getattr(OmniSchedulerMixin, name), scheduler))
|
||||
scheduler._cleanup_kv_tracking = MethodType(OmniARScheduler._cleanup_kv_tracking, scheduler)
|
||||
scheduler._emit_streaming_context_overflow_outputs = MethodType(
|
||||
OmniARScheduler._emit_streaming_context_overflow_outputs, scheduler
|
||||
)
|
||||
scheduler.make_spec_decoding_stats = lambda *args, **kwargs: None
|
||||
scheduler.make_stats = lambda *args, **kwargs: None
|
||||
return scheduler
|
||||
|
||||
@@ -5,6 +5,7 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections import defaultdict
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
@@ -17,6 +18,8 @@ import vllm_omni # noqa: F401 - import for side effects (patch vLLM)
|
||||
from vllm.sampling_params import SamplingParams
|
||||
from vllm.v1.core.sched.output import SchedulerOutput
|
||||
from vllm.v1.outputs import ModelRunnerOutput
|
||||
from vllm.v1.engine import FinishReason
|
||||
from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue
|
||||
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 (
|
||||
@@ -958,3 +961,303 @@ def test_chunk_segment_cleanup_keeps_explicit_update_stage_parked(
|
||||
assert session.request_id not in sched.chunk_transfer_adapter.segment_finished_requests
|
||||
sched.skipped_waiting.remove_requests.assert_not_called()
|
||||
sched._enqueue_waiting_request.assert_not_called()
|
||||
|
||||
|
||||
def _make_live_session_scheduler(*, max_model_len: int) -> OmniARScheduler:
|
||||
"""A scheduler whose finish path is real: the queues, ``finish_requests``,
|
||||
``_free_request`` and the finished-request bookkeeping that wakes an idle
|
||||
engine all run. Only block freeing and the encoder cache are mocked."""
|
||||
sched = _make_admission_scheduler(max_model_len=max_model_len)
|
||||
del sched._free_request # the admission helper mocks it; run the real one
|
||||
sched._omits_kv_transfer_cache = {}
|
||||
sched.connector = None
|
||||
sched.perf_metrics = None
|
||||
sched.recompute_kv_load_failures = False
|
||||
sched.finished_req_ids = set()
|
||||
sched.finished_req_ids_dict = defaultdict(set)
|
||||
sched.transfer_triggered_requests = set()
|
||||
sched.active_kv_transfers = set()
|
||||
sched.waiting_for_transfer_free = set()
|
||||
sched.pending_stop_after_extraction = set()
|
||||
sched.requests_needing_kv_transfer = {}
|
||||
sched._kv_wait_start_ts = {}
|
||||
sched.kv_cache_manager = SimpleNamespace(take_events=lambda: None, estimate_cached_tokens=lambda _request: 0)
|
||||
return sched
|
||||
|
||||
|
||||
def _run_idle_step(sched: OmniARScheduler):
|
||||
"""One ``update_from_output`` with nothing scheduled: the step an idle
|
||||
engine runs once ``has_finished_requests()`` reports a freed session."""
|
||||
scheduler_output = SimpleNamespace(
|
||||
num_scheduled_tokens={},
|
||||
scheduled_spec_decode_tokens={},
|
||||
num_invalid_spec_tokens=0,
|
||||
)
|
||||
model_runner_output = SimpleNamespace(
|
||||
sampled_token_ids=[],
|
||||
logprobs=None,
|
||||
prompt_logprobs_dict={},
|
||||
pooler_output=None,
|
||||
num_nans_in_logits=None,
|
||||
kv_connector_output=None,
|
||||
cudagraph_stats=None,
|
||||
req_id_to_index={},
|
||||
routed_experts=None,
|
||||
)
|
||||
return OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output)
|
||||
|
||||
|
||||
def _park_session(sched: OmniARScheduler, session: Request) -> None:
|
||||
"""Put a session where a 1 fps video stream leaves it between appends:
|
||||
prompt partly computed, waiting for streaming input, in admission."""
|
||||
sched.requests[session.request_id] = session
|
||||
session.append_output_token_ids([7, 8, 9])
|
||||
session.num_computed_tokens = 6
|
||||
session.num_output_placeholders = 0
|
||||
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
|
||||
sched.skipped_waiting.add_request(session)
|
||||
sched.num_waiting_for_streaming_input = 1
|
||||
|
||||
|
||||
def test_stage0_streaming_update_that_overflows_max_model_len_finishes_the_session() -> None:
|
||||
"""A native duplex session grows its prompt on every append; once the
|
||||
extension cannot fit the model the worker crashes copying the prompt.
|
||||
The scheduler must drop the update and fail only this session, at once,
|
||||
through the real finish path: a parked session does not make the engine
|
||||
schedule, so what wakes the engine is the freed request showing up in
|
||||
``has_finished_requests()``, and the step that follows must carry one
|
||||
explicit ERROR rather than the synthesized ABORT."""
|
||||
sched = _make_live_session_scheduler(max_model_len=8)
|
||||
session = _make_request()
|
||||
_park_session(sched, session)
|
||||
|
||||
sched._update_request_as_session(session, _make_update([10, 20, 30]))
|
||||
|
||||
assert session.prompt_token_ids == [1, 2, 3]
|
||||
assert session.num_prompt_tokens == 3
|
||||
assert session.status == RequestStatus.FINISHED_ERROR
|
||||
assert session.request_id not in sched.requests
|
||||
assert len(sched.waiting) == 0
|
||||
assert len(sched.skipped_waiting) == 0
|
||||
assert sched.num_waiting_for_streaming_input == 0
|
||||
sched._free_request_blocks.assert_called_once_with(session)
|
||||
assert sched.finished_req_ids == {session.request_id}
|
||||
assert sched.has_finished_requests()
|
||||
client_index, reason = sched._streaming_context_overflow[session.request_id]
|
||||
assert client_index == session.client_index
|
||||
assert reason.startswith("context_length_exceeded: ")
|
||||
assert "9 tokens" in reason and "max_model_len 8" in reason
|
||||
|
||||
engine_core_outputs = _run_idle_step(sched)
|
||||
|
||||
(output,) = engine_core_outputs[session.client_index].outputs
|
||||
assert output.request_id == session.request_id
|
||||
assert output.finish_reason == FinishReason.ERROR
|
||||
assert output.stop_reason == reason
|
||||
assert output.new_token_ids == []
|
||||
assert engine_core_outputs[session.client_index].finished_requests == {session.request_id}
|
||||
assert sched._streaming_context_overflow == {}
|
||||
# Nothing is left to report once the error went out.
|
||||
assert _run_idle_step(sched) == {}
|
||||
|
||||
|
||||
def test_stage0_streaming_update_that_fills_max_model_len_exactly_finishes_the_session() -> None:
|
||||
"""An exactly full prompt leaves no room for the token every duplex step
|
||||
samples: upstream's running budget drops to -1, which ``schedule()`` does
|
||||
not catch, so the exact fill is an overflow too."""
|
||||
sched = _make_scheduler(stage_id=0)
|
||||
sched.max_model_len = 8
|
||||
sched.finish_requests = MagicMock()
|
||||
session = _make_request()
|
||||
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
|
||||
session.append_output_token_ids([7, 8, 9])
|
||||
session.num_computed_tokens = 6
|
||||
session.num_output_placeholders = 0
|
||||
|
||||
sched._update_request_as_session(session, _make_update([10, 20]))
|
||||
|
||||
assert session.prompt_token_ids == [1, 2, 3]
|
||||
assert session.num_prompt_tokens == 3
|
||||
sched.finish_requests.assert_called_once_with((session.request_id,), RequestStatus.FINISHED_ERROR)
|
||||
_, reason = sched._streaming_context_overflow[session.request_id]
|
||||
assert "8 tokens" in reason and "max_model_len 8" in reason
|
||||
|
||||
|
||||
def test_stage0_streaming_update_that_leaves_room_to_sample_is_applied() -> None:
|
||||
sched = _make_scheduler(stage_id=0)
|
||||
sched.max_model_len = 9
|
||||
sched.finish_requests = MagicMock()
|
||||
session = _make_request()
|
||||
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
|
||||
session.append_output_token_ids([7, 8, 9])
|
||||
session.num_computed_tokens = 6
|
||||
session.num_output_placeholders = 0
|
||||
|
||||
sched._update_request_as_session(session, _make_update([10, 20]))
|
||||
|
||||
assert session.prompt_token_ids == [1, 2, 3, 7, 8, 9, 10, 20]
|
||||
assert session.num_prompt_tokens == 8
|
||||
assert session.status == RequestStatus.WAITING
|
||||
sched.finish_requests.assert_not_called()
|
||||
assert not getattr(sched, "_streaming_context_overflow", {})
|
||||
|
||||
|
||||
def test_stage0_streaming_update_keeps_room_for_every_token_sampled_per_step() -> None:
|
||||
"""With speculative decoding a step samples more than one token, and the
|
||||
prompt must leave room for all of them."""
|
||||
sched = _make_scheduler(stage_id=0)
|
||||
sched.max_model_len = 9
|
||||
sched.num_sampled_tokens_per_step = 2
|
||||
sched.finish_requests = MagicMock()
|
||||
session = _make_request()
|
||||
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
|
||||
session.append_output_token_ids([7, 8, 9])
|
||||
session.num_computed_tokens = 6
|
||||
session.num_output_placeholders = 0
|
||||
|
||||
sched._update_request_as_session(session, _make_update([10, 20]))
|
||||
|
||||
assert session.prompt_token_ids == [1, 2, 3]
|
||||
sched.finish_requests.assert_called_once_with((session.request_id,), RequestStatus.FINISHED_ERROR)
|
||||
assert "8 tokens" in sched._streaming_context_overflow[session.request_id][1]
|
||||
|
||||
|
||||
def test_stage0_streaming_update_overflow_counts_an_uncomputed_prompt() -> None:
|
||||
"""A preempted session keeps its prompt even though num_computed_tokens
|
||||
fell behind it; the projection must use the longer of the two."""
|
||||
sched = _make_scheduler(stage_id=0)
|
||||
sched.max_model_len = 4
|
||||
sched.finish_requests = MagicMock()
|
||||
session = _make_request()
|
||||
session.status = RequestStatus.WAITING_FOR_STREAMING_REQ
|
||||
session.num_computed_tokens = 0
|
||||
|
||||
sched._update_request_as_session(session, _make_update([10, 20]))
|
||||
|
||||
assert session.prompt_token_ids == [1, 2, 3]
|
||||
sched.finish_requests.assert_called_once()
|
||||
assert "5 tokens" in sched._streaming_context_overflow[session.request_id][1]
|
||||
|
||||
|
||||
def _make_admission_scheduler(*, max_model_len: int) -> OmniARScheduler:
|
||||
"""A scheduler with real admission queues, so nothing about the terminal
|
||||
handling below is mocked away except block freeing."""
|
||||
sched = _make_scheduler(stage_id=0)
|
||||
sched.max_model_len = max_model_len
|
||||
sched.policy = SchedulingPolicy.FCFS
|
||||
sched.waiting = create_request_queue(sched.policy)
|
||||
sched.skipped_waiting = create_request_queue(sched.policy)
|
||||
sched.running = []
|
||||
sched.requests = {}
|
||||
sched._free_request = MagicMock() # needs a KV manager; not what this covers
|
||||
return sched
|
||||
|
||||
|
||||
def _make_queued_stop(session: Request, update: StreamingUpdate) -> None:
|
||||
"""Put the request in the state upstream hands to _handle_stopped_request:
|
||||
an append that arrived during generation is queued, and the stop status is
|
||||
still on the request."""
|
||||
from collections import deque
|
||||
|
||||
session.resumable = True
|
||||
session.streaming_queue = deque([update])
|
||||
session.status = RequestStatus.FINISHED_STOPPED
|
||||
|
||||
|
||||
def test_queued_streaming_update_that_overflows_does_not_return_to_admission() -> None:
|
||||
sched = _make_admission_scheduler(max_model_len=4)
|
||||
session = _make_request()
|
||||
sched.requests[session.request_id] = session
|
||||
session.num_computed_tokens = 3
|
||||
_make_queued_stop(session, _make_update([10, 20]))
|
||||
|
||||
finished = sched._handle_stopped_request(session)
|
||||
|
||||
assert finished is True
|
||||
assert len(sched.waiting) == 0
|
||||
assert len(sched.skipped_waiting) == 0
|
||||
assert session.status == RequestStatus.FINISHED_ERROR
|
||||
assert session.resumable is False
|
||||
assert "context_length_exceeded: " in sched._streaming_context_overflow[session.request_id][1]
|
||||
# update_from_output frees every request reported as finished; freeing here
|
||||
# as well deleted it from self.requests twice (KeyError in _free_blocks).
|
||||
sched._free_request.assert_not_called()
|
||||
|
||||
|
||||
def test_queued_streaming_update_that_fits_still_resumes_the_session() -> None:
|
||||
sched = _make_admission_scheduler(max_model_len=64)
|
||||
session = _make_request()
|
||||
sched.requests[session.request_id] = session
|
||||
session.num_computed_tokens = 3
|
||||
_make_queued_stop(session, _make_update([10, 20]))
|
||||
|
||||
finished = sched._handle_stopped_request(session)
|
||||
|
||||
assert finished is False
|
||||
assert len(sched.waiting) == 1
|
||||
assert session.status == RequestStatus.WAITING
|
||||
assert not getattr(sched, "_streaming_context_overflow", {})
|
||||
sched._free_request.assert_not_called()
|
||||
|
||||
|
||||
def test_queued_stop_with_a_recorded_overflow_still_leaves_admission() -> None:
|
||||
"""The overflow may already be on record when the stop is handled (an
|
||||
earlier update of the same session). Upstream still re-enqueues the
|
||||
request, so the override must take it back out either way instead of
|
||||
leaving a FINISHED_ERROR request for admission to trip over."""
|
||||
sched = _make_admission_scheduler(max_model_len=64)
|
||||
session = _make_request()
|
||||
sched.requests[session.request_id] = session
|
||||
session.num_computed_tokens = 3
|
||||
sched._streaming_context_overflow = {session.request_id: (session.client_index, "context_length_exceeded: earlier")}
|
||||
_make_queued_stop(session, _make_update([10, 20])) # this update itself fits
|
||||
|
||||
finished = sched._handle_stopped_request(session)
|
||||
|
||||
assert finished is True
|
||||
assert len(sched.waiting) == 0
|
||||
assert len(sched.skipped_waiting) == 0
|
||||
assert session.status == RequestStatus.FINISHED_ERROR
|
||||
assert session.resumable is False
|
||||
sched._free_request.assert_not_called()
|
||||
|
||||
|
||||
def test_parked_stop_with_a_recorded_overflow_keeps_the_streaming_counter_balanced() -> None:
|
||||
"""Same as above with nothing queued: upstream parks the request as
|
||||
waiting for streaming input and counts it; the dequeue must uncount it."""
|
||||
from collections import deque
|
||||
|
||||
sched = _make_admission_scheduler(max_model_len=64)
|
||||
session = _make_request()
|
||||
sched.requests[session.request_id] = session
|
||||
session.num_computed_tokens = 3
|
||||
sched._streaming_context_overflow = {session.request_id: (session.client_index, "context_length_exceeded: earlier")}
|
||||
session.resumable = True
|
||||
session.streaming_queue = deque()
|
||||
session.status = RequestStatus.FINISHED_STOPPED
|
||||
|
||||
finished = sched._handle_stopped_request(session)
|
||||
|
||||
assert finished is True
|
||||
assert len(sched.skipped_waiting) == 0
|
||||
assert sched.num_waiting_for_streaming_input == 0
|
||||
assert session.status == RequestStatus.FINISHED_ERROR
|
||||
sched._free_request.assert_not_called()
|
||||
|
||||
|
||||
def test_context_overflow_emits_an_error_output_with_the_reason() -> None:
|
||||
sched = _make_scheduler(stage_id=0)
|
||||
sched._streaming_context_overflow = {"req-a": (2, "context_length_exceeded: too long")}
|
||||
outputs: dict[int, list] = {}
|
||||
|
||||
sched._emit_streaming_context_overflow_outputs(outputs)
|
||||
|
||||
(output,) = outputs[2]
|
||||
assert output.request_id == "req-a"
|
||||
assert output.finish_reason == FinishReason.ERROR
|
||||
assert output.stop_reason == "context_length_exceeded: too long"
|
||||
assert output.new_token_ids == []
|
||||
assert sched._streaming_context_overflow == {}
|
||||
|
||||
sched._emit_streaming_context_overflow_outputs(outputs)
|
||||
assert len(outputs[2]) == 1
|
||||
@@ -26,8 +26,8 @@
|
||||
{
|
||||
"dataset_name": "seed-tts",
|
||||
"dataset_path": "zhaochenyang20/seed-tts-eval",
|
||||
"backend": "openai-chat-omni",
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"backend": "openai-realtime-tts",
|
||||
"endpoint": "/v1/realtime",
|
||||
"num_prompts": [
|
||||
32,
|
||||
64,
|
||||
@@ -42,73 +42,7 @@
|
||||
"disable_shuffle": true,
|
||||
"trust_remote_code": true,
|
||||
"seed_tts_locale": "en",
|
||||
"extra_body": {
|
||||
"modalities": [
|
||||
"text",
|
||||
"audio"
|
||||
],
|
||||
"chat_template_kwargs": {
|
||||
"enable_thinking": false,
|
||||
"use_tts_template": true
|
||||
}
|
||||
},
|
||||
"percentile-metrics": "ttft,e2el,audio_ttfp,audio_rtf,audio_duration",
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"request_throughput": [
|
||||
0.9731,
|
||||
1.1986,
|
||||
1.686
|
||||
],
|
||||
"mean_ttft_ms": [
|
||||
152.5665,
|
||||
270.2559,
|
||||
301.455
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
1027.2953,
|
||||
3306.3735,
|
||||
4686.5366
|
||||
],
|
||||
"mean_audio_ttfp_ms": [
|
||||
458.1553,
|
||||
2030.2335,
|
||||
2068.6233
|
||||
],
|
||||
"mean_audio_rtf": [
|
||||
0.2357,
|
||||
0.791,
|
||||
1.0324
|
||||
]
|
||||
},
|
||||
"A3": {
|
||||
"request_throughput": [
|
||||
0.5383,
|
||||
0.6042,
|
||||
0.7547
|
||||
],
|
||||
"mean_ttft_ms": [
|
||||
333.2633,
|
||||
533.8656,
|
||||
547.3388
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
1857.2154,
|
||||
6600.5101,
|
||||
10450.6811
|
||||
],
|
||||
"mean_audio_ttfp_ms": [
|
||||
986.4666,
|
||||
3411.0754,
|
||||
3352.7515
|
||||
],
|
||||
"mean_audio_rtf": [
|
||||
0.4423,
|
||||
1.5734,
|
||||
2.3024
|
||||
]
|
||||
}
|
||||
}
|
||||
"percentile-metrics": "ttft,e2el,audio_ttfp,audio_rtf,audio_duration"
|
||||
}
|
||||
]
|
||||
},
|
||||
@@ -138,8 +72,8 @@
|
||||
{
|
||||
"dataset_name": "seed-tts",
|
||||
"dataset_path": "zhaochenyang20/seed-tts-eval",
|
||||
"backend": "openai-chat-omni",
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"backend": "openai-realtime-tts",
|
||||
"endpoint": "/v1/realtime",
|
||||
"num_prompts": [
|
||||
32,
|
||||
64,
|
||||
@@ -154,16 +88,6 @@
|
||||
"disable_shuffle": true,
|
||||
"trust_remote_code": true,
|
||||
"seed_tts_locale": "en",
|
||||
"extra_body": {
|
||||
"modalities": [
|
||||
"text",
|
||||
"audio"
|
||||
],
|
||||
"chat_template_kwargs": {
|
||||
"enable_thinking": false,
|
||||
"use_tts_template": true
|
||||
}
|
||||
},
|
||||
"percentile-metrics": "ttft,e2el,audio_ttfp,audio_rtf,audio_duration"
|
||||
}
|
||||
]
|
||||
|
||||
@@ -43,43 +43,7 @@
|
||||
"extra_body": {
|
||||
"save_duplex_request_metrics": true
|
||||
},
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_ttfp,audio_rtf,audio_duration",
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"request_throughput": [
|
||||
0.3038
|
||||
],
|
||||
"mean_ttft_ms": [
|
||||
139.997
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
3269.8807
|
||||
],
|
||||
"mean_audio_ttfp_ms": [
|
||||
456.073
|
||||
],
|
||||
"mean_audio_rtf": [
|
||||
0.2254
|
||||
]
|
||||
},
|
||||
"A3": {
|
||||
"request_throughput": [
|
||||
0.1594
|
||||
],
|
||||
"mean_ttft_ms": [
|
||||
274.164
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
6250.7696
|
||||
],
|
||||
"mean_audio_ttfp_ms": [
|
||||
994.422
|
||||
],
|
||||
"mean_audio_rtf": [
|
||||
0.46163
|
||||
]
|
||||
}
|
||||
}
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_ttfp,audio_rtf,audio_duration"
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
@@ -26,8 +26,8 @@
|
||||
{
|
||||
"dataset_name": "seed-tts",
|
||||
"dataset_path": "zhaochenyang20/seed-tts-eval",
|
||||
"backend": "openai-chat-omni",
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"backend": "openai-realtime-tts",
|
||||
"endpoint": "/v1/realtime",
|
||||
"num_prompts": [
|
||||
10,
|
||||
40
|
||||
@@ -40,41 +40,7 @@
|
||||
"disable_shuffle": true,
|
||||
"trust_remote_code": true,
|
||||
"seed_tts_locale": "en",
|
||||
"extra_body": {
|
||||
"modalities": [
|
||||
"text",
|
||||
"audio"
|
||||
],
|
||||
"chat_template_kwargs": {
|
||||
"enable_thinking": false,
|
||||
"use_tts_template": true
|
||||
}
|
||||
},
|
||||
"percentile-metrics": "ttft,e2el,audio_ttfp,audio_rtf,audio_duration",
|
||||
"baseline": {
|
||||
"H100": {
|
||||
"request_throughput": [
|
||||
0.9731,
|
||||
1.1986
|
||||
],
|
||||
"mean_ttft_ms": [
|
||||
152.5665,
|
||||
270.2559
|
||||
],
|
||||
"mean_e2el_ms": [
|
||||
1027.2953,
|
||||
3306.3735
|
||||
],
|
||||
"mean_audio_ttfp_ms": [
|
||||
458.1553,
|
||||
2030.2335
|
||||
],
|
||||
"mean_audio_rtf": [
|
||||
0.2357,
|
||||
0.791
|
||||
]
|
||||
}
|
||||
}
|
||||
"percentile-metrics": "ttft,e2el,audio_ttfp,audio_rtf,audio_duration"
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
@@ -157,14 +157,8 @@ def test_speech_missing_required_fields(omni_server: OmniServer, online_client:
|
||||
"ref_audio": None,
|
||||
"ref_text": None,
|
||||
},
|
||||
# NOTE: This is delegating down the TTS adapter, so the message is both specific
|
||||
# to the Qwen3TTS adapter and dependent on whether or not the model has any speakers
|
||||
# (as opposed to having valida speakers, but getting an invalid speaker). it would
|
||||
# be more ideal to validate that this is landing on `.validate()` on the adapter
|
||||
# and ensuring that the returned message matches what we would get from calling the
|
||||
# adapter validate directly.
|
||||
("no speakers configured",),
|
||||
id="customvoice_invalid_voice",
|
||||
("Base checkpoint does not support task_type='CustomVoice'",),
|
||||
id="base_checkpoint_customvoice_task_mismatch",
|
||||
),
|
||||
pytest.param(
|
||||
{"task_type": "InvalidEnum"}, ("task_type", "literal_error", "CustomVoice"), id="task_type_invalid"
|
||||
|
||||
@@ -0,0 +1,126 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
import sys
|
||||
import types
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from torch._subclasses.fake_tensor import FakeTensorMode
|
||||
|
||||
from vllm_omni.diffusion.attention.ops.block_sparse import (
|
||||
build_prefix_dense_block_map,
|
||||
fastvideo_block_sparse_attn_bshd,
|
||||
mean_pool_tiles,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("prefix_blocks,topk", [(0, 1), (2, 0), (2, 1), (2, 3), (2, 5)])
|
||||
def test_block_map_keeps_prefix_dense_and_selects_non_prefix_keys(prefix_blocks, topk):
|
||||
sparse_blocks = 3
|
||||
blocks = prefix_blocks + sparse_blocks
|
||||
scores = torch.arange(2 * blocks * blocks, dtype=torch.float32).reshape(1, 2, blocks, blocks)
|
||||
block_map = build_prefix_dense_block_map(scores, prefix_blocks, sparse_blocks, topk)
|
||||
expected = torch.zeros_like(scores, dtype=torch.bool)
|
||||
# Scores increase with key index, so top-k always selects the last keys.
|
||||
for query_block in range(blocks):
|
||||
for key_block in range(blocks):
|
||||
if (
|
||||
query_block < prefix_blocks
|
||||
or key_block < prefix_blocks
|
||||
or key_block >= blocks - min(topk, sparse_blocks)
|
||||
):
|
||||
expected[..., query_block, key_block] = True
|
||||
assert torch.equal(block_map, expected)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("block_size", [4, 32, 64])
|
||||
@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
|
||||
def test_pooling_uses_valid_tile_sizes_and_fp32_accumulation(block_size, dtype):
|
||||
sizes = torch.tensor([block_size, block_size - 1, 0], dtype=torch.int32)
|
||||
tiles = torch.arange(2 * 3 * block_size * 2 * 3).reshape(2, 3, block_size, 2, 3).to(dtype)
|
||||
tiles[:, 1, -1] = 0
|
||||
tiles[:, 2] = 0
|
||||
expected = torch.stack(
|
||||
[tiles[:, 0].float().mean(dim=1), tiles[:, 1, :-1].float().mean(dim=1), torch.zeros(2, 2, 3)], dim=2
|
||||
)
|
||||
result = mean_pool_tiles(tiles.flatten(1, 2), sizes, block_size)
|
||||
assert result.dtype == torch.float32
|
||||
torch.testing.assert_close(result, expected, rtol=0, atol=0)
|
||||
|
||||
|
||||
def test_provider_wrapper_restores_transport_rows_and_bshd_layout(monkeypatch):
|
||||
monkeypatch.delenv("FASTVIDEO_VSA_SM100A", raising=False)
|
||||
calls = []
|
||||
|
||||
def block_sparse_attn(q, k, v, block_map, sizes):
|
||||
calls.append((q.shape, block_map.shape, sizes.tolist()))
|
||||
assert q.is_contiguous() and k.is_contiguous() and v.is_contiguous()
|
||||
assert sizes.dtype == torch.int32
|
||||
return q + k + v, None
|
||||
|
||||
module = types.ModuleType("fastvideo_kernel.block_sparse_attn")
|
||||
setattr(module, "block_sparse_attn", block_sparse_attn)
|
||||
monkeypatch.setitem(sys.modules, "fastvideo_kernel", types.ModuleType("fastvideo_kernel"))
|
||||
monkeypatch.setitem(sys.modules, "fastvideo_kernel.block_sparse_attn", module)
|
||||
query = torch.randn(1, 4 * 64, 2, 8)
|
||||
block_map = torch.ones(1, 2, 4, 4, dtype=torch.bool)
|
||||
sizes = torch.tensor([64, 17, 64, 0])
|
||||
output = fastvideo_block_sparse_attn_bshd(query, query, query, block_map, sizes, 3)
|
||||
expected = torch.zeros_like(query)
|
||||
expected[:, : 3 * 64] = query[:, : 3 * 64] * 3
|
||||
torch.testing.assert_close(output, expected)
|
||||
assert calls == [(torch.Size([1, 2, 3 * 64, 8]), torch.Size([1, 2, 3, 3]), [64, 17, 64])]
|
||||
|
||||
|
||||
def test_custom_op_fake_preserves_metadata_without_loading_provider(monkeypatch):
|
||||
monkeypatch.setitem(sys.modules, "fastvideo_kernel", None)
|
||||
with FakeTensorMode():
|
||||
query = torch.empty(1, 4 * 64, 2, 8, dtype=torch.bfloat16)
|
||||
block_map = torch.empty(1, 2, 4, 4, dtype=torch.bool)
|
||||
sizes = torch.empty(4, dtype=torch.int32)
|
||||
output = fastvideo_block_sparse_attn_bshd(query, query, query, block_map, sizes, 3)
|
||||
assert output.shape == query.shape
|
||||
assert output.dtype == query.dtype
|
||||
assert output.device == query.device
|
||||
|
||||
|
||||
@pytest.mark.parametrize("error_type", [ImportError, RuntimeError, torch.AcceleratorError])
|
||||
def test_native_provider_failure_only_recovers_without_accelerator_fault(monkeypatch, mocker, error_type):
|
||||
monkeypatch.setenv("FASTVIDEO_VSA_SM100A", "1")
|
||||
error = error_type("native provider failure")
|
||||
native = mocker.Mock()
|
||||
native.is_supported.return_value = True
|
||||
native.block_sparse_attn_sm100a.side_effect = error
|
||||
provider = types.ModuleType("fastvideo_kernel")
|
||||
setattr(provider, "block_sparse_attn_sm100a", native)
|
||||
index = types.ModuleType("fastvideo_kernel.triton_kernels.index")
|
||||
setattr(
|
||||
index,
|
||||
"map_to_index",
|
||||
mocker.Mock(return_value=(torch.zeros(1, dtype=torch.int32), torch.ones(1, dtype=torch.int32))),
|
||||
)
|
||||
triton = types.ModuleType("fastvideo_kernel.block_sparse_attn")
|
||||
fallback = mocker.Mock(side_effect=lambda q, *args: (q.clone(), None))
|
||||
setattr(triton, "block_sparse_attn", fallback)
|
||||
monkeypatch.setitem(sys.modules, "fastvideo_kernel", provider)
|
||||
monkeypatch.setitem(
|
||||
sys.modules, "fastvideo_kernel.triton_kernels", types.ModuleType("fastvideo_kernel.triton_kernels")
|
||||
)
|
||||
monkeypatch.setitem(sys.modules, "fastvideo_kernel.triton_kernels.index", index)
|
||||
monkeypatch.setitem(sys.modules, "fastvideo_kernel.block_sparse_attn", triton)
|
||||
query = torch.arange(128 * 2 * 8, dtype=torch.float32).reshape(1, 128, 2, 8)
|
||||
block_map = torch.ones(1, 2, 2, 2, dtype=torch.bool)
|
||||
sizes = torch.tensor([64, 64], dtype=torch.int32)
|
||||
|
||||
if error_type is torch.AcceleratorError:
|
||||
with pytest.raises(torch.AcceleratorError, match="native provider failure") as raised:
|
||||
fastvideo_block_sparse_attn_bshd(query, query, query, block_map, sizes, 2)
|
||||
assert raised.value is error
|
||||
fallback.assert_not_called()
|
||||
else:
|
||||
output = fastvideo_block_sparse_attn_bshd(query, query, query, block_map, sizes, 2)
|
||||
torch.testing.assert_close(output, query, rtol=0, atol=0)
|
||||
fallback.assert_called_once()
|
||||
@@ -1,7 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
import math
|
||||
import sys
|
||||
import types
|
||||
from typing import Any
|
||||
@@ -12,14 +11,10 @@ import torch
|
||||
from vllm_omni.diffusion.attention.backends.abstract import (
|
||||
AttentionMetadata,
|
||||
PackedPaddingMetadata,
|
||||
VideoTokenLayout,
|
||||
VideoTokenSpan,
|
||||
)
|
||||
from vllm_omni.diffusion.attention.backends.fastvideo_vsa import (
|
||||
FastVideoVSABackend,
|
||||
FastVideoVSAImpl,
|
||||
_build_h3_block_map,
|
||||
_get_h3_tile_metadata,
|
||||
)
|
||||
from vllm_omni.diffusion.attention.backends.registry import (
|
||||
DiffusionAttentionBackendEnum,
|
||||
@@ -218,144 +213,8 @@ def test_fastvideo_vsa_allows_topk_equal_to_num_blocks():
|
||||
assert too_many._fallback_reason(query, query, query, metadata) == "topk 3 > num_blocks 2"
|
||||
|
||||
|
||||
def _h3_metadata(prefix_segments, video_shape, *, gate=None, packed_padding=None):
|
||||
extra: dict[str, Any] = {
|
||||
"vsa_h3_prefix_segments": prefix_segments,
|
||||
}
|
||||
if gate is not None:
|
||||
extra["gate_compress"] = gate
|
||||
return AttentionMetadata(
|
||||
packed_padding=packed_padding,
|
||||
video_layout=VideoTokenLayout(
|
||||
used_len=sum(prefix_segments) + math.prod(video_shape),
|
||||
video_spans=(VideoTokenSpan(start=sum(prefix_segments), latent_grid=video_shape, role="target"),),
|
||||
),
|
||||
extra=extra,
|
||||
)
|
||||
|
||||
|
||||
def _fake_block_sparse_kernel(monkeypatch, calls):
|
||||
"""Stand in for fastvideo_kernel.block_sparse_attn with a mask-obeying SDPA."""
|
||||
|
||||
def block_sparse_attn(q, k, v, block_map, variable_block_sizes):
|
||||
calls["block_map"] = block_map.clone()
|
||||
calls["variable_block_sizes"] = variable_block_sizes.tolist()
|
||||
calls["q_shape"] = tuple(q.shape)
|
||||
blocks = block_map.shape[-1]
|
||||
assert blocks * 64 == q.shape[2]
|
||||
# Expand the per-block map to a per-token mask so the stand-in honours
|
||||
# the sparsity contract the real kernel implements, and drop the rows
|
||||
# each partial block pads with, which the real kernel skips through
|
||||
# variable_block_sizes.
|
||||
token_mask = block_map.repeat_interleave(64, dim=-1).repeat_interleave(64, dim=-2)
|
||||
within = torch.arange(64, device=q.device)[None, :] < variable_block_sizes[:, None]
|
||||
token_mask = token_mask & within.reshape(1, 1, 1, -1)
|
||||
out = torch.nn.functional.scaled_dot_product_attention(q.float(), k.float(), v.float(), attn_mask=token_mask)
|
||||
return out.to(q.dtype), None
|
||||
|
||||
module = types.ModuleType("fastvideo_kernel.block_sparse_attn")
|
||||
setattr(module, "block_sparse_attn", block_sparse_attn)
|
||||
monkeypatch.setitem(sys.modules, "fastvideo_kernel", types.ModuleType("fastvideo_kernel"))
|
||||
monkeypatch.setitem(sys.modules, "fastvideo_kernel.block_sparse_attn", module)
|
||||
|
||||
|
||||
def _h3_impl(**backend_kwargs):
|
||||
kwargs = {"topk": 1, "min_seq_len": 1, "disable_when_sp_active": False}
|
||||
kwargs.update(backend_kwargs)
|
||||
return FastVideoVSAImpl(num_heads=2, head_size=8, softmax_scale=8**-0.5, backend_kwargs=kwargs)
|
||||
|
||||
|
||||
def test_h3_forward_tiles_untiles_and_restores_the_original_row_order(monkeypatch):
|
||||
calls: dict[str, Any] = {}
|
||||
_fake_block_sparse_kernel(monkeypatch, calls)
|
||||
impl = _h3_impl(topk=64)
|
||||
prefix_segments, video_shape = (5, 70), (2, 4, 4)
|
||||
seq_len = sum(prefix_segments) + 2 * 4 * 4
|
||||
torch.manual_seed(0)
|
||||
query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16)
|
||||
|
||||
output = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
|
||||
|
||||
assert output.shape == query.shape
|
||||
# topk covers every video block, so the block map is dense and the tiled
|
||||
# kernel must reproduce plain attention over the same rows.
|
||||
reference = torch.nn.functional.scaled_dot_product_attention(
|
||||
*(x.float().transpose(1, 2) for x in (query, query, query))
|
||||
).transpose(1, 2)
|
||||
torch.testing.assert_close(output.float(), reference, atol=2e-2, rtol=2e-2)
|
||||
# 5 + 70 splits into segment-pure 5 | 64 | 6, then one (2, 4, 4) video tile.
|
||||
assert calls["variable_block_sizes"] == [5, 64, 6, 32]
|
||||
|
||||
|
||||
def test_h3_forward_keeps_prefix_dense_and_selects_top_k_video_tiles(monkeypatch):
|
||||
calls: dict[str, Any] = {}
|
||||
_fake_block_sparse_kernel(monkeypatch, calls)
|
||||
impl = _h3_impl(topk=1)
|
||||
prefix_segments, video_shape = (5,), (8, 4, 4)
|
||||
seq_len = 5 + 8 * 4 * 4
|
||||
torch.manual_seed(0)
|
||||
query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16)
|
||||
|
||||
impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
|
||||
|
||||
# 1 prefix block + 2 video blocks. The odd count is padded to an even one
|
||||
# for the paired-CTA kernel contract, and the Triton route slices that
|
||||
# transport-only partner off again before launching.
|
||||
block_map = calls["block_map"]
|
||||
assert block_map.shape[-1] == 3
|
||||
assert calls["variable_block_sizes"] == [5, 64, 64]
|
||||
assert block_map[:, :, :1, :].all(), "prefix queries stay dense"
|
||||
assert block_map[..., :1].all(), "prefix keys are exempt from selection"
|
||||
assert (block_map[:, :, 1:, 1:].sum(dim=-1) == 1).all(), "one video tile per video query"
|
||||
|
||||
# The same odd geometry must still untile back onto the original rows.
|
||||
dense = _h3_impl(topk=64).forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
|
||||
reference = torch.nn.functional.scaled_dot_product_attention(
|
||||
*(x.float().transpose(1, 2) for x in (query, query, query))
|
||||
).transpose(1, 2)
|
||||
torch.testing.assert_close(dense.float(), reference, atol=2e-2, rtol=2e-2)
|
||||
|
||||
|
||||
def test_h3_forward_applies_the_learned_compression_gate(monkeypatch):
|
||||
_fake_block_sparse_kernel(monkeypatch, {})
|
||||
impl = _h3_impl(topk=64)
|
||||
prefix_segments, video_shape = (5,), (2, 4, 4)
|
||||
seq_len = 5 + 2 * 4 * 4
|
||||
torch.manual_seed(0)
|
||||
query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16)
|
||||
|
||||
without_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
|
||||
zero_gate = torch.zeros_like(query)
|
||||
with_zero_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, gate=zero_gate))
|
||||
live_gate = torch.full_like(query, 0.5)
|
||||
with_live_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, gate=live_gate))
|
||||
|
||||
torch.testing.assert_close(with_zero_gate.float(), without_gate.float())
|
||||
assert not torch.allclose(with_live_gate.float(), without_gate.float())
|
||||
|
||||
|
||||
def test_h3_forward_restores_rows_the_packed_padding_excludes(monkeypatch):
|
||||
_fake_block_sparse_kernel(monkeypatch, {})
|
||||
impl = _h3_impl(topk=64)
|
||||
prefix_segments, video_shape = (5,), (2, 4, 4)
|
||||
valid = 5 + 2 * 4 * 4
|
||||
torch.manual_seed(0)
|
||||
query = torch.randn(1, valid + 24, 2, 8, dtype=torch.bfloat16)
|
||||
padding = PackedPaddingMetadata(
|
||||
q_length=valid,
|
||||
kv_length=valid,
|
||||
cu_seqlens_q=torch.tensor([0, valid], dtype=torch.int32),
|
||||
cu_seqlens_k=torch.tensor([0, valid], dtype=torch.int32),
|
||||
)
|
||||
|
||||
output = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, packed_padding=padding))
|
||||
|
||||
assert output.shape == query.shape
|
||||
assert torch.count_nonzero(output[:, valid:]) == 0
|
||||
|
||||
|
||||
def test_sdpa_fallback_never_attends_the_structural_padding(monkeypatch):
|
||||
# The backend advertises supports_packed_mask_free, so MiniMax-H3 skips
|
||||
# The backend advertises supports_packed_mask_free, so a model may skip
|
||||
# building the padding mask. Every fallback out of the VSA path must honour
|
||||
# that contract itself; SDPA reads attn_mask and nothing else.
|
||||
impl = FastVideoVSAImpl(
|
||||
@@ -374,6 +233,8 @@ def test_sdpa_fallback_never_attends_the_structural_padding(monkeypatch):
|
||||
cu_seqlens_k=torch.tensor([0, valid], dtype=torch.int32),
|
||||
)
|
||||
|
||||
monkeypatch.setattr(impl.sdpa_fallback, "forward", impl.sdpa_fallback.forward_cuda)
|
||||
|
||||
baseline = impl.forward_cuda(query, query, query, AttentionMetadata(packed_padding=padding, extra={}))
|
||||
perturbed_input = query.clone()
|
||||
perturbed_input[:, valid:] += 100.0
|
||||
@@ -384,22 +245,3 @@ def test_sdpa_fallback_never_attends_the_structural_padding(monkeypatch):
|
||||
assert baseline.shape == query.shape
|
||||
torch.testing.assert_close(baseline[:, :valid], perturbed[:, :valid])
|
||||
assert torch.count_nonzero(baseline[:, valid:]) == 0
|
||||
|
||||
|
||||
def test_h3_geometry_keeps_prefix_segments_pure_and_tiles_video_3d():
|
||||
partition, sizes, non_pad, untile, prefix_blocks, video_blocks = _get_h3_tile_metadata(
|
||||
(5, 70, 9), (5, 6, 6), torch.device("cpu")
|
||||
)
|
||||
assert sizes[:4].tolist() == [5, 64, 6, 9]
|
||||
assert prefix_blocks == 4
|
||||
assert video_blocks == 8
|
||||
assert int(sizes.sum()) == 5 + 70 + 9 + 5 * 6 * 6
|
||||
assert partition.numel() == non_pad.numel() == untile.numel() == int(sizes.sum())
|
||||
|
||||
|
||||
def test_h3_block_map_makes_prefix_queries_dense_and_prefix_keys_exempt():
|
||||
scores = torch.arange(1 * 2 * 5 * 5, dtype=torch.float32).reshape(1, 2, 5, 5)
|
||||
block_map = _build_h3_block_map(scores, num_prefix_blocks=2, num_video_blocks=3, topk=1)
|
||||
assert block_map[:, :, :2].all()
|
||||
assert block_map[..., :2].all()
|
||||
assert (block_map[:, :, 2:, 2:].sum(dim=-1) == 1).all()
|
||||
@@ -0,0 +1,136 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.diffusion.attention import layer
|
||||
from vllm_omni.diffusion.attention.backends.fastvideo_vsa import FastVideoVSABackend, FastVideoVSAImpl
|
||||
from vllm_omni.diffusion.attention.backends.sdpa import SDPABackend
|
||||
from vllm_omni.diffusion.attention.parallel.base import NoParallelAttention
|
||||
from vllm_omni.diffusion.data import AttentionSpec
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
|
||||
class _SpecializedImpl(FastVideoVSAImpl):
|
||||
pass
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def local_attention(monkeypatch):
|
||||
monkeypatch.setattr(layer, "get_current_diffusion_config_or_none", lambda: None)
|
||||
monkeypatch.setattr(layer, "build_parallel_attention_strategy", lambda **kwargs: NoParallelAttention())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("explicit", [False, True])
|
||||
@pytest.mark.parametrize("override", [False, True])
|
||||
def test_specialization_preserves_selected_backend_and_options(mocker, local_attention, explicit, override):
|
||||
spec = AttentionSpec(backend="FASTVIDEO_VSA", fastvideo_vsa_topk=7) if explicit else None
|
||||
select = mocker.patch.object(layer, "get_attn_backend_for_role", return_value=(FastVideoVSABackend, spec))
|
||||
attention = layer.Attention(
|
||||
num_heads=2,
|
||||
head_size=8,
|
||||
softmax_scale=8**-0.5,
|
||||
causal=False,
|
||||
role="video.self",
|
||||
role_category="self",
|
||||
qkv_layout="BSND",
|
||||
impl_overrides={"FASTVIDEO_VSA": _SpecializedImpl} if override else None,
|
||||
)
|
||||
|
||||
assert attention.attn_backend is FastVideoVSABackend
|
||||
assert type(attention.attention) is (_SpecializedImpl if override else FastVideoVSAImpl)
|
||||
assert attention.attention.topk == (7 if explicit else 64)
|
||||
assert attention.attention.qkv_layout == "BSND"
|
||||
assert attention.backend_pref == "FASTVIDEO_VSA"
|
||||
assert attention.backend_explicit is explicit
|
||||
assert attention.attn_spec is spec
|
||||
assert isinstance(attention.parallel_strategy, NoParallelAttention)
|
||||
select.assert_called_once_with(
|
||||
role="video.self", head_size=8, attention_config=None, role_category="self", allow_trtllm_default=False
|
||||
)
|
||||
|
||||
|
||||
def test_unselected_specialization_does_not_replace_dense_backend(mocker, local_attention):
|
||||
mocker.patch.object(layer, "get_attn_backend_for_role", return_value=(SDPABackend, None))
|
||||
attention = layer.Attention(
|
||||
num_heads=2,
|
||||
head_size=8,
|
||||
softmax_scale=8**-0.5,
|
||||
causal=False,
|
||||
impl_overrides={"FASTVIDEO_VSA": _SpecializedImpl},
|
||||
)
|
||||
assert attention.attn_backend is SDPABackend
|
||||
assert type(attention.attention) is SDPABackend.get_impl_cls()
|
||||
|
||||
|
||||
def test_specialization_does_not_bypass_selection_errors(mocker, local_attention):
|
||||
mocker.patch.object(layer, "get_attn_backend_for_role", side_effect=ImportError("optional kernel unavailable"))
|
||||
with pytest.raises(ImportError, match="optional kernel unavailable"):
|
||||
layer.Attention(
|
||||
num_heads=2,
|
||||
head_size=8,
|
||||
softmax_scale=8**-0.5,
|
||||
causal=False,
|
||||
impl_overrides={"FASTVIDEO_VSA": _SpecializedImpl},
|
||||
)
|
||||
|
||||
|
||||
def test_specialization_must_extend_selected_implementation(mocker, local_attention):
|
||||
mocker.patch.object(layer, "get_attn_backend_for_role", return_value=(FastVideoVSABackend, None))
|
||||
with pytest.raises(TypeError, match="must subclass the selected implementation"):
|
||||
layer.Attention(
|
||||
num_heads=2,
|
||||
head_size=8,
|
||||
softmax_scale=8**-0.5,
|
||||
causal=False,
|
||||
impl_overrides={"FASTVIDEO_VSA": SDPABackend.get_impl_cls()},
|
||||
)
|
||||
|
||||
|
||||
def test_platform_backend_capabilities_survive_model_specialization(mocker, local_attention):
|
||||
from vllm_omni.diffusion.models.minimax_h3.attention.fastvideo_h3 import MiniMaxH3VSAImpl
|
||||
|
||||
class PlatformBackend(FastVideoVSABackend):
|
||||
@classmethod
|
||||
def supports_packed_mask_free(cls):
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def get_supported_head_sizes():
|
||||
return [128]
|
||||
|
||||
# A registry/platform may refine capabilities without changing the
|
||||
# implementation. The model must not replace those capability decisions.
|
||||
mocker.patch.object(layer, "get_attn_backend_for_role", return_value=(PlatformBackend, None))
|
||||
attention = layer.Attention(
|
||||
num_heads=2,
|
||||
head_size=128,
|
||||
softmax_scale=128**-0.5,
|
||||
causal=False,
|
||||
impl_overrides={"FASTVIDEO_VSA": MiniMaxH3VSAImpl},
|
||||
)
|
||||
assert attention.attn_backend is PlatformBackend
|
||||
assert not attention.attn_backend.supports_packed_mask_free()
|
||||
assert attention.attn_backend.get_supported_head_sizes() == [128]
|
||||
assert type(attention.attention) is MiniMaxH3VSAImpl
|
||||
|
||||
|
||||
def test_incompatible_platform_implementation_is_not_discarded(mocker, local_attention):
|
||||
class PlatformImpl(FastVideoVSAImpl):
|
||||
pass
|
||||
|
||||
class PlatformBackend(FastVideoVSABackend):
|
||||
@staticmethod
|
||||
def get_impl_cls():
|
||||
return PlatformImpl
|
||||
|
||||
mocker.patch.object(layer, "get_attn_backend_for_role", return_value=(PlatformBackend, None))
|
||||
with pytest.raises(TypeError, match="selected implementation .*PlatformImpl"):
|
||||
layer.Attention(
|
||||
num_heads=2,
|
||||
head_size=8,
|
||||
softmax_scale=8**-0.5,
|
||||
causal=False,
|
||||
impl_overrides={"FASTVIDEO_VSA": _SpecializedImpl},
|
||||
)
|
||||
+635
@@ -0,0 +1,635 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import datetime
|
||||
from contextlib import nullcontext
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.helpers.runtime import get_distributed_init_method
|
||||
from vllm_omni.diffusion.cache.seacache import (
|
||||
SeaCacheBackend,
|
||||
SeaCacheConfig,
|
||||
SeaCacheRootHook,
|
||||
apply_sea_cache_hook,
|
||||
)
|
||||
from vllm_omni.diffusion.cache.seacache.sea_filter import (
|
||||
apply_sea_filter,
|
||||
extrapolate_residual,
|
||||
indicator_distance,
|
||||
)
|
||||
from vllm_omni.diffusion.cache.selector import get_cache_backend
|
||||
from vllm_omni.diffusion.cache.teacache.extractors import EXTRACTOR_REGISTRY, CacheContext, get_extractor
|
||||
from vllm_omni.diffusion.data import DiffusionCacheConfig
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class TinyCosmos3Transformer(torch.nn.Module):
|
||||
"""Small model that exposes a cache-neutral execution boundary."""
|
||||
|
||||
def _run_gen_layers(self, hidden_gen: torch.Tensor) -> torch.Tensor:
|
||||
return hidden_gen * 2
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
timestep: torch.Tensor,
|
||||
text_ids: torch.Tensor | None = None,
|
||||
text_mask: torch.Tensor | None = None,
|
||||
video_shape: tuple[int, int, int] | None = None,
|
||||
noisy_frame_mask: torch.Tensor | None = None,
|
||||
control_latents: list[torch.Tensor] | torch.Tensor | None = None,
|
||||
**kwargs,
|
||||
) -> torch.Tensor:
|
||||
del timestep, text_ids, text_mask, video_shape, noisy_frame_mask
|
||||
if kwargs:
|
||||
raise TypeError(f"Unexpected tiny transformer kwargs: {sorted(kwargs)}")
|
||||
controls = (
|
||||
[]
|
||||
if control_latents is None
|
||||
else [control_latents]
|
||||
if isinstance(control_latents, torch.Tensor)
|
||||
else list(control_latents)
|
||||
)
|
||||
inputs = [*controls, hidden_states]
|
||||
gen_input = torch.cat(
|
||||
[value.movedim(1, -1).flatten(1, 3) for value in inputs],
|
||||
dim=1,
|
||||
)
|
||||
return self._run_gen_layers(gen_input)
|
||||
|
||||
|
||||
def _extract_tiny_cosmos3_context(
|
||||
module: TinyCosmos3Transformer,
|
||||
hidden_states: torch.Tensor,
|
||||
timestep: torch.Tensor,
|
||||
text_ids: torch.Tensor | None = None,
|
||||
text_mask: torch.Tensor | None = None,
|
||||
video_shape: tuple[int, int, int] | None = None,
|
||||
noisy_frame_mask: torch.Tensor | None = None,
|
||||
control_latents: list[torch.Tensor] | torch.Tensor | None = None,
|
||||
**kwargs,
|
||||
) -> CacheContext:
|
||||
del timestep, text_ids, text_mask, video_shape
|
||||
if kwargs:
|
||||
raise TypeError(f"Unexpected tiny transformer kwargs: {sorted(kwargs)}")
|
||||
controls = (
|
||||
[]
|
||||
if control_latents is None
|
||||
else [control_latents]
|
||||
if isinstance(control_latents, torch.Tensor)
|
||||
else list(control_latents)
|
||||
)
|
||||
vision_items = [*controls, hidden_states]
|
||||
gen_input = torch.cat(
|
||||
[value.movedim(1, -1).flatten(1, 3) for value in vision_items],
|
||||
dim=1,
|
||||
)
|
||||
return CacheContext(
|
||||
modulated_input=gen_input,
|
||||
hidden_states=gen_input,
|
||||
encoder_hidden_states=None,
|
||||
temb=torch.zeros_like(gen_input[:, 0]),
|
||||
run_transformer_blocks=lambda: (module._run_gen_layers(gen_input),),
|
||||
postprocess=lambda output: output,
|
||||
extra_states={
|
||||
"sea_cache_latents": vision_items,
|
||||
"sea_cache_noisy_frame_mask": noisy_frame_mask,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
class Cosmos3OmniDiffusersPipeline:
|
||||
def __init__(self) -> None:
|
||||
self.transformer = TinyCosmos3Transformer()
|
||||
self._current_step_index: int | None = None
|
||||
self._current_sigma: float | None = None
|
||||
self._num_timesteps: int | None = None
|
||||
|
||||
@property
|
||||
def current_step_index(self) -> int | None:
|
||||
return self._current_step_index
|
||||
|
||||
@property
|
||||
def current_sigma(self) -> float | None:
|
||||
return self._current_sigma
|
||||
|
||||
@property
|
||||
def num_timesteps(self) -> int | None:
|
||||
return self._num_timesteps
|
||||
|
||||
|
||||
def _latent(value: float) -> torch.Tensor:
|
||||
return torch.full((1, 2, 2, 2, 2), value)
|
||||
|
||||
|
||||
def _run_step(
|
||||
transformer: TinyCosmos3Transformer,
|
||||
timestep: int,
|
||||
value: float,
|
||||
*,
|
||||
hook: SeaCacheRootHook | None = None,
|
||||
context: str = "cond",
|
||||
control: float | None = None,
|
||||
noisy_frame_mask: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
controls = None if control is None else [_latent(control)]
|
||||
cache_context = hook.cache_context(context) if hook is not None else nullcontext()
|
||||
with torch.inference_mode(), cache_context:
|
||||
return transformer(
|
||||
hidden_states=_latent(value),
|
||||
timestep=torch.tensor([timestep]),
|
||||
noisy_frame_mask=noisy_frame_mask,
|
||||
control_latents=controls,
|
||||
)
|
||||
|
||||
|
||||
def _apply_test_hook(
|
||||
transformer: TinyCosmos3Transformer,
|
||||
metadata: SimpleNamespace,
|
||||
config: SeaCacheConfig | None = None,
|
||||
) -> SeaCacheRootHook:
|
||||
return apply_sea_cache_hook(
|
||||
transformer,
|
||||
config or SeaCacheConfig(threshold=100.0),
|
||||
current_step_callback=lambda: metadata.step,
|
||||
current_sigma_callback=lambda: metadata.sigma,
|
||||
num_inference_steps_callback=lambda: metadata.num_steps,
|
||||
extractor_fn=_extract_tiny_cosmos3_context,
|
||||
)
|
||||
|
||||
|
||||
def test_config_validation() -> None:
|
||||
assert SeaCacheConfig().threshold == 0.25
|
||||
with pytest.raises(ValueError, match="residual_order"):
|
||||
SeaCacheConfig(residual_order=-1)
|
||||
with pytest.raises(ValueError, match="max_consecutive_cached"):
|
||||
SeaCacheConfig(max_consecutive_cached=-1)
|
||||
|
||||
|
||||
def test_sea_filter_matches_reference_equation() -> None:
|
||||
hidden = torch.randn(3, 4, 5, 2, dtype=torch.float32)
|
||||
sigma = 0.4
|
||||
power_exp = 3.0
|
||||
|
||||
spectrum = torch.fft.fftn(hidden, dim=(0, 1, 2))
|
||||
gain = None
|
||||
for axis in (0, 1, 2):
|
||||
frequencies = torch.fft.fftfreq(hidden.shape[axis], dtype=torch.float32)
|
||||
clean_power = 1.0 / (frequencies.abs().pow(power_exp) + 1e-16)
|
||||
axis_gain = (1.0 - sigma) * clean_power / ((1.0 - sigma) ** 2 * clean_power + sigma**2 + 1e-16)
|
||||
shape = [1] * hidden.ndim
|
||||
shape[axis] = hidden.shape[axis]
|
||||
gain = axis_gain.reshape(shape) if gain is None else gain * axis_gain.reshape(shape)
|
||||
assert gain is not None
|
||||
gain = gain / gain.mean()
|
||||
expected = torch.fft.ifftn(spectrum * gain, dim=(0, 1, 2)).real
|
||||
|
||||
torch.testing.assert_close(
|
||||
apply_sea_filter(hidden, sigma=sigma, power_exp=power_exp),
|
||||
expected,
|
||||
)
|
||||
|
||||
|
||||
def test_indicator_distance_and_linear_extrapolation() -> None:
|
||||
previous = [torch.ones(2, 2)]
|
||||
current = [torch.full((2, 2), 1.5)]
|
||||
assert indicator_distance(current, previous) == pytest.approx(0.5)
|
||||
assert indicator_distance([torch.ones(3)], previous) == float("inf")
|
||||
|
||||
history = [
|
||||
(0, torch.full((2, 2), 2.0)),
|
||||
(2, torch.full((2, 2), 6.0)),
|
||||
]
|
||||
torch.testing.assert_close(
|
||||
extrapolate_residual(history, step=3, order=1),
|
||||
torch.full((2, 2), 8.0),
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
extrapolate_residual(history, step=3, order=0),
|
||||
torch.full((2, 2), 6.0),
|
||||
)
|
||||
quadratic_history = [
|
||||
(0, torch.zeros(2, 2)),
|
||||
(1, torch.ones(2, 2)),
|
||||
(2, torch.full((2, 2), 4.0)),
|
||||
]
|
||||
torch.testing.assert_close(
|
||||
extrapolate_residual(quadratic_history, step=3, order=2),
|
||||
torch.full((2, 2), 9.0),
|
||||
)
|
||||
|
||||
|
||||
def test_hook_skips_middle_steps_and_forces_endpoints() -> None:
|
||||
transformer = TinyCosmos3Transformer()
|
||||
metadata = SimpleNamespace(step=0, sigma=1.0, num_steps=4)
|
||||
hook = _apply_test_hook(
|
||||
transformer,
|
||||
metadata,
|
||||
SeaCacheConfig(threshold=100.0, max_consecutive_cached=2),
|
||||
)
|
||||
hook.refresh(transformer)
|
||||
|
||||
for step, timestep in enumerate((1000, 750, 500, 250)):
|
||||
metadata.step = step
|
||||
metadata.sigma = timestep / 1000
|
||||
_run_step(transformer, timestep, 1.0 - step * 0.01, hook=hook)
|
||||
|
||||
assert hook.full_count == 2
|
||||
assert hook.skip_count == 2
|
||||
assert [step for step, _ in hook.state_manager._states["cond"].history] == [0, 3]
|
||||
|
||||
|
||||
def test_parameter_sharded_hook_skips_when_all_shard_ranks_agree(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
from vllm_omni.diffusion.distributed import parallel_state
|
||||
|
||||
transformer = TinyCosmos3Transformer()
|
||||
metadata = SimpleNamespace(step=0, sigma=1.0, num_steps=3)
|
||||
hook = _apply_test_hook(transformer, metadata)
|
||||
hook._parameter_sharded = True
|
||||
fs_group = object()
|
||||
reduced_groups: list[object] = []
|
||||
|
||||
monkeypatch.setattr(torch.distributed, "is_available", lambda: True)
|
||||
monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True)
|
||||
monkeypatch.setattr(
|
||||
parallel_state,
|
||||
"get_fs_group",
|
||||
lambda: SimpleNamespace(world_size=2, device_group=fs_group),
|
||||
)
|
||||
monkeypatch.setattr(parallel_state, "get_sequence_parallel_world_size", lambda: 1)
|
||||
|
||||
def all_reduce(decision, *, op, group):
|
||||
assert op == torch.distributed.ReduceOp.MAX
|
||||
reduced_groups.append(group)
|
||||
|
||||
monkeypatch.setattr(torch.distributed, "all_reduce", all_reduce)
|
||||
|
||||
_run_step(transformer, 1000, 1.0, hook=hook)
|
||||
metadata.step = 1
|
||||
metadata.sigma = 0.5
|
||||
_run_step(transformer, 500, 0.99, hook=hook)
|
||||
|
||||
assert hook.full_count == 1
|
||||
assert hook.skip_count == 1
|
||||
assert reduced_groups == [fs_group, fs_group]
|
||||
|
||||
|
||||
def test_parameter_sharded_peer_forces_full_compute(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
from vllm_omni.diffusion.distributed import parallel_state
|
||||
|
||||
transformer = TinyCosmos3Transformer()
|
||||
metadata = SimpleNamespace(step=0, sigma=1.0, num_steps=3)
|
||||
hook = _apply_test_hook(transformer, metadata)
|
||||
hook._parameter_sharded = True
|
||||
|
||||
monkeypatch.setattr(torch.distributed, "is_available", lambda: True)
|
||||
monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True)
|
||||
monkeypatch.setattr(
|
||||
parallel_state,
|
||||
"get_fs_group",
|
||||
lambda: SimpleNamespace(world_size=2, device_group=object()),
|
||||
)
|
||||
monkeypatch.setattr(parallel_state, "get_sequence_parallel_world_size", lambda: 1)
|
||||
reduce_count = 0
|
||||
|
||||
def all_reduce(decision, *, op, group):
|
||||
nonlocal reduce_count
|
||||
del op, group
|
||||
reduce_count += 1
|
||||
if reduce_count == 2:
|
||||
decision.fill_(1)
|
||||
|
||||
monkeypatch.setattr(torch.distributed, "all_reduce", all_reduce)
|
||||
|
||||
_run_step(transformer, 1000, 1.0, hook=hook)
|
||||
metadata.step = 1
|
||||
metadata.sigma = 0.5
|
||||
_run_step(transformer, 500, 0.99, hook=hook)
|
||||
|
||||
assert hook.full_count == 2
|
||||
assert hook.skip_count == 0
|
||||
assert hook.state_manager._states["cond"].accumulated_distance == 0.0
|
||||
|
||||
|
||||
def _uneven_cfg_sharded_worker(
|
||||
rank: int,
|
||||
init_method: str,
|
||||
result_queue: torch.multiprocessing.Queue,
|
||||
) -> None:
|
||||
from vllm_omni.diffusion.distributed import parallel_state
|
||||
|
||||
torch.distributed.init_process_group(
|
||||
"gloo",
|
||||
init_method=init_method,
|
||||
rank=rank,
|
||||
world_size=4,
|
||||
timeout=datetime.timedelta(seconds=10),
|
||||
)
|
||||
try:
|
||||
fs_groups = [
|
||||
(ranks, torch.distributed.new_group(ranks, timeout=datetime.timedelta(seconds=10)))
|
||||
for ranks in ([0, 1], [2, 3])
|
||||
]
|
||||
sp_groups = [
|
||||
(ranks, torch.distributed.new_group(ranks, timeout=datetime.timedelta(seconds=10)))
|
||||
for ranks in ([0, 1], [2, 3])
|
||||
]
|
||||
cfg_groups = [
|
||||
(ranks, torch.distributed.new_group(ranks, timeout=datetime.timedelta(seconds=10)))
|
||||
for ranks in ([0, 2], [1, 3])
|
||||
]
|
||||
fs_group = next(group for ranks, group in fs_groups if rank in ranks)
|
||||
sp_group = next(group for ranks, group in sp_groups if rank in ranks)
|
||||
cfg_group = next(group for ranks, group in cfg_groups if rank in ranks)
|
||||
parallel_state._FS = SimpleNamespace(world_size=2, device_group=fs_group)
|
||||
parallel_state._SP = SimpleNamespace(world_size=2, device_group=sp_group)
|
||||
|
||||
hook = SeaCacheRootHook(SeaCacheConfig())
|
||||
hook._parameter_sharded = True
|
||||
hook._synchronize_compute(False, torch.device("cpu"))
|
||||
if rank in (0, 1):
|
||||
hook._synchronize_compute(False, torch.device("cpu"))
|
||||
|
||||
gathered = [torch.zeros(1) for _ in range(2)]
|
||||
torch.distributed.all_gather(
|
||||
gathered,
|
||||
torch.tensor([rank], dtype=torch.float32),
|
||||
group=cfg_group,
|
||||
)
|
||||
result_queue.put((rank, [int(value.item()) for value in gathered]))
|
||||
finally:
|
||||
parallel_state._FS = None
|
||||
parallel_state._SP = None
|
||||
torch.distributed.destroy_process_group()
|
||||
|
||||
|
||||
def test_parameter_sharded_hook_handles_uneven_cfg_branch_dispatch() -> None:
|
||||
mp_context = torch.multiprocessing.get_context("spawn")
|
||||
manager = mp_context.Manager()
|
||||
result_queue = manager.Queue()
|
||||
torch.multiprocessing.spawn(
|
||||
_uneven_cfg_sharded_worker,
|
||||
args=(get_distributed_init_method("seacache_uneven_cfg_"), result_queue),
|
||||
nprocs=4,
|
||||
)
|
||||
|
||||
results = sorted(result_queue.get(timeout=2) for _ in range(4))
|
||||
assert results == [
|
||||
(0, [0, 2]),
|
||||
(1, [1, 3]),
|
||||
(2, [0, 2]),
|
||||
(3, [1, 3]),
|
||||
]
|
||||
|
||||
|
||||
def _hybrid_sp_worker(
|
||||
rank: int,
|
||||
init_method: str,
|
||||
result_queue: torch.multiprocessing.Queue,
|
||||
) -> None:
|
||||
from vllm_omni.diffusion.distributed import parallel_state
|
||||
|
||||
torch.distributed.init_process_group(
|
||||
"gloo",
|
||||
init_method=init_method,
|
||||
rank=rank,
|
||||
world_size=4,
|
||||
timeout=datetime.timedelta(seconds=10),
|
||||
)
|
||||
try:
|
||||
fs_groups = [
|
||||
(ranks, torch.distributed.new_group(ranks, timeout=datetime.timedelta(seconds=10)))
|
||||
for ranks in ([0, 1], [2, 3])
|
||||
]
|
||||
ulysses_groups = [
|
||||
(ranks, torch.distributed.new_group(ranks, timeout=datetime.timedelta(seconds=10)))
|
||||
for ranks in ([0, 1], [2, 3])
|
||||
]
|
||||
ring_groups = [
|
||||
(ranks, torch.distributed.new_group(ranks, timeout=datetime.timedelta(seconds=10)))
|
||||
for ranks in ([0, 2], [1, 3])
|
||||
]
|
||||
sp_group = torch.distributed.new_group(
|
||||
[0, 1, 2, 3],
|
||||
timeout=datetime.timedelta(seconds=10),
|
||||
)
|
||||
fs_group = next(group for ranks, group in fs_groups if rank in ranks)
|
||||
ulysses_group = next(group for ranks, group in ulysses_groups if rank in ranks)
|
||||
ring_group = next(group for ranks, group in ring_groups if rank in ranks)
|
||||
parallel_state._FS = SimpleNamespace(world_size=2, device_group=fs_group)
|
||||
parallel_state._SP = SimpleNamespace(
|
||||
world_size=4,
|
||||
device_group=sp_group,
|
||||
ulysses_world_size=2,
|
||||
ulysses_group=ulysses_group,
|
||||
ring_world_size=2,
|
||||
ring_group=ring_group,
|
||||
)
|
||||
|
||||
compute = rank == 3
|
||||
non_sharded_hook = SeaCacheRootHook(SeaCacheConfig())
|
||||
non_sharded_result = non_sharded_hook._synchronize_compute(compute, torch.device("cpu"))
|
||||
|
||||
sharded_hook = SeaCacheRootHook(SeaCacheConfig())
|
||||
sharded_hook._parameter_sharded = True
|
||||
sharded_result = sharded_hook._synchronize_compute(compute, torch.device("cpu"))
|
||||
result_queue.put((rank, non_sharded_result, sharded_result))
|
||||
finally:
|
||||
parallel_state._FS = None
|
||||
parallel_state._SP = None
|
||||
torch.distributed.destroy_process_group()
|
||||
|
||||
|
||||
def test_hook_synchronizes_full_hybrid_sequence_parallel_group() -> None:
|
||||
mp_context = torch.multiprocessing.get_context("spawn")
|
||||
manager = mp_context.Manager()
|
||||
result_queue = manager.Queue()
|
||||
torch.multiprocessing.spawn(
|
||||
_hybrid_sp_worker,
|
||||
args=(get_distributed_init_method("seacache_hybrid_sp_"), result_queue),
|
||||
nprocs=4,
|
||||
)
|
||||
|
||||
results = sorted(result_queue.get(timeout=2) for _ in range(4))
|
||||
assert results == [
|
||||
(0, True, True),
|
||||
(1, True, True),
|
||||
(2, True, True),
|
||||
(3, True, True),
|
||||
]
|
||||
|
||||
|
||||
def test_parameter_sharded_hook_fails_open_without_distributed_state(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
hook = SeaCacheRootHook(SeaCacheConfig())
|
||||
hook._parameter_sharded = True
|
||||
monkeypatch.setattr(torch.distributed, "is_available", lambda: True)
|
||||
monkeypatch.setattr(torch.distributed, "is_initialized", lambda: False)
|
||||
|
||||
assert hook._synchronize_compute(False, torch.device("cpu")) is True
|
||||
|
||||
|
||||
def test_hook_keeps_three_transfer_branches_separate() -> None:
|
||||
transformer = TinyCosmos3Transformer()
|
||||
metadata = SimpleNamespace(step=0, sigma=1.0, num_steps=3)
|
||||
hook = _apply_test_hook(transformer, metadata)
|
||||
hook.refresh(transformer)
|
||||
|
||||
contexts = ("cond", "cond_no_control", "uncond")
|
||||
for step, (timestep, value) in enumerate(((1000, 1.0), (500, 0.9))):
|
||||
metadata.step = step
|
||||
metadata.sigma = timestep / 1000
|
||||
_run_step(transformer, timestep, value, hook=hook, context=contexts[0], control=0.5)
|
||||
_run_step(transformer, timestep, value, hook=hook, context=contexts[1])
|
||||
_run_step(transformer, timestep, value, hook=hook, context=contexts[2], control=0.5)
|
||||
|
||||
assert set(hook.state_manager._states) == set(contexts)
|
||||
assert hook.full_count == 3
|
||||
assert hook.skip_count == 3
|
||||
assert len(hook.state_manager._states["cond"].previous_indicator) == 2
|
||||
assert len(hook.state_manager._states["cond_no_control"].previous_indicator) == 1
|
||||
|
||||
|
||||
def test_hook_fails_open_without_noisy_vision() -> None:
|
||||
transformer = TinyCosmos3Transformer()
|
||||
metadata = SimpleNamespace(step=0, sigma=1.0, num_steps=3)
|
||||
hook = _apply_test_hook(transformer, metadata)
|
||||
hook.refresh(transformer)
|
||||
all_clean = torch.zeros(1, 1, 2, 1, 1)
|
||||
|
||||
_run_step(transformer, 1000, 1.0, hook=hook, noisy_frame_mask=all_clean)
|
||||
metadata.step = 1
|
||||
metadata.sigma = 0.5
|
||||
_run_step(transformer, 500, 0.9, hook=hook, noisy_frame_mask=all_clean)
|
||||
|
||||
assert hook.full_count == 0
|
||||
assert hook.skip_count == 0
|
||||
|
||||
|
||||
def test_conditioning_only_model_failure_runs_forward_once(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
transformer = TinyCosmos3Transformer()
|
||||
metadata = SimpleNamespace(step=0, sigma=1.0, num_steps=3)
|
||||
hook = _apply_test_hook(transformer, metadata)
|
||||
hook.refresh(transformer)
|
||||
forward_calls = 0
|
||||
|
||||
def fail_forward(hidden_gen: torch.Tensor) -> torch.Tensor:
|
||||
nonlocal forward_calls
|
||||
forward_calls += 1
|
||||
raise RuntimeError("model forward failed")
|
||||
|
||||
monkeypatch.setattr(transformer, "_run_gen_layers", fail_forward)
|
||||
all_clean = torch.zeros(1, 1, 2, 1, 1)
|
||||
|
||||
with pytest.raises(RuntimeError, match="model forward failed"):
|
||||
_run_step(transformer, 1000, 1.0, hook=hook, noisy_frame_mask=all_clean)
|
||||
|
||||
assert forward_calls == 1
|
||||
|
||||
|
||||
def test_hook_fails_open_without_explicit_context() -> None:
|
||||
transformer = TinyCosmos3Transformer()
|
||||
metadata = SimpleNamespace(step=0, sigma=1.0, num_steps=2)
|
||||
hook = _apply_test_hook(transformer, metadata)
|
||||
hook.refresh(transformer)
|
||||
|
||||
_run_step(transformer, 1000, 1.0)
|
||||
|
||||
assert hook.full_count == 0
|
||||
assert hook.skip_count == 0
|
||||
assert hook.state_manager._states == {}
|
||||
|
||||
|
||||
def test_hook_uses_exact_sigma_callback(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
from vllm_omni.diffusion.cache.seacache import hook as hook_module
|
||||
|
||||
transformer = TinyCosmos3Transformer()
|
||||
metadata = SimpleNamespace(step=0, sigma=0.37, num_steps=2)
|
||||
observed_sigmas: list[float] = []
|
||||
original_filter = hook_module.apply_sea_filter
|
||||
|
||||
def recording_filter(hidden_states: torch.Tensor, sigma: float, power_exp: float) -> torch.Tensor:
|
||||
observed_sigmas.append(sigma)
|
||||
return original_filter(hidden_states, sigma, power_exp)
|
||||
|
||||
monkeypatch.setattr(hook_module, "apply_sea_filter", recording_filter)
|
||||
hook = _apply_test_hook(transformer, metadata)
|
||||
hook.refresh(transformer)
|
||||
_run_step(transformer, timestep=999, value=1.0, hook=hook)
|
||||
|
||||
assert observed_sigmas
|
||||
assert all(sigma == pytest.approx(0.37) for sigma in observed_sigmas)
|
||||
|
||||
|
||||
def test_backend_selector_and_refresh(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setitem(EXTRACTOR_REGISTRY, "TinyCosmos3Transformer", _extract_tiny_cosmos3_context)
|
||||
|
||||
class FSDPTinyCosmos3Transformer(TinyCosmos3Transformer):
|
||||
pass
|
||||
|
||||
assert get_extractor(FSDPTinyCosmos3Transformer) is _extract_tiny_cosmos3_context
|
||||
backend = get_cache_backend(
|
||||
"sea_cache",
|
||||
{
|
||||
"sea_threshold": 0.4,
|
||||
"sea_residual_order": 0,
|
||||
},
|
||||
)
|
||||
assert isinstance(backend, SeaCacheBackend)
|
||||
assert backend.config.sea_threshold == 0.4
|
||||
|
||||
pipeline = Cosmos3OmniDiffusersPipeline()
|
||||
backend.enable(pipeline)
|
||||
hook = pipeline.transformer._hook_registry.get_hook(SeaCacheRootHook._HOOK_NAME)
|
||||
assert isinstance(hook, SeaCacheRootHook)
|
||||
assert callable(getattr(pipeline, "_cache_context_factory", None))
|
||||
for name in ("_seacache_skip", "_seacache_record", "_seacache_residual", "_seacache_last_residual"):
|
||||
assert not hasattr(pipeline.transformer, name)
|
||||
pipeline._current_step_index = 0
|
||||
pipeline._current_sigma = 1.0
|
||||
pipeline._num_timesteps = 7
|
||||
_run_step(pipeline.transformer, 1000, 1.0, hook=hook)
|
||||
assert hook.full_count == 1
|
||||
|
||||
backend.refresh(pipeline, num_inference_steps=7)
|
||||
assert hook.full_count == 0
|
||||
assert hook.skip_count == 0
|
||||
assert hook.state_manager._states == {}
|
||||
|
||||
|
||||
def test_backend_uses_resolved_pipeline_metadata_not_refresh_argument(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setitem(EXTRACTOR_REGISTRY, "TinyCosmos3Transformer", _extract_tiny_cosmos3_context)
|
||||
pipeline = Cosmos3OmniDiffusersPipeline()
|
||||
backend = SeaCacheBackend(DiffusionCacheConfig())
|
||||
backend.enable(pipeline)
|
||||
hook = pipeline.transformer._hook_registry.get_hook(SeaCacheRootHook._HOOK_NAME)
|
||||
assert isinstance(hook, SeaCacheRootHook)
|
||||
backend.refresh(pipeline, num_inference_steps=35)
|
||||
pipeline._num_timesteps = 50
|
||||
pipeline._current_step_index = 0
|
||||
pipeline._current_sigma = 1.0
|
||||
_run_step(pipeline.transformer, 1000, 1.0, hook=hook)
|
||||
assert hook.full_count == 1
|
||||
assert hook.num_inference_steps_callback is not None
|
||||
assert hook.num_inference_steps_callback() == 50
|
||||
|
||||
pipeline._current_step_index = 1
|
||||
pipeline._current_sigma = 0.98
|
||||
_run_step(pipeline.transformer, 980, 0.99, hook=hook)
|
||||
assert hook.full_count == 1
|
||||
assert hook.skip_count == 1
|
||||
|
||||
|
||||
def test_shared_config_defaults() -> None:
|
||||
config = DiffusionCacheConfig()
|
||||
assert config.sea_threshold == 0.25
|
||||
assert config.sea_residual_order == 1
|
||||
assert config.sea_max_consecutive_cached == 2
|
||||
assert config.sea_power_exp == 3.0
|
||||
+5
-5
@@ -68,12 +68,12 @@ class BaseExtractorTest(ABC):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_module(self):
|
||||
def get_module(self, module, /):
|
||||
"""Return model module instance."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_sample_inputs(self):
|
||||
def get_sample_inputs(self, sample_inputs, /):
|
||||
"""Return sample inputs for model."""
|
||||
pass
|
||||
|
||||
@@ -737,10 +737,10 @@ class TestMiniMaxH3Extractor(BaseExtractorTest):
|
||||
monkeypatch,
|
||||
):
|
||||
"""Strict SP must not mix local TeaCache state with gathered block rows."""
|
||||
from vllm_omni.diffusion.attention.ops import minimax_h3_modulation
|
||||
from vllm_omni.diffusion.cache.teacache.config import TeaCacheConfig
|
||||
from vllm_omni.diffusion.cache.teacache.hook import TeaCacheHook
|
||||
from vllm_omni.diffusion.distributed import parallel_state
|
||||
from vllm_omni.diffusion.layers import indexed_modulation
|
||||
|
||||
seq_len = sample_inputs["x"].shape[1]
|
||||
local_len = seq_len // 2
|
||||
@@ -844,13 +844,13 @@ class TestMiniMaxH3Extractor(BaseExtractorTest):
|
||||
)
|
||||
monkeypatch.setattr(parallel_state, "get_sp_group", lambda: sp_group)
|
||||
|
||||
original_indexed_scale_shift = minimax_h3_modulation.indexed_scale_shift_
|
||||
original_indexed_scale_shift = indexed_modulation.indexed_scale_shift_
|
||||
|
||||
def capture_indexed_scale_shift(hidden, shift, scale, indices):
|
||||
captured_modulation_indices.append(indices.clone())
|
||||
return original_indexed_scale_shift(hidden, shift, scale, indices)
|
||||
|
||||
monkeypatch.setattr(minimax_h3_modulation, "indexed_scale_shift_", capture_indexed_scale_shift)
|
||||
monkeypatch.setattr(indexed_modulation, "indexed_scale_shift_", capture_indexed_scale_shift)
|
||||
|
||||
local_img_pos = torch.tensor([0, 1])
|
||||
local_audio_pos = torch.tensor([2, 3])
|
||||
|
||||
@@ -0,0 +1,299 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from vllm_omni.diffusion.layers import qwen_select01_modulation
|
||||
from vllm_omni.diffusion.layers.qwen_select01_modulation import (
|
||||
can_use_qwen_select01_triton,
|
||||
fused_layernorm_select01,
|
||||
fused_residual_layernorm_select01,
|
||||
select01_modulation_native,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
@pytest.mark.parametrize(
|
||||
("platform_is_cuda", "tensor_is_cuda", "has_triton", "expected"),
|
||||
[
|
||||
(True, True, True, True),
|
||||
(False, True, True, False),
|
||||
(True, False, True, False),
|
||||
(True, True, False, False),
|
||||
],
|
||||
)
|
||||
def test_qwen_select01_triton_requires_cuda_platform(
|
||||
monkeypatch,
|
||||
platform_is_cuda: bool,
|
||||
tensor_is_cuda: bool,
|
||||
has_triton: bool,
|
||||
expected: bool,
|
||||
):
|
||||
class FakePlatform:
|
||||
@staticmethod
|
||||
def is_cuda() -> bool:
|
||||
return platform_is_cuda
|
||||
|
||||
class FakeTensor:
|
||||
is_cuda = tensor_is_cuda
|
||||
|
||||
monkeypatch.setattr(qwen_select01_modulation, "current_platform", FakePlatform())
|
||||
monkeypatch.setattr(qwen_select01_modulation, "HAS_TRITON", has_triton)
|
||||
|
||||
assert can_use_qwen_select01_triton(FakeTensor()) is expected
|
||||
|
||||
|
||||
def _make_inputs(
|
||||
batch_size: int,
|
||||
seq_len: int,
|
||||
hidden_size: int,
|
||||
dtype: torch.dtype,
|
||||
device: torch.device,
|
||||
) -> tuple[torch.Tensor, ...]:
|
||||
generator = torch.Generator(device=device)
|
||||
generator.manual_seed(202817)
|
||||
x = torch.randn(batch_size, seq_len, hidden_size, dtype=dtype, device=device, generator=generator)
|
||||
residual = torch.randn(batch_size, seq_len, hidden_size, dtype=dtype, device=device, generator=generator)
|
||||
residual_gate = torch.randn(batch_size, 1, hidden_size, dtype=dtype, device=device, generator=generator)
|
||||
mod_params = torch.randn(batch_size * 2, hidden_size * 3, dtype=dtype, device=device, generator=generator)
|
||||
index = torch.randint(0, 2, (batch_size, seq_len), dtype=torch.int64, device=device, generator=generator)
|
||||
return x, residual, residual_gate, mod_params, index
|
||||
|
||||
|
||||
def _reference_layernorm_select01(
|
||||
x: torch.Tensor,
|
||||
mod_params: torch.Tensor,
|
||||
index: torch.Tensor,
|
||||
eps: float,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
scale, shift, gate = select01_modulation_native(mod_params, index)
|
||||
out = F.layer_norm(x.float(), (x.shape[-1],), eps=eps).to(x.dtype)
|
||||
return out * (1 + scale) + shift, gate
|
||||
|
||||
|
||||
def _reference_residual_layernorm_select01(
|
||||
x: torch.Tensor,
|
||||
residual: torch.Tensor,
|
||||
residual_gate: torch.Tensor,
|
||||
mod_params: torch.Tensor,
|
||||
index: torch.Tensor,
|
||||
eps: float,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
scale, shift, gate = select01_modulation_native(mod_params, index)
|
||||
residual_out = residual + residual_gate * x
|
||||
out = F.layer_norm(residual_out.float(), (residual_out.shape[-1],), eps=eps).to(residual_out.dtype)
|
||||
return out * (1 + scale) + shift, residual_out, gate
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
|
||||
@pytest.mark.parametrize("broadcast_index", [False, True])
|
||||
def test_qwen_select01_native_fallback_matches_reference(dtype: torch.dtype, broadcast_index: bool):
|
||||
device = torch.device("cpu")
|
||||
x, residual, residual_gate, mod_params, index = _make_inputs(2, 17, 64, dtype, device)
|
||||
if broadcast_index:
|
||||
index = index[:1]
|
||||
eps = 1e-6
|
||||
|
||||
actual_norm, actual_gate = fused_layernorm_select01(x, mod_params, index, eps)
|
||||
ref_norm, ref_gate = _reference_layernorm_select01(x, mod_params, index, eps)
|
||||
torch.testing.assert_close(actual_norm, ref_norm, atol=1e-5, rtol=1e-5)
|
||||
torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0)
|
||||
|
||||
actual_norm, actual_residual, actual_gate = fused_residual_layernorm_select01(
|
||||
x,
|
||||
residual,
|
||||
residual_gate,
|
||||
mod_params,
|
||||
index,
|
||||
eps,
|
||||
)
|
||||
ref_norm, ref_residual, ref_gate = _reference_residual_layernorm_select01(
|
||||
x,
|
||||
residual,
|
||||
residual_gate,
|
||||
mod_params,
|
||||
index,
|
||||
eps,
|
||||
)
|
||||
torch.testing.assert_close(actual_norm, ref_norm, atol=1e-5, rtol=1e-5)
|
||||
torch.testing.assert_close(actual_residual, ref_residual, atol=1e-5, rtol=1e-5)
|
||||
torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
|
||||
@pytest.mark.parametrize(
|
||||
("dtype", "atol", "rtol"),
|
||||
[
|
||||
(torch.float32, 1e-5, 1e-5),
|
||||
(torch.bfloat16, 5e-2, 5e-2),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("seq_len", [1, 7, 128, 1024])
|
||||
def test_qwen_select01_triton_matches_reference(dtype: torch.dtype, atol: float, rtol: float, seq_len: int):
|
||||
device = torch.device("cuda:0")
|
||||
x, residual, residual_gate, mod_params, index = _make_inputs(2, seq_len, 256, dtype, device)
|
||||
eps = 1e-6
|
||||
|
||||
actual_norm, actual_gate = fused_layernorm_select01(x, mod_params, index, eps)
|
||||
ref_norm, ref_gate = _reference_layernorm_select01(x, mod_params, index, eps)
|
||||
torch.testing.assert_close(actual_norm, ref_norm, atol=atol, rtol=rtol)
|
||||
torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0)
|
||||
|
||||
actual_norm, actual_residual, actual_gate = fused_residual_layernorm_select01(
|
||||
x,
|
||||
residual,
|
||||
residual_gate.expand_as(residual),
|
||||
mod_params,
|
||||
index,
|
||||
eps,
|
||||
)
|
||||
ref_norm, ref_residual, ref_gate = _reference_residual_layernorm_select01(
|
||||
x,
|
||||
residual,
|
||||
residual_gate.expand_as(residual),
|
||||
mod_params,
|
||||
index,
|
||||
eps,
|
||||
)
|
||||
torch.testing.assert_close(actual_norm, ref_norm, atol=atol, rtol=rtol)
|
||||
torch.testing.assert_close(actual_residual, ref_residual, atol=atol, rtol=rtol)
|
||||
torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
|
||||
@pytest.mark.parametrize(
|
||||
("dtype", "atol", "rtol"),
|
||||
[
|
||||
(torch.float32, 1e-5, 1e-5),
|
||||
(torch.bfloat16, 5e-2, 5e-2),
|
||||
],
|
||||
)
|
||||
def test_qwen_select01_triton_matches_reference_for_structured_multi_batch_index(
|
||||
dtype: torch.dtype,
|
||||
atol: float,
|
||||
rtol: float,
|
||||
):
|
||||
device = torch.device("cuda:0")
|
||||
x, residual, residual_gate, mod_params, _ = _make_inputs(3, 16, 256, dtype, device)
|
||||
index = torch.tensor(
|
||||
[
|
||||
[0] * 16,
|
||||
[1] * 16,
|
||||
[0, 1] * 8,
|
||||
],
|
||||
dtype=torch.int64,
|
||||
device=device,
|
||||
)
|
||||
eps = 1e-6
|
||||
|
||||
actual_norm, actual_gate = fused_layernorm_select01(x, mod_params, index, eps)
|
||||
ref_norm, ref_gate = _reference_layernorm_select01(x, mod_params, index, eps)
|
||||
torch.testing.assert_close(actual_norm, ref_norm, atol=atol, rtol=rtol)
|
||||
torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0)
|
||||
|
||||
actual_norm, actual_residual, actual_gate = fused_residual_layernorm_select01(
|
||||
x,
|
||||
residual,
|
||||
residual_gate.expand_as(residual),
|
||||
mod_params,
|
||||
index,
|
||||
eps,
|
||||
)
|
||||
ref_norm, ref_residual, ref_gate = _reference_residual_layernorm_select01(
|
||||
x,
|
||||
residual,
|
||||
residual_gate.expand_as(residual),
|
||||
mod_params,
|
||||
index,
|
||||
eps,
|
||||
)
|
||||
torch.testing.assert_close(actual_norm, ref_norm, atol=atol, rtol=rtol)
|
||||
torch.testing.assert_close(actual_residual, ref_residual, atol=atol, rtol=rtol)
|
||||
torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
|
||||
@pytest.mark.parametrize(
|
||||
("dtype", "atol", "rtol"),
|
||||
[
|
||||
(torch.float32, 1e-5, 1e-5),
|
||||
(torch.bfloat16, 5e-2, 5e-2),
|
||||
],
|
||||
)
|
||||
def test_qwen_select01_triton_broadcasts_single_index_batch(dtype: torch.dtype, atol: float, rtol: float):
|
||||
device = torch.device("cuda:0")
|
||||
x, residual, residual_gate, mod_params, index = _make_inputs(2, 17, 256, dtype, device)
|
||||
index = index[:1]
|
||||
eps = 1e-6
|
||||
|
||||
actual_norm, actual_gate = fused_layernorm_select01(x, mod_params, index, eps)
|
||||
ref_norm, ref_gate = _reference_layernorm_select01(x, mod_params, index, eps)
|
||||
torch.testing.assert_close(actual_norm, ref_norm, atol=atol, rtol=rtol)
|
||||
torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0)
|
||||
|
||||
actual_norm, actual_residual, actual_gate = fused_residual_layernorm_select01(
|
||||
x,
|
||||
residual,
|
||||
residual_gate.expand_as(residual),
|
||||
mod_params,
|
||||
index,
|
||||
eps,
|
||||
)
|
||||
ref_norm, ref_residual, ref_gate = _reference_residual_layernorm_select01(
|
||||
x,
|
||||
residual,
|
||||
residual_gate.expand_as(residual),
|
||||
mod_params,
|
||||
index,
|
||||
eps,
|
||||
)
|
||||
torch.testing.assert_close(actual_norm, ref_norm, atol=atol, rtol=rtol)
|
||||
torch.testing.assert_close(actual_residual, ref_residual, atol=atol, rtol=rtol)
|
||||
torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
|
||||
def test_qwen_select01_cuda_falls_back_without_triton(monkeypatch):
|
||||
device = torch.device("cuda:0")
|
||||
x, residual, residual_gate, mod_params, index = _make_inputs(2, 7, 256, torch.bfloat16, device)
|
||||
residual_gate = residual_gate.expand_as(residual)
|
||||
eps = 1e-6
|
||||
|
||||
def fail_if_launched(*args, **kwargs):
|
||||
pytest.fail("Triton launcher must not run when HAS_TRITON is false")
|
||||
|
||||
monkeypatch.setattr(qwen_select01_modulation, "HAS_TRITON", False)
|
||||
monkeypatch.setattr(qwen_select01_modulation, "_launch_layernorm_select01", fail_if_launched)
|
||||
monkeypatch.setattr(qwen_select01_modulation, "_launch_residual_layernorm_select01", fail_if_launched)
|
||||
|
||||
actual_norm, actual_gate = fused_layernorm_select01(x, mod_params, index, eps)
|
||||
ref_norm, ref_gate = _reference_layernorm_select01(x, mod_params, index, eps)
|
||||
torch.testing.assert_close(actual_norm, ref_norm, atol=0, rtol=0)
|
||||
torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0)
|
||||
|
||||
actual_norm, actual_residual, actual_gate = fused_residual_layernorm_select01(
|
||||
x,
|
||||
residual,
|
||||
residual_gate,
|
||||
mod_params,
|
||||
index,
|
||||
eps,
|
||||
)
|
||||
ref_norm, ref_residual, ref_gate = _reference_residual_layernorm_select01(
|
||||
x,
|
||||
residual,
|
||||
residual_gate,
|
||||
mod_params,
|
||||
index,
|
||||
eps,
|
||||
)
|
||||
torch.testing.assert_close(actual_norm, ref_norm, atol=0, rtol=0)
|
||||
torch.testing.assert_close(actual_residual, ref_residual, atol=0, rtol=0)
|
||||
torch.testing.assert_close(actual_gate, ref_gate, atol=0, rtol=0)
|
||||
@@ -440,29 +440,86 @@ def test_transformer_rejects_prompt_tuning():
|
||||
BooguImageTransformer2DModel(od_config=_tiny_od_config(prompt_tuning_configs={"use_prompt_tuning": True}))
|
||||
|
||||
|
||||
def _native_to_checkpoint_name(name: str) -> str:
|
||||
"""Inverse of ``load_weights`` remapping: native param name -> diffusers name.
|
||||
_QKV_FANOUT = {
|
||||
"to_qkv": ("to_q", "to_k", "to_v"),
|
||||
"img_to_qkv": ("img_to_q", "img_to_k", "img_to_v"),
|
||||
"instruct_to_qkv": ("instruct_to_q", "instruct_to_k", "instruct_to_v"),
|
||||
}
|
||||
|
||||
|
||||
def _native_to_checkpoint_weights(name: str, param: torch.Tensor) -> list[tuple[str, torch.Tensor]]:
|
||||
"""Inverse of ``load_weights`` remapping: native param -> diffusers weights.
|
||||
|
||||
The diffusers checkpoint stores fused projections as separate matrices, so a
|
||||
single merged native param (QKV or FFN gate/up) fans out into several
|
||||
checkpoint weights. Returns ``(diffusers_name, value)`` pairs whose values,
|
||||
when fed through ``load_weights``, reassemble into the original ``param``.
|
||||
|
||||
- ``.to_out.<suffix>`` -> ``.to_out.0.<suffix>`` (diffusers ModuleList wrap).
|
||||
- promoted joint-attention projections move back under ``.processor.``.
|
||||
"""
|
||||
q_size = NUM_HEADS * HEAD_DIM
|
||||
kv_size = NUM_KV_HEADS * HEAD_DIM
|
||||
|
||||
# Fused QKV projections split back into per-matrix q/k/v weights.
|
||||
for token, (q_name, k_name, v_name) in _QKV_FANOUT.items():
|
||||
marker = f".{token}."
|
||||
if marker in name:
|
||||
q, k, v = param[:q_size], param[q_size : q_size + kv_size], param[q_size + kv_size :]
|
||||
results = []
|
||||
for sub_name, value in ((q_name, q), (k_name, k), (v_name, v)):
|
||||
ckpt_name = name.replace(marker, f".{sub_name}.")
|
||||
if token != "to_qkv":
|
||||
# Joint-attention projections live under `.processor.` upstream.
|
||||
ckpt_name = ckpt_name.replace(".img_instruct_attn.", ".img_instruct_attn.processor.")
|
||||
results.append((ckpt_name, value))
|
||||
return results
|
||||
|
||||
# Fused FFN gate/up splits into linear_1 (gate) / linear_3 (input).
|
||||
if ".gate_up_proj." in name:
|
||||
inner = param.shape[0] // 2
|
||||
return [
|
||||
(name.replace(".gate_up_proj.", ".linear_1."), param[:inner]),
|
||||
(name.replace(".gate_up_proj.", ".linear_3."), param[inner:]),
|
||||
]
|
||||
|
||||
# Default: 1:1 with the existing diffusers name promotions.
|
||||
if ".to_out." in name:
|
||||
name = name.replace(".to_out.", ".to_out.0.")
|
||||
for proj in (
|
||||
"img_to_q",
|
||||
"img_to_k",
|
||||
"img_to_v",
|
||||
"instruct_to_q",
|
||||
"instruct_to_k",
|
||||
"instruct_to_v",
|
||||
"instruct_out",
|
||||
"img_out",
|
||||
):
|
||||
for proj in ("instruct_out", "img_out"):
|
||||
token = f".img_instruct_attn.{proj}."
|
||||
if token in name:
|
||||
name = name.replace(token, f".img_instruct_attn.processor.{proj}.")
|
||||
break
|
||||
return name
|
||||
return [(name, param)]
|
||||
|
||||
|
||||
def test_transformer_exposes_stacked_params_mapping():
|
||||
"""The packed -> sub-layer mapping must be discoverable from the module tree.
|
||||
|
||||
``diffusion/lora/loader.py`` and the quantized weight loaders read
|
||||
``stacked_params_mapping`` off the model, so it has to exist before (and
|
||||
independently of) ``load_weights``.
|
||||
"""
|
||||
from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
|
||||
_BOOGU_STACKED_PARAMS_MAPPING,
|
||||
BooguImageTransformer2DModel,
|
||||
)
|
||||
|
||||
model = BooguImageTransformer2DModel(od_config=_tiny_od_config())
|
||||
|
||||
assert tuple(model.stacked_params_mapping) == _BOOGU_STACKED_PARAMS_MAPPING
|
||||
# A per-instance copy, so consumers cannot mutate the module constant.
|
||||
assert model.stacked_params_mapping is not _BOOGU_STACKED_PARAMS_MAPPING
|
||||
|
||||
# Every entry is a (param, shard, shard_id) triple.
|
||||
for param_name, shard_name, shard_id in model.stacked_params_mapping:
|
||||
assert param_name.startswith(".") and shard_name.startswith(".")
|
||||
assert shard_id in {"q", "k", "v", 0, 1}
|
||||
|
||||
# The mapping targets exactly the fused projections this port creates.
|
||||
mapped_leaves = {param.strip(".").split(".")[-1] for param, _, _ in model.stacked_params_mapping}
|
||||
assert mapped_leaves == {"to_qkv", "img_to_qkv", "instruct_to_qkv", "gate_up_proj"}
|
||||
|
||||
|
||||
def test_transformer_load_weights_round_trip():
|
||||
@@ -473,24 +530,59 @@ def test_transformer_load_weights_round_trip():
|
||||
model = BooguImageTransformer2DModel(od_config=_tiny_od_config())
|
||||
native_params = dict(model.named_parameters())
|
||||
|
||||
# Build synthetic diffusers-named weights (one per native parameter).
|
||||
checkpoint_weights = {}
|
||||
# Build synthetic diffusers-named weights. Fused native projections fan out
|
||||
# into the separate matrices the diffusers checkpoint stores.
|
||||
expected: dict[str, torch.Tensor] = {}
|
||||
checkpoint_weights: dict[str, torch.Tensor] = {}
|
||||
for native_name, param in native_params.items():
|
||||
checkpoint_weights[_native_to_checkpoint_name(native_name)] = torch.randn_like(param)
|
||||
full = torch.randn_like(param)
|
||||
expected[native_name] = full
|
||||
for ckpt_name, value in _native_to_checkpoint_weights(native_name, full):
|
||||
checkpoint_weights[ckpt_name] = value
|
||||
|
||||
# The remapping must be a bijection over the parameter set.
|
||||
assert len(checkpoint_weights) == len(native_params)
|
||||
# Every checkpoint name is unique; merged params fan out to >=1 weight.
|
||||
assert len(checkpoint_weights) >= len(native_params)
|
||||
|
||||
loaded = model.load_weights(list(checkpoint_weights.items()))
|
||||
|
||||
# No missing / unexpected parameters.
|
||||
assert loaded == set(native_params.keys())
|
||||
|
||||
# Values landed on the right parameters (TP=1: weight_loader copies verbatim).
|
||||
# Values landed on the right parameters (TP=1: weight_loader copies verbatim;
|
||||
# fused shards are placed at their q/k/v / gate/up offsets).
|
||||
reloaded = dict(model.named_parameters())
|
||||
for native_name in native_params:
|
||||
expected = checkpoint_weights[_native_to_checkpoint_name(native_name)]
|
||||
assert torch.allclose(reloaded[native_name], expected)
|
||||
assert torch.allclose(reloaded[native_name], expected[native_name])
|
||||
|
||||
|
||||
def test_transformer_load_weights_rejects_partial_fused_params():
|
||||
"""A fused parameter must not count as loaded until all its shards arrive.
|
||||
|
||||
The loader compares parameter *names*, so reporting ``to_qkv`` complete
|
||||
after a single ``to_q`` would let a checkpoint carrying only ``to_q`` start
|
||||
up with the ``k``/``v`` slices left uninitialized.
|
||||
"""
|
||||
from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
|
||||
BooguImageTransformer2DModel,
|
||||
)
|
||||
|
||||
model = BooguImageTransformer2DModel(od_config=_tiny_od_config())
|
||||
native_params = dict(model.named_parameters())
|
||||
|
||||
complete_name = "noise_refiner.0.attn.to_qkv.weight"
|
||||
partial_name = "noise_refiner.0.feed_forward.gate_up_proj.weight"
|
||||
assert complete_name in native_params
|
||||
assert partial_name in native_params
|
||||
|
||||
weights = list(_native_to_checkpoint_weights(complete_name, native_params[complete_name]))
|
||||
# Only the gate half of the FFN, so its fused parameter stays incomplete.
|
||||
gate_ckpt, gate_value = _native_to_checkpoint_weights(partial_name, native_params[partial_name])[0]
|
||||
weights.append((gate_ckpt, gate_value))
|
||||
|
||||
loaded = model.load_weights(weights)
|
||||
|
||||
assert complete_name in loaded
|
||||
assert partial_name not in loaded
|
||||
|
||||
|
||||
def test_transformer_load_weights_warns_for_unexpected_and_unloaded():
|
||||
@@ -513,12 +605,15 @@ def test_transformer_load_weights_warns_for_unexpected_and_unloaded():
|
||||
try:
|
||||
model = BooguImageTransformer2DModel(od_config=_tiny_od_config())
|
||||
native_params = dict(model.named_parameters())
|
||||
loaded_name = next(iter(native_params))
|
||||
checkpoint_name = _native_to_checkpoint_name(loaded_name)
|
||||
# A non-fused parameter, so one checkpoint entry loads it completely.
|
||||
loaded_name = "x_embedder.weight"
|
||||
assert loaded_name in native_params
|
||||
(checkpoint_name, checkpoint_value), *_ = _native_to_checkpoint_weights(loaded_name, native_params[loaded_name])
|
||||
checkpoint_value = torch.randn_like(checkpoint_value)
|
||||
|
||||
loaded = model.load_weights(
|
||||
[
|
||||
(checkpoint_name, torch.randn_like(native_params[loaded_name])),
|
||||
(checkpoint_name, checkpoint_value),
|
||||
("unexpected.weight", torch.ones(1)),
|
||||
]
|
||||
)
|
||||
|
||||
@@ -1,11 +1,12 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import sys
|
||||
import types
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
@@ -85,6 +86,7 @@ class StubScheduler:
|
||||
flow_shift: float = 1.0,
|
||||
) -> None:
|
||||
self.timesteps = torch.tensor(timesteps or [9, 3], dtype=torch.int64)
|
||||
self.sigmas = self.timesteps.float() / 10
|
||||
self.config = SimpleNamespace(
|
||||
num_train_timesteps=1000,
|
||||
flow_shift=flow_shift,
|
||||
@@ -113,6 +115,7 @@ class StubScheduler:
|
||||
else:
|
||||
assert num_inference_steps is not None
|
||||
self.timesteps = torch.arange(num_inference_steps, 0, -1, dtype=torch.int64, device=device)
|
||||
self.sigmas = self.timesteps.float()
|
||||
|
||||
def step(self, noise_pred: torch.Tensor, timestep: torch.Tensor, latents: torch.Tensor, **kwargs):
|
||||
del kwargs
|
||||
@@ -250,7 +253,7 @@ def passthrough_progress_bar(iterable):
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def fake_cosmos3_guardrails(monkeypatch: pytest.MonkeyPatch):
|
||||
module = types.ModuleType("vllm_omni.diffusion.models.cosmos3.guardrails")
|
||||
module: Any = types.ModuleType("vllm_omni.diffusion.models.cosmos3.guardrails")
|
||||
module.is_guardrails_enabled = lambda od_config, sampling_params=None: False
|
||||
module.ensure_initialized = lambda od_config: None
|
||||
module.check_text_safety = lambda text: None
|
||||
@@ -283,6 +286,8 @@ def make_cosmos3_pipeline():
|
||||
pipeline.is_edge_model = False
|
||||
pipeline._guidance_scale = None
|
||||
pipeline._num_timesteps = None
|
||||
pipeline._current_step_index = None
|
||||
pipeline._current_sigma = None
|
||||
pipeline._cosmos3_branch_caches = None
|
||||
pipeline._cache_dit_requires_paired_cfg = False
|
||||
pipeline._sound_tokenizer = None
|
||||
@@ -300,7 +305,7 @@ def sequential_cfg_parallel(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
|
||||
|
||||
def make_sampling_params(**overrides: Any) -> SimpleNamespace:
|
||||
values = {
|
||||
values: dict[str, Any] = {
|
||||
"height": None,
|
||||
"width": None,
|
||||
"num_frames": None,
|
||||
@@ -814,7 +819,7 @@ def test_pipeline_resolves_scheduler_class_from_checkpoint_file(
|
||||
t_list = [1.0, 0.75, 0.5, 0.25]
|
||||
scheduler_dir = tmp_path / "scheduler"
|
||||
scheduler_dir.mkdir()
|
||||
scheduler_config = {"_class_name": scheduler_class_name}
|
||||
scheduler_config: dict[str, Any] = {"_class_name": scheduler_class_name}
|
||||
if expected_distilled:
|
||||
scheduler_config["fixed_step_sampler_config"] = {"sample_type": "sde", "t_list": t_list}
|
||||
(scheduler_dir / "scheduler_config.json").write_text(json.dumps(scheduler_config))
|
||||
@@ -1531,7 +1536,7 @@ def test_ir_op_priority_hook_preserves_platform_fields(monkeypatch: pytest.Monke
|
||||
fused_add_rms_norm: list[str]
|
||||
custom_op: list[str]
|
||||
|
||||
fake_kernel = types.ModuleType("vllm.config.kernel")
|
||||
fake_kernel: Any = types.ModuleType("vllm.config.kernel")
|
||||
fake_kernel.IrOpPriorityConfig = FakeIrOpPriorityConfig
|
||||
monkeypatch.setitem(sys.modules, fake_kernel.__name__, fake_kernel)
|
||||
|
||||
@@ -2154,6 +2159,47 @@ def test_diffuse_covers_cfg_i2v_and_multimodal_steps(make_cosmos3_pipeline) -> N
|
||||
torch.testing.assert_close(action_result, torch.full((), 44.0).expand_as(action_result))
|
||||
|
||||
|
||||
def test_diffuse_publishes_exact_seacache_metadata_and_cfg_contexts(make_cosmos3_pipeline) -> None:
|
||||
pipeline = make_cosmos3_pipeline()
|
||||
pipeline.scheduler.sigmas = torch.tensor([0.91, 0.17])
|
||||
observations: list[tuple[str, int | None, float | None, int | None]] = []
|
||||
|
||||
class RecordingHook:
|
||||
@contextmanager
|
||||
def cache_context(self, name: str):
|
||||
sigma = pipeline.current_sigma
|
||||
observations.append(
|
||||
(
|
||||
name,
|
||||
pipeline.current_step_index,
|
||||
None if sigma is None else float(sigma),
|
||||
pipeline.num_timesteps,
|
||||
)
|
||||
)
|
||||
yield
|
||||
|
||||
pipeline._cache_context_factory = RecordingHook().cache_context
|
||||
pipeline.diffuse(
|
||||
latents=torch.zeros(1, 2, 1, 1, 1),
|
||||
timesteps=torch.tensor([900, 100]),
|
||||
cond_ids=_ids(2),
|
||||
cond_mask=_mask(),
|
||||
uncond_ids=_ids(1),
|
||||
uncond_mask=_mask(),
|
||||
guidance_scale=3.0,
|
||||
shared_kwargs={"video_shape": (1, 1, 1), "fps": 24.0},
|
||||
)
|
||||
|
||||
assert observations == [
|
||||
("cond", 0, pytest.approx(0.91), 2),
|
||||
("uncond", 0, pytest.approx(0.91), 2),
|
||||
("cond", 1, pytest.approx(0.17), 2),
|
||||
("uncond", 1, pytest.approx(0.17), 2),
|
||||
]
|
||||
assert pipeline.current_step_index is None
|
||||
assert pipeline.current_sigma is None
|
||||
|
||||
|
||||
def test_diffuse_drops_session_when_progress_iteration_fails(make_cosmos3_pipeline) -> None:
|
||||
pipeline = make_cosmos3_pipeline()
|
||||
pipeline._use_session_state = True
|
||||
@@ -2213,6 +2259,53 @@ def test_diffuse_transfer_applies_control_cfg(make_cosmos3_pipeline, sequential_
|
||||
torch.testing.assert_close(result, torch.full_like(latents, 254.0))
|
||||
|
||||
|
||||
def test_diffuse_transfer_uses_named_seacache_contexts(make_cosmos3_pipeline, sequential_cfg_parallel) -> None:
|
||||
pipeline = make_cosmos3_pipeline()
|
||||
pipeline.scheduler.sigmas = torch.tensor([0.42])
|
||||
contexts: list[tuple[str, int | None, float | None, int | None]] = []
|
||||
|
||||
class RecordingHook:
|
||||
@contextmanager
|
||||
def cache_context(self, name: str):
|
||||
sigma = pipeline.current_sigma
|
||||
contexts.append(
|
||||
(
|
||||
name,
|
||||
pipeline.current_step_index,
|
||||
None if sigma is None else float(sigma),
|
||||
pipeline.num_timesteps,
|
||||
)
|
||||
)
|
||||
yield
|
||||
|
||||
pipeline._cache_context_factory = RecordingHook().cache_context
|
||||
latents = torch.zeros(1, 2, 1, 1, 1)
|
||||
velocity_mask = torch.ones(1, 1, 1, 1, 1)
|
||||
pipeline.diffuse_transfer(
|
||||
latents=latents,
|
||||
timesteps=torch.tensor([7]),
|
||||
cond_ids=_ids(2),
|
||||
cond_mask=_mask(),
|
||||
uncond_ids=_ids(1),
|
||||
uncond_mask=_mask(),
|
||||
guidance_scale=3.0,
|
||||
control_guidance=1.5,
|
||||
control_guidance_interval=None,
|
||||
control_latents=[torch.zeros_like(latents)],
|
||||
shared_kwargs={"video_shape": (1, 1, 1), "fps": 24.0, "noisy_frame_mask": velocity_mask},
|
||||
velocity_mask=velocity_mask,
|
||||
condition_latents=torch.zeros_like(latents),
|
||||
)
|
||||
|
||||
assert contexts == [
|
||||
("cond", 0, pytest.approx(0.42), 1),
|
||||
("cond_no_control", 0, pytest.approx(0.42), 1),
|
||||
("uncond", 0, pytest.approx(0.42), 1),
|
||||
]
|
||||
assert pipeline.current_step_index is None
|
||||
assert pipeline.current_sigma is None
|
||||
|
||||
|
||||
def test_diffuse_transfer_rejects_session_state_manager(make_cosmos3_pipeline) -> None:
|
||||
pipeline = make_cosmos3_pipeline()
|
||||
pipeline._use_session_state = True
|
||||
@@ -2583,7 +2676,7 @@ def test_diffuse_keeps_paired_cfg_when_cache_dit_active(make_cosmos3_pipeline) -
|
||||
|
||||
class TestForwardRouting:
|
||||
def _install_forward_stubs(self, pipeline):
|
||||
captured: dict[str, object] = {"diffuse_calls": [], "prepare_calls": []}
|
||||
captured: dict[str, Any] = {"diffuse_calls": [], "prepare_calls": []}
|
||||
|
||||
def fake_format(
|
||||
prompt,
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
@@ -513,6 +513,108 @@ def test_forward_returns_video_prediction(monkeypatch: pytest.MonkeyPatch) -> No
|
||||
assert tuple(output.shape) == (1, 2, 1, 2, 2)
|
||||
|
||||
|
||||
def test_cache_execution_residual_spans_final_gen_norm(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
from vllm_omni.diffusion.cache.teacache.extractors import extract_cosmos3_context
|
||||
from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
|
||||
|
||||
class TrackingNorm(nn.Module):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.calls = 0
|
||||
|
||||
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
||||
self.calls += 1
|
||||
return hidden_states + 5.0
|
||||
|
||||
monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None))
|
||||
model = transformer_cosmos3.Cosmos3VFMTransformer(
|
||||
SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32)
|
||||
)
|
||||
norm = TrackingNorm()
|
||||
model.norm_moe_gen = norm
|
||||
captured: dict[str, torch.Tensor] = {}
|
||||
|
||||
def run_gen_layers(hidden_gen: torch.Tensor, **kwargs) -> torch.Tensor:
|
||||
del kwargs
|
||||
captured["input"] = hidden_gen.detach().clone()
|
||||
return hidden_gen + 2.0
|
||||
|
||||
monkeypatch.setattr(model, "_run_gen_layers", run_gen_layers)
|
||||
forward_kwargs = {
|
||||
"hidden_states": torch.zeros(1, 2, 1, 2, 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, 2, 2),
|
||||
"fps": 24.0,
|
||||
}
|
||||
|
||||
full_output = model(**forward_kwargs)
|
||||
|
||||
assert norm.calls == 1
|
||||
norm.calls = 0
|
||||
|
||||
ctx = extract_cosmos3_context(model, **forward_kwargs)
|
||||
execution_input = ctx.hidden_states.detach().clone()
|
||||
execution_output = ctx.run_transformer_blocks()[0]
|
||||
residual = execution_output - execution_input
|
||||
extracted_output = ctx.postprocess(execution_output)
|
||||
|
||||
assert norm.calls == 1
|
||||
expected_residual = ((captured["input"] + 2.0) + 5.0) - captured["input"]
|
||||
assert torch.equal(residual, expected_residual)
|
||||
assert torch.equal(extracted_output, full_output)
|
||||
|
||||
norm.calls = 0
|
||||
|
||||
def fail_if_gen_layers_run(*args, **kwargs):
|
||||
del args, kwargs
|
||||
raise AssertionError("SeaCache hit unexpectedly executed GEN layers")
|
||||
|
||||
monkeypatch.setattr(model, "_run_gen_layers", fail_if_gen_layers_run)
|
||||
|
||||
cached_ctx = extract_cosmos3_context(model, **forward_kwargs)
|
||||
cached_output = cached_ctx.postprocess(cached_ctx.hidden_states + residual)
|
||||
|
||||
assert norm.calls == 0
|
||||
assert torch.equal(cached_output, full_output)
|
||||
for name in ("_seacache_skip", "_seacache_record", "_seacache_residual", "_seacache_last_residual"):
|
||||
assert not hasattr(model, name)
|
||||
|
||||
|
||||
def test_no_cache_still_runs_final_gen_norm_once(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
|
||||
|
||||
class CountingNorm(nn.Module):
|
||||
def __init__(self, original: nn.Module) -> None:
|
||||
super().__init__()
|
||||
self.original = original
|
||||
self.calls = 0
|
||||
|
||||
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
||||
self.calls += 1
|
||||
return self.original(hidden_states)
|
||||
|
||||
monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None))
|
||||
model = transformer_cosmos3.Cosmos3VFMTransformer(
|
||||
SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32)
|
||||
)
|
||||
norm = CountingNorm(model.norm_moe_gen)
|
||||
model.norm_moe_gen = norm
|
||||
|
||||
output = model(
|
||||
hidden_states=torch.zeros(1, 2, 1, 2, 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, 2, 2),
|
||||
fps=24.0,
|
||||
)
|
||||
|
||||
assert norm.calls == 1
|
||||
assert tuple(output.shape) == (1, 2, 1, 2, 2)
|
||||
|
||||
|
||||
def test_qwen_gen_mlp_remains_gated() -> None:
|
||||
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3GatedMLP, Cosmos3VFMTransformer
|
||||
|
||||
|
||||
@@ -2711,3 +2711,125 @@ def test_registry_and_model_exports_resolve_official_pipeline_class_name() -> No
|
||||
assert "from .pipeline import" in lingbot_init
|
||||
assert '"LingBotWorldCausalDMDPipeline"' in lingbot_init
|
||||
assert '"CausalLingBotWorldTransformer3DModel"' in lingbot_init
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# a prompt is encoded once, whichever path asks for it
|
||||
# ---------------------------------------------------------------------------
|
||||
def _counting_text_encoder(pipeline) -> list[str]:
|
||||
"""Give ``pipeline`` its real encode_prompt over stubs that record each encode."""
|
||||
del pipeline.encode_prompt # _pipeline() replaces it; these tests need the real one
|
||||
encoded: list[str] = []
|
||||
|
||||
def tokenizer(texts, **kwargs):
|
||||
del kwargs
|
||||
encoded.extend(texts)
|
||||
return SimpleNamespace(
|
||||
input_ids=torch.arange(512).view(1, 512),
|
||||
attention_mask=torch.ones(1, 512, dtype=torch.long),
|
||||
)
|
||||
|
||||
def text_encoder(input_ids, attention_mask):
|
||||
del input_ids, attention_mask
|
||||
return SimpleNamespace(last_hidden_state=torch.full((1, 512, 8), float(len(encoded))))
|
||||
|
||||
pipeline.tokenizer = tokenizer
|
||||
pipeline.text_encoder = text_encoder
|
||||
return encoded
|
||||
|
||||
|
||||
def test_the_text_encoder_runs_once_per_distinct_prompt() -> None:
|
||||
module = _load_pipeline_module()
|
||||
pipeline = _pipeline(module)
|
||||
encoded = _counting_text_encoder(pipeline)
|
||||
|
||||
def encode(prompt, *, length=512, dtype=torch.float32):
|
||||
return pipeline.encode_prompt(prompt, max_sequence_length=length, dtype=dtype)
|
||||
|
||||
first = encode("move through the room")
|
||||
assert encode("move through the room") is first
|
||||
# The key is the text the tokenizer actually sees, so spacing alone is not new.
|
||||
assert encode(" move through the room ") is first
|
||||
assert encoded == ["move through the room"]
|
||||
|
||||
encode("enter the snowy valley")
|
||||
# dtype and sequence length change the tensor, so each is its own entry.
|
||||
encode("move through the room", dtype=torch.float16)
|
||||
encode("move through the room", length=256)
|
||||
assert encoded == [
|
||||
"move through the room",
|
||||
"enter the snowy valley",
|
||||
"move through the room",
|
||||
"move through the room",
|
||||
]
|
||||
|
||||
|
||||
def test_the_prompt_encode_cache_is_bounded() -> None:
|
||||
module = _load_pipeline_module()
|
||||
pipeline = _pipeline(module)
|
||||
encoded = _counting_text_encoder(pipeline)
|
||||
size = module._PROMPT_EMBEDS_CACHE_SIZE
|
||||
|
||||
for index in range(size + 1):
|
||||
pipeline.encode_prompt(f"prompt {index}", max_sequence_length=512, dtype=torch.float32)
|
||||
assert len(pipeline._prompt_embeds_cache) == size
|
||||
|
||||
# The oldest entry was the one evicted, so it is the one encoded again.
|
||||
pipeline.encode_prompt("prompt 0", max_sequence_length=512, dtype=torch.float32)
|
||||
pipeline.encode_prompt(f"prompt {size}", max_sequence_length=512, dtype=torch.float32)
|
||||
assert encoded == [f"prompt {index}" for index in range(size + 1)] + ["prompt 0"]
|
||||
|
||||
|
||||
def test_realtime_ticks_reuse_the_encode_and_still_invalidate_on_a_prompt_change(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
module = _load_pipeline_module()
|
||||
pipeline = _pipeline(module)
|
||||
pipeline._ar_height = 16
|
||||
pipeline._ar_width = 16
|
||||
pipeline._ar_diffusion_kv_state = object()
|
||||
encoded = _counting_text_encoder(pipeline)
|
||||
invalidations: list[bool] = []
|
||||
|
||||
def ar_text_caches(prompt_embeds, *, invalidate):
|
||||
del prompt_embeds
|
||||
invalidations.append(invalidate)
|
||||
return [SimpleNamespace()]
|
||||
|
||||
monkeypatch.setattr(pipeline, "_ar_text_caches", ar_text_caches)
|
||||
monkeypatch.setattr(
|
||||
pipeline,
|
||||
"_generate_block",
|
||||
lambda **kwargs: torch.randn((1, 16, 3, 2, 2), generator=kwargs["generator"]),
|
||||
)
|
||||
|
||||
for chunk_index in range(3):
|
||||
pipeline(_request(sampling=_SamplingParams(extra_args=_tick_extra_args(chunk_index=chunk_index))))
|
||||
switched = _prompt()
|
||||
switched["prompt"] = "enter the snowy valley"
|
||||
pipeline(
|
||||
_request(
|
||||
sampling=_SamplingParams(extra_args=_tick_extra_args(chunk_index=3, prompt=switched["prompt"])),
|
||||
prompt=switched,
|
||||
)
|
||||
)
|
||||
|
||||
assert encoded == ["move through the room", "enter the snowy valley"]
|
||||
# A new prompt must still drop the cross-attention K/V built from the old one.
|
||||
assert invalidations == [False, False, False, True]
|
||||
|
||||
|
||||
def test_stepwise_requests_with_the_same_prompt_share_one_encode() -> None:
|
||||
module = _load_pipeline_module()
|
||||
pipeline = _pipeline(module, transformer=_RecordingTransformer())
|
||||
pipeline._ar_height = 16
|
||||
pipeline._ar_width = 16
|
||||
encoded = _counting_text_encoder(pipeline)
|
||||
|
||||
states = [_stepwise_state(request_id=request_id, num_frames=21) for request_id in ("req-a", "req-b")]
|
||||
for state in states:
|
||||
with pipeline.bind_ar_diffusion_state(state.request_id, _FakeARState(state.request_id)):
|
||||
pipeline.prepare_encode(state)
|
||||
|
||||
assert encoded == ["move through the room"]
|
||||
assert states[0].prompt_embeds is states[1].prompt_embeds
|
||||
@@ -0,0 +1,137 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.diffusion.models.magi2.test_native_packing import (
|
||||
_longer_text_tensors,
|
||||
_sampler_tensors,
|
||||
_tiny_model,
|
||||
_tiny_sampler,
|
||||
)
|
||||
from vllm_omni.diffusion.models.magi2.sampler_magi2 import CFGConfig
|
||||
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
|
||||
|
||||
|
||||
def test_layer_regions_bracket_the_eager_kernels() -> None:
|
||||
moe_layer, dense_layer = _tiny_model().block.layers
|
||||
|
||||
assert moe_layer.region_methods == ("_attention_input", "_moe_input", "_moe_output")
|
||||
assert dense_layer.region_methods == ("_attention_input", "_dense_output")
|
||||
|
||||
|
||||
def test_compile_regions_fullgraph_matches_eager_bitwise() -> None:
|
||||
model = _tiny_model()
|
||||
sampler = _tiny_sampler(model)
|
||||
steps = [
|
||||
sampler.prepare_model_input(**_sampler_tensors(seed), t=t, cfg_config=CFGConfig())
|
||||
for seed, t in enumerate((torch.tensor([900.0]), torch.tensor([450.0])))
|
||||
]
|
||||
expected = [sampler.forward(step) for step in steps]
|
||||
|
||||
torch._dynamo.reset()
|
||||
graphs_before = torch._dynamo.utils.counters["stats"]["unique_graphs"]
|
||||
model.compile_regions(fullgraph=True, backend="eager", dynamic=True)
|
||||
first = sampler.forward(steps[0])
|
||||
graphs = torch._dynamo.utils.counters["stats"]["unique_graphs"] - graphs_before
|
||||
with torch._dynamo.config.patch(error_on_recompile=True):
|
||||
second = sampler.forward(steps[1])
|
||||
|
||||
for (video, audio), (video_ref, audio_ref) in zip((first, second), expected, strict=True):
|
||||
assert torch.equal(video, video_ref)
|
||||
assert torch.equal(audio, audio_ref)
|
||||
# pre adapter, post adapter, three MoE-layer regions, two dense-layer regions
|
||||
assert graphs == 7
|
||||
|
||||
|
||||
def test_pipeline_setup_compile_forwards_config_dynamic() -> None:
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.models.magi2.pipeline_magi2 import Magi2Pipeline
|
||||
|
||||
calls: list[dict[str, object]] = []
|
||||
|
||||
class _Transformer(torch.nn.Module):
|
||||
def compile_regions(self, **compile_kwargs: object) -> None:
|
||||
calls.append(compile_kwargs)
|
||||
|
||||
pipeline = object.__new__(Magi2Pipeline)
|
||||
torch.nn.Module.__init__(pipeline)
|
||||
pipeline.od_config = OmniDiffusionConfig(
|
||||
model="sand-ai/MAGI-2-preview",
|
||||
model_class_name="Magi2Pipeline",
|
||||
diffusion_compile_dynamic=False,
|
||||
)
|
||||
pipeline.transformer = _Transformer()
|
||||
|
||||
pipeline.setup_compile()
|
||||
|
||||
assert calls == [{"fullgraph": True, "dynamic": False, "options": {"emulate_precision_casts": True}}]
|
||||
|
||||
|
||||
def test_pipeline_setup_compile_warns_when_granularity_is_not_regional(caplog: pytest.LogCaptureFixture) -> None:
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.models.magi2.pipeline_magi2 import Magi2Pipeline
|
||||
|
||||
calls: list[dict[str, object]] = []
|
||||
|
||||
class _Transformer(torch.nn.Module):
|
||||
def compile_regions(self, **compile_kwargs: object) -> None:
|
||||
calls.append(compile_kwargs)
|
||||
|
||||
pipeline = object.__new__(Magi2Pipeline)
|
||||
torch.nn.Module.__init__(pipeline)
|
||||
pipeline.od_config = OmniDiffusionConfig(
|
||||
model="sand-ai/MAGI-2-preview",
|
||||
model_class_name="Magi2Pipeline",
|
||||
diffusion_compile_granularity="full",
|
||||
)
|
||||
pipeline.transformer = _Transformer()
|
||||
|
||||
pipeline.setup_compile()
|
||||
|
||||
assert calls == [{"fullgraph": True, "dynamic": True, "options": {"emulate_precision_casts": True}}]
|
||||
assert "diffusion_compile_granularity='full'" in caplog.text
|
||||
|
||||
|
||||
def test_compiled_graphs_are_shared_across_layers_of_one_kind() -> None:
|
||||
model = _tiny_model(moe_layers=2)
|
||||
sampler = _tiny_sampler(model)
|
||||
step = sampler.prepare_model_input(**_sampler_tensors(0), t=torch.tensor([900.0]), cfg_config=CFGConfig())
|
||||
expected = sampler.forward(step)
|
||||
|
||||
torch._dynamo.reset()
|
||||
graphs_before = torch._dynamo.utils.counters["stats"]["unique_graphs"]
|
||||
model.compile_regions(fullgraph=True, backend="eager", dynamic=True)
|
||||
actual = sampler.forward(step)
|
||||
|
||||
assert torch.equal(actual[0], expected[0])
|
||||
# Two MoE layers and two dense layers still compile one graph per region
|
||||
# kind: pre adapter, post adapter, three MoE regions, two dense regions.
|
||||
assert torch._dynamo.utils.counters["stats"]["unique_graphs"] - graphs_before == 7
|
||||
|
||||
|
||||
def test_compiled_regions_survive_new_requests_without_recompiling() -> None:
|
||||
model = _tiny_model()
|
||||
sampler = _tiny_sampler(model)
|
||||
t = torch.tensor([900.0])
|
||||
first_request = sampler.prepare_model_input(**_sampler_tensors(0), t=t, cfg_config=CFGConfig())
|
||||
second_request = sampler.prepare_model_input(**_sampler_tensors(1), t=t, cfg_config=CFGConfig())
|
||||
longer_request = sampler.prepare_model_input(**_longer_text_tensors(2), t=t, cfg_config=CFGConfig())
|
||||
expected = [sampler.forward(request) for request in (first_request, second_request, longer_request)]
|
||||
|
||||
torch._dynamo.reset()
|
||||
graphs_before = torch._dynamo.utils.counters["stats"]["unique_graphs"]
|
||||
model.compile_regions(fullgraph=True, backend="eager", dynamic=True)
|
||||
actual = [sampler.forward(first_request)]
|
||||
with torch._dynamo.config.patch(error_on_recompile=True):
|
||||
actual.append(sampler.forward(second_request))
|
||||
actual.append(sampler.forward(longer_request))
|
||||
|
||||
for (video, audio), (video_ref, audio_ref) in zip(actual, expected, strict=True):
|
||||
assert torch.equal(video, video_ref)
|
||||
assert torch.equal(audio, audio_ref)
|
||||
assert torch._dynamo.utils.counters["stats"]["unique_graphs"] - graphs_before == 7
|
||||
@@ -0,0 +1,77 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.diffusion.models.magi2.test_native_packing import (
|
||||
_longer_text_tensors,
|
||||
_sampler_tensors,
|
||||
_tiny_model,
|
||||
_tiny_sampler,
|
||||
)
|
||||
from tests.helpers.mark import hardware_marks
|
||||
from vllm_omni.diffusion.models.magi2.sampler_magi2 import CFGConfig
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.diffusion,
|
||||
pytest.mark.core_model,
|
||||
*hardware_marks(res={"cuda": "L4"}, num_cards=1),
|
||||
pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA"),
|
||||
]
|
||||
|
||||
|
||||
def _cuda_steps(sampler, tensors=_sampler_tensors):
|
||||
timesteps = (torch.tensor([900.0]), torch.tensor([600.0]), torch.tensor([300.0]))
|
||||
return [
|
||||
sampler.prepare_model_input(
|
||||
**{name: value.cuda() for name, value in tensors(seed).items()},
|
||||
t=t,
|
||||
cfg_config=CFGConfig(),
|
||||
)
|
||||
for seed, t in enumerate(timesteps)
|
||||
]
|
||||
|
||||
|
||||
def test_inductor_regions_match_eager_in_deterministic_mode(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("MAGI2_DETERMINISTIC", "1")
|
||||
model = _tiny_model(params_dtype=torch.bfloat16).cuda()
|
||||
sampler = _tiny_sampler(model)
|
||||
steps = _cuda_steps(sampler)
|
||||
expected = [sampler.forward(step) for step in steps]
|
||||
new_request_step = _cuda_steps(sampler)[0]
|
||||
longer_text_step = _cuda_steps(sampler, _longer_text_tensors)[0]
|
||||
expected_new_request = sampler.forward(new_request_step)
|
||||
expected_longer_text = sampler.forward(longer_text_step)
|
||||
|
||||
torch._dynamo.reset()
|
||||
graphs_before = torch._dynamo.utils.counters["stats"]["unique_graphs"]
|
||||
# The pipeline compiles with inductor emulating eager's intermediate
|
||||
# low-precision casts, so fused bf16 chains keep eager's rounding.
|
||||
model.compile_regions(fullgraph=True, dynamic=True, options={"emulate_precision_casts": True})
|
||||
first = sampler.forward(steps[0])
|
||||
with torch._dynamo.config.patch(error_on_recompile=True):
|
||||
compiled = [first, sampler.forward(steps[1]), sampler.forward(steps[2])]
|
||||
repeated = [sampler.forward(step) for step in steps]
|
||||
# A second request repacks the same shapes; a third changes the
|
||||
# text length and therefore every token count.
|
||||
compiled_new_request = sampler.forward(new_request_step)
|
||||
compiled_longer_text = sampler.forward(longer_text_step)
|
||||
|
||||
# Emulated precision casts restore eager's rounding boundaries; what is
|
||||
# left is one-ulp transcendental noise that some Triton backends carry.
|
||||
for (video, audio), (video_ref, audio_ref) in zip(compiled, expected, strict=True):
|
||||
torch.testing.assert_close(video, video_ref, atol=1e-6, rtol=0)
|
||||
torch.testing.assert_close(audio, audio_ref, atol=1e-6, rtol=0)
|
||||
for (video, audio), (video_again, audio_again) in zip(compiled, repeated, strict=True):
|
||||
assert torch.equal(video, video_again)
|
||||
assert torch.equal(audio, audio_again)
|
||||
for (video, audio), (video_ref, audio_ref) in (
|
||||
(compiled_new_request, expected_new_request),
|
||||
(compiled_longer_text, expected_longer_text),
|
||||
):
|
||||
torch.testing.assert_close(video, video_ref, atol=1e-6, rtol=0)
|
||||
torch.testing.assert_close(audio, audio_ref, atol=1e-6, rtol=0)
|
||||
assert torch._dynamo.utils.counters["stats"]["unique_graphs"] - graphs_before == 7
|
||||
@@ -0,0 +1,179 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
from contextlib import ExitStack, contextmanager
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
import torch.multiprocessing as mp
|
||||
import torch.nn as nn
|
||||
from torch.distributed import init_device_mesh
|
||||
from torch.distributed.fsdp import MixedPrecisionPolicy
|
||||
|
||||
import vllm_omni.diffusion.offloader.distributed_layerwise_backend as dlo_module
|
||||
from tests.diffusion.models.magi2.test_native_distributed_parity import (
|
||||
_current_group,
|
||||
_new_groups,
|
||||
_patched_groups,
|
||||
)
|
||||
from tests.diffusion.models.magi2.test_native_packing import _tiny_config, _tiny_model, _tiny_sampler
|
||||
from tests.diffusion.offloader.helpers import patch_offload_runtime
|
||||
from vllm_omni.diffusion.distributed.hsdp import shard_model
|
||||
from vllm_omni.diffusion.models.magi2.modeling_magi2 import Magi2PreviewTransformer
|
||||
from vllm_omni.diffusion.models.magi2.parallel import Magi2ParallelGroup
|
||||
from vllm_omni.diffusion.models.magi2.pipeline_magi2 import Magi2Pipeline
|
||||
from vllm_omni.diffusion.models.magi2.sampler_magi2 import CFGConfig
|
||||
from vllm_omni.diffusion.offloader.base import OffloadConfig
|
||||
from vllm_omni.diffusion.offloader.config import OffloadStrategy
|
||||
|
||||
_WORLD_SIZE = 4
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _cpu_offload_runtime():
|
||||
with pytest.MonkeyPatch.context() as monkeypatch:
|
||||
patch_offload_runtime(monkeypatch, dlo_module.current_omni_platform, synchronize=True)
|
||||
yield
|
||||
|
||||
|
||||
def _enable_dlo(
|
||||
model: Magi2PreviewTransformer,
|
||||
*,
|
||||
dp_group: dist.ProcessGroup | None,
|
||||
dp_size: int,
|
||||
dp_rank: int,
|
||||
) -> None:
|
||||
pipeline = Magi2Pipeline.__new__(Magi2Pipeline)
|
||||
nn.Module.__init__(pipeline)
|
||||
pipeline.transformer = model
|
||||
backend = dlo_module.DistributedLayerwiseOffloadBackend(
|
||||
OffloadConfig(
|
||||
strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
|
||||
pin_cpu_memory=False,
|
||||
dp_size=dp_size,
|
||||
dlo_use_allgather=dp_size > 1,
|
||||
),
|
||||
torch.device("cpu"),
|
||||
)
|
||||
backend.dp_group = dp_group
|
||||
backend.rank = dp_rank
|
||||
backend.enable(pipeline)
|
||||
assert backend.enabled
|
||||
|
||||
|
||||
def _enable_hsdp(model: Magi2PreviewTransformer) -> None:
|
||||
mesh = init_device_mesh("cpu", mesh_shape=(1, _WORLD_SIZE), mesh_dim_names=("replicate", "shard"))
|
||||
ignored = {
|
||||
parameter for name in model._hsdp_ignored_modules for parameter in model.get_submodule(name).parameters()
|
||||
}
|
||||
shard_model(
|
||||
model,
|
||||
mesh=mesh,
|
||||
mp_policy=MixedPrecisionPolicy(cast_forward_inputs=False),
|
||||
hsdp_shard_conditions=model._hsdp_shard_conditions,
|
||||
ignored_params=ignored,
|
||||
)
|
||||
model.requires_grad_(False)
|
||||
|
||||
|
||||
def _sampler_tensors(seed: int) -> dict[str, torch.Tensor]:
|
||||
# 14 tokens per CFG branch, so the packed sequence splits evenly across
|
||||
# two Ulysses ranks; gloo cannot gather uneven shards.
|
||||
generator = torch.Generator(device="cpu").manual_seed(seed)
|
||||
return {
|
||||
"latent": torch.randn(1, 4, 2, 1, 3, generator=generator),
|
||||
"audio_latent": torch.randn(1, 5, 4, generator=generator),
|
||||
"txt_feat": torch.randn(1, 3, 4, generator=generator),
|
||||
"null_txt_feat": torch.randn(1, 3, 4, generator=generator),
|
||||
}
|
||||
|
||||
|
||||
def _steps(sampler):
|
||||
return [
|
||||
sampler.prepare_model_input(**_sampler_tensors(seed), t=t, cfg_config=CFGConfig())
|
||||
for seed, t in enumerate((torch.tensor([900.0]), torch.tensor([450.0])))
|
||||
]
|
||||
|
||||
|
||||
def _run(model: Magi2PreviewTransformer, *, compile_regions: bool) -> list[tuple[torch.Tensor, torch.Tensor]]:
|
||||
sampler = _tiny_sampler(model)
|
||||
steps = _steps(sampler)
|
||||
if not compile_regions:
|
||||
return [sampler.forward(step) for step in steps]
|
||||
torch._dynamo.reset()
|
||||
model.compile_regions(fullgraph=True, backend="eager", dynamic=True)
|
||||
outputs = [sampler.forward(steps[0])]
|
||||
with torch._dynamo.config.patch(error_on_recompile=True):
|
||||
outputs.append(sampler.forward(steps[1]))
|
||||
return outputs
|
||||
|
||||
|
||||
def _assert_same(actual, expected, *, exact: bool, label: str) -> None:
|
||||
for (video, audio), (video_ref, audio_ref) in zip(actual, expected, strict=True):
|
||||
if exact:
|
||||
assert torch.equal(video, video_ref), label
|
||||
assert torch.equal(audio, audio_ref), label
|
||||
else:
|
||||
torch.testing.assert_close(video, video_ref, atol=2e-5, rtol=2e-4, msg=label)
|
||||
torch.testing.assert_close(audio, audio_ref, atol=2e-5, rtol=2e-4, msg=label)
|
||||
|
||||
|
||||
def _distributed_worker(rank: int, rendezvous: str) -> None:
|
||||
torch.set_num_threads(1)
|
||||
dist.init_process_group("gloo", init_method=rendezvous, rank=rank, world_size=_WORLD_SIZE)
|
||||
try:
|
||||
singleton = Magi2ParallelGroup(None, world_size=1, rank=0)
|
||||
with _patched_groups(singleton, singleton):
|
||||
oracle = _tiny_model()
|
||||
checkpoint = [(name, value.detach().clone()) for name, value in oracle.state_dict().items()]
|
||||
expected = _run(oracle, compile_regions=False)
|
||||
|
||||
sp2_groups = _new_groups(((0, 1), (2, 3)))
|
||||
dp2_groups = _new_groups(((0, 2), (1, 3)))
|
||||
sp_group = _current_group(rank, sp2_groups)
|
||||
dp_ranks, dp_process_group = next((ranks, group) for ranks, group in dp2_groups if rank in ranks)
|
||||
|
||||
def dlo_rank_local(model: Magi2PreviewTransformer) -> None:
|
||||
_enable_dlo(model, dp_group=None, dp_size=1, dp_rank=0)
|
||||
|
||||
def dlo_allgather(model: Magi2PreviewTransformer) -> None:
|
||||
_enable_dlo(model, dp_group=dp_process_group, dp_size=len(dp_ranks), dp_rank=dp_ranks.index(rank))
|
||||
|
||||
variants = (
|
||||
("dlo_rank_local_sp2", sp_group, dlo_rank_local),
|
||||
("dlo_allgather_dp2sp2", sp_group, dlo_allgather),
|
||||
("hsdp4", singleton, _enable_hsdp),
|
||||
)
|
||||
for name, ulysses_group, enable in variants:
|
||||
with ExitStack() as stack:
|
||||
stack.enter_context(_patched_groups(singleton, ulysses_group))
|
||||
stack.enter_context(_cpu_offload_runtime())
|
||||
outputs = {}
|
||||
for compile_regions in (False, True):
|
||||
model = Magi2PreviewTransformer(_tiny_config())
|
||||
assert model.load_weights(checkpoint) == set(model.state_dict())
|
||||
enable(model)
|
||||
outputs[compile_regions] = _run(model, compile_regions=compile_regions)
|
||||
_assert_same(outputs[False], expected, exact=False, label=f"{name} eager vs single-rank oracle")
|
||||
_assert_same(outputs[True], outputs[False], exact=True, label=f"{name} compiled vs eager")
|
||||
dist.barrier()
|
||||
finally:
|
||||
dist.destroy_process_group()
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not dist.is_available() or not dist.is_gloo_available(),
|
||||
reason="requires torch.distributed gloo",
|
||||
)
|
||||
def test_compiled_regions_match_eager_under_dlo_and_hsdp() -> None:
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
rendezvous = f"file://{os.path.join(temp_dir, 'gloo-rendezvous')}"
|
||||
with patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": ""}):
|
||||
mp.spawn(_distributed_worker, args=(rendezvous,), nprocs=_WORLD_SIZE, join=True)
|
||||
@@ -0,0 +1,118 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.models.magi2.configuration_magi2 import (
|
||||
Magi2MHCConfig,
|
||||
Magi2MoEConfig,
|
||||
Magi2PreviewConfig,
|
||||
)
|
||||
from vllm_omni.diffusion.models.magi2.layers import MultiModalityRMSNorm
|
||||
from vllm_omni.diffusion.models.magi2.mh_moe import Magi2MultiHeadMoE
|
||||
from vllm_omni.diffusion.models.magi2.modeling_magi2 import Magi2PreviewTransformer
|
||||
from vllm_omni.diffusion.models.magi2.preview_data_proxy import (
|
||||
Magi2DataProxy,
|
||||
Magi2PreviewDataProxyConfig,
|
||||
)
|
||||
from vllm_omni.diffusion.models.magi2.sampler_magi2 import CFGConfig, Magi2PreviewSampler
|
||||
|
||||
pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
|
||||
|
||||
|
||||
def _tiny_config(params_dtype: torch.dtype = torch.float32, *, moe_layers: int = 1) -> Magi2PreviewConfig:
|
||||
# The leading layers are multimodal with MoE, the rest single-modality
|
||||
# dense, so every layer belongs to one of two kinds.
|
||||
moe_indices = tuple(range(moe_layers))
|
||||
return Magi2PreviewConfig(
|
||||
num_layers=2 * moe_layers,
|
||||
hidden_size=16,
|
||||
head_dim=8,
|
||||
num_query_groups=2,
|
||||
video_in_channels=4,
|
||||
audio_in_channels=4,
|
||||
text_in_channels=4,
|
||||
intermediate_factor=2,
|
||||
multimodal_layers=moe_indices,
|
||||
params_dtype=params_dtype,
|
||||
mhc=Magi2MHCConfig(num_streams=2),
|
||||
moe=Magi2MoEConfig(
|
||||
num_heads=2,
|
||||
num_experts=4,
|
||||
top_k=2,
|
||||
expert_intermediate_size=8,
|
||||
shared_expert_intermediate_size=8,
|
||||
modality_shared_expert_intermediate_size=8,
|
||||
layers=moe_indices,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _tiny_model(
|
||||
seed: int = 11,
|
||||
params_dtype: torch.dtype = torch.float32,
|
||||
*,
|
||||
moe_layers: int = 1,
|
||||
) -> Magi2PreviewTransformer:
|
||||
model = Magi2PreviewTransformer(_tiny_config(params_dtype, moe_layers=moe_layers))
|
||||
generator = torch.Generator(device="cpu").manual_seed(seed)
|
||||
with torch.no_grad():
|
||||
for parameter in model.parameters():
|
||||
parameter.copy_(torch.randn(parameter.shape, generator=generator, dtype=parameter.dtype) * 0.02)
|
||||
for module in model.modules():
|
||||
if isinstance(module, MultiModalityRMSNorm):
|
||||
module.weight.zero_()
|
||||
elif isinstance(module, Magi2MultiHeadMoE):
|
||||
module.router.expert_bias.zero_()
|
||||
module.router.expert_bias_ema.zero_()
|
||||
return model
|
||||
|
||||
|
||||
def _tiny_sampler(model: torch.nn.Module) -> Magi2PreviewSampler:
|
||||
return Magi2PreviewSampler(model, Magi2DataProxy(Magi2PreviewDataProxyConfig(time_channel_dim=8)))
|
||||
|
||||
|
||||
def _sampler_tensors(seed: int) -> dict[str, torch.Tensor]:
|
||||
generator = torch.Generator(device="cpu").manual_seed(seed)
|
||||
return {
|
||||
"latent": torch.randn(1, 4, 2, 1, 3, generator=generator),
|
||||
"audio_latent": torch.randn(1, 5, 4, generator=generator),
|
||||
"txt_feat": torch.randn(1, 3, 4, generator=generator),
|
||||
"null_txt_feat": torch.randn(1, 2, 4, generator=generator),
|
||||
}
|
||||
|
||||
|
||||
def _longer_text_tensors(seed: int) -> dict[str, torch.Tensor]:
|
||||
tensors = _sampler_tensors(seed)
|
||||
generator = torch.Generator(device="cpu").manual_seed(seed + 100)
|
||||
tensors["txt_feat"] = torch.randn(1, 4, 4, generator=generator)
|
||||
return tensors
|
||||
|
||||
|
||||
def test_prepare_model_input_keeps_lengths_on_host() -> None:
|
||||
sampler = Magi2PreviewSampler(torch.nn.Identity())
|
||||
model_input = sampler.prepare_model_input(**_sampler_tensors(0), t=torch.tensor([500.0]), cfg_config=CFGConfig())
|
||||
|
||||
assert model_input.audio_feat_len == [5, 5]
|
||||
assert model_input.txt_feat_len == [3, 2]
|
||||
assert model_input.ref_audio_feat_len == [0, 0]
|
||||
assert model_input.ref_video_feat_len == [0, 0]
|
||||
|
||||
positive, negative = Magi2PreviewSampler._split_cfg_model_input(model_input)
|
||||
assert positive.txt_feat_len == [3]
|
||||
assert negative.txt_feat_len == [2]
|
||||
|
||||
|
||||
def test_pre_adapter_embeds_directly_in_checkpoint_dtype() -> None:
|
||||
model = _tiny_model(params_dtype=torch.bfloat16)
|
||||
packed = torch.randn(6, 4)
|
||||
indices = torch.tensor([0, 1]), torch.tensor([2, 3]), torch.tensor([4, 5])
|
||||
|
||||
with torch.no_grad():
|
||||
hidden = model.pre_adapter(packed, *indices)
|
||||
|
||||
assert hidden.dtype == torch.bfloat16
|
||||
assert hidden.shape == (6, model.config.virtual_width)
|
||||
@@ -0,0 +1,157 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
import os
|
||||
import time
|
||||
from datetime import timedelta
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
import torch.multiprocessing as mp
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
|
||||
|
||||
|
||||
def _collective_worker(rank, rendezvous, backend):
|
||||
from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
|
||||
from vllm_omni.errors import OmniClientError
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.model_executor.models.minimax_h3.conditioning import MiniMaxH3EncoderMediaInput
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
torch.set_num_threads(1)
|
||||
device = torch.device("cuda", rank) if backend == "nccl" else torch.device("cpu")
|
||||
if backend == "nccl":
|
||||
current_omni_platform.set_device(device)
|
||||
dist.init_process_group(
|
||||
backend, init_method=f"file://{rendezvous}", rank=rank, world_size=2, timeout=timedelta(seconds=45)
|
||||
)
|
||||
try:
|
||||
module._dit_rank_world = lambda: (dist.group.WORLD, rank, 2)
|
||||
|
||||
class VisualCodec:
|
||||
fail = False
|
||||
|
||||
def is_distributed_enabled(self):
|
||||
return True
|
||||
|
||||
def collective(self):
|
||||
value = torch.ones(1, device=device)
|
||||
dist.all_reduce(value)
|
||||
assert value.item() == 2
|
||||
|
||||
def encode_image(self, image):
|
||||
self.collective()
|
||||
assert image.size == (32, 32)
|
||||
return torch.ones(1, 96)
|
||||
|
||||
def encode_video(self, frames):
|
||||
self.collective()
|
||||
assert frames.shape == (4, 32, 32, 3)
|
||||
if self.fail and rank == 1:
|
||||
raise ValueError("codec failed on rank 1")
|
||||
return torch.full((1, 96), 2.0), (1, 2, 2)
|
||||
|
||||
class AudioCodec:
|
||||
def encode_waveform(self, waveform, sample_rate):
|
||||
assert rank == 0
|
||||
assert sample_rate == 32000
|
||||
assert waveform.shape == (64000,)
|
||||
return torch.full((160, 32), 3.0), 80
|
||||
|
||||
model = object.__new__(module.MiniMaxH3Pipeline)
|
||||
torch.nn.Module.__init__(model)
|
||||
model.device = device
|
||||
model.od_config = SimpleNamespace(enable_layerwise_offload=False)
|
||||
model.video_vae = VisualCodec()
|
||||
model.audio_vae = AudioCodec()
|
||||
model.partition = "combined"
|
||||
model.supported_tasks = frozenset({"t2va", "fl2va", "ref2va"})
|
||||
model.text_encoder_tp_size = 2
|
||||
model.text_encoder_group = SimpleNamespace(
|
||||
world_size=2, ranks=[0, 1], rank_in_group=rank, device_group=dist.group.WORLD
|
||||
)
|
||||
model.processor = SimpleNamespace(image_processor=SimpleNamespace(merge_size=2))
|
||||
model.tokenizer = lambda text, **kwargs: {"input_ids": [1, 2, 3]}
|
||||
|
||||
def encode_text(ids, kwargs):
|
||||
assert ids.tolist() == [1, 2, 3]
|
||||
assert not kwargs
|
||||
hidden = torch.ones(3, 5120, device=device, dtype=torch.bfloat16)
|
||||
dist.all_reduce(hidden)
|
||||
return hidden
|
||||
|
||||
model._encode_text_hidden = encode_text
|
||||
sampling = OmniDiffusionSamplingParams(
|
||||
height=32, width=32, num_frames=96, extra_args={"task": "t2va", "aspect_ratio": "1:1"}
|
||||
)
|
||||
text_conditioning = model._prepare_local_conditioning({"prompt": "test"}, sampling)
|
||||
torch.testing.assert_close(text_conditioning.hidden_states, torch.full_like(text_conditioning.hidden_states, 2))
|
||||
assert text_conditioning.visual_condition is None
|
||||
|
||||
media = (
|
||||
MiniMaxH3EncoderMediaInput(
|
||||
task="ref2va",
|
||||
height=32,
|
||||
width=32,
|
||||
num_frames=107,
|
||||
latent_t=32,
|
||||
audio_t=178,
|
||||
images=(torch.zeros(32, 32, 3, dtype=torch.uint8),),
|
||||
videos=(torch.zeros(4, 32, 32, 3, dtype=torch.uint8),),
|
||||
video_audios=(None,),
|
||||
audios=((torch.zeros(64000), 32000),),
|
||||
)
|
||||
if rank == 0
|
||||
else None
|
||||
)
|
||||
conditioning = model._encode_local_media(media)
|
||||
assert [block["kind"] for block in conditioning.ref_blocks] == ["image", "video", "audio"]
|
||||
torch.testing.assert_close(
|
||||
conditioning.visual_condition.cpu(), torch.cat([torch.ones(1, 96), torch.full((1, 96), 2.0)])
|
||||
)
|
||||
torch.testing.assert_close(conditioning.audio_condition.cpu(), torch.full((160, 32), 3.0))
|
||||
assert conditioning.audio_condition_lengths == (80,)
|
||||
|
||||
model.video_vae.fail = True
|
||||
with pytest.raises(OmniClientError, match="codec failed on rank 1"):
|
||||
model._encode_local_media(media)
|
||||
with pytest.raises(OmniClientError, match="non-empty prompt"):
|
||||
model._prepare_local_conditioning({"prompt": ""}, sampling)
|
||||
model.video_vae.fail = False
|
||||
assert model._encode_local_media(media).visual_condition_shapes == ((1, 2, 2), (1, 2, 2))
|
||||
finally:
|
||||
dist.destroy_process_group()
|
||||
|
||||
|
||||
def _run_collectives(tmp_path, backend):
|
||||
context = mp.spawn(_collective_worker, args=(str(tmp_path / "rendezvous"), backend), nprocs=2, join=False)
|
||||
deadline = time.monotonic() + 150
|
||||
try:
|
||||
while not context.join(timeout=1):
|
||||
if time.monotonic() > deadline:
|
||||
pytest.fail("MiniMax H3 collectives did not finish within 150 seconds")
|
||||
finally:
|
||||
for process in context.processes:
|
||||
if process.is_alive():
|
||||
process.terminate()
|
||||
process.join(timeout=5)
|
||||
|
||||
|
||||
@pytest.mark.cpu
|
||||
def test_two_rank_conditioning_and_errors_gloo(tmp_path):
|
||||
_run_collectives(tmp_path, "gloo")
|
||||
|
||||
|
||||
@pytest.mark.cuda
|
||||
@pytest.mark.parallel
|
||||
@pytest.mark.skipif(
|
||||
os.environ.get("VLLM_TEST_MINIMAX_H3_COLLECTIVES") != "1",
|
||||
reason="set VLLM_TEST_MINIMAX_H3_COLLECTIVES=1 with two reserved GPUs",
|
||||
)
|
||||
def test_two_rank_conditioning_and_errors_nccl(tmp_path):
|
||||
if torch.accelerator.device_count() < 2:
|
||||
pytest.skip("requires two GPUs")
|
||||
_run_collectives(tmp_path, "nccl")
|
||||
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,133 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def pipeline():
|
||||
model = object.__new__(MiniMaxH3Pipeline)
|
||||
torch.nn.Module.__init__(model)
|
||||
model.load_text_encoder = True
|
||||
model.load_vae_encoder = True
|
||||
model.device = torch.device("cpu")
|
||||
model.partition = "combined"
|
||||
model.supported_tasks = frozenset({"t2va", "fl2va", "ref2va"})
|
||||
model.default_video_shift = 12.0
|
||||
model.default_audio_shift = 3.0
|
||||
model.od_config = SimpleNamespace(enable_layerwise_offload=False)
|
||||
model._quality_policy = SimpleNamespace(resolve=Mock(return_value=SimpleNamespace(cache_dit=None)))
|
||||
model._cache_dit_runtime = SimpleNamespace(prepare=Mock())
|
||||
model.encode_prompt = Mock(
|
||||
return_value=(torch.ones(2, 5120, dtype=torch.bfloat16), torch.ones(2, dtype=torch.long))
|
||||
)
|
||||
model.video_vae = Mock()
|
||||
model.video_vae.is_distributed_enabled.return_value = False
|
||||
model.video_vae.encode_image.side_effect = lambda image: torch.ones((image.height // 32) * (image.width // 32), 96)
|
||||
model.audio_vae = Mock()
|
||||
model.transformer = torch.nn.Identity()
|
||||
model.diffuse = Mock(return_value=(torch.zeros(1), torch.zeros(1)))
|
||||
model.decode = Mock(return_value=(torch.zeros(1, 3, 1, 32, 32), torch.zeros(1)))
|
||||
return model
|
||||
|
||||
|
||||
def _request(task):
|
||||
prompt: dict[str, Any] = {"prompt": "A person waves."}
|
||||
if task != "t2va":
|
||||
prompt["multi_modal_data"] = {"image": Image.new("RGB", (256, 256))}
|
||||
sampling = OmniDiffusionSamplingParams(
|
||||
height=32, width=32, num_frames=96, num_inference_steps=2, extra_args={"task": task, "aspect_ratio": "1:1"}
|
||||
)
|
||||
return SimpleNamespace(prompts=[prompt], sampling_params=sampling)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("task", ["t2va", "fl2va", "ref2va"])
|
||||
def test_single_stage_encodes_without_stage_zero(pipeline, task):
|
||||
output = pipeline.forward(_request(task))
|
||||
pipeline.encode_prompt.assert_called_once()
|
||||
assert pipeline.video_vae.encode_image.call_count == int(task != "t2va")
|
||||
pipeline.diffuse.assert_called_once()
|
||||
assert pipeline.diffuse.call_args.kwargs["task"] == task
|
||||
video, audio = output.output
|
||||
torch.testing.assert_close(video, torch.zeros(1, 1, 32, 32, 3, dtype=torch.uint8))
|
||||
torch.testing.assert_close(audio, pipeline.decode.return_value[1])
|
||||
|
||||
|
||||
def test_mixed_reference_conditioning_matches_stage_zero_handoff(pipeline, monkeypatch):
|
||||
import numpy as np
|
||||
|
||||
from vllm_omni.data_entry_keys import flatten_payload
|
||||
from vllm_omni.model_executor.models.minimax_h3 import encoder_processing as processing
|
||||
from vllm_omni.model_executor.models.minimax_h3.conditioning import (
|
||||
MiniMaxH3EncoderConditioning,
|
||||
MiniMaxH3TextConditioning,
|
||||
)
|
||||
from vllm_omni.model_executor.models.minimax_h3.encoder import MiniMaxH3Encoder
|
||||
from vllm_omni.model_executor.stage_input_processors.minimax_h3 import encoder2diffusion, prepare_encoder_prompt
|
||||
|
||||
frames = np.zeros((48, 32, 32, 3), dtype=np.uint8)
|
||||
decode_frames = Mock(return_value=frames)
|
||||
sampler = Mock(wraps=processing.sample_reference_video_frames)
|
||||
monkeypatch.setattr(processing, "load_video_frames", decode_frames)
|
||||
monkeypatch.setattr(processing, "sample_reference_video_frames", sampler)
|
||||
monkeypatch.setattr(
|
||||
processing,
|
||||
"prepare_reference_videos",
|
||||
Mock(
|
||||
return_value=[{"prepared_path": "prepared.mp4", "original_path": "original.mp4", "input_has_audio": True}]
|
||||
),
|
||||
)
|
||||
monkeypatch.setattr(processing, "load_video_audio", Mock(return_value=(torch.full((64000,), 3.0), 32000)))
|
||||
pipeline.video_vae.encode_video.return_value = (torch.full((1, 96), 2.0), (1, 2, 2))
|
||||
|
||||
def encode_audio(waveform, sample_rate):
|
||||
length = round(waveform.shape[-1] / sample_rate * 40)
|
||||
return torch.full((length * 2, 32), float(waveform[0])), length
|
||||
|
||||
pipeline.audio_vae.encode_waveform.side_effect = encode_audio
|
||||
request = _request("ref2va")
|
||||
request.prompts[0]["multi_modal_data"].update(video="original.mp4", audio=(torch.full((192000,), 4.0), 32000))
|
||||
pipeline.forward(request)
|
||||
local_kwargs = pipeline.diffuse.call_args.kwargs
|
||||
|
||||
transformed = prepare_encoder_prompt(request.prompts[0], [request.sampling_params])
|
||||
stage_zero = object.__new__(MiniMaxH3Encoder)
|
||||
torch.nn.Module.__init__(stage_zero)
|
||||
stage_zero._component_leader = True
|
||||
stage_zero.video_vae = pipeline.video_vae
|
||||
stage_zero.audio_vae = pipeline.audio_vae
|
||||
media = stage_zero._encode_media(stage_zero._media_input(transformed["additional_information"]))
|
||||
text = MiniMaxH3TextConditioning(*pipeline.encode_prompt.return_value)
|
||||
payload = MiniMaxH3EncoderConditioning.from_components(text, media).to_omni_payload()
|
||||
source = SimpleNamespace(finished=True, outputs=[SimpleNamespace(multimodal_output=flatten_payload(payload))])
|
||||
bridged = encoder2diffusion([source], request.prompts[0])
|
||||
pipeline.load_text_encoder = False
|
||||
pipeline.load_vae_encoder = False
|
||||
pipeline.encode_prompt.reset_mock()
|
||||
pipeline.forward(SimpleNamespace(prompts=[bridged], sampling_params=request.sampling_params))
|
||||
external_kwargs = pipeline.diffuse.call_args.kwargs
|
||||
for key, value in local_kwargs.items():
|
||||
if isinstance(value, torch.Tensor):
|
||||
torch.testing.assert_close(value, external_kwargs[key])
|
||||
else:
|
||||
assert value == external_kwargs[key]
|
||||
pipeline.encode_prompt.assert_not_called()
|
||||
assert decode_frames.call_count == sampler.call_count == 2
|
||||
assert all(call.kwargs["decoded_frames"] is frames for call in sampler.call_args_list)
|
||||
assert [block["kind"] for block in local_kwargs["ref_blocks"]] == ["image", "video_audio", "audio"]
|
||||
assert local_kwargs["audio_condition_lengths"] == [80, round(107 / 24 * 40)]
|
||||
assert all(
|
||||
call.args[0].shape[-1] == round(107 / 24 * 32000)
|
||||
for call in pipeline.audio_vae.encode_waveform.call_args_list[1::2]
|
||||
)
|
||||
@@ -6,7 +6,7 @@ import torch
|
||||
from vllm.triton_utils import HAS_TRITON
|
||||
|
||||
from tests.helpers.mark import hardware_marks
|
||||
from vllm_omni.diffusion.attention.ops.minimax_h3_modulation import (
|
||||
from vllm_omni.diffusion.layers.indexed_modulation import (
|
||||
_MAX_1D_GRID_SIZE,
|
||||
_iter_row_chunks,
|
||||
_launch_row_chunks,
|
||||
|
||||
@@ -9,72 +9,25 @@ from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from PIL import Image
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
|
||||
def test_h3_profiler_targets_cover_reference_encoders_and_vae_decoders():
|
||||
def test_h3_full_pipeline_profiler_includes_local_encoding():
|
||||
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
|
||||
|
||||
targets = MiniMaxH3Pipeline._PROFILER_TARGETS
|
||||
assert "_encode_visual_conditions" in targets
|
||||
assert "_encode_reference_audio_conditions" in targets
|
||||
assert "video_vae.decode_latent" in targets
|
||||
assert "audio_vae.decode_latent" in targets
|
||||
assert "_encode_video_conditions" not in targets
|
||||
assert "_encode_video_audio_conditions" not in targets
|
||||
assert "_encode_audio_conditions" not in targets
|
||||
|
||||
|
||||
@pytest.mark.parametrize("distributed", [False, True])
|
||||
def test_manual_component_offload_wins_over_requested_model_offload(distributed):
|
||||
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
|
||||
|
||||
pipeline = object.__new__(MiniMaxH3Pipeline)
|
||||
torch.nn.Module.__init__(pipeline)
|
||||
pipeline.od_config = SimpleNamespace(
|
||||
enable_cpu_offload=True,
|
||||
enable_layerwise_offload=not distributed,
|
||||
enable_distributed_layerwise_offload=distributed,
|
||||
)
|
||||
pipeline._model_cpu_offload_modules = []
|
||||
pipeline.text_encoder = Mock()
|
||||
expected = torch.ones(2, 3)
|
||||
pipeline.text_encoder.encode_ids.return_value = expected
|
||||
component = Mock()
|
||||
|
||||
actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {})
|
||||
with pipeline._component_on_device(component):
|
||||
component.load_to_device.assert_called_once_with()
|
||||
component.offload_to_cpu.assert_not_called()
|
||||
|
||||
assert actual is expected
|
||||
pipeline.text_encoder.load_to_device.assert_called_once_with()
|
||||
pipeline.text_encoder.offload_to_cpu.assert_called_once_with()
|
||||
component.offload_to_cpu.assert_called_once_with()
|
||||
|
||||
|
||||
def test_requested_model_offload_without_backend_uses_resident_components():
|
||||
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
|
||||
|
||||
pipeline = object.__new__(MiniMaxH3Pipeline)
|
||||
torch.nn.Module.__init__(pipeline)
|
||||
pipeline.od_config = SimpleNamespace(
|
||||
enable_cpu_offload=True,
|
||||
enable_layerwise_offload=False,
|
||||
enable_distributed_layerwise_offload=False,
|
||||
)
|
||||
pipeline._model_cpu_offload_modules = []
|
||||
pipeline.text_encoder = Mock()
|
||||
expected = torch.ones(2, 3)
|
||||
pipeline.text_encoder.encode_ids.return_value = expected
|
||||
|
||||
actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {})
|
||||
|
||||
assert actual is expected
|
||||
pipeline.text_encoder.load_to_device.assert_called_once_with()
|
||||
pipeline.text_encoder.encode_ids.assert_called_once()
|
||||
assert targets == [
|
||||
"encode_prompt",
|
||||
"_encode_local_media",
|
||||
"diffuse",
|
||||
"decode",
|
||||
"video_vae.decode_latent",
|
||||
"audio_vae.decode_latent",
|
||||
"prepare_encode",
|
||||
"denoise_step",
|
||||
"post_decode",
|
||||
]
|
||||
|
||||
|
||||
def test_decoded_output_is_unchanged_without_active_model_offload():
|
||||
@@ -194,7 +147,7 @@ def test_forward_releases_video_after_quantization_and_before_next_seed(monkeypa
|
||||
_extract_prompt=lambda raw: ("a prompt", {}),
|
||||
_extract_text_conditioning=lambda raw: None,
|
||||
_extract_prepared_reference_videos=lambda raw: None,
|
||||
_prepare_request_inputs=lambda **kwargs: context,
|
||||
_prepare_request_inputs=lambda *args, **kwargs: context,
|
||||
_denoise_kwargs=lambda ctx: {},
|
||||
diffuse=diffuse,
|
||||
decode=decode,
|
||||
@@ -243,7 +196,7 @@ def test_forward_drops_the_decoded_video_before_the_next_seed_diffuses(monkeypat
|
||||
_extract_prompt=lambda raw: ("a prompt", {}),
|
||||
_extract_text_conditioning=lambda raw: None,
|
||||
_extract_prepared_reference_videos=lambda raw: None,
|
||||
_prepare_request_inputs=lambda **kwargs: context,
|
||||
_prepare_request_inputs=lambda *args, **kwargs: context,
|
||||
_denoise_kwargs=lambda ctx: {},
|
||||
diffuse=diffuse,
|
||||
decode=decode,
|
||||
@@ -318,7 +271,8 @@ def test_output_owner_rank_matches_executor_contract(monkeypatch, available, ini
|
||||
assert MiniMaxH3Pipeline._is_output_owner_rank() is expected
|
||||
|
||||
|
||||
def test_h3_model_cpu_offload_registers_direct_vae_stages(monkeypatch):
|
||||
@pytest.mark.parametrize("load_text_encoder", [True, False])
|
||||
def test_h3_model_cpu_offload_registers_direct_vae_stages(monkeypatch, load_text_encoder):
|
||||
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
|
||||
from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
|
||||
|
||||
@@ -326,9 +280,11 @@ def test_h3_model_cpu_offload_registers_direct_vae_stages(monkeypatch):
|
||||
torch.nn.Module.__init__(pipeline)
|
||||
pipeline.transformer = torch.nn.Linear(2, 2)
|
||||
pipeline.transformers_ref = torch.nn.Linear(2, 2)
|
||||
pipeline.text_encoder = torch.nn.Linear(2, 2)
|
||||
pipeline.video_vae = torch.nn.Linear(2, 2)
|
||||
pipeline.audio_vae = torch.nn.Linear(2, 2)
|
||||
pipeline._encoder_modules = ["text_encoder"] if load_text_encoder else []
|
||||
if load_text_encoder:
|
||||
pipeline.text_encoder = torch.nn.Linear(2, 2)
|
||||
apply_offload = Mock()
|
||||
remove_offload = Mock()
|
||||
monkeypatch.setattr(module, "apply_sequential_offload", apply_offload)
|
||||
@@ -341,7 +297,7 @@ def test_h3_model_cpu_offload_registers_direct_vae_stages(monkeypatch):
|
||||
)
|
||||
|
||||
dits = [pipeline.transformer, pipeline.transformers_ref]
|
||||
stages = [pipeline.text_encoder, pipeline.video_vae, pipeline.audio_vae]
|
||||
stages = [*([pipeline.text_encoder] if load_text_encoder else []), pipeline.video_vae, pipeline.audio_vae]
|
||||
apply_offload.assert_called_once_with(
|
||||
dit_modules=dits,
|
||||
encoder_modules=stages,
|
||||
@@ -356,6 +312,46 @@ def test_h3_model_cpu_offload_registers_direct_vae_stages(monkeypatch):
|
||||
remove_offload.assert_called_once_with([*dits, *stages])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("decode_fails", [False, True])
|
||||
def test_h3_model_cpu_offload_scopes_direct_vae_call(monkeypatch, decode_fails):
|
||||
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
|
||||
from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
|
||||
|
||||
pipeline = object.__new__(MiniMaxH3Pipeline)
|
||||
torch.nn.Module.__init__(pipeline)
|
||||
component = torch.nn.Linear(2, 2)
|
||||
events: list[tuple[Any, ...]] = []
|
||||
|
||||
@contextmanager
|
||||
def record_component(value):
|
||||
events.append(("activate", value))
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
events.append(("offload", value))
|
||||
|
||||
monkeypatch.setattr(module, "sequential_offload_component", record_component)
|
||||
pipeline._model_cpu_offload_modules = [component]
|
||||
|
||||
def decode():
|
||||
with pipeline._component_on_device(component):
|
||||
events.append(("decode", component))
|
||||
if decode_fails:
|
||||
raise RuntimeError("decode failed")
|
||||
|
||||
if decode_fails:
|
||||
with pytest.raises(RuntimeError, match="decode failed"):
|
||||
decode()
|
||||
else:
|
||||
decode()
|
||||
|
||||
assert events == [
|
||||
("activate", component),
|
||||
("decode", component),
|
||||
("offload", component),
|
||||
]
|
||||
|
||||
|
||||
def test_h3_model_cpu_offload_keeps_unselected_vaes_resident(monkeypatch):
|
||||
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
|
||||
from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
|
||||
@@ -416,44 +412,20 @@ def test_h3_model_cpu_offload_rejects_unloaded_selected_encoder(monkeypatch):
|
||||
apply_offload.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("decode_fails", [False, True])
|
||||
def test_h3_model_cpu_offload_scopes_direct_vae_call(monkeypatch, decode_fails):
|
||||
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
|
||||
from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as module
|
||||
def _encode_media(pipeline, **kwargs):
|
||||
from vllm_omni.model_executor.models.minimax_h3.conditioning import MiniMaxH3EncoderMediaInput
|
||||
from vllm_omni.model_executor.models.minimax_h3.encoder_processing import encode_media
|
||||
|
||||
pipeline = object.__new__(MiniMaxH3Pipeline)
|
||||
torch.nn.Module.__init__(pipeline)
|
||||
component = torch.nn.Linear(2, 2)
|
||||
events: list[tuple[Any, ...]] = []
|
||||
|
||||
@contextmanager
|
||||
def record_component(value):
|
||||
events.append(("activate", value))
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
events.append(("offload", value))
|
||||
|
||||
monkeypatch.setattr(module, "sequential_offload_component", record_component)
|
||||
pipeline._model_cpu_offload_modules = [component]
|
||||
|
||||
def decode():
|
||||
with pipeline._component_on_device(component):
|
||||
events.append(("decode", component))
|
||||
if decode_fails:
|
||||
raise RuntimeError("decode failed")
|
||||
|
||||
if decode_fails:
|
||||
with pytest.raises(RuntimeError, match="decode failed"):
|
||||
decode()
|
||||
else:
|
||||
decode()
|
||||
|
||||
assert events == [
|
||||
("activate", component),
|
||||
("decode", component),
|
||||
("offload", component),
|
||||
]
|
||||
media = MiniMaxH3EncoderMediaInput(
|
||||
task="ref2va", height=32, width=32, num_frames=124, latent_t=31, audio_t=208, **kwargs
|
||||
)
|
||||
return encode_media(
|
||||
media,
|
||||
video_vae=getattr(pipeline, "video_vae", None),
|
||||
audio_vae=getattr(pipeline, "audio_vae", None),
|
||||
emit_conditioning=True,
|
||||
component_scope=pipeline._component_on_device,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("encode_fails", [False, True])
|
||||
@@ -463,7 +435,6 @@ def test_h3_model_cpu_offload_batches_visual_reference_scope(monkeypatch, encode
|
||||
|
||||
pipeline = object.__new__(MiniMaxH3Pipeline)
|
||||
torch.nn.Module.__init__(pipeline)
|
||||
pipeline.device = torch.device("cpu")
|
||||
pipeline.video_vae = Mock()
|
||||
pipeline._model_cpu_offload_modules = [pipeline.video_vae]
|
||||
events = []
|
||||
@@ -484,17 +455,15 @@ def test_h3_model_cpu_offload_batches_visual_reference_scope(monkeypatch, encode
|
||||
|
||||
pipeline.video_vae.encode_image.side_effect = encode_image
|
||||
monkeypatch.setattr(module, "sequential_offload_component", record_component)
|
||||
monkeypatch.setattr(module, "_dit_rank_world", lambda: (None, 0, 1))
|
||||
monkeypatch.setattr(module, "_broadcast_tensor", lambda value, **kwargs: value)
|
||||
images = [Image.new("RGB", (16, 16)), Image.new("RGB", (32, 16))]
|
||||
images = (torch.zeros(16, 16, 3, dtype=torch.uint8), torch.zeros(16, 32, 3, dtype=torch.uint8))
|
||||
|
||||
if encode_fails:
|
||||
with pytest.raises(RuntimeError, match="encode failed"):
|
||||
pipeline._encode_visual_conditions(images, None, video_count=0)
|
||||
_encode_media(pipeline, images=images)
|
||||
else:
|
||||
rows, shapes = pipeline._encode_visual_conditions(images, None, video_count=0)
|
||||
assert rows.shape == (3, 4)
|
||||
assert shapes == [(1, 1, 1), (1, 1, 2)]
|
||||
result = _encode_media(pipeline, images=images)
|
||||
assert result.visual_condition.shape == (3, 4)
|
||||
assert result.visual_condition_shapes == ((1, 1, 1), (1, 1, 2))
|
||||
|
||||
assert events == [
|
||||
("activate", pipeline.video_vae),
|
||||
@@ -510,9 +479,7 @@ def test_h3_model_cpu_offload_shares_image_and_video_scope(monkeypatch):
|
||||
|
||||
pipeline = object.__new__(MiniMaxH3Pipeline)
|
||||
torch.nn.Module.__init__(pipeline)
|
||||
pipeline.device = torch.device("cpu")
|
||||
pipeline.video_vae = Mock()
|
||||
pipeline.video_vae.is_distributed_enabled.return_value = False
|
||||
pipeline._model_cpu_offload_modules = [pipeline.video_vae]
|
||||
events = []
|
||||
|
||||
@@ -529,31 +496,26 @@ def test_h3_model_cpu_offload_shares_image_and_video_scope(monkeypatch):
|
||||
return torch.ones(1, 4)
|
||||
|
||||
def encode_video(frames):
|
||||
events.append(("video", frames))
|
||||
events.append(("video", frames.shape))
|
||||
return torch.ones(3, 4), (3, 2, 2)
|
||||
|
||||
pipeline.video_vae.encode_image.side_effect = encode_image
|
||||
pipeline.video_vae.encode_video.side_effect = encode_video
|
||||
monkeypatch.setattr(module, "sequential_offload_component", record_component)
|
||||
monkeypatch.setattr(module, "_dit_rank_world", lambda: (None, 0, 1))
|
||||
monkeypatch.setattr(module, "_broadcast_tensor", lambda value, **kwargs: value)
|
||||
monkeypatch.setattr(module, "load_video_frames", lambda path: f"frames:{path}")
|
||||
images = [Image.new("RGB", (16, 16)), Image.new("RGB", (32, 16))]
|
||||
prepared_videos = [{"prepared_path": "reference.mp4"}]
|
||||
|
||||
rows, shapes = pipeline._encode_visual_conditions(
|
||||
images,
|
||||
prepared_videos,
|
||||
video_count=1,
|
||||
result = _encode_media(
|
||||
pipeline,
|
||||
images=(torch.zeros(16, 16, 3, dtype=torch.uint8), torch.zeros(16, 32, 3, dtype=torch.uint8)),
|
||||
videos=(torch.zeros(12, 32, 32, 3, dtype=torch.uint8),),
|
||||
video_audios=(None,),
|
||||
)
|
||||
|
||||
assert rows.shape == (5, 4)
|
||||
assert shapes == [(1, 1, 1), (1, 1, 2), (3, 2, 2)]
|
||||
assert result.visual_condition.shape == (5, 4)
|
||||
assert result.visual_condition_shapes == ((1, 1, 1), (1, 1, 2), (3, 2, 2))
|
||||
assert events == [
|
||||
("activate", pipeline.video_vae),
|
||||
("image", (16, 16)),
|
||||
("image", (32, 16)),
|
||||
("video", "frames:reference.mp4"),
|
||||
("video", (12, 32, 32, 3)),
|
||||
("offload", pipeline.video_vae),
|
||||
]
|
||||
|
||||
@@ -567,7 +529,7 @@ def test_h3_model_cpu_offload_shares_embedded_and_standalone_audio_scope(monkeyp
|
||||
torch.nn.Module.__init__(pipeline)
|
||||
pipeline.audio_vae = Mock()
|
||||
pipeline._model_cpu_offload_modules = [pipeline.audio_vae]
|
||||
events: list[tuple[Any, ...]] = []
|
||||
events = []
|
||||
|
||||
@contextmanager
|
||||
def record_component(value):
|
||||
@@ -577,26 +539,20 @@ def test_h3_model_cpu_offload_shares_embedded_and_standalone_audio_scope(monkeyp
|
||||
finally:
|
||||
events.append(("offload", value))
|
||||
|
||||
def encode_embedded(*args, **kwargs):
|
||||
events.append(("embedded",))
|
||||
def encode_waveform(waveform, sample_rate):
|
||||
events.append(("audio", int(waveform[0]), sample_rate))
|
||||
if encode_fails:
|
||||
raise RuntimeError("audio encode failed")
|
||||
return torch.ones(1, 2), [1]
|
||||
return torch.ones(80, 32), 80
|
||||
|
||||
def encode_standalone(*args, **kwargs):
|
||||
events.append(("standalone",))
|
||||
return torch.ones(1, 2), [1]
|
||||
|
||||
pipeline._encode_video_audio_conditions_resident = encode_embedded
|
||||
pipeline._encode_audio_conditions_resident = encode_standalone
|
||||
pipeline.audio_vae.encode_waveform.side_effect = encode_waveform
|
||||
monkeypatch.setattr(module, "sequential_offload_component", record_component)
|
||||
|
||||
def call():
|
||||
return pipeline._encode_reference_audio_conditions(
|
||||
[{"input_has_audio": True}],
|
||||
has_audio=[True],
|
||||
standalone_audios=[(torch.ones(1), 16_000)],
|
||||
max_duration_seconds=1.0,
|
||||
return _encode_media(
|
||||
pipeline,
|
||||
video_audios=(),
|
||||
audios=((torch.ones(32_000), 16_000),),
|
||||
)
|
||||
|
||||
if encode_fails:
|
||||
@@ -604,10 +560,61 @@ def test_h3_model_cpu_offload_shares_embedded_and_standalone_audio_scope(monkeyp
|
||||
call()
|
||||
else:
|
||||
result = call()
|
||||
assert result[0].shape == (1, 2)
|
||||
assert result[2].shape == (1, 2)
|
||||
assert result.audio_condition.shape == (80, 32)
|
||||
assert result.audio_condition_lengths == (80,)
|
||||
assert events[1:-1] == [("audio", 1, 16_000)]
|
||||
|
||||
assert events[0] == ("activate", pipeline.audio_vae)
|
||||
assert events[-1] == ("offload", pipeline.audio_vae)
|
||||
assert events.count(("activate", pipeline.audio_vae)) == 1
|
||||
assert events.count(("offload", pipeline.audio_vae)) == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("distributed", [False, True])
|
||||
def test_manual_component_offload_wins_over_requested_model_offload(distributed):
|
||||
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
|
||||
|
||||
pipeline = object.__new__(MiniMaxH3Pipeline)
|
||||
torch.nn.Module.__init__(pipeline)
|
||||
pipeline.od_config = SimpleNamespace(
|
||||
enable_cpu_offload=True,
|
||||
enable_layerwise_offload=not distributed,
|
||||
enable_distributed_layerwise_offload=distributed,
|
||||
)
|
||||
pipeline._model_cpu_offload_modules = []
|
||||
pipeline.text_encoder = Mock()
|
||||
expected = torch.ones(2, 3)
|
||||
pipeline.text_encoder.encode_ids.return_value = expected
|
||||
component = Mock()
|
||||
|
||||
actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {})
|
||||
with pipeline._component_on_device(component):
|
||||
component.load_to_device.assert_called_once_with()
|
||||
component.offload_to_cpu.assert_not_called()
|
||||
|
||||
assert actual is expected
|
||||
pipeline.text_encoder.load_to_device.assert_called_once_with()
|
||||
pipeline.text_encoder.offload_to_cpu.assert_called_once_with()
|
||||
component.offload_to_cpu.assert_called_once_with()
|
||||
|
||||
|
||||
def test_requested_model_offload_without_backend_uses_resident_components():
|
||||
from vllm_omni.diffusion.models.minimax_h3 import MiniMaxH3Pipeline
|
||||
|
||||
pipeline = object.__new__(MiniMaxH3Pipeline)
|
||||
torch.nn.Module.__init__(pipeline)
|
||||
pipeline.od_config = SimpleNamespace(
|
||||
enable_cpu_offload=True,
|
||||
enable_layerwise_offload=False,
|
||||
enable_distributed_layerwise_offload=False,
|
||||
)
|
||||
pipeline._model_cpu_offload_modules = []
|
||||
pipeline.text_encoder = Mock()
|
||||
expected = torch.ones(2, 3)
|
||||
pipeline.text_encoder.encode_ids.return_value = expected
|
||||
|
||||
actual = pipeline._encode_text_hidden(torch.tensor([1, 2]), {})
|
||||
|
||||
assert actual is expected
|
||||
pipeline.text_encoder.load_to_device.assert_called_once_with()
|
||||
pipeline.text_encoder.encode_ids.assert_called_once()
|
||||
@@ -186,9 +186,8 @@ def test_minimax_h3_quantization_quality(config: QualityTestConfig):
|
||||
common_kwargs = {
|
||||
"model": model_ref,
|
||||
"enforce_eager": True,
|
||||
"tensor_parallel_size": 2,
|
||||
# The fused BF16 baseline only fits on H100-80GB with encoder TP.
|
||||
"text_encoder_tp_size": 2,
|
||||
"stage_0_tensor_parallel_size": 2,
|
||||
"vae_use_tiling": True,
|
||||
}
|
||||
|
||||
|
||||
@@ -110,6 +110,8 @@ def _step_pipeline(model, *, packed_batch_supported: bool = True):
|
||||
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline
|
||||
|
||||
pipeline = object.__new__(MiniMaxH3Pipeline)
|
||||
pipeline.load_text_encoder = False
|
||||
pipeline.load_vae_encoder = False
|
||||
pipeline.transformer = model
|
||||
pipeline.device = torch.device("cpu")
|
||||
pipeline._transformer_for_task = lambda task: model
|
||||
@@ -304,8 +306,17 @@ def test_prepare_encode_seeds_runner_visible_state(monkeypatch, batch_frames):
|
||||
}
|
||||
|
||||
pipeline = _step_pipeline(_SegmentMeanModel())
|
||||
monkeypatch.setattr(mod.MiniMaxH3Pipeline, "_extract_prompt", staticmethod(lambda _: ("a prompt", {})))
|
||||
monkeypatch.setattr(mod.MiniMaxH3Pipeline, "_prepare_request_inputs", lambda self, **_: context)
|
||||
conditioning = object()
|
||||
monkeypatch.setattr(
|
||||
mod.MiniMaxH3Pipeline,
|
||||
"_extract_encoder_conditioning",
|
||||
staticmethod(lambda _: conditioning),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
mod.MiniMaxH3Pipeline,
|
||||
"_prepare_encoder_conditioning_inputs",
|
||||
lambda self, value, sampling: context,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
mod.MiniMaxH3Pipeline,
|
||||
"_build_denoise_inputs",
|
||||
@@ -439,50 +450,3 @@ def test_packed_batch_rejects_backends_that_cannot_isolate_requests(attention):
|
||||
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import MiniMaxH3Pipeline
|
||||
|
||||
assert MiniMaxH3Pipeline._packed_batch_supported(_FakeTransformer([attention])) is False
|
||||
|
||||
|
||||
def test_broadcast_rank0_exception_single_rank_reraises():
|
||||
"""Single-rank execution has no group; the helper just reraises."""
|
||||
from vllm_omni.diffusion.models.minimax_h3.pipeline_minimax_h3 import (
|
||||
_broadcast_rank0_exception,
|
||||
)
|
||||
from vllm_omni.errors import OmniClientError
|
||||
|
||||
_broadcast_rank0_exception(None)
|
||||
with pytest.raises(OmniClientError, match="bad ref"):
|
||||
_broadcast_rank0_exception(OmniClientError("bad ref"))
|
||||
|
||||
|
||||
def test_broadcast_rank0_exception_propagates_to_non_zero_ranks(monkeypatch):
|
||||
"""A rank-0 error becomes a matching client error on every other DiT rank."""
|
||||
from vllm_omni.diffusion.models.minimax_h3 import pipeline_minimax_h3 as mod
|
||||
from vllm_omni.errors import OmniClientError
|
||||
|
||||
def fake_rank_world(rank):
|
||||
return lambda: (object(), rank, 4)
|
||||
|
||||
def make_broadcast(rank0_payload):
|
||||
def fake_broadcast(payload_list, *, src, group):
|
||||
payload_list[0] = rank0_payload
|
||||
|
||||
return fake_broadcast
|
||||
|
||||
monkeypatch.setattr(mod, "_dit_rank_world", fake_rank_world(0))
|
||||
err = OmniClientError("invalid reference-video file", status_code=422, error_type="UnprocessableEntityError")
|
||||
rank0_payload = {
|
||||
"type": type(err).__name__,
|
||||
"message": str(err),
|
||||
"status_code": err.status_code,
|
||||
"error_type": err.error_type,
|
||||
}
|
||||
monkeypatch.setattr(mod.dist, "broadcast_object_list", make_broadcast(rank0_payload))
|
||||
with pytest.raises(OmniClientError) as rank0_info:
|
||||
mod._broadcast_rank0_exception(err)
|
||||
assert rank0_info.value is err
|
||||
|
||||
monkeypatch.setattr(mod, "_dit_rank_world", fake_rank_world(2))
|
||||
with pytest.raises(OmniClientError) as rank2_info:
|
||||
mod._broadcast_rank0_exception(None)
|
||||
assert rank2_info.value.status_code == 422
|
||||
assert rank2_info.value.error_type == "UnprocessableEntityError"
|
||||
assert "invalid reference-video file" in str(rank2_info.value)
|
||||
@@ -0,0 +1,221 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
import math
|
||||
import sys
|
||||
import types
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.attention.backends.abstract import (
|
||||
AttentionMetadata,
|
||||
PackedPaddingMetadata,
|
||||
VideoTokenLayout,
|
||||
VideoTokenSpan,
|
||||
)
|
||||
from vllm_omni.diffusion.models.minimax_h3.attention.fastvideo_h3 import (
|
||||
MiniMaxH3VSAImpl,
|
||||
_get_h3_tile_metadata,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _h3_metadata(prefix_segments, video_shape, *, gate=None, packed_padding=None):
|
||||
extra: dict[str, Any] = {
|
||||
"vsa_h3_prefix_segments": prefix_segments,
|
||||
}
|
||||
if gate is not None:
|
||||
extra["gate_compress"] = gate
|
||||
return AttentionMetadata(
|
||||
packed_padding=packed_padding,
|
||||
video_layout=VideoTokenLayout(
|
||||
used_len=sum(prefix_segments) + math.prod(video_shape),
|
||||
video_spans=(VideoTokenSpan(start=sum(prefix_segments), latent_grid=video_shape, role="target"),),
|
||||
),
|
||||
extra=extra,
|
||||
)
|
||||
|
||||
|
||||
def _fake_block_sparse_kernel(monkeypatch, calls):
|
||||
"""Stand in for fastvideo_kernel.block_sparse_attn with a mask-obeying SDPA."""
|
||||
|
||||
def block_sparse_attn(q, k, v, block_map, variable_block_sizes):
|
||||
calls["block_map"] = block_map.clone()
|
||||
calls["variable_block_sizes"] = variable_block_sizes.tolist()
|
||||
calls["q_shape"] = tuple(q.shape)
|
||||
blocks = block_map.shape[-1]
|
||||
assert blocks * 64 == q.shape[2]
|
||||
# Expand the per-block map to a per-token mask so the stand-in honours
|
||||
# the sparsity contract the real kernel implements, and drop the rows
|
||||
# each partial block pads with, which the real kernel skips through
|
||||
# variable_block_sizes.
|
||||
token_mask = block_map.repeat_interleave(64, dim=-1).repeat_interleave(64, dim=-2)
|
||||
within = torch.arange(64, device=q.device)[None, :] < variable_block_sizes[:, None]
|
||||
token_mask = token_mask & within.reshape(1, 1, 1, -1)
|
||||
out = torch.nn.functional.scaled_dot_product_attention(q.float(), k.float(), v.float(), attn_mask=token_mask)
|
||||
return out.to(q.dtype), None
|
||||
|
||||
module = types.ModuleType("fastvideo_kernel.block_sparse_attn")
|
||||
setattr(module, "block_sparse_attn", block_sparse_attn)
|
||||
monkeypatch.setitem(sys.modules, "fastvideo_kernel", types.ModuleType("fastvideo_kernel"))
|
||||
monkeypatch.setitem(sys.modules, "fastvideo_kernel.block_sparse_attn", module)
|
||||
|
||||
|
||||
def _h3_impl(**backend_kwargs):
|
||||
kwargs = {"topk": 1, "min_seq_len": 1, "disable_when_sp_active": False}
|
||||
kwargs.update(backend_kwargs)
|
||||
return MiniMaxH3VSAImpl(num_heads=2, head_size=8, softmax_scale=8**-0.5, backend_kwargs=kwargs)
|
||||
|
||||
|
||||
def test_h3_forward_tiles_untiles_and_restores_the_original_row_order(monkeypatch):
|
||||
calls: dict[str, Any] = {}
|
||||
_fake_block_sparse_kernel(monkeypatch, calls)
|
||||
impl = _h3_impl(topk=64)
|
||||
prefix_segments, video_shape = (5, 70), (2, 4, 4)
|
||||
seq_len = sum(prefix_segments) + 2 * 4 * 4
|
||||
torch.manual_seed(0)
|
||||
query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16)
|
||||
|
||||
output = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
|
||||
|
||||
assert output.shape == query.shape
|
||||
# topk covers every video block, so the block map is dense and the tiled
|
||||
# kernel must reproduce plain attention over the same rows.
|
||||
reference = torch.nn.functional.scaled_dot_product_attention(
|
||||
*(x.float().transpose(1, 2) for x in (query, query, query))
|
||||
).transpose(1, 2)
|
||||
torch.testing.assert_close(output.float(), reference, atol=2e-2, rtol=2e-2)
|
||||
# 5 + 70 splits into segment-pure 5 | 64 | 6, then one (2, 4, 4) video tile.
|
||||
assert calls["variable_block_sizes"] == [5, 64, 6, 32]
|
||||
|
||||
|
||||
def test_h3_forward_keeps_prefix_dense_and_selects_top_k_video_tiles(monkeypatch):
|
||||
calls: dict[str, Any] = {}
|
||||
_fake_block_sparse_kernel(monkeypatch, calls)
|
||||
impl = _h3_impl(topk=1)
|
||||
prefix_segments, video_shape = (5,), (8, 4, 4)
|
||||
seq_len = 5 + 8 * 4 * 4
|
||||
torch.manual_seed(0)
|
||||
query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16)
|
||||
|
||||
impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
|
||||
|
||||
# 1 prefix block + 2 video blocks. The odd count is padded to an even one
|
||||
# for the paired-CTA kernel contract, and the Triton route slices that
|
||||
# transport-only partner off again before launching.
|
||||
block_map = calls["block_map"]
|
||||
assert block_map.shape[-1] == 3
|
||||
assert calls["variable_block_sizes"] == [5, 64, 64]
|
||||
assert block_map[:, :, :1, :].all(), "prefix queries stay dense"
|
||||
assert block_map[..., :1].all(), "prefix keys are exempt from selection"
|
||||
assert (block_map[:, :, 1:, 1:].sum(dim=-1) == 1).all(), "one video tile per video query"
|
||||
|
||||
# The same odd geometry must still untile back onto the original rows.
|
||||
dense = _h3_impl(topk=64).forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
|
||||
reference = torch.nn.functional.scaled_dot_product_attention(
|
||||
*(x.float().transpose(1, 2) for x in (query, query, query))
|
||||
).transpose(1, 2)
|
||||
torch.testing.assert_close(dense.float(), reference, atol=2e-2, rtol=2e-2)
|
||||
|
||||
|
||||
def test_h3_forward_applies_the_learned_compression_gate(monkeypatch):
|
||||
_fake_block_sparse_kernel(monkeypatch, {})
|
||||
impl = _h3_impl(topk=64)
|
||||
prefix_segments, video_shape = (5,), (2, 4, 4)
|
||||
seq_len = 5 + 2 * 4 * 4
|
||||
torch.manual_seed(0)
|
||||
query = torch.randn(1, seq_len, 2, 8, dtype=torch.bfloat16)
|
||||
|
||||
without_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape))
|
||||
zero_gate = torch.zeros_like(query)
|
||||
with_zero_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, gate=zero_gate))
|
||||
live_gate = torch.full_like(query, 0.5)
|
||||
with_live_gate = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, gate=live_gate))
|
||||
|
||||
torch.testing.assert_close(with_zero_gate.float(), without_gate.float())
|
||||
assert not torch.allclose(with_live_gate.float(), without_gate.float())
|
||||
|
||||
|
||||
def test_h3_forward_restores_rows_the_packed_padding_excludes(monkeypatch):
|
||||
_fake_block_sparse_kernel(monkeypatch, {})
|
||||
impl = _h3_impl(topk=64)
|
||||
prefix_segments, video_shape = (5,), (2, 4, 4)
|
||||
valid = 5 + 2 * 4 * 4
|
||||
torch.manual_seed(0)
|
||||
query = torch.randn(1, valid + 24, 2, 8, dtype=torch.bfloat16)
|
||||
padding = PackedPaddingMetadata(
|
||||
q_length=valid,
|
||||
kv_length=valid,
|
||||
cu_seqlens_q=torch.tensor([0, valid], dtype=torch.int32),
|
||||
cu_seqlens_k=torch.tensor([0, valid], dtype=torch.int32),
|
||||
)
|
||||
|
||||
output = impl.forward_cuda(query, query, query, _h3_metadata(prefix_segments, video_shape, packed_padding=padding))
|
||||
|
||||
assert output.shape == query.shape
|
||||
assert torch.count_nonzero(output[:, valid:]) == 0
|
||||
|
||||
|
||||
def test_h3_geometry_keeps_prefix_segments_pure_and_tiles_video_3d():
|
||||
partition, sizes, non_pad, untile, prefix_blocks, video_blocks = _get_h3_tile_metadata(
|
||||
(5, 70, 9), (5, 6, 6), torch.device("cpu")
|
||||
)
|
||||
assert sizes[:4].tolist() == [5, 64, 6, 9]
|
||||
assert prefix_blocks == 4
|
||||
assert video_blocks == 8
|
||||
assert int(sizes.sum()) == 5 + 70 + 9 + 5 * 6 * 6
|
||||
assert partition.numel() == non_pad.numel() == untile.numel() == int(sizes.sum())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fallback_on_error", [False, True])
|
||||
def test_h3_kernel_failure_preserves_fallback_policy(monkeypatch, fallback_on_error):
|
||||
impl = _h3_impl(fallback_on_error=fallback_on_error)
|
||||
query = torch.randn(1, 37, 2, 8)
|
||||
metadata = _h3_metadata((5,), (2, 4, 4))
|
||||
|
||||
def fail(*args, **kwargs):
|
||||
raise RuntimeError("kernel failure")
|
||||
|
||||
monkeypatch.setattr(impl, "_forward_h3", fail)
|
||||
monkeypatch.setattr(impl.sdpa_fallback, "forward", impl.sdpa_fallback.forward_cuda)
|
||||
if fallback_on_error:
|
||||
output = impl.forward_cuda(query, query, query, metadata)
|
||||
expected = impl.sdpa_fallback.forward_cuda(query, query, query, metadata)
|
||||
torch.testing.assert_close(output, expected, rtol=0, atol=0)
|
||||
else:
|
||||
with pytest.raises(RuntimeError, match="kernel failure"):
|
||||
impl.forward_cuda(query, query, query, metadata)
|
||||
|
||||
|
||||
def test_h3_accelerator_failure_is_not_swallowed(monkeypatch):
|
||||
impl = _h3_impl(fallback_on_error=True)
|
||||
query = torch.randn(1, 37, 2, 8)
|
||||
|
||||
def fail(*args, **kwargs):
|
||||
raise torch.AcceleratorError("device failure")
|
||||
|
||||
monkeypatch.setattr(impl, "_forward_h3", fail)
|
||||
with pytest.raises(torch.AcceleratorError, match="device failure"):
|
||||
impl.forward_cuda(query, query, query, _h3_metadata((5,), (2, 4, 4)))
|
||||
|
||||
|
||||
def test_h3_without_layout_delegates_to_shared_backend(monkeypatch):
|
||||
from vllm_omni.diffusion.attention.backends.fastvideo_vsa import FastVideoVSAImpl
|
||||
|
||||
impl = _h3_impl()
|
||||
query = torch.randn(1, 37, 2, 8)
|
||||
metadata = AttentionMetadata()
|
||||
expected = torch.empty_like(query)
|
||||
calls = []
|
||||
|
||||
def forward(self, q, k, v, attn_metadata):
|
||||
calls.append((self, q, k, v, attn_metadata))
|
||||
return expected
|
||||
|
||||
monkeypatch.setattr(FastVideoVSAImpl, "forward_cuda", forward)
|
||||
assert impl.forward_cuda(query, query, query, metadata) is expected
|
||||
assert len(calls) == 1
|
||||
assert all(actual is wanted for actual, wanted in zip(calls[0], (impl, query, query, query, metadata)))
|
||||
@@ -0,0 +1,87 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from vllm_omni.diffusion.models.qwen_image import qwen_image_transformer
|
||||
from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import QwenImageTransformerBlock
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
|
||||
class _StaticOutput(nn.Module):
|
||||
def __init__(self, output: torch.Tensor):
|
||||
super().__init__()
|
||||
self.output = output
|
||||
|
||||
def forward(self, _: torch.Tensor) -> torch.Tensor:
|
||||
return self.output
|
||||
|
||||
|
||||
class _RecordingNorm(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = []
|
||||
|
||||
def forward(self, x: torch.Tensor, scale: torch.Tensor, shift: torch.Tensor) -> torch.Tensor:
|
||||
self.calls.append((x, scale, shift))
|
||||
return x
|
||||
|
||||
|
||||
class _ZeroAttention(nn.Module):
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
encoder_hidden_states: torch.Tensor,
|
||||
**_: object,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
return torch.zeros_like(hidden_states), torch.zeros_like(encoder_hidden_states)
|
||||
|
||||
|
||||
class _ZeroMLP(nn.Module):
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return torch.zeros_like(x)
|
||||
|
||||
|
||||
def test_non_cuda_select01_uses_adalayernorm_dispatch(monkeypatch):
|
||||
hidden_size = 4
|
||||
block = QwenImageTransformerBlock.__new__(QwenImageTransformerBlock)
|
||||
nn.Module.__init__(block)
|
||||
block.zero_cond_t = True
|
||||
block.img_mod = _StaticOutput(torch.zeros(2, 6 * hidden_size))
|
||||
block.txt_mod = _StaticOutput(torch.zeros(1, 6 * hidden_size))
|
||||
block.img_norm1 = _RecordingNorm()
|
||||
block.img_norm2 = _RecordingNorm()
|
||||
block.txt_norm1 = _RecordingNorm()
|
||||
block.txt_norm2 = _RecordingNorm()
|
||||
block.attn = _ZeroAttention()
|
||||
block.img_mlp = _ZeroMLP()
|
||||
block.txt_mlp = _ZeroMLP()
|
||||
|
||||
monkeypatch.setattr(qwen_image_transformer, "can_use_qwen_select01_triton", lambda _: False)
|
||||
|
||||
def fail_if_called(*args, **kwargs):
|
||||
pytest.fail("non-CUDA select01 must preserve the AdaLayerNorm dispatch")
|
||||
|
||||
monkeypatch.setattr(qwen_image_transformer, "fused_layernorm_select01", fail_if_called)
|
||||
monkeypatch.setattr(qwen_image_transformer, "fused_residual_layernorm_select01", fail_if_called)
|
||||
|
||||
hidden_states = torch.randn(1, 3, hidden_size)
|
||||
encoder_hidden_states = torch.randn(1, 2, hidden_size)
|
||||
modulate_index = torch.tensor([[0, 1, 0]])
|
||||
block(
|
||||
hidden_states,
|
||||
encoder_hidden_states,
|
||||
torch.ones(1, 2),
|
||||
torch.zeros(2, hidden_size),
|
||||
(torch.empty(0), torch.empty(0)),
|
||||
modulate_index=modulate_index,
|
||||
)
|
||||
|
||||
assert len(block.img_norm1.calls) == 1
|
||||
assert len(block.img_norm2.calls) == 1
|
||||
for _, scale, shift in (*block.img_norm1.calls, *block.img_norm2.calls):
|
||||
assert scale.shape == hidden_states.shape
|
||||
assert shift.shape == hidden_states.shape
|
||||
@@ -1,3 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from contextlib import contextmanager
|
||||
@@ -33,7 +36,7 @@ class StubScheduler:
|
||||
self.config = SimpleNamespace(num_train_timesteps=1000)
|
||||
self.set_timesteps_calls: list[tuple[int, torch.device]] = []
|
||||
|
||||
def set_timesteps(self, num_steps: int, device: torch.device) -> None:
|
||||
def set_timesteps(self, num_steps: int, device: torch.device, shift: float | None = None) -> None:
|
||||
self.set_timesteps_calls.append((num_steps, device))
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,97 @@
|
||||
{
|
||||
"source_url": "https://github.com/Wan-Video/Wan2.2/blob/42bf4cfaa384bc21833865abc2f9e6c0e67233dc/wan/utils/fm_solvers_unipc.py",
|
||||
"source_sha256": "0dec8c7ed17f6f2049275c6848113314da6ccec1c8db5bdc89df43c05c6038d9",
|
||||
"git_head": "42bf4cfaa384bc21833865abc2f9e6c0e67233dc",
|
||||
"num_train_timesteps": 1000,
|
||||
"use_dynamic_shifting": false,
|
||||
"reproduction": "Load FlowUniPCMultistepScheduler from source_url; construct with num_train_timesteps=1000, shift=1.0, use_dynamic_shifting=False; call set_timesteps(40, device='cpu', shift=12.0); record timesteps.tolist() and sigmas.tolist().",
|
||||
"constructor_shift": 1.0,
|
||||
"runtime_shift": 12.0,
|
||||
"steps": 40,
|
||||
"timesteps": [
|
||||
999,
|
||||
997,
|
||||
995,
|
||||
993,
|
||||
990,
|
||||
988,
|
||||
985,
|
||||
982,
|
||||
979,
|
||||
976,
|
||||
972,
|
||||
969,
|
||||
965,
|
||||
961,
|
||||
956,
|
||||
952,
|
||||
947,
|
||||
941,
|
||||
936,
|
||||
929,
|
||||
922,
|
||||
915,
|
||||
907,
|
||||
898,
|
||||
888,
|
||||
877,
|
||||
865,
|
||||
852,
|
||||
837,
|
||||
819,
|
||||
799,
|
||||
776,
|
||||
749,
|
||||
717,
|
||||
678,
|
||||
631,
|
||||
571,
|
||||
492,
|
||||
386,
|
||||
235
|
||||
],
|
||||
"sigmas": [
|
||||
0.9999166131019592,
|
||||
0.9977826476097107,
|
||||
0.9955461621284485,
|
||||
0.9931996464729309,
|
||||
0.9907346963882446,
|
||||
0.988142192363739,
|
||||
0.9854119420051575,
|
||||
0.9825326800346375,
|
||||
0.9794917702674866,
|
||||
0.9762753844261169,
|
||||
0.9728676676750183,
|
||||
0.9692512154579163,
|
||||
0.9654061794281006,
|
||||
0.961310088634491,
|
||||
0.9569376707077026,
|
||||
0.95225989818573,
|
||||
0.9472436308860779,
|
||||
0.9418508410453796,
|
||||
0.9360373020172119,
|
||||
0.9297519326210022,
|
||||
0.9229347705841064,
|
||||
0.9155154228210449,
|
||||
0.9074103832244873,
|
||||
0.8985199332237244,
|
||||
0.888724148273468,
|
||||
0.8778772950172424,
|
||||
0.865800678730011,
|
||||
0.8522725701332092,
|
||||
0.8370144367218018,
|
||||
0.8196719288825989,
|
||||
0.7997865080833435,
|
||||
0.7767546772956848,
|
||||
0.7497654557228088,
|
||||
0.717703104019165,
|
||||
0.6789888143539429,
|
||||
0.6313128471374512,
|
||||
0.5711563229560852,
|
||||
0.49288034439086914,
|
||||
0.38684695959091187,
|
||||
0.23510953783988953,
|
||||
0.0
|
||||
],
|
||||
"implementation": "Wan-Video/Wan2.2 wan/utils/fm_solvers_unipc.py"
|
||||
}
|
||||
@@ -1,15 +1,16 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from PIL import Image
|
||||
from torch import nn
|
||||
|
||||
from tests.diffusion.models.wan2_2.conftest import StubScheduler, StubTransformer, StubVAE, noop_progress_bar
|
||||
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import _WAN_TEXT_ENCODER_OFFLOAD_PLAN
|
||||
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import _WAN_TEXT_ENCODER_OFFLOAD_PLAN, build_wan_scheduler
|
||||
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v import (
|
||||
Wan22I2VPipeline,
|
||||
get_wan22_i2v_post_process_func,
|
||||
@@ -50,7 +51,7 @@ def _make_i2v_pipeline(*, expand_timesteps: bool) -> Wan22I2VPipeline:
|
||||
|
||||
|
||||
def _make_i2v_sampling(**overrides):
|
||||
values = {
|
||||
values: dict[str, object] = {
|
||||
"height": 16,
|
||||
"width": 16,
|
||||
"num_frames": 5,
|
||||
@@ -275,16 +276,18 @@ def test_i2v_prepare_latents_preserves_batched_image_conditions() -> None:
|
||||
assert first_frame_mask.shape == (2, 1, 2, 2, 2)
|
||||
|
||||
|
||||
def test_i2v_forward_batches_conditions_random_inputs_and_outputs(monkeypatch) -> None:
|
||||
@pytest.mark.parametrize("solver", ["unipc", "euler"])
|
||||
@pytest.mark.parametrize("shift", [3.0, 5.0, 12.0])
|
||||
def test_i2v_forward_batches_conditions_random_inputs_and_outputs(monkeypatch, solver: str, shift: float) -> None:
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v.current_omni_platform",
|
||||
SimpleNamespace(is_available=lambda: False),
|
||||
)
|
||||
pipeline = _make_i2v_pipeline(expand_timesteps=True)
|
||||
pipeline.scheduler = StubScheduler([9])
|
||||
pipeline.od_config = SimpleNamespace(flow_shift=5.0)
|
||||
pipeline._sample_solver = "unipc"
|
||||
pipeline._flow_shift = 5.0
|
||||
pipeline.scheduler = build_wan_scheduler(solver, shift)
|
||||
pipeline.od_config = SimpleNamespace(flow_shift=shift)
|
||||
pipeline._sample_solver = solver
|
||||
pipeline._flow_shift = shift
|
||||
pipeline.boundary_ratio = 0.875
|
||||
pipeline.has_image_encoder = False
|
||||
pipeline._guidance_scale = None
|
||||
@@ -327,6 +330,8 @@ def test_i2v_forward_batches_conditions_random_inputs_and_outputs(monkeypatch) -
|
||||
generator=gen_a,
|
||||
latents=latents_a,
|
||||
num_outputs_per_prompt=2,
|
||||
num_inference_steps=50,
|
||||
extra_args={"sample_solver": solver, "flow_shift": shift},
|
||||
),
|
||||
),
|
||||
SimpleNamespace(
|
||||
@@ -336,12 +341,25 @@ def test_i2v_forward_batches_conditions_random_inputs_and_outputs(monkeypatch) -
|
||||
generator=gen_b,
|
||||
latents=latents_b,
|
||||
num_outputs_per_prompt=2,
|
||||
num_inference_steps=50,
|
||||
extra_args={"sample_solver": solver, "flow_shift": shift},
|
||||
),
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
outputs = pipeline.forward(batch)
|
||||
if solver == "unipc":
|
||||
sigmas = np.linspace(float(np.float32(0.999)), 0.0, 51)[:-1]
|
||||
sigmas = shift * sigmas / (1.0 + (shift - 1.0) * sigmas)
|
||||
torch.testing.assert_close(
|
||||
pipeline.scheduler.sigmas, torch.tensor(np.append(sigmas, 0.0), dtype=torch.float32), rtol=0, atol=0
|
||||
)
|
||||
assert pipeline.scheduler.timesteps.tolist() == (sigmas * 1000).astype(np.int64).tolist()
|
||||
else:
|
||||
reference = build_wan_scheduler("euler", shift)
|
||||
reference.set_timesteps(50, device="cpu")
|
||||
torch.testing.assert_close(pipeline.scheduler.sigmas, reference.sigmas, rtol=0, atol=0)
|
||||
|
||||
assert prepare_call["batch_size"] == 4
|
||||
assert prepare_call["generator"] == [gen_a, gen_a, gen_b, gen_b]
|
||||
|
||||
@@ -10,7 +10,7 @@ import torch
|
||||
from torch import nn
|
||||
|
||||
from vllm_omni.diffusion.media import VideoTensorEncoding, VideoTensorLayout, VideoValueRange
|
||||
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import Wan22Pipeline
|
||||
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import Wan22Pipeline, build_wan_scheduler
|
||||
from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanSelfAttention
|
||||
from vllm_omni.diffusion.request import OmniDiffusionRequest
|
||||
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
|
||||
@@ -19,6 +19,12 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _cpu_platform(monkeypatch):
|
||||
module = importlib.import_module("vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2")
|
||||
monkeypatch.setattr(module.current_omni_platform, "is_available", lambda: False)
|
||||
|
||||
|
||||
class _StubTransformer(nn.Module):
|
||||
@property
|
||||
def dtype(self) -> torch.dtype:
|
||||
@@ -51,10 +57,10 @@ class _StubScheduler:
|
||||
def __init__(self, timesteps: list[int]) -> None:
|
||||
self.timesteps = torch.tensor(timesteps, dtype=torch.int64)
|
||||
self.config = SimpleNamespace(num_train_timesteps=1000)
|
||||
self.set_timesteps_calls: list[tuple[int, torch.device]] = []
|
||||
self.set_timesteps_calls: list[tuple[int, torch.device, float | None]] = []
|
||||
|
||||
def set_timesteps(self, num_steps: int, device: torch.device) -> None:
|
||||
self.set_timesteps_calls.append((num_steps, device))
|
||||
def set_timesteps(self, num_steps: int, device: torch.device, shift: float | None = None) -> None:
|
||||
self.set_timesteps_calls.append((num_steps, device, shift))
|
||||
|
||||
|
||||
@contextmanager
|
||||
@@ -184,7 +190,24 @@ def test_forward_delegates_denoising_to_diffuse(
|
||||
assert captured["boundary_timestep"] == pytest.approx(875.0)
|
||||
assert captured["latent_condition"] is None
|
||||
assert captured["first_frame_mask"] is None
|
||||
assert pipeline.scheduler.set_timesteps_calls == [(2, torch.device("cpu"))]
|
||||
assert pipeline.scheduler.set_timesteps_calls == [(2, torch.device("cpu"), 5.0)]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("solver", ["unipc", "euler"])
|
||||
def test_forward_passes_request_shift_without_mutating_scheduler_config(solver: str) -> None:
|
||||
pipeline = _make_pipeline()
|
||||
pipeline.diffuse = lambda **kwargs: kwargs["latents"]
|
||||
for shift in (3.0, 12.0, 5.0):
|
||||
sampling = _make_sampling(num_inference_steps=5, extra_args={"sample_solver": solver, "flow_shift": shift})
|
||||
request = OmniDiffusionRequest(prompt="prompt", request_id="schedule", sampling_params=sampling)
|
||||
pipeline.forward(DiffusionRequestBatch(requests=[request]))
|
||||
reference = build_wan_scheduler(solver, shift)
|
||||
if solver == "unipc":
|
||||
reference.set_timesteps(5, device="cpu", shift=shift)
|
||||
assert pipeline.scheduler.config.shift == pipeline.scheduler.config["shift"] == 1.0
|
||||
else:
|
||||
reference.set_timesteps(5, device="cpu")
|
||||
torch.testing.assert_close(pipeline.scheduler.sigmas, reference.sigmas, rtol=0, atol=0)
|
||||
|
||||
|
||||
def test_forward_batches_text_generators_latents_and_splits_outputs() -> None:
|
||||
|
||||
@@ -1,3 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
@@ -7,6 +10,7 @@ import pytest
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler
|
||||
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v import (
|
||||
Wan22S2VPipeline,
|
||||
_make_clip_generators,
|
||||
@@ -14,11 +18,32 @@ from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v import (
|
||||
from vllm_omni.diffusion.models.wan2_2.wan2_2_s2v_transformer import WanS2VTransformer3DModel
|
||||
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("configured_shift", [None, 12.0])
|
||||
def test_s2v_constructor_preserves_unshifted_endpoints(configured_shift: float | None) -> None:
|
||||
module = "vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_s2v"
|
||||
config = SimpleNamespace(model="unused", flow_shift=configured_shift, enable_diffusion_pipeline_profiler=False)
|
||||
|
||||
def init_components(pipeline, *args):
|
||||
pipeline.vae = SimpleNamespace(config=SimpleNamespace(scale_factor_temporal=4, scale_factor_spatial=8))
|
||||
|
||||
with (
|
||||
patch(f"{module}.get_local_device", return_value=torch.device("cpu")),
|
||||
patch(f"{module}._resolve_model_path", return_value="unused"),
|
||||
patch(f"{module}._is_diffusers_format", return_value=True),
|
||||
patch.object(Wan22S2VPipeline, "_init_diffusers_format", init_components),
|
||||
patch.object(Wan22S2VPipeline, "setup_diffusion_pipeline_profiler"),
|
||||
):
|
||||
pipeline = Wan22S2VPipeline(od_config=config)
|
||||
assert pipeline._flow_shift == (3.0 if configured_shift is None else configured_shift)
|
||||
assert pipeline.scheduler.config.shift == pipeline.scheduler.config["shift"] == 1.0
|
||||
assert pipeline.scheduler.sigma_max == float(np.float32(0.999))
|
||||
|
||||
|
||||
def _make_s2v_sampling(**overrides):
|
||||
values = {
|
||||
values: dict[str, object] = {
|
||||
"height": 16,
|
||||
"width": 16,
|
||||
"num_frames": 8,
|
||||
@@ -211,7 +236,8 @@ def test_s2v_forward_rejects_different_raw_audio_shapes() -> None:
|
||||
pipeline.forward(batch)
|
||||
|
||||
|
||||
def test_s2v_forward_batches_request_local_inputs_and_splits_outputs() -> None:
|
||||
@pytest.mark.parametrize("shift", [3.0, 12.0])
|
||||
def test_s2v_forward_batches_request_local_inputs_and_splits_outputs(shift: float) -> None:
|
||||
pipeline = object.__new__(Wan22S2VPipeline)
|
||||
nn.Module.__init__(pipeline)
|
||||
pipeline.device = torch.device("cpu")
|
||||
@@ -227,7 +253,8 @@ def test_s2v_forward_batches_request_local_inputs_and_splits_outputs() -> None:
|
||||
enable_cpu_offload=False,
|
||||
parallel_config=SimpleNamespace(use_hsdp=False),
|
||||
)
|
||||
pipeline.scheduler = MagicMock(timesteps=torch.tensor([1.0]))
|
||||
pipeline._flow_shift = shift
|
||||
pipeline.scheduler = FlowUniPCMultistepScheduler(shift=1.0)
|
||||
pipeline.vae_scale_factor_spatial = 8
|
||||
pipeline.resolution_divisor = 16
|
||||
pipeline.motion_frames = 7
|
||||
@@ -262,6 +289,7 @@ def test_s2v_forward_batches_request_local_inputs_and_splits_outputs() -> None:
|
||||
"multi_modal_data": {"image": image, "audio": audio_a},
|
||||
},
|
||||
sampling_params=_make_s2v_sampling(
|
||||
num_inference_steps=5,
|
||||
num_outputs_per_prompt=2,
|
||||
generator=generators_a,
|
||||
latents=latents_a,
|
||||
@@ -275,6 +303,7 @@ def test_s2v_forward_batches_request_local_inputs_and_splits_outputs() -> None:
|
||||
"multi_modal_data": {"image": image, "audio": audio_b},
|
||||
},
|
||||
sampling_params=_make_s2v_sampling(
|
||||
num_inference_steps=5,
|
||||
num_outputs_per_prompt=2,
|
||||
generator=generators_b,
|
||||
latents=latents_b,
|
||||
@@ -287,6 +316,12 @@ def test_s2v_forward_batches_request_local_inputs_and_splits_outputs() -> None:
|
||||
platform.is_available.return_value = False
|
||||
outputs = pipeline.forward(batch)
|
||||
|
||||
expected = np.linspace(float(np.float32(0.999)), 0.0, 6)[:-1]
|
||||
expected = shift * expected / (1.0 + (shift - 1.0) * expected)
|
||||
torch.testing.assert_close(
|
||||
pipeline.scheduler.sigmas, torch.tensor(np.append(expected, 0.0), dtype=torch.float32), rtol=0, atol=0
|
||||
)
|
||||
assert pipeline.scheduler.config.shift == pipeline.scheduler.config["shift"] == 1.0
|
||||
assert len(outputs) == 2
|
||||
assert outputs[0].output[0].shape[0] == 2
|
||||
assert outputs[1].output[0].shape[0] == 2
|
||||
@@ -519,6 +554,7 @@ def test_s2v_pipeline_hsdp_forward_complete_process():
|
||||
mock_scheduler.timesteps = torch.linspace(999, 0, 5)
|
||||
mock_scheduler.step = MagicMock(return_value=(torch.zeros(1, 16, 20, 88, 128),))
|
||||
pipeline.scheduler = mock_scheduler
|
||||
pipeline._flow_shift = 3.0
|
||||
|
||||
# -- Bind methods from the real class --
|
||||
pipeline.encode_prompt = Wan22S2VPipeline.encode_prompt.__get__(pipeline)
|
||||
|
||||
@@ -0,0 +1,54 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Compare actual Wan schedules to native Wan outputs and independent sigma arithmetic."""
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import numpy.typing as npt
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
|
||||
FASTWAN_DMD_SCHEDULER_SHIFT,
|
||||
build_wan_scheduler,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
|
||||
def test_wan_unipc_matches_captured_native_full_schedule() -> None:
|
||||
# The fixture pins the official Wan source and records how to regenerate it.
|
||||
golden = json.loads((Path(__file__).parent / "fixtures/native_unipc_shift12.json").read_text())
|
||||
scheduler = build_wan_scheduler("unipc", golden["runtime_shift"])
|
||||
for steps in (golden["steps"], 20, golden["steps"]):
|
||||
scheduler.set_timesteps(steps, device="cpu", shift=golden["runtime_shift"])
|
||||
if steps == golden["steps"]:
|
||||
assert scheduler.timesteps.tolist() == golden["timesteps"]
|
||||
torch.testing.assert_close(scheduler.sigmas, torch.tensor(golden["sigmas"]), rtol=0, atol=0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shift", [1.0, 3.0, 5.0, 12.0])
|
||||
@pytest.mark.parametrize("steps", [20, 40])
|
||||
def test_wan_unipc_shifts_unmodified_training_endpoints(shift: float, steps: int) -> None:
|
||||
# Native Wan: float32 training sigma_max=(1-1/1000), sigma_min=0;
|
||||
# interpolate in float64, apply the runtime shift, then cast sigma to FP32.
|
||||
unshifted = np.linspace(float(np.float32(0.999)), 0.0, steps + 1)[:-1]
|
||||
shifted = shift * unshifted / (1.0 + (shift - 1.0) * unshifted)
|
||||
expected_timesteps = (shifted * 1000).astype(np.int64)
|
||||
expected_sigmas: npt.NDArray[np.float32] = np.append(shifted, 0.0).astype(np.float32)
|
||||
scheduler = build_wan_scheduler("unipc", shift)
|
||||
scheduler.set_timesteps(steps, device="cpu", shift=shift)
|
||||
assert scheduler.config.shift == scheduler.config["shift"] == 1.0
|
||||
torch.testing.assert_close(scheduler.timesteps, torch.from_numpy(expected_timesteps), rtol=0, atol=0)
|
||||
torch.testing.assert_close(scheduler.sigmas, torch.from_numpy(expected_sigmas), rtol=0, atol=0)
|
||||
|
||||
|
||||
def test_wan_euler_keeps_dmd_shift_semantics() -> None:
|
||||
scheduler = build_wan_scheduler("euler", FASTWAN_DMD_SCHEDULER_SHIFT)
|
||||
scheduler.set_timesteps(3, device="cpu")
|
||||
# The unchanged Euler branch starts at sigma=1, with shift=8 for DMD.
|
||||
expected = torch.tensor([1.0, 16.0 / 17.0, 0.8, 0.0])
|
||||
torch.testing.assert_close(scheduler.sigmas, expected, rtol=0, atol=1e-7)
|
||||
torch.testing.assert_close(scheduler.timesteps, expected[:-1] * 1000, rtol=0, atol=1e-4)
|
||||
@@ -1,15 +1,17 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from types import SimpleNamespace
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from PIL import Image
|
||||
from torch import nn
|
||||
|
||||
from tests.diffusion.models.wan2_2.conftest import StubScheduler, StubTransformer, StubVAE, noop_progress_bar
|
||||
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import build_wan_scheduler
|
||||
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace import (
|
||||
Wan22VACEPipeline,
|
||||
create_vace_transformer_from_config,
|
||||
@@ -33,6 +35,8 @@ def test_vace_default_flow_shift_preserves_compact_offload_runtime_state(monkeyp
|
||||
)
|
||||
enable_layerwise_offload: bool = True
|
||||
post_init_calls: int = 0
|
||||
_resolved_diffusion_offload: object = None
|
||||
_diffusion_offload_flags_materialized: bool = False
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
self.post_init_calls += 1
|
||||
@@ -41,7 +45,7 @@ def test_vace_default_flow_shift_preserves_compact_offload_runtime_state(monkeyp
|
||||
resolved_offload = object()
|
||||
config._resolved_diffusion_offload = resolved_offload
|
||||
config._diffusion_offload_flags_materialized = True
|
||||
captured = {}
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
def capture_base_init(self, *, od_config, prefix="") -> None:
|
||||
captured.update(od_config=od_config, prefix=prefix)
|
||||
@@ -51,6 +55,7 @@ def test_vace_default_flow_shift_preserves_compact_offload_runtime_state(monkeyp
|
||||
Wan22VACEPipeline(od_config=config, prefix="stage")
|
||||
|
||||
forwarded = captured["od_config"]
|
||||
assert isinstance(forwarded, RuntimeConfig)
|
||||
assert forwarded is not config
|
||||
assert config.flow_shift is None
|
||||
assert forwarded.flow_shift == 3.0
|
||||
@@ -75,7 +80,7 @@ def _make_vace_pipeline() -> Wan22VACEPipeline:
|
||||
|
||||
|
||||
def _make_vace_sampling(**overrides):
|
||||
values = {
|
||||
values: dict[str, object] = {
|
||||
"height": 16,
|
||||
"width": 16,
|
||||
"num_frames": 5,
|
||||
@@ -123,7 +128,7 @@ def test_vace_preprocess_collects_reference_video_and_mask_inputs() -> None:
|
||||
|
||||
|
||||
def test_create_vace_transformer_from_config_maps_vace_specific_keys(monkeypatch) -> None:
|
||||
captured = {}
|
||||
captured: dict[str, object] = {}
|
||||
|
||||
class FakeVACETransformer:
|
||||
def __init__(self, **kwargs) -> None:
|
||||
@@ -203,7 +208,9 @@ def test_vace_diffuse_passes_context_and_scale_to_cfg_branches() -> None:
|
||||
torch.testing.assert_close(result, torch.ones_like(latents))
|
||||
|
||||
|
||||
def test_vace_forward_batches_random_inputs_and_splits_outputs(monkeypatch) -> None:
|
||||
@pytest.mark.parametrize("solver", ["unipc", "euler"])
|
||||
@pytest.mark.parametrize("shift", [3.0, 5.0, 12.0])
|
||||
def test_vace_forward_batches_random_inputs_and_splits_outputs(monkeypatch, solver: str, shift: float) -> None:
|
||||
monkeypatch.setattr(
|
||||
"vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace.current_omni_platform",
|
||||
SimpleNamespace(is_available=lambda: False),
|
||||
@@ -211,10 +218,10 @@ def test_vace_forward_batches_random_inputs_and_splits_outputs(monkeypatch) -> N
|
||||
pipeline = _make_vace_pipeline()
|
||||
pipeline.transformer.vace_patch_embedding = None
|
||||
pipeline.transformer_2 = None
|
||||
pipeline.scheduler = StubScheduler([9])
|
||||
pipeline.od_config = SimpleNamespace(flow_shift=3.0)
|
||||
pipeline._sample_solver = "unipc"
|
||||
pipeline._flow_shift = 3.0
|
||||
pipeline.scheduler = build_wan_scheduler(solver, shift)
|
||||
pipeline.od_config = SimpleNamespace(flow_shift=shift)
|
||||
pipeline._sample_solver = solver
|
||||
pipeline._flow_shift = shift
|
||||
pipeline.boundary_ratio = None
|
||||
pipeline._guidance_scale = None
|
||||
pipeline._num_timesteps = None
|
||||
@@ -247,6 +254,8 @@ def test_vace_forward_batches_random_inputs_and_splits_outputs(monkeypatch) -> N
|
||||
generator=gen_a,
|
||||
latents=latents_a,
|
||||
num_outputs_per_prompt=2,
|
||||
num_inference_steps=50,
|
||||
extra_args={"sample_solver": solver, "flow_shift": shift},
|
||||
),
|
||||
),
|
||||
SimpleNamespace(
|
||||
@@ -256,12 +265,25 @@ def test_vace_forward_batches_random_inputs_and_splits_outputs(monkeypatch) -> N
|
||||
generator=gen_b,
|
||||
latents=latents_b,
|
||||
num_outputs_per_prompt=2,
|
||||
num_inference_steps=50,
|
||||
extra_args={"sample_solver": solver, "flow_shift": shift},
|
||||
),
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
outputs = pipeline.forward(batch)
|
||||
if solver == "unipc":
|
||||
sigmas = np.linspace(float(np.float32(0.999)), 0.0, 51)[:-1]
|
||||
sigmas = shift * sigmas / (1.0 + (shift - 1.0) * sigmas)
|
||||
torch.testing.assert_close(
|
||||
pipeline.scheduler.sigmas, torch.tensor(np.append(sigmas, 0.0), dtype=torch.float32), rtol=0, atol=0
|
||||
)
|
||||
assert pipeline.scheduler.timesteps.tolist() == (sigmas * 1000).astype(np.int64).tolist()
|
||||
else:
|
||||
reference = build_wan_scheduler("euler", shift)
|
||||
reference.set_timesteps(50, device="cpu")
|
||||
torch.testing.assert_close(pipeline.scheduler.sigmas, reference.sigmas, rtol=0, atol=0)
|
||||
|
||||
assert prepare_call["batch_size"] == 4
|
||||
assert prepare_call["generator"] == [gen_a, gen_a, gen_b, gen_b]
|
||||
|
||||
@@ -39,7 +39,8 @@ from vllm_omni.distributed.omni_connectors.utils.serialization import (
|
||||
from vllm_omni.engine.async_omni_engine import StageRuntimeInfo
|
||||
from vllm_omni.engine.messages import ShutdownRequestMessage, StageSubmissionMessage
|
||||
from vllm_omni.engine.stage_init_utils import StageMetadata
|
||||
from vllm_omni.entrypoints.async_omni import AsyncEventResolver, AsyncOmni
|
||||
from vllm_omni.entrypoints.async_omni import AsyncOmni
|
||||
from vllm_omni.entrypoints.async_omni_base import AsyncEventResolver
|
||||
from vllm_omni.entrypoints.openai.serving_video_output_stream import OmniStreamingVideoOutputHandler
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
|
||||
@@ -42,7 +42,8 @@ from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState
|
||||
from vllm_omni.engine.stage_client import StagePoolClient
|
||||
from vllm_omni.engine.stage_init_utils import StageMetadata
|
||||
from vllm_omni.engine.stage_pool import StagePool
|
||||
from vllm_omni.entrypoints.async_omni import AsyncEventResolver, AsyncOmni
|
||||
from vllm_omni.entrypoints.async_omni import AsyncOmni
|
||||
from vllm_omni.entrypoints.async_omni_base import AsyncEventResolver
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniInteractionPrompt
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
|
||||
|
||||
Loaded 100 of 371 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user