in ddim, noise should be predicted after x0 is maybe clipped, thanks to @lukovnikov for pointing this out in another repository

README.md

@@ -1126,6 +1126,7 @@ For detailed information on training the diffusion prior, please refer to the [d
 - [x] add inpainting ability using resampler from repaint paper https://arxiv.org/abs/2201.09865
 - [x] add the final combination of upsample feature maps, used in unet squared, seems to have an effect in local experiments
 - [ ] consider elucidated dalle2 https://arxiv.org/abs/2206.00364
+- [ ] add simple outpainting, text-guided 2x size the image for starters
 - [ ] interface out the vqgan-vae so a pretrained one can be pulled off the shelf to validate latent diffusion + DALL-E2
 
 ## Citations
@@ -1285,4 +1286,14 @@ For detailed information on training the diffusion prior, please refer to the [d
 }
 ```
 
+```bibtex
+@article{Sunkara2022NoMS,
+    title   = {No More Strided Convolutions or Pooling: A New CNN Building Block for Low-Resolution Images and Small Objects},
+    author  = {Raja Sunkara and Tie Luo},
+    journal = {ArXiv},
+    year    = {2022},
+    volume  = {abs/2208.03641}
+}
+```
+
 *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

@@ -1070,7 +1070,7 @@ class DiffusionPriorNetwork(nn.Module):
 
         null_text_embeds = self.null_text_embeds.to(text_embed.dtype)
 
-        text_embeds = torch.where(
+        text_embed = torch.where(
             text_keep_mask,
             text_embed,
             null_text_embeds
@@ -1166,6 +1166,10 @@ class DiffusionPrior(nn.Module):
 
         self.net = net
         self.image_embed_dim = default(image_embed_dim, lambda: clip.dim_latent)
+
+        assert net.dim == self.image_embed_dim, f'your diffusion prior network has a dimension of {net.dim}, but you set your image embedding dimension (keyword image_embed_dim) on DiffusionPrior to {self.image_embed_dim}'
+        assert not exists(clip) or clip.dim_latent == self.image_embed_dim, f'you passed in a CLIP to the diffusion prior with latent dimensions of {clip.dim_latent}, but your image embedding dimension (keyword image_embed_dim) for the DiffusionPrior was set to {self.image_embed_dim}'
+
         self.channels = default(image_channels, lambda: clip.image_channels)
 
         self.text_cond_drop_prob = default(text_cond_drop_prob, cond_drop_prob)
@@ -1255,7 +1259,7 @@ class DiffusionPrior(nn.Module):
     def p_sample_loop_ddim(self, shape, text_cond, *, timesteps, eta = 1., cond_scale = 1.):
         batch, device, alphas, total_timesteps = shape[0], self.device, self.noise_scheduler.alphas_cumprod_prev, self.noise_scheduler.num_timesteps
 
-        times = torch.linspace(0., total_timesteps, steps = timesteps + 2)[:-1]
+        times = torch.linspace(-1., total_timesteps, steps = timesteps + 1)[:-1]
 
         times = list(reversed(times.int().tolist()))
         time_pairs = list(zip(times[:-1], times[1:]))
@@ -1277,12 +1281,14 @@ class DiffusionPrior(nn.Module):
 
             pred = self.net.forward_with_cond_scale(image_embed, time_cond, self_cond = self_cond, cond_scale = cond_scale, **text_cond)
 
+            # derive x0
+
             if self.predict_x_start:
                 x_start = pred
-                pred_noise = self.noise_scheduler.predict_noise_from_start(image_embed, t = time_cond, x0 = pred)
             else:
-                x_start = self.noise_scheduler.predict_start_from_noise(image_embed, t = time_cond, noise = pred)
-                pred_noise = pred
+                x_start = self.noise_scheduler.predict_start_from_noise(image_embed, t = time_cond, noise = pred_noise)
+
+            # clip x0 before maybe predicting noise
 
             if not self.predict_x_start:
                 x_start.clamp_(-1., 1.)
@@ -1290,6 +1296,17 @@ class DiffusionPrior(nn.Module):
             if self.predict_x_start and self.sampling_clamp_l2norm:
                 x_start = self.l2norm_clamp_embed(x_start)
 
+            # predict noise
+
+            if self.predict_x_start:
+                pred_noise = self.noise_scheduler.predict_noise_from_start(image_embed, t = time_cond, x0 = x_start)
+            else:
+                pred_noise = pred
+
+            if time_next < 0:
+                image_embed = x_start
+                continue
+
             c1 = eta * ((1 - alpha / alpha_next) * (1 - alpha_next) / (1 - alpha)).sqrt()
             c2 = ((1 - alpha_next) - torch.square(c1)).sqrt()
             noise = torch.randn_like(image_embed) if time_next > 0 else 0.
@@ -1479,9 +1496,14 @@ class PixelShuffleUpsample(nn.Module):
     def forward(self, x):
         return self.net(x)
 
-def Downsample(dim, *, dim_out = None):
+def Downsample(dim, dim_out = None):
+    # https://arxiv.org/abs/2208.03641 shows this is the most optimal way to downsample
+    # named SP-conv in the paper, but basically a pixel unshuffle
     dim_out = default(dim_out, dim)
-    return nn.Conv2d(dim, dim_out, 4, 2, 1)
+    return nn.Sequential(
+        Rearrange('b c (h s1) (w s2) -> b (c s1 s2) h w', s1 = 2, s2 = 2),
+        nn.Conv2d(dim * 4, dim_out, 1)
+    )
 
 class WeightStandardizedConv2d(nn.Conv2d):
     """
@@ -2836,12 +2858,13 @@ class Decoder(nn.Module):
         inpaint_mask = None,
         inpaint_resample_times = 5
     ):
-        batch, device, total_timesteps, alphas, eta = shape[0], self.device, noise_scheduler.num_timesteps, noise_scheduler.alphas_cumprod_prev, self.ddim_sampling_eta
+        batch, device, total_timesteps, alphas, eta = shape[0], self.device, noise_scheduler.num_timesteps, noise_scheduler.alphas_cumprod, self.ddim_sampling_eta
 
         times = torch.linspace(0., total_timesteps, steps = timesteps + 2)[:-1]
 
         times = list(reversed(times.int().tolist()))
         time_pairs = list(zip(times[:-1], times[1:]))
+        time_pairs = list(filter(lambda t: t[0] > t[1], time_pairs))
 
         is_inpaint = exists(inpaint_image)
         resample_times = inpaint_resample_times if is_inpaint else 1
@@ -2883,16 +2906,25 @@ class Decoder(nn.Module):
 
                 pred, _ = self.parse_unet_output(learned_variance, unet_output)
 
+                # predict x0
+
                 if predict_x_start:
                     x_start = pred
-                    pred_noise = noise_scheduler.predict_noise_from_start(img, t = time_cond, x0 = pred)
                 else:
                     x_start = noise_scheduler.predict_start_from_noise(img, t = time_cond, noise = pred)
-                    pred_noise = pred
+
+                # maybe clip x0
 
                 if clip_denoised:
                     x_start = self.dynamic_threshold(x_start)
 
+                # predict noise
+
+                if predict_x_start:
+                    pred_noise = noise_scheduler.predict_noise_from_start(img, t = time_cond, x0 = pred)
+                else:
+                    pred_noise = pred
+
                 c1 = eta * ((1 - alpha / alpha_next) * (1 - alpha_next) / (1 - alpha)).sqrt()
                 c2 = ((1 - alpha_next) - torch.square(c1)).sqrt()
                 noise = torch.randn_like(img) if not is_last_timestep else 0.

dalle2_pytorch/trainer.py

@@ -519,7 +519,7 @@ class DecoderTrainer(nn.Module):
             clip = decoder.clip
             clip.to(precision_type)
 
-        decoder, train_dataloader, *optimizers = list(self.accelerator.prepare(decoder, dataloaders['train'], *optimizers))
+        decoder, *optimizers = list(self.accelerator.prepare(decoder, *optimizers))
 
         self.decoder = decoder
 

dalle2_pytorch/version.py

@@ -1 +1 @@
-__version__ = '1.9.0'
+__version__ = '1.10.5'