Create analogy_input_processor.py

analogy_input_processor/analogy_input_processor.py

@@ -0,0 +1,96 @@
+"""
+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 torchvision import transforms
+from diffusers import ModelMixin
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+
+DINO_SIZE = 224
+DINO_MEAN = [0.485, 0.456, 0.406]
+DINO_STD = [0.229, 0.224, 0.225]
+
+SIGLIP_SIZE = 256
+SIGLIP_MEAN = [0.5]
+SIGLIP_STD = [0.5]
+
+    
+class AnalogyInputProcessor(ModelMixin, ConfigMixin):
+    
+    @register_to_config
+    def __init__(self,):
+        super(AnalogyInputProcessor, self).__init__()
+        
+        self.dino_transform = transforms.Compose(
+            [
+                transforms.Resize((DINO_SIZE, DINO_SIZE)),
+                transforms.ToTensor(),
+                transforms.Normalize(DINO_MEAN, DINO_STD), 
+            ]
+        )
+        
+        self.siglip_transform = transforms.Compose(
+            [
+                transforms.Resize((SIGLIP_SIZE, SIGLIP_SIZE)),
+                transforms.ToTensor(),
+                transforms.Normalize(SIGLIP_MEAN, SIGLIP_STD), 
+            ]
+        )
+        
+        dino_mean = th.tensor(DINO_MEAN).view(1, 3, 1, 1)
+        dino_std = th.tensor(DINO_STD).view(1, 3, 1, 1)
+        siglip_mean = [SIGLIP_MEAN[0],] * 3
+        siglip_std = [SIGLIP_STD[0],] * 3
+        siglip_mean = th.tensor(siglip_mean).view(1, 3, 1, 1)
+        siglip_std = th.tensor(siglip_std).view(1, 3, 1, 1)
+        self.register_buffer("dino_mean", dino_mean)
+        self.register_buffer("dino_std", dino_std)
+        self.register_buffer("siglip_mean", siglip_mean)
+        self.register_buffer("siglip_std", siglip_std)
+        
+    def __call__(self, analogy_prompt):
+        # List of tuples of (A, A*, B)
+        img_a_dino = []
+        img_a_siglip = []
+        img_a_star_dino = []
+        img_a_star_siglip = []
+        img_b_dino = []
+        img_b_siglip = []
+        
+        for im_set in analogy_prompt:
+            img_a, img_a_star, img_b = im_set
+            img_a_dino.append(self.dino_transform(img_a))
+            img_a_siglip.append(self.siglip_transform(img_a))
+            img_a_star_dino.append(self.dino_transform(img_a_star))
+            img_a_star_siglip.append(self.siglip_transform(img_a_star))
+            img_b_dino.append(self.dino_transform(img_b))
+            img_b_siglip.append(self.siglip_transform(img_b))
+        
+        img_a_dino = th.stack(img_a_dino, 0)
+        img_a_siglip = th.stack(img_a_siglip, 0)
+        img_a_star_dino = th.stack(img_a_star_dino, 0)
+        img_a_star_siglip = th.stack(img_a_star_siglip, 0)
+        img_b_dino = th.stack(img_b_dino, 0)
+        img_b_siglip = th.stack(img_b_siglip, 0)
+        
+        dino_combined_input = th.stack([img_b_dino, img_a_dino, img_a_star_dino], 0)
+        siglip_combined_input = th.stack([img_b_siglip, img_a_siglip, img_a_star_siglip], 0)
+        
+        return dino_combined_input, siglip_combined_input
+    def get_negative(self, dino_in, siglip_in):
+        
+        dino_i = ((dino_in * 0 + 0.5) - self.dino_mean) / self.dino_std
+        siglip_i = ((siglip_in * 0 + 0.5) - self.siglip_mean) / self.siglip_std
+        return dino_i, siglip_i
+