diff --git a/docs/design/ar_diffusion_pipeline_capability.md b/docs/design/ar_diffusion_pipeline_capability.md index 3c1b4b9cd..f96c80b11 100644 --- a/docs/design/ar_diffusion_pipeline_capability.md +++ b/docs/design/ar_diffusion_pipeline_capability.md @@ -22,11 +22,15 @@ A pipeline opts in by implementing `SupportsARDiffusionPipeline`. Its immutable - named fixed-length cross-attention KV, model-specific scratch geometry, and the pipeline's requested session-capacity upper bound. -The runner binds an `ARDiffusionKVState` only around one `forward()` call. The -pipeline may use it during that context but must not retain it. The session -survives subsequent requests with the same `session_id`. Request reset, explicit -close, LRU eviction, and forward exceptions all release runner-owned KV and call -the pipeline's reset/close lifecycle hook. Cross-attention KV is allocated lazily +The runner binds an `ARDiffusionKVState` only around one runner invocation +(`execute_model` or `execute_stepwise`). The pipeline may use it during that +context but must not retain it. Scratch and committed pages stay on the +runner-owned session object, so unbinding does not drop an in-flight stepwise +chunk. Tick-mode sessions survive subsequent requests with the same +`session_id`. On the stepwise path the session id is the request id and is +released when that request finishes or fails. Request reset, explicit close, +LRU eviction, and forward exceptions all release runner-owned KV and call the +pipeline's reset/close lifecycle hook. Cross-attention KV is allocated lazily per session and logical KV branch, and is released on the same path. A pipeline publishes one named cache transactionally by yielding exactly one `(k, v)` pair per layer: @@ -48,11 +52,12 @@ successful forward. The runner also exposes `reset_session()` and ## Current execution limits -The experimental runner supports request-mode execution with -`max_num_seqs=1`. It rejects step execution and request-batch execution at load -time, and its inherited batch/step entry points also fail explicitly. Supporting -those paths requires a batch-aware state-binding contract rather than silently -bypassing `bind_ar_diffusion_state()`. +The experimental runner supports `max_num_seqs=1`. Request-batch execution is +still rejected. Step execution is allowed only when the loaded pipeline +implements `SupportsStepExecution`; other AR-Diffusion pipelines still fail at +load if `step_execution=True`. Capable pipelines bind KV around each +`execute_stepwise` invocation instead of inheriting the generic runner entry +point unwrapped. The implementation may retain multiple sessions per model instance. It selects the largest feasible capacity up to the pipeline's `session_capacity`, based on diff --git a/examples/offline_inference/diffusion/lingbot_world_v2.py b/examples/offline_inference/diffusion/lingbot_world_v2.py index b3a4ae0d1..590442c4d 100644 --- a/examples/offline_inference/diffusion/lingbot_world_v2.py +++ b/examples/offline_inference/diffusion/lingbot_world_v2.py @@ -30,7 +30,6 @@ from typing import Any import numpy as np _MODEL = "robbyant/lingbot-world-v2-14b-causal-fast-diffusers" -_NUM_INFERENCE_STEPS = 4 _MAX_SEQUENCE_LENGTH = 512 _MAX_PIXEL_AREA = 480 * 832 _MAX_RAW_FRAMES = 117 @@ -209,7 +208,6 @@ def build_request( "height": args.height, "width": args.width, "num_frames": args.num_frames, - "num_inference_steps": _NUM_INFERENCE_STEPS, "max_sequence_length": _MAX_SEQUENCE_LENGTH, "seed": args.seed, "fps": args.fps, diff --git a/examples/offline_inference/diffusion/lingbot_world_v2_realtime.py b/examples/offline_inference/diffusion/lingbot_world_v2_realtime.py index 86461f9f8..d98fe06f9 100644 --- a/examples/offline_inference/diffusion/lingbot_world_v2_realtime.py +++ b/examples/offline_inference/diffusion/lingbot_world_v2_realtime.py @@ -155,7 +155,6 @@ async def run(argv: Sequence[str] | None = None) -> Path: height=args.height, width=args.width, num_frames=9, - num_inference_steps=4, max_sequence_length=512, seed=args.seed, output_type="latent", diff --git a/examples/online_serving/streaming_video_generation/README.md b/examples/online_serving/streaming_video_generation/README.md index f2d465797..8e535ce0e 100644 --- a/examples/online_serving/streaming_video_generation/README.md +++ b/examples/online_serving/streaming_video_generation/README.md @@ -1,7 +1,7 @@ # Streaming Video Generation This example uses the custom WebSocket endpoint `WS /v1/realtime/video` to receive a video byte stream as chunks are produced. -It covers text-only video generation. Image/reference input is intentionally not included for now. +It covers text-to-video, and image-conditioned models through `--image-reference`. ## Start The Server @@ -20,7 +20,7 @@ The `--diffusion-streaming-output` CLI flag is forwarded as `streaming_output=Tr | Direction | Message | Format | Description | | --- | --- | --- | --- | -| Client to server | `session.start` | JSON text: `{"type":"session.start","model":"...","prompt":"...","format":"m4s"}` | Starts generation. `format` is optional and accepts `m4s` (default). Sampling fields such as `width`, `height`, `fps`, `num_frames`, and `extra_params` may be included. | +| Client to server | `session.start` | JSON text: `{"type":"session.start","model":"...","prompt":"...","format":"m4s"}` | Starts generation. `format` is optional and accepts `m4s` (default). Sampling fields such as `width`, `height`, `fps`, `num_frames`, and `extra_params` may be included. Image-conditioned models take a first frame as `image_reference`: `{"image_url":"https://..."}` or a `data:` URL. | | Client to server | `session.interaction` | JSON text: `{"type":"session.interaction","interaction":{"event_id":"xxx","event":{"prompt":"..."},"transition_chunks":3}}` | Updates the active prompt midway through generation. `event_id` is optional and `transition_chunks` defaults to the model setting. | | Server to client | `video.start` | JSON text: `{"type":"video.start","request_id":"...","format":"m4s","config":{...}}` | Confirms the session and mirrors the accepted `format`. | | Server to client | `video.chunk_metadata` | JSON text: `{"type":"video.chunk_metadata","request_id":"...","kind":"media","transport_chunk_index":0,"generation_chunk_index":0,"num_frames":9,"byte_length":1234,"started_event_ids":[],"active_event_ids":[],"completed_event_ids":[]}` | Precedes each binary frame and describes the immediately following payload. | @@ -62,7 +62,7 @@ python streaming_video_client.py \ The client sends one `session.start` message, prints each received binary video chunk with its byte size and elapsed time, and saves the received bytes to `--output` after `session.done`. The client remuxes the gathered stream to a regular progressive MP4 file so that local playback knows the video duration. -Schedule midway prompt updates with `--prompt-updates`. Each entry uses `"at"` as seconds on the client clock after the server sends `video.start`: +Schedule midway prompt updates with `--prompt-updates`. None are sent unless you ask for them, and only pipelines that implement midway prompt updates accept them (LingBot-World, for one, does not, and rejects the whole session). Each entry uses `"at"` as seconds on the client clock after the server sends `video.start`: ```bash python streaming_video_client.py \ @@ -73,6 +73,19 @@ python streaming_video_client.py \ ]' ``` +The Helios demo continues the default prompt into an underwater storm: + +```bash +python streaming_video_client.py \ + --model BestWishYsh/Helios-Distilled \ + --prompt-updates '[ + {"at": 4, + "prompt": "An underwater tornado appears and affects the ocean floor in a dramatic and chaotic scene. The water is murky, swirling violently, carrying debris and marine life into the vortex. The tropical fish on the scene all swim in panic, trying to avoid the powerful currents. The camera remains stationary, capturing the intensity of the underwater tornado as it disrupts the serene ocean floor. Close-up shot emphasizing the turbulent motion and destruction."}, + {"at": 11, + "prompt": "The swirling underwater vortex now seizes a heavy, encrusted treasure chest, its lid flapping open as it is smashed onto the ocean floor. Gold coins and silver trinkets spill out, glittering briefly in the murky water before being swept instantly into the violent funnel. The heavy wooden box tumbles end over end, colliding with floating rocks and adding to the debris field. Swirling sediment and bubbles surround the spilling fortune, highlighting the chaotic power of the storm as it ravages the seabed. Close-up shot emphasizing the turbulent motion and destruction."} + ]' +``` + ## Run The Gradio Demo ```bash @@ -106,3 +119,31 @@ Disable it in the CLI example with `--no-helios-distilled-preset`, or override/e python streaming_video_client.py \ --extra-params '{"pyramid_num_inference_steps_list":[2, 2, 2]}' ``` + +### LingBot-World 2.0 + +`robbyant/lingbot-world-v2-14b-causal-fast-diffusers` is image-conditioned and runs on the +AR-Diffusion engine, so it needs the stepwise deploy config and a first frame. See +[the recipe](../../../recipes/Robbyant/LingBot-World-2.0.md) for the constraints on +`width`/`height` and `num_frames`. + +```bash +vllm serve robbyant/lingbot-world-v2-14b-causal-fast-diffusers \ + --omni \ + --deploy-config vllm_omni/deploy/lingbot_world_v2_stepwise.yaml \ + --port 8000 +``` + +```bash +python streaming_video_client.py \ + --model robbyant/lingbot-world-v2-14b-causal-fast-diffusers \ + --prompt "The camera moves slowly forward through the scene." \ + --image-reference /path/to/first_frame.png \ + --width 832 --height 480 --num-frames 33 --fps 16 --seed 42 \ + --extra-params '{"camera_action_script":[[["w"],["w"],["w"]],[["a"],[],[]],[[],[],[]]]}' \ + --output lingbot_world_v2_stream.mp4 +``` + +`camera_action_script` carries one three-latent-frame WASD action list per generated +chunk. Mid-session `session.interaction` only updates the prompt, so a LingBot rollout +follows the camera script it started with. diff --git a/examples/online_serving/streaming_video_generation/streaming_video_client.py b/examples/online_serving/streaming_video_generation/streaming_video_client.py index d1a9052ed..268891284 100644 --- a/examples/online_serving/streaming_video_generation/streaming_video_client.py +++ b/examples/online_serving/streaming_video_generation/streaming_video_client.py @@ -1,11 +1,12 @@ #!/usr/bin/env python3 # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Text-only WebSocket client for `/v1/realtime/video`. +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""WebSocket client for `/v1/realtime/video`. The client prints a line for every binary video chunk it receives and writes the received bytes to disk when the session finishes or is interrupted. -Image/reference input is intentionally not supported in this example yet. +``--image-reference`` sends a first frame for image-conditioned models such as +LingBot-World 2.0; text-only models need no image. Requirements: pip install av websockets @@ -13,9 +14,11 @@ Requirements: import argparse import asyncio +import base64 import contextlib import io import json +import mimetypes import os import time from dataclasses import dataclass @@ -53,6 +56,18 @@ class ScheduledPromptUpdate: transition_chunks: int +def _image_reference(value: str) -> dict[str, Any]: + """Accept an http(s) URL as-is, or inline a local image as a data URL.""" + if value.startswith(("http://", "https://", "data:")): + return {"image_url": value} + path = Path(value) + if not path.is_file(): + raise argparse.ArgumentTypeError(f"--image-reference is neither a URL nor a readable file: {value}") + media_type = mimetypes.guess_type(path.name)[0] or "image/png" + encoded = base64.b64encode(path.read_bytes()).decode() + return {"image_url": f"data:{media_type};base64,{encoded}"} + + def _parse_extra_params(value: str) -> dict[str, Any]: try: parsed = json.loads(value) @@ -116,6 +131,7 @@ def build_session_start(args: argparse.Namespace) -> dict[str, Any]: "prompt": args.prompt, } + _maybe_set(payload, "image_reference", args.image_reference) _maybe_set(payload, "negative_prompt", args.negative_prompt) _maybe_set(payload, "width", args.width) _maybe_set(payload, "height", args.height) @@ -362,6 +378,12 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--fps", type=int, default=16, help="Frames per second") parser.add_argument("--num-frames", type=int, default=429, help="Number of generated frames") + parser.add_argument( + "--image-reference", + type=_image_reference, + default=None, + help="First frame for image-conditioned models: a local image path or an http(s) URL.", + ) parser.add_argument("--negative-prompt", default=None, help="Negative prompt") parser.add_argument("--num-inference-steps", type=int, default=None, help="Number of diffusion steps") parser.add_argument("--guidance-scale", type=float, default=1.0, help="Classifier-free guidance scale") @@ -383,36 +405,16 @@ def parse_args() -> argparse.Namespace: default=None, help="JSON object merged into extra_params; overrides preset keys on conflict.", ) + # Midway prompt updates are pipeline-specific: a model that does not implement + # them (LingBot-World, for one) rejects the whole session on the first one, so + # nothing is scheduled unless the caller asks for it. parser.add_argument( "--prompt-updates", type=_parse_prompt_updates, - default=json.dumps( - [ - { - "at": 4, - "prompt": ( - "An underwater tornado appears and affects the ocean floor in a dramatic and chaotic scene. " - "The water is murky, swirling violently, carrying debris and marine life into the vortex. " - "The tropical fish on the scene all swim in panic, trying to avoid the powerful currents. " - "The camera remains stationary, capturing the intensity of the underwater tornado as it disrupts the serene ocean floor. " - "Close-up shot emphasizing the turbulent motion and destruction." - ), - }, - { - "at": 11, - "prompt": ( - "The swirling underwater vortex now seizes a heavy, encrusted treasure chest, its lid flapping open as it is smashed onto the ocean floor. " - "Gold coins and silver trinkets spill out, glittering briefly in the murky water before being swept instantly into the violent funnel. " - "The heavy wooden box tumbles end over end, colliding with floating rocks and adding to the debris field. " - "Swirling sediment and bubbles surround the spilling fortune, highlighting the chaotic power of the storm as it ravages the seabed. " - "Close-up shot emphasizing the turbulent motion and destruction." - ), - # "transition_chunks": 3, - }, - ] - ), + default=[], help=( - "JSON array of scheduled prompt updates. Each object requires " + "JSON array of scheduled prompt updates; none are sent by default. " + "Each object requires " "'at' (seconds after video.start) and 'prompt'. Optional " f"'transition_chunks' defaults to {DEFAULT_TRANSITION_CHUNKS}." ), diff --git a/recipes/Robbyant/LingBot-World-2.0.md b/recipes/Robbyant/LingBot-World-2.0.md index 851aeb930..c38e24fd2 100644 --- a/recipes/Robbyant/LingBot-World-2.0.md +++ b/recipes/Robbyant/LingBot-World-2.0.md @@ -7,13 +7,31 @@ - Vendor: Robbyant - Model: `robbyant/lingbot-world-v2-14b-causal-fast-diffusers` - Task: image-conditioned interactive world generation -- Modes: offline trajectory replay and in-process realtime AR-Diffusion ticks +- Modes: offline trajectory replay, stepwise one-request streaming, and in-process realtime AR-Diffusion ticks (deprecated) - Hardware validated: NVIDIA H200 and B200 - Maintainer: Community The checkpoint is separately licensed under CC BY-NC-SA and restricted to non-commercial use. The vLLM-Omni integration code remains Apache-2.0. +## Which path to use + +There are three ways to drive this model, each documented in its own section +below. + +- **Realtime stepwise** — see [Streaming video serving](#streaming-video-serving). + The suggested path: one `WS /v1/realtime/video` session produces the whole + rollout, one video chunk per AR block. It cannot take realtime camera + interaction yet; that is being added, and this path is the one that will + remain. +- **Realtime tick** *(deprecated)* — see + [Realtime in-process generation](#realtime-in-process-generation-deprecated). + The only path that takes realtime camera interaction today, so it stays until + the stepwise path can, but it takes no new work and will be removed. +- **Offline** — see [Offline generation](#offline-generation). Replays a fixed + pose/intrinsics trajectory in one request and writes an MP4. Use it when the + camera path is known up front and streaming is not needed. + ## Offline generation The offline path consumes one source image and an action directory containing @@ -34,7 +52,12 @@ counts must be `9 + 12k`, up to 117 frames. Camera files may contain a longer bounded source trajectory (the official example contains 269 frames); the runtime consumes only the prefix needed by the request. -## Realtime in-process generation +## Realtime in-process generation (deprecated) + +Use [Streaming video serving](#streaming-video-serving) instead unless you need +mid-session camera interaction, which only this path has. It still works and +emits no runtime warning; removal is tracked as B4 of the LingBot World 2.0 +roadmap ([#6672](https://github.com/vllm-project/vllm-omni/issues/6672)). The realtime example keeps AR-Diffusion KV and model-owned state across requests. Each JSONL line describes the prompt and/or three latent-frame camera @@ -61,12 +84,81 @@ exercises the same `ARDiffusionSessionManager -> ARDiffusionOmniTickConsumer -> AsyncOmni -> ARDiffusionEngine` path used by a future HTTP or WebSocket transport. -This PR intentionally does not define a public realtime HTTP/WebSocket schema. -An online serving client should be added together with the public transport -API rather than exposing LingBot-specific event fields from the generic model -runtime. The structured camera-interaction frontend and transport work is -tracked separately in -[vllm-project/vllm-omni#5527](https://github.com/vllm-project/vllm-omni/pull/5527). +Realtime serving now goes through the generic `WS /v1/realtime/video` +transport, which exposes no LingBot-specific event fields; the tick loop above +stays available only for the older control plane. + +## Streaming video serving + +The stepwise path keeps AR-Diffusion paged KV but issues **one** request for +the whole rollout: `prepare_encode` runs once, then every AR block is four DMD +steps followed by one streamed chunk. Identity metadata uses +`session_id = request_id` with contiguous `chunk_index` values from zero. +Camera motion for this landing is request-scoped: one three-frame action list +per chunk, fixed when the rollout starts. + +Serve it with the AR-Diffusion deploy config, which selects the AR-Diffusion +engine and enables streamed step execution: + +```bash +vllm serve robbyant/lingbot-world-v2-14b-causal-fast-diffusers \ + --omni \ + --deploy-config vllm_omni/deploy/lingbot_world_v2_stepwise.yaml \ + --port 8000 +``` + +`--deploy-config` is required: the `lingbot_world` pipeline deliberately +registers no default deploy config, so offline replay and the tick example +keep their request-mode topology when no deploy config is given, and the +stepwise serving topology is only ever an explicit choice. + +Clients then use the generic WebSocket protocol documented in +[`docs/serving/video_stream_api.md`](../../docs/serving/video_stream_api.md): +`session.start` begins one rollout and each AR block arrives as a binary video +chunk. The model is image-conditioned, so the first frame is required and rides +on `image_reference` as an `http(s)` or `data:` URL; per-chunk camera actions +ride on `extra_params`: + +```json +{"type": "session.start", "model": "robbyant/lingbot-world-v2-14b-causal-fast-diffusers", + "prompt": "The camera moves slowly forward through the scene.", + "image_reference": {"image_url": "data:image/png;base64,"}, + "width": 832, "height": 480, "num_frames": 33, + "extra_params": {"flow_shift": 5.0, + "camera_action_script": [[["w"], ["w"], ["w"]], [["a"], [], []], [[], [], []]]}} +``` + +`camera_action_script` must hold exactly one action list per generated +chunk, and a request generates `((num_frames - 1) // 4 + 1) // 3` chunks: +three for `num_frames: 33`, seven for `num_frames: 81`. + +The bundled client speaks this protocol: + +```bash +python examples/online_serving/streaming_video_generation/streaming_video_client.py \ + --model robbyant/lingbot-world-v2-14b-causal-fast-diffusers \ + --prompt "The camera moves slowly forward through the scene." \ + --image-reference /path/to/first_frame.png \ + --width 832 --height 480 --num-frames 33 --fps 16 --seed 42 \ + --extra-params '{"camera_action_script":[[["w"],["w"],["w"]],[["a"],[],[]],[[],[],[]]]}' \ + --output lingbot_world_v2_stream.mp4 +``` + +A served request may instead point at a pose/intrinsics trajectory with +`extra_params.action_path`, which is resolved inside the trusted root set by +`model_config.lingbot_action_root` or `VLLM_OMNI_LINGBOT_ACTION_ROOT`; a server +started without that root configured accepts only `camera_action_script`. + +Requested `width`/`height` must match `ar_diffusion_width`/`ar_diffusion_height` +in the deploy config, because the AR cache geometry is fixed at load time. +Blocks are decoded independently, so seams between chunks are possible; a +session-owned streaming decoder is tracked separately. + +Mid-session `session.interaction` for camera control is not wired yet, so a +served rollout follows the script it started with. + +The deprecated tick example remains available for the older +one-block-per-`generate()` control plane. ## Realtime identity and controls @@ -104,10 +196,16 @@ tested commit. ## Current limitations - Only the 14B causal-fast checkpoint is supported. -- The realtime control plane is internal; there is no public server transport yet. +- The tick control plane is internal; the public transport is the stepwise + `WS /v1/realtime/video` path, which cannot take mid-session camera + interaction yet. +- Stepwise serving requires an explicit + `--deploy-config vllm_omni/deploy/lingbot_world_v2_stepwise.yaml`; there is + no default deploy config for this model. - AR-Diffusion stages currently require one replica because session-affine routing across replicas is not implemented. -- One AR block is generated per request and `max_num_seqs` must be one. +- Tick mode generates one AR block per request. Stepwise mode generates many + AR blocks in one request. `max_num_seqs` must be one in both cases. - Stateful streaming VAE decode is not implemented; the realtime example emits latent chunks. - SP/USP, pipeline/CFG parallelism, HSDP, VAE parallelism, quantization, @@ -119,4 +217,7 @@ tested commit. - Checkpoint: - Official implementation: - Offline example: [`examples/offline_inference/diffusion/lingbot_world_v2.py`](../../examples/offline_inference/diffusion/lingbot_world_v2.py) +- Streaming serving client: [`examples/online_serving/streaming_video_generation/`](../../examples/online_serving/streaming_video_generation/README.md) +- Streaming serving end-to-end test: [`tests/e2e/online_serving/test_lingbot_world_v2_stepwise.py`](../../tests/e2e/online_serving/test_lingbot_world_v2_stepwise.py) +- Stepwise latent-level end-to-end test: [`tests/e2e/offline_inference/test_lingbot_world_v2_stepwise.py`](../../tests/e2e/offline_inference/test_lingbot_world_v2_stepwise.py) - Realtime design: [`docs/design/feature/realtime_ar_diffusion.md`](../../docs/design/feature/realtime_ar_diffusion.md) diff --git a/tests/config/test_lingbot_world_stepwise_deploy.py b/tests/config/test_lingbot_world_stepwise_deploy.py new file mode 100644 index 000000000..29d355c1c --- /dev/null +++ b/tests/config/test_lingbot_world_stepwise_deploy.py @@ -0,0 +1,52 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""The LingBot stepwise deploy config must keep serving on the AR-Diffusion path. + +`WS /v1/realtime/video` only reaches the stepwise code when the stage selects the +AR-Diffusion engine and asks for streamed step execution, so those fields are +pinned here rather than left to drift. +""" + +from __future__ import annotations + +import pytest + +from tests.helpers.stage_config import get_deploy_config_path +from vllm_omni.config.stage_config import load_deploy_config + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_DEPLOY = "lingbot_world_v2_stepwise.yaml" + + +def _stage(): + deploy = load_deploy_config(get_deploy_config_path(_DEPLOY)) + assert len(deploy.stages) == 1, "stepwise serving is a single diffusion stage" + return deploy.stages[0] + + +def test_stage_selects_the_ar_diffusion_engine_and_pipeline() -> None: + stage = _stage() + assert stage.model_class_name == "LingBotWorldCausalDMDPipeline" + # engine_backend/model_config reach the engine through the passthrough, the + # same way the DreamZero AR-Diffusion deploy config does. + assert stage.engine_extras["engine_backend"].endswith("ar_diffusion.engine.ARDiffusionEngine") + + +def test_stage_requests_streamed_step_execution_on_one_sequence() -> None: + stage = _stage() + assert stage.step_execution is True + assert stage.engine_extras["streaming_output"] is True + # AR-Diffusion binds one session per runner invocation; batching requests + # would break that contract, so the runner rejects max_num_seqs > 1. + assert stage.max_num_seqs == 1 + + +def test_stage_declares_the_fixed_ar_cache_geometry() -> None: + stage = _stage() + model_config = stage.engine_extras["model_config"] + # The KV cache geometry is fixed at load time, so a served request must ask + # for exactly this resolution. + assert model_config["ar_diffusion_height"] == 480 + assert model_config["ar_diffusion_width"] == 832 + assert model_config["ar_diffusion_kv_config"]["gpu_memory_fraction"] == 0.6 diff --git a/tests/diffusion/ar_diffusion/test_capability_runner.py b/tests/diffusion/ar_diffusion/test_capability_runner.py index 10451c134..70166189b 100644 --- a/tests/diffusion/ar_diffusion/test_capability_runner.py +++ b/tests/diffusion/ar_diffusion/test_capability_runner.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """CPU contracts for capability-driven AR-Diffusion sessions.""" from __future__ import annotations @@ -10,6 +11,7 @@ from types import SimpleNamespace import pytest import torch +from vllm_omni.diffusion.sched.interface import CachedRequestData, DiffusionSchedulerOutput from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner from vllm_omni.experimental.ar_diffusion.capability import ( ARDiffusionCrossAttentionKVSpec, @@ -27,6 +29,22 @@ NEG = "negative" pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +@pytest.fixture(autouse=True) +def platform_synchronize_calls(monkeypatch) -> list[None]: + """Stub the platform barrier the runner takes at every chunk boundary. + + These tests run on a CPU host, where the resolved platform is + ``UnspecifiedOmniPlatform`` and ``synchronize()`` raises. Tests that need + a failing barrier override this with their own fake. + """ + calls: list[None] = [] + monkeypatch.setattr( + "vllm_omni.experimental.ar_diffusion.runner.current_omni_platform", + SimpleNamespace(synchronize=lambda: calls.append(None)), + ) + return calls + + def lingbot_like_spec(*, capacity: int = 2) -> ARDiffusionKVCacheSpec: """Single-kv_branch causal DMD: 3 latent frames/block + sink/window + text KV.""" return ARDiffusionKVCacheSpec( @@ -114,6 +132,43 @@ class BatchCapablePipeline(CapablePipeline): supports_request_batch = True +class StepCapablePipeline(CapablePipeline): + supports_step_execution = True + + def prepare_encode(self, state, **kwargs): + del kwargs + return state + + def denoise_step(self, input_batch, *, states=None, **kwargs): + del input_batch, states, kwargs + return None + + def step_scheduler(self, state, noise_pred, **kwargs): + del state, noise_pred, kwargs + + def post_decode(self, state, **kwargs): + del state, kwargs + from vllm_omni.diffusion.data import DiffusionOutput + + return DiffusionOutput() + + +def scheduler_output( + request_ids: tuple[str, ...] = ("req-1",), + *, + finished: tuple[str, ...] = (), +) -> DiffusionSchedulerOutput: + """Build the real scheduler payload; stepwise requests arrive as cached rows.""" + return DiffusionSchedulerOutput( + step_id=0, + scheduled_new_reqs=[], + scheduled_cached_reqs=CachedRequestData(request_ids=list(request_ids)), + finished_req_ids=set(finished), + num_running_reqs=len(request_ids), + num_waiting_reqs=0, + ) + + def make_runner( pipeline: object, *, @@ -140,6 +195,7 @@ def make_runner( runner._sessions = OrderedDict() runner._session_capacity = 0 runner._perf_e2e_times = [] + runner._stepwise_chunk_started = {} runner._preallocate_kv_cache(available_bytes=available_bytes) return runner @@ -395,12 +451,16 @@ def test_synchronize_exception_uses_forward_cleanup_path(monkeypatch): ctx.ensure_video_slots(torch.device("cpu")) return object() - def synchronize_boom(device): + def synchronize_boom(): raise RuntimeError("asynchronous kernel failed") monkeypatch.setattr(DiffusionModelRunner, "execute_model", return_after_allocation) - monkeypatch.setattr(torch.accelerator, "synchronize", synchronize_boom) - runner.device = torch.device("cuda") + # The runner defers the device barrier to the platform, so the failure is + # injected there rather than on a device-specific torch entry point. + monkeypatch.setattr( + "vllm_omni.experimental.ar_diffusion.runner.current_omni_platform", + SimpleNamespace(synchronize=synchronize_boom), + ) request = SimpleNamespace( request_id="broken-request", sampling_params=SimpleNamespace(extra_args={"session_id": "broken"}), @@ -424,6 +484,11 @@ def test_ar_runner_rejects_step_and_request_batch_modes(): make_runner(BatchCapablePipeline(lingbot_like_spec())) +def test_ar_runner_allows_step_execution_when_pipeline_implements_the_contract(): + runner = make_runner(StepCapablePipeline(lingbot_like_spec()), step_execution=True) + assert runner.kv_cache is not None + + def test_ar_runner_defensively_rejects_inherited_batch_and_step_entrypoints(): runner = object.__new__(ARDiffusionModelRunner) with pytest.raises(RuntimeError, match="request-batch execution"): @@ -432,6 +497,144 @@ def test_ar_runner_defensively_rejects_inherited_batch_and_step_entrypoints(): runner.execute_stepwise(None) +def test_execute_stepwise_binds_request_id_session_and_unbinds_after_call(monkeypatch): + from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput + + pipeline = StepCapablePipeline(lingbot_like_spec()) + runner = make_runner(pipeline, step_execution=True) + bound_during = [] + + def fake_execute(self, scheduler_output): + del self, scheduler_output + bound_during.append(pipeline.bound_state is not None) + return BatchRunnerOutput.from_list([RunnerOutput(request_id="req-1", finished=False)]) + + monkeypatch.setattr(DiffusionModelRunner, "execute_stepwise", fake_execute) + output = runner.execute_stepwise(scheduler_output()) + + assert bound_during == [True] + assert pipeline.bound_state is None + assert pipeline.binds == ["req-1"] + assert "req-1" in runner._sessions + assert output.get_request_output("req-1") is not None + assert not pipeline.closes + + +def test_execute_stepwise_closes_session_when_request_finishes(monkeypatch): + from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput + + pipeline = StepCapablePipeline(lingbot_like_spec()) + runner = make_runner(pipeline, step_execution=True) + + def fake_execute(self, scheduler_output): + del self, scheduler_output + return BatchRunnerOutput.from_list([RunnerOutput(request_id="req-1", finished=True)]) + + monkeypatch.setattr(DiffusionModelRunner, "execute_stepwise", fake_execute) + runner.execute_stepwise(scheduler_output()) + + assert pipeline.bound_state is None + assert pipeline.closes == ["req-1"] + assert "req-1" not in runner._sessions + + +def test_execute_stepwise_exception_releases_kv_and_does_not_resume_pages(monkeypatch): + pipeline = StepCapablePipeline(lingbot_like_spec()) + runner = make_runner(pipeline, step_execution=True) + first = runner._get_or_create_session("req-1") + + def boom(self, scheduler_output): + del self, scheduler_output + raise RuntimeError("step failed") + + monkeypatch.setattr(DiffusionModelRunner, "execute_stepwise", boom) + with pytest.raises(RuntimeError, match="step failed"): + runner.execute_stepwise(scheduler_output()) + + assert pipeline.bound_state is None + assert pipeline.closes == ["req-1"] + assert "req-1" not in runner._sessions + second = runner._get_or_create_session("req-1") + assert second is not first + + +def test_execute_stepwise_closes_session_for_scheduler_aborted_request(monkeypatch): + from vllm_omni.diffusion.worker.utils import BatchRunnerOutput + + pipeline = StepCapablePipeline(lingbot_like_spec()) + runner = make_runner(pipeline, step_execution=True) + runner._get_or_create_session("req-1") + + def fake_execute(self, scheduler_output): + del self, scheduler_output + return BatchRunnerOutput.from_list([]) + + monkeypatch.setattr(DiffusionModelRunner, "execute_stepwise", fake_execute) + # An aborted request is never scheduled again and never reports finished, + # so the retire path is the only thing that can free its KV. + runner.execute_stepwise(scheduler_output(request_ids=(), finished=("req-1",))) + + assert "req-1" not in runner._sessions + assert pipeline.closes == ["req-1"] + + +def test_execute_stepwise_times_once_per_chunk_not_once_per_step(monkeypatch, platform_synchronize_calls): + from vllm_omni.diffusion.data import DiffusionOutput + from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput + + pipeline = StepCapablePipeline(lingbot_like_spec()) + runner = make_runner(pipeline, step_execution=True) + results = [None, None, None, DiffusionOutput()] + + def fake_execute(self, scheduler_output): + del self, scheduler_output + return BatchRunnerOutput.from_list([RunnerOutput(request_id="req-1", finished=False, result=results.pop(0))]) + + monkeypatch.setattr(DiffusionModelRunner, "execute_stepwise", fake_execute) + for _ in range(4): + runner.execute_stepwise(scheduler_output()) + + # Four denoise steps, one emitted AR block: one timing sample, matching the + # one-sample-per-block meaning request mode already has. + assert len(runner._perf_e2e_times) == 1 + assert len(platform_synchronize_calls) == 1 + assert "req-1" not in runner._stepwise_chunk_started + + +def test_execute_stepwise_synchronize_exception_fails_closed(monkeypatch): + """The chunk barrier is where an asynchronous device error surfaces. It must + clean up like an in-step failure instead of leaving the session, the cached + step state and the timing entry to the scheduler's next cycle.""" + from vllm_omni.diffusion.worker.utils import BatchRunnerOutput, RunnerOutput + + pipeline = StepCapablePipeline(lingbot_like_spec()) + runner = make_runner(pipeline, step_execution=True) + runner.state_cache = {"req-1": object()} + + def fake_execute(self, scheduler_output): + del self, scheduler_output + return BatchRunnerOutput.from_list([RunnerOutput(request_id="req-1", finished=False, result=object())]) + + def synchronize_boom(): + raise RuntimeError("asynchronous kernel failed") + + monkeypatch.setattr(DiffusionModelRunner, "execute_stepwise", fake_execute) + monkeypatch.setattr( + "vllm_omni.experimental.ar_diffusion.runner.current_omni_platform", + SimpleNamespace(synchronize=synchronize_boom), + ) + + with pytest.raises(RuntimeError, match="asynchronous kernel failed"): + runner.execute_stepwise(scheduler_output()) + + assert pipeline.bound_state is None + assert pipeline.closes == ["req-1"] + assert "req-1" not in runner._sessions + assert "req-1" not in runner.state_cache + assert "req-1" not in runner._stepwise_chunk_started + assert not runner._perf_e2e_times + + def test_model_specific_warmup_provider_is_consumed(monkeypatch): requests = [object(), object()] pipeline = WarmupPipeline(lingbot_like_spec(), requests) diff --git a/tests/diffusion/models/lingbot_world/test_lingbot_world_actions.py b/tests/diffusion/models/lingbot_world/test_lingbot_world_actions.py index df012777f..e26e44c2c 100644 --- a/tests/diffusion/models/lingbot_world/test_lingbot_world_actions.py +++ b/tests/diffusion/models/lingbot_world/test_lingbot_world_actions.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations @@ -13,6 +14,7 @@ from vllm_omni.diffusion.models.lingbot_world.actions import ( LingBotCameraControlReducer, integrate_lingbot_camera_actions, parse_lingbot_camera_action_frames, + parse_lingbot_camera_action_script, ) from vllm_omni.experimental.ar_diffusion.session import ( ARDiffusionSession, @@ -292,3 +294,8 @@ def test_session_reducer_keeps_prompt_and_actions_atomic_and_resets() -> None: assert consumer.ticks[1].controls[0].data["frames"] == ((), (), ()) asyncio.run(exercise()) + + +def test_parse_camera_action_script_rejects_wrong_chunk_width() -> None: + with pytest.raises(ValueError, match="exactly 3 per-latent-frame"): + parse_lingbot_camera_action_script([[["w"], ["w"]]], frames_per_chunk=3) diff --git a/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py b/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py index 1a08e8270..52113208a 100644 --- a/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py +++ b/tests/diffusion/models/lingbot_world/test_pipeline_lingbot_world.py @@ -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 @@ -15,8 +15,10 @@ from diffusers.utils.torch_utils import randn_tensor as _diffusers_randn_tensor from PIL import Image from torch import nn +import vllm_omni.diffusion.models.lingbot_world.dmd_block as lingbot_dmd_block import vllm_omni.diffusion.models.lingbot_world.pipeline as lingbot_pipeline from tests.diffusion.models.wan2_2.conftest import noop_progress_bar +from vllm_omni.diffusion.models.interface import SupportsStepExecution, supports_step_execution from vllm_omni.diffusion.models.lingbot_world.actions import ( integrate_lingbot_camera_actions, ) @@ -24,6 +26,8 @@ from vllm_omni.diffusion.models.lingbot_world.camera import CameraTrajectory as from vllm_omni.diffusion.models.lingbot_world.camera import ( build_plucker_embedding as _real_build_plucker_embedding, ) +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.worker.utils import StepRequestState from vllm_omni.experimental.ar_diffusion.tick_protocol import ( ARDiffusionControlInput, ARDiffusionTickRequest, @@ -253,6 +257,8 @@ class _SamplingParams: if extra_args is not None: self.extra_args.update(extra_args) self.latents = None + self.guidance_scale = None + self.guidance_scale_2 = None class _RequestBatch: @@ -345,7 +351,6 @@ def _stub_pipeline_dependencies(monkeypatch: pytest.MonkeyPatch) -> None: "FlowUniPCMultistepScheduler": _FakeScheduler, "CausalLingBotWorldTransformer3DModel": _FakeTransformerFactory, "get_local_device": lambda: torch.device("cpu"), - "set_forward_context_denoise_step_idx": lambda index: None, "prefetch_subfolders": prefetch_subfolders, "from_pretrained_with_prefetch": from_pretrained_with_prefetch, "load_transformer_config": load_transformer_config, @@ -358,6 +363,9 @@ def _stub_pipeline_dependencies(monkeypatch: pytest.MonkeyPatch) -> None: } for name, value in replacements.items(): monkeypatch.setattr(lingbot_pipeline, name, value) + # The DMD block math lives in its own module; stub the symbols it imports. + monkeypatch.setattr(lingbot_dmd_block, "set_forward_context_denoise_step_idx", lambda index: None) + monkeypatch.setattr(lingbot_dmd_block, "randn_tensor", _diffusers_randn_tensor) monkeypatch.setattr(lingbot_pipeline, "_loader_state", loader_state, raising=False) @@ -774,6 +782,7 @@ def test_denoise_state_stays_fp32_while_transformer_inputs_use_model_dtype() -> return torch.zeros(shape, device=device, dtype=dtype) module.randn_tensor = randn + lingbot_dmd_block.randn_tensor = randn result = pipeline(_request()) assert requested_noise_dtypes == [torch.float32] * 4 @@ -1139,7 +1148,7 @@ def test_first_frame_condition_and_camera_fold_match_transformer_contract() -> N module = _load_pipeline_module() transformer = _RecordingTransformer() pipeline = _pipeline(module, transformer=transformer) - module.randn_tensor = lambda shape, **kwargs: torch.full( + module.randn_tensor = lingbot_dmd_block.randn_tensor = lambda shape, **kwargs: torch.full( shape, -99.0, device=kwargs["device"], @@ -1745,6 +1754,288 @@ def test_request_cache_becomes_unreachable_after_transformer_error() -> None: assert not hasattr(pipeline, "transformer_cache") +class _FakeARState: + def __init__(self, session_id: str = "req-1") -> None: + self.session_id = session_id + self.commits: list[str] = [] + + def get_kv_caches(self, branch, *, seq_len, commit_current): + del branch, seq_len, commit_current + return [SimpleNamespace()] + + def commit_paged_context(self, branch): + self.commits.append(branch) + + def clear_cross_attention(self): + return None + + def is_cross_attention_populated(self, branch, name): + del branch, name + return True + + def get_cross_attention_kv(self, branch, name): + del branch, name + zeros = torch.zeros(1, 1, 2, 4) + return [{"k": zeros, "v": zeros} for _ in range(2)] + + +def _empty_action_script(num_chunks: int): + return tuple(((), (), ()) for _ in range(num_chunks)) + + +def _stepwise_state( + *, + request_id: str = "req-1", + num_frames: int = 21, + script=None, + seed: int = 17, +): + num_chunks = ((num_frames - 1) // 4 + 1) // 3 + sampling = _SamplingParams(num_frames=num_frames, seed=seed) + sampling.extra_args["_lingbot_camera_trajectory"] = None + sampling.extra_args["_lingbot_camera_actions"] = None + sampling.extra_args["_lingbot_camera_action_script"] = ( + script if script is not None else _empty_action_script(num_chunks) + ) + return StepRequestState( + request_id=request_id, + sampling=sampling, + prompt=_prompt(), + ) + + +def _run_stepwise(pipeline, state): + outputs = [] + pipeline.prepare_encode(state) + while not state.request_denoise_completed: + noise = pipeline.denoise_step(None, states=[state]) + pipeline.step_scheduler(state, noise) + if state.chunk_denoise_completed: + outputs.append(pipeline.post_decode(state)) + return outputs + + +def test_pipeline_declares_step_execution_support() -> None: + module = _load_pipeline_module() + pipeline = _pipeline(module) + + assert pipeline.supports_step_execution is True + assert isinstance(pipeline, SupportsStepExecution) + assert supports_step_execution(pipeline) is True + + +def test_preprocess_materializes_camera_action_script_without_action_path() -> None: + module = _load_pipeline_module() + sampling = _SamplingParams(include_action=False) + sampling.extra_args["camera_action_script"] = [[["w"], ["w"], ["w"]], [["a"], [], []]] + request = OmniDiffusionRequest(prompt=_prompt(), sampling_params=sampling, request_id="req-1") + + result = module.get_lingbot_world_pre_process_func(_od_config())(request) + + assert result.sampling_params.extra_args["_lingbot_camera_trajectory"] is None + assert result.sampling_params.extra_args["_lingbot_camera_actions"] is None + assert result.sampling_params.extra_args["_lingbot_camera_action_script"] == ( + (("w",), ("w",), ("w",)), + (("a",), (), ()), + ) + + +@pytest.mark.parametrize( + "extra_args", + [ + pytest.param({}, id="request_mode"), + pytest.param(_tick_extra_args(chunk_index=0), id="tick_mode"), + ], +) +def test_forward_rejects_a_stepwise_camera_action_script(extra_args) -> None: + """A script only steers step execution, so request mode must not drop it silently.""" + module = _load_pipeline_module() + pipeline = _pipeline(module) + sampling = _SamplingParams( + include_action=False, + extra_args={ + **extra_args, + "_lingbot_camera_trajectory": None, + "_lingbot_camera_action_script": _empty_action_script(1), + }, + ) + + with pytest.raises(ValueError, match="camera_action_script is read only by LingBot step execution"): + pipeline(_request(sampling=sampling)) + + +def test_stepwise_progress_metadata_and_commit_trace() -> None: + module = _load_pipeline_module() + transformer = _RecordingTransformer() + pipeline = _pipeline(module, transformer=transformer) + pipeline._ar_height = 16 + pipeline._ar_width = 16 + state = _stepwise_state(num_frames=21) + fake = _FakeARState(state.request_id) + + with pipeline.bind_ar_diffusion_state(state.request_id, fake): + outputs = _run_stepwise(pipeline, state) + + assert pipeline._ar_diffusion_kv_state is None + assert len(outputs) == 2 + assert [output.chunk_index for output in outputs] == [0, 1] + assert all(output.total_chunks == 2 for output in outputs) + assert outputs[-1].finished is True + assert [call["update_cache"] for call in transformer.calls] == [False, False, False, False, True] * 2 + assert [call["start_frame"] for call in transformer.calls] == [0] * 5 + [3] * 5 + assert fake.commits == ["main", "main"] + for chunk_index, output in enumerate(outputs): + metadata = output.output["metadata"]["ar_diffusion"] + assert metadata == { + "session_id": "req-1", + "request_id": "req-1", + "chunk_index": chunk_index, + "applied_event_ids": [], + } + assert output.output["payload"]["latents"].shape == (1, 16, 3, 2, 2) + + +def test_stepwise_matches_tick_transformer_trace_and_latents() -> None: + module = _load_pipeline_module() + transformer = _RecordingTransformer() + pipeline = _pipeline(module, transformer=transformer) + pipeline._ar_height = 16 + pipeline._ar_width = 16 + script = ((("w",), ("w",), ("w",)), (("d",), ("d",), ("d",))) + fake = _FakeARState("world-1") + + with pipeline.bind_ar_diffusion_state("world-1", fake): + tick_latents = [] + for chunk_index, frames in enumerate(script): + tick = ARDiffusionTickRequest( + session_id="world-1", + request_id="legacy-lingbot-request", + chunk_index=chunk_index, + controls=( + ARDiffusionControlInput( + track="camera", + schema="lingbot.camera_actions.v1", + data={"mode": "frames", "frames": [list(frame) for frame in frames]}, + ), + ), + ) + sampling = _SamplingParams(extra_args=tick.to_extra_args(), seed=17) + sampling.extra_args["_lingbot_camera_trajectory"] = None + sampling.extra_args["_lingbot_camera_actions"] = frames + result = pipeline(_request(sampling=sampling)) + tick_latents.append(result.output["payload"]["latents"].clone()) + tick_calls = list(transformer.calls) + transformer.calls.clear() + pipeline.close_ar_diffusion_session("world-1") + + stepwise_state = _stepwise_state(request_id="req-1", num_frames=21, script=script, seed=17) + stepwise_fake = _FakeARState(stepwise_state.request_id) + with pipeline.bind_ar_diffusion_state(stepwise_state.request_id, stepwise_fake): + stepwise_outputs = _run_stepwise(pipeline, stepwise_state) + + assert [call["update_cache"] for call in transformer.calls] == [call["update_cache"] for call in tick_calls] + assert [call["start_frame"] for call in transformer.calls] == [call["start_frame"] for call in tick_calls] + for tick_latent, output in zip(tick_latents, stepwise_outputs, strict=True): + torch.testing.assert_close(output.output["payload"]["latents"], tick_latent) + + +def test_stepwise_trajectory_camera_matches_request_mode_under_non_uniform_speed(monkeypatch) -> None: + """A trajectory that slows down between blocks must condition every stepwise + chunk exactly as request mode does. Request mode embeds the whole trajectory + once; embedding per chunk would re-normalize each block's framewise + translations by its own largest step and present a slow block as full-speed + motion.""" + from vllm_omni.diffusion.models.lingbot_world import camera as camera_module + + module = _load_pipeline_module() + # The fixture stubs the ray embedding with a pose-blind ramp; parity needs + # the real geometry so a scale mismatch actually shows up. + monkeypatch.setattr(module, "build_plucker_embedding", camera_module.build_plucker_embedding) + transformer = _RecordingTransformer() + pipeline = _pipeline(module, transformer=transformer) + pipeline._ar_height = 16 + pipeline._ar_width = 16 + + # Six latent frames -> two 3-frame blocks. The first block moves at full + # speed and everything after it at a tenth of that, so a per-chunk + # normalization would rescale the second block by 10x. + steps = torch.tensor([1.0] * 3 + [0.1] * 18) + poses = torch.eye(4).repeat(21, 1, 1) + poses[:, 2, 3] = torch.cumsum(steps, dim=0) - steps[0] + trajectory = _CameraTrajectory( + poses=poses, + intrinsics=torch.tensor([[100.0, 100.0, 8.0, 8.0]]).repeat(21, 1), + ) + + sampling = _SamplingParams(num_frames=21) + sampling.extra_args["_lingbot_camera_trajectory"] = trajectory + pipeline(_request(sampling=sampling)) + request_cameras = [call["camera_hidden_states"] for call in transformer.calls] + transformer.calls.clear() + + state = _stepwise_state(num_frames=21) + state.sampling.extra_args["_lingbot_camera_trajectory"] = trajectory + state.sampling.extra_args.pop("_lingbot_camera_action_script") + with pipeline.bind_ar_diffusion_state(state.request_id, _FakeARState(state.request_id)): + _run_stepwise(pipeline, state) + stepwise_cameras = [call["camera_hidden_states"] for call in transformer.calls] + + # Two blocks x (four probes + one commit), in the same order on both paths. + assert len(request_cameras) == len(stepwise_cameras) == 10 + for stepwise, request in zip(stepwise_cameras, request_cameras, strict=True): + torch.testing.assert_close(stepwise, request) + + +def test_stepwise_emits_latents_without_touching_the_vae() -> None: + """Latent mode must stay latent: streaming decode is opt-in per request.""" + module = _load_pipeline_module() + pipeline = _pipeline(module, transformer=_RecordingTransformer()) + pipeline._ar_height = 16 + pipeline._ar_width = 16 + state = _stepwise_state(num_frames=21) + fake = _FakeARState(state.request_id) + + with pipeline.bind_ar_diffusion_state(state.request_id, fake): + outputs = _run_stepwise(pipeline, state) + + assert [set(output.output["payload"]) for output in outputs] == [{"latents"}, {"latents"}] + assert pipeline.vae.decode_inputs == [] + + +def test_stepwise_decodes_each_chunk_for_streaming_consumers(monkeypatch) -> None: + """A non-latent request must stream pixels under the primary "video" key.""" + import diffusers.video_processor as video_processor_module + + module = _load_pipeline_module() + processed = [] + + class VideoProcessor: + def __init__(self, *, vae_scale_factor): + assert vae_scale_factor == 8 + + def postprocess_video(self, video, *, output_type): + processed.append((tuple(video.shape), output_type)) + return f"frames-{len(processed)}" + + monkeypatch.setattr(video_processor_module, "VideoProcessor", VideoProcessor) + pipeline = _pipeline(module, transformer=_RecordingTransformer()) + pipeline._ar_height = 16 + pipeline._ar_width = 16 + state = _stepwise_state(num_frames=21) + state.sampling.output_type = "np" + fake = _FakeARState(state.request_id) + + with pipeline.bind_ar_diffusion_state(state.request_id, fake): + outputs = _run_stepwise(pipeline, state) + + # One decode per AR block, each seeing only that block's three latent frames. + assert [tuple(latents.shape) for latents in pipeline.vae.decode_inputs] == [(1, 16, 3, 2, 2)] * 2 + assert [output_type for _, output_type in processed] == ["np", "np"] + assert [output.output["payload"] for output in outputs] == [{"video": "frames-1"}, {"video": "frames-2"}] + # Identity metadata survives the switch to pixel output. + assert [output.output["metadata"]["ar_diffusion"]["chunk_index"] for output in outputs] == [0, 1] + + def test_registry_and_model_exports_resolve_official_pipeline_class_name() -> None: module = _load_pipeline_module() resolved, entry, cache_acceleration_disabled, preprocess_name, preprocess = _resolve_pipeline_through_real_registry( @@ -1755,7 +2046,7 @@ def test_registry_and_model_exports_resolve_official_pipeline_class_name() -> No assert resolved is module.LingBotWorldCausalDMDPipeline assert cache_acceleration_disabled assert preprocess_name == "get_lingbot_world_pre_process_func" - request = SimpleNamespace(prompt=_prompt(), sampling_params=_SamplingParams()) + request = OmniDiffusionRequest(prompt=_prompt(), sampling_params=_SamplingParams(), request_id="req-1") assert preprocess(request) is request assert isinstance(request.sampling_params.extra_args["_lingbot_camera_trajectory"], _CameraTrajectory) assert module.LingBotWorldCausalDMDPipeline.__name__ == "LingBotWorldCausalDMDPipeline" diff --git a/tests/e2e/offline_inference/test_lingbot_world_v2.py b/tests/e2e/offline_inference/test_lingbot_world_v2.py index 39d9b9d95..64578af01 100644 --- a/tests/e2e/offline_inference/test_lingbot_world_v2.py +++ b/tests/e2e/offline_inference/test_lingbot_world_v2.py @@ -62,7 +62,6 @@ def test_lingbot_world_v2_tp1_one_block(omni_runner_handler: OmniRunnerHandler) height=480, width=832, num_frames=9, - num_inference_steps=4, max_sequence_length=512, seed=42, fps=16, diff --git a/tests/e2e/offline_inference/test_lingbot_world_v2_stepwise.py b/tests/e2e/offline_inference/test_lingbot_world_v2_stepwise.py new file mode 100644 index 000000000..f677398ce --- /dev/null +++ b/tests/e2e/offline_inference/test_lingbot_world_v2_stepwise.py @@ -0,0 +1,149 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""End-to-end stepwise smoke test for LingBot-World v2. + +One ``AsyncOmni.generate()`` must stream one latent chunk per AR block while +keeping a single request identity, which is what separates this path from the +tick control plane exercised by ``test_lingbot_world_v2.py``. +""" + +from __future__ import annotations + +import asyncio +import os +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import pytest +import torch +from vllm.assets.image import ImageAsset + +from tests.helpers.mark import hardware_test + +if TYPE_CHECKING: + from vllm_omni.outputs import OmniRequestOutput + +MODEL = os.environ.get( + "VLLM_OMNI_LINGBOT_WORLD_V2_CHECKPOINT_PATH", + "robbyant/lingbot-world-v2-14b-causal-fast-diffusers", +) +# The pipeline resizes the first frame to the requested geometry, so any real +# photograph conditions the rollout; the shared vLLM asset keeps the test +# runnable with no environment setup and adds no binary to this repository. +_IMAGE_ASSET = "2560px-Gfp-wisconsin-madison-the-nature-boardwalk" +_IMAGE_PATH_OVERRIDE = os.environ.get("VLLM_OMNI_LINGBOT_WORLD_V2_IMAGE_PATH") + +_HEIGHT = 480 +_WIDTH = 832 +_FRAMES_PER_BLOCK = 3 +_TEMPORAL_COMPRESSION = 4 +_REQUEST_ID = "lingbot-world-v2-stepwise-e2e" +# One three-frame action list per generated chunk. +_CAMERA_ACTION_SCRIPT = [ + [["w"], ["w"], ["w"]], + [["a"], [], []], + [[], [], []], +] +_NUM_CHUNKS = len(_CAMERA_ACTION_SCRIPT) +_NUM_FRAMES = (_NUM_CHUNKS * _FRAMES_PER_BLOCK - 1) * _TEMPORAL_COMPRESSION + 1 + +pytestmark = [ + pytest.mark.slow, + pytest.mark.diffusion, +] + + +def _first_frame() -> Path: + if _IMAGE_PATH_OVERRIDE: + return Path(_IMAGE_PATH_OVERRIDE).expanduser().resolve() + return Path(ImageAsset(_IMAGE_ASSET).get_path("jpg")) + + +def _chunk_metadata(output: OmniRequestOutput) -> dict[str, Any]: + multimodal = getattr(output, "multimodal_output", None) or {} + metadata = multimodal.get("metadata") if isinstance(multimodal, dict) else None + assert isinstance(metadata, dict), "streamed chunk is missing its metadata envelope" + ar_diffusion = metadata.get("ar_diffusion") + assert isinstance(ar_diffusion, dict), "streamed chunk is missing ar_diffusion metadata" + return dict(ar_diffusion) + + +async def _stream_chunks(image: Path) -> list[tuple[torch.Tensor, dict[str, Any], bool]]: + from vllm_omni.entrypoints.async_omni import AsyncOmni + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + engine = AsyncOmni( + model=MODEL, + engine_backend="vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine", + enforce_eager=True, + tensor_parallel_size=1, + max_num_seqs=1, + step_execution=True, + diffusion_streaming_output=True, + model_config={ + "ar_diffusion_height": _HEIGHT, + "ar_diffusion_width": _WIDTH, + "ar_diffusion_kv_config": {"gpu_memory_fraction": 0.6, "warmup_cudagraph": True}, + }, + ) + sampling = OmniDiffusionSamplingParams( + height=_HEIGHT, + width=_WIDTH, + num_frames=_NUM_FRAMES, + num_inference_steps=4, + max_sequence_length=512, + seed=42, + output_type="latent", + extra_args={"flow_shift": 5.0, "camera_action_script": _CAMERA_ACTION_SCRIPT}, + ) + prompt = { + "prompt": "The camera moves slowly forward through the scene.", + "multi_modal_data": {"image": str(image)}, + } + + chunks: list[tuple[torch.Tensor, dict[str, Any], bool]] = [] + try: + async for output in engine.generate(prompt, sampling, request_id=_REQUEST_ID): + images = getattr(output, "images", None) + if not images: + continue + assert len(images) == 1, "each AR block yields exactly one latent tensor" + latent = images[0].detach().float().cpu() + chunks.append((latent, _chunk_metadata(output), bool(getattr(output, "finished", False)))) + finally: + engine.shutdown() + return chunks + + +@hardware_test(res={"cuda": "H100"}, num_cards=1) +def test_lingbot_world_v2_stepwise_streams_one_chunk_per_block() -> None: + """Load the real checkpoint and stream N chunks from a single request.""" + + chunks = asyncio.run(_stream_chunks(_first_frame())) + + assert len(chunks) == _NUM_CHUNKS + expected_shape = ( + 1, + 16, + _FRAMES_PER_BLOCK, + _HEIGHT // 8, + _WIDTH // 8, + ) + for chunk_index, (latent, metadata, finished) in enumerate(chunks): + assert tuple(latent.shape) == expected_shape + assert torch.isfinite(latent).all() + # The request is the session, so both ids must agree on every chunk, + # and chunk indices must be contiguous from zero. + assert metadata["chunk_index"] == chunk_index + assert metadata["session_id"] == metadata["request_id"] + assert metadata["applied_event_ids"] == [] + assert finished is (chunk_index == _NUM_CHUNKS - 1) + + # The engine suffixes the caller's id to keep it unique; what this path + # guarantees is that one rollout keeps one identity from start to finish. + request_ids = {metadata["request_id"] for _, metadata, _ in chunks} + assert len(request_ids) == 1 + assert request_ids.pop().startswith(_REQUEST_ID) + + concatenated = torch.cat([latent for latent, _, _ in chunks], dim=2) + assert concatenated.shape[2] == _NUM_CHUNKS * _FRAMES_PER_BLOCK diff --git a/tests/e2e/online_serving/test_lingbot_world_v2_stepwise.py b/tests/e2e/online_serving/test_lingbot_world_v2_stepwise.py new file mode 100644 index 000000000..0f078049f --- /dev/null +++ b/tests/e2e/online_serving/test_lingbot_world_v2_stepwise.py @@ -0,0 +1,144 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""End-to-end streaming serving smoke test for LingBot-World v2. + +One ``WS /v1/realtime/video`` session must drive the whole rollout and return +one video chunk per AR block. This is the entry point the recipe documents; +``tests/e2e/offline_inference/test_lingbot_world_v2_stepwise.py`` covers the +same step-execution contract at the latent level, below the transport. +""" + +from __future__ import annotations + +import os + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + +import asyncio +import base64 +import json +import mimetypes +from pathlib import Path +from typing import Any + +import pytest +from vllm.assets.image import ImageAsset + +from tests.helpers.mark import hardware_test +from tests.helpers.runtime import OmniServerParams +from tests.helpers.stage_config import get_deploy_config_path + +MODEL = os.environ.get( + "VLLM_OMNI_LINGBOT_WORLD_V2_CHECKPOINT_PATH", + "robbyant/lingbot-world-v2-14b-causal-fast-diffusers", +) +# The pipeline resizes the first frame to the requested geometry, so any real +# photograph conditions the rollout; the shared vLLM asset keeps the test +# runnable with no environment setup and adds no binary to this repository. +_IMAGE_ASSET = "2560px-Gfp-wisconsin-madison-the-nature-boardwalk" +_IMAGE_PATH_OVERRIDE = os.environ.get("VLLM_OMNI_LINGBOT_WORLD_V2_IMAGE_PATH") + +# The AR cache geometry is fixed at load time, so a request must ask for +# exactly the resolution the deploy config declares. +_HEIGHT = 480 +_WIDTH = 832 +_FRAMES_PER_BLOCK = 3 +_TEMPORAL_COMPRESSION = 4 +# One three-frame action list per generated chunk. +_CAMERA_ACTION_SCRIPT = [ + [["w"], ["w"], ["w"]], + [["a"], [], []], + [[], [], []], +] +_NUM_CHUNKS = len(_CAMERA_ACTION_SCRIPT) +_NUM_FRAMES = (_NUM_CHUNKS * _FRAMES_PER_BLOCK - 1) * _TEMPORAL_COMPRESSION + 1 +_PROMPT = "The camera moves slowly forward through the scene." + +pytestmark = [ + pytest.mark.slow, + pytest.mark.diffusion, +] + + +def _first_frame() -> Path: + if _IMAGE_PATH_OVERRIDE: + return Path(_IMAGE_PATH_OVERRIDE).expanduser().resolve() + return Path(ImageAsset(_IMAGE_ASSET).get_path("jpg")) + + +lingbot_world_v2_stepwise_server_params = [ + pytest.param( + OmniServerParams( + model=MODEL, + stage_config_path=get_deploy_config_path("lingbot_world_v2_stepwise.yaml"), + ), + id="lingbot_world_v2_stepwise", + ) +] + + +def _session_start(model: str, image: Path) -> dict[str, Any]: + media_type = mimetypes.guess_type(image.name)[0] or "image/png" + encoded = base64.b64encode(image.read_bytes()).decode() + return { + "type": "session.start", + "model": model, + "prompt": _PROMPT, + # LingBot is image-conditioned: without a first frame the request is + # rejected during pre-processing. + "image_reference": {"image_url": f"data:{media_type};base64,{encoded}"}, + "width": _WIDTH, + "height": _HEIGHT, + "num_frames": _NUM_FRAMES, + "fps": 16, + "seed": 42, + "extra_params": {"flow_shift": 5.0, "camera_action_script": _CAMERA_ACTION_SCRIPT}, + } + + +async def _stream_session(url: str, model: str, image: Path) -> tuple[list[dict[str, Any]], bytes, dict[str, Any]]: + import websockets + + events: list[dict[str, Any]] = [] + media = bytearray() + done: dict[str, Any] = {} + async with websockets.connect(url, max_size=None, ping_interval=None) as websocket: + await websocket.send(json.dumps(_session_start(model, image))) + while True: + message = await asyncio.wait_for(websocket.recv(), timeout=900) + if isinstance(message, bytes): + media.extend(message) + continue + event = json.loads(message) + events.append(event) + if event["type"] == "error": + pytest.fail(f"server returned an error event: {event}") + if event["type"] == "session.done": + done = event + break + return events, bytes(media), done + + +@hardware_test(res={"cuda": "H100"}, num_cards=1) +@pytest.mark.parametrize("omni_server", lingbot_world_v2_stepwise_server_params, indirect=True) +def test_lingbot_world_v2_stepwise_streams_one_chunk_per_block(omni_server) -> None: + """Serve the real checkpoint and stream N chunks from a single WS session.""" + + url = f"ws://{omni_server.host}:{omni_server.port}/v1/realtime/video" + events, media, done = asyncio.run(_stream_session(url, omni_server.model, _first_frame())) + + started = [event for event in events if event["type"] == "video.start"] + assert len(started) == 1, "one session.start must open exactly one rollout" + + # A trailer chunk may follow the media chunks; only media chunks carry a + # generation index, and there is exactly one per AR block. + media_chunks = [event for event in events if event["type"] == "video.chunk_metadata" and event["kind"] == "media"] + assert [event["generation_chunk_index"] for event in media_chunks] == list(range(_NUM_CHUNKS)) + for event in media_chunks: + assert event["request_id"] == started[0]["request_id"] + assert event["num_frames"] > 0 + assert event["byte_length"] > 0 + + assert len(media) == sum(event["byte_length"] for event in events if event["type"] == "video.chunk_metadata") + assert done["stopped"] is False + assert done["chunks"] >= _NUM_CHUNKS diff --git a/tests/engine/test_diffusion_streaming_output_kwarg.py b/tests/engine/test_diffusion_streaming_output_kwarg.py new file mode 100644 index 000000000..a6f3dff02 --- /dev/null +++ b/tests/engine/test_diffusion_streaming_output_kwarg.py @@ -0,0 +1,79 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""``diffusion_streaming_output`` must reach a diffusion stage as ``streaming_output``. + +The public AsyncOmni/serve kwarg and the stage schema field have different +names, and the two config paths translate it in different places: the +unregistered single-stage fallback writes it into the default stage config, +while a registered pipeline resolves through ``StageConfigFactory``, which drops +keys the stage schema does not know. Without the mirror in +``_resolve_stage_configs`` a registered pipeline (LingBot World stepwise) serves +only its final chunk instead of streaming one per step. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any + +import pytest + +from vllm_omni.engine.async_omni_engine import AsyncOmniEngine + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _resolve(monkeypatch, kwargs: dict[str, Any]) -> dict[str, Any]: + """Run ``_resolve_stage_configs`` and return the kwargs it forwarded.""" + + seen: dict[str, Any] = {} + + def _fake_load(model, forwarded_kwargs, **_: Any): + seen.update(forwarded_kwargs) + return "config.yaml", [], None + + monkeypatch.setattr( + "vllm_omni.engine.async_omni_engine.load_and_resolve_stage_configs", + _fake_load, + ) + engine = SimpleNamespace( + _create_default_diffusion_stage_cfg=AsyncOmniEngine._create_default_diffusion_stage_cfg, + _apply_strategy_lb_policy=lambda *_args, **_kw: None, + ) + AsyncOmniEngine._resolve_stage_configs( + engine, + "some/model", + kwargs, + trust_remote_code=None, + ) + return seen + + +def test_registered_pipeline_receives_streaming_output(monkeypatch) -> None: + forwarded = _resolve(monkeypatch, {"diffusion_streaming_output": True}) + + assert forwarded["streaming_output"] is True + + +def test_disabled_kwarg_does_not_override_the_deploy_config(monkeypatch) -> None: + # ``serve`` always carries the flag's ``False`` default; mirroring it would + # silently turn streaming off for a deploy YAML that asks for it. + forwarded = _resolve(monkeypatch, {"diffusion_streaming_output": False}) + + assert "streaming_output" not in forwarded + + +def test_explicit_streaming_output_wins(monkeypatch) -> None: + forwarded = _resolve( + monkeypatch, + {"diffusion_streaming_output": True, "streaming_output": False}, + ) + + assert forwarded["streaming_output"] is False + + +@pytest.mark.parametrize("enabled", [True, False]) +def test_unregistered_fallback_stage_carries_the_flag(enabled: bool) -> None: + stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg({"diffusion_streaming_output": enabled})[0] + + assert stage_cfg["engine_args"]["streaming_output"] is enabled diff --git a/vllm_omni/config/pipeline_registry.py b/vllm_omni/config/pipeline_registry.py index 83f8de383..5ac47dddc 100644 --- a/vllm_omni/config/pipeline_registry.py +++ b/vllm_omni/config/pipeline_registry.py @@ -69,6 +69,7 @@ from vllm_omni.model_executor.models.indextts2.pipeline import ( INDEXTTS25_PIPELINE, ) from vllm_omni.model_executor.models.lance.pipeline import LANCE_PIPELINE +from vllm_omni.model_executor.models.lingbot_world.pipeline import LINGBOT_WORLD_PIPELINE from vllm_omni.model_executor.models.mammoth_moda2.pipeline import ( MAMMOTH_MODA2_AR_PIPELINE, MAMMOTH_MODA2_PIPELINE, @@ -141,6 +142,7 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = { "bagel_single_stage": BAGEL_SINGLE_STAGE_PIPELINE, "lance": LANCE_PIPELINE, "dreamzero": DREAMZERO_PIPELINE, + "lingbot_world": LINGBOT_WORLD_PIPELINE, "Gr00tN1d7": GR00T_N1D7_PIPELINE, "pi0": PI0_PIPELINE, "gepard": GEPARD_PIPELINE, diff --git a/vllm_omni/deploy/lingbot_world_v2_stepwise.yaml b/vllm_omni/deploy/lingbot_world_v2_stepwise.yaml new file mode 100644 index 000000000..3fc950e4c --- /dev/null +++ b/vllm_omni/deploy/lingbot_world_v2_stepwise.yaml @@ -0,0 +1,41 @@ +# LingBot-World 2.0 stepwise deploy: single diffusion stage for WS /v1/realtime/video. +# +# One AsyncOmni request is one rollout: prepare_encode runs once, then every +# four DMD steps stream one AR block. Step execution keeps AR-Diffusion paged KV +# on ARDiffusionModelRunner, so this stage must stay at max_num_seqs=1. +# +# The cache geometry is fixed at load time, so ar_diffusion_height/width must +# match the height/width a client asks for. + +pipeline: lingbot_world +async_chunk: false +distributed_executor_backend: mp +dtype: bfloat16 + +stages: + - stage_id: 0 + devices: "0" + max_num_seqs: 1 + enforce_eager: true + model_class_name: LingBotWorldCausalDMDPipeline + # LingBot runs on the AR-Diffusion engine (engine-level paged KV cache). + engine_backend: vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine + # One request, many streamed chunks. streaming_output implies step_execution, + # but both are stated here so the topology is readable on its own. + step_execution: true + streaming_output: true + # num_inference_steps is for the step scheduler, not the pipeline: it sizes a + # request's step budget before the pipeline is reached, so a served request + # does not have to carry it. The pipeline itself pins the causal-fast + # checkpoint's four-step DMD schedule and rejects any other count. + default_sampling_params: + num_inference_steps: 4 + output_type: pil + extra_args: + flow_shift: 5.0 + model_config: + ar_diffusion_height: 480 + ar_diffusion_width: 832 + ar_diffusion_kv_config: + gpu_memory_fraction: 0.6 + warmup_cudagraph: true diff --git a/vllm_omni/diffusion/models/lingbot_world/actions.py b/vllm_omni/diffusion/models/lingbot_world/actions.py index e725c698a..4720b041e 100644 --- a/vllm_omni/diffusion/models/lingbot_world/actions.py +++ b/vllm_omni/diffusion/models/lingbot_world/actions.py @@ -1,13 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Realtime keyboard controls for LingBot World 2.0.""" from __future__ import annotations import math -from collections.abc import Mapping, Sequence +from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass -from typing import Any +from typing import Any, TypeAlias import numpy as np import torch @@ -42,7 +42,31 @@ def _normalize_actions(value: object, *, field: str) -> tuple[str, ...]: return tuple(action for action in _ACTION_ORDER if action in actions) -def _normalize_frames(value: object, *, field: str) -> tuple[tuple[str, ...], ...]: +#: Key states for one latent frame each, e.g. ``(("w",), ("w", "j"), ())`` for +#: the three latent frames of one AR block. +LingBotCameraActionFrames: TypeAlias = tuple[tuple[str, ...], ...] + +#: One :data:`LingBotCameraActionFrames` per generated chunk, carried on the +#: request for the whole rollout. +LingBotCameraActionScript: TypeAlias = tuple[LingBotCameraActionFrames, ...] + + +def as_camera_action_frames(value: Iterable[Iterable[str]]) -> LingBotCameraActionFrames: + """Restore the tuple form of one chunk's per-latent-frame key states.""" + return tuple(tuple(frame) for frame in value) + + +def as_camera_action_script(value: Iterable[Iterable[Iterable[str]]]) -> LingBotCameraActionScript: + """Restore the tuple form of a whole request's script. + + ``sampling_params.extra_args`` round-trips through JSON, so an already + validated script comes back as lists; this rebuilds the hashable tuple form + without re-running validation. + """ + return tuple(as_camera_action_frames(chunk) for chunk in value) + + +def _normalize_frames(value: object, *, field: str) -> LingBotCameraActionFrames: if not isinstance(value, Sequence) or isinstance(value, (str, bytes)): raise ValueError(f"{field} must be a sequence of per-frame action lists.") return tuple(_normalize_actions(actions, field=f"{field}[{index}]") for index, actions in enumerate(value)) @@ -52,7 +76,7 @@ def parse_lingbot_camera_action_frames( data: Mapping[str, Any], *, expected_frames: int, -) -> tuple[tuple[str, ...], ...]: +) -> LingBotCameraActionFrames: """Validate the chunk-sized payload emitted by the session reducer.""" if data.get("mode") != "frames": @@ -66,6 +90,30 @@ def parse_lingbot_camera_action_frames( return frames +def parse_lingbot_camera_action_script( + script: object, + *, + frames_per_chunk: int, +) -> LingBotCameraActionScript: + """Validate a request-scoped list of per-chunk camera action frames.""" + + if not isinstance(script, Sequence) or isinstance(script, (str, bytes)): + raise ValueError("camera_action_script must be a sequence of per-chunk action lists.") + if not script: + raise ValueError("camera_action_script must contain at least one chunk.") + chunks: list[LingBotCameraActionFrames] = [] + for index, chunk in enumerate(script): + frames = _normalize_frames(chunk, field=f"camera_action_script[{index}]") + if len(frames) != frames_per_chunk: + raise ValueError( + "camera_action_script chunks must contain exactly " + f"{frames_per_chunk} per-latent-frame action lists; " + f"chunk {index} has {len(frames)}." + ) + chunks.append(frames) + return tuple(chunks) + + @dataclass(frozen=True) class _LingBotCameraReducerState: mode: str | None = None diff --git a/vllm_omni/diffusion/models/lingbot_world/dmd_block.py b/vllm_omni/diffusion/models/lingbot_world/dmd_block.py new file mode 100644 index 000000000..eace73a41 --- /dev/null +++ b/vllm_omni/diffusion/models/lingbot_world/dmd_block.py @@ -0,0 +1,239 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""One AR block of LingBot World's causal DMD sampler. + +The pipeline drives this in two ways -- request mode loops ``generate_block`` +over every block, stepwise execution calls ``probe_step`` / ``apply_transition`` +one denoise step at a time and ``commit_block_kv`` from ``post_decode`` -- so the +math lives here once and cannot drift between the two modes. + +A block is four probes that must not touch KV, then one clean-x0 commit that +does. Whether a call sees paged (session-bound) or request-local KV is decided +by the ``ar`` argument, so this class never reaches back into the pipeline. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, cast + +import torch +from diffusers.utils.torch_utils import randn_tensor + +from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx +from vllm_omni.diffusion.models.lingbot_world.transformer import ( + LingBotAttentionCache, + LingBotTransformerCache, +) + +if TYPE_CHECKING: + from vllm_omni.diffusion.models.lingbot_world.transformer import CausalLingBotWorldTransformer3DModel + from vllm_omni.diffusion.models.progress_bar import TqdmProgressBar + from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState + + +@dataclass(frozen=True) +class ARBlockContext: + """The bound AR-Diffusion session one block runs against. + + Built by the pipeline from the state the runner bound for this invocation; + ``None`` in place of a context means request-local contiguous KV. + """ + + state: ARDiffusionKVState + cross_attention: list[LingBotAttentionCache] + branch: str + + +class LingBotDMDBlockRunner: + """Runs the DMD probes and the KV commit for one latent-frame block.""" + + def __init__( + self, + transformer: CausalLingBotWorldTransformer3DModel, + *, + device: torch.device, + enforce_eager: bool, + ) -> None: + self.transformer = transformer + self.device = device + self.enforce_eager = enforce_eager + + # ── cache selection ────────────────────────────────────────────────── + + def _paged_cache( + self, + *, + condition: torch.Tensor, + ar: ARBlockContext, + commit_current: bool, + ) -> LingBotTransformerCache: + patch_frames, patch_height, patch_width = self.transformer.config.patch_size + frames = condition.shape[2] // patch_frames + height = condition.shape[3] // patch_height + width = condition.shape[4] // patch_width + return LingBotTransformerCache( + self_attention=ar.state.get_kv_caches( + ar.branch, + seq_len=frames * height * width, + commit_current=commit_current, + ), + cross_attention=ar.cross_attention, + ) + + def _block_cache( + self, + *, + condition: torch.Tensor, + cache: LingBotTransformerCache | None, + ar: ARBlockContext | None, + commit_current: bool, + ) -> LingBotTransformerCache: + """Pick the cache one transformer call sees: paged when a session is + bound, else the caller's request-local cache.""" + if ar is not None: + return self._paged_cache(condition=condition, ar=ar, commit_current=commit_current) + return cast(LingBotTransformerCache, cache) + + # ── one block ──────────────────────────────────────────────────────── + + def probe_step( + self, + *, + current_latents: torch.Tensor, + condition: torch.Tensor, + camera: torch.Tensor, + prompt_embeds: torch.Tensor, + cache: LingBotTransformerCache | None, + ar: ARBlockContext | None, + start_frame: int, + timestep_value: float, + step_index: int, + ) -> torch.Tensor: + """Predict flow for one denoise step. + + A probe never writes KV: only the clean x0 of a finished block may + enter the cache, which is ``commit_block_kv``'s job. + """ + if not self.enforce_eager: + torch.compiler.cudagraph_mark_step_begin() + set_forward_context_denoise_step_idx(step_index) + timestep = torch.full((1,), float(timestep_value), device=self.device, dtype=torch.float32) + # Checkpoint channel contract: + # [noise/x_t(16), temporal_mask(4), image_latent(16)] -> 36. + model_input = torch.cat((current_latents.to(dtype=condition.dtype), condition), dim=1) + flow_prediction = self.transformer( + hidden_states=model_input, + timestep=timestep, + encoder_hidden_states=prompt_embeds, + camera_hidden_states=camera, + cache=self._block_cache(condition=condition, cache=cache, ar=ar, commit_current=False), + start_frame=start_frame, + update_cache=False, + ) + if flow_prediction.shape != current_latents.shape: + raise RuntimeError( + "transformer flow prediction shape must match the 16-channel noise latent, " + f"got {tuple(flow_prediction.shape)} and {tuple(current_latents.shape)}." + ) + return flow_prediction + + def apply_transition( + self, + current_latents: torch.Tensor, + flow_prediction: torch.Tensor, + sigma: float, + *, + next_sigma: float | None, + generator: torch.Generator, + ) -> torch.Tensor: + """Invert the flow to x0, then re-noise at ``next_sigma``. + + ``next_sigma=None`` marks the final step, so where the caller sits in + its own loop stays out of this function. + """ + # The checkpoint's flow parameterization is inverted by + # x0 = x_t - sigma * flow. Intermediate steps re-noise that x0 + # estimate at the next sigma; the final step keeps x0 as this + # block's generated latent. + x0 = current_latents - sigma * flow_prediction.float() + if next_sigma is None: + return x0 + noise = randn_tensor(current_latents.shape, generator=generator, device=self.device, dtype=torch.float32) + return (1.0 - next_sigma) * x0 + next_sigma * noise + + def commit_block_kv( + self, + *, + latents: torch.Tensor, + condition: torch.Tensor, + camera: torch.Tensor, + prompt_embeds: torch.Tensor, + cache: LingBotTransformerCache | None, + ar: ARBlockContext | None, + start_frame: int, + ) -> None: + """Write the finished block's clean x0 into KV and commit its pages. + + The fifth transformer call of a block, deliberately not a denoise step: + on the stepwise path it belongs to ``post_decode()``. + """ + # Commit K/V only for the final clean block, never for noisy probes. + cache_input = torch.cat((latents.to(dtype=condition.dtype), condition), dim=1) + if not self.enforce_eager: + torch.compiler.cudagraph_mark_step_begin() + self.transformer( + hidden_states=cache_input, + timestep=torch.zeros(1, device=self.device, dtype=torch.float32), + encoder_hidden_states=prompt_embeds, + camera_hidden_states=camera, + cache=self._block_cache(condition=condition, cache=cache, ar=ar, commit_current=True), + start_frame=start_frame, + update_cache=True, + ) + if ar is not None: + ar.state.commit_paged_context(ar.branch) + + def generate_block( + self, + *, + condition: torch.Tensor, + camera: torch.Tensor, + prompt_embeds: torch.Tensor, + cache: LingBotTransformerCache | None, + ar: ARBlockContext | None, + start_frame: int, + schedule: tuple[tuple[float, float], ...], + generator: torch.Generator, + progress_bar: TqdmProgressBar[Any], + ) -> torch.Tensor: + """Request mode: all probes of one block, then its commit.""" + block_shape = (1, self.transformer.config.out_channels, *condition.shape[2:5]) + current_latents = randn_tensor(block_shape, generator=generator, device=self.device, dtype=torch.float32) + for step_index, (timestep_value, sigma) in enumerate(schedule): + flow_prediction = self.probe_step( + current_latents=current_latents, + condition=condition, + camera=camera, + prompt_embeds=prompt_embeds, + cache=cache, + ar=ar, + start_frame=start_frame, + timestep_value=timestep_value, + step_index=step_index, + ) + next_sigma = schedule[step_index + 1][1] if step_index + 1 < len(schedule) else None + current_latents = self.apply_transition( + current_latents, flow_prediction, sigma, next_sigma=next_sigma, generator=generator + ) + progress_bar.update() + self.commit_block_kv( + latents=current_latents, + condition=condition, + camera=camera, + prompt_embeds=prompt_embeds, + cache=cache, + ar=ar, + start_frame=start_frame, + ) + return current_latents diff --git a/vllm_omni/diffusion/models/lingbot_world/pipeline.py b/vllm_omni/diffusion/models/lingbot_world/pipeline.py index 5f71979e8..b41abd861 100644 --- a/vllm_omni/diffusion/models/lingbot_world/pipeline.py +++ b/vllm_omni/diffusion/models/lingbot_world/pipeline.py @@ -1,12 +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 """Request-scoped LingBot-World v2 causal DMD pipeline.""" from __future__ import annotations import math import os -from collections.abc import Callable, Iterable, Iterator +from collections.abc import Callable, Iterable, Iterator, Sequence from contextlib import contextmanager from dataclasses import dataclass, replace from typing import TYPE_CHECKING, Any, ClassVar, cast @@ -24,15 +24,19 @@ from vllm.model_executor.models.utils import AutoWeightsLoader from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import DistributedAutoencoderKLWan from vllm_omni.diffusion.distributed.utils import get_local_device -from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders -from vllm_omni.diffusion.models.interface import SupportImageInput, SupportsComponentDiscovery +from vllm_omni.diffusion.models.interface import SupportImageInput, SupportsComponentDiscovery, SupportsStepExecution from vllm_omni.diffusion.models.lingbot_world.actions import ( LINGBOT_CAMERA_ACTION_SCHEMA, LINGBOT_CAMERA_TRAJECTORY_SCHEMA, + LingBotCameraActionFrames, + LingBotCameraActionScript, + as_camera_action_frames, + as_camera_action_script, integrate_lingbot_camera_actions, parse_lingbot_camera_action_frames, + parse_lingbot_camera_action_script, ) from vllm_omni.diffusion.models.lingbot_world.camera import ( CameraTrajectory, @@ -41,6 +45,7 @@ from vllm_omni.diffusion.models.lingbot_world.camera import ( load_camera_trajectory, resolve_trusted_action_directory, ) +from vllm_omni.diffusion.models.lingbot_world.dmd_block import ARBlockContext, LingBotDMDBlockRunner from vllm_omni.diffusion.models.lingbot_world.transformer import ( CausalLingBotWorldTransformer3DModel, LingBotAttentionCache, @@ -64,8 +69,11 @@ from vllm_omni.experimental.ar_diffusion.tick_protocol import ( ) if TYPE_CHECKING: + from diffusers.video_processor import VideoProcessor from tqdm.std import tqdm as TqdmProgressBar + from vllm_omni.diffusion.worker.input_batch import InputBatch + from vllm_omni.diffusion.worker.utils import StepRequestState from vllm_omni.experimental.ar_diffusion.kv_cache.state import ( ARDiffusionKVState, ) @@ -79,6 +87,7 @@ _MAX_SEQUENCE_LENGTH = 512 _ACTION_ROOT_ENV = "VLLM_OMNI_LINGBOT_ACTION_ROOT" _PREPROCESSED_CAMERA_KEY = "_lingbot_camera_trajectory" _PREPROCESSED_CAMERA_ACTIONS_KEY = "_lingbot_camera_actions" +_PREPROCESSED_CAMERA_ACTION_SCRIPT_KEY = "_lingbot_camera_action_script" _SOURCE_IMAGE_ERROR = ( "Unable to load multi_modal_data.image; expected a decodable image within 4096 * 4096 source pixels." ) @@ -86,12 +95,34 @@ _SOURCE_IMAGE_ERROR = ( @dataclass(frozen=True) class _LingBotRequestInputs: - """Validated model-specific request boundary.""" + """Validated model-specific request boundary. + + Pre-processing materializes exactly one of the three camera fields, chosen + from the shape of the incoming control. They are not alternative spellings + of one input: ``camera_action_script`` is a *container* whose element is a + ``camera_actions``, and the two are live at once during step execution. + + - ``camera_trajectory`` is explicit per-frame poses and intrinsics. Offline + replay loads it from ``extra_args["action_path"]``; the tick entry point + receives the same content as a ``lingbot.camera_trajectory.v1`` control. + Integrating WASD actions also produces one, written back onto this field. + - ``camera_actions`` is WASD key state for *one* three-latent-frame block: + the unit every path reduces to before the camera is embedded. The tick + entry point fills it directly, because one of its requests is one block; + step execution derives it per block by slicing ``camera_action_script``. + ``_prepare_camera`` also reads it as the signal that the trajectory it + was handed is already latent-aligned. + - ``camera_action_script`` holds one such action list *per generated chunk* + for the whole rollout, which only step execution needs, because only it + runs every block inside a single ``generate()``. The tick entry point has + no use for it and ``forward()`` rejects a request that carries it. + """ prompt: str image: PIL.Image.Image | torch.Tensor camera_trajectory: CameraTrajectory | None - camera_actions: tuple[tuple[str, ...], ...] | None + camera_actions: LingBotCameraActionFrames | None + camera_action_script: LingBotCameraActionScript | None height: int width: int num_frames: int @@ -303,6 +334,16 @@ def get_lingbot_world_pre_process_func( "LingBot realtime camera controls require schema " "'lingbot.camera_trajectory.v1' or 'lingbot.camera_actions.v1'." ) + camera_action_script = None + elif extra_args.get("camera_action_script") is not None: + if extra_args.get("action_path"): + raise ValueError("camera_action_script cannot be combined with action_path.") + camera_action_script = parse_lingbot_camera_action_script( + extra_args.get("camera_action_script"), + frames_per_chunk=3, + ) + trajectory = None + camera_actions = None else: action_path = extra_args.get("action_path") if not isinstance(action_path, (str, os.PathLike)) or not str(action_path): @@ -323,6 +364,7 @@ def get_lingbot_world_pre_process_func( "Unable to load camera trajectory from action_path; expected poses.npy and intrinsics.npy." ) from None camera_actions = None + camera_action_script = None updated_prompt = dict(prompt) updated_multi_modal_data = dict(multi_modal_data) @@ -333,6 +375,7 @@ def get_lingbot_world_pre_process_func( **extra_args, _PREPROCESSED_CAMERA_KEY: trajectory, _PREPROCESSED_CAMERA_ACTIONS_KEY: camera_actions, + _PREPROCESSED_CAMERA_ACTION_SCRIPT_KEY: camera_action_script, } return request @@ -377,11 +420,13 @@ class LingBotWorldCausalDMDPipeline( nn.Module, SupportImageInput, SupportsComponentDiscovery, + SupportsStepExecution, ProgressBarMixin, DiffusionPipelineProfilerMixin, ): """LingBot-World v2 I2V generation with a request-local causal cache.""" + supports_step_execution: ClassVar[bool] = True _dit_modules: ClassVar[list[str]] = ["transformer"] _encoder_modules: ClassVar[list[str]] = ["text_encoder"] _vae_modules: ClassVar[list[str]] = ["vae"] @@ -617,11 +662,17 @@ class LingBotWorldCausalDMDPipeline( raise ValueError("sampling_params.extra_args must be a mapping.") camera_trajectory = extra_args.get(_PREPROCESSED_CAMERA_KEY) camera_actions = extra_args.get(_PREPROCESSED_CAMERA_ACTIONS_KEY) + camera_action_script = extra_args.get(_PREPROCESSED_CAMERA_ACTION_SCRIPT_KEY) if camera_trajectory is not None and not isinstance(camera_trajectory, CameraTrajectory): raise ValueError("LingBot camera trajectory must be materialized by the pre-process function.") if camera_actions is not None: - camera_actions = tuple(tuple(actions) for actions in camera_actions) - if (camera_trajectory is None) == (camera_actions is None): + camera_actions = as_camera_action_frames(camera_actions) + if camera_action_script is not None: + camera_action_script = as_camera_action_script(camera_action_script) + camera_sources = ( + int(camera_trajectory is not None) + int(camera_actions is not None) + int(camera_action_script is not None) + ) + if camera_sources != 1: raise ValueError("LingBot pre-processing must materialize exactly one camera input.") request_flow_shift = ( @@ -720,6 +771,7 @@ class LingBotWorldCausalDMDPipeline( image=image, camera_trajectory=camera_trajectory, camera_actions=camera_actions, + camera_action_script=camera_action_script, height=height, width=width, num_frames=num_frames, @@ -818,6 +870,7 @@ class LingBotWorldCausalDMDPipeline( *, dtype: torch.dtype, previous: CameraTrajectory | None = None, + latent_aligned: bool = False, ) -> tuple[torch.Tensor, CameraTrajectory]: """Convert raw camera frames to a latent-aligned ray tensor.""" @@ -825,7 +878,7 @@ class LingBotWorldCausalDMDPipeline( if trajectory is None: raise RuntimeError("LingBot camera trajectory was not prepared before camera embedding.") available_frames = int(trajectory.poses.shape[0]) - if inputs.camera_actions is not None: + if inputs.camera_actions is not None or latent_aligned: if available_frames != inputs.num_latent_frames: raise ValueError( "camera actions must produce exactly one pose per latent frame; " @@ -942,29 +995,38 @@ class LingBotWorldCausalDMDPipeline( ) ] - def _ar_transformer_cache( + # ── DMD block helpers ───────────────────────────────────────────────── + # + # The DMD math for one AR block lives in ``LingBotDMDBlockRunner`` so that + # request mode (``_generate_block`` in a loop) and stepwise execution + # (``denoise_step`` / ``step_scheduler`` / ``post_decode``) share one copy. + # The pipeline's only jobs here are building the runner and translating + # the bound AR session into an ``ARBlockContext``. + + @property + def _dmd_blocks(self) -> LingBotDMDBlockRunner: + runner = getattr(self, "_dmd_block_runner", None) + if runner is None: + runner = LingBotDMDBlockRunner( + self.transformer, + device=self.device, + enforce_eager=bool(self.od_config.enforce_eager), + ) + self._dmd_block_runner = runner + return runner + + def _ar_block_context( self, - *, - condition: torch.Tensor, - cross_attention: list[LingBotAttentionCache], - commit_current: bool, - ) -> LingBotTransformerCache: + cross_attention: list[LingBotAttentionCache] | None, + ) -> ARBlockContext | None: + """Wrap the bound session for the block runner; ``None`` keeps + request-local KV.""" + if cross_attention is None: + return None state = self._ar_diffusion_kv_state if state is None: raise RuntimeError("LingBot AR cache requested without a bound state.") - patch_frames, patch_height, patch_width = self.transformer.config.patch_size - frames = condition.shape[2] // patch_frames - height = condition.shape[3] // patch_height - width = condition.shape[4] // patch_width - seq_len = frames * height * width - return LingBotTransformerCache( - self_attention=state.get_kv_caches( - self._AR_BRANCH, - seq_len=seq_len, - commit_current=commit_current, - ), - cross_attention=cross_attention, - ) + return ARBlockContext(state=state, cross_attention=cross_attention, branch=self._AR_BRANCH) def _generate_block( self, @@ -979,102 +1041,17 @@ class LingBotWorldCausalDMDPipeline( generator: torch.Generator, progress_bar: TqdmProgressBar[Any], ) -> torch.Tensor: - block_shape = ( - 1, - self.transformer.config.out_channels, - condition.shape[2], - condition.shape[3], - condition.shape[4], - ) - current_latents = randn_tensor( - block_shape, - generator=generator, - device=self.device, - dtype=torch.float32, - ) - for step_index, (timestep_value, sigma) in enumerate(schedule): - if not self.od_config.enforce_eager: - torch.compiler.cudagraph_mark_step_begin() - set_forward_context_denoise_step_idx(step_index) - timestep = torch.full((1,), float(timestep_value), device=self.device, dtype=torch.float32) - # Checkpoint channel contract: - # [noise/x_t(16), temporal_mask(4), image_latent(16)] -> 36. - model_input = torch.cat((current_latents.to(dtype=condition.dtype), condition), dim=1) - current_cache = ( - self._ar_transformer_cache( - condition=condition, - cross_attention=cast( - list[LingBotAttentionCache], - ar_cross_attention, - ), - commit_current=False, - ) - if ar_cross_attention is not None - else cast(LingBotTransformerCache, cache) - ) - flow_prediction = self.transformer( - hidden_states=model_input, - timestep=timestep, - encoder_hidden_states=prompt_embeds, - camera_hidden_states=camera, - cache=current_cache, - start_frame=start_frame, - update_cache=False, - ) - if flow_prediction.shape != current_latents.shape: - raise RuntimeError( - "transformer flow prediction shape must match the 16-channel noise latent, " - f"got {tuple(flow_prediction.shape)} and {tuple(current_latents.shape)}." - ) - # The checkpoint's flow parameterization is inverted by - # x0 = x_t - sigma * flow. Intermediate steps re-noise that x0 - # estimate at the next sigma; the final step keeps x0 as this - # block's generated latent. - x0 = current_latents - sigma * flow_prediction.float() - if step_index + 1 < len(schedule): - next_sigma = schedule[step_index + 1][1] - noise = randn_tensor( - current_latents.shape, - generator=generator, - device=self.device, - dtype=torch.float32, - ) - current_latents = (1.0 - next_sigma) * x0 + next_sigma * noise - else: - current_latents = x0 - progress_bar.update() - - # Commit K/V only for the final clean block, never for noisy probes. - cache_input = torch.cat((current_latents.to(dtype=condition.dtype), condition), dim=1) - commit_cache = ( - self._ar_transformer_cache( - condition=condition, - cross_attention=cast( - list[LingBotAttentionCache], - ar_cross_attention, - ), - commit_current=True, - ) - if ar_cross_attention is not None - else cast(LingBotTransformerCache, cache) - ) - if not self.od_config.enforce_eager: - torch.compiler.cudagraph_mark_step_begin() - self.transformer( - hidden_states=cache_input, - timestep=torch.zeros(1, device=self.device, dtype=torch.float32), - encoder_hidden_states=prompt_embeds, - camera_hidden_states=camera, - cache=commit_cache, + return self._dmd_blocks.generate_block( + condition=condition, + camera=camera, + prompt_embeds=prompt_embeds, + cache=cache, + ar=self._ar_block_context(ar_cross_attention), start_frame=start_frame, - update_cache=True, + schedule=schedule, + generator=generator, + progress_bar=progress_bar, ) - if ar_cross_attention is not None: - state = self._ar_diffusion_kv_state - if state is None: - raise RuntimeError("LingBot AR state disappeared before KV commit.") - state.commit_paged_context(self._AR_BRANCH) - return current_latents def encode_prompt( self, @@ -1100,6 +1077,15 @@ class LingBotWorldCausalDMDPipeline( def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: inputs = self._parse_request(req) + if inputs.camera_action_script is not None: + # Only step execution walks a whole rollout inside one request, so + # a per-chunk script here would be silently dropped rather than + # steering the camera. + raise ValueError( + "camera_action_script is read only by LingBot step execution; " + "request mode takes one three-frame camera_actions control per " + "block, or a camera_trajectory for the request." + ) tick = ARDiffusionTickRequest.from_extra_args(req.sampling_params.extra_args) if tick is not None and self._ar_diffusion_kv_state is None: raise RuntimeError("LingBot typed ticks require ARDiffusionEngine session binding.") @@ -1287,6 +1273,317 @@ class LingBotWorldCausalDMDPipeline( stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, ) + def _decode_chunk_to_pixels( + self, latents: torch.Tensor, *, output_type: str + ) -> torch.Tensor | np.ndarray | list[list[PIL.Image.Image]]: + """Decode one AR block so streaming consumers receive pixels, not latents. + + Blocks are decoded independently, which keeps the shared VAE stateless + and lets a rollout be served without a per-session temporal cache. That + is the same shape Helios streams today; cross-block ``feat_cache`` + continuity is a separate decision and is not made here. + + The registered post-process hook returns a ``{"payload", "metadata"}`` + envelope untouched, so the pixel conversion it would do for a bare + tensor has to happen here instead. + """ + latent_mean, latent_std = self._vae_latent_stats(latents) + vae_latents = (latents * latent_std + latent_mean).to(dtype=self.vae.dtype) + video = self.vae.decode(vae_latents, return_dict=False)[0] + return self._video_processor().postprocess_video(video, output_type=output_type) + + def _video_processor(self) -> VideoProcessor: + processor = getattr(self, "_cached_video_processor", None) + if processor is None: + from diffusers.video_processor import VideoProcessor + + processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) + self._cached_video_processor = processor + return processor + + def _horizon_latent_frames(self) -> int: + return (_MAX_RAW_FRAMES - 1) // self.vae_scale_factor_temporal + 1 + + def _max_realtime_chunks(self) -> int: + return self._horizon_latent_frames() // int(self.transformer.config.num_frames_per_block) + + def _require_bound_ar_state(self) -> None: + if self._ar_diffusion_kv_state is None: + raise RuntimeError("LingBot step execution requires AR-Diffusion session binding.") + + def prepare_encode(self, state: StepRequestState, **kwargs: Any) -> StepRequestState: + del kwargs + self._require_bound_ar_state() + req = DiffusionRequestBatch( + requests=[ + OmniDiffusionRequest( + prompt=state.prompt, + sampling_params=state.sampling, + request_id=state.request_id, + ) + ] + ) + inputs = self._parse_request(req) + if ARDiffusionTickRequest.from_extra_args(state.sampling.extra_args) is not None: + raise ValueError("LingBot step execution does not accept AR-Diffusion ticks.") + if (inputs.height, inputs.width) != (self._ar_height, self._ar_width): + raise ValueError( + "LingBot AR-Diffusion request resolution must match the " + f"fixed cache geometry {self._ar_height}x{self._ar_width}." + ) + block_frames = int(self.transformer.config.num_frames_per_block) + total_chunks = inputs.num_latent_frames // block_frames + max_chunks = self._max_realtime_chunks() + if total_chunks > max_chunks: + raise ValueError( + "LingBot step execution currently supports at most " + f"{max_chunks} chunks because the image-condition horizon is " + f"{_MAX_RAW_FRAMES} pixel frames." + ) + if inputs.camera_action_script is not None and len(inputs.camera_action_script) != total_chunks: + raise ValueError( + "camera_action_script must contain one action list per generated chunk; " + f"got {len(inputs.camera_action_script)} chunks for {total_chunks} requested chunks." + ) + dtype = self.transformer.dtype + prompt_embeds = self.encode_prompt( + inputs.prompt, + max_sequence_length=inputs.max_sequence_length, + dtype=dtype, + ) + image_condition = self._prepare_condition( + replace( + inputs, + num_frames=_MAX_RAW_FRAMES, + num_latent_frames=self._horizon_latent_frames(), + ), + dtype=dtype, + ) + camera_trajectory_cache = None + camera_embedding_cache = None + if inputs.camera_trajectory is not None: + available_frames = int(inputs.camera_trajectory.poses.shape[0]) + if available_frames < inputs.num_frames: + raise ValueError( + "camera trajectory frames must be at least num_frames; " + f"got camera_frames={available_frames}, num_frames={inputs.num_frames}." + ) + truncated = CameraTrajectory( + poses=inputs.camera_trajectory.poses[: inputs.num_frames], + intrinsics=inputs.camera_trajectory.intrinsics[: inputs.num_frames], + ) + camera_trajectory_cache = interpolate_camera_trajectory(truncated, inputs.num_latent_frames) + # Embed the whole trajectory once, exactly as request mode does, and + # slice it per block in ``_prepare_next_chunk``. Embedding each chunk + # on its own would normalize that block's framewise translations by + # its own largest step, so a slow block would be conditioned as + # full-speed motion and drift from offline replay of the same path. + camera_embedding_cache, _ = self._prepare_camera(inputs, dtype=dtype) + state.prompt_embeds = prompt_embeds + state.chunk_index = 0 + state.step_index = 0 + state.step_in_chunk = 0 + state.total_chunks = total_chunks + state.chunk_num_steps = len(LINGBOT_DMD_TIMESTEPS) + state.extra = { + "inputs": inputs, + "image_condition": image_condition, + "schedule": _build_shifted_flow_schedule(flow_shift=inputs.flow_shift), + "dtype": dtype, + "block_frames": block_frames, + "camera_pitch": 0.0, + "camera_tail": None, + "camera_action_script": inputs.camera_action_script, + "camera_trajectory_cache": camera_trajectory_cache, + "camera_embedding_cache": camera_embedding_cache, + } + self._ar_text_caches(prompt_embeds, invalidate=False) + self._prepare_next_chunk(state) + return state + + def _prepare_next_chunk(self, state: StepRequestState) -> None: + extra = state.extra + inputs: _LingBotRequestInputs = extra["inputs"] + block_frames = int(extra["block_frames"]) + start_frame = state.chunk_index * block_frames + stop_frame = start_frame + block_frames + image_condition = extra["image_condition"] + if stop_frame > image_condition.shape[2]: + raise ValueError("LingBot chunk_index exceeds the configured causal image condition horizon.") + condition = image_condition[:, :, start_frame:stop_frame] + previous = extra.get("camera_tail") + if extra.get("camera_action_script") is not None: + chunk_actions = extra["camera_action_script"][state.chunk_index] + action_trajectory, camera_pitch = integrate_lingbot_camera_actions( + chunk_actions, + width=inputs.width, + height=inputs.height, + initial_pose=(previous.poses[-1] if previous is not None else None), + initial_pitch=float(extra.get("camera_pitch", 0.0)), + ) + extra["camera_pitch"] = camera_pitch + chunk_inputs = replace( + inputs, + camera_trajectory=action_trajectory, + camera_actions=chunk_actions, + num_frames=(block_frames - 1) * self.vae_scale_factor_temporal + 1, + num_latent_frames=block_frames, + ) + camera, camera_tail = self._prepare_camera( + chunk_inputs, + dtype=extra["dtype"], + previous=previous, + ) + elif extra.get("camera_embedding_cache") is not None: + # Same slice request mode takes from its one full-trajectory + # embedding, so both paths condition a block identically. + trajectory = extra["camera_trajectory_cache"] + camera = extra["camera_embedding_cache"][:, :, start_frame:stop_frame] + camera_tail = CameraTrajectory( + poses=trajectory.poses[stop_frame - 1 : stop_frame].clone(), + intrinsics=trajectory.intrinsics[stop_frame - 1 : stop_frame].clone(), + ) + else: + raise RuntimeError("LingBot step execution is missing a camera trajectory or action script.") + if camera.shape[2:] != condition.shape[2:]: + raise RuntimeError( + "folded camera and image condition must share latent frame/height/width geometry; " + f"got camera={tuple(camera.shape)}, condition={tuple(condition.shape)}." + ) + generator = getattr(state.sampling, "generator", None) + if not isinstance(generator, torch.Generator): + raise ValueError("LingBot World requires the runner-provided torch.Generator.") + extra["condition"] = condition + extra["camera"] = camera + extra["camera_tail"] = camera_tail + extra["start_frame"] = start_frame + extra["ar_cross_attention"] = self._ar_text_caches( + cast(torch.Tensor, state.prompt_embeds), + invalidate=False, + ) + state.latents = randn_tensor( + ( + 1, + self.transformer.config.out_channels, + condition.shape[2], + condition.shape[3], + condition.shape[4], + ), + generator=generator, + device=self.device, + dtype=torch.float32, + ) + state.timesteps = torch.tensor( + [timestep for timestep, _ in extra["schedule"]], + device=self.device, + dtype=torch.float32, + ) + state.step_in_chunk = 0 + state.step_index = 0 + + def denoise_step( + self, + input_batch: InputBatch, + *, + states: Sequence[StepRequestState] | None = None, + **kwargs: Any, + ) -> torch.Tensor | None: + del input_batch, kwargs + if states is None or len(states) != 1: + raise ValueError("LingBot step execution supports a single request, not a batched request.") + self._require_bound_ar_state() + state = states[0] + extra = state.extra + latents = state.latents + if latents is None: + raise RuntimeError("LingBot step execution requires latents before denoise_step.") + timestep = state.current_timestep + if timestep is None: + raise RuntimeError("LingBot step execution requires a current timestep before denoise_step.") + schedule = extra["schedule"] + step_in_chunk = state.step_in_chunk + return self._dmd_blocks.probe_step( + current_latents=latents, + condition=extra["condition"], + camera=extra["camera"], + prompt_embeds=cast(torch.Tensor, state.prompt_embeds), + cache=None, + ar=self._ar_block_context(extra["ar_cross_attention"]), + start_frame=int(extra["start_frame"]), + timestep_value=float(schedule[step_in_chunk][0]), + step_index=step_in_chunk, + ) + + def step_scheduler(self, state: StepRequestState, noise_pred: torch.Tensor, **kwargs: Any) -> None: + del kwargs + extra = state.extra + latents = state.latents + if latents is None: + raise RuntimeError("LingBot step execution requires latents before step_scheduler.") + generator = getattr(state.sampling, "generator", None) + if not isinstance(generator, torch.Generator): + raise ValueError("LingBot World requires the runner-provided torch.Generator.") + schedule = extra["schedule"] + step_in_chunk = state.step_in_chunk + next_sigma = schedule[step_in_chunk + 1][1] if step_in_chunk + 1 < len(schedule) else None + state.latents = self._dmd_blocks.apply_transition( + latents, + noise_pred, + schedule[step_in_chunk][1], + next_sigma=next_sigma, + generator=generator, + ) + state.step_in_chunk += 1 + state.step_index = state.step_in_chunk + + def post_decode(self, state: StepRequestState, **kwargs: Any) -> DiffusionOutput: + del kwargs + self._require_bound_ar_state() + extra = state.extra + latents = state.latents + if latents is None: + raise RuntimeError("LingBot step execution requires latents before post_decode.") + self._dmd_blocks.commit_block_kv( + latents=latents, + condition=extra["condition"], + camera=extra["camera"], + prompt_embeds=cast(torch.Tensor, state.prompt_embeds), + cache=None, + ar=self._ar_block_context(extra["ar_cross_attention"]), + start_frame=int(extra["start_frame"]), + ) + completed_chunk_index = state.chunk_index + inputs: _LingBotRequestInputs = extra["inputs"] + # "video" is the payload key the output formatter treats as primary, so + # a streaming client receives frames on ``images`` like any other video + # pipeline; latent mode keeps the tensor for offline consumers. + if inputs.output_type == "latent": + payload: dict[str, Any] = {"latents": latents} + else: + payload = {"video": self._decode_chunk_to_pixels(latents, output_type=inputs.output_type)} + output = { + "payload": payload, + "metadata": { + "ar_diffusion": ARDiffusionChunkMetadata( + session_id=state.request_id, + request_id=state.request_id, + chunk_index=completed_chunk_index, + applied_event_ids=(), + ).to_dict() + }, + } + state.chunk_index += 1 + finished = state.request_denoise_completed + if not finished: + self._prepare_next_chunk(state) + return DiffusionOutput( + output=output, + stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, + chunk_index=completed_chunk_index, + total_chunks=state.total_chunks, + finished=finished, + ) + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self) return cast(set[str], loader.load_weights(weights)) diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 608d605e9..4cec16382 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1268,6 +1268,17 @@ class AsyncOmniEngine: strategy_config_path = kwargs.pop("strategy_config", None) stage_overrides_json = kwargs.pop("stage_overrides", None) + # ``diffusion_streaming_output`` is the public AsyncOmni/serve kwarg; + # stage configs know the field as ``streaming_output``. The unregistered + # single-stage fallback translates it in + # ``_create_default_diffusion_stage_cfg``, but a registered pipeline + # resolves through StageConfigFactory, which only passes through keys the + # stage schema recognizes — so mirror the translation here. Only a truthy + # value is mirrored: the serve CLI always carries the flag's ``False`` + # default, which must not override a deploy YAML's ``streaming_output``. + if kwargs.get("diffusion_streaming_output") and kwargs.get("streaming_output") is None: + kwargs["streaming_output"] = True + # Parse --stage-overrides JSON string if provided stage_overrides = parse_stage_overrides(stage_overrides_json) diff --git a/vllm_omni/experimental/ar_diffusion/capability.py b/vllm_omni/experimental/ar_diffusion/capability.py index 64eb62cf0..145261513 100644 --- a/vllm_omni/experimental/ar_diffusion/capability.py +++ b/vllm_omni/experimental/ar_diffusion/capability.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Typed pipeline contract for the model-neutral AR-Diffusion runtime.""" from __future__ import annotations @@ -139,10 +140,13 @@ class SupportsARDiffusionPipeline(Protocol): """Required pipeline capability for :class:`ARDiffusionModelRunner`. A session begins when the runner first sees its ``session_id`` and persists - across requests until reset, explicit close, LRU eviction, or a failed - forward. ``bind_ar_diffusion_state`` exposes the runner-owned KV state only - for the duration of one request. The pipeline must not retain the state - after the context exits. + until reset, explicit close, LRU eviction, a failed forward, or request + completion on the stepwise path. ``bind_ar_diffusion_state`` exposes the + runner-owned KV state only for the duration of one runner invocation + (``execute_model`` or ``execute_stepwise``). The pipeline must not retain + the state after the context exits. Uncommitted scratch lives on the + runner-owned session object and survives across stepwise invocations of + the same request. """ def ar_diffusion_kv_cache_spec(self) -> ARDiffusionKVCacheSpec: @@ -154,7 +158,7 @@ class SupportsARDiffusionPipeline(Protocol): session_id: str, state: ARDiffusionKVState, ) -> AbstractContextManager[None]: - """Bind ``state`` to model execution for one request.""" + """Bind ``state`` to model execution for one runner invocation.""" ... def reset_ar_diffusion_session(self, session_id: str) -> None: diff --git a/vllm_omni/experimental/ar_diffusion/runner.py b/vllm_omni/experimental/ar_diffusion/runner.py index 410c45fb1..23da43054 100644 --- a/vllm_omni/experimental/ar_diffusion/runner.py +++ b/vllm_omni/experimental/ar_diffusion/runner.py @@ -1,14 +1,18 @@ # SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from __future__ import annotations import dataclasses import time from collections import OrderedDict +from collections.abc import Iterator +from contextlib import contextmanager import torch from vllm.logger import init_logger from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.models.interface import supports_step_execution from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput, KVPrefetchJob from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner @@ -22,6 +26,7 @@ from vllm_omni.experimental.ar_diffusion.kv_cache.config import ARDiffusionKVCon from vllm_omni.experimental.ar_diffusion.kv_cache.manager import ARDiffusionKVCache from vllm_omni.experimental.ar_diffusion.kv_cache.state import ARDiffusionKVState from vllm_omni.experimental.ar_diffusion.tick_protocol import ARDiffusionTickRequest +from vllm_omni.platforms import current_omni_platform logger = init_logger(__name__) @@ -62,6 +67,10 @@ class ARDiffusionModelRunner(DiffusionModelRunner): self._sessions: OrderedDict[str, ARDiffusionKVState] = OrderedDict() self._session_capacity = 0 self._perf_e2e_times: list[float] = [] + # Wall-clock start of each request's in-flight stepwise chunk, so + # ``_perf_e2e_times`` keeps one entry per AR block in both execution + # modes instead of one entry per denoise step. + self._stepwise_chunk_started: dict[str, float] = {} @staticmethod def _require_capability(pipeline: object) -> SupportsARDiffusionPipeline: @@ -90,7 +99,7 @@ class ARDiffusionModelRunner(DiffusionModelRunner): def _preallocate_kv_cache(self, *, available_bytes: int | None = None) -> None: """Build pools solely from the pipeline capability and runner config.""" - if bool(getattr(self.od_config, "step_execution", False)): + if bool(getattr(self.od_config, "step_execution", False)) and not supports_step_execution(self.pipeline): raise ValueError( "ARDiffusionModelRunner currently supports request-mode execution only; " "step_execution=True would bypass per-request AR session binding." @@ -240,29 +249,29 @@ class ARDiffusionModelRunner(DiffusionModelRunner): return tick.session_id, extra_args, tick return str(extra_args.get("session_id") or "default"), extra_args, None - def execute_model( + @contextmanager + def _bound_ar_session( self, - req: OmniDiffusionRequest, - kv_prefetch_job: KVPrefetchJob | None = None, - ) -> DiffusionOutput: - if self.kv_cache is None: - return super().execute_model(req, kv_prefetch_job=kv_prefetch_job) + session_id: str, + *, + description: str, + synchronize: bool = True, + ) -> Iterator[None]: + """Bind runner-owned KV for one runner invocation and fail closed. + + The pipeline sees the state only inside this context. On any failure the + session is released before the exception propagates, so a later request + for the same id cannot resume half-written pages. + """ capability = self._ar_diffusion_capability if capability is None: raise RuntimeError("AR-Diffusion capability missing after KV cache initialization") - - session_id, extra_args, tick = self._request_session(req) - reset = tick.reset if tick is not None else bool(extra_args.get("reset", False)) - close_session = tick.close_session if tick is not None else bool(extra_args.get("close_session", False)) - if reset: - self.reset_session(session_id) state = self._get_or_create_session(session_id) - started = time.perf_counter() try: with capability.bind_ar_diffusion_state(session_id, state): - output = super().execute_model(req, kv_prefetch_job=kv_prefetch_job) - if self.device is not None and torch.device(self.device).type == "cuda": - torch.accelerator.synchronize(self.device) + yield + if synchronize: + current_omni_platform.synchronize() except Exception: self._release_session( session_id, @@ -271,10 +280,49 @@ class ARDiffusionModelRunner(DiffusionModelRunner): suppress_errors=True, ) logger.warning( - "AR-Diffusion forward failed for session=%s; KV and model state were released", + "AR-Diffusion %s failed for session=%s; KV and model state were released", + description, session_id, ) raise + + def _release_scheduler_finished_sessions(self, scheduler_output: DiffusionSchedulerOutput) -> None: + """Close sessions the scheduler retired, including aborted requests. + + A stepwise request that runs to completion is closed as soon as its + final chunk is emitted; an aborted one never produces that output, so + its KV would otherwise linger until LRU eviction. Session ids equal + request ids on this path, and tick sessions are keyed by their own ids, + so the membership check leaves them untouched. + """ + for request_id in getattr(scheduler_output, "finished_req_ids", None) or (): + self._stepwise_chunk_started.pop(request_id, None) + if request_id in self._sessions: + self._release_session( + request_id, + reset_model=False, + reason="close", + suppress_errors=True, + ) + + def execute_model( + self, + req: OmniDiffusionRequest, + kv_prefetch_job: KVPrefetchJob | None = None, + ) -> DiffusionOutput: + if self.kv_cache is None: + return super().execute_model(req, kv_prefetch_job=kv_prefetch_job) + if self._ar_diffusion_capability is None: + raise RuntimeError("AR-Diffusion capability missing after KV cache initialization") + + session_id, extra_args, tick = self._request_session(req) + reset = tick.reset if tick is not None else bool(extra_args.get("reset", False)) + close_session = tick.close_session if tick is not None else bool(extra_args.get("close_session", False)) + if reset: + self.reset_session(session_id) + started = time.perf_counter() + with self._bound_ar_session(session_id, description="forward"): + output = super().execute_model(req, kv_prefetch_job=kv_prefetch_job) self._perf_e2e_times.append(time.perf_counter() - started) if close_session: self.close_session(session_id) @@ -291,10 +339,81 @@ class ARDiffusionModelRunner(DiffusionModelRunner): ) def execute_stepwise(self, scheduler_output: DiffusionSchedulerOutput) -> BatchRunnerOutput: - """Reject step execution until step-aware AR state binding exists.""" - raise RuntimeError( - "ARDiffusionModelRunner does not support step execution; use request mode with step_execution=False." - ) + """Bind runner-owned KV for one stepwise invocation, then inherit the step loop.""" + pipeline = getattr(self, "pipeline", None) + if pipeline is None or not supports_step_execution(pipeline): + raise RuntimeError( + "ARDiffusionModelRunner does not support step execution; use request mode with step_execution=False." + ) + if self.kv_cache is None: + return super().execute_stepwise(scheduler_output) + if self._ar_diffusion_capability is None: + raise RuntimeError("AR-Diffusion capability missing after KV cache initialization") + # Retire sessions the scheduler dropped before serving this cycle: a + # completed request is closed when its final chunk is emitted, an + # aborted one only here. + self._release_scheduler_finished_sessions(scheduler_output) + request_ids = list(getattr(scheduler_output, "scheduled_request_ids", ())) + if not request_ids: + return super().execute_stepwise(scheduler_output) + if len(request_ids) != 1: + raise RuntimeError( + "AR-Diffusion step execution supports one request at a time; " + f"got {len(request_ids)} scheduled requests." + ) + request_id = request_ids[0] + session_id = request_id + chunk_started = self._stepwise_chunk_started.setdefault(request_id, time.perf_counter()) + try: + # Denoise steps inside one chunk need no device barrier; the chunk + # boundary below is the meaningful one, and syncing per step would + # cost one synchronize per denoise step instead of per AR block. + with self._bound_ar_session(session_id, description="stepwise", synchronize=False): + output = super().execute_stepwise(scheduler_output) + runner_output = output.get_request_output(request_id) + finished = bool(runner_output is not None and runner_output.finished) + result = None if runner_output is None else runner_output.result + errored = bool(result is not None and getattr(result, "error", None)) + chunk_emitted = result is not None or finished or errored + if chunk_emitted: + # The barrier is where an asynchronous device error surfaces, so + # it must sit inside the same fail-closed scope as the step + # itself rather than leave cleanup to the scheduler's next cycle. + current_omni_platform.synchronize() + except Exception: + self._fail_closed_stepwise(request_id) + raise + # One entry per emitted AR block keeps ``_perf_e2e_times`` comparable + # with request mode, where one execute_model() is one block. + if chunk_emitted: + self._perf_e2e_times.append(time.perf_counter() - chunk_started) + self._stepwise_chunk_started.pop(request_id, None) + if finished or errored: + self._release_session( + session_id, + reset_model=False, + reason="forward_exception" if errored else "close", + suppress_errors=True, + ) + return output + + def _fail_closed_stepwise(self, request_id: str) -> None: + """Drop everything a failed stepwise invocation left behind. + + A failure inside the step is already released by ``_bound_ar_session``; + a failure at the chunk barrier happens after that context closed, so the + session is still present and is released here. + """ + self._stepwise_chunk_started.pop(request_id, None) + state_cache = getattr(self, "state_cache", None) + if isinstance(state_cache, dict): + state_cache.pop(request_id, None) + if request_id in self._sessions: + self._release_session(request_id, reset_model=False, reason="forward_exception", suppress_errors=True) + logger.warning( + "AR-Diffusion stepwise chunk barrier failed for session=%s; KV and model state were released", + request_id, + ) def _warmup_ar_rollout(self) -> None: """Run model-provided warmup requests, or safely skip when absent.""" diff --git a/vllm_omni/model_executor/models/lingbot_world/__init__.py b/vllm_omni/model_executor/models/lingbot_world/__init__.py new file mode 100644 index 000000000..4244ce4b4 --- /dev/null +++ b/vllm_omni/model_executor/models/lingbot_world/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project diff --git a/vllm_omni/model_executor/models/lingbot_world/pipeline.py b/vllm_omni/model_executor/models/lingbot_world/pipeline.py new file mode 100644 index 000000000..82d8258f9 --- /dev/null +++ b/vllm_omni/model_executor/models/lingbot_world/pipeline.py @@ -0,0 +1,37 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""LingBot-World 2.0 single-stage diffusion topology. + +Serving resolves a deploy config through this registration, so the stepwise +stage cannot be launched with ``--deploy-config`` until the pipeline name is +registered here. + +Deliberately no ``default_deploy_config_name``: model-type inference matches +this checkpoint by name, so a default would silently retarget every ``Omni`` +and ``AsyncOmni`` caller that passes no deploy config (the offline replay +example, the tick example) onto the stepwise serving topology. Stepwise +serving stays opt-in through +``--deploy-config vllm_omni/deploy/lingbot_world_v2_stepwise.yaml``. +""" + +from vllm_omni.config.stage_config import ( + PipelineConfig, + StageExecutionType, + StagePipelineConfig, +) + +LINGBOT_WORLD_PIPELINE = PipelineConfig( + model_type="lingbot_world", + model_arch="LingBotWorldCausalDMDPipeline", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="diffusion", + execution_type=StageExecutionType.DIFFUSION, + input_sources=(), + final_output=True, + final_output_type="video", + model_arch="LingBotWorldCausalDMDPipeline", + ), + ), +)