Upload 23 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* 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
+model.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,3 +1,72 @@
----
-license: mit
----
+# PixelModel 🖼️
+
+> A neural network where the weights **are** the image.
+
+## What is this?
+
+`model.png` is not a picture of anything. It *is* the model.  
+Every pixel's RGB values store neural network weights:
+- **R channel** → weight magnitude
+- **B channel** → weight sign (≥128 = positive)
+- **G channel** → bias values
+
+At inference, pixels are parsed into 3 weight matrices (a tiny MLP),  
+your prompt is embedded into a small vector, and a forward pass  
+produces a 32×32 output image.
+
+Training = gradient descent directly on pixel values. The PNG gets  
+updated until it can reproduce all images in the dataset.
+
+## Files
+
+```
+model.png       ← THE MODEL (64×3200 px, ~284 KB)
+main.py         ← inference
+train.py        ← training
+model.py        ← architecture (pixels → weights → forward pass)
+dataset/        ← your training data
+  cat.png
+  cat.txt       ← contains the prompt: "a cat"
+  ...
+```
+
+## Usage
+
+**Train:**
+```bash
+# Add image+prompt pairs to dataset/, then:
+python train.py
+python train.py --epochs 500 --lr 0.05
+```
+
+**Generate:**
+```bash
+python main.py "red"
+python main.py "a cat" --out cat_out.png --scale 8
+```
+
+`--scale 8` upscales 32×32 → 256×256 with nearest-neighbour (keeps the pixel look).
+
+## Architecture
+
+```
+prompt string
+  → char-level embedding → 32-dim vector
+  → W1 (64×32) → tanh
+  → W2 (64×64) → tanh  
+  → W3 (3072×64) → sigmoid
+  → reshape → 32×32×3 image
+```
+
+All weights live in `model.png`. The file you open in Preview **is** the neural net.
+
+## Dataset tips
+
+- 6–20 pairs works fine
+- Simpler targets = faster convergence (solid colors, gradients, simple shapes)
+- Run 200–500 epochs, loss should drop below 0.001 for simple datasets
+- The model capacity is fixed (64×3200 pixels = ~600K "parameters")
+
+## Why?
+
+It's a toy. It's not useful. But it's cool that it works.

dataset/blue.txt

@@ -0,0 +1 @@
+blue

dataset/dark.txt

@@ -0,0 +1 @@
+dark

dataset/green.txt

@@ -0,0 +1 @@
+green

dataset/red.txt

@@ -0,0 +1 @@
+red

dataset/white.txt

@@ -0,0 +1 @@
+white

dataset/yellow.txt

@@ -0,0 +1 @@
+yellow

main.py

@@ -0,0 +1,53 @@
+"""
+main.py  –  run the pixel model.
+
+Usage:
+  python main.py "a red circle"
+  python main.py "a red circle" --model model.png --out out.png --scale 8
+"""
+
+import argparse
+import sys
+import os
+import numpy as np
+from PIL import Image
+import torch
+
+from model import load_model, forward, OUT_SIZE
+
+
+def generate(prompt: str, model_path: str, out_path: str, scale: int = 8):
+    if not os.path.exists(model_path):
+        sys.exit(f"Model not found: {model_path}\n"
+                 f"Run: python train.py   to create one first.")
+
+    pixels = load_model(model_path)             # (H, W, 3)
+
+    with torch.no_grad():
+        result = forward(pixels, prompt)        # (32, 32, 3)
+
+    arr = (result.numpy() * 255).clip(0, 255).astype(np.uint8)
+    img = Image.fromarray(arr, mode="RGB")
+
+    # upscale for visibility (nearest-neighbour keeps pixel look)
+    if scale > 1:
+        img = img.resize(
+            (OUT_SIZE * scale, OUT_SIZE * scale),
+            Image.NEAREST
+        )
+
+    img.save(out_path)
+    print(f"prompt : '{prompt}'")
+    print(f"model  : {model_path}  ({os.path.getsize(model_path)//1024} KB)")
+    print(f"output : {out_path}  ({OUT_SIZE*scale}×{OUT_SIZE*scale} px)")
+
+
+if __name__ == "__main__":
+    p = argparse.ArgumentParser(description="PixelModel inference")
+    p.add_argument("prompt",              help="Text prompt")
+    p.add_argument("--model", default="model.png", help="Path to model PNG")
+    p.add_argument("--out",   default="out.png",   help="Output image path")
+    p.add_argument("--scale", type=int, default=8,
+                   help="Upscale factor for output (default 8 → 256×256)")
+    args = p.parse_args()
+    generate(args.prompt, args.model, args.out, args.scale)

model.py

