AuroraIMG2-50M-Preview

AuroraIMG2-50M-Preview is a 256×256 text-to-image model built around a 50.4M-parameter flow-matching diffusion transformer and a 12.6M-parameter text encoder. Both were trained from scratch in 10 hours 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/).

This is a Preview release. On GenEval 2 it scores 6.00 Soft-TIFA GM and 36.03 Soft-TIFA AM (see Results).

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-50M-Preview",
    custom_pipeline="AuroraAI-Research/AuroraIMG2-50M-Preview",
    trust_remote_code=True,
).to(device)

generator = torch.Generator(device=device).manual_seed(0)
image = pipe(
    prompt="A tiny wooden house on top of a giant red mushroom in an enchanted forest, warm glowing windows, soft morning sunlight, little flowers and moss, magical storybook illustration, vibrant colors, charming and detailed.",
    num_inference_steps=50,
    guidance_scale=3.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 and the decoder in bfloat16. The pipeline applies autocast automatically on CUDA, and pipe.vae.to(torch.bfloat16) reproduces the decoder precision. Loading with torch_dtype=torch.bfloat16 also works and uses less memory, but it was not benchmarked separately.

Samples

Eight fixed prompts, seed 0, 50 steps, CFG 3.0. These are the first samples generated, not a cherry-picked selection. Row-major order, left to right:

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

How it works

Tokenizer and text encoder (12.6M). A byte-level BPE tokenizer (8192 tokens) trained on the training captions feeds a 6-layer pre-norm transformer (width 384, 6 query heads with 2 key/value heads, head dim 64) with rotary position embeddings, SwiGLU MLPs and RMSNorm. It is trained jointly with the denoiser, so it learns whatever representation the image model finds most useful. Prompts are truncated to 62 tokens plus start and end tokens.

Denoiser (50.4M). A DiT with 8 blocks of width 640 and 10 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, and a SwiGLU feed-forward layer. Timestep conditioning uses PixArt-style adaLN-single: one shared time projection plus a learned modulation table per block, zero-initialised so every block starts as the identity.

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. The encoder half of that checkpoint is included but not used at inference.

Sampling. Plain Euler integration of the learned velocity field from t = 1 (noise) to t = 0 on a uniform grid, with classifier-free guidance against the learned empty prompt.

Training

Data. QLNI/FLUX-Reason-6M-flux2-latents provides precomputed FLUX.2 latents of 256px center crops, with several captions per image (one is sampled at each step). A few shards were held out for evaluation and the rest were used for training, streamed through a rotating GPU-resident pool of shards. See the dataset card and FLUX-Reason-6M for provenance and licensing.

Objective. Velocity prediction with an MSE loss on z_t = (1 - t)·z + t·ε with target ε - z. Timesteps are sampled from a logit-normal distribution (mean 0, std 1). 10% of captions are replaced by a learned null prompt for classifier-free guidance. The text encoder also receives an auxiliary masked-language-modelling loss (weight 0.05, applied every 4th step, 15% masking, tied to the token embedding).

Optimisation. Muon (lr 1e-2, momentum 0.95, 5 Newton–Schulz steps) on the hidden weight matrices and AdamW (betas 0.9 / 0.95) on the token embedding (lr 5e-4) and the input/output projections, norms, tables and biases (lr 3e-4). Weight decay 0.01 on the Muon matrices and the input/output projections. Linear warm-up over the first 2% of training followed by cosine decay to zero, gradient clipping at 1.0, bfloat16 autocast and torch.compile. Weights are tracked with an EMA (decay 0.9995, with warm-up) and the EMA weights are the ones released here.

Budget. 10 hours of training time on 1× NVIDIA RTX PRO 6000: 247,805 steps at batch size 256, about 63.4M training samples seen. Evaluation, sample generation and compilation warm-up were not counted against the budget.

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-50M-Preview was run at 256×256 with 50 Euler steps, CFG 3.0 and one image per prompt.

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

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 12.6M, and its frozen decoder has 27.9M decoder parameters.
• Stable Diffusion 2.1 and SDXL numbers are taken from Table 4 of the GenEval 2 paper. 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 these two models, 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 (256px, 50 steps, CFG 3.0, seed 42, one image per prompt). They were not re-run here.
• Sampling settings, seeds and image counts differ between sources, so small gaps should not be over-interpreted.
• "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, Muon, adaLN-single DiTs and joint text-encoder training, all on a 10-hour budget.
• Fast experiments and fine-tuning on modest hardware (the trained parts are only 63.0M).

Limitations

• Fixed 256×256 output. Small details, faces, hands and rendered text are weak.
• Compositional prompts are hard for it. Counting, spatial relations and verbs are the weakest skills (see the per-skill table).
• English only, 62 caption tokens at most.
• Data bias. The model inherits the biases and artifacts of its training data. No separate safety evaluation of outputs has been done for this preview.
• 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.

Exported with diffusers 0.42.1, transformers 5.21.0, torch 2.14.1+cu130.

Citation

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