(self, x, seq_lens, freqs)
| 282 | self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity() |
| 283 | |
| 284 | def forward(self, x, seq_lens, freqs): |
| 285 | b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim |
| 286 | |
| 287 | # query, key, value function |
| 288 | def qkv_fn(x): |
| 289 | q = self.norm_q(self.q(x)).view(b, s, n, d) |
| 290 | k = self.norm_k(self.k(x)).view(b, s, n, d) |
| 291 | v = self.v(x).view(b, s, n, d) |
| 292 | return q, k, v |
| 293 | |
| 294 | # import pdb; pdb.set_trace() |
| 295 | q, k, v = qkv_fn(x) |
| 296 | x = flash_attention( |
| 297 | q=q, |
| 298 | k=k, |
| 299 | v=v, |
| 300 | k_lens=seq_lens, |
| 301 | window_size=self.window_size) |
| 302 | |
| 303 | # output |
| 304 | x = x.flatten(2) |
| 305 | x = self.o(x) |
| 306 | return x |
| 307 | |
| 308 | |
| 309 | class WanT2VCrossAttention(WanSelfAttention): |
nothing calls this directly
no test coverage detected