| 261 | |
| 262 | |
| 263 | class AdaLayerNormContinuous(torch.nn.Module): |
| 264 | def __init__(self, dim): |
| 265 | super().__init__() |
| 266 | self.silu = torch.nn.SiLU() |
| 267 | self.linear = torch.nn.Linear(dim, dim * 2, bias=True) |
| 268 | self.norm = torch.nn.LayerNorm(dim, eps=1e-6, elementwise_affine=False) |
| 269 | |
| 270 | def forward(self, x, conditioning): |
| 271 | emb = self.linear(self.silu(conditioning)) |
| 272 | scale, shift = torch.chunk(emb, 2, dim=1) |
| 273 | x = self.norm(x) * (1 + scale)[:, None] + shift[:, None] |
| 274 | return x |
| 275 | |
| 276 | |
| 277 |