| 208 | |
| 209 | |
| 210 | class TimestepEmbedder(nn.Module): |
| 211 | |
| 212 | def __init__(self, latent_dim, sequence_pos_encoder): |
| 213 | super().__init__() |
| 214 | self.latent_dim = latent_dim |
| 215 | self.sequence_pos_encoder = sequence_pos_encoder |
| 216 | |
| 217 | time_embed_dim = self.latent_dim |
| 218 | self.time_embed = nn.Sequential( |
| 219 | nn.Linear(self.latent_dim, time_embed_dim), |
| 220 | nn.SiLU(), |
| 221 | nn.Linear(time_embed_dim, time_embed_dim), |
| 222 | ) |
| 223 | |
| 224 | def forward(self, timesteps): |
| 225 | output = self.time_embed(self.sequence_pos_encoder.pe[timesteps]) |
| 226 | output = output.permute(1, 0, 2) |
| 227 | return output |