@jacobwjs reports dynamic thresholding works very well and 0.95 is a better value

README.md

@@ -628,6 +628,82 @@ images = dalle2(
 
 Now you'll just have to worry about training the Prior and the Decoder!
 
+## Inpainting
+
+Inpainting is also built into the `Decoder`. You simply have to pass in the `inpaint_image` and `inpaint_mask` (boolean tensor where `True` indicates which regions of the inpaint image to keep)
+
+This repository uses the formulation put forth by <a href="https://arxiv.org/abs/2201.09865">Lugmayr et al. in Repaint</a>
+
+```python
+import torch
+from dalle2_pytorch import Unet, Decoder, CLIP
+
+# trained clip from step 1
+
+clip = CLIP(
+    dim_text = 512,
+    dim_image = 512,
+    dim_latent = 512,
+    num_text_tokens = 49408,
+    text_enc_depth = 6,
+    text_seq_len = 256,
+    text_heads = 8,
+    visual_enc_depth = 6,
+    visual_image_size = 256,
+    visual_patch_size = 32,
+    visual_heads = 8
+).cuda()
+
+# 2 unets for the decoder (a la cascading DDPM)
+
+unet = Unet(
+    dim = 16,
+    image_embed_dim = 512,
+    cond_dim = 128,
+    channels = 3,
+    dim_mults = (1, 1, 1, 1)
+).cuda()
+
+
+# decoder, which contains the unet(s) and clip
+
+decoder = Decoder(
+    clip = clip,
+    unet = (unet,),               # insert both unets in order of low resolution to highest resolution (you can have as many stages as you want here)
+    image_sizes = (256,),         # resolutions, 256 for first unet, 512 for second. these must be unique and in ascending order (matches with the unets passed in)
+    timesteps = 1000,
+    image_cond_drop_prob = 0.1,
+    text_cond_drop_prob = 0.5
+).cuda()
+
+# mock images (get a lot of this)
+
+images = torch.randn(4, 3, 256, 256).cuda()
+
+# feed images into decoder, specifying which unet you want to train
+# each unet can be trained separately, which is one of the benefits of the cascading DDPM scheme
+
+loss = decoder(images, unet_number = 1)
+loss.backward()
+
+# do the above for many steps for both unets
+
+mock_image_embed = torch.randn(1, 512).cuda()
+
+# then to do inpainting
+
+inpaint_image = torch.randn(1, 3, 256, 256).cuda()      # (batch, channels, height, width)
+inpaint_mask = torch.ones(1, 256, 256).bool().cuda()    # (batch, height, width)
+
+inpainted_images = decoder.sample(
+    image_embed = mock_image_embed,
+    inpaint_image = inpaint_image,    # just pass in the inpaint image
+    inpaint_mask = inpaint_mask       # and the mask
+)
+
+inpainted_images.shape # (1, 3, 256, 256)
+```
+
 ## Experimental
 
 ### DALL-E2 with Latent Diffusion
@@ -991,26 +1067,12 @@ dataset = ImageEmbeddingDataset(
 )
 ```
 
-### Scripts (wip)
+### Scripts
 
 #### `train_diffusion_prior.py`
 
 For detailed information on training the diffusion prior, please refer to the [dedicated readme](prior.md)
 
-## CLI (wip)
-
-```bash
-$ dream 'sharing a sunset at the summit of mount everest with my dog'
-```
-
-Once built, images will be saved to the same directory the command is invoked
-
-<a href="https://github.com/lucidrains/big-sleep">template</a>
-
-## Training CLI (wip)
-
-<a href="https://github.com/lucidrains/stylegan2-pytorch">template</a>
-
 ## Todo
 
 - [x] finish off gaussian diffusion class for latent embedding - allow for prediction of epsilon
@@ -1049,8 +1111,8 @@ Once built, images will be saved to the same directory the command is invoked
 - [x] test out grid attention in cascading ddpm locally, decide whether to keep or remove https://arxiv.org/abs/2204.01697 (keeping, seems to be fine)
 - [x] allow for unet to be able to condition non-cross attention style as well
 - [x] speed up inference, read up on papers (ddim)
-- [ ] add inpainting ability using resampler from repaint paper https://arxiv.org/abs/2201.09865
-- [ ] become an expert with unets, cleanup unet code, make it fully configurable, port all learnings over to https://github.com/lucidrains/x-unet (test out unet² in ddpm repo) - consider https://github.com/lucidrains/uformer-pytorch attention-based unet
+- [x] add inpainting ability using resampler from repaint paper https://arxiv.org/abs/2201.09865
+- [ ] try out the nested unet from https://arxiv.org/abs/2005.09007 after hearing several positive testimonies from researchers, for segmentation anyhow
 - [ ] interface out the vqgan-vae so a pretrained one can be pulled off the shelf to validate latent diffusion + DALL-E2
 
 ## Citations
@@ -1169,4 +1231,14 @@ Once built, images will be saved to the same directory the command is invoked
 }
 ```
 
+```bibtex
+@article{Lugmayr2022RePaintIU,
+    title   = {RePaint: Inpainting using Denoising Diffusion Probabilistic Models},
+    author  = {Andreas Lugmayr and Martin Danelljan and Andr{\'e}s Romero and Fisher Yu and Radu Timofte and Luc Van Gool},
+    journal = {ArXiv},
+    year    = {2022},
+    volume  = {abs/2201.09865}
+}
+```
+
 *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

@@ -522,7 +522,7 @@ class NoiseScheduler(nn.Module):
 
     def predict_noise_from_start(self, x_t, t, x0):
         return (
-            (x0 - extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t) / \
+            (extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - x0) / \
             extract(self.sqrt_recipm1_alphas_cumprod, t, x_t.shape)
         )
 
@@ -1542,10 +1542,10 @@ class Unet(nn.Module):
         self_attn = False,
         attn_dim_head = 32,
         attn_heads = 16,
-        lowres_cond = False, # for cascading diffusion - https://cascaded-diffusion.github.io/
-        lowres_noise_cond = False, # for conditioning on low resolution noising, based on Imagen
+        lowres_cond = False,             # for cascading diffusion - https://cascaded-diffusion.github.io/
+        lowres_noise_cond = False,       # for conditioning on low resolution noising, based on Imagen
         sparse_attn = False,
-        attend_at_middle = True, # whether to have a layer of attention at the bottleneck (can turn off for higher resolution in cascading DDPM, before bringing in efficient attention)
+        attend_at_middle = True,         # whether to have a layer of attention at the bottleneck (can turn off for higher resolution in cascading DDPM, before bringing in efficient attention)
         cond_on_text_encodings = False,
         max_text_len = 256,
         cond_on_image_embeds = False,
@@ -2100,7 +2100,7 @@ class Decoder(nn.Module):
         image_sizes = None,                         # for cascading ddpm, image size at each stage
         random_crop_sizes = None,                   # whether to random crop the image at that stage in the cascade (super resoluting convolutions at the end may be able to generalize on smaller crops)
         use_noise_for_lowres_cond = False,          # whether to use Imagen-like noising for low resolution conditioning  
-        use_blur_for_lowres_cond = True,                # whether to use the blur conditioning used in the original cascading ddpm paper, as well as DALL-E2
+        use_blur_for_lowres_cond = True,            # whether to use the blur conditioning used in the original cascading ddpm paper, as well as DALL-E2
         lowres_downsample_first = True,             # cascading ddpm - resizes to lower resolution, then to next conditional resolution + blur
         blur_prob = 0.5,                            # cascading ddpm - when training, the gaussian blur is only applied 50% of the time
         blur_sigma = 0.6,                           # cascading ddpm - blur sigma
@@ -2115,7 +2115,7 @@ class Decoder(nn.Module):
         unconditional = False,                      # set to True for generating images without conditioning
         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,
+        dynamic_thres_percentile = 0.95,
         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,
         ddim_sampling_eta = 1.                      # can be set to 0. for deterministic sampling afaict
@@ -2371,10 +2371,10 @@ class Decoder(nn.Module):
         x = x.clamp(-s, s) / s
         return x
 
-    def p_mean_variance(self, unet, x, t, image_embed, noise_scheduler, text_encodings = None, lowres_cond_img = None, clip_denoised = True, predict_x_start = False, learned_variance = False, cond_scale = 1., model_output = None):
+    def p_mean_variance(self, unet, x, t, image_embed, noise_scheduler, text_encodings = None, lowres_cond_img = None, clip_denoised = True, predict_x_start = False, learned_variance = False, cond_scale = 1., model_output = None, lowres_noise_level = None):
         assert not (cond_scale != 1. and not self.can_classifier_guidance), 'the decoder was not trained with conditional dropout, and thus one cannot use classifier free guidance (cond_scale anything other than 1)'
 
-        pred = default(model_output, lambda: unet.forward_with_cond_scale(x, t, image_embed = image_embed, text_encodings = text_encodings, cond_scale = cond_scale, lowres_cond_img = lowres_cond_img))
+        pred = default(model_output, lambda: unet.forward_with_cond_scale(x, t, image_embed = image_embed, text_encodings = text_encodings, cond_scale = cond_scale, lowres_cond_img = lowres_cond_img, lowres_noise_level = lowres_noise_level))
 
         if learned_variance:
             pred, var_interp_frac_unnormalized = pred.chunk(2, dim = 1)
@@ -2406,29 +2406,60 @@ class Decoder(nn.Module):
         return model_mean, posterior_variance, posterior_log_variance
 
     @torch.no_grad()
-    def p_sample(self, unet, x, t, image_embed, noise_scheduler, text_encodings = 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, noise_scheduler, text_encodings = None, cond_scale = 1., lowres_cond_img = None, predict_x_start = False, learned_variance = False, clip_denoised = True, lowres_noise_level = None):
         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, cond_scale = cond_scale, lowres_cond_img = lowres_cond_img, clip_denoised = clip_denoised, predict_x_start = predict_x_start, noise_scheduler = noise_scheduler, learned_variance = learned_variance)
