File size: 4,736 Bytes
edbe9c6 | 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 114 115 116 117 118 119 120 121 122 123 | import torch
import torch.nn as nn
import numpy as np
import math
from .modules import FinalLayer, DiTBlock
from x_transformers.x_transformers import RotaryEmbedding
class SinusPositionEmbedding(nn.Module):
def __init__(self, dim):
super().__init__()
self.dim = dim
def forward(self, x, scale=1000):
device = x.device
half_dim = self.dim // 2
emb = math.log(10000) / (half_dim - 1)
emb = torch.exp(torch.arange(half_dim, device=device).float() * -emb)
emb = scale * x.unsqueeze(1) * emb.unsqueeze(0)
emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
return emb
class TimestepEmbedder(nn.Module):
def __init__(self, dim, freq_embed_dim=256):
super().__init__()
self.time_embed = SinusPositionEmbedding(freq_embed_dim)
self.time_mlp = nn.Sequential(nn.Linear(freq_embed_dim, dim), nn.SiLU(), nn.Linear(dim, dim))
def forward(self, timestep):
time_hidden = self.time_embed(timestep)
time_hidden = time_hidden.to(timestep.dtype)
time = self.time_mlp(time_hidden) # b d
return time
class CondEmbedder(nn.Module):
def __init__(self, input_feature_size, hidden_size, dropout_prob):
super().__init__()
self.dropout_prob = dropout_prob
self.cond_embedder = nn.Linear(input_feature_size, hidden_size)
def cond_drop(self, llm_cond, fake_latent):
bsz = llm_cond.shape[0]
drop_latent_mask = torch.rand(bsz) < self.dropout_prob
drop_latent_mask = drop_latent_mask.unsqueeze(-1).unsqueeze(-1).to(llm_cond.dtype).to(llm_cond.device)
# fake_latent = torch.zeros(llm_cond.shape).to(llm_cond.device)
llm_cond = drop_latent_mask * fake_latent.to(llm_cond.device) + (1 - drop_latent_mask) * llm_cond
return llm_cond
def forward(self, llm_cond, train, fake_latent):
use_dropout = self.dropout_prob > 0
if train and use_dropout:
llm_cond = self.cond_drop(llm_cond, fake_latent)
llm_cond = self.cond_embedder(llm_cond)
return llm_cond
class DiT(nn.Module):
def __init__(
self,
in_channels=4,
hidden_size=1024,
depth=28,
num_heads=16,
mlp_ratio=4.0,
llm_cond_dim=896,
cfg_dropout_prob=0.1,
cfg_trainable=0,
**kwargs,
):
super().__init__()
self.in_channels = in_channels
self.out_channels = in_channels
self.num_heads = num_heads
self.t_embedder = TimestepEmbedder(hidden_size)
self.x_embedder = nn.Linear(in_channels, hidden_size)
self.c_embedder = CondEmbedder(llm_cond_dim, hidden_size, cfg_dropout_prob)
self.hidden_size = hidden_size
self.rotary_embed = RotaryEmbedding(hidden_size // num_heads)
self.blocks = nn.ModuleList([
DiTBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio, **kwargs) for _ in range(depth)
])
self.final_layer = FinalLayer(hidden_size, self.out_channels)
if cfg_trainable:
self.fake_latent = nn.Parameter(torch.empty(1, 1, llm_cond_dim))
torch.nn.init.normal_(self.fake_latent, std=.02)
else:
self.fake_latent = torch.zeros(1, 1, llm_cond_dim)
def forward(self, x, t, c, latent_history, mask=None):
t = self.t_embedder(t).unsqueeze(1)
x_now = self.x_embedder(x)
x_history = self.x_embedder(latent_history)
x = torch.cat([x_history, x_now], dim=1)
c = self.c_embedder(c, self.training, self.fake_latent)
y = t + c
x = torch.cat([y, x], dim=1)
rope = self.rotary_embed.forward_from_seq_len(x.shape[1])
if mask is not None:
mask_pad = mask.clone().detach()[:, :1].expand(-1, x_history.shape[1] + c.shape[1])
mask = torch.cat([mask_pad, mask], dim=-1)
for block in self.blocks:
x = block(x, mask, rope)
x = self.final_layer(x)
return x
def forward_with_cfg(self, x, t, c, cfg_scale, latent_history, patch_size):
if not cfg_scale == 1:
x = torch.cat([x, x], dim=0)
latent_history = torch.cat([latent_history, latent_history], dim=0)
# fake_latent = torch.zeros(c.shape).to(c.device)
fake_latent = self.fake_latent.repeat(c.shape[0], c.shape[1], 1).to(c.device)
c = torch.cat([c, fake_latent], dim=0)
if t.ndim == 0:
t = t.repeat(x.shape[0])
model_out = self.forward(x, t, c, latent_history)
return model_out[:, -patch_size:, :]
|