skinmap_runtime/_compat.py

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