does not seem right to clip for the prior diffusion part

README.md

@@ -383,9 +383,11 @@ You can also train the decoder on images of greater than the size (say 512x512)
 
 For the layperson, no worries, training will all be automated into a CLI tool, at least for small scale training.
 
-## Experimental - DALL-E2 with Latent Diffusion
+## Experimental
 
-This repository decides to take the next step and offer DALL-E2 combined with latent diffusion, from Rombach et al.
+### DALL-E2 with Latent Diffusion
+
+This repository decides to take the next step and offer DALL-E2 combined with <a href="https://huggingface.co/spaces/multimodalart/latentdiffusion">latent diffusion</a>, from Rombach et al.
 
 You can use it as follows. Latent diffusion can be limited to just the first U-Net in the cascade, or to any number you wish.
 
@@ -409,9 +411,10 @@ clip = CLIP(
     visual_heads = 8
 )
 
-# 2 unets for the decoder (a la cascading DDPM)
+# 3 unets for the decoder (a la cascading DDPM)
 
-# 1st unet is doing latent diffusion
+# first two unets are doing latent diffusion
+# vqgan-vae must be trained before hand
 
 vae1 = VQGanVAE(
     dim = 32,
@@ -469,7 +472,7 @@ decoder = Decoder(
 
 # mock images (get a lot of this)
 
-images = torch.randn(1, 3, 512, 512).cuda()
+images = torch.randn(1, 3, 1024, 1024).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
@@ -482,6 +485,10 @@ with decoder.one_unet_in_gpu(2):
     loss = decoder(images, unet_number = 2)
     loss.backward()
 
+with decoder.one_unet_in_gpu(3):
+    loss = decoder(images, unet_number = 3)
+    loss.backward()
+
 # do the above for many steps for both unets
 
 # then it will learn to generate images based on the CLIP image embeddings
@@ -492,7 +499,11 @@ mock_image_embed = torch.randn(1, 512).cuda()
 images = decoder.sample(mock_image_embed) # (1, 3, 1024, 1024)
 ```
 
-## CLI Usage (work in progress)
+## Training wrapper (wip)
+
+Offer training wrappers
+
+## CLI (wip)
 
 ```bash
 $ dream 'sharing a sunset at the summit of mount everest with my dog'
@@ -500,9 +511,7 @@ $ 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
 
-## Training wrapper (wip)
-
-Offer training wrappers
+<a href="https://github.com/lucidrains/big-sleep">template</a>
 
 ## Training CLI (wip)
 
@@ -525,6 +534,9 @@ Offer training wrappers
 - [ ] copy the cascading ddpm code to a separate repo (perhaps https://github.com/lucidrains/denoising-diffusion-pytorch) as the main contribution of dalle2 really is just the prior network
 - [ ] train on a toy task, offer in colab
 - [ ] extend diffusion head to use diffusion-gan (potentially using lightweight-gan) to speed up inference
+- [ ] bring in tools to train vqgan-vae
+- [ ] for decoder, allow ability to customize objective (predict epsilon vs x0), in case latent diffusion does better with prediction of x0
+- [ ] bring in vit-vqgan https://arxiv.org/abs/2110.04627 for the latent diffusion
 
 ## Citations
 
@@ -562,17 +574,6 @@ Offer training wrappers
 }
 ```
 
-```bibtex
-@misc{zhang2019root,
-    title   = {Root Mean Square Layer Normalization},
-    author  = {Biao Zhang and Rico Sennrich},
-    year    = {2019},
-    eprint  = {1910.07467},
-    archivePrefix = {arXiv},
-    primaryClass = {cs.LG}
-}
-```
-
 ```bibtex
 @inproceedings{Tu2022MaxViTMV,
     title   = {MaxViT: Multi-Axis Vision Transformer},

dalle2_pytorch/cli.py

@@ -1,9 +1,51 @@
 import click
+import torch
+import torchvision.transforms as T
+from pathlib import Path
+
+from dalle2_pytorch import DALLE2, Decoder, DiffusionPrior
+
+def safeget(dictionary, keys, default = None):
+    return reduce(lambda d, key: d.get(key, default) if isinstance(d, dict) else default, keys.split('.'), dictionary)
+
+def simple_slugify(text, max_length = 255):
+    return text.replace("-", "_").replace(",", "").replace(" ", "_").replace("|", "--").strip('-_')[:max_length]
+
+def get_pkg_version():
+    from pkg_resources import get_distribution
+    return get_distribution('dalle2_pytorch').version
 
 def main():
     pass
 
 @click.command()
+@click.option('--model', default = './dalle2.pt', help = 'path to trained DALL-E2 model')
+@click.option('--cond_scale', default = 2, help = 'conditioning scale (classifier free guidance) in decoder')
 @click.argument('text')
-def dream(text):
-    return 'not ready yet'
+def dream(
+    model,
+    cond_scale,
+    text
+):
+    model_path = Path(model)
+    full_model_path = str(model_path.resolve())
+    assert model_path.exists(), f'model not found at {full_model_path}'
+    loaded = torch.load(str(model_path))
+
+    version = safeget(loaded, 'version')
+    print(f'loading DALL-E2 from {full_model_path}, saved at version {version} - current package version is {get_pkg_version()}')
+
+    prior_init_params = safeget(loaded, 'init_params.prior')
+    decoder_init_params = safeget(loaded, 'init_params.decoder')
+    model_params = safeget(loaded, 'model_params')
+
+    prior = DiffusionPrior(**prior_init_params)
+    decoder = Decoder(**decoder_init_params)
+
+    dalle2 = DALLE2(prior, decoder)
+    dalle2.load_state_dict(model_params)
+
+    image = dalle2(text, cond_scale = cond_scale)
+
+    pil_image = T.ToPILImage()(image)
+    return pil_image.save(f'./{simple_slugify(text)}.png')

dalle2_pytorch/dalle2_pytorch.py

@@ -16,6 +16,7 @@ from einops_exts.torch import EinopsToAndFrom
 from kornia.filters import gaussian_blur2d
 
 from dalle2_pytorch.tokenizer import tokenizer
+from dalle2_pytorch.vqgan_vae import NullVQGanVAE, VQGanVAE
 
 # use x-clip
 
@@ -48,11 +49,11 @@ def is_list_str(x):
         return False
     return all([type(el) == str for el in x])
 
-def pad_tuple_to_length(t, length):
+def pad_tuple_to_length(t, length, fillvalue = None):
     remain_length = length - len(t)
     if remain_length <= 0:
         return t
-    return (*t, *((None,) * remain_length))
+    return (*t, *((fillvalue,) * remain_length))
 
 # for controlling freezing of CLIP
 
@@ -590,7 +591,7 @@ class DiffusionPrior(nn.Module):
         else:
             x_recon = self.predict_start_from_noise(x, t = t, noise = self.net(x, t, **text_cond))
 
-        if clip_denoised:
+        if clip_denoised and not self.predict_x0:
             x_recon.clamp_(-1., 1.)
 
         model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t)
@@ -1135,12 +1136,15 @@ class Decoder(nn.Module):
         # while the rest of the unets are conditioned on the low resolution image produced by previous unet
 
         unets = cast_tuple(unet)
-        vaes = pad_tuple_to_length(cast_tuple(vae), len(unets))
+        vaes = pad_tuple_to_length(cast_tuple(vae), len(unets), fillvalue = NullVQGanVAE(channels = self.channels))
 
         self.unets = nn.ModuleList([])
         self.vaes = nn.ModuleList([])
 
         for ind, (one_unet, one_vae) in enumerate(zip(unets, vaes)):
+            assert isinstance(one_unet, Unet)
+            assert isinstance(one_vae, (VQGanVAE, NullVQGanVAE))
+
             is_first = ind == 0
             latent_dim = one_vae.encoded_dim if exists(one_vae) else None
 
@@ -1152,7 +1156,7 @@ class Decoder(nn.Module):
             )
 
             self.unets.append(one_unet)
-            self.vaes.append(one_vae.copy_for_eval() if exists(one_vae) else None)
+            self.vaes.append(one_vae.copy_for_eval())
 
         # unet image sizes
 
@@ -1362,12 +1366,11 @@ class Decoder(nn.Module):
                 if unet.lowres_cond:
                     lowres_cond_img = self.to_lowres_cond(img, target_image_size = image_size)
 
-                if exists(vae):
-                    image_size //= (2 ** vae.layers)
-                    shape = (batch_size, vae.encoded_dim, image_size, image_size)
+                image_size = vae.get_encoded_fmap_size(image_size)
+                shape = (batch_size, vae.encoded_dim, image_size, image_size)
 
-                    if exists(lowres_cond_img):
-                        lowres_cond_img = vae.encode(lowres_cond_img)
+                if exists(lowres_cond_img):
+                    lowres_cond_img = vae.encode(lowres_cond_img)
 
                 img = self.p_sample_loop(
                     unet,
@@ -1378,8 +1381,7 @@ class Decoder(nn.Module):
                     lowres_cond_img = lowres_cond_img
                 )
 
-                if exists(vae):
-                    img = vae.decode(img)
+                img = vae.decode(img)
 
         return img
 
@@ -1415,13 +1417,12 @@ class Decoder(nn.Module):
         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)
 
