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
| 113 | |
| 114 | |
| 115 | class AdaGroupNorm(nn.Module): |
| 116 | r""" |
| 117 | GroupNorm layer modified to incorporate timestep embeddings. |
| 118 | |
| 119 | Parameters: |
| 120 | embedding_dim (`int`): The size of each embedding vector. |
| 121 | num_embeddings (`int`): The size of the embeddings dictionary. |
| 122 | num_groups (`int`): The number of groups to separate the channels into. |
| 123 | act_fn (`str`, *optional*, defaults to `None`): The activation function to use. |
| 124 | eps (`float`, *optional*, defaults to `1e-5`): The epsilon value to use for numerical stability. |
| 125 | """ |
| 126 | |
| 127 | def __init__( |
| 128 | self, embedding_dim: int, out_dim: int, num_groups: int, act_fn: Optional[str] = None, eps: float = 1e-5 |
| 129 | ): |
| 130 | super().__init__() |
| 131 | self.num_groups = num_groups |
| 132 | self.eps = eps |
| 133 | |
| 134 | if act_fn is None: |
| 135 | self.act = None |
| 136 | else: |
| 137 | self.act = get_activation(act_fn) |
| 138 | |
| 139 | self.linear = nn.Linear(embedding_dim, out_dim * 2) |
| 140 | |
| 141 | def forward(self, x: torch.Tensor, emb: torch.Tensor) -> torch.Tensor: |
| 142 | if self.act: |
| 143 | emb = self.act(emb) |
| 144 | emb = self.linear(emb) |
| 145 | emb = emb[:, :, None, None] |
| 146 | scale, shift = emb.chunk(2, dim=1) |
| 147 | |
| 148 | x = F.group_norm(x, self.num_groups, eps=self.eps) |
| 149 | x = x * (1 + scale) + shift |
| 150 | return x |
| 151 | |
| 152 | |
| 153 | class AdaLayerNormContinuous(nn.Module): |