Publish 55K parameter class-conditional diffusion model

README.md

@@ -0,0 +1,41 @@
+---
+license: apache-2.0
+task_categories:
+- unconditional-image-generation
+tags:
+- diffusion
+- classifier-free-guidance
+- tiny-model
+- pytorch
+---
+
+# PocketDiffusion
+
+PocketDiffusion is a compact class-conditional denoising diffusion model for 8x8
+handwritten digits. It learns to predict Gaussian noise over 50 diffusion steps and
+uses classifier-free guidance during sampling.
+
+The same frozen Tiny Vision classifier used for GlyphForge evaluates conditional
+recognizability, making the VAE and diffusion results directly comparable under one
+judge.
+
+## Reproduce
+
+```powershell
+uv run python projects/tiny-vision-foundry/prepare_data.py
+uv run python projects/pocket-diffusion/train.py
+```
+
+## Verified results
+
+- Parameters: **55,608**
+- Diffusion steps: **50**
+- Training epochs: **300**
+- Generated samples: **1,000**
+- Selected classifier-free guidance: **3.0**
+- Frozen-judge class fidelity: **96.20%**
+
+Guidance search improved fidelity monotonically from 46.10% at `1.0` to 96.20% at
+`3.0`. Per-class fidelity ranged from 83% for digit `8` to 100% for digits `0` and
+`6`. Mean within-class pixel variance ranged from 0.0209 to 0.0456, noticeably higher
+than the CVAE's 0.0066 to 0.0168 range under the same 100-samples-per-class protocol.

evaluation.json

