From ddb890cb813de32cea5edce41b7b69339deb0b3d Mon Sep 17 00:00:00 2001
From: Yuan-Man <68322456+Yuan-ManX@users.noreply.github.com>
Date: Tue, 13 May 2025 14:29:40 +0800
Subject: [PATCH] Create motion_module.py

---
 .../matrixgame_dit_src/motion_module.py       | 331 ++++++++++++++++++
 1 file changed, 331 insertions(+)
 create mode 100644 matrixgame/model_variants/matrixgame_dit_src/motion_module.py

diff --git a/matrixgame/model_variants/matrixgame_dit_src/motion_module.py b/matrixgame/model_variants/matrixgame_dit_src/motion_module.py
new file mode 100644
index 0000000..d9e4e42
--- /dev/null
+++ b/matrixgame/model_variants/matrixgame_dit_src/motion_module.py
@@ -0,0 +1,331 @@
+from typing import Any, List, Tuple, Optional, Union, Dict
+from einops import rearrange
+from flash_attn import flash_attn_func
+import torch
+import torch.nn as nn
+from .mlp_layers import MLP, MLPEmbedder, FinalLayer
+from .attenion import attention, get_cu_seqlens
+from .norm_layers import get_norm_layer
+from .posemb_layers import apply_rotary_emb, get_nd_rotary_pos_embed, get_1d_rotary_pos_embed
+
+from .activation_layers import get_activation_layer
+import torch.nn.functional as F
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+    """
+    This function generates 1D positional embeddings from a grid.
+
+    Args:
+        embed_dim (`int`): The embedding dimension `D`
+        pos (`torch.Tensor`): 1D tensor of positions with shape `(M,)`
+
+    Returns:
+        `torch.Tensor`: Sinusoidal positional embeddings of shape `(M, D)`.
+    """
+    if embed_dim % 2 != 0:
+        raise ValueError("embed_dim must be divisible by 2")
+
+    omega = torch.arange(embed_dim // 2, device=pos.device, dtype=torch.float64)
+    omega /= embed_dim / 2.0
+    omega = 1.0 / 10000**omega  # (D/2,)
+
+    pos = pos.reshape(-1)  # (M,)
+    out = torch.outer(pos, omega)  # (M, D/2), outer product
+
+    emb_sin = torch.sin(out)  # (M, D/2)
+    emb_cos = torch.cos(out)  # (M, D/2)
+
+    emb = torch.concat([emb_sin, emb_cos], dim=1)  # (M, D)
+    return emb
+
+class Condition_in(nn.Module):
+    def __init__(self, in_channels, hidden_size_mid, hidden_size, windows_size, vae_time_compression_ratio , device=None, dtype=None):
+        super().__init__()
+        factory_kwargs = {"device": device, "dtype": dtype}
+        self.in_channels = in_channels
+        self.hidden_size = hidden_size
+        self.windows_size = windows_size
+        self.vae_time_compression_ratio = vae_time_compression_ratio
+        self.embd = MLPEmbedder(in_channels, hidden_size_mid, **factory_kwargs)
+        self.out = MLP(in_channels = hidden_size_mid * windows_size * vae_time_compression_ratio, hidden_channels = hidden_size * 2 ,out_features = hidden_size, **factory_kwargs)
+
+    def forward(self, x, length):
+        x = self.embd(x) # b l d
+        # pos embedding
+        pos_embed = get_1d_sincos_pos_embed_from_grid(x.shape[-1], torch.arange(x.shape[1]).to(x.device)).unsqueeze(0)
+        x = x + pos_embed.repeat(x.shape[0], 1, 1).to(x.dtype)
+        # group
+        b, num_frames, dim = x.shape
+        pad_t = self.vae_time_compression_ratio * (self.windows_size)
+        feature_time_dim = feature_time_dim = int((num_frames - 1) / self.vae_time_compression_ratio) + 1
+        x = F.pad(x, (0, 0, pad_t, 0), value=0)
+        x = [x[:, self.vae_time_compression_ratio*(i - self.windows_size) + pad_t:i * self.vae_time_compression_ratio + pad_t,:] for i in range(feature_time_dim)]
+        x = torch.stack(x, dim = 1) # b n l d
+        x = x.reshape(shape=(x.shape[0],x.shape[1],-1)) # b n l*d
+        # out
+        x = self.out(x)
+        mask = torch.ones((x.shape[0],x.shape[1])).to(device=x.device, dtype=x.dtype)
+        padding = length - x.shape[1]
+        if padding < 0:
+            raise ValueError("action length is too long !")
+        mask = F.pad(mask,(0, padding, 0, 0), value=0)
+        x =  F.pad(x,(0, 0, 0, padding), value=0)
+        return x, mask
+
+class ActionModule(nn.Module):
+    """
+    action module from https://arxiv.org/pdf/2501.08325
+    鼠标控制信号的输入是一个 L*D 的向量
+    键盘同样
+    """
+
+    def __init__(
+        self, 
+        mouse_dim_in: int,
+        keyboard_dim_in: int,
+        hidden_size: int,
+        img_hidden_size: int,
+        vae_time_compression_ratio:int, 
+        windows_size:int,
+        heads_num: int = 24,
+        patch_size: list = [1, 2, 2],
+        mlp_act_type: str = "gelu_tanh",
+        qk_norm: bool = True,
+        qk_norm_type: str = "rms",
+        qkv_bias: bool = False,
+        drop: float = 0.0,
+        rope_dim_list: list = [16, 56, 56],
+        dtype: Optional[torch.dtype] = None,
+        device: Optional[torch.device] = None,
+        rope_theta = 256,
+        mouse_qk_dim_list = [24, 4, 4],
+        enable_mouse = True,
+        enable_keyboard = False,
+    ):
+        device = None
+        dtype = None
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+        self.enable_mouse = enable_mouse
+        self.enable_keyboard = enable_keyboard
+        # self.group = 
+        self.rope_dim_list = rope_dim_list
+        self.rope_theta = rope_theta
+        if self.enable_keyboard:
+            self.keyboard_embed = MLPEmbedder(in_dim=keyboard_dim_in, hidden_dim=hidden_size, **factory_kwargs)
+        mouse_dim_out = img_hidden_size // (patch_size[0]*patch_size[1]*patch_size[2])
+        self.mouse_qk_dim_list = mouse_qk_dim_list
+        qk_norm_layer = get_norm_layer(qk_norm_type)
+        self.heads_num = heads_num
+        if self.enable_mouse:
+            # self.mouse_embed = MLPEmbedder(in_dim=mouse_dim_in, hidden_dim=hidden_size, **factory_kwargs) # b rn c
+            self.mouse_mlp = MLP(
+                    mouse_dim_in * vae_time_compression_ratio * windows_size + img_hidden_size,
+                    out_features=img_hidden_size,
+                    act_layer=get_activation_layer(mlp_act_type),
+                    bias=True,
+                    **factory_kwargs,
+                )
+            c = img_hidden_size
+            head_dim = c // heads_num
+            self.t_qkv = nn.Linear(c, c*3, bias=qkv_bias, **factory_kwargs)
+            self.img_attn_q_norm = (
+                qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs)
+                if qk_norm
+                else nn.Identity()
+            )
+            self.img_attn_k_norm = (
+                qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs)
+                if qk_norm
+                else nn.Identity()
+            )
+            self.proj = nn.Linear(c, c, bias=qkv_bias, **factory_kwargs)
+            self.drop = nn.Dropout(drop)
+        if self.enable_keyboard:
+            head_dim_key = img_hidden_size // heads_num
+            self.key_attn_q_norm = (
+                qk_norm_layer(head_dim_key, elementwise_affine=True, eps=1e-6, **factory_kwargs)
+                if qk_norm
+                else nn.Identity()
+            )
+            self.key_attn_k_norm = (
+                qk_norm_layer(head_dim_key, elementwise_affine=True, eps=1e-6, **factory_kwargs)
+                if qk_norm
+                else nn.Identity()
+            )
+            self.mouse_attn_q = nn.Linear(img_hidden_size, img_hidden_size, bias=qkv_bias, **factory_kwargs)
+            self.keyboard_attn_kv = nn.Linear(hidden_size * windows_size * vae_time_compression_ratio, img_hidden_size * 2, bias=qkv_bias, **factory_kwargs)
+            self.proj = nn.Linear(img_hidden_size, img_hidden_size, bias=qkv_bias, **factory_kwargs)
+            self.drop = nn.Dropout(drop)
+        self.vae_time_compression_ratio = vae_time_compression_ratio
+        self.windows_size = windows_size
+        self.patch_size = patch_size
+        # self.unpatchify_channels = unpatchify_channels # 16 VAE.config.latent_channels * patch_size ** 2
+
+    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
+
+    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 #self.unpatchify_channels
+        pt, ph, pw = self.patch_size
+        assert t * h * w == x.shape[1]
+
+        x = x.reshape(shape=(x.shape[0], t, h, w, c, pt, ph, pw))
+        x = torch.einsum("nthwcopq->nctohpwq", x)
+        imgs = x.reshape(shape=(x.shape[0], 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
+        # 884 
+        latents_size = [video_length, height, width]
+
+        if isinstance(self.patch_size, int):
+            assert all(s % self.patch_size == 0 for s in latents_size), (
+                f"Latent size(last {ndim} dimensions) should be divisible by patch size({self.patch_size}), "
+                f"but got {latents_size}."
+            )
+            rope_sizes = [s // self.patch_size for s in latents_size]
+        elif isinstance(self.patch_size, list):
+            assert all(
+                s % self.patch_size[idx] == 0
+                for idx, s in enumerate(latents_size)
+            ), (
+                f"Latent size(last {ndim} dimensions) should be divisible by patch size({self.patch_size}), "
+                f"but got {latents_size}."
+            )
+            rope_sizes = [
+                s // self.patch_size[idx] for idx, s in enumerate(latents_size)
+            ]
+
+        if len(rope_sizes) != target_ndim:
+            rope_sizes = [1] * (target_ndim - len(rope_sizes)) + rope_sizes  # time axis
+        
+        if rope_dim_list is None:
+            rope_dim_list = [head_dim // target_ndim for _ in range(target_ndim)]
+        assert (
+            sum(rope_dim_list) == head_dim
+        ), "sum(rope_dim_list) should equal to head_dim of attention layer"
+        freqs_cos, freqs_sin = get_nd_rotary_pos_embed(
+            rope_dim_list,
+            rope_sizes,
+            theta=self.rope_theta,
+            use_real=True,
+            theta_rescale_factor=1,
+        )
+        return freqs_cos, freqs_sin
+
+    def forward(self, mouse_condition, keyboard_condition, x, tt, th, tw):
+        '''
+        hidden_states: b l 3072
+        mouse_condition: b t d
+        keyboard_condition:b t d
+        '''
+        b, num_frames, dim = keyboard_condition.shape
+        assert (num_frames + (self.vae_time_compression_ratio - 1)) % self.vae_time_compression_ratio == 0
+        feature_time_dim = int((num_frames - 1) / self.vae_time_compression_ratio) + 1
+        ## unpathch
+        if self.enable_mouse:
+            hidden_states = rearrange(x, "B (T S) C -> (B S) T C", T=tt, S=th*tw)
+            # hidden_states = self.unpatchify(x, tt, th, tw, self.patch_size[0]*self.patch_size[1]*self.patch_size[2])
+            # mouse_condition = self.mouse_embed(mouse_condition)  # b t hidden_size
+            b, num_frames, dim = mouse_condition.shape
+        else:
+            hidden_states = x
+        # padding
+        pad_t = self.vae_time_compression_ratio * (self.windows_size)
+        if self.enable_mouse:
+            mouse_condition = F.pad(mouse_condition, (0, 0, pad_t, 0))
+            group_mouse = [mouse_condition[:, self.vae_time_compression_ratio*(i - self.windows_size) + pad_t:i * self.vae_time_compression_ratio + pad_t,:] for i in range(feature_time_dim)]
+            group_mouse = torch.stack(group_mouse, dim = 1)
+            # group
+            group_mouse = torch.stack([group_mouse] * th * tw, dim = -1) # b t win*scale d s
+            group_mouse = rearrange(group_mouse, 'b t window d s -> (b s) t (window d)')
+            # concat hidden_states, group_mouse
+            group_mouse = torch.cat([hidden_states, group_mouse], dim = -1)
+            group_mouse = self.mouse_mlp(group_mouse)
+            # qkv
+            mouse_qkv = self.t_qkv(group_mouse)
+            q,k,v = rearrange(mouse_qkv, "B L (K H D) -> K B L H D",K=3,H=self.heads_num) # b l h d
+            q = self.img_attn_q_norm(q).to(v)
+            k = self.img_attn_k_norm(k).to(v)        
+            # rope embd
+            freqs_cos, freqs_sin = self.get_rotary_pos_embed(tt, self.patch_size[1], self.patch_size[2], k.shape[-1], self.mouse_qk_dim_list)
+            freqs_cis = (freqs_cos, freqs_sin)
+            if freqs_cis is not None:
+                qq, kk = apply_rotary_emb(q, k, freqs_cis, head_first=False)
+                assert (
+                    qq.shape == q.shape and kk.shape == k.shape
+                ), f"qq: {qq.shape}, q: {q.shape}, kk: {kk.shape}, k: {k.shape}"
+                q, k = qq, kk
+            attn = flash_attn_func(
+                q,
+                k,
+                v,
+                softmax_scale=q.shape[-1]**-0.5,
+                causal=False,
+            )
+            # Compute cu_squlens and max_seqlen for flash attention
+            # qk norm
+            attn = rearrange(attn, '(b S) T h d -> b (T S) (h d)',b=b)
+            hidden_states = rearrange(x, "(B S) T C -> B (T S) C", B=b)
+            attn = self.drop(self.proj(attn))
+            hidden_states = hidden_states + attn
+        if self.enable_keyboard:
+            keyboard_condition = self.keyboard_embed(keyboard_condition) # b t hidden_size
+            # add pos embed
+            pos_embed = get_1d_sincos_pos_embed_from_grid(keyboard_condition.shape[-1], torch.arange(keyboard_condition.shape[1]).to(keyboard_condition.device)).unsqueeze(0)
+            keyboard_condition = keyboard_condition + pos_embed.repeat((keyboard_condition.shape[0], 1, 1)).to(keyboard_condition.dtype)
+            keyboard_condition = F.pad(keyboard_condition, (0, 0, pad_t, 0))
+            ## group:  b t hidden_size --- > (n + 1) * rw * hidden_size
+            group_keyboard = [keyboard_condition[:, self.vae_time_compression_ratio*(i - self.windows_size) + pad_t:i * self.vae_time_compression_ratio + pad_t,:] for i in range(feature_time_dim)]
+            group_keyboard = torch.stack(group_keyboard, dim = 1)
+            group_keyboard = group_keyboard.reshape(shape=(group_keyboard.shape[0],group_keyboard.shape[1],-1))
+            ## repeat
+        
+            # apply corss attn
+            mouse_q = self.mouse_attn_q(hidden_states)
+            keyboard_kv = self.keyboard_attn_kv(group_keyboard)
+            q = rearrange(mouse_q, "B L (H D) -> B L H D",H=self.heads_num)
+            k, v = rearrange(keyboard_kv, "B L (K H D) -> K B L H D",K=2,H=self.heads_num)
+            # Compute cu_squlens and max_seqlen for flash attention
+            # qk norm
+            q = self.key_attn_q_norm(q).to(v)
+            k = self.key_attn_k_norm(k).to(v)
+            # position embed 
+            freqs_cos, freqs_sin = self.get_rotary_pos_embed(tt * self.patch_size[0], th * self.patch_size[1], tw * self.patch_size[2], k.shape[-1], self.rope_dim_list)
+            freqs_cis = (freqs_cos, freqs_sin)
+            if freqs_cis is not None:
+                qq, kk = apply_rotary_emb(q, k, freqs_cis, head_first=False)
+                assert (
+                    qq.shape == q.shape and kk.shape == k.shape
+                ), f"img_kk: {qq.shape}, img_q: {q.shape}, img_kk: {kk.shape}, img_k: {k.shape}"
+                q, k = qq, kk
+            attn = flash_attn_func(
+                q,
+                k,
+                v,
+                softmax_scale=q.shape[-1]**-0.5,
+                causal=False,
+            )
+            attn = rearrange(attn, 'B L H D -> B L (H D)')
+            attn = self.drop(self.proj(attn))
+            hidden_states = hidden_states + attn
+        return hidden_states
+