### Setup and Installation for ELDM Plugin

Source: https://github.com/probabilists/azula/blob/master/docs/api/azula.plugins.eldm.md

Instructions for cloning the required NVlabs/edm2 repository and installing necessary Python dependencies like diffusers and accelerate.

```console
git clone https://github.com/edm2
pip install diffusers accelerate
```

```python
import sys
sys.path.append("path/to/edm2")
from azula.plugins import eldm
```

--------------------------------

### Setup and Imports for Azula Diffusion

Source: https://github.com/probabilists/azula/blob/master/docs/tutorials/mnist.md

Initializes the environment by installing dependencies and importing essential PyTorch and Azula modules for diffusion modeling.

```python
import torch
import torch.nn as nn

from datasets import load_dataset
from PIL import Image
from torch.utils.data import DataLoader
from torchvision.transforms.functional import to_pil_image, to_tensor
from torchvision.utils import make_grid
from tqdm import tqdm

from azula.denoise import KarrasDenoiser
from azula.nn.unet import UNet
from azula.noise import VPSchedule
from azula.sample import DDIMSampler

device = "cuda"
```

--------------------------------

### Install Azula via pip

Source: https://github.com/probabilists/azula/blob/master/README.md

Commands to install the Azula package from PyPI or directly from the GitHub repository.

```bash
pip install azula
pip install git+https://github.com/probabilists/azula
```

--------------------------------

### Install Flux plugin dependencies

Source: https://github.com/probabilists/azula/blob/master/docs/api/azula.plugins.flux.md

Install the necessary Python packages required to run the Flux plugin functionality.

```console
pip install diffusers transformers accelerate protobuf sentencepiece
```

--------------------------------

### Install Dependencies for Latent Diffusion

Source: https://github.com/probabilists/azula/blob/master/docs/tutorials/latent.ipynb

Installs necessary Python libraries for using latent diffusion models, including diffusers, transformers, and accelerate.

```python
# !pip install diffusers transformers accelerate
```

--------------------------------

### Install Dependencies

Source: https://github.com/probabilists/azula/blob/master/docs/tutorials/latent.md

Installs the necessary Python packages for running latent diffusion models, including diffusers, transformers, and accelerate.

```ipython3
!pip install diffusers transformers accelerate
```

--------------------------------

### Install Dependencies for Azula

Source: https://github.com/probabilists/azula/blob/master/docs/tutorials/mnist.ipynb

Installs the 'datasets' library, which is a dependency for loading datasets in the Azula diffusion model tutorial. This is typically run in a notebook environment.

```python
# !pip install datasets
```

--------------------------------

### Install and Configure Pre-commit Hooks

Source: https://github.com/probabilists/azula/blob/master/CONTRIBUTING.md

Installs pre-commit hooks to automatically enforce code conventions before each commit. It uses a specified configuration file ('pre-commit.yaml') for hook definitions.

```bash
pre-commit install --config pre-commit.yaml
```

--------------------------------

### Install Stable Diffusion Plugin Dependencies

Source: https://github.com/probabilists/azula/blob/master/docs/api/azula.plugins.sd.md

Installs the necessary Python packages for the Stable Diffusion plugin, including diffusers, transformers, and accelerate.

```console
pip install diffusers transformers accelerate
```

--------------------------------

### Perform Pseudo-inverse Guided Diffusion (PGDM)

Source: https://github.com/probabilists/azula/blob/master/docs/tutorials/guidance.ipynb

Applies the PGDMSampler for guided sampling using pseudo-inverse approximations for the inverse problem.

```python
from azula.guidance import PGDMSampler

cond_sampler = PGDMSampler(denoiser, y=y, A=A, A_inv=A_inv, steps=64, eta=1.0)
x1 = cond_sampler.init((4, 3, 256, 256), device=device)
x0 = cond_sampler(x1)
to_pil_image(make_grid(postprocess(x0)))
```

--------------------------------

### Install Azula in Editable Mode with Development Dependencies

Source: https://github.com/probabilists/azula/blob/master/CONTRIBUTING.md

Installs the 'azula' package in editable mode, including development, documentation, and testing dependencies. This is recommended for contributing code to the project.

```bash
pip install -e .[dev,docs,test]
```

--------------------------------

### Stable Diffusion Plugin for Loading Models (Python)

Source: https://context7.com/probabilists/azula/llms.txt

Loads pre-trained Stable Diffusion models. This example demonstrates loading the 'sd_v1_5' model and preparing it for evaluation on CUDA.

```python
from azula.plugins import sd

denoiser = sd.load_model("sd_v1_5")
denoiser.to("cuda").eval()
```

--------------------------------

### Perform Moment Matching Posterior Sampling (MMPS)

Source: https://github.com/probabilists/azula/blob/master/docs/tutorials/guidance.ipynb

Uses the MMPSDenoiser to perform guided sampling via moment matching with diagonal covariance.

```python
from azula.guidance import MMPSDenoiser
from azula.linalg.covariance import DiagonalCovariance

cond_denoiser = MMPSDenoiser(denoiser, y=y, A=A, cov_y=DiagonalCovariance(sigma_y**2), iterations=3)
cond_sampler = DDIMSampler(cond_denoiser, steps=64, eta=1.0)
x1 = cond_sampler.init((4, 3, 256, 256), device=device)
x0 = cond_sampler(x1)
to_pil_image(make_grid(postprocess(x0)))
```

--------------------------------

### GET /azula/plugins/eldm/load_model

Source: https://github.com/probabilists/azula/blob/master/docs/api/azula.plugins.eldm.md

Utility function to load pre-trained ELDM/EDM2 models.

```APIDOC
## load_model(name)

### Description
Loads a pre-trained ELDM or EDM2 latent denoiser from the library.

### Parameters
- **name** (str) - Required - The identifier of the pre-trained model.

### Response
- **tuple**: Returns a tuple containing (Denoiser, AutoEncoder).
```

