| 247 | |
| 248 | |
| 249 | class WanLayerNorm(nn.LayerNorm): |
| 250 | |
| 251 | def __init__(self, dim, eps=1e-6, elementwise_affine=False): |
| 252 | super().__init__(dim, elementwise_affine=elementwise_affine, eps=eps) |
| 253 | |
| 254 | def forward(self, x): |
| 255 | return super().forward(x.float()).type_as(x) |
| 256 | |
| 257 | |
| 258 | class WanSelfAttention(nn.Module): |