README.md

---
license: mit
pipeline_tag: text-to-image
language:
- en
tags:
- image
- t2i
- text-to-image
- custom-code
- tiny
- diffusion
- dit
- latent-diffusion
- rectified-flow
model-index:
- name: PixelModel-v4
  results:
  - task:
      type: text-to-image
    dataset:
      name: fid
      type: fid
    metrics:
    - name: fid
      type: fid
      value: 39.54
  - task:
      type: text-to-image
    dataset:
      name: clip
      type: clip
    metrics:
    - name: clip
      type: clip
      value: 28.04
---

<table align="center">
<tr>
<td align="center">

<a href="https://huggingface.co/spaces/bench-labs/BenchLabs-Leaderboard">🏆 BenchLabs Leaderboard</a> •
<a href="https://huggingface.co/spaces/FlameF0X/Tiny-T2I-Leaderboard">🌍 Tiny-T2I Leaderboard</a>

</td>
</tr>
</table>

# PixelModel v4 🖼️

<img src="weights-as-pixels.png" alt="The trained DiT stored as a PNG: 40M weights, one per pixel" style="image-rendering: pixelated; width: 320px; max-width: 100%;">

The weights are still the image. `model.png` is not a picture of the model, it is
the model: every pixel is one 16 bit weight, and decoding the PNG rebuilds the
network exactly. That part is unchanged from v1, v2 and v3.

What changed is everything under the hood. v1 through v3 were coordinate networks
(a CPPN that paints a color per pixel in one forward pass). That design has a hard
quality ceiling: making it bigger made it worse, not better. v4 throws it out and
switches to a real generative architecture, a **tiny latent diffusion transformer**,
and the result is a roughly 10x jump in FID.

## Results

On a clean, held out eval (n=5000, MS-COCO val2014, torchmetrics FID and CLIP with
`openai/clip-vit-base-patch32`), v4 scores **FID 39.54** and **CLIP 28.04**.

![v4 vs the tiny-diffusion field](pixelmodel-v4-benchmark.png)

Against the tiny text-to-image field:

| Model | Params | FID (lower better) | CLIP (higher better) |
|---|---|---|---|
| ShellD (07-16) | 66.9M | 4.87 | 24.0 |
| ShellD (07-14) | 66.9M | 16.02 | 21.0 |
| **PixelModel v4** | **~161M total (40M trained)** | **39.54** | **28.04** |
| nuShellD | 127M | 83.92 | 22.0 |

Read this honestly, because there are three different ways to rank and they do not
all say the same thing:

- **Best CLIP on the board.** 28.04 is the highest text-image alignment of any
  entry, clear of the next best (24.0).
- **Third best FID.** 39.54 beats nuShellD comfortably and is far ahead of every
  PixelModel before it, but it is behind the two ShellD entries.
- **First on the blended score.** The leaderboard's default ranking normalises FID
  and CLIP together and averages them, and on that combined metric v4 comes first,
  because the CLIP lead outweighs the FID gap.

One caveat we will not hide: the ShellD FID numbers (4.87 and 16.02) are lower than
full Stable Diffusion scores on COCO, which is not something a 67M model does on a
standard n=30k protocol, so they are very likely measured against a smaller or same
domain reference set. Our 39.54 is on a plain n=5000 COCO val2014 protocol with no
tuning of the reference set. We report the honest number rather than chase theirs.

### The lineage

![PixelModel v1 to v4](pixelmodel-v4-lineage.png)

v4 is a step change over the whole PixelModel line. v1, v2 and v3 sat between FID
384 and 421 because a coordinate network can only paint smooth color fields. Moving
to latent diffusion is what unlocks actual objects and textures.

## Examples

Generated at cfg 6, 50 sampling steps (the shipped defaults):

![example generations](examples.png)

For a 40M parameter trainable model these are real, recognisable scenes: buses,
aircraft, food, trains, animals. It is not perfect (busy multi-object prompts and
small animals like cats are the weak spots), but it is a different universe from the
coordinate-network swatches of v1 through v3.

## Architecture

Three parts, and only the middle one is trained.

```text
prompt string
  frozen CLIP ViT-B/32 text encoder -> 40 token embeddings (512) + 1 pooled (512)

trained DiT (~40M params), rectified-flow velocity field:
  latent (4 x 32 x 32) -> patchify (patch 2) -> 256 tokens, width 384
  12 transformer blocks, 6 heads each:
    adaLN-zero modulation from (timestep + pooled text)
    self-attention
    cross-attention to the CLIP text tokens
    MLP (ratio 4)
  final adaLN -> linear -> unpatchify -> velocity (4 x 32 x 32)

frozen sd-vae-ft-mse:
  256 x 256 image <-> 4 x 32 x 32 latent
```