--------------------------------

### Initialize Azula Environment and Pre-trained Model

Source: https://github.com/probabilists/azula/blob/master/docs/tutorials/guidance.md

Imports necessary libraries, sets the device, and loads a pre-trained ADM diffusion model for image generation.

```ipython3
import io
import requests
import torch

from PIL import Image
from torchvision.transforms.functional import to_pil_image, to_tensor
from torchvision.utils import make_grid

from azula.guidance import (
    DiffPIRDenoiser,
    DPSSampler,
    JFPSDenoiser,
    MMPSDenoiser,
    PGDMSampler,
    TMPDenoiser,
)
from azula.linalg.covariance import DiagonalCovariance, KroneckerCovariance
from azula.plugins import adm
from azula.sample import DDIMSampler

device = "cuda"
_ = torch.manual_seed(42)

def preprocess(x):
    return 2 * x - 1

def postprocess(x):
    return torch.clip((x + 1) / 2, min=0, max=1)

denoiser = adm.load_model("imagenet_256x256").to(device)
denoiser = denoiser.requires_grad_(False)
```

--------------------------------

### Perform Conditional Sampling with Various Methods

Source: https://github.com/probabilists/azula/blob/master/docs/tutorials/guidance.md

Demonstrates how to use different samplers and denoisers (DPS, PGDM, DiffPIR, TMPD, MMPS, JFPS) to restore images based on the defined measurement operators.

```ipython3
# DPS
cond_sampler = DPSSampler(denoiser, y=y, A=A, steps=64)

# PGDM
cond_sampler = PGDMSampler(denoiser, y=y, A=A, A_inv=A_inv, steps=64, eta=1.0)

# DiffPIR
cond_denoiser = DiffPIRDenoiser(denoiser, y=y, A=A, var_y=sigma_y**2, iterations=1)
cond_sampler = DDIMSampler(cond_denoiser, steps=64, eta=1.0)

# TMPD
cond_denoiser = TMPDenoiser(denoiser, y=y, A=A, var_y=sigma_y**2)
cond_sampler = DDIMSampler(cond_denoiser, steps=64, eta=1.0)

# MMPS
cond_denoiser = MMPSDenoiser(denoiser, y=y, A=A, cov_y=DiagonalCovariance(sigma_y**2), iterations=3)
cond_sampler = DDIMSampler(cond_denoiser, steps=64, eta=1.0)

# JFPS
x0_uncond = sampler(sampler.init((4, 3, 256, 256), device=device))
cov_x = KroneckerCovariance.from_data(x0_uncond)
cond_denoiser = JFPSDenoiser(denoiser, y=y, A=A, cov_y=DiagonalCovariance(sigma_y**2), cov_x=cov_x, iterations=11)
cond_sampler = DDIMSampler(cond_denoiser, steps=64, eta=1.0)
```

--------------------------------

### Initialize and Use JITDenoiser

Source: https://github.com/probabilists/azula/blob/master/docs/api/azula.plugins.jit.md

Demonstrates the initialization of the JITDenoiser class and the structure of its forward pass for denoising operations.

```python
from azula.plugins.jit import JITDenoiser

# Initialize the denoiser
denoiser = JITDenoiser(backbone=my_backbone, schedule=my_schedule, num_classes=1000)

# Perform forward pass
# x_t: noisy tensor, t: time, label: class label
output = denoiser.forward(x_t, t, label=label)
```

--------------------------------

### Initialize Latent Diffusion Components (Python)

Source: https://github.com/probabilists/azula/blob/master/docs/tutorials/latent.ipynb

Imports core libraries and components for latent diffusion, including PyTorch, tqdm, accelerate, torchvision, and specific modules from azula for guidance, plugins, and sampling. It also sets up device and disables HTML progress bars.

```python
import torch
import tqdm.auto

tqdm.auto.tqdm = tqdm.asyncio.tqdm  # disable HTML progress bars

from accelerate import cpu_offload
from torchvision.transforms.functional import to_pil_image
from torchvision.utils import make_grid

from azula.guidance import CFGDenoiser
from azula.plugins import flux, sana, sd
from azula.sample import zABSampler  # Adams-Bashforth sampler

device = "cuda"
```

--------------------------------

### Build Documentation Locally with Sphinx

Source: https://github.com/probabilists/azula/blob/master/CONTRIBUTING.md

Builds the project's documentation locally using Sphinx. This command should be executed from the 'docs' directory and outputs the built documentation to the 'html' directory.

```bash
cd docs
sphinx-build . html
```

--------------------------------

### Initialize Diffusion Model and Sampler

Source: https://github.com/probabilists/azula/blob/master/docs/tutorials/guidance.ipynb

Sets up the environment, loads a pre-trained ADM model, and initializes a DDIM sampler for unconditional image generation.

```python
import torch
from azula.plugins import adm
from azula.sample import DDIMSampler
from torchvision.transforms.functional import to_pil_image, to_pil_image, make_grid

device = "cuda"
def postprocess(x):
    return torch.clip((x + 1) / 2, min=0, max=1)

denoiser = adm.load_model("imagenet_256x256").to(device).requires_grad_(False)
sampler = DDIMSampler(denoiser, steps=64)
x1 = sampler.init((4, 3, 256, 256), device=device)
x0_uncond = sampler(x1)
to_pil_image(make_grid(postprocess(x0_uncond)))
```

--------------------------------

### Implement and Train a Diffusion Model

Source: https://github.com/probabilists/azula/blob/master/docs/index.md

Demonstrates how to define a noise schedule, initialize a Karras denoiser with a custom backbone, and perform a training loop. It also shows how to use a DDPM sampler to generate data points.

