1497 lines
61 KiB
Python
1497 lines
61 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
"""DreamZero pipeline for vllm-omni.
|
|
|
|
Entry point for DiffusionEngine.step() → pipeline.forward(req)
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import copy
|
|
import json
|
|
import logging
|
|
import os
|
|
import re as re_module
|
|
from collections import OrderedDict
|
|
from collections.abc import Iterable
|
|
|
|
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.platforms import current_omni_platform
|
|
|
|
logger = logging.getLogger(__name__)
|
|
MAX_DREAMZERO_SESSIONS = 64
|
|
|
|
|
|
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 positive branch only.
|
|
State: DreamZeroState manages KV cache + frame buffer across forward() calls.
|
|
"""
|
|
|
|
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__()
|
|
|
|
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,
|
|
)
|
|
|
|
self._states: OrderedDict[str, DreamZeroState] = OrderedDict()
|
|
self._max_session_states = MAX_DREAMZERO_SESSIONS
|
|
self.state = self._get_or_create_state("default")
|
|
|
|
# DiT step cache is configured by StepCacheBackend
|
|
# (cache_backend="step_cache") via pipeline._stepcache_config.
|
|
|
|
# Keep runtime inference settings separate from the training-time config.
|
|
self.num_inference_steps: int = model_config.get(
|
|
"num_inference_steps",
|
|
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.num_frame_per_block: int = ah_config["num_frame_per_block"]
|
|
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)
|
|
|
|
# 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:
|
|
session_key = str(session_id or "default")
|
|
state = self._states.get(session_key)
|
|
if state is None:
|
|
state = DreamZeroState()
|
|
self._states[session_key] = state
|
|
max_states = getattr(self, "_max_session_states", MAX_DREAMZERO_SESSIONS)
|
|
while len(self._states) > max_states:
|
|
self._states.popitem(last=False)
|
|
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) and updated_kv_caches:
|
|
state = kwargs.get("dreamzero_state", self.state)
|
|
is_neg = kwargs.get("is_negative", False)
|
|
for i, kv in enumerate(updated_kv_caches):
|
|
state.update_kv_cache(i, kv, is_negative=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).
|
|
DiT blocks use per-block ``fullgraph=True``.
|
|
"""
|
|
if not torch.cuda.is_available():
|
|
logger.info("DreamZero setup_compile skipped: CUDA not available.")
|
|
return
|
|
|
|
from vllm_omni.diffusion.models.dreamzero.wan_vae_feat_cache_patch import (
|
|
apply_wan_vae_feat_cache_tensor_patch,
|
|
)
|
|
|
|
apply_wan_vae_feat_cache_tensor_patch()
|
|
|
|
compile_ro = {"mode": "reduce-overhead", "fullgraph": True, "dynamic": False}
|
|
# DiT blocks: default avoids CUDAGraph overwrite on modulation tensors; encoders use reduce-overhead.
|
|
dit_compile = {"mode": "default", "fullgraph": True, "dynamic": False}
|
|
|
|
logger.info(
|
|
"DreamZero: torch.compile text/image/VAE encode + per-block DiT (encoders reduce-overhead, DiT default)."
|
|
)
|
|
|
|
try:
|
|
self.text_encoder.forward = torch.compile(self.text_encoder.forward, **compile_ro)
|
|
except Exception as exc:
|
|
logger.warning("DreamZero: text_encoder compile failed (%s); skipping.", exc)
|
|
|
|
try:
|
|
self.image_encoder.model.visual.forward = torch.compile(
|
|
self.image_encoder.model.visual.forward,
|
|
**compile_ro,
|
|
)
|
|
except Exception as exc:
|
|
logger.warning("DreamZero: image_encoder compile failed (%s); skipping.", exc)
|
|
|
|
try:
|
|
self.vae.encode = torch.compile(self.vae.encode, **compile_ro)
|
|
except Exception as exc:
|
|
logger.warning("DreamZero: vae.encode compile failed (%s); skipping.", exc)
|
|
|
|
compiled_blocks = 0
|
|
for block in self.transformer.blocks:
|
|
try:
|
|
block.forward = torch.compile(block.forward, **dit_compile)
|
|
compiled_blocks += 1
|
|
except Exception as exc:
|
|
logger.warning(
|
|
"DreamZero: transformer block %d compile failed (%s); leaving remaining eager.",
|
|
compiled_blocks,
|
|
exc,
|
|
)
|
|
break
|
|
if compiled_blocks:
|
|
logger.info(
|
|
"DreamZero: compiled %d/%d transformer blocks.",
|
|
compiled_blocks,
|
|
len(self.transformer.blocks),
|
|
)
|
|
|
|
self.warmup_compile()
|
|
|
|
def warmup_compile(self) -> None:
|
|
"""Warm up compiled text/image/VAE paths before timed inference."""
|
|
if not torch.cuda.is_available():
|
|
return
|
|
|
|
device = next(self.text_encoder.parameters()).device
|
|
with torch.inference_mode():
|
|
try:
|
|
text_tokens = torch.zeros(1, 16, dtype=torch.long, device=device)
|
|
attention_mask = torch.ones_like(text_tokens)
|
|
self._encode_text(text_tokens, attention_mask)
|
|
except Exception as exc:
|
|
logger.warning("DreamZero compile warmup (text_encoder) skipped: %s", exc)
|
|
|
|
try:
|
|
image = torch.zeros(1, 1, 3, 180, 320, dtype=torch.bfloat16, device=device)
|
|
self._encode_image(image, self.num_frames, 180, 320)
|
|
except Exception as exc:
|
|
logger.warning("DreamZero compile warmup (image_encoder) skipped: %s", exc)
|
|
|
|
try:
|
|
latent_h, latent_w = 180 // 8, 320 // 8
|
|
dummy_latent = torch.zeros(
|
|
1,
|
|
16,
|
|
self.num_frame_per_block,
|
|
latent_h * 2,
|
|
latent_w * 2,
|
|
dtype=self.vae.dtype,
|
|
device=device,
|
|
)
|
|
mean = self.vae_latents_mean.to(device=device, dtype=self.vae.dtype)
|
|
inv_std = self.vae_latents_inv_std.to(device=device, dtype=self.vae.dtype)
|
|
dummy_latent_denorm = dummy_latent / inv_std + mean
|
|
self._cudagraph_mark_step_begin()
|
|
self.vae.decode(dummy_latent_denorm, return_dict=False)
|
|
except Exception as exc:
|
|
logger.warning("DreamZero compile warmup (vae decode) skipped: %s", exc)
|
|
|
|
torch.accelerator.synchronize(device)
|
|
logger.info("DreamZero compile warmup finished (text / image / vae decode).")
|
|
|
|
def combine_cfg_noise(
|
|
self,
|
|
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() + platform synchronize."""
|
|
latents = tuple(t.contiguous() for t in latents)
|
|
if do_true_cfg and get_classifier_free_guidance_world_size() > 1:
|
|
current_omni_platform.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,
|
|
) -> 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)
|
|
|
|
return clip_context, y, new_image
|
|
|
|
def _encode_vae_latents(self, videos: torch.Tensor) -> torch.Tensor:
|
|
"""Encode videos into normalized VAE latents."""
|
|
input_dtype = videos.dtype
|
|
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: DreamZeroState,
|
|
) -> 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:
|
|
state.create_kv_caches(
|
|
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=state.get_kv_caches(False),
|
|
crossattn_cache=state.get_crossattn_caches(False),
|
|
is_negative=False,
|
|
**common,
|
|
)
|
|
negative_kwargs = (
|
|
dict(
|
|
encoder_hidden_states=negative_prompt_embeds,
|
|
kv_cache=state.get_kv_caches(True),
|
|
crossattn_cache=state.get_crossattn_caches(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=state.get_kv_caches(False),
|
|
crossattn_cache=state.get_crossattn_caches(False),
|
|
is_negative=False,
|
|
**common,
|
|
)
|
|
negative_kwargs = (
|
|
dict(
|
|
encoder_hidden_states=negative_prompt_embeds,
|
|
kv_cache=state.get_kv_caches(True),
|
|
crossattn_cache=state.get_crossattn_caches(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: DreamZeroState,
|
|
**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 cond & uncond 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
|
|
if csf + self.num_frame_per_block <= state.ys.shape[2]:
|
|
y = state.ys[:, :, csf : csf + self.num_frame_per_block]
|
|
else:
|
|
y = state.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=state.get_kv_caches(False),
|
|
crossattn_cache=state.get_crossattn_caches(False),
|
|
y=y,
|
|
clip_feature=state.clip_feas,
|
|
action=noisy_input_action,
|
|
timestep_action=timestep_action,
|
|
is_negative=False,
|
|
**common_kwargs,
|
|
)
|
|
|
|
if do_true_cfg and negative_prompt_embeds is not None:
|
|
negative_kwargs = dict(
|
|
hidden_states=noisy_input.transpose(1, 2),
|
|
timestep_video=timestep,
|
|
encoder_hidden_states=negative_prompt_embeds,
|
|
kv_cache=state.get_kv_caches(True),
|
|
crossattn_cache=state.get_crossattn_caches(True),
|
|
y=y,
|
|
clip_feature=state.clip_feas,
|
|
action=noisy_input_action,
|
|
timestep_action=timestep_action,
|
|
is_negative=True,
|
|
**common_kwargs,
|
|
)
|
|
else:
|
|
negative_kwargs = None
|
|
|
|
noise_pred = self.predict_noise_maybe_with_cfg(
|
|
positive_kwargs=positive_kwargs,
|
|
negative_kwargs=negative_kwargs,
|
|
do_true_cfg=do_true_cfg,
|
|
true_cfg_scale=self.cfg_scale,
|
|
cfg_normalize=False,
|
|
)
|
|
flow_pred, flow_pred_action = noise_pred
|
|
_cached_flow_pred = flow_pred
|
|
_cached_flow_pred_action = flow_pred_action
|
|
|
|
_prev_predictions.append((flow_pred,))
|
|
if _step_cache is not None:
|
|
_step_cache.trim_history(_prev_predictions)
|
|
else:
|
|
# Reuse previous prediction (DiT step-skipping cache).
|
|
assert _cached_flow_pred is not None, "DiT cache: no cached prediction available for step skip."
|
|
flow_pred = _cached_flow_pred
|
|
flow_pred_action = _cached_flow_pred_action
|
|
|
|
latents = (noisy_input, noisy_input_action)
|
|
t = (video_timestep, action_timestep)
|
|
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: OmniDiffusionRequest, **kwargs) -> DiffusionOutput:
|
|
"""Full inference step. Called by DiffusionEngine.step()."""
|
|
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 reset from OpenPI serving is carried by `extra_args["reset"]`
|
|
# on the next inference request after websocket reset/session switch.
|
|
if extra_args.get("reset", False):
|
|
state.reset()
|
|
else:
|
|
reset_reason = state.reset_reason(text_tokens, 0, self.transformer.local_attn_size)
|
|
if reset_reason == "session":
|
|
state.reset()
|
|
elif reset_reason == "inference":
|
|
state.reset_inference_state()
|
|
state.language = text_tokens
|
|
|
|
# Frame accumulation: stitched single frame → multi-frame video
|
|
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
|
|
|
|
prompt_embeds = self._encode_text(text_tokens, attention_mask)
|
|
# Negative prompt for CFG uncond branch (model constant)
|
|
negative_prompt_embeds = None
|
|
if self.cfg_scale > 1.0:
|
|
neg_inputs = self.tokenizer(
|
|
self.negative_prompt,
|
|
max_length=512,
|
|
padding="max_length",
|
|
truncation=True,
|
|
return_tensors="pt",
|
|
add_special_tokens=True,
|
|
)
|
|
negative_prompt_embeds = self._encode_text(
|
|
neg_inputs["input_ids"].to(device),
|
|
neg_inputs["attention_mask"].to(device),
|
|
)
|
|
|
|
# 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.clip_feas = clip_feas.to(dtype=image.dtype)
|
|
state.ys = ys.to(dtype=image.dtype)
|
|
|
|
if state.current_start_frame != 0:
|
|
# Subsequent calls: encode current observation via VAE
|
|
if (num_frames_raw - 1) // 4 == self.num_frame_per_block:
|
|
pass
|
|
elif 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]
|
|
videos = torch.cat([first_frame, videos], dim=2)
|
|
else:
|
|
first_frame = videos[:, :, 0:1]
|
|
videos = torch.cat([first_frame, videos], dim=2)
|
|
|
|
latent_dtype = videos.dtype
|
|
with torch.no_grad():
|
|
image = self._encode_vae_latents(videos)
|
|
image = image.to(dtype=latent_dtype)
|
|
|
|
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,
|
|
)
|
|
|
|
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 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)
|