some cleanup

README.md

@@ -1207,14 +1207,4 @@ This library would not have gotten to this working state without the help of
 }
 ```
 
-```bibtex
-@article{Choi2022PerceptionPT,
-    title   = {Perception Prioritized Training of Diffusion Models},
-    author  = {Jooyoung Choi and Jungbeom Lee and Chaehun Shin and Sungwon Kim and Hyunwoo J. Kim and Sung-Hoon Yoon},
-    journal = {ArXiv},
-    year    = {2022},
-    volume  = {abs/2204.00227}
-}
-```
-
 *Creating noise from data is easy; creating data from noise is generative modeling.* - <a href="https://arxiv.org/abs/2011.13456">Yang Song's paper</a>

dalle2_pytorch/dalle2_pytorch.py

@@ -1,6 +1,7 @@
 import math
 import random
 from tqdm import tqdm
+from inspect import isfunction
 from functools import partial, wraps
 from contextlib import contextmanager
 from collections import namedtuple
@@ -11,7 +12,7 @@ import torch.nn.functional as F
 from torch import nn, einsum
 import torchvision.transforms as T
 
-from einops import rearrange, repeat, reduce
+from einops import rearrange, repeat
 from einops.layers.torch import Rearrange
 from einops_exts import rearrange_many, repeat_many, check_shape
 from einops_exts.torch import EinopsToAndFrom
@@ -56,7 +57,7 @@ def maybe(fn):
 def default(val, d):
     if exists(val):
         return val
-    return d() if callable(d) else d
+    return d() if isfunction(d) else d
 
 def cast_tuple(val, length = 1):
     if isinstance(val, list):
@@ -313,6 +314,11 @@ def extract(a, t, x_shape):
     out = a.gather(-1, t)
     return out.reshape(b, *((1,) * (len(x_shape) - 1)))
 
+def noise_like(shape, device, repeat=False):
+    repeat_noise = lambda: torch.randn((1, *shape[1:]), device=device).repeat(shape[0], *((1,) * (len(shape) - 1)))
+    noise = lambda: torch.randn(shape, device=device)
+    return repeat_noise() if repeat else noise()
+
 def meanflat(x):
     return x.mean(dim = tuple(range(1, len(x.shape))))
 
@@ -367,7 +373,7 @@ def quadratic_beta_schedule(timesteps):
     scale = 1000 / timesteps
     beta_start = scale * 0.0001
     beta_end = scale * 0.02
-    return torch.linspace(beta_start**0.5, beta_end**0.5, timesteps, dtype = torch.float64) ** 2
+    return torch.linspace(beta_start**2, beta_end**2, timesteps, dtype = torch.float64) ** 2
 
 
 def sigmoid_beta_schedule(timesteps):
@@ -379,7 +385,7 @@ def sigmoid_beta_schedule(timesteps):
 
 
 class BaseGaussianDiffusion(nn.Module):
-    def __init__(self, *, beta_schedule, timesteps, loss_type, p2_loss_weight_gamma = 0., p2_loss_weight_k = 1):
+    def __init__(self, *, beta_schedule, timesteps, loss_type):
         super().__init__()
 
         if beta_schedule == "cosine":
@@ -444,11 +450,6 @@ class BaseGaussianDiffusion(nn.Module):
         register_buffer('posterior_mean_coef1', betas * torch.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod))
         register_buffer('posterior_mean_coef2', (1. - alphas_cumprod_prev) * torch.sqrt(alphas) / (1. - alphas_cumprod))
 
-        # p2 loss reweighting
-
-        self.has_p2_loss_reweighting = p2_loss_weight_gamma > 0.
-        register_buffer('p2_loss_weight', (p2_loss_weight_k + alphas_cumprod / (1 - alphas_cumprod)) ** -p2_loss_weight_gamma)
-
     def q_posterior(self, x_start, x_t, t):
         posterior_mean = (
             extract(self.posterior_mean_coef1, t, x_t.shape) * x_start +
@@ -945,10 +946,10 @@ class DiffusionPrior(BaseGaussianDiffusion):
         return model_mean, posterior_variance, posterior_log_variance
 
     @torch.no_grad()
-    def p_sample(self, x, t, text_cond = None, clip_denoised = True, cond_scale = 1.):
+    def p_sample(self, x, t, text_cond = None, clip_denoised = True, repeat_noise = False, cond_scale = 1.):
         b, *_, device = *x.shape, x.device
         model_mean, _, model_log_variance = self.p_mean_variance(x = x, t = t, text_cond = text_cond, clip_denoised = clip_denoised, cond_scale = cond_scale)
-        noise = torch.randn_like(x)
+        noise = noise_like(x.shape, device, repeat_noise)
         # no noise when t == 0
         nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1)))
         return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise
@@ -1084,9 +1085,8 @@ class DiffusionPrior(BaseGaussianDiffusion):
 def Upsample(dim):
     return nn.ConvTranspose2d(dim, dim, 4, 2, 1)
 
-def Downsample(dim, *, dim_out = None):
-    dim_out = default(dim_out, dim)
-    return nn.Conv2d(dim, dim_out, 4, 2, 1)
+def Downsample(dim):
+    return nn.Conv2d(dim, dim, 4, 2, 1)
 
 class SinusoidalPosEmb(nn.Module):
     def __init__(self, dim):
@@ -1352,7 +1352,6 @@ class Unet(nn.Module):
         init_cross_embed_kernel_sizes = (3, 7, 15),
         cross_embed_downsample = False,
         cross_embed_downsample_kernel_sizes = (2, 4),
-        memory_efficient = False,
         **kwargs
     ):
         super().__init__()
@@ -1372,7 +1371,7 @@ class Unet(nn.Module):
         self.channels_out = default(channels_out, channels)
 
         init_channels = channels if not lowres_cond else channels * 2 # in cascading diffusion, one concats the low resolution image, blurred, for conditioning the higher resolution synthesis
-        init_dim = default(init_dim, dim)
+        init_dim = default(init_dim, dim // 3 * 2)
 
         self.init_conv = CrossEmbedLayer(init_channels, dim_out = init_dim, kernel_sizes = init_cross_embed_kernel_sizes, stride = 1)
 
@@ -1429,7 +1428,6 @@ class Unet(nn.Module):
         # for classifier free guidance
 
         self.null_image_embed = nn.Parameter(torch.randn(1, num_image_tokens, cond_dim))
-        self.null_image_hiddens = nn.Parameter(torch.randn(1, time_cond_dim))
 
         self.max_text_len = max_text_len
         self.null_text_embed = nn.Parameter(torch.randn(1, max_text_len, cond_dim))
@@ -1463,11 +1461,10 @@ class Unet(nn.Module):
             layer_cond_dim = cond_dim if not is_first else None
 
             self.downs.append(nn.ModuleList([
-                downsample_klass(dim_in, dim_out = dim_out) if memory_efficient else None,
-                ResnetBlock(dim_out if memory_efficient else dim_in, dim_out, time_cond_dim = time_cond_dim, groups = groups),
+                ResnetBlock(dim_in, dim_out, time_cond_dim = time_cond_dim, groups = groups),
                 Residual(LinearAttention(dim_out, **attn_kwargs)) if sparse_attn else nn.Identity(),
                 nn.ModuleList([ResnetBlock(dim_out, dim_out, cond_dim = layer_cond_dim, time_cond_dim = time_cond_dim, groups = groups) for _ in range(layer_num_resnet_blocks)]),
-                downsample_klass(dim_out) if not is_last and not memory_efficient else None
+                downsample_klass(dim_out) if not is_last else nn.Identity()
             ]))
 
         mid_dim = dims[-1]
@@ -1476,9 +1473,7 @@ class Unet(nn.Module):
         self.mid_attn = EinopsToAndFrom('b c h w', 'b (h w) c', Residual(Attention(mid_dim, **attn_kwargs))) if attend_at_middle else None
         self.mid_block2 = ResnetBlock(mid_dim, mid_dim, cond_dim = cond_dim, time_cond_dim = time_cond_dim, groups = resnet_groups[-1])
 
-        up_in_out_slice = slice(1 if not memory_efficient else None, None)
-
-        for ind, ((dim_in, dim_out), groups, layer_num_resnet_blocks) in enumerate(zip(reversed(in_out[up_in_out_slice]), reversed(resnet_groups), reversed(num_resnet_blocks))):
+        for ind, ((dim_in, dim_out), groups, layer_num_resnet_blocks) in enumerate(zip(reversed(in_out[1:]), reversed(resnet_groups), reversed(num_resnet_blocks))):
             is_last = ind >= (num_resolutions - 2)
             layer_cond_dim = cond_dim if not is_last else None
 
@@ -1489,10 +1484,8 @@ class Unet(nn.Module):
                 Upsample(dim_in)
             ]))
 
-        final_dim_in = dim * (1 if memory_efficient else 2)
-
         self.final_conv = nn.Sequential(
-            ResnetBlock(final_dim_in, dim, groups = resnet_groups[0]),
+            ResnetBlock(dim, dim, groups = resnet_groups[0]),
             nn.Conv2d(dim, self.channels_out, 1)
         )
 
@@ -1572,41 +1565,31 @@ class Unet(nn.Module):
         time_tokens = self.to_time_tokens(time_hiddens)
         t = self.to_time_cond(time_hiddens)
 
-        # conditional dropout
-
-        image_keep_mask = prob_mask_like((batch_size,), 1 - image_cond_drop_prob, device = device)
-        text_keep_mask = prob_mask_like((batch_size,), 1 - text_cond_drop_prob, device = device)
-
-        text_keep_mask = rearrange(text_keep_mask, 'b -> b 1 1')
-
         # image embedding to be summed to time embedding
         # discovered by @mhh0318 in the paper
 
         if exists(image_embed) and exists(self.to_image_hiddens):
             image_hiddens = self.to_image_hiddens(image_embed)
-            image_keep_mask_hidden = rearrange(image_keep_mask, 'b -> b 1')
-            null_image_hiddens = self.null_image_hiddens.to(image_hiddens.dtype)
-
-            image_hiddens = torch.where(
-                image_keep_mask_hidden,
-                image_hiddens,
-                null_image_hiddens
-            )
-
             t = t + image_hiddens
 
+        # conditional dropout
+
+        image_keep_mask = prob_mask_like((batch_size,), 1 - image_cond_drop_prob, device = device)
+        text_keep_mask = prob_mask_like((batch_size,), 1 - text_cond_drop_prob, device = device)
+
+        image_keep_mask, text_keep_mask = rearrange_many((image_keep_mask, text_keep_mask), 'b -> b 1 1')
+
         # mask out image embedding depending on condition dropout
         # for classifier free guidance
 
         image_tokens = None
 
         if self.cond_on_image_embeds:
-            image_keep_mask_embed = rearrange(image_keep_mask, 'b -> b 1 1')
             image_tokens = self.image_to_tokens(image_embed)
             null_image_embed = self.null_image_embed.to(image_tokens.dtype) # for some reason pytorch AMP not working
 
             image_tokens = torch.where(
-                image_keep_mask_embed,
+                image_keep_mask,
                 image_tokens,
                 null_image_embed
             )
@@ -1661,10 +1644,7 @@ class Unet(nn.Module):
 
         hiddens = []
 
-        for pre_downsample, init_block, sparse_attn, resnet_blocks, post_downsample in self.downs:
-            if exists(pre_downsample):
-                x = pre_downsample(x)
-
+        for init_block, sparse_attn, resnet_blocks, downsample in self.downs:
             x = init_block(x, c, t)
             x = sparse_attn(x)
 
@@ -1672,9 +1652,7 @@ class Unet(nn.Module):
                 x = resnet_block(x, c, t)
 
             hiddens.append(x)
-
-            if exists(post_downsample):
-                x = post_downsample(x)
+            x = downsample(x)
 
         x = self.mid_block1(x, mid_c, t)
 
@@ -1684,7 +1662,7 @@ class Unet(nn.Module):
         x = self.mid_block2(x, mid_c, t)
 
         for init_block, sparse_attn, resnet_blocks, upsample in self.ups:
-            x = torch.cat((x, hiddens.pop()), dim = 1)
+            x = torch.cat((x, hiddens.pop()), dim=1)
             x = init_block(x, c, t)
             x = sparse_attn(x)
 
@@ -1693,9 +1671,6 @@ class Unet(nn.Module):
 
             x = upsample(x)
 
-        if len(hiddens) > 0:
-            x = torch.cat((x, hiddens.pop()), dim = 1)
-
         return self.final_conv(x)
 
 class LowresConditioner(nn.Module):
@@ -1775,16 +1750,12 @@ class Decoder(BaseGaussianDiffusion):
         unconditional = False,
         auto_normalize_img = True,                  # whether to take care of normalizing the image from [0, 1] to [-1, 1] and back automatically - you can turn this off if you want to pass in the [-1, 1] ranged image yourself from the dataloader
         use_dynamic_thres = False,                  # from the Imagen paper
-        dynamic_thres_percentile = 0.9,
-        p2_loss_weight_gamma = 0.,                  # p2 loss weight, from https://arxiv.org/abs/2204.00227 - 0 is equivalent to weight of 1 across time - 1. is recommended
-        p2_loss_weight_k = 1
+        dynamic_thres_percentile = 0.9
     ):
         super().__init__(
             beta_schedule = beta_schedule,
             timesteps = timesteps,
-            loss_type = loss_type,
-            p2_loss_weight_gamma = p2_loss_weight_gamma,
-            p2_loss_weight_k = p2_loss_weight_k
+            loss_type = loss_type
         )
 
         self.unconditional = unconditional
@@ -1985,10 +1956,10 @@ class Decoder(BaseGaussianDiffusion):
         return model_mean, posterior_variance, posterior_log_variance
 
     @torch.no_grad()
-    def p_sample(self, unet, x, t, image_embed, text_encodings = None, text_mask = None, cond_scale = 1., lowres_cond_img = None, predict_x_start = False, learned_variance = False, clip_denoised = True):
+    def p_sample(self, unet, x, t, image_embed, text_encodings = None, text_mask = None, cond_scale = 1., lowres_cond_img = None, predict_x_start = False, learned_variance = False, clip_denoised = True, repeat_noise = False):
         b, *_, device = *x.shape, x.device
         model_mean, _, model_log_variance = self.p_mean_variance(unet, x = x, t = t, image_embed = image_embed, text_encodings = text_encodings, text_mask = text_mask, cond_scale = cond_scale, lowres_cond_img = lowres_cond_img, clip_denoised = clip_denoised, predict_x_start = predict_x_start, learned_variance = learned_variance)
-        noise = torch.randn_like(x)
+        noise = noise_like(x.shape, device, repeat_noise)
         # no noise when t == 0
         nonzero_mask = (1 - (t == 0).float()).reshape(b, *((1,) * (len(x.shape) - 1)))
         return model_mean + nonzero_mask * (0.5 * model_log_variance).exp() * noise
@@ -2052,13 +2023,7 @@ class Decoder(BaseGaussianDiffusion):
 
         target = noise if not predict_x_start else x_start
 
-        loss = self.loss_fn(pred, target, reduction = 'none')
-        loss = reduce(loss, 'b ... -> b (...)', 'mean')
-
-        if self.has_p2_loss_reweighting:
-            loss = loss * extract(self.p2_loss_weight, times, loss.shape)
-
-        loss = loss.mean()
+        loss = self.loss_fn(pred, target)
 
         if not learned_variance:
             # return simple loss if not using learned variance

dalle2_pytorch/trainer.py

@@ -238,7 +238,7 @@ class EMA(nn.Module):
             ma_buffer.data.copy_(current_buffer.data)
 
     def get_current_decay(self):
-        epoch = clamp(self.step.item() - self.update_after_step - 1, min_value = 0)
+        epoch = clamp(self.step.item() - self.update_after_step - 1, min = 0)
         value = 1 - (1 + epoch / self.inv_gamma) ** - self.power
 
         if epoch <= 0:

dalle2_pytorch/version.py

@@ -1 +1 @@
-__version__ = '0.9.0'
+__version__ = '0.6.14'

dalle2_pytorch/vqgan_vae.py

@@ -68,8 +68,8 @@ def group_dict_by_key(cond, d):
         return_val[ind][key] = d[key]
     return (*return_val,)
 
-def string_begins_with(prefix, string_input):
-    return string_input.startswith(prefix)
+def string_begins_with(prefix, str):
+    return str.startswith(prefix)
 
 def group_by_key_prefix(prefix, d):
     return group_dict_by_key(partial(string_begins_with, prefix), d)

train_decoder.py

@@ -211,7 +211,7 @@ def recall_trainer(tracker, trainer, recall_source=None, **load_config):
     Loads the model with an appropriate method depending on the tracker
     """
     print(print_ribbon(f"Loading model from {recall_source}"))
-    state_dict = tracker.recall_state_dict(recall_source, **load_config.dict())
+    state_dict = tracker.recall_state_dict(recall_source, **load_config)
     trainer.load_state_dict(state_dict["trainer"])
     print("Model loaded")
     return state_dict["epoch"], state_dict["step"], state_dict["validation_losses"]