| 265 | class AdaLayerNormContinuous(Module): |
| 266 | |
| 267 | def __init__(self, |
| 268 | embedding_dim: int, |
| 269 | conditioning_embedding_dim: int, |
| 270 | elementwise_affine: bool = True, |
| 271 | eps: float = 1e-5, |
| 272 | bias: bool = True, |
| 273 | norm_type: str = "layer_norm", |
| 274 | mapping=Mapping(), |
| 275 | dtype=None): |
| 276 | super().__init__() |
| 277 | self.silu = ACT2FN['silu'] |
| 278 | self.linear = Linear(conditioning_embedding_dim, |
| 279 | embedding_dim * 2, |
| 280 | bias=bias, |
| 281 | tp_group=mapping.tp_group, |
| 282 | tp_size=mapping.tp_size, |
| 283 | dtype=dtype) |
| 284 | if norm_type == "layer_norm": |
| 285 | self.norm = LayerNorm(embedding_dim, |
| 286 | eps=eps, |
| 287 | elementwise_affine=elementwise_affine, |
| 288 | bias=bias, |
| 289 | dtype=dtype) |
| 290 | elif norm_type == "rms_norm": |
| 291 | self.norm = RmsNorm(embedding_dim, |
| 292 | eps=eps, |
| 293 | elementwise_affine=elementwise_affine, |
| 294 | dtype=dtype) |
| 295 | else: |
| 296 | raise ValueError(f"unknown norm_type {norm_type}") |
| 297 | |
| 298 | def forward(self, x: Tensor, conditioning_embedding: Tensor): |
| 299 | # convert back to the original dtype in case `conditioning_embedding`` is upcasted to float32 (needed for hunyuanDiT) |