File size: 4,418 Bytes
9894238 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 | ---
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
```
|