make open clip available for use with dalle2 pytorch

dalle2_pytorch/dalle2_pytorch.py

@@ -339,6 +339,75 @@ 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):