modeling_skinmap.py

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