autoregressive-diffusion-pytorch

Autoregressive Diffusion - Pytorch

Autoregressive Diffusion in Pytorch is a notable implementation of the architecture put forth in the paper "Autoregressive Image Generation without Vector Quantization." This project provides tools and methodologies to facilitate advanced image generation processes without relying on vector quantization. For those interested in the finer details or official updates, the project hosts an official repository online. An alternative project route is also available for enthusiasts seeking different perspectives or additional functionalities.

Installation

The package can be conveniently installed via pip with a single command:

$ pip install autoregressive-diffusion-pytorch

Basic Usage

One of the core elements of this package is the ease of use it offers for implementing autoregressive models. Presented below is a sample usage scenario, illustrating how to define and train a basic autoregressive model:

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

Image-Based Models

For handling images as sequences of tokens, the package provides a model specifically tailored for image data:

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

This approach allows the transformation of images into a format that can be processed by autoregressive models, as described in the associated research paper.

Image Training

To engage with images in a more comprehensive fashion, users can utilize the ImageTrainer utility, which facilitates the training of models on given image datasets:

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

This block of code underscores how to set up and train a model, ensuring that the desired outcomes can be achieved efficiently.

Flow Matching for Improvisation

The package also offers an alternative approach to image modeling through flow matching:

from autoregressive_diffusion_pytorch import (
    ImageDataset,
    ImageTrainer,
    ImageAutoregressiveFlow,
)

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

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

trainer()

By importing different components, as shown above, users can experiment with various strategies for image generation and enhancement tasks.

Citations

For academic purposes or further reading, here are some of the principal papers and their citations associated with the developments and ideas in this project:

@article{Li2024AutoregressiveIG,
    title   = {Autoregressive Image Generation without Vector Quantization},
    author  = {Tianhong Li and Yonglong Tian and He Li and Mingyang Deng and Kaiming He},
    journal = {ArXiv},
    year    = {2024},
    volume  = {abs/2406.11838},
    url     = {https://api.semanticscholar.org/CorpusID:270560593}
}

@article{Wu2023ARDiffusionAD,
    title     = {AR-Diffusion: Auto-Regressive Diffusion Model for Text Generation},
    author    = {Tong Wu and Zhihao Fan and Xiao Liu and Yeyun Gong and Yelong Shen and Jian Jiao and Haitao Zheng and Juntao Li and Zhongyu Wei and Jian Guo and Nan Duan and Weizhu Chen},
    journal   = {ArXiv},
    year      = {2023},
    volume    = {abs/2305.09515},
    url       = {https://api.semanticscholar.org/CorpusID:258714669}
}

These references offer valuable insights into the theoretical foundations and potential applications of autoregressive diffusion models.