diff --git a/README.md b/README.md index d2863b7..0ffbbc4 100644 --- a/README.md +++ b/README.md @@ -1,85 +1,69 @@ # Modular Diffusion -Modular Diffusion provides an easy-to-use modular API to design and train custom Diffusion Models with PyTorch. +[![PyPI version](https://badge.fury.io/py/modular-diffusion.svg)](https://badge.fury.io/py/modular-diffusion) +[![Documentation](https://img.shields.io/badge/docs-stable-blue.svg)](https://cabralpinto.github.io/modular-diffusion/) +[![MIT license](https://img.shields.io/badge/license-MIT-blue.svg)](https://lbesson.mit-license.org/) -> **Warning**: The code for Modular Diffusion is not yet available on PyPI, as it is still undergoing significant development. A stable release will be published there once the code base, unit tests, and documentation reach a satisfactory state of completion. Stay tuned for more updates, and thank you for your patience! +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**: Tailor and build your own diffusion models by leveraging our flexible `Model` class. This class allows you to choose and implement various components of the diffusion process such as noise type, schedule type, denoising network, and loss function. + +- ⚙️ **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. -- 🤝 **Seamless Integration with PyTorch**: This project harmoniously integrates with PyTorch, allowing users to leverage the plethora of features and community support the framework has to offer. -- 🌈 **Broad Range of Applications**: From generating high-quality images to implementing non-autoregressive audio and text synthesis pipelines, the potential uses of Modular Diffusion are vast. - -## Usage - -With Modular Diffusion, you can build and train a custom Diffusion Model in just a few lines: - -1. Load and normalize your dataset. We are using [MNIST](http://yann.lecun.com/exdb/mnist/). - - ```python - x, y = zip(*MNIST(str(input), transform=ToTensor(), download=True)) - x, y = torch.stack(x) * 2 - 1, torch.tensor(y) + 1 - ``` +- 🤝 **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. -2. Build your custom model using Modular Diffusion's prebuilt modules ([or create your own!](https://cabralpinto.github.io/modular-diffusion/guides/custom-modules/)). +## Installation - ```python - model = diffusion.Model( - data=Identity(x, y, batch=128, shuffle=True), - schedule=Cosine(steps=1000), - noise=Gaussian(parameter="epsilon", variance="fixed"), - net=UNet(channels=(1, 64, 128, 256), labels=10), - guidance=ClassifierFree(dropout=0.1, strength=2), - loss=Simple(parameter="epsilon"), - ) - ``` +Modular Diffusion officially supports Python 3.10+ and is available on PyPI: -3. Train and sample from the model. +```bash +pip install modular-diffusion +``` - ```python - losses = [*model.train(epochs=10)] - z = model.sample(torch.arange(1, 11)) - 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") - ``` +> **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. - 4. Marvel at the results. +## Usage - ![](docs/public/images/guides/getting-started/conditional.png) +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(str(input), ToTensor())) +x = resize(x, [h, w], antialias=False) +x = torch.stack(x) * 2 - 1 +``` -Check out the [examples](examples) folder in our repository for more examples and refer to our [documentation](https://cabralpinto.github.io/modular-diffusion/guides/getting-started/) for more information. +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/). -## Installation +```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"), +) +``` -Modular Diffusion officially supports Python 3.10+ and is available on PyPI: +Now, train and sample from the model. -```bash -pip install modular-diffusion +```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") ``` -> **Note**: Although Modular Diffusion works with later Python versions, we highly 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. +Finally, marvel at the results. + +Modular Diffusion teaser ## Contributing -Contributions to this project are very much welcome! Feel free to raise issues or submit pull requests here on GitHub. +We appreciate your support and welcome your contributions! Please fell 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. ## License This project is licensed under the [MIT License](LICENSE). - -## Citation - -If you use this library in your work, please cite it using the following BibTeX entry: - -```bibtex -@misc{ModularDiffusion, - author = {João Cabral Pinto}, - title = {Modular Diffusion}, - year = {2023}, - publisher = {GitHub}, - howpublished = {\url{https://github.com/cabralpinto/modular-diffusion}}, -} -```