| 176 | |
| 177 | |
| 178 | class AdaRMSNorm(nn.Module): |
| 179 | def __init__(self, features, cond_features, eps=1e-6): |
| 180 | super().__init__() |
| 181 | self.eps = eps |
| 182 | self.linear = apply_wd(zero_init(Linear(cond_features, features, bias=False))) |
| 183 | tag_module(self.linear, "mapping") |
| 184 | |
| 185 | def extra_repr(self): |
| 186 | return f"eps={self.eps}," |
| 187 | |
| 188 | def forward(self, x, cond): |
| 189 | return rms_norm(x, self.linear(cond)[:, None, None, :] + 1, self.eps) |
| 190 | |
| 191 | |
| 192 | # Rotary position embeddings |