---
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

One epoch of [FLUX-Reason-6M](https://huggingface.co/datasets/LucasFang/FLUX-Reason-6M)
(Aesthetics-Part01/02, Imaginative, Text — 5.89M images): **560 142 steps** at batch 16, latent MSE
**0.2274** (pixel 0.0169), lr cosine 1e-4 → 1e-5, ~34 h on a single RTX 4080. The bundled
`transformer/` is the final checkpoint of that run.

## Examples

Validated on the butterfly set first — step 20 003, 970 images (63% of the AE's own high-frequency
ceiling at 256²):

<img src="docs/img/samples_butterflies_20k.png" width="100%"/>

Final checkpoint on FLUX-Reason-6M, step 560 142, 256 px, 16 body steps, 8 head steps, cfg 2.5,
seed 0 (fisherman · studio portrait · charcoal · anatomical eye · ink bird · impressionism ·
art nouveau · dance photograph):

<img src="docs/img/samples.png" width="100%"/>

## Usage

```bash
python3 src/sample.py --prompt "A portrait photograph of an elderly fisherman in a wool cap." --out portrait.png
```

## Limitations

- Hands, limbs and skeletons fail (the anatomical plate renders an eye, not a real cross-section);
  anything finer than 8 px is invented by the decoder.
- Loose caption adherence, no real glyphs (text inside images is decoration), 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`.