| 236 | |
| 237 | |
| 238 | class WanAttentionBlock(nn.Module): |
| 239 | |
| 240 | def __init__(self, |
| 241 | cross_attn_type, |
| 242 | dim, |
| 243 | ffn_dim, |
| 244 | num_heads, |
| 245 | window_size=(-1, -1), |
| 246 | qk_norm=True, |
| 247 | cross_attn_norm=False, |
| 248 | eps=1e-6): |
| 249 | super().__init__() |
| 250 | self.dim = dim |
| 251 | self.ffn_dim = ffn_dim |
| 252 | self.num_heads = num_heads |
| 253 | self.window_size = window_size |
| 254 | self.qk_norm = qk_norm |
| 255 | self.cross_attn_norm = cross_attn_norm |
| 256 | self.eps = eps |
| 257 | |
| 258 | # layers |
| 259 | self.norm1 = WanLayerNorm(dim, eps) |
| 260 | self.self_attn = WanSelfAttention(dim, num_heads, window_size, qk_norm, |
| 261 | eps) |
| 262 | self.norm3 = WanLayerNorm( |
| 263 | dim, eps, |
| 264 | elementwise_affine=True) if cross_attn_norm else nn.Identity() |
| 265 | self.cross_attn = WAN_CROSSATTENTION_CLASSES[cross_attn_type](dim, |
| 266 | num_heads, |
| 267 | (-1, -1), |
| 268 | qk_norm, |
| 269 | eps) |
| 270 | self.norm2 = WanLayerNorm(dim, eps) |
| 271 | self.ffn = nn.Sequential( |
| 272 | nn.Linear(dim, ffn_dim), nn.GELU(approximate='tanh'), |
| 273 | nn.Linear(ffn_dim, dim)) |
| 274 | |
| 275 | # modulation |
| 276 | self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) |
| 277 | |
| 278 | def forward( |
| 279 | self, |
| 280 | x, |
| 281 | e, |
| 282 | seq_lens, |
| 283 | grid_sizes, |
| 284 | freqs, |
| 285 | context, |
| 286 | context_lens, |
| 287 | ): |
| 288 | r""" |
| 289 | Args: |
| 290 | x(Tensor): Shape [B, L, C] |
| 291 | e(Tensor): Shape [B, 6, C] |
| 292 | seq_lens(Tensor): Shape [B], length of each sequence in batch |
| 293 | grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W) |
| 294 | freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2] |
| 295 | """ |