Method forward

models/transformer/wan/modules/tm2m_model.py:382–436

(
        self,
        x,        
        e,
        seq_lens,
        freqs,
        context,
        context_lens,
        ref_motion,    # [B, T, C]
        ref_motion_lens,   # [B]
        attend_to_text_mask=None,  # [B]: 1 for attend to text, 0 for attend to ref motion
    )

    def forward(
        self,
        x,        
        e,
        seq_lens,
        freqs,
        context,
        context_lens,
        ref_motion,    # [B, T, C]
        ref_motion_lens,   # [B]
        attend_to_text_mask=None,  # [B]: 1 for attend to text, 0 for attend to ref motion
    ):
        assert e.dtype == torch.float32
        with amp.autocast(dtype=torch.float32, device_type="cuda"):
            e = (self.modulation.to(dtype=e.dtype, device=e.device) + e).chunk(6, dim=1)
        assert e[0].dtype == torch.float32

        # self-attention
        y = self.self_attn(
            self.norm1(x).float() * (1 + e[1]) + e[0], seq_lens, 
            freqs)

        with amp.autocast(dtype=torch.float32, device_type="cuda"):
            x = x + y * e[2]

        # Cross-attention based on attend_to_text_mask
        text_mask = attend_to_text_mask == 1
        ref_motion_mask = ~text_mask

        # Dummy way to get output shape
        b, s, d = ref_motion.shape
        # combined_out = self.cross_attn(self.norm3(x), context, None)
        combined_out = torch.zeros([b, s, self.dim], dtype=context.dtype, device=context.device)

        # Process text and reference motion separately
        attend_to_text_mask = attend_to_text_mask.to(dtype=bool)
        have_text = attend_to_text_mask.sum() > 0
        have_ref_motion = (~attend_to_text_mask).sum() > 0 
        if have_text:
            text_out = self.cross_attn(self.norm3(x[text_mask]), context[text_mask], None)
            # print('text_out shape', text_out.shape)
            combined_out[text_mask] = text_out

        if have_ref_motion:
            ref_motion_out = self.ref_motion_cross_attn(self.norm4(x[ref_motion_mask]), ref_motion[ref_motion_mask], ref_motion_lens[ref_motion_mask])
            # print('ref_motion_out shape', ref_motion_out.shape)
            combined_out[ref_motion_mask] = ref_motion_out if x[ref_motion_mask].size(0) > 0 else 0
        
        x += combined_out

        # ffn function
        y = self.ffn(self.norm2(x).float() * (1 + e[4]) + e[3])
        with amp.autocast(dtype=torch.float32, device_type="cuda"):
            x = x + y * e[5]
        return x


class WanAttentionBlock_SAonly(nn.Module):