fix repaint

README.md

@@ -371,7 +371,6 @@ loss.backward()
 unet1 = Unet(
     dim = 128,
     image_embed_dim = 512,
-    text_embed_dim = 512,
     cond_dim = 128,
     channels = 3,
     dim_mults=(1, 2, 4, 8),
@@ -396,7 +395,7 @@ decoder = Decoder(
 ).cuda()
 
 for unet_number in (1, 2):
-    loss = decoder(images, text = text, unet_number = unet_number) # this can optionally be decoder(images, text) if you wish to condition on the text encodings as well, though it was hinted in the paper it didn't do much
+    loss = decoder(images, unet_number = unet_number) # this can optionally be decoder(images, text) if you wish to condition on the text encodings as well, though it was hinted in the paper it didn't do much
     loss.backward()
 
 # do above for many steps
@@ -861,23 +860,25 @@ unet1 = Unet(
     text_embed_dim = 512,
     cond_dim = 128,
     channels = 3,
-    dim_mults=(1, 2, 4, 8),
-    cond_on_text_encodings = True,
+    dim_mults=(1, 2, 4, 8)
 ).cuda()
 
 unet2 = Unet(
     dim = 16,
     image_embed_dim = 512,
+    text_embed_dim = 512,
     cond_dim = 128,
     channels = 3,
     dim_mults = (1, 2, 4, 8, 16),
+    cond_on_text_encodings = True
 ).cuda()
 
 decoder = Decoder(
     unet = (unet1, unet2),
     image_sizes = (128, 256),
     clip = clip,
-    timesteps = 1000
+    timesteps = 1000,
+    condition_on_text_encodings = True
 ).cuda()
 
 decoder_trainer = DecoderTrainer(
@@ -902,8 +903,8 @@ for unet_number in (1, 2):
 # after much training
 # you can sample from the exponentially moving averaged unets as so
 
-mock_image_embed = torch.randn(32, 512).cuda()
-images = decoder_trainer.sample(image_embed = mock_image_embed, text = text) # (4, 3, 256, 256)
+mock_image_embed = torch.randn(4, 512).cuda()
+images = decoder_trainer.sample(mock_image_embed, text = text) # (4, 3, 256, 256)
 ```
 
 ### Diffusion Prior Training
@@ -1111,8 +1112,7 @@ For detailed information on training the diffusion prior, please refer to the [d
 - [x] allow for unet to be able to condition non-cross attention style as well
 - [x] speed up inference, read up on papers (ddim)
 - [x] add inpainting ability using resampler from repaint paper https://arxiv.org/abs/2201.09865
-- [x] add the final combination of upsample feature maps, used in unet squared, seems to have an effect in local experiments
-- [ ] consider elucidated dalle2 https://arxiv.org/abs/2206.00364
+- [ ] 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

dalle2_pytorch/dalle2_pytorch.py

@@ -547,34 +547,28 @@ class NoiseScheduler(nn.Module):
 # diffusion prior
 
 class LayerNorm(nn.Module):
-    def __init__(self, dim, eps = 1e-5, fp16_eps = 1e-3, stable = False):
+    def __init__(self, dim, eps = 1e-5, stable = False):
         super().__init__()
         self.eps = eps
-        self.fp16_eps = fp16_eps
         self.stable = stable
         self.g = nn.Parameter(torch.ones(dim))
 
     def forward(self, x):
-        eps = self.eps if x.dtype == torch.float32 else self.fp16_eps
-
         if self.stable:
             x = x / x.amax(dim = -1, keepdim = True).detach()
 
         var = torch.var(x, dim = -1, unbiased = False, keepdim = True)
         mean = torch.mean(x, dim = -1, keepdim = True)
-        return (x - mean) * (var + eps).rsqrt() * self.g
+        return (x - mean) * (var + self.eps).rsqrt() * self.g
 
 class ChanLayerNorm(nn.Module):
-    def __init__(self, dim, eps = 1e-5, fp16_eps = 1e-3, stable = False):
+    def __init__(self, dim, eps = 1e-5, stable = False):
         super().__init__()
         self.eps = eps
-        self.fp16_eps = fp16_eps
         self.stable = stable
         self.g = nn.Parameter(torch.ones(1, dim, 1, 1))
 
     def forward(self, x):
-        eps = self.eps if x.dtype == torch.float32 else self.fp16_eps
-
         if self.stable:
             x = x / x.amax(dim = 1, keepdim = True).detach()
 
@@ -1363,8 +1357,7 @@ class ResnetBlock(nn.Module):
         *,
         cond_dim = None,
         time_cond_dim = None,
-        groups = 8,
-        cosine_sim_cross_attn = False
+        groups = 8
     ):
         super().__init__()
 
@@ -1384,8 +1377,7 @@ class ResnetBlock(nn.Module):
                 'b (h w) c',
                 CrossAttention(
                     dim = dim_out,
-                    context_dim = cond_dim,
-                    cosine_sim = cosine_sim_cross_attn
+                    context_dim = cond_dim
                 )
             )
 
@@ -1420,12 +1412,11 @@ class CrossAttention(nn.Module):
         heads = 8,
         dropout = 0.,
         norm_context = False,
-        cosine_sim = False,
-        cosine_sim_scale = 16
+        pb_relax_alpha = 32 ** 2
     ):
         super().__init__()
-        self.cosine_sim = cosine_sim
-        self.scale = cosine_sim_scale if cosine_sim else (dim_head ** -0.5)
+        self.pb_relax_alpha = pb_relax_alpha
+        self.scale = dim_head ** -0.5 * (pb_relax_alpha ** -1)
         self.heads = heads
         inner_dim = dim_head * heads
 
@@ -1461,10 +1452,7 @@ class CrossAttention(nn.Module):
         k = torch.cat((nk, k), dim = -2)
         v = torch.cat((nv, v), dim = -2)
 
-        if self.cosine_sim:
-            q, k = map(l2norm, (q, k))
-
-        q, k = map(lambda t: t * math.sqrt(self.scale), (q, k))
+        q = q * self.scale
 
         sim = einsum('b h i d, b h j d -> b h i j', q, k)
         max_neg_value = -torch.finfo(sim.dtype).max
@@ -1474,6 +1462,9 @@ class CrossAttention(nn.Module):
             mask = rearrange(mask, 'b j -> b 1 1 j')
             sim = sim.masked_fill(~mask, max_neg_value)
 
+        sim = sim - sim.amax(dim = -1, keepdim = True).detach()
+        sim = sim * self.pb_relax_alpha
+
         attn = sim.softmax(dim = -1)
 
         out = einsum('b h i j, b h j d -> b h i d', attn, v)
@@ -1503,7 +1494,6 @@ class LinearAttention(nn.Module):
 
     def forward(self, fmap):
         h, x, y = self.heads, *fmap.shape[-2:]
-        seq_len = x * y
 
         fmap = self.norm(fmap)
         q, k, v = self.to_qkv(fmap).chunk(3, dim = 1)
@@ -1513,9 +1503,6 @@ class LinearAttention(nn.Module):
         k = k.softmax(dim = -2)
 
         q = q * self.scale
-        v = l2norm(v)
-
-        k, v = map(lambda t: t / math.sqrt(seq_len), (k, v))
 
         context = einsum('b n d, b n e -> b d e', k, v)
         out = einsum('b n d, b d e -> b n e', q, context)
@@ -1551,38 +1538,6 @@ class CrossEmbedLayer(nn.Module):
         fmaps = tuple(map(lambda conv: conv(x), self.convs))
         return torch.cat(fmaps, dim = 1)
 
-class UpsampleCombiner(nn.Module):
-    def __init__(
-        self,
-        dim,
-        *,
-        enabled = False,
-        dim_ins = tuple(),
-        dim_outs = tuple()
-    ):
-        super().__init__()
-        assert len(dim_ins) == len(dim_outs)
-        self.enabled = enabled
-
-        if not self.enabled:
-            self.dim_out = dim
-            return
-
-        self.fmap_convs = nn.ModuleList([Block(dim_in, dim_out) for dim_in, dim_out in zip(dim_ins, dim_outs)])
-        self.dim_out = dim + (sum(dim_outs) if len(dim_outs) > 0 else 0)
-
-    def forward(self, x, fmaps = None):
-        target_size = x.shape[-1]
-
-        fmaps = default(fmaps, tuple())
-
-        if not self.enabled or len(fmaps) == 0 or len(self.fmap_convs) == 0:
-            return x
-
-        fmaps = [resize_image_to(fmap, target_size) for fmap in fmaps]
-        outs = [conv(fmap) for fmap, conv in zip(fmaps, self.fmap_convs)]
-        return torch.cat((x, *outs), dim = 1)
-
 class Unet(nn.Module):
     def __init__(
         self,
@@ -1603,7 +1558,6 @@ class Unet(nn.Module):
         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,
-        cosine_sim_cross_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)
         cond_on_text_encodings = False,
         max_text_len = 256,
@@ -1621,7 +1575,6 @@ class Unet(nn.Module):
         scale_skip_connection = False,
         pixel_shuffle_upsample = True,
         final_conv_kernel_size = 1,
-        combine_upsample_fmaps = False, # whether to combine the outputs of all upsample blocks, as in unet squared paper
         **kwargs
     ):
         super().__init__()
@@ -1747,13 +1700,9 @@ class Unet(nn.Module):
 
         upsample_klass = NearestUpsample if not pixel_shuffle_upsample else PixelShuffleUpsample
 
-        # prepare resnet klass
-
-        resnet_block = partial(ResnetBlock, cosine_sim_cross_attn = cosine_sim_cross_attn)
-
         # give memory efficient unet an initial resnet block
 
-        self.init_resnet_block = resnet_block(init_dim, init_dim, time_cond_dim = time_cond_dim, groups = top_level_resnet_group) if memory_efficient else None
+        self.init_resnet_block = ResnetBlock(init_dim, init_dim, time_cond_dim = time_cond_dim, groups = top_level_resnet_group) if memory_efficient else None
 
         # layers
 
@@ -1761,8 +1710,7 @@ class Unet(nn.Module):
         self.ups = nn.ModuleList([])
         num_resolutions = len(in_out)
 
-        skip_connect_dims = []          # keeping track of skip connection dimensions
-        upsample_combiner_dims = []     # keeping track of dimensions for final upsample feature map combiner
+        skip_connect_dims = [] # keeping track of skip connection dimensions
 
         for ind, ((dim_in, dim_out), groups, layer_num_resnet_blocks, layer_self_attn) in enumerate(zip(in_out, resnet_groups, num_resnet_blocks, self_attn)):
             is_first = ind == 0
@@ -1780,17 +1728,17 @@ class Unet(nn.Module):
 
             self.downs.append(nn.ModuleList([
                 downsample_klass(dim_in, dim_out = dim_out) if memory_efficient else None,
-                resnet_block(dim_layer, dim_layer, time_cond_dim = time_cond_dim, groups = groups),
-                nn.ModuleList([resnet_block(dim_layer, dim_layer, cond_dim = layer_cond_dim, time_cond_dim = time_cond_dim, groups = groups) for _ in range(layer_num_resnet_blocks)]),
+                ResnetBlock(dim_layer, dim_layer, time_cond_dim = time_cond_dim, groups = groups),
+                nn.ModuleList([ResnetBlock(dim_layer, dim_layer, cond_dim = layer_cond_dim, time_cond_dim = time_cond_dim, groups = groups) for _ in range(layer_num_resnet_blocks)]),
                 attention,
                 downsample_klass(dim_layer, dim_out = dim_out) if not is_last and not memory_efficient else nn.Conv2d(dim_layer, dim_out, 1)
             ]))
 
         mid_dim = dims[-1]
 
-        self.mid_block1 = resnet_block(mid_dim, mid_dim, cond_dim = cond_dim, time_cond_dim = time_cond_dim, groups = resnet_groups[-1])
+        self.mid_block1 = ResnetBlock(mid_dim, mid_dim, cond_dim = cond_dim, time_cond_dim = time_cond_dim, groups = resnet_groups[-1])
         self.mid_attn = create_self_attn(mid_dim)
-        self.mid_block2 = resnet_block(mid_dim, mid_dim, cond_dim = cond_dim, time_cond_dim = time_cond_dim, groups = resnet_groups[-1])
+        self.mid_block2 = ResnetBlock(mid_dim, mid_dim, cond_dim = cond_dim, time_cond_dim = time_cond_dim, groups = resnet_groups[-1])
 
         for ind, ((dim_in, dim_out), groups, layer_num_resnet_blocks, layer_self_attn) in enumerate(zip(reversed(in_out), reversed(resnet_groups), reversed(num_resnet_blocks), reversed(self_attn))):
             is_last = ind >= (len(in_out) - 1)
@@ -1804,27 +1752,14 @@ class Unet(nn.Module):
             elif sparse_attn:
                 attention = Residual(LinearAttention(dim_out, **attn_kwargs))
 
-            upsample_combiner_dims.append(dim_out)
-
             self.ups.append(nn.ModuleList([
-                resnet_block(dim_out + skip_connect_dim, dim_out, cond_dim = layer_cond_dim, time_cond_dim = time_cond_dim, groups = groups),
-                nn.ModuleList([resnet_block(dim_out + skip_connect_dim, dim_out, cond_dim = layer_cond_dim, time_cond_dim = time_cond_dim, groups = groups)  for _ in range(layer_num_resnet_blocks)]),
+                ResnetBlock(dim_out + skip_connect_dim, dim_out, cond_dim = layer_cond_dim, time_cond_dim = time_cond_dim, groups = groups),
+                nn.ModuleList([ResnetBlock(dim_out + skip_connect_dim, dim_out, cond_dim = layer_cond_dim, time_cond_dim = time_cond_dim, groups = groups)  for _ in range(layer_num_resnet_blocks)]),
                 attention,
                 upsample_klass(dim_out, dim_in) if not is_last or memory_efficient else nn.Identity()
             ]))
 
-        # whether to combine outputs from all upsample blocks for final resnet block
-
-        self.upsample_combiner = UpsampleCombiner(
-            dim = dim,
-            enabled = combine_upsample_fmaps,
-            dim_ins = upsample_combiner_dims,
-            dim_outs = (dim,) * len(upsample_combiner_dims)
-        )
-
-        # a final resnet block
-
-        self.final_resnet_block = resnet_block(self.upsample_combiner.dim_out + dim, dim, time_cond_dim = time_cond_dim, groups = top_level_resnet_group)
+        self.final_resnet_block = ResnetBlock(dim * 2, dim, time_cond_dim = time_cond_dim, groups = top_level_resnet_group)
 
         out_dim_in = dim + (channels if lowres_cond else 0)
 
@@ -1848,7 +1783,7 @@ class Unet(nn.Module):
             channels == self.channels and \
             cond_on_image_embeds == self.cond_on_image_embeds and \
             cond_on_text_encodings == self.cond_on_text_encodings and \
-            lowres_noise_cond == self.lowres_noise_cond and \
+            cond_on_lowres_noise == self.cond_on_lowres_noise and \
             channels_out == self.channels_out:
             return self
 
@@ -2018,8 +1953,7 @@ class Unet(nn.Module):
 
         # go through the layers of the unet, down and up
 
-        down_hiddens = []
-        up_hiddens = []
+        hiddens = []
 
         for pre_downsample, init_block, resnet_blocks, attn, post_downsample in self.downs:
             if exists(pre_downsample):
@@ -2029,10 +1963,10 @@ class Unet(nn.Module):
 
             for resnet_block in resnet_blocks:
                 x = resnet_block(x, t, c)
-                down_hiddens.append(x.contiguous())
+                hiddens.append(x)
 
             x = attn(x)
-            down_hiddens.append(x.contiguous())
+            hiddens.append(x.contiguous())
 
             if exists(post_downsample):
                 x = post_downsample(x)
@@ -2044,7 +1978,7 @@ class Unet(nn.Module):
 
         x = self.mid_block2(x, t, mid_c)
 
-        connect_skip = lambda fmap: torch.cat((fmap, down_hiddens.pop() * self.skip_connect_scale), dim = 1)
+        connect_skip = lambda fmap: torch.cat((fmap, hiddens.pop() * self.skip_connect_scale), dim = 1)
 
         for init_block, resnet_blocks, attn, upsample in self.ups:
             x = connect_skip(x)
@@ -2055,12 +1989,8 @@ class Unet(nn.Module):
                 x = resnet_block(x, t, c)
 
             x = attn(x)
-
-            up_hiddens.append(x.contiguous())
             x = upsample(x)
 
-        x = self.upsample_combiner(x, up_hiddens)
-
         x = torch.cat((x, r), dim = 1)
 
         x = self.final_resnet_block(x, t)
@@ -2659,7 +2589,7 @@ class Decoder(nn.Module):
                 if is_inpaint and not (is_last_timestep or is_last_resample_step):
                     # in repaint, you renoise and resample up to 10 times every step
                     time_next_cond = torch.full((batch,), time_next, device = device, dtype = torch.long)
-                    img = noise_scheduler.q_sample_from_to(img, time_next_cond, time_cond)
+                    img = noise_scheduler.q_sample_from_to(img, time_cond, time_next_cond)
 
         if exists(inpaint_image):
             img = (img * ~inpaint_mask) + (inpaint_image * inpaint_mask)
@@ -2955,7 +2885,7 @@ class DALLE2(nn.Module):
         image_embed = self.prior.sample(text, num_samples_per_batch = self.prior_num_samples, cond_scale = prior_cond_scale)
 
         text_cond = text if self.decoder_need_text_cond else None
-        images = self.decoder.sample(image_embed = image_embed, text = text_cond, cond_scale = cond_scale)
+        images = self.decoder.sample(image_embed, text = text_cond, cond_scale = cond_scale)
 
         if return_pil_images:
             images = list(map(self.to_pil, images.unbind(dim = 0)))

dalle2_pytorch/trainer.py

@@ -174,7 +174,7 @@ class DiffusionPriorTrainer(nn.Module):
     def __init__(
         self,
         diffusion_prior,
-        accelerator = None,
+        accelerator,
         use_ema = True,
         lr = 3e-4,
         wd = 1e-2,
@@ -186,12 +186,8 @@ class DiffusionPriorTrainer(nn.Module):
     ):
         super().__init__()
         assert isinstance(diffusion_prior, DiffusionPrior)
-
+        assert isinstance(accelerator, Accelerator)
         ema_kwargs, kwargs = groupby_prefix_and_trim('ema_', kwargs)
-        accelerator_kwargs, kwargs = groupby_prefix_and_trim('accelerator_', kwargs)
-
-        if not exists(accelerator):
-            accelerator = Accelerator(**accelerator_kwargs)
 
         # assign some helpful member vars
 
@@ -304,7 +300,7 @@ class DiffusionPriorTrainer(nn.Module):
 
         # all processes need to load checkpoint. no restriction here
         if isinstance(path_or_state, str):
-            path = Path(path_or_state)
+            path = Path(path)
             assert path.exists()
             loaded_obj = torch.load(str(path), map_location=self.device)
 

dalle2_pytorch/version.py

@@ -1 +1 @@
-__version__ = '1.4.1'
+__version__ = '1.0.4'