AuroraIMG-6M

AuroraIMG-6M is a text-to-image model that's just 6.82M parameters, and that number includes everything: the text encoder, the denoiser and the VAE. It generates 64×64 images from short English captions, and every part of it was trained from scratch on CC12M. No pretrained text encoder or VAE was used.

It reaches an FID of 33 and a CLIP of 24.8 on a 10k-image CC12M holdout set.

At a glance

Quick start

from diffusers import DiffusionPipeline
import torch

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

pipe = DiffusionPipeline.from_pretrained(
    "AuroraAI-Research/AuroraIMG-6M",
    custom_pipeline="AuroraAI-Research/AuroraIMG-6M",
    trust_remote_code=True,
    torch_dtype=dtype
).to(device)

prompt = "a sunset over the ocean"
print(f"Generating: '{prompt}'...")

generator = torch.Generator(device=device).manual_seed(42)
image = pipe(
    prompt=prompt,
    num_inference_steps=50,
    guidance_scale=6.0,
    generator=generator
).images[0]

filename = "Aurora-Image.png"
image.save(filename)
print(f"\n Saved to {filename}!")

The model is small enough to run on CPU.

Samples

Same eight prompts, same seed, increasing guidance (CFG 1, 2, 3 and 5, top to bottom):

Final checkpoint at CFG 3 for the fixed prompt set ("a red car on a street", "a dog running on green grass", "a sunset over the ocean", "a snowy mountain with blue sky", "a plate of food on a table", "a man standing in front of a building", "a small boat on a calm lake", "a vase of flowers on a wooden table"):

Real COCO val2014 images (top row) next to generated images from their captions (bottom row):

And the VAE on its own: originals on top, reconstructions below.

How it works

VAE (1.08M). A small convolutional VAE with residual blocks and GroupNorm that squeezes a 64×64×3 image into an 8-channel 16×16 latent. It was trained with an L1 loss, a gradient-based sharpness loss, a Gram-matrix loss and a tiny KL term (5e-5). There's no adversarial or LPIPS loss. On holdout images it reconstructs at 31.9 dB PSNR.

Text encoder (1.32M). A 4-layer pre-LN transformer (width 128, 4 heads) with rotary position embeddings, sitting on a byte-level BPE tokenizer trained on the CC12M captions. It's trained jointly with the denoiser, so it learns whatever representation the image model finds most useful.

Denoiser (4.42M). A compact U-Net that works on the latent at 16×16, 8×8 and 4×4. Residual blocks use AdaLN-Zero conditioning (scale, shift and gate from the time embedding, zero-initialised so each block starts as identity). The time embedding uses random Fourier features, with a mean-pooled text embedding added on top. Each level has an attention block with self-attention, cross-attention to the text tokens and a feed-forward layer.

Training

Data: The whole CC12M set (pixparse/cc12m-wds) was streamed and filtered. Captions were lowercased and cleaned of HTML and URLs, and generic ones ("image", "stock photo", …) or anything under 3 words were dropped. Images under 128px on the short side, or not RGB, were dropped too. The rest were resized to 64px on the short side (bicubic), center-cropped, and de-duplicated by pixel hash. That left 9,882,283 training pairs plus a 50k holdout. Random horizontal flips are the only augmentation.

Stage 1: VAE. 25k steps at batch size 2048 with AdamW (lr 5e-4, cosine decay). Final loss 0.0880 (L1 0.0343, KL 5.91).

Stage 2: flow matching. 300k steps at batch size 2048, which is roughly 614M samples or ~62 epochs. The model predicts velocity (ε − z) with an MSE loss, and timesteps come from a logit-normal distribution mixed with 15% uniform. Cached latents are re-noised with the VAE's own posterior variance during training. Weight matrices and conv kernels are optimised with Muon (lr 0.02), and everything else (embeddings, input/output layers, norms, biases, time and pool projections) with AdamW (lr 5e-4). On top of that: 10k warmup steps, cosine decay to 1%, gradient clipping at 1.0, EMA 0.9998, 10% caption dropout for CFG, bf16 autocast and torch.compile. Final training loss: 0.6208. All samples and metrics below use the EMA weights.

Results

Everything for AuroraIMG-6M is measured at the native 64×64 resolution. FID uses torchmetrics (2048-d Inception features) and CLIP score uses openai/clip-vit-base-patch32. Both upsample the images internally, so treat them as numbers for tracking this model rather than for lining up against 256px+ models.

CC12M holdout (10k images)

I swept CFG from 1.0 to 9.0, and 6.0 gave the best results, achieving nearly the maximum CLIP it can, without sacrificing FID, so that's the recommended setting.

The generated images land at about 88% of the CLIP score of the real photos (24.80 vs 28.25) when compared against their captions.

VAE only

GenEval

Overall score (averaged over tasks): 0.1387. 12.88% of images and 20.61% of prompts were fully correct (2212 images, 553 prompts).

Comparison with Stable Diffusion 1.5

GenEval 2

AuroraIMG-6M was run at CFG 6.0, 50 Euler steps, 4 samples per prompt.

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

Prompting tips. It was trained on CC12M alt-text, so short, lowercase, caption-style prompts work best ("a dog running on green grass", "a snowy mountain with blue sky").

Good for

• Experimenting with tiny generative models: Muon, flow matching, AdaLN-Zero, small VAEs, joint text-encoder training.
• Running text-to-image on very limited hardware.

Things to keep in mind

• Resolution. Output is 64×64, so small details like faces, hands and text aren't going to be visible.
• Prompt complexity. It does best with a single main subject and a simple scene. Counting, positions and multi-attribute prompts (see GenEval and GenEval 2) are where it falls short.
• Language. English only, lowercased, up to 62 caption tokens.
• Data bias. CC12M is web-scraped alt-text, so the model picks up its biases and its stock-photo look. Only the quality filters described above were applied, with no extra safety filtering of data or outputs.
• Not for generating realistic images of real people or events, or anywhere accuracy and fairness matter.

License

The model weights are released under the Apache 2.0 license. The model was trained on CC12M, which is distributed as image URLs and captions; the underlying images carry their own licenses, so check the dataset terms for your use case.

Citation

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