Create sinusoidal timestep embeddings. :param timesteps: a 1-D Tensor of N indices, one per batch element. These may be fractional. :param dim: the dimension of the output. :param max_period: controls the minimum frequency of the embeddings. :return: an [N
(timesteps, dim, max_period=10000, repeat_only=False, dtype=torch.float32)
| 178 | |
| 179 | |
| 180 | def timestep_embedding(timesteps, dim, max_period=10000, repeat_only=False, dtype=torch.float32): |
| 181 | """ |
| 182 | Create sinusoidal timestep embeddings. |
| 183 | :param timesteps: a 1-D Tensor of N indices, one per batch element. |
| 184 | These may be fractional. |
| 185 | :param dim: the dimension of the output. |
| 186 | :param max_period: controls the minimum frequency of the embeddings. |
| 187 | :return: an [N x dim] Tensor of positional embeddings. |
| 188 | """ |
| 189 | if not repeat_only: |
| 190 | half = dim // 2 |
| 191 | freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to( |
| 192 | device=timesteps.device |
| 193 | ) |
| 194 | args = timesteps[:, None].float() * freqs[None] |
| 195 | embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) |
| 196 | if dim % 2: |
| 197 | embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) |
| 198 | else: |
| 199 | embedding = repeat(timesteps, "b -> b d", d=dim) |
| 200 | return embedding.to(dtype) |
| 201 | |
| 202 | |
| 203 | def zero_module(module): |