| 152 | |
| 153 | class AdaLayerNormContinuous(nn.Module): |
| 154 | def __init__( |
| 155 | self, |
| 156 | embedding_dim: int, |
| 157 | conditioning_embedding_dim: int, |
| 158 | # NOTE: It is a bit weird that the norm layer can be configured to have scale and shift parameters |
| 159 | # because the output is immediately scaled and shifted by the projected conditioning embeddings. |
| 160 | # Note that AdaLayerNorm does not let the norm layer have scale and shift parameters. |
| 161 | # However, this is how it was implemented in the original code, and it's rather likely you should |
| 162 | # set `elementwise_affine` to False. |
| 163 | elementwise_affine=True, |
| 164 | eps=1e-5, |
| 165 | bias=True, |
| 166 | norm_type="layer_norm", |
| 167 | ): |
| 168 | super().__init__() |
| 169 | self.silu = nn.SiLU() |
| 170 | self.linear = nn.Linear(conditioning_embedding_dim, embedding_dim * 2, bias=bias) |
| 171 | if norm_type == "layer_norm": |
| 172 | self.norm = LayerNorm(embedding_dim, eps, elementwise_affine, bias) |
| 173 | elif norm_type == "rms_norm": |
| 174 | self.norm = RMSNorm(embedding_dim, eps, elementwise_affine) |
| 175 | else: |
| 176 | raise ValueError(f"unknown norm_type {norm_type}") |
| 177 | |
| 178 | def forward(self, x: torch.Tensor, conditioning_embedding: torch.Tensor) -> torch.Tensor: |
| 179 | emb = self.linear(self.silu(conditioning_embedding)) |