| 1363 | |
| 1364 | |
| 1365 | class Timesteps(nn.Module): |
| 1366 | def __init__(self, num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float, scale: int = 1): |
| 1367 | super().__init__() |
| 1368 | self.num_channels = num_channels |
| 1369 | self.flip_sin_to_cos = flip_sin_to_cos |
| 1370 | self.downscale_freq_shift = downscale_freq_shift |
| 1371 | self.scale = scale |
| 1372 | |
| 1373 | def forward(self, timesteps): |
| 1374 | t_emb = get_timestep_embedding( |
| 1375 | timesteps, |
| 1376 | self.num_channels, |
| 1377 | flip_sin_to_cos=self.flip_sin_to_cos, |
| 1378 | downscale_freq_shift=self.downscale_freq_shift, |
| 1379 | scale=self.scale, |
| 1380 | ) |
| 1381 | return t_emb |
| 1382 | |
| 1383 | |
| 1384 | class GaussianFourierProjection(nn.Module): |