Skip to content

Commit afdb6fb

Browse files
H1yori233SolitaryThinker
H1yori233
and
SolitaryThinker
authored
[feat] Add MatrixGame3.0 (#1201)
Co-authored-by: SolitaryThinker <wlsaidhi@gmail.com>
1 parent ba4c02d commit afdb6fb

25 files changed

Lines changed: 2688 additions & 32 deletions

‎docs/inference/support_matrix.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,7 @@ pipeline initialization and sampling.
7373
| Matrix Game 2.0 Base Distilled | `FastVideo/Matrix-Game-2.0-Base-Distilled-Diffusers` | 352x640 | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
7474
| Matrix Game 2.0 GTA Distilled | `FastVideo/Matrix-Game-2.0-GTA-Distilled-Diffusers` | 352x640 | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
7575
| Matrix Game 2.0 TempleRun Distilled | `FastVideo/Matrix-Game-2.0-TempleRun-Distilled-Diffusers` | 352x640 | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
76+
| Matrix Game 3.0 Base Distilled | `FastVideo/Matrix-Game-3.0-Base-Distilled-Diffusers` | 720x1280 | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
7677
| GEN3C Cosmos 7B | `FastVideo/GEN3C-Cosmos-7B-Diffusers` | 704px1280p | ❌ | ❌ | ❌ | ⭕ | ⭕ |
7778

7879
**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‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
from fastvideo import VideoGenerator
2+
3+
MODEL_PATH = "FastVideo/Matrix-Game-3.0-Base-Distilled-Diffusers"
4+
IMAGE_URL = "https://raw.githubusercontent.com/SkyworkAI/Matrix-Game/main/Matrix-Game-3/demo_images/001/image.png"
5+
PROMPT = "A colorful, animated cityscape with a gas station and various buildings."
6+
OUTPUT_PATH = "video_samples_matrixgame3"
7+
8+
9+
def main():
10+
generator = VideoGenerator.from_pretrained(
11+
MODEL_PATH,
12+
num_gpus=1,
13+
use_fsdp_inference=False,
14+
dit_cpu_offload=False,
15+
vae_cpu_offload=False,
16+
text_encoder_cpu_offload=True,
17+
pin_cpu_memory=True,
18+
)
19+
20+
generator.generate_video(
21+
prompt=PROMPT,
22+
image_path=IMAGE_URL,
23+
height=720,
24+
width=1280,
25+
num_frames=57,
26+
num_inference_steps=3,
27+
guidance_scale=1.0,
28+
seed=42,
29+
output_path=OUTPUT_PATH,
30+
save_video=True,
31+
)
32+
33+
34+
if __name__ == "__main__":
35+
main()

‎fastvideo/api/matrixgame3.py‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
# SPDX-License-Identifier: Apache-2.0
2+
from dataclasses import dataclass
3+
4+
from fastvideo.api.sampling_param import SamplingParam
5+
6+
7+
@dataclass
8+
class MatrixGame3SamplingParam(SamplingParam):
9+
height: int = 720
10+
width: int = 1280
11+
num_frames: int = 57
12+
fps: int = 25
13+
guidance_scale: float = 1.0
14+
num_inference_steps: int = 3
15+
negative_prompt: str = ""
16+
num_iterations: int | None = None
17+
use_base_model: bool = False

‎fastvideo/configs/models/dits/matrixgame3.py‎

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,79 @@
1+
from dataclasses import dataclass, field
2+
3+
import torch
4+
from fastvideo.configs.models.dits.wanvideo import WanVideoArchConfig, WanVideoConfig
5+
6+
7+
def _is_transformer_block(param_name: str, module: torch.nn.Module) -> bool:
8+
return bool("blocks" in param_name and param_name.split(".")[-1].isdigit())
9+
10+
11+
@dataclass
12+
class MatrixGame3WanVideoArchConfig(WanVideoArchConfig):
13+
param_names_mapping: dict = field(
14+
default_factory=lambda: {
15+
r"^patch_embedding\.(weight|bias)$": r"patch_embedding.proj.\1",
16+
r"^patch_embedding_wancamctrl\.(.*)$": r"camera_patch_embedding.proj.\1",
17+
r"^time_embedding\.0\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_in.\1",
18+
r"^time_embedding\.2\.(.*)$": r"condition_embedder.time_embedder.mlp.fc_out.\1",
19+
r"^time_projection\.1\.(.*)$": r"condition_embedder.time_modulation.linear.\1",
20+
r"^head\.head\.(.*)$": r"proj_out.\1",
21+
r"^head\.modulation$": r"scale_shift_table",
22+
r"^blocks\.(\d+)\.self_attn\.q\.(.*)$": r"blocks.\1.to_q.\2",
23+
r"^blocks\.(\d+)\.self_attn\.k\.(.*)$": r"blocks.\1.to_k.\2",
24+
r"^blocks\.(\d+)\.self_attn\.v\.(.*)$": r"blocks.\1.to_v.\2",
25+
r"^blocks\.(\d+)\.self_attn\.o\.(.*)$": r"blocks.\1.to_out.\2",
26+
r"^blocks\.(\d+)\.self_attn\.norm_q\.(.*)$": r"blocks.\1.norm_q.\2",
27+
r"^blocks\.(\d+)\.self_attn\.norm_k\.(.*)$": r"blocks.\1.norm_k.\2",
28+
r"^blocks\.(\d+)\.cross_attn\.q\.(.*)$": r"blocks.\1.attn2.to_q.\2",
29+
r"^blocks\.(\d+)\.cross_attn\.k\.(.*)$": r"blocks.\1.attn2.to_k.\2",
30+
r"^blocks\.(\d+)\.cross_attn\.v\.(.*)$": r"blocks.\1.attn2.to_v.\2",
31+
r"^blocks\.(\d+)\.cross_attn\.o\.(.*)$": r"blocks.\1.attn2.to_out.\2",
32+
r"^blocks\.(\d+)\.cross_attn\.norm_q\.(.*)$": r"blocks.\1.attn2.norm_q.\2",
33+
r"^blocks\.(\d+)\.cross_attn\.norm_k\.(.*)$": r"blocks.\1.attn2.norm_k.\2",
34+
r"^blocks\.(\d+)\.ffn\.0\.(.*)$": r"blocks.\1.ffn.fc_in.\2",
35+
r"^blocks\.(\d+)\.ffn\.2\.(.*)$": r"blocks.\1.ffn.fc_out.\2",
36+
r"^blocks\.(\d+)\.norm3\.(.*)$": r"blocks.\1.self_attn_residual_norm.norm.\2",
37+
r"^blocks\.(\d+)\.modulation$": r"blocks.\1.scale_shift_table",
38+
})
39+
patch_size: tuple[int, int, int] = (1, 2, 2)
40+
in_channels: int = 48
41+
out_channels: int = 48
42+
num_attention_heads: int = 24
43+
attention_head_dim: int = 128
44+
ffn_dim: int = 14336
45+
num_layers: int = 30
46+
text_len: int = 512
47+
image_dim: int = 0
48+
use_text_crossattn: bool = True
49+
use_memory: bool = True
50+
sigma_theta: float = 0.8
51+
camera_embed_in_channels: int = 1536
52+
action_config: dict = field(
53+
default_factory=lambda: {
54+
"blocks": list(range(15)),
55+
"enable_mouse": True,
56+
"enable_keyboard": True,
57+
"heads_num": 16,
58+
"hidden_size": 128,
59+
"img_hidden_size": 3072,
60+
"keyboard_dim_in": 6,
61+
"keyboard_hidden_dim": 1024,
62+
"mouse_dim_in": 2,
63+
"mouse_hidden_dim": 1024,
64+
"mouse_qk_dim_list": [8, 28, 28],
65+
"patch_size": [1, 2, 2],
66+
"qk_norm": True,
67+
"qkv_bias": False,
68+
"rope_dim_list": [8, 28, 28],
69+
"rope_theta": 256,
70+
"vae_time_compression_ratio": 4,
71+
"windows_size": 3,
72+
})
73+
74+
75+
@dataclass
76+
class MatrixGame3WanVideoConfig(WanVideoConfig):
77+
arch_config: MatrixGame3WanVideoArchConfig = field(default_factory=MatrixGame3WanVideoArchConfig)
78+
prefix: str = "Wan"
79+
_compile_conditions: list = field(default_factory=lambda: [_is_transformer_block])

‎fastvideo/configs/models/vaes/wanvae.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,5 +76,7 @@ class WanVAEConfig(VAEConfig):
7676
use_temporal_tiling: bool = False
7777
use_parallel_tiling: bool = False
7878

79+
use_light_vae: bool = False
80+
7981
def __post_init__(self):
8082
self.blend_num_frames = (self.tile_sample_min_num_frames - self.tile_sample_stride_num_frames) * 2

‎fastvideo/configs/pipelines/__init__.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
from fastvideo.configs.pipelines.hunyuangamecraft import HunyuanGameCraftPipelineConfig
77
from fastvideo.configs.pipelines.hyworld import HYWorldConfig
88
from fastvideo.configs.pipelines.matrixgame2 import MatrixGame2I2V480PConfig
9+
from fastvideo.configs.pipelines.matrixgame3 import MatrixGame3I2V720PConfig
910
from fastvideo.pipelines.basic.ltx2.pipeline_configs import LTX2T2VConfig
1011
from fastvideo.registry import get_pipeline_config_cls_from_name
1112
from fastvideo.configs.pipelines.wan import (SelfForcingWanT2V480PConfig, WanI2V480PConfig, WanI2V720PConfig,
@@ -15,5 +16,5 @@
1516
"HunyuanConfig", "FastHunyuanConfig", "HunyuanGameCraftPipelineConfig", "PipelineConfig", "Hunyuan15T2V480PConfig",
1617
"Hunyuan15T2V720PConfig", "WanT2V480PConfig", "WanI2V480PConfig", "WanT2V720PConfig", "WanI2V720PConfig",
1718
"SelfForcingWanT2V480PConfig", "CosmosConfig", "Cosmos25Config", "LTX2T2VConfig", "HYWorldConfig",
18-
"MatrixGame2I2V480PConfig", "get_pipeline_config_cls_from_name"
19+
"MatrixGame2I2V480PConfig", "MatrixGame3I2V720PConfig", "get_pipeline_config_cls_from_name"
1920
]

‎fastvideo/configs/pipelines/matrixgame3.py‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
# SPDX-License-Identifier: Apache-2.0
2+
from dataclasses import dataclass, field
3+
4+
from fastvideo.configs.models import DiTConfig
5+
from fastvideo.configs.models.dits.matrixgame3 import MatrixGame3WanVideoConfig
6+
from fastvideo.configs.pipelines.wan import WanT2V480PConfig
7+
8+
9+
@dataclass
10+
class MatrixGame3I2V720PConfig(WanT2V480PConfig):
11+
dit_config: DiTConfig = field(default_factory=MatrixGame3WanVideoConfig)
12+
flow_shift: float | None = 5.0
13+
vae_precision: str = "fp32"
14+
15+
def __post_init__(self) -> None:
16+
self.vae_config.load_encoder = True
17+
self.vae_config.load_decoder = True
18+
self.vae_config.use_light_vae = True

‎fastvideo/models/dits/matrixgame3/__init__.py‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
from .action_module import MatrixGame3ActionModule
2+
from .model import MatrixGame3CrossAttention, MatrixGame3TransformerBlock, MatrixGame3WanModel
3+
4+
__all__ = [
5+
"MatrixGame3WanModel",
6+
"MatrixGame3TransformerBlock",
7+
"MatrixGame3CrossAttention",
8+
"MatrixGame3ActionModule",
9+
]
10+
11+
EntryClass = [MatrixGame3WanModel]

0 commit comments

Comments
 (0)