# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """DreamZero pipeline for vllm-omni. Entry point for DiffusionEngine.step_streaming() -> pipeline.forward(req) """ from __future__ import annotations import copy import json import logging import math import os import re as re_module from collections import OrderedDict from collections.abc import Iterable from contextlib import contextmanager from typing import ClassVar import numpy as np import torch import torch.nn as nn 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, ) from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin from vllm_omni.diffusion.distributed.parallel_state import get_classifier_free_guidance_world_size from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.models.dreamzero.causal_wan_model import CausalWanModel from vllm_omni.diffusion.models.dreamzero.image_encoder import DreamZeroImageEncoder from vllm_omni.diffusion.models.dreamzero.state_dreamzero import DreamZeroState from vllm_omni.diffusion.models.dreamzero.transform import ( DEFAULT_EMBODIMENT, ensure_transforms_loaded, ) from vllm_omni.diffusion.models.dreamzero.transform.base import get_transform from vllm_omni.diffusion.models.dreamzero.utils import ( DEFAULT_CFG_SCALE, DEFAULT_EMBODIMENT_NAME_TO_ID, DEFAULT_NEGATIVE_PROMPT, DEFAULT_NUM_INFERENCE_STEPS, DEFAULT_SEED, DEFAULT_SIGMA_SHIFT, ) from vllm_omni.diffusion.models.schedulers.scheduling_flow_unipc_multistep import FlowUniPCMultistepScheduler from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.experimental.ar_diffusion.capability import ( ARDiffusionCrossAttentionKVSpec, ARDiffusionKVBranchSpec, ARDiffusionKVCacheSpec, ) from vllm_omni.experimental.world_models.adapters.state_dreamzero_adapter import DreamZeroStateAdapter from vllm_omni.experimental.world_models.session_state import ( SessionStateManager, resolve_session_state_config, ) from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.transformers_utils.repo_utils import hf_api logger = logging.getLogger(__name__) MAX_DREAMZERO_SESSIONS = 64 # Shipped DreamZero geometry retains 24 Wan VAE causal-convolution cache # entries. This is the measured persistent CUDA upper bound per live session. DREAMZERO_MODEL_OWNED_STATE_BYTES_PER_SESSION = 603 * 1024 * 1024 # The pipeline's per-session state is a bespoke ``DreamZeroState`` by default, or # a ``DreamZeroStateAdapter`` view when the opt-in session manager is enabled. # The adapter mirrors ``DreamZeroState``'s surface, so every helper that reads or # writes session state accepts either. DreamZeroSessionState = DreamZeroState | DreamZeroStateAdapter class VideoActionScheduler: """Wraps video + action schedulers into single .step() interface.""" def __init__(self, video_scheduler, action_scheduler): self.video_scheduler = video_scheduler self.action_scheduler = action_scheduler def step(self, noise_pred, t, latents, return_dict=False, generator=None): video_out = self.video_scheduler.step( noise_pred[0], t[0], latents[0], return_dict=False, generator=generator, )[0] action_out = self.action_scheduler.step( noise_pred[1], t[1], latents[1], return_dict=False, generator=generator, )[0] return ((video_out, action_out),) # --------------------------------------------------------------------------- # DreamZeroPipeline # --------------------------------------------------------------------------- class DreamZeroPipeline(nn.Module, CFGParallelMixin): """DreamZero world model pipeline. Multi-output: predict_noise() returns (video_pred, action_pred). CFG: video gets standard CFG, action takes the positive KV branch only. KV is managed by the AR-Diffusion engine through the explicit capability methods below. The runner binds one session state only for ``forward()``. """ # Generic warmup cannot synthesize robot observations. AR-Diffusion uses # ar_diffusion_warmup_requests() for model-specific warmup instead. dummy_run_num_frames: ClassVar[int] = 0 _POSITIVE_BRANCH = "positive" _NEGATIVE_BRANCH = "negative" _ar_diffusion_kv_state = None state: DreamZeroSessionState | None def ar_diffusion_kv_cache_spec(self) -> ARDiffusionKVCacheSpec: """Describe DreamZero's local KV geometry to the generic runner.""" transformer = self.transformer frame_tokens = int(transformer.frame_seqlen) max_attention_tokens = int(transformer.blocks[0].self_attn.max_attention_size) cfg_world = int(get_classifier_free_guidance_world_size()) negative_local_index = 0 if cfg_world >= 2 else 1 cross_attention = [ ARDiffusionCrossAttentionKVSpec("text", int(transformer.text_len)), ] if transformer.model_type == "i2v": cross_attention.append(ARDiffusionCrossAttentionKVSpec("image", 257)) return ARDiffusionKVCacheSpec( num_layers=int(transformer.num_layers), num_kv_heads=int(transformer.blocks[0].self_attn.tp_num_heads), head_size=int(transformer.dim // transformer.num_heads), tokens_per_frame=frame_tokens, frames_per_block=int(transformer.num_frame_per_block), window_frames=max_attention_tokens // frame_tokens, kv_branches=( ARDiffusionKVBranchSpec(self._POSITIVE_BRANCH, 0), ARDiffusionKVBranchSpec(self._NEGATIVE_BRANCH, negative_local_index), ), session_capacity=MAX_DREAMZERO_SESSIONS, cross_attention=tuple(cross_attention), max_scratch_tokens_per_branch=int(transformer.num_action_per_block + transformer.num_state_per_block), model_owned_state_bytes_per_session=DREAMZERO_MODEL_OWNED_STATE_BYTES_PER_SESSION, ) @contextmanager def bind_ar_diffusion_state(self, session_id, state): """Bind runner-owned KV only for the duration of one forward.""" if self._ar_diffusion_kv_state is not None: raise RuntimeError("DreamZero AR-Diffusion state is already bound") if state.session_id != session_id: raise ValueError(f"DreamZero bound session mismatch: {state.session_id!r} != {session_id!r}") self._ar_diffusion_kv_state = state try: yield finally: self._ar_diffusion_kv_state = None def reset_ar_diffusion_session(self, session_id: str) -> None: """Reset DreamZero-owned state after the runner releases session KV.""" self._drop_ar_diffusion_session_state(session_id) def close_ar_diffusion_session(self, session_id: str) -> None: """Drop DreamZero-owned state after close, eviction, or failed forward.""" self._drop_ar_diffusion_session_state(session_id) def _drop_ar_diffusion_session_state(self, session_id: str) -> None: """Remove model state and clear the compatibility alias when it points there.""" key = str(session_id or "default") # Local binding narrows the Optional and guards lightweight test fixtures # that build the pipeline via __new__ without setting _memory_manager. manager = getattr(self, "_memory_manager", None) if manager is not None: # Manager-backed path: the session lives in the manager, not in # ``_states``. The runner's explicit close/reset is the session's # end-of-life signal, so release it there (freeing its buffers) and # drop the alias if it still views this session. manager.drop_session(key) current = getattr(self, "state", None) if isinstance(current, DreamZeroStateAdapter) and str(current.session_id or "default") == key: self.state = None return removed = self._states.pop(key, None) if removed is not None and getattr(self, "state", None) is removed: self.state = None @staticmethod def _ar_warmup_robot_obs(height: int, width: int, n_frames: int, session_id: str) -> dict: image = ( np.zeros((height, width, 3), dtype=np.uint8) if n_frames == 1 else np.zeros((n_frames, height, width, 3), dtype=np.uint8) ) return { "observation/exterior_image_0_left": image, "observation/exterior_image_1_left": image, "observation/wrist_image_left": image, "observation/joint_position": np.zeros(7, dtype=np.float32), "observation/cartesian_position": np.zeros(6, dtype=np.float32), "observation/gripper_position": np.zeros(1, dtype=np.float32), "prompt": "warmup", "session_id": session_id, } def ar_diffusion_warmup_requests(self, session_id: str) -> Iterable[OmniDiffusionRequest]: """Yield DreamZero-valid requests covering each resident window shape.""" spec = self.ar_diffusion_kv_cache_spec() n_forwards = 1 + math.ceil(max(0, spec.window_frames - 1) / spec.frames_per_block) raw_kv_config = getattr(self.od_config, "ar_diffusion_kv_config", None) if raw_kv_config is None and isinstance(self.od_config.model_config, dict): raw_kv_config = self.od_config.model_config.get("ar_diffusion_kv_config") capture_reset = ( bool(raw_kv_config.get("warmup_capture_reset", False)) if isinstance(raw_kv_config, dict) else bool(getattr(raw_kv_config, "warmup_capture_reset", False)) ) if capture_reset: n_forwards += 1 policy_config = self.od_config.model_config.get("policy_server_config", {}) height, width = (int(value) for value in policy_config.get("image_resolution", [180, 320])) for index in range(n_forwards): n_frames = 1 if index == 0 else 4 robot_obs = self._ar_warmup_robot_obs(height, width, n_frames, session_id) sampling_params = OmniDiffusionSamplingParams( extra_args={ "reset": index == 0, "session_id": session_id, "robot_obs": robot_obs, } ) yield OmniDiffusionRequest( prompt="warmup", sampling_params=sampling_params, request_id=f"ardiffusion-warmup-{index}", ) def _ar_branch(self, is_negative: bool) -> str: return self._NEGATIVE_BRANCH if is_negative else self._POSITIVE_BRANCH def _kv_get(self, state, is_negative, seq_len=None, update_kv_cache=False): return self._ar_diffusion_kv_state.get_kv_caches( self._ar_branch(is_negative), seq_len=seq_len, commit_current=update_kv_cache, ) def _kv_create(self, state, batch_size, dtype, device, num_layers, num_heads, head_dim): # The engine owns all KV allocation: self-attn is allocated lazily from # paged contexts, and cross-attn is populated eagerly in _kv_populate_cross. # Nothing is created model-side. return def _kv_commit(self, is_negative: bool): self._ar_diffusion_kv_state.commit_paged_context(self._ar_branch(is_negative)) def _kv_get_cross(self, state, is_negative): """Cross-attn cache from the engine pool (text k/v + I2V image k_img/v_img).""" kv_branch = self._ar_branch(is_negative) text_caches = self._ar_diffusion_kv_state.get_cross_attention_kv(kv_branch, "text") if "image" in self._ar_diffusion_kv_state.kv_cache.cross_attention_lengths: image_caches = self._ar_diffusion_kv_state.get_cross_attention_kv(kv_branch, "image") for text_cache, image_cache in zip(text_caches, image_caches, strict=True): text_cache["k_img"] = image_cache["k"] text_cache["v_img"] = image_cache["v"] return text_caches def _kv_populate_cross(self, context: torch.Tensor, clip_feature, is_negative: bool) -> None: """Eagerly project cross-attn K/V for all layers into the AR-Diffusion pool. Caches the session-invariant cross-attn projections once, per half: the text ``k``/``v`` from ``text_embedding(context)`` survive window-boundary resets (prompt unchanged within a session — only session resets clear them), while the I2V image-token ``k_img``/``v_img`` from ``img_emb(clip_feature)`` (the 257 image tokens the forward splits off, cached model-side by #4154) are re-projected on every window restart from the fresh CLIP features. Must run after the image is encoded so ``clip_feature`` is available. """ s = self._ar_diffusion_kv_state kv_branch = self._ar_branch(is_negative) need_text = not s.is_cross_attention_populated(kv_branch, "text") need_img = ( clip_feature is not None and self.transformer.model_type == "i2v" and not s.is_cross_attention_populated(kv_branch, "image") ) if not need_text and not need_img: return projected = self.transformer.text_embedding(context) if need_text else None img_ctx = self.transformer.img_emb(clip_feature) if need_img else None if projected is not None: def text_layer_kv(): for block in self.transformer.blocks: ca = block.cross_attn n, d = ca.tp_num_heads, ca.head_dim yield ( ca.norm_k(ca.k(projected)).unflatten(2, (n, d)), ca.v(projected).unflatten(2, (n, d)), ) s.populate_cross_attention(kv_branch, "text", text_layer_kv()) if img_ctx is not None: def image_layer_kv(): for block in self.transformer.blocks: ca = block.cross_attn n, d = ca.tp_num_heads, ca.head_dim yield ( ca.norm_k_img(ca.k_img(img_ctx)).unflatten(2, (n, d)), ca.v_img(img_ctx).unflatten(2, (n, d)), ) s.populate_cross_attention(kv_branch, "image", image_layer_kv()) logger.info( "AR-Diffusion CROSS POPULATE [%s]: %d layers, text=%s img=%s", "neg" if is_negative else "pos", len(self.transformer.blocks), "kept" if projected is None else tuple(context.shape), None if img_ctx is None else tuple(img_ctx.shape), ) def _kv_reset(self, state, *, clear_video_latents: bool = True): """Reset the engine's pooled session window plus the model's non-KV state. DreamZero resets at the attention-window boundary; the engine pool drops the same window so the next forward starts fresh. ``clear_video_latents=False`` keeps the accumulated video latents for export. ``clear_video_latents=False`` also marks a window ("inference") reset: the prompt is unchanged, so the pool keeps the text cross-attn K/V and only the image half repopulates on the restart forward. """ state.reset(clear_video_latents=clear_video_latents) keep_cross = ("text",) if not clear_video_latents else () self._ar_diffusion_kv_state.reset(keep_cross_attention=keep_cross) def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = "") -> None: """Initialize pipeline components. DreamZero root checkpoint layout (GEAR-Dreams/DreamZero-DROID): config.json -- root config (action_head_cfg, architectures, etc.) model-*.safetensors -- all learned weights (action_head.{model,text_encoder,image_encoder,vae}.*) experiment_cfg/metadata.json -- per-embodiment action normalization stats vae/ -- symlink to Wan2.1 VAE (diffusers-compatible) Components are instantiated from config (not from_pretrained), then filled by load_weights() which reads root safetensors and remaps key prefixes. Exceptions: - tokenizer loads from `google/umt5-xxl` - VAE uses `DistributedAutoencoderKLWan` as the local execution module. It can be bootstrapped either from an explicit diffusers source (`od_config.model_paths["vae"]`) or directly from constructor defaults that match Wan2.1 VAE, after which DreamZero root `action_head.vae.*` weights are remapped onto that module in `load_weights()` """ super().__init__() # DreamZero is engine-only: every KV access in forward() routes through # the AR-Diffusion engine's pool-backed state. Fail fast here — a stale # or programmatic config that leaves engine_backend="default" would # otherwise only crash mid-forward on the first KV access. engine_backend = str(getattr(od_config, "engine_backend", "") or "") if "ar_diffusion" not in engine_backend.lower().replace("-", "_"): raise ValueError( "DreamZeroPipeline requires the AR-Diffusion engine; set " "engine_backend: vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine " f"in the deploy config (got engine_backend={engine_backend!r})." ) model_path = od_config.model model_config = od_config.model_config local_files_only = os.path.exists(model_path) self.od_config = od_config ensure_transforms_loaded() self.default_robot_embodiment = model_config.get( "default_robot_embodiment", DEFAULT_EMBODIMENT, ) root_cfg = self._load_repo_json(model_path, "config.json", local_files_only) if root_cfg is None: raise ValueError(f"DreamZero requires root config.json in {model_path}.") action_head_cfg = root_cfg["action_head_cfg"] ah_config = action_head_cfg["config"] diffusion_model_cfg = ah_config["diffusion_model_cfg"] # ---- Tokenizer ---- tokenizer_source = od_config.model_paths.get("tokenizer", "google/umt5-xxl") self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_source) # Instantiate from config; weights load through `load_weights()`. umt5_config = UMT5Config( d_model=4096, d_ff=10240, num_heads=64, num_layers=24, vocab_size=256384, relative_attention_num_buckets=32, relative_attention_max_distance=128, dense_act_fn="gelu_new", feed_forward_proj="gated-gelu", is_encoder_decoder=False, ) self.text_encoder = UMT5EncoderModel(umt5_config) self.image_encoder = DreamZeroImageEncoder() # Build a compatible VAE module, then fill it through `load_weights()`. vae_source = od_config.model_paths.get("vae") if vae_source: self.vae = DistributedAutoencoderKLWan.from_pretrained( vae_source, torch_dtype=torch.float32, ) elif local_files_only and os.path.isdir(os.path.join(model_path, "vae")): self.vae = DistributedAutoencoderKLWan.from_pretrained( model_path, subfolder="vae", torch_dtype=torch.float32, ) else: self.vae = DistributedAutoencoderKLWan() self.vae.init_distributed() if not ( getattr(od_config, "enable_cpu_offload", False) or getattr(od_config, "enable_layerwise_offload", False) ): self.vae = self.vae.to(device=get_local_device(), dtype=od_config.dtype) self.register_buffer( "vae_latents_mean", torch.tensor(self.vae.config.latents_mean, dtype=torch.float32).view(1, -1, 1, 1, 1), persistent=False, ) self.register_buffer( "vae_latents_inv_std", (1.0 / torch.tensor(self.vae.config.latents_std, dtype=torch.float32)).view(1, -1, 1, 1, 1), persistent=False, ) # Filter out keys not accepted by `CausalWanModel.__init__`. transformer_kwargs = {k: v for k, v in diffusion_model_cfg.items() if k not in ("_convert_", "_target_")} transformer_kwargs["action_dim"] = ah_config["action_dim"] transformer_kwargs["max_state_dim"] = ah_config["max_state_dim"] transformer_kwargs["num_frame_per_block"] = ah_config["num_frame_per_block"] self.transformer = CausalWanModel(**transformer_kwargs) self.scheduler = FlowUniPCMultistepScheduler( num_train_timesteps=1000, shift=1, use_dynamic_shifting=False, ) # Read before the first `_get_or_create_state` below: the manager-backed # state bounds its VAE encoder history to this many latent frames. self.num_frame_per_block: int = ah_config["num_frame_per_block"] self._states: OrderedDict[str, DreamZeroState] = OrderedDict() # Opt-in: back per-session state with the shared SessionStateManager # (RFC #4480). Default off -> the bespoke DreamZeroState path above. self._use_memory_manager, mm_max_sessions = resolve_session_state_config( enable=od_config.enable_session_state_manager, max_sessions=MAX_DREAMZERO_SESSIONS, ) self._memory_manager: SessionStateManager | None = ( SessionStateManager(max_sessions=mm_max_sessions) if self._use_memory_manager else None ) if self._use_memory_manager: logger.info("DreamZero: session state manager enabled (max_sessions=%d)", mm_max_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", DEFAULT_NUM_INFERENCE_STEPS, ) self.cfg_scale: float = model_config.get("cfg_scale", DEFAULT_CFG_SCALE) self.sigma_shift: float = model_config.get("sigma_shift", DEFAULT_SIGMA_SHIFT) self.num_frames: int = ah_config["num_frames"] self.action_horizon: int = ah_config["action_horizon"] self.decouple_inference_noise: bool = ah_config["decouple_inference_noise"] self.video_inference_final_noise: float = ah_config["video_inference_final_noise"] self.seed: int = model_config.get("seed", DEFAULT_SEED) # Model-level constants for state/action padding. self.max_state_dim: int = ah_config["max_state_dim"] self.max_action_dim: int = ah_config["max_action_dim"] self.negative_prompt: str = model_config.get("negative_prompt", DEFAULT_NEGATIVE_PROMPT) # The negative prompt is a model constant: encode it once, lazily, and # reuse across every forward/session (UMT5 encode is deterministic). self._negative_prompt_embeds_cache: torch.Tensor | None = None # Embodiment name -> numeric ID mapping (model knowledge) self.embodiment_name_to_id: dict[str, int] = model_config.get( "embodiment_name_to_id", DEFAULT_EMBODIMENT_NAME_TO_ID, ) # Prefer root `experiment_cfg/metadata.json`, then `model_config`. stats_path = model_config.get("action_norm_stats_path") metadata = self._load_repo_json(model_path, "experiment_cfg/metadata.json", local_files_only) if metadata is not None: self.action_norm_stats = self._parse_action_norm_stats(metadata) self.state_norm_stats = self._parse_state_norm_stats(metadata) elif stats_path: self.action_norm_stats = self._load_action_norm_stats(stats_path) self.state_norm_stats = {} else: self.action_norm_stats: dict[str, dict[str, torch.Tensor]] = {} self.state_norm_stats: dict[str, dict[str, torch.Tensor]] = {} # Whether model uses relative actions (need to add back last state) self.relative_action: bool = model_config.get("relative_action", True) # Number of action dims that are relative (DROID: 7 = joint only, gripper is absolute) self.relative_action_dim: int = model_config.get("relative_action_dim", 7) self._weights_sources = [ DiffusersPipelineLoader.ComponentSource( model_or_path=model_path, subfolder=None, revision=None, prefix="", fall_back_to_pt=False, allow_patterns_overrides=[ "model-*.safetensors", "model.safetensors", ], ), ] def _get_or_create_state(self, session_id: str | None) -> DreamZeroState | DreamZeroStateAdapter: # getattr guards lightweight test fixtures that build the pipeline via # __new__ and seed only the bespoke fields, never setting _memory_manager. if getattr(self, "_memory_manager", None) is not None: # The manager owns session lifecycle and LRU; the adapter is a thin # per-call view over it (no second LRU to drift). return DreamZeroStateAdapter( session_id, self._memory_manager, vae_encoder_window=self.num_frame_per_block, ) session_key = str(session_id or "default") state = self._states.get(session_key) if state is None: state = DreamZeroState() self._states[session_key] = state else: self._states.move_to_end(session_key) return state # ----------------------------------------------------------------------- # Root config loading # ----------------------------------------------------------------------- @staticmethod def _load_repo_json(model_path: str, relative_path: str, local_files_only: bool) -> dict | None: """Load a JSON file from a local checkpoint directory or HF repo.""" if local_files_only and os.path.isdir(model_path): json_path = os.path.join(model_path, relative_path) if not os.path.exists(json_path): return None with open(json_path) as f: return json.load(f) try: json_path = hf_api().hf_hub_download(model_path, relative_path) with open(json_path) as f: return json.load(f) except Exception: logger.warning("Failed to load %s from %s", relative_path, model_path) return None # ----------------------------------------------------------------------- # CFGParallelMixin overrides # ----------------------------------------------------------------------- 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"], context=kwargs["encoder_hidden_states"], seq_len=kwargs["seq_len"], kv_cache=kwargs["kv_cache"], crossattn_cache=kwargs["crossattn_cache"], current_start_frame=kwargs["current_start_frame"], y=kwargs.get("y"), clip_feature=kwargs.get("clip_feature"), action=kwargs.get("action"), timestep_action=kwargs.get("timestep_action"), state=kwargs.get("state_features"), embodiment_id=kwargs.get("embodiment_id"), ) if kwargs.get("update_kv_cache", False): is_neg = kwargs.get("is_negative", False) logger.debug( "AR-Diffusion pipeline predict_noise -> commit paged context: " "is_neg=%s seq_len=%s current_start_frame=%s layers=%d", is_neg, kwargs.get("seq_len"), kwargs.get("current_start_frame"), len(updated_kv_caches), ) self._kv_commit(is_neg) 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). Incremental VAE encode (``_vae_encode_encoder_chunk``) stays eager because Wan ``feat_cache`` mutation is incompatible with CUDAGraph capture. DiT blocks use per-block ``fullgraph=True``. """ 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. # The AR-Diffusion paged self-attention is a registered custom op, so the # block stays fullgraph even on that path. 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.accelerator.is_available(): return state = self.state if state is None: state = self._get_or_create_state("default") self.state = state 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, state=state) except Exception as exc: logger.warning("DreamZero compile warmup (image_encoder) skipped: %s", exc) try: state.reset_vae_encoder_stream() dummy_video = torch.zeros(1, 3, 1, 180, 320, dtype=torch.bfloat16, device=device) self._vae_stream_seed(state, dummy_video[:, :, :1]) for _ in range(4): self._vae_stream_append_frame(state, dummy_video[:, :, :1]) self._vae_stream_get_observation_latents( state, self.num_frame_per_block, dtype=torch.bfloat16, ) except Exception as exc: logger.warning("DreamZero compile warmup (vae encode stream) 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, positive_noise_pred: torch.Tensor | tuple[torch.Tensor, ...], negative_noise_pred: torch.Tensor | tuple[torch.Tensor, ...], true_cfg_scale: float, cfg_normalize: bool = False, ) -> torch.Tensor | tuple[torch.Tensor, ...]: """Video: standard CFG. Action: positive only (no CFG). action = cond only (no uncond blending) """ (video_pos, action_pos) = positive_noise_pred (video_neg, _) = negative_noise_pred video_combined = super().combine_cfg_noise(video_pos, video_neg, true_cfg_scale, cfg_normalize) return (video_combined, action_pos) # ----------------------------------------------------------------------- # ----------------------------------------------------------------------- def _synchronize_cfg_parallel_step_output( self, latents: tuple[torch.Tensor, torch.Tensor], do_true_cfg: bool, ) -> tuple[torch.Tensor, torch.Tensor]: """Post-step sync: .contiguous() + accelerator stream synchronize""" latents = tuple(t.contiguous() for t in latents) if do_true_cfg and get_classifier_free_guidance_world_size() > 1: device = next((t.device for t in latents if t.device.type != "cpu"), None) if device is not None: torch.accelerator.current_stream(device).synchronize() return latents # ----------------------------------------------------------------------- # Video preprocessing # ----------------------------------------------------------------------- def _preprocess_video(self, videos: torch.Tensor) -> torch.Tensor: """uint8 [B,T,H,W,C] -> bfloat16 [B,C,T,H,W] normalized to [-1,1].""" videos = videos.permute(0, 4, 1, 2, 3) if videos.dtype == torch.uint8: videos = videos.float() / 255.0 # Cast to bf16 before normalization to preserve input rounding. videos = videos.to(dtype=torch.bfloat16) b, c, t, h, w = videos.shape videos = videos.permute(0, 2, 1, 3, 4) videos = videos.reshape(b * t, c, h, w) videos = videos * 2.0 - 1.0 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, attention_mask, ).last_hidden_state prompt_emb = prompt_emb.clone().to(dtype=torch.bfloat16) for i, v in enumerate(seq_lens): prompt_emb[:, v:] = 0 return prompt_emb # ----------------------------------------------------------------------- # Image encoding # ----------------------------------------------------------------------- def _encode_image( self, image: torch.Tensor, num_frames: int, height: int, width: int, *, state: DreamZeroSessionState | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Encode first frame via CLIP + VAE. Returns: (clip_feas, ys, image_latent) """ device = image.device 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) msk[:, 1:] = 0 msk = torch.concat([torch.repeat_interleave(msk[:, 0:1], repeats=4, dim=1), msk[:, 1:]], dim=1) msk = msk.view(batch_size, msk.shape[1] // 4, 4, height // 8, width // 8) msk = msk.transpose(1, 2) latent_dtype = image.dtype image_input = image.transpose(1, 2) image_zeros = torch.zeros( batch_size, 3, num_frames - 1, height, width, dtype=latent_dtype, device=device, ) vae_input = torch.concat([image_input, image_zeros], dim=2) y = self._encode_vae_latents(vae_input) y = y.to(dtype=latent_dtype) new_image = y[:, :, 0:1] y = torch.concat([msk, y], dim=1) if state is not None: if not state.vae_stream_initialized: # Seed the AR streaming encoder after the compiled full encode above; # ``cudagraph_mark_step_begin`` isolates eager feat_cache work from CUDAGraph. self._vae_stream_seed(state, image_input[:, :, :1]) else: # Window ("inference") restart with a live stream: keep the real # frame history in the Wan feat_cache and append the restart # observation instead of reseeding from scratch. self._cudagraph_mark_step_begin() self._vae_stream_append_frame(state, image_input[:, :, :1]) return clip_context, y, new_image def _vae_patchify(self, videos: torch.Tensor) -> torch.Tensor: if self.vae.config.patch_size is not None: from diffusers.models.autoencoders.autoencoder_kl_wan import patchify return patchify(videos, patch_size=self.vae.config.patch_size) return videos @staticmethod def _vae_clone_feat_map(feat_map: list[torch.Tensor | None]) -> list[torch.Tensor | None]: return [entry.clone() if isinstance(entry, torch.Tensor) else entry for entry in feat_map] def _vae_init_enc_feat_map(self) -> list[torch.Tensor | None]: self.vae.clear_cache() return self._vae_clone_feat_map(self.vae._enc_feat_map) def _vae_encode_encoder_chunk( self, chunk: torch.Tensor, feat_map: list[torch.Tensor | None], ) -> tuple[torch.Tensor, list[torch.Tensor | None]]: """Run one Wan encoder chunk while mutating ``feat_map`` in place. Must stay eager (not ``torch.compile``): Wan causal ``feat_cache`` updates conflict with ``reduce-overhead`` CUDAGraph capture. """ self.vae._enc_feat_map = feat_map self.vae._enc_conv_idx = [0] out = self.vae.encoder( chunk, feat_cache=self.vae._enc_feat_map, feat_idx=self.vae._enc_conv_idx, ) return out, self._vae_clone_feat_map(self.vae._enc_feat_map) def _vae_quantize_encoder_out(self, encoder_out: torch.Tensor) -> torch.Tensor: enc = self.vae.quant_conv(encoder_out) mu, _ = enc.chunk(2, dim=1) mean = self.vae_latents_mean.to(device=mu.device, dtype=mu.dtype) inv_std = self.vae_latents_inv_std.to(device=mu.device, dtype=mu.dtype) return (mu - mean) * inv_std def _vae_stream_seed(self, state: DreamZeroSessionState, first_frame: torch.Tensor) -> None: """Seed incremental VAE encode with the first observation frame.""" state.reset_vae_encoder_stream() feat_map = self._vae_init_enc_feat_map() chunk = self._vae_patchify(first_frame.to(dtype=self.vae.dtype)) self._cudagraph_mark_step_begin() encoder_out, feat_map = self._vae_encode_encoder_chunk(chunk[:, :, :1], feat_map) state.vae_enc_feat_map = feat_map state.vae_encoder_out = encoder_out state.vae_pending_body_frames = None state.vae_stream_initialized = True def _vae_stream_append_frame(self, state: DreamZeroSessionState, new_frame: torch.Tensor) -> None: """Append one pixel frame and encode a 4-frame body chunk when ready.""" if not state.vae_stream_initialized or state.vae_enc_feat_map is None or state.vae_encoder_out is None: raise RuntimeError("VAE encoder stream is not initialized.") frame = new_frame.to(dtype=self.vae.dtype) if state.vae_pending_body_frames is None: state.vae_pending_body_frames = frame else: state.vae_pending_body_frames = torch.cat([state.vae_pending_body_frames, frame], dim=2) while state.vae_pending_body_frames is not None and state.vae_pending_body_frames.shape[2] >= 4: body = state.vae_pending_body_frames[:, :, :4] if state.vae_pending_body_frames.shape[2] > 4: state.vae_pending_body_frames = state.vae_pending_body_frames[:, :, 4:] else: state.vae_pending_body_frames = None chunk = self._vae_patchify(body) encoder_chunk, feat_map = self._vae_encode_encoder_chunk(chunk, state.vae_enc_feat_map) state.append_vae_encoder_chunk(encoder_chunk) state.vae_enc_feat_map = feat_map def _vae_stream_get_observation_latents( self, state: DreamZeroSessionState, num_latent_frames: int, *, dtype: torch.dtype, ) -> torch.Tensor: # Read once: the manager-backed state rebuilds this from its ring. encoder_out = state.vae_encoder_out if encoder_out is None: raise RuntimeError("VAE encoder stream has no accumulated encoder output.") # A state that bounds its encoder history keeps only the most recent # frames. Asking for more than it retains would silently return a # shorter (and differently padded) window, so fail instead. window = state.vae_encoder_window if window is not None and num_latent_frames > window: raise ValueError(f"Requested {num_latent_frames} latent frames but the session retains at most {window}.") latents = self._vae_quantize_encoder_out(encoder_out).to(dtype=dtype) if latents.shape[2] >= num_latent_frames: return latents[:, :, -num_latent_frames:] pad_count = num_latent_frames - latents.shape[2] pad = latents[:, :, -1:].expand(-1, -1, pad_count, -1, -1) return torch.cat([pad, latents], dim=2) def _preprocess_vae_observation_window(self, videos: torch.Tensor) -> torch.Tensor: _, _, num_frames_raw, _, _ = videos.shape if (num_frames_raw - 1) // 4 == self.num_frame_per_block: return videos if num_frames_raw // 4 != self.num_frame_per_block: repeat_factor = self.num_frame_per_block // (num_frames_raw // 4) videos = torch.repeat_interleave(videos, repeat_factor, dim=2) first_frame = videos[:, :, 0:1] return torch.cat([first_frame, videos], dim=2) first_frame = videos[:, :, 0:1] return torch.cat([first_frame, videos], dim=2) def _encode_observation_latents( self, state: DreamZeroSessionState, videos: torch.Tensor, *, latent_dtype: torch.dtype, ) -> torch.Tensor: """Encode current robot observation into normalized VAE latents.""" if state.vae_stream_initialized: self._cudagraph_mark_step_begin() self._vae_stream_append_frame(state, videos[:, :, -1:]) return self._vae_stream_get_observation_latents( state, self.num_frame_per_block, dtype=latent_dtype, ) videos = self._preprocess_vae_observation_window(videos) return self._encode_vae_latents(videos).to(dtype=latent_dtype) def _encode_vae_latents(self, videos: torch.Tensor) -> torch.Tensor: """Encode videos into normalized VAE latents.""" input_dtype = videos.dtype self._cudagraph_mark_step_begin() hidden = self.vae._encode(videos.to(dtype=self.vae.dtype)) mu, _ = hidden.chunk(2, dim=1) mean = self.vae_latents_mean.to(device=mu.device, dtype=mu.dtype) inv_std = self.vae_latents_inv_std.to(device=mu.device, dtype=mu.dtype) mu = (mu - mean) * inv_std return mu.to(dtype=input_dtype) def decode_video_latents(self, video_latents: torch.Tensor) -> torch.Tensor: """Decode normalized VAE latents into RGB video tensors.""" vae_dtype = self.vae.dtype vae_device = next(self.vae.parameters()).device latents = video_latents.to(device=vae_device, dtype=vae_dtype) mean = self.vae_latents_mean.to(device=vae_device, dtype=vae_dtype) 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 # ----------------------------------------------------------------------- def _prefill_kv_cache( self, image_latents: torch.Tensor, prompt_embeds: torch.Tensor, negative_prompt_embeds: torch.Tensor | None, frame_seqlen: int, seq_len: int, do_true_cfg: bool, state: DreamZeroSessionState, ) -> None: """Prefill KV cache with first frame and/or current observation. Uses predict_noise_maybe_with_cfg() for CFG parallel -- same path as the denoise loop. The mixin handles rank dispatch automatically. KV cache update happens as a side effect inside predict_noise(). """ batch_size = image_latents.shape[0] device = image_latents.device dtype = image_latents.dtype num_heads = getattr(self.transformer.blocks[0].self_attn, "tp_num_heads", self.transformer.num_heads) head_dim = self.transformer.dim // self.transformer.num_heads if state.current_start_frame == 0: self._kv_create( state, batch_size, dtype, device, self.transformer.num_layers, num_heads, head_dim, ) zero_t = torch.zeros([batch_size, 1], device=device, dtype=torch.long) y_first = state.ys[:, :, 0:1] if state.ys is not None else None # KV cache update is a side effect in predict_noise() common = dict( hidden_states=image_latents.transpose(1, 2), timestep_video=zero_t, seq_len=frame_seqlen, current_start_frame=0, y=y_first, clip_feature=state.clip_feas, update_kv_cache=True, dreamzero_state=state, ) positive_kwargs = dict( encoder_hidden_states=prompt_embeds, kv_cache=self._kv_get(state, False, seq_len=frame_seqlen, update_kv_cache=True), crossattn_cache=self._kv_get_cross(state, False), is_negative=False, **common, ) negative_kwargs = ( dict( encoder_hidden_states=negative_prompt_embeds, kv_cache=self._kv_get(state, True, seq_len=frame_seqlen, update_kv_cache=True), crossattn_cache=self._kv_get_cross(state, True), is_negative=True, **common, ) if negative_prompt_embeds is not None else None ) 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, ) state.current_start_frame = 1 if state.current_start_frame != 1: csf = state.current_start_frame nfpb = self.num_frame_per_block current_ref = image_latents[:, -nfpb:] if state.ys is not None and csf <= state.ys.shape[2]: y = state.ys[:, :, csf - nfpb : csf] elif state.ys is not None: y = state.ys[:, :, -nfpb:] else: y = None zero_t = torch.zeros([batch_size, nfpb], device=device, dtype=torch.long) common = dict( hidden_states=current_ref.transpose(1, 2), timestep_video=zero_t, seq_len=seq_len, current_start_frame=csf - nfpb, y=y, clip_feature=state.clip_feas, update_kv_cache=True, dreamzero_state=state, ) positive_kwargs = dict( encoder_hidden_states=prompt_embeds, kv_cache=self._kv_get(state, False, seq_len=seq_len, update_kv_cache=True), crossattn_cache=self._kv_get_cross(state, False), is_negative=False, **common, ) negative_kwargs = ( dict( encoder_hidden_states=negative_prompt_embeds, kv_cache=self._kv_get(state, True, seq_len=seq_len, update_kv_cache=True), crossattn_cache=self._kv_get_cross(state, True), is_negative=True, **common, ) if negative_prompt_embeds is not None else None ) 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, ) def diffuse( self, video_latents: torch.Tensor, action_latents: torch.Tensor, timesteps_video: torch.Tensor, timesteps_action: torch.Tensor, prompt_embeds: torch.Tensor, negative_prompt_embeds: torch.Tensor | None, video_action_scheduler: VideoActionScheduler, do_true_cfg: bool, state: DreamZeroSessionState, **kwargs, ) -> tuple[torch.Tensor, torch.Tensor]: """Denoising loop with CFG parallel support. For each timestep: 1. Build positive_kwargs / negative_kwargs 2. predict_noise_maybe_with_cfg() -> (video_pred, action_pred) 3. scheduler_step_maybe_with_cfg() -> VideoActionScheduler 4. _synchronize_cfg_parallel_step_output() """ seq_len = kwargs["seq_len"] state_features = kwargs.get("state_features") embodiment_id = kwargs.get("embodiment_id") # Shared kwargs for predict_noise (both conditional and unconditional branches) common_kwargs = dict( seq_len=seq_len, current_start_frame=state.current_start_frame, state_features=state_features, embodiment_id=embodiment_id, update_kv_cache=False, dreamzero_state=state, ) 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] batch_size = noisy_input.shape[0] timestep = ( torch.ones( [batch_size, self.num_frame_per_block], device=noisy_input.device, dtype=torch.int64, ) * video_timestep ) timestep_action = ( torch.ones( [batch_size, self.action_horizon], device=noisy_input.device, dtype=torch.int64, ) * action_timestep ) csf = state.current_start_frame ys = state.ys if ys is None: raise RuntimeError("diffuse() requires state.ys, populated by the forward pass before denoising") if csf + self.num_frame_per_block <= ys.shape[2]: y = ys[:, :, csf : csf + self.num_frame_per_block] else: y = ys[:, :, -self.num_frame_per_block :] 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=prompt_embeds, kv_cache=self._kv_get(state, False, seq_len=seq_len, update_kv_cache=False), crossattn_cache=self._kv_get_cross(state, False), y=y, clip_feature=state.clip_feas, action=noisy_input_action, timestep_action=timestep_action, is_negative=False, **common_kwargs, ) 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=self._kv_get(state, True, seq_len=seq_len, update_kv_cache=False), crossattn_cache=self._kv_get_cross(state, 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) noise_pred_tuple = (flow_pred.transpose(1, 2), flow_pred_action) step_output = video_action_scheduler.step( noise_pred_tuple, t, latents, generator=kwargs.get("generator"), ) noisy_input, noisy_input_action = step_output[0] noisy_input, noisy_input_action = self._synchronize_cfg_parallel_step_output( (noisy_input, noisy_input_action), do_true_cfg, ) return noisy_input, noisy_input_action # ----------------------------------------------------------------------- # Main entry point # ----------------------------------------------------------------------- def _transform_robot_obs(self, robot_obs: dict): """Select DreamZero robot transform and convert raw obs to model input.""" embodiment = robot_obs.get("embodiment", self.default_robot_embodiment) transform = get_transform(embodiment) return transform, transform.transform_input(robot_obs) @torch.no_grad() def forward(self, req: DiffusionRequestBatch, **kwargs) -> DiffusionOutput: """Full inference step. Called by DiffusionEngine.step_streaming().""" extra_args = req.sampling_params.extra_args or {} robot_obs = extra_args.get("robot_obs") if robot_obs is None: first_prompt = req.prompts[0] if req.prompts else "" prompt = first_prompt if isinstance(first_prompt, str) else (first_prompt.get("prompt") or "") is_dummy_warmup = prompt == "dummy run" and req.sampling_params.num_inference_steps == 1 if is_dummy_warmup: logger.info("Skipping DreamZero dummy warmup request without robot_obs.") return DiffusionOutput( output={ "actions": np.zeros( (self.action_horizon, self.max_action_dim), dtype=np.float32, ), }, ) raise KeyError("robot_obs") session_id = str(extra_args.get("session_id") or "default") state = self._get_or_create_state(session_id) self.state = state transform, unified_obs = self._transform_robot_obs(robot_obs) device = get_local_device() # ---- Step 1: Extract inputs from unified observation ---- prompt_str = unified_obs["prompt"] # str (templated) stitched = unified_obs["images"] # ndarray (T,H,W,C) from transform if not isinstance(stitched, np.ndarray): stitched = np.asarray(stitched) embodiment_name = unified_obs["embodiment_name"] embodiment_id = torch.tensor( # (B,) tensor for CategorySpecificMLP [self.embodiment_name_to_id[embodiment_name]], dtype=torch.long, device=device, ) # State: raw from transform -> pad to (B, state_horizon=1, max_state_dim) raw_state = unified_obs["state"] state_for_postprocess = None if raw_state is not None: if not isinstance(raw_state, np.ndarray): raw_state = np.asarray(raw_state, dtype=np.float64) raw_state = raw_state.flatten() padded = np.zeros(self.max_state_dim, dtype=np.float64) n = min(len(raw_state), self.max_state_dim) padded[:n] = raw_state[:n] state_for_postprocess = ( torch.from_numpy(padded) .reshape(1, 1, self.max_state_dim) .to( device=device, dtype=torch.float32, ) ) state_features = self._normalize_state( state_for_postprocess, embodiment_name, ).to(dtype=torch.bfloat16) else: state_features = None # ---- Step 1b: Tokenize ---- (wan2_2 convention: pipeline owns tokenizer) text_inputs = self.tokenizer( prompt_str, max_length=512, padding="max_length", truncation=True, return_tensors="pt", add_special_tokens=True, ) text_tokens = text_inputs["input_ids"].to(device) attention_mask = text_inputs["attention_mask"].to(device) # Explicit request reset is handled by ARDiffusionModelRunner before it # binds this forward. Model-detected prompt/window resets still happen # here because they depend on tokenized DreamZero state. reset_reason = state.reset_reason(text_tokens, 0, self.transformer.local_attn_size) if reset_reason == "session": self._kv_reset(state) elif reset_reason == "inference": self._kv_reset(state, clear_video_latents=False) state.language = text_tokens # Frame accumulation: stitched single frame -> multi-frame video video_frames = state.accumulate_frames(stitched) # (T, H, W, C) videos = torch.from_numpy(video_frames).unsqueeze(0).to(device) # (B=1, T, H, W, C) videos = self._preprocess_video(videos) # -> [B,C,T,H,W] bf16 _, _, num_frames_raw, height, width = videos.shape # Optional phase timing (DZ_PHASE_TIMING=1): logs per-forward stage costs # (text encode / obs VAE encode / KV prefill / denoise) at INFO. Each mark # synchronizes CUDA, so leave it off for timed benchmark runs. _pt = None if os.environ.get("DZ_PHASE_TIMING"): import time as _time torch.accelerator.synchronize() _pt = {"time": _time, "t0": _time.perf_counter(), "marks": []} def _pt_mark(name: str) -> None: if _pt is not None: torch.accelerator.synchronize() _pt["marks"].append((name, _pt["time"].perf_counter())) # Prompt embeds are constant within a session (a prompt change triggers a # "session" reset above, which clears this cache alongside state.language). if state.prompt_embeds is None: state.prompt_embeds = self._encode_text(text_tokens, attention_mask) prompt_embeds = state.prompt_embeds # Negative prompt for the unconditional CFG branch (model constant) negative_prompt_embeds = None if self.cfg_scale > 1.0: if self._negative_prompt_embeds_cache is None: neg_inputs = self.tokenizer( self.negative_prompt, max_length=512, padding="max_length", truncation=True, return_tensors="pt", add_special_tokens=True, ) self._negative_prompt_embeds_cache = self._encode_text( neg_inputs["input_ids"].to(device), neg_inputs["attention_mask"].to(device), ) negative_prompt_embeds = self._negative_prompt_embeds_cache _pt_mark("text_encode") # Extract first/last frame for CLIP + VAE encoding if num_frames_raw == 4 or num_frames_raw == 9: image = videos[:, :, -1:].transpose(1, 2) else: image = videos[:, :, :1].transpose(1, 2) if state.current_start_frame == 0: clip_feas, ys, image = self._encode_image( image, self.num_frames, height, width, state=state, ) state.clip_feas = clip_feas.to(dtype=image.dtype) state.ys = ys.to(dtype=image.dtype) # Eager cross-attn population (AR-Diffusion only): cache text + image-token K/V # into the pool now that the image is encoded (clip_feas available). # Runs on the first forward of a session and after each window-boundary # reset (current_start_frame returns to 0); the populate guards (text/img halves) gate re-entry. self._kv_populate_cross(prompt_embeds, state.clip_feas, is_negative=False) if negative_prompt_embeds is not None: self._kv_populate_cross(negative_prompt_embeds, state.clip_feas, is_negative=True) if state.current_start_frame != 0: latent_dtype = videos.dtype with torch.no_grad(): image = self._encode_observation_latents(state, videos, latent_dtype=latent_dtype) _pt_mark("obs_vae_encode") batch_size = image.shape[0] generator = torch.Generator(device=device).manual_seed(self.seed) noise_obs = torch.randn( batch_size, 16, self.num_frame_per_block, height // 8, width // 8, device=device, dtype=torch.bfloat16, generator=generator, ) generator = torch.Generator(device=device).manual_seed(self.seed) noise_action = torch.randn( batch_size, self.action_horizon, self.transformer.action_dim, device=device, dtype=torch.bfloat16, generator=generator, ) _, num_channels, num_frames, h_latent, w_latent = noise_obs.shape frame_seqlen = int(h_latent * w_latent / 4) seq_len = frame_seqlen * num_frames image = image.transpose(1, 2) noise_obs = noise_obs.transpose(1, 2) do_true_cfg = self.cfg_scale > 1.0 and negative_prompt_embeds is not None self._prefill_kv_cache( image, prompt_embeds, negative_prompt_embeds, frame_seqlen, seq_len, do_true_cfg, state, ) _pt_mark("prefill_kv") sample_scheduler = copy.deepcopy(self.scheduler) sample_scheduler_action = copy.deepcopy(self.scheduler) sample_scheduler.set_timesteps( self.num_inference_steps, device=device, shift=self.sigma_shift, ) sample_scheduler_action.set_timesteps( self.num_inference_steps, device=device, shift=self.sigma_shift, ) if self.decouple_inference_noise: video_final_noise = self.video_inference_final_noise sigma_max = sample_scheduler.sigmas[0].item() sample_scheduler.sigmas = ( sample_scheduler.sigmas * (sigma_max - video_final_noise) / sigma_max + video_final_noise ) sample_scheduler.timesteps = (sample_scheduler.sigmas[:-1] * 1000).to(torch.int64) video_action_scheduler = VideoActionScheduler( sample_scheduler, sample_scheduler_action, ) video_out, action_out = self.diffuse( video_latents=noise_obs, action_latents=noise_action, timesteps_video=sample_scheduler.timesteps, timesteps_action=sample_scheduler_action.timesteps, prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, video_action_scheduler=video_action_scheduler, do_true_cfg=do_true_cfg, state=state, seq_len=seq_len, state_features=state_features, embodiment_id=embodiment_id, ) if _pt is not None: _pt_mark("diffuse") prev_t = _pt["t0"] parts = [] for name, t in _pt["marks"]: parts.append(f"{name}={1000 * (t - prev_t):.1f}ms") prev_t = t logger.info( "DZ_PHASE_TIMING csf=%s total=%.1fms %s", state.current_start_frame, 1000 * (prev_t - _pt["t0"]), " ".join(parts), ) if state.current_start_frame == 1: 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) # Relative -> absolute: only for relative_action_keys (joint_position only) # gripper_position is NOT relative, so don't add state back to it if self.relative_action and state_for_postprocess is not None: n_relative = self.relative_action_dim # 7 for DROID (joint only) # Use original state precision for post-denorm absolute recovery. # Upstream adds obs state after `eval_transform.unapply()` # the bf16 denoising path. last_state = state_for_postprocess[:, 0, :n_relative] # (B, n_relative) action_out[..., :n_relative] = ( action_out[..., :n_relative] + last_state.unsqueeze(1) # broadcast over horizon ) # Squeeze batch dim for output: (B, horizon, dim) -> (horizon, dim) actions_np = action_out.squeeze(0).float().cpu().numpy() # (horizon, max_action_dim) actions_np = transform.transform_action_output(actions_np) return DiffusionOutput( output={ "actions": actions_np, # Source `video_pred` is normalized VAE latent output, not RGB. # Use `decode_video_latents()` for DreamZero-equivalent debug # video decoding. "video": video_out.transpose(1, 2).cpu(), }, ) # ----------------------------------------------------------------------- # Action denormalization # ----------------------------------------------------------------------- def _load_action_norm_stats(self, stats_path: str) -> dict[str, dict[str, torch.Tensor]]: """Load per-embodiment action normalization stats from metadata.json. Returns: {embodiment_name: {"q01": Tensor(action_dim,), "q99": Tensor(action_dim,)}} """ with open(stats_path) as f: metadata = json.load(f) return self._parse_action_norm_stats(metadata) @staticmethod def _parse_action_norm_stats(metadata: dict) -> dict[str, dict[str, torch.Tensor]]: result = {} for emb_name, emb_data in metadata.items(): action_stats = emb_data.get("statistics", {}).get("action", {}) q01_parts, q99_parts = [], [] # Concatenate joint_position + gripper_position stats for key in ["joint_position", "gripper_position"]: if key in action_stats: q01_parts.extend(action_stats[key]["q01"]) q99_parts.extend(action_stats[key]["q99"]) if q01_parts: result[emb_name] = { "q01": torch.tensor(q01_parts, dtype=torch.float32), "q99": torch.tensor(q99_parts, dtype=torch.float32), } return result @staticmethod def _parse_state_norm_stats(metadata: dict) -> dict[str, dict[str, torch.Tensor]]: """Load per-embodiment state normalization stats from metadata.json.""" result = {} for emb_name, emb_data in metadata.items(): state_stats = emb_data.get("statistics", {}).get("state", {}) q01_parts, q99_parts = [], [] for key in ["joint_position", "gripper_position"]: if key in state_stats: q01_parts.extend(state_stats[key]["q01"]) q99_parts.extend(state_stats[key]["q99"]) if q01_parts: result[emb_name] = { "q01": torch.tensor(q01_parts, dtype=torch.float32), "q99": torch.tensor(q99_parts, dtype=torch.float32), } return result def _normalize_state( self, state: torch.Tensor, embodiment_name: str, ) -> torch.Tensor: """Normalize state with q99 stats before feeding the model.""" state_norm_stats = getattr(self, "state_norm_stats", {}) if embodiment_name not in state_norm_stats: return state stats = state_norm_stats[embodiment_name] q01 = stats["q01"].to(device=state.device, dtype=state.dtype) q99 = stats["q99"].to(device=state.device, dtype=state.dtype) actual_dim = q01.shape[0] normalized = state.clone() range_vals = q99 - q01 mask = range_vals != 0 normalized_slice = normalized[..., :actual_dim] normalized_slice[..., mask] = 2 * (normalized_slice[..., mask] - q01[mask]) / range_vals[mask] - 1 normalized_slice = torch.clamp(normalized_slice, -1, 1) normalized[..., :actual_dim] = normalized_slice return normalized def _denormalize_action( self, action: torch.Tensor, embodiment_name: str, ) -> torch.Tensor: """Denormalize action from [-1,1] to real values using q99 mode. Formula: real = (normalized + 1) / 2 * (q99 - q01) + q01 """ if embodiment_name not in self.action_norm_stats: return action stats = self.action_norm_stats[embodiment_name] q01 = stats["q01"].to(device=action.device, dtype=action.dtype) q99 = stats["q99"].to(device=action.device, dtype=action.dtype) # action shape: (B, horizon, action_dim) or (B, horizon, max_action_dim) # q01/q99 shape: (actual_action_dim,) -- only denorm actual dims actual_dim = q01.shape[0] action_real = action.clone() action_real[..., :actual_dim] = (action[..., :actual_dim] + 1) / 2 * (q99 - q01) + q01 return action_real # ----------------------------------------------------------------------- # Weight loading # ----------------------------------------------------------------------- @property def weights_sources(self): """ComponentSource list for DiffusersPipelineLoader.""" return self._weights_sources def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Load checkpoint weights with key remapping.""" loaded: set[str] = set() params = dict(self.named_parameters()) buffers = dict(self.named_buffers()) for name, tensor in weights: if name.startswith("action_head.model."): new_name = "transformer." + name[len("action_head.model.") :] new_name = ( new_name.replace("img_emb.proj.0.", "img_emb.norm1.") .replace("img_emb.proj.1.", "img_emb.fc1.") .replace("img_emb.proj.3.", "img_emb.fc2.") .replace("img_emb.proj.4.", "img_emb.norm2.") ) # Self-attn q/k/v are fused into a single QKVParallelLinear; route # each separate checkpoint weight/bias to the packed `qkv` param # with its shard id. cross_attn keeps separate q/k/v (q from x, # k/v from context — not fusible), so it is left untouched here. qkv_shard_id: str | None = None for shard_id in ("q", "k", "v"): needle = f".self_attn.{shard_id}." if needle in new_name: new_name = new_name.replace(needle, ".self_attn.qkv.") qkv_shard_id = shard_id break if new_name in params: param = params[new_name] if qkv_shard_id is not None: # QKVParallelLinear.weight_loader needs the shard id. param.weight_loader(param, tensor, qkv_shard_id) else: weight_loader = getattr(param, "weight_loader", default_weight_loader) weight_loader(param, tensor) loaded.add(new_name) elif new_name in buffers: buffers[new_name].data.copy_(tensor) loaded.add(new_name) elif name.startswith("action_head.text_encoder."): mapped = self._remap_text_encoder_key(name) if mapped is None: continue for new_name in mapped if isinstance(mapped, list) else [mapped]: full_name = "text_encoder." + new_name if full_name in params: params[full_name].data.copy_(tensor) loaded.add(full_name) elif name.startswith("action_head.image_encoder."): self._remap_image_encoder_key(name, tensor, params, loaded) elif name.startswith("action_head.vae."): mapped = self._remap_vae_key(name) if mapped is None: continue full_name = "vae." + mapped if full_name in params: params[full_name].data.copy_(tensor) loaded.add(full_name) logger.info( "DreamZero load_weights: loaded %d parameters from root checkpoint", len(loaded), ) return loaded # ----------------------------------------------------------------------- # Text encoder key remapping # ----------------------------------------------------------------------- @staticmethod def _remap_text_encoder_key(name: str) -> str | list[str] | None: """Remap a single text encoder key.""" subkey = name[len("action_head.text_encoder.") :] if subkey == "token_embedding.weight": return "shared.weight" if subkey == "norm.weight": return "encoder.final_layer_norm.weight" m = re_module.match(r"blocks\.(\d+)\.(.*)", subkey) if not m: return None block_idx = m.group(1) rest = m.group(2) prefix = f"encoder.block.{block_idx}" if rest == "attn.q.weight": return f"{prefix}.layer.0.SelfAttention.q.weight" if rest == "attn.k.weight": return f"{prefix}.layer.0.SelfAttention.k.weight" if rest == "attn.v.weight": return f"{prefix}.layer.0.SelfAttention.v.weight" if rest == "attn.o.weight": return f"{prefix}.layer.0.SelfAttention.o.weight" if rest == "pos_embedding.embedding.weight": return f"{prefix}.layer.0.SelfAttention.relative_attention_bias.weight" if rest == "norm1.weight": return f"{prefix}.layer.0.layer_norm.weight" if rest == "ffn.gate.0.weight": return f"{prefix}.layer.1.DenseReluDense.wi_0.weight" if rest == "ffn.fc1.weight": return f"{prefix}.layer.1.DenseReluDense.wi_1.weight" if rest == "ffn.fc2.weight": return f"{prefix}.layer.1.DenseReluDense.wo.weight" if rest == "norm2.weight": return f"{prefix}.layer.1.layer_norm.weight" return None # ----------------------------------------------------------------------- # VAE key remapping # ----------------------------------------------------------------------- @staticmethod def _remap_vae_key(name: str) -> str | None: """Remap DreamZero VAE keys to `DistributedAutoencoderKLWan` keys.""" if not name.startswith("action_head.vae.model."): return None rest = name[len("action_head.vae.model.") :] direct_prefix_map = { "encoder.conv1.": "encoder.conv_in.", "encoder.head.0.": "encoder.norm_out.", "encoder.head.2.": "encoder.conv_out.", "decoder.conv1.": "decoder.conv_in.", "decoder.head.0.": "decoder.norm_out.", "decoder.head.2.": "decoder.conv_out.", "conv1.": "quant_conv.", "conv2.": "post_quant_conv.", } for src_prefix, dst_prefix in direct_prefix_map.items(): if rest.startswith(src_prefix): return dst_prefix + rest[len(src_prefix) :] resnet_leaf_map = { "residual.0.gamma": "norm1.gamma", "residual.2.weight": "conv1.weight", "residual.2.bias": "conv1.bias", "residual.3.gamma": "norm2.gamma", "residual.6.weight": "conv2.weight", "residual.6.bias": "conv2.bias", } block_leaf_map = { **resnet_leaf_map, "shortcut.weight": "conv_shortcut.weight", "shortcut.bias": "conv_shortcut.bias", "resample.1.weight": "resample.1.weight", "resample.1.bias": "resample.1.bias", "time_conv.weight": "time_conv.weight", "time_conv.bias": "time_conv.bias", } m = re_module.match(r"encoder\.middle\.(\d+)\.(.*)", rest) if m: idx = int(m.group(1)) tail = m.group(2) if idx in (0, 2) and tail in resnet_leaf_map: res_idx = 0 if idx == 0 else 1 return f"encoder.mid_block.resnets.{res_idx}.{resnet_leaf_map[tail]}" if idx == 1: return f"encoder.mid_block.attentions.0.{tail}" return None m = re_module.match(r"decoder\.middle\.(\d+)\.(.*)", rest) if m: idx = int(m.group(1)) tail = m.group(2) if idx in (0, 2) and tail in resnet_leaf_map: res_idx = 0 if idx == 0 else 1 return f"decoder.mid_block.resnets.{res_idx}.{resnet_leaf_map[tail]}" if idx == 1: return f"decoder.mid_block.attentions.0.{tail}" return None m = re_module.match(r"encoder\.downsamples\.(\d+)\.(.*)", rest) if m: idx = int(m.group(1)) tail = m.group(2) if tail in block_leaf_map: return f"encoder.down_blocks.{idx}.{block_leaf_map[tail]}" return None m = re_module.match(r"decoder\.upsamples\.(\d+)\.(.*)", rest) if m: idx = int(m.group(1)) tail = m.group(2) if tail not in block_leaf_map: return None if idx <= 2: prefix = f"decoder.up_blocks.0.resnets.{idx}." elif idx == 3: prefix = "decoder.up_blocks.0.upsamplers.0." elif 4 <= idx <= 6: prefix = f"decoder.up_blocks.1.resnets.{idx - 4}." elif idx == 7: prefix = "decoder.up_blocks.1.upsamplers.0." elif 8 <= idx <= 10: prefix = f"decoder.up_blocks.2.resnets.{idx - 8}." elif idx == 11: prefix = "decoder.up_blocks.2.upsamplers.0." elif 12 <= idx <= 14: prefix = f"decoder.up_blocks.3.resnets.{idx - 12}." else: return None return prefix + block_leaf_map[tail] return None # ----------------------------------------------------------------------- # Image encoder key remapping # ----------------------------------------------------------------------- def _remap_image_encoder_key( self, name: str, tensor: torch.Tensor, params: dict[str, torch.nn.Parameter], loaded: set[str], ) -> None: """Map an image encoder key onto the local module.""" if not name.startswith("action_head.image_encoder."): return full_name = "image_encoder." + name[len("action_head.image_encoder.") :] if full_name in params: params[full_name].data.copy_(tensor) loaded.add(full_name)