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