| """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) |
|
|
|
|
|
|
| |
|
|
| """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 |
|
|
| |
|
|
| |
| |
| 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 |
|
|
| |
|
|
| |
| |
| 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 |
|
|
| |
|
|
| |
| |
| 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() |
|
|
| |
|
|
| |
| |
| 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 |
|
|
| |
|
|
| |
| |
| 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): |
| |
| 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): |
| |
| 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) |
| 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_num_query_tokens = num_temporal_groups * num_query_per_group |
| total_num_key_tokens = num_temporal_groups * num_key_per_group |
|
|
| |
| query_time_indices = torch.arange(num_temporal_groups).repeat_interleave(num_query_per_group) |
| key_time_indices = torch.arange(num_temporal_groups).repeat_interleave(num_key_per_group) |
|
|
| |
| query_time_indices = query_time_indices.unsqueeze(1) |
| key_time_indices = key_time_indices.unsqueeze(0) |
|
|
| if mode == "causal": |
| |
| attn_mask = query_time_indices >= key_time_indices |
| elif mode == "group_diagonal": |
| |
| attn_mask = query_time_indices == key_time_indices |
|
|
| 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) |
| |
| |
|
|
| 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) |
| |
| |
|
|
| 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) |
| has_seq = len(t_mod.shape) == 4 |
| chunk_dim = 2 if has_seq else 1 |
| |
| 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: |
| |
| 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": |
| |
| |
| 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) |
| |
| 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) |
| |
| |
| 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) |
| 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) |
| action_pos_embed = sinusoidal_embedding_1d(self.hidden_dim, |
| torch.arange(action_len, device=action_emb.device)) |
| action_emb = action_emb + action_pos_embed.unsqueeze(0) |
| context = torch.cat([context, action_emb], dim=1) |
|
|
| |
| num_temporal_groups = f - 1 |
| 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}" |
| |
| 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) |
|
|
| seq_len = f * h * w |
| final_context_mask = torch.zeros((batch_size, seq_len, context.shape[1]), dtype=torch.bool, device=context.device) |
| |
| final_context_mask[:, :, :context_len] = context_mask.unsqueeze(1).expand(-1, seq_len, -1) |
| |
| final_context_mask[:, tokens_per_frame:, context_len:] = action_group_mask.unsqueeze(0).expand(batch_size, -1, -1) |
| 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) |
| else: |
| context_mask = context_mask.unsqueeze(1).expand(-1, f * h * w, -1) |
|
|
| 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 |
|
|
| 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) |
|
|
| |
|
|
| |
| |
| 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): |
| |
| b, n, s, _, device = *x.size(), x.device |
| assert s % block_size == 0 |
| num_blocks = s // block_size |
|
|
| |
| 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 |
|
|
| |
| 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_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 |
|
|
| |
| 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 |
|
|
| |
| 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_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 |
|
|
| |
| self.norm = RMS_norm(dim) |
| self.to_qkv = nn.Conv2d(dim, dim * 3, 1) |
| self.proj = nn.Conv2d(dim, dim, 1) |
|
|
| |
| 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) |
| |
| q, k, v = self.to_qkv(x).reshape(b * t, 1, c * 3, -1).permute( |
| 0, 1, 3, 2).contiguous().chunk(3, dim=-1) |
|
|
| |
| x = F.scaled_dot_product_attention( |
| q, |
| k, |
| v, |
| |
| ) |
| x = x.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w) |
|
|
| |
| 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__() |
|
|
| |
| self.avg_shortcut = AvgDown3D( |
| in_dim, |
| out_dim, |
| factor_t=2 if temperal_downsample else 1, |
| factor_s=2 if down_flag else 1, |
| ) |
|
|
| |
| downsamples = [] |
| for _ in range(mult): |
| downsamples.append(ResidualBlock(in_dim, out_dim, dropout)) |
| in_dim = out_dim |
|
|
| |
| 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__() |
| |
| 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 |
|
|
| |
| upsamples = [] |
| for _ in range(mult): |
| upsamples.append(ResidualBlock(in_dim, out_dim, dropout)) |
| in_dim = out_dim |
|
|
| |
| 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 |
|
|
| |
| dims = [dim * u for u in [1] + dim_mult] |
| scale = 1.0 |
|
|
| |
| self.conv1 = CausalConv3d(3, dims[0], 3, padding=1) |
|
|
| |
| downsamples = [] |
| for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): |
| |
| 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 |
|
|
| |
| 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) |
|
|
| |
| self.middle = nn.Sequential(ResidualBlock(out_dim, out_dim, dropout), |
| AttentionBlock(out_dim), |
| ResidualBlock(out_dim, out_dim, dropout)) |
|
|
| |
| 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) |
|
|
| |
| 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) |
|
|
| |
| 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) |
|
|
| |
| 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 |
|
|
|
|
| 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 |
|
|
| |
| dims = [dim * u for u in [1] + dim_mult] |
| scale = 1.0 |
|
|
| |
| self.conv1 = CausalConv3d(12, dims[0], 3, padding=1) |
|
|
| |
| 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) |
|
|
| |
| self.middle = nn.Sequential( |
| ResidualBlock(out_dim, out_dim, dropout), |
| AttentionBlock(out_dim), |
| ResidualBlock(out_dim, out_dim, dropout), |
| ) |
|
|
| |
| 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) |
|
|
| |
| 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) |
|
|
| |
| 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) |
|
|
| |
| 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 |
|
|
| |
| dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]] |
| scale = 1.0 / 2**(len(dim_mult) - 2) |
|
|
| |
| self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1) |
|
|
| |
| self.middle = nn.Sequential(ResidualBlock(dims[0], dims[0], dropout), |
| AttentionBlock(dims[0]), |
| ResidualBlock(dims[0], dims[0], dropout)) |
|
|
| |
| upsamples = [] |
| for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])): |
| |
| 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 |
|
|
| |
| 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) |
|
|
| |
| 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]): |
| |
| 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) |
|
|
| |
| 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) |
|
|
| |
| 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 |
|
|
|
|
|
|
| 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 |
|
|
| |
| dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]] |
| scale = 1.0 / 2 ** (len(dim_mult) - 2) |
| |
| self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1) |
|
|
| |
| self.middle = nn.Sequential(ResidualBlock(dims[0], dims[0], dropout), |
| AttentionBlock(dims[0]), |
| ResidualBlock(dims[0], dims[0], dropout)) |
|
|
| |
| 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) |
|
|
| |
| 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) |
|
|
| |
| 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) |
|
|
| |
| 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] |
|
|
| |
| 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() |
| |
| 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() |
| |
| 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) |
| 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 |
| |
| 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] |
|
|
| |
| 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 |
|
|
| |
| 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 |
|
|
| |
| 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)): |
| |
| 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] |
|
|
| |
| 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] |
|
|
| |
| 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 |
|
|
| |
|
|
| |
| |
| 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 |
|
|
| |
| 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. |
| """ |
| |
| context = x if context is None else context |
| b, n, c = x.size(0), self.num_heads, self.head_dim |
|
|
| |
| 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) |
|
|
| |
| 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) |
|
|
| |
| 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) |
|
|
| |
| 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 |
|
|
| |
| 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 |
|
|
| |
| 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 |
|
|
| |
| self.embedding = nn.Embedding(num_buckets, num_heads) |
|
|
| def forward(self, lq, lk): |
| device = self.embedding.weight.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) |
| return rel_pos_embeds.contiguous() |
|
|
| def _relative_position_bucket(self, rel_pos): |
| |
| 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)) |
|
|
| |
| 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 |
|
|
| |
| 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) |
|
|
| |
| |
| |
|
|
| 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 |
|
|
| |
| 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) |
|
|
| |
| _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) |
|
|
| |
| if isinstance(sequence, str): |
| sequence = [sequence] |
| if self.clean: |
| sequence = [self._clean(u) for u in sequence] |
| ids = self.tokenizer(sequence, **_kwargs) |
|
|
| |
| 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 |
|
|
| |
|
|
| |
| |
| 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: |
| |
| 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) |
| 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]}." |
| ) |
|
|
| |
| 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)) |
|
|
| 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) |
|
|
| |
|
|
| |
| |
|
|
| 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: |
| |
| 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] |
| |
| ( |
| 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, |
| ) |
| |
| mixed = self._mixed_attention( |
| q_cat=q, |
| k_cat=k, |
| v_cat=v, |
| attention_mask=video_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=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}" |
| ) |
| |
| 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] |
| |
| ( |
| 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}." |
| ) |
|
|
| |
| 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, |
| } |
|
|
| |
| 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): |
| |
| 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 |
|
|
| |
|
|
| |
| |
| 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 = [ |
| { |
| |
| "model_hash": "9c8818c2cbea55eca56c7b447df170da", |
| "model_name": "wan_video_text_encoder", |
| "model_class": WanTextEncoder, |
| }, |
| { |
| |
| "model_hash": "1f5ab7703c6fc803fdded85ff040c316", |
| "model_name": "wan_video_dit", |
| "model_class": WanVideoDiT, |
| }, |
| { |
| |
| "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, |
| ) |
|
|
| |
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| 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)): |
| |
| |
| |
| |
| |
| |
| images = self._coerce_input_images(input_image) |
| videos = [ |
| images[b].to(device=self.device).permute(1, 0, 2, 3) |
| 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) |
|
|
| |
| 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, |
| ) |
| |
| mask[video_seq_len:, video_seq_len:] = True |
| |
| |
| 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, |
| 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]: |
| |
| |
| 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]) |
| |
| |
| _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}" |
| ) |
| |
| 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) |
| |
| |
| 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 |
| ) |
| |
| 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) |
| |
| if vcf >= 1: |
| latents_video[:, :, :vcf] = cond_latents[:, :, :vcf] |
| if vcf_v >= 1: |
| latents_value[:, :vcf_v] = value[:, :vcf_v] |
|
|
| |
| |
| |
| 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})" |
| ) |
| |
| 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." |
| ) |
|
|
| |
| 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 |
|
|
| |
| 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) |
| |
| 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 { |
| |
| |
| |
| "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 |
|
|
