vllm-omni/vllm_omni/diffusion/models/dreamzero/pipeline_dreamzero.py

# 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)