@@ -0,0 +1,124 @@
+"""
+PixelModel - the weights ARE the image.
+
+model.png stores all weights as pixel RGB values (0-255 → floats).
+Architecture:
+  - Row 0..H//3-1      → Layer 1 weights  (prompt_dim → hidden)
+  - Row H//3..2*H//3-1 → Layer 2 weights  (hidden → hidden)
+  - Row 2*H//3..H-1    → Layer 3 weights  (hidden → output_flat)
+
+Width W = max(prompt_dim, hidden, output_flat), padded/sliced as needed.
+RGB channels encode sign & magnitude:
+  R = weight magnitude (0..255 → 0..1)
+  G = reserved (bias values, layer 1)
+  B = sign bit: <128 = negative, >=128 = positive
+"""
+
+import torch
+import torch.nn.functional as F
+import numpy as np
+from PIL import Image
+
+
+# ── config ────────────────────────────────────────────────────────────────────
+PROMPT_DIM  = 32    # prompt embedding size
+HIDDEN      = 64    # hidden layer width
+OUT_SIZE    = 32    # output image side length
+OUT_FLAT    = OUT_SIZE * OUT_SIZE * 3   # 3072
+
+# model.png dimensions that fit all weights
+# Layer sizes: (PROMPT_DIM→HIDDEN), (HIDDEN→HIDDEN), (HIDDEN→OUT_FLAT)
+# We stack weight matrices vertically, each row = one output neuron's weights
+# pad width to W = max of all input dims
+MODEL_W = max(PROMPT_DIM, HIDDEN)          # 64
+MODEL_H = HIDDEN + HIDDEN + OUT_FLAT       # 64+64+3072 = 3200
+# → model.png is 64×3200 px  (~600 KB uncompressed, tiny compressed)
+
+
+def prompt_to_embedding(prompt: str) -> torch.Tensor:
+    """Deterministic char-level embedding → PROMPT_DIM vector."""
+    vec = torch.zeros(PROMPT_DIM)
+    for i, ch in enumerate(prompt.lower()):
+        idx = i % PROMPT_DIM
+        vec[idx] += (ord(ch) / 127.0)
+    # normalise
+    norm = vec.norm()
+    if norm > 0:
+        vec = vec / norm
+    return vec  # shape (PROMPT_DIM,)
+
+
+def pixels_to_weights(pixels: torch.Tensor) -> tuple:
+    """
+    pixels: (H, W, 3) float32 tensor, values 0..1
+
+    Returns W1, W2, W3 weight matrices.
+    """
+    H, W, _ = pixels.shape
+
+    r = pixels[:, :, 0]   # magnitude 0..1
+    b = pixels[:, :, 2]   # sign: <0.5 = neg, >=0.5 = pos
+
+    # sign-magnitude → signed float in roughly [-1, 1]
+    sign  = torch.where(b >= 0.5, torch.ones_like(b), -torch.ones_like(b))
+    vals  = sign * r  # signed weights
+
+    row = 0
+
+    # W1: HIDDEN × PROMPT_DIM
+    W1_rows = HIDDEN
+    W1 = vals[row: row + W1_rows, :PROMPT_DIM]          # (64, 32)
+    row += W1_rows
+
+    # W2: HIDDEN × HIDDEN
+    W2_rows = HIDDEN
+    W2 = vals[row: row + W2_rows, :HIDDEN]               # (64, 64)
+    row += W2_rows
+
+    # W3: OUT_FLAT × HIDDEN
+    W3_rows = OUT_FLAT
+    W3 = vals[row: row + W3_rows, :HIDDEN]               # (3072, 64)
+    row += W3_rows
+
+    return W1, W2, W3
+
+
+def forward(pixels: torch.Tensor, prompt: str) -> torch.Tensor:
+    """
+    pixels : (H, W, 3) float32, values 0..1  ← the model
+    prompt : str
+    returns: (OUT_SIZE, OUT_SIZE, 3) float32, values 0..1
+    """
+    emb = prompt_to_embedding(prompt)           # (32,)
+    W1, W2, W3 = pixels_to_weights(pixels)
+
+    x = torch.tanh(W1 @ emb)                   # (64,)
+    x = torch.tanh(W2 @ x)                     # (64,)
+    x = torch.sigmoid(W3 @ x)                  # (3072,)  values 0..1
+
+    img = x.reshape(OUT_SIZE, OUT_SIZE, 3)
+    return img
+
+
+def load_model(path: str) -> torch.Tensor:
+    """Load model.png → float tensor (H, W, 3) in [0,1]."""
+    img = Image.open(path).convert("RGB")
+    arr = np.array(img, dtype=np.float32) / 255.0
+    return torch.tensor(arr)
+
+
+def save_model(pixels: torch.Tensor, path: str):
+    """Save float tensor (H, W, 3) in [0,1] → model.png."""
+    arr = (pixels.detach().clamp(0, 1).numpy() * 255).astype(np.uint8)
+    Image.fromarray(arr, mode="RGB").save(path)
+
+
+def init_model(path: str):
+    """Create a fresh random model.png."""
+    # small random weights around 0 → R~0.1, B random sign
+    r = torch.rand(MODEL_H, MODEL_W) * 0.2          # small magnitude
+    g = torch.zeros(MODEL_H, MODEL_W)
+    b = (torch.rand(MODEL_H, MODEL_W) > 0.5).float()  # random sign
+    pixels = torch.stack([r, g, b], dim=2)
+    save_model(pixels, path)
+    print(f"Initialised model: {path}  ({MODEL_W}×{MODEL_H} px)")

