make space happen

.editorconfig

@@ -0,0 +1,34 @@
+# http://editorconfig.org
+
+root = true
+
+[*]
+indent_style = space
+indent_size = 4
+trim_trailing_whitespace = true
+insert_final_newline = true
+charset = utf-8
+end_of_line = lf
+
+[*.bat]
+indent_style = tab
+end_of_line = crlf
+
+[*.{json,jsonc}]
+indent_style = space
+indent_size = 2
+
+[.vscode/*.{json,jsonc}]
+indent_style = space
+indent_size = 4
+
+[*.{yml,yaml,toml}]
+indent_style = space
+indent_size = 2
+
+[*.md]
+trim_trailing_whitespace = false
+
+[Makefile]
+indent_style = tab
+indent_size = 8

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,264 @@
+# Created by https://www.toptal.com/developers/gitignore/api/linux,windows,macos,visualstudiocode,python
+# Edit at https://www.toptal.com/developers/gitignore?templates=linux,windows,macos,visualstudiocode,python
+
+### Linux ###
+*~
+
+# temporary files which can be created if a process still has a handle open of a deleted file
+.fuse_hidden*
+
+# KDE directory preferences
+.directory
+
+# Linux trash folder which might appear on any partition or disk
+.Trash-*
+
+# .nfs files are created when an open file is removed but is still being accessed
+.nfs*
+
+### macOS ###
+# General
+.DS_Store
+.AppleDouble
+.LSOverride
+
+# Icon must end with two \r
+Icon
+
+
+# Thumbnails
+._*
+
+# Files that might appear in the root of a volume
+.DocumentRevisions-V100
+.fseventsd
+.Spotlight-V100
+.TemporaryItems
+.Trashes
+.VolumeIcon.icns
+.com.apple.timemachine.donotpresent
+
+# Directories potentially created on remote AFP share
+.AppleDB
+.AppleDesktop
+Network Trash Folder
+Temporary Items
+.apdisk
+
+### Python ###
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+### VisualStudioCode ###
+.vscode/*
+!.vscode/settings.json
+!.vscode/tasks.json
+!.vscode/launch.json
+!.vscode/extensions.json
+*.code-workspace
+
+# Local History for Visual Studio Code
+.history/
+
+### VisualStudioCode Patch ###
+# Ignore all local history of files
+.history
+.ionide
+
+### Windows ###
+# Windows thumbnail cache files
+Thumbs.db
+Thumbs.db:encryptable
+ehthumbs.db
+ehthumbs_vista.db
+
+# Dump file
+*.stackdump
+
+# Folder config file
+[Dd]esktop.ini
+
+# Recycle Bin used on file shares
+$RECYCLE.BIN/
+
+# Windows Installer files
+*.cab
+*.msi
+*.msix
+*.msm
+*.msp
+
+# Windows shortcuts
+*.lnk
+
+# End of https://www.toptal.com/developers/gitignore/api/linux,windows,macos,visualstudiocode,python
+
+# setuptools-scm _version file
+src/neurosis/_version.py
+
+# temp and misc
+/misc/
+/temp/
+
+# external repos
+/repos/
+
+# wandb
+/wandb/
+
+# outputs and such
+/logs/
+/cache/
+/outputs/
+/projects/
+
+# direnv
+.envrc
+.envrc.*
+
+# dotenv
+.env
+.env.*
+
+# temp files
+**/tmp_*.*
+**/*.tmp.*
+
+# but keep examples
+!*.example

.pre-commit-config.yaml

@@ -0,0 +1,27 @@
+# See https://pre-commit.com for more information
+ci:
+  autofix_prs: true
+  autoupdate_branch: "main"
+  autoupdate_commit_msg: "[pre-commit.ci] pre-commit autoupdate"
+  autoupdate_schedule: weekly
+
+repos:
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.2.0
+    hooks:
+      # Run the linter.
+      - id: ruff
+        types_or: [python, pyi, jupyter]
+        args: [--fix, --exit-non-zero-on-fix]
+      # Run the formatter.
+      - id: ruff-format
+        types_or: [python, pyi, jupyter]
+
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.5.0
+    hooks:
+      - id: trailing-whitespace
+        exclude_types:
+          - "markdown"
+      - id: end-of-file-fixer
+      - id: check-yaml

.vscode/settings.json

@@ -0,0 +1,94 @@
+{
+    "editor.insertSpaces": true,
+    "editor.tabSize": 4,
+    "files.trimTrailingWhitespace": true,
+    "editor.rulers": [100, 120],
+
+    "files.associations": {
+        "*.yaml": "yaml"
+    },
+    "files.exclude": {
+        "**/.git": true,
+        "**/.svn": true,
+        "**/.hg": true,
+        "**/CVS": true,
+        "**/.DS_Store": true,
+        "**/Thumbs.db": true,
+        "**/.ruff_cache": true,
+        "**/__pycache__": true,
+        "**/*.egg-info": true
+    },
+
+    "[shellscript]": {
+        "files.eol": "\n",
+        "editor.tabSize": 4,
+        "editor.detectIndentation": false
+    },
+
+    "[python]": {
+        "editor.wordBasedSuggestions": "off",
+        "editor.formatOnSave": true,
+        "editor.defaultFormatter": "charliermarsh.ruff",
+        "editor.codeActionsOnSave": {
+            "source.organizeImports": "always"
+        }
+    },
+    "python.analysis.include": ["./src", "./scripts", "./tests"],
+
+    "[json]": {
+        "editor.defaultFormatter": "esbenp.prettier-vscode",
+        "editor.detectIndentation": false,
+        "editor.formatOnSaveMode": "file",
+        "editor.formatOnSave": true,
+        "editor.tabSize": 2
+    },
+    "[jsonc]": {
+        "editor.defaultFormatter": "esbenp.prettier-vscode",
+        "editor.detectIndentation": false,
+        "editor.formatOnSaveMode": "file",
+        "editor.formatOnSave": true,
+        "editor.tabSize": 2
+    },
+
+    "[toml]": {
+        "editor.tabSize": 2,
+        "editor.detectIndentation": false,
+        "editor.formatOnSave": true,
+        "editor.formatOnSaveMode": "file",
+        "editor.defaultFormatter": "tamasfe.even-better-toml",
+        "editor.rulers": [80, 100]
+    },
+    "evenBetterToml.formatter.columnWidth": 88,
+
+    "[yaml]": {
+        "editor.detectIndentation": false,
+        "editor.tabSize": 2,
+        "editor.formatOnSave": true,
+        "editor.formatOnSaveMode": "file",
+        "diffEditor.ignoreTrimWhitespace": false,
+        "editor.defaultFormatter": "redhat.vscode-yaml"
+    },
+    "yaml.format.bracketSpacing": true,
+    "yaml.format.proseWrap": "preserve",
+    "yaml.format.singleQuote": false,
+    "yaml.format.printWidth": 110,
+
+    "[hcl]": {
+        "editor.detectIndentation": false,
+        "editor.formatOnSave": true,
+        "editor.formatOnSaveMode": "file",
+        "editor.defaultFormatter": "fredwangwang.vscode-hcl-format"
+    },
+
+    "[markdown]": {
+        "files.trimTrailingWhitespace": false
+    },
+
+    "css.lint.validProperties": ["dock", "content-align", "content-justify"],
+    "[css]": {
+        "editor.formatOnSave": true
+    },
+
+    "remote.autoForwardPorts": false,
+    "remote.autoForwardPortsSource": "process"
+}

