| """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 |
| |
| 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 |
|
|
| |
| |
| 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 |
|
|
| |
| 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 |
|
|
| |
| 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) |
|
|
| |
| 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) |
|
|
| |
| table = torch.tensor([RAMPS[name] for name in RAMP_NAMES], dtype=torch.long) |
| self.register_buffer("ramp_table", table) |
|
|
| |
| self.outline = nn.Conv2d(1, 1, 3, padding=1, bias=False) |
| with torch.no_grad(): |
| self.outline.weight.fill_(1.0) |
|
|
| |
| 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 |
|
|
| 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) |
|
|
| dil = self.outline(body[:, None])[:, 0].clamp(0.0, 1.0) |
| outline = dil * (1.0 - body) |
|
|
| ramp_idx = self.ramp_table[ramp].float() |
| 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) |
| 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) |
|
|
