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

@@ -647,12 +647,9 @@ class DiffusionPrior(BaseGaussianDiffusion):
         )
 
         if exists(clip):
-            if isinstance(clip, CLIP):
+            assert isinstance(clip, CLIP)
-                clip = XClipAdapter(clip)
-
-            assert isinstance(clip, BaseClipAdapter)
             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
@@ -743,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']
@@ -786,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
@@ -796,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
 
@@ -1251,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
@@ -1423,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
 
@@ -1491,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)
@@ -1547,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:

setup.py

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