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.

Downloads last month
-
Safetensors
Model size
0.2B params
Tensor type
F16
·
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Dataset used to train AiArtLab/Ra