train.py

@@ -0,0 +1,110 @@
+"""
+train.py  –  optimize pixel values via gradient descent.
+
+Dataset: dataset/ folder containing pairs:
+  dataset/cat.png + dataset/cat.txt   (image + its prompt)
+  dataset/sun.png + dataset/sun.txt
+  ...
+
+Each .txt holds one line: the prompt for that image.
+Target images are resized to OUT_SIZE×OUT_SIZE automatically.
+
+Usage:
+  python train.py                     # train with defaults
+  python train.py --epochs 500 --lr 0.05
+"""
+
+import argparse
+import os
+import sys
+import torch
+import torch.nn.functional as F
+import numpy as np
+from PIL import Image
+
+from model import (
+    MODEL_H, MODEL_W, OUT_SIZE,
+    forward, load_model, save_model, init_model,
+    pixels_to_weights, prompt_to_embedding,
+)
+
+DATASET_DIR  = "dataset"
+MODEL_PATH   = "model.png"
+
+
+def load_dataset():
+    pairs = []
+    for fname in sorted(os.listdir(DATASET_DIR)):
+        if not fname.lower().endswith((".png", ".jpg", ".jpeg")):
+            continue
+        stem = os.path.splitext(fname)[0]
+        txt_path = os.path.join(DATASET_DIR, stem + ".txt")
+        img_path = os.path.join(DATASET_DIR, fname)
+        if not os.path.exists(txt_path):
+            print(f"  [skip] no .txt for {fname}")
+            continue
+        with open(txt_path) as f:
+            prompt = f.read().strip()
+        img = Image.open(img_path).convert("RGB").resize(
+            (OUT_SIZE, OUT_SIZE), Image.BILINEAR
+        )
+        target = torch.tensor(
+            np.array(img, dtype=np.float32) / 255.0
+        )  # (32,32,3)
+        pairs.append((prompt, target))
+        print(f"  loaded: '{prompt}'  ← {fname}")
+    return pairs
+
+
+def train(epochs=300, lr=0.03, save_every=50):
+    # init model if missing
+    if not os.path.exists(MODEL_PATH):
+        init_model(MODEL_PATH)
+
+    raw = load_model(MODEL_PATH)               # (H, W, 3)
+
+    # pixels are the parameters — wrap in a leaf tensor
+    pixels = raw.clone().requires_grad_(True)
+    optimizer = torch.optim.Adam([pixels], lr=lr)
+
+    print(f"\nLoading dataset from '{DATASET_DIR}/'…")
+    dataset = load_dataset()
+    if not dataset:
+        sys.exit("No dataset pairs found. Add image+txt pairs to dataset/")
+
+    print(f"\nTraining  epochs={epochs}  lr={lr}  pairs={len(dataset)}\n")
+
+    for epoch in range(1, epochs + 1):
+        optimizer.zero_grad()
+        total_loss = torch.tensor(0.0)
+
+        for prompt, target in dataset:
+            # pixels must stay in [0,1] for the sign/magnitude encoding
+            p_clamped = pixels.clamp(0, 1)
+            pred = forward(p_clamped, prompt)       # (32,32,3)
+            loss = F.mse_loss(pred, target)
+            total_loss = total_loss + loss
+
+        total_loss.backward()
+        optimizer.step()
+
+        # clamp pixels back to valid range after update
+        with torch.no_grad():
+            pixels.clamp_(0, 1)
+
+        if epoch % save_every == 0 or epoch == 1:
+            save_model(pixels, MODEL_PATH)
+            avg = total_loss.item() / len(dataset)
+            print(f"  epoch {epoch:>5}/{epochs}  loss={avg:.5f}  → saved model.png")
+
+    save_model(pixels, MODEL_PATH)
+    print(f"\nDone. Final model saved to {MODEL_PATH}")
+
+
+if __name__ == "__main__":
+    p = argparse.ArgumentParser()
+    p.add_argument("--epochs",     type=int,   default=300)
+    p.add_argument("--lr",         type=float, default=0.03)
+    p.add_argument("--save-every", type=int,   default=50)
+    args = p.parse_args()
+    train(args.epochs, args.lr, args.save_every)