license: mit
pipeline_tag: text-to-image
tags:
- text-to-image
- flow-matching
- pyramidal-unet
- 256x256
- ettin-150m
- from-scratch
- ra
- udit
datasets:
- LucasFang/FLUX-Reason-6M
Ra
Ra is a 256×256 text-to-image model trained from scratch: 177.3M parameters — a pyramidal U-transformer (143.3M) flowing in an 8×/32-channel latent, a frozen Ettin-150M text encoder with a learnable layer-slice mixer (10.0M), and the AE decoder retrained as a pixel-space flow head (23.9M).
| output | 256×256, sides multiple of 64 |
| sampler | 8–16 Euler steps, cfg 1.5–2.5, 4–8 head steps |
| latent | 8× spatial / 32 channels from Muinez/f8c32-ae, bundled in vae/ |
| text encoder | jhu-clsp/ettin-encoder-150m, frozen, bundled in text_encoder/ |
| weights | fp16, 803 MiB bundle |
| training | FLUX-Reason-6M, one RTX 4080 |
| licence | MIT |
Why the pyramid
At the same parameter count, a flat classic DiT (12 blocks at width 1024, one resolution) costs 232 GFLOP per forward pass of the body; the pyramid costs 20.4 — 11× less compute. The step itself runs 97 ms vs 196 ms because it is memory-bound, and the detail metric is 0.54 vs 0.18 (ground truth 0.913, AE ceiling 0.859).
- Width grows as tokens shrink: expensive layers run on 16–64 tokens, while the 1024-token grid keeps cheap 3×3 convolutions — convolutions are local, so anatomy and texture survive.
- Additive skip connections at every level: fine structure is built at its own level instead of being pulled through one single-resolution attention stack.
- Text reaches all four levels; the 16-token bottleneck (width 1536, full attention) holds composition and symmetry.
Training in progress
The bundled transformer/ is a snapshot of step 212 000; lr was restarted at 1e-4 on step 192 000,
so it is mid-recovery and will be replaced.
Examples
Validated on the butterfly set first — step 20 003, the stage that proved the architecture (63% of the AE's own high-frequency ceiling at 256²):
Now training on FLUX-Reason-6M — bundled snapshot, step 212 000, 256 px, 16 body steps, 8 head steps, cfg 2.5, seed 0 (fox · sailboat · fisherman · sunflowers · blue morpho · white butterfly · swallowtail · monarch):
Usage
python3 src/sample.py --prompt "A macro photograph of a blue morpho butterfly." --out butterfly.png
Limitations
- Animal anatomy fails (see the fox), hands fail; detail finer than 8 px is invented by the decoder.
- Loose caption adherence, no glyphs, 256² only. Research artifact — no safety tuning, no filter.
MIT for our code and weights (src/, transformer/); the bundled text encoder
Ettin-150M is MIT, the autoencoder
f8c32-ae is Apache-2.0. Architecture notes: docs/ARCH.md,
Russian slides: docs/presentation.pdf.