SkinMap: 12-teacher ensemble + predict_meta (validated release)

.gitattributes

@@ -0,0 +1,8 @@
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.faiss filter=lfs diff=lfs merge=lfs -text
+*.bin.index filter=lfs diff=lfs merge=lfs -text
+analysis/**/*.bin filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,97 @@
+---
+license: cc-by-nc-4.0
+library_name: transformers
+tags:
+  - dermatology
+  - medical-imaging
+  - embeddings
+  - image-feature-extraction
+  - clip
+pipeline_tag: image-feature-extraction
+---
+
+# A Global Atlas of Digital Dermatology to Map Innovation and Disparities
+
+![SkinMap method](https://fabiangroeger96.github.io/assets/img/projects/skinmap/Method.svg)
+
+SkinMap is a dermatology image/text embedding model. It fuses **12 frozen teacher
+encoders** (multi-modal and self-supervised models) through a trained projector
+into a single **1024-d shared image/text space**, that can be used to infer
+missing metadata, solve downstream tasks, and get relational information about the
+samples.
+
+**Compute.** This is the faithful full ensemble: 12 backbones load into memory at
+once and need roughly **16 GB of GPU VRAM**. It runs on CPU but slowly.
+
+## Usage
+
+```python
+from transformers import AutoModel
+
+model = AutoModel.from_pretrained(
+    "Digital-Dermatology/SkinMap",
+    trust_remote_code=True,
+    device="cuda",            # "cpu" also works
+)
+
+# 1) Encode an image into the 1024-d SkinMap space (L2-normalized)
+img_emb = model.encode_image("lesion.jpg")            # np.ndarray, shape (1024,)
+
+# 2) Encode a text query into the same space
+txt_emb = model.encode_text("melanoma on the back")   # np.ndarray, shape (1024,)
+
+# 3) Predict clinical metadata from the image
+meta = model.predict_meta("lesion.jpg")
+# {'fitzpatrick': 3, 'fitzpatrick_grouped': '3-4', 'age': 42.1,
+#  'gender': 'female', 'origin': ..., 'body_region': ..., ...}
+
+import numpy as np
+sim = float(np.dot(img_emb, txt_emb))                 # cosine in the shared space
+```
+
+`encode_image` accepts a file path, raw bytes, or a `PIL.Image`. `encode_text`
+expands the query into clinical prompt templates and averages the results; pass
+`templates=["{}"]` to encode the raw text as-is.
+
+### Metadata prediction
+
+`predict_meta` applies bundled linear probes to the SkinMap embedding to estimate
+Fitzpatrick skin type, age, sex, geographic origin, body region and more. By
+default it uses **global** probes, which reproduce the demographic estimates
+reported in the paper. If you know the imaging modality, pass it to switch the
+Fitzpatrick prediction to a **modality-specific** probe (more accurate when the
+modality is known); the other attributes stay global:
+
+```python
+meta = model.predict_meta("lesion.jpg", modality="dermoscopy")  # clinical | dermoscopy | TBP
+```
+
+`predict_meta` also accepts a precomputed 1024-d vector instead of an image, and a
+subset via `attributes=[...]`.
+
+### Nearest-neighbor search (optional)
+
+If this repo was built with an atlas index bundled, you can retrieve neighbors
+directly:
+
+```python
+indices, distances = model.search("lesion.jpg", k=10)
+```
+
+Without a bundled index, `search()` raises a clear error. Encoding still works.
+
+## What's inside
+
+| Component | Role |
+|---|---|
+| 9 multi-modal teachers (CLIP & SigLIP; ViT-L/14 & ViT-B/32) | image + text features |
+| 3 self-supervised teachers (DINO, iBOT, MAE; ViT-B/16) | dermatology image features |
+| Per-teacher whitening (`whitening_stats.npz`) | decorrelate before fusion |
+| Trained projector (`projector_model.pth`) | fuse to 1024-d image/text heads |
+| Metadata probes (`probes/`) | linear probes for `predict_meta` |
+
+## Licensing
+
+SkinMap is released under the Creative Commons Attribution-NonCommercial 4.0
+International license, for research use only. **Not a medical device. Not for
+clinical decision-making.**

config.json

@@ -0,0 +1,13 @@
+{
+  "model_type": "skinmap",
+  "architectures": [
+    "SkinMapModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_skinmap.SkinMapConfig",
+    "AutoModel": "modeling_skinmap.SkinMapModel"
+  },
+  "embedding_dim": 1024,
+  "pipeline_config": "weights/embedding_pipeline_config.json",
+  "probes_dir": "probes"
+}

configuration_skinmap.py

@@ -0,0 +1,30 @@
+"""HuggingFace config for the packaged SkinMap multi-teacher embedding model."""
+
+from transformers import PretrainedConfig
+
+
+class SkinMapConfig(PretrainedConfig):
+    """Minimal config for SkinMap.
+
+    SkinMap is not a single set of HF weights but an ensemble of teacher encoders
+    plus a trained projector, whose loading is driven by a bundled pipeline-config
+    JSON. This config only points at that JSON and records the output
+    dimensionality. The heavy lifting happens in `SkinMapModel`.
+    """
+
+    model_type = "skinmap"
+
+    def __init__(
+        self,
+        embedding_dim: int = 1024,
+        pipeline_config: str = "weights/embedding_pipeline_config.json",
+        probes_dir: str = "probes",
+        **kwargs,
+    ):
+        self.embedding_dim = embedding_dim
+        # Paths are relative to the repo snapshot root.
+        self.pipeline_config = pipeline_config
+        # Directory of bundled metadata probes (enables predict_meta); may be
+        # absent in an embedding-only package.
+        self.probes_dir = probes_dir
+        super().__init__(**kwargs)

meta_predictor.py

@@ -0,0 +1,102 @@
+"""Metadata prediction from SkinMap embeddings via bundled linear probes.
+
+Default uses the global probes, which reproduce the demographic estimates reported
+in the paper (e.g. Fitzpatrick V--VI ~ 11%). Passing ``modality`` switches the
+Fitzpatrick prediction to a modality-specific probe (clinical vs. dermoscopy/TBP),
+which is more accurate when the imaging modality is known.
+"""
+from __future__ import annotations
+
+import json
+import pickle
+from pathlib import Path
+from typing import Dict, List, Optional, Union
+
+import numpy as np
+
+# imaging modality -> probe group used for the modality-specific Fitzpatrick model
+_MODALITY_GROUP = {"clinical": "clinical", "dermoscopy": "derm", "tbp": "derm"}
+_FST_GROUP = {1.0: "1-2", 2.0: "1-2", 3.0: "3-4", 4.0: "3-4", 5.0: "5-6", 6.0: "5-6"}
+
+
+class MetaPredictor:
+    """Loads a bundled probe directory and predicts metadata from embeddings."""
+
+    def __init__(self, probes_dir: Union[str, Path]):
+        self.dir = Path(probes_dir)
+        manifest = self.dir / "manifest.json"
+        if not manifest.exists():
+            raise FileNotFoundError(f"probe manifest not found at {manifest}")
+        self.manifest = json.loads(manifest.read_text())
+        self._cache: Dict[str, object] = {}
+
+    def _load(self, name: str):
+        if name not in self._cache:
+            with open(self.dir / f"{name}.pkl", "rb") as f:
+                self._cache[name] = pickle.load(f)
+        return self._cache[name]
+
+    @property
+    def attributes(self) -> List[str]:
+        return list(self.manifest["attributes"].keys())
+
+    def predict(
+        self,
+        embedding: np.ndarray,
+        modality: Optional[str] = None,
+        attributes: Optional[List[str]] = None,
+    ) -> Dict[str, object]:
+        """Predict metadata for one embedding (1024,) or a batch (N, 1024).
+
+        Args:
+            embedding: L2-normalized SkinMap embedding(s).
+            modality: optional {'clinical','dermoscopy','TBP'}; switches Fitzpatrick
+                to the modality-specific probe. Other attributes stay global.
+            attributes: subset to predict (default: all available).
+        Returns:
+            dict attribute -> scalar (single input) or np.ndarray (batch).
+        """
+        emb = np.asarray(embedding, dtype=np.float32)
+        single = emb.ndim == 1
+        if single:
+            emb = emb[None, :]
+        grp = _MODALITY_GROUP.get(str(modality).lower()) if modality else None
+        attrs = attributes or self.attributes
+        out: Dict[str, object] = {}
+        for a in attrs:
+            spec = self.manifest["attributes"].get(a)
+            if spec is None:
+                continue
+            if spec["task"] == "regression":
+                pred = self._load(f"model_{a}").predict(emb).astype(float)
+                clip = spec.get("clip")
+                if clip:
+                    pred = np.clip(pred, clip[0], clip[1])
+                vals = pred
+            else:
+                if a == "fitzpatrick" and grp and spec.get("modality_specific"):
+                    clf = self._load(f"model_fitzpatrick__{grp}")
+                    le = self._load("label_encoder_fitzpatrick_modspecific")
+                else:
+                    clf = self._load(f"model_{a}")
+                    le = self._load(f"label_encoder_{a}")
+                vals = le.inverse_transform(clf.predict(emb))
+                try:
+                    vals = vals.astype(float)
+                except (ValueError, TypeError):
+                    pass
+            out[a] = vals
+
+        # derived: grouped Fitzpatrick (I-II / III-IV / V-VI)
+        if "fitzpatrick" in out:
+            fst = np.asarray(out["fitzpatrick"], dtype=float)
+            out["fitzpatrick_grouped"] = np.array(
+                [_FST_GROUP.get(round(v), "unknown") for v in np.atleast_1d(fst)]
+            )
+
+        if single:
+            out = {
+                k: (v[0].item() if hasattr(v[0], "item") else v[0])
+                for k, v in out.items()
+            }
+        return out

modeling_skinmap.py

@@ -0,0 +1,301 @@
+"""HuggingFace wrapper for the packaged SkinMap multi-teacher embedding model.
+
+SkinMap fuses ~12 frozen teacher encoders (CLIP fine-tunes plus dermatology SSL
+backbones) through a trained projector into a single 1024-d image/text space.
+The public API is ``encode_image``, ``encode_text`` and ``search``.
+
+SkinMap is not a single state-dict, so the standard PreTrainedModel weight
+loading is not used. ``from_pretrained`` resolves the repo snapshot, puts the
+bundled ``skinmap_runtime`` package on ``sys.path``, and builds the underlying
+``CombinedEmbeddingPipeline`` from the bundled pipeline-config JSON. Top-level
+imports stay light (stdlib, torch, transformers, hub) because the snapshot
+location, and therefore ``skinmap_runtime``, is only known at call time.
+"""
+
+from __future__ import annotations
+
+import io
+import sys
+import warnings
+from pathlib import Path
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+from transformers import PreTrainedModel
+
+from .configuration_skinmap import SkinMapConfig
+from .meta_predictor import MetaPredictor
+
+ImageInput = Union[str, "Path", bytes, "object"]  # path | bytes | PIL.Image.Image
+
+
+class SkinMapModel(PreTrainedModel):
+    config_class = SkinMapConfig
+    # SkinMap manages its own weights from the bundled files. There is no HF
+    # state-dict to load or split for device_map / low-cpu-mem paths.
+    _supports_device_map = False
+
+    def __init__(self, config: SkinMapConfig):
+        super().__init__(config)
+        # Keep a trivial buffer so .device / .to() behave on a Module with no
+        # registered parameters of its own.
+        self.register_buffer("_anchor", torch.zeros(1), persistent=False)
+        self._pipeline = None
+        self._predictor = None
+        self._skinmap_device = torch.device("cpu")
+
+    # ------------------------------------------------------------------ load
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_model_name_or_path,
+        *model_args,
+        device: Optional[str] = None,
+        **kwargs,
+    ):
+        """Resolve the repo snapshot and build the SkinMap pipeline.
+
+        Accepts a local directory or a (private) hub repo id. Hub/auth kwargs
+        (`token`, `revision`, `cache_dir`) are forwarded to `snapshot_download`.
+        Standard weight-loading kwargs that SkinMap cannot honor (`torch_dtype`,
+        `device_map`, `low_cpu_mem_usage`) are warned about rather than silently
+        ignored. Use `device=...` to choose the device.
+        """
+        token = kwargs.pop("token", kwargs.pop("use_auth_token", None))
+        revision = kwargs.pop("revision", None)
+        cache_dir = kwargs.pop("cache_dir", None)
+        config = kwargs.pop("config", None)
+
+        # SkinMap loads its own weights from the bundled files, so the standard
+        # dtype/placement/low-mem machinery does not apply. Warn instead of
+        # silently dropping these. A user passing torch_dtype=fp16 to fit the
+        # ensemble in limited VRAM would otherwise get fp32 and OOM with no signal.
+        for unsupported in ("torch_dtype", "device_map", "low_cpu_mem_usage"):
+            if kwargs.pop(unsupported, None) is not None:
+                warnings.warn(
+                    f"SkinMapModel.from_pretrained ignores `{unsupported}`: the "
+                    "ensemble runs in fp32 on a single device. Pass `device=...` "
+                    "to choose the device. fp16/sharded loading is unsupported.",
+                    stacklevel=2,
+                )
+
+        path = Path(pretrained_model_name_or_path)
+        if path.is_dir():
+            snapshot = path
+        else:
+            from huggingface_hub import snapshot_download
+
+            snapshot = Path(
+                snapshot_download(
+                    repo_id=str(pretrained_model_name_or_path),
+                    revision=revision,
+                    cache_dir=cache_dir,
+                    token=token,
+                )
+            )
+
+        if config is None:
+            config = SkinMapConfig.from_pretrained(snapshot)
+
+        model = cls(config)
+        model._load_pipeline(snapshot, device=device)
+        return model
+
+    def _load_pipeline(self, snapshot, device: Optional[str] = None):
+        snapshot = Path(snapshot).resolve()
+        if str(snapshot) not in sys.path:
+            sys.path.insert(0, str(snapshot))
+
+        # `skinmap_runtime` is a fixed top-level package name resolved off sys.path,
+        # so only one snapshot's runtime can be live per process. If a different
+        # snapshot already imported it, that one is reused, so warn rather than
+        # silently run model B on model A's code.
+        already = sys.modules.get("skinmap_runtime")
+        if already is not None and getattr(already, "__path__", None):
+            loaded_from = Path(list(already.__path__)[0]).resolve().parent
+            if loaded_from != snapshot:
+                warnings.warn(
+                    "A different SkinMap snapshot's runtime is already loaded in "
+                    f"this process (from {loaded_from}). Reusing it. Loading two "
+                    "different SkinMap revisions in one process is unsupported.",
+                    stacklevel=2,
+                )
+
+        # Import the bundled runtime via importlib (NOT `from skinmap_runtime ...`):
+        # transformers' check_imports statically scans this file for import
+        # statements and would otherwise treat `skinmap_runtime` as a missing
+        # PyPI dependency, breaking from_pretrained for every user.
+        import importlib
+
+        CombinedEmbeddingPipeline = importlib.import_module(
+            "skinmap_runtime.combined_embedder"
+        ).CombinedEmbeddingPipeline
+
+        cfg_path = snapshot / self.config.pipeline_config
+        if not cfg_path.exists():
+            raise FileNotFoundError(
+                f"SkinMap pipeline config not found at {cfg_path}. "
+                "The repo snapshot may be incomplete."
+            )
+
+        resolved = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        self._pipeline = CombinedEmbeddingPipeline.from_config(
+            cfg_path, device=resolved
+        )
+        self._skinmap_device = torch.device(resolved)
+        # Keep the module's own buffer on the resolved device so standard
+        # introspection (next(model.buffers()).device) agrees with .device.
+        self._anchor = self._anchor.to(self._skinmap_device)
+        # Bundled metadata probes (optional): enables predict_meta(). Absent in
+        # embedding-only snapshots, so load lazily and stay quiet if not present.
+        probes_path = snapshot / getattr(self.config, "probes_dir", "probes")
+        if (probes_path / "manifest.json").exists():
+            self._predictor = MetaPredictor(probes_path)
+        return self
+
+    def _require_pipeline(self):
+        if self._pipeline is None:
+            raise RuntimeError(
+                "SkinMap pipeline is not loaded. Use "
+                "SkinMapModel.from_pretrained(...) rather than constructing the "
+                "model directly."
+            )
+        return self._pipeline
+
+    # ------------------------------------------------------------- device
+    @property
+    def device(self) -> torch.device:  # type: ignore[override]
+        return self._skinmap_device
+
+    def to(self, *args, **kwargs):  # type: ignore[override]
+        """Move every teacher + the projector onto the requested device."""
+        device = None
+        if args and (isinstance(args[0], (str, torch.device)) or args[0] is None):
+            device = args[0]
+        device = kwargs.get("device", device)
+        if device is None:
+            return self
+        device = torch.device(device)
+        super().to(device)
+        self._skinmap_device = device
+
+        pipe = self._pipeline
+        if pipe is not None:
+            pipe.device = device
+            for wrapper in getattr(pipe, "clip_models", []):
+                wrapper.model.to(device)
+            for wrapper in getattr(pipe, "ssl_models", []):
+                wrapper.model.to(device)
+            if getattr(pipe, "projector_model", None) is not None:
+                pipe.projector_model.to(device)
+        return self
+
+    def cuda(self, device=None):  # type: ignore[override]
+        return self.to(f"cuda:{device}" if isinstance(device, int) else "cuda")
+
+    def cpu(self):  # type: ignore[override]
+        return self.to("cpu")
+
+    # ------------------------------------------------------------- encode
+    @staticmethod
+    def _as_pil(image: ImageInput):
+        from PIL import Image
+
+        if isinstance(image, Image.Image):
+            return image.convert("RGB")
+        if isinstance(image, (bytes, bytearray)):
+            return Image.open(io.BytesIO(image)).convert("RGB")
+        if isinstance(image, (str, Path)):
+            return Image.open(image).convert("RGB")
+        raise TypeError(
+            f"Unsupported image input type: {type(image).__name__}. "
+            "Pass a file path, raw bytes, or a PIL.Image."
+        )
+
+    @torch.inference_mode()
+    def encode_image(self, image: ImageInput) -> np.ndarray:
+        """Encode one image into the 1024-d SkinMap space (L2-normalized)."""
+        pipe = self._require_pipeline()
+        return pipe.embed_image(self._as_pil(image))
+
+    @torch.inference_mode()
+    def encode_text(
+        self, text: str, templates: Optional[List[str]] = None
+    ) -> np.ndarray:
+        """Encode a text query into the shared 1024-d space.
+
+        By default the query is expanded into clinical prompt templates and the
+        results averaged. Pass ``templates=["{}"]`` to encode the raw text.
+        """
+        pipe = self._require_pipeline()
+        return pipe.encode_text(text, templates=templates)
+
+    # ------------------------------------------------------------- predict_meta
+    def _require_predictor(self):
+        if self._predictor is None:
+            raise RuntimeError(
+                "No metadata probes are bundled with this SkinMap repo, so "
+                "predict_meta() is unavailable. Rebuild the package with the "
+                "probes/ directory to enable metadata prediction."
+            )
+        return self._predictor
+
+    @torch.inference_mode()
+    def predict_meta(
+        self,
+        image: Union[ImageInput, np.ndarray],
+        modality: Optional[str] = None,
+        attributes: Optional[List[str]] = None,
+    ) -> dict:
+        """Predict metadata from an image or a precomputed 1024-d SkinMap vector.
+
+        Returns a dict of attributes (Fitzpatrick skin type, age, sex, geographic
+        origin, body region, ...). By default the global probes are used, which
+        reproduce the demographic estimates reported in the paper. Pass
+        ``modality`` (``"clinical"``, ``"dermoscopy"`` or ``"TBP"``) to use the
+        modality-specific Fitzpatrick probe, which is more accurate when the
+        imaging modality is known; the other attributes stay global.
+        """
+        predictor = self._require_predictor()
+        emb = image if isinstance(image, np.ndarray) else self.encode_image(image)
+        return predictor.predict(emb, modality=modality, attributes=attributes)
+
+    def forward(self, image: Optional[ImageInput] = None, **kwargs) -> torch.Tensor:
+        """Convenience: ``forward(image)`` returns ``encode_image(image)`` as a tensor.
+
+        SkinMap is an encoder, not a trainable head, so prefer the explicit
+        ``encode_image`` / ``encode_text``. HF-style tensor calls such as
+        ``model(pixel_values=...)`` raise a clear error instead of crashing
+        cryptically.
+        """
+        if image is None:
+            raise TypeError(
+                "SkinMapModel is an encoder: call model.encode_image(img) or "
+                "model.encode_text(txt). forward() accepts a single image "
+                f"(path/bytes/PIL), got keyword args {sorted(kwargs)} instead."
+            )
+        return torch.from_numpy(np.ascontiguousarray(self.encode_image(image)))
+
+    # ------------------------------------------------------------- search
+    def search(
+        self, query: ImageInput, k: int = 10
+    ) -> Tuple[List[int], List[float]]:
+        """Nearest-neighbor search against the bundled atlas index.
+
+        `query` may be an image (path/bytes/PIL) or a precomputed 1024-d vector.
+        Raises a clear error if no atlas/index was bundled with this repo.
+        """
+        pipe = self._require_pipeline()
+        if isinstance(query, np.ndarray):
+            vector = query
+        else:
+            vector = self.encode_image(query)
+        try:
+            return pipe.find_nearest_neighbors(np.asarray(vector), k=k)
+        except RuntimeError as exc:
+            raise RuntimeError(
+                "Nearest-neighbor search needs a bundled atlas index, which is "
+                "not present in this repo. Rebuild the package with --with-atlas "
+                "to enable search()."
+            ) from exc

