https://github.com/cabralpinto/modular-diffusion
Python library for designing and training your own Diffusion Models with PyTorch
https://github.com/cabralpinto/modular-diffusion
audio-generation deep-learning diffusion-models image-generation machine-learning modular-design python pytorch text-generation transformer u-net
Last synced: 5 months ago
JSON representation
Python library for designing and training your own Diffusion Models with PyTorch
- Host: GitHub
- URL: https://github.com/cabralpinto/modular-diffusion
- Owner: cabralpinto
- License: mit
- Created: 2023-06-22T10:22:19.000Z (over 2 years ago)
- Default Branch: main
- Last Pushed: 2025-06-17T11:53:06.000Z (7 months ago)
- Last Synced: 2025-08-22T15:08:05.448Z (5 months ago)
- Topics: audio-generation, deep-learning, diffusion-models, image-generation, machine-learning, modular-design, python, pytorch, text-generation, transformer, u-net
- Language: Python
- Homepage: https://cabralpinto.github.io/modular-diffusion/
- Size: 31.1 MB
- Stars: 287
- Watchers: 8
- Forks: 14
- Open Issues: 12
-
Metadata Files:
- Readme: README.md
- Contributing: CONTRIBUTING.md
- License: LICENSE
- Code of conduct: CODE_OF_CONDUCT.md
- Citation: CITATION.cff
Awesome Lists containing this project
README
# Modular Diffusion
[](https://badge.fury.io/py/modular-diffusion)
[](https://cabralpinto.github.io/modular-diffusion/)
[](https://lbesson.mit-license.org/)
[](https://discord.gg/mYJWQATfTV)
> ⚠️ **This project is currently unmaintained.**
> I'm no longer able to actively maintain this repository due to other commitments. If you’re interested in taking over as a maintainer and helping the project grow, please open an issue or reach out with a brief overview of your background and interest.
Modular Diffusion provides an easy-to-use modular API to design and train custom Diffusion Models with PyTorch. Whether you're an enthusiast exploring Diffusion Models or a hardcore ML researcher, **this framework is for you**.
## Features
- ⚙️ **Highly Modular Design**: Effortlessly swap different components of the diffusion process, including noise type, schedule type, denoising network, and loss function.
- 📚 **Growing Library of Pre-built Modules**: Get started right away with our comprehensive selection of pre-built modules.
- 🔨 **Custom Module Creation Made Easy**: Craft your own original modules by inheriting from a base class and implementing the required methods.
- 🤝 **Integration with PyTorch**: Built on top of PyTorch, Modular Diffusion enables you to develop custom modules using a familiar syntax.
- 🌈 **Broad Range of Applications**: From generating high-quality images to implementing non-autoregressive text synthesis pipelines, the possiblities are endless.
## Installation
Modular Diffusion officially supports Python 3.10+ and is available on PyPI:
```bash
pip install modular-diffusion
```
You also need to install the correct [PyTorch distribution](https://pytorch.org/get-started/locally/) for your system.
> **Note**: Although Modular Diffusion works with later Python versions, we currently recommend using Python 3.10. This is because `torch.compile`, which significantly improves the speed of the models, is not currently available for versions above Python 3.10.
## Usage
With Modular Diffusion, you can build and train a custom Diffusion Model in just a few lines. First, load and normalize your dataset. We are using the dog pictures from [AFHQ](https://paperswithcode.com/dataset/afhq).
```python
x, _ = zip(*ImageFolder("afhq", ToTensor()))
x = resize(x, [h, w], antialias=False)
x = torch.stack(x) * 2 - 1
```
Next, build your custom model using either Modular Diffusion's prebuilt modules or [your custom modules](https://cabralpinto.github.io/modular-diffusion/guides/custom-modules/).
```python
model = diffusion.Model(
data=Identity(x, batch=128, shuffle=True),
schedule=Cosine(steps=1000),
noise=Gaussian(parameter="epsilon", variance="fixed"),
net=UNet(channels=(1, 64, 128, 256)),
loss=Simple(parameter="epsilon"),
)
```
Now, train and sample from the model.
```python
losses = [*model.train(epochs=400)]
z = model.sample(batch=10)
z = z[torch.linspace(0, z.shape[0] - 1, 10).long()]
z = rearrange(z, "t b c h w -> c (b h) (t w)")
save_image((z + 1) / 2, "output.png")
```
Finally, marvel at the results.
Check out the [Getting Started Guide](https://cabralpinto.github.io/modular-diffusion/guides/getting-started/) to learn more and find more examples [here](https://github.com/cabralpinto/modular-diffusion/tree/main/examples).
## Contributing
We appreciate your support and welcome your contributions! Please feel free to submit pull requests if you found a bug or typo you want to fix. If you want to contribute a new prebuilt module or feature, please start by opening an issue and discussing it with us. If you don't know where to begin, take a look at the [open issues](https://github.com/cabralpinto/modular-diffusion/issues). Please read our [Contributing Guide](https://github.com/cabralpinto/modular-diffusion/blob/main/CONTRIBUTING.md) for more details.
## License
This project is licensed under the [MIT License](LICENSE).