```python
from azula.denoise import KarrasDenoiser
from azula.noise import VPSchedule
from azula.sample import DDPMSampler

schedule = VPSchedule()
denoiser = KarrasDenoiser(
   backbone=CustomNN(in_features=5, out_features=5),
   schedule=schedule,
)

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()

sampler = DDPMSampler(denoiser.eval(), steps=1000)
x1 = sampler.init((64, 5))
x0 = sampler(x1)
```

--------------------------------

### Train and sample with custom diffusion models

Source: https://github.com/probabilists/azula/blob/master/README.md

Demonstrates how to define a noise schedule, initialize a Karras denoiser with a custom backbone, perform training, and generate samples using a DDPM sampler.

```python
from azula.denoise import KarrasDenoiser
from azula.noise import VPSchedule
from azula.sample import DDPMSampler

schedule = VPSchedule()
denoiser = KarrasDenoiser(
    backbone=CustomNN(in_features=5, out_features=5),
    schedule=schedule,
)

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()

sampler = DDPMSampler(denoiser.eval(), steps=1000)
x1 = sampler.init((64, 5))
x0 = sampler(x1)
```

--------------------------------

### Perform Tweedie Moment Projected Diffusion (TMPD)

Source: https://github.com/probabilists/azula/blob/master/docs/tutorials/guidance.ipynb

Uses the TMPDenoiser to perform guided sampling based on Tweedie moment projection.

```python
from azula.guidance import TMPDenoiser

cond_denoiser = TMPDenoiser(denoiser, y=y, A=A, var_y=sigma_y**2)
cond_sampler = DDIMSampler(cond_denoiser, steps=64, eta=1.0)
x1 = cond_sampler.init((4, 3, 256, 256), device=device)
x0 = cond_sampler(x1)
to_pil_image(make_grid(postprocess(x0)))
```

--------------------------------

### Import JIT Plugin

Source: https://github.com/probabilists/azula/blob/master/docs/api/azula.plugins.jit.md

Basic import statement to access the JIT plugin functionality within the azula library.

```python
from azula.plugins import jit
```

--------------------------------

### Class: PGDMSampler

Source: https://github.com/probabilists/azula/blob/master/docs/api/azula.guidance.pgdm.md

Initializes the PGDM sampler for inverse problem solving using a denoiser and forward/pseudo-inverse operators.

```APIDOC
## CLASS PGDMSampler

### Description
Creates a PGDM sampler for Pseudo-inverse Guided Diffusion Models as described by Song et al. (2023).

### Constructor
`azula.guidance.pgdm.PGDMSampler(denoiser, y, A, A_inv, **kwargs)`

### Parameters
- **denoiser** (Denoiser) - Required - A denoiser q_phi(X | X_t).
- **y** (Tensor) - Required - An observation y ~ N(A(x), Sigma_y).
- **A** (Callable) - Required - The forward operator x -> A(x).
- **A_inv** (Callable) - Required - The pseudo-inverse operator y -> A_dagger(y), such that A(A_dagger(A(x))) = A(x).
- **kwargs** (dict) - Optional - Keyword arguments passed to azula.sample.DDIMSampler.

### Usage Example
```python
sampler = PGDMSampler(
    denoiser=my_denoiser,
    y=observation_tensor,
    A=forward_operator,
    A_inv=pseudo_inverse_operator
)
```
```

--------------------------------

### Perform Diffusion Posterior Sampling (DPS)

Source: https://github.com/probabilists/azula/blob/master/docs/tutorials/guidance.ipynb

Uses the DPSSampler to perform guided image restoration based on a measurement y and forward operator A.

```python
from azula.guidance import DPSSampler

cond_sampler = DPSSampler(denoiser, y=y, A=A, steps=64)
x1 = cond_sampler.init((4, 3, 256, 256), device=device)
x0 = cond_sampler(x1)
to_pil_image(make_grid(postprocess(x0)))
```

--------------------------------

### Get Module Device (Python)

Source: https://github.com/probabilists/azula/blob/master/docs/api/azula.nn.utils.md

Retrieves the execution device of a PyTorch module. It checks the module's parameters and buffers for the first device. Returns None if no device is found.

```python
import torch.nn as nn

def get_module_device(module: nn.Module) -> torch.device | None:
    """Returns the execution device of a module."""
    for param in module.parameters():
        return param.device
    for buffer in module.buffers():
        return buffer.device
    return None
```

--------------------------------

### POST /azula/sample/EABSampler

Source: https://github.com/probabilists/azula/blob/master/docs/api/azula.sample.md

Initializes an Exponential Adams-Bashforth (EAB) sampler, a multi-step generalization of the DPM-Solver++ algorithm.

```APIDOC
## POST /azula/sample/EABSampler

### Description
Creates an Exponential Adams-Bashforth (EAB) multi-step sampler based on DPM-Solver++.

### Parameters
#### Request Body
- **denoiser** (Denoiser) - Required - A denoiser q_phi(X | X_t).
- **order** (int) - Required - The order n of the multi-step method.
- **kwargs** (dict) - Optional - Keyword arguments passed to Sampler.

### Request Example
{
  "denoiser": "<DenoiserObject>",
  "order": 2
}
```

--------------------------------

### Initialize SimpleDenoiser for training

Source: https://context7.com/probabilists/azula/llms.txt

Creates a SimpleDenoiser with a custom backbone and VPSchedule. This denoiser directly predicts clean data and supports loss weight clipping via the max_weight parameter.

```python
import torch
import torch.nn as nn
from azula.denoise import SimpleDenoiser
from azula.noise import VPSchedule

class Backbone(nn.Module):
    def __init__(self, dim=64):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(dim, 128),
            nn.SiLU(),
            nn.Linear(128, dim),
        )

    def forward(self, x, t):
        return self.net(x)

denoiser = SimpleDenoiser(backbone=Backbone(dim=64), schedule=VPSchedule())

x = torch.randn(16, 64)
t = torch.rand(16)
loss = denoiser.loss(x, t, max_weight=1e4)
```

