| 437 | |
| 438 | |
| 439 | class WanAttentionBlock_SAonly(nn.Module): |
| 440 | |
| 441 | def __init__(self, |
| 442 | dim, |
| 443 | ffn_dim, |
| 444 | num_heads, |
| 445 | window_size=(-1, -1), |
| 446 | qk_norm=True, |
| 447 | eps=1e-6, |
| 448 | ): |
| 449 | super().__init__() |
| 450 | self.dim = dim |
| 451 | self.ffn_dim = ffn_dim |
| 452 | self.num_heads = num_heads |
| 453 | self.window_size = window_size |
| 454 | self.qk_norm = qk_norm |
| 455 | self.eps = eps |
| 456 | |
| 457 | # layers |
| 458 | self.norm1 = WanLayerNorm(dim, eps) |
| 459 | self.self_attn = WanSelfAttention(dim, num_heads, window_size, qk_norm, eps) |
| 460 | self.norm2 = WanLayerNorm(dim, eps) |
| 461 | self.ffn = nn.Sequential( |
| 462 | nn.Linear(dim, ffn_dim), nn.GELU(approximate='tanh'), |
| 463 | nn.Linear(ffn_dim, dim)) |
| 464 | |
| 465 | # modulation |
| 466 | self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) |
| 467 | |
| 468 | |
| 469 | def forward( |
| 470 | self, |
| 471 | x, |
| 472 | e, |
| 473 | seq_lens, |
| 474 | freqs, |
| 475 | ): |
| 476 | assert e.dtype == torch.float32 |
| 477 | with amp.autocast(dtype=torch.float32, device_type="cuda"): |
| 478 | e = (self.modulation.to(dtype=e.dtype, device=e.device) + e).chunk(6, dim=1) |
| 479 | assert e[0].dtype == torch.float32 |
| 480 | |
| 481 | # self-attention |
| 482 | y = self.self_attn( |
| 483 | self.norm1(x).float() * (1 + e[1]) + e[0], seq_lens, |
| 484 | freqs) |
| 485 | |
| 486 | with amp.autocast(dtype=torch.float32, device_type="cuda"): |
| 487 | x = x + y * e[2] |
| 488 | |
| 489 | # ffn function |
| 490 | y = self.ffn(self.norm2(x).float() * (1 + e[4]) + e[3]) |
| 491 | with amp.autocast(dtype=torch.float32, device_type="cuda"): |
| 492 | x = x + y * e[5] |
| 493 | return x |
| 494 | |
| 495 | |
| 496 | class Head(nn.Module): |