File size: 5,254 Bytes
1540c35 | 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 | # Klein-style single-stream DiT for 32x32 RGBA sprites, shared by train.py and prof_train.py; 2026-10-17.
import math
import torch, torch.nn as nn, torch.nn.functional as F
HEAD = 128
NCAT, NFP = 8, 3
# prof_train.py flips these to measure them; 'off' settings exist only for timing.
FAST = {'norm': True, 'rope': 'stack', 'qknorm': True, 'joint_qk': False, 'addcmul': False}
class RMSNorm(nn.Module):
def __init__(s, d):
super().__init__(); s.w = nn.Parameter(torch.ones(d))
def forward(s, x):
if FAST['norm']:
return F.rms_norm(x, (x.shape[-1],), s.w, 1e-6)
return x * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + 1e-6).to(x.dtype) * s.w
def rope_table(side):
# Half of each head encodes the row, half the column (2D RoPE, as in klein).
f = 1.0 / (10000 ** (torch.arange(0, HEAD // 4) / (HEAD // 4)))
ys, xs = torch.meshgrid(torch.arange(side), torch.arange(side), indexing='ij')
ang = torch.cat([ys.reshape(-1, 1) * f, xs.reshape(-1, 1) * f], 1)
return torch.cos(ang), torch.sin(ang)
def rope(x, cos, sin):
if FAST['rope'] == 'off':
return x
if FAST['rope'] == 'half':
# Rotate-half pairing (dim i with i+64): fewer passes over q/k than the interleaved form.
c, s_ = torch.cat([cos, cos], -1).to(x.dtype), torch.cat([sin, sin], -1).to(x.dtype)
x1, x2 = x.chunk(2, -1)
return x * c + torch.cat([-x2, x1], -1) * s_
if FAST['rope'] == 'stack':
a, b = x.unflatten(-1, (-1, 2)).unbind(-1)
cos, sin = cos.to(x.dtype), sin.to(x.dtype)
return torch.stack([a * cos - b * sin, a * sin + b * cos], -1).flatten(-2)
a, b = x.float().unflatten(-1, (-1, 2)).unbind(-1)
return torch.stack([a * cos - b * sin, a * sin + b * cos], -1).flatten(-2).to(x.dtype)
class Block(nn.Module):
def __init__(s, h):
super().__init__()
s.h, s.nh, s.mlp = h, h // HEAD, 3 * h
s.norm = nn.LayerNorm(h, elementwise_affine=False, eps=1e-6)
s.mod = nn.Linear(h, 3 * h)
s.lin1 = nn.Linear(h, 3 * h + 2 * s.mlp)
s.qn, s.kn = RMSNorm(HEAD), RMSNorm(HEAD)
s.lin2 = nn.Linear(h + s.mlp, h)
nn.init.zeros_(s.mod.weight); nn.init.zeros_(s.mod.bias)
def forward(s, x, c, cos, sin):
B, T, _ = x.shape
sh, sc, g = s.mod(F.silu(c)).unsqueeze(1).chunk(3, -1)
xm = torch.addcmul(sh, s.norm(x), 1 + sc) if FAST['addcmul'] else s.norm(x) * (1 + sc) + sh
qkv, a, b = s.lin1(xm).split([3 * s.h, s.mlp, s.mlp], -1)
qkv = qkv.view(B, T, 3, s.nh, HEAD).permute(2, 0, 3, 1, 4)
v = qkv[2]
if FAST['joint_qk']:
qk = qkv[:2]
if FAST['qknorm']:
qk = F.rms_norm(qk, (HEAD,), eps=1e-6) * torch.stack([s.qn.w, s.kn.w]).view(2, 1, 1, 1, HEAD).to(qk.dtype)
q, k = rope(qk, cos, sin).unbind(0)
else:
q, k = qkv[0], qkv[1]
if FAST['qknorm']:
q, k = s.qn(q), s.kn(k)
q, k = rope(q, cos, sin), rope(k, cos, sin)
o = F.scaled_dot_product_attention(q, k, v).transpose(1, 2).reshape(B, T, s.h)
out = s.lin2(torch.cat([o, F.silu(a) * b], -1))
return torch.addcmul(x, g, out) if FAST['addcmul'] else x + g * out
class DiT(nn.Module):
def __init__(s, h, depth, patch, vocab=0):
super().__init__()
s.p, s.side = patch, 32 // patch
s.inp = nn.Linear(4 * patch * patch, h)
s.t1, s.t2 = nn.Linear(256, h), nn.Linear(h, h)
s.cat = nn.Embedding(NCAT + 1, h) # last index = no label
s.fp = nn.Embedding(NFP + 1, h)
if vocab: # caption words: learned embeddings, mean-pooled, 0 = padding/none
s.wemb = nn.Embedding(vocab, h, padding_idx=0)
s.wproj = nn.Sequential(nn.Linear(h, h), nn.SiLU(), nn.Linear(h, h))
s.blocks = nn.ModuleList([Block(h) for _ in range(depth)])
s.fnorm = nn.LayerNorm(h, elementwise_affine=False, eps=1e-6)
s.fmod = nn.Linear(h, 2 * h)
s.outp = nn.Linear(h, 4 * patch * patch)
for m in (s.fmod, s.outp):
nn.init.zeros_(m.weight); nn.init.zeros_(m.bias)
cos, sin = rope_table(s.side)
s.register_buffer('cos', cos, persistent=False); s.register_buffer('sin', sin, persistent=False)
def forward(s, x, t, cat, fp, txt=None):
B, p, n = x.shape[0], s.p, s.side
x = x.unfold(2, p, p).unfold(3, p, p).permute(0, 2, 3, 1, 4, 5).reshape(B, n * n, 4 * p * p)
half = torch.exp(-math.log(10000) * torch.arange(128, device=x.device) / 128)
te = torch.cat([torch.cos(t[:, None] * 1000 * half), torch.sin(t[:, None] * 1000 * half)], -1)
c = s.t2(F.silu(s.t1(te))) + s.cat(cat) + s.fp(fp)
if txt is not None:
m = (txt > 0).to(c.dtype)
e = (s.wemb(txt) * m[..., None]).sum(1) / m.sum(1, keepdim=True).clamp(min=1)
c = c + s.wproj(e) * (m.sum(1, keepdim=True) > 0)
x = s.inp(x)
for blk in s.blocks:
x = blk(x, c, s.cos, s.sin)
sh, sc = s.fmod(F.silu(c)).unsqueeze(1).chunk(2, -1)
x = s.outp(s.fnorm(x) * (1 + sc) + sh)
return x.view(B, n, n, 4, p, p).permute(0, 3, 1, 4, 2, 5).reshape(B, 4, 32, 32)
|