--------------------------------

### Get Module Data Type (Python)

Source: https://github.com/probabilists/azula/blob/master/docs/api/azula.nn.utils.md

Retrieves the data type of a PyTorch module. It checks the module's parameters and buffers for the first floating-point type. Returns None if no floating-point type is found.

```python
import torch.nn as nn

def get_module_dtype(module: nn.Module) -> torch.dtype | None:
    """Returns the data type of a module."""
    for param in module.parameters():
        return param.dtype
    for buffer in module.buffers():
        return buffer.dtype
    return None
```

--------------------------------

### Load Sana Model Components (Python)

Source: https://github.com/probabilists/azula/blob/master/docs/tutorials/latent.ipynb

Loads the denoiser, autoencoder, and text encoder for the Sana 1.6B model with 512 resolution. It moves the denoiser to the specified device and offloads the other components to CPU.

```python
denoiser, autoencoder, textencoder = sana.load_model("sana_1.6b_512")
denoiser, autoencoder, textencoder = (
    denoiser.to(device),
    cpu_offload(autoencoder, device),
    cpu_offload(textencoder, device),
)
```

--------------------------------

### EDM Plugin for Pre-trained Diffusion Models (Python)

Source: https://context7.com/probabilists/azula/llms.txt

Loads pre-trained Elucidating Diffusion Models (EDM) models. This example demonstrates loading an unconditional CIFAR-10 model and preparing it for use.

```python
from azula.plugins import edm

denoiser = edm.load_model("cifar10_uncond")
denoiser.to("cuda").eval()
```

--------------------------------

### Initialize Environment and Post-processing

Source: https://github.com/probabilists/azula/blob/master/docs/tutorials/latent.md

Sets up the torch environment, configures progress bars, and defines a utility function to normalize and clip latent outputs for image visualization.

```ipython3
import torch
import tqdm.auto

tqdm.auto.tqdm = tqdm.asyncio.tqdm

from accelerate import cpu_offload
from torchvision.transforms.functional import to_pil_image
from torchvision.utils import make_grid

from azula.guidance import CFGDenoiser
from azula.plugins import flux, sana, sd
from azula.sample import zABSampler

device = "cuda"
_ = torch.manual_seed(42)

def postprocess(x):
    return torch.clip((x + 1) / 2, min=0, max=1)
```

--------------------------------

### Function: load_model

Source: https://github.com/probabilists/azula/blob/master/docs/api/azula.plugins.jit.md

Utility function to load pre-trained JIT denoiser models from storage.

```APIDOC
## FUNCTION azula.plugins.jit.load_model

### Description
Loads a pre-trained JIT denoiser model based on the provided name.

### Parameters
- **name** (str) - Required - The identifier of the pre-trained model.
- **ema** (bool) - Optional - Whether to load Exponential Moving Average weights.
- **kwargs** (dict) - Optional - Additional arguments passed to torch.load.

### Response
- **Returns** (Denoiser) - An instance of the pre-trained denoiser.
```

--------------------------------

### Train a diffusion model on MNIST with Azula

Source: https://context7.com/probabilists/azula/llms.txt

A complete implementation of a conditional diffusion model using a UNet backbone, VPSchedule, and KarrasDenoiser. Demonstrates the training loop, optimizer setup, and sample generation using the DDIM sampler.

```python
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

from azula.denoise import KarrasDenoiser
from azula.nn.unet import UNet
from azula.noise import VPSchedule
from azula.sample import DDIMSampler

# Custom backbone wrapping UNet with embeddings
class DiffusionBackbone(nn.Module):
    def __init__(self, channels=1, emb_features=64, num_classes=10):
        super().__init__()
        self.unet = UNet(
            in_channels=channels,
            out_channels=channels,
            hid_channels=[32, 64, 128],
            hid_blocks=[2, 2, 2],
            mod_features=emb_features,
        )
        self.label_emb = nn.Embedding(num_classes, emb_features)
        self.time_mlp = nn.Sequential(
            nn.Linear(1, emb_features),
            nn.SiLU(),
            nn.Linear(emb_features, emb_features),
        )

    def forward(self, x_t, log_snr, label=None):
        emb = self.time_mlp(log_snr[..., None])
        if label is not None:
            emb = emb + self.label_emb(label)
        return self.unet(x_t, mod=emb)

# Setup
device = "cuda" if torch.cuda.is_available() else "cpu"
schedule = VPSchedule()
backbone = DiffusionBackbone()
denoiser = KarrasDenoiser(backbone=backbone, schedule=schedule).to(device)

# Data
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Lambda(lambda x: 2 * x - 1),  # Scale to [-1, 1]
])
dataset = datasets.MNIST(root="./data", train=True, download=True, transform=transform)
loader = DataLoader(dataset, batch_size=64, shuffle=True)

# Training
optimizer = torch.optim.AdamW(denoiser.parameters(), lr=3e-4)

for epoch in range(10):
    total_loss = 0
    for images, labels in loader:
        images, labels = images.to(device), labels.to(device)
        t = torch.rand(len(images), device=device)

        loss = denoiser.loss(images, t, label=labels)
        loss.backward()

        optimizer.step()
        optimizer.zero_grad()
        total_loss += loss.item()

    print(f"Epoch {epoch+1}, Loss: {total_loss / len(loader):.4f}")

# Generation
denoiser.eval()
sampler = DDIMSampler(denoiser, steps=50, silent=True)

# Generate digit "7"
with torch.no_grad():
    label = torch.tensor([7] * 8, device=device)
    x1 = sampler.init((8, 1, 28, 28), device=device)
    x0 = sampler(x1, label=label)

    # Post-process
    images = torch.clip((x0 + 1) / 2, 0, 1)
```

