Skip to content
FastVideo
Esc
↑↓navigate↵open⌘Jpreview
On this page

action_module

fastvideo.models.dits.matrixgame3.action_module.LocalAttention

LocalAttention(num_heads: int, head_size: int, num_kv_heads: int | None = None, softmax_scale: float | None = None, causal: bool = False, supported_attention_backends: tuple[AttentionBackendEnum, ...] | None = None, default_backend: AttentionBackendEnum | None = None, **extra_impl_args)

Attention layer.

fastvideo.models.dits.matrixgame3.action_module.LocalAttention.forward

forward(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, freqs_cis: tuple[torch.Tensor, torch.Tensor] | None = None) -> torch.Tensor

Apply local attention between query, key and value tensors.

Parameters

q: torch.Tensor

Query tensor of shape [batch_size, seq_len, num_heads, head_dim]

k: torch.Tensor

Key tensor of shape [batch_size, seq_len, num_heads, head_dim]

v: torch.Tensor

Value tensor of shape [batch_size, seq_len, num_heads, head_dim]

Returns

torch.Tensor

torch.Tensor: Output tensor after local attention

fastvideo.models.dits.matrixgame3.action_module.ReplicatedLinear

ReplicatedLinear(input_size: int, output_size: int, bias: bool = True, skip_bias_add: bool = False, params_dtype: torch.dtype | None = None, quant_config: QuantizationConfig | None = None, prefix: str = '')

Replicated linear layer.

Parameters

input_size: int

input dimension of the linear layer.

output_size: int

output dimension of the linear layer.

bias: bool

If true, add bias.

skip_bias_add: bool

If true, skip adding bias but instead return it.

params_dtype: torch.dtype | None

Data type for the parameters.

quant_config: QuantizationConfig | None

Quantization configure.

prefix: str

The name of the layer in the state dict, including all parents (e.g. model.layers.0.qkv_proj)

fastvideo.models.dits.matrixgame3.action_module.ReplicatedLinear.weight_loader

weight_loader(param: Parameter, loaded_weight: torch.Tensor) -> None

fastvideo.models.dits.matrixgame3.action_module.ReplicatedLinear.forward

forward(x: torch.Tensor) -> tuple[torch.Tensor, Parameter | None]

fastvideo.models.dits.matrixgame3.action_module.ReplicatedLinear.extra_repr

extra_repr() -> str

fastvideo.models.dits.matrixgame3.action_module.ReplicatedLinear.get_shape_mapping

get_shape_mapping() -> dict

Get the mapping from (input_shape, output_shape) to layer types.

fastvideo.models.dits.matrixgame3.action_module.ReplicatedLinear.reset_shape_tracking

reset_shape_tracking() -> None

Clear tracked shapes and layer type mappings.

fastvideo.models.dits.matrixgame3.action_module.ReplicatedLinear.print_shape_summary

print_shape_summary() -> None

Log a summary of all unique shapes and their layer types.

fastvideo.models.dits.matrixgame3.action_module.FP32LayerNorm

FP32LayerNorm()

fastvideo.models.dits.matrixgame3.action_module.FP32LayerNorm.forward

forward(inputs: torch.Tensor) -> torch.Tensor

fastvideo.models.dits.matrixgame3.action_module.AttentionBackendEnum

AttentionBackendEnum()

fastvideo.models.dits.matrixgame3.action_module.WanRMSNorm

WanRMSNorm(dim: int, eps: float = 1e-05)
Source
class WanRMSNorm(nn.Module):

def __init__(self, dim: int, eps: float = 1e-5):
    super().__init__()
    self.dim = dim
    self.eps = eps
    self.weight = nn.Parameter(torch.ones(dim))

def forward(self, x: torch.Tensor) -> torch.Tensor:
    return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)</code></pre></details>


fastvideo.models.dits.matrixgame3.action_module.WanRMSNorm.forward


forward(x: torch.Tensor) -> torch.Tensor


Source
def forward(self, x: torch.Tensor) -> torch.Tensor:
    return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)


fastvideo.models.dits.matrixgame3.action_module.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


