r""" GroupNorm layer modified to incorporate timestep embeddings. Parameters: embedding_dim (`int`): The size of each embedding vector. num_embeddings (`int`): The size of the embeddings dictionary. num_groups (`int`): The number of groups to separate the channels in
| 267 | |
| 268 | |
| 269 | class AdaGroupNorm(nn.Module): |
| 270 | r""" |
| 271 | GroupNorm layer modified to incorporate timestep embeddings. |
| 272 | |
| 273 | Parameters: |
| 274 | embedding_dim (`int`): The size of each embedding vector. |
| 275 | num_embeddings (`int`): The size of the embeddings dictionary. |
| 276 | num_groups (`int`): The number of groups to separate the channels into. |
| 277 | act_fn (`str`, *optional*, defaults to `None`): The activation function to use. |
| 278 | eps (`float`, *optional*, defaults to `1e-5`): The epsilon value to use for numerical stability. |
| 279 | """ |
| 280 | |
| 281 | def __init__( |
| 282 | self, embedding_dim: int, out_dim: int, num_groups: int, act_fn: str | None = None, eps: float = 1e-5 |
| 283 | ): |
| 284 | super().__init__() |
| 285 | self.num_groups = num_groups |
| 286 | self.eps = eps |
| 287 | |
| 288 | if act_fn is None: |
| 289 | self.act = None |
| 290 | else: |
| 291 | self.act = get_activation(act_fn) |
| 292 | |
| 293 | self.linear = nn.Linear(embedding_dim, out_dim * 2) |
| 294 | |
| 295 | def forward(self, x: torch.Tensor, emb: torch.Tensor) -> torch.Tensor: |
| 296 | if self.act: |
| 297 | emb = self.act(emb) |
| 298 | emb = self.linear(emb) |
| 299 | emb = emb[:, :, None, None] |
| 300 | scale, shift = emb.chunk(2, dim=1) |
| 301 | |
| 302 | x = F.group_norm(x, self.num_groups, eps=self.eps) |
| 303 | x = x * (1 + scale) + shift |
| 304 | return x |
| 305 | |
| 306 | |
| 307 | class AdaLayerNormContinuous(nn.Module): |