Create analogy_encoder.py

analogy_encoder/analogy_encoder.py

@@ -0,0 +1,68 @@
+"""
+ADOBE CONFIDENTIAL
+Copyright 2024 Adobe
+All Rights Reserved.
+NOTICE: All information contained herein is, and remains
+the property of Adobe and its suppliers, if any. The intellectual
+and technical concepts contained herein are proprietary to Adobe
+and its suppliers and are protected by all applicable intellectual
+property laws, including trade secret and copyright laws.
+Dissemination of this information or reproduction of this material
+is strictly forbidden unless prior written permission is obtained
+from Adobe.
+"""
+
+import torch as th
+from diffusers import ModelMixin
+from transformers import AutoModel, SiglipVisionConfig, Dinov2Config
+from transformers import SiglipVisionModel
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+    
+class AnalogyEncoder(ModelMixin, ConfigMixin):
+    @register_to_config
+    def __init__(self, load_pretrained=False, 
+                 dino_config_dict=None, siglip_config_dict=None):
+        super().__init__()
+        if load_pretrained:
+            image_encoder_dino = AutoModel.from_pretrained('facebook/dinov2-large', torch_dtype=th.float16)
+            image_encoder_siglip = SiglipVisionModel.from_pretrained("google/siglip-large-patch16-256", torch_dtype=th.float16, attn_implementation="sdpa")
+        else:
+            image_encoder_dino = AutoModel.from_config(Dinov2Config.from_dict(dino_config_dict))
+            image_encoder_siglip = AutoModel.from_config(SiglipVisionConfig.from_dict(siglip_config_dict))
+            
+        image_encoder_dino.requires_grad_(False)
+        image_encoder_dino = image_encoder_dino.to(memory_format=th.channels_last)
+
+        image_encoder_siglip.requires_grad_(False)
+        image_encoder_siglip = image_encoder_siglip.to(memory_format=th.channels_last)
+        self.image_encoder_dino = image_encoder_dino
+        self.image_encoder_siglip = image_encoder_siglip
+
+
+    def dino_normalization(self, encoder_output):
+        embeds = encoder_output.last_hidden_state
+        embeds_pooled = embeds[:, 0:1]
+        embeds = embeds / th.norm(embeds_pooled, dim=-1, keepdim=True)
+        return embeds
+    
+    def siglip_normalization(self, encoder_output):
+        embeds = th.cat ([encoder_output.pooler_output[:, None, :], encoder_output.last_hidden_state], dim=1)
+        embeds_pooled = embeds[:, 0:1]
+        embeds = embeds / th.norm(embeds_pooled, dim=-1, keepdim=True)
+        return embeds
+    
+    def forward(self, dino_in, siglip_in):
+
+        x_1 = self.image_encoder_dino(dino_in, output_hidden_states=True)
+        x_1_first = x_1.hidden_states[0]
+        x_1 = self.dino_normalization(x_1)
+        x_2 = self.image_encoder_siglip(siglip_in, output_hidden_states=True)
+        x_2_first = x_2.hidden_states[0]
+        x_2_first_pool = th.mean(x_2_first, dim=1, keepdim=True)
+        x_2_first = th.cat([x_2_first_pool, x_2_first], 1)
+        x_2 = self.siglip_normalization(x_2)
+        dino_embd = th.cat([x_1, x_1_first], -1)
+        siglip_embd = th.cat([x_2, x_2_first], -1)
+        return dino_embd, siglip_embd
+