| 1 | import math |
| 2 | |
| 3 | import torch |
| 4 | |
| 5 | |
| 6 | def get_timestep_embedding( |
| 7 | timesteps: torch.Tensor, |
| 8 | embedding_dim: int, |
| 9 | flip_sin_to_cos: bool = False, |
| 10 | downscale_freq_shift: float = 1, |
| 11 | scale: float = 1, |
| 12 | max_period: int = 10000, |
| 13 | ) -> torch.Tensor: |
| 14 | """ |
| 15 | This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings. |
| 16 | Args |
| 17 | timesteps (torch.Tensor): |
| 18 | a 1-D Tensor of N indices, one per batch element. These may be fractional. |
| 19 | embedding_dim (int): |
| 20 | the dimension of the output. |
| 21 | flip_sin_to_cos (bool): |
| 22 | Whether the embedding order should be `cos, sin` (if True) or `sin, cos` (if False) |
| 23 | downscale_freq_shift (float): |
| 24 | Controls the delta between frequencies between dimensions |
| 25 | scale (float): |
| 26 | Scaling factor applied to the embeddings. |
| 27 | max_period (int): |
| 28 | Controls the maximum frequency of the embeddings |
| 29 | Returns |
| 30 | torch.Tensor: an [N x dim] Tensor of positional embeddings. |
| 31 | """ |
| 32 | assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array" |
| 33 | |
| 34 | half_dim = embedding_dim // 2 |
| 35 | exponent = -math.log(max_period) * torch.arange(start=0, end=half_dim, dtype=torch.float32, device=timesteps.device) |
| 36 | exponent = exponent / (half_dim - downscale_freq_shift) |
| 37 | |
| 38 | emb = torch.exp(exponent) |
| 39 | emb = timesteps[:, None].float() * emb[None, :] |
| 40 | |
| 41 | # scale embeddings |
| 42 | emb = scale * emb |
| 43 | |
| 44 | # concat sine and cosine embeddings |
| 45 | emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1) |
| 46 | |
| 47 | # flip sine and cosine embeddings |
| 48 | if flip_sin_to_cos: |
| 49 | emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1) |
| 50 | |
| 51 | # zero pad |
| 52 | if embedding_dim % 2 == 1: |
| 53 | emb = torch.nn.functional.pad(emb, (0, 1, 0, 0)) |
| 54 | return emb |
| 55 | |
| 56 | |
| 57 | class TimestepEmbedding(torch.nn.Module): |
| 58 | def __init__( |
| 59 | self, |
| 60 | in_channels: int, |
| 61 | time_embed_dim: int, |
| 62 | out_dim: int | None = None, |
| 63 | post_act_fn: str | None = None, |
| 64 | cond_proj_dim: int | None = None, |
| 65 | sample_proj_bias: bool = True, |
| 66 | ): |
| 67 | super().__init__() |
| 68 | |
| 69 | self.linear_1 = torch.nn.Linear(in_channels, time_embed_dim, sample_proj_bias) |
| 70 | |
| 71 | if cond_proj_dim is not None: |
| 72 | self.cond_proj = torch.nn.Linear(cond_proj_dim, in_channels, bias=False) |
| 73 | else: |
| 74 | self.cond_proj = None |
| 75 | |
| 76 | self.act = torch.nn.SiLU() |
| 77 | time_embed_dim_out = out_dim if out_dim is not None else time_embed_dim |
| 78 | |
| 79 | self.linear_2 = torch.nn.Linear(time_embed_dim, time_embed_dim_out, sample_proj_bias) |
| 80 | |
| 81 | if post_act_fn is None: |
| 82 | self.post_act = None |
| 83 | |
| 84 | def forward(self, sample: torch.Tensor, condition: torch.Tensor | None = None) -> torch.Tensor: |
| 85 | if condition is not None: |
| 86 | sample = sample + self.cond_proj(condition) |
| 87 | sample = self.linear_1(sample) |
| 88 | |
| 89 | if self.act is not None: |
| 90 | sample = self.act(sample) |
| 91 | |
| 92 | sample = self.linear_2(sample) |
| 93 | |
| 94 | if self.post_act is not None: |
| 95 | sample = self.post_act(sample) |
| 96 | return sample |
| 97 | |
| 98 | |
| 99 | class Timesteps(torch.nn.Module): |
| 100 | def __init__(self, num_channels: int, flip_sin_to_cos: bool, downscale_freq_shift: float, scale: int = 1): |
| 101 | super().__init__() |
| 102 | self.num_channels = num_channels |
| 103 | self.flip_sin_to_cos = flip_sin_to_cos |
| 104 | self.downscale_freq_shift = downscale_freq_shift |
| 105 | self.scale = scale |
| 106 | |
| 107 | def forward(self, timesteps: torch.Tensor) -> torch.Tensor: |
| 108 | t_emb = get_timestep_embedding( |
| 109 | timesteps, |
| 110 | self.num_channels, |
| 111 | flip_sin_to_cos=self.flip_sin_to_cos, |
| 112 | downscale_freq_shift=self.downscale_freq_shift, |
| 113 | scale=self.scale, |
| 114 | ) |
| 115 | return t_emb |
| 116 | |
| 117 | |
| 118 | class PixArtAlphaCombinedTimestepSizeEmbeddings(torch.nn.Module): |
| 119 | """ |
| 120 | For PixArt-Alpha. |
| 121 | Reference: |
| 122 | https://github.com/PixArt-alpha/PixArt-alpha/blob/0f55e922376d8b797edd44d25d0e7464b260dcab/diffusion/model/nets/PixArtMS.py#L164C9-L168C29 |
| 123 | """ |
| 124 | |
| 125 | def __init__( |
| 126 | self, |
| 127 | embedding_dim: int, |
| 128 | size_emb_dim: int, |
| 129 | ): |
| 130 | super().__init__() |
| 131 | |
| 132 | self.outdim = size_emb_dim |
| 133 | self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0) |
| 134 | self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) |
| 135 | |
| 136 | def forward( |
| 137 | self, |
| 138 | timestep: torch.Tensor, |
| 139 | hidden_dtype: torch.dtype, |
| 140 | ) -> torch.Tensor: |
| 141 | timesteps_proj = self.time_proj(timestep) |
| 142 | timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype)) # (N, D) |
| 143 | return timesteps_emb |
| 144 |