-        if exists(vae):
-            vae.eval()
-            with torch.no_grad():
-                image = vae.encode(image)
+        vae.eval()
+        with torch.no_grad():
+            image = vae.encode(image)
 
-                if exists(lowres_cond_img):
-                    lowres_cond_img = vae.encode(lowres_cond_img)
+            if exists(lowres_cond_img):
+                lowres_cond_img = vae.encode(lowres_cond_img)
 
         return self.p_losses(unet, image, times, image_embed = image_embed, text_encodings = text_encodings, lowres_cond_img = lowres_cond_img)
 
@@ -1450,6 +1451,7 @@ class DALLE2(nn.Module):
         cond_scale = 1.
     ):
         device = next(self.parameters()).device
+        one_text = isinstance(text, str) or (not is_list_str(text) and text.shape[0] == 1)
 
         if isinstance(text, str) or is_list_str(text):
             text = [text] if not isinstance(text, (list, tuple)) else text
@@ -1457,4 +1459,8 @@ class DALLE2(nn.Module):
 
         image_embed = self.prior.sample(text, num_samples_per_batch = self.prior_num_samples)
         images = self.decoder.sample(image_embed, cond_scale = cond_scale)
+
+        if one_text:
+            return images[0]
+
         return images

dalle2_pytorch/vqgan_vae.py

