| 1316 | |
| 1317 | class TimestepEmbedding(nn.Module): |
| 1318 | def __init__( |
| 1319 | self, |
| 1320 | in_channels: int, |
| 1321 | time_embed_dim: int, |
| 1322 | act_fn: str = "silu", |
| 1323 | out_dim: int = None, |
| 1324 | post_act_fn: Optional[str] = None, |
| 1325 | cond_proj_dim=None, |
| 1326 | sample_proj_bias=True, |
| 1327 | ): |
| 1328 | super().__init__() |
| 1329 | |
| 1330 | self.linear_1 = nn.Linear(in_channels, time_embed_dim, sample_proj_bias) |
| 1331 | |
| 1332 | if cond_proj_dim is not None: |
| 1333 | self.cond_proj = nn.Linear(cond_proj_dim, in_channels, bias=False) |
| 1334 | else: |
| 1335 | self.cond_proj = None |
| 1336 | |
| 1337 | self.act = get_activation(act_fn) |
| 1338 | |
| 1339 | if out_dim is not None: |
| 1340 | time_embed_dim_out = out_dim |
| 1341 | else: |
| 1342 | time_embed_dim_out = time_embed_dim |
| 1343 | self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim_out, sample_proj_bias) |
| 1344 | |
| 1345 | if post_act_fn is None: |
| 1346 | self.post_act = None |
| 1347 | else: |
| 1348 | self.post_act = get_activation(post_act_fn) |
| 1349 | |
| 1350 | def forward(self, sample, condition=None): |
| 1351 | if condition is not None: |