| --- |
| 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`](https://huggingface.co/Muinez/f8c32-ae), bundled in `vae/` | |
| | text encoder | [`jhu-clsp/ettin-encoder-150m`](https://huggingface.co/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²): |
|
|
| <img src="docs/img/samples_butterflies_20k.png" width="100%"/> |
|
|
| 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): |
|
|
| <img src="docs/img/samples.png" width="100%"/> |
|
|
| ## Usage |
|
|
| ```bash |
| 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](https://huggingface.co/jhu-clsp/ettin-encoder-150m) is MIT, the autoencoder |
| [f8c32-ae](https://huggingface.co/Muinez/f8c32-ae) is Apache-2.0. Architecture notes: `docs/ARCH.md`, |
| Russian slides: `docs/presentation.pdf`. |
|
|
