Add QuickDraw diffusion model and app code

README.md

@@ -0,0 +1,148 @@
+---
+license: apache-2.0
+tags:
+  - quickdraw
+  - sketch-generation
+  - diffusion
+  - ddpm
+  - class-conditional
+  - classifier-free-guidance
+library_name: pytorch
+pipeline_tag: image-to-image
+---
+
+# QuickDraw Sketch Diffusion
+
+This repository contains a small class-conditional DDPM trained on rasterized Google QuickDraw sketches. The model generates 64x64 grayscale sketches with a black background and white strokes.
+
+It also includes a local FastAPI web app for selecting a class, generating sketches with classifier-free guidance, and optionally recognizing the generated image with the bundled 100-class CNN classifier.
+
+## Files
+
+- `models/diffusion/checkpoint_step_500000.pt`: 500k-step class-conditional diffusion checkpoint.
+- `train_quickdraw_ddpm.py`: training, sampling, QuickDraw download, rasterization, and model definitions.
+- `quickdraw_app/server.py`: FastAPI generation and recognition API.
+- `quickdraw_app/static/`: browser UI.
+- `quickdraw_app/cnn_classifier.py`: optional CNN recognition helper.
+- `quickdraw_app/models/cnn_residual_100cls_64/best_model.pt`: bundled 100-class 64x64 CNN classifier.
+- `assets/samples_step_500000.png`: sample grid from the diffusion checkpoint.
+
+## Model Details
+
+- Dataset: Google QuickDraw `full/simplified` drawings.
+- Classes: 100 built-in QuickDraw classes.
+- Training set used: 50,000 samples per class, about 5 million sketches total.
+- Image size: 64x64.
+- Channels: 1 grayscale channel.
+- Diffusion steps: 200.
+- Architecture: compact conditional U-Net.
+- Conditioning: class embedding plus classifier-free guidance.
+- Checkpoint step: 500,000.
+- Default guidance scale: 3.0.
+- Output convention: black background, white sketch strokes.
+
+## Quick Start
+
+Install dependencies:
+
+```bash
+pip install -r requirements.txt
+```
+
+Run the local web app:
+
+```bash
+uvicorn quickdraw_app.server:app --host 127.0.0.1 --port 7860
+```
+
+Then open:
+
+```text
+http://127.0.0.1:7860
+```
+
+The first generation loads the checkpoint and may take longer. Later requests reuse the loaded model.
+
+## Python Sampling
+
+```python
+import torch
+from torchvision.utils import save_image
+
+from train_quickdraw_ddpm import SmallConditionalUNet, make_schedule, pick_device, sample
+
+checkpoint_path = "models/diffusion/checkpoint_step_500000.pt"
+device = pick_device()
+checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
+
+classes = checkpoint["classes"]
+image_size = int(checkpoint["image_size"])
+timesteps = int(checkpoint["timesteps"])
+base_channels = int(checkpoint["base_channels"])
+
+model = SmallConditionalUNet(len(classes), base_channels=base_channels).to(device)
+model.load_state_dict(checkpoint.get("model_unwrapped") or checkpoint["model"])
+model.eval()
+
+schedule = make_schedule(timesteps, device)
+class_name = "cat"
+label = torch.tensor([classes.index(class_name)], device=device)
+
+with torch.no_grad():
+    image = sample(
+        model,
+        label,
+        image_size,
+        schedule,
+        timesteps,
+        device,
+        guidance_scale=3.0,
+    )
+
+save_image((image + 1) / 2, "cat.png")
+```
+
+## Training Example
+
+The final run used 100 classes, 50,000 samples per class, 64x64 images, CFG dropout, and 500k optimization steps. A similar run can be started with:
+
+```bash
+python train_quickdraw_ddpm.py \
+  --num-classes 100 \
+  --samples-per-class 50000 \
+  --image-size 64 \
+  --line-width 2 \
+  --batch-size 512 \
+  --steps 500000 \
+  --timesteps 200 \
+  --base-channels 64 \
+  --cfg-drop-prob 0.1 \
+  --guidance-scale 3.0 \
+  --data-parallel
+```
+
+To resume:
+
+```bash
+python train_quickdraw_ddpm.py \
+  --resume models/diffusion/checkpoint_step_500000.pt \
+  --steps 550000 \
+  --batch-size 512 \
+  --data-parallel
+```
+
+## Evaluation Note
+
+Using the bundled 100-class 64x64 CNN classifier on 400 generated samples, the diffusion checkpoint reached about 36.00% top-1 and 71.75% top-5 agreement. This is a rough model-to-model sanity check, not a human preference score.
+
+## Limitations
+
+This is a compact 64x64 sketch model. It is useful as a QuickDraw-style benchmark and interactive demo, but it is not a high-resolution image generator. Some classes are visually ambiguous, and CNN agreement can be poor for categories with similar silhouettes.
+
+## Data
+
+The training script downloads examples from the public Google QuickDraw dataset:
+
+```text
+https://storage.googleapis.com/quickdraw_dataset/full/simplified/{word}.ndjson
+```

models/diffusion/checkpoint_step_500000.pt

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

quickdraw_app/README.md

@@ -0,0 +1,25 @@
+# QuickDraw Diffusion Web App
+
+Run the local web UI from the repository root:
+
+```bash
+uvicorn quickdraw_app.server:app --host 127.0.0.1 --port 7860
+```
+
+By default it loads:
+
+```text
+models/diffusion/checkpoint_step_500000.pt
+```
+
+To use another checkpoint:
+
+```bash
+QUICKDRAW_CHECKPOINT=/absolute/path/to/checkpoint.pt uvicorn quickdraw_app.server:app --host 127.0.0.1 --port 7860
+```
+
+The recognition button uses the bundled 100-class 64x64 residual CNN:
+
+```text
+quickdraw_app/models/cnn_residual_100cls_64/best_model.pt
+```

quickdraw_app/cnn_classifier.py