Source
class MatrixGame3ActionModule(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,
    blocks: list | None = None,
):
    super().__init__()
    patch_size = patch_size if patch_size is not None else &#91;1, 2, 2&#93;
    rope_dim_list = (rope_dim_list if rope_dim_list is not None else &#91;8, 28, 28&#93;)
    mouse_qk_dim_list = (mouse_qk_dim_list if mouse_qk_dim_list is not None else &#91;8, 28, 28&#93;)
    blocks = blocks if blocks is not None else &#91;&#93;
    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=&quot;tanh&quot;),
            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 = (WanRMSNorm(head_dim, eps=1e-6) if qk_norm else nn.Identity())
        self.img_attn_k_norm = (WanRMSNorm(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 = (WanRMSNorm(head_dim_key, eps=1e-6) if qk_norm else nn.Identity())
        self.key_attn_k_norm = (WanRMSNorm(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):
    &quot;&quot;&quot;
    x : (N C T H W)
    &quot;&quot;&quot;
    pt, ph, pw = self.patch_size
    t, h, w = x.shape&#91;2&#93; // pt, x.shape&#91;3&#93; // ph, x.shape&#91;4&#93; // pw
    c = x.shape&#91;1&#93;
    x = x.reshape(shape=(x.shape&#91;0&#93;, c, t, pt, h, ph, w, pw))
    x = torch.einsum(&quot;nctohpwq-&gt;nthwcopq&quot;, x)
    x = x.reshape(shape=(x.shape&#91;0&#93;, t * h * w, c * pt * ph * pw))
    return x

def unpatchify(self, x, t, h, w, patch_size):
    &quot;&quot;&quot;
    x: (N, T, patch_size**2 * C)
    imgs: (N, H, W, C)
    &quot;&quot;&quot;
    c = x.shape&#91;2&#93; // patch_size
    pt, ph, pw = self.patch_size
    assert t * h * w == x.shape&#91;1&#93;

    x = x.reshape(shape=(x.shape&#91;0&#93;, t, h, w, c, pt, ph, pw))
    x = torch.einsum(&quot;nthwcopq-&gt;nctohpwq&quot;, x)
    imgs = x.reshape(shape=(x.shape&#91;0&#93;, 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,
):
    target_ndim = 3
    ndim = 5 - 2

    latents_size = &#91;video_length, height, width&#93;

    if isinstance(self.patch_size, int):
        assert all(s % self.patch_size == 0 for s in latents_size), (
            f&quot;Latent size(last &#123;ndim&#125; dimensions) should be divisible by patch size(&#123;self.patch_size&#125;), &quot;
            f&quot;but got &#123;latents_size&#125;.&quot;)
        rope_sizes = &#91;s // self.patch_size for s in latents_size&#93;
    elif isinstance(self.patch_size, list):
        assert all(s % self.patch_size&#91;idx&#93; == 0 for idx, s in enumerate(latents_size)), (
            f&quot;Latent size(last &#123;ndim&#125; dimensions) should be divisible by patch size(&#123;self.patch_size&#125;), &quot;
            f&quot;but got &#123;latents_size&#125;.&quot;)
        rope_sizes = &#91;s // self.patch_size&#91;idx&#93; for idx, s in enumerate(latents_size)&#93;

    if len(rope_sizes) != target_ndim:
        rope_sizes = &#91;1&#93; * (target_ndim - len(rope_sizes)) + rope_sizes  # time axis

    if rope_dim_list is None:
        rope_dim_list = &#91;head_dim // target_ndim for _ in range(target_ndim)&#93;
    assert sum(rope_dim_list) == head_dim, (&quot;sum(rope_dim_list) should equal to head_dim of attention layer&quot;)
    # Use Matrix-Game wrapper for FastVideo&#x27;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&#91;-video_length * rope_sizes&#91;1&#93; * rope_sizes&#91;2&#93; //
                     self.patch_size&#91;0&#93;:&#93;, freqs_sin&#91;-video_length * rope_sizes&#91;1&#93; * rope_sizes&#91;2&#93; //
                                                     self.patch_size&#91;0&#93;:&#93;

def _forward_mouse(
    self,
    hidden_states: torch.Tensor,
    mouse_condition: torch.Tensor,
    mouse_cond_memory: torch.Tensor | None,
    *,
    pad_t: int,
    freqs_cis: tuple&#91;torch.Tensor, torch.Tensor&#93;,
    N_feats: int,
    B: int,
    C: int,
    th: int,
    tw: int,
) -&gt; torch.Tensor:
    pad = mouse_condition&#91;:, 0:1, :&#93;.expand(-1, pad_t, -1)
    mouse_condition = torch.cat(&#91;pad, mouse_condition&#93;, dim=1)
    group_mouse = &#91;
        mouse_condition&#91;
            :,
            self.vae_time_compression_ratio * (i - self.windows_size) + pad_t:i * self.vae_time_compression_ratio +
            pad_t,
            :,
        &#93; for i in range(N_feats)
    &#93;
    group_mouse = torch.stack(group_mouse, dim=1)
    mem_len = 0
    if mouse_cond_memory is not None:
        mem_len = mouse_cond_memory.shape&#91;1&#93;
        mem = mouse_cond_memory.to(device=group_mouse.device, dtype=group_mouse.dtype)
        mem = mem.unsqueeze(2).expand(-1, -1, pad_t, -1)
        group_mouse = torch.cat(&#91;mem, group_mouse&#93;, dim=1)
    actual_num_frames = group_mouse.shape&#91;1&#93;

    S = th * tw
    group_mouse = group_mouse.unsqueeze(-1).expand(B, actual_num_frames, pad_t, C, S)
    group_mouse = group_mouse.permute(0, 4, 1, 2, 3).reshape(B * S, actual_num_frames, pad_t * C)

    group_mouse = torch.cat(&#91;hidden_states, group_mouse&#93;, dim=-1)
    group_mouse = self.mouse_mlp(group_mouse)
    mouse_qkv, _ = self.t_qkv(group_mouse)
    q, k, v = rearrange(mouse_qkv, &quot;B L (K H D) -&gt; K B L H D&quot;, K=3, H=self.heads_num)
    q = self.img_attn_q_norm(q).to(v)
    k = self.img_attn_k_norm(k).to(v)
    if mem_len &gt; 0:
        q_mem, k_mem = _apply_rotary_emb_qk(q&#91;:, :mem_len&#93;, k&#91;:, :mem_len&#93;, freqs_cis&#91;0&#93;, freqs_cis&#91;1&#93;)
        q_pred, k_pred = _apply_rotary_emb_qk(q&#91;:, mem_len:&#93;, k&#91;:, mem_len:&#93;, freqs_cis&#91;0&#93;, freqs_cis&#91;1&#93;)
        q = torch.cat(&#91;q_mem, q_pred&#93;, dim=1)
        k = torch.cat(&#91;k_mem, k_pred&#93;, dim=1)
    else:
        q, k = _apply_rotary_emb_qk(q, k, freqs_cis&#91;0&#93;, freqs_cis&#91;1&#93;)

    attn = self.mouse_attn_layer(q, k, v)
    attn = rearrange(attn, &quot;(b S) T h d -&gt; b (T S) (h d)&quot;, b=B)
    attn, _ = self.proj_mouse(attn)
    return attn

def _forward_keyboard(
    self,
    hidden_states: torch.Tensor,
    keyboard_condition: torch.Tensor,
    keyboard_cond_memory: torch.Tensor | None,
    *,
    pad_t: int,
    freqs_cis: tuple&#91;torch.Tensor, torch.Tensor&#93;,
    N_feats: int,
    B: int,
    th: int,
    tw: int,
) -&gt; torch.Tensor:
    pad = keyboard_condition&#91;:, 0:1, :&#93;.expand(-1, pad_t, -1)
    keyboard_condition = torch.cat(&#91;pad, keyboard_condition&#93;, dim=1)
    keyboard_condition = self.keyboard_embed(keyboard_condition)
    group_keyboard = &#91;
        keyboard_condition&#91;
            :,
            self.vae_time_compression_ratio * (i - self.windows_size) + pad_t:i * self.vae_time_compression_ratio +
            pad_t,
            :,
        &#93; for i in range(N_feats)
    &#93;
    group_keyboard = torch.stack(group_keyboard, dim=1)
    mem_len = 0
    if keyboard_cond_memory is not None:
        mem_len = keyboard_cond_memory.shape&#91;1&#93;
        k_mem = keyboard_cond_memory.to(device=group_keyboard.device, dtype=group_keyboard.dtype)
        k_mem = self.keyboard_embed(k_mem)
        k_mem = k_mem.unsqueeze(2).expand(-1, -1, pad_t, -1)
        group_keyboard = torch.cat(&#91;k_mem, group_keyboard&#93;, dim=1)
    group_keyboard = group_keyboard.reshape(shape=(group_keyboard.shape&#91;0&#93;, group_keyboard.shape&#91;1&#93;, -1))

    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)

    q = self.key_attn_q_norm(q).to(v)
    k = self.key_attn_k_norm(k).to(v)
    S = th * tw

    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)
    if mem_len &gt; 0:
        q_mem, k_mem_r = _apply_rotary_emb_qk(q&#91;:, :mem_len&#93;, k&#91;:, :mem_len&#93;, freqs_cis&#91;0&#93;, freqs_cis&#91;1&#93;)
        q_pred, k_pred_r = _apply_rotary_emb_qk(q&#91;:, mem_len:&#93;, k&#91;:, mem_len:&#93;, freqs_cis&#91;0&#93;, freqs_cis&#91;1&#93;)
        q = torch.cat(&#91;q_mem, q_pred&#93;, dim=1)
        k = torch.cat(&#91;k_mem_r, k_pred_r&#93;, dim=1)
    else:
        q, k = _apply_rotary_emb_qk(q, k, freqs_cis&#91;0&#93;, freqs_cis&#91;1&#93;)

    k = k.repeat_interleave(S, dim=0)
    v = v.repeat_interleave(S, dim=0)

    attn = self.keyboard_attn_layer(q, k, v)
    attn = rearrange(attn, &quot;(B S) T H D -&gt; B (T S) (H D)&quot;, S=S)
    attn, _ = self.proj_keyboard(attn)
    return attn

def forward(
    self,
    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,
):
    &quot;&quot;&quot;
    hidden_states: B, tt*th*tw, C
    mouse_condition: B, N_frames, C1
    keyboard_condition: B, N_frames, C2
    &quot;&quot;&quot;
    target_device = x.device
    target_dtype = x.dtype
    if mouse_condition is not None:
        mouse_condition = mouse_condition.to(device=target_device, dtype=target_dtype)
    if keyboard_condition is not None:
        keyboard_condition = keyboard_condition.to(device=target_device, dtype=target_dtype)
    else:
        return x

    B, N_frames, C = keyboard_condition.shape
    assert tt * th * tw == x.shape&#91;1&#93;
    assert (((N_frames - 1) + self.vae_time_compression_ratio) % self.vae_time_compression_ratio == 0
            or (N_frames % self.vae_time_compression_ratio == 0))
    if ((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
    else:
        N_feats = N_frames // self.vae_time_compression_ratio

    # Lazy initialization of freqs
    if (self._freqs_cos is None or self._freqs_sin is None or self._freqs_cos.device != x.device):
        fc, fs = self.get_rotary_pos_embed(
            7500,
            self.patch_size&#91;1&#93;,
            self.patch_size&#91;2&#93;,
            64,
            self.mouse_qk_dim_list,
        )
        self._freqs_cos = fc.to(x.device)
        self._freqs_sin = fs.to(x.device)

    freqs_cis = (self._freqs_cos, self._freqs_sin)

    if self.enable_mouse and mouse_condition is not None:
        hidden_states = rearrange(x, &quot;B (T S) C -&gt; (B S) T C&quot;, T=tt, S=th * tw)
        B, N_frames, C = mouse_condition.shape
    else:
        hidden_states = x

    pad_t = self.vae_time_compression_ratio * self.windows_size
    if self.enable_mouse and mouse_condition is not None:
        attn = self._forward_mouse(
            hidden_states,
            mouse_condition,
            mouse_cond_memory,
            pad_t=pad_t,
            freqs_cis=freqs_cis,
            N_feats=N_feats,
            B=B,
            C=C,
            th=th,
            tw=tw,
        )
        hidden_states = rearrange(x, &quot;(B S) T C -&gt; B (T S) C&quot;, B=B)
        hidden_states = hidden_states + attn

    if self.enable_keyboard and keyboard_condition is not None:
        attn = self._forward_keyboard(
            hidden_states,
            keyboard_condition,
            keyboard_cond_memory,
            pad_t=pad_t,
            freqs_cis=freqs_cis,
            N_feats=N_feats,
            B=B,
            th=th,
            tw=tw,
        )
        hidden_states = hidden_states + attn
    return hidden_states</code></pre></details>


fastvideo.models.dits.matrixgame3.action_module.MatrixGame3ActionModule.patchify


patchify(x, patch_size)


x : (N C T H W)


Source
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


fastvideo.models.dits.matrixgame3.action_module.MatrixGame3ActionModule.unpatchify


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


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


Source
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
    pt, ph, pw = self.patch_size
    assert t * h * w == x.shape[1]

x = x.reshape(shape=(x.shape&#91;0&#93;, t, h, w, c, pt, ph, pw))
x = torch.einsum(&quot;nthwcopq-&gt;nctohpwq&quot;, x)
imgs = x.reshape(shape=(x.shape&#91;0&#93;, c, t * pt, h * ph, w * pw))

return imgs</code></pre></details>


fastvideo.models.dits.matrixgame3.action_module.MatrixGame3ActionModule.get_rotary_pos_embed


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


Source
def get_rotary_pos_embed(
    self,
    video_length,
    height,
    width,
    head_dim,
    rope_dim_list=None,
):
    target_ndim = 3
    ndim = 5 - 2

latents_size = &#91;video_length, height, width&#93;

if isinstance(self.patch_size, int):
    assert all(s % self.patch_size == 0 for s in latents_size), (
        f&quot;Latent size(last &#123;ndim&#125; dimensions) should be divisible by patch size(&#123;self.patch_size&#125;), &quot;
        f&quot;but got &#123;latents_size&#125;.&quot;)
    rope_sizes = &#91;s // self.patch_size for s in latents_size&#93;
elif isinstance(self.patch_size, list):
    assert all(s % self.patch_size&#91;idx&#93; == 0 for idx, s in enumerate(latents_size)), (
        f&quot;Latent size(last &#123;ndim&#125; dimensions) should be divisible by patch size(&#123;self.patch_size&#125;), &quot;
        f&quot;but got &#123;latents_size&#125;.&quot;)
    rope_sizes = &#91;s // self.patch_size&#91;idx&#93; for idx, s in enumerate(latents_size)&#93;

if len(rope_sizes) != target_ndim:
    rope_sizes = &#91;1&#93; * (target_ndim - len(rope_sizes)) + rope_sizes  # time axis

if rope_dim_list is None:
    rope_dim_list = &#91;head_dim // target_ndim for _ in range(target_ndim)&#93;
assert sum(rope_dim_list) == head_dim, (&quot;sum(rope_dim_list) should equal to head_dim of attention layer&quot;)
# Use Matrix-Game wrapper for FastVideo&#x27;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&#91;-video_length * rope_sizes&#91;1&#93; * rope_sizes&#91;2&#93; //
                 self.patch_size&#91;0&#93;:&#93;, freqs_sin&#91;-video_length * rope_sizes&#91;1&#93; * rope_sizes&#91;2&#93; //
                                                 self.patch_size&#91;0&#93;:&#93;</code></pre></details>


fastvideo.models.dits.matrixgame3.action_module.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


Source
def forward(
    self,
    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, tt*th*tw, C
    mouse_condition: B, N_frames, C1
    keyboard_condition: B, N_frames, C2
    """
    target_device = x.device
    target_dtype = x.dtype
    if mouse_condition is not None:
        mouse_condition = mouse_condition.to(device=target_device, dtype=target_dtype)
    if keyboard_condition is not None:
        keyboard_condition = keyboard_condition.to(device=target_device, dtype=target_dtype)
    else:
        return x

B, N_frames, C = keyboard_condition.shape
assert tt * th * tw == x.shape&#91;1&#93;
assert (((N_frames - 1) + self.vae_time_compression_ratio) % self.vae_time_compression_ratio == 0
        or (N_frames % self.vae_time_compression_ratio == 0))
if ((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
else:
    N_feats = N_frames // self.vae_time_compression_ratio

# Lazy initialization of freqs
if (self._freqs_cos is None or self._freqs_sin is None or self._freqs_cos.device != x.device):
    fc, fs = self.get_rotary_pos_embed(
        7500,
        self.patch_size&#91;1&#93;,
        self.patch_size&#91;2&#93;,
        64,
        self.mouse_qk_dim_list,
    )
    self._freqs_cos = fc.to(x.device)
    self._freqs_sin = fs.to(x.device)

freqs_cis = (self._freqs_cos, self._freqs_sin)

if self.enable_mouse and mouse_condition is not None:
    hidden_states = rearrange(x, &quot;B (T S) C -&gt; (B S) T C&quot;, T=tt, S=th * tw)
    B, N_frames, C = mouse_condition.shape
else:
    hidden_states = x

pad_t = self.vae_time_compression_ratio * self.windows_size
if self.enable_mouse and mouse_condition is not None:
    attn = self._forward_mouse(
        hidden_states,
        mouse_condition,
        mouse_cond_memory,
        pad_t=pad_t,
        freqs_cis=freqs_cis,
        N_feats=N_feats,
        B=B,
        C=C,
        th=th,
        tw=tw,
    )
    hidden_states = rearrange(x, &quot;(B S) T C -&gt; B (T S) C&quot;, B=B)
    hidden_states = hidden_states + attn

if self.enable_keyboard and keyboard_condition is not None:
    attn = self._forward_keyboard(
        hidden_states,
        keyboard_condition,
        keyboard_cond_memory,
        pad_t=pad_t,
        freqs_cis=freqs_cis,
        N_feats=N_feats,
        B=B,
        th=th,
        tw=tw,
    )
    hidden_states = hidden_states + attn
return hidden_states</code></pre></details>

Was this page helpful?