pxg_tiny/config.py

"""PXG-Tiny shared configuration: master palette, tokenizer, classes, model dims.

Single source of truth imported by both the training pipeline and the
inference runtime. All sprites are represented as 16x16 grids of *master
palette indices* (0 = fully transparent alpha, 1..31 = opaque colors).
This makes reconstruction byte-exact by construction and keeps the
generation vocabulary tiny (32 visual symbols).
"""

IMG = 16                 # sprite resolution (16x16, even, inside the 8..256 schema band)
N_COLORS = 32            # visual vocabulary size: index 0 = transparent, 1..31 colors

# ---------------------------------------------------------------------------
# Master palette (index -> RGBA). Hand-crafted to span the classic
# warm_fantasy / nes / cool_sci_fi flavor families while staying <= 31 opaque
# colors so a single byte encodes any pixel.
# ---------------------------------------------------------------------------
PALETTE = {
    1:  (248, 248, 248),  # white / specular
    2:  (217, 221, 230),  # pale gray
    3:  (183, 191, 204),  # silver
    4:  (120, 129, 142),  # gray
    5:  (69, 76, 87),     # charcoal
    6:  (23, 23, 31),     # ink (outline black)
    7:  (212, 61, 51),    # red
    8:  (140, 35, 32),    # maroon
    9:  (232, 137, 46),   # orange / terracotta
    10: (242, 193, 46),   # gold
    11: (255, 232, 107),  # lemon (gold highlight)
    12: (127, 201, 79),   # lime
    13: (76, 154, 68),    # green mid
    14: (44, 106, 53),    # forest dark
    15: (55, 179, 171),   # teal
    16: (159, 223, 240),  # ice / cyan highlight
    17: (79, 147, 220),   # sky blue
    18: (45, 71, 127),    # navy deep
    19: (124, 80, 173),   # purple
    20: (180, 140, 224),  # light violet
    21: (168, 123, 35),   # bronze (dark gold)
    22: (242, 203, 160),  # peach
    23: (221, 182, 113),  # sand
    24: (244, 230, 192),  # cream ivory
    25: (165, 116, 74),   # tan wood
    26: (122, 78, 42),    # brown mid
    27: (76, 48, 26),     # dark wood
    28: (111, 117, 57),   # moss olive
    29: (185, 230, 197),  # mint pale (slime)
    30: (93, 114, 136),   # steel blue-gray
    31: (238, 173, 174),  # blush soft pink
}

# ---------------------------------------------------------------------------
# Text tokenizer (char-level with byte-fallback buckets).
# id 0 = <pad>, 1 = <bos>, 2 = <sep>, 3..6 punctuation, 7..32 a-z,
# 33..62 byte-fallback buckets (ord % 30 + 33), 63 spare.
# Deterministic, lossless round-trip on the supported charset; any other
# Unicode character folds into its byte-fallback bucket (documented lossy
# fold for exotic input, mirrors the "byte fallback" philosophy of the big
# sibling project).
# ---------------------------------------------------------------------------
CHARSET = "abcdefghijklmnopqrstuvwxyz ,:-:'"
PAD_ID, BOS_ID, SEP_ID = 0, 1, 2
_TEXT_VOCAB = 64         # total text-side embedding width
CAP_LEN = 32             # fixed padded caption slot count (incl. <pad>)
PREFIX_LEN = 8           # intent-encoder output prefix vectors fed to decoder


def encode_char(c: str) -> int:
    if c == " ": return 3
    if c == ",": return 4
    if c == "-": return 5
    if c == ":": return 6
    if "a" <= c <= "z":
        return 7 + ord(c) - ord("a")
    if c == "'": return 39
    # byte fallback bucket for anything else (digits, unicode, etc.)
    fb = c.encode("utf-8", "replace")[0]        # first UTF-8 byte as int
    return 33 + (fb % 30)


def decode_char(i: int):
    if i == 3: return " "
    if i == 4: return ","
    if i == 5: return "-"
    if i == 6: return ":"
    if 7 <= i <= 32: return chr(ord("a") + i - 7)
    if i == 39: return "'"
    if 33 <= i <= 38 or 40 <= i <= 62: return "?"   # fallback bucket: unknown glyph
    return None                                      # control specials


def encode_caption(text: str):
    """text -> list of CAP_LEN ids, left-padded slot layout:
    [<pad>...]<chars>. No <bos>/<sep> needed: the prefix encoder consumes the
    whole 32-slot window positionally."""
    ids = [encode_char(c) for c in text.lower()][:CAP_LEN]
    return [PAD_ID] * (CAP_LEN - len(ids)) + ids


def decode_caption(ids) -> str:
    out = []
    for i in ids:
        c = decode_char(int(i))
        if c is not None:
            out.append(c)
    return "".join(out).strip()


# ---------------------------------------------------------------------------
# Generation model geometry (extremely tiny on purpose).
# ---------------------------------------------------------------------------
D_MODEL = 96
N_LAYERS = 4
N_HEADS = 4              # 24-dim heads
FF_HIDDEN = 256          # ffn_mult ~= 2.7 (lean by design)
SEQ_VIS = IMG * IMG      # 256 autoregressive visual positions
POS_TOTAL = PREFIX_LEN + SEQ_VIS   # 264 learned positions

SAMPLING = {"temperature": 0.85, "top_k": 10}

APPROX_PARAMS = 483_040  # exact: encoder + 4-block decoder + heads