fix bug with misnamed variable in diffusion prior network

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

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
@@ -1259,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:]))
@@ -1294,6 +1294,10 @@ class DiffusionPrior(nn.Module):
             if self.predict_x_start and self.sampling_clamp_l2norm:
                 x_start = self.l2norm_clamp_embed(x_start)
 
+            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.
@@ -2845,12 +2849,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

dalle2_pytorch/version.py

@@ -1 +1 @@
-__version__ = '1.10.1'
+__version__ = '1.10.4'