PixelGen2 demo

.gitignore

@@ -0,0 +1,2 @@
+__pycache__/
+*.pyc

README.md

@@ -1,13 +1,18 @@
 ---
-title: Pixelgen2
-emoji: 🌖
-colorFrom: yellow
-colorTo: yellow
+title: PixelGen2
+emoji: 🕹️
+colorFrom: green
+colorTo: pink
 sdk: gradio
 sdk_version: 6.31.0
-python_version: '3.13'
 app_file: app.py
 pinned: false
+license: mit
+short_description: 6,747 hand-written params, 0 training steps
+models:
+  - DJLougen/pixelgen2
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# PixelGen2
+
+Text-to-pixel-art with hand-written weights and zero training steps. Model card, code and experiments: [DJLougen/pixelgen2](https://huggingface.co/DJLougen/pixelgen2).

app.py

@@ -0,0 +1,74 @@
+"""PixelGen2 demo Space: hand-written weights, zero training."""
+
+import gradio as gr
+from PIL import Image
+
+from pixelgen.compose import PixelGen2
+from pixelgen.model import PixelGen
+from pixelgen.palette import RAMP_NAMES
+from pixelgen.parts import ACCESSORIES
+from pixelgen.templates import CLASSES
+
+V2, V1 = PixelGen2(), PixelGen()
+N_PARAMS = sum(p.numel() for p in V2.parameters())
+
+
+def render(prompt, seed, scale, model):
+    m = V1 if model == "v1 (original)" else V2
+    return Image.fromarray(m.generate([prompt], [int(seed)], scale=int(scale))[0])
+
+
+def seeds(prompt, scale):
+    return [Image.fromarray(V2.generate([prompt], [s], scale=int(scale))[0]) for s in range(8)]
+
+
+def build(body, colour, acc, acc_colour, small):
+    p = f"{'small ' if small else ''}{colour} {body}"
+    if acc != "none":
+        p += f" with {acc_colour} {acc}"
+    return p
+
+
+EXAMPLES = [
+    ["green slime with gold crown", 0], ["purple ghost with horns", 3], ["red potion", 0],
+    ["small blue invader", 1], ["angel skull with white halo", 2], ["knight robot", 0],
+    ["bat with orange fire", 5], ["martian alien", 0], ["pink cherry", 4],
+]
+
+with gr.Blocks(title="PixelGen2") as demo:
+    gr.Markdown(
+        f"# 🕹️ PixelGen2\n"
+        f"A text-to-pixel-art model with **{N_PARAMS:,} parameters, zero of them trained**. "
+        f"Every weight was written by hand; it composes 21 bodies × 7 accessories × 9 colours into 21,866 distinct sprites. "
+        f"[Model card & experiments](https://huggingface.co/DJLougen/pixelgen2) · "
+        f"[weight-space programming](https://rikiyat.github.io/blog/weight-space-programming/)\n\n"
+        f"Grammar: `[small] [colour] <body> [with [colour] <accessory>]`"
+    )
+    with gr.Row():
+        with gr.Column(scale=1):
+            prompt = gr.Textbox(label="Prompt", value="green slime with gold crown")
+            with gr.Accordion("…or build one", open=False):
+                body = gr.Dropdown(CLASSES, value="slime", label="body")
+                colour = gr.Dropdown(RAMP_NAMES, value="green", label="colour")
+                acc = gr.Dropdown(ACCESSORIES, value="crown", label="accessory")
+                acc_colour = gr.Dropdown(RAMP_NAMES, value="yellow", label="accessory colour")
+                small = gr.Checkbox(label="small")
+                for c in (body, colour, acc, acc_colour, small):
+                    c.change(build, [body, colour, acc, acc_colour, small], prompt)
+            seed = gr.Slider(0, 99, value=0, step=1, label="seed")
+            scale = gr.Slider(4, 16, value=12, step=1, label="scale")
+            model = gr.Radio(["v2 (compositional)", "v1 (original)"], value="v2 (compositional)", label="model")
+            go = gr.Button("Generate", variant="primary")
+        with gr.Column(scale=1):
+            out = gr.Image(label="sprite", type="pil", height=300)
+    gr.Markdown("**Same prompt, seeds 0–7** — variation only toggles the hand-marked `?` pixels.")
+    strip = gr.Gallery(columns=8, height=140, label="seeds")
+    gr.Examples(EXAMPLES, [prompt, seed])
+
+    inputs = [prompt, seed, scale, model]
+    go.click(render, inputs, out).then(seeds, [prompt, scale], strip)
+    prompt.submit(render, inputs, out).then(seeds, [prompt, scale], strip)
+    demo.load(render, inputs, out).then(seeds, [prompt, scale], strip)
+
+if __name__ == "__main__":
+    demo.launch()

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