| 275 | |
| 276 | |
| 277 | class AdaLayerNormContinuous(nn.Module): |
| 278 | def __init__( |
| 279 | self, |
| 280 | embedding_dim: int, |
| 281 | conditioning_embedding_dim: int, |
| 282 | # NOTE: It is a bit weird that the norm layer can be configured to have scale and shift parameters |
| 283 | # because the output is immediately scaled and shifted by the projected conditioning embeddings. |
| 284 | # Note that AdaLayerNorm does not let the norm layer have scale and shift parameters. |
| 285 | # However, this is how it was implemented in the original code, and it's rather likely you should |
| 286 | # set `elementwise_affine` to False. |
| 287 | elementwise_affine=True, |
| 288 | eps=1e-5, |
| 289 | bias=True, |
| 290 | norm_type="layer_norm", |
| 291 | ): |
| 292 | super().__init__() |
| 293 | self.silu = nn.SiLU() |
| 294 | self.linear = nn.Linear(conditioning_embedding_dim, embedding_dim * 2, bias=bias) |
| 295 | if norm_type == "layer_norm": |
| 296 | self.norm = LayerNorm(embedding_dim, eps, elementwise_affine, bias) |
| 297 | elif norm_type == "rms_norm": |
| 298 | self.norm = RMSNorm(embedding_dim, eps, elementwise_affine) |
| 299 | else: |
| 300 | raise ValueError(f"unknown norm_type {norm_type}") |
| 301 | |
| 302 | def forward(self, x: torch.Tensor, conditioning_embedding: torch.Tensor) -> torch.Tensor: |
| 303 | # convert back to the original dtype in case `conditioning_embedding`` is upcasted to float32 (needed for hunyuanDiT) |
| 304 | emb = self.linear(self.silu(conditioning_embedding).to(x.dtype)) |
| 305 | scale, shift = torch.chunk(emb, 2, dim=1) |
| 306 | x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :] |
| 307 | return x |
| 308 | |
| 309 | |
| 310 | class LuminaLayerNormContinuous(nn.Module): |