🎨 Rust Drawing Generator — Burn (Rust)

An interactive drawing generator utilizing a Denoising Diffusion Probabilistic Model (DDPM/DDIM) built with the Burn deep learning framework in Rust. Generate handwritten digits, letters, or doodles right in your browser!

🚀 Try the Live WebAssembly Demo!

---

📑 Table of Contents

• Features
• Model Architecture
• Project Structure
• Getting Started
  • Train the Model
  • Run Tests
  • CLI Generation
  • Interactive Web Server (Axum backend)
  • Client-Side WebAssembly App (WASM)
• Quick, Draw! Generation Details
• References
• License

---

✨ Features

• Conditional U-Net Architecture — Sinusoidal time embedding module, class conditioning embedding module, self-attention modules, skip connections, and residual blocks.
• DDPM & Flow Matching (Rectified Flow) Support — Dual support for both standard noise-prediction (DDPM/DDIM) and straight-line velocity-prediction (Flow Matching) generative paradigms.
• DDIM/Euler & Heun Scheduler Denoising — Accelerated reverse sampling supporting both DDIM/Euler (1st-Order) and Heun (2nd-Order) samplers, with customizable Polynomial Spacing Schedules (exponents 0.25–7.0) to achieve superior drawing quality in as few as 5–10 steps.
• Classifier-Free Guidance (CFG) with Dynamic Linear Decay — Extrapolates conditional predictions from an unconditional baseline to enhance semantic shape accuracy, utilizing a dynamic decay schedule to eliminate late-stage drawing artifacts.
• Interactive Educational Tooltips — Sleek inline SVG information toggles that explain the parameter spaces on click or hover to guide user selections.
• Cosine Annealing Learning Rate Scheduler — Smoothly decays the learning rate to zero during training to optimize final output convergence.
• Progressive Denoising Animation — Visual progressive rendering showing the drawing emerge frame-by-frame from pure Gaussian noise.
• Dual Inference Modes — Server-side Axum streaming via Server-Sent Events (SSE) and client-side WebAssembly local browser execution.
• Fully in Rust — Training, scheduler, inference engine, and web frontend in a unified workspace.

---

🏗️ Model Architecture

The generator relies on a lightweight Denoising Diffusion Probabilistic Model (DDPM/DDIM) with a conditional U-Net structure:

Inputs: Latent State x_t [B×1×28×28], Timestep t [B], Class ID c [B]
  → Time Embedding: Sinusoidal Positional Encoding (48) + MLP (48→192)
  → Class Embedding: Embedding Lookup (48) + Linear Projection (48→192)
  → Merged Embedding: Addition of Time & Class Embeddings [B×192]
  → U-Net Encoder:
      → Stem: Conv2d(1→48 channels) [B×48×28×28]
      → Down Block 1: UNetBlock + Time/Class Injection [B×48×28×28]
      → Downsample 1: Conv2d(48→96 channels, Stride 2) [B×96×14×14]
      → Down Block 2: UNetBlock + Time/Class Injection [B×96×14×14]
      → Downsample 2: Conv2d(96→192 channels, Stride 2) [B×192×7×7]
  → Bottleneck Middle Block: UNetBlock + Time/Class Injection + Self-Attention [B×192×7×7]
  → U-Net Decoder:
      → Upsample 1: ConvTranspose2d(192→96 channels, Stride 2) [B×96×14×14]
      → Up Block 1: Concatenate(Upsample 1, Skip 2) → UNetBlock(192→96 channels) + Self-Attention [B×96×14×14]
      → Upsample 2: ConvTranspose2d(96→48 channels, Stride 2) [B×48×28×28]
      → Up Block 2: Concatenate(Upsample 2, Skip 1) → UNetBlock(96→48 channels) [B×48×28×28]
      → Output Layer: Conv2d(48→1 channels) [B×1×28×28]

---

📁 Project Structure

