action_module ¶

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 [1, 2, 2]
    rope_dim_list = (
        rope_dim_list if rope_dim_list is not None else [8, 28, 28]
    )
    mouse_qk_dim_list = (
        mouse_qk_dim_list if mouse_qk_dim_list is not None else [8, 28, 28]
    )
    blocks = blocks if blocks is not None else []
    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="tanh"),
            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