+        model_mean, _, model_log_variance = self.p_mean_variance(unet, x = x, t = t, image_embed = image_embed, text_encodings = text_encodings, cond_scale = cond_scale, lowres_cond_img = lowres_cond_img, clip_denoised = clip_denoised, predict_x_start = predict_x_start, noise_scheduler = noise_scheduler, learned_variance = learned_variance, lowres_noise_level = lowres_noise_level)
         noise = torch.randn_like(x)
         # 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
 
     @torch.no_grad()
-    def p_sample_loop_ddpm(self, unet, shape, image_embed, noise_scheduler, predict_x_start = False, learned_variance = False, clip_denoised = True, lowres_cond_img = None, text_encodings = None, cond_scale = 1, is_latent_diffusion = False, lowres_noise_level = None):
+    def p_sample_loop_ddpm(
+        self,
+        unet,
+        shape,
+        image_embed,
+        noise_scheduler,
+        predict_x_start = False,
+        learned_variance = False,
+        clip_denoised = True,
+        lowres_cond_img = None,
+        text_encodings = None,
+        cond_scale = 1,
+        is_latent_diffusion = False,
+        lowres_noise_level = None,
+        inpaint_image = None,
+        inpaint_mask = None
+    ):
         device = self.device
 
         b = shape[0]
         img = torch.randn(shape, device = device)
 