--------------------------------

### Load Flux Model Components (Python)

Source: https://github.com/probabilists/azula/blob/master/docs/tutorials/latent.ipynb

Loads the denoiser, autoencoder, and text encoder for the Flux model. It moves the denoiser to the specified device and offloads the other components to CPU.

```python
denoiser, autoencoder, textencoder = flux.load_model()
denoiser, autoencoder, textencoder = (
    denoiser.to(device),
    cpu_offload(autoencoder, device),
    cpu_offload(textencoder, device),
)
```

--------------------------------

### UNet for Diffusion Models with Modulation (Python)

Source: https://context7.com/probabilists/azula/llms.txt

Implements a flexible U-Net architecture with modulation support, suitable for diffusion models. It can handle conditional inputs and time/class embeddings. The example shows initialization and forward passes with and without conditioning.

```python
import torch
from azula.nn.unet import UNet

# Create U-Net for image generation
unet = UNet(
    in_channels=3,
    out_channels=3,
    cond_channels=0,
    hid_channels=[64, 128, 256],
    hid_blocks=[2, 3, 4],
    kernel_size=3,
    stride=2,
    spatial=2,
    mod_features=128,  # Modulation features for time/class embedding
)

# Forward pass
x = torch.randn(4, 3, 64, 64)  # Input image
mod = torch.randn(4, 128)  # Time/class embedding

output = unet(x, mod=mod)  # Shape: (4, 3, 64, 64)

# With condition tensor
cond = torch.randn(4, 16, 64, 64)
unet_cond = UNet(
    in_channels=3,
    out_channels=3,
    cond_channels=16,
    hid_channels=[64, 128, 256],
    hid_blocks=[2, 3, 4],
)
output = unet_cond(x, cond=cond)
```

--------------------------------

### POST /azula/sample/REABSampler

Source: https://github.com/probabilists/azula/blob/master/docs/api/azula.sample.md

Initializes a Rosenbrock-type exponential Adams-Bashforth (REAB) sampler, a multi-step generalization of the DPM-Solver-v3 algorithm.

```APIDOC
## POST /azula/sample/REABSampler

### Description
Creates a Rosenbrock-type exponential Adams-Bashforth (REAB) multi-step sampler based on DPM-Solver-v3.

### Parameters
#### Request Body
- **denoiser** (Denoiser) - Required - A denoiser q_phi(X | X_t).
- **order** (int) - Required - The order n of the multi-step method.
- **kwargs** (dict) - Optional - Keyword arguments passed to Sampler.

### Request Example
{
  "denoiser": "<DenoiserObject>",
  "order": 2
}
```

--------------------------------

### Generate Image with CFG using Sana Model (Python)

Source: https://github.com/probabilists/azula/blob/master/docs/tutorials/latent.ipynb

Generates an image using the Sana model with classifier-free guidance (CFG). It initializes a sampler with CFGDenoiser, generates latent representations, decodes them, and applies post-processing.

```python
sampler = zABSampler(CFGDenoiser(denoiser), steps=16)

z1 = sampler.init((2, 32, 16, 16), device=device)
z0 = sampler(z1, positive=prompt, negative=null, guidance=3.0)

with torch.no_grad():
    x = autoencoder.decode(z0)

to_pil_image(make_grid(postprocess(x)))
```

--------------------------------

### Enable Classifier-Free Guidance

Source: https://context7.com/probabilists/azula/llms.txt

Wraps a base denoiser with CFGDenoiser to enable conditional generation capabilities during the sampling process.

```python
import torch
from azula.guidance import CFGDenoiser
from azula.sample import DDIMSampler

cfg_denoiser = CFGDenoiser(denoiser)
sampler = DDIMSampler(cfg_denoiser, steps=50)
x1 = sampler.init((4, 3, 256, 256), device="cuda")
x0 = sampler(x1)
```

--------------------------------

### Generate Image with CFG using Stable Diffusion (Python)

Source: https://github.com/probabilists/azula/blob/master/docs/tutorials/latent.ipynb

Generates an image using Stable Diffusion with classifier-free guidance (CFG). It utilizes a CFGDenoiser, initializes a sampler, and generates latent representations with positive and negative prompts.

```python
sampler = zABSampler(CFGDenoiser(denoiser), steps=16)

z1 = sampler.init((2, 4, 64, 64), device=device)
z0 = sampler(z1, positive=prompt, negative=null, guidance=3.0)

with torch.no_grad():
    x = autoencoder.decode(z0)

to_pil_image(make_grid(postprocess(x)))
```

--------------------------------

### Load and sample with pre-trained diffusion models

Source: https://github.com/probabilists/azula/blob/master/README.md

Shows how to use the plugin interface to load pre-trained models, such as ADM, and perform image generation using a DDIM sampler.

```python
from azula.plugins import adm
from azula.sample import DDIMSampler

denoiser = adm.load_model("imagenet_256x256")
denoiser.to("cuda")

sampler = DDIMSampler(denoiser, steps=64)
x1 = sampler.init((4, 3, 256, 256), device="cuda")
x0 = sampler(x1)

images = torch.clip((x0 + 1) / 2, min=0, max=1)
```

--------------------------------

### Generate Image using Flux Model (Python)

Source: https://github.com/probabilists/azula/blob/master/docs/tutorials/latent.ipynb

Generates an image using the Flux model. It initializes a sampler, generates latent representations based on the prompt, decodes them into an image, and applies post-processing.

```python
sampler = zABSampler(denoiser, steps=16)

z1 = sampler.init((2, 32, 32, 64), device=device)  # B H W C
z0 = sampler(z1, **prompt)

with torch.no_grad():
    x = autoencoder.decode(z0)

to_pil_image(make_grid(postprocess(x)))
```

