[feat] Add MatrixGame3.0

docs/inference/support_matrix.md

@@ -73,6 +73,7 @@ pipeline initialization and sampling.
 | Matrix Game 2.0 Base Distilled | `FastVideo/Matrix-Game-2.0-Base-Distilled-Diffusers` | 352x640 | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
 | Matrix Game 2.0 GTA Distilled | `FastVideo/Matrix-Game-2.0-GTA-Distilled-Diffusers` | 352x640 | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
 | Matrix Game 2.0 TempleRun Distilled | `FastVideo/Matrix-Game-2.0-TempleRun-Distilled-Diffusers` | 352x640 | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
+| Matrix Game 3.0 Base Distilled | `FastVideo/Matrix-Game-3.0-Base-Distilled-Diffusers` | 720x1280 | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
 | GEN3C Cosmos 7B | `FastVideo/GEN3C-Cosmos-7B-Diffusers` | 704px1280p | ❌ | ❌ | ❌ | ⭕ | ⭕ |
 
 **Note**: Wan2.2 TI2V 5B has some quality issues when performing I2V generation. We are working on fixing this issue.

examples/inference/basic/basic_matrixgame3.py

@@ -0,0 +1,35 @@
+from fastvideo import VideoGenerator
+
+MODEL_PATH = "FastVideo/Matrix-Game-3.0-Base-Distilled-Diffusers"
+IMAGE_URL = "https://raw.githubusercontent.com/SkyworkAI/Matrix-Game/main/Matrix-Game-3/demo_images/001/image.png"
+PROMPT = "A colorful, animated cityscape with a gas station and various buildings."
+OUTPUT_PATH = "video_samples_matrixgame3"
+
+
+def main():
+    generator = VideoGenerator.from_pretrained(
+        MODEL_PATH,
+        num_gpus=1,
+        use_fsdp_inference=False,
+        dit_cpu_offload=False,
+        vae_cpu_offload=False,
+        text_encoder_cpu_offload=True,
+        pin_cpu_memory=True,
+    )
+
+    generator.generate_video(
+        prompt=PROMPT,
+        image_path=IMAGE_URL,
+        height=720,
+        width=1280,
+        num_frames=57,
+        num_inference_steps=3,
+        guidance_scale=1.0,
+        seed=42,
+        output_path=OUTPUT_PATH,
+        save_video=True,
+    )
+
+
+if __name__ == "__main__":
+    main()

fastvideo/api/matrixgame3.py

@@ -0,0 +1,17 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass
+
+from fastvideo.api.sampling_param import SamplingParam
+
+
+@dataclass
+class MatrixGame3SamplingParam(SamplingParam):
+    height: int = 720
+    width: int = 1280
+    num_frames: int = 57
+    fps: int = 25
+    guidance_scale: float = 1.0
+    num_inference_steps: int = 3
+    negative_prompt: str = ""
+    num_iterations: int | None = None
+    use_base_model: bool = False

fastvideo/configs/models/dits/matrixgame3.py

@@ -0,0 +1,79 @@
+from dataclasses import dataclass, field
+
+import torch
+from fastvideo.configs.models.dits.wanvideo import WanVideoArchConfig, WanVideoConfig
+
+
+def _is_transformer_block(param_name: str, module: torch.nn.Module) -> bool:
+    return bool("blocks" in param_name and param_name.split(".")[-1].isdigit())
+
+
+@dataclass
+class MatrixGame3WanVideoArchConfig(WanVideoArchConfig):
+    param_names_mapping: dict = field(
+        default_factory=lambda: {
+            r"^patch_embedding\.(weight|bias)$": r"patch_embedding.proj.\1",
+            r"^patch_embedding_wancamctrl\.(.*)$": r"camera_patch_embedding.proj.\1",
+            r"^time_embedding\.0\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_in.\1",
+            r"^time_embedding\.2\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_out.\1",
+            r"^time_projection\.1\.(.*)$": r"condition_embedder.time_modulation.linear.\1",
+            r"^head\.head\.(.*)$": r"proj_out.\1",
+            r"^head\.modulation$": r"scale_shift_table",
+            r"^blocks\.(\d+)\.self_attn\.q\.(.*)$": r"blocks.\1.to_q.\2",
+            r"^blocks\.(\d+)\.self_attn\.k\.(.*)$": r"blocks.\1.to_k.\2",
+            r"^blocks\.(\d+)\.self_attn\.v\.(.*)$": r"blocks.\1.to_v.\2",
+            r"^blocks\.(\d+)\.self_attn\.o\.(.*)$": r"blocks.\1.to_out.\2",
+            r"^blocks\.(\d+)\.self_attn\.norm_q\.(.*)$": r"blocks.\1.norm_q.\2",
+            r"^blocks\.(\d+)\.self_attn\.norm_k\.(.*)$": r"blocks.\1.norm_k.\2",
+            r"^blocks\.(\d+)\.cross_attn\.q\.(.*)$": r"blocks.\1.attn2.to_q.\2",
+            r"^blocks\.(\d+)\.cross_attn\.k\.(.*)$": r"blocks.\1.attn2.to_k.\2",
+            r"^blocks\.(\d+)\.cross_attn\.v\.(.*)$": r"blocks.\1.attn2.to_v.\2",
+            r"^blocks\.(\d+)\.cross_attn\.o\.(.*)$": r"blocks.\1.attn2.to_out.\2",
+            r"^blocks\.(\d+)\.cross_attn\.norm_q\.(.*)$": r"blocks.\1.attn2.norm_q.\2",
+            r"^blocks\.(\d+)\.cross_attn\.norm_k\.(.*)$": r"blocks.\1.attn2.norm_k.\2",
+            r"^blocks\.(\d+)\.ffn\.0\.(.*)$": r"blocks.\1.ffn.fc_in.\2",
+            r"^blocks\.(\d+)\.ffn\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2",
+            r"^blocks\.(\d+)\.norm3\.(.*)$": r"blocks.\1.self_attn_residual_norm.norm.\2",
+            r"^blocks\.(\d+)\.modulation$": r"blocks.\1.scale_shift_table",
+        })
+    patch_size: tuple[int, int, int] = (1, 2, 2)
+    in_channels: int = 48
+    out_channels: int = 48
+    num_attention_heads: int = 24
+    attention_head_dim: int = 128
+    ffn_dim: int = 14336
+    num_layers: int = 30
+    text_len: int = 512
+    image_dim: int = 0
+    use_text_crossattn: bool = True
+    use_memory: bool = True
+    sigma_theta: float = 0.8
+    camera_embed_in_channels: int = 1536
+    action_config: dict = field(
+        default_factory=lambda: {
+            "blocks": list(range(15)),
+            "enable_mouse": True,
+            "enable_keyboard": True,
+            "heads_num": 16,
+            "hidden_size": 128,
+            "img_hidden_size": 3072,
+            "keyboard_dim_in": 6,
+            "keyboard_hidden_dim": 1024,
+            "mouse_dim_in": 2,
+            "mouse_hidden_dim": 1024,
+            "mouse_qk_dim_list": [8, 28, 28],
+            "patch_size": [1, 2, 2],
+            "qk_norm": True,
+            "qkv_bias": False,
+            "rope_dim_list": [8, 28, 28],
+            "rope_theta": 256,
+            "vae_time_compression_ratio": 4,
+            "windows_size": 3,
+        })
+
+
+@dataclass
+class MatrixGame3WanVideoConfig(WanVideoConfig):
+    arch_config: MatrixGame3WanVideoArchConfig = field(default_factory=MatrixGame3WanVideoArchConfig)
+    prefix: str = "Wan"
+    _compile_conditions: list = field(default_factory=lambda: [_is_transformer_block])

