Upload 17 files

.gitattributes

@@ -1,35 +1,2 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,2 @@
+__pycache__/
+*.pyc

EVAL_REPRODUCTION.md

@@ -0,0 +1,98 @@
+# Reproducing the eval.md numbers
+
+Step-by-step guide to reproduce the FID / CLIP Score numbers in `eval.md` from
+scratch. Unlike v0 (whose eval scripts lived in a throwaway env), everything
+needed is **in this repo**: `fetch_coco_subset.py` builds the data,
+`eval/run_eval.py` computes the metrics with torchmetrics.
+
+CPU-only is fine (that's how the published numbers were produced). Budget
+~7 GB disk: ~5 GB dataset shards + ~2 GB python deps and model weights
+(InceptionV3, CLIP ViT-B/32).
+
+## 1. Environment
+
+```bash
+python3.13 -m venv eval-venv
+# Windows: eval-venv\Scripts\activate    macOS/Linux: source eval-venv/bin/activate
+
+pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
+pip install torchmetrics torch-fidelity transformers "datasets<3" safetensors pillow numpy
+```
+
+`torch-fidelity` is the backend torchmetrics uses for the FID InceptionV3
+features; `transformers` backs the CLIPScore metric.
+
+## 2. Build the data (train + eval split)
+
+```bash
+python fetch_coco_subset.py --out ../pm-work
+```
+
+This streams [`sayakpaul/coco-30-val-2014`](https://huggingface.co/datasets/sayakpaul/coco-30-val-2014)
+(a public 30K-pair subset of MS-COCO val2014) and produces:
+
+- `eval_real/` — stream rows 0–4999: real images, 256×256 center-crop JPEG (FID real set)
+- `eval_captions.json` — the 5,000 captions for those rows (generation + CLIP Score)
+- `coco_train.npz` + `coco_train.captions.json` — training pairs from stream
+  rows 10,000+, at 64×64
+
+**Leakage control:** the same COCO image appears in multiple rows with
+different captions, so training rows are dropped if their image content hash
+(md5 of a 64×64 bilinear thumbnail) appears in the eval set, and deduped
+against each other. The counts printed at the end of the script report both.
+
+Streaming row-by-row from the Hub is slow (~hours on a normal connection). The
+much faster path is to download the 10 parquet shards first (~4.9 GB), then:
+
+```bash
+python fetch_coco_subset.py --out ../pm-work --parquet-dir ../pm-work/shards
+```
+
+Sorted shard filename order equals hub streaming order, so both paths produce
+identical splits.
+
+## 3. (Optional) retrain
+
+```bash
+python train.py --data ../pm-work/coco_train.npz --epochs 50
+```
+
+Deterministic given `--seed` (default 0): same data + same seed reproduces the
+published `model.png` up to floating-point ordering. Skip this step to
+evaluate the shipped `model.png` as-is.
+
+## 4. Run the eval
+
+```bash
+python eval/run_eval.py --work ../pm-work --model model.png --n 5000
+```
+
+This:
+1. Generates one 64×64 image per eval caption (native resolution — no
+   upscaling before metrics).
+2. **FID** via `torchmetrics.image.fid.FrechetInceptionDistance`
+   (`feature=2048`, InceptionV3 pool3): real = 5,000 COCO val2014 images at
+   256×256 center-crop, fake = the 5,000 generated 64×64 images. The metric's
+   Inception backbone resizes both sets to 299×299 internally.
+3. **CLIP Score** via `torchmetrics.multimodal.CLIPScore` with
+   `openai/clip-vit-base-patch32` (the leaderboard default): each generated
+   image against the caption that produced it, averaged.
+
+Results are printed and written to `eval_results.json` in the `--work` dir.
+First run downloads InceptionV3 (~95 MB) and CLIP ViT-B/32 (~600 MB).
+
+## Notes
+
+- n=5000 vs the standard n=30000: FID's 2048-dim covariance estimate benefits
+  from more samples; n=5000 was chosen as a compromise that is far past the
+  singular-covariance regime that made v0's n=40 number directional-only,
+  while staying CPU-friendly. Raising `--n` (and `EVAL_N` in
+  `fetch_coco_subset.py`) toward 30K tightens the estimate at proportional
+  cost — but note this dataset has 30K *rows*, not 30K unique images, and
+  rows past index 10,000 are used for training, so a larger eval set requires
+  re-splitting.
+- The generated set's resolution (64×64) is reported as the native generation
+  resolution per the leaderboard rules; FID/CLIP preprocessing upsamples
+  internally as part of the metric, not as part of generation.
+- `INFERENCE.py` (safetensors) and `main.py` (model.png) are verified
+  byte-identical, so either weights file reproduces the same metrics.

INFERENCE.py

@@ -0,0 +1,104 @@
+"""
+INFERENCE.py - run PixelModel v2 from model.safetensors, self-contained.
+
+Only needs torch + safetensors + pillow + numpy; no other file from this
+repo. Produces byte-identical output to main.py (which loads model.png)
+for the same prompt and resolution.
+
+Usage:
+  python INFERENCE.py "a red double decker bus"
+  python INFERENCE.py "a cat on a couch" --model model.safetensors --out cat.png --res 64 --scale 4
+"""
+
+import argparse
+import os
+import sys
+
+import numpy as np
+import torch
+from PIL import Image
+from safetensors.torch import load_file
+
+EMB_DIM    = 64
+NATIVE_RES = 64
+FREQS      = (1.0, 2.0, 4.0, 8.0)
+
+
+def _fnv1a(data: bytes) -> int:
+    h = 0x811C9DC5
+    for byte in data:
+        h ^= byte
+        h = (h * 0x01000193) & 0xFFFFFFFF
+    return h
+
+
+def prompt_to_embedding(prompt: str) -> torch.Tensor:
+    text = "".join(c if c.isalnum() or c == " " else " " for c in prompt.lower())
+    text = " ".join(text.split())
+    vec = np.zeros(EMB_DIM, dtype=np.float32)
+    padded = f" {text} "
+    for i in range(len(padded) - 2):
+        h = _fnv1a(padded[i:i + 3].encode("utf-8"))
+        vec[h % EMB_DIM] += 1.0 if (h >> 16) & 1 else -1.0
+    for word in text.split():
+        h = _fnv1a(b"w:" + word.encode("utf-8"))
+        vec[h % EMB_DIM] += 2.0 if (h >> 16) & 1 else -2.0
+    norm = np.linalg.norm(vec)
+    if norm > 0:
+        vec /= norm
+    return torch.from_numpy(vec)
+
+
+def coord_features(res: int) -> torch.Tensor:
+    axis = torch.linspace(-1.0, 1.0, res)
+    yy, xx = torch.meshgrid(axis, axis, indexing="ij")
+    x, y = xx.reshape(-1), yy.reshape(-1)
+    feats = [x, y]
+    for f in FREQS:
+        feats += [torch.sin(f * torch.pi * x), torch.cos(f * torch.pi * x),
+                  torch.sin(f * torch.pi * y), torch.cos(f * torch.pi * y)]
+    return torch.stack(feats, dim=1)
+
+
+def forward(w: dict, prompt: str, res: int) -> torch.Tensor:
+    emb = prompt_to_embedding(prompt).unsqueeze(0)
+    z = torch.tanh(emb @ w["T1"].T + w["b1"])
+    z = torch.tanh(z @ w["T2"].T + w["b2"])
+    feats = coord_features(res)
+    P = feats.shape[0]
+    inp = torch.cat([z.expand(P, -1), feats], dim=1)
+    h = torch.tanh(inp @ w["D1"].T + w["bd1"])
+    h = torch.tanh(h @ w["D2"].T + w["bd2"])
+    rgb = torch.sigmoid(h @ w["D3"].T + w["bd3"])
+    return rgb.reshape(res, res, 3)
+
+
+def main():
+    p = argparse.ArgumentParser(description="PixelModel v2 inference (safetensors)")
+    p.add_argument("prompt")
+    p.add_argument("--model", default="model.safetensors")
+    p.add_argument("--out",   default="out.png")
+    p.add_argument("--res",   type=int, default=NATIVE_RES)
+    p.add_argument("--scale", type=int, default=4)
+    args = p.parse_args()
+
+    if not os.path.exists(args.model):
+        sys.exit(f"Model not found: {args.model}\n"
+                 f"Run: python convert_to_safetensors.py   to create it from model.png.")
+
+    weights = load_file(args.model)
+    with torch.no_grad():
+        result = forward(weights, args.prompt, args.res)
+
+    arr = (result.numpy() * 255).clip(0, 255).astype(np.uint8)
+    img = Image.fromarray(arr, mode="RGB")
+    if args.scale > 1:
+        img = img.resize((args.res * args.scale,) * 2, Image.NEAREST)
+    img.save(args.out)
+    print(f"prompt : '{args.prompt}'")
+    print(f"model  : {args.model}  ({os.path.getsize(args.model)} bytes, safetensors)")
+    print(f"output : {args.out}  ({args.res}x{args.res} native, x{args.scale} view)")
+
+
+if __name__ == "__main__":
+    main()

README.md

@@ -1,3 +1,210 @@
----
-license: mit
----
+---
+license: mit
+pipeline_tag: text-to-image
+language:
+- en
+tags:
+- image
+- t2i
+- text-to-image
+- custom-code
+- tiny
+model-index:
+  - name: PixelModel-v1
+    results:
+      - task:
+          type: text-to-image
+        dataset:
+          name: fid
+          type: fid
+        metrics:
+          - name: fid
+            type: fid
+            value: 439.46
+      - task:
+          type: text-to-image
+        dataset:
+          name: clip
+          type: clip
+        metrics:
+          - name: clip
+            type: clip
+            value: 0.2002
+---
+
+<table>
+<tr>
+<td>
+<br>
+  
+# See us on the [Leaderboard ⇗](https://huggingface.co/spaces/FlameF0X/Tiny-T2I-Leaderboard)
+  
+</td>
+</tr>
+</table>
+
+<table>
+<tr>
+<td>
+<br>
+  
+# Blog post [read it here ⇗](https://huggingface.co/spaces/bench-labs/blog?post=pixelmodel-v1.html)
+  
+</td>
+</tr>
+</table>
+
+# PixelModel v1 🖼️
+
+<img src="model.png" alt="The entire PixelModel v1: a 160x149 PNG containing all 23,747 weights" style="image-rendering: pixelated; width: 320px; max-width: 100%;">
+
+
+A neural network where the weights **are** the image. Now the model is a thumbnail.
+
+## 📌 What is this?
+
+`model.png` is not a picture — it *is* the model.
+
+Every pixel encodes neural network weights. At inference, the PNG is decoded
+into weight matrices, the prompt is hashed into an embedding, and a
+coordinate-conditioned decoder paints an image — at **any resolution**.
+
+[v0](https://huggingface.co/bench-labs/pixelmodel) was a 202,752-parameter MLP
+welded to 32×32 output, trained on 6 color swatches. v1 is **8.5× smaller
+(23,747 parameters)**, renders at any resolution (native 64×64), and is trained
+on ~20K real MS-COCO caption/image pairs.
+
+# cherry on top 🍒
+The model generates 600 images (cpu) in 5 (five) seconds.
+Thats 5000 images in 24 seconds on cpu.
+The model trained on cpu for just 30 minutes.
+
+## 🆕 v0 → v1
+
+| | v0 | v1 |
+|---|---|---|
+| Parameters | 202,752 | **23,747** |
+| `model.png` | 64×3200 px | **160×149 px** (a thumbnail) |
+| Output head | 1 weight row per pixel (196K params, 97% of model) | CPPN decoder on (x, y) — resolution-free |
+| Resolution | fixed 32×32 | any; native 64×64 |
+| Prompt embedding | char-sum (order-blind, collision-heavy) | hashed trigrams + words (FNV-1a), still 0 params |
+| Biases | none | yes |
+| Weight precision | 8-bit (G channel wasted) | **16-bit** (R=high byte, G=low byte) |
+| Training data | 6 solid-color swatches | ~20K MS-COCO caption/image pairs |
+
+<img src="pixelmodel-v1-loss.png" alt="Training loss curve crossing below both naive-predictor baselines" style="width: 100%; max-width: 720px;">
+
+## 🎨 Weight Encoding
+
+Each pixel stores one weight at 16-bit precision, mapped from [-2, 2]:
+
+- **R channel** → high byte
+- **G channel** → low byte
+- **B channel** → reserved
+
+Round-trip quantization error ≈ 3×10⁻⁵ per weight.
+
+## 🧠 Architecture
+
+```text
+prompt string
+  → hashed char-trigram + word embedding (64-dim, deterministic, 0 params)
+  → T1 (80×64)+b → tanh
+  → T2 (64×80)+b → tanh          = latent z (64)
+
+for every pixel (x, y):
+  concat(z, fourier features of (x, y))   ← 18 coord dims, freqs 1/2/4/8
+  → D1 (80×82)+b → tanh
+  → D2 (80×80)+b → tanh
+  → D3 (3×80)+b  → sigmoid  = RGB
+```
+
+Because pixels are decoded from coordinates, parameter count is independent of
+resolution — `--res 256` works with the same 23,747 weights.
+
+All weights live inside `model.png` (160×149 px).
+
+## 📦 Standard weights (safetensors)
+
+`model.png` is the canonical model — training writes to it directly. For
+tooling that expects standard weight files, the same 10 tensors are exported as
+`model.safetensors` (23,747 parameters total):
+
+```bash
+python convert_to_safetensors.py            # model.png -> model.safetensors
+```
+
+Parameter count is verifiable without running code: `config.json`
+(`total_parameters: 23747`, full per-layer breakdown) and the safetensors
+header metadata (`total_parameters`, `param_breakdown`, `has_bias`,
+`text_encoder_parameters: 0`, `vae_parameters: 0`).
+
+## 📊 Benchmark results
+
+Measured per the [Tiny-T2I-Leaderboard](https://huggingface.co/spaces/FlameF0X/Tiny-T2I-Leaderboard) protocol — see `eval.md` for full method and `EVAL_REPRODUCTION.md` to re-run:
+
+| Metric | v0 | v1 | Tooling |
+|---|---|---|---|
+| FID ↓ | 566.84 (n=40) | **439.46** (n=5000) | `torchmetrics.image.fid.FrechetInceptionDistance` |
+| CLIP Score ↑ | 18.60 (n=40) | **20.02** (n=5000) | `torchmetrics.multimodal.CLIPScore`, `openai/clip-vit-base-patch32` |
+| Native resolution | 32×32 | 64×64 | |
+
+Real set: MS-COCO val2014 pairs from `sayakpaul/coco-30-val-2014` (256×256
+center-crop). Training images are hash-checked to be disjoint from the eval
+images.
+
+Expectation management: a 23K-parameter model does not draw recognizable
+objects. It learns caption-conditioned color, layout, and texture statistics of
+COCO photos — the outputs are blurry impressions, not pictures. The point is
+the size:quality ratio and the PNG-as-model gimmick, executed honestly.
+
+<img src="pixelmodel-v1-probe.png" alt="Target photos vs model outputs: the model produces caption-conditioned color washes" style="width: 100%; max-width: 640px;">
+
+## ⚙️ Usage
+
+```bash
+# training data (COCO subset)
+python fetch_coco_subset.py --out ../pm-work
+
+# train (writes model.png every epoch)
+python train.py --data ../pm-work/coco_train.npz --epochs 50
+
+# inference from model.png (canonical)
+python main.py "a red double decker bus" --out bus.png
+
+# inference from model.safetensors (standalone, needs only this one file + torch)
+python convert_to_safetensors.py
+python INFERENCE.py "a red double decker bus" --out bus.png
+
+# any resolution from the same 23,747 weights
+python main.py "a beach with palm trees" --res 256 --scale 1
+
+# benchmark eval
+python eval/run_eval.py --work ../pm-work --model model.png --n 5000
+```
+
+`main.py` and `INFERENCE.py` produce identical output for the same prompt and
+resolution.
+
+## 📁 Files
+
+```text
+model.png                   ← THE MODEL (160×149 px)
+model.safetensors           ← same weights, standard format + param metadata
+config.json                 ← architecture + parameter-count metadata
+main.py                     ← inference, loads model.png
+INFERENCE.py                ← inference, loads model.safetensors (standalone)
+convert_to_safetensors.py   ← model.png -> model.safetensors
+train.py                    ← training
+model.py                    ← architecture + PNG weight codec
+fetch_coco_subset.py        ← builds train/eval data from COCO
+eval/run_eval.py            ← FID + CLIP Score via torchmetrics
+eval.md                     ← benchmark results + method
+EVAL_REPRODUCTION.md        ← step-by-step reproduction guide
+```
+
+---
+
+*Still a toy. Slightly less useless. The model is literally a thumbnail.*
+
+Bench Labs · Simple, Reliable, Open sourced

config.json

@@ -0,0 +1,35 @@
+{
+  "model_type": "pixelmodel-v2",
+  "architecture": "hashed trigram/word embedding -> MLP trunk -> CPPN decoder (coordinate-conditioned, resolution-free)",
+  "weights_file": "model.safetensors",
+  "total_parameters": 200259,
+  "parameter_breakdown": {
+    "text_encoder_parameters": 0,
+    "generative_backbone_parameters": 200259,
+    "vae_parameters": 0
+  },
+  "layers": {
+    "T1":  { "shape": [256, 64],  "role": "prompt embedding -> trunk hidden", "parameters": 16384 },
+    "b1":  { "shape": [256],      "role": "trunk hidden bias",                "parameters": 256 },
+    "T2":  { "shape": [128, 256], "role": "trunk hidden -> latent",           "parameters": 32768 },
+    "b2":  { "shape": [128],      "role": "latent bias",                      "parameters": 128 },
+    "D1":  { "shape": [320, 146], "role": "latent + fourier(x,y) -> decoder hidden", "parameters": 46720 },
+    "bd1": { "shape": [320],      "role": "decoder hidden bias",              "parameters": 320 },
+    "D2":  { "shape": [320, 320], "role": "decoder hidden -> decoder hidden", "parameters": 102400 },
+    "bd2": { "shape": [320],      "role": "decoder hidden bias",              "parameters": 320 },
+    "D3":  { "shape": [3, 320],   "role": "decoder hidden -> RGB",            "parameters": 960 },
+    "bd3": { "shape": [3],      "role": "RGB bias",                         "parameters": 3 }
+  },
+  "has_bias": true,
+  "prompt_embedding": "hashed character trigrams + words (FNV-1a), 64-dim, deterministic, 0 parameters",
+  "embedding_dim": 64,
+  "trunk_hidden_dim": 256,
+  "latent_dim": 128,
+  "decoder_hidden_dim": 320,
+  "coordinate_features": "x, y + sin/cos at frequencies (1, 2, 4, 8) per axis (18 dims)",
+  "native_resolution": "64x64",
+  "resolution_free": true,
+  "output_channels": 3,
+  "model_png": "256x783 px, 16-bit weight codec (R = high byte, G = low byte, B = reserved)",
+  "weight_range": [-2.0, 2.0]
+}

convert_to_safetensors.py

@@ -0,0 +1,48 @@
+"""
+convert_to_safetensors.py - model.png -> model.safetensors.
+
+model.png stays the canonical model; this exports the same weights in the
+standard format, with parameter-count metadata embedded in the safetensors
+header so the count is verifiable without running any code.
+
+Usage:
+  python convert_to_safetensors.py
+  python convert_to_safetensors.py --model model.png --out model.safetensors
+"""
+
+import argparse
+import json
+import os
+
+from safetensors.torch import save_file
+
+from model import MODEL_H, MODEL_W, N_PARAMS, NATIVE_RES, PARAM_SPECS, load_model
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--model", default="model.png")
+    p.add_argument("--out",   default="model.safetensors")
+    args = p.parse_args()
+
+    weights = load_model(args.model)
+    counts = {name: int(weights[name].numel()) for name, _ in PARAM_SPECS}
+    assert sum(counts.values()) == N_PARAMS
+
+    metadata = {
+        "model_type": "pixelmodel-v2",
+        "total_parameters": str(N_PARAMS),
+        "param_breakdown": json.dumps(counts),
+        "text_encoder_parameters": "0",
+        "vae_parameters": "0",
+        "has_bias": "true",
+        "native_resolution": f"{NATIVE_RES}x{NATIVE_RES}",
+        "source_png": f"{MODEL_W}x{MODEL_H} px, 16-bit codec (R=high byte, G=low byte)",
+    }
+    save_file(weights, args.out, metadata=metadata)
+    print(f"{args.model} -> {args.out}  ({os.path.getsize(args.out)} bytes, "
+          f"{N_PARAMS} parameters)")
+
+
+if __name__ == "__main__":
+    main()

fetch_coco_subset.py

@@ -0,0 +1,124 @@
+"""
+fetch_coco_subset.py - build train and eval sets from sayakpaul/coco-30-val-2014.
+
+The stream of 30K (image, caption) rows is split by row index:
+  rows 0..EVAL_N-1        -> eval set (real images for FID + captions for generation)
+  rows TRAIN_START..end   -> training candidates
+
+Training images whose content hash also appears in the eval set are dropped,
+so the model is never trained on an image it is evaluated against (the same
+image can appear in multiple rows with different captions). Training rows are
+also deduped against each other by image hash.
+
+Outputs (in --out dir):
+  coco_train.npz               images uint8 (N, 64, 64, 3), center-cropped
+  coco_train.captions.json     N caption strings
+  eval_real/real_00000.jpg...  EVAL_N real images, 256x256 center-crop, q95
+  eval_captions.json           EVAL_N caption strings (for generation + CLIP)
+
+Usage:
+  python fetch_coco_subset.py --out ../pm-work
+  # much faster: pre-download the 10 parquet shards, then
+  python fetch_coco_subset.py --out ../pm-work --parquet-dir ../pm-work/shards
+"""
+
+import argparse
+import glob
+import hashlib
+import json
+import os
+
+import numpy as np
+from datasets import load_dataset
+from PIL import Image
+
+EVAL_N      = 5000
+TRAIN_START = 10000     # comfortable gap after the eval rows
+TRAIN_RES   = 64
+EVAL_RES    = 256
+
+
+def center_crop_resize(img: Image.Image, size: int) -> Image.Image:
+    w, h = img.size
+    side = min(w, h)
+    left, top = (w - side) // 2, (h - side) // 2
+    return img.crop((left, top, left + side, top + side)).resize(
+        (size, size), Image.BICUBIC
+    )
+
+
+def img_hash(img: Image.Image) -> str:
+    """Content hash of the image, robust to caption-duplicated rows."""
+    return hashlib.md5(
+        img.convert("RGB").resize((64, 64), Image.BILINEAR).tobytes()
+    ).hexdigest()
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--out", required=True)
+    p.add_argument("--parquet-dir", default=None,
+                   help="dir with the dataset's parquet shards already downloaded "
+                        "(sorted filename order == hub streaming row order)")
+    args = p.parse_args()
+
+    real_dir = os.path.join(args.out, "eval_real")
+    os.makedirs(real_dir, exist_ok=True)
+
+    if args.parquet_dir:
+        files = sorted(glob.glob(os.path.join(args.parquet_dir, "*.parquet")))
+        assert files, f"no parquet files in {args.parquet_dir}"
+        ds = load_dataset("parquet", data_files=files, split="train", streaming=True)
+    else:
+        ds = load_dataset("sayakpaul/coco-30-val-2014", split="train", streaming=True)
+
+    eval_hashes = set()
+    eval_captions = []
+    train_images, train_captions = [], []
+    train_hashes = set()
+    skipped_eval_overlap = skipped_dupe = 0
+
+    for i, row in enumerate(ds):
+        img, caption = row["image"], row["caption"]
+
+        if i < EVAL_N:
+            eval_hashes.add(img_hash(img))
+            eval_captions.append(caption)
+            center_crop_resize(img.convert("RGB"), EVAL_RES).save(
+                os.path.join(real_dir, f"real_{i:05d}.jpg"), quality=95
+            )
+            if (i + 1) % 1000 == 0:
+                print(f"  eval: {i + 1}/{EVAL_N}", flush=True)
+            continue
+
+        if i < TRAIN_START:
+            continue
+
+        h = img_hash(img)
+        if h in eval_hashes:
+            skipped_eval_overlap += 1
+            continue
+        if h in train_hashes:
+            skipped_dupe += 1
+            continue
+        train_hashes.add(h)
+        thumb = center_crop_resize(img.convert("RGB"), TRAIN_RES)
+        train_images.append(np.array(thumb, dtype=np.uint8))
+        train_captions.append(caption)
+        if len(train_images) % 2000 == 0:
+            print(f"  train: {len(train_images)} (row {i})", flush=True)
+
+    images = np.stack(train_images)
+    np.savez_compressed(os.path.join(args.out, "coco_train.npz"), images=images)
+    with open(os.path.join(args.out, "coco_train.captions.json"), "w", encoding="utf-8") as f:
+        json.dump(train_captions, f)
+    with open(os.path.join(args.out, "eval_captions.json"), "w", encoding="utf-8") as f:
+        json.dump(eval_captions, f)
+
+    print(f"\neval:  {len(eval_captions)} pairs @ {EVAL_RES}x{EVAL_RES} -> {real_dir}")
+    print(f"train: {len(train_captions)} pairs @ {TRAIN_RES}x{TRAIN_RES} "
+          f"(skipped: {skipped_eval_overlap} eval-overlap, {skipped_dupe} dupes)")
+
+
+if __name__ == "__main__":
+    main()

main.py

@@ -0,0 +1,43 @@
+"""
+main.py - PixelModel v1 inference from model.png (the canonical model).
+
+Usage:
+  python main.py "a red double decker bus"
+  python main.py "a cat on a couch" --out cat.png --res 64 --scale 4
+"""
+
+import argparse
+
+import numpy as np
+import torch
+from PIL import Image
+
+from model import NATIVE_RES, forward, load_model
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("prompt")
+    p.add_argument("--model", default="model.png")
+    p.add_argument("--out",   default="out.png")
+    p.add_argument("--res",   type=int, default=NATIVE_RES,
+                   help="generation resolution (decoder is resolution-free)")
+    p.add_argument("--scale", type=int, default=4, help="nearest-neighbor upscale for viewing")
+    args = p.parse_args()
+
+    weights = load_model(args.model)
+    with torch.no_grad():
+        result = forward(weights, args.prompt, res=args.res)[0]
+
+    arr = (result.numpy() * 255).clip(0, 255).astype(np.uint8)
+    img = Image.fromarray(arr, mode="RGB")
+    if args.scale > 1:
+        img = img.resize((args.res * args.scale,) * 2, Image.NEAREST)
+    img.save(args.out)
+    print(f"prompt : '{args.prompt}'")
+    print(f"model  : {args.model}")
+    print(f"output : {args.out}  ({args.res}x{args.res} native, x{args.scale} view)")
+
+
+if __name__ == "__main__":
+    main()

model.py

@@ -0,0 +1,210 @@
+"""
+PixelModel v2 - the weights ARE the image.
+
+model.png stores every parameter of the network as pixel values.
+v2 changes vs v1:
+  - Wider trunk and decoder: 200,259 parameters rather than 23,747, while
+    keeping the same prompt encoder, coordinate features, 64x64 native
+    resolution, and 16-bit PNG codec.
+  - The larger capacity is intended for the same ~20K COCO training rows.
+
+v1 retained these changes from v0:
+  - Coordinate-conditioned decoder (CPPN-style): parameters no longer scale
+    with output resolution. v0 spent 196,608 of its 202,752 params on the
+    output head; v1 has 23,747 params total and renders at ANY resolution
+    (native 64x64 for eval).
+  - Hashed character-trigram + word prompt embedding (deterministic, 0 params)
+    instead of v0's char-sum, so different prompts actually get different,
+    partially compositional embeddings.
+  - Biases everywhere (v0 had none).
+  - 16-bit weight codec: R = high byte, G = low byte, B = reserved.
+    v0 wasted G and stored weights at 8 bits; v1's quantization error is
+    ~6e-5 per weight. model.png is 160x149 px - the model is a thumbnail.
+
+Architecture:
+  prompt -> hashed trigram/word embedding (64)
+    -> T1 (256x64)+b tanh -> T2 (128x256)+b tanh -> latent z (128)
+  per pixel: concat(z, fourier(x,y) (18))
+    -> D1 (320x146)+b tanh -> D2 (320x320)+b tanh -> D3 (3x320)+b sigmoid -> RGB
+"""
+
+import numpy as np
+import torch
+from PIL import Image
+
+# -- config -------------------------------------------------------------------
+EMB_DIM      = 64     # prompt embedding size (hashed, no learned params)
+TRUNK_HIDDEN = 256
+LATENT       = 128
+COORD_FEATS  = 18     # x, y + sin/cos at freqs (1,2,4,8) for each axis
+DEC_HIDDEN   = 320
+NATIVE_RES   = 64     # native eval resolution (decoder works at any res)
+
+FREQS = (1.0, 2.0, 4.0, 8.0)
+
+# every parameter tensor, in canonical flattening order
+PARAM_SPECS = [
+    ("T1",  (TRUNK_HIDDEN, EMB_DIM)),
+    ("b1",  (TRUNK_HIDDEN,)),
+    ("T2",  (LATENT, TRUNK_HIDDEN)),
+    ("b2",  (LATENT,)),
+    ("D1",  (DEC_HIDDEN, LATENT + COORD_FEATS)),
+    ("bd1", (DEC_HIDDEN,)),
+    ("D2",  (DEC_HIDDEN, DEC_HIDDEN)),
+    ("bd2", (DEC_HIDDEN,)),
+    ("D3",  (3, DEC_HIDDEN)),
+    ("bd3", (3,)),
+]
+N_PARAMS = sum(int(np.prod(s)) for _, s in PARAM_SPECS)   # 200,259
+
+# model.png geometry: one weight per pixel, 16-bit (R=high, G=low, B=reserved)
+MODEL_W = 256
+MODEL_H = -(-N_PARAMS // MODEL_W)                          # 783
+WMAX    = 2.0                                              # weights live in [-2, 2]
+
+
+# -- prompt embedding (deterministic, zero parameters) ------------------------
+
+def _fnv1a(data: bytes) -> int:
+    h = 0x811C9DC5
+    for byte in data:
+        h ^= byte
+        h = (h * 0x01000193) & 0xFFFFFFFF
+    return h
+
+
+def prompt_to_embedding(prompt: str) -> torch.Tensor:
+    """Hashed bag of character trigrams + words -> EMB_DIM vector, L2-normed."""
+    text = "".join(c if c.isalnum() or c == " " else " " for c in prompt.lower())
+    text = " ".join(text.split())
+    vec = np.zeros(EMB_DIM, dtype=np.float32)
+
+    padded = f" {text} "
+    for i in range(len(padded) - 2):
+        h = _fnv1a(padded[i:i + 3].encode("utf-8"))
+        sign = 1.0 if (h >> 16) & 1 else -1.0
+        vec[h % EMB_DIM] += sign
+
+    for word in text.split():
+        h = _fnv1a(b"w:" + word.encode("utf-8"))
+        sign = 1.0 if (h >> 16) & 1 else -1.0
+        vec[h % EMB_DIM] += 2.0 * sign      # words weigh more than trigrams
+
+    norm = np.linalg.norm(vec)
+    if norm > 0:
+        vec /= norm
+    return torch.from_numpy(vec)
+
+
+def prompts_to_embeddings(prompts) -> torch.Tensor:
+    return torch.stack([prompt_to_embedding(p) for p in prompts])
+
+
+# -- coordinate features ------------------------------------------------------
+
+_coord_cache = {}
+
+
+def coord_features(res: int) -> torch.Tensor:
+    """(res*res, COORD_FEATS) fourier features of the pixel grid, cached."""
+    if res not in _coord_cache:
+        axis = torch.linspace(-1.0, 1.0, res)
+        yy, xx = torch.meshgrid(axis, axis, indexing="ij")
+        x = xx.reshape(-1)
+        y = yy.reshape(-1)
+        feats = [x, y]
+        for f in FREQS:
+            feats += [torch.sin(f * torch.pi * x), torch.cos(f * torch.pi * x),
+                      torch.sin(f * torch.pi * y), torch.cos(f * torch.pi * y)]
+        _coord_cache[res] = torch.stack(feats, dim=1)     # (res*res, 18)
+    return _coord_cache[res]
+
+
+# -- forward ------------------------------------------------------------------
+
+def encode_prompt(weights: dict, emb: torch.Tensor) -> torch.Tensor:
+    """emb (B, EMB_DIM) -> latent z (B, LATENT)."""
+    z = torch.tanh(emb @ weights["T1"].T + weights["b1"])
+    z = torch.tanh(z @ weights["T2"].T + weights["b2"])
+    return z
+
+
+def decode_pixels(weights: dict, z: torch.Tensor, feats: torch.Tensor) -> torch.Tensor:
+    """z (B, LATENT), feats (P, COORD_FEATS) -> RGB (B, P, 3) in [0, 1]."""
+    B, P = z.shape[0], feats.shape[0]
+    inp = torch.cat([z.unsqueeze(1).expand(B, P, LATENT),
+                     feats.unsqueeze(0).expand(B, P, COORD_FEATS)], dim=2)
+    h = torch.tanh(inp @ weights["D1"].T + weights["bd1"])
+    h = torch.tanh(h @ weights["D2"].T + weights["bd2"])
+    return torch.sigmoid(h @ weights["D3"].T + weights["bd3"])
+
+
+def forward(weights: dict, prompts, res: int = NATIVE_RES) -> torch.Tensor:
+    """prompts: str or list of str -> (B, res, res, 3) float in [0, 1]."""
+    if isinstance(prompts, str):
+        prompts = [prompts]
+    emb = prompts_to_embeddings(prompts)
+    z = encode_prompt(weights, emb)
+    rgb = decode_pixels(weights, z, coord_features(res))
+    return rgb.reshape(len(prompts), res, res, 3)
+
+
+# -- 16-bit PNG weight codec --------------------------------------------------
+
+def weights_to_pixels(weights: dict) -> np.ndarray:
+    """weight dict -> (MODEL_H, MODEL_W, 3) uint8. R=high byte, G=low, B=0."""
+    flat = torch.cat([weights[n].detach().reshape(-1) for n, _ in PARAM_SPECS])
+    q = ((flat.clamp(-WMAX, WMAX) / WMAX + 1.0) / 2.0 * 65535.0).round()
+    q = q.to(torch.int64).cpu().numpy()
+    q = np.pad(q, (0, MODEL_W * MODEL_H - N_PARAMS))
+    img = np.zeros((MODEL_H * MODEL_W, 3), dtype=np.uint8)
+    img[:, 0] = q >> 8
+    img[:, 1] = q & 0xFF
+    return img.reshape(MODEL_H, MODEL_W, 3)
+
+
+def pixels_to_weights(arr: np.ndarray) -> dict:
+    """(MODEL_H, MODEL_W, 3) uint8 -> weight dict."""
+    flat = arr.reshape(-1, 3).astype(np.int64)
+    q = (flat[:, 0] << 8) | flat[:, 1]
+    vals = (q.astype(np.float32) / 65535.0 * 2.0 - 1.0) * WMAX
+    vals = torch.from_numpy(vals[:N_PARAMS])
+    weights, i = {}, 0
+    for name, shape in PARAM_SPECS:
+        n = int(np.prod(shape))
+        weights[name] = vals[i:i + n].reshape(shape).clone()
+        i += n
+    return weights
+
+
+def save_model(weights: dict, path: str):
+    Image.fromarray(weights_to_pixels(weights), mode="RGB").save(path)
+
+
+def load_model(path: str) -> dict:
+    img = Image.open(path).convert("RGB")
+    arr = np.array(img, dtype=np.uint8)
+    assert arr.shape == (MODEL_H, MODEL_W, 3), \
+        f"expected {MODEL_H}x{MODEL_W} model.png, got {arr.shape[0]}x{arr.shape[1]}"
+    return pixels_to_weights(arr)
+
+
+def init_weights(seed: int = 0) -> dict:
+    g = torch.Generator().manual_seed(seed)
+    weights = {}
+    for name, shape in PARAM_SPECS:
+        if len(shape) == 1:
+            weights[name] = torch.zeros(shape)
+        else:
+            fan_in = shape[1]
+            scale = 1.0 / fan_in ** 0.5
+            if name == "D3":
+                scale *= 0.1        # start near grey, not saturated
+            weights[name] = torch.randn(shape, generator=g) * scale
+    return weights
+
+
+if __name__ == "__main__":
+    print(f"params: {N_PARAMS}   model.png: {MODEL_W}x{MODEL_H} px")
+    for name, shape in PARAM_SPECS:
+        print(f"  {name:>4}  {str(tuple(shape)):>12}  {int(np.prod(shape)):>6}")

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:26b6f43f85a826fcf0c7f1c013bb391e9a73169ab2fe18cb0a96bd62664a0a43
+size 96028

pixelmodel_v2_kaggle.ipynb

@@ -0,0 +1,103 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# PixelModel v2 — Kaggle training and evaluation\n",
+    "\n",
+    "This notebook trains the 200,259-parameter v2 model on the same approximately 20K disjoint COCO caption/image pairs as v1, then evaluates it on the held-out 5K COCO rows using FID and CLIP Score.\n",
+    "\n",
+    "Before running: enable **Internet** and **GPU** in Kaggle. Add this repository as a Kaggle Dataset named `pixelmodel-v2`, so it is mounted at `/kaggle/input/pixelmodel-v2`. The notebook copies it to the writable working directory; the trained `model.png`, `model.safetensors`, and evaluation results will be in `/kaggle/working/pixelmodel-v2-output`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!pip -q install datasets safetensors 'torchmetrics[image]'\n",
+    "!pip -q install --upgrade 'torch-fidelity'\n",
+    "!nvidia-smi"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import shutil\n",
+    "from pathlib import Path\n",
+    "\n",
+    "SOURCE_DIR = Path('/kaggle/input/pixelmodel-v2')\n",
+    "PROJECT_DIR = Path('/kaggle/working/pixelmodel-v2')\n",
+    "WORK_DIR = Path('/kaggle/working/pm-work')\n",
+    "OUTPUT_DIR = Path('/kaggle/working/pixelmodel-v2-output')\n",
+    "\n",
+    "assert SOURCE_DIR.exists(), f'Attach the v2 repository as a Kaggle Dataset: {SOURCE_DIR}'\n",
+    "shutil.rmtree(PROJECT_DIR, ignore_errors=True)\n",
+    "shutil.copytree(SOURCE_DIR, PROJECT_DIR)\n",
+    "WORK_DIR.mkdir(parents=True, exist_ok=True)\n",
+    "OUTPUT_DIR.mkdir(parents=True, exist_ok=True)\n",
+    "os.chdir(PROJECT_DIR)\n",
+    "!python model.py\n",
+    "!cat config.json"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Creates the v1-compatible split: ~20K 64x64 training pairs and 5K held-out eval pairs.\n",
+    "# This downloads COCO through Hugging Face, so Kaggle Internet must stay enabled.\n",
+    "!python fetch_coco_subset.py --out {WORK_DIR}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Start fresh: do not pass --resume for the first run.\n",
+    "# The v2 model is larger than v1, so use a GPU and give it a longer run.\n",
+    "!python train.py --data {WORK_DIR / 'coco_train.npz'} --epochs 50 --batch 128 --pixels 1024 --lr 0.002\n",
+    "!python convert_to_safetensors.py --model model.png --out model.safetensors\n",
+    "# Persist the trained model immediately; Kaggle exposes /kaggle/working as notebook output.\n",
+    "!cp model.png model.safetensors config.json {OUTPUT_DIR}/\n",
+    "!ls -lh {OUTPUT_DIR}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Held-out evaluation: generate 5K images, then report FID (lower is better) and CLIP Score (higher is better).\n",
+    "!python eval/run_eval.py --work {WORK_DIR} --model model.png --n 5000\n",
+    "!cp {WORK_DIR / 'eval_results.json'} {OUTPUT_DIR}/\n",
+    "!ls -lh {OUTPUT_DIR}\n",
+    "!cat {WORK_DIR / 'eval_results.json'}"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.x"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

run_eval.py

@@ -0,0 +1,136 @@
+"""
+run_eval.py - benchmark eval: COCO FID + CLIP Score via torchmetrics.
+
+Protocol (matches the Tiny-T2I leaderboard requirements):
+  - FID: torchmetrics.image.fid.FrechetInceptionDistance (InceptionV3,
+    2048-dim pool3 features). Real set: n COCO val2014 images (256x256
+    center-crop) from sayakpaul/coco-30-val-2014, rows 0..n-1 of the
+    stream — disjoint by image hash from the training set (see
+    fetch_coco_subset.py). Generated set: model output at native 64x64
+    for those same n captions.
+  - CLIP Score: torchmetrics.multimodal.CLIPScore with
+    openai/clip-vit-base-patch32 (the default), generated image vs the
+    caption that produced it.
+
+Usage (after fetch_coco_subset.py has populated --work):
+  python eval/run_eval.py --work ../pm-work --model model.png --n 5000
+"""
+
+import argparse
+import json
+import os
+import sys
+import time
+
+import numpy as np
+import torch
+from PIL import Image
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from model import NATIVE_RES, coord_features, decode_pixels, encode_prompt, load_model, prompts_to_embeddings  # noqa: E402
+
+
+def generate(model_path: str, captions, out_dir: str, device: torch.device, batch: int = 64):
+    os.makedirs(out_dir, exist_ok=True)
+    weights = load_model(model_path)
+    weights = {name: value.to(device) for name, value in weights.items()}
+    feats = coord_features(NATIVE_RES).to(device)
+    t0 = time.time()
+    for start in range(0, len(captions), batch):
+        chunk = captions[start:start + batch]
+        with torch.no_grad():
+            emb = prompts_to_embeddings(chunk).to(device)
+            z = encode_prompt(weights, emb)
+            rgb = decode_pixels(weights, z, feats)
+        arr = (rgb.reshape(len(chunk), NATIVE_RES, NATIVE_RES, 3).cpu().numpy()
+               * 255).clip(0, 255).astype(np.uint8)
+        for j in range(len(chunk)):
+            Image.fromarray(arr[j], mode="RGB").save(
+                os.path.join(out_dir, f"gen_{start + j:05d}.png"))
+        if (start // batch) % 20 == 0:
+            print(f"  gen {start + len(chunk)}/{len(captions)}  "
+                  f"({time.time() - t0:.0f}s)", flush=True)
+    print(f"  generated {len(captions)} images @ {NATIVE_RES}x{NATIVE_RES} "
+          f"in {time.time() - t0:.0f}s", flush=True)
+
+
+def load_batch(paths):
+    imgs = [np.array(Image.open(p).convert("RGB"), dtype=np.uint8) for p in paths]
+    return torch.from_numpy(np.stack(imgs)).permute(0, 3, 1, 2)   # (B,3,H,W) uint8
+
+
+def compute_fid(real_dir: str, gen_dir: str, n: int, device: torch.device, batch: int = 32) -> float:
+    from torchmetrics.image.fid import FrechetInceptionDistance
+    fid = FrechetInceptionDistance(feature=2048, normalize=False).to(device)
+    t0 = time.time()
+    for label, dir_, real in (("real", real_dir, True), ("gen", gen_dir, False)):
+        files = sorted(os.listdir(dir_))[:n]
+        for start in range(0, len(files), batch):
+            imgs = load_batch([os.path.join(dir_, f) for f in files[start:start + batch]])
+            fid.update(imgs.to(device), real=real)
+            if (start // batch) % 25 == 0:
+                print(f"  fid/{label}: {start + imgs.shape[0]}/{len(files)}  "
+                      f"({time.time() - t0:.0f}s)", flush=True)
+    return float(fid.compute())
+
+
+def compute_clip_score(gen_dir: str, captions, device: torch.device, batch: int = 32):
+    from torchmetrics.multimodal import CLIPScore
+    metric = CLIPScore(model_name_or_path="openai/clip-vit-base-patch32").to(device)
+    files = sorted(os.listdir(gen_dir))[:len(captions)]
+    t0 = time.time()
+    for start in range(0, len(files), batch):
+        imgs = load_batch([os.path.join(gen_dir, f) for f in files[start:start + batch]])
+        metric.update(imgs.to(device), captions[start:start + imgs.shape[0]])
+        if (start // batch) % 25 == 0:
+            print(f"  clip: {start + imgs.shape[0]}/{len(files)}  "
+                  f"({time.time() - t0:.0f}s)", flush=True)
+    return float(metric.compute())
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--work",  required=True, help="dir from fetch_coco_subset.py")
+    p.add_argument("--model", default="model.png")
+    p.add_argument("--n",     type=int, default=5000)
+    p.add_argument("--device", default="auto", help="auto, cpu, cuda, or a PyTorch device string")
+    p.add_argument("--skip-gen", action="store_true")
+    p.add_argument("--skip-fid", action="store_true")
+    args = p.parse_args()
+    device = torch.device("cuda" if args.device == "auto" and torch.cuda.is_available()
+                          else "cpu" if args.device == "auto" else args.device)
+    print(f"device: {device}")
+
+    with open(os.path.join(args.work, "eval_captions.json"), encoding="utf-8") as f:
+        captions = json.load(f)[:args.n]
+    real_dir = os.path.join(args.work, "eval_real")
+    gen_dir = os.path.join(args.work, "eval_gen")
+
+    if not args.skip_gen:
+        print(f"[1/3] generating {len(captions)} images from '{args.model}'...")
+        generate(args.model, captions, gen_dir, device)
+    fid = None
+    if not args.skip_fid:
+        print("[2/3] FID (torchmetrics.image.fid, InceptionV3 2048)...")
+        fid = compute_fid(real_dir, gen_dir, args.n, device)
+        print(f"FID = {fid:.4f}", flush=True)
+    print("[3/3] CLIP Score (torchmetrics, openai/clip-vit-base-patch32)...")
+    clip = compute_clip_score(gen_dir, captions, device)
+    print(f"CLIP Score = {clip:.4f}  (cosine {clip / 100:.4f})")
+
+    print(f"\nRESULTS  n={args.n}  native_res={NATIVE_RES}x{NATIVE_RES}")
+    if fid is not None:
+        print(f"  FID        = {fid:.2f}")
+    print(f"  CLIP Score = {clip:.2f}")
+    out_path = os.path.join(args.work, "eval_results.json")
+    results = {"n": args.n, "native_resolution": f"{NATIVE_RES}x{NATIVE_RES}",
+               "fid": fid, "clip_score": clip}
+    if fid is None and os.path.exists(out_path):
+        old = json.load(open(out_path))
+        results["fid"] = old.get("fid")
+    with open(out_path, "w") as f:
+        json.dump(results, f, indent=2)
+
+
+if __name__ == "__main__":
+    main()

train.py

@@ -0,0 +1,121 @@
+"""
+train.py - train PixelModel v1 on caption/image pairs, save into model.png.
+
+Expects the npz produced by fetch_coco_subset.py:
+  images   uint8 (N, 64, 64, 3)
+  captions json list of N strings (stored alongside as captions.json)
+
+Each step samples a batch of captions and a random subset of pixel
+coordinates (the CPPN decoder makes per-pixel training natural), computes
+MSE against the target pixels, and Adam-steps the weights. Weights are
+clamped to [-WMAX, WMAX] so they always round-trip through the 16-bit
+PNG codec. model.png is written every epoch.
+
+Usage:
+  python train.py --data ../pm-work/coco_train.npz
+  python train.py --data ../pm-work/coco_train.npz --epochs 40 --lr 2e-3
+"""
+
+import argparse
+import json
+import os
+import time
+
+import numpy as np
+import torch
+
+from model import (
+    NATIVE_RES, N_PARAMS, PARAM_SPECS, WMAX,
+    coord_features, decode_pixels, encode_prompt,
+    init_weights, load_model, prompts_to_embeddings, save_model,
+)
+
+MODEL_PATH = "model.png"
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--data",     required=True, help="npz with images (N,64,64,3) uint8")
+    p.add_argument("--captions", default=None,  help="json list of captions (default: <data>.captions.json)")
+    p.add_argument("--epochs",   type=int,   default=30)
+    p.add_argument("--batch",    type=int,   default=128)
+    p.add_argument("--pixels",   type=int,   default=768, help="pixel coords sampled per step")
+    p.add_argument("--lr",       type=float, default=2e-3)
+    p.add_argument("--seed",     type=int,   default=0)
+    p.add_argument("--device",   default="auto", help="auto, cpu, cuda, or a PyTorch device string")
+    p.add_argument("--resume",   action="store_true", help="continue from existing model.png")
+    args = p.parse_args()
+
+    torch.manual_seed(args.seed)
+    rng = np.random.default_rng(args.seed)
+    device = torch.device("cuda" if args.device == "auto" and torch.cuda.is_available()
+                          else "cpu" if args.device == "auto" else args.device)
+    print(f"device: {device}")
+
+    data = np.load(args.data)
+    images = data["images"]                                  # (N, 64, 64, 3) uint8
+    cap_path = args.captions or args.data.replace(".npz", ".captions.json")
+    with open(cap_path, encoding="utf-8") as f:
+        captions = json.load(f)
+    N = len(captions)
+    assert images.shape[0] == N
+    res = images.shape[1]
+    print(f"dataset: {N} pairs @ {res}x{res}")
+
+    print("precomputing prompt embeddings...")
+    embs = prompts_to_embeddings(captions).to(device)        # (N, 64)
+
+    targets = torch.from_numpy(images.reshape(N, res * res, 3).astype(np.float32) / 255.0).to(device)
+    feats_all = coord_features(res).to(device)               # (res*res, 18)
+
+    if args.resume and os.path.exists(MODEL_PATH):
+        weights = load_model(MODEL_PATH)
+        print(f"resumed from {MODEL_PATH}")
+    else:
+        weights = init_weights(args.seed)
+    for w in weights.values():
+        w.data = w.data.to(device)
+        w.requires_grad_(True)
+    params = [weights[n] for n, _ in PARAM_SPECS]
+    opt = torch.optim.Adam(params, lr=args.lr)
+
+    steps_per_epoch = N // args.batch
+    print(f"training: {args.epochs} epochs x {steps_per_epoch} steps "
+          f"(batch={args.batch}, pixels/step={args.pixels}, lr={args.lr}, "
+          f"params={N_PARAMS})\n")
+
+    t0 = time.time()
+    for epoch in range(1, args.epochs + 1):
+        order = rng.permutation(N)
+        ep_loss, ep_steps = 0.0, 0
+        for s in range(steps_per_epoch):
+            idx = order[s * args.batch:(s + 1) * args.batch]
+            pix = torch.from_numpy(rng.choice(res * res, size=args.pixels, replace=False)).to(device)
+
+            emb = embs[idx]
+            tgt = targets[idx][:, pix, :]                    # (B, P, 3)
+
+            z = encode_prompt(weights, emb)
+            pred = decode_pixels(weights, z, feats_all[pix])
+            loss = torch.nn.functional.mse_loss(pred, tgt)
+
+            opt.zero_grad()
+            loss.backward()
+            opt.step()
+            with torch.no_grad():
+                for w in params:
+                    w.clamp_(-WMAX, WMAX)
+
+            ep_loss += loss.item()
+            ep_steps += 1
+
+        save_model(weights, MODEL_PATH)
+        elapsed = time.time() - t0
+        print(f"epoch {epoch:>3}/{args.epochs}  loss={ep_loss / ep_steps:.5f}  "
+              f"elapsed={elapsed:.0f}s  -> saved {MODEL_PATH}", flush=True)
+
+    print(f"\nDone in {time.time() - t0:.0f}s. Final model saved to {MODEL_PATH}")
+
+
+if __name__ == "__main__":
+    main()