autoregressive-diffusion-pytorch 0.2.0__tar.gz → 0.2.2__tar.gz

{autoregressive_diffusion_pytorch-0.2.0 → autoregressive_diffusion_pytorch-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: autoregressive-diffusion-pytorch
-Version: 0.2.0
+Version: 0.2.2
 Summary: Autoregressive Diffusion - Pytorch
 Project-URL: Homepage, https://pypi.org/project/autoregressive-diffusion-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/autoregressive-diffusion-pytorch
@@ -52,9 +52,9 @@ Implementation of the architecture behind <a href="https://arxiv.org/abs/2406.11
 
 Official repository has been released <a href="https://github.com/LTH14/mar">here</a>
 
-<img src="./images/sample.flowers.59000.png" width="400px"></img>
+<img src="./images/results.96600.png" width="400px"></img>
 
-*oxford flowers at 59k steps*
+*oxford flowers at 96k steps*
 
 ## Install
 
@@ -115,6 +115,74 @@ assert sampled.shape == images.shape
 
 ```
 
+An images trainer
+
+```python
+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, just import `ImageAutoregressiveFlow` and `AutoregressiveFlow` instead
+
+The rest is the same
+
+ex.
+
+```python
+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()
+```
+
 ## Citations
 
 ```bibtex

{autoregressive_diffusion_pytorch-0.2.0 → autoregressive_diffusion_pytorch-0.2.2}/README.md

@@ -6,9 +6,9 @@ Implementation of the architecture behind <a href="https://arxiv.org/abs/2406.11
 
 Official repository has been released <a href="https://github.com/LTH14/mar">here</a>
 
-<img src="./images/sample.flowers.59000.png" width="400px"></img>
+<img src="./images/results.96600.png" width="400px"></img>
 
-*oxford flowers at 59k steps*
+*oxford flowers at 96k steps*
 
 ## Install
 
@@ -69,6 +69,74 @@ assert sampled.shape == images.shape
 
 ```
 
+An images trainer
+
+```python
+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, just import `ImageAutoregressiveFlow` and `AutoregressiveFlow` instead
+
+The rest is the same
+
+ex.
+
+```python
+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()
+```
+
 ## Citations
 
 ```bibtex

autoregressive_diffusion_pytorch-0.2.2/autoregressive_diffusion_pytorch/__init__.py

@@ -0,0 +1,15 @@
+from autoregressive_diffusion_pytorch.autoregressive_diffusion import (
+    MLP,
+    AutoregressiveDiffusion,
+    ImageAutoregressiveDiffusion
+)
+
+from autoregressive_diffusion_pytorch.autoregressive_flow import (
+    AutoregressiveFlow,
+    ImageAutoregressiveFlow
+)
+
+from autoregressive_diffusion_pytorch.image_trainer import (
+    ImageDataset,
+    ImageTrainer
+)

{autoregressive_diffusion_pytorch-0.2.0 → autoregressive_diffusion_pytorch-0.2.2}/autoregressive_diffusion_pytorch/autoregressive_diffusion.py

@@ -231,7 +231,7 @@ class ElucidatedDiffusion(Module):
         if isinstance(sigma, float):
             sigma = torch.full((batch,), sigma, device = device)
 
-        padded_sigma = rearrange(sigma, 'b -> b 1')
+        padded_sigma = right_pad_dims_to(noised_seq, sigma)
 
         net_out = self.net(
             self.c_in(padded_sigma) * noised_seq,
@@ -331,7 +331,7 @@ class ElucidatedDiffusion(Module):
         assert dim == self.dim, f'dimension of sequence being passed in must be {self.dim} but received {dim}'
 
         sigmas = self.noise_distribution(batch_size)
-        padded_sigmas = rearrange(sigmas, 'b -> b 1')
+        padded_sigmas = right_pad_dims_to(seq, sigmas)
 
         noise = torch.randn_like(seq)
 

autoregressive_diffusion_pytorch-0.2.2/autoregressive_diffusion_pytorch/autoregressive_flow.py

@@ -0,0 +1,335 @@
+from __future__ import annotations
+
+import math
+from math import sqrt
+from typing import Literal
+from functools import partial
+
+import torch
+from torch import nn, pi
+import torch.nn.functional as F
+from torch.nn import Module, ModuleList
+
+from torchdiffeq import odeint
+
+import einx
+from einops import rearrange, repeat, reduce, pack, unpack
+from einops.layers.torch import Rearrange
+
+from tqdm import tqdm
+
+from x_transformers import Decoder
+
+from autoregressive_diffusion_pytorch.autoregressive_diffusion import MLP
+
+# helpers
+
+def exists(v):
+    return v is not None
+
+def default(v, d):
+    return v if exists(v) else d
+
+def divisible_by(num, den):
+    return (num % den) == 0
+
+# tensor helpers
+
+def log(t, eps = 1e-20):
+    return torch.log(t.clamp(min = eps))
+
+def safe_div(num, den, eps = 1e-5):
+    return num / den.clamp(min = eps)
+
+def right_pad_dims_to(x, t):
+    padding_dims = x.ndim - t.ndim
+
+    if padding_dims <= 0:
+        return t
+
+    return t.view(*t.shape, *((1,) * padding_dims))
+
+def pack_one(t, pattern):
+    packed, ps = pack([t], pattern)
+
+    def unpack_one(to_unpack, unpack_pattern = None):
+        unpacked, = unpack(to_unpack, ps, default(unpack_pattern, pattern))
+        return unpacked
+
+    return packed, unpack_one
+
+# sinusoidal embedding
+
+class AdaptiveLayerNorm(Module):
+    def __init__(
+        self,
+        dim,
+        dim_condition = None
+    ):
+        super().__init__()
+        dim_condition = default(dim_condition, dim)
+
+        self.ln = nn.LayerNorm(dim, elementwise_affine = False)
+        self.to_gamma = nn.Linear(dim_condition, dim, bias = False)
+        nn.init.zeros_(self.to_gamma.weight)
+
+    def forward(self, x, *, condition):
+        normed = self.ln(x)
+        gamma = self.to_gamma(condition)
+        return normed * (gamma + 1.)
+
+class LearnedSinusoidalPosEmb(Module):
+    def __init__(self, dim):
+        super().__init__()
+        assert divisible_by(dim, 2)
+        half_dim = dim // 2
+        self.weights = nn.Parameter(torch.randn(half_dim))
+
+    def forward(self, x):
+        x = rearrange(x, 'b -> b 1')
+        freqs = x * rearrange(self.weights, 'd -> 1 d') * 2 * pi
+        fouriered = torch.cat((freqs.sin(), freqs.cos()), dim = -1)
+        fouriered = torch.cat((x, fouriered), dim = -1)
+        return fouriered
+
+# gaussian diffusion
+
+class Flow(Module):
+    def __init__(
+        self,
+        dim: int,
+        net: MLP,
+        *,
+        atol = 1e-5,
+        rtol = 1e-5,
+        method = 'midpoint'
+    ):
+        super().__init__()
+        self.net = net
+        self.dim = dim
+
+        self.odeint_kwargs = dict(
+            atol = atol,
+            rtol = rtol,
+            method = method
+        )
+
+    @property
+    def device(self):
+        return next(self.net.parameters()).device
+
+    @torch.no_grad()
+    def sample(
+        self,
+        cond,
+        num_sample_steps = 16
+    ):
+
+        batch = cond.shape[0]
+
+        sampled_data_shape = (batch, self.dim)
+
+        # start with random gaussian noise - y0
+
+        noise = torch.randn(sampled_data_shape, device = self.device)
+
+        # time steps
+
+        times = torch.linspace(0., 1., num_sample_steps, device = self.device)
+
+        # ode
+
+        def ode_fn(t, x):
+            t = repeat(t, '-> b', b = batch)
+            flow = self.net(x, times = t, cond = cond)
+            return flow
+
+        trajectory = odeint(ode_fn, noise, times, **self.odeint_kwargs)
+
+        sampled = trajectory[-1]
+
+        return sampled
+
+    # training
+
+    def forward(self, seq, *, cond):
+        batch_size, dim, device = *seq.shape, self.device
+
+        assert dim == self.dim, f'dimension of sequence being passed in must be {self.dim} but received {dim}'
+
+        times = torch.rand(batch_size, device = device)
+        noise = torch.randn_like(seq)
+        padded_times = right_pad_dims_to(seq, times)
+
+        flow = seq - noise
+
+        noised = (1.- padded_times) * noise + padded_times * seq
+
+        pred_flow = self.net(noised, times = times, cond = cond)
+
+        return F.mse_loss(pred_flow, flow)
+
+# main model, a decoder with continuous wrapper + small denoising mlp
+
+class AutoregressiveFlow(Module):
+    def __init__(
+        self,
+        dim,
+        *,
+        max_seq_len,
+        depth = 8,
+        dim_head = 64,
+        heads = 8,
+        mlp_depth = 3,
+        mlp_width = None,
+        dim_input = None,
+        decoder_kwargs: dict = dict(),
+        mlp_kwargs: dict = dict(),
+        flow_kwargs: dict = dict()
+    ):
+        super().__init__()
+
+        self.start_token = nn.Parameter(torch.zeros(dim))
+        self.max_seq_len = max_seq_len
+        self.abs_pos_emb = nn.Embedding(max_seq_len, dim)
+
+        dim_input = default(dim_input, dim)
+        self.dim_input = dim_input
+        self.proj_in = nn.Linear(dim_input, dim)
+
+        self.transformer = Decoder(
+            dim = dim,
+            depth = depth,
+            heads = heads,
+            attn_dim_head = dim_head,
+            **decoder_kwargs
+        )
+
+        self.denoiser = MLP(
+            dim_cond = dim,
+            dim_input = dim_input,
+            depth = mlp_depth,
+            width = default(mlp_width, dim),
+            **mlp_kwargs
+        )
+
+        self.flow = Flow(
+            dim_input,
+            self.denoiser,
+            **flow_kwargs
+        )
+
+    @property
+    def device(self):
+        return next(self.transformer.parameters()).device
+
+    @torch.no_grad()
+    def sample(
+        self,
+        batch_size = 1,
+        prompt  = None
+    ):
+        self.eval()
+
+        start_tokens = repeat(self.start_token, 'd -> b 1 d', b = batch_size)
+
+        if not exists(prompt):
+            out = torch.empty((batch_size, 0, self.dim_input), device = self.device, dtype = torch.float32)
+        else:
+            out = prompt
+
+        cache = None
+
+        for _ in tqdm(range(self.max_seq_len - out.shape[1]), desc = 'tokens'):
+
+            cond = self.proj_in(out)
+
+            cond = torch.cat((start_tokens, cond), dim = 1)
+            cond = cond + self.abs_pos_emb(torch.arange(cond.shape[1], device = self.device))
+
+            cond, cache = self.transformer(cond, cache = cache, return_hiddens = True)
+
+            last_cond = cond[:, -1]
+
+            denoised_pred = self.flow.sample(cond = last_cond)
+
+            denoised_pred = rearrange(denoised_pred, 'b d -> b 1 d')
+            out = torch.cat((out, denoised_pred), dim = 1)
+
+        return out
+
+    def forward(
+        self,
+        seq
+    ):
+        b, seq_len, dim = seq.shape
+
+        assert dim == self.dim_input
+        assert seq_len == self.max_seq_len
+
+        # break into seq and the continuous targets to be predicted
+
+        seq, target = seq[:, :-1], seq
+
+        # append start tokens
+
+        seq = self.proj_in(seq)
+        start_token = repeat(self.start_token, 'd -> b 1 d', b = b)
+
+        seq = torch.cat((start_token, seq), dim = 1)
+        seq = seq + self.abs_pos_emb(torch.arange(seq_len, device = self.device))
+
+        cond = self.transformer(seq)
+
+        # pack batch and sequence dimensions, so to train each token with different noise levels
+
+        target, _ = pack_one(target, '* d')
+        cond, _ = pack_one(cond, '* d')
+
+        return self.flow(target, cond = cond)
+
+# image wrapper
+
+def normalize_to_neg_one_to_one(img):
+    return img * 2 - 1
+
+def unnormalize_to_zero_to_one(t):
+    return (t + 1) * 0.5
+
+class ImageAutoregressiveFlow(Module):
+    def __init__(
+        self,
+        *,
+        image_size,
+        patch_size,
+        channels = 3,
+        model: dict = dict(),
+    ):
+        super().__init__()
+        assert divisible_by(image_size, patch_size)
+
+        num_patches = (image_size // patch_size) ** 2
+        dim_in = channels * patch_size ** 2
+
+        self.image_size = image_size
+        self.patch_size = patch_size
+
+        self.to_tokens = Rearrange('b c (h p1) (w p2) -> b (h w) (c p1 p2)', p1 = patch_size, p2 = patch_size)
+
+        self.model = AutoregressiveFlow(
+            **model,
+            dim_input = dim_in,
+            max_seq_len = num_patches
+        )
+
+        self.to_image = Rearrange('b (h w) (c p1 p2) -> b c (h p1) (w p2)', p1 = patch_size, p2 = patch_size, h = int(math.sqrt(num_patches)))
+
+    def sample(self, batch_size = 1):
+        tokens = self.model.sample(batch_size = batch_size)
+        images = self.to_image(tokens)
+        return unnormalize_to_zero_to_one(images)
+
+    def forward(self, images):
+        images = normalize_to_neg_one_to_one(images)
+        tokens = self.to_tokens(images)
+        return self.model(tokens)

autoregressive_diffusion_pytorch-0.2.2/autoregressive_diffusion_pytorch/image_trainer.py

@@ -0,0 +1,194 @@
+from typing import List
+
+import math
+from pathlib import Path
+
+from accelerate import Accelerator
+from ema_pytorch import EMA
+
+import torch
+from torch import nn
+from torch.optim import Adam
+from torch.utils.data import DataLoader
+from torch.nn import Module, ModuleList
+from torch.utils.data import Dataset
+
+from torchvision.utils import save_image
+import torchvision.transforms as T
+
+from PIL import Image
+
+# functions
+
+def exists(v):
+    return v is not None
+
+def default(v, d):
+    return v if exists(v) else d
+
+def divisible_by(num, den):
+    return (num % den) == 0
+
+def cycle(dl):
+    while True:
+        for batch in dl:
+            yield batch
+
+# dataset classes
+
+class ImageDataset(Dataset):
+    def __init__(
+        self,
+        folder: str | Path,
+        image_size: int,
+        exts: List[str] = ['jpg', 'jpeg', 'png', 'tiff'],
+        augment_horizontal_flip = False,
+        convert_image_to = None
+    ):
+        super().__init__()
+        if isinstance(folder, str):
+            folder = Path(folder)
+
+        assert folder.is_dir()
+
+        self.folder = folder
+        self.image_size = image_size
+
+        self.paths = [p for ext in exts for p in folder.glob(f'**/*.{ext}')]
+
+        def convert_image_to_fn(img_type, image):
+            if image.mode == img_type:
+                return image
+
+            return image.convert(img_type)
+
+        maybe_convert_fn = partial(convert_image_to_fn, convert_image_to) if exists(convert_image_to) else nn.Identity()
+
+        self.transform = T.Compose([
+            T.Lambda(maybe_convert_fn),
+            T.Resize(image_size),
+            T.RandomHorizontalFlip() if augment_horizontal_flip else nn.Identity(),
+            T.CenterCrop(image_size),
+            T.ToTensor()
+        ])
+
+    def __len__(self):
+        return len(self.paths)
+
+    def __getitem__(self, index):
+        path = self.paths[index]
+        img = Image.open(path)
+        return self.transform(img)
+
+# trainer
+
+class ImageTrainer(Module):
+    def __init__(
+        self,
+        model,
+        *,
+        dataset: Dataset,
+        num_train_steps = 70_000,
+        learning_rate = 3e-4,
+        batch_size = 16,
+        checkpoints_folder: str = './checkpoints',
+        results_folder: str = './results',
+        save_results_every: int = 100,
+        checkpoint_every: int = 1000,
+        num_samples: int = 16,
+        adam_kwargs: dict = dict(),
+        accelerate_kwargs: dict = dict(),
+        ema_kwargs: dict = dict()
+    ):
+        super().__init__()
+        self.accelerator = Accelerator(**accelerate_kwargs)
+
+        self.model = model
+
+        if self.is_main:
+            self.ema_model = EMA(
+                self.model,
+                forward_method_names = ('sample',),
+                **ema_kwargs
+            )
+
+            self.ema_model.to(self.accelerator.device)
+
+        self.optimizer = Adam(model.parameters(), lr = learning_rate, **adam_kwargs)
+        self.dl = DataLoader(dataset, batch_size = batch_size, shuffle = True, drop_last = True)
+
+        self.model, self.optimizer, self.dl = self.accelerator.prepare(self.model, self.optimizer, self.dl)
+
+        self.num_train_steps = num_train_steps
+
+        self.checkpoints_folder = Path(checkpoints_folder)
+        self.results_folder = Path(results_folder)
+
+        self.checkpoints_folder.mkdir(exist_ok = True, parents = True)
+        self.results_folder.mkdir(exist_ok = True, parents = True)
+
+        self.checkpoint_every = checkpoint_every
+        self.save_results_every = save_results_every
+
+        self.num_sample_rows = int(math.sqrt(num_samples))
+        assert (self.num_sample_rows ** 2) == num_samples, f'{num_samples} must be a square'
+        self.num_samples = num_samples
+
+        assert self.checkpoints_folder.is_dir()
+        assert self.results_folder.is_dir()
+
+    @property
+    def is_main(self):
+        return self.accelerator.is_main_process
+
+    def save(self, path):
+        if not self.is_main:
+            return
+
+        save_package = dict(
+            model = self.accelerator.unwrap_model(self.model).state_dict(),
+            ema_model = self.ema_model.state_dict(),
+            optimizer = self.accelerator.unwrap_model(self.optimizer).state_dict(),
+        )
+
+        torch.save(save_package, str(self.checkpoints_folder / path))
+
+    def forward(self):
+
+        dl = cycle(self.dl)
+
+        for ind in range(self.num_train_steps):
+            step = ind + 1
+
+            self.model.train()
+
+            data = next(dl)
+            loss = self.model(data)
+
+            self.accelerator.print(f'[{step}] loss: {loss.item():.3f}')
+            self.accelerator.backward(loss)
+
+            self.optimizer.step()
+            self.optimizer.zero_grad()
+
+            if self.is_main:
+                self.ema_model.update()
+
+            self.accelerator.wait_for_everyone()
+
+            if self.is_main:
+                if divisible_by(step, self.save_results_every):
+
+                    with torch.no_grad():
+                        sampled = self.ema_model.sample(batch_size = self.num_samples)
+
+                    sampled.clamp_(0., 1.)
+                    save_image(sampled, str(self.results_folder / f'results.{step}.png'), nrow = self.num_sample_rows)
+
+                if divisible_by(step, self.checkpoint_every):
+                    self.save(f'checkpoint.{step}.pt')
+
+            self.accelerator.wait_for_everyone()
+
+
+        print('training complete')

{autoregressive_diffusion_pytorch-0.2.0 → autoregressive_diffusion_pytorch-0.2.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "autoregressive-diffusion-pytorch"
-version = "0.2.0"
+version = "0.2.2"
 description = "Autoregressive Diffusion - Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

autoregressive_diffusion_pytorch-0.2.0/autoregressive_diffusion_pytorch/__init__.py

@@ -1,5 +0,0 @@
-from autoregressive_diffusion_pytorch.autoregressive_diffusion import (
-    AutoregressiveDiffusion,
-    MLP,
-    ImageAutoregressiveDiffusion
-)