Create inference.py

inference.py

@@ -0,0 +1,411 @@
+#!/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
+
+If ./model_final_ema.pt is missing, it is downloaded from Hugging Face:
+    SupraLabs/Supra2-IMG
+"""
+
+from __future__ import annotations
+
+import argparse
+import math
+import os
+import sys
+import time
+
+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
+T5_NAME = "google/flan-t5-base"
+VAE_NAME = "stabilityai/sd-vae-ft-mse"
+VAE_SCALE = 0.18215
+
+HF_REPO = "SupraLabs/Supra2-IMG"
+DEFAULT_CKPT = os.path.join(".", "model_final_ema.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:
+    from transformers import AutoTokenizer, T5EncoderModel
+    from diffusers import AutoencoderKL
+    import torchvision.utils as vutils
+
+    ckpt_path = ensure_checkpoint(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()
+    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 tokenizer + {T5_NAME} ...")
+    tokenizer = AutoTokenizer.from_pretrained(T5_NAME)
+    text_model = T5EncoderModel.from_pretrained(T5_NAME).to(device).eval()
+    for p in text_model.parameters():
+        p.requires_grad = False
+
+    log(f"[vae] loading {VAE_NAME} ...")
+    vae = AutoencoderKL.from_pretrained(VAE_NAME).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 range(args.steps):
+        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
+        if (i + 1) % max(1, args.steps // 10) == 0 or i == 0:
+            log(f"  step {i + 1}/{args.steps}")
+
+    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(
+        "--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()