| 77 | |
| 78 | |
| 79 | class WanRMSNorm(nn.Module): |
| 80 | |
| 81 | def __init__(self, dim, eps=1e-5): |
| 82 | super().__init__() |
| 83 | self.dim = dim |
| 84 | self.eps = eps |
| 85 | self.weight = nn.Parameter(torch.ones(dim)) |
| 86 | |
| 87 | def forward(self, x): |
| 88 | r""" |
| 89 | Args: |
| 90 | x(Tensor): Shape [B, L, C] |
| 91 | """ |
| 92 | return self._norm(x.float()).type_as(x) * self.weight |
| 93 | |
| 94 | def _norm(self, x): |
| 95 | return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps) |
| 96 | |
| 97 | |
| 98 | class WanLayerNorm(nn.LayerNorm): |