| 153 | |
| 154 | |
| 155 | class TimestepEmbedding(nn.Module): |
| 156 | def __init__( |
| 157 | self, |
| 158 | in_channels: int, |
| 159 | time_embed_dim: int, |
| 160 | act_fn: str = "silu", |
| 161 | out_dim: int = None, |
| 162 | post_act_fn: Optional[str] = None, |
| 163 | cond_proj_dim=None, |
| 164 | ): |
| 165 | super().__init__() |
| 166 | |
| 167 | self.linear_1 = nn.Linear(in_channels, time_embed_dim) |
| 168 | |
| 169 | if cond_proj_dim is not None: |
| 170 | self.cond_proj = nn.Linear(cond_proj_dim, in_channels, bias=False) |
| 171 | else: |
| 172 | self.cond_proj = None |
| 173 | |
| 174 | if act_fn == "silu": |
| 175 | self.act = nn.SiLU() |
| 176 | elif act_fn == "mish": |
| 177 | self.act = nn.Mish() |
| 178 | elif act_fn == "gelu": |
| 179 | self.act = nn.GELU() |
| 180 | else: |
| 181 | raise ValueError(f"{act_fn} does not exist. Make sure to define one of 'silu', 'mish', or 'gelu'") |
| 182 | |
| 183 | if out_dim is not None: |
| 184 | time_embed_dim_out = out_dim |
| 185 | else: |
| 186 | time_embed_dim_out = time_embed_dim |
| 187 | self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim_out) |
| 188 | |
| 189 | if post_act_fn is None: |
| 190 | self.post_act = None |
| 191 | elif post_act_fn == "silu": |
| 192 | self.post_act = nn.SiLU() |
| 193 | elif post_act_fn == "mish": |
| 194 | self.post_act = nn.Mish() |
| 195 | elif post_act_fn == "gelu": |
| 196 | self.post_act = nn.GELU() |
| 197 | else: |
| 198 | raise ValueError(f"{post_act_fn} does not exist. Make sure to define one of 'silu', 'mish', or 'gelu'") |
| 199 | |
| 200 | def forward(self, sample, condition=None): |
| 201 | if condition is not None: |
| 202 | sample = sample + self.cond_proj(condition) |
| 203 | sample = self.linear_1(sample) |
| 204 | |
| 205 | if self.act is not None: |
| 206 | sample = self.act(sample) |
| 207 | |
| 208 | sample = self.linear_2(sample) |
| 209 | |
| 210 | if self.post_act is not None: |
| 211 | sample = self.post_act(sample) |
| 212 | return sample |