probes/manifest.json

@@ -0,0 +1,100 @@
+{
+  "attributes": {
+    "fitzpatrick": {
+      "task": "classification",
+      "n_train": 33082,
+      "classes": [
+        "1.0",
+        "2.0",
+        "3.0",
+        "4.0",
+        "5.0",
+        "6.0"
+      ],
+      "modality_specific": true,
+      "modality_groups": {
+        "clinical": "clinical",
+        "dermoscopy": "derm",
+        "TBP": "derm"
+      }
+    },
+    "gender": {
+      "task": "classification",
+      "n_train": 523878,
+      "classes": [
+        "female",
+        "male"
+      ],
+      "modality_specific": false
+    },
+    "origin": {
+      "task": "classification",
+      "n_train": 518420,
+      "classes": [
+        "Argentina",
+        "Australia",
+        "Austria",
+        "Brazil",
+        "Germany",
+        "Greece",
+        "Guinea",
+        "India",
+        "Italy",
+        "Madagascar",
+        "Malawi",
+        "Netherlands",
+        "New Zealand",
+        "Portugal",
+        "Spain",
+        "Switzerland",
+        "Tanzania",
+        "United States"
+      ],
+      "modality_specific": false
+    },
+    "body_region": {
+      "task": "classification",
+      "n_train": 613385,
+      "classes": [
+        "acral",
+        "gluteal_buttocks",
+        "groin_genital_perineum",
+        "head_neck",
+        "lower_limb",
+        "trunk_back",
+        "trunk_front",
+        "upper_limb"
+      ],
+      "modality_specific": false
+    },
+    "laterality": {
+      "task": "classification",
+      "n_train": 173585,
+      "classes": [
+        "left",
+        "right"
+      ],
+      "modality_specific": false
+    },
+    "modality": {
+      "task": "classification",
+      "n_train": 596759,
+      "classes": [
+        "TBP",
+        "clinical",
+        "dermoscopy"
+      ],
+      "modality_specific": false
+    },
+    "age": {
+      "task": "regression",
+      "n_train": 534789,
+      "clip": [
+        0,
+        100
+      ],
+      "modality_specific": false
+    }
+  },
+  "note": "Global probes are the original frozen Jun-17 (paper) prediction models that produced the atlas *_pred imputations; modality-specific Fitzpatrick probes are fit on the frozen atlas. Input = 1024-d L2-normalized SkinMap embedding."
+}

probes/model_origin.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6402dd67061dc178ee158a64842e426a57f62f95f201e3d1142488dabe112e7d
+size 148444

requirements.txt

@@ -0,0 +1,13 @@
+# Inference dependencies for the packaged SkinMap model.
+# (Training-only deps — wandb, umap, matplotlib, plotly, glasbey — are intentionally absent.)
+torch>=2.0
+torchvision>=0.15
+transformers>=4.40
+huggingface_hub>=0.23
+numpy>=1.24
+scipy>=1.10
+tqdm>=4.65
+loguru>=0.7
+faiss-cpu>=1.7
+joblib>=1.3
+Pillow>=9.0

skinmap_runtime/_compat.py

