Autoregressive Diffusion PyTorch

Install the package using pip:

$ pip install autoregressive-diffusion-pytorch

To train an image model, combine ImageDataset, ImageAutoregressiveDiffusion, and ImageTrainer. The ImageDataset loads images from a directory, and the ImageTrainer handles the training loop.

import torch

from autoregressive_diffusion_pytorch import (
    ImageDataset,
    ImageAutoregressiveDiffusion,
    ImageTrainer
)

dataset = ImageDataset(
    '/path/to/your/images',
    image_size = 128
)

model = ImageAutoregressiveDiffusion(
    model = dict(
        dim = 512
    ),
    image_size = 128,
    patch_size = 16
)

trainer = ImageTrainer(
    model = model,
    dataset = dataset
)

trainer()

For an improvised version using flow matching instead of standard diffusion, use ImageAutoregressiveFlow and AutoregressiveFlow. The training workflow with ImageDataset and ImageTrainer remains identical.

import torch

from autoregressive_diffusion_pytorch import (
    ImageDataset,
    ImageTrainer,
    ImageAutoregressiveFlow,
)

dataset = ImageDataset(
    '/path/to/your/images',
    image_size = 128
)

model = ImageAutoregressiveFlow(
    model = dict(
        dim = 512
    ),
    image_size = 128,
    patch_size = 16
)

trainer = ImageTrainer(
    model = model,
    dataset = dataset
)

trainer()

The AutoregressiveDiffusion class is used for processing sequences of tokens. You can pass a sequence to the model to compute the loss and use the .sample() method to generate new sequences.

import torch
from autoregressive_diffusion_pytorch import AutoregressiveDiffusion

model = AutoregressiveDiffusion(
    dim_input = 512,
    dim = 1024,
    max_seq_len = 32,
    depth = 8,
    mlp_depth = 3,
    mlp_width = 1024
)

seq = torch.randn(3, 32, 512)

loss = model(seq)
loss.backward()

sampled = model.sample(batch_size = 3)

assert sampled.shape == seq.shape

For treating images as a sequence of tokens (as described in the paper), use ImageAutoregressiveDiffusion. It requires an image_size and a patch_size to handle the spatial dimensions.

import torch
from autoregressive_diffusion_pytorch import ImageAutoregressiveDiffusion

model = ImageAutoregressiveDiffusion(
    model = dict(
        dim = 1024,
        depth = 12,
        heads = 12,
    ),
    image_size = 64,
    patch_size = 8
)

images = torch.randn(3, 3, 64, 64)

loss = model(images)
loss.backward()

sampled = model.sample(batch_size = 3)

assert sampled.shape == images.shape

When initializing AutoregressiveDiffusion, you can pass specific keyword arguments to its sub-components:

• dim: The dimension of the transformer's hidden states.
• max_seq_len: The maximum sequence length the model can handle.
• depth: Number of transformer layers.
• dim_head: Dimension of each attention head.
• heads: Number of attention heads.
• mlp_depth: Depth of the diffusion MLP.
• mlp_width: Width of the diffusion MLP (defaults to dim).
• dim_input: The dimension of the continuous tokens being predicted (must match the patch dimension for images).
• decoder_kwargs: Dictionary of arguments passed to the x-transformers Decoder.
• mlp_kwargs: Dictionary of arguments passed to the diffusion MLP.
• diffusion_kwargs: Dictionary of arguments passed to ElucidatedDiffusion (e.g., clamp_during_sampling=True).

The ImageAutoregressiveDiffusion constructor accepts the following arguments:

• image_size: The height/width of the square image.
• patch_size: The size of the square patches.
• channels: Number of color channels (e.g., 3 for RGB).
• model: A dictionary containing the configuration for the underlying AutoregressiveDiffusion model. This must include dim and max_seq_len (calculated as (image_size // patch_size) ** 2), and dim_input (calculated as channels * patch_size ** 2).

When initializing ImageTrainer, you can pass keyword arguments to underlying components:

• adam_kwargs: Dictionary of arguments passed to torch.optim.Adam.
• accelerate_kwargs: Dictionary of arguments passed to accelerate.Accelerator.
• ema_kwargs: Dictionary of arguments passed to ema_pytorch.EMA (note: forward_method_names is hardcoded to ('sample',)).
• num_samples: The number of images to sample for results. This must be a square number (e.g., 16, 64, 100) because the trainer arranges them in a square grid.

The autoregressive_diffusion_pytorch package provides classes for autoregressive diffusion models. You can import the core diffusion classes directly from the package root:

• AutoregressiveDiffusion: The base autoregressive diffusion model.
• ImageAutoregressiveDiffusion: A specialized version of the autoregressive diffusion model designed for image data.
• MLP: A Multi-Layer Perceptron utility used within the models.

from autoregressive_diffusion_pytorch import (
    AutoregressiveDiffusion,
    ImageAutoregressiveDiffusion,
    MLP
)

For training models on image datasets, use the following utilities exported from the package root:

• ImageTrainer: The main class for managing the training process.
• ImageDataset: A dataset abstraction for handling image data during training.

from autoregressive_diffusion_pytorch import (
    ImageTrainer,
    ImageDataset
)