README.md

---
license: mit
pipeline_tag: text-to-image
language:
- en
tags:
- text-to-image
- diffusion
- rectified-flow
- dit
- tiny
- custom-code
---

<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> •
<a href="https://huggingface.co/spaces/bench-labs/blog?post=Pm5.html">📝 Blog Post</a>

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

# PixelModel v5 🖼️

<img src="examples.png" alt="Nine 256px images generated by PixelModel v5" style="width: 100%; max-width: 860px;">

PixelModel v5 is a 40,013,980 parameter text-to-image diffusion transformer. It generates 256x256
images in 50 sampling steps and it trained in about 7 hours on one A100.

The architecture is identical to [v4](https://huggingface.co/bench-labs/PixelModel-v4). Same DiT,
same width, same depth, same VAE, same text encoder, same objective. Nothing in the model changed.
The only difference is that v5 saw 36 times more data.

That was the whole experiment, and it is worth being blunt about why.

## v4 was starving, not undersized

v4 trained on MS-COCO, which is about 83,000 images. At batch 256 for 80,000 steps that works out to
roughly 247 passes over the same photographs. A model that sees the same 83,000 images 247 times is
not learning what the world looks like, it is memorising a small album.

We spent a while looking at architecture changes first. Linear attention, deeper stacks, a
higher compression autoencoder, representation alignment. All of them are real techniques and none
of them address the actual problem, because the bottleneck was never the shape of the network.

<img src="pixelmodel-v5-data.png" alt="Training set size, v4 at 83k images versus v5 at 2.99M" style="width: 100%; max-width: 780px;">

v5 trains on 2,986,571 image and caption pairs, which is about 10 epochs instead of 247.

## The captions matter as much as the count

The images come from CC12M. Raw CC12M alt text is close to unusable for this, because it is scraped
SEO copy. A typical caption reads like "This is the best recipe I have ever tried for Cuban bread.
Cuban Recipes, Bread Recipes, Cooking Recipes, Pan Cubano".

v5 uses a recaptioned version where a vision model rewrote every caption in COCO style:

> A purple-lit subway train sits at a station platform.
>
> An old red tractor with gold fenders and red wheels stands in a snowy field.
>
> A compact kitchenette in a studio apartment features a stainless steel refrigerator, a microwave,
> and a sink with a glass faucet.

Those describe what is in the frame, which is what a text-to-image model needs. They also match the
caption distribution that FID and CLIP score are computed against, so the training signal and the
evaluation are speaking the same language.

## Training

<img src="pixelmodel-v5-loss.png" alt="Training and held-out loss over 120,000 steps" style="width: 100%; max-width: 860px;">

120,000 steps at batch 256. AdamW at 2e-4 with betas (0.9, 0.99), no weight decay, gradient clipping
at 1.0, cosine decay to a tenth of peak after a 1,000 step warmup, EMA at 0.9999 starting after
warmup, and 10% caption dropout for classifier free guidance. It held 4.6 steps per second at
44.2 GiB the whole way.

The held-out curve climbs for the first 10,000 steps or so, which looks alarming and is not. EMA at
0.9999 has a time constant of roughly 10,000 steps, so early on those weights are still mostly
initialisation. The number only becomes meaningful once EMA has caught up, and from there it falls
without interruption.

<img src="pixelmodel-v5-val.png" alt="Held-out loss falling monotonically with no turn-up" style="width: 100%; max-width: 860px;">

**v4 had no validation split at all**, so it shipped whichever weights the last step happened to
produce. v5 holds out 4,096 pairs, evaluates every 2,000 steps, and writes the best checkpoint to a
separate file. Across the whole run the held-out loss improved at every single check and never once
turned upward, which is the clearest evidence that the data change worked. Our sibling model
[VoxelModel v1](https://huggingface.co/bench-labs/VoxelModel-v1), trained on 28,415 meshes for 811
epochs, peaked at step 64,000 and got 3.08% worse by the end. Same recipe, same team, different data
regime, opposite outcome.

## Results

Measured on 5,000 MS-COCO val2014 images at 256 center crop, using `torchmetrics` FID at 2048
features and `torchmetrics` CLIPScore with `openai/clip-vit-base-patch32`. This is v4's protocol.

| | FID ↓ | CLIP ↑ |
|---|---|---|
| PixelModel v4 | 39.54 | 28.04 |
| **PixelModel v5, cfg 5.0** | **34.06** | **29.27** |
| PixelModel v5, cfg 3.0 | 35.03 | 28.59 |
| PixelModel v5, cfg 6.0 | 35.18 | 29.28 |
| real photographs | n/a | 30.05 |

FID is 13.9% lower and CLIP score is 1.23 higher, at an identical parameter count. Nothing in the
network changed between v4 and v5, so the entire difference is the training data.

The last row is there because a metric without a reference point is hard to read. Feeding the real
COCO photographs and their own captions through the same CLIPScore gives 30.05, so v5 at 29.27 is
within 0.78 of the ceiling this metric can award. That number also served as a sanity check on the
harness, since `torchmetrics` CLIPScore is incompatible with transformers 5.x and needed patching
before it would run at all.

Guidance scale was swept rather than assumed. FID bottoms out at cfg 5.0 and CLIP keeps climbing
slightly to cfg 6.0, which is the usual tension between the two. Anyone reproducing this should
also know that FID is heavily biased at small sample sizes. The same checkpoint at cfg 5.0 scores
74.35 on 1,000 images and 34.06 on 5,000, so sample count has to match before two FID numbers can be
compared at all.

One caveat worth stating plainly. v4's cached evaluation set is not published in its repository, so
the reference images here were rebuilt from COCO val2014 with a fixed seed. The protocol is the same
but the exact 5,000 photographs may not be. CLIP score is unaffected by this since it compares each
generated image against its own caption, but the FID comparison carries that asterisk.

## Specification

| | |
|---|---|
| trainable parameters | 40,013,980 |
| frozen VAE, `sd-vae-ft-mse` | 83,653,863 |
| frozen text encoder, CLIP ViT-B/32 | 63,165,952 |
| **combined total** | **186,833,795** |
| architecture | DiT, dim 384, depth 12, heads 6, patch 2 |
| tokens | 256, over a 32x32x4 latent |
| objective | rectified flow, target `x1 - x0`, logit-normal timesteps |
| conditioning | cross attention on 40 CLIP text tokens, plus adaLN-zero |
| resolution | 256x256 |
| sampling | 50 steps, cfg 5.0 for best FID, cfg 6.0 for best CLIP |

A note on that combined number, because it is higher than what v4 published. v4's card reports
161,500,000 total. That figure counts the CLIP text transformer but leaves out its 49,408 by 512
token embedding table, which is another 25.3M parameters. The encoder does not run without it, so it
counts. The correct total for both v4 and v5 is 186,833,795.

## What it does well and what it does not

After a few hundred samples the boundary is consistent.

It is good at material and light. Food, landscapes, skies, water, snow, foliage, stone, interiors,
depth of field. A bowl of ramen, a wet cobblestone street, a wheat field under storm light, a library
of wooden shelves. These come back looking like photographs.

<img src="hero.png" alt="Cobblestone street, library, wheat field, and snowy mountain generated by PixelModel v5" style="width: 100%; max-width: 640px;">

It is bad at structure. Faces and animal heads deform. Limbs merge. Wings on aircraft are stubby or
fused to the fuselage. Any text it renders, on a shop sign or an airliner livery, is convincing
looking gibberish. Prompts asking for one isolated object in an empty scene often return the scene
without the object, so "a lighthouse on a cliff at sunset" gives a very nice sunset over a cliff and
no lighthouse.

Frequency in the training data matters more than complexity. Aircraft are structurally complicated
and come out fine because CC12M is full of them. Hot air balloons are simpler and fail, because they
are rare.

That failure profile is what a 40M model looks like once data is no longer the constraint. It has
learned appearance thoroughly and topology poorly, and topology is what capacity buys. v6 is where
that gets spent.

## Usage

```bash
python main.py "a bowl of ramen with a soft boiled egg" --out ramen.png --cfg 6.0 --steps 50
```

As with every model in this line the weights are also stored as a PNG. `model.png` is not a picture
of the network, it is the network, with each parameter packed across the red and green channels of
one pixel.

## Files

| file | contents |
|---|---|
| `model.safetensors` | combined weights, backbone and VAE and text encoder |
| `model.png` | the same backbone weights as an image |
| `dit.py` | the model |
| `train_v5.py` | training loop with the held-out split |
| `prep_v5.py` | streams the dataset, encodes latents, caches captions |
| `main.py` | inference |