@@ -287,6 +287,28 @@ class VQGanAttention(nn.Module):
 
         return out + residual
 
+class NullVQGanVAE(nn.Module):
+    def __init__(
+        self,
+        *,
+        channels
+    ):
+        super().__init__()
+        self.encoded_dim = channels
+        self.layers = 0
+
+    def get_encoded_fmap_size(self, size):
+        return size
+
+    def copy_for_eval(self):
+        return self
+
+    def encode(self, x):
+        return x
+
+    def decode(self, x):
+        return x
+
 class VQGanVAE(nn.Module):
     def __init__(
         self,
@@ -407,6 +429,9 @@ class VQGanVAE(nn.Module):
         self.discr_loss = hinge_discr_loss if use_hinge_loss else bce_discr_loss
         self.gen_loss = hinge_gen_loss if use_hinge_loss else bce_gen_loss
 
+    def get_encoded_fmap_size(self, image_size):
+        return image_size // (2 ** self.layers)
+
     def copy_for_eval(self):
         device = next(self.parameters()).device
         vae_copy = copy.deepcopy(self.cpu())
@@ -452,7 +477,8 @@ class VQGanVAE(nn.Module):
         img,
         return_loss = False,
         return_discr_loss = False,
-        return_recons = False
+        return_recons = False,
+        add_gradient_penalty = True
     ):
         batch, channels, height, width, device = *img.shape, img.device
         assert height == self.image_size and width == self.image_size, 'height and width of input image must be equal to {self.image_size}'
@@ -477,11 +503,11 @@ class VQGanVAE(nn.Module):
 
             fmap_discr_logits, img_discr_logits = map(self.discr, (fmap, img))
 
-            gp = gradient_penalty(img, img_discr_logits)
-
             discr_loss = self.discr_loss(fmap_discr_logits, img_discr_logits)
 
-            loss = discr_loss + gp
+            if add_gradient_penalty:
+                gp = gradient_penalty(img, img_discr_logits)
+                loss = discr_loss + gp
 
             if return_recons:
                 return loss, fmap

setup.py

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