86 lines
3.0 KiB
Python
86 lines
3.0 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
|
|
# https://github.com/facebookresearch/DiT
|
|
|
|
from typing import Union
|
|
|
|
import torch
|
|
from einops import rearrange
|
|
from torch import Tensor
|
|
|
|
# Ref: https://github.com/black-forest-labs/flux/blob/main/src/flux/math.py
|
|
# Ref: https://github.com/lucidrains/rotary-embedding-torch
|
|
|
|
|
|
def compute_rope_rotations(length: int,
|
|
dim: int,
|
|
theta: int,
|
|
*,
|
|
freq_scaling: float = 1.0,
|
|
device: Union[torch.device, str] = 'cpu') -> Tensor:
|
|
assert dim % 2 == 0
|
|
|
|
with torch.amp.autocast(device_type='cuda', enabled=False):
|
|
pos = torch.arange(length, dtype=torch.float32, device=device)
|
|
freqs = 1.0 / (theta**(torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim))
|
|
freqs *= freq_scaling
|
|
|
|
rot = torch.einsum('..., f -> ... f', pos, freqs)
|
|
rot = torch.stack([torch.cos(rot), -torch.sin(rot), torch.sin(rot), torch.cos(rot)], dim=-1)
|
|
rot = rearrange(rot, 'n d (i j) -> 1 n d i j', i=2, j=2)
|
|
return rot
|
|
|
|
|
|
def apply_rope(x: Tensor, rot: Tensor) -> tuple[Tensor, Tensor]:
|
|
with torch.amp.autocast(device_type='cuda', enabled=False):
|
|
_x = x.float()
|
|
_x = _x.view(*_x.shape[:-1], -1, 1, 2)
|
|
x_out = rot[..., 0] * _x[..., 0] + rot[..., 1] * _x[..., 1]
|
|
return x_out.reshape(*x.shape).to(dtype=x.dtype)
|
|
|
|
|
|
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.
|
|
"""
|
|
# https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
|
|
|
|
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
|