Release category-conditioned sprite EMA weights and browser ONNX exports

README.md

@@ -0,0 +1,46 @@
+---
+language:
+- en
+tags:
+- image-generation
+- pixel-art
+- rectified-flow
+- onnx
+- webgpu
+pipeline_tag: unconditional-image-generation
+library_name: onnxruntime
+---
+# Sprite diffusion
+
+A 24.2M-parameter, category-conditioned rectified-flow transformer for **32×32 RGBA pixel-art sprites**. This repository distributes inference-only EMA weights for the original 100k model and its 150k continuation, plus ONNX exports for a browser generator. It is not a text-to-image model: it accepts eight category labels, not natural-language prompts.
+
+[Code and training write-up](https://github.com/Cyronius/model-experiments/tree/main/sprite-diffusion) · [Browser generator](https://cyronius.github.io/model-experiments/sprite-diffusion/)
+
+## Files and use
+
+- `sprite-100k.safetensors`, `sprite-150k.safetensors`: float32 EMA state dictionaries, no optimizer, training arguments, sampler state or datasets.
+- `sprite-100k.onnx`, `sprite-150k.onnx`: float32 velocity networks, opset 17, fixed batch two for conditional/unconditional classifier-free guidance.
+- Matching JSON files record architecture, artifact hashes and local PyTorch/ONNX parity checks.
+- `dit.py` defines the PyTorch architecture. Use `DiT(384, 10, 2)`, load a safetensors state dictionary, then run Euler integration from Gaussian noise at t=1 to t=0.
+
+ONNX inputs: `x` float32 `[2,4,32,32]`, `t` float32 `[2]`, `category` int64 `[2]`, `footprint` int64 `[2]`. The output is `velocity` float32 `[2,4,32,32]`. For a category c, duplicate the latent into both batch rows, use category `[c,8]`, footprint `[3,3]`, and combine velocities as `unconditional + 3 * (conditional - unconditional)`. Take a negative Euler step of `1 / steps`. Scale the final clamped [-1,1] RGBA to [0,255], with alpha threshold 128.
+
+Categories, in order: creature, item, weapon, food, effect, terrain, structure, nature. The food category includes drinks and potions.
+
+**Sampling:** 32 steps is the owner's preferred default; 16 is balanced and 8 is fast. 1–2 steps produce underdeveloped outputs. 64 is not reliably better. Keep both checkpoints: a qualitative same-seed comparison did not establish a universal winner. The 150k model showed more nature variety, while the original retained more varied characters in the inspected seeds.
+
+## Training and evaluation
+
+Hidden width 384, 10 blocks, three 128-dimensional attention heads, patch size 2 and 256 tokens. The original checkpoint ran for 100k updates. The continuation used an expanded 383,660-image/frame dataset, 3e-5 learning rate, category/subtype-balanced sampling and palette augmentation: a 10k-update pilot reset the optimizer, then another 40k resumed that pilot. Creature subtypes affect exposure only; they are not output labels.
+
+The comparison used 100k/110k/150k checkpoints at 1,2,4,8,12,16,32,64 steps, 64 sprites per setting, two fixed seed sets, CFG 3, EMA weights and no palette snapping. One conversational vision review inspected 1,536 generated sprites. It was not an independent blinded evaluation. Unseen validation and a matched uniform control remain outstanding; no new formal copying or diversity evaluation is claimed.
+
+## Data provenance, licensing and limitations
+
+**No training data is distributed.** The corpus combined game-sprite extractions and art packs from mixed sources. The expanded archive contains 20,542 TSR entries, 352,465 OGA entries, 9,597 pack entries and 1,056 curated CC0 entries. These are stored image/frame counts, not distinct designs or training exposure. The supplement's CC0 records do not establish licenses for the rest of the corpus. Source-permission and attribution records are incomplete, and some historical pack licenses were not verified.
+
+**No verified redistribution license for these weights is asserted by this model card.** The source repository's Apache-2.0 license covers its code; do not assume it grants rights in the training artwork, weights, or generated outputs. Public availability is not a representation of third-party permission or commercial-use clearance. Consider the underlying rights and attribution obligations before reuse.
+
+Outputs can be incomplete, noisy or resemble familiar game-art conventions. Creature output favors people; food often resembles potions. Terrain is not guaranteed to tile. The model does not produce coherent animation sequences. Select and edit results rather than treating every output as a production-ready asset. No output non-infringement guarantee is provided.
+
+The original 100k model was tested on AMD XDNA2 NPU hardware: approximately 5.0 s for a 64-sprite sheet at 8 steps and 9.9 s at 16 steps. These historical timings are not browser benchmarks or measurements of the 150k continuation. Browser inference support and speed vary by device.

dit.py

@@ -0,0 +1,113 @@
+# 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)

sprite-100k.json

@@ -0,0 +1,30 @@
+{
+  "name": "sprite-100k",
+  "architecture": {
+    "hidden": 384,
+    "depth": 10,
+    "patch": 2,
+    "channels": 4,
+    "canvas": 32
+  },
+  "weights": "EMA",
+  "cfg_batch": 2,
+  "dtype": "float32",
+  "opset": 17,
+  "velocity_relative_errors": [
+    5.293890126267797e-07,
+    2.8245901830814546e-06,
+    1.3860703802492935e-06
+  ],
+  "euler_16_relative_error": 7.527040725108236e-05,
+  "files": {
+    "sprite-100k.onnx": {
+      "bytes": 102321347,
+      "sha256": "3a4cc151d353ebdcee2a01e9b7cd999eb2eb760f545585c6b065ccd8ee138c06"
+    },
+    "sprite-100k.safetensors": {
+      "bytes": 96829272,
+      "sha256": "c2198937ababc6f50701206079ae4fdefa8a5f6fc358bc2333dadb3190198388"
+    }
+  }
+}

sprite-100k.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3a4cc151d353ebdcee2a01e9b7cd999eb2eb760f545585c6b065ccd8ee138c06
+size 102321347

sprite-100k.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c2198937ababc6f50701206079ae4fdefa8a5f6fc358bc2333dadb3190198388
+size 96829272

sprite-150k.json

@@ -0,0 +1,30 @@
+{
+  "name": "sprite-150k",
+  "architecture": {
+    "hidden": 384,
+    "depth": 10,
+    "patch": 2,
+    "channels": 4,
+    "canvas": 32
+  },
+  "weights": "EMA",
+  "cfg_batch": 2,
+  "dtype": "float32",
+  "opset": 17,
+  "velocity_relative_errors": [
+    9.889696457321406e-07,
+    2.4362118438148173e-06,
+    1.3656183455168502e-06
+  ],
+  "euler_16_relative_error": 1.3164431948098354e-05,
+  "files": {
+    "sprite-150k.onnx": {
+      "bytes": 102321347,
+      "sha256": "f6d6f3949dddd34de027a62f76e082cc7f12d63a6545feedcc0f4ba13263d3d3"
+    },
+    "sprite-150k.safetensors": {
+      "bytes": 96829272,
+      "sha256": "bbe7cedcc7716b98d5dcef6db8088115b6698c31cebce6e80917c888211bc820"
+    }
+  }
+}

sprite-150k.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f6d6f3949dddd34de027a62f76e082cc7f12d63a6545feedcc0f4ba13263d3d3
+size 102321347

sprite-150k.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bbe7cedcc7716b98d5dcef6db8088115b6698c31cebce6e80917c888211bc820
+size 96829272