README.md

metadata

license: mit
library_name: diffusers
pipeline_tag: text-to-image
tags:
  - diffusers
  - ddpm
  - conditional-image-generation
  - mnist
  - diffusion-models-class

DiffusionNumbers DDPM

DiffusionNumbers is a DDPM trained from scratch on the 60k training images of MNIST on Apple Silicon. Training took approximately 5 hours. It generates 64×64 grayscale images of a requested handwritten digit from 0 to 9.

Examples

Usage

Install PyTorch, Diffusers, Safetensors, and Pillow, then run the model as follows. Class labels have to be passed to the U-Net during every denoising step:

import torch
from diffusers import DDPMPipeline
from PIL import Image

if torch.backends.mps.is_available():
    device = "mps"
elif torch.cuda.is_available():
    device = "cuda"
else:
    device = "cpu"

pipe = DDPMPipeline.from_pretrained("DD-65/diffusionnumbers").to(device)

digit = 7
labels = torch.tensor([digit], device=device)
images = torch.randn((1, 1, 64, 64), device=device)

pipe.scheduler.set_timesteps(1000)
for timestep in pipe.scheduler.timesteps:
    with torch.no_grad():
        noise_prediction = pipe.unet(
            images,
            timestep,
            class_labels=labels,
        ).sample

    images = pipe.scheduler.step(
        noise_prediction,
        timestep,
        images,
    ).prev_sample

image = (images[0, 0] / 2 + 0.5).clamp(0, 1)
image = image.mul(255).round().to(torch.uint8).cpu().numpy()
Image.fromarray(image).save("digit-7.png")

Change digit to any integer from 0 through 9. See also the Github repo for the complete training code and a more elaborate inference script.

Model details

• Architecture: UNet2DModel with a DDPMScheduler
• Parameters: 3.94 million
• Resolution: 64×64 grayscale (MNIST images have been subsampled)
• Conditioning: class labels for digits 0–9
• Training data: 60,000 MNIST training images
• Training: 40 epochs, batch size 128, AdamW with a 2e-4 learning rate
• Training objective: epsilon/noise prediction
• Scheduler training timesteps: 1,000

The original 28×28 images were resized to 64×64 and randomly scaled between 90% and 110% during training.

Limitations

The model is designed for experimental image generation. It only generates MNIST-like digits and is not a classifier or OCR model. Outputs can be malformed or ambiguous, and their style and diversity are limited by the training data.

License

The model weights are released under the MIT license.

The model was trained on MNIST, which is also distributed under the MIT license.