| import torch
|
| import torch.nn as nn
|
|
|
|
|
|
|
|
|
| class TimestepEmbedder(nn.Module):
|
| """
|
| Embeds scalar timesteps into vector representations.
|
| """
|
|
|
| def __init__(self, dim, frequency_embedding_size, max_period):
|
| super().__init__()
|
| self.mlp = nn.Sequential(
|
| nn.Linear(frequency_embedding_size, dim),
|
| nn.SiLU(),
|
| nn.Linear(dim, dim),
|
| )
|
| self.dim = dim
|
| self.max_period = max_period
|
| assert dim % 2 == 0, 'dim must be even.'
|
|
|
| with torch.autocast('cuda', enabled=False):
|
| self.freqs = nn.Buffer(
|
| 1.0 / (10000**(torch.arange(0, frequency_embedding_size, 2, dtype=torch.float32) /
|
| frequency_embedding_size)),
|
| persistent=False)
|
| freq_scale = 10000 / max_period
|
| self.freqs = freq_scale * self.freqs
|
|
|
| def timestep_embedding(self, t):
|
| """
|
| Create sinusoidal timestep embeddings.
|
| :param t: 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, D) Tensor of positional embeddings.
|
| """
|
|
|
|
|
| args = t[:, None].float() * self.freqs[None]
|
| embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
|
| return embedding
|
|
|
| def forward(self, t):
|
| t_freq = self.timestep_embedding(t).to(t.dtype)
|
| t_emb = self.mlp(t_freq)
|
| return t_emb
|
|
|
