Supra2-IMG — single-file safetensors

The three components of SupraLabs/Supra2-IMG — DiT, text encoder and VAE — merged into one .safetensors file. No need to download google/flan-t5-base and stabilityai/sd-vae-ft-mse separately anymore.

File	Size	Precision
`Supra2-IMG-fp32.safetensors`	1.19 GB	all float32 — lossless, values identical to the original checkpoint
`Supra2-IMG-fp16.safetensors`	0.81 GB	DiT + VAE float16, text encoder float32 — drop-in, no extra steps

598 tensors · 297.5M parameters in total.

Key layout

Prefix	Component	Params	fp32 file	fp16 file	Source
`dit.`	SupraDiT (EMA weights)	104.1M	float32	float16	`model_final_ema.pt`
`text_encoder.`	T5 encoder	109.6M	float32	float32	`google/flan-t5-base`
`vae.`	AutoencoderKL	83.7M	float32	float16	`stabilityai/sd-vae-ft-mse`
`cfg.`	`uncond_text` `[128, 768]`, `uncond_mask` `[128]`	0.1M	float32	float16	checkpoint config

Metadata

The safetensors header carries everything needed to rebuild the model without any config.json:

arch — architecture constants (d_model=576, depth=14, n_heads=9, patch=2, latent_size=32, vae_scale=0.18215, ctx_len=128, ...)
dit_config, text_encoder_config, vae_config — component configs as JSON
tied_weights — map of deduplicated shared tensors
tokenizer — tokenizer pointer (not embeddable, see Notes)

Why the fp16 file is 0.81 GB and not 0.60 GB

The text encoder is deliberately stored in float32. T5 overflows in float16: its hidden states drift far from the fp32 result (cosine similarity 0.892, max abs difference 1.38), which produces wildly oversaturated, confetti-like images that no longer follow the prompt.

Measured on a Tesla T4, same prompt and seed, each component in fp16 against an all-float32 reference:

Component in fp16	Mean abs diff vs fp32	Verdict
DiT	0.0001	safe
VAE	0.0001	safe
T5 text encoder	0.1849	broken — never use fp16

An all-fp16 variant was published first and had to be withdrawn for this reason. Both files in this repo now keep the text encoder in float32, so just load them and run — there is nothing to work around.

Load

import json
import torch
from safetensors import safe_open
from safetensors.torch import load_file
from transformers import AutoTokenizer, T5Config, T5EncoderModel
from diffusers import AutoencoderKL

PATH = "Supra2-IMG-fp16.safetensors"

sd = load_file(PATH)
with safe_open(PATH, framework="pt") as f:
    meta = f.metadata()

arch = json.loads(meta["arch"])
dt = sd["dit.x_embed.weight"].dtype   # torch.float16 for the fp16 file

# --- text encoder: already float32 in both files ---
text_encoder = T5EncoderModel(T5Config.from_dict(json.loads(meta["text_encoder_config"])))
te_sd = {k[len("text_encoder."):]: v for k, v in sd.items() if k.startswith("text_encoder.")}
te_sd["encoder.embed_tokens.weight"] = te_sd["shared.weight"]   # tied weight
text_encoder.load_state_dict(te_sd, strict=True)

# --- VAE ---
vae = AutoencoderKL.from_config(json.loads(meta["vae_config"]))
vae.load_state_dict({k[len("vae."):]: v for k, v in sd.items() if k.startswith("vae.")}, strict=True)
vae = vae.to(dt)

# --- DiT: use the SupraDiT class from the original repo's inference.py ---
model = SupraDiT().to(dt)
model.load_state_dict({k[len("dit."):]: v for k, v in sd.items() if k.startswith("dit.")}, strict=True)

# --- stored unconditional embeddings for CFG ---
uncond_text = sd["cfg.uncond_text"].float()   # [128, 768]
uncond_mask = sd["cfg.uncond_mask"].float()   # [128]

tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-base")

device = "cuda"
model.to(device).eval()
text_encoder.to(device).eval()
vae.to(device).eval()

Generate

import torch
import torchvision.utils as vutils

prompt = "a sea jellyfish floating in the pitch-black ocean depths"
steps, cfg_scale, seed = 50, 3.0, 0

tok = tokenizer([prompt], padding="max_length", truncation=True,
                max_length=arch["ctx_len"], return_tensors="pt")
tok = {k: v.to(device) for k, v in tok.items()}

with torch.no_grad():
    # text encoder stays fp32, outside autocast
    ctx = text_encoder(**tok).last_hidden_state.float()
    mask = tok["attention_mask"].float()

    g = torch.Generator(device=device).manual_seed(seed)
    z = torch.randn(1, arch["latent_ch"], arch["latent_size"], arch["latent_size"],
                    device=device, generator=g)
    step = 1.0 / steps

    # DiT + VAE are safe in fp16
    with torch.autocast("cuda", dtype=torch.float16):
        for i in range(steps):
            t = torch.full((1,), i * step, device=device)
            v = model(torch.cat([z, z]), torch.cat([t, t]),
                      torch.cat([ctx, uncond_text[None].to(device)]),
                      torch.cat([mask, uncond_mask[None].to(device)])).float()
            v_cond, v_uncond = v.chunk(2, 0)
            z = z + step * (v_uncond + cfg_scale * (v_cond - v_uncond))

        img = vae.decode(z / arch["vae_scale"]).sample.float()

vutils.save_image((img.clamp(-1, 1) + 1) / 2, "out.png")

Notes

The tokenizer is not embedded. spiece.model is a binary file and cannot live inside a safetensors header. Load it with AutoTokenizer.from_pretrained("google/flan-t5-base").
text_encoder.encoder.embed_tokens.weight is tied to text_encoder.shared.weight and is stored only once. Restore the alias on load (one line, see above); the mapping is recorded in tied_weights.
Never cast the text encoder to fp16. A blanket model.half() over the whole file reintroduces the oversaturation bug described above.
Sampling settings: 256x256 output, 32x32 latent, patch size 2, Euler flow, 50 steps, CFG 3.0, vae_scale = 0.18215. Fewer than ~40 steps gets noticeably blurry on this 104M model.

Source

DiT + config: SupraLabs/Supra2-IMG (Apache-2.0)
Text encoder: google/flan-t5-base (Apache-2.0)
VAE: stabilityai/sd-vae-ft-mse

Merge produced on Google Colab (Tesla T4), verified by loading every component with strict=True, by matching sha256 against the local build, and by generating images from the merged file alone.

Downloads last month: -; Downloads are not tracked for this model. How to track