https://github.com/probabilists/azula
Diffusion models in PyTorch
https://github.com/probabilists/azula
deep-learning diffusion-models generative-model torch
Last synced: about 2 months ago
JSON representation
Diffusion models in PyTorch
- Host: GitHub
- URL: https://github.com/probabilists/azula
- Owner: probabilists
- License: mit
- Created: 2024-06-21T09:14:37.000Z (11 months ago)
- Default Branch: master
- Last Pushed: 2025-03-26T10:31:38.000Z (about 2 months ago)
- Last Synced: 2025-03-29T06:01:33.386Z (about 2 months ago)
- Topics: deep-learning, diffusion-models, generative-model, torch
- Language: Python
- Homepage: https://azula.readthedocs.io
- Size: 3.63 MB
- Stars: 99
- Watchers: 6
- Forks: 4
- Open Issues: 3
-
Metadata Files:
- Readme: README.md
- Contributing: CONTRIBUTING.md
- License: LICENSE
Awesome Lists containing this project
README
# Azula - Diffusion models in PyTorch
Azula is a Python package that implements diffusion models in [PyTorch](https://pytorch.org). Its goal is to unify the different formalisms and notations of the generative diffusion models literature into a single, convenient and hackable interface.
> In the [Avatar](https://wikipedia.org/wiki/Avatar:_The_Last_Airbender) cartoon, [Azula](https://wikipedia.org/wiki/Azula) is a powerful fire and lightning bender ⚡️
## Installation
The `azula` package is available on [PyPI](https://pypi.org/project/azula), which means it is installable via `pip`.
```
pip install azula
```
Alternatively, if you need the latest features, you can install it from the repository.
```
pip install git+https://github.com/probabilists/azula
```
## Getting started
In Azula's formalism, a diffusion model is the composition of three elements: a noise schedule, a denoiser and a sampler.
* A noise schedule is a mapping from a time $t \in [0, 1]$ to the signal scale $\alpha_t$ and the noise scale $\sigma_t$ in a perturbation kernel $p(X_t \mid X) = \mathcal{N}(X_t \mid \alpha_t X_t, \sigma_t^2 I)$ from a "clean" random variable $X \sim p(X)$ to a "noisy" random variable $X_t$.
* A denoiser is a neural network trained to predict $X$ given $X_t$.
* A sampler defines a series of transition kernels $q(X_s \mid X_t)$ from $t$ to $s < t$ based on a noise schedule and a denoiser. Simulating these transitions from $t = 1$ to $0$ samples approximately from $p(X)$.
This formalism is closely followed by Azula's API.
```python
from azula.denoise import PreconditionedDenoiser
from azula.noise import VPSchedule
from azula.sample import DDPMSampler
# Choose the variance preserving (VP) noise schedule
schedule = VPSchedule()
# Initialize a denoiser
denoiser = PreconditionedDenoiser(
backbone=CustomNN(in_features=5, out_features=5),
schedule=schedule,
)
# Train to predict x given x_t
optimizer = torch.optim.Adam(denoiser.parameters(), lr=1e-3)
for x in train_loader:
t = torch.rand((batch_size,))
loss = denoiser.loss(x, t)
loss.backward()
optimizer.step()
optimizer.zero_grad()
# Generate 64 points in 1000 steps
sampler = DDPMSampler(denoiser.eval(), steps=1000)
x1 = sampler.init((64, 5))
x0 = sampler(x1)
```
Alternatively, Azula's plugin interface allows to load pre-trained models and use them with the same convenient interface.
```python
import sys
sys.path.append("path/to/guided-diffusion")
from azula.plugins import adm
from azula.sample import DDIMSampler
# Download weights from openai/guided-diffusion
denoiser = adm.load_model("imagenet_256x256")
# Generate a batch of 4 images
sampler = DDIMSampler(denoiser, steps=64).cuda()
x1 = sampler.init((4, 3, 256, 256))
x0 = sampler(x1)
images = torch.clip((x0 + 1) / 2, min=0, max=1)
```
For more information, check out the documentation and tutorials at [azula.readthedocs.io](https://azula.readthedocs.io).
## Contributing
If you have a question, an issue or would like to contribute, please read our [contributing guidelines](https://github.com/probabilists/azula/blob/master/CONTRIBUTING.md).