"""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 # 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)