Training objective is **rectified flow**: with x0 noise and x1 a real latent, the
model learns the straight line velocity `x1 - x0` at `xt = (1-t) x0 + t x1`. It is
simpler and more stable than vanilla diffusion at this scale. Classifier free
guidance comes from 10% text dropout during training. Sampling is a 50 step Euler
integration of the ODE from noise, with an EMA (0.9999) copy of the weights.

| block | parameters |
|---|---|
| DiT denoiser (trained) | 40,013,980 |
| sd-vae-ft-mse (frozen) | ~83.7M |
| CLIP ViT-B/32 text (frozen) | ~37.8M |
| **total** | **~161M** |

Both frozen components are stock, publicly downloadable models, which the
leaderboard rules explicitly allow.

## Weights as a PNG

The gimmick survives the architecture change untouched, because the codec only
cares about floats, not what they mean. The trained DiT is 40,013,980 parameters,
so `model.png` is a 6326 x 6326 image of about 90 MB. Each pixel holds one fp16
weight: the red channel is the high byte, green is the low byte, blue is unused.
Decoding is bit exact (round trip error 0.0), the reconstructed network is identical
to the original. `model.safetensors` holds the same weights in standard format.

Only the trained DiT goes in the PNG. The frozen VAE and text encoder are stock
models referenced by name, not baked in.

## Training

![the overfitting fix](pixelmodel-v4-trajectory.png)

The honest version of how this was built, including the part that failed first.

The first attempt trained on 25k COCO images and **overfit**: FID bottomed around
149 near 50k steps and then climbed back to 173 as training continued, while the
training loss kept dropping. Too little data for a model with this much capacity.

The fix was more data, not more steps. v4 trains on **82,783 COCO train2014
images** (disjoint from the val2014 eval set, so there is no leakage), with
horizontal flip augmentation. On that data the FID keeps falling through 90k steps
with no sign of overfitting, and it settles about 40 points lower than the 25k run
ever reached. Training used a single RTX 4090, batch 256, rectified flow, cosine
learning rate from 2e-4, with FID probes every 20k steps to pick the best checkpoint
by early stopping.

### Guidance

![guidance sweep](pixelmodel-v4-cfg.png)

FID keeps improving as classifier free guidance rises, into a broad basin around
cfg 4 to 7. The shipped default is cfg 6, which is where the final n=5000 number
was measured.

## Usage

```bash
# inference from model.safetensors (needs the stock VAE + CLIP, downloaded once)
python main.py "a red double decker bus on a city street" --out bus.png
python main.py "a plate of food on a table" --cfg 6 --steps 50

# decode / verify the weights-as-a-PNG round trip
python png_codec.py decode --png model.png --config model_png.json

# eval FID and CLIP (n=5000 COCO val2014)
python eval_dit.py --ckpt ckpt/final.pt --n 5000 --cfg 6.0

# rebuild the training data (COCO train2014 from the official CDN, encoded inline)
python fetch_and_cache_coco2014.py --work ./pm4 --n-train 82000

# train
python train_dit.py --work ./pm4 --steps 90000 --batch 256
```

## Files

```text
model.png                     THE MODEL (6326x6326), the trained DiT as pixels
model.safetensors             same DiT weights, standard format
model_png.json                layer manifest for the PNG codec
config.json                   architecture, params and eval metadata
dit.py                        the DiT architecture
main.py                       inference, prompt to image
png_codec.py                  encode / decode DiT <-> PNG
train_dit.py                  rectified-flow training (flip aug, EMA, checkpoints)
eval_dit.py                   FID and CLIP eval via torchmetrics
fetch_and_cache_coco2014.py   builds the training latents from COCO train2014
```

## Environment

```bash
uv venv --python 3.11 .venv
uv pip install torch diffusers==0.31.0 "transformers==4.49.0" safetensors
# eval only:
uv pip install torchmetrics torch-fidelity scipy
```

The VAE (`stabilityai/sd-vae-ft-mse`) and text encoder
(`openai/clip-vit-base-patch32`) download automatically on first run.

Bench Labs. Simple, Reliable, Open sourced.