@@ -0,0 +1,153 @@
+"""Self-contained inference-only shims for the packaged SkinMap model.
+
+The runtime `combined_embedder.py` normally pulls three helpers out of the
+training code base (`src.create_skinmap`, `src.train_clip`). Those modules carry
+heavy, training-only top-level imports (wandb, umap, matplotlib, plotly, DDP,
+torch.compile) that have no place in a portable inference package. This module
+re-implements exactly the three symbols the runtime needs, with nothing else.
+
+Keep this file dependency-light: stdlib + torch + transformers only.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from loguru import logger
+from torchvision import transforms
+from torchvision.transforms import InterpolationMode
+
+# Mirrors src/skinmap/models/loaders.py: names the SSL teacher checkpoints are
+# known by. Only the first three are used by the shipped SkinMap ensemble.
+SSL_MODEL_NAMES = [
+    "dino_qderma",
+    "ibot_qderma",
+    "mae_qderma",
+    "simclr_qderma",
+    "byol_qderma",
+    "colorme_qderma",
+    "panderm_base",
+    "panderm_large",
+]
+
+
+def get_imagenet_transform():
+    """Standard ImageNet preprocessing for the SSL teachers.
+
+    Verbatim copy of src/skinmap/data/transforms.py:get_imagenet_transform so the
+    packaged pipeline preprocesses SSL inputs identically to training.
+    """
+    return transforms.Compose(
+        [
+            transforms.Resize(256, interpolation=InterpolationMode.BICUBIC),
+            transforms.CenterCrop(224),
+            transforms.ToTensor(),
+            transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
+        ]
+    )
+
+
+def _processor_name_fallback(model_name: str) -> str:
+    """Infer the upstream processor repo for a local CLIP checkpoint name.
+
+    Only a fallback: the shipped CLIP checkpoints are self-contained HF dirs, so
+    the processor loads locally first. Match on the architecture substring so it
+    stays robust to clean teacher directory names (e.g. ``siglip_vit-base-patch32``).
+    """
+    if "monet" in model_name:
+        return "suinleelab/monet"
+    if "vit-base-patch32" in model_name:
+        return "openai/clip-vit-base-patch32"
+    if "vit-large-patch14" in model_name:
+        return "openai/clip-vit-large-patch14"
+    return model_name
+
+
+def load_model_and_processor(model_name, device, *args, **kwargs):
+    """Load a CLIP teacher (model + processor) for inference.
+
+    A trimmed equivalent of src/train_clip.py:load_model_and_processor: no
+    DDP wrapping, no torch.compile, no random-init / siglip-bias training paths.
+    The shipped CLIP checkpoints are self-contained HF directories (config +
+    safetensors + tokenizer + preprocessor), so we load them locally first and
+    only reach out to the hub if a processor artifact is somehow missing.
+    """
+    from transformers import (
+        CLIPImageProcessor,
+        CLIPModel,
+        CLIPProcessor,
+        CLIPTokenizer,
+    )
+
+    is_local = Path(model_name).exists()
+
+    # Candidate processor sources, local checkpoint dir first.
+    candidates = [model_name]
+    fallback = _processor_name_fallback(model_name)
+    if fallback != model_name:
+        candidates.append(fallback)
+
+    def _load_processor(candidate: str, local_only: bool):
+        try:
+            return CLIPProcessor.from_pretrained(
+                candidate, local_files_only=local_only
+            )
+        except Exception as base_exc:
+            # Some checkpoints ship tokenizer + image processor separately.
+            try:
+                tok = CLIPTokenizer.from_pretrained(
+                    candidate, local_files_only=local_only
+                )
+                img = CLIPImageProcessor.from_pretrained(
+                    candidate, local_files_only=local_only
+                )
+                return CLIPProcessor(tokenizer=tok, image_processor=img)
+            except Exception:
+                raise base_exc
+
+    processor = None
+    errors = []
+    for candidate in candidates:
+        local_only = Path(candidate).exists()
+        try:
+            processor = _load_processor(candidate, local_only=local_only)
+            break
+        except Exception as exc:  # pragma: no cover - runtime safeguard
+            errors.append(exc)
+            logger.warning(f"Could not load CLIP processor from {candidate}: {exc}")
+    if processor is None:
+        raise RuntimeError(
+            f"Failed to load a CLIP processor for {model_name}: {errors[-1]}"
+        ) from errors[-1]
+
+    model = CLIPModel.from_pretrained(model_name, local_files_only=is_local)
+    model.to(device)
+    model.eval()
+    return model, processor
+
+
+def load_ssl_local(ckpt_path, n_head_layers: int = 0):
+    """Load an SSL teacher from a bundled local checkpoint.
+
+    `Embedder.load_pretrained` always downloads from the upstream vm02 URL, so we
+    bypass it: the per-family loader functions (`load_dino`/`load_ibot`/
+    `load_mae`) accept a local `ckp_path` and read their architecture config out
+    of the checkpoint itself. The teacher family (dino/ibot/mae) is inferred from
+    the file stem, robustly to clean names (``mae_vit-base-patch16.pth`` as well as
+    ``mae_qderma.pth``); the family key maps to the same loader either way.
+    """
+    from skinmap_runtime.core.pkg.embedder import Embedder
+
+    stem = Path(ckpt_path).stem
+    parts = set(stem.split("-")) | set(stem.split("_")) | {stem}
+    family = next((f for f in ("dino", "ibot", "mae") if f in parts), None)
+    # Map the family to its dermatology-pretrained registry key
+    # (dino_qderma/ibot_qderma/mae_qderma); all three exist and select
+    # load_dino/load_ibot/load_mae, which read the architecture from the checkpoint.
+    key = f"{family}_qderma" if family else stem
+    loader_func = Embedder.get_model_func(key)
+    return loader_func(
+        ckp_path=str(ckpt_path),
+        return_info=True,
+        n_head_layers=n_head_layers,
+    )

skinmap_runtime/combined_embedder.py

@@ -0,0 +1,926 @@
+import io
+import json
+import logging
+import os
+import sys
+import threading
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Iterable, List, Optional, Sequence, Tuple
+
+import faiss  # type: ignore
+import numpy as np
+import torch
+from PIL import Image
+from torch import nn
+
+from skinmap_runtime.core.pkg.embedder import Embedder
+from skinmap_runtime._compat import SSL_MODEL_NAMES, get_imagenet_transform
+from skinmap_runtime._compat import load_model_and_processor
+
+try:
+    import joblib
+except ImportError:  # pragma: no cover - fallback for older sklearn versions
+    from sklearn.externals import joblib  # type: ignore
+
+
+logger = logging.getLogger(__name__)
+
+
+def _coerce_clip_features(feats, embeds_attr: str):
+    """Return the projected feature tensor from a CLIP ``get_*_features`` call.
+
+    Some transformers versions return a ``BaseModelOutputWithPooling`` rather
+    than a tensor. In that case the projected embedding lives in the ``*_embeds``
+    field, or in ``pooler_output`` (which holds the projected vector for the
+    shipped CLIP checkpoints). ``last_hidden_state`` is not used as a fallback:
+    its CLS token has the encoder hidden dim, not the projection dim, so it would
+    silently corrupt the embedding. Shared by the image and text paths.
+    """
+    if isinstance(feats, torch.Tensor):
+        return feats
+    for attr in (embeds_attr, "pooler_output"):
+        val = getattr(feats, attr, None)
+        if val is not None:
+            return val
+    kind = embeds_attr.split("_")[0]
+    raise TypeError(
+        f"get_{kind}_features returned {type(feats).__name__} with neither "
+        f"'{embeds_attr}' nor 'pooler_output'; cannot extract the projected "
+        f"embedding (unsupported transformers version?)."
+    )
+
+
+@dataclass
+class ClipModelWrapper:
+    name: str
+    model: nn.Module
+    processor: Any  # CLIPProcessor
+
+    def embed(self, image: Image.Image, device: torch.device) -> np.ndarray:
+        prepared = image.copy().convert("RGB")
+        prepared.thumbnail((512, 512), Image.Resampling.LANCZOS)
+        inputs = self.processor(images=[prepared], return_tensors="pt", padding=True)
+        inputs = {k: v.to(device) for k, v in inputs.items()}
+        self.model.eval()
+        with torch.inference_mode():
+            feats = _coerce_clip_features(
+                self.model.get_image_features(**inputs), "image_embeds"
+            )
+            # Normalize to match training pipeline
+            feats = feats / feats.norm(p=2, dim=-1, keepdim=True)
+        return feats.detach().cpu().numpy()
+
+
+@dataclass
+class SSLModelWrapper:
+    name: str
+    model: nn.Module
+    transform: Any  # torchvision.transforms.Compose
+
+    def embed(self, image: Image.Image, device: torch.device) -> np.ndarray:
+        prepared = image.copy().convert("RGB")
+        tensor = self.transform(prepared).unsqueeze(0).to(device)
+        self.model.eval()
+        with torch.inference_mode():
+            feats = self.model(tensor)
+            feats = torch.nn.functional.normalize(feats, dim=-1, p=2)
+        return feats.detach().cpu().numpy()
+
+
+class CombinedEmbeddingPipeline:
+    """
+    Runtime helper that mirrors the embedding combination used to build SkinMap.
+
+    It loads the individual models, the saved SVD projection (if present),
+    and the FAISS index so that uploaded images can be embedded and searched
+    in the same latent space.
+    """
+
+    def __init__(self, config: dict, config_path: Path, device: Optional[str] = None):
+        self.config = config
+        self.config_path = config_path.resolve()
+        self.config_dir = self.config_path.parent.resolve()
+        self.device = torch.device(device or "cpu")
+        self._faiss_lock = threading.Lock()
+
+        self.skinmap_root_path = self._resolve_skinmap_root(config.get("skinmap_root"))
+        self.output_dir = self._resolve_output_dir(config.get("output_dir"))
+        self._model_search_roots = self._build_model_search_roots()
+
+        self.clip_models: List[ClipModelWrapper] = []
+        self.ssl_models: List[SSLModelWrapper] = []
+        self.models_in_order: List[Tuple[str, Any]] = (
+            []
+        )  # (type, wrapper) preserving config order
+
+        self._load_models()
+
+        # Load SVD (legacy pipeline) or projector (new pipeline)
+        self.svd_image = self._load_joblib(config.get("svd", {}).get("image"))
+        self.projector_model = None
+        self.whitening_stats = None
+        self._load_projector_if_available()
+
+        artifacts = config.get("artifacts", {})
+        self._embeddings_npz_path = self._resolve_artifact(
+            artifacts.get("embeddings_npz")
+        )
+        self._faiss_index_path = self._resolve_artifact(artifacts.get("faiss_index"))
+        self.faiss_index = self._load_faiss(self._faiss_index_path)
+        self.vector_dim = config.get("vector", {}).get("dimension")
+        self.index_metric = "cosine"  # Matches build_faiss_index default
+        self.umap_model = self._load_umap(config.get("umap", {}).get("model"))
+
+    @classmethod
+    def from_config(cls, path: str | os.PathLike, device: Optional[str] = None):
+        path = Path(path)
+        with open(path, "r") as f:
+            config = json.load(f)
+        return cls(config=config, config_path=path, device=device)
+
+    def _resolve_artifact(self, rel_path: Optional[str]) -> Optional[Path]:
+        if rel_path is None:
+            return None
+        return (self.output_dir / rel_path).resolve()
+
+    def _resolve_skinmap_root(self, raw_root: Optional[str]) -> Optional[Path]:
+        candidate_roots: List[Path] = []
+
+        def _add_candidate(path_like: Optional[str | os.PathLike]):
+            if not path_like:
+                return
+            path_obj = Path(path_like)
+            if not path_obj.is_absolute():
+                path_obj = (self.config_dir / path_obj).resolve()
+            else:
+                path_obj = path_obj.resolve()
+            candidate_roots.append(path_obj)
+
+        if raw_root:
+            _add_candidate(raw_root)
+            raw_path = Path(raw_root)
+            if not raw_path.is_absolute():
+                _add_candidate(self.config_dir.parent / raw_path)
+
+        # Walk up the config directory hierarchy to find the project root.
+        for ancestor in self.config_dir.parents:
+            candidate_roots.append(ancestor.resolve())
+
+        seen: set[str] = set()
+        for candidate in candidate_roots:
+            key = str(candidate)
+            if key in seen:
+                continue
+            seen.add(key)
+            if not candidate.exists():
+                continue
+            # Heuristic: the SkinMap repo root contains src/combined_embedder.py
+            if (candidate / "src" / "combined_embedder.py").exists():
+                if key not in sys.path:
+                    sys.path.append(key)
+                return candidate
+
+        logger.warning(
+            "Unable to resolve SkinMap project root from config %s; "
+            "nearest-neighbor upload search may fail.",
+            self.config_path,
+        )
+        return None
+
+    def _resolve_with_bases(
+        self,
+        path_value: str | os.PathLike,
+        bases: Iterable[Optional[Path]],
+        *,
+        must_exist: bool = False,
+    ) -> Optional[Path]:
+        path_obj = Path(path_value)
+        candidates: List[Path] = []
+
+        if path_obj.is_absolute():
+            candidates.append(path_obj)
+        else:
+            candidates.append((Path.cwd() / path_obj).resolve())
+            for base in bases:
+                if base is None:
+                    continue
+                candidates.append((base / path_obj).resolve())
+
+        for candidate in candidates:
+            if candidate.exists():
+                return candidate
+        if must_exist:
+            return None
+        return candidates[0] if candidates else None
+
+    def _resolve_output_dir(self, raw_output: Optional[str]) -> Path:
+        if not raw_output:
+            return self.config_dir
+
+        bases: List[Optional[Path]] = [
+            self.skinmap_root_path,
+            self.config_dir,
+            (
+                self.config_dir.parent
+                if self.config_dir.parent != self.config_dir
+                else None
+            ),
+        ]
+        resolved = self._resolve_with_bases(raw_output, bases, must_exist=True)
+        if resolved is None or not resolved.exists():
+            logger.warning(
+                "Configured output_dir '%s' could not be resolved relative to %s; "
+                "falling back to %s",
+                raw_output,
+                self.config_path,
+                self.config_dir,
+            )
+            return self.config_dir
+        return resolved
+
+    def _build_model_search_roots(self) -> List[Path]:
+        roots: List[Path] = []
+        for candidate in (
+            self.output_dir,
+            self.output_dir.parent,
+            self.skinmap_root_path,
+            self.skinmap_root_path / "assets" if self.skinmap_root_path else None,
+            self.config_dir,
+        ):
+            if candidate is None:
+                continue
+            resolved = candidate.resolve()
+            if resolved not in roots:
+                roots.append(resolved)
+        return roots
+
+    def _resolve_model_path(self, source: str) -> Optional[Path]:
+        resolved = self._resolve_with_bases(
+            source, self._model_search_roots, must_exist=False
+        )
+        if resolved is not None and resolved.exists():
+            return resolved
+        return None
+
+    def _load_models(self):
+        models_config: Sequence[dict] = self.config.get("models", [])
+        if not models_config:
+            raise ValueError("No model definitions found in configuration.")
+
+        ssl_transform = get_imagenet_transform()
+        for entry in models_config:
+            source = entry.get("source")
+            if source is None:
+                continue
+            resolved_path = self._resolve_model_path(source)
+            resolved_source = str(resolved_path) if resolved_path else source
+            model_type = entry.get("type") or (
+                "ssl" if source in SSL_MODEL_NAMES else "clip"
+            )
+            if model_type == "ssl":
+                if resolved_path is not None and entry.get("is_local"):
+                    from skinmap_runtime._compat import load_ssl_local
+
+                    model, _info, _cfg = load_ssl_local(
+                        str(resolved_path), n_head_layers=0
+                    )
+                else:
+                    model, _info, _cfg = Embedder.load_pretrained(
+                        source, return_info=True, n_head_layers=0
+                    )
+                model.to(self.device)
+                wrapper = SSLModelWrapper(
+                    name=source, model=model, transform=ssl_transform
+                )
+                self.ssl_models.append(wrapper)
+                self.models_in_order.append(("ssl", wrapper))
+            else:
+                model, processor = load_model_and_processor(
+                    resolved_source, self.device
+                )
+                wrapper = ClipModelWrapper(
+                    name=source, model=model, processor=processor
+                )
+                self.clip_models.append(wrapper)
+                self.models_in_order.append(("clip", wrapper))
+
+    def _load_joblib(self, rel_path: Optional[str]):
+        resolved = self._resolve_artifact(rel_path)
+        if resolved is None or not resolved.exists():
+            return None
+        try:
+            logger.info("Loading model from %s", resolved)
+            return joblib.load(resolved)
+        except Exception as exc:
+            logger.warning("Failed to load model from %s: %s", resolved, exc)
+            return None
+
+    def _load_umap(self, rel_path: Optional[str]):
+        model = self._load_joblib(rel_path)
+        if model is None:
+            # Retry with NNDescent monkey-patch for pynndescent version mismatch.
+            # Newer pynndescent adds a `quantization` attribute that older versions
+            # lack, causing joblib.load to fail.
+            resolved = self._resolve_artifact(rel_path)
+            if resolved is not None and resolved.exists():
+                try:
+                    os.environ.setdefault("NUMBA_CACHE_DIR", "/tmp/numba_cache")
+                    import pynndescent
+
+                    NNDescent = pynndescent.NNDescent
+                    if not hasattr(NNDescent, "quantization"):
+                        NNDescent.quantization = None  # type: ignore[attr-defined]
+                        logger.info(
+                            "Patched NNDescent with missing 'quantization' attribute"
+                        )
+                    if not hasattr(NNDescent, "_min_distance"):
+                        NNDescent._min_distance = None  # type: ignore[attr-defined]
+                        logger.info(
+                            "Patched NNDescent with missing '_min_distance' attribute"
+                        )
+                    model = joblib.load(resolved)
+                except Exception as exc:
+                    logger.warning(
+                        "Failed to load UMAP model even after NNDescent patch: %s", exc
+                    )
+                    model = None
+        if model is not None:
+            logger.info("UMAP model loaded for upload projections")
+        return model
+
+    def _load_projector_if_available(self):
+        """Load trained projector model if available in config."""
+        projector_cfg = self.config.get("projector", {})
+        if not projector_cfg.get("used", False):
+            logger.info("No trained projector in config; using legacy SVD pipeline")
+            return
+
+        # Load projector config
+        projector_config_path = self._resolve_artifact(projector_cfg.get("config"))
+        if projector_config_path is None or not projector_config_path.exists():
+            logger.warning("Projector config not found; falling back to SVD pipeline")
+            return
+
+        try:
+            with open(projector_config_path, "r") as f:
+                proj_config = json.load(f)
+        except Exception as exc:
+            logger.warning(f"Failed to load projector config: {exc}")
+            return
+
+        # Load whitening stats
+        whitening_path = self._resolve_artifact(projector_cfg.get("whitening_stats"))
+        if whitening_path is None or not whitening_path.exists():
+            logger.warning("Whitening stats not found; falling back to SVD pipeline")
+            return
+
+        try:
+            with np.load(whitening_path, allow_pickle=True) as data:
+                files = set(data.files)
+                dims = list(data["dims"])
+                n_models = int(data["n_models"]) if "n_models" in files else len(dims)
+
+                # Load mu/W arrays if they exist (skip_whitening=False case)
+                mu_list = []
+                W_list = []
+                if n_models > 0:
+                    if "mu" in files:
+                        mu_list = list(data["mu"])
+                    else:
+                        # Try to load individual mu_i files
+                        try:
+                            mu_list = [data[f"mu_{i}"] for i in range(n_models)]
+                        except KeyError:
+                            logger.warning(
+                                "No mu whitening parameters found (skip_whitening=True?)"
+                            )
+
+                    if "W" in files:
+                        W_list = list(data["W"])
+                    else:
+                        # Try to load individual W_i files
+                        try:
+                            W_list = [data[f"W_{i}"] for i in range(n_models)]
+                        except KeyError:
+                            logger.warning(
+                                "No W whitening parameters found (skip_whitening=True?)"
+                            )
+
+                clip_indices = (
+                    list(data["clip_indices"]) if "clip_indices" in files else []
+                )
+
+                # Load text whitening if available
+                text_mu_list = []
+                text_W_list = []
+                text_dims = []
+                if "text_mu" in files and "text_W" in files:
+                    text_mu_list = list(data["text_mu"])
+                    text_W_list = list(data["text_W"])
+                    text_dims = (
+                        list(data["text_dims"])
+                        if "text_dims" in files
+                        else [w.shape[0] for w in text_W_list]
+                    )
+                elif "n_text_models" in files:
+                    n_text = int(data["n_text_models"])
+                    if n_text > 0:
+                        try:
+                            text_mu_list = [data[f"text_mu_{i}"] for i in range(n_text)]
+                            text_W_list = [data[f"text_W_{i}"] for i in range(n_text)]
+                            text_dims = (
+                                list(data["text_dims"])
+                                if "text_dims" in files
+                                else [w.shape[0] for w in text_W_list]
+                            )
+                        except KeyError:
+                            logger.warning("No text whitening parameters found")
+
+                self.whitening_stats = {
+                    "mu": mu_list,
+                    "W": W_list,
+                    "dims": dims,
+                    "clip_indices": clip_indices,
+                }
+
+                if text_mu_list and text_W_list:
+                    self.whitening_stats["text_whitening"] = {
+                        "mu": text_mu_list,
+                        "W": text_W_list,
+                        "dims": text_dims,
+                    }
+                    logger.info(
+                        f"Loaded text whitening stats for {len(text_mu_list)} CLIP models"
+                    )
+                else:
+                    self.whitening_stats["text_whitening"] = None
+
+            if len(mu_list) > 0:
+                logger.info(
+                    f"Loaded whitening stats for {len(self.whitening_stats['mu'])} teachers"
+                )
+            else:
+                logger.info("Whitening was skipped during training (--skip_whitening)")
+            logger.info(
+                f"CLIP model indices for text: {self.whitening_stats['clip_indices']}"
+            )
+        except Exception as exc:
+            logger.warning(f"Failed to load whitening stats: {exc}")
+            return
+
+        # Load projector model
+        projector_model_path = self._resolve_artifact(projector_cfg.get("model"))
+        if projector_model_path is None or not projector_model_path.exists():
+            logger.warning("Projector model not found; falling back to SVD pipeline")
+            return
+
+        try:
+            from skinmap_runtime.embedding_fusion import BuildSpec, build_model
+
+            # Calculate actual text input dimension from CLIP models
+            # text_dim in config is the OUTPUT dimension, but we need INPUT dimension
+            clip_indices = self.whitening_stats.get("clip_indices", [])
+            if clip_indices and "text_dims" in self.whitening_stats:
+                # Use text_dims from whitening stats (these are the actual input dims)
+                text_input_dim = sum(self.whitening_stats["text_dims"])
+                logger.info(f"Calculated text input dimension: {text_input_dim}")
+            else:
+                # Fallback: calculate from teacher_dims at clip_indices
+                teacher_dims = proj_config["teacher_dims"]
+                text_input_dim = (
+                    sum(teacher_dims[i] for i in clip_indices)
+                    if clip_indices
+                    else teacher_dims[0]
+                )
+                logger.info(
+                    f"Calculated text input dimension from teacher_dims: {text_input_dim}"
+                )
+
+            spec = BuildSpec(
+                teacher_names=proj_config.get("model_paths", []),
+                teacher_dims=proj_config["teacher_dims"],
+                text_dim=text_input_dim,  # Use calculated INPUT dimension, not output
+                out_dim=proj_config["projector_dim"],
+                pca_image=None,
+                pca_text=None,
+                kind=proj_config["projector_type"],
+            )
+
+            self.projector_model = build_model(spec)
+            self.projector_model.load_state_dict(
+                torch.load(
+                    projector_model_path, map_location=self.device, weights_only=False
+                )
+            )
+            self.projector_model.to(self.device)
+            self.projector_model.eval()
+            logger.info(
+                f"Loaded trained projector ({proj_config['projector_type']}, dim={proj_config['projector_dim']})"
+            )
+        except Exception as exc:
+            logger.warning(f"Failed to load projector model: {exc}")
+            self.projector_model = None
+            self.whitening_stats = None
+
+    def _load_faiss(self, path: Optional[Path]):
+        if path is None or not path.exists():
+            return None
+        try:
+            logger.info("Loading FAISS index from %s", path)
+            return faiss.read_index(str(path))
+        except Exception as exc:
+            logger.warning("Failed to load FAISS index from %s: %s", path, exc)
+            return None
+
+    def _ensure_faiss_index(self) -> faiss.Index:
+        index = self.faiss_index
+        if index is not None:
+            return index
+        if self._embeddings_npz_path is None or not self._embeddings_npz_path.exists():
+            raise RuntimeError(
+                "FAISS index not available and embeddings_npz artifact is missing."
+            )
+        with self._faiss_lock:
+            if self.faiss_index is not None:
+                return self.faiss_index
+            self.faiss_index = self._build_faiss_from_embeddings(
+                self._embeddings_npz_path
+            )
+            return self.faiss_index
+
+    def _build_faiss_from_embeddings(self, npz_path: Path) -> faiss.Index:
+        logger.info(
+            "Building FAISS index on the fly from %s; this may take several minutes.",
+            npz_path,
+        )
+        try:
+            data = np.load(npz_path, mmap_mode="r")
+        except Exception as exc:
+            raise RuntimeError(
+                f"Failed to load embeddings NPZ from {npz_path}: {exc}"
+            ) from exc
+
+        if "image_embeddings" not in data:
+            data.close()
+            raise RuntimeError(
+                f"Embeddings NPZ at {npz_path} does not contain 'image_embeddings'."
+            )
+
+        embeddings = data["image_embeddings"]
+        if embeddings.ndim != 2:
+            data.close()
+            raise RuntimeError("Embeddings array must be 2-dimensional.")
+
+        dim = embeddings.shape[1]
+        if self.vector_dim is not None and self.vector_dim != dim:
+            logger.warning(
+                "Configured vector dimension (%s) does not match embeddings (%s); using %s.",
+                self.vector_dim,
+                dim,
+                dim,
+            )
+
+        index = faiss.IndexFlatIP(dim)
+        batch_size = 8192
+        total = embeddings.shape[0]
+        try:
+            for start in range(0, total, batch_size):
+                end = min(start + batch_size, total)
+                batch = np.array(embeddings[start:end], dtype=np.float32, copy=False)
+                norms = np.linalg.norm(batch, axis=1, keepdims=True)
+                norms[norms == 0] = 1.0
+                batch = batch / norms
+                index.add(batch)
+                if (start // batch_size) % 50 == 0:
+                    logger.info(
+                        "FAISS build progress: %d/%d vectors (%.1f%%)",
+                        end,
+                        total,
+                        (end / total) * 100.0,
+                    )
+        finally:
+            data.close()
+        logger.info("FAISS index build complete (%d vectors).", total)
+
+        target_path = self._faiss_index_path
+        if target_path is not None:
+            try:
+                target_path.parent.mkdir(parents=True, exist_ok=True)
+                faiss.write_index(index, str(target_path))
+                logger.info("Cached FAISS index to %s", target_path)
+                self._faiss_index_path = target_path
+            except Exception as exc:
+                logger.warning(
+                    "Failed to write FAISS index to %s: %s. Continuing without persistence.",
+                    target_path,
+                    exc,
+                )
+        return index
+
+    def embed_image(self, image: Image.Image) -> np.ndarray:
+        """
+        Embed a PIL image into the combined latent space.
+
+        Uses trained projector if available, otherwise falls back to SVD.
+        """
+        if not self.models_in_order:
+            raise RuntimeError("No models loaded for embedding.")
+
+        features: List[np.ndarray] = []
+
+        # Iterate models in the same order as in config (critical for projector)
+        for model_type, model_wrapper in self.models_in_order:
+            if model_type == "clip":
+                features.append(model_wrapper.embed(image, self.device))
+            else:  # ssl
+                # SSL transforms expect tensor input; ensure consistent preprocessing
+                prepared = image.copy().convert("RGB")
+                features.append(model_wrapper.embed(prepared, self.device))
+
+        # Now apply either projector or SVD based on what's available
+        if self.projector_model is not None and self.whitening_stats is not None:
+            # Use trained projector pipeline
+            combined = self._apply_projector_pipeline(features)
+        else:
+            # Legacy SVD pipeline
+            combined = np.concatenate(
+                [feat.astype(np.float32) for feat in features], axis=1
+            )
+
+            if self.svd_image is not None:
+                combined = self.svd_image.transform(combined)
+
+        if combined.ndim == 2:
+            return combined[0]
+        return combined
+
+    def _apply_projector_pipeline(self, features: List[np.ndarray]) -> np.ndarray:
+        """Apply whitening + projector to per-model features."""
+        from skinmap_runtime.embedding_fusion import apply_whiten, l2_normalize
+
+        # Check if whitening was used (mu/W arrays present)
+        has_whitening = (
+            self.whitening_stats is not None
+            and "mu" in self.whitening_stats
+            and "W" in self.whitening_stats
+            and len(self.whitening_stats["mu"]) > 0
+            and len(self.whitening_stats["W"]) > 0
+        )
+
+        # Step 1: L2 normalize and optionally whiten each model's embeddings
+        processed = []
+        for i, feat in enumerate(features):
+            # L2 normalize
+            feat_norm = l2_normalize(feat, axis=1)
+
+            # Whiten if whitening stats are available
+            if has_whitening:
+                feat_processed = apply_whiten(
+                    feat_norm,
+                    self.whitening_stats["mu"][i],
+                    self.whitening_stats["W"][i],
+                )
+            else:
+                # No whitening - projector was trained on raw L2-normalized embeddings
+                feat_processed = feat_norm
+
+            processed.append(feat_processed)
+
+        # Step 2: Concatenate
+        z_cat = np.concatenate(processed, axis=1).astype(np.float32)
+
+        # Step 3: L2-normalize each teacher block (matches training pipeline)
+        dims = self.whitening_stats["dims"]
+        start = 0
+        for dim in dims:
+            block = z_cat[:, start : start + dim]
+            norms = np.linalg.norm(block, axis=1, keepdims=True)
+            norms[norms == 0] = 1.0
+            z_cat[:, start : start + dim] = block / norms
+            start += dim
+
+        # Step 4: Project through trained model
+        z_cat_tensor = torch.from_numpy(z_cat).to(self.device)
+        with torch.inference_mode():
+            projected = self.projector_model.img_head(z_cat_tensor)
+
+        return projected.detach().cpu().numpy()
+
+    def embed_bytes(self, data: bytes) -> np.ndarray:
+        image = Image.open(io.BytesIO(data)).convert("RGB")
+        return self.embed_image(image)
+
+    # Default prompt templates for text search.  CLIP models are trained on
+    # image-caption pairs so bare keywords ("melanoma") live in a very
+    # different part of the embedding space than descriptive captions.
+    # Wrapping the query in templates and averaging the resulting embeddings
+    # brings the query much closer to the image distribution.
+    DEFAULT_TEXT_TEMPLATES: List[str] = [
+        "This image shows {}",
+        "a clinical photograph of {}",
+        "a dermatoscopic image of {}",
+        "a medical image showing {}",
+        "a close-up photograph of {} on skin",
+    ]
+
+    def encode_text(
+        self,
+        text: str,
+        templates: Optional[List[str]] = None,
+    ) -> np.ndarray:
+        """Encode a text query into the shared embedding space.
+
+        When *templates* is provided (or defaults are used), the query is
+        formatted into each template, encoded separately, and the resulting
+        embeddings are averaged and re-normalized.  Pass ``templates=["{}"]``
+        to encode the raw text without any prompt engineering.
+
+        For projector pipeline: extracts text embeddings from CLIP models only,
+        concatenates them, and projects through the text head.
+
+        For SVD pipeline: uses first CLIP model only.
+        """
+        if not self.clip_models:
+            raise RuntimeError("No CLIP models loaded. Cannot encode text.")
+
+        if templates is None:
+            templates = self.DEFAULT_TEXT_TEMPLATES
+
+        # Expand the query into all templates
+        prompts = [t.format(text) for t in templates]
+
+        # Encode all prompts and average. With a trained projector, all prompts
+        # go through each CLIP teacher in one batched forward (C forwards, not
+        # N*C). The legacy pipeline falls back to per-prompt encoding.
+        if self.projector_model is not None and self.whitening_stats is not None:
+            per_prompt = self._encode_texts_projector(prompts)
+        else:
+            per_prompt = np.stack(
+                [self._encode_single_text(p) for p in prompts], axis=0
+            )
+        avg = per_prompt.mean(axis=0)
+        avg = avg / (np.linalg.norm(avg) + 1e-8)
+        return avg.astype(np.float32)
+
+    def _encode_texts_projector(self, texts: List[str]) -> np.ndarray:
+        """Batched projector-path text encoding.
+
+        Returns an ``(N, d_out)`` array of projected, L2-normalized text
+        embeddings, one per input string. All prompts pass through each CLIP
+        teacher in a single padded forward. CLIP's EOS-token pooling is
+        right-padding invariant, so this matches per-prompt encoding.
+        """
+        from skinmap_runtime.embedding_fusion import apply_whiten, l2_normalize
+
+        clip_indices = self.whitening_stats.get("clip_indices", [])
+        if not clip_indices:
+            clip_indices = list(range(len(self.clip_models)))
+        clip_index_set = set(clip_indices)
+        text_whitening = self.whitening_stats.get("text_whitening")
+
+        per_model = []  # each (N, d_i)
+        clip_idx_counter = 0
+        for i, (_model_type, model_wrapper) in enumerate(self.models_in_order):
+            if i not in clip_index_set:
+                continue
+            inputs = model_wrapper.processor(
+                text=list(texts), return_tensors="pt", padding=True
+            )
+            inputs = {k: v.to(self.device) for k, v in inputs.items()}
+            with torch.inference_mode():
+                text_emb = _coerce_clip_features(
+                    model_wrapper.model.get_text_features(**inputs), "text_embeds"
+                )
+                text_emb = text_emb / text_emb.norm(p=2, dim=-1, keepdim=True)
+            feat = text_emb.cpu().numpy()  # (N, d_i)
+            if text_whitening is not None:
+                feat = apply_whiten(
+                    l2_normalize(feat, axis=1),
+                    text_whitening["mu"][clip_idx_counter],
+                    text_whitening["W"][clip_idx_counter],
+                )
+            per_model.append(feat)
+            clip_idx_counter += 1
+
+        # Concatenate per-teacher blocks, L2-normalize each row, project.
+        text_concat = np.concatenate(per_model, axis=1).astype(np.float32)
+        norms = np.linalg.norm(text_concat, axis=1, keepdims=True)
+        norms[norms == 0] = 1.0
+        text_concat = text_concat / norms
+        text_tensor = torch.from_numpy(text_concat).to(self.device)
+        with torch.inference_mode():
+            projected = self.projector_model.txt_head(text_tensor)
+        return projected.cpu().numpy()  # (N, d_out)
+
+    def _encode_single_text(self, text: str) -> np.ndarray:
+        """Encode a single text string (no template expansion)."""
+        if self.projector_model is not None and self.whitening_stats is not None:
+            # Projector pipeline: concatenate text from CLIP models only
+            from skinmap_runtime.embedding_fusion import apply_whiten, l2_normalize
+
+            clip_indices = self.whitening_stats.get("clip_indices", [])
+            if not clip_indices:
+                # Fallback: use all CLIP models if indices not available
+                clip_indices = list(range(len(self.clip_models)))
+
+            text_features = []
+            clip_idx_counter = 0
+            for i, (model_type, model_wrapper) in enumerate(self.models_in_order):
+                if i in clip_indices:
+                    # This is a CLIP model, extract text embedding
+                    inputs = model_wrapper.processor(text=[text], return_tensors="pt")
+                    inputs = {k: v.to(self.device) for k, v in inputs.items()}
+                    with torch.inference_mode():
+                        text_emb = _coerce_clip_features(
+                            model_wrapper.model.get_text_features(**inputs),
+                            "text_embeds",
+                        )
+                        # Normalize
+                        text_emb = text_emb / text_emb.norm(p=2, dim=-1, keepdim=True)
+                        text_feat = text_emb.cpu().numpy()
+
+                        # Apply text whitening if available
+                        text_whitening = self.whitening_stats.get("text_whitening")
+                        if text_whitening is not None:
+                            text_feat_normed = l2_normalize(text_feat, axis=1)
+                            text_feat_whitened = apply_whiten(
+                                text_feat_normed,
+                                text_whitening["mu"][clip_idx_counter],
+                                text_whitening["W"][clip_idx_counter],
+                            )
+                            text_features.append(text_feat_whitened[0])
+                        else:
+                            text_features.append(text_feat[0])
+
+                        clip_idx_counter += 1
+
+            # Concatenate text features from all CLIP models
+            text_concat = np.concatenate(text_features).astype(np.float32)
+
+            # L2-normalize the concatenated text vector (matches training pipeline)
+            norm = np.linalg.norm(text_concat)
+            if norm > 0:
+                text_concat = text_concat / norm
+
+            # Project through text head
+            text_tensor = torch.from_numpy(text_concat).unsqueeze(0).to(self.device)
+            with torch.inference_mode():
+                projected = self.projector_model.txt_head(text_tensor)
+
+            return projected.cpu().numpy()[0]
+        else:
+            # Legacy SVD pipeline or no projector: use first CLIP model
+            clip_model = self.clip_models[0]
+            inputs = clip_model.processor(text=[text], return_tensors="pt")
+            inputs = {k: v.to(self.device) for k, v in inputs.items()}
+            with torch.inference_mode():
+                text_emb = _coerce_clip_features(
+                    clip_model.model.get_text_features(**inputs), "text_embeds"
+                )
+                text_emb = text_emb / text_emb.norm(p=2, dim=-1, keepdim=True)
+
+            # If SVD available, apply it
+            text_feat = text_emb.cpu().numpy()
+            if self.svd_image is not None:
+                # Note: this assumes SVD was fitted on concatenated features
+                # For single model, this might not work well
+                pass
+
+            return text_feat[0]
+
+    def find_nearest_neighbors(
+        self, vector: np.ndarray, k: int = 16
+    ) -> Tuple[List[int], List[float]]:
+        index = self._ensure_faiss_index()
+        vec = vector.astype(np.float32)
+        norm = np.linalg.norm(vec)
+        if norm > 0:
+            vec = vec / norm
+        similarities, indices = index.search(vec.reshape(1, -1), k)
+        flat_indices = indices.reshape(-1).tolist()
+        flat_distances = (1.0 - similarities.reshape(-1)).tolist()
+        return flat_indices, flat_distances
+
+    def project_vector(self, vector: np.ndarray) -> Optional[np.ndarray]:
+        if self.umap_model is None:
+            logger.warning(
+                "Upload embedding pipeline has no UMAP reducer; skipping projection. "
+                "Ensure the embedding pipeline config includes a saved UMAP model."
+            )
+            return None
+        try:
+            projected = self.umap_model.transform(vector.reshape(1, -1))
+            if projected.ndim == 2 and projected.shape[1] >= 2:
+                return projected[0, :2].astype(float)
+        except Exception as exc:
+            logger.warning("Failed to project vector with UMAP: %s", exc)
+        return None
+
+    def search_image(
+        self, image_bytes: bytes, k: int = 16
+    ) -> Tuple[List[int], List[float], Optional[np.ndarray]]:
+        vector = self.embed_bytes(image_bytes)
+        indices, distances = self.find_nearest_neighbors(vector, k=k)
+        coords = self.project_vector(vector)
+        return indices, distances, coords

skinmap_runtime/core/models/byol/model.py

@@ -0,0 +1,42 @@
+from torch import nn
+
+from ..encoders.utils import get_encoder_class
+from ..utils import ModelType
+from .predictor import MLP
+
+
+class BYOLModel(nn.Module):
+    def __init__(
+        self,
+        base_model: str,
+        projection_size=256,
+        projection_hidden_size=4096,
+        **kwargs,
+    ):
+        super(BYOLModel, self).__init__()
+        encoder_cls, model_type = get_encoder_class(base_model)
+        if model_type is ModelType.VIT:
+            self.backbone = encoder_cls(**kwargs)
+            n_feat = self.backbone.embed_dim
+        elif model_type is ModelType.CNN:
+            encoder = encoder_cls(**kwargs)
+            n_feat = encoder.fc.in_features
+            self.backbone = nn.Sequential(*list(encoder.children())[:-1])
+        else:
+            raise ValueError(f"Unknown model type: {model_type}")
+        # projection head
+        self.projection = MLP(
+            in_channels=n_feat,
+            projection_size=projection_size,
+            hidden_size=projection_hidden_size,
+        )
+
+    def forward(self, x, return_embedding=False):
+        # embedding
+        e = self.backbone(x)
+        e = e.squeeze()
+        if return_embedding:
+            return e
+        # project
+        z = self.projection(e)
+        return z

skinmap_runtime/core/models/byol/predictor.py

@@ -0,0 +1,15 @@
+from torch import nn
+
+
+class MLP(nn.Module):
+    def __init__(self, in_channels, hidden_size=256, projection_size=4096):
+        super(MLP, self).__init__()
+        self.net = nn.Sequential(
+            nn.Linear(in_channels, hidden_size),
+            nn.BatchNorm1d(hidden_size),
+            nn.ReLU(inplace=True),
+            nn.Linear(hidden_size, projection_size),
+        )
+
+    def forward(self, x):
+        return self.net(x)

skinmap_runtime/core/models/colorme/model.py

@@ -0,0 +1,49 @@
+import segmentation_models_pytorch as smp
+from torch import nn
+
+
+class MLPHead(nn.Module):
+    def __init__(self, out_features: int = 512):
+        super(MLPHead, self).__init__()
+        self.head = nn.Sequential(
+            nn.Linear(out_features, 10),
+            nn.Softmax(dim=-1),
+        )
+
+    def forward(self, x):
+        return self.head(x)
+
+
+class ColorMeModel(nn.Module):
+    def __init__(self, num_classes=2, encoder_name="resnet18"):
+        super(ColorMeModel, self).__init__()
+        # create our encoder decoder model
+        self.enc_dec_model = smp.Unet(
+            encoder_name=encoder_name,
+            encoder_weights="imagenet",
+            in_channels=1,
+            classes=num_classes,
+        )
+
+        # color distribution MLP head
+        if encoder_name == "resnet18":
+            self.color_dist_mlp = MLPHead(out_features=512)
+        elif encoder_name == "resnet50":
+            self.color_dist_mlp = MLPHead(out_features=2048)
+        else:
+            raise ValueError("Unrecognized encoder")
+
+    def forward(self, x, return_embedding=False):
+        # color reconstruction
+        rec = self.enc_dec_model(x)
+        # retreive embedding
+        emb = self.enc_dec_model.encoder(x)[-1]
+        emb = nn.AdaptiveAvgPool2d((1, 1))(emb)
+        emb = emb.squeeze()
+        # return the embedding if needed
+        if return_embedding:
+            return emb
+        # color distribution
+        dist = self.color_dist_mlp(emb)
+
+        return rec, dist

skinmap_runtime/core/models/dino/head.py

@@ -0,0 +1,52 @@
+import torch.nn as nn
+
+from ...models.utils import trunc_normal_
+
+
+class DINOHead(nn.Module):
+    def __init__(
+        self,
+        in_dim,
+        out_dim,
+        use_bn=False,
+        norm_last_layer=True,
+        n_layers=3,
+        hidden_dim=2048,
+        bottleneck_dim=256,
+    ):
+        super().__init__()
+
+        n_layers = max(n_layers, 1)
+        if n_layers == 1:
+            self.mlp = nn.Linear(in_dim, bottleneck_dim)
+        else:
+            layers = [nn.Linear(in_dim, hidden_dim)]
+            if use_bn:
+                layers.append(nn.BatchNorm1d(hidden_dim))
+            layers.append(nn.GELU())
+            for _ in range(n_layers - 2):
+                layers.append(nn.Linear(hidden_dim, hidden_dim))
+                if use_bn:
+                    layers.append(nn.BatchNorm1d(hidden_dim))
+                layers.append(nn.GELU())
+            layers.append(nn.Linear(hidden_dim, bottleneck_dim))
+            self.mlp = nn.Sequential(*layers)
+        self.apply(self._init_weights)
+        self.last_layer = nn.utils.weight_norm(
+            nn.Linear(bottleneck_dim, out_dim, bias=False)
+        )
+        self.last_layer.weight_g.data.fill_(1)
+        if norm_last_layer:
+            self.last_layer.weight_g.requires_grad = False
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=0.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+
+    def forward(self, x):
+        x = self.mlp(x)
+        x = nn.functional.normalize(x, dim=-1, p=2)
+        x = self.last_layer(x)
+        return x

skinmap_runtime/core/models/dino/multi_crop_wrapper.py

@@ -0,0 +1,70 @@
+import torch
+from torch import nn
+
+
+class MultiCropWrapper(nn.Module):
+    """
+    Perform forward pass separately on each resolution input.
+    The inputs corresponding to a single resolution are clubbed and single
+    forward is run on the same resolution inputs. Hence we do several
+    forward passes = number of different resolutions used. We then
+    concatenate all the output features and run the head forward on these
+    concatenated features.
+    """
+
+    def __init__(
+        self,
+        backbone: torch.nn.Module,
+        head: torch.nn.Module,
+        apply_l2_norm: float = True,
+    ):
+        super(MultiCropWrapper, self).__init__()
+        self.apply_l2_norm = apply_l2_norm
+        # disable layers dedicated to ImageNet labels classification
+        backbone.fc, backbone.head = nn.Identity(), nn.Identity()
+        self.backbone = backbone
+        if head is None:
+            self.head = nn.Identity()
+        else:
+            self.head = head
+
+    def forward(self, x, mask=None, return_backbone_feat=False, **kwargs):
+        # convert to list
+        if not isinstance(x, list):
+            x = [x]
+            mask = [mask] if mask is not None else None
+        idx_crops = torch.cumsum(
+            torch.unique_consecutive(
+                torch.tensor([inp.shape[-1] for inp in x]),
+                return_counts=True,
+            )[1],
+            0,
+        )
+
+        # start concatenating the features
+        start_idx = 0
+        output = torch.empty(0).to(x[0].device)
+        for end_idx in idx_crops:
+            inp_x = torch.cat(x[start_idx:end_idx])
+
+            if mask is not None:
+                inp_m = torch.cat(mask[start_idx:end_idx])
+                kwargs.update(dict(mask=inp_m))
+
+            _out = self.backbone(inp_x, **kwargs)
+            if start_idx == 0:
+                output = _out
+            else:
+                output = torch.cat((output, _out))
+            start_idx = end_idx
+
+        # SelfClean: apply L2 norm if requested
+        if self.apply_l2_norm:
+            output = torch.nn.functional.normalize(output, dim=-1, p=2)
+
+        # run the head forward on the concatenated features
+        output_ = self.head(output)
+        if return_backbone_feat:
+            return output, output_
+
+        return output_

skinmap_runtime/core/models/encoders/__init__.py

@@ -0,0 +1 @@
+

skinmap_runtime/core/models/encoders/swin_transformer.py

@@ -0,0 +1,1050 @@
+"""
+Mostly copy-paste from Swin-Transformer library:
+https://github.com/facebookresearch/dino
+https://github.com/microsoft/Swin-Transformer/blob/main/models/swin_transformer.py
+"""
+
+import os
+from functools import partial
+from math import sqrt
+
+import numpy as np
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+import torch.nn.functional as F
+from loguru import logger
+
+from ..utils import DropPath, to_2tuple, trunc_normal_
+
+
+class Mlp(nn.Module):
+    def __init__(
+        self,
+        in_features,
+        hidden_features=None,
+        out_features=None,
+        act_layer=nn.GELU,
+        drop=0.0,
+    ):
+        super(Mlp, self).__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)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+def window_partition(x, window_size):
+    """
+    Args:
+        x: (B, H, W, C)
+        window_size (int): window size
+    Returns:
+        windows: (num_windows*B, window_size, window_size, C)
+    """
+    B, H, W, C = x.shape
+    x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
+    windows = (
+        x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)
+    )
+    return windows
+
+
+def window_reverse(windows, window_size, H, W):
+    """
+    Args:
+        windows: (num_windows*B, window_size, window_size, C)
+        window_size (int): Window size
+        H (int): Height of image
+        W (int): Width of image
+    Returns:
+        x: (B, H, W, C)
+    """
+    B = int(windows.shape[0] / (H * W / window_size / window_size))
+    x = windows.view(
+        B, H // window_size, W // window_size, window_size, window_size, -1
+    )
+    x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)
+    return x
+
+
+class WindowAttention(nn.Module):
+    r"""Window based multi-head self attention (W-MSA) module with relative position bias.
+    It supports both of shifted and non-shifted window.
+    Args:
+        dim (int): Number of input channels.
+        window_size (tuple[int]): The height and width of the window.
+        num_heads (int): Number of attention heads.
+        qkv_bias (bool, optional):  If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set
+        attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0
+        proj_drop (float, optional): Dropout ratio of output. Default: 0.0
+    """
+
+    def __init__(
+        self,
+        dim,
+        window_size,
+        num_heads,
+        qkv_bias=True,
+        qk_scale=None,
+        attn_drop=0.0,
+        proj_drop=0.0,
+    ):
+        super(WindowAttention, self).__init__()
+        self.dim = dim
+        self.window_size = window_size  # Wh, Ww
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim**-0.5
+
+        # define a parameter table of relative position bias
+        self.relative_position_bias_table = nn.Parameter(
+            torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads)
+        )  # 2*Wh-1 * 2*Ww-1, nH
+
+        # get pair-wise relative position index for each token inside the window
+        coords_h = torch.arange(self.window_size[0])
+        coords_w = torch.arange(self.window_size[1])
+        coords = torch.stack(torch.meshgrid([coords_h, coords_w]))  # 2, Wh, Ww
+        coords_flatten = torch.flatten(coords, 1)  # 2 Wh*Ww
+        relative_coords = (
+            coords_flatten[:, :, None] - coords_flatten[:, None, :]
+        )  # 2, Wh*Ww, Wh*Ww
+        relative_coords = relative_coords.permute(
+            1, 2, 0
+        ).contiguous()  # Wh*Ww, Wh*Ww, 2
+        relative_coords[:, :, 0] += self.window_size[0] - 1  # shift to start from 0
+        relative_coords[:, :, 1] += self.window_size[1] - 1
+        relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
+        relative_position_index = relative_coords.sum(-1)  # Wh*Ww, Wh*Ww
+        self.register_buffer("relative_position_index", relative_position_index)
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+        trunc_normal_(self.relative_position_bias_table, std=0.02)
+        self.softmax = nn.Softmax(dim=-1)
+
+    def forward(self, x, mask=None):
+        """
+        Args:
+            x: input features with shape of (num_windows*B, N, C)
+            mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None
+        """
+        B_, N, C = x.shape
+        qkv = (
+            self.qkv(x)
+            .reshape(B_, N, 3, self.num_heads, C // self.num_heads)
+            .permute(2, 0, 3, 1, 4)
+        )
+        q, k, v = (
+            qkv[0],
+            qkv[1],
+            qkv[2],
+        )  # make torchscript happy (cannot use tensor as tuple)
+
+        q = q * self.scale
+        attn = q @ k.transpose(-2, -1)
+
+        relative_position_bias = self.relative_position_bias_table[
+            self.relative_position_index.view(-1)
+        ].view(
+            self.window_size[0] * self.window_size[1],
+            self.window_size[0] * self.window_size[1],
+            -1,
+        )  # Wh*Ww,Wh*Ww,nH
+        relative_position_bias = relative_position_bias.permute(
+            2, 0, 1
+        ).contiguous()  # nH, Wh*Ww, Wh*Ww
+        attn = attn + relative_position_bias.unsqueeze(0)
+
+        if mask is not None:
+            nW = mask.shape[0]
+            attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(
+                1
+            ).unsqueeze(0)
+            attn = attn.view(-1, self.num_heads, N, N)
+            attn = self.softmax(attn)
+        else:
+            attn = self.softmax(attn)
+
+        attn_out = attn
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B_, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x, attn_out
+
+    def extra_repr(self) -> str:
+        return f"dim={self.dim}, window_size={self.window_size}, num_heads={self.num_heads}"
+
+    def flops(self, N):
+        # calculate flops for 1 window with token length of N
+        flops = 0
+        # qkv = self.qkv(x)
+        flops += N * self.dim * 3 * self.dim
+        # attn = (q @ k.transpose(-2, -1))
+        flops += self.num_heads * N * (self.dim // self.num_heads) * N
+        #  x = (attn @ v)
+        flops += self.num_heads * N * N * (self.dim // self.num_heads)
+        # x = self.proj(x)
+        flops += N * self.dim * self.dim
+        return flops
+
+    @staticmethod
+    def compute_macs(module, input, output):
+        B, N, C = input[0].shape
+
+        module.__flops__ += module.flops(N) * B
+
+
+class SwinTransformerBlock(nn.Module):
+    r"""Swin Transformer Block.
+    Args:
+        dim (int): Number of input channels.
+        input_resolution (tuple[int]): Input resolution.
+        num_heads (int): Number of attention heads.
+        window_size (int): Window size.
+        shift_size (int): Shift size for SW-MSA.
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
+        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
+        drop (float, optional): Dropout rate. Default: 0.0
+        attn_drop (float, optional): Attention dropout rate. Default: 0.0
+        drop_path (float, optional): Stochastic depth rate. Default: 0.0
+        act_layer (nn.Module, optional): Activation layer. Default: nn.GELU
+        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
+    """
+
+    def __init__(
+        self,
+        dim,
+        input_resolution,
+        num_heads,
+        window_size=7,
+        shift_size=0,
+        mlp_ratio=4.0,
+        qkv_bias=True,
+        qk_scale=None,
+        drop=0.0,
+        attn_drop=0.0,
+        drop_path=0.0,
+        act_layer=nn.GELU,
+        norm_layer=nn.LayerNorm,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.input_resolution = input_resolution
+        self.num_heads = num_heads
+        self.window_size = window_size
+        self.shift_size = shift_size
+        self.mlp_ratio = mlp_ratio
+        if min(self.input_resolution) <= self.window_size:
+            # if window size is larger than input resolution, we don't partition windows
+            self.shift_size = 0
+            self.window_size = min(self.input_resolution)
+        assert (
+            0 <= self.shift_size < self.window_size
+        ), "shift_size must in 0-window_size"
+
+        self.norm1 = norm_layer(dim)
+        self.attn = WindowAttention(
+            dim,
+            window_size=to_2tuple(self.window_size),
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            qk_scale=qk_scale,
+            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,
+        )
+
+        self.H = input_resolution[0]
+        self.W = input_resolution[1]
+
+        self.attn_mask_dict = {}  # {self.H: self.create_attn_mask(self.H, self.W)}
+
+    def create_attn_mask(self, H, W):
+        # calculate attention mask for SW-MSA
+
+        Hp = int(np.ceil(H / self.window_size)) * self.window_size
+        Wp = int(np.ceil(W / self.window_size)) * self.window_size
+        img_mask = torch.zeros((1, Hp, Wp, 1))  # 1 Hp Wp 1
+        h_slices = (
+            slice(0, -self.window_size),
+            slice(-self.window_size, -self.shift_size),
+            slice(-self.shift_size, None),
+        )
+        w_slices = (
+            slice(0, -self.window_size),
+            slice(-self.window_size, -self.shift_size),
+            slice(-self.shift_size, None),
+        )
+        cnt = 0
+        for h in h_slices:
+            for w in w_slices:
+                img_mask[:, h, w, :] = cnt
+                cnt += 1
+
+        mask_windows = window_partition(
+            img_mask, self.window_size
+        )  # nW, window_size, window_size, 1
+        mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
+        attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
+        attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(
+            attn_mask == 0, float(0.0)
+        )
+
+        return attn_mask
+
+    def forward(self, x):
+        B, L, C = x.shape
+        H = int(sqrt(L))
+        W = H
+
+        shortcut = x
+        x = self.norm1(x)
+        x = x.view(B, H, W, C)
+
+        # pad feature maps to multiples of window size
+        pad_l = pad_t = 0
+        pad_r = (self.window_size - W % self.window_size) % self.window_size
+        pad_b = (self.window_size - H % self.window_size) % self.window_size
+        x = F.pad(x, (0, 0, pad_l, pad_r, pad_t, pad_b))
+        _, Hp, Wp, _ = x.shape
+
+        # cyclic shift
+        if self.shift_size > 0:
+            shifted_x = torch.roll(
+                x,
+                shifts=(-self.shift_size, -self.shift_size),
+                dims=(1, 2),
+            )
+
+            if H is self.attn_mask_dict.keys():
+                attn_mask = self.attn_mask_dict[H]
+            else:
+                self.attn_mask_dict[H] = self.create_attn_mask(H, W).to(x.device)
+                attn_mask = self.attn_mask_dict[H]
+
+        else:
+            shifted_x = x
+            attn_mask = None
+
+        # partition windows
+        x_windows = window_partition(
+            shifted_x, self.window_size
+        )  # nW*B, window_size, window_size, C
+        x_windows = x_windows.view(
+            -1, self.window_size * self.window_size, C
+        )  # nW*B, window_size*window_size, C
+
+        # W-MSA/SW-MSA
+        attn_windows, attn = self.attn(
+            x_windows, attn_mask
+        )  # nW*B, window_size*window_size, C
+
+        # merge windows
+        attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)
+        shifted_x = window_reverse(attn_windows, self.window_size, Hp, Wp)  # B H' W' C
+
+        # reverse cyclic shift
+        if self.shift_size > 0:
+            x = torch.roll(
+                shifted_x,
+                shifts=(self.shift_size, self.shift_size),
+                dims=(1, 2),
+            )
+        else:
+            x = shifted_x
+
+        if pad_r > 0 or pad_b > 0:
+            x = x[:, :H, :W, :].contiguous()
+
+        x = x.view(B, H * W, C)
+
+        # FFN
+        x = shortcut + self.drop_path(x)
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+
+        return x, attn
+
+    def extra_repr(self) -> str:
+        return (
+            f"dim={self.dim}, input_resolution={self.input_resolution}, num_heads={self.num_heads}, "
+            f"window_size={self.window_size}, shift_size={self.shift_size} mlp_ratio={self.mlp_ratio}"
+        )
+
+    def flops(self):
+        flops = 0
+        H, W = self.input_resolution
+        # norm1
+        flops += self.dim * H * W
+        # W-MSA/SW-MSA
+        nW = H * W / self.window_size / self.window_size
+        flops += nW * self.attn.flops(self.window_size * self.window_size)
+        # mlp
+        flops += 2 * H * W * self.dim * self.dim * self.mlp_ratio
+        # norm2
+        flops += self.dim * H * W
+        return flops
+
+
+class PatchMerging(nn.Module):
+    r"""Patch Merging Layer.
+    Args:
+        input_resolution (tuple[int]): Resolution of input feature.
+        dim (int): Number of input channels.
+        norm_layer (nn.Module, optional): Normalization layer.  Default: nn.LayerNorm
+    """
+
+    def __init__(self, input_resolution, dim, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.input_resolution = input_resolution
+        self.dim = dim
+        self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
+        self.norm = norm_layer(4 * dim)
+
+    def forward(self, x):
+        """Forward function.
+        Args:
+            x: Input feature, tensor size (B, H*W, C).
+            H, W: Spatial resolution of the input feature.
+        """
+        B, L, C = x.shape
+        H = int(sqrt(L))
+        W = H
+
+        x = x.view(B, H, W, C)
+
+        # padding
+        pad_input = (H % 2 == 1) or (W % 2 == 1)
+        if pad_input:
+            x = F.pad(x, (0, 0, 0, W % 2, 0, H % 2))
+
+        x0 = x[:, 0::2, 0::2, :]  # B H/2 W/2 C
+        x1 = x[:, 1::2, 0::2, :]  # B H/2 W/2 C
+        x2 = x[:, 0::2, 1::2, :]  # B H/2 W/2 C
+        x3 = x[:, 1::2, 1::2, :]  # B H/2 W/2 C
+        x = torch.cat([x0, x1, x2, x3], -1)  # B H/2 W/2 4*C
+        x = x.view(B, -1, 4 * C)  # B H/2*W/2 4*C
+
+        x = self.norm(x)
+        x = self.reduction(x)
+
+        return x
+
+    def extra_repr(self) -> str:
+        return f"input_resolution={self.input_resolution}, dim={self.dim}"
+
+    def flops(self):
+        H, W = self.input_resolution
+        flops = H * W * self.dim
+        flops += (H // 2) * (W // 2) * 4 * self.dim * 2 * self.dim
+        return flops
+
+
+class BasicLayer(nn.Module):
+    """A basic Swin Transformer layer for one stage.
+    Args:
+        dim (int): Number of input channels.
+        input_resolution (tuple[int]): Input resolution.
+        depth (int): Number of blocks.
+        num_heads (int): Number of attention heads.
+        window_size (int): Window size.
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
+        qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
+        drop (float, optional): Dropout rate. Default: 0.0
+        attn_drop (float, optional): Attention dropout rate. Default: 0.0
+        drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
+        norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
+        downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
+    """
+
+    def __init__(
+        self,
+        dim,
+        input_resolution,
+        depth,
+        num_heads,
+        window_size,
+        mlp_ratio=4.0,
+        qkv_bias=True,
+        qk_scale=None,
+        drop=0.0,
+        attn_drop=0.0,
+        drop_path=0.0,
+        norm_layer=nn.LayerNorm,
+        downsample=None,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.input_resolution = input_resolution
+        self.depth = depth
+
+        self.blocks = nn.ModuleList(
+            [
+                SwinTransformerBlock(
+                    dim=dim,
+                    input_resolution=input_resolution,
+                    num_heads=num_heads,
+                    window_size=window_size,
+                    shift_size=0 if (i % 2 == 0) else window_size // 2,
+                    mlp_ratio=mlp_ratio,
+                    qkv_bias=qkv_bias,
+                    qk_scale=qk_scale,
+                    drop=drop,
+                    attn_drop=attn_drop,
+                    drop_path=(
+                        drop_path[i] if isinstance(drop_path, list) else drop_path
+                    ),
+                    norm_layer=norm_layer,
+                )
+                for i in range(depth)
+            ]
+        )
+        if downsample is not None:
+            self.downsample = downsample(
+                input_resolution,
+                dim=dim,
+                norm_layer=norm_layer,
+            )
+        else:
+            self.downsample = None
+
+    def forward(self, x):
+        for blk in self.blocks:
+            x, _ = blk(x)
+        if self.downsample is not None:
+            x = self.downsample(x)
+        return x
+
+    def forward_with_features(self, x):
+        fea = []
+        for blk in self.blocks:
+            x, _ = blk(x)
+            fea.append(x)
+        if self.downsample is not None:
+            x = self.downsample(x)
+        return x, fea
+
+    def forward_with_attention(self, x):
+        attns = []
+        for blk in self.blocks:
+            x, attn = blk(x)
+            attns.append(attn)
+        if self.downsample is not None:
+            x = self.downsample(x)
+        return x, attns
+
+    def extra_repr(self) -> str:
+        return f"dim={self.dim}, input_resolution={self.input_resolution}, depth={self.depth}"
+
+    def flops(self):
+        flops = 0
+        for blk in self.blocks:
+            flops += blk.flops()
+        if self.downsample is not None:
+            flops += self.downsample.flops()
+        return flops
+
+
+class PatchEmbed(nn.Module):
+    """Image to Patch Embedding"""
+
+    def __init__(
+        self,
+        img_size=224,
+        patch_size=16,
+        in_chans=3,
+        embed_dim=768,
+        norm_layer=None,
+    ):
+        super().__init__()
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        patches_resolution = [
+            img_size[0] // patch_size[0],
+            img_size[1] // patch_size[1],
+        ]
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.patches_resolution = patches_resolution
+        self.num_patches = patches_resolution[0] * patches_resolution[1]
+
+        self.in_chans = in_chans
+        self.embed_dim = embed_dim
+
+        self.proj = nn.Conv2d(
+            in_chans,
+            embed_dim,
+            kernel_size=patch_size,
+            stride=patch_size,
+        )
+        if norm_layer is not None:
+            self.norm = norm_layer(embed_dim)
+        else:
+            self.norm = None
+
+    def forward(self, x):
+        # # FIXME look at relaxing size constraints
+        # assert H == self.img_size[0] and W == self.img_size[1], \
+        #     f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
+
+        x = self.proj(x)
+        B, C, H, W = x.shape
+        x = x.flatten(2).transpose(1, 2)  # B Ph*Pw C
+        if self.norm is not None:
+            x = self.norm(x)
+        return x.transpose(1, 2).reshape(B, C, H, W)
+
+    def flops(self):
+        Ho, Wo = self.patches_resolution
+        flops = (
+            Ho
+            * Wo
+            * self.embed_dim
+            * self.in_chans
+            * (self.patch_size[0] * self.patch_size[1])
+        )
+        if self.norm is not None:
+            flops += Ho * Wo * self.embed_dim
+        return flops
+
+
+class SwinTransformer(nn.Module):
+    r"""Swin Transformer
+        A PyTorch impl of : `Swin Transformer: Hierarchical Vision Transformer using Shifted Windows`  -
+          https://arxiv.org/pdf/2103.14030
+    Args:
+        img_size (int | tuple(int)): Input image size.
+        patch_size (int | tuple(int)): Patch size.
+        in_chans (int): Number of input channels.
+        num_classes (int): Number of classes for classification head.
+        embed_dim (int): Embedding dimension.
+        depths (tuple(int)): Depth of Swin Transformer layers.
+        num_heads (tuple(int)): Number of attention heads in different layers.
+        window_size (int): Window size.
+        mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
+        qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True
+        qk_scale (float): Override default qk scale of head_dim ** -0.5 if set.
+        drop_rate (float): Dropout rate.
+        attn_drop_rate (float): Attention dropout rate.
+        drop_path_rate (float): Stochastic depth rate.
+        norm_layer (nn.Module): normalization layer.
+        ape (bool): If True, add absolute position embedding to the patch embedding.
+        patch_norm (bool): If True, add normalization after patch embedding.
+    """
+
+    def __init__(
+        self,
+        img_size=224,
+        patch_size=4,
+        in_chans=3,
+        num_classes=1000,
+        embed_dim=96,
+        depths=[2, 2, 6, 2],
+        num_heads=[3, 6, 12, 24],
+        window_size=7,
+        mlp_ratio=4.0,
+        qkv_bias=True,
+        qk_scale=None,
+        drop_rate=0.0,
+        attn_drop_rate=0.0,
+        drop_path_rate=0.0,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),
+        ape=False,
+        patch_norm=True,
+        return_all_tokens=False,
+        use_mean_pooling=True,
+        masked_im_modeling=False,
+    ):
+        super().__init__()
+
+        self.num_classes = num_classes
+        self.depths = depths
+        self.num_layers = len(depths)
+        self.embed_dim = embed_dim
+        self.ape = ape
+        self.patch_norm = patch_norm
+        self.num_features = int(embed_dim * 2 ** (self.num_layers - 1))
+        self.mlp_ratio = mlp_ratio
+        self.return_all_tokens = return_all_tokens
+
+        self.patch_embed = PatchEmbed(
+            img_size=img_size,
+            patch_size=patch_size,
+            in_chans=in_chans,
+            embed_dim=embed_dim,
+            norm_layer=norm_layer if self.patch_norm else None,
+        )
+        num_patches = self.patch_embed.num_patches
+        patches_resolution = self.patch_embed.patches_resolution
+        self.patches_resolution = patches_resolution
+
+        if self.ape:
+            self.absolute_pos_embed = nn.Parameter(
+                torch.zeros(1, num_patches, embed_dim)
+            )
+            trunc_normal_(self.absolute_pos_embed, std=0.02)
+
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        dpr = [
+            x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))
+        ]  # stochastic depth decay rule
+        self.layers = nn.ModuleList()
+        for i_layer in range(self.num_layers):
+            layer = BasicLayer(
+                dim=int(embed_dim * 2**i_layer),
+                input_resolution=(
+                    patches_resolution[0] // (2**i_layer),
+                    patches_resolution[1] // (2**i_layer),
+                ),
+                depth=depths[i_layer],
+                num_heads=num_heads[i_layer],
+                window_size=window_size,
+                mlp_ratio=self.mlp_ratio,
+                qkv_bias=qkv_bias,
+                qk_scale=qk_scale,
+                drop=drop_rate,
+                attn_drop=attn_drop_rate,
+                drop_path=dpr[sum(depths[:i_layer]) : sum(depths[: i_layer + 1])],
+                norm_layer=norm_layer,
+                downsample=PatchMerging if (i_layer < self.num_layers - 1) else None,
+            )
+            self.layers.append(layer)
+
+        self.norm = norm_layer(self.num_features)
+        self.avgpool = nn.AdaptiveAvgPool1d(1)
+        self.head = (
+            nn.Linear(self.num_features, num_classes)
+            if num_classes > 0
+            else nn.Identity()
+        )
+
+        self.apply(self._init_weights)
+
+        # masked image modeling
+        self.masked_im_modeling = masked_im_modeling
+        if masked_im_modeling:
+            self.masked_embed = nn.Parameter(torch.zeros(1, embed_dim))
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=0.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {"absolute_pos_embed"}
+
+    @torch.jit.ignore
+    def no_weight_decay_keywords(self):
+        # todo: to be implemented
+        return {"relative_position_bias_table"}
+
+    def forward(self, x, return_all_tokens=None, mask=None):
+        # patch linear embedding
+        x = self.patch_embed(x)
+        # mask image modeling
+        if mask is not None:
+            x = self.mask_model(x, mask)
+        x = x.flatten(2).transpose(1, 2)
+
+        if self.ape:
+            x = x + self.absolute_pos_embed
+        x = self.pos_drop(x)
+
+        for layer in self.layers:
+            x = layer(x)
+
+        x_region = self.norm(x)  # B L C
+        x = self.avgpool(x_region.transpose(1, 2))  # B C 1
+        x = torch.flatten(x, 1)
+
+        return_all_tokens = (
+            self.return_all_tokens if return_all_tokens is None else return_all_tokens
+        )
+        if return_all_tokens:
+            return torch.cat([x.unsqueeze(1), x_region], dim=1)
+        return x
+
+    def prepare_tokens(self, x):
+        x = self.patch_embed(x)
+        x = x.flatten(2).transpose(1, 2)
+        if self.ape:
+            x = x + self.absolute_pos_embed
+        return self.pos_drop(x)
+
+    def get_last_selfattention(self, x):
+        x = self.prepare_tokens(x)
+        for i, layer in enumerate(self.layers):
+            if i < len(self.layers) - 1:
+                x = layer(x)
+            else:
+                x, attns = layer.forward_with_attention(x)
+                return attns[-1]
+
+    def get_all_selfattention(self, x):
+        x = self.prepare_tokens(x)
+        attn_out = []
+        for layer in self.layers:
+            x, attns = layer.forward_with_attention(x)
+            attn_out += attns
+        return attn_out
+
+    def get_intermediate_layers(self, x, n=1, return_patch_avgpool=False):
+        num_blks = sum(self.depths)
+        start_idx = num_blks - n
+
+        sum_cur = 0
+        for i, d in enumerate(self.depths):
+            sum_cur_new = sum_cur + d
+            if start_idx >= sum_cur and start_idx < sum_cur_new:
+                start_stage = i
+                start_blk = start_idx - sum_cur
+            sum_cur = sum_cur_new
+
+        x = self.patch_embed(x)
+        x = x.flatten(2).transpose(1, 2)
+        if self.ape:
+            x = x + self.absolute_pos_embed
+        x = self.pos_drop(x)
+
+        # we will return the averaged token features from the `n` last blocks
+        # note: there is no [CLS] token in Swin Transformer
+        output = []
+        s = 0
+        for i, layer in enumerate(self.layers):
+            x, fea = layer.forward_with_features(x)
+
+            if i >= start_stage:
+                for x_ in fea[start_blk:]:
+                    if i == len(self.layers) - 1:  # use the norm in the last stage
+                        x_ = self.norm(x_)
+
+                    x_avg = torch.flatten(self.avgpool(x_.transpose(1, 2)), 1)  # B C
+                    if return_patch_avgpool:
+                        x_o = x_avg
+                    else:
+                        x_o = torch.cat((x_avg.unsqueeze(1), x_), dim=1)
+                    logger.debug(f"Stage {i},  x_o {x_o.shape}")
+                    output.append(x_o)
+
+                start_blk = 0
+
+        # return torch.cat(output, dim=-1)
+        return output
+
+    def flops(self):
+        flops = 0
+        flops += self.patch_embed.flops()
+        for i, layer in enumerate(self.layers):
+            flops += layer.flops()
+            if dist.get_rank() == 0:
+                logger.debug(f"GFLOPs layer_{i}: {layer.flops() / 1e9}")
+        flops += (
+            self.num_features
+            * self.patches_resolution[0]
+            * self.patches_resolution[1]
+            // (2**self.num_layers)
+        )
+        flops += self.num_features * self.num_classes
+        return flops
+
+    def init_weights(self, pretrained="", pretrained_layers=[], verbose=True):
+        if os.path.isfile(pretrained):
+            pretrained_dict = torch.load(pretrained, map_location="cpu")
+            logger.info(f"=> loading pretrained model {pretrained}")
+            model_dict = self.state_dict()
+            pretrained_dict = {
+                k: v for k, v in pretrained_dict.items() if k in model_dict.keys()
+            }
+            need_init_state_dict = {}
+            for k, v in pretrained_dict.items():
+                need_init = (
+                    k.split(".")[0] in pretrained_layers
+                    or pretrained_layers[0] == "*"
+                    or "relative_position_index" not in k
+                    or "attn_mask" not in k
+                )
+
+                if need_init:
+                    logger.info(f"=> init {k} from {pretrained}")
+                    if (
+                        "relative_position_bias_table" in k
+                        and v.size() != model_dict[k].size()
+                    ):
+                        relative_position_bias_table_pretrained = v
+                        relative_position_bias_table_current = model_dict[k]
+                        L1, nH1 = relative_position_bias_table_pretrained.size()
+                        L2, nH2 = relative_position_bias_table_current.size()
+                        if nH1 != nH2:
+                            logger.info(f"Error in loading {k}, passing")
+                        else:
+                            if L1 != L2:
+                                logger.info(
+                                    "=> load_pretrained: resized variant: {} to {}".format(
+                                        (L1, nH1), (L2, nH2)
+                                    )
+                                )
+                                S1 = int(L1**0.5)
+                                S2 = int(L2**0.5)
+                                relative_position_bias_table_pretrained_resized = (
+                                    torch.nn.functional.interpolate(
+                                        relative_position_bias_table_pretrained.permute(
+                                            1, 0
+                                        ).view(1, nH1, S1, S1),
+                                        size=(S2, S2),
+                                        mode="bicubic",
+                                    )
+                                )
+                                v = relative_position_bias_table_pretrained_resized.view(
+                                    nH2, L2
+                                ).permute(
+                                    1, 0
+                                )
+
+                    if "absolute_pos_embed" in k and v.size() != model_dict[k].size():
+                        absolute_pos_embed_pretrained = v
+                        absolute_pos_embed_current = model_dict[k]
+                        _, L1, C1 = absolute_pos_embed_pretrained.size()
+                        _, L2, C2 = absolute_pos_embed_current.size()
+                        if C1 != C1:
+                            logger.info(f"Error in loading {k}, passing")
+                        else:
+                            if L1 != L2:
+                                logger.info(
+                                    "=> load_pretrained: resized variant: {} to {}".format(
+                                        (1, L1, C1), (1, L2, C2)
+                                    )
+                                )
+                                S1 = int(L1**0.5)
+                                S2 = int(L2**0.5)
+                                absolute_pos_embed_pretrained = (
+                                    absolute_pos_embed_pretrained.reshape(
+                                        -1, S1, S1, C1
+                                    )
+                                )
+                                absolute_pos_embed_pretrained = (
+                                    absolute_pos_embed_pretrained.permute(0, 3, 1, 2)
+                                )
+                                absolute_pos_embed_pretrained_resized = (
+                                    torch.nn.functional.interpolate(
+                                        absolute_pos_embed_pretrained,
+                                        size=(S2, S2),
+                                        mode="bicubic",
+                                    )
+                                )
+                                v = absolute_pos_embed_pretrained_resized.permute(
+                                    0, 2, 3, 1
+                                ).flatten(1, 2)
+
+                    need_init_state_dict[k] = v
+            self.load_state_dict(need_init_state_dict, strict=False)
+
+    def freeze_pretrained_layers(self, frozen_layers=[]):
+        for name, module in self.named_modules():
+            if (
+                name.split(".")[0] in frozen_layers
+                or ".".join(name.split(".")[0:2]) in frozen_layers
+                or (len(frozen_layers) > 0 and frozen_layers[0] == "*")
+            ):
+                for _name, param in module.named_parameters():
+                    param.requires_grad = False
+                logger.info("=> set param {} requires grad to False".format(name))
+        for name, param in self.named_parameters():
+            if (
+                name.split(".")[0] in frozen_layers
+                or (len(frozen_layers) > 0 and frozen_layers[0] == "*")
+                and param.requires_grad is True
+            ):
+                param.requires_grad = False
+                logger.info("=> set param {} requires grad to False".format(name))
+        return self
+
+    def get_num_layers(self):
+        # return len(self.layers)
+        return sum(self.depths)
+
+    def mask_model(self, x, mask):
+        # extend mask for hierarchical features
+        if x.shape[-2:] != mask.shape[-2:]:
+            htimes, wtimes = np.array(x.shape[-2:]) // np.array(mask.shape[-2:])
+            mask = mask.repeat_interleave(htimes, -2).repeat_interleave(wtimes, -1)
+
+        # mask embed
+        x.permute(0, 2, 3, 1)[mask, :] = self.masked_embed.to(x.dtype)
+
+        return x
+
+
+def swin_tiny(window_size=7, **kwargs):
+    model = SwinTransformer(
+        window_size=window_size,
+        embed_dim=96,
+        depths=[2, 2, 6, 2],
+        num_heads=[3, 6, 12, 24],
+        mlp_ratio=4,
+        qkv_bias=True,
+        drop_path_rate=kwargs.pop("drop_path_rate", 0.1),
+        **kwargs,
+    )
+    return model
+
+
+def swin_small(window_size=7, **kwargs):
+    model = SwinTransformer(
+        window_size=window_size,
+        embed_dim=96,
+        depths=[2, 2, 18, 2],
+        num_heads=[3, 6, 12, 24],
+        mlp_ratio=4,
+        qkv_bias=True,
+        drop_path_rate=kwargs.pop("drop_path_rate", 0.2),
+        **kwargs,
+    )
+    return model
+
+
+def swin_base(window_size=7, **kwargs):
+    model = SwinTransformer(
+        window_size=window_size,
+        embed_dim=128,
+        depths=[2, 2, 18, 2],
+        num_heads=[4, 8, 16, 32],
+        mlp_ratio=4,
+        qkv_bias=True,
+        drop_path_rate=kwargs.pop("drop_path_rate", 0.2),
+        **kwargs,
+    )
+    return model

skinmap_runtime/core/models/encoders/utils.py

@@ -0,0 +1,34 @@
+from typing import Callable, Tuple
+
+from torchvision import models as torchvision_models
+
+from ....src.models.encoders.swin_transformer import swin_base, swin_small, swin_tiny
+from ....src.models.encoders.vision_transformer import (
+    vit_base,
+    vit_large,
+    vit_small,
+    vit_tiny,
+)
+from ....src.models.utils import ModelType
+
+VIT_DICT = {
+    "vit_tiny": vit_tiny,
+    "vit_small": vit_small,
+    "vit_base": vit_base,
+    "vit_large": vit_large,
+    "swin_tiny": swin_tiny,
+    "swin_small": swin_small,
+    "swin_base": swin_base,
+}
+
+
+def get_encoder_class(base_model_name: str) -> Tuple[Callable, ModelType]:
+    encoder_cls = VIT_DICT.get(base_model_name, None)
+    model_type = ModelType.VIT
+    if encoder_cls is None:
+        if base_model_name in torchvision_models.__dict__.keys():
+            encoder_cls = torchvision_models.__dict__[base_model_name]
+            model_type = ModelType.CNN
+        else:
+            raise ValueError(f"Invalid base model name: {base_model_name}")
+    return encoder_cls, model_type

skinmap_runtime/core/models/encoders/vision_transformer.py

@@ -0,0 +1,374 @@
+"""
+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 List
+
+import torch
+import torch.nn as nn
+
+from ..utils import DropPath, trunc_normal_
+
+
+class Mlp(nn.Module):
+    def __init__(
+        self,
+        in_features,
+        hidden_features=None,
+        out_features=None,
+        act_layer=nn.GELU,
+        drop=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)
+
+    def forward(self, x):
+        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):
+    def __init__(
+        self,
+        dim,
+        num_heads=8,
+        qkv_bias=False,
+        qk_scale=None,
+        attn_drop=0.0,
+        proj_drop=0.0,
+    ):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim**-0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x):
+        B, N, C = x.shape
+        qkv = (
+            self.qkv(x)
+            .reshape(B, N, 3, self.num_heads, C // self.num_heads)
+            .permute(2, 0, 3, 1, 4)
+        )
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x, attn
+
+
+class Block(nn.Module):
+    def __init__(
+        self,
+        dim,
+        num_heads,
+        mlp_ratio=4.0,
+        qkv_bias=False,
+        qk_scale=None,
+        drop=0.0,
+        attn_drop=0.0,
+        drop_path=0.0,
+        act_layer=nn.GELU,
+        norm_layer=nn.LayerNorm,
+    ):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            qk_scale=qk_scale,
+            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, return_attention=False):
+        y, attn = self.attn(self.norm1(x))
+        if return_attention:
+            return attn
+        x = x + self.drop_path(y)
+        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=224, patch_size=16, in_channels=3, embed_dim=768):
+        super().__init__()
+        num_patches = (img_size // patch_size) * (img_size // patch_size)
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.num_patches = num_patches
+
+        self.proj = nn.Conv2d(
+            in_channels, embed_dim, kernel_size=patch_size, stride=patch_size
+        )
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        return self.proj(x)
+
+
+class VisionTransformer(nn.Module):
+    """Vision Transformer"""
+
+    def __init__(
+        self,
+        img_size: List[int] = [224],
+        patch_size: int = 16,
+        in_channels: 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,
+        qk_scale=None,
+        drop_rate: float = 0.0,
+        attn_drop_rate: float = 0.0,
+        drop_path_rate: float = 0.0,
+        norm_layer: nn.Module = nn.LayerNorm,
+        return_all_tokens: bool = False,
+        masked_im_modeling: bool = False,
+        **kwargs
+    ):
+        super().__init__()
+        self.num_features = self.embed_dim = embed_dim
+        self.return_all_tokens = return_all_tokens
+
+        self.patch_embed = PatchEmbed(
+            img_size=img_size[0],
+            patch_size=patch_size,
+            in_channels=in_channels,
+            embed_dim=embed_dim,
+        )
+        num_patches = self.patch_embed.num_patches
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
+        self.pos_drop = nn.Dropout(p=drop_rate)
+        # stochastic depth decay rule
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]
+        self.blocks = nn.ModuleList(
+            [
+                Block(
+                    dim=embed_dim,
+                    num_heads=num_heads,
+                    mlp_ratio=mlp_ratio,
+                    qkv_bias=qkv_bias,
+                    qk_scale=qk_scale,
+                    drop=drop_rate,
+                    attn_drop=attn_drop_rate,
+                    drop_path=dpr[i],
+                    norm_layer=norm_layer,
+                )
+                for i in range(depth)
+            ]
+        )
+        self.norm = norm_layer(embed_dim)
+
+        # Classifier head
+        if num_classes > 0:
+            self.head = nn.Linear(embed_dim, num_classes)
+        else:
+            self.head = nn.Identity()
+
+        trunc_normal_(self.pos_embed, std=0.02)
+        trunc_normal_(self.cls_token, std=0.02)
+        self.apply(self._init_weights)
+
+        # Masked image modeling (MIM)
+        self.masked_im_modeling = masked_im_modeling
+        if masked_im_modeling:
+            self.masked_embed = nn.Parameter(torch.zeros(1, embed_dim))
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=0.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def interpolate_pos_encoding(self, x, w, h):
+        n_patch = x.shape[1] - 1
+        N = self.pos_embed.shape[1] - 1
+        if n_patch == 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
+        h0 = h // self.patch_embed.patch_size
+        # 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",
+        )
+        assert (
+            int(w0) == patch_pos_embed.shape[-2]
+            and int(h0) == patch_pos_embed.shape[-1]
+        )
+        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, mask=None):
+        B, nc, w, h = x.shape
+        # patch linear embedding
+        x = self.patch_embed(x)
+
+        # MIM
+        if mask is not None:
+            x = self.mask_model(x, mask)
+        x = x.flatten(2).transpose(1, 2)
+
+        # 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: torch.Tensor, return_all_tokens=None, mask=None, **kwargs):
+        # MIM
+        if self.masked_im_modeling:
+            assert mask is not None
+            x = self.prepare_tokens(x, mask=mask)
+        else:
+            x = self.prepare_tokens(x)
+        for blk in self.blocks:
+            x = blk(x)
+        x = self.norm(x)
+        return_all_tokens = (
+            self.return_all_tokens if return_all_tokens is None else return_all_tokens
+        )
+        if return_all_tokens:
+            return x
+
+        return x[:, 0]
+
+    def get_last_selfattention(self, x: torch.Tensor):
+        x = self.prepare_tokens(x)
+        for i, blk in enumerate(self.blocks):
+            if i < len(self.blocks) - 1:
+                x = blk(x)
+            else:
+                # return attention of the last block
+                return blk(x, return_attention=True)
+
+    def get_all_selfattention(self, x: torch.Tensor):
+        x = self.prepare_tokens(x)
+        attn_out = []
+        for blk in self.blocks:
+            attn_out.append(blk(x, return_attention=True))
+            x = blk(x)
+        return attn_out
+
+    def get_intermediate_layers(self, x: torch.Tensor, n: int = 1):
+        x = self.prepare_tokens(x)
+        # we return the output tokens from the `n` last blocks
+        output = []
+        for i, blk in enumerate(self.blocks):
+            x = blk(x)
+            if len(self.blocks) - i <= n:
+                output.append(self.norm(x))
+        return output
+
+    def mask_model(self, x, mask):
+        x.permute(0, 2, 3, 1)[mask, :] = self.masked_embed.to(x.dtype)
+        return x
+
+
+def vit_tiny(patch_size: int = 16, **kwargs):
+    model = VisionTransformer(
+        patch_size=patch_size,
+        embed_dim=192,
+        depth=12,
+        num_heads=3,
+        mlp_ratio=4,
+        qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),
+        **kwargs
+    )
+    return model
+
+
+def vit_small(patch_size: int = 16, **kwargs):
+    model = VisionTransformer(
+        patch_size=patch_size,
+        embed_dim=384,
+        depth=12,
+        num_heads=6,
+        mlp_ratio=4,
+        qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),
+        **kwargs
+    )
+    return model
+
+
+def vit_base(patch_size: int = 16, **kwargs):
+    model = VisionTransformer(
+        patch_size=patch_size,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        mlp_ratio=4,
+        qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),
+        **kwargs
+    )
+    return model
+
+
+def vit_large(patch_size: int = 16, **kwargs):
+    model = VisionTransformer(
+        patch_size=patch_size,
+        embed_dim=1024,
+        depth=24,
+        num_heads=16,
+        mlp_ratio=4,
+        qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),
+        **kwargs
+    )
+    return model