fastvideo/configs/models/vaes/wanvae.py

@@ -76,5 +76,7 @@ class WanVAEConfig(VAEConfig):
     use_temporal_tiling: bool = False
     use_parallel_tiling: bool = False
 
+    use_light_vae: bool = False
+
     def __post_init__(self):
         self.blend_num_frames = (self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames) * 2

fastvideo/configs/pipelines/__init__.py

@@ -6,6 +6,7 @@
 from fastvideo.configs.pipelines.hunyuangamecraft import HunyuanGameCraftPipelineConfig
 from fastvideo.configs.pipelines.hyworld import HYWorldConfig
 from fastvideo.configs.pipelines.matrixgame2 import MatrixGame2I2V480PConfig
+from fastvideo.configs.pipelines.matrixgame3 import MatrixGame3I2V720PConfig
 from fastvideo.pipelines.basic.ltx2.pipeline_configs import LTX2T2VConfig
 from fastvideo.registry import get_pipeline_config_cls_from_name
 from fastvideo.configs.pipelines.wan import (SelfForcingWanT2V480PConfig, WanI2V480PConfig, WanI2V720PConfig,
@@ -15,5 +16,5 @@
     "HunyuanConfig", "FastHunyuanConfig", "HunyuanGameCraftPipelineConfig", "PipelineConfig", "Hunyuan15T2V480PConfig",
     "Hunyuan15T2V720PConfig", "WanT2V480PConfig", "WanI2V480PConfig", "WanT2V720PConfig", "WanI2V720PConfig",
     "SelfForcingWanT2V480PConfig", "CosmosConfig", "Cosmos25Config", "LTX2T2VConfig", "HYWorldConfig",
-    "MatrixGame2I2V480PConfig", "get_pipeline_config_cls_from_name"
+    "MatrixGame2I2V480PConfig", "MatrixGame3I2V720PConfig", "get_pipeline_config_cls_from_name"
 ]

fastvideo/configs/pipelines/matrixgame3.py

@@ -0,0 +1,18 @@
+# SPDX-License-Identifier: Apache-2.0
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models import DiTConfig
+from fastvideo.configs.models.dits.matrixgame3 import MatrixGame3WanVideoConfig
+from fastvideo.configs.pipelines.wan import WanT2V480PConfig
+
+
+@dataclass
+class MatrixGame3I2V720PConfig(WanT2V480PConfig):
+    dit_config: DiTConfig = field(default_factory=MatrixGame3WanVideoConfig)
+    flow_shift: float | None = 5.0
+    vae_precision: str = "fp32"
+
+    def __post_init__(self) -> None:
+        self.vae_config.load_encoder = True
+        self.vae_config.load_decoder = True
+        self.vae_config.use_light_vae = True

fastvideo/models/dits/matrixgame3/__init__.py

@@ -0,0 +1,11 @@
+from .action_module import MatrixGame3ActionModule
+from .model import MatrixGame3CrossAttention, MatrixGame3TransformerBlock, MatrixGame3WanModel
+
+__all__ = [
+    "MatrixGame3WanModel",
+    "MatrixGame3TransformerBlock",
+    "MatrixGame3CrossAttention",
+    "MatrixGame3ActionModule",
+]
+
+EntryClass = [MatrixGame3WanModel]