Minimal DDPM/DiT-based generation of MNIST digits. Heavily commented, self-contained script with all non-essential complexities stripped.

Added an extra coefficient into denoising function to adjust the strength of extra noise term, enabling to tradeoff the variety of generated samples with their quality.

High variety / low quality:

No variety / high quality:

Decent variety / decent quality:

References:

• https://arxiv.org/abs/2006.11239
• https://arxiv.org/abs/2212.09748
• https://www.deeplearning.ai/short-courses/how-diffusion-models-work/
• https://github.com/facebookresearch/DiT
• https://github.com/lucidrains/linformer