You can deploy this model and start earning money today!

jamesaasher/
quickdraw-text-diffusion

Xenova/quickdraw-small

mit

Views

No views yet

Text-Conditional QuickDraw Diffusion Model

A text-conditional diffusion model for generating Google QuickDraw-style sketches from text prompts. This model uses DDPM (Denoising Diffusion Probabilistic Models) with CLIP text encoding and classifier-free guidance to generate 64x64 grayscale sketches.

Model Description

This is a U-Net based diffusion model that generates sketches conditioned on text prompts. It uses:

CLIP text encoder (openai/clip-vit-base-patch32) for text conditioning
DDPM for the diffusion process (1000 timesteps)
Classifier-free guidance for improved text-image alignment
Trained on Google QuickDraw dataset

Model Details

Model Type: Text-conditional DDPM diffusion model
Architecture: U-Net with cross-attention for text conditioning
Image Size: 64x64 grayscale
Base Channels: 256
Text Encoder: CLIP ViT-B/32 (frozen)
Training Steps: 100 epochs
Diffusion Timesteps: 1000
Guidance Scale: 5.0 (default)

Training Configuration

Dataset: Xenova/quickdraw-small (5 classes)
Batch Size: 128 (32 per GPU × 4 GPUs)
Learning Rate: 1e-4
CFG Drop Probability: 0.15
Optimizer: Adam

Usage

Installation

pip install torch torchvision transformers diffusers datasets matplotlib pillow tqdm

Generate Images

python

import torch
from model import TextConditionedUNet
from scheduler import SimpleDDPMScheduler
from text_encoder import CLIPTextEncoder
from generate import generate_samples

# Load checkpoint
checkpoint_path = "text_diffusion_final_epoch_100.pt"
checkpoint = torch.load(checkpoint_path)

# Initialize model
model = TextConditionedUNet(text_dim=512).cuda()
model.load_state_dict(checkpoint['model_state_dict'])
model.eval()

# Initialize text encoder
text_encoder = CLIPTextEncoder(model_name="openai/clip-vit-base-patch32", freeze=True).cuda()
text_encoder.eval()

# Generate samples
scheduler = SimpleDDPMScheduler(1000)
prompt = "a drawing of a cat"
num_samples = 4
guidance_scale = 5.0

with torch.no_grad():
    text_embedding = text_encoder(prompt)
    text_embeddings = text_embedding.repeat(num_samples, 1)

    shape = (num_samples, 1, 64, 64)
    samples = scheduler.sample_text(model, shape, text_embeddings, 'cuda', guidance_scale)

Command Line Usage

bash

# Generate samples
python generate.py --checkpoint text_diffusion_final_epoch_100.pt \
                   --prompt "a drawing of a fire truck" \
                   --num-samples 4 \
                   --guidance-scale 5.0

# Visualize denoising process
python visualize_generation.py --checkpoint text_diffusion_final_epoch_100.pt \
                               --prompt "a drawing of a cat" \
                               --num-steps 10

Example Prompts

Try these prompts for best results:

"a drawing of a cat"
"a drawing of a fire truck"
"a drawing of an airplane"
"a drawing of a house"
"a drawing of a tree"

Note: The model is trained on a limited set of QuickDraw classes, so it works best with simple object descriptions in the format "a drawing of a [object]".

Classifier-Free Guidance

The model supports classifier-free guidance to improve text-image alignment:

guidance_scale = 1.0: No guidance (pure conditional generation)
guidance_scale = 3.0-7.0: Recommended range (default: 5.0)
Higher values: Stronger adherence to text prompt (may reduce diversity)

Model Architecture

U-Net Structure

Input: (batch, 1, 64, 64)
├── Down Block 1: 1 → 256 channels
├── Down Block 2: 256 → 512 channels
├── Down Block 3: 512 → 512 channels
├── Middle Block: 512 channels
├── Up Block 3: 1024 → 512 channels (with skip connections)
├── Up Block 2: 768 → 256 channels (with skip connections)
└── Up Block 1: 512 → 1 channel (with skip connections)
Output: (batch, 1, 64, 64) - predicted noise

Text Conditioning

Text prompts encoded via CLIP ViT-B/32
512-dimensional text embeddings
Injected into U-Net via cross-attention
Classifier-free guidance with 15% dropout during training

Training Details

Framework: PyTorch 2.0+
Hardware: 4x NVIDIA GPUs
Training Time: ~100 epochs
Dataset: Google QuickDraw sketches (5 classes)
Noise Schedule: Linear (β from 0.0001 to 0.02)

Limitations

Limited to 64x64 resolution
Grayscale output only
Best performance on simple objects from QuickDraw classes
May not generalize well to complex or out-of-distribution prompts

Citation

If you use this model, please cite:

bibtex

@misc{quickdraw-text-diffusion,
  title={Text-Conditional QuickDraw Diffusion Model},
  author={Your Name},
  year={2024},
  howpublished={\url{https://huggingface.co/YOUR_USERNAME/quickdraw-text-diffusion}}
}

License

MIT License

Acknowledgments

Google QuickDraw dataset
OpenAI CLIP
DDPM paper: "Denoising Diffusion Probabilistic Models" (Ho et al., 2020)
Classifier-free guidance: "Classifier-Free Diffusion Guidance" (Ho & Salimans, 2022)

jamesaasher/quickdraw-text-diffusion