| 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))) |
|
|
