# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM 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 import numpy as np import torch import torch.nn as nn 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, ) 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 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()``. """ _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_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``. """ 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. # 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.cuda.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() + cuda.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.is_cuda), None) if device is not None: torch.cuda.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)