File size: 4,474 Bytes
28555c1 | 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 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 | import math
import torch
from torch import nn
def timestep_embedding(t, dim, max_period=10000):
"""
Create sinusoidal timestep embeddings.
Args:
t (torch.Tensor): a 1-D Tensor of N indices, one per batch element. These may be fractional.
dim (int): the dimension of the output.
max_period (int): controls the minimum frequency of the embeddings.
Returns:
embedding (torch.Tensor): An (N, D) Tensor of positional embeddings.
.. ref_link: https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
"""
half = dim // 2
freqs = torch.exp(
-math.log(max_period)
* torch.arange(start=0, end=half, dtype=torch.float32)
/ half
).to(device=t.device)
args = t[:, None].float() * freqs[None]
embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
if dim % 2:
embedding = torch.cat(
[embedding, torch.zeros_like(embedding[:, :1])], dim=-1
)
return embedding
class TimestepEmbedder(nn.Module):
"""
Embeds scalar timesteps into vector representations.
"""
def __init__(self,
hidden_size,
act_layer=nn.GELU,
frequency_embedding_size=256,
max_period=10000,
out_size=None,
dtype=None,
device=None
):
factory_kwargs = {'dtype': dtype, 'device': device}
super().__init__()
self.frequency_embedding_size = frequency_embedding_size
self.max_period = max_period
if out_size is None:
out_size = hidden_size
self.mlp = nn.Sequential(
nn.Linear(frequency_embedding_size, hidden_size, bias=True, **factory_kwargs),
act_layer(),
nn.Linear(hidden_size, out_size, bias=True, **factory_kwargs),
)
nn.init.normal_(self.mlp[0].weight, std=0.02)
nn.init.normal_(self.mlp[2].weight, std=0.02)
def forward(self, t):
t_freq = timestep_embedding(t, self.frequency_embedding_size, self.max_period).type(self.mlp[0].weight.dtype)
t_emb = self.mlp(t_freq)
return t_emb
class LinearPatchEmbed(nn.Module):
"""Linear projection for SigLIP latent tokens: (B, N, in_chans) -> (B, N, embed_dim)"""
def __init__(self, in_chans, embed_dim, act_layer=nn.SiLU, bias=True,
device=None, dtype=None):
factory_kwargs = {'dtype': dtype, 'device': device}
super().__init__()
self.proj = nn.Linear(in_chans, embed_dim, bias=bias, **factory_kwargs)
self.norm_final = nn.LayerNorm(embed_dim, elementwise_affine=False, eps=1e-6,
**factory_kwargs)
self.adaLN_modulation = nn.Sequential(
act_layer(),
nn.Linear(embed_dim, 2 * embed_dim, bias=True, **factory_kwargs)
)
nn.init.zeros_(self.adaLN_modulation[1].weight)
nn.init.zeros_(self.adaLN_modulation[1].bias)
def forward(self, x, t):
x = self.proj(x)
shift, scale = self.adaLN_modulation(t).chunk(2, dim=1)
dtype = x.dtype
x = self.norm_final(x) * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
return x.to(dtype), 0, 0
class LinearFinalLayer(nn.Module):
"""Linear projection back to SigLIP latent space: (B, N, hidden) -> (B, N, out_channels)"""
def __init__(self, hidden_size, out_channels, act_layer=nn.SiLU,
device=None, dtype=None):
factory_kwargs = {'device': device, 'dtype': dtype}
super().__init__()
self.out_channels = out_channels
self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6,
**factory_kwargs)
self.linear = nn.Linear(hidden_size, out_channels, bias=True, **factory_kwargs)
nn.init.zeros_(self.linear.weight)
nn.init.zeros_(self.linear.bias)
self.adaLN_modulation = nn.Sequential(
act_layer(),
nn.Linear(hidden_size, 2 * hidden_size, bias=True, **factory_kwargs)
)
nn.init.zeros_(self.adaLN_modulation[1].weight)
nn.init.zeros_(self.adaLN_modulation[1].bias)
def forward(self, x, t, token_h=0, token_w=0):
shift, scale = self.adaLN_modulation(t).chunk(2, dim=1)
x = self.norm_final(x) * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
x = self.linear(x)
return x
|