make sure non-latent diffusion still works

README.md

@@ -10,11 +10,9 @@ The main novelty seems to be an extra layer of indirection with the prior networ
 
 This model is SOTA for text-to-image for now.
 
-It may also explore an extension of using <a href="https://huggingface.co/spaces/multimodalart/latentdiffusion">latent diffusion</a> in the decoder from Rombach et al.
-
 Please join <a href="https://discord.gg/xBPBXfcFHd"><img alt="Join us on Discord" src="https://img.shields.io/discord/823813159592001537?color=5865F2&logo=discord&logoColor=white"></a> if you are interested in helping out with the replication
 
-There was enough interest for a Jax version. It will be completed after the Pytorch version shows signs of life on my toy tasks. <a href="https://github.com/lucidrains/dalle2-jax">Placeholder repository</a>. I will also eventually extend this to <a href="https://github.com/lucidrains/dalle2-video">text to video</a>, once the repository is in a good place.
+There was enough interest for a <a href="https://github.com/lucidrains/dalle2-jax">Jax version</a>. I will also eventually extend this to <a href="https://github.com/lucidrains/dalle2-video">text to video</a>, once the repository is in a good place.
 
 ## Install
 
@@ -248,13 +246,6 @@ loss = decoder(images, unet_number = 2)
 loss.backward()
 
 # do the above for many steps for both unets
-
-# then it will learn to generate images based on the CLIP image embeddings
-
-# chaining the unets from lowest resolution to highest resolution (thus cascading)
-
-mock_image_embed = torch.randn(1, 512).cuda()
-images = decoder.sample(mock_image_embed) # (1, 3, 512, 512)
 ```
 
 Finally, to generate the DALL-E2 images from text. Insert the trained `DiffusionPrior` as well as the `Decoder` (which wraps `CLIP`, the causal transformer, and unet(s))
@@ -385,7 +376,127 @@ 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.
 
-## CLI Usage (work in progress)
+## Experimental
+
+### 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.
+
+```python
+import torch
+from dalle2_pytorch import Unet, Decoder, CLIP, VQGanVAE
+
+# trained clip from step 1
+
+clip = CLIP(
+    dim_text = 512,
+    dim_image = 512,
+    dim_latent = 512,
+    num_text_tokens = 49408,
+    text_enc_depth = 1,
+    text_seq_len = 256,
+    text_heads = 8,
+    visual_enc_depth = 1,
+    visual_image_size = 256,
+    visual_patch_size = 32,
+    visual_heads = 8
+)
+
+# 3 unets for the decoder (a la cascading DDPM)
+
+# first two unets are doing latent diffusion
+# vqgan-vae must be trained before hand
+
+vae1 = VQGanVAE(
+    dim = 32,
+    image_size = 256,
+    layers = 3,
+    layer_mults = (1, 2, 4)
+)
+
+vae2 = VQGanVAE(
+    dim = 32,
+    image_size = 512,
+    layers = 3,
+    layer_mults = (1, 2, 4)
+)
+
+unet1 = Unet(
+    dim = 32,
+    image_embed_dim = 512,
+    cond_dim = 128,
+    channels = 3,
+    sparse_attn = True,
+    sparse_attn_window = 2,
+    dim_mults = (1, 2, 4, 8)
+)
+
+unet2 = Unet(
+    dim = 32,
+    image_embed_dim = 512,
+    channels = 3,
+    dim_mults = (1, 2, 4, 8, 16),
+    cond_on_image_embeds = True,
+    cond_on_text_encodings = False
+)
+
+unet3 = Unet(
+    dim = 32,
+    image_embed_dim = 512,
+    channels = 3,
+    dim_mults = (1, 2, 4, 8, 16),
+    cond_on_image_embeds = True,
+    cond_on_text_encodings = False,
+    attend_at_middle = False
+)
+
+# decoder, which contains the unet(s) and clip
+
+decoder = Decoder(
+    clip = clip,
+    vae = (vae1, vae2),                # latent diffusion for unet1 (vae1) and unet2 (vae2), but not for the last unet3
+    unet = (unet1, unet2, unet3),      # insert unets in order of low resolution to highest resolution (you can have as many stages as you want here)
+    image_sizes = (256, 512, 1024),    # resolutions, 256 for first unet, 512 for second, 1024 for third
+    timesteps = 100,
+    cond_drop_prob = 0.2
+).cuda()
+
+# mock images (get a lot of this)
+
+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
+
+with decoder.one_unet_in_gpu(1):
+    loss = decoder(images, unet_number = 1)
+    loss.backward()
+
+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
+
+# chaining the unets from lowest resolution to highest resolution (thus cascading)
+
+mock_image_embed = torch.randn(1, 512).cuda()
+images = decoder.sample(mock_image_embed) # (1, 3, 1024, 1024)
+```
+
+## Training wrapper (wip)
+
+Offer training wrappers
+
+## CLI (wip)
 
 ```bash
 $ dream 'sharing a sunset at the summit of mount everest with my dog'
