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