diff --git a/examples/offline_inference/dreamzero/client_schedule.py b/examples/offline_inference/dreamzero/client_schedule.py new file mode 100644 index 000000000..e4dbad55a --- /dev/null +++ b/examples/offline_inference/dreamzero/client_schedule.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""DreamZero AR client frame schedule aligned with upstream ``test_client_AR.py``.""" + +from __future__ import annotations + +import logging +from typing import Any + +import numpy as np + +# Frame schedule constants (matching upstream debug_inference.py / test_client_AR.py). +ACTION_HORIZON = 24 +DEFAULT_ACTION_DIM = 8 +RELATIVE_OFFSETS = [-23, -16, -8, 0] +DEFAULT_NUM_AR_CHUNKS = 15 + +CAMERA_FILES = { + "observation/exterior_image_0_left": "exterior_image_1_left.mp4", + "observation/exterior_image_1_left": "exterior_image_2_left.mp4", + "observation/wrist_image_left": "wrist_image_left.mp4", +} + + +def build_frame_schedule( + total_frames: int, + num_chunks: int, + *, + logger: logging.Logger | None = None, +) -> list[list[int]]: + """Build the frame index schedule for multi-frame chunks. + + Args: + total_frames: Number of frames available in each camera video. + num_chunks: Number of 4-frame chunks to schedule after the initial frame. + + Returns: + A list of frame-index lists. Each inner list has four indices. The + returned list may be shorter than ``num_chunks`` when the videos run + out of frames (upstream stops early instead of erroring). + """ + chunks: list[list[int]] = [] + current_frame = 23 + for _ in range(num_chunks): + indices = [max(current_frame + offset, 0) for offset in RELATIVE_OFFSETS] + if indices[-1] >= total_frames: + if logger is not None: + logger.info( + "Frame %s >= %s, stopping at %s chunks", + indices[-1], + total_frames, + len(chunks), + ) + break + chunks.append(indices) + current_frame += ACTION_HORIZON + return chunks + + +def make_obs_from_video( + camera_frames: dict[str, np.ndarray], + frame_indices: list[int], + *, + prompt: str, + session_id: str, +) -> dict[str, Any]: + """Build an observation dict from real video frames. + + For one frame each image key is ``(H, W, 3)``. For four frames each key is + ``(4, H, W, 3)``. + """ + obs: dict[str, Any] = {} + for camera_key, all_frames in camera_frames.items(): + selected = all_frames[frame_indices] + obs[camera_key] = selected[0] if len(frame_indices) == 1 else selected + + obs["observation/joint_position"] = np.zeros(7, dtype=np.float32) + obs["observation/cartesian_position"] = np.zeros(6, dtype=np.float32) + obs["observation/gripper_position"] = np.zeros(1, dtype=np.float32) + obs["prompt"] = prompt + obs["session_id"] = session_id + return obs + + +def build_ar_observations( + camera_frames: dict[str, np.ndarray], + *, + prompt: str, + session_id: str, + num_chunks: int = DEFAULT_NUM_AR_CHUNKS, + repeat_chunk_observations: bool = False, + logger: logging.Logger | None = None, +) -> list[dict[str, Any]]: + """Build the AR observation sequence used by upstream ``test_client_AR.py``. + + Step 0 sends frame ``[0]``. Each subsequent step sends one 4-frame chunk. + + ``num_chunks`` counts only the 4-frame chunks after the initial frame, not + the total number of inferences. + """ + if num_chunks < 0: + raise ValueError("num_chunks must be non-negative") + + total_frames = min(frames.shape[0] for frames in camera_frames.values()) + observations = [ + make_obs_from_video( + camera_frames, + [0], + prompt=prompt, + session_id=session_id, + ) + ] + + chunk_schedule = build_frame_schedule(total_frames, num_chunks, logger=logger) + if repeat_chunk_observations and chunk_schedule and len(chunk_schedule) < num_chunks: + while len(chunk_schedule) < num_chunks: + chunk_schedule.append(chunk_schedule[-1]) + + for frame_indices in chunk_schedule: + observations.append( + make_obs_from_video( + camera_frames, + frame_indices, + prompt=prompt, + session_id=session_id, + ) + ) + return observations diff --git a/examples/offline_inference/dreamzero/export_prediction_video.py b/examples/offline_inference/dreamzero/export_prediction_video.py index 82f56ff26..4c560418c 100644 --- a/examples/offline_inference/dreamzero/export_prediction_video.py +++ b/examples/offline_inference/dreamzero/export_prediction_video.py @@ -11,6 +11,11 @@ from pathlib import Path import cv2 import numpy as np import torch +from client_schedule import ( + CAMERA_FILES, + DEFAULT_NUM_AR_CHUNKS, + build_ar_observations, +) from PIL import Image from vllm_omni import Omni @@ -20,13 +25,6 @@ from vllm_omni.outputs import OmniRequestOutput WORKER_EXTENSION = "vllm_omni.diffusion.models.dreamzero.video_export_worker.DreamZeroVideoExportWorkerExtension" ASSET_REPO_ID = "YangshenDeng/vllm-omni-dreamzero-assets" DEFAULT_SESSION_PREFIX = "dreamzero-export" -RELATIVE_OFFSETS = [-23, -16, -8, 0] -ACTION_HORIZON = 24 -CAMERA_FILES = { - "observation/exterior_image_0_left": "exterior_image_1_left.mp4", - "observation/exterior_image_1_left": "exterior_image_2_left.mp4", - "observation/wrist_image_left": "wrist_image_left.mp4", -} def _parse_args() -> argparse.Namespace: @@ -49,6 +47,17 @@ def _parse_args() -> argparse.Namespace: default="Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan", ) parser.add_argument("--session-id", default=None) + parser.add_argument( + "--num-chunks", + type=int, + default=DEFAULT_NUM_AR_CHUNKS, + help="Number of 4-frame chunks to send after the initial frame (default: 15).", + ) + parser.add_argument( + "--repeat-chunk-observations", + action="store_true", + help="Repeat the last valid chunk observation when bundled assets run out of frames.", + ) parser.add_argument("--save-input-video", action="store_true") parser.add_argument("--save-gif", action="store_true") parser.add_argument("--save-actions", action="store_true") @@ -84,57 +93,25 @@ def _load_camera_frames(video_dir: Path) -> dict[str, np.ndarray]: return camera_frames -def _build_frame_schedule(total_frames: int, num_chunks: int) -> list[list[int]]: - chunks: list[list[int]] = [] - current_frame = 23 - for _ in range(num_chunks): - indices = [max(current_frame + offset, 0) for offset in RELATIVE_OFFSETS] - if indices[-1] >= total_frames: - break - chunks.append(indices) - current_frame += ACTION_HORIZON - return chunks - - -def _make_obs_from_video( - camera_frames: dict[str, np.ndarray], - frame_indices: list[int], - *, +def _build_observations( + video_dir: Path, prompt: str, session_id: str, -) -> dict: - obs: dict = {} - for camera_key, all_frames in camera_frames.items(): - selected = all_frames[frame_indices] - obs[camera_key] = selected[0] if len(frame_indices) == 1 else selected - - obs["observation/joint_position"] = np.zeros(7, dtype=np.float32) - obs["observation/cartesian_position"] = np.zeros(6, dtype=np.float32) - obs["observation/gripper_position"] = np.zeros(1, dtype=np.float32) - obs["prompt"] = prompt - obs["session_id"] = session_id - return obs - - -def _build_observations(video_dir: Path, prompt: str, session_id: str) -> tuple[dict[str, np.ndarray], list[dict]]: + *, + num_chunks: int, + repeat_chunk_observations: bool, +) -> tuple[dict[str, np.ndarray], list[dict]]: camera_frames = _load_camera_frames(video_dir) - total_frames = min(frames.shape[0] for frames in camera_frames.values()) - chunks = _build_frame_schedule(total_frames, 1) - observations = [ - _make_obs_from_video(camera_frames, [0], prompt=prompt, session_id=session_id), - ] - if chunks: - observations.append( - _make_obs_from_video( - camera_frames, - chunks[0], - prompt=prompt, - session_id=session_id, - ) - ) + observations = build_ar_observations( + camera_frames, + prompt=prompt, + session_id=session_id, + num_chunks=num_chunks, + repeat_chunk_observations=repeat_chunk_observations, + ) if len(observations) < 2: - raise RuntimeError("Need at least two DreamZero example observations to export a prediction video.") - return camera_frames, observations[:2] + raise RuntimeError("Need at least two DreamZero observations to export a prediction video.") + return camera_frames, observations def _extract_latents(output: OmniRequestOutput) -> torch.Tensor: @@ -229,6 +206,26 @@ def _run_generation( return omni, outputs +def _decode_accumulated_with_worker(omni: Omni, session_id: str) -> np.ndarray: + """Decode server-side accumulated latents when the pipeline supports it.""" + stage_client = omni.engine.stage_clients[0] + engine = getattr(stage_client, "_engine", None) + if engine is None: + raise RuntimeError("DreamZero export requires inline diffusion stage access.") + + decoded = engine.executor.collective_rpc( + "decode_accumulated_video_latents_to_uint8", + args=(session_id,), + unique_reply_rank=0, + exec_all_ranks=True, + ) + if isinstance(decoded, torch.Tensor): + decoded = decoded.numpy() + if not isinstance(decoded, np.ndarray): + raise TypeError(f"Unexpected decoded output type: {type(decoded)!r}") + return decoded + + def _decode_with_worker(omni: Omni, full_latents: torch.Tensor) -> np.ndarray: stage_client = omni.engine.stage_clients[0] engine = getattr(stage_client, "_engine", None) @@ -248,6 +245,15 @@ def _decode_with_worker(omni: Omni, full_latents: torch.Tensor) -> np.ndarray: return decoded +def _decode_video(omni: Omni, session_id: str, outputs: list[OmniRequestOutput]) -> np.ndarray: + """Decode full AR rollout by concatenating per-step latents from all outputs.""" + latent_steps = [_extract_latents(output) for output in outputs] + latent_t = [step.shape[2] for step in latent_steps] + print(f"Decoding {len(latent_steps)} latent steps, latent_t={latent_t}, total_t={sum(latent_t)}") + full_latents = torch.cat(latent_steps, dim=2) + return _decode_with_worker(omni, full_latents) + + def main() -> None: args = _parse_args() session_id = args.session_id or f"{DEFAULT_SESSION_PREFIX}-{uuid.uuid4()}" @@ -256,6 +262,12 @@ def main() -> None: video_dir=args.video_dir, prompt=args.prompt, session_id=session_id, + num_chunks=args.num_chunks, + repeat_chunk_observations=args.repeat_chunk_observations, + ) + print( + f"Running {len(observations)} inferences " + f"(1 initial + {len(observations) - 1} chunks; requested {args.num_chunks} post-initial chunks)" ) args.output_dir.mkdir(parents=True, exist_ok=True) @@ -273,17 +285,19 @@ def main() -> None: deploy_config_path=args.deploy_config, observations=observations, ) - latent_steps = [_extract_latents(output) for output in outputs] - full_latents = torch.cat(latent_steps, dim=2) - frames = _decode_with_worker(omni, full_latents) + frames = _decode_video(omni, session_id, outputs) finally: if omni is not None: omni.close() + expected_frames = 8 * len(observations) + 1 mp4_path = args.output_dir / f"{args.output_stem}.mp4" _write_mp4(mp4_path, frames, fps=args.fps) - print(f"SAVED_MP4={mp4_path}") + print( + f"SAVED_MP4={mp4_path} frames={len(frames)} " + f"expected≈{expected_frames} duration≈{len(frames) / max(args.fps, 1):.1f}s" + ) if args.save_gif: gif_path = args.output_dir / f"{args.output_stem}.gif" @@ -294,8 +308,10 @@ def main() -> None: npz_path = args.output_dir / f"{args.output_stem}_actions.npz" np.savez( npz_path, - step0=np.asarray(outputs[0].multimodal_output.get("actions")), - step1=np.asarray(outputs[1].multimodal_output.get("actions")), + **{ + f"step{index}": np.asarray(output.multimodal_output.get("actions")) + for index, output in enumerate(outputs) + }, ) print(f"SAVED_ACTIONS={npz_path}") diff --git a/tests/diffusion/cache/test_stepcache.py b/tests/diffusion/cache/test_stepcache.py new file mode 100644 index 000000000..1fc3bbb59 --- /dev/null +++ b/tests/diffusion/cache/test_stepcache.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Unit tests for StepCacheBackend and StepCacheState.""" + +from unittest.mock import Mock + +import pytest +import torch + +from vllm_omni.diffusion.cache.selector import get_cache_backend +from vllm_omni.diffusion.cache.stepcache import ( + STEP_CACHE_CONFIG_ATTR, + STEP_CACHE_STATE_ATTR, + StepCacheBackend, + StepCacheConfig, + StepCacheState, + is_stepcache_active, +) +from vllm_omni.diffusion.data import DiffusionCacheConfig + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _make_prev_predictions(sim: float, dim: int = 8) -> list[tuple[torch.Tensor]]: + v_prev = torch.randn(1, dim) + if sim >= 0.99: + v_last = v_prev.clone() + else: + v_last = torch.randn(1, dim) + return [(v_prev,), (v_last,)] + + +class TestStepCacheConfig: + def test_from_diffusion_cache_config_respects_step_cache_dit_enabled(self): + enabled = StepCacheConfig.from_diffusion_cache_config(DiffusionCacheConfig(step_cache_dit_enabled=True)) + disabled = StepCacheConfig.from_diffusion_cache_config(DiffusionCacheConfig(step_cache_dit_enabled=False)) + assert enabled.enabled is True + assert disabled.enabled is False + + +class TestStepCacheState: + def test_disabled_always_runs(self): + state = StepCacheState(StepCacheConfig(enabled=False)) + assert state.should_run_step(_make_prev_predictions(0.99)) is True + + def test_insufficient_history_always_runs(self): + state = StepCacheState(StepCacheConfig(min_history_steps=2)) + assert state.should_run_step([(torch.randn(1, 4),)]) is True + + def test_high_similarity_skips_with_countdown(self): + state = StepCacheState( + StepCacheConfig( + sim_thresholds=(0.95,), + skip_countdowns=(2,), + ) + ) + preds = _make_prev_predictions(0.99) + assert state.should_run_step(preds) is False + assert state.skip_countdown == 2 + assert state.should_run_step(preds) is False + assert state.should_run_step(preds) is True + + def test_low_similarity_runs(self): + state = StepCacheState(StepCacheConfig()) + preds = _make_prev_predictions(0.5) + assert state.should_run_step(preds) is True + + def test_trim_history(self): + state = StepCacheState(StepCacheConfig(max_history=2)) + history: list[tuple[torch.Tensor]] = [(torch.randn(1, 2),) for _ in range(4)] + state.trim_history(history) + assert len(history) == 2 + + +class TestStepCacheBackend: + def test_enable_dreamzero_pipeline(self): + mock_pipeline = Mock() + mock_pipeline.__class__.__name__ = "DreamZeroPipeline" + + backend = StepCacheBackend(DiffusionCacheConfig()) + backend.enable(mock_pipeline) + + assert backend.enabled is True + cfg = getattr(mock_pipeline, STEP_CACHE_CONFIG_ATTR) + assert isinstance(cfg, StepCacheConfig) + assert is_stepcache_active(mock_pipeline) + state = getattr(mock_pipeline, STEP_CACHE_STATE_ATTR) + assert isinstance(state, StepCacheState) + + def test_refresh_resets_state(self): + mock_pipeline = Mock() + mock_pipeline.__class__.__name__ = "DreamZeroPipeline" + backend = StepCacheBackend(DiffusionCacheConfig()) + backend.enable(mock_pipeline) + state = getattr(mock_pipeline, STEP_CACHE_STATE_ATTR) + state.skip_countdown = 3 + backend.refresh(mock_pipeline, num_inference_steps=16, verbose=False) + assert state.skip_countdown == 0 + + def test_enable_unsupported_pipeline_raises(self): + mock_pipeline = Mock() + mock_pipeline.__class__.__name__ = "QwenImagePipeline" + backend = StepCacheBackend(DiffusionCacheConfig()) + + with pytest.raises(ValueError, match="step_cache backend does not support"): + backend.enable(mock_pipeline) + + def test_selector_returns_backend(self): + backend = get_cache_backend("step_cache", {"velocity_sim_thresholds": [0.9]}) + assert isinstance(backend, StepCacheBackend) + assert backend.config.velocity_sim_thresholds == [0.9] + + def test_legacy_backend_aliases(self): + assert isinstance(get_cache_backend("stepcache", {}), StepCacheBackend) + assert isinstance(get_cache_backend("step_cache_dit", {}), StepCacheBackend) diff --git a/tests/e2e/online_serving/test_dreamzero_expansion.py b/tests/e2e/online_serving/test_dreamzero_expansion.py index bc3215891..79429c123 100644 --- a/tests/e2e/online_serving/test_dreamzero_expansion.py +++ b/tests/e2e/online_serving/test_dreamzero_expansion.py @@ -28,8 +28,6 @@ test_params = [ server_args=[ "--deploy-config", "vllm_omni/deploy/dreamzero_tp1_cfg2.yaml", - "--enforce-eager", - "--disable-log-stats", ], env_dict={ "ATTENTION_BACKEND": "torch", diff --git a/tests/entrypoints/openai_api/test_openpi_connection.py b/tests/entrypoints/openai_api/test_openpi_connection.py index 83d17a28c..398ae6494 100644 --- a/tests/entrypoints/openai_api/test_openpi_connection.py +++ b/tests/entrypoints/openai_api/test_openpi_connection.py @@ -88,6 +88,47 @@ def test_unpack_accepts_msgpack_numpy_marker_dicts(): np.testing.assert_allclose(decoded[b"actions"], np.asarray([[1.0, 0.0]], dtype=np.float32)) +def test_unpack_accepts_openpi_client_ndarray_markers(): + image = np.arange(6, dtype=np.uint8).reshape(2, 3) + payload = { + "observation/exterior_image_0_left": { + b"__ndarray__": True, + b"dtype": image.dtype.str, + b"shape": image.shape, + b"data": image.tobytes(), + } + } + + decoded = openpi_connection._unpack_numpy(payload) + + np.testing.assert_array_equal(decoded["observation/exterior_image_0_left"], image) + assert decoded["observation/exterior_image_0_left"].dtype == np.uint8 + + +def test_unpack_openpi_client_packed_observation(): + openpi_msgpack = pytest.importorskip("openpi_client.msgpack_numpy") + image = np.zeros((180, 320, 3), dtype=np.uint8) + obs = { + "observation/exterior_image_0_left": image, + "observation/joint_position": np.zeros(7, dtype=np.float32), + } + + decoded = openpi_connection._unpack(openpi_msgpack.packb(obs)) + + np.testing.assert_array_equal(decoded["observation/exterior_image_0_left"], image) + np.testing.assert_allclose(decoded["observation/joint_position"], obs["observation/joint_position"]) + + +def test_pack_uses_openpi_client_ndarray_markers(): + openpi_msgpack = pytest.importorskip("openpi_client.msgpack_numpy") + actions = np.asarray([[1.0, 2.0]], dtype=np.float32) + + packed = openpi_connection._pack(actions) + decoded = openpi_msgpack.unpackb(packed) + + np.testing.assert_allclose(decoded, actions) + + def test_unpack_leaves_user_dict_without_numpy_kind_marker_unchanged(): payload = { "metadata": { diff --git a/vllm_omni/deploy/dreamzero.yaml b/vllm_omni/deploy/dreamzero.yaml index e77afad84..079101ba0 100644 --- a/vllm_omni/deploy/dreamzero.yaml +++ b/vllm_omni/deploy/dreamzero.yaml @@ -12,7 +12,7 @@ stages: - stage_id: 0 devices: "0" max_num_seqs: 1 - enforce_eager: true + enforce_eager: false model_class_name: DreamZeroPipeline model_config: default_robot_embodiment: roboarena @@ -23,3 +23,8 @@ stages: needs_stereo_camera: false needs_session_id: true action_space: joint_position + cache_backend: step_cache + cache_config: + step_cache_dit_enabled: true + velocity_sim_thresholds: [0.95, 0.93] + velocity_skip_countdowns: [4, 2] diff --git a/vllm_omni/deploy/dreamzero_tp1_cfg2.yaml b/vllm_omni/deploy/dreamzero_tp1_cfg2.yaml index ed9988c92..a453d43f3 100644 --- a/vllm_omni/deploy/dreamzero_tp1_cfg2.yaml +++ b/vllm_omni/deploy/dreamzero_tp1_cfg2.yaml @@ -9,7 +9,7 @@ stages: - stage_id: 0 devices: "0,1" max_num_seqs: 1 - enforce_eager: true + enforce_eager: false model_class_name: DreamZeroPipeline parallel_config: tensor_parallel_size: 1 @@ -23,3 +23,8 @@ stages: needs_stereo_camera: false needs_session_id: true action_space: joint_position + cache_backend: step_cache + cache_config: + step_cache_dit_enabled: true + velocity_sim_thresholds: [0.95, 0.93] + velocity_skip_countdowns: [4, 2] diff --git a/vllm_omni/diffusion/cache/base.py b/vllm_omni/diffusion/cache/base.py index d78076224..aafbf0ba7 100644 --- a/vllm_omni/diffusion/cache/base.py +++ b/vllm_omni/diffusion/cache/base.py @@ -13,6 +13,8 @@ Main cache backend implementations: the cache-dit library. Inherits from CacheBackend. Used via cache_backend="cache_dit". 2. TeaCacheBackend: Hook-based backend for TeaCache acceleration. Inherits from CacheBackend. Used via cache_backend="tea_cache". +3. StepCacheBackend: Velocity cosine step-skipping for DreamZero. Inherits from + CacheBackend. Used via cache_backend="step_cache". All backends implement the same interface: - enable(pipeline): Enable cache on the pipeline diff --git a/vllm_omni/diffusion/cache/selector.py b/vllm_omni/diffusion/cache/selector.py index e6c2e70e3..3b002abe0 100644 --- a/vllm_omni/diffusion/cache/selector.py +++ b/vllm_omni/diffusion/cache/selector.py @@ -3,6 +3,7 @@ from typing import Any from vllm_omni.diffusion.cache.base import CacheBackend from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTBackend from vllm_omni.diffusion.cache.magcache.backend import MagCacheBackend +from vllm_omni.diffusion.cache.stepcache import StepCacheBackend from vllm_omni.diffusion.cache.teacache.backend import TeaCacheBackend from vllm_omni.diffusion.data import DiffusionCacheConfig @@ -14,14 +15,15 @@ def get_cache_backend(cache_backend: str | None, cache_config: Any) -> CacheBack - cache_dit: Uses CacheDiTBackend with enable()/refresh() interface - tea_cache: Uses TeaCacheBackend with enable()/refresh() interface - mag_cache: Uses MagCacheBackend with enable()/refresh() interface + - step_cache: Uses StepCacheBackend for DreamZero velocity step skip Args: - cache_backend: Cache backend name ("cache_dit", "tea_cache", "mag_cache", or None). + cache_backend: Cache backend name ("cache_dit", "tea_cache", + "mag_cache", "step_cache", or None). cache_config: Cache configuration (dict or DiffusionCacheConfig instance). Returns: - Cache backend instance (CacheDiTBackend, TeaCacheBackend, or MagCacheBackend) - if cache_backend is set, None otherwise. + Cache backend instance if cache_backend is set, None otherwise. Raises: ValueError: If cache_backend is unsupported. @@ -38,7 +40,10 @@ def get_cache_backend(cache_backend: str | None, cache_config: Any) -> CacheBack return TeaCacheBackend(cache_config) elif cache_backend == "mag_cache": return MagCacheBackend(cache_config) + elif cache_backend in ("step_cache", "stepcache", "step_cache_dit"): + return StepCacheBackend(cache_config) else: raise ValueError( - f"Unsupported cache backend: {cache_backend}. Supported: 'cache_dit', 'tea_cache', 'mag_cache'" + f"Unsupported cache backend: {cache_backend}. " + "Supported: 'cache_dit', 'tea_cache', 'mag_cache', 'step_cache'" ) diff --git a/vllm_omni/diffusion/cache/stepcache/__init__.py b/vllm_omni/diffusion/cache/stepcache/__init__.py new file mode 100644 index 000000000..a7e1852b0 --- /dev/null +++ b/vllm_omni/diffusion/cache/stepcache/__init__.py @@ -0,0 +1,66 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +StepCache: velocity-based step skipping for DreamZero-style DiT denoising. + +Unlike block-level Cache-DiT, this backend skips entire DiT forwards during the +pipeline denoise loop when successive velocity predictions are highly aligned. + +Usage: + from vllm_omni import Omni + + omni = Omni( + model="...", + cache_backend="step_cache", + ) + + # Alternative: environment variable + # export DIFFUSION_CACHE_BACKEND=step_cache +""" + +from vllm_omni.diffusion.cache.stepcache.backend import ( + CUSTOM_STEP_CACHE_DIT_ENABLERS, + CUSTOM_STEPCACHE_ENABLERS, + STEP_CACHE_CONFIG_ATTR, + STEP_CACHE_DIT_CONFIG_ATTR, + STEP_CACHE_DIT_STATE_ATTR, + STEP_CACHE_STATE_ATTR, + StepCacheBackend, + StepCacheDiTBackend, + StepCacheDiTConfig, + StepCacheDiTState, + enable_dreamzero_step_cache_dit, + enable_dreamzero_stepcache, + get_step_cache_dit_config, + get_step_cache_dit_state, + get_stepcache_config, + get_stepcache_state, + is_step_cache_dit_active, + is_stepcache_active, +) +from vllm_omni.diffusion.cache.stepcache.config import StepCacheConfig +from vllm_omni.diffusion.cache.stepcache.state import StepCacheState + +__all__ = [ + "CUSTOM_STEP_CACHE_DIT_ENABLERS", + "CUSTOM_STEPCACHE_ENABLERS", + "STEP_CACHE_CONFIG_ATTR", + "STEP_CACHE_DIT_CONFIG_ATTR", + "STEP_CACHE_DIT_STATE_ATTR", + "STEP_CACHE_STATE_ATTR", + "StepCacheBackend", + "StepCacheConfig", + "StepCacheDiTBackend", + "StepCacheDiTConfig", + "StepCacheDiTState", + "StepCacheState", + "enable_dreamzero_step_cache_dit", + "enable_dreamzero_stepcache", + "get_step_cache_dit_config", + "get_step_cache_dit_state", + "get_stepcache_config", + "get_stepcache_state", + "is_step_cache_dit_active", + "is_stepcache_active", +] diff --git a/vllm_omni/diffusion/cache/stepcache/backend.py b/vllm_omni/diffusion/cache/stepcache/backend.py new file mode 100644 index 000000000..517d2369a --- /dev/null +++ b/vllm_omni/diffusion/cache/stepcache/backend.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +StepCache backend implementation. + +This module provides the stepcache backend that implements the CacheBackend +interface for DreamZero-style velocity-based step skipping in the denoise loop. +""" + +from __future__ import annotations + +from typing import Any + +from vllm.logger import init_logger + +from vllm_omni.diffusion.cache.base import CacheBackend +from vllm_omni.diffusion.cache.stepcache.config import ( + STEP_CACHE_CONFIG_ATTR, + STEP_CACHE_STATE_ATTR, + StepCacheConfig, +) +from vllm_omni.diffusion.cache.stepcache.state import StepCacheState +from vllm_omni.diffusion.data import DiffusionCacheConfig + +logger = init_logger(__name__) + +# Backward-compatible aliases for older imports and pipeline attribute names. +STEP_CACHE_DIT_CONFIG_ATTR = STEP_CACHE_CONFIG_ATTR +STEP_CACHE_DIT_STATE_ATTR = STEP_CACHE_STATE_ATTR +StepCacheDiTConfig = StepCacheConfig +StepCacheDiTState = StepCacheState + + +def is_stepcache_active(pipeline: Any) -> bool: + """Return True when stepcache is enabled on *pipeline*.""" + cfg = getattr(pipeline, STEP_CACHE_CONFIG_ATTR, None) + return isinstance(cfg, StepCacheConfig) and cfg.enabled + + +def get_stepcache_config(pipeline: Any) -> StepCacheConfig | None: + cfg = getattr(pipeline, STEP_CACHE_CONFIG_ATTR, None) + return cfg if isinstance(cfg, StepCacheConfig) else None + + +def get_stepcache_state(pipeline: Any) -> StepCacheState | None: + state = getattr(pipeline, STEP_CACHE_STATE_ATTR, None) + return state if isinstance(state, StepCacheState) else None + + +def is_step_cache_dit_active(pipeline: Any) -> bool: + return is_stepcache_active(pipeline) + + +def get_step_cache_dit_config(pipeline: Any) -> StepCacheConfig | None: + return get_stepcache_config(pipeline) + + +def get_step_cache_dit_state(pipeline: Any) -> StepCacheState | None: + return get_stepcache_state(pipeline) + + +def enable_dreamzero_stepcache(pipeline: Any, config: DiffusionCacheConfig) -> None: + """Enable stepcache for DreamZeroPipeline.""" + cache_config = StepCacheConfig.from_diffusion_cache_config(config) + setattr(pipeline, STEP_CACHE_CONFIG_ATTR, cache_config) + setattr(pipeline, STEP_CACHE_STATE_ATTR, StepCacheState(cache_config)) + + logger.info( + "stepcache enabled on DreamZeroPipeline (enabled=%s, thresholds=%s, countdowns=%s)", + cache_config.enabled, + cache_config.sim_thresholds, + cache_config.skip_countdowns, + ) + + +enable_dreamzero_step_cache_dit = enable_dreamzero_stepcache + +CUSTOM_STEPCACHE_ENABLERS = { + "DreamZeroPipeline": enable_dreamzero_stepcache, +} + +CUSTOM_STEP_CACHE_DIT_ENABLERS = CUSTOM_STEPCACHE_ENABLERS + + +class StepCacheBackend(CacheBackend): + """ + Velocity cosine step-skipping cache backend for DreamZero. + + Attaches :class:`StepCacheConfig` and :class:`StepCacheState` to supported + pipelines. The denoise loop calls :meth:`StepCacheState.should_run_step` to + decide whether to run ``predict_noise``. + + Example: + >>> from vllm_omni.diffusion.data import DiffusionCacheConfig + >>> backend = StepCacheBackend(DiffusionCacheConfig()) + >>> backend.enable(pipeline) + >>> backend.refresh(pipeline, num_inference_steps=16) + """ + + def enable(self, pipeline: Any) -> None: + """Enable stepcache on a supported pipeline.""" + pipeline_type = pipeline.__class__.__name__ + + if pipeline_type in CUSTOM_STEPCACHE_ENABLERS: + logger.info("Using custom stepcache enabler for model: %s", pipeline_type) + CUSTOM_STEPCACHE_ENABLERS[pipeline_type](pipeline, self.config) + else: + raise ValueError( + f"step_cache backend does not support {pipeline_type}. Supported: {sorted(CUSTOM_STEPCACHE_ENABLERS)}" + ) + + self.enabled = True + + def refresh(self, pipeline: Any, num_inference_steps: int, verbose: bool = True) -> None: + """Refresh stepcache state for a new generation.""" + state = get_stepcache_state(pipeline) + if state is not None: + state.reset() + if verbose: + logger.debug( + "stepcache state refreshed (num_inference_steps=%d)", + num_inference_steps, + ) + elif verbose and is_stepcache_active(pipeline): + logger.warning("stepcache config active but state not found on pipeline") + + +StepCacheDiTBackend = StepCacheBackend diff --git a/vllm_omni/diffusion/cache/stepcache/config.py b/vllm_omni/diffusion/cache/stepcache/config.py new file mode 100644 index 000000000..00a7984ae --- /dev/null +++ b/vllm_omni/diffusion/cache/stepcache/config.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Configuration for step-level DiT velocity caching.""" + +from __future__ import annotations + +from dataclasses import dataclass + +from vllm_omni.diffusion.data import DiffusionCacheConfig + +# Pipeline attributes set by StepCacheBackend.enable(). +STEP_CACHE_CONFIG_ATTR = "_stepcache_config" +STEP_CACHE_STATE_ATTR = "_stepcache_state" + + +@dataclass(frozen=True) +class StepCacheConfig: + """Runtime config for velocity-based step skipping. + + Skips entire DiT forwards when successive video velocity predictions are + highly aligned (cosine similarity above configured thresholds). + + Reference: DreamZero paper DiT Caching / dreamzero.git ``should_run_model``. + """ + + enabled: bool = True + min_history_steps: int = 2 + max_history: int = 2 + sim_thresholds: tuple[float, ...] = (0.95, 0.93) + skip_countdowns: tuple[int, ...] = (4, 2) + + @classmethod + def from_diffusion_cache_config(cls, config: DiffusionCacheConfig) -> StepCacheConfig: + enabled = config.step_cache_dit_enabled + + thresholds = tuple(config.velocity_sim_thresholds) + countdowns = tuple(config.velocity_skip_countdowns) + if len(thresholds) != len(countdowns): + raise ValueError( + "velocity_sim_thresholds and velocity_skip_countdowns must have the same length; " + f"got {len(thresholds)} and {len(countdowns)}." + ) + + return cls( + enabled=enabled, + min_history_steps=config.step_cache_dit_min_history, + max_history=config.step_cache_dit_max_history, + sim_thresholds=thresholds, + skip_countdowns=countdowns, + ) diff --git a/vllm_omni/diffusion/cache/stepcache/state.py b/vllm_omni/diffusion/cache/stepcache/state.py new file mode 100644 index 000000000..db786d4f2 --- /dev/null +++ b/vllm_omni/diffusion/cache/stepcache/state.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Mutable state for step-level DiT velocity caching.""" + +from __future__ import annotations + +import torch + +from vllm_omni.diffusion.cache.stepcache.config import StepCacheConfig + + +class StepCacheState: + """Per-generation mutable state (skip countdown).""" + + def __init__(self, config: StepCacheConfig) -> None: + self.config = config + self.skip_countdown = 0 + + def reset(self) -> None: + self.skip_countdown = 0 + + def should_run_step(self, prev_predictions: list[tuple[torch.Tensor, ...]]) -> bool: + """Return True when the DiT forward should execute at this scheduler step.""" + if not self.config.enabled: + return True + + if len(prev_predictions) < self.config.min_history_steps: + return True + + if self.skip_countdown > 1: + self.skip_countdown -= 1 + return False + if self.skip_countdown == 1: + self.skip_countdown = 0 + return True + + v_last = prev_predictions[-1][0].flatten(1).float() + v_prev = prev_predictions[-2][0].flatten(1).float() + sim = torch.nn.functional.cosine_similarity(v_last, v_prev, dim=1).mean() + + for threshold, countdown in zip(self.config.sim_thresholds, self.config.skip_countdowns): + if sim > threshold: + self.skip_countdown = countdown + return False + + return True + + def trim_history(self, prev_predictions: list[tuple[torch.Tensor, ...]]) -> None: + while len(prev_predictions) > self.config.max_history: + prev_predictions.pop(0) diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index c4cca97bb..bf294fdfe 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -302,6 +302,8 @@ class DiffusionCacheConfig: scm_steps_mask_policy, scm_steps_policy - MagCache: mag_threshold, mag_max_skip_steps, mag_retention_ratio, mag_ratios, mag_calibrate + - step_cache: step_cache_dit_enabled, velocity_sim_thresholds, + velocity_skip_countdowns, step_cache_dit_min_history Example: >>> # From dict (user-facing API) - partial config uses defaults for missing keys @@ -375,6 +377,13 @@ class DiffusionCacheConfig: # "repeat" refreshes every force_refresh_step_hint steps. force_refresh_step_policy: str = "once" + # step_cache parameters [step_cache only] — DreamZero velocity schedule + step_cache_dit_enabled: bool = True + velocity_sim_thresholds: list[float] = field(default_factory=lambda: [0.95, 0.93]) + velocity_skip_countdowns: list[int] = field(default_factory=lambda: [4, 2]) + step_cache_dit_min_history: int = 2 + step_cache_dit_max_history: int = 2 + # Additional parameters that may be passed but not explicitly defined _extra_params: dict[str, Any] = field(default_factory=dict, repr=False) diff --git a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py index 1d1429534..bbff12373 100644 --- a/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py +++ b/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py @@ -23,6 +23,10 @@ from huggingface_hub import hf_hub_download from transformers import AutoTokenizer, UMT5Config, UMT5EncoderModel from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm_omni.diffusion.cache.stepcache import ( + get_stepcache_state, + is_stepcache_active, +) from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import ( DistributedAutoencoderKLWan, @@ -55,10 +59,6 @@ logger = logging.getLogger(__name__) MAX_DREAMZERO_SESSIONS = 64 -# --------------------------------------------------------------------------- -# --------------------------------------------------------------------------- - - class VideoActionScheduler: """Wraps video + action schedulers into single .step() interface.""" @@ -205,6 +205,9 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin): self._max_session_states = MAX_DREAMZERO_SESSIONS self.state = self._get_or_create_state("default") + # DiT step cache is configured by StepCacheBackend + # (cache_backend="step_cache") via pipeline._stepcache_config. + # Keep runtime inference settings separate from the training-time config. self.num_inference_steps: int = model_config.get( "num_inference_steps", @@ -306,6 +309,27 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin): def predict_noise(self, **kwargs) -> tuple[torch.Tensor, torch.Tensor]: """Call CausalWanModel, return (video_pred, action_pred).""" + video_pred, action_pred = self._predict_noise_eager(kwargs) + + if is_stepcache_active(self): + video_pred = video_pred.clone() + if action_pred is not None: + action_pred = action_pred.clone() + + if action_pred is None: + batch_size = kwargs["hidden_states"].shape[0] + action_pred = torch.empty( + batch_size, + 0, + self.transformer.action_dim, + device=video_pred.device, + dtype=video_pred.dtype, + ) + return (video_pred, action_pred) + + def _predict_noise_eager(self, kwargs: dict) -> tuple[torch.Tensor, torch.Tensor | None]: + """Eager DiT forward; also handles KV-cache write-back on prefill.""" + self._cudagraph_mark_step_begin() video_pred, action_pred, updated_kv_caches = self.transformer( x=kwargs["hidden_states"], timestep=kwargs["timestep_video"], @@ -327,19 +351,118 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin): for i, kv in enumerate(updated_kv_caches): state.update_kv_cache(i, kv, is_negative=is_neg) - video_pred = video_pred.clone() - if action_pred is not None: - action_pred = action_pred.clone() - else: - batch_size = kwargs["hidden_states"].shape[0] - action_pred = torch.empty( - batch_size, - 0, - self.transformer.action_dim, - device=video_pred.device, - dtype=video_pred.dtype, - ) # CFG-parallel-safe dummy action pred - return (video_pred, action_pred) + return video_pred, action_pred + + # ----------------------------------------------------------------------- + # torch.compile setup (paper D.2 extended: encoders + VAE + DiT blocks) + # ----------------------------------------------------------------------- + + def setup_compile(self) -> None: + """Compile DreamZero encoders, VAE, and per-block DiT for inference. + + Paper D.2 uses ``mode=reduce-overhead``, ``fullgraph=True``, ``dynamic=False`` + on text/image/VAE and DiT. VAE decode uses a tensor feat_cache patch and + compiles ``decoder.forward`` (not ``_decode``, which has a Python frame loop). + DiT blocks use per-block ``fullgraph=True``. + """ + if not torch.cuda.is_available(): + logger.info("DreamZero setup_compile skipped: CUDA not available.") + return + + from vllm_omni.diffusion.models.dreamzero.wan_vae_feat_cache_patch import ( + apply_wan_vae_feat_cache_tensor_patch, + ) + + apply_wan_vae_feat_cache_tensor_patch() + + compile_ro = {"mode": "reduce-overhead", "fullgraph": True, "dynamic": False} + # DiT blocks: default avoids CUDAGraph overwrite on modulation tensors; encoders use reduce-overhead. + dit_compile = {"mode": "default", "fullgraph": True, "dynamic": False} + + logger.info( + "DreamZero: torch.compile text/image/VAE encode + per-block DiT (encoders reduce-overhead, DiT default)." + ) + + try: + self.text_encoder.forward = torch.compile(self.text_encoder.forward, **compile_ro) + except Exception as exc: + logger.warning("DreamZero: text_encoder compile failed (%s); skipping.", exc) + + try: + self.image_encoder.model.visual.forward = torch.compile( + self.image_encoder.model.visual.forward, + **compile_ro, + ) + except Exception as exc: + logger.warning("DreamZero: image_encoder compile failed (%s); skipping.", exc) + + try: + self.vae.encode = torch.compile(self.vae.encode, **compile_ro) + except Exception as exc: + logger.warning("DreamZero: vae.encode compile failed (%s); skipping.", exc) + + compiled_blocks = 0 + for block in self.transformer.blocks: + try: + block.forward = torch.compile(block.forward, **dit_compile) + compiled_blocks += 1 + except Exception as exc: + logger.warning( + "DreamZero: transformer block %d compile failed (%s); leaving remaining eager.", + compiled_blocks, + exc, + ) + break + if compiled_blocks: + logger.info( + "DreamZero: compiled %d/%d transformer blocks.", + compiled_blocks, + len(self.transformer.blocks), + ) + + self.warmup_compile() + + def warmup_compile(self) -> None: + """Warm up compiled text/image/VAE paths before timed inference.""" + if not torch.cuda.is_available(): + return + + device = next(self.text_encoder.parameters()).device + with torch.inference_mode(): + try: + text_tokens = torch.zeros(1, 16, dtype=torch.long, device=device) + attention_mask = torch.ones_like(text_tokens) + self._encode_text(text_tokens, attention_mask) + except Exception as exc: + logger.warning("DreamZero compile warmup (text_encoder) skipped: %s", exc) + + try: + image = torch.zeros(1, 1, 3, 180, 320, dtype=torch.bfloat16, device=device) + self._encode_image(image, self.num_frames, 180, 320) + except Exception as exc: + logger.warning("DreamZero compile warmup (image_encoder) skipped: %s", exc) + + try: + latent_h, latent_w = 180 // 8, 320 // 8 + dummy_latent = torch.zeros( + 1, + 16, + self.num_frame_per_block, + latent_h * 2, + latent_w * 2, + dtype=self.vae.dtype, + device=device, + ) + mean = self.vae_latents_mean.to(device=device, dtype=self.vae.dtype) + inv_std = self.vae_latents_inv_std.to(device=device, dtype=self.vae.dtype) + dummy_latent_denorm = dummy_latent / inv_std + mean + self._cudagraph_mark_step_begin() + self.vae.decode(dummy_latent_denorm, return_dict=False) + except Exception as exc: + logger.warning("DreamZero compile warmup (vae decode) skipped: %s", exc) + + torch.accelerator.synchronize(device) + logger.info("DreamZero compile warmup finished (text / image / vae decode).") def combine_cfg_noise( self, @@ -388,12 +511,20 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin): videos = videos.reshape(b, t, c, h, w).permute(0, 2, 1, 3, 4) return videos.to(dtype=torch.bfloat16) + @staticmethod + def _cudagraph_mark_step_begin() -> None: + try: + torch.compiler.cudagraph_mark_step_begin() + except Exception: + pass + # ----------------------------------------------------------------------- # Text encoding # ----------------------------------------------------------------------- def _encode_text(self, text_tokens: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: """Encode text prompt via UMT5.""" + self._cudagraph_mark_step_begin() seq_lens = attention_mask.gt(0).sum(dim=1).long() prompt_emb = self.text_encoder( text_tokens, @@ -422,6 +553,7 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin): batch_size = image.shape[0] with torch.amp.autocast(dtype=torch.bfloat16, device_type=device.type): + self._cudagraph_mark_step_begin() clip_context = self.image_encoder.encode_image(image) msk = torch.ones(batch_size, num_frames, height // 8, width // 8, device=device) @@ -469,8 +601,23 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin): inv_std = self.vae_latents_inv_std.to(device=vae_device, dtype=vae_dtype) latents = latents / inv_std + mean with torch.no_grad(): + self._cudagraph_mark_step_begin() return self.vae.decode(latents, return_dict=False)[0] + def decode_accumulated_video_latents(self, session_id: str | None = None) -> torch.Tensor: + """Decode all AR-chunk latents accumulated for ``session_id``.""" + state = self._get_or_create_state(session_id) + latents = state.get_concatenated_video_latents() + if latents is None: + session_key = str(session_id or "default") + raise RuntimeError(f"No accumulated video latents for session {session_key!r}.") + return self.decode_video_latents(latents) + + def clear_accumulated_video_latents(self, session_id: str | None = None) -> None: + """Clear accumulated video latents for ``session_id`` without resetting KV state.""" + state = self._get_or_create_state(session_id) + state.clear_video_latents() + # ----------------------------------------------------------------------- # KV cache prefill # ----------------------------------------------------------------------- @@ -635,6 +782,12 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin): noisy_input = video_latents noisy_input_action = action_latents + # ---- step_cache (StepCacheBackend / dreamzero.git) ---- + _cached_flow_pred: torch.Tensor | None = None + _cached_flow_pred_action: torch.Tensor | None = None + _prev_predictions: list[tuple[torch.Tensor]] = [] + _step_cache = get_stepcache_state(self) if is_stepcache_active(self) else None + for index in range(len(timesteps_video)): video_timestep = timesteps_video[index] action_timestep = timesteps_action[index] @@ -663,46 +816,58 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin): else: y = state.ys[:, :, -self.num_frame_per_block :] - positive_kwargs = dict( - hidden_states=noisy_input.transpose(1, 2), - timestep_video=timestep, - encoder_hidden_states=prompt_embeds, - kv_cache=state.get_kv_caches(False), - crossattn_cache=state.get_crossattn_caches(False), - y=y, - clip_feature=state.clip_feas, - action=noisy_input_action, - timestep_action=timestep_action, - is_negative=False, - **common_kwargs, - ) - - # Negative (uncond) kwargs - if do_true_cfg and negative_prompt_embeds is not None: - negative_kwargs = dict( + run_dit = _step_cache is None or _step_cache.should_run_step(_prev_predictions) + if run_dit: + positive_kwargs = dict( hidden_states=noisy_input.transpose(1, 2), timestep_video=timestep, - encoder_hidden_states=negative_prompt_embeds, - kv_cache=state.get_kv_caches(True), - crossattn_cache=state.get_crossattn_caches(True), + encoder_hidden_states=prompt_embeds, + kv_cache=state.get_kv_caches(False), + crossattn_cache=state.get_crossattn_caches(False), y=y, clip_feature=state.clip_feas, action=noisy_input_action, timestep_action=timestep_action, - is_negative=True, + is_negative=False, **common_kwargs, ) - else: - negative_kwargs = None - noise_pred = self.predict_noise_maybe_with_cfg( - positive_kwargs=positive_kwargs, - negative_kwargs=negative_kwargs, - do_true_cfg=do_true_cfg, - true_cfg_scale=self.cfg_scale, - cfg_normalize=False, - ) - flow_pred, flow_pred_action = noise_pred + if do_true_cfg and negative_prompt_embeds is not None: + negative_kwargs = dict( + hidden_states=noisy_input.transpose(1, 2), + timestep_video=timestep, + encoder_hidden_states=negative_prompt_embeds, + kv_cache=state.get_kv_caches(True), + crossattn_cache=state.get_crossattn_caches(True), + y=y, + clip_feature=state.clip_feas, + action=noisy_input_action, + timestep_action=timestep_action, + is_negative=True, + **common_kwargs, + ) + else: + negative_kwargs = None + + noise_pred = self.predict_noise_maybe_with_cfg( + positive_kwargs=positive_kwargs, + negative_kwargs=negative_kwargs, + do_true_cfg=do_true_cfg, + true_cfg_scale=self.cfg_scale, + cfg_normalize=False, + ) + flow_pred, flow_pred_action = noise_pred + _cached_flow_pred = flow_pred + _cached_flow_pred_action = flow_pred_action + + _prev_predictions.append((flow_pred,)) + if _step_cache is not None: + _step_cache.trim_history(_prev_predictions) + else: + # Reuse previous prediction (DiT step-skipping cache). + assert _cached_flow_pred is not None, "DiT cache: no cached prediction available for step skip." + flow_pred = _cached_flow_pred + flow_pred_action = _cached_flow_pred_action latents = (noisy_input, noisy_input_action) t = (video_timestep, action_timestep) @@ -811,9 +976,12 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin): # on the next inference request after websocket reset/session switch. if extra_args.get("reset", False): state.reset() - # Auto-reset based on model state (before accumulation) - if state.should_reset(text_tokens, 0, self.transformer.local_attn_size): - state.reset() + else: + reset_reason = state.reset_reason(text_tokens, 0, self.transformer.local_attn_size) + if reset_reason == "session": + state.reset() + elif reset_reason == "inference": + state.reset_inference_state() state.language = text_tokens # Frame accumulation: stitched single frame → multi-frame video @@ -959,6 +1127,8 @@ class DreamZeroPipeline(nn.Module, CFGParallelMixin): video_out = torch.cat([image, video_out], dim=1) state.current_start_frame += self.num_frame_per_block + state.append_video_latents(video_out) + # q99 denorm: [-1,1] → real values action_out = self._denormalize_action(action_out.float(), embodiment_name) diff --git a/vllm_omni/diffusion/models/dreamzero/state_dreamzero.py b/vllm_omni/diffusion/models/dreamzero/state_dreamzero.py index 9c1e7873d..5815d93d4 100644 --- a/vllm_omni/diffusion/models/dreamzero/state_dreamzero.py +++ b/vllm_omni/diffusion/models/dreamzero/state_dreamzero.py @@ -23,7 +23,8 @@ class DreamZeroState: Lifecycle: - Created once in DreamZeroPipeline.__init__() - Mutated every forward() call (frame append, KV cache grow) - - reset() on new session / language change / local_attn_size exceeded + - reset() on new session / language change + - reset_inference_state() on local_attn_size exceeded (KV only) """ def __init__(self) -> None: @@ -65,50 +66,106 @@ class DreamZeroState: self.call_count += 1 return np.stack(frames, axis=0) # (T, H, W, C) + def append_video_latents(self, video_out: torch.Tensor) -> None: + """Append one AR chunk of normalized video latents for later decode. + + Args: + video_out: ``(B, T, C, H, W)`` tensor returned by the denoise loop + (before the ``transpose(1, 2)`` stored in ``DiffusionOutput``). + """ + if video_out.dim() != 5: + raise ValueError(f"Expected 5D video_out, got shape {tuple(video_out.shape)}") + # Upstream ``torch.cat(self.video_across_time, dim=2)`` uses (B, C, T, H, W). + chunk = video_out.transpose(1, 2).detach().cpu() + self.video_latents_across_time.append(chunk) + logger.info( + "append_video_latents: chunk_shape=%s total_chunks=%d total_latent_t=%d", + tuple(chunk.shape), + len(self.video_latents_across_time), + int(sum(c.shape[2] for c in self.video_latents_across_time)), + ) + + def get_concatenated_video_latents(self) -> torch.Tensor | None: + """Return all accumulated chunks concatenated along the time dimension.""" + if not self.video_latents_across_time: + return None + if len(self.video_latents_across_time) == 1: + return self.video_latents_across_time[0] + return torch.cat(self.video_latents_across_time, dim=2) + + def clear_video_latents(self) -> None: + """Drop accumulated video latents without resetting KV/frame state.""" + self.video_latents_across_time = [] + # ------------------------------------------------------------------ # Reset / should_reset # ------------------------------------------------------------------ - def reset(self) -> None: - """Clear all state.""" - self.stitched_buffer: deque[np.ndarray] = deque(maxlen=FRAMES_PER_CHUNK) - self.call_count: int = 0 + def reset(self, *, clear_video_latents: bool = True) -> None: + """Clear session state. + + Args: + clear_video_latents: When ``False``, keep ``video_latents_across_time`` + so offline/online export can decode the full AR rollout. + """ + saved_video_latents = [] if clear_video_latents else list(self.video_latents_across_time) + self.stitched_buffer = deque(maxlen=FRAMES_PER_CHUNK) + self.call_count = 0 + # Normalized VAE latents per AR chunk, each (B, C, T, H, W) on CPU. + # Concatenated along T before decode (matches upstream video_across_time). + self.video_latents_across_time = saved_video_latents # KV cache once robot-policy diffusion supports that integration. - self.kv_cache: list[torch.Tensor] | None = None - self.kv_cache_neg: list[torch.Tensor] | None = None - self.crossattn_cache: list[dict[str, bool | torch.Tensor | None]] | None = None - self.crossattn_cache_neg: list[dict[str, bool | torch.Tensor | None]] | None = None - self.current_start_frame: int = 0 + self.kv_cache = None + self.kv_cache_neg = None + self.crossattn_cache = None + self.crossattn_cache_neg = None + self.current_start_frame = 0 - self.clip_feas: torch.Tensor | None = None - self.ys: torch.Tensor | None = None - self.language: torch.Tensor | None = None + self.clip_feas = None + self.ys = None + if clear_video_latents: + self.language = None - def should_reset(self, text_tokens: torch.Tensor | None, num_video_frames: int, local_attn_size: int) -> bool: - """Determine if state should be reset before this forward().""" + def reset_inference_state(self) -> None: + """Reset KV/frame state after local attention rolls without dropping video latents.""" + self.reset(clear_video_latents=False) + + def reset_reason( + self, + text_tokens: torch.Tensor | None, + num_video_frames: int, + local_attn_size: int, + ) -> str | None: + """Return why state should reset before the next forward(), if any.""" if self.language is None: logger.info("language is None, resetting") - return True + return "session" if text_tokens is not None and not torch.equal(self.language, text_tokens): logger.info("language changed, resetting") - return True + return "session" # NOTE: after accumulate_frames, num_video_frames is the accumulated T # (1 for first call, 4 for subsequent). Only reset on true single-frame # which happens when the stitched_buffer was cleared externally. if num_video_frames == 1 and self.call_count > 1: logger.info("single frame input after first call, resetting") - return True + return "session" if local_attn_size != -1 and self.current_start_frame >= local_attn_size: logger.info( - "current_start_frame %d >= local_attn_size %d, resetting", self.current_start_frame, local_attn_size + "current_start_frame %d >= local_attn_size %d, resetting inference state", + self.current_start_frame, + local_attn_size, ) - return True + return "inference" - return False + return None + + def should_reset(self, text_tokens: torch.Tensor | None, num_video_frames: int, local_attn_size: int) -> bool: + """Determine if state should be reset before this forward().""" + return self.reset_reason(text_tokens, num_video_frames, local_attn_size) is not None # ------------------------------------------------------------------ # KV cache management diff --git a/vllm_omni/diffusion/models/dreamzero/video_export_worker.py b/vllm_omni/diffusion/models/dreamzero/video_export_worker.py index 20c8cd8cd..d039e53c8 100644 --- a/vllm_omni/diffusion/models/dreamzero/video_export_worker.py +++ b/vllm_omni/diffusion/models/dreamzero/video_export_worker.py @@ -9,13 +9,30 @@ import torch class DreamZeroVideoExportWorkerExtension: """DreamZero worker RPCs used by offline example video export.""" + @staticmethod + def _latents_to_uint8_frames(decoded: torch.Tensor) -> torch.Tensor: + decoded = decoded.squeeze(0).permute(1, 2, 3, 0).contiguous() + decoded = decoded.clamp(-1, 1) * 0.5 + 0.5 + return (decoded * 255.0).round().to(torch.uint8).cpu() + def decode_video_latents_to_uint8(self, video_latents: torch.Tensor) -> torch.Tensor: if self.model_runner is None or self.model_runner.pipeline is None: raise RuntimeError("DreamZero pipeline is not initialized on this worker.") with torch.inference_mode(): decoded = self.model_runner.pipeline.decode_video_latents(video_latents) - decoded = decoded.squeeze(0).permute(1, 2, 3, 0).contiguous() - decoded = decoded.clamp(-1, 1) * 0.5 + 0.5 - decoded = (decoded * 255.0).round().to(torch.uint8).cpu() - return decoded + return self._latents_to_uint8_frames(decoded) + + def decode_accumulated_video_latents_to_uint8(self, session_id: str = "default") -> torch.Tensor: + """Decode all AR-chunk latents accumulated on the server for one session.""" + if self.model_runner is None or self.model_runner.pipeline is None: + raise RuntimeError("DreamZero pipeline is not initialized on this worker.") + + with torch.inference_mode(): + decoded = self.model_runner.pipeline.decode_accumulated_video_latents(session_id) + return self._latents_to_uint8_frames(decoded) + + def clear_accumulated_video_latents(self, session_id: str = "default") -> None: + if self.model_runner is None or self.model_runner.pipeline is None: + raise RuntimeError("DreamZero pipeline is not initialized on this worker.") + self.model_runner.pipeline.clear_accumulated_video_latents(session_id) diff --git a/vllm_omni/diffusion/models/dreamzero/wan_vae_feat_cache_patch.py b/vllm_omni/diffusion/models/dreamzero/wan_vae_feat_cache_patch.py new file mode 100644 index 000000000..57d9b866c --- /dev/null +++ b/vllm_omni/diffusion/models/dreamzero/wan_vae_feat_cache_patch.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Make diffusers Wan VAE feat_cache compile-friendly. + +diffusers ``WanResample`` stores the string sentinel ``"Rep"`` in feat_cache slots. +``tensor == "Rep"`` returns a non-Tensor and breaks ``torch.compile(fullgraph=True)``. +We replace it with a scalar int8 tensor marker (value 0). +""" + +from __future__ import annotations + +from typing import Any + +import torch +from vllm.logger import init_logger + +logger = init_logger(__name__) + +_PATCHED = False + + +def _is_feat_rep(entry: Any) -> bool: + if not isinstance(entry, torch.Tensor): + return False + return entry.numel() == 1 and entry.dtype == torch.int8 + + +def _make_feat_rep(x: torch.Tensor) -> torch.Tensor: + return torch.zeros(1, device=x.device, dtype=torch.int8) + + +def _wan_resample_forward_tensorized(self, x, feat_cache=None, feat_idx=[0]): + """``WanResample.forward`` with tensor feat_cache sentinel (no ``"Rep"`` string).""" + from diffusers.models.autoencoders.autoencoder_kl_wan import CACHE_T + + b, c, t, h, w = x.size() + if self.mode == "upsample3d": + if feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = _make_feat_rep(x) + feat_idx[0] += 1 + else: + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None and not _is_feat_rep(feat_cache[idx]): + cache_x = torch.cat( + [feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], + dim=2, + ) + if cache_x.shape[2] < 2 and feat_cache[idx] is not None and _is_feat_rep(feat_cache[idx]): + cache_x = torch.cat([torch.zeros_like(cache_x).to(cache_x.device), cache_x], dim=2) + if _is_feat_rep(feat_cache[idx]): + x = self.time_conv(x) + else: + x = self.time_conv(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + + x = x.reshape(b, 2, c, t, h, w) + x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), 3) + x = x.reshape(b, c, t * 2, h, w) + t = x.shape[2] + x = x.permute(0, 2, 1, 3, 4).reshape(b * t, c, h, w) + x = self.resample(x) + x = x.view(b, t, x.size(1), x.size(2), x.size(3)).permute(0, 2, 1, 3, 4) + + if self.mode == "downsample3d": + if feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = x.clone() + feat_idx[0] += 1 + else: + cache_x = x[:, :, -1:, :, :].clone() + x = self.time_conv(torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2)) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + return x + + +def apply_wan_vae_feat_cache_tensor_patch() -> None: + """Monkeypatch diffusers WanResample once per process.""" + global _PATCHED + if _PATCHED: + return + + from diffusers.models.autoencoders.autoencoder_kl_wan import WanResample + + WanResample.forward = _wan_resample_forward_tensorized + _PATCHED = True + logger.info("DreamZero: applied Wan VAE feat_cache tensor sentinel patch.") diff --git a/vllm_omni/diffusion/worker/diffusion_model_runner.py b/vllm_omni/diffusion/worker/diffusion_model_runner.py index 278f1df40..288973637 100644 --- a/vllm_omni/diffusion/worker/diffusion_model_runner.py +++ b/vllm_omni/diffusion/worker/diffusion_model_runner.py @@ -168,8 +168,17 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): # Apply torch.compile if not in eager mode if not self.od_config.enforce_eager: if current_omni_platform.supports_torch_inductor(): - self._compile_transformer("transformer") - self._compile_transformer("transformer_2") + if hasattr(self.pipeline, "setup_compile"): + try: + self.pipeline.setup_compile() + except Exception as exc: + logger.warning( + "Model runner: setup_compile() failed (%s); running without compile.", + exc, + ) + else: + self._compile_transformer("transformer") + self._compile_transformer("transformer_2") else: logger.warning( "Model runner: Platform %s does not support torch inductor, skipping torch.compile.", @@ -306,8 +315,15 @@ class DiffusionModelRunner(OmniConnectorModelRunnerMixin): # num_inference_steps at all (i.e., just resets state and clears # stale residuals). num_inference_steps = req.sampling_params.num_inference_steps - if self.od_config.cache_backend == "tea_cache" and num_inference_steps is None: - num_inference_steps = 0 + if num_inference_steps is None and self.od_config.cache_backend in ( + "tea_cache", + "step_cache", + ): + # TeaCache refresh ignores the value; step_cache refresh is a + # no-op (per-chunk state resets in the denoise loop). DreamZero often + # leaves sampling_params.num_inference_steps unset and uses the + # pipeline default instead. + num_inference_steps = getattr(self.pipeline, "num_inference_steps", 0) or 0 if num_inference_steps is not None: self.cache_backend.refresh(self.pipeline, num_inference_steps) diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 721e58c7b..ac6140291 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -520,7 +520,7 @@ class AsyncOmniEngine: return mm_data = prompt.get("multi_modal_data") - if not isinstance(mm_data, Mapping) or not mm_data: + if not isinstance(mm_data, dict) or not mm_data: return from vllm.multimodal.hasher import MultiModalHasher @@ -842,6 +842,14 @@ class AsyncOmniEngine: "mag_max_skip_steps": 5, "mag_retention_ratio": 0.1, } + if cache_backend in ("step_cache"): + return { + "step_cache_dit_enabled": True, + "velocity_sim_thresholds": [0.95, 0.93], + "velocity_skip_countdowns": [4, 2], + "step_cache_dit_min_history": 2, + "step_cache_dit_max_history": 2, + } return None @staticmethod diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 905ff4ede..984706ae0 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -512,7 +512,7 @@ class OmniServeCommand(CLISubcommand): "--cache-backend", type=str, default="none", - help="Cache backend for diffusion models, options: 'tea_cache', 'cache_dit', 'mag_cache'", + help="Cache backend for diffusion models, options: 'tea_cache', 'cache_dit', 'mag_cache', 'step_cache'", ) omni_config_group.add_argument( "--cache-config", diff --git a/vllm_omni/entrypoints/openpi/connection.py b/vllm_omni/entrypoints/openpi/connection.py index 69f722ac8..b65264583 100644 --- a/vllm_omni/entrypoints/openpi/connection.py +++ b/vllm_omni/entrypoints/openpi/connection.py @@ -8,7 +8,13 @@ Protocol (compatible with OpenPI policy clients): Infer -> client sends msgpack(obs), server sends msgpack(ndarray) Reset -> client sends msgpack({endpoint:reset}), server sends msgpack(status) -NumPy values use the msgpack-numpy marker mapping: +NumPy values use msgpack-numpy marker mappings. Outbound payloads follow the +openpi-client format so official OpenPI clients can decode action responses: + ndarray -> {__ndarray__: true, dtype, shape, data} + scalar -> {__npgeneric__: true, dtype, data} + +Inbound payloads accept both the openpi-client markers above and the legacy +vLLM-native markers: ndarray -> {nd: true, type, kind, shape, data} scalar -> {nd: false, type, kind, data} """ @@ -41,18 +47,16 @@ def _pack_numpy(obj: Any) -> Any: if not obj.flags.c_contiguous: obj = np.ascontiguousarray(obj) return { - b"nd": True, + b"__ndarray__": True, b"data": obj.tobytes(), - b"type": obj.dtype.str, - b"kind": obj.dtype.kind, + b"dtype": obj.dtype.str, b"shape": obj.shape, } if isinstance(obj, np.generic): return { - b"nd": False, - b"data": obj.tobytes(), - b"type": obj.dtype.str, - b"kind": obj.dtype.kind, + b"__npgeneric__": True, + b"data": obj.item(), + b"dtype": obj.dtype.str, } raise TypeError(f"Unsupported type: {type(obj)!r}") @@ -67,27 +71,67 @@ def _decode_marker_text(value: Any) -> str: return str(value) +def _is_truthy_marker(value: Any) -> bool: + return value is True or value == 1 + + +def _decode_vllm_numpy_marker(obj: dict[Any, Any]) -> Any: + nd = _mapping_get(obj, "nd", _MISSING) + dtype = _mapping_get(obj, "type", _MISSING) + kind = _mapping_get(obj, "kind", _MISSING) + data = _mapping_get(obj, "data", _MISSING) + if nd is _MISSING or dtype is _MISSING or kind is _MISSING or data is _MISSING: + return _MISSING + + dtype_obj = np.dtype(_decode_marker_text(dtype)) + kind_text = _decode_marker_text(kind) + if dtype_obj.kind != kind_text: + raise ValueError(f"NumPy dtype marker kind mismatch: {dtype_obj.kind!r} != {kind_text!r}") + if dtype_obj.kind in ("V", "O", "c"): + raise ValueError(f"Unsupported dtype: {dtype_obj}") + + array = np.frombuffer(data, dtype=dtype_obj).copy() + if nd: + shape = _mapping_get(obj, "shape", _MISSING) + if shape is _MISSING: + raise ValueError("NumPy ndarray marker is missing shape") + return array.reshape(tuple(shape)) + return array[0] + + +def _decode_openpi_numpy_marker(obj: dict[Any, Any]) -> Any: + if _is_truthy_marker(_mapping_get(obj, "__ndarray__", False)): + data = _mapping_get(obj, "data", _MISSING) + dtype = _mapping_get(obj, "dtype", _MISSING) + shape = _mapping_get(obj, "shape", _MISSING) + if data is _MISSING or dtype is _MISSING or shape is _MISSING: + raise ValueError("OpenPI ndarray marker is missing required fields") + + dtype_obj = np.dtype(_decode_marker_text(dtype)) + if dtype_obj.kind in ("V", "O", "c"): + raise ValueError(f"Unsupported dtype: {dtype_obj}") + array = np.frombuffer(data, dtype=dtype_obj).copy() + return array.reshape(tuple(shape)) + + if _is_truthy_marker(_mapping_get(obj, "__npgeneric__", False)): + data = _mapping_get(obj, "data", _MISSING) + dtype = _mapping_get(obj, "dtype", _MISSING) + if data is _MISSING or dtype is _MISSING: + raise ValueError("OpenPI scalar marker is missing required fields") + dtype_obj = np.dtype(_decode_marker_text(dtype)) + if dtype_obj.kind in ("V", "O", "c"): + raise ValueError(f"Unsupported dtype: {dtype_obj}") + return dtype_obj.type(data) + + return _MISSING + + def _unpack_numpy(obj: Any) -> Any: if isinstance(obj, dict): - nd = _mapping_get(obj, "nd", _MISSING) - dtype = _mapping_get(obj, "type", _MISSING) - kind = _mapping_get(obj, "kind", _MISSING) - data = _mapping_get(obj, "data", _MISSING) - if nd is not _MISSING and dtype is not _MISSING and kind is not _MISSING and data is not _MISSING: - dtype_obj = np.dtype(_decode_marker_text(dtype)) - kind_text = _decode_marker_text(kind) - if dtype_obj.kind != kind_text: - raise ValueError(f"NumPy dtype marker kind mismatch: {dtype_obj.kind!r} != {kind_text!r}") - if dtype_obj.kind in ("V", "O", "c"): - raise ValueError(f"Unsupported dtype: {dtype_obj}") - - array = np.frombuffer(data, dtype=dtype_obj).copy() - if nd: - shape = _mapping_get(obj, "shape", _MISSING) - if shape is _MISSING: - raise ValueError("NumPy ndarray marker is missing shape") - return array.reshape(tuple(shape)) - return array[0] + for decoder in (_decode_vllm_numpy_marker, _decode_openpi_numpy_marker): + decoded = decoder(obj) + if decoded is not _MISSING: + return decoded return {key: _unpack_numpy(value) for key, value in obj.items()} if isinstance(obj, list): return [_unpack_numpy(value) for value in obj]