diffusionnumbers / README.md
DD-65's picture
Updated model card with examples
e8965ee verified
|
Raw
History Blame Contribute Delete
2.83 kB
---
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](https://huggingface.co/datasets/ylecun/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
| | | |
|---|---|---|
| ![](samples/generated-2.png) | ![](samples/generated-4.png) | ![](samples/generated-8.png) |
## 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](https://github.com/DD-65/diffusionnumbers-code) 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](https://huggingface.co/datasets/ylecun/mnist),
which is also distributed under the MIT license.