File size: 1,504 Bytes
a04730e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 | from __future__ import annotations
import math
import torch
from torch import nn
import torch.nn.functional as F
def sinusoidal_timestep_embedding(
timesteps: torch.Tensor,
dim: int,
max_period: int = 10_000,
) -> torch.Tensor:
"""
Create sinusoidal timestep embeddings
"""
half_dim = dim // 2
frequencies = torch.exp(
-math.log(max_period)
* torch.arange(
start=0,
end=half_dim,
dtype=torch.float32,
device=timesteps.device,
)
/ half_dim
)
args = timesteps.float()[:, None] * frequencies[None]
embedding = torch.cat(
[
torch.cos(args),
torch.sin(args),
],
dim=-1,
)
if dim % 2 == 1:
embedding = F.pad(embedding, (0, 1))
return embedding
class TimestepEmbedding(nn.Module):
def __init__(
self,
embedding_dim: int,
time_embed_dim: int,
):
super().__init__()
self.embedding_dim = embedding_dim
self.time_embed_dim = time_embed_dim
self.mlp = nn.Sequential(
nn.Linear(embedding_dim, time_embed_dim),
nn.SiLU(),
nn.Linear(time_embed_dim, time_embed_dim),
)
def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
emb = sinusoidal_timestep_embedding(
timesteps=timesteps,
dim=self.embedding_dim,
)
emb = self.mlp(emb)
return emb |