HomeIMG-641M-epoch15

Experimental 256×256 text-to-image model trained locally from scratch on 8 consumer GPUs: a 641M-parameter flow-matching DiT conditioned on a frozen small language model (SlayerLab/Slayer149) through a 35.7M-parameter trainable refiner, decoding through the frozen FLUX.2 small decoder. 15 epochs over 5.88M images (88.2M samples). The successor of HomeIMG-50M and HomeIMG-253M; an intermediate checkpoint of the same run (10.8 epochs) is HomeIMG-641M-epoch10.8.

Samples

Eight prompts from the AuroraIMG2 card, each with seed 0, 50 Euler steps, CFG 3.0, decoder in bfloat16. Row-major order:

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

Future work

Next release will be bilingual with support for both English and Polish. I will continue training this DiT but will swap out TE and will retrain the refiner. Then, I'll try to distill bigger models or bring the performance up by training on 512x512 imagery or doing preference optimization.

Evaluation

FD-DINOv2 (Stein et al. 2023, lower is better): Fréchet distance between DINOv2 CLS features of generated and real images. 4,096 held-out captions (the first caption of the first 4,096 images of the 4 held-out shards, one per dataset subset), 50 Euler steps, CFG 3.0, fixed seeds; features from vit_base_patch14_reg4_dinov2.lvd142m; the reference is the real held-out latents decoded by the same FLUX.2 decoder. All models below were measured with the same code on the same captions. The absolute values are not comparable with FD-DINOv2 numbers from papers (different encoder, dataset and sample count).

model	denoiser	training samples	FD-DINOv2
IMG-50M (WSM)	50.4M + 12.6M text encoder from scratch	63.4M	69.61
IMG-253M (EMA)	253M + Slayer149	29.4M	37.64
IMG-641M-epoch10.8 (EMA)	641M + Slayer149	63.3M	29.50
IMG-641M-epoch15, EMA weights	641M + Slayer149	88.2M	28.27
IMG-641M-epoch15 (this repo, WSM)	641M + Slayer149	88.2M	27.66

Held-out velocity MSE (same 4,096 images, first caption, fixed noise and t ~ logit-normal(0, 1)):

model	val/loss
AuroraIMG2-50M-Preview	0.6929
HomeIMG-50M (WSM)	0.7044
HomeIMG-641M-epoch15, raw weights at the last step	0.6329
HomeIMG-641M-epoch15, EMA	0.6304
HomeIMG-641M-epoch15 (this repo, WSM)	0.6297

Architecture

• Denoiser (641.1M): DiT, 6 blocks, width 2560, 20 heads (head dim 128), SwiGLU MLP with hidden size 6912, one token per latent pixel (16×16 grid, no patchify). Each block has self-attention with 2D RoPE and QK-norm, cross-attention to the text tokens and a SwiGLU MLP; timestep conditioning is PixArt-style adaLN-single. The wide and shallow shape (width/depth ≈ 427, against ~40–80 for common DiTs) was chosen for throughput on this hardware: one model per GPU with large matrix multiplies. It is outside the range of published DiT shape studies.
• Text encoder: SlayerLab/Slayer149 (148.9M, English, 12,288-token vocabulary), frozen; the residual stream before its last block (18 of 19 blocks) feeds a trainable refiner (35.7M): input projection 768 → 1280, 2 bidirectional transformer blocks (width 1280, 20 query / 5 KV heads, 1D RoPE, SwiGLU 3456), RMSNorm, and a learned null embedding for classifier-free guidance. Up to 95 tokens + BOS.
• Decoder (frozen): FLUX.2 small decoder. Latents are 128×16×16 (BatchNorm-normalised, 2×2-packed FLUX.2 latents).

Training

• Data: QLNI/FLUX-Reason-6M-flux2-latents, 5.88M training images (all subsets). One random caption per image per step. 15 epochs = 88.2M samples.

• Objective: rectified flow, z_t = (1−t)·x + t·ε, velocity target ε − x, MSE, t ~ logit-normal(0, 1); 10% of captions replaced by the null prompt.

• Hardware: 8× RTX 3090 Ti (6 on PCIe 3.0 x4, 2 on x8), local hardware, bf16 autocast, torch.compile. 50–100 TFLOPS per GPU depending on the phase (up to ~70% of the dense bf16 peak), about two days of training in total.

• MuLoCo (DiLoCo with a Muon inner optimiser, arXiv:2505.23725): every GPU trains its own replica on its own share of the batch; every 64 steps the 8 replicas are averaged (outer lr 1.0, momentum 0).

• Optimiser: Muon on all hidden weight matrices with two DiT-specific variants: CMuon (arXiv:2608.02502) orthogonalises the blocks of fused matrices (q/k/v, cross-attention k/v, SwiGLU gate/up) separately; periodic row-wise Muon (arXiv:2608.20818, Newton–Schulz every 3rd step, row-normalised momentum in between) in the first two thirds, full Newton–Schulz every step afterwards. Momentum 0.95, weight decay 0.01. AdamW on the rest.

• REPA: iREPA (arXiv:2512.10794; REPA, arXiv:2410.06940): the hidden state after block 2 is aligned (3×3 conv projector) with spatially normalised DINOv2-B/14 (registers) patch features of the decoded image, on 25% of each micro-batch, weight 0.5. Switched off for the last 20% of training once the alignment loss had saturated (in line with HASTE).

• Run history: the run was resumed several times with different settings; the data stream always continued from the same image:

  steps	batch	Muon lr	REPA
  0 – 2,112	4096	1e-2	iREPA until 1,435
  2,112 – 34,944	512	3e-3	off
  34,944 – 45,056	1024	3e-3	iREPA
  45,056 – 48,128	4096	1e-2	every 32nd step
  48,128 – 69,888	1024	3e-3	iREPA
  69,888 – 86,997	1024	3e-3	off

  The jump back to lr 1e-2 at step 45,056 temporarily worsened FD-DINOv2 (33.5 → 43.2) while barely moving the loss; returning to 3e-3 with iREPA recovered it within a few million samples.

• Schedule and released weights: 2% warmup, then constant lr with an EMA of the averaged weights. Released weights: WSM (Warmup-Stable-Merge, arXiv:2507.17634): uniform average of the last 4 checkpoints (steps 77,440 / 80,960 / 84,480 / 86,997, the last 12% of training), chosen over EMA and other WSM/QAM variants by held-out loss and FD-DINOv2.

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

Usage

import torch
from diffusers import DiffusionPipeline

REPO = "cpral/IMG-641M-epoch15"  # this repository
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, up to 95 tokens.

Limitations

A research model: 256×256 only, trained on ~6M synthetic (FLUX-generated) images, far below the data and compute of production text-to-image models, and with an unusually shallow architecture. Expect weak text rendering, limited compositional accuracy and the biases of its training data.

License

Weights and code: Apache-2.0. Text encoder: Slayer149 (Apache-2.0), unchanged blocks plus the trained refiner. vae/: FLUX.2 small decoder © Black Forest Labs, Apache-2.0, unchanged. Training data has its own terms, see the dataset card.