+        if exists(inpaint_image):
+            inpaint_image = self.normalize_img(inpaint_image)
+            inpaint_image = resize_image_to(inpaint_image, shape[-1], nearest = True)
+            inpaint_mask = rearrange(inpaint_mask, 'b h w -> b 1 h w').float()
+            inpaint_mask = resize_image_to(inpaint_mask, shape[-1], nearest = True)
+            inpaint_mask = inpaint_mask.bool()
+
         if not is_latent_diffusion:
             lowres_cond_img = maybe(self.normalize_img)(lowres_cond_img)
 
         for i in tqdm(reversed(range(0, noise_scheduler.num_timesteps)), desc = 'sampling loop time step', total = noise_scheduler.num_timesteps):
+            times = torch.full((b,), i, device = device, dtype = torch.long)
+
+            if exists(inpaint_image):
+                # following the repaint paper
+                # https://arxiv.org/abs/2201.09865
+                noised_inpaint_image = noise_scheduler.q_sample(inpaint_image, t = times)
+                img = (img * ~inpaint_mask) + (noised_inpaint_image * inpaint_mask)
+
             img = self.p_sample(
                 unet,
                 img,
-                torch.full((b,), i, device = device, dtype = torch.long),
+                times,
                 image_embed = image_embed,
                 text_encodings = text_encodings,
                 cond_scale = cond_scale,
@@ -2440,11 +2471,32 @@ class Decoder(nn.Module):
                 clip_denoised = clip_denoised
             )
 
+        if exists(inpaint_image):
+            img = (img * ~inpaint_mask) + (inpaint_image * inpaint_mask)
+
         unnormalize_img = self.unnormalize_img(img)
         return unnormalize_img
 
     @torch.no_grad()
-    def p_sample_loop_ddim(self, unet, shape, image_embed, noise_scheduler, timesteps, eta = 1., predict_x_start = False, learned_variance = False, clip_denoised = True, lowres_cond_img = None, text_encodings = None, cond_scale = 1, is_latent_diffusion = False, lowres_noise_level = None):
+    def p_sample_loop_ddim(
+        self,
+        unet,
+        shape,
+        image_embed,
+        noise_scheduler,
+        timesteps,
+        eta = 1.,
+        predict_x_start = False,
+        learned_variance = False,
+        clip_denoised = True,
+        lowres_cond_img = None,
+        text_encodings = None,
+        cond_scale = 1,
+        is_latent_diffusion = False,
+        lowres_noise_level = None,
+        inpaint_image = None,
+        inpaint_mask = None
+    ):
         batch, device, total_timesteps, alphas, eta = shape[0], self.device, noise_scheduler.num_timesteps, noise_scheduler.alphas_cumprod_prev, self.ddim_sampling_eta
 
         times = torch.linspace(0., total_timesteps, steps = timesteps + 2)[:-1]
