AuroraIMG2-100M

AuroraIMG2-100M is a 256×256 text-to-image model built around a 100.6M-parameter flow-matching diffusion transformer and a 44.9M-parameter text encoder. Both were trained from scratch on a single NVIDIA RTX PRO 6000. The model works in the FLUX.2 VAE latent space and turns latents into pixels with the frozen, pretrained FLUX.2 small decoder (Apache-2.0, bundled unchanged in vae/).

It is the successor of AuroraIMG2-50M-Preview and scores 10.80 Soft-TIFA GM and 48.62 Soft-TIFA AM on GenEval 2 (see Results), up from 6.00 / 36.03 for the Preview.

At a glance

Quick start

This repository uses a custom pipeline, so trust_remote_code=True is required.`

import torch
from diffusers import DiffusionPipeline

device = "cuda" if torch.cuda.is_available() else "cpu"

pipe = DiffusionPipeline.from_pretrained(
    "AuroraAI-Research/AuroraIMG2-100M",
    custom_pipeline="AuroraAI-Research/AuroraIMG2-100M",
    trust_remote_code=True,
).to(device)

generator = torch.Generator(device=device).manual_seed(0)
image = pipe(
    prompt="Vintage Yellow Enamel Tea Kettle with Black Handle and Spout",
    num_inference_steps=50,
    guidance_scale=4.0,
    generator=generator,
).images[0]
image.save("auroraimg2.png")

Pass a list of prompts for batches, or num_images_per_prompt for several images per prompt.

Recommended settings

The benchmark numbers were measured with fp32 weights under bfloat16 autocast, which is what the pipeline applies automatically (the same numerics as in training). Loading with torch_dtype=torch.bfloat16 also works and uses less memory, but it was not benchmarked separately.

Samples

36 fixed prompts, seed 0, 50 steps, CFG 4.0. These are the first samples generated, not a cherry-picked selection:

How it works

Tokenizer and text encoder (44.9M). A byte-level BPE tokenizer (16384 tokens) trained on the training captions feeds a 12-layer pre-norm bidirectional transformer (width 512, 8 query heads with 2 key/value heads, head dim 64) with rotary position embeddings, SwiGLU MLPs (hidden 1536) and RMSNorm. Prompts are truncated to 190 tokens plus start and end tokens (192 total). The encoder was first pretrained on its own with masked language modelling and then trained jointly with the denoiser. A learned null prompt (192 × 512) is the unconditional input for classifier-free guidance.

Denoiser (100.6M). A deep-thin DiT with 24 blocks of width 512 and 8 attention heads (head dim 64). Each latent pixel of the 16×16 latent grid is one token, so there is no patchify step. Every block has:

• self-attention with 2D rotary embeddings and QK-norm,
• cross-attention to the text tokens (8 query heads, 2 key/value heads, padding-masked, zero-initialised output),
• a Mix-FFN: a SwiGLU feed-forward layer (hidden 1536) with a zero-initialised depthwise 3×3 convolution on the value branch.

Timestep conditioning uses PixArt-style adaLN-single (one shared projection plus a learned, zero-initialised modulation table per block). The conditioning vector is the timestep embedding plus an MLP of the masked-mean-pooled text features. U-ViT-style long skip connections link block i (for the second half of the network) to block 23 − i through learned per-channel gates initialised to zero.

Decoder (frozen). The pretrained FLUX.2 small decoder maps the 128×16×16 latents (after undoing the dataset's BatchNorm normalisation and the 2×2 channel packing) to a 256×256 image. It is stored in vae/ exactly as published by Black Forest Labs, and was not trained or modified here.

Sampling. Euler integration of the learned velocity field from t = 1 (noise) to t = 0 on a shifted time grid (t = s·u / (1 + (s − 1)·u) with shift s = 1.5), with classifier-free guidance against the learned null prompt, switched off for t ≤ 0.05.

Training

Data. QLNI/FLUX-Reason-6M-flux2-latents provides precomputed FLUX.2 latents of 256px center crops of the 5.6M images of FLUX-Reason-6M, with several captions per image. At each step one caption per image is sampled: 80% from the longest descriptive caption and 20% from the shorter ones. No augmentation was used (flipping would corrupt left/right words). One shard of each of the four subsets was held out for validation. See the dataset card and FLUX-Reason-6M for provenance and licensing.

Stage 1: text-encoder pretraining. 16,000 steps of masked language modelling at batch 1024 (about 16.4M caption sequences): 25% masking with span masking (mean span 2.5 tokens) and a boosted masking rate for counts, positions, colours, sizes and materials. The final held-out masked-token accuracy was 65.5%.

Stage 2: joint training (235,000 steps, batch 512).

• Objective. Velocity prediction with an MSE loss on z_t = (1 − t)·z + t·ε with target ε − z. Timesteps come from a mixture of 90% logit-normal (mean 0, std 1) and 10% uniform, so low-noise steps are not starved. 10% of captions are replaced by the learned null prompt.
• Representation alignment. A Self-Representation Alignment loss with no external encoder: block-6 features of the denoiser are matched to block-16 features of an EMA copy of itself at a slightly lower noise level (smooth-L1 on layer-normalised features). Weight 0.2 from step 2,000, annealed to zero between steps 90,000 and 150,000.
• Text encoder. Frozen for the first 3,000 steps so the denoiser learns to read the features before it can distort them, then trained at half the learning rate with a small auxiliary masked-language-modelling loss (weight 0.05, every 4th step, 15% masking).
• Optimisation. NorMuon (lr 1e-3, momentum 0.95) on the hidden weight matrices of both networks; AdamW (betas 0.9 / 0.95) on the token embedding (lr 5e-4) and on the input/output projections, norms, tables and biases (lr 3e-4). Weight decay 0.01 on the Muon matrices and the input/output projections. Warmup-stable-decay learning-rate schedule: 3,000 warmup steps, constant until step 188,000, then linear decay to zero. Gradient clipping at 1.0, fp32 weights with bfloat16 autocast, torch.compile.
• EMA. Three EMAs (0.999, 0.9995, 0.9998) were tracked. The EMA 0.999 weights are the ones released here.

Budget. 235,000 steps at batch size 512, about 120M training samples (roughly 21 passes over the 5.6M images), on 1× NVIDIA RTX PRO 6000, around 80 Hours

Results

All numbers below are GenEval 2 (Kamath et al., 2025): 800 prompts scored with Soft-TIFA using the benchmark's official evaluation code (Qwen3-VL-8B as the VQA model). GM (geometric mean over the atoms of a prompt, averaged over prompts) measures prompt-level correctness and AM (arithmetic mean) measures atom-level correctness. All scores are in percent and higher is better.

AuroraIMG2-100M was run at 256×256 with 50 Euler steps, CFG 4.0, time-grid shift 1.5, EMA 0.999 weights, and one image per prompt with fixed seeds.

Per-skill breakdown (Soft-TIFA AM, atom level):

Per-atomicity breakdown (Soft-TIFA GM, prompt level; atomicity is the number of objects, attributes, counts and relations in a prompt, 100 prompts each):

Notes on the comparison.

• Parameter counts are for the denoiser only (U-Net or DiT). The competitors' text encoders and VAEs are separate pretrained models that are not counted. AuroraIMG2's from-scratch text encoder adds 44.9M, and its frozen decoder has 27.9M decoder parameters.
• Against the 50M-Preview, GM rises from 6.00 to 10.80 and AM from 36.03 to 48.62. The biggest relative gains are in counting (17.52 → 33.26) and position (18.16 → 35.11).
• Compared with Stable Diffusion XL's numbers from Table 4 of the GenEval 2 paper, AuroraIMG2-100M is higher on prompt-level GM (10.80 vs 9.1) and slightly lower on atom-level AM (48.62 vs 50.1). SDXL's U-Net has 2.6B parameters and its text encoders add 817M. The paper does not report per-skill Soft-TIFA scores for SDXL and SD 2.1, so they are missing from the per-skill table.
• Supra2-IMG and Surjo-Image-Preview numbers are as reported on the Surjo-Image-Preview model card. They were not re-run here, and sampling settings, seeds and image counts differ between sources, so small gaps should not be over-interpreted.
• Much larger models score far higher on this benchmark (for example Qwen-Image reaches 33.8 GM and 80.8 AM in the GenEval 2 paper). This model is a small-scale research model, not a competitor to them.
• "verb" is very low for every small model. The GenEval 2 authors note that the VQA judge is stricter than human annotators on this skill.

Good for

• Studying small-scale text-to-image training: flow matching, NorMuon, adaLN-single DiTs, self-representation alignment and a separately pretrained small text encoder, all on a single GPU.
• Fast experiments and fine-tuning on modest hardware (the trained parts are only 145.5M parameters).

Limitations

• Fixed 256×256 output. Small details, faces, hands and rendered text were not evaluated and are expected to be weak at this scale.
• Compositional prompts are hard for it. Verbs are almost never correct (3.18 AM), counting and spatial relations are about a third correct at the atom level, and prompt-level correctness falls from 12.6–20.7 GM at 3–5 atoms to 2.4 GM at 10 atoms.
• English only, 190 caption tokens at most.
• No negative prompts. The unconditional branch is a learned null prompt.
• Data bias. The model inherits the biases and artifacts of its training data. No separate safety evaluation of outputs has been done for this release.
• Not for generating realistic images of real people or events, or for settings where accuracy and fairness matter.

License

The model weights (text_encoder/, transformer/) and code are released under the Apache 2.0 license. The bundled FLUX.2 small decoder in vae/ is © Black Forest Labs and is also released under Apache 2.0. The training data has its own terms: check the dataset cards before using the model in a commercial setting.

Acknowledgements

• Black Forest Labs for the FLUX.2 VAE and small decoder.
• QLNI for the precomputed latents and the FLUX-Reason-6M authors for the underlying dataset.
• The GenEval 2 authors for the benchmark and the Soft-TIFA evaluation code.
• The authors of NorMuon, Self-Representation Alignment, SR-DiT and the Photoroom PRX training-design write-up whose findings shaped the training recipe.

Citation

@misc{auroraimg2_100m,
  title        = {AuroraIMG2-100M: a 100M-parameter text-to-image model},
  author       = {AuroraAI-Research},
  year         = {2026},
  howpublished = {\url{https://huggingface.co/AuroraAI-Research/AuroraIMG2-100M}}
}