|
from diffusers.configuration_utils import ConfigMixin, register_to_config |
|
from diffusers.models.modeling_utils import ModelMixin |
|
import einops |
|
import numpy as np |
|
import torch |
|
import torch.nn as nn |
|
import torch.nn.functional as F |
|
from .layers import LLamaFeedForward, RMSNorm |
|
|
|
|
|
def modulate(x, scale): |
|
return x * (1 + scale) |
|
|
|
class TimestepEmbedder(nn.Module): |
|
""" |
|
Embeds scalar timesteps into vector representations. |
|
""" |
|
def __init__(self, hidden_size, frequency_embedding_size=256): |
|
super().__init__() |
|
self.hidden_size = hidden_size |
|
self.frequency_embedding_size = frequency_embedding_size |
|
self.mlp = nn.Sequential( |
|
nn.Linear(self.frequency_embedding_size, self.hidden_size), |
|
nn.SiLU(), |
|
nn.Linear(self.hidden_size, self.hidden_size), |
|
) |
|
|
|
@staticmethod |
|
def timestep_embedding(t, dim, max_period=10000): |
|
""" |
|
Create sinusoidal timestep embeddings. |
|
:param t: a 1-D Tensor of N indices, one per batch element. |
|
: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. |
|
""" |
|
half = dim // 2 |
|
freqs = torch.exp( |
|
-np.log(max_period) * torch.arange(0, half, dtype=t.dtype) / half |
|
).to(t.device) |
|
args = t[:, :, None] * 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 |
|
|
|
def forward(self, t): |
|
t_freq = self.timestep_embedding(t, self.frequency_embedding_size) |
|
t_freq = t_freq.to(self.mlp[0].weight.dtype) |
|
return self.mlp(t_freq) |
|
|
|
class FinalLayer(nn.Module): |
|
def __init__(self, hidden_size, num_patches, out_channels): |
|
super().__init__() |
|
self.norm_final = nn.LayerNorm(hidden_size, eps=1e-6, elementwise_affine=False) |
|
self.linear = nn.Linear(hidden_size, num_patches * out_channels) |
|
self.adaLN_modulation = nn.Sequential( |
|
nn.SiLU(), |
|
nn.Linear(min(hidden_size, 1024), hidden_size), |
|
) |
|
|
|
def forward(self, x, c): |
|
scale = self.adaLN_modulation(c) |
|
x = modulate(self.norm_final(x), scale) |
|
x = self.linear(x) |
|
return x |
|
|
|
class Attention(nn.Module): |
|
def __init__( |
|
self, |
|
dim, |
|
n_heads, |
|
n_kv_heads=None, |
|
qk_norm=False, |
|
y_dim=0, |
|
base_seqlen=None, |
|
proportional_attn=False, |
|
attention_dropout=0.0, |
|
max_position_embeddings=384, |
|
): |
|
super().__init__() |
|
self.dim = dim |
|
self.n_heads = n_heads |
|
self.n_kv_heads = n_kv_heads or n_heads |
|
self.qk_norm = qk_norm |
|
self.y_dim = y_dim |
|
self.base_seqlen = base_seqlen |
|
self.proportional_attn = proportional_attn |
|
self.attention_dropout = attention_dropout |
|
self.max_position_embeddings = max_position_embeddings |
|
|
|
self.head_dim = dim // n_heads |
|
|
|
self.wq = nn.Linear(dim, n_heads * self.head_dim, bias=False) |
|
self.wk = nn.Linear(dim, self.n_kv_heads * self.head_dim, bias=False) |
|
self.wv = nn.Linear(dim, self.n_kv_heads * self.head_dim, bias=False) |
|
|
|
if y_dim > 0: |
|
self.wk_y = nn.Linear(y_dim, self.n_kv_heads * self.head_dim, bias=False) |
|
self.wv_y = nn.Linear(y_dim, self.n_kv_heads * self.head_dim, bias=False) |
|
self.gate = nn.Parameter(torch.zeros(n_heads)) |
|
|
|
self.wo = nn.Linear(n_heads * self.head_dim, dim, bias=False) |
|
|
|
if qk_norm: |
|
self.q_norm = nn.LayerNorm(self.n_heads * self.head_dim) |
|
self.k_norm = nn.LayerNorm(self.n_kv_heads * self.head_dim) |
|
if y_dim > 0: |
|
self.ky_norm = nn.LayerNorm(self.n_kv_heads * self.head_dim, eps=1e-6) |
|
else: |
|
self.ky_norm = nn.Identity() |
|
else: |
|
self.q_norm = nn.Identity() |
|
self.k_norm = nn.Identity() |
|
self.ky_norm = nn.Identity() |
|
|
|
|
|
@staticmethod |
|
def apply_rotary_emb(xq, xk, freqs_cis): |
|
|
|
|
|
xq_ = xq.float().reshape(*xq.shape[:-1], -1, 2) |
|
xk_ = xk.float().reshape(*xk.shape[:-1], -1, 2) |
|
|
|
xq_complex = torch.view_as_complex(xq_) |
|
xk_complex = torch.view_as_complex(xk_) |
|
|
|
freqs_cis = freqs_cis.unsqueeze(2) |
|
|
|
|
|
xq_out = xq_complex * freqs_cis |
|
xk_out = xk_complex * freqs_cis |
|
|
|
|
|
xq_out = torch.view_as_real(xq_out).flatten(-2) |
|
xk_out = torch.view_as_real(xk_out).flatten(-2) |
|
|
|
return xq_out.type_as(xq), xk_out.type_as(xk) |
|
|
|
def forward( |
|
self, |
|
x, |
|
x_mask, |
|
freqs_cis, |
|
y=None, |
|
y_mask=None, |
|
init_cache=False, |
|
): |
|
bsz, seqlen, _ = x.size() |
|
xq = self.wq(x) |
|
xk = self.wk(x) |
|
xv = self.wv(x) |
|
|
|
if x_mask is None: |
|
x_mask = torch.ones(bsz, seqlen, dtype=torch.bool, device=x.device) |
|
inp_dtype = xq.dtype |
|
|
|
xq = self.q_norm(xq) |
|
xk = self.k_norm(xk) |
|
|
|
xq = xq.view(bsz, seqlen, self.n_heads, self.head_dim) |
|
xk = xk.view(bsz, seqlen, self.n_kv_heads, self.head_dim) |
|
xv = xv.view(bsz, seqlen, self.n_kv_heads, self.head_dim) |
|
|
|
if self.n_kv_heads != self.n_heads: |
|
n_rep = self.n_heads // self.n_kv_heads |
|
xk = xk.repeat_interleave(n_rep, dim=2) |
|
xv = xv.repeat_interleave(n_rep, dim=2) |
|
|
|
freqs_cis = freqs_cis.to(xq.device) |
|
xq, xk = self.apply_rotary_emb(xq, xk, freqs_cis) |
|
|
|
output = ( |
|
F.scaled_dot_product_attention( |
|
xq.permute(0, 2, 1, 3), |
|
xk.permute(0, 2, 1, 3), |
|
xv.permute(0, 2, 1, 3), |
|
attn_mask=x_mask.bool().view(bsz, 1, 1, seqlen).expand(-1, self.n_heads, seqlen, -1), |
|
scale=None, |
|
) |
|
.permute(0, 2, 1, 3) |
|
.to(inp_dtype) |
|
) |
|
|
|
|
|
if hasattr(self, "wk_y"): |
|
yk = self.ky_norm(self.wk_y(y)).view(bsz, -1, self.n_kv_heads, self.head_dim) |
|
yv = self.wv_y(y).view(bsz, -1, self.n_kv_heads, self.head_dim) |
|
n_rep = self.n_heads // self.n_kv_heads |
|
|
|
|
|
|
|
if n_rep >= 1: |
|
yk = einops.repeat(yk, "b l h d -> b l (repeat h) d", repeat=n_rep) |
|
yv = einops.repeat(yv, "b l h d -> b l (repeat h) d", repeat=n_rep) |
|
output_y = F.scaled_dot_product_attention( |
|
xq.permute(0, 2, 1, 3), |
|
yk.permute(0, 2, 1, 3), |
|
yv.permute(0, 2, 1, 3), |
|
y_mask.view(bsz, 1, 1, -1).expand(bsz, self.n_heads, seqlen, -1).to(torch.bool), |
|
).permute(0, 2, 1, 3) |
|
output_y = output_y * self.gate.tanh().view(1, 1, -1, 1) |
|
output = output + output_y |
|
|
|
output = output.flatten(-2) |
|
output = self.wo(output) |
|
|
|
return output.to(inp_dtype) |
|
|
|
class TransformerBlock(nn.Module): |
|
""" |
|
Corresponds to the Transformer block in the JAX code. |
|
""" |
|
def __init__( |
|
self, |
|
dim, |
|
n_heads, |
|
n_kv_heads, |
|
multiple_of, |
|
ffn_dim_multiplier, |
|
norm_eps, |
|
qk_norm, |
|
y_dim, |
|
max_position_embeddings, |
|
): |
|
super().__init__() |
|
self.attention = Attention(dim, n_heads, n_kv_heads, qk_norm, y_dim=y_dim, max_position_embeddings=max_position_embeddings) |
|
self.feed_forward = LLamaFeedForward( |
|
dim=dim, |
|
hidden_dim=4 * dim, |
|
multiple_of=multiple_of, |
|
ffn_dim_multiplier=ffn_dim_multiplier, |
|
) |
|
self.attention_norm1 = RMSNorm(dim, eps=norm_eps) |
|
self.attention_norm2 = RMSNorm(dim, eps=norm_eps) |
|
self.ffn_norm1 = RMSNorm(dim, eps=norm_eps) |
|
self.ffn_norm2 = RMSNorm(dim, eps=norm_eps) |
|
self.adaLN_modulation = nn.Sequential( |
|
nn.SiLU(), |
|
nn.Linear(min(dim, 1024), 4 * dim), |
|
) |
|
self.attention_y_norm = RMSNorm(y_dim, eps=norm_eps) |
|
|
|
def forward( |
|
self, |
|
x, |
|
x_mask, |
|
freqs_cis, |
|
y, |
|
y_mask, |
|
adaln_input=None, |
|
): |
|
if adaln_input is not None: |
|
scales_gates = self.adaLN_modulation(adaln_input) |
|
|
|
|
|
scale_msa, gate_msa, scale_mlp, gate_mlp = scales_gates.chunk(4, dim=-1) |
|
x = x + torch.tanh(gate_msa) * self.attention_norm2( |
|
self.attention( |
|
modulate(self.attention_norm1(x), scale_msa), |
|
x_mask, |
|
freqs_cis, |
|
self.attention_y_norm(y), |
|
y_mask, |
|
) |
|
) |
|
x = x + torch.tanh(gate_mlp) * self.ffn_norm2( |
|
self.feed_forward( |
|
modulate(self.ffn_norm1(x), scale_mlp), |
|
) |
|
) |
|
else: |
|
x = x + self.attention_norm2( |
|
self.attention( |
|
self.attention_norm1(x), |
|
x_mask, |
|
freqs_cis, |
|
self.attention_y_norm(y), |
|
y_mask, |
|
) |
|
) |
|
x = x + self.ffn_norm2(self.feed_forward(self.ffn_norm1(x))) |
|
return x |
|
|
|
|
|
class NextDiT(ModelMixin, ConfigMixin): |
|
""" |
|
Diffusion model with a Transformer backbone for joint image-video training. |
|
""" |
|
@register_to_config |
|
def __init__( |
|
self, |
|
input_size=(1, 32, 32), |
|
patch_size=(1, 2, 2), |
|
in_channels=16, |
|
hidden_size=4096, |
|
depth=32, |
|
num_heads=32, |
|
num_kv_heads=None, |
|
multiple_of=256, |
|
ffn_dim_multiplier=None, |
|
norm_eps=1e-5, |
|
pred_sigma=False, |
|
caption_channels=4096, |
|
qk_norm=False, |
|
norm_type="rms", |
|
model_max_length=120, |
|
rotary_max_length=384, |
|
rotary_max_length_t=None |
|
): |
|
super().__init__() |
|
self.input_size = input_size |
|
self.patch_size = patch_size |
|
self.in_channels = in_channels |
|
self.hidden_size = hidden_size |
|
self.depth = depth |
|
self.num_heads = num_heads |
|
self.num_kv_heads = num_kv_heads or num_heads |
|
self.multiple_of = multiple_of |
|
self.ffn_dim_multiplier = ffn_dim_multiplier |
|
self.norm_eps = norm_eps |
|
self.pred_sigma = pred_sigma |
|
self.caption_channels = caption_channels |
|
self.qk_norm = qk_norm |
|
self.norm_type = norm_type |
|
self.model_max_length = model_max_length |
|
self.rotary_max_length = rotary_max_length |
|
self.rotary_max_length_t = rotary_max_length_t |
|
self.out_channels = in_channels * 2 if pred_sigma else in_channels |
|
|
|
self.x_embedder = nn.Linear(np.prod(self.patch_size) * in_channels, hidden_size) |
|
|
|
self.t_embedder = TimestepEmbedder(min(hidden_size, 1024)) |
|
self.y_embedder = nn.Sequential( |
|
nn.LayerNorm(caption_channels, eps=1e-6), |
|
nn.Linear(caption_channels, min(hidden_size, 1024)), |
|
) |
|
|
|
self.layers = nn.ModuleList([ |
|
TransformerBlock( |
|
dim=hidden_size, |
|
n_heads=num_heads, |
|
n_kv_heads=self.num_kv_heads, |
|
multiple_of=multiple_of, |
|
ffn_dim_multiplier=ffn_dim_multiplier, |
|
norm_eps=norm_eps, |
|
qk_norm=qk_norm, |
|
y_dim=caption_channels, |
|
max_position_embeddings=rotary_max_length, |
|
) |
|
for _ in range(depth) |
|
]) |
|
|
|
self.final_layer = FinalLayer( |
|
hidden_size=hidden_size, |
|
num_patches=np.prod(patch_size), |
|
out_channels=self.out_channels, |
|
) |
|
|
|
assert (hidden_size // num_heads) % 6 == 0, "3d rope needs head dim to be divisible by 6" |
|
|
|
self.freqs_cis = self.precompute_freqs_cis( |
|
hidden_size // num_heads, |
|
self.rotary_max_length, |
|
end_t=self.rotary_max_length_t |
|
) |
|
|
|
def to(self, *args, **kwargs): |
|
self = super().to(*args, **kwargs) |
|
|
|
return self |
|
|
|
@staticmethod |
|
def precompute_freqs_cis( |
|
dim: int, |
|
end: int, |
|
end_t: int = None, |
|
theta: float = 10000.0, |
|
scale_factor: float = 1.0, |
|
scale_watershed: float = 1.0, |
|
timestep: float = 1.0, |
|
): |
|
if timestep < scale_watershed: |
|
linear_factor = scale_factor |
|
ntk_factor = 1.0 |
|
else: |
|
linear_factor = 1.0 |
|
ntk_factor = scale_factor |
|
|
|
theta = theta * ntk_factor |
|
freqs = 1.0 / (theta ** (torch.arange(0, dim, 6)[: (dim // 6)] / dim)) / linear_factor |
|
|
|
timestep = torch.arange(end, dtype=torch.float32) |
|
freqs = torch.outer(timestep, freqs).float() |
|
freqs_cis = torch.exp(1j * freqs) |
|
|
|
if end_t is not None: |
|
freqs_t = 1.0 / (theta ** (torch.arange(0, dim, 6)[: (dim // 6)] / dim)) / linear_factor |
|
timestep_t = torch.arange(end_t, dtype=torch.float32) |
|
freqs_t = torch.outer(timestep_t, freqs_t).float() |
|
freqs_cis_t = torch.exp(1j * freqs_t) |
|
freqs_cis_t = freqs_cis_t.view(end_t, 1, 1, dim // 6).repeat(1, end, end, 1) |
|
else: |
|
end_t = end |
|
freqs_cis_t = freqs_cis.view(end_t, 1, 1, dim // 6).repeat(1, end, end, 1) |
|
|
|
freqs_cis_h = freqs_cis.view(1, end, 1, dim // 6).repeat(end_t, 1, end, 1) |
|
freqs_cis_w = freqs_cis.view(1, 1, end, dim // 6).repeat(end_t, end, 1, 1) |
|
freqs_cis = torch.cat([freqs_cis_t, freqs_cis_h, freqs_cis_w], dim=-1).view(end_t, end, end, -1) |
|
return freqs_cis |
|
|
|
def forward( |
|
self, |
|
samples, |
|
timesteps, |
|
encoder_hidden_states, |
|
encoder_attention_mask, |
|
scale_factor: float = 1.0, |
|
scale_watershed: float = 1.0, |
|
): |
|
if samples.ndim == 4: |
|
samples = samples[:, None, ...] |
|
|
|
precomputed_freqs_cis = None |
|
if scale_factor != 1 or scale_watershed != 1: |
|
precomputed_freqs_cis = self.precompute_freqs_cis( |
|
self.hidden_size // self.num_heads, |
|
self.rotary_max_length, |
|
end_t=self.rotary_max_length_t, |
|
scale_factor=scale_factor, |
|
scale_watershed=scale_watershed, |
|
timestep=torch.max(timesteps.cpu()).item() |
|
) |
|
|
|
if len(timesteps.shape) == 5: |
|
t, *_ = self.patchify(timesteps, precomputed_freqs_cis) |
|
timesteps = t.mean(dim=-1) |
|
elif len(timesteps.shape) == 1: |
|
timesteps = timesteps[:, None, None, None, None].expand_as(samples) |
|
t, *_ = self.patchify(timesteps, precomputed_freqs_cis) |
|
timesteps = t.mean(dim=-1) |
|
samples, T, H, W, freqs_cis = self.patchify(samples, precomputed_freqs_cis) |
|
samples = self.x_embedder(samples) |
|
t = self.t_embedder(timesteps) |
|
|
|
encoder_attention_mask_float = encoder_attention_mask[..., None].float() |
|
encoder_hidden_states_pool = (encoder_hidden_states * encoder_attention_mask_float).sum(dim=1) / (encoder_attention_mask_float.sum(dim=1) + 1e-8) |
|
encoder_hidden_states_pool = encoder_hidden_states_pool.to(samples.dtype) |
|
y = self.y_embedder(encoder_hidden_states_pool) |
|
y = y.unsqueeze(1).expand(-1, samples.size(1), -1) |
|
|
|
adaln_input = t + y |
|
|
|
for block in self.layers: |
|
samples = block(samples, None, freqs_cis, encoder_hidden_states, encoder_attention_mask, adaln_input) |
|
|
|
samples = self.final_layer(samples, adaln_input) |
|
samples = self.unpatchify(samples, T, H, W) |
|
|
|
return samples |
|
|
|
def patchify(self, x, precompute_freqs_cis=None): |
|
|
|
B, T, C, H, W = x.size() |
|
pT, pH, pW = self.patch_size |
|
x = x.view(B, T // pT, pT, C, H // pH, pH, W // pW, pW) |
|
x = x.permute(0, 1, 4, 6, 2, 5, 7, 3) |
|
x = x.reshape(B, -1, pT * pH * pW * C) |
|
if precompute_freqs_cis is None: |
|
freqs_cis = self.freqs_cis[: T // pT, :H // pH, :W // pW].reshape(-1, * self.freqs_cis.shape[3:])[None].to(x.device) |
|
else: |
|
freqs_cis = precompute_freqs_cis[: T // pT, :H // pH, :W // pW].reshape(-1, * precompute_freqs_cis.shape[3:])[None].to(x.device) |
|
return x, T // pT, H // pH, W // pW, freqs_cis |
|
|
|
def unpatchify(self, x, T, H, W): |
|
B = x.size(0) |
|
C = self.out_channels |
|
pT, pH, pW = self.patch_size |
|
x = x.view(B, T, H, W, pT, pH, pW, C) |
|
x = x.permute(0, 1, 4, 7, 2, 5, 3, 6) |
|
x = x.reshape(B, T * pT, C, H * pH, W * pW) |
|
return x |
|
|