inference.py

#!/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()