Update README.md

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.
+
+<img width="360" alt="Modular Diffusion teaser" src="https://github.com/cabralpinto/modular-diffusion/assets/47889626/2756f798-8037-460e-b827-255812f203b6">
 
 ## 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}},
-}
-```