@@ -393,9 +504,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)
 
@@ -412,11 +521,17 @@ Offer training wrappers
 - [x] add efficient attention in unet
 - [x] be able to finely customize what to condition on (text, image embed) for specific unet in the cascade (super resolution ddpms near the end may not need too much conditioning)
 - [x] offload unets not being trained on to CPU for memory efficiency (for training each resolution unets separately)
-- [ ] build out latent diffusion architecture, with the vq-reg variant (vqgan-vae), make it completely optional
+- [x] build out latent diffusion architecture, with the vq-reg variant (vqgan-vae), make it completely optional and compatible with cascading ddpms
+- [x] for decoder, allow ability to customize objective (predict epsilon vs x0), in case latent diffusion does better with prediction of x0
+- [x] use attention-based upsampling https://arxiv.org/abs/2112.11435
+- [ ] spend one day cleaning up tech debt in decoder
 - [ ] become an expert with unets, cleanup unet code, make it fully configurable, port all learnings over to https://github.com/lucidrains/x-unet
 - [ ] 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
+- [ ] transcribe code to Jax, which lowers the activation energy for distributed training, given access to TPUs
 - [ ] 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
+- [ ] bring in vit-vqgan https://arxiv.org/abs/2110.04627 for the latent diffusion
 
 ## Citations
 
@@ -448,23 +563,12 @@ Offer training wrappers
 
 ```bibtex
 @inproceedings{Liu2022ACF,
-    title   = {A ConvNet for the 2020s},
+    title   = {A ConvNet for the 2020https://arxiv.org/abs/2112.11435s},
     author  = {Zhuang Liu and Hanzi Mao and Chaozheng Wu and Christoph Feichtenhofer and Trevor Darrell and Saining Xie},
     year    = {2022}
 }
 ```
 