LICENSE.md

@@ -0,0 +1,25 @@
+The MIT License (MIT)
+=====================
+
+Copyright © 2024 Andi Powers-Holmes <aholmes@omnom.net>
+
+Permission is hereby granted, free of charge, to any person
+obtaining a copy of this software and associated documentation
+files (the “Software”), to deal in the Software without
+restriction, including without limitation the rights to use,
+copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the
+Software is furnished to do so, subject to the following
+conditions:
+
+The above copyright notice and this permission notice shall be
+included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND,
+EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
+OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
+NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
+HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
+WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
+OTHER DEALINGS IN THE SOFTWARE.

app.py

@@ -0,0 +1,161 @@
+from os import getenv
+from typing import Optional
+
+import gradio as gr
+import torch
+from PIL import Image
+from torchvision.transforms import v2 as T
+
+from dreamsim import DreamsimBackbone, DreamsimEnsemble, DreamsimModel
+
+_ = torch.set_grad_enabled(False)
+torchdev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+torch.set_float32_matmul_precision("high")
+
+HF_TOKEN = getenv("HF_TOKEN", None)
+MODEL_REPO = "neggles/dreamsim"
+MODEL_VARIANTS: dict[str, str] = {
+    "Ensemble": "ensemble_vitb16",
+    "CLIP ViT-B/32": "clip_vitb32",
+    "OpenCLIP ViT-B/32": "open_clip_vitb32",
+    "DINO ViT-B/16": "dino_vitb16",
+}
+
+loaded_models: dict[str, Optional[DreamsimBackbone]] = {
+    "ensemble_vitb16": None,
+    "clip_vitb32": None,
+    "open_clip_vitb32": None,
+    "dino_vitb16": None,
+}
+
+
+def pil_ensure_rgb(image: Image.Image) -> Image.Image:
+    # convert to RGB/RGBA if not already (deals with palette images etc.)
+    if image.mode not in ["RGB", "RGBA"]:
+        image = image.convert("RGBA") if "transparency" in image.info else image.convert("RGB")
+    # convert RGBA to RGB with white background
+    if image.mode == "RGBA":
+        canvas = Image.new("RGBA", image.size, (255, 255, 255))
+        canvas.alpha_composite(image)
+        image = canvas.convert("RGB")
+    return image
+
+
+def pil_pad_square(
+    image: Image.Image,
+    fill: tuple[int, int, int] = (255, 255, 255),
+) -> Image.Image:
+    w, h = image.size
+    # get the largest dimension so we can pad to a square
+    px = max(image.size)
+    # pad to square with white background
+    canvas = Image.new("RGB", (px, px), fill)
+    canvas.paste(image, ((px - w) // 2, (px - h) // 2))
+    return canvas
+
+
+def load_model(variant: str) -> DreamsimBackbone:
+    global loaded_models
+
+    if variant in MODEL_VARIANTS:
+        # resolve the repo branch for the model variant
+        variant = MODEL_VARIANTS[variant]
+
+    match variant:
+        case "ensemble_vitb16":
+            if loaded_models[variant] is None:
+                model: DreamsimEnsemble = DreamsimEnsemble.from_pretrained(
+                    MODEL_REPO,
+                    token=HF_TOKEN,
+                    revision=variant,
+                )
+                model.do_resize = False
+                loaded_models[variant] = model
+
+        case "clip_vitb32" | "open_clip_vitb32" | "dino_vitb16":
+            if loaded_models[variant] is None:
+                model: DreamsimModel = DreamsimModel.from_pretrained(
+                    MODEL_REPO,
+                    token=HF_TOKEN,
+                    revision=variant,
+                )
+                model.do_resize = False
+                loaded_models[variant] = model
+
+        case _:
+            raise ValueError(f"Unknown model variant: {variant}")
+
+    return loaded_models[variant]
+
+
+def predict(
+    variant: str,
+    resize_to: Optional[int],
+    image_a: Image.Image,
+    image_b: Image.Image,
+):
+    # Load model
+    model: DreamsimModel | DreamsimEnsemble = load_model(variant)
+    model = model.eval().to(torchdev)
+
+    # yeet alpha, make white background
+    image_a, image_b = pil_ensure_rgb(image_a), pil_ensure_rgb(image_b)
+    # pad to square
+    image_a, image_b = pil_pad_square(image_a), pil_pad_square(image_b)
+
+    # Resize images, if necessary
+    if resize_to is not None:
+        image_a.thumbnail((resize_to, resize_to), resample=Image.Resampling.BICUBIC)
+        image_b.thumbnail((resize_to, resize_to), resample=Image.Resampling.BICUBIC)
+
+    # Preprocess images
+    transforms = T.Compose([T.ToImage(), T.ToDtype(torch.float32, scale=True)])
+    batch = torch.stack([transforms(image_a).unsqueeze(0), transforms(image_b).unsqueeze(0)], dim=0)
+
+    loss = model(batch.to(model.device, model.dtype)).cpu().item()
+    score = 1.0 - loss
+    return score, variant
+
+
+def main():
+    with gr.Blocks(title="DreamSIM Perceptual Similarity") as demo:
+        with gr.Row():
+            with gr.Column():
+                img_input = gr.Image(label="Input", type="pil", image_mode="RGB", scale=1)
+            with gr.Column():
+                img_target = gr.Image(label="Target", type="pil", image_mode="RGB", scale=1)
+        with gr.Row(equal_height=True):
+            with gr.Column():
+                variant = gr.Radio(
+                    choices=list(MODEL_VARIANTS.keys()), label="Model Variant", value="Ensemble"
+                )
+                resize_to = gr.Dropdown(label="Resize To", choices=[224, 384, 512, None], value=224)
+            with gr.Column():
+                score = gr.Number(label="Similarity Score", precision=8, minimum=0, maximum=1)
+                variant_out = gr.Textbox(label="Variant", interactive=False)
+                with gr.Row():
+                    clear = gr.ClearButton(
+                        components=[img_input, img_target, score], variant="secondary", size="lg"
+                    )
+                    submit = gr.Button(value="Submit", variant="primary", size="lg")
+
+        submit.click(
+            predict,
+            inputs=[variant, resize_to, img_input, img_target],
+            outputs=[score, variant_out],
+            api_name=False,
+        )
+        examples = gr.Examples(
+            [
+                ["examples/img_a_1.png", "examples/ref_1.png", "Ensemble", 224],
+                ["examples/img_b_1.png", "examples/ref_1.png", "Ensemble", 224],
+            ],
+            inputs=[img_input, img_target, variant, resize_to],
+        )
+
+    demo.queue(max_size=10)
+    demo.launch()
+
+
+if __name__ == "__main__":
+    main()

dreamsim/.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

dreamsim/README.md

@@ -0,0 +1,10 @@
+---
+license: mit
+datasets:
+- PerceptionEval/DreamSim
+library_name: transformers
+---
+
+dreamsim! now in quasi-transformers quasi-diffusers form.
+
+this probably won't work for you! but if it works for what i'm experimenting with, i'll try to get it upstreamed.

dreamsim/__init__.py

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

dreamsim/common.py

@@ -0,0 +1,38 @@
+from typing import Callable
+
+import torch
+from torch import Tensor, nn
+from torch.nn import functional as F
+
+
+def ensure_tuple(val: int | tuple[int, ...], n: int = 2) -> tuple[int, ...]:
+    if isinstance(val, int):
+        return (val,) * n
+    elif len(val) != n:
+        raise ValueError(f"Expected a tuple of {n} values, but got {len(val)}: {val}")
+    return val
+
+
+def use_fused_attn():
+    if hasattr(F, "scaled_dot_product_attention"):
+        return True
+    return False
+
+
+class QuickGELU(nn.Module):
+    """
+    Applies GELU approximation that is fast but somewhat inaccurate. See: https://github.com/hendrycks/GELUs
+    """
+
+    def forward(self, input: Tensor) -> Tensor:
+        return input * torch.sigmoid(1.702 * input)
+
+
+def get_act_layer(name: str) -> Callable[[], nn.Module]:
+    match name:
+        case "gelu":
+            return nn.GELU
+        case "quick_gelu":
+            return QuickGELU
+        case _:
+            raise ValueError(f"Activation layer {name} not supported.")

dreamsim/model.py

@@ -0,0 +1,188 @@
+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)
+
+        # these expect to always receive a batch, and will return a batch
+        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

dreamsim/utils.py

@@ -0,0 +1,160 @@
+"""
+Functions in this file are courtesty of @ashen-sensored on GitHub - thankyou so much! <3
+
+Used to merge DreamSim LoRA weights into the base ViT models manually, so we don't need
+to use an ancient version of PeFT that is no longer supported (and kind of broken)
+"""
+import logging
+from os import PathLike
+from pathlib import Path
+
+import torch
+from safetensors.torch import load_file
+from torch import Tensor, nn
+
+from .model import DreamsimModel
+
+logger = logging.getLogger(__name__)
+
+
+@torch.no_grad()
+def calculate_merged_weight(
+    lora_a: Tensor,
+    lora_b: Tensor,
+    base: Tensor,
+    scale: float,
+    qkv_switches: list[bool],
+) -> Tensor:
+    n_switches = len(qkv_switches)
+    n_groups = sum(qkv_switches)
+
+    qkv_mask = torch.tensor(qkv_switches, dtype=torch.bool).reshape(len(qkv_switches), -1)
+    qkv_mask = qkv_mask.broadcast_to((-1, base.shape[0] // n_switches)).reshape(-1)
+
+    lora_b = lora_b.squeeze()
+    delta_w = base.new_zeros(lora_b.shape[0], base.shape[1])
+
+    grp_in_ch = lora_a.shape[0] // n_groups
+    grp_out_ch = lora_b.shape[0] // n_groups
+    for i in range(n_groups):
+        islice = slice(i * grp_in_ch, (i + 1) * grp_in_ch)
+        oslice = slice(i * grp_out_ch, (i + 1) * grp_out_ch)
+        delta_w[oslice, :] = lora_b[oslice, :] @ lora_a[islice, :]
+
+    delta_w_full = base.new_zeros(base.shape)
+    delta_w_full[qkv_mask, :] = delta_w
+
+    merged = base + scale * delta_w_full
+    return merged.to(base)
+
+
+@torch.no_grad()
+def merge_dreamsim_lora(
+    base_model: nn.Module,
+    lora_path: PathLike,
+    torch_device: torch.device | str = torch.device("cpu"),
+):
+    lora_path = Path(lora_path)
+    # make sure model is on device
+    base_model = base_model.eval().requires_grad_(False).to(torch_device)
+
+    # load the lora
+    if lora_path.suffix.lower() in [".pt", ".pth", ".bin"]:
+        lora_sd = torch.load(lora_path, map_location=torch_device, weights_only=True)
+    elif lora_path.suffix.lower() == ".safetensors":
+        lora_sd = load_file(lora_path)
+    else:
+        raise ValueError(f"Unsupported file extension '{lora_path.suffix}'")
+
+    # these loras were created by a cursed PEFT version, okay? so we have to do some crimes.
+    group_prefix = "base_model.model.base_model.model.model."
+    # get all lora weights for qkv layers, stripping the insane prefix
+    group_weights = {k.replace(group_prefix, ""): v for k, v in lora_sd.items() if k.startswith(group_prefix)}
+    # strip ".lora_X.weight" from keys to match against base model keys
+    group_layers = set([k.rsplit(".", 2)[0] for k in group_weights.keys()])
+
+    base_weights = base_model.state_dict()
+    for key in [x for x in base_weights.keys() if "attn.qkv.weight" in x]:
+        param_name = key.rsplit(".", 1)[0]
+        if param_name not in group_layers:
+            logger.warning(f"QKV param '{param_name}' not found in lora weights")
+            continue
+        new_weight = calculate_merged_weight(
+            group_weights[f"{param_name}.lora_A.weight"],
+            group_weights[f"{param_name}.lora_B.weight"],
+            base_weights[key],
+            0.5 / 16,
+            [True, False, True],
+        )
+        base_weights[key] = new_weight
+
+    base_model.load_state_dict(base_weights)
+    return base_model.requires_grad_(False)
+
+
+def remap_clip(state_dict: dict[str, Tensor], variant: str) -> dict[str, Tensor]:
+    """Remap keys from the original DreamSim checkpoint to match new model structure."""
+
+    def prepend_extractor(state_dict: dict[str, Tensor]) -> dict[str, Tensor]:
+        if variant.endswith("single"):
+            return {f"extractor.{k}": v for k, v in state_dict.items()}
+        return state_dict
+
+    if "clip" not in variant:
+        return prepend_extractor(state_dict)
+
+    if "patch_embed.proj.bias" in state_dict:
+        _ = state_dict.pop("patch_embed.proj.bias", None)
+    if "pos_drop.weight" in state_dict:
+        state_dict["norm_pre.weight"] = state_dict.pop("pos_drop.weight")
+        state_dict["norm_pre.bias"] = state_dict.pop("pos_drop.bias")
+    if "head.weight" in state_dict and "head.bias" not in state_dict:
+        state_dict["head.bias"] = torch.zeros(state_dict["head.weight"].shape[0])
+
+    return prepend_extractor(state_dict)
+
+
+def convert_dreamsim_single(
+    ckpt_path: PathLike,
+    variant: str,
+    ensemble: bool = False,
+) -> DreamsimModel:
+    ckpt_path = Path(ckpt_path)
+    if ckpt_path.exists():
+        if ckpt_path.is_dir():
+            ckpt_path = ckpt_path.joinpath("ensemble" if ensemble else variant)
+            ckpt_path = ckpt_path.joinpath(f"{variant}_merged.safetensors")
+
+    # defaults are for dino, overridden as needed below
+    patch_size = 16
+    layer_norm_eps = 1e-6
+    pre_norm = False
+    act_layer = "gelu"
+
+    match variant:
+        case "open_clip_vitb16" | "open_clip_vitb32" | "clip_vitb16" | "clip_vitb32":
+            patch_size = 32 if "b32" in variant else 16
+            layer_norm_eps = 1e-5
+            pre_norm = True
+            img_mean = (0.48145466, 0.4578275, 0.40821073)
+            img_std = (0.26862954, 0.26130258, 0.27577711)
+            act_layer = "quick_gelu" if variant.startswith("clip_") else "gelu"
+        case "dino_vitb16":
+            img_mean = (0.485, 0.456, 0.406)
+            img_std = (0.229, 0.224, 0.225)
+        case _:
+            raise NotImplementedError(f"Unsupported model variant '{variant}'")
+
+    model: DreamsimModel = DreamsimModel(
+        image_size=224,
+        patch_size=patch_size,
+        layer_norm_eps=layer_norm_eps,
+        pre_norm=pre_norm,
+        act_layer=act_layer,
+        img_mean=img_mean,
+        img_std=img_std,
+    )
+    state_dict = load_file(ckpt_path, device="cpu")
+    state_dict = remap_clip(state_dict)
+    model.extractor.load_state_dict(state_dict)
+    return model

dreamsim/vit.py

@@ -0,0 +1,375 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Mostly copy-paste from timm library.
+https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py
+"""
+import math
+from functools import partial
+from typing import Callable, Final, Optional, Sequence
+
+import torch
+from torch import Tensor, nn
+from torch.nn import functional as F
+
+from .common import ensure_tuple, get_act_layer, use_fused_attn
+
+
+def vit_weights_init(module: nn.Module) -> None:
+    if isinstance(module, nn.Linear):
+        nn.init.trunc_normal_(module.weight, std=0.02)
+        if module.bias is not None:
+            nn.init.zeros_(module.bias)
+    elif isinstance(module, nn.LayerNorm):
+        nn.init.ones_(module.weight)
+        nn.init.zeros_(module.bias)
+
+
+class DropPath(nn.Module):
+    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
+
+    def __init__(self, drop_prob: float = 0.0, scale_by_keep: bool = True):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+        self.scale_by_keep = scale_by_keep
+
+    def forward(self, x: Tensor) -> Tensor:
+        if self.drop_prob == 0 or not self.training:
+            return x
+        keep_prob = 1 - self.drop_prob
+        shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
+        random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
+        if keep_prob > 0.0 and self.scale_by_keep:
+            random_tensor.div_(keep_prob)
+        return x * random_tensor
+
+    def extra_repr(self):
+        return f"drop_prob={self.drop_prob:0.3f}"
+
+
+class Mlp(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        hidden_features: Optional[int] = None,
+        out_features: Optional[int] = None,
+        act_layer: Callable[[], nn.Module] = nn.GELU,
+        drop: float = 0.0,
+    ):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop) if drop > 0.0 else nn.Identity()
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class Attention(nn.Module):
+    fused_attn: Final[bool]
+
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int = 8,
+        qkv_bias: bool = False,
+        qk_scale: Optional[float] = None,
+        attn_drop: float = 0.0,
+        proj_drop: float = 0.0,
+    ):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.scale = qk_scale or self.head_dim**-0.5
+        self.fused_attn = use_fused_attn()
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop) if attn_drop > 0.0 else nn.Identity()
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop) if proj_drop > 0.0 else nn.Identity()
+
+    def forward(self, x: Tensor) -> Tensor:
+        B, N, C = x.shape
+        qkv: Tensor = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv.unbind(0)
+
+        if self.fused_attn:
+            dropout_p = getattr(self.attn_drop, "p", 0.0) if self.training else 0.0
+            x = F.scaled_dot_product_attention(q, k, v, dropout_p=dropout_p)
+        else:
+            q = q * self.scale
+            attn = q @ k.transpose(-2, -1)
+            attn = attn.softmax(dim=-1)
+            attn = self.attn_drop(attn)
+            x = attn @ v
+
+        x = x.transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class Block(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        qkv_bias: bool = False,
+        drop: float = 0.0,
+        attn_drop: float = 0.0,
+        drop_path: float = 0.0,
+        act_layer: Callable[[], nn.Module] = nn.GELU,
+        norm_layer: Callable[[], nn.Module] = nn.LayerNorm,
+    ):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            attn_drop=attn_drop,
+            proj_drop=drop,
+        )
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(
+            in_features=dim,
+            hidden_features=mlp_hidden_dim,
+            act_layer=act_layer,
+            drop=drop,
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = x + self.drop_path(self.attn(self.norm1(x)))
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+        return x
+
+
+class PatchEmbed(nn.Module):
+    """Image to Patch Embedding"""
+
+    def __init__(
+        self,
+        img_size: int | tuple[int, int] = 224,
+        patch_size: int | tuple[int, int] = 16,
+        in_chans: int = 3,
+        embed_dim: int = 768,
+        bias: bool = True,
+        dynamic_pad: bool = False,
+    ):
+        super().__init__()
+        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
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size, bias=bias)
+
+    def forward(self, x: Tensor) -> Tensor:
+        _, _, H, W = x.shape
+        if self.dynamic_pad:
+            pad_h = (self.patch_size[0] - H % self.patch_size[0]) % self.patch_size[0]
+            pad_w = (self.patch_size[1] - W % self.patch_size[1]) % self.patch_size[1]
+            x = F.pad(x, (0, pad_w, 0, pad_h))
+        x = self.proj(x)
+        x = x.flatten(2).transpose(1, 2)  # NCHW -> NLC
+        return x
+
+
+class VisionTransformer(nn.Module):
+    """Vision Transformer"""
+
+    def __init__(
+        self,
+        img_size: int | tuple[int, int] = 224,
+        patch_size: int | tuple[int, int] = 16,
+        in_chans: int = 3,
+        num_classes: int = 0,
+        embed_dim: int = 768,
+        depth: int = 12,
+        num_heads: int = 12,
+        mlp_ratio: float = 4.0,
+        qkv_bias: bool = False,
+        pre_norm: bool = False,
+        drop_rate: float = 0.0,
+        attn_drop_rate: float = 0.0,
+        drop_path_rate: float = 0.0,
+        norm_layer: Callable[[], nn.Module] = nn.LayerNorm,
+        act_layer: Callable[[], nn.Module] = nn.GELU,
+        skip_init: bool = False,
+        dynamic_pad: bool = False,
+        **kwargs,
+    ):
+        super().__init__()
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.num_classes = num_classes
+        self.num_features = self.embed_dim = embed_dim
+        self.depth = depth
+
+        self.patch_embed = PatchEmbed(
+            img_size=img_size,
+            patch_size=patch_size,
+            in_chans=in_chans,
+            embed_dim=embed_dim,
+            bias=not pre_norm,  # disable bias if pre-norm is used (e.g. CLIP)
+            dynamic_pad=dynamic_pad,
+        )
+        num_patches = self.patch_embed.num_patches
+        embed_len = num_patches + 1  # num_patches + 1 for the [CLS] token
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.pos_embed = nn.Parameter(torch.zeros(1, embed_len, embed_dim))
+        self.pos_drop = nn.Dropout(p=drop_rate) if drop_rate > 0.0 else nn.Identity()
+        self.norm_pre = norm_layer(embed_dim) if pre_norm else nn.Identity()
+
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, self.depth)]  # stochastic depth decay rule
+        self.blocks: list[Block] = nn.ModuleList(
+            [
+                Block(
+                    dim=embed_dim,
+                    num_heads=num_heads,
+                    mlp_ratio=mlp_ratio,
+                    qkv_bias=qkv_bias,
+                    drop=drop_rate,
+                    attn_drop=attn_drop_rate,
+                    drop_path=dpr[i],
+                    act_layer=act_layer,
+                    norm_layer=norm_layer,
+                )
+                for i in range(self.depth)
+            ]
+        )
+        self.norm = norm_layer(embed_dim)
+
+        # Classifier head
+        self.head = nn.Linear(embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+
+        if not skip_init:
+            self.reset_parameters()
+
+    def reset_parameters(self):
+        nn.init.trunc_normal_(self.cls_token, std=0.02)
+        nn.init.trunc_normal_(self.pos_embed, std=0.02)
+        self.apply(vit_weights_init)
+
+    def interpolate_pos_encoding(self, x: Tensor, w: Tensor, h: Tensor) -> Tensor:
+        npatch = x.shape[1] - 1
+        N = self.pos_embed.shape[1] - 1
+        if npatch == N and w == h:
+            return self.pos_embed
+        class_pos_embed = self.pos_embed[:, 0]
+        patch_pos_embed = self.pos_embed[:, 1:]
+        dim = x.shape[-1]
+        w0 = w // self.patch_embed.patch_size[0]
+        h0 = h // self.patch_embed.patch_size[0]
+        # we add a small number to avoid floating point error in the interpolation
+        # see discussion at https://github.com/facebookresearch/dino/issues/8
+        w0, h0 = w0 + 0.1, h0 + 0.1
+        patch_pos_embed = nn.functional.interpolate(
+            patch_pos_embed.reshape(1, int(math.sqrt(N)), int(math.sqrt(N)), dim).permute(0, 3, 1, 2),
+            scale_factor=(w0 / math.sqrt(N), h0 / math.sqrt(N)),
+            mode="bicubic",
+        )
+        if int(w0) != patch_pos_embed.shape[-2] or int(h0) != patch_pos_embed.shape[-1]:
+            raise ValueError("Error in positional encoding interpolation.")
+        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+        return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1)
+
+    def prepare_tokens(self, x: Tensor) -> Tensor:
+        B, _, W, H = x.shape
+        x = self.patch_embed(x)  # patch linear embedding
+
+        # add the [CLS] token to the embed patch tokens
+        cls_tokens = self.cls_token.expand(B, -1, -1)
+        x = torch.cat((cls_tokens, x), dim=1)
+
+        # add positional encoding to each token
+        x = x + self.interpolate_pos_encoding(x, W, H)
+
+        return self.pos_drop(x)
+
+    def forward(self, x: Tensor, norm: bool = True) -> Tensor:
+        x = self.forward_features(x, norm=norm)
+        x = self.forward_head(x)
+        return x
+
+    def forward_features(self, x: Tensor, norm: bool = True) -> Tensor:
+        x = self.prepare_tokens(x)
+        x = self.norm_pre(x)
+        for blk in self.blocks:
+            x = blk(x)
+        if norm:
+            x = self.norm(x)
+        return x[:, 0]
+
+    def forward_head(self, x: Tensor) -> Tensor:
+        x = self.head(x)
+        return x
+
+    def get_intermediate_layers(
+        self,
+        x: Tensor,
+        n: int | Sequence[int] = 1,
+        norm: bool = True,
+    ) -> list[Tensor]:
+        # we return the output tokens from the `n` last blocks
+        outputs = []
+        layer_indices = set(range(self.depth - n, self.depth) if isinstance(n, int) else n)
+        x = self.prepare_tokens(x)
+        x = self.norm_pre(x)
+
+        for idx, blk in enumerate(self.blocks):
+            x = blk(x)
+            if idx in layer_indices:
+                outputs.append(x)
+        if norm:
+            outputs = [self.norm(x) for x in outputs]
+        return outputs
+
+
+def vit_base_dreamsim(
+    patch_size: int = 16,
+    layer_norm_eps: float = 1e-6,
+    num_classes: int = 512,
+    act_layer: str | Callable[[], nn.Module] = "gelu",
+    **kwargs,
+):
+    if isinstance(act_layer, str):
+        act_layer = get_act_layer(act_layer)
+
+    model = VisionTransformer(
+        patch_size=patch_size,
+        num_classes=num_classes,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        mlp_ratio=4,
+        qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=layer_norm_eps),
+        act_layer=act_layer,
+        **kwargs,
+    )
+    return model

pyproject.toml

@@ -0,0 +1,98 @@
+[project]
+name = "dreamsim-space"
+version = "0.1.0"
+authors = [
+  { name = "Stephanie Fu" },
+  { name = "Netanel Tamir" },
+  { name = "Shobhita Sundaram" },
+  { name = "Lucy Chai" },
+  { name = "Richard Zhang" },
+  { name = "Tali Dekel" },
+  { name = "Phillip Isola" },
+]
+maintainers = [
+  { name = "Andi Powers-Holmes", email = "aholmes@omnom.net" },
+]
+description = "DreamSim Gradio Space"
+readme = "README.md"
+requires-python = ">=3.9, <3.11"
+keywords = [
+  "deep-learning",
+  "machine-learning",
+  "pytorch",
+]
+license = { file = "LICENSE.md" }
+classifiers = [
+  "Programming Language :: Python :: 3",
+  "License :: OSI Approved :: MIT License",
+]
+dependencies = [
+  "accelerate",
+  "diffusers",
+  "gradio >=4.19.1, < 5.0.0",
+  "numpy",
+  "pandas",
+  "Pillow",
+  "PyYAML",
+  "safetensors",
+  "simple-parsing >= 0.1.0",
+  "torch",
+  "torchvision",
+  "transformers",
+  'xformers; sys_platform != "win32"',
+]
+
+[project.urls]
+Repository = "https://huggingface.co/spaces/neggles/dreamsim"
+
+[project.optional-dependencies]
+dev = [
+  "ruff >=0.0.289",
+  "setuptools-scm >= 8.0.0",
+  "pre-commit >= 3.0.0",     # remember to run `pre-commit install` after installing
+  "tabulate >= 0.8.9",       # for inductor log prettyprinting
+]
+all = [
+  "dreamsim-space[dev]",
+]
+
+[build-system]
+build-backend = "setuptools.build_meta"
+requires = ["setuptools>=64", "wheel"]
+
+[tool.setuptools.packages.find]
+namespaces = true
+where = ["."]
+include = ["dreamsim"]
+
+
+[tool.ruff]
+line-length = 110
+target-version = "py310"
+extend-exclude = ["/usr/lib/*"]
+
+[tool.ruff.lint]
+ignore = [
+  "F841", # local variable assigned but never used
+  "F842", # local variable annotated but never used
+  "E501", # line too long - will be fixed in format
+]
+
+[tool.ruff.format]
+quote-style = "double"
+indent-style = "space"
+line-ending = "auto"
+skip-magic-trailing-comma = false
+docstring-code-format = true
+
+[tool.ruff.lint.isort]
+combine-as-imports = true
+force-wrap-aliases = true
+known-local-folder = ["dreamsim"]
+known-first-party = ["dreamsim"]
+
+
+[tool.pyright]
+include = ["src/**"]
+exclude = ["/usr/lib/**"]
+stubPath = "./typings"

requirements.txt

@@ -0,0 +1 @@
+-e .[all]