skinmap_runtime/core/models/fine_tuning/__init__.py

@@ -0,0 +1 @@
+

skinmap_runtime/core/models/fine_tuning/classifiers.py

@@ -0,0 +1,69 @@
+from torch import nn
+
+
+class LinearClassifier(nn.Module):
+    """Linear layer to train on top of frozen features"""
+
+    def __init__(
+        self,
+        dim: int,
+        num_labels: int = 1000,
+        use_dropout_in_head: bool = False,
+        dropout_rate: float = 0.3,
+        large_head: bool = True,
+        use_bn: bool = False,
+        log_softmax: bool = False,
+    ):
+        super(LinearClassifier, self).__init__()
+        self.num_labels = num_labels
+        self.large_head = large_head
+        self.use_bn = use_bn
+        self.log_softmax = log_softmax
+        self.use_dropout_in_head = use_dropout_in_head
+
+        if self.use_dropout_in_head:
+            self.dropout = nn.Dropout(dropout_rate)
+        if self.use_bn:
+            self.bn = nn.BatchNorm1d(dim)
+
+        if self.large_head:
+            self.linear = nn.Linear(dim, 128)
+            self.linear.weight.data.normal_(mean=0.0, std=0.01)
+            self.linear.bias.data.zero_()
+            self.relu = nn.ReLU()
+
+            self.dropout2 = nn.Dropout(dropout_rate)
+            if self.use_bn:
+                self.bn2 = nn.BatchNorm1d(128)
+
+            self.linear2 = nn.Linear(128, num_labels)
+            self.linear2.weight.data.normal_(mean=0.0, std=0.01)
+            self.linear2.bias.data.zero_()
+        else:
+            self.linear = nn.Linear(dim, num_labels)
+            self.linear.weight.data.normal_(mean=0.0, std=0.01)
+            self.linear.bias.data.zero_()
+
+    def forward(self, x):
+        # flatten
+        x = x.view(x.size(0), -1)
+        # dropout
+        if self.use_dropout_in_head:
+            x = self.dropout(x)
+        if self.use_bn:
+            x = self.bn(x)
+        # 1. linear layer
+        x = self.linear(x)
+        # smaller version of head
+        if self.large_head:
+            x = self.relu(x)
+            x = self.dropout2(x)
+            if self.use_bn:
+                x = self.bn2(x)
+            # 2. linear layer
+            x = self.linear2(x)
+        # output
+        if self.log_softmax:
+            return nn.LogSoftmax(dim=1)(x)
+        else:
+            return x

