File size: 3,731 Bytes
e995b65 | 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 77 78 79 80 81 | ---
license: mit
pipeline_tag: text-to-image
tags: [text-to-image, flow-matching, pyramidal-unet, latent-diffusion, asym-vae, from-scratch, ra, udit]
---
# Ra2
**Ra2** — second-generation layout of Ra: the transformer no longer carries a pixel-space refiner, so the
model is a plain latent diffusion model and the image is decoded by a **separate asymmetric VAE**
(f8 encoder / f16 decoder → 512×512 from a 32×32 latent).
It is a 256→512 text-to-image model trained from scratch: a pyramidal U-transformer (240M, 40 blocks)
flow-matching in an 8×/32-channel latent, a frozen Ettin-150M text encoder with a learnable layer-slice
mixer (10.0M), and a **separate asymmetric VAE** that decodes the latent into pixels.
| | |
|---|---|
| output | 512×512 from a 32×32 latent (encoder f8, decoder f16) |
| sampler | 16 Euler steps, cfg 2.5 |
| latent | 32 channels from the bundled `vae/` (AsymmetricAutoencoderKL) |
| text encoder | [`jhu-clsp/ettin-encoder-150m`](https://huggingface.co/jhu-clsp/ettin-encoder-150m), frozen, bundled in `text_encoder/` |
| body | 40 blocks — L0: 6 conv + 8 attention, L1: 9, L2: 9, L3: 4 (widths 192→512→1024→1536) |
| training | latent flow matching only (MSE in latent space); the VAE decodes, no pixel-space refiner |
| licence | MIT |
## Design
The body is a U-stack: width grows as tokens shrink, so expensive attention runs on 16–64 tokens while
the 1024-token grid keeps 3×3 convolutions, and skip connections add every level back in the decoder.
One L0-attention block costs 0.86 GMAC (3.9% of the body) at 0.37M parameters, which is why extra blocks
were added at full resolution rather than in depth.
Positions: 2D RoPE inside attention plus an absolute 2D sin/cos embedding projected into the level width
by a zero-initialised linear layer, so new blocks can be grafted onto a trained model without changing
its output (verified bit-identical: `Body.graft`).
Decoding is plain: the latent goes through the VAE decoder once. That keeps the objective a clean latent
MSE — the price is that fine detail must live in the 32×32 latent.
## Training
| stage | what | steps |
|---|---|---|
| 1 | FLUX-Reason-6M, 256², old f8c32 AE + pixel-space refiner (23.9M) | 560 142 |
| 2 | same data, capacity graft 26 → 36 blocks, absolute-PE branch | 404 000 |
| 3 | switch to the asymmetric VAE (32×32 latent → 512²), latent-only loss, MJ shard | 5 000 |
| 4 | butterfly fine-tune, 966 images, zero-init graft 36 → 40 blocks (3+3 L0 attention) + EMA | 25 000 |
The bundled `transformer/` is the stage-4 checkpoint (step 61 000 of the current counter). It is a
**research stage**, not a general-purpose model: it was trained to fit one narrow dataset.
## Examples
The current checkpoint on its training domain (butterfly specimens, 512² downscaled to 256²), and the
same species in the source data:
<img src="docs/img/samples.png" width="100%"/>
## Usage
```bash
python3 src/sample.py --prompt "butterfly, Troides hypolitus" --out out.png
```
## Training / fine-tuning
```bash
bash run.sh /path/to/tfrecord_shards 10 # 10 эпох: TFRecord-шарды i1 (PNG 256² + подписи)
python3 src/train.py # дефолты: --init transformer, --out tmp/run, --target 708000
```
## Limitations
- Butterflies only, right now: 966 training images, 3-word captions.
- Colour drifts towards saturated species stereotypes; vein pattern is approximate, not faithful.
- No text rendering, no faces, no general scenes. Research artifact — no safety tuning.
MIT for our code and weights (`src/`, `transformer/`). The bundled VAE is `AsymmetricAutoencoderKL`
(32 latent channels, encoder f8 / decoder f16); the text encoder is MIT.
|