Method forward

wan/modules/animate/face_blocks.py:334–383

(
        self,
        x: torch.Tensor,
        motion_vec: torch.Tensor,
        motion_mask: Optional[torch.Tensor] = None,
        use_context_parallel=False,
    )

        self.pre_norm_motion = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs)

    def forward(
        self,
        x: torch.Tensor,
        motion_vec: torch.Tensor,
        motion_mask: Optional[torch.Tensor] = None,
        use_context_parallel=False,
    ) -> torch.Tensor:
        
        B, T, N, C = motion_vec.shape
        T_comp = T

        x_motion = self.pre_norm_motion(motion_vec)
        x_feat = self.pre_norm_feat(x)

        kv = self.linear1_kv(x_motion)
        q = self.linear1_q(x_feat)

        k, v = rearrange(kv, "B L N (K H D) -> K B L N H D", K=2, H=self.heads_num)
        q = rearrange(q, "B S (H D) -> B S H D", H=self.heads_num)

        # Apply QK-Norm if needed.
        q = self.q_norm(q).to(v)
        k = self.k_norm(k).to(v)

        k = rearrange(k, "B L N H D -> (B L) N H D")  
        v = rearrange(v, "B L N H D -> (B L) N H D") 

        if use_context_parallel:
            q = gather_forward(q, dim=1)

        q = rearrange(q, "B (L S) H D -> (B L) S H D", L=T_comp)  
        # Compute attention.
        attn = attention(
            q,
            k,
            v,
            max_seqlen_q=q.shape[1],
            batch_size=q.shape[0],
        )

        attn = rearrange(attn, "(B L) S C -> B (L S) C", L=T_comp)
        if use_context_parallel:
            attn = torch.chunk(attn, get_world_size(), dim=1)[get_rank()]

        output = self.linear2(attn)

        if motion_mask is not None:
            output = output * rearrange(motion_mask, "B T H W -> B (T H W)").unsqueeze(-1)

        return output