Skip to content

model ¶

Classes¶

fastvideo.models.dits.matrixgame3.model.MatrixGame3CrossAttention ¶

MatrixGame3CrossAttention(dim: int, num_heads: int, window_size=(-1, -1), qk_norm=True, eps=1e-06, parallel_attention=False, quant_config: QuantizationConfig | None = None, prefix: str = '')

Bases: WanSelfAttention

Source code in fastvideo/models/dits/wanvideo.py

def __init__(self,
             dim: int,
             num_heads: int,
             window_size=(-1, -1),
             qk_norm=True,
             eps=1e-6,
             parallel_attention=False,
             quant_config: QuantizationConfig | None = None,
             prefix: str = "") -> None:
    assert dim % num_heads == 0
    super().__init__()
    self.dim = dim
    self.num_heads = num_heads
    self.head_dim = dim // num_heads
    self.window_size = window_size
    self.qk_norm = qk_norm
    self.eps = eps
    self.parallel_attention = parallel_attention

    # layers
    self.to_q = ReplicatedLinear(dim, dim, quant_config=quant_config, prefix=f"{prefix}.to_q")
    self.to_k = ReplicatedLinear(dim, dim, quant_config=quant_config, prefix=f"{prefix}.to_k")
    self.to_v = ReplicatedLinear(dim, dim, quant_config=quant_config, prefix=f"{prefix}.to_v")
    self.to_out = ReplicatedLinear(dim, dim, quant_config=quant_config, prefix=f"{prefix}.to_out")
    self.norm_q = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
    self.norm_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity()

    # Scaled dot product attention
    self.attn = LocalAttention(
        num_heads=num_heads,
        head_size=self.head_dim,
        dropout_rate=0,
        softmax_scale=None,
        causal=False,
        supported_attention_backends=(AttentionBackendEnum.FLASH_ATTN,
                                      AttentionBackendEnum.TORCH_SDPA))

Methods:¶

fastvideo.models.dits.matrixgame3.model.MatrixGame3CrossAttention.forward ¶

forward(x, context, context_lens=None)

Parameters:

Name	Type	Description	Default
`x`	`Tensor`	Shape [B, L1, C]	required
`context`	`Tensor`	Shape [B, L2, C] - typically 257 image tokens	required
`context_lens`	`Tensor`	Shape [B]	`None`

Source code in fastvideo/models/dits/matrixgame3/model.py

def forward(self, x, context, context_lens=None):
    r"""
    Args:
        x(Tensor): Shape [B, L1, C]
        context(Tensor): Shape [B, L2, C] - typically 257 image tokens
        context_lens(Tensor): Shape [B]
    """
    b, n, d = x.size(0), self.num_heads, self.head_dim

    q_input = x.to(self.to_q.weight.dtype)
    q = self.norm_q(self.to_q(q_input)[0]).view(b, -1, n, d)

    context_input = context.to(self.to_k.weight.dtype)
    k = self.norm_k(self.to_k(context_input)[0]).view(b, -1, n, d)
    v = self.to_v(context_input)[0].view(b, -1, n, d)

    x = self.attn(q, k, v)

    x = x.flatten(2)
    x, _ = self.to_out(x)
    return x

Functions:¶