README.md

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

# PixelModel v6 🖼️
<a href="https://huggingface.co/spaces/bench-labs/BenchLabs-Leaderboard">🏆 BenchLabs Leaderboard</a> •
<a href="https://huggingface.co/spaces/bench-labs/blog?post=Pm6.html">📝 Blog Post</a>

<img src="exampl1.png" alt="Two 256px images generated by PixelModel v6" style="width: 100%; max-width: 600px;">
<img src="example2.png" alt="Two 256px images generated by PixelModel v6" style="width: 100%; max-width: 600px;">

*cfg 4.0, single draw each, not
cherry-picked from multiple seeds*

PixelModel v6 is a 155,090,448 parameter text-to-image diffusion transformer. It generates 256x256
images in 50 sampling steps and it trained in about 55 hours across two A100s.

Where [v5](https://huggingface.co/bench-labs/PixelModel-v5) changed the data and kept v4's
architecture byte for byte, v6 changes the architecture and keeps v5's data. v5's own README ended
with a diagnosis: the model had learned appearance thoroughly and topology poorly, and topology is
what capacity buys. This is that bet.

## What actually changed

Six real architecture changes, not one, which is a departure from this line's usual practice of
isolating a single variable per release. Each change has independent literature support; the
combination does not, and that is worth being upfront about rather than implying six papers'
results simply add up.

**Joint attention instead of cross-attention.** v4 and v5 read text through cross-attention, the way
PixArt-alpha does it. v6 concatenates image and text tokens into one sequence and runs joint
self-attention over both, the way SD3's MMDiT does it, so text tokens can attend to image tokens and
not just the reverse.

**T5 instead of CLIP-only conditioning.** The joint stream is built from `google/flan-t5-base`
hidden states rather than CLIP token embeddings. CLIP's pooled vector is kept, but only for the
global adaLN conditioning signal. T5 conditioning is the specific change PixArt-alpha, DeepFloyd IF,
and SD3 all made for prompt adherence, and v5's README named exactly this failure mode: prompts
asking for one object in an empty scene often returned the scene without the object.

**2D RoPE instead of additive sincos position embeddings.** Image tokens get axial rotary position
encoding, text tokens get standard 1D rotary encoding. This is also what makes a future higher
resolution fine-tune cheap: RoPE carries no learned parameters, so a checkpoint trained at 256px can
extend to a larger token grid without retraining or interpolating a position embedding table.

**QK-norm and SwiGLU.** RMSNorm on queries and keys before attention, and a SwiGLU MLP in place of
GELU. Both are standard modern stabilizers, included because a bigger model trains less forgivingly
than a 40M one.

**REPA.** An auxiliary loss projecting mid-stack image token hidden states to align with frozen
DINOv2-small features of the same image, cosine similarity, weight scheduled to peak early and decay
to exactly zero by step 105,000. This is the one item on this list without solid precedent at this
scale: the original REPA paper is class-conditional ImageNet at 675M parameters, and v5's own
planning notes flagged that transfer to text-conditional generation at a fraction of that size as
genuinely unverified. The decay schedule follows the follow-up finding that the alignment loss helps
early and can hurt if left on too long.

**More capacity.** dim 512, depth 16, heads 8: 155.1M trainable, versus v5's 40M. Still well short of
SD1.5's 860M UNet or SDXL's 2.6B, but the largest this line has been.

<img src="pixelmodel-v6-params.png" alt="Trainable backbone size by version, v4 and v5 at 40M, v6 at 155.1M" style="width: 100%; max-width: 700px;">

The frozen VAE also moved from `sd-vae-ft-mse` to `madebyollin/sdxl-vae-fp16-fix`, a drop-in upgrade
at the same latent shape.

## Training

<img src="pixelmodel-v6-loss.png" alt="Training loss, diffusion and REPA components, over 150,000 steps" style="width: 100%; max-width: 860px;">

150,000 steps at batch 192 on the same 2,986,571 image CC12M-recaptioned set v5 used, about 9.7
epochs. AdamW at 2e-4, cosine decay after a 1,500 step warmup, EMA at 0.9999, 10% caption dropout,
gradient checkpointing throughout. The REPA loss ran at peak weight for the first 40% of training,
decayed linearly through the next 30%, and sat at exactly zero for the final 30%, both to follow the
early-stopping finding above and because computing DINOv2 targets costs real throughput: the run held
0.86 steps/s while REPA was active and 1.39 steps/s once it turned off.

<img src="pixelmodel-v6-val.png" alt="Held-out validation loss falling monotonically from 0.938 to 0.798" style="width: 100%; max-width: 860px;">

Held-out loss on a fixed 4,096-pair split, evaluated every 2,000 steps, fell without interruption
from 0.938 to 0.798 across the entire run. No overfitting turn-up at any point, same clean shape v5
reported.



## Results

Measured on 5,000 MS-COCO val2014 images at 256 center crop, `torchmetrics` FID at 2048 features and
`torchmetrics` CLIPScore with `openai/clip-vit-base-patch32`. Identical protocol to v4 and v5.

| | FID ↓ | CLIP ↑ |
|---|---|---|
| PixelModel v4 | 39.54 | 28.04 |
| PixelModel v5 | 34.06 | 29.27 |
| **PixelModel v6, cfg 3.0** | **23.62** | 30.43 |
| real photographs | n/a | 30.05 |

v6 beats v5's best FID by over 10 points at every guidance scale tested, not just at one cherry-picked
value. Full sweep, same 5,000 image set and 50 sampling steps throughout:

| cfg | FID ↓ | CLIP ↑ |
|---|---|---|
| 2.0 | 25.65 | 29.65 |
| **3.0** | **23.62** | 30.43 |
| 4.0 | 23.77 | 30.81 |
| 5.0 | 24.42 | 31.06 |
| 6.0 | 25.10 | 31.12 |
| 7.0 | 26.26 | 31.16 |
| 8.0 | 26.87 | 31.17 |

FID bottoms at cfg 3.0 and rises steadily past it, the usual guidance tradeoff: more weight on the
text condition pulls generations away from the real-image manifold even as it pulls them closer to
the prompt. cfg 3.0 is the reported default. Worth flagging rather than celebrating: CLIP score
crosses the real-photograph baseline (30.05) starting around cfg 4.0 and keeps climbing to 31.17.
That is not v6 out-photographing real photographs, it is high guidance making generations more
literal about the prompt than an actual photo ever is, which CLIP score rewards and FID does not. Use
FID, not CLIP, to pick a guidance scale.

## What it does well and what it does not

Judged across several hundred samples during training, not just the checkpoint above.

It is good at material, light, and increasingly at structure that v5 could not hold: food, faces of
common animals, latte foam texture, boat masts and sails, neon-lit streets. Where v5's README singled
out faces and animal heads as a weak point, v6 renders recognizable dog and fox faces reliably by the
back half of training, something no earlier version in this line managed.

It is still uneven on rare subjects. Across many samples of the same fox prompt, output quality
varied far more than for common subjects like coffee or strawberries, which is a direct read on
training data frequency in CC12M rather than a model defect: the architecture change did not fix data
scarcity, it just makes better use of what is there.

A watermark and ghost-text artifact shows up intermittently, most often on the lighthouse and boat
prompts: faint illegible text-like marks bled in from CC12M's watermarked stock photography, the same
failure family v5 documented for actual text rendering. It did not fully train out by step 150,000.

## Specification

| | |
|---|---|
| trainable parameters | 155,090,448 |
| frozen VAE, `sdxl-vae-fp16-fix` | 83,653,863 |
| frozen text encoder, CLIP ViT-B/32 | 63,165,952 |
| frozen text encoder, T5-base | 109,628,544 |
| **combined total** | **411,538,807** |
| architecture | MMDiT, dim 512, depth 16, heads 8, patch 2 |
| tokens | 256 image (16x16 over a 32x32x4 latent) + 32 text |
| objective | rectified flow, target `x1 - x0`, logit-normal timesteps |
| conditioning | joint attention on T5 tokens, plus adaLN-zero from CLIP pooled + timestep |
| resolution | 256x256 |
| sampling | 50 steps, cfg 3.0 default (23.62 FID); see cfg sweep above for the tradeoff |

All parameter counts above are the released backbone. Training also carried a small REPA projection
head (459,648 params, 155,550,096 total while training) that only ever fed the auxiliary alignment
loss; it has no role in sampling and is dropped from the published weights.

## Usage

```bash
python main.py "a bowl of ramen with a soft boiled egg" --out ramen.png --cfg 3.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` | the trainable MMDiT backbone only, fp32. VAE, CLIP, and T5 load fresh from their own HF repos at inference time, same convention as v4 and v5 despite what their READMEs said |
| `model.png` | the same backbone weights as an image, fp16, packed via `png_codec.py`. Verified bit-exact against `model.safetensors` (rounded to fp16) before upload |
| `model_png.json` | manifest `png_codec.py` needs to decode `model.png` back into a state dict |
| `png_codec.py` | the PNG weight codec, encode and decode |
| `dit_v6.py` | the model |
| `train_v6.py` | training loop, REPA schedule, held-out split |
| `prep_v6.py` | streams CC12M-recaptioned, encodes latents, caches tokens |
| `build_eval_set.py` | builds the 5,000 image COCO val2014 eval set |
| `eval_v6.py` | FID / CLIPScore protocol and cfg sweep |
| `preview.py` | quick sampling CLI for spot checks |
| `main.py` | inference |