Spaces:
Running on Zero
Running on Zero
| # Trimmed port of GenFirst's Prior_only/models/mmdit.py (from Stability-AI/sd3-ref, MIT). | |
| # Only the configuration used by the released MMDiT-XL/2 text-to-image prior is kept: | |
| # LayerNorm, GELU MLP, qkv bias, no qk-norm, no registers, CLIP-L (768-d) context. | |
| # The REPA projector heads are training-only and are dropped. | |
| import math | |
| import numpy as np | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| def attention(q, k, v, heads): | |
| b, _, dim = q.shape | |
| dh = dim // heads | |
| q, k, v = (t.view(b, -1, heads, dh).transpose(1, 2) for t in (q, k, v)) | |
| out = F.scaled_dot_product_attention(q, k, v) | |
| return out.transpose(1, 2).reshape(b, -1, heads * dh) | |
| def modulate(x, shift, scale): | |
| return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) | |
| def get_2d_sincos_pos_embed(embed_dim, grid_size): | |
| grid_h = np.arange(grid_size, dtype=np.float32) | |
| grid_w = np.arange(grid_size, dtype=np.float32) | |
| grid = np.stack(np.meshgrid(grid_w, grid_h), axis=0).reshape([2, 1, grid_size, grid_size]) | |
| def _1d(d, pos): | |
| omega = np.arange(d // 2, dtype=np.float64) / (d / 2.0) | |
| omega = 1.0 / 10000**omega | |
| out = np.einsum("m,d->md", pos.reshape(-1), omega) | |
| return np.concatenate([np.sin(out), np.cos(out)], axis=1) | |
| return np.concatenate([_1d(embed_dim // 2, grid[0]), _1d(embed_dim // 2, grid[1])], axis=1) | |
| class PatchEmbed(nn.Module): | |
| def __init__(self, img_size, patch_size, in_chans, embed_dim): | |
| super().__init__() | |
| self.patch_size = (patch_size, patch_size) | |
| self.num_patches = (img_size // patch_size) ** 2 | |
| self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size, bias=True) | |
| def forward(self, x): | |
| return self.proj(x).flatten(2).transpose(1, 2) | |
| class TimestepEmbedder(nn.Module): | |
| def __init__(self, hidden_size, frequency_embedding_size=256): | |
| super().__init__() | |
| self.mlp = nn.Sequential( | |
| nn.Linear(frequency_embedding_size, hidden_size, bias=True), | |
| nn.SiLU(), | |
| nn.Linear(hidden_size, hidden_size, bias=True), | |
| ) | |
| self.frequency_embedding_size = frequency_embedding_size | |
| def timestep_embedding(t, dim, max_period=10000): | |
| half = dim // 2 | |
| freqs = torch.exp( | |
| -math.log(max_period) * torch.arange(0, half, dtype=torch.float32) / half | |
| ).to(device=t.device) | |
| args = t[:, None].float() * freqs[None] | |
| emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) | |
| if torch.is_floating_point(t): | |
| emb = emb.to(dtype=t.dtype) | |
| return emb | |
| def forward(self, t, dtype): | |
| return self.mlp(self.timestep_embedding(t, self.frequency_embedding_size).to(dtype)) | |
| class Mlp(nn.Module): | |
| def __init__(self, dim, hidden): | |
| super().__init__() | |
| self.fc1 = nn.Linear(dim, hidden) | |
| self.act = nn.GELU(approximate="tanh") | |
| self.fc2 = nn.Linear(hidden, dim) | |
| def forward(self, x): | |
| return self.fc2(self.act(self.fc1(x))) | |
| class SelfAttention(nn.Module): | |
| def __init__(self, dim, num_heads, pre_only=False): | |
| super().__init__() | |
| self.num_heads = num_heads | |
| self.head_dim = dim // num_heads | |
| self.qkv = nn.Linear(dim, dim * 3, bias=True) | |
| if not pre_only: | |
| self.proj = nn.Linear(dim, dim) | |
| def pre_attention(self, x): | |
| qkv = self.qkv(x) | |
| qkv = qkv.reshape(qkv.shape[0], qkv.shape[1], 3, -1, self.head_dim).movedim(2, 0) | |
| q, k, v = qkv[0], qkv[1], qkv[2] | |
| B, L = q.shape[0], q.shape[1] | |
| return q.reshape(B, L, -1), k.reshape(B, L, -1), v | |
| class DismantledBlock(nn.Module): | |
| def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, pre_only=False): | |
| super().__init__() | |
| self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) | |
| self.attn = SelfAttention(hidden_size, num_heads, pre_only=pre_only) | |
| if not pre_only: | |
| self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) | |
| self.mlp = Mlp(hidden_size, int(hidden_size * mlp_ratio)) | |
| n_mods = 2 if pre_only else 6 | |
| self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, n_mods * hidden_size, bias=True)) | |
| self.pre_only = pre_only | |
| def pre_attention(self, x, c): | |
| if not self.pre_only: | |
| shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(6, dim=1) | |
| qkv = self.attn.pre_attention(modulate(self.norm1(x), shift_msa, scale_msa)) | |
| return qkv, (x, gate_msa, shift_mlp, scale_mlp, gate_mlp) | |
| shift_msa, scale_msa = self.adaLN_modulation(c).chunk(2, dim=1) | |
| return self.attn.pre_attention(modulate(self.norm1(x), shift_msa, scale_msa)), None | |
| def post_attention(self, attn, x, gate_msa, shift_mlp, scale_mlp, gate_mlp): | |
| x = x + gate_msa.unsqueeze(1) * self.attn.proj(attn) | |
| x = x + gate_mlp.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp)) | |
| return x | |
| class JointBlock(nn.Module): | |
| def __init__(self, hidden_size, num_heads, pre_only): | |
| super().__init__() | |
| self.context_block = DismantledBlock(hidden_size, num_heads, pre_only=pre_only) | |
| self.x_block = DismantledBlock(hidden_size, num_heads, pre_only=False) | |
| def forward(self, context, x, c): | |
| cqkv, cint = self.context_block.pre_attention(context, c) | |
| xqkv, xint = self.x_block.pre_attention(x, c) | |
| q, k, v = (torch.cat((cqkv[t], xqkv[t]), dim=1) for t in range(3)) | |
| attn = attention(q, k, v, self.x_block.attn.num_heads) | |
| n_ctx = cqkv[0].shape[1] | |
| context_attn, x_attn = attn[:, :n_ctx], attn[:, n_ctx:] | |
| context = None if self.context_block.pre_only else self.context_block.post_attention(context_attn, *cint) | |
| x = self.x_block.post_attention(x_attn, *xint) | |
| return context, x | |
| class FinalLayer(nn.Module): | |
| def __init__(self, hidden_size, patch_size, out_channels): | |
| super().__init__() | |
| self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) | |
| self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True) | |
| self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True)) | |
| def forward(self, x, c): | |
| shift, scale = self.adaLN_modulation(c).chunk(2, dim=1) | |
| return self.linear(modulate(self.norm_final(x), shift, scale)) | |
| class MMDiT(nn.Module): | |
| def __init__(self, input_size=64, patch_size=2, in_channels=16, depth=28, | |
| hidden_size=1152, num_heads=16, context_dim=768): | |
| super().__init__() | |
| self.out_channels = in_channels | |
| self.x_embedder = PatchEmbed(input_size, patch_size, in_channels, hidden_size) | |
| self.t_embedder = TimestepEmbedder(hidden_size) | |
| self.context_embedder = nn.Linear(context_dim, hidden_size) | |
| self.register_buffer("pos_embed", torch.zeros(1, self.x_embedder.num_patches, hidden_size)) | |
| self.joint_blocks = nn.ModuleList( | |
| [JointBlock(hidden_size, num_heads, pre_only=(i == depth - 1)) for i in range(depth)] | |
| ) | |
| self.final_layer = FinalLayer(hidden_size, patch_size, in_channels) | |
| pe = get_2d_sincos_pos_embed(hidden_size, int(self.x_embedder.num_patches ** 0.5)) | |
| self.pos_embed.data.copy_(torch.from_numpy(pe).float().unsqueeze(0)) | |
| def unpatchify(self, x, h, w): | |
| c, p = self.out_channels, self.x_embedder.patch_size[0] | |
| x = x.reshape(x.shape[0], h, w, p, p, c) | |
| x = torch.einsum("nhwpqc->nchpwq", x) | |
| return x.reshape(x.shape[0], c, h * p, w * p) | |
| def forward(self, x, t, context): | |
| H, W = x.shape[-2:] | |
| p = self.x_embedder.patch_size[0] | |
| x = self.x_embedder(x) + self.pos_embed | |
| c = self.t_embedder(t, dtype=x.dtype) | |
| context = self.context_embedder(context) | |
| for block in self.joint_blocks: | |
| context, x = block(context, x, c=c) | |
| x = self.final_layer(x, c) | |
| return self.unpatchify(x, H // p, W // p) | |