(self,
cross_attn_type,
dim,
ffn_dim,
num_heads,
window_size=(-1, -1),
qk_norm=True,
cross_attn_norm=False,
eps=1e-6)
| 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, |
nothing calls this directly
no test coverage detected