File size: 3,395 Bytes
3dd752f
 
 
fa36b8b
3dd752f
 
 
 
 
 
fa36b8b
 
 
3dd752f
 
 
1b38e93
 
cfbcc16
 
1b38e93
 
 
3dd752f
0f8003c
 
 
 
 
 
 
 
 
 
 
 
 
 
6695e9b
0f8003c
6695e9b
 
 
 
3dd752f
6695e9b
3dd752f
e17aaaf
 
3dd752f
e17aaaf
4ad8e7e
e17aaaf
 
 
4ad8e7e
e17aaaf
4ad8e7e
3dd752f
 
 
6695e9b
3dd752f
 
1b38e93
3dd752f
6695e9b
 
 
 
3dd752f
cfbcc16
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
---
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`.