diff --git a/examples/dreamzero/README.md b/examples/dreamzero/README.md new file mode 100644 index 0000000000..59e43685a9 --- /dev/null +++ b/examples/dreamzero/README.md @@ -0,0 +1,162 @@ +# DreamZero: Training + +This guide covers how to train DreamZero 14B using FlagScale with FSDP2 distributed training. + +## Overview + +DreamZero is a World Action Model — a Wan2.1 14B video DiT repurposed as a zero-shot robot policy. Architecture: +- **T5-XXL** text encoder (~4.7B, frozen) +- **CLIP** image encoder (~0.6B, frozen) +- **VAE** video encoder (~0.1B, frozen) +- **CausalWan DiT** video diffusion transformer (~14B, trainable) +- Action/state encoder-decoder projectors (trainable) + +FlagScale uses FSDP2 (ZeRO-3 style sharding) to distribute the 14B trainable DiT across GPUs. + +## Installation + +### Clone Repository + +```sh +git clone https://github.com/FlagOpen/FlagScale.git +cd FlagScale/ +``` + +### Setup Environment + +Create a conda environment with PyTorch 2.4+ and FSDP2 support: + +```sh +conda create -n flagscale-dreamzero python=3.11 +conda activate flagscale-dreamzero +pip install ".[cuda-train]" --verbose +``` + +Install additional dependencies: + +```sh +pip install safetensors peft decord +``` + +DreamZero also requires the reference `groot` package for model architecture and data transforms: + +```sh +git clone /workspace/dreamzero +``` + +Add both FlagScale and the reference repo to your PYTHONPATH: + +```sh +export PYTHONPATH=/path/to/FlagScale:/workspace/dreamzero:$PYTHONPATH +``` + +## Download Models + +Download the DreamZero-AgiBot checkpoint (contains `config.json` + safetensors shards): + +```sh +# Place at /workspace/models/DreamZero-AgiBot/ +# Expected contents: +# config.json +# model.safetensors.index.json +# model-00001-of-00010.safetensors +# ... +# model-00010-of-00010.safetensors +``` + +Download the Wan2.1-I2V-14B-480P pretrained weights (for tokenizer, and optionally T5/CLIP/VAE): + +```sh +# Place at /workspace/models/Wan2.1-I2V-14B-480P/ +# Required for tokenizer: +# google/umt5-xxl/tokenizer.json +# google/umt5-xxl/spiece.model +``` + +## Training + +### Prepare Dataset + +DreamZero uses a LeRobot-format dataset with `data/`, `meta/`, and `videos/` subdirectories. + +```sh +# Place at /workspace/datasets/vla_arena_dreamzero/ +# Expected structure: +# data/chunk-000/, data/chunk-001/, ... +# meta/info.json, meta/episodes.jsonl, ... +# videos/chunk-000/, ... +``` + +### Edit Config + +FlagScale uses a two-level configuration system: + +1. **Experiment-level config** (`examples/dreamzero/conf/train.yaml`): Experiment settings, environment variables, and resource allocation +2. **Task-level config** (`examples/dreamzero/conf/train/dreamzero_14b.yaml`): Model, dataset, and training hyperparameters + +#### Experiment-Level Config + +```sh +vim examples/dreamzero/conf/train.yaml +``` + +Configure the following fields: + +- `experiment.exp_name` — Experiment name +- `experiment.exp_dir` — Output directory for checkpoints and logs +- `experiment.envs.CUDA_VISIBLE_DEVICES` — GPU devices to use (e.g., `"0,1,2,3,4,5,6,7"`) +- `experiment.runner.nproc_per_node` — Number of GPUs + +#### Task-Level Config + +```sh +vim examples/dreamzero/conf/train/dreamzero_14b.yaml +``` + +Configure the following fields: + +**System settings:** +- `system.batch_size` — Per-GPU micro batch size (default: `1`) +- `system.train_steps` — Total training steps (default: `5000`) +- `system.checkpoint.save_freq` — Steps between checkpoints + +**Model settings:** +- `model.pretrained_model_path` — Path to DreamZero-AgiBot checkpoint (e.g., `/workspace/models/DreamZero-AgiBot`) +- `model.train_architecture` — `"full"` for full fine-tuning (recommended), `"lora"` for LoRA on DiT +- `model.optimizer.lr` — Learning rate (default: `1.0e-5`) +- `model.optimizer.betas` — Adam betas (default: `[0.95, 0.999]`) +- `model.optimizer.scheduler.warmup_ratio` — Warmup ratio (default: `0.05`) + +**Data settings:** +- `data.data_path` — Path to LeRobot dataset (e.g., `/workspace/datasets/vla_arena_dreamzero`) +- `data.tokenizer_path` — Path to UMT5-XXL tokenizer (e.g., `/workspace/models/Wan2.1-I2V-14B-480P/google/umt5-xxl`) +- `data.embodiment_tag` — Embodiment tag (e.g., `"libero"`) + +### Start Training + +```sh +cd FlagScale/ +python flagscale/run.py --config-path examples/dreamzero/conf --config-name train +``` + +Training logs are saved to `outputs/dreamzero_train/logs/host_0_localhost.output` by default. + +### Stop Training + +```sh +cd FlagScale/ +python flagscale/run.py --config-path examples/dreamzero/conf --config-name train action=stop +``` + +## Hardware Requirements + +With 8x H100/H800 80GB GPUs and `batch_size=1`: +- Full fine-tuning (FSDP2): ~65 GB per GPU +- Model loading takes ~30 minutes (14B params from 10 safetensors shards) +- Training throughput: ~8.7s per step (effective batch size 8) + +## Known Issues + +1. **torch.compile + FSDP2**: Compiling attention sub-methods with `mode="reduce-overhead"` causes CUDA graph tensor deallocation mismatches during backward. Disabled by default via `system.disable_attention_compile: true`. +2. **LoRA mode NaN**: `train_architecture: lora` can produce NaN due to bf16 backward overflow at certain timestep/data combinations. Use `train_architecture: full` to avoid this. +3. **batch_size=2 OOM**: Full fine-tuning with `batch_size=2` exceeds 80GB GPU memory. Use `batch_size=1` with gradient accumulation if larger effective batches are needed. diff --git a/examples/dreamzero/conf/train.yaml b/examples/dreamzero/conf/train.yaml new file mode 100644 index 0000000000..b3c7825c9c --- /dev/null +++ b/examples/dreamzero/conf/train.yaml @@ -0,0 +1,39 @@ +defaults: + - _self_ + - train: dreamzero_14b + +experiment: + exp_name: dreamzero_train + seed: 42 + save_steps: 500 + load: null + exp_dir: outputs/${experiment.exp_name} + ckpt_format: torch + task: + type: train + backend: native + entrypoint: flagscale/train/train_dreamzero.py + runner: + per_node_task: false + no_shared_fs: false + rdzv_backend: static + hostfile: null + nproc_per_node: 8 + cmds: + before_start: echo "Starting DreamZero Training" + envs: + LOGLEVEL: "INFO" + CUDA_VISIBLE_DEVICES: "0,1,2,3,4,5,6,7" + CUDA_DEVICE_MAX_CONNECTIONS: 1 + WANDB_MODE: offline + OTEL_SDK_DISABLED: true + TOKENIZERS_PARALLELISM: false + NCCL_TIMEOUT: "1800" + TORCH_NCCL_ASYNC_ERROR_HANDLING: "1" + PYTHONUNBUFFERED: "1" + +action: run + +hydra: + run: + dir: ${experiment.exp_dir}/hydra diff --git a/examples/dreamzero/conf/train/dreamzero_14b.yaml b/examples/dreamzero/conf/train/dreamzero_14b.yaml new file mode 100644 index 0000000000..011df39d5a --- /dev/null +++ b/examples/dreamzero/conf/train/dreamzero_14b.yaml @@ -0,0 +1,92 @@ +# DreamZero 14B Training Config for FlagScale +# LIBERO fine-tuning with full DIT fine-tuning via FSDP2 on 8xH100/H800 + +system: + train_steps: 5000 + log_freq: 10 + grad_clip_norm: 1.0 + use_amp: true + shuffle: true + num_workers: 4 + batch_size: 1 + # Disable torch.compile on attention sub-methods — CUDA graph replay + # crashes with FSDP2 reshard-after-forward (tensor deallocation mismatch). + disable_attention_compile: true + + checkpoint: + save_freq: 500 + output_directory: ${experiment.exp_dir} + save_total_limit: 10 + +model: + model_name: dreamzero + checkpoint_dir: null + # Path to DreamZero-AgiBot checkpoint (contains config.json + safetensors shards) + pretrained_model_path: /workspace/models/DreamZero-AgiBot + # Optional: explicit paths for T5/CLIP/VAE (null = loaded from checkpoint) + text_encoder_pretrained_path: null + image_encoder_pretrained_path: null + vae_pretrained_path: null + compute_dtype: bfloat16 + use_gradient_checkpointing: true + + # Training mode: "full" (full fine-tuning) or "lora" (LoRA on DIT) + # "full" recommended — avoids LoRA-on-fine-tuned-weights NaN issue. + train_architecture: full + # LoRA config (only used if train_architecture=lora) + lora_rank: 4 + lora_alpha: 4 + lora_target_modules: "q,k,v,o,ffn.0,ffn.2" + + # Dimensions matching LIBERO data + num_frames: 33 + action_horizon: 24 + state_horizon: 1 + action_dim: 32 + max_state_dim: 64 + + # Frame/block settings + num_frame_per_block: 2 + num_action_per_block: 24 + num_state_per_block: 1 + # frame_seqlen: 2 views tiled to 352x640 → VAE 8x → 44x80 → patch(1,2,2) → 22x40 = 880 + frame_seqlen: 880 + num_views: 2 + + optimizer: + name: AdamW + lr: 1.0e-5 + weight_decay: 1.0e-5 + betas: [0.95, 0.999] + eps: 1.0e-8 + scheduler: + name: cosine + warmup_ratio: 0.05 + scheduler_kwargs: null + +data: + # Path to LeRobot-format dataset (with data/, meta/, videos/ subdirs) + data_path: /workspace/datasets/vla_arena_dreamzero + # Path to UMT5-XXL tokenizer directory + tokenizer_path: /workspace/models/Wan2.1-I2V-14B-480P/google/umt5-xxl + # Per-view resolution: 176x320 (two views tiled vertically: 352x640 before VAE encoding) + image_size: [176, 320] + max_text_length: 512 + embodiment_tag: libero + # Native multi-chunk dataloader + use_aligned_dataloader: false + max_chunk_size: 4 + macro_stride: 24 + relative_action: true + crop_ratio: 0.95 + color_jitter: true + # Embodiment ID mapping + embodiment_tag_mapping: + agibot: 26 + oxe_droid: 17 + bridge: 1 + libero: 17 + gr1_unified: 24 + mecka_hands: 27 + xdof: 22 + yam: 31 diff --git a/flagscale/train/datasets/dreamzero/__init__.py b/flagscale/train/datasets/dreamzero/__init__.py new file mode 100644 index 0000000000..8ad86a1938 --- /dev/null +++ b/flagscale/train/datasets/dreamzero/__init__.py @@ -0,0 +1,812 @@ +# Copyright (c) 2025, FlagScale Authors. All rights reserved. +""" +DreamZero Data Pipeline for FlagScale native training backend. + +Multi-anchor temporal sampling with rejection guarantees fixed-shape outputs +for batch_size > 1. Samples exactly max_chunk_size chunks per sample; if the +trajectory is too short, the sample is rejected and a new index is drawn. + +Data flow: + LeRobot parquet + video files + -> DreamZeroDataset.__getitem__ (multi-anchor sampling, load, transform) + -> DreamZeroCollator.__call__ (batch, tokenize text) + -> get_batch (move to device, dtype cast) + -> model.forward() + +Expected batch keys from get_batch: + - images: (B, 33, 2*H, 2*W, C) uint8 video frames (2-view tiled grid) + - action: (B, 96, 32) float32 normalized [-1, 1] + - action_mask: (B, 96, 32) bool + - state: (B, 4, 64) float32 + - state_mask: (B, 4, 64) bool + - text: (B, max_text_len) int64 tokenized text + - text_attention_mask: (B, max_text_len) int64 + - embodiment_id: (B,) int64 + - has_real_action: (B,) bool +""" + +import json +import logging +import random +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from torch.utils.data import DataLoader, Dataset, DistributedSampler + +logger = logging.getLogger(__name__) + +# Default video frame offsets within each chunk (8 frames per chunk) +DEFAULT_VIDEO_IN_CHUNK_OFFSETS = (0, 3, 6, 9, 12, 15, 18, 21) + + +class EmptyTemporalSampleError(ValueError): + """Raised when multi-anchor sampling cannot fill max_chunk_size chunks.""" + pass + + +class DreamZeroDataset(Dataset): + """Map-style dataset for DreamZero training with multi-anchor temporal sampling. + + Reads LeRobot v2 format: per-episode parquet files for low-dim data and + mp4 video files decoded via decord. Each sample produces exactly + max_chunk_size temporal chunks with fixed output shapes. + """ + + def __init__( + self, + data_path: str, + max_chunk_size: int = 4, + macro_stride: int = 24, + action_horizon: int = 24, + state_horizon: int = 1, + action_dim: int = 32, + max_state_dim: int = 64, + video_in_chunk_offsets: tuple[int, ...] = DEFAULT_VIDEO_IN_CHUNK_OFFSETS, + embodiment_tag: str = "libero", + embodiment_tag_mapping: dict[str, int] | None = None, + image_size: tuple[int, int] = (176, 320), + num_views: int = 2, + resample_attempts: int = 8, + relative_action: bool = True, + relative_action_keys: list[str] | None = None, + crop_ratio: float = 0.95, + color_jitter: bool = True, + training: bool = True, + ): + self.data_path = Path(data_path) + self.max_chunk_size = max_chunk_size + self.macro_stride = macro_stride + self.action_horizon = action_horizon + self.state_horizon = state_horizon + self.action_dim = action_dim + self.max_state_dim = max_state_dim + self.video_in_chunk_offsets = video_in_chunk_offsets + self.embodiment_tag = embodiment_tag + self.embodiment_tag_mapping = embodiment_tag_mapping or { + "agibot": 26, "oxe_droid": 17, "gr1_unified": 2, + "mecka_hands": 27, "libero": 17, "dream": 100, "lapa": 101, + } + self.image_size = image_size + self.num_views = num_views + self.resample_attempts = resample_attempts + self.relative_action = relative_action + self.relative_action_keys = relative_action_keys or ["joint_pos"] + self.crop_ratio = crop_ratio + self.color_jitter = color_jitter + self.training = training + + # Derived constants + self.frames_per_chunk = len(video_in_chunk_offsets) # 8 + self.total_video_frames = max_chunk_size * self.frames_per_chunk + 1 # 33 + self.total_action_steps = max_chunk_size * action_horizon # 96 + self.total_state_steps = max_chunk_size * state_horizon # 4 + + # Load dataset metadata + self._load_metadata() + self._load_normalization_stats() + + logger.info( + f"DreamZeroDataset: {len(self)} samples, {len(self._episodes)} episodes, " + f"max_chunk_size={max_chunk_size}, macro_stride={macro_stride}, " + f"data_path={data_path}" + ) + + def _load_metadata(self): + """Load LeRobot v2 dataset metadata and build sample index.""" + import pyarrow.parquet as pq + + meta_dir = self.data_path / "meta" + with open(meta_dir / "info.json") as f: + info = json.load(f) + + self._chunks_size = info.get("chunks_size", 1000) + self._fps = info.get("fps", 10) + self._features = info.get("features", {}) + + # Video keys + video_keys = [k for k, v in self._features.items() if v.get("dtype") == "video"] + self._video_keys = video_keys if video_keys else ["observation.images.image"] + + # Load modality.json for state/action slicing + modality_path = meta_dir / "modality.json" + if modality_path.exists(): + with open(modality_path) as f: + self._modality = json.load(f) + else: + self._modality = {} + + # Load task texts + self._tasks = {} + tasks_path = meta_dir / "tasks.jsonl" + if tasks_path.exists(): + with open(tasks_path) as f: + for line in f: + entry = json.loads(line) + self._tasks[entry["task_index"]] = entry["task"] + + # Load episodes metadata + episodes_path = meta_dir / "episodes.jsonl" + self._episodes = [] + with open(episodes_path) as f: + for line in f: + self._episodes.append(json.loads(line)) + + # Build sample index: (episode_idx, frame_in_ep) for all valid frames + # A frame is valid if it can anchor multi-chunk sampling + self._sample_index = [] + for ep in self._episodes: + ep_idx = ep["episode_index"] + ep_len = ep["length"] + for frame_idx in range(ep_len): + self._sample_index.append((ep_idx, frame_idx, ep_len)) + + # Pre-load parquet tables (they're small for LIBERO) + self._episode_tables = {} + data_dir = self.data_path / "data" + for ep in self._episodes: + ep_idx = ep["episode_index"] + chunk_idx = ep_idx // self._chunks_size + pf = data_dir / f"chunk-{chunk_idx:03d}" / f"episode_{ep_idx:06d}.parquet" + if pf.exists(): + self._episode_tables[ep_idx] = pq.read_table(pf) + + def _load_normalization_stats(self): + """Load q99 normalization stats for actions and states.""" + meta_dir = self.data_path / "meta" + stats_path = meta_dir / "relative_stats_dreamzero.json" + if stats_path.exists(): + with open(stats_path) as f: + self._norm_stats = json.load(f) + else: + # Fallback: try stats.json + stats_path = meta_dir / "stats.json" + if stats_path.exists(): + with open(stats_path) as f: + self._norm_stats = json.load(f) + else: + self._norm_stats = {} + + def __len__(self): + return len(self._sample_index) + + def __getitem__(self, idx: int) -> dict[str, Any]: + """Load a sample with multi-anchor rejection loop.""" + last_error = None + for attempt in range(self.resample_attempts): + try: + return self._build_sample(idx) + except EmptyTemporalSampleError as e: + last_error = e + idx = random.randint(0, len(self) - 1) + raise RuntimeError( + f"Failed to sample valid multi-anchor window after " + f"{self.resample_attempts} attempts: {last_error}" + ) + + def _sample_anchors(self, frame_in_ep: int, ep_len: int) -> list[int]: + """Sample exactly max_chunk_size anchors via multi-anchor expansion. + + Expands outward from frame_in_ep in steps of macro_stride. + Each anchor must allow a full action window (anchor + action_horizon - 1 < ep_len). + Raises EmptyTemporalSampleError if fewer than max_chunk_size anchors found. + """ + anchors = [] + + def try_add(anchor: int): + if len(anchors) >= self.max_chunk_size: + return + # Anchor must allow full action window and video offsets + max_offset = max(self.video_in_chunk_offsets) + if anchor < 0 or anchor + max_offset >= ep_len: + return False + if anchor + self.action_horizon - 1 >= ep_len: + return False + anchors.append(anchor) + return True + + # Start with the given frame + try_add(frame_in_ep) + + step = 1 + back_done = False + fwd_done = False + while len(anchors) < self.max_chunk_size and (not back_done or not fwd_done): + if not back_done: + back = frame_in_ep - self.macro_stride * step + if back < 0: + back_done = True + else: + try_add(back) + + if len(anchors) >= self.max_chunk_size: + break + + if not fwd_done: + fwd = frame_in_ep + self.macro_stride * step + if fwd >= ep_len: + fwd_done = True + else: + try_add(fwd) + + step += 1 + + if len(anchors) < self.max_chunk_size: + raise EmptyTemporalSampleError( + f"Only found {len(anchors)} anchors (need {self.max_chunk_size}) " + f"at frame {frame_in_ep}, ep_len={ep_len}" + ) + + return sorted(anchors) + + def _compute_indices(self, anchors: list[int]) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Compute video, action, and state frame indices from anchors. + + Returns: + video_indices: (33,) - 8 per chunk + 1 boundary + action_indices: (96,) - 24 per chunk + state_indices: (4,) - 1 per chunk + """ + video_idx = [] + action_idx = [] + state_idx = [] + + for anchor in anchors: + # Video: 8 frames at offsets within each chunk + for offset in self.video_in_chunk_offsets: + video_idx.append(anchor + offset) + # Action: 24 contiguous frames starting at anchor + for a in range(self.action_horizon): + action_idx.append(anchor + a) + # State: anchor frame + state_idx.append(anchor) + + # Add boundary frame: last video frame + 3 + last_video = video_idx[-1] + video_idx.append(last_video + 3) + + return ( + np.array(video_idx, dtype=np.int64), + np.array(action_idx, dtype=np.int64), + np.array(state_idx, dtype=np.int64), + ) + + def _build_sample(self, idx: int) -> dict[str, Any]: + """Build a single training sample.""" + ep_idx, frame_in_ep, ep_len = self._sample_index[idx] + + # 1. Multi-anchor temporal sampling + anchors = self._sample_anchors(frame_in_ep, ep_len) + video_indices, action_indices, state_indices = self._compute_indices(anchors) + + # Validate boundary frame + if video_indices[-1] >= ep_len: + raise EmptyTemporalSampleError( + f"Boundary frame {video_indices[-1]} >= ep_len {ep_len}" + ) + + # 2. Load video frames + images = self._load_video(ep_idx, video_indices) + + # 3. Load actions and states from parquet + table = self._episode_tables[ep_idx] + actions, action_mask = self._load_actions(table, action_indices, state_indices) + states, state_mask = self._load_states(table, state_indices) + + # 4. Load language + language = self._load_language(table, frame_in_ep, ep_idx) + + # 5. Embodiment ID + embodiment_id = self.embodiment_tag_mapping.get(self.embodiment_tag, 17) + + return { + "video": images, # (33, 2H, 2W, 3) uint8 + "action": actions, # (96, 32) float32 + "action_mask": action_mask, # (96, 32) bool + "state": states, # (4, 64) float32 + "state_mask": state_mask, # (4, 64) bool + "language": language, # str + "embodiment_id": embodiment_id, # int + "has_real_action": True, + } + + def _load_video(self, ep_idx: int, frame_indices: np.ndarray) -> np.ndarray: + """Load video frames, apply transforms, tile multiview. + + Returns: (T, 2*H, 2*W, 3) uint8 + """ + import decord + decord.bridge.set_bridge("native") + + episode_chunk = ep_idx // self._chunks_size + H, W = self.image_size # target per-view size + + video_keys_to_use = self._video_keys[:self.num_views] + views = [] + + for video_key in video_keys_to_use: + video_path = ( + self.data_path / "videos" + / f"chunk-{episode_chunk:03d}" + / video_key + / f"episode_{ep_idx:06d}.mp4" + ) + + if video_path.exists(): + vr = decord.VideoReader(str(video_path), num_threads=1) + n_total = len(vr) + safe_indices = [min(int(i), n_total - 1) for i in frame_indices] + batch = vr.get_batch(safe_indices).asnumpy() # (T, H_orig, W_orig, C) + + # Apply crop and resize + frames = self._apply_video_transforms(batch) + views.append(frames) + else: + logger.warning(f"Video not found: {video_path}") + views.append(np.zeros((len(frame_indices), H, W, 3), dtype=np.uint8)) + + # Pad missing views with black + while len(views) < self.num_views: + views.append(np.zeros((len(frame_indices), H, W, 3), dtype=np.uint8)) + + # Tile into 2x2 grid: (T, 2*H, 2*W, C) + T = len(frame_indices) + tiled = np.zeros((T, 2 * H, 2 * W, 3), dtype=np.uint8) + + # View 0 -> top-left (head/main camera) + tiled[:, :H, :W, :] = views[0] + # View 1 -> bottom-left (wrist camera) + if len(views) > 1: + tiled[:, H:, :W, :] = views[1] + # Top-right and bottom-right stay black (matching reference for 2-view) + + return tiled + + def _apply_video_transforms(self, frames: np.ndarray) -> np.ndarray: + """Apply crop, resize, and optional color jitter. Returns uint8 (T, H, W, 3).""" + from PIL import Image + + H, W = self.image_size + T, h_orig, w_orig, C = frames.shape + + result = np.empty((T, H, W, C), dtype=np.uint8) + + # Random crop params (same for all frames in sample) + if self.training and self.crop_ratio < 1.0: + crop_h = int(h_orig * self.crop_ratio) + crop_w = int(w_orig * self.crop_ratio) + top = random.randint(0, h_orig - crop_h) + left = random.randint(0, w_orig - crop_w) + else: + # Center crop + crop_h = int(h_orig * self.crop_ratio) + crop_w = int(w_orig * self.crop_ratio) + top = (h_orig - crop_h) // 2 + left = (w_orig - crop_w) // 2 + + for i in range(T): + frame = frames[i] + # Crop + frame = frame[top:top+crop_h, left:left+crop_w] + # Resize + img = Image.fromarray(frame) + img = img.resize((W, H), Image.BILINEAR) + result[i] = np.array(img) + + # Color jitter (simple brightness/contrast/saturation) + if self.training and self.color_jitter: + result = self._apply_color_jitter(result) + + return result + + def _apply_color_jitter(self, frames: np.ndarray) -> np.ndarray: + """Simple color jitter matching reference VideoColorJitter defaults.""" + # brightness=0.1, contrast=0.1, saturation=0.1, hue=0.05 + brightness = 1.0 + random.uniform(-0.1, 0.1) + contrast = 1.0 + random.uniform(-0.1, 0.1) + + frames = frames.astype(np.float32) + # Brightness + frames = frames * brightness + # Contrast + mean = frames.mean(axis=(1, 2), keepdims=True) + frames = (frames - mean) * contrast + mean + # Clip and convert back + frames = np.clip(frames, 0, 255).astype(np.uint8) + return frames + + def _load_actions( + self, table, action_indices: np.ndarray, state_indices: np.ndarray + ) -> tuple[np.ndarray, np.ndarray]: + """Load actions, apply relative action and q99 normalization. + + Returns: (96, action_dim) float32, (96, action_dim) bool + """ + n_rows = table.num_rows + action_modality = self._modality.get("action", {}) + + # Gather raw action components + raw_actions = [] + for idx in action_indices: + safe_idx = min(int(idx), n_rows - 1) + row_val = table.column("action")[safe_idx].as_py() + raw_actions.append(row_val) + + actions = np.array(raw_actions, dtype=np.float32) # (96, raw_action_dim) + raw_dim = actions.shape[1] + + # Apply relative action (per-chunk anchor subtraction) + if self.relative_action: + # Load states at anchor frames for relative computation + state_arr = [] + for idx in state_indices: + safe_idx = min(int(idx), n_rows - 1) + row_val = table.column("observation.state")[safe_idx].as_py() + state_arr.append(row_val) + state_arr = np.array(state_arr, dtype=np.float32) # (4, state_dim) + + # Per-chunk subtraction for relative_action_keys + for key in self.relative_action_keys: + if key not in action_modality: + continue + action_slice = slice( + action_modality[key]["start"], + action_modality[key]["end"], + ) + # Find corresponding state slice + state_modality = self._modality.get("state", {}) + if key in state_modality: + state_slice = slice( + state_modality[key]["start"], + state_modality[key]["end"], + ) + else: + continue + + a_start = action_modality[key]["start"] + a_end = action_modality[key]["end"] + s_start = state_modality[key]["start"] + s_end = state_modality[key]["end"] + d = min(a_end - a_start, s_end - s_start) + + for c in range(self.max_chunk_size): + rs = c * self.action_horizon + re = rs + self.action_horizon + ref = state_arr[c, s_start:s_start + d] + actions[rs:re, a_start:a_start + d] -= ref + + # Q99 normalization + actions = self._normalize_q99(actions, action_modality) + + # Pad to action_dim + if raw_dim < self.action_dim: + pad = np.zeros((actions.shape[0], self.action_dim - raw_dim), dtype=np.float32) + actions = np.concatenate([actions, pad], axis=1) + elif raw_dim > self.action_dim: + actions = actions[:, :self.action_dim] + + # Clip to [-1, 1] + actions = np.clip(actions, -1.0, 1.0) + + # Action mask: True for real dims, False for padding + mask = np.zeros((self.total_action_steps, self.action_dim), dtype=bool) + mask[:, :raw_dim] = True + + return actions, mask + + def _load_states( + self, table, state_indices: np.ndarray + ) -> tuple[np.ndarray, np.ndarray]: + """Load states and normalize. + + Returns: (4, max_state_dim) float32, (4, max_state_dim) bool + """ + n_rows = table.num_rows + + raw_states = [] + for idx in state_indices: + safe_idx = min(int(idx), n_rows - 1) + row_val = table.column("observation.state")[safe_idx].as_py() + raw_states.append(row_val) + + states = np.array(raw_states, dtype=np.float32) # (4, raw_state_dim) + raw_dim = states.shape[1] + + # Q99 normalization for states + state_modality = self._modality.get("state", {}) + states = self._normalize_q99_state(states, state_modality) + + # Pad to max_state_dim + if raw_dim < self.max_state_dim: + pad = np.zeros((states.shape[0], self.max_state_dim - raw_dim), dtype=np.float32) + states = np.concatenate([states, pad], axis=1) + elif raw_dim > self.max_state_dim: + states = states[:, :self.max_state_dim] + + # Clip + states = np.clip(states, -1.0, 1.0) + + # State mask + mask = np.zeros((self.total_state_steps, self.max_state_dim), dtype=bool) + mask[:, :raw_dim] = True + + return states, mask + + def _normalize_q99(self, data: np.ndarray, modality_info: dict) -> np.ndarray: + """Apply q99 normalization: maps [q01, q99] -> [-1, 1].""" + if not self._norm_stats: + return data + + for key, meta in modality_info.items(): + if key not in self._norm_stats: + continue + stats = self._norm_stats[key] + if "q01" not in stats or "q99" not in stats: + continue + + start = meta["start"] + end = meta["end"] + q01 = np.array(stats["q01"], dtype=np.float32) + q99 = np.array(stats["q99"], dtype=np.float32) + + denom = q99 - q01 + valid = denom > 1e-8 + + # Normalize: 2 * (x - q01) / (q99 - q01) - 1 + d = end - start + for i in range(d): + if valid[i]: + data[:, start + i] = 2.0 * (data[:, start + i] - q01[i]) / denom[i] - 1.0 + + return data + + def _normalize_q99_state(self, data: np.ndarray, modality_info: dict) -> np.ndarray: + """Apply q99 normalization to states (uses main stats.json if available).""" + # States use absolute stats, not relative + meta_dir = self.data_path / "meta" + stats_path = meta_dir / "stats.json" + if not stats_path.exists(): + return data + + with open(stats_path) as f: + stats = json.load(f) + + # Look for observation.state stats + obs_key = "observation.state" + if obs_key not in stats: + return data + + obs_stats = stats[obs_key] + if "q01" not in obs_stats or "q99" not in obs_stats: + return data + + q01 = np.array(obs_stats["q01"], dtype=np.float32) + q99 = np.array(obs_stats["q99"], dtype=np.float32) + denom = q99 - q01 + valid = denom > 1e-8 + + raw_dim = min(data.shape[1], len(q01)) + for i in range(raw_dim): + if valid[i]: + data[:, i] = 2.0 * (data[:, i] - q01[i]) / denom[i] - 1.0 + + return data + + def _load_language(self, table, frame_idx: int, ep_idx: int) -> str: + """Load language instruction for this sample.""" + if "annotation.task" in table.column_names: + val = table.column("annotation.task")[frame_idx].as_py() + if val: + return str(val) + if "task_index" in table.column_names: + task_idx = table.column("task_index")[frame_idx].as_py() + return self._tasks.get(task_idx, "") + return "" + + def _load_language(self, table, frame_idx: int, ep_idx: int) -> str: + """Load language instruction for the sample.""" + if "annotation.task" in table.column_names: + val = table.column("annotation.task")[frame_idx].as_py() + if val: + return str(val) + + if "task_index" in table.column_names: + task_idx = table.column("task_index")[frame_idx].as_py() + return self._tasks.get(task_idx, "") + + # Fallback from episodes metadata + for ep in self._episodes: + if ep["episode_index"] == ep_idx: + tasks = ep.get("tasks", []) + if tasks: + return tasks[0] + return "" + + +class DreamZeroCollator: + """Collate DreamZero samples into batches with text tokenization. + + Handles: + - Stacking numpy arrays into tensors + - Tokenizing language instructions + - Creating embodiment_id tensor + """ + + def __init__( + self, + tokenizer_path: str, + max_text_length: int = 512, + ): + from transformers import AutoTokenizer + self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path) + self.max_text_length = max_text_length + + def __call__(self, features: list[dict[str, Any]]) -> dict[str, torch.Tensor]: + """Collate a list of samples into a batch.""" + # Stack numpy arrays -> tensors + images = torch.from_numpy(np.stack([f["video"] for f in features])) + actions = torch.from_numpy(np.stack([f["action"] for f in features])) + action_masks = torch.from_numpy(np.stack([f["action_mask"] for f in features])) + states = torch.from_numpy(np.stack([f["state"] for f in features])) + state_masks = torch.from_numpy(np.stack([f["state_mask"] for f in features])) + + has_real_action = torch.tensor( + [f["has_real_action"] for f in features], dtype=torch.bool + ) + + # Tokenize language instructions + languages = [f["language"] for f in features] + tokenized = self.tokenizer( + languages, + max_length=self.max_text_length, + padding="max_length", + truncation=True, + return_tensors="pt", + ) + + # Embodiment IDs + embodiment_ids = torch.tensor( + [f["embodiment_id"] for f in features], dtype=torch.long + ) + + return { + "images": images, # (B, 33, 2H, 2W, 3) uint8 + "action": actions, # (B, 96, 32) float32 + "action_mask": action_masks, # (B, 96, 32) bool + "has_real_action": has_real_action, # (B,) bool + "state": states, # (B, 4, 64) float32 + "state_mask": state_masks, # (B, 4, 64) bool + "text": tokenized["input_ids"], # (B, max_text_len) int64 + "text_attention_mask": tokenized["attention_mask"], # (B, max_text_len) int64 + "embodiment_id": embodiment_ids, # (B,) int64 + } + + +def build_dataloader( + data_path: str, + tokenizer_path: str, + batch_size: int = 1, + num_workers: int = 8, + max_chunk_size: int = 4, + macro_stride: int = 24, + action_horizon: int = 24, + state_horizon: int = 1, + action_dim: int = 32, + max_state_dim: int = 64, + embodiment_tag: str = "libero", + embodiment_tag_mapping: dict[str, int] | None = None, + image_size: tuple[int, int] = (176, 320), + max_text_length: int = 512, + num_views: int = 2, + shuffle: bool = True, + distributed: bool = True, + relative_action: bool = True, + crop_ratio: float = 0.95, + color_jitter: bool = True, + training: bool = True, +) -> DataLoader: + """Build DataLoader for DreamZero training with multi-anchor sampling. + + All samples are guaranteed to have exactly max_chunk_size chunks, so + np.stack in the collator always succeeds for batch_size > 1. + """ + dataset = DreamZeroDataset( + data_path=data_path, + max_chunk_size=max_chunk_size, + macro_stride=macro_stride, + action_horizon=action_horizon, + state_horizon=state_horizon, + action_dim=action_dim, + max_state_dim=max_state_dim, + embodiment_tag=embodiment_tag, + embodiment_tag_mapping=embodiment_tag_mapping, + image_size=image_size, + num_views=num_views, + relative_action=relative_action, + crop_ratio=crop_ratio, + color_jitter=color_jitter, + training=training, + ) + + collator = DreamZeroCollator( + tokenizer_path=tokenizer_path, + max_text_length=max_text_length, + ) + + sampler = None + if distributed and torch.distributed.is_initialized(): + sampler = DistributedSampler( + dataset, + num_replicas=torch.distributed.get_world_size(), + rank=torch.distributed.get_rank(), + shuffle=shuffle, + ) + shuffle = False + + dataloader = DataLoader( + dataset, + batch_size=batch_size, + shuffle=shuffle if sampler is None else False, + num_workers=num_workers, + collate_fn=collator, + sampler=sampler, + pin_memory=True, + drop_last=True, + prefetch_factor=2 if num_workers > 0 else None, + ) + + logger.info( + f"Built DreamZero DataLoader: {len(dataset)} samples, " + f"batch_size={batch_size}, max_chunk_size={max_chunk_size}, " + f"distributed={distributed}" + ) + return dataloader + + +def get_batch( + batch: dict[str, torch.Tensor], + device: torch.device, + compute_dtype: torch.dtype = torch.bfloat16, +) -> dict[str, torch.Tensor]: + """Move batch to device and cast dtypes for model consumption.""" + result = {} + + # Images stay as uint8 on device (VAE handles normalization internally) + result["images"] = batch["images"].to(device, non_blocking=True) + + # Actions and states in compute dtype + result["action"] = batch["action"].to(device, dtype=compute_dtype, non_blocking=True) + result["action_mask"] = batch["action_mask"].to(device, non_blocking=True) + result["state"] = batch["state"].to(device, dtype=compute_dtype, non_blocking=True) + result["state_mask"] = batch["state_mask"].to(device, non_blocking=True) + + # Text tokens as int64 + result["text"] = batch["text"].to(device, non_blocking=True) + result["text_attention_mask"] = batch["text_attention_mask"].to(device, non_blocking=True) + + # Scalar tensors + result["has_real_action"] = batch["has_real_action"].to(device, non_blocking=True) + result["embodiment_id"] = batch["embodiment_id"].to(device, non_blocking=True) + + return result diff --git a/flagscale/train/datasets/dreamzero/aligned_dataloader.py b/flagscale/train/datasets/dreamzero/aligned_dataloader.py new file mode 100644 index 0000000000..705e2cea85 --- /dev/null +++ b/flagscale/train/datasets/dreamzero/aligned_dataloader.py @@ -0,0 +1,298 @@ +# Copyright (c) 2025, FlagScale Authors. All rights reserved. +""" +DreamZero Aligned Data Pipeline for FlagScale native training backend. + +Uses reference (groot) ShardedLeRobotMixtureDataset with the FULL transform +chain to produce identical data to the reference training pipeline. + +Transform chain: + VideoToTensor -> VideoCrop(0.95) -> VideoResize(176,320) -> VideoColorJitter + -> VideoToNumpy -> StateActionToTensor -> StateActionTransform(q99) + -> ConcatTransform -> DreamTransform + +Output per sample (after collation): + - images: (B, 33, 352, 640, 3) uint8 + - action: (B, 96, 32) float32 + - action_mask: (B, 96, 32) bool + - state: (B, 4, 64) float32 + - state_mask: (B, 4, 64) bool + - text: (B, max_len) int64 (tokenized) + - text_attention_mask: (B, max_len) int64 + - embodiment_id: (B,) int64 + - has_real_action: (B,) bool +""" + +import logging +import os +import sys +from typing import Any + +import torch +from torch.utils.data import DataLoader + +logger = logging.getLogger(__name__) + +# Ensure reference repo is importable +_REFERENCE_REPO = "/public-mixed/fengyupu/github/dreamzero" +if _REFERENCE_REPO not in sys.path: + sys.path.insert(0, _REFERENCE_REPO) + + +def build_dataloader_aligned( + data_path: str, + tokenizer_path: str, + batch_size: int = 1, + num_workers: int = 1, + num_frames: int = 25, + action_horizon: int = 24, + state_horizon: int = 1, + action_dim: int = 32, + max_state_dim: int = 64, + embodiment_tag: str = "libero", + embodiment_tag_mapping: dict[str, int] | None = None, + deterministic: bool = False, + image_size: tuple[int, int] = (176, 320), + max_text_length: int = 512, + num_views: int = 2, + seed: int = 42, + shard_sampling_rate: float = 0.1, + max_chunk_size: int = 1, +) -> DataLoader: + """Build DataLoader using the reference pipeline for data alignment. + + Uses groot's ShardedLeRobotMixtureDataset (IterableDataset) with the full + transform chain. NO external sampler — dataset handles distributed splitting. + """ + from groot.vla.data.dataset.lerobot_sharded import ( + ShardedLeRobotMixtureDataset, + ShardedLeRobotSubLangSingleActionChunkDatasetDROID, + ) + from groot.vla.data.dataset.lerobot import ModalityConfig + from groot.vla.data.transform.base import ComposedModalityTransform + from groot.vla.data.transform.video import ( + VideoToTensor, VideoCrop, VideoResize, VideoColorJitter, VideoToNumpy, + ) + from groot.vla.data.transform.state_action import ( + StateActionToTensor, StateActionTransform, + ) + from groot.vla.data.transform.concat import ConcatTransform + from groot.vla.model.dreamzero.transform.dreamzero_cotrain import ( + DefaultDataCollator, DreamTransform, + ) + + if embodiment_tag_mapping is None: + embodiment_tag_mapping = { + "agibot": 26, "oxe_droid": 17, "gr1_unified": 2, + "mecka_hands": 27, "xdof": 30, "yam": 31, + "dream": 100, "lapa": 101, "libero": 17, + } + + # --- Modality keys --- + video_keys = ["video.image", "video.wrist_image"] + state_keys = ["state.joint_pos", "state.gripper_pos"] + action_keys = ["action.joint_pos", "action.gripper_pos"] + + # --- Modality configs --- + # NOTE: delta_indices must match the reference config (base_48_wan_fine_aug_relative.yaml) + # The reference uses 25 video frames (0..24) and 24 action steps (0..23) in its delta_indices, + # regardless of num_frames passed to the model. num_frames controls the model's sequence length, + # but delta_indices controls which timesteps are loaded from the dataset. + video_delta_indices = list(range(25)) # Reference hardcodes [0..24] + action_delta_indices = list(range(action_horizon)) # [0..23] + state_delta_indices = list(range(state_horizon)) # [0] + + modality_configs = { + embodiment_tag: { + "video": ModalityConfig( + delta_indices=video_delta_indices, + modality_keys=video_keys, + ), + "state": ModalityConfig( + delta_indices=state_delta_indices, + modality_keys=state_keys, + ), + "action": ModalityConfig( + delta_indices=action_delta_indices, + modality_keys=action_keys, + ), + "language": ModalityConfig( + delta_indices=[0], + modality_keys=["annotation.task"], + ), + } + } + + # --- Full transform pipeline (matching reference config exactly) --- + h, w = image_size + if deterministic: + # Deterministic mode: disable random crop and color jitter for exact reproducibility. + # Use center crop and no jitter so that the only source of randomness is + # the shard/trajectory sampling (controlled by seed). + # NOTE: VideoCrop.apply() uses self.training to choose RandomCrop vs CenterCrop. + # The dataset's __init__ calls transforms.train() which resets training=True. + # We subclass VideoCrop to lock it in eval mode (CenterCrop always). + class DeterministicCrop(VideoCrop): + """VideoCrop that always uses CenterCrop regardless of train/eval mode.""" + def train(self): + pass # ignore — stay in eval mode + + def apply(self, data): + # Force eval path (CenterCrop) + self.training = False + return super().apply(data) + + crop = DeterministicCrop(apply_to=video_keys, scale=0.95) + transform_list = [ + VideoToTensor(apply_to=video_keys), + crop, + VideoResize(apply_to=video_keys, height=h, width=w, interpolation="linear"), + # No VideoColorJitter in deterministic mode + VideoToNumpy(apply_to=video_keys), + ] + else: + transform_list = [ + # Video transforms + VideoToTensor(apply_to=video_keys), + VideoCrop(apply_to=video_keys, scale=0.95, mode="random"), + VideoResize(apply_to=video_keys, height=h, width=w, interpolation="linear"), + VideoColorJitter( + apply_to=video_keys, + brightness=0.3, contrast=0.4, saturation=0.5, hue=0.08, + ), + VideoToNumpy(apply_to=video_keys), + ] + transform_list += [ + # State transforms + StateActionToTensor(apply_to=state_keys), + StateActionTransform( + apply_to=state_keys, + normalization_modes={ + "state.joint_pos": "q99", + "state.gripper_pos": "q99", + }, + ), + # Action transforms + StateActionToTensor(apply_to=action_keys), + StateActionTransform( + apply_to=action_keys, + normalization_modes={ + "action.joint_pos": "q99", + "action.gripper_pos": "q99", + }, + ), + # Consolidation: stack views, concat state/action dims + ConcatTransform( + video_concat_order=video_keys, + state_concat_order=state_keys, + action_concat_order=action_keys, + ), + # Model-specific final transform + DreamTransform( + num_views=num_views, + action_horizon=action_horizon, + state_horizon=state_horizon, + max_action_dim=action_dim, + max_state_dim=max_state_dim, + embodiment_tag_mapping=embodiment_tag_mapping, + training=True, + default_instruction="", + tokenizer_path=tokenizer_path, + ), + ] + + transforms = { + embodiment_tag: ComposedModalityTransform(transforms=transform_list) + } + + # --- Mixture spec --- + mixture_spec = [{ + "dataset_path": {embodiment_tag: [data_path]}, + "dataset_weight": 1.0, + "distribute_weights": True, + }] + + # --- Instantiate dataset --- + logger.info(f"Building ALIGNED ShardedLeRobotMixtureDataset from {data_path}") + train_dataset = ShardedLeRobotMixtureDataset.from_mixture_spec( + mixture_spec=mixture_spec, + dataset_class=ShardedLeRobotSubLangSingleActionChunkDatasetDROID, + all_modality_configs=modality_configs, + all_transforms=transforms, + metadata_versions={embodiment_tag: None}, + fps={embodiment_tag: None}, + dataset_kwargs={ + "video_backend": "decord", + "use_global_metadata": False, + "max_chunk_size": max_chunk_size, + "relative_action": True, + "relative_action_keys": ["joint_pos"], + "relative_action_per_horizon": False, + }, + mixture_kwargs={ + "training": True, + "balance_dataset_weights": False, + "seed": seed, + "shard_sampling_rate": shard_sampling_rate, + }, + ) + + # --- Collator --- + collator = DefaultDataCollator( + tokenizer_path=tokenizer_path, + max_length=max_text_length, + num_views=num_views, + embodiment_tag_mapping=embodiment_tag_mapping, + ) + + # --- DataLoader (no sampler for IterableDataset) --- + dataloader = DataLoader( + train_dataset, + batch_size=batch_size, + num_workers=num_workers, + collate_fn=collator, + pin_memory=False, + persistent_workers=num_workers > 0, + ) + + logger.info( + f"Built ALIGNED DataLoader: IterableDataset, " + f"batch_size={batch_size}, num_workers={num_workers}, seed={seed}" + ) + return dataloader + + +def get_batch_aligned( + batch: dict[str, torch.Tensor], + device: torch.device, + compute_dtype: torch.dtype = torch.bfloat16, +) -> dict[str, torch.Tensor]: + """Move batch to device and cast dtypes for model consumption.""" + result = {} + + # Images stay uint8 (VAE handles conversion) + if "images" in batch: + result["images"] = batch["images"].to(device, non_blocking=True) + + # Actions/states to compute dtype + for key in ["action", "state"]: + if key in batch: + result[key] = batch[key].to(device=device, dtype=compute_dtype, non_blocking=True) + + # Masks + for key in ["action_mask", "state_mask"]: + if key in batch: + result[key] = batch[key].to(device, non_blocking=True) + + # Text tokens (int64) + if "text" in batch: + result["text"] = batch["text"].to(device, non_blocking=True) + if "text_attention_mask" in batch: + result["text_attention_mask"] = batch["text_attention_mask"].to(device, non_blocking=True) + + # Scalars + if "has_real_action" in batch: + result["has_real_action"] = batch["has_real_action"].to(device, non_blocking=True) + if "embodiment_id" in batch: + result["embodiment_id"] = batch["embodiment_id"].to(device, non_blocking=True) + + return result diff --git a/flagscale/train/models/dreamzero/__init__.py b/flagscale/train/models/dreamzero/__init__.py new file mode 100644 index 0000000000..54f8eccab4 --- /dev/null +++ b/flagscale/train/models/dreamzero/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) 2025, FlagScale Authors. All rights reserved. +"""DreamZero model for FlagScale native training backend.""" + +from flagscale.train.models.dreamzero.dreamzero_model import DreamZeroPolicy + +__all__ = ["DreamZeroPolicy"] diff --git a/flagscale/train/models/dreamzero/action_head/__init__.py b/flagscale/train/models/dreamzero/action_head/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/flagscale/train/models/dreamzero/action_head/wan_flow_matching_action_tf.py b/flagscale/train/models/dreamzero/action_head/wan_flow_matching_action_tf.py new file mode 100644 index 0000000000..0fd4e381b3 --- /dev/null +++ b/flagscale/train/models/dreamzero/action_head/wan_flow_matching_action_tf.py @@ -0,0 +1,1413 @@ +from dataclasses import dataclass, field +import logging +import time +from typing import TypeAlias, cast +import os + +from accelerate import load_checkpoint_and_dispatch + +from einops import rearrange +from hydra.utils import instantiate +from peft import LoraConfig, get_peft_model +import torch +from torch import nn +import torch.distributed as dist +from torch.distributed.device_mesh import DeviceMesh +from safetensors.torch import load_file +import json +from huggingface_hub import hf_hub_download + + +logger = logging.getLogger(__name__) + +WAN_HF_REPO_ID = "Wan-AI/Wan2.1-I2V-14B-480P" +WAN22_HF_REPO_ID = "Wan-AI/Wan2.2-TI2V-5B" + + +def hf_download(filename: str, repo_id: str = WAN_HF_REPO_ID) -> str: + """Download a file from the specified HuggingFace repo to HF cache.""" + path = hf_hub_download(repo_id=repo_id, filename=filename) + return path + + +def ensure_file(path: str | None, hf_filename: str, repo_id: str = WAN_HF_REPO_ID) -> str: + """Return a valid local path: use `path` if it exists, otherwise download from HuggingFace.""" + if path is not None and os.path.exists(path): + return path + return hf_download(hf_filename, repo_id) + +from torch.distributions import Beta +import torch.distributed as dist +from torch.distributed.device_mesh import DeviceMesh +from torchvision.transforms import v2 +from transformers import PretrainedConfig +from transformers.feature_extraction_utils import BatchFeature + +from flagscale.train.models.dreamzero.n1_5.action_head.base_action_head import ActionHead +from flagscale.train.models.dreamzero.modules.flow_match_scheduler import FlowMatchScheduler +from flagscale.train.models.dreamzero.modules.vram_management import enable_vram_management, AutoWrappedModule, AutoWrappedLinear +from flagscale.train.models.dreamzero.modules.wan_video_text_encoder import T5RelativeEmbedding, T5LayerNorm +from flagscale.train.models.dreamzero.modules.flow_unipc_multistep_scheduler import FlowUniPCMultistepScheduler + + +KVCacheType: TypeAlias = torch.Tensor + +@dataclass +class WANPolicyHeadConfig(PretrainedConfig): + add_pos_embed: bool = field( + default=True, metadata={"help": "Whether to add positional embedding"} + ) + model_dtype: str = field(default="float32", metadata={"help": "Model data type."}) + diffusion_model_cfg: dict = field( + default=None, metadata={"help": "Diffusion model configuration."} + ) + input_embedding_dim: int = field( + default=1536, metadata={"help": "Input embedding channel dimension."} + ) + backbone_embedding_dim: int = field( + default=1536, metadata={"help": "Backbone embedding channel dimension."} + ) + tiled: bool = field(default=True, metadata={"help": "Whether to use tiled input."}) + tile_size_height: int = field(default=34, metadata={"help": "Tile size height."}) + tile_size_width: int = field(default=34, metadata={"help": "Tile size width."}) + tile_stride_height: int = field(default=18, metadata={"help": "Tile stride height."}) + tile_stride_width: int = field(default=16, metadata={"help": "Tile stride width."}) + num_frame_per_block: int = field(default=1, metadata={"help": "Number of frames per block."}) + # Target video (H, W) for Wan22 resize. When set, videos are resized to this before VAE so latent + # spatial size matches. Use height/width divisible by 32 for WanVideoVAE38 (16x) so latent H,W are even. + target_video_height: int | None = field(default=None, metadata={"help": "Target video height for resize (e.g. 160 for even latent with VAE38)."}) + target_video_width: int | None = field(default=None, metadata={"help": "Target video width for resize (e.g. 320)."}) + + lora_rank: int = field(default=4, metadata={"help": "LoRA rank."}) + lora_alpha: int = field(default=4, metadata={"help": "LoRA alpha."}) + lora_target_modules: str = field(default="q,k,v,o,ffn.0,ffn.2") + init_lora_weights: str = field(default="kaiming", metadata={"help": "LoRA initialization method."}) + train_architecture: str= field(default="lora", metadata={"help": "Train architecture."}) + skip_component_loading: bool = field(default=False, metadata={"help": "Skip loading individual component weights (used when loading from full pretrained model)."}) + + use_gradient_checkpointing: bool = field(default=True, metadata={"help": "Whether to use gradient checkpointing."}) + qformer_cfg: dict = field(default=None, metadata={"help": "Qformer configuration."}) + hidden_size: int = field(default=1024, metadata={"help": "Input embedding dimension."}) + max_seq_len: int = field(default=1024, metadata={"help": "Maxium Sequence Length"}) + action_dim: int = field(default=None, metadata={"help": "Action dimension."}) + action_horizon: int = field(default=None, metadata={"help": "Action horizon."}) + noise_beta_alpha: float = field(default=1.5, metadata={"help": ""}) + noise_beta_beta: float = field(default=1.0, metadata={"help": ""}) + noise_s: float = field( + default=0.999, metadata={"help": "Flow matching noise Beta distribution s."} + ) + # High noise emphasis for BASE (coupled) training - applies Beta distribution to BOTH video and action together + use_high_noise_emphasis: bool = field( + default=False, metadata={"help": "Use Beta distribution for noise sampling (biases BOTH video and action towards high noise levels together)."} + ) + high_noise_beta_alpha: float = field( + default=3.0, metadata={"help": "Beta alpha for high noise emphasis. Beta(3,1): mean=0.75, Beta(5,1): mean=0.83. Higher = more high noise bias."} + ) + # Decoupled noise sampling config for training-inference alignment + # When enabled: video uses Beta(alpha,beta) biased towards high noise, action uses independent uniform + decouple_video_action_noise: bool = field( + default=False, metadata={"help": "Decouple video/action noise: video uses Beta distribution (high noise bias), action uses independent uniform."} + ) + video_noise_beta_alpha: float = field( + default=3.0, metadata={"help": "Beta alpha for video noise. Beta(3,1): mean=0.75, Beta(5,1): mean=0.83. Higher alpha = more bias to high noise."} + ) + video_noise_beta_beta: float = field( + default=1.0, metadata={"help": "Beta beta for video noise. Keep at 1.0."} + ) + # Decoupled inference config - allows video to stay noisy while action fully denoises + decouple_inference_noise: bool = field( + default=False, metadata={"help": "Use decoupled noise schedules during inference (video stays noisy, action fully denoises)."} + ) + video_inference_final_noise: float = field( + default=0.8, metadata={"help": "Final noise level for video during decoupled inference (0.0-1.0). E.g., 0.8 means video ends at 80% noise."} + ) + num_timestep_buckets: int = field( + default=1000, metadata={"help": "Number of timestep discretization buckets."} + ) + num_inference_timesteps: int = field( + default=None, + metadata={"help": "Number of inference steps for noise diffusion."}, + ) + max_num_embodiments: int = field(default=32, metadata={"help": "Number of embodiments."}) + tune_projector: bool = field(default=True, metadata={"help": "Whether to tune the projector."}) + tune_diffusion_model: bool = field( + default=True, metadata={"help": "Whether to tune the diffusion model."} + ) + load_pretrained_det_decode_layer_path: str = field( + default=None, metadata={"help": "Path to pretrained detection model."} + ) + detection_coeff: float = field(default=1.0, metadata={"help": "Detection coefficient."}) + + freeze_decode_layer: bool = field(default=False) + expand_batch: int = field(default=None) + use_vlln: bool = field(default=True) + defer_lora_injection: bool = field(default=False, metadata={"help": "Defer LoRA injection until after loading pretrained weights."}) + + vl_self_attention_cfg: dict = field(default=None) + text_encoder_cfg: dict = field(default=None) + image_encoder_cfg: dict = field(default=None) + vae_cfg: dict = field(default=None) + + def __init__(self, **kwargs): + super().__init__(**kwargs) + for key, value in kwargs.items(): + setattr(self, key, value) + + +class WANPolicyHead(ActionHead): + config_class = WANPolicyHeadConfig + supports_gradient_checkpointing = True + + def __init__( + self, + config: WANPolicyHeadConfig, + ): + super().__init__() + self.tiled = config.tiled + self.tile_size_height = config.tile_size_height + self.tile_size_width = config.tile_size_width + self.tile_stride_height = config.tile_stride_height + self.tile_stride_width = config.tile_stride_width + self.num_frame_per_block = config.num_frame_per_block + self.hidden_size = config.hidden_size + self.num_frames = config.num_frames + self.text_encoder = instantiate(config.text_encoder_cfg) + self.image_encoder = instantiate(config.image_encoder_cfg) + self.vae = instantiate(config.vae_cfg) + self.scheduler = FlowMatchScheduler(shift=5, sigma_min=0.0, extra_one_step=True) + self.model_names = ['text_encoder'] + + self.num_inference_steps = 16 + self.seed = 1140 + self.cfg_scale = 5.0 + self.denoising_strength = 1.0 + self.sigma_shift = 5.0 + self.kv_cache1: KVCacheType | None = None + self.kv_cache_neg: KVCacheType | None = None + self.crossattn_cache: KVCacheType | None = None + self.crossattn_cache_neg: KVCacheType | None = None + + self.global_step = 0 + self.max_steps = 0 + self.lora_rank = config.lora_rank + self.lora_alpha = config.lora_alpha + self.lora_target_modules = config.lora_target_modules + self.init_lora_weights = config.init_lora_weights + self.train_architecture = config.train_architecture + self.clip_feas = None + self.ys = None + self.current_start_frame = 0 + self.language = None + + self.ip_rank = 0 + self.ip_size = 1 + self.ip_group = None + + self._device = "cuda" + self.dynamic_cache_schedule = os.getenv("DYNAMIC_CACHE_SCHEDULE", "False").lower() == "true" + + + num_dit_steps = 8 + if os.getenv("NUM_DIT_STEPS") is not None: + num_dit_steps = int(os.getenv("NUM_DIT_STEPS")) + if num_dit_steps == 5: + self.dit_step_mask = [True, True, True, False, False, False, False, True, False, False, False, False, True, False, False, False] + elif num_dit_steps == 6: + self.dit_step_mask = [True, True, False, False, False, True, False, False, False, False, True, False, False, False, True, True] + elif num_dit_steps == 7: + self.dit_step_mask = [True, True, True, False, False, False, True, False, False, False, True, False, False, False, True, True] + elif num_dit_steps == 8: + self.dit_step_mask = [True, True, True, False, False, False, True, False, False, False, True, False, False, True, True, True] + else: + self.dit_step_mask = [True, True, True, True, True, True, True, True, True, True, True, True, True, True, True, True] + assert self.dit_step_mask[0] == True, "first step must be True" + + self.normalize_video = v2.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]) + + + self.use_gradient_checkpointing = config.use_gradient_checkpointing + if self.training: + self.scheduler.set_timesteps(1000, training=True) + + + self.input_embedding_dim = config.input_embedding_dim + + self.cpu_offload = False + + self.model = instantiate(config.diffusion_model_cfg) + self.action_dim = config.action_dim + self.action_horizon = config.action_horizon + self.num_inference_timesteps = config.num_inference_timesteps + + text_enc_path = ensure_file( + self.text_encoder.text_encoder_pretrained_path, + "models_t5_umt5-xxl-enc-bf16.pth", + ) + self.text_encoder.load_state_dict(torch.load(text_enc_path, map_location='cpu')) + + img_enc_path = ensure_file( + self.image_encoder.image_encoder_pretrained_path, + "models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth", + ) + self.image_encoder.model.load_state_dict(torch.load(img_enc_path, map_location='cpu'), strict=False) + + # Wan2.2 (WanVideoVAE38, z_dim=48) uses Wan2.2_VAE.pth; Wan2.1 uses Wan2.1_VAE.pth + vae_hf_filename = "Wan2.2_VAE.pth" if getattr(self.vae, "z_dim", 16) == 48 else "Wan2.1_VAE.pth" + vae_repo_id = WAN22_HF_REPO_ID if getattr(self.vae, "z_dim", 16) == 48 else WAN_HF_REPO_ID + vae_path = ensure_file( + self.vae.vae_pretrained_path, + vae_hf_filename, + repo_id=vae_repo_id, + ) + self.vae.model.load_state_dict(torch.load(vae_path, map_location='cpu')) + + if not config.skip_component_loading: + dit_dir = self.model.diffusion_model_pretrained_path + # Wan2.2 (in_dim=48) uses Wan2.2-TI2V-5B repo; Wan2.1 uses Wan2.1-I2V-14B-480P + dit_repo_id = WAN22_HF_REPO_ID if getattr(self.model, "in_dim", 16) == 48 else WAN_HF_REPO_ID + if dit_dir is None or not os.path.isdir(dit_dir): + index_path = hf_hub_download(repo_id=dit_repo_id, filename="diffusion_pytorch_model.safetensors.index.json") + dit_dir = os.path.dirname(index_path) + with open(index_path, 'r') as f: + index = json.load(f) + for shard_file in set(index["weight_map"].values()): + hf_hub_download(repo_id=dit_repo_id, filename=shard_file) + + if dit_dir is not None: + safetensors_path = os.path.join(dit_dir, "diffusion_pytorch_model.safetensors") + safetensors_index_path = os.path.join(dit_dir, "diffusion_pytorch_model.safetensors.index.json") + state_dict = {} + + if os.path.exists(safetensors_index_path): + # Handle sharded safetensors + print(f"Loading sharded safetensors using index: {safetensors_index_path}") + + with open(safetensors_index_path, 'r') as f: + index = json.load(f) + + # Load each shard + for shard_file in set(index["weight_map"].values()): + shard_path = os.path.join(dit_dir, shard_file) + print(f"Loading shard: {shard_path}") + shard_state_dict = load_file(shard_path) + state_dict.update(shard_state_dict) + + elif os.path.exists(safetensors_path): + # Handle single safetensors file + print(f"Loading weights from safetensors: {safetensors_path}") + state_dict = load_file(safetensors_path) + + else: + raise ValueError(f"No safetensors file found at {safetensors_path} or {safetensors_index_path}") + + missing_keys, unexpected_keys = self.model.load_state_dict(state_dict, strict=False) + + if missing_keys: + print(f"Missing keys when loading pretrained weights: {missing_keys}") + if unexpected_keys: + print(f"Unexpected keys when loading pretrained weights: {unexpected_keys}") + + print("Successfully loaded pretrained weights") + else: + print("Skipping individual component loading (loading from full pretrained model)") + self.beta_dist = Beta(config.noise_beta_alpha, config.noise_beta_beta) + # Video noise Beta distribution (biased towards high noise levels when enabled) + self.video_beta_dist = Beta(config.video_noise_beta_alpha, config.video_noise_beta_beta) + # High noise emphasis Beta distribution for coupled training (applies to both video and action) + self.high_noise_beta_dist = Beta(config.high_noise_beta_alpha, 1.0) + # self.num_timestep_buckets = config.num_timestep_buckets + self.config = config + self._noise_logged = False + self.defer_lora_injection = config.defer_lora_injection + print("defer_lora_injection@@", self.defer_lora_injection) + self.set_trainable_parameters(config.tune_projector, config.tune_diffusion_model) + + def set_trainable_parameters(self, tune_projector: bool, tune_diffusion_model: bool): + self.tune_projector = tune_projector + self.tune_diffusion_model = tune_diffusion_model + for p in self.parameters(): + p.requires_grad = True + if not tune_diffusion_model: + self.model.requires_grad_(False) + print(f"Tune action head projector: {self.tune_projector}") + print(f"Tune action head diffusion model: {self.tune_diffusion_model}") + # Check if any parameters are still trainable. If not, print a warning. + if not tune_projector and not tune_diffusion_model: + for name, p in self.named_parameters(): + if p.requires_grad: + print(f"Action head trainable parameter: {name}") + if not any(p.requires_grad for p in self.parameters()): + print("Warning: No action head trainable parameters found.") + + if self.train_architecture == "lora" and not self.defer_lora_injection: + print("Adding LoRA to model") + for p in self.parameters(): + p.requires_grad = False + self.model = self.add_lora_to_model( + self.model, + lora_rank=self.lora_rank, + lora_alpha=self.lora_alpha, + lora_target_modules=self.lora_target_modules, + init_lora_weights=self.init_lora_weights, + ) + self.model.state_encoder.requires_grad_(True) + self.model.action_encoder.requires_grad_(True) + self.model.action_decoder.requires_grad_(True) + elif self.train_architecture == "lora" and self.defer_lora_injection: + print("Deferring LoRA injection until after pretrained weights are loaded") + else: + self.print_trainable_params() + + self.text_encoder.requires_grad_(False) + self.image_encoder.requires_grad_(False) + self.vae.requires_grad_(False) + if not self.defer_lora_injection: + self.print_trainable_params() + + + def print_trainable_params(self): + """Print trainable parameters of the diffusion model.""" + trainable_params = [] + total_params = 0 + trainable_total = 0 + + for name, param in self.model.named_parameters(): + total_params += param.numel() + if param.requires_grad: + trainable_params.append(name) + trainable_total += param.numel() + + print(f"Total parameters in diffusion model: {total_params:,}") + print(f"Trainable parameters in diffusion model: {trainable_total:,}") + # print(trainable_params) + + + def inject_lora_after_loading(self): + """ + Inject LoRA adapters after pretrained weights have been loaded. + This should be called when defer_lora_injection=True. + """ + if self.train_architecture == "lora": + print("Injecting LoRA after loading pretrained weights") + for p in self.parameters(): + p.requires_grad = False + self.model = self.add_lora_to_model( + self.model, + lora_rank=self.lora_rank, + lora_alpha=self.lora_alpha, + lora_target_modules=self.lora_target_modules, + init_lora_weights=self.init_lora_weights, + ) + self.model.state_encoder.requires_grad_(True) + self.model.action_encoder.requires_grad_(True) + self.model.action_decoder.requires_grad_(True) + # self.model.registers.requires_grad_(True) + # self.model.time_modality_projection.requires_grad_(True) + + self.text_encoder.requires_grad_(False) + self.image_encoder.requires_grad_(False) + self.vae.requires_grad_(False) + self.print_trainable_params() + else: + print("LoRA injection not needed (train_architecture != 'lora')") + + def set_frozen_modules_to_eval_mode(self): + """ + Huggingface will call model.train() at each training_step. To ensure + the expected behaviors for modules like dropout, batchnorm, etc., we + need to call model.eval() for the frozen modules. + """ + if self.training: + if not self.tune_diffusion_model: + self.model.eval() + self.text_encoder.eval() + self.image_encoder.eval() + self.vae.eval() + + + def enable_vram_management(self, num_persistent_param_in_dit=None): + dtype = next(iter(self.text_encoder.parameters())).dtype + enable_vram_management( + self.text_encoder, + module_map = { + torch.nn.Linear: AutoWrappedLinear, + torch.nn.Embedding: AutoWrappedModule, + T5RelativeEmbedding: AutoWrappedModule, + T5LayerNorm: AutoWrappedModule, + }, + module_config = dict( + offload_dtype=dtype, + offload_device="cpu", + onload_dtype=dtype, + onload_device="cpu", + computation_dtype=self.dtype, + computation_device='cuda', + ), + ) + + self.cpu_offload = True + + def load_models_to_device(self, loadmodel_names=[]): + # only load models to device if cpu_offload is enabled + if not self.cpu_offload: + return + # offload the unneeded models to cpu + for model_name in self.model_names: + if model_name not in loadmodel_names: + model = getattr(self, model_name) + if model is not None: + if hasattr(model, "vram_management_enabled") and model.vram_management_enabled: + print("offloadd") + for module in model.modules(): + if hasattr(module, "offload"): + # print("offload", module) + module.offload() + else: + print("tocpu") + model.cpu() + # load the needed models to device + for model_name in loadmodel_names: + model = getattr(self, model_name) + if model is not None: + if hasattr(model, "vram_management_enabled") and model.vram_management_enabled: + print("onload") + for module in model.modules(): + if hasattr(module, "onload"): + # print("onload", module) + module.onload() + else: + print("togpu") + model.to(self._device) + # fresh the cuda cache + torch.cuda.empty_cache() + + def _create_kv_caches( + self, + batch_size: int, + dtype: torch.dtype, + device: torch.device, + frame_seqlen: int, + ) -> tuple[KVCacheType, KVCacheType]: + """ + Initialize a Per-GPU KV cache for the Wan model. + Use the model's num_heads and head_dim (5B has 24 heads, 14B has 40). + """ + num_heads = self.model.num_heads + head_dim = self.model.dim // num_heads + kv_cache1: KVCacheType = [] + kv_cache_neg: KVCacheType = [] + for _ in range(self.model.num_layers): + kv_cache1.append( + torch.zeros([2, batch_size, 0, num_heads, head_dim], dtype=dtype, device=device), + ) + kv_cache_neg.append( + torch.zeros([2, batch_size, 0, num_heads, head_dim], dtype=dtype, device=device), + ) + + return kv_cache1, kv_cache_neg + + def _create_crossattn_caches( + self, batch_size: int, dtype: torch.dtype, device: torch.device, + ) -> tuple[KVCacheType, KVCacheType]: + """ + Initialize a Per-GPU cross-attention cache for the Wan model. + Use the model's num_heads and head_dim (5B has 24 heads, 14B has 40). + """ + num_heads = self.model.num_heads + head_dim = self.model.dim // num_heads + crossattn_cache: KVCacheType = [] + crossattn_cache_neg: KVCacheType = [] + + for _ in range(self.model.num_layers): + crossattn_cache.append( + torch.zeros([2, batch_size, 512, num_heads, head_dim], dtype=dtype, device=device), + ) + crossattn_cache_neg.append( + torch.zeros([2, batch_size, 512, num_heads, head_dim], dtype=dtype, device=device), + ) + + return crossattn_cache, crossattn_cache_neg + + def sample_time(self, batch_size, device, dtype): + sample = self.beta_dist.sample([batch_size]).to(device, dtype=dtype) + return (self.config.noise_s - sample) / self.config.noise_s + + def prepare_input(self, batch: dict) -> BatchFeature: + return BatchFeature(data=batch) + + def preprocess_image(self, image): + image = (image * (2 / 255) - 1).permute(0, 1, 4, 2, 3) + return image + + def encode_prompt(self, input_ids, attention_mask): + seq_lens = attention_mask.gt(0).sum(dim=1).long() + prompt_emb = self.text_encoder(input_ids, attention_mask) + prompt_emb = prompt_emb.clone().to(dtype=torch.bfloat16) + for i, v in enumerate(seq_lens): + prompt_emb[:, v:] = 0 + return prompt_emb + + def _ensure_vae_on_device(self, ref_tensor): + """Lazily move the VAE to the correct device/dtype on first use.""" + if not getattr(self, '_vae_device_ready', False): + self.vae.to(device=ref_tensor.device, dtype=torch.bfloat16) + self.vae.eval() + self._vae_device_ready = True + + def encode_video(self, input_video, tiled=True, tile_size=(34, 34), tile_stride=(18, 16)): + self._ensure_vae_on_device(input_video) + with torch.no_grad(): + latents = self.vae.encode(input_video, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride) + return latents + + def encode_image(self, image, num_frames, height, width): + with torch.amp.autocast(dtype=torch.bfloat16, device_type=torch.device(self._device).type): + batch_size = image.shape[0] + clip_context = self.image_encoder.encode_image(image) + image_input = image.transpose(1, 2) + image_zeros = torch.zeros(batch_size, 3, num_frames-1, height, width, dtype=torch.bfloat16, device=self._device) + self._ensure_vae_on_device(image_input) + with torch.no_grad(): + y = self.vae.encode(torch.concat([image_input, image_zeros], dim=2)) + # Build mask to match VAE output shape (VAE may use different spatial downsampling, e.g. WanVideoVAE38 uses patch_size=2 -> height/16) + # y shape is B * 16 * (1+(T-1)/4) * H_latent * W_latent + num_t = y.shape[2] + h_latent, w_latent = y.shape[3], y.shape[4] + msk = torch.zeros(batch_size, 4, num_t, h_latent, w_latent, dtype=y.dtype, device=self._device) + msk[:, :, 0:1, :, :] = 1 + new_image = y[:, :, 0:1] + # concat: B * (4+16) * (1+(T-1)/4) * H_latent * W_latent + y = torch.concat([msk, y], dim=1) + return clip_context, y, new_image + + def prepare_extra_input(self, latents=None): + return {} + + def add_lora_to_model(self, model, lora_rank=4, lora_alpha=4, lora_target_modules="q,k,v,o,ffn.0,ffn.2", init_lora_weights="kaiming") -> nn.Module: + # Add LoRA to UNet + self.lora_alpha = lora_alpha + if init_lora_weights == "kaiming": + init_lora_weights = True + + lora_config = LoraConfig( + r=lora_rank, + lora_alpha=lora_alpha, + init_lora_weights=init_lora_weights, + target_modules=lora_target_modules.split(","), + ) + model = get_peft_model(model, lora_config) + for param in model.parameters(): + param.data = param.to(torch.float32) + return model + + def forward(self, backbone_output: BatchFeature, action_input: BatchFeature) -> BatchFeature: + # Set frozen modules to eval + self.set_frozen_modules_to_eval_mode() + + data = action_input + # Get embodiment ID. + embodiment_id = action_input.embodiment_id + # print("embodiment_id", embodiment_id) + has_real_action = action_input.has_real_action + action_mask = action_input.action_mask + + state_features = action_input.state + + actions = action_input.action + # assert the values of action is in between -1 and 1 + if actions.numel() > 0: + assert actions.min() >= -1.0 and actions.max() <= 1.0, "actions must be in [-1,1] range" + videos = data["images"] + + videos = rearrange(videos, "b t h w c -> b c t h w") + + if videos.dtype == torch.uint8: + videos = videos.float() / 255.0 + b, c, t, h, w = videos.shape + videos = videos.permute(0, 2, 1, 3, 4) # [b, t, c, h, w] + videos = videos.reshape(b * t, c, h, w) + videos = self.normalize_video(videos) + videos = videos.reshape(b, t, c, h, w).permute(0, 2, 1, 3, 4) # back to [b, c, t, h, w] + assert videos.min() >= -1.0 and videos.max() <= 1.0, "videos must be in [-1,1] range" + # Match reference precision: VAE encoding runs in fp32 (reference has all params + # in fp32 due to add_lora_to_model casting, so self.dtype=fp32 there). + # Under FSDP2, self.dtype=bf16, but we force fp32 here for alignment. + videos = videos.to(dtype=torch.float32) + + # shape of B * max_length * dim + prompt_embs = self.encode_prompt(data["text"], data["text_attention_mask"]) + + # Wan 5B: resize to target resolution so latent tokens/frame matches DiT. Use config target when set + # (e.g. 160x320 so latent is 10x20 with VAE38 16x → even H,W, no crop in dynamics loss); else 176x320. + target_h = getattr(self.config, "target_video_height", None) + target_w = getattr(self.config, "target_video_width", None) + if target_h is None or target_w is None: + if getattr(self.model, "frame_seqlen", None) in (50, 55): + target_h, target_w = 176, 320 + else: + target_h, target_w = None, None + if target_h is not None and target_w is not None: + _, _, _, h, w = videos.shape + if (h, w) != (target_h, target_w): + b, c, t, _, _ = videos.shape + videos = torch.nn.functional.interpolate( + videos.reshape(b * t, c, h, w), + size=(target_h, target_w), + mode="bilinear", + align_corners=False, + ).reshape(b, c, t, target_h, target_w) + + latents = self.encode_video(videos, self.tiled, (self.tile_size_height, self.tile_size_width), (self.tile_stride_height, self.tile_stride_width)) + + # print("latents shape", latents.shape, self.dtype) + _, _, num_frames, height, width = videos.shape + image = videos[:, :, :1].transpose(1, 2) + + clip_feas, ys, _ = self.encode_image(image, num_frames, height, width) + + latents = latents.to(self._device) + clip_feas = clip_feas.to(self._device) + ys = ys.to(self._device) + prompt_embs = prompt_embs.to(self._device) + + # AGENT_DEBUG: Fix random state for Level 4 alignment comparison + import os as _os_align # AGENT_DEBUG + if _os_align.environ.get("DREAMZERO_ALIGN_SEED"): # AGENT_DEBUG + _align_seed = int(_os_align.environ["DREAMZERO_ALIGN_SEED"]) # AGENT_DEBUG + torch.manual_seed(_align_seed) # AGENT_DEBUG + torch.cuda.manual_seed(_align_seed) # AGENT_DEBUG + # END AGENT_DEBUG + + # Loss + noise = torch.randn_like(latents) + + # specific to autoregressive + noise = noise.transpose(1, 2) + latents = latents.transpose(1, 2) + + # ============ VIDEO TIMESTEP SAMPLING ============ + if self.config.decouple_video_action_noise: + # Decoupled mode: sample video from Beta distribution biased towards HIGH noise + video_noise_ratio = self.video_beta_dist.sample([noise.shape[0], noise.shape[1]]) + timestep_id = ((1.0 - video_noise_ratio) * self.scheduler.num_train_timesteps).long() + timestep_id = torch.clamp(timestep_id, 0, self.scheduler.num_train_timesteps - 1) + noise_mode = "DECOUPLED" + elif self.config.use_high_noise_emphasis: + # High noise emphasis mode (coupled): BOTH video and action use Beta distribution + noise_ratio = self.high_noise_beta_dist.sample([noise.shape[0], noise.shape[1]]) + timestep_id = ((1.0 - noise_ratio) * self.scheduler.num_train_timesteps).long() + timestep_id = torch.clamp(timestep_id, 0, self.scheduler.num_train_timesteps - 1) + noise_mode = "HIGH_NOISE_EMPHASIS" + else: + # Original: uniform sampling over full range + timestep_id = torch.randint(0, self.scheduler.num_train_timesteps, (noise.shape[0], noise.shape[1])) + noise_mode = "STANDARD" + + timestep_id_block = timestep_id[:, 1:].reshape( + timestep_id.shape[0], -1, self.num_frame_per_block) + timestep_id_block[:, :, 1:] = timestep_id_block[:, :, 0:1] + + if actions.numel() > 0: + noise_action = torch.randn_like(actions) + # NOTE: These assertions are invalid for the DreamZero-AgiBot checkpoint config + # (num_frames=33, num_frame_per_block=2, num_action_per_block=48, action_horizon=48) + # 48/8=6 != 48//2=24. Commented out for training. + # assert actions.shape[1] / (noise.shape[1]-1) == (self.model.num_action_per_block // self.num_frame_per_block), f"actions.shape, {actions.shape}, noise.shape, {noise.shape}, video.shape, {videos.shape}, latents.shape, {latents.shape}" + # assert (noise.shape[1]-1) / state_features.shape[1] == (self.num_frame_per_block // self.model.num_state_per_block), f"state_features.shape, {state_features.shape}, noise.shape, {noise.shape}, video.shape, {videos.shape}, latents.shape, {latents.shape}" + + # ============ ACTION TIMESTEP SAMPLING ============ + if self.config.decouple_video_action_noise: + # Decoupled: sample action timestep independently with full range + timestep_action_id = torch.randint( + 0, + self.scheduler.num_train_timesteps, + (actions.shape[0], actions.shape[1]) + ) + action_mode = "INDEPENDENT" + else: + # Original coupled: action timestep derived from video timestep + timestep_action_id = timestep_id_block.repeat(1, 1, actions.shape[1]//(noise.shape[1]-1)) + timestep_action_id = timestep_action_id.reshape(timestep_action_id.shape[0], -1) + action_mode = "COUPLED" + + # Log noise mode once + if not self._noise_logged: + video_mean = timestep_id.float().mean().item() + action_mean = timestep_action_id.float().mean().item() + if noise_mode == "DECOUPLED": + print(f"[NOISE] Mode={noise_mode} | Video: Beta({self.config.video_noise_beta_alpha},1) mean_t={video_mean:.0f} | Action: {action_mode} Uniform mean_t={action_mean:.0f}") + elif noise_mode == "HIGH_NOISE_EMPHASIS": + print(f"[NOISE] Mode={noise_mode} | Video+Action: Beta({self.config.high_noise_beta_alpha},1) mean_t={video_mean:.0f} | Action: {action_mode}") + else: + print(f"[NOISE] Mode={noise_mode} | Video+Action: Uniform mean_t={video_mean:.0f} | Action: {action_mode}") + self._noise_logged = True + else: + noise_action = None + timestep_action_id = None + + timestep_id_block = timestep_id_block.reshape(timestep_id_block.shape[0], -1) + timestep_id = torch.concat([timestep_id[:, :1], timestep_id_block], dim=1) + _, num_frames, num_channels, height, width = noise.shape + # DiT patch_embedding uses stride (1,2,2), so sequence length is num_frames * (H//2) * (W//2) + tokens_per_frame = (height // 2) * (width // 2) + seq_len = num_frames * tokens_per_frame + + timestep = self.scheduler.timesteps[timestep_id].to(self._device) + noisy_latents = self.scheduler.add_noise(latents.flatten(0, 1), noise.flatten(0, 1), timestep.flatten(0, 1)).unflatten(0, (noise.shape[0], noise.shape[1])) + training_target = self.scheduler.training_target(latents, noise, timestep).transpose(1, 2) + + if actions.numel() > 0: + timestep_action = self.scheduler.timesteps[timestep_action_id].to(self._device) + noisy_actions = self.scheduler.add_noise( + actions.flatten(0, 1), + noise_action.flatten(0, 1), + timestep_action.flatten(0, 1), + ).unflatten(0, (noise_action.shape[0], noise_action.shape[1])) + training_target_action = self.scheduler.training_target(actions, noise_action, timestep_action) + else: + timestep_action = None + noisy_actions = None + training_target_action = None + + # Compute loss + with torch.amp.autocast(dtype=torch.bfloat16, device_type=torch.device(self._device).type): + if actions.numel() > 0: + video_noise_pred, action_noise_pred = self.model( + noisy_latents.transpose(1, 2), timestep=timestep, clip_feature=clip_feas, y=ys, context=prompt_embs, seq_len=seq_len, + state=state_features, embodiment_id=embodiment_id, + action=noisy_actions, timestep_action=timestep_action, + clean_x=latents.transpose(1, 2), + ) + else: + video_noise_pred, action_noise_pred = self.model( + noisy_latents.transpose(1, 2), timestep=timestep, timestep_action=timestep_action, + clip_feature=clip_feas, y=ys, context=prompt_embs, seq_len=seq_len, + state=state_features, embodiment_id=embodiment_id, + clean_x=latents.transpose(1, 2), + ) + + # Per-sample dynamics loss + # DiT patch_embedding uses stride (1,2,2), so output spatial size can be smaller than + # latent when H or W is odd (e.g. latent 11x20 -> model output 10x20). Crop target to match. + if training_target.shape != video_noise_pred.shape: + training_target = training_target[ + ..., : video_noise_pred.shape[3], : video_noise_pred.shape[4] + ] + dynamics_loss_per_sample = torch.nn.functional.mse_loss( + video_noise_pred.float(), training_target.float(), reduction='none' + ).mean(dim=(1,3,4)) # shape: [B, ...] + + weight_dynamics = dynamics_loss_per_sample * self.scheduler.training_weight(timestep.flatten(0, 1)).unflatten(0, (noise.shape[0], noise.shape[1])).to(self._device) + weighted_dynamics_loss = weight_dynamics.mean() + + if actions.numel() > 0: + action_loss_per_sample = torch.nn.functional.mse_loss( + action_noise_pred.float(), training_target_action.float(), reduction='none' + ) * action_mask # shape: [B, ...] + action_loss_per_sample = has_real_action[:, None, None].float() * action_loss_per_sample # apply has_real_action + weight_action = action_loss_per_sample.mean(dim=2) * self.scheduler.training_weight( + timestep_action.flatten(0, 1), + ).unflatten(0, (noise_action.shape[0], noise_action.shape[1])).to(self._device) + weighted_action_loss = weight_action.mean() + loss = weighted_dynamics_loss + weighted_action_loss + else: + weighted_action_loss = torch.tensor(0.0, device=self._device) + loss = weighted_dynamics_loss + # loss = dynamics_loss_per_sample.mean() + + # Record log + output_dict = { + "loss": loss, + "dynamics_loss": weighted_dynamics_loss, + "action_loss": weighted_action_loss, + } + + return BatchFeature(data=output_dict) + + def generate_noise(self, shape, seed=None, device="cpu", dtype=torch.float16): + generator = None if seed is None else torch.Generator(device).manual_seed(seed) + noise = torch.randn(shape, generator=generator, device=device, dtype=dtype) + return noise + + def _get_caches( + self, kv_caches_input: list[KVCacheType], + ) -> list[KVCacheType]: + if self.ip_size > 1: + assert self.cfg_scale != 1.0, "cfg_scale must be != 1.0 when ip_size > 1" + assert len(kv_caches_input) == 2 + if self.ip_rank == 0: + kv_caches = [kv_caches_input[0]] + else: + kv_caches = [kv_caches_input[1]] + else: + assert len(kv_caches_input) <= 2 + kv_caches = [kv_caches_input[0]] + if self.cfg_scale != 1.0: + kv_caches.append(kv_caches_input[1]) + return kv_caches + + def _prepare_text_inputs(self, data: BatchFeature) -> list[tuple[torch.Tensor, torch.Tensor]]: + + if self.ip_size > 1: + assert self.cfg_scale != 1.0, "cfg_scale must be != 1.0 when ip_size > 1" + if self.ip_rank == 0: + text_inputs = [(data["text"], data["text_attention_mask"])] + else: + text_inputs = [(data["text_negative"], data["text_attention_mask_negative"])] + else: + text_inputs = [(data["text"], data["text_attention_mask"])] + if self.cfg_scale != 1.0: + text_inputs.append((data["text_negative"], data["text_attention_mask_negative"])) + return text_inputs + + + def _run_diffusion_steps( + self, + noisy_input: torch.Tensor, + timestep: torch.Tensor, + action: torch.Tensor, + timestep_action: torch.Tensor, + state: torch.Tensor, + embodiment_id: torch.Tensor, + context: torch.Tensor, + seq_len: int, + y: torch.Tensor, + clip_feature: torch.Tensor, + kv_caches: list[KVCacheType], + crossattn_caches: list[KVCacheType], + kv_cache_metadata: dict[str, bool | int], + ) -> list[tuple[torch.Tensor, torch.Tensor]]: + predictions = [] + for index, prompt_emb in enumerate(context): + kv_cache = kv_caches[index] + crossattn_cache = crossattn_caches[index] + if not kv_cache_metadata["update_kv_cache"] and self.trt_engine is not None: + obs_noise_pred, action_noise_pred = self.trt_engine( + noisy_input, + timestep, + action=action, + timestep_action=timestep_action, + state=state, + context=prompt_emb, + y=y, + clip_feature=clip_feature, + kv_cache=kv_cache, + ) + else: + obs_noise_pred, action_noise_pred, updated_kv_caches = self.model( + noisy_input, + timestep, + action=action, + timestep_action=timestep_action, + state=state, + embodiment_id=embodiment_id, + context=prompt_emb, + seq_len=seq_len, + y=y, + clip_feature=clip_feature, + kv_cache=kv_cache, + crossattn_cache=crossattn_cache, + current_start_frame=kv_cache_metadata["start_frame"], + ) + if kv_cache_metadata["update_kv_cache"]: + for block_index, updated_kv_cache in enumerate(updated_kv_caches): + kv_cache[block_index] = updated_kv_cache.clone() + obs_noise_pred = obs_noise_pred.clone() + if action_noise_pred is not None: + action_noise_pred = action_noise_pred.clone() + else: + action_noise_pred = torch.tensor(0.0, device=obs_noise_pred.device) # dummy action noise prediction + predictions.append((obs_noise_pred, action_noise_pred)) + return self._exchange_predictions(predictions) + + def _exchange_predictions( + self, + predictions: list[tuple[torch.Tensor, torch.Tensor]], + ) -> list[tuple[torch.Tensor, torch.Tensor]]: + if self.ip_size == 1: + return predictions + + assert len(predictions) == 1 + my_predictions = list(predictions[0]) + + other_predictions = [torch.empty_like(pred) for pred in my_predictions] + + send_ops = [ + dist.P2POp(op=dist.isend, tensor=pred, group_peer=(self.ip_rank + 1) % self.ip_size, group=self.ip_group) + for pred in my_predictions + ] + recv_ops = [ + dist.P2POp(op=dist.irecv, tensor=other_pred, group_peer=(self.ip_rank + 1) % self.ip_size, group=self.ip_group) + for other_pred in other_predictions + ] + ops = send_ops + recv_ops + + reqs = dist.batch_isend_irecv(ops) + for req in reqs: + req.wait() + + output_predictions: list[tuple[torch.Tensor, torch.Tensor] | None] = [None for _ in range(self.ip_size)] + output_predictions[self.ip_rank] = tuple(my_predictions) + output_predictions[(self.ip_rank + 1) % self.ip_size] = tuple(other_predictions) + assert all(isinstance(pred, tuple) for pred in output_predictions) + return cast(list[tuple[torch.Tensor, torch.Tensor]], output_predictions) + + def should_run_model(self, index, current_timestep, prev_predictions): + + if not self.dynamic_cache_schedule: + return self.dit_step_mask[index] + + # Always run first 2 steps to establish history + if len(prev_predictions) < 2: + return True + + if self.skip_countdown > 1: + self.skip_countdown -= 1 + return False + elif self.skip_countdown == 1: + self.skip_countdown = 0 + return True + + v_last = prev_predictions[-1][1].flatten(1).float() + v_prev = prev_predictions[-2][1].flatten(1).float() + sim = torch.nn.functional.cosine_similarity(v_last, v_prev, dim=1).mean() + + thresholds = [0.95, 0.93] + countdowns = [4, 2] + + for threshold, countdown in zip(thresholds, countdowns): + if sim > threshold: + self.skip_countdown = countdown + return False + + return True + + def lazy_joint_video_action(self, backbone_output: BatchFeature, action_input: BatchFeature, latent_video: torch.Tensor | None = None) -> BatchFeature: + start_time = time.perf_counter() + + # Tracking time taken on GPU for various operations. + start_text_encoder_event = torch.cuda.Event(enable_timing=True) + end_text_encoder_event = torch.cuda.Event(enable_timing=True) + start_image_encoder_event = torch.cuda.Event(enable_timing=True) + end_image_encoder_event = torch.cuda.Event(enable_timing=True) + start_vae_event = torch.cuda.Event(enable_timing=True) + end_vae_event = torch.cuda.Event(enable_timing=True) + start_kv_event = torch.cuda.Event(enable_timing=True) + end_kv_event = torch.cuda.Event(enable_timing=True) + start_diffusion_events = [torch.cuda.Event(enable_timing=True) for _ in range(self.num_inference_steps)] + end_diffusion_events = [torch.cuda.Event(enable_timing=True) for _ in range(self.num_inference_steps)] + + self.set_frozen_modules_to_eval_mode() + data = action_input + + videos = data["images"] + + embodiment_id = action_input.embodiment_id + state_features = action_input.state + + videos = rearrange(videos, "b t h w c -> b c t h w") + + if videos.dtype == torch.uint8: + videos = videos.float() / 255.0 + videos = videos.to(dtype=self.dtype) + b, c, t, h, w = videos.shape + videos = videos.permute(0, 2, 1, 3, 4) # [b, t, c, h, w] + videos = videos.reshape(b * t, c, h, w) + videos = self.normalize_video(videos) + videos = videos.reshape(b, t, c, h, w).permute(0, 2, 1, 3, 4) # back to [b, c, t, h, w] + assert videos.min() >= -1.0 and videos.max() <= 1.0, "videos must be in [-1,1] range" + videos = videos.to(dtype=self.dtype) + + state_features = state_features.to(dtype=torch.bfloat16) + videos = videos.to(dtype=torch.bfloat16) + + # Wan 5B: same as training — resize to target resolution so latent matches DiT + target_h = getattr(self.config, "target_video_height", None) + target_w = getattr(self.config, "target_video_width", None) + if target_h is None or target_w is None: + if getattr(self.model, "frame_seqlen", None) in (50, 55): + target_h, target_w = 176, 320 + else: + target_h, target_w = None, None + if target_h is not None and target_w is not None: + _, _, _, h, w = videos.shape + if (h, w) != (target_h, target_w): + b, c, t, _, _ = videos.shape + videos = torch.nn.functional.interpolate( + videos.reshape(b * t, c, h, w), + size=(target_h, target_w), + mode="bilinear", + align_corners=False, + ).reshape(b, c, t, target_h, target_w) + + if self.language is None: + print("language is None, reset current_start_frame to 0") + self.language = data["text"] + self.current_start_frame = 0 + elif not torch.equal(self.language, data["text"]): + print("language changed, reset current_start_frame to 0") + self.current_start_frame = 0 + self.language = data["text"] + elif videos.shape[2] == 1: + print("videos.shape[2] == 1, reset current_start_frame to 0") + self.current_start_frame = 0 + elif self.current_start_frame >= self.model.local_attn_size: + print("current_start_frame >= local_attn_size, reset current_start_frame to 0") + self.current_start_frame = 0 + + if self.ip_rank == 0: + print("videos shape", videos.shape, self.num_frames) + + start_text_encoder_event.record() + + text_inputs = self._prepare_text_inputs(data) + prompt_embs = [self.encode_prompt(text, attention_mask) for text, attention_mask in text_inputs] + + end_text_encoder_event.record() + + start_image_encoder_event.record() + + _, _, num_frames, height, width = videos.shape + if videos.shape[2] == 4 or videos.shape[2] == 9: + # special case for real-world eval where language is updated + image = videos[:, :, -1:].transpose(1, 2) + else: + image = videos[:, :, :1].transpose(1, 2) + + if self.current_start_frame == 0: + clip_feas, ys, image = self.encode_image(image, self.num_frames, height, width) + self.clip_feas = clip_feas.to(dtype=image.dtype) + self.ys = ys.to(dtype=image.dtype) + + assert self.clip_feas is not None and self.ys is not None, "clip_feas and ys must be set" + + end_image_encoder_event.record() + + start_vae_event.record() + + if latent_video is not None and self.current_start_frame != 0: + image = latent_video + if self.ip_rank == 0: + print("image shape@@", image.shape) + elif self.current_start_frame != 0: + # this is for real world execution + if (videos.shape[2] - 1) // 4 == self.num_frame_per_block: + print("no further action") + elif videos.shape[2] // 4 != self.num_frame_per_block: + # Repeating videos along dim 2. + repeat_factor = self.num_frame_per_block // (videos.shape[2] // 4) + videos = torch.repeat_interleave(videos, repeat_factor, dim=2) + + first_frame = videos[:, :, 0:1] # Extract first frame + videos = torch.cat([first_frame, videos], dim=2) + else: + first_frame = videos[:, :, 0:1] # Extract first frame + videos = torch.cat([first_frame, videos], dim=2) + + image = self.vae.encode( + videos, + tiled=self.tiled, + tile_size=(self.tile_size_height, self.tile_size_width), + tile_stride=(self.tile_stride_height, self.tile_stride_width), + ) + + end_vae_event.record() + + noise_obs = self.generate_noise((image.shape[0], image.shape[1], self.num_frame_per_block, image.shape[3], image.shape[4]), seed=self.seed, device='cuda', dtype=torch.bfloat16) + noise_action = self.generate_noise((image.shape[0], self.action_horizon, self.model.action_dim), seed=self.seed, device='cuda', dtype=torch.bfloat16) + batch_size, num_channels, num_frames, height, width = noise_obs.shape + ######### Generate video ######### + # DiT patch_embedding uses stride (1,2,2), so tokens per frame = (H//2)*(W//2) + tokens_per_frame = (height // 2) * (width // 2) + frame_seqlen = tokens_per_frame + seq_len = num_frames * frame_seqlen + + image = image.transpose(1, 2) + noise_obs = noise_obs.transpose(1, 2) + + if self.current_start_frame == 0: + # Reinitialize KV cache and crossattn cache for each new sequence. + self.kv_cache1, self.kv_cache_neg = self._create_kv_caches( + batch_size=batch_size, + dtype=noise_obs.dtype, + device=noise_obs.device, + frame_seqlen=frame_seqlen, + ) + self.crossattn_cache, self.crossattn_cache_neg = self._create_crossattn_caches( + batch_size=batch_size, + dtype=noise_obs.dtype, + device=noise_obs.device, + ) + + assert self.kv_cache1 is not None + assert self.kv_cache_neg is not None + assert self.crossattn_cache is not None + assert self.crossattn_cache_neg is not None + kv_caches = self._get_caches( + [self.kv_cache1, self.kv_cache_neg], + ) + crossattn_caches = self._get_caches( + [self.crossattn_cache, self.crossattn_cache_neg], + ) + + start_kv_event.record() + + if self.current_start_frame == 0: + timestep = torch.ones([batch_size, 1], device=noise_obs.device, dtype=torch.int64) * 0 + self._run_diffusion_steps( + noisy_input=image.transpose(1, 2), + timestep=timestep * 0, + action=None, + timestep_action=None, + state=None, + embodiment_id=None, + context=prompt_embs, + seq_len=frame_seqlen, + y=self.ys[:, :, 0:1], + clip_feature=self.clip_feas, + kv_caches=kv_caches, + crossattn_caches=crossattn_caches, + kv_cache_metadata=dict( + start_frame=0, + update_kv_cache=True, + ), + ) + self.current_start_frame += 1 + + timestep = torch.ones([batch_size, self.num_frame_per_block], device=noise_obs.device, dtype=torch.int64) * 0 + + if self.current_start_frame != 1: + current_ref_latents = image[:, -self.num_frame_per_block:] + if self.current_start_frame <= self.ys.shape[2]: + y = self.ys[:, :, self.current_start_frame - self.num_frame_per_block : self.current_start_frame] + else: + y = self.ys[:, :, -self.num_frame_per_block:] + self._run_diffusion_steps( + noisy_input=current_ref_latents.transpose(1, 2), + timestep=timestep * 0, + action=None, + timestep_action=None, + state=None, + embodiment_id=None, + context=prompt_embs, + seq_len=seq_len, + y=y, + clip_feature=self.clip_feas, + kv_caches=kv_caches, + crossattn_caches=crossattn_caches, + kv_cache_metadata=dict( + start_frame=self.current_start_frame - self.num_frame_per_block, + update_kv_cache=True, + ), + ) + + end_kv_event.record() + + noisy_input = noise_obs + noisy_input_action = noise_action + + # Step 3.1: Spatial denoising loop + + sample_scheduler = FlowUniPCMultistepScheduler( + num_train_timesteps=self.scheduler.num_train_timesteps, + shift=1, + use_dynamic_shifting=False) + sample_scheduler_action = FlowUniPCMultistepScheduler( + num_train_timesteps=self.scheduler.num_train_timesteps, + shift=1, + use_dynamic_shifting=False) + sample_scheduler.set_timesteps( + self.num_inference_steps, device=noise_obs.device, shift=self.sigma_shift) + sample_scheduler_action.set_timesteps( + self.num_inference_steps, device=noise_obs.device, shift=self.sigma_shift) + + # Decoupled inference: video sigmas end at video_final_noise instead of 0 + # This rescales the schedule so video still takes all denoising steps, + # but ends at a higher noise level (e.g., 1.0 → 0.9 → 0.8 instead of 1.0 → 0.5 → 0.0) + if self.config.decouple_inference_noise: + video_final_noise = self.config.video_inference_final_noise + # Rescale video sigmas: map [sigma_max, 0] -> [sigma_max, video_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) + if self.ip_rank == 0: + print(f"Decoupled inference: video sigmas {sigma_max:.3f} -> {sample_scheduler.sigmas[-1].item():.3f}") + + start_diffusion_events = [torch.cuda.Event(enable_timing=True) for _ in sample_scheduler.timesteps] + end_diffusion_events = [torch.cuda.Event(enable_timing=True) for _ in sample_scheduler.timesteps] + prev_predictions = [] + self.skip_countdown = 0 + dit_compute_steps = 0 + for index, current_timestep in enumerate(sample_scheduler.timesteps): + start_diffusion_events[index].record() + + # Get timesteps from respective schedulers + action_timestep = sample_scheduler_action.timesteps[index] + video_timestep = sample_scheduler.timesteps[index] # Already rescaled if decoupled + + # set current timestep + timestep = torch.ones( + [batch_size, self.num_frame_per_block], + device=noise_obs.device, + dtype=torch.int64, + ) * video_timestep + timestep_action = torch.ones( + [batch_size, self.action_horizon], + device=noise_obs.device, + dtype=torch.int64, + ) * action_timestep + + # check if we need to run the DIT step + should_run_model = self.should_run_model(index, current_timestep, prev_predictions) + if should_run_model: + dit_compute_steps += 1 + if self.current_start_frame + self.num_frame_per_block <= self.ys.shape[2]: + y = self.ys[:, :, self.current_start_frame : self.current_start_frame + self.num_frame_per_block] + else: + y = self.ys[:, :, -self.num_frame_per_block:] + predictions = self._run_diffusion_steps( + noisy_input=noisy_input.transpose(1, 2), + timestep=timestep, + action=noisy_input_action, + timestep_action=timestep_action, + state=state_features, + embodiment_id=embodiment_id, + context=prompt_embs, + seq_len=seq_len, + y=y, + clip_feature=self.clip_feas, + kv_caches=kv_caches, + crossattn_caches=crossattn_caches, + kv_cache_metadata=dict( + start_frame=self.current_start_frame, + update_kv_cache=False, + ), + ) + flow_pred_cond, flow_pred_cond_action = predictions[0] + flow_pred_uncond, flow_pred_uncond_action = predictions[1] + + flow_pred = flow_pred_uncond + self.cfg_scale * (flow_pred_cond - flow_pred_uncond) + prev_predictions.append((current_timestep, flow_pred, flow_pred_cond_action)) + max_cache_size = 2 + if len(prev_predictions) > max_cache_size: + prev_predictions.pop(0) + + else: + assert len(prev_predictions) > 0, "prev_predictions must be set when skipping" + _, flow_pred, flow_pred_cond_action = prev_predictions[-1] + + end_diffusion_events[index].record() + + # Video: denoising step (uses rescaled schedule if decoupled) + noisy_input = sample_scheduler.step( + model_output=flow_pred.transpose(1, 2), + timestep=video_timestep, + sample=noisy_input, + step_index=index, + return_dict=False, + )[0] + + # Action: always fully denoises with standard schedule (1000->0) + noisy_input_action = sample_scheduler_action.step( + model_output=flow_pred_cond_action, + timestep=action_timestep, + sample=noisy_input_action, + step_index=index, + return_dict=False, + )[0] + + latents = noisy_input + latents_action = noisy_input_action + output = latents + + if self.current_start_frame == 1: + output = torch.cat([image, output], dim=1) + self.current_start_frame += self.num_frame_per_block + + # Do torch.cuda.synchronize() to ensure all operations are completed before timing. + # This isn't expected to affect inference performance since it's at the end of an inference step. + torch.cuda.synchronize() + + total_time = time.perf_counter() - start_time + text_encoder_time = start_text_encoder_event.elapsed_time(end_text_encoder_event) / 1000 + image_encoder_time = start_image_encoder_event.elapsed_time(end_image_encoder_event) / 1000 + vae_time = start_vae_event.elapsed_time(end_vae_event) / 1000 + kv_creation_time = start_kv_event.elapsed_time(end_kv_event) / 1000 + diffusion_times = [s.elapsed_time(e) for s, e in zip(start_diffusion_events, end_diffusion_events)] + diffusion_time = sum(diffusion_times) / 1000 + scheduler_time = total_time - kv_creation_time - diffusion_time - text_encoder_time - image_encoder_time - vae_time + + if self.ip_rank == 0: + print(f"Time taken: Total {total_time:.2f} seconds, " + f"Text Encoder {text_encoder_time:.2f} seconds, " + f"Image Encoder {image_encoder_time:.2f} seconds, " + f"VAE {vae_time:.2f} seconds, " + f"KV Cache Creation {kv_creation_time:.2f} seconds, " + f"Diffusion {diffusion_time:.2f} seconds, " + f"DIT Compute Steps {dit_compute_steps} steps, " + f"Scheduler {scheduler_time:.2f} seconds") + + return BatchFeature(data={"action_pred": latents_action, "video_pred": output.transpose(1, 2)}) + + def cache_predict_order1(self, current_timestep, timestep_1, f1, timestep_2, f2): + h_curr = current_timestep - timestep_1 + h_past = timestep_1 - timestep_2 + + v_prime = (f1 - f2) / h_past + + # Prediction + damping_factor = 0.25 + flow_pred = f1 + (v_prime * h_curr) * damping_factor + return flow_pred + + def post_initialize(self): + # Move models to the cuda device and set the dtype to bfloat16. + print("Moving models to the cuda device and setting the dtype to bfloat16.") + self.model.to(device=self._device, dtype=torch.bfloat16) + self.text_encoder.to(device=self._device, dtype=torch.bfloat16) + self.image_encoder.to(device=self._device, dtype=torch.bfloat16) + self.vae.to(device=self._device, dtype=torch.bfloat16) + import os + ENABLE_TENSORRT = os.getenv("ENABLE_TENSORRT", "False").lower() == "true" + LOAD_TRT_ENGINE = os.getenv("LOAD_TRT_ENGINE", None) + + # Torch compile the modules. Skip _forward_blocks: Dynamo with fullgraph can fail on + # shape variation (e.g. x [1,50,C] vs e [1,200,C]); the block aligns e to x at runtime. + if not ENABLE_TENSORRT: + print("Torch compiling the TextEncoder, ImageEncoder, and VAE modules (Wan _forward_blocks not compiled).") + + self.text_encoder.forward = torch.compile( + mode="reduce-overhead", fullgraph=True, dynamic=False, + )(self.text_encoder.forward) + + self.image_encoder.model.visual.forward = torch.compile( + mode="reduce-overhead", fullgraph=True, dynamic=False, + )(self.image_encoder.model.visual.forward) + + self.vae.model.encode = torch.compile( + mode="reduce-overhead", fullgraph=True, dynamic=False, + )(self.vae.model.encode) + + self.trt_engine = None + if LOAD_TRT_ENGINE is not None: + print(f"Loading TRT engine from {LOAD_TRT_ENGINE}") + import flagscale.train.models.dreamzero.modules.tensorrt_utils as trt_utils + model_path = LOAD_TRT_ENGINE + self.trt_engine = trt_utils.load_tensorrt_engine(model_path, model_type="ar_14B") + + def parallelize(self, device_mesh: DeviceMesh) -> None: + ip_mesh = device_mesh["ip"] + self.ip_rank = ip_mesh.get_local_rank() + self.ip_size = ip_mesh.size() + self.ip_group = ip_mesh.get_group() + + assert self.ip_size == 1 or self.ip_size == 2, "ip_size must be 1 or 2" + assert self.ip_rank >= 0 and self.ip_rank < self.ip_size, "ip_rank must be in [0, ip_size)" + + @property + def device(self): + return next(iter(self.parameters())).device + + @property + def dtype(self): + return next(iter(self.parameters())).dtype diff --git a/flagscale/train/models/dreamzero/backbone/__init__.py b/flagscale/train/models/dreamzero/backbone/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/flagscale/train/models/dreamzero/backbone/base_backbone.py b/flagscale/train/models/dreamzero/backbone/base_backbone.py new file mode 100644 index 0000000000..b768bcdef5 --- /dev/null +++ b/flagscale/train/models/dreamzero/backbone/base_backbone.py @@ -0,0 +1,16 @@ +from abc import ABC, abstractmethod + +from torch import nn +from transformers.feature_extraction_utils import BatchFeature + + +class Backbone(ABC, nn.Module): + def __init__(self): + super(Backbone, self).__init__() + + @abstractmethod + def forward(self, backbone_input: BatchFeature) -> BatchFeature: + pass + + def prepare_input(self, batch: dict) -> BatchFeature: + pass diff --git a/flagscale/train/models/dreamzero/backbone/identity.py b/flagscale/train/models/dreamzero/backbone/identity.py new file mode 100644 index 0000000000..09d894bff1 --- /dev/null +++ b/flagscale/train/models/dreamzero/backbone/identity.py @@ -0,0 +1,49 @@ +import torch +from transformers.feature_extraction_utils import BatchFeature + +from flagscale.train.models.dreamzero.backbone.base_backbone import Backbone + + +class IdentityBackbone(Backbone): + """ + This class allows pretraining the action head without depending on any backbone. + That's why it's called "identity" — it preserves the action head to be a standalone trainable model. + """ + + def set_trainable_parameters(self, **kwargs): + return + + def forward(self, backbone_input: BatchFeature) -> BatchFeature: + backbone_input_first_value = next(iter(backbone_input.values())) + B = backbone_input_first_value.shape[0] + + backbone_features = torch.empty( + B, 1, 0, dtype=torch.float32, device=backbone_input_first_value.device + ) + output_dict = { + "backbone_features": backbone_features, + } + + return BatchFeature(data=output_dict) + + def prepare_input(self, batch: dict) -> BatchFeature: + """ + Args: + batch: dict + Must contain at least one key-value pair to inform the batch size. + Expects the first dimension to be the batch size. See `forward`. + """ + if "action" in batch: + return BatchFeature(data={"action": batch["action"]}) + else: + # at inference time, we have to use either state or video + if "state" in batch: + return BatchFeature(data={"state": batch["state"]}) + elif "video" in batch: + # For video, it's tricky because it's a numpy array, which isn't compatible with BatchFeature's `to` method + # So instead, we make it a tensor and return it + video = batch["video"] + video_tensor = torch.from_numpy(video) + return BatchFeature(data={"video": video_tensor}) + else: + return BatchFeature(data=batch) diff --git a/flagscale/train/models/dreamzero/base_vla.py b/flagscale/train/models/dreamzero/base_vla.py new file mode 100644 index 0000000000..274b8640f4 --- /dev/null +++ b/flagscale/train/models/dreamzero/base_vla.py @@ -0,0 +1,614 @@ +from dataclasses import dataclass, field +from typing import Tuple + +from hydra.utils import instantiate +import numpy as np +import torch +import torch.distributed as dist +from torch.distributed.device_mesh import DeviceMesh + +from transformers import AutoConfig, AutoModel, PretrainedConfig, PreTrainedModel +from transformers.feature_extraction_utils import BatchFeature +import tree + +BACKBONE_FEATURE_KEY = "backbone_features" +ACTION_KEY = "action_pred" +LOSS_KEY = "loss" +ERROR_MSG = "Error: unexpected input/output" +N_COLOR_CHANNELS = 3 + + +@dataclass +class VLAConfig(PretrainedConfig): + model_type = "vla" + backbone_cfg: PretrainedConfig = field( + default=None, metadata={"help": "Backbone configuration."} + ) + + action_head_cfg: PretrainedConfig = field( + default=None, metadata={"help": "Action head configuration."} + ) + + action_horizon: int = field(default=None, metadata={"help": "Action horizon."}) + + action_dim: int = field(default=None, metadata={"help": "Action dimension."}) + compute_dtype: str = field(default="float32", metadata={"help": "Compute dtype."}) + + def __init__(self, **kwargs): + super().__init__(**kwargs) + for key, value in kwargs.items(): + setattr(self, key, value) + + +class VLA(PreTrainedModel): + supports_gradient_checkpointing = True + config_class = VLAConfig + """ + we expect the backbone output to have a key 'backbone_features' with shape (batch_size, n, hidden_size) + here n is variable and can be e.g. time, 1 or user specified + we expect the action head output to have a key 'action_pred' with shape (batch_size, time, action_dim) during inference time + we expect these to have type BatchFeature, and they can of course have many other user specified keys too + see discussion at https://nvidia.slack.com/archives/C07T1V7L886/p1732550624654139 + """ + + def __init__( + self, + config: VLAConfig, + ): + assert isinstance(config.backbone_cfg, dict) + assert isinstance(config.action_head_cfg, dict) + super().__init__(config) + self.backbone = instantiate(config.backbone_cfg) + self.action_head = instantiate(config.action_head_cfg) + self.action_horizon = config.action_horizon + self.action_dim = config.action_dim + self.compute_dtype = config.compute_dtype + + self.rank = dist.get_rank() if dist.is_initialized() else 0 + + def validate_inputs(self, inputs): + detected_error = False + error_msg = ERROR_MSG + if "action" in inputs: + action = inputs["action"] + type_ok = isinstance(action, torch.Tensor) + shape_ok = ( + len(action.shape) == 3 + and action.shape[1] % self.action_horizon == 0 + and action.shape[2] == self.action_dim + ) + if not type_ok: + error_msg += f"\n{action.dtype=}" + detected_error = True + if not shape_ok: + error_msg += f"\n{action.shape=}" + detected_error = True + + if "video" in inputs: + video = inputs["video"] + type_ok = isinstance(video, np.ndarray) + dtype_ok = video.dtype == np.uint8 + shape_ok = len(video.shape) == 6 and video.shape[3] == N_COLOR_CHANNELS + if not type_ok: + error_msg += f"\n{type(video)=}" + detected_error = True + if not dtype_ok: + error_msg += f"\n{video.dtype=}" + detected_error = True + if not shape_ok: + error_msg += f"\n{video.shape=}" + detected_error = True + + if detected_error: + raise ValueError(error_msg) + + def validate_data(self, action_head_outputs, backbone_outputs, is_training): + + fail_backbone = ( + not isinstance(backbone_outputs, BatchFeature) + or BACKBONE_FEATURE_KEY not in backbone_outputs + ) + + if fail_backbone: + error_msg = ERROR_MSG + error_msg += f"\n{isinstance(backbone_outputs, BatchFeature)=}" + error_msg += f"\n{BACKBONE_FEATURE_KEY in backbone_outputs=}" + error_msg += f"\n{backbone_outputs[BACKBONE_FEATURE_KEY].shape=}" + raise ValueError(error_msg) + + fail_action_head = (not isinstance(action_head_outputs, BatchFeature)) or not ( + ( + LOSS_KEY in action_head_outputs and is_training + ) # there might not be an action prediction during training + or ( + ACTION_KEY in action_head_outputs + and action_head_outputs[ACTION_KEY].shape[1] == self.action_horizon + and action_head_outputs[ACTION_KEY].shape[2] == self.action_dim + ) + ) + + if fail_action_head: + error_msg = ERROR_MSG + error_msg += f"\n{isinstance(action_head_outputs, BatchFeature)=}" + error_msg += f"\n{LOSS_KEY in action_head_outputs=}" + error_msg += f"\n{action_head_outputs[ACTION_KEY].shape=}" + error_msg += f"\n{self.action_horizon=}" + error_msg += f"\n{self.action_dim=}" + raise ValueError(error_msg) + + def forward( + self, + inputs: dict, + ) -> BatchFeature: + + backbone_inputs, action_inputs = self.prepare_input(inputs) + backbone_outputs = self.backbone(backbone_inputs) + action_head_outputs = self.action_head(backbone_outputs, action_inputs) + + return action_head_outputs + + def get_action( + self, + inputs: dict, + ) -> BatchFeature: + backbone_inputs, action_inputs = self.prepare_input(inputs) + backbone_outputs = self.backbone(backbone_inputs) + action_head_outputs = self.action_head.get_action(backbone_outputs, action_inputs) + self.validate_data(action_head_outputs, backbone_outputs, is_training=False) + return action_head_outputs + + def joint_video_action( + self, + inputs: dict, + ) -> BatchFeature: + backbone_inputs, action_inputs = self.prepare_input(inputs) + backbone_outputs = self.backbone(backbone_inputs) + action_head_outputs = self.action_head.joint_video_action(backbone_outputs, action_inputs) + self.validate_data(action_head_outputs, backbone_outputs, is_training=False) + return action_head_outputs + + def lazy_joint_video_action( + self, + inputs: dict, + ) -> BatchFeature: + backbone_inputs, action_inputs = self.prepare_input(inputs) + backbone_outputs = self.backbone(backbone_inputs) + action_head_outputs = self.action_head.lazy_joint_video_action(backbone_outputs, action_inputs) + self.validate_data(action_head_outputs, backbone_outputs, is_training=False) + return action_head_outputs + + def lazy_joint_video_action_causal( + self, + inputs: dict, + latent_video: torch.Tensor | None = None, + ) -> BatchFeature: + backbone_inputs, action_inputs = self.prepare_input(inputs) + backbone_outputs = self.backbone(backbone_inputs) + action_head_outputs = self.action_head.lazy_joint_video_action(backbone_outputs, action_inputs, latent_video=latent_video) + self.validate_data(action_head_outputs, backbone_outputs, is_training=False) + return action_head_outputs + + def lazy_joint_video_action_causal_gt_cond( + self, + inputs: dict, + latent_video: torch.Tensor | None = None, + ) -> BatchFeature: + backbone_inputs, action_inputs = self.prepare_input(inputs) + backbone_outputs = self.backbone(backbone_inputs) + + action_head_outputs = self.action_head.lazy_joint_video_action_causal_gt_cond(backbone_outputs, action_inputs, latent_video=latent_video) + self.validate_data(action_head_outputs, backbone_outputs, is_training=False) + return action_head_outputs + + def lazy_joint_video_action_efficient( + self, + inputs: dict, + prompt_embs: torch.Tensor | None = None, + prompt_emb_nega: torch.Tensor | None = None, + ) -> BatchFeature: + backbone_inputs, action_inputs = self.prepare_input(inputs) + backbone_outputs = self.backbone(backbone_inputs) + action_head_outputs = self.action_head.lazy_joint_video_action_efficient(backbone_outputs, action_inputs, prompt_embs=prompt_embs, prompt_emb_nega=prompt_emb_nega) + self.validate_data(action_head_outputs, backbone_outputs, is_training=False) + return action_head_outputs + + def gt_video_action_pred( + self, + inputs: dict, + ) -> BatchFeature: + backbone_inputs, action_inputs = self.prepare_input(inputs) + backbone_outputs = self.backbone(backbone_inputs) + action_head_outputs = self.action_head.gt_video_action_pred(backbone_outputs, action_inputs) + self.validate_data(action_head_outputs, backbone_outputs, is_training=False) + return action_head_outputs + + def get_language( + self, + inputs: dict, + ) -> BatchFeature: + backbone_inputs, action_inputs = self.prepare_input(inputs) + # Because the behavior of backbones remains the same for training and inference, we can use `forward` for backbones. + backbone_outputs = self.backbone.generate(backbone_inputs) + return backbone_outputs + + def get_video( + self, + inputs: dict, + ) -> BatchFeature: + _, video_inputs = self.prepare_input(inputs) + video_outputs = self.action_head.get_video(video_inputs) + return video_outputs + + def prepare_input(self, inputs) -> Tuple[BatchFeature, BatchFeature]: + self.validate_inputs(inputs) + backbone_inputs = self.backbone.prepare_input(inputs) + action_inputs = self.action_head.prepare_input(inputs) + + def to_device_with_maybe_dtype(x): + # Only cast to self.compute_dtype if the tensor is floating + if torch.is_floating_point(x): + return x.to(self.device, dtype=self.action_head.dtype) + else: + # Keep original dtype + return x.to(self.device) + + backbone_inputs = tree.map_structure(to_device_with_maybe_dtype, backbone_inputs) + action_inputs = tree.map_structure(to_device_with_maybe_dtype, action_inputs) + return backbone_inputs, action_inputs + + + @classmethod + def from_pretrained_for_tuning( + cls, + pretrained_model_name_or_path: str, + config: VLAConfig = None, # This config will now be USED + device_map: str = "auto", + dtype: torch.dtype = torch.bfloat16, + offload_state_dict: bool = True, + lora_weights_path: str | None = None, + ): + if config is None: + raise ValueError( + "A `config` object must be provided to build the model structure." + ) + + import os + import json + import gc + from safetensors.torch import load_file + + model = cls(config) + + safetensors_path = os.path.join(pretrained_model_name_or_path, "model.safetensors") + safetensors_index_path = os.path.join(pretrained_model_name_or_path, "model.safetensors.index.json") + + if os.path.exists(safetensors_index_path): + with open(safetensors_index_path, 'r') as f: + index = json.load(f) + missing_keys_accum = set() + unexpected_keys_accum = set() + shard_files = sorted(set(index["weight_map"].values())) + for shard_file in shard_files: + shard_path = os.path.join(pretrained_model_name_or_path, shard_file) + print(f"Loading shard: {shard_path}") + shard_state_dict = load_file(shard_path) + missing_keys, unexpected_keys = model.load_state_dict(shard_state_dict, strict=False) + if missing_keys: + missing_keys_accum.update(missing_keys) + if unexpected_keys: + unexpected_keys_accum.update(unexpected_keys) + # Free shard immediately + del shard_state_dict + gc.collect() + if missing_keys_accum: + print(f"Missing keys when loading sharded pretrained weights: {sorted(missing_keys_accum)} ... total={len(missing_keys_accum)}") + if unexpected_keys_accum: + print(f"Unexpected keys when loading sharded pretrained weights: {sorted(unexpected_keys_accum)} ... total={len(unexpected_keys_accum)}") + if not missing_keys_accum and not unexpected_keys_accum: + print("Successfully loaded pretrained base weights (sharded)") + elif os.path.exists(safetensors_path): + # Handle single safetensors file + print(f"Loading weights from safetensors: {safetensors_path}") + state_dict = load_file(safetensors_path) + missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False) + if missing_keys: + print(f"Missing keys when loading pretrained weights: {missing_keys}") + if unexpected_keys: + print(f"Unexpected keys when loading pretrained weights: {unexpected_keys}") + if not missing_keys and not unexpected_keys: + print("Successfully loaded pretrained base weights") + else: + raise FileNotFoundError( + f"No weights found at '{pretrained_model_name_or_path}'. " + "Expected 'model.safetensors' or 'model.safetensors.index.json'." + ) + + if lora_weights_path is not None: + print(f"Loading LoRA weights from: {lora_weights_path}") + model.load_lora_weight(lora_weights_path) + else: + if hasattr(model, 'action_head') and hasattr(model.action_head, 'inject_lora_after_loading') and model.action_head.config.defer_lora_injection: + print("Injecting LoRA adapters into action_head after loading pretrained weights") + model.action_head.inject_lora_after_loading() + + print(f"{cls}\n") + return model + + @classmethod + def load_lora( + cls, + pretrained_model_name_or_path: str + ): + from safetensors.torch import load_file + import os + import json + print("loading lora@@@@@") + + # Check for different checkpoint formats + safetensors_path = os.path.join(pretrained_model_name_or_path, "model.safetensors") + safetensors_index_path = os.path.join(pretrained_model_name_or_path, "model.safetensors.index.json") + + state_dict = {} + if os.path.exists(safetensors_index_path): + # Handle sharded safetensors + print(f"Loading sharded safetensors using index: {safetensors_index_path}") + + with open(safetensors_index_path, 'r') as f: + index = json.load(f) + + # Load each shard + for shard_file in set(index["weight_map"].values()): + shard_path = os.path.join(pretrained_model_name_or_path, shard_file) + print(f"Loading shard: {shard_path}") + shard_state_dict = load_file(shard_path) + state_dict.update(shard_state_dict) + + elif os.path.exists(safetensors_path): + # Handle single safetensors file + print(f"Loading weights from safetensors: {safetensors_path}") + state_dict.update(load_file(safetensors_path)) + + # Load config + print("loading config@@") + config_path = os.path.join(pretrained_model_name_or_path, "config.json") + with open(config_path, "r") as f: + config_dict = json.load(f) + config = VLAConfig(**config_dict) + print("loading model") + + # Disable defer_lora_injection so LoRA layers are created during init, + # matching the PEFT key hierarchy (base_model.model.*) in the checkpoint. + ah_cfg = config.action_head_cfg + inner = ah_cfg.get('config', ah_cfg) if isinstance(ah_cfg.get('config'), dict) else ah_cfg + if 'defer_lora_injection' in inner: + inner['defer_lora_injection'] = False + print("defer_lora_injection disabled for load_lora") + # Enable component loading so DiT base weights are loaded from pretrained + if 'skip_component_loading' in inner: + inner['skip_component_loading'] = False + print("skip_component_loading disabled for load_lora") + + # Instantiate model (LoRA layers now exist from init) + model = cls(config) + + # Remove .base_layer from keys if present + has_base_layer = any(".base_layer." in key for key in state_dict.keys()) + if has_base_layer: + print("Removing '.base_layer' from state dict keys") + state_dict = {k.replace(".base_layer.", "."): v for k, v in state_dict.items()} + + # Load weights + missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False) + + if missing_keys: + print(f"Missing keys when loading pretrained weights: {missing_keys}") + if unexpected_keys: + print(f"Unexpected keys when loading pretrained weights: {unexpected_keys}") + + print("Successfully loaded pretrained weights") + + print(f"{cls}\n") + return model + + def load_lora_weight(self, pretrained_model_name_or_path: str): + """Load only LoRA weights from a pretrained model without loading config.""" + from safetensors.torch import load_file + import os + import json + + print(f"Loading LoRA weights from {pretrained_model_name_or_path}") + + # Check for different checkpoint formats + safetensors_path = os.path.join(pretrained_model_name_or_path, "model.safetensors") + safetensors_index_path = os.path.join(pretrained_model_name_or_path, "model.safetensors.index.json") + + state_dict = {} + if os.path.exists(safetensors_index_path): + # Handle sharded safetensors + print(f"Loading sharded safetensors using index: {safetensors_index_path}") + + with open(safetensors_index_path, 'r') as f: + index = json.load(f) + + # Load each shard + for shard_file in set(index["weight_map"].values()): + shard_path = os.path.join(pretrained_model_name_or_path, shard_file) + print(f"Loading shard: {shard_path}") + shard_state_dict = load_file(shard_path) + state_dict.update(shard_state_dict) + + elif os.path.exists(safetensors_path): + # Handle single safetensors file + print(f"Loading weights from safetensors: {safetensors_path}") + state_dict.update(load_file(safetensors_path)) + else: + raise FileNotFoundError(f"No valid checkpoint found at {pretrained_model_name_or_path}") + + print("Loading LoRA weights into existing model") + + def rewrite_lora_state_dict_keys(state_dict, pattern, repl): + new_state_dict = {} + for k, v in state_dict.items(): + new_k = k.replace(pattern, repl) + new_state_dict[new_k] = v + return new_state_dict + + has_target_pattern = any("action_head.model.base_model.model" in key for key in state_dict.keys()) + + if not has_target_pattern: + print("Rewriting LoRA state dict keys from 'action_head.model' to 'action_head.model.base_model.model'") + state_dict = rewrite_lora_state_dict_keys( + state_dict, + pattern="action_head.model", + repl="action_head.model.base_model.model", + ) + else: + print("State dict already has 'action_head.model.base_model.model' pattern, skipping key rewrite") + + # Load only the weights into the existing model + missing_keys, unexpected_keys = self.load_state_dict(state_dict, strict=False) + + print("Successfully loaded LoRA state dict") + + if missing_keys: + print(f"Missing keys when loading LoRA weights: {missing_keys}") + if unexpected_keys: + print(f"Unexpected keys when loading LoRA weights: {unexpected_keys}") + + print("Successfully loaded LoRA weights") + + @classmethod + def from_config_with_lora_weights( + cls, + config: VLAConfig, + pretrained_model_path: str, + ): + """Create VLA model from config and then load LoRA weights from pretrained model.""" + print(f"Creating VLA model from config and loading LoRA weights from {pretrained_model_path}") + + # 1. Create model from config (similar to vla.yaml) + model = cls(config) + print("Model created from config") + + # 2. Load LoRA weights into the created model + model.load_lora_weight(pretrained_model_path) + + return model + + @classmethod + def from_pretrained( + cls, + pretrained_model_name_or_path: str, + config: VLAConfig = None + ): + del config + + from safetensors.torch import load_file + import os + import json + print("loading pretrained@@@@@") + # Check for different checkpoint formats + safetensors_path = os.path.join(pretrained_model_name_or_path, "model.safetensors") + safetensors_index_path = os.path.join(pretrained_model_name_or_path, "model.safetensors.index.json") + + state_dict = {} + if os.path.exists(safetensors_index_path): + # Handle sharded safetensors + print(f"Loading sharded safetensors using index: {safetensors_index_path}") + + with open(safetensors_index_path, 'r') as f: + index = json.load(f) + + # Load each shard + for shard_file in set(index["weight_map"].values()): + shard_path = os.path.join(pretrained_model_name_or_path, shard_file) + print(f"Loading shard: {shard_path}") + shard_state_dict = load_file(shard_path) + state_dict.update(shard_state_dict) + + elif os.path.exists(safetensors_path): + # Handle single safetensors file + print(f"Loading weights from safetensors: {safetensors_path}") + state_dict.update(load_file(safetensors_path)) + + # Load config + print("loading config@@") + config_path = os.path.join(pretrained_model_name_or_path, "config.json") + with open(config_path, "r") as f: + config_dict = json.load(f) + config = VLAConfig(**config_dict) + print("loading model") + print("config.action_head_cfg", config.action_head_cfg) + # Always disable defer_lora_injection + # config.action_head_cfg is a dict, and defer_lora_injection is nested in config.action_head_cfg['config'] + if 'config' in config.action_head_cfg and isinstance(config.action_head_cfg['config'], dict): + if 'defer_lora_injection' in config.action_head_cfg['config']: + config.action_head_cfg['config']['defer_lora_injection'] = False + print("config.action_head_cfg['config']['defer_lora_injection'] disabled (set to False)") + elif 'defer_lora_injection' in config.action_head_cfg: + config.action_head_cfg['defer_lora_injection'] = False + print("config.action_head_cfg['defer_lora_injection'] disabled (set to False)") + + # Instantiate model + model = cls(config) + print("model", model) + # Remove .base_layer from keys (e.g., 'action_head.model.base_model.model.blocks.19.self_attn.v.base_layer.bias' -> 'action_head.model.base_model.model.blocks.19.self_attn.v.bias') + has_base_layer = any(".base_layer." in key for key in state_dict.keys()) + if has_base_layer: + print("Removing '.base_layer' from state dict keys") + new_state_dict = {} + for k, v in state_dict.items(): + new_k = k.replace(".base_layer.", ".") + new_state_dict[new_k] = v + state_dict = new_state_dict + + missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False) + + if missing_keys: + print(f"Missing keys when loading pretrained weights: {missing_keys}") + if unexpected_keys: + print(f"Unexpected keys when loading pretrained weights: {unexpected_keys}") + + print("Successfully loaded pretrained weights") + + print(f"{cls}\n") + return model + + def post_initialize(self): + self.action_head.post_initialize() + + def parallelize(self, device_mesh: DeviceMesh): + self.action_head.parallelize(device_mesh=device_mesh) + + +class CotrainVLA(VLA): + + def forward( + self, + inputs: dict, + ) -> BatchFeature: + if "cotrain" in inputs and inputs["cotrain"]: + return self.backbone.cotrain(inputs) + return super().forward(inputs) + + +def create_vla_with_pretrained_action_head(pretrained_vla_path: str, config: VLAConfig): + # 1. Instantiate a new VLAModel + vla = VLA(config) + + # 2. Load the pretrained VLAModel + pretrained_vla = VLA.from_pretrained(pretrained_vla_path) + + # 3. Replace the action head in the new VLAModel with the pretrained action head + vla.action_head = pretrained_vla.action_head + + # 4. Replace the action head config in the new VLAModel with the pretrained action head config + vla.config.action_head_cfg = pretrained_vla.config.action_head_cfg + + # 5. Return the new VLAModel + return vla + + +# register +AutoConfig.register("vla", VLAConfig) +AutoModel.register(VLAConfig, VLA) diff --git a/flagscale/train/models/dreamzero/dreamzero_model.py b/flagscale/train/models/dreamzero/dreamzero_model.py new file mode 100644 index 0000000000..569d438127 --- /dev/null +++ b/flagscale/train/models/dreamzero/dreamzero_model.py @@ -0,0 +1,414 @@ +# Copyright (c) 2025, FlagScale Authors. All rights reserved. +""" +DreamZero Model Module for FlagScale native training backend. + +DreamZero is a World Action Model (WAM) - a Wan2.1 video DiT repurposed as a +zero-shot robot policy. Architecture: + - IdentityBackbone (no LLM backbone) + - WANPolicyHead: CausalWanModel (DiT) + VAE + T5 + CLIP + action/state encoders + +Training forward: + 1. Encode video frames -> latents (VAE, frozen) + 2. Encode text -> prompt embeddings (T5, frozen) + 3. Encode first frame -> CLIP features (CLIP, frozen) + 4. Add noise to latents and actions (flow matching scheduler) + 5. DiT predicts noise (trainable, with LoRA) + 6. Compute MSE loss (dynamics + action) +""" + +import logging +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import torch +import torch.nn as nn +from transformers import PreTrainedModel, PretrainedConfig +from transformers.feature_extraction_utils import BatchFeature + +logger = logging.getLogger(__name__) + + +@dataclass +class DreamZeroConfig(PretrainedConfig): + """DreamZero model configuration.""" + model_type = "dreamzero" + + # DiT backbone (14B default from DreamZero-AgiBot) + dim: int = 5120 + ffn_dim: int = 13824 + num_heads: int = 40 + num_layers: int = 40 + freq_dim: int = 256 + in_dim: int = 36 + out_dim: int = 16 + frame_seqlen: int = 220 + max_chunk_size: int = 4 + num_frame_per_block: int = 2 + num_action_per_block: int = 48 + num_state_per_block: int = 1 + eps: float = 1e-6 + model_subtype: str = "i2v" + + # Action head + action_dim: int = 32 + action_horizon: int = 48 + max_state_dim: int = 64 + max_action_dim: int = 32 + hidden_size: int = 64 + input_embedding_dim: int = 1536 + + # LoRA + lora_rank: int = 4 + lora_alpha: int = 4 + lora_target_modules: str = "q,k,v,o,ffn.0,ffn.2" + + # Training + num_frames: int = 33 + use_gradient_checkpointing: bool = True + train_architecture: str = "full" + tune_diffusion_model: bool = True + tune_projector: bool = True + use_vlln: bool = True + + # Noise scheduling + noise_beta_alpha: float = 1.5 + noise_beta_beta: float = 1.0 + noise_s: float = 0.999 + num_timestep_buckets: int = 1000 + + # REPA + repa_layer: int = 8 + repa_coeff: float = 1.0 + + # VL self-attention + vl_num_layers: int = 4 + vl_num_heads: int = 24 + vl_head_dim: int = 64 + vl_dropout: float = 0.2 + + # Compute + compute_dtype: str = "bfloat16" + + # Paths (set at runtime) + model_path: str = None + tokenizer_path: str = None + dit_version: str = None # Path to Wan2.1-I2V-14B-480P dir (DIT pretrained weights) + text_encoder_pretrained_path: str = None + image_encoder_pretrained_path: str = None + vae_pretrained_path: str = None + + # Embodiment + embodiment_tag: str = "agibot" + + def __init__(self, **kwargs): + super().__init__(**kwargs) + for key, value in kwargs.items(): + setattr(self, key, value) + + +class DreamZeroPolicy(PreTrainedModel): + """DreamZero policy model for FlagScale native backend (FSDP2). + + Wraps the original DreamZero VLA (IdentityBackbone + WANPolicyHead). + All weights live under action_head.* (2146 keys total). + """ + + supports_gradient_checkpointing = True + config_class = DreamZeroConfig + + # FSDP2 no-split modules + _no_split_modules = [ + "T5SelfAttention", + "AttentionBlock", + "CausalWanModel", + "CausalWanAttentionBlock", + ] + + def __init__(self, config: DreamZeroConfig): + super().__init__(config) + self.config = config + self._model_loaded = False + + @classmethod + def from_pretrained_dreamzero(cls, config: DreamZeroConfig) -> "DreamZeroPolicy": + """Load DreamZero for FSDP2 training. + + Supports two training modes: + - train_architecture="full": Full fine-tuning of DIT + action/state projectors. + Loads ALL weights from DreamZero-AgiBot checkpoint directly. No LoRA. + VAE/T5/CLIP frozen, DIT + encoders/decoders fully trainable. + - train_architecture="lora": LoRA on DIT with Wan2.1 base weights. + Loads Wan2.1 DIT from pretrained, skips DIT keys from checkpoint, + then injects LoRA. (Matches reference accidental behavior.) + + Full fine-tuning is preferred for FSDP2 since FSDP2 handles memory sharding + natively and avoids the LoRA-on-fine-tuned-weights NaN issue. + """ + from flagscale.train.models.dreamzero.base_vla import VLA, VLAConfig + import json + import gc + + model_path = Path(config.model_path) + if not model_path.exists(): + raise FileNotFoundError(f"Model path not found: {model_path}") + + # Step 1: Build model config. Use skip_component_loading=False so __init__ + # loads all pretrained weights (T5, CLIP, VAE, Wan2.1 DIT). + # Use defer_lora_injection=True so we can load non-DIT weights from checkpoint + # before injecting LoRA (unlike the reference which relies on PEFT key mismatch). + config_json = model_path / "config.json" + if not config_json.exists(): + raise FileNotFoundError( + f"config.json not found at {model_path}. " + "DreamZero checkpoint must contain config.json." + ) + with open(config_json) as f: + vla_dict = json.load(f) + + # skip_component_loading controls whether __init__ loads T5/CLIP/VAE/DIT individually. + # For full fine-tuning: skip it (we load everything from checkpoint afterward). + # For LoRA mode: load components so DIT gets Wan2.1 base weights. + use_lora = config.train_architecture == "lora" + load_components_separately = use_lora and config.dit_version is not None + if "action_head_cfg" in vla_dict and "config" in vla_dict["action_head_cfg"]: + vla_dict["action_head_cfg"]["config"]["defer_lora_injection"] = True + vla_dict["action_head_cfg"]["config"]["skip_component_loading"] = not load_components_separately + + # Inject pretrained paths for component loading (LoRA mode only) + if load_components_separately and "action_head_cfg" in vla_dict: + ah_cfg = vla_dict["action_head_cfg"].get("config", {}) + # DIT pretrained path (Wan2.1 weights) + if "diffusion_model_cfg" in ah_cfg: + ah_cfg["diffusion_model_cfg"]["diffusion_model_pretrained_path"] = config.dit_version + # Text encoder + if config.text_encoder_pretrained_path and "text_encoder_cfg" in ah_cfg: + ah_cfg["text_encoder_cfg"]["text_encoder_pretrained_path"] = config.text_encoder_pretrained_path + # Image encoder + if config.image_encoder_pretrained_path and "image_encoder_cfg" in ah_cfg: + ah_cfg["image_encoder_cfg"]["image_encoder_pretrained_path"] = config.image_encoder_pretrained_path + # VAE + if config.vae_pretrained_path and "vae_cfg" in ah_cfg: + ah_cfg["vae_cfg"]["vae_pretrained_path"] = config.vae_pretrained_path + + # Instantiate VLA + vla_config = VLAConfig(**vla_dict) + vla_model = VLA(vla_config) + + if load_components_separately: + logger.info("VLA instantiated with Wan2.1 DIT from pretrained (LoRA mode)") + else: + logger.info("VLA instantiated (full fine-tuning mode, loading all from checkpoint)") + + # Step 2: Load DreamZero-AgiBot checkpoint. + # If use_component_loading=True: filter out DIT keys (DIT keeps Wan2.1 base weights). + # If use_component_loading=False: load ALL keys from checkpoint. + from safetensors.torch import load_file + import os + + safetensors_index_path = os.path.join(str(model_path), "model.safetensors.index.json") + safetensors_path = os.path.join(str(model_path), "model.safetensors") + + # In LoRA mode: filter DIT + encoder/decoder keys (DIT keeps Wan2.1, encoders start random) + # In full mode: load everything from checkpoint + dit_prefixes = ( + "action_head.model.blocks.", + "action_head.model.head.", + "action_head.model.img_emb.", + "action_head.model.patch_embedding.", + "action_head.model.text_embedding.", + "action_head.model.time_embedding.", + "action_head.model.time_projection.", + "action_head.model.state_encoder.", + "action_head.model.action_encoder.", + "action_head.model.action_decoder.", + ) + + def maybe_filter(state_dict): + """Filter DIT keys in LoRA mode; pass everything in full fine-tuning mode.""" + if not load_components_separately: + return state_dict, 0 + filtered = {} + skipped = 0 + for k, v in state_dict.items(): + if any(k.startswith(p) for p in dit_prefixes): + skipped += 1 + else: + filtered[k] = v + return filtered, skipped + + if os.path.exists(safetensors_index_path): + with open(safetensors_index_path, 'r') as f: + index = json.load(f) + total_skipped = 0 + for shard_file in sorted(set(index["weight_map"].values())): + shard_path = os.path.join(str(model_path), shard_file) + logger.info(f"Loading shard: {shard_path}") + shard_state_dict = load_file(shard_path) + shard_state_dict, skipped = maybe_filter(shard_state_dict) + total_skipped += skipped + if shard_state_dict: + vla_model.load_state_dict(shard_state_dict, strict=False) + del shard_state_dict + gc.collect() + if total_skipped: + logger.info(f"Skipped {total_skipped} DIT keys (keeping Wan2.1 base weights)") + elif os.path.exists(safetensors_path): + state_dict = load_file(safetensors_path) + state_dict, skipped = maybe_filter(state_dict) + if skipped: + logger.info(f"Skipped {skipped} DIT keys (keeping Wan2.1 base weights)") + vla_model.load_state_dict(state_dict, strict=False) + del state_dict + gc.collect() + else: + raise FileNotFoundError( + f"No weights at '{model_path}'. " + "Expected 'model.safetensors' or 'model.safetensors.index.json'." + ) + logger.info("DreamZero-AgiBot checkpoint loaded") + + # Step 3: Transfer action_head to our policy wrapper + policy = cls(config) + policy.action_head = vla_model.action_head + policy._model_loaded = True + del vla_model + gc.collect() + + # Override frame_seqlen on all attention blocks to match our resolution. + # With num_views=2, tiled to 352x640 → VAE 8x → 44x80 → patch (1,2,2) → 22x40 = 880 + correct_frame_seqlen = config.frame_seqlen + if hasattr(policy.action_head, 'model'): + for module in policy.action_head.model.modules(): + if hasattr(module, 'frame_seqlen'): + module.frame_seqlen = correct_frame_seqlen + logger.info(f"Override frame_seqlen={correct_frame_seqlen} on all attention blocks") + + # Override action_horizon and num_action_per_block from our config. + correct_action_horizon = config.action_horizon + correct_num_action_per_block = config.num_action_per_block + if hasattr(policy.action_head, 'model'): + dit = policy.action_head.model + if hasattr(dit, 'num_action_per_block'): + dit.num_action_per_block = correct_num_action_per_block + if hasattr(dit, 'action_horizon'): + dit.action_horizon = correct_action_horizon + for module in dit.modules(): + if hasattr(module, 'num_action_per_block'): + module.num_action_per_block = correct_num_action_per_block + if hasattr(policy.action_head, 'action_horizon'): + policy.action_head.action_horizon = correct_action_horizon + logger.info( + f"Override action params: action_horizon={correct_action_horizon}, " + f"num_action_per_block={correct_num_action_per_block}" + ) + + # Step 4: Inject LoRA ONCE (same as reference base.py:734) + if config.train_architecture == "lora": + policy.action_head.lora_rank = config.lora_rank + policy.action_head.lora_alpha = config.lora_alpha + policy.action_head.lora_target_modules = config.lora_target_modules + policy.action_head.init_lora_weights = getattr(config, "init_lora_weights", "kaiming") + policy.action_head.train_architecture = "lora" + # Reset RNG to fixed state before LoRA init for reproducibility. + # This ensures lora_A (Kaiming init) is identical regardless of + # code path between set_seed() and here (which differs between + # FSDP2 and DeepSpeed/HF Trainer pipelines). + rng_state = torch.random.get_rng_state() + torch.manual_seed(1234) + policy.action_head.inject_lora_after_loading() + torch.random.set_rng_state(rng_state) + logger.info( + f"LoRA injected: rank={config.lora_rank}, alpha={config.lora_alpha}, " + f"targets={config.lora_target_modules}" + ) + else: + policy._apply_freeze_config() + + trainable = sum(p.numel() for p in policy.parameters() if p.requires_grad) + total = sum(p.numel() for p in policy.parameters()) + logger.info( + f"DreamZero loaded. Total: {total/1e9:.2f}B, " + f"Trainable: {trainable/1e9:.2f}B ({100*trainable/total:.1f}%)" + ) + + return policy + + def _apply_freeze_config(self): + """Freeze VAE, T5 text encoder, CLIP image encoder. Keep DiT + action head trainable.""" + if not hasattr(self, "action_head"): + return + + # Freeze VAE + if hasattr(self.action_head, "vae"): + for p in self.action_head.vae.parameters(): + p.requires_grad = False + self.action_head.vae.eval() + logger.info("Froze VAE encoder") + + # Freeze T5 text encoder + if hasattr(self.action_head, "text_encoder"): + for p in self.action_head.text_encoder.parameters(): + p.requires_grad = False + self.action_head.text_encoder.eval() + logger.info("Froze T5 text encoder") + + # Freeze CLIP image encoder + if hasattr(self.action_head, "image_encoder"): + for p in self.action_head.image_encoder.parameters(): + p.requires_grad = False + self.action_head.image_encoder.eval() + logger.info("Froze CLIP image encoder") + + # If LoRA mode, freeze DiT base weights and only train LoRA + if self.config.train_architecture == "lora": + for name, p in self.action_head.model.named_parameters(): + if "lora" not in name.lower(): + p.requires_grad = False + logger.info("LoRA mode: froze DiT base weights, training LoRA only") + + def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None): + """Enable gradient checkpointing on the DiT backbone.""" + if gradient_checkpointing_kwargs is None: + gradient_checkpointing_kwargs = {} + use_reentrant = gradient_checkpointing_kwargs.get("use_reentrant", False) + + if hasattr(self, "action_head") and hasattr(self.action_head, "model"): + diffusion_model = self.action_head.model + if hasattr(diffusion_model, "gradient_checkpointing"): + diffusion_model.gradient_checkpointing = True + setattr(diffusion_model, "gradient_checkpointing_use_reentrant", use_reentrant) + logger.info(f"Enabled gradient checkpointing (use_reentrant={use_reentrant})") + + def forward(self, batch: dict) -> dict[str, torch.Tensor]: + """Training forward pass. + + Args: + batch: Dict from get_batch containing: + - images: (B, T, H, W, C) uint8 video frames + - action: (B, action_horizon, action_dim) normalized actions [-1, 1] + - state: (B, state_dim) robot proprioceptive state + - text: (B, seq_len) tokenized text + - text_attention_mask: (B, seq_len) attention mask + - embodiment_id: (B,) embodiment category IDs + - has_real_action: (B,) bool mask for valid actions + - action_mask: (B, action_horizon, action_dim) action validity mask + + Returns: + Dict with 'loss', 'dynamics_loss', 'action_loss' + """ + # The action_head.forward handles the full training pipeline: + # encode video, add noise, predict, compute loss + backbone_output = BatchFeature(data={}) # Identity backbone + action_input = BatchFeature(data=batch) + output = self.action_head(backbone_output, action_input) + + if hasattr(output, "data"): + return dict(output.data) + return dict(output) + + def set_frozen_modules_to_eval(self): + """Set frozen modules to eval mode (called before each forward).""" + if hasattr(self, "action_head"): + if hasattr(self.action_head, "set_frozen_modules_to_eval_mode"): + self.action_head.set_frozen_modules_to_eval_mode() diff --git a/flagscale/train/models/dreamzero/modules/__init__.py b/flagscale/train/models/dreamzero/modules/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/flagscale/train/models/dreamzero/modules/attention.py b/flagscale/train/models/dreamzero/modules/attention.py new file mode 100644 index 0000000000..0ebc62a64a --- /dev/null +++ b/flagscale/train/models/dreamzero/modules/attention.py @@ -0,0 +1,258 @@ +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. +import torch +import os + +try: + import flash_attn_interface + FLASH_ATTN_3_AVAILABLE = True +except ModuleNotFoundError: + FLASH_ATTN_3_AVAILABLE = False + +try: + import flash_attn + FLASH_ATTN_2_AVAILABLE = True +except ModuleNotFoundError: + FLASH_ATTN_2_AVAILABLE = False + +import warnings + + +__all__ = [ + 'flash_attention', + 'attention', +] + + +def _gpu_supports_flash_attention(): + """FlashAttention requires Ampere (compute capability 8.0) or newer.""" + if not (FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE): + return False + try: + if not torch.cuda.is_available(): + return False + cap = torch.cuda.get_device_capability() + return cap[0] >= 8 + except Exception: + return False + + +def _sdpa_attention_fallback( + q, k, v, + q_lens=None, + k_lens=None, + dropout_p=0., + softmax_scale=None, + q_scale=None, + causal=False, + dtype=torch.bfloat16, +): + """PyTorch SDPA fallback for GPUs that don't support FlashAttention (e.g. pre-Ampere).""" + if q_lens is not None or k_lens is not None: + warnings.warn( + 'Padding mask is disabled when using scaled_dot_product_attention on this GPU. ' + 'It can have a slight impact on quality.' + ) + q = q.transpose(1, 2).to(dtype) + k = k.transpose(1, 2).to(dtype) + v = v.transpose(1, 2).to(dtype) + if q_scale is not None: + q = q * q_scale + if softmax_scale is not None: + q = q * softmax_scale + out = torch.nn.functional.scaled_dot_product_attention( + q, k, v, attn_mask=None, is_causal=causal, dropout_p=dropout_p + ) + return out.transpose(1, 2).contiguous() + + +def flash_attention( + q, + k, + v, + q_lens=None, + k_lens=None, + dropout_p=0., + softmax_scale=None, + q_scale=None, + causal=False, + window_size=(-1, -1), + deterministic=False, + dtype=torch.bfloat16, + version=None, +): + """ + q: [B, Lq, Nq, C1]. + k: [B, Lk, Nk, C1]. + v: [B, Lk, Nk, C2]. Nq must be divisible by Nk. + q_lens: [B]. + k_lens: [B]. + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + causal: bool. Whether to apply causal attention mask. + window_size: (left right). If not (-1, -1), apply sliding window local attention. + deterministic: bool. If True, slightly slower and uses more memory. + dtype: torch.dtype. Apply when dtype of q/k/v is not float16/bfloat16. + """ + half_dtypes = (torch.float16, torch.bfloat16) + assert dtype in half_dtypes + assert q.device.type == 'cuda' and q.size(-1) <= 256 + + # Use PyTorch SDPA on pre-Ampere GPUs (FlashAttention requires Ampere or newer) + if not _gpu_supports_flash_attention(): + return _sdpa_attention_fallback( + q, k, v, + q_lens=q_lens, + k_lens=k_lens, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + q_scale=q_scale, + causal=causal, + dtype=dtype, + ) + + # params + b, lq, lk, out_dtype = q.size(0), q.size(1), k.size(1), q.dtype + + def half(x): + return x if x.dtype in half_dtypes else x.to(dtype) + + # preprocess query + if q_lens is None: + q = half(q.flatten(0, 1)) + q_lens = torch.tensor( + [lq] * b, dtype=torch.int32).to( + device=q.device, non_blocking=True) + else: + q = half(torch.cat([u[:v] for u, v in zip(q, q_lens)])) + + # preprocess key, value + if k_lens is None: + k = half(k.flatten(0, 1)) + v = half(v.flatten(0, 1)) + k_lens = torch.tensor( + [lk] * b, dtype=torch.int32).to( + device=k.device, non_blocking=True) + else: + k = half(torch.cat([u[:v] for u, v in zip(k, k_lens)])) + v = half(torch.cat([u[:v] for u, v in zip(v, k_lens)])) + + q = q.to(v.dtype) + k = k.to(v.dtype) + + if q_scale is not None: + q = q * q_scale + + if version is not None and version == 3 and not FLASH_ATTN_3_AVAILABLE: + warnings.warn( + 'Flash attention 3 is not available, use flash attention 2 instead.' + ) + + # Check for TensorRT at runtime, not import time + if os.getenv("ENABLE_TENSORRT", "False").lower() == "true": + # use torch.nn.functional.scaled_dot_product_attention for tensorrt export + + # The input is (s, n, d), but sdpa needs (b, n, s, d). + # We add a batch dimension and transpose. + q = q.unsqueeze(0).transpose(1, 2) + k = k.unsqueeze(0).transpose(1, 2) + v = v.unsqueeze(0).transpose(1, 2) + + # Fix for ONNX export: repeat k and v to match q's batch size in cross-attention + if q.shape[0] != k.shape[0] and k.shape[0] == 1: + k = k.repeat(q.shape[0], 1, 1, 1) + v = v.repeat(q.shape[0], 1, 1, 1) + + attn_mask = None + out = torch.nn.functional.scaled_dot_product_attention( + q, k, v, attn_mask=attn_mask, is_causal=causal, dropout_p=dropout_p) + + # Transpose back to (b, s, n, d) format. + out = out.transpose(1, 2).contiguous() + return out + + elif (version is None or version == 3) and FLASH_ATTN_3_AVAILABLE: + # Note: dropout_p, window_size are not supported in FA3 now. + x = flash_attn_interface.flash_attn_varlen_func( + q=q, + k=k, + v=v, + cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens]).cumsum( + 0, dtype=torch.int32).to(q.device, non_blocking=True), + cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens]).cumsum( + 0, dtype=torch.int32).to(q.device, non_blocking=True), + seqused_q=None, + seqused_k=None, + max_seqlen_q=lq, + max_seqlen_k=lk, + softmax_scale=softmax_scale, + causal=causal, + deterministic=deterministic)[0].unflatten(0, (b, lq)) + else: + assert FLASH_ATTN_2_AVAILABLE + x = flash_attn.flash_attn_varlen_func( + q=q, + k=k, + v=v, + cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens]).cumsum( + 0, dtype=torch.int32).to(q.device, non_blocking=True), + cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens]).cumsum( + 0, dtype=torch.int32).to(q.device, non_blocking=True), + max_seqlen_q=lq, + max_seqlen_k=lk, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + causal=causal, + window_size=window_size, + deterministic=deterministic).unflatten(0, (b, lq)) + + # output + return x.type(out_dtype) + + +def attention( + q, + k, + v, + q_lens=None, + k_lens=None, + dropout_p=0., + softmax_scale=None, + q_scale=None, + causal=False, + window_size=(-1, -1), + deterministic=False, + dtype=torch.bfloat16, + fa_version=None, +): + if FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE: + return flash_attention( + q=q, + k=k, + v=v, + q_lens=q_lens, + k_lens=k_lens, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + q_scale=q_scale, + causal=causal, + window_size=window_size, + deterministic=deterministic, + dtype=dtype, + version=fa_version, + ) + else: + if q_lens is not None or k_lens is not None: + warnings.warn( + 'Padding mask is disabled when using scaled_dot_product_attention. It can have a significant impact on performance.' + ) + attn_mask = None + + q = q.transpose(1, 2).to(dtype) + k = k.transpose(1, 2).to(dtype) + v = v.transpose(1, 2).to(dtype) + + out = torch.nn.functional.scaled_dot_product_attention( + q, k, v, attn_mask=attn_mask, is_causal=causal, dropout_p=dropout_p) + + out = out.transpose(1, 2).contiguous() + return out diff --git a/flagscale/train/models/dreamzero/modules/flow_match_scheduler.py b/flagscale/train/models/dreamzero/modules/flow_match_scheduler.py new file mode 100644 index 0000000000..03b53b3b26 --- /dev/null +++ b/flagscale/train/models/dreamzero/modules/flow_match_scheduler.py @@ -0,0 +1,92 @@ +import torch + + + +class FlowMatchScheduler(): + + def __init__(self, num_inference_steps=100, num_train_timesteps=1000, shift=3.0, sigma_max=1.0, sigma_min=0.003/1.002, inverse_timesteps=False, extra_one_step=False, reverse_sigmas=False): + self.num_train_timesteps = num_train_timesteps + self.shift = shift + self.sigma_max = sigma_max + self.sigma_min = sigma_min + self.inverse_timesteps = inverse_timesteps + self.extra_one_step = extra_one_step + self.reverse_sigmas = reverse_sigmas + self.set_timesteps(num_inference_steps) + + + def set_timesteps(self, num_inference_steps=100, denoising_strength=1.0, training=False, shift=None): + if shift is not None: + self.shift = shift + sigma_start = self.sigma_min + (self.sigma_max - self.sigma_min) * denoising_strength + if self.extra_one_step: + self.sigmas = torch.linspace(sigma_start, self.sigma_min, num_inference_steps + 1)[:-1] + else: + self.sigmas = torch.linspace(sigma_start, self.sigma_min, num_inference_steps) + if self.inverse_timesteps: + self.sigmas = torch.flip(self.sigmas, dims=[0]) + self.sigmas = self.shift * self.sigmas / (1 + (self.shift - 1) * self.sigmas) + if self.reverse_sigmas: + self.sigmas = 1 - self.sigmas + self.timesteps = self.sigmas * self.num_train_timesteps + if training: + x = self.timesteps + y = torch.exp(-2 * ((x - num_inference_steps / 2) / num_inference_steps) ** 2) + y_shifted = y - y.min() + bsmntw_weighing = y_shifted * (num_inference_steps / y_shifted.sum()) + self.linear_timesteps_weights = bsmntw_weighing + self.training = True + else: + self.training = False + + + def step(self, model_output, timestep, sample, to_final=False, **kwargs): + if isinstance(timestep, torch.Tensor): + timestep = timestep.cpu() + timestep_id = torch.argmin((self.timesteps - timestep).abs()) + sigma = self.sigmas[timestep_id] + if to_final or timestep_id + 1 >= len(self.timesteps): + sigma_ = 1 if (self.inverse_timesteps or self.reverse_sigmas) else 0 + else: + sigma_ = self.sigmas[timestep_id + 1] + prev_sample = sample + model_output * (sigma_ - sigma) + return prev_sample + + + def return_to_timestep(self, timestep, sample, sample_stablized): + if isinstance(timestep, torch.Tensor): + timestep = timestep.cpu() + timestep_id = torch.argmin((self.timesteps - timestep).abs()) + sigma = self.sigmas[timestep_id] + model_output = (sample - sample_stablized) / sigma + return model_output + + + # def add_noise(self, original_samples, noise, timestep): + # if isinstance(timestep, torch.Tensor): + # timestep = timestep.cpu() + # timestep_id = torch.argmin((self.timesteps - timestep).abs()) + # sigma = self.sigmas[timestep_id] + # sample = (1 - sigma) * original_samples + sigma * noise + # return sample + + def add_noise(self, original_samples, noise, timestep): + if isinstance(timestep, torch.Tensor): + timestep = timestep.cpu() + timestep_id = torch.argmin((self.timesteps.unsqueeze(1) - timestep.unsqueeze(0)).abs(), dim = 0) + sigma = self.sigmas[timestep_id].to(device=original_samples.device, dtype=original_samples.dtype) + while len(sigma.shape) < len(original_samples.shape): + sigma = sigma.unsqueeze(-1) + sample = (1 - sigma) * original_samples + sigma * noise + return sample + + def training_target(self, sample, noise, timestep): + target = noise - sample + return target + + + def training_weight(self, timestep): + # timestep_id = torch.argmin((self.timesteps - timestep.to(self.timesteps.device)).abs()) + timestep_id = torch.argmin((self.timesteps.unsqueeze(1) - timestep.unsqueeze(0).to(self.timesteps.device)).abs(), dim = 0) + weights = self.linear_timesteps_weights[timestep_id] + return weights \ No newline at end of file diff --git a/flagscale/train/models/dreamzero/modules/flow_unipc_multistep_scheduler.py b/flagscale/train/models/dreamzero/modules/flow_unipc_multistep_scheduler.py new file mode 100644 index 0000000000..df7ea56e27 --- /dev/null +++ b/flagscale/train/models/dreamzero/modules/flow_unipc_multistep_scheduler.py @@ -0,0 +1,680 @@ +# Copied from https://github.com/huggingface/diffusers/blob/v0.31.0/src/diffusers/schedulers/scheduling_unipc_multistep.py +# Convert unipc for flow matching +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. + +import math +from typing import List, Optional, Tuple, Union + +import numpy as np +import torch +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.schedulers.scheduling_utils import ( + KarrasDiffusionSchedulers, + SchedulerMixin, + SchedulerOutput, +) + + +class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin): + """ + `UniPCMultistepScheduler` is a training-free framework designed for the fast sampling of diffusion models. + + This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic + methods the library implements for all schedulers such as loading and saving. + + Args: + num_train_timesteps (`int`, defaults to 1000): + The number of diffusion steps to train the model. + solver_order (`int`, default `2`): + The UniPC order which can be any positive integer. The effective order of accuracy is `solver_order + 1` + due to the UniC. It is recommended to use `solver_order=2` for guided sampling, and `solver_order=3` for + unconditional sampling. + prediction_type (`str`, defaults to "flow_prediction"): + Prediction type of the scheduler function; must be `flow_prediction` for this scheduler, which predicts + the flow of the diffusion process. + thresholding (`bool`, defaults to `False`): + Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such + as Stable Diffusion. + dynamic_thresholding_ratio (`float`, defaults to 0.995): + The ratio for the dynamic thresholding method. Valid only when `thresholding=True`. + sample_max_value (`float`, defaults to 1.0): + The threshold value for dynamic thresholding. Valid only when `thresholding=True` and `predict_x0=True`. + predict_x0 (`bool`, defaults to `True`): + Whether to use the updating algorithm on the predicted x0. + solver_type (`str`, default `bh2`): + Solver type for UniPC. It is recommended to use `bh1` for unconditional sampling when steps < 10, and `bh2` + otherwise. + lower_order_final (`bool`, default `True`): + Whether to use lower-order solvers in the final steps. Only valid for < 15 inference steps. This can + stabilize the sampling of DPMSolver for steps < 15, especially for steps <= 10. + disable_corrector (`list`, default `[]`): + Decides which step to disable the corrector to mitigate the misalignment between `epsilon_theta(x_t, c)` + and `epsilon_theta(x_t^c, c)` which can influence convergence for a large guidance scale. Corrector is + usually disabled during the first few steps. + solver_p (`SchedulerMixin`, default `None`): + Any other scheduler that if specified, the algorithm becomes `solver_p + UniC`. + use_karras_sigmas (`bool`, *optional*, defaults to `False`): + Whether to use Karras sigmas for step sizes in the noise schedule during the sampling process. If `True`, + the sigmas are determined according to a sequence of noise levels {σi}. + use_exponential_sigmas (`bool`, *optional*, defaults to `False`): + Whether to use exponential sigmas for step sizes in the noise schedule during the sampling process. + timestep_spacing (`str`, defaults to `"linspace"`): + The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and + Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information. + steps_offset (`int`, defaults to 0): + An offset added to the inference steps, as required by some model families. + final_sigmas_type (`str`, defaults to `"zero"`): + The final `sigma` value for the noise schedule during the sampling process. If `"sigma_min"`, the final + sigma is the same as the last sigma in the training schedule. If `zero`, the final sigma is set to 0. + """ + + _compatibles = [e.name for e in KarrasDiffusionSchedulers] + order = 1 + + @register_to_config + def __init__( + self, + num_train_timesteps: int = 1000, + solver_order: int = 2, + prediction_type: str = "flow_prediction", + shift: Optional[float] = 1.0, + use_dynamic_shifting=False, + thresholding: bool = False, + dynamic_thresholding_ratio: float = 0.995, + sample_max_value: float = 1.0, + predict_x0: bool = True, + solver_type: str = "bh2", + lower_order_final: bool = True, + disable_corrector: List[int] = [], + solver_p: SchedulerMixin = None, + timestep_spacing: str = "linspace", + steps_offset: int = 0, + final_sigmas_type: Optional[str] = "zero", # "zero", "sigma_min" + ): + + if solver_type not in ["bh1", "bh2"]: + if solver_type in ["midpoint", "heun", "logrho"]: + self.register_to_config(solver_type="bh2") + else: + raise NotImplementedError( + f"{solver_type} is not implemented for {self.__class__}") + + self.predict_x0 = predict_x0 + # setable values + self.num_inference_steps = None + alphas = np.linspace(1, 1 / num_train_timesteps, + num_train_timesteps)[::-1].copy() + sigmas = 1.0 - alphas + sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32, device='cuda') + + if not use_dynamic_shifting: + # when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution + assert shift is not None, "shift must be not None when use_dynamic_shifting is False" + sigmas = shift * sigmas / (1 + (shift - 1) * sigmas) + + self.sigmas = sigmas + self.timesteps = sigmas * num_train_timesteps + + self.model_outputs = [None] * solver_order + self.timestep_list = [None] * solver_order + self.lower_order_nums = 0 + self.disable_corrector = disable_corrector + self.solver_p = solver_p + self.last_sample = None + + self.sigma_min = self.sigmas[-1].item() + self.sigma_max = self.sigmas[0].item() + + # Modified from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler.set_timesteps + def set_timesteps( + self, + num_inference_steps: int, + device: str | torch.device = None, + sigmas: np.ndarray[float] | None = None, + mu: float | None = None, + shift: float | None = None, + ): + """ + Sets the discrete timesteps used for the diffusion chain (to be run before inference). + Args: + num_inference_steps (`int`): + Total number of the spacing of the time steps. + device (`str` or `torch.device`, *optional*): + The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. + """ + self.num_inference_steps = num_inference_steps + + if self.config.use_dynamic_shifting and mu is None: + raise ValueError( + " you have to pass a value for `mu` when `use_dynamic_shifting` is set to be `True`" + ) + + if sigmas is None: + sigmas = np.linspace( + self.sigma_max, + self.sigma_min, + num_inference_steps + 1, + ).copy()[:-1] + + if self.config.use_dynamic_shifting: + assert mu is not None + sigmas = self.time_shift(mu, 1.0, sigmas) + else: + if shift is None: + shift = self.config.shift + assert isinstance(shift, float) + assert sigmas is not None + sigmas = shift * sigmas / (1 + (shift - 1) * sigmas) + + if self.config.final_sigmas_type == "sigma_min": + sigma_last = ((1 - self.alphas_cumprod[0]) / self.alphas_cumprod[0])**0.5 + elif self.config.final_sigmas_type == "zero": + sigma_last = 0 + else: + raise ValueError( + f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}" + ) + + timesteps = sigmas * self.config.num_train_timesteps + sigmas = np.concatenate([sigmas, [sigma_last]]).astype(np.float32) + + self.sigmas = torch.from_numpy(sigmas).to(device=device) + self.timesteps = torch.from_numpy(timesteps).to(device=device, dtype=torch.int64) + + self.model_outputs = [None] * self.config.solver_order + self.lower_order_nums = 0 + self.last_sample = None + if self.solver_p: + self.solver_p.set_timesteps(self.num_inference_steps, device=device) + + # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample + def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor: + """ + "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the + prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by + s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing + pixels from saturation at each step. We find that dynamic thresholding results in significantly better + photorealism as well as better image-text alignment, especially when using very large guidance weights." + + https://arxiv.org/abs/2205.11487 + """ + dtype = sample.dtype + batch_size, channels, *remaining_dims = sample.shape + + if dtype not in (torch.float32, torch.float64): + sample = sample.float( + ) # upcast for quantile calculation, and clamp not implemented for cpu half + + # Flatten sample for doing quantile calculation along each image + sample = sample.reshape(batch_size, channels * np.prod(remaining_dims)) + + abs_sample = sample.abs() # "a certain percentile absolute pixel value" + + s = torch.quantile( + abs_sample, self.config.dynamic_thresholding_ratio, dim=1) + s = torch.clamp( + s, min=1, max=self.config.sample_max_value + ) # When clamped to min=1, equivalent to standard clipping to [-1, 1] + s = s.unsqueeze( + 1) # (batch_size, 1) because clamp will broadcast along dim=0 + sample = torch.clamp( + sample, -s, s + ) / s # "we threshold xt0 to the range [-s, s] and then divide by s" + + sample = sample.reshape(batch_size, channels, *remaining_dims) + sample = sample.to(dtype) + + return sample + + # Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler._sigma_to_t + def _sigma_to_t(self, sigma): + return sigma * self.config.num_train_timesteps + + @staticmethod + def _sigma_to_alpha_sigma_t(sigma): + return 1 - sigma, sigma + + # Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.set_timesteps + def time_shift(self, mu: float, sigma: float, t: torch.Tensor): + return math.exp(mu) / (math.exp(mu) + (1 / t - 1)**sigma) + + def convert_model_output( + self, + model_output: torch.Tensor, + sample: torch.Tensor, + step_index: int, + ) -> torch.Tensor: + r""" + Convert the model output to the corresponding type the UniPC algorithm needs. + + Args: + model_output (`torch.Tensor`): + The direct output from the learned diffusion model. + timestep (`int`): + The current discrete timestep in the diffusion chain. + sample (`torch.Tensor`): + A current instance of a sample created by the diffusion process. + step_index (`int`): + The current timestep index. + + Returns: + `torch.Tensor`: + The converted model output. + """ + if self.predict_x0: + if self.config.prediction_type == "flow_prediction": + sigma_t = self.sigmas[step_index] + x0_pred = sample - sigma_t * model_output + else: + raise ValueError( + f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`," + " `v_prediction` or `flow_prediction` for the UniPCMultistepScheduler." + ) + + if self.config.thresholding: + x0_pred = self._threshold_sample(x0_pred) + + return x0_pred + else: + if self.config.prediction_type == "flow_prediction": + sigma_t = self.sigmas[step_index] + epsilon = sample - (1 - sigma_t) * model_output + else: + raise ValueError( + f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`," + " `v_prediction` or `flow_prediction` for the UniPCMultistepScheduler." + ) + + if self.config.thresholding: + sigma_t = self.sigmas[step_index] + x0_pred = sample - sigma_t * model_output + x0_pred = self._threshold_sample(x0_pred) + epsilon = model_output + x0_pred + + return epsilon + + @torch.compile(mode="reduce-overhead", fullgraph=True, dynamic=False) + def multistep_uni_p_bh_update( + self, + model_output: torch.Tensor, + sample: torch.Tensor, + order: int, + step_index: int, + ) -> torch.Tensor: + """ + One step for the UniP (B(h) version). Alternatively, `self.solver_p` is used if is specified. + + Args: + model_output (`torch.Tensor`): + The direct output from the learned diffusion model at the current timestep. + sample (`torch.Tensor`): + A current instance of a sample created by the diffusion process. + order (`int`): + The order of UniP at this timestep (corresponds to the *p* in UniPC-p). + step_index (`int`): + The current timestep index. + + Returns: + `torch.Tensor`: + The sample tensor at the previous timestep. + """ + model_output_list = self.model_outputs + + s0 = self.timestep_list[-1] + m0 = model_output_list[-1] + x = sample + + if self.solver_p: + x_t = self.solver_p.step(model_output, s0, x).prev_sample + return x_t + + sigma_t, sigma_s0 = self.sigmas[step_index + 1], self.sigmas[step_index] + alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) + alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) + + lambda_t = torch.log(alpha_t) - torch.log(sigma_t) + lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) + + h = lambda_t - lambda_s0 + + rks = [] + D1s = [] + for i in range(1, order): + si = step_index - i + mi = model_output_list[-(i + 1)] + alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si]) + lambda_si = torch.log(alpha_si) - torch.log(sigma_si) + rk = (lambda_si - lambda_s0) / h + rks.append(rk) + D1s.append((mi - m0) / rk) + + rks.append(torch.ones((), dtype=self.sigmas.dtype, device=self.sigmas.device)) + rks = torch.stack(rks, dim=0) + + R = [] + b = [] + + hh = -h if self.predict_x0 else h + h_phi_1 = torch.expm1(hh) # h\phi_1(h) = e^h - 1 + h_phi_k = h_phi_1 / hh - 1 + + factorial_i = 1 + + if self.config.solver_type == "bh1": + B_h = hh + elif self.config.solver_type == "bh2": + B_h = torch.expm1(hh) + else: + raise NotImplementedError() + + for i in range(1, order + 1): + R.append(torch.pow(rks, i - 1)) + b.append(h_phi_k * factorial_i / B_h) + factorial_i *= i + 1 + h_phi_k = h_phi_k / hh - 1 / factorial_i + + R = torch.stack(R, dim=0) + b = torch.stack(b, dim=0) + + if len(D1s) > 0: + D1s = torch.stack(D1s, dim=1) # (B, K) + # for order 2, we use a simplified version + if order == 2: + rhos_p = torch.full((1,), 0.5, dtype=x.dtype, device=self.sigmas.device) + else: + rhos_p = torch.linalg.solve_ex(R[:-1, :-1], b[:-1])[0].to(x.dtype) + else: + D1s = None + rhos_p = None + + if self.predict_x0: + x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0 + if D1s is not None: + pred_res = torch.einsum("k,bkc...->bc...", rhos_p, D1s) + else: + pred_res = 0 + x_t = x_t_ - alpha_t * B_h * pred_res + else: + x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0 + if D1s is not None: + pred_res = torch.einsum("k,bkc...->bc...", rhos_p, D1s) + else: + pred_res = 0 + x_t = x_t_ - sigma_t * B_h * pred_res + + x_t = x_t.to(x.dtype) + return x_t + + @torch.compile(mode="reduce-overhead", fullgraph=True, dynamic=False) + def multistep_uni_c_bh_update( + self, + this_model_output: torch.Tensor, + last_sample: torch.Tensor, + this_sample: torch.Tensor, + order: int, + step_index: int, + ) -> torch.Tensor: + """ + One step for the UniC (B(h) version). + + Args: + this_model_output (`torch.Tensor`): + The model outputs at `x_t`. + this_timestep (`int`): + The current timestep `t`. + last_sample (`torch.Tensor`): + The generated sample before the last predictor `x_{t-1}`. + this_sample (`torch.Tensor`): + The generated sample after the last predictor `x_{t}`. + order (`int`): + The `p` of UniC-p at this step. The effective order of accuracy should be `order + 1`. + step_index (`int`): + The current timestep index. + + Returns: + `torch.Tensor`: + The corrected sample tensor at the current timestep. + """ + model_output_list = self.model_outputs + + m0 = model_output_list[-1] + x = last_sample + x_t = this_sample + model_t = this_model_output + + sigma_t, sigma_s0 = self.sigmas[step_index], self.sigmas[step_index - 1] + alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t) + alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0) + + lambda_t = torch.log(alpha_t) - torch.log(sigma_t) + lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0) + + h = lambda_t - lambda_s0 + + rks = [] + D1s = [] + for i in range(1, order): + si = step_index - (i + 1) + mi = model_output_list[-(i + 1)] + alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si]) + lambda_si = torch.log(alpha_si) - torch.log(sigma_si) + rk = (lambda_si - lambda_s0) / h + rks.append(rk) + D1s.append((mi - m0) / rk) + + rks.append(torch.ones((), dtype=self.sigmas.dtype, device=self.sigmas.device)) + rks = torch.stack(rks, dim=0) + + R = [] + b = [] + + hh = -h if self.predict_x0 else h + h_phi_1 = torch.expm1(hh) # h\phi_1(h) = e^h - 1 + h_phi_k = h_phi_1 / hh - 1 + + factorial_i = 1 + + if self.config.solver_type == "bh1": + B_h = hh + elif self.config.solver_type == "bh2": + B_h = torch.expm1(hh) + else: + raise NotImplementedError() + + for i in range(1, order + 1): + R.append(torch.pow(rks, i - 1)) + b.append(h_phi_k * factorial_i / B_h) + factorial_i *= i + 1 + h_phi_k = h_phi_k / hh - 1 / factorial_i + + R = torch.stack(R, dim=0) + b = torch.stack(b, dim=0) + + if len(D1s) > 0: + D1s = torch.stack(D1s, dim=1) + else: + D1s = None + + # for order 1, we use a simplified version + if order == 1: + rhos_c = torch.full((1,), 0.5, dtype=x.dtype, device=self.sigmas.device) + else: + rhos_c = torch.linalg.solve_ex(R, b)[0].to(x.dtype) + + if self.predict_x0: + x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0 + if D1s is not None: + corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s) + else: + corr_res = 0 + D1_t = model_t - m0 + x_t = x_t_ - alpha_t * B_h * (corr_res + rhos_c[-1] * D1_t) + else: + x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0 + if D1s is not None: + corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s) + else: + corr_res = 0 + D1_t = model_t - m0 + x_t = x_t_ - sigma_t * B_h * (corr_res + rhos_c[-1] * D1_t) + x_t = x_t.to(x.dtype) + return x_t + + def index_for_timestep(self, timestep, schedule_timesteps): + indices = (schedule_timesteps == timestep).nonzero() + + # The sigma index that is taken for the **very** first `step` + # is always the second index (or the last index if there is only 1) + # This way we can ensure we don't accidentally skip a sigma in + # case we start in the middle of the denoising schedule (e.g. for image-to-image) + pos = 1 if len(indices) > 1 else 0 + + return indices[pos].item() + + def step( + self, + model_output: torch.Tensor, + timestep: torch.Tensor, + sample: torch.Tensor, + step_index: int, + return_dict: bool = True, + ) -> SchedulerOutput | tuple: + """ + Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with + the multistep UniPC. + + Args: + model_output (`torch.Tensor`): + The direct output from learned diffusion model. + timestep (`torch.Tensor`): + The current discrete timestep in the diffusion chain. + sample (`torch.Tensor`): + A current instance of a sample created by the diffusion process. + step_index (`int`): + The current timestep index. + return_dict (`bool`): + Whether or not to return a [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`. + + Returns: + [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`: + If return_dict is `True`, [`~schedulers.scheduling_utils.SchedulerOutput`] is returned, otherwise a + tuple is returned where the first element is the sample tensor. + + """ + if self.num_inference_steps is None: + raise ValueError( + "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler" + ) + + use_corrector = ( + step_index > 0 and + step_index - 1 not in self.disable_corrector and + self.last_sample is not None + ) + + model_output_convert = self.convert_model_output( + model_output=model_output, + sample=sample, + step_index=step_index, + ) + if use_corrector: + sample = self.multistep_uni_c_bh_update( + this_model_output=model_output_convert, + last_sample=self.last_sample, + this_sample=sample, + order=self.this_order, + step_index=step_index, + ) + # We must clone the outputs of a CUDA graph'd computation. + sample = sample.clone() + + for i in range(self.config.solver_order - 1): + self.model_outputs[i] = self.model_outputs[i + 1] + self.timestep_list[i] = self.timestep_list[i + 1] + + self.model_outputs[-1] = model_output_convert + self.timestep_list[-1] = timestep + + if self.config.lower_order_final: + this_order = min( + self.config.solver_order, + len(self.timesteps) - step_index, + ) + else: + this_order = self.config.solver_order + + # Warmup for multistep. + self.this_order = min(this_order, self.lower_order_nums + 1) + assert self.this_order > 0 + + self.last_sample = sample + # Pass the original non-converted model output, in case solver-p is used. + prev_sample = self.multistep_uni_p_bh_update( + model_output=model_output, + sample=sample, + order=self.this_order, + step_index=step_index, + ) + # We must clone the outputs of a CUDA graph'd computation. + prev_sample = prev_sample.clone() + + if self.lower_order_nums < self.config.solver_order: + self.lower_order_nums += 1 + + if not return_dict: + return (prev_sample,) + + return SchedulerOutput(prev_sample=prev_sample) + + def scale_model_input(self, sample: torch.Tensor, *args, + **kwargs) -> torch.Tensor: + """ + Ensures interchangeability with schedulers that need to scale the denoising model input depending on the + current timestep. + + Args: + sample (`torch.Tensor`): + The input sample. + + Returns: + `torch.Tensor`: + A scaled input sample. + """ + return sample + + # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.add_noise + def add_noise( + self, + original_samples: torch.Tensor, + noise: torch.Tensor, + timesteps: torch.IntTensor, + ) -> torch.Tensor: + # Make sure sigmas and timesteps have the same device and dtype as original_samples + sigmas = self.sigmas.to( + device=original_samples.device, dtype=original_samples.dtype) + if original_samples.device.type == "mps" and torch.is_floating_point(timesteps): + # mps does not support float64 + schedule_timesteps = self.timesteps.to( + original_samples.device, dtype=torch.float32) + timesteps = timesteps.to( + original_samples.device, dtype=torch.float32) + else: + schedule_timesteps = self.timesteps.to(original_samples.device) + timesteps = timesteps.to(original_samples.device) + + # begin_index is None when the scheduler is used for training or pipeline does not implement set_begin_index + step_indices = [ + self.index_for_timestep(t, schedule_timesteps) + for t in timesteps + ] + + sigma = sigmas[step_indices].flatten() + while len(sigma.shape) < len(original_samples.shape): + sigma = sigma.unsqueeze(-1) + + alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma) + noisy_samples = alpha_t * original_samples + sigma_t * noise + return noisy_samples + + def __len__(self): + return self.config.num_train_timesteps diff --git a/flagscale/train/models/dreamzero/modules/utils.py b/flagscale/train/models/dreamzero/modules/utils.py new file mode 100644 index 0000000000..0d58e4e11d --- /dev/null +++ b/flagscale/train/models/dreamzero/modules/utils.py @@ -0,0 +1,182 @@ +import torch, os +from safetensors import safe_open +from contextlib import contextmanager +import hashlib + +@contextmanager +def init_weights_on_device(device = torch.device("meta"), include_buffers :bool = False): + + old_register_parameter = torch.nn.Module.register_parameter + if include_buffers: + old_register_buffer = torch.nn.Module.register_buffer + + def register_empty_parameter(module, name, param): + old_register_parameter(module, name, param) + if param is not None: + param_cls = type(module._parameters[name]) + kwargs = module._parameters[name].__dict__ + kwargs["requires_grad"] = param.requires_grad + module._parameters[name] = param_cls(module._parameters[name].to(device), **kwargs) + + def register_empty_buffer(module, name, buffer, persistent=True): + old_register_buffer(module, name, buffer, persistent=persistent) + if buffer is not None: + module._buffers[name] = module._buffers[name].to(device) + + def patch_tensor_constructor(fn): + def wrapper(*args, **kwargs): + kwargs["device"] = device + return fn(*args, **kwargs) + + return wrapper + + if include_buffers: + tensor_constructors_to_patch = { + torch_function_name: getattr(torch, torch_function_name) + for torch_function_name in ["empty", "zeros", "ones", "full"] + } + else: + tensor_constructors_to_patch = {} + + try: + torch.nn.Module.register_parameter = register_empty_parameter + if include_buffers: + torch.nn.Module.register_buffer = register_empty_buffer + for torch_function_name in tensor_constructors_to_patch.keys(): + setattr(torch, torch_function_name, patch_tensor_constructor(getattr(torch, torch_function_name))) + yield + finally: + torch.nn.Module.register_parameter = old_register_parameter + if include_buffers: + torch.nn.Module.register_buffer = old_register_buffer + for torch_function_name, old_torch_function in tensor_constructors_to_patch.items(): + setattr(torch, torch_function_name, old_torch_function) + +def load_state_dict_from_folder(file_path, torch_dtype=None): + state_dict = {} + for file_name in os.listdir(file_path): + if "." in file_name and file_name.split(".")[-1] in [ + "safetensors", "bin", "ckpt", "pth", "pt" + ]: + state_dict.update(load_state_dict(os.path.join(file_path, file_name), torch_dtype=torch_dtype)) + return state_dict + + +def load_state_dict(file_path, torch_dtype=None, device="cpu"): + if file_path.endswith(".safetensors"): + return load_state_dict_from_safetensors(file_path, torch_dtype=torch_dtype, device=device) + else: + return load_state_dict_from_bin(file_path, torch_dtype=torch_dtype, device=device) + + +def load_state_dict_from_safetensors(file_path, torch_dtype=None, device="cpu"): + state_dict = {} + with safe_open(file_path, framework="pt", device=device) as f: + for k in f.keys(): + state_dict[k] = f.get_tensor(k) + if torch_dtype is not None: + state_dict[k] = state_dict[k].to(torch_dtype) + return state_dict + + +def load_state_dict_from_bin(file_path, torch_dtype=None, device="cpu"): + state_dict = torch.load(file_path, map_location=device, weights_only=True) + if torch_dtype is not None: + for i in state_dict: + if isinstance(state_dict[i], torch.Tensor): + state_dict[i] = state_dict[i].to(torch_dtype) + return state_dict + + +def search_for_embeddings(state_dict): + embeddings = [] + for k in state_dict: + if isinstance(state_dict[k], torch.Tensor): + embeddings.append(state_dict[k]) + elif isinstance(state_dict[k], dict): + embeddings += search_for_embeddings(state_dict[k]) + return embeddings + + +def search_parameter(param, state_dict): + for name, param_ in state_dict.items(): + if param.numel() == param_.numel(): + if param.shape == param_.shape: + if torch.dist(param, param_) < 1e-3: + return name + else: + if torch.dist(param.flatten(), param_.flatten()) < 1e-3: + return name + return None + + +def build_rename_dict(source_state_dict, target_state_dict, split_qkv=False): + matched_keys = set() + with torch.no_grad(): + for name in source_state_dict: + rename = search_parameter(source_state_dict[name], target_state_dict) + if rename is not None: + print(f'"{name}": "{rename}",') + matched_keys.add(rename) + elif split_qkv and len(source_state_dict[name].shape)>=1 and source_state_dict[name].shape[0]%3==0: + length = source_state_dict[name].shape[0] // 3 + rename = [] + for i in range(3): + rename.append(search_parameter(source_state_dict[name][i*length: i*length+length], target_state_dict)) + if None not in rename: + print(f'"{name}": {rename},') + for rename_ in rename: + matched_keys.add(rename_) + for name in target_state_dict: + if name not in matched_keys: + print("Cannot find", name, target_state_dict[name].shape) + + +def search_for_files(folder, extensions): + files = [] + if os.path.isdir(folder): + for file in sorted(os.listdir(folder)): + files += search_for_files(os.path.join(folder, file), extensions) + elif os.path.isfile(folder): + for extension in extensions: + if folder.endswith(extension): + files.append(folder) + break + return files + + +def convert_state_dict_keys_to_single_str(state_dict, with_shape=True): + keys = [] + for key, value in state_dict.items(): + if isinstance(key, str): + if isinstance(value, torch.Tensor): + if with_shape: + shape = "_".join(map(str, list(value.shape))) + keys.append(key + ":" + shape) + keys.append(key) + elif isinstance(value, dict): + keys.append(key + "|" + convert_state_dict_keys_to_single_str(value, with_shape=with_shape)) + keys.sort() + keys_str = ",".join(keys) + return keys_str + + +def split_state_dict_with_prefix(state_dict): + keys = sorted([key for key in state_dict if isinstance(key, str)]) + prefix_dict = {} + for key in keys: + prefix = key if "." not in key else key.split(".")[0] + if prefix not in prefix_dict: + prefix_dict[prefix] = [] + prefix_dict[prefix].append(key) + state_dicts = [] + for prefix, keys in prefix_dict.items(): + sub_state_dict = {key: state_dict[key] for key in keys} + state_dicts.append(sub_state_dict) + return state_dicts + + +def hash_state_dict_keys(state_dict, with_shape=True): + keys_str = convert_state_dict_keys_to_single_str(state_dict, with_shape=with_shape) + keys_str = keys_str.encode(encoding="UTF-8") + return hashlib.md5(keys_str).hexdigest() \ No newline at end of file diff --git a/flagscale/train/models/dreamzero/modules/vram_management.py b/flagscale/train/models/dreamzero/modules/vram_management.py new file mode 100644 index 0000000000..7e2fcd3b3c --- /dev/null +++ b/flagscale/train/models/dreamzero/modules/vram_management.py @@ -0,0 +1,212 @@ +import torch, copy +from flagscale.train.models.dreamzero.modules.utils import init_weights_on_device + + +def cast_to(weight, dtype, device): + r = torch.empty_like(weight, dtype=dtype, device=device) + r.copy_(weight) + return r + + +class AutoTorchModule(torch.nn.Module): + def __init__(self): + super().__init__() + + def check_free_vram(self): + gpu_mem_state = torch.cuda.mem_get_info(self.computation_device) + used_memory = (gpu_mem_state[1] - gpu_mem_state[0]) / (1024 ** 3) + return used_memory < self.vram_limit + + def offload(self): + if self.state != 0: + self.to(dtype=self.offload_dtype, device=self.offload_device) + self.state = 0 + + def onload(self): + if self.state != 1: + self.to(dtype=self.onload_dtype, device=self.onload_device) + self.state = 1 + + def keep(self): + if self.state != 2: + self.to(dtype=self.computation_dtype, device=self.computation_device) + self.state = 2 + + +class AutoWrappedModule(AutoTorchModule): + def __init__(self, module: torch.nn.Module, offload_dtype, offload_device, onload_dtype, onload_device, computation_dtype, computation_device, vram_limit, **kwargs): + super().__init__() + self.module = module.to(dtype=offload_dtype, device=offload_device) + self.offload_dtype = offload_dtype + self.offload_device = offload_device + self.onload_dtype = onload_dtype + self.onload_device = onload_device + self.computation_dtype = computation_dtype + self.computation_device = computation_device + self.vram_limit = vram_limit + self.state = 0 + + def forward(self, *args, **kwargs): + if self.state == 2: + module = self.module + else: + if self.onload_dtype == self.computation_dtype and self.onload_device == self.computation_device: + module = self.module + elif self.vram_limit is not None and self.check_free_vram(): + self.keep() + module = self.module + else: + module = copy.deepcopy(self.module).to(dtype=self.computation_dtype, device=self.computation_device) + return module(*args, **kwargs) + + +class WanAutoCastLayerNorm(torch.nn.LayerNorm, AutoTorchModule): + def __init__(self, module: torch.nn.LayerNorm, offload_dtype, offload_device, onload_dtype, onload_device, computation_dtype, computation_device, vram_limit, **kwargs): + with init_weights_on_device(device=torch.device("meta")): + super().__init__(module.normalized_shape, eps=module.eps, elementwise_affine=module.elementwise_affine, bias=module.bias is not None, dtype=offload_dtype, device=offload_device) + self.weight = module.weight + self.bias = module.bias + self.offload_dtype = offload_dtype + self.offload_device = offload_device + self.onload_dtype = onload_dtype + self.onload_device = onload_device + self.computation_dtype = computation_dtype + self.computation_device = computation_device + self.vram_limit = vram_limit + self.state = 0 + + def forward(self, x, *args, **kwargs): + if self.state == 2: + weight, bias = self.weight, self.bias + else: + if self.onload_dtype == self.computation_dtype and self.onload_device == self.computation_device: + weight, bias = self.weight, self.bias + elif self.vram_limit is not None and self.check_free_vram(): + self.keep() + weight, bias = self.weight, self.bias + else: + weight = None if self.weight is None else cast_to(self.weight, self.computation_dtype, self.computation_device) + bias = None if self.bias is None else cast_to(self.bias, self.computation_dtype, self.computation_device) + with torch.amp.autocast(device_type=x.device.type): + x = torch.nn.functional.layer_norm(x.float(), self.normalized_shape, weight, bias, self.eps).type_as(x) + return x + + +class AutoWrappedLinear(torch.nn.Linear, AutoTorchModule): + def __init__(self, module: torch.nn.Linear, offload_dtype, offload_device, onload_dtype, onload_device, computation_dtype, computation_device, vram_limit, name="", **kwargs): + with init_weights_on_device(device=torch.device("meta")): + super().__init__(in_features=module.in_features, out_features=module.out_features, bias=module.bias is not None, dtype=offload_dtype, device=offload_device) + self.weight = module.weight + self.bias = module.bias + self.offload_dtype = offload_dtype + self.offload_device = offload_device + self.onload_dtype = onload_dtype + self.onload_device = onload_device + self.computation_dtype = computation_dtype + self.computation_device = computation_device + self.vram_limit = vram_limit + self.state = 0 + self.name = name + self.lora_A_weights = [] + self.lora_B_weights = [] + self.lora_merger = None + self.enable_fp8 = computation_dtype in [torch.float8_e4m3fn, torch.float8_e4m3fnuz] + + def fp8_linear( + self, + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + device = input.device + origin_dtype = input.dtype + origin_shape = input.shape + input = input.reshape(-1, origin_shape[-1]) + + x_max = torch.max(torch.abs(input), dim=-1, keepdim=True).values + fp8_max = 448.0 + # For float8_e4m3fnuz, the maximum representable value is half of that of e4m3fn. + # To avoid overflow and ensure numerical compatibility during FP8 computation, + # we scale down the input by 2.0 in advance. + # This scaling will be compensated later during the final result scaling. + if self.computation_dtype == torch.float8_e4m3fnuz: + fp8_max = fp8_max / 2.0 + scale_a = torch.clamp(x_max / fp8_max, min=1.0).float().to(device=device) + scale_b = torch.ones((weight.shape[0], 1)).to(device=device) + input = input / (scale_a + 1e-8) + input = input.to(self.computation_dtype) + weight = weight.to(self.computation_dtype) + bias = bias.to(torch.bfloat16) + + result = torch._scaled_mm( + input, + weight.T, + scale_a=scale_a, + scale_b=scale_b.T, + bias=bias, + out_dtype=origin_dtype, + ) + new_shape = origin_shape[:-1] + result.shape[-1:] + result = result.reshape(new_shape) + return result + + def forward(self, x, *args, **kwargs): + # VRAM management + if self.state == 2: + weight, bias = self.weight, self.bias + else: + if self.onload_dtype == self.computation_dtype and self.onload_device == self.computation_device: + weight, bias = self.weight, self.bias + elif self.vram_limit is not None and self.check_free_vram(): + self.keep() + weight, bias = self.weight, self.bias + else: + weight = cast_to(self.weight, self.computation_dtype, self.computation_device) + bias = None if self.bias is None else cast_to(self.bias, self.computation_dtype, self.computation_device) + + # Linear forward + if self.enable_fp8: + out = self.fp8_linear(x, weight, bias) + else: + out = torch.nn.functional.linear(x, weight, bias) + + # LoRA + if len(self.lora_A_weights) == 0: + # No LoRA + return out + elif self.lora_merger is None: + # Native LoRA inference + for lora_A, lora_B in zip(self.lora_A_weights, self.lora_B_weights): + out = out + x @ lora_A.T @ lora_B.T + else: + # LoRA fusion + lora_output = [] + for lora_A, lora_B in zip(self.lora_A_weights, self.lora_B_weights): + lora_output.append(x @ lora_A.T @ lora_B.T) + lora_output = torch.stack(lora_output) + out = self.lora_merger(out, lora_output) + return out + + +def enable_vram_management_recursively(model: torch.nn.Module, module_map: dict, module_config: dict, max_num_param=None, overflow_module_config: dict = None, total_num_param=0, vram_limit=None, name_prefix=""): + for name, module in model.named_children(): + layer_name = name if name_prefix == "" else name_prefix + "." + name + for source_module, target_module in module_map.items(): + if isinstance(module, source_module): + num_param = sum(p.numel() for p in module.parameters()) + if max_num_param is not None and total_num_param + num_param > max_num_param: + module_config_ = overflow_module_config + else: + module_config_ = module_config + module_ = target_module(module, **module_config_, vram_limit=vram_limit, name=layer_name) + setattr(model, name, module_) + total_num_param += num_param + break + else: + total_num_param = enable_vram_management_recursively(module, module_map, module_config, max_num_param, overflow_module_config, total_num_param, vram_limit=vram_limit, name_prefix=layer_name) + return total_num_param + + +def enable_vram_management(model: torch.nn.Module, module_map: dict, module_config: dict, max_num_param=None, overflow_module_config: dict = None, vram_limit=None): + enable_vram_management_recursively(model, module_map, module_config, max_num_param, overflow_module_config, total_num_param=0, vram_limit=vram_limit) + model.vram_management_enabled = True \ No newline at end of file diff --git a/flagscale/train/models/dreamzero/modules/wan2_1_attention.py b/flagscale/train/models/dreamzero/modules/wan2_1_attention.py new file mode 100644 index 0000000000..b6ac19a7ac --- /dev/null +++ b/flagscale/train/models/dreamzero/modules/wan2_1_attention.py @@ -0,0 +1,348 @@ +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. + +import contextlib +import torch +from torch.profiler import profile, ProfilerActivity +import time +from typing import Optional +import os + +try: + import flash_attn_interface + + def is_hopper_gpu(): + if not torch.cuda.is_available(): + return False + device_name = torch.cuda.get_device_name(0).lower() + return "h100" in device_name or "hopper" in device_name + FLASH_ATTN_3_AVAILABLE = is_hopper_gpu() +except ModuleNotFoundError: + FLASH_ATTN_3_AVAILABLE = False + +try: + import flash_attn + FLASH_ATTN_2_AVAILABLE = True +except ModuleNotFoundError: + FLASH_ATTN_2_AVAILABLE = False + +try: + import transformer_engine + from flagscale.train.models.dreamzero.modules.cudnn_attention import DotProductAttention + TRANSFORMER_ENGINE_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + TRANSFORMER_ENGINE_AVAILABLE = False + +import warnings + + +def _gpu_supports_flash_attention(): + """FlashAttention requires Ampere (compute capability 8.0) or newer.""" + if not (FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE): + return False + try: + if not torch.cuda.is_available(): + return False + cap = torch.cuda.get_device_capability() + return cap[0] >= 8 + except Exception: + return False + + +def _sdpa_attention_fallback( + q, k, v, + q_lens=None, + k_lens=None, + dropout_p=0., + softmax_scale=None, + q_scale=None, + causal=False, + dtype=torch.bfloat16, +): + """PyTorch SDPA fallback for GPUs that don't support FlashAttention (e.g. pre-Ampere).""" + if q_lens is not None or k_lens is not None: + warnings.warn( + 'Padding mask is disabled when using scaled_dot_product_attention on this GPU. ' + 'It can have a slight impact on quality.' + ) + q = q.transpose(1, 2).to(dtype) + k = k.transpose(1, 2).to(dtype) + v = v.transpose(1, 2).to(dtype) + if q_scale is not None: + q = q * q_scale + if softmax_scale is not None: + q = q * softmax_scale + out = torch.nn.functional.scaled_dot_product_attention( + q, k, v, attn_mask=None, is_causal=causal, dropout_p=dropout_p + ) + return out.transpose(1, 2).contiguous() + + +def flash_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + q_lens: Optional[torch.Tensor] = None, + k_lens: Optional[torch.Tensor] = None, + dropout_p: float = 0., + softmax_scale: Optional[float] = None, + q_scale: Optional[float] = None, + causal: bool = False, + window_size: Optional[tuple[int, int]] = None, + deterministic: bool = False, + dtype: torch.dtype = torch.bfloat16, + version: Optional[int] = None, +): + """ + q: [B, Lq, Nq, C1]. + k: [B, Lk, Nk, C1]. + v: [B, Lk, Nk, C2]. Nq must be divisible by Nk. + q_lens: [B]. + k_lens: [B]. + dropout_p: float. Dropout probability. + softmax_scale: float. The scaling of QK^T before applying softmax. + causal: bool. Whether to apply causal attention mask. + window_size: (left right). If not (-1, -1), apply sliding window local attention. + deterministic: bool. If True, slightly slower and uses more memory. + dtype: torch.dtype. Apply when dtype of q/k/v is not float16/bfloat16. + version: int. 2 for flash attention 2, 3 for flash attention 3. + + Returns: + x: [B, Lq, Nq, C2]. + """ + if window_size is None: + window_size = (-1, -1) + if version is None: + version = 3 + + half_dtypes = (torch.float16, torch.bfloat16) + assert dtype in half_dtypes + assert q.device.type == 'cuda' and q.size(-1) <= 256 + + # Use PyTorch SDPA on pre-Ampere GPUs (FlashAttention requires Ampere or newer) + if not _gpu_supports_flash_attention(): + return _sdpa_attention_fallback( + q, k, v, + q_lens=q_lens, + k_lens=k_lens, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + q_scale=q_scale, + causal=causal, + dtype=dtype, + ) + + # params + b, lq, lk, out_dtype = q.size(0), q.size(1), k.size(1), q.dtype + + def half(x): + return x if x.dtype in half_dtypes else x.to(dtype) + + # preprocess query + if q_lens is None: + q = half(q.flatten(0, 1)) + q_lens = torch.tensor([lq] * b, dtype=torch.int32, device=q.device) + else: + q = half(torch.cat([u[:v] for u, v in zip(q, q_lens)])) + + # preprocess key, value + if k_lens is None: + k = half(k.flatten(0, 1)) + v = half(v.flatten(0, 1)) + k_lens = torch.tensor([lk] * b, dtype=torch.int32, device=k.device) + else: + k = half(torch.cat([u[:v] for u, v in zip(k, k_lens)])) + v = half(torch.cat([u[:v] for u, v in zip(v, k_lens)])) + + q = q.to(v.dtype) + k = k.to(v.dtype) + + if q_scale is not None: + q = q * q_scale + + if version == 3 and not FLASH_ATTN_3_AVAILABLE: + warnings.warn( + 'Flash attention 3 is not available, use flash attention 2 instead.' + ) + zeros = torch.zeros([1], dtype=torch.int32, device=q.device) + cu_seqlens_q = torch.cat([zeros, q_lens]).cumsum(0).to(torch.int32) + cu_seqlens_k = torch.cat([zeros, k_lens]).cumsum(0).to(torch.int32) + + # apply attention + if version == 3 and FLASH_ATTN_3_AVAILABLE: + # Note: dropout_p, window_size are not supported in FA3 now. + x = flash_attn_interface.flash_attn_varlen_func( + q=q, + k=k, + v=v, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=lq, + max_seqlen_k=lk, + softmax_scale=softmax_scale, + causal=causal, + deterministic=deterministic)[0].unflatten(0, (b, lq)) + elif FLASH_ATTN_2_AVAILABLE: + x = flash_attn.flash_attn_varlen_func( + q=q, + k=k, + v=v, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=lq, + max_seqlen_k=lk, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + causal=causal, + window_size=window_size, + deterministic=deterministic).unflatten(0, (b, lq)) + else: + raise ValueError(f"Invalid version: {version}") + + # output + return x.type(out_dtype) + + +class AttentionModule(torch.nn.Module): + def __init__( + self, + num_heads: int, + head_dim: int, + dropout_p: float = 0., + softmax_scale: Optional[float] = None, + q_scale: Optional[float] = None, + causal: bool = False, + window_size: Optional[tuple[int, int]] = None, + deterministic: bool = False, + dtype: torch.dtype = torch.bfloat16, + backend: Optional[str] = None, + ): + super().__init__() + if backend is None: + backend = "torch" + + if os.getenv("ATTENTION_BACKEND") is not None: + backend = os.getenv("ATTENTION_BACKEND") + else: + backend = "FA2" + + # Check for TensorRT at runtime, not import time + if os.getenv("ENABLE_TENSORRT", "False").lower() == "true": + backend = "torch" + + # Fall back to FA backend if TE is specified but not available + if backend == "TE" and not TRANSFORMER_ENGINE_AVAILABLE: + print("Warning: Transformer Engine is not available. Falling back to FA2 backend.") + backend = "FA2" + + assert backend in ["torch", "FA2", "FA3", "TE", "torch_onnx"] + self.backend = backend + + if backend == "torch": + def _torch_impl(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + out_dtype = q.dtype + q = q.transpose(1, 2).to(dtype) + k = k.transpose(1, 2).to(dtype) + v = v.transpose(1, 2).to(dtype) + + out = torch.nn.functional.scaled_dot_product_attention( + q, k, v, + attn_mask=None, + is_causal=causal, + dropout_p=dropout_p, + scale=softmax_scale, + ) + + out = out.transpose(1, 2).contiguous() + return out.to(out_dtype) + self.attn_func = _torch_impl + + elif backend == "torch_onnx": + def _torch_onnx_impl(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + out_dtype = q.dtype + # use torch.nn.functional.scaled_dot_product_attention for tensorrt export + + # The input is (s, n, d), but sdpa needs (b, h, s, d). + # We add a batch dimension and transpose. + q = q.unsqueeze(0).transpose(1, 2).to(dtype) + k = k.unsqueeze(0).transpose(1, 2).to(dtype) + v = v.unsqueeze(0).transpose(1, 2).to(dtype) + + # Fix for ONNX export: repeat k and v to match q's batch size in cross-attention + if q.shape[0] != k.shape[0] and k.shape[0] == 1: + k = k.repeat(q.shape[0], 1, 1, 1) + v = v.repeat(q.shape[0], 1, 1, 1) + + out = torch.nn.functional.scaled_dot_product_attention( + q, k, v, + attn_mask=None, + is_causal=causal, + dropout_p=dropout_p, + scale=softmax_scale, + ) + + # Transpose back to (b, s, n, d) format. + out = out.transpose(1, 2).contiguous() + return out.to(out_dtype) + self.attn_func = _torch_onnx_impl + + elif backend == "TE" and TRANSFORMER_ENGINE_AVAILABLE: + self.attn_backend = DotProductAttention( + num_attention_heads=num_heads, + kv_channels=head_dim, + qkv_format="bshd", + attn_mask_type="causal" if causal else "no_mask", + window_size=window_size, + attention_dropout=dropout_p, + ) + + def _te_impl(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + out_dtype = q.dtype + return self.attn_backend( + query_layer=q.to(dtype), + key_layer=k.to(dtype), + value_layer=v.to(dtype), + ).to(out_dtype) + self.attn_func = _te_impl + + elif backend == "FA2" or backend == "FA3": + def _flash_attn_impl( + q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, + q_lens: Optional[torch.Tensor], k_lens: Optional[torch.Tensor], + ) -> torch.Tensor: + return flash_attention( + q=q, k=k, v=v, + q_lens=q_lens, k_lens=k_lens, + dropout_p=dropout_p, + softmax_scale=softmax_scale, + q_scale=q_scale, + causal=causal, + window_size=window_size, + deterministic=deterministic, + dtype=dtype, + version=3 if backend == "FA3" else 2, + ) + self.attn_func = _flash_attn_impl + + else: + raise ValueError(f"Invalid backend: {backend}") + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + q_lens: Optional[torch.Tensor] = None, + k_lens: Optional[torch.Tensor] = None, + ): + if ( + self.backend == "torch" or + self.backend == "torch_onnx" or + (self.backend == "TE" and TRANSFORMER_ENGINE_AVAILABLE) + ): + if q_lens is not None or k_lens is not None: + warnings.warn( + 'Padding mask is disabled when using scaled_dot_product_attention. It can have a significant impact on performance.' + ) + return self.attn_func(q, k, v) # type: ignore[call-arg] + else: + return self.attn_func(q, k, v, q_lens, k_lens) # type: ignore[call-arg] diff --git a/flagscale/train/models/dreamzero/modules/wan2_1_submodule.py b/flagscale/train/models/dreamzero/modules/wan2_1_submodule.py new file mode 100644 index 0000000000..2b7ba92323 --- /dev/null +++ b/flagscale/train/models/dreamzero/modules/wan2_1_submodule.py @@ -0,0 +1,905 @@ +# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. +import math + +import torch +import torch.nn as nn +import os +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.modeling_utils import ModelMixin +from einops import repeat +from flagscale.train.models.dreamzero.modules.attention import flash_attention + +__all__ = ['WanModel'] + +ENABLE_TENSORRT = os.getenv("ENABLE_TENSORRT", "False").lower() == "true" + +def sinusoidal_embedding_1d(dim: int, position: torch.Tensor) -> torch.Tensor: + # preprocess + assert dim % 2 == 0 + half = dim // 2 + position = position.type(torch.float64) + + # calculation + sinusoid = torch.outer( + position, torch.pow(10000, -torch.arange(half, dtype=position.dtype, device=position.device).div(half))) + x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1) + return x + + +def rope_params(max_seq_len, dim, theta=10000): + if ENABLE_TENSORRT: + return rope_params_no_polar(max_seq_len, dim, theta) + else: + return rope_params_polar(max_seq_len, dim, theta) + + +# @amp.autocast(enabled=False) +def rope_params_polar(max_seq_len: int, dim: int, theta: float = 10000) -> torch.Tensor: + assert dim % 2 == 0 + freqs = torch.outer( + torch.arange(max_seq_len), + 1.0 / torch.pow(theta, + torch.arange(0, dim, 2).to(torch.float64).div(dim))) + freqs = torch.polar(torch.ones_like(freqs), freqs) + return freqs + +def rope_params_no_polar(max_seq_len: int, dim: int, theta: float = 10000) -> torch.Tensor: + assert dim % 2 == 0 + inv_freq = 1.0 / torch.pow( + theta, + torch.arange(0, dim, 2).to(torch.float32) / dim + ) + t = torch.arange(max_seq_len, dtype=inv_freq.dtype) + freqs = torch.outer(t, inv_freq) + emb = torch.stack((freqs.cos(), freqs.sin()), dim=-1).flatten(-2) + return emb + +def rope_apply(x, grid_sizes, freqs): + if ENABLE_TENSORRT: + return rope_apply_no_polar(x, freqs) + else: + return rope_apply_polar(x, freqs) + +# @amp.autocast(enabled=False) +def rope_apply_polar(x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor: + B, seq_len, n, _ = x.shape + + # precompute multipliers + x = torch.view_as_complex( + x.to(torch.float64).reshape(B, seq_len, n, -1, 2) + ) + + # apply rotary embedding + freqs = freqs.unsqueeze(0) + x = torch.view_as_real(x * freqs).flatten(3) + return x + +def rope_apply_no_polar(x: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor: + B, seq_len, n, D = x.shape + + # Reshape freqs to be broadcastable: (1, seq_len, 1, D) + freqs = freqs.unsqueeze(0).unsqueeze(2) + + x0, x1 = x.chunk(2, dim=-1) + freqs_cos, freqs_sin = freqs.chunk(2, dim=-1) + + rotated_x0 = x0 * freqs_cos - x1 * freqs_sin + rotated_x1 = x1 * freqs_cos + x0 * freqs_sin + x_rotated = torch.cat((rotated_x0, rotated_x1), dim=-1) + return x_rotated + + +def rope_action_apply(x, freqs, freqs_action, freqs_state, action_register_length, num_action_per_block=32, num_state_per_block=1): + if ENABLE_TENSORRT: + return rope_action_apply_no_polar(x, freqs, freqs_action, freqs_state, action_register_length, num_action_per_block, num_state_per_block) + else: + return rope_action_apply_polar(x, freqs, freqs_action, freqs_state, action_register_length, num_action_per_block, num_state_per_block) + + +def rope_action_apply_no_polar( + x: torch.Tensor, + freqs: torch.Tensor, + freqs_action: torch.Tensor, + freqs_state: torch.Tensor, + action_register_length: int, + num_action_per_block: int = 32, + num_state_per_block: int = 1, +) -> torch.Tensor: + B, seq_len, n, D = x.shape + + if action_register_length is not None: + chunk_size = action_register_length // (num_action_per_block + num_state_per_block) + freqs_1d_action = freqs_action[:chunk_size * num_action_per_block] + freqs_1d_state = freqs_state[:chunk_size * num_state_per_block] + freqs = torch.cat([freqs, freqs_1d_action, freqs_1d_state], dim=0) + + # Reshape freqs to be broadcastable: (1, seq_len, 1, D) + freqs = freqs.unsqueeze(0).unsqueeze(2) + + x0, x1 = x.chunk(2, dim=-1) + freqs_cos, freqs_sin = freqs.chunk(2, dim=-1) + + rotated_x0 = x0 * freqs_cos - x1 * freqs_sin + rotated_x1 = x1 * freqs_cos + x0 * freqs_sin + x_rotated = torch.cat((rotated_x0, rotated_x1), dim=-1) + + return x_rotated + + +# @amp.autocast(enabled=False) +def rope_action_apply_polar( + x: torch.Tensor, + freqs: torch.Tensor, + freqs_action: torch.Tensor, + freqs_state: torch.Tensor, + action_register_length: int | None, + num_action_per_block: int | None = None, + num_state_per_block: int | None = None, +) -> torch.Tensor: + B, seq_len, n, _ = x.shape + + # precompute multipliers + x = torch.view_as_complex( + x.to(torch.float64).reshape(B, seq_len, n, -1, 2) + ) + + if action_register_length is not None: + assert num_action_per_block is not None + assert num_state_per_block is not None + + + chunk_size = action_register_length // (num_action_per_block + num_state_per_block) + + freqs_1d_action = freqs_action[:chunk_size * num_action_per_block].view(chunk_size * num_action_per_block, 1, -1) + freqs_1d_state = freqs_state[:chunk_size * num_state_per_block].view(chunk_size * num_state_per_block, 1, -1) + freqs = torch.cat([freqs, freqs_1d_action, freqs_1d_state], dim=0) + + # apply rotary embedding + freqs = freqs.unsqueeze(0) + x = torch.view_as_real(x * freqs).flatten(3) + return x + + +class WanRMSNorm(nn.Module): + + def __init__(self, dim, eps=1e-5): + super().__init__() + self.dim = dim + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x): + r""" + Args: + x(Tensor): Shape [B, L, C] + """ + return self._norm(x.float()).type_as(x) * self.weight + + def _norm(self, x): + return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps) + + +class WanLayerNorm(nn.LayerNorm): + + def __init__(self, dim, eps=1e-6, elementwise_affine=False): + super().__init__(dim, elementwise_affine=elementwise_affine, eps=eps) + + +class WanSelfAttention(nn.Module): + + def __init__(self, + dim, + num_heads, + window_size=(-1, -1), + qk_norm=True, + eps=1e-6): + assert dim % num_heads == 0 + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.window_size = window_size + self.qk_norm = qk_norm + self.eps = eps + + # layers + self.q = nn.Linear(dim, dim) + self.k = nn.Linear(dim, dim) + self.v = nn.Linear(dim, dim) + self.o = nn.Linear(dim, dim) + self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + + def forward(self, x, seq_lens, freqs): + r""" + Args: + x(Tensor): Shape [B, L, num_heads, C / num_heads] + seq_lens(Tensor): Shape [B] + freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2] + """ + b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim + + # query, key, value function + def qkv_fn(x): + q = self.norm_q(self.q(x)).view(b, s, n, d) + k = self.norm_k(self.k(x)).view(b, s, n, d) + v = self.v(x).view(b, s, n, d) + return q, k, v + + q, k, v = qkv_fn(x) + + x = flash_attention( + q=rope_apply(q, freqs), + k=rope_apply(k, freqs), + v=v, + k_lens=seq_lens, + window_size=self.window_size) + + # output + x = x.flatten(2) + x = self.o(x) + return x + + +class WanT2VCrossAttention(WanSelfAttention): + + def forward(self, x, context, context_lens, crossattn_cache=None): + r""" + Args: + x(Tensor): Shape [B, L1, C] + context(Tensor): Shape [B, L2, C] + context_lens(Tensor): Shape [B] + crossattn_cache (List[dict], *optional*): Contains the cached key and value tensors for context embedding. + """ + b, n, d = x.size(0), self.num_heads, self.head_dim + + # compute query, key, value + q = self.norm_q(self.q(x)).view(b, -1, n, d) + + if crossattn_cache is not None: + if not crossattn_cache["is_init"]: + crossattn_cache["is_init"] = True + k = self.norm_k(self.k(context)).view(b, -1, n, d) + v = self.v(context).view(b, -1, n, d) + crossattn_cache["k"] = k + crossattn_cache["v"] = v + else: + k = crossattn_cache["k"] + v = crossattn_cache["v"] + else: + k = self.norm_k(self.k(context)).view(b, -1, n, d) + v = self.v(context).view(b, -1, n, d) + + # compute attention + x = flash_attention(q, k, v, k_lens=context_lens) + + # output + x = x.flatten(2) + x = self.o(x) + return x + + +class WanGanCrossAttention(WanSelfAttention): + + def forward(self, x, context, crossattn_cache=None): + r""" + Args: + x(Tensor): Shape [B, L1, C] + context(Tensor): Shape [B, L2, C] + context_lens(Tensor): Shape [B] + crossattn_cache (List[dict], *optional*): Contains the cached key and value tensors for context embedding. + """ + b, n, d = x.size(0), self.num_heads, self.head_dim + + # compute query, key, value + qq = self.norm_q(self.q(context)).view(b, 1, -1, d) + + kk = self.norm_k(self.k(x)).view(b, -1, n, d) + vv = self.v(x).view(b, -1, n, d) + + # compute attention + x = flash_attention(qq, kk, vv) + + # output + x = x.flatten(2) + x = self.o(x) + return x + + +class WanI2VCrossAttention(WanSelfAttention): + + def __init__(self, + dim, + num_heads, + window_size=(-1, -1), + qk_norm=True, + eps=1e-6): + super().__init__(dim, num_heads, window_size, qk_norm, eps) + + self.k_img = nn.Linear(dim, dim) + self.v_img = nn.Linear(dim, dim) + # self.alpha = nn.Parameter(torch.zeros((1, ))) + self.norm_k_img = WanRMSNorm( + dim, eps=eps) if qk_norm else nn.Identity() + + def forward(self, x, context, crossattn_cache=None): + r""" + Args: + x(Tensor): Shape [B, L1, C] + context(Tensor): Shape [B, L2, C] + """ + context_img = context[:, :257] + context = context[:, 257:] + b, n, d = x.size(0), self.num_heads, self.head_dim + + q = self.norm_q(self.q(x)).view(b, -1, n, d) + + if crossattn_cache is not None: + if not crossattn_cache["is_init"]: + crossattn_cache["is_init"] = True + k = self.norm_k(self.k(context)).view(b, -1, n, d) + v = self.v(context).view(b, -1, n, d) + crossattn_cache["k"] = k + crossattn_cache["v"] = v + else: + k = crossattn_cache["k"] + v = crossattn_cache["v"] + else: + # compute query, key, value + k = self.norm_k(self.k(context)).view(b, -1, n, d) + v = self.v(context).view(b, -1, n, d) + x = flash_attention(q, k, v, k_lens=None) + + k_img = self.norm_k_img(self.k_img(context_img)).view(b, -1, n, d) + v_img = self.v_img(context_img).view(b, -1, n, d) + img_x = flash_attention(q, k_img, v_img, k_lens=None) + + # output + x = x.flatten(2) + img_x = img_x.flatten(2) + x = x + img_x + x = self.o(x) + return x + + +WAN_CROSSATTENTION_CLASSES = { + 't2v_cross_attn': WanT2VCrossAttention, + 'i2v_cross_attn': WanI2VCrossAttention, +} + + +class WanAttentionBlock(nn.Module): + + def __init__(self, + cross_attn_type, + dim, + ffn_dim, + num_heads, + window_size=(-1, -1), + qk_norm=True, + cross_attn_norm=False, + eps=1e-6): + super().__init__() + self.dim = dim + self.ffn_dim = ffn_dim + self.num_heads = num_heads + self.window_size = window_size + self.qk_norm = qk_norm + self.cross_attn_norm = cross_attn_norm + self.eps = eps + + # layers + self.norm1 = WanLayerNorm(dim, eps) + self.self_attn = WanSelfAttention(dim, num_heads, window_size, qk_norm, + eps) + self.norm3 = WanLayerNorm( + dim, eps, + elementwise_affine=True) if cross_attn_norm else nn.Identity() + self.cross_attn = WAN_CROSSATTENTION_CLASSES[cross_attn_type](dim, + num_heads, + (-1, -1), + qk_norm, + eps) + self.norm2 = WanLayerNorm(dim, eps) + self.ffn = nn.Sequential( + nn.Linear(dim, ffn_dim), nn.GELU(approximate='tanh'), + nn.Linear(ffn_dim, dim)) + + # modulation + self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + + def forward( + self, + x, + e, + seq_lens, + grid_sizes, + freqs, + context, + context_lens, + ): + r""" + Args: + x(Tensor): Shape [B, L, C] + e(Tensor): Shape [B, 6, C] + seq_lens(Tensor): Shape [B], length of each sequence in batch + grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W) + freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2] + """ + # assert e.dtype == torch.float32 + # with amp.autocast(dtype=torch.float32): + e = (self.modulation + e).chunk(6, dim=1) + # assert e[0].dtype == torch.float32 + + # self-attention + y = self.self_attn( + self.norm1(x) * (1 + e[1]) + e[0], seq_lens, grid_sizes, + freqs, + ) + # with amp.autocast(dtype=torch.float32): + x = x + y * e[2] + + # cross-attention & ffn function + def cross_attn_ffn(x, context, context_lens, e): + x = x + self.cross_attn(self.norm3(x), context, context_lens) + y = self.ffn(self.norm2(x) * (1 + e[4]) + e[3]) + # with amp.autocast(dtype=torch.float32): + x = x + y * e[5] + return x + + x = cross_attn_ffn(x, context, context_lens, e) + return x + + +class GanAttentionBlock(nn.Module): + + def __init__(self, + dim=1536, + ffn_dim=8192, + num_heads=12, + window_size=(-1, -1), + qk_norm=True, + cross_attn_norm=True, + eps=1e-6): + super().__init__() + self.dim = dim + self.ffn_dim = ffn_dim + self.num_heads = num_heads + self.window_size = window_size + self.qk_norm = qk_norm + self.cross_attn_norm = cross_attn_norm + self.eps = eps + + # layers + # self.norm1 = WanLayerNorm(dim, eps) + # self.self_attn = WanSelfAttention(dim, num_heads, window_size, qk_norm, + # eps) + self.norm3 = WanLayerNorm( + dim, eps, + elementwise_affine=True) if cross_attn_norm else nn.Identity() + + self.norm2 = WanLayerNorm(dim, eps) + self.ffn = nn.Sequential( + nn.Linear(dim, ffn_dim), nn.GELU(approximate='tanh'), + nn.Linear(ffn_dim, dim)) + + self.cross_attn = WanGanCrossAttention(dim, num_heads, + (-1, -1), + qk_norm, + eps) + + # modulation + # self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + + def forward( + self, + x, + context, + # seq_lens, + # grid_sizes, + # freqs, + # context, + # context_lens, + ): + r""" + Args: + x(Tensor): Shape [B, L, C] + e(Tensor): Shape [B, 6, C] + seq_lens(Tensor): Shape [B], length of each sequence in batch + grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W) + freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2] + """ + # assert e.dtype == torch.float32 + # with amp.autocast(dtype=torch.float32): + # e = (self.modulation + e).chunk(6, dim=1) + # assert e[0].dtype == torch.float32 + + # # self-attention + # y = self.self_attn( + # self.norm1(x) * (1 + e[1]) + e[0], seq_lens, grid_sizes, + # freqs) + # # with amp.autocast(dtype=torch.float32): + # x = x + y * e[2] + + # cross-attention & ffn function + def cross_attn_ffn(x, context): + token = context + self.cross_attn(self.norm3(x), context) + y = self.ffn(self.norm2(token)) + token # * (1 + e[4]) + e[3]) + # with amp.autocast(dtype=torch.float32): + # x = x + y * e[5] + return y + + x = cross_attn_ffn(x, context) + return x + + +class Head(nn.Module): + + def __init__(self, dim, out_dim, patch_size, eps=1e-6): + super().__init__() + self.dim = dim + self.out_dim = out_dim + self.patch_size = patch_size + self.eps = eps + + # layers + out_dim = math.prod(patch_size) * out_dim + self.norm = WanLayerNorm(dim, eps) + self.head = nn.Linear(dim, out_dim) + + # modulation + self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5) + + def forward(self, x, e): + r""" + Args: + x(Tensor): Shape [B, L1, C] + e(Tensor): Shape [B, C] + """ + # assert e.dtype == torch.float32 + # with amp.autocast(dtype=torch.float32): + e = (self.modulation + e.unsqueeze(1)).chunk(2, dim=1) + x = (self.head(self.norm(x) * (1 + e[1]) + e[0])) + return x + + +class MLPProj(torch.nn.Module): + + def __init__(self, in_dim, out_dim): + super().__init__() + + self.proj = torch.nn.Sequential( + torch.nn.LayerNorm(in_dim), torch.nn.Linear(in_dim, in_dim), + torch.nn.GELU(), torch.nn.Linear(in_dim, out_dim), + torch.nn.LayerNorm(out_dim)) + + def forward(self, image_embeds): + clip_extra_context_tokens = self.proj(image_embeds) + return clip_extra_context_tokens + + +class RegisterTokens(nn.Module): + def __init__(self, num_registers: int, dim: int): + super().__init__() + self.register_tokens = nn.Parameter(torch.randn(num_registers, dim) * 0.02) + self.rms_norm = WanRMSNorm(dim, eps=1e-6) + + def forward(self): + return self.rms_norm(self.register_tokens) + + def reset_parameters(self): + nn.init.normal_(self.register_tokens, std=0.02) + + +class WanModel(ModelMixin, ConfigMixin): + r""" + Wan diffusion backbone supporting both text-to-video and image-to-video. + """ + + ignore_for_config = [ + 'patch_size', 'cross_attn_norm', 'qk_norm', 'text_dim', 'window_size' + ] + _no_split_modules = ['WanAttentionBlock'] + _supports_gradient_checkpointing = True + + @register_to_config + def __init__(self, + model_type='t2v', + patch_size=(1, 2, 2), + text_len=512, + in_dim=16, + dim=2048, + ffn_dim=8192, + freq_dim=256, + text_dim=4096, + out_dim=16, + num_heads=16, + num_layers=32, + window_size=(-1, -1), + qk_norm=True, + cross_attn_norm=True, + eps=1e-6): + r""" + Initialize the diffusion model backbone. + + Args: + model_type (`str`, *optional*, defaults to 't2v'): + Model variant - 't2v' (text-to-video) or 'i2v' (image-to-video) + patch_size (`tuple`, *optional*, defaults to (1, 2, 2)): + 3D patch dimensions for video embedding (t_patch, h_patch, w_patch) + text_len (`int`, *optional*, defaults to 512): + Fixed length for text embeddings + in_dim (`int`, *optional*, defaults to 16): + Input video channels (C_in) + dim (`int`, *optional*, defaults to 2048): + Hidden dimension of the transformer + ffn_dim (`int`, *optional*, defaults to 8192): + Intermediate dimension in feed-forward network + freq_dim (`int`, *optional*, defaults to 256): + Dimension for sinusoidal time embeddings + text_dim (`int`, *optional*, defaults to 4096): + Input dimension for text embeddings + out_dim (`int`, *optional*, defaults to 16): + Output video channels (C_out) + num_heads (`int`, *optional*, defaults to 16): + Number of attention heads + num_layers (`int`, *optional*, defaults to 32): + Number of transformer blocks + window_size (`tuple`, *optional*, defaults to (-1, -1)): + Window size for local attention (-1 indicates global attention) + qk_norm (`bool`, *optional*, defaults to True): + Enable query/key normalization + cross_attn_norm (`bool`, *optional*, defaults to False): + Enable cross-attention normalization + eps (`float`, *optional*, defaults to 1e-6): + Epsilon value for normalization layers + """ + + super().__init__() + + assert model_type in ['t2v', 'i2v'] + self.model_type = model_type + + self.patch_size = patch_size + self.text_len = text_len + self.in_dim = in_dim + self.dim = dim + self.ffn_dim = ffn_dim + self.freq_dim = freq_dim + self.text_dim = text_dim + self.out_dim = out_dim + self.num_heads = num_heads + self.num_layers = num_layers + self.window_size = window_size + self.qk_norm = qk_norm + self.cross_attn_norm = cross_attn_norm + self.eps = eps + self.local_attn_size = 21 + + # embeddings + self.patch_embedding = nn.Conv3d( + in_dim, dim, kernel_size=patch_size, stride=patch_size) + self.text_embedding = nn.Sequential( + nn.Linear(text_dim, dim), nn.GELU(approximate='tanh'), + nn.Linear(dim, dim)) + + self.time_embedding = nn.Sequential( + nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim)) + self.time_projection = nn.Sequential( + nn.SiLU(), nn.Linear(dim, dim * 6)) + + # blocks + cross_attn_type = 't2v_cross_attn' if model_type == 't2v' else 'i2v_cross_attn' + self.blocks = nn.ModuleList([ + WanAttentionBlock(cross_attn_type, dim, ffn_dim, num_heads, + window_size, qk_norm, cross_attn_norm, eps) + for _ in range(num_layers) + ]) + + # head + self.head = Head(dim, out_dim, patch_size, eps) + + # buffers (don't use register_buffer otherwise dtype will be changed in to()) + assert (dim % num_heads) == 0 and (dim // num_heads) % 2 == 0 + d = dim // num_heads + self.freqs = [ + rope_params(1024, d - 4 * (d // 6)), + rope_params(1024, 2 * (d // 6)), + rope_params(1024, 2 * (d // 6)) + ] + + if model_type == 'i2v': + self.img_emb = MLPProj(1280, dim) + + # initialize weights + self.init_weights() + + self.gradient_checkpointing = False + + def _set_gradient_checkpointing(self, module, value=False): + self.gradient_checkpointing = value + + def forward( + self, + x, + t, + context, + seq_len, + classify_mode=False, + concat_time_embeddings=False, + register_tokens=None, + cls_pred_branch=None, + gan_ca_blocks=None, + clip_fea=None, + y=None, + ): + r""" + Forward pass through the diffusion model + + Args: + x (List[Tensor]): + List of input video tensors, each with shape [C_in, F, H, W] + t (Tensor): + Diffusion timesteps tensor of shape [B] + context (List[Tensor]): + List of text embeddings each with shape [L, C] + seq_len (`int`): + Maximum sequence length for positional encoding + clip_fea (Tensor, *optional*): + CLIP image features for image-to-video mode + y (List[Tensor], *optional*): + Conditional video inputs for image-to-video mode, same shape as x + + Returns: + List[Tensor]: + List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8] + """ + if self.model_type == 'i2v': + assert clip_fea is not None and y is not None + # params + device = self.patch_embedding.weight.device + if any(freqs_i.device != device for freqs_i in self.freqs): + self.freqs = [freqs_i.to(device) for freqs_i in self.freqs] + + if y is not None: + x = [torch.cat([u, v.to(dtype=u.dtype)], dim=0) for u, v in zip(x, y)] + + # embeddings + x = [self.patch_embedding(u.unsqueeze(0)) for u in x] + grid_sizes = torch.stack( + [torch.tensor(u.shape[2:], dtype=torch.long) for u in x]) + x = [u.flatten(2).transpose(1, 2) for u in x] + seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long) + assert seq_lens.max() <= seq_len + x = torch.cat([ + torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))], + dim=1) for u in x + ]) + + # time embeddings + # with amp.autocast(dtype=torch.float32): + e = self.time_embedding( + sinusoidal_embedding_1d(self.freq_dim, t).type_as(x)) + e0 = self.time_projection(e).unflatten(1, (6, self.dim)) + # assert e.dtype == torch.float32 and e0.dtype == torch.float32 + + # context + context_lens = None + context = self.text_embedding( + torch.stack([ + torch.cat( + [u, u.new_zeros(self.text_len - u.size(0), u.size(1))]) + for u in context + ])) + + if clip_fea is not None: + context_clip = self.img_emb(clip_fea) # bs x 257 x dim + context = torch.concat([context_clip, context], dim=1) + + # arguments + kwargs = dict( + e=e0, + seq_lens=seq_lens, + grid_sizes=grid_sizes, + freqs=self.freqs, + context=context, + context_lens=context_lens) + + def create_custom_forward(module): + def custom_forward(*inputs, **kwargs): + return module(*inputs, **kwargs) + return custom_forward + + # TODO: Tune the number of blocks for feature extraction + final_x = None + if classify_mode: + assert register_tokens is not None + assert gan_ca_blocks is not None + assert cls_pred_branch is not None + + final_x = [] + registers = repeat(register_tokens(), "n d -> b n d", b=x.shape[0]) + # x = torch.cat([registers, x], dim=1) + + gan_idx = 0 + for ii, block in enumerate(self.blocks): + if torch.is_grad_enabled() and self.gradient_checkpointing: + x = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + x, **kwargs, + use_reentrant=False, + ) + else: + x = block(x, **kwargs) + + if classify_mode and ii in [13, 21, 29]: + gan_token = registers[:, gan_idx: gan_idx + 1] + final_x.append(gan_ca_blocks[gan_idx](x, gan_token)) + gan_idx += 1 + + if classify_mode: + final_x = torch.cat(final_x, dim=1) + if concat_time_embeddings: + final_x = cls_pred_branch(torch.cat([final_x, 10 * e[:, None, :]], dim=1).view(final_x.shape[0], -1)) + else: + final_x = cls_pred_branch(final_x.view(final_x.shape[0], -1)) + + # head + x = self.head(x, e) + + # unpatchify + x = self.unpatchify(x, grid_sizes) + + if classify_mode: + return torch.stack(x), final_x + + return torch.stack(x) + + def unpatchify(self, x, grid_sizes, c=None): + r""" + Reconstruct video tensors from patch embeddings. + + Args: + x (List[Tensor]): + List of patchified features, each with shape [L, C_out * prod(patch_size)] + grid_sizes (Tensor): + Original spatial-temporal grid dimensions before patching, + shape [B, 3] (3 dimensions correspond to F_patches, H_patches, W_patches) + + Returns: + List[Tensor]: + Reconstructed video tensors with shape [C_out, F, H / 8, W / 8] + """ + + c = self.out_dim if c is None else c + out = [] + for u, v in zip(x, grid_sizes.tolist()): + u = u[:math.prod(v)].view(*v, *self.patch_size, c) + u = torch.einsum('fhwpqrc->cfphqwr', u) + u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)]) + out.append(u) + return out + + def init_weights(self): + r""" + Initialize model parameters using Xavier initialization. + """ + + # basic init + for m in self.modules(): + if isinstance(m, nn.Linear): + nn.init.xavier_uniform_(m.weight) + if m.bias is not None: + nn.init.zeros_(m.bias) + + # init embeddings + nn.init.xavier_uniform_(self.patch_embedding.weight.flatten(1)) + for m in self.text_embedding.modules(): + if isinstance(m, nn.Linear): + nn.init.normal_(m.weight, std=.02) + for m in self.time_embedding.modules(): + if isinstance(m, nn.Linear): + nn.init.normal_(m.weight, std=.02) + + # init output layer + nn.init.zeros_(self.head.head.weight) diff --git a/flagscale/train/models/dreamzero/modules/wan_video_camera_controller.py b/flagscale/train/models/dreamzero/modules/wan_video_camera_controller.py new file mode 100644 index 0000000000..8cbe658699 --- /dev/null +++ b/flagscale/train/models/dreamzero/modules/wan_video_camera_controller.py @@ -0,0 +1,202 @@ +import torch +import torch.nn as nn +import numpy as np +from einops import rearrange +import os +from typing_extensions import Literal + +class SimpleAdapter(nn.Module): + def __init__(self, in_dim, out_dim, kernel_size, stride, num_residual_blocks=1): + super(SimpleAdapter, self).__init__() + + # Pixel Unshuffle: reduce spatial dimensions by a factor of 8 + self.pixel_unshuffle = nn.PixelUnshuffle(downscale_factor=8) + + # Convolution: reduce spatial dimensions by a factor + # of 2 (without overlap) + self.conv = nn.Conv2d(in_dim * 64, out_dim, kernel_size=kernel_size, stride=stride, padding=0) + + # Residual blocks for feature extraction + self.residual_blocks = nn.Sequential( + *[ResidualBlock(out_dim) for _ in range(num_residual_blocks)] + ) + + def forward(self, x): + # Reshape to merge the frame dimension into batch + bs, c, f, h, w = x.size() + x = x.permute(0, 2, 1, 3, 4).contiguous().view(bs * f, c, h, w) + + # Pixel Unshuffle operation + x_unshuffled = self.pixel_unshuffle(x) + + # Convolution operation + x_conv = self.conv(x_unshuffled) + + # Feature extraction with residual blocks + out = self.residual_blocks(x_conv) + + # Reshape to restore original bf dimension + out = out.view(bs, f, out.size(1), out.size(2), out.size(3)) + + # Permute dimensions to reorder (if needed), e.g., swap channels and feature frames + out = out.permute(0, 2, 1, 3, 4) + + return out + + def process_camera_coordinates( + self, + direction: Literal["Left", "Right", "Up", "Down", "LeftUp", "LeftDown", "RightUp", "RightDown"], + length: int, + height: int, + width: int, + speed: float = 1/54, + origin=(0, 0.532139961, 0.946026558, 0.5, 0.5, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0) + ): + if origin is None: + origin = (0, 0.532139961, 0.946026558, 0.5, 0.5, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0) + coordinates = generate_camera_coordinates(direction, length, speed, origin) + plucker_embedding = process_pose_file(coordinates, width, height) + return plucker_embedding + + + +class ResidualBlock(nn.Module): + def __init__(self, dim): + super(ResidualBlock, self).__init__() + self.conv1 = nn.Conv2d(dim, dim, kernel_size=3, padding=1) + self.relu = nn.ReLU(inplace=True) + self.conv2 = nn.Conv2d(dim, dim, kernel_size=3, padding=1) + + def forward(self, x): + residual = x + out = self.relu(self.conv1(x)) + out = self.conv2(out) + out += residual + return out + +class Camera(object): + """Copied from https://github.com/hehao13/CameraCtrl/blob/main/inference.py + """ + def __init__(self, entry): + fx, fy, cx, cy = entry[1:5] + self.fx = fx + self.fy = fy + self.cx = cx + self.cy = cy + w2c_mat = np.array(entry[7:]).reshape(3, 4) + w2c_mat_4x4 = np.eye(4) + w2c_mat_4x4[:3, :] = w2c_mat + self.w2c_mat = w2c_mat_4x4 + self.c2w_mat = np.linalg.inv(w2c_mat_4x4) + +def get_relative_pose(cam_params): + """Copied from https://github.com/hehao13/CameraCtrl/blob/main/inference.py + """ + abs_w2cs = [cam_param.w2c_mat for cam_param in cam_params] + abs_c2ws = [cam_param.c2w_mat for cam_param in cam_params] + cam_to_origin = 0 + target_cam_c2w = np.array([ + [1, 0, 0, 0], + [0, 1, 0, -cam_to_origin], + [0, 0, 1, 0], + [0, 0, 0, 1] + ]) + abs2rel = target_cam_c2w @ abs_w2cs[0] + ret_poses = [target_cam_c2w, ] + [abs2rel @ abs_c2w for abs_c2w in abs_c2ws[1:]] + ret_poses = np.array(ret_poses, dtype=np.float32) + return ret_poses + +def custom_meshgrid(*args): + # torch>=2.0.0 only + return torch.meshgrid(*args, indexing='ij') + + +def ray_condition(K, c2w, H, W, device): + """Copied from https://github.com/hehao13/CameraCtrl/blob/main/inference.py + """ + # c2w: B, V, 4, 4 + # K: B, V, 4 + + B = K.shape[0] + + j, i = custom_meshgrid( + torch.linspace(0, H - 1, H, device=device, dtype=c2w.dtype), + torch.linspace(0, W - 1, W, device=device, dtype=c2w.dtype), + ) + i = i.reshape([1, 1, H * W]).expand([B, 1, H * W]) + 0.5 # [B, HxW] + j = j.reshape([1, 1, H * W]).expand([B, 1, H * W]) + 0.5 # [B, HxW] + + fx, fy, cx, cy = K.chunk(4, dim=-1) # B,V, 1 + + zs = torch.ones_like(i) # [B, HxW] + xs = (i - cx) / fx * zs + ys = (j - cy) / fy * zs + zs = zs.expand_as(ys) + + directions = torch.stack((xs, ys, zs), dim=-1) # B, V, HW, 3 + directions = directions / directions.norm(dim=-1, keepdim=True) # B, V, HW, 3 + + rays_d = directions @ c2w[..., :3, :3].transpose(-1, -2) # B, V, 3, HW + rays_o = c2w[..., :3, 3] # B, V, 3 + rays_o = rays_o[:, :, None].expand_as(rays_d) # B, V, 3, HW + # c2w @ dirctions + rays_dxo = torch.linalg.cross(rays_o, rays_d) + plucker = torch.cat([rays_dxo, rays_d], dim=-1) + plucker = plucker.reshape(B, c2w.shape[1], H, W, 6) # B, V, H, W, 6 + # plucker = plucker.permute(0, 1, 4, 2, 3) + return plucker + + +def process_pose_file(cam_params, width=672, height=384, original_pose_width=1280, original_pose_height=720, device='cpu', return_poses=False): + if return_poses: + return cam_params + else: + cam_params = [Camera(cam_param) for cam_param in cam_params] + + sample_wh_ratio = width / height + pose_wh_ratio = original_pose_width / original_pose_height # Assuming placeholder ratios, change as needed + + if pose_wh_ratio > sample_wh_ratio: + resized_ori_w = height * pose_wh_ratio + for cam_param in cam_params: + cam_param.fx = resized_ori_w * cam_param.fx / width + else: + resized_ori_h = width / pose_wh_ratio + for cam_param in cam_params: + cam_param.fy = resized_ori_h * cam_param.fy / height + + intrinsic = np.asarray([[cam_param.fx * width, + cam_param.fy * height, + cam_param.cx * width, + cam_param.cy * height] + for cam_param in cam_params], dtype=np.float32) + + K = torch.as_tensor(intrinsic)[None] # [1, 1, 4] + c2ws = get_relative_pose(cam_params) # Assuming this function is defined elsewhere + c2ws = torch.as_tensor(c2ws)[None] # [1, n_frame, 4, 4] + plucker_embedding = ray_condition(K, c2ws, height, width, device=device)[0].permute(0, 3, 1, 2).contiguous() # V, 6, H, W + plucker_embedding = plucker_embedding[None] + plucker_embedding = rearrange(plucker_embedding, "b f c h w -> b f h w c")[0] + return plucker_embedding + + + +def generate_camera_coordinates( + direction: Literal["Left", "Right", "Up", "Down", "LeftUp", "LeftDown", "RightUp", "RightDown"], + length: int, + speed: float = 1/54, + origin=(0, 0.532139961, 0.946026558, 0.5, 0.5, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0) +): + coordinates = [list(origin)] + while len(coordinates) < length: + coor = coordinates[-1].copy() + if "Left" in direction: + coor[9] += speed + if "Right" in direction: + coor[9] -= speed + if "Up" in direction: + coor[13] += speed + if "Down" in direction: + coor[13] -= speed + coordinates.append(coor) + return coordinates \ No newline at end of file diff --git a/flagscale/train/models/dreamzero/modules/wan_video_dit.py b/flagscale/train/models/dreamzero/modules/wan_video_dit.py new file mode 100644 index 0000000000..fde7784b35 --- /dev/null +++ b/flagscale/train/models/dreamzero/modules/wan_video_dit.py @@ -0,0 +1,818 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +import math +import os +from typing import Tuple, Optional +from einops import rearrange +from flagscale.train.models.dreamzero.modules.utils import hash_state_dict_keys +from flagscale.train.models.dreamzero.modules.wan_video_camera_controller import SimpleAdapter +try: + import flash_attn_interface + FLASH_ATTN_3_AVAILABLE = True +except ModuleNotFoundError: + FLASH_ATTN_3_AVAILABLE = False + +try: + import flash_attn + FLASH_ATTN_2_AVAILABLE = True +except ModuleNotFoundError: + FLASH_ATTN_2_AVAILABLE = False + +try: + from sageattention import sageattn + SAGE_ATTN_AVAILABLE = True +except ModuleNotFoundError: + SAGE_ATTN_AVAILABLE = False + + +def _gpu_supports_flash_attention(): + """FlashAttention requires Ampere (compute capability 8.0) or newer.""" + if not (FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE): + return False + try: + if not torch.cuda.is_available(): + return False + cap = torch.cuda.get_device_capability() + return cap[0] >= 8 + except Exception: + return False + + +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.modeling_utils import ModelMixin + +ENABLE_TENSORRT = os.getenv("ENABLE_TENSORRT", "False").lower() == "true" +if ENABLE_TENSORRT: + # disable torch compile and transformer engine and flash attention for onnx/tensorrt export + FLASH_ATTN_COMPATIBILITY_MODE = True + DISABLE_TORCH_COMPILE = True +else: + DISABLE_TORCH_COMPILE = False + FLASH_ATTN_COMPATIBILITY_MODE = False +def flash_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, num_heads: int, compatibility_mode=False): + # Use PyTorch SDPA on pre-Ampere GPUs or when compatibility_mode (FlashAttention requires Ampere or newer) + if compatibility_mode or not _gpu_supports_flash_attention(): + q = rearrange(q, "b s (n d) -> b n s d", n=num_heads) + k = rearrange(k, "b s (n d) -> b n s d", n=num_heads) + v = rearrange(v, "b s (n d) -> b n s d", n=num_heads) + x = F.scaled_dot_product_attention(q, k, v) + x = rearrange(x, "b n s d -> b s (n d)", n=num_heads) + elif FLASH_ATTN_3_AVAILABLE: + q = rearrange(q, "b s (n d) -> b s n d", n=num_heads) + k = rearrange(k, "b s (n d) -> b s n d", n=num_heads) + v = rearrange(v, "b s (n d) -> b s n d", n=num_heads) + x = flash_attn_interface.flash_attn_func(q, k, v) + if isinstance(x,tuple): + x = x[0] + x = rearrange(x, "b s n d -> b s (n d)", n=num_heads) + elif FLASH_ATTN_2_AVAILABLE: + q = rearrange(q, "b s (n d) -> b s n d", n=num_heads) + k = rearrange(k, "b s (n d) -> b s n d", n=num_heads) + v = rearrange(v, "b s (n d) -> b s n d", n=num_heads) + x = flash_attn.flash_attn_func(q, k, v) + x = rearrange(x, "b s n d -> b s (n d)", n=num_heads) + elif SAGE_ATTN_AVAILABLE: + q = rearrange(q, "b s (n d) -> b n s d", n=num_heads) + k = rearrange(k, "b s (n d) -> b n s d", n=num_heads) + v = rearrange(v, "b s (n d) -> b n s d", n=num_heads) + x = sageattn(q, k, v) + x = rearrange(x, "b n s d -> b s (n d)", n=num_heads) + else: + q = rearrange(q, "b s (n d) -> b n s d", n=num_heads) + k = rearrange(k, "b s (n d) -> b n s d", n=num_heads) + v = rearrange(v, "b s (n d) -> b n s d", n=num_heads) + x = F.scaled_dot_product_attention(q, k, v) + x = rearrange(x, "b n s d -> b s (n d)", n=num_heads) + return x + + +@torch.compile(disable=DISABLE_TORCH_COMPILE) +def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor): + return (x * (1 + scale) + shift) + + +@torch.compile(disable=DISABLE_TORCH_COMPILE) +def sinusoidal_embedding_1d(dim, position): + sinusoid = torch.outer( + position.type(torch.float64), + torch.pow(10000, -torch.arange(dim//2, dtype=torch.float64, device=position.device).div(dim//2)) + ) + x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1) + return x.to(position.dtype) + +@torch.compile(disable=DISABLE_TORCH_COMPILE) +def rope_apply_no_polar_op(x, freqs_cis, num_heads): + freqs_cos, freqs_sin = freqs_cis + + x_reshaped = rearrange(x, "b s (h d) -> b s h d", h=num_heads) + x_as_complex = x_reshaped.float().reshape(*x_reshaped.shape[:-1], -1, 2) + x_real, x_imag = x_as_complex.unbind(-1) + + freqs_cos = freqs_cos.to(x.device).squeeze(1).unsqueeze(0).unsqueeze(2) + freqs_sin = freqs_sin.to(x.device).squeeze(1).unsqueeze(0).unsqueeze(2) + + x_out_real = x_real * freqs_cos - x_imag * freqs_sin + x_out_imag = x_real * freqs_sin + x_imag * freqs_cos + + x_out = torch.stack([x_out_real, x_out_imag], dim=-1).flatten(start_dim=-2) + + return rearrange(x_out, "b s h d -> b s (h d)").to(x.dtype) + + +@torch.compile(disable=DISABLE_TORCH_COMPILE) +def rope_apply_polar_op(x, freqs, num_heads): + x = rearrange(x, "b s (n d) -> b s n d", n=num_heads) + x_out = torch.view_as_complex(x.to(torch.float64).reshape( + x.shape[0], x.shape[1], x.shape[2], -1, 2)) + freqs = freqs.to(x_out.device) + x_out = torch.view_as_real(x_out * freqs).flatten(2) + return x_out.to(x.dtype) + + +class RMSNorm(nn.Module): + def __init__(self, dim, eps=1e-5): + super().__init__() + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def norm(self, x): + return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps) + + def forward(self, x): + dtype = x.dtype + return self.norm(x.float()).to(dtype) * self.weight + + +def _RMSNorm(normalized_shape, eps): + return RMSNorm(normalized_shape, eps=eps) + + +def RotaryPositionEmbedding(num_heads, head_dim): + if ENABLE_TENSORRT: + return RotaryPositionEmbeddingNoPolarOp(num_heads, head_dim) + else: + return RotaryPositionEmbeddingWithPolarOp(num_heads, head_dim) + + +def rope_apply(x, freqs, num_heads): + if ENABLE_TENSORRT: + return rope_apply_no_polar_op(x, freqs, num_heads) + else: + return rope_apply_polar_op(x, freqs, num_heads) + + +class RotaryPositionEmbeddingNoPolarOp(nn.Module): + def __init__(self, num_heads: int, head_dim: int): + super().__init__() + self.num_heads = num_heads + self.head_dim = head_dim + + self.freqs = self.precompute_freqs_cis_3d(head_dim) + + def precompute_freqs_cis_3d(self, dim: int, end: int = 1024, theta: float = 10000.0): + # 3d rope precompute + f_freqs_cis = self.precompute_freqs_cis(dim - 2 * (dim // 3), end, theta) + h_freqs_cis = self.precompute_freqs_cis(dim // 3, end, theta) + w_freqs_cis = self.precompute_freqs_cis(dim // 3, end, theta) + return {"f": f_freqs_cis, "h": h_freqs_cis, "w": w_freqs_cis} + + def precompute_freqs_cis(self, dim: int, end: int = 1024, theta: float = 10000.0): + # 1d rope precompute + freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)) + freqs = torch.outer(torch.arange(end, device=freqs.device), freqs) + freqs_real = torch.cos(freqs) + freqs_imag = torch.sin(freqs) + return (freqs_real, freqs_imag) + + def forward(self, f: int, h: int, w: int, a: int) -> torch.Tensor: + freqs_cos_3d = torch.cat( + [ + self.freqs["f"][0][:f].view(f, 1, 1, -1).expand(f, h, w, -1), + self.freqs["h"][0][:h].view(1, h, 1, -1).expand(f, h, w, -1), + self.freqs["w"][0][:w].view(1, 1, w, -1).expand(f, h, w, -1), + ], + dim=-1, + ).reshape(f * h * w, 1, -1) + freqs_sin_3d = torch.cat( + [ + self.freqs["f"][1][:f].view(f, 1, 1, -1).expand(f, h, w, -1), + self.freqs["h"][1][:h].view(1, h, 1, -1).expand(f, h, w, -1), + self.freqs["w"][1][:w].view(1, 1, w, -1).expand(f, h, w, -1), + ], + dim=-1, + ).reshape(f * h * w, 1, -1) + + + return torch.cat([freqs_cos_3d, freqs_sin_3d], dim=0) + + def post_initialize(self): + self.freqs = { + key: (value[0].to("cuda"), value[1].to("cuda")) for key, value in self.freqs.items() + } + + +class RotaryPositionEmbeddingWithPolarOp(nn.Module): + def __init__(self, num_heads: int, head_dim: int): + super().__init__() + self.num_heads = num_heads + self.head_dim = head_dim + + self.freqs = self.precompute_freqs_cis_3d(head_dim) + + def precompute_freqs_cis_3d(self, dim: int, end: int = 1024, theta: float = 10000.0): + # 3d rope precompute + f_freqs_cis = self.precompute_freqs_cis(dim - 2 * (dim // 3), end, theta) + h_freqs_cis = self.precompute_freqs_cis(dim // 3, end, theta) + w_freqs_cis = self.precompute_freqs_cis(dim // 3, end, theta) + return {"f": f_freqs_cis, "h": h_freqs_cis, "w": w_freqs_cis} + + def precompute_freqs_cis(self, dim: int, end: int = 1024, theta: float = 10000.0): + # 1d rope precompute + freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].double() / dim)) + freqs = torch.outer(torch.arange(end), freqs) + freqs_cis = torch.polar(torch.ones_like(freqs), freqs) # complex64 + return freqs_cis + + def forward(self, f: int, h: int, w: int, a: int) -> torch.Tensor: + freqs = torch.cat( + [ + self.freqs["f"][:f].view(f, 1, 1, -1).expand(f, h, w, -1), + self.freqs["h"][:h].view(1, h, 1, -1).expand(f, h, w, -1), + self.freqs["w"][:w].view(1, 1, w, -1).expand(f, h, w, -1), + ], + dim=-1, + ).reshape(f * h * w, 1, -1) + return freqs + + def post_initialize(self): + self.freqs = {key: value.to(device="cuda") for key, value in self.freqs.items()} + + +class AttentionModule(nn.Module): + def __init__(self, num_heads): + super().__init__() + self.num_heads = num_heads + + def forward(self, q, k, v): + x = flash_attention(q=q, k=k, v=v, num_heads=self.num_heads, compatibility_mode=FLASH_ATTN_COMPATIBILITY_MODE) + return x + + +class SelfAttention(nn.Module): + def __init__(self, dim: int, num_heads: int, eps: float = 1e-6): + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + + self.q = nn.Linear(dim, dim) + self.k = nn.Linear(dim, dim) + self.v = nn.Linear(dim, dim) + self.o = nn.Linear(dim, dim) + self.norm_q = _RMSNorm(dim, eps=eps) + self.norm_k = _RMSNorm(dim, eps=eps) + + self.attn = AttentionModule(self.num_heads) + + def forward(self, x, freqs): + q = self.norm_q(self.q(x)) + k = self.norm_k(self.k(x)) + v = self.v(x) + q = rope_apply(q, freqs, self.num_heads) + k = rope_apply(k, freqs, self.num_heads) + x = self.attn(q, k, v) + return self.o(x) + + +class CrossAttention(nn.Module): + def __init__(self, dim: int, num_heads: int, eps: float = 1e-6, has_image_input: bool = False): + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + + self.q = nn.Linear(dim, dim) + self.k = nn.Linear(dim, dim) + self.v = nn.Linear(dim, dim) + self.o = nn.Linear(dim, dim) + self.norm_q = _RMSNorm(dim, eps=eps) + self.norm_k = _RMSNorm(dim, eps=eps) + self.has_image_input = has_image_input + if has_image_input: + self.k_img = nn.Linear(dim, dim) + self.v_img = nn.Linear(dim, dim) + self.norm_k_img = _RMSNorm(dim, eps=eps) + + self.attn = AttentionModule(self.num_heads) + + def forward(self, x: torch.Tensor, y: torch.Tensor): + if self.has_image_input: + img = y[:, :257] + ctx = y[:, 257:] + else: + ctx = y + q = self.norm_q(self.q(x)) + k = self.norm_k(self.k(ctx)) + v = self.v(ctx) + x = self.attn(q, k, v) + if self.has_image_input: + k_img = self.norm_k_img(self.k_img(img)) + v_img = self.v_img(img) + y = flash_attention(q, k_img, v_img, num_heads=self.num_heads, compatibility_mode=FLASH_ATTN_COMPATIBILITY_MODE) + x = x + y + return self.o(x) + + +class GateModule(nn.Module): + def __init__(self,): + super().__init__() + + @torch.compile(disable=DISABLE_TORCH_COMPILE) + def forward(self, x, gate, residual): + return x + gate * residual + +class DiTBlock(nn.Module): + def __init__(self, has_image_input: bool, dim: int, num_heads: int, ffn_dim: int, eps: float = 1e-6): + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.ffn_dim = ffn_dim + + self.self_attn = SelfAttention(dim, num_heads, eps) + self.cross_attn = CrossAttention( + dim, num_heads, eps, has_image_input=has_image_input) + self.norm1 = nn.LayerNorm(dim, eps=eps, elementwise_affine=False) + self.norm2 = nn.LayerNorm(dim, eps=eps, elementwise_affine=False) + self.norm3 = nn.LayerNorm(dim, eps=eps) + self.ffn = nn.Sequential(nn.Linear(dim, ffn_dim), nn.GELU( + approximate='tanh'), nn.Linear(ffn_dim, dim)) + self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + self.gate = GateModule() + + def forward(self, x, context, t_mod, freqs): + # msa: multi-head self-attention mlp: multi-layer perceptron + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + t_mod).chunk(6, dim=1) + input_x = modulate(self.norm1(x), shift_msa, scale_msa).to(dtype=t_mod.dtype) + x = self.gate(x, gate_msa, self.self_attn(input_x, freqs)) + x = x + self.cross_attn(self.norm3(x).to(dtype=t_mod.dtype), context.to(dtype=t_mod.dtype)) + input_x = modulate(self.norm2(x), shift_mlp, scale_mlp) + x = self.gate(x, gate_mlp, self.ffn(input_x)) + return x + + +class MLP(torch.nn.Module): + def __init__(self, in_dim, out_dim, has_pos_emb=False): + super().__init__() + self.proj = torch.nn.Sequential( + nn.LayerNorm(in_dim), + nn.Linear(in_dim, in_dim), + nn.GELU(), + nn.Linear(in_dim, out_dim), + nn.LayerNorm(out_dim) + ) + self.has_pos_emb = has_pos_emb + if has_pos_emb: + self.emb_pos = torch.nn.Parameter(torch.zeros((1, 514, 1280))) + + def forward(self, x): + if self.has_pos_emb: + x = x + self.emb_pos.to(dtype=x.dtype, device=x.device) + return self.proj(x) + + +class Head(nn.Module): + def __init__(self, dim: int, out_dim: int, patch_size: Tuple[int, int, int], eps: float): + super().__init__() + self.dim = dim + self.patch_size = patch_size + self.norm = nn.LayerNorm(dim, eps=eps, elementwise_affine=False) + self.head = nn.Linear(dim, out_dim * math.prod(patch_size)) + self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5) + + def forward(self, x, t_mod): + batch_size = x.shape[0] + shift, scale = (self.modulation.repeat(batch_size,1,1).to(dtype=t_mod.dtype, device=t_mod.device) + t_mod.unsqueeze(1)).chunk(2, dim=1) + x = (self.head(self.norm(x) * (1 + scale) + shift)) + return x + + +class WanModel(ModelMixin, ConfigMixin): + @register_to_config + def __init__( + self, + dim: int, + in_dim: int, + ffn_dim: int, + out_dim: int, + freq_dim: int, + eps: float, + num_heads: int, + num_layers: int, + text_dim: int = 4096, + patch_size: Tuple[int, int, int] = [1, 2, 2], + has_image_input: bool = True, + has_image_pos_emb: bool = False, + has_ref_conv: bool = False, + add_control_adapter: bool = False, + in_dim_control_adapter: int = 24, + diffusion_model_pretrained_path: str = None, + ): + super().__init__() + self.dim = dim + self.freq_dim = freq_dim + self.has_image_input = has_image_input + self.patch_size = patch_size + self.diffusion_model_pretrained_path = diffusion_model_pretrained_path + + + self.patch_embedding = nn.Conv3d( + in_dim, dim, kernel_size=patch_size, stride=patch_size) + self.text_embedding = nn.Sequential( + nn.Linear(text_dim, dim), + nn.GELU(approximate='tanh'), + nn.Linear(dim, dim) + ) + self.time_embedding = nn.Sequential( + nn.Linear(freq_dim, dim), + nn.SiLU(), + nn.Linear(dim, dim) + ) + self.time_projection = nn.Sequential( + nn.SiLU(), nn.Linear(dim, dim * 6)) + self.blocks = nn.ModuleList([ + DiTBlock(has_image_input, dim, num_heads, ffn_dim, eps) + for _ in range(num_layers) + ]) + self.head = Head(dim, out_dim, patch_size, eps) + head_dim = dim // num_heads + + self.rope = RotaryPositionEmbedding(num_heads=num_heads, head_dim=head_dim) + + if has_image_input: + self.img_emb = MLP(1280, dim, has_pos_emb=has_image_pos_emb) # clip_feature_dim = 1280 + if has_ref_conv: + self.ref_conv = nn.Conv2d(16, dim, kernel_size=(2, 2), stride=(2, 2)) + self.has_image_pos_emb = has_image_pos_emb + self.has_ref_conv = has_ref_conv + if add_control_adapter: + self.control_adapter = SimpleAdapter(in_dim_control_adapter, dim, kernel_size=patch_size[1:], stride=patch_size[1:]) + else: + self.control_adapter = None + + self.use_gradient_checkpointing = False + self.use_gradient_checkpointing_offload = False + + def patchify(self, x: torch.Tensor,control_camera_latents_input: torch.Tensor = None): + x = self.patch_embedding(x) + if self.control_adapter is not None and control_camera_latents_input is not None: + y_camera = self.control_adapter(control_camera_latents_input) + x = [u + v for u, v in zip(x, y_camera)] + x = x[0].unsqueeze(0) + grid_size = x.shape[2:] + x = rearrange(x, 'b c f h w -> b (f h w) c').contiguous() + return x, grid_size # x, grid_size: (f, h, w) + + def unpatchify(self, x: torch.Tensor, grid_size: torch.Tensor): + return rearrange( + x, 'b (f h w) (x y z c) -> b c (f x) (h y) (w z)', + f=grid_size[0], h=grid_size[1], w=grid_size[2], + x=self.patch_size[0], y=self.patch_size[1], z=self.patch_size[2] + ) + + def forward(self, + x: torch.Tensor, + timestep: torch.Tensor, + context: torch.Tensor, + clip_feature: Optional[torch.Tensor] = None, + y: Optional[torch.Tensor] = None, + use_gradient_checkpointing: bool = False, + use_gradient_checkpointing_offload: bool = False, + **kwargs, + ): + t = self.time_embedding( + sinusoidal_embedding_1d(self.freq_dim, timestep)) + t_mod = self.time_projection(t).unflatten(1, (6, self.dim)) + context = self.text_embedding(context) + + if self.has_image_input: + x = torch.cat([x, y], dim=1) # (b, c_x + c_y, f, h, w) + clip_embdding = self.img_emb(clip_feature) + context = torch.cat([clip_embdding, context], dim=1) + # print("clip embedding shape", clip_embdding.shape) + # print("x before patchify", x.shape) + # clip embedding shape of B * 257 * dim (5120) + # context shape of B * 769 * dim (5120) - 257 + 512 + # x before patchify B * 36 * l_t * l_h * l_w + x, (f, h, w) = self.patchify(x) + # x after patchify B * 512 * 5120 + # f = 2, h = 16, w = 16 + # print("x and context shape", x.shape, context.shape, f,h,w) + + freqs = self.rope(f=f, h=h, w=w, a=x.shape[1]) + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + return custom_forward + + for block in self.blocks: + if self.training and self.use_gradient_checkpointing: + if self.use_gradient_checkpointing_offload: + with torch.autograd.graph.save_on_cpu(): + x = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + x, context, t_mod, freqs, + use_reentrant=False, + ) + else: + x = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + x, context, t_mod, freqs, + use_reentrant=False, + ) + else: + x = block(x, context, t_mod, freqs) + + # print("x before head", x.shape) + x = self.head(x, t) + # print("x before unpatchify", x.shape) + x = self.unpatchify(x, (f, h, w)) + # print("x after patchify", x.shape) + return x + def post_initialize(self): + self.rope.post_initialize() + @staticmethod + def state_dict_converter(): + return WanModelStateDictConverter() + + +class WanModelStateDictConverter: + def __init__(self): + pass + + def from_diffusers(self, state_dict): + rename_dict = { + "blocks.0.attn1.norm_k.weight": "blocks.0.self_attn.norm_k.weight", + "blocks.0.attn1.norm_q.weight": "blocks.0.self_attn.norm_q.weight", + "blocks.0.attn1.to_k.bias": "blocks.0.self_attn.k.bias", + "blocks.0.attn1.to_k.weight": "blocks.0.self_attn.k.weight", + "blocks.0.attn1.to_out.0.bias": "blocks.0.self_attn.o.bias", + "blocks.0.attn1.to_out.0.weight": "blocks.0.self_attn.o.weight", + "blocks.0.attn1.to_q.bias": "blocks.0.self_attn.q.bias", + "blocks.0.attn1.to_q.weight": "blocks.0.self_attn.q.weight", + "blocks.0.attn1.to_v.bias": "blocks.0.self_attn.v.bias", + "blocks.0.attn1.to_v.weight": "blocks.0.self_attn.v.weight", + "blocks.0.attn2.norm_k.weight": "blocks.0.cross_attn.norm_k.weight", + "blocks.0.attn2.norm_q.weight": "blocks.0.cross_attn.norm_q.weight", + "blocks.0.attn2.to_k.bias": "blocks.0.cross_attn.k.bias", + "blocks.0.attn2.to_k.weight": "blocks.0.cross_attn.k.weight", + "blocks.0.attn2.to_out.0.bias": "blocks.0.cross_attn.o.bias", + "blocks.0.attn2.to_out.0.weight": "blocks.0.cross_attn.o.weight", + "blocks.0.attn2.to_q.bias": "blocks.0.cross_attn.q.bias", + "blocks.0.attn2.to_q.weight": "blocks.0.cross_attn.q.weight", + "blocks.0.attn2.to_v.bias": "blocks.0.cross_attn.v.bias", + "blocks.0.attn2.to_v.weight": "blocks.0.cross_attn.v.weight", + "blocks.0.ffn.net.0.proj.bias": "blocks.0.ffn.0.bias", + "blocks.0.ffn.net.0.proj.weight": "blocks.0.ffn.0.weight", + "blocks.0.ffn.net.2.bias": "blocks.0.ffn.2.bias", + "blocks.0.ffn.net.2.weight": "blocks.0.ffn.2.weight", + "blocks.0.norm2.bias": "blocks.0.norm3.bias", + "blocks.0.norm2.weight": "blocks.0.norm3.weight", + "blocks.0.scale_shift_table": "blocks.0.modulation", + "condition_embedder.text_embedder.linear_1.bias": "text_embedding.0.bias", + "condition_embedder.text_embedder.linear_1.weight": "text_embedding.0.weight", + "condition_embedder.text_embedder.linear_2.bias": "text_embedding.2.bias", + "condition_embedder.text_embedder.linear_2.weight": "text_embedding.2.weight", + "condition_embedder.time_embedder.linear_1.bias": "time_embedding.0.bias", + "condition_embedder.time_embedder.linear_1.weight": "time_embedding.0.weight", + "condition_embedder.time_embedder.linear_2.bias": "time_embedding.2.bias", + "condition_embedder.time_embedder.linear_2.weight": "time_embedding.2.weight", + "condition_embedder.time_proj.bias": "time_projection.1.bias", + "condition_embedder.time_proj.weight": "time_projection.1.weight", + "patch_embedding.bias": "patch_embedding.bias", + "patch_embedding.weight": "patch_embedding.weight", + "scale_shift_table": "head.modulation", + "proj_out.bias": "head.head.bias", + "proj_out.weight": "head.head.weight", + } + state_dict_ = {} + for name, param in state_dict.items(): + if name in rename_dict: + state_dict_[rename_dict[name]] = param + else: + name_ = ".".join(name.split(".")[:1] + ["0"] + name.split(".")[2:]) + if name_ in rename_dict: + name_ = rename_dict[name_] + name_ = ".".join(name_.split(".")[:1] + [name.split(".")[1]] + name_.split(".")[2:]) + state_dict_[name_] = param + if hash_state_dict_keys(state_dict) == "cb104773c6c2cb6df4f9529ad5c60d0b": + config = { + "model_type": "t2v", + "patch_size": (1, 2, 2), + "text_len": 512, + "in_dim": 16, + "dim": 5120, + "ffn_dim": 13824, + "freq_dim": 256, + "text_dim": 4096, + "out_dim": 16, + "num_heads": 40, + "num_layers": 40, + "window_size": (-1, -1), + "qk_norm": True, + "cross_attn_norm": True, + "eps": 1e-6, + } + else: + config = {} + return state_dict_, config + + def from_civitai(self, state_dict): + state_dict = {name: param for name, param in state_dict.items() if not name.startswith("vace")} + if hash_state_dict_keys(state_dict) == "9269f8db9040a9d860eaca435be61814": + config = { + "has_image_input": False, + "patch_size": [1, 2, 2], + "in_dim": 16, + "dim": 1536, + "ffn_dim": 8960, + "freq_dim": 256, + "text_dim": 4096, + "out_dim": 16, + "num_heads": 12, + "num_layers": 30, + "eps": 1e-6 + } + elif hash_state_dict_keys(state_dict) == "aafcfd9672c3a2456dc46e1cb6e52c70": + config = { + "has_image_input": False, + "patch_size": [1, 2, 2], + "in_dim": 16, + "dim": 5120, + "ffn_dim": 13824, + "freq_dim": 256, + "text_dim": 4096, + "out_dim": 16, + "num_heads": 40, + "num_layers": 40, + "eps": 1e-6 + } + elif hash_state_dict_keys(state_dict) == "6bfcfb3b342cb286ce886889d519a77e": + config = { + "has_image_input": True, + "patch_size": [1, 2, 2], + "in_dim": 36, + "dim": 5120, + "ffn_dim": 13824, + "freq_dim": 256, + "text_dim": 4096, + "out_dim": 16, + "num_heads": 40, + "num_layers": 40, + "eps": 1e-6 + } + elif hash_state_dict_keys(state_dict) == "6d6ccde6845b95ad9114ab993d917893": + config = { + "has_image_input": True, + "patch_size": [1, 2, 2], + "in_dim": 36, + "dim": 1536, + "ffn_dim": 8960, + "freq_dim": 256, + "text_dim": 4096, + "out_dim": 16, + "num_heads": 12, + "num_layers": 30, + "eps": 1e-6 + } + elif hash_state_dict_keys(state_dict) == "6bfcfb3b342cb286ce886889d519a77e": + config = { + "has_image_input": True, + "patch_size": [1, 2, 2], + "in_dim": 36, + "dim": 5120, + "ffn_dim": 13824, + "freq_dim": 256, + "text_dim": 4096, + "out_dim": 16, + "num_heads": 40, + "num_layers": 40, + "eps": 1e-6 + } + elif hash_state_dict_keys(state_dict) == "349723183fc063b2bfc10bb2835cf677": + # 1.3B PAI control + config = { + "has_image_input": True, + "patch_size": [1, 2, 2], + "in_dim": 48, + "dim": 1536, + "ffn_dim": 8960, + "freq_dim": 256, + "text_dim": 4096, + "out_dim": 16, + "num_heads": 12, + "num_layers": 30, + "eps": 1e-6 + } + elif hash_state_dict_keys(state_dict) == "efa44cddf936c70abd0ea28b6cbe946c": + # 14B PAI control + config = { + "has_image_input": True, + "patch_size": [1, 2, 2], + "in_dim": 48, + "dim": 5120, + "ffn_dim": 13824, + "freq_dim": 256, + "text_dim": 4096, + "out_dim": 16, + "num_heads": 40, + "num_layers": 40, + "eps": 1e-6 + } + elif hash_state_dict_keys(state_dict) == "3ef3b1f8e1dab83d5b71fd7b617f859f": + config = { + "has_image_input": True, + "patch_size": [1, 2, 2], + "in_dim": 36, + "dim": 5120, + "ffn_dim": 13824, + "freq_dim": 256, + "text_dim": 4096, + "out_dim": 16, + "num_heads": 40, + "num_layers": 40, + "eps": 1e-6, + "has_image_pos_emb": True + } + elif hash_state_dict_keys(state_dict) == "70ddad9d3a133785da5ea371aae09504": + # 1.3B PAI control v1.1 + config = { + "has_image_input": True, + "patch_size": [1, 2, 2], + "in_dim": 48, + "dim": 1536, + "ffn_dim": 8960, + "freq_dim": 256, + "text_dim": 4096, + "out_dim": 16, + "num_heads": 12, + "num_layers": 30, + "eps": 1e-6, + "has_ref_conv": True + } + elif hash_state_dict_keys(state_dict) == "26bde73488a92e64cc20b0a7485b9e5b": + # 14B PAI control v1.1 + config = { + "has_image_input": True, + "patch_size": [1, 2, 2], + "in_dim": 48, + "dim": 5120, + "ffn_dim": 13824, + "freq_dim": 256, + "text_dim": 4096, + "out_dim": 16, + "num_heads": 40, + "num_layers": 40, + "eps": 1e-6, + "has_ref_conv": True + } + elif hash_state_dict_keys(state_dict) == "ac6a5aa74f4a0aab6f64eb9a72f19901": + # 1.3B PAI control-camera v1.1 + config = { + "has_image_input": True, + "patch_size": [1, 2, 2], + "in_dim": 32, + "dim": 1536, + "ffn_dim": 8960, + "freq_dim": 256, + "text_dim": 4096, + "out_dim": 16, + "num_heads": 12, + "num_layers": 30, + "eps": 1e-6, + "has_ref_conv": False, + "add_control_adapter": True, + "in_dim_control_adapter": 24, + } + elif hash_state_dict_keys(state_dict) == "b61c605c2adbd23124d152ed28e049ae": + # 14B PAI control-camera v1.1 + config = { + "has_image_input": True, + "patch_size": [1, 2, 2], + "in_dim": 32, + "dim": 5120, + "ffn_dim": 13824, + "freq_dim": 256, + "text_dim": 4096, + "out_dim": 16, + "num_heads": 40, + "num_layers": 40, + "eps": 1e-6, + "has_ref_conv": False, + "add_control_adapter": True, + "in_dim_control_adapter": 24, + } + else: + config = {} + return state_dict, config diff --git a/flagscale/train/models/dreamzero/modules/wan_video_dit_action_casual_chunk.py b/flagscale/train/models/dreamzero/modules/wan_video_dit_action_casual_chunk.py new file mode 100644 index 0000000000..b7a301ae61 --- /dev/null +++ b/flagscale/train/models/dreamzero/modules/wan_video_dit_action_casual_chunk.py @@ -0,0 +1,2286 @@ +from typing import Any, TypeAlias + +from flagscale.train.models.dreamzero.modules.wan2_1_attention import AttentionModule +from flagscale.train.models.dreamzero.n1_5.modules.action_encoder import ( + SinusoidalPositionalEncoding, + swish, +) +from flagscale.train.models.dreamzero.modules.wan2_1_submodule import ( + WanRMSNorm, + rope_action_apply, + WanLayerNorm, + WAN_CROSSATTENTION_CLASSES, + rope_params, + MLPProj, + sinusoidal_embedding_1d +) +from torch.nn.attention.flex_attention import create_block_mask, create_mask +from torch.nn.attention.flex_attention import BlockMask +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.models.modeling_utils import ModelMixin +import torch.nn as nn +import torch.nn.functional as F +import torch +import math +import torch.distributed as dist +import os + +ENABLE_TENSORRT = os.getenv("ENABLE_TENSORRT", "False").lower() == "true" + + +class CategorySpecificLinear(nn.Module): + def __init__(self, num_categories, input_dim, hidden_dim): + super().__init__() + self.num_categories = num_categories + # For each category, we have separate weights and biases. + self.W = nn.Parameter(0.02 * torch.randn(num_categories, input_dim, hidden_dim)) + self.b = nn.Parameter(torch.zeros(num_categories, hidden_dim)) + + def forward(self, x, cat_ids): + selected_W = self.W[cat_ids] + selected_b = self.b[cat_ids] + return torch.bmm(x, selected_W) + selected_b.unsqueeze(1) + + +class CategorySpecificMLP(nn.Module): + def __init__(self, num_categories, input_dim, hidden_dim, output_dim): + super().__init__() + self.num_categories = num_categories + self.layer1 = CategorySpecificLinear(num_categories, input_dim, hidden_dim) + self.layer2 = CategorySpecificLinear(num_categories, hidden_dim, output_dim) + + def forward(self, x, cat_ids): + hidden = F.relu(self.layer1(x, cat_ids)) + return self.layer2(hidden, cat_ids) + + +class MultiEmbodimentActionEncoder(nn.Module): + def __init__(self, action_dim, hidden_size, num_embodiments): + super().__init__() + self.hidden_size = hidden_size + self.num_embodiments = num_embodiments + + # W1: R^{w x d}, W2: R^{w x 2w}, W3: R^{w x w} + self.W1 = CategorySpecificLinear(num_embodiments, action_dim, hidden_size) # (d -> w) + self.W2 = CategorySpecificLinear(num_embodiments, 2 * hidden_size, hidden_size) # (2w -> w) + self.W3 = CategorySpecificLinear(num_embodiments, hidden_size, hidden_size) # (w -> w) + self.pos_encoding = SinusoidalPositionalEncoding(hidden_size) + + def forward(self, actions, timesteps, cat_ids): + """ + actions: shape (B, T, action_dim) + timesteps: shape (B,) -- a single scalar per batch item + cat_ids: shape (B,) + returns: shape (B, T, hidden_size) + """ + B, T, _ = actions.shape + + # Standard action MLP step for shape => (B, T, w) + a_emb = self.W1(actions, cat_ids) + + # 3) Get the sinusoidal encoding (B, T, w) + tau_emb = self.pos_encoding(timesteps).to(dtype=a_emb.dtype) + + # 4) Concat along last dim => (B, T, 2w), then W2 => (B, T, w), swish + x = torch.cat([a_emb, tau_emb], dim=-1) + x = swish(self.W2(x, cat_ids)) + + # 5) Finally W3 => (B, T, w) + x = self.W3(x, cat_ids) + return x + + +def causal_rope_action_apply(x, freqs, freqs_action, freqs_state, action_register_length, num_action_per_block, num_state_per_block, action_state_index): + if ENABLE_TENSORRT: + return causal_rope_action_apply_no_polar(x, freqs, freqs_action, freqs_state, action_register_length, num_action_per_block, num_state_per_block, action_state_index) + else: + return causal_rope_action_apply_polar(x, freqs, freqs_action, freqs_state, action_register_length, num_action_per_block, num_state_per_block, action_state_index) + + +def causal_rope_action_apply_no_polar( + x: torch.Tensor, + freqs: torch.Tensor, + freqs_action: torch.Tensor, + freqs_state: torch.Tensor, + action_register_length: int | None, + num_action_per_block: int, + num_state_per_block: int, + action_state_index: int, +): + B, seq_len, n, d = x.shape + + # (B, seq_len, n, d) -> (B, seq_len, n, d/2, 2) + x = x.reshape(B, seq_len, n, -1, 2) + x_real = x[..., 0] + x_imag = x[..., 1] + + # Split freqs into cos and sin components + freqs = freqs.unsqueeze(0).view(1, freqs.shape[0], 1, -1, 2) + freqs_cos = freqs[..., 0] # Shape: (1, seq_len', 1, d/2) + freqs_sin = freqs[..., 1] # Shape: (1, seq_len', 1, d/2) + + # Handle the Action/State Register Frequencies + if action_register_length is not None: + assert action_register_length == (num_action_per_block + num_state_per_block) + + freqs_action_slice = freqs_action[ + action_state_index * num_action_per_block:(action_state_index + 1) * num_action_per_block + ] + freqs_state_slice = freqs_state[ + action_state_index * num_state_per_block:(action_state_index + 1) * num_state_per_block + ] + + # Combine the action/state tokens for this frame + freqs_1d = torch.cat([freqs_action_slice, freqs_state_slice], dim=0).view( + action_register_length, 1, -1, 2 + ) + + # Split the new action/state frequencies + freqs_cos_1d = freqs_1d[..., 0] + freqs_sin_1d = freqs_1d[..., 1] + + # Append the action/state register sin/cos to the main sequence sin/cos + freqs_cos = torch.cat([freqs_cos[0], freqs_cos_1d], dim=0).unsqueeze(0) + freqs_sin = torch.cat([freqs_sin[0], freqs_sin_1d], dim=0).unsqueeze(0) + + x_real_rotated = x_real * freqs_cos - x_imag * freqs_sin + x_imag_rotated = x_real * freqs_sin + x_imag * freqs_cos + + x_rotated = torch.stack((x_real_rotated, x_imag_rotated), dim=-1) + + return x_rotated.flatten(3) + +def causal_rope_action_apply_polar( + x: torch.Tensor, + freqs: torch.Tensor, + freqs_action: torch.Tensor, + freqs_state: torch.Tensor, + action_register_length: int | None, + num_action_per_block: int, + num_state_per_block: int, + action_state_index: int, +): + B, seq_len, n, _ = x.shape + + # precompute multipliers + x = torch.view_as_complex( + x.to(torch.float64).reshape(B, seq_len, n, -1, 2) + ) + + if action_register_length is not None: + assert action_register_length == (num_action_per_block + num_state_per_block) + freqs_action = freqs_action[ + action_state_index * num_action_per_block:(action_state_index + 1) * num_action_per_block + ] + freqs_state = freqs_state[ + action_state_index * num_state_per_block:(action_state_index + 1) * num_state_per_block + ] + freqs_1d = torch.cat([freqs_action, freqs_state], dim=0).view(action_register_length, 1, -1) + freqs = torch.cat([freqs, freqs_1d], dim=0) + + # apply rotary embedding + freqs = freqs.unsqueeze(0) + x = torch.view_as_real(x * freqs).flatten(3) + + return x + + +class CausalWanSelfAttention(nn.Module): + + def __init__(self, + dim, + num_heads, + frame_seqlen, + local_attn_size=-1, + sink_size=0, + num_frame_per_block=1, + qk_norm=True, + eps=1e-6, + num_action_per_block=32, + num_state_per_block=1): + assert dim % num_heads == 0 + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.local_attn_size = local_attn_size + self.sink_size = sink_size + self.num_frame_per_block = num_frame_per_block + self.qk_norm = qk_norm + self.eps = eps + self.max_attention_size = 21 * frame_seqlen if local_attn_size == -1 else local_attn_size * frame_seqlen + self.frame_seqlen = frame_seqlen + self.num_action_per_block = num_action_per_block + self.num_state_per_block = num_state_per_block + # layers + self.q = nn.Linear(dim, dim) + self.k = nn.Linear(dim, dim) + self.v = nn.Linear(dim, dim) + self.o = nn.Linear(dim, dim) + self.norm_q = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity() + self.attn = AttentionModule(num_heads=self.num_heads, head_dim=self.head_dim) + self.causal_attn = AttentionModule(num_heads=self.num_heads, head_dim=self.head_dim, causal=True) + + def _visualize_attention_mask(self, total_len, first_image_len, image_blocks_len, + action_len, state_len, num_image_blocks, + num_action_blocks, num_state_blocks, + num_frame_per_block, frame_seqlen, + num_action_per_block, num_state_per_block): + """ + Create and print a visualization of the attention mask pattern. + Returns a binary mask [total_len, total_len] where 1 = can attend, 0 = cannot attend. + """ + # Token ranges + first_image_start = 0 + first_image_end = first_image_len + image_blocks_start = first_image_end + image_blocks_end = image_blocks_start + image_blocks_len + action_start = image_blocks_end + action_end = action_start + action_len + state_start = action_end + state_end = state_start + state_len + + # Create mask tensor + mask = torch.zeros(total_len, total_len, dtype=torch.bool) + + # First image: self-attention only + mask[first_image_start:first_image_end, first_image_start:first_image_end] = True + + # Image blocks + for block_idx in range(num_image_blocks): + block_start = image_blocks_start + block_idx * num_frame_per_block * frame_seqlen + block_end = image_blocks_start + (block_idx + 1) * num_frame_per_block * frame_seqlen + + # Attend to first image + mask[block_start:block_end, first_image_start:first_image_end] = True + + # Attend to previous and current image blocks + if self.local_attn_size != -1: + image_kv_start = max(image_blocks_start, block_end - self.local_attn_size * frame_seqlen) + else: + image_kv_start = image_blocks_start + mask[block_start:block_end, image_kv_start:block_end] = True + + # Attend to current action block + action_block_start = action_start + block_idx * num_action_per_block + action_block_end = action_start + (block_idx + 1) * num_action_per_block + mask[block_start:block_end, action_block_start:action_block_end] = True + + # Attend to current state block + state_block_start = state_start + block_idx * num_state_per_block + state_block_end = state_start + (block_idx + 1) * num_state_per_block + mask[block_start:block_end, state_block_start:state_block_end] = True + + # Action blocks + for block_idx in range(num_action_blocks): + action_block_start = action_start + block_idx * num_action_per_block + action_block_end = action_start + (block_idx + 1) * num_action_per_block + + # Attend to first image + mask[action_block_start:action_block_end, first_image_start:first_image_end] = True + + # Attend to previous and current image blocks + image_block_end = image_blocks_start + (block_idx + 1) * num_frame_per_block * frame_seqlen + if self.local_attn_size != -1: + image_kv_start = max(image_blocks_start, image_block_end - self.local_attn_size * frame_seqlen) + else: + image_kv_start = image_blocks_start + mask[action_block_start:action_block_end, image_kv_start:image_block_end] = True + + # Self-attention + mask[action_block_start:action_block_end, action_block_start:action_block_end] = True + + # Attend to current state block + state_block_start = state_start + block_idx * num_state_per_block + state_block_end = state_start + (block_idx + 1) * num_state_per_block + mask[action_block_start:action_block_end, state_block_start:state_block_end] = True + + # State blocks: self-attention only + for block_idx in range(num_state_blocks): + state_block_start = state_start + block_idx * num_state_per_block + state_block_end = state_start + (block_idx + 1) * num_state_per_block + mask[state_block_start:state_block_end, state_block_start:state_block_end] = True + + return mask + + def _blockwise_causal_flash_attn(self, q, k, v, frame_seqlen, num_frame_per_block=1, + action_horizon=None, state_horizon=None, + num_action_per_block=None, num_state_per_block=None, + visualize_mask=False): + """ + Implement blockwise causal attention using flash_attention. + Matches the pattern from _prepare_blockwise_causal_attn_mask: + + Structure: + - First image: conditioning only, cannot attend to anything + - Image blocks: can attend to first image + previous image blocks + current action block + current state block + - Action blocks: can attend to previous image blocks + current image block + current state block + first image + - State blocks: conditioning only, cannot attend to anything + + Args: + q, k, v: Query, key, value tensors [B, L, num_heads, head_dim] + frame_seqlen: Number of tokens per frame + num_frame_per_block: Number of frames per attention block + action_horizon: Total number of action tokens (if None, no action/state tokens) + state_horizon: Total number of state tokens (if None, no action/state tokens) + num_action_per_block: Number of action tokens per block + num_state_per_block: Number of state tokens per block + visualize_mask: If True, print the attention mask pattern + + Returns: + Attention output [B, L, num_heads, head_dim] + """ + b, total_len, n, d = q.shape + + # Check if we have action/state tokens + has_action_state = (action_horizon is not None and state_horizon is not None) + + if not has_action_state: + # OPTIMIZED: Simple blockwise causal attention (without action/state tokens) + num_frames = total_len // frame_seqlen + block_size = frame_seqlen * num_frame_per_block + num_blocks = (num_frames - 1) // num_frame_per_block + + # Handle edge case when sequence is too short (no blocks to process) + if num_blocks <= 0: + # Process entire sequence as a single block + return self.attn(q, k, v) + + # OPTIMIZATION: For global attention, process all blocks in one call with causal masking + if self.local_attn_size == -1: + # Single flash_attention call with causal=True for all blocks at once + # This is much faster than looping! + return self.causal_attn(q, k, v) + + # With local attention, still need loop but optimize it + # Pre-allocate output tensor + output = torch.empty_like(q) + + # Pre-compute block boundaries + block_starts = [frame_seqlen + i * block_size for i in range(num_blocks)] + block_ends = [min(start + block_size, total_len) for start in block_starts] + kv_starts = [max(0, end - self.local_attn_size * frame_seqlen) for end in block_ends] + + for block_idx in range(num_blocks): + block_start = block_starts[block_idx] + block_end = block_ends[block_idx] + kv_start = kv_starts[block_idx] + + output[:, block_start:block_end] = self.attn( + q[:, block_start:block_end], + k[:, kv_start:block_end], + v[:, kv_start:block_end] + ) + + return output + + assert action_horizon is not None and state_horizon is not None + assert num_action_per_block is not None and num_state_per_block is not None + + # Multi-modal structure: [first image] [image blocks] [action blocks] [state blocks] + # Calculate block structure + first_image_len = frame_seqlen + action_len = action_horizon + state_len = state_horizon + image_blocks_len = total_len - first_image_len - action_len - state_len + + num_image_blocks = image_blocks_len // (num_frame_per_block * frame_seqlen) + num_action_blocks = action_horizon // num_action_per_block + num_state_blocks = state_horizon // num_state_per_block + + assert num_image_blocks == num_action_blocks == num_state_blocks + + # Token ranges + first_image_start = 0 + first_image_end = first_image_len + image_blocks_start = first_image_end + image_blocks_end = image_blocks_start + image_blocks_len + action_start = image_blocks_end + action_end = action_start + action_len + state_start = action_end + state_end = state_start + state_len + + # Visualize attention mask if requested + if visualize_mask: + mask = self._visualize_attention_mask( + total_len, first_image_len, image_blocks_len, + action_len, state_len, num_image_blocks, + num_action_blocks, num_state_blocks, + num_frame_per_block, frame_seqlen, + num_action_per_block, num_state_per_block + ) + + print("\n" + "="*80) + print("ATTENTION MASK VISUALIZATION") + print("="*80) + print(f"Total length: {total_len}") + print(f"First image: [{first_image_start}:{first_image_end}] (len={first_image_len})") + print(f"Image blocks: [{image_blocks_start}:{image_blocks_end}] (len={image_blocks_len}, num_blocks={num_image_blocks})") + print(f"Action tokens: [{action_start}:{action_end}] (len={action_len}, num_blocks={num_action_blocks})") + print(f"State tokens: [{state_start}:{state_end}] (len={state_len}, num_blocks={num_state_blocks})") + print(f"Local attention size: {self.local_attn_size}") + print("-"*80) + + # Print a downsampled version of the mask if it's too large + if total_len <= 100: + # Print full mask for small sequences + print("Attention mask (1=can attend, 0=cannot attend):") + print("Rows=Query tokens, Cols=Key tokens") + for i in range(total_len): + row = "".join(["1" if mask[i, j] else "." for j in range(total_len)]) + print(f"{i:4d}: {row}") + else: + # Print downsampled version for large sequences + downsample = max(1, total_len // 100) + print(f"Attention mask (downsampled by {downsample}x):") + print("Rows=Query tokens, Cols=Key tokens (1=can attend, .=cannot attend)") + for i in range(0, total_len, downsample): + row = "".join(["1" if mask[i, j] else "." for j in range(0, total_len, downsample)]) + print(f"{i:4d}: {row}") + + # Save mask as image + try: + import cv2 + import numpy as np + mask_np = mask.cpu().float().numpy() + # Resize for visualization if needed + if total_len > 1000: + mask_np = cv2.resize(mask_np, (1000, 1000), interpolation=cv2.INTER_NEAREST) + mask_img = (mask_np * 255).astype(np.uint8) + cv2.imwrite("attention_mask_blockwise_flash.png", mask_img) + print(f"\nMask saved to: attention_mask_blockwise_flash.png") + except Exception as e: + print(f"Could not save mask image: {e}") + + print("="*80 + "\n") + + # OPTIMIZED: Pre-allocate output tensor and pre-compute all indices + output = torch.empty_like(q) + + # Process first image (conditioning, can only self-attend) + output[:, first_image_start:first_image_end] = self.attn( + q[:, first_image_start:first_image_end], + k[:, first_image_start:first_image_end], + v[:, first_image_start:first_image_end] + ) + + # Pre-compute all block indices for image blocks + image_block_starts = [image_blocks_start + i * num_frame_per_block * frame_seqlen for i in range(num_image_blocks)] + image_block_ends = [image_blocks_start + (i + 1) * num_frame_per_block * frame_seqlen for i in range(num_image_blocks)] + if self.local_attn_size != -1: + image_kv_starts = [max(image_blocks_start, end - self.local_attn_size * frame_seqlen) for end in image_block_ends] + else: + image_kv_starts = [image_blocks_start] * num_image_blocks + + # Pre-compute action and state block indices + action_block_starts = [action_start + i * num_action_per_block for i in range(num_action_blocks)] + action_block_ends = [action_start + (i + 1) * num_action_per_block for i in range(num_action_blocks)] + state_block_starts = [state_start + i * num_state_per_block for i in range(num_state_blocks)] + state_block_ends = [state_start + (i + 1) * num_state_per_block for i in range(num_state_blocks)] + + # Process each image block + for block_idx in range(num_image_blocks): + block_start = image_block_starts[block_idx] + block_end = image_block_ends[block_idx] + image_kv_start = image_kv_starts[block_idx] + action_block_start = action_block_starts[block_idx] + action_block_end = action_block_ends[block_idx] + state_block_start = state_block_starts[block_idx] + state_block_end = state_block_ends[block_idx] + + # Build context: first image + relevant image blocks + current action + current state + k_context = torch.cat([ + k[:, first_image_start:first_image_end], # First image + k[:, image_kv_start:block_end], # Image blocks + k[:, action_block_start:action_block_end], # Current action block + k[:, state_block_start:state_block_end] # Current state block + ], dim=1) + v_context = torch.cat([ + v[:, first_image_start:first_image_end], + v[:, image_kv_start:block_end], + v[:, action_block_start:action_block_end], + v[:, state_block_start:state_block_end] + ], dim=1) + + output[:, block_start:block_end] = self.attn( + q[:, block_start:block_end], k_context, v_context + ) + + # Process each action block + for block_idx in range(num_action_blocks): + action_block_start = action_block_starts[block_idx] + action_block_end = action_block_ends[block_idx] + image_block_end = image_block_ends[block_idx] + state_block_start = state_block_starts[block_idx] + state_block_end = state_block_ends[block_idx] + + # Determine image context range + if self.local_attn_size != -1: + image_kv_start = max(image_blocks_start, image_block_end - self.local_attn_size * frame_seqlen) + else: + image_kv_start = image_blocks_start + + # Build context + k_context = torch.cat([ + k[:, first_image_start:first_image_end], # First image + k[:, image_kv_start:image_block_end], # Image blocks + k[:, action_block_start:action_block_end], # Current action block + k[:, state_block_start:state_block_end] # Current state block + ], dim=1) + v_context = torch.cat([ + v[:, first_image_start:first_image_end], + v[:, image_kv_start:image_block_end], + v[:, action_block_start:action_block_end], + v[:, state_block_start:state_block_end] + ], dim=1) + + output[:, action_block_start:action_block_end] = self.attn( + q[:, action_block_start:action_block_end], k_context, v_context + ) + + # Process state blocks (conditioning, can only self-attend) + for block_idx in range(num_state_blocks): + state_block_start = state_block_starts[block_idx] + state_block_end = state_block_ends[block_idx] + + output[:, state_block_start:state_block_end] = self.attn( + q[:, state_block_start:state_block_end], + k[:, state_block_start:state_block_end], + v[:, state_block_start:state_block_end] + ) + + return output + + def _process_clean_image_only(self, clean_image_q, clean_image_k, clean_image_v, clean_frames): + """Process clean image blocks with causal attention pattern - OPTIMIZED + + First frame: conditioning, cannot attend to anything (self-attention only) + Block i: attends to first frame + previous blocks (0 to i-1) + current block + + OPTIMIZATION: Instead of looping through blocks, we batch process them together + by using a single flash_attention call with properly structured KV cache. + """ + block_size = self.frame_seqlen * self.num_frame_per_block + num_blocks = (clean_frames - 1) // self.num_frame_per_block + + if num_blocks == 0: + # Only first frame - single attention call + return self.attn( + clean_image_q[:, :self.frame_seqlen], + clean_image_k[:, :self.frame_seqlen], + clean_image_v[:, :self.frame_seqlen] + ) + + # Pre-allocate output tensor (avoids list append + cat overhead) + b, total_len, n, d = clean_image_q.shape + output = torch.empty_like(clean_image_q) + + # First frame: conditioning, self-attention only + output[:, :self.frame_seqlen] = self.attn( + clean_image_q[:, :self.frame_seqlen], + clean_image_k[:, :self.frame_seqlen], + clean_image_v[:, :self.frame_seqlen] + ) + + # OPTIMIZATION: Process all blocks together with causal masking + # For global attention (no local_attn_size), we can process all blocks in one call + if self.local_attn_size == -1: + # Single attention call for all blocks! + # Each position can attend to first_frame + everything up to itself + blocks_q = clean_image_q[:, self.frame_seqlen:] + blocks_k = clean_image_k # Can attend to everything including first frame + blocks_v = clean_image_v + + # Use causal masking: each block token can see first frame + all previous tokens + output[:, self.frame_seqlen:] = self.causal_attn( + blocks_q, blocks_k, blocks_v + ) + else: + # With local attention, we still need to loop but with optimizations + # Pre-compute all block boundaries to reduce overhead + block_starts = [self.frame_seqlen + i * block_size for i in range(num_blocks)] + block_ends = [min(start + block_size, total_len) for start in block_starts] + + for block_idx in range(num_blocks): + block_start = block_starts[block_idx] + block_end = block_ends[block_idx] + + q_block = clean_image_q[:, block_start:block_end] + + # Context: first frame + recent blocks within local_attn_size + image_kv_start = max(self.frame_seqlen, block_end - self.local_attn_size * self.frame_seqlen) + k_context = torch.cat([ + clean_image_k[:, :self.frame_seqlen], # First frame + clean_image_k[:, image_kv_start:block_end] # Recent blocks + current + ], dim=1) + v_context = torch.cat([ + clean_image_v[:, :self.frame_seqlen], + clean_image_v[:, image_kv_start:block_end] + ], dim=1) + + output[:, block_start:block_end] = self.attn(q_block, k_context, v_context) + + return output + + def _process_state_blocks(self, state_q, state_k, state_v, state_horizon): + """Process state blocks: self-attention only - OPTIMIZED + + OPTIMIZATION: State blocks only do self-attention within each block. + Instead of looping, we can process all blocks in a single call with block-diagonal masking, + or even simpler: just one attention call since they're independent. + """ + num_blocks = state_horizon // self.num_state_per_block + + if num_blocks == 1: + # Single block - one attention call + return self.attn(state_q, state_k, state_v) + + # OPTIMIZATION: Since each state block only attends to itself (no cross-block attention), + # we can process all blocks in a single batched call. Flash attention will handle this + # efficiently. The blocks are independent, so this is safe. + # Alternative: reshape and process as separate batch items + + # Pre-allocate output + output = torch.empty_like(state_q) + + # Process all blocks (keeping loop for now due to block-diagonal pattern) + # This could be further optimized with custom masking + for block_idx in range(num_blocks): + state_block_start = block_idx * self.num_state_per_block + state_block_end = state_block_start + self.num_state_per_block + + output[:, state_block_start:state_block_end] = self.attn( + state_q[:, state_block_start:state_block_end], + state_k[:, state_block_start:state_block_end], + state_v[:, state_block_start:state_block_end] + ) + + return output + + def _process_noisy_image_blocks(self, noisy_image_q, noisy_image_k, noisy_image_v, + clean_image_k, clean_image_v, + noisy_action_k, noisy_action_v, noisy_state_k, noisy_state_v, + half_frames, action_horizon, state_horizon): + """Process noisy image blocks with teacher forcing pattern. + + Matches reference behavior: per-block slicing of action/state tokens. + With action_horizon=24 and num_action_per_block=24, only block 0 gets + action/state context (blocks 1+ slice beyond tensor bounds → empty). + + First frame: conditioning, self-attention only. + Block i: attends to first_clean_frame + clean_blocks[0:i] + current_noisy_block + + action[i*num_action_per_block:(i+1)*num_action_per_block] + + state[i*num_state_per_block:(i+1)*num_state_per_block] + """ + block_size = self.frame_seqlen * self.num_frame_per_block + num_blocks = (half_frames - 1) // self.num_frame_per_block + + # Pre-allocate output tensor + output = torch.empty_like(noisy_image_q) + + # First noisy frame: conditioning, self-attention only + output[:, :self.frame_seqlen] = self.attn( + noisy_image_q[:, :self.frame_seqlen], + noisy_image_k[:, :self.frame_seqlen], + noisy_image_v[:, :self.frame_seqlen] + ) + + if num_blocks == 0: + return output + + # Pre-compute all block indices (matches reference) + noisy_block_starts = [self.frame_seqlen + i * block_size for i in range(num_blocks)] + noisy_block_ends = [min(start + block_size, noisy_image_q.shape[1]) for start in noisy_block_starts] + clean_context_ends = [self.frame_seqlen + i * block_size for i in range(num_blocks)] + action_block_starts = [i * self.num_action_per_block for i in range(num_blocks)] + action_block_ends = [start + self.num_action_per_block for start in action_block_starts] + state_block_starts = [i * self.num_state_per_block for i in range(num_blocks)] + state_block_ends = [start + self.num_state_per_block for start in state_block_starts] + + # Process noisy image blocks + for block_idx in range(num_blocks): + noisy_start = noisy_block_starts[block_idx] + noisy_end = noisy_block_ends[block_idx] + clean_end = clean_context_ends[block_idx] + action_start = action_block_starts[block_idx] + action_end = action_block_ends[block_idx] + state_start = state_block_starts[block_idx] + state_end = state_block_ends[block_idx] + + q_block = noisy_image_q[:, noisy_start:noisy_end] + + # Build context: first_clean_frame + clean_blocks[0:i] + current_noisy_block + # + action[i] + state[i] + # NOTE: For blocks > 0, action/state slices are beyond tensor bounds + # (empty tensors) — this matches reference behavior exactly. + k_context = torch.cat([ + clean_image_k[:, :clean_end], + noisy_image_k[:, noisy_start:noisy_end], + noisy_action_k[:, action_start:action_end], + noisy_state_k[:, state_start:state_end], + ], dim=1) + v_context = torch.cat([ + clean_image_v[:, :clean_end], + noisy_image_v[:, noisy_start:noisy_end], + noisy_action_v[:, action_start:action_end], + noisy_state_v[:, state_start:state_end], + ], dim=1) + + output[:, noisy_start:noisy_end] = self.attn(q_block, k_context, v_context) + + return output + + def _process_noisy_action_blocks(self, noisy_action_q, noisy_action_k, noisy_action_v, + clean_image_k, clean_image_v, + noisy_image_k, noisy_image_v, + noisy_state_k, noisy_state_v, + half_frames, action_horizon, state_horizon): + """Process noisy action blocks with teacher forcing pattern. + + Matches reference behavior: per-block slicing. + Action block i: attends to first_clean_frame + clean_blocks[0:i] + + noisy_image[i] + action[i] + state[i] + + With action_horizon=24 and num_action_per_block=24, only block 0 has + actual tokens; blocks 1+ slice beyond bounds → empty queries, producing + empty output (no contribution to final sequence). + """ + num_blocks = (half_frames - 1) // self.num_frame_per_block + + if num_blocks == 0: + return torch.empty_like(noisy_action_q) + + # Pre-allocate output tensor + output = torch.empty_like(noisy_action_q) + + # Pre-compute all block indices (matches reference) + action_block_starts = [i * self.num_action_per_block for i in range(num_blocks)] + action_block_ends = [start + self.num_action_per_block for start in action_block_starts] + clean_context_ends = [self.frame_seqlen + i * self.frame_seqlen * self.num_frame_per_block for i in range(num_blocks)] + noisy_image_block_starts = [self.frame_seqlen + i * self.frame_seqlen * self.num_frame_per_block for i in range(num_blocks)] + noisy_image_block_ends = [start + self.frame_seqlen * self.num_frame_per_block for start in noisy_image_block_starts] + state_block_starts = [i * self.num_state_per_block for i in range(num_blocks)] + state_block_ends = [start + self.num_state_per_block for start in state_block_starts] + + # Process noisy action blocks + for block_idx in range(num_blocks): + action_start = action_block_starts[block_idx] + action_end = action_block_ends[block_idx] + clean_end = clean_context_ends[block_idx] + noisy_img_start = noisy_image_block_starts[block_idx] + noisy_img_end = noisy_image_block_ends[block_idx] + state_start = state_block_starts[block_idx] + state_end = state_block_ends[block_idx] + + q_block = noisy_action_q[:, action_start:action_end] + + # Skip if this block has no query tokens (OOB slice → empty) + if q_block.shape[1] == 0: + continue + + # Build context: first_clean_frame + clean_blocks[0:i] + noisy_image[i] + # + action[i] + state[i] + k_context = torch.cat([ + clean_image_k[:, :clean_end], + noisy_image_k[:, noisy_img_start:noisy_img_end], + noisy_action_k[:, action_start:action_end], + noisy_state_k[:, state_start:state_end], + ], dim=1) + v_context = torch.cat([ + clean_image_v[:, :clean_end], + noisy_image_v[:, noisy_img_start:noisy_img_end], + noisy_action_v[:, action_start:action_end], + noisy_state_v[:, state_start:state_end], + ], dim=1) + + output[:, action_start:action_end] = self.attn(q_block, k_context, v_context) + + return output + + def forward( + self, + x: torch.Tensor, + freqs: torch.Tensor, + freqs_action: torch.Tensor, + freqs_state: torch.Tensor, + action_register_length: int | None, + kv_cache: torch.Tensor | None = None, + current_start_frame: int = 0, + is_tf: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + r""" + Args: + x(Tensor): Shape [B, L, num_heads, C / num_heads] + seq_lens(Tensor): Shape [B] + freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2] + """ + b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim + + # query, key, value function + def qkv_fn(x): + q = self.norm_q(self.q(x)).view(b, s, n, d) + k = self.norm_k(self.k(x)).view(b, s, n, d) + v = self.v(x).view(b, s, n, d) + return q, k, v + + q, k, v = qkv_fn(x) + + updated_kv_cache: torch.Tensor | None = None + + if kv_cache is None: + if is_tf: + # Teacher forcing training. + if action_register_length is not None: + q_context = q[:, :(s-action_register_length)//2] + k_context = k[:, :(s-action_register_length)//2] + q_noisy = q[:, (s-action_register_length)//2:] + k_noisy = k[:, (s-action_register_length)//2:] + else: + q_context = q[:, :s//2] + k_context = k[:, :s//2] + q_noisy = q[:, s//2:] + k_noisy = k[:, s//2:] + roped_query = [] + roped_key = [] + + # rope should be same for clean and noisy parts + rq_context = rope_action_apply( + x=q_context, + freqs=freqs, + freqs_action=freqs_action, + freqs_state=freqs_state, + action_register_length=None, + ).type_as(v) + rk_context = rope_action_apply( + x=k_context, + freqs=freqs, + freqs_action=freqs_action, + freqs_state=freqs_state, + action_register_length=None, + ).type_as(v) + + rq_noisy = rope_action_apply( + x=q_noisy, + freqs=freqs, + freqs_action=freqs_action, + freqs_state=freqs_state, + action_register_length=action_register_length, + num_action_per_block=self.num_action_per_block, + num_state_per_block=self.num_state_per_block, + ).type_as(v) + rk_noisy = rope_action_apply( + x=k_noisy, + freqs=freqs, + freqs_action=freqs_action, + freqs_state=freqs_state, + action_register_length=action_register_length, + num_action_per_block=self.num_action_per_block, + num_state_per_block=self.num_state_per_block, + ).type_as(v) + + roped_query.append(rq_context) + roped_key.append(rk_context) + roped_query.append(rq_noisy) + roped_key.append(rk_noisy) + + roped_query = torch.cat(roped_query, dim=1) + roped_key = torch.cat(roped_key, dim=1) + # Calculate sequence dimensions + half_seq_len = (s - (action_register_length if action_register_length is not None else 0)) // 2 + + if action_register_length is not None: + # Teacher forcing structure: + # Clean half: [image tokens only] + # Noisy half: [image tokens][action tokens][state tokens] + # Causality only applies to image blocks! + + # Clean half contains ONLY image tokens + clean_image_seq_len = half_seq_len + clean_frames = clean_image_seq_len // self.frame_seqlen + + # Noisy half contains image + action + state tokens + noisy_image_seq_len = half_seq_len + noisy_frames = noisy_image_seq_len // self.frame_seqlen + + # action_register_length is the TOTAL register tokens appended: + # = num_action_tokens + num_state_tokens + # Per-block slicing means only block 0 gets actual tokens; + # blocks 1+ slice beyond bounds → empty (matches reference). + action_horizon = action_register_length - self.num_state_per_block + state_horizon = self.num_state_per_block + + # Validate sequence length matches layout + expected_total = half_seq_len + noisy_image_seq_len + action_horizon + state_horizon + if roped_query.shape[1] != expected_total: + raise ValueError( + f"Sequence length does not match block layout. " + f"action_register_length={action_register_length}, " + f"action_horizon={action_horizon}, state_horizon={state_horizon}. " + f"Expected total={expected_total}, got={roped_query.shape[1]}." + ) + + # Split clean and noisy parts + # Clean: [image tokens only] + clean_image_q = roped_query[:, :clean_image_seq_len] + clean_image_k = roped_key[:, :clean_image_seq_len] + clean_image_v = v[:, :clean_image_seq_len] + + # Noisy: [image tokens][action tokens][state tokens] + noisy_image_q = roped_query[:, half_seq_len:half_seq_len + noisy_image_seq_len] + noisy_action_q = roped_query[:, half_seq_len + noisy_image_seq_len:half_seq_len + noisy_image_seq_len + action_horizon] + noisy_state_q = roped_query[:, half_seq_len + noisy_image_seq_len + action_horizon:] + + noisy_image_k = roped_key[:, half_seq_len:half_seq_len + noisy_image_seq_len] + noisy_action_k = roped_key[:, half_seq_len + noisy_image_seq_len:half_seq_len + noisy_image_seq_len + action_horizon] + noisy_state_k = roped_key[:, half_seq_len + noisy_image_seq_len + action_horizon:] + + noisy_image_v = v[:, half_seq_len:half_seq_len + noisy_image_seq_len] + noisy_action_v = v[:, half_seq_len + noisy_image_seq_len:half_seq_len + noisy_image_seq_len + action_horizon] + noisy_state_v = v[:, half_seq_len + noisy_image_seq_len + action_horizon:] + + # ========== Process CLEAN (context) image tokens ========== + # Clean images: simple blockwise causal attention (no action/state) + clean_image_outputs = self._process_clean_image_only( + clean_image_q, clean_image_k, clean_image_v, clean_frames) + + # ========== Process NOISY tokens ========== + # Noisy image blocks: attend to previous clean image blocks + current noisy image + current noisy action + current noisy state + noisy_image_outputs = self._process_noisy_image_blocks( + noisy_image_q, noisy_image_k, noisy_image_v, + clean_image_k, clean_image_v, + noisy_action_k, noisy_action_v, noisy_state_k, noisy_state_v, + noisy_frames, action_horizon, state_horizon) + + # Noisy action blocks: attend to previous clean image blocks (including first) + current noisy image + current noisy action + same state + noisy_action_outputs = self._process_noisy_action_blocks( + noisy_action_q, noisy_action_k, noisy_action_v, + clean_image_k, clean_image_v, + noisy_image_k, noisy_image_v, + noisy_state_k, noisy_state_v, + noisy_frames, action_horizon, state_horizon) + + # Noisy state blocks: self-attention only + noisy_state_outputs = self._process_state_blocks( + noisy_state_q, noisy_state_k, noisy_state_v, state_horizon) + + # Concatenate all outputs in order: clean_img, noisy_img, noisy_act, noisy_state + x = torch.cat([ + clean_image_outputs, + noisy_image_outputs, noisy_action_outputs, noisy_state_outputs + ], dim=1) + else: + # No action/state tokens, fall back to simple image-only teacher forcing + half_frames = half_seq_len // self.frame_seqlen + clean_q = roped_query[:, :half_seq_len] + clean_k = roped_key[:, :half_seq_len] + clean_v = v[:, :half_seq_len] + noisy_q = roped_query[:, half_seq_len:] + noisy_k = roped_key[:, half_seq_len:] + noisy_v = v[:, half_seq_len:] + + # Process clean frames with blockwise causal attention + x_clean = self._blockwise_causal_flash_attn( + clean_q, clean_k, clean_v, self.frame_seqlen, self.num_frame_per_block, + action_horizon=None, state_horizon=None, + num_action_per_block=None, num_state_per_block=None, + visualize_mask=False) + + # Process noisy frames: attend to all clean frames + themselves + full_k = torch.cat([clean_k, noisy_k], dim=1) + full_v = torch.cat([clean_v, noisy_v], dim=1) + x_noisy = self.attn(noisy_q, full_k, full_v) + + x = torch.cat([x_clean, x_noisy], dim=1) + + else: + roped_query = rope_action_apply( + x=q, + freqs=freqs, + freqs_action=freqs_action, + freqs_state=freqs_state, + action_register_length=action_register_length, + num_action_per_block=self.num_action_per_block, + num_state_per_block=self.num_state_per_block, + ).type_as(v) + roped_key = rope_action_apply( + x=k, + freqs=freqs, + freqs_action=freqs_action, + freqs_state=freqs_state, + action_register_length=action_register_length, + num_action_per_block=self.num_action_per_block, + num_state_per_block=self.num_state_per_block, + ).type_as(v) + + # Calculate dynamic action and state horizons + if action_register_length is not None: + chunk_size = action_register_length // (self.num_action_per_block + self.num_state_per_block) + action_horizon = chunk_size * self.num_action_per_block + state_horizon = chunk_size * self.num_state_per_block + else: + action_horizon = None + state_horizon = None + + # Use blockwise causal flash attention without massive padding + visualize = False + x = self._blockwise_causal_flash_attn( + roped_query, roped_key, v, self.frame_seqlen, self.num_frame_per_block, + action_horizon=action_horizon, + state_horizon=state_horizon, + num_action_per_block=self.num_action_per_block if action_register_length else None, + num_state_per_block=self.num_state_per_block if action_register_length else None, + visualize_mask=visualize) + + else: + action_state_index = (current_start_frame - 1) // self.num_frame_per_block + + roped_query = causal_rope_action_apply( + x=q, + freqs=freqs, + freqs_action=freqs_action, + freqs_state=freqs_state, + action_register_length=action_register_length, + num_action_per_block=self.num_action_per_block, + num_state_per_block=self.num_state_per_block, + action_state_index=action_state_index, + ).type_as(v) + roped_key = causal_rope_action_apply( + x=k, + freqs=freqs, + freqs_action=freqs_action, + freqs_state=freqs_state, + action_register_length=action_register_length, + num_action_per_block=self.num_action_per_block, + num_state_per_block=self.num_state_per_block, + action_state_index=action_state_index, + ).type_as(v) + + # split roped_query and roped_action_query (the last action_register_length tokens) + roped_action_query: torch.Tensor | None = None + roped_action_key: torch.Tensor | None = None + action_v: torch.Tensor | None = None + + if action_register_length is not None: + roped_action_query = roped_query[:, -action_register_length:] + roped_query = roped_query[:, :-action_register_length] + roped_action_key = roped_key[:, -action_register_length:] + roped_key = roped_key[:, :-action_register_length] + action_v = v[:, -action_register_length:] + v = v[:, :-action_register_length] + assert roped_action_query is not None + assert roped_action_key is not None + assert action_v is not None + + num_new_tokens = roped_query.shape[1] + assert roped_key.shape[1] == num_new_tokens + assert v.shape[1] == num_new_tokens + + # If we are using local attention and the current KV cache size is larger + # than the local attention size, we need to truncate the KV cache + + updated_kv_cache = kv_cache + updated_k = updated_kv_cache[0] + updated_v = updated_kv_cache[1] + # Assign new keys/values directly up to current_end + new_k = torch.cat([updated_k, roped_key], dim=1) + new_v = torch.cat([updated_v, v], dim=1) + + # We may need to truncate the KV cache if it's size is larger than the max attention size. + new_k = new_k[:, -self.max_attention_size:] + new_v = new_v[:, -self.max_attention_size:] + + if action_register_length is not None: + x = self.attn( + torch.cat([roped_query, roped_action_query], dim=1), + torch.cat([new_k, roped_action_key], dim=1), + torch.cat([new_v, action_v], dim=1), + ) + else: + x = self.attn( + roped_query, + new_k, + new_v, + ) + updated_kv_cache = torch.stack([new_k, new_v], dim=0) + + + # output + x = x.flatten(2) + x = self.o(x) + return x, updated_kv_cache + + +class CausalWanAttentionBlock(nn.Module): + + def __init__(self, + cross_attn_type, + dim, + ffn_dim, + num_heads, + frame_seqlen, + local_attn_size=-1, + sink_size=0, + num_frame_per_block=1, + qk_norm=True, + cross_attn_norm=False, + eps=1e-6, + num_action_per_block=32, + num_state_per_block=1): + super().__init__() + self.dim = dim + self.ffn_dim = ffn_dim + self.num_heads = num_heads + self.local_attn_size = local_attn_size + self.qk_norm = qk_norm + self.cross_attn_norm = cross_attn_norm + self.eps = eps + + # layers + self.norm1 = WanLayerNorm(dim, eps) + self.self_attn = CausalWanSelfAttention( + dim=dim, + num_heads=num_heads, + frame_seqlen=frame_seqlen, + local_attn_size=local_attn_size, + sink_size=sink_size, + num_frame_per_block=num_frame_per_block, + qk_norm=qk_norm, + eps=eps, + num_action_per_block=num_action_per_block, + num_state_per_block=num_state_per_block, + ) + self.norm3 = WanLayerNorm( + dim, eps, + elementwise_affine=True) if cross_attn_norm else nn.Identity() + self.cross_attn = WAN_CROSSATTENTION_CLASSES[cross_attn_type](dim, + num_heads, + (-1, -1), + qk_norm, + eps) + self.norm2 = WanLayerNorm(dim, eps) + self.ffn = nn.Sequential( + nn.Linear(dim, ffn_dim), nn.GELU(approximate='tanh'), + nn.Linear(ffn_dim, dim)) + + # modulation + self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + + def forward( + self, + x: torch.Tensor, + e: torch.Tensor, + freqs: torch.Tensor, + freqs_action: torch.Tensor, + freqs_state: torch.Tensor, + action_register_length: int | None, + context: torch.Tensor, + kv_cache: torch.Tensor | None = None, + crossattn_cache: torch.Tensor | None = None, + current_start_frame: int = 0, + is_tf: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + r""" + Args: + x(Tensor): Shape [B, L, C] + e(Tensor): Shape [B, F, 6, C] + freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2] + """ + e = (self.modulation.unsqueeze(1) + e).chunk(6, dim=2) + + # Align modulation sequence length to x so mul/add broadcast (e.g. when F != L under compile) + L = x.shape[1] + aligned = [] + for part in e: + L_e = part.shape[1] + if L_e == L: + aligned.append(part) + elif L_e >= L: + aligned.append(part[:, :L]) + else: + repeat = (L + L_e - 1) // L_e + aligned.append(part.repeat_interleave(repeat, dim=1)[:, :L]) + e = tuple(aligned) + + # self-attention + y, updated_kv_cache = self.self_attn( + x=(self.norm1(x) * (1 + e[1].squeeze(2)) + e[0].squeeze(2)), + freqs=freqs, + freqs_action=freqs_action, + freqs_state=freqs_state, + action_register_length=action_register_length, + kv_cache=kv_cache, + is_tf=is_tf, + current_start_frame=current_start_frame, + ) + x = x + (y * e[2].squeeze(2)) + + # cross-attention & ffn function + def cross_attn_ffn(x, context, e): + x = x + self.cross_attn(self.norm3(x), context) + y = self.ffn( + (self.norm2(x) * (1 + e[4].squeeze(2)) + e[3].squeeze(2)) + ) + x = x + (y * e[5].squeeze(2)) + return x + + x = cross_attn_ffn(x, context, e) + return x, updated_kv_cache + + +class CausalHead(nn.Module): + + def __init__(self, dim, out_dim, patch_size, eps=1e-6): + super().__init__() + self.dim = dim + self.out_dim = out_dim + self.patch_size = patch_size + self.eps = eps + + # layers + out_dim = math.prod(patch_size) * out_dim + self.norm = WanLayerNorm(dim, eps) + self.head = nn.Linear(dim, out_dim) + + # modulation + self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5) + + def forward(self, x, e): + r""" + Args: + x(Tensor): Shape [B, L1, C] + e(Tensor): Shape [B, F, 1, C] + """ + e = (self.modulation.unsqueeze(1) + e).chunk(2, dim=2) + # Align modulation sequence length to x (e.g. when F != L1 under compile) + L = x.shape[1] + aligned = [] + for part in e: + L_e = part.shape[1] + if L_e == L: + aligned.append(part) + elif L_e >= L: + aligned.append(part[:, :L]) + else: + repeat = (L + L_e - 1) // L_e + aligned.append(part.repeat_interleave(repeat, dim=1)[:, :L]) + e = tuple(aligned) + x = (self.head(self.norm(x) * (1 + e[1].squeeze(2)) + e[0].squeeze(2))) + return x + + +class CausalWanModel(ModelMixin, ConfigMixin): + r""" + Wan diffusion backbone supporting both text-to-video and image-to-video. + """ + + ignore_for_config = [ + 'patch_size', 'cross_attn_norm', 'qk_norm', 'text_dim' + ] + _no_split_modules = ['WanAttentionBlock'] + _supports_gradient_checkpointing = True + + @register_to_config + def __init__(self, + model_type='t2v', + patch_size=(1, 2, 2), + frame_seqlen=220, + text_len=512, + in_dim=16, + dim=2048, + ffn_dim=8192, + freq_dim=256, + text_dim=4096, + out_dim=16, + num_heads=16, + num_layers=32, + max_chunk_size=-1, + sink_size=0, + qk_norm=True, + cross_attn_norm=True, + eps=1e-6, + num_frame_per_block=1, + action_dim=32, + num_registers=8, + max_state_dim=64, + max_num_embodiments=32, + hidden_size=1024, + diffusion_model_pretrained_path=None, + num_action_per_block=32, + num_state_per_block=1, + concat_first_frame_latent=True): + r""" + Initialize the diffusion model backbone. + + Args: + concat_first_frame_latent (`bool`, *optional*, defaults to True): + If True, concat [x; y] before patch_embedding (14B I2V style). If False, latent only (5B pretrained style; first-frame via CLIP). + model_type (`str`, *optional*, defaults to 't2v'): + Model variant - 't2v' (text-to-video) or 'i2v' (image-to-video) + patch_size (`tuple`, *optional*, defaults to (1, 2, 2)): + 3D patch dimensions for video embedding (t_patch, h_patch, w_patch) + text_len (`int`, *optional*, defaults to 512): + Fixed length for text embeddings + in_dim (`int`, *optional*, defaults to 16): + Input video channels (C_in) + dim (`int`, *optional*, defaults to 2048): + Hidden dimension of the transformer + ffn_dim (`int`, *optional*, defaults to 8192): + Intermediate dimension in feed-forward network + freq_dim (`int`, *optional*, defaults to 256): + Dimension for sinusoidal time embeddings + text_dim (`int`, *optional*, defaults to 4096): + Input dimension for text embeddings + out_dim (`int`, *optional*, defaults to 16): + Output video channels (C_out) + num_heads (`int`, *optional*, defaults to 16): + Number of attention heads + num_layers (`int`, *optional*, defaults to 32): + Number of transformer blocks + local_attn_size (`int`, *optional*, defaults to -1): + Window size for temporal local attention (-1 indicates global attention) + sink_size (`int`, *optional*, defaults to 0): + Size of the attention sink, we keep the first `sink_size` frames unchanged when rolling the KV cache + qk_norm (`bool`, *optional*, defaults to True): + Enable query/key normalization + cross_attn_norm (`bool`, *optional*, defaults to False): + Enable cross-attention normalization + eps (`float`, *optional*, defaults to 1e-6): + Epsilon value for normalization layers + """ + + super().__init__() + + assert model_type in ['t2v', 'i2v', 'ti2v'] + self.model_type = model_type + + self.patch_size = patch_size + self.frame_seqlen = frame_seqlen + self.text_len = text_len + self.in_dim = in_dim + self.dim = dim + self.ffn_dim = ffn_dim + self.freq_dim = freq_dim + self.text_dim = text_dim + self.out_dim = out_dim + self.num_heads = num_heads + self.num_layers = num_layers + self.local_attn_size = max_chunk_size * num_frame_per_block + 1 if max_chunk_size != -1 else -1 + self.qk_norm = qk_norm + self.cross_attn_norm = cross_attn_norm + self.eps = eps + self.num_frame_per_block = num_frame_per_block + self.diffusion_model_pretrained_path = diffusion_model_pretrained_path + self.action_dim = action_dim + self.num_registers = num_registers + self.max_state_dim = max_state_dim + self.max_num_embodiments = max_num_embodiments + self.hidden_size = hidden_size + self.num_action_per_block = num_action_per_block + self.num_state_per_block = num_state_per_block + self.concat_first_frame_latent = concat_first_frame_latent + + max_num_embodiments = 1 + + self.state_encoder = CategorySpecificMLP( + num_categories=max_num_embodiments, + input_dim=max_state_dim, + hidden_dim=self.hidden_size, + output_dim=self.dim, + ) + self.action_encoder = MultiEmbodimentActionEncoder( + action_dim=action_dim, + hidden_size=self.dim, + num_embodiments=max_num_embodiments, + ) + self.action_decoder = CategorySpecificMLP( + num_categories=max_num_embodiments, + input_dim=dim, + hidden_dim=self.hidden_size, + output_dim=action_dim, + ) + + # embeddings + self.patch_embedding = nn.Conv3d( + in_dim, dim, kernel_size=patch_size, stride=patch_size) + self.text_embedding = nn.Sequential( + nn.Linear(text_dim, dim), nn.GELU(approximate='tanh'), + nn.Linear(dim, dim)) + + self.time_embedding = nn.Sequential( + nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim)) + self.time_projection = nn.Sequential( + nn.SiLU(), nn.Linear(dim, dim * 6)) + + # blocks + cross_attn_type = 't2v_cross_attn' if model_type == 't2v' else 'i2v_cross_attn' + self.blocks = nn.ModuleList([ + CausalWanAttentionBlock(cross_attn_type, dim, ffn_dim, num_heads, frame_seqlen, + self.local_attn_size, sink_size, num_frame_per_block, qk_norm, cross_attn_norm, eps, + num_action_per_block, num_state_per_block) + for _ in range(num_layers) + ]) + + # head + self.head = CausalHead(dim, out_dim, patch_size, eps) + + # buffers (don't use register_buffer otherwise dtype will be changed in to()) + assert (dim % num_heads) == 0 and (dim // num_heads) % 2 == 0 + d = dim // num_heads + + self.freqs_action = rope_params(1024*10, d) + self.freqs_state = rope_params(1024, d) + self.freqs = [ + rope_params(1024, d - 4 * (d // 6)), + rope_params(1024, 2 * (d // 6)), + rope_params(1024, 2 * (d // 6)), + ] + if model_type in ('i2v', 'ti2v'): + self.img_emb = MLPProj(1280, dim) + + # initialize weights + self.init_weights() + + self.gradient_checkpointing = True + self.independent_first_frame = False if self.num_frame_per_block == 1 else True + + + def _set_gradient_checkpointing(self, module, value=False): + self.gradient_checkpointing = value + + @staticmethod + def _prepare_blockwise_causal_attn_mask( + device: torch.device | str, num_frames: int = 21, + frame_seqlen: int = 1560, num_frame_per_block=1, local_attn_size=-1, action_horizon=1, state_horizon=1, num_action_per_block=30, num_state_per_block=1 + ) -> BlockMask: + """ + We will divide the token sequence into the following format: + [first image (conditioning)] [image blocks] [action blocks] [state blocks] + + Structure: + - First image: conditioning only, cannot attend to anything + - Image blocks: can attend to first image + previous image block + current action block + current state block + - Action blocks: can attend to previous image block + current image block + current state block + - State blocks: conditioning only, cannot attend to anything + + Block alignment: + - num_image_blocks = (num_frames - 1) // num_frame_per_block + - num_action_blocks = action_horizon // num_action_per_block + - num_state_blocks = state_horizon // num_state_per_block + - num_image_blocks = num_action_blocks + 1 = num_state_blocks + 1 + """ + # Calculate block structure + num_image_blocks = (num_frames - 1) // num_frame_per_block + num_action_blocks = action_horizon // num_action_per_block + num_state_blocks = state_horizon // num_state_per_block + + # Verify the relationship: num_image_blocks = num_action_blocks + 1 = num_state_blocks + 1 + assert num_image_blocks == num_action_blocks, \ + f"image_blocks mismatch: {num_image_blocks} != {num_action_blocks}" + assert num_image_blocks == num_state_blocks, \ + f"image_blocks mismatch: {num_image_blocks} != {num_state_blocks}" + + # Token ranges + first_image_len = frame_seqlen # First image (conditioning) + image_blocks_len = num_image_blocks * num_frame_per_block * frame_seqlen + action_len = action_horizon + state_len = state_horizon + total_length = first_image_len + image_blocks_len + action_len + state_len + + # print("total_length", total_length, first_image_len, image_blocks_len, action_len, state_len) + # Padding to multiple of 128 + # padded_length = math.ceil(total_length / 128) * 128 - total_length + padded_length = math.ceil((local_attn_size * frame_seqlen + (local_attn_size - 1) + 32 * (local_attn_size - 1))/128) * 128 - total_length + total_padded_length = total_length + padded_length + # print("total_padded_length", total_padded_length, total_length, padded_length) + + # Define token ranges for each modality + first_image_start = 0 + first_image_end = first_image_len + image_blocks_start = first_image_end + image_blocks_end = image_blocks_start + image_blocks_len + action_start = image_blocks_end + action_end = action_start + action_len + state_start = action_end + state_end = state_start + state_len + + # Precompute block indices for each token + block_indices = torch.zeros(total_padded_length, device=device, dtype=torch.long) + + # First image gets special block index -1 (conditioning, cannot attend to anything) + block_indices[first_image_start:first_image_end] = -1 + + # Assign block indices for image blocks (0 to num_image_blocks-1) + for block_idx in range(num_image_blocks): + start_idx = image_blocks_start + block_idx * num_frame_per_block * frame_seqlen + end_idx = image_blocks_start + (block_idx + 1) * num_frame_per_block * frame_seqlen + block_indices[start_idx:end_idx] = block_idx + + # Assign block indices for action tokens (0 to num_action_blocks-1) + for block_idx in range(num_action_blocks): + start_idx = action_start + block_idx * num_action_per_block + end_idx = action_start + (block_idx + 1) * num_action_per_block + block_indices[start_idx:end_idx] = block_idx + + # Assign block indices for state tokens (0 to num_state_blocks-1) + for block_idx in range(num_state_blocks): + start_idx = state_start + block_idx * num_state_per_block + end_idx = state_start + (block_idx + 1) * num_state_per_block + block_indices[start_idx:end_idx] = block_idx + + # Padding tokens get block index of last block + 1 (won't attend to anything) + block_indices[total_length:] = num_image_blocks + + def attention_mask(b, h, q_idx, kv_idx): + # Self-attention + self_attn = (q_idx == kv_idx) + + # Determine which modality q and kv belong to + q_is_first_image = (q_idx >= first_image_start) & (q_idx < first_image_end) + q_is_image_block = (q_idx >= image_blocks_start) & (q_idx < image_blocks_end) + q_is_action = (q_idx >= action_start) & (q_idx < action_end) + q_is_state = (q_idx >= state_start) & (q_idx < state_end) + + kv_is_first_image = (kv_idx >= first_image_start) & (kv_idx < first_image_end) + kv_is_image_block = (kv_idx >= image_blocks_start) & (kv_idx < image_blocks_end) + kv_is_action = (kv_idx >= action_start) & (kv_idx < action_end) + kv_is_state = (kv_idx >= state_start) & (kv_idx < state_end) + + q_block = block_indices[q_idx] + kv_block = block_indices[kv_idx] + + # First image query (conditioning) - cannot attend to anything + first_image_mask = q_is_first_image & False + + # Image block query + image_to_first = q_is_image_block & kv_is_first_image # Image block to first image: always allowed + image_to_image = q_is_image_block & kv_is_image_block & (kv_block <= q_block) # Image block to image block: can attend to current and previous image blocks + image_to_action = q_is_image_block & kv_is_action & (kv_block == q_block) # Image block to action: can attend to current action block + image_to_state = q_is_image_block & kv_is_state & (kv_block == q_block) # Image block to state: can attend to current state block + + image_block_mask = image_to_first | image_to_image | image_to_action | image_to_state + + # Action query + action_to_image = q_is_action & kv_is_image_block & (kv_block <= q_block) # Action to image block: can attend to current and all previous image blocks + action_to_action = q_is_action & kv_is_action & (kv_block == q_block) # Action to action: only same block + action_to_state = q_is_action & kv_is_state & (kv_block == q_block) # Action to state: only same block + action_to_first = q_is_action & kv_is_first_image # Action to first image: always allowed + + action_mask = action_to_image | action_to_action | action_to_state | action_to_first + + # State query (conditioning) - cannot attend to anything + state_mask = q_is_state & False + + # Combine all masks + return self_attn | first_image_mask | image_block_mask | action_mask | state_mask + + block_mask = create_block_mask( + attention_mask, B=None, H=None, + Q_LEN=total_padded_length, + KV_LEN=total_padded_length, + _compile=False, device=device + ) + + if not dist.is_initialized() or dist.get_rank() == 0: + print(f"Created blockwise causal attention mask:") + print(f" first_image_tokens={first_image_len} (conditioning)") + print(f" num_image_blocks={num_image_blocks} (blocks of {num_frame_per_block * frame_seqlen})") + print(f" num_action_blocks={num_action_blocks} (blocks of {num_action_per_block})") + print(f" num_state_blocks={num_state_blocks} (blocks of {num_state_per_block})") + print(f" total_length={total_length}, padded_length={padded_length}") + print(block_mask) + + # Debug: materialize a small slice of the mask into 0/1 strings + try: + dense_mask = create_mask( + attention_mask, + B=None, + H=None, + Q_LEN=total_padded_length, + KV_LEN=total_padded_length, + device=device, + )[0, 0] # [Q, K] + preview_q = min(979, dense_mask.shape[0]) + preview_k = min(979, dense_mask.shape[1]) + print("Block mask (preview):") + for qi in range(preview_q): + row = dense_mask[qi, :preview_k].to(torch.int8).tolist() + print(" ".join(str(int(v)) for v in row)) + except Exception as err: + print("[warn] Failed to materialize block mask preview:", err) + + return block_mask + + @staticmethod + def _prepare_teacher_forcing_mask( + device: torch.device | str, num_frames: int = 21, + frame_seqlen: int = 1560, num_frame_per_block=1 + ) -> BlockMask: + """ + we will divide the token sequence into the following format + [1 latent frame] [1 latent frame] ... [1 latent frame] + We use flexattention to construct the attention mask + """ + total_length = num_frames * frame_seqlen * 2 + + # we do right padding to get to a multiple of 128 + padded_length = math.ceil(self.local_attn_size * frame_seqlen/128) * 128 - total_length + # padded_length = math.ceil(total_length / 128) * 128 - total_length + + clean_ends = num_frames * frame_seqlen + # for clean context frames, we can construct their flex attention mask based on a [start, end] interval + context_ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long) + # for noisy frames, we need two intervals to construct the flex attention mask [context_start, context_end] [noisy_start, noisy_end] + noise_context_starts = torch.zeros(total_length + padded_length, device=device, dtype=torch.long) + noise_context_ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long) + noise_noise_starts = torch.zeros(total_length + padded_length, device=device, dtype=torch.long) + noise_noise_ends = torch.zeros(total_length + padded_length, device=device, dtype=torch.long) + + # Block-wise causal mask will attend to all elements that are before the end of the current chunk + attention_block_size = frame_seqlen * num_frame_per_block + frame_indices = torch.arange( + start=0, + end=num_frames * frame_seqlen, + step=attention_block_size, + device=device, dtype=torch.long + ) + + # attention for clean context frames + for start in frame_indices: + context_ends[start:start + attention_block_size] = start + attention_block_size + + noisy_image_start_list = torch.arange( + num_frames * frame_seqlen, total_length, + step=attention_block_size, + device=device, dtype=torch.long + ) + noisy_image_end_list = noisy_image_start_list + attention_block_size + + # attention for noisy frames + for block_index, (start, end) in enumerate(zip(noisy_image_start_list, noisy_image_end_list)): + # attend to noisy tokens within the same block + noise_noise_starts[start:end] = start + noise_noise_ends[start:end] = end + # attend to context tokens in previous blocks + # noise_context_starts[start:end] = 0 + noise_context_ends[start:end] = block_index * attention_block_size + + def attention_mask(b, h, q_idx, kv_idx): + # first design the mask for clean frames + clean_mask = (q_idx < clean_ends) & (kv_idx < context_ends[q_idx]) + # then design the mask for noisy frames + # noisy frames will attend to all clean preceeding clean frames + itself + C1 = (kv_idx < noise_noise_ends[q_idx]) & (kv_idx >= noise_noise_starts[q_idx]) + C2 = (kv_idx < noise_context_ends[q_idx]) & (kv_idx >= noise_context_starts[q_idx]) + noise_mask = (q_idx >= clean_ends) & (C1 | C2) + + eye_mask = q_idx == kv_idx + return eye_mask | clean_mask | noise_mask + + block_mask = create_block_mask(attention_mask, B=None, H=None, Q_LEN=total_length + padded_length, + KV_LEN=total_length + padded_length, _compile=False, device=device) + + if DEBUG: + print(block_mask) + import imageio + import numpy as np + from torch.nn.attention.flex_attention import create_mask + + mask = create_mask(attention_mask, B=None, H=None, Q_LEN=total_length + + padded_length, KV_LEN=total_length + padded_length, device=device) + import cv2 + mask = cv2.resize(mask[0, 0].cpu().float().numpy(), (1024, 1024)) + imageio.imwrite("mask_%d.jpg" % (0), np.uint8(255. * mask)) + + return block_mask + + @staticmethod + def _prepare_blockwise_causal_attn_mask_i2v( + device: torch.device | str, num_frames: int = 21, + frame_seqlen: int = 1560, num_frame_per_block=4, local_attn_size=-1 + ) -> BlockMask: + """ + we will divide the token sequence into the following format + [1 latent frame] [N latent frame] ... [N latent frame] + The first frame is separated out to support I2V generation + We use flexattention to construct the attention mask + """ + total_length = num_frames * frame_seqlen + + # we do right padding to get to a multiple of 128 + padded_length = math.ceil(local_attn_size * frame_seqlen/128) * 128 - total_length + # padded_length = math.ceil(total_length / 128) * 128 - total_length + + ends = torch.zeros(total_length + padded_length, + device=device, dtype=torch.long) + + # special handling for the first frame + ends[:frame_seqlen] = frame_seqlen + + # Block-wise causal mask will attend to all elements that are before the end of the current chunk + frame_indices = torch.arange( + start=frame_seqlen, + end=total_length, + step=frame_seqlen * num_frame_per_block, + device=device + ) + + for idx, tmp in enumerate(frame_indices): + ends[tmp:tmp + frame_seqlen * num_frame_per_block] = tmp + \ + frame_seqlen * num_frame_per_block + + def attention_mask(b, h, q_idx, kv_idx): + if local_attn_size == -1: + return (kv_idx < ends[q_idx]) | (q_idx == kv_idx) + else: + return ((kv_idx < ends[q_idx]) & (kv_idx >= (ends[q_idx] - local_attn_size * frame_seqlen))) | \ + (q_idx == kv_idx) + + block_mask = create_block_mask(attention_mask, B=None, H=None, Q_LEN=total_length + padded_length, + KV_LEN=total_length + padded_length, _compile=False, device=device) + + if not dist.is_initialized() or dist.get_rank() == 0: + print( + f" cache a block wise causal mask with block size of {num_frame_per_block} frames") + print(block_mask) + + return block_mask + + def _forward_blocks( + self, + x: torch.Tensor, + seq_len: int, + freqs: torch.Tensor, + timestep: torch.Tensor, + context: torch.Tensor, + clip_feature: torch.Tensor | None, + embodiment_id: torch.Tensor | None, + action: torch.Tensor | None, + timestep_action: torch.Tensor | None, + state: torch.Tensor | None, + kv_cache: list[torch.Tensor], + current_start_frame: int, + ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor]]: + r""" + Forward pass through the diffusion model blocks. + """ + x = x.flatten(start_dim=2).transpose(1, 2) + + B = x.shape[0] + F = timestep.shape[1] + + if action is not None: + embodiment_id = torch.tensor([0], device=x.device).repeat(x.shape[0]) + action_features = self.action_encoder(action, timestep_action, embodiment_id) + state_features = self.state_encoder(state, embodiment_id) + action_register = torch.cat([action_features, state_features], dim=1) + action_length = action_features.shape[1] + action_register_length = action_register.shape[1] + x = torch.cat([x, action_register], dim=1) + else: + action_features = None + state_features = None + action_length = 0 + action_register_length = None + + # time embeddings: expand to exactly seq_len so e matches x (5B: frame_seqlen=50, 1 frame -> 50 tokens) + if F <= seq_len: + repeat = (seq_len + F - 1) // F + timestep = timestep.repeat_interleave(repeat, dim=1)[:, :seq_len] + else: + indices = torch.linspace(0, F - 1, seq_len, device=timestep.device, dtype=torch.long) + timestep = timestep[:, indices] + + if action is not None: + assert timestep_action is not None + assert state_features is not None + stride = timestep_action.shape[1] // state_features.shape[1] + timestep_state = timestep_action[:, ::stride] + timestep = torch.cat([timestep, timestep_action, timestep_state], dim=1) + + e = self.time_embedding( + sinusoidal_embedding_1d(self.freq_dim, timestep.flatten()).type_as(x)) + e = e.unflatten(dim=0, sizes=(B, -1)) + e0 = self.time_projection(e) + e0 = e0.unflatten(dim=2, sizes=(6, self.dim)) + + # context + context = self.text_embedding(context) + + if clip_feature is not None: + clip_embedding = self.img_emb(clip_feature) + context = torch.cat([clip_embedding, context], dim=1) + + updated_kv_caches: list[torch.Tensor] = [] + for block_index, block in enumerate(self.blocks): + x, updated_kv_cache = block( + x=x, + e=e0, + freqs=freqs, + freqs_action=self.freqs_action, + freqs_state=self.freqs_state, + context=context, + action_register_length=action_register_length, + kv_cache=kv_cache[block_index], + current_start_frame=current_start_frame, + ) + updated_kv_caches.append(updated_kv_cache) + + if action is not None: + action_noise_pred = x[:, seq_len: seq_len + action_length] + action_noise_pred = self.action_decoder(action_noise_pred, embodiment_id) + else: + action_noise_pred = None + + # Build a tensor that contains only video tokens per sample with length = max(video_lens) + x_video = x[:, :seq_len] + e_video = e[:, :seq_len] + + # Unpatchify video-only tokens + x_video = self.head(x_video, e_video.unsqueeze(2)) + + return x_video, action_noise_pred, updated_kv_caches + + + def _forward_inference_trt( + self, + x, + timestep, + context, + kv_cache_packed: torch.Tensor, + y, + clip_feature, + action, + timestep_action, + state, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + + + frame_seqlen = 880 + seq_len = 2*frame_seqlen + kv_cache_seq_len = kv_cache_packed.shape[3] + current_start_frame = kv_cache_seq_len // frame_seqlen + + kv_cache_list = [] + for block_index in range(len(self.blocks)): + kv_cache_list.append(kv_cache_packed[block_index]) + + x_video, action_noise_pred, _ = self._forward_inference( + x=x, + timestep=timestep, + context=context, + seq_len=int(seq_len), + kv_cache=kv_cache_list, + crossattn_cache=None, + y=y, + clip_feature=clip_feature, + action=action, + timestep_action=timestep_action, + state=state, + current_start_frame = current_start_frame, + ) + + return x_video, action_noise_pred + + def _forward_inference_trt_droid( + self, + x, + timestep, + context, + kv_cache_packed: torch.Tensor, + y, + clip_feature, + action, + timestep_action, + state, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + + + frame_seqlen = 880 + seq_len = 2*frame_seqlen + kv_cache_seq_len = kv_cache_packed.shape[3] + current_start_frame = kv_cache_seq_len // frame_seqlen + + kv_cache_list = [] + for block_index in range(len(self.blocks)): + kv_cache_list.append(kv_cache_packed[block_index]) + + x_video, action_noise_pred, _ = self._forward_inference( + x=x, + timestep=timestep, + context=context, + seq_len=int(seq_len), + kv_cache=kv_cache_list, + crossattn_cache=None, + y=y, + clip_feature=clip_feature, + action=action, + timestep_action=timestep_action, + state=state, + current_start_frame = current_start_frame, + ) + + return x_video, action_noise_pred + + + def _forward_inference( + self, + x, + timestep, + context, + seq_len, + kv_cache: list[torch.Tensor], + crossattn_cache: list[torch.Tensor], + current_start_frame: int, + y=None, + clip_feature=None, + action=None, + timestep_action=None, + state=None, + embodiment_id=None, + ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor]]: + r""" + Run the diffusion model with kv caching. + See Algorithm 2 of CausVid paper https://arxiv.org/abs/2412.07772 for details. + This function will be run for num_frame times. + Process the latent frames one by one (1560 tokens each) + + Args: + x (List[Tensor]): + List of input video tensors, each with shape [C_in, F, H, W] + timestep (Tensor): + Diffusion timesteps tensor of shape [B] + context (List[Tensor]): + List of text embeddings each with shape [L, C] + seq_len (`int`): + Maximum sequence length for positional encoding + action (Tensor, *optional*): + Action tensor of shape [B, H, D] + state (Tensor, *optional*): + State tensor of shape [B, H, D] + embodiment_id (Tensor, *optional*): + Embodiment ID tensor of shape [B] + y (List[Tensor], *optional*): + Conditional video inputs for image-to-video mode, same shape as x + clip_feature (Tensor, *optional*): + CLIP image features for image-to-video mode + timestep_action (Tensor, *optional*): + Action timestep tensor of shape [B] + Returns: + List[Tensor]: + List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8] + """ + if self.model_type == 'i2v': + assert clip_feature is not None and y is not None + assert context.shape[1] == self.text_len + + # Concat [x; y] only when pretrained that way (14B). 5B uses latent only, first-frame via CLIP. + if y is not None and self.concat_first_frame_latent: + x = torch.cat([x, y.to(dtype=x.dtype)], dim=1) + + # embeddings + x = self.patch_embedding(x) + grid_size = torch.tensor(x.shape[2:], dtype=torch.long) + + freqs = self._create_freqs( + grid_size=grid_size, + start_frame=current_start_frame, + ) + + x_video, action_noise_pred, updated_kv_caches = self._forward_blocks( + x=x, + seq_len=seq_len, + freqs=freqs, + timestep=timestep, + context=context, + clip_feature=clip_feature, + embodiment_id=embodiment_id, + action=action, + timestep_action=timestep_action, + state=state, + kv_cache=kv_cache, + current_start_frame=current_start_frame, + ) + + # Copy the updated KV caches back to the original KV cache. + x_video = x_video.clone() + if action_noise_pred is not None: + action_noise_pred = action_noise_pred.clone() + #for block_index, updated_kv_cache in enumerate(updated_kv_caches): + # kv_cache[block_index] = updated_kv_cache.clone() + + video_noise_pred = self.unpatchify(x_video, grid_size) + + return video_noise_pred, action_noise_pred, updated_kv_caches + + def _forward_train( + self, + x, + timestep, + timestep_action, + context, + seq_len, + clean_x=None, + aug_t=None, + y=None, + clip_feature=None, + action=None, + state=None, + embodiment_id=None, + ): + r""" + Forward pass through the diffusion model + + Args: + x (List[Tensor]): + List of input video tensors, each with shape [C_in, F, H, W] + t (Tensor): + Diffusion timesteps tensor of shape [B] + context (List[Tensor]): + List of text embeddings each with shape [L, C] + seq_len (`int`): + Maximum sequence length for positional encoding + clip_fea (Tensor, *optional*): + CLIP image features for image-to-video mode + y (List[Tensor], *optional*): + Conditional video inputs for image-to-video mode, same shape as x + + Returns: + List[Tensor]: + List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8] + """ + if self.model_type == 'i2v': + assert clip_feature is not None and y is not None + + # Concat [x; y] only when pretrained that way (14B). 5B uses latent only, first-frame via CLIP. + if y is not None and self.concat_first_frame_latent: + x = torch.cat([x, y.to(dtype=x.dtype)], dim=1) + + # embeddings + x = self.patch_embedding(x) + + grid_size = torch.tensor(x.shape[2:], dtype=torch.long) + freqs = self._create_freqs( + grid_size=grid_size, + start_frame=0, + ) + + x = x.flatten(start_dim=2).transpose(1, 2) + assert x.shape[1] == seq_len + + B = x.shape[0] + F = timestep.shape[1] + + # time embeddings + if action is not None: + embodiment_id = torch.tensor([0]).repeat(x.shape[0]).to(device=embodiment_id.device) + action_features = self.action_encoder(action, timestep_action, embodiment_id) + action_length = action_features.shape[1] + state_features = self.state_encoder(state, embodiment_id) + action_register = torch.cat([action_features, state_features], dim=1) + action_register_length = action_register.shape[1] + x = torch.cat([x, action_register], dim=1) + else: + action_features = None + action_length = None + state_features = None + action_register = None + action_register_length = None + + # time embeddings + timestep = timestep.unsqueeze(-1).expand(B, F, seq_len // F).reshape(B, -1) + timestep_original = timestep.clone() + + if action is not None: + assert timestep_action is not None + assert state_features is not None + stride = timestep_action.shape[1] // state_features.shape[1] + timestep_state = timestep_action[:, ::stride] + timestep = torch.cat([timestep, timestep_action, timestep_state], dim=1) + + e = self.time_embedding( + sinusoidal_embedding_1d(self.freq_dim, timestep.flatten()).type_as(x)) + e = e.unflatten(dim=0, sizes=(B, -1)) + e0 = self.time_projection(e) + e0 = e0.unflatten(dim=2, sizes=(6, self.dim)) + + # context + assert context.shape[1] == self.text_len + context = self.text_embedding(context) + + if clip_feature is not None: + clip_embedding = self.img_emb(clip_feature) + context = torch.cat([clip_embedding, context], dim=1) + + if clean_x is not None: + if y is not None and self.concat_first_frame_latent: + clean_x = torch.cat([clean_x, y.to(dtype=clean_x.dtype)], dim=1) + clean_x = self.patch_embedding(clean_x) + clean_x = clean_x.flatten(start_dim=2).transpose(1, 2) + assert clean_x.shape[1] == seq_len + + x = torch.cat([clean_x, x], dim=1) + + if aug_t is None: + aug_t = torch.zeros_like(timestep_original) + assert aug_t is not None + + e_clean = self.time_embedding( + sinusoidal_embedding_1d(self.freq_dim, aug_t.flatten()).type_as(x)) + e_clean = e_clean.unflatten(dim=0, sizes=timestep_original.shape) + e0_clean = self.time_projection(e_clean) + e0_clean = e0_clean.unflatten(dim=2, sizes=(6, self.dim)) + e0 = torch.cat([e0_clean, e0], dim=1) + + # arguments + context_casted = context.to(dtype=x.dtype) if context is not None else context + kwargs = dict( + e=e0, + freqs=freqs, + freqs_action=self.freqs_action, + freqs_state=self.freqs_state, + action_register_length=action_register_length, + context=context_casted, + is_tf=clean_x is not None, + ) + + def create_custom_forward(module): + def custom_forward(*inputs, **kwargs): + outputs, updated_kv_cache = module(*inputs, **kwargs) + assert updated_kv_cache is None + return outputs + return custom_forward + + for block in self.blocks: + if torch.is_grad_enabled() and self.gradient_checkpointing: + ckpt_use_reentrant = getattr( + self, "gradient_checkpointing_use_reentrant", False + ) + if ckpt_use_reentrant: + # use_reentrant=True requires positional args only (no kwargs). + # Faster than non-reentrant (no tensor pack/unpack hooks). + x, _ = torch.utils.checkpoint.checkpoint( + block, + x, + e0, + freqs, + self.freqs_action, + self.freqs_state, + action_register_length, + context_casted, + None, # kv_cache + None, # crossattn_cache + 0, # current_start_frame + clean_x is not None, # is_tf + use_reentrant=True, + ) + else: + x = torch.utils.checkpoint.checkpoint( + create_custom_forward(block), + x, **kwargs, + use_reentrant=False, + ) + else: + x, _kv_cache = block(x, **kwargs) + + if clean_x is not None: + x = x[:, clean_x.shape[1]:] + + if action is not None: + action_noise_pred = x[:, seq_len: seq_len + action_length] + action_noise_pred = self.action_decoder(action_noise_pred, embodiment_id) + else: + action_noise_pred = None + + # Build a tensor that contains only video tokens per sample with length = max(video_lens) + x_video = x[:, :seq_len] + e_video = e[:, :seq_len] + + # Unpatchify video-only tokens + x_video = self.head(x_video, e_video.unsqueeze(2)) + video_noise_pred = self.unpatchify(x_video, grid_size) + + return video_noise_pred, action_noise_pred + + def forward( + self, + *args, + **kwargs + ): + if kwargs.get('kv_cache', None) is not None: + return self._forward_inference(*args, **kwargs) + else: + return self._forward_train(*args, **kwargs) + + def unpatchify(self, x, grid_size): + r""" + Reconstruct video tensors from patch embeddings. + + Args: + x (Tensor): + Patchified features, with shape [B, L, C_out * prod(patch_size)]. + grid_size (Tensor): + Spatial-temporal grid dimensions before patching, with shape [3] + (3 dimensions correspond to F_patches, H_patches, W_patches) + + Returns: + Tensor: + Reconstructed video tensors with shape [B, C_out, F, H / 8, W / 8] + """ + B = x.shape[0] + c = self.out_dim + grid_size = grid_size.tolist() + assert x.shape[1] == math.prod(grid_size) + x = x.view(B, *grid_size, *self.patch_size, c) + x = torch.einsum('bfhwpqrc->bcfphqwr', x) + x = x.reshape(B, c, *[i * j for i, j in zip(grid_size, self.patch_size)]) + return x + + def _create_freqs( + self, + grid_size: torch.Tensor, + start_frame: int, + ): + device = self.patch_embedding.weight.device + if any(freq.device != device for freq in self.freqs): + self.freqs = [freq.to(device) for freq in self.freqs] + if self.freqs_action.device != device: + self.freqs_action = self.freqs_action.to(device) + if self.freqs_state.device != device: + self.freqs_state = self.freqs_state.to(device) + + f, h, w = grid_size.tolist() + freqs = torch.cat( + [ + self.freqs[0][start_frame:start_frame + f].view(f, 1, 1, -1).expand(f, h, w, -1), + self.freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1), + self.freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1), + ], + dim=-1 + ).reshape(f * h * w, 1, -1) + + return freqs + + def init_weights(self): + r""" + Initialize model parameters using Xavier initialization. + """ + + # basic init + for m in self.modules(): + if isinstance(m, nn.Linear): + nn.init.xavier_uniform_(m.weight) + if m.bias is not None: + nn.init.zeros_(m.bias) + + # init embeddings + nn.init.xavier_uniform_(self.patch_embedding.weight.flatten(1)) + for m in self.text_embedding.modules(): + if isinstance(m, nn.Linear): + nn.init.normal_(m.weight, std=.02) + for m in self.time_embedding.modules(): + if isinstance(m, nn.Linear): + nn.init.normal_(m.weight, std=.02) + + # init output layer + nn.init.zeros_(self.head.head.weight) diff --git a/flagscale/train/models/dreamzero/modules/wan_video_image_encoder.py b/flagscale/train/models/dreamzero/modules/wan_video_image_encoder.py new file mode 100644 index 0000000000..c5e47788c3 --- /dev/null +++ b/flagscale/train/models/dreamzero/modules/wan_video_image_encoder.py @@ -0,0 +1,908 @@ +""" +Concise re-implementation of +``https://github.com/openai/CLIP'' and +``https://github.com/mlfoundations/open_clip''. +""" +import math +import torch +import torch.nn as nn +import torch.nn.functional as F +import torchvision.transforms as T +from flagscale.train.models.dreamzero.modules.wan_video_dit import flash_attention + + +class SelfAttention(nn.Module): + + def __init__(self, dim, num_heads, dropout=0.1, eps=1e-5): + assert dim % num_heads == 0 + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.eps = eps + + # layers + self.q = nn.Linear(dim, dim) + self.k = nn.Linear(dim, dim) + self.v = nn.Linear(dim, dim) + self.o = nn.Linear(dim, dim) + self.dropout = nn.Dropout(dropout) + + def forward(self, x, mask): + """ + x: [B, L, C]. + """ + b, s, c, n, d = *x.size(), self.num_heads, self.head_dim + + # compute query, key, value + q = self.q(x).reshape(b, s, n, d).permute(0, 2, 1, 3) + k = self.k(x).reshape(b, s, n, d).permute(0, 2, 1, 3) + v = self.v(x).reshape(b, s, n, d).permute(0, 2, 1, 3) + + # compute attention + p = self.dropout.p if self.training else 0.0 + x = F.scaled_dot_product_attention(q, k, v, mask, p) + x = x.permute(0, 2, 1, 3).reshape(b, s, c) + + # output + x = self.o(x) + x = self.dropout(x) + return x + + +class AttentionBlock(nn.Module): + + def __init__(self, dim, num_heads, post_norm, dropout=0.1, eps=1e-5): + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.post_norm = post_norm + self.eps = eps + + # layers + self.attn = SelfAttention(dim, num_heads, dropout, eps) + self.norm1 = nn.LayerNorm(dim, eps=eps) + self.ffn = nn.Sequential( + nn.Linear(dim, dim * 4), nn.GELU(), nn.Linear(dim * 4, dim), + nn.Dropout(dropout)) + self.norm2 = nn.LayerNorm(dim, eps=eps) + + def forward(self, x, mask): + if self.post_norm: + x = self.norm1(x + self.attn(x, mask)) + x = self.norm2(x + self.ffn(x)) + else: + x = x + self.attn(self.norm1(x), mask) + x = x + self.ffn(self.norm2(x)) + return x + + +class XLMRoberta(nn.Module): + """ + XLMRobertaModel with no pooler and no LM head. + """ + + def __init__(self, + vocab_size=250002, + max_seq_len=514, + type_size=1, + pad_id=1, + dim=1024, + num_heads=16, + num_layers=24, + post_norm=True, + dropout=0.1, + eps=1e-5): + super().__init__() + self.vocab_size = vocab_size + self.max_seq_len = max_seq_len + self.type_size = type_size + self.pad_id = pad_id + self.dim = dim + self.num_heads = num_heads + self.num_layers = num_layers + self.post_norm = post_norm + self.eps = eps + + # embeddings + self.token_embedding = nn.Embedding(vocab_size, dim, padding_idx=pad_id) + self.type_embedding = nn.Embedding(type_size, dim) + self.pos_embedding = nn.Embedding(max_seq_len, dim, padding_idx=pad_id) + self.dropout = nn.Dropout(dropout) + + # blocks + self.blocks = nn.ModuleList([ + AttentionBlock(dim, num_heads, post_norm, dropout, eps) + for _ in range(num_layers) + ]) + + # norm layer + self.norm = nn.LayerNorm(dim, eps=eps) + + def forward(self, ids): + """ + ids: [B, L] of torch.LongTensor. + """ + b, s = ids.shape + mask = ids.ne(self.pad_id).long() + + # embeddings + x = self.token_embedding(ids) + \ + self.type_embedding(torch.zeros_like(ids)) + \ + self.pos_embedding(self.pad_id + torch.cumsum(mask, dim=1) * mask) + if self.post_norm: + x = self.norm(x) + x = self.dropout(x) + + # blocks + mask = torch.where( + mask.view(b, 1, 1, s).gt(0), 0.0, + torch.finfo(x.dtype).min) + for block in self.blocks: + x = block(x, mask) + + # output + if not self.post_norm: + x = self.norm(x) + return x + + +def xlm_roberta_large(pretrained=False, + return_tokenizer=False, + device='cpu', + **kwargs): + """ + XLMRobertaLarge adapted from Huggingface. + """ + # params + cfg = dict( + vocab_size=250002, + max_seq_len=514, + type_size=1, + pad_id=1, + dim=1024, + num_heads=16, + num_layers=24, + post_norm=True, + dropout=0.1, + eps=1e-5) + cfg.update(**kwargs) + + # init model + if pretrained: + from sora import DOWNLOAD_TO_CACHE + + # init a meta model + with torch.device('meta'): + model = XLMRoberta(**cfg) + + # load checkpoint + model.load_state_dict( + torch.load( + DOWNLOAD_TO_CACHE('models/xlm_roberta/xlm_roberta_large.pth'), + map_location=device), + assign=True) + else: + # init a model on device + with torch.device(device): + model = XLMRoberta(**cfg) + + # init tokenizer + if return_tokenizer: + from sora.data import HuggingfaceTokenizer + tokenizer = HuggingfaceTokenizer( + name='xlm-roberta-large', + seq_len=model.text_len, + clean='whitespace') + return model, tokenizer + else: + return model + + + +def pos_interpolate(pos, seq_len): + if pos.size(1) == seq_len: + return pos + else: + src_grid = int(math.sqrt(pos.size(1))) + tar_grid = int(math.sqrt(seq_len)) + n = pos.size(1) - src_grid * src_grid + return torch.cat([ + pos[:, :n], + F.interpolate( + pos[:, n:].float().reshape(1, src_grid, src_grid, -1).permute( + 0, 3, 1, 2), + size=(tar_grid, tar_grid), + mode='bicubic', + align_corners=False).flatten(2).transpose(1, 2) + ], + dim=1) + + +class QuickGELU(nn.Module): + + def forward(self, x): + return x * torch.sigmoid(1.702 * x) + + +class LayerNorm(nn.LayerNorm): + + def forward(self, x): + return super().forward(x).type_as(x) + + +class SelfAttention(nn.Module): + + def __init__(self, + dim, + num_heads, + causal=False, + attn_dropout=0.0, + proj_dropout=0.0): + assert dim % num_heads == 0 + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.causal = causal + self.attn_dropout = attn_dropout + self.proj_dropout = proj_dropout + + # layers + self.to_qkv = nn.Linear(dim, dim * 3) + self.proj = nn.Linear(dim, dim) + + def forward(self, x): + """ + x: [B, L, C]. + """ + # compute query, key, value + q, k, v = self.to_qkv(x).chunk(3, dim=-1) + + # compute attention + x = flash_attention(q, k, v, num_heads=self.num_heads, compatibility_mode=True) + + # output + x = self.proj(x) + x = F.dropout(x, self.proj_dropout, self.training) + return x + + +class SwiGLU(nn.Module): + + def __init__(self, dim, mid_dim): + super().__init__() + self.dim = dim + self.mid_dim = mid_dim + + # layers + self.fc1 = nn.Linear(dim, mid_dim) + self.fc2 = nn.Linear(dim, mid_dim) + self.fc3 = nn.Linear(mid_dim, dim) + + def forward(self, x): + x = F.silu(self.fc1(x)) * self.fc2(x) + x = self.fc3(x) + return x + + +class AttentionBlock(nn.Module): + + def __init__(self, + dim, + mlp_ratio, + num_heads, + post_norm=False, + causal=False, + activation='quick_gelu', + attn_dropout=0.0, + proj_dropout=0.0, + norm_eps=1e-5): + assert activation in ['quick_gelu', 'gelu', 'swi_glu'] + super().__init__() + self.dim = dim + self.mlp_ratio = mlp_ratio + self.num_heads = num_heads + self.post_norm = post_norm + self.causal = causal + self.norm_eps = norm_eps + + # layers + self.norm1 = LayerNorm(dim, eps=norm_eps) + self.attn = SelfAttention(dim, num_heads, causal, attn_dropout, + proj_dropout) + self.norm2 = LayerNorm(dim, eps=norm_eps) + if activation == 'swi_glu': + self.mlp = SwiGLU(dim, int(dim * mlp_ratio)) + else: + self.mlp = nn.Sequential( + nn.Linear(dim, int(dim * mlp_ratio)), + QuickGELU() if activation == 'quick_gelu' else nn.GELU(), + nn.Linear(int(dim * mlp_ratio), dim), nn.Dropout(proj_dropout)) + + def forward(self, x): + if self.post_norm: + x = x + self.norm1(self.attn(x)) + x = x + self.norm2(self.mlp(x)) + else: + x = x + self.attn(self.norm1(x)) + x = x + self.mlp(self.norm2(x)) + return x + + +class AttentionPool(nn.Module): + + def __init__(self, + dim, + mlp_ratio, + num_heads, + activation='gelu', + proj_dropout=0.0, + norm_eps=1e-5): + assert dim % num_heads == 0 + super().__init__() + self.dim = dim + self.mlp_ratio = mlp_ratio + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.proj_dropout = proj_dropout + self.norm_eps = norm_eps + + # layers + gain = 1.0 / math.sqrt(dim) + self.cls_embedding = nn.Parameter(gain * torch.randn(1, 1, dim)) + self.to_q = nn.Linear(dim, dim) + self.to_kv = nn.Linear(dim, dim * 2) + self.proj = nn.Linear(dim, dim) + self.norm = LayerNorm(dim, eps=norm_eps) + self.mlp = nn.Sequential( + nn.Linear(dim, int(dim * mlp_ratio)), + QuickGELU() if activation == 'quick_gelu' else nn.GELU(), + nn.Linear(int(dim * mlp_ratio), dim), nn.Dropout(proj_dropout)) + + def forward(self, x): + """ + x: [B, L, C]. + """ + b, s, c, n, d = *x.size(), self.num_heads, self.head_dim + + # compute query, key, value + q = self.to_q(self.cls_embedding).view(1, 1, n*d).expand(b, -1, -1) + k, v = self.to_kv(x).chunk(2, dim=-1) + + # compute attention + x = flash_attention(q, k, v, num_heads=self.num_heads, compatibility_mode=True) + x = x.reshape(b, 1, c) + + # output + x = self.proj(x) + x = F.dropout(x, self.proj_dropout, self.training) + + # mlp + x = x + self.mlp(self.norm(x)) + return x[:, 0] + + +class VisionTransformer(nn.Module): + + def __init__(self, + image_size=224, + patch_size=16, + dim=768, + mlp_ratio=4, + out_dim=512, + num_heads=12, + num_layers=12, + pool_type='token', + pre_norm=True, + post_norm=False, + activation='quick_gelu', + attn_dropout=0.0, + proj_dropout=0.0, + embedding_dropout=0.0, + norm_eps=1e-5): + if image_size % patch_size != 0: + print( + '[WARNING] image_size is not divisible by patch_size', + flush=True) + assert pool_type in ('token', 'token_fc', 'attn_pool') + out_dim = out_dim or dim + super().__init__() + self.image_size = image_size + self.patch_size = patch_size + self.num_patches = (image_size // patch_size)**2 + self.dim = dim + self.mlp_ratio = mlp_ratio + self.out_dim = out_dim + self.num_heads = num_heads + self.num_layers = num_layers + self.pool_type = pool_type + self.post_norm = post_norm + self.norm_eps = norm_eps + + # embeddings + gain = 1.0 / math.sqrt(dim) + self.patch_embedding = nn.Conv2d( + 3, + dim, + kernel_size=patch_size, + stride=patch_size, + bias=not pre_norm) + if pool_type in ('token', 'token_fc'): + self.cls_embedding = nn.Parameter(gain * torch.randn(1, 1, dim)) + self.pos_embedding = nn.Parameter(gain * torch.randn( + 1, self.num_patches + + (1 if pool_type in ('token', 'token_fc') else 0), dim)) + self.dropout = nn.Dropout(embedding_dropout) + + # transformer + self.pre_norm = LayerNorm(dim, eps=norm_eps) if pre_norm else None + self.transformer = nn.Sequential(*[ + AttentionBlock(dim, mlp_ratio, num_heads, post_norm, False, + activation, attn_dropout, proj_dropout, norm_eps) + for _ in range(num_layers) + ]) + self.post_norm = LayerNorm(dim, eps=norm_eps) + + # head + if pool_type == 'token': + self.head = nn.Parameter(gain * torch.randn(dim, out_dim)) + elif pool_type == 'token_fc': + self.head = nn.Linear(dim, out_dim) + elif pool_type == 'attn_pool': + self.head = AttentionPool(dim, mlp_ratio, num_heads, activation, + proj_dropout, norm_eps) + + def forward(self, x, interpolation=False, use_31_block=False): + b = x.size(0) + + # embeddings + x = self.patch_embedding(x).flatten(2).permute(0, 2, 1) + # print("x in forward: ", x[0,0,100:105], x.shape) + # print("patch_embedding: ", self.patch_embedding.module.weight[0:10, 0, 0, 3], self.patch_embedding.module.weight.shape) + if self.pool_type in ('token', 'token_fc'): + x = torch.cat([self.cls_embedding.expand(b, -1, -1).to(dtype=x.dtype, device=x.device), x], dim=1) + if interpolation: + e = pos_interpolate(self.pos_embedding, x.size(1)) + else: + e = self.pos_embedding + e = e.to(dtype=x.dtype, device=x.device) + x = self.dropout(x + e) + if self.pre_norm is not None: + x = self.pre_norm(x) + + # transformer + if use_31_block: + # print("x before transformer: ", x[0,0,100:105], x.shape) + x = self.transformer[:-1](x) + return x + else: + # print("x before transformer: ", x[0,0,100:105], x.shape) + x = self.transformer(x) + return x + + +class CLIP(nn.Module): + + def __init__(self, + embed_dim=512, + image_size=224, + patch_size=16, + vision_dim=768, + vision_mlp_ratio=4, + vision_heads=12, + vision_layers=12, + vision_pool='token', + vision_pre_norm=True, + vision_post_norm=False, + vocab_size=49408, + text_len=77, + text_dim=512, + text_mlp_ratio=4, + text_heads=8, + text_layers=12, + text_causal=True, + text_pool='argmax', + text_head_bias=False, + logit_bias=None, + activation='quick_gelu', + attn_dropout=0.0, + proj_dropout=0.0, + embedding_dropout=0.0, + norm_eps=1e-5): + super().__init__() + self.embed_dim = embed_dim + self.image_size = image_size + self.patch_size = patch_size + self.vision_dim = vision_dim + self.vision_mlp_ratio = vision_mlp_ratio + self.vision_heads = vision_heads + self.vision_layers = vision_layers + self.vision_pool = vision_pool + self.vision_pre_norm = vision_pre_norm + self.vision_post_norm = vision_post_norm + self.vocab_size = vocab_size + self.text_len = text_len + self.text_dim = text_dim + self.text_mlp_ratio = text_mlp_ratio + self.text_heads = text_heads + self.text_layers = text_layers + self.text_causal = text_causal + self.text_pool = text_pool + self.text_head_bias = text_head_bias + self.norm_eps = norm_eps + + # models + self.visual = VisionTransformer( + image_size=image_size, + patch_size=patch_size, + dim=vision_dim, + mlp_ratio=vision_mlp_ratio, + out_dim=embed_dim, + num_heads=vision_heads, + num_layers=vision_layers, + pool_type=vision_pool, + pre_norm=vision_pre_norm, + post_norm=vision_post_norm, + activation=activation, + attn_dropout=attn_dropout, + proj_dropout=proj_dropout, + embedding_dropout=embedding_dropout, + norm_eps=norm_eps) + self.textual = TextTransformer( + vocab_size=vocab_size, + text_len=text_len, + dim=text_dim, + mlp_ratio=text_mlp_ratio, + out_dim=embed_dim, + num_heads=text_heads, + num_layers=text_layers, + causal=text_causal, + pool_type=text_pool, + head_bias=text_head_bias, + activation=activation, + attn_dropout=attn_dropout, + proj_dropout=proj_dropout, + embedding_dropout=embedding_dropout, + norm_eps=norm_eps) + self.log_scale = nn.Parameter(math.log(1 / 0.07) * torch.ones([])) + if logit_bias is not None: + self.logit_bias = nn.Parameter(logit_bias * torch.ones([])) + + # initialize weights + self.init_weights() + + def forward(self, imgs, txt_ids): + """ + imgs: [B, 3, H, W] of torch.float32. + - mean: [0.48145466, 0.4578275, 0.40821073] + - std: [0.26862954, 0.26130258, 0.27577711] + txt_ids: [B, L] of torch.long. Encoded by data.CLIPTokenizer. + """ + xi = self.visual(imgs) + xt = self.textual(txt_ids) + return xi, xt + + def init_weights(self): + # embeddings + nn.init.normal_(self.textual.token_embedding.weight, std=0.02) + nn.init.normal_(self.visual.patch_embedding.weight, std=0.1) + + # attentions + for modality in ['visual', 'textual']: + dim = self.vision_dim if modality == 'visual' else self.text_dim + transformer = getattr(self, modality).transformer + proj_gain = (1.0 / math.sqrt(dim)) * ( + 1.0 / math.sqrt(2 * len(transformer))) + attn_gain = 1.0 / math.sqrt(dim) + mlp_gain = 1.0 / math.sqrt(2.0 * dim) + for block in transformer: + nn.init.normal_(block.attn.to_qkv.weight, std=attn_gain) + nn.init.normal_(block.attn.proj.weight, std=proj_gain) + nn.init.normal_(block.mlp[0].weight, std=mlp_gain) + nn.init.normal_(block.mlp[2].weight, std=proj_gain) + + def param_groups(self): + groups = [{ + 'params': [ + p for n, p in self.named_parameters() + if 'norm' in n or n.endswith('bias') + ], + 'weight_decay': 0.0 + }, { + 'params': [ + p for n, p in self.named_parameters() + if not ('norm' in n or n.endswith('bias')) + ] + }] + return groups + + +class XLMRobertaWithHead(XLMRoberta): + + def __init__(self, **kwargs): + self.out_dim = kwargs.pop('out_dim') + super().__init__(**kwargs) + + # head + mid_dim = (self.dim + self.out_dim) // 2 + self.head = nn.Sequential( + nn.Linear(self.dim, mid_dim, bias=False), nn.GELU(), + nn.Linear(mid_dim, self.out_dim, bias=False)) + + def forward(self, ids): + # xlm-roberta + x = super().forward(ids) + + # average pooling + mask = ids.ne(self.pad_id).unsqueeze(-1).to(x) + x = (x * mask).sum(dim=1) / mask.sum(dim=1) + + # head + x = self.head(x) + return x + + +class XLMRobertaCLIP(nn.Module): + + def __init__(self, + embed_dim=1024, + image_size=224, + patch_size=14, + vision_dim=1280, + vision_mlp_ratio=4, + vision_heads=16, + vision_layers=32, + vision_pool='token', + vision_pre_norm=True, + vision_post_norm=False, + activation='gelu', + vocab_size=250002, + max_text_len=514, + type_size=1, + pad_id=1, + text_dim=1024, + text_heads=16, + text_layers=24, + text_post_norm=True, + text_dropout=0.1, + attn_dropout=0.0, + proj_dropout=0.0, + embedding_dropout=0.0, + norm_eps=1e-5): + super().__init__() + self.embed_dim = embed_dim + self.image_size = image_size + self.patch_size = patch_size + self.vision_dim = vision_dim + self.vision_mlp_ratio = vision_mlp_ratio + self.vision_heads = vision_heads + self.vision_layers = vision_layers + self.vision_pre_norm = vision_pre_norm + self.vision_post_norm = vision_post_norm + self.activation = activation + self.vocab_size = vocab_size + self.max_text_len = max_text_len + self.type_size = type_size + self.pad_id = pad_id + self.text_dim = text_dim + self.text_heads = text_heads + self.text_layers = text_layers + self.text_post_norm = text_post_norm + self.norm_eps = norm_eps + + # models + self.visual = VisionTransformer( + image_size=image_size, + patch_size=patch_size, + dim=vision_dim, + mlp_ratio=vision_mlp_ratio, + out_dim=embed_dim, + num_heads=vision_heads, + num_layers=vision_layers, + pool_type=vision_pool, + pre_norm=vision_pre_norm, + post_norm=vision_post_norm, + activation=activation, + attn_dropout=attn_dropout, + proj_dropout=proj_dropout, + embedding_dropout=embedding_dropout, + norm_eps=norm_eps) + self.textual = None + self.log_scale = nn.Parameter(math.log(1 / 0.07) * torch.ones([])) + + def forward(self, imgs, txt_ids): + """ + imgs: [B, 3, H, W] of torch.float32. + - mean: [0.48145466, 0.4578275, 0.40821073] + - std: [0.26862954, 0.26130258, 0.27577711] + txt_ids: [B, L] of torch.long. + Encoded by data.CLIPTokenizer. + """ + xi = self.visual(imgs) + xt = self.textual(txt_ids) + return xi, xt + + def param_groups(self): + groups = [{ + 'params': [ + p for n, p in self.named_parameters() + if 'norm' in n or n.endswith('bias') + ], + 'weight_decay': 0.0 + }, { + 'params': [ + p for n, p in self.named_parameters() + if not ('norm' in n or n.endswith('bias')) + ] + }] + return groups + + +def _clip(pretrained=False, + pretrained_name=None, + model_cls=CLIP, + return_transforms=False, + return_tokenizer=False, + tokenizer_padding='eos', + dtype=torch.float32, + device='cpu', + **kwargs): + # init model + if pretrained and pretrained_name: + from sora import BUCKET, DOWNLOAD_TO_CACHE + + # init a meta model + with torch.device('meta'): + model = model_cls(**kwargs) + + # checkpoint path + checkpoint = f'models/clip/{pretrained_name}' + if dtype in (torch.float16, torch.bfloat16): + suffix = '-' + { + torch.float16: 'fp16', + torch.bfloat16: 'bf16' + }[dtype] + if object_exists(BUCKET, f'{checkpoint}{suffix}.pth'): + checkpoint = f'{checkpoint}{suffix}' + checkpoint += '.pth' + + # load + model.load_state_dict( + torch.load(DOWNLOAD_TO_CACHE(checkpoint), map_location=device), + assign=True, + strict=False) + else: + # init a model on device + with torch.device(device): + model = model_cls(**kwargs) + + # set device + output = (model,) + + # init transforms + if return_transforms: + # mean and std + if 'siglip' in pretrained_name.lower(): + mean, std = [0.5, 0.5, 0.5], [0.5, 0.5, 0.5] + else: + mean = [0.48145466, 0.4578275, 0.40821073] + std = [0.26862954, 0.26130258, 0.27577711] + + # transforms + transforms = T.Compose([ + T.Resize((model.image_size, model.image_size), + interpolation=T.InterpolationMode.BICUBIC), + T.ToTensor(), + T.Normalize(mean=mean, std=std) + ]) + output += (transforms,) + + # init tokenizer + if return_tokenizer: + from sora import data + if 'siglip' in pretrained_name.lower(): + tokenizer = data.HuggingfaceTokenizer( + name=f'timm/{pretrained_name}', + seq_len=model.text_len, + clean='canonicalize') + elif 'xlm' in pretrained_name.lower(): + tokenizer = data.HuggingfaceTokenizer( + name='xlm-roberta-large', + seq_len=model.max_text_len - 2, + clean='whitespace') + elif 'mba' in pretrained_name.lower(): + tokenizer = data.HuggingfaceTokenizer( + name='facebook/xlm-roberta-xl', + seq_len=model.max_text_len - 2, + clean='whitespace') + else: + tokenizer = data.CLIPTokenizer( + seq_len=model.text_len, padding=tokenizer_padding) + output += (tokenizer,) + return output[0] if len(output) == 1 else output + + +def clip_xlm_roberta_vit_h_14( + pretrained=False, + pretrained_name='open-clip-xlm-roberta-large-vit-huge-14', + **kwargs): + cfg = dict( + embed_dim=1024, + image_size=224, + patch_size=14, + vision_dim=1280, + vision_mlp_ratio=4, + vision_heads=16, + vision_layers=32, + vision_pool='token', + activation='gelu', + vocab_size=250002, + max_text_len=514, + type_size=1, + pad_id=1, + text_dim=1024, + text_heads=16, + text_layers=24, + text_post_norm=True, + text_dropout=0.1, + attn_dropout=0.0, + proj_dropout=0.0, + embedding_dropout=0.0) + cfg.update(**kwargs) + return _clip(pretrained, pretrained_name, XLMRobertaCLIP, **cfg) + + +class WanImageEncoder(torch.nn.Module): + + def __init__(self, image_encoder_pretrained_path: str=None): + super().__init__() + # init model + self.model, self.transforms = clip_xlm_roberta_vit_h_14( + pretrained=False, + return_transforms=True, + return_tokenizer=False, + dtype=torch.float32, + device="cpu") + self.image_encoder_pretrained_path = image_encoder_pretrained_path + + def encode_image(self, videos): + # preprocess + size = (self.model.image_size,) * 2 + videos = torch.cat([ + F.interpolate( + u, + size=size, + mode='bicubic', + align_corners=False) for u in videos + ]) + videos = self.transforms.transforms[-1](videos.mul_(0.5).add_(0.5)) + + # forward + dtype = next(iter(self.model.visual.parameters())).dtype + videos = videos.to(dtype) + out = self.model.visual(videos, use_31_block=True) + # The outputs of torch compile always need to be cloned before being used. + out = out.clone() + return out + + @staticmethod + def state_dict_converter(): + return WanImageEncoderStateDictConverter() + + +class WanImageEncoderStateDictConverter: + def __init__(self): + pass + + def from_diffusers(self, state_dict): + return state_dict + + def from_civitai(self, state_dict): + state_dict_ = {} + for name, param in state_dict.items(): + if name.startswith("textual."): + continue + name = "model." + name + state_dict_[name] = param + return state_dict_ diff --git a/flagscale/train/models/dreamzero/modules/wan_video_text_encoder.py b/flagscale/train/models/dreamzero/modules/wan_video_text_encoder.py new file mode 100644 index 0000000000..abeb018333 --- /dev/null +++ b/flagscale/train/models/dreamzero/modules/wan_video_text_encoder.py @@ -0,0 +1,280 @@ +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +def fp16_clamp(x): + if x.dtype == torch.float16 and torch.isinf(x).any(): + clamp = torch.finfo(x.dtype).max - 1000 + x = torch.clamp(x, min=-clamp, max=clamp) + return x + + +class GELU(nn.Module): + + def forward(self, x): + return 0.5 * x * (1.0 + torch.tanh( + math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0)))) + + +class T5LayerNorm(nn.Module): + + def __init__(self, dim, eps=1e-6): + super(T5LayerNorm, self).__init__() + self.dim = dim + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x): + x = x * torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + + self.eps) + if self.weight.dtype in [torch.float16, torch.bfloat16]: + x = x.type_as(self.weight) + return self.weight * x + + +class T5Attention(nn.Module): + + def __init__(self, dim, dim_attn, num_heads, dropout=0.1): + assert dim_attn % num_heads == 0 + super(T5Attention, self).__init__() + self.dim = dim + self.dim_attn = dim_attn + self.num_heads = num_heads + self.head_dim = dim_attn // num_heads + + # layers + self.q = nn.Linear(dim, dim_attn, bias=False) + self.k = nn.Linear(dim, dim_attn, bias=False) + self.v = nn.Linear(dim, dim_attn, bias=False) + self.o = nn.Linear(dim_attn, dim, bias=False) + self.dropout = nn.Dropout(dropout) + + def forward(self, x, context=None, mask=None, pos_bias=None): + """ + x: [B, L1, C]. + context: [B, L2, C] or None. + mask: [B, L2] or [B, L1, L2] or None. + """ + # check inputs + context = x if context is None else context + b, n, c = x.size(0), self.num_heads, self.head_dim + + # compute query, key, value + q = self.q(x).view(b, -1, n, c) + k = self.k(context).view(b, -1, n, c) + v = self.v(context).view(b, -1, n, c) + + # attention bias + attn_bias = x.new_zeros(b, n, q.size(1), k.size(1)) + if pos_bias is not None: + attn_bias += pos_bias + if mask is not None: + assert mask.ndim in [2, 3] + mask = mask.view(b, 1, 1, + -1) if mask.ndim == 2 else mask.unsqueeze(1) + attn_bias.masked_fill_(mask == 0, torch.finfo(x.dtype).min) + + # compute attention (T5 does not use scaling) + attn = torch.einsum('binc,bjnc->bnij', q, k) + attn_bias + attn = F.softmax(attn.float(), dim=-1).type_as(attn) + x = torch.einsum('bnij,bjnc->binc', attn, v) + + # output + x = x.reshape(b, -1, n * c) + x = self.o(x) + x = self.dropout(x) + return x + + +class T5FeedForward(nn.Module): + + def __init__(self, dim, dim_ffn, dropout=0.1): + super(T5FeedForward, self).__init__() + self.dim = dim + self.dim_ffn = dim_ffn + + # layers + self.gate = nn.Sequential(nn.Linear(dim, dim_ffn, bias=False), GELU()) + self.fc1 = nn.Linear(dim, dim_ffn, bias=False) + self.fc2 = nn.Linear(dim_ffn, dim, bias=False) + self.dropout = nn.Dropout(dropout) + + def forward(self, x): + x = self.fc1(x) * self.gate(x) + x = self.dropout(x) + x = self.fc2(x) + x = self.dropout(x) + return x + + +class T5SelfAttention(nn.Module): + + def __init__(self, + dim, + dim_attn, + dim_ffn, + num_heads, + num_buckets, + shared_pos=True, + dropout=0.1): + super(T5SelfAttention, self).__init__() + self.dim = dim + self.dim_attn = dim_attn + self.dim_ffn = dim_ffn + self.num_heads = num_heads + self.num_buckets = num_buckets + self.shared_pos = shared_pos + + # layers + self.norm1 = T5LayerNorm(dim) + self.attn = T5Attention(dim, dim_attn, num_heads, dropout) + self.norm2 = T5LayerNorm(dim) + self.ffn = T5FeedForward(dim, dim_ffn, dropout) + self.pos_embedding = None if shared_pos else T5RelativeEmbedding( + num_buckets, num_heads, bidirectional=True) + + def forward(self, x, mask=None, pos_bias=None): + e = pos_bias if self.shared_pos else self.pos_embedding( + x.size(1), x.size(1)) + x = fp16_clamp(x + self.attn(self.norm1(x), mask=mask, pos_bias=e)) + x = fp16_clamp(x + self.ffn(self.norm2(x))) + # print("x after attn: ", x[0, 0:10], x.shape) + return x + + +class T5RelativeEmbedding(nn.Module): + + def __init__(self, num_buckets, num_heads, bidirectional, max_dist=128): + super(T5RelativeEmbedding, self).__init__() + self.num_buckets = num_buckets + self.num_heads = num_heads + self.bidirectional = bidirectional + self.max_dist = max_dist + + # layers + self.embedding = nn.Embedding(num_buckets, num_heads) + + def forward(self, lq, lk): + device = self.embedding.weight.device + # rel_pos = torch.arange(lk).unsqueeze(0).to(device) - \ + # torch.arange(lq).unsqueeze(1).to(device) + rel_pos = torch.arange(lk, device=device).unsqueeze(0) - \ + torch.arange(lq, device=device).unsqueeze(1) + rel_pos = self._relative_position_bucket(rel_pos) + rel_pos_embeds = self.embedding(rel_pos) + rel_pos_embeds = rel_pos_embeds.permute(2, 0, 1).unsqueeze( + 0) # [1, N, Lq, Lk] + return rel_pos_embeds.contiguous() + + def _relative_position_bucket(self, rel_pos): + # preprocess + if self.bidirectional: + num_buckets = self.num_buckets // 2 + rel_buckets = (rel_pos > 0).long() * num_buckets + rel_pos = torch.abs(rel_pos) + else: + num_buckets = self.num_buckets + rel_buckets = 0 + rel_pos = -torch.min(rel_pos, torch.zeros_like(rel_pos)) + + # embeddings for small and large positions + max_exact = num_buckets // 2 + rel_pos_large = max_exact + (torch.log(rel_pos.float() / max_exact) / + math.log(self.max_dist / max_exact) * + (num_buckets - max_exact)).long() + rel_pos_large = torch.min( + rel_pos_large, torch.full_like(rel_pos_large, num_buckets - 1)) + rel_buckets += torch.where(rel_pos < max_exact, rel_pos, rel_pos_large) + return rel_buckets + +def init_weights(m): + if isinstance(m, T5LayerNorm): + nn.init.ones_(m.weight) + elif isinstance(m, T5FeedForward): + nn.init.normal_(m.gate[0].weight, std=m.dim**-0.5) + nn.init.normal_(m.fc1.weight, std=m.dim**-0.5) + nn.init.normal_(m.fc2.weight, std=m.dim_ffn**-0.5) + elif isinstance(m, T5Attention): + nn.init.normal_(m.q.weight, std=(m.dim * m.dim_attn)**-0.5) + nn.init.normal_(m.k.weight, std=m.dim**-0.5) + nn.init.normal_(m.v.weight, std=m.dim**-0.5) + nn.init.normal_(m.o.weight, std=(m.num_heads * m.dim_attn)**-0.5) + elif isinstance(m, T5RelativeEmbedding): + nn.init.normal_( + m.embedding.weight, std=(2 * m.num_buckets * m.num_heads)**-0.5) + + +class WanTextEncoder(torch.nn.Module): + + def __init__(self, + vocab: int | nn.Embedding = 256384, + dim=4096, + dim_attn=4096, + dim_ffn=10240, + num_heads=64, + num_layers=24, + num_buckets=32, + shared_pos=False, + dropout=0.1, + text_encoder_pretrained_path: str=None): + super(WanTextEncoder, self).__init__() + self.dim = dim + self.dim_attn = dim_attn + self.dim_ffn = dim_ffn + self.num_heads = num_heads + self.num_layers = num_layers + self.num_buckets = num_buckets + self.shared_pos = shared_pos + self.text_encoder_pretrained_path = text_encoder_pretrained_path + + # layers + if isinstance(vocab, int): + self.token_embedding = nn.Embedding(vocab, dim) + else: + self.token_embedding = vocab + if shared_pos: + self.pos_embedding = T5RelativeEmbedding( + num_buckets, num_heads, bidirectional=True) + else: + self.pos_embedding = None + self.dropout = nn.Dropout(dropout) + self.blocks = nn.ModuleList([ + T5SelfAttention(dim, dim_attn, dim_ffn, num_heads, num_buckets, + shared_pos, dropout) for _ in range(num_layers) + ]) + self.norm = T5LayerNorm(dim) + + # initialize weights + self.apply(init_weights) + + def forward(self, ids, mask=None): + x = self.token_embedding(ids) + x = self.dropout(x) + if self.shared_pos: + assert self.pos_embedding is not None + e = self.pos_embedding(x.size(1), x.size(1)) + else: + e = None + for block in self.blocks: + x = block(x, mask, pos_bias=e) + x = self.norm(x) + x = self.dropout(x) + return x + + @staticmethod + def state_dict_converter(): + return WanTextEncoderStateDictConverter() + + +class WanTextEncoderStateDictConverter: + def __init__(self): + pass + + def from_diffusers(self, state_dict): + return state_dict + + def from_civitai(self, state_dict): + return state_dict \ No newline at end of file diff --git a/flagscale/train/models/dreamzero/modules/wan_video_vae.py b/flagscale/train/models/dreamzero/modules/wan_video_vae.py new file mode 100644 index 0000000000..4e2ff44f0d --- /dev/null +++ b/flagscale/train/models/dreamzero/modules/wan_video_vae.py @@ -0,0 +1,1373 @@ +from einops import rearrange, repeat + +import torch +import torch.nn as nn +import torch.nn.functional as F +from tqdm import tqdm + +CACHE_T = 2 + + +def check_is_instance(model, module_class): + if isinstance(model, module_class): + return True + if hasattr(model, "module") and isinstance(model.module, module_class): + return True + return False + + +def block_causal_mask(x, block_size): + # params + b, n, s, _ = x.shape + assert s % block_size == 0 + num_blocks = s // block_size + + # build mask + mask = torch.zeros(b, n, s, s, dtype=torch.bool, device=x.device) + for i in range(num_blocks): + mask[:, :, + i * block_size:(i + 1) * block_size, :(i + 1) * block_size] = 1 + return mask + + +class CausalConv3d(nn.Conv3d): + """ + Causal 3d convolusion. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._padding = (self.padding[2], self.padding[2], self.padding[1], + self.padding[1], 2 * self.padding[0], 0) + self.padding = (0, 0, 0) + + def forward(self, x, cache_x=None): + padding = list(self._padding) + if cache_x is not None and self._padding[4] > 0: + x = torch.cat([cache_x, x], dim=2) + padding[4] -= cache_x.shape[2] + x = F.pad(x, padding) + + return super().forward(x) + + +class RMS_norm(nn.Module): + + def __init__(self, dim, channel_first=True, images=True, bias=False): + super().__init__() + broadcastable_dims = (1, 1, 1) if not images else (1, 1) + shape = (dim, *broadcastable_dims) if channel_first else (dim,) + + self.channel_first = channel_first + self.scale = dim**0.5 + self.gamma = nn.Parameter(torch.ones(shape)) + self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0. + + def forward(self, x): + return F.normalize( + x, dim=(1 if self.channel_first else + -1)) * self.scale * self.gamma + self.bias + + +class Upsample(nn.Upsample): + + def forward(self, x): + """ + Fix bfloat16 support for nearest neighbor interpolation. + """ + return super().forward(x.float()).type_as(x) + + +class Resample(nn.Module): + + def __init__(self, dim, mode): + assert mode in ('none', 'upsample2d', 'upsample3d', 'downsample2d', + 'downsample3d') + super().__init__() + self.dim = dim + self.mode = mode + + # layers + if mode == 'upsample2d': + self.resample = nn.Sequential( + Upsample(scale_factor=(2., 2.), mode='nearest-exact'), + nn.Conv2d(dim, dim // 2, 3, padding=1)) + elif mode == 'upsample3d': + self.resample = nn.Sequential( + Upsample(scale_factor=(2., 2.), mode='nearest-exact'), + nn.Conv2d(dim, dim // 2, 3, padding=1)) + self.time_conv = CausalConv3d(dim, + dim * 2, (3, 1, 1), + padding=(1, 0, 0)) + + elif mode == 'downsample2d': + self.resample = nn.Sequential( + nn.ZeroPad2d((0, 1, 0, 1)), + nn.Conv2d(dim, dim, 3, stride=(2, 2))) + elif mode == 'downsample3d': + self.resample = nn.Sequential( + nn.ZeroPad2d((0, 1, 0, 1)), + nn.Conv2d(dim, dim, 3, stride=(2, 2))) + self.time_conv = CausalConv3d(dim, + dim, (3, 1, 1), + stride=(2, 1, 1), + padding=(0, 0, 0)) + + else: + self.resample = nn.Identity() + + def forward(self, x, feat_cache=None, feat_idx=[0]): + b, c, t, h, w = x.size() + if self.mode == 'upsample3d': + if feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = 'Rep' + feat_idx[0] += 1 + else: + + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[ + idx] is not None and feat_cache[idx] != 'Rep': + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2), + cache_x, + ], dim=2) + if cache_x.shape[2] < 2 and feat_cache[ + idx] is not None and feat_cache[idx] == 'Rep': + cache_x = torch.cat([ + torch.zeros_like(cache_x), + cache_x, + ], dim=2) + if feat_cache[idx] == 'Rep': + x = self.time_conv(x) + else: + x = self.time_conv(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + + x = x.reshape(b, 2, c, t, h, w) + x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), + 3) + x = x.reshape(b, c, t * 2, h, w) + t = x.shape[2] + x = rearrange(x, 'b c t h w -> (b t) c h w') + x = self.resample(x) + x = rearrange(x, '(b t) c h w -> b c t h w', t=t) + + if self.mode == 'downsample3d': + if feat_cache is not None: + idx = feat_idx[0] + if feat_cache[idx] is None: + feat_cache[idx] = x.clone() + feat_idx[0] += 1 + else: + cache_x = x[:, :, -1:, :, :].clone() + x = self.time_conv( + torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2)) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + return x + + def init_weight(self, conv): + conv_weight = conv.weight + nn.init.zeros_(conv_weight) + c1, c2, t, h, w = conv_weight.size() + one_matrix = torch.eye(c1, c2) + init_matrix = one_matrix + nn.init.zeros_(conv_weight) + conv_weight.data[:, :, 1, 0, 0] = init_matrix + conv.weight.data.copy_(conv_weight) + nn.init.zeros_(conv.bias.data) + + def init_weight2(self, conv): + conv_weight = conv.weight.data + nn.init.zeros_(conv_weight) + c1, c2, t, h, w = conv_weight.size() + init_matrix = torch.eye(c1 // 2, c2) + conv_weight[:c1 // 2, :, -1, 0, 0] = init_matrix + conv_weight[c1 // 2:, :, -1, 0, 0] = init_matrix + conv.weight.data.copy_(conv_weight) + nn.init.zeros_(conv.bias.data) + + + +def patchify(x, patch_size): + if patch_size == 1: + return x + if x.dim() == 4: + x = rearrange(x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size, r=patch_size) + elif x.dim() == 5: + x = rearrange(x, + "b c f (h q) (w r) -> b (c r q) f h w", + q=patch_size, + r=patch_size) + else: + raise ValueError(f"Invalid input shape: {x.shape}") + return x + + +def unpatchify(x, patch_size): + if patch_size == 1: + return x + if x.dim() == 4: + x = rearrange(x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size, r=patch_size) + elif x.dim() == 5: + x = rearrange(x, + "b (c r q) f h w -> b c f (h q) (w r)", + q=patch_size, + r=patch_size) + return x + + +class Resample38(Resample): + + def __init__(self, dim, mode): + assert mode in ( + "none", + "upsample2d", + "upsample3d", + "downsample2d", + "downsample3d", + ) + super(Resample, self).__init__() + self.dim = dim + self.mode = mode + + # layers + if mode == "upsample2d": + self.resample = nn.Sequential( + Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), + nn.Conv2d(dim, dim, 3, padding=1), + ) + elif mode == "upsample3d": + self.resample = nn.Sequential( + Upsample(scale_factor=(2.0, 2.0), mode="nearest-exact"), + nn.Conv2d(dim, dim, 3, padding=1), + ) + self.time_conv = CausalConv3d(dim, dim * 2, (3, 1, 1), padding=(1, 0, 0)) + elif mode == "downsample2d": + self.resample = nn.Sequential( + nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)) + ) + elif mode == "downsample3d": + self.resample = nn.Sequential( + nn.ZeroPad2d((0, 1, 0, 1)), nn.Conv2d(dim, dim, 3, stride=(2, 2)) + ) + self.time_conv = CausalConv3d( + dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0) + ) + else: + self.resample = nn.Identity() + +class ResidualBlock(nn.Module): + + def __init__(self, in_dim, out_dim, dropout=0.0): + super().__init__() + self.in_dim = in_dim + self.out_dim = out_dim + + # layers + self.residual = nn.Sequential( + RMS_norm(in_dim, images=False), nn.SiLU(), + CausalConv3d(in_dim, out_dim, 3, padding=1), + RMS_norm(out_dim, images=False), nn.SiLU(), nn.Dropout(dropout), + CausalConv3d(out_dim, out_dim, 3, padding=1)) + self.shortcut = CausalConv3d(in_dim, out_dim, 1) \ + if in_dim != out_dim else nn.Identity() + + def forward(self, x, feat_cache=None, feat_idx=[0]): + h = self.shortcut(x) + for layer in self.residual: + if check_is_instance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2), + cache_x, + ], dim=2) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + h + + +class AttentionBlock(nn.Module): + """ + Causal self-attention with a single head. + """ + + def __init__(self, dim): + super().__init__() + self.dim = dim + + # layers + self.norm = RMS_norm(dim) + self.to_qkv = nn.Conv2d(dim, dim * 3, 1) + self.proj = nn.Conv2d(dim, dim, 1) + + # zero out the last layer params + nn.init.zeros_(self.proj.weight) + + def forward(self, x): + identity = x + b, c, t, h, w = x.size() + x = rearrange(x, 'b c t h w -> (b t) c h w') + x = self.norm(x) + # compute query, key, value + q, k, v = self.to_qkv(x).reshape(b * t, 1, c * 3, -1).permute( + 0, 1, 3, 2).contiguous().chunk(3, dim=-1) + + # apply attention + x = F.scaled_dot_product_attention( + q, + k, + v, + #attn_mask=block_causal_mask(q, block_size=h * w) + ) + x = x.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w) + + # output + x = self.proj(x) + x = rearrange(x, '(b t) c h w-> b c t h w', t=t) + return x + identity + + +class AvgDown3D(nn.Module): + def __init__( + self, + in_channels, + out_channels, + factor_t, + factor_s=1, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.factor_t = factor_t + self.factor_s = factor_s + self.factor = self.factor_t * self.factor_s * self.factor_s + + assert in_channels * self.factor % out_channels == 0 + self.group_size = in_channels * self.factor // out_channels + + def forward(self, x: torch.Tensor) -> torch.Tensor: + pad_t = (self.factor_t - x.shape[2] % self.factor_t) % self.factor_t + pad = (0, 0, 0, 0, pad_t, 0) + x = F.pad(x, pad) + B, C, T, H, W = x.shape + x = x.view( + B, + C, + T // self.factor_t, + self.factor_t, + H // self.factor_s, + self.factor_s, + W // self.factor_s, + self.factor_s, + ) + x = x.permute(0, 1, 3, 5, 7, 2, 4, 6).contiguous() + x = x.view( + B, + C * self.factor, + T // self.factor_t, + H // self.factor_s, + W // self.factor_s, + ) + x = x.view( + B, + self.out_channels, + self.group_size, + T // self.factor_t, + H // self.factor_s, + W // self.factor_s, + ) + x = x.mean(dim=2) + return x + + +class DupUp3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + factor_t, + factor_s=1, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + + self.factor_t = factor_t + self.factor_s = factor_s + self.factor = self.factor_t * self.factor_s * self.factor_s + + assert out_channels * self.factor % in_channels == 0 + self.repeats = out_channels * self.factor // in_channels + + def forward(self, x: torch.Tensor, first_chunk=False) -> torch.Tensor: + x = x.repeat_interleave(self.repeats, dim=1) + x = x.view( + x.size(0), + self.out_channels, + self.factor_t, + self.factor_s, + self.factor_s, + x.size(2), + x.size(3), + x.size(4), + ) + x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous() + x = x.view( + x.size(0), + self.out_channels, + x.size(2) * self.factor_t, + x.size(4) * self.factor_s, + x.size(6) * self.factor_s, + ) + if first_chunk: + x = x[:, :, self.factor_t - 1 :, :, :] + return x + + +class Down_ResidualBlock(nn.Module): + def __init__( + self, in_dim, out_dim, dropout, mult, temperal_downsample=False, down_flag=False + ): + super().__init__() + + # Shortcut path with downsample + self.avg_shortcut = AvgDown3D( + in_dim, + out_dim, + factor_t=2 if temperal_downsample else 1, + factor_s=2 if down_flag else 1, + ) + + # Main path with residual blocks and downsample + downsamples = [] + for _ in range(mult): + downsamples.append(ResidualBlock(in_dim, out_dim, dropout)) + in_dim = out_dim + + # Add the final downsample block + if down_flag: + mode = "downsample3d" if temperal_downsample else "downsample2d" + downsamples.append(Resample38(out_dim, mode=mode)) + + self.downsamples = nn.Sequential(*downsamples) + + def forward(self, x, feat_cache=None, feat_idx=[0]): + x_copy = x.clone() + for module in self.downsamples: + x = module(x, feat_cache, feat_idx) + + return x + self.avg_shortcut(x_copy) + + +class Up_ResidualBlock(nn.Module): + def __init__( + self, in_dim, out_dim, dropout, mult, temperal_upsample=False, up_flag=False + ): + super().__init__() + # Shortcut path with upsample + if up_flag: + self.avg_shortcut = DupUp3D( + in_dim, + out_dim, + factor_t=2 if temperal_upsample else 1, + factor_s=2 if up_flag else 1, + ) + else: + self.avg_shortcut = None + + # Main path with residual blocks and upsample + upsamples = [] + for _ in range(mult): + upsamples.append(ResidualBlock(in_dim, out_dim, dropout)) + in_dim = out_dim + + # Add the final upsample block + if up_flag: + mode = "upsample3d" if temperal_upsample else "upsample2d" + upsamples.append(Resample38(out_dim, mode=mode)) + + self.upsamples = nn.Sequential(*upsamples) + + def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False): + x_main = x.clone() + for module in self.upsamples: + x_main = module(x_main, feat_cache, feat_idx) + if self.avg_shortcut is not None: + x_shortcut = self.avg_shortcut(x, first_chunk) + return x_main + x_shortcut + else: + return x_main + + +class Encoder3d(nn.Module): + + def __init__(self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[True, True, False], + dropout=0.0): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_downsample = temperal_downsample + + # dimensions + dims = [dim * u for u in [1] + dim_mult] + scale = 1.0 + + # init block + self.conv1 = CausalConv3d(3, dims[0], 3, padding=1) + + # downsample blocks + downsamples = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + # residual (+attention) blocks + for _ in range(num_res_blocks): + downsamples.append(ResidualBlock(in_dim, out_dim, dropout)) + if scale in attn_scales: + downsamples.append(AttentionBlock(out_dim)) + in_dim = out_dim + + # downsample block + if i != len(dim_mult) - 1: + mode = 'downsample3d' if temperal_downsample[ + i] else 'downsample2d' + downsamples.append(Resample(out_dim, mode=mode)) + scale /= 2.0 + self.downsamples = nn.Sequential(*downsamples) + + # middle blocks + self.middle = nn.Sequential(ResidualBlock(out_dim, out_dim, dropout), + AttentionBlock(out_dim), + ResidualBlock(out_dim, out_dim, dropout)) + + # output blocks + self.head = nn.Sequential(RMS_norm(out_dim, images=False), nn.SiLU(), + CausalConv3d(out_dim, z_dim, 3, padding=1)) + + def forward(self, x, feat_cache=None, feat_idx=[0]): + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2), + cache_x, + ], dim=2) + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + ## downsamples + for layer in self.downsamples: + if feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## middle + for layer in self.middle: + if check_is_instance(layer, ResidualBlock) and feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## head + for layer in self.head: + if check_is_instance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2), + cache_x, + ], dim=2) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + + +class Encoder3d_38(nn.Module): + + def __init__(self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[False, True, True], + dropout=0.0): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_downsample = temperal_downsample + + # dimensions + dims = [dim * u for u in [1] + dim_mult] + scale = 1.0 + + # init block + self.conv1 = CausalConv3d(12, dims[0], 3, padding=1) + + # downsample blocks + downsamples = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + t_down_flag = ( + temperal_downsample[i] if i < len(temperal_downsample) else False + ) + downsamples.append( + Down_ResidualBlock( + in_dim=in_dim, + out_dim=out_dim, + dropout=dropout, + mult=num_res_blocks, + temperal_downsample=t_down_flag, + down_flag=i != len(dim_mult) - 1, + ) + ) + scale /= 2.0 + self.downsamples = nn.Sequential(*downsamples) + + # middle blocks + self.middle = nn.Sequential( + ResidualBlock(out_dim, out_dim, dropout), + AttentionBlock(out_dim), + ResidualBlock(out_dim, out_dim, dropout), + ) + + # # output blocks + self.head = nn.Sequential( + RMS_norm(out_dim, images=False), + nn.SiLU(), + CausalConv3d(out_dim, z_dim, 3, padding=1), + ) + + + def forward(self, x, feat_cache=None, feat_idx=[0]): + + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2), + cache_x, + ], + dim=2, + ) + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + ## downsamples + for layer in self.downsamples: + if feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## middle + for layer in self.middle: + if isinstance(layer, ResidualBlock) and feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## head + for layer in self.head: + if isinstance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2), + cache_x, + ], + dim=2, + ) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + + return x + + +class Decoder3d(nn.Module): + + def __init__(self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_upsample=[False, True, True], + dropout=0.0): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_upsample = temperal_upsample + + # dimensions + dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]] + scale = 1.0 / 2**(len(dim_mult) - 2) + + # init block + self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1) + + # middle blocks + self.middle = nn.Sequential(ResidualBlock(dims[0], dims[0], dropout), + AttentionBlock(dims[0]), + ResidualBlock(dims[0], dims[0], dropout)) + + # upsample blocks + upsamples = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + # residual (+attention) blocks + if i == 1 or i == 2 or i == 3: + in_dim = in_dim // 2 + for _ in range(num_res_blocks + 1): + upsamples.append(ResidualBlock(in_dim, out_dim, dropout)) + if scale in attn_scales: + upsamples.append(AttentionBlock(out_dim)) + in_dim = out_dim + + # upsample block + if i != len(dim_mult) - 1: + mode = 'upsample3d' if temperal_upsample[i] else 'upsample2d' + upsamples.append(Resample(out_dim, mode=mode)) + scale *= 2.0 + self.upsamples = nn.Sequential(*upsamples) + + # output blocks + self.head = nn.Sequential(RMS_norm(out_dim, images=False), nn.SiLU(), + CausalConv3d(out_dim, 3, 3, padding=1)) + + def forward(self, x, feat_cache=None, feat_idx=[0]): + ## conv1 + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2), + cache_x, + ], dim=2) + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + ## middle + for layer in self.middle: + if check_is_instance(layer, ResidualBlock) and feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## upsamples + for layer in self.upsamples: + if feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## head + for layer in self.head: + if check_is_instance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + # cache last frame of last two chunk + cache_x = torch.cat([ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2), + cache_x, + ], dim=2) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + + + +class Decoder3d_38(nn.Module): + + def __init__(self, + dim=128, + z_dim=4, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_upsample=[False, True, True], + dropout=0.0): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_upsample = temperal_upsample + + # dimensions + dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]] + scale = 1.0 / 2 ** (len(dim_mult) - 2) + # init block + self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1) + + # middle blocks + self.middle = nn.Sequential(ResidualBlock(dims[0], dims[0], dropout), + AttentionBlock(dims[0]), + ResidualBlock(dims[0], dims[0], dropout)) + + # upsample blocks + upsamples = [] + for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): + t_up_flag = temperal_upsample[i] if i < len(temperal_upsample) else False + upsamples.append( + Up_ResidualBlock(in_dim=in_dim, + out_dim=out_dim, + dropout=dropout, + mult=num_res_blocks + 1, + temperal_upsample=t_up_flag, + up_flag=i != len(dim_mult) - 1)) + self.upsamples = nn.Sequential(*upsamples) + + # output blocks + self.head = nn.Sequential(RMS_norm(out_dim, images=False), nn.SiLU(), + CausalConv3d(out_dim, 12, 3, padding=1)) + + + def forward(self, x, feat_cache=None, feat_idx=[0], first_chunk=False): + if feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2), + cache_x, + ], + dim=2, + ) + x = self.conv1(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = self.conv1(x) + + for layer in self.middle: + if check_is_instance(layer, ResidualBlock) and feat_cache is not None: + x = layer(x, feat_cache, feat_idx) + else: + x = layer(x) + + ## upsamples + for layer in self.upsamples: + if feat_cache is not None: + x = layer(x, feat_cache, feat_idx, first_chunk) + else: + x = layer(x) + + ## head + for layer in self.head: + if check_is_instance(layer, CausalConv3d) and feat_cache is not None: + idx = feat_idx[0] + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and feat_cache[idx] is not None: + cache_x = torch.cat( + [ + feat_cache[idx][:, :, -1, :, :].unsqueeze(2), + cache_x, + ], + dim=2, + ) + x = layer(x, feat_cache[idx]) + feat_cache[idx] = cache_x + feat_idx[0] += 1 + else: + x = layer(x) + return x + + +def count_conv3d(model): + count = 0 + for m in model.modules(): + if isinstance(m, CausalConv3d): + count += 1 + return count + + +class VideoVAE_(nn.Module): + + def __init__(self, + dim=96, + z_dim=16, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[False, True, True], + dropout=0.0): + super().__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_downsample = temperal_downsample + self.temperal_upsample = temperal_downsample[::-1] + + # modules + self.encoder = Encoder3d(dim, z_dim * 2, dim_mult, num_res_blocks, + attn_scales, self.temperal_downsample, dropout) + self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1) + self.conv2 = CausalConv3d(z_dim, z_dim, 1) + self.decoder = Decoder3d(dim, z_dim, dim_mult, num_res_blocks, + attn_scales, self.temperal_upsample, dropout) + + self._enc_conv_num = count_conv3d(self.encoder) + self._dec_conv_num = count_conv3d(self.decoder) + + def encode(self, x, scale): + feat_map = [None] * self._enc_conv_num + + t = x.shape[2] + iter_ = 1 + (t - 1) // 4 + + out = self.encoder( + x[:, :, :1, :, :], + feat_cache=feat_map, + feat_idx=[0], + ) + + for i in range(1, iter_): + out_ = self.encoder( + x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :], + feat_cache=feat_map, + feat_idx=[0], + ) + out = torch.cat([out, out_], dim=2) + mu, _ = self.conv1(out).chunk(2, dim=1) + if isinstance(scale[0], torch.Tensor): + scale = [s.to(dtype=mu.dtype, device=mu.device) for s in scale] + mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view( + 1, self.z_dim, 1, 1, 1) + else: + scale = scale.to(dtype=mu.dtype, device=mu.device) + mu = (mu - scale[0]) * scale[1] + return mu + + def decode(self, z, scale): + feat_map = [None] * self._dec_conv_num + + # z: [b,c,t,h,w] + if isinstance(scale[0], torch.Tensor): + scale = [s.to(dtype=z.dtype, device=z.device) for s in scale] + z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view( + 1, self.z_dim, 1, 1, 1) + else: + scale = scale.to(dtype=z.dtype, device=z.device) + z = z / scale[1] + scale[0] + iter_ = z.shape[2] + x = self.conv2(z) + + out = self.decoder( + x[:, :, 0:1, :, :], + feat_cache=feat_map, + feat_idx=[0], + ) + + for i in range(1, iter_): + out_ = self.decoder( + x[:, :, i:i + 1, :, :], + feat_cache=feat_map, + feat_idx=[0], + ) + out = torch.cat([out, out_], dim=2) + return out + + def reparameterize(self, mu, log_var): + std = torch.exp(0.5 * log_var) + eps = torch.randn_like(std) + return eps * std + mu + + def sample(self, imgs, deterministic=False): + mu, log_var = self.encode(imgs) + if deterministic: + return mu + std = torch.exp(0.5 * log_var.clamp(-30.0, 20.0)) + return mu + std * torch.randn_like(std) + + +class WanVideoVAE(nn.Module): + + def __init__(self, z_dim=16, vae_pretrained_path: str | None = None): + super().__init__() + + mean = [ + -0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508, + 0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921 + ] + std = [ + 2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743, + 3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.9160 + ] + self.mean = torch.tensor(mean, device='cuda') + self.std = torch.tensor(std, device='cuda') + self.scale = [self.mean, 1.0 / self.std] + + # init model + self.model = VideoVAE_(z_dim=z_dim).eval().requires_grad_(False) + self.upsampling_factor = 8 + self.z_dim = z_dim + self.vae_pretrained_path = vae_pretrained_path + + + def build_1d_mask(self, length, left_bound, right_bound, border_width, device): + x = torch.ones((length,), device=device) + border = (torch.arange(border_width, device=device) + 1) + if not left_bound: + x[:border_width] = border / border_width + if not right_bound: + x[-border_width:] = torch.flip(border / border_width, dims=(0,)) + return x + + + def build_mask(self, data, is_bound, border_width): + _, _, _, H, W = data.shape + h = self.build_1d_mask(H, is_bound[0], is_bound[1], border_width[0], device=data.device) + w = self.build_1d_mask(W, is_bound[2], is_bound[3], border_width[1], device=data.device) + + h = repeat(h, "H -> H W", H=H, W=W) + w = repeat(w, "W -> H W", H=H, W=W) + + mask = torch.stack([h, w]).min(dim=0).values + mask = rearrange(mask, "H W -> 1 1 1 H W") + return mask + + + def tiled_decode(self, hidden_states, tile_size, tile_stride): + _, _, T, H, W = hidden_states.shape + size_h, size_w = tile_size + stride_h, stride_w = tile_stride + + # Split tasks + tasks = [] + for h in range(0, H, stride_h): + if (h-stride_h >= 0 and h-stride_h+size_h >= H): continue + for w in range(0, W, stride_w): + if (w-stride_w >= 0 and w-stride_w+size_w >= W): continue + h_, w_ = h + size_h, w + size_w + tasks.append((h, h_, w, w_)) + + out_T = T * 4 - 3 + weight = torch.zeros( + (1, 1, out_T, H * self.upsampling_factor, W * self.upsampling_factor), + dtype=hidden_states.dtype, + device=hidden_states.device + ) + values = torch.zeros( + (1, 3, out_T, H * self.upsampling_factor, W * self.upsampling_factor), + dtype=hidden_states.dtype, + device=hidden_states.device, + ) + + for h, h_, w, w_ in tqdm(tasks, desc="VAE decoding"): + hidden_states_batch = hidden_states[:, :, :, h:h_, w:w_] + hidden_states_batch = self.model.decode(hidden_states_batch, self.scale) + + mask = self.build_mask( + hidden_states_batch, + is_bound=(h==0, h_>=H, w==0, w_>=W), + border_width=((size_h - stride_h) * self.upsampling_factor, (size_w - stride_w) * self.upsampling_factor) + ).to(dtype=hidden_states.dtype) + + target_h = h * self.upsampling_factor + target_w = w * self.upsampling_factor + values[ + :, + :, + :, + target_h:target_h + hidden_states_batch.shape[3], + target_w:target_w + hidden_states_batch.shape[4], + ] += hidden_states_batch * mask + weight[ + :, + :, + :, + target_h: target_h + hidden_states_batch.shape[3], + target_w: target_w + hidden_states_batch.shape[4], + ] += mask + values = values / weight + values = values.clamp_(-1, 1) + return values + + def tiled_encode(self, video, tile_size, tile_stride): + _, _, T, H, W = video.shape + size_h, size_w = tile_size + stride_h, stride_w = tile_stride + + # Split tasks + tasks = [] + for h in range(0, H, stride_h): + if (h-stride_h >= 0 and h-stride_h+size_h >= H): continue + for w in range(0, W, stride_w): + if (w-stride_w >= 0 and w-stride_w+size_w >= W): continue + h_, w_ = h + size_h, w + size_w + tasks.append((h, h_, w, w_)) + + out_T = (T + 3) // 4 + weight = torch.zeros( + (1, 1, out_T, H // self.upsampling_factor, W // self.upsampling_factor), + dtype=video.dtype, + device=video.device, + ) + values = torch.zeros( + (1, self.z_dim, out_T, H // self.upsampling_factor, W // self.upsampling_factor), + dtype=video.dtype, + device=video.device, + ) + + for h, h_, w, w_ in tqdm(tasks, desc="VAE encoding"): + hidden_states_batch = video[:, :, :, h:h_, w:w_] + hidden_states_batch = self.model.encode(hidden_states_batch, self.scale) + + mask = self.build_mask( + hidden_states_batch, + is_bound=(h==0, h_>=H, w==0, w_>=W), + border_width=((size_h - stride_h) // self.upsampling_factor, (size_w - stride_w) // self.upsampling_factor) + ).to(dtype=video.dtype) + + target_h = h // self.upsampling_factor + target_w = w // self.upsampling_factor + values[ + :, + :, + :, + target_h:target_h + hidden_states_batch.shape[3], + target_w:target_w + hidden_states_batch.shape[4], + ] += hidden_states_batch * mask + weight[ + :, + :, + :, + target_h: target_h + hidden_states_batch.shape[3], + target_w: target_w + hidden_states_batch.shape[4], + ] += mask + values = values / weight + return values + + def single_encode(self, video): + x = self.model.encode(video, self.scale) + # The outputs of torch compile always need to be cloned before being used. + x = x.clone() + return x + + def single_decode(self, hidden_state): + video = self.model.decode(hidden_state, self.scale) + return video.clamp_(-1, 1) + + def encode(self, videos, tiled=False, tile_size=(34, 34), tile_stride=(18, 16)): + hidden_states = [] + for video in videos: + video = video.unsqueeze(0) + if tiled: + tile_size = (tile_size[0] * self.upsampling_factor, tile_size[1] * self.upsampling_factor) + tile_stride = (tile_stride[0] * self.upsampling_factor, tile_stride[1] * self.upsampling_factor) + hidden_state = self.tiled_encode(video, tile_size, tile_stride) + else: + hidden_state = self.single_encode(video) + hidden_state = hidden_state.squeeze(0) + hidden_states.append(hidden_state) + hidden_states = torch.stack(hidden_states) + return hidden_states + + def decode(self, hidden_states, tiled=False, tile_size=(34, 34), tile_stride=(18, 16)): + if tiled: + video = self.tiled_decode(hidden_states, tile_size, tile_stride) + else: + video = self.single_decode(hidden_states) + return video + + + @staticmethod + def state_dict_converter(): + return WanVideoVAEStateDictConverter() + + +class WanVideoVAEStateDictConverter: + + def __init__(self): + pass + + def from_civitai(self, state_dict): + state_dict_ = {} + if 'model_state' in state_dict: + state_dict = state_dict['model_state'] + for name in state_dict: + state_dict_['model.' + name] = state_dict[name] + return state_dict_ + + +class VideoVAE38_(VideoVAE_): + + def __init__(self, + dim=160, + z_dim=48, + dec_dim=256, + dim_mult=[1, 2, 4, 4], + num_res_blocks=2, + attn_scales=[], + temperal_downsample=[False, True, True], + dropout=0.0): + super(VideoVAE_, self).__init__() + self.dim = dim + self.z_dim = z_dim + self.dim_mult = dim_mult + self.num_res_blocks = num_res_blocks + self.attn_scales = attn_scales + self.temperal_downsample = temperal_downsample + self.temperal_upsample = temperal_downsample[::-1] + + # modules + self.encoder = Encoder3d_38(dim, z_dim * 2, dim_mult, num_res_blocks, + attn_scales, self.temperal_downsample, dropout) + self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1) + self.conv2 = CausalConv3d(z_dim, z_dim, 1) + self.decoder = Decoder3d_38(dec_dim, z_dim, dim_mult, num_res_blocks, + attn_scales, self.temperal_upsample, dropout) + + self._enc_conv_num = count_conv3d(self.encoder) + self._dec_conv_num = count_conv3d(self.decoder) + + def encode(self, x, scale): + feat_map = [None] * self._enc_conv_num + + x = patchify(x, patch_size=2) + t = x.shape[2] + iter_ = 1 + (t - 1) // 4 + + out = self.encoder( + x[:, :, :1, :, :], + feat_cache=feat_map, + feat_idx=[0], + ) + + for i in range(1, iter_): + out_ = self.encoder( + x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :], + feat_cache=feat_map, + feat_idx=[0], + ) + out = torch.cat([out, out_], dim=2) + + mu, _ = self.conv1(out).chunk(2, dim=1) + if isinstance(scale[0], torch.Tensor): + scale = [s.to(dtype=mu.dtype, device=mu.device) for s in scale] + mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view( + 1, self.z_dim, 1, 1, 1) + else: + scale = scale.to(dtype=mu.dtype, device=mu.device) + mu = (mu - scale[0]) * scale[1] + return mu + + def decode(self, z, scale): + feat_map = [None] * self._dec_conv_num + + if isinstance(scale[0], torch.Tensor): + scale = [s.to(dtype=z.dtype, device=z.device) for s in scale] + z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view( + 1, self.z_dim, 1, 1, 1) + else: + scale = scale.to(dtype=z.dtype, device=z.device) + z = z / scale[1] + scale[0] + iter_ = z.shape[2] + x = self.conv2(z) + + out = self.decoder( + x[:, :, 0:1, :, :], + feat_cache=feat_map, + feat_idx=[0], + first_chunk=True, + ) + + for i in range(1, iter_): + out_ = self.decoder( + x[:, :, i:i + 1, :, :], + feat_cache=feat_map, + feat_idx=[0], + ) + out = torch.cat([out, out_], 2) + + out = unpatchify(out, patch_size=2) + return out + + +class WanVideoVAE38(WanVideoVAE): + + def __init__(self, z_dim=48, dim=160, vae_pretrained_path: str | None = None): + super(WanVideoVAE, self).__init__() + + mean = [ + -0.2289, -0.0052, -0.1323, -0.2339, -0.2799, 0.0174, 0.1838, 0.1557, + -0.1382, 0.0542, 0.2813, 0.0891, 0.1570, -0.0098, 0.0375, -0.1825, + -0.2246, -0.1207, -0.0698, 0.5109, 0.2665, -0.2108, -0.2158, 0.2502, + -0.2055, -0.0322, 0.1109, 0.1567, -0.0729, 0.0899, -0.2799, -0.1230, + -0.0313, -0.1649, 0.0117, 0.0723, -0.2839, -0.2083, -0.0520, 0.3748, + 0.0152, 0.1957, 0.1433, -0.2944, 0.3573, -0.0548, -0.1681, -0.0667 + ] + std = [ + 0.4765, 1.0364, 0.4514, 1.1677, 0.5313, 0.4990, 0.4818, 0.5013, + 0.8158, 1.0344, 0.5894, 1.0901, 0.6885, 0.6165, 0.8454, 0.4978, + 0.5759, 0.3523, 0.7135, 0.6804, 0.5833, 1.4146, 0.8986, 0.5659, + 0.7069, 0.5338, 0.4889, 0.4917, 0.4069, 0.4999, 0.6866, 0.4093, + 0.5709, 0.6065, 0.6415, 0.4944, 0.5726, 1.2042, 0.5458, 1.6887, + 0.3971, 1.0600, 0.3943, 0.5537, 0.5444, 0.4089, 0.7468, 0.7744 + ] + self.mean = torch.tensor(mean, device='cuda') + self.std = torch.tensor(std, device='cuda') + self.scale = [self.mean, 1.0 / self.std] + + # init model + self.model = VideoVAE38_(z_dim=z_dim, dim=dim).eval().requires_grad_(False) + self.upsampling_factor = 16 + self.z_dim = z_dim + self.vae_pretrained_path = vae_pretrained_path \ No newline at end of file diff --git a/flagscale/train/models/dreamzero/n1_5/__init__.py b/flagscale/train/models/dreamzero/n1_5/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/flagscale/train/models/dreamzero/n1_5/action_head/__init__.py b/flagscale/train/models/dreamzero/n1_5/action_head/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/flagscale/train/models/dreamzero/n1_5/action_head/base_action_head.py b/flagscale/train/models/dreamzero/n1_5/action_head/base_action_head.py new file mode 100644 index 0000000000..e8398622d4 --- /dev/null +++ b/flagscale/train/models/dreamzero/n1_5/action_head/base_action_head.py @@ -0,0 +1,32 @@ +from abc import ABC, abstractmethod + +from torch import nn +from transformers.feature_extraction_utils import BatchFeature + + +class ActionHead(ABC, nn.Module): + def __init__(self): + super(ActionHead, self).__init__() + + @abstractmethod + def forward(self, backbone_output: BatchFeature, action_input: BatchFeature) -> BatchFeature: + pass + + def get_action( + self, + backbone_output: BatchFeature, + action_input: BatchFeature, + num_action_samples: int = 1, + inference_batch_size: int = 32, + ) -> BatchFeature: + # Used for predicting actions during inference + # By default, the action head does the same thing as a normal forward pass + return self.forward(backbone_output, action_input) + + def prepare_input(self, batch: dict) -> BatchFeature: + pass + + def set_override_kwargs(self, **kwargs): + for key, value in kwargs.items(): + setattr(self.config, key, value) + setattr(self, key, value) diff --git a/flagscale/train/models/dreamzero/n1_5/modules/__init__.py b/flagscale/train/models/dreamzero/n1_5/modules/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/flagscale/train/models/dreamzero/n1_5/modules/action_encoder.py b/flagscale/train/models/dreamzero/n1_5/modules/action_encoder.py new file mode 100644 index 0000000000..55a39303c2 --- /dev/null +++ b/flagscale/train/models/dreamzero/n1_5/modules/action_encoder.py @@ -0,0 +1,87 @@ +### Implemented based on Pi0 Action Encoding +import torch +import torch.nn as nn + + +def swish(x): + return x * torch.sigmoid(x) + + +class SinusoidalPositionalEncoding(nn.Module): + """ + Produces a sinusoidal encoding of shape (B, T, w) + given timesteps of shape (B, T). + """ + + def __init__(self, embedding_dim): + super().__init__() + self.embedding_dim = embedding_dim + + def forward(self, timesteps): + # timesteps: shape (B, T) + # We'll compute sin/cos frequencies across dim T + timesteps = timesteps.float() # ensure float + + B, T = timesteps.shape + device = timesteps.device + + half_dim = self.embedding_dim // 2 + # typical log space frequencies for sinusoidal encoding + exponent = -torch.arange(half_dim, dtype=torch.float, device=device) * ( + torch.log(torch.tensor(10000.0)) / half_dim + ) + # Expand timesteps to (B, T, 1) then multiply + freqs = timesteps.unsqueeze(-1) * exponent.exp() # (B, T, half_dim) + + sin = torch.sin(freqs) + cos = torch.cos(freqs) + enc = torch.cat([sin, cos], dim=-1) # (B, T, w) + + return enc + + +class ActionEncoder(nn.Module): + def __init__(self, action_dim, hidden_size): + super().__init__() + self.hidden_size = hidden_size + + # W1: R^{w x d}, W2: R^{w x 2w}, W3: R^{w x w} + self.W1 = nn.Linear(action_dim, hidden_size) # (d -> w) + self.W2 = nn.Linear(2 * hidden_size, hidden_size) # (2w -> w) + self.W3 = nn.Linear(hidden_size, hidden_size) # (w -> w) + + self.pos_encoding = SinusoidalPositionalEncoding(hidden_size) + + def forward(self, actions, timesteps): + """ + actions: shape (B, T, action_dim) + timesteps: shape (B,) -- a single scalar per batch item + returns: shape (B, T, hidden_size) + """ + B, T, _ = actions.shape + + # 1) Expand each batch's single scalar time 'tau' across all T steps + # so that shape => (B, T) + # e.g. if timesteps is (B,), replicate across T + if timesteps.dim() == 1 and timesteps.shape[0] == B: + # shape (B,) => (B,T) + timesteps = timesteps.unsqueeze(1).expand(-1, T) + else: + raise ValueError( + "Expected `timesteps` to have shape (B,) so we can replicate across T." + ) + + # 2) Standard action MLP step for shape => (B, T, w) + a_emb = self.W1(actions) + + # 3) Get the sinusoidal encoding (B, T, w) + tau_emb = self.pos_encoding(timesteps).to(dtype=a_emb.dtype) + + # 4) Concat along last dim => (B, T, 2w), then W2 => (B, T, w), swish + x = torch.cat([a_emb, tau_emb], dim=-1) + x = swish(self.W2(x)) + + # 5) Finally W3 => (B, T, w) + x = self.W3(x) + + return x diff --git a/flagscale/train/train_config.py b/flagscale/train/train_config.py index 107fc72f6a..281065c968 100644 --- a/flagscale/train/train_config.py +++ b/flagscale/train/train_config.py @@ -44,10 +44,12 @@ class SchedulerConfig(BaseModel): Uses warmup_steps, scheduler_kwargs. For backward compatibility with pi0/pi0.5, the legacy fields (decay_steps, decay_lr) are kept. + warmup_ratio takes precedence over warmup_steps if both are set. """ name: str | None = None warmup_steps: int = 1000 + warmup_ratio: float | None = None # If set, overrides warmup_steps as ratio * train_steps scheduler_kwargs: dict[str, Any] | None = None # Used by cosine_decay_with_warmup and legacy pi0/pi0.5 @@ -241,7 +243,7 @@ def __getattr__(self, name): @field_validator("model_name") @classmethod def validate_model_name(cls, v): - valid_names = {"pi0", "pi0.5", "qwen_gr00t", "gr00t_n1_5"} + valid_names = {"pi0", "pi0.5", "qwen_gr00t", "gr00t_n1_5", "dreamzero"} if v not in valid_names: raise ValueError(f"Invalid model_name: {v}. Must be one of {valid_names}") return v diff --git a/flagscale/train/train_dreamzero.py b/flagscale/train/train_dreamzero.py new file mode 100644 index 0000000000..d33e0fd6b0 --- /dev/null +++ b/flagscale/train/train_dreamzero.py @@ -0,0 +1,565 @@ +#!/usr/bin/env python3 +# Copyright (c) 2025, FlagScale Authors. All rights reserved. +""" +DreamZero Training Entrypoint for FlagScale native backend. + +Uses FSDP2 for distributed training (matching original DeepSpeed ZeRO-2 strategy). +No TP/PP — pure data parallelism with sharded parameters. + +Training loop: load data -> forward (VAE encode + noise + DiT predict) -> loss -> backward. +""" + +import os +import random +import time +from collections.abc import Iterator +from pathlib import Path + +import numpy as np +import torch +import torch.distributed as dist +from torch.distributed._composable.fsdp import fully_shard, MixedPrecisionPolicy +from torch.distributed.device_mesh import init_device_mesh + +from omegaconf import OmegaConf, DictConfig +from flagscale.logger import logger +from flagscale.train.train_config import TrainConfig +from flagscale.train.utils.logging_utils import AverageMeter +from flagscale.train.utils.train_utils import ( + get_step_checkpoint_dir, + save_checkpoint, + update_last_checkpoint, +) +from flagscale.train.utils.optim_setup import setup_optimizer_and_scheduler + + +def set_seed(seed: int): + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + torch.backends.cudnn.enabled = False + torch.backends.cudnn.benchmark = False + torch.backends.cudnn.deterministic = False + torch.backends.cuda.matmul.allow_tf32 = True + # Raise Dynamo limits for compiled attention sub-methods. + # The blockwise causal attention pattern has multiple code paths + # that trigger recompiles — need higher limits to avoid FailOnRecompileLimitHit. + torch._dynamo.config.cache_size_limit = 256 + torch._dynamo.config.accumulated_cache_size_limit = 512 + + +def apply_fsdp2(policy, device_mesh): + """Apply FSDP2 sharding to DreamZero. + + Strategy: shard every large sub-module individually to minimize peak memory. + During forward, FSDP2 all-gathers ONE module's params at a time and reshards after. + This is equivalent to ZeRO-3 (param + grad + optimizer sharding). + + Model structure: + policy.action_head.text_encoder — T5-XXL (~4.7B, frozen) + policy.action_head.image_encoder — CLIP (~0.6B, frozen) + policy.action_head.vae — VAE (~0.1B, frozen) + policy.action_head.model — DiT (40 blocks, ~14B, trainable) + policy.action_head.{action,state}_{encoder,decoder} — small projectors (trainable) + """ + # Mixed precision policy for FSDP2. + # param_dtype=bfloat16: store and all-gather sharded params in bf16 for memory efficiency. + # With 8-GPU sharding: 30.7GB model / 8 = 3.8GB per GPU (vs 7.7GB with fp32). + # The optimizer (AdamW) maintains fp32 master weights, matching DeepSpeed ZeRO-2. + # reduce_dtype=float32: reduce-scatter gradients in fp32 for numerical accuracy. + # Tested bf16 reduce — no speed gain on single-node NVLink (bandwidth-saturated). + # cast_forward_inputs=False: activations handled by autocast, not FSDP boundaries. + mp_policy = MixedPrecisionPolicy( + param_dtype=torch.bfloat16, + reduce_dtype=torch.float32, + cast_forward_inputs=False, + ) + + # Fix scalar parameters (FSDP2 requires ndim >= 1) + for name, module in policy.named_modules(): + for pname, param in list(module.named_parameters(recurse=False)): + if param.ndim == 0: + new_param = torch.nn.Parameter( + param.data.reshape(1), requires_grad=param.requires_grad + ) + setattr(module, pname, new_param) + + ah = policy.action_head + + # FROZEN encoders (VAE, T5, CLIP): do NOT shard with FSDP2. + # They have no gradients/optimizer states, so sharding gives no memory benefit + # for optimizer. Also, VAE has 3D convolutions that are incompatible with DTensor. + # Keep them replicated (same as ZeRO-2 behavior for frozen params). + logger.info("Frozen encoders (text/image/vae) kept replicated (not sharded)") + + # 4. Shard DiT blocks individually (each ~400M params) + # In full mode: action_head.model.blocks directly + # In LoRA mode: action_head.model = PeftModel -> base_model -> model -> blocks + blocks = None + if hasattr(ah, "model") and hasattr(ah.model, "blocks"): + blocks = ah.model.blocks + elif hasattr(ah, "model") and hasattr(ah.model, "base_model"): + base = ah.model.base_model + if hasattr(base, "model") and hasattr(base.model, "blocks"): + blocks = base.model.blocks + elif hasattr(base, "blocks"): + blocks = base.blocks + + if blocks is not None: + logger.info(f"Sharding {len(blocks)} DiT blocks with FSDP2...") + for i, block in enumerate(blocks): + fully_shard(block, mesh=device_mesh, mp_policy=mp_policy, reshard_after_forward=True) + if i == 0: + logger.info(f" Block 0 sharded successfully") + logger.info(f" All {len(blocks)} blocks sharded") + else: + logger.warning("Could not find DiT blocks for per-block FSDP2 sharding") + + # 5. Shard the DiT model wrapper (ah.model) — NOT the full policy + # This catches remaining DiT flat params (head, embeddings) without sharding frozen encoders. + # With blocks already sharded individually, this only processes root-level params (~200MB). + + # Debug: check what params are NOT yet sharded (managed by outer fully_shard) + from torch.distributed._tensor import DTensor + unsharded_size = 0 + sharded_size = 0 + for name, p in ah.model.named_parameters(): + if isinstance(p, DTensor): + sharded_size += p.numel() * p.element_size() + else: + unsharded_size += p.numel() * p.element_size() + logger.info(f"Before outer shard: unsharded={unsharded_size/1024**3:.2f}GB, already_sharded={sharded_size/1024**3:.2f}GB") + + logger.info("Sharding DiT model (PeftModel wrapper)...") + fully_shard(ah.model, mesh=device_mesh, mp_policy=mp_policy) + + return policy + + +def apply_compile_and_reentrant(policy): + """Apply torch.compile to attention sub-methods and switch to use_reentrant=True. + + Key optimizations from RLinf's DreamZero training: + 1. torch.compile(mode="reduce-overhead") on the 4 attention processing functions + inside CausalWanSelfAttention — enables CUDA graphs for these hot inner loops. + 2. use_reentrant=True gradient checkpointing — avoids tensor pack/unpack hooks overhead + (requires passing block args as positional, not keyword). + """ + ah = policy.action_head + dit = ah.model + + # 1. Compile attention sub-methods on each block's self_attn + compiled_count = 0 + for block in dit.blocks: + # The attention module is at block.self_attn (CausalWanSelfAttention) + self_attn = getattr(block, "self_attn", None) + if self_attn is None: + continue + for method_name in ( + "_process_clean_image_only", + "_process_state_blocks", + "_process_noisy_image_blocks", + "_process_noisy_action_blocks", + ): + if hasattr(self_attn, method_name): + original = getattr(self_attn, method_name) + compiled = torch.compile(original, mode="reduce-overhead") + setattr(self_attn, method_name, compiled) + compiled_count += 1 + + logger.info(f"Compiled {compiled_count} attention sub-methods with mode='reduce-overhead'") + + # 2. Switch to use_reentrant=True gradient checkpointing + # This is faster (no pack/unpack hooks) and compatible with CUDA graphs. + if hasattr(dit, "gradient_checkpointing"): + dit.gradient_checkpointing = True + if hasattr(dit, "gradient_checkpointing_use_reentrant"): + dit.gradient_checkpointing_use_reentrant = True + else: + dit.gradient_checkpointing_use_reentrant = True + logger.info("Gradient checkpointing set to use_reentrant=True") + + +def safe_cycle(iterable) -> Iterator: + """Cycle over iterable safely.""" + iterator = iter(iterable) + while True: + try: + yield next(iterator) + except StopIteration: + iterator = iter(iterable) + + +def get_model(config: TrainConfig): + """Instantiate DreamZero model from config.""" + from flagscale.train.models.dreamzero import DreamZeroPolicy + from flagscale.train.models.dreamzero.dreamzero_model import DreamZeroConfig + + model_cfg = config.model + # pretrained_model_path is where the DreamZero checkpoint lives + ckpt_dir = getattr(model_cfg, "pretrained_model_path", None) or model_cfg.checkpoint_dir + dreamzero_config = DreamZeroConfig( + model_path=ckpt_dir, + tokenizer_path=getattr(model_cfg, "tokenizer_path", ckpt_dir), + action_horizon=getattr(model_cfg, "action_horizon", 24), + action_dim=getattr(model_cfg, "action_dim", 32), + max_state_dim=getattr(model_cfg, "max_state_dim", 64), + num_frames=getattr(model_cfg, "num_frames", 33), + num_frame_per_block=getattr(model_cfg, "num_frame_per_block", 2), + num_action_per_block=getattr(model_cfg, "num_action_per_block", 24), + num_state_per_block=getattr(model_cfg, "num_state_per_block", 1), + frame_seqlen=getattr(model_cfg, "frame_seqlen", 880), + use_gradient_checkpointing=getattr(model_cfg, "use_gradient_checkpointing", True), + train_architecture=getattr(model_cfg, "train_architecture", "full"), + tune_diffusion_model=getattr(model_cfg, "tune_diffusion_model", True), + tune_projector=getattr(model_cfg, "tune_projector", True), + compute_dtype=getattr(model_cfg, "compute_dtype", "bfloat16"), + embodiment_tag=getattr(model_cfg, "embodiment_tag", "libero"), + lora_rank=getattr(model_cfg, "lora_rank", 4), + lora_alpha=getattr(model_cfg, "lora_alpha", 4), + lora_target_modules=getattr(model_cfg, "lora_target_modules", "q,k,v,o,ffn.0,ffn.2"), + # Pretrained paths for component loading (Wan2.1 DIT, T5, CLIP, VAE) + dit_version=getattr(model_cfg, "dit_version", None), + text_encoder_pretrained_path=getattr(model_cfg, "text_encoder_pretrained_path", None), + image_encoder_pretrained_path=getattr(model_cfg, "image_encoder_pretrained_path", None), + vae_pretrained_path=getattr(model_cfg, "vae_pretrained_path", None), + ) + policy = DreamZeroPolicy.from_pretrained_dreamzero(dreamzero_config) + return policy + + +def get_dataloader(config: TrainConfig, seed: int = 42): + """Build DreamZero DataLoader from config. + + Uses aligned_dataloader (reference transform chain) by default to ensure + data pipeline matches the reference implementation exactly. + """ + data_cfg = config.data + model_cfg = config.model + system_cfg = config.system + + use_aligned = getattr(data_cfg, "use_aligned_dataloader", True) + + if use_aligned: + from flagscale.train.datasets.dreamzero.aligned_dataloader import build_dataloader_aligned + + dataloader = build_dataloader_aligned( + data_path=data_cfg.data_path, + tokenizer_path=getattr(data_cfg, "tokenizer_path", None) or getattr(model_cfg, "tokenizer_path", None), + batch_size=system_cfg.batch_size, + num_workers=system_cfg.num_workers, + num_frames=getattr(model_cfg, "num_frames", 33), + action_horizon=getattr(model_cfg, "action_horizon", 24), + state_horizon=getattr(model_cfg, "state_horizon", 1), + action_dim=getattr(model_cfg, "action_dim", 32), + max_state_dim=getattr(model_cfg, "max_state_dim", 64), + embodiment_tag=getattr(data_cfg, "embodiment_tag", "oxe_droid"), + embodiment_tag_mapping=getattr(data_cfg, "embodiment_tag_mapping", None), + image_size=tuple(getattr(data_cfg, "image_size", [176, 320])), + max_text_length=getattr(data_cfg, "max_text_length", 512), + num_views=getattr(model_cfg, "num_views", 3), + max_chunk_size=getattr(data_cfg, "max_chunk_size", 4), + seed=seed, + ) + else: + from flagscale.train.datasets.dreamzero import build_dataloader + + dataloader = build_dataloader( + data_path=data_cfg.data_path, + tokenizer_path=getattr(data_cfg, "tokenizer_path", None) or getattr(model_cfg, "tokenizer_path", None), + batch_size=system_cfg.batch_size, + num_workers=system_cfg.num_workers, + max_chunk_size=getattr(data_cfg, "max_chunk_size", 4), + macro_stride=getattr(data_cfg, "macro_stride", 24), + action_horizon=getattr(model_cfg, "action_horizon", 24), + state_horizon=getattr(model_cfg, "state_horizon", 1), + action_dim=getattr(model_cfg, "action_dim", 32), + max_state_dim=getattr(model_cfg, "max_state_dim", 64), + embodiment_tag=getattr(data_cfg, "embodiment_tag", "libero"), + embodiment_tag_mapping=getattr(data_cfg, "embodiment_tag_mapping", None), + image_size=tuple(getattr(data_cfg, "image_size", [176, 320])), + max_text_length=getattr(data_cfg, "max_text_length", 512), + num_views=getattr(model_cfg, "num_views", 2), + shuffle=system_cfg.shuffle, + distributed=dist.is_initialized(), + relative_action=getattr(data_cfg, "relative_action", True), + crop_ratio=getattr(data_cfg, "crop_ratio", 0.95), + color_jitter=getattr(data_cfg, "color_jitter", True), + training=True, + ) + return dataloader + + +def main(train_config: TrainConfig, seed: int = 42): + """Main training loop for DreamZero.""" + set_seed(seed) + + if not dist.is_initialized(): + dist.init_process_group(backend="nccl") + + rank = dist.get_rank() + world_size = dist.get_world_size() + local_rank = int(os.environ.get("LOCAL_RANK", 0)) + device = torch.device(f"cuda:{local_rank}") + torch.cuda.set_device(device) + + if rank == 0: + logger.info(f"Starting DreamZero training: world_size={world_size}") + + # Build model (loads on CPU) + policy = get_model(train_config) + logger.info(f"[Rank {rank}] Model loaded on CPU. Starting FSDP2 sharding...") + import sys; sys.stdout.flush() + + # AGENT_DEBUG: Weight dump disabled for performance (was writing 45GB to disk) + # from flagscale.train.models.dreamzero.dump_init_weights import dump_trainable_weights # AGENT_DEBUG + # dump_trainable_weights(policy, _dump_path, rank=rank) # AGENT_DEBUG + + # Apply FSDP2 BEFORE moving to GPU — model is too large for single GPU + # FSDP2 shards parameters across ranks, so each rank only holds 1/N of params on GPU + try: + device_mesh = init_device_mesh("cuda", (world_size,)) + logger.info(f"[Rank {rank}] Device mesh created. Calling apply_fsdp2...") + sys.stdout.flush() + policy = apply_fsdp2(policy, device_mesh) + logger.info(f"[Rank {rank}] FSDP2 applied successfully") + sys.stdout.flush() + except Exception as e: + logger.error(f"[Rank {rank}] FSDP2 FAILED: {type(e).__name__}: {e}") + import traceback; traceback.print_exc() + sys.stdout.flush() + raise + + # Move frozen encoders to GPU in bf16 (matching reference _initialize behavior at line 1367-1368) + # These are inference-only under autocast, so bf16 is correct and saves ~50% memory vs fp32. + # The DiT (ah.model) stays fp32 via FSDP2 param_dtype — matching DeepSpeed ZeRO-2 master weights. + ah = policy.action_head + ah.text_encoder.to(device=device, dtype=torch.bfloat16) + ah.image_encoder.to(device=device, dtype=torch.bfloat16) + ah.vae.to(device=device, dtype=torch.bfloat16) + # Move small projectors/encoders that live outside ah.model + for name, module in ah.named_children(): + if name not in ("model", "text_encoder", "image_encoder", "vae"): + module.to(device=device) + if rank == 0: + logger.info(f"Frozen encoders and projectors moved to {device}") + + # torch.compile: tested, gives <5% gain due to graph breaks in blockwise attention. + # Keeping code for reference but disabled by default. + if getattr(train_config.system, "use_torch_compile", False): + import torch._inductor.config as inductor_config + inductor_config.reorder_for_compute_comm_overlap = True + if rank == 0: + logger.info("Enabling torch.compile on DiT (mode=max-autotune-no-cudagraphs)") + logger.info("Enabled inductor reorder_for_compute_comm_overlap=True") + ah.model = torch.compile(ah.model, mode="max-autotune-no-cudagraphs") + if rank == 0: + logger.info("torch.compile applied to action_head.model") + + # RLinf-style optimization: compile attention sub-methods + use_reentrant=True + # This is much more effective than whole-model compile because: + # 1. The 4 attention methods have fixed shapes → CUDA graphs work + # 2. use_reentrant=True avoids pack/unpack hooks overhead in gradient checkpointing + if not getattr(train_config.system, "disable_attention_compile", False): + apply_compile_and_reentrant(policy) + else: + if rank == 0: + logger.info("Attention compile disabled by config") + + # Optimizer and scheduler + optimizer, scheduler = setup_optimizer_and_scheduler(policy, train_config) + + # Dataloader with reference transforms (each rank gets different micro-batch via IterableDataset sharding) + dataloader = get_dataloader(train_config, seed=seed) + data_iter = safe_cycle(dataloader) + + # Training params + train_steps = train_config.system.train_steps + log_freq = train_config.system.log_freq + grad_accum_steps = getattr(train_config.system, "gradient_accumulation_steps", 1) + grad_clip_norm = getattr(train_config.system, "grad_clip_norm", 1.0) + save_freq = getattr(train_config.system, "save_steps", 1000) + output_dir = Path(train_config.system.checkpoint.output_directory) + if rank == 0: + output_dir.mkdir(parents=True, exist_ok=True) + + # Metrics + loss_meter = AverageMeter("loss") + step_time_meter = AverageMeter("step_time") + + if rank == 0: + eff_batch = train_config.system.batch_size * world_size * grad_accum_steps + logger.info( + f"Config: steps={train_steps}, micro_batch={train_config.system.batch_size}, " + f"grad_accum={grad_accum_steps}, effective_batch={eff_batch}" + ) + + use_aligned = getattr(train_config.data, "use_aligned_dataloader", True) + if use_aligned: + from flagscale.train.datasets.dreamzero.aligned_dataloader import get_batch_aligned as get_batch + else: + from flagscale.train.datasets.dreamzero import get_batch + + # Barrier: ensure all ranks finished init before any forward pass + dist.barrier() + if rank == 0: + logger.info("All ranks synchronized — starting training loop") + + # Training loop + policy.train() + optimizer.zero_grad(set_to_none=True) + + # DebugHooks for Level 4 alignment (capture fwd/bwd intermediates on step 1) + _debug_hooks_enabled = os.environ.get("DREAMZERO_DEBUG_HOOKS", "") in ("1", "true", "True") + if _debug_hooks_enabled and rank == 0: + import sys + sys.path.insert(0, "/public-mixed/fengyupu/github/null-space") + from null_space.hamster.debug.hooks import DebugHooks + _hook_log_path = os.path.join(str(output_dir), "debug_hooks_step1.log") + _hook_log_file = open(_hook_log_path, "w") + def _hook_print(msg): + _hook_log_file.write(msg + "\n") + _hook_log_file.flush() + _debug_hooks = DebugHooks(policy.action_head.model, print_fn=_hook_print) + _debug_hooks.register() + logger.info(f"[DEBUG_HOOKS] Registered on action_head.model, logging to {_hook_log_path}") + + for step in range(1, train_steps + 1): + step_start = time.time() + accumulated_loss = 0.0 + # Signal to PyTorch that a new training iteration is starting. + # Helps manage CUDA graph memory from compiled attention sub-methods. + torch.compiler.cudagraph_mark_step_begin() + + for _ in range(grad_accum_steps): + raw_batch = next(data_iter) + batch = get_batch(raw_batch, device=device, compute_dtype=torch.bfloat16) + + # Debug: print batch shapes on first step + if step == 1 and rank == 0: + logger.info("[BATCH DEBUG] Batch shapes and stats:") + for k, v in batch.items(): + if torch.is_tensor(v): + logger.info(f" {k}: shape={tuple(v.shape)} dtype={v.dtype} min={v.min().item():.4f} max={v.max().item():.4f}") + else: + logger.info(f" {k}: type={type(v)} value={v}") + + with torch.amp.autocast("cuda", dtype=torch.bfloat16): + outputs = policy(batch) + loss = outputs["loss"] / grad_accum_steps + loss.backward() + accumulated_loss += loss.item() + # Track component losses for logging + if rank == 0: + dyn_loss = outputs.get("dynamics_loss", torch.tensor(0.0)).item() + act_loss = outputs.get("action_loss", torch.tensor(0.0)).item() + + # DebugHooks: remove after step 1 + if _debug_hooks_enabled and step == 1 and rank == 0: + _debug_hooks.remove() + _hook_log_file.close() + logger.info(f"[DEBUG_HOOKS] Step 1 complete, hooks removed. Log: {_hook_log_path}") + + if grad_clip_norm > 0: + total_norm = torch.nn.utils.clip_grad_norm_(policy.parameters(), grad_clip_norm) + else: + total_norm = 0.0 + + optimizer.step() + scheduler.step() + optimizer.zero_grad(set_to_none=True) + policy.set_frozen_modules_to_eval() + + step_time = time.time() - step_start + loss_meter.update(accumulated_loss) + step_time_meter.update(step_time) + + # Logging + if rank == 0 and step % log_freq == 0: + lr = scheduler.get_last_lr()[0] + grad_norm_val = total_norm.item() if torch.is_tensor(total_norm) else total_norm + logger.info( + f"step={step}/{train_steps} | loss={accumulated_loss:.4f} | " + f"dynamics={dyn_loss:.4f} | action={act_loss:.4f} | " + f"grad_norm={grad_norm_val:.4f} | " + f"lr={lr:.2e} | time={step_time:.2f}s/step" + ) + # JSON log for easy comparison parsing + import json as _json + _eff_batch = train_config.system.batch_size * world_size * grad_accum_steps + _samples_per_sec = _eff_batch / step_time + _log_entry = { + "step": step, "loss": accumulated_loss, + "dynamics_loss": dyn_loss, "action_loss": act_loss, + "grad_norm": grad_norm_val, "learning_rate": lr, + "step_time": round(step_time, 3), + "samples_per_sec": round(_samples_per_sec, 3), + "samples_per_sec_per_gpu": round(_samples_per_sec / world_size, 4), + } + _log_path = os.path.join(output_dir, "loss_log.jsonl") + with open(_log_path, "a") as _f: + _f.write(_json.dumps(_log_entry) + "\n") + loss_meter.reset() + step_time_meter.reset() + + # Checkpoint + if step % save_freq == 0 and rank == 0: + ckpt_dir = get_step_checkpoint_dir(output_dir, train_steps, step) + save_checkpoint( + checkpoint_dir=ckpt_dir, + step=step, + config=train_config, + policy=policy, + optimizer_state_dict=optimizer.state_dict(), + lr_scheduler=scheduler, + ) + update_last_checkpoint(ckpt_dir) + logger.info(f"Checkpoint saved: {ckpt_dir}") + + # Final checkpoint (disabled for alignment testing — FSDP2 save not yet supported) + # if rank == 0: + # ckpt_dir = get_step_checkpoint_dir(output_dir, train_steps, train_steps) + # save_checkpoint( + # checkpoint_dir=ckpt_dir, + # step=train_steps, + # config=train_config, + # policy=policy, + # optimizer_state_dict=optimizer.state_dict(), + # lr_scheduler=scheduler, + # ) + # update_last_checkpoint(ckpt_dir) + # logger.info(f"Training complete. Final checkpoint: {ckpt_dir}") + + if rank == 0: + logger.info(f"Training complete after {train_steps} steps (checkpoint save disabled).") + + dist.barrier() + dist.destroy_process_group() + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser( + description="Train DreamZero model. Called by FlagScale runner, not directly." + ) + parser.add_argument( + "--config-file", type=str, required=True, help="Path to the configuration YAML file" + ) + args = parser.parse_args() + + # Load config from Hydra-generated YAML + config = OmegaConf.load(args.config_file) + + logger.info(f"Full config: {OmegaConf.to_yaml(config)}") + + # Convert to Pydantic TrainConfig + train_config = TrainConfig.from_hydra_config(config) + + # Extract seed from experiment config + experiment_config = OmegaConf.to_container(config.experiment, resolve=True) + seed = experiment_config.get("seed", 42) + + logger.info(f"Experiment: {experiment_config}") + main(train_config, seed) diff --git a/flagscale/train/utils/optim_setup.py b/flagscale/train/utils/optim_setup.py index d8dfc02756..04f2580cf8 100644 --- a/flagscale/train/utils/optim_setup.py +++ b/flagscale/train/utils/optim_setup.py @@ -409,12 +409,18 @@ def setup_scheduler( if scheduler_config.name is None: raise ValueError("scheduler_config.name must be specified to use setup_scheduler") + # Resolve warmup_steps: warmup_ratio takes precedence if set + warmup_steps = scheduler_config.warmup_steps + if scheduler_config.warmup_ratio is not None: + warmup_steps = int(scheduler_config.warmup_ratio * num_training_steps) + logger.info(f"Scheduler: warmup_ratio={scheduler_config.warmup_ratio} → warmup_steps={warmup_steps} (of {num_training_steps} total)") + if scheduler_config.name == "cosine_decay_with_warmup": peak_lr = scheduler_config.peak_lr if peak_lr is None: peak_lr = optimizer.defaults.get("lr", optimizer.param_groups[0]["lr"]) config = CosineDecayWithWarmupSchedulerConfig( - num_warmup_steps=scheduler_config.warmup_steps, + num_warmup_steps=warmup_steps, num_decay_steps=scheduler_config.decay_steps, peak_lr=peak_lr, decay_lr=scheduler_config.decay_lr, @@ -424,7 +430,7 @@ def setup_scheduler( return get_scheduler( name=scheduler_config.name, optimizer=optimizer, - num_warmup_steps=scheduler_config.warmup_steps, + num_warmup_steps=warmup_steps, num_training_steps=num_training_steps, scheduler_specific_kwargs=scheduler_config.scheduler_kwargs, ) diff --git a/flagscale/train/utils/random_utils.py b/flagscale/train/utils/random_utils.py index 629098af2f..065a7d739b 100644 --- a/flagscale/train/utils/random_utils.py +++ b/flagscale/train/utils/random_utils.py @@ -29,7 +29,6 @@ cur_platform = get_platform() - def serialize_python_rng_state() -> dict[str, torch.Tensor]: """ Returns the rng state for `random` in the form of a flat dict[str, torch.Tensor] to be saved using