7190edb verified 1 day ago

7.55 kB

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

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