---
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](https://huggingface.co/SlayerLab/Slayer149) language model** (first 18 of 19 blocks) through a
12.6M-parameter trainable refiner, decoding with the frozen
[FLUX.2 small decoder](https://huggingface.co/black-forest-labs/FLUX.2-small-decoder).
Trained in ~12.5 h on 8× RTX 3090 Ti. Released weights: EMA.

![samples](samples/grid.png)

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:

| model | val/loss |
| --- | --- |
| [AuroraIMG2-50M-Preview](https://huggingface.co/AuroraAI-Research/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](https://huggingface.co/datasets/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](https://huggingface.co/facebook/PE-Spatial-B16-512) (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

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