sketch-model

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:

pip install -r requirements.txt

Run the local web app:

uvicorn quickdraw_app.server:app --host 127.0.0.1 --port 7860

Then open:

http://127.0.0.1:7860

The first generation loads the checkpoint and may take longer. Later requests reuse the loaded model.

Python Sampling

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:

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:

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:

https://storage.googleapis.com/quickdraw_dataset/full/simplified/{word}.ndjson