| --- |
| 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> |
| <tr> |
| <td> |
| <br> |
| |
| # See us on the [Leaderboard ⇗](https://huggingface.co/spaces/FlameF0X/Tiny-T2I-Leaderboard) |
| |
| </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**. |
|
|
|  |
|
|
| 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 |
|
|
|  |
|
|
| 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): |
|
|
|  |
|
|
| 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 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 |
|
|
|  |
|
|
| 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. |
|
|
