---
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
```