diff --git a/fastvideo/demo/gradio_web_demo.py b/fastvideo/demo/gradio_web_demo.py index bb4bd810d..2c663c563 100644 --- a/fastvideo/demo/gradio_web_demo.py +++ b/fastvideo/demo/gradio_web_demo.py @@ -1,13 +1,15 @@ import gradio as gr import torch -from fastvideo.model.pipeline_mochi import MochiPipeline -from fastvideo.model.modeling_mochi import MochiTransformer3DModel +from fastvideo.models.mochi_hf.pipeline_mochi import MochiPipeline +from fastvideo.models.mochi_hf.modeling_mochi import MochiTransformer3DModel +from fastvideo.models.mochi_genmo.mochi_preview.dit.joint_model.asymm_models_joint import AsymmDiTJoint from diffusers import FlowMatchEulerDiscreteScheduler from diffusers.utils import export_to_video from fastvideo.distill.solver import PCMFMScheduler import tempfile import os import argparse +from safetensors.torch import load_file def init_args(): parser = argparse.ArgumentParser() @@ -36,10 +38,21 @@ def load_model(args): else: scheduler = PCMFMScheduler(1000, args.shift, args.num_euler_timesteps, False, args.linear_threshold, args.linear_range) - if args.transformer_path: - transformer = MochiTransformer3DModel.from_pretrained(args.transformer_path) + mochi_genmo = True + if mochi_genmo: + model_path = "/root/weights/dit.safetensors" + state_dcit = load_file(model_path) + transformer = AsymmDiTJoint() + transformer.load_state_dict(state_dcit) + # from IPython import embed + # embed() + transformer.config.in_channels = 12 + print("load gennmo mochi successfully") else: - transformer = MochiTransformer3DModel.from_pretrained(args.model_path, subfolder='transformer/') + if args.transformer_path: + transformer = MochiTransformer3DModel.from_pretrained(args.transformer_path) + else: + transformer = MochiTransformer3DModel.from_pretrained(args.model_path, subfolder='transformer/') pipe = MochiPipeline.from_pretrained(args.model_path, transformer=transformer, scheduler=scheduler) pipe.enable_vae_tiling() diff --git a/fastvideo/models/mochi_genmo/lib/attn_imports.py b/fastvideo/models/mochi_genmo/lib/attn_imports.py new file mode 100644 index 000000000..4d4a8c101 --- /dev/null +++ b/fastvideo/models/mochi_genmo/lib/attn_imports.py @@ -0,0 +1,29 @@ +from contextlib import contextmanager + +import torch + + +try: + from flash_attn import flash_attn_varlen_func as flash_varlen_attn +except ImportError: + flash_varlen_attn = None + +try: + from sageattention import sageattn as sage_attn +except ImportError: + sage_attn = None + +from torch.nn.attention import SDPBackend, sdpa_kernel + +training_backends = [SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION] +eval_backends = list(training_backends) +if torch.cuda.get_device_properties(0).major >= 9.0: + # Enable fast CuDNN attention on Hopper. + # This gives NaN on the backward pass for some reason, + # so only use it for evaluation. + eval_backends.append(SDPBackend.CUDNN_ATTENTION) + +@contextmanager +def sdpa_attn_ctx(training: bool = False): + with sdpa_kernel(training_backends if training else eval_backends): + yield diff --git a/fastvideo/models/mochi_genmo/lib/progress.py b/fastvideo/models/mochi_genmo/lib/progress.py new file mode 100644 index 000000000..6477c6a15 --- /dev/null +++ b/fastvideo/models/mochi_genmo/lib/progress.py @@ -0,0 +1,87 @@ +import contextlib +from typing import Any, Iterable, Iterator, Optional + +try: + from tqdm import tqdm +except ImportError: + tqdm = None + +try: + from ray.experimental.tqdm_ray import tqdm as ray_tqdm +except: + ray_tqdm = None + +# Global state +_current_progress_type = "none" +_is_progress_bar_active = False + + +class DummyProgressBar: + """A no-op progress bar that mimics tqdm interface""" + + def __init__(self, iterable=None, **kwargs): + self.iterable = iterable + + def __iter__(self): + return iter(self.iterable) + + def update(self, n=1): + pass + + def close(self): + pass + + def set_description(self, desc): + pass + + +def get_new_progress_bar(iterable: Optional[Iterable] = None, **kwargs) -> Any: + if not _is_progress_bar_active: + return DummyProgressBar(iterable=iterable, **kwargs) + + if _current_progress_type == "tqdm": + if tqdm is None: + raise ImportError("tqdm is required but not installed. Please install tqdm to use the tqdm progress bar.") + return tqdm(iterable=iterable, **kwargs) + elif _current_progress_type == "ray_tqdm": + if ray_tqdm is None: + raise ImportError("ray is required but not installed. Please install ray to use the ray_tqdm progress bar.") + return ray_tqdm(iterable=iterable, **kwargs) + return DummyProgressBar(iterable=iterable, **kwargs) + + +@contextlib.contextmanager +def progress_bar(type: str = "none", enabled=True): + """ + Context manager for setting progress bar type and options. + + Args: + type: Type of progress bar ("none" or "tqdm") + **options: Options to pass to the progress bar (e.g., total, desc) + + Raises: + ValueError: If progress bar type is invalid + RuntimeError: If progress bars are nested + + Example: + with progress_bar(type="tqdm", total=100): + for i in get_new_progress_bar(range(100)): + process(i) + """ + if type not in ("none", "tqdm", "ray_tqdm"): + raise ValueError("Progress bar type must be 'none' or 'tqdm' or 'ray_tqdm'") + if not enabled: + type = "none" + global _current_progress_type, _is_progress_bar_active + + if _is_progress_bar_active: + raise RuntimeError("Nested progress bars are not supported") + + _is_progress_bar_active = True + _current_progress_type = type + + try: + yield + finally: + _is_progress_bar_active = False + _current_progress_type = "none" diff --git a/fastvideo/models/mochi_genmo/lib/utils.py b/fastvideo/models/mochi_genmo/lib/utils.py new file mode 100644 index 000000000..b8ebbd611 --- /dev/null +++ b/fastvideo/models/mochi_genmo/lib/utils.py @@ -0,0 +1,67 @@ +import os +import subprocess +import tempfile +import time + +import numpy as np +from moviepy.editor import ImageSequenceClip +from PIL import Image + +from genmo.lib.progress import get_new_progress_bar + + +class Timer: + def __init__(self): + self.times = {} # Dictionary to store times per stage + + def __call__(self, name): + print(f"Timing {name}") + return self.TimerContextManager(self, name) + + def print_stats(self): + total_time = sum(self.times.values()) + # Print table header + print("{:<20} {:>10} {:>10}".format("Stage", "Time(s)", "Percent")) + for name, t in self.times.items(): + percent = (t / total_time) * 100 if total_time > 0 else 0 + print("{:<20} {:>10.2f} {:>9.2f}%".format(name, t, percent)) + + class TimerContextManager: + def __init__(self, outer, name): + self.outer = outer # Reference to the Timer instance + self.name = name + self.start_time = None + + def __enter__(self): + self.start_time = time.perf_counter() + return self + + def __exit__(self, exc_type, exc_value, traceback): + end_time = time.perf_counter() + elapsed = end_time - self.start_time + self.outer.times[self.name] = self.outer.times.get(self.name, 0) + elapsed + + +def save_video(final_frames, output_path, fps=30): + assert final_frames.ndim == 4 and final_frames.shape[3] == 3, f"invalid shape: {final_frames} (need t h w c)" + if final_frames.dtype != np.uint8: + final_frames = (final_frames * 255).astype(np.uint8) + ImageSequenceClip(list(final_frames), fps=fps).write_videofile(output_path) + + +def create_memory_tracker(): + import torch + + previous = [None] # Use list for mutable closure state + + def track(label="all2all"): + current = torch.cuda.memory_allocated() / 1e9 + if previous[0] is not None: + diff = current - previous[0] + sign = "+" if diff >= 0 else "" + print(f"GPU memory ({label}): {current:.2f} GB ({sign}{diff:.2f} GB)") + else: + print(f"GPU memory ({label}): {current:.2f} GB") + previous[0] = current # type: ignore + + return track diff --git a/fastvideo/models/mochi_genmo/mochi_preview/__init__.py b/fastvideo/models/mochi_genmo/mochi_preview/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/fastvideo/models/mochi_genmo/mochi_preview/dit/joint_model/__init__.py b/fastvideo/models/mochi_genmo/mochi_preview/dit/joint_model/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/fastvideo/models/mochi_genmo/mochi_preview/dit/joint_model/asymm_models_joint.py b/fastvideo/models/mochi_genmo/mochi_preview/dit/joint_model/asymm_models_joint.py new file mode 100644 index 000000000..e6d69dddd --- /dev/null +++ b/fastvideo/models/mochi_genmo/mochi_preview/dit/joint_model/asymm_models_joint.py @@ -0,0 +1,752 @@ +import os +from typing import Dict, List, Optional, Tuple, Any +import warnings +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange +from torch.nn.attention import sdpa_kernel +import fastvideo.models.mochi_genmo.mochi_preview.dit.joint_model.context_parallel as cp +from fastvideo.models.mochi_genmo.lib.attn_imports import flash_varlen_attn, sage_attn, sdpa_attn_ctx +from fastvideo.models.mochi_genmo.mochi_preview.dit.joint_model.layers import ( + FeedForward, + PatchEmbed, + RMSNorm, + TimestepEmbedder, +) +from fastvideo.models.mochi_genmo.mochi_preview.dit.joint_model.lora import LoraLinear +from fastvideo.models.mochi_genmo.mochi_preview.dit.joint_model.mod_rmsnorm import modulated_rmsnorm +from fastvideo.models.mochi_genmo.mochi_preview.dit.joint_model.residual_tanh_gated_rmsnorm import ( + residual_tanh_gated_rmsnorm, +) +from fastvideo.models.mochi_genmo.mochi_preview.dit.joint_model.rope_mixed import ( + compute_mixed_rotation, + create_position_matrix, +) +from fastvideo.models.mochi_genmo.mochi_preview.dit.joint_model.temporal_rope import apply_rotary_emb_qk_real +from fastvideo.models.mochi_genmo.mochi_preview.dit.joint_model.utils import ( + AttentionPool, + modulate, + pad_and_split_xy, +) +from fastvideo.models.mochi_genmo.mochi_preview.pipelines import compute_packed_indices +from diffusers.models.modeling_utils import ModelMixin +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.loaders import PeftAdapterMixin + +COMPILE_FINAL_LAYER = os.environ.get("COMPILE_DIT") == "1" +COMPILE_MMDIT_BLOCK = os.environ.get("COMPILE_DIT") == "1" + + +def ck(fn, *args, enabled=True, **kwargs) -> torch.Tensor: + if enabled: + return torch.utils.checkpoint.checkpoint(fn, *args, **kwargs, use_reentrant=False) + + return fn(*args, **kwargs) + + +class AsymmetricAttention(nn.Module): + def __init__( + self, + dim_x: int, + dim_y: int, + num_heads: int = 8, + qkv_bias: bool = False, + qk_norm: bool = True, + update_y: bool = True, + out_bias: bool = True, + attention_mode: str = "flash", + softmax_scale: Optional[float] = None, + device: Optional[torch.device] = None, + # Disable LoRA by default ... + qkv_proj_lora_rank: int = 0, + qkv_proj_lora_alpha: int = 0, + qkv_proj_lora_dropout: float = 0.0, + out_proj_lora_rank: int = 0, + out_proj_lora_alpha: int = 0, + out_proj_lora_dropout: float = 0.0, + ): + super().__init__() + self.attention_mode = attention_mode + self.dim_x = dim_x + self.dim_y = dim_y + self.num_heads = num_heads + self.head_dim = dim_x // num_heads + self.update_y = update_y + self.softmax_scale = softmax_scale + if dim_x % num_heads != 0: + raise ValueError(f"dim_x={dim_x} should be divisible by num_heads={num_heads}") + + # Input layers. + self.qkv_bias = qkv_bias + qkv_lora_kwargs = dict( + bias=qkv_bias, + device=device, + r=qkv_proj_lora_rank, + lora_alpha=qkv_proj_lora_alpha, + lora_dropout=qkv_proj_lora_dropout, + ) + self.qkv_x = LoraLinear(dim_x, 3 * dim_x, **qkv_lora_kwargs) + # Project text features to match visual features (dim_y -> dim_x) + self.qkv_y = LoraLinear(dim_y, 3 * dim_x, **qkv_lora_kwargs) + + # Query and key normalization for stability. + assert qk_norm + self.q_norm_x = RMSNorm(self.head_dim, device=device) + self.k_norm_x = RMSNorm(self.head_dim, device=device) + self.q_norm_y = RMSNorm(self.head_dim, device=device) + self.k_norm_y = RMSNorm(self.head_dim, device=device) + + # Output layers. y features go back down from dim_x -> dim_y. + proj_lora_kwargs = dict( + bias=out_bias, + device=device, + r=out_proj_lora_rank, + lora_alpha=out_proj_lora_alpha, + lora_dropout=out_proj_lora_dropout, + ) + self.proj_x = LoraLinear(dim_x, dim_x, **proj_lora_kwargs) + self.proj_y = LoraLinear(dim_x, dim_y, **proj_lora_kwargs) if update_y else nn.Identity() + + def run_qkv_y(self, y): + cp_rank, cp_size = cp.get_cp_rank_size() + local_heads = self.num_heads // cp_size + + if cp.is_cp_active(): + # Only predict local heads. + assert not self.qkv_bias + W_qkv_y = self.qkv_y.weight.view(3, self.num_heads, self.head_dim, self.dim_y) + W_qkv_y = W_qkv_y.narrow(1, cp_rank * local_heads, local_heads) + W_qkv_y = W_qkv_y.reshape(3 * local_heads * self.head_dim, self.dim_y) + qkv_y = F.linear(y, W_qkv_y, None) # (B, L, 3 * local_h * head_dim) + else: + qkv_y = self.qkv_y(y) # (B, L, 3 * dim) + + qkv_y = qkv_y.view(qkv_y.size(0), qkv_y.size(1), 3, local_heads, self.head_dim) + q_y, k_y, v_y = qkv_y.unbind(2) + + q_y = self.q_norm_y(q_y) + k_y = self.k_norm_y(k_y) + return q_y, k_y, v_y + + def prepare_qkv( + self, + x: torch.Tensor, # (B, M, dim_x) + y: torch.Tensor, # (B, L, dim_y) + *, + scale_x: torch.Tensor, + scale_y: torch.Tensor, + rope_cos: torch.Tensor, + rope_sin: torch.Tensor, + valid_token_indices: torch.Tensor, + max_seqlen_in_batch: int, + ): + # Process visual features + x = modulated_rmsnorm(x, scale_x) # (B, M, dim_x) where M = N / cp_group_size + qkv_x = self.qkv_x(x) # (B, M, 3 * dim_x) + assert qkv_x.dtype == torch.bfloat16 + + qkv_x = cp.all_to_all_collect_tokens(qkv_x, self.num_heads) # (3, B, N, local_h, head_dim) + + # Split qkv_x into q, k, v + q_x, k_x, v_x = qkv_x.unbind(0) # (B, N, local_h, head_dim) + q_x = self.q_norm_x(q_x) + q_x = apply_rotary_emb_qk_real(q_x, rope_cos, rope_sin) + k_x = self.k_norm_x(k_x) + k_x = apply_rotary_emb_qk_real(k_x, rope_cos, rope_sin) + + # Concatenate streams + B, N, num_heads, head_dim = q_x.size() + D = num_heads * head_dim + + # Process text features + if B == 1: + text_seqlen = max_seqlen_in_batch - N + if text_seqlen > 0: + y = y[:, :text_seqlen] # Remove padding tokens. + y = modulated_rmsnorm(y, scale_y) # (B, L, dim_y) + q_y, k_y, v_y = self.run_qkv_y(y) # (B, L, local_heads, head_dim) + + q = torch.cat([q_x, q_y], dim=1) + k = torch.cat([k_x, k_y], dim=1) + v = torch.cat([v_x, v_y], dim=1) + else: + q, k, v = q_x, k_x, v_x + else: + y = modulated_rmsnorm(y, scale_y) # (B, L, dim_y) + q_y, k_y, v_y = self.run_qkv_y(y) # (B, L, local_heads, head_dim) + + indices = valid_token_indices[:, None].expand(-1, D) + q = torch.cat([q_x, q_y], dim=1).view(-1, D).gather(0, indices) # (total, D) + k = torch.cat([k_x, k_y], dim=1).view(-1, D).gather(0, indices) # (total, D) + v = torch.cat([v_x, v_y], dim=1).view(-1, D).gather(0, indices) # (total, D) + + q = q.view(-1, num_heads, head_dim) + k = k.view(-1, num_heads, head_dim) + v = v.view(-1, num_heads, head_dim) + return q, k, v + + @torch.autocast("cuda", enabled=False) + def flash_attention(self, q, k, v, cu_seqlens, max_seqlen_in_batch, total, local_dim): + out: torch.Tensor = flash_varlen_attn( + q, k, v, + cu_seqlens_q=cu_seqlens, + cu_seqlens_k=cu_seqlens, + max_seqlen_q=max_seqlen_in_batch, + max_seqlen_k=max_seqlen_in_batch, + dropout_p=0.0, + softmax_scale=self.softmax_scale, + ) # (total, local_heads, head_dim) + return out.view(total, local_dim) + + def sdpa_attention(self, q, k, v): + with sdpa_attn_ctx(training=self.training): + out = F.scaled_dot_product_attention( + q, k, v, + attn_mask=None, + dropout_p=0.0, + is_causal=False, + ) + return out + + @torch.autocast("cuda", enabled=False) + def sage_attention(self, q, k, v): + return sage_attn(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=False) + + def run_attention( + self, + q: torch.Tensor, # (total <= B * (N + L), num_heads, head_dim) + k: torch.Tensor, # (total <= B * (N + L), num_heads, head_dim) + v: torch.Tensor, # (total <= B * (N + L), num_heads, head_dim) + *, + B: int, + cu_seqlens: Optional[torch.Tensor] = None, + max_seqlen_in_batch: Optional[int] = None, + ): + _, cp_size = cp.get_cp_rank_size() + assert self.num_heads % cp_size == 0 + local_heads = self.num_heads // cp_size + local_dim = local_heads * self.head_dim + + # Check shapes + assert q.ndim == 3 and k.ndim == 3 and v.ndim == 3 + total = q.size(0) + assert k.size(0) == total and v.size(0) == total + + if self.attention_mode == "flash": + out = self.flash_attention( + q, k, v, cu_seqlens, max_seqlen_in_batch, total, local_dim) # (total, local_dim) + else: + assert B == 1, \ + f"Non-flash attention mode {self.attention_mode} only supports batch size 1, got {B}" + + q = rearrange(q, "(b s) h d -> b h s d", b=B) + k = rearrange(k, "(b s) h d -> b h s d", b=B) + v = rearrange(v, "(b s) h d -> b h s d", b=B) + + if self.attention_mode == "sdpa": + out = self.sdpa_attention(q, k, v) # (B, local_heads, seq_len, head_dim) + elif self.attention_mode == "sage": + out = self.sage_attention(q, k, v) # (B, local_heads, seq_len, head_dim) + else: + raise ValueError(f"Unknown attention mode: {self.attention_mode}") + + out = rearrange(out, "b h s d -> (b s) (h d)") + + return out + + def post_attention( + self, + out: torch.Tensor, + B: int, + M: int, + L: int, + dtype: torch.dtype, + valid_token_indices: torch.Tensor, + ): + """ + Args: + out: (total <= B * (N + L), local_dim) + valid_token_indices: (total <= B * (N + L),) + B: Batch size + M: Number of visual tokens per context parallel rank + L: Number of text tokens + dtype: Data type of the input and output tensors + + Returns: + x: (B, N, dim_x) tensor of visual tokens where N = M * cp_size + y: (B, L, dim_y) tensor of text token features + """ + _, cp_size = cp.get_cp_rank_size() + local_heads = self.num_heads // cp_size + local_dim = local_heads * self.head_dim + N = M * cp_size + + # Split sequence into visual and text tokens, adding back padding. + if B == 1: + out = out.view(B, -1, local_dim) + if out.size(1) > N: + x, y = torch.tensor_split(out, (N,), dim=1) # (B, N, local_dim), (B, <= L, local_dim) + y = F.pad(y, (0, 0, 0, L - y.size(1))) # (B, L, local_dim) + else: + # Empty prompt. + x, y = out, out.new_zeros(B, L, local_dim) + else: + x, y = pad_and_split_xy(out, valid_token_indices, B, N, L, dtype) + assert x.size() == (B, N, local_dim) + assert y.size() == (B, L, local_dim) + + # Communicate across context parallel ranks. + x = x.view(B, N, local_heads, self.head_dim) + x = cp.all_to_all_collect_heads(x) # (B, M, dim_x = num_heads * head_dim) + if cp.is_cp_active(): + y = cp.all_gather(y) # (cp_size * B, L, local_heads * head_dim) + y = rearrange(y, "(G B) L D -> B L (G D)", G=cp_size, D=local_dim) # (B, L, dim_x) + + x = self.proj_x(x) + y = self.proj_y(y) + return x, y + + def forward( + self, + x: torch.Tensor, # (B, M, dim_x) + y: torch.Tensor, # (B, L, dim_y) + *, + scale_x: torch.Tensor, # (B, dim_x), modulation for pre-RMSNorm. + scale_y: torch.Tensor, # (B, dim_y), modulation for pre-RMSNorm. + packed_indices: Dict[str, torch.Tensor] = None, + checkpoint_qkv: bool = False, + checkpoint_post_attn: bool = False, + **rope_rotation, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Forward pass of asymmetric multi-modal attention. + + Args: + x: (B, M, dim_x) tensor of visual tokens + y: (B, L, dim_y) tensor of text token features + packed_indices: Dict with keys for Flash Attention + num_frames: Number of frames in the video. N = num_frames * num_spatial_tokens + + Returns: + x: (B, M, dim_x) tensor of visual tokens after multi-modal attention + y: (B, L, dim_y) tensor of text token features after multi-modal attention + """ + B, L, _ = y.shape + _, M, _ = x.shape + + # Predict a packed QKV tensor from visual and text features. + q, k, v = ck(self.prepare_qkv, + x=x, + y=y, + scale_x=scale_x, + scale_y=scale_y, + rope_cos=rope_rotation.get("rope_cos"), + rope_sin=rope_rotation.get("rope_sin"), + valid_token_indices=packed_indices["valid_token_indices_kv"], + max_seqlen_in_batch=packed_indices["max_seqlen_in_batch_kv"], + enabled=checkpoint_qkv, + ) # (total <= B * (N + L), 3, local_heads, head_dim) + + # Self-attention is expensive, so don't checkpoint it. + out = self.run_attention( + q, k, v, B=B, + cu_seqlens=packed_indices["cu_seqlens_kv"], + max_seqlen_in_batch=packed_indices["max_seqlen_in_batch_kv"], + ) + + x, y = ck(self.post_attention, + out, + B=B, M=M, L=L, + dtype=v.dtype, + valid_token_indices=packed_indices["valid_token_indices_kv"], + enabled=checkpoint_post_attn, + ) + + return x, y + + +@torch.compile(disable=not COMPILE_MMDIT_BLOCK) +class AsymmetricJointBlock(nn.Module): + def __init__( + self, + hidden_size_x: int, + hidden_size_y: int, + num_heads: int, + *, + mlp_ratio_x: float = 8.0, # Ratio of hidden size to d_model for MLP for visual tokens. + mlp_ratio_y: float = 4.0, # Ratio of hidden size to d_model for MLP for text tokens. + update_y: bool = True, # Whether to update text tokens in this block. + device: Optional[torch.device] = None, + **block_kwargs, + ): + super().__init__() + self.update_y = update_y + self.hidden_size_x = hidden_size_x + self.hidden_size_y = hidden_size_y + self.mod_x = nn.Linear(hidden_size_x, 4 * hidden_size_x, device=device) + if self.update_y: + self.mod_y = nn.Linear(hidden_size_x, 4 * hidden_size_y, device=device) + else: + self.mod_y = nn.Linear(hidden_size_x, hidden_size_y, device=device) + + # Self-attention: + self.attn = AsymmetricAttention( + hidden_size_x, + hidden_size_y, + num_heads=num_heads, + update_y=update_y, + device=device, + **block_kwargs, + ) + + # MLP. + mlp_hidden_dim_x = int(hidden_size_x * mlp_ratio_x) + assert mlp_hidden_dim_x == int(1536 * 8) + self.mlp_x = FeedForward( + in_features=hidden_size_x, + hidden_size=mlp_hidden_dim_x, + multiple_of=256, + ffn_dim_multiplier=None, + device=device, + ) + + # MLP for text not needed in last block. + if self.update_y: + mlp_hidden_dim_y = int(hidden_size_y * mlp_ratio_y) + self.mlp_y = FeedForward( + in_features=hidden_size_y, + hidden_size=mlp_hidden_dim_y, + multiple_of=256, + ffn_dim_multiplier=None, + device=device, + ) + + def forward( + self, + x: torch.Tensor, + c: torch.Tensor, + y: torch.Tensor, + # TODO: These could probably just go into attn_kwargs + checkpoint_ff: bool = False, + checkpoint_qkv: bool = False, + checkpoint_post_attn: bool = False, + **attn_kwargs, + ): + """Forward pass of a block. + + Args: + x: (B, N, dim) tensor of visual tokens + c: (B, dim) tensor of conditioned features + y: (B, L, dim) tensor of text tokens + num_frames: Number of frames in the video. N = num_frames * num_spatial_tokens + + Returns: + x: (B, N, dim) tensor of visual tokens after block + y: (B, L, dim) tensor of text tokens after block + """ + N = x.size(1) + + c = F.silu(c) + mod_x = self.mod_x(c) + scale_msa_x, gate_msa_x, scale_mlp_x, gate_mlp_x = mod_x.chunk(4, dim=1) + mod_y = self.mod_y(c) + + if self.update_y: + scale_msa_y, gate_msa_y, scale_mlp_y, gate_mlp_y = mod_y.chunk(4, dim=1) + else: + scale_msa_y = mod_y + + # Self-attention block. + x_attn, y_attn = self.attn( + x, + y, + scale_x=scale_msa_x, + scale_y=scale_msa_y, + checkpoint_qkv=checkpoint_qkv, + checkpoint_post_attn=checkpoint_post_attn, + **attn_kwargs, + ) + + assert x_attn.size(1) == N + x = residual_tanh_gated_rmsnorm(x, x_attn, gate_msa_x) + + if self.update_y: + y = residual_tanh_gated_rmsnorm(y, y_attn, gate_msa_y) + + # MLP block. + x = ck(self.ff_block_x, x, scale_mlp_x, gate_mlp_x, enabled=checkpoint_ff) + if self.update_y: + y = ck(self.ff_block_y, y, scale_mlp_y, gate_mlp_y, enabled=checkpoint_ff) # type: ignore + return x, y + + def ff_block_x(self, x, scale_x, gate_x): + x_mod = modulated_rmsnorm(x, scale_x) + x_res = self.mlp_x(x_mod) + x = residual_tanh_gated_rmsnorm(x, x_res, gate_x) # Sandwich norm + return x + + def ff_block_y(self, y, scale_y, gate_y): + y_mod = modulated_rmsnorm(y, scale_y) + y_res = self.mlp_y(y_mod) + y = residual_tanh_gated_rmsnorm(y, y_res, gate_y) # Sandwich norm + return y + + +@torch.compile(disable=not COMPILE_FINAL_LAYER) +class FinalLayer(nn.Module): + """ + The final layer of DiT. + """ + + def __init__( + self, + hidden_size, + patch_size, + out_channels, + device: Optional[torch.device] = None, + ): + super().__init__() + self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, device=device) + self.mod = nn.Linear(hidden_size, 2 * hidden_size, device=device) + self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, device=device) + + def forward(self, x, c): + c = F.silu(c) + shift, scale = self.mod(c).chunk(2, dim=1) + x = modulate(self.norm_final(x), shift, scale) + x = self.linear(x) + return x + + +class AsymmDiTJoint(ModelMixin, ConfigMixin, PeftAdapterMixin): + """ + Diffusion model with a Transformer backbone. + + Ingests text embeddings instead of a label. + """ + @register_to_config + def __init__( + self, + *, + patch_size=2, + in_channels=12, + hidden_size_x=3072, + hidden_size_y=1536, + depth=48, + num_heads=24, + mlp_ratio_x=4.0, + mlp_ratio_y=4.0, + t5_feat_dim: int = 4096, + t5_token_length: int = 256, + patch_embed_bias: bool = True, + timestep_mlp_bias: bool = True, + timestep_scale: float = 1000.0, + use_extended_posenc: bool = False, + rope_theta: float = 10000.0, + device: Optional[torch.device] = None, + **block_kwargs, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = in_channels + self.patch_size = patch_size + self.num_heads = num_heads + self.hidden_size_x = hidden_size_x + self.hidden_size_y = hidden_size_y + self.head_dim = hidden_size_x // num_heads # Head dimension and count is determined by visual. + self.use_extended_posenc = use_extended_posenc + self.t5_token_length = t5_token_length + self.t5_feat_dim = t5_feat_dim + self.rope_theta = rope_theta # Scaling factor for frequency computation for temporal RoPE. + self.timestep_scale = timestep_scale + + self.x_embedder = PatchEmbed( + patch_size=patch_size, + in_chans=in_channels, + embed_dim=hidden_size_x, + bias=patch_embed_bias, + device=device, + ) + # Conditionings + # Timestep + self.t_embedder = TimestepEmbedder(hidden_size_x, bias=timestep_mlp_bias, timestep_scale=timestep_scale) + + # Caption Pooling (T5) + self.t5_y_embedder = AttentionPool(t5_feat_dim, num_heads=8, output_dim=hidden_size_x, device=device) + + # Dense Embedding Projection (T5) + self.t5_yproj = nn.Linear(t5_feat_dim, hidden_size_y, bias=True, device=device) + + # Initialize pos_frequencies as an empty parameter. + self.pos_frequencies = nn.Parameter(torch.empty(3, self.num_heads, self.head_dim // 2, device=device)) + + # for depth 48: + # b = 0: AsymmetricJointBlock, update_y=True + # b = 1: AsymmetricJointBlock, update_y=True + # ... + # b = 46: AsymmetricJointBlock, update_y=True + # b = 47: AsymmetricJointBlock, update_y=False. No need to update text features. + blocks = [] + for b in range(depth): + # Joint multi-modal block + update_y = b < depth - 1 + block = AsymmetricJointBlock( + hidden_size_x, + hidden_size_y, + num_heads, + mlp_ratio_x=mlp_ratio_x, + mlp_ratio_y=mlp_ratio_y, + update_y=update_y, + device=device, + **block_kwargs, + ) + + blocks.append(block) + self.blocks = nn.ModuleList(blocks) + + self.final_layer = FinalLayer(hidden_size_x, patch_size, self.out_channels, device=device) + + def embed_x(self, x: torch.Tensor) -> torch.Tensor: + """ + Args: + x: (B, C=12, T, H, W) tensor of visual tokens + + Returns: + x: (B, C=3072, N) tensor of visual tokens with positional embedding. + """ + return self.x_embedder(x) # Convert BcTHW to BCN + + @torch.compile(disable=not COMPILE_MMDIT_BLOCK) + def prepare( + self, + x: torch.Tensor, + sigma: torch.Tensor, + t5_feat: torch.Tensor, + t5_mask: torch.Tensor, + ): + """Prepare input and conditioning embeddings.""" + + # Visual patch embeddings with positional encoding. + T, H, W = x.shape[-3:] + pH, pW = H // self.patch_size, W // self.patch_size + x = self.embed_x(x) # (B, N, D), where N = T * H * W / patch_size ** 2 + assert x.ndim == 3 + B = x.size(0) + + # Construct position array of size [N, 3]. + # pos[:, 0] is the frame index for each location, + # pos[:, 1] is the row index for each location, and + # pos[:, 2] is the column index for each location. + N = T * pH * pW + assert x.size(1) == N + pos = create_position_matrix(T, pH=pH, pW=pW, device=x.device, dtype=torch.float32) # (N, 3) + rope_cos, rope_sin = compute_mixed_rotation( + freqs=self.pos_frequencies, pos=pos + ) # Each are (N, num_heads, dim // 2) + + # Global vector embedding for conditionings. + c_t = self.t_embedder(1 - sigma) # (B, D) + + # Pool T5 tokens using attention pooler + # Note encoder_hidden_states[1] contains T5 token features. + assert ( + t5_feat.size(1) == self.t5_token_length + ), f"Expected L={self.t5_token_length}, got {t5_feat.shape} for encoder_hidden_states." + t5_y_pool = self.t5_y_embedder(t5_feat, t5_mask) # (B, D) + assert t5_y_pool.size(0) == B, f"Expected B={B}, got {t5_y_pool.shape} for t5_y_pool." + + c = c_t + t5_y_pool + + encoder_hidden_states = self.t5_yproj(t5_feat) # (B, L, t5_feat_dim) --> (B, L, D) + + return x, c, encoder_hidden_states, rope_cos, rope_sin + + def forward( + self, + hidden_states: torch.Tensor, # [1, 12, 7, 60, 106] + timestep: torch.Tensor, # [1] + encoder_hidden_states: torch.Tensor, # [1, 256, 4096] + encoder_attention_mask: torch.Tensor, # [1, 256] + attention_kwargs: Optional[Dict[str, Any]] = None, + return_dict: bool = False, + rope_cos: torch.Tensor = None, + rope_sin: torch.Tensor = None, + num_ff_checkpoint: int = 48, # 48 + num_qkv_checkpoint: int = 48, # 48 + num_post_attn_checkpoint: int = 0, # 0 + ): + """Forward pass of DiT. + + Args: + hidden_states: (B, C, T, H, W) tensor of spatial inputs (images or latent representations of images) + sigma: (B,) tensor of noise standard deviations + encoder_hidden_states: List((B, L, encoder_hidden_states_dim) tensor of caption token features. For SDXL text encoders: L=77, encoder_hidden_states_dim=2048) + encoder_attention_mask: List((B, L) boolean tensor indicating which tokens are not padding) + packed_indices: Dict with keys for Flash Attention. Result of compute_packed_indices. + {'cu_seqlens_kv': tensor([ 0, 11230], device='cuda:0', dtype=torch.int32), + 'max_seqlen_in_batch_kv': 11230, + 'valid_token_indices_kv': tensor([ 0, 1, 2, ..., 11227, 11228, 11229], device='cuda:0')} + """ + sigma = timestep / self.timestep_scale + num_latent_toks = np.prod(hidden_states.shape[-3:]) + + packed_indices = compute_packed_indices(hidden_states.device, encoder_attention_mask, int(num_latent_toks)) + + _, _, T, H, W = hidden_states.shape + + if self.pos_frequencies.dtype != torch.float32: + warnings.warn(f"pos_frequencies dtype {self.pos_frequencies.dtype} != torch.float32") + + # Use EFFICIENT_ATTENTION backend for T5 pooling, since we have a mask. + # Have to call sdpa_kernel outside of a torch.compile region. + with sdpa_kernel(torch.nn.attention.SDPBackend.EFFICIENT_ATTENTION): + hidden_states, c, encoder_hidden_states, rope_cos, rope_sin = self.prepare(hidden_states, sigma, encoder_hidden_states, encoder_attention_mask) # [1, 11130, 3072], [1, 3072], [1, 256, 1536], [11130, 24, 64] + del encoder_attention_mask + + cp_rank, cp_size = cp.get_cp_rank_size() + N = hidden_states.size(1) + M = N // cp_size + assert N % cp_size == 0, f"Visual sequence length ({hidden_states.shape[1]}) must be divisible by cp_size ({cp_size})." + + if cp_size > 1: + hidden_states = hidden_states.narrow(1, cp_rank * M, M) + + assert self.num_heads % cp_size == 0 + local_heads = self.num_heads // cp_size + rope_cos = rope_cos.narrow(1, cp_rank * local_heads, local_heads) + rope_sin = rope_sin.narrow(1, cp_rank * local_heads, local_heads) + + for i, block in enumerate(self.blocks): + hidden_states, encoder_hidden_states = block( # [1, 11130, 3072], [1, 256, 1536] + hidden_states, + c, + encoder_hidden_states, + rope_cos=rope_cos, + rope_sin=rope_sin, + packed_indices=packed_indices, + checkpoint_ff=i < num_ff_checkpoint, + checkpoint_qkv=i < num_qkv_checkpoint, + checkpoint_post_attn=i < num_post_attn_checkpoint, + ) # (B, M, D), (B, L, D) + del encoder_hidden_states # Final layers don't use dense text features. + + hidden_states = self.final_layer(hidden_states, c) # (B, M, patch_size ** 2 * out_channels) [1, 11130, 48] + + patch = hidden_states.size(2) + + hidden_states = cp.all_gather(hidden_states) + hidden_states = rearrange(hidden_states, "(G B) M P -> B (G M) P", G=cp_size, P=patch) # [1, 11130, 48] + hidden_states = rearrange( # [1, 12, 7, 60, 106] + hidden_states, + "B (T hp wp) (p1 p2 c) -> B c T (hp p1) (wp p2)", + T=T, + hp=H // self.patch_size, + wp=W // self.patch_size, + p1=self.patch_size, + p2=self.patch_size, + c=self.out_channels, + ) + + attn_outputs_list = None + return (-hidden_states, attn_outputs_list) diff --git a/fastvideo/models/mochi_genmo/mochi_preview/dit/joint_model/asymm_models_joint_raw.py b/fastvideo/models/mochi_genmo/mochi_preview/dit/joint_model/asymm_models_joint_raw.py new file mode 100644 index 000000000..1e83d37d9 --- /dev/null +++ b/fastvideo/models/mochi_genmo/mochi_preview/dit/joint_model/asymm_models_joint_raw.py @@ -0,0 +1,737 @@ +import os +from typing import Dict, List, Optional, Tuple +import warnings + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange +from torch.nn.attention import sdpa_kernel + +import genmo.mochi_preview.dit.joint_model.context_parallel as cp +from genmo.lib.attn_imports import flash_varlen_attn, sage_attn, sdpa_attn_ctx +from genmo.mochi_preview.dit.joint_model.layers import ( + FeedForward, + PatchEmbed, + RMSNorm, + TimestepEmbedder, +) +from genmo.mochi_preview.dit.joint_model.lora import LoraLinear +from genmo.mochi_preview.dit.joint_model.mod_rmsnorm import modulated_rmsnorm +from genmo.mochi_preview.dit.joint_model.residual_tanh_gated_rmsnorm import ( + residual_tanh_gated_rmsnorm, +) +from genmo.mochi_preview.dit.joint_model.rope_mixed import ( + compute_mixed_rotation, + create_position_matrix, +) +from genmo.mochi_preview.dit.joint_model.temporal_rope import apply_rotary_emb_qk_real +from genmo.mochi_preview.dit.joint_model.utils import ( + AttentionPool, + modulate, + pad_and_split_xy, +) + +COMPILE_FINAL_LAYER = os.environ.get("COMPILE_DIT") == "1" +COMPILE_MMDIT_BLOCK = os.environ.get("COMPILE_DIT") == "1" + + +def ck(fn, *args, enabled=True, **kwargs) -> torch.Tensor: + if enabled: + return torch.utils.checkpoint.checkpoint(fn, *args, **kwargs, use_reentrant=False) + + return fn(*args, **kwargs) + + +class AsymmetricAttention(nn.Module): + def __init__( + self, + dim_x: int, + dim_y: int, + num_heads: int = 8, + qkv_bias: bool = True, + qk_norm: bool = False, + update_y: bool = True, + out_bias: bool = True, + attention_mode: str = "flash", + softmax_scale: Optional[float] = None, + device: Optional[torch.device] = None, + # Disable LoRA by default ... + qkv_proj_lora_rank: int = 0, + qkv_proj_lora_alpha: int = 0, + qkv_proj_lora_dropout: float = 0.0, + out_proj_lora_rank: int = 0, + out_proj_lora_alpha: int = 0, + out_proj_lora_dropout: float = 0.0, + ): + super().__init__() + self.attention_mode = attention_mode + self.dim_x = dim_x + self.dim_y = dim_y + self.num_heads = num_heads + self.head_dim = dim_x // num_heads + self.update_y = update_y + self.softmax_scale = softmax_scale + if dim_x % num_heads != 0: + raise ValueError(f"dim_x={dim_x} should be divisible by num_heads={num_heads}") + + # Input layers. + self.qkv_bias = qkv_bias + qkv_lora_kwargs = dict( + bias=qkv_bias, + device=device, + r=qkv_proj_lora_rank, + lora_alpha=qkv_proj_lora_alpha, + lora_dropout=qkv_proj_lora_dropout, + ) + self.qkv_x = LoraLinear(dim_x, 3 * dim_x, **qkv_lora_kwargs) + # Project text features to match visual features (dim_y -> dim_x) + self.qkv_y = LoraLinear(dim_y, 3 * dim_x, **qkv_lora_kwargs) + + # Query and key normalization for stability. + assert qk_norm + self.q_norm_x = RMSNorm(self.head_dim, device=device) + self.k_norm_x = RMSNorm(self.head_dim, device=device) + self.q_norm_y = RMSNorm(self.head_dim, device=device) + self.k_norm_y = RMSNorm(self.head_dim, device=device) + + # Output layers. y features go back down from dim_x -> dim_y. + proj_lora_kwargs = dict( + bias=out_bias, + device=device, + r=out_proj_lora_rank, + lora_alpha=out_proj_lora_alpha, + lora_dropout=out_proj_lora_dropout, + ) + self.proj_x = LoraLinear(dim_x, dim_x, **proj_lora_kwargs) + self.proj_y = LoraLinear(dim_x, dim_y, **proj_lora_kwargs) if update_y else nn.Identity() + + def run_qkv_y(self, y): + cp_rank, cp_size = cp.get_cp_rank_size() + local_heads = self.num_heads // cp_size + + if cp.is_cp_active(): + # Only predict local heads. + assert not self.qkv_bias + W_qkv_y = self.qkv_y.weight.view(3, self.num_heads, self.head_dim, self.dim_y) + W_qkv_y = W_qkv_y.narrow(1, cp_rank * local_heads, local_heads) + W_qkv_y = W_qkv_y.reshape(3 * local_heads * self.head_dim, self.dim_y) + qkv_y = F.linear(y, W_qkv_y, None) # (B, L, 3 * local_h * head_dim) + else: + qkv_y = self.qkv_y(y) # (B, L, 3 * dim) + + qkv_y = qkv_y.view(qkv_y.size(0), qkv_y.size(1), 3, local_heads, self.head_dim) + q_y, k_y, v_y = qkv_y.unbind(2) + + q_y = self.q_norm_y(q_y) + k_y = self.k_norm_y(k_y) + return q_y, k_y, v_y + + def prepare_qkv( + self, + x: torch.Tensor, # (B, M, dim_x) + y: torch.Tensor, # (B, L, dim_y) + *, + scale_x: torch.Tensor, + scale_y: torch.Tensor, + rope_cos: torch.Tensor, + rope_sin: torch.Tensor, + valid_token_indices: torch.Tensor, + max_seqlen_in_batch: int, + ): + # Process visual features + x = modulated_rmsnorm(x, scale_x) # (B, M, dim_x) where M = N / cp_group_size + qkv_x = self.qkv_x(x) # (B, M, 3 * dim_x) + assert qkv_x.dtype == torch.bfloat16 + + qkv_x = cp.all_to_all_collect_tokens(qkv_x, self.num_heads) # (3, B, N, local_h, head_dim) + + # Split qkv_x into q, k, v + q_x, k_x, v_x = qkv_x.unbind(0) # (B, N, local_h, head_dim) + q_x = self.q_norm_x(q_x) + q_x = apply_rotary_emb_qk_real(q_x, rope_cos, rope_sin) + k_x = self.k_norm_x(k_x) + k_x = apply_rotary_emb_qk_real(k_x, rope_cos, rope_sin) + + # Concatenate streams + B, N, num_heads, head_dim = q_x.size() + D = num_heads * head_dim + + # Process text features + if B == 1: + text_seqlen = max_seqlen_in_batch - N + if text_seqlen > 0: + y = y[:, :text_seqlen] # Remove padding tokens. + y = modulated_rmsnorm(y, scale_y) # (B, L, dim_y) + q_y, k_y, v_y = self.run_qkv_y(y) # (B, L, local_heads, head_dim) + + q = torch.cat([q_x, q_y], dim=1) + k = torch.cat([k_x, k_y], dim=1) + v = torch.cat([v_x, v_y], dim=1) + else: + q, k, v = q_x, k_x, v_x + else: + y = modulated_rmsnorm(y, scale_y) # (B, L, dim_y) + q_y, k_y, v_y = self.run_qkv_y(y) # (B, L, local_heads, head_dim) + + indices = valid_token_indices[:, None].expand(-1, D) + q = torch.cat([q_x, q_y], dim=1).view(-1, D).gather(0, indices) # (total, D) + k = torch.cat([k_x, k_y], dim=1).view(-1, D).gather(0, indices) # (total, D) + v = torch.cat([v_x, v_y], dim=1).view(-1, D).gather(0, indices) # (total, D) + + q = q.view(-1, num_heads, head_dim) + k = k.view(-1, num_heads, head_dim) + v = v.view(-1, num_heads, head_dim) + return q, k, v + + @torch.autocast("cuda", enabled=False) + def flash_attention(self, q, k, v, cu_seqlens, max_seqlen_in_batch, total, local_dim): + out: torch.Tensor = flash_varlen_attn( + q, k, v, + cu_seqlens_q=cu_seqlens, + cu_seqlens_k=cu_seqlens, + max_seqlen_q=max_seqlen_in_batch, + max_seqlen_k=max_seqlen_in_batch, + dropout_p=0.0, + softmax_scale=self.softmax_scale, + ) # (total, local_heads, head_dim) + return out.view(total, local_dim) + + def sdpa_attention(self, q, k, v): + with sdpa_attn_ctx(training=self.training): + out = F.scaled_dot_product_attention( + q, k, v, + attn_mask=None, + dropout_p=0.0, + is_causal=False, + ) + return out + + @torch.autocast("cuda", enabled=False) + def sage_attention(self, q, k, v): + return sage_attn(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=False) + + def run_attention( + self, + q: torch.Tensor, # (total <= B * (N + L), num_heads, head_dim) + k: torch.Tensor, # (total <= B * (N + L), num_heads, head_dim) + v: torch.Tensor, # (total <= B * (N + L), num_heads, head_dim) + *, + B: int, + cu_seqlens: Optional[torch.Tensor] = None, + max_seqlen_in_batch: Optional[int] = None, + ): + _, cp_size = cp.get_cp_rank_size() + assert self.num_heads % cp_size == 0 + local_heads = self.num_heads // cp_size + local_dim = local_heads * self.head_dim + + # Check shapes + assert q.ndim == 3 and k.ndim == 3 and v.ndim == 3 + total = q.size(0) + assert k.size(0) == total and v.size(0) == total + + if self.attention_mode == "flash": + out = self.flash_attention( + q, k, v, cu_seqlens, max_seqlen_in_batch, total, local_dim) # (total, local_dim) + else: + assert B == 1, \ + f"Non-flash attention mode {self.attention_mode} only supports batch size 1, got {B}" + + q = rearrange(q, "(b s) h d -> b h s d", b=B) + k = rearrange(k, "(b s) h d -> b h s d", b=B) + v = rearrange(v, "(b s) h d -> b h s d", b=B) + + if self.attention_mode == "sdpa": + out = self.sdpa_attention(q, k, v) # (B, local_heads, seq_len, head_dim) + elif self.attention_mode == "sage": + out = self.sage_attention(q, k, v) # (B, local_heads, seq_len, head_dim) + else: + raise ValueError(f"Unknown attention mode: {self.attention_mode}") + + out = rearrange(out, "b h s d -> (b s) (h d)") + + return out + + def post_attention( + self, + out: torch.Tensor, + B: int, + M: int, + L: int, + dtype: torch.dtype, + valid_token_indices: torch.Tensor, + ): + """ + Args: + out: (total <= B * (N + L), local_dim) + valid_token_indices: (total <= B * (N + L),) + B: Batch size + M: Number of visual tokens per context parallel rank + L: Number of text tokens + dtype: Data type of the input and output tensors + + Returns: + x: (B, N, dim_x) tensor of visual tokens where N = M * cp_size + y: (B, L, dim_y) tensor of text token features + """ + _, cp_size = cp.get_cp_rank_size() + local_heads = self.num_heads // cp_size + local_dim = local_heads * self.head_dim + N = M * cp_size + + # Split sequence into visual and text tokens, adding back padding. + if B == 1: + out = out.view(B, -1, local_dim) + if out.size(1) > N: + x, y = torch.tensor_split(out, (N,), dim=1) # (B, N, local_dim), (B, <= L, local_dim) + y = F.pad(y, (0, 0, 0, L - y.size(1))) # (B, L, local_dim) + else: + # Empty prompt. + x, y = out, out.new_zeros(B, L, local_dim) + else: + x, y = pad_and_split_xy(out, valid_token_indices, B, N, L, dtype) + assert x.size() == (B, N, local_dim) + assert y.size() == (B, L, local_dim) + + # Communicate across context parallel ranks. + x = x.view(B, N, local_heads, self.head_dim) + x = cp.all_to_all_collect_heads(x) # (B, M, dim_x = num_heads * head_dim) + if cp.is_cp_active(): + y = cp.all_gather(y) # (cp_size * B, L, local_heads * head_dim) + y = rearrange(y, "(G B) L D -> B L (G D)", G=cp_size, D=local_dim) # (B, L, dim_x) + + x = self.proj_x(x) + y = self.proj_y(y) + return x, y + + def forward( + self, + x: torch.Tensor, # (B, M, dim_x) + y: torch.Tensor, # (B, L, dim_y) + *, + scale_x: torch.Tensor, # (B, dim_x), modulation for pre-RMSNorm. + scale_y: torch.Tensor, # (B, dim_y), modulation for pre-RMSNorm. + packed_indices: Dict[str, torch.Tensor] = None, + checkpoint_qkv: bool = False, + checkpoint_post_attn: bool = False, + **rope_rotation, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Forward pass of asymmetric multi-modal attention. + + Args: + x: (B, M, dim_x) tensor of visual tokens + y: (B, L, dim_y) tensor of text token features + packed_indices: Dict with keys for Flash Attention + num_frames: Number of frames in the video. N = num_frames * num_spatial_tokens + + Returns: + x: (B, M, dim_x) tensor of visual tokens after multi-modal attention + y: (B, L, dim_y) tensor of text token features after multi-modal attention + """ + B, L, _ = y.shape + _, M, _ = x.shape + + # Predict a packed QKV tensor from visual and text features. + q, k, v = ck(self.prepare_qkv, + x=x, + y=y, + scale_x=scale_x, + scale_y=scale_y, + rope_cos=rope_rotation.get("rope_cos"), + rope_sin=rope_rotation.get("rope_sin"), + valid_token_indices=packed_indices["valid_token_indices_kv"], + max_seqlen_in_batch=packed_indices["max_seqlen_in_batch_kv"], + enabled=checkpoint_qkv, + ) # (total <= B * (N + L), 3, local_heads, head_dim) + + # Self-attention is expensive, so don't checkpoint it. + out = self.run_attention( + q, k, v, B=B, + cu_seqlens=packed_indices["cu_seqlens_kv"], + max_seqlen_in_batch=packed_indices["max_seqlen_in_batch_kv"], + ) + + x, y = ck(self.post_attention, + out, + B=B, M=M, L=L, + dtype=v.dtype, + valid_token_indices=packed_indices["valid_token_indices_kv"], + enabled=checkpoint_post_attn, + ) + + return x, y + + +@torch.compile(disable=not COMPILE_MMDIT_BLOCK) +class AsymmetricJointBlock(nn.Module): + def __init__( + self, + hidden_size_x: int, + hidden_size_y: int, + num_heads: int, + *, + mlp_ratio_x: float = 8.0, # Ratio of hidden size to d_model for MLP for visual tokens. + mlp_ratio_y: float = 4.0, # Ratio of hidden size to d_model for MLP for text tokens. + update_y: bool = True, # Whether to update text tokens in this block. + device: Optional[torch.device] = None, + **block_kwargs, + ): + super().__init__() + self.update_y = update_y + self.hidden_size_x = hidden_size_x + self.hidden_size_y = hidden_size_y + self.mod_x = nn.Linear(hidden_size_x, 4 * hidden_size_x, device=device) + if self.update_y: + self.mod_y = nn.Linear(hidden_size_x, 4 * hidden_size_y, device=device) + else: + self.mod_y = nn.Linear(hidden_size_x, hidden_size_y, device=device) + + # Self-attention: + self.attn = AsymmetricAttention( + hidden_size_x, + hidden_size_y, + num_heads=num_heads, + update_y=update_y, + device=device, + **block_kwargs, + ) + + # MLP. + mlp_hidden_dim_x = int(hidden_size_x * mlp_ratio_x) + assert mlp_hidden_dim_x == int(1536 * 8) + self.mlp_x = FeedForward( + in_features=hidden_size_x, + hidden_size=mlp_hidden_dim_x, + multiple_of=256, + ffn_dim_multiplier=None, + device=device, + ) + + # MLP for text not needed in last block. + if self.update_y: + mlp_hidden_dim_y = int(hidden_size_y * mlp_ratio_y) + self.mlp_y = FeedForward( + in_features=hidden_size_y, + hidden_size=mlp_hidden_dim_y, + multiple_of=256, + ffn_dim_multiplier=None, + device=device, + ) + + def forward( + self, + x: torch.Tensor, + c: torch.Tensor, + y: torch.Tensor, + # TODO: These could probably just go into attn_kwargs + checkpoint_ff: bool = False, + checkpoint_qkv: bool = False, + checkpoint_post_attn: bool = False, + **attn_kwargs, + ): + """Forward pass of a block. + + Args: + x: (B, N, dim) tensor of visual tokens + c: (B, dim) tensor of conditioned features + y: (B, L, dim) tensor of text tokens + num_frames: Number of frames in the video. N = num_frames * num_spatial_tokens + + Returns: + x: (B, N, dim) tensor of visual tokens after block + y: (B, L, dim) tensor of text tokens after block + """ + N = x.size(1) + + c = F.silu(c) + mod_x = self.mod_x(c) + scale_msa_x, gate_msa_x, scale_mlp_x, gate_mlp_x = mod_x.chunk(4, dim=1) + mod_y = self.mod_y(c) + + if self.update_y: + scale_msa_y, gate_msa_y, scale_mlp_y, gate_mlp_y = mod_y.chunk(4, dim=1) + else: + scale_msa_y = mod_y + + # Self-attention block. + x_attn, y_attn = self.attn( + x, + y, + scale_x=scale_msa_x, + scale_y=scale_msa_y, + checkpoint_qkv=checkpoint_qkv, + checkpoint_post_attn=checkpoint_post_attn, + **attn_kwargs, + ) + + assert x_attn.size(1) == N + x = residual_tanh_gated_rmsnorm(x, x_attn, gate_msa_x) + + if self.update_y: + y = residual_tanh_gated_rmsnorm(y, y_attn, gate_msa_y) + + # MLP block. + x = ck(self.ff_block_x, x, scale_mlp_x, gate_mlp_x, enabled=checkpoint_ff) + if self.update_y: + y = ck(self.ff_block_y, y, scale_mlp_y, gate_mlp_y, enabled=checkpoint_ff) # type: ignore + return x, y + + def ff_block_x(self, x, scale_x, gate_x): + x_mod = modulated_rmsnorm(x, scale_x) + x_res = self.mlp_x(x_mod) + x = residual_tanh_gated_rmsnorm(x, x_res, gate_x) # Sandwich norm + return x + + def ff_block_y(self, y, scale_y, gate_y): + y_mod = modulated_rmsnorm(y, scale_y) + y_res = self.mlp_y(y_mod) + y = residual_tanh_gated_rmsnorm(y, y_res, gate_y) # Sandwich norm + return y + + +@torch.compile(disable=not COMPILE_FINAL_LAYER) +class FinalLayer(nn.Module): + """ + The final layer of DiT. + """ + + def __init__( + self, + hidden_size, + patch_size, + out_channels, + device: Optional[torch.device] = None, + ): + super().__init__() + self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, device=device) + self.mod = nn.Linear(hidden_size, 2 * hidden_size, device=device) + self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, device=device) + + def forward(self, x, c): + c = F.silu(c) + shift, scale = self.mod(c).chunk(2, dim=1) + x = modulate(self.norm_final(x), shift, scale) + x = self.linear(x) + return x + + +class AsymmDiTJoint(nn.Module): + """ + Diffusion model with a Transformer backbone. + + Ingests text embeddings instead of a label. + """ + + def __init__( + self, + *, + patch_size=2, + in_channels=4, + hidden_size_x=1152, + hidden_size_y=1152, + depth=48, + num_heads=16, + mlp_ratio_x=8.0, + mlp_ratio_y=4.0, + t5_feat_dim: int = 4096, + t5_token_length: int = 256, + patch_embed_bias: bool = True, + timestep_mlp_bias: bool = True, + timestep_scale: Optional[float] = None, + use_extended_posenc: bool = False, + rope_theta: float = 10000.0, + device: Optional[torch.device] = None, + **block_kwargs, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = in_channels + self.patch_size = patch_size + self.num_heads = num_heads + self.hidden_size_x = hidden_size_x + self.hidden_size_y = hidden_size_y + self.head_dim = hidden_size_x // num_heads # Head dimension and count is determined by visual. + self.use_extended_posenc = use_extended_posenc + self.t5_token_length = t5_token_length + self.t5_feat_dim = t5_feat_dim + self.rope_theta = rope_theta # Scaling factor for frequency computation for temporal RoPE. + + self.x_embedder = PatchEmbed( + patch_size=patch_size, + in_chans=in_channels, + embed_dim=hidden_size_x, + bias=patch_embed_bias, + device=device, + ) + # Conditionings + # Timestep + self.t_embedder = TimestepEmbedder(hidden_size_x, bias=timestep_mlp_bias, timestep_scale=timestep_scale) + + # Caption Pooling (T5) + self.t5_y_embedder = AttentionPool(t5_feat_dim, num_heads=8, output_dim=hidden_size_x, device=device) + + # Dense Embedding Projection (T5) + self.t5_yproj = nn.Linear(t5_feat_dim, hidden_size_y, bias=True, device=device) + + # Initialize pos_frequencies as an empty parameter. + self.pos_frequencies = nn.Parameter(torch.empty(3, self.num_heads, self.head_dim // 2, device=device)) + + # for depth 48: + # b = 0: AsymmetricJointBlock, update_y=True + # b = 1: AsymmetricJointBlock, update_y=True + # ... + # b = 46: AsymmetricJointBlock, update_y=True + # b = 47: AsymmetricJointBlock, update_y=False. No need to update text features. + blocks = [] + for b in range(depth): + # Joint multi-modal block + update_y = b < depth - 1 + block = AsymmetricJointBlock( + hidden_size_x, + hidden_size_y, + num_heads, + mlp_ratio_x=mlp_ratio_x, + mlp_ratio_y=mlp_ratio_y, + update_y=update_y, + device=device, + **block_kwargs, + ) + + blocks.append(block) + self.blocks = nn.ModuleList(blocks) + + self.final_layer = FinalLayer(hidden_size_x, patch_size, self.out_channels, device=device) + + def embed_x(self, x: torch.Tensor) -> torch.Tensor: + """ + Args: + x: (B, C=12, T, H, W) tensor of visual tokens + + Returns: + x: (B, C=3072, N) tensor of visual tokens with positional embedding. + """ + return self.x_embedder(x) # Convert BcTHW to BCN + + @torch.compile(disable=not COMPILE_MMDIT_BLOCK) + def prepare( + self, + x: torch.Tensor, + sigma: torch.Tensor, + t5_feat: torch.Tensor, + t5_mask: torch.Tensor, + ): + """Prepare input and conditioning embeddings.""" + + # Visual patch embeddings with positional encoding. + T, H, W = x.shape[-3:] + pH, pW = H // self.patch_size, W // self.patch_size + x = self.embed_x(x) # (B, N, D), where N = T * H * W / patch_size ** 2 + assert x.ndim == 3 + B = x.size(0) + + # Construct position array of size [N, 3]. + # pos[:, 0] is the frame index for each location, + # pos[:, 1] is the row index for each location, and + # pos[:, 2] is the column index for each location. + N = T * pH * pW + assert x.size(1) == N + pos = create_position_matrix(T, pH=pH, pW=pW, device=x.device, dtype=torch.float32) # (N, 3) + rope_cos, rope_sin = compute_mixed_rotation( + freqs=self.pos_frequencies, pos=pos + ) # Each are (N, num_heads, dim // 2) + + # Global vector embedding for conditionings. + c_t = self.t_embedder(1 - sigma) # (B, D) + + # Pool T5 tokens using attention pooler + # Note y_feat[1] contains T5 token features. + assert ( + t5_feat.size(1) == self.t5_token_length + ), f"Expected L={self.t5_token_length}, got {t5_feat.shape} for y_feat." + t5_y_pool = self.t5_y_embedder(t5_feat, t5_mask) # (B, D) + assert t5_y_pool.size(0) == B, f"Expected B={B}, got {t5_y_pool.shape} for t5_y_pool." + + c = c_t + t5_y_pool + + y_feat = self.t5_yproj(t5_feat) # (B, L, t5_feat_dim) --> (B, L, D) + + return x, c, y_feat, rope_cos, rope_sin + + def forward( + self, + x: torch.Tensor, # [1, 12, 7, 60, 106] + sigma: torch.Tensor, # [1] + y_feat: List[torch.Tensor], # [0][1, 256, 4096] + y_mask: List[torch.Tensor], # [0][1, 256] + packed_indices: Dict[str, torch.Tensor] = None, + rope_cos: torch.Tensor = None, + rope_sin: torch.Tensor = None, + num_ff_checkpoint: int = 0, # 48 + num_qkv_checkpoint: int = 0, # 48 + num_post_attn_checkpoint: int = 0, # 0 + ): + """Forward pass of DiT. + + Args: + x: (B, C, T, H, W) tensor of spatial inputs (images or latent representations of images) + sigma: (B,) tensor of noise standard deviations + y_feat: List((B, L, y_feat_dim) tensor of caption token features. For SDXL text encoders: L=77, y_feat_dim=2048) + y_mask: List((B, L) boolean tensor indicating which tokens are not padding) + packed_indices: Dict with keys for Flash Attention. Result of compute_packed_indices. + """ + _, _, T, H, W = x.shape + + if self.pos_frequencies.dtype != torch.float32: + warnings.warn(f"pos_frequencies dtype {self.pos_frequencies.dtype} != torch.float32") + + # Use EFFICIENT_ATTENTION backend for T5 pooling, since we have a mask. + # Have to call sdpa_kernel outside of a torch.compile region. + with sdpa_kernel(torch.nn.attention.SDPBackend.EFFICIENT_ATTENTION): + x, c, y_feat, rope_cos, rope_sin = self.prepare(x, sigma, y_feat[0], y_mask[0]) # [1, 11130, 3072], [1, 3072], [1, 256, 1536], [11130, 24, 64] + del y_mask + + cp_rank, cp_size = cp.get_cp_rank_size() + N = x.size(1) + M = N // cp_size + assert N % cp_size == 0, f"Visual sequence length ({x.shape[1]}) must be divisible by cp_size ({cp_size})." + + if cp_size > 1: + x = x.narrow(1, cp_rank * M, M) + + assert self.num_heads % cp_size == 0 + local_heads = self.num_heads // cp_size + rope_cos = rope_cos.narrow(1, cp_rank * local_heads, local_heads) + rope_sin = rope_sin.narrow(1, cp_rank * local_heads, local_heads) + + for i, block in enumerate(self.blocks): + x, y_feat = block( # [1, 11130, 3072], [1, 256, 1536] + x, + c, + y_feat, + rope_cos=rope_cos, + rope_sin=rope_sin, + packed_indices=packed_indices, + checkpoint_ff=i < num_ff_checkpoint, + checkpoint_qkv=i < num_qkv_checkpoint, + checkpoint_post_attn=i < num_post_attn_checkpoint, + ) # (B, M, D), (B, L, D) + del y_feat # Final layers don't use dense text features. + + x = self.final_layer(x, c) # (B, M, patch_size ** 2 * out_channels) [1, 11130, 48] + + patch = x.size(2) + x = cp.all_gather(x) + x = rearrange(x, "(G B) M P -> B (G M) P", G=cp_size, P=patch) # [1, 11130, 48] + x = rearrange( # [1, 12, 7, 60, 106] + x, + "B (T hp wp) (p1 p2 c) -> B c T (hp p1) (wp p2)", + T=T, + hp=H // self.patch_size, + wp=W // self.patch_size, + p1=self.patch_size, + p2=self.patch_size, + c=self.out_channels, + ) + + return x diff --git a/fastvideo/models/mochi_genmo/mochi_preview/dit/joint_model/context_parallel.py b/fastvideo/models/mochi_genmo/mochi_preview/dit/joint_model/context_parallel.py new file mode 100644 index 000000000..7b21244c7 --- /dev/null +++ b/fastvideo/models/mochi_genmo/mochi_preview/dit/joint_model/context_parallel.py @@ -0,0 +1,158 @@ +from typing import Tuple + +import torch +import torch.distributed as dist +from einops import rearrange + +_CONTEXT_PARALLEL_GROUP = None +_CONTEXT_PARALLEL_RANK = None +_CONTEXT_PARALLEL_GROUP_SIZE = None +_CONTEXT_PARALLEL_GROUP_RANKS = None + + +def get_cp_rank_size() -> Tuple[int, int]: + if _CONTEXT_PARALLEL_GROUP: + assert isinstance(_CONTEXT_PARALLEL_RANK, int) and isinstance(_CONTEXT_PARALLEL_GROUP_SIZE, int) + return _CONTEXT_PARALLEL_RANK, _CONTEXT_PARALLEL_GROUP_SIZE + else: + return 0, 1 + + +def local_shard(x: torch.Tensor, dim: int = 2) -> torch.Tensor: + if not _CONTEXT_PARALLEL_GROUP: + return x + + cp_rank, cp_size = get_cp_rank_size() + return x.tensor_split(cp_size, dim=dim)[cp_rank] + + +def set_cp_group(cp_group, ranks, global_rank): + global _CONTEXT_PARALLEL_GROUP, _CONTEXT_PARALLEL_RANK, _CONTEXT_PARALLEL_GROUP_SIZE, _CONTEXT_PARALLEL_GROUP_RANKS + if _CONTEXT_PARALLEL_GROUP is not None: + raise RuntimeError("CP group already initialized.") + _CONTEXT_PARALLEL_GROUP = cp_group + _CONTEXT_PARALLEL_RANK = dist.get_rank(cp_group) + _CONTEXT_PARALLEL_GROUP_SIZE = dist.get_world_size(cp_group) + _CONTEXT_PARALLEL_GROUP_RANKS = ranks + + assert _CONTEXT_PARALLEL_RANK == ranks.index( + global_rank + ), f"Rank mismatch: {global_rank} in {ranks} does not have position {_CONTEXT_PARALLEL_RANK} " + assert _CONTEXT_PARALLEL_GROUP_SIZE == len( + ranks + ), f"Group size mismatch: {_CONTEXT_PARALLEL_GROUP_SIZE} != len({ranks})" + + +def get_cp_group(): + if _CONTEXT_PARALLEL_GROUP is None: + raise RuntimeError("CP group not initialized") + return _CONTEXT_PARALLEL_GROUP + + +def is_cp_active(): + return _CONTEXT_PARALLEL_GROUP is not None + + +class AllGatherIntoTensorFunction(torch.autograd.Function): + @staticmethod + def forward(ctx, x: torch.Tensor, reduce_dtype, group: dist.ProcessGroup): + ctx.reduce_dtype = reduce_dtype + ctx.group = group + ctx.batch_size = x.size(0) + group_size = dist.get_world_size(group) + + x = x.contiguous() + output = torch.empty(group_size * x.size(0), *x.shape[1:], dtype=x.dtype, device=x.device) + dist.all_gather_into_tensor(output, x, group=group) + return output + + +def all_gather(tensor: torch.Tensor) -> torch.Tensor: + if not _CONTEXT_PARALLEL_GROUP: + return tensor + + return AllGatherIntoTensorFunction.apply(tensor, torch.float32, _CONTEXT_PARALLEL_GROUP) + + +@torch.compiler.disable() +def _all_to_all_single(output, input, group): + # Disable compilation since torch compile changes contiguity. + assert input.is_contiguous(), "Input tensor must be contiguous." + assert output.is_contiguous(), "Output tensor must be contiguous." + return dist.all_to_all_single(output, input, group=group) + + +class CollectTokens(torch.autograd.Function): + @staticmethod + def forward(ctx, qkv: torch.Tensor, group: dist.ProcessGroup, num_heads: int): + """Redistribute heads and receive tokens. + + Args: + qkv: query, key or value. Shape: [B, M, 3 * num_heads * head_dim] + + Returns: + qkv: shape: [3, B, N, local_heads, head_dim] + + where M is the number of local tokens, + N = cp_size * M is the number of global tokens, + local_heads = num_heads // cp_size is the number of local heads. + """ + ctx.group = group + ctx.num_heads = num_heads + cp_size = dist.get_world_size(group) + assert num_heads % cp_size == 0 + ctx.local_heads = num_heads // cp_size + + qkv = rearrange( + qkv, + "B M (qkv G h d) -> G M h B (qkv d)", + qkv=3, + G=cp_size, + h=ctx.local_heads, + ).contiguous() + + output_chunks = torch.empty_like(qkv) + _all_to_all_single(output_chunks, qkv, group=group) + + return rearrange(output_chunks, "G M h B (qkv d) -> qkv B (G M) h d", qkv=3) + + +def all_to_all_collect_tokens(x: torch.Tensor, num_heads: int) -> torch.Tensor: + if not _CONTEXT_PARALLEL_GROUP: + # Move QKV dimension to the front. + # B M (3 H d) -> 3 B M H d + B, M, _ = x.size() + x = x.view(B, M, 3, num_heads, -1) + return x.permute(2, 0, 1, 3, 4) + + return CollectTokens.apply(x, _CONTEXT_PARALLEL_GROUP, num_heads) + + +class CollectHeads(torch.autograd.Function): + @staticmethod + def forward(ctx, x: torch.Tensor, group: dist.ProcessGroup): + """Redistribute tokens and receive heads. + + Args: + x: Output of attention. Shape: [B, N, local_heads, head_dim] + + Returns: + Shape: [B, M, num_heads * head_dim] + """ + ctx.group = group + ctx.local_heads = x.size(2) + ctx.head_dim = x.size(3) + group_size = dist.get_world_size(group) + x = rearrange(x, "B (G M) h D -> G h M B D", G=group_size).contiguous() + output = torch.empty_like(x) + _all_to_all_single(output, x, group=group) + del x + return rearrange(output, "G h M B D -> B M (G h D)") + + +def all_to_all_collect_heads(x: torch.Tensor) -> torch.Tensor: + if not _CONTEXT_PARALLEL_GROUP: + # Merge heads. + return x.view(x.size(0), x.size(1), x.size(2) * x.size(3)) + + return CollectHeads.apply(x, _CONTEXT_PARALLEL_GROUP) diff --git a/fastvideo/models/mochi_genmo/mochi_preview/dit/joint_model/layers.py b/fastvideo/models/mochi_genmo/mochi_preview/dit/joint_model/layers.py new file mode 100644 index 000000000..05c685847 --- /dev/null +++ b/fastvideo/models/mochi_genmo/mochi_preview/dit/joint_model/layers.py @@ -0,0 +1,179 @@ +import collections.abc +import math +from itertools import repeat +from typing import Callable, Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange + + +# From PyTorch internals +def _ntuple(n): + def parse(x): + if isinstance(x, collections.abc.Iterable) and not isinstance(x, str): + return tuple(x) + return tuple(repeat(x, n)) + + return parse + + +to_2tuple = _ntuple(2) + + +class TimestepEmbedder(nn.Module): + def __init__( + self, + hidden_size: int, + frequency_embedding_size: int = 256, + *, + bias: bool = True, + timestep_scale: Optional[float] = None, + device: Optional[torch.device] = None, + ): + super().__init__() + self.mlp = nn.Sequential( + nn.Linear(frequency_embedding_size, hidden_size, bias=bias, device=device), + nn.SiLU(), + nn.Linear(hidden_size, hidden_size, bias=bias, device=device), + ) + self.frequency_embedding_size = frequency_embedding_size + self.timestep_scale = timestep_scale + + @staticmethod + def timestep_embedding(t, dim, max_period=10000): + half = dim // 2 + freqs = torch.arange(start=0, end=half, dtype=torch.float32, device=t.device) + freqs.mul_(-math.log(max_period) / half).exp_() + args = t[:, None].float() * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return embedding + + def forward(self, t): + if self.timestep_scale is not None: + t = t * self.timestep_scale + t_freq = self.timestep_embedding(t, self.frequency_embedding_size) + t_emb = self.mlp(t_freq) + return t_emb + + +class PooledCaptionEmbedder(nn.Module): + def __init__( + self, + caption_feature_dim: int, + hidden_size: int, + *, + bias: bool = True, + device: Optional[torch.device] = None, + ): + super().__init__() + self.caption_feature_dim = caption_feature_dim + self.hidden_size = hidden_size + self.mlp = nn.Sequential( + nn.Linear(caption_feature_dim, hidden_size, bias=bias, device=device), + nn.SiLU(), + nn.Linear(hidden_size, hidden_size, bias=bias, device=device), + ) + + def forward(self, x): + return self.mlp(x) + + +class FeedForward(nn.Module): + def __init__( + self, + in_features: int, + hidden_size: int, + multiple_of: int, + ffn_dim_multiplier: Optional[float], + device: Optional[torch.device] = None, + ): + super().__init__() + # keep parameter count and computation constant compared to standard FFN + hidden_size = int(2 * hidden_size / 3) + # custom dim factor multiplier + if ffn_dim_multiplier is not None: + hidden_size = int(ffn_dim_multiplier * hidden_size) + hidden_size = multiple_of * ((hidden_size + multiple_of - 1) // multiple_of) + + self.hidden_dim = hidden_size + self.w1 = nn.Linear(in_features, 2 * hidden_size, bias=False, device=device) + self.w2 = nn.Linear(hidden_size, in_features, bias=False, device=device) + + def forward(self, x): + # assert self.w1.weight.dtype == torch.bfloat16, f"FFN weight dtype {self.w1.weight.dtype} != bfloat16" + x, gate = self.w1(x).chunk(2, dim=-1) + x = self.w2(F.silu(x) * gate) + return x + + +class PatchEmbed(nn.Module): + def __init__( + self, + patch_size: int = 16, + in_chans: int = 3, + embed_dim: int = 768, + norm_layer: Optional[Callable] = None, + flatten: bool = True, + bias: bool = True, + dynamic_img_pad: bool = False, + device: Optional[torch.device] = None, + ): + super().__init__() + self.patch_size = to_2tuple(patch_size) + self.flatten = flatten + self.dynamic_img_pad = dynamic_img_pad + + self.proj = nn.Conv2d( + in_chans, + embed_dim, + kernel_size=patch_size, + stride=patch_size, + bias=bias, + device=device, + ) + assert norm_layer is None + self.norm = norm_layer(embed_dim, device=device) if norm_layer else nn.Identity() + + def forward(self, x): + B, _C, T, H, W = x.shape + if not self.dynamic_img_pad: + assert ( + H % self.patch_size[0] == 0 + ), f"Input height ({H}) should be divisible by patch size ({self.patch_size[0]})." + assert ( + W % self.patch_size[1] == 0 + ), f"Input width ({W}) should be divisible by patch size ({self.patch_size[1]})." + else: + pad_h = (self.patch_size[0] - H % self.patch_size[0]) % self.patch_size[0] + pad_w = (self.patch_size[1] - W % self.patch_size[1]) % self.patch_size[1] + x = F.pad(x, (0, pad_w, 0, pad_h)) + + x = rearrange(x, "B C T H W -> (B T) C H W", B=B, T=T) + x = self.proj(x) + + # Flatten temporal and spatial dimensions. + if not self.flatten: + raise NotImplementedError("Must flatten output.") + x = rearrange(x, "(B T) C H W -> B (T H W) C", B=B, T=T) + + x = self.norm(x) + return x + + +class RMSNorm(torch.nn.Module): + def __init__(self, hidden_size, eps=1e-5, device=None): + super().__init__() + self.eps = eps + self.weight = torch.nn.Parameter(torch.empty(hidden_size, device=device)) + self.register_parameter("bias", None) + + def forward(self, x): + # assert self.weight.dtype == torch.float32, f"RMSNorm weight dtype {self.weight.dtype} != float32" + + x_fp32 = x.float() + x_normed = x_fp32 * torch.rsqrt(x_fp32.pow(2).mean(-1, keepdim=True) + self.eps) + return (x_normed * self.weight).type_as(x) diff --git a/fastvideo/models/mochi_genmo/mochi_preview/dit/joint_model/lora.py b/fastvideo/models/mochi_genmo/mochi_preview/dit/joint_model/lora.py new file mode 100644 index 000000000..da6a7b511 --- /dev/null +++ b/fastvideo/models/mochi_genmo/mochi_preview/dit/joint_model/lora.py @@ -0,0 +1,112 @@ +#! /usr/bin/env python3 +import math +from typing import Dict, List, Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class LoRALayer: + def __init__( + self, + r: int, + lora_alpha: int, + lora_dropout: float, + merge_weights: bool, + ): + self.r = r + self.lora_alpha = lora_alpha + if lora_dropout > 0.0: + self.lora_dropout = nn.Dropout(p=lora_dropout) + else: + self.lora_dropout = lambda x: x + self.merged = False + self.merge_weights = merge_weights + + +def mark_only_lora_as_trainable(model: nn.Module, bias: str = "none") -> None: + assert bias == "none", f"Only bias='none' is supported" + for n, p in model.named_parameters(): + if "lora_" not in n: + p.requires_grad = False + + +def lora_state_dict(model: nn.Module, bias: str = "none") -> Dict[str, torch.Tensor]: + assert bias == "none", f"Only bias='none' is supported" + my_state_dict = model.state_dict() + return {k: my_state_dict[k] for k in my_state_dict if "lora_" in k} + + +class LoraLinear(nn.Linear, LoRALayer): + # LoRA implemented in a dense layer + def __init__( + self, + in_features: int, + out_features: int, + r: int = 0, + lora_alpha: int = 1, + lora_dropout: float = 0.0, + fan_in_fan_out: bool = False, # Set this to True if the layer to replace stores weight like (fan_in, fan_out) + merge_weights: bool = True, + **kwargs, + ): + nn.Linear.__init__(self, in_features, out_features, **kwargs) + LoRALayer.__init__(self, r=r, lora_alpha=lora_alpha, lora_dropout=lora_dropout, merge_weights=merge_weights) + + self.fan_in_fan_out = fan_in_fan_out + # Actual trainable parameters + if r > 0: + self.lora_A = nn.Parameter(self.weight.new_zeros((r, in_features)).to(torch.float32)) + self.lora_B = nn.Parameter(self.weight.new_zeros((out_features, r)).to(torch.float32)) + self.scaling = self.lora_alpha / self.r + + # Freezing the pre-trained weight matrix + self.weight.requires_grad = False + + self.reset_parameters() + + if fan_in_fan_out: + self.weight.data = self.weight.data.transpose(0, 1) + + def reset_parameters(self): + nn.Linear.reset_parameters(self) + if hasattr(self, "lora_A"): + # initialize B the same way as the default for nn.Linear and A to zero + # this is different than what is described in the paper but should not affect performance + nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5)) + nn.init.zeros_(self.lora_B) + + def train(self, mode: bool = True): + def T(w): + return w.transpose(0, 1) if self.fan_in_fan_out else w + + nn.Linear.train(self, mode) + if mode: + if self.merge_weights and self.merged: + # Make sure that the weights are not merged + if self.r > 0: + self.weight.data -= T(self.lora_B @ self.lora_A) * self.scaling + self.merged = False + else: + if self.merge_weights and not self.merged: + # Merge the weights and mark it + if self.r > 0: + self.weight.data += T(self.lora_B @ self.lora_A) * self.scaling + self.merged = True + + def forward(self, x: torch.Tensor): + def T(w): + return w.transpose(0, 1) if self.fan_in_fan_out else w + + if self.r > 0 and not self.merged: + result = F.linear(x, T(self.weight), bias=self.bias) + + x = self.lora_dropout(x) + x = x @ self.lora_A.transpose(0, 1) + x = x @ self.lora_B.transpose(0, 1) + x = x * self.scaling + + return result + x + else: + return F.linear(x, T(self.weight), bias=self.bias) diff --git a/fastvideo/models/mochi_genmo/mochi_preview/dit/joint_model/mod_rmsnorm.py b/fastvideo/models/mochi_genmo/mochi_preview/dit/joint_model/mod_rmsnorm.py new file mode 100644 index 000000000..dacf0d91f --- /dev/null +++ b/fastvideo/models/mochi_genmo/mochi_preview/dit/joint_model/mod_rmsnorm.py @@ -0,0 +1,15 @@ +import torch + + +def modulated_rmsnorm(x, scale, eps=1e-6): + dtype = x.dtype + x = x.float() + + # Compute RMS + mean_square = x.pow(2).mean(-1, keepdim=True) + inv_rms = torch.rsqrt(mean_square + eps) + + # Normalize and modulate + x_normed = x * inv_rms + x_modulated = x_normed * (1 + scale.unsqueeze(1).float()) + return x_modulated.to(dtype) diff --git a/fastvideo/models/mochi_genmo/mochi_preview/dit/joint_model/residual_tanh_gated_rmsnorm.py b/fastvideo/models/mochi_genmo/mochi_preview/dit/joint_model/residual_tanh_gated_rmsnorm.py new file mode 100644 index 000000000..3c6d8b21e --- /dev/null +++ b/fastvideo/models/mochi_genmo/mochi_preview/dit/joint_model/residual_tanh_gated_rmsnorm.py @@ -0,0 +1,20 @@ +import torch + + +def residual_tanh_gated_rmsnorm(x, x_res, gate, eps=1e-6): + # Convert to fp32 for precision + x_res = x_res.float() + + # Compute RMS + mean_square = x_res.pow(2).mean(-1, keepdim=True) + scale = torch.rsqrt(mean_square + eps) + + # Apply tanh to gate + tanh_gate = torch.tanh(gate).unsqueeze(1) + + # Normalize and apply gated scaling + x_normed = x_res * scale * tanh_gate + + # Apply residual connection + output = x + x_normed.type_as(x) + return output diff --git a/fastvideo/models/mochi_genmo/mochi_preview/dit/joint_model/rope_mixed.py b/fastvideo/models/mochi_genmo/mochi_preview/dit/joint_model/rope_mixed.py new file mode 100644 index 000000000..f2952bd6b --- /dev/null +++ b/fastvideo/models/mochi_genmo/mochi_preview/dit/joint_model/rope_mixed.py @@ -0,0 +1,88 @@ +import functools +import math + +import torch + + +def centers(start: float, stop, num, dtype=None, device=None): + """linspace through bin centers. + + Args: + start (float): Start of the range. + stop (float): End of the range. + num (int): Number of points. + dtype (torch.dtype): Data type of the points. + device (torch.device): Device of the points. + + Returns: + centers (Tensor): Centers of the bins. Shape: (num,). + """ + edges = torch.linspace(start, stop, num + 1, dtype=dtype, device=device) + return (edges[:-1] + edges[1:]) / 2 + + +@functools.lru_cache(maxsize=1) +def create_position_matrix( + T: int, + pH: int, + pW: int, + device: torch.device, + dtype: torch.dtype, + *, + target_area: float = 36864, +): + """ + Args: + T: int - Temporal dimension + pH: int - Height dimension after patchify + pW: int - Width dimension after patchify + + Returns: + pos: [T * pH * pW, 3] - position matrix + """ + with torch.no_grad(): + # Create 1D tensors for each dimension + t = torch.arange(T, dtype=dtype) + + # Positionally interpolate to area 36864. + # (3072x3072 frame with 16x16 patches = 192x192 latents). + # This automatically scales rope positions when the resolution changes. + # We use a large target area so the model is more sensitive + # to changes in the learned pos_frequencies matrix. + scale = math.sqrt(target_area / (pW * pH)) + w = centers(-pW * scale / 2, pW * scale / 2, pW) + h = centers(-pH * scale / 2, pH * scale / 2, pH) + + # Use meshgrid to create 3D grids + grid_t, grid_h, grid_w = torch.meshgrid(t, h, w, indexing="ij") + + # Stack and reshape the grids. + pos = torch.stack([grid_t, grid_h, grid_w], dim=-1) # [T, pH, pW, 3] + pos = pos.view(-1, 3) # [T * pH * pW, 3] + pos = pos.to(dtype=dtype, device=device) + + return pos + + +def compute_mixed_rotation( + freqs: torch.Tensor, + pos: torch.Tensor, +): + """ + Project each 3-dim position into per-head, per-head-dim 1D frequencies. + + Args: + freqs: [3, num_heads, num_freqs] - learned rotation frequency (for t, row, col) for each head position + pos: [N, 3] - position of each token + num_heads: int + + Returns: + freqs_cos: [N, num_heads, num_freqs] - cosine components + freqs_sin: [N, num_heads, num_freqs] - sine components + """ + with torch.autocast("cuda", enabled=False): + assert freqs.ndim == 3 + freqs_sum = torch.einsum("Nd,dhf->Nhf", pos.to(freqs), freqs) + freqs_cos = torch.cos(freqs_sum) + freqs_sin = torch.sin(freqs_sum) + return freqs_cos, freqs_sin diff --git a/fastvideo/models/mochi_genmo/mochi_preview/dit/joint_model/temporal_rope.py b/fastvideo/models/mochi_genmo/mochi_preview/dit/joint_model/temporal_rope.py new file mode 100644 index 000000000..a8276db18 --- /dev/null +++ b/fastvideo/models/mochi_genmo/mochi_preview/dit/joint_model/temporal_rope.py @@ -0,0 +1,34 @@ +# Based on Llama3 Implementation. +import torch + + +def apply_rotary_emb_qk_real( + xqk: torch.Tensor, + freqs_cos: torch.Tensor, + freqs_sin: torch.Tensor, +) -> torch.Tensor: + """ + Apply rotary embeddings to input tensors using the given frequency tensor without complex numbers. + + Args: + xqk (torch.Tensor): Query and/or Key tensors to apply rotary embeddings. Shape: (B, S, *, num_heads, D) + Can be either just query or just key, or both stacked along some batch or * dim. + freqs_cos (torch.Tensor): Precomputed cosine frequency tensor. + freqs_sin (torch.Tensor): Precomputed sine frequency tensor. + + Returns: + torch.Tensor: The input tensor with rotary embeddings applied. + """ + assert xqk.dtype == torch.bfloat16 + # Split the last dimension into even and odd parts + xqk_even = xqk[..., 0::2] + xqk_odd = xqk[..., 1::2] + + # Apply rotation + cos_part = (xqk_even * freqs_cos - xqk_odd * freqs_sin).type_as(xqk) + sin_part = (xqk_even * freqs_sin + xqk_odd * freqs_cos).type_as(xqk) + + # Interleave the results back into the original shape + out = torch.stack([cos_part, sin_part], dim=-1).flatten(-2) + assert out.dtype == torch.bfloat16 + return out diff --git a/fastvideo/models/mochi_genmo/mochi_preview/dit/joint_model/utils.py b/fastvideo/models/mochi_genmo/mochi_preview/dit/joint_model/utils.py new file mode 100644 index 000000000..bde2406f2 --- /dev/null +++ b/fastvideo/models/mochi_genmo/mochi_preview/dit/joint_model/utils.py @@ -0,0 +1,109 @@ +from typing import Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +def modulate(x, shift, scale): + return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) + + +def pool_tokens(x: torch.Tensor, mask: torch.Tensor, *, keepdim=False) -> torch.Tensor: + """ + Pool tokens in x using mask. + + NOTE: We assume x does not require gradients. + + Args: + x: (B, L, D) tensor of tokens. + mask: (B, L) boolean tensor indicating which tokens are not padding. + + Returns: + pooled: (B, D) tensor of pooled tokens. + """ + assert x.size(1) == mask.size(1) # Expected mask to have same length as tokens. + assert x.size(0) == mask.size(0) # Expected mask to have same batch size as tokens. + mask = mask[:, :, None].to(dtype=x.dtype) + mask = mask / mask.sum(dim=1, keepdim=True).clamp(min=1) + pooled = (x * mask).sum(dim=1, keepdim=keepdim) + return pooled + + +class AttentionPool(nn.Module): + def __init__( + self, + embed_dim: int, + num_heads: int, + output_dim: int = None, + device: Optional[torch.device] = None, + ): + """ + Args: + spatial_dim (int): Number of tokens in sequence length. + embed_dim (int): Dimensionality of input tokens. + num_heads (int): Number of attention heads. + output_dim (int): Dimensionality of output tokens. Defaults to embed_dim. + """ + super().__init__() + self.num_heads = num_heads + self.to_kv = nn.Linear(embed_dim, 2 * embed_dim, device=device) + self.to_q = nn.Linear(embed_dim, embed_dim, device=device) + self.to_out = nn.Linear(embed_dim, output_dim or embed_dim, device=device) + + def forward(self, x, mask): + """ + Args: + x (torch.Tensor): (B, L, D) tensor of input tokens. + mask (torch.Tensor): (B, L) boolean tensor indicating which tokens are not padding. + + NOTE: We assume x does not require gradients. + + Returns: + x (torch.Tensor): (B, D) tensor of pooled tokens. + """ + D = x.size(2) + + # Construct attention mask, shape: (B, 1, num_queries=1, num_keys=1+L). + attn_mask = mask[:, None, None, :].bool() # (B, 1, 1, L). + attn_mask = F.pad(attn_mask, (1, 0), value=True) # (B, 1, 1, 1+L). + + # Average non-padding token features. These will be used as the query. + x_pool = pool_tokens(x, mask, keepdim=True) # (B, 1, D) + + # Concat pooled features to input sequence. + x = torch.cat([x_pool, x], dim=1) # (B, L+1, D) + + # Compute queries, keys, values. Only the mean token is used to create a query. + kv = self.to_kv(x) # (B, L+1, 2 * D) + q = self.to_q(x[:, 0]) # (B, D) + + # Extract heads. + head_dim = D // self.num_heads + kv = kv.unflatten(2, (2, self.num_heads, head_dim)) # (B, 1+L, 2, H, head_dim) + kv = kv.transpose(1, 3) # (B, H, 2, 1+L, head_dim) + k, v = kv.unbind(2) # (B, H, 1+L, head_dim) + q = q.unflatten(1, (self.num_heads, head_dim)) # (B, H, head_dim) + q = q.unsqueeze(2) # (B, H, 1, head_dim) + + # Compute attention. + x = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, dropout_p=0.0) # (B, H, 1, head_dim) + + # Concatenate heads and run output. + x = x.squeeze(2).flatten(1, 2) # (B, D = H * head_dim) + x = self.to_out(x) + return x + + +def pad_and_split_xy(xy, indices, B, N, L, dtype) -> Tuple[torch.Tensor, torch.Tensor]: + D = xy.size(1) + + # Pad sequences to (B, N + L, dim). + assert indices.ndim == 1 + indices = indices.unsqueeze(1).expand(-1, D) # (total,) -> (total, num_heads * head_dim) + output = torch.zeros(B * (N + L), D, device=xy.device, dtype=dtype) + output = torch.scatter(output, 0, indices, xy) + xy = output.view(B, N + L, D) + + # Split visual and text tokens along the sequence length. + return torch.tensor_split(xy, (N,), dim=1) diff --git a/fastvideo/models/mochi_genmo/mochi_preview/pipelines.py b/fastvideo/models/mochi_genmo/mochi_preview/pipelines.py new file mode 100644 index 000000000..d9ec1d8b1 --- /dev/null +++ b/fastvideo/models/mochi_genmo/mochi_preview/pipelines.py @@ -0,0 +1,682 @@ +import json +import os +import random +from abc import ABC, abstractmethod +from contextlib import contextmanager +from functools import partial +from typing import Any, Dict, List, Literal, Optional, Union, cast + +import numpy as np +import ray +import torch +import torch.distributed as dist +import torch.nn as nn +import torch.nn.functional as F +from einops import repeat +from safetensors import safe_open +from safetensors.torch import load_file +from torch import nn +from torch.distributed.fsdp import ( + BackwardPrefetch, + MixedPrecision, + ShardingStrategy, +) +from torch.distributed.fsdp import FullyShardedDataParallel as FSDP +from torch.distributed.fsdp.wrap import ( + lambda_auto_wrap_policy, + transformer_auto_wrap_policy, +) +from transformers import T5EncoderModel, T5Tokenizer +from transformers.models.t5.modeling_t5 import T5Block + +import genmo.mochi_preview.dit.joint_model.context_parallel as cp +from genmo.lib.progress import get_new_progress_bar, progress_bar +from genmo.lib.utils import Timer +from genmo.mochi_preview.vae.models import ( + Decoder, + Encoder, + decode_latents, + decode_latents_tiled_full, + decode_latents_tiled_spatial, +) +from genmo.mochi_preview.vae.vae_stats import dit_latents_to_vae_latents + + +def load_to_cpu(p, weights_only=True): + if p.endswith(".safetensors"): + return load_file(p) + else: + assert p.endswith(".pt") + return torch.load(p, map_location="cpu", weights_only=weights_only) + + +def linear_quadratic_schedule(num_steps, threshold_noise, linear_steps=None): + if linear_steps is None: + linear_steps = num_steps // 2 + linear_sigma_schedule = [i * threshold_noise / linear_steps for i in range(linear_steps)] + threshold_noise_step_diff = linear_steps - threshold_noise * num_steps + quadratic_steps = num_steps - linear_steps + quadratic_coef = threshold_noise_step_diff / (linear_steps * quadratic_steps**2) + linear_coef = threshold_noise / linear_steps - 2 * threshold_noise_step_diff / (quadratic_steps**2) + const = quadratic_coef * (linear_steps**2) + quadratic_sigma_schedule = [ + quadratic_coef * (i**2) + linear_coef * i + const for i in range(linear_steps, num_steps) + ] + sigma_schedule = linear_sigma_schedule + quadratic_sigma_schedule + [1.0] + sigma_schedule = [1.0 - x for x in sigma_schedule] + return sigma_schedule + + +T5_MODEL = "google/t5-v1_1-xxl" +MAX_T5_TOKEN_LENGTH = 256 + + +def setup_fsdp_sync(model, device_id, *, param_dtype, auto_wrap_policy) -> FSDP: + model = FSDP( + model, + sharding_strategy=ShardingStrategy.FULL_SHARD, + mixed_precision=MixedPrecision( + param_dtype=param_dtype, + reduce_dtype=torch.float32, + buffer_dtype=torch.float32, + ), + auto_wrap_policy=auto_wrap_policy, + backward_prefetch=BackwardPrefetch.BACKWARD_PRE, + limit_all_gathers=True, + device_id=device_id, + sync_module_states=True, + use_orig_params=True, + ) + torch.cuda.synchronize() + return model + + +class ModelFactory(ABC): + def __init__(self, **kwargs): + self.kwargs = kwargs + + @abstractmethod + def get_model(self, *, local_rank: int, device_id: Union[int, Literal["cpu"]], world_size: int) -> Any: + assert isinstance(device_id, int) or device_id == "cpu", "device_id must be an integer or 'cpu'" + # FSDP does not work when the model is on the CPU + if device_id == "cpu": + assert world_size == 1, "CPU offload only supports single-GPU inference" + + +class T5ModelFactory(ModelFactory): + def __init__(self, model_dir=None): + super().__init__() + self.model_dir = model_dir or T5_MODEL + + def get_model(self, *, local_rank, device_id, world_size): + super().get_model(local_rank=local_rank, device_id=device_id, world_size=world_size) + model = T5EncoderModel.from_pretrained(self.model_dir) + if world_size > 1: + model = setup_fsdp_sync( + model, + device_id=device_id, + param_dtype=torch.float32, + auto_wrap_policy=partial( + transformer_auto_wrap_policy, + transformer_layer_cls={ + T5Block, + }, + ), + ) + elif isinstance(device_id, int): + model = model.to(torch.device(f"cuda:{device_id}")) # type: ignore + return model.eval() + + +class DitModelFactory(ModelFactory): + def __init__( + self, *, + model_path: str, + model_dtype: str, + lora_path: Optional[str] = None, + attention_mode: Optional[str] = None + ): + # Infer attention mode if not specified + if attention_mode is None: + from genmo.lib.attn_imports import flash_varlen_attn # type: ignore + attention_mode = "sdpa" if flash_varlen_attn is None else "flash" + print(f"Attention mode: {attention_mode}") + + super().__init__( + model_path=model_path, + lora_path=lora_path, + model_dtype=model_dtype, + attention_mode=attention_mode + ) + + def get_model( + self, + *, + local_rank, + device_id, + world_size, + model_kwargs=None, + patch_model_fns=None, + strict_load=True, + load_checkpoint=True, + fast_init=True, + ): + from genmo.mochi_preview.dit.joint_model.asymm_models_joint import AsymmDiTJoint + + if not model_kwargs: + model_kwargs = {} + + lora_sd = None + lora_path = self.kwargs["lora_path"] + if lora_path is not None: + if lora_path.endswith(".safetensors"): + lora_sd = {} + with safe_open(lora_path, framework="pt") as f: + for k in f.keys(): + lora_sd[k] = f.get_tensor(k) + lora_kwargs = json.loads(f.metadata()["kwargs"]) + print(f"Loaded LoRA kwargs: {lora_kwargs}") + else: + lora = load_to_cpu(lora_path, weights_only=False) + lora_sd, lora_kwargs = lora["state_dict"], lora["kwargs"] + + model_kwargs.update(cast(dict, lora_kwargs)) + + model_args = dict( + depth=48, + patch_size=2, + num_heads=24, + hidden_size_x=3072, + hidden_size_y=1536, + mlp_ratio_x=4.0, + mlp_ratio_y=4.0, + in_channels=12, + qk_norm=True, + qkv_bias=False, + out_bias=True, + patch_embed_bias=True, + timestep_mlp_bias=True, + timestep_scale=1000.0, + t5_feat_dim=4096, + t5_token_length=256, + rope_theta=10000.0, + attention_mode=self.kwargs["attention_mode"], + **model_kwargs, + ) + + if fast_init: + model: nn.Module = torch.nn.utils.skip_init(AsymmDiTJoint, **model_args) + else: + model: nn.Module = AsymmDiTJoint(**model_args) + + for fn in patch_model_fns or []: + model = fn(model) + + # FSDP syncs weights from rank 0 to all other ranks + if local_rank == 0 and load_checkpoint: + model_path = self.kwargs["model_path"] + sd = load_to_cpu(model_path) + + # Load the state dictionary and capture the return value + load_result = model.load_state_dict(sd, strict=strict_load) + if not strict_load: + # Print mismatched keys + missing_keys = [k for k in load_result.missing_keys if ".lora_" not in k] + if missing_keys: + print(f"Missing keys from {model_path}: {missing_keys}") + if load_result.unexpected_keys: + print(f"Unexpected keys from {model_path}: {load_result.unexpected_keys}") + + if lora_sd: + model.load_state_dict(lora_sd, strict=strict_load) # type: ignore + + if world_size > 1: + assert self.kwargs["model_dtype"] == "bf16", "FP8 is not supported for multi-GPU inference" + + model = setup_fsdp_sync( + model, + device_id=device_id, + param_dtype=torch.float32, + auto_wrap_policy=partial( + lambda_auto_wrap_policy, + lambda_fn=lambda m: m in model.blocks, + ), + ) + elif isinstance(device_id, int): + model = model.to(torch.device(f"cuda:{device_id}")) + return model.eval() + + +class DecoderModelFactory(ModelFactory): + def __init__(self, *, model_path: str): + super().__init__(model_path=model_path) + + def get_model(self, *, local_rank=0, device_id=0, world_size=1): + # TODO(ved): Set flag for torch.compile + # TODO(ved): Use skip_init + + decoder = Decoder( + out_channels=3, + base_channels=128, + channel_multipliers=[1, 2, 4, 6], + temporal_expansions=[1, 2, 3], + spatial_expansions=[2, 2, 2], + num_res_blocks=[3, 3, 4, 6, 3], + latent_dim=12, + has_attention=[False, False, False, False, False], + output_norm=False, + nonlinearity="silu", + output_nonlinearity="silu", + causal=True, + ) + # VAE is not FSDP-wrapped + state_dict = load_file(self.kwargs["model_path"]) + decoder.load_state_dict(state_dict, strict=True) + device = torch.device(f"cuda:{device_id}") if isinstance(device_id, int) else "cpu" + decoder.eval().to(device) + return decoder + + +class EncoderModelFactory(ModelFactory): + def __init__(self, *, model_path: str): + super().__init__(model_path=model_path) + + def get_model(self, *, local_rank=0, device_id=0, world_size=1): + # TODO(ved): Set flag for torch.compile + # TODO(ved): Use skip_init + + # We don't FSDP the encoder b/c it is small + encoder = Encoder( + in_channels=15, + base_channels=64, + channel_multipliers=[1, 2, 4, 6], + num_res_blocks=[3, 3, 4, 6, 3], + latent_dim=12, + temporal_reductions=[1, 2, 3], + spatial_reductions=[2, 2, 2], + prune_bottlenecks=[False, False, False, False, False], + has_attentions=[False, True, True, True, True], + affine=True, + bias=True, + input_is_conv_1x1=True, + padding_mode="replicate", + ) + state_dict = load_file(self.kwargs["model_path"]) + encoder.load_state_dict(state_dict, strict=True) + device = torch.device(f"cuda:{device_id}") if isinstance(device_id, int) else "cpu" + encoder.eval().to(device) + return encoder + + +def get_conditioning( + tokenizer: T5Tokenizer, + encoder: Encoder, + device: torch.device, + batch_inputs: bool, + *, + prompt: str, + negative_prompt: str, +): + if batch_inputs: + return dict( + batched=get_conditioning_for_prompts( + tokenizer, encoder, device, [prompt, negative_prompt] + ) + ) + else: + cond_input = get_conditioning_for_prompts(tokenizer, encoder, device, [prompt]) + null_input = get_conditioning_for_prompts(tokenizer, encoder, device, [negative_prompt]) + return dict(cond=cond_input, null=null_input) + + +def get_conditioning_for_prompts(tokenizer, encoder, device, prompts: List[str]): + assert len(prompts) in [1, 2] # [neg] or [pos] or [pos, neg] + B = len(prompts) + t5_toks = tokenizer( + prompts, + padding="max_length", + truncation=True, + max_length=MAX_T5_TOKEN_LENGTH, + return_tensors="pt", + return_attention_mask=True, + ) + caption_input_ids_t5 = t5_toks["input_ids"] + caption_attention_mask_t5 = t5_toks["attention_mask"].bool() + del t5_toks + + assert caption_input_ids_t5.shape == (B, MAX_T5_TOKEN_LENGTH) + assert caption_attention_mask_t5.shape == (B, MAX_T5_TOKEN_LENGTH) + + # Special-case empty negative prompt by zero-ing it + if prompts[-1] == "": + caption_input_ids_t5[-1] = 0 + caption_attention_mask_t5[-1] = False + + caption_input_ids_t5 = caption_input_ids_t5.to(device, non_blocking=True) + caption_attention_mask_t5 = caption_attention_mask_t5.to(device, non_blocking=True) + + y_mask = [caption_attention_mask_t5] + y_feat = [encoder(caption_input_ids_t5, caption_attention_mask_t5).last_hidden_state.detach()] + # Sometimes returns a tensor, othertimes a tuple, not sure why + # See: https://huggingface.co/genmo/mochi-1-preview/discussions/3 + assert tuple(y_feat[-1].shape) == (B, MAX_T5_TOKEN_LENGTH, 4096) + assert y_feat[-1].dtype == torch.float32 + + return dict(y_mask=y_mask, y_feat=y_feat) + + +def compute_packed_indices( + device: torch.device, text_mask: torch.Tensor, num_latents: int +) -> Dict[str, Union[torch.Tensor, int]]: + """ + Based on https://github.com/Dao-AILab/flash-attention/blob/765741c1eeb86c96ee71a3291ad6968cfbf4e4a1/flash_attn/bert_padding.py#L60-L80 + + Args: + num_latents: Number of latent tokens + text_mask: (B, L) List of boolean tensor indicating which text tokens are not padding. + + Returns: + packed_indices: Dict with keys for Flash Attention: + - valid_token_indices_kv: up to (B * (N + L),) tensor of valid token indices (non-padding) + in the packed sequence. + - cu_seqlens_kv: (B + 1,) tensor of cumulative sequence lengths in the packed sequence. + - max_seqlen_in_batch_kv: int of the maximum sequence length in the batch. + """ + # Create an expanded token mask saying which tokens are valid across both visual and text tokens. + PATCH_SIZE = 2 + num_visual_tokens = num_latents // (PATCH_SIZE**2) + assert num_visual_tokens > 0 + + mask = F.pad(text_mask, (num_visual_tokens, 0), value=True) # (B, N + L) + seqlens_in_batch = mask.sum(dim=-1, dtype=torch.int32) # (B,) + valid_token_indices = torch.nonzero(mask.flatten(), as_tuple=False).flatten() # up to (B * (N + L),) + assert valid_token_indices.size(0) >= text_mask.size(0) * num_visual_tokens # At least (B * N,) + cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) + max_seqlen_in_batch = seqlens_in_batch.max().item() + + return { + "cu_seqlens_kv": cu_seqlens.to(device, non_blocking=True), + "max_seqlen_in_batch_kv": cast(int, max_seqlen_in_batch), + "valid_token_indices_kv": valid_token_indices.to(device, non_blocking=True), + } + + +def assert_eq(x, y, msg=None): + assert x == y, f"{msg or 'Assertion failed'}: {x} != {y}" + + +def sample_model(device, dit, conditioning, **args): + random.seed(args["seed"]) + np.random.seed(args["seed"]) + torch.manual_seed(args["seed"]) + + generator = torch.Generator(device=device) + generator.manual_seed(args["seed"]) + + w, h, t = args["width"], args["height"], args["num_frames"] + sample_steps = args["num_inference_steps"] + cfg_schedule = args["cfg_schedule"] + sigma_schedule = args["sigma_schedule"] + + assert_eq(len(cfg_schedule), sample_steps, "cfg_schedule must have length sample_steps") + assert_eq((t - 1) % 6, 0, "t - 1 must be divisible by 6") + assert_eq( + len(sigma_schedule), + sample_steps + 1, + "sigma_schedule must have length sample_steps + 1", + ) + + B = 1 + SPATIAL_DOWNSAMPLE = 8 + TEMPORAL_DOWNSAMPLE = 6 + IN_CHANNELS = 12 + latent_t = ((t - 1) // TEMPORAL_DOWNSAMPLE) + 1 + latent_w, latent_h = w // SPATIAL_DOWNSAMPLE, h // SPATIAL_DOWNSAMPLE + + z = torch.randn( + (B, IN_CHANNELS, latent_t, latent_h, latent_w), + device=device, + dtype=torch.float32, + ) + + num_latents = latent_t * latent_h * latent_w + cond_batched = cond_text = cond_null = None + if "cond" in conditioning: + cond_text = conditioning["cond"] + cond_null = conditioning["null"] + cond_text["packed_indices"] = compute_packed_indices(device, cond_text["y_mask"][0], num_latents) + cond_null["packed_indices"] = compute_packed_indices(device, cond_null["y_mask"][0], num_latents) + else: + cond_batched = conditioning["batched"] + cond_batched["packed_indices"] = compute_packed_indices(device, cond_batched["y_mask"][0], num_latents) + z = repeat(z, "b ... -> (repeat b) ...", repeat=2) + + def model_fn(*, z, sigma, cfg_scale): + if cond_batched: + with torch.autocast("cuda", dtype=torch.bfloat16): + out = dit(z, sigma, **cond_batched) + out_cond, out_uncond = torch.chunk(out, chunks=2, dim=0) + else: + nonlocal cond_text, cond_null + with torch.autocast("cuda", dtype=torch.bfloat16): + out_cond = dit(z, sigma, **cond_text) + out_uncond = dit(z, sigma, **cond_null) + assert out_cond.shape == out_uncond.shape + out_uncond = out_uncond.to(z) + out_cond = out_cond.to(z) + return out_uncond + cfg_scale * (out_cond - out_uncond) + + # Euler sampler w/ customizable sigma schedule & cfg scale + for i in get_new_progress_bar(range(0, sample_steps), desc="Sampling"): + sigma = sigma_schedule[i] + dsigma = sigma - sigma_schedule[i + 1] + + # `pred` estimates `z_0 - eps`. + pred = model_fn( + z=z, + sigma=torch.full([B] if cond_text else [B * 2], sigma, device=z.device), + cfg_scale=cfg_schedule[i], + ) + assert pred.dtype == torch.float32 + z = z + dsigma * pred + + z = z[:B] if cond_batched else z + return dit_latents_to_vae_latents(z) + + +@contextmanager +def move_to_device(model: nn.Module, target_device, *, enabled=True): + if not enabled: + yield + return + + og_device = next(model.parameters()).device + if og_device == target_device: + print(f"move_to_device is a no-op model is already on {target_device}") + else: + print(f"moving model from {og_device} -> {target_device}") + + model.to(target_device) + yield + if og_device != target_device: + print(f"moving model from {target_device} -> {og_device}") + model.to(og_device) + + +def t5_tokenizer(model_dir=None): + return T5Tokenizer.from_pretrained(model_dir or T5_MODEL, legacy=False) + + +class MochiSingleGPUPipeline: + def __init__( + self, + *, + text_encoder_factory: ModelFactory, + dit_factory: ModelFactory, + decoder_factory: ModelFactory, + cpu_offload: Optional[bool] = False, + decode_type: str = "full", + decode_args: Optional[Dict[str, Any]] = None, + fast_init=True, + strict_load=True + ): + self.device = torch.device("cuda:0") + self.tokenizer = t5_tokenizer(text_encoder_factory.model_dir) + t = Timer() + self.cpu_offload = cpu_offload + self.decode_args = decode_args or {} + self.decode_type = decode_type + init_id = "cpu" if cpu_offload else 0 + with t("load_text_encoder"): + self.text_encoder = text_encoder_factory.get_model( + local_rank=0, + device_id=init_id, + world_size=1, + ) + with t("load_dit"): + self.dit = dit_factory.get_model(local_rank=0, device_id=init_id, world_size=1, fast_init=fast_init, strict_load=strict_load) # type: ignore + with t("load_vae"): + self.decoder = decoder_factory.get_model(local_rank=0, device_id=init_id, world_size=1) + t.print_stats() + + def __call__(self, batch_cfg, prompt, negative_prompt, **kwargs): + with torch.inference_mode(): + print_max_memory = lambda: print( + f"Max memory reserved: {torch.cuda.max_memory_reserved() / 1024**3:.2f} GB" + ) + print_max_memory() + + with move_to_device(self.text_encoder, self.device): + conditioning = get_conditioning( + tokenizer=self.tokenizer, + encoder=self.text_encoder, + device=self.device, + batch_inputs=batch_cfg, + prompt=prompt, + negative_prompt=negative_prompt, + ) + print_max_memory() + + with move_to_device(self.dit, self.device): + latents = sample_model(self.device, self.dit, conditioning, **kwargs) + print_max_memory() + + with move_to_device(self.decoder, self.device): + if self.decode_type == "tiled_full": + frames = decode_latents_tiled_full( + self.decoder, latents, **self.decode_args) + elif self.decode_type == "tiled_spatial": + frames = decode_latents_tiled_spatial( + self.decoder, latents, **self.decode_args, + num_tiles_w=4, num_tiles_h=2) + else: + frames = decode_latents(self.decoder, latents) + print_max_memory() + return frames.cpu().numpy() + + +def cast_dit(model, dtype): + for name, module in model.named_modules(): + if isinstance(module, nn.Linear): + assert any( + n in name for n in ["mlp", "t5", "mod_", "attn.qkv_", "attn.proj_", "final_layer"] + ), f"Unexpected linear layer: {name}" + module.to(dtype=dtype) + elif isinstance(module, nn.Conv2d): + assert "x_embedder.proj" in name, f"Unexpected conv2d layer: {name}" + module.to(dtype=dtype) + return model + + +### ALL CODE BELOW HERE IS FOR MULTI-GPU MODE ### + + +# In multi-gpu mode, all models must belong to a device which has a predefined context parallel group +# So it doesn't make sense to work with models individually +class MultiGPUContext: + def __init__( + self, + *, + text_encoder_factory, + dit_factory, + decoder_factory, + device_id, + local_rank, + world_size, + ): + t = Timer() + self.device = torch.device(f"cuda:{device_id}") + print(f"Initializing rank {local_rank+1}/{world_size}") + assert world_size > 1, f"Multi-GPU mode requires world_size > 1, got {world_size}" + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = "29500" + with t("init_process_group"): + dist.init_process_group( + "nccl", + rank=local_rank, + world_size=world_size, + device_id=self.device, # force non-lazy init + ) + pg = dist.group.WORLD + cp.set_cp_group(pg, list(range(world_size)), local_rank) + distributed_kwargs = dict(local_rank=local_rank, device_id=device_id, world_size=world_size) + self.world_size = world_size + self.tokenizer = t5_tokenizer(text_encoder_factory.model_dir) + with t("load_text_encoder"): + self.text_encoder = text_encoder_factory.get_model(**distributed_kwargs) + with t("load_dit"): + self.dit = dit_factory.get_model(**distributed_kwargs) + with t("load_vae"): + self.decoder = decoder_factory.get_model(**distributed_kwargs) + self.local_rank = local_rank + t.print_stats() + + def run(self, *, fn, **kwargs): + return fn(self, **kwargs) + + +class MochiMultiGPUPipeline: + def __init__( + self, + *, + text_encoder_factory: ModelFactory, + dit_factory: ModelFactory, + decoder_factory: ModelFactory, + world_size: int, + ): + ray.init() + RemoteClass = ray.remote(MultiGPUContext) + self.ctxs = [ + RemoteClass.options(num_gpus=1).remote( + text_encoder_factory=text_encoder_factory, + dit_factory=dit_factory, + decoder_factory=decoder_factory, + world_size=world_size, + device_id=0, + local_rank=i, + ) + for i in range(world_size) + ] + for ctx in self.ctxs: + ray.get(ctx.__ray_ready__.remote()) + + def __call__(self, **kwargs): + def sample(ctx, *, batch_cfg, prompt, negative_prompt, **kwargs): + with progress_bar(type="ray_tqdm", enabled=ctx.local_rank == 0), torch.inference_mode(): + conditioning = get_conditioning( + ctx.tokenizer, + ctx.text_encoder, + ctx.device, + batch_cfg, + prompt=prompt, + negative_prompt=negative_prompt, + ) + latents = sample_model(ctx.device, ctx.dit, conditioning=conditioning, **kwargs) + if ctx.local_rank == 0: + torch.save(latents, "latents.pt") + frames = decode_latents(ctx.decoder, latents) + return frames.cpu().numpy() + + return ray.get([ctx.run.remote(fn=sample, **kwargs, show_progress=i == 0) for i, ctx in enumerate(self.ctxs)])[ + 0 + ] diff --git a/fastvideo/models/mochi_genmo/mochi_preview/vae/__init__.py b/fastvideo/models/mochi_genmo/mochi_preview/vae/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/fastvideo/models/mochi_genmo/mochi_preview/vae/cp_conv.py b/fastvideo/models/mochi_genmo/mochi_preview/vae/cp_conv.py new file mode 100644 index 000000000..71b86d82c --- /dev/null +++ b/fastvideo/models/mochi_genmo/mochi_preview/vae/cp_conv.py @@ -0,0 +1,155 @@ +from typing import Tuple, Union + +import torch +import torch.distributed as dist +import torch.nn.functional as F + +import genmo.mochi_preview.dit.joint_model.context_parallel as cp + + +def cast_tuple(t, length=1): + return t if isinstance(t, tuple) else ((t,) * length) + + +def cp_pass_frames(x: torch.Tensor, frames_to_send: int) -> torch.Tensor: + """ + Forward pass that handles communication between ranks for inference. + Args: + x: Tensor of shape (B, C, T, H, W) + frames_to_send: int, number of frames to communicate between ranks + Returns: + output: Tensor of shape (B, C, T', H, W) + """ + cp_rank, cp_world_size = cp.get_cp_rank_size() + if frames_to_send == 0 or cp_world_size == 1: + return x + + group = cp.get_cp_group() + global_rank = dist.get_rank() + + # Send to next rank + if cp_rank < cp_world_size - 1: + assert x.size(2) >= frames_to_send + tail = x[:, :, -frames_to_send:].contiguous() + dist.send(tail, global_rank + 1, group=group) + + # Receive from previous rank + if cp_rank > 0: + B, C, _, H, W = x.shape + recv_buffer = torch.empty( + (B, C, frames_to_send, H, W), + dtype=x.dtype, + device=x.device, + ) + dist.recv(recv_buffer, global_rank - 1, group=group) + x = torch.cat([recv_buffer, x], dim=2) + + return x + + +def _pad_to_max(x: torch.Tensor, max_T: int) -> torch.Tensor: + if max_T > x.size(2): + pad_T = max_T - x.size(2) + pad_dims = (0, 0, 0, 0, 0, pad_T) + return F.pad(x, pad_dims) + return x + + +def gather_all_frames(x: torch.Tensor) -> torch.Tensor: + """ + Gathers all frames from all processes for inference. + Args: + x: Tensor of shape (B, C, T, H, W) + Returns: + output: Tensor of shape (B, C, T_total, H, W) + """ + cp_rank, cp_size = cp.get_cp_rank_size() + if cp_size == 1: + return x + + cp_group = cp.get_cp_group() + + # Ensure the tensor is contiguous for collective operations + x = x.contiguous() + + # Get the local time dimension size + local_T = x.size(2) + local_T_tensor = torch.tensor([local_T], device=x.device, dtype=torch.int64) + + # Gather all T sizes from all processes + all_T = [torch.zeros(1, dtype=torch.int64, device=x.device) for _ in range(cp_size)] + dist.all_gather(all_T, local_T_tensor, group=cp_group) + all_T = [t.item() for t in all_T] + + # Pad the tensor at the end of the time dimension to match max_T + max_T = max(all_T) + x = _pad_to_max(x, max_T).contiguous() + + # Prepare a list to hold the gathered tensors + gathered_x = [torch.zeros_like(x).contiguous() for _ in range(cp_size)] + + # Perform the all_gather operation + dist.all_gather(gathered_x, x, group=cp_group) + + # Slice each gathered tensor back to its original T size + for idx, t_size in enumerate(all_T): + gathered_x[idx] = gathered_x[idx][:, :, :t_size] + + return torch.cat(gathered_x, dim=2) + + +def excessive_memory_usage(input: torch.Tensor, max_gb: float = 2.0) -> bool: + """Estimate memory usage based on input tensor size and data type.""" + element_size = input.element_size() # Size in bytes of each element + memory_bytes = input.numel() * element_size + memory_gb = memory_bytes / 1024**3 + return memory_gb > max_gb + + +class ContextParallelCausalConv3d(torch.nn.Conv3d): + def __init__( + self, + in_channels, + out_channels, + kernel_size: Union[int, Tuple[int, int, int]], + stride: Union[int, Tuple[int, int, int]], + **kwargs, + ): + kernel_size = cast_tuple(kernel_size, 3) + stride = cast_tuple(stride, 3) + height_pad = (kernel_size[1] - 1) // 2 + width_pad = (kernel_size[2] - 1) // 2 + + super().__init__( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + stride=stride, + dilation=(1, 1, 1), + padding=(0, height_pad, width_pad), + **kwargs, + ) + + def forward(self, x: torch.Tensor): + cp_rank, cp_world_size = cp.get_cp_rank_size() + + context_size = self.kernel_size[0] - 1 + if cp_rank == 0: + mode = "constant" if self.padding_mode == "zeros" else self.padding_mode + x = F.pad(x, (0, 0, 0, 0, context_size, 0), mode=mode) + + if cp_world_size == 1: + return super().forward(x) + + if all(s == 1 for s in self.stride): + # Receive some frames from previous rank. + x = cp_pass_frames(x, context_size) + return super().forward(x) + + # Less efficient implementation for strided convs. + # All gather x, infer and chunk. + x = gather_all_frames(x) # [B, C, k - 1 + global_T, H, W] + x = super().forward(x) + x_chunks = x.tensor_split(cp_world_size, dim=2) + assert len(x_chunks) == cp_world_size + return x_chunks[cp_rank] diff --git a/fastvideo/models/mochi_genmo/mochi_preview/vae/latent_dist.py b/fastvideo/models/mochi_genmo/mochi_preview/vae/latent_dist.py new file mode 100644 index 000000000..d99eac0c0 --- /dev/null +++ b/fastvideo/models/mochi_genmo/mochi_preview/vae/latent_dist.py @@ -0,0 +1,35 @@ +"""Container for latent space posterior.""" + +import torch + + +class LatentDistribution: + def __init__(self, mean: torch.Tensor, logvar: torch.Tensor): + """Initialize latent distribution. + + Args: + mean: Mean of the distribution. Shape: [B, C, T, H, W]. + logvar: Logarithm of variance of the distribution. Shape: [B, C, T, H, W]. + """ + assert mean.shape == logvar.shape + self.mean = mean + self.logvar = logvar + + def sample(self, temperature=1.0, generator: torch.Generator = None, noise=None): + if temperature == 0.0: + return self.mean + + if noise is None: + noise = torch.randn(self.mean.shape, device=self.mean.device, dtype=self.mean.dtype, generator=generator) + else: + assert noise.device == self.mean.device + noise = noise.to(self.mean.dtype) + + if temperature != 1.0: + raise NotImplementedError(f"Temperature {temperature} is not supported.") + + # Just Gaussian sample with no scaling of variance. + return noise * torch.exp(self.logvar * 0.5) + self.mean + + def mode(self): + return self.mean diff --git a/fastvideo/models/mochi_genmo/mochi_preview/vae/models.py b/fastvideo/models/mochi_genmo/mochi_preview/vae/models.py new file mode 100644 index 000000000..568bcabfa --- /dev/null +++ b/fastvideo/models/mochi_genmo/mochi_preview/vae/models.py @@ -0,0 +1,1021 @@ +from functools import partial +from typing import Any, Callable, Dict, List, Optional, Tuple, Union + +import math +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange + +import genmo.mochi_preview.dit.joint_model.context_parallel as cp +from genmo.lib.progress import get_new_progress_bar +from genmo.mochi_preview.vae.cp_conv import cp_pass_frames, gather_all_frames +from genmo.mochi_preview.vae.latent_dist import LatentDistribution +import genmo.mochi_preview.vae.cp_conv as cp_conv + + +def cast_tuple(t, length=1): + return t if isinstance(t, tuple) else ((t,) * length) + + +class GroupNormSpatial(nn.GroupNorm): + """ + GroupNorm applied per-frame. + """ + + def forward(self, x: torch.Tensor, *, chunk_size: int = 8): + B, C, T, H, W = x.shape + x = rearrange(x, "B C T H W -> (B T) C H W") + # Run group norm in chunks. + output = torch.empty_like(x) + for b in range(0, B * T, chunk_size): + output[b : b + chunk_size] = super().forward(x[b : b + chunk_size]) + return rearrange(output, "(B T) C H W -> B C T H W", B=B, T=T) + + +class SafeConv3d(torch.nn.Conv3d): + """ + NOTE: No support for padding along time dimension. + Input must already be padded along time. + """ + + def forward(self, input): + memory_count = torch.prod(torch.tensor(input.shape)).item() * 2 / 1024**3 + if memory_count > 2 and self.stride[0] == 1: + part_num = int(memory_count / 2) + 1 + k = self.kernel_size[0] + input_idx = torch.arange(k - 1, input.size(2)) + input_chunks_idx = torch.chunk(input_idx, part_num, dim=0) + + # assert self.stride[0] == 1, f"stride {self.stride}" + assert self.dilation[0] == 1, f"dilation {self.dilation}" + assert self.padding[0] == 0, f"padding {self.padding}" + + # Comptue output size + assert not input.requires_grad + B, _, T_in, H_in, W_in = input.shape + output_size = ( + B, + self.out_channels, + T_in - k + 1, + H_in // self.stride[1], + W_in // self.stride[2], + ) + output = torch.empty(output_size, dtype=input.dtype, device=input.device) + for input_chunk_idx in input_chunks_idx: + input_s = input_chunk_idx[0] - k + 1 + input_e = input_chunk_idx[-1] + 1 + input_chunk = input[:, :, input_s:input_e, :, :] + output_chunk = super(SafeConv3d, self).forward(input_chunk) + + output_s = input_s + output_e = output_s + output_chunk.size(2) + output[:, :, output_s:output_e, :, :] = output_chunk + + return output + else: + return super(SafeConv3d, self).forward(input) + + +class StridedSafeConv3d(torch.nn.Conv3d): + def forward(self, input, local_shard: bool = False): + assert self.stride[0] == self.kernel_size[0] + assert self.dilation[0] == 1 + assert self.padding[0] == 0 + + kernel_size = self.kernel_size[0] + stride = self.stride[0] + T_in = input.size(2) + T_out = T_in // kernel_size + + # Parallel implementation. + if local_shard: + idx = torch.arange(T_out) + idx = cp.local_shard(idx, dim=0) + start = idx.min() * stride + end = idx.max() * stride + kernel_size + local_input = input[:, :, start:end, :, :] + return torch.nn.Conv3d.forward(self, local_input) + + raise NotImplementedError + + +class ContextParallelConv3d(SafeConv3d): + def __init__( + self, + in_channels, + out_channels, + kernel_size: Union[int, Tuple[int, int, int]], + stride: Union[int, Tuple[int, int, int]], + causal: bool = True, + context_parallel: bool = True, + **kwargs, + ): + self.causal = causal + self.context_parallel = context_parallel + kernel_size = cast_tuple(kernel_size, 3) + stride = cast_tuple(stride, 3) + height_pad = (kernel_size[1] - 1) // 2 + width_pad = (kernel_size[2] - 1) // 2 + + super().__init__( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + stride=stride, + dilation=(1, 1, 1), + padding=(0, height_pad, width_pad), + **kwargs, + ) + + def forward(self, x: torch.Tensor): + cp_rank, cp_world_size = cp.get_cp_rank_size() + + # Compute padding amounts. + context_size = self.kernel_size[0] - 1 + if self.causal: + pad_front = context_size + pad_back = 0 + else: + pad_front = context_size // 2 + pad_back = context_size - pad_front + + # Apply padding. + mode = "constant" if self.padding_mode == "zeros" else self.padding_mode + if self.context_parallel and cp_world_size == 1: + x = F.pad(x, (0, 0, 0, 0, pad_front, pad_back), mode=mode) + else: + if cp_rank == 0: + x = F.pad(x, (0, 0, 0, 0, pad_front, 0), mode=mode) + elif cp_rank == cp_world_size - 1 and pad_back: + x = F.pad(x, (0, 0, 0, 0, 0, pad_back), mode=mode) + + if self.context_parallel and cp_world_size == 1: + return super().forward(x) + + if self.stride[0] == 1: + # Receive some frames from previous rank. + x = cp_pass_frames(x, context_size) + return super().forward(x) + + # Less efficient implementation for strided convs. + # All gather x, infer and chunk. + assert x.dtype == torch.bfloat16, f"Expected x to be of type torch.bfloat16, got {x.dtype}" + + x = gather_all_frames(x) # [B, C, k - 1 + global_T, H, W] + return StridedSafeConv3d.forward(self, x, local_shard=True) + + +class Conv1x1(nn.Linear): + """*1x1 Conv implemented with a linear layer.""" + + def __init__(self, in_features: int, out_features: int, *args, **kwargs): + super().__init__(in_features, out_features, *args, **kwargs) + + def forward(self, x: torch.Tensor): + """Forward pass. + + Args: + x: Input tensor. Shape: [B, C, *] or [B, *, C]. + + Returns: + x: Output tensor. Shape: [B, C', *] or [B, *, C']. + """ + x = x.movedim(1, -1) + x = super().forward(x) + x = x.movedim(-1, 1) + return x + + +class DepthToSpaceTime(nn.Module): + def __init__( + self, + temporal_expansion: int, + spatial_expansion: int, + ): + super().__init__() + self.temporal_expansion = temporal_expansion + self.spatial_expansion = spatial_expansion + + # When printed, this module should show the temporal and spatial expansion factors. + def extra_repr(self): + return f"texp={self.temporal_expansion}, sexp={self.spatial_expansion}" + + def forward(self, x: torch.Tensor): + """Forward pass. + + Args: + x: Input tensor. Shape: [B, C, T, H, W]. + + Returns: + x: Rearranged tensor. Shape: [B, C/(st*s*s), T*st, H*s, W*s]. + """ + x = rearrange( + x, + "B (C st sh sw) T H W -> B C (T st) (H sh) (W sw)", + st=self.temporal_expansion, + sh=self.spatial_expansion, + sw=self.spatial_expansion, + ) + + cp_rank, _ = cp.get_cp_rank_size() + if self.temporal_expansion > 1 and cp_rank == 0: + # Drop the first self.temporal_expansion - 1 frames. + # This is because we always want the 3x3x3 conv filter to only apply + # to the first frame, and the first frame doesn't need to be repeated. + assert all(x.shape) + x = x[:, :, self.temporal_expansion - 1 :] + assert all(x.shape) + + return x + + +def norm_fn( + in_channels: int, + affine: bool = True, +): + return GroupNormSpatial(affine=affine, num_groups=32, num_channels=in_channels) + + +class ResBlock(nn.Module): + """Residual block that preserves the spatial dimensions.""" + + def __init__( + self, + channels: int, + *, + affine: bool = True, + attn_block: Optional[nn.Module] = None, + causal: bool = True, + prune_bottleneck: bool = False, + padding_mode: str, + bias: bool = True, + ): + super().__init__() + self.channels = channels + + assert causal + self.stack = nn.Sequential( + norm_fn(channels, affine=affine), + nn.SiLU(inplace=True), + ContextParallelConv3d( + in_channels=channels, + out_channels=channels // 2 if prune_bottleneck else channels, + kernel_size=(3, 3, 3), + stride=(1, 1, 1), + padding_mode=padding_mode, + bias=bias, + causal=causal, + ), + norm_fn(channels, affine=affine), + nn.SiLU(inplace=True), + ContextParallelConv3d( + in_channels=channels // 2 if prune_bottleneck else channels, + out_channels=channels, + kernel_size=(3, 3, 3), + stride=(1, 1, 1), + padding_mode=padding_mode, + bias=bias, + causal=causal, + ), + ) + + self.attn_block = attn_block if attn_block else nn.Identity() + + def forward(self, x: torch.Tensor): + """Forward pass. + + Args: + x: Input tensor. Shape: [B, C, T, H, W]. + """ + residual = x + x = self.stack(x) + x = x + residual + del residual + + return self.attn_block(x) + + +def prepare_for_attention(qkv: torch.Tensor, head_dim: int, qk_norm: bool = True): + """Prepare qkv tensor for attention and normalize qk. + + Args: + qkv: Input tensor. Shape: [B, L, 3 * num_heads * head_dim]. + + Returns: + q, k, v: qkv tensor split into q, k, v. Shape: [B, num_heads, L, head_dim]. + """ + assert qkv.ndim == 3 # [B, L, 3 * num_heads * head_dim] + assert qkv.size(2) % (3 * head_dim) == 0 + num_heads = qkv.size(2) // (3 * head_dim) + qkv = qkv.unflatten(2, (3, num_heads, head_dim)) + + q, k, v = qkv.unbind(2) # [B, L, num_heads, head_dim] + q = q.transpose(1, 2) # [B, num_heads, L, head_dim] + k = k.transpose(1, 2) # [B, num_heads, L, head_dim] + v = v.transpose(1, 2) # [B, num_heads, L, head_dim] + + if qk_norm: + q = F.normalize(q, p=2, dim=-1) + k = F.normalize(k, p=2, dim=-1) + + # Mixed precision can change the dtype of normed q/k to float32. + q = q.to(dtype=qkv.dtype) + k = k.to(dtype=qkv.dtype) + + return q, k, v + + +class Attention(nn.Module): + def __init__( + self, + dim: int, + head_dim: int = 32, + qkv_bias: bool = False, + out_bias: bool = True, + qk_norm: bool = True, + ) -> None: + super().__init__() + self.head_dim = head_dim + self.num_heads = dim // head_dim + self.qk_norm = qk_norm + + self.qkv = nn.Linear(dim, 3 * dim, bias=qkv_bias) + self.out = nn.Linear(dim, dim, bias=out_bias) + + def forward( + self, + x: torch.Tensor, + *, + chunk_size=2**15, + ) -> torch.Tensor: + """Compute temporal self-attention. + + Args: + x: Input tensor. Shape: [B, C, T, H, W]. + chunk_size: Chunk size for large tensors. + + Returns: + x: Output tensor. Shape: [B, C, T, H, W]. + """ + B, _, T, H, W = x.shape + + if T == 1: + # No attention for single frame. + x = x.movedim(1, -1) # [B, C, T, H, W] -> [B, T, H, W, C] + qkv = self.qkv(x) + _, _, x = qkv.chunk(3, dim=-1) # Throw away queries and keys. + x = self.out(x) + return x.movedim(-1, 1) # [B, T, H, W, C] -> [B, C, T, H, W] + + # 1D temporal attention. + x = rearrange(x, "B C t h w -> (B h w) t C") + qkv = self.qkv(x) + + # Input: qkv with shape [B, t, 3 * num_heads * head_dim] + # Output: x with shape [B, num_heads, t, head_dim] + + attn_kwargs = dict( + attn_mask=None, + dropout_p=0.0, + is_causal=True, + scale=self.head_dim**-0.5, + ) + + if qkv.size(0) <= chunk_size: + q, k, v = prepare_for_attention(qkv, self.head_dim, qk_norm=self.qk_norm) + x = F.scaled_dot_product_attention(q, k, v, **attn_kwargs) # [B, num_heads, t, head_dim] + assert x.size(0) == q.size(0) + else: + # Evaluate in chunks to avoid `RuntimeError: CUDA error: invalid configuration argument.` + # Chunks of 2**16 and up cause an error. + x = torch.empty((qkv.size(0), self.num_heads, qkv.size(1), self.head_dim), dtype=qkv.dtype, device=qkv.device) + for i in range(0, qkv.size(0), chunk_size): + qkv_chunk = qkv[i:i+chunk_size] + qc, kc, vc = prepare_for_attention(qkv_chunk, self.head_dim, qk_norm=self.qk_norm) + chunk = F.scaled_dot_product_attention(qc, kc, vc, **attn_kwargs) + assert chunk.size(0) == qc.size(0) + x[i:i+chunk_size].copy_(chunk) + + x = x.transpose(1, 2) # [B, t, num_heads, head_dim] + x = x.flatten(2) # [B, t, num_heads * head_dim] + + x = self.out(x) + x = rearrange(x, "(B h w) t C -> B C t h w", B=B, h=H, w=W) + return x + + +class AttentionBlock(nn.Module): + def __init__( + self, + dim: int, + **attn_kwargs, + ) -> None: + super().__init__() + self.norm = norm_fn(dim) + self.attn = Attention(dim, **attn_kwargs) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + self.attn(self.norm(x)) + + +class CausalUpsampleBlock(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + num_res_blocks: int, + *, + temporal_expansion: int = 2, + spatial_expansion: int = 2, + **block_kwargs, + ): + super().__init__() + + blocks = [] + for _ in range(num_res_blocks): + blocks.append(block_fn(in_channels, **block_kwargs)) + self.blocks = nn.Sequential(*blocks) + + self.temporal_expansion = temporal_expansion + self.spatial_expansion = spatial_expansion + + # Change channels in the final convolution layer. + self.proj = Conv1x1( + in_channels, + out_channels * temporal_expansion * (spatial_expansion**2), + ) + + self.d2st = DepthToSpaceTime(temporal_expansion=temporal_expansion, spatial_expansion=spatial_expansion) + + def forward(self, x): + x = self.blocks(x) + x = self.proj(x) + x = self.d2st(x) + return x + + +def block_fn(channels, *, affine: bool = True, has_attention: bool = False, **block_kwargs): + attn_block = AttentionBlock(channels) if has_attention else None + return ResBlock(channels, affine=affine, attn_block=attn_block, **block_kwargs) + + +def add_fourier_features(inputs: torch.Tensor, start=6, stop=8, step=1): + num_freqs = (stop - start) // step + assert inputs.ndim == 5 + C = inputs.size(1) + + # Create Base 2 Fourier features. + freqs = torch.arange(start, stop, step, dtype=inputs.dtype, device=inputs.device) + assert num_freqs == len(freqs) + w = torch.pow(2.0, freqs) * (2 * torch.pi) # [num_freqs] + C = inputs.shape[1] + w = w.repeat(C)[None, :, None, None, None] # [1, C * num_freqs, 1, 1, 1] + + # Interleaved repeat of input channels to match w. + h = inputs.repeat_interleave(num_freqs, dim=1) # [B, C * num_freqs, T, H, W] + # Scale channels by frequency. + h = w * h + + return torch.cat( + [ + inputs, + torch.sin(h), + torch.cos(h), + ], + dim=1, + ) + + +class FourierFeatures(nn.Module): + def __init__(self, start: int = 6, stop: int = 8, step: int = 1): + super().__init__() + self.start = start + self.stop = stop + self.step = step + + def forward(self, inputs): + """Add Fourier features to inputs. + + Args: + inputs: Input tensor. Shape: [B, C, T, H, W] + + Returns: + h: Output tensor. Shape: [B, (1 + 2 * num_freqs) * C, T, H, W] + """ + assert inputs.dtype == torch.float32, f"Fourier features must be computed in float32 to avoid artifacts." + return add_fourier_features(inputs, self.start, self.stop, self.step) + + +class Decoder(nn.Module): + def __init__( + self, + *, + out_channels: int = 3, + latent_dim: int, + base_channels: int, + channel_multipliers: List[int], + num_res_blocks: List[int], + temporal_expansions: Optional[List[int]] = None, + spatial_expansions: Optional[List[int]] = None, + has_attention: List[bool], + output_norm: bool = True, + nonlinearity: str = "silu", + output_nonlinearity: str = "silu", + causal: bool = True, + **block_kwargs, + ): + super().__init__() + self.input_channels = latent_dim + self.base_channels = base_channels + self.channel_multipliers = channel_multipliers + self.num_res_blocks = num_res_blocks + self.output_nonlinearity = output_nonlinearity + assert nonlinearity == "silu" + assert causal + + ch = [mult * base_channels for mult in channel_multipliers] + self.num_up_blocks = len(ch) - 1 + assert len(num_res_blocks) == self.num_up_blocks + 2 + + blocks = [] + new_block_fn = partial(block_fn, padding_mode="replicate") + + first_block = [nn.Conv3d(latent_dim, ch[-1], kernel_size=(1, 1, 1))] # Input layer. + # First set of blocks preserve channel count. + for _ in range(num_res_blocks[-1]): + first_block.append( + new_block_fn( + ch[-1], + has_attention=has_attention[-1], + causal=causal, + **block_kwargs, + ) + ) + blocks.append(nn.Sequential(*first_block)) + + assert len(temporal_expansions) == len(spatial_expansions) == self.num_up_blocks + assert len(num_res_blocks) == len(has_attention) == self.num_up_blocks + 2 + + upsample_block_fn = CausalUpsampleBlock + + for i in range(self.num_up_blocks): + block = upsample_block_fn( + ch[-i - 1], + ch[-i - 2], + num_res_blocks=num_res_blocks[-i - 2], + has_attention=has_attention[-i - 2], + temporal_expansion=temporal_expansions[-i - 1], + spatial_expansion=spatial_expansions[-i - 1], + causal=causal, + padding_mode="replicate", + **block_kwargs, + ) + blocks.append(block) + + assert not output_norm + + # Last block. Preserve channel count. + last_block = [] + for _ in range(num_res_blocks[0]): + last_block.append(new_block_fn(ch[0], has_attention=has_attention[0], causal=causal, **block_kwargs)) + blocks.append(nn.Sequential(*last_block)) + + self.blocks = nn.ModuleList(blocks) + self.output_proj = Conv1x1(ch[0], out_channels) + + def forward(self, x): + """Forward pass. + + Args: + x: Latent tensor. Shape: [B, input_channels, t, h, w]. Scaled [-1, 1]. + + Returns: + x: Reconstructed video tensor. Shape: [B, C, T, H, W]. Scaled to [-1, 1]. + T + 1 = (t - 1) * 4. + H = h * 16, W = w * 16. + """ + for block in self.blocks: + x = block(x) + + if self.output_nonlinearity == "silu": + x = F.silu(x, inplace=not self.training) + else: + assert not self.output_nonlinearity # StyleGAN3 omits the to-RGB nonlinearity. + + return self.output_proj(x).contiguous() + + +def make_broadcastable( + tensor: torch.Tensor, + axis: int, + ndim: int, +) -> torch.Tensor: + """ + Reshapes the input tensor to have singleton dimensions in all axes except the specified axis. + + Args: + tensor (torch.Tensor): The tensor to reshape. Typically 1D. + axis (int): The axis along which the tensor should retain its original size. + ndim (int): The total number of dimensions the reshaped tensor should have. + + Returns: + torch.Tensor: The reshaped tensor with shape suitable for broadcasting. + """ + if tensor.dim() != 1: + raise ValueError(f"Expected tensor to be 1D, but got {tensor.dim()}D tensor.") + + axis = (axis + ndim) % ndim # Ensure the axis is within the tensor dimensions + shape = [1] * ndim # Start with all dimensions as 1 + shape[axis] = tensor.size(0) # Set the specified axis to the size of the tensor + return tensor.view(*shape) + + +def blend(a: torch.Tensor, b: torch.Tensor, axis: int) -> torch.Tensor: + """ + Blends two tensors `a` and `b` along the specified axis using linear interpolation. + + Args: + a (torch.Tensor): The first tensor. + b (torch.Tensor): The second tensor. Must have the same shape as `a`. + axis (int): The axis along which to perform the blending. + + Returns: + torch.Tensor: The blended tensor. + """ + assert a.shape == b.shape, f"Tensors must have the same shape, got {a.shape} and {b.shape}" + steps = a.size(axis) + + # Create a weight tensor that linearly interpolates from 0 to 1 + start = 1 / (steps + 1) + end = steps / (steps + 1) + weight = torch.linspace(start, end, steps=steps, device=a.device, dtype=a.dtype) + + # Make the weight tensor broadcastable across all dimensions + weight = make_broadcastable(weight, axis, a.dim()) + + # Perform the blending + return a * (1 - weight) + b * weight + + +def blend_horizontal(a: torch.Tensor, b: torch.Tensor, overlap: int) -> torch.Tensor: + if overlap == 0: + return torch.cat([a, b], dim=-1) + + assert a.size(-1) >= overlap + assert b.size(-1) >= overlap + a_left, a_overlap = a[..., :-overlap], a[..., -overlap:] + b_overlap, b_right = b[..., :overlap], b[..., overlap:] + return torch.cat([a_left, blend(a_overlap, b_overlap, -1), b_right], dim=-1) + + +def blend_vertical(a: torch.Tensor, b: torch.Tensor, overlap: int) -> torch.Tensor: + if overlap == 0: + return torch.cat([a, b], dim=-2) + + assert a.size(-2) >= overlap + assert b.size(-2) >= overlap + a_top, a_overlap = a[..., :-overlap, :], a[..., -overlap:, :] + b_overlap, b_bottom = b[..., :overlap, :], b[..., overlap:, :] + return torch.cat([a_top, blend(a_overlap, b_overlap, -2), b_bottom], dim=-2) + + +def nearest_multiple(x: int, multiple: int) -> int: + return round(x / multiple) * multiple + + +def apply_tiled( + fn: Callable[[torch.Tensor], torch.Tensor], + x: torch.Tensor, + num_tiles_w: int, + num_tiles_h: int, + overlap: int = 0, # Number of pixel of overlap between adjacent tiles. + # Use a factor of 2 times the latent downsample factor. + min_block_size: int = 1, # Minimum number of pixels in each dimension when subdividing. +) -> Optional[torch.Tensor]: + if num_tiles_w == 1 and num_tiles_h == 1: + return fn(x) + + assert num_tiles_w & (num_tiles_w - 1) == 0, f"num_tiles_w={num_tiles_w} must be a power of 2" + assert num_tiles_h & (num_tiles_h - 1) == 0, f"num_tiles_h={num_tiles_h} must be a power of 2" + + H, W = x.shape[-2:] + assert H % min_block_size == 0 + assert W % min_block_size == 0 + ov = overlap // 2 + assert ov % min_block_size == 0 + + if num_tiles_w >= 2: + # Subdivide horizontally. + half_W = nearest_multiple(W // 2, min_block_size) + left = x[..., :, : half_W + ov] + right = x[..., :, half_W - ov :] + + assert num_tiles_w % 2 == 0, f"num_tiles_w={num_tiles_w} must be even" + left = apply_tiled(fn, left, num_tiles_w // 2, num_tiles_h, overlap, min_block_size) + right = apply_tiled(fn, right, num_tiles_w // 2, num_tiles_h, overlap, min_block_size) + if left is None or right is None: + return None + + # If `fn` changed the resolution, adjust the overlap. + resample_factor = left.size(-1) / (half_W + ov) + out_overlap = int(overlap * resample_factor) + + return blend_horizontal(left, right, out_overlap) + + if num_tiles_h >= 2: + # Subdivide vertically. + half_H = nearest_multiple(H // 2, min_block_size) + top = x[..., : half_H + ov, :] + bottom = x[..., half_H - ov :, :] + + assert num_tiles_h % 2 == 0, f"num_tiles_h={num_tiles_h} must be even" + top = apply_tiled(fn, top, num_tiles_w, num_tiles_h // 2, overlap, min_block_size) + bottom = apply_tiled(fn, bottom, num_tiles_w, num_tiles_h // 2, overlap, min_block_size) + if top is None or bottom is None: + return None + + # If `fn` changed the resolution, adjust the overlap. + resample_factor = top.size(-2) / (half_H + ov) + out_overlap = int(overlap * resample_factor) + + return blend_vertical(top, bottom, out_overlap) + + raise ValueError(f"Invalid num_tiles_w={num_tiles_w} and num_tiles_h={num_tiles_h}") + + +class DownsampleBlock(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + num_res_blocks, + *, + temporal_reduction=2, + spatial_reduction=2, + **block_kwargs, + ): + """ + Downsample block for the VAE encoder. + + Args: + in_channels: Number of input channels. + out_channels: Number of output channels. + num_res_blocks: Number of residual blocks. + temporal_reduction: Temporal reduction factor. + spatial_reduction: Spatial reduction factor. + """ + super().__init__() + layers = [] + + assert in_channels != out_channels + layers.append( + ContextParallelConv3d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=(temporal_reduction, spatial_reduction, spatial_reduction), + stride=(temporal_reduction, spatial_reduction, spatial_reduction), + # First layer in each block always uses replicate padding + padding_mode="replicate", + bias=block_kwargs["bias"], + ) + ) + + for _ in range(num_res_blocks): + layers.append(block_fn(out_channels, **block_kwargs)) + + self.layers = nn.Sequential(*layers) + + def forward(self, x): + return self.layers(x) + + +class Encoder(nn.Module): + def __init__( + self, + *, + in_channels: int, + base_channels: int, + channel_multipliers: List[int], + num_res_blocks: List[int], + latent_dim: int, + temporal_reductions: List[int], + spatial_reductions: List[int], + prune_bottlenecks: List[bool], + has_attentions: List[bool], + affine: bool = True, + bias: bool = True, + input_is_conv_1x1: bool = False, + padding_mode: str, + ): + super().__init__() + self.temporal_reductions = temporal_reductions + self.spatial_reductions = spatial_reductions + self.base_channels = base_channels + self.channel_multipliers = channel_multipliers + self.num_res_blocks = num_res_blocks + self.latent_dim = latent_dim + + ch = [mult * base_channels for mult in channel_multipliers] + num_down_blocks = len(ch) - 1 + assert len(num_res_blocks) == num_down_blocks + 2 + + layers = ( + [nn.Conv3d(in_channels, ch[0], kernel_size=(1, 1, 1), bias=True)] + if not input_is_conv_1x1 + else [Conv1x1(in_channels, ch[0])] + ) + + assert len(prune_bottlenecks) == num_down_blocks + 2 + assert len(has_attentions) == num_down_blocks + 2 + block = partial(block_fn, padding_mode=padding_mode, affine=affine, bias=bias) + + for _ in range(num_res_blocks[0]): + layers.append(block(ch[0], has_attention=has_attentions[0], prune_bottleneck=prune_bottlenecks[0])) + prune_bottlenecks = prune_bottlenecks[1:] + has_attentions = has_attentions[1:] + + assert len(temporal_reductions) == len(spatial_reductions) == len(ch) - 1 + for i in range(num_down_blocks): + layer = DownsampleBlock( + ch[i], + ch[i + 1], + num_res_blocks=num_res_blocks[i + 1], + temporal_reduction=temporal_reductions[i], + spatial_reduction=spatial_reductions[i], + prune_bottleneck=prune_bottlenecks[i], + has_attention=has_attentions[i], + affine=affine, + bias=bias, + padding_mode=padding_mode, + ) + + layers.append(layer) + + # Additional blocks. + for _ in range(num_res_blocks[-1]): + layers.append(block(ch[-1], has_attention=has_attentions[-1], prune_bottleneck=prune_bottlenecks[-1])) + + self.layers = nn.Sequential(*layers) + + # Output layers. + self.output_norm = norm_fn(ch[-1]) + self.output_proj = Conv1x1(ch[-1], 2 * latent_dim, bias=False) + + @property + def temporal_downsample(self): + return math.prod(self.temporal_reductions) + + @property + def spatial_downsample(self): + return math.prod(self.spatial_reductions) + + def forward(self, x) -> LatentDistribution: + """Forward pass. + + Args: + x: Input video tensor. Shape: [B, C, T, H, W]. Scaled to [-1, 1] + + Returns: + means: Latent tensor. Shape: [B, latent_dim, t, h, w]. Scaled [-1, 1]. + h = H // 8, w = W // 8, t - 1 = (T - 1) // 6 + logvar: Shape: [B, latent_dim, t, h, w]. + """ + assert x.ndim == 5, f"Expected 5D input, got {x.shape}" + + x = self.layers(x) + + x = self.output_norm(x) + x = F.silu(x, inplace=True) + x = self.output_proj(x) + + means, logvar = torch.chunk(x, 2, dim=1) + + assert means.ndim == 5 + assert logvar.shape == means.shape + assert means.size(1) == self.latent_dim + + return LatentDistribution(means, logvar) + + +def normalize_decoded_frames(samples): + samples = samples.float() + samples = (samples + 1.0) / 2.0 + samples.clamp_(0.0, 1.0) + frames = rearrange(samples, "b c t h w -> b t h w c") + return frames + + +@torch.inference_mode() +def decode_latents_tiled_full( + decoder, + z, + *, + tile_sample_min_height: int = 240, + tile_sample_min_width: int = 424, + tile_overlap_factor_height: float = 0.1666, + tile_overlap_factor_width: float = 0.2, + auto_tile_size: bool = True, + frame_batch_size: int = 6, +): + B, C, T, H, W = z.shape + assert frame_batch_size <= T, f"frame_batch_size must be <= T, got {frame_batch_size} > {T}" + + tile_sample_min_height = tile_sample_min_height if not auto_tile_size else H // 2 * 8 + tile_sample_min_width = tile_sample_min_width if not auto_tile_size else W // 2 * 8 + + tile_latent_min_height = int(tile_sample_min_height / 8) + tile_latent_min_width = int(tile_sample_min_width / 8) + + def blend_v(a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: + blend_extent = min(a.shape[3], b.shape[3], blend_extent) + for y in range(blend_extent): + b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, :, y, :] * ( + y / blend_extent + ) + return b + + def blend_h(a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: + blend_extent = min(a.shape[4], b.shape[4], blend_extent) + for x in range(blend_extent): + b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, :, x] * ( + x / blend_extent + ) + return b + + overlap_height = int(tile_latent_min_height * (1 - tile_overlap_factor_height)) + overlap_width = int(tile_latent_min_width * (1 - tile_overlap_factor_width)) + blend_extent_height = int(tile_sample_min_height * tile_overlap_factor_height) + blend_extent_width = int(tile_sample_min_width * tile_overlap_factor_width) + row_limit_height = tile_sample_min_height - blend_extent_height + row_limit_width = tile_sample_min_width - blend_extent_width + + # Split z into overlapping tiles and decode them separately. + # The tiles have an overlap to avoid seams between tiles. + pbar = get_new_progress_bar( + desc="Decoding latent tiles", + total=len(range(0, H, overlap_height)) * len(range(0, W, overlap_width)) * len(range(T // frame_batch_size)), + ) + rows = [] + for i in range(0, H, overlap_height): + row = [] + for j in range(0, W, overlap_width): + temporal = [] + for k in range(T // frame_batch_size): + remaining_frames = T % frame_batch_size + start_frame = frame_batch_size * k + (0 if k == 0 else remaining_frames) + end_frame = frame_batch_size * (k + 1) + remaining_frames + tile = z[ + :, + :, + start_frame:end_frame, + i : i + tile_latent_min_height, + j : j + tile_latent_min_width, + ] + tile = decoder(tile) + temporal.append(tile) + pbar.update(1) + row.append(torch.cat(temporal, dim=2)) + rows.append(row) + + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + # blend the above tile and the left tile + # to the current tile and add the current tile to the result row + if i > 0: + tile = blend_v(rows[i - 1][j], tile, blend_extent_height) + if j > 0: + tile = blend_h(row[j - 1], tile, blend_extent_width) + result_row.append(tile[:, :, :, :row_limit_height, :row_limit_width]) + result_rows.append(torch.cat(result_row, dim=4)) + + return normalize_decoded_frames(torch.cat(result_rows, dim=3)) + + +@torch.inference_mode() +def decode_latents_tiled_spatial( + decoder, + z, + *, + num_tiles_w: int, + num_tiles_h: int, + overlap: int = 0, # Number of pixel of overlap between adjacent tiles. + # Use a factor of 2 times the latent downsample factor. + min_block_size: int = 1, # Minimum number of pixels in each dimension when subdividing. +): + decoded = apply_tiled(decoder, z, num_tiles_w, num_tiles_h, overlap, min_block_size) + assert decoded is not None, f"Failed to decode latents with tiled spatial method" + return normalize_decoded_frames(decoded) + + +@torch.inference_mode() +def decode_latents(decoder, z): + assert z.ndim == 5 + cp_rank, cp_size = cp.get_cp_rank_size() + z = z.tensor_split(cp_size, dim=2)[cp_rank] # split along temporal dim + with torch.autocast("cuda", dtype=torch.bfloat16): + samples = decoder(z) + samples = gather_all_frames(samples) + return normalize_decoded_frames(samples) diff --git a/fastvideo/models/mochi_genmo/mochi_preview/vae/vae_stats.py b/fastvideo/models/mochi_genmo/mochi_preview/vae/vae_stats.py new file mode 100644 index 000000000..b0e8f69a8 --- /dev/null +++ b/fastvideo/models/mochi_genmo/mochi_preview/vae/vae_stats.py @@ -0,0 +1,67 @@ +import torch + +# Channel-wise mean and standard deviation of VAE encoder latents +STATS = { + "mean": torch.Tensor( + [ + -0.06730895953510081, + -0.038011381506090416, + -0.07477820912866141, + -0.05565264470995561, + 0.012767231469026969, + -0.04703542746246419, + 0.043896967884726704, + -0.09346305707025976, + -0.09918314763016893, + -0.008729793427399178, + -0.011931556316503654, + -0.0321993391887285, + ] + ), + "std": torch.Tensor( + [ + 0.9263795028493863, + 0.9248894543193766, + 0.9393059390890617, + 0.959253732819592, + 0.8244560132752793, + 0.917259975397747, + 0.9294154431013696, + 1.3720942357788521, + 0.881393668867029, + 0.9168315692124348, + 0.9185249279345552, + 0.9274757570805041, + ] + ), +} + + +def dit_latents_to_vae_latents(dit_outputs: torch.Tensor) -> torch.Tensor: + """Unnormalize latents output by Mochi's DiT to be compatible with VAE. + Run this on sampled latents before calling the VAE decoder. + + Args: + latents (torch.Tensor): [B, C_z, T_z, H_z, W_z], float + + Returns: + torch.Tensor: [B, C_z, T_z, H_z, W_z], float + """ + mean = STATS["mean"][:, None, None, None] + std = STATS["std"][:, None, None, None] + + assert dit_outputs.ndim == 5 + assert dit_outputs.size(1) == mean.size(0) == std.size(0) + return dit_outputs * std.to(dit_outputs) + mean.to(dit_outputs) + + +def vae_latents_to_dit_latents(vae_latents: torch.Tensor): + """Normalize latents output by the VAE encoder to be compatible with Mochi's DiT. + E.g, for fine-tuning or video-to-video. + """ + mean = STATS["mean"][:, None, None, None] + std = STATS["std"][:, None, None, None] + + assert vae_latents.ndim == 5 + assert vae_latents.size(1) == mean.size(0) == std.size(0) + return (vae_latents - mean.to(vae_latents)) / std.to(vae_latents) diff --git a/fastvideo/models/mochi_hf/modeling_mochi.py b/fastvideo/models/mochi_hf/modeling_mochi.py index 2d674be8d..49cdaf27a 100644 --- a/fastvideo/models/mochi_hf/modeling_mochi.py +++ b/fastvideo/models/mochi_hf/modeling_mochi.py @@ -205,8 +205,6 @@ def sage_attn_processor(query, key, value, attn_mask): batch_q = query.narrow(0, b, 1).index_select(1, true_indices) batch_k = key.narrow(0, b, 1).index_select(1, true_indices) batch_v = value.narrow(0, b, 1).index_select(1, true_indices) - from IPython import embed - embed() # Process with sage attention result = sageattn( @@ -325,12 +323,13 @@ class MochiAttnProcessor2_0: attn_mask = encoder_attention_mask[:, :].bool() attn_mask = F.pad(attn_mask, (sequence_length, 0), value=True) - sage_attn = True + + sage_attn = False if sage_attn: hidden_states = sage_attn_processor(query, key, value, attn_mask) else: hidden_states = flash_attn_no_pad(qkv, attn_mask, causal=False, dropout_p=0.0, softmax_scale=None) - + # hidden_states = F.scaled_dot_product_attention(query, key, value, attn_mask = None, dropout_p=0.0, is_causal=False) # valid_lengths = encoder_attention_mask.sum(dim=1) + sequence_length @@ -690,11 +689,11 @@ class MochiTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin): def forward( self, - hidden_states: torch.Tensor, - encoder_hidden_states: torch.Tensor, - timestep: torch.LongTensor, - encoder_attention_mask: torch.Tensor, - output_attn=False, + hidden_states: torch.Tensor, # [2, 12, 28, 60, 106] + encoder_hidden_states: torch.Tensor, # [2, 256, 4096] + timestep: torch.LongTensor, # [2] + encoder_attention_mask: torch.Tensor, #[2, 256] + output_attn = False, attention_kwargs: Optional[Dict[str, Any]] = None, return_dict: bool = False, ) -> torch.Tensor: @@ -719,28 +718,28 @@ class MochiTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin): logger.warning( "Passing `scale` via `attention_kwargs` when not using the PEFT backend is ineffective." ) - - batch_size, num_channels, num_frames, height, width = hidden_states.shape + + batch_size, num_channels, num_frames, height, width = hidden_states.shape p = self.config.patch_size post_patch_height = height // p post_patch_width = width // p # Peiyuan: This is hacked to force mochi to follow the behaviour of SD3 and Flux timestep = 1000 - timestep - temb, encoder_hidden_states = self.time_embed( + temb, encoder_hidden_states = self.time_embed( # [2, 3072], [2, 256, 1536] timestep, encoder_hidden_states, encoder_attention_mask, hidden_dtype=hidden_states.dtype, ) - hidden_states = hidden_states.permute(0, 2, 1, 3, 4).flatten(0, 1) - hidden_states = self.patch_embed(hidden_states) - hidden_states = hidden_states.unflatten(0, (batch_size, -1)).flatten(1, 2) + hidden_states = hidden_states.permute(0, 2, 1, 3, 4).flatten(0, 1) # [56, 12, 60, 106] + hidden_states = self.patch_embed(hidden_states) # [56, 1590, 3072] + hidden_states = hidden_states.unflatten(0, (batch_size, -1)).flatten(1, 2) # [2, 44520, 3072] - image_rotary_emb = self.rope( - self.pos_frequencies, - num_frames, + image_rotary_emb = self.rope( #[0][44520, 24, 64] + self.pos_frequencies, #[3, 24, 64] + num_frames, # 28 post_patch_height, post_patch_width, device=hidden_states.device, @@ -774,7 +773,7 @@ class MochiTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin): **ckpt_kwargs, ) else: - hidden_states, encoder_hidden_states, attn_outputs = block( + hidden_states, encoder_hidden_states, attn_outputs = block( # [2, 44520, 3072], [2, 256, 1536], hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, encoder_attention_mask=encoder_attention_mask, @@ -783,16 +782,14 @@ class MochiTransformer3DModel(ModelMixin, ConfigMixin, PeftAdapterMixin): output_attn=output_attn, ) attn_outputs_list.append(attn_outputs) - + hidden_states = self.norm_out(hidden_states, temb) - hidden_states = self.proj_out(hidden_states) - - hidden_states = hidden_states.reshape( - batch_size, num_frames, post_patch_height, post_patch_width, p, p, -1 - ) - hidden_states = hidden_states.permute(0, 6, 1, 2, 4, 3, 5) - output = hidden_states.reshape(batch_size, -1, num_frames, height, width) + hidden_states = self.proj_out(hidden_states) #[2, 44520, 48] + hidden_states = hidden_states.reshape(batch_size, num_frames, post_patch_height, post_patch_width, p, p, -1) # [2, 28, 30, 53, 2, 2, 12] + hidden_states = hidden_states.permute(0, 6, 1, 2, 4, 3, 5) # [2, 12, 28, 30, 2, 53, 2] + output = hidden_states.reshape(batch_size, -1, num_frames, height, width) # [2, 12, 28, 60, 106] + if USE_PEFT_BACKEND: # remove `lora_scale` from each PEFT layer unscale_lora_layers(self, lora_scale) diff --git a/fastvideo/train.py b/fastvideo/train.py index a1d48d680..93252a9b6 100644 --- a/fastvideo/train.py +++ b/fastvideo/train.py @@ -217,6 +217,7 @@ def main(args): main_print(f"--> loading model from {args.pretrained_model_name_or_path}") # keep the master weight to float32 + weight_type = torch.float32 if args.master_weight_type == 'fp32' else torch.bfloat16 transformer = MochiTransformer3DModel.from_pretrained( args.pretrained_model_name_or_path, subfolder="transformer",