--------------------------------

### Load Pre-trained JIT Model

Source: https://github.com/probabilists/azula/blob/master/docs/api/azula.plugins.jit.md

Utility function to load a pre-trained JIT denoiser model by name, with options for EMA weights.

```python
from azula.plugins.jit import load_model

# Load a model by name
model = load_model("jit-model-name", ema=True)
```

--------------------------------

### Initialize GaussianDenoiser for analytical inference

Source: https://context7.com/probabilists/azula/llms.txt

Creates an analytical Gaussian denoiser for scenarios where the prior distribution is known. It performs denoising without requiring a training phase.

```python
import torch
from azula.denoise import GaussianDenoiser
from azula.noise import VPSchedule
from azula.linalg.covariance import DiagonalCovariance

mean = torch.zeros(64)
cov = DiagonalCovariance(torch.ones(64))

denoiser = GaussianDenoiser(mean=mean, cov=cov, schedule=VPSchedule())

x_t = torch.randn(8, 64)
t = torch.tensor(0.5)
posterior = denoiser(x_t, t)
x_denoised = posterior.mean
```

--------------------------------

### POST /azula/sample/PCSampler

Source: https://github.com/probabilists/azula/blob/master/docs/api/azula.sample.md

Initializes a Predictor-Corrector (PC) sampler for diffusion model sampling.

```APIDOC
## POST /azula/sample/PCSampler

### Description
Creates a predictor-corrector (PC) sampler to refine diffusion sampling steps.

### Parameters
#### Request Body
- **denoiser** (Denoiser) - Required - A denoiser q_phi(X | X_t).
- **corrections** (int) - Optional - The number of corrector steps for each predictor step.
- **delta** (float) - Optional - The amplitude of corrector steps delta in [0,1].
- **kwargs** (dict) - Optional - Keyword arguments passed to Sampler.

### Request Example
{
  "denoiser": "<DenoiserObject>",
  "corrections": 1,
  "delta": 0.01
}
```

--------------------------------

### POST /azula.denoise.GaussianPosterior

Source: https://github.com/probabilists/azula/blob/master/docs/api/azula.denoise.md

Initializes a Gaussian posterior distribution and provides methods for log-density calculation.

```APIDOC
## POST /azula.denoise.GaussianPosterior

### Description
Creates a Gaussian posterior $\mathcal{N}(X \mid \mu, \sigma^2)$ and provides the `log_prob` method.

### Method
POST

### Endpoint
azula.denoise.GaussianPosterior

### Parameters
#### Request Body
- **mean** (Tensor) - Required - The mean $\mu$, with shape (*).
- **var** (Tensor) - Required - The variance $\sigma^2$, with shape (*).

### Response
#### Success Response (200)
- **log_prob(x)** (Tensor) - Returns the log-density $\log \mathcal{N}(x \mid \mu, \sigma^2)$.
```

--------------------------------

### Import Stable Diffusion Plugin

Source: https://github.com/probabilists/azula/blob/master/docs/api/azula.plugins.sd.md

Imports the Stable Diffusion plugin module from the azula.plugins package.

```python
from azula.plugins import sd
```

--------------------------------

### Encode Prompt for Flux Model (Python)

Source: https://github.com/probabilists/azula/blob/master/docs/tutorials/latent.ipynb

Encodes a text prompt using the Flux text encoder. This is used for generating images.

```python
with torch.no_grad():
    prompt = textencoder("a forest with a big warning sign that says 'Flux'")
```

--------------------------------

### Class: TDSSampler

Source: https://github.com/probabilists/azula/blob/master/docs/api/azula.guidance.tds.md

Initializes and executes the Twisted Posterior Sampling sampler for diffusion models.

```APIDOC
## POST /azula/guidance/tds/TDSSampler

### Description
Creates a TDS sampler instance to perform conditional sampling using a denoiser and a twisting function.

### Method
POST

### Endpoint
azula.guidance.tds.TDSSampler(denoiser, twist, **kwargs)

### Parameters
#### Request Body
- **denoiser** (Denoiser) - Required - A denoiser q_phi(X | X_t).
- **twist** (Callable) - Required - A twisting function log p(y | hat{x}, t).
- **kwargs** (dict) - Optional - Keyword arguments passed to the base Sampler.

### Request Example
{
  "denoiser": "<Denoiser_Object>",
  "twist": "<Callable_Function>"
}

### Response
#### Success Response (200)
- **sampler** (TDSSampler) - An initialized TDS sampler instance.

---

## POST /azula/guidance/tds/TDSSampler/__call__

### Description
Simulates the reverse diffusion process from t_T to t_0 to generate clean tensors.

### Method
POST

### Endpoint
TDSSampler.__call__(x, **kwargs)

### Parameters
#### Request Body
- **x** (Tensor) - Required - A set of K noisy tensors x_tT, with shape (K, *).
- **kwargs** (dict) - Optional - Additional keyword arguments.

### Request Example
{
  "x": "torch.Tensor(K, ...)"
}

### Response
#### Success Response (200)
- **output** (Tensor) - The clean(er) tensors x_t0, with shape (K, *).

#### Response Example
{
  "output": "torch.Tensor(K, ...)"
}
```

--------------------------------

### Encode Prompt for Sana Model (Python)

Source: https://github.com/probabilists/azula/blob/master/docs/tutorials/latent.ipynb

Encodes a text prompt and an empty string using the Sana text encoder. This is used for generating images with classifier-free guidance.

```python
with torch.no_grad():
    prompt = textencoder("a cyberpunk cat with a neon sign that says 'Sana'")
    null = textencoder("")
```

--------------------------------

### Import Flux plugin

Source: https://github.com/probabilists/azula/blob/master/docs/api/azula.plugins.flux.md

