diff --git a/docs/inference/support_matrix.md b/docs/inference/support_matrix.md index e0358b05e..2b2508822 100644 --- a/docs/inference/support_matrix.md +++ b/docs/inference/support_matrix.md @@ -73,6 +73,7 @@ pipeline initialization and sampling. | Matrix Game 2.0 Base Distilled | `FastVideo/Matrix-Game-2.0-Base-Distilled-Diffusers` | 352x640 | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | | Matrix Game 2.0 GTA Distilled | `FastVideo/Matrix-Game-2.0-GTA-Distilled-Diffusers` | 352x640 | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | | Matrix Game 2.0 TempleRun Distilled | `FastVideo/Matrix-Game-2.0-TempleRun-Distilled-Diffusers` | 352x640 | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | +| Matrix Game 3.0 Base Distilled | `FastVideo/Matrix-Game-3.0-Base-Distilled-Diffusers` | 720x1280 | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | | GEN3C Cosmos 7B | `FastVideo/GEN3C-Cosmos-7B-Diffusers` | 704px1280p | ❌ | ❌ | ❌ | ⭕ | ⭕ | **Note**: Wan2.2 TI2V 5B has some quality issues when performing I2V generation. We are working on fixing this issue. diff --git a/examples/inference/basic/basic_matrixgame3.py b/examples/inference/basic/basic_matrixgame3.py new file mode 100644 index 000000000..37783fdab --- /dev/null +++ b/examples/inference/basic/basic_matrixgame3.py @@ -0,0 +1,35 @@ +from fastvideo import VideoGenerator + +MODEL_PATH = "FastVideo/Matrix-Game-3.0-Base-Distilled-Diffusers" +IMAGE_URL = "https://raw.githubusercontent.com/SkyworkAI/Matrix-Game/main/Matrix-Game-3/demo_images/001/image.png" +PROMPT = "A colorful, animated cityscape with a gas station and various buildings." +OUTPUT_PATH = "video_samples_matrixgame3" + + +def main(): + generator = VideoGenerator.from_pretrained( + MODEL_PATH, + num_gpus=1, + use_fsdp_inference=False, + dit_cpu_offload=False, + vae_cpu_offload=False, + text_encoder_cpu_offload=True, + pin_cpu_memory=True, + ) + + generator.generate_video( + prompt=PROMPT, + image_path=IMAGE_URL, + height=720, + width=1280, + num_frames=57, + num_inference_steps=3, + guidance_scale=1.0, + seed=42, + output_path=OUTPUT_PATH, + save_video=True, + ) + + +if __name__ == "__main__": + main() diff --git a/fastvideo/api/matrixgame3.py b/fastvideo/api/matrixgame3.py new file mode 100644 index 000000000..cc3360289 --- /dev/null +++ b/fastvideo/api/matrixgame3.py @@ -0,0 +1,17 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass + +from fastvideo.api.sampling_param import SamplingParam + + +@dataclass +class MatrixGame3SamplingParam(SamplingParam): + height: int = 720 + width: int = 1280 + num_frames: int = 57 + fps: int = 25 + guidance_scale: float = 1.0 + num_inference_steps: int = 3 + negative_prompt: str = "" + num_iterations: int | None = None + use_base_model: bool = False diff --git a/fastvideo/configs/models/dits/matrixgame3.py b/fastvideo/configs/models/dits/matrixgame3.py new file mode 100644 index 000000000..bfe63654f --- /dev/null +++ b/fastvideo/configs/models/dits/matrixgame3.py @@ -0,0 +1,79 @@ +from dataclasses import dataclass, field + +import torch +from fastvideo.configs.models.dits.wanvideo import WanVideoArchConfig, WanVideoConfig + + +def _is_transformer_block(param_name: str, module: torch.nn.Module) -> bool: + return bool("blocks" in param_name and param_name.split(".")[-1].isdigit()) + + +@dataclass +class MatrixGame3WanVideoArchConfig(WanVideoArchConfig): + param_names_mapping: dict = field( + default_factory=lambda: { + r"^patch_embedding\.(weight|bias)$": r"patch_embedding.proj.\1", + r"^patch_embedding_wancamctrl\.(.*)$": r"camera_patch_embedding.proj.\1", + r"^time_embedding\.0\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_in.\1", + r"^time_embedding\.2\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_out.\1", + r"^time_projection\.1\.(.*)$": r"condition_embedder.time_modulation.linear.\1", + r"^head\.head\.(.*)$": r"proj_out.\1", + r"^head\.modulation$": r"scale_shift_table", + r"^blocks\.(\d+)\.self_attn\.q\.(.*)$": r"blocks.\1.to_q.\2", + r"^blocks\.(\d+)\.self_attn\.k\.(.*)$": r"blocks.\1.to_k.\2", + r"^blocks\.(\d+)\.self_attn\.v\.(.*)$": r"blocks.\1.to_v.\2", + r"^blocks\.(\d+)\.self_attn\.o\.(.*)$": r"blocks.\1.to_out.\2", + r"^blocks\.(\d+)\.self_attn\.norm_q\.(.*)$": r"blocks.\1.norm_q.\2", + r"^blocks\.(\d+)\.self_attn\.norm_k\.(.*)$": r"blocks.\1.norm_k.\2", + r"^blocks\.(\d+)\.cross_attn\.q\.(.*)$": r"blocks.\1.attn2.to_q.\2", + r"^blocks\.(\d+)\.cross_attn\.k\.(.*)$": r"blocks.\1.attn2.to_k.\2", + r"^blocks\.(\d+)\.cross_attn\.v\.(.*)$": r"blocks.\1.attn2.to_v.\2", + r"^blocks\.(\d+)\.cross_attn\.o\.(.*)$": r"blocks.\1.attn2.to_out.\2", + r"^blocks\.(\d+)\.cross_attn\.norm_q\.(.*)$": r"blocks.\1.attn2.norm_q.\2", + r"^blocks\.(\d+)\.cross_attn\.norm_k\.(.*)$": r"blocks.\1.attn2.norm_k.\2", + r"^blocks\.(\d+)\.ffn\.0\.(.*)$": r"blocks.\1.ffn.fc_in.\2", + r"^blocks\.(\d+)\.ffn\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2", + r"^blocks\.(\d+)\.norm3\.(.*)$": r"blocks.\1.self_attn_residual_norm.norm.\2", + r"^blocks\.(\d+)\.modulation$": r"blocks.\1.scale_shift_table", + }) + patch_size: tuple[int, int, int] = (1, 2, 2) + in_channels: int = 48 + out_channels: int = 48 + num_attention_heads: int = 24 + attention_head_dim: int = 128 + ffn_dim: int = 14336 + num_layers: int = 30 + text_len: int = 512 + image_dim: int = 0 + use_text_crossattn: bool = True + use_memory: bool = True + sigma_theta: float = 0.8 + camera_embed_in_channels: int = 1536 + action_config: dict = field( + default_factory=lambda: { + "blocks": list(range(15)), + "enable_mouse": True, + "enable_keyboard": True, + "heads_num": 16, + "hidden_size": 128, + "img_hidden_size": 3072, + "keyboard_dim_in": 6, + "keyboard_hidden_dim": 1024, + "mouse_dim_in": 2, + "mouse_hidden_dim": 1024, + "mouse_qk_dim_list": [8, 28, 28], + "patch_size": [1, 2, 2], + "qk_norm": True, + "qkv_bias": False, + "rope_dim_list": [8, 28, 28], + "rope_theta": 256, + "vae_time_compression_ratio": 4, + "windows_size": 3, + }) + + +@dataclass +class MatrixGame3WanVideoConfig(WanVideoConfig): + arch_config: MatrixGame3WanVideoArchConfig = field(default_factory=MatrixGame3WanVideoArchConfig) + prefix: str = "Wan" + _compile_conditions: list = field(default_factory=lambda: [_is_transformer_block]) diff --git a/fastvideo/configs/models/vaes/wanvae.py b/fastvideo/configs/models/vaes/wanvae.py index be08f8102..41b6da73c 100644 --- a/fastvideo/configs/models/vaes/wanvae.py +++ b/fastvideo/configs/models/vaes/wanvae.py @@ -76,5 +76,7 @@ class WanVAEConfig(VAEConfig): use_temporal_tiling: bool = False use_parallel_tiling: bool = False + use_light_vae: bool = False + def __post_init__(self): self.blend_num_frames = (self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames) * 2 diff --git a/fastvideo/configs/pipelines/__init__.py b/fastvideo/configs/pipelines/__init__.py index e60b5cf03..5d48639e1 100644 --- a/fastvideo/configs/pipelines/__init__.py +++ b/fastvideo/configs/pipelines/__init__.py @@ -6,6 +6,7 @@ from fastvideo.configs.pipelines.hunyuan15 import Hunyuan15T2V480PConfig, Hunyua from fastvideo.configs.pipelines.hunyuangamecraft import HunyuanGameCraftPipelineConfig from fastvideo.configs.pipelines.hyworld import HYWorldConfig from fastvideo.configs.pipelines.matrixgame2 import MatrixGame2I2V480PConfig +from fastvideo.configs.pipelines.matrixgame3 import MatrixGame3I2V720PConfig from fastvideo.pipelines.basic.ltx2.pipeline_configs import LTX2T2VConfig from fastvideo.registry import get_pipeline_config_cls_from_name from fastvideo.configs.pipelines.wan import (SelfForcingWanT2V480PConfig, WanI2V480PConfig, WanI2V720PConfig, @@ -15,5 +16,5 @@ __all__ = [ "HunyuanConfig", "FastHunyuanConfig", "HunyuanGameCraftPipelineConfig", "PipelineConfig", "Hunyuan15T2V480PConfig", "Hunyuan15T2V720PConfig", "WanT2V480PConfig", "WanI2V480PConfig", "WanT2V720PConfig", "WanI2V720PConfig", "SelfForcingWanT2V480PConfig", "CosmosConfig", "Cosmos25Config", "LTX2T2VConfig", "HYWorldConfig", - "MatrixGame2I2V480PConfig", "get_pipeline_config_cls_from_name" + "MatrixGame2I2V480PConfig", "MatrixGame3I2V720PConfig", "get_pipeline_config_cls_from_name" ] diff --git a/fastvideo/configs/pipelines/matrixgame3.py b/fastvideo/configs/pipelines/matrixgame3.py new file mode 100644 index 000000000..e226710b6 --- /dev/null +++ b/fastvideo/configs/pipelines/matrixgame3.py @@ -0,0 +1,18 @@ +# SPDX-License-Identifier: Apache-2.0 +from dataclasses import dataclass, field + +from fastvideo.configs.models import DiTConfig +from fastvideo.configs.models.dits.matrixgame3 import MatrixGame3WanVideoConfig +from fastvideo.configs.pipelines.wan import WanT2V480PConfig + + +@dataclass +class MatrixGame3I2V720PConfig(WanT2V480PConfig): + dit_config: DiTConfig = field(default_factory=MatrixGame3WanVideoConfig) + flow_shift: float | None = 5.0 + vae_precision: str = "fp32" + + def __post_init__(self) -> None: + self.vae_config.load_encoder = True + self.vae_config.load_decoder = True + self.vae_config.use_light_vae = True diff --git a/fastvideo/models/dits/matrixgame3/__init__.py b/fastvideo/models/dits/matrixgame3/__init__.py new file mode 100644 index 000000000..3105caf38 --- /dev/null +++ b/fastvideo/models/dits/matrixgame3/__init__.py @@ -0,0 +1,11 @@ +from .action_module import MatrixGame3ActionModule +from .model import MatrixGame3CrossAttention, MatrixGame3TransformerBlock, MatrixGame3WanModel + +__all__ = [ + "MatrixGame3WanModel", + "MatrixGame3TransformerBlock", + "MatrixGame3CrossAttention", + "MatrixGame3ActionModule", +] + +EntryClass = [MatrixGame3WanModel] diff --git a/fastvideo/models/dits/matrixgame3/action_module.py b/fastvideo/models/dits/matrixgame3/action_module.py new file mode 100644 index 000000000..126ac815a --- /dev/null +++ b/fastvideo/models/dits/matrixgame3/action_module.py @@ -0,0 +1,515 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from Matrix-Game-3.0: https://github.com/SkyworkAI/Matrix-Game/blob/main/Matrix-Game-3/wan/modules/action_module.py + +from einops import rearrange +import torch +import torch.nn as nn + +from fastvideo.attention import LocalAttention +from fastvideo.layers.linear import ReplicatedLinear +from fastvideo.layers.layernorm import FP32LayerNorm +from fastvideo.layers.rotary_embedding import ( + get_nd_rotary_pos_embed as _fv_get_nd_rotary_pos_embed, + _apply_rotary_emb, +) +from fastvideo.platforms import AttentionBackendEnum + + +class WanRMSNorm(nn.Module): + + def __init__(self, dim: int, eps: float = 1e-5): + super().__init__() + self.dim = dim + self.eps = eps + self.weight = nn.Parameter(torch.ones(dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps) + + +def _get_nd_rotary_pos_embed_matrixgame( + rope_dim_list, + rope_sizes, + theta: float = 10000.0, + theta_rescale_factor: float = 1.0, +): + cos, sin = _fv_get_nd_rotary_pos_embed( + rope_dim_list, + rope_sizes, + theta=theta, + theta_rescale_factor=theta_rescale_factor, + dtype=torch.float32, + ) + # convert from [S, D/2] to [S, D] format + cos = cos.repeat_interleave(2, dim=1) + sin = sin.repeat_interleave(2, dim=1) + return cos, sin + + +def _apply_rotary_emb_qk( + xq: torch.Tensor, + xk: torch.Tensor, + freqs_cos: torch.Tensor, + freqs_sin: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + seq_len = xq.shape[1] + cos = freqs_cos[:seq_len] + sin = freqs_sin[:seq_len] + cos_half = cos[:, ::2] + sin_half = sin[:, ::2] + xq_out = _apply_rotary_emb(xq, cos_half, sin_half, is_neox_style=False) + xk_out = _apply_rotary_emb(xk, cos_half, sin_half, is_neox_style=False) + return xq_out, xk_out + + +class MatrixGame3ActionModule(nn.Module): + """ + action module from https://arxiv.org/pdf/2501.08325 + """ + + def __init__( + self, + mouse_dim_in: int = 2, + keyboard_dim_in: int = 6, + hidden_size: int = 128, + img_hidden_size: int = 1536, + keyboard_hidden_dim: int = 1024, + mouse_hidden_dim: int = 1024, + vae_time_compression_ratio: int = 4, + windows_size: int = 3, + heads_num: int = 16, + patch_size: list | None = None, + qk_norm: bool = True, + qkv_bias: bool = False, + rope_dim_list: list | None = None, + rope_theta=256, + mouse_qk_dim_list: list | None = None, + enable_mouse=True, + enable_keyboard=True, + blocks: list | None = None, + ): + super().__init__() + patch_size = patch_size if patch_size is not None else [1, 2, 2] + rope_dim_list = ( + rope_dim_list if rope_dim_list is not None else [8, 28, 28] + ) + mouse_qk_dim_list = ( + mouse_qk_dim_list if mouse_qk_dim_list is not None else [8, 28, 28] + ) + blocks = blocks if blocks is not None else [] + self.enable_mouse = enable_mouse + self.enable_keyboard = enable_keyboard + + self.rope_dim_list = rope_dim_list + self.rope_theta = rope_theta + if self.enable_keyboard: + self.keyboard_embed = nn.Sequential( + nn.Linear(keyboard_dim_in, hidden_size, bias=True), + nn.SiLU(), + nn.Linear(hidden_size, hidden_size, bias=True), + ) + + self.mouse_qk_dim_list = mouse_qk_dim_list + self.heads_num = heads_num + if self.enable_mouse: + c = mouse_hidden_dim + self.mouse_mlp = nn.Sequential( + nn.Linear( + mouse_dim_in * vae_time_compression_ratio * windows_size + + img_hidden_size, + c, + bias=True, + ), + nn.GELU(approximate="tanh"), + nn.Linear(c, c), + FP32LayerNorm(c, elementwise_affine=True), + ) + + head_dim = c // heads_num + self.t_qkv = ReplicatedLinear(c, c * 3, bias=qkv_bias) + self.img_attn_q_norm = ( + WanRMSNorm(head_dim, eps=1e-6) if qk_norm else nn.Identity() + ) + self.img_attn_k_norm = ( + WanRMSNorm(head_dim, eps=1e-6) if qk_norm else nn.Identity() + ) + self.proj_mouse = ReplicatedLinear( + c, img_hidden_size, bias=qkv_bias + ) + + if self.enable_keyboard: + head_dim_key = keyboard_hidden_dim // heads_num + self.key_attn_q_norm = ( + WanRMSNorm(head_dim_key, eps=1e-6) if qk_norm else nn.Identity() + ) + self.key_attn_k_norm = ( + WanRMSNorm(head_dim_key, eps=1e-6) if qk_norm else nn.Identity() + ) + + self.mouse_attn_q = ReplicatedLinear( + img_hidden_size, keyboard_hidden_dim, bias=qkv_bias + ) + self.keyboard_attn_kv = ReplicatedLinear( + hidden_size * windows_size * vae_time_compression_ratio, + keyboard_hidden_dim * 2, + bias=qkv_bias, + ) + self.proj_keyboard = ReplicatedLinear( + keyboard_hidden_dim, img_hidden_size, bias=qkv_bias + ) + + self.mouse_attn_layer = ( + LocalAttention( + num_heads=heads_num, + head_size=mouse_hidden_dim // heads_num, + causal=False, + supported_attention_backends=( + AttentionBackendEnum.FLASH_ATTN, + AttentionBackendEnum.TORCH_SDPA, + ), + ) + if self.enable_mouse + else None + ) + + self.keyboard_attn_layer = ( + LocalAttention( + num_heads=heads_num, + head_size=keyboard_hidden_dim // heads_num, + causal=False, + supported_attention_backends=( + AttentionBackendEnum.FLASH_ATTN, + AttentionBackendEnum.TORCH_SDPA, + ), + ) + if self.enable_keyboard + else None + ) + + self.vae_time_compression_ratio = vae_time_compression_ratio + self.windows_size = windows_size + self.patch_size = patch_size + # Lazy initialization: freqs will be created on first forward pass + self._freqs_cos = None + self._freqs_sin = None + + def patchify(self, x, patch_size): + """ + x : (N C T H W) + """ + pt, ph, pw = self.patch_size + t, h, w = x.shape[2] // pt, x.shape[3] // ph, x.shape[4] // pw + c = x.shape[1] + x = x.reshape(shape=(x.shape[0], c, t, pt, h, ph, w, pw)) + x = torch.einsum("nctohpwq->nthwcopq", x) + x = x.reshape(shape=(x.shape[0], t * h * w, c * pt * ph * pw)) + return x + + def unpatchify(self, x, t, h, w, patch_size): + """ + x: (N, T, patch_size**2 * C) + imgs: (N, H, W, C) + """ + c = x.shape[2] // patch_size + pt, ph, pw = self.patch_size + assert t * h * w == x.shape[1] + + x = x.reshape(shape=(x.shape[0], t, h, w, c, pt, ph, pw)) + x = torch.einsum("nthwcopq->nctohpwq", x) + imgs = x.reshape(shape=(x.shape[0], c, t * pt, h * ph, w * pw)) + + return imgs + + def get_rotary_pos_embed( + self, + video_length, + height, + width, + head_dim, + rope_dim_list=None, + ): + target_ndim = 3 + ndim = 5 - 2 + + latents_size = [video_length, height, width] + + if isinstance(self.patch_size, int): + assert all(s % self.patch_size == 0 for s in latents_size), ( + f"Latent size(last {ndim} dimensions) should be divisible by patch size({self.patch_size}), " + f"but got {latents_size}." + ) + rope_sizes = [s // self.patch_size for s in latents_size] + elif isinstance(self.patch_size, list): + assert all( + s % self.patch_size[idx] == 0 + for idx, s in enumerate(latents_size) + ), ( + f"Latent size(last {ndim} dimensions) should be divisible by patch size({self.patch_size}), " + f"but got {latents_size}." + ) + rope_sizes = [ + s // self.patch_size[idx] for idx, s in enumerate(latents_size) + ] + + if len(rope_sizes) != target_ndim: + rope_sizes = [1] * ( + target_ndim - len(rope_sizes) + ) + rope_sizes # time axis + + if rope_dim_list is None: + rope_dim_list = [ + head_dim // target_ndim for _ in range(target_ndim) + ] + assert sum(rope_dim_list) == head_dim, ( + "sum(rope_dim_list) should equal to head_dim of attention layer" + ) + # Use Matrix-Game wrapper for FastVideo's function + freqs_cos, freqs_sin = _get_nd_rotary_pos_embed_matrixgame( + rope_dim_list, + rope_sizes, + theta=self.rope_theta, + theta_rescale_factor=1, + ) + return freqs_cos[ + -video_length + * rope_sizes[1] + * rope_sizes[2] + // self.patch_size[0] : + ], freqs_sin[ + -video_length + * rope_sizes[1] + * rope_sizes[2] + // self.patch_size[0] : + ] + + def _forward_mouse( + self, + hidden_states: torch.Tensor, + mouse_condition: torch.Tensor, + mouse_cond_memory: torch.Tensor | None, + *, + pad_t: int, + freqs_cis: tuple[torch.Tensor, torch.Tensor], + N_feats: int, + B: int, + C: int, + th: int, + tw: int, + ) -> torch.Tensor: + pad = mouse_condition[:, 0:1, :].expand(-1, pad_t, -1) + mouse_condition = torch.cat([pad, mouse_condition], dim=1) + group_mouse = [ + mouse_condition[ + :, + self.vae_time_compression_ratio * (i - self.windows_size) + + pad_t : i * self.vae_time_compression_ratio + pad_t, + :, + ] + for i in range(N_feats) + ] + group_mouse = torch.stack(group_mouse, dim=1) + mem_len = 0 + if mouse_cond_memory is not None: + mem_len = mouse_cond_memory.shape[1] + mem = mouse_cond_memory.to(device=group_mouse.device, dtype=group_mouse.dtype) + mem = mem.unsqueeze(2).expand(-1, -1, pad_t, -1) + group_mouse = torch.cat([mem, group_mouse], dim=1) + actual_num_frames = group_mouse.shape[1] + + S = th * tw + group_mouse = group_mouse.unsqueeze(-1).expand( + B, actual_num_frames, pad_t, C, S + ) + group_mouse = group_mouse.permute(0, 4, 1, 2, 3).reshape( + B * S, actual_num_frames, pad_t * C + ) + + group_mouse = torch.cat([hidden_states, group_mouse], dim=-1) + group_mouse = self.mouse_mlp(group_mouse) + mouse_qkv, _ = self.t_qkv(group_mouse) + q, k, v = rearrange( + mouse_qkv, "B L (K H D) -> K B L H D", K=3, H=self.heads_num + ) + q = self.img_attn_q_norm(q).to(v) + k = self.img_attn_k_norm(k).to(v) + if mem_len > 0: + q_mem, k_mem = _apply_rotary_emb_qk(q[:, :mem_len], k[:, :mem_len], freqs_cis[0], freqs_cis[1]) + q_pred, k_pred = _apply_rotary_emb_qk(q[:, mem_len:], k[:, mem_len:], freqs_cis[0], freqs_cis[1]) + q = torch.cat([q_mem, q_pred], dim=1) + k = torch.cat([k_mem, k_pred], dim=1) + else: + q, k = _apply_rotary_emb_qk(q, k, freqs_cis[0], freqs_cis[1]) + + attn = self.mouse_attn_layer(q, k, v) + attn = rearrange(attn, "(b S) T h d -> b (T S) (h d)", b=B) + attn, _ = self.proj_mouse(attn) + return attn + + def _forward_keyboard( + self, + hidden_states: torch.Tensor, + keyboard_condition: torch.Tensor, + keyboard_cond_memory: torch.Tensor | None, + *, + pad_t: int, + freqs_cis: tuple[torch.Tensor, torch.Tensor], + N_feats: int, + B: int, + th: int, + tw: int, + ) -> torch.Tensor: + pad = keyboard_condition[:, 0:1, :].expand(-1, pad_t, -1) + keyboard_condition = torch.cat([pad, keyboard_condition], dim=1) + keyboard_condition = self.keyboard_embed(keyboard_condition) + group_keyboard = [ + keyboard_condition[ + :, + self.vae_time_compression_ratio * (i - self.windows_size) + + pad_t : i * self.vae_time_compression_ratio + pad_t, + :, + ] + for i in range(N_feats) + ] + group_keyboard = torch.stack(group_keyboard, dim=1) + mem_len = 0 + if keyboard_cond_memory is not None: + mem_len = keyboard_cond_memory.shape[1] + k_mem = keyboard_cond_memory.to( + device=group_keyboard.device, dtype=group_keyboard.dtype + ) + k_mem = self.keyboard_embed(k_mem) + k_mem = k_mem.unsqueeze(2).expand(-1, -1, pad_t, -1) + group_keyboard = torch.cat([k_mem, group_keyboard], dim=1) + group_keyboard = group_keyboard.reshape( + shape=(group_keyboard.shape[0], group_keyboard.shape[1], -1) + ) + + mouse_q, _ = self.mouse_attn_q(hidden_states) + keyboard_kv, _ = self.keyboard_attn_kv(group_keyboard) + + B, L, HD = mouse_q.shape + D = HD // self.heads_num + q = mouse_q.view(B, L, self.heads_num, D) + + B, L, KHD = keyboard_kv.shape + k, v = keyboard_kv.view(B, L, 2, self.heads_num, D).permute( + 2, 0, 1, 3, 4 + ) + + q = self.key_attn_q_norm(q).to(v) + k = self.key_attn_k_norm(k).to(v) + S = th * tw + + B, TS, H, D = q.shape + T_ = TS // S + q = q.view(B, T_, S, H, D).transpose(1, 2).reshape(B * S, T_, H, D) + if mem_len > 0: + q_mem, k_mem_r = _apply_rotary_emb_qk(q[:, :mem_len], k[:, :mem_len], freqs_cis[0], freqs_cis[1]) + q_pred, k_pred_r = _apply_rotary_emb_qk(q[:, mem_len:], k[:, mem_len:], freqs_cis[0], freqs_cis[1]) + q = torch.cat([q_mem, q_pred], dim=1) + k = torch.cat([k_mem_r, k_pred_r], dim=1) + else: + q, k = _apply_rotary_emb_qk(q, k, freqs_cis[0], freqs_cis[1]) + + k = k.repeat_interleave(S, dim=0) + v = v.repeat_interleave(S, dim=0) + + attn = self.keyboard_attn_layer(q, k, v) + attn = rearrange(attn, "(B S) T H D -> B (T S) (H D)", S=S) + attn, _ = self.proj_keyboard(attn) + return attn + + def forward( + self, + x: torch.Tensor, + tt: int, + th: int, + tw: int, + mouse_condition: torch.Tensor | None = None, + keyboard_condition: torch.Tensor | None = None, + mouse_cond_memory: torch.Tensor | None = None, + keyboard_cond_memory: torch.Tensor | None = None, + ): + """ + hidden_states: B, tt*th*tw, C + mouse_condition: B, N_frames, C1 + keyboard_condition: B, N_frames, C2 + """ + target_device = x.device + target_dtype = x.dtype + if mouse_condition is not None: + mouse_condition = mouse_condition.to( + device=target_device, dtype=target_dtype + ) + if keyboard_condition is not None: + keyboard_condition = keyboard_condition.to( + device=target_device, dtype=target_dtype + ) + else: + return x + + B, N_frames, C = keyboard_condition.shape + assert tt * th * tw == x.shape[1] + assert ( + ((N_frames - 1) + self.vae_time_compression_ratio) % self.vae_time_compression_ratio == 0 + or (N_frames % self.vae_time_compression_ratio == 0) + ) + if ((N_frames - 1) + self.vae_time_compression_ratio) % self.vae_time_compression_ratio == 0: + N_feats = int((N_frames - 1) / self.vae_time_compression_ratio) + 1 + else: + N_feats = N_frames // self.vae_time_compression_ratio + + # Lazy initialization of freqs + if (self._freqs_cos is None or self._freqs_sin is None + or self._freqs_cos.device != x.device): + fc, fs = self.get_rotary_pos_embed( + 7500, + self.patch_size[1], + self.patch_size[2], + 64, + self.mouse_qk_dim_list, + ) + self._freqs_cos = fc.to(x.device) + self._freqs_sin = fs.to(x.device) + + freqs_cis = (self._freqs_cos, self._freqs_sin) + + if self.enable_mouse and mouse_condition is not None: + hidden_states = rearrange( + x, "B (T S) C -> (B S) T C", T=tt, S=th * tw + ) + B, N_frames, C = mouse_condition.shape + else: + hidden_states = x + + pad_t = self.vae_time_compression_ratio * self.windows_size + if self.enable_mouse and mouse_condition is not None: + attn = self._forward_mouse( + hidden_states, + mouse_condition, + mouse_cond_memory, + pad_t=pad_t, + freqs_cis=freqs_cis, + N_feats=N_feats, + B=B, + C=C, + th=th, + tw=tw, + ) + hidden_states = rearrange(x, "(B S) T C -> B (T S) C", B=B) + hidden_states = hidden_states + attn + + if self.enable_keyboard and keyboard_condition is not None: + attn = self._forward_keyboard( + hidden_states, + keyboard_condition, + keyboard_cond_memory, + pad_t=pad_t, + freqs_cis=freqs_cis, + N_feats=N_feats, + B=B, + th=th, + tw=tw, + ) + hidden_states = hidden_states + attn + return hidden_states diff --git a/fastvideo/models/dits/matrixgame3/model.py b/fastvideo/models/dits/matrixgame3/model.py new file mode 100644 index 000000000..bd800c558 --- /dev/null +++ b/fastvideo/models/dits/matrixgame3/model.py @@ -0,0 +1,797 @@ +# SPDX-License-Identifier: Apache-2.0 + +import math +from typing import Any + +import torch +import torch.nn as nn + +from fastvideo.attention import DistributedAttention +from fastvideo.configs.models.dits.matrixgame3 import MatrixGame3WanVideoConfig +from fastvideo.layers.layernorm import ( + FP32LayerNorm, + LayerNormScaleShift, + RMSNorm, + ScaleResidual, + ScaleResidualLayerNormScaleShift, +) +from fastvideo.layers.linear import ReplicatedLinear +from fastvideo.layers.mlp import MLP +from fastvideo.layers.visual_embedding import ( + PatchEmbed, + TimestepEmbedder, + ModulateProjection, + WanCamControlPatchEmbedding, +) +from fastvideo.logger import init_logger +from fastvideo.models.dits.base import BaseDiT +from fastvideo.models.dits.wanvideo import WanSelfAttention +from fastvideo.platforms import AttentionBackendEnum + +# Import ActionModule +from .action_module import MatrixGame3ActionModule + +logger = init_logger(__name__) + + +def _build_rope_freqs( + *, + max_seq_len: int, + head_dim: int, + num_heads: int, + sigma_theta: float = 0.0, + device: torch.device | None = None, +) -> torch.Tensor: + c = head_dim // 2 + c_t = c - 2 * (c // 3) + c_h = c // 3 + c_w = c // 3 + + if sigma_theta > 0: + rope_epsilon = torch.linspace(-1, 1, num_heads, dtype=torch.float64, device=device) + theta_hat = 10000.0 * (1 + sigma_theta * rope_epsilon) + + def build_axis_freqs(seq_len: int, c_part: int) -> torch.Tensor: + exp = torch.arange(c_part, dtype=torch.float64, device=device) / c_part + omega = 1.0 / torch.pow(theta_hat.unsqueeze(1), exp.unsqueeze(0)) + pos = torch.arange(seq_len, dtype=torch.float64, device=device) + angles = pos.view(1, -1, 1) * omega.unsqueeze(1) + return torch.polar(torch.ones_like(angles), angles) + + return torch.cat( + [ + build_axis_freqs(max_seq_len, c_t), + build_axis_freqs(max_seq_len, c_h), + build_axis_freqs(max_seq_len, c_w), + ], + dim=2, + ) + + def build_axis_freqs(seq_len: int, c_part: int) -> torch.Tensor: + exp = torch.arange(c_part, dtype=torch.float64, device=device) / c_part + omega = 1.0 / torch.pow(torch.tensor(10000.0, dtype=torch.float64, device=device), exp) + pos = torch.arange(seq_len, dtype=torch.float64, device=device) + angles = pos[:, None] * omega[None, :] + return torch.polar(torch.ones_like(angles), angles) + + return torch.cat( + [ + build_axis_freqs(max_seq_len, c_t), + build_axis_freqs(max_seq_len, c_h), + build_axis_freqs(max_seq_len, c_w), + ], + dim=1, + ) + + +@torch.amp.autocast("cuda", enabled=False) +def _apply_rope_with_frame_indices( + x: torch.Tensor, + freqs: torch.Tensor, + *, + height: int, + width: int, + frame_indices: list[int] | tuple[int, ...] | torch.Tensor, +) -> torch.Tensor: + num_heads = x.shape[2] + half_dim = x.shape[-1] // 2 + c_t = half_dim - 2 * (half_dim // 3) + c_h = half_dim // 3 + c_w = half_dim // 3 + + if not torch.is_tensor(frame_indices): + frame_indices = torch.tensor(frame_indices, device=freqs.device, dtype=torch.long) + else: + frame_indices = frame_indices.to(device=freqs.device, dtype=torch.long) + + f = int(frame_indices.numel()) + seq_len = f * height * width + x_complex = torch.view_as_complex(x.to(torch.float32).reshape(x.shape[0], seq_len, num_heads, -1, 2)) + + if freqs.dim() == 3: + t_freqs = freqs[:, frame_indices, :c_t] + h_freqs = freqs[:, :height, c_t:c_t + c_h] + w_freqs = freqs[:, :width, c_t + c_h:c_t + c_h + c_w] + freqs_i = torch.cat( + [ + t_freqs.permute(1, 0, 2).view(f, 1, 1, num_heads, -1).expand(f, height, width, num_heads, -1), + h_freqs.permute(1, 0, 2).view(1, height, 1, num_heads, -1).expand(f, height, width, num_heads, -1), + w_freqs.permute(1, 0, 2).view(1, 1, width, num_heads, -1).expand(f, height, width, num_heads, -1), + ], + dim=-1, + ).reshape(seq_len, num_heads, -1) + else: + t_freqs = freqs[frame_indices, :c_t] + h_freqs = freqs[:height, c_t:c_t + c_h] + w_freqs = freqs[:width, c_t + c_h:c_t + c_h + c_w] + freqs_i = torch.cat( + [ + t_freqs.view(f, 1, 1, -1).expand(f, height, width, -1), + h_freqs.view(1, height, 1, -1).expand(f, height, width, -1), + w_freqs.view(1, 1, width, -1).expand(f, height, width, -1), + ], + dim=-1, + ).reshape(seq_len, 1, -1) + + out = torch.view_as_real(x_complex * freqs_i.unsqueeze(0).to(x_complex.dtype)).flatten(3) + return out.to(dtype=x.dtype) + + +class MatrixGame3TimeImageEmbedding(nn.Module): + def __init__( + self, + dim: int, + time_freq_dim: int, + ): + super().__init__() + + self.time_embedder = TimestepEmbedder( + dim, frequency_embedding_size=time_freq_dim, act_layer="silu" + ) + self.time_modulation = ModulateProjection( + dim, factor=6, act_layer="silu" + ) + + def forward( + self, + timestep: torch.Tensor, + encoder_hidden_states: torch.Tensor | None, + timestep_seq_len: int | None = None, + ): + temb = self.time_embedder(timestep, timestep_seq_len) + timestep_proj = self.time_modulation(temb) + + if encoder_hidden_states is None: + batch_size = temb.shape[0] if temb.dim() > 1 else timestep.shape[0] + encoder_hidden_states = torch.zeros( + (batch_size, 0, temb.shape[-1]), + device=temb.device, + dtype=temb.dtype, + ) + + return temb, timestep_proj, encoder_hidden_states + + +class MatrixGame3CrossAttention(WanSelfAttention): + def forward(self, x, context, context_lens=None): + r""" + Args: + x(Tensor): Shape [B, L1, C] + context(Tensor): Shape [B, L2, C] - typically 257 image tokens + context_lens(Tensor): Shape [B] + """ + b, n, d = x.size(0), self.num_heads, self.head_dim + + q_input = x.to(self.to_q.weight.dtype) + q = self.norm_q(self.to_q(q_input)[0]).view(b, -1, n, d) + + context_input = context.to(self.to_k.weight.dtype) + k = self.norm_k(self.to_k(context_input)[0]).view(b, -1, n, d) + v = self.to_v(context_input)[0].view(b, -1, n, d) + + x = self.attn(q, k, v) + + x = x.flatten(2) + x, _ = self.to_out(x) + return x + + +class MatrixGame3TransformerBlock(nn.Module): + def __init__( + self, + dim: int, + ffn_dim: int, + num_heads: int, + qk_norm: str = "rms_norm_across_heads", + cross_attn_norm: bool = False, + eps: float = 1e-6, + supported_attention_backends: tuple[AttentionBackendEnum, ...] + | None = None, + prefix: str = "", + action_config: dict | None = None, + block_id: int | None = None, + use_memory: bool = False, + ): + super().__init__() + action_config = action_config or {} + self.use_memory = use_memory + + # 1. Self-attention + self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) + self.to_q = ReplicatedLinear(dim, dim, bias=True) + self.to_k = ReplicatedLinear(dim, dim, bias=True) + self.to_v = ReplicatedLinear(dim, dim, bias=True) + + self.to_out = ReplicatedLinear(dim, dim, bias=True) + self.attn1 = DistributedAttention( + num_heads=num_heads, + head_size=dim // num_heads, + causal=False, + supported_attention_backends=supported_attention_backends, + prefix=f"{prefix}.attn1", + ) + self.hidden_dim = dim + self.num_attention_heads = num_heads + dim_head = dim // num_heads + if qk_norm == "rms_norm": + self.norm_q = RMSNorm(dim_head, eps=eps) + self.norm_k = RMSNorm(dim_head, eps=eps) + elif qk_norm == "rms_norm_across_heads": + # LTX applies qk norm across all heads + self.norm_q = RMSNorm(dim, eps=eps) + self.norm_k = RMSNorm(dim, eps=eps) + else: + print("QK Norm type not supported") + raise Exception + assert cross_attn_norm is True + self.self_attn_residual_norm = ScaleResidualLayerNormScaleShift( + dim, + norm_type="layer", + eps=eps, + elementwise_affine=True, + dtype=torch.float32, + compute_dtype=torch.float32, + ) + + # 2. Cross-attention + self.attn2 = MatrixGame3CrossAttention( + dim, num_heads, qk_norm=qk_norm, eps=eps + ) + self.cross_attn_residual_norm = ScaleResidualLayerNormScaleShift( + dim, + norm_type="layer", + eps=eps, + elementwise_affine=False, + dtype=torch.float32, + compute_dtype=torch.float32, + ) + + # 2.1. Action Module Integration + enabled_blocks = set(action_config.get("blocks", [])) if action_config else set() + self.use_action_module = len(action_config) > 0 and (block_id is None or block_id in enabled_blocks) + if self.use_action_module: + self.action_model = MatrixGame3ActionModule(**action_config) + else: + self.action_model = None + + # 3. Feed-forward + self.ffn = MLP(dim, ffn_dim, act_type="gelu_pytorch_tanh") + self.mlp_residual = ScaleResidual() + + self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + if self.use_memory: + self.cam_injector_layer1 = nn.Linear(dim, dim) + self.cam_injector_layer2 = nn.Linear(dim, dim) + self.cam_scale_layer = nn.Linear(dim, dim) + self.cam_shift_layer = nn.Linear(dim, dim) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor, + freqs: torch.Tensor, + # Action Module specific args + grid_sizes: tuple[int, int, int], + mouse_cond: torch.Tensor | None = None, + keyboard_cond: torch.Tensor | None = None, + mouse_cond_memory: torch.Tensor | None = None, + keyboard_cond_memory: torch.Tensor | None = None, + plucker_emb: torch.Tensor | None = None, + memory_length: int = 0, + memory_latent_idx: list[int] | tuple[int, ...] | torch.Tensor | None = None, + predict_latent_idx: tuple[int, int] | list[int] | tuple[int, ...] | torch.Tensor | None = None, + ) -> torch.Tensor: + if hidden_states.dim() == 4: + hidden_states = hidden_states.squeeze(1) + orig_dtype = hidden_states.dtype + + if temb.dim() == 4: + # temb: batch_size, seq_len, 6, inner_dim (wan2.2 ti2v) + ( + shift_msa, + scale_msa, + gate_msa, + c_shift_msa, + c_scale_msa, + c_gate_msa, + ) = (self.scale_shift_table.unsqueeze(0) + temb.float()).chunk( + 6, dim=2 + ) + # batch_size, seq_len, 1, inner_dim + shift_msa = shift_msa.squeeze(2) + scale_msa = scale_msa.squeeze(2) + gate_msa = gate_msa.squeeze(2) + c_shift_msa = c_shift_msa.squeeze(2) + c_scale_msa = c_scale_msa.squeeze(2) + c_gate_msa = c_gate_msa.squeeze(2) + else: + # temb: batch_size, 6, inner_dim (wan2.1/wan2.2 14B) + e = self.scale_shift_table + temb.float() + ( + shift_msa, + scale_msa, + gate_msa, + c_shift_msa, + c_scale_msa, + c_gate_msa, + ) = e.chunk(6, dim=1) + assert shift_msa.dtype == torch.float32 + + # 1. Self-attention + norm_hidden_states = ( + self.norm1(hidden_states).float() * (1 + scale_msa) + shift_msa + ).to(orig_dtype) + query, _ = self.to_q(norm_hidden_states) + key, _ = self.to_k(norm_hidden_states) + value, _ = self.to_v(norm_hidden_states) + + if self.norm_q is not None: + query = self.norm_q(query) + if self.norm_k is not None: + key = self.norm_k(key) + + query = query.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) + key = key.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) + value = value.squeeze(1).unflatten(2, (self.num_attention_heads, -1)) + + # Apply rotary embeddings + grid_frames, grid_height, grid_width = grid_sizes + + if memory_length > 0: + hw = grid_height * grid_width + query_memory = query[:, :memory_length * hw] + key_memory = key[:, :memory_length * hw] + query_pred = query[:, memory_length * hw:] + key_pred = key[:, memory_length * hw:] + + # Indices are pre-built device tensors by MatrixGame3WanModel.forward. + mem_indices = memory_latent_idx + if mem_indices is None: + mem_indices = torch.arange(memory_length, device=query.device, dtype=torch.long) + pred_indices = predict_latent_idx + if pred_indices is None: + pred_indices = torch.arange(grid_frames - memory_length, device=query.device, dtype=torch.long) + + query_memory = _apply_rope_with_frame_indices( + query_memory, freqs, height=grid_height, width=grid_width, frame_indices=mem_indices + ) + key_memory = _apply_rope_with_frame_indices( + key_memory, freqs, height=grid_height, width=grid_width, frame_indices=mem_indices + ) + query_pred = _apply_rope_with_frame_indices( + query_pred, freqs, height=grid_height, width=grid_width, frame_indices=pred_indices + ) + key_pred = _apply_rope_with_frame_indices( + key_pred, freqs, height=grid_height, width=grid_width, frame_indices=pred_indices + ) + query = torch.cat([query_memory, query_pred], dim=1) + key = torch.cat([key_memory, key_pred], dim=1) + else: + pred_indices = predict_latent_idx + if pred_indices is None: + pred_indices = torch.arange(grid_frames, device=query.device, dtype=torch.long) + query = _apply_rope_with_frame_indices( + query, freqs, height=grid_height, width=grid_width, frame_indices=pred_indices + ) + key = _apply_rope_with_frame_indices( + key, freqs, height=grid_height, width=grid_width, frame_indices=pred_indices + ) + + attn_output, _ = self.attn1(query, key, value) + attn_output = attn_output.flatten(2) + attn_output = attn_output.to(self.to_out.weight.dtype) + attn_output, _ = self.to_out(attn_output) + attn_output = attn_output.squeeze(1) + + hidden_states = hidden_states.float() + attn_output.float() * gate_msa.float() + if self.use_memory and plucker_emb is not None: + c2ws_hidden_states = self.cam_injector_layer2( + torch.nn.functional.silu(self.cam_injector_layer1(plucker_emb)) + ) + c2ws_hidden_states = c2ws_hidden_states + plucker_emb + cam_scale = self.cam_scale_layer(c2ws_hidden_states) + cam_shift = self.cam_shift_layer(c2ws_hidden_states) + hidden_states = (1.0 + cam_scale) * hidden_states + cam_shift + norm_hidden_states = self.self_attn_residual_norm.norm(hidden_states) + + # 2. Cross-attention + attn_output = self.attn2( + norm_hidden_states, context=encoder_hidden_states, context_lens=None + ) + + # - use_memory / action path: x = norm3(x); x = x + cross_attn(x) + # - plain path: x = x + cross_attn(norm3(x)) + cross_attn_residual_base = ( + norm_hidden_states if (mouse_cond is not None or self.use_memory) else hidden_states + ) + hidden_states = cross_attn_residual_base + attn_output + norm_hidden_states = self.cross_attn_residual_norm.norm( + hidden_states.float() + ) + norm_hidden_states = norm_hidden_states * (1 + c_scale_msa) + c_shift_msa + + # ================= Action Module ================= + if self.action_model is not None and (mouse_cond is not None or keyboard_cond is not None): + action_dtype = self.ffn.fc_in.weight.dtype + # grid_sizes is expected to be [F, H, W] + hidden_states = self.action_model( + hidden_states.to(action_dtype), + grid_sizes[0], + grid_sizes[1], + grid_sizes[2], + mouse_cond, + keyboard_cond, + mouse_cond_memory=mouse_cond_memory, + keyboard_cond_memory=keyboard_cond_memory, + ) + + norm_hidden_states = self.cross_attn_residual_norm.norm( + hidden_states.float() + ) + norm_hidden_states = norm_hidden_states * (1 + c_scale_msa) + c_shift_msa + # ================================================= + + # 3. Feed-forward + ff_output = self.ffn(norm_hidden_states.to(self.ffn.fc_in.weight.dtype)) + hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa) + hidden_states = hidden_states.to(orig_dtype) + + return hidden_states + + +_DEFAULT_MATRIXGAME3_CONFIG = MatrixGame3WanVideoConfig() + + +class MatrixGame3WanModel(BaseDiT): + # Marker for action input support (Matrix-Game) + supports_action_input = True + + _fsdp_shard_conditions = _DEFAULT_MATRIXGAME3_CONFIG._fsdp_shard_conditions + _compile_conditions = _DEFAULT_MATRIXGAME3_CONFIG._compile_conditions + _supported_attention_backends = ( + _DEFAULT_MATRIXGAME3_CONFIG._supported_attention_backends + ) + param_names_mapping = _DEFAULT_MATRIXGAME3_CONFIG.param_names_mapping + reverse_param_names_mapping = ( + _DEFAULT_MATRIXGAME3_CONFIG.reverse_param_names_mapping + ) + lora_param_names_mapping = ( + _DEFAULT_MATRIXGAME3_CONFIG.lora_param_names_mapping + ) + + def __init__( + self, + config: MatrixGame3WanVideoConfig, + hf_config: dict[str, Any], + **kwargs, + ) -> None: + super().__init__(config=config, hf_config=hf_config) + + inner_dim = config.num_attention_heads * config.attention_head_dim + self.hidden_size = config.hidden_size + self.num_attention_heads = config.num_attention_heads + self.num_channels_latents = config.out_channels + self.patch_size = config.patch_size + + # 1. Patch & position embedding + self.patch_embedding = PatchEmbed( + in_chans=config.in_channels, + embed_dim=inner_dim, + patch_size=config.patch_size, + flatten=False, + ) + + # 2. Condition embeddings + self.condition_embedder = MatrixGame3TimeImageEmbedding( + dim=inner_dim, + time_freq_dim=config.freq_dim, + ) + self.text_embedding = nn.Sequential( + nn.Linear(config.text_dim, inner_dim), + nn.GELU(approximate="tanh"), + nn.Linear(inner_dim, inner_dim), + ) + self.use_memory = getattr(config, "use_memory", False) + self.sigma_theta = float(getattr(config, "sigma_theta", 0.0)) + self.camera_patch_embedding = None + self.c2ws_hidden_states_layer1 = None + self.c2ws_hidden_states_layer2 = None + if self.use_memory: + camera_in_channels = getattr(config, "camera_embed_in_channels", 1536) + self.camera_patch_embedding = WanCamControlPatchEmbedding( + patch_size=config.patch_size, + in_chans=camera_in_channels, + embed_dim=inner_dim, + ) + self.c2ws_hidden_states_layer1 = nn.Linear(inner_dim, inner_dim) + self.c2ws_hidden_states_layer2 = nn.Linear(inner_dim, inner_dim) + + # 2.1. Get action config + self.action_config = getattr(config, "action_config", {}) + + # 3. Transformer blocks + self.blocks = nn.ModuleList( + [ + MatrixGame3TransformerBlock( + inner_dim, + config.ffn_dim, + config.num_attention_heads, + config.qk_norm, + config.cross_attn_norm, + config.eps, + self._supported_attention_backends, + prefix=f"{getattr(config, 'prefix', 'Wan')}.blocks.{i}", + action_config=self.action_config, + block_id=i, + use_memory=self.use_memory, + ) + for i in range(config.num_layers) + ] + ) + + # 4. Output norm & projection + self.norm_out = LayerNormScaleShift( + inner_dim, + norm_type="layer", + eps=config.eps, + elementwise_affine=False, + dtype=torch.float32, + compute_dtype=torch.float32, + ) + self.proj_out = nn.Linear( + inner_dim, config.out_channels * math.prod(config.patch_size) + ) + self.scale_shift_table = nn.Parameter( + torch.randn(1, 2, inner_dim) / inner_dim**0.5 + ) + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | list[torch.Tensor], + timestep: torch.LongTensor, + # Action inputs + mouse_cond: torch.Tensor | None = None, + keyboard_cond: torch.Tensor | None = None, + x_memory: torch.Tensor | None = None, + timestep_memory: torch.Tensor | None = None, + mouse_cond_memory: torch.Tensor | None = None, + keyboard_cond_memory: torch.Tensor | None = None, + c2ws_plucker_emb: torch.Tensor | None = None, + memory_latent_idx: list[int] | tuple[int, ...] | torch.Tensor | None = None, + predict_latent_idx: tuple[int, int] | list[int] | tuple[int, ...] | torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + if encoder_hidden_states is not None and not isinstance( + encoder_hidden_states, torch.Tensor + ): + encoder_hidden_states = encoder_hidden_states[0] + + memory_length = 0 + if x_memory is not None: + memory_length = x_memory.shape[2] + hidden_states = torch.cat([x_memory.to(hidden_states.dtype), hidden_states], dim=2) + + batch_size, num_channels, num_frames, height, width = hidden_states.shape + p_t, p_h, p_w = self.patch_size + post_patch_num_frames = num_frames // p_t + post_patch_height = height // p_h + post_patch_width = width // p_w + max_frame_index = post_patch_num_frames + if isinstance(predict_latent_idx, tuple) and len(predict_latent_idx) == 2: + max_frame_index = max(max_frame_index, int(predict_latent_idx[1])) + elif predict_latent_idx is not None: + max_frame_index = max(max_frame_index, int(max(predict_latent_idx)) + 1) + if memory_latent_idx is not None and len(memory_latent_idx) > 0: + max_frame_index = max(max_frame_index, int(max(memory_latent_idx)) + 1) + + # Pre-build RoPE frame index tensors once per forward so blocks don't + # do `torch.tensor(list, device=cuda)` per call (host sync). + device_idx = hidden_states.device + if predict_latent_idx is None: + predict_latent_idx = None + elif torch.is_tensor(predict_latent_idx): + predict_latent_idx = predict_latent_idx.to(device=device_idx, dtype=torch.long) + elif isinstance(predict_latent_idx, tuple) and len(predict_latent_idx) == 2: + predict_latent_idx = torch.arange( + predict_latent_idx[0], predict_latent_idx[1], + device=device_idx, dtype=torch.long, + ) + else: + predict_latent_idx = torch.tensor( + list(predict_latent_idx), device=device_idx, dtype=torch.long, + ) + + if memory_latent_idx is None or ( + not torch.is_tensor(memory_latent_idx) and len(memory_latent_idx) == 0 + ): + memory_latent_idx = None + elif torch.is_tensor(memory_latent_idx): + memory_latent_idx = memory_latent_idx.to(device=device_idx, dtype=torch.long) + else: + memory_latent_idx = torch.tensor( + list(memory_latent_idx), device=device_idx, dtype=torch.long, + ) + + head_dim = self.hidden_size // self.num_attention_heads + freqs = _build_rope_freqs( + max_seq_len=(2048 if self.use_memory else 1024), + head_dim=head_dim, + num_heads=self.num_attention_heads, + sigma_theta=(self.sigma_theta if self.use_memory else 0.0), + device=hidden_states.device, + ) + max_seq_len = freqs.shape[1] if freqs.dim() == 3 else freqs.shape[0] + if max_frame_index > max_seq_len or post_patch_height > max_seq_len or post_patch_width > max_seq_len: + raise ValueError( + f"got frames={max_frame_index}, height={post_patch_height}, width={post_patch_width}, " + f"max_seq_len={max_seq_len}" + ) + + hidden_states = self.patch_embedding(hidden_states) + hidden_states = hidden_states.flatten(2).transpose(1, 2) + + plucker_emb = None + if c2ws_plucker_emb is not None and self.camera_patch_embedding is not None: + if memory_length > 0 and c2ws_plucker_emb.shape[2] == post_patch_num_frames - memory_length: + zeros = torch.zeros(c2ws_plucker_emb.shape[0], + c2ws_plucker_emb.shape[1], + memory_length, + c2ws_plucker_emb.shape[3], + c2ws_plucker_emb.shape[4], + device=c2ws_plucker_emb.device, + dtype=c2ws_plucker_emb.dtype) + c2ws_plucker_emb = torch.cat([zeros, c2ws_plucker_emb], dim=2) + target_frames = post_patch_num_frames * p_t + target_height = post_patch_height * p_h + target_width = post_patch_width * p_w + if c2ws_plucker_emb.shape[2] > target_frames or c2ws_plucker_emb.shape[3] > target_height or c2ws_plucker_emb.shape[4] > target_width: + c2ws_plucker_emb = c2ws_plucker_emb[:, :, :target_frames, :target_height, :target_width] + if c2ws_plucker_emb.shape[2:] != (target_frames, target_height, target_width): + raise ValueError( + "Camera conditioning must align with the patchified latent grid, " + f"got camera shape {tuple(c2ws_plucker_emb.shape[2:])} vs target " + f"{(target_frames, target_height, target_width)}" + ) + plucker_emb = self.camera_patch_embedding( + c2ws_plucker_emb.to(hidden_states.dtype) + ) + plucker_hidden = self.c2ws_hidden_states_layer2( + torch.nn.functional.silu(self.c2ws_hidden_states_layer1(plucker_emb)) + ) + plucker_emb = plucker_emb + plucker_hidden + + timestep_tokens = timestep + if timestep_tokens.dim() == 0: + timestep_tokens = timestep_tokens.unsqueeze(0) + if timestep_tokens.dim() == 1: + timestep_tokens = timestep_tokens.unsqueeze(1).repeat( + 1, post_patch_num_frames * post_patch_height * post_patch_width + ) + elif timestep_tokens.dim() == 2 and timestep_tokens.shape[1] == num_frames: + timestep_tokens = ( + timestep_tokens[:, :, None, None] + .repeat(1, 1, post_patch_height, post_patch_width) + .reshape(timestep_tokens.shape[0], -1) + ) + if memory_length > 0: + if timestep_memory is None: + raise ValueError("MatrixGame3 requires timestep_memory when x_memory is provided") + timestep_tokens = torch.cat( + [timestep_memory.to(timestep_tokens.dtype), timestep_tokens], dim=1 + ) + + temb, timestep_proj, encoder_hidden_states = self.condition_embedder( + timestep_tokens.flatten(), + encoder_hidden_states, + timestep_seq_len=timestep_tokens.shape[1], + ) + if timestep_proj.dim() == 3 and timestep_proj.shape[-1] % 6 == 0: + timestep_proj = timestep_proj.unflatten(2, (6, -1)) + + if encoder_hidden_states is not None: + if isinstance(encoder_hidden_states, list): + encoder_hidden_states = encoder_hidden_states[0] + elif encoder_hidden_states.ndim == 2: + encoder_hidden_states = encoder_hidden_states.unsqueeze(0) + if ( + encoder_hidden_states.ndim == 3 + and encoder_hidden_states.shape[-1] == self.text_embedding[0].in_features + ): + encoder_hidden_states = self.text_embedding( + encoder_hidden_states.to(hidden_states.dtype) + ) + + # [F, H, W] for the ActionModule + grid_sizes = (post_patch_num_frames, post_patch_height, post_patch_width) + + # Blocks + for block in self.blocks: + if torch.is_grad_enabled() and self.gradient_checkpointing: + hidden_states = self._gradient_checkpointing_func( + block, + hidden_states, + encoder_hidden_states, + timestep_proj, + freqs, + grid_sizes=grid_sizes, + mouse_cond=mouse_cond, + keyboard_cond=keyboard_cond, + mouse_cond_memory=mouse_cond_memory, + keyboard_cond_memory=keyboard_cond_memory, + plucker_emb=plucker_emb, + memory_length=memory_length, + memory_latent_idx=memory_latent_idx, + predict_latent_idx=predict_latent_idx, + ) + else: + hidden_states = block( + hidden_states, + encoder_hidden_states, + timestep_proj, + freqs, + grid_sizes=grid_sizes, + mouse_cond=mouse_cond, + keyboard_cond=keyboard_cond, + mouse_cond_memory=mouse_cond_memory, + keyboard_cond_memory=keyboard_cond_memory, + plucker_emb=plucker_emb, + memory_length=memory_length, + memory_latent_idx=memory_latent_idx, + predict_latent_idx=predict_latent_idx, + ) + + # Output + if temb.dim() == 3: + shift, scale = ( + self.scale_shift_table.unsqueeze(0) + temb.unsqueeze(2) + ).chunk(2, dim=2) + shift = shift.squeeze(2) + scale = scale.squeeze(2) + else: + shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk( + 2, dim=1 + ) + hidden_states = self.norm_out(hidden_states, shift, scale) + hidden_states = self.proj_out(hidden_states) + + hidden_states = hidden_states.reshape( + batch_size, + post_patch_num_frames, + post_patch_height, + post_patch_width, + p_t, + p_h, + p_w, + -1, + ) + hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6) + output = hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3) + + if memory_length > 0: + output = output[:, :, memory_length:] + + return output diff --git a/fastvideo/models/dits/matrixgame3/utils.py b/fastvideo/models/dits/matrixgame3/utils.py new file mode 100644 index 000000000..f51d018f3 --- /dev/null +++ b/fastvideo/models/dits/matrixgame3/utils.py @@ -0,0 +1,364 @@ +from __future__ import annotations + +import random + +import numpy as np +import torch + +from fastvideo.models.dits.lingbotworld.cam_utils import ( + compute_relative_poses, + get_plucker_embeddings, + interpolate_camera_poses, +) + + +WSAD_OFFSET = 12.35 +DIAGONAL_OFFSET = 8.73 +MOUSE_PITCH_SENSITIVITY = 15.0 +MOUSE_YAW_SENSITIVITY = 15.0 +MOUSE_THRESHOLD = 0.02 + + +def compute_next_pose_from_action(current_pose, keyboard_action, mouse_action): + x, y, z, pitch, yaw = current_pose + w, s, a, d = keyboard_action[:4] + mouse_x, mouse_y = mouse_action[:2] + + delta_pitch = MOUSE_PITCH_SENSITIVITY * mouse_x if abs(mouse_x) >= MOUSE_THRESHOLD else 0.0 + delta_yaw = MOUSE_YAW_SENSITIVITY * mouse_y if abs(mouse_y) >= MOUSE_THRESHOLD else 0.0 + new_pitch = pitch + delta_pitch + new_yaw = yaw + delta_yaw + + new_yaw = (new_yaw + 180) % 360 - 180 + + local_forward = 0.0 + if w > 0.5 and s < 0.5: + local_forward = WSAD_OFFSET + elif s > 0.5 and w < 0.5: + local_forward = -WSAD_OFFSET + + local_right = 0.0 + if d > 0.5 and a < 0.5: + local_right = WSAD_OFFSET + elif a > 0.5 and d < 0.5: + local_right = -WSAD_OFFSET + + if abs(local_forward) > 0.1 and abs(local_right) > 0.1: + local_forward = np.sign(local_forward) * DIAGONAL_OFFSET + local_right = np.sign(local_right) * DIAGONAL_OFFSET + + avg_yaw = float((yaw + new_yaw) / 2.0) + yaw_rad = float(np.deg2rad(avg_yaw)) + cos_yaw = np.cos(yaw_rad) + sin_yaw = np.sin(yaw_rad) + + delta_x = cos_yaw * local_forward - sin_yaw * local_right + delta_y = sin_yaw * local_forward + cos_yaw * local_right + return np.array([x + delta_x, y + delta_y, z, new_pitch, new_yaw], dtype=np.float32) + + +def compute_all_poses_from_actions(keyboard_conditions, mouse_conditions): + poses = np.zeros((len(keyboard_conditions), 5), dtype=np.float32) + for idx in range(len(keyboard_conditions) - 1): + poses[idx + 1] = compute_next_pose_from_action(poses[idx], keyboard_conditions[idx], mouse_conditions[idx]) + return poses + + +def build_intrinsics(height: int, width: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor: + fov_rad = float(np.deg2rad(90.0)) + fx = float(width) / (2.0 * float(np.tan(fov_rad / 2.0))) + fy = float(height) / (2.0 * float(np.tan(fov_rad / 2.0))) + return torch.tensor([[fx, fy, width / 2.0, height / 2.0]], device=device, dtype=dtype) + + +def build_matrixgame3_action_preset(num_frames: int, seed: int | None = None) -> tuple[torch.Tensor, torch.Tensor]: + presets = create_action_presets(num_frames=num_frames, keyboard_dim=6, seed=seed) + return presets["keyboard"], presets["mouse"] + + +def interpolate_camera_poses_handedness( + src_indices: np.ndarray, + src_rot_mat: np.ndarray, + src_trans_vec: np.ndarray, + tgt_indices: np.ndarray, +) -> torch.Tensor: + dets = np.linalg.det(src_rot_mat) + flip_handedness = dets.size > 0 and np.median(dets) < 0.0 + if flip_handedness: + flip_mat = np.diag([1.0, 1.0, -1.0]).astype(src_rot_mat.dtype) + src_rot_mat = src_rot_mat @ flip_mat + + c2ws = interpolate_camera_poses( + src_indices=src_indices, + src_rot_mat=src_rot_mat, + src_trans_vec=src_trans_vec, + tgt_indices=tgt_indices, + ) + if flip_handedness: + flip_mat_t = torch.from_numpy(flip_mat).to(c2ws.device, dtype=c2ws.dtype) + c2ws[:, :3, :3] = c2ws[:, :3, :3] @ flip_mat_t + return c2ws + + +def build_extrinsics_from_actions( + keyboard_conditions: torch.Tensor, + mouse_conditions: torch.Tensor, +) -> torch.Tensor: + keyboard_np = keyboard_conditions.detach().to(dtype=torch.float32).cpu().numpy() + mouse_np = mouse_conditions.detach().to(dtype=torch.float32).cpu().numpy() + poses = compute_all_poses_from_actions(keyboard_np, mouse_np) + rotations = np.concatenate( + [np.zeros((poses.shape[0], 1), dtype=np.float32), poses[:, 3:5]], + axis=1, + ) + positions = poses[:, :3] + return build_extrinsics(rotations, positions) + + +def build_extrinsics( + video_rotation: np.ndarray, + video_position: np.ndarray, +) -> torch.Tensor: + extrinsics = [] + for frame_rotation, frame_position in zip(video_rotation, video_position, strict=False): + roll_deg, pitch_deg, yaw_deg = frame_rotation + roll, pitch, yaw = np.radians([roll_deg, pitch_deg, yaw_deg]) + + rot_z = np.array( + [[np.cos(yaw), -np.sin(yaw), 0], [np.sin(yaw), np.cos(yaw), 0], [0, 0, 1]], + dtype=np.float32, + ) + rot_y = np.array( + [[np.cos(pitch), 0, np.sin(pitch)], [0, 1, 0], [-np.sin(pitch), 0, np.cos(pitch)]], + dtype=np.float32, + ) + rot_x = np.array( + [[1, 0, 0], [0, np.cos(roll), -np.sin(roll)], [0, np.sin(roll), np.cos(roll)]], + dtype=np.float32, + ) + rot = rot_z @ rot_y @ rot_x + ext = np.eye(4, dtype=np.float32) + ext[:3, :3] = rot + ext[:3, 3] = np.asarray(frame_position, dtype=np.float32) + extrinsics.append(ext) + + r_init = np.array([[0, 0, 1], [1, 0, 0], [0, -1, 0]], dtype=np.float32) + extrinsics_tensor = torch.from_numpy(np.stack(extrinsics, axis=0)).float() + extrinsics_tensor[:, :3, :3] = extrinsics_tensor[:, :3, :3] @ torch.from_numpy(r_init) + extrinsics_tensor[:, :3, 3] = extrinsics_tensor[:, :3, 3] * 0.01 + return extrinsics_tensor + + +def build_plucker_from_c2ws( + c2ws_seq: torch.Tensor, + src_indices: np.ndarray, + tgt_indices: np.ndarray, + *, + target_h: int, + target_w: int, + latent_h: int, + latent_w: int, + framewise: bool = True, +) -> torch.Tensor: + c2ws_np = c2ws_seq.detach().to(dtype=torch.float32).cpu().numpy() + c2ws_infer = interpolate_camera_poses_handedness( + src_indices=src_indices, + src_rot_mat=c2ws_np[:, :3, :3], + src_trans_vec=c2ws_np[:, :3, 3], + tgt_indices=tgt_indices, + ).to(device=c2ws_seq.device, dtype=c2ws_seq.dtype) + c2ws_infer = compute_relative_poses(c2ws_infer, framewise=framewise) + return build_plucker_from_pose( + c2ws_infer, + target_h=target_h, + target_w=target_w, + latent_h=latent_h, + latent_w=latent_w, + ) + + +def build_plucker_from_pose( + c2ws_pose: torch.Tensor, + *, + target_h: int, + target_w: int, + latent_h: int, + latent_w: int, +) -> torch.Tensor: + ks = build_intrinsics(target_h, target_w, c2ws_pose.device, c2ws_pose.dtype).repeat(c2ws_pose.shape[0], 1) + plucker = get_plucker_embeddings(c2ws_pose, ks, target_h, target_w) + c1 = target_h // latent_h + c2 = target_w // latent_w + plucker = plucker.view(c2ws_pose.shape[0], latent_h, c1, latent_w, c2, 6) + plucker = plucker.permute(0, 1, 3, 5, 2, 4).reshape(c2ws_pose.shape[0], latent_h, latent_w, 6 * c1 * c2) + plucker = plucker.permute(3, 0, 1, 2).unsqueeze(0) + return plucker + + +def select_memory_idx_fov( + extrinsics_all: torch.Tensor | np.ndarray, + current_start_frame_idx: int, + selected_index_base: list[int], + *, + height: int = 720, + width: int = 1280, + return_confidence: bool = False, +) -> list[int] | tuple[list[int], list[float]]: + if isinstance(extrinsics_all, np.ndarray): + extrinsics_tensor = torch.from_numpy(extrinsics_all).float() + else: + extrinsics_tensor = extrinsics_all.float() + + device = extrinsics_tensor.device + if current_start_frame_idx <= 1: + selected = [0] * len(selected_index_base) + if return_confidence: + return selected, [0.0] * len(selected) + return selected + + fov_rad = np.deg2rad(90.0) + fx = width / (2 * np.tan(fov_rad / 2)) + fy = height / (2 * np.tan(fov_rad / 2)) + near, far = 0.1, 30.0 + + candidate_indices = torch.arange(1, current_start_frame_idx, device=device) + r_cand = extrinsics_tensor[candidate_indices, :3, :3] + t_cand = extrinsics_tensor[candidate_indices, :3, 3:4] + r_cand_inv = r_cand.transpose(1, 2) + t_cand_inv = -torch.bmm(r_cand_inv, t_cand) + + num_side = 10 + z_samples = torch.linspace(near, far, num_side, device=device) + x_samples = torch.linspace(-1, 1, num_side, device=device) + y_samples = torch.linspace(-1, 1, num_side, device=device) + grid_x, grid_y, grid_z = torch.meshgrid(x_samples, y_samples, z_samples, indexing="ij") + points_cam_base = torch.stack([ + grid_x.reshape(-1) * grid_z.reshape(-1) * (width / (2 * fx)), + grid_y.reshape(-1) * grid_z.reshape(-1) * (height / (2 * fy)), + grid_z.reshape(-1), + ], dim=0) + + selected_index: list[int] = [] + selected_confidence: list[float] = [] + for frame_idx in selected_index_base: + base_pose = extrinsics_tensor[frame_idx] + points_world = base_pose[:3, :3] @ points_cam_base + base_pose[:3, 3:4] + points_world_batched = points_world.unsqueeze(0) + points_in_cands = torch.bmm(r_cand_inv, points_world_batched.expand(len(candidate_indices), -1, -1)) + t_cand_inv + + x = points_in_cands[:, 0, :] + y = points_in_cands[:, 1, :] + z = points_in_cands[:, 2, :] + u = (x * fx / torch.clamp(z, min=1e-6)) + width / 2 + v = (y * fy / torch.clamp(z, min=1e-6)) + height / 2 + + in_view = (z > near) & (z < far) & (u >= 0) & (u <= width) & (v >= 0) & (v <= height) + ratios = in_view.float().mean(dim=1) + best_idx = torch.argmax(ratios) + selected_index.append(candidate_indices[best_idx].item()) + selected_confidence.append(ratios[best_idx].item()) + + if return_confidence: + return selected_index, selected_confidence + return selected_index + + +def create_action_presets(num_frames: int, keyboard_dim: int = 4, seed: int = None): + if keyboard_dim not in (2, 4, 6, 7): + raise ValueError(f"keyboard_dim must be 2, 4, 6, or 7, got {keyboard_dim}") + if num_frames % 4 != 1: + raise ValueError("Matrix-Game conditioning expects num_frames to be 4k+1.") + + if seed is not None: + random.seed(seed) + + num_samples_per_action = 4 + + if keyboard_dim == 4: + actions_single_action = ["forward", "left", "right"] + actions_double_action = ["forward_left", "forward_right"] + actions_single_camera = ["camera_l", "camera_r"] + keyboard_idx = {"forward": 0, "back": 1, "left": 2, "right": 3} + elif keyboard_dim == 2: + actions_single_action = ["forward", "back"] + actions_double_action = [] + actions_single_camera = ["camera_l", "camera_r"] + keyboard_idx = {"forward": 0, "back": 1} + elif keyboard_dim == 6: + actions_single_action = ["forward", "back", "left", "right"] + actions_double_action = ["forward_left", "forward_right"] + actions_single_camera = ["camera_l", "camera_r"] + keyboard_idx = {"forward": 0, "back": 1, "left": 2, "right": 3, "t1": 4, "t2": 5} + else: # keyboard_dim == 7 + actions_single_action = ["forward", "back", "left", "right"] + actions_double_action = [] + actions_single_camera = [] + keyboard_idx = {"still": 0, "forward": 1, "back": 2, "left": 3, "right": 4, "a": 5, "d": 6} + + actions_to_test = ( + actions_double_action * 5 + actions_single_camera * 5 + actions_single_action * 5 + ) + for action in (actions_single_action + actions_double_action): + for camera in actions_single_camera: + actions_to_test.append(f"{action}_{camera}") + + if not actions_to_test: + actions_to_test = actions_single_action * 5 + + base_action = actions_single_action + actions_single_camera + cam_value = 0.1 + camera_value_map = { + "camera_up": [cam_value, 0], + "camera_down": [-cam_value, 0], + "camera_l": [0, -cam_value], + "camera_r": [0, cam_value], + "camera_ur": [cam_value, cam_value], + "camera_ul": [cam_value, -cam_value], + "camera_dr": [-cam_value, cam_value], + "camera_dl": [-cam_value, -cam_value], + } + + data = [] + for action_name in actions_to_test: + keyboard_condition = torch.zeros((num_samples_per_action, keyboard_dim)) + mouse_condition = torch.zeros((num_samples_per_action, 2)) + + for sub_act in base_action: + if sub_act not in action_name: + continue + if sub_act in camera_value_map: + mouse_condition = torch.tensor( + [camera_value_map[sub_act] for _ in range(num_samples_per_action)], + dtype=mouse_condition.dtype, + ) + elif sub_act in keyboard_idx: + keyboard_condition[:, keyboard_idx[sub_act]] = 1 + + data.append({ + "keyboard_condition": keyboard_condition, + "mouse_condition": mouse_condition, + }) + + keyboard_condition = torch.zeros((num_frames, keyboard_dim)) + mouse_condition = torch.zeros((num_frames, 2)) + current_frame = 0 + selections = [12] + + while current_frame < num_frames: + rd_frame = selections[random.randint(0, len(selections) - 1)] + entry = data[random.randint(0, len(data) - 1)] + key_seq = entry["keyboard_condition"] + mouse_seq = entry["mouse_condition"] + + if current_frame == 0: + keyboard_condition[:1] = key_seq[:1] + mouse_condition[:1] = mouse_seq[:1] + current_frame = 1 + else: + rd_frame = min(rd_frame, num_frames - current_frame) + repeat_time = rd_frame // 4 + keyboard_condition[current_frame:current_frame + rd_frame] = key_seq.repeat(repeat_time, 1) + mouse_condition[current_frame:current_frame + rd_frame] = mouse_seq.repeat(repeat_time, 1) + current_frame += rd_frame + + return {"keyboard": keyboard_condition, "mouse": mouse_condition} diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index 973a76c7d..e3b087adc 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/fastvideo/models/loader/component_loader.py @@ -82,6 +82,7 @@ class ComponentLoader(ABC): "transformer_2": (TransformerLoader, "diffusers"), "transformer_3": (TransformerLoader, "diffusers"), "vae": (VAELoader, "diffusers"), + "light_vae": (VAELoader, "diffusers"), "audio_vae": (AudioDecoderLoader, "diffusers"), "audio_decoder": (AudioDecoderLoader, "diffusers"), "vocoder": (VocoderLoader, "diffusers"), @@ -791,6 +792,11 @@ class VAELoader(ComponentLoader): # strictly so missing/unexpected keys are surfaced early. strict_load = class_name == "AutoencoderKL" vae.load_state_dict(loaded, strict=strict_load) + if (class_name == "AutoencoderKLWan" + and getattr(vae.config, "use_light_vae", False) + and target_device.type == "cuda" + and hasattr(vae, "optimize_memory_format")): + vae.optimize_memory_format() return vae.eval() diff --git a/fastvideo/models/registry.py b/fastvideo/models/registry.py index 4ada68a49..a525baa51 100644 --- a/fastvideo/models/registry.py +++ b/fastvideo/models/registry.py @@ -53,6 +53,7 @@ _IMAGE_TO_VIDEO_DIT_MODELS = { # Legacy aliases for older HF model_index.json files "MatrixGameWanModel": ("dits", "matrixgame2", "MatrixGame2WanModel"), "CausalMatrixGameWanModel": ("dits", "matrixgame2", "CausalMatrixGame2WanModel"), + "MatrixGame3WanModel": ("dits", "matrixgame3", "MatrixGame3WanModel"), } _TEXT_ENCODER_MODELS = { diff --git a/fastvideo/models/vaes/wanvae.py b/fastvideo/models/vaes/wanvae.py index 77bcf987a..55d28a70a 100644 --- a/fastvideo/models/vaes/wanvae.py +++ b/fastvideo/models/vaes/wanvae.py @@ -33,17 +33,20 @@ is_first_frame = contextvars.ContextVar("is_first_frame", default=False) feat_cache = contextvars.ContextVar("feat_cache", default=None) feat_idx = contextvars.ContextVar("feat_idx", default=0) first_chunk = contextvars.ContextVar("first_chunk", default=None) +use_light_vae = contextvars.ContextVar("use_light_vae", default=False) @contextmanager def forward_context(first_frame_arg=False, feat_cache_arg=None, feat_idx_arg=None, - first_chunk_arg=None): + first_chunk_arg=None, + use_light_vae_arg=False): is_first_frame_token = is_first_frame.set(first_frame_arg) feat_cache_token = feat_cache.set(feat_cache_arg) feat_idx_token = feat_idx.set(feat_idx_arg) first_chunk_token = first_chunk.set(first_chunk_arg) + use_light_vae_token = use_light_vae.set(use_light_vae_arg) try: yield finally: @@ -51,6 +54,7 @@ def forward_context(first_frame_arg=False, feat_cache.reset(feat_cache_token) feat_idx.reset(feat_idx_token) first_chunk.reset(first_chunk_token) + use_light_vae.reset(use_light_vae_token) class AvgDown3D(nn.Module): def __init__( @@ -302,6 +306,8 @@ class WanResample(nn.Module): def forward(self, x): b, c, t, h, w = x.size() first_frame = is_first_frame.get() + _first_chunk = first_chunk.get() + _use_light_vae = use_light_vae.get() if first_frame: assert t == 1 _feat_cache = feat_cache.get() @@ -309,34 +315,60 @@ class WanResample(nn.Module): if self.mode == "upsample3d": if _feat_cache is not None: idx = _feat_idx - if _feat_cache[idx] is None: - _feat_cache[idx] = "Rep" - _feat_idx += 1 - else: + if _use_light_vae: + if _feat_cache[idx] is None: + _feat_cache[idx] = "Rep" cache_x = x[:, :, -CACHE_T:, :, :].clone() - if cache_x.shape[2] < 2 and _feat_cache[ - idx] is not None and _feat_cache[idx] != "Rep": - # cache last frame of last two chunk - cache_x = torch.cat([ - _feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( - cache_x.device), cache_x - ], - dim=2) - if cache_x.shape[2] < 2 and _feat_cache[ - idx] is not None and _feat_cache[idx] == "Rep": - cache_x = torch.cat([ - torch.zeros_like(cache_x).to(cache_x.device), - cache_x - ], - dim=2) - x = self.time_conv(x) if _feat_cache[idx] == "Rep" else self.time_conv(x, _feat_cache[idx]) + # "Rep" sentinel on first chunk, cache tensor after; isinstance avoids `tensor == str` version surprises. + if isinstance(_feat_cache[idx], str): + x = self.time_conv(x) + else: + if cache_x.shape[2] < 2: + cache_x = torch.cat([ + _feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + x.device + ), + cache_x, + ], dim=2) + x = self.time_conv(x, _feat_cache[idx]) + _feat_cache[idx] = cache_x _feat_idx += 1 x = x.reshape(b, 2, c, t, h, w) - x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), - 3) + x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), 3) x = x.reshape(b, c, t * 2, h, w) + + if _first_chunk: + x = x[:, :, 1:, :, :] + else: + if _feat_cache[idx] is None: + _feat_cache[idx] = "Rep" + _feat_idx += 1 + else: + cache_x = x[:, :, -CACHE_T:, :, :].clone() + if cache_x.shape[2] < 2 and _feat_cache[ + idx] is not None and _feat_cache[idx] != "Rep": + cache_x = torch.cat([ + _feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to( + cache_x.device), cache_x + ], + dim=2) + if cache_x.shape[2] < 2 and _feat_cache[ + idx] is not None and _feat_cache[idx] == "Rep": + cache_x = torch.cat([ + torch.zeros_like(cache_x).to(cache_x.device), + cache_x + ], + dim=2) + x = self.time_conv(x) if _feat_cache[idx] == "Rep" else self.time_conv(x, _feat_cache[idx]) + _feat_cache[idx] = cache_x + _feat_idx += 1 + + x = x.reshape(b, 2, c, t, h, w) + x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]), + 3) + x = x.reshape(b, c, t * 2, h, w) feat_cache.set(_feat_cache) feat_idx.set(_feat_idx) elif not first_frame and hasattr(self, "time_conv"): @@ -1129,13 +1161,28 @@ class AutoencoderKLWan(nn.Module, ParallelTiledVAE): self._enc_conv_idx = 0 self._enc_feat_map = [None] * self._enc_conv_num + def optimize_memory_format(self) -> None: + """Optional CUDA-only weight layout optimization used by MG3 LightVAE.""" + + def _convert(module: nn.Module) -> None: + for child in module.children(): + if isinstance(child, nn.Conv3d): + child.weight.data = child.weight.data.to( + memory_format=torch.channels_last_3d + ) + else: + _convert(child) + + _convert(self) + def encode(self, x: torch.Tensor) -> torch.Tensor: if self.use_feature_cache: self.clear_cache() if self.config.patch_size is not None: x = patchify(x, patch_size=self.config.patch_size) with forward_context(feat_cache_arg=self._enc_feat_map, - feat_idx_arg=self._enc_conv_idx): + feat_idx_arg=self._enc_conv_idx, + use_light_vae_arg=self.config.use_light_vae): t = x.shape[2] iter_ = 1 + (t - 1) // 4 for i in range(iter_): @@ -1164,7 +1211,8 @@ class AutoencoderKLWan(nn.Module, ParallelTiledVAE): return enc def _encode(self, x: torch.Tensor, first_frame=False) -> torch.Tensor: - with forward_context(first_frame_arg=first_frame): + with forward_context(first_frame_arg=first_frame, + use_light_vae_arg=self.config.use_light_vae): out = self.encoder(x) enc = self.quant_conv(out) mu, logvar = enc[:, :self.z_dim, :, :, :], enc[:, self.z_dim:, :, :, :] @@ -1195,7 +1243,8 @@ class AutoencoderKLWan(nn.Module, ParallelTiledVAE): iter_ = z.shape[2] x = self.post_quant_conv(z) with forward_context(feat_cache_arg=self._feat_map, - feat_idx_arg=self._conv_idx): + feat_idx_arg=self._conv_idx, + use_light_vae_arg=self.config.use_light_vae): for i in range(iter_): feat_idx.set(0) if i == 0: @@ -1219,7 +1268,8 @@ class AutoencoderKLWan(nn.Module, ParallelTiledVAE): def _decode(self, z: torch.Tensor, first_frame=False) -> torch.Tensor: x = self.post_quant_conv(z) - with forward_context(first_frame_arg=first_frame): + with forward_context(first_frame_arg=first_frame, + use_light_vae_arg=self.config.use_light_vae): out = self.decoder(x) out = torch.clamp(out, min=-1.0, max=1.0) @@ -1275,7 +1325,8 @@ class AutoencoderKLWan(nn.Module, ParallelTiledVAE): iter_ = z.shape[2] x = self.post_quant_conv(z) - with forward_context(feat_cache_arg=cache, feat_idx_arg=0): + with forward_context(feat_cache_arg=cache, feat_idx_arg=0, + use_light_vae_arg=self.config.use_light_vae): outputs = [] for i in range(iter_): feat_idx.set(0) diff --git a/fastvideo/models/vision_utils.py b/fastvideo/models/vision_utils.py index 74fcffb79..82ee2cfec 100644 --- a/fastvideo/models/vision_utils.py +++ b/fastvideo/models/vision_utils.py @@ -361,6 +361,23 @@ def resize( if resize_mode == "default": image = image.resize((width, height), resample=PIL_INTERPOLATION[resample]) + elif resize_mode == "crop": + src_w, src_h = image.size + target_ratio = height / width + src_ratio = src_h / src_w + + if src_ratio > target_ratio: + crop_h = int(src_w * target_ratio) + crop_w = src_w + else: + crop_h = src_h + crop_w = int(src_h / target_ratio) + + top = max(0, (src_h - crop_h) // 2) + left = max(0, (src_w - crop_w) // 2) + image = image.crop((left, top, left + crop_w, top + crop_h)) + image = image.resize((width, height), + resample=PIL_INTERPOLATION[resample]) else: raise ValueError(f"resize_mode {resize_mode} is not supported") return image @@ -410,4 +427,3 @@ def preprocess_reference_image_for_clip(image: PIL.Image.Image, device: torch.de denormalized_tensor = resized_tensor.mul_(0.5).add_(0.5) return TF.to_pil_image(denormalized_tensor) - diff --git a/fastvideo/pipelines/basic/matrixgame3/__init__.py b/fastvideo/pipelines/basic/matrixgame3/__init__.py new file mode 100644 index 000000000..89b6b8f94 --- /dev/null +++ b/fastvideo/pipelines/basic/matrixgame3/__init__.py @@ -0,0 +1,4 @@ +# SPDX-License-Identifier: Apache-2.0 +from fastvideo.pipelines.basic.matrixgame3.matrixgame3_i2v_pipeline import (MatrixGame3I2VPipeline) + +__all__ = ["MatrixGame3I2VPipeline"] diff --git a/fastvideo/pipelines/basic/matrixgame3/matrixgame3_i2v_pipeline.py b/fastvideo/pipelines/basic/matrixgame3/matrixgame3_i2v_pipeline.py new file mode 100644 index 000000000..8bf03604c --- /dev/null +++ b/fastvideo/pipelines/basic/matrixgame3/matrixgame3_i2v_pipeline.py @@ -0,0 +1,52 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Matrix-Game 3.0 I2V pipeline implementation.""" + +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.models.schedulers.scheduling_flow_unipc_multistep import FlowUniPCMultistepScheduler +from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase +from fastvideo.pipelines.lora_pipeline import LoRAPipeline +from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage, InputValidationStage, LatentPreparationStage, + MatrixGame3DenoisingStage, TextEncodingStage) +from fastvideo.pipelines.stages.image_encoding import MatrixGame3ImageVAEEncodingStage + +logger = init_logger(__name__) + + +class MatrixGame3I2VPipeline(LoRAPipeline, ComposedPipelineBase): + _required_config_modules = ["vae", "transformer", "scheduler", "text_encoder", "tokenizer"] + _extra_config_module_map = {"vae": "light_vae"} + + def initialize_pipeline(self, fastvideo_args: FastVideoArgs): + self.modules["scheduler"] = FlowUniPCMultistepScheduler(shift=fastvideo_args.pipeline_config.flow_shift) + + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs): + self.add_stage(stage_name="input_validation_stage", stage=InputValidationStage()) + + self.add_stage(stage_name="prompt_encoding_stage", + stage=TextEncodingStage( + text_encoders=[self.get_module("text_encoder")], + tokenizers=[self.get_module("tokenizer")], + )) + + self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage()) + + self.add_stage(stage_name="latent_preparation_stage", + stage=LatentPreparationStage(scheduler=self.get_module("scheduler"), + transformer=self.get_module("transformer"))) + + self.add_stage(stage_name="image_latent_preparation_stage", + stage=MatrixGame3ImageVAEEncodingStage(vae=self.get_module("vae"))) + + self.add_stage(stage_name="denoising_stage", + stage=MatrixGame3DenoisingStage(transformer=self.get_module("transformer"), + scheduler=self.get_module("scheduler"), + pipeline=self, + vae=self.get_module("vae"))) + + self.add_stage(stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"))) + + logger.info("MatrixGame3I2VPipeline initialized with text, action, and camera support") + + +EntryClass = [MatrixGame3I2VPipeline] diff --git a/fastvideo/pipelines/basic/matrixgame3/presets.py b/fastvideo/pipelines/basic/matrixgame3/presets.py new file mode 100644 index 000000000..059d3c2df --- /dev/null +++ b/fastvideo/pipelines/basic/matrixgame3/presets.py @@ -0,0 +1,33 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Matrix-Game 3.0 pipeline presets.""" +from fastvideo.api.presets import InferencePreset, PresetStageSpec + +_DENOISE_STAGE = PresetStageSpec( + name="denoise", + kind="denoising", + description="Iterative denoising pass with action and camera conditioning", + allowed_overrides=frozenset({ + "num_inference_steps", + "guidance_scale", + }), +) + +MATRIXGAME3_I2V = InferencePreset( + name="matrixgame3_i2v", + version=1, + model_family="matrixgame", + description="Matrix-Game 3.0 I2V (5B, 720p)", + workload_type="i2v", + stage_schemas=(_DENOISE_STAGE, ), + defaults={ + "height": 720, + "width": 1280, + "num_frames": 57, + "fps": 25, + "guidance_scale": 1.0, + "num_inference_steps": 3, + "negative_prompt": "", + }, +) + +ALL_PRESETS = (MATRIXGAME3_I2V, ) diff --git a/fastvideo/pipelines/pipeline_batch_info.py b/fastvideo/pipelines/pipeline_batch_info.py index 7b16ac82f..7af701ea1 100644 --- a/fastvideo/pipelines/pipeline_batch_info.py +++ b/fastvideo/pipelines/pipeline_batch_info.py @@ -130,6 +130,8 @@ class ForwardBatch: mouse_cond: torch.Tensor | None = None # Shape: (B, T, 2) keyboard_cond: torch.Tensor | None = None # Shape: (B, T, K) grid_sizes: torch.Tensor | None = None # Shape: (3,) [F,H,W] + num_iterations: int | None = None + use_base_model: bool = False # Camera control inputs (HYWorld) pose: str | None = None # Camera trajectory: pose string (e.g., 'w-31') or JSON file path diff --git a/fastvideo/pipelines/stages/__init__.py b/fastvideo/pipelines/stages/__init__.py index 94e419c83..acc4c5f67 100644 --- a/fastvideo/pipelines/stages/__init__.py +++ b/fastvideo/pipelines/stages/__init__.py @@ -16,7 +16,8 @@ from fastvideo.pipelines.stages.denoising import (Cosmos25AutoDenoisingStage, Co from fastvideo.pipelines.stages.sr_denoising import SRDenoisingStage from fastvideo.pipelines.stages.encoding import EncodingStage from fastvideo.pipelines.stages.image_encoding import (ImageEncodingStage, MatrixGame2ImageEncodingStage, - MatrixGame2ImageVAEEncodingStage, RefImageEncodingStage, + MatrixGame2ImageVAEEncodingStage, + MatrixGame3ImageVAEEncodingStage, RefImageEncodingStage, ImageVAEEncodingStage, VideoVAEEncodingStage, Hy15ImageEncodingStage, HYWorldImageEncodingStage) from fastvideo.pipelines.stages.gamecraft_image_encoding import (GameCraftImageVAEEncodingStage) @@ -30,6 +31,7 @@ from fastvideo.pipelines.basic.ltx2.stages import ( # noqa: F401 LTX2TextEncodingStage, LTX2UpsampleStage, STAGE_2_DISTILLED_SIGMA_VALUES, ) from fastvideo.pipelines.stages.matrixgame2_denoising import MatrixGame2CausalDenoisingStage +from fastvideo.pipelines.stages.matrixgame3_denoising import MatrixGame3DenoisingStage from fastvideo.pipelines.stages.hyworld_denoising import HYWorldDenoisingStage from fastvideo.pipelines.stages.gamecraft_denoising import GameCraftDenoisingStage from fastvideo.pipelines.stages.gen3c_stages import (Gen3CCFGPolicyStage, Gen3CConditioningStage, Gen3CDenoisingStage, @@ -61,6 +63,7 @@ __all__ = [ "CausalDMDDenosingStage", "CausalDenoisingStage", "MatrixGame2CausalDenoisingStage", + "MatrixGame3DenoisingStage", "HYWorldDenoisingStage", "GameCraftDenoisingStage", "Gen3CCFGPolicyStage", @@ -80,6 +83,7 @@ __all__ = [ "ImageEncodingStage", "MatrixGame2ImageEncodingStage", "MatrixGame2ImageVAEEncodingStage", + "MatrixGame3ImageVAEEncodingStage", "Hy15ImageEncodingStage", "HYWorldImageEncodingStage", "RefImageEncodingStage", diff --git a/fastvideo/pipelines/stages/image_encoding.py b/fastvideo/pipelines/stages/image_encoding.py index ecedd416f..7aa0109a4 100644 --- a/fastvideo/pipelines/stages/image_encoding.py +++ b/fastvideo/pipelines/stages/image_encoding.py @@ -863,3 +863,100 @@ class MatrixGame2ImageVAEEncodingStage(ImageVAEEncodingStage): self.vae.to("cpu") return batch + + +class MatrixGame3ImageVAEEncodingStage(ImageVAEEncodingStage): + + def preprocess( + self, + image: torch.Tensor | PIL.Image.Image, + vae_scale_factor: int, + height: int | None = None, + width: int | None = None, + resize_mode: str = "crop", + ) -> torch.Tensor: + return super().preprocess( + image, + vae_scale_factor=vae_scale_factor, + height=height, + width=width, + resize_mode=resize_mode, + ) + + def forward( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> ForwardBatch: + assert batch.pil_image is not None + + if fastvideo_args.mode == ExecutionMode.INFERENCE: + assert batch.height is not None and isinstance(batch.height, int) + assert batch.width is not None and isinstance(batch.width, int) + height = batch.height + width = batch.width + elif fastvideo_args.mode == ExecutionMode.PREPROCESS: + assert batch.height is not None and isinstance(batch.height, list) + assert batch.width is not None and isinstance(batch.width, list) + height = batch.height[0] + width = batch.width[0] + else: + height = batch.height if isinstance(batch.height, int) else batch.height[0] + width = batch.width if isinstance(batch.width, int) else batch.width[0] + + self.vae = self.vae.to(get_local_torch_device()) + + image = batch.pil_image + if isinstance(image, torch.Tensor): + if image.dim() == 5: + image = image[:, :, :1] + elif image.dim() == 4: + image = image.unsqueeze(2) + else: + raise ValueError(f"Unexpected tensor dimensions for MatrixGame3 image: {image.shape}") + video_condition = image.to(get_local_torch_device(), dtype=torch.float32) + else: + image = self.preprocess(image, + vae_scale_factor=self.vae.spatial_compression_ratio, + height=height, + width=width).to(get_local_torch_device(), dtype=torch.float32) + video_condition = image.unsqueeze(2).to(get_local_torch_device(), dtype=torch.float32) + + vae_dtype = PRECISION_TO_TYPE[fastvideo_args.pipeline_config.vae_precision] + vae_autocast_enabled = (vae_dtype != torch.float32) and not fastvideo_args.disable_autocast + + with torch.autocast(device_type="cuda", dtype=vae_dtype, enabled=vae_autocast_enabled): + if fastvideo_args.pipeline_config.vae_tiling: + self.vae.enable_tiling() + if not vae_autocast_enabled: + video_condition = video_condition.to(vae_dtype) + encoder_output = self.vae.encode(video_condition) + + img_cond = encoder_output.mode() + + if (hasattr(self.vae.config, 'latents_mean') and hasattr(self.vae.config, 'latents_std')): + latents_mean = torch.tensor(self.vae.config.latents_mean, device=img_cond.device, + dtype=img_cond.dtype).view(1, -1, 1, 1, 1) + latents_std = torch.tensor(self.vae.config.latents_std, device=img_cond.device, + dtype=img_cond.dtype).view(1, -1, 1, 1, 1) + img_cond = (img_cond - latents_mean) / latents_std + elif (hasattr(self.vae, "shift_factor") and self.vae.shift_factor is not None): + if isinstance(self.vae.shift_factor, torch.Tensor): + img_cond -= self.vae.shift_factor.to(img_cond.device, img_cond.dtype) + else: + img_cond -= self.vae.shift_factor + + if hasattr(self.vae, 'scaling_factor'): + if isinstance(self.vae.scaling_factor, torch.Tensor): + img_cond = img_cond * self.vae.scaling_factor.to(img_cond.device, img_cond.dtype) + else: + img_cond = img_cond * self.vae.scaling_factor + + batch.image_latent = img_cond + + if hasattr(self, 'maybe_free_model_hooks'): + self.maybe_free_model_hooks() + + if fastvideo_args.vae_cpu_offload: + self.vae.to("cpu") + return batch diff --git a/fastvideo/pipelines/stages/matrixgame3_denoising.py b/fastvideo/pipelines/stages/matrixgame3_denoising.py new file mode 100644 index 000000000..088fd2eb6 --- /dev/null +++ b/fastvideo/pipelines/stages/matrixgame3_denoising.py @@ -0,0 +1,328 @@ +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations + +import numpy as np +import torch # type: ignore + +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.forward_context import set_forward_context +from fastvideo.logger import init_logger +from fastvideo.models.dits.lingbotworld.cam_utils import compute_relative_poses +from fastvideo.models.dits.matrixgame3.utils import ( + build_extrinsics_from_actions, + build_matrixgame3_action_preset, + build_plucker_from_c2ws, + build_plucker_from_pose, + interpolate_camera_poses_handedness, + select_memory_idx_fov, +) +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.pipelines.stages.denoising import DenoisingStage +from fastvideo.pipelines.stages.validators import StageValidators as V +from fastvideo.pipelines.stages.validators import VerificationResult + +logger = init_logger(__name__) + + +class MatrixGame3DenoisingStage(DenoisingStage): + + def _infer_num_iterations(self, batch: ForwardBatch) -> int: + if batch.num_iterations is not None: + return batch.num_iterations + if isinstance(batch.num_frames, int) and batch.num_frames > 57: + return 1 + max(0, (batch.num_frames - 57 + 39) // 40) + return 1 + + def verify_input(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: + result = VerificationResult() + result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)]) + result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty) + result.add_check("image_embeds", batch.image_embeds, V.is_list) + result.add_check("image_latent", batch.image_latent, V.none_or_tensor_with_dims(5)) + result.add_check("num_inference_steps", batch.num_inference_steps, V.positive_int) + result.add_check("guidance_scale", batch.guidance_scale, V.positive_float) + result.add_check("eta", batch.eta, V.non_negative_float) + result.add_check("generator", batch.generator, V.generator_or_list_generators) + result.add_check("do_classifier_free_guidance", batch.do_classifier_free_guidance, V.bool_value) + result.add_check("negative_prompt_embeds", batch.negative_prompt_embeds, + lambda x: not batch.do_classifier_free_guidance or V.list_not_empty(x)) + return result + + def forward( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> ForwardBatch: + assert batch.latents is not None, "latents must be prepared before MatrixGame3 denoising" + assert batch.image_latent is not None, "MatrixGame3 requires first-frame VAE latents" + assert batch.prompt_embeds, "MatrixGame3 requires text embeddings" + + target_dtype = torch.bfloat16 + autocast_enabled = (target_dtype != torch.float32) and not fastvideo_args.disable_autocast + device = batch.latents.device + + extra_step_kwargs = self.prepare_extra_func_kwargs( + self.scheduler.step, + { + "generator": batch.generator, + "eta": batch.eta, + }, + ) + + prompt_embeds = batch.prompt_embeds + negative_prompt_embeds = batch.negative_prompt_embeds or [] + use_base_model = batch.use_base_model + + latents = batch.latents + img_cond = batch.image_latent.to(device=device, dtype=target_dtype) + latent_h = latents.shape[-2] + latent_w = latents.shape[-1] + patch_h = int(getattr(self.transformer, "patch_size", (1, 2, 2))[1]) + patch_w = int(getattr(self.transformer, "patch_size", (1, 2, 2))[2]) + latent_h_aligned = (latent_h // patch_h) * patch_h + latent_w_aligned = (latent_w // patch_w) * patch_w + if latent_h_aligned != latent_h or latent_w_aligned != latent_w: + logger.warning( + "Cropping MatrixGame3 latents to patch-aligned size: (%d, %d) -> (%d, %d), patch=(%d,%d)", + latent_h, + latent_w, + latent_h_aligned, + latent_w_aligned, + patch_h, + patch_w, + ) + latents = latents[:, :, :, :latent_h_aligned, :latent_w_aligned] + img_cond = img_cond[:, :, :, :latent_h_aligned, :latent_w_aligned] + latent_h = latent_h_aligned + latent_w = latent_w_aligned + spatial_ratio = fastvideo_args.pipeline_config.vae_config.arch_config.spatial_compression_ratio + target_h = latent_h * spatial_ratio + target_w = latent_w * spatial_ratio + num_iterations = self._infer_num_iterations(batch) + clip_frame = 56 # hardcode for now + first_clip_frame = clip_frame + 1 + past_frame = 16 + + def align_frame_to_block(frame_idx: int) -> int: + return (frame_idx - 1) // 4 * 4 + 1 if frame_idx > 0 else 1 + + def get_latent_idx(frame_idx: int) -> int: + return (frame_idx - 1) // 4 + 1 + + total_video_frames = first_clip_frame + max(0, num_iterations - 1) * (clip_frame - past_frame) + if batch.keyboard_cond is None or batch.mouse_cond is None: + keyboard_cond, mouse_cond = build_matrixgame3_action_preset(total_video_frames, seed=batch.seed) + batch.keyboard_cond = keyboard_cond.unsqueeze(0).to(device=device, dtype=target_dtype) + batch.mouse_cond = mouse_cond.unsqueeze(0).to(device=device, dtype=target_dtype) + else: + batch.keyboard_cond = batch.keyboard_cond.to(device=device, dtype=target_dtype) + batch.mouse_cond = batch.mouse_cond.to(device=device, dtype=target_dtype) + + # MG3 currently assumes batch_size=1 for camera-trajectory derivation. If a + # future caller batches multiple action streams, build_extrinsics_from_actions + # must be made per-batch-item; until then, fail loudly rather than silently + # using item-0's trajectory for the whole batch. + assert batch.keyboard_cond.shape[0] == 1, ( + "MatrixGame3 denoising stage requires batch_size=1 for action conditioning; " + f"got batch.keyboard_cond.shape[0]={batch.keyboard_cond.shape[0]}. " + "If batched action streams are needed, update build_extrinsics_from_actions " + "to handle per-batch-item trajectories.") + extrinsics_all = build_extrinsics_from_actions(batch.keyboard_cond[0], batch.mouse_cond[0]).to(device) + all_latents: list[torch.Tensor] = [] + + for clip_idx in range(num_iterations): + try: + self.scheduler.set_timesteps( + batch.num_inference_steps, + device=device, + shift=fastvideo_args.pipeline_config.flow_shift, + ) + except TypeError: + self.scheduler.set_timesteps(batch.num_inference_steps, device=device) + timesteps = self.scheduler.timesteps + + first_clip = clip_idx == 0 + current_end_frame_idx = first_clip_frame if first_clip else first_clip_frame + clip_idx * (clip_frame - + past_frame) + current_start_frame_idx = 0 if first_clip else current_end_frame_idx - clip_frame + # 15 for first clip, 14 for later clips + current_latent_frames = (first_clip_frame - 1) // 4 + 1 if first_clip else clip_frame // 4 + cond_frames = 1 if first_clip else 4 + latent_start_idx = get_latent_idx(current_start_frame_idx) + latent_end_idx = get_latent_idx(current_end_frame_idx) + + clip_keyboard = batch.keyboard_cond[:, current_start_frame_idx:current_end_frame_idx] + clip_mouse = batch.mouse_cond[:, current_start_frame_idx:current_end_frame_idx] + + cond_frames = min(cond_frames, img_cond.shape[2]) + clip_generator = batch.generator[0] if isinstance(batch.generator, list) else batch.generator + current_latents = torch.randn( + (latents.shape[0], latents.shape[1], latent_end_idx - latent_start_idx, latent_h, latent_w), + generator=clip_generator if isinstance(clip_generator, torch.Generator) else None, + dtype=target_dtype).to(device) + current_latents[:, :, :cond_frames] = img_cond[:, :, :cond_frames] + + c2ws_chunk = extrinsics_all[current_start_frame_idx:current_end_frame_idx] + src_indices = np.linspace(current_start_frame_idx, current_end_frame_idx - 1, + current_end_frame_idx - current_start_frame_idx) + tgt_indices = np.linspace(0 if first_clip else current_start_frame_idx + 3, current_end_frame_idx - 1, + current_latent_frames) + + if batch.c2ws_plucker_emb is not None: + plucker_no_mem = batch.c2ws_plucker_emb[:, :, current_start_frame_idx:current_end_frame_idx] + plucker_no_mem = plucker_no_mem.to(device=device, dtype=target_dtype) + else: + plucker_no_mem = build_plucker_from_c2ws( + c2ws_chunk, + src_indices, + tgt_indices, + target_h=target_h, + target_w=target_w, + latent_h=latent_h, + latent_w=latent_w, + framewise=True, + ).to(device=device, dtype=target_dtype) + + x_memory = None + timestep_memory = None + mouse_cond_memory = None + keyboard_cond_memory = None + memory_latent_idx = None + c2ws_plucker_emb = plucker_no_mem + if all_latents: + # t-1, t-9, t-17, t-25, t-33 + selected_index_base = [current_end_frame_idx - offset for offset in range(1, 34, 8)] + selected_index = select_memory_idx_fov( + extrinsics_all, + current_start_frame_idx, + selected_index_base, + height=target_h, + width=target_w, + ) + if len(selected_index) > 0: + selected_index[-1] = 4 + + memory_pluckers: list[torch.Tensor] = [] + memory_latent_idx = [] + for mem_idx, reference_idx in zip(selected_index, selected_index_base, strict=False): + memory_latent_idx.append(get_latent_idx(mem_idx)) + + mem_idx_aligned = align_frame_to_block(mem_idx) + mem_block = extrinsics_all[mem_idx_aligned:mem_idx_aligned + 4] + mem_src = np.linspace(mem_idx_aligned, mem_idx_aligned + mem_block.shape[0] - 1, mem_block.shape[0]) + mem_tgt = np.array([mem_idx_aligned + 3], dtype=np.float32) + mem_pose = interpolate_camera_poses_handedness( + src_indices=mem_src, + src_rot_mat=mem_block[:, :3, :3].detach().cpu().numpy(), + src_trans_vec=mem_block[:, :3, 3].detach().cpu().numpy(), + tgt_indices=mem_tgt, + ).to(device=device, dtype=target_dtype) + reference_pose = extrinsics_all[reference_idx:reference_idx + 1].to(device=device, + dtype=target_dtype) + rel_pair = torch.cat([reference_pose, mem_pose], dim=0) + rel_pose = compute_relative_poses(rel_pair, framewise=False)[1:2] + memory_pluckers.append( + build_plucker_from_pose( + rel_pose, + target_h=target_h, + target_w=target_w, + latent_h=latent_h, + latent_w=latent_w, + ).to(device=device, dtype=target_dtype)) + + if memory_pluckers: + c2ws_plucker_emb = torch.cat(memory_pluckers + [plucker_no_mem], dim=2) + history = torch.cat(all_latents, dim=2) + x_memory = history[:, :, memory_latent_idx] + mouse_cond_memory = torch.ones( + (clip_keyboard.shape[0], len(memory_latent_idx), clip_mouse.shape[-1]), + device=device, + dtype=target_dtype) + keyboard_cond_memory = -torch.ones( + (clip_keyboard.shape[0], len(memory_latent_idx), clip_keyboard.shape[-1]), + device=device, + dtype=target_dtype) + timestep_memory = x_memory.new_zeros( + (x_memory.shape[0], x_memory.shape[2] * x_memory.shape[3] * x_memory.shape[4] // 4)) + with self.progress_bar(total=len(timesteps)) as progress_bar: + for timestep in timesteps: + latent_model_input = current_latents + if hasattr(self.scheduler, "scale_model_input"): + latent_model_input = self.scheduler.scale_model_input(latent_model_input, timestep) + + timestep_tokens = latent_model_input.new_full( + ( + latent_model_input.shape[2], + latent_model_input.shape[3] * latent_model_input.shape[4] // 4, + ), + timestep, + ) + timestep_tokens[:cond_frames].zero_() + timestep_tokens = timestep_tokens.flatten().unsqueeze(0) + with torch.autocast(device_type="cuda", dtype=target_dtype, enabled=autocast_enabled), \ + set_forward_context(current_timestep=timestep, + attn_metadata=None, + forward_batch=batch): + noise_pred = self.transformer( + latent_model_input, + prompt_embeds, + timestep_tokens, + mouse_cond=clip_mouse, + keyboard_cond=clip_keyboard, + x_memory=x_memory, + timestep_memory=timestep_memory, + mouse_cond_memory=mouse_cond_memory, + keyboard_cond_memory=keyboard_cond_memory, + c2ws_plucker_emb=c2ws_plucker_emb, + memory_latent_idx=memory_latent_idx, + predict_latent_idx=(latent_start_idx, latent_end_idx), + ) + + if use_base_model and batch.do_classifier_free_guidance and negative_prompt_embeds: + null_mouse_cond = torch.ones_like(clip_mouse) + null_keyboard_cond = -torch.ones_like(clip_keyboard) + noise_pred_uncond = self.transformer( + latent_model_input, + negative_prompt_embeds, + timestep_tokens, + mouse_cond=null_mouse_cond, + keyboard_cond=null_keyboard_cond, + x_memory=None, + timestep_memory=None, + mouse_cond_memory=None, + keyboard_cond_memory=None, + c2ws_plucker_emb=plucker_no_mem, + memory_latent_idx=None, + predict_latent_idx=(latent_start_idx, latent_end_idx), + ) + noise_pred = noise_pred_uncond + batch.guidance_scale * (noise_pred - noise_pred_uncond) + + if noise_pred.shape != current_latents.shape: + aligned_t = min(noise_pred.shape[2], current_latents.shape[2]) + aligned_h = min(noise_pred.shape[3], current_latents.shape[3]) + aligned_w = min(noise_pred.shape[4], current_latents.shape[4]) + logger.warning( + "Aligning noise/sample shapes before scheduler.step: noise=%s sample=%s -> (*,*,%d,%d,%d)", + tuple(noise_pred.shape), + tuple(current_latents.shape), + aligned_t, + aligned_h, + aligned_w, + ) + noise_pred = noise_pred[:, :, :aligned_t, :aligned_h, :aligned_w] + current_latents = current_latents[:, :, :aligned_t, :aligned_h, :aligned_w] + + current_latents = self.scheduler.step(noise_pred, + timestep, + current_latents, + **extra_step_kwargs, + return_dict=False)[0] + current_latents[:, :, :cond_frames] = img_cond[:, :, :cond_frames] + progress_bar.update() + + img_cond = current_latents[:, :, -4:].detach() + denoised_pred = current_latents if first_clip else current_latents[:, :, -10:] + all_latents.append(denoised_pred.detach()) + + batch.latents = torch.cat(all_latents, dim=2) + return batch diff --git a/fastvideo/registry.py b/fastvideo/registry.py index 81fbde072..af516c113 100644 --- a/fastvideo/registry.py +++ b/fastvideo/registry.py @@ -31,6 +31,7 @@ from fastvideo.configs.pipelines.lingbotworld import LingBotWorldI2V480PConfig from fastvideo.configs.pipelines.longcat import LongCatT2V480PConfig from fastvideo.pipelines.basic.ltx2.pipeline_configs import LTX2T2VConfig from fastvideo.configs.pipelines.matrixgame2 import MatrixGame2I2V480PConfig +from fastvideo.configs.pipelines.matrixgame3 import MatrixGame3I2V720PConfig from fastvideo.configs.pipelines.turbodiffusion import ( TurboDiffusionI2V_A14B_Config, TurboDiffusionT2V_14B_Config, @@ -54,6 +55,7 @@ from fastvideo.configs.pipelines.sd35 import SD35Config from fastvideo.configs.pipelines.stable_audio import (StableAudioOpenSmallConfig, StableAudioT2AConfig) from fastvideo.api.sampling_param import SamplingParam from fastvideo.api.matrixgame2 import MatrixGame2SamplingParam +from fastvideo.api.matrixgame3 import MatrixGame3SamplingParam from fastvideo.fastvideo_args import WorkloadType from fastvideo.logger import init_logger @@ -483,11 +485,33 @@ def _register_configs() -> None: "FastVideo/Matrix-Game-2.0-TempleRun-Diffusers", ], model_detectors=[ - lambda path: "matrix-game" in path.lower() or "matrixgame" in path.lower(), + lambda path: any(token in path.lower() for token in ( + "matrix-game-2", + "matrixgame2", + "matrix-game-2.0", + )), ], model_family="matrixgame", default_preset="matrixgame2_i2v", ) + # MatrixGame 3.0 (I2V) + register_configs( + sampling_param_cls=MatrixGame3SamplingParam, + pipeline_config_cls=MatrixGame3I2V720PConfig, + workload_types=(WorkloadType.I2V, ), + hf_model_paths=[ + "FastVideo/Matrix-Game-3.0-Base-Distilled-Diffusers", + ], + model_detectors=[ + lambda path: any(token in path.lower() for token in ( + "matrix-game-3", + "matrixgame3", + "matrix-game-3.0", + )), + ], + model_family="matrixgame", + default_preset="matrixgame3_i2v", + ) # GEN3C (must register before generic Cosmos detector) register_configs( @@ -859,6 +883,8 @@ def _register_presets() -> None: ALL_PRESETS as LTX2_PRESETS, ) from fastvideo.pipelines.basic.matrixgame2.presets import ( ALL_PRESETS as MATRIXGAME2_PRESETS, ) + from fastvideo.pipelines.basic.matrixgame3.presets import ( + ALL_PRESETS as MATRIXGAME3_PRESETS, ) from fastvideo.pipelines.basic.sd35.presets import ( ALL_PRESETS as SD35_PRESETS, ) from fastvideo.pipelines.basic.stable_audio.presets import ( @@ -879,6 +905,7 @@ def _register_presets() -> None: LONGCAT_PRESETS, LTX2_PRESETS, MATRIXGAME2_PRESETS, + MATRIXGAME3_PRESETS, SD35_PRESETS, STABLE_AUDIO_PRESETS, TURBODIFFUSION_PRESETS, diff --git a/fastvideo/tests/api/test_presets.py b/fastvideo/tests/api/test_presets.py index cb149d79b..8e9491997 100644 --- a/fastvideo/tests/api/test_presets.py +++ b/fastvideo/tests/api/test_presets.py @@ -563,6 +563,9 @@ class TestSinglePresetFamilies: p2 = get_preset("matrixgame2_i2v", "matrixgame") assert p2.defaults["num_inference_steps"] == 3 assert p2.defaults["fps"] == 25 + p3 = get_preset("matrixgame3_i2v", "matrixgame") + assert p3.defaults["num_inference_steps"] == 3 + assert p3.defaults["height"] == 720 def test_longcat_presets_registered(self) -> None: import fastvideo.registry # noqa: F401 diff --git a/fastvideo/tests/ssim/test_matrixgame3_similarity.py b/fastvideo/tests/ssim/test_matrixgame3_similarity.py new file mode 100644 index 000000000..20092b624 --- /dev/null +++ b/fastvideo/tests/ssim/test_matrixgame3_similarity.py @@ -0,0 +1,192 @@ +# SPDX-License-Identifier: Apache-2.0 +import os + +import pytest + +from fastvideo import VideoGenerator +from fastvideo.api.sampling_param import SamplingParam +from fastvideo.logger import init_logger +from fastvideo.tests.ssim.reference_utils import ( + build_generated_output_dir, + build_reference_folder_path, + get_cuda_device_name, + resolve_device_reference_folder, + select_ssim_params, +) +from fastvideo.tests.utils import ( + compute_video_ssim_torchvision, + write_ssim_results, +) +from fastvideo.worker.multiproc_executor import MultiprocExecutor + +logger = init_logger(__name__) + +REQUIRED_GPUS = 1 + +device_reference_folder = resolve_device_reference_folder( + ( + ("A40", "A40"), + ("L40S", "L40S"), + ("H200", "H200"), + ), + device_name=get_cuda_device_name(), + fallback_device_prefix="L40S", + logger=logger, +) + + +MATRIXGAME3_PARAMS = { + "num_gpus": 1, + "model_path": "FastVideo/Matrix-Game-3.0-Base-Distilled-Diffusers", + "height": 352, + "width": 640, + "num_frames": 57, + "num_inference_steps": 3, + "guidance_scale": 5.0, + "seed": 42, +} +_MATRIXGAME3_FULL_QUALITY_DEFAULTS = SamplingParam.from_pretrained( + MATRIXGAME3_PARAMS["model_path"]) +MATRIXGAME3_FULL_QUALITY_PARAMS = { + "num_gpus": MATRIXGAME3_PARAMS["num_gpus"], + "model_path": MATRIXGAME3_PARAMS["model_path"], + "height": _MATRIXGAME3_FULL_QUALITY_DEFAULTS.height, + "width": _MATRIXGAME3_FULL_QUALITY_DEFAULTS.width, + "num_frames": MATRIXGAME3_PARAMS["num_frames"], + "num_inference_steps": _MATRIXGAME3_FULL_QUALITY_DEFAULTS.num_inference_steps, + "guidance_scale": _MATRIXGAME3_FULL_QUALITY_DEFAULTS.guidance_scale, + "seed": _MATRIXGAME3_FULL_QUALITY_DEFAULTS.seed, +} + + +MODEL_TO_PARAMS = { + "Matrix-Game-3.0-Diffusers-Base": MATRIXGAME3_PARAMS, +} +FULL_QUALITY_MODEL_TO_PARAMS = { + "Matrix-Game-3.0-Diffusers-Base": MATRIXGAME3_FULL_QUALITY_PARAMS, +} + +I2V_MODEL_TO_PARAMS = {} + +TEST_PROMPTS = [ + "A colorful, animated cityscape with a gas station and various buildings.", +] + +TEST_IMAGE_PATHS = [ + "https://raw.githubusercontent.com/SkyworkAI/Matrix-Game/main/Matrix-Game-3/demo_images/001/image.png", +] + + +@pytest.mark.parametrize("prompt", TEST_PROMPTS) +@pytest.mark.parametrize("ATTENTION_BACKEND", ["FLASH_ATTN"]) +@pytest.mark.parametrize("model_id", list(MODEL_TO_PARAMS.keys())) +def test_matrixgame3_similarity(prompt, ATTENTION_BACKEND, model_id): + """ + Test that runs MG3 inference (action conditions auto-generated from seed) + and compares the output to reference videos using SSIM. + """ + os.environ["FASTVIDEO_ATTENTION_BACKEND"] = ATTENTION_BACKEND + + script_dir = os.path.dirname(os.path.abspath(__file__)) + + output_dir = build_generated_output_dir( + script_dir, + device_reference_folder, + model_id, + ATTENTION_BACKEND, + ) + + os.makedirs(output_dir, exist_ok=True) + + params_map = select_ssim_params( + MODEL_TO_PARAMS, + FULL_QUALITY_MODEL_TO_PARAMS, + ) + BASE_PARAMS = params_map[model_id] + num_inference_steps = BASE_PARAMS["num_inference_steps"] + + init_kwargs = { + "num_gpus": BASE_PARAMS["num_gpus"], + "use_fsdp_inference": True, + "dit_layerwise_offload": False, + "dit_cpu_offload": False, + "vae_cpu_offload": False, + "text_encoder_cpu_offload": True, + "pin_cpu_memory": True, + } + + generation_kwargs = { + "num_inference_steps": num_inference_steps, + "output_path": output_dir, + "image_path": TEST_IMAGE_PATHS[0], + "height": BASE_PARAMS["height"], + "width": BASE_PARAMS["width"], + "num_frames": BASE_PARAMS["num_frames"], + "guidance_scale": BASE_PARAMS["guidance_scale"], + "seed": BASE_PARAMS["seed"], + "save_video": True, + } + + generator = VideoGenerator.from_pretrained(model_path=BASE_PARAMS["model_path"], **init_kwargs) + generator.generate_video(prompt, **generation_kwargs) + + if isinstance(generator.executor, MultiprocExecutor): + generator.executor.shutdown() + + assert os.path.exists(output_dir), f"Output video was not generated at {output_dir}" + + reference_folder = build_reference_folder_path( + script_dir, + device_reference_folder, + model_id, + ATTENTION_BACKEND, + ) + + if not os.path.exists(reference_folder): + logger.error("Reference folder missing") + raise FileNotFoundError(f"Reference video folder does not exist: {reference_folder}") + + prompt_prefix = prompt[:100].strip().rstrip(".") + + def _find_mp4(folder: str) -> str | None: + for filename in sorted(os.listdir(folder)): + if filename.endswith(".mp4") and prompt_prefix in filename: + return filename + return None + + reference_video_name = _find_mp4(reference_folder) + if not reference_video_name: + logger.error(f"Reference video not found for model: {model_id} with backend: {ATTENTION_BACKEND}") + raise FileNotFoundError("Reference video missing") + + generated_video_name = _find_mp4(output_dir) + if not generated_video_name: + logger.error(f"Generated video not found for model: {model_id} with backend: {ATTENTION_BACKEND}") + raise FileNotFoundError(f"Generated video missing in {output_dir}") + + reference_video_path = os.path.join(reference_folder, reference_video_name) + generated_video_path = os.path.join(output_dir, generated_video_name) + + logger.info(f"Computing SSIM between {reference_video_path} and {generated_video_path}") + ssim_values = compute_video_ssim_torchvision(reference_video_path, generated_video_path, use_ms_ssim=True) + + mean_ssim = ssim_values[0] + logger.info(f"SSIM mean value: {mean_ssim}") + logger.info(f"Writing SSIM results to directory: {output_dir}") + + success = write_ssim_results( + output_dir, + ssim_values, + reference_video_path, + generated_video_path, + num_inference_steps, + prompt, + ) + + if not success: + logger.error("Failed to write SSIM results to file") + + min_acceptable_ssim = 0.98 + assert mean_ssim >= min_acceptable_ssim, ( + f"SSIM value {mean_ssim} is below threshold {min_acceptable_ssim} for {model_id} with backend {ATTENTION_BACKEND}" + )