This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings. Args timesteps (torch.Tensor): a 1-D Tensor of N indices, one per batch element. These may be fractional. embedding_dim (int): the dim
(
timesteps: torch.Tensor,
embedding_dim: int,
flip_sin_to_cos: bool = False,
downscale_freq_shift: float = 1,
scale: float = 1,
max_period: int = 10000,
)
| 25 | |
| 26 | |
| 27 | def get_timestep_embedding( |
| 28 | timesteps: torch.Tensor, |
| 29 | embedding_dim: int, |
| 30 | flip_sin_to_cos: bool = False, |
| 31 | downscale_freq_shift: float = 1, |
| 32 | scale: float = 1, |
| 33 | max_period: int = 10000, |
| 34 | ): |
| 35 | """ |
| 36 | This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings. |
| 37 | |
| 38 | Args |
| 39 | timesteps (torch.Tensor): |
| 40 | a 1-D Tensor of N indices, one per batch element. These may be fractional. |
| 41 | embedding_dim (int): |
| 42 | the dimension of the output. |
| 43 | flip_sin_to_cos (bool): |
| 44 | Whether the embedding order should be `cos, sin` (if True) or `sin, cos` (if False) |
| 45 | downscale_freq_shift (float): |
| 46 | Controls the delta between frequencies between dimensions |
| 47 | scale (float): |
| 48 | Scaling factor applied to the embeddings. |
| 49 | max_period (int): |
| 50 | Controls the maximum frequency of the embeddings |
| 51 | Returns |
| 52 | torch.Tensor: an [N x dim] Tensor of positional embeddings. |
| 53 | """ |
| 54 | assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array" |
| 55 | |
| 56 | half_dim = embedding_dim // 2 |
| 57 | exponent = -math.log(max_period) * torch.arange( |
| 58 | start=0, end=half_dim, dtype=torch.float32, device=timesteps.device |
| 59 | ) |
| 60 | exponent = exponent / (half_dim - downscale_freq_shift) |
| 61 | |
| 62 | emb = torch.exp(exponent) |
| 63 | emb = timesteps[:, None].float() * emb[None, :] |
| 64 | |
| 65 | # scale embeddings |
| 66 | emb = scale * emb |
| 67 | |
| 68 | # concat sine and cosine embeddings |
| 69 | emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1) |
| 70 | |
| 71 | # flip sine and cosine embeddings |
| 72 | if flip_sin_to_cos: |
| 73 | emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1) |
| 74 | |
| 75 | # zero pad |
| 76 | if embedding_dim % 2 == 1: |
| 77 | emb = torch.nn.functional.pad(emb, (0, 1, 0, 0)) |
| 78 | return emb |
| 79 | |
| 80 | |
| 81 | def get_3d_sincos_pos_embed( |