| 307 | |
| 308 | |
| 309 | class WanT2VCrossAttention(WanSelfAttention): |
| 310 | |
| 311 | def forward(self, x, context, context_lens): |
| 312 | """ |
| 313 | x: [B, L1, C]. |
| 314 | context: [B, L2, C]. |
| 315 | context_lens: [B]. |
| 316 | """ |
| 317 | b, n, d = x.size(0), self.num_heads, self.head_dim |
| 318 | |
| 319 | # compute query, key, value |
| 320 | q = self.norm_q(self.q(x)).view(b, -1, n, d) |
| 321 | k = self.norm_k(self.k(context)).view(b, -1, n, d) |
| 322 | v = self.v(context).view(b, -1, n, d) |
| 323 | |
| 324 | # compute attention |
| 325 | x = flash_attention(q, k, v, k_lens=context_lens) |
| 326 | |
| 327 | # output |
| 328 | x = x.flatten(2) |
| 329 | x = self.o(x) |
| 330 | return x |
| 331 | |
| 332 | |
| 333 | WANX_CROSSATTENTION_CLASSES = { |