| """ |
| T.I.T.S.2 — a Diffusion Transformer (DiT), i.e. this one actually is a transformer. |
| |
| Works on 4x32x32 SD-VAE latents (256x256 images). Patchify 2x2 -> 256 tokens, then a |
| stack of transformer blocks. Each block does: |
| |
| self-attention : every patch talks to every other patch |
| cross-attention : every patch reads the prompt, token by token |
| MLP |
| |
| Timestep + pooled text conditioning goes in via adaLN-Zero (the DiT paper's trick: |
| each block starts as an identity function and learns its way out, which makes deep |
| stacks train stably from scratch). |
| |
| ~93M params at defaults (dim=576, depth=12, heads=9). |
| """ |
| import math |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from torch.utils.checkpoint import checkpoint |
|
|
|
|
| def timestep_embedding(t: torch.Tensor, dim: int) -> torch.Tensor: |
| """Sinusoidal embedding. t is continuous in [0,1] here (flow matching), so scale up.""" |
| half = dim // 2 |
| freqs = torch.exp(-math.log(10000) * torch.arange(half, device=t.device, dtype=torch.float32) / half) |
| args = t.float()[:, None] * 1000.0 * freqs[None] |
| return torch.cat([torch.cos(args), torch.sin(args)], dim=-1) |
|
|
|
|
| def modulate(x, shift, scale): |
| return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) |
|
|
|
|
| class Attention(nn.Module): |
| """Self-attention, or cross-attention when `context_dim` is given.""" |
|
|
| def __init__(self, dim: int, heads: int, context_dim: int = None): |
| super().__init__() |
| self.heads = heads |
| self.head_dim = dim // heads |
| self.q = nn.Linear(dim, dim, bias=False) |
| self.kv = nn.Linear(context_dim or dim, dim * 2, bias=False) |
| self.out = nn.Linear(dim, dim) |
| self.cross = context_dim is not None |
|
|
| def forward(self, x, context=None, context_mask=None): |
| b, n, _ = x.shape |
| ctx = context if self.cross else x |
| q = self.q(x).view(b, n, self.heads, self.head_dim).transpose(1, 2) |
| k, v = self.kv(ctx).view(b, ctx.shape[1], 2, self.heads, self.head_dim).permute(2, 0, 3, 1, 4) |
| mask = None |
| if self.cross and context_mask is not None: |
| mask = context_mask[:, None, None, :].bool() |
| x = F.scaled_dot_product_attention(q, k, v, attn_mask=mask) |
| return self.out(x.transpose(1, 2).reshape(b, n, -1)) |
|
|
|
|
| class DiTBlock(nn.Module): |
| def __init__(self, dim: int, heads: int, text_dim: int, mlp_ratio: float = 4.0): |
| super().__init__() |
| self.norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) |
| self.attn = Attention(dim, heads) |
| self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) |
| self.cross_attn = Attention(dim, heads, context_dim=text_dim) |
| self.norm3 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6) |
| hidden = int(dim * mlp_ratio) |
| self.mlp = nn.Sequential(nn.Linear(dim, hidden), nn.GELU(approximate="tanh"), nn.Linear(hidden, dim)) |
| |
| self.ada = nn.Sequential(nn.SiLU(), nn.Linear(dim, 6 * dim)) |
| nn.init.zeros_(self.ada[1].weight) |
| nn.init.zeros_(self.ada[1].bias) |
|
|
| def forward(self, x, cond, context, context_mask): |
| shift1, scale1, gate1, shift2, scale2, gate2 = self.ada(cond).chunk(6, dim=-1) |
| x = x + gate1.unsqueeze(1) * self.attn(modulate(self.norm1(x), shift1, scale1)) |
| x = x + self.cross_attn(self.norm2(x), context, context_mask) |
| x = x + gate2.unsqueeze(1) * self.mlp(modulate(self.norm3(x), shift2, scale2)) |
| return x |
|
|
|
|
| class TITS2(nn.Module): |
| def __init__( |
| self, |
| latent_size: int = 32, |
| in_ch: int = 4, |
| patch_size: int = 2, |
| dim: int = 576, |
| depth: int = 12, |
| heads: int = 9, |
| text_dim: int = 768, |
| use_checkpoint: bool = False, |
| ): |
| super().__init__() |
| self.patch_size = patch_size |
| self.in_ch = in_ch |
| self.latent_size = latent_size |
| self.use_checkpoint = use_checkpoint |
| n_patches = (latent_size // patch_size) ** 2 |
|
|
| self.x_embed = nn.Conv2d(in_ch, dim, patch_size, stride=patch_size) |
| self.pos = nn.Parameter(torch.zeros(1, n_patches, dim)) |
| nn.init.normal_(self.pos, std=0.02) |
|
|
| self.t_mlp = nn.Sequential(nn.Linear(dim, dim), nn.SiLU(), nn.Linear(dim, dim)) |
| self.pooled_text = nn.Sequential(nn.Linear(text_dim, dim), nn.SiLU(), nn.Linear(dim, dim)) |
| self.text_proj = nn.Linear(text_dim, text_dim) |
|
|
| self.blocks = nn.ModuleList([DiTBlock(dim, heads, text_dim) 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.proj_out = nn.Linear(dim, patch_size * patch_size * in_ch) |
| nn.init.zeros_(self.ada_out[1].weight) |
| nn.init.zeros_(self.ada_out[1].bias) |
| nn.init.zeros_(self.proj_out.weight) |
| nn.init.zeros_(self.proj_out.bias) |
|
|
| def unpatchify(self, x): |
| p, c = self.patch_size, self.in_ch |
| h = w = self.latent_size // p |
| x = x.reshape(x.shape[0], h, w, p, p, c).permute(0, 5, 1, 3, 2, 4) |
| return x.reshape(x.shape[0], c, h * p, w * p) |
|
|
| def forward(self, x, t, text_tokens, text_pooled, text_mask=None): |
| """x: (b,4,32,32) noisy latent. t: (b,) in [0,1]. text_tokens: (b,77,768).""" |
| h = self.x_embed(x).flatten(2).transpose(1, 2) + self.pos |
| cond = self.t_mlp(timestep_embedding(t, h.shape[-1])) + self.pooled_text(text_pooled) |
| context = self.text_proj(text_tokens) |
|
|
| for blk in self.blocks: |
| if self.use_checkpoint and self.training: |
| h = checkpoint(blk, h, cond, context, text_mask, use_reentrant=False) |
| else: |
| h = blk(h, cond, context, text_mask) |
|
|
| shift, scale = self.ada_out(cond).chunk(2, dim=-1) |
| h = modulate(self.norm_out(h), shift, scale) |
| return self.unpatchify(self.proj_out(h)) |
|
|
|
|
| if __name__ == "__main__": |
| net = TITS2() |
| print(f"TITS2 params: {sum(p.numel() for p in net.parameters()) / 1e6:.2f}M") |
| x = torch.randn(2, 4, 32, 32) |
| out = net(x, torch.rand(2), torch.randn(2, 77, 768), torch.randn(2, 768)) |
| print("output:", out.shape) |
|
|