Basic import statement to initialize the Flux plugin within the azula framework.

```python
from azula.plugins import flux
```

--------------------------------

### Skip Weight Initialization Context (Python)

Source: https://github.com/probabilists/azula/blob/master/docs/api/azula.nn.utils.md

Creates a context manager that skips weight initialization for modules created within it. This is useful for layers where custom initialization or pre-trained weights are used.

```python
import torch.nn as nn

class skip_init:
    """Creates a context in which weight initialization is skipped."""
    def __enter__(self):
        # Logic to skip initialization
        pass
    def __exit__(self, exc_type, exc_val, exc_tb):
        # Logic to re-enable initialization if needed
        pass

# Example Usage:
# with skip_init():
#    layer = nn.Linear(3, 5)
```

--------------------------------

### POST /azula.denoise.DiracPosterior

Source: https://github.com/probabilists/azula/blob/master/docs/api/azula.denoise.md

Initializes a Dirac delta posterior distribution, representing a point estimate.

```APIDOC
## POST /azula.denoise.DiracPosterior

### Description
Creates a Dirac delta posterior distribution $\delta(X - \mu)$ for a given mean.

### Method
POST

### Endpoint
azula.denoise.DiracPosterior

### Parameters
#### Request Body
- **mean** (Tensor) - Required - The mean $\mu$, with shape (*).

### Request Example
{
  "mean": "torch.Tensor([...])"
}

### Response
#### Success Response (200)
- **object** (DiracPosterior) - An instance of the DiracPosterior class.
```

--------------------------------

### Set Random Seed (Python)

Source: https://github.com/probabilists/azula/blob/master/docs/tutorials/latent.ipynb

Sets the manual seed for PyTorch's random number generator to ensure reproducibility of results.

```python
_ = torch.manual_seed(42)
```

--------------------------------

### azula.plugins.utils.load_cards

Source: https://github.com/probabilists/azula/blob/master/docs/api/azula.plugins.utils.md

Returns the name-card mapping of pre-trained models available in a plugin.

```APIDOC
## azula.plugins.utils.load_cards

### Description
Returns the name-card mapping of pre-trained models available in a plugin.

### Method
GET

### Endpoint
/probabilists/azula/plugins/utils/load_cards

### Parameters
#### Query Parameters
- **plugin** (ModuleType) - Required - The plugin module.

### Request Example
```json
{
  "plugin": "azula.plugins.adm"
}
```

### Response
#### Success Response (200)
- **cards** (dict) - A dictionary mapping model names to their card information.

#### Response Example
```json
{
  "cards": {
    "imagenet_64x64_cond": { ... },
    "imagenet_128x128_cond": { ... },
    "imagenet_256x256": { ... },
    "imagenet_256x256_cond": { ... },
    "imagenet_512x512_cond": { ... },
    "ffhq_256x256": { ... }
  }
}
```
```

--------------------------------

### RePaintSampler Class

Source: https://github.com/probabilists/azula/blob/master/docs/api/azula.guidance.repaint.md

Initializes the RePaint sampler for performing inpainting tasks using Denoising Diffusion Probabilistic Models.

```APIDOC
## POST /azula/guidance/repaint/RePaintSampler

### Description
Creates a RePaint sampler instance to perform inpainting using Denoising Diffusion Probabilistic Models as described by Lugmayr et al. (2022).

### Method
POST

### Endpoint
azula.guidance.repaint.RePaintSampler

### Parameters
#### Request Body
- **denoiser** (Denoiser) - Required - A denoiser q_phi(X | X_t).
- **y** (Tensor) - Required - An observation y = m * x.
- **mask** (BoolTensor) - Required - The observation mask m.
- **iterations** (int) - Optional - The number of RePaint iterations per step (default: 3).
- **kwargs** (dict) - Optional - Additional keyword arguments passed to DDIMSampler.

### Request Example
{
  "denoiser": "<DenoiserInstance>",
  "y": "<Tensor>",
  "mask": "<BoolTensor>",
  "iterations": 3
}

### Response
#### Success Response (200)
- **instance** (RePaintSampler) - The initialized RePaint sampler object.

#### Response Example
{
  "status": "success",
  "message": "RePaintSampler initialized"
}
```

--------------------------------

### DPSSampler for Inverse Problems with Gradient Guidance (Python)

Source: https://context7.com/probabilists/azula/llms.txt

Implements the DPSSampler for solving inverse problems using gradient guidance. It requires a denoiser, an observation 'y', and a forward measurement operator 'A'. The sampler is initialized with image dimensions and then used to reconstruct the image.

```python
import torch
from azula.guidance import DPSSampler

# Define forward measurement operator
def A(x):
    # Example: downsampling operator
    return torch.nn.functional.interpolate(x, scale_factor=0.25)

# Observation
y = A(original_image)  # Low-res observation

# DPS sampler for super-resolution
sampler = DPSSampler(
    denoiser,
    y=y,
    A=A,
    zeta=1.0,  # Guidance strength
    steps=1000,
)

x1 = sampler.init((1, 3, 256, 256), device="cuda")
x0 = sampler(x1)  # Reconstructed high-res image
```

--------------------------------

### Implement Noise Schedules in Azula

Source: https://context7.com/probabilists/azula/llms.txt

Demonstrates the initialization and usage of various noise schedules including Variance Preserving (VP), Variance Exploding (VE), Cosine, and Rectified schedules to compute signal and noise scales over time.

