| 233 | |
| 234 | |
| 235 | class WanRMSNorm(nn.Module): |
| 236 | |
| 237 | def __init__(self, dim, eps=1e-5): |
| 238 | super().__init__() |
| 239 | self.dim = dim |
| 240 | self.eps = eps |
| 241 | self.weight = nn.Parameter(torch.ones(dim)) |
| 242 | |
| 243 | def forward(self, x): |
| 244 | return self._norm(x.float()).type_as(x) * self.weight |
| 245 | |
| 246 | def _norm(self, x): |
| 247 | return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps) |
| 248 | |
| 249 | |
| 250 | class WanLayerNorm(nn.LayerNorm): |