| 160 | |
| 161 | |
| 162 | class WanT2VCrossAttention(WanSelfAttention): |
| 163 | |
| 164 | def forward(self, x, context, context_lens): |
| 165 | r""" |
| 166 | Args: |
| 167 | x(Tensor): Shape [B, L1, C] |
| 168 | context(Tensor): Shape [B, L2, C] |
| 169 | context_lens(Tensor): Shape [B] |
| 170 | """ |
| 171 | b, n, d = x.size(0), self.num_heads, self.head_dim |
| 172 | |
| 173 | # compute query, key, value |
| 174 | q = self.norm_q(self.q(x)).view(b, -1, n, d) |
| 175 | k = self.norm_k(self.k(context)).view(b, -1, n, d) |
| 176 | v = self.v(context).view(b, -1, n, d) |
| 177 | |
| 178 | # compute attention |
| 179 | x = flash_attention(q, k, v, k_lens=context_lens) |
| 180 | |
| 181 | # output |
| 182 | x = x.flatten(2) |
| 183 | x = self.o(x) |
| 184 | return x |
| 185 | |
| 186 | |
| 187 | class WanI2VCrossAttention(WanSelfAttention): |
nothing calls this directly
no outgoing calls
no test coverage detected