app.py

#!/usr/bin/env python3
"""
DiT-MNIST — Gradio demo for Hugging Face Spaces.

Place the trained checkpoint (default name: dit_mnist.pt) in the same
directory as this file before deploying. The checkpoint is the one saved
by the training script (contains a 'model' and/or 'ema' state dict).
"""

import math
import os

import gradio as gr
import torch
import torch.nn as nn
from torchvision.utils import make_grid
from PIL import Image
import numpy as np

# ─── Config (must match training config) ──────────────────────────────────

IMAGE_SIZE = 28
PATCH_SIZE = 2
HIDDEN = 128
DEPTH = 6
HEADS = 4
MLP_RATIO = 4.0
NUM_CLASSES = 10
TIMESTEPS = 1000

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
CKPT_PATH = os.environ.get("CKPT_PATH", "dit_mnist.pt")

# ─── Model (identical architecture to training script) ────────────────────

class DiT(nn.Module):
    def __init__(self):
        super().__init__()
        self.num_patches = (IMAGE_SIZE // PATCH_SIZE) ** 2

        self.patchify = nn.Conv2d(1, HIDDEN, PATCH_SIZE, PATCH_SIZE)
        self.pos = nn.Parameter(torch.zeros(1, self.num_patches, HIDDEN))

        self.t_embed = nn.Sequential(
            nn.Linear(128, HIDDEN),
            nn.SiLU(),
            nn.Linear(HIDDEN, HIDDEN),
        )
        self.y_embed = nn.Embedding(NUM_CLASSES + 1, HIDDEN)

        self.blocks = nn.ModuleList([
            nn.ModuleDict({
                'norm1': nn.LayerNorm(HIDDEN, elementwise_affine=False),
                'attn': nn.MultiheadAttention(HIDDEN, HEADS, batch_first=True),
                'norm2': nn.LayerNorm(HIDDEN, elementwise_affine=False),
                'mlp': nn.Sequential(
                    nn.Linear(HIDDEN, int(HIDDEN * MLP_RATIO)),
                    nn.GELU(),
                    nn.Linear(int(HIDDEN * MLP_RATIO), HIDDEN),
                ),
                'adaLN': nn.Sequential(nn.SiLU(), nn.Linear(HIDDEN, 6 * HIDDEN)),
            })
            for _ in range(DEPTH)
        ])

        self.out = nn.Sequential(
            nn.LayerNorm(HIDDEN, elementwise_affine=False),
            nn.Linear(HIDDEN, PATCH_SIZE * PATCH_SIZE),
        )
        self.out_adaLN = nn.Sequential(nn.SiLU(), nn.Linear(HIDDEN, 2 * HIDDEN))

        nn.init.normal_(self.pos, std=0.02)

    def timestep_embed(self, t):
        half = 64
        freqs = torch.exp(-math.log(10000) * torch.arange(half, dtype=torch.float32, device=t.device) / half)
        emb = torch.cat([torch.cos(t[:, None] * freqs), torch.sin(t[:, None] * freqs)], dim=-1)
        return self.t_embed(emb)

    def forward(self, x, t, y):
        x = self.patchify(x).flatten(2).transpose(1, 2) + self.pos
        c = self.timestep_embed(t) + self.y_embed(y)

        for b in self.blocks:
            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = b['adaLN'](c).chunk(6, dim=-1)

            h = b['norm1'](x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
            h, _ = b['attn'](h, h, h, need_weights=False)
            x = x + gate_msa[:, None] * h

            h = b['norm2'](x) * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
            x = x + gate_mlp[:, None] * b['mlp'](h)

        shift, scale = self.out_adaLN(c).chunk(2, dim=-1)
        x = self.out[0](x) * (1 + scale[:, None]) + shift[:, None]
        x = self.out[1](x)

        B = x.shape[0]
        x = x.reshape(B, 14, 14, 2, 2).permute(0, 1, 3, 2, 4).reshape(B, 1, 28, 28)
        return x

# ─── Sampler (rectified-flow style Euler integration, matches training) ───

@torch.no_grad()
def generate(model, y, device, steps=50, cfg_scale=2.0, seed=None):
    model.eval()
    g = torch.Generator(device=device)
    if seed is not None:
        g.manual_seed(int(seed))
    z = torch.randn(y.shape[0], 1, 28, 28, device=device, generator=g)
    null_y = torch.full_like(y, NUM_CLASSES)
    dt = 1 / steps

    for i in range(steps):
        t = torch.full((y.shape[0],), i / steps, device=device)
        v = model(z, t, y)
        if cfg_scale != 1.0:
            v_null = model(z, t, null_y)
            v = v_null + cfg_scale * (v - v_null)
        z = z + dt * v

    return z

# ─── Load model once at startup ────────────────────────────────────────────

_model = DiT().to(DEVICE)

def _load_checkpoint(path):
    if not os.path.exists(path):
        print(f"[WARN] checkpoint '{path}' not found — using randomly initialized weights.")
        return
    ckpt = torch.load(path, map_location=DEVICE)
    state = ckpt.get("ema") or ckpt.get("model") or ckpt
    _model.load_state_dict(state)
    print(f"[OK] loaded checkpoint: {path}")

_load_checkpoint(CKPT_PATH)
_model.eval()

# ─── Inference helper for Gradio ───────────────────────────────────────────

def sample_digits(digit, num_samples, cfg_scale, steps, seed):
    digit = int(digit)
    num_samples = int(num_samples)
    steps = int(steps)
    seed = int(seed) if seed is not None and seed != -1 else None

    y = torch.full((num_samples,), digit, dtype=torch.long, device=DEVICE)
    imgs = generate(_model, y, DEVICE, steps=steps, cfg_scale=cfg_scale, seed=seed)
    imgs = (imgs.clamp(-1, 1) + 1) / 2  # [-1,1] -> [0,1]

    grid = make_grid(imgs, nrow=min(num_samples, 8), padding=2)
    grid_np = (grid.permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8)
    if grid_np.shape[-1] == 1:
        grid_np = grid_np[:, :, 0]
    return Image.fromarray(grid_np)

# ─── Gradio UI ──────────────────────────────────────────────────────────────

with gr.Blocks(title="DiT-MNIST") as demo:
    gr.Markdown(
        "# DiT-MNIST\n"
        "A small class-conditional Diffusion Transformer trained on MNIST. "
        "Pick a digit and sample new handwritten-digit images."
    )
    with gr.Row():
        with gr.Column():
            digit = gr.Slider(0, 9, value=3, step=1, label="Digit (0-9)")
            num_samples = gr.Slider(1, 16, value=8, step=1, label="Number of samples")
            cfg_scale = gr.Slider(0.0, 5.0, value=2.0, step=0.1, label="Classifier-free guidance scale")
            steps = gr.Slider(10, 100, value=50, step=5, label="Sampling steps")
            seed = gr.Number(value=-1, label="Seed (-1 = random)")
            btn = gr.Button("Generate", variant="primary")
        with gr.Column():
            out = gr.Image(label="Generated digits", image_mode="L")

    btn.click(sample_digits, inputs=[digit, num_samples, cfg_scale, steps, seed], outputs=out)

if __name__ == "__main__":
    demo.launch()