Supra2-IMG / inference.py

loras and requirements

89cd62c 29 days ago

19.9 kB

#!/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 _position_embed(self, grid_h: int, grid_w: int) -> torch.Tensor:
        """Learned 16x16 positions, bicubic-resized when sampling above 256px."""
        pos = self.pos_embed
        side = int(math.isqrt(pos.shape[1]))
        if grid_h == side and grid_w == side:
            return pos
        weight = pos.reshape(1, side, side, pos.shape[-1]).permute(0, 3, 1, 2).float()
        weight = F.interpolate(weight, size=(grid_h, grid_w), mode="bicubic", align_corners=False)
        return weight.permute(0, 2, 3, 1).reshape(1, grid_h * grid_w, pos.shape[-1]).to(dtype=pos.dtype)

    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._position_embed(h, w)
        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 apply_lora(model: torch.nn.Module, path: str) -> None:
    """Merge a maltrainer Supra LoRA into the loaded weights. W = W + B @ A * (alpha / rank)."""
    from safetensors.torch import load_file

    tensors = load_file(path, device="cpu")
    grouped: dict[str, dict[str, torch.Tensor]] = {}
    for key, tensor in tensors.items():
        if key.endswith(".alpha"):
            grouped.setdefault(key[: -len(".alpha")], {})["alpha"] = tensor
        elif key.endswith(".lora_A.weight"):
            grouped.setdefault(key[: -len(".lora_A.weight")], {})["A"] = tensor
        elif key.endswith(".lora_B.weight"):
            grouped.setdefault(key[: -len(".lora_B.weight")], {})["B"] = tensor
    applied = 0
    for name, parts in grouped.items():
        if "A" not in parts or "B" not in parts:
            raise SystemExit(f"LoRA module {name} is missing lora_A or lora_B")
        module = model.get_submodule(name)
        down = parts["A"].to(device=module.weight.device, dtype=module.weight.dtype)
        up = parts["B"].to(device=module.weight.device, dtype=module.weight.dtype)
        rank = down.shape[0]
        alpha = float(parts["alpha"].item()) if "alpha" in parts else float(rank)
        module.weight.data.add_((up @ down) * (alpha / rank))
        applied += 1
    log(f"[lora] merged {applied} layers from {path} (alpha/rank scale)")


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")
    if getattr(args, "lora", None):
        apply_lora(model, args.lora)

    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)

    image_size = int(getattr(args, "size", IMG_SIZE))
    if image_size % 16 != 0:
        raise SystemExit(f"--size must be a multiple of 16, got {image_size}")
    latent_size = image_size // 8
    log(f"[sample] image {image_size}px, latent {latent_size}")
    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")
    p.add_argument("--size", type=int, default=IMG_SIZE, help="Output side in pixels. 256 is native; 512 interpolates positions.")
    p.add_argument("--lora", default=None, help="Optional Supra LoRA .safetensors to merge before sampling")
    return p.parse_args()


def main() -> None:
    args = parse_args()
    log("=== Supra2-IMG inference ===")
    device = pick_device()
    generate(args, device)


if __name__ == "__main__":
    main()

	#!/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 _position_embed(self, grid_h: int, grid_w: int) -> torch.Tensor:
	"""Learned 16x16 positions, bicubic-resized when sampling above 256px."""
	pos = self.pos_embed
	side = int(math.isqrt(pos.shape[1]))
	if grid_h == side and grid_w == side:
	return pos
	weight = pos.reshape(1, side, side, pos.shape[-1]).permute(0, 3, 1, 2).float()
	weight = F.interpolate(weight, size=(grid_h, grid_w), mode="bicubic", align_corners=False)
	return weight.permute(0, 2, 3, 1).reshape(1, grid_h * grid_w, pos.shape[-1]).to(dtype=pos.dtype)

	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._position_embed(h, w)
	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 apply_lora(model: torch.nn.Module, path: str) -> None:
	"""Merge a maltrainer Supra LoRA into the loaded weights. W = W + B @ A * (alpha / rank)."""
	from safetensors.torch import load_file

	tensors = load_file(path, device="cpu")
	grouped: dict[str, dict[str, torch.Tensor]] = {}
	for key, tensor in tensors.items():
	if key.endswith(".alpha"):
	grouped.setdefault(key[: -len(".alpha")], {})["alpha"] = tensor
	elif key.endswith(".lora_A.weight"):
	grouped.setdefault(key[: -len(".lora_A.weight")], {})["A"] = tensor
	elif key.endswith(".lora_B.weight"):
	grouped.setdefault(key[: -len(".lora_B.weight")], {})["B"] = tensor
	applied = 0
	for name, parts in grouped.items():
	if "A" not in parts or "B" not in parts:
	raise SystemExit(f"LoRA module {name} is missing lora_A or lora_B")
	module = model.get_submodule(name)
	down = parts["A"].to(device=module.weight.device, dtype=module.weight.dtype)
	up = parts["B"].to(device=module.weight.device, dtype=module.weight.dtype)
	rank = down.shape[0]
	alpha = float(parts["alpha"].item()) if "alpha" in parts else float(rank)
	module.weight.data.add_((up @ down) * (alpha / rank))
	applied += 1
	log(f"[lora] merged {applied} layers from {path} (alpha/rank scale)")


	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")
	if getattr(args, "lora", None):
	apply_lora(model, args.lora)

	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)

	image_size = int(getattr(args, "size", IMG_SIZE))
	if image_size % 16 != 0:
	raise SystemExit(f"--size must be a multiple of 16, got {image_size}")
	latent_size = image_size // 8
	log(f"[sample] image {image_size}px, latent {latent_size}")
	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")
	p.add_argument("--size", type=int, default=IMG_SIZE, help="Output side in pixels. 256 is native; 512 interpolates positions.")
	p.add_argument("--lora", default=None, help="Optional Supra LoRA .safetensors to merge before sampling")
	return p.parse_args()


	def main() -> None:
	args = parse_args()
	log("=== Supra2-IMG inference ===")
	device = pick_device()
	generate(args, device)


	if __name__ == "__main__":
	main()