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