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, 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:

Usage

python3 src/sample.py --prompt "butterfly, Troides hypolitus" --out out.png

Training / fine-tuning

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.