"""NimCritic inference runtime. Video DiT, value DiT, VAE, text encoder, and the sliding-window denoisers. This file and infer_nimcritic.py are the whole release. No package layout. """ from __future__ import annotations import logging def get_logger(name: str) -> logging.Logger: if not name.startswith("nimcritic"): name = f"nimcritic.{name}" if name else "nimcritic" return logging.getLogger(name) # ----- value representation (value_repr.py) ----- """NimCritic progress representation: absolute [0, 1] vs relative delta in [-1, 1]. Condition frames are always absolute [0, 1] for both progress and success. ``relative`` only changes the predicted progress channel (already in [-1, 1], no 2x-1 / denorm); success still uses the absolute latent mapping. """ import torch VALUE_PROGRESS_REPRS = ("absolute", "relative") def parse_value_progress_repr(raw: object, *, default: str = "absolute") -> str: value = default if raw is None else str(raw) if value not in VALUE_PROGRESS_REPRS: raise ValueError( f"value_progress_repr must be one of {VALUE_PROGRESS_REPRS}, got {value!r}" ) return value def value_to_latent(value: torch.Tensor, value_progress_repr: str) -> torch.Tensor: """Map adapter-space value to model latent space. * ``absolute``: both channels ``[0, 1] → [-1, 1]`` via ``2x-1``. * ``relative``: progress (col 0) stays as-is (condition is 0-1 absolute); success still ``2x-1``. """ repr_ = parse_value_progress_repr(value_progress_repr) if repr_ == "absolute": return 2.0 * value - 1.0 out = value.clone() if out.shape[-1] >= 2: out[..., 1:] = 2.0 * out[..., 1:] - 1.0 return out def value_from_latent(latents: torch.Tensor, value_progress_repr: str) -> torch.Tensor: """Map model latent to adapter-space (CPU float32). * ``absolute``: both channels ``[-1, 1] → [0, 1]``. * ``relative``: progress stays as-is (predicted steps are ``[-1, 1]`` delta); success still ``(x+1)/2``. """ repr_ = parse_value_progress_repr(value_progress_repr) out = latents.detach().to(device="cpu", dtype=torch.float32) if repr_ == "absolute": return (out + 1.0) / 2.0 if out.shape[-1] >= 2: out = out.clone() out[..., 1:] = (out[..., 1:] + 1.0) / 2.0 return out def apply_progress_delta( anchor: torch.Tensor, absolute: torch.Tensor, *, eps: float = 1e-6 ) -> torch.Tensor: """Training-side forward: absolute ``[0, 1]`` → relative delta ``[-1, 1]`` vs ``a``.""" a = anchor.clamp(0.0, 1.0) b = absolute.clamp(0.0, 1.0) ge = b >= a return torch.where( ge, (b - a) / (1.0 - a).clamp(min=eps), (b - a) / a.clamp(min=eps), ).clamp(-1.0, 1.0) def invert_progress_delta(anchor: torch.Tensor, delta: torch.Tensor) -> torch.Tensor: """Recover absolute progress in ``[0, 1]`` from relative delta in ``[-1, 1]``. ``d >= 0`` → ``a + d * (1 - a)``; ``d < 0`` → ``a * (1 + d)``. """ a = anchor.clamp(0.0, 1.0) d = delta.clamp(-1.0, 1.0) return torch.where(d >= 0, a + d * (1.0 - a), a * (1.0 + d)).clamp(0.0, 1.0) def decode_window_progress( window_value: torch.Tensor, gt_value: torch.Tensor, vcf_v: int, value_progress_repr: str, ) -> torch.Tensor: """Invert predicted progress against the window's first condition frame. ``window_value`` / ``gt_value`` are ``[B, H, D]``. Condition prefix stays absolute; only ``[:, vcf_v:, 0]`` is treated as delta. """ if parse_value_progress_repr(value_progress_repr) != "relative": return window_value if int(vcf_v) < 1: raise ValueError(f"value_condition_frames must be >= 1, got {vcf_v}") if window_value.shape[1] <= vcf_v: return window_value decoded = window_value.clone() anchor = gt_value[:, 0, 0] decoded[:, vcf_v:, 0] = invert_progress_delta(anchor[:, None], window_value[:, vcf_v:, 0]) return decoded # ----- gradient checkpoint (gradient.py) ----- # Vendored Wan2.2 / DiffSynth-Studio inference code: helpers/gradient.py # Local changes: none. import torch def create_custom_forward(module): def custom_forward(*inputs, **kwargs): return module(*inputs, **kwargs) return custom_forward def gradient_checkpoint_forward( model, use_gradient_checkpointing, *args, **kwargs, ): if use_gradient_checkpointing: model_output = torch.utils.checkpoint.checkpoint( create_custom_forward(model), *args, **kwargs, use_reentrant=False, ) else: model_output = model(*args, **kwargs) return model_output # ----- state dict converters (state_dict_converters.py) ----- # Vendored Wan2.2 / DiffSynth-Studio inference code: helpers/state_dict_converters.py # Local changes: none. def wan_video_vae_state_dict_converter(state_dict): converted = {} if "model_state" in state_dict: state_dict = state_dict["model_state"] for name, value in state_dict.items(): converted[f"model.{name}"] = value return converted def wan_video_dit_from_diffusers(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.attn2.add_k_proj.bias": "blocks.0.cross_attn.k_img.bias", "blocks.0.attn2.add_k_proj.weight": "blocks.0.cross_attn.k_img.weight", "blocks.0.attn2.add_v_proj.bias": "blocks.0.cross_attn.v_img.bias", "blocks.0.attn2.add_v_proj.weight": "blocks.0.cross_attn.v_img.weight", "blocks.0.attn2.norm_added_k.weight": "blocks.0.cross_attn.norm_k_img.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", "condition_embedder.image_embedder.ff.net.0.proj.bias": "img_emb.proj.1.bias", "condition_embedder.image_embedder.ff.net.0.proj.weight": "img_emb.proj.1.weight", "condition_embedder.image_embedder.ff.net.2.bias": "img_emb.proj.3.bias", "condition_embedder.image_embedder.ff.net.2.weight": "img_emb.proj.3.weight", "condition_embedder.image_embedder.norm1.bias": "img_emb.proj.0.bias", "condition_embedder.image_embedder.norm1.weight": "img_emb.proj.0.weight", "condition_embedder.image_embedder.norm2.bias": "img_emb.proj.4.bias", "condition_embedder.image_embedder.norm2.weight": "img_emb.proj.4.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", } converted = {} for name, value in state_dict.items(): if name in rename_dict: converted[rename_dict[name]] = value else: probe_name = ".".join(name.split(".")[:1] + ["0"] + name.split(".")[2:]) if probe_name in rename_dict: mapped = rename_dict[probe_name] mapped = ".".join(mapped.split(".")[:1] + [name.split(".")[1]] + mapped.split(".")[2:]) converted[mapped] = value return converted def wan_video_dit_state_dict_converter(state_dict): converted = {} for name, value in state_dict.items(): if name.startswith("vace"): continue if name.split(".")[0] in ["pose_patch_embedding", "face_adapter", "face_encoder", "motion_encoder"]: continue key = name[6:] if name.startswith("model.") else name converted[key] = value return converted # ----- checkpoint io (io.py) ----- # Vendored Wan2.2 / DiffSynth-Studio inference code: helpers/io.py # Local changes: none. import glob import hashlib import os from dataclasses import dataclass from typing import Dict, Optional, Union import torch from safetensors import safe_open @dataclass class ModelConfig: path: Union[str, list[str], None] = None model_id: Optional[str] = None origin_file_pattern: Union[str, list[str], None] = None download_source: Optional[str] = None local_model_path: Optional[str] = None skip_download: Optional[bool] = None state_dict: Optional[Dict[str, torch.Tensor]] = None def check_input(self): if self.path is None and self.model_id is None: raise ValueError("ModelConfig requires either `path` or (`model_id`, `origin_file_pattern`).") def parse_original_file_pattern(self): if self.origin_file_pattern in [None, "", "./"]: return "*" if isinstance(self.origin_file_pattern, list): return self.origin_file_pattern if self.origin_file_pattern.endswith("/"): return self.origin_file_pattern + "*" return self.origin_file_pattern def parse_download_source(self): if self.download_source is not None: return self.download_source env = os.environ.get("DIFFSYNTH_DOWNLOAD_SOURCE") return env if env is not None else "modelscope" def parse_skip_download(self): if self.skip_download is not None: return self.skip_download env = os.environ.get("DIFFSYNTH_SKIP_DOWNLOAD") if env is None: return False return env.lower() == "true" def reset_local_model_path(self): if os.environ.get("DIFFSYNTH_MODEL_BASE_PATH") is not None: self.local_model_path = os.environ.get("DIFFSYNTH_MODEL_BASE_PATH") elif self.local_model_path is None: self.local_model_path = "./checkpoints/" def require_downloading(self): if self.path is not None: return False origin_file_pattern = self.parse_original_file_pattern() local_root = os.path.join(self.local_model_path, self.model_id) if isinstance(origin_file_pattern, list): all_exist = True for pattern in origin_file_pattern: matches = glob.glob(os.path.join(local_root, pattern)) if len(matches) == 0: all_exist = False break if all_exist: return False else: if len(glob.glob(os.path.join(local_root, origin_file_pattern))) > 0: return False return not self.parse_skip_download() def download(self): origin_file_pattern = self.parse_original_file_pattern() root = os.path.join(self.local_model_path, self.model_id) downloaded_files = glob.glob(origin_file_pattern, root_dir=root) download_source = self.parse_download_source().lower() if download_source == "modelscope": from modelscope import snapshot_download snapshot_download( self.model_id, local_dir=root, allow_file_pattern=origin_file_pattern, ignore_file_pattern=downloaded_files, local_files_only=False, ) elif download_source == "huggingface": from huggingface_hub import snapshot_download as hf_snapshot_download hf_snapshot_download( self.model_id, local_dir=root, allow_patterns=origin_file_pattern, ignore_patterns=downloaded_files, local_files_only=False, ) else: raise ValueError("`download_source` should be `modelscope` or `huggingface`.") def download_if_necessary(self): self.check_input() self.reset_local_model_path() if self.require_downloading(): self.download() if self.path is None: if self.origin_file_pattern in [None, "", "./"]: self.path = os.path.join(self.local_model_path, self.model_id) else: matches = glob.glob(os.path.join(self.local_model_path, self.model_id, self.origin_file_pattern)) matches.sort() self.path = matches if isinstance(self.path, list) and len(self.path) == 1: self.path = self.path[0] def load_state_dict(file_path, torch_dtype=None, device="cpu"): if isinstance(file_path, list): state_dict = {} for file_path_ in file_path: state_dict.update(load_state_dict(file_path_, torch_dtype=torch_dtype, device=device)) return state_dict if file_path.endswith(".safetensors"): return load_state_dict_from_safetensors(file_path, torch_dtype=torch_dtype, device=device) 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=str(device)) as f: for key in f.keys(): value = f.get_tensor(key) if torch_dtype is not None: value = value.to(torch_dtype) state_dict[key] = value 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 len(state_dict) == 1: if "state_dict" in state_dict: state_dict = state_dict["state_dict"] elif "module" in state_dict: state_dict = state_dict["module"] elif "model_state" in state_dict: state_dict = state_dict["model_state"] if torch_dtype is not None: for key in state_dict: if isinstance(state_dict[key], torch.Tensor): state_dict[key] = state_dict[key].to(torch_dtype) return state_dict def _load_keys_dict_from_safetensors(file_path): keys_dict = {} with safe_open(file_path, framework="pt", device="cpu") as f: for key in f.keys(): keys_dict[key] = f.get_slice(key).get_shape() return keys_dict def _convert_state_dict_to_keys_dict(state_dict): keys_dict = {} for key, value in state_dict.items(): if isinstance(value, torch.Tensor): keys_dict[key] = list(value.shape) else: keys_dict[key] = _convert_state_dict_to_keys_dict(value) return keys_dict def _load_keys_dict_from_bin(file_path): state_dict = load_state_dict_from_bin(file_path) return _convert_state_dict_to_keys_dict(state_dict) def _load_keys_dict(file_path): if isinstance(file_path, list): merged = {} for path in file_path: merged.update(_load_keys_dict(path)) return merged if file_path.endswith(".safetensors"): return _load_keys_dict_from_safetensors(file_path) return _load_keys_dict_from_bin(file_path) def _convert_keys_dict_to_single_str(keys_dict, with_shape=True): keys = [] for key, value in keys_dict.items(): if isinstance(key, str): if isinstance(value, dict): keys.append(key + "|" + _convert_keys_dict_to_single_str(value, with_shape=with_shape)) else: if with_shape: shape = "_".join(map(str, list(value))) keys.append(key + ":" + shape) keys.append(key) keys.sort() return ",".join(keys) def hash_model_file(path, with_shape=True): keys_dict = _load_keys_dict(path) keys_str = _convert_keys_dict_to_single_str(keys_dict, with_shape=with_shape).encode("UTF-8") return hashlib.md5(keys_str).hexdigest() # ----- flow scheduler (scheduler_continuous.py) ----- # Vendored Wan2.2 / DiffSynth-Studio inference code: schedulers/scheduler_continuous.py # Local changes: none. import torch class WanContinuousFlowMatchScheduler: """Continuous-time Flow-Matching scheduler with shift-based sampling.""" def __init__(self, num_train_timesteps: int = 1000, shift: float = 5.0, eps: float = 1e-10): if num_train_timesteps <= 0: raise ValueError(f"`num_train_timesteps` must be positive, got {num_train_timesteps}") if shift <= 0: raise ValueError(f"`shift` must be positive, got {shift}") self.num_train_timesteps = int(num_train_timesteps) self.shift = float(shift) self.eps = float(eps) self._y_min, self._weight_norm_const = self._precompute_training_weight_stats() @staticmethod def _phi(u: torch.Tensor, shift: float) -> torch.Tensor: return shift * u / (1.0 + (shift - 1.0) * u) def _precompute_training_weight_stats(self) -> tuple[float, float]: steps = self.num_train_timesteps u_grid = torch.linspace(1.0, 0.0, steps + 1, dtype=torch.float64)[:-1] t_grid = self._phi(u_grid, self.shift) * float(steps) y_grid = torch.exp(-2.0 * ((t_grid - (steps / 2.0)) / steps) ** 2) y_min = float(y_grid.min().item()) y_shifted_grid = y_grid - y_min norm_const = float(y_shifted_grid.mean().item()) return y_min, norm_const def sample_training_t(self, batch_size: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor: if batch_size <= 0: raise ValueError(f"`batch_size` must be positive, got {batch_size}") u = torch.rand((batch_size,), device=device, dtype=torch.float32) sigma = self._phi(u, self.shift) timestep = sigma * float(self.num_train_timesteps) return timestep.to(dtype=dtype) def training_weight(self, timestep: torch.Tensor) -> torch.Tensor: t = timestep.to(dtype=torch.float32) steps = float(self.num_train_timesteps) y = torch.exp(-2.0 * ((t - (steps / 2.0)) / steps) ** 2) y_shifted = y - self._y_min weight = y_shifted / (self._weight_norm_const + self.eps) if weight.numel() == 1: return weight.reshape(()) return weight def add_noise(self, original_samples: torch.Tensor, noise: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor: sigma = (timestep / float(self.num_train_timesteps)).to( original_samples.device, dtype=original_samples.dtype ) if sigma.ndim == 0: return (1 - sigma) * original_samples + sigma * noise sigma = sigma.view(-1, *([1] * (original_samples.ndim - 1))) return (1 - sigma) * original_samples + sigma * noise @staticmethod def training_target(sample: torch.Tensor, noise: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor: del timestep return noise - sample def build_inference_schedule( self, num_inference_steps: int, device: torch.device, dtype: torch.dtype, shift_override: float | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: if num_inference_steps <= 0: raise ValueError(f"`num_inference_steps` must be positive, got {num_inference_steps}") shift = self.shift if shift_override is None else float(shift_override) if shift <= 0: raise ValueError(f"`shift` must be positive, got {shift}") u_steps = torch.linspace(1.0, 0.0, num_inference_steps + 1, device=device, dtype=torch.float32) sigma_steps = self._phi(u_steps, shift) timesteps = sigma_steps[:-1] * float(self.num_train_timesteps) deltas = sigma_steps[1:] - sigma_steps[:-1] return timesteps.to(dtype=dtype), deltas.to(dtype=dtype) @staticmethod def step(model_output: torch.Tensor, delta: torch.Tensor, sample: torch.Tensor) -> torch.Tensor: delta = delta.to(sample.device, dtype=sample.dtype) if delta.ndim == 0: return sample + model_output * delta delta = delta.view(-1, *([1] * (sample.ndim - 1))) return sample + model_output * delta # ----- video DiT (wan_video_dit.py) ----- # Vendored Wan2.2 / DiffSynth-Studio inference code: wan_video_dit.py # Local change: logging import -> nimcritic.utils.logging. import torch import torch.nn as nn import torch.nn.functional as F import math from typing import Any, Dict, Tuple, Optional from einops import rearrange logger = get_logger(__name__) def flash_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, num_heads: int, ctx_mask: Optional[torch.Tensor] = None, compatibility_mode=True): if compatibility_mode: 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, attn_mask=ctx_mask) x = rearrange(x, "b n s d -> b s (n d)", n=num_heads) return x else: raise NotImplementedError("Only compatibility mode is implemented for flash attention. Please set compatibility_mode=True.") def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor): return (x * (1 + scale) + shift) 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) def precompute_freqs_cis_3d(dim: int, end: int = 1024, theta: float = 10000.0): # 3d rope precompute f_freqs_cis = precompute_freqs_cis(dim - 2 * (dim // 3), end, theta) h_freqs_cis = precompute_freqs_cis(dim // 3, end, theta) w_freqs_cis = precompute_freqs_cis(dim // 3, end, theta) return f_freqs_cis, h_freqs_cis, w_freqs_cis def precompute_freqs_cis(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, device=freqs.device), freqs) freqs_cis = torch.polar(torch.ones_like(freqs), freqs) # complex64 return freqs_cis def rope_apply(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(torch.complex64) if freqs.device.type == "npu" else freqs x_out = torch.view_as_real(x_out * freqs).flatten(2) return x_out.to(x.dtype) def create_group_causal_attn_mask( num_temporal_groups: int, num_query_per_group: int, num_key_per_group: int, mode: str = "causal" ) -> torch.Tensor: """ Creates a group-based attention mask for scaled dot-product attention with two modes: 'causal' and 'group_diagonal'. Parameters: - num_temporal_groups (int): The number of temporal groups (e.g., frames in a video sequence). - num_query_per_group (int): The number of query tokens per temporal group. (e.g., latent tokens in a frame, H x W). - num_key_per_group (int): The number of key tokens per temporal group. (e.g., action tokens per frame). - mode (str): The mode of the attention mask. Options are: - 'causal': Query tokens can attend to key tokens from the same or previous temporal groups. - 'group_diagonal': Query tokens can attend only to key tokens from the same temporal group. Returns: - attn_mask (torch.Tensor): A boolean tensor of shape (L, S), where: - L = num_temporal_groups * num_query_per_group (total number of query tokens) - S = num_temporal_groups * num_key_per_group (total number of key tokens) The mask indicates where attention is allowed (True) and disallowed (False). Example: Input: num_temporal_groups = 3 num_query_per_group = 4 num_key_per_group = 2 Output: Causal Mask Shape: torch.Size([12, 6]) Group Diagonal Mask Shape: torch.Size([12, 6]) if mode='causal': tensor([[ True, True, False, False, False, False], [ True, True, False, False, False, False], [ True, True, False, False, False, False], [ True, True, False, False, False, False], [ True, True, True, True, False, False], [ True, True, True, True, False, False], [ True, True, True, True, False, False], [ True, True, True, True, False, False], [ True, True, True, True, True, True], [ True, True, True, True, True, True], [ True, True, True, True, True, True], [ True, True, True, True, True, True]]) if mode='group_diagonal': tensor([[ True, True, False, False, False, False], [ True, True, False, False, False, False], [ True, True, False, False, False, False], [ True, True, False, False, False, False], [False, False, True, True, False, False], [False, False, True, True, False, False], [False, False, True, True, False, False], [False, False, True, True, False, False], [False, False, False, False, True, True], [False, False, False, False, True, True], [False, False, False, False, True, True], [False, False, False, False, True, True]]) """ assert mode in ["causal", "group_diagonal"], f"Mode {mode} must be 'causal' or 'group_diagonal'" # Total number of query and key tokens total_num_query_tokens = num_temporal_groups * num_query_per_group # Total number of query tokens (L) total_num_key_tokens = num_temporal_groups * num_key_per_group # Total number of key tokens (S) # Generate time indices for query and key tokens (shape: [L] and [S]) query_time_indices = torch.arange(num_temporal_groups).repeat_interleave(num_query_per_group) # Shape: [L] key_time_indices = torch.arange(num_temporal_groups).repeat_interleave(num_key_per_group) # Shape: [S] # Expand dimensions to compute outer comparison query_time_indices = query_time_indices.unsqueeze(1) # Shape: [L, 1] key_time_indices = key_time_indices.unsqueeze(0) # Shape: [1, S] if mode == "causal": # Causal Mode: Query can attend to keys where key_time <= query_time attn_mask = query_time_indices >= key_time_indices # Shape: [L, S] elif mode == "group_diagonal": # Group Diagonal Mode: Query can attend only to keys where key_time == query_time attn_mask = query_time_indices == key_time_indices # Shape: [L, S] assert attn_mask.shape == (total_num_query_tokens, total_num_key_tokens), "Attention mask shape mismatch" return attn_mask 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 class AttentionModule(nn.Module): def __init__(self, num_heads): super().__init__() self.num_heads = num_heads def forward(self, q, k, v, ctx_mask=None): x = flash_attention(q=q, k=k, v=v, num_heads=self.num_heads, ctx_mask=ctx_mask) return x class SelfAttention(nn.Module): def __init__(self, hidden_dim: int, attn_head_dim: int, num_heads: int, eps: float = 1e-6): super().__init__() self.hidden_dim = hidden_dim self.num_heads = num_heads self.attn_head_dim = attn_head_dim self.attn_hidden_dim = self.num_heads * self.attn_head_dim self.q = nn.Linear(hidden_dim, self.attn_hidden_dim) self.k = nn.Linear(hidden_dim, self.attn_hidden_dim) self.v = nn.Linear(hidden_dim, self.attn_hidden_dim) self.o = nn.Linear(self.attn_hidden_dim, hidden_dim) self.norm_q = RMSNorm(self.attn_hidden_dim, eps=eps) self.norm_k = RMSNorm(self.attn_hidden_dim, eps=eps) # self.attn = AttentionModule(self.num_heads) def forward(self, x, freqs, self_attn_mask: Optional[torch.Tensor] = None): 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 = flash_attention(q=q, k=k, v=v, num_heads=self.num_heads, ctx_mask=self_attn_mask) return self.o(x) class CrossAttention(nn.Module): def __init__(self, hidden_dim: int, attn_head_dim: int, num_heads: int, eps: float = 1e-6,): super().__init__() self.hidden_dim = hidden_dim self.num_heads = num_heads self.attn_head_dim = attn_head_dim self.attn_hidden_dim = self.num_heads * self.attn_head_dim self.q = nn.Linear(hidden_dim, self.attn_hidden_dim) self.k = nn.Linear(hidden_dim, self.attn_hidden_dim) self.v = nn.Linear(hidden_dim, self.attn_hidden_dim) self.o = nn.Linear(self.attn_hidden_dim, hidden_dim) self.norm_q = RMSNorm(self.attn_hidden_dim, eps=eps) self.norm_k = RMSNorm(self.attn_hidden_dim, eps=eps) # self.attn = AttentionModule(self.num_heads) def forward(self, x: torch.Tensor, ctx: torch.Tensor, ctx_mask: Optional[torch.Tensor] = None): q = self.norm_q(self.q(x)) k = self.norm_k(self.k(ctx)) v = self.v(ctx) x = flash_attention(q=q, k=k, v=v, num_heads=self.num_heads, ctx_mask=ctx_mask) return self.o(x) class GateModule(nn.Module): def __init__(self,): super().__init__() def forward(self, x, gate, residual): return x + gate * residual class DiTBlock(nn.Module): def __init__(self, hidden_dim: int, attn_head_dim: int, num_heads: int, ffn_dim: int, eps: float = 1e-6): super().__init__() self.hidden_dim = hidden_dim self.attn_head_dim = attn_head_dim self.num_heads = num_heads self.ffn_dim = ffn_dim self.self_attn = SelfAttention(hidden_dim, attn_head_dim, num_heads, eps) self.cross_attn = CrossAttention( hidden_dim, attn_head_dim, num_heads, eps) self.norm1 = nn.LayerNorm(hidden_dim, eps=eps, elementwise_affine=False) self.norm2 = nn.LayerNorm(hidden_dim, eps=eps, elementwise_affine=False) self.norm3 = nn.LayerNorm(hidden_dim, eps=eps) self.ffn = nn.Sequential(nn.Linear(hidden_dim, ffn_dim), nn.GELU( approximate='tanh'), nn.Linear(ffn_dim, hidden_dim)) self.modulation = nn.Parameter(torch.randn(1, 6, hidden_dim) / hidden_dim**0.5) self.gate = GateModule() def forward(self, x, context, t_mod, freqs, context_mask=None, self_attn_mask: Optional[torch.Tensor] = None): if context_mask is not None and context_mask.dim() == 3: context_mask = context_mask.unsqueeze(1) # (B, 1, seq_len, context_len), 1 for heads has_seq = len(t_mod.shape) == 4 chunk_dim = 2 if has_seq else 1 # 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=chunk_dim) if has_seq: # means t_mod has separate modulation for each token, otherwise same modulation for all tokens in the block shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( shift_msa.squeeze(2), scale_msa.squeeze(2), gate_msa.squeeze(2), shift_mlp.squeeze(2), scale_mlp.squeeze(2), gate_mlp.squeeze(2), ) input_x = modulate(self.norm1(x), shift_msa, scale_msa) x = self.gate(x, gate_msa, self.self_attn(input_x, freqs, self_attn_mask=self_attn_mask)) x = x + self.cross_attn(self.norm3(x), context, ctx_mask=context_mask) 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): if len(t_mod.shape) == 3: shift, scale = (self.modulation.unsqueeze(0).to(dtype=t_mod.dtype, device=t_mod.device) + t_mod.unsqueeze(2)).chunk(2, dim=2) x = (self.head(self.norm(x) * (1 + scale.squeeze(2)) + shift.squeeze(2))) else: shift, scale = (self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) + t_mod).chunk(2, dim=1) x = (self.head(self.norm(x) * (1 + scale) + shift)) return x class WanVideoDiT(torch.nn.Module): def __init__( self, hidden_dim: int, in_dim: int, ffn_dim: int, out_dim: int, text_dim: int, freq_dim: int, eps: float, patch_size: Tuple[int, int, int], num_heads: int, attn_head_dim: int, num_layers: int, has_image_input: bool, has_image_pos_emb: bool = False, has_ref_conv: bool = False, add_control_adapter: bool = False, in_dim_control_adapter: int = 24, seperated_timestep: bool = False, require_vae_embedding: bool = False, require_clip_embedding: bool = False, fuse_vae_embedding_in_latents: bool = True, action_conditioned: bool = False, action_dim: int = 7, action_group_causal_mask_mode = "causal", video_attention_mask_mode: str = "bidirectional", use_gradient_checkpointing: bool = False, ): super().__init__() self.hidden_dim = hidden_dim self.in_dim = in_dim self.freq_dim = freq_dim self.patch_size = patch_size self.num_heads = num_heads self.attn_head_dim = attn_head_dim self.seperated_timestep = seperated_timestep self.require_vae_embedding = require_vae_embedding self.require_clip_embedding = require_clip_embedding self.fuse_vae_embedding_in_latents = fuse_vae_embedding_in_latents self.video_attention_mask_mode = str(video_attention_mask_mode) if num_heads <= 0: raise ValueError(f"`num_heads` must be > 0, got {num_heads}") if attn_head_dim <= 0: raise ValueError(f"`attn_head_dim` must be > 0, got {attn_head_dim}") if attn_head_dim % 2 != 0: raise ValueError( f"`attn_head_dim` must be even for RoPE, got {attn_head_dim}" ) self.action_conditioned = action_conditioned self.action_dim = action_dim assert has_image_input == False assert require_clip_embedding == False assert require_vae_embedding == False and fuse_vae_embedding_in_latents == True, "Only support fusing vae embedding in latents" self.patch_embedding = nn.Conv3d( in_dim, hidden_dim, kernel_size=patch_size, stride=patch_size) self.text_embedding = nn.Sequential( nn.Linear(text_dim, hidden_dim), nn.GELU(approximate='tanh'), nn.Linear(hidden_dim, hidden_dim) ) self.time_embedding = nn.Sequential( nn.Linear(freq_dim, hidden_dim), nn.SiLU(), nn.Linear(hidden_dim, hidden_dim) ) self.time_projection = nn.Sequential( nn.SiLU(), nn.Linear(hidden_dim, hidden_dim * 6)) self.blocks = nn.ModuleList([ DiTBlock(hidden_dim, attn_head_dim, num_heads, ffn_dim, eps) for _ in range(num_layers) ]) self.head = Head(hidden_dim, out_dim, patch_size, eps) self.freqs = precompute_freqs_cis_3d(attn_head_dim) if has_ref_conv: self.ref_conv = nn.Conv2d(16, hidden_dim, kernel_size=(2, 2), stride=(2, 2)) self.has_image_pos_emb = has_image_pos_emb self.has_ref_conv = has_ref_conv self.control_adapter = None if self.action_conditioned: self.action_embedding = nn.Linear(action_dim, hidden_dim) self.action_group_causal_mask_mode = action_group_causal_mask_mode self.use_gradient_checkpointing = use_gradient_checkpointing if self.use_gradient_checkpointing: logger.info("Using gradient checkpointing for DiT blocks. This will save memory but use more computation.") def patchify(self, x: torch.Tensor, control_camera_latents_input: Optional[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) return x 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 _validate_forward_inputs( self, x: torch.Tensor, timestep: torch.Tensor, context: torch.Tensor, context_mask: Optional[torch.Tensor], action: Optional[torch.Tensor], ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: if x.ndim != 5: raise ValueError(f"`latents` must be 5D [B, C, T, H, W], got shape {tuple(x.shape)}") num_latent_frames = x.shape[2] if context.ndim != 3: raise ValueError(f"`context` must be 3D [B, L, D], got shape {tuple(context.shape)}") if timestep.ndim != 1: raise ValueError(f"`timestep` must be 1D [B] or [1], got shape {tuple(timestep.shape)}") if self.action_conditioned: allow_text_only_single_frame = (num_latent_frames == 1 and action is None) if not allow_text_only_single_frame: assert action is not None, "Action input is required for action-conditioned model." if action.ndim != 3: raise ValueError(f"`action` must be 3D [B, action_horizon, action_dim], got shape {tuple(action.shape)}") if action.shape[2] != self.action_dim: raise ValueError(f"`action` last dimension must be {self.action_dim}, got {action.shape[2]}") if num_latent_frames <= 1: raise ValueError(f"video length must be > 1 for action-conditioned model, got {num_latent_frames}") if action.shape[1] % (num_latent_frames - 1) != 0: raise ValueError( f"action horizon must be divisible by (num_latent_frames - 1), got action_horizon={action.shape[1]}" ) if context_mask is None: context_mask = torch.ones((context.shape[0], context.shape[1]), dtype=torch.bool, device=context.device) else: if context_mask.ndim != 2: raise ValueError(f"`context_mask` must be 2D [B, L], got shape {tuple(context_mask.shape)}") if context_mask.shape[0] != context.shape[0] or context_mask.shape[1] != context.shape[1]: raise ValueError(f"`context_mask` shape must match `context` shape [B, L], got {tuple(context_mask.shape)} vs {tuple(context.shape)}") batch_size = x.shape[0] if batch_size != context.shape[0]: if not self.training and batch_size == 1: x = x.expand(context.shape[0], -1, -1, -1, -1) batch_size = context.shape[0] else: raise ValueError( f"Batch mismatch between latents and context: {batch_size} vs {context.shape[0]}." ) if timestep.shape[0] not in (1, batch_size): raise ValueError( f"`timestep` length must be 1 or batch_size({batch_size}), got {timestep.shape[0]}" ) if timestep.shape[0] == 1 and batch_size > 1: assert not self.training, "During training, timestep length must match batch_size." timestep = timestep.expand(batch_size) return x, timestep, context_mask def build_video_to_video_mask( self, video_seq_len: int, video_tokens_per_frame: int, device: torch.device, condition_frames: int = 1, ) -> torch.Tensor: if video_seq_len <= 0: raise ValueError(f"`video_seq_len` must be positive, got {video_seq_len}") if video_tokens_per_frame <= 0: raise ValueError(f"`video_tokens_per_frame` must be positive, got {video_tokens_per_frame}") if condition_frames < 1: raise ValueError(f"`condition_frames` must be >= 1, got {condition_frames}") if self.video_attention_mask_mode == "bidirectional": return torch.ones((video_seq_len, video_seq_len), dtype=torch.bool, device=device) if self.video_attention_mask_mode == "per_frame_causal": if video_seq_len % video_tokens_per_frame != 0: raise ValueError( "`video_seq_len` must be divisible by `video_tokens_per_frame` in `per_frame_causal` mode, " f"got {video_seq_len} and {video_tokens_per_frame}" ) num_video_frames = video_seq_len // video_tokens_per_frame frame_causal = torch.tril( torch.ones((num_video_frames, num_video_frames), dtype=torch.bool, device=device) ) return frame_causal.repeat_interleave(video_tokens_per_frame, dim=0).repeat_interleave( video_tokens_per_frame, dim=1 ) if self.video_attention_mask_mode == "first_frame_causal": # First `condition_frames` frames are condition (visible to all video tokens); # remaining frames are causal among themselves and can attend to condition frames. video_mask = torch.ones((video_seq_len, video_seq_len), dtype=torch.bool, device=device) cond_tokens = min(condition_frames * video_tokens_per_frame, video_seq_len) # Condition frames must not attend to non-condition (future) frames. video_mask[:cond_tokens, cond_tokens:] = False return video_mask raise ValueError(f"Unsupported video attention mask mode: {self.video_attention_mask_mode}") def pre_dit( self, x: torch.Tensor, timestep: torch.Tensor, context: torch.Tensor, context_mask: Optional[torch.Tensor] = None, action: Optional[torch.Tensor] = None, fuse_vae_embedding_in_latents: bool = False, control_camera_latents_input: Optional[torch.Tensor] = None, condition_frames: int = 1, ) -> Dict[str, Any]: x, timestep, context_mask = self._validate_forward_inputs( x=x, timestep=timestep, context=context, context_mask=context_mask, action=action, ) batch_size = x.shape[0] patch_h = int(self.patch_size[1]) patch_w = int(self.patch_size[2]) if x.shape[3] % patch_h != 0 or x.shape[4] % patch_w != 0: raise ValueError( "Latent spatial shape must be divisible by DiT patch size, " f"got HxW=({x.shape[3]}, {x.shape[4]}), patch=({patch_h}, {patch_w})" ) tokens_per_frame = (x.shape[3] // patch_h) * (x.shape[4] // patch_w) if self.seperated_timestep and fuse_vae_embedding_in_latents: if not hasattr(self, "patch_size") or len(self.patch_size) < 3: raise ValueError(f"Invalid dit.patch_size: {getattr(self, 'patch_size', None)}") token_timesteps = torch.ones( (batch_size, x.shape[2], tokens_per_frame), dtype=timestep.dtype, device=timestep.device, ) * timestep.view(batch_size, 1, 1) # First `condition_frames` latent frames are GT condition (clean, t=0); # remaining frames use the sampled timestep. vcf = max(0, int(condition_frames)) if vcf > 0: if vcf > x.shape[2]: raise ValueError( f"`condition_frames`={vcf} exceeds latent T={x.shape[2]}." ) token_timesteps[:, :vcf, :] = 0 token_timesteps = token_timesteps.reshape(batch_size, -1) token_t_emb = sinusoidal_embedding_1d(self.freq_dim, token_timesteps.reshape(-1)) t = self.time_embedding(token_t_emb).reshape(batch_size, -1, self.hidden_dim) t_mod = self.time_projection(t).unflatten(2, (6, self.hidden_dim)) else: raise NotImplementedError("Only support seperated_timestep with fuse_vae_embedding_in_latents for now.") t = self.time_embedding(sinusoidal_embedding_1d(self.freq_dim, timestep)) t_mod = self.time_projection(t).unflatten(1, (6, self.hidden_dim)) x = self.patchify(x, control_camera_latents_input=control_camera_latents_input) f, h, w = x.shape[2:] context = self.text_embedding(context) # (B, L, dim) context_len = context.shape[1] if self.action_conditioned and action is not None: action_len = action.shape[1] action_emb = self.action_embedding(action) # (B, action_len, dim) action_pos_embed = sinusoidal_embedding_1d(self.hidden_dim, torch.arange(action_len, device=action_emb.device)) # (action_len, dim) action_emb = action_emb + action_pos_embed.unsqueeze(0) # (B, action_len, dim) context = torch.cat([context, action_emb], dim=1) # (B, context_len + action_len, dim) # new mask num_temporal_groups = f - 1 # first latent frame do not attend to actions if num_temporal_groups <= 0: raise ValueError( "Action-conditioned context mask requires at least 2 latent frames when `action` is provided." ) assert action_emb.shape[1] % num_temporal_groups == 0, \ f"Action embedding length {action_emb.shape[1]} must be divisible by number of temporal groups {num_temporal_groups}" # Each latent frame (from the 2nd one) attends to the corresponding group of action tokens action_group_mask = create_group_causal_attn_mask( num_temporal_groups=num_temporal_groups, num_query_per_group=tokens_per_frame, num_key_per_group=action_len // num_temporal_groups, mode=self.action_group_causal_mask_mode, ).to(context.device) # ((f-1)*tokens_per_frame, action_len) seq_len = f * h * w # query length final_context_mask = torch.zeros((batch_size, seq_len, context.shape[1]), dtype=torch.bool, device=context.device) # (B, seq_len, L + action_len) # all latent frames attend to text tokens final_context_mask[:, :, :context_len] = context_mask.unsqueeze(1).expand(-1, seq_len, -1) # (B, seq_len, L) # latent frames from the 2nd one attend to action tokens final_context_mask[:, tokens_per_frame:, context_len:] = action_group_mask.unsqueeze(0).expand(batch_size, -1, -1) # (B, seq_len, action_len) context_mask = final_context_mask elif self.action_conditioned and action is None: if f != 1: raise ValueError( "Action-conditioned model requires `action` unless running single-frame text-only mode with num_latent_frames=1." ) context_mask = context_mask.unsqueeze(1).expand(-1, f * h * w, -1) # (B, seq_len, L) else: context_mask = context_mask.unsqueeze(1).expand(-1, f * h * w, -1) # (B, seq_len, L) x_tokens = rearrange(x, "b c f h w -> b (f h w) c").contiguous() freqs = torch.cat([ self.freqs[0][: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).to(x_tokens.device) return { "tokens": x_tokens, "freqs": freqs, "t": t, "t_mod": t_mod, "context": context, "context_mask": context_mask, "meta": { "grid_size": (f, h, w), "tokens_per_frame": tokens_per_frame, "batch_size": batch_size, }, } def post_dit(self, x_tokens: torch.Tensor, pre_state: Dict[str, Any]) -> torch.Tensor: f, h, w = pre_state["meta"]["grid_size"] x = self.head(x_tokens, pre_state["t"]) x = self.unpatchify(x, (f, h, w)) return x def forward( self, x: torch.Tensor, timestep: torch.Tensor, context: torch.Tensor, context_mask: Optional[torch.Tensor] = None, action: Optional[torch.Tensor] = None, fuse_vae_embedding_in_latents: bool = False, ): pre_state = self.pre_dit( x=x, timestep=timestep, context=context, context_mask=context_mask, action=action, fuse_vae_embedding_in_latents=fuse_vae_embedding_in_latents, ) x_tokens = pre_state["tokens"] context_emb = pre_state["context"] t_mod = pre_state["t_mod"] freqs = pre_state["freqs"] context_attn_mask = pre_state["context_mask"] self_attn_mask = self.build_video_to_video_mask( video_seq_len=x_tokens.shape[1], video_tokens_per_frame=int(pre_state["meta"]["tokens_per_frame"]), device=x_tokens.device, ) if self.video_attention_mask_mode != "bidirectional" else None # special rule for faster speed for block in self.blocks: if self.use_gradient_checkpointing: x_tokens = gradient_checkpoint_forward( block, self.use_gradient_checkpointing, x_tokens, context_emb, t_mod, freqs, context_mask=context_attn_mask, self_attn_mask=self_attn_mask ) else: x_tokens = block(x_tokens, context_emb, t_mod, freqs, context_mask=context_attn_mask, self_attn_mask=self_attn_mask) return self.post_dit(x_tokens, pre_state) # ----- VAE (wan_video_vae.py) ----- # Vendored Wan2.2 / DiffSynth-Studio inference code: wan_video_vae.py # Local changes: none. 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, _, device = *x.size(), x.device assert s % block_size == 0 num_blocks = s // block_size # build mask mask = torch.zeros(b, n, s, s, dtype=torch.bool, device=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: cache_x = cache_x.to(x.device) 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).to( cache_x.device), 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).to(cache_x.device), 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, feat_cache, feat_idx 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).to( cache_x.device), 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, feat_cache, feat_idx 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, feat_cache, feat_idx = module(x, feat_cache, feat_idx) return x + self.avg_shortcut(x_copy), feat_cache, feat_idx 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, feat_cache, feat_idx = module(x_main, feat_cache, feat_idx) if self.avg_shortcut is not None: x_shortcut = self.avg_shortcut(x, first_chunk) x_main = x_main + x_shortcut return x_main, feat_cache, feat_idx 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).to( cache_x.device), 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, feat_cache, feat_idx = 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, feat_cache, feat_idx = 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).to( cache_x.device), 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, feat_cache, feat_idx 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).to(cache_x.device), 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, feat_cache, feat_idx = 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, feat_cache, feat_idx = 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) .to(cache_x.device), 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, feat_cache, feat_idx 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).to( cache_x.device), 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, feat_cache, feat_idx = layer(x, feat_cache, feat_idx) else: x = layer(x) ## upsamples for layer in self.upsamples: if feat_cache is not None: x, feat_cache, feat_idx = 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).to( cache_x.device), 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, feat_cache, feat_idx 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).to(cache_x.device), 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, feat_cache, feat_idx = layer(x, feat_cache, feat_idx) else: x = layer(x) ## upsamples for layer in self.upsamples: if feat_cache is not None: x, feat_cache, feat_idx = 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) .to(cache_x.device), 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, feat_cache, feat_idx 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) def forward(self, x): mu, log_var = self.encode(x) z = self.reparameterize(mu, log_var) x_recon = self.decode(z) return x_recon, mu, log_var def encode(self, x, scale): self.clear_cache() ## cache t = x.shape[2] iter_ = 1 + (t - 1) // 4 for i in range(iter_): self._enc_conv_idx = [0] if i == 0: out, self._enc_feat_map, self._enc_conv_idx = self.encoder(x[:, :, :1, :, :], feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx) else: out_, self._enc_feat_map, self._enc_conv_idx = self.encoder(x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :], feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx) out = torch.cat([out, out_], 2) mu, log_var = 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): self.clear_cache() # 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) for i in range(iter_): self._conv_idx = [0] if i == 0: out, self._feat_map, self._conv_idx = self.decoder(x[:, :, i:i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx) else: out_, self._feat_map, self._conv_idx = self.decoder(x[:, :, i:i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx) out = torch.cat([out, out_], 2) # may add tensor offload 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) def clear_cache(self): self._conv_num = count_conv3d(self.decoder) self._conv_idx = [0] self._feat_map = [None] * self._conv_num # cache encode self._enc_conv_num = count_conv3d(self.encoder) self._enc_conv_idx = [0] self._enc_feat_map = [None] * self._enc_conv_num class WanVideoVAE(nn.Module): def __init__(self, z_dim=16): 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) self.std = torch.tensor(std) 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.temporal_downsample_factor = 4 self.z_dim = z_dim def build_1d_mask(self, length, left_bound, right_bound, border_width): x = torch.ones((length,)) if not left_bound: x[:border_width] = (torch.arange(border_width) + 1) / border_width if not right_bound: x[-border_width:] = torch.flip((torch.arange(border_width) + 1) / 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]) w = self.build_1d_mask(W, is_bound[2], is_bound[3], border_width[1]) 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, device, 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_)) data_device = "cpu" computation_device = device 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=data_device) values = torch.zeros((1, 3, out_T, H * self.upsampling_factor, W * self.upsampling_factor), dtype=hidden_states.dtype, device=data_device) for h, h_, w, w_ in tqdm(tasks, desc="VAE decoding"): hidden_states_batch = hidden_states[:, :, :, h:h_, w:w_].to(computation_device) hidden_states_batch = self.model.decode(hidden_states_batch, self.scale).to(data_device) 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, device=data_device) 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, device, 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_)) data_device = "cpu" computation_device = device out_T = (T + 3) // 4 weight = torch.zeros((1, 1, out_T, H // self.upsampling_factor, W // self.upsampling_factor), dtype=video.dtype, device=data_device) values = torch.zeros((1, self.z_dim, out_T, H // self.upsampling_factor, W // self.upsampling_factor), dtype=video.dtype, device=data_device) for h, h_, w, w_ in tqdm(tasks, desc="VAE encoding"): hidden_states_batch = video[:, :, :, h:h_, w:w_].to(computation_device) hidden_states_batch = self.model.encode(hidden_states_batch, self.scale).to(data_device) 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, device=data_device) 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, device): video = video.to(device) x = self.model.encode(video, self.scale) return x def single_decode(self, hidden_state, device): hidden_state = hidden_state.to(device) video = self.model.decode(hidden_state, self.scale) return video.clamp_(-1, 1) def encode(self, videos, device, tiled=False, tile_size=(34, 34), tile_stride=(18, 16)): # videos = [video.to("cpu") for video in videos] hidden_states = [] for video in videos: video = video.unsqueeze(0) if tiled: raise NotImplementedError("Tiled encoding is not allowed yet.") 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, device, tile_size, tile_stride) else: hidden_state = self.single_encode(video, device) 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, device, tiled=False, tile_size=(34, 34), tile_stride=(18, 16)): hidden_states = [hidden_state.to("cpu") for hidden_state in hidden_states] videos = [] for hidden_state in hidden_states: hidden_state = hidden_state.unsqueeze(0) if tiled: video = self.tiled_decode(hidden_state, device, tile_size, tile_stride) else: video = self.single_decode(hidden_state, device) video = video.squeeze(0) videos.append(video) videos = torch.stack(videos) return videos @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) def encode(self, x, scale): self.clear_cache() x = patchify(x, patch_size=2) t = x.shape[2] iter_ = 1 + (t - 1) // 4 for i in range(iter_): self._enc_conv_idx = [0] if i == 0: out, self._enc_feat_map, self._enc_conv_idx = self.encoder(x[:, :, :1, :, :], feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx) else: out_, self._enc_feat_map, self._enc_conv_idx = self.encoder(x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :], feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx) out = torch.cat([out, out_], 2) mu, log_var = 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] self.clear_cache() return mu def decode(self, z, scale): self.clear_cache() 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) for i in range(iter_): self._conv_idx = [0] if i == 0: out, self._feat_map, self._conv_idx = self.decoder(x[:, :, i:i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx, first_chunk=True) else: out_, self._feat_map, self._conv_idx = self.decoder(x[:, :, i:i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx) out = torch.cat([out, out_], 2) out = unpatchify(out, patch_size=2) self.clear_cache() return out class WanVideoVAE38(WanVideoVAE): def __init__(self, z_dim=48, dim=160): 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) self.std = torch.tensor(std) 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.temporal_downsample_factor = 4 self.z_dim = z_dim # ----- text encoder (wan_video_text_encoder.py) ----- # Vendored Wan2.2 / DiffSynth-Studio inference code: wan_video_text_encoder.py # Local changes: none. import math import torch import torch.nn as nn import torch.nn.functional as F from transformers import AutoTokenizer import html import string try: import ftfy except ModuleNotFoundError: class _FTFY: @staticmethod def fix_text(text): return text ftfy = _FTFY() try: import regex as re except ModuleNotFoundError: import re 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))) 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=256384, dim=4096, dim_attn=4096, dim_ffn=10240, num_heads=64, num_layers=24, num_buckets=32, shared_pos=False, dropout=0.1): 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 # layers self.token_embedding = vocab if isinstance(vocab, nn.Embedding) \ else nn.Embedding(vocab, dim) self.pos_embedding = T5RelativeEmbedding( num_buckets, num_heads, bidirectional=True) if shared_pos else 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) # Skip costly random init when loading pretrained checkpoints. # Verified: checkpoint fully covers this module (missing/unexpected keys are both 0). # self.apply(init_weights) def forward(self, ids, mask=None): x = self.token_embedding(ids) x = self.dropout(x) e = self.pos_embedding(x.size(1), x.size(1)) if self.shared_pos else None for block in self.blocks: x = block(x, mask, pos_bias=e) x = self.norm(x) x = self.dropout(x) return x def basic_clean(text): text = ftfy.fix_text(text) text = html.unescape(html.unescape(text)) return text.strip() def whitespace_clean(text): text = re.sub(r'\s+', ' ', text) text = text.strip() return text def canonicalize(text, keep_punctuation_exact_string=None): text = text.replace('_', ' ') if keep_punctuation_exact_string: text = keep_punctuation_exact_string.join( part.translate(str.maketrans('', '', string.punctuation)) for part in text.split(keep_punctuation_exact_string)) else: text = text.translate(str.maketrans('', '', string.punctuation)) text = text.lower() text = re.sub(r'\s+', ' ', text) return text.strip() class HuggingfaceTokenizer: def __init__(self, name, seq_len=None, clean=None, **kwargs): assert clean in (None, 'whitespace', 'lower', 'canonicalize') self.name = name self.seq_len = seq_len self.clean = clean # init tokenizer self.tokenizer = AutoTokenizer.from_pretrained(name, **kwargs) self.vocab_size = self.tokenizer.vocab_size def __call__(self, sequence, **kwargs): return_mask = kwargs.pop('return_mask', False) # arguments _kwargs = {'return_tensors': 'pt'} if self.seq_len is not None: _kwargs.update({ 'padding': 'max_length', 'truncation': True, 'max_length': self.seq_len }) _kwargs.update(**kwargs) # tokenization if isinstance(sequence, str): sequence = [sequence] if self.clean: sequence = [self._clean(u) for u in sequence] ids = self.tokenizer(sequence, **_kwargs) # output if return_mask: return ids.input_ids, ids.attention_mask else: return ids.input_ids def _clean(self, text): if self.clean == 'whitespace': text = whitespace_clean(basic_clean(text)) elif self.clean == 'lower': text = whitespace_clean(basic_clean(text)).lower() elif self.clean == 'canonicalize': text = canonicalize(basic_clean(text)) return text # ----- value DiT (value_dit.py) ----- # Vendored Wan2.2 / DiffSynth-Studio inference code: value_dit.py # Local change: logging import -> nimcritic.utils.logging. import os import torch import torch.nn as nn from typing import Any, Dict, Optional logger = get_logger(__name__) class ValueHead(nn.Module): def __init__(self, hidden_dim: int, out_dim: int, eps: float): super().__init__() self.norm = nn.LayerNorm(hidden_dim, eps=eps, elementwise_affine=False) self.proj = nn.Linear(hidden_dim, out_dim) self.modulation = nn.Parameter(torch.randn(1, 2, hidden_dim) / hidden_dim**0.5) def forward(self, x: torch.Tensor, t: torch.Tensor) -> torch.Tensor: # Supports both per-sample t ([B, hidden_dim]) and per-token t ([B, seq, hidden_dim]). if t.ndim == 3: shift, scale = (self.modulation.unsqueeze(0).to(dtype=t.dtype, device=t.device) + t.unsqueeze(2)).chunk(2, dim=2) shift = shift.squeeze(2) # [B, seq, hidden_dim] scale = scale.squeeze(2) return self.proj(self.norm(x) * (1 + scale) + shift) shift, scale = (self.modulation.to(dtype=t.dtype, device=t.device) + t.unsqueeze(1)).chunk(2, dim=1) shift = shift.squeeze(1) scale = scale.squeeze(1) return self.proj(self.norm(x) * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)) class ValueDiT(nn.Module): VALUE_BACKBONE_SKIP_PREFIXES = ("value_encoder.", "head.") VALUE_BACKBONE_META_KEYS = ( "hidden_dim", "ffn_dim", "num_layers", "num_heads", "attn_head_dim", "text_dim", "freq_dim", "eps", ) def __init__( self, hidden_dim: int, value_dim: int, ffn_dim: int, text_dim: int, freq_dim: int, eps: float, num_heads: int, attn_head_dim: int, num_layers: int, use_gradient_checkpointing: bool = False, ): super().__init__() self.hidden_dim = hidden_dim self.value_dim = value_dim self.ffn_dim = ffn_dim self.text_dim = text_dim self.freq_dim = freq_dim self.num_heads = num_heads self.attn_head_dim = attn_head_dim if num_heads <= 0: raise ValueError(f"`num_heads` must be > 0, got {num_heads}") if attn_head_dim <= 0: raise ValueError(f"`attn_head_dim` must be > 0, got {attn_head_dim}") if attn_head_dim % 2 != 0: raise ValueError(f"`attn_head_dim` must be even for RoPE, got {attn_head_dim}") self.value_encoder = nn.Linear(value_dim, hidden_dim) self.text_embedding = nn.Sequential( nn.Linear(text_dim, hidden_dim), nn.GELU(approximate="tanh"), nn.Linear(hidden_dim, hidden_dim), ) self.time_embedding = nn.Sequential( nn.Linear(freq_dim, hidden_dim), nn.SiLU(), nn.Linear(hidden_dim, hidden_dim), ) self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(hidden_dim, hidden_dim * 6)) self.blocks = nn.ModuleList( [ DiTBlock( hidden_dim=hidden_dim, attn_head_dim=attn_head_dim, num_heads=num_heads, ffn_dim=ffn_dim, eps=eps, ) for _ in range(num_layers) ] ) self.head = ValueHead(hidden_dim, value_dim, eps) self.freqs = precompute_freqs_cis(attn_head_dim, end=1024) self.use_gradient_checkpointing = use_gradient_checkpointing @classmethod def backbone_key_set(cls, keys) -> set[str]: return { key for key in keys if not any(key.startswith(prefix) for prefix in cls.VALUE_BACKBONE_SKIP_PREFIXES) } @classmethod def from_pretrained( cls, value_dit_config: dict[str, Any], value_dit_pretrained_path: str | None = None, skip_dit_load_from_pretrain: bool = False, device: str = "cuda", torch_dtype: torch.dtype = torch.bfloat16, ) -> "ValueDiT": if value_dit_config is None: raise ValueError("`value_dit_config` is required for ValueDiT.from_pretrained().") if skip_dit_load_from_pretrain: logger.info( "Skipping ValueDiT pretrained load (`skip_dit_load_from_pretrain=True`); " "initializing value expert randomly and expecting checkpoint override." ) return cls(**value_dit_config).to(device=device, dtype=torch_dtype) if not value_dit_pretrained_path: logger.info("No `value_dit_pretrained_path` provided, initializing ValueDiT with random weights.") return cls(**value_dit_config).to(device=device, dtype=torch_dtype) from pathlib import Path p = Path(value_dit_pretrained_path) if not p.is_absolute(): p = Path(__file__).resolve().parent / p value_dit_pretrained_path = str(p) if not os.path.isfile(value_dit_pretrained_path): raise FileNotFoundError( f"`value_dit_pretrained_path` does not exist: {value_dit_pretrained_path}" ) value_cfg = dict(value_dit_config) value_expert = cls(**value_cfg).to(device=device, dtype=torch_dtype) value_state = value_expert.state_dict() expected_backbone_keys = cls.backbone_key_set(value_state.keys()) payload = torch.load(value_dit_pretrained_path, map_location="cpu") if not isinstance(payload, dict): raise ValueError( f"Invalid value backbone payload type from {value_dit_pretrained_path}: {type(payload)}" ) policy = payload.get("policy", {}) if policy: logger.info(f"ValueDiT backbone payload policy: {policy}") meta = payload.get("meta") expected_meta = { "hidden_dim": int(value_cfg["hidden_dim"]), "ffn_dim": int(value_cfg["ffn_dim"]), "num_layers": int(value_cfg["num_layers"]), "num_heads": int(value_cfg["num_heads"]), "attn_head_dim": int(value_cfg["attn_head_dim"]), "text_dim": int(value_cfg["text_dim"]), "freq_dim": int(value_cfg["freq_dim"]), "eps": float(value_cfg["eps"]), } for key in cls.VALUE_BACKBONE_META_KEYS: if key not in meta: raise ValueError(f"`meta.{key}` missing in {value_dit_pretrained_path}") expected_value = expected_meta[key] got_value = meta[key] if key == "eps": if abs(float(got_value) - float(expected_value)) > 1e-12: raise ValueError( f"`meta.{key}` mismatch in {value_dit_pretrained_path}: " f"expected {expected_value}, got {got_value}" ) elif int(got_value) != int(expected_value): raise ValueError( f"`meta.{key}` mismatch in {value_dit_pretrained_path}: " f"expected {expected_value}, got {got_value}" ) backbone_state_dict = payload.get("backbone_state_dict") if not isinstance(backbone_state_dict, dict): raise ValueError( f"`backbone_state_dict` must be a dict in {value_dit_pretrained_path}, " f"got {type(backbone_state_dict)}" ) provided_keys = set(backbone_state_dict.keys()) missing_keys = sorted(expected_backbone_keys - provided_keys) unexpected_keys = sorted(provided_keys - expected_backbone_keys) if missing_keys or unexpected_keys: raise ValueError( "Value backbone key mismatch in preprocessed payload. " f"missing={missing_keys[:10]}{'...' if len(missing_keys) > 10 else ''}, " f"unexpected={unexpected_keys[:10]}{'...' if len(unexpected_keys) > 10 else ''}" ) merged_state = dict(value_state) for key in expected_backbone_keys: value = backbone_state_dict[key] if not isinstance(value, torch.Tensor): raise ValueError( f"`backbone_state_dict[{key}]` must be torch.Tensor in {value_dit_pretrained_path}, " f"got {type(value)}" ) target = merged_state[key] if tuple(value.shape) != tuple(target.shape): raise ValueError( f"Shape mismatch for `{key}` in {value_dit_pretrained_path}: " f"expected {tuple(target.shape)}, got {tuple(value.shape)}" ) merged_state[key] = value.to(device=target.device, dtype=target.dtype) value_expert.load_state_dict(merged_state, strict=True) logger.info( "Loaded ValueDiT backbone from %s (keys=%d; random_kept_prefixes=%s).", value_dit_pretrained_path, len(expected_backbone_keys), list(cls.VALUE_BACKBONE_SKIP_PREFIXES), ) return value_expert.to(device=device, dtype=torch_dtype) def pre_dit( self, value_tokens: torch.Tensor, timestep: torch.Tensor, context: torch.Tensor, context_mask: Optional[torch.Tensor] = None, condition_frames: int = 0, ) -> Dict[str, Any]: if value_tokens.ndim != 3: raise ValueError( f"`value_tokens` must be 3D [B, T, value_dim], got shape {tuple(value_tokens.shape)}" ) if value_tokens.shape[2] != self.value_dim: raise ValueError( f"`value_tokens` last dim must be {self.value_dim}, got {value_tokens.shape[2]}" ) if timestep.ndim != 1: raise ValueError(f"`timestep` must be 1D [B] or [1], got shape {tuple(timestep.shape)}") if context.ndim != 3: raise ValueError( f"`context` must be 3D [B, L, D], got shape {tuple(context.shape)}" ) batch_size = value_tokens.shape[0] if context.shape[0] != batch_size: raise ValueError( f"Batch mismatch between value tokens and text context: {batch_size} vs {context.shape[0]}" ) if timestep.shape[0] not in (1, batch_size): raise ValueError( f"`timestep` length must be 1 or batch_size({batch_size}), got {timestep.shape[0]}" ) if timestep.shape[0] == 1 and batch_size > 1: if self.training: raise ValueError("During training, value timestep length must match batch_size.") timestep = timestep.expand(batch_size) if context_mask is None: context_mask = torch.ones( (batch_size, context.shape[1]), dtype=torch.bool, device=context.device ) else: if context_mask.ndim != 2: raise ValueError(f"`context_mask` must be 2D [B, L], got shape {tuple(context_mask.shape)}") if context_mask.shape[0] != batch_size or context_mask.shape[1] != context.shape[1]: raise ValueError( f"`context_mask` shape must match `context` shape [B, L], got {tuple(context_mask.shape)} vs {tuple(context.shape)}" ) seq_len = value_tokens.shape[1] if seq_len > self.freqs.shape[0]: raise ValueError( f"Value token length {seq_len} exceeds RoPE cache {self.freqs.shape[0]}." ) # Per-token timestep: condition timesteps are GT (clean, t=0); the rest use the sampled timestep. vcf = max(0, int(condition_frames)) if vcf > seq_len: raise ValueError(f"`condition_frames`={vcf} exceeds value horizon={seq_len}.") token_timesteps = torch.ones( (batch_size, seq_len), dtype=timestep.dtype, device=timestep.device ) * timestep.view(batch_size, 1) if vcf > 0: token_timesteps[:, :vcf] = 0 token_t_emb = sinusoidal_embedding_1d(self.freq_dim, token_timesteps.reshape(-1)) t = self.time_embedding(token_t_emb).reshape(batch_size, seq_len, self.hidden_dim) t_mod = self.time_projection(t).unflatten(2, (6, self.hidden_dim)) # [B, seq, 6, hdim] tokens = self.value_encoder(value_tokens) context_emb = self.text_embedding(context) context_attn_mask = context_mask.unsqueeze(1).expand(-1, seq_len, -1) freqs = self.freqs[:seq_len].view(seq_len, 1, -1).to(tokens.device) return { "tokens": tokens, "freqs": freqs, "t": t, "t_mod": t_mod, "context": context_emb, "context_mask": context_attn_mask, "meta": { "batch_size": batch_size, "seq_len": seq_len, }, } def post_dit(self, tokens: torch.Tensor, pre_state: Dict[str, Any]) -> torch.Tensor: return self.head(tokens, pre_state["t"]) def forward( self, value_tokens: torch.Tensor, timestep: torch.Tensor, context: torch.Tensor, context_mask: Optional[torch.Tensor] = None, ) -> torch.Tensor: pre_state = self.pre_dit( value_tokens=value_tokens, timestep=timestep, context=context, context_mask=context_mask, ) x = pre_state["tokens"] context = pre_state["context"] t_mod = pre_state["t_mod"] freqs = pre_state["freqs"] context_mask = pre_state["context_mask"] for block in self.blocks: if self.use_gradient_checkpointing: x = gradient_checkpoint_forward( block, self.use_gradient_checkpointing, x, context, t_mod, freqs, context_mask=context_mask, ) else: x = block(x, context, t_mod, freqs, context_mask=context_mask) return self.post_dit(x, pre_state) # ----- mixture of transformers (mot.py) ----- # Vendored Wan2.2 / DiffSynth-Studio inference code: mot.py # Local change: logging import -> nimcritic.utils.logging. from typing import Dict, Optional import torch import torch.nn as nn logger = get_logger(__name__) class MoT(nn.Module): def __init__( self, mixtures: Dict[str, nn.Module], mot_checkpoint_mixed_attn: bool = True, secondary_expert_name: str = "action", ): super().__init__() if not mixtures: raise ValueError("`mixtures` cannot be empty.") if "video" not in mixtures or secondary_expert_name not in mixtures: raise ValueError( f"`mixtures` must include both 'video' and '{secondary_expert_name}' experts." ) self.mixtures = nn.ModuleDict(mixtures) self.expert_order = list(self.mixtures.keys()) self.mot_checkpoint_mixed_attn = mot_checkpoint_mixed_attn self.secondary_expert_name = secondary_expert_name if mot_checkpoint_mixed_attn: logger.info("Using gradient checkpointing for mixture attention. This will save memory but use more computation.") first_expert = self.mixtures[self.expert_order[0]] self.num_layers = len(first_expert.blocks) self.num_heads = first_expert.num_heads self.attn_head_dim = first_expert.attn_head_dim for name in self.expert_order[1:]: expert = self.mixtures[name] if len(expert.blocks) != self.num_layers: raise ValueError( f"All experts must have same number of layers; got {self.num_layers} and {len(expert.blocks)}" ) if expert.num_heads != self.num_heads: raise ValueError( f"All experts must have same num_heads; got {self.num_heads} and {expert.num_heads}" ) if expert.attn_head_dim != self.attn_head_dim: raise ValueError( "All experts must have same attn_head_dim; " f"got {self.attn_head_dim} and {expert.attn_head_dim}" ) logger.info(f"Initialized MoT with experts: {self.expert_order}, num_layers={self.num_layers}") for name in self.expert_order: expert = self.mixtures[name] logger.info(f" Expert '{name}': num_params={sum(p.numel() for p in expert.parameters()) / 1e9:.2f} B") @staticmethod def _split_modulation(block, t_mod: torch.Tensor): has_seq = len(t_mod.shape) == 4 chunk_dim = 2 if has_seq else 1 base_mod = block.modulation.to(dtype=t_mod.dtype, device=t_mod.device) shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (base_mod + t_mod).chunk(6, dim=chunk_dim) if has_seq: # means t_mod has separate modulation for each token, otherwise same modulation for all tokens in the block shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( shift_msa.squeeze(2), scale_msa.squeeze(2), gate_msa.squeeze(2), shift_mlp.squeeze(2), scale_mlp.squeeze(2), gate_mlp.squeeze(2), ) return shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp def _mixed_attention( self, q_cat: torch.Tensor, k_cat: torch.Tensor, v_cat: torch.Tensor, attention_mask: torch.Tensor, ) -> torch.Tensor: attn_mask = attention_mask.to(device=q_cat.device) def _forward(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: return flash_attention(q=q, k=k, v=v, num_heads=self.num_heads, ctx_mask=attn_mask) if self.mot_checkpoint_mixed_attn and self.training: return torch.utils.checkpoint.checkpoint( _forward, q_cat, k_cat, v_cat, use_reentrant=False, ) return _forward(q_cat, k_cat, v_cat) @staticmethod def _apply_expert_post_block( block, residual_x: torch.Tensor, mixed_attn_out: torch.Tensor, gate_msa: torch.Tensor, shift_mlp: torch.Tensor, scale_mlp: torch.Tensor, gate_mlp: torch.Tensor, context_payload: Optional[dict], ) -> torch.Tensor: x = block.gate(residual_x, gate_msa, block.self_attn.o(mixed_attn_out)) if context_payload is not None: context = context_payload.get("context") if context is not None: context_mask = context_payload.get("mask") if context_mask is not None and context_mask.dim() == 3: context_mask = context_mask.unsqueeze(1) x = x + block.cross_attn(block.norm3(x), context, ctx_mask=context_mask) mlp_input = modulate(block.norm2(x), shift_mlp, scale_mlp) x = block.gate(x, gate_mlp, block.ffn(mlp_input)) return x def _build_expert_attention_io( self, expert, block, x: torch.Tensor, freqs: torch.Tensor, t_mod: torch.Tensor, ) -> tuple[ torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, bool, ]: """Build per-expert attention tensors and post-block states. Args: expert: Expert module that owns this `block`; only used to read `use_gradient_checkpointing`. block: Transformer block for current layer (`expert.blocks[layer_idx]`). x: Current expert tokens, shape [B, S, D]. freqs: RoPE frequencies aligned with token sequence, shape [S, 1, rope_dim]. t_mod: Time modulation tensor for this expert/layer. Returns: q: Query after q-proj, RMSNorm, and RoPE, shape [B, S, H*Dh]. k: Key after k-proj, RMSNorm, and RoPE, shape [B, S, H*Dh]. v: Value after v-proj, shape [B, S, H*Dh]. residual_x: Original input `x` for residual path in post block. gate_msa: Gating tensor for self-attention residual branch. shift_mlp: Shift tensor for MLP modulation. scale_mlp: Scale tensor for MLP modulation. gate_mlp: Gating tensor for MLP residual branch. use_gradient_checkpointing: Whether this expert enables checkpointing. """ shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self._split_modulation(block, t_mod) attn_input = modulate(block.norm1(x), shift_msa, scale_msa) q = block.self_attn.norm_q(block.self_attn.q(attn_input)) k = block.self_attn.norm_k(block.self_attn.k(attn_input)) v = block.self_attn.v(attn_input) q = rope_apply(q, freqs, block.num_heads) k = rope_apply(k, freqs, block.num_heads) use_gradient_checkpointing = bool(getattr(expert, "use_gradient_checkpointing", False)) return ( q, k, v, x, gate_msa, shift_mlp, scale_mlp, gate_mlp, use_gradient_checkpointing, ) def _apply_post_with_optional_checkpoint( self, block, residual_x: torch.Tensor, gate_msa: torch.Tensor, shift_mlp: torch.Tensor, scale_mlp: torch.Tensor, gate_mlp: torch.Tensor, use_gradient_checkpointing: bool, mixed_slice: torch.Tensor, context_payload: Optional[dict], ) -> torch.Tensor: """Apply post-attention computations, with optional checkpointing. Args: block: Transformer block for current layer. residual_x: Residual input tokens before attention update, shape [B, S, D]. gate_msa: Gating tensor used after mixed self-attention. shift_mlp: Shift tensor for MLP input modulation. scale_mlp: Scale tensor for MLP input modulation. gate_mlp: Gating tensor used after MLP. use_gradient_checkpointing: If True and training, checkpoint this post block. mixed_slice: Mixed-attention output for this expert, shape [B, S, H*Dh]. context_payload: Optional dict for cross-attention. - `context`: encoder states [B, L, D] - `mask`: attention mask [B, S, L] or [B, 1, S, L] Returns: Updated expert tokens after self-attn residual, optional cross-attn, and MLP. """ def _post_fn( _mixed_slice: torch.Tensor, _x: torch.Tensor, _gate_msa: torch.Tensor, _shift_mlp: torch.Tensor, _scale_mlp: torch.Tensor, _gate_mlp: torch.Tensor, _block=block, _context_payload=context_payload, ) -> torch.Tensor: return self._apply_expert_post_block( block=_block, residual_x=_x, mixed_attn_out=_mixed_slice, gate_msa=_gate_msa, shift_mlp=_shift_mlp, scale_mlp=_scale_mlp, gate_mlp=_gate_mlp, context_payload=_context_payload, ) if use_gradient_checkpointing and self.training: return torch.utils.checkpoint.checkpoint( _post_fn, mixed_slice, residual_x, gate_msa, shift_mlp, scale_mlp, gate_mlp, use_reentrant=False, ) return _post_fn( mixed_slice, residual_x, gate_msa, shift_mlp, scale_mlp, gate_mlp, ) def prefill_video_cache( self, video_tokens: torch.Tensor, video_freqs: torch.Tensor, video_t_mod: torch.Tensor, video_context_payload: Optional[dict], video_attention_mask: torch.Tensor, ) -> list[dict[str, torch.Tensor]]: """Prefill video branch once and cache per-layer K/V for action denoising. Args: video_tokens: Video tokens before layer 0, shape [B, Sv, D]. video_freqs: Video RoPE frequencies, shape [Sv, 1, rope_dim]. video_t_mod: Video time modulation tensor. video_context_payload: Optional dict for video cross-attention. - `context`: encoder states [B, L, D] - `mask`: attention mask [B, Sv, L] or [B, 1, Sv, L] video_attention_mask: Video self-attention mask, shape [Sv, Sv]. Returns: Layer-wise cache list with length `num_layers`. Each entry contains: - `k`: video key tensor [B, Sv, H*Dh] - `v`: video value tensor [B, Sv, H*Dh] """ if "video" not in self.mixtures: raise ValueError("MoT requires `video` expert for `prefill_video_cache`.") if video_attention_mask.ndim != 2: raise ValueError( f"`video_attention_mask` must be 2D [S,S], got shape {tuple(video_attention_mask.shape)}" ) if video_attention_mask.shape[0] != video_attention_mask.shape[1]: raise ValueError( f"`video_attention_mask` must be square, got shape {tuple(video_attention_mask.shape)}" ) if video_attention_mask.shape[0] != video_tokens.shape[1]: raise ValueError( "`video_attention_mask` seq length mismatch: " f"mask={video_attention_mask.shape[0]} vs tokens={video_tokens.shape[1]}" ) expert = self.mixtures["video"] x = video_tokens kv_cache: list[dict[str, torch.Tensor]] = [] for layer_idx in range(self.num_layers): block = expert.blocks[layer_idx] # Build video Q/K/V from current layer input tokens. ( q, k, v, residual_x, gate_msa, shift_mlp, scale_mlp, gate_mlp, use_gradient_checkpointing, ) = self._build_expert_attention_io( expert=expert, block=block, x=x, freqs=video_freqs, t_mod=video_t_mod, ) # Video prefill uses only video self-attention mask. mixed = self._mixed_attention( q_cat=q, k_cat=k, v_cat=v, attention_mask=video_attention_mask, ) # Update video tokens for the next layer and persist current layer K/V. x = self._apply_post_with_optional_checkpoint( block=block, residual_x=residual_x, gate_msa=gate_msa, shift_mlp=shift_mlp, scale_mlp=scale_mlp, gate_mlp=gate_mlp, use_gradient_checkpointing=use_gradient_checkpointing, mixed_slice=mixed, context_payload=video_context_payload, ) kv_cache.append({"k": k, "v": v}) return kv_cache def forward_secondary_with_video_cache( self, secondary_tokens: torch.Tensor, secondary_freqs: torch.Tensor, secondary_t_mod: torch.Tensor, secondary_context_payload: Optional[dict], video_kv_cache: list[dict[str, torch.Tensor]], attention_mask: torch.Tensor, video_seq_len: int, expert_name: Optional[str] = None, ) -> torch.Tensor: """Run secondary branch (action/value) with cached video K/V. Args: secondary_tokens: Secondary tokens before layer 0, shape [B, Ss, D]. secondary_freqs: Secondary RoPE frequencies, shape [Ss, 1, rope_dim]. secondary_t_mod: Secondary time modulation tensor. secondary_context_payload: Optional dict for secondary cross-attention. - `context`: encoder states [B, L, D] - `mask`: attention mask [B, Ss, L] or [B, 1, Ss, L] video_kv_cache: Layer-wise cached video K/V from `prefill_video_cache`. attention_mask: Joint [video+secondary] mask, shape [Sv+Ss, Sv+Ss]. video_seq_len: Video token count `Sv` in the joint sequence prefix. expert_name: Secondary expert name (defaults to `self.secondary_expert_name`). Returns: Updated secondary tokens after all layers, shape [B, Ss, D]. """ if expert_name is None: expert_name = self.secondary_expert_name if expert_name not in self.mixtures: raise ValueError( f"MoT requires `{expert_name}` expert for `forward_secondary_with_video_cache`." ) if len(video_kv_cache) != self.num_layers: raise ValueError( f"`video_kv_cache` must contain {self.num_layers} layers, got {len(video_kv_cache)}." ) if attention_mask.ndim != 2: raise ValueError(f"`attention_mask` must be 2D [S,S], got shape {tuple(attention_mask.shape)}") if attention_mask.shape[0] != attention_mask.shape[1]: raise ValueError(f"`attention_mask` must be square, got shape {tuple(attention_mask.shape)}") secondary_seq_len = int(secondary_tokens.shape[1]) total_seq_len = int(video_seq_len) + secondary_seq_len if attention_mask.shape[0] != total_seq_len: raise ValueError( "`attention_mask` seq length mismatch: " f"mask={attention_mask.shape[0]} vs expected_total={total_seq_len}" ) # Use the secondary query rows from the joint [video+secondary] mask. secondary_attention_mask = attention_mask[video_seq_len:total_seq_len, :total_seq_len] expert = self.mixtures[expert_name] x = secondary_tokens for layer_idx in range(self.num_layers): block = expert.blocks[layer_idx] # Secondary query/key/value are still step-dependent and must be recomputed each step. ( q_secondary, k_secondary, v_secondary, residual_x, gate_msa, shift_mlp, scale_mlp, gate_mlp, use_gradient_checkpointing, ) = self._build_expert_attention_io( expert=expert, block=block, x=x, freqs=secondary_freqs, t_mod=secondary_t_mod, ) layer_cache = video_kv_cache[layer_idx] if "k" not in layer_cache or "v" not in layer_cache: raise ValueError( f"`video_kv_cache[{layer_idx}]` must contain `k` and `v`." ) k_video = layer_cache["k"] v_video = layer_cache["v"] if k_video.shape[1] != video_seq_len or v_video.shape[1] != video_seq_len: raise ValueError( f"`video_kv_cache[{layer_idx}]` seq len mismatch, expected {video_seq_len}." ) # Mixed attention: secondary queries attend to cached video K/V plus current secondary K/V. k_cat = torch.cat([k_video, k_secondary], dim=1) v_cat = torch.cat([v_video, v_secondary], dim=1) mixed = self._mixed_attention( q_cat=q_secondary, k_cat=k_cat, v_cat=v_cat, attention_mask=secondary_attention_mask, ) x = self._apply_post_with_optional_checkpoint( block=block, residual_x=residual_x, gate_msa=gate_msa, shift_mlp=shift_mlp, scale_mlp=scale_mlp, gate_mlp=gate_mlp, use_gradient_checkpointing=use_gradient_checkpointing, mixed_slice=mixed, context_payload=secondary_context_payload, ) return x def forward_action_with_video_cache( self, action_tokens: torch.Tensor, action_freqs: torch.Tensor, action_t_mod: torch.Tensor, action_context_payload: Optional[dict], video_kv_cache: list[dict[str, torch.Tensor]], attention_mask: torch.Tensor, video_seq_len: int, ) -> torch.Tensor: """Backward-compatible wrapper around `forward_secondary_with_video_cache`.""" return self.forward_secondary_with_video_cache( secondary_tokens=action_tokens, secondary_freqs=action_freqs, secondary_t_mod=action_t_mod, secondary_context_payload=action_context_payload, video_kv_cache=video_kv_cache, attention_mask=attention_mask, video_seq_len=video_seq_len, expert_name="action", ) def forward( self, embeds_all: Dict[str, torch.Tensor], attention_mask: torch.Tensor, freqs_all: Dict[str, torch.Tensor], context_all: Dict[str, Optional[dict]], t_mod_all: Dict[str, torch.Tensor], ): missing = [k for k in self.expert_order if k not in embeds_all] if missing: raise ValueError(f"Missing expert tokens for {missing}") missing = [k for k in self.expert_order if k not in freqs_all] if missing: raise ValueError(f"Missing expert freqs for {missing}") missing = [k for k in self.expert_order if k not in t_mod_all] if missing: raise ValueError(f"Missing expert t_mod for {missing}") if attention_mask.ndim != 2: raise ValueError(f"`attention_mask` must be 2D [S, S], got shape {tuple(attention_mask.shape)}") if attention_mask.shape[0] != attention_mask.shape[1]: raise ValueError(f"`attention_mask` must be square, got shape {tuple(attention_mask.shape)}") tokens_all = {k: v for k, v in embeds_all.items()} for layer_idx in range(self.num_layers): q_chunks = [] k_chunks = [] v_chunks = [] cached = {} seq_lens = [] for name in self.expert_order: expert = self.mixtures[name] block = expert.blocks[layer_idx] x = tokens_all[name] freqs = freqs_all[name] t_mod = t_mod_all[name] ( q, k, v, residual_x, gate_msa, shift_mlp, scale_mlp, gate_mlp, use_gradient_checkpointing, ) = self._build_expert_attention_io( expert=expert, block=block, x=x, freqs=freqs, t_mod=t_mod, ) q_chunks.append(q) k_chunks.append(k) v_chunks.append(v) seq_lens.append(x.shape[1]) cached[name] = { "block": block, "residual_x": residual_x, "gate_msa": gate_msa, "shift_mlp": shift_mlp, "scale_mlp": scale_mlp, "gate_mlp": gate_mlp, "use_gradient_checkpointing": use_gradient_checkpointing, } # 3. concat all tokens for mixed attention q_cat = torch.cat(q_chunks, dim=1) k_cat = torch.cat(k_chunks, dim=1) v_cat = torch.cat(v_chunks, dim=1) total_seq = q_cat.shape[1] if attention_mask.shape[0] != total_seq: raise ValueError( "Attention mask seq length mismatch: " f"mask={attention_mask.shape[0]} vs tokens={total_seq}" ) mixed = self._mixed_attention(q_cat=q_cat, k_cat=k_cat, v_cat=v_cat, attention_mask=attention_mask) start = 0 for name, seq_len in zip(self.expert_order, seq_lens): # 4. split mixed attention output and apply post-attention blocks for each expert end = start + seq_len mixed_slice = mixed[:, start:end, :] cached_expert = cached[name] block = cached_expert["block"] context_payload = context_all.get(name) updated_tokens = self._apply_post_with_optional_checkpoint( block=block, residual_x=cached_expert["residual_x"], gate_msa=cached_expert["gate_msa"], shift_mlp=cached_expert["shift_mlp"], scale_mlp=cached_expert["scale_mlp"], gate_mlp=cached_expert["gate_mlp"], use_gradient_checkpointing=cached_expert["use_gradient_checkpointing"], mixed_slice=mixed_slice, context_payload=context_payload, ) tokens_all[name] = updated_tokens start = end return tokens_all # ----- component loader (loader.py) ----- # Vendored Wan2.2 / DiffSynth-Studio inference code: helpers/loader.py # Local change: logging import -> nimcritic.utils.logging. from dataclasses import dataclass import inspect from typing import Any import torch import time logger = get_logger(__name__) SKIPPED_PRETRAIN_SENTINEL = "SKIPPED_PRETRAIN" @dataclass class Wan22LoadedComponents: dit: WanVideoDiT vae: WanVideoVAE38 text_encoder: WanTextEncoder | None tokenizer: HuggingfaceTokenizer | None dit_path: str vae_path: str text_encoder_path: str | None tokenizer_path: str | None WAN22_MODEL_REGISTRY = [ { # Example: ModelConfig(model_id="Wan-AI/Wan2.1-T2V-14B", origin_file_pattern="models_t5_umt5-xxl-enc-bf16.pth") "model_hash": "9c8818c2cbea55eca56c7b447df170da", "model_name": "wan_video_text_encoder", "model_class": WanTextEncoder, }, { # Example: ModelConfig(model_id="Wan-AI/Wan2.2-TI2V-5B", origin_file_pattern="diffusion_pytorch_model*.safetensors") "model_hash": "1f5ab7703c6fc803fdded85ff040c316", "model_name": "wan_video_dit", "model_class": WanVideoDiT, }, { # Example: ModelConfig(model_id="Wan-AI/Wan2.2-TI2V-5B", origin_file_pattern="Wan2.2_VAE.pth") "model_hash": "e1de6c02cdac79f8b739f4d3698cd216", "model_name": "wan_video_vae", "model_class": WanVideoVAE38, "state_dict_converter": wan_video_vae_state_dict_converter, }, ] def _validate_dit_config(dit_config: dict[str, Any]) -> dict[str, Any]: if not isinstance(dit_config, dict): raise ValueError(f"`dit_config` must be a dict, got {type(dit_config)}") validated = dict(dit_config) signature = inspect.signature(WanVideoDiT.__init__) allowed_keys = set() required_keys = set() for name, param in signature.parameters.items(): if name == "self": continue allowed_keys.add(name) if param.default is inspect.Signature.empty: required_keys.add(name) unknown_keys = sorted(set(validated) - allowed_keys) if unknown_keys: raise ValueError( f"Unknown keys in `dit_config`: {unknown_keys}. " f"Allowed keys: {sorted(allowed_keys)}" ) missing_keys = sorted(required_keys - set(validated)) if missing_keys: raise ValueError( f"Missing required keys in `dit_config`: {missing_keys}. " "Please specify all required WanVideoDiT constructor args." ) return validated def _load_registered_model( path, model_name: str, torch_dtype: torch.dtype, device: str, model_kwargs_override: dict[str, Any] | None = None, ): model_hash = hash_model_file(path) matched_config = None for config in WAN22_MODEL_REGISTRY: if config["model_hash"] == model_hash and config["model_name"] == model_name: matched_config = config break if matched_config is None: raise ValueError( f"Cannot detect model type for {model_name}. File: {path}. " f"Model hash: {model_hash}. This standalone package follows DiffSynth hash-based loading." ) model_class = matched_config["model_class"] model_kwargs = dict(matched_config.get("extra_kwargs", {})) if model_kwargs_override is not None: model_kwargs.update(model_kwargs_override) state_dict_converter = matched_config.get("state_dict_converter") model = model_class(**model_kwargs) state_dict = load_state_dict(path, torch_dtype=torch_dtype, device="cpu") if state_dict_converter is not None: state_dict = state_dict_converter(state_dict) model.load_state_dict(state_dict, strict=False) model = model.to(device=device, dtype=torch_dtype) return model def _resolve_configs(model_id: str, tokenizer_model_id: str, redirect_common_files: bool = True): dit_config = ModelConfig(model_id=model_id, origin_file_pattern="diffusion_pytorch_model*.safetensors") text_config = ModelConfig(model_id=model_id, origin_file_pattern="models_t5_umt5-xxl-enc-bf16.pth") vae_config = ModelConfig(model_id=model_id, origin_file_pattern="Wan2.2_VAE.pth") tokenizer_config = ModelConfig(model_id=tokenizer_model_id, origin_file_pattern="google/umt5-xxl/") if redirect_common_files: redirect_dict = { "models_t5_umt5-xxl-enc-bf16.pth": ("DiffSynth-Studio/Wan-Series-Converted-Safetensors", "models_t5_umt5-xxl-enc-bf16.safetensors"), "Wan2.2_VAE.pth": ("DiffSynth-Studio/Wan-Series-Converted-Safetensors", "Wan2.2_VAE.safetensors"), } text_config.model_id, text_config.origin_file_pattern = redirect_dict[text_config.origin_file_pattern] vae_config.model_id, vae_config.origin_file_pattern = redirect_dict[vae_config.origin_file_pattern] return dit_config, text_config, vae_config, tokenizer_config def load_wan22_ti2v_5b_components( device: str = "cuda", torch_dtype: torch.dtype = torch.bfloat16, model_id: str = "Wan-AI/Wan2.2-TI2V-5B", tokenizer_model_id: str = "Wan-AI/Wan2.1-T2V-1.3B", tokenizer_max_len: int = 512, redirect_common_files: bool = True, dit_config: dict[str, Any] | None = None, skip_dit_load_from_pretrain: bool = False, load_text_encoder: bool = True, ): logger.info("Loading Wan2.2-TI2V-5B components...") start = time.time() if dit_config is None: raise ValueError("`dit_config` is required for Wan2.2-TI2V-5B loading.") validated_dit_config = _validate_dit_config(dit_config) dit_model_config, text_config, vae_config, tokenizer_config = _resolve_configs( model_id=model_id, tokenizer_model_id=tokenizer_model_id, redirect_common_files=redirect_common_files, ) vae_config.download_if_necessary() if load_text_encoder: text_config.download_if_necessary() tokenizer_config.download_if_necessary() if skip_dit_load_from_pretrain: logger.info( "Skipping pretrained video DiT load (`skip_dit_load_from_pretrain=True`); " "initializing video expert randomly and expecting checkpoint override." ) dit: WanVideoDiT = WanVideoDiT(**validated_dit_config).to(device=device, dtype=torch_dtype) dit_path = SKIPPED_PRETRAIN_SENTINEL else: dit_model_config.download_if_necessary() dit = _load_registered_model( dit_model_config.path, "wan_video_dit", torch_dtype=torch_dtype, device=device, model_kwargs_override=validated_dit_config, ) dit_path = str(dit_model_config.path) text_encoder: WanTextEncoder | None = None tokenizer: HuggingfaceTokenizer | None = None text_encoder_path: str | None = None tokenizer_path: str | None = None if load_text_encoder: text_encoder = _load_registered_model( text_config.path, "wan_video_text_encoder", torch_dtype=torch_dtype, device=device, ) tokenizer = HuggingfaceTokenizer( name=tokenizer_config.path, seq_len=int(tokenizer_max_len), clean="whitespace", ) text_encoder_path = str(text_config.path) tokenizer_path = str(tokenizer_config.path) else: logger.info( "Skipping pretrained text encoder/tokenizer load (`load_text_encoder=False`); " "training must provide cached `context/context_mask`." ) vae: WanVideoVAE38 = _load_registered_model(vae_config.path, "wan_video_vae", torch_dtype=torch_dtype, device=device) logger.info("Finished loading Wan2.2-TI2V-5B components in %.2f seconds.", time.time() - start) return Wan22LoadedComponents( dit=dit, vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, dit_path=dit_path, vae_path=str(vae_config.path), text_encoder_path=text_encoder_path, tokenizer_path=tokenizer_path, ) # ----- NimCritic (nimcritic.py) ----- # Vendored Wan2.2 / DiffSynth-Studio inference code: fast_wvm.py (trimmed). # Kept: construction (from_wan22_pretrained), load_checkpoint, the fast # value-inference path (encode_prompt / _encode_input_image_latents_tensor / # _build_mot_attention_mask / _predict_value_noise_with_cache / infer_value), # and the joint video+value generation path (_check_resize_height_width / # _decode_latents / _predict_joint_noise / infer_joint / infer). # Removed: training (build_inputs / training_loss / _compute_video_loss_per_sample / # forward / save_checkpoint), dead code never called upstream # (_predict_value_noise / _encode_video_latents), proprio dead code, and the # trainer-compat `dit` alias. # Patched: `_encode_input_image_latents_tensor` accepts multi-frame [N,3,H,W] # (upstream single-image check contradicts `infer_value`'s multi-frame contract). # # Patched: infer_value / infer_joint accept batch>1 (5D input_image, [B,T,v_dim] # value, list[str] prompt). # Patched: value_progress_repr absolute/relative — relative skips 2x-1/denorm on # progress (col 0) so predicted progress stays a [-1, 1] delta; success still # maps [0, 1] ↔ [-1, 1]. from collections.abc import Sequence from typing import Any import torch from PIL import Image logger = get_logger(__name__) class NimCritic(torch.nn.Module): """MoT world model with video/value experts (value replaces action).""" def __init__( self, video_expert, value_expert: ValueDiT, mot: MoT, vae, text_encoder=None, tokenizer=None, text_dim: int | None = None, device: str = "cpu", torch_dtype: torch.dtype = torch.float32, video_infer_shift: float = 5.0, video_num_train_timesteps: int = 1000, value_infer_shift: float = 5.0, value_num_train_timesteps: int = 1000, video_condition_frames: int = 1, value_condition_frames: int = 1, ): super().__init__() self.video_expert = video_expert self.value_expert = value_expert self.mot = mot if int(video_condition_frames) < 1: raise ValueError(f"`video_condition_frames` must be >= 1, got {video_condition_frames}") if int(value_condition_frames) < 1: raise ValueError(f"`value_condition_frames` must be >= 1, got {value_condition_frames}") self.video_condition_frames = int(video_condition_frames) self.value_condition_frames = int(value_condition_frames) self.vae = vae self.text_encoder = text_encoder self.tokenizer = tokenizer if text_dim is None: if self.text_encoder is None: raise ValueError("`text_dim` is required when `text_encoder` is not loaded.") text_dim = int(self.text_encoder.dim) self.text_dim = int(text_dim) self.infer_video_scheduler = WanContinuousFlowMatchScheduler( num_train_timesteps=video_num_train_timesteps, shift=video_infer_shift, ) self.infer_value_scheduler = WanContinuousFlowMatchScheduler( num_train_timesteps=value_num_train_timesteps, shift=value_infer_shift, ) self.device = torch.device(device) self.torch_dtype = torch_dtype self.to(self.device) @classmethod def from_wan22_pretrained( cls, device: str = "cuda", torch_dtype: torch.dtype = torch.bfloat16, model_id: str = "Wan-AI/Wan2.2-TI2V-5B", tokenizer_model_id: str = "Wan-AI/Wan2.1-T2V-1.3B", tokenizer_max_len: int = 512, load_text_encoder: bool = True, redirect_common_files: bool = True, video_dit_config: dict[str, Any] | None = None, value_dit_config: dict[str, Any] | None = None, value_dit_pretrained_path: str | None = None, skip_dit_load_from_pretrain: bool = False, mot_checkpoint_mixed_attn: bool = True, video_infer_shift: float = 5.0, video_num_train_timesteps: int = 1000, value_infer_shift: float = 5.0, value_num_train_timesteps: int = 1000, video_condition_frames: int = 1, value_condition_frames: int = 1, ): if video_dit_config is None: raise ValueError("`video_dit_config` is required for NimCritic.from_wan22_pretrained().") if "text_dim" not in video_dit_config: raise ValueError("`video_dit_config['text_dim']` is required for NimCritic.") components = load_wan22_ti2v_5b_components( device=device, torch_dtype=torch_dtype, model_id=model_id, tokenizer_model_id=tokenizer_model_id, tokenizer_max_len=tokenizer_max_len, redirect_common_files=redirect_common_files, dit_config=video_dit_config, skip_dit_load_from_pretrain=skip_dit_load_from_pretrain, load_text_encoder=load_text_encoder, ) video_expert = components.dit value_expert = ValueDiT.from_pretrained( value_dit_config=value_dit_config, value_dit_pretrained_path=value_dit_pretrained_path, skip_dit_load_from_pretrain=skip_dit_load_from_pretrain, device=device, torch_dtype=torch_dtype, ) if int(value_expert.num_heads) != int(video_expert.num_heads): raise ValueError("ValueDiT `num_heads` must match video expert for MoT mixed attention.") if int(value_expert.attn_head_dim) != int(video_expert.attn_head_dim): raise ValueError("ValueDiT `attn_head_dim` must match video expert for MoT mixed attention.") if len(value_expert.blocks) != len(video_expert.blocks): raise ValueError("ValueDiT `num_layers` must match video expert.") mot = MoT( mixtures={"video": video_expert, "value": value_expert}, mot_checkpoint_mixed_attn=mot_checkpoint_mixed_attn, secondary_expert_name="value", ) model = cls( video_expert=video_expert, value_expert=value_expert, mot=mot, vae=components.vae, text_encoder=components.text_encoder, tokenizer=components.tokenizer, text_dim=int(video_dit_config["text_dim"]), device=device, torch_dtype=torch_dtype, video_infer_shift=video_infer_shift, video_num_train_timesteps=video_num_train_timesteps, value_infer_shift=value_infer_shift, value_num_train_timesteps=value_num_train_timesteps, video_condition_frames=video_condition_frames, value_condition_frames=value_condition_frames, ) model.model_paths = { "video_dit": components.dit_path, "vae": components.vae_path, "text_encoder": components.text_encoder_path, "tokenizer": components.tokenizer_path, "value_dit_backbone": ( "SKIPPED_PRETRAIN" if skip_dit_load_from_pretrain else value_dit_pretrained_path ), } return model def to(self, *args, **kwargs): super().to(*args, **kwargs) self.mot.to(*args, **kwargs) if self.text_encoder is not None: self.text_encoder.to(*args, **kwargs) self.vae.to(*args, **kwargs) return self @staticmethod def _check_resize_height_width(height, width, num_frames): if height % 16 != 0: height = (height + 15) // 16 * 16 if width % 16 != 0: width = (width + 15) // 16 * 16 if num_frames % 4 != 1: num_frames = (num_frames + 3) // 4 * 4 + 1 return height, width, num_frames @torch.no_grad() def encode_prompt(self, prompt: str | Sequence[str]): if self.text_encoder is None or self.tokenizer is None: raise ValueError( "Prompt encoding requires loaded text encoder/tokenizer. " "Set `load_text_encoder=true` or provide precomputed `context/context_mask`." ) ids, mask = self.tokenizer(prompt, return_mask=True, add_special_tokens=True) ids = ids.to(self.device) mask = mask.to(self.device, dtype=torch.bool) prompt_emb = self.text_encoder(ids, mask) seq_lens = mask.gt(0).sum(dim=1).long() for i, v in enumerate(seq_lens): prompt_emb[i, v:] = 0 mask = torch.ones_like(mask) return prompt_emb.to(device=self.device), mask def _coerce_input_images(self, input_image: torch.Tensor) -> torch.Tensor: """Normalize pixel condition to ``[B, N, 3, H, W]``. Accepted layouts: ``[3,H,W]``, ``[N,3,H,W]``, ``[1,N,3,H,W]``, ``[B,N,3,H,W]``. 4D ``[N,3,H,W]`` stays batch=1 so it is not confused with a batched single-frame ``[B,3,H,W]``. """ if input_image.ndim == 3: if input_image.shape[0] != 3: raise ValueError( f"`input_image` must have shape [3,H,W], [N,3,H,W], " f"[1,N,3,H,W] or [B,N,3,H,W], got {tuple(input_image.shape)}" ) return input_image.unsqueeze(0).unsqueeze(0) if input_image.ndim == 4: if input_image.shape[1] == 3: return input_image.unsqueeze(0) if input_image.shape[0] == 1 and input_image.shape[2] == 3: return input_image raise ValueError( f"`input_image` must have shape [3,H,W], [N,3,H,W], " f"[1,N,3,H,W] or [B,N,3,H,W], got {tuple(input_image.shape)}" ) if input_image.ndim == 5: if input_image.shape[2] != 3: raise ValueError( f"`input_image` must have shape [3,H,W], [N,3,H,W], " f"[1,N,3,H,W] or [B,N,3,H,W], got {tuple(input_image.shape)}" ) return input_image raise ValueError( f"`input_image` must have shape [3,H,W], [N,3,H,W], " f"[1,N,3,H,W] or [B,N,3,H,W], got {tuple(input_image.shape)}" ) def _coerce_value( self, value: torch.Tensor | None, horizon: int, batch: int, name: str, ) -> torch.Tensor | None: if value is None: return None if value.ndim == 2: value = value.unsqueeze(0) if value.ndim != 3 or value.shape[1] != horizon: raise ValueError( f"`{name}` must have shape [T, v_dim], [1, T, v_dim] or " f"[B, T, v_dim], got {tuple(value.shape)} with {name.replace('gt_', '')}_horizon={horizon}" ) if value.shape[0] not in (1, batch): raise ValueError( f"`{name}` batch {value.shape[0]} does not match input_image batch {batch}" ) if value.shape[0] == 1 and batch > 1: value = value.expand(batch, -1, -1) return value def _coerce_prompt(self, prompt: str | Sequence[str] | None, batch: int) -> str | list[str] | None: if prompt is None: return None if isinstance(prompt, str): return [prompt] * batch if batch > 1 else prompt prompts = [str(p) for p in prompt] if len(prompts) != batch: raise ValueError(f"`prompt` length {len(prompts)} does not match batch {batch}") return prompts[0] if batch == 1 else prompts @torch.no_grad() def _encode_input_image_latents_tensor(self, input_image: torch.Tensor, tiled=False, tile_size=(30, 52), tile_stride=(15, 26)): # The upstream encoder only accepted a single image while # `infer_value` already parses/validates multi-frame input ([N,3,H,W] with # N=(vcf-1)*tf+1). Encode N frames as one video so temporal latents match # the training-side `vae.encode` conditioning; N=1 is bit-identical to the # original single-image path. # 5D [B,N,3,H,W] encodes B videos in one VAE call. images = self._coerce_input_images(input_image) videos = [ images[b].to(device=self.device).permute(1, 0, 2, 3) # [N,3,H,W] -> [3,N,H,W] for b in range(images.shape[0]) ] z = self.vae.encode(videos, device=self.device, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride) if isinstance(z, list): stacked = [] for zi in z: if zi.ndim == 4: stacked.append(zi) elif zi.ndim == 5 and zi.shape[0] == 1: stacked.append(zi[0]) else: raise ValueError(f"VAE encode item must be [C,T,H,W] or [1,C,T,H,W], got {tuple(zi.shape)}") z = torch.stack(stacked, dim=0) elif z.ndim != 5: raise ValueError(f"VAE encode must return [B,C,T,H,W] or a list, got {tuple(z.shape)}") return z def _frames_from_video_tensor(self, video: torch.Tensor) -> list[Image.Image]: """``video`` is ``[3, T, H, W]`` in [-1, 1].""" video = video.detach().float().clamp(-1, 1) video = ((video + 1.0) * 127.5).to(torch.uint8).cpu() frames = [] for t in range(video.shape[1]): frame = video[:, t].permute(1, 2, 0).numpy() frames.append(Image.fromarray(frame)) return frames def _decode_latents(self, latents, tiled=False, tile_size=(30, 52), tile_stride=(15, 26)): video_tensor = self.vae.decode(latents, device=self.device, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride) if video_tensor.ndim != 5: raise ValueError( f"VAE decode must return [B,3,T,H,W], got {tuple(video_tensor.shape)}" ) batch = int(video_tensor.shape[0]) videos = [self._frames_from_video_tensor(video_tensor[b]) for b in range(batch)] return videos[0] if batch == 1 else videos @torch.no_grad() def _build_mot_attention_mask( self, video_seq_len: int, value_seq_len: int, video_tokens_per_frame: int, device: torch.device, video_condition_frames: int | None = None, ) -> torch.Tensor: vcf = self.video_condition_frames if video_condition_frames is None else int(video_condition_frames) total_seq_len = video_seq_len + value_seq_len mask = torch.zeros((total_seq_len, total_seq_len), dtype=torch.bool, device=device) # video -> video (first `vcf` frames are condition frames) mask[:video_seq_len, :video_seq_len] = self.video_expert.build_video_to_video_mask( video_seq_len=video_seq_len, video_tokens_per_frame=video_tokens_per_frame, device=device, condition_frames=vcf, ) # value -> value mask[video_seq_len:, video_seq_len:] = True # value -> first `vcf` video condition frames only # (value conditions on the observed video frames, which are video's condition frames) value_cond_tokens = min(vcf * video_tokens_per_frame, video_seq_len) mask[video_seq_len:, :value_cond_tokens] = True return mask @torch.no_grad() def _predict_joint_noise( self, latents_video: torch.Tensor, latents_value: torch.Tensor, timestep_video: torch.Tensor, timestep_value: torch.Tensor, context: torch.Tensor, context_mask: torch.Tensor, fuse_vae_embedding_in_latents: bool, gt_value: torch.Tensor | None = None, video_condition_frames: int | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: video_pre = self.video_expert.pre_dit( x=latents_video, timestep=timestep_video, context=context, context_mask=context_mask, action=gt_value, fuse_vae_embedding_in_latents=fuse_vae_embedding_in_latents, condition_frames=video_condition_frames if video_condition_frames is not None else self.video_condition_frames, ) value_pre = self.value_expert.pre_dit( value_tokens=latents_value, timestep=timestep_value, context=context, context_mask=context_mask, condition_frames=self.value_condition_frames, ) attention_mask = self._build_mot_attention_mask( video_seq_len=video_pre["tokens"].shape[1], value_seq_len=value_pre["tokens"].shape[1], video_tokens_per_frame=int(video_pre["meta"]["tokens_per_frame"]), device=video_pre["tokens"].device, video_condition_frames=video_condition_frames, ) tokens_out = self.mot( embeds_all={ "video": video_pre["tokens"], "value": value_pre["tokens"], }, attention_mask=attention_mask, freqs_all={ "video": video_pre["freqs"], "value": value_pre["freqs"], }, context_all={ "video": { "context": video_pre["context"], "mask": video_pre["context_mask"], }, "value": { "context": value_pre["context"], "mask": value_pre["context_mask"], }, }, t_mod_all={ "video": video_pre["t_mod"], "value": value_pre["t_mod"], }, ) pred_video = self.video_expert.post_dit(tokens_out["video"], video_pre) pred_value = self.value_expert.post_dit(tokens_out["value"], value_pre) return pred_video, pred_value @torch.no_grad() def _predict_value_noise_with_cache( self, latents_value: torch.Tensor, timestep_value: torch.Tensor, context: torch.Tensor, context_mask: torch.Tensor, video_kv_cache: list[dict[str, torch.Tensor]], attention_mask: torch.Tensor, video_seq_len: int, value_condition_frames: int | None = None, ) -> torch.Tensor: value_pre = self.value_expert.pre_dit( value_tokens=latents_value, timestep=timestep_value, context=context, context_mask=context_mask, condition_frames=self.value_condition_frames if value_condition_frames is None else value_condition_frames, ) value_tokens = self.mot.forward_secondary_with_video_cache( secondary_tokens=value_pre["tokens"], secondary_freqs=value_pre["freqs"], secondary_t_mod=value_pre["t_mod"], secondary_context_payload={ "context": value_pre["context"], "mask": value_pre["context_mask"], }, video_kv_cache=video_kv_cache, attention_mask=attention_mask, video_seq_len=video_seq_len, expert_name="value", ) return self.value_expert.post_dit(value_tokens, value_pre) @torch.no_grad() def infer_joint( self, prompt: str | Sequence[str] | None, input_image: torch.Tensor, num_video_frames: int, value_horizon: int, value: torch.Tensor | None = None, # GT value: first `value_condition_frames` timesteps used as condition context: torch.Tensor | None = None, context_mask: torch.Tensor | None = None, negative_prompt: str | None = None, text_cfg_scale: float = 1.0, num_inference_steps: int = 20, sigma_shift: float | None = None, seed: int | None = None, rand_device: str = "cpu", tiled: bool = False, test_value_with_infer_value: bool = True, video_condition_frames: int | None = None, value_condition_frames: int | None = None, value_progress_repr: str = "absolute", ) -> dict[str, Any]: # upstream signature also carries `proprio`; it is dead # code upstream ("proprio is not used in NimCritic") and stays removed here. self.eval() vcf = int(self.video_condition_frames) if video_condition_frames is None else int(video_condition_frames) vcf_v = int(self.value_condition_frames) if value_condition_frames is None else int(value_condition_frames) progress_repr = parse_value_progress_repr(value_progress_repr) if vcf < 1: raise ValueError(f"`video_condition_frames` must be >= 1, got {vcf}") if vcf_v < 1: raise ValueError(f"`value_condition_frames` must be >= 1, got {vcf_v}") images = self._coerce_input_images(input_image) batch = int(images.shape[0]) num_cond_frames = int(images.shape[1]) height, width = int(images.shape[3]), int(images.shape[4]) # `infer_value` only supports single-frame video condition; auto-disable the # cross-check for multi-frame video condition or B>1 (shape of value_only_out). _is_single_frame_input = batch == 1 and num_cond_frames == 1 if test_value_with_infer_value and (vcf != 1 or not _is_single_frame_input): logger.info( "Disabling `test_value_with_infer_value` because video_condition_frames=%d " "or multi-frame/batched input (single-frame B=1 required for infer_value).", vcf, ) test_value_with_infer_value = False if test_value_with_infer_value: if seed is None: raise ValueError("`test_value_with_infer_value=True` requires non-null `seed`.") value_only_out = self.infer_value( prompt=prompt, input_image=input_image.clone(), value_horizon=value_horizon, gt_value=value.clone() if value is not None else None, video_condition_frames=vcf, value_condition_frames=vcf_v, context=context.clone() if context is not None else None, context_mask=context_mask.clone() if context_mask is not None else None, num_inference_steps=num_inference_steps, sigma_shift=sigma_shift, seed=seed, rand_device=rand_device, tiled=tiled, value_progress_repr=progress_repr, )["value"] input_image = images if batch > 1 else images[0] prompt = self._coerce_prompt(prompt, batch) checked_h, checked_w, checked_t = self._check_resize_height_width(height, width, num_video_frames) if (checked_h, checked_w) != (height, width): raise ValueError( f"`input_image` must be resized before infer, expected multiples of 16 but got HxW=({height},{width})" ) if checked_t != num_video_frames: raise ValueError( f"`num_video_frames` must satisfy T % 4 == 1, got {num_video_frames}" ) # Validate that provided condition frames match `vcf` (in latent frames). tf = int(self.vae.temporal_downsample_factor) cond_latent_frames = (num_cond_frames - 1) // tf + 1 if tf > 0 else num_cond_frames if cond_latent_frames != vcf: raise ValueError( f"Provided {num_cond_frames} condition frame(s) -> {cond_latent_frames} latent frame(s), " f"but `video_condition_frames`={vcf}. They must match." ) if value_horizon < vcf_v: raise ValueError( f"`value_horizon`={value_horizon} < value_condition_frames={vcf_v}; need at least vcf_v GT timesteps." ) value = self._coerce_value(value, value_horizon, batch, "value") if value is not None: value = value.to(device=self.device, dtype=self.torch_dtype) value = value_to_latent(value, progress_repr) elif vcf_v > 0: raise ValueError( f"`value_condition_frames`={vcf_v} > 0 but no `value` (GT value) provided for conditioning." ) latent_t = (num_video_frames - 1) // self.vae.temporal_downsample_factor + 1 latent_h = height // self.vae.upsampling_factor latent_w = width // self.vae.upsampling_factor if latent_t < vcf: raise ValueError( f"video latent T={latent_t} < video_condition_frames={vcf}; need more video frames." ) video_generator = None if seed is None else torch.Generator(device=rand_device).manual_seed(seed) value_generator = None if seed is None else torch.Generator(device=rand_device).manual_seed(seed) # Build latents by concatenating GT condition part + pure-noise prediction part. # Video: first `vcf` latent frames from encoded GT input, rest from noise. input_image = input_image.to(device=self.device, dtype=self.torch_dtype) cond_latents = self._encode_input_image_latents_tensor(input_image=input_image, tiled=tiled) if cond_latents.shape[2] < vcf: raise ValueError( f"Encoded condition latents T={cond_latents.shape[2]} < video_condition_frames={vcf}." ) noise_video = torch.randn( (batch, self.vae.model.z_dim, latent_t, latent_h, latent_w), generator=video_generator, device=rand_device, dtype=torch.float32, ).to(device=self.device, dtype=self.torch_dtype) latents_video = torch.cat( [cond_latents[:, :, :vcf], noise_video[:, :, vcf:]], dim=2 ) # Value: first `vcf_v` timesteps from GT value, rest from noise. noise_value = torch.randn( (batch, value_horizon, self.value_expert.value_dim), generator=value_generator, device=rand_device, dtype=torch.float32, ).to(device=self.device, dtype=self.torch_dtype) if vcf_v > 0: latents_value = torch.cat([value[:, :vcf_v], noise_value[:, vcf_v:]], dim=1) else: latents_value = noise_value fuse_flag = bool(getattr(self.video_expert, "fuse_vae_embedding_in_latents", False)) use_prompt = prompt is not None use_context = context is not None or context_mask is not None if use_prompt and use_context: raise ValueError("`prompt` and `context/context_mask` are mutually exclusive.") if not use_prompt and not use_context: raise ValueError("Either `prompt` or both `context/context_mask` must be provided.") if use_prompt: context, context_mask = self.encode_prompt(prompt) else: if context is None or context_mask is None: raise ValueError("`context` and `context_mask` must be both provided together.") if context.ndim == 2: context = context.unsqueeze(0) if context_mask.ndim == 1: context_mask = context_mask.unsqueeze(0) if context.ndim != 3 or context_mask.ndim != 2: raise ValueError( f"`context/context_mask` must be [B,L,D]/[B,L], got {tuple(context.shape)} and {tuple(context_mask.shape)}" ) context = context.to(device=self.device, dtype=self.torch_dtype, non_blocking=True) context_mask = context_mask.to(device=self.device, dtype=torch.bool, non_blocking=True) if int(context.shape[0]) == 1 and batch > 1: context = context.expand(batch, -1, -1).contiguous() context_mask = context_mask.expand(batch, -1).contiguous() elif int(context.shape[0]) != batch: raise ValueError( f"`context` batch {context.shape[0]} does not match input_image batch {batch}" ) infer_timesteps_video, infer_deltas_video = self.infer_video_scheduler.build_inference_schedule( num_inference_steps=num_inference_steps, device=self.device, dtype=latents_video.dtype, shift_override=sigma_shift, ) infer_timesteps_value, infer_deltas_value = self.infer_value_scheduler.build_inference_schedule( num_inference_steps=num_inference_steps, device=self.device, dtype=latents_value.dtype, shift_override=sigma_shift, ) for step_t_video, step_delta_video, step_t_value, step_delta_value in zip( infer_timesteps_video, infer_deltas_video, infer_timesteps_value, infer_deltas_value, ): timestep_video = step_t_video.reshape(()).expand(batch).to(dtype=latents_video.dtype, device=self.device) timestep_value = step_t_value.reshape(()).expand(batch).to(dtype=latents_value.dtype, device=self.device) pred_video_posi, pred_value_posi = self._predict_joint_noise( latents_video=latents_video, latents_value=latents_value, timestep_video=timestep_video, timestep_value=timestep_value, context=context, context_mask=context_mask, fuse_vae_embedding_in_latents=fuse_flag, gt_value=value, video_condition_frames=vcf, ) pred_video = pred_video_posi pred_value = pred_value_posi latents_video = self.infer_video_scheduler.step(pred_video, step_delta_video, latents_video) latents_value = self.infer_value_scheduler.step(pred_value, step_delta_value, latents_value) # Re-fix condition parts after each denoising step (scheduler updates all positions). if vcf >= 1: latents_video[:, :, :vcf] = cond_latents[:, :, :vcf] if vcf_v >= 1: latents_value[:, :vcf_v] = value[:, :vcf_v] # Returns the full trajectory: first `vcf_v` timesteps are GT condition, rest are predicted. # absolute: both channels denorm [-1,1] → [0,1]; relative: progress stays # latent-scale (predicted delta in [-1,1]), success still denorms. value_out = value_from_latent(latents_value, progress_repr) if batch == 1: value_out = value_out[0] if test_value_with_infer_value: if not torch.allclose(value_out, value_only_out, atol=1e-2, rtol=1e-2): max_abs_diff = (value_out - value_only_out).abs().max().item() logger.warning( f"Value from infer_joint and infer_value differ with max abs diff {max_abs_diff:.6f}. " ) return { "video": self._decode_latents(latents_video, tiled=tiled), "value": value_out, } @torch.no_grad() def infer( self, prompt: str | None, input_image: torch.Tensor, num_frames: int, value: torch.Tensor | None = None, value_horizon: int | None = None, context: torch.Tensor | None = None, context_mask: torch.Tensor | None = None, negative_prompt: str | None = None, text_cfg_scale: float = 5.0, value_cfg_scale: float = 1.0, num_inference_steps: int = 20, sigma_shift: float | None = None, seed: int | None = None, rand_device: str = "cpu", tiled: bool = False, video_condition_frames: int | None = None, value_condition_frames: int | None = None, value_progress_repr: str = "absolute", ): return self.infer_joint( prompt=prompt, input_image=input_image, num_video_frames=num_frames, value_horizon=value_horizon, value=value, context=context, context_mask=context_mask, negative_prompt=negative_prompt, text_cfg_scale=text_cfg_scale, num_inference_steps=num_inference_steps, sigma_shift=sigma_shift, seed=seed, rand_device=rand_device, tiled=tiled, video_condition_frames=video_condition_frames, value_condition_frames=value_condition_frames, value_progress_repr=value_progress_repr, ) @torch.no_grad() def infer_value( self, prompt: str | Sequence[str] | None, input_image: torch.Tensor, value_horizon: int, gt_value: torch.Tensor | None = None, context: torch.Tensor | None = None, context_mask: torch.Tensor | None = None, num_inference_steps: int = 20, sigma_shift: float | None = None, seed: int | None = None, rand_device: str = "cpu", tiled: bool = False, video_condition_frames: int | None = None, value_condition_frames: int | None = None, value_progress_repr: str = "absolute", ) -> dict[str, Any]: self.eval() if str(getattr(self.video_expert, "video_attention_mask_mode", "")) != "first_frame_causal": raise ValueError( "`infer_value` requires `video_attention_mask_mode='first_frame_causal'`." ) vcf = int(self.video_condition_frames) if video_condition_frames is None else int(video_condition_frames) vcf_v = int(self.value_condition_frames) if value_condition_frames is None else int(value_condition_frames) progress_repr = parse_value_progress_repr(value_progress_repr) if vcf < 1: raise ValueError(f"`video_condition_frames` must be >= 1, got {vcf}") if vcf_v < 1: raise ValueError(f"`value_condition_frames` must be >= 1, got {vcf_v}") if value_horizon < vcf_v: raise ValueError( f"`value_horizon`={value_horizon} < value_condition_frames={vcf_v}; need at least vcf_v timesteps." ) images = self._coerce_input_images(input_image) batch = int(images.shape[0]) num_cond_frames = int(images.shape[1]) height, width = int(images.shape[3]), int(images.shape[4]) input_image = images if batch > 1 else images[0] prompt = self._coerce_prompt(prompt, batch) if height % 16 != 0 or width % 16 != 0: raise ValueError( f"`input_image` must be resized before infer, expected multiples of 16 but got HxW=({height},{width})" ) # Validate provided condition frames match vcf (in latent frames). tf = int(self.vae.temporal_downsample_factor) cond_latent_frames = (num_cond_frames - 1) // tf + 1 if tf > 0 else num_cond_frames if cond_latent_frames != vcf: raise ValueError( f"Provided {num_cond_frames} condition frame(s) -> {cond_latent_frames} latent frame(s), " f"but `video_condition_frames`={vcf}. They must match." ) gt_value = self._coerce_value(gt_value, value_horizon, batch, "gt_value") if gt_value is not None: gt_value = gt_value.to(device=self.device, dtype=self.torch_dtype) gt_value = value_to_latent(gt_value, progress_repr) elif vcf_v > 0: raise ValueError( f"`value_condition_frames`={vcf_v} > 0 but no `gt_value` provided for value conditioning." ) # Build value latents by concatenating GT condition part + pure-noise prediction part. generator = None if seed is None else torch.Generator(device=rand_device).manual_seed(seed) noise_value = torch.randn( (batch, value_horizon, self.value_expert.value_dim), generator=generator, device=rand_device, dtype=torch.float32, ).to(device=self.device, dtype=self.torch_dtype) if vcf_v > 0: latents_value = torch.cat([gt_value[:, :vcf_v], noise_value[:, vcf_v:]], dim=1) else: latents_value = noise_value # Encode video condition frames and build video KV cache (video is conditioning only here). input_image = input_image.to(device=self.device, dtype=self.torch_dtype) cond_latents = self._encode_input_image_latents_tensor(input_image=input_image, tiled=tiled) if cond_latents.shape[2] < vcf: raise ValueError( f"Encoded condition latents T={cond_latents.shape[2]} < video_condition_frames={vcf}." ) cond_latents = cond_latents[:, :, :vcf] fuse_flag = bool(getattr(self.video_expert, "fuse_vae_embedding_in_latents", False)) use_prompt = prompt is not None use_context = context is not None or context_mask is not None if use_prompt and use_context: raise ValueError("`prompt` and `context/context_mask` are mutually exclusive.") if not use_prompt and not use_context: raise ValueError("Either `prompt` or both `context/context_mask` must be provided.") if use_prompt: context, context_mask = self.encode_prompt(prompt) else: if context is None or context_mask is None: raise ValueError("`context` and `context_mask` must be both provided together.") if context.ndim == 2: context = context.unsqueeze(0) if context_mask.ndim == 1: context_mask = context_mask.unsqueeze(0) if context.ndim != 3 or context_mask.ndim != 2: raise ValueError( f"`context/context_mask` must be [B,L,D]/[B,L], got {tuple(context.shape)} and {tuple(context_mask.shape)}" ) context = context.to(device=self.device, dtype=self.torch_dtype, non_blocking=True) context_mask = context_mask.to(device=self.device, dtype=torch.bool, non_blocking=True) if int(context.shape[0]) == 1 and batch > 1: context = context.expand(batch, -1, -1).contiguous() context_mask = context_mask.expand(batch, -1).contiguous() elif int(context.shape[0]) != batch: raise ValueError( f"`context` batch {context.shape[0]} does not match input_image batch {batch}" ) timestep_video = torch.zeros( (cond_latents.shape[0],), dtype=cond_latents.dtype, device=self.device, ) video_pre = self.video_expert.pre_dit( x=cond_latents, timestep=timestep_video, context=context, context_mask=context_mask, action=None, fuse_vae_embedding_in_latents=fuse_flag, condition_frames=vcf, ) video_seq_len = int(video_pre["tokens"].shape[1]) attention_mask = self._build_mot_attention_mask( video_seq_len=video_seq_len, value_seq_len=latents_value.shape[1], video_tokens_per_frame=int(video_pre["meta"]["tokens_per_frame"]), device=video_pre["tokens"].device, video_condition_frames=vcf, ) video_kv_cache = self.mot.prefill_video_cache( video_tokens=video_pre["tokens"], video_freqs=video_pre["freqs"], video_t_mod=video_pre["t_mod"], video_context_payload={ "context": video_pre["context"], "mask": video_pre["context_mask"], }, video_attention_mask=attention_mask[:video_seq_len, :video_seq_len], ) infer_timesteps_value, infer_deltas_value = self.infer_value_scheduler.build_inference_schedule( num_inference_steps=num_inference_steps, device=self.device, dtype=latents_value.dtype, shift_override=sigma_shift, ) for step_t_value, step_delta_value in zip(infer_timesteps_value, infer_deltas_value): timestep_value = step_t_value.reshape(()).expand(batch).to(dtype=latents_value.dtype, device=self.device) pred_value_posi = self._predict_value_noise_with_cache( latents_value=latents_value, timestep_value=timestep_value, context=context, context_mask=context_mask, video_kv_cache=video_kv_cache, attention_mask=attention_mask, video_seq_len=video_seq_len, value_condition_frames=vcf_v, ) pred_value = pred_value_posi latents_value = self.infer_value_scheduler.step(pred_value, step_delta_value, latents_value) # Re-fix value condition timesteps after each denoising step (scheduler updates all positions). if vcf_v > 0: latents_value[:, :vcf_v] = gt_value[:, :vcf_v] value_out = value_from_latent(latents_value, progress_repr) if batch == 1: value_out = value_out[0] return { # absolute: both channels denorm [-1,1] → [0,1]; relative: progress stays # latent-scale (predicted delta in [-1,1]), success still denorms. # Returns the full trajectory: first `vcf_v` timesteps are GT condition, rest are predicted. "value": value_out, } def load_checkpoint(self, path): payload = torch.load(path, map_location="cpu") if "mot" in payload: self.mot.load_state_dict(payload["mot"], strict=False) elif "dit" in payload: logger.warning("Loading legacy `dit` checkpoint into video expert only.") self.video_expert.load_state_dict(payload["dit"], strict=False) else: raise ValueError(f"Checkpoint missing both `mot` and `dit` keys: {path}") return payload