add ability for DALL-E2 to return PIL images with return_pil_images = True on forward, for those who have no clue about deep learning

dalle2_pytorch/dalle2_pytorch.py

@@ -223,18 +223,7 @@ class BaseGaussianDiffusion(nn.Module):
 
         timesteps, = betas.shape
         self.num_timesteps = int(timesteps)
-
-        if loss_type == 'l1':
-            loss_fn = F.l1_loss
-        elif loss_type == 'l2':
-            loss_fn = F.mse_loss
-        elif loss_type == 'huber':
-            loss_fn = F.smooth_l1_loss
-        else:
-            raise NotImplementedError()
-
         self.loss_type = loss_type
-        self.loss_fn = loss_fn
 
         self.register_buffer('betas', betas)
         self.register_buffer('alphas_cumprod', alphas_cumprod)
@@ -658,12 +647,9 @@ class DiffusionPrior(BaseGaussianDiffusion):
         )
 
         if exists(clip):
-            if isinstance(clip, CLIP):
-                clip = XClipAdapter(clip)
-
-            assert isinstance(clip, BaseClipAdapter)
+            assert isinstance(clip, CLIP)
             freeze_model_and_make_eval_(clip)
-            self.clip = clip
+            self.clip = XClipAdapter(clip)
         else:
             assert exists(image_embed_dim), 'latent dimension must be given, if training prior network without CLIP given'
             self.clip = None
@@ -714,21 +700,29 @@ class DiffusionPrior(BaseGaussianDiffusion):
             img = self.p_sample(img, torch.full((b,), i, device = device, dtype = torch.long), text_cond = text_cond)
         return img
 
-    def p_losses(self, image_embed, times, text_cond, noise = None):
+    def p_losses(self, image_embed, t, text_cond, noise = None):
         noise = default(noise, lambda: torch.randn_like(image_embed))
 
-        image_embed_noisy = self.q_sample(x_start = image_embed, t = times, noise = noise)
+        image_embed_noisy = self.q_sample(x_start = image_embed, t = t, noise = noise)
 
-        pred = self.net(
+        x_recon = self.net(
             image_embed_noisy,
-            times,
+            t,
             cond_drop_prob = self.cond_drop_prob,
             **text_cond
         )
 
-        target = noise if not self.predict_x_start else image_embed
+        to_predict = noise if not self.predict_x_start else image_embed
+
+        if self.loss_type == 'l1':
+            loss = F.l1_loss(to_predict, x_recon)
+        elif self.loss_type == 'l2':
+            loss = F.mse_loss(to_predict, x_recon)
+        elif self.loss_type == "huber":
+            loss = F.smooth_l1_loss(to_predict, x_recon)
+        else:
+            raise NotImplementedError()
 
-        loss = self.loss_fn(pred, target)
         return loss
 
     @torch.no_grad()
@@ -746,7 +740,7 @@ class DiffusionPrior(BaseGaussianDiffusion):
         text_cond = dict(text_embed = text_embed)
 
         if self.condition_on_text_encodings:
-            text_cond = {**text_cond, 'text_encodings': text_encodings, 'mask': text != 0}
+            text_cond = {**text_cond, 'text_encodings': text_encodings, 'mask': text_mask}
 
         image_embeds = self.p_sample_loop((batch_size, image_embed_dim), text_cond = text_cond)
         text_embeds = text_cond['text_embed']
@@ -789,7 +783,6 @@ class DiffusionPrior(BaseGaussianDiffusion):
         text_cond = dict(text_embed = text_embed)
 
         if self.condition_on_text_encodings:
-            assert exists(text_encodings), 'text encodings must be present for diffusion prior if specified'
             text_cond = {**text_cond, 'text_encodings': text_encodings, 'mask': text_mask}
 
         # timestep conditioning from ddpm
@@ -799,7 +792,8 @@ class DiffusionPrior(BaseGaussianDiffusion):
 
         # calculate forward loss
 
-        return self.p_losses(image_embed, times, text_cond = text_cond, *args, **kwargs)
+        loss = self.p_losses(image_embed, times, text_cond = text_cond, *args, **kwargs)
+        return loss
 
 # decoder
 
@@ -1069,14 +1063,13 @@ class Unet(nn.Module):
         self,
         *,
         lowres_cond,
-        channels,
-        cond_on_image_embeds
+        channels
     ):
-        if lowres_cond == self.lowres_cond and channels == self.channels and cond_on_image_embeds == self.cond_on_image_embeds:
+        if lowres_cond == self.lowres_cond and channels == self.channels:
             return self
 
