This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings. :param timesteps: a 1-D Tensor of N indices, one per batch element. These may be fractional. :param embedding_dim: the dimension of the output. :par
(
timesteps: torch.Tensor,
embedding_dim: int,
flip_sin_to_cos: bool = False,
downscale_freq_shift: float = 1,
scale: float = 1,
max_period: int = 10000,
computation_device = None,
)
| 38 | |
| 39 | |
| 40 | def get_timestep_embedding( |
| 41 | timesteps: torch.Tensor, |
| 42 | embedding_dim: int, |
| 43 | flip_sin_to_cos: bool = False, |
| 44 | downscale_freq_shift: float = 1, |
| 45 | scale: float = 1, |
| 46 | max_period: int = 10000, |
| 47 | computation_device = None, |
| 48 | ): |
| 49 | """ |
| 50 | This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings. |
| 51 | |
| 52 | :param timesteps: a 1-D Tensor of N indices, one per batch element. |
| 53 | These may be fractional. |
| 54 | :param embedding_dim: the dimension of the output. :param max_period: controls the minimum frequency of the |
| 55 | embeddings. :return: an [N x dim] Tensor of positional embeddings. |
| 56 | """ |
| 57 | assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array" |
| 58 | |
| 59 | half_dim = embedding_dim // 2 |
| 60 | exponent = -math.log(max_period) * torch.arange( |
| 61 | start=0, end=half_dim, dtype=torch.float32, device=timesteps.device if computation_device is None else computation_device |
| 62 | ) |
| 63 | exponent = exponent / (half_dim - downscale_freq_shift) |
| 64 | |
| 65 | emb = torch.exp(exponent).to(timesteps.device) |
| 66 | emb = timesteps[:, None].float() * emb[None, :] |
| 67 | |
| 68 | # scale embeddings |
| 69 | emb = scale * emb |
| 70 | |
| 71 | # concat sine and cosine embeddings |
| 72 | emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1) |
| 73 | |
| 74 | # flip sine and cosine embeddings |
| 75 | if flip_sin_to_cos: |
| 76 | emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1) |
| 77 | |
| 78 | # zero pad |
| 79 | if embedding_dim % 2 == 1: |
| 80 | emb = torch.nn.functional.pad(emb, (0, 1, 0, 0)) |
| 81 | return emb |
| 82 | |
| 83 | |
| 84 | class TemporalTimesteps(torch.nn.Module): |