IMG-253M

256×256 text-to-image model: a 253.4M-parameter flow-matching DiT trained from scratch, conditioned on a frozen SlayerLab/Slayer149 language model (first 18 of 19 blocks) through a 12.6M-parameter trainable refiner, decoding with the frozen FLUX.2 small decoder. Trained in ~12.5 h on 8× RTX 3090 Ti. Released weights: EMA.

Eight prompts from the AuroraIMG2 card, seed 0, 50 Euler steps, CFG 3.0:

1. a red fox sitting in a snowy forest at sunrise, photograph
2. a futuristic city skyline at night with neon lights reflected in the rain
3. a bowl of fresh strawberries on a wooden table, soft natural light
4. a watercolor painting of a lighthouse on a rocky coast
5. a small robot reading a book in a cozy library
6. an astronaut riding a horse on the moon, digital art
7. a street sign that says OPEN in neon letters
8. two children flying a kite on a green hill under a blue sky

Held-out loss

Velocity MSE (lower is better) on 4,096 held-out images (last shard of each dataset subset), fixed noise and timesteps, same evaluation code for both models:

Architecture

• Denoiser (253.4M): DiT, 16 blocks, width 1024, 16 heads, SwiGLU MLP 2432, one token per latent pixel (16×16 FLUX.2 latents, 128 channels, no patchify). Self-attention with 2D RoPE and QK-norm, cross-attention to the text tokens, PixArt-style adaLN-single timestep conditioning.
• Text encoder: frozen Slayer149 (GoLLeM-style causal GPT, 768 wide), hidden states before its last block → RMSNorm → 2 bidirectional refiner blocks (768 wide, GQA 12/4 heads, 1D RoPE, SwiGLU) trained with the denoiser. A learned vector serves as the empty prompt for classifier-free guidance. Prompts up to 95 tokens.
• Decoder (frozen): FLUX.2 small decoder, unchanged.

Training

• Data: QLNI/FLUX-Reason-6M-flux2-latents, 5.88M images, 5 epochs: 114,844 steps × batch 256 = 29.4M samples.
• Objective: rectified flow, velocity target, logit-normal timesteps, 10% caption dropout for CFG.
• REPA with PE-Spatial-B16 (Apache-2.0): DiT block 5 aligned to the encoder's 16×16 patch features of the decoded image (weight 0.5, 50% of each batch, first 25% of training, then switched off). The encoder is not part of this model.
• Optimiser: Muon (lr 3.5e-3, momentum 0.95, wd 0.01) for hidden matrices, AdamW for the rest; 2% warmup, then constant lr; EMA 0.9995 per step.
• Distributed: MuLoCo (DiLoCo with a Muon inner optimiser): each of 8 GPUs trains its own replica on a batch of 32, and the replicas are averaged every 64 steps (plain averaging). This avoids gradient all-reduce over PCIe 3.0 x4.
• bf16 autocast with fp32 weights, torch.compile.

Training code: https://github.com/adamo1139/trenowanie_generatora_obrazkow

Usage

import torch
from diffusers import DiffusionPipeline

REPO = "cpral/IMG-253M"
pipe = DiffusionPipeline.from_pretrained(REPO, custom_pipeline=REPO, trust_remote_code=True).to("cuda")
image = pipe("a red fox sitting in a snowy forest at sunrise, photograph",
             num_inference_steps=50, guidance_scale=3.0,
             generator=torch.Generator("cuda").manual_seed(0)).images[0]
image.save("out.png")

Fixed 256×256 output, English prompts.

License

Apache-2.0. text_encoder/ contains layers of SlayerLab/Slayer149 (Apache-2.0, see text_encoder/LICENSE and text_encoder/NOTICE). vae/: FLUX.2 small decoder © Black Forest Labs, Apache-2.0, unchanged. The pipeline code is based on AuroraIMG2-50M-Preview (Apache-2.0). Training data has its own terms, see the dataset cards.