matrixgame2 ¶

def forward(
    self,
    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, tt*th*tw, C
    mouse_condition: B, N_frames, C1
    keyboard_condition: B, N_frames, C2
    """
    assert use_rope_keyboard

    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[1]
    assert (
        (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

    # Lazy initialization of freqs on first forward pass
    if self._freqs_cos is None or self._freqs_sin is None:
        self._freqs_cos, self._freqs_sin = self.get_rotary_pos_embed(
            7500,
            self.patch_size[1],
            self.patch_size[2],
            64,
            self.mouse_qk_dim_list,
            start_offset=0,
        )

    # Defined freqs_cis early so it's available for both mouse and keyboard
    freqs_cis = (self._freqs_cos, self._freqs_sin)

    if is_causal:
        assert (N_feats == tt and kv_cache_mouse is None) or (
            (N_frames - 1) // self.vae_time_compression_ratio + 1
            == start_frame + num_frame_per_block
        )
    # For non-causal (training), we trust that the caller provides correctly shaped inputs

    if self.enable_mouse and mouse_condition is not None:
        hidden_states = rearrange(
            x, "B (T S) C -> (B S) T C", T=tt, S=th * tw
        )  # 65*272*480 -> 17*(272//16)*(480//16) -> 8670
        B, N_frames, C = mouse_condition.shape
    else:
        hidden_states = x
    # padding

    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,
            is_causal=is_causal,
            kv_cache_mouse=kv_cache_mouse,
            pad_t=pad_t,
            num_frame_per_block=num_frame_per_block,
            block_mask_mouse=block_mask_mouse,
            start_frame=start_frame,
            freqs_cis=freqs_cis,
            N_feats=N_feats,
            B=B,
            C=C,
            tt=tt,
            th=th,
            tw=tw,
        )
        hidden_states = rearrange(x, "(B S) T C -> B (T S) C", 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,
            is_causal=is_causal,
            use_rope_keyboard=use_rope_keyboard,
            kv_cache_keyboard=kv_cache_keyboard,
            pad_t=pad_t,
            num_frame_per_block=num_frame_per_block,
            block_mask_keyboard=block_mask_keyboard,
            start_frame=start_frame,
            freqs_cis=freqs_cis,
            N_feats=N_feats,
            B=B,
            tt=tt,
            th=th,
            tw=tw,
        )
        hidden_states = hidden_states + attn
    return hidden_states

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