burn-drawing-generator/
├── model_shared/           # Shared library workspace crate
│   ├── src/lib.rs          # Model architecture & DDIMScheduler definition
│   ├── src/unet.rs         # Conditional U-Net blocks and modules
│   └── src/scheduler.rs    # DDIM forward process and reverse sampling math
├── web/                    # Rust WASM crate (wasm-pack entry point)
│   └── src/lib.rs          # Stateful GeneratorWasm wrapper exposing .step()
├── src/                    # Training, CLI generation, & serving (Burn backend)
│   ├── main.rs             # CLI router & Axum API serve endpoint
│   ├── model.rs            # Re-exports shared Model wrapper
│   ├── training.rs         # Autodiff training loop and MSE loss wrapper
│   ├── inference.rs        # Iterative DDIM/Heun progressive sampling & ASCII art
│   ├── data.rs             # Noise collator & normalized dataset batcher
│   ├── emnist.rs           # EMNIST classes mapping helper
│   ├── quickdraw.rs        # QuickDraw dataset loading and class mappings
│   └── bin/                # Executable runner scripts
│       ├── build_web.rs    # Script to build, optimize, and bundle WASM
│       ├── convert.rs      # Utility to convert trained weights to binary format
│       └── publish_weights.rs # Utility to bundle and publish weights to release targets
├── docs/                   # Static web frontend (served by GitHub Pages)
│   ├── index.html          # Drawing generator UI with Developer Console
│   └── pkg/                # Compiled WASM output (gitignored, built by CI)
└── assets/                 # README showcase demo assets

---

🚀 Getting Started

Train the Model

Warning

Migration Warning: If you are migrating from an older checkpoint model (without Classifier-Free Guidance support), you must delete the existing checkpoint directory (e.g., ./target/mnist-model/) before starting training. Otherwise, the model loading step will crash due to class embedding size shape mismatch (num_classes + 1 rows).

By default, training runs on the CPU (NdArray backend):

cargo run --release

To train on your GPU (Wgpu backend):

cargo run --release -- --gpu

Training Configuration Flags:

• --epochs <num>: Customize the number of epochs to train (defaults to 5).
• --lr <num>: Customize the peak learning rate (defaults to 2e-4, decayed to 0 via Cosine Annealing).
• --prediction-type <noise|velocity>: Choose the target type (noise for DDPM/DDIM, velocity for Flow Matching).

(Note: These flags can be combined, for example: cargo run --release -- --epochs 40 --lr 0.0003 --prediction-type velocity --gpu)

---

💡 Hyperparameter Recommendations & Training Guides

Objective	Paradigm	Recommended Command
Quick Run (~10 Mins)	DDPM (Noise)	cargo run --release -- --epochs 6 --lr 2.5e-4 --prediction-type noise --gpu
Quick Run (~10 Mins)	Flow Matching	cargo run --release -- --epochs 6 --lr 8e-4 --prediction-type velocity --gpu
Fully Converged	DDPM (Noise)	cargo run --release -- --epochs 25 --lr 2e-4 --prediction-type noise --gpu
Fully Converged	Flow Matching	cargo run --release -- --epochs 30 --lr 5e-4 --prediction-type velocity --gpu

---

To train on the EMNIST Letters dataset (26 classes):

cargo run --release -- --dataset emnist --gpu

To train on the Google Quick, Draw! dataset (25 classes):

cargo run --release -- --dataset quickdraw --gpu

Dataset Cache: Dataset files are downloaded once and cached at target/emnist_dataset/ and target/quickdraw_dataset/.

---

Run Tests

Run the mathematical tests verifying the forward scheduling, time embeddings, and U-Net blocks:

cargo test

---

CLI Generation

Generate drawings from random Gaussian noise and watch the progressive ASCII rendering directly in your terminal:

# Generate MNIST digits
cargo run --release -- --predict

# Generate EMNIST Letters
cargo run --release -- --predict --dataset emnist

# Generate Quick, Draw! doodles
cargo run --release -- --predict --dataset quickdraw

📝 Example Terminal Output (MNIST digit 4 progress)

Loading model for generation (dataset: mnist)...
Generating drawing for class: '4' (class ID: 4) using 20 DDIM steps...

Generated Output:
                            
                            
               ..           
               ##.          
              .###          
              ####          
             #####          
            .#####          
            #####           
           .####.           
           #####            
          ######            
         #######            
         ####.##            
        ##### ##            
       .#### .##            
       ####  .##            
      ######.###.           
      ##########.           
      ##########            
          .  ##             
             ##.            
             ##.            
             ##.            
                            
                            

Generation complete!

---

Interactive Web Server (Axum backend)

Start the browser-based generator UI backed by the Rust Axum server streaming denoising steps via Server-Sent Events (SSE):

