change dreamsim class hierarchy a bit

__init__.py

@@ -1,9 +1,10 @@
-from .model import DreamsimEnsemble, DreamsimModel
+from .model import DreamsimBackbone, DreamsimEnsemble, DreamsimModel
 from .vit import VisionTransformer, vit_base_dreamsim
 
 __all__ = [
-    "DreamsimModel",
+    "DreamsimBackbone",
     "DreamsimEnsemble",
+    "DreamsimModel",
     "VisionTransformer",
     "vit_base_dreamsim",
 ]

model.py

@@ -1,3 +1,5 @@
+from abc import abstractmethod
+
 import torch
 from diffusers.configuration_utils import ConfigMixin, register_to_config
 from diffusers.models.modeling_utils import ModelMixin
@@ -9,7 +11,31 @@ from .common import ensure_tuple
 from .vit import VisionTransformer, vit_base_dreamsim
 
 
-class DreamsimModel(ModelMixin, ConfigMixin):
+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,
@@ -25,7 +51,7 @@ class DreamsimModel(ModelMixin, ConfigMixin):
         super().__init__()
 
         self.image_size = ensure_tuple(image_size, 2)
-        self.patch_size = patch_size
+        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
@@ -49,6 +75,12 @@ class DreamsimModel(ModelMixin, ConfigMixin):
         )
         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)
@@ -60,42 +92,29 @@ class DreamsimModel(ModelMixin, ConfigMixin):
         x = self.transforms(x)
         x = self.extractor.forward(x, norm=self.pre_norm)
 
-        x.div_(x.norm(dim=1, keepdim=True))
-        x.sub_(x.mean(dim=1, keepdim=True))
+        x = x.div(x.norm(dim=1, keepdim=True))
+        x = x.sub(x.mean(dim=1, keepdim=True))
         return x
 
-    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].
-        """
-        all_images = x.view(-1, 3, *x.shape[-2:])
-
-        x = self.forward_features(all_images)
-        x = x.view(*x.shape[:2], -1)
-
-        return 1 - F.cosine_similarity(x[0], x[1], dim=1)
-
-
-class DreamsimEnsemble(ModelMixin, ConfigMixin):
+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: tuple[int, int, int] = (0, 512, 512),
+        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 = patch_size
+        self.patch_size = ensure_tuple(patch_size, 2)
         self.do_resize = do_resize
 
         self.dino: VisionTransformer = vit_base_dreamsim(
@@ -137,10 +156,21 @@ class DreamsimEnsemble(ModelMixin, ConfigMixin):
             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)
-        return self.dino_norm(x), self.clip_norm(x), self.clip_norm(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:
@@ -153,21 +183,6 @@ class DreamsimEnsemble(ModelMixin, ConfigMixin):
         x_clip2 = self.clip2.forward(x_clip2, norm=True)
 
         z: Tensor = torch.cat([x_dino, x_clip1, x_clip2], dim=1)
-        z.div_(z.norm(dim=1, keepdim=True))
-        z.sub_(z.mean(dim=1, keepdim=True))
+        z = z.div(z.norm(dim=1, keepdim=True))
+        z = z.sub(z.mean(dim=1, keepdim=True))
         return z
-
-    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].
-        """
-        all_images = x.view(-1, 3, *x.shape[-2:])
-
-        x = self.forward_features(all_images)
-        x = x.view(*x.shape[:2], -1)
-
-        return 1 - F.cosine_similarity(x[0], x[1], dim=1)

vit.py

@@ -179,9 +179,9 @@ class PatchEmbed(nn.Module):
         dynamic_pad: bool = False,
     ):
         super().__init__()
-        self.img_size = ensure_tuple(img_size)
-        self.patch_size = ensure_tuple(patch_size)
-        self.num_patches = (img_size // patch_size) ** 2
+        self.img_size = ensure_tuple(img_size, 2)
+        self.patch_size = ensure_tuple(patch_size, 2)
+        self.num_patches = (self.img_size[0] // self.patch_size[0]) * (self.img_size[1] // self.patch_size[1])
 
         self.dynamic_pad = dynamic_pad