supra2 safetensors

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,2 @@
+__pycache__/
+*.py[cod]

README.md

@@ -1,3 +1,63 @@
 ---
 license: apache-2.0
+base_model: SupraLabs/Supra2-IMG
+pipeline_tag: text-to-image
+library_name: transformers
+datasets:
+- LucasFang/FLUX-Reason-6M
+language:
+- en
+tags:
+- safetensors
+- text-to-image
+- supra
 ---
+
+# Supra2-IMG safetensors
+
+Safetensors build of [SupraLabs/Supra2-IMG](https://huggingface.co/SupraLabs/Supra2-IMG). The announcement is [Massive Release: Supra2-IMG, a tiny 100M text-to-image model](https://www.reddit.com/r/StableDiffusion/comments/1wmg1pb/massive_release_supra2img_a_tiny_100m_texttoimage/) on r/StableDiffusion.
+
+The model is SupraLabs' work: a 104.1M-parameter text-to-image DiT, 256×256, trained from scratch for 10 epochs on [LucasFang/FLUX-Reason-6M](https://huggingface.co/datasets/LucasFang/FLUX-Reason-6M). The text encoder is frozen [Flan-T5-Base](https://huggingface.co/google/flan-t5-base). The VAE is [SD-VAE-FT-MSE](https://huggingface.co/stabilityai/sd-vae-ft-mse). License is Apache-2.0. Architecture and training details are on the original model card.
+
+## This repo
+
+The original checkpoint is `model_final_ema.pt`. This repo hosts that checkpoint converted to FP32 safetensors.
+
+`Supra2-IMG_fp32.safetensors` is a safetensors file, not a renamed PyTorch checkpoint. It holds 237 FP32 weight tensors plus the stored unconditional text embedding and mask used for classifier-free guidance. Metadata keeps `patch=2` and `ctx_len=128`.
+
+Flan-T5-Base and SD-VAE-FT-MSE are included under their own names. They are the original models, not part of Supra2. `inference.py` loads the DiT, the text encoder, the tokenizer, and the VAE from this repo.
+
+Files the script does not read were left out: the SentencePiece model, `tokenizer_config.json`, `special_tokens_map.json`, and the original repo `config.json`. The tokenizer vocab is in `flan-t5-base.tokenizer.json`.
+
+`generate.bat` and `generate.sh` take a prompt and an output filename.
+
+## Files
+
+| File | Role |
+| --- | --- |
+| `inference.py` | Sampling script |
+| `generate.bat`, `generate.sh` | Launchers |
+| `Supra2-IMG_fp32.safetensors` | DiT weights, FP32 |
+| `flan-t5-base.safetensors` | Text encoder weights |
+| `flan-t5-base.config.json` | Text encoder config |
+| `flan-t5-base.tokenizer.json` | Tokenizer |
+| `sd-vae-ft-mse.safetensors` | VAE weights |
+| `sd-vae-ft-mse.config.json` | VAE config |
+
+## Generate
+
+```bat
+generate.bat "a boat" boat.png
+```
+
+```bash
+./generate.sh "a boat" boat.png
+```
+
+A bare filename is written next to the script. A full path is used as given.
+
+```bash
+python inference.py --prompt "a boat" --seed 0 --cfg 3.0 --steps 50 --out boat.png
+```
+
+Recommended settings from the original model card: seed `0`, CFG `3.0`, steps `50`.

Supra2-IMG_fp32.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:89851684d1b01e7020732d3c860e7bc4f7000edb94447c22260daa6a5b06c1ca
+size 416599856

flan-t5-base.config.json

@@ -0,0 +1,57 @@
+{
+  "architectures": [
+    "T5ForConditionalGeneration"
+  ],
+  "d_ff": 2048,
+  "d_kv": 64,
+  "d_model": 768,
+  "decoder_start_token_id": 0,
+  "dropout_rate": 0.1,
+  "eos_token_id": 1,
+  "feed_forward_proj": "gated-gelu",
+  "initializer_factor": 1.0,
+  "is_encoder_decoder": true,
+  "layer_norm_epsilon": 1e-06,
+  "model_type": "t5",
+  "n_positions": 512,
+  "num_decoder_layers": 12,
+  "num_heads": 12,
+  "num_layers": 12,
+  "output_past": true,
+  "pad_token_id": 0,
+  "relative_attention_max_distance": 128,
+  "relative_attention_num_buckets": 32,
+  "task_specific_params": {
+    "summarization": {
+      "early_stopping": true,
+      "length_penalty": 2.0,
+      "max_length": 200,
+      "min_length": 30,
+      "no_repeat_ngram_size": 3,
+      "num_beams": 4,
+      "prefix": "summarize: "
+    },
+    "translation_en_to_de": {
+      "early_stopping": true,
+      "max_length": 300,
+      "num_beams": 4,
+      "prefix": "translate English to German: "
+    },
+    "translation_en_to_fr": {
+      "early_stopping": true,
+      "max_length": 300,
+      "num_beams": 4,
+      "prefix": "translate English to French: "
+    },
+    "translation_en_to_ro": {
+      "early_stopping": true,
+      "max_length": 300,
+      "num_beams": 4,
+      "prefix": "translate English to Romanian: "
+    }
+  },
+  "tie_word_embeddings": false,
+  "transformers_version": "4.23.1",
+  "use_cache": true,
+  "vocab_size": 32128
+}

flan-t5-base.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1dfb70afdcedceb9f9fae2f9b68e004ad934361fb35b9b2bd50b45ea90790fc8
+size 990345061

generate.bat

@@ -0,0 +1,16 @@
+@echo off
+setlocal
+cd /d "%~dp0"
+
+if "%~1"=="" goto :usage
+if "%~2"=="" goto :usage
+
+set "OUT=%~2"
+if not "%OUT:~1,1%"==":" if not "%OUT:~0,1%"=="\" set "OUT=%~dp0%OUT%"
+
+"C:\Development\trainers\maltrainer\venv\Scripts\python.exe" "%~dp0inference.py" --ckpt "%~dp0Supra2-IMG_fp32.safetensors" --prompt "%~1" --out "%OUT%"
+exit /b %ERRORLEVEL%
+
+:usage
+echo Usage: generate.bat "your prompt" output.png
+exit /b 1

generate.sh

@@ -0,0 +1,28 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+ROOT="$(cd "$(dirname "$0")" && pwd)"
+
+if [[ $# -lt 2 ]]; then
+  echo 'Usage: ./generate.sh "your prompt" output.png'
+  exit 1
+fi
+
+if [[ -x "/c/Development/trainers/maltrainer/venv/Scripts/python.exe" ]]; then
+  PY="/c/Development/trainers/maltrainer/venv/Scripts/python.exe"
+elif [[ -x "C:/Development/trainers/maltrainer/venv/Scripts/python.exe" ]]; then
+  PY="C:/Development/trainers/maltrainer/venv/Scripts/python.exe"
+else
+  PY="${PYTHON:-python3}"
+fi
+
+OUT="$2"
+case "$OUT" in
+  /* | [A-Za-z]:* | [A-Za-z]:\\*) ;;
+  *) OUT="$ROOT/$OUT" ;;
+esac
+
+"$PY" "$ROOT/inference.py" \
+  --ckpt "$ROOT/Supra2-IMG_fp32.safetensors" \
+  --prompt "$1" \
+  --out "$OUT"

inference.py

@@ -0,0 +1,456 @@
+#!/usr/bin/env python3
+"""
+Supra2-IMG inference — standalone text-to-image (DiT + Flan-T5-Base + SD VAE).
+
+Works on Linux and Windows. No imports from other project files.
+
+Usage:
+    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
+
+    generate.bat "a red rose in daylight" rose.png
+    ./generate.sh "a red rose in daylight" rose.png
+
+Default checkpoint is Supra2-IMG_fp32.safetensors next to this script.
+"""
+
+from __future__ import annotations
+
+import argparse
+import math
+import os
+import sys
+import time
+from tqdm import tqdm
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+# ---------------------------------------------------------------------------
+# Architecture constants (must match the trained checkpoint)
+# ---------------------------------------------------------------------------
+IMG_SIZE = 256
+LATENT_SIZE = 32  # 256 / 8 (f8 VAE)
+LATENT_CH = 4
+PATCH = 2
+NUM_TOKENS = (LATENT_SIZE // PATCH) ** 2
+
+D_MODEL = 576
+DEPTH = 14
+N_HEADS = 9
+MLP_RATIO = 4.0
+D_CTX = 768  # Flan-T5-Base
+MAX_CTX_LEN = 128
+LOCAL_DIR = os.path.dirname(os.path.abspath(__file__))
+FLAN_WEIGHTS = os.path.join(LOCAL_DIR, "flan-t5-base.safetensors")
+FLAN_CONFIG = os.path.join(LOCAL_DIR, "flan-t5-base.config.json")
+FLAN_TOKENIZER = os.path.join(LOCAL_DIR, "flan-t5-base.tokenizer.json")
+VAE_WEIGHTS = os.path.join(LOCAL_DIR, "sd-vae-ft-mse.safetensors")
+VAE_CONFIG = os.path.join(LOCAL_DIR, "sd-vae-ft-mse.config.json")
+VAE_SCALE = 0.18215
+
+HF_REPO = "SupraLabs/Supra2-IMG"
+DEFAULT_CKPT_SAFETENSORS = os.path.join(LOCAL_DIR, "Supra2-IMG_fp32.safetensors")
+DEFAULT_CKPT_PT = os.path.join(LOCAL_DIR, "model_final_ema.pt")
+DEFAULT_CKPT = DEFAULT_CKPT_SAFETENSORS if os.path.isfile(DEFAULT_CKPT_SAFETENSORS) else DEFAULT_CKPT_PT
+
+
+def log(msg: str) -> None:
+    """Print immediately so the user sees live progress."""
+    print(msg, flush=True)
+
+
+def pick_device() -> torch.device:
+    if torch.cuda.is_available():
+        torch.cuda.set_device(0)
+        name = torch.cuda.get_device_name(0)
+        mem = torch.cuda.get_device_properties(0).total_memory / 1e9
+        log(f"[device] CUDA: {name} ({mem:.1f} GB)")
+        return torch.device("cuda:0")
+    if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+        log("[device] Apple MPS")
+        return torch.device("mps")
+    log("[device] CPU (this will be slow)")
+    return torch.device("cpu")
+
+
+# ===========================================================================
+# MODEL
+# ===========================================================================
+def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
+    """AdaLN: x * (1 + scale) + shift."""
+    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+class TimestepEmbedder(nn.Module):
+    """Sinusoidal timestep embedding followed by an MLP."""
+
+    def __init__(self, hidden_size: int, freq_dim: int = 256) -> None:
+        super().__init__()
+        self.freq_dim = freq_dim
+        self.mlp = nn.Sequential(
+            nn.Linear(freq_dim, hidden_size),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size),
+        )
+
+    def _sinusoidal(self, t: torch.Tensor) -> torch.Tensor:
+        half = self.freq_dim // 2
+        freqs = torch.exp(
+            -math.log(10000.0) * torch.arange(half, device=t.device) / half
+        )
+        args = t[:, None].float() * freqs[None] * 1000.0
+        emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if self.freq_dim % 2:
+            emb = F.pad(emb, (0, 1))
+        return emb
+
+    def forward(self, t: torch.Tensor) -> torch.Tensor:
+        return self.mlp(self._sinusoidal(t))
+
+
+class Attention(nn.Module):
+    """Multi-head self- or cross-attention."""
+
+    def __init__(self, dim: int, n_heads: int, ctx_dim: int | None = None) -> None:
+        super().__init__()
+        self.n_heads = n_heads
+        self.head_dim = dim // n_heads
+        self.is_self = ctx_dim is None
+        if self.is_self:
+            self.qkv = nn.Linear(dim, dim * 3, bias=True)
+        else:
+            self.q = nn.Linear(dim, dim, bias=True)
+            self.kv = nn.Linear(ctx_dim, dim * 2, bias=True)
+        self.proj = nn.Linear(dim, dim, bias=True)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        ctx: torch.Tensor | None = None,
+        ctx_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        B, N, C = x.shape
+        if self.is_self:
+            qkv = self.qkv(x).view(B, N, 3, self.n_heads, self.head_dim)
+            q, k, v = (qkv[:, :, i].transpose(1, 2) for i in range(3))
+        else:
+            M = ctx.shape[1]
+            q = self.q(x).view(B, N, self.n_heads, self.head_dim).transpose(1, 2)
+            kv = self.kv(ctx).view(B, M, 2, self.n_heads, self.head_dim)
+            k, v = kv[:, :, 0].transpose(1, 2), kv[:, :, 1].transpose(1, 2)
+
+        attn_mask = None
+        if ctx_mask is not None:
+            attn_mask = ctx_mask.bool()[:, None, None, :]
+
+        out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
+        out = out.transpose(1, 2).reshape(B, N, C)
+        return self.proj(out)
+
+
+class DiTBlock(nn.Module):
+    """DiT block: AdaLN-Zero self-attn + cross-attn + MLP."""
+
+    def __init__(self, dim: int, n_heads: int, ctx_dim: int, mlp_ratio: float) -> None:
+        super().__init__()
+        self.norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.self_attn = Attention(dim, n_heads)
+        self.norm_ca = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.cross_attn = Attention(dim, n_heads, ctx_dim=dim)
+        self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        hidden = int(dim * mlp_ratio)
+        self.mlp = nn.Sequential(
+            nn.Linear(dim, hidden),
+            nn.GELU(approximate="tanh"),
+            nn.Linear(hidden, dim),
+        )
+        self.adaln = nn.Sequential(nn.SiLU(), nn.Linear(dim, 6 * dim, bias=True))
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        c: torch.Tensor,
+        ctx: torch.Tensor,
+        ctx_mask: torch.Tensor | None,
+    ) -> torch.Tensor:
+        shift_sa, scale_sa, gate_sa, shift_mlp, scale_mlp, gate_mlp = self.adaln(c).chunk(6, dim=1)
+        x = x + gate_sa.unsqueeze(1) * self.self_attn(
+            modulate(self.norm1(x), shift_sa, scale_sa)
+        )
+        x = x + self.cross_attn(self.norm_ca(x), ctx=ctx, ctx_mask=ctx_mask)
+        x = x + gate_mlp.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp))
+        return x
+
+
+class FinalLayer(nn.Module):
+    def __init__(self, dim: int, out_ch: int) -> None:
+        super().__init__()
+        self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+        self.linear = nn.Linear(dim, out_ch, bias=True)
+        self.adaln = nn.Sequential(nn.SiLU(), nn.Linear(dim, 2 * dim, bias=True))
+
+    def forward(self, x: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
+        shift, scale = self.adaln(c).chunk(2, dim=1)
+        return self.linear(modulate(self.norm(x), shift, scale))
+
+
+class SupraDiT(nn.Module):
+    """~100M parameter DiT for rectified-flow text-to-image."""
+
+    def __init__(
+        self,
+        latent_ch: int = LATENT_CH,
+        d_model: int = D_MODEL,
+        depth: int = DEPTH,
+        n_heads: int = N_HEADS,
+        ctx_dim: int = D_CTX,
+        mlp_ratio: float = MLP_RATIO,
+        num_tokens: int = NUM_TOKENS,
+    ) -> None:
+        super().__init__()
+        self.num_tokens = num_tokens
+        self.patch = PATCH
+        self.x_embed = nn.Linear(latent_ch * PATCH * PATCH, d_model)
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_tokens, d_model))
+        self.t_embed = TimestepEmbedder(d_model)
+        self.ctx_proj = nn.Linear(ctx_dim, d_model)
+        self.blocks = nn.ModuleList(
+            [DiTBlock(d_model, n_heads, d_model, mlp_ratio) for _ in range(depth)]
+        )
+        self.final = FinalLayer(d_model, latent_ch * PATCH * PATCH)
+
+    def forward(
+        self,
+        z: torch.Tensor,
+        t: torch.Tensor,
+        ctx: torch.Tensor,
+        ctx_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        B, C, H, W = z.shape
+        P = self.patch
+        h, w = H // P, W // P
+        x = z.view(B, C, h, P, w, P).permute(0, 2, 4, 1, 3, 5).reshape(B, h * w, C * P * P)
+        x = self.x_embed(x) + self.pos_embed
+        c = self.t_embed(t)
+        ctx = self.ctx_proj(ctx)
+        for blk in self.blocks:
+            x = blk(x, c, ctx, ctx_mask)
+        x = self.final(x, c)
+        return x.view(B, h, w, C, P, P).permute(0, 3, 1, 4, 2, 5).reshape(B, C, H, W)
+
+
+# ===========================================================================
+# Checkpoint download
+# ===========================================================================
+def ensure_checkpoint(path: str) -> str:
+    """Use local checkpoint or download model_final_ema.pt from Hugging Face."""
+    if os.path.isfile(path):
+        log(f"[ckpt] found {path}")
+        return path
+
+    log(f"[ckpt] {path} not found — downloading from {HF_REPO} ...")
+    try:
+        from huggingface_hub import hf_hub_download
+    except ImportError:
+        log("[ckpt] installing huggingface_hub ...")
+        import subprocess
+
+        subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "huggingface_hub"])
+        from huggingface_hub import hf_hub_download
+
+    filename = os.path.basename(path) or "model_final_ema.pt"
+    downloaded = hf_hub_download(
+        repo_id=HF_REPO,
+        filename=filename,
+        local_dir=".",
+        local_dir_use_symlinks=False,
+    )
+    # Prefer the expected local path if the hub placed it elsewhere
+    if os.path.isfile(filename) and os.path.abspath(filename) != os.path.abspath(path):
+        import shutil
+
+        shutil.copy2(filename, path)
+        log(f"[ckpt] copied to {path}")
+        return path
+    log(f"[ckpt] downloaded: {downloaded}")
+    return downloaded if os.path.isfile(downloaded) else path
+
+
+# ===========================================================================
+# Sampling (Euler integration of the flow ODE + optional CFG)
+# ===========================================================================
+@torch.no_grad()
+def generate(args: argparse.Namespace, device: torch.device) -> None:
+    import json
+
+    from transformers import T5Config, T5EncoderModel, T5TokenizerFast
+    from diffusers import AutoencoderKL
+    import torchvision.utils as vutils
+    from safetensors import safe_open
+    from safetensors.torch import load_file
+
+    ckpt_path = ensure_checkpoint(args.ckpt or DEFAULT_CKPT)
+
+    log("[model] building SupraDiT ...")
+    model = SupraDiT().to(device).eval()
+    n_params = sum(p.numel() for p in model.parameters())
+    log(f"[model] {n_params / 1e6:.1f}M parameters")
+
+    log(f"[model] loading weights from {ckpt_path} ...")
+    t0 = time.perf_counter()
+
+    cfg = {}
+    if ckpt_path.endswith(".safetensors"):
+        weights = load_file(ckpt_path, device=str(device))
+        with safe_open(ckpt_path, framework="pt", device="cpu") as f:
+            meta = f.metadata() or {}
+            cfg = {k: int(v) if v.isdigit() else v for k, v in meta.items()}
+        if "__uncond_text__" in weights:
+            cfg["uncond_text"] = weights.pop("__uncond_text__")
+        if "__uncond_mask__" in weights:
+            cfg["uncond_mask"] = weights.pop("__uncond_mask__")
+    else:
+        state = torch.load(ckpt_path, map_location=device, weights_only=False)
+        cfg = state.get("config", {}) if isinstance(state, dict) else {}
+        if isinstance(state, dict):
+            if cfg.get("patch", PATCH) != PATCH:
+                raise SystemExit(f"Checkpoint PATCH={cfg['patch']} != script PATCH={PATCH}")
+            weights = state["ema"] if "ema" in state else state.get("model", state)
+        else:
+            weights = state
+
+    model.load_state_dict(weights, strict=True)
+    log(f"[model] weights loaded in {time.perf_counter() - t0:.1f}s")
+
+    ctx_len = int(cfg.get("ctx_len", MAX_CTX_LEN))
+    log(f"[text] ctx_len={ctx_len}")
+
+    log(f"[text] loading {os.path.basename(FLAN_WEIGHTS)} ...")
+    tokenizer = T5TokenizerFast(tokenizer_file=FLAN_TOKENIZER, extra_ids=100)
+    text_model = T5EncoderModel(T5Config.from_json_file(FLAN_CONFIG))
+    text_weights = load_file(FLAN_WEIGHTS)
+    text_incompat = text_model.load_state_dict(text_weights, strict=False)
+    bad_missing = [k for k in text_incompat.missing_keys if "embed_tokens" not in k]
+    bad_unexpected = [
+        k for k in text_incompat.unexpected_keys if not k.startswith("decoder.") and k != "lm_head.weight"
+    ]
+    if bad_missing or bad_unexpected:
+        raise SystemExit(f"Flan-T5 load mismatch missing={bad_missing} unexpected={bad_unexpected}")
+    del text_weights
+    text_model = text_model.to(device).eval()
+    for p in text_model.parameters():
+        p.requires_grad = False
+
+    log(f"[vae] loading {os.path.basename(VAE_WEIGHTS)} ...")
+    with open(VAE_CONFIG, encoding="utf-8") as f:
+        vae = AutoencoderKL.from_config(json.load(f))
+    vae_weights = load_file(VAE_WEIGHTS)
+    vae._fix_state_dict_keys_on_load(vae_weights)
+    vae.load_state_dict(vae_weights, strict=True)
+    del vae_weights
+    vae = vae.to(device).eval()
+
+    torch.manual_seed(args.seed)
+    if device.type == "cuda":
+        torch.cuda.manual_seed_all(args.seed)
+
+    prompts = [args.prompt] * args.n
+    n_tok = len(tokenizer(args.prompt)["input_ids"])
+    log(f"[text] prompt tokens={n_tok}  n={args.n}  seed={args.seed}  cfg={args.cfg}  steps={args.steps}")
+    if n_tok > ctx_len:
+        log(f"[text] WARNING: prompt has {n_tok} tokens, truncated to ctx_len={ctx_len}")
+
+    tok = tokenizer(
+        prompts,
+        padding="max_length",
+        truncation=True,
+        max_length=ctx_len,
+        return_tensors="pt",
+    ).to(device)
+
+    use_amp = device.type == "cuda"
+    with torch.autocast("cuda", dtype=torch.bfloat16, enabled=use_amp):
+        ctx = text_model(**tok).last_hidden_state.float()
+    cmask = tok["attention_mask"].float()
+
+    use_cfg = args.cfg > 1.0
+    if use_cfg:
+        if isinstance(cfg, dict) and "uncond_text" in cfg:
+            uncond_ctx = cfg["uncond_text"].to(device).float().unsqueeze(0).expand(args.n, -1, -1)
+            uncond_mask = cfg["uncond_mask"].to(device).float().unsqueeze(0).expand(args.n, -1)
+            log("[cfg] using stored unconditional embeddings")
+        else:
+            u_tok = tokenizer(
+                [""] * args.n,
+                padding="max_length",
+                truncation=True,
+                max_length=ctx_len,
+                return_tensors="pt",
+            ).to(device)
+            with torch.autocast("cuda", dtype=torch.bfloat16, enabled=use_amp):
+                uncond_ctx = text_model(**u_tok).last_hidden_state.float()
+            uncond_mask = u_tok["attention_mask"].float()
+            log("[cfg] encoded empty string as unconditional")
+        ctx_all = torch.cat([ctx, uncond_ctx], 0)
+        mask_all = torch.cat([cmask, uncond_mask], 0)
+
+    z = torch.randn(args.n, LATENT_CH, LATENT_SIZE, LATENT_SIZE, device=device)
+    dt = 1.0 / args.steps
+    log(f"[sample] Euler flow, {args.steps} steps ...")
+    t_sample = time.perf_counter()
+
+    for i in tqdm(range(args.steps), desc="Generating", unit="step", dynamic_ncols=True):
+        t = torch.full((args.n,), i * dt, device=device)
+        with torch.autocast("cuda", dtype=torch.bfloat16, enabled=use_amp):
+            if use_cfg:
+                v_both = model(torch.cat([z, z], 0), torch.cat([t, t], 0), ctx_all, mask_all)
+                v_cond, v_uncond = v_both.float().chunk(2, 0)
+                v = v_uncond + args.cfg * (v_cond - v_uncond)
+            else:
+                v = model(z, t, ctx, cmask).float()
+        z = z + dt * v
+
+    log(f"[sample] denoising done in {time.perf_counter() - t_sample:.1f}s")
+    log("[vae] decoding latents ...")
+    with torch.autocast("cuda", dtype=torch.bfloat16, enabled=use_amp):
+        imgs = vae.decode(z / VAE_SCALE).sample
+    imgs = (imgs.clamp(-1, 1) + 1) / 2
+
+    out_dir = os.path.dirname(os.path.abspath(args.out))
+    if out_dir:
+        os.makedirs(out_dir, exist_ok=True)
+    vutils.save_image(imgs, args.out, nrow=int(math.ceil(math.sqrt(args.n))))
+    log(f"[done] saved {args.n} image(s) -> {args.out}")
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(description="Supra2-IMG standalone inference")
+    p.add_argument(
+        "--ckpt",
+        default=None,
+        help="Path to checkpoint (.safetensors or .pt)",
+    )
+    p.add_argument(
+        "--prompt",
+        default="a sea jellyfish floating in the pitch-black ocean depths",
+        help="Text prompt",
+    )
+    p.add_argument("--seed", type=int, default=0, help="RNG seed")
+    p.add_argument("--cfg", type=float, default=3.0, help="Classifier-free guidance scale")
+    p.add_argument("--steps", type=int, default=50, help="Euler ODE steps")
+    p.add_argument("--n", type=int, default=1, help="Number of images")
+    p.add_argument("--out", default="jellyfish.png", help="Output image path")
+    return p.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    log("=== Supra2-IMG inference ===")
+    device = pick_device()
+    generate(args, device)
+
+
+if __name__ == "__main__":
+    main()

sd-vae-ft-mse.config.json

@@ -0,0 +1,31 @@
+{
+  "_class_name": "AutoencoderKL",
+  "_diffusers_version": "0.27.0",
+  "act_fn": "silu",
+  "block_out_channels": [
+    128,
+    256,
+    512,
+    512
+  ],
+  "down_block_types": [
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D",
+    "DownEncoderBlock2D"
+  ],
+  "force_upcast": true,
+  "in_channels": 3,
+  "latent_channels": 4,
+  "layers_per_block": 2,
+  "norm_num_groups": 32,
+  "out_channels": 3,
+  "sample_size": 256,
+  "scaling_factor": 0.18215,
+  "up_block_types": [
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D",
+    "UpDecoderBlock2D"
+  ]
+}

sd-vae-ft-mse.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a1d993488569e928462932c8c38a0760b874d166399b14414135bd9c42df5815
+size 334643276