@@ -2452,6 +2504,13 @@ class Decoder(nn.Module):
         times = list(reversed(times.int().tolist()))
         time_pairs = list(zip(times[:-1], times[1:]))
 
+        if exists(inpaint_image):
+            inpaint_image = self.normalize_img(inpaint_image)
+            inpaint_image = resize_image_to(inpaint_image, shape[-1], nearest = True)
+            inpaint_mask = rearrange(inpaint_mask, 'b h w -> b 1 h w').float()
+            inpaint_mask = resize_image_to(inpaint_mask, shape[-1], nearest = True)
+            inpaint_mask = inpaint_mask.bool()
+
         img = torch.randn(shape, device = device)
 
         if not is_latent_diffusion:
@@ -2463,6 +2522,12 @@ class Decoder(nn.Module):
 
             time_cond = torch.full((batch,), time, device = device, dtype = torch.long)
 
+            if exists(inpaint_image):
+                # following the repaint paper
+                # https://arxiv.org/abs/2201.09865
+                noised_inpaint_image = noise_scheduler.q_sample(inpaint_image, t = time_cond)
+                img = (img * ~inpaint_mask) + (noised_inpaint_image * inpaint_mask)
+
             pred = unet.forward_with_cond_scale(img, time_cond, image_embed = image_embed, text_encodings = text_encodings, cond_scale = cond_scale, lowres_cond_img = lowres_cond_img, lowres_noise_level = lowres_noise_level)
 
             if learned_variance:
@@ -2486,6 +2551,9 @@ class Decoder(nn.Module):
                   c1 * noise + \
                   c2 * pred_noise
 
+        if exists(inpaint_image):
+            img = (img * ~inpaint_mask) + (inpaint_image * inpaint_mask)
+
         img = self.unnormalize_img(img)
         return img
 
@@ -2585,6 +2653,8 @@ class Decoder(nn.Module):
         cond_scale = 1.,
         stop_at_unet_number = None,
         distributed = False,
+        inpaint_image = None,
+        inpaint_mask = None
     ):
         assert self.unconditional or exists(image_embed), 'image embed must be present on sampling from decoder unless if trained unconditionally'
 
@@ -2598,6 +2668,8 @@ class Decoder(nn.Module):
         assert not (self.condition_on_text_encodings and not exists(text_encodings)), 'text or text encodings must be passed into decoder if specified'
         assert not (not self.condition_on_text_encodings and exists(text_encodings)), 'decoder specified not to be conditioned on text, yet it is presented'
 
+        assert not (exists(inpaint_image) ^ exists(inpaint_mask)), 'inpaint_image and inpaint_mask (boolean mask of [batch, height, width]) must be both given for inpainting'
+
         img = None
         is_cuda = next(self.parameters()).is_cuda
 
@@ -2609,6 +2681,8 @@ class Decoder(nn.Module):
             context = self.one_unet_in_gpu(unet = unet) if is_cuda and not distributed else null_context()
 
             with context:
+                # prepare low resolution conditioning for upsamplers
+
                 lowres_cond_img = lowres_noise_level = None
                 shape = (batch_size, channel, image_size, image_size)
 
@@ -2619,12 +2693,16 @@ class Decoder(nn.Module):
                         lowres_noise_level = torch.full((batch_size,), int(self.lowres_noise_sample_level * 1000), dtype = torch.long, device = self.device)
                         lowres_cond_img, _ = lowres_cond.noise_image(lowres_cond_img, lowres_noise_level)
 
+                # latent diffusion
+
                 is_latent_diffusion = isinstance(vae, VQGanVAE)
                 image_size = vae.get_encoded_fmap_size(image_size)
                 shape = (batch_size, vae.encoded_dim, image_size, image_size)
 
                 lowres_cond_img = maybe(vae.encode)(lowres_cond_img)
 
+                # denoising loop for image
+
                 img = self.p_sample_loop(
                     unet,
                     shape,
@@ -2638,7 +2716,9 @@ class Decoder(nn.Module):
                     lowres_noise_level = lowres_noise_level,
                     is_latent_diffusion = is_latent_diffusion,
                     noise_scheduler = noise_scheduler,
-                    timesteps = sample_timesteps
+                    timesteps = sample_timesteps,
+                    inpaint_image = inpaint_image,
+                    inpaint_mask = inpaint_mask
                 )
 
                 img = vae.decode(img)

dalle2_pytorch/version.py

@@ -1 +1 @@
-__version__ = '0.25.0'
+__version__ = '0.26.1'