license: apache-2.0
library_name: diffusers
pipeline_tag: text-to-image
language:
- en
datasets:
- QLNI/FLUX-Reason-6M-flux2-latents
base_model:
- SlayerLab/Slayer149
tags:
- text-to-image
- flow-matching
- rectified-flow
- diffusion-transformer
- from-scratch
- muon
- diloco
- repa
- custom_code
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:
- a red fox sitting in a snowy forest at sunrise, photograph
- a futuristic city skyline at night with neon lights reflected in the rain
- a bowl of fresh strawberries on a wooden table, soft natural light
- a watercolor painting of a lighthouse on a rocky coast
- a small robot reading a book in a cozy library
- an astronaut riding a horse on the moon, digital art
- a street sign that says OPEN in neon letters
- 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:
| model | val/loss |
|---|---|
| AuroraIMG2-50M-Preview | 0.6929 |
| IMG-253M (EMA, released) | 0.6499 |
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.