Supra2-IMG

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: 576
• DEPTH: 14
• N_HEADS: 9
• HEAD_DIM: 64
• MLP_RATIO: 4.0
• D_CTX: 768
• VAE_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 &arrowright; caption_entity &arrowright; caption_text &arrowright; caption_style &arrowright; 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.

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!