| 96 | |
| 97 | |
| 98 | class WanLayerNorm(nn.LayerNorm): |
| 99 | |
| 100 | def __init__(self, dim, eps=1e-6, elementwise_affine=False): |
| 101 | super().__init__(dim, elementwise_affine=elementwise_affine, eps=eps) |
| 102 | |
| 103 | def forward(self, inputs: torch.Tensor) -> torch.Tensor: |
| 104 | origin_dtype = inputs.dtype |
| 105 | out = F.layer_norm( |
| 106 | inputs.float(), |
| 107 | self.normalized_shape, |
| 108 | None if self.weight is None else self.weight.float(), |
| 109 | None if self.bias is None else self.bias.float() , |
| 110 | self.eps |
| 111 | ).to(origin_dtype) |
| 112 | return out |
| 113 | |
| 114 | |
| 115 | class WanSelfAttention(nn.Module): |