You can deploy this model and start earning money today!

DD-65/
diffusionnumbers

Name: diffusionnumbers
Rating: 2.5 (25 reviews)

diffusers

safetensors

ddpm

conditional-image-generation

Views

No views yet

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:

python

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.