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