skinmap_runtime/core/models/ibot/head.py

@@ -0,0 +1,80 @@
+import torch.nn as nn
+
+from ...models.dino.head import DINOHead
+
+
+class iBOTHead(DINOHead):
+    def __init__(
+        self,
+        *args,
+        patch_out_dim=8192,
+        n_layers=3,
+        hidden_dim=2048,
+        bottleneck_dim=256,
+        norm_last_layer=True,
+        shared_head=False,
+        **kwargs
+    ):
+        super(iBOTHead, self).__init__(
+            *args,
+            n_layers=n_layers,
+            hidden_dim=hidden_dim,
+            bottleneck_dim=bottleneck_dim,
+            norm_last_layer=norm_last_layer,
+            **kwargs
+        )
+
+        if not shared_head:
+            if bottleneck_dim > 0:
+                self.last_layer2 = nn.utils.weight_norm(
+                    nn.Linear(bottleneck_dim, patch_out_dim, bias=False)
+                )
+                self.last_layer2.weight_g.data.fill_(1)
+                if norm_last_layer:
+                    self.last_layer2.weight_g.requires_grad = False
+            else:
+                self.mlp2 = nn.Linear(hidden_dim, patch_out_dim)
+                self.last_layer2 = None
+
+        else:
+            if bottleneck_dim > 0:
+                self.last_layer2 = self.last_layer
+            else:
+                self.mlp2 = self.mlp[-1]
+                self.last_layer2 = None
+
+    def forward(self, x):
+        if len(x.shape) == 2:
+            return super(iBOTHead, self).forward(x)
+
+        if self.last_layer is not None:
+            x = self.mlp(x)
+            x = nn.functional.normalize(x, dim=-1, p=2)
+            x1 = self.last_layer(x[:, 0])
+            x2 = self.last_layer2(x[:, 1:])
+        else:
+            x = self.mlp[:-1](x)
+            x1 = self.mlp[-1](x[:, 0])
+            x2 = self.mlp2(x[:, 1:])
+
+        return x1, x2
+
+    def _build_norm(self, norm, hidden_dim, **kwargs):
+        if norm == "bn":
+            norm = nn.BatchNorm1d(hidden_dim, **kwargs)
+        elif norm == "syncbn":
+            norm = nn.SyncBatchNorm(hidden_dim, **kwargs)
+        elif norm == "ln":
+            norm = nn.LayerNorm(hidden_dim, **kwargs)
+        else:
+            assert norm is None, "unknown norm type {}".format(norm)
+        return norm
+
+    def _build_act(self, act):
+        if act == "relu":
+            act = nn.ReLU()
+        elif act == "gelu":
+            act = nn.GELU()
+        else:
+            assert False, "unknown act type {}".format(act)
+        return act