-```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/attention.py

@@ -0,0 +1,125 @@
+import torch
+from torch import nn, einsum
+import torch.nn.functional as F
+
+from einops import rearrange, repeat
+
+class LayerNormChan(nn.Module):
+    def __init__(
+        self,
+        dim,
+        eps = 1e-5
+    ):
+        super().__init__()
+        self.eps = eps
+        self.gamma = nn.Parameter(torch.ones(1, dim, 1, 1))
+
+    def forward(self, x):
+        var = torch.var(x, dim = 1, unbiased = False, keepdim = True)
+        mean = torch.mean(x, dim = 1, keepdim = True)
+        return (x - mean) / (var + self.eps).sqrt() * self.gamma
+
+# attention-based upsampling
+# from https://arxiv.org/abs/2112.11435
+
+class QueryAndAttend(nn.Module):
+    def __init__(
+        self,
+        *,
+        dim,
+        num_queries = 1,
+        dim_head = 32,
+        heads = 8,
+        window_size = 3
+    ):
+        super().__init__()
+        self.scale = dim_head ** -0.5
+        inner_dim = dim_head * heads
+        self.heads = heads
+        self.dim_head = dim_head
+        self.window_size = window_size
+        self.num_queries = num_queries
+
+        self.rel_pos_bias = nn.Parameter(torch.randn(heads, num_queries, window_size * window_size, 1, 1))
+
+        self.queries = nn.Parameter(torch.randn(heads, num_queries, dim_head))
+        self.to_kv = nn.Conv2d(dim, dim_head * 2, 1, bias = False)
+        self.to_out = nn.Conv2d(inner_dim, dim, 1, bias = False)
+
+    def forward(self, x):
+        """
+        einstein notation
+        b - batch
+        h - heads
+        l - num queries
+        d - head dimension
+        x - height
+        y - width
+        j - source sequence for attending to (kernel size squared in this case)
+        """
+
+        wsz, heads, dim_head, num_queries = self.window_size, self.heads, self.dim_head, self.num_queries
+        batch, _, height, width = x.shape
+
+        is_one_query = self.num_queries == 1
+
+        # queries, keys, values
+
+        q = self.queries * self.scale
+        k, v = self.to_kv(x).chunk(2, dim = 1)
+
+        # similarities
+
+        sim = einsum('h l d, b d x y -> b h l x y', q, k)
+        sim = rearrange(sim, 'b ... x y -> b (...) x y')
+
+        # unfold the similarity scores, with float(-inf) as padding value
+
+        mask_value = -torch.finfo(sim.dtype).max
+        sim = F.pad(sim, ((wsz // 2,) * 4), value = mask_value)
+        sim = F.unfold(sim, kernel_size = wsz)
+        sim = rearrange(sim, 'b (h l j) (x y) -> b h l j x y', h = heads, l = num_queries, x = height, y = width)
+
+        # rel pos bias
+
+        sim = sim + self.rel_pos_bias
+
+        # numerically stable attention
+
+        sim = sim - sim.amax(dim = -3, keepdim = True).detach()
+        attn = sim.softmax(dim = -3)
+
+        # unfold values
+
+        v = F.pad(v, ((wsz // 2,) * 4), value = 0.)
+        v = F.unfold(v, kernel_size = wsz)
+        v = rearrange(v, 'b (d j) (x y) -> b d j x y', d = dim_head, x = height, y = width)
+
+        # aggregate values
+
+        out = einsum('b h l j x y, b d j x y -> b l h d x y', attn, v)
+
+        # combine heads
+
+        out = rearrange(out, 'b l h d x y -> (b l) (h d) x y')
+        out = self.to_out(out)
+        out = rearrange(out, '(b l) d x y -> b l d x y', b = batch)
+
+        # return original input if one query
+
+        if is_one_query:
+            out = rearrange(out, 'b 1 ... -> b ...')
+
+        return out
+
+class QueryAttnUpsample(nn.Module):
+    def __init__(self, dim, **kwargs):
+        super().__init__()
+        self.norm = LayerNormChan(dim)
+        self.qna = QueryAndAttend(dim = dim, num_queries = 4, **kwargs)
+
+    def forward(self, x):
+        x = self.norm(x)
+        out = self.qna(x)
+        out = rearrange(out, 'b (w1 w2) c h w -> b c (h w1) (w w2)', w1 = 2, w2 = 2)
+        return out

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,8 @@ 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
+from dalle2_pytorch.attention import QueryAttnUpsample
 
 # use x-clip
 
@@ -48,6 +50,12 @@ def is_list_str(x):
         return False
     return all([type(el) == str for el in x])
 
+def pad_tuple_to_length(t, length, fillvalue = None):
+    remain_length = length - len(t)
+    if remain_length <= 0:
+        return t
+    return (*t, *((fillvalue,) * remain_length))
+
 # for controlling freezing of CLIP
 
 def set_module_requires_grad_(module, requires_grad):
@@ -476,7 +484,7 @@ class DiffusionPrior(nn.Module):
         timesteps = 1000,
         cond_drop_prob = 0.2,
         loss_type = "l1",
-        predict_x0 = True,
+        predict_x_start = True,
         beta_schedule = "cosine",
     ):
         super().__init__()
@@ -490,7 +498,7 @@ class DiffusionPrior(nn.Module):
         self.image_size = clip.image_size
         self.cond_drop_prob = cond_drop_prob
 
-        self.predict_x0 = predict_x0
+        self.predict_x_start = predict_x_start
         # in paper, they do not predict the noise, but predict x0 directly for image embedding, claiming empirically better results. I'll just offer both.
 
         if beta_schedule == "cosine":
@@ -540,12 +548,14 @@ class DiffusionPrior(nn.Module):
         self.register_buffer('posterior_mean_coef1', betas * torch.sqrt(alphas_cumprod_prev) / (1. - alphas_cumprod))
         self.register_buffer('posterior_mean_coef2', (1. - alphas_cumprod_prev) * torch.sqrt(alphas) / (1. - alphas_cumprod))
 
+    @torch.no_grad()
     def get_image_embed(self, image):
         image_encoding = self.clip.visual_transformer(image)
         image_cls = image_encoding[:, 0]
         image_embed = self.clip.to_visual_latent(image_cls)
         return l2norm(image_embed)
 
+    @torch.no_grad()
     def get_text_cond(self, text):
         text_encodings = self.clip.text_transformer(text)
         text_cls, text_encodings = text_encodings[:, 0], text_encodings[:, 1:]
@@ -575,14 +585,16 @@ class DiffusionPrior(nn.Module):
         return posterior_mean, posterior_variance, posterior_log_variance_clipped
 
     def p_mean_variance(self, x, t, text_cond, clip_denoised: bool):
-        if self.predict_x0:
-            x_recon = self.net(x, t, **text_cond)
+        pred = self.net(x, t, **text_cond)
+
+        if self.predict_x_start:
+            x_recon = pred
             # not 100% sure of this above line - for any spectators, let me know in the github issues (or through a pull request) if you know how to correctly do this
             # i'll be rereading https://arxiv.org/abs/2111.14822, where i think a similar approach is taken
         else:
-            x_recon = self.predict_start_from_noise(x, t = t, noise = self.net(x, t, **text_cond))
+            x_recon = self.predict_start_from_noise(x, t = t, noise = pred)
 
-        if clip_denoised:
+        if clip_denoised and not self.predict_x_start:
             x_recon.clamp_(-1., 1.)
 
         model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t)
@@ -628,7 +640,7 @@ class DiffusionPrior(nn.Module):
             **text_cond
         )
 
-        to_predict = noise if not self.predict_x0 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)
@@ -940,11 +952,16 @@ class Unet(nn.Module):
 
     # if the current settings for the unet are not correct
     # for cascading DDPM, then reinit the unet with the right settings
-    def force_lowres_cond(self, lowres_cond):
-        if lowres_cond == self.lowres_cond:
+    def cast_model_parameters(
+        self,
+        *,
+        lowres_cond,
+        channels
+    ):
+        if lowres_cond == self.lowres_cond and channels == self.channels:
             return self
 
-        updated_kwargs = {**self._locals, 'lowres_cond': lowres_cond}
+        updated_kwargs = {**self._locals, 'lowres_cond': lowres_cond, 'channels': channels}
         return self.__class__(**updated_kwargs)
 
     def forward_with_cond_scale(
@@ -1100,10 +1117,13 @@ class Decoder(nn.Module):
         unet,
         *,
         clip,
+        vae = tuple(),
         timesteps = 1000,
         cond_drop_prob = 0.2,
         loss_type = 'l1',
         beta_schedule = 'cosine',
+        predict_x_start = False,
+        predict_x_start_for_latent_diffusion = False,
         image_sizes = None,                         # for cascading ddpm, image size at each stage
         lowres_cond_upsample_mode = 'bilinear',     # cascading ddpm - low resolution upsample mode
         lowres_downsample_first = True,             # cascading ddpm - resizes to lower resolution, then to next conditional resolution + blur
@@ -1120,11 +1140,28 @@ class Decoder(nn.Module):
         # automatically take care of ensuring that first unet is unconditional
         # 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), fillvalue = NullVQGanVAE(channels = self.channels))
+
         self.unets = nn.ModuleList([])
-        for ind, one_unet in enumerate(cast_tuple(unet)):
+        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
-            one_unet = one_unet.force_lowres_cond(not is_first)
+            latent_dim = one_vae.encoded_dim if exists(one_vae) else None
+
+            unet_channels = default(latent_dim, self.channels)
+
+            one_unet = one_unet.cast_model_parameters(
+                lowres_cond = not is_first,
+                channels = unet_channels
+            )
+
             self.unets.append(one_unet)
+            self.vaes.append(one_vae.copy_for_eval())
 
         # unet image sizes
 
@@ -1135,6 +1172,10 @@ class Decoder(nn.Module):
         self.image_sizes = image_sizes
         self.sample_channels = cast_tuple(self.channels, len(image_sizes))
 
+        # predict x0 config
+
+        self.predict_x_start = cast_tuple(predict_x_start, len(unets)) if not predict_x_start_for_latent_diffusion else tuple(map(lambda t: isinstance(t, VQGanVAE), self.vaes))
+
         # cascading ddpm related stuff
 
         lowres_conditions = tuple(map(lambda t: t.lowres_cond, self.unets))
@@ -1219,10 +1260,12 @@ class Decoder(nn.Module):
         yield
         unet.cpu()
 
+    @torch.no_grad()
     def get_text_encodings(self, text):
         text_encodings = self.clip.text_transformer(text)
         return text_encodings[:, 1:]
 
+    @torch.no_grad()
     def get_image_embed(self, image):
         image = resize_image_to(image, self.clip_image_size)
         image_encoding = self.clip.visual_transformer(image)
@@ -1251,34 +1294,47 @@ class Decoder(nn.Module):
         posterior_log_variance_clipped = extract(self.posterior_log_variance_clipped, t, x_t.shape)
         return posterior_mean, posterior_variance, posterior_log_variance_clipped
 
-    def p_mean_variance(self, unet, x, t, image_embed, text_encodings = None, lowres_cond_img = None, clip_denoised = True, cond_scale = 1.):
-        pred_noise = 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)
-        x_recon = self.predict_start_from_noise(x, t = t, noise = pred_noise)
+    def p_mean_variance(self, unet, x, t, image_embed, text_encodings = None, lowres_cond_img = None, clip_denoised = True, predict_x_start = False, cond_scale = 1.):
+        pred = 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)
 
-        if clip_denoised:
+        if predict_x_start:
+            x_recon = pred
+        else:
+            x_recon = self.predict_start_from_noise(x, t = t, noise = pred)
+
+        if clip_denoised and not predict_x_start:
             x_recon.clamp_(-1., 1.)
 
         model_mean, posterior_variance, posterior_log_variance = self.q_posterior(x_start=x_recon, x_t=x, t=t)
         return model_mean, posterior_variance, posterior_log_variance
 
     @torch.no_grad()
-    def p_sample(self, unet, x, t, image_embed, text_encodings = None, cond_scale = 1., lowres_cond_img = None, clip_denoised = True, repeat_noise = False):
+    def p_sample(self, unet, x, t, image_embed, text_encodings = None, cond_scale = 1., lowres_cond_img = None, predict_x_start = False, clip_denoised = True, repeat_noise = False):
         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)
+        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 = noise_like(x.shape, device, repeat_noise)
         # 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(self, unet, shape, image_embed, lowres_cond_img = None, text_encodings = None, cond_scale = 1):
+    def p_sample_loop(self, unet, shape, image_embed, predict_x_start = False, lowres_cond_img = None, text_encodings = None, cond_scale = 1):
         device = self.betas.device
 
         b = shape[0]
         img = torch.randn(shape, device = device)
 
         for i in tqdm(reversed(range(0, self.num_timesteps)), desc = 'sampling loop time step', total = self.num_timesteps):
-            img = self.p_sample(unet, img, torch.full((b,), i, device = device, dtype = torch.long), image_embed = image_embed, text_encodings = text_encodings, cond_scale = cond_scale, lowres_cond_img = lowres_cond_img)
+            img = self.p_sample(
+                unet,
+                img,
+                torch.full((b,), i, device = device, dtype = torch.long),
+                image_embed = image_embed,
+                text_encodings = text_encodings,
+                cond_scale = cond_scale,
+                lowres_cond_img = lowres_cond_img,
+                predict_x_start = predict_x_start
+            )
 
         return img
 
@@ -1290,7 +1346,7 @@ class Decoder(nn.Module):
             extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise
         )
 
-    def p_losses(self, unet, x_start, t, *, image_embed, lowres_cond_img = None, text_encodings = None, 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 = t, noise = noise)
@@ -1304,12 +1360,14 @@ class Decoder(nn.Module):
             cond_drop_prob = self.cond_drop_prob
         )
 
+        target = noise if not predict_x_start else x_start
+
         if self.loss_type == 'l1':
-            loss = F.l1_loss(noise, x_recon)
+            loss = F.l1_loss(target, x_recon)
         elif self.loss_type == 'l2':
-            loss = F.mse_loss(noise, x_recon)
+            loss = F.mse_loss(target, x_recon)
         elif self.loss_type == "huber":
-            loss = F.smooth_l1_loss(noise, x_recon)
+            loss = F.smooth_l1_loss(target, x_recon)
         else:
             raise NotImplementedError()
 
@@ -1324,25 +1382,42 @@ class Decoder(nn.Module):
 
         img = None
 
-        for unet, channel, image_size in tqdm(zip(self.unets, self.sample_channels, self.image_sizes)):
+        for unet, vae, channel, image_size, predict_x_start in tqdm(zip(self.unets, self.vaes, self.sample_channels, self.image_sizes, self.predict_x_start)):
             with self.one_unet_in_gpu(unet = unet):
-                lowres_cond_img = self.to_lowres_cond(
-                    img,
-                    target_image_size = image_size
-                ) if unet.lowres_cond else None
+                lowres_cond_img = None
+                shape = (batch_size, channel, image_size, image_size)
+
+                if unet.lowres_cond:
+                    lowres_cond_img = self.to_lowres_cond(img, target_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)
 
                 img = self.p_sample_loop(
                     unet,
-                    (batch_size, channel, image_size, image_size),
+                    shape,
                     image_embed = image_embed,
                     text_encodings = text_encodings,
                     cond_scale = cond_scale,
+                    predict_x_start = predict_x_start,
                     lowres_cond_img = lowres_cond_img
                 )
 
+                img = vae.decode(img)
+
         return img
 
-    def forward(self, image, text = None, image_embed = None, text_encodings = None, unet_number = None):
+    def forward(
+        self,
+        image,
+        text = None,
+        image_embed = None,
+        text_encodings = None,
+        unet_number = None
+    ):
         assert not (len(self.unets) > 1 and not exists(unet_number)), f'you must specify which unet you want trained, from a range of 1 to {len(self.unets)}, if you are training cascading DDPM (multiple unets)'
         unet_number = default(unet_number, 1)
         unet_index = unet_number - 1
@@ -1350,6 +1425,8 @@ class Decoder(nn.Module):
         unet = self.get_unet(unet_number)
 
         target_image_size = self.image_sizes[unet_index]
+        vae = self.vaes[unet_index]
+        predict_x_start = self.predict_x_start[unet_index]
 
         b, c, h, w, device, = *image.shape, image.device
 
@@ -1364,8 +1441,16 @@ class Decoder(nn.Module):
         text_encodings = self.get_text_encodings(text) if exists(text) and not exists(text_encodings) else None
 
         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
-        ddpm_image = resize_image_to(image, target_image_size)
-        return self.p_losses(unet, ddpm_image, times, image_embed = image_embed, text_encodings = text_encodings, lowres_cond_img = lowres_cond_img)
+        image = resize_image_to(image, target_image_size)
+
+        vae.eval()
+        with torch.no_grad():
+            image = vae.encode(image)
+
+            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, predict_x_start = predict_x_start)
 
 # main class
 
@@ -1392,6 +1477,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
@@ -1399,4 +1485,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/train.py

@@ -0,0 +1,53 @@
+import copy
+import torch
+from torch import nn
+
+# exponential moving average wrapper
+
+class EMA(nn.Module):
+    def __init__(
+        self,
+        model,
+        beta = 0.99,
+        ema_update_after_step = 1000,
+        ema_update_every = 10,
+    ):
+        super().__init__()
+        self.beta = beta
+        self.online_model = model
+        self.ema_model = copy.deepcopy(model)
+
+        self.ema_update_after_step = ema_update_after_step # only start EMA after this step number, starting at 0
+        self.ema_update_every = ema_update_every
+
+        self.register_buffer('initted', torch.Tensor([False]))
+        self.register_buffer('step', torch.tensor([0.]))
+
+    def update(self):
+        self.step += 1
+
+        if self.step <= self.ema_update_after_step or (self.step % self.ema_update_every) != 0:
+            return
+
+        if not self.initted:
+            self.ema_model.state_dict(self.online_model.state_dict())
+            self.initted.data.copy_(torch.Tensor([True]))
+
+        self.update_moving_average(self.ema_model, self.online_model)
+
+    def update_moving_average(ma_model, current_model):
+        def calculate_ema(beta, old, new):
+            if not exists(old):
+                return new
+            return old * beta + (1 - beta) * new
+
+        for current_params, ma_params in zip(current_model.parameters(), ma_model.parameters()):
+            old_weight, up_weight = ma_params.data, current_params.data
+            ma_params.data = calculate_ema(self.beta, old_weight, up_weight)
+
+        for current_buffer, ma_buffer in zip(current_model.buffers(), ma_model.buffers()):
+            new_buffer_value = calculate_ema(self.beta, ma_buffer, current_buffer)
+            ma_buffer.copy_(new_buffer_value)
+
+    def __call__(self, *args, **kwargs):
+        return self.ema_model(*args, **kwargs)

dalle2_pytorch/vqgan_vae.py

@@ -13,6 +13,8 @@ import torchvision
 
 from einops import rearrange, reduce, repeat
 
+from dalle2_pytorch.attention import QueryAttnUpsample
+
 # constants
 
 MList = nn.ModuleList
@@ -243,6 +245,7 @@ class ResBlock(nn.Module):
     def forward(self, x):
         return self.net(x) + x
 
+# vqgan attention layer
 class VQGanAttention(nn.Module):
     def __init__(
         self,
@@ -287,6 +290,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,
@@ -294,7 +319,7 @@ class VQGanVAE(nn.Module):
         dim,
         image_size,
         channels = 3,
-        num_layers = 4,
+        layers = 4,
         layer_mults = None,
         l2_recon_loss = False,
         use_hinge_loss = True,
@@ -321,37 +346,39 @@ class VQGanVAE(nn.Module):
 
         self.image_size = image_size
         self.channels = channels
-        self.num_layers = num_layers
-        self.fmap_size = image_size // (num_layers ** 2)
+        self.layers = layers
+        self.fmap_size = image_size // (layers ** 2)
         self.codebook_size = vq_codebook_size
 
         self.encoders = MList([])
         self.decoders = MList([])
 
-        layer_mults = default(layer_mults, list(map(lambda t: 2 ** t, range(num_layers))))
-        assert len(layer_mults) == num_layers, 'layer multipliers must be equal to designated number of layers'
+        layer_mults = default(layer_mults, list(map(lambda t: 2 ** t, range(layers))))
+        assert len(layer_mults) == layers, 'layer multipliers must be equal to designated number of layers'
 
         layer_dims = [dim * mult for mult in layer_mults]
         dims = (dim, *layer_dims)
         codebook_dim = layer_dims[-1]
 
+        self.encoded_dim = dims[-1]
+
         dim_pairs = zip(dims[:-1], dims[1:])
 
         append = lambda arr, t: arr.append(t)
         prepend = lambda arr, t: arr.insert(0, t)
 
         if not isinstance(num_resnet_blocks, tuple):
-            num_resnet_blocks = (*((0,) * (num_layers - 1)), num_resnet_blocks)
+            num_resnet_blocks = (*((0,) * (layers - 1)), num_resnet_blocks)
 
         if not isinstance(use_attn, tuple):
-            use_attn = (*((False,) * (num_layers - 1)), use_attn)
+            use_attn = (*((False,) * (layers - 1)), use_attn)
 
-        assert len(num_resnet_blocks) == num_layers, 'number of resnet blocks config must be equal to number of layers'
-        assert len(use_attn) == num_layers
+        assert len(num_resnet_blocks) == layers, 'number of resnet blocks config must be equal to number of layers'
+        assert len(use_attn) == layers
 
-        for layer_index, (dim_in, dim_out), layer_num_resnet_blocks, layer_use_attn in zip(range(num_layers), dim_pairs, num_resnet_blocks, use_attn):
+        for layer_index, (dim_in, dim_out), layer_num_resnet_blocks, layer_use_attn in zip(range(layers), dim_pairs, num_resnet_blocks, use_attn):
             append(self.encoders, nn.Sequential(nn.Conv2d(dim_in, dim_out, 4, stride = 2, padding = 1), leaky_relu()))
-            prepend(self.decoders, nn.Sequential(nn.Upsample(scale_factor = 2, mode = 'bilinear', align_corners = False), nn.Conv2d(dim_out, dim_in, 3, padding = 1), leaky_relu()))
+            prepend(self.decoders, nn.Sequential(nn.ConvTranspose2d(dim_out, dim_in, 4, 2, 1), leaky_relu()))
 
             if layer_use_attn:
                 prepend(self.decoders, VQGanAttention(dim = dim_out, heads = attn_heads, dim_head = attn_dim_head, dropout = attn_dropout))
@@ -405,6 +432,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())
@@ -434,28 +464,32 @@ class VQGanVAE(nn.Module):
 
         return fmap
 
-    def decode(self, fmap):
-        fmap = self.vq(fmap)
+    def decode(self, fmap, return_indices_and_loss = False):
+        fmap, indices, commit_loss = self.vq(fmap)
 
         for dec in self.decoders:
             fmap = dec(fmap)
 
-        return fmap
+        if not return_indices_and_loss:
+            return fmap
+
+        return fmap, indices, commit_loss
 
     def forward(
         self,
         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}'
         assert channels == self.channels, 'number of channels on image or sketch is not equal to the channels set on this VQGanVAE'
 
-        fmap, indices, commit_loss = self.encode(img)
+        fmap = self.encode(img)
 
-        fmap = self.decode(fmap)
+        fmap, indices, commit_loss = self.decode(fmap, return_indices_and_loss = True)
 
         if not return_loss and not return_discr_loss:
             return fmap
@@ -472,11 +506,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.35',
+  version = '0.0.46',
   license='MIT',
   description = 'DALL-E 2',
   author = 'Phil Wang',