How to use from the
Use from the
Diffusers library
pip install -U diffusers transformers accelerate
import torch
from diffusers import DiffusionPipeline

# switch to "mps" for apple devices
pipe = DiffusionPipeline.from_pretrained("AiArtLab/Ra", dtype=torch.bfloat16, device_map="cuda")

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
image = pipe(prompt).images[0]

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

One epoch of 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Β²):

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

Usage

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