File size: 7,109 Bytes
949ccb2 | 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 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 | from __future__ import annotations
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
def modulate(x, shift, scale):
return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
def timestep_embedding(t, dim, max_period=10000):
half = dim // 2
freqs = torch.exp(-math.log(max_period) * torch.arange(half, device=t.device) / half)
args = t[:, None].float() * freqs[None]
emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
return torch.cat([emb, torch.zeros_like(emb[:, :1])], dim=-1) if dim % 2 else emb
def rope_freqs(positions, dim, base=10000.0):
inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
return torch.outer(positions.float(), inv_freq)
def rope_cos_sin(freqs):
emb = torch.cat([freqs, freqs], dim=-1)
return emb.cos(), emb.sin()
def rotate_half(x):
x1, x2 = x.chunk(2, dim=-1)
return torch.cat([-x2, x1], dim=-1)
def apply_rope(x, cos, sin):
return x * cos + rotate_half(x) * sin
def apply_rope_2d(x, row_cos, row_sin, col_cos, col_sin):
x1, x2 = x.chunk(2, dim=-1)
return torch.cat([apply_rope(x1, row_cos, row_sin), apply_rope(x2, col_cos, col_sin)], dim=-1)
class RMSNormHead(nn.Module):
def __init__(self, head_dim, eps=1e-6):
super().__init__()
self.weight = nn.Parameter(torch.ones(head_dim))
self.eps = eps
def forward(self, x):
return x * x.pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt() * self.weight
class SwiGLU(nn.Module):
def __init__(self, dim, hidden):
super().__init__()
self.gate = nn.Linear(dim, hidden)
self.up = nn.Linear(dim, hidden)
self.down = nn.Linear(hidden, dim)
def forward(self, x):
return self.down(F.silu(self.gate(x)) * self.up(x))
class JointBlock(nn.Module):
def __init__(self, dim, heads, mlp_hidden):
super().__init__()
self.heads, self.head_dim = heads, dim // heads
self.norm1_img = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
self.norm1_txt = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
self.qkv_img = nn.Linear(dim, 3 * dim)
self.qkv_txt = nn.Linear(dim, 3 * dim)
self.qn_img, self.kn_img = RMSNormHead(self.head_dim), RMSNormHead(self.head_dim)
self.qn_txt, self.kn_txt = RMSNormHead(self.head_dim), RMSNormHead(self.head_dim)
self.proj_img, self.proj_txt = nn.Linear(dim, dim), nn.Linear(dim, dim)
self.norm2_img = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
self.norm2_txt = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
self.mlp_img, self.mlp_txt = SwiGLU(dim, mlp_hidden), SwiGLU(dim, mlp_hidden)
self.ada_img = nn.Sequential(nn.SiLU(), nn.Linear(dim, 6 * dim))
self.ada_txt = nn.Sequential(nn.SiLU(), nn.Linear(dim, 6 * dim))
def forward(self, img, txt, c, rope_img, rope_txt, key_valid):
s1i, sc1i, g1i, s2i, sc2i, g2i = self.ada_img(c).chunk(6, dim=-1)
s1t, sc1t, g1t, s2t, sc2t, g2t = self.ada_txt(c).chunk(6, dim=-1)
xi, xt = modulate(self.norm1_img(img), s1i, sc1i), modulate(self.norm1_txt(txt), s1t, sc1t)
b, ni, cdim = xi.shape; nt = xt.shape[1]; h, d = self.heads, self.head_dim
qi, ki, vi = self.qkv_img(xi).reshape(b, ni, 3, h, d).permute(2, 0, 3, 1, 4)
qt, kt, vt = self.qkv_txt(xt).reshape(b, nt, 3, h, d).permute(2, 0, 3, 1, 4)
qi, ki, qt, kt = self.qn_img(qi), self.kn_img(ki), self.qn_txt(qt), self.kn_txt(kt)
rc, rs, cc, cs = rope_img; tc, ts = rope_txt
qi, ki = apply_rope_2d(qi, rc, rs, cc, cs), apply_rope_2d(ki, rc, rs, cc, cs)
qt, kt = apply_rope(qt, tc, ts), apply_rope(kt, tc, ts)
q, k, v = torch.cat([qi, qt], 2), torch.cat([ki, kt], 2), torch.cat([vi, vt], 2)
o = F.scaled_dot_product_attention(q, k, v, attn_mask=key_valid[:, None, None, :])
o = o.transpose(1, 2).reshape(b, ni + nt, cdim)
oi, ot = o[:, :ni], o[:, ni:]
img = img + g1i[:, None] * self.proj_img(oi)
txt = txt + g1t[:, None] * self.proj_txt(ot)
img = img + g2i[:, None] * self.mlp_img(modulate(self.norm2_img(img), s2i, sc2i))
txt = txt + g2t[:, None] * self.mlp_txt(modulate(self.norm2_txt(txt), s2t, sc2t))
return img, txt
class MMDiT(nn.Module):
def __init__(self, latent_ch=4, latent_size=32, patch=2, dim=512, depth=16, heads=8,
t5_dim=768, clip_dim=512, t5_len=32, mlp_hidden=1408, repa_dim=384, repa_layer=8):
super().__init__()
self.latent_ch, self.latent_size, self.patch = latent_ch, latent_size, patch
self.grid, self.patch_dim, self.dim = latent_size // patch, latent_ch * patch * patch, dim
self.t5_len, self.repa_layer = t5_len, repa_layer
self.x_embed = nn.Linear(self.patch_dim, dim)
self.t_mlp = nn.Sequential(nn.Linear(dim, dim), nn.SiLU(), nn.Linear(dim, dim))
self.clip_proj, self.t5_proj = nn.Linear(clip_dim, dim), nn.Linear(t5_dim, dim)
self.blocks = nn.ModuleList([JointBlock(dim, heads, mlp_hidden) for _ in range(depth)])
self.norm_out = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
self.ada_out, self.head = nn.Sequential(nn.SiLU(), nn.Linear(dim, 2 * dim)), nn.Linear(dim, self.patch_dim)
self.repa_head = nn.Sequential(nn.Linear(dim, dim), nn.GELU(approximate="tanh"), nn.Linear(dim, repa_dim))
hd2 = (dim // heads) // 2
rows = torch.arange(self.grid).repeat_interleave(self.grid); cols = torch.arange(self.grid).repeat(self.grid)
for name, value in zip(("row_cos", "row_sin", "col_cos", "col_sin"), (*rope_cos_sin(rope_freqs(rows, hd2)), *rope_cos_sin(rope_freqs(cols, hd2)))):
self.register_buffer(name, value, persistent=False)
self.register_buffer("t_cos", rope_cos_sin(rope_freqs(torch.arange(t5_len), dim // heads))[0], persistent=False)
self.register_buffer("t_sin", rope_cos_sin(rope_freqs(torch.arange(t5_len), dim // heads))[1], persistent=False)
def patchify(self, x):
b, c, h, w = x.shape; p = self.patch
return x.reshape(b, c, h // p, p, w // p, p).permute(0, 2, 4, 1, 3, 5).reshape(b, (h // p) * (w // p), c * p * p)
def unpatchify(self, x):
b, _, _ = x.shape; p, g, c = self.patch, self.grid, self.latent_ch
return x.reshape(b, g, g, c, p, p).permute(0, 3, 1, 4, 2, 5).reshape(b, c, g * p, g * p)
def forward(self, x, t, t5_seq, t5_mask, clip_pool):
b = x.shape[0]; img = self.x_embed(self.patchify(x)); txt = self.t5_proj(t5_seq)
c = self.t_mlp(timestep_embedding(t, self.dim)) + self.clip_proj(clip_pool)
valid = torch.cat([torch.ones(b, img.shape[1], dtype=torch.bool, device=x.device), t5_mask.bool()], dim=1)
ri = (self.row_cos, self.row_sin, self.col_cos, self.col_sin); rt = (self.t_cos, self.t_sin)
for block in self.blocks: img, txt = block(img, txt, c, ri, rt, valid)
shift, scale = self.ada_out(c).chunk(2, dim=-1)
return self.unpatchify(self.head(modulate(self.norm_out(img), shift, scale)))
|