• MNIST Digits (default):

  cargo run --release -- --serve

  Open http://127.0.0.1:3000 to generate digits (0–9).

• EMNIST Letters:

  cargo run --release -- --serve --dataset emnist

  Open http://127.0.0.1:3000 to generate letters (A–Z).

• Quick, Draw! Doodles:

  cargo run --release -- --serve --dataset quickdraw

  Open http://127.0.0.1:3000 to generate doodles (25 classes).

GET /api/generate SSE Stream

The Axum server exposes a GET endpoint /api/generate that streams the progressive denoising steps using Server-Sent Events (SSE):

• Query Parameters:
  • class_id (integer, required): The target class ID to generate.
  • steps (integer, optional, default: 16): Number of denoising steps (clamped between 1 and 128).
  • schedule (string, optional, default: "linear"): Denoising schedule type or power exponent (e.g. "linear", "quadratic", or a custom float string like "3.0", "7.0").
  • sampler (string, optional, default: "ddim"): Denoising sampler algorithm: "ddim" or "heun".

---

Client-Side WebAssembly App (WASM)

The models compile to WebAssembly for fully client-side generation. Model weights (*-model.bin) are downloaded on-demand in the browser and cached locally.

1. Build the WASM bundle locally

Make sure you have trained the models first, then:

1. Install wasm-pack:

   cargo install wasm-pack

2. Build the WebAssembly module:

   cargo run --bin build_web

3. Install a local static file server:

   cargo install basic-http-server

4. Serve locally:

   basic-http-server docs

   Navigate to http://localhost:4000.

2. Automatic Deployments & Release Management

The CI workflow (.github/workflows/deploy.yml) automatically builds and deploys to GitHub Pages on every push to master.

To update the model weights used by the CI runner:

1. Upload your local weights to a GitHub Release:

   cargo run --bin publish_weights
   # or with custom tag
   cargo run --bin publish_weights -- v2.0.0

2. Commit and push the updated version files:

   git add docs/weights-version.txt
   git commit -m "Update model weights version"
   git push origin master

---

💡 Recommended Parameters Setup

To get the highest quality digit and doodle drawings at the best inference speeds, we recommend using the following parameters:

• Classifier-Free Guidance (CFG): On (3.0) (Amplifies class details during early stages)
• Denoising Schedule (Power): Early Focus (0.5) (Spends 70% of the step budget on global shape layouts)
• Sampler Algorithm: Euler / DDIM (1st-Order) (Fastest inference mode, requiring only 2 forward passes per step)
• Denoising Steps: 16 or 32

🧠 Dynamic Classifier-Free Guidance Decay

This workspace implements a Linear CFG Decay Schedule in both the backend server (src/inference.rs) and local browser WebAssembly crate (web/src/lib.rs): $$\text{Scale}{\text{step}} = 1.0 + (\text{Scale}{\text{initial}} - 1.0) \times (1.0 - \text{progress})$$

• At early steps (noise to outline): CFG scale is at its maximum (e.g. 3.0), strongly forcing the model to establish correct digit and category structures.
• At late steps (smoothing & sharpening): CFG decays to 1.0 (conditional-only), completely eliminating high-frequency stroke distortions and artifacts commonly caused by constant guidance.

---

🎨 Quick, Draw! Generation Details

This project supports doodle generation using a subset of the Quick, Draw! Dataset.

Selected Categories (25 classes)

We train on a curated subset of 25 doodle classes:

• Nature / Weather: sun, moon, star, tree, flower
• Animals: cat, dog, fish, butterfly
• Common Objects: cup, key, umbrella, hat, clock, envelope, toothbrush
• Structures / Vehicles: house, car
• Shapes: circle, triangle, square, smiley face
• Clothing: pants, t-shirt
• Food: apple

---

📚 References

• Burn — Deep Learning Framework for Rust
• The Burn Book
• wasm-pack — Rust WebAssembly Packager
• Rust Drawing Classifier Web (Original Fork Base)
• Denoising Diffusion Probabilistic Models (DDPM)
• Denoising Diffusion Implicit Models (DDIM)
• Classifier-Free Diffusion Guidance (Ho & Salimans, 2022)
• Flow Matching for Generative Modeling (Lipman et al., 2022)
• EMNIST Dataset (NIST Special Database 19)
• Google Quick, Draw! Dataset

---

📄 License

This project is licensed under the MIT License.