SupraLabs/Supra2-IMG

🤗 Hugging Face sourcetext-to-imageapache-2.0417 MBother✓ 1 checksumupdated today
Submit in one command

Run it next to your model folder. It makes the torrent, checks your files against Hugging Face, and submits it. You just start seeding and paste your key from your account. It only reads your files and never changes them. Read the script first if you like.

curl -fsSL https://pirateface.co/package.sh | bash -s -- --repo SupraLabs/Supra2-IMG ./model-folder
Needs a seeder →

Supra2-IMG

Text-To-Image • 100M Parameters • SOTA quality

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: 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!