change up epsilon in layernorm the case of using fp16, thanks to @Veldrovive for figuring out this stabilizes training

README.md

@@ -627,18 +627,6 @@ images = dalle2(
 # save your image (in this example, of size 256x256)
 ```
 
-Alternatively, you can also use <a href="https://github.com/mlfoundations/open_clip">Open Clip</a>
-
-```bash
-$ pip install open-clip-torch
-```
-
-```python
-from dalle2_pytorch import OpenClipAdapter
-
-clip = OpenClipAdapter()
-```
-
 Now you'll just have to worry about training the Prior and the Decoder!
 
 ## Inpainting

dalle2_pytorch/dalle2_pytorch.py

@@ -339,75 +339,6 @@ class OpenAIClipAdapter(BaseClipAdapter):
         image_embed = self.clip.encode_image(image)
         return EmbeddedImage(l2norm(image_embed.float()), None)
 
-class OpenClipAdapter(BaseClipAdapter):
-    def __init__(
-        self,
-        name = 'ViT-B/32',
-        pretrained = 'laion400m_e32'
-    ):
-        import open_clip
-        clip, _, preprocess = open_clip.create_model_and_transforms(name, pretrained = pretrained)
-
-        super().__init__(clip)
-        self.eos_id = 49407
-
-        text_attention_final = self.find_layer('ln_final')
-        self.handle = text_attention_final.register_forward_hook(self._hook)
-        self.clip_normalize = preprocess.transforms[-1]
-        self.cleared = False
-
-    def find_layer(self,  layer):
-        modules = dict([*self.clip.named_modules()])
-        return modules.get(layer, None)
-
-    def clear(self):
-        if self.cleared:
-            return
-
-        self.handle()
-
-    def _hook(self, _, inputs, outputs):
-        self.text_encodings = outputs
-
-    @property
-    def dim_latent(self):
-        return 512
-
-    @property
-    def image_size(self):
-        return self.clip.visual.image_size
-
-    @property
-    def image_channels(self):
-        return 3
-
-    @property
-    def max_text_len(self):
-        return self.clip.context_length
-
-    @torch.no_grad()
-    def embed_text(self, text):
-        text = text[..., :self.max_text_len]
-
-        is_eos_id = (text == self.eos_id)
-        text_mask_excluding_eos = is_eos_id.cumsum(dim = -1) == 0
-        text_mask = F.pad(text_mask_excluding_eos, (1, -1), value = True)
-        assert not self.cleared
-
-        text_embed = self.clip.encode_text(text)
-        text_encodings = self.text_encodings
-        text_encodings = text_encodings.masked_fill(~text_mask[..., None], 0.)
-        del self.text_encodings
-        return EmbeddedText(l2norm(text_embed.float()), text_encodings.float())
-
-    @torch.no_grad()
-    def embed_image(self, image):
-        assert not self.cleared
-        image = self.validate_and_resize_image(image)
-        image = self.clip_normalize(image)
-        image_embed = self.clip.encode_image(image)
-        return EmbeddedImage(l2norm(image_embed.float()), None)
-
 # classifier free guidance functions
 
 def prob_mask_like(shape, prob, device):
@@ -649,7 +580,7 @@ class ChanLayerNorm(nn.Module):
 
         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 Residual(nn.Module):
     def __init__(self, fn):
@@ -770,12 +701,11 @@ class Attention(nn.Module):
         dropout = 0.,
         causal = False,
         rotary_emb = None,
-        cosine_sim = True,
-        cosine_sim_scale = 16
+        pb_relax_alpha = 128
     ):
         super().__init__()
-        self.scale = cosine_sim_scale if cosine_sim else (dim_head ** -0.5)
-        self.cosine_sim = cosine_sim
+        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
@@ -815,13 +745,6 @@ class Attention(nn.Module):
         k = torch.cat((nk, k), dim = -2)
         v = torch.cat((nv, v), dim = -2)
 
-        # whether to use cosine sim
-
-        if self.cosine_sim:
-            q, k = map(l2norm, (q, k))
-
-        q, k = map(lambda t: t * math.sqrt(self.scale), (q, k))
-
         # calculate query / key similarities
 
         sim = einsum('b h i d, b j d -> b h i j', q, k)
@@ -847,6 +770,9 @@ class Attention(nn.Module):
 
         # attention
 
+        sim = sim - sim.amax(dim = -1, keepdim = True).detach()
+        sim = sim * self.pb_relax_alpha
+
         attn = sim.softmax(dim = -1)
         attn = self.dropout(attn)
 
@@ -1559,8 +1485,7 @@ class LinearAttention(nn.Module):
         self,
         dim,
         dim_head = 32,
-        heads = 8,
-        **kwargs
+        heads = 8
     ):
         super().__init__()
         self.scale = dim_head ** -0.5
@@ -1679,7 +1604,6 @@ class Unet(nn.Module):
         lowres_noise_cond = False,       # for conditioning on low resolution noising, based on Imagen
         sparse_attn = False,
         cosine_sim_cross_attn = False,
-        cosine_sim_self_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,
@@ -1800,7 +1724,7 @@ class Unet(nn.Module):
 
         # attention related params
 
-        attn_kwargs = dict(heads = attn_heads, dim_head = attn_dim_head, cosine_sim = cosine_sim_self_attn)
+        attn_kwargs = dict(heads = attn_heads, dim_head = attn_dim_head)
 
         self_attn = cast_tuple(self_attn, num_stages)
 

dalle2_pytorch/version.py

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