PixelGen2: compositional hand-weighted pixel-art generator + 4 scaling experiments

.gitignore

@@ -0,0 +1,5 @@
+.venv/
+__pycache__/
+*.pyc
+.pytest_cache/
+out/

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 DJ Lougen
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,256 @@
+---
+license: mit
+library_name: pytorch
+pipeline_tag: text-to-image
+tags:
+  - pixel-art
+  - weight-space-programming
+  - hand-written-weights
+  - no-training
+  - tiny-model
+  - pico-8
+  - toy
+language:
+  - en
+---
+
+<div align="center">
+
+# 🕹️ PixelGen2
+
+**A text-to-pixel-art model with 6,747 parameters. Zero of them were trained.**
+
+Every weight was typed in by hand. It draws 21,866 distinct sprites from 28 templates.
+
+<img src="showcase2.png" alt="12 bodies × 8 accessories, two colours each, generated by PixelGen2" width="100%"/>
+
+<sub>Every sprite above is one prompt: <code>"&lt;colour&gt; &lt;body&gt; with &lt;colour&gt; &lt;accessory&gt;"</code>. No gradients were harmed.</sub>
+
+<br/>
+
+![params](https://img.shields.io/badge/parameters-6%2C747-ff004d?style=flat-square)
+![outputs](https://img.shields.io/badge/distinct%20sprites-21%2C866-29adff?style=flat-square)
+![training](https://img.shields.io/badge/training%20steps-0-00e436?style=flat-square)
+![palette](https://img.shields.io/badge/palette-PICO--8-ffec27?style=flat-square)
+![license](https://img.shields.io/badge/license-MIT-83769c?style=flat-square)
+
+</div>
+
+---
+
+## What this is
+
+A tiny PyTorch `nn.Module` that turns a prompt like `"red potion"` or `"purple ghost with gold crown"` into a 16×16 pixel-art sprite in the PICO-8 palette, upscaled to 128×128.
+
+It's an experiment in **weight-space programming**, inspired by [this post](https://rikiyat.github.io/blog/weight-space-programming/): instead of training a network, you write a program and *compile it into weights*. Linear layers become lookup tables, argmax becomes a decision, a 1×1 conv becomes a palette. The model is a real `nn.Module` — you can call `.forward()`, save a `state_dict`, load it back — but no optimizer ever touched it.
+
+There are two models in the repo:
+
+| | `PixelGen` (v1) | `PixelGen2` (v2) |
+|---|---|---|
+| prompt | `"<colour> <body>"` | `"[small] <colour> <body> [with <colour> <accessory>]"` |
+| parameters | 4,227 | 6,747 |
+| hand-drawn templates | 21 bodies | 21 bodies + 7 accessories |
+| distinct sprites (seed 0) | 189 | **21,866** |
+| shading | fixed row bands | structural rim light (two 3×3 convs) |
+
+v1 is the minimal version, published separately as [DJLougen/pixelgen](https://huggingface.co/DJLougen/pixelgen); v2 is the same thing plus four compositional layers, and is what the CLI uses by default. This repo contains both. Everything below was measured, not estimated — scripts in `experiments/`.
+
+## How it works
+
+Nine steps for v1. Each is an ordinary tensor op; together they're a program.
+
+```
+"red potion"
+    │
+    ▼
+ 1. bag of words      prompt → count vector over a 138-word vocab
+ 2. class WTA         Linear(138→21) → argmax → one-hot body           [potion]
+ 3. shape bank        one-hot selects a 16×8 half-template;
+                      '?' cells flip on when seed-noise ≥ 0.5
+ 4. mirror            half + horizontal flip → 16×16 body            (always symmetric)
+ 5. outline           3×3 dilate(body) − body → outline pixels       (palette index 1)
+ 6. shading           rows 0–6 light, 7–10 mid, 11–15 dark
+ 7. ramp WTA          Linear(138→9) → argmax → colour ramp            [red]
+ 8. one-hot palette   every pixel → one-hot over 16 PICO-8 colours
+ 9. palette conv      Conv2d(16→3, 1×1) → RGB
+    │
+    ▼
+ 16×16 sprite → nearest-neighbour ×8 → 128×128 PNG
+```
+
+v2 inserts four more fixed-weight layers (`pixelgen/compose.py`):
+
+| layer | mechanism | params |
+|---|---|---|
+| **accessory** | `Linear(138→8)` WTA selects a 16×8 part (crown, halo, horns, helmet, antenna, wings, torch) overlaid on the body with `max`. `none` wins ties. | 1,112 |
+| **second colour** | the prompt is split at `with`/`wearing`/`holding`/`and`; a second `Linear(138→9)` colours the accessory from the words after it. No colour word there → body colour wins via a 0.5 one-hot bonus. | 1,251 |
+| **structural shading** | two 3×3 convs with a single tap each: a pixel is *light* if its up-left neighbour is empty, *dark* if its down-right neighbour is. Light from the top-left, for every shape, for 18 parameters. | 18 |
+| **size** | `small`/`tiny`/`mini` → nearest-neighbour 0.75× resample of the body map, bottom-centred. | 139 |
+
+The "generative" part is step 3: templates have `?` cells gated by per-seed noise, so the same prompt gives small variations across seeds:
+
+<div align="center">
+<img src="seeds.png" alt="blue invader, seeds 0–7" width="80%"/>
+<sub><code>"blue invader"</code>, seeds 0 through 7</sub>
+</div>
+
+## Numbers
+
+| | v1 | v2 |
+|---|---|---|
+| Registered parameters | 4,227 | **6,747** |
+| Non-trainable buffers (template / shade / ramp / part banks) | 6,171 | 7,333 |
+| Total values | 10,398 (≈41.6 KB fp32) | 14,080 (≈56.3 KB fp32) |
+| Training steps | **0** | **0** |
+| Optimizers, backward passes, `requires_grad=True` in the repo | 0 (grep-verified) | 0 |
+| Bodies / accessories / colour ramps | 21 / – / 9 | 21 / 7 / 9 |
+| Distinct images at seed 0 | 189 of 189 prompts | 21,866 of 24,192 prompts |
+| Output | 16×16 → 128×128 RGB, every pixel ∈ PICO-8 | same |
+| `weights/pixelgen.pt` (v1) | 45.7 KB, SHA-256 `4c5dffd1…540eb`, identical across rebuilds | – |
+
+<details>
+<summary>Parameter breakdown (v2; v1 is the first six rows)</summary>
+
+| Tensor | Shape | Count | What it is |
+|---|---|---|---|
+| `word_to_class.weight` | 21 × 138 | 2,898 | vocab → body scores (synonym table) |
+| `word_to_class.bias` | 21 | 21 | default body (potion) |
+| `word_to_ramp.weight` | 9 × 138 | 1,242 | vocab → colour-ramp scores |
+| `word_to_ramp.bias` | 9 | 9 | default ramp (green) |
+| `outline.weight` | 1×1×3×3 | 9 | dilation kernel |
+| `palette_conv.weight` | 3×16×1×1 | 48 | one-hot colour → RGB (exact PICO-8 ÷ 255) |
+| `word_to_acc.weight` / `.bias` | 8 × 138, 8 | 1,112 | vocab → accessory scores, `none` biased |
+| `word_to_ramp2.weight` / `.bias` | 9 × 138, 9 | 1,251 | accessory-half vocab → ramp |
+| `lit.weight`, `shadow.weight` | 1×1×3×3 ×2 | 18 | up-left tap, down-right tap |
+| `word_to_size.weight` / `.bias` | 1 × 138, 1 | 139 | size words → small flag |
+| **Parameters** | | **6,747** | |
+| `on_bank` / `free_bank` | 21×16×8 ×2 | 5,376 | body templates (`#` solid, `?` noise-controlled) |
+| `acc_bank` | 8×16×8 | 1,024 | accessory templates |
+| `shade_bank` | 3×16×16 | 768 | row-band base shading |
+| `ramp_table` | 9×3 | 27 | palette indices per ramp |
+| `acc_mask_a` | 138 | 138 | body-class words never mean a hat in the body half |
+| **Buffers** | | **7,333** | |
+
+</details>
+
+## How to use
+
+```bash
+git clone https://huggingface.co/DJLougen/pixelgen2 && cd pixelgen2
+pip install torch numpy pillow            # that's the whole dependency list
+python generate.py "red potion" --out out/red_potion.png
+python generate.py "green slime with gold crown" --seed 3 --scale 16
+python generate.py "small angel skull" --out out/skull.png
+python generate.py "blue invader" --v1  # original single-part model
+```
+
+Or from Python:
+
+```python
+from pixelgen.compose import PixelGen2
+
+model = PixelGen2()                     # weights are built in __init__, no file needed
+img = model.sprite("purple ghost with horns", seed=4, scale=8)   # PIL.Image, 128×128
+img.save("ghost.png")
+```
+
+**Prompt grammar.** `[small] [colour] <body> [with [colour] <accessory>]`. Words are matched by a bag-of-words, so order doesn't matter within each half.
+
+<details>
+<summary>Vocabulary</summary>
+
+**Bodies:** potion (bottle, flask, elixir, vial) · sword (blade, dagger) · shield (buckler) · ship (rocket, spaceship, starship) · invader (alien, monster) · slime (blob, jelly) · skull (bone, skeleton) · tree (pine, forest) · key · heart (love, life) · star (sparkle) · mushroom (fungus, toadstool, shroom) · ghost (spirit, phantom, spooky) · bomb (grenade, explosive) · coin (token, money, medal) · gem (diamond, jewel, crystal) · cherry (cherries, berry, fruit) · crown · bat (vampire) · cactus (desert, succulent) · robot (droid, android, mech)
+
+**Accessories:** crown (king, queen, royal, crowned, prince, princess) · halo (angel, angelic, holy, saint) · horns (horned, devil, demon, evil) · helmet (knight, armor, armour, armored) · antenna (antennae, bug, martian) · wings (winged, fly, flying, bird, fairy, butterfly) · torch (torches, fire, flame, burning)
+
+**Colours:** red (crimson, ruby) · blue (azure, sapphire) · green (emerald, lime) · yellow (gold, golden) · purple (violet, amethyst) · orange (amber) · pink (rose) · brown (wood, wooden) · white (grey, gray, silver, bone)
+
+**Size:** small, tiny, mini, little. **Separators:** with, wearing, holding, and, plus.
+
+`"crown"` in the body half is the crown object; after a separator it's a hat. `"bone"` is both a skull synonym and a white synonym, so `"bone"` alone is a white skull. `"alien"` is an invader body and `"martian"` is an antenna accessory, so `"martian alien"` is what you'd hope.
+
+</details>
+
+## How far can you push it?
+
+Four experiments, all in `experiments/` with the scripts that produced them. Full notes in [`experiments/README.md`](experiments/README.md).
+
+**1. Adding a class is a pure append.** Freeze the 11-class model, add `mushroom`, diff. +400 values. Diffed by *index*, 84 old values changed; diffed by *meaning* (class name × word), **zero** — the raw diff was entirely a column permutation from a `sorted()` vocab. 17/17 old prompts bit-identical over 3 seeds. The same check still passes after +10 classes and the whole accessory vocabulary.
+
+**2. Capacity is unbounded; addressing is quadratic.** Synthetic random classes, rebuilt and probed at each size:
+
+| classes | params | dense layer share | isolated |
+|---|---|---|---|
+| 21 | 2,937 | 22% | ✓ |
+| 121 | 51,537 | 57% | ✓ |
+| 1,021 | 3.2M | 92% | ✓ |
+| 10,021 | 302M | 99.1% | ✓ |
+
+It never breaks — the original 21 classes stay bit-identical at 10k. But vocab grows with classes, so `word_to_class` is ≈3C² and 99.997% zeros. Under 100k params the dense design tops out near **170 classes**.
+
+**3. Compressing the class layer costs semantics.** A binary code (`d = ⌈log₂C⌉`, words emit ±1 codes, score = `d − 2·hamming`) shrinks the class layer **528×** at 10k classes with no loss of isolation — but 61% of two-body-word prompts (`"sword shield"`) now resolve to a *third* class, vs 0% dense. The additive bag-of-words semantics were the thing that got compressed away.
+
+**4. Composition is the lever.** v1 → v2: 1.6× params, **115× distinct outputs**. 781 distinct images per hand-drawn template instead of 9.
+
+<div align="center">
+<img src="experiments/shading_before_after.png" alt="row-band shading vs structural rim lighting" width="100%"/>
+<sub>Top: v1 row bands. Bottom: v2 structural lighting, 18 params.</sub>
+</div>
+
+Quality, honestly: structural lighting helps blobby shapes (potion, heart, mushroom, tree, gem) and **hurts thin ones** (invader, bat, robot) — when every pixel is an edge, rim light becomes speckle. A fixed 3×3 kernel has no way to know a shape is too thin to light. That's the limit of hand-written rules without a learned prior, and exactly the kind of thing a few hundred trained parameters would fix.
+
+## Limitations
+
+- **It only knows what was drawn.** 21 bodies, 7 accessories. Prompt for a dragon and you get a potion (the default). Capacity is unbounded; *content* is a human with a text editor.
+- **Every shape is horizontally symmetric.** The mirror step makes 16×8 templates cheap; it also means `"torch"` is two torches and the key is a top-down key. Lighting (v2) is deliberately *not* symmetric.
+- **Accessories sit at fixed rows.** A crown floats above short bodies; a helmet merges into tall ones. Fixing that needs per-body anchor points — more drawing, not more mechanism.
+- **One accessory per sprite.** `"devil knight"` picks one by index order.
+- **Variation is cosmetic.** Seeds only toggle the `?` cells — a few pixels per sprite.
+- **Bag of words.** `"sword shield"` is a tie broken by index; `"not red"` is red; within each half, order is ignored.
+- **Structural shading speckles thin shapes.** See above.
+- **It's PyTorch for the bit, not for the need.** You could rewrite this in NumPy. Using `nn.Linear` and `nn.Conv2d` is the point of the exercise.
+- **Art quality is programmer art**, drawn in a text editor by someone who is not a pixel artist.
+
+## Why bother?
+
+Because it's a nice demonstration that a lot of what we call "a model" is just *structure*, and structure can be written directly. A linear layer with the right weights *is* a dictionary. A conv with the right kernel *is* an edge detector. Winner-take-all *is* an `if`. The blog post argues this is a useful lens for understanding what trained networks are doing; this repo is the version where you skip the training and just write the weights — and then push on it to see where hand-writing stops being enough.
+
+Where it stops: addressing (quadratic), compression (breaks additivity), and quality rules at the edges (shading thin shapes). The natural next step is the hybrid — hand-write the ~95% that's structure, train only the small word→part map.
+
+Also it's cool that it worked.
+
+## Repo contents
+
+```
+pixelgen/
+  model.py       PixelGen (v1) — all weights set in __init__
+  compose.py     PixelGen2 (v2) — v1 + accessory, second colour, structural shading, size
+  templates.py   21 body half-templates ('#' solid, '?' noise, '.' empty)
+  parts.py       7 accessory half-templates
+  palette.py     PICO-8 palette + 9 colour ramps
+  tokenizer.py   vocab, synonyms, prompt splitting
+  weights.py     save/load helpers
+build.py         writes weights/pixelgen.pt (v1, deterministic)
+generate.py      CLI (v2 default, --v1 for the original)
+tests/           16 property tests
+experiments/     the four experiments above: scripts, frozen baselines, README
+showcase2.png    the grid at the top
+seeds.png        seed variation strip
+```
+
+Tests: `python -m pytest tests -q` → 16 passed.
+
+## Citation
+
+```bibtex
+@misc{pixelgen2_2025,
+  title  = {PixelGen2: a compositional hand-weighted pixel-art generator},
+  author = {Lougen, D. J.},
+  year   = {2025},
+  url    = {https://huggingface.co/DJLougen/pixelgen2},
+  note   = {Weight-space programming experiment; no parameters were trained.}
+}
+```
+
+Inspired by *[Weight-Space Programming](https://rikiyat.github.io/blog/weight-space-programming/)*.

build.py

@@ -0,0 +1,17 @@
+"""Build PixelGen: write the hand-written weights to weights/pixelgen.pt."""
+
+import sys
+
+from pixelgen import weights
+
+
+def main(argv=None):
+    path = weights.save()
+    print(f"wrote {path}")
+    model = weights.load(path)
+    print(f"loaded back: {sum(p.numel() for p in model.buffers())} buffer values")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

experiments/README.md

@@ -0,0 +1,94 @@
+# Experiments: how far does hand-written capacity go?
+
+Every number below was produced by a script in this folder. No training anywhere.
+
+| script | question | answer |
+|---|---|---|
+| `add_class.py` | what changes when you add one class? | +400 values, 0 semantic edits to old weights, 0 behaviour change on old prompts |
+| `scale_classes.py` | how many classes before it breaks? | it doesn't; 10,021 classes still bit-isolated — but the word→class matrix is O(C²) |
+| `coded_classes.py` | can the class layer be compressed? | 528× at 10k classes via binary codes, at the cost of 61% multi-word interference |
+| `count_outputs.py` | what does composition buy? | 21,866 distinct images from 28 templates; 115× output for 1.6× params |
+
+## 1. Adding one class is a pure append (in function space)
+
+`before.pt` is the frozen 11-class model. After adding `mushroom`:
+
+```
+tensor shapes      on_bank (11,16,8)->(12,16,8)   word_to_class (11,59)->(12,63)   +400 values
+raw index diff     84 old values "changed"
+semantic diff      0 old (class,word) entries changed; shade/ramp/outline/palette identical
+old prompts        17/17 bit-identical over 3 seeds; only mushroom-word prompts moved (potion -> mushroom)
+```
+
+The 84-value raw diff is entirely a **column permutation**: the vocab is `sorted()`, so inserting `fungus` at column 11 shifts 48 of 59 old words right. Diffed by meaning, nothing moved. Lesson: if you want weight diffs to be legible, make index assignment append-only.
+
+The same check still passes after +10 hand-drawn classes and the accessory vocabulary: v1 renders the original 11 classes bit-identically.
+
+## 2. Capacity is unbounded; addressing is quadratic
+
+Synthetic random classes (3 nonsense words each), rebuilt and probed at each size:
+
+```
+ classes  vocab      params   buffers      MB  dense%  gen ms  isolated
+      21     95       2,937     6,171    0.04   21.9%    0.12      True
+     121    395      51,537    31,771    0.33   57.4%    0.13      True
+   1,021  3,095   3,188,937   262,171   13.80   91.6%    0.25      True
+   3,021  9,095  27,560,937   774,171  113.34   97.0%    1.45      True
+  10,021 30,095 301,862,937 2,566,171 1217.72   99.1%   11.95      True
+```
+
+`isolated` = the 21 real classes and probe prompts (`red potion`, `dragon`, `""`, `bone`) are bit-identical to the 21-class model, and the newest synthetic class resolves by all its synonyms. It never fails.
+
+But vocab grows with classes, so `word_to_class` is V×C ≈ 3C² — at 10k classes it's 99.1% of the model and 99.997% zeros. Under a 100k-parameter budget the dense design tops out around **170 classes**.
+
+## 3. Compressing the class layer buys size, costs semantics
+
+Replace `Linear(V,C)` with a binary code: each word emits its class's d=⌈log₂C⌉-bit code as ±1, a `Linear(d,C)` of code rows scores `d − 2·hamming`.
+
+```
+ classes   d   dense params   coded params   ratio  isolated  interference
+      21   5          2,016            601    3.4x      True     61%
+     121   7         47,916          3,733   12.8x      True     61%
+   1,021  10      3,161,016         42,181   74.9x      True     61%
+   3,021  12     27,479,016        148,413  185.2x      True     61%
+  10,021  14    301,592,016        571,645  527.6x      True     61%
+```
+
+Isolation survives. But `interference` — the share of `"wordA wordB"` prompts (over the 21 real classes) that resolve to a **third** class — jumps from 0% (dense; ties go to lower index) to 61%, flat across scale. The sum of two codes is nearest to some unrelated code. The bag-of-words layer's additive semantics were the thing being compressed away. Open question: do error-correcting or sparse k-hot codes recover 0% interference, and at what parameter cost?
+
+## 4. Composition: the lever that actually matters
+
+`PixelGen2` (`pixelgen/compose.py`) adds four fixed-weight layers to the same model:
+
+| layer | mechanism | params |
+|---|---|---|
+| accessory | `Linear(V,8)` WTA → 16×8 part overlaid with `max` | 1,112 + 1,162 buffer |
+| second colour | `Linear(V,9)` over words after `with`/`wearing`/…; body ramp wins ties via a 0.5 one-hot bonus | 1,251 |
+| structural shading | two 3×3 convs: light where up-left neighbour is empty, dark where down-right is | 18 |
+| size | `Linear(V,1)` → nearest 0.75× resample, bottom-centred | 139 |
+
+```
+PixelGen :  4,227 params + 6,171 buffers |    189 prompts ->    189 distinct images
+PixelGen2:  6,747 params + 7,333 buffers | 24,192 prompts -> 21,866 distinct images
+hand-drawn content: 21 bodies + 7 parts = 28 templates
+distinct images per template:  9  ->  781
+param growth 1.60x   output growth 115.7x
+```
+
+(The 2,326 duplicates are prompts whose accessory colour equals the body colour, or whose accessory is hidden behind the body.)
+
+### Quality, honestly
+
+![shading before/after](shading_before_after.png)
+
+Top: row-band shading (v1). Bottom: structural rim lighting (v2), 18 parameters. It helps blobby shapes — potion, heart, mushroom, tree, gem get a readable highlight and shadow — and **hurts thin ones**: invader, bat, robot turn to speckle because every pixel is an edge. A fixed 3×3 kernel has no notion of "this shape is too thin to light". That's the limit of rule-based shading without a learned prior, and it's exactly the kind of thing a few hundred trained parameters would fix.
+
+Accessories are placed at fixed rows, so a crown floats above short bodies and a helmet merges into tall ones. Fixing that properly needs per-body anchor points — more hand-written content, not more mechanism.
+
+## What this says about weight-space programming
+
+- **Adding capacity is free and modular.** Hand-written weights are block-diagonal by construction; trained weights are not. Add-a-class is `dict.append`.
+- **The expensive part is the index, not the content.** Dense lookup is O(C²); the content (templates) is O(C). Compression trades additivity for size.
+- **Composition beats enumeration by two orders of magnitude**, and it's the only route to "many things" that doesn't scale a human's drawing time linearly.
+- **Quality rules written by hand are brittle at the edges** — the shading result is the clearest example. The honest endpoint for pure weight-space programming is a very capable procedural generator with legible, auditable failure modes, and zero creativity outside what was drawn.
+- The natural next step is the hybrid the blog post gestures at: hand-write the structure (parts, compositing, palette, transforms — ~95% of params), and train only the small word→part mapping, so new vocabulary becomes learnable while everything downstream stays readable.

experiments/add_class.py

@@ -0,0 +1,133 @@
+"""Measure exactly how PixelGen changed after adding one class.
+
+Compares the frozen before-state (experiments/before.pt, before_outputs.npz,
+before_meta.json) against the live model built from current source.
+
+Reports:
+  1. tensor shapes before -> after
+  2. raw parameter diff (how many values differ by index)
+  3. semantic diff (how many values differ once rows/cols are aligned by
+     class name / word, i.e. the change that actually matters)
+  4. vocab index remap (sorted vocab means new words shift old columns)
+  5. output equivalence on every pre-existing prompt (bit-exact or not)
+  6. what the new prompts do now vs before
+"""
+
+import json
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from pixelgen.model import PixelGen
+
+HERE = Path(__file__).resolve().parent
+before_sd = torch.load(HERE / "before.pt", map_location="cpu", weights_only=True)
+before_out = np.load(HERE / "before_outputs.npz")
+meta = json.load(open(HERE / "before_meta.json"))
+b_vocab, b_classes, prompts = meta["vocab"], meta["classes"], meta["prompts"]
+
+m = PixelGen()
+after_sd = m.state_dict()
+a_vocab, a_classes = m.vocab, m.classes
+
+new_classes = [c for c in a_classes if c not in b_classes]
+new_words = [w for w in a_vocab if w not in b_vocab]
+print("=" * 70)
+print(f"new classes: {new_classes}   new words: {new_words}")
+print(f"vocab {len(b_vocab)} -> {len(a_vocab)}   classes {len(b_classes)} -> {len(a_classes)}")
+
+# 1. shapes -------------------------------------------------------------
+print("\n[1] tensor shapes")
+tot_b = tot_a = 0
+for k in after_sd:
+    sb, sa = tuple(before_sd[k].shape), tuple(after_sd[k].shape)
+    tot_b += before_sd[k].numel(); tot_a += after_sd[k].numel()
+    flag = "  <-- grew" if sb != sa else ""
+    print(f"  {k:24} {str(sb):>16} -> {str(sa):<16}{flag}")
+print(f"  total values {tot_b} -> {tot_a}  (+{tot_a - tot_b})")
+
+# 2. raw diff ------------------------------------------------------------
+print("\n[2] raw diff by index (naive view)")
+for k in after_sd:
+    b, a = before_sd[k], after_sd[k]
+    if b.shape == a.shape:
+        n = int((b != a).sum())
+        print(f"  {k:24} {n:5d} / {a.numel()} values differ")
+    else:
+        # compare overlapping prefix region
+        sl = tuple(slice(0, s) for s in b.shape)
+        n = int((b != a[sl]).sum())
+        print(f"  {k:24} {n:5d} / {b.numel()} values differ in the old-shaped prefix (+{a.numel()-b.numel()} new)")
+
+# 3. semantic diff (align by names) ---------------------------------------
+print("\n[3] semantic diff (rows/cols aligned by class name / word)")
+def aligned_linear(name):
+    b, a = before_sd[f"{name}.weight"], after_sd[f"{name}.weight"]
+    rows = b_classes if name == "word_to_class" else m.ramps
+    changed = []
+    for ri, r in enumerate(rows):
+        for w in b_vocab:
+            bv = b[ri, b_vocab[w]].item()
+            av = a[a_classes.index(r) if name == "word_to_class" else ri, a_vocab[w]].item()
+            if bv != av:
+                changed.append((r, w, bv, av))
+    return changed
+for name in ("word_to_class", "word_to_ramp"):
+    ch = aligned_linear(name)
+    print(f"  {name:16} old (row,word) entries changed: {len(ch)}  {ch[:5]}")
+bb, ab = before_sd["word_to_class.bias"], after_sd["word_to_class.bias"]
+print("  word_to_class.bias old-class entries changed:",
+      sum(bb[i].item() != ab[a_classes.index(c)].item() for i, c in enumerate(b_classes)))
+for bank in ("on_bank", "free_bank"):
+    b, a = before_sd[bank], after_sd[bank]
+    n = sum(int((b[i] != a[a_classes.index(c)]).sum()) for i, c in enumerate(b_classes))
+    print(f"  {bank:16} old-class template cells changed: {n}")
+for k in ("shade_bank", "ramp_table", "outline.weight", "palette_conv.weight"):
+    print(f"  {k:16} identical: {torch.equal(before_sd[k], after_sd[k])}")
+
+# 4. vocab remap ----------------------------------------------------------
+print("\n[4] vocab index remap caused by sorted insertion")
+moved = [(w, b_vocab[w], a_vocab[w]) for w in b_vocab if b_vocab[w] != a_vocab[w]]
+print(f"  {len(moved)} of {len(b_vocab)} old words changed column index")
+print("  first few:", moved[:6])
+print("  new word columns:", {w: a_vocab[w] for w in new_words})
+
+# 5. output equivalence ---------------------------------------------------
+print("\n[5] output equivalence on pre-existing prompts (3 seeds each)")
+same = diff = 0
+changed_prompts = []
+for i, p in enumerate(prompts):
+    b = before_out[f"p{i}"]
+    a = m.generate([p] * 3, [0, 1, 2], scale=1)
+    if np.array_equal(a, b):
+        same += 1
+    else:
+        diff += 1
+        changed_prompts.append(p)
+print(f"  identical: {same}   changed: {diff}   -> {changed_prompts}")
+
+# 6. new prompt behaviour ---------------------------------------------------
+print("\n[6] what the new-class prompts resolve to now")
+ref = {c: (m.generate([c], [0], scale=1)[0].sum(-1) > 0) for c in a_classes}
+for p in ["mushroom", "red mushroom", "purple fungus", "toadstool", "dragon", ""]:
+    mask = m.generate([p], [0], scale=1)[0].sum(-1) > 0
+    cls = [c for c, r in ref.items() if np.array_equal(r, mask)]
+    i = prompts.index(p) if p in prompts else None
+    was = "n/a"
+    if i is not None:
+        bm = before_out[f"p{i}"][0].sum(-1) > 0
+        was = [c for c, r in ref.items() if np.array_equal(r, bm)]
+    print(f"  {p!r:18} before -> {was}   after -> {cls}")
+
+# image ---------------------------------------------------------------------
+from PIL import Image
+S = 6; cell = 16 * S; pad = 4
+rows = [("before", b_classes), ("after", a_classes)]
+W = max(len(c) for _, c in rows) * (cell + pad) + pad
+img = Image.new("RGB", (W, 2 * (cell + pad) + pad), (29, 43, 83))
+for r, (_, classes) in enumerate(rows):
+    for c, cls in enumerate(classes):
+        img.paste(Image.fromarray(m.generate([cls], [0], scale=S)[0]), (pad + c * (cell + pad), pad + r * (cell + pad)))
+img.save(HERE / "before_after.png")
+print(f"\nwrote {HERE/'before_after.png'}  (top: old classes rendered by new model; bottom: all classes)")

experiments/before.pt

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

experiments/before_meta.json

@@ -0,0 +1 @@
+{"prompts": ["potion", "sword", "shield", "ship", "invader", "slime", "skull", "tree", "key", "heart", "star", "red potion", "blue sword", "bone", "dragon", "golden key", "mushroom", "red mushroom", "purple fungus", ""], "vocab": {"<pad>": 0, "alien": 1, "blade": 2, "blob": 3, "bone": 4, "bottle": 5, "buckler": 6, "dagger": 7, "elixir": 8, "flask": 9, "forest": 10, "heart": 11, "invader": 12, "jelly": 13, "key": 14, "life": 15, "love": 16, "monster": 17, "pine": 18, "potion": 19, "rocket": 20, "shield": 21, "ship": 22, "skeleton": 23, "skull": 24, "slime": 25, "spaceship": 26, "sparkle": 27, "star": 28, "starship": 29, "sword": 30, "tree": 31, "vial": 32, "amber": 33, "amethyst": 34, "azure": 35, "blue": 36, "brown": 37, "crimson": 38, "emerald": 39, "gold": 40, "golden": 41, "gray": 42, "green": 43, "grey": 44, "lime": 45, "orange": 46, "pink": 47, "purple": 48, "red": 49, "rose": 50, "ruby": 51, "sapphire": 52, "silver": 53, "violet": 54, "white": 55, "wood": 56, "wooden": 57, "yellow": 58}, "classes": ["potion", "sword", "shield", "ship", "invader", "slime", "skull", "tree", "key", "heart", "star"]}

experiments/before_outputs.npz

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

experiments/coded_classes.py

@@ -0,0 +1,67 @@
+"""Replace the dense V x C word->class matrix with a hand-written binary code.
+
+  word -> d-bit code of its class, emitted as +-1      Linear(V, d, bias=False)
+  class rows = +-1 codes; score_c = d - 2*hamming      Linear(d, C)  (+0.5 bias on default)
+
+Params drop from V*C to V*d + d*C. Price: multi-word prompts sum two codes
+and can land on a third class. Both are measured here.
+"""
+
+import itertools
+import sys
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+sys.path.insert(0, ".")
+from experiments.scale_classes import build, REAL_CLASSES, REAL_SYN  # noqa: E402
+import pixelgen.model as M  # noqa: E402
+
+
+def code_layer(m):
+    C, V = len(m.classes), len(m.vocab)
+    d = max(1, int(np.ceil(np.log2(C))))
+    codes = torch.tensor([[(ci >> b) & 1 for b in range(d)] for ci in range(C)], dtype=torch.float32) * 2 - 1
+    enc = nn.Linear(V, d, bias=False); dec = nn.Linear(d, C)
+    with torch.no_grad():
+        enc.weight.zero_(); dec.bias.zero_()
+        for ci, cls in enumerate(m.classes):
+            for w in M.CLASS_SYNONYMS[cls]:
+                enc.weight[:, m.vocab[w]] = codes[ci]
+        dec.weight.copy_(codes)
+        dec.bias[m.classes.index(M.DEFAULT_CLASS)] = 0.5
+    seq = nn.Sequential(enc, dec)
+    for p in seq.parameters(): p.requires_grad_(False)
+    return seq, d
+
+
+def interference(m):
+    """Fraction of 2-class-word prompts whose argmax is neither word's class."""
+    words = [(c, M.CLASS_SYNONYMS[c][0]) for c in REAL_CLASSES]
+    bad = 0; n = 0
+    for (ca, wa), (cb, wb) in itertools.combinations(words, 2):
+        toks = m.tokenize(f"{wa} {wb}")
+        bag = torch.zeros(1, len(m.vocab)); bag[0, toks] += 1  # noqa (tokenize returns ids)
+        cls = m.classes[m.word_to_class(bag).argmax(-1).item()]
+        bad += cls not in (ca, cb); n += 1
+    return bad / n
+
+
+Ns = [int(a) for a in sys.argv[1:]] or [0, 100, 1000, 3000, 10000]
+print(f"{'classes':>8} {'d':>3} {'dense params':>13} {'coded params':>13} {'ratio':>7} {'isolated':>8} {'interf':>7}")
+for n in Ns:
+    m, _ = build(n)
+    ref = {p: m.generate([p] * 2, [0, 1], scale=1) for p in REAL_CLASSES + ["red potion", "", "dragon", "bone"]}
+    dense_p = sum(p.numel() for p in m.word_to_class.parameters())
+    seq, d = code_layer(m)
+    m.word_to_class = seq
+    coded_p = sum(p.numel() for p in seq.parameters())
+    ok = all(np.array_equal(m.generate([p] * 2, [0, 1], scale=1), ref[p]) for p in ref)
+    if n:
+        nm = f"syn{n-1:05d}"
+        ok = ok and np.array_equal(m.generate([nm + "b"], [0], scale=1)[0], m.generate([nm], [0], scale=1)[0])
+    print(f"{len(m.classes):>8} {d:>3} {dense_p:>13,} {coded_p:>13,} {dense_p/coded_p:>6.1f}x {str(ok):>8} {interference(m):>6.0%}")
+
+print("\n(interf = share of 'wordA wordB' prompts over the 21 real classes that resolve to a THIRD class;")
+print(" dense layer: 0% by construction, ties go to lower index)")

experiments/count_outputs.py

@@ -0,0 +1,46 @@
+"""Enumerate PixelGen vs PixelGen2 output space at seed 0 and count distinct images."""
+
+import itertools
+import sys
+import time
+
+import numpy as np
+
+sys.path.insert(0, ".")
+from pixelgen.compose import PixelGen2  # noqa: E402
+from pixelgen.model import PixelGen  # noqa: E402
+from pixelgen.palette import RAMP_NAMES  # noqa: E402
+from pixelgen.parts import ACCESSORIES  # noqa: E402
+from pixelgen.templates import CLASSES  # noqa: E402
+from pixelgen.tokenizer import ACCESSORY_SYNONYMS  # noqa: E402
+
+
+def distinct(model, prompts, bs=512):
+    seen = set()
+    for i in range(0, len(prompts), bs):
+        chunk = prompts[i:i + bs]
+        imgs = model.generate(chunk, [0] * len(chunk), scale=1)
+        for im in imgs:
+            seen.add(im.tobytes())
+    return len(seen)
+
+
+m1, m2 = PixelGen(), PixelGen2()
+p1 = sum(p.numel() for p in m1.parameters()); b1 = sum(b.numel() for b in m1.buffers())
+p2 = sum(p.numel() for p in m2.parameters()); b2 = sum(b.numel() for b in m2.buffers())
+
+base = [f"{r} {c}" for c in CLASSES for r in RAMP_NAMES]
+t = time.perf_counter(); d1 = distinct(m1, base)
+print(f"PixelGen : {p1:,} params + {b1:,} buffers | {len(base):,} prompts -> {d1:,} distinct images ({time.perf_counter()-t:.1f}s)")
+
+acc_word = {a: ws[0] for a, ws in ACCESSORY_SYNONYMS.items() if ws}
+comp = []
+for c, r, size in itertools.product(CLASSES, RAMP_NAMES, ["", "small "]):
+    comp.append(f"{size}{r} {c}")
+    for a, r2 in itertools.product(acc_word, RAMP_NAMES):
+        comp.append(f"{size}{r} {c} with {r2} {acc_word[a]}")
+t = time.perf_counter(); d2 = distinct(m2, comp)
+print(f"PixelGen2: {p2:,} params + {b2:,} buffers | {len(comp):,} prompts -> {d2:,} distinct images ({time.perf_counter()-t:.1f}s)")
+print(f"\nhand-drawn content: {len(CLASSES)} bodies + {len(ACCESSORIES)-1} parts = {len(CLASSES)+len(ACCESSORIES)-1} templates")
+print(f"distinct images per template: PixelGen {d1/len(CLASSES):.0f}   PixelGen2 {d2/(len(CLASSES)+len(ACCESSORIES)-1):.0f}")
+print(f"param growth {p2/p1:.2f}x   output growth {d2/d1:.1f}x")

experiments/scale_classes.py

@@ -0,0 +1,83 @@
+"""How many classes can PixelGen absorb before something gives?
+
+Synthesises N extra classes (random valid 16x8 templates + 3 unique nonsense
+words each), rebuilds the model at each N, and measures:
+
+  params / buffers / bytes       - the capacity cost
+  build time                     - cost of "compiling" the weights
+  generate latency (1 prompt)    - inference cost
+  isolation                      - the real 21 classes still bit-identical?
+  dense-layer share              - how much of the model is the word->class matrix
+
+Nothing is trained. The synthetic classes exist only in this process.
+"""
+
+import random
+import sys
+import time
+
+import numpy as np
+import torch
+
+import pixelgen.model as M
+import pixelgen.templates as T
+import pixelgen.tokenizer as K
+
+REAL_TEMPLATES = dict(T.TEMPLATES)
+REAL_SYN = dict(K.CLASS_SYNONYMS)
+REAL_CLASSES = list(REAL_TEMPLATES)
+
+
+def random_template(rng):
+    rows = ["." * 8, "." * 8]
+    for _ in range(12):
+        row = ".." + "".join(rng.choice("##.?") for _ in range(6))
+        rows.append(row)
+    rows += ["." * 8, "." * 8]
+    on = sum(r.count("#") for r in rows); free = sum(r.count("?") for r in rows)
+    if on < 8 or on + free > 60:
+        return random_template(rng)
+    return rows
+
+
+def build(n_extra, seed=0):
+    rng = random.Random(seed)
+    tmpl = dict(REAL_TEMPLATES); syn = dict(REAL_SYN)
+    for i in range(n_extra):
+        name = f"syn{i:05d}"
+        tmpl[name] = random_template(rng)
+        syn[name] = [name, f"{name}a", f"{name}b"]
+    T.TEMPLATES = tmpl; T.CLASSES = list(tmpl)
+    K.CLASS_SYNONYMS = syn; K.VOCAB = K.build_vocab()
+    M.TEMPLATES, M.CLASSES, M.CLASS_SYNONYMS, M.VOCAB = tmpl, T.CLASSES, syn, K.VOCAB
+    t0 = time.perf_counter(); m = M.PixelGen(); dt = time.perf_counter() - t0
+    return m, dt
+
+
+# reference outputs from the real model
+ref_model, _ = build(0)
+probe = REAL_CLASSES + ["red potion", "blue invader", "dragon", "", "golden key", "bone"]
+ref_out = {p: ref_model.generate([p] * 2, [0, 1], scale=1) for p in probe}
+
+Ns = [int(a) for a in sys.argv[1:]] or [0, 10, 30, 100, 300, 1000, 3000]
+print(f"{'classes':>8} {'vocab':>6} {'params':>11} {'buffers':>9} {'MB':>7} {'dense%':>6} {'build s':>8} {'gen ms':>7} {'isolated':>8}")
+for n in Ns:
+    m, dt = build(n)
+    P = sum(p.numel() for p in m.parameters()); B = sum(b.numel() for b in m.buffers())
+    dense = m.word_to_class.weight.numel()
+    t0 = time.perf_counter()
+    for _ in range(5): m.generate(["red potion"], [0], scale=1)
+    gen = (time.perf_counter() - t0) / 5 * 1000
+    ok = all(np.array_equal(m.generate([p] * 2, [0, 1], scale=1), ref_out[p]) for p in probe)
+    # also check a synthetic class actually resolves
+    if n:
+        nm = f"syn{n-1:05d}"
+        got = m.generate([f"{nm}b"], [0], scale=1)[0].sum(-1) > 0
+        want = m.generate([nm], [0], scale=1)[0].sum(-1) > 0
+        ok = ok and np.array_equal(got, want)
+    print(f"{len(m.classes):>8} {len(m.vocab):>6} {P:>11,} {B:>9,} {(P+B)*4/1e6:>7.2f} {100*dense/(P+B):>5.1f}% {dt:>8.2f} {gen:>7.2f} {str(ok):>8}")
+
+print("\nparam model: word_to_class = V*C where V≈3C  ->  O(C^2); banks = 256*C  -> O(C)")
+for C in (21, 100, 300, 1000, 3000, 10000):
+    V = 59 - 0 + 3 * (C - 21) if C >= 21 else 59
+    print(f"  C={C:>6}: dense {V*C:>12,}  banks {256*C:>10,}  dense/banks = {V*C/(256*C):.1f}x")

generate.py

@@ -0,0 +1,35 @@
+"""Generate a sprite from a prompt.
+
+    python generate.py "red potion" --seed 3 --scale 8 --out sprite.png
+"""
+
+import argparse
+import sys
+from pathlib import Path
+
+from pixelgen import weights
+from pixelgen.compose import PixelGen2
+from pixelgen.model import PixelGen
+
+
+def main(argv=None):
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("prompt")
+    ap.add_argument("--seed", type=int, default=0)
+    ap.add_argument("--scale", type=int, default=8)
+    ap.add_argument("--out", default="sprite.png")
+    ap.add_argument("--weights", default=None, help="load a saved weights file instead of building")
+    ap.add_argument("--v1", action="store_true", help="use the original single-part model")
+    args = ap.parse_args(argv)
+
+    model = weights.load(args.weights) if args.weights else (PixelGen() if args.v1 else PixelGen2())
+    img = model.sprite(args.prompt, seed=args.seed, scale=args.scale)
+    out = Path(args.out)
+    out.parent.mkdir(parents=True, exist_ok=True)
+    img.save(out)
+    print(f"wrote {out} ({img.size[0]}x{img.size[1]})")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

pixelgen/compose.py

@@ -0,0 +1,170 @@
+"""PixelGen2: the same hand-written model, plus four compositional layers.
+
+Still zero training. New fixed weights on top of PixelGen:
+
+  A. accessory WTA      Linear(V, A) -> argmax -> one-hot part, overlaid
+                        on the body with max(). 'none' wins ties (bias 0.5).
+  B. second ramp        Linear(V, R) over the words AFTER a separator
+                        ('with', 'wearing', ...). If that half has no colour
+                        word, the body's ramp wins via a 0.5 one-hot bonus.
+  C. structural shade   light where the up-left neighbour is background,
+                        dark where the down-right neighbour is background.
+                        Two fixed 3x3 convs. Replaces the fixed row bands,
+                        which stay only as the interior base.
+  D. size               'small' -> nearest-neighbour 0.75x of the body map,
+                        bottom-centred. One fixed resampling op.
+
+Output space = classes x accessories x ramps x ramps x sizes, all composed
+from parts a human drew once.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .model import N_COLORS, PixelGen, SHADE_ROWS, noise_for_seed
+from .palette import COLOR_WORDS, OUTLINE, RAMP_NAMES
+from .parts import ACCESSORIES, PARTS
+from .templates import CLASSES, GRID, HALF
+from .tokenizer import ACCESSORY_SYNONYMS, CLASS_SYNONYMS, SIZE_WORDS, VOCAB, pad_batch_split
+
+import numpy as np
+
+SMALL = 12  # 'small' sprites are 12x12 inside the 16x16 grid
+
+
+class PixelGen2(PixelGen):
+    def __init__(self):
+        super().__init__()
+        V, A, R = len(VOCAB), len(ACCESSORIES), len(RAMP_NAMES)
+        self.accessories = ACCESSORIES
+
+        # A. word -> accessory. 'none' biased so plain prompts stay plain.
+        self.word_to_acc = nn.Linear(V, A)
+        with torch.no_grad():
+            self.word_to_acc.weight.zero_(); self.word_to_acc.bias.zero_()
+            for ai, acc in enumerate(ACCESSORIES):
+                for w in ACCESSORY_SYNONYMS[acc]:
+                    self.word_to_acc.weight[ai, VOCAB[w]] = 1.0
+            self.word_to_acc.bias[ACCESSORIES.index("none")] = 0.5
+        # In the body half, a word that names a body class means the body
+        # ('crown' = the object). Only the accessory half can use it as a hat.
+        mask = torch.ones(V)
+        for ws in CLASS_SYNONYMS.values():
+            for w in ws:
+                mask[VOCAB[w]] = 0.0
+        self.register_buffer("acc_mask_a", mask)
+
+        bank = torch.zeros(A, GRID, HALF)
+        for ai, acc in enumerate(ACCESSORIES):
+            for r, row in enumerate(PARTS[acc]):
+                for c, ch in enumerate(row):
+                    bank[ai, r, c] = float(ch == "#")
+        self.register_buffer("acc_bank", bank)
+
+        # B. second ramp from the accessory half of the prompt, no default bias
+        #    (the body ramp is added as a 0.5 bonus in forward).
+        self.word_to_ramp2 = nn.Linear(V, R)
+        with torch.no_grad():
+            self.word_to_ramp2.weight.zero_(); self.word_to_ramp2.bias.zero_()
+            for w, name in COLOR_WORDS.items():
+                self.word_to_ramp2.weight[RAMP_NAMES.index(name), VOCAB[w]] = 1.0
+
+        # C. structural shading: does the up-left / down-right neighbour exist?
+        self.lit = nn.Conv2d(1, 1, 3, padding=1, bias=False)    # up-left tap
+        self.shadow = nn.Conv2d(1, 1, 3, padding=1, bias=False)  # down-right tap
+        with torch.no_grad():
+            self.lit.weight.zero_(); self.lit.weight[0, 0, 0, 0] = 1.0
+            self.shadow.weight.zero_(); self.shadow.weight[0, 0, 2, 2] = 1.0
+
+        # D. size: 'small' words -> 1, else 0.
+        self.word_to_size = nn.Linear(V, 1)
+        with torch.no_grad():
+            self.word_to_size.weight.zero_(); self.word_to_size.bias.zero_()
+            for w in SIZE_WORDS:
+                self.word_to_size.weight[0, VOCAB[w]] = 1.0
+
+        for p in self.parameters():
+            p.requires_grad_(False)
+        self.eval()
+
+    # ------------------------------------------------------------------
+    def _bag(self, tokens):
+        B = tokens.shape[0]
+        bag = torch.zeros(B, len(self.vocab))
+        bag.scatter_add_(1, tokens, torch.ones(tokens.shape, dtype=torch.float32))
+        bag[:, 0] = 0.0
+        return bag
+
+    def _shade(self, mask):
+        """(B,H,W) mask -> (B,H,W) shade index 0 light / 1 mid / 2 dark."""
+        band = torch.zeros(GRID, GRID)
+        for s, rows in SHADE_ROWS.items():
+            for r in rows:
+                band[r, :] = float(s)
+        m = mask[:, None]
+        lit = (self.lit(m)[:, 0] < 0.5) & (mask > 0.5)       # up-left is empty
+        dark = (self.shadow(m)[:, 0] < 0.5) & (mask > 0.5)   # down-right is empty
+        shade = band.expand_as(mask).clone()
+        shade[lit] = 0.0
+        shade[dark & ~lit] = 2.0
+        return shade
+
+    def _colour(self, mask, ramp, shade):
+        ramp_idx = self.ramp_table[ramp].float()  # (B,3)
+        oh = F.one_hot(shade.long(), 3).float()   # (B,H,W,3)
+        return mask * torch.einsum("bhws,bs->bhw", oh, ramp_idx)
+
+    def forward2(self, tok_a, tok_b, noise):
+        bag_a, bag_b = self._bag(tok_a), self._bag(tok_b)
+        bag = bag_a + bag_b
+        B = bag.shape[0]
+
+        cls = self.word_to_class(bag_a).argmax(-1)
+        ramp = self.word_to_ramp(bag_a).argmax(-1)
+        acc = (self.word_to_acc(bag_a * self.acc_mask_a) + self.word_to_acc(bag_b) - self.word_to_acc.bias).argmax(-1)
+        ramp2 = (self.word_to_ramp2(bag_b) + 0.5 * F.one_hot(ramp, len(RAMP_NAMES))).argmax(-1)
+        small = (self.word_to_size(bag)[:, 0] > 0.5)
+
+        oh_c = F.one_hot(cls, len(CLASSES)).float()
+        half = torch.einsum("bc,crh->brh", oh_c, self.on_bank) \
+             + torch.einsum("bc,crh->brh", oh_c, self.free_bank) * (noise >= 0.5).float()
+        body = torch.cat([half, half.flip(-1)], dim=-1)
+
+        # D. size
+        if small.any():
+            s = F.interpolate(body[:, None], size=(SMALL, SMALL), mode="nearest")[:, 0]
+            pad = GRID - SMALL
+            s = F.pad(s, (pad // 2, pad - pad // 2, pad, 0))  # bottom-centred
+            body = torch.where(small[:, None, None], s, body)
+
+        oh_a = F.one_hot(acc, len(ACCESSORIES)).float()
+        ahalf = torch.einsum("ba,arh->brh", oh_a, self.acc_bank)
+        accm = torch.cat([ahalf, ahalf.flip(-1)], dim=-1)  # accessory sits on top
+        body_only = body * (1.0 - accm)
+
+        union = torch.maximum(body, accm)
+        dil = self.outline(union[:, None])[:, 0].clamp(0, 1)
+        outline = dil * (1.0 - union)
+
+        colour = self._colour(body_only, ramp, self._shade(body)) \
+               + self._colour(accm, ramp2, self._shade(accm)) \
+               + outline * float(OUTLINE)
+        idx = colour.round().long()
+        return F.one_hot(idx, N_COLORS).permute(0, 3, 1, 2).float()
+
+    @torch.no_grad()
+    def generate(self, prompts, seeds, scale=8, noise=None):
+        if isinstance(prompts, str):
+            prompts = [prompts]
+        prompts = list(prompts)
+        seeds = [seeds] * len(prompts) if isinstance(seeds, int) else list(seeds)
+        if len(seeds) != len(prompts):
+            raise ValueError("prompts and seeds must have the same length")
+        tok_a, tok_b = pad_batch_split(prompts, self.vocab)
+        if noise is None:
+            noise = torch.stack([noise_for_seed(s) for s in seeds])
+        img = self.rgb(self.forward2(tok_a, tok_b, noise))
+        if scale > 1:
+            img = np.repeat(np.repeat(img, scale, axis=1), scale, axis=2)
+        return img

pixelgen/model.py

@@ -0,0 +1,177 @@
+"""PixelGen: a tiny pixel-art generator whose weights are written by hand.
+
+Nothing here is trained. Every weight is set explicitly in __init__ from
+the templates and palette in this package. The forward pass is:
+
+  1. bag of words        prompt -> word counts (one vector over the vocab)
+  2. class WTA           linear(word counts) -> argmax -> one-hot class
+  3. shape bank          one-hot class selects a 16x8 half-template;
+                         '?' cells take the noise bit (noise >= 0.5)
+  4. mirror              half-template + its horizontal flip -> 16x16 body
+  5. outline             3x3 dilation of body, minus body -> palette 1
+  6. shading             body rows get light / mid / dark ramp colours
+  7. ramp WTA            linear(colour word counts) -> argmax -> ramp
+  8. one-hot palette     each pixel becomes a one-hot over the 16 PICO-8 colours
+  9. palette conv        1x1 conv from one-hot to RGB (exact PICO-8 values)
+"""
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .palette import COLOR_WORDS, DEFAULT_RAMP, OUTLINE, PICO8, RAMP_NAMES, RAMPS
+from .templates import CLASSES, GRID, HALF, TEMPLATES
+from .tokenizer import CLASS_SYNONYMS, VOCAB, pad_batch, tokenize
+
+N_COLORS = len(PICO8)
+N_SHADES = 3
+# rows of the 16-row grid that get each shade (body only lives in rows 2..13)
+SHADE_ROWS = {0: range(0, 7), 1: range(7, 11), 2: range(11, GRID)}
+DEFAULT_CLASS = "potion"
+
+
+class PixelGen(nn.Module):
+    def __init__(self):
+        super().__init__()
+        V, C, R = len(VOCAB), len(CLASSES), len(RAMP_NAMES)
+        self.vocab = VOCAB
+        self.classes = CLASSES
+        self.ramps = RAMP_NAMES
+        self.palette = PICO8
+        self.grid = GRID
+
+        # 1-2. word -> class scores. Each class word adds 1 to its class.
+        #      A 0.5 bias on the default class wins ties and empty prompts.
+        self.word_to_class = nn.Linear(V, C)
+        with torch.no_grad():
+            self.word_to_class.weight.zero_()
+            self.word_to_class.bias.zero_()
+            for ci, cls in enumerate(CLASSES):
+                for w in CLASS_SYNONYMS[cls]:
+                    self.word_to_class.weight[ci, VOCAB[w]] = 1.0
+            self.word_to_class.bias[CLASSES.index(DEFAULT_CLASS)] = 0.5
+
+        # 7. colour word -> ramp scores, same scheme, default ramp biased.
+        self.word_to_ramp = nn.Linear(V, R)
+        with torch.no_grad():
+            self.word_to_ramp.weight.zero_()
+            self.word_to_ramp.bias.zero_()
+            for w, name in COLOR_WORDS.items():
+                self.word_to_ramp.weight[RAMP_NAMES.index(name), VOCAB[w]] = 1.0
+            self.word_to_ramp.bias[RAMP_NAMES.index(DEFAULT_RAMP)] = 0.5
+
+        # 3. shape bank: per class, a 16x8 on-mask and a free-mask.
+        on = torch.zeros(C, GRID, HALF)
+        free = torch.zeros(C, GRID, HALF)
+        for ci, cls in enumerate(CLASSES):
+            for r, row in enumerate(TEMPLATES[cls]):
+                for c, ch in enumerate(row):
+                    if ch == "#":
+                        on[ci, r, c] = 1.0
+                    elif ch == "?":
+                        free[ci, r, c] = 1.0
+        self.register_buffer("on_bank", on)
+        self.register_buffer("free_bank", free)
+
+        # 6. shade bands: one-hot (3,16,16) for light / mid / dark by row.
+        shade = torch.zeros(N_SHADES, GRID, GRID)
+        for s, rows in SHADE_ROWS.items():
+            for r in rows:
+                shade[s, r, :] = 1.0
+        self.register_buffer("shade_bank", shade)
+
+        # 7. ramp table: ramp index -> [light, mid, dark] palette indices.
+        table = torch.tensor([RAMPS[name] for name in RAMP_NAMES], dtype=torch.long)
+        self.register_buffer("ramp_table", table)
+
+        # 5. outline: 3x3 all-ones conv counts body neighbours (incl. centre).
+        self.outline = nn.Conv2d(1, 1, 3, padding=1, bias=False)
+        with torch.no_grad():
+            self.outline.weight.fill_(1.0)
+
+        # 9. palette conv: 1x1, 16 one-hot channels -> RGB in [0,1].
+        self.palette_conv = nn.Conv2d(N_COLORS, 3, 1, bias=False)
+        with torch.no_grad():
+            self.palette_conv.weight.zero_()
+            for idx, rgb in enumerate(PICO8):
+                for ch in range(3):
+                    self.palette_conv.weight[ch, idx, 0, 0] = rgb[ch] / 255.0
+
+        for p in self.parameters():
+            p.requires_grad_(False)
+        self.eval()
+
+    # ------------------------------------------------------------------
+    def tokenize(self, prompt):
+        return tokenize(prompt, self.vocab)
+
+    def forward(self, tokens, noise):
+        """tokens: (B, L) long word ids. noise: (B, 16, 8) floats.
+        Returns a one-hot palette map of shape (B, 16, 16, 16)
+        (batch, palette index, row, col)."""
+        B = tokens.shape[0]
+        V = len(self.vocab)
+        bag = torch.zeros(B, V)
+        bag.scatter_add_(1, tokens, torch.ones(tokens.shape, dtype=torch.float32))
+        bag[:, 0] = 0.0  # <pad> never counts
+
+        cls = self.word_to_class(bag).argmax(-1)
+        ramp = self.word_to_ramp(bag).argmax(-1)
+
+        oh_c = F.one_hot(cls, len(CLASSES)).float()
+        on = torch.einsum("bc,crh->brh", oh_c, self.on_bank)
+        free = torch.einsum("bc,crh->brh", oh_c, self.free_bank)
+        bit = (noise >= 0.5).float()
+        half = on + free * bit
+
+        body = torch.cat([half, half.flip(-1)], dim=-1)  # exact mirror
+
+        dil = self.outline(body[:, None])[:, 0].clamp(0.0, 1.0)
+        outline = dil * (1.0 - body)
+
+        ramp_idx = self.ramp_table[ramp].float()  # (B, 3)
+        body_colour = torch.einsum("bs,shw->bhw", ramp_idx, self.shade_bank)
+        colour = body * body_colour + outline * float(OUTLINE)
+
+        idx = colour.round().long()
+        oh = F.one_hot(idx, N_COLORS).permute(0, 3, 1, 2).float()
+        return oh
+
+    @torch.no_grad()
+    def rgb(self, onehot):
+        """(B,16,H,W) one-hot -> (B,H,W,3) uint8."""
+        x = self.palette_conv(onehot) * 255.0
+        return x.round().clamp(0, 255).to(torch.uint8).permute(0, 2, 3, 1).numpy()
+
+    @torch.no_grad()
+    def generate(self, prompts, seeds, scale=8, noise=None):
+        """Return a uint8 array of shape (B, 16*scale, 16*scale, 3)."""
+        if isinstance(prompts, str):
+            prompts = [prompts]
+        prompts = list(prompts)
+        seeds = [seeds] * len(prompts) if isinstance(seeds, int) else list(seeds)
+        if len(seeds) != len(prompts):
+            raise ValueError("prompts and seeds must have the same length")
+        tokens = pad_batch(prompts, self.vocab)
+        if noise is None:
+            noise = torch.stack([noise_for_seed(s) for s in seeds])
+        else:
+            noise = torch.as_tensor(noise, dtype=torch.float32).expand(len(prompts), GRID, HALF)
+        oh = self.forward(tokens, noise)
+        img = self.rgb(oh)  # (B,16,16,3)
+        if scale > 1:
+            img = np.repeat(np.repeat(img, scale, axis=1), scale, axis=2)
+        return img
+
+    def sprite(self, prompt, seed=0, scale=8):
+        from PIL import Image
+
+        arr = self.generate([prompt], [seed], scale=scale)[0]
+        return Image.fromarray(arr, mode="RGB")
+
+
+def noise_for_seed(seed):
+    """Deterministic per-seed noise in [0,1), shape (GRID, HALF)."""
+    g = torch.Generator().manual_seed(int(seed))
+    return torch.rand(GRID, HALF, generator=g)

pixelgen/palette.py

@@ -0,0 +1,63 @@
+"""PICO-8 palette, colour ramps and the words that select them.
+
+The network's final layer can only emit one-hot vectors over these 16
+colours, so nothing outside this palette can ever appear in an image.
+"""
+
+PICO8 = [
+    (0x00, 0x00, 0x00),  # 0  black        (background)
+    (0x1D, 0x2B, 0x53),  # 1  dark blue    (outline)
+    (0x7E, 0x25, 0x53),  # 2  dark purple
+    (0x00, 0x87, 0x51),  # 3  dark green
+    (0xAB, 0x52, 0x36),  # 4  brown
+    (0x5F, 0x57, 0x4F),  # 5  dark grey
+    (0xC2, 0xC3, 0xC7),  # 6  light grey
+    (0xFF, 0xF1, 0xE8),  # 7  white
+    (0xFF, 0x00, 0x4D),  # 8  red
+    (0xFF, 0xA3, 0x00),  # 9  orange
+    (0xFF, 0xEC, 0x27),  # 10 yellow
+    (0x00, 0xE4, 0x36),  # 11 green
+    (0x29, 0xAD, 0xFF),  # 12 blue
+    (0x83, 0x76, 0x9C),  # 13 lavender
+    (0xFF, 0x77, 0xA8),  # 14 pink
+    (0xFF, 0xCC, 0xAA),  # 15 peach
+]
+
+BACKGROUND = 0
+OUTLINE = 1
+
+# Each ramp is [light, mid, dark] palette indices. Indices 0 and 1 are
+# reserved for background and outline so a sprite never has more than
+# 2 + 3 = 5 colours.
+RAMPS = {
+    "red": [14, 8, 2],
+    "blue": [12, 13, 2],
+    "green": [11, 3, 5],
+    "yellow": [10, 9, 4],
+    "purple": [14, 13, 2],
+    "orange": [15, 9, 4],
+    "pink": [15, 14, 2],
+    "brown": [15, 4, 5],
+    "white": [7, 6, 5],
+}
+RAMP_NAMES = list(RAMPS)
+
+# Ramp used when the prompt names no colour word.
+DEFAULT_RAMP = "green"
+
+# Prompt words that force a ramp.
+COLOR_WORDS = {
+    "red": "red", "crimson": "red", "ruby": "red",
+    "blue": "blue", "azure": "blue", "sapphire": "blue",
+    "green": "green", "emerald": "green", "lime": "green",
+    "yellow": "yellow", "gold": "yellow", "golden": "yellow",
+    "purple": "purple", "violet": "purple", "amethyst": "purple",
+    "orange": "orange", "amber": "orange",
+    "pink": "pink", "rose": "pink",
+    "brown": "brown", "wood": "brown", "wooden": "brown",
+    "white": "white", "grey": "white", "gray": "white", "silver": "white", "bone": "white",
+}
+
+for _name, _ramp in RAMPS.items():
+    assert len(_ramp) == 3, _name
+    assert all(2 <= i <= 15 for i in _ramp), _name

pixelgen/parts.py

@@ -0,0 +1,153 @@
+"""Accessory parts: 16x8 half-templates overlaid on any body.
+
+Same alphabet as templates.py ('#' on, '.' off). Mirrored like bodies, so a
+'torch' is two torches. Rows 0 and 15 stay empty so the outline has room.
+Hats live in rows 1-4 (most bodies start at row 3); side parts in rows 5-10,
+columns 2-3 (outside most bodies).
+"""
+
+from .templates import GRID, HALF
+
+PARTS = {
+    "none": ["........"] * GRID,
+    "crown": [
+        "........",
+        "...#.#.#",
+        "...#####",
+        "...#####",
+        "........",
+        "........",
+        "........",
+        "........",
+        "........",
+        "........",
+        "........",
+        "........",
+        "........",
+        "........",
+        "........",
+        "........",
+    ],
+    "halo": [
+        ".....###",
+        "....#...",
+        ".....###",
+        "........",
+        "........",
+        "........",
+        "........",
+        "........",
+        "........",
+        "........",
+        "........",
+        "........",
+        "........",
+        "........",
+        "........",
+        "........",
+    ],
+    "horns": [
+        "........",
+        "...#....",
+        "...#....",
+        "....#...",
+        "....##..",
+        "........",
+        "........",
+        "........",
+        "........",
+        "........",
+        "........",
+        "........",
+        "........",
+        "........",
+        "........",
+        "........",
+    ],
+    "helmet": [
+        "........",
+        "........",
+        ".....###",
+        "....####",
+        "...#####",
+        "...##...",
+        "........",
+        "........",
+        "........",
+        "........",
+        "........",
+        "........",
+        "........",
+        "........",
+        "........",
+        "........",
+    ],
+    "antenna": [
+        "........",
+        ".....#..",
+        "......#.",
+        "......#.",
+        "........",
+        "........",
+        "........",
+        "........",
+        "........",
+        "........",
+        "........",
+        "........",
+        "........",
+        "........",
+        "........",
+        "........",
+    ],
+    "wings": [
+        "........",
+        "........",
+        "........",
+        "........",
+        "........",
+        "..#.....",
+        "..##....",
+        "..###...",
+        "..###...",
+        "..##....",
+        "..#.....",
+        "........",
+        "........",
+        "........",
+        "........",
+        "........",
+    ],
+    "torch": [
+        "........",
+        "........",
+        "........",
+        "..#.....",
+        "..##....",
+        "..#.....",
+        "..#.....",
+        "..#.....",
+        "..#.....",
+        "..#.....",
+        "........",
+        "........",
+        "........",
+        "........",
+        "........",
+        "........",
+    ],
+}
+
+ACCESSORIES = list(PARTS)
+
+
+def _check():
+    for name, rows in PARTS.items():
+        assert len(rows) == GRID, name
+        for r, row in enumerate(rows):
+            assert len(row) == HALF and set(row) <= set("#."), (name, r, row)
+            assert row[:2] == ".." or name in ("wings", "torch"), (name, r)
+        assert rows[-1] == "." * HALF, name
+
+
+_check()

pixelgen/templates.py

@@ -0,0 +1,421 @@
+"""Hand-drawn sprite templates: the "program" for the shape layer.
+
+Each template is the LEFT half of a 16x16 sprite: 16 rows x 8 columns.
+Column 7 is the column just left of the centre line; the network mirrors
+the half to produce a symmetric sprite.
+
+    '#'  forced ON   (always body)
+    '.'  forced OFF  (never body)
+    '?'  free        (decided by one noise bit, then smoothed)
+
+Rules the shape layer relies on:
+  * rows 0,1,14,15 and columns 0,1 are '.', so the body sits in rows/cols
+    2..13 of the full grid and a 1px outline fits inside the border.
+  * at least 8 '#' per half  -> at least 16 body cells.
+  * at most 60 '#'+'?' per half -> at most 120 body cells.
+"""
+
+GRID = 16
+HALF = GRID // 2
+
+TEMPLATES = {
+    "potion": [
+        "........",
+        "........",
+        "......##",
+        "......##",
+        ".......#",
+        ".......#",
+        ".....?##",
+        "...?####",
+        "..?#####",
+        "..?#####",
+        "..?#####",
+        "..?#####",
+        "...?####",
+        "....?###",
+        "........",
+        "........",
+    ],
+    "sword": [
+        "........",
+        "........",
+        ".......#",
+        ".......#",
+        "......?#",
+        "......?#",
+        "......?#",
+        "......?#",
+        "......?#",
+        "....####",
+        "....?###",
+        ".......#",
+        ".......#",
+        "......##",
+        "........",
+        "........",
+    ],
+    "shield": [
+        "........",
+        "........",
+        "...#####",
+        "..??####",
+        "..#?####",
+        "..#?#?##",
+        "..#?####",
+        "..#?#?##",
+        "..#?####",
+        "...#####",
+        "...?####",
+        "....?###",
+        ".....?##",
+        "......?#",
+        "........",
+        "........",
+    ],
+    "ship": [
+        "........",
+        "........",
+        ".......#",
+        ".......#",
+        "......##",
+        "......##",
+        ".....?##",
+        ".....###",
+        "....?###",
+        "...?####",
+        "..?#?###",
+        "..##?###",
+        "..#?..##",
+        "......?#",
+        "........",
+        "........",
+    ],
+    "invader": [
+        "........",
+        "........",
+        "........",
+        "........",
+        "....#...",
+        ".....#..",
+        "....####",
+        "...##.##",
+        "..######",
+        "..#.####",
+        "..#.#..?",
+        ".....##.",
+        "........",
+        "........",
+        "........",
+        "........",
+    ],
+    "slime": [
+        "........",
+        "........",
+        "........",
+        "........",
+        "........",
+        "......??",
+        ".....?##",
+        "....?###",
+        "...?####",
+        "...##.##",
+        "..?#####",
+        "..######",
+        "..?#####",
+        "...?####",
+        "........",
+        "........",
+    ],
+    "skull": [
+        "........",
+        "........",
+        "........",
+        ".....?##",
+        "....####",
+        "...#####",
+        "...#####",
+        "...#..##",
+        "...#..##",
+        "...#####",
+        "....####",
+        "....#.#.",
+        "....#.#.",
+        "........",
+        "........",
+        "........",
+    ],
+    "tree": [
+        "........",
+        "........",
+        ".......#",
+        "......?#",
+        "......##",
+        ".....?##",
+        ".....###",
+        "....?###",
+        "....####",
+        "...?####",
+        "...#####",
+        "..?#####",
+        "......##",
+        "......##",
+        "........",
+        "........",
+    ],
+    "key": [
+        "........",
+        "........",
+        "........",
+        "......##",
+        ".....#..",
+        ".....#..",
+        "......##",
+        ".......#",
+        ".......#",
+        ".......#",
+        "......##",
+        ".......#",
+        ".....###",
+        ".......#",
+        "........",
+        "........",
+    ],
+    "heart": [
+        "........",
+        "........",
+        "........",
+        "........",
+        "..?###..",
+        "..#####?",
+        "..######",
+        "..######",
+        "..######",
+        "...#####",
+        "....####",
+        ".....###",
+        "......##",
+        ".......#",
+        "........",
+        "........",
+    ],
+    "star": [
+        "........",
+        "........",
+        ".......#",
+        ".......#",
+        "......##",
+        "......##",
+        "..?#####",
+        "...#####",
+        "....####",
+        "....####",
+        "...?####",
+        "...##?.#",
+        "..##...?",
+        "..#.....",
+        "........",
+        "........",
+    ],
+    "mushroom": [
+        "........",
+        "........",
+        "........",
+        ".....###",
+        "....##?#",
+        "...##.##",
+        "..##?###",
+        "..######",
+        "..######",
+        ".....###",
+        ".....#.#",
+        ".....###",
+        ".....###",
+        "........",
+        "........",
+        "........",
+    ],
+    "ghost": [
+        "........",
+        "........",
+        "........",
+        ".....###",
+        "....####",
+        "...#####",
+        "...##.##",
+        "...##.##",
+        "...#####",
+        "...#####",
+        "...#####",
+        "...#####",
+        "...#?#.#",
+        "...#.#..",
+        "........",
+        "........",
+    ],
+    "bomb": [
+        "........",
+        "........",
+        ".......#",
+        "......#.",
+        "......##",
+        ".....###",
+        "....####",
+        "...#####",
+        "...#####",
+        "...##?##",
+        "...#####",
+        "....####",
+        ".....###",
+        "........",
+        "........",
+        "........",
+    ],
+    "coin": [
+        "........",
+        "........",
+        "........",
+        ".....###",
+        "....####",
+        "...#####",
+        "...##..#",
+        "...#..##",
+        "...#..##",
+        "...##..#",
+        "...#####",
+        "....####",
+        ".....###",
+        "........",
+        "........",
+        "........",
+    ],
+    "gem": [
+        "........",
+        "........",
+        "........",
+        "....####",
+        "...#####",
+        "..######",
+        "..#?####",
+        "...#####",
+        "....####",
+        ".....###",
+        "......##",
+        ".......#",
+        "........",
+        "........",
+        "........",
+        "........",
+    ],
+    "cherry": [
+        "........",
+        "........",
+        "........",
+        ".......#",
+        "......#.",
+        ".....#..",
+        "....#...",
+        "....#...",
+        "...###..",
+        "..####..",
+        "..####..",
+        "..#?##..",
+        "...##...",
+        "........",
+        "........",
+        "........",
+    ],
+    "crown": [
+        "........",
+        "........",
+        "........",
+        "..#....#",
+        "..#..#.#",
+        "..##.###",
+        "..######",
+        "..######",
+        "..######",
+        "..######",
+        "..######",
+        "..######",
+        "........",
+        "........",
+        "........",
+        "........",
+    ],
+    "bat": [
+        "........",
+        "........",
+        "........",
+        "......##",
+        ".....###",
+        ".......#",
+        "..#..###",
+        "..##.###",
+        "..######",
+        "...#####",
+        "....####",
+        "....#.##",
+        "........",
+        "........",
+        "........",
+        "........",
+    ],
+    "cactus": [
+        "........",
+        "........",
+        "........",
+        "......##",
+        "......##",
+        "..#...##",
+        "..#...##",
+        "..##..##",
+        "...#####",
+        "......##",
+        "......##",
+        "......##",
+        ".....###",
+        "........",
+        "........",
+        "........",
+    ],
+    "robot": [
+        "........",
+        "........",
+        "........",
+        "......##",
+        ".....###",
+        "...#####",
+        "...#.###",
+        "...#####",
+        "...#####",
+        "....####",
+        "..?#####",
+        "...#####",
+        "...##.##",
+        "...##.##",
+        "........",
+        "........",
+    ],
+}
+
+CLASSES = list(TEMPLATES)
+
+
+def _check():
+    for name, rows in TEMPLATES.items():
+        assert len(rows) == GRID, name
+        on = free = 0
+        for r, row in enumerate(rows):
+            assert len(row) == HALF, (name, r, row)
+            assert set(row) <= set("#.?"), (name, r, row)
+            if r in (0, 1, GRID - 2, GRID - 1):
+                assert row == "." * HALF, (name, r)
+            assert row[:2] == "..", (name, r)
+            on += row.count("#")
+            free += row.count("?")
+        assert on >= 8, (name, on)
+        assert on + free <= 60, (name, on + free)
+
+
+_check()

pixelgen/tokenizer.py

@@ -0,0 +1,108 @@
+"""Word-level tokenizer. Unknown words map to id 0 (<pad>), which the
+embedding layer ignores, so garbage prompts are harmless."""
+
+import re
+
+import torch
+
+from .palette import COLOR_WORDS
+from .templates import CLASSES
+
+CLASS_SYNONYMS = {
+    "potion": ["potion", "bottle", "flask", "elixir", "vial"],
+    "sword": ["sword", "blade", "dagger"],
+    "shield": ["shield", "buckler"],
+    "ship": ["ship", "rocket", "spaceship", "starship"],
+    "invader": ["invader", "alien", "monster"],
+    "slime": ["slime", "blob", "jelly"],
+    "skull": ["skull", "bone", "skeleton"],
+    "tree": ["tree", "pine", "forest"],
+    "key": ["key"],
+    "heart": ["heart", "love", "life"],
+    "star": ["star", "sparkle"],
+    "mushroom": ["mushroom", "fungus", "toadstool", "shroom"],
+    "ghost": ["ghost", "spirit", "phantom", "spooky"],
+    "bomb": ["bomb", "grenade", "explosive"],
+    "coin": ["coin", "token", "money", "medal"],
+    "gem": ["gem", "diamond", "jewel", "crystal"],
+    "cherry": ["cherry", "cherries", "berry", "fruit"],
+    "crown": ["crown"],
+    "bat": ["bat", "vampire"],
+    "cactus": ["cactus", "desert", "succulent"],
+    "robot": ["robot", "droid", "android", "mech"],
+}
+
+PAD = "<pad>"
+
+# Accessory words select a second, overlaid part (see parts.py).
+ACCESSORY_SYNONYMS = {
+    "none": [],
+    "crown": ["crown", "king", "queen", "royal", "crowned", "prince", "princess"],
+    "halo": ["halo", "angel", "angelic", "holy", "saint"],
+    "horns": ["horns", "horned", "devil", "demon", "evil"],
+    "helmet": ["helmet", "knight", "armor", "armour", "armored"],
+    "antenna": ["antenna", "antennae", "bug", "martian"],
+    "wings": ["wings", "winged", "fly", "flying", "bird", "fairy", "butterfly"],
+    "torch": ["torch", "torches", "fire", "flame", "burning"],
+}
+SIZE_WORDS = {"small": 1, "tiny": 1, "mini": 1, "little": 1}
+# These split a prompt in two: words after them describe the accessory.
+SEPARATORS = ["with", "wearing", "holding", "and", "plus"]
+
+
+def build_vocab():
+    vocab = {PAD: 0}
+    groups = [
+        set(w for ws in CLASS_SYNONYMS.values() for w in ws),
+        set(COLOR_WORDS),
+        set(w for ws in ACCESSORY_SYNONYMS.values() for w in ws),
+        set(SIZE_WORDS),
+        set(SEPARATORS),
+    ]
+    for g in groups:
+        for word in sorted(g):
+            if word not in vocab:
+                vocab[word] = len(vocab)
+    return vocab
+
+
+VOCAB = build_vocab()
+
+_SPLIT = re.compile(r"[^a-z]+")
+
+
+def tokenize(prompt, vocab=VOCAB):
+    words = [w for w in _SPLIT.split(prompt.lower()) if w]
+    ids = [vocab.get(w, 0) for w in words]
+    return torch.tensor(ids if ids else [0], dtype=torch.long)
+
+
+def pad_batch(prompts, vocab=VOCAB):
+    toks = [tokenize(p, vocab) for p in prompts]
+    L = max(t.numel() for t in toks)
+    out = torch.zeros(len(toks), L, dtype=torch.long)
+    for i, t in enumerate(toks):
+        out[i, : t.numel()] = t
+    return out
+
+
+def split_prompt(prompt):
+    """'green slime with gold crown' -> ('green slime', 'gold crown').
+    No separator -> whole prompt is the body half, accessory half empty."""
+    words = [w for w in _SPLIT.split(prompt.lower()) if w]
+    for i, w in enumerate(words):
+        if w in SEPARATORS:
+            return " ".join(words[:i]), " ".join(words[i + 1:])
+    return " ".join(words), ""
+
+
+def pad_batch_split(prompts, vocab=VOCAB):
+    halves = [split_prompt(p) for p in prompts]
+    return pad_batch([a for a, _ in halves], vocab), pad_batch([b for _, b in halves], vocab)
+
+
+def word_class(word):
+    for cls, ws in CLASS_SYNONYMS.items():
+        if word in ws:
+            return CLASSES.index(cls)
+    return None

pixelgen/weights.py

@@ -0,0 +1,30 @@
+"""Save and load the hand-written PixelGen state dict.
+
+The weights are set in code, so saving is only a convenience: it lets
+generate.py run without re-building the model and gives a file to check
+into version control.
+"""
+
+from pathlib import Path
+
+import torch
+
+from .model import PixelGen
+
+DEFAULT_PATH = Path(__file__).resolve().parent.parent / "weights" / "pixelgen.pt"
+
+
+def save(path=DEFAULT_PATH, model=None):
+    model = model or PixelGen()
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    torch.save(model.state_dict(), path)
+    return path
+
+
+def load(path=DEFAULT_PATH):
+    model = PixelGen()
+    state = torch.load(path, map_location="cpu", weights_only=True)
+    model.load_state_dict(state)
+    model.eval()
+    return model

requirements.txt

@@ -0,0 +1,3 @@
+torch>=2.0
+numpy
+pillow

tests/test_pixel_art.py

@@ -0,0 +1,121 @@
+import numpy as np
+import pytest
+
+from pixelgen import weights
+from pixelgen.model import PixelGen
+from pixelgen.palette import COLOR_WORDS, PICO8, RAMP_NAMES
+from pixelgen.templates import CLASSES
+
+
+@pytest.fixture(scope="module")
+def model():
+    return PixelGen()
+
+
+def colours(img):
+    return {tuple(int(v) for v in c) for c in img.reshape(-1, 3)}
+
+
+def test_output_shape_and_dtype(model):
+    img = model.generate(["red potion"], [0], scale=8)
+    assert img.shape == (1, 128, 128, 3)
+    assert img.dtype == np.uint8
+
+
+def test_only_pico8_colours(model):
+    pico = set(PICO8)
+    for prompt in ["red potion", "blue sword", "green shield", "gold crown", ""]:
+        for seed in range(3):
+            img = model.generate([prompt], [seed], scale=1)[0]
+            assert colours(img) <= pico, (prompt, seed)
+
+
+def test_deterministic_for_same_seed(model):
+    a = model.generate(["red potion"], [7], scale=1)
+    b = model.generate(["red potion"], [7], scale=1)
+    assert np.array_equal(a, b)
+
+
+def test_sprite_is_horizontally_symmetric(model):
+    img = model.generate(["red potion"], [0], scale=1)[0]
+    assert np.array_equal(img, img[:, ::-1])
+
+
+@pytest.mark.parametrize("word,ramp", [("red", "red"), ("blue", "blue"), ("green", "green")])
+def test_colour_word_selects_ramp(model, word, ramp):
+    # The ramp's light/mid/dark colours should dominate the body pixels.
+    from pixelgen.palette import RAMPS
+
+    img = model.generate([f"{word} potion"], [0], scale=1)[0]
+    ramp_cols = {PICO8[i] for i in RAMPS[ramp]}
+    assert ramp_cols & colours(img), f"{word} produced none of {ramp} ramp colours"
+
+
+def test_every_colour_word_maps_to_a_known_ramp():
+    assert set(COLOR_WORDS.values()) <= set(RAMP_NAMES)
+
+
+def test_every_class_renders(model):
+    for cls in CLASSES:
+        img = model.generate([cls], [0], scale=1)[0]
+        assert img.shape == (16, 16, 3)
+        assert img.any(), f"{cls} rendered empty"
+
+
+def test_save_and_load_round_trip(tmp_path):
+    path = weights.save(tmp_path / "pixelgen.pt")
+    loaded = weights.load(path)
+    fresh = PixelGen()
+    a = fresh.generate(["red potion"], [1], scale=1)
+    b = loaded.generate(["red potion"], [1], scale=1)
+    assert np.array_equal(a, b)
+
+
+# ---------------------------------------------------------------- PixelGen2
+from pixelgen.compose import PixelGen2  # noqa: E402
+from pixelgen.palette import RAMPS  # noqa: E402
+
+
+@pytest.fixture(scope="module")
+def model2():
+    return PixelGen2()
+
+
+def test_v2_only_pico8_and_symmetric(model2):
+    img = model2.generate(["purple ghost with gold crown"], [0], scale=1)[0]
+    assert set(map(tuple, img.reshape(-1, 3))) <= set(map(tuple, PICO8))
+    mask = img.sum(-1) > 0  # shape is mirrored; lighting is deliberately not
+    assert np.array_equal(mask, mask[:, ::-1])
+
+
+def test_v2_plain_prompt_has_no_accessory(model2):
+    a = model2.generate(["green slime"], [0], scale=1)[0]
+    b = model2.generate(["green slime with crown"], [0], scale=1)[0]
+    assert not np.array_equal(a, b)
+    # top rows of a plain slime are background; crown fills them
+    assert (a[:3].sum(-1) > 0).sum() == 0 and (b[:3].sum(-1) > 0).sum() > 0
+
+
+def test_v2_accessory_takes_second_colour(model2):
+    img = model2.generate(["green slime with red crown"], [0], scale=1)[0]
+    red = {tuple(PICO8[i]) for i in RAMPS["red"]}
+    green = {tuple(PICO8[i]) for i in RAMPS["green"]}
+    px = set(map(tuple, img.reshape(-1, 3)))
+    assert px & red and px & green
+
+
+def test_v2_crown_word_means_object_in_body_half(model2):
+    crown = model2.generate(["crown"], [0], scale=1)[0]
+    king = model2.generate(["slime with crown"], [0], scale=1)[0]
+    assert not np.array_equal(crown, king)
+    assert (crown[:2].sum(-1) > 0).sum() == 0  # the crown object has no hat on it
+
+
+def test_v2_small_is_smaller(model2):
+    big = (model2.generate(["skull"], [0], scale=1)[0].sum(-1) > 0).sum()
+    small = (model2.generate(["small skull"], [0], scale=1)[0].sum(-1) > 0).sum()
+    assert small < big
+
+
+def test_v2_no_trainable_params(model2):
+    assert all(not p.requires_grad for p in model2.parameters())

weights/pixelgen.pt

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