diff --git a/examples/inference/basic/basic_matrixgame.py b/examples/inference/basic/basic_matrixgame.py new file mode 100644 index 000000000..bae6fba3f --- /dev/null +++ b/examples/inference/basic/basic_matrixgame.py @@ -0,0 +1,75 @@ +from fastvideo import VideoGenerator +from fastvideo.configs.pipelines.wan import MatrixGameI2V480PConfig +from fastvideo.models.dits.matrix_game.utils import create_action_presets + +import torch + +# Available variants: "base_distilled_model", "gta_distilled_model", "templerun_distilled_model" +# Each variant has different keyboard_dim: +# - base_distilled_model: keyboard_dim=4 +# - gta_distilled_model: keyboard_dim=2 +# - templerun_distilled_model: keyboard_dim=7 (keyboard only, no mouse) +MODEL_VARIANT = "base_distilled_model" + +# Variant-specific settings +VARIANT_CONFIG = { + "base_distilled_model": { + "model_path": "FastVideo/Matrix-Game-2.0-Base-Diffusers", + "keyboard_dim": 4, + "image_url": "https://raw.githubusercontent.com/SkyworkAI/Matrix-Game/main/Matrix-Game-2/demo_images/universal/0000.png", + }, + "gta_distilled_model": { + "model_path": "FastVideo/Matrix-Game-2.0-GTA-Diffusers", + "keyboard_dim": 2, + "image_url": "https://raw.githubusercontent.com/SkyworkAI/Matrix-Game/main/Matrix-Game-2/demo_images/gta_drive/0000.png", + }, + "templerun_distilled_model": { + "model_path": "FastVideo/Matrix-Game-2.0-TempleRun-Diffusers", + "keyboard_dim": 7, + "image_url": "https://raw.githubusercontent.com/SkyworkAI/Matrix-Game/main/Matrix-Game-2/demo_images/temple_run/0000.png", + }, +} + + +OUTPUT_PATH = "video_samples_matrixgame2" +def main(): + # FastVideo will automatically use the optimal default arguments for the + # model. + # If a local path is provided, FastVideo will make a best effort + # attempt to identify the optimal arguments. + config = VARIANT_CONFIG[MODEL_VARIANT] + + generator = VideoGenerator.from_pretrained( + config["model_path"], + # FastVideo will automatically handle distributed setup + num_gpus=1, + use_fsdp_inference=True, + dit_cpu_offload=True, # DiT need to be offloaded for MoE + vae_cpu_offload=False, + text_encoder_cpu_offload=True, + # Set pin_cpu_memory to false if CPU RAM is limited and there're no frequent CPU-GPU transfer + pin_cpu_memory=True, + # image_encoder_cpu_offload=False, + ) + + num_frames = 597 + actions = create_action_presets(num_frames, keyboard_dim=config["keyboard_dim"]) + grid_sizes = torch.tensor([150, 44, 80]) + + generator.generate_video( + prompt="", + image_path=config["image_url"], + mouse_cond=actions["mouse"].unsqueeze(0), + keyboard_cond=actions["keyboard"].unsqueeze(0), + grid_sizes=grid_sizes, + num_frames=num_frames, + height=352, + width=640, + num_inference_steps=50, + output_path=OUTPUT_PATH, + save_video=True, + ) + + +if __name__ == "__main__": + main() diff --git a/fastvideo/configs/models/dits/matrixgame.py b/fastvideo/configs/models/dits/matrixgame.py new file mode 100644 index 000000000..f0f7e0e24 --- /dev/null +++ b/fastvideo/configs/models/dits/matrixgame.py @@ -0,0 +1,83 @@ +from dataclasses import dataclass, field +from fastvideo.configs.models.dits.wanvideo import WanVideoArchConfig, WanVideoConfig + + +@dataclass +class MatrixGameWanVideoArchConfig(WanVideoArchConfig): + # Override param_names_mapping to remove patch_embedding transformation + # because MatrixGame checkpoints already have patch_embedding.proj format + param_names_mapping: dict = field( + default_factory=lambda: { + # Removed: r"^patch_embedding\.(.*)$": r"patch_embedding.proj.\1" + # because checkpoint already has correct format + r"^condition_embedder\.text_embedder\.linear_1\.(.*)$": + r"condition_embedder.text_embedder.fc_in.\1", + r"^condition_embedder\.text_embedder\.linear_2\.(.*)$": + r"condition_embedder.text_embedder.fc_out.\1", + r"^condition_embedder\.time_embedder\.linear_1\.(.*)$": + r"condition_embedder.time_embedder.mlp.fc_in.\1", + r"^condition_embedder\.time_embedder\.linear_2\.(.*)$": + r"condition_embedder.time_embedder.mlp.fc_out.\1", + r"^condition_embedder\.time_proj\.(.*)$": + r"condition_embedder.time_modulation.linear.\1", + r"^condition_embedder\.image_embedder\.ff\.net\.0\.proj\.(.*)$": + r"condition_embedder.image_embedder.ff.fc_in.\1", + r"^condition_embedder\.image_embedder\.ff\.net\.2\.(.*)$": + r"condition_embedder.image_embedder.ff.fc_out.\1", + r"^blocks\.(\d+)\.attn1\.to_q\.(.*)$": + r"blocks.\1.to_q.\2", + r"^blocks\.(\d+)\.attn1\.to_k\.(.*)$": + r"blocks.\1.to_k.\2", + r"^blocks\.(\d+)\.attn1\.to_v\.(.*)$": + r"blocks.\1.to_v.\2", + r"^blocks\.(\d+)\.attn1\.to_out\.0\.(.*)$": + r"blocks.\1.to_out.\2", + r"^blocks\.(\d+)\.attn1\.norm_q\.(.*)$": + r"blocks.\1.norm_q.\2", + r"^blocks\.(\d+)\.attn1\.norm_k\.(.*)$": + r"blocks.\1.norm_k.\2", + r"^blocks\.(\d+)\.attn2\.to_out\.0\.(.*)$": + r"blocks.\1.attn2.to_out.\2", + r"^blocks\.(\d+)\.ffn\.net\.0\.proj\.(.*)$": + r"blocks.\1.ffn.fc_in.\2", + r"^blocks\.(\d+)\.ffn\.net\.2\.(.*)$": + r"blocks.\1.ffn.fc_out.\2", + r"^blocks\.(\d+)\.norm2\.(.*)$": + r"blocks.\1.self_attn_residual_norm.norm.\2", + }) + + 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": 1536, + "keyboard_dim_in": 4, + "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, + }) + + local_attn_size: int = -1 + sink_size: int = 0 + num_frames_per_block: int = 3 + text_len: int = 512 + text_dim: int = 0 + image_dim: int = 1280 + + +@dataclass +class MatrixGameWanVideoConfig(WanVideoConfig): + arch_config: MatrixGameWanVideoArchConfig = field( + default_factory=MatrixGameWanVideoArchConfig) + prefix: str = "Wan" diff --git a/fastvideo/configs/pipelines/registry.py b/fastvideo/configs/pipelines/registry.py index ddd8de3aa..d3de0b872 100644 --- a/fastvideo/configs/pipelines/registry.py +++ b/fastvideo/configs/pipelines/registry.py @@ -14,7 +14,8 @@ from fastvideo.configs.pipelines.wan import ( FastWan2_1_T2V_480P_Config, FastWan2_2_TI2V_5B_Config, Wan2_2_I2V_A14B_Config, Wan2_2_T2V_A14B_Config, Wan2_2_TI2V_5B_Config, WanI2V480PConfig, WanI2V720PConfig, WanT2V480PConfig, WanT2V720PConfig, - SelfForcingWanT2V480PConfig, WANV2VConfig, SelfForcingWan2_2_T2V480PConfig) + SelfForcingWanT2V480PConfig, WANV2VConfig, SelfForcingWan2_2_T2V480PConfig, + MatrixGameI2V480PConfig) # isort: on from fastvideo.logger import init_logger from fastvideo.utils import (maybe_download_model_index, @@ -45,18 +46,30 @@ PIPE_NAME_TO_CONFIG: dict[str, type[PipelineConfig]] = { "Wan-AI/Wan2.2-T2V-A14B-Diffusers": Wan2_2_T2V_A14B_Config, "Wan-AI/Wan2.2-I2V-A14B-Diffusers": Wan2_2_I2V_A14B_Config, "nvidia/Cosmos-Predict2-2B-Video2World": CosmosConfig, + "FastVideo/Matrix-Game-2.0-Base-Diffusers": MatrixGameI2V480PConfig, + "FastVideo/Matrix-Game-2.0-GTA-Diffusers": MatrixGameI2V480PConfig, + "FastVideo/Matrix-Game-2.0-TempleRun-Diffusers": MatrixGameI2V480PConfig, # Add other specific weight variants } # For determining pipeline type from model ID PIPELINE_DETECTOR: dict[str, Callable[[str], bool]] = { - "hunyuan": lambda id: "hunyuan" in id.lower(), - "wanpipeline": lambda id: "wanpipeline" in id.lower(), - "wanimagetovideo": lambda id: "wanimagetovideo" in id.lower(), - "wandmdpipeline": lambda id: "wandmdpipeline" in id.lower(), - "wancausaldmdpipeline": lambda id: "wancausaldmdpipeline" in id.lower(), - "stepvideo": lambda id: "stepvideo" in id.lower(), - "cosmos": lambda id: "cosmos" in id.lower(), + "hunyuan": + lambda id: "hunyuan" in id.lower(), + "matrixgame": + lambda id: "matrix-game" in id.lower() or "matrixgame" in id.lower(), + "wanpipeline": + lambda id: "wanpipeline" in id.lower(), + "wanimagetovideo": + lambda id: "wanimagetovideo" in id.lower(), + "wandmdpipeline": + lambda id: "wandmdpipeline" in id.lower(), + "wancausaldmdpipeline": + lambda id: "wancausaldmdpipeline" in id.lower(), + "stepvideo": + lambda id: "stepvideo" in id.lower(), + "cosmos": + lambda id: "cosmos" in id.lower(), # Add other pipeline architecture detectors } @@ -64,6 +77,7 @@ PIPELINE_DETECTOR: dict[str, Callable[[str], bool]] = { PIPELINE_FALLBACK_CONFIG: dict[str, type[PipelineConfig]] = { "hunyuan": HunyuanConfig, # Base Hunyuan config as fallback for any Hunyuan variant + "matrixgame": MatrixGameI2V480PConfig, "wanpipeline": WanT2V480PConfig, # Base Wan config as fallback for any Wan variant "wanimagetovideo": WanI2V480PConfig, diff --git a/fastvideo/configs/pipelines/wan.py b/fastvideo/configs/pipelines/wan.py index a3092bfc2..55d917fb0 100644 --- a/fastvideo/configs/pipelines/wan.py +++ b/fastvideo/configs/pipelines/wan.py @@ -6,6 +6,7 @@ import torch from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig from fastvideo.configs.models.dits import WanVideoConfig +from fastvideo.configs.models.dits.matrixgame import MatrixGameWanVideoConfig from fastvideo.configs.models.encoders import (BaseEncoderOutput, CLIPVisionConfig, T5Config, WAN2_1ControlCLIPVisionConfig) @@ -190,3 +191,23 @@ class SelfForcingWan2_2_T2V480PConfig(Wan2_2_T2V_A14B_Config): def __post_init__(self) -> None: self.vae_config.load_encoder = True self.vae_config.load_decoder = True + + +# ============================================= +# ============= Matrix Game =================== +# ============================================= +@dataclass +class MatrixGameI2V480PConfig(WanI2V480PConfig): + dit_config: DiTConfig = field(default_factory=MatrixGameWanVideoConfig) + + image_encoder_config: EncoderConfig = field( + default_factory=WAN2_1ControlCLIPVisionConfig) + + is_causal: bool = True + flow_shift: float | None = 5.0 + dmd_denoising_steps: list[int] | None = field( + default_factory=lambda: [1000, 666, 333]) + warp_denoising_step: bool = True + context_noise: int = 0 + num_frames_per_block: int = 3 + # sliding_window_num_frames: int = 15 diff --git a/fastvideo/configs/sample/base.py b/fastvideo/configs/sample/base.py index 873777561..8d62f95a1 100644 --- a/fastvideo/configs/sample/base.py +++ b/fastvideo/configs/sample/base.py @@ -17,10 +17,16 @@ class SamplingParam: # Image inputs image_path: str | None = None + pil_image: Any | None = None # Video inputs video_path: str | None = None + # Action control inputs (Matrix-Game) + mouse_cond: Any | None = None # Shape: (B, T, 2) + keyboard_cond: Any | None = None # Shape: (B, T, K) + grid_sizes: Any | None = None # Shape: (3,) [F,H,W] + # Text inputs prompt: str | list[str] | None = None negative_prompt: str = "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" diff --git a/fastvideo/configs/sample/registry.py b/fastvideo/configs/sample/registry.py index ada25bf0e..58969f2f0 100644 --- a/fastvideo/configs/sample/registry.py +++ b/fastvideo/configs/sample/registry.py @@ -23,6 +23,7 @@ from fastvideo.configs.sample.wan import ( Wan2_1_Fun_1_3B_Control_SamplingParam, SelfForcingWan2_1_T2V_1_3B_480P_SamplingParam, SelfForcingWan2_2_T2V_A14B_480P_SamplingParam, + MatrixGame2_SamplingParam, ) # isort: on from fastvideo.logger import init_logger @@ -32,44 +33,32 @@ from fastvideo.utils import (maybe_download_model_index, logger = init_logger(__name__) # Registry maps specific model weights to their config classes SAMPLING_PARAM_REGISTRY: dict[str, Any] = { - "FastVideo/FastHunyuan-diffusers": - FastHunyuanSamplingParam, - "hunyuanvideo-community/HunyuanVideo": - HunyuanSamplingParam, - "FastVideo/stepvideo-t2v-diffusers": - StepVideoT2VSamplingParam, + "FastVideo/FastHunyuan-diffusers": FastHunyuanSamplingParam, + "hunyuanvideo-community/HunyuanVideo": HunyuanSamplingParam, + "FastVideo/stepvideo-t2v-diffusers": StepVideoT2VSamplingParam, # Wan2.1 - "Wan-AI/Wan2.1-T2V-1.3B-Diffusers": - WanT2V_1_3B_SamplingParam, - "Wan-AI/Wan2.1-T2V-14B-Diffusers": - WanT2V_14B_SamplingParam, - "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": - WanI2V_14B_480P_SamplingParam, - "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers": - WanI2V_14B_720P_SamplingParam, + "Wan-AI/Wan2.1-T2V-1.3B-Diffusers": WanT2V_1_3B_SamplingParam, + "Wan-AI/Wan2.1-T2V-14B-Diffusers": WanT2V_14B_SamplingParam, + "Wan-AI/Wan2.1-I2V-14B-480P-Diffusers": WanI2V_14B_480P_SamplingParam, + "Wan-AI/Wan2.1-I2V-14B-720P-Diffusers": WanI2V_14B_720P_SamplingParam, "weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers": Wan2_1_Fun_1_3B_InP_SamplingParam, "IRMChen/Wan2.1-Fun-1.3B-Control-Diffusers": Wan2_1_Fun_1_3B_Control_SamplingParam, # Wan2.2 - "Wan-AI/Wan2.2-TI2V-5B-Diffusers": - Wan2_2_TI2V_5B_SamplingParam, + "Wan-AI/Wan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_SamplingParam, "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers": Wan2_2_TI2V_5B_SamplingParam, - "Wan-AI/Wan2.2-T2V-A14B-Diffusers": - Wan2_2_T2V_A14B_SamplingParam, - "Wan-AI/Wan2.2-I2V-A14B-Diffusers": - Wan2_2_I2V_A14B_SamplingParam, + "Wan-AI/Wan2.2-T2V-A14B-Diffusers": Wan2_2_T2V_A14B_SamplingParam, + "Wan-AI/Wan2.2-I2V-A14B-Diffusers": Wan2_2_I2V_A14B_SamplingParam, # FastWan2.1 - "FastVideo/FastWan2.1-T2V-1.3B-Diffusers": - FastWanT2V480P_SamplingParam, + "FastVideo/FastWan2.1-T2V-1.3B-Diffusers": FastWanT2V480P_SamplingParam, # FastWan2.2 - "FastVideo/FastWan2.2-TI2V-5B-Diffusers": - Wan2_2_TI2V_5B_SamplingParam, + "FastVideo/FastWan2.2-TI2V-5B-Diffusers": Wan2_2_TI2V_5B_SamplingParam, # Causal Self-Forcing Wan2.1 "wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers": @@ -85,17 +74,30 @@ SAMPLING_PARAM_REGISTRY: dict[str, Any] = { "nvidia/Cosmos-Predict2-2B-Video2World": Cosmos_Predict2_2B_Video2World_SamplingParam, + # MatrixGame2.0 models + "FastVideo/Matrix-Game-2.0-Base-Diffusers": MatrixGame2_SamplingParam, + "FastVideo/Matrix-Game-2.0-GTA-Diffusers": MatrixGame2_SamplingParam, + "FastVideo/Matrix-Game-2.0-TempleRun-Diffusers": MatrixGame2_SamplingParam, + # Add other specific weight variants } # For determining pipeline type from model ID SAMPLING_PARAM_DETECTOR: dict[str, Callable[[str], bool]] = { - "hunyuan": lambda id: "hunyuan" in id.lower(), - "wanpipeline": lambda id: "wanpipeline" in id.lower(), - "wanimagetovideo": lambda id: "wanimagetovideo" in id.lower(), - "stepvideo": lambda id: "stepvideo" in id.lower(), - "wandmdpipeline": lambda id: "wandmdpipeline" in id.lower(), - "wancausaldmdpipeline": lambda id: "wancausaldmdpipeline" in id.lower(), + "hunyuan": + lambda id: "hunyuan" in id.lower(), + "wanpipeline": + lambda id: "wanpipeline" in id.lower(), + "wanimagetovideo": + lambda id: "wanimagetovideo" in id.lower(), + "stepvideo": + lambda id: "stepvideo" in id.lower(), + "wandmdpipeline": + lambda id: "wandmdpipeline" in id.lower(), + "wancausaldmdpipeline": + lambda id: "wancausaldmdpipeline" in id.lower(), + "matrixgame": + lambda id: "matrixgame" in id.lower() or "matrix-game" in id.lower(), # Add other pipeline architecture detectors } @@ -109,6 +111,7 @@ SAMPLING_FALLBACK_PARAM: dict[str, Any] = { "wandmdpipeline": FastWanT2V480P_SamplingParam, "wancausaldmdpipeline": SelfForcingWan2_1_T2V_1_3B_480P_SamplingParam, "stepvideo": StepVideoT2VSamplingParam, + "matrixgame": MatrixGame2_SamplingParam, # Other fallbacks by architecture } @@ -116,6 +119,20 @@ SAMPLING_FALLBACK_PARAM: dict[str, Any] = { def get_sampling_param_cls_for_name(pipeline_name_or_path: str) -> Any | None: """Get the appropriate sampling param for specific pretrained weights.""" + # First try exact match for specific weights + if pipeline_name_or_path in SAMPLING_PARAM_REGISTRY: + return SAMPLING_PARAM_REGISTRY[pipeline_name_or_path] + + # Try partial matches (for local paths that might include the weight ID) + for registered_id, config_class in SAMPLING_PARAM_REGISTRY.items(): + if registered_id in pipeline_name_or_path: + return config_class + + matrixgame_patterns = ["Matrix-Game", "Skywork--Matrix-Game", "matrixgame"] + for pattern in matrixgame_patterns: + if pattern.lower() in pipeline_name_or_path.lower(): + return MatrixGame2_SamplingParam + if os.path.exists(pipeline_name_or_path): config = verify_model_config_and_directory(pipeline_name_or_path) logger.warning( @@ -126,15 +143,6 @@ def get_sampling_param_cls_for_name(pipeline_name_or_path: str) -> Any | None: pipeline_name = config["_class_name"] - # First try exact match for specific weights - if pipeline_name_or_path in SAMPLING_PARAM_REGISTRY: - return SAMPLING_PARAM_REGISTRY[pipeline_name_or_path] - - # Try partial matches (for local paths that might include the weight ID) - for registered_id, config_class in SAMPLING_PARAM_REGISTRY.items(): - if registered_id in pipeline_name_or_path: - return config_class - # If no match, try to use the fallback config fallback_config = None # Try to determine pipeline architecture for fallback diff --git a/fastvideo/configs/sample/wan.py b/fastvideo/configs/sample/wan.py index 120235055..85f76a7ac 100644 --- a/fastvideo/configs/sample/wan.py +++ b/fastvideo/configs/sample/wan.py @@ -196,3 +196,14 @@ class SelfForcingWan2_2_T2V_A14B_480P_SamplingParam( height: int = 448 width: int = 832 fps: int = 16 + + +@dataclass +class MatrixGame2_SamplingParam(SamplingParam): + height: int = 352 + width: int = 640 + num_frames: int = 57 + fps: int = 25 + guidance_scale: float = 1.0 + num_inference_steps: int = 3 + negative_prompt: str | None = None diff --git a/fastvideo/entrypoints/video_generator.py b/fastvideo/entrypoints/video_generator.py index f9e50ea21..d14b7a73d 100644 --- a/fastvideo/entrypoints/video_generator.py +++ b/fastvideo/entrypoints/video_generator.py @@ -102,6 +102,11 @@ class VideoGenerator: self, prompt: str | None = None, sampling_param: SamplingParam | None = None, + # Action control inputs (Matrix-Game) + mouse_cond: torch.Tensor | None = None, + keyboard_cond: torch.Tensor | None = None, + grid_sizes: tuple[int, int, int] | list[int] | torch.Tensor + | None = None, **kwargs, ) -> dict[str, Any] | list[np.ndarray] | list[dict[str, Any]]: """ @@ -131,6 +136,15 @@ class VideoGenerator: if sampling_param is None: sampling_param = SamplingParam.from_pretrained( self.fastvideo_args.model_path) + + # Add action control inputs to kwargs if provided + if mouse_cond is not None: + kwargs['mouse_cond'] = mouse_cond + if keyboard_cond is not None: + kwargs['keyboard_cond'] = keyboard_cond + if grid_sizes is not None: + kwargs['grid_sizes'] = grid_sizes + sampling_param.update(kwargs) if self.fastvideo_args.prompt_txt is not None or sampling_param.prompt_path is not None: diff --git a/fastvideo/models/dits/matrix_game/__init__.py b/fastvideo/models/dits/matrix_game/__init__.py new file mode 100644 index 000000000..0da9c954c --- /dev/null +++ b/fastvideo/models/dits/matrix_game/__init__.py @@ -0,0 +1,11 @@ +from .model import MatrixGameWanModel, MatrixGameTransformerBlock +from .causal_model import CausalMatrixGameWanModel, CausalMatrixGameTransformerBlock +from .action_module import ActionModule + +__all__ = [ + "MatrixGameWanModel", + "MatrixGameTransformerBlock", + "CausalMatrixGameWanModel", + "CausalMatrixGameTransformerBlock", + "ActionModule", +] diff --git a/fastvideo/models/dits/matrix_game/action_module.py b/fastvideo/models/dits/matrix_game/action_module.py new file mode 100644 index 000000000..ae4d01b4b --- /dev/null +++ b/fastvideo/models/dits/matrix_game/action_module.py @@ -0,0 +1,567 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from Matrix-Game: https://github.com/SkyworkAI/Matrix-Game/blob/main/Matrix-Game-2/wan/modules/action_module.py + +from einops import rearrange +import torch +import torch.nn as nn +import math +from torch.nn.attention.flex_attention import flex_attention + +from fastvideo.attention import LocalAttention +from fastvideo.layers.linear import ReplicatedLinear +from fastvideo.layers.layernorm import FP32LayerNorm, RMSNorm +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 + + +DISABLE_COMPILE = False +flex_attention = torch.compile( + flex_attention, dynamic=False, mode="max-autotune-no-cudagraphs") + + +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, + start_offset: int = 0, +) -> tuple[torch.Tensor, torch.Tensor]: + seq_len = xq.shape[1] + + # Slice frequencies based on offset + cos = freqs_cos[start_offset:start_offset + seq_len] # [S, D] + sin = freqs_sin[start_offset:start_offset + seq_len] # [S, D] + + # Move to device + cos = cos.to(xq.device) + sin = sin.to(xq.device) + + # Convert from [S, D] (interleaved) back to [S, D/2] + cos_half = cos[:, ::2] # [S, D/2] + sin_half = sin[:, ::2] # [S, D/2] + + # xq/xk are [B, S, H, D], need to reshape for each batch + B, S, H, D = xq.shape + + 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 ActionModule(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, + local_attn_size = 6, + blocks: list | None = None, + ): + super().__init__() + # Initialize mutable defaults + 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.local_attn_size = local_attn_size + 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 = ( + RMSNorm(head_dim, eps=1e-6) + if qk_norm + else nn.Identity() + ) + self.img_attn_k_norm = ( + RMSNorm(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 = ( + RMSNorm(head_dim_key, eps=1e-6) + if qk_norm + else nn.Identity() + ) + self.key_attn_k_norm = ( + RMSNorm(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 #self.unpatchify_channels + 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, start_offset=0): + target_ndim = 3 + ndim = 5 - 2 + + latents_size = [video_length+start_offset, 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(self, x, tt, th, tw, mouse_condition=None, keyboard_condition=None, block_mask_mouse=None, block_mask_keyboard=None, is_causal=False, kv_cache_mouse=None, kv_cache_keyboard=None, start_frame=0, use_rope_keyboard=True, num_frame_per_block=3): + ''' + hidden_states: B, tt*th*tw, C + mouse_condition: B, N_frames, C1 + keyboard_condition: B, N_frames, C2 + ''' + assert use_rope_keyboard + + 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 + N_feats = int((N_frames - 1) / self.vae_time_compression_ratio) + 1 + + # Lazy initialization of freqs on first forward pass + if self._freqs_cos is None or self._freqs_sin is None: + self._freqs_cos, self._freqs_sin = self.get_rotary_pos_embed( + 7500, self.patch_size[1], self.patch_size[2], 64, + self.mouse_qk_dim_list, start_offset=0 + ) + + # Defined freqs_cis early so it's available for both mouse and keyboard + freqs_cis = (self._freqs_cos, self._freqs_sin) + + assert (N_feats == tt and ((is_causal and kv_cache_mouse is None) or not is_causal)) or ((N_frames - 1) // self.vae_time_compression_ratio + 1 == start_frame + num_frame_per_block and is_causal) + + 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) # 65*272*480 -> 17*(272//16)*(480//16) -> 8670 + B, N_frames, C = mouse_condition.shape + else: + hidden_states = x + # padding + + pad_t = self.vae_time_compression_ratio * self.windows_size + if self.enable_mouse and mouse_condition is not None: + pad = mouse_condition[:, 0:1, :].expand(-1, pad_t, -1) + mouse_condition = torch.cat([pad, mouse_condition], dim=1) + if is_causal and kv_cache_mouse is not None: + mouse_condition = mouse_condition[:, self.vae_time_compression_ratio*(N_feats - num_frame_per_block - self.windows_size) + pad_t:, :] + 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(num_frame_per_block)] + else: + 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) + + S = th * tw + group_mouse = group_mouse.unsqueeze(-1).expand(B, num_frame_per_block, pad_t, C, S) + group_mouse = group_mouse.permute(0, 4, 1, 2, 3).reshape(B * S, num_frame_per_block, pad_t * C) + + group_mouse = torch.cat([hidden_states, group_mouse], dim = -1) + group_mouse = self.mouse_mlp(group_mouse) + # qkv + 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) # BHW F H C + q = self.img_attn_q_norm(q).to(v) + k = self.img_attn_k_norm(k).to(v) + # rope embd + + + # freqs_cis = (self.freqs_cos, self.freqs_sin) + + + q, k = _apply_rotary_emb_qk(q, k, freqs_cis[0], freqs_cis[1], start_offset=start_frame) + ## TODO: adding cache here + if is_causal: + if kv_cache_mouse is None: + assert q.shape[0] == k.shape[0] and q.shape[0] % 880 == 0 # == 880, f"{q.shape[0]},{k.shape[0]}" + padded_length = math.ceil(q.shape[1] / 32) * 32 - q.shape[1] + padded_q = torch.cat( + [q, + torch.zeros([q.shape[0], padded_length, q.shape[2], q.shape[3]], + device=q.device, dtype=v.dtype)], + dim=1 + ) + padded_k = torch.cat( + [k, torch.zeros([k.shape[0], padded_length, k.shape[2], k.shape[3]], + device=k.device, dtype=v.dtype)], + dim=1 + ) + padded_v = torch.cat( + [v, torch.zeros([v.shape[0], padded_length, v.shape[2], v.shape[3]], + device=v.device, dtype=v.dtype)], + dim=1 + ) + attn = flex_attention( + query=padded_q.transpose(2, 1), # after: B, HW, F, C + key=padded_k.transpose(2, 1), + value=padded_v.transpose(2, 1), + block_mask=block_mask_mouse + )[:, :, :-padded_length].transpose(2, 1) + else: + current_start = start_frame + current_end = current_start + q.shape[1] + + assert q.shape[1] == num_frame_per_block + sink_size = 0 + max_attention_size = self.local_attn_size + sink_tokens = sink_size * 1 + kv_cache_size = kv_cache_mouse["k"].shape[1] + num_new_tokens = q.shape[1] + + + if (current_end > kv_cache_mouse["global_end_index"].item()) and ( + num_new_tokens + kv_cache_mouse["local_end_index"].item() > kv_cache_size): + num_evicted_tokens = num_new_tokens + kv_cache_mouse["local_end_index"].item() - kv_cache_size + num_rolled_tokens = kv_cache_mouse["local_end_index"].item() - num_evicted_tokens - sink_tokens + kv_cache_mouse["k"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \ + kv_cache_mouse["k"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone() + kv_cache_mouse["v"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \ + kv_cache_mouse["v"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone() + # Insert the new keys/values at the end + local_end_index = kv_cache_mouse["local_end_index"].item() + current_end - \ + kv_cache_mouse["global_end_index"].item() - num_evicted_tokens + local_start_index = local_end_index - num_new_tokens + else: + local_end_index = kv_cache_mouse["local_end_index"].item() + current_end - kv_cache_mouse["global_end_index"].item() + local_start_index = local_end_index - num_new_tokens + kv_cache_mouse["k"][:, local_start_index:local_end_index] = k + kv_cache_mouse["v"][:, local_start_index:local_end_index] = v + + attn = self.mouse_attn_layer( + q, + kv_cache_mouse["k"][:, max(0, local_end_index - max_attention_size):local_end_index], + kv_cache_mouse["v"][:, max(0, local_end_index - max_attention_size):local_end_index], + ) + kv_cache_mouse["global_end_index"].fill_(current_end) + kv_cache_mouse["local_end_index"].fill_(local_end_index) + else: + attn = self.mouse_attn_layer(q, k, v) + # Compute cu_squlens and max_seqlen for flash attention + # qk norm + attn = rearrange(attn, '(b S) T h d -> b (T S) (h d)',b=B) + + hidden_states = rearrange(x, "(B S) T C -> B (T S) C", B=B) + attn, _ = self.proj_mouse(attn) + + hidden_states = hidden_states + attn + + if self.enable_keyboard and keyboard_condition is not None: + pad = keyboard_condition[:, 0:1, :].expand(-1, pad_t, -1) + keyboard_condition = torch.cat([pad, keyboard_condition], dim=1) + if is_causal and kv_cache_keyboard is not None: + keyboard_condition = keyboard_condition[:, self.vae_time_compression_ratio*(N_feats - num_frame_per_block - self.windows_size) + pad_t:, :] # keyboard_condition[:, self.vae_time_compression_ratio*(start_frame - self.windows_size) + pad_t:start_frame * self.vae_time_compression_ratio + pad_t,:] + 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(num_frame_per_block)] + else: + 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) # B F RW C + group_keyboard = group_keyboard.reshape(shape=(group_keyboard.shape[0],group_keyboard.shape[1],-1)) + # apply cross attn + 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) + + # Compute cu_squlens and max_seqlen for flash attention + # qk norm + + q = self.key_attn_q_norm(q).to(v) + k = self.key_attn_k_norm(k).to(v) + S = th * tw + assert S == 880 + # position embed + if use_rope_keyboard: + 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) + q, k = _apply_rotary_emb_qk(q, k, freqs_cis[0], freqs_cis[1], start_offset=start_frame) + + k1, k2, k3, k4 = k.shape + k = k.expand(S, k2, k3, k4) + v = v.expand(S, k2, k3, k4) + + + if is_causal: + if kv_cache_keyboard is None: + assert q.shape[0] == k.shape[0] and q.shape[0] % 880 == 0 + + padded_length = math.ceil(q.shape[1] / 32) * 32 - q.shape[1] + padded_q = torch.cat( + [q, + torch.zeros([q.shape[0], padded_length, q.shape[2], q.shape[3]], + device=q.device, dtype=v.dtype)], + dim=1 + ) + padded_k = torch.cat( + [k, torch.zeros([k.shape[0], padded_length, k.shape[2], k.shape[3]], + device=k.device, dtype=v.dtype)], + dim=1 + ) + padded_v = torch.cat( + [v, torch.zeros([v.shape[0], padded_length, v.shape[2], v.shape[3]], + device=v.device, dtype=v.dtype)], + dim=1 + ) + attn = flex_attention( + query=padded_q.transpose(2, 1), # after: B, HW, F, C + key=padded_k.transpose(2, 1), + value=padded_v.transpose(2, 1), + block_mask=block_mask_keyboard + )[:, :, :-padded_length].transpose(2, 1) + else: + current_start = start_frame + current_end = current_start + k.shape[1] + assert k.shape[1] == num_frame_per_block + sink_size = 0 + max_attention_size = self.local_attn_size + sink_tokens = sink_size * 1 + kv_cache_size = kv_cache_keyboard["k"].shape[1] + num_new_tokens = k.shape[1] + + if (current_end > kv_cache_keyboard["global_end_index"].item()) and ( + num_new_tokens + kv_cache_keyboard["local_end_index"].item() > kv_cache_size): + num_evicted_tokens = num_new_tokens + kv_cache_keyboard["local_end_index"].item() - kv_cache_size + num_rolled_tokens = kv_cache_keyboard["local_end_index"].item() - num_evicted_tokens - sink_tokens + kv_cache_keyboard["k"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \ + kv_cache_keyboard["k"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone() + kv_cache_keyboard["v"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \ + kv_cache_keyboard["v"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone() + # Insert the new keys/values at the end + local_end_index = kv_cache_keyboard["local_end_index"].item() + current_end - \ + kv_cache_keyboard["global_end_index"].item() - num_evicted_tokens + local_start_index = local_end_index - num_new_tokens + else: + local_end_index = kv_cache_keyboard["local_end_index"].item() + current_end - kv_cache_keyboard["global_end_index"].item() + local_start_index = local_end_index - num_new_tokens + assert k.shape[0] == 880 # BS == 1 or the cache should not be saved/ load method should be modified + kv_cache_keyboard["k"][:, local_start_index:local_end_index] = k[:1] + kv_cache_keyboard["v"][:, local_start_index:local_end_index] = v[:1] + + attn = self.keyboard_attn_layer( + q, + kv_cache_keyboard["k"][:, max(0, local_end_index - max_attention_size):local_end_index].repeat(S, 1, 1, 1), + kv_cache_keyboard["v"][:, max(0, local_end_index - max_attention_size):local_end_index].repeat(S, 1, 1, 1), + ) + + kv_cache_keyboard["global_end_index"].fill_(current_end) + kv_cache_keyboard["local_end_index"].fill_(local_end_index) + else: + attn = self.keyboard_attn_layer(q, k, v) + attn = rearrange(attn, '(B S) T H D -> B (T S) (H D)', S=S) + else: + if is_causal: + if kv_cache_keyboard is None: + + padded_length = math.ceil(q.shape[1] / 32) * 32 - q.shape[1] + padded_q = torch.cat( + [q, + torch.zeros([q.shape[0], padded_length, q.shape[2], q.shape[3]], + device=q.device, dtype=v.dtype)], + dim=1 + ) + padded_k = torch.cat( + [k, torch.zeros([k.shape[0], padded_length, k.shape[2], k.shape[3]], + device=k.device, dtype=v.dtype)], + dim=1 + ) + padded_v = torch.cat( + [v, torch.zeros([v.shape[0], padded_length, v.shape[2], v.shape[3]], + device=v.device, dtype=v.dtype)], + dim=1 + ) + attn = flex_attention( + query=padded_q.transpose(2, 1), # after: B, HW, F, C + key=padded_k.transpose(2, 1), + value=padded_v.transpose(2, 1), + block_mask=block_mask_keyboard + )[:, :, :-padded_length].transpose(2, 1) + else: + current_start = start_frame + current_end = current_start + k.shape[1] + assert k.shape[1] == num_frame_per_block + sink_size = 0 + max_attention_size = self.local_attn_size + sink_tokens = sink_size * 1 + kv_cache_size = kv_cache_keyboard["k"].shape[1] + num_new_tokens = k.shape[1] + + + if (current_end > kv_cache_keyboard["global_end_index"].item()) and ( + num_new_tokens + kv_cache_keyboard["local_end_index"].item() > kv_cache_size): + num_evicted_tokens = num_new_tokens + kv_cache_keyboard["local_end_index"].item() - kv_cache_size + num_rolled_tokens = kv_cache_keyboard["local_end_index"].item() - num_evicted_tokens - sink_tokens + kv_cache_keyboard["k"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \ + kv_cache_keyboard["k"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone() + kv_cache_keyboard["v"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \ + kv_cache_keyboard["v"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone() + # Insert the new keys/values at the end + local_end_index = kv_cache_keyboard["local_end_index"].item() + current_end - \ + kv_cache_keyboard["global_end_index"].item() - num_evicted_tokens + local_start_index = local_end_index - num_new_tokens + + + else: + local_end_index = kv_cache_keyboard["local_end_index"].item() + current_end - kv_cache_keyboard["global_end_index"].item() + local_start_index = local_end_index - num_new_tokens + kv_cache_keyboard["k"][:, local_start_index:local_end_index] = k + kv_cache_keyboard["v"][:, local_start_index:local_end_index] = v + attn = self.keyboard_attn_layer( + q, + kv_cache_keyboard["k"][:, max(0, local_end_index - max_attention_size):local_end_index], + kv_cache_keyboard["v"][:, max(0, local_end_index - max_attention_size):local_end_index], + ) + kv_cache_keyboard["global_end_index"].fill_(current_end) + kv_cache_keyboard["local_end_index"].fill_(local_end_index) + else: + attn = self.keyboard_attn_layer(q, k, v) + attn = rearrange(attn, 'B L H D -> B L (H D)') + attn, _ = self.proj_keyboard(attn) + hidden_states = hidden_states + attn + return hidden_states \ No newline at end of file diff --git a/fastvideo/models/dits/matrix_game/causal_model.py b/fastvideo/models/dits/matrix_game/causal_model.py new file mode 100644 index 000000000..fc859e5a8 --- /dev/null +++ b/fastvideo/models/dits/matrix_game/causal_model.py @@ -0,0 +1,1099 @@ +# SPDX-License-Identifier: Apache-2.0 + +import math +from typing import Any + +PYTEST_DONT_REWRITE = True + +import torch +import torch.nn as nn +import torch.distributed as dist + +from torch.nn.attention.flex_attention import create_block_mask, flex_attention +from torch.nn.attention.flex_attention import BlockMask +from fastvideo.forward_context import get_forward_context + +flex_attention = torch.compile( + flex_attention, dynamic=False, mode="max-autotune-no-cudagraphs") + +from fastvideo.attention import LocalAttention +from fastvideo.configs.models.dits.matrixgame import MatrixGameWanVideoConfig +from fastvideo.distributed.parallel_state import get_sp_world_size +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.rotary_embedding import (get_rotary_pos_embed) +from fastvideo.layers.visual_embedding import PatchEmbed, TimestepEmbedder, ModulateProjection +from fastvideo.logger import init_logger +from fastvideo.models.dits.base import BaseDiT +from fastvideo.models.dits.wanvideo import (WanT2VCrossAttention, + WanImageEmbedding) +from fastvideo.platforms import AttentionBackendEnum, current_platform + +from .action_module import ActionModule +from .model import MatrixGameCrossAttention + +logger = init_logger(__name__) + + +def causal_rope_apply(x, grid_sizes, freqs, start_frame=0): + n, c = x.size(2), x.size(3) // 2 + + # split freqs + freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1) + + # loop over samples + output = [] + f, h, w = grid_sizes.tolist() + + for i in range(len(x)): + seq_len = f * h * w + + # precompute multipliers + x_i = torch.view_as_complex(x[i, :seq_len].to(torch.float64).reshape( + seq_len, n, -1, 2)) + freqs_i = torch.cat([ + freqs[0][start_frame:start_frame + f].view(f, 1, 1, -1).expand(f, h, w, -1), + freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1), + freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1) + ], + dim=-1).reshape(seq_len, 1, -1) + + # apply rotary embedding + x_i = torch.view_as_real(x_i * freqs_i).flatten(2) + x_i = torch.cat([x_i, x[i, seq_len:]]) + + # append to collection + output.append(x_i) + return torch.stack(output).type_as(x) + + +def rope_params(max_seq_len, dim, theta=10000): + assert dim % 2 == 0 + freqs = torch.outer( + torch.arange(max_seq_len), + 1.0 / torch.pow(theta, + torch.arange(0, dim, 2).to(torch.float64).div(dim))) + freqs = torch.polar(torch.ones_like(freqs), freqs) + return freqs + + +class CausalMatrixGameTimeImageEmbedding(nn.Module): + def __init__( + self, + dim: int, + time_freq_dim: int, + image_embed_dim: int | None = None, + ): + 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") + + self.image_embedder = None + if image_embed_dim is not None: + self.image_embedder = WanImageEmbedding(image_embed_dim, dim) + + def forward( + self, + timestep: torch.Tensor, + encoder_hidden_states: torch.Tensor, + encoder_hidden_states_image: torch.Tensor | None = 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_image is not None: + assert self.image_embedder is not None + encoder_hidden_states_image = self.image_embedder( + encoder_hidden_states_image) + + return temb, timestep_proj, None, encoder_hidden_states_image + + +class CausalMatrixGameSelfAttention(nn.Module): + + def __init__(self, + dim: int, + num_heads: int, + local_attn_size: int = -1, + sink_size: int = 0, + qk_norm: bool = True, + eps: float = 1e-6) -> None: + assert dim % num_heads == 0 + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.local_attn_size = local_attn_size + self.sink_size = sink_size + self.qk_norm = qk_norm + self.eps = eps + self._freqs_cache = None + + self.attn = LocalAttention( + num_heads=num_heads, + head_size=self.head_dim, + dropout_rate=0, + softmax_scale=None, + causal=False, + supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN, + AttentionBackendEnum.TORCH_SDPA)) + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + freqs_cis: tuple[torch.Tensor, torch.Tensor], + block_mask: BlockMask, + kv_cache: dict | None = None, + current_start: int = 0, + cache_start: int | None = None, + grid_sizes: torch.Tensor | None = None + ): + if cache_start is None: + cache_start = current_start + + # Calculate start_frame for causal mode + if kv_cache is not None and grid_sizes is not None: + frame_seqlen = int(grid_sizes[1] * grid_sizes[2]) + start_frame = current_start // frame_seqlen + else: + start_frame = 0 + + if self._freqs_cache is None or self._freqs_cache.device != q.device: + d = self.head_dim + self._freqs_cache = torch.cat([ + rope_params(1024, d - 4 * (d // 6)), + rope_params(1024, 2 * (d // 6)), + rope_params(1024, 2 * (d // 6)) + ], dim=1).to(q.device) + + freqs = self._freqs_cache + + roped_query = causal_rope_apply(q, grid_sizes, freqs, start_frame=start_frame).type_as(v) + roped_key = causal_rope_apply(k, grid_sizes, freqs, start_frame=start_frame).type_as(v) + + if kv_cache is None: + padded_length = math.ceil(q.shape[1] / 128) * 128 - q.shape[1] + padded_roped_query = torch.cat( + [roped_query, + torch.zeros([q.shape[0], padded_length, q.shape[2], q.shape[3]], + device=q.device, dtype=v.dtype)], + dim=1 + ) + + padded_roped_key = torch.cat( + [roped_key, torch.zeros([k.shape[0], padded_length, k.shape[2], k.shape[3]], + device=k.device, dtype=v.dtype)], + dim=1 + ) + + padded_v = torch.cat( + [v, torch.zeros([v.shape[0], padded_length, v.shape[2], v.shape[3]], + device=v.device, dtype=v.dtype)], + dim=1 + ) + + x = flex_attention( + query=padded_roped_query.transpose(2, 1), + key=padded_roped_key.transpose(2, 1), + value=padded_v.transpose(2, 1), + block_mask=block_mask + )[:, :, :-padded_length].transpose(2, 1) + else: + # Calculate frame_seqlen correctly from grid_sizes (single frame token count) + if grid_sizes is not None: + frame_seqlen = int(grid_sizes[1] * grid_sizes[2]) + else: + # Fallback: assume q.shape[1] is single frame (shouldn't happen in causal mode) + frame_seqlen = q.shape[1] + logger.warning("grid_sizes not provided, using q.shape[1] as frame_seqlen") + + current_end = current_start + roped_query.shape[1] + sink_tokens = self.sink_size * frame_seqlen + + # Compute max_attention_size dynamically based on actual frame_seqlen + max_attention_size = (15 * frame_seqlen if self.local_attn_size == -1 + else self.local_attn_size * frame_seqlen) + + kv_cache_size = kv_cache["k"].shape[1] + num_new_tokens = roped_query.shape[1] + + if (current_end > kv_cache["global_end_index"].item()) and ( + num_new_tokens + kv_cache["local_end_index"].item() > kv_cache_size): + + num_evicted_tokens = num_new_tokens + kv_cache["local_end_index"].item() - kv_cache_size + num_rolled_tokens = kv_cache["local_end_index"].item() - num_evicted_tokens - sink_tokens + kv_cache["k"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \ + kv_cache["k"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone() + kv_cache["v"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \ + kv_cache["v"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone() + local_end_index = kv_cache["local_end_index"].item() + current_end - \ + kv_cache["global_end_index"].item() - num_evicted_tokens + local_start_index = local_end_index - num_new_tokens + kv_cache["k"][:, local_start_index:local_end_index] = roped_key + kv_cache["v"][:, local_start_index:local_end_index] = v + else: + local_end_index = kv_cache["local_end_index"].item() + current_end - kv_cache["global_end_index"].item() + local_start_index = local_end_index - num_new_tokens + + kv_cache["k"][:, local_start_index:local_end_index] = roped_key + kv_cache["v"][:, local_start_index:local_end_index] = v + + kv_start = max(0, local_end_index - max_attention_size) + k_for_attn = kv_cache["k"][:, kv_start:local_end_index] + v_for_attn = kv_cache["v"][:, kv_start:local_end_index] + + x = torch.nn.functional.scaled_dot_product_attention( + roped_query.transpose(1, 2), + k_for_attn.transpose(1, 2), + v_for_attn.transpose(1, 2), + ).transpose(1, 2) + + kv_cache["global_end_index"].fill_(current_end) + kv_cache["local_end_index"].fill_(local_end_index) + + return x + + +class CausalMatrixGameTransformerBlock(nn.Module): + + def __init__(self, + dim: int, + ffn_dim: int, + num_heads: int, + local_attn_size: int = -1, + sink_size: int = 0, + qk_norm: str = "rms_norm_across_heads", + cross_attn_norm: bool = False, + eps: float = 1e-6, + added_kv_proj_dim: int | None = None, + supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, + prefix: str = "", + action_config: dict | None = None, + block_idx: int = 0): + super().__init__() + action_config = action_config or {} + + 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 = CausalMatrixGameSelfAttention( + dim, + num_heads, + local_attn_size=local_attn_size, + sink_size=sink_size, + qk_norm=qk_norm, + eps=eps) + self.hidden_dim = dim + self.num_attention_heads = num_heads + self.local_attn_size = local_attn_size + 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": + 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) + + if added_kv_proj_dim is not None: + self.attn2 = MatrixGameCrossAttention(dim, + num_heads, + qk_norm=qk_norm, + eps=eps) + else: + self.attn2 = WanT2VCrossAttention(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) + + + if len(action_config) > 0 and block_idx in action_config.get("blocks", []): + self.action_model = ActionModule( + heads_num=action_config["heads_num"], + hidden_size=action_config["hidden_size"], + img_hidden_size=action_config["img_hidden_size"], + mouse_hidden_dim=action_config.get("mouse_hidden_dim", 0), + keyboard_hidden_dim=action_config["keyboard_hidden_dim"], + mouse_dim_in=action_config.get("mouse_dim_in", 0), + keyboard_dim_in=action_config["keyboard_dim_in"], + enable_mouse=action_config["enable_mouse"], + enable_keyboard=action_config["enable_keyboard"], + windows_size=action_config["windows_size"], + rope_theta=action_config["rope_theta"], + rope_dim_list=action_config["rope_dim_list"], + mouse_qk_dim_list=action_config.get("mouse_qk_dim_list", [8, 28, 28]), + patch_size=action_config["patch_size"], + local_attn_size=local_attn_size, + qk_norm=action_config["qk_norm"], + qkv_bias=action_config["qkv_bias"], + vae_time_compression_ratio=action_config["vae_time_compression_ratio"], + ) + else: + self.action_model = None + + 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) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor, + freqs_cis: tuple[torch.Tensor, torch.Tensor], + block_mask: BlockMask, + grid_sizes: torch.Tensor | None = None, + mouse_cond: torch.Tensor | None = None, + keyboard_cond: torch.Tensor | None = None, + block_mask_mouse: BlockMask | None = None, + block_mask_keyboard: BlockMask | None = None, + num_frame_per_block: int = 1, + use_rope_keyboard: bool = True, + kv_cache: dict | None = None, + kv_cache_mouse: dict | None = None, + kv_cache_keyboard: dict | None = None, + crossattn_cache: dict | None = None, + current_start: int = 0, + cache_start: int | None = None, + ) -> torch.Tensor: + if hidden_states.dim() == 4: + hidden_states = hidden_states.squeeze(1) + + assert temb.ndim == 4 + num_frames = temb.shape[1] + frame_seqlen = hidden_states.shape[1] // num_frames + bs, seq_length, _ = hidden_states.shape + orig_dtype = hidden_states.dtype + # hidden_states = hidden_states.float() + + e = self.scale_shift_table + temb + assert e.shape == (bs, num_frames, 6, self.hidden_dim) + shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = e.chunk(6, dim=2) + + norm_hidden_states = (self.norm1(hidden_states).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * + (1 + scale_msa) + shift_msa).flatten(1, 2) + 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.forward_native(query) + if self.norm_k is not None: + key = self.norm_k.forward_native(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)) + + attn_output = self.attn1(query, key, value, freqs_cis, block_mask, kv_cache, current_start, cache_start, grid_sizes) + attn_output = attn_output.flatten(2) + attn_output, _ = self.to_out(attn_output) + attn_output = attn_output.squeeze(1) + + null_shift = null_scale = torch.tensor([0], device=hidden_states.device) + norm_hidden_states, hidden_states = self.self_attn_residual_norm( + hidden_states, attn_output, gate_msa, null_shift, null_scale) + + # norm3 weights are loaded into self_attn_residual_norm.norm + attn_output = self.attn2(norm_hidden_states, + context=encoder_hidden_states, + context_lens=None, + crossattn_cache=crossattn_cache) + # residual connection + hidden_states = hidden_states + attn_output + + if self.action_model is not None: + if mouse_cond is not None or keyboard_cond is not None: + start_frame = current_start // frame_seqlen if grid_sizes is not None else 0 + + hidden_states = self.action_model( + hidden_states, + grid_sizes[0], grid_sizes[1], grid_sizes[2], + mouse_cond, keyboard_cond, + block_mask_mouse, block_mask_keyboard, + is_causal=True, + kv_cache_mouse=kv_cache_mouse, + kv_cache_keyboard=kv_cache_keyboard, + start_frame=start_frame, + use_rope_keyboard=use_rope_keyboard, + num_frame_per_block=num_frame_per_block + ) + + + ff_output = self.ffn( + (self.norm1(hidden_states).unflatten(dim=1, sizes=(num_frames, frame_seqlen)) * + + (1 + c_scale_msa) + c_shift_msa).flatten(1, 2) + ) + hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa) + + # Convert back to original dtype before return + return hidden_states.to(orig_dtype) + + +_DEFAULT_MATRIXGAME_CONFIG = MatrixGameWanVideoConfig() + + +class CausalMatrixGameWanModel(BaseDiT): + supports_action_input = True + + _fsdp_shard_conditions = _DEFAULT_MATRIXGAME_CONFIG._fsdp_shard_conditions + _compile_conditions = _DEFAULT_MATRIXGAME_CONFIG._compile_conditions + _supported_attention_backends = _DEFAULT_MATRIXGAME_CONFIG._supported_attention_backends + param_names_mapping = _DEFAULT_MATRIXGAME_CONFIG.param_names_mapping + reverse_param_names_mapping = _DEFAULT_MATRIXGAME_CONFIG.reverse_param_names_mapping + lora_param_names_mapping = _DEFAULT_MATRIXGAME_CONFIG.lora_param_names_mapping + + def __init__(self, + config: MatrixGameWanVideoConfig, + 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.attention_head_dim = config.attention_head_dim + self.in_channels = config.in_channels + self.out_channels = config.out_channels + self.num_channels_latents = config.num_channels_latents + self.patch_size = config.patch_size + + arch_cfg = getattr(config, "arch_config", None) + self.local_attn_size = getattr(arch_cfg, 'local_attn_size', getattr(config, 'local_attn_size', -1)) if arch_cfg else getattr(config, 'local_attn_size', -1) + self.sink_size = getattr(arch_cfg, 'sink_size', getattr(config, 'sink_size', 0)) if arch_cfg else getattr(config, 'sink_size', 0) + + # 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 = CausalMatrixGameTimeImageEmbedding( + dim=inner_dim, + time_freq_dim=config.freq_dim, + image_embed_dim=config.image_dim, + ) + + # 2.1. Get action config + arch_cfg = getattr(config, "arch_config", None) + self.action_config = getattr(arch_cfg, "action_config", {}) if arch_cfg is not None else {} + + # 3. Transformer blocks + self.blocks = nn.ModuleList([ + CausalMatrixGameTransformerBlock( + inner_dim, + config.ffn_dim, + config.num_attention_heads, + self.local_attn_size, + self.sink_size, + config.qk_norm, + config.cross_attn_norm, + config.eps, + config.added_kv_proj_dim, + self._supported_attention_backends, + prefix=f"{getattr(config, 'prefix', 'Wan')}.blocks.{i}", + action_config=self.action_config, + block_idx=i) + 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) + 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 + + self.block_mask = None + self.block_mask_keyboard = None + self.block_mask_mouse = None + self.use_rope_keyboard = True + + self.num_frame_per_block = getattr(arch_cfg, 'num_frames_per_block', getattr(config, 'num_frame_per_block', 1)) if arch_cfg else getattr(config, 'num_frame_per_block', 1) + + self.__post_init__() + + @staticmethod + def _prepare_blockwise_causal_attn_mask( + device: torch.device | str, num_frames: int = 9, + frame_seqlen: int = 880, num_frame_per_block: int = 1, local_attn_size: int = -1 + ) -> BlockMask: + total_length = num_frames * frame_seqlen + padded_length = math.ceil(total_length / 128) * 128 - total_length + + ends = torch.zeros(total_length + padded_length, + device=device, dtype=torch.long) + + frame_indices = torch.arange( + start=0, + end=total_length, + step=frame_seqlen * num_frame_per_block, + device=device + ) + + for tmp in frame_indices: + ends[tmp:tmp + frame_seqlen * num_frame_per_block] = tmp + \ + frame_seqlen * num_frame_per_block + + def attention_mask(b, h, q_idx, kv_idx): + if local_attn_size == -1: + return (kv_idx < ends[q_idx]) | (q_idx == kv_idx) + else: + return ((kv_idx < ends[q_idx]) & (kv_idx >= (ends[q_idx] - local_attn_size * frame_seqlen))) | (q_idx == kv_idx) + + block_mask = create_block_mask(attention_mask, B=None, H=None, Q_LEN=total_length + padded_length, + KV_LEN=total_length + padded_length, _compile=False, device=device) + + if not dist.is_initialized() or dist.get_rank() == 0: + print(f" cache a block wise causal mask with block size of {num_frame_per_block} frames") + + return block_mask + + @staticmethod + def _prepare_blockwise_causal_attn_mask_keyboard( + device: torch.device | str, num_frames: int = 9, + frame_seqlen: int = 880, num_frame_per_block: int = 1, local_attn_size: int = -1 + ) -> BlockMask: + total_length2 = num_frames * frame_seqlen + padded_length2 = math.ceil(total_length2 / 32) * 32 - total_length2 + padded_length_kv2 = math.ceil(num_frames / 32) * 32 - num_frames + ends2 = torch.zeros(total_length2 + padded_length2, + device=device, dtype=torch.long) + + frame_indices2 = torch.arange( + start=0, + end=total_length2, + step=frame_seqlen * num_frame_per_block, + device=device + ) + cnt = num_frame_per_block + for tmp in frame_indices2: + ends2[tmp:tmp + frame_seqlen * num_frame_per_block] = cnt + cnt += num_frame_per_block + + def attention_mask2(b, h, q_idx, kv_idx): + if local_attn_size == -1: + return (kv_idx < ends2[q_idx]) | (q_idx == kv_idx) + else: + return ((kv_idx < ends2[q_idx]) & (kv_idx >= (ends2[q_idx] - local_attn_size))) | (q_idx == kv_idx) + + block_mask2 = create_block_mask(attention_mask2, B=None, H=None, Q_LEN=total_length2 + padded_length2, + KV_LEN=num_frames + padded_length_kv2, _compile=False, device=device) + + if not dist.is_initialized() or dist.get_rank() == 0: + print(f" cache a block wise causal mask for keyboard with block size of {num_frame_per_block} frames") + + return block_mask2 + + @staticmethod + def _prepare_blockwise_causal_attn_mask_action( + device: torch.device | str, num_frames: int = 9, + frame_seqlen: int = 1, num_frame_per_block: int = 1, local_attn_size: int = -1 + ) -> BlockMask: + total_length2 = num_frames * frame_seqlen + padded_length2 = math.ceil(total_length2 / 32) * 32 - total_length2 + padded_length_kv2 = math.ceil(num_frames / 32) * 32 - num_frames + ends2 = torch.zeros(total_length2 + padded_length2, + device=device, dtype=torch.long) + + frame_indices2 = torch.arange( + start=0, + end=total_length2, + step=frame_seqlen * num_frame_per_block, + device=device + ) + cnt = num_frame_per_block + for tmp in frame_indices2: + ends2[tmp:tmp + frame_seqlen * num_frame_per_block] = cnt + cnt += num_frame_per_block + + def attention_mask2(b, h, q_idx, kv_idx): + if local_attn_size == -1: + return (kv_idx < ends2[q_idx]) | (q_idx == kv_idx) + else: + return ((kv_idx < ends2[q_idx]) & (kv_idx >= (ends2[q_idx] - local_attn_size))) | (q_idx == kv_idx) + + block_mask2 = create_block_mask(attention_mask2, B=None, H=None, Q_LEN=total_length2 + padded_length2, + KV_LEN=num_frames + padded_length_kv2, _compile=False, device=device) + + if not dist.is_initialized() or dist.get_rank() == 0: + print(f" cache a block wise causal mask for action with block size of {num_frame_per_block} frames") + + return block_mask2 + + def _forward_inference( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | list[torch.Tensor], + timestep: torch.LongTensor, + encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None, + mouse_cond: torch.Tensor | None = None, + keyboard_cond: torch.Tensor | None = None, + kv_cache: dict = None, + kv_cache_mouse: dict = None, + kv_cache_keyboard: dict = None, + crossattn_cache: dict = None, + current_start: int = 0, + cache_start: int = 0, + start_frame: int = 0, + **kwargs + ) -> torch.Tensor: + # Extract num_frame_per_block from kwargs + effective_num_frame_per_block = kwargs.pop('num_frame_per_block', self.num_frame_per_block) + + if mouse_cond is not None or keyboard_cond is not None: + assert self.action_config is not None and len(self.action_config) > 0 + + if encoder_hidden_states is not None and not isinstance(encoder_hidden_states, torch.Tensor): + encoder_hidden_states = encoder_hidden_states[0] + if isinstance(encoder_hidden_states_image, list) and len(encoder_hidden_states_image) > 0: + encoder_hidden_states_image = encoder_hidden_states_image[0] + else: + encoder_hidden_states_image = None + + ctx = get_forward_context() + batch = getattr(ctx, "forward_batch", None) + + if batch is not None and getattr(batch, "image_latent", None) is not None: + image_latent = batch.image_latent + if isinstance(image_latent, torch.Tensor) and image_latent.ndim == 5 and hidden_states.ndim == 5: + _, _, num_frames, _, _ = hidden_states.shape + start = start_frame + end = start + num_frames + cond = image_latent + if cond.shape[2] >= end: + cond = cond[:, :, start:end] + elif cond.shape[2] > start: + cond = cond[:, :, start:] + pad_frames = num_frames - cond.shape[2] + if pad_frames > 0: + pad = torch.zeros( + cond.shape[0], + cond.shape[1], + pad_frames, + cond.shape[3], + cond.shape[4], + device=cond.device, + dtype=cond.dtype, + ) + cond = torch.cat([cond, pad], dim=2) + else: + pad = torch.zeros( + cond.shape[0], + cond.shape[1], + num_frames, + cond.shape[3], + cond.shape[4], + device=cond.device, + dtype=cond.dtype, + ) + cond = pad + hidden_states = torch.cat( + [hidden_states, cond.to(dtype=hidden_states.dtype)], dim=1 + ) + + 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 + + d = self.hidden_size // self.num_attention_heads + rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] + freqs_cos, freqs_sin = get_rotary_pos_embed( + (post_patch_num_frames * get_sp_world_size(), post_patch_height, post_patch_width), + self.hidden_size, + self.num_attention_heads, + rope_dim_list, + dtype=torch.float32 if current_platform.is_mps() else torch.float64, + rope_theta=10000, + start_frame=start_frame + ) + freqs_cos = freqs_cos.to(hidden_states.device) + freqs_sin = freqs_sin.to(hidden_states.device) + freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None + + hidden_states = self.patch_embedding(hidden_states) + grid_sizes = torch.tensor([post_patch_num_frames, post_patch_height, post_patch_width], + device=hidden_states.device) + hidden_states = hidden_states.flatten(2).transpose(1, 2) + + if timestep.dim() == 2: + timestep = timestep.flatten() + + temb, timestep_proj, encoder_hidden_states_text, encoder_hidden_states_image = self.condition_embedder( + timestep, encoder_hidden_states, encoder_hidden_states_image) + timestep_proj = timestep_proj.unflatten(1, (6, self.hidden_size)).unflatten(dim=0, sizes=(batch_size, post_patch_num_frames)) + + if encoder_hidden_states_text is not None: + if isinstance(encoder_hidden_states_text, list): + encoder_hidden_states_text = encoder_hidden_states_text[0] + elif encoder_hidden_states_text.ndim == 2: + encoder_hidden_states_text = encoder_hidden_states_text.unsqueeze(0) + + if encoder_hidden_states_image is not None: + if encoder_hidden_states_text is not None: + encoder_hidden_states = torch.concat( + [encoder_hidden_states_image, encoder_hidden_states_text], dim=1) + else: + encoder_hidden_states = encoder_hidden_states_image + + if self.block_mask is None: + self.block_mask = self._prepare_blockwise_causal_attn_mask( + device=hidden_states.device, + num_frames=num_frames, + frame_seqlen=post_patch_height * post_patch_width, + num_frame_per_block=self.num_frame_per_block, + local_attn_size=self.local_attn_size + ) + if self.block_mask_keyboard is None: + if not self.use_rope_keyboard: + self.block_mask_keyboard = self._prepare_blockwise_causal_attn_mask_keyboard( + device=hidden_states.device, + num_frames=num_frames, + frame_seqlen=post_patch_height * post_patch_width, + num_frame_per_block=self.num_frame_per_block, + local_attn_size=self.local_attn_size + ) + else: + self.block_mask_keyboard = self._prepare_blockwise_causal_attn_mask_action( + device=hidden_states.device, + num_frames=num_frames, + frame_seqlen=1, + num_frame_per_block=self.num_frame_per_block, + local_attn_size=self.local_attn_size + ) + if self.block_mask_mouse is None: + self.block_mask_mouse = self._prepare_blockwise_causal_attn_mask_action( + device=hidden_states.device, + num_frames=num_frames, + frame_seqlen=1, + num_frame_per_block=self.num_frame_per_block, + local_attn_size=self.local_attn_size + ) + if kv_cache is None: + kv_cache = [None] * len(self.blocks) + if kv_cache_mouse is None: + kv_cache_mouse = [None] * len(self.blocks) + if kv_cache_keyboard is None: + kv_cache_keyboard = [None] * len(self.blocks) + if crossattn_cache is None: + crossattn_cache = [None] * len(self.blocks) + + for block_index, block in enumerate(self.blocks): + if torch.is_grad_enabled() and self.gradient_checkpointing: + kwargs.update({ + "kv_cache": kv_cache[block_index], + "kv_cache_mouse": kv_cache_mouse[block_index] if kv_cache_mouse else None, + "kv_cache_keyboard": kv_cache_keyboard[block_index] if kv_cache_keyboard else None, + "crossattn_cache": crossattn_cache[block_index] if crossattn_cache else None, + "current_start": current_start, + "cache_start": cache_start, + "num_frame_per_block": effective_num_frame_per_block, + }) + hidden_states = self._gradient_checkpointing_func( + block, hidden_states, **kwargs) + else: + kwargs.update({ + "kv_cache": kv_cache[block_index], + "kv_cache_mouse": kv_cache_mouse[block_index] if kv_cache_mouse else None, + "kv_cache_keyboard": kv_cache_keyboard[block_index] if kv_cache_keyboard else None, + "crossattn_cache": crossattn_cache[block_index] if crossattn_cache else None, + "current_start": current_start, + "cache_start": cache_start, + }) + hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, freqs_cis, + block_mask=self.block_mask, + grid_sizes=grid_sizes, + mouse_cond=mouse_cond, + keyboard_cond=keyboard_cond, + block_mask_mouse=self.block_mask_mouse, + block_mask_keyboard=self.block_mask_keyboard, + num_frame_per_block=effective_num_frame_per_block, + use_rope_keyboard=self.use_rope_keyboard, + **kwargs) + + temb = temb.unflatten(dim=0, sizes=(batch_size, post_patch_num_frames)).unsqueeze(2) + shift, scale = (self.scale_shift_table.unsqueeze(1) + temb).chunk(2, dim=2) + hidden_states = self.norm_out(hidden_states, shift, scale) + hidden_states = self.proj_out(hidden_states) + + hidden_states = self.unpatchify(hidden_states, grid_sizes) + + return hidden_states + + + def _forward_train( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor | list[torch.Tensor], + timestep: torch.LongTensor, + encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None, + mouse_cond: torch.Tensor | None = None, + keyboard_cond: torch.Tensor | None = None, + **kwargs + ) -> torch.Tensor: + if mouse_cond is not None or keyboard_cond is not None: + assert self.action_config is not None and len(self.action_config) > 0 + + if encoder_hidden_states is not None and not isinstance(encoder_hidden_states, torch.Tensor): + encoder_hidden_states = encoder_hidden_states[0] + if isinstance(encoder_hidden_states_image, list) and len(encoder_hidden_states_image) > 0: + encoder_hidden_states_image = encoder_hidden_states_image[0] + else: + encoder_hidden_states_image = None + + ctx = get_forward_context() + batch = getattr(ctx, "forward_batch", None) + if batch is not None and getattr(batch, "image_latent", None) is not None: + image_latent = batch.image_latent + if isinstance(image_latent, torch.Tensor) and image_latent.ndim == 5 and hidden_states.ndim == 5: + _, _, num_frames, _, _ = hidden_states.shape + cond = image_latent + if cond.shape[2] >= num_frames: + cond = cond[:, :, :num_frames] + else: + pad_frames = num_frames - cond.shape[2] + pad = torch.zeros( + cond.shape[0], + cond.shape[1], + pad_frames, + cond.shape[3], + cond.shape[4], + device=cond.device, + dtype=cond.dtype, + ) + cond = torch.cat([cond, pad], dim=2) + hidden_states = torch.cat( + [hidden_states, cond.to(dtype=hidden_states.dtype)], dim=1 + ) + + 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 + + d = self.hidden_size // self.num_attention_heads + rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] + freqs_cos, freqs_sin = get_rotary_pos_embed( + (post_patch_num_frames * get_sp_world_size(), post_patch_height, post_patch_width), + self.hidden_size, + self.num_attention_heads, + rope_dim_list, + dtype=torch.float32 if current_platform.is_mps() else torch.float64, + rope_theta=10000 + ) + freqs_cos = freqs_cos.to(hidden_states.device) + freqs_sin = freqs_sin.to(hidden_states.device) + freqs_cis = (freqs_cos, freqs_sin) if freqs_cos is not None else None + + if self.block_mask is None: + self.block_mask = self._prepare_blockwise_causal_attn_mask( + device=hidden_states.device, + num_frames=num_frames, + frame_seqlen=post_patch_height * post_patch_width, + num_frame_per_block=self.num_frame_per_block, + local_attn_size=self.local_attn_size + ) + if self.block_mask_keyboard is None: + if not self.use_rope_keyboard: + self.block_mask_keyboard = self._prepare_blockwise_causal_attn_mask_keyboard( + device=hidden_states.device, + num_frames=num_frames, + frame_seqlen=post_patch_height * post_patch_width, + num_frame_per_block=self.num_frame_per_block, + local_attn_size=self.local_attn_size + ) + else: + self.block_mask_keyboard = self._prepare_blockwise_causal_attn_mask_action( + device=hidden_states.device, + num_frames=num_frames, + frame_seqlen=1, + num_frame_per_block=self.num_frame_per_block, + local_attn_size=self.local_attn_size + ) + if self.block_mask_mouse is None: + self.block_mask_mouse = self._prepare_blockwise_causal_attn_mask_action( + device=hidden_states.device, + num_frames=num_frames, + frame_seqlen=1, + num_frame_per_block=self.num_frame_per_block, + local_attn_size=self.local_attn_size + ) + + hidden_states = self.patch_embedding(hidden_states) + grid_sizes = torch.tensor([post_patch_num_frames, post_patch_height, post_patch_width], + device=hidden_states.device) + hidden_states = hidden_states.flatten(2).transpose(1, 2) + + if timestep.dim() == 2: + timestep = timestep.flatten() + + temb, timestep_proj, encoder_hidden_states_text, encoder_hidden_states_image = self.condition_embedder( + timestep, encoder_hidden_states, encoder_hidden_states_image) + timestep_proj = timestep_proj.unflatten(1, (6, self.hidden_size)).unflatten(dim=0, sizes=(batch_size, post_patch_num_frames)) + + if encoder_hidden_states_text is not None: + if isinstance(encoder_hidden_states_text, list): + encoder_hidden_states_text = encoder_hidden_states_text[0] + elif encoder_hidden_states_text.ndim == 2: + encoder_hidden_states_text = encoder_hidden_states_text.unsqueeze(0) + + if encoder_hidden_states_image is not None: + if encoder_hidden_states_text is not None: + encoder_hidden_states = torch.concat( + [encoder_hidden_states_image, encoder_hidden_states_text], dim=1) + else: + encoder_hidden_states = encoder_hidden_states_image + else: + encoder_hidden_states = encoder_hidden_states_text + + kwargs = dict( + e=timestep_proj, + grid_sizes=grid_sizes, + freqs_cis=freqs_cis, + encoder_hidden_states=encoder_hidden_states, + mouse_cond=mouse_cond, + keyboard_cond=keyboard_cond, + block_mask=self.block_mask, + block_mask_mouse=self.block_mask_mouse, + block_mask_keyboard=self.block_mask_keyboard, + use_rope_keyboard=self.use_rope_keyboard, + num_frame_per_block=self.num_frame_per_block + ) + + # Initialize kv_cache, kv_cache_mouse, kv_cache_keyboard, crossattn_cache if not provided + kv_cache = kwargs.pop("kv_cache", [None] * len(self.blocks)) + kv_cache_mouse = kwargs.pop("kv_cache_mouse", [None] * len(self.blocks)) + kv_cache_keyboard = kwargs.pop("kv_cache_keyboard", [None] * len(self.blocks)) + crossattn_cache = kwargs.pop("crossattn_cache", [None] * len(self.blocks)) + + for i, block in enumerate(self.blocks): + block_crossattn_cache = None + if crossattn_cache is not None: + if crossattn_cache[i] is None: + crossattn_cache[i] = {"is_init": False} + block_crossattn_cache = crossattn_cache[i] + + if torch.is_grad_enabled() and self.gradient_checkpointing: + hidden_states = self._gradient_checkpointing_func( + block, hidden_states, encoder_hidden_states, timestep_proj, freqs_cis, + kv_cache=kv_cache[i], + kv_cache_mouse=kv_cache_mouse[i], + kv_cache_keyboard=kv_cache_keyboard[i], + crossattn_cache=block_crossattn_cache, + block_mask=self.block_mask, + grid_sizes=grid_sizes, + mouse_cond=mouse_cond, + keyboard_cond=keyboard_cond, + block_mask_mouse=self.block_mask_mouse, + block_mask_keyboard=self.block_mask_keyboard, + num_frame_per_block=self.num_frame_per_block, + use_rope_keyboard=self.use_rope_keyboard) + else: + hidden_states = block( + hidden_states, + encoder_hidden_states, + timestep_proj, + freqs_cis, + kv_cache=kv_cache[i], + kv_cache_mouse=kv_cache_mouse[i], + kv_cache_keyboard=kv_cache_keyboard[i], + crossattn_cache=block_crossattn_cache, + block_mask=self.block_mask, + grid_sizes=grid_sizes, + mouse_cond=mouse_cond, + keyboard_cond=keyboard_cond, + block_mask_mouse=self.block_mask_mouse, + block_mask_keyboard=self.block_mask_keyboard, + num_frame_per_block=self.num_frame_per_block, + use_rope_keyboard=self.use_rope_keyboard) + + temb = temb.unflatten(dim=0, sizes=(batch_size, post_patch_num_frames)).unsqueeze(2) + shift, scale = (self.scale_shift_table.unsqueeze(1) + temb).chunk(2, dim=2) + hidden_states = self.norm_out(hidden_states, shift, scale) + hidden_states = self.proj_out(hidden_states) + + hidden_states = self.unpatchify(hidden_states, grid_sizes) + + return hidden_states + + def forward(self, *args, **kwargs): + if kwargs.get('kv_cache') is not None: + return self._forward_inference(*args, **kwargs) + else: + return self._forward_train(*args, **kwargs) + + def unpatchify(self, x, grid_sizes): + c = self.proj_out.out_features // math.prod(self.patch_size) + p_t, p_h, p_w = self.patch_size + f, h, w = grid_sizes.tolist() + + x = x[:, :f * h * w].view(-1, f, h, w, p_t, p_h, p_w, c) + x = x.permute(0, 7, 1, 4, 2, 5, 3, 6) + x = x.reshape(-1, c, f * p_t, h * p_h, w * p_w) + return x + + def init_weights(self): + for m in self.modules(): + if isinstance(m, nn.Linear): + nn.init.xavier_uniform_(m.weight) + if m.bias is not None: + nn.init.zeros_(m.bias) + + nn.init.xavier_uniform_(self.patch_embedding.proj.weight.flatten(1)) + for m in self.condition_embedder.time_embedder.mlp.modules(): + if isinstance(m, nn.Linear): + nn.init.normal_(m.weight, std=.02) + + nn.init.zeros_(self.proj_out.weight) + + for block in self.blocks: + if block.action_model is not None: + try: + nn.init.zeros_(block.action_model.proj_mouse.weight) + if block.action_model.proj_mouse.bias is not None: + nn.init.zeros_(block.action_model.proj_mouse.bias) + nn.init.zeros_(block.action_model.proj_keyboard.weight) + if block.action_model.proj_keyboard.bias is not None: + nn.init.zeros_(block.action_model.proj_keyboard.bias) + except AttributeError: + pass diff --git a/fastvideo/models/dits/matrix_game/model.py b/fastvideo/models/dits/matrix_game/model.py new file mode 100644 index 000000000..30891164d --- /dev/null +++ b/fastvideo/models/dits/matrix_game/model.py @@ -0,0 +1,469 @@ +# 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.matrixgame import MatrixGameWanVideoConfig +from fastvideo.distributed.parallel_state import get_sp_world_size +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.rotary_embedding import (_apply_rotary_emb, + get_rotary_pos_embed) +from fastvideo.layers.visual_embedding import PatchEmbed, TimestepEmbedder, ModulateProjection +from fastvideo.logger import init_logger +from fastvideo.models.dits.base import BaseDiT +from fastvideo.models.dits.wanvideo import (WanSelfAttention, + WanI2VCrossAttention, + WanT2VCrossAttention, + WanImageEmbedding) +from fastvideo.platforms import AttentionBackendEnum, current_platform + +# Import ActionModule +from .action_module import ActionModule + +logger = init_logger(__name__) + + +class MatrixGameTimeImageEmbedding(nn.Module): + def __init__( + self, + dim: int, + time_freq_dim: int, + image_embed_dim: int | None = None, + ): + 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") + + self.image_embedder = None + if image_embed_dim is not None: + self.image_embedder = WanImageEmbedding(image_embed_dim, dim) + + def forward( + self, + timestep: torch.Tensor, + encoder_hidden_states: torch.Tensor, # Kept for interface compatibility + encoder_hidden_states_image: torch.Tensor | None = None, + timestep_seq_len: int | None = None, + ): + temb = self.time_embedder(timestep, timestep_seq_len) + timestep_proj = self.time_modulation(temb) + + # MatrixGame does not use text embeddings, so we ignore encoder_hidden_states + # and return None for the text embedding part + + if encoder_hidden_states_image is not None: + assert self.image_embedder is not None + encoder_hidden_states_image = self.image_embedder( + encoder_hidden_states_image) + + return temb, timestep_proj, None, encoder_hidden_states_image + + +class MatrixGameCrossAttention(WanSelfAttention): + def forward(self, x, context, context_lens=None, crossattn_cache=None): + r""" + Args: + x(Tensor): Shape [B, L1, C] + context(Tensor): Shape [B, L2, C] - typically 257 image tokens + context_lens(Tensor): Shape [B] + crossattn_cache(dict): Optional cache for k/v during inference + """ + b, n, d = x.size(0), self.num_heads, self.head_dim + + # compute query, key, value + q = self.norm_q(self.to_q(x)[0]).view(b, -1, n, d) + + if crossattn_cache is not None: + if not crossattn_cache["is_init"]: + crossattn_cache["is_init"] = True + k = self.norm_k(self.to_k(context)[0]).view(b, -1, n, d) + v = self.to_v(context)[0].view(b, -1, n, d) + crossattn_cache["k"] = k + crossattn_cache["v"] = v + else: + k = crossattn_cache["k"] + v = crossattn_cache["v"] + else: + k = self.norm_k(self.to_k(context)[0]).view(b, -1, n, d) + v = self.to_v(context)[0].view(b, -1, n, d) + + # compute attention + x = self.attn(q, k, v) + + # output + x = x.flatten(2) + x, _ = self.to_out(x) + return x + + +class MatrixGameTransformerBlock(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, + added_kv_proj_dim: int | None = None, + supported_attention_backends: tuple[AttentionBackendEnum, ...] + | None = None, + prefix: str = "", + action_config: dict | None = None): + super().__init__() + action_config = action_config or {} + + # 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 + if added_kv_proj_dim is not None: + # I2V + self.attn2 = WanI2VCrossAttention(dim, + num_heads, + qk_norm=qk_norm, + eps=eps) + else: + # T2V + self.attn2 = WanT2VCrossAttention(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 + self.use_action_module = len(action_config) > 0 + if self.use_action_module: + self.action_model = ActionModule(**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) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor, + freqs_cis: tuple[torch.Tensor, torch.Tensor], + # Action Module specific args + grid_sizes: torch.Tensor | None = None, + mouse_cond: torch.Tensor | None = None, + keyboard_cond: 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 + cos, sin = freqs_cis + query, key = _apply_rotary_emb(query, cos, sin, + is_neox_style=False), _apply_rotary_emb( + key, cos, sin, is_neox_style=False) + + attn_output, _ = self.attn1(query, key, value) + attn_output = attn_output.flatten(2) + attn_output, _ = self.to_out(attn_output) + attn_output = attn_output.squeeze(1) + + null_shift = null_scale = torch.tensor([0], device=hidden_states.device) + norm_hidden_states, hidden_states = self.self_attn_residual_norm( + hidden_states, attn_output, gate_msa, null_shift, null_scale) + norm_hidden_states, hidden_states = norm_hidden_states.to( + orig_dtype), hidden_states.to(orig_dtype) + + # 2. Cross-attention + attn_output = self.attn2(norm_hidden_states, + context=encoder_hidden_states, + context_lens=None) + norm_hidden_states, hidden_states = self.cross_attn_residual_norm( + hidden_states, attn_output, 1, c_shift_msa, c_scale_msa) + norm_hidden_states, hidden_states = norm_hidden_states.to( + orig_dtype), hidden_states.to(orig_dtype) + + # ================= Action Module ================= + if self.action_model is not None: + if mouse_cond is not None or keyboard_cond is not None: + # grid_sizes is expected to be [F, H, W] + # ActionModule implementation takes hidden_states directly + hidden_states = self.action_model( + hidden_states, + grid_sizes[0], grid_sizes[1], grid_sizes[2], + mouse_cond, keyboard_cond, + num_frame_per_block=grid_sizes[0], + ) + # ================================================= + + # 3. Feed-forward + ff_output = self.ffn(norm_hidden_states) + hidden_states = self.mlp_residual(hidden_states, ff_output, c_gate_msa) + hidden_states = hidden_states.to(orig_dtype) + + return hidden_states + +_DEFAULT_MATRIXGAME_CONFIG = MatrixGameWanVideoConfig() + + +class MatrixGameWanModel(BaseDiT): + # Marker for action input support (Matrix-Game) + supports_action_input = True + + _fsdp_shard_conditions = _DEFAULT_MATRIXGAME_CONFIG._fsdp_shard_conditions + _compile_conditions = _DEFAULT_MATRIXGAME_CONFIG._compile_conditions + _supported_attention_backends = _DEFAULT_MATRIXGAME_CONFIG._supported_attention_backends + param_names_mapping = _DEFAULT_MATRIXGAME_CONFIG.param_names_mapping + reverse_param_names_mapping = _DEFAULT_MATRIXGAME_CONFIG.reverse_param_names_mapping + lora_param_names_mapping = _DEFAULT_MATRIXGAME_CONFIG.lora_param_names_mapping + + def __init__(self, + config: MatrixGameWanVideoConfig, + 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.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 = MatrixGameTimeImageEmbedding( + dim=inner_dim, + time_freq_dim=config.freq_dim, + image_embed_dim=config.image_dim, + ) + + # 2.1. Get action config + self.action_config = getattr(config, 'action_config', {}) + + # 3. Transformer blocks + self.blocks = nn.ModuleList([ + MatrixGameTransformerBlock( + inner_dim, + config.ffn_dim, + config.num_attention_heads, + config.qk_norm, + config.cross_attn_norm, + config.eps, + config.added_kv_proj_dim, + self._supported_attention_backends, + prefix=f"{getattr(config, 'prefix', 'Wan')}.blocks.{i}", + action_config=self.action_config) + 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, + encoder_hidden_states_image: torch.Tensor + | list[torch.Tensor] | None = None, + # Action inputs + mouse_cond: torch.Tensor | None = None, + keyboard_cond: 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] + if isinstance(encoder_hidden_states_image, + list) and len(encoder_hidden_states_image) > 0: + encoder_hidden_states_image = encoder_hidden_states_image[0] + else: + encoder_hidden_states_image = None + + 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 + + # Get rotary embeddings + d = self.hidden_size // self.num_attention_heads + rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] + freqs_cos, freqs_sin = get_rotary_pos_embed( + (post_patch_num_frames * get_sp_world_size(), post_patch_height, + post_patch_width), + self.hidden_size, + self.num_attention_heads, + rope_dim_list, + dtype=torch.float32 if current_platform.is_mps() else torch.float64, + rope_theta=10000) + freqs_cos = freqs_cos.to(hidden_states.device) + freqs_sin = freqs_sin.to(hidden_states.device) + freqs_cis = (freqs_cos.float(), + freqs_sin.float()) if freqs_cos is not None else None + + hidden_states = self.patch_embedding(hidden_states) + hidden_states = hidden_states.flatten(2).transpose(1, 2) + + if timestep.dim() == 2: + timestep = timestep.flatten() + + temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder( + timestep, encoder_hidden_states, encoder_hidden_states_image) + timestep_proj = timestep_proj.unflatten(1, (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) + else: + # encoder_hidden_states is None (e.g. no text encoder) + # MatrixGame uses image-action cross-attn. + pass + + if encoder_hidden_states_image is not None: + if encoder_hidden_states is not None: + encoder_hidden_states = torch.concat( + [encoder_hidden_states_image, encoder_hidden_states], dim=1) + else: + encoder_hidden_states = encoder_hidden_states_image + + # This is [F, H, W] for the ActionModule + grid_sizes = torch.tensor([ + post_patch_num_frames, post_patch_height, post_patch_width + ], + device=hidden_states.device) + + # 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_cis, + grid_sizes=grid_sizes, + mouse_cond=mouse_cond, + keyboard_cond=keyboard_cond) + else: + hidden_states = block( + hidden_states, + encoder_hidden_states, + timestep_proj, + freqs_cis, + grid_sizes=grid_sizes, + mouse_cond=mouse_cond, + keyboard_cond=keyboard_cond) + + # Output + 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) + + return output \ No newline at end of file diff --git a/fastvideo/models/dits/matrix_game/utils.py b/fastvideo/models/dits/matrix_game/utils.py new file mode 100644 index 000000000..d3ef3667a --- /dev/null +++ b/fastvideo/models/dits/matrix_game/utils.py @@ -0,0 +1,307 @@ +from __future__ import annotations + +import os +import random + +# import cv2 +import numpy as np +import torch +from diffusers.utils import export_to_video +from PIL import Image + +from fastvideo.utils import logger + + +CAM_VALUE = 0.1 +CAMERA_MAP = { + "i": [CAM_VALUE, 0], "k": [-CAM_VALUE, 0], + "j": [0, -CAM_VALUE], "l": [0, CAM_VALUE], "u": [0, 0] +} + +KEYBOARD_MAP_4 = { # base_distilled_model (universal): W/S/A/D + "w": [1, 0, 0, 0], "s": [0, 1, 0, 0], + "a": [0, 0, 1, 0], "d": [0, 0, 0, 1], "q": [0, 0, 0, 0] +} +KEYBOARD_MAP_2 = { # gta_distilled_model: W/S only (steering via mouse) + "w": [1, 0], "s": [0, 1], "q": [0, 0] +} +KEYBOARD_MAP_7 = { # templerun_distilled_model: still/w/s/left/right/a/d + "q": [1, 0, 0, 0, 0, 0, 0], # still + "w": [0, 1, 0, 0, 0, 0, 0], # forward + "s": [0, 0, 1, 0, 0, 0, 0], # back + "j": [0, 0, 0, 1, 0, 0, 0], # left (swipe) + "l": [0, 0, 0, 0, 1, 0, 0], # right (swipe) + "a": [0, 0, 0, 0, 0, 1, 0], # a + "d": [0, 0, 0, 0, 0, 0, 1], # d +} +KEYBOARD_MAP = KEYBOARD_MAP_4 # Default for backward compatibility + + +def load_initial_image(image_path: str = None) -> Image.Image: + if image_path and os.path.exists(image_path): + return Image.open(image_path).convert("RGB") + logger.warning("No image provided, creating placeholder...") + return Image.new("RGB", (640, 352), (128, 128, 128)) + + +def create_action_presets(num_frames: int, keyboard_dim: int = 4, seed: int = None): + if keyboard_dim not in (2, 4, 7): + raise ValueError(f"keyboard_dim must be 2, 4, or 7, got {keyboard_dim}") + if num_frames % 4 != 1: + raise ValueError("Matrix-Game conditioning expects num_frames to be 4k+1.") + + # Set seed for reproducibility if provided + if seed is not None: + random.seed(seed) + + num_samples_per_action = 4 + + # Define actions based on keyboard_dim + if keyboard_dim == 4: + # Universal model: W, S, A, D + 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: + # GTA model: W, S only (steering via mouse) + actions_single_action = ["forward", "back"] + actions_double_action = [] + actions_single_camera = ["camera_l", "camera_r"] + keyboard_idx = {"forward": 0, "back": 1} + else: # keyboard_dim == 7 + # Temple Run model: still, w, s, left, right, a, d (no mouse) + actions_single_action = ["forward", "back", "left", "right"] + actions_double_action = [] + actions_single_camera = [] # No mouse for Temple Run + 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}") + + # Ensure we have at least some actions + 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} + + +def parse_config(config, mode="universal"): + assert mode in ['universal', 'gta_drive', 'templerun'] + key_data = {} + mouse_data = {} + if mode != 'templerun': + key, mouse = config + else: + key = config + + for i in range(len(key)): + if mode == 'templerun': + still, w, s, left, right, a, d = key[i] + elif mode == 'universal': + w, s, a, d = key[i] + else: + w, s, a, d = key[i][0], key[i][1], mouse[i][1] < 0, mouse[i][1] > 0 + if mode == 'universal': + mouse_y, mouse_x = mouse[i] + mouse_y = -1 * mouse_y + + key_data[i] = {"W": bool(w), "A": bool(a), "S": bool(s), "D": bool(d)} + if mode == 'templerun': + key_data[i].update({"left": bool(left), "right": bool(right)}) + + if mode == 'universal': + if i == 0: + mouse_data[i] = (320, 352 // 2) + else: + global_scale_factor = 0.1 + mouse_scale_x = 15 * global_scale_factor + mouse_scale_y = 15 * 4 * global_scale_factor + mouse_data[i] = ( + mouse_data[i - 1][0] + mouse_x * mouse_scale_x, + mouse_data[i - 1][1] + mouse_y * mouse_scale_y, + ) + return key_data, mouse_data + + +# NOTE: drawing functions are commented out to avoid cv2/libGL dependency. +# +# def draw_rounded_rectangle(image, top_left, bottom_right, color, radius=10, alpha=0.5): +# overlay = image.copy() +# x1, y1 = top_left +# x2, y2 = bottom_right +# +# cv2.rectangle(overlay, (x1 + radius, y1), (x2 - radius, y2), color, -1) +# cv2.rectangle(overlay, (x1, y1 + radius), (x2, y2 - radius), color, -1) +# cv2.ellipse(overlay, (x1 + radius, y1 + radius), (radius, radius), 180, 0, 90, color, -1) +# cv2.ellipse(overlay, (x2 - radius, y1 + radius), (radius, radius), 270, 0, 90, color, -1) +# cv2.ellipse(overlay, (x1 + radius, y2 - radius), (radius, radius), 90, 0, 90, color, -1) +# cv2.ellipse(overlay, (x2 - radius, y2 - radius), (radius, radius), 0, 0, 90, color, -1) +# cv2.addWeighted(overlay, alpha, image, 1 - alpha, 0, image) +# +# +# def draw_keys_on_frame(frame, keys, key_size=(80, 50), spacing=20, bottom_margin=30, mode='universal'): +# h, w, _ = frame.shape +# horison_shift = 90 +# vertical_shift = -20 +# horizon_shift_all = 50 +# key_positions = { +# "W": (w // 2 - key_size[0] // 2 - horison_shift - horizon_shift_all, +# h - bottom_margin - key_size[1] * 2 + vertical_shift - 20), +# "A": (w // 2 - key_size[0] * 2 + 5 - horison_shift - horizon_shift_all, +# h - bottom_margin - key_size[1] + vertical_shift), +# "S": (w // 2 - key_size[0] // 2 - horison_shift - horizon_shift_all, +# h - bottom_margin - key_size[1] + vertical_shift), +# "D": (w // 2 + key_size[0] - 5 - horison_shift - horizon_shift_all, +# h - bottom_margin - key_size[1] + vertical_shift), +# } +# key_icon = {"W": "W", "A": "A", "S": "S", "D": "D", "left": "left", "right": "right"} +# if mode == 'templerun': +# key_positions.update({ +# "left": (w // 2 + key_size[0] * 2 + spacing * 2 - horison_shift - horizon_shift_all, +# h - bottom_margin - key_size[1] + vertical_shift), +# "right": (w // 2 + key_size[0] * 3 + spacing * 7 - horison_shift - horizon_shift_all, +# h - bottom_margin - key_size[1] + vertical_shift) +# }) +# +# for key, (x, y) in key_positions.items(): +# is_pressed = keys.get(key, False) +# top_left = (x, y) +# if key in ["left", "right"]: +# bottom_right = (x + key_size[0] + 40, y + key_size[1]) +# else: +# bottom_right = (x + key_size[0], y + key_size[1]) +# +# color = (0, 255, 0) if is_pressed else (200, 200, 200) +# alpha = 0.8 if is_pressed else 0.5 +# draw_rounded_rectangle(frame, top_left, bottom_right, color, radius=10, alpha=alpha) +# +# text_size = cv2.getTextSize(key, cv2.FONT_HERSHEY_SIMPLEX, 0.8, 2)[0] +# if key in ["left", "right"]: +# text_x = x + (key_size[0] + 40 - text_size[0]) // 2 +# else: +# text_x = x + (key_size[0] - text_size[0]) // 2 +# text_y = y + (key_size[1] + text_size[1]) // 2 +# cv2.putText(frame, key_icon[key], (text_x, text_y), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 0), 2) +# +# +# def overlay_icon(frame, icon, position, scale=1.0, rotation=0): +# x, y = position +# h, w, _ = icon.shape +# +# scaled_width = int(w * scale) +# scaled_height = int(h * scale) +# icon_resized = cv2.resize(icon, (scaled_width, scaled_height), interpolation=cv2.INTER_AREA) +# +# center = (scaled_width // 2, scaled_height // 2) +# rotation_matrix = cv2.getRotationMatrix2D(center, rotation, 1.0) +# icon_rotated = cv2.warpAffine( +# icon_resized, rotation_matrix, (scaled_width, scaled_height), +# flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT, borderValue=(0, 0, 0, 0) +# ) +# +# h, w, _ = icon_rotated.shape +# frame_h, frame_w, _ = frame.shape +# +# top_left_x = max(0, int(x - w // 2)) +# top_left_y = max(0, int(y - h // 2)) +# bottom_right_x = min(frame_w, int(x + w // 2)) +# bottom_right_y = min(frame_h, int(y + h // 2)) +# +# icon_x_start = max(0, int(-x + w // 2)) +# icon_y_start = max(0, int(-y + h // 2)) +# icon_x_end = icon_x_start + (bottom_right_x - top_left_x) +# icon_y_end = icon_y_start + (bottom_right_y - top_left_y) +# +# icon_region = icon_rotated[icon_y_start:icon_y_end, icon_x_start:icon_x_end] +# alpha = icon_region[:, :, 3] / 255.0 +# icon_rgb = icon_region[:, :, :3] +# +# frame_region = frame[top_left_y:bottom_right_y, top_left_x:bottom_right_x] +# for c in range(3): +# frame_region[:, :, c] = (1 - alpha) * frame_region[:, :, c] + alpha * icon_rgb[:, :, c] +# frame[top_left_y:bottom_right_y, top_left_x:bottom_right_x] = frame_region +# +# +# def process_video(input_video, output_video, config, mouse_icon_path, +# mouse_scale=1.0, mouse_rotation=0, process_icon=True, mode='universal'): +# key_data, mouse_data = parse_config(config, mode=mode) +# fps = 12 +# +# mouse_icon = cv2.imread(mouse_icon_path, cv2.IMREAD_UNCHANGED) +# +# out_video = [] +# for frame_idx, frame in enumerate(input_video): +# frame = np.ascontiguousarray(frame) +# if process_icon: +# keys = key_data.get(frame_idx, {"W": False, "A": False, "S": False, "D": False, "left": False, "right": False}) +# draw_keys_on_frame(frame, keys, key_size=(50, 50), spacing=10, bottom_margin=20, mode=mode) +# if mode == 'universal': +# frame_width = frame.shape[1] +# frame_height = frame.shape[0] +# mouse_position = mouse_data.get(frame_idx, (frame_width // 2, frame_height // 2)) +# overlay_icon(frame, mouse_icon, mouse_position, scale=mouse_scale, rotation=mouse_rotation) +# out_video.append(frame / 255) +# +# export_to_video(out_video, output_video, fps=fps) +# logger.info(f"Video saved to {output_video}") diff --git a/fastvideo/models/registry.py b/fastvideo/models/registry.py index 964acb836..04eb6936d 100644 --- a/fastvideo/models/registry.py +++ b/fastvideo/models/registry.py @@ -34,6 +34,8 @@ _IMAGE_TO_VIDEO_DIT_MODELS = { # "HunyuanVideoTransformer3DModel": ("dits", "hunyuanvideo", "HunyuanVideoDiT"), "WanTransformer3DModel": ("dits", "wanvideo", "WanTransformer3DModel"), "CausalWanTransformer3DModel": ("dits", "causal_wanvideo", "CausalWanTransformer3DModel"), + "MatrixGameWanModel": ("dits", "matrix_game", "MatrixGameWanModel"), + "CausalMatrixGameWanModel": ("dits", "matrix_game", "CausalMatrixGameWanModel"), } _TEXT_ENCODER_MODELS = { @@ -48,6 +50,7 @@ _TEXT_ENCODER_MODELS = { _IMAGE_ENCODER_MODELS: dict[str, tuple] = { # "HunyuanVideoTransformer3DModel": ("image_encoder", "hunyuanvideo", "HunyuanVideoImageEncoder"), "CLIPVisionModelWithProjection": ("encoders", "clip", "CLIPVisionModel"), + "CLIPVisionModel": ("encoders", "clip", "CLIPVisionModel"), } _VAE_MODELS = { diff --git a/fastvideo/pipelines/basic/matrixgame/__init__.py b/fastvideo/pipelines/basic/matrixgame/__init__.py new file mode 100644 index 000000000..cd0ab7b61 --- /dev/null +++ b/fastvideo/pipelines/basic/matrixgame/__init__.py @@ -0,0 +1,7 @@ +# SPDX-License-Identifier: Apache-2.0 +from fastvideo.pipelines.basic.matrixgame.matrixgame_i2v_pipeline import ( + MatrixGamePipeline) +from fastvideo.pipelines.basic.matrixgame.matrixgame_causal_dmd_pipeline import ( + MatrixGameCausalDMDPipeline) + +__all__ = ["MatrixGamePipeline", "MatrixGameCausalDMDPipeline"] diff --git a/fastvideo/pipelines/basic/matrixgame/matrixgame_causal_dmd_pipeline.py b/fastvideo/pipelines/basic/matrixgame/matrixgame_causal_dmd_pipeline.py new file mode 100644 index 000000000..e2f8ce6bb --- /dev/null +++ b/fastvideo/pipelines/basic/matrixgame/matrixgame_causal_dmd_pipeline.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Matrix-Game causal DMD pipeline implementation.""" + +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.pipelines import ComposedPipelineBase, LoRAPipeline + +from fastvideo.pipelines.stages import (ConditioningStage, DecodingStage, + InputValidationStage, + LatentPreparationStage, + TextEncodingStage, + MatrixGameImageEncodingStage, + MatrixGameCausalDenoisingStage) +from fastvideo.pipelines.stages.image_encoding import ( + MatrixGameImageVAEEncodingStage) + +logger = init_logger(__name__) + + +class MatrixGameCausalDMDPipeline(LoRAPipeline, ComposedPipelineBase): + _required_config_modules = [ + "vae", "transformer", "scheduler", "image_encoder", "image_processor" + ] + + def create_pipeline_stages(self, fastvideo_args: FastVideoArgs) -> None: + self.add_stage(stage_name="input_validation_stage", + stage=InputValidationStage()) + + if (self.get_module("text_encoder", None) is not None + and self.get_module("tokenizer", None) is not None): + self.add_stage(stage_name="prompt_encoding_stage", + stage=TextEncodingStage( + text_encoders=[self.get_module("text_encoder")], + tokenizers=[self.get_module("tokenizer")], + )) + + if (self.get_module("image_encoder", None) is not None + and self.get_module("image_processor", None) is not None): + self.add_stage( + stage_name="image_encoding_stage", + stage=MatrixGameImageEncodingStage( + image_encoder=self.get_module("image_encoder"), + image_processor=self.get_module("image_processor"), + )) + + 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", None))) + + self.add_stage( + stage_name="image_latent_preparation_stage", + stage=MatrixGameImageVAEEncodingStage(vae=self.get_module("vae"))) + + self.add_stage(stage_name="denoising_stage", + stage=MatrixGameCausalDenoisingStage( + transformer=self.get_module("transformer"), + transformer_2=self.get_module("transformer_2", None), + 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( + "MatrixGameCausalDMDPipeline initialized with action support") + + +EntryClass = [MatrixGameCausalDMDPipeline] diff --git a/fastvideo/pipelines/basic/matrixgame/matrixgame_i2v_pipeline.py b/fastvideo/pipelines/basic/matrixgame/matrixgame_i2v_pipeline.py new file mode 100644 index 000000000..f403cbf48 --- /dev/null +++ b/fastvideo/pipelines/basic/matrixgame/matrixgame_i2v_pipeline.py @@ -0,0 +1,79 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Matrix-Game I2V pipeline implementation.""" + +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.logger import init_logger +from fastvideo.pipelines.composed_pipeline_base import ComposedPipelineBase +from fastvideo.pipelines.lora_pipeline import LoRAPipeline + +from fastvideo.pipelines.stages import (MatrixGameImageEncodingStage, + ConditioningStage, DecodingStage, + DenoisingStage, InputValidationStage, + LatentPreparationStage, + TextEncodingStage, + TimestepPreparationStage) +from fastvideo.pipelines.stages.image_encoding import ( + MatrixGameImageVAEEncodingStage) +from fastvideo.models.schedulers.scheduling_flow_unipc_multistep import ( + FlowUniPCMultistepScheduler) + +logger = init_logger(__name__) + + +class MatrixGamePipeline(LoRAPipeline, ComposedPipelineBase): + _required_config_modules = [ + "vae", "transformer", "scheduler", "image_encoder", "image_processor" + ] + + 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()) + + if (self.get_module("text_encoder", None) is not None + and self.get_module("tokenizer", None) is not None): + self.add_stage(stage_name="prompt_encoding_stage", + stage=TextEncodingStage( + text_encoders=[self.get_module("text_encoder")], + tokenizers=[self.get_module("tokenizer")], + )) + + if (self.get_module("image_encoder", None) is not None + and self.get_module("image_processor", None) is not None): + self.add_stage( + stage_name="image_encoding_stage", + stage=MatrixGameImageEncodingStage( + image_encoder=self.get_module("image_encoder"), + image_processor=self.get_module("image_processor"), + )) + + self.add_stage(stage_name="conditioning_stage", + stage=ConditioningStage()) + + self.add_stage(stage_name="timestep_preparation_stage", + stage=TimestepPreparationStage( + scheduler=self.get_module("scheduler"))) + + 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=MatrixGameImageVAEEncodingStage(vae=self.get_module("vae"))) + + self.add_stage(stage_name="denoising_stage", + stage=DenoisingStage( + transformer=self.get_module("transformer"), + transformer_2=self.get_module("transformer_2", None), + scheduler=self.get_module("scheduler"))) + + self.add_stage(stage_name="decoding_stage", + stage=DecodingStage(vae=self.get_module("vae"))) + + +EntryClass = [MatrixGamePipeline] diff --git a/fastvideo/pipelines/pipeline_batch_info.py b/fastvideo/pipelines/pipeline_batch_info.py index c46469be5..6b9de8d73 100644 --- a/fastvideo/pipelines/pipeline_batch_info.py +++ b/fastvideo/pipelines/pipeline_batch_info.py @@ -119,6 +119,11 @@ class ForwardBatch: noise_pred: torch.Tensor | None = None image_latent: torch.Tensor | None = None + # Action control inputs (Matrix-Game) + 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] + # Latent dimensions height_latents: list[int] | int | None = None width_latents: list[int] | int | None = None diff --git a/fastvideo/pipelines/pipeline_registry.py b/fastvideo/pipelines/pipeline_registry.py index e6678e1cd..64b7699a5 100644 --- a/fastvideo/pipelines/pipeline_registry.py +++ b/fastvideo/pipelines/pipeline_registry.py @@ -25,7 +25,9 @@ _PIPELINE_NAME_TO_ARCHITECTURE_NAME: dict[str, str] = { "WanCausalDMDPipeline": "wan", "StepVideoPipeline": "stepvideo", "HunyuanVideoPipeline": "hunyuan", - "Cosmos2VideoToWorldPipeline": "cosmos" + "Cosmos2VideoToWorldPipeline": "cosmos", + "MatrixGamePipeline": "matrixgame", + "MatrixGameCausalDMDPipeline": "matrixgame", } _PREPROCESS_WORKLOAD_TYPE_TO_PIPELINE_NAME: dict[WorkloadType, str] = { diff --git a/fastvideo/pipelines/stages/__init__.py b/fastvideo/pipelines/stages/__init__.py index 7371f955d..998d6eb09 100644 --- a/fastvideo/pipelines/stages/__init__.py +++ b/fastvideo/pipelines/stages/__init__.py @@ -14,13 +14,14 @@ from fastvideo.pipelines.stages.denoising import (CosmosDenoisingStage, DenoisingStage, DmdDenoisingStage) from fastvideo.pipelines.stages.encoding import EncodingStage -from fastvideo.pipelines.stages.image_encoding import (ImageEncodingStage, - RefImageEncodingStage, - ImageVAEEncodingStage, - VideoVAEEncodingStage) +from fastvideo.pipelines.stages.image_encoding import ( + ImageEncodingStage, MatrixGameImageEncodingStage, RefImageEncodingStage, + ImageVAEEncodingStage, VideoVAEEncodingStage) from fastvideo.pipelines.stages.input_validation import InputValidationStage from fastvideo.pipelines.stages.latent_preparation import ( CosmosLatentPreparationStage, LatentPreparationStage) +from fastvideo.pipelines.stages.matrixgame_denoising import ( + MatrixGameCausalDenoisingStage) from fastvideo.pipelines.stages.stepvideo_encoding import ( StepvideoPromptEncodingStage) from fastvideo.pipelines.stages.text_encoding import TextEncodingStage @@ -37,10 +38,12 @@ __all__ = [ "DenoisingStage", "DmdDenoisingStage", "CausalDMDDenosingStage", + "MatrixGameCausalDenoisingStage", "CosmosDenoisingStage", "EncodingStage", "DecodingStage", "ImageEncodingStage", + "MatrixGameImageEncodingStage", "RefImageEncodingStage", "ImageVAEEncodingStage", "VideoVAEEncodingStage", diff --git a/fastvideo/pipelines/stages/conditioning.py b/fastvideo/pipelines/stages/conditioning.py index 3eda6a123..8cc4c0a69 100644 --- a/fastvideo/pipelines/stages/conditioning.py +++ b/fastvideo/pipelines/stages/conditioning.py @@ -77,19 +77,32 @@ class ConditioningStage(PipelineStage): fastvideo_args: FastVideoArgs) -> VerificationResult: """Verify conditioning stage inputs.""" result = VerificationResult() + if not batch.prompt_embeds: + # No text encoder/prompt embeddings: skip checks and effectively disable CFG. + batch.do_classifier_free_guidance = False + return result result.add_check("do_classifier_free_guidance", batch.do_classifier_free_guidance, V.bool_value) result.add_check("guidance_scale", batch.guidance_scale, V.positive_float) - result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty) - result.add_check( - "negative_prompt_embeds", batch.negative_prompt_embeds, lambda x: - not batch.do_classifier_free_guidance or V.list_not_empty(x)) + # Matrix-Game allow empty prompt + # embeddings when CFG isn't enabled. + if batch.do_classifier_free_guidance or batch.prompt_embeds: + result.add_check("prompt_embeds", batch.prompt_embeds, + V.list_not_empty) + 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 verify_output(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> VerificationResult: """Verify conditioning stage outputs.""" result = VerificationResult() - result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty) + if batch.prompt_embeds is None or not batch.prompt_embeds: + batch.do_classifier_free_guidance = False + return result + if batch.do_classifier_free_guidance or batch.prompt_embeds: + result.add_check("prompt_embeds", batch.prompt_embeds, + V.list_not_empty) return result diff --git a/fastvideo/pipelines/stages/decoding.py b/fastvideo/pipelines/stages/decoding.py index 830dbf34b..75f3326c8 100644 --- a/fastvideo/pipelines/stages/decoding.py +++ b/fastvideo/pipelines/stages/decoding.py @@ -76,21 +76,39 @@ class DecodingStage(PipelineStage): vae_autocast_enabled = ( vae_dtype != torch.float32) and not fastvideo_args.disable_autocast - if hasattr(self.vae, 'scaling_factor'): + # denormalization for MatrixGame VAE + # z = z * std + mean during decode + if (hasattr(self.vae.config, 'latents_mean') + and hasattr(self.vae.config, 'latents_std')): + # Convert config values to tensors + latents_mean = torch.tensor(self.vae.config.latents_mean, + device=latents.device, + dtype=latents.dtype).view( + 1, -1, 1, 1, 1) + + latents_std = torch.tensor(self.vae.config.latents_std, + device=latents.device, + dtype=latents.dtype).view( + 1, -1, 1, 1, 1) + + # Apply denormalization: z = z * std + mean + latents = latents * latents_std + latents_mean + elif hasattr(self.vae, 'scaling_factor'): + # Standard VAE scaling if isinstance(self.vae.scaling_factor, torch.Tensor): latents = latents / self.vae.scaling_factor.to( latents.device, latents.dtype) else: latents = latents / self.vae.scaling_factor - # Apply shifting if needed - if (hasattr(self.vae, "shift_factor") - and self.vae.shift_factor is not None): - if isinstance(self.vae.shift_factor, torch.Tensor): - latents += self.vae.shift_factor.to(latents.device, - latents.dtype) - else: - latents += self.vae.shift_factor + # Apply shifting if needed + if (hasattr(self.vae, "shift_factor") + and self.vae.shift_factor is not None): + if isinstance(self.vae.shift_factor, torch.Tensor): + latents += self.vae.shift_factor.to(latents.device, + latents.dtype) + else: + latents += self.vae.shift_factor # Decode latents with torch.autocast(device_type="cuda", diff --git a/fastvideo/pipelines/stages/denoising.py b/fastvideo/pipelines/stages/denoising.py index d14ee9699..4ab4e1b76 100644 --- a/fastvideo/pipelines/stages/denoising.py +++ b/fastvideo/pipelines/stages/denoising.py @@ -168,6 +168,14 @@ class DenoisingStage(PipelineStage): }, ) + action_kwargs = self.prepare_extra_func_kwargs( + self.transformer.forward, + { + "mouse_cond": batch.mouse_cond, + "keyboard_cond": batch.keyboard_cond, + }, + ) + # Prepare STA parameters if st_attn_available and self.attn_backend == SlidingTileAttentionBackend: self.prepare_sta_param(batch, fastvideo_args) @@ -379,6 +387,7 @@ class DenoisingStage(PipelineStage): guidance=guidance_expand, **image_kwargs, **pos_cond_kwargs, + **action_kwargs, ) if batch.do_classifier_free_guidance: @@ -395,6 +404,7 @@ class DenoisingStage(PipelineStage): guidance=guidance_expand, **image_kwargs, **neg_cond_kwargs, + **action_kwargs, ) noise_pred_text = noise_pred @@ -1175,4 +1185,4 @@ class DmdDenoisingStage(DenoisingStage): # Update batch with final latents batch.latents = latents - return batch + return batch \ No newline at end of file diff --git a/fastvideo/pipelines/stages/image_encoding.py b/fastvideo/pipelines/stages/image_encoding.py index 3c6ff9b7a..66d8c7db6 100644 --- a/fastvideo/pipelines/stages/image_encoding.py +++ b/fastvideo/pipelines/stages/image_encoding.py @@ -100,6 +100,65 @@ class ImageEncodingStage(PipelineStage): return result +class MatrixGameImageEncodingStage(ImageEncodingStage): + CLIP_MEAN = [0.48145466, 0.4578275, 0.40821073] + CLIP_STD = [0.26862954, 0.26130258, 0.27577711] + + @torch.no_grad() + def forward( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> ForwardBatch: + assert batch.pil_image is not None + self.image_encoder = self.image_encoder.to(get_local_torch_device()) + + image = batch.pil_image + + if isinstance(image, torch.Tensor): + if image.dim() == 5: + image = image[:, :, + 0] # Extract first frame: [B, C, T, H, W] -> [B, C, H, W] + else: + from torchvision import transforms as T + transform = T.Compose([ + T.ToTensor(), # [0, 1] + T.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, + 0.5]), # -> [-1, 1] + ]) + image = transform(image).unsqueeze(0) # [1, C, H, W] + + device = get_local_torch_device() + image = image.to(device) + + # F.interpolate with bicubic + image = torch.nn.functional.interpolate(image, + size=(224, 224), + mode='bicubic', + align_corners=False) + + # [-1, 1] to [0, 1] + image = image * 0.5 + 0.5 + + # CLIP normalization + mean = torch.tensor(self.CLIP_MEAN, device=device, + dtype=image.dtype).view(1, 3, 1, 1) + std = torch.tensor(self.CLIP_STD, device=device, + dtype=image.dtype).view(1, 3, 1, 1) + image = (image - mean) / std + + with set_forward_context(current_timestep=0, attn_metadata=None): + # WAN2_1ControlCLIPVisionConfig sets num_hidden_layers_override=31 + # so last_hidden_state is the second-to-last layer output + outputs = self.image_encoder(pixel_values=image) + image_embeds = outputs.last_hidden_state + + batch.image_embeds.append(image_embeds) + if fastvideo_args.image_encoder_cpu_offload: + self.image_encoder.to('cpu') + return batch + + class RefImageEncodingStage(ImageEncodingStage): """ Stage for encoding reference image prompts into embeddings for Wan2.1 Control models. @@ -526,3 +585,172 @@ class VideoVAEEncodingStage(ImageVAEEncodingStage): result.add_check("video_latent", batch.video_latent, [V.is_tensor, V.with_dims(5)]) return result + + +class MatrixGameImageVAEEncodingStage(ImageVAEEncodingStage): + + def forward( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> ForwardBatch: + assert batch.pil_image is not None + + if fastvideo_args.mode == ExecutionMode.INFERENCE: + # Accept both PIL.Image and torch.Tensor + # Causal pipeline (is_causal=True) converts PIL to tensor in InputValidationStage + assert batch.pil_image is not None and isinstance( + batch.pil_image, PIL.Image.Image | torch.Tensor) + assert batch.height is not None and isinstance(batch.height, int) + assert batch.width is not None and isinstance(batch.width, int) + assert batch.num_frames is not None and isinstance( + batch.num_frames, int) + height = batch.height + width = batch.width + num_frames = batch.num_frames + elif fastvideo_args.mode == ExecutionMode.PREPROCESS: + assert batch.pil_image is not None and isinstance( + batch.pil_image, torch.Tensor) + assert batch.height is not None and isinstance(batch.height, list) + assert batch.width is not None and isinstance(batch.width, list) + assert batch.num_frames is not None and isinstance( + batch.num_frames, list) + num_frames = batch.num_frames[0] + height = batch.height[0] + width = batch.width[0] + else: + # Fallback for other modes + height = batch.height if isinstance(batch.height, + int) else batch.height[0] + width = batch.width if isinstance(batch.width, + int) else batch.width[0] + num_frames = batch.num_frames if isinstance( + batch.num_frames, int) else batch.num_frames[0] + + self.vae = self.vae.to(get_local_torch_device()) + + # Process single image for I2V (latent dimensions computed but not used directly) + + image = batch.pil_image + + # Handle tensor input from causal pipeline + if isinstance(image, torch.Tensor): + # Causal pipeline provides tensor in [B, C, F, H, W] format + if image.dim() == 5: + # Already 5D, extract first frame for conditioning + # Shape: [B, C, F, H, W] -> use first frame [B, C, 1, H, W] + first_frame = image[:, :, :1] # Keep dim, [B, C, 1, H, W] + # Create video condition with first frame + zeros + video_condition = torch.cat([ + first_frame, + first_frame.new_zeros( + first_frame.shape[0], first_frame.shape[1], num_frames - + 1, first_frame.shape[3], first_frame.shape[4]) + ], + dim=2) + elif image.dim() == 4: + # [B, C, H, W] -> need to add frame dim + image = image.unsqueeze(2) # [B, C, 1, H, W] + video_condition = torch.cat([ + image, + image.new_zeros(image.shape[0], image.shape[1], num_frames - + 1, image.shape[3], image.shape[4]) + ], + dim=2) + else: + raise ValueError(f"Unexpected tensor dimensions: {image.dim()}") + video_condition = video_condition.to(get_local_torch_device(), + dtype=torch.float32) + else: + # PIL Image input - use preprocess + image = self.preprocess( + image, + vae_scale_factor=self.vae.spatial_compression_ratio, + height=height, + width=width).to(get_local_torch_device(), dtype=torch.float32) + + # (B, C, H, W) -> (B, C, 1, H, W) + image = image.unsqueeze(2) + + # Create video tensor with first frame as image, rest as zeros + video_condition = torch.cat([ + image, + image.new_zeros(image.shape[0], image.shape[1], num_frames - 1, + image.shape[3], image.shape[4]) + ], + dim=2) + video_condition = video_condition.to( + device=get_local_torch_device(), dtype=torch.float32) + + # Setup VAE precision + vae_dtype = PRECISION_TO_TYPE[ + fastvideo_args.pipeline_config.vae_precision] + vae_autocast_enabled = ( + vae_dtype != torch.float32) and not fastvideo_args.disable_autocast + + # Encode Image + 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) + + # MatrixGame uses deterministic VAE encode for the first-frame conditioning. + # Sampling would inject random noise into the cond_concat tensor and destroy the action guidance. + img_cond = encoder_output.mode() + + # manually using latents_mean and latents_std from config... + if (hasattr(self.vae.config, 'latents_mean') + and hasattr(self.vae.config, 'latents_std')): + # Convert config values to tensors + 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) + + # Apply normalization: (latent - mean) * (1/std) + img_cond = (img_cond - latents_mean) / latents_std + elif (hasattr(self.vae, "shift_factor") + and self.vae.shift_factor is not None): + # Fallback to shift_factor/scaling_factor if available + 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 + + # Create mask_cond: ones for first frame, zeros for rest + # Shape: (B, 16, latent_frames, latent_height, latent_width) + mask_cond = torch.ones_like(img_cond) + mask_cond[:, :, 1:] = 0 # Set all frames except first to 0 + + # Create cond_concat: first 4 channels of mask + all 16 channels of img_cond + # Shape: (B, 20, latent_frames, latent_height, latent_width) + cond_concat = torch.cat([mask_cond[:, :4], img_cond], dim=1) + + # Store cond_concat in batch.image_latent + # This will be concatenated with noise latents in DenoisingStage + batch.image_latent = cond_concat + + # Offload models if needed + if hasattr(self, 'maybe_free_model_hooks'): + self.maybe_free_model_hooks() + + self.vae.to("cpu") + + return batch diff --git a/fastvideo/pipelines/stages/input_validation.py b/fastvideo/pipelines/stages/input_validation.py index 100199ec8..089fd7192 100644 --- a/fastvideo/pipelines/stages/input_validation.py +++ b/fastvideo/pipelines/stages/input_validation.py @@ -39,6 +39,7 @@ class InputValidationStage(PipelineStage): assert seed is not None seeds = [seed + i for i in range(num_videos_per_prompt)] batch.seeds = seeds + # Peiyuan: using GPU seed will cause A100 and H100 to generate different results... batch.generator = [ torch.Generator("cpu").manual_seed(seed) for seed in seeds @@ -111,20 +112,28 @@ class InputValidationStage(PipelineStage): ) and batch.pil_image is not None: img = batch.pil_image ih, iw = img.height, img.width - patch_size = fastvideo_args.pipeline_config.dit_config.arch_config.patch_size - vae_stride = fastvideo_args.pipeline_config.vae_config.arch_config.scale_factor_spatial - dh, dw = patch_size[1] * vae_stride, patch_size[2] * vae_stride - max_area = 480 * 832 - ow, oh = best_output_size(iw, ih, dw, dh, max_area) - scale = max(ow / iw, oh / ih) - img = img.resize((round(iw * scale), round(ih * scale)), - Image.LANCZOS) + pipeline_class_name = type(fastvideo_args.pipeline_config).__name__ + if 'MatrixGame' in pipeline_class_name or 'MatrixCausal' in pipeline_class_name: + oh, ow = batch.height, batch.width + img = img.resize((ow, oh), Image.LANCZOS) + else: + # Standard Wan logic + patch_size = fastvideo_args.pipeline_config.dit_config.arch_config.patch_size + vae_stride = fastvideo_args.pipeline_config.vae_config.arch_config.scale_factor_spatial + dh, dw = patch_size[1] * vae_stride, patch_size[2] * vae_stride + max_area = 480 * 832 + ow, oh = best_output_size(iw, ih, dw, dh, max_area) + + scale = max(ow / iw, oh / ih) + img = img.resize((round(iw * scale), round(ih * scale)), + Image.LANCZOS) + + # center-crop + x1 = (img.width - ow) // 2 + y1 = (img.height - oh) // 2 + img = img.crop((x1, y1, x1 + ow, y1 + oh)) - # center-crop - x1 = (img.width - ow) // 2 - y1 = (img.height - oh) // 2 - img = img.crop((x1, y1, x1 + ow, y1 + oh)) assert img.width == ow and img.height == oh logger.info("final processed img height: %s, img width: %s", img.height, img.width) @@ -186,6 +195,43 @@ class InputValidationStage(PipelineStage): batch.video_latent = input_video + # Validate action control inputs (Matrix-Game) + if batch.mouse_cond is not None: + if batch.mouse_cond.dim() != 3 or batch.mouse_cond.shape[-1] != 2: + raise ValueError( + f"mouse_cond must have shape (B, T, 2), but got {batch.mouse_cond.shape}" + ) + logger.info("Action control: mouse_cond validated - shape %s", + batch.mouse_cond.shape) + + if batch.keyboard_cond is not None: + if batch.keyboard_cond.dim() != 3: + raise ValueError( + f"keyboard_cond must have 3 dimensions (B, T, K), but got {batch.keyboard_cond.dim()}" + ) + keyboard_dim = batch.keyboard_cond.shape[-1] + if keyboard_dim not in {2, 4, 6, 7}: + raise ValueError( + f"keyboard_cond last dimension must be 2, 4, 6, or 7, but got {keyboard_dim}" + ) + logger.info( + "Action control: keyboard_cond validated - shape %s (dim=%d)", + batch.keyboard_cond.shape, keyboard_dim) + + if batch.grid_sizes is not None: + if not isinstance(batch.grid_sizes, list | tuple | torch.Tensor): + raise ValueError("grid_sizes must be a list, tuple, or tensor") + if isinstance(batch.grid_sizes, torch.Tensor): + if batch.grid_sizes.numel() != 3: + raise ValueError( + "grid_sizes must have 3 elements [F, H, W]") + else: + if len(batch.grid_sizes) != 3: + raise ValueError( + "grid_sizes must have 3 elements [F, H, W]") + logger.info("Action control: grid_sizes validated - %s", + batch.grid_sizes) + return batch def verify_input(self, batch: ForwardBatch, diff --git a/fastvideo/pipelines/stages/latent_preparation.py b/fastvideo/pipelines/stages/latent_preparation.py index e726e3070..125e76207 100644 --- a/fastvideo/pipelines/stages/latent_preparation.py +++ b/fastvideo/pipelines/stages/latent_preparation.py @@ -57,8 +57,17 @@ class LatentPreparationStage(PipelineStage): # Adjust video length based on VAE version if needed if hasattr(self, 'adjust_video_length'): latent_num_frames = self.adjust_video_length(batch, fastvideo_args) - # Determine batch size - if isinstance(batch.prompt, list): + # Determine batch size; fall back to action/image inputs when no text encoder is present + if not batch.prompt_embeds: + if batch.keyboard_cond is not None: + batch_size = batch.keyboard_cond.shape[0] + elif batch.mouse_cond is not None: + batch_size = batch.mouse_cond.shape[0] + elif batch.image_embeds: + batch_size = batch.image_embeds[0].shape[0] + else: + batch_size = 1 + elif isinstance(batch.prompt, list): batch_size = len(batch.prompt) elif batch.prompt is not None: batch_size = 1 @@ -69,6 +78,19 @@ class LatentPreparationStage(PipelineStage): batch_size *= batch.num_videos_per_prompt # Get required parameters + if not batch.prompt_embeds: + # Create a dummy zero-length text embedding to satisfy downstream checks. + # Matrix-Game models have text_dim=0 and ignore encoder_hidden_states. + transformer_dtype = next(self.transformer.parameters()).dtype + device = get_local_torch_device() + dummy_prompt = torch.zeros(batch_size, + 0, + self.transformer.hidden_size, + device=device, + dtype=transformer_dtype) + batch.prompt_embeds = [dummy_prompt] + batch.negative_prompt_embeds = [] + batch.do_classifier_free_guidance = False dtype = batch.prompt_embeds[0].dtype device = get_local_torch_device() generator = batch.generator @@ -157,10 +179,12 @@ class LatentPreparationStage(PipelineStage): """Verify latent preparation stage inputs.""" result = VerificationResult() result.add_check( - "prompt_or_embeds", None, lambda _: V.string_or_list_strings( - batch.prompt) or V.list_not_empty(batch.prompt_embeds)) - result.add_check("prompt_embeds", batch.prompt_embeds, - V.list_of_tensors) + "prompt_or_embeds", None, + lambda _: V.string_or_list_strings(batch.prompt) or not batch. + prompt_embeds or V.list_not_empty(batch.prompt_embeds)) + if batch.prompt_embeds: + result.add_check("prompt_embeds", batch.prompt_embeds, + V.list_of_tensors) result.add_check("num_videos_per_prompt", batch.num_videos_per_prompt, V.positive_int) result.add_check("generator", batch.generator, diff --git a/fastvideo/pipelines/stages/matrixgame_denoising.py b/fastvideo/pipelines/stages/matrixgame_denoising.py new file mode 100644 index 000000000..70e17127d --- /dev/null +++ b/fastvideo/pipelines/stages/matrixgame_denoising.py @@ -0,0 +1,566 @@ +from __future__ import annotations +from typing import Any + +import torch # type: ignore + +from fastvideo.distributed import get_local_torch_device +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.forward_context import set_forward_context +from fastvideo.logger import init_logger +from fastvideo.models.utils import pred_noise_to_pred_video, pred_noise_to_x_bound +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 + +try: + from fastvideo.attention.backends.sliding_tile_attn import ( + SlidingTileAttentionBackend) + st_attn_available = True +except ImportError: + st_attn_available = False + SlidingTileAttentionBackend = None # type: ignore + +try: + from fastvideo.attention.backends.video_sparse_attn import ( + VideoSparseAttentionBackend) + vsa_available = True +except ImportError: + vsa_available = False + VideoSparseAttentionBackend = None # type: ignore + +logger = init_logger(__name__) + + +class MatrixGameCausalDenoisingStage(DenoisingStage): + + def __init__(self, + transformer, + scheduler, + pipeline=None, + transformer_2=None, + vae=None) -> None: + super().__init__(transformer, scheduler, pipeline, transformer_2, vae) + self.transformer = transformer + self.transformer_2 = transformer_2 + self.vae = vae + self.num_transformer_blocks = len(self.transformer.blocks) + + if hasattr(self.transformer, 'config') and hasattr( + self.transformer.config, 'arch_config'): + self.num_frame_per_block = getattr( + self.transformer.config.arch_config, 'num_frames_per_block', + getattr(self.transformer, 'num_frame_per_block', 1)) + self.sliding_window_num_frames = getattr( + self.transformer.config.arch_config, + 'sliding_window_num_frames', 15) + else: + self.num_frame_per_block = getattr(self.transformer, + 'num_frame_per_block', 1) + self.sliding_window_num_frames = 15 + + try: + self.local_attn_size = getattr(self.transformer, "local_attn_size", + -1) + except Exception: + self.local_attn_size = -1 + + assert self.local_attn_size != -1, ( + f"local_attn_size must be set for Matrix-Game causal inference, " + f"got {self.local_attn_size}. Check MatrixGameWanVideoArchConfig.") + assert self.num_frame_per_block > 0, ( + f"num_frame_per_block must be positive, got {self.num_frame_per_block}" + ) + + logger.info( + "MatrixGame causal inference initialized: " + "local_attn_size=%s, num_frame_per_block=%s", self.local_attn_size, + self.num_frame_per_block) + + self.action_config = getattr(self.transformer, 'action_config', {}) + self.use_action_module = len(self.action_config) > 0 + + def forward( + self, + batch: ForwardBatch, + fastvideo_args: FastVideoArgs, + ) -> ForwardBatch: + target_dtype = torch.bfloat16 + autocast_enabled = (target_dtype != torch.float32 + ) and not fastvideo_args.disable_autocast + + latent_seq_length = batch.latents.shape[-1] * batch.latents.shape[-2] + patch_size = self.transformer.patch_size + patch_ratio = patch_size[-1] * patch_size[-2] + self.frame_seq_length = latent_seq_length // patch_ratio + + independent_first_frame = getattr(self.transformer, + 'independent_first_frame', False) + timesteps = torch.tensor( + fastvideo_args.pipeline_config.dmd_denoising_steps, + dtype=torch.long).cpu() + if fastvideo_args.pipeline_config.warp_denoising_step: + scheduler_timesteps = torch.cat((self.scheduler.timesteps.cpu(), + torch.tensor([0], + dtype=torch.float32))) + timesteps = scheduler_timesteps[1000 - timesteps] + timesteps = timesteps.to(get_local_torch_device()) + + boundary_ratio = getattr(fastvideo_args.pipeline_config.dit_config, + 'boundary_ratio', None) + if boundary_ratio is not None: + boundary_timestep = boundary_ratio * self.scheduler.num_train_timesteps + high_noise_timesteps = timesteps[timesteps >= boundary_timestep] + else: + boundary_timestep = None + high_noise_timesteps = None + + image_embeds = batch.image_embeds + if len(image_embeds) > 0: + assert torch.isnan(image_embeds[0]).sum() == 0 + image_embeds = [ + image_embed.to(target_dtype) for image_embed in image_embeds + ] + + # directly set the kwarg. + image_kwargs = {"encoder_hidden_states_image": image_embeds} + pos_cond_kwargs: dict[str, Any] = {} + + if st_attn_available and self.attn_backend == SlidingTileAttentionBackend: + self.prepare_sta_param(batch, fastvideo_args) + + assert batch.latents is not None, "latents must be provided" + latents = batch.latents + b, c, t, h, w = latents.shape + prompt_embeds = batch.prompt_embeds + assert torch.isnan(prompt_embeds[0]).sum() == 0 + + kv_cache1 = self._initialize_kv_cache(batch_size=latents.shape[0], + dtype=target_dtype, + device=latents.device) + kv_cache2 = None + if boundary_timestep is not None: + kv_cache2 = self._initialize_kv_cache(batch_size=latents.shape[0], + dtype=target_dtype, + device=latents.device) + + kv_cache_mouse = None + kv_cache_keyboard = None + if self.use_action_module: + kv_cache_mouse, kv_cache_keyboard = self._initialize_action_kv_cache( + batch_size=latents.shape[0], + dtype=target_dtype, + device=latents.device) + + def _get_kv_cache(timestep_val: float) -> list[dict]: + if boundary_timestep is not None: + if timestep_val >= boundary_timestep: + return kv_cache1 + else: + assert kv_cache2 is not None, "kv_cache2 is not initialized" + return kv_cache2 + return kv_cache1 + + crossattn_cache = self._initialize_crossattn_cache( + batch_size=latents.shape[0], + max_text_len=257, # 1 CLS + 256 patch tokens + dtype=target_dtype, + device=latents.device) + + pos_start_base = 0 + + if t % self.num_frame_per_block != 0: + raise ValueError( + "num_frames must be divisible by num_frame_per_block for causal denoising" + ) + num_blocks = t // self.num_frame_per_block + block_sizes = [self.num_frame_per_block] * num_blocks + start_index = 0 + + if boundary_timestep is not None: + block_sizes[0] = 1 + + # NOTE: MatrixGame does NOT process the first frame separately. + # The first frame information is already encoded in batch.image_latent (cond_concat) + # and will be used by the model via channel concatenation: torch.cat([x, cond_concat], dim=1) + + with self.progress_bar(total=len(block_sizes) * + len(timesteps)) as progress_bar: + for block_idx, current_num_frames in enumerate(block_sizes): + current_latents = latents[:, :, start_index:start_index + + current_num_frames, :, :] + noise_latents_btchw = current_latents.permute(0, 2, 1, 3, 4) + _video_raw_latent_shape = noise_latents_btchw.shape # noqa: F841 + + # NOTE: crossattn_cache should NOT be reset between blocks! + + action_kwargs = self._prepare_action_kwargs( + batch, start_index, current_num_frames) + + for i, t_cur in enumerate(timesteps): + if boundary_timestep is not None and t_cur < boundary_timestep: + current_model = self.transformer_2 if self.transformer_2 is not None else self.transformer + else: + current_model = self.transformer + + noise_latents = noise_latents_btchw.clone() + latent_model_input = current_latents.to(target_dtype) + + if batch.image_latent is not None and independent_first_frame and start_index == 0: + latent_model_input = torch.cat([ + latent_model_input, + batch.image_latent.to(target_dtype) + ], + dim=2) + + # t_expand needs to be [batch * frames] to match flattened pred_noise/noise_latents + t_expand = t_cur.repeat(latent_model_input.shape[0] * + current_num_frames) + + if vsa_available and self.attn_backend == VideoSparseAttentionBackend: + self.attn_metadata_builder_cls = self.attn_backend.get_builder_cls( + ) + if self.attn_metadata_builder_cls is not None: + self.attn_metadata_builder = self.attn_metadata_builder_cls( + ) + attn_metadata = self.attn_metadata_builder.build( + current_timestep=i, + raw_latent_shape=(current_num_frames, h, w), + patch_size=fastvideo_args.pipeline_config. + dit_config.patch_size, + STA_param=batch.STA_param, + VSA_sparsity=fastvideo_args.VSA_sparsity, + device=get_local_torch_device(), + ) + assert attn_metadata is not None, "attn_metadata cannot be None" + else: + attn_metadata = None + else: + attn_metadata = None + + with torch.autocast(device_type="cuda", + dtype=target_dtype, + enabled=autocast_enabled), \ + set_forward_context(current_timestep=i, + attn_metadata=attn_metadata, + forward_batch=batch): + # Expand timestep to per-frame format [batch, num_frames] for causal model + t_expanded_noise = t_cur * torch.ones( + (latent_model_input.shape[0], current_num_frames), + device=latent_model_input.device, + dtype=torch.long) + + model_kwargs = { + "kv_cache": _get_kv_cache(t_cur), + "crossattn_cache": crossattn_cache, + "current_start": (pos_start_base + start_index) * + self.frame_seq_length, + "start_frame": start_index, + } + + if self.use_action_module and current_model == self.transformer: + model_kwargs.update({ + "kv_cache_mouse": + kv_cache_mouse, + "kv_cache_keyboard": + kv_cache_keyboard, + }) + model_kwargs.update(action_kwargs) + + pred_noise_btchw = current_model( + latent_model_input, + prompt_embeds, + t_expanded_noise, + **image_kwargs, + **pos_cond_kwargs, + **model_kwargs, + ).permute(0, 2, 1, 3, 4) + + if boundary_timestep is not None and t_cur >= boundary_timestep: + pred_video_btchw = pred_noise_to_x_bound( + pred_noise=pred_noise_btchw.flatten(0, 1), + noise_input_latent=noise_latents.flatten(0, 1), + timestep=t_expand, + boundary_timestep=torch.ones_like(t_expand) * + boundary_timestep, + scheduler=self.scheduler).unflatten( + 0, pred_noise_btchw.shape[:2]) + else: + pred_video_btchw = pred_noise_to_pred_video( + pred_noise=pred_noise_btchw.flatten(0, 1), + noise_input_latent=noise_latents.flatten(0, 1), + timestep=t_expand, + scheduler=self.scheduler).unflatten( + 0, pred_noise_btchw.shape[:2]) + + if i < len(timesteps) - 1: + next_timestep = timesteps[i + 1] * torch.ones( + [1], + dtype=torch.long, + device=pred_video_btchw.device) + noise = torch.randn( + pred_video_btchw.shape, + dtype=pred_video_btchw.dtype, + generator=(batch.generator[0] if isinstance( + batch.generator, list) else + batch.generator)).to( + pred_video_btchw.device) + noise_btchw = noise + if boundary_timestep is not None and high_noise_timesteps is not None and i < len( + high_noise_timesteps) - 1: + noise_latents_btchw = self.scheduler.add_noise_high( + pred_video_btchw.flatten(0, 1), + noise_btchw.flatten(0, 1), next_timestep, + torch.ones_like(next_timestep) * + boundary_timestep).unflatten( + 0, pred_video_btchw.shape[:2]) + elif boundary_timestep is not None and high_noise_timesteps is not None and i == len( + high_noise_timesteps) - 1: + noise_latents_btchw = pred_video_btchw + else: + noise_latents_btchw = self.scheduler.add_noise( + pred_video_btchw.flatten(0, 1), + noise_btchw.flatten(0, 1), + next_timestep).unflatten( + 0, pred_video_btchw.shape[:2]) + current_latents = noise_latents_btchw.permute( + 0, 2, 1, 3, 4) + else: + current_latents = pred_video_btchw.permute( + 0, 2, 1, 3, 4) + + if progress_bar is not None: + progress_bar.update() + + latents[:, :, start_index:start_index + + current_num_frames, :, :] = current_latents + + context_noise = getattr(fastvideo_args.pipeline_config, + "context_noise", 0) + # Expand context timestep to per-frame format [batch, num_frames] for causal model + t_context = torch.ones([latents.shape[0], current_num_frames], + device=latents.device, + dtype=torch.long) * int(context_noise) + context_bcthw = current_latents.to(target_dtype) + + with torch.autocast(device_type="cuda", + dtype=target_dtype, + enabled=autocast_enabled), \ + set_forward_context(current_timestep=0, + attn_metadata=attn_metadata, + forward_batch=batch): + + context_model_kwargs = { + "kv_cache": kv_cache1, + "crossattn_cache": crossattn_cache, + "current_start": + (pos_start_base + start_index) * self.frame_seq_length, + "start_frame": start_index, + } + + if self.use_action_module: + context_model_kwargs.update({ + "kv_cache_mouse": + kv_cache_mouse, + "kv_cache_keyboard": + kv_cache_keyboard, + }) + context_model_kwargs.update(action_kwargs) + + if boundary_timestep is not None and self.transformer_2 is not None: + self.transformer_2( + context_bcthw, + prompt_embeds, + t_context, + kv_cache=kv_cache2, + crossattn_cache=crossattn_cache, + current_start=(pos_start_base + start_index) * + self.frame_seq_length, + start_frame=start_index, + **image_kwargs, + **pos_cond_kwargs, + ) + + self.transformer( + context_bcthw, + prompt_embeds, + t_context, + **image_kwargs, + **pos_cond_kwargs, + **context_model_kwargs, + ) + + start_index += current_num_frames + + if boundary_timestep is not None: + num_frames_to_remove = self.num_frame_per_block - 1 + if num_frames_to_remove > 0: + latents = latents[:, :, :-num_frames_to_remove, :, :] + + batch.latents = latents + return batch + + def _prepare_action_kwargs(self, batch: ForwardBatch, start_index: int, + num_frames: int) -> dict[str, Any]: + action_kwargs: dict[str, Any] = {} + if not self.use_action_module: + return action_kwargs + + vae_time_compression_ratio = 4 + end_frame_idx = 1 + vae_time_compression_ratio * (start_index + + num_frames - 1) + if hasattr(batch, 'mouse_cond') and batch.mouse_cond is not None: + action_kwargs['mouse_cond'] = batch.mouse_cond[:, :end_frame_idx] + if hasattr(batch, 'keyboard_cond') and batch.keyboard_cond is not None: + action_kwargs[ + 'keyboard_cond'] = batch.keyboard_cond[:, :end_frame_idx] + + # CRITICAL: Pass num_frame_per_block to model - this should be num_frames (current block size) + action_kwargs['num_frame_per_block'] = num_frames + return action_kwargs + + def _initialize_kv_cache(self, batch_size: int, dtype: torch.dtype, + device: torch.device) -> list[dict]: + kv_cache = [] + num_attention_heads = self.transformer.num_attention_heads + attention_head_dim = getattr( + self.transformer, 'attention_head_dim', + self.transformer.hidden_size // num_attention_heads) + if self.local_attn_size != -1: + kv_cache_size = self.local_attn_size * self.frame_seq_length + else: + kv_cache_size = self.frame_seq_length * self.sliding_window_num_frames + + for _ in range(self.num_transformer_blocks): + kv_cache.append({ + "k": + torch.zeros([ + batch_size, kv_cache_size, num_attention_heads, + attention_head_dim + ], + dtype=dtype, + device=device), + "v": + torch.zeros([ + batch_size, kv_cache_size, num_attention_heads, + attention_head_dim + ], + dtype=dtype, + device=device), + "global_end_index": + torch.tensor([0], dtype=torch.long, device=device), + "local_end_index": + torch.tensor([0], dtype=torch.long, device=device), + }) + + return kv_cache + + def _initialize_action_kv_cache(self, batch_size: int, dtype: torch.dtype, + device: torch.device): + kv_cache_mouse = [] + kv_cache_keyboard = [] + + action_heads = self.action_config.get('heads_num', 16) + mouse_head_dim = self.action_config.get('mouse_hidden_dim', + 1024) // action_heads + keyboard_head_dim = self.action_config.get('keyboard_hidden_dim', + 1024) // action_heads + + if self.local_attn_size != -1: + kv_cache_size = self.local_attn_size + else: + kv_cache_size = 15 + + for _ in range(self.num_transformer_blocks): + kv_cache_keyboard.append({ + "k": + torch.zeros([ + batch_size, kv_cache_size, action_heads, keyboard_head_dim + ], + dtype=dtype, + device=device), + "v": + torch.zeros([ + batch_size, kv_cache_size, action_heads, keyboard_head_dim + ], + dtype=dtype, + device=device), + "global_end_index": + torch.tensor([0], dtype=torch.long, device=device), + "local_end_index": + torch.tensor([0], dtype=torch.long, device=device), + }) + kv_cache_mouse.append({ + "k": + torch.zeros([ + batch_size * self.frame_seq_length, kv_cache_size, + action_heads, mouse_head_dim + ], + dtype=dtype, + device=device), + "v": + torch.zeros([ + batch_size * self.frame_seq_length, kv_cache_size, + action_heads, mouse_head_dim + ], + dtype=dtype, + device=device), + "global_end_index": + torch.tensor([0], dtype=torch.long, device=device), + "local_end_index": + torch.tensor([0], dtype=torch.long, device=device), + }) + + return kv_cache_mouse, kv_cache_keyboard + + def _initialize_crossattn_cache(self, batch_size: int, max_text_len: int, + dtype: torch.dtype, + device: torch.device) -> list[dict]: + crossattn_cache = [] + num_attention_heads = self.transformer.num_attention_heads + attention_head_dim = getattr( + self.transformer, 'attention_head_dim', + self.transformer.hidden_size // num_attention_heads) + for _ in range(self.num_transformer_blocks): + crossattn_cache.append({ + "k": + torch.zeros([ + batch_size, max_text_len, num_attention_heads, + attention_head_dim + ], + dtype=dtype, + device=device), + "v": + torch.zeros([ + batch_size, max_text_len, num_attention_heads, + attention_head_dim + ], + dtype=dtype, + device=device), + "is_init": + False, + }) + return crossattn_cache + + 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 diff --git a/fastvideo/tests/ssim/L40S_reference_videos/Matrix-Game-2.0-Diffusers-Base/FLASH_ATTN/video.mp4 b/fastvideo/tests/ssim/L40S_reference_videos/Matrix-Game-2.0-Diffusers-Base/FLASH_ATTN/video.mp4 new file mode 100644 index 000000000..deca1faad Binary files /dev/null and b/fastvideo/tests/ssim/L40S_reference_videos/Matrix-Game-2.0-Diffusers-Base/FLASH_ATTN/video.mp4 differ diff --git a/fastvideo/tests/ssim/test_matrixgame_similarity.py b/fastvideo/tests/ssim/test_matrixgame_similarity.py new file mode 100644 index 000000000..87cfc0030 --- /dev/null +++ b/fastvideo/tests/ssim/test_matrixgame_similarity.py @@ -0,0 +1,176 @@ +# SPDX-License-Identifier: Apache-2.0 +import os + +import torch +import pytest + +from fastvideo import VideoGenerator +from fastvideo.models.dits.matrix_game.utils import create_action_presets +from fastvideo.logger import init_logger +from fastvideo.tests.utils import ( + compute_video_ssim_torchvision, + write_ssim_results, +) +from fastvideo.worker.multiproc_executor import MultiprocExecutor + +logger = init_logger(__name__) + +device_name = torch.cuda.get_device_name() +device_reference_folder_suffix = "_reference_videos" + +if "A40" in device_name: + device_reference_folder = "A40" + device_reference_folder_suffix +elif "L40S" in device_name: + device_reference_folder = "L40S" + device_reference_folder_suffix +else: + # device_reference_folder = "L40S" + device_reference_folder_suffix + logger.warning(f"Unsupported device for ssim tests: {device_name}") + # raise ValueError(f"Unsupported device for ssim tests: {device_name}") + +# Base parameters from the shell script + +MATRIXGAME_PARAMS = { + "num_gpus": 1, + "model_path": "FastVideo/Matrix-Game-2.0-Base-Diffusers", + "height": 352, + "width": 640, + "num_frames": 117, + "num_inference_steps": 10, + "seed": 1024, + "keyboard_dim": 4, +} + + +MODEL_TO_PARAMS = { + "Matrix-Game-2.0-Diffusers-Base": MATRIXGAME_PARAMS, +} + +I2V_MODEL_TO_PARAMS = {} + +TEST_PROMPTS = [ + "", # MatrixGame is I2V with action conditions, no text prompt +] + +TEST_IMAGE_PATHS = [ + "https://raw.githubusercontent.com/SkyworkAI/Matrix-Game/main/Matrix-Game-2/demo_images/universal/0002.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_matrixgame_similarity(prompt, ATTENTION_BACKEND, model_id): + """ + Test that runs inference with different parameters 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__)) + + base_output_dir = os.path.join(script_dir, "generated_videos", model_id) + output_dir = os.path.join(base_output_dir, ATTENTION_BACKEND) + output_video_name = "video.mp4" + + os.makedirs(output_dir, exist_ok=True) + + BASE_PARAMS = MODEL_TO_PARAMS[model_id] + num_inference_steps = BASE_PARAMS["num_inference_steps"] + + # Create action conditions for MatrixGame + actions = create_action_presets( + BASE_PARAMS["num_frames"], keyboard_dim=BASE_PARAMS["keyboard_dim"], + seed=BASE_PARAMS["seed"] + ) + latent_frames = (BASE_PARAMS["num_frames"] - 1) // 4 + 1 + grid_sizes = torch.tensor([latent_frames, 44, 80]) + + init_kwargs = { + "num_gpus": BASE_PARAMS["num_gpus"], + "use_fsdp_inference": True, + "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"], + "seed": BASE_PARAMS["seed"], + "mouse_cond": actions["mouse"].unsqueeze(0), + "keyboard_cond": actions["keyboard"].unsqueeze(0), + "grid_sizes": grid_sizes, + "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 = os.path.join( + 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}" + ) + + # Find the matching reference video based on the prompt + reference_video_name = None + + for filename in os.listdir(reference_folder): + if filename.endswith(".mp4"): + reference_video_name = filename + break + + 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") + + reference_video_path = os.path.join(reference_folder, reference_video_name) + generated_video_path = os.path.join(output_dir, output_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}" + ) diff --git a/fastvideo/tests/transformers/test_matrixgame.py b/fastvideo/tests/transformers/test_matrixgame.py new file mode 100644 index 000000000..d7c1253d0 --- /dev/null +++ b/fastvideo/tests/transformers/test_matrixgame.py @@ -0,0 +1,133 @@ +# SPDX-License-Identifier: Apache-2.0 +import os + +import pytest +import torch + +from fastvideo.configs.models.dits.matrixgame import MatrixGameWanVideoConfig +from fastvideo.configs.pipelines import PipelineConfig +from fastvideo.distributed.parallel_state import ( + get_sp_parallel_rank, + get_sp_world_size, +) +from fastvideo.fastvideo_args import FastVideoArgs +from fastvideo.forward_context import set_forward_context +from fastvideo.logger import init_logger +from fastvideo.models.loader.component_loader import TransformerLoader +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch +from fastvideo.utils import maybe_download_model + +logger = init_logger(__name__) + +os.environ["MASTER_ADDR"] = "localhost" +os.environ["MASTER_PORT"] = "29505" +os.environ.setdefault("TORCHDYNAMO_DISABLE", "1") + +MODEL_PATH = maybe_download_model("FastVideo/Matrix-Game-2.0-Base-Diffusers") +TRANSFORMER_PATH = os.path.join(MODEL_PATH, "transformer") +REFERENCE_LATENT = -111835.58895772696 + + +@pytest.mark.skip(reason="Not reliably reproducible") +@pytest.mark.usefixtures("distributed_setup") +def test_matrixgame_transformer(): + transformer_path = TRANSFORMER_PATH + + sp_rank = get_sp_parallel_rank() + sp_world_size = get_sp_world_size() + + device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + precision = torch.bfloat16 + precision_str = "bf16" + + args = FastVideoArgs( + model_path=transformer_path, + dit_cpu_offload=False, + use_fsdp_inference=False, + pipeline_config=PipelineConfig( + dit_config=MatrixGameWanVideoConfig(), dit_precision=precision_str + ), + ) + args.device = device + + loader = TransformerLoader() + model = loader.load(transformer_path, args).to(device, dtype=precision) + model.eval() + + total_params = sum(p.numel() for p in model.parameters()) + weight_sum = sum(p.to(torch.float64).sum().item() for p in model.parameters()) + weight_mean = weight_sum / total_params + logger.info("Total parameters: %s", total_params) + logger.info("Weight sum: %s", weight_sum) + logger.info("Weight mean: %s", weight_mean) + + torch.manual_seed(42) + + batch_size = 1 + latent_frames = 15 + latent_height = 44 + latent_width = 80 + image_dim = 1280 + image_seq_len = 257 + + # Set num_frame_per_block to match actual latent_frames for action_model + model.num_frame_per_block = latent_frames + # Force full attention mode + model.local_attn_size = -1 + # Clear cached block masks so they regenerate with correct settings + model.block_mask = None + model.block_mask_mouse = None + model.block_mask_keyboard = None + + # Video latents [B, C, T, H, W] + x = torch.randn(batch_size, 16, latent_frames, latent_height, latent_width, + device=device, dtype=precision) + cond_concat = torch.randn(batch_size, 20, latent_frames, latent_height, latent_width, + device=device, dtype=precision) + hidden_states = torch.cat([x, cond_concat], dim=1) + + if sp_world_size > 1: + chunk_per_rank = hidden_states.shape[2] // sp_world_size + hidden_states = hidden_states[:, :, sp_rank * chunk_per_rank:(sp_rank + 1) * chunk_per_rank] + + # Image embeddings + encoder_hidden_states_image = [ + torch.randn(batch_size, image_seq_len, image_dim, device=device, dtype=precision) + ] + + # Timestep + timestep = torch.full((batch_size, latent_frames), 500, device=device, dtype=precision) + + # Action conditions + num_pixel_frames = (latent_frames - 1) * 4 + 1 + mouse_cond = torch.randn(batch_size, num_pixel_frames, 2, device=device, dtype=precision) + keyboard_cond = torch.randn(batch_size, num_pixel_frames, 4, device=device, dtype=precision) + + # Text embeddings (empty for i2v) + encoder_hidden_states = torch.zeros(batch_size, 0, model.config.hidden_size, + device=device, dtype=precision) + + forward_batch = ForwardBatch(data_type="dummy") + model.num_frame_per_block = latent_frames + + with torch.no_grad(): + with torch.amp.autocast('cuda', dtype=precision): + with set_forward_context(current_timestep=0, attn_metadata=None, forward_batch=forward_batch): + output = model( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + encoder_hidden_states_image=encoder_hidden_states_image, + timestep=timestep, + mouse_cond=mouse_cond, + keyboard_cond=keyboard_cond, + ) + + latent = output.double().sum().item() + + diff = abs(REFERENCE_LATENT - latent) + relative_diff = diff / abs(REFERENCE_LATENT) + logger.info(f"Reference latent: {REFERENCE_LATENT}, Current latent: {latent}") + logger.info(f"Absolute diff: {diff}, Relative diff: {relative_diff * 100:.4f}%") + + # The 0.19% diff now + assert relative_diff < 0.005, f"Output latents differ significantly: relative diff = {relative_diff * 100:.4f}% (max allowed: 0.5%)"