@@ -0,0 +1,271 @@
+"""Embedded CNN image classifier from the sketch-model project."""
+
+from __future__ import annotations
+
+import base64
+import json
+from functools import lru_cache
+from io import BytesIO
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+from fastapi import HTTPException
+from PIL import Image, UnidentifiedImageError
+from torch import nn
+from torch.nn import functional as F
+from torchvision.models import ResNet18_Weights, resnet18
+
+APP_DIR = Path(__file__).resolve().parent
+MODEL_DIR = APP_DIR / "models" / "cnn_resnet18_20cls"
+MODEL_PATH = MODEL_DIR / "best_model.pt"
+LABEL_MAP_PATH = MODEL_DIR / "label_map.json"
+QUICKDRAW100_MODEL_DIR = APP_DIR / "models" / "cnn_residual_100cls_64"
+QUICKDRAW100_MODEL_PATH = QUICKDRAW100_MODEL_DIR / "best_model.pt"
+QUICKDRAW100_LABEL_MAP_PATH = QUICKDRAW100_MODEL_DIR / "label_map.json"
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, channels: int, dropout: float = 0.0):
+        super().__init__()
+        self.block = nn.Sequential(
+            nn.Conv2d(channels, channels, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(channels),
+            nn.ReLU(),
+            nn.Dropout2d(dropout),
+            nn.Conv2d(channels, channels, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(channels),
+        )
+        self.activation = nn.ReLU()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.activation(x + self.block(x))
+
+
+class ResidualCNN(nn.Module):
+    def __init__(self, num_classes: int):
+        super().__init__()
+        self.features = nn.Sequential(
+            nn.Conv2d(1, 48, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(48),
+            nn.ReLU(),
+            ResidualBlock(48, dropout=0.05),
+            nn.MaxPool2d(2),
+            nn.Conv2d(48, 96, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(96),
+            nn.ReLU(),
+            ResidualBlock(96, dropout=0.05),
+            nn.MaxPool2d(2),
+            nn.Conv2d(96, 192, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(192),
+            nn.ReLU(),
+            ResidualBlock(192, dropout=0.08),
+            nn.MaxPool2d(2),
+            nn.Conv2d(192, 256, kernel_size=3, padding=1, bias=False),
+            nn.BatchNorm2d(256),
+            nn.ReLU(),
+            ResidualBlock(256, dropout=0.08),
+            nn.AdaptiveAvgPool2d((1, 1)),
+        )
+        self.classifier = nn.Sequential(
+            nn.Flatten(),
+            nn.Dropout(0.35),
+            nn.Linear(256, num_classes),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.classifier(self.features(x))
+
+
+def build_model(
+    model_name: str,
+    num_classes: int,
+    pretrained: bool = False,
+    freeze_backbone: bool = False,
+) -> nn.Module:
+    if model_name == "resnet18":
+        weights = ResNet18_Weights.DEFAULT if pretrained else None
+        model = resnet18(weights=weights)
+        model.conv1 = nn.Conv2d(1, 64, kernel_size=3, stride=1, padding=1, bias=False)
+        model.maxpool = nn.Identity()
+        model.fc = nn.Linear(model.fc.in_features, num_classes)
+        if freeze_backbone:
+            for name, param in model.named_parameters():
+                param.requires_grad = name.startswith("fc.")
+        return model
+    if model_name == "residual-cnn":
+        return ResidualCNN(num_classes=num_classes)
+    raise ValueError(f"Unsupported embedded CNN model: {model_name}")
+
+
+def pick_cnn_device() -> torch.device:
+    if torch.backends.mps.is_available():
+        return torch.device("mps")
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    return torch.device("cpu")
+
+
+def load_label_map(label_map_path: Path = LABEL_MAP_PATH) -> dict[str, int]:
+    if not label_map_path.exists():
+        raise FileNotFoundError(f"Missing CNN label map: {label_map_path}")
+    return json.loads(label_map_path.read_text(encoding="utf-8"))
+
+
+def labels_from_map(label_map: dict[str, int]) -> list[str]:
+    return [name for name, _ in sorted(label_map.items(), key=lambda item: item[1])]
+
+
+@lru_cache(maxsize=1)
+def image_predictor() -> tuple[nn.Module, list[str], torch.device, int]:
+    if not MODEL_PATH.exists():
+        raise FileNotFoundError(f"Missing CNN model: {MODEL_PATH}")
+    return load_predictor(MODEL_PATH, LABEL_MAP_PATH)
+
+
+@lru_cache(maxsize=1)
+def quickdraw100_predictor() -> tuple[nn.Module, list[str], torch.device, int]:
+    return load_predictor(QUICKDRAW100_MODEL_PATH, QUICKDRAW100_LABEL_MAP_PATH)
+
+
+def load_predictor(model_path: Path, label_map_path: Path) -> tuple[nn.Module, list[str], torch.device, int]:
+    if not model_path.exists():
+        raise FileNotFoundError(f"Missing CNN model: {model_path}")
+    label_map = load_label_map(label_map_path)
+    labels = labels_from_map(label_map)
+    device = pick_cnn_device()
+
+    checkpoint = torch.load(model_path, map_location=device, weights_only=False)
+    config = checkpoint.get("config", {})
+    image_size = int(config.get("image_size", 96))
+    model = build_model(
+        config.get("model", "resnet18"),
+        num_classes=len(labels),
+        pretrained=bool(config.get("pretrained", False)),
+        freeze_backbone=bool(config.get("freeze_backbone", False)),
+    ).to(device)
+    model.load_state_dict(checkpoint["model_state"])
+    model.eval()
+    return model, labels, device, image_size
+
+
+def decode_data_url(data_url: str) -> bytes:
+    if "," in data_url:
+        _, encoded = data_url.split(",", 1)
+    else:
+        encoded = data_url
+    try:
+        return base64.b64decode(encoded, validate=True)
+    except ValueError as exc:
+        raise HTTPException(status_code=400, detail="image must be a base64 PNG data URL.") from exc
+
+
+def uploaded_image_to_tensor(image_bytes: bytes, image_size: int) -> torch.Tensor:
+    try:
+        image = Image.open(BytesIO(image_bytes)).convert("L")
+    except UnidentifiedImageError as exc:
+        raise HTTPException(status_code=400, detail="Unsupported image file. Use PNG, JPG, or WEBP.") from exc
+
+    array = np.asarray(image, dtype=np.float32) / 255.0
+    if array.mean() > 0.5:
+        array = 1.0 - array
+
+    ink = array > 0.15
+    if not np.any(ink):
+        raise HTTPException(status_code=400, detail="No visible sketch stroke detected in the image.")
+
+    ys, xs = np.where(ink)
+    y0, y1 = int(ys.min()), int(ys.max()) + 1
+    x0, x1 = int(xs.min()), int(xs.max()) + 1
+    crop = array[y0:y1, x0:x1]
+
+    side = max(crop.shape)
+    pad = max(2, int(side * 0.12))
+    canvas = np.zeros((side + pad * 2, side + pad * 2), dtype=np.float32)
+    offset_y = (canvas.shape[0] - crop.shape[0]) // 2
+    offset_x = (canvas.shape[1] - crop.shape[1]) // 2
+    canvas[offset_y : offset_y + crop.shape[0], offset_x : offset_x + crop.shape[1]] = crop
+
+    resample = Image.Resampling.BILINEAR if hasattr(Image, "Resampling") else Image.BILINEAR
+    resized = Image.fromarray(np.uint8(np.clip(canvas, 0.0, 1.0) * 255)).resize(
+        (image_size, image_size),
+        resample,
+    )
+    tensor_array = np.asarray(resized, dtype=np.float32) / 255.0
+    return torch.from_numpy(tensor_array[None, :, :]).unsqueeze(0)
+
+
+def tensor_image_to_input(image: torch.Tensor, image_size: int) -> torch.Tensor:
+    if image.ndim == 2:
+        image = image[None, :, :]
+    if image.ndim == 3:
+        image = image.unsqueeze(0)
+    if image.shape[1] != 1:
+        image = image.mean(dim=1, keepdim=True)
+    image = image.float().clamp(0, 1)
+    if image.shape[-2:] != (image_size, image_size):
+        image = F.interpolate(image, size=(image_size, image_size), mode="bilinear", align_corners=False)
+    return image
+
+
+def top_predictions(probs: torch.Tensor, labels: list[str], top_k: int) -> list[dict[str, Any]]:
+    top_k = max(1, min(int(top_k), len(labels)))
+    scores, indices = torch.topk(probs, k=top_k)
+    predictions = []
+    for score, idx in zip(scores.tolist(), indices.tolist()):
+        label = labels[idx]
+        predictions.append(
+            {
+                "label": label,
+                "confidence": float(score),
+                "reference": f"/api/reference/{label}.svg",
+            }
+        )
+    return predictions
+
+
+def classes(predictor: str = "legacy20") -> list[str]:
+    if predictor == "quickdraw100":
+        return labels_from_map(load_label_map(QUICKDRAW100_LABEL_MAP_PATH))
+    return labels_from_map(load_label_map(LABEL_MAP_PATH))
+
+
+def get_predictor(predictor: str) -> tuple[nn.Module, list[str], torch.device, int]:
+    if predictor == "quickdraw100":
+        return quickdraw100_predictor()
+    return image_predictor()
+
+
+def predict_tensor(image: torch.Tensor, top_k: int = 5, predictor: str = "legacy20") -> dict[str, Any]:
+    model, labels, device, image_size = get_predictor(predictor)
+    tensor = tensor_image_to_input(image, image_size=image_size).to(device)
+    with torch.no_grad():
+        logits = model(tensor)
+        probs = torch.softmax(logits, dim=1).squeeze(0).detach().cpu()
+    predictions = top_predictions(probs, labels, top_k)
+    return {
+        "model": "cnn",
+        "predictor": predictor,
+        "input": "tensor",
+        "image_size": image_size,
+        "prediction": predictions[0],
+        "top": predictions,
+    }
+
+
+def predict_image(image_bytes: bytes, top_k: int = 5, predictor: str = "legacy20") -> dict[str, Any]:
+    model, labels, device, image_size = get_predictor(predictor)
+    tensor = uploaded_image_to_tensor(image_bytes, image_size=image_size).to(device)
+    with torch.no_grad():
+        logits = model(tensor)
+        probs = torch.softmax(logits, dim=1).squeeze(0).detach().cpu()
+    predictions = top_predictions(probs, labels, top_k)
+    return {
+        "model": "cnn",
+        "predictor": predictor,
+        "input": "image",
+        "image_size": image_size,
+        "prediction": predictions[0],
+        "top": predictions,
+    }

quickdraw_app/models/cnn_residual_100cls_64/best_model.pt

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

quickdraw_app/models/cnn_residual_100cls_64/label_map.json

@@ -0,0 +1,102 @@
+{
+  "aircraft carrier": 0,
+  "airplane": 1,
+  "alarm clock": 2,
+  "ambulance": 3,
+  "angel": 4,
+  "animal migration": 5,
+  "ant": 6,
+  "anvil": 7,
+  "apple": 8,
+  "arm": 9,
+  "asparagus": 10,
+  "axe": 11,
+  "backpack": 12,
+  "banana": 13,
+  "bandage": 14,
+  "barn": 15,
+  "baseball": 16,
+  "baseball bat": 17,
+  "basket": 18,
+  "basketball": 19,
+  "bat": 20,
+  "bathtub": 21,
+  "beach": 22,
+  "bear": 23,
+  "beard": 24,
+  "bed": 25,
+  "bee": 26,
+  "belt": 27,
+  "bench": 28,
+  "bicycle": 29,
+  "binoculars": 30,
+  "bird": 31,
+  "birthday cake": 32,
+  "blackberry": 33,
+  "blueberry": 34,
+  "book": 35,
+  "boomerang": 36,
+  "bottlecap": 37,
+  "bowtie": 38,
+  "bracelet": 39,
+  "brain": 40,
+  "bread": 41,
+  "bridge": 42,
+  "broccoli": 43,
+  "broom": 44,
+  "bucket": 45,
+  "bulldozer": 46,
+  "bus": 47,
+  "bush": 48,
+  "butterfly": 49,
+  "cactus": 50,
+  "cake": 51,
+  "calculator": 52,
+  "calendar": 53,
+  "camel": 54,
+  "camera": 55,
+  "camouflage": 56,
+  "campfire": 57,
+  "candle": 58,
+  "cannon": 59,
+  "canoe": 60,
+  "car": 61,
+  "carrot": 62,
+  "castle": 63,
+  "cat": 64,
+  "ceiling fan": 65,
+  "cello": 66,
+  "cell phone": 67,
+  "chair": 68,
+  "chandelier": 69,
+  "church": 70,
+  "circle": 71,
+  "clarinet": 72,
+  "clock": 73,
+  "cloud": 74,
+  "coffee cup": 75,
+  "compass": 76,
+  "computer": 77,
+  "cookie": 78,
+  "cooler": 79,
+  "couch": 80,
+  "cow": 81,
+  "crab": 82,
+  "crayon": 83,
+  "crocodile": 84,
+  "crown": 85,
+  "cruise ship": 86,
+  "cup": 87,
+  "diamond": 88,
+  "dishwasher": 89,
+  "diving board": 90,
+  "dog": 91,
+  "dolphin": 92,
+  "donut": 93,
+  "door": 94,
+  "dragon": 95,
+  "dresser": 96,
+  "drill": 97,
+  "drums": 98,
+  "duck": 99
+}

quickdraw_app/server.py

@@ -0,0 +1,192 @@
+#!/usr/bin/env python3
+"""Local web API for sampling the QuickDraw conditional DDPM."""
+
+from __future__ import annotations
+
+import base64
+import io
+import os
+import random
+import sys
+import threading
+from pathlib import Path
+
+import torch
+from fastapi import FastAPI, HTTPException
+from fastapi.responses import FileResponse
+from fastapi.staticfiles import StaticFiles
+from pydantic import BaseModel, Field
+from torchvision.utils import make_grid
+
+ROOT = Path(__file__).resolve().parents[1]
+if str(ROOT) not in sys.path:
+    sys.path.insert(0, str(ROOT))
+
+from train_quickdraw_ddpm import (  # noqa: E402
+    QUICKDRAW_100_CLASSES,
+    SmallConditionalUNet,
+    make_schedule,
+    pick_device,
+    sample,
+)
+
+from quickdraw_app import cnn_classifier  # noqa: E402
+
+
+DEFAULT_CHECKPOINT = ROOT / "models" / "diffusion" / "checkpoint_step_500000.pt"
+
+CHECKPOINT_PATH = Path(os.environ.get("QUICKDRAW_CHECKPOINT", DEFAULT_CHECKPOINT)).expanduser()
+STATIC_DIR = Path(__file__).resolve().parent / "static"
+
+
+class GenerateRequest(BaseModel):
+    class_name: str = Field(..., min_length=1)
+    count: int = Field(4, ge=1, le=8)
+    guidance_scale: float = Field(3.0, ge=0.0, le=8.0)
+    seed: int | None = Field(None, ge=0, le=2**31 - 1)
+
+
+class RecognizeRequest(BaseModel):
+    image: str = Field(..., min_length=1)
+    top_k: int = Field(5, ge=1, le=10)
+    predictor: str = Field("quickdraw100")
+
+
+class ModelState:
+    def __init__(self) -> None:
+        self.lock = threading.Lock()
+        self.loaded = False
+        self.device = pick_device()
+        self.model: SmallConditionalUNet | None = None
+        self.schedule = None
+        self.classes: list[str] = QUICKDRAW_100_CLASSES
+        self.image_size = 64
+        self.timesteps = 200
+        self.base_channels = 64
+        self.step: int | None = None
+
+    def load(self) -> None:
+        if self.loaded:
+            return
+        if not CHECKPOINT_PATH.exists():
+            raise FileNotFoundError(f"Checkpoint not found: {CHECKPOINT_PATH}")
+
+        checkpoint = torch.load(CHECKPOINT_PATH, map_location=self.device, weights_only=False)
+        self.classes = list(checkpoint.get("classes", QUICKDRAW_100_CLASSES))
+        self.image_size = int(checkpoint.get("image_size", 64))
+        self.timesteps = int(checkpoint.get("timesteps", 200))
+        self.base_channels = int(checkpoint.get("base_channels", 64))
+        self.step = checkpoint.get("step")
+
+        model = SmallConditionalUNet(len(self.classes), base_channels=self.base_channels).to(self.device)
+        state_dict = checkpoint.get("model_unwrapped") or checkpoint.get("model")
+        if state_dict is None:
+            raise RuntimeError("Checkpoint does not contain a model state dict.")
+        model.load_state_dict(state_dict)
+        model.eval()
+
+        self.model = model
+        self.schedule = make_schedule(self.timesteps, self.device)
+        self.loaded = True
+
+    def generate(self, request: GenerateRequest) -> str:
+        with self.lock:
+            self.load()
+            assert self.model is not None
+            assert self.schedule is not None
+
+            try:
+                class_index = self.classes.index(request.class_name)
+            except ValueError as exc:
+                raise HTTPException(status_code=400, detail="Unknown class name.") from exc
+
+            seed = request.seed if request.seed is not None else random.randrange(0, 2**31)
+            torch.manual_seed(seed)
+            if self.device.type == "cuda":
+                torch.cuda.manual_seed_all(seed)
+
+            labels = torch.full((request.count,), class_index, dtype=torch.long, device=self.device)
+            images = sample(
+                self.model,
+                labels,
+                self.image_size,
+                self.schedule,
+                self.timesteps,
+                self.device,
+                guidance_scale=request.guidance_scale,
+            )
+            grid = make_grid((images + 1) / 2, nrow=min(request.count, 4), padding=8, pad_value=0)
+            image = grid.mul(255).byte().permute(1, 2, 0).cpu().numpy()
+
+        from PIL import Image
+
+        buffer = io.BytesIO()
+        Image.fromarray(image).save(buffer, format="PNG")
+        encoded = base64.b64encode(buffer.getvalue()).decode("ascii")
+        return f"data:image/png;base64,{encoded}"
+
+
+state = ModelState()
+app = FastAPI(title="QuickDraw Diffusion")
+app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")
+
+
+@app.get("/")
+def index() -> FileResponse:
+    return FileResponse(STATIC_DIR / "index.html")
+
+
+@app.get("/api/status")
+def status() -> dict:
+    loaded = state.loaded
+    return {
+        "loaded": loaded,
+        "checkpoint": str(CHECKPOINT_PATH),
+        "checkpoint_exists": CHECKPOINT_PATH.exists(),
+        "device": str(state.device),
+        "step": state.step,
+        "image_size": state.image_size,
+        "timesteps": state.timesteps,
+        "cnn_model": str(cnn_classifier.MODEL_PATH),
+        "cnn_model_exists": cnn_classifier.MODEL_PATH.exists(),
+        "quickdraw100_cnn_model": str(cnn_classifier.QUICKDRAW100_MODEL_PATH),
+        "quickdraw100_cnn_model_exists": cnn_classifier.QUICKDRAW100_MODEL_PATH.exists(),
+    }
+
+
+@app.get("/api/classes")
+def classes() -> dict:
+    return {"classes": state.classes}
+
+
+@app.get("/api/cnn/classes")
+def cnn_classes(predictor: str = "quickdraw100") -> dict:
+    try:
+        return {"classes": cnn_classifier.classes(predictor=predictor), "predictor": predictor}
+    except FileNotFoundError as exc:
+        raise HTTPException(status_code=500, detail=str(exc)) from exc
+
+
+@app.post("/api/generate")
+def generate(request: GenerateRequest) -> dict:
+    try:
+        image = state.generate(request)
+    except FileNotFoundError as exc:
+        raise HTTPException(status_code=500, detail=str(exc)) from exc
+    return {
+        "image": image,
+        "class_name": request.class_name,
+        "count": request.count,
+        "guidance_scale": request.guidance_scale,
+        "step": state.step,
+        "device": str(state.device),
+    }
+
+
+@app.post("/api/recognize")
+def recognize(request: RecognizeRequest) -> dict:
+    image = cnn_classifier.decode_data_url(request.image)
+    try:
+        return cnn_classifier.predict_image(image, request.top_k, predictor=request.predictor)
+    except FileNotFoundError as exc:
+        raise HTTPException(status_code=500, detail=str(exc)) from exc

quickdraw_app/static/app.js

@@ -0,0 +1,167 @@
+const classSelect = document.querySelector("#class-select");
+const form = document.querySelector("#generate-form");
+const guidance = document.querySelector("#guidance");
+const guidanceValue = document.querySelector("#guidance-value");
+const statusEl = document.querySelector("#status");
+const imageStage = document.querySelector("#image-stage");
+const resultTitle = document.querySelector("#result-title");
+const message = document.querySelector("#message");
+const button = document.querySelector("#generate-button");
+const downloadLink = document.querySelector("#download-link");
+const recognizeButton = document.querySelector("#recognize-button");
+const recognitionList = document.querySelector("#recognition-list");
+const cnnStatus = document.querySelector("#cnn-status");
+
+let currentImage = null;
+
+function setMessage(text, isError = false) {
+  message.textContent = text;
+  message.classList.toggle("error", isError);
+}
+
+function setLoading(isLoading) {
+  button.disabled = isLoading;
+  button.querySelector("span:last-child").textContent = isLoading ? "生成中" : "生成";
+}
+
+async function loadStatus() {
+  const response = await fetch("/api/status");
+  const status = await response.json();
+  statusEl.textContent = status.checkpoint_exists
+    ? `${status.device} · ${status.image_size}x${status.image_size} · step ${status.step ?? "not loaded"}`
+    : "checkpoint missing";
+}
+
+async function loadCnnStatus() {
+  try {
+    const response = await fetch("/api/cnn/classes");
+    const data = await response.json();
+    if (!response.ok) {
+      throw new Error(data.detail || "CNN offline");
+    }
+    const classes = Array.isArray(data) ? data : data.classes;
+    cnnStatus.textContent = `CNN: ${classes.length} classes`;
+    cnnStatus.title = classes.join(", ");
+  } catch (error) {
+    cnnStatus.textContent = "CNN: offline";
+    cnnStatus.title = error.message;
+  }
+}
+
+async function loadClasses() {
+  const response = await fetch("/api/classes");
+  const data = await response.json();
+  classSelect.innerHTML = "";
+  for (const name of data.classes) {
+    const option = document.createElement("option");
+    option.value = name;
+    option.textContent = name;
+    classSelect.append(option);
+  }
+  classSelect.value = data.classes.includes("cat") ? "cat" : data.classes[0];
+}
+
+guidance.addEventListener("input", () => {
+  guidanceValue.value = Number(guidance.value).toFixed(2);
+});
+
+form.addEventListener("submit", async (event) => {
+  event.preventDefault();
+  const formData = new FormData(form);
+  const seedValue = formData.get("seed");
+  const payload = {
+    class_name: formData.get("class_name"),
+    count: Number(formData.get("count")),
+    guidance_scale: Number(formData.get("guidance_scale")),
+    seed: seedValue ? Number(seedValue) : null,
+  };
+
+  setLoading(true);
+  setMessage("模型采样大约需要几十秒，MPS/CPU 会更慢。");
+
+  try {
+    const response = await fetch("/api/generate", {
+      method: "POST",
+      headers: { "Content-Type": "application/json" },
+      body: JSON.stringify(payload),
+    });
+    const data = await response.json();
+    if (!response.ok) {
+      throw new Error(data.detail || "生成失败");
+    }
+
+    imageStage.innerHTML = "";
+    const image = document.createElement("img");
+    image.src = data.image;
+    image.alt = `${data.class_name} generated sketch samples`;
+    imageStage.append(image);
+    currentImage = data.image;
+    resultTitle.textContent = data.class_name;
+    downloadLink.href = data.image;
+    downloadLink.setAttribute("aria-disabled", "false");
+    recognizeButton.disabled = false;
+    recognitionList.innerHTML = "";
+    setMessage(`完成 · ${data.device} · step ${data.step} · CFG ${data.guidance_scale}`);
+    await loadStatus();
+  } catch (error) {
+    setMessage(error.message, true);
+  } finally {
+    setLoading(false);
+  }
+});
+
+recognizeButton.addEventListener("click", async () => {
+  if (!currentImage) {
+    return;
+  }
+
+  recognizeButton.disabled = true;
+  recognizeButton.textContent = "识别中";
+  setMessage("正在请求 CNN 图片识别服务。");
+
+  try {
+    const response = await fetch("/api/recognize", {
+      method: "POST",
+      headers: { "Content-Type": "application/json" },
+    body: JSON.stringify({ image: currentImage, top_k: 5, predictor: "quickdraw100" }),
+    });
+    const data = await response.json();
+    if (!response.ok) {
+      throw new Error(data.detail || "识别失败");
+    }
+
+    recognitionList.innerHTML = "";
+    for (const item of data.top || []) {
+      const row = document.createElement("div");
+      row.className = "recognition-item";
+
+      const label = document.createElement("strong");
+      label.textContent = item.label;
+
+      const track = document.createElement("div");
+      track.className = "confidence-track";
+      const fill = document.createElement("div");
+      fill.className = "confidence-fill";
+      fill.style.width = `${Math.max(0, Math.min(1, item.confidence)) * 100}%`;
+      track.append(fill);
+
+      const value = document.createElement("span");
+      value.className = "confidence-value";
+      value.textContent = `${Math.round(item.confidence * 100)}%`;
+
+      row.append(label, track, value);
+      recognitionList.append(row);
+    }
+    const top = data.prediction;
+    setMessage(top ? `CNN top-1: ${top.label} · ${(top.confidence * 100).toFixed(1)}%` : "CNN 返回为空");
+  } catch (error) {
+    setMessage(error.message, true);
+  } finally {
+    recognizeButton.disabled = false;
+    recognizeButton.textContent = "识别当前图片";
+  }
+});
+
+Promise.all([loadStatus(), loadClasses(), loadCnnStatus()]).catch((error) => {
+  setMessage(error.message, true);
+});

quickdraw_app/static/index.html

@@ -0,0 +1,81 @@
+<!doctype html>
+<html lang="zh-CN">
+  <head>
+    <meta charset="utf-8" />
+    <meta name="viewport" content="width=device-width, initial-scale=1" />
+    <title>QuickDraw Diffusion</title>
+    <link rel="stylesheet" href="/static/styles.css" />
+  </head>
+  <body>
+    <main class="app-shell">
+      <section class="workspace" aria-labelledby="title">
+        <div class="toolbar">
+          <div>
+            <p class="eyebrow">QuickDraw Diffusion</p>
+            <h1 id="title">Sketch Generator</h1>
+          </div>
+          <div class="status" id="status">Checking model...</div>
+        </div>
+
+        <div class="layout">
+          <form class="panel controls" id="generate-form">
+            <label class="field">
+              <span>类别</span>
+              <select id="class-select" name="class_name" required></select>
+            </label>
+
+            <div class="field">
+              <span>数量</span>
+              <div class="segmented" role="radiogroup" aria-label="生成数量">
+                <label><input type="radio" name="count" value="1" />1</label>
+                <label><input type="radio" name="count" value="4" checked />4</label>
+                <label><input type="radio" name="count" value="8" />8</label>
+              </div>
+            </div>
+
+            <label class="field">
+              <span>CFG</span>
+              <div class="range-row">
+                <input id="guidance" name="guidance_scale" type="range" min="0" max="8" step="0.25" value="3" />
+                <output id="guidance-value" for="guidance">3.00</output>
+              </div>
+            </label>
+
+            <label class="field">
+              <span>Seed</span>
+              <input id="seed" name="seed" type="number" min="0" max="2147483647" placeholder="随机" />
+            </label>
+
+            <button class="generate-button" type="submit" id="generate-button">
+              <span class="button-icon" aria-hidden="true"></span>
+              <span>生成</span>
+            </button>
+          </form>
+
+          <section class="panel preview" aria-live="polite">
+            <div class="preview-header">
+              <div>
+                <p class="eyebrow">Result</p>
+                <h2 id="result-title">选择类别后生成</h2>
+              </div>
+              <a id="download-link" class="download-link" href="#" download="quickdraw-sample.png" aria-disabled="true">下载</a>
+            </div>
+            <div class="image-stage" id="image-stage">
+              <div class="empty-state">等待生成</div>
+            </div>
+            <div class="recognition">
+              <div class="recognition-actions">
+                <button class="recognize-button" type="button" id="recognize-button" disabled>识别当前图片</button>
+                <span id="cnn-status">CNN: checking...</span>
+              </div>
+              <div class="recognition-list" id="recognition-list"></div>
+            </div>
+            <p class="message" id="message"></p>
+          </section>
+        </div>
+      </section>
+    </main>
+
+    <script src="/static/app.js"></script>
+  </body>
+</html>

quickdraw_app/static/styles.css

@@ -0,0 +1,386 @@
+:root {
+  color-scheme: light;
+  --bg: #f6f7f4;
+  --surface: #ffffff;
+  --surface-soft: #eef1ea;
+  --ink: #18201d;
+  --muted: #66726d;
+  --line: #d8ded6;
+  --primary: #16615a;
+  --primary-strong: #0d443f;
+  --accent: #b7472a;
+  --shadow: 0 18px 50px rgba(31, 43, 37, 0.12);
+}
+
+* {
+  box-sizing: border-box;
+}
+
+body {
+  margin: 0;
+  min-height: 100dvh;
+  background: var(--bg);
+  color: var(--ink);
+  font-family: Inter, ui-sans-serif, system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif;
+  letter-spacing: 0;
+}
+
+button,
+input,
+select {
+  font: inherit;
+}
+
+.app-shell {
+  min-height: 100dvh;
+  padding: 28px;
+}
+
+.workspace {
+  max-width: 1180px;
+  margin: 0 auto;
+}
+
+.toolbar {
+  display: flex;
+  align-items: end;
+  justify-content: space-between;
+  gap: 20px;
+  margin-bottom: 20px;
+}
+
+.eyebrow {
+  margin: 0 0 6px;
+  color: var(--accent);
+  font-size: 12px;
+  font-weight: 700;
+  text-transform: uppercase;
+}
+
+h1,
+h2 {
+  margin: 0;
+  line-height: 1.1;
+}
+
+h1 {
+  font-size: clamp(32px, 5vw, 56px);
+}
+
+h2 {
+  font-size: 24px;
+}
+
+.status {
+  min-height: 38px;
+  padding: 9px 13px;
+  border: 1px solid var(--line);
+  border-radius: 8px;
+  background: var(--surface);
+  color: var(--muted);
+  font-size: 14px;
+}
+
+.layout {
+  display: grid;
+  grid-template-columns: 340px minmax(0, 1fr);
+  gap: 18px;
+  align-items: start;
+}
+
+.panel {
+  border: 1px solid var(--line);
+  border-radius: 8px;
+  background: var(--surface);
+  box-shadow: var(--shadow);
+}
+
+.controls {
+  display: grid;
+  gap: 18px;
+  padding: 18px;
+}
+
+.field {
+  display: grid;
+  gap: 8px;
+  color: var(--muted);
+  font-size: 14px;
+  font-weight: 700;
+}
+
+select,
+input[type="number"] {
+  width: 100%;
+  min-height: 46px;
+  border: 1px solid var(--line);
+  border-radius: 8px;
+  background: var(--surface);
+  color: var(--ink);
+  padding: 0 12px;
+}
+
+select:focus,
+input:focus,
+button:focus-visible,
+a:focus-visible {
+  outline: 3px solid rgba(22, 97, 90, 0.22);
+  outline-offset: 2px;
+}
+
+.segmented {
+  display: grid;
+  grid-template-columns: repeat(3, 1fr);
+  gap: 6px;
+  padding: 4px;
+  border: 1px solid var(--line);
+  border-radius: 8px;
+  background: var(--surface-soft);
+}
+
+.segmented label {
+  min-height: 40px;
+  display: grid;
+  place-items: center;
+  border-radius: 6px;
+  color: var(--muted);
+  cursor: pointer;
+}
+
+.segmented input {
+  position: absolute;
+  opacity: 0;
+}
+
+.segmented label:has(input:checked) {
+  background: var(--surface);
+  color: var(--ink);
+  box-shadow: 0 1px 6px rgba(31, 43, 37, 0.12);
+}
+
+.range-row {
+  display: grid;
+  grid-template-columns: minmax(0, 1fr) 58px;
+  gap: 10px;
+  align-items: center;
+}
+
+input[type="range"] {
+  width: 100%;
+  accent-color: var(--primary);
+}
+
+output {
+  color: var(--ink);
+  font-variant-numeric: tabular-nums;
+}
+
+.generate-button {
+  min-height: 48px;
+  display: inline-flex;
+  align-items: center;
+  justify-content: center;
+  gap: 10px;
+  border: 0;
+  border-radius: 8px;
+  background: var(--primary);
+  color: white;
+  font-weight: 800;
+  cursor: pointer;
+  transition: transform 160ms ease, background 160ms ease, opacity 160ms ease;
+}
+
+.generate-button:hover {
+  background: var(--primary-strong);
+}
+
+.generate-button:active {
+  transform: scale(0.98);
+}
+
+.generate-button:disabled {
+  cursor: wait;
+  opacity: 0.64;
+}
+
+.recognize-button {
+  min-height: 40px;
+  display: inline-flex;
+  align-items: center;
+  justify-content: center;
+  border: 1px solid var(--line);
+  border-radius: 8px;
+  background: var(--surface);
+  color: var(--primary);
+  padding: 0 14px;
+  font-weight: 800;
+  cursor: pointer;
+}
+
+.recognize-button:disabled {
+  cursor: not-allowed;
+  color: var(--muted);
+  opacity: 0.55;
+}
+
+.button-icon {
+  width: 16px;
+  height: 16px;
+  border: 2px solid currentColor;
+  border-left-color: transparent;
+  border-radius: 50%;
+}
+
+.generate-button:not(:disabled) .button-icon {
+  border-left-color: currentColor;
+  border-radius: 4px;
+}
+
+.generate-button:disabled .button-icon {
+  animation: spin 800ms linear infinite;
+}
+
+.preview {
+  min-height: 650px;
+  padding: 18px;
+}
+
+.preview-header {
+  display: flex;
+  align-items: center;
+  justify-content: space-between;
+  gap: 16px;
+  margin-bottom: 16px;
+}
+
+.download-link {
+  min-height: 40px;
+  display: inline-flex;
+  align-items: center;
+  justify-content: center;
+  padding: 0 14px;
+  border: 1px solid var(--line);
+  border-radius: 8px;
+  color: var(--primary);
+  text-decoration: none;
+  font-weight: 800;
+}
+
+.download-link[aria-disabled="true"] {
+  pointer-events: none;
+  color: var(--muted);
+  opacity: 0.5;
+}
+
+.image-stage {
+  min-height: 540px;
+  display: grid;
+  place-items: center;
+  border: 1px solid #151916;
+  border-radius: 8px;
+  background: #050706;
+  overflow: hidden;
+}
+
+.image-stage img {
+  max-width: 100%;
+  height: auto;
+  image-rendering: pixelated;
+}
+
+.empty-state {
+  color: #b9c5bf;
+  font-weight: 700;
+}
+
+.message {
+  min-height: 24px;
+  margin: 12px 0 0;
+  color: var(--muted);
+  font-size: 14px;
+}
+
+.recognition {
+  display: grid;
+  gap: 10px;
+  margin-top: 14px;
+}
+
+.recognition-actions {
+  display: flex;
+  align-items: center;
+  justify-content: space-between;
+  gap: 12px;
+  color: var(--muted);
+  font-size: 14px;
+}
+
+.recognition-list {
+  display: grid;
+  gap: 8px;
+}
+
+.recognition-item {
+  display: grid;
+  grid-template-columns: 110px minmax(0, 1fr) 56px;
+  align-items: center;
+  gap: 10px;
+  min-height: 36px;
+  color: var(--ink);
+  font-size: 14px;
+}
+
+.confidence-track {
+  height: 8px;
+  border-radius: 999px;
+  background: var(--surface-soft);
+  overflow: hidden;
+}
+
+.confidence-fill {
+  height: 100%;
+  border-radius: inherit;
+  background: var(--primary);
+}
+
+.confidence-value {
+  color: var(--muted);
+  font-variant-numeric: tabular-nums;
+  text-align: right;
+}
+
+.message.error {
+  color: #a12a22;
+}
+
+@keyframes spin {
+  to {
+    transform: rotate(360deg);
+  }
+}
+
+@media (max-width: 820px) {
+  .app-shell {
+    padding: 16px;
+  }
+
+  .toolbar,
+  .preview-header,
+  .recognition-actions {
+    align-items: stretch;
+    flex-direction: column;
+  }
+
+  .recognition-item {
+    grid-template-columns: 88px minmax(0, 1fr) 50px;
+  }
+
+  .layout {
+    grid-template-columns: 1fr;
+  }
+
+  .preview,
+  .image-stage {
+    min-height: 420px;
+  }
+}

requirements.txt

@@ -0,0 +1,7 @@
+fastapi
+uvicorn
+numpy
+pillow
+torch
+torchvision
+tqdm

train_quickdraw_ddpm.py

@@ -0,0 +1,434 @@
+#!/usr/bin/env python3
+"""Train a small class-conditional DDPM on rasterized QuickDraw sketches."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import random
+import time
+import urllib.parse
+import urllib.request
+from dataclasses import dataclass
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from PIL import Image, ImageDraw
+from torch.utils.data import DataLoader, Dataset
+from torchvision.utils import save_image
+from tqdm import tqdm
+
+
+QUICKDRAW_URL = "https://storage.googleapis.com/quickdraw_dataset/full/simplified/{word}.ndjson"
+QUICKDRAW_100_CLASSES = [
+    "aircraft carrier", "airplane", "alarm clock", "ambulance", "angel",
+    "animal migration", "ant", "anvil", "apple", "arm", "asparagus", "axe",
+    "backpack", "banana", "bandage", "barn", "baseball", "baseball bat",
+    "basket", "basketball", "bat", "bathtub", "beach", "bear", "beard",
+    "bed", "bee", "belt", "bench", "bicycle", "binoculars", "bird",
+    "birthday cake", "blackberry", "blueberry", "book", "boomerang",
+    "bottlecap", "bowtie", "bracelet", "brain", "bread", "bridge",
+    "broccoli", "broom", "bucket", "bulldozer", "bus", "bush", "butterfly",
+    "cactus", "cake", "calculator", "calendar", "camel", "camera",
+    "camouflage", "campfire", "candle", "cannon", "canoe", "car", "carrot",
+    "castle", "cat", "ceiling fan", "cello", "cell phone", "chair",
+    "chandelier", "church", "circle", "clarinet", "clock", "cloud",
+    "coffee cup", "compass", "computer", "cookie", "cooler", "couch",
+    "cow", "crab", "crayon", "crocodile", "crown", "cruise ship", "cup",
+    "diamond", "dishwasher", "diving board", "dog", "dolphin", "donut",
+    "door", "dragon", "dresser", "drill", "drums", "duck",
+]
+
+
+def unwrap_model(model: nn.Module) -> nn.Module:
+    return model.module if isinstance(model, nn.DataParallel) else model
+
+
+def pick_device() -> torch.device:
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    if torch.backends.mps.is_available():
+        return torch.device("mps")
+    return torch.device("cpu")
+
+
+def render_drawing(drawing: list, image_size: int, line_width: int) -> torch.Tensor:
+    image = Image.new("L", (image_size, image_size), 255)
+    draw = ImageDraw.Draw(image)
+    scale = image_size / 256.0
+
+    for stroke in drawing:
+        xs, ys = stroke
+        points = [(round(x * scale), round(y * scale)) for x, y in zip(xs, ys)]
+        if len(points) >= 2:
+            draw.line(points, fill=0, width=line_width)
+        elif len(points) == 1:
+            x, y = points[0]
+            r = max(1, line_width // 2)
+            draw.ellipse((x - r, y - r, x + r, y + r), fill=0)
+
+    data = torch.tensor(list(image.tobytes()), dtype=torch.uint8).view(1, image_size, image_size)
+    return 255 - data
+
+
+class QuickDrawSketches(Dataset):
+    def __init__(
+        self,
+        classes: list[str],
+        samples_per_class: int,
+        image_size: int,
+        line_width: int,
+        recognized_only: bool = True,
+        download_retries: int = 5,
+    ) -> None:
+        self.classes = classes
+        total_samples = len(classes) * samples_per_class
+        images = torch.empty(total_samples, 1, image_size, image_size, dtype=torch.uint8)
+        labels = torch.empty(total_samples, dtype=torch.long)
+
+        for label, word in enumerate(classes):
+            quoted = urllib.parse.quote(word, safe="")
+            url = QUICKDRAW_URL.format(word=quoted)
+            for attempt in range(1, download_retries + 1):
+                loaded = 0
+                try:
+                    with urllib.request.urlopen(url, timeout=60) as response:
+                        for raw_line in response:
+                            item = json.loads(raw_line)
+                            if recognized_only and not item.get("recognized", False):
+                                continue
+                            index = label * samples_per_class + loaded
+                            images[index] = render_drawing(item["drawing"], image_size, line_width)
+                            labels[index] = label
+                            loaded += 1
+                            if loaded >= samples_per_class:
+                                break
+                    if loaded >= samples_per_class:
+                        print(f"loaded class {label + 1}/{len(classes)}: {word} ({loaded})", flush=True)
+                        break
+                    raise RuntimeError(f"Only loaded {loaded} samples for class {word!r}")
+                except Exception:
+                    if attempt == download_retries:
+                        raise
+                    time.sleep(min(2 ** attempt, 30))
+        self.images = images
+        self.labels = labels
+
+    def __len__(self) -> int:
+        return self.images.shape[0]
+
+    def __getitem__(self, index: int) -> tuple[torch.Tensor, torch.Tensor]:
+        image = self.images[index].float() / 127.5 - 1.0
+        return image, self.labels[index]
+
+
+class SinusoidalTimeEmbedding(nn.Module):
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, t: torch.Tensor) -> torch.Tensor:
+        half = self.dim // 2
+        freqs = torch.exp(
+            -math.log(10000) * torch.arange(half, device=t.device).float() / max(half - 1, 1)
+        )
+        args = t.float().unsqueeze(1) * freqs.unsqueeze(0)
+        emb = torch.cat([args.sin(), args.cos()], dim=1)
+        if self.dim % 2:
+            emb = F.pad(emb, (0, 1))
+        return emb
+
+
+class ResBlock(nn.Module):
+    def __init__(self, in_ch: int, out_ch: int, emb_dim: int) -> None:
+        super().__init__()
+        self.norm1 = nn.GroupNorm(min(8, in_ch), in_ch)
+        self.conv1 = nn.Conv2d(in_ch, out_ch, 3, padding=1)
+        self.emb = nn.Linear(emb_dim, out_ch)
+        self.norm2 = nn.GroupNorm(min(8, out_ch), out_ch)
+        self.conv2 = nn.Conv2d(out_ch, out_ch, 3, padding=1)
+        self.skip = nn.Conv2d(in_ch, out_ch, 1) if in_ch != out_ch else nn.Identity()
+
+    def forward(self, x: torch.Tensor, emb: torch.Tensor) -> torch.Tensor:
+        h = self.conv1(F.silu(self.norm1(x)))
+        h = h + self.emb(F.silu(emb))[:, :, None, None]
+        h = self.conv2(F.silu(self.norm2(h)))
+        return h + self.skip(x)
+
+
+class SmallConditionalUNet(nn.Module):
+    def __init__(self, num_classes: int, base_channels: int = 64, emb_dim: int = 256) -> None:
+        super().__init__()
+        self.num_classes = num_classes
+        self.null_label = num_classes
+        self.time_mlp = nn.Sequential(
+            SinusoidalTimeEmbedding(emb_dim),
+            nn.Linear(emb_dim, emb_dim),
+            nn.SiLU(),
+            nn.Linear(emb_dim, emb_dim),
+        )
+        self.class_emb = nn.Embedding(num_classes + 1, emb_dim)
+
+        c = base_channels
+        self.in_conv = nn.Conv2d(1, c, 3, padding=1)
+        self.down1 = ResBlock(c, c, emb_dim)
+        self.downsample1 = nn.Conv2d(c, c * 2, 4, stride=2, padding=1)
+        self.down2 = ResBlock(c * 2, c * 2, emb_dim)
+        self.downsample2 = nn.Conv2d(c * 2, c * 4, 4, stride=2, padding=1)
+        self.mid1 = ResBlock(c * 4, c * 4, emb_dim)
+        self.mid2 = ResBlock(c * 4, c * 4, emb_dim)
+        self.upsample2 = nn.ConvTranspose2d(c * 4, c * 2, 4, stride=2, padding=1)
+        self.up2 = ResBlock(c * 4, c * 2, emb_dim)
+        self.upsample1 = nn.ConvTranspose2d(c * 2, c, 4, stride=2, padding=1)
+        self.up1 = ResBlock(c * 2, c, emb_dim)
+        self.out_norm = nn.GroupNorm(min(8, c), c)
+        self.out_conv = nn.Conv2d(c, 1, 3, padding=1)
+
+    def forward(self, x: torch.Tensor, t: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+        emb = self.time_mlp(t) + self.class_emb(y)
+        x0 = self.in_conv(x)
+        x1 = self.down1(x0, emb)
+        x2 = self.down2(self.downsample1(x1), emb)
+        x3 = self.mid2(self.mid1(self.downsample2(x2), emb), emb)
+        x = self.upsample2(x3)
+        x = self.up2(torch.cat([x, x2], dim=1), emb)
+        x = self.upsample1(x)
+        x = self.up1(torch.cat([x, x1], dim=1), emb)
+        return self.out_conv(F.silu(self.out_norm(x)))
+
+
+@dataclass
+class DiffusionSchedule:
+    betas: torch.Tensor
+    alphas: torch.Tensor
+    alphas_cumprod: torch.Tensor
+    alphas_cumprod_prev: torch.Tensor
+    sqrt_alphas_cumprod: torch.Tensor
+    sqrt_one_minus_alphas_cumprod: torch.Tensor
+    posterior_variance: torch.Tensor
+
+
+def make_schedule(timesteps: int, device: torch.device) -> DiffusionSchedule:
+    steps = timesteps + 1
+    x = torch.linspace(0, timesteps, steps, device=device)
+    alphas_cumprod = torch.cos(((x / timesteps) + 0.008) / 1.008 * math.pi * 0.5) ** 2
+    alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
+    betas = 1.0 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
+    betas = betas.clamp(1e-4, 0.999)
+    alphas = 1.0 - betas
+    alphas_cumprod = torch.cumprod(alphas, dim=0)
+    alphas_cumprod_prev = F.pad(alphas_cumprod[:-1], (1, 0), value=1.0)
+    posterior_variance = betas * (1.0 - alphas_cumprod_prev) / (1.0 - alphas_cumprod)
+    return DiffusionSchedule(
+        betas=betas,
+        alphas=alphas,
+        alphas_cumprod=alphas_cumprod,
+        alphas_cumprod_prev=alphas_cumprod_prev,
+        sqrt_alphas_cumprod=torch.sqrt(alphas_cumprod),
+        sqrt_one_minus_alphas_cumprod=torch.sqrt(1.0 - alphas_cumprod),
+        posterior_variance=posterior_variance,
+    )
+
+
+def extract(values: torch.Tensor, t: torch.Tensor, x_shape: torch.Size) -> torch.Tensor:
+    return values.gather(0, t).view(t.shape[0], *((1,) * (len(x_shape) - 1)))
+
+
+def q_sample(x0: torch.Tensor, t: torch.Tensor, noise: torch.Tensor, schedule: DiffusionSchedule) -> torch.Tensor:
+    return (
+        extract(schedule.sqrt_alphas_cumprod, t, x0.shape) * x0
+        + extract(schedule.sqrt_one_minus_alphas_cumprod, t, x0.shape) * noise
+    )
+
+
+@torch.no_grad()
+def sample(
+    model: nn.Module,
+    labels: torch.Tensor,
+    image_size: int,
+    schedule: DiffusionSchedule,
+    timesteps: int,
+    device: torch.device,
+    guidance_scale: float = 1.0,
+) -> torch.Tensor:
+    model.eval()
+    x = torch.randn(labels.shape[0], 1, image_size, image_size, device=device)
+    null_labels = torch.full_like(labels, unwrap_model(model).null_label)
+    for step in tqdm(reversed(range(timesteps)), total=timesteps, desc="sample"):
+        t = torch.full((labels.shape[0],), step, device=device, dtype=torch.long)
+        if guidance_scale == 1.0:
+            pred_noise = model(x, t, labels)
+        else:
+            pred_uncond = model(x, t, null_labels)
+            pred_cond = model(x, t, labels)
+            pred_noise = pred_uncond + guidance_scale * (pred_cond - pred_uncond)
+        alpha_bar_t = extract(schedule.alphas_cumprod, t, x.shape)
+        alpha_bar_prev = extract(schedule.alphas_cumprod_prev, t, x.shape)
+        beta_t = extract(schedule.betas, t, x.shape)
+        alpha_t = extract(schedule.alphas, t, x.shape)
+        pred_x0 = (x - torch.sqrt(1.0 - alpha_bar_t) * pred_noise) / torch.sqrt(alpha_bar_t)
+        pred_x0 = pred_x0.clamp(-1, 1)
+        coef_x0 = beta_t * torch.sqrt(alpha_bar_prev) / (1.0 - alpha_bar_t)
+        coef_xt = (1.0 - alpha_bar_prev) * torch.sqrt(alpha_t) / (1.0 - alpha_bar_t)
+        mean = coef_x0 * pred_x0 + coef_xt * x
+        if step > 0:
+            variance = extract(schedule.posterior_variance, t, x.shape)
+            x = mean + torch.sqrt(variance.clamp_min(1e-20)) * torch.randn_like(x)
+        else:
+            x = mean
+    return x.clamp(-1, 1)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--classes", nargs="+", default=["cat", "dog", "house", "airplane"])
+    parser.add_argument("--num-classes", type=int, default=0)
+    parser.add_argument("--samples-per-class", type=int, default=1000)
+    parser.add_argument("--image-size", type=int, default=64)
+    parser.add_argument("--line-width", type=int, default=2)
+    parser.add_argument("--batch-size", type=int, default=64)
+    parser.add_argument("--steps", type=int, default=1000)
+    parser.add_argument("--timesteps", type=int, default=200)
+    parser.add_argument("--lr", type=float, default=2e-4)
+    parser.add_argument("--base-channels", type=int, default=48)
+    parser.add_argument("--seed", type=int, default=7)
+    parser.add_argument("--out-dir", type=Path, default=Path("runs/quickdraw-ddpm"))
+    parser.add_argument("--sample-every", type=int, default=250)
+    parser.add_argument("--save-every", type=int, default=500)
+    parser.add_argument("--cfg-drop-prob", type=float, default=0.1)
+    parser.add_argument("--guidance-scale", type=float, default=3.0)
+    parser.add_argument("--download-retries", type=int, default=5)
+    parser.add_argument("--data-parallel", action="store_true")
+    parser.add_argument("--sample-num-classes", type=int, default=16)
+    parser.add_argument("--resume", type=Path, default=None)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    random.seed(args.seed)
+    torch.manual_seed(args.seed)
+    if args.num_classes:
+        if args.num_classes > len(QUICKDRAW_100_CLASSES):
+            raise ValueError(f"--num-classes supports at most {len(QUICKDRAW_100_CLASSES)} built-in classes")
+        args.classes = QUICKDRAW_100_CLASSES[: args.num_classes]
+    resume_checkpoint = None
+    if args.resume is not None:
+        resume_checkpoint = torch.load(args.resume, map_location="cpu", weights_only=False)
+        args.classes = list(resume_checkpoint["classes"])
+        args.image_size = int(resume_checkpoint["image_size"])
+        args.timesteps = int(resume_checkpoint["timesteps"])
+        args.base_channels = int(resume_checkpoint["base_channels"])
+
+    run_dir = args.out_dir / time.strftime("%Y%m%d-%H%M%S")
+    run_dir.mkdir(parents=True, exist_ok=True)
+    device = pick_device()
+
+    print(f"device: {device}")
+    print(f"classes: {args.classes}")
+    print("loading and rasterizing QuickDraw samples...")
+    dataset = QuickDrawSketches(
+        classes=args.classes,
+        samples_per_class=args.samples_per_class,
+        image_size=args.image_size,
+        line_width=args.line_width,
+        download_retries=args.download_retries,
+    )
+    loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=True, drop_last=True)
+
+    model = SmallConditionalUNet(len(args.classes), base_channels=args.base_channels).to(device)
+    if args.data_parallel:
+        if device.type != "cuda" or torch.cuda.device_count() < 2:
+            raise RuntimeError("--data-parallel requires at least two visible CUDA devices")
+        model = nn.DataParallel(model)
+        print(f"data_parallel_devices: {torch.cuda.device_count()}")
+    schedule = make_schedule(args.timesteps, device)
+    opt = torch.optim.AdamW(model.parameters(), lr=args.lr)
+    start_step = 0
+    if resume_checkpoint is not None:
+        state_dict = resume_checkpoint.get("model_unwrapped") or resume_checkpoint["model"]
+        unwrap_model(model).load_state_dict(state_dict)
+        opt.load_state_dict(resume_checkpoint["optimizer"])
+        start_step = int(resume_checkpoint["step"])
+        print(f"resumed checkpoint: {args.resume} at step {start_step}", flush=True)
+
+    with (run_dir / "config.json").open("w") as f:
+        json.dump(
+            vars(args) | {"device": str(device), "run_dir": str(run_dir), "start_step": start_step},
+            f,
+            indent=2,
+            default=str,
+        )
+
+    data_iter = iter(loader)
+    pbar = tqdm(range(start_step + 1, args.steps + 1), desc="train")
+    last_loss = None
+    for step in pbar:
+        try:
+            x0, labels = next(data_iter)
+        except StopIteration:
+            data_iter = iter(loader)
+            x0, labels = next(data_iter)
+
+        x0 = x0.to(device)
+        labels = labels.to(device)
+        if args.cfg_drop_prob > 0:
+            drop_mask = torch.rand(labels.shape, device=device) < args.cfg_drop_prob
+            labels_for_model = labels.masked_fill(drop_mask, unwrap_model(model).null_label)
+        else:
+            labels_for_model = labels
+        t = torch.randint(0, args.timesteps, (x0.shape[0],), device=device)
+        noise = torch.randn_like(x0)
+        xt = q_sample(x0, t, noise, schedule)
+        pred_noise = model(xt, t, labels_for_model)
+        loss = F.mse_loss(pred_noise, noise)
+
+        opt.zero_grad(set_to_none=True)
+        loss.backward()
+        nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+        opt.step()
+
+        last_loss = float(loss.item())
+        pbar.set_postfix(loss=f"{last_loss:.4f}")
+
+        if step % args.sample_every == 0 or step == args.steps:
+            sample_class_count = min(args.sample_num_classes, len(args.classes))
+            sample_labels = torch.arange(sample_class_count, device=device).repeat_interleave(4)
+            images = sample(
+                model,
+                sample_labels,
+                args.image_size,
+                schedule,
+                args.timesteps,
+                device,
+                guidance_scale=args.guidance_scale,
+            )
+            save_image((images + 1) / 2, run_dir / f"samples_step_{step:06d}.png", nrow=4)
+            model.train()
+
+        if step % args.save_every == 0 or step == args.steps:
+            torch.save(
+                {
+                    "model": model.state_dict(),
+                    "model_unwrapped": unwrap_model(model).state_dict(),
+                    "optimizer": opt.state_dict(),
+                    "step": step,
+                    "classes": args.classes,
+                    "image_size": args.image_size,
+                    "timesteps": args.timesteps,
+                    "base_channels": args.base_channels,
+                    "cfg_drop_prob": args.cfg_drop_prob,
+                    "guidance_scale": args.guidance_scale,
+                    "loss": last_loss,
+                },
+                run_dir / f"checkpoint_step_{step:06d}.pt",
+            )
+
+    print(f"done: {run_dir}")
+
+
+if __name__ == "__main__":
+    main()