-        updated_kwargs = {'lowres_cond': lowres_cond, 'channels': channels, 'cond_on_image_embeds': cond_on_image_embeds}
-        return self.__class__(**{**self._locals, **updated_kwargs})
+        updated_kwargs = {**self._locals, 'lowres_cond': lowres_cond, 'channels': channels}
+        return self.__class__(**updated_kwargs)
 
     def forward_with_cond_scale(
         self,
@@ -1213,7 +1206,7 @@ class LowresConditioner(nn.Module):
         target_image_size = cast_tuple(target_image_size, 2)
 
         if self.training and self.downsample_first and exists(downsample_image_size):
-            cond_fmap = resize_image_to(cond_fmap, downsample_image_size, mode = self.cond_upsample_mode)
+            cond_fmap = resize_image_to(cond_fmap, target_image_size, mode = self.cond_upsample_mode)
 
         if self.training:
             # when training, blur the low resolution conditional image
@@ -1255,8 +1248,6 @@ class Decoder(BaseGaussianDiffusion):
             clip = XClipAdapter(clip)
 
         freeze_model_and_make_eval_(clip)
-        assert isinstance(clip, BaseClipAdapter)
-
         self.clip = clip
         self.clip_image_size = clip.image_size
         self.channels = clip.image_channels
@@ -1283,7 +1274,6 @@ class Decoder(BaseGaussianDiffusion):
 
             one_unet = one_unet.cast_model_parameters(
                 lowres_cond = not is_first,
-                cond_on_image_embeds = is_first,
                 channels = unet_channels
             )
 
@@ -1391,14 +1381,14 @@ class Decoder(BaseGaussianDiffusion):
 
         return img
 
-    def p_losses(self, unet, x_start, times, *, image_embed, lowres_cond_img = None, text_encodings = None, predict_x_start = False, noise = None):
+    def p_losses(self, unet, x_start, t, *, image_embed, lowres_cond_img = None, text_encodings = None, predict_x_start = False, noise = None):
         noise = default(noise, lambda: torch.randn_like(x_start))
 
-        x_noisy = self.q_sample(x_start = x_start, t = times, noise = noise)
+        x_noisy = self.q_sample(x_start = x_start, t = t, noise = noise)
 
-        pred = unet(
+        x_recon = unet(
             x_noisy,
-            times,
+            t,
             image_embed = image_embed,
             text_encodings = text_encodings,
             lowres_cond_img = lowres_cond_img,
@@ -1407,7 +1397,15 @@ class Decoder(BaseGaussianDiffusion):
 
         target = noise if not predict_x_start else x_start
 
-        loss = self.loss_fn(pred, target)
+        if self.loss_type == 'l1':
+            loss = F.l1_loss(target, x_recon)
+        elif self.loss_type == 'l2':
+            loss = F.mse_loss(target, x_recon)
+        elif self.loss_type == "huber":
+            loss = F.smooth_l1_loss(target, x_recon)
+        else:
+            raise NotImplementedError()
+
         return loss
 
     @torch.no_grad()
@@ -1420,7 +1418,6 @@ class Decoder(BaseGaussianDiffusion):
             _, text_encodings = self.clip.embed_text(text)
 
         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'
 
         img = None
 
@@ -1488,7 +1485,6 @@ class Decoder(BaseGaussianDiffusion):
             _, text_encodings = self.clip.embed_text(text)
 
         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'
 
         lowres_cond_img = self.to_lowres_cond(image, target_image_size = target_image_size, downsample_image_size = self.image_sizes[unet_index - 1]) if unet_number > 1 else None
         image = resize_image_to(image, target_image_size)
@@ -1544,6 +1540,8 @@ class DALLE2(nn.Module):
         images = self.decoder.sample(image_embed, text = text_cond, cond_scale = cond_scale)
 
         if return_pil_images:
+            # do some magic - if the user passed in a string text, or a list of strings
+            # assume they do not know anything about tensors and return PIL Image(s)
             images = list(map(self.to_pil, images.unbind(dim = 0)))
 
         if one_text:

dalle2_pytorch/vqgan_vae.py

@@ -545,7 +545,6 @@ class VQGanVAE(nn.Module):
         l2_recon_loss = False,
         use_hinge_loss = True,
         vgg = None,
-        vq_codebook_dim = 256,
         vq_codebook_size = 512,
         vq_decay = 0.8,
         vq_commitment_weight = 1.,
@@ -580,7 +579,6 @@ class VQGanVAE(nn.Module):
 
         self.vq = VQ(
             dim = self.enc_dec.encoded_dim,
-            codebook_dim = vq_codebook_dim,
             codebook_size = vq_codebook_size,
             decay = vq_decay,
             commitment_weight = vq_commitment_weight,

setup.py

@@ -10,7 +10,7 @@ setup(
       'dream = dalle2_pytorch.cli:dream'
     ],
   },
-  version = '0.0.65',
+  version = '0.0.59',
   license='MIT',
   description = 'DALL-E 2',
   author = 'Phil Wang',