skinmap_runtime/core/models/mae/model.py

@@ -0,0 +1,297 @@
+from functools import partial
+from typing import List
+
+import torch
+import torch.nn as nn
+
+from ..encoders.vision_transformer import Block, VisionTransformer
+from ..utils import get_2d_sincos_pos_embed
+
+
+class MaskedAutoencoderViT(VisionTransformer):
+    """Masked Autoencoder with VisionTransformer backbone."""
+
+    def __init__(
+        self,
+        img_size: List[int] = [224],
+        patch_size: int = 16,
+        in_channels: 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,
+        qk_scale=None,
+        drop_rate: float = 0.0,
+        attn_drop_rate: float = 0.0,
+        drop_path_rate: float = 0.0,
+        norm_layer: nn.Module = nn.LayerNorm,
+        return_all_tokens: bool = False,
+        masked_im_modeling: bool = False,
+        decoder_embed_dim: int = 512,
+        decoder_depth: int = 8,
+        decoder_num_heads: int = 16,
+    ):
+        super().__init__(
+            img_size=img_size,
+            patch_size=patch_size,
+            in_channels=in_channels,
+            num_classes=num_classes,
+            embed_dim=embed_dim,
+            depth=depth,
+            num_heads=num_heads,
+            mlp_ratio=mlp_ratio,
+            qkv_bias=qkv_bias,
+            qk_scale=qk_scale,
+            drop_rate=drop_rate,
+            attn_drop_rate=attn_drop_rate,
+            drop_path_rate=drop_path_rate,
+            norm_layer=norm_layer,
+            return_all_tokens=return_all_tokens,
+            masked_im_modeling=masked_im_modeling,
+        )
+        num_patches = self.patch_embed.num_patches
+        # fixed sin-cos embedding
+        self.pos_embed.requires_grad = False
+
+        # MAE decoder specifics
+        self.decoder_embed = nn.Linear(embed_dim, decoder_embed_dim, bias=True)
+        self.mask_token = nn.Parameter(torch.zeros(1, 1, decoder_embed_dim))
+        # fixed sin-cos embedding
+        self.decoder_pos_embed = nn.Parameter(
+            data=torch.zeros(1, num_patches + 1, decoder_embed_dim),
+            requires_grad=False,
+        )
+        self.decoder_blocks = nn.ModuleList(
+            [
+                Block(
+                    dim=decoder_embed_dim,
+                    num_heads=decoder_num_heads,
+                    mlp_ratio=mlp_ratio,
+                    qkv_bias=True,
+                    qk_scale=None,
+                    norm_layer=norm_layer,
+                )
+                for _ in range(decoder_depth)
+            ]
+        )
+        # decoder to patch
+        self.decoder_norm = norm_layer(decoder_embed_dim)
+        self.decoder_pred = nn.Linear(
+            decoder_embed_dim,
+            patch_size**2 * in_channels,
+            bias=True,
+        )
+        self.initialize_weights()
+
+    def initialize_weights(self):
+        # initialize (and freeze) pos_embed by sin-cos embedding
+        pos_embed = get_2d_sincos_pos_embed(
+            self.pos_embed.shape[-1],
+            int(self.patch_embed.num_patches**0.5),
+            cls_token=True,
+        )
+        self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))
+
+        decoder_pos_embed = get_2d_sincos_pos_embed(
+            self.decoder_pos_embed.shape[-1],
+            int(self.patch_embed.num_patches**0.5),
+            cls_token=True,
+        )
+        self.decoder_pos_embed.data.copy_(
+            torch.from_numpy(decoder_pos_embed).float().unsqueeze(0)
+        )
+
+        # initialize patch_embed like nn.Linear (instead of nn.Conv2d)
+        w = self.patch_embed.proj.weight.data
+        torch.nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
+
+        # timm's trunc_normal_(std=.02) is effectively normal_(std=0.02) as cutoff is too big (2.)
+        torch.nn.init.normal_(self.cls_token, std=0.02)
+        torch.nn.init.normal_(self.mask_token, std=0.02)
+
+        # initialize nn.Linear and nn.LayerNorm
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            # we use xavier_uniform following official JAX ViT:
+            torch.nn.init.xavier_uniform_(m.weight)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def prepare_tokens(
+        self,
+        x: torch.Tensor,
+        mask_ratio: float = 0.0,
+        return_helpers: bool = False,
+        **kwargs,
+    ):
+        # embed patches
+        x = self.patch_embed(x)
+        x = x.flatten(2).transpose(1, 2)
+
+        # add pos embed w/o cls token
+        x = x + self.pos_embed[:, 1:, :]
+
+        # Perform per-sample random masking by per-sample shuffling.
+        # Per-sample shuffling is done by argsort random noise.
+        N, N_PATCHES, D = x.shape
+        len_keep = int(N_PATCHES * (1 - mask_ratio))
+
+        # generate random noise for the selection (in [0, 1])
+        noise = torch.rand(N, N_PATCHES, device=x.device)
+        # sort noise for each sample (ascend: small is keep, large is remove)
+        ids_shuffle = torch.argsort(noise, dim=1)
+        # IDs to restore the mask (used for the decoder)
+        ids_restore = torch.argsort(ids_shuffle, dim=1)
+
+        # keep the first subset (random ids to keep from every sample)
+        ids_keep = ids_shuffle[:, :len_keep]
+        # reshape for alignment with "x"
+        ids_keep = ids_keep.unsqueeze(-1).repeat(1, 1, D)
+        # masked input
+        x_masked = torch.gather(x, dim=1, index=ids_keep)
+
+        # generate the binary mask: 0 is keep, 1 is remove
+        mask = torch.ones([N, N_PATCHES], device=x.device)
+        mask[:, :len_keep] = 0
+        # unshuffle to get the binary mask
+        # binary mask used to create "x_masked"
+        mask = torch.gather(mask, dim=1, index=ids_restore)
+
+        # append cls token
+        cls_token = self.cls_token + self.pos_embed[:, :1, :]
+        cls_tokens = cls_token.expand(x_masked.shape[0], -1, -1)
+        x_masked = torch.cat((cls_tokens, x_masked), dim=1)
+
+        if return_helpers:
+            return x_masked, mask, ids_restore
+        else:
+            return x_masked
+
+    def forward_decoder(self, x: torch.Tensor, ids_restore: torch.Tensor):
+        # embed tokens
+        x = self.decoder_embed(x)
+
+        # append mask tokens to sequence
+        mask_tokens = self.mask_token.repeat(
+            x.shape[0], ids_restore.shape[1] + 1 - x.shape[1], 1
+        )
+        x_ = torch.cat([x[:, 1:, :], mask_tokens], dim=1)  # no cls token
+        x_ = torch.gather(
+            x_, dim=1, index=ids_restore.unsqueeze(-1).repeat(1, 1, x.shape[2])
+        )  # unshuffle
+        x = torch.cat([x[:, :1, :], x_], dim=1)  # append cls token
+
+        # add pos embed
+        x = x + self.decoder_pos_embed
+
+        # apply Transformer blocks
+        for blk in self.decoder_blocks:
+            x = blk(x)
+        x = self.decoder_norm(x)
+
+        # predictor projection
+        x = self.decoder_pred(x)
+
+        # remove cls token
+        x = x[:, 1:, :]
+
+        return x
+
+    def forward(
+        self,
+        imgs: torch.Tensor,
+        mask_ratio: float = 0.75,
+        return_all_tokens=None,
+        **kwargs,
+    ):
+        x, mask, ids_restore = self.prepare_tokens(
+            x=imgs,
+            mask_ratio=mask_ratio,
+            return_helpers=True,
+        )
+        for blk in self.blocks:
+            x = blk(x)
+        latent = self.norm(x)
+        # [N, N_PATCHES, p*p*3]
+        pred = self.forward_decoder(x=latent, ids_restore=ids_restore)
+        return_all_tokens = (
+            self.return_all_tokens if return_all_tokens is None else return_all_tokens
+        )
+        if return_all_tokens:
+            return imgs, pred, mask, latent
+        return imgs, pred, mask, latent[:, 0]
+
+
+def masked_vit_tiny(patch_size: int = 16, **kwargs):
+    model = MaskedAutoencoderViT(
+        patch_size=patch_size,
+        embed_dim=192,
+        depth=12,
+        num_heads=3,
+        mlp_ratio=4,
+        qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),
+        decoder_embed_dim=512,
+        decoder_depth=8,
+        decoder_num_heads=16,
+        **kwargs,
+    )
+    return model
+
+
+def masked_vit_small(patch_size: int = 16, **kwargs):
+    model = MaskedAutoencoderViT(
+        patch_size=patch_size,
+        embed_dim=384,
+        depth=12,
+        num_heads=6,
+        mlp_ratio=4,
+        qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),
+        decoder_embed_dim=512,
+        decoder_depth=8,
+        decoder_num_heads=16,
+        **kwargs,
+    )
+    return model
+
+
+def masked_vit_base(patch_size: int = 16, **kwargs):
+    model = MaskedAutoencoderViT(
+        patch_size=patch_size,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        mlp_ratio=4,
+        qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),
+        decoder_embed_dim=512,
+        decoder_depth=8,
+        decoder_num_heads=16,
+        **kwargs,
+    )
+    return model
+
+
+def masked_vit_large(patch_size: int = 16, **kwargs):
+    model = MaskedAutoencoderViT(
+        patch_size=patch_size,
+        embed_dim=1024,
+        depth=24,
+        num_heads=16,
+        mlp_ratio=4,
+        qkv_bias=True,
+        norm_layer=partial(nn.LayerNorm, eps=1e-6),
+        decoder_embed_dim=512,
+        decoder_depth=8,
+        decoder_num_heads=16,
+        **kwargs,
+    )
+    return model

