| from __future__ import annotations |
|
|
| import math |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| import torch.utils.checkpoint |
|
|
| 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) |
| if dim % 2: |
| emb = torch.cat([emb, torch.zeros_like(emb[:, :1])], dim=-1) |
| return 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) |
| x1 = apply_rope(x1, row_cos, row_sin) |
| x2 = apply_rope(x2, col_cos, col_sin) |
| return torch.cat([x1, x2], 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): |
| n = x.pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt() |
| return x * n * 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 = heads |
| self.head_dim = 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 = RMSNormHead(self.head_dim) |
| self.kn_img = RMSNormHead(self.head_dim) |
| self.qn_txt = RMSNormHead(self.head_dim) |
| self.kn_txt = RMSNormHead(self.head_dim) |
| self.proj_img = nn.Linear(dim, dim) |
| self.proj_txt = 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 = SwiGLU(dim, mlp_hidden) |
| self.mlp_txt = 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 = modulate(self.norm1_img(img), s1i, sc1i) |
| xt = modulate(self.norm1_txt(txt), s1t, sc1t) |
|
|
| B, Ni, C = 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 = self.qn_img(qi), self.kn_img(ki) |
| qt, kt = self.qn_txt(qt), self.kn_txt(kt) |
|
|
| row_cos, row_sin, col_cos, col_sin = rope_img |
| qi = apply_rope_2d(qi, row_cos, row_sin, col_cos, col_sin) |
| ki = apply_rope_2d(ki, row_cos, row_sin, col_cos, col_sin) |
|
|
| t_cos, t_sin = rope_txt |
| qt = apply_rope(qt, t_cos, t_sin) |
| kt = apply_rope(kt, t_cos, t_sin) |
|
|
| q = torch.cat([qi, qt], dim=2) |
| k = torch.cat([ki, kt], dim=2) |
| v = torch.cat([vi, vt], dim=2) |
|
|
| mask = key_valid[:, None, None, :] |
| o = F.scaled_dot_product_attention(q, k, v, attn_mask=mask) |
| o = o.transpose(1, 2).reshape(B, Ni + Nt, C) |
| oi, ot = o[:, :Ni], o[:, Ni:] |
|
|
| img = img + g1i.unsqueeze(1) * self.proj_img(oi) |
| txt = txt + g1t.unsqueeze(1) * self.proj_txt(ot) |
|
|
| img = img + g2i.unsqueeze(1) * self.mlp_img(modulate(self.norm2_img(img), s2i, sc2i)) |
| txt = txt + g2t.unsqueeze(1) * 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 = latent_ch |
| self.latent_size = latent_size |
| self.patch = patch |
| self.grid = latent_size // patch |
| self.patch_dim = latent_ch * patch * patch |
| self.dim = dim |
| self.depth = depth |
| self.heads = heads |
| self.head_dim = dim // heads |
| self.t5_len = t5_len |
| self.repa_layer = 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 = nn.Linear(clip_dim, dim) |
| self.t5_proj = 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 = nn.Sequential(nn.SiLU(), nn.Linear(dim, 2 * dim)) |
| self.head = 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 = self.head_dim // 2 |
| rows = torch.arange(self.grid).repeat_interleave(self.grid) |
| cols = torch.arange(self.grid).repeat(self.grid) |
| row_cos, row_sin = rope_cos_sin(rope_freqs(rows, hd2)) |
| col_cos, col_sin = rope_cos_sin(rope_freqs(cols, hd2)) |
| self.register_buffer("row_cos", row_cos, persistent=False) |
| self.register_buffer("row_sin", row_sin, persistent=False) |
| self.register_buffer("col_cos", col_cos, persistent=False) |
| self.register_buffer("col_sin", col_sin, persistent=False) |
| t_cos, t_sin = rope_cos_sin(rope_freqs(torch.arange(t5_len), self.head_dim)) |
| self.register_buffer("t_cos", t_cos, persistent=False) |
| self.register_buffer("t_sin", t_sin, persistent=False) |
|
|
| self._init() |
|
|
| def _init(self): |
| for m in self.modules(): |
| if isinstance(m, nn.Linear): |
| nn.init.xavier_uniform_(m.weight) |
| if m.bias is not None: |
| nn.init.zeros_(m.bias) |
| for b in self.blocks: |
| nn.init.zeros_(b.ada_img[-1].weight); nn.init.zeros_(b.ada_img[-1].bias) |
| nn.init.zeros_(b.ada_txt[-1].weight); nn.init.zeros_(b.ada_txt[-1].bias) |
| nn.init.zeros_(self.ada_out[-1].weight); nn.init.zeros_(self.ada_out[-1].bias) |
| nn.init.zeros_(self.head.weight); nn.init.zeros_(self.head.bias) |
|
|
| def patchify(self, x): |
| B, C, H, W = x.shape |
| p = self.patch |
| x = x.reshape(B, C, H // p, p, W // p, p) |
| x = x.permute(0, 2, 4, 1, 3, 5).reshape(B, (H // p) * (W // p), C * p * p) |
| return x |
|
|
| def unpatchify(self, x): |
| B, N, _ = x.shape |
| p = self.patch |
| g = self.grid |
| C = self.latent_ch |
| x = x.reshape(B, g, g, C, p, p).permute(0, 3, 1, 4, 2, 5) |
| return x.reshape(B, C, g * p, g * p) |
|
|
| def forward(self, x, t, t5_seq, t5_mask, clip_pool, return_repa=False, use_checkpoint=False): |
| 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) |
|
|
| key_valid = torch.cat([ |
| torch.ones(B, img.shape[1], dtype=torch.bool, device=x.device), |
| t5_mask.bool(), |
| ], dim=1) |
|
|
| rope_img = (self.row_cos, self.row_sin, self.col_cos, self.col_sin) |
| rope_txt = (self.t_cos, self.t_sin) |
|
|
| repa_hidden = None |
| for i, blk in enumerate(self.blocks): |
| if use_checkpoint and self.training: |
| img, txt = torch.utils.checkpoint.checkpoint( |
| blk, img, txt, c, rope_img, rope_txt, key_valid, use_reentrant=False) |
| else: |
| img, txt = blk(img, txt, c, rope_img, rope_txt, key_valid) |
| if return_repa and i == self.repa_layer: |
| repa_hidden = img |
|
|
| shift, scale = self.ada_out(c).chunk(2, dim=-1) |
| img = modulate(self.norm_out(img), shift, scale) |
| out = self.unpatchify(self.head(img)) |
|
|
| if return_repa: |
| return out, self.repa_head(repa_hidden) |
| return out |
|
|
| def num_params(self): |
| return sum(p.numel() for p in self.parameters()) |
|
|
| def num_backbone_params(self): |
| return sum(p.numel() for n, p in self.named_parameters() if not n.startswith("repa_head")) |
|
|