@@ -0,0 +1,43 @@
+{
+  "model": "PocketDiffusion",
+  "parameters": 55608,
+  "diffusion_steps": 50,
+  "epochs": 300,
+  "guidance_search": {
+    "1.0": 0.460999995470047,
+    "1.5": 0.7149999737739563,
+    "2.0": 0.8560000061988831,
+    "2.5": 0.9359999895095825,
+    "3.0": 0.9620000123977661
+  },
+  "generation": {
+    "judge_accuracy": 0.9620000123977661,
+    "judge_accuracy_by_class": {
+      "0": 1.0,
+      "1": 0.9200000166893005,
+      "2": 0.9900000095367432,
+      "3": 0.9800000190734863,
+      "4": 0.9599999785423279,
+      "5": 0.9800000190734863,
+      "6": 1.0,
+      "7": 0.9900000095367432,
+      "8": 0.8299999833106995,
+      "9": 0.9700000286102295
+    },
+    "mean_pixel_variance_by_class": {
+      "0": 0.020940322428941727,
+      "1": 0.03187673166394234,
+      "2": 0.03836727514863014,
+      "3": 0.03142565116286278,
+      "4": 0.04158321022987366,
+      "5": 0.034002821892499924,
+      "6": 0.022990796715021133,
+      "7": 0.036991652101278305,
+      "8": 0.04558330774307251,
+      "9": 0.03457583114504814
+    },
+    "samples": 1000,
+    "guidance": 3.0
+  },
+  "judge": "Tiny Vision labels-only student, 98.52% real-image test accuracy"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f4f48f2fbc3b5dd6af4c631c3d86cff984ffbd50e624064c5cf929054bd7d813
+size 223088

source/app.py

@@ -0,0 +1,64 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import gradio as gr
+import numpy as np
+import torch
+from model import PocketDenoiser
+from PIL import Image
+from safetensors.torch import load_file
+
+STEPS = 50
+ARTIFACT = Path(__file__).resolve().parent / "artifacts" / "pocket-diffusion"
+MODEL = PocketDenoiser(diffusion_steps=STEPS)
+MODEL.load_state_dict(load_file(ARTIFACT / "model.safetensors"))
+MODEL.eval()
+
+
+def generate_digit(label: int, seed: int, guidance: float) -> Image.Image:
+    generator = torch.Generator().manual_seed(seed)
+    betas = torch.linspace(1e-4, 0.025, STEPS)
+    alphas = 1.0 - betas
+    cumulative = torch.cumprod(alphas, dim=0)
+    pixels = torch.randn(1, 64, generator=generator)
+    labels = torch.tensor([label])
+    null_labels = torch.tensor([10])
+    with torch.no_grad():
+        for step in reversed(range(STEPS)):
+            timesteps = torch.tensor([step])
+            conditional = MODEL(pixels, timesteps, labels)
+            unconditional = MODEL(pixels, timesteps, null_labels)
+            noise_prediction = unconditional + guidance * (conditional - unconditional)
+            alpha = alphas[step]
+            mean = (
+                pixels - (1 - alpha) / torch.sqrt(1 - cumulative[step]) * noise_prediction
+            ) / torch.sqrt(alpha)
+            if step:
+                pixels = mean + torch.sqrt(betas[step]) * torch.randn(
+                    pixels.shape,
+                    generator=generator,
+                )
+            else:
+                pixels = mean
+    image = torch.clamp((pixels[0] + 1) / 2, 0, 1).reshape(8, 8).numpy()
+    array = np.clip(image * 255, 0, 255).astype(np.uint8)
+    return Image.fromarray(array, mode="L").resize(
+        (512, 512),
+        Image.Resampling.NEAREST,
+    )
+
+
+with gr.Blocks(title="PocketDiffusion") as demo:
+    gr.Markdown("# PocketDiffusion\nGenerate a digit through 50 reverse-denoising steps.")
+    with gr.Row():
+        label = gr.Slider(0, 9, value=8, step=1, label="Digit")
+        seed = gr.Slider(0, 100_000, value=2032, step=1, label="Noise seed")
+        guidance = gr.Slider(1.0, 4.0, value=3.0, step=0.1, label="Guidance")
+    output = gr.Image(value=generate_digit(8, 2032, 3.0), label="Generated glyph")
+    button = gr.Button("Denoise", variant="primary")
+    button.click(generate_digit, inputs=[label, seed, guidance], outputs=output)
+
+
+if __name__ == "__main__":
+    demo.launch()

source/model.py

@@ -0,0 +1,60 @@
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+
+class PocketDenoiser(nn.Module):
+    def __init__(self, diffusion_steps: int = 50) -> None:
+        super().__init__()
+        self.diffusion_steps = diffusion_steps
+        self.time_embedding = nn.Embedding(diffusion_steps, 24)
+        self.label_embedding = nn.Embedding(11, 24)
+        self.network = nn.Sequential(
+            nn.Linear(64 + 24 + 24, 160),
+            nn.GELU(),
+            nn.Linear(160, 160),
+            nn.GELU(),
+            nn.Linear(160, 64),
+        )
+
+    def forward(
+        self,
+        noisy_pixels: torch.Tensor,
+        timesteps: torch.Tensor,
+        labels: torch.Tensor,
+    ) -> torch.Tensor:
+        features = torch.cat(
+            [
+                noisy_pixels,
+                self.time_embedding(timesteps),
+                self.label_embedding(labels),
+            ],
+            dim=1,
+        )
+        return self.network(features)
+
+
+class TinyVisionJudge(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.features = nn.Sequential(
+            nn.Conv2d(1, 8, kernel_size=3, padding=1),
+            nn.GELU(),
+            nn.Conv2d(8, 8, kernel_size=3, padding=1, groups=8),
+            nn.GELU(),
+            nn.Conv2d(8, 12, kernel_size=1),
+            nn.GELU(),
+            nn.MaxPool2d(2),
+        )
+        self.classifier = nn.Sequential(
+            nn.Flatten(),
+            nn.Linear(12 * 4 * 4, 10),
+        )
+
+    def forward(self, pixels: torch.Tensor) -> torch.Tensor:
+        return self.classifier(self.features(pixels))
+
+
+def parameter_count(model: nn.Module) -> int:
+    return sum(parameter.numel() for parameter in model.parameters())

source/requirements.txt

@@ -0,0 +1,5 @@
+gradio>=5,<7
+numpy>=2,<3
+pillow>=11,<13
+safetensors>=0.6,<1
+torch>=2.7,<3

source/train.py

@@ -0,0 +1,208 @@
+from __future__ import annotations
+
+import json
+import random
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import torch
+import trackio
+from model import PocketDenoiser, TinyVisionJudge, parameter_count
+from PIL import Image
+from safetensors.torch import load_file, save_file
+from torch.nn import functional as F
+from torch.utils.data import DataLoader, TensorDataset
+
+PROJECT_DIR = Path(__file__).resolve().parent
+ROOT_DIR = PROJECT_DIR.parents[1]
+VISION_DIR = ROOT_DIR / "projects" / "tiny-vision-foundry"
+DATA_DIR = VISION_DIR / "data"
+JUDGE_WEIGHTS = VISION_DIR / "artifacts" / "tiny-student-scratch" / "model.safetensors"
+ARTIFACT_DIR = PROJECT_DIR / "artifacts" / "pocket-diffusion"
+STEPS = 50
+
+
+def seed_everything(seed: int) -> None:
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+
+
+def load_training_data() -> DataLoader:
+    frame = pd.read_parquet(DATA_DIR / "train.parquet")
+    pixels = np.stack(frame["image"].to_numpy()).astype(np.float32) / 8.0 - 1.0
+    labels = frame["label"].to_numpy(dtype=np.int64, copy=True)
+    return DataLoader(
+        TensorDataset(torch.from_numpy(pixels), torch.from_numpy(labels)),
+        batch_size=128,
+        shuffle=True,
+        generator=torch.Generator().manual_seed(2032),
+    )
+
+
+def schedule() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    betas = torch.linspace(1e-4, 0.025, STEPS)
+    alphas = 1.0 - betas
+    cumulative = torch.cumprod(alphas, dim=0)
+    return betas, alphas, cumulative
+
+
+@torch.inference_mode()
+def sample(
+    model: PocketDenoiser,
+    labels: torch.Tensor,
+    guidance: float,
+    seed: int,
+) -> torch.Tensor:
+    generator = torch.Generator().manual_seed(seed)
+    betas, alphas, cumulative = schedule()
+    pixels = torch.randn(len(labels), 64, generator=generator)
+    null_labels = torch.full_like(labels, 10)
+    model.eval()
+    for step in reversed(range(STEPS)):
+        timesteps = torch.full((len(labels),), step, dtype=torch.long)
+        conditional = model(pixels, timesteps, labels)
+        unconditional = model(pixels, timesteps, null_labels)
+        predicted_noise = unconditional + guidance * (conditional - unconditional)
+        alpha = alphas[step]
+        cumulative_alpha = cumulative[step]
+        mean = (
+            pixels - (1 - alpha) / torch.sqrt(1 - cumulative_alpha) * predicted_noise
+        ) / torch.sqrt(alpha)
+        if step:
+            noise = torch.randn(pixels.shape, generator=generator)
+            pixels = mean + torch.sqrt(betas[step]) * noise
+        else:
+            pixels = mean
+    return torch.clamp((pixels + 1) / 2, 0, 1)
+
+
+@torch.inference_mode()
+def generation_metrics(
+    model: PocketDenoiser,
+    judge: TinyVisionJudge,
+    guidance: float,
+) -> tuple[dict, torch.Tensor, torch.Tensor]:
+    labels = torch.arange(10).repeat_interleave(100)
+    generated = sample(model, labels, guidance=guidance, seed=2032)
+    predictions = judge(generated.reshape(-1, 1, 8, 8)).argmax(dim=1)
+    accuracy_by_class = {
+        str(label): float(
+            (predictions[labels == label] == labels[labels == label]).float().mean()
+        )
+        for label in range(10)
+    }
+    diversity = {
+        str(label): float(generated[labels == label].var(dim=0).mean())
+        for label in range(10)
+    }
+    return (
+        {
+            "judge_accuracy": float((predictions == labels).float().mean()),
+            "judge_accuracy_by_class": accuracy_by_class,
+            "mean_pixel_variance_by_class": diversity,
+            "samples": len(labels),
+            "guidance": guidance,
+        },
+        generated,
+        labels,
+    )
+
+
+def save_grid(generated: torch.Tensor, labels: torch.Tensor, path: Path) -> None:
+    selected = [generated[labels == label][:10] for label in range(10)]
+    images = torch.cat(selected).reshape(10, 10, 8, 8).cpu().numpy()
+    canvas = np.zeros((80, 80), dtype=np.uint8)
+    for row in range(10):
+        for column in range(10):
+            canvas[
+                row * 8 : (row + 1) * 8,
+                column * 8 : (column + 1) * 8,
+            ] = np.clip(images[row, column] * 255, 0, 255).astype(np.uint8)
+    Image.fromarray(canvas, mode="L").resize((800, 800), Image.Resampling.NEAREST).save(
+        path
+    )
+
+
+def main() -> None:
+    seed_everything(2032)
+    loader = load_training_data()
+    model = PocketDenoiser()
+    judge = TinyVisionJudge()
+    judge.load_state_dict(load_file(JUDGE_WEIGHTS))
+    judge.eval()
+    betas, _, cumulative = schedule()
+    optimizer = torch.optim.AdamW(model.parameters(), lr=0.0015, weight_decay=0.001)
+    epochs = 300
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)
+    trackio.init(
+        project="pocket-diffusion",
+        name="conditional-ddpm-cfg-v1",
+        config={
+            "parameters": parameter_count(model),
+            "diffusion_steps": STEPS,
+            "epochs": epochs,
+            "label_dropout": 0.12,
+        },
+    )
+    for epoch in range(1, epochs + 1):
+        model.train()
+        running_loss = 0.0
+        examples = 0
+        for pixels, labels in loader:
+            timesteps = torch.randint(0, STEPS, (len(labels),))
+            noise = torch.randn_like(pixels)
+            cumulative_alpha = cumulative[timesteps].unsqueeze(1)
+            noisy = (
+                torch.sqrt(cumulative_alpha) * pixels
+                + torch.sqrt(1 - cumulative_alpha) * noise
+            )
+            conditioned_labels = labels.clone()
+            drop = torch.rand(len(labels)) < 0.12
+            conditioned_labels[drop] = 10
+            prediction = model(noisy, timesteps, conditioned_labels)
+            loss = F.mse_loss(prediction, noise)
+            optimizer.zero_grad(set_to_none=True)
+            loss.backward()
+            optimizer.step()
+            running_loss += loss.item() * len(labels)
+            examples += len(labels)
+        scheduler.step()
+        if epoch == 1 or epoch % 10 == 0:
+            trackio.log(
+                {
+                    "epoch": epoch,
+                    "noise_prediction_mse": running_loss / examples,
+                    "learning_rate": scheduler.get_last_lr()[0],
+                }
+            )
+    trackio.finish()
+
+    guidance_candidates = {}
+    for guidance in [1.0, 1.5, 2.0, 2.5, 3.0]:
+        metrics, _, _ = generation_metrics(model, judge, guidance)
+        guidance_candidates[str(guidance)] = metrics["judge_accuracy"]
+    best_guidance = float(max(guidance_candidates, key=guidance_candidates.get))
+    generation, generated, labels = generation_metrics(model, judge, best_guidance)
+    results = {
+        "model": "PocketDiffusion",
+        "parameters": parameter_count(model),
+        "diffusion_steps": STEPS,
+        "epochs": epochs,
+        "guidance_search": guidance_candidates,
+        "generation": generation,
+        "judge": "Tiny Vision labels-only student, 98.52% real-image test accuracy",
+    }
+    ARTIFACT_DIR.mkdir(parents=True, exist_ok=True)
+    save_file(model.state_dict(), ARTIFACT_DIR / "model.safetensors")
+    save_grid(generated, labels, ARTIFACT_DIR / "samples.png")
+    (ARTIFACT_DIR / "evaluation.json").write_text(
+        json.dumps(results, indent=2),
+        encoding="utf-8",
+    )
+    print(json.dumps(results, indent=2))
+
+
+if __name__ == "__main__":
+    main()