skinmap_runtime/core/models/mae/utils.py

@@ -0,0 +1,58 @@
+from typing import Callable, Tuple
+
+import torch
+
+from ....src.models.utils import ModelType
+from .model import masked_vit_base, masked_vit_large, masked_vit_small, masked_vit_tiny
+
+
+def patch_images(images: torch.Tensor, patch_size: int):
+    """
+    Transforms images into patched images.
+
+    imgs: (N, C, H, W)
+    x: (N, #patches, patch_size**2 * C)
+    """
+    # make sure the image properties are correct
+    img_is_square = images.shape[2] == images.shape[3]
+    img_is_patchable = images.shape[2] % patch_size == 0
+    assert img_is_square and img_is_patchable
+
+    channels = images.shape[1]
+    h = w = images.shape[2] // patch_size
+    x = images.reshape(shape=(images.shape[0], channels, h, patch_size, w, patch_size))
+    x = torch.einsum("nchpwq->nhwpqc", x)
+    x = x.reshape(shape=(images.shape[0], h * w, patch_size**2 * channels))
+    return x
+
+
+def unpatch_images(x: torch.Tensor, patch_size: int):
+    """
+    Transforms patched images into images.
+
+    x: (N, #patches, patch_size**2 * C)
+    imgs: (N, C, H, W)
+    """
+    h = w = int(x.shape[1] ** 0.5)
+    assert h * w == x.shape[1]
+
+    x = x.reshape(shape=(x.shape[0], h, w, patch_size, patch_size, 3))
+    x = torch.einsum("nhwpqc->nchpwq", x)
+    imgs = x.reshape(shape=(x.shape[0], 3, h * patch_size, h * patch_size))
+    return imgs
+
+
+MASKED_VIT_DICT = {
+    "masked_vit_tiny": masked_vit_tiny,
+    "masked_vit_small": masked_vit_small,
+    "masked_vit_base": masked_vit_base,
+    "masked_vit_large": masked_vit_large,
+}
+
+
+def get_model_class(base_model_name: str) -> Tuple[Callable, ModelType]:
+    encoder_cls = MASKED_VIT_DICT.get(base_model_name, None)
+    model_type = ModelType.VIT
+    if encoder_cls is None:
+        raise ValueError(f"Invalid base model name: {base_model_name}")
+    return encoder_cls, model_type

