|
from abc import abstractmethod |
|
|
|
import torch |
|
from diffusers.configuration_utils import ConfigMixin, register_to_config |
|
from diffusers.models.modeling_utils import ModelMixin |
|
from torch import Tensor |
|
from torch.nn import functional as F |
|
from torchvision.transforms import v2 as T |
|
|
|
from .common import ensure_tuple |
|
from .vit import VisionTransformer, vit_base_dreamsim |
|
|
|
|
|
class DreamsimBackbone(ModelMixin, ConfigMixin): |
|
@abstractmethod |
|
def forward_features(self, x: Tensor) -> Tensor: |
|
raise NotImplementedError("abstract base class was called ;_;") |
|
|
|
def forward(self, x: Tensor) -> Tensor: |
|
"""Dreamsim forward pass for similarity computation. |
|
Args: |
|
x (Tensor): Input tensor of shape [2, B, 3, H, W]. |
|
|
|
Returns: |
|
sim (torch.Tensor): dreamsim similarity score of shape [B]. |
|
""" |
|
inputs = x.view(-1, 3, *x.shape[-2:]) |
|
|
|
x = self.forward_features(inputs).view(*x.shape[:2], -1) |
|
|
|
return 1 - F.cosine_similarity(x[0], x[1], dim=1) |
|
|
|
def compile(self, *args, **kwargs): |
|
"""Compile the model with Inductor. This is a no-op unless overridden by a subclass.""" |
|
return self |
|
|
|
|
|
class DreamsimModel(DreamsimBackbone): |
|
@register_to_config |
|
def __init__( |
|
self, |
|
image_size: int = 224, |
|
patch_size: int = 16, |
|
layer_norm_eps: float = 1e-6, |
|
pre_norm: bool = False, |
|
act_layer: str = "gelu", |
|
img_mean: tuple[float, float, float] = (0.485, 0.456, 0.406), |
|
img_std: tuple[float, float, float] = (0.229, 0.224, 0.225), |
|
do_resize: bool = False, |
|
) -> None: |
|
super().__init__() |
|
|
|
self.image_size = ensure_tuple(image_size, 2) |
|
self.patch_size = ensure_tuple(patch_size, 2) |
|
self.layer_norm_eps = layer_norm_eps |
|
self.pre_norm = pre_norm |
|
self.do_resize = do_resize |
|
self.img_mean = img_mean |
|
self.img_std = img_std |
|
|
|
num_classes = 512 if self.pre_norm else 0 |
|
self.extractor: VisionTransformer = vit_base_dreamsim( |
|
image_size=image_size, |
|
patch_size=patch_size, |
|
layer_norm_eps=layer_norm_eps, |
|
num_classes=num_classes, |
|
pre_norm=pre_norm, |
|
act_layer=act_layer, |
|
) |
|
|
|
self.resize = T.Resize( |
|
self.image_size, |
|
interpolation=T.InterpolationMode.BICUBIC, |
|
antialias=True, |
|
) |
|
self.img_norm = T.Normalize(mean=self.img_mean, std=self.img_std) |
|
|
|
def compile(self, *, mode: str = "reduce-overhead", force: bool = False, **kwargs): |
|
if (not self._compiled) or force: |
|
self.extractor = torch.compile(self.extractor, mode=mode, **kwargs) |
|
self._compiled = True |
|
return self |
|
|
|
def transforms(self, x: Tensor) -> Tensor: |
|
if self.do_resize: |
|
x = self.resize(x) |
|
return self.img_norm(x) |
|
|
|
def forward_features(self, x: Tensor) -> Tensor: |
|
if x.ndim == 3: |
|
x = x.unsqueeze(0) |
|
x = self.transforms(x) |
|
x = self.extractor.forward(x, norm=self.pre_norm) |
|
|
|
x = x.div(x.norm(dim=1, keepdim=True)) |
|
x = x.sub(x.mean(dim=1, keepdim=True)) |
|
return x |
|
|
|
|
|
class DreamsimEnsemble(DreamsimBackbone): |
|
@register_to_config |
|
def __init__( |
|
self, |
|
image_size: int = 224, |
|
patch_size: int = 16, |
|
layer_norm_eps: float | tuple[float, ...] = (1e-6, 1e-5, 1e-5), |
|
num_classes: int | tuple[int, ...] = (0, 512, 512), |
|
do_resize: bool = False, |
|
) -> None: |
|
super().__init__() |
|
if isinstance(layer_norm_eps, float): |
|
layer_norm_eps = (layer_norm_eps,) * 3 |
|
if isinstance(num_classes, int): |
|
num_classes = (num_classes,) * 3 |
|
|
|
self.image_size = ensure_tuple(image_size, 2) |
|
self.patch_size = ensure_tuple(patch_size, 2) |
|
self.do_resize = do_resize |
|
|
|
self.dino: VisionTransformer = vit_base_dreamsim( |
|
image_size=self.image_size, |
|
patch_size=self.patch_size, |
|
layer_norm_eps=layer_norm_eps[0], |
|
num_classes=num_classes[0], |
|
pre_norm=False, |
|
act_layer="gelu", |
|
) |
|
self.clip1: VisionTransformer = vit_base_dreamsim( |
|
image_size=self.image_size, |
|
patch_size=self.patch_size, |
|
layer_norm_eps=layer_norm_eps[1], |
|
num_classes=num_classes[1], |
|
pre_norm=True, |
|
act_layer="quick_gelu", |
|
) |
|
self.clip2: VisionTransformer = vit_base_dreamsim( |
|
image_size=self.image_size, |
|
patch_size=self.patch_size, |
|
layer_norm_eps=layer_norm_eps[2], |
|
num_classes=num_classes[2], |
|
pre_norm=True, |
|
act_layer="gelu", |
|
) |
|
|
|
self.resize = T.Resize( |
|
self.image_size, |
|
interpolation=T.InterpolationMode.BICUBIC, |
|
antialias=True, |
|
) |
|
self.dino_norm = T.Normalize( |
|
mean=(0.485, 0.456, 0.406), |
|
std=(0.229, 0.224, 0.225), |
|
) |
|
self.clip_norm = T.Normalize( |
|
mean=(0.48145466, 0.4578275, 0.40821073), |
|
std=(0.26862954, 0.26130258, 0.27577711), |
|
) |
|
|
|
self._compiled = False |
|
|
|
def compile(self, *, mode: str = "reduce-overhead", force: bool = False, **kwargs): |
|
if (not self._compiled) or force: |
|
self.dino = torch.compile(self.dino, mode=mode, **kwargs) |
|
self.clip1 = torch.compile(self.clip1, mode=mode, **kwargs) |
|
self.clip2 = torch.compile(self.clip2, mode=mode, **kwargs) |
|
self._compiled = True |
|
return self |
|
|
|
def transforms(self, x: Tensor, resize: bool = False) -> tuple[Tensor, Tensor, Tensor]: |
|
if resize: |
|
x = self.resize(x) |
|
x = self.dino_norm(x), self.clip_norm(x), self.clip_norm(x) |
|
return x |
|
|
|
def forward_features(self, x: Tensor) -> Tensor: |
|
if x.ndim == 3: |
|
x = x.unsqueeze(0) |
|
x_dino, x_clip1, x_clip2 = self.transforms(x, self.do_resize) |
|
|
|
|
|
x_dino = self.dino.forward(x_dino, norm=False) |
|
x_clip1 = self.clip1.forward(x_clip1, norm=True) |
|
x_clip2 = self.clip2.forward(x_clip2, norm=True) |
|
|
|
z: Tensor = torch.cat([x_dino, x_clip1, x_clip2], dim=1) |
|
z = z.div(z.norm(dim=1, keepdim=True)) |
|
z = z.sub(z.mean(dim=1, keepdim=True)) |
|
return z |
|
|
