matrixgame2

fastvideo.models.dits.matrixgame2.MatrixGame2WanModel

MatrixGame2WanModel(config: MatrixGame2WanVideoConfig, hf_config: dict[str, Any], **kwargs)

fastvideo.models.dits.matrixgame2.MatrixGame2WanModel.forward

forward(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

fastvideo.models.dits.matrixgame2.MatrixGame2TransformerBlock

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

fastvideo.models.dits.matrixgame2.MatrixGame2TransformerBlock.forward

forward(hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, freqs_cis: tuple[torch.Tensor, torch.Tensor], grid_sizes: torch.Tensor, mouse_cond: torch.Tensor | None = None, keyboard_cond: torch.Tensor | None = None) -> torch.Tensor

fastvideo.models.dits.matrixgame2.CausalMatrixGame2WanModel

CausalMatrixGame2WanModel(config: MatrixGame2WanVideoConfig, hf_config: dict[str, Any], **kwargs)

fastvideo.models.dits.matrixgame2.CausalMatrixGame2WanModel.set_causal_attention_window

set_causal_attention_window(*, local_attn_size: int | None = None, sink_size: int | None = None) -> None

fastvideo.models.dits.matrixgame2.CausalMatrixGame2WanModel.set_sink_size

set_sink_size(sink_size: int) -> None

fastvideo.models.dits.matrixgame2.CausalMatrixGame2WanModel.forward

forward(*args, **kwargs)

fastvideo.models.dits.matrixgame2.CausalMatrixGame2WanModel.unpatchify

unpatchify(x: torch.Tensor, grid_sizes: tuple[int, int, int]) -> torch.Tensor

fastvideo.models.dits.matrixgame2.CausalMatrixGame2WanModel.init_weights

init_weights()

fastvideo.models.dits.matrixgame2.CausalMatrixGame2TransformerBlock

CausalMatrixGame2TransformerBlock(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-06, 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, rope_cache_policy: str = 'absolute')

fastvideo.models.dits.matrixgame2.CausalMatrixGame2TransformerBlock.forward

forward(hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor, freqs_cis: tuple[torch.Tensor, torch.Tensor], block_mask: BlockMask, grid_sizes: tuple[int, int, int], 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

fastvideo.models.dits.matrixgame2.ActionModule

ActionModule(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, sink_size = 0, blocks: list | None = None)

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

fastvideo.models.dits.matrixgame2.ActionModule.patchify

patchify(x, patch_size)

x : (N C T H W)

fastvideo.models.dits.matrixgame2.ActionModule.unpatchify

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

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

fastvideo.models.dits.matrixgame2.ActionModule.get_rotary_pos_embed

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

fastvideo.models.dits.matrixgame2.ActionModule.forward

forward(x: torch.Tensor, tt: int, th: int, tw: int, mouse_condition: torch.Tensor | None = None, keyboard_condition: torch.Tensor | None = None, block_mask_mouse: BlockMask | None = None, block_mask_keyboard: BlockMask | None = None, is_causal: bool = False, kv_cache_mouse: dict[str, torch.Tensor] | None = None, kv_cache_keyboard: dict[str, torch.Tensor] | None = None, start_frame: int = 0, use_rope_keyboard: bool = True, num_frame_per_block: int = 3)

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