Show HN: Modular Diffusion – A modular Python library for diffusion models

40 points

3 years ago

Hello everyone! I've been working on this project for a few months as part of my thesis in Machine Learning. It's meant to be a library that provides an easy-to-use but flexible API to design and train Diffusion Models. I decided to make it because I wanted to quickly prototype a Diffusion Model but there were no good tools to do it with. I think it really can help people prototype their own Diffusion Models a lot faster and only in a few lines of code.

The base idea is to have a Model class that takes different modules corresponding to the different aspects of the Diffusion Model process (noise schedule, noise type, denoising network, loss function, guidance, etc.) and allow the user to mix and match different modules to achieve different results. The library ships with a bunch of prebuilt modules and the plan is to add many more. I also made it super easy to implement your own modules, you just need to extend from one of the base classes available.

Contrary to HuggingFace Diffusers, this library is focused on designing and training your own Diffusion Models rather than finetuning pretrained ones (although this is possible).

I would really appreciate your feedback.

2 comments