```python
import torch
from azula.noise import VPSchedule, VESchedule, CosineSchedule, RectifiedSchedule

# VP Schedule
schedule_vp = VPSchedule(alpha_min=1e-3, sigma_min=1e-3)
t = torch.tensor([0.0, 0.5, 1.0])
alpha_vp, sigma_vp = schedule_vp(t)

# VE Schedule
schedule_ve = VESchedule(sigma_min=1e-3, sigma_max=1e3)
alpha_ve, sigma_ve = schedule_ve(t)

# Cosine Schedule
schedule_cos = CosineSchedule(alpha_min=1e-3, sigma_min=1e-3)
alpha_cos, sigma_cos = schedule_cos(t)

# Rectified Schedule
schedule_rect = RectifiedSchedule(alpha_min=1e-3, sigma_min=1e-3)
alpha_rect, sigma_rect = schedule_rect(t)
```

--------------------------------

### Encode Prompt for Stable Diffusion (Python)

Source: https://github.com/probabilists/azula/blob/master/docs/tutorials/latent.ipynb

Encodes a given text prompt and an empty string using the Stable Diffusion text encoder. This is used for generating images with and without classifier-free guidance.

```python
with torch.no_grad():
    prompt = textencoder("an astronaut riding a horse in space")
    null = textencoder("")
```

--------------------------------

### Class: JITDenoiser

Source: https://github.com/probabilists/azula/blob/master/docs/api/azula.plugins.jit.md

The JITDenoiser class provides the core functionality for creating and executing a time-conditional denoising network.

```APIDOC
## CLASS azula.plugins.jit.JITDenoiser

### Description
Creates a JIT denoiser instance for processing noisy tensors using a time-conditional backbone.

### Parameters
- **backbone** (Module) - Required - A time conditional network.
- **schedule** (Schedule) - Optional - A noise schedule (defaults to RectifiedSchedule).
- **num_classes** (int) - Optional - The number of classes for conditional generation.

### Forward Method
- **x_t** (Tensor) - Required - Noisy tensor with shape (B, 3, H, W).
- **t** (Tensor) - Required - Time value with shape () or (B).
- **label** (Tensor) - Optional - Class label integer with shape (B).

### Response
- **Returns** (DiracPosterior) - The Dirac delta distribution representing the denoised output.
```

--------------------------------

### Train and Use KarrasDenoiser

Source: https://context7.com/probabilists/azula/llms.txt

Shows how to wrap a neural network backbone with KarrasDenoiser for EDM-style training and inference. It includes a simple backbone definition, loss calculation, and clean image prediction.

```python
import torch
import torch.nn as nn
from azula.denoise import KarrasDenoiser
from azula.noise import VPSchedule

class SimpleBackbone(nn.Module):
    def __init__(self, channels=3, features=64):
        super().__init__()
        self.net = nn.Sequential(nn.Conv2d(channels, features, 3, padding=1), nn.SiLU(), nn.Conv2d(features, channels, 3, padding=1))

    def forward(self, x, t):
        return self.net(x)

schedule = VPSchedule()
denoiser = KarrasDenoiser(backbone=SimpleBackbone(), schedule=schedule)

# Training step
x = torch.randn(32, 3, 32, 32)
t = torch.rand(32)
loss = denoiser.loss(x, t)

# Inference
x_t = torch.randn(1, 3, 32, 32)
posterior = denoiser(x_t, torch.tensor(0.5))
x_pred = posterior.mean
```

--------------------------------

### Define Measurement Operators for Image Restoration

Source: https://github.com/probabilists/azula/blob/master/docs/tutorials/guidance.md

Defines forward and inverse operators (A and A_inv) for downsampling and restoration, and prepares a sample image for testing.

```ipython3
image = requests.get("https://upload.wikimedia.org/wikipedia/commons/3/3a/Cat03.jpg", headers={"User-Agent": "Azula"}).content
image = io.BytesIO(image)
image = Image.open(image).convert("RGB")
image = image.crop((0, 0, min(image.size), min(image.size))).resize((256, 256))

x = preprocess(to_tensor(image)).to(device)

def A(x):
    return torch.nn.functional.interpolate(x, (32, 32), mode="bicubic", antialias=True).flatten(-3)

def A_inv(y):
    return torch.nn.functional.interpolate(
        y.unflatten(-1, (3, 32, 32)), (256, 256), mode="nearest"
    )

sigma_y = 0.01
y = A(x.unsqueeze(0))
y = y + sigma_y * torch.randn_like(y)
```

--------------------------------

### azula.nn.layers.Unpatchify

Source: https://github.com/probabilists/azula/blob/master/docs/api/azula.nn.layers.md

Creates a layer that transforms channel dimensions back into patches.

```APIDOC
## azula.nn.layers.Unpatchify(patch_shape, channel_last=False)

### Description
Returns a channel-to-patch layer. This layer reverses the operation of Patchify, reshaping channel dimensions back into spatial patches.

### Method
N/A (This is a class constructor for a layer)

### Endpoint
N/A

### Parameters
#### Path Parameters
None

#### Query Parameters
None

#### Request Body
None

### Parameters
* **patch_shape** (Sequence[int]) - Required - The shape of the patches to reconstruct into the input.
* **channel_last** (bool) - Optional - Defaults to False. Specifies whether the channel dimension is the last dimension (True) or the first dimension (False) in the input and output tensors.
```

--------------------------------

### Import Sana Plugin

Source: https://github.com/probabilists/azula/blob/master/docs/api/azula.plugins.sana.md

Import the Sana plugin module into your Python environment.

```python
from azula.plugins import sana
```

--------------------------------

### ElucidatedSchedule Class Initialization

Source: https://github.com/probabilists/azula/blob/master/docs/api/azula.plugins.edm.md

Initializes an ElucidatedSchedule object, which defines a noise schedule for diffusion models. It takes parameters for the minimum and maximum noise scales (sigma_min, sigma_max) and a hyper-parameter rho, influencing the schedule's shape.

```python
schedule = edm.ElucidatedSchedule(sigma_min=0.002, sigma_max=80.0, rho=7.0)
```