| """ |
| 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): |
| |
| 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 |
| |
|
|
