| 52 | |
| 53 | |
| 54 | class TimestepEmbeddings(torch.nn.Module): |
| 55 | def __init__(self, dim_in, dim_out, computation_device=None): |
| 56 | super().__init__() |
| 57 | self.time_proj = TemporalTimesteps(num_channels=dim_in, flip_sin_to_cos=True, downscale_freq_shift=0, computation_device=computation_device) |
| 58 | self.timestep_embedder = torch.nn.Sequential( |
| 59 | torch.nn.Linear(dim_in, dim_out), torch.nn.SiLU(), torch.nn.Linear(dim_out, dim_out) |
| 60 | ) |
| 61 | |
| 62 | def forward(self, timestep, dtype): |
| 63 | time_emb = self.time_proj(timestep).to(dtype) |
| 64 | time_emb = self.timestep_embedder(time_emb) |
| 65 | return time_emb |
| 66 | |
| 67 | |
| 68 | |