See us on the Leaderboard โ‡—

PixelModel v4 ๐Ÿ–ผ๏ธ

The trained DiT stored as a PNG: 40M weights, one per pixel

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

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

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

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.

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

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

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

# 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

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

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.

Downloads last month
-
Safetensors
Model size
40.1M params
Tensor type
F32
ยท
Inference Providers NEW
This model isn't deployed by any Inference Provider. ๐Ÿ™‹ Ask for provider support

Spaces using bench-labs/PixelModel-v4 2

Collections including bench-labs/PixelModel-v4

Evaluation results