| 231 | class AdaLayerNormZeroSingle(Module): |
| 232 | |
| 233 | def __init__(self, |
| 234 | embedding_dim: int, |
| 235 | norm_type: str = "layer_norm", |
| 236 | bias: bool = True, |
| 237 | mapping=Mapping(), |
| 238 | dtype=None): |
| 239 | super().__init__() |
| 240 | |
| 241 | self.silu = ACT2FN['silu'] |
| 242 | self.linear = Linear(embedding_dim, |
| 243 | 3 * embedding_dim, |
| 244 | bias=bias, |
| 245 | tp_group=mapping.tp_group, |
| 246 | tp_size=mapping.tp_size, |
| 247 | dtype=dtype) |
| 248 | if norm_type == "layer_norm": |
| 249 | self.norm = LayerNorm(embedding_dim, |
| 250 | elementwise_affine=False, |
| 251 | eps=1e-6) |
| 252 | else: |
| 253 | raise ValueError( |
| 254 | f"Unsupported `norm_type` ({norm_type}) provided. Supported ones are: 'layer_norm', 'fp32_layer_norm'." |
| 255 | ) |
| 256 | |
| 257 | def forward(self, x: Tensor, emb: Optional[Tensor] = None): |
| 258 | emb = self.linear(self.silu(emb)) |