diff --git a/README.md b/README.md index 83d44a0..7d6a46a 100644 --- a/README.md +++ b/README.md @@ -119,7 +119,7 @@ The package writes explicit artifact directories: ### Supported Models -`slide2vec` currently ships preset configs for 18 tile-level models and 3 slide-level models. +`slide2vec` currently ships preset configs for 20 tile-level models and 3 slide-level models. For the full catalog and preset names, see [`docs/models.md`](docs/models.md). ## CLI diff --git a/docs/models.rst b/docs/models.rst index 35ed977..248fe6f 100644 --- a/docs/models.rst +++ b/docs/models.rst @@ -72,7 +72,12 @@ Tile-level encoders - `mSTAR `_ - 1024 - ``0.5`` - - Xu et al. (2024); gated weights, requires ``HF_TOKEN`` + - Xu et al. (2024) + * - ``gpfm`` + - `GPFM `_ + - 1024 + - ``0.5`` + - Ma et al. (2024) * - ``uni`` - `UNI `_ - 1024 @@ -118,6 +123,11 @@ Tile-level encoders - 3072 - ``0.25``, ``0.5``, ``1.0``, ``2.0`` - Karasikov et al. (2025) + * - ``genbio-pathfm`` + - `GenBio-PathFM `_ + - 4608 + - ``0.5`` + - GenBio AI (2024) Dense tile grids ~~~~~~~~~~~~~~~~ @@ -127,10 +137,10 @@ Dense tile feature extraction is available on tile encoders that implement instead of the pooled ``(B, D)`` tensor returned by ``encode_tiles``. The following built-in tile presets are covered by the dense encoder interface: -``conch``, ``conchv15``, ``gigapath``, ``h0-mini``, ``h-optimus-0``, -``h-optimus-1``, ``hibou-b``, ``hibou-l``, ``lunit``, ``midnight``, ``mstar``, -``musk``, ``phikon``, ``phikonv2``, ``prost40m``, ``uni``, ``uni2``, -``virchow``, and ``virchow2``. +``conch``, ``conchv15``, ``genbio-pathfm``, ``gigapath``, ``gpfm``, ``h0-mini``, +``h-optimus-0``, ``h-optimus-1``, ``hibou-b``, ``hibou-l``, ``lunit``, +``midnight``, ``mstar``, ``musk``, ``phikon``, ``phikonv2``, ``prost40m``, +``uni``, ``uni2``, ``virchow``, and ``virchow2``. Notes: @@ -148,6 +158,9 @@ Notes: - H-Optimus dense extraction at non-native input sizes requires ``dynamic_img_size=True`` and ``allow_non_recommended_settings=True`` when constructing the encoder. +- ``genbio-pathfm`` is a single-channel ViT: its dense grid (via + ``forward_with_patches``) concatenates the three per-colour-channel patch grids + along the feature dim, so ``d`` = 4608 matches the pooled output dimension. Dense attention maps ~~~~~~~~~~~~~~~~~~~~~ @@ -158,15 +171,20 @@ prefix-token self-attention as a spatial grid ``(B, K, h, w)`` — see the contract and knobs. The following built-in tile presets are covered: ``conch``, ``conchv15``, -``gigapath``, ``h0-mini``, ``h-optimus-0``, ``h-optimus-1``, ``hibou-b``, -``hibou-l``, ``lunit``, ``midnight``, ``mstar``, ``phikon``, ``phikonv2``, -``prost40m``, ``uni``, ``uni2``, ``virchow``, and ``virchow2``. +``gigapath``, ``gpfm``, ``h0-mini``, ``h-optimus-0``, ``h-optimus-1``, +``hibou-b``, ``hibou-l``, ``lunit``, ``midnight``, ``mstar``, ``phikon``, +``phikonv2``, ``prost40m``, ``uni``, ``uni2``, ``virchow``, and ``virchow2``. Notes: - ``musk`` is **not** covered: its BEiT3 backbone uses a non-timm attention module, so attention extraction raises ``NotImplementedError`` (dense patch-token extraction is still available). +- ``genbio-pathfm`` is **not** covered: it computes attention with a fused + ``scaled_dot_product_attention`` (no materialized weights) and encodes the three + colour channels as independent single-channel images, so there is no single + coherent CLS-over-patches attention to extract; ``encode_tiles_attention`` raises + ``NotImplementedError`` (dense patch-token extraction is still available). - ``hibou-b`` / ``hibou-l`` carry register tokens; pass ``include_registers=True`` to add their query rows as extra channels. - ``conch`` / ``conchv15`` recover attention from their inner timm ViT trunk, the diff --git a/slide2vec/encoders/models/__init__.py b/slide2vec/encoders/models/__init__.py index 964eca1..a012419 100644 --- a/slide2vec/encoders/models/__init__.py +++ b/slide2vec/encoders/models/__init__.py @@ -5,7 +5,9 @@ from . import ( conch, + genbio, gigapath, + gpfm, hibou, hoptimus, lunit, @@ -23,7 +25,9 @@ __all__ = [ "conch", + "genbio", "gigapath", + "gpfm", "hibou", "hoptimus", "lunit", diff --git a/slide2vec/encoders/models/genbio.py b/slide2vec/encoders/models/genbio.py new file mode 100644 index 0000000..f3aa452 --- /dev/null +++ b/slide2vec/encoders/models/genbio.py @@ -0,0 +1,155 @@ +"""GenBio-PathFM tile encoder. + +GenBio-PathFM (GenBio AI, 2024; ``genbio-ai/genbio-pathfm``) is a 1.1B-param ViT +histopathology tile encoder (JEDI = JEPA + DINO training on public data). It is +loaded via HF ``AutoModel(trust_remote_code=True)`` (auto_map -> +``GenBioPathFMModel``), so it is a custom :class:`TileEncoder`, not a +``TimmTileEncoder``. Weights are openly downloadable (safetensors); released +under the custom **GenBio AI Community License** (not OSI-approved — read it for +acceptable-use / commercial restrictions; slide2vec only wraps user-downloaded +weights and does not redistribute them). + +Output dim (4608, verified on real weights): the backbone is a *single-channel* +ViT (``in_chans=1``, ``embed_dim=1536``). The model's canonical ``forward`` takes +an RGB ``[B, 3, H, W]`` tensor, treats each colour channel as a separate +single-channel image, encodes all three, and concatenates the three per-channel +CLS tokens -> ``[B, embed_dim * 3] = [B, 4608]``. This per-channel-CLS +concatenation is the model's intrinsic design (matching the HF card's advertised +feature dimension of 4608), not an ad-hoc CLS+patch pooling; ``encode_tiles`` +therefore returns the model's default ``forward`` output directly. (Confirmed by +running a ``(1, 3, 224, 224)`` dummy through the real weights -> shape +``(1, 4608)``; see ``tests/test_gpfm_genbio_heavy.py``.) + +Normalization is **non-ImageNet** (``config.json`` ``image_mean`` / ``image_std``) +and must be set explicitly. + +Dense (spatial-grid) extraction is supported via the model's +``forward_with_patches`` (a DINOv2-style ``x_norm_patchtokens`` path): it returns +the fused per-channel patch tokens ``(B, T, 4608)`` — the three single-channel +patch-token grids concatenated along the feature dim, with the prefix tokens +(CLS + storage tokens) already stripped — which fold straight into a +``(B, 4608, h, w)`` grid. The patch size is 16 (a 224 tile -> a 14x14 = 196 grid). + +Attention-map extraction is **not** supported. The backbone computes attention +with a fused ``F.scaled_dot_product_attention`` (no materialized weights, no +``output_attentions``), and — more fundamentally — it encodes the three colour +channels as three independent single-channel images, so there is no single +coherent CLS-over-patches attention to extract: any "attention grid" would be +three separate grayscale-channel attentions. Recovering it would need a bespoke +per-channel recompute path that diverges from the shared timm/HF attention +helpers, so GenBio deliberately opts out of ``encode_tiles_attention``. +""" + +from typing import Callable + +import torch +from torch import Tensor +from torchvision.transforms import v2 + +from slide2vec.encoders.base import ( + TileEncoder, + preferred_default_device, + reshape_tokens_to_grid, + resolve_requested_output_variant, +) +from slide2vec.encoders.registry import register_encoder + +_HF_REPO_ID = "genbio-ai/genbio-pathfm" +# Non-ImageNet normalization from the model's config.json (image_mean / image_std). +_GENBIO_MEAN = (0.697, 0.575, 0.728) +_GENBIO_STD = (0.188, 0.240, 0.187) + + +@register_encoder( + "genbio-pathfm", + # encode_dim 4608 = embed_dim (1536) x 3 colour channels; patch_size 16 (a 224 + # tile -> a 14x14 = 196 patch-token grid) — see module docstring. + output_variants={"default": {"encode_dim": 4608}}, + default_output_variant="default", + input_size=224, + patch_size=16, + supported_spacing_um=0.5, # 20x; card states no magnification, house default + precision="fp32", # upstream runs plain fp32, no autocast + source="genbio-ai/genbio-pathfm", +) +class GenBioPathFM(TileEncoder): + def __init__(self, *, output_variant: str | None = None): + from transformers import AutoModel + + self._model = AutoModel.from_pretrained(_HF_REPO_ID, trust_remote_code=True).eval() + self._device = preferred_default_device() + self._output_variant = resolve_requested_output_variant(output_variant) + + def get_transform(self) -> Callable: + # Mirrors the model card: Resize to an exact 224x224 (tuple form, as the + # card's ``Resize((224, 224))``) + custom Normalize. No CenterCrop: the card + # has none, and after a square resize it would be a no-op anyway. + return v2.Compose([ + v2.ToImage(), + v2.Resize((224, 224), interpolation=v2.InterpolationMode.BICUBIC, antialias=True), + v2.ToDtype(torch.float32, scale=True), + v2.Normalize(mean=_GENBIO_MEAN, std=_GENBIO_STD), + ]) + + def get_dense_transform(self) -> Callable: + # Normalization only — no Resize/CenterCrop (see TileEncoder.get_dense_transform), + # so the dense grid stays registered to the full source tile. + return v2.Compose([ + v2.ToImage(), + v2.ToDtype(torch.float32, scale=True), + v2.Normalize(mean=_GENBIO_MEAN, std=_GENBIO_STD), + ]) + + def encode_tiles(self, batch: Tensor) -> Tensor: + # Canonical forward: per-channel CLS tokens concatenated -> (B, 4608). + return self._model(batch) + + @property + def patch_size(self) -> tuple[int, int]: + return (16, 16) + + def encode_tiles_dense(self, batch: Tensor) -> Tensor: + """Encode tiles into a dense spatial grid. (B, C, H, W) -> (B, d, h, w). + + Uses the model's ``forward_with_patches`` (DINOv2-style), which returns the + fused per-channel patch tokens ``(B, T, 4608)`` — the three single-channel + patch-token grids concatenated along the feature dim, with the prefix + tokens (CLS + storage tokens) already stripped — then folds that token + sequence back into its spatial grid (``num_prefix_tokens=0``). ``H, W`` must + be divisible by the patch size; non-224 inputs rely on the backbone's + DINOv2 positional-embedding interpolation. + """ + if batch.ndim != 4: + raise ValueError( + "encode_tiles_dense expects a (B, C, H, W) batch, got shape " + f"{tuple(batch.shape)}." + ) + _, _, height, width = batch.shape + patch_h, patch_w = self.patch_size + if height % patch_h != 0 or width % patch_w != 0: + raise ValueError( + f"Dense extraction for '{type(self).__name__}' requires input " + f"divisible by the patch size: got {height}x{width}, patch " + f"{patch_h}x{patch_w}. Pad the tile up to a patch multiple first." + ) + _, patch_tokens = self._model.forward_with_patches(batch) + return reshape_tokens_to_grid( + patch_tokens, + grid_h=height // patch_h, + grid_w=width // patch_w, + num_prefix_tokens=0, + encoder_name=type(self).__name__, + ) + + @property + def encode_dim(self) -> int: + return 4608 + + @property + def device(self) -> torch.device: + return self._device + + def to(self, device: torch.device | str) -> "GenBioPathFM": + self._device = torch.device(device) + self._model = self._model.to(self._device) + return self diff --git a/slide2vec/encoders/models/gpfm.py b/slide2vec/encoders/models/gpfm.py new file mode 100644 index 0000000..567bdab --- /dev/null +++ b/slide2vec/encoders/models/gpfm.py @@ -0,0 +1,89 @@ +"""GPFM tile encoder (Generalizable Pathology Foundation Model). + +GPFM (Ma et al., 2024; ``birkhoffkiki/GPFM``) is a ``ViT-L/14`` DINOv2 tile +encoder (embedding dim 1024, patch_size 14). Weights are openly available under +the MIT license at ``majiabo/GPFM`` but ship as a standalone ``GPFM.pth`` +checkpoint, so — unlike ``hf-hub:`` timm presets — we build the timm arch +unpretrained and ``load_state_dict`` the downloaded checkpoint (the MOOZY +loading pattern). + +The published ``GPFM.pth`` is a bare DINOv2 ``state_dict`` (``cls_token``, +``pos_embed``, ``blocks.*``, …) that loads into +``vit_large_patch14_dinov2.lvd142m`` with ``strict=True`` and zero missing / +unexpected keys (verified on the real weights; see +``tests/test_gpfm_genbio_heavy.py``). The loader below still defensively unwraps +common checkpoint wrappers (``{"model": ...}`` / ``{"teacher": ...}`` / +``{"student": ...}`` / ``{"state_dict": ...}``) and strips ``module.`` / +``backbone.`` prefixes so a re-exported checkpoint keeps loading strictly. +""" + +from typing import Mapping + +import torch +from huggingface_hub import hf_hub_download + +from slide2vec.encoders.base import TimmTileEncoder +from slide2vec.encoders.registry import register_encoder + +_HF_REPO_ID = "majiabo/GPFM" +_HF_CHECKPOINT = "GPFM.pth" +_CHECKPOINT_WRAPPER_KEYS = ("model", "teacher", "student", "state_dict", "teacher_backbone") +_STATE_DICT_PREFIXES = ("module.", "backbone.") + + +def _unwrap_gpfm_state_dict(payload: Mapping[str, object]) -> dict[str, torch.Tensor]: + """Reduce a GPFM checkpoint payload to a bare ``state_dict``. + + Unwraps a single common wrapper key if present, then strips ``module.`` / + ``backbone.`` prefixes off every key. A no-op on the published bare + ``GPFM.pth`` (no wrapper, no prefixes), so it preserves strict loading there + while tolerating a re-wrapped/prefixed re-export. + """ + state: object = payload + for wrapper in _CHECKPOINT_WRAPPER_KEYS: + if isinstance(state, Mapping) and wrapper in state and isinstance(state[wrapper], Mapping): + state = state[wrapper] + break + if not isinstance(state, Mapping): + raise ValueError( + f"Unexpected GPFM checkpoint payload: expected a state_dict mapping, got {type(state)}" + ) + cleaned: dict[str, torch.Tensor] = {} + for key, value in state.items(): + name = key + for prefix in _STATE_DICT_PREFIXES: + if name.startswith(prefix): + name = name[len(prefix):] + cleaned[name] = value + return cleaned + + +@register_encoder( + "gpfm", + output_variants={"default": {"encode_dim": 1024}}, + default_output_variant="default", + input_size=224, + patch_size=14, + supported_spacing_um=0.5, # 512px@0.25um (128um FOV) resized to 224 => ~0.5um/px effective (20x), as for UNI + precision="fp32", # upstream runs plain fp32, no autocast + source="majiabo/GPFM", +) +class GPFM(TimmTileEncoder): + def __init__(self, *, output_variant: str | None = None): + super().__init__( + "vit_large_patch14_dinov2.lvd142m", + output_variant=output_variant, + pretrained=False, + img_size=224, + init_values=1e-5, + dynamic_img_size=True, + ) + checkpoint_path = hf_hub_download(repo_id=_HF_REPO_ID, filename=_HF_CHECKPOINT) + payload = torch.load(checkpoint_path, map_location="cpu", weights_only=False) + if not isinstance(payload, Mapping): + raise ValueError( + f"Invalid GPFM checkpoint payload: expected a dict, got {type(payload)}" + ) + state_dict = _unwrap_gpfm_state_dict(payload) + self._model.load_state_dict(state_dict, strict=True) + self._model.eval() diff --git a/tests/test_encoder_registry.py b/tests/test_encoder_registry.py index 72bbaa9..8ac1eaf 100644 --- a/tests/test_encoder_registry.py +++ b/tests/test_encoder_registry.py @@ -12,7 +12,9 @@ "virchow2", "conch", "conchv15", + "genbio-pathfm", "gigapath", + "gpfm", "h-optimus-0", "h-optimus-1", "h0-mini", @@ -131,6 +133,35 @@ def test_mstar_slide_alias_resolves_to_mstar(): assert canonicalize_model_name("mstar-slide") == "mstar" +def test_gpfm_metadata_contract(): + info = encoder_registry.info("gpfm") + assert info["level"] == "tile" + assert info["input_size"] == 224 + assert info["patch_size"] == 14 + assert info["supported_spacing_um"] == pytest.approx(0.5) + assert info["precision"] == "fp32" + assert info["source"] == "majiabo/GPFM" + assert info["output_variants"]["default"]["encode_dim"] == 1024 + assert info["default_output_variant"] == "default" + + +def test_genbio_pathfm_metadata_contract(): + info = encoder_registry.info("genbio-pathfm") + assert info["level"] == "tile" + assert info["input_size"] == 224 + # GenBio exposes a dense patch grid via forward_with_patches (single-channel + # ViT, patch size 16: a 224 tile -> a 14x14 grid). + assert info["patch_size"] == 16 + assert info["supported_spacing_um"] == pytest.approx(0.5) + assert info["precision"] == "fp32" + assert info["source"] == "genbio-ai/genbio-pathfm" + # GenBio is a single-channel ViT (in_chans=1): forward() splits RGB into 3 + # single-channel images and concatenates their per-channel CLS tokens, so the + # canonical pooled output is embed_dim*3 = 1536*3 = 4608 (verified on real weights). + assert info["output_variants"]["default"]["encode_dim"] == 4608 + assert info["default_output_variant"] == "default" + + def test_resolve_preprocessing_requirements_for_tile_encoder(): from slide2vec.encoders.registry import resolve_preprocessing_requirements diff --git a/tests/test_gpfm_genbio_heavy.py b/tests/test_gpfm_genbio_heavy.py new file mode 100644 index 0000000..63e44c6 --- /dev/null +++ b/tests/test_gpfm_genbio_heavy.py @@ -0,0 +1,64 @@ +"""Real-weight forward-shape contracts for GPFM and GenBio-PathFM. + +These load multi-GB foundation-model weights on CPU (minutes each), so they are +marked ``heavy`` and excluded from the PR suite (``-m 'not heavy'``); they run on +the scheduled/manual heavy workflow. Each skips cleanly when weights / optional +deps / network are unavailable so a developer run never hard-fails. + +They pin the two registration choices that can only be confirmed against the real +checkpoints: + +* GPFM: ``GPFM.pth`` loads into the timm DINOv2 ViT-L/14 (``strict=True``) and the + pooled forward returns the registered 1024-d embedding; its dense path folds a + 224 tile into a ``(1, 1024, 16, 16)`` patch-token grid (patch size 14). +* GenBio-PathFM: the single-channel backbone's canonical forward returns the + registered 4608-d (= embed_dim*3) per-channel-CLS concatenation; its dense path + (``forward_with_patches``) folds a 224 tile into a ``(1, 4608, 14, 14)`` grid + (patch size 16). +""" + +from __future__ import annotations + +import pytest +import torch + +from slide2vec.encoders import encoder_registry + + +@pytest.mark.heavy +def test_gpfm_loads_strict_and_forward_is_1024_d(): + try: + encoder = encoder_registry.require("gpfm")(output_variant="default") + except Exception as exc: # network / weights / optional dep unavailable + pytest.skip(f"gpfm weights unavailable: {type(exc).__name__}: {exc}") + encoder = encoder.to("cpu") + assert encoder.encode_dim == 1024 + assert encoder.patch_size == (14, 14) + batch = torch.randn(1, 3, 224, 224) + with torch.inference_mode(): + out = encoder.encode_tiles(batch) + dense = encoder.encode_tiles_dense(batch) + assert out.shape == (1, 1024) + # patch 14: 224 / 14 = 16 -> a 16x16 grid of 1024-d patch tokens. + assert dense.shape == (1, 1024, 16, 16) + + +@pytest.mark.heavy +def test_genbio_pathfm_forward_is_4608_d(): + try: + encoder = encoder_registry.require("genbio-pathfm")(output_variant="default") + except Exception as exc: # network / weights / optional dep unavailable + pytest.skip(f"genbio-pathfm weights unavailable: {type(exc).__name__}: {exc}") + encoder = encoder.to("cpu") + assert encoder.encode_dim == 4608 + assert encoder.patch_size == (16, 16) + batch = torch.randn(1, 3, 224, 224) + with torch.inference_mode(): + out = encoder.encode_tiles(batch) + dense = encoder.encode_tiles_dense(batch) + # Single-channel ViT (in_chans=1): forward concatenates the 3 per-channel CLS + # tokens (embed_dim 1536) -> 4608. + assert out.shape == (1, 4608) + # Dense via forward_with_patches: the 3 per-channel patch grids concatenated + # along the feature dim; patch 16: 224 / 16 = 14 -> a 14x14 grid of 4608-d tokens. + assert dense.shape == (1, 4608, 14, 14) diff --git a/tests/test_patch_size_metadata.py b/tests/test_patch_size_metadata.py index e60c3e7..2b64c82 100644 --- a/tests/test_patch_size_metadata.py +++ b/tests/test_patch_size_metadata.py @@ -27,7 +27,9 @@ "uni2": (14, 14), "virchow": (14, 14), "virchow2": (14, 14), + "genbio-pathfm": (16, 16), "gigapath": (14, 14), + "gpfm": (14, 14), "h-optimus-0": (14, 14), "h-optimus-1": (14, 14), "h0-mini": (14, 14), @@ -44,8 +46,15 @@ "hibou-l": (14, 14), } +# Non-dense tile encoders: real tile encoders with no recoverable patch grid +# (e.g. a custom AutoModel backbone whose forward exposes no spatial patch-token +# sequence), so they declare no ``patch_size`` and opt out of the dense interface. +# Currently empty — every built-in tile encoder is dense-capable — but the bucket +# is kept so a future non-dense tile encoder has a tracked home. +NON_DENSE_TILE_ENCODERS: list[str] = [] + # Non-dense encoders: no recoverable patch grid, so no declared patch_size. -NON_DENSE_ENCODERS = ["gigapath-slide", "titan", "prism"] +NON_DENSE_ENCODERS = ["gigapath-slide", "titan", "prism"] + NON_DENSE_TILE_ENCODERS def test_normalize_patch_size_int_and_tuple(): @@ -95,11 +104,22 @@ def test_resolve_patch_size_raises_for_non_dense_encoder(name): def test_every_tile_encoder_declares_patch_size(): - """All registered tile-level encoders are dense-capable and declare patch_size.""" + """Dense-capable tile encoders declare patch_size; non-dense ones opt out. + + Most tile encoders are dense-capable ViTs and must declare a ``patch_size``. + A few (``NON_DENSE_TILE_ENCODERS``, e.g. custom AutoModel backbones) expose no + recoverable patch grid and deliberately declare none. Every tile encoder must + fall into exactly one of those two tracked buckets. + """ for info in encoder_registry.list_with_metadata(): if info["level"] != "tile": continue name = info["name"] + if name in NON_DENSE_TILE_ENCODERS: + assert info.get("patch_size") is None, ( + f"non-dense tile encoder '{name}' must not declare a patch_size" + ) + continue assert info.get("patch_size") is not None, ( f"tile encoder '{name}' is dense-capable but declares no patch_size" ) diff --git a/tests/test_regression_core.py b/tests/test_regression_core.py index 62b5723..63b21ec 100644 --- a/tests/test_regression_core.py +++ b/tests/test_regression_core.py @@ -189,7 +189,9 @@ def test_list_models_can_filter_by_level(): assert list_models("tile") == [ "conch", "conchv15", + "genbio-pathfm", "gigapath", + "gpfm", "h-optimus-0", "h-optimus-1", "h0-mini",