autoregressive-diffusion-pytorch 0.2.4__tar.gz → 0.2.7__tar.gz

{autoregressive_diffusion_pytorch-0.2.4 → autoregressive_diffusion_pytorch-0.2.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: autoregressive-diffusion-pytorch
-Version: 0.2.4
+Version: 0.2.7
 Summary: Autoregressive Diffusion - Pytorch
 Project-URL: Homepage, https://pypi.org/project/autoregressive-diffusion-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/autoregressive-diffusion-pytorch

{autoregressive_diffusion_pytorch-0.2.4 → autoregressive_diffusion_pytorch-0.2.7}/autoregressive_diffusion_pytorch/autoregressive_flow.py

@@ -33,6 +33,9 @@ def default(v, d):
 def divisible_by(num, den):
     return (num % den) == 0
 
+def cast_tuple(t):
+    return (t,) if not isinstance(t, tuple) else t
+
 # tensor helpers
 
 def log(t, eps = 1e-20):
@@ -58,41 +61,7 @@ def pack_one(t, pattern):
 
     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
+# rectified flow
 
 class Flow(Module):
     def __init__(
@@ -176,12 +145,12 @@ class AutoregressiveFlow(Module):
         self,
         dim,
         *,
-        max_seq_len,
+        max_seq_len: int | tuple[int, ...],
         depth = 8,
         dim_head = 64,
         heads = 8,
         mlp_depth = 3,
-        mlp_width = None,
+        mlp_width = 1024,
         dim_input = None,
         decoder_kwargs: dict = dict(),
         mlp_kwargs: dict = dict(),
@@ -190,8 +159,11 @@ class AutoregressiveFlow(Module):
         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)
+
+        max_seq_len = cast_tuple(max_seq_len)
+        self.abs_pos_emb = nn.ParameterList([nn.Parameter(torch.zeros(seq_len, dim)) for seq_len in max_seq_len])
+
+        self.max_seq_len = math.prod(max_seq_len)
 
         dim_input = default(dim_input, dim)
         self.dim_input = dim_input
@@ -205,11 +177,13 @@ class AutoregressiveFlow(Module):
             **decoder_kwargs
         )
 
+        self.to_cond_emb = nn.Linear(dim, dim, bias = False)
+
         self.denoiser = MLP(
             dim_cond = dim,
             dim_input = dim_input,
             depth = mlp_depth,
-            width = default(mlp_width, dim),
+            width = mlp_width,
             **mlp_kwargs
         )
 
@@ -223,6 +197,16 @@ class AutoregressiveFlow(Module):
     def device(self):
         return next(self.transformer.parameters()).device
 
+    def axial_pos_emb(self):
+        # prepare maybe axial positional embedding
+
+        pos_emb, *rest_pos_embs = self.abs_pos_emb
+
+        for rest_pos_emb in rest_pos_embs:
+            pos_emb = einx.add('i d, j d -> (i j) d', pos_emb, rest_pos_emb)
+
+        return F.pad(pos_emb, (0, 0, 1, 0), value = 0.)
+
     @torch.no_grad()
     def sample(
         self,
@@ -245,12 +229,18 @@ class AutoregressiveFlow(Module):
             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))
+
+            seq_len = cond.shape[-2]
+            axial_pos_emb = self.axial_pos_emb()
+            cond += axial_pos_emb[:seq_len]
 
             cond, cache = self.transformer(cond, cache = cache, return_hiddens = True)
 
             last_cond = cond[:, -1]
 
+            last_cond += axial_pos_emb[seq_len]
+            last_cond = self.to_cond_emb(last_cond)
+
             denoised_pred = self.flow.sample(cond = last_cond)
 
             denoised_pred = rearrange(denoised_pred, 'b d -> b 1 d')
@@ -281,10 +271,15 @@ class AutoregressiveFlow(Module):
         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))
+
+        axial_pos_emb = self.axial_pos_emb()
+        seq = seq + axial_pos_emb[:seq_len]
 
         cond = self.transformer(seq)
 
+        cond = cond + axial_pos_emb[1:(seq_len + 1)]
+        cond = self.to_cond_emb(cond)
+
         # pack batch and sequence dimensions, so to train each token with different noise levels
 
         target, _ = pack_one(target, '* d')
@@ -313,7 +308,8 @@ class ImageAutoregressiveFlow(Module):
         super().__init__()
         assert divisible_by(image_size, patch_size)
 
-        num_patches = (image_size // patch_size) ** 2
+        patch_height_width = image_size // patch_size
+        num_patches = patch_height_width ** 2
         dim_in = channels * patch_size ** 2
 
         self.image_size = image_size
@@ -328,7 +324,7 @@ class ImageAutoregressiveFlow(Module):
         self.model = AutoregressiveFlow(
             **model,
             dim_input = dim_in,
-            max_seq_len = num_patches
+            max_seq_len = (patch_height_width, patch_height_width)
         )
 
         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)))

{autoregressive_diffusion_pytorch-0.2.4 → autoregressive_diffusion_pytorch-0.2.7}/pyproject.toml

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