This repository serves as an educational resource, providing implementations of various generative modeling algorithms for unconditional image generation on the MNIST dataset. Each algorithm is implemented in a separate trainer class within the trainers
module, allowing for easy comparison and study of different approaches.
- networks/: Contains model architectures used by the trainers.
- trainers/: Implements various generative modeling algorithms for MNIST generation.
- main.py: Script to run training and sampling for selected models.
- README.md: This file, providing an overview and citations.
- Trainer:
trainers/ddpm.py
- Citation: Ho, J., Jain, A., & Abbeel, P. (2020). Denoising Diffusion Probabilistic Models. arXiv preprint arXiv:2006.11239. arXiv
- Trainer:
trainers/edm.py
- Citation: Karras, T., Aittala, M., Aila, T., & Laine, S. (2022). Elucidating the Design Space of Diffusion-Based Generative Models. arXiv preprint arXiv:2206.00364. arXiv
- Trainer:
trainers/rf.py
- Citation: Liu, X., Gong, C., & Liu, Q. (2023). Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow. arXiv preprint arXiv:2209.03003. arXiv
- Trainer:
trainers/vrf.py
- Citation: Guo, P., & Schwing, A.G. (2025). Variational Rectified Flow. arXiv preprint arXiv:2502.09616. arXiv
- Trainer:
trainers/pfgmpp.py
- Citation: Xu, Y., Liu, Z., Tian, Y., Tong, S., Tegmark, M., & Jaakola, T. (2023). PFGM++: Unlocking the Potential of Physics-Inspired Generative Models. arXiv preprint arXiv:2206.06910. arXiv
- Trainer:
trainers/cm.py
- Citation: Song, Y., Dhariwal, P., Chen, M., & Sutskever, I. (2023). Consistency Models. arXiv preprint arXiv:2303.01469. arXiv
- Trainer:
trainers/cd.py
- Citation: Bansal, A., Borgnia, E., Chu, H.M., Li, J.S., Kazemi, H., Huang, F., Goldblum, M., Geiping, J., & Goldstein, T. (2022). Cold Diffusion: Inverting Arbitrary Image Transforms Without Noise. arXiv preprint arXiv:2208.09392. arXiv
- Trainer:
trainers/vdm.py
- Citation: Kingma, D. P., Salimans, T., Poole, B., & Ho, J. (2021). Variational Diffusion Models. arXiv preprint arXiv:2107.00630. arXiv
To run the training and sampling for a specific model:
python main.py --model <model_name> [additional arguments]
Replace <model_name>
with one of the implemented algorithms (e.g., ddpm
, edm
, rf
, etc.). Use --help
to see all available options.
- Variational Encoder is based on black-forest-labs/flux
- DiT is from facebookresearch/DiT
- EDM is based on lucidrains/denoising-diffusion-pytorch
This repository was built for personal use and for teaching a section on generative modeling. A few bugs crept in while refactoring so if you happen to find any, or see any errors, please let me know!
This README was generated with assistance from Claude-3.5-Sonnet