Supra2-IMG is a tiny 100M parameters text-to-image (T2I) model that has been trained from scratch on high-quality synthetic data and delivers state-of-the-art image quality for its size.
Samples
Model
About the model
The model is a tiny diffusion transformer (DiT) with ~105M parameters.
- Pipeline: text-to-image
- Parameter Count: 104.1M
- Encoder: frozen Flan-T5-Base
- VAE: SD-VAE-FT-MSE
- Image resolution: 256²
- Latent size: 32²
- Patch: 2
- Context length (Flan-T5): 128 tokens
Model config
D_MODEL: 576DEPTH: 14N_HEADS: 9HEAD_DIM: 64MLP_RATIO: 4.0D_CTX: 768VAE_SCALE: 0.18215
Training
Dataset
The model was trained for 10 epochs on the full LucasFang/FLUX-Reason-6M dataset.
Data preparation
All train data images were downloaded as parquets + metadata and prepared by first chosing the prompt.
This was done in the following order (each next prompt is a fallback for the previous prompt): caption_composition → caption_entity → caption_text → caption_style → caption_imaginative.
That way, we ensured only using the highest quality data for pretraining the model.
Exact image count
5.6M images
Epochs count
10 epochs
Hardware
The training ran on a single Nvidia H100 SXM 80GB Runpod Pod for 9 hours (incl. data preparation) with a 2.5TB disk.
Benchmarks
FID ↓ (Lower is better): 27.9
CLIP ↑: 29.4
How to run the model
First, run:
# Create project directory
mkdir Supra2-IMG
cd Supra2-IMG
# Download the inference script
wget https://huggingface.co/SupraLabs/Supra2-IMG/resolve/main/inference.py
Then, you can generate images by running:
python inference.py --prompt "a sea jellyfish floating in the pitch-black ocean depths" --seed 0 --cfg 3.0 --steps 50 --n 1 --out jellyfish.png
Recommended settings for sampling
--seed: 0--cfg: 3.0--steps: 50
Expected Output
The script will output something like:
=== Supra2-IMG inference ===
[device] ...
[ckpt] found ./model_final_ema.pt
[model] building SupraDiT ...
[model] 104.1M parameters
[model] loading weights from ./model_final_ema.pt ...
[model] weights loaded in 0.7s
[text] ctx_len=128
[text] loading tokenizer + google/flan-t5-base ...
...
[text] prompt tokens=15 n=1 seed=0 cfg=3.0 steps=50
[cfg] using stored unconditional embeddings
[sample] Euler flow, 50 steps ...
Generating: ...
[sample] denoising done in ...s
[vae] decoding latents ...
[done] saved 1 image(s) -> ...png
Possible output
Acknowledgments
- Thank you, FLUX-Reason-6M team, for proving the LucasFang/FLUX-Reason-6M dataset. Without your work, this would never be possible!
- Thanks to the creator of HobbyLM-Image for inspiration
- Thanks to Runpod for providing the compute!
- Also a big Thank You To Stability AI and Google for providing SD-VAE and Flan-T5
- Thanks to all the other people inspiring us to make great open-source models for the community!
What's next
We will keep improving Supra2-IMG, maybe for a next-gen like Supra2.5-IMG, and we will share our progress and findings on the way to the best open-source T2I model 🤗 Please give us a like and a follow on Hugging Face if you want to support our work!