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