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
| 237 | |
| 238 | |
| 239 | class AdaGroupNorm(nn.Module): |
| 240 | r""" |
| 241 | GroupNorm layer modified to incorporate timestep embeddings. |
| 242 | |
| 243 | Parameters: |
| 244 | embedding_dim (`int`): The size of each embedding vector. |
| 245 | num_embeddings (`int`): The size of the embeddings dictionary. |
| 246 | num_groups (`int`): The number of groups to separate the channels into. |
| 247 | act_fn (`str`, *optional*, defaults to `None`): The activation function to use. |
| 248 | eps (`float`, *optional*, defaults to `1e-5`): The epsilon value to use for numerical stability. |
| 249 | """ |
| 250 | |
| 251 | def __init__( |
| 252 | self, embedding_dim: int, out_dim: int, num_groups: int, act_fn: Optional[str] = None, eps: float = 1e-5 |
| 253 | ): |
| 254 | super().__init__() |
| 255 | self.num_groups = num_groups |
| 256 | self.eps = eps |
| 257 | |
| 258 | if act_fn is None: |
| 259 | self.act = None |
| 260 | else: |
| 261 | self.act = get_activation(act_fn) |
| 262 | |
| 263 | self.linear = nn.Linear(embedding_dim, out_dim * 2) |
| 264 | |
| 265 | def forward(self, x: torch.Tensor, emb: torch.Tensor) -> torch.Tensor: |
| 266 | if self.act: |
| 267 | emb = self.act(emb) |
| 268 | emb = self.linear(emb) |
| 269 | emb = emb[:, :, None, None] |
| 270 | scale, shift = emb.chunk(2, dim=1) |
| 271 | |
| 272 | x = F.group_norm(x, self.num_groups, eps=self.eps) |
| 273 | x = x * (1 + scale) + shift |
| 274 | return x |
| 275 | |
| 276 | |
| 277 | class AdaLayerNormContinuous(nn.Module): |