matrixgame3

fastvideo.models.dits.matrixgame3.MatrixGame3ActionModule

MatrixGame3ActionModule(mouse_dim_in: int = 2, keyboard_dim_in: int = 6, hidden_size: int = 128, img_hidden_size: int = 1536, keyboard_hidden_dim: int = 1024, mouse_hidden_dim: int = 1024, vae_time_compression_ratio: int = 4, windows_size: int = 3, heads_num: int = 16, patch_size: list | None = None, qk_norm: bool = True, qkv_bias: bool = False, rope_dim_list: list | None = None, rope_theta = 256, mouse_qk_dim_list: list | None = None, enable_mouse = True, enable_keyboard = True, blocks: list | None = None)

action module from https://arxiv.org/pdf/2501.08325

fastvideo.models.dits.matrixgame3.MatrixGame3ActionModule.patchify

patchify(x, patch_size)

x : (N C T H W)

fastvideo.models.dits.matrixgame3.MatrixGame3ActionModule.unpatchify

unpatchify(x, t, h, w, patch_size)

x: (N, T, patch_size**2 * C) imgs: (N, H, W, C)

fastvideo.models.dits.matrixgame3.MatrixGame3ActionModule.get_rotary_pos_embed

get_rotary_pos_embed(video_length, height, width, head_dim, rope_dim_list = None)

fastvideo.models.dits.matrixgame3.MatrixGame3ActionModule.forward

forward(x: torch.Tensor, tt: int, th: int, tw: int, mouse_condition: torch.Tensor | None = None, keyboard_condition: torch.Tensor | None = None, mouse_cond_memory: torch.Tensor | None = None, keyboard_cond_memory: torch.Tensor | None = None)

hidden_states: B, ttthtw, C mouse_condition: B, N_frames, C1 keyboard_condition: B, N_frames, C2

fastvideo.models.dits.matrixgame3.MatrixGame3CrossAttention

MatrixGame3CrossAttention(dim: int, num_heads: int, window_size = (-1, -1), qk_norm = True, eps = 1e-06, parallel_attention = False, quant_config: QuantizationConfig | None = None, prefix: str = '')

fastvideo.models.dits.matrixgame3.MatrixGame3CrossAttention.forward

forward(x, context, context_lens = None)

Parameters

x: Tensor

Shape [B, L1, C]

context: Tensor

Shape [B, L2, C] - typically 257 image tokens

context_lens: Tensor

Shape [B]

fastvideo.models.dits.matrixgame3.MatrixGame3TransformerBlock

MatrixGame3TransformerBlock(dim: int, ffn_dim: int, num_heads: int, qk_norm: str = 'rms_norm_across_heads', cross_attn_norm: bool = False, eps: float = 1e-06, supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, prefix: str = '', action_config: dict | None = None, block_id: int | None = None, use_memory: bool = False)

fastvideo.models.dits.matrixgame3.MatrixGame3TransformerBlock.forward

forward(hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, freqs: torch.Tensor, grid_sizes: tuple[int, int, int], mouse_cond: torch.Tensor | None = None, keyboard_cond: torch.Tensor | None = None, mouse_cond_memory: torch.Tensor | None = None, keyboard_cond_memory: torch.Tensor | None = None, plucker_emb: torch.Tensor | None = None, memory_length: int = 0, memory_latent_idx: list[int] | tuple[int, ...] | torch.Tensor | None = None, predict_latent_idx: tuple[int, int] | list[int] | tuple[int, ...] | torch.Tensor | None = None) -> torch.Tensor

fastvideo.models.dits.matrixgame3.MatrixGame3WanModel

MatrixGame3WanModel(config: MatrixGame3WanVideoConfig, hf_config: dict[str, Any], **kwargs)

fastvideo.models.dits.matrixgame3.MatrixGame3WanModel.forward

forward(hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | list[torch.Tensor], timestep: torch.LongTensor, mouse_cond: torch.Tensor | None = None, keyboard_cond: torch.Tensor | None = None, x_memory: torch.Tensor | None = None, timestep_memory: torch.Tensor | None = None, mouse_cond_memory: torch.Tensor | None = None, keyboard_cond_memory: torch.Tensor | None = None, c2ws_plucker_emb: torch.Tensor | None = None, memory_latent_idx: list[int] | tuple[int, ...] | torch.Tensor | None = None, predict_latent_idx: tuple[int, int] | list[int] | tuple[int, ...] | torch.Tensor | None = None, **kwargs) -> torch.Tensor