skinmap_runtime/core/models/simclr/__init__.py

@@ -0,0 +1 @@
+

skinmap_runtime/core/models/simclr/model.py

@@ -0,0 +1,33 @@
+import torch.nn as nn
+import torch.nn.functional as F
+
+from ..encoders.utils import get_encoder_class
+from ..utils import ModelType
+
+
+class ResNetSimCLR(nn.Module):
+    def __init__(self, base_model: str, out_dim: int, **kwargs):
+        super(ResNetSimCLR, self).__init__()
+        encoder_cls, model_type = get_encoder_class(base_model)
+        if model_type is ModelType.VIT:
+            self.backbone = encoder_cls(**kwargs)
+            n_feat = self.backbone.embed_dim
+        elif model_type is ModelType.CNN:
+            encoder = encoder_cls(**kwargs)
+            n_feat = encoder.fc.in_features
+            self.backbone = nn.Sequential(*list(encoder.children())[:-1])
+        else:
+            raise ValueError(f"Unknown model type: {model_type}")
+        # projection MLP
+        self.dense1 = nn.Linear(n_feat, n_feat)
+        self.dense2 = nn.Linear(n_feat, out_dim)
+
+    def forward(self, z):
+        # embed
+        e = self.backbone(z)
+        e = e.squeeze()
+        # project
+        z = self.dense1(e)
+        z = F.relu(z)
+        z = self.dense2(z)
+        return e, z