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