slide2vec 5.8.0__tar.gz → 5.8.2__tar.gz

{slide2vec-5.8.0 → slide2vec-5.8.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: slide2vec
-Version: 5.8.0
+Version: 5.8.2
 Summary: Embedding of whole slide images with Foundation Models
 Author-email: Clément Grisi <clement.grisi@radboudumc.nl>
 License-Expression: Apache-2.0
@@ -15,7 +15,7 @@ Classifier: Programming Language :: Python :: 3.13
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: hs2p[asap,cucim,openslide,sam2,vips]>=4.4.1
+Requires-Dist: hs2p[asap,cucim,openslide,sam2,vips]>=4.4.2
 Requires-Dist: omegaconf
 Requires-Dist: matplotlib
 Requires-Dist: numpy<2
@@ -75,7 +75,7 @@ Requires-Dist: numpy<2; extra == "fm"
 Requires-Dist: pandas; extra == "fm"
 Requires-Dist: pillow; extra == "fm"
 Requires-Dist: rich; extra == "fm"
-Requires-Dist: hs2p[asap,cucim,openslide,sam2,vips]>=4.4.1; extra == "fm"
+Requires-Dist: hs2p[asap,cucim,openslide,sam2,vips]>=4.4.2; extra == "fm"
 Requires-Dist: wandb; extra == "fm"
 Requires-Dist: torch<2.8,>=2.3; extra == "fm"
 Requires-Dist: torchvision>=0.18.0; extra == "fm"

{slide2vec-5.8.0 → slide2vec-5.8.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "slide2vec"
-version = "5.8.0"
+version = "5.8.2"
 description = "Embedding of whole slide images with Foundation Models"
 readme = "README.md"
 requires-python = ">=3.10"
@@ -21,7 +21,7 @@ classifiers = [
     "Programming Language :: Python :: 3.13",
 ]
 dependencies = [
-    "hs2p[asap,cucim,openslide,sam2,vips]>=4.4.1",
+    "hs2p[asap,cucim,openslide,sam2,vips]>=4.4.2",
     "omegaconf",
     "matplotlib",
     "numpy<2",
@@ -100,7 +100,7 @@ fm = [
     "pandas",
     "pillow",
     "rich",
-    "hs2p[asap,cucim,openslide,sam2,vips]>=4.4.1",
+    "hs2p[asap,cucim,openslide,sam2,vips]>=4.4.2",
     "wandb",
     "torch>=2.3,<2.8",
     "torchvision>=0.18.0",
@@ -180,7 +180,7 @@ no_implicit_reexport = true
 max-line-length = 160
 
 [tool.bumpver]
-current_version = "5.8.0"
+current_version = "5.8.2"
 version_pattern = "MAJOR.MINOR.PATCH"
 commit = false       # We do version bumping in CI, not as a commit
 tag = false          # Git tag already exists — we don't auto-tag

{slide2vec-5.8.0 → slide2vec-5.8.2}/slide2vec/__init__.py

@@ -27,7 +27,7 @@ from slide2vec.encoders import (
 )
 
 
-__version__ = "5.8.0"
+__version__ = "5.8.2"
 
 __all__ = [
     "Model",

{slide2vec-5.8.0 → slide2vec-5.8.2}/slide2vec/encoders/models/gigapath.py

@@ -123,6 +123,9 @@ class GigaPathSlideEncoder(SlideEncoder):
             tile_features = tile_features.unsqueeze(0)
         if coordinates.ndim == 2:
             coordinates = coordinates.unsqueeze(0)
+        # prov-gigapath's reference usage feeds fp32 tile embeddings under autocast;
+        # fp16 input risks dtype mismatches in ops autocast does not cover.
+        tile_features = tile_features.float()
         # gigapath_slide_enc12l768d.forward always returns a list of per-layer
         # embeddings (a single element when all_layer_embed is False); the final
         # slide embedding is the last layer, matching prov-gigapath's own usage.

slide2vec-5.8.2/slide2vec/encoders/models/titan.py

@@ -0,0 +1,180 @@
+"""TITAN slide encoder implementation."""
+
+import math
+import sys
+
+import numpy as np
+import torch
+from transformers import AutoModel
+
+from slide2vec.encoders.base import SlideEncoder, preferred_default_device, resolve_requested_output_variant
+from slide2vec.encoders.registry import register_encoder
+
+# Pinned so the remote-code patches below stay valid (and so runs are reproducible —
+# an unpinned from_pretrained re-downloads whatever is at the repo HEAD).
+_TITAN_REVISION = "dac6773d9961cfc75503440676ff157a2c6e8d2e"
+
+
+def _alibi_slopes(n: int) -> list[float]:
+    # verbatim ALiBi slope schedule from TITAN's get_alibi at the pinned revision
+    if math.log2(n).is_integer():
+        p = 2 ** (-(2 ** -(math.log2(n) - 3)))
+        return [p * (p ** i) for i in range(n)]
+    nearest = 2 ** math.floor(math.log2(n))
+    base = _alibi_slopes(nearest)
+    extra = _alibi_slopes(2 * nearest)[0::2][: n - nearest]
+    return base + extra
+
+
+def _lean_alibi_bias(module, w, h, bg_mask, device, dtype):
+    """Bitwise-identical replacement for TITAN's get_alibi + the caller's cast/move.
+
+    The reference builds the [1, heads, N, N] bias through N x N float64 numpy on
+    the CPU (>100 GB of host RAM for slides in the tens of thousands of tiles) and
+    hands SDPA an unaligned bias that forces the math kernel, which materializes a
+    second N^2 tensor on the GPU. This builds the same values on-device in fp16, in
+    row chunks, with rows padded to a multiple of 8 elements so the memory-efficient
+    SDPA kernel accepts the bias — peak memory drops from ~4x to ~1x the bias size.
+    """
+    ii, jj = torch.meshgrid(
+        torch.arange(w, device=device), torch.arange(h, device=device), indexing="ij"
+    )
+    if bg_mask is not None:
+        mask = bg_mask.to(device).squeeze(0)
+        ii, jj = ii[mask], jj[mask]
+    points = torch.stack([ii.reshape(-1), jj.reshape(-1)], dim=1).float()
+    n = points.shape[0]
+    length = n + 1  # +1 for the cls token; its bias row/col stays zero
+    padded = ((length + 7) // 8) * 8
+    slopes = torch.tensor(
+        _alibi_slopes(module.num_heads), device=device, dtype=torch.float32
+    ).view(module.num_heads, 1, 1)
+    bias = torch.zeros(1, module.num_heads, length, padded, device=device, dtype=dtype)
+    # chunk intermediate is (heads, step, n) fp32 — stays under ~1.5 GB even at 53k tiles
+    step = 512
+    for start in range(0, n, step):
+        diff = points[start : start + step].unsqueeze(1) - points.unsqueeze(0)
+        dist = diff.square().sum(-1).sqrt()
+        bias[0, :, 1 + start : 1 + start + dist.shape[0], 1:length] = (
+            dist.unsqueeze(0) * slopes * -1
+        ).to(dtype)
+    return bias[..., :length]
+
+
+def _patch_titan_remote_code(model) -> bool:
+    """Patch TITAN's remote code for fp16 input and bounded bias memory.
+
+    Two patches on the dynamically loaded module (it exists only after
+    from_pretrained): preprocess_features runs its grid index_add_ in fp32 (the op
+    rejects fp16 features) but returns the grid in the features' own dtype — an
+    exact roundtrip that keeps the whole forward on the fp16 path — and
+    forward_features' single-slide alibi branch swaps in _lean_alibi_bias.
+    Returns False without patching if the module does not look like the pinned
+    revision; the caller then falls back to fp32 features (correct, but with the
+    reference implementation's memory behavior).
+    """
+    try:
+        vision_encoder = model.vision_encoder
+        vit = sys.modules[type(vision_encoder).__module__]
+        if getattr(vit, "_slide2vec_titan_patched", False):
+            return True
+        for attr in ("pos_encode_type", "num_heads", "patch_embed", "_pos_embed", "norm_pre", "blocks", "norm"):
+            if not hasattr(vision_encoder, attr):
+                return False
+        if not callable(getattr(vit, "preprocess_features", None)):
+            return False
+    except Exception:
+        return False
+
+    orig_preprocess = vit.preprocess_features
+
+    def preprocess_features(features, coords, patch_size_lv0):
+        grid, coords_grid, bg_mask = orig_preprocess(features.float(), coords, patch_size_lv0)
+        return grid.to(features.dtype), coords_grid, bg_mask
+
+    orig_forward_features = type(vision_encoder).forward_features
+
+    def forward_features(self, x, coords=None, mask=None, bg_mask=None):
+        # single-slide alibi path only; anything else falls through to the original
+        if self.pos_encode_type != "alibi" or x.shape[0] != 1 or self.masked_im_modeling:
+            return orig_forward_features(self, x, coords=coords, mask=mask, bg_mask=bg_mask)
+        B, nc, w, h = x.shape
+        # bias dtype = input grid dtype, as in the original, which builds the bias
+        # before patch_embed; post-norm activations can be fp32 under autocast
+        in_dtype = x.dtype
+        x = x.flatten(2, 3).transpose(1, 2)
+        x = self.patch_embed(x)
+        x = self._pos_embed(x, coords, w, h)
+        x = self.norm_pre(x)
+        if bg_mask is not None:
+            keep = torch.cat(
+                (torch.ones((1, 1), dtype=torch.bool, device=x.device), bg_mask.view(1, -1)),
+                dim=1,
+            )
+            x = x[keep].unsqueeze(0)
+        attn_bias = _lean_alibi_bias(self, w, h, bg_mask, device=x.device, dtype=in_dtype)
+        x = self.blocks(x, attn_bias, bg_mask)
+        x = self.norm(x)
+        return x
+
+    vit.preprocess_features = preprocess_features
+    type(vision_encoder).forward_features = forward_features
+    vit._slide2vec_titan_patched = True
+    return True
+
+
+@register_encoder(
+    "titan",
+    level="slide",
+    tile_encoder="conchv15",
+    tile_encoder_output_variant="default",
+    output_variants={"default": {"encode_dim": 768}},
+    default_output_variant="default",
+    supported_spacing_um=0.5,
+    precision="fp16",
+    source="MahmoodLab/TITAN",
+)
+class TitanSlideEncoder(SlideEncoder):
+    def __init__(self, *, output_variant: str | None = None):
+        self._model = AutoModel.from_pretrained(
+            "MahmoodLab/TITAN", revision=_TITAN_REVISION, trust_remote_code=True
+        ).eval()
+        self._remote_code_patched = _patch_titan_remote_code(self._model)
+        self._device = preferred_default_device()
+        self._output_variant = resolve_requested_output_variant(output_variant)
+
+    @property
+    def encode_dim(self) -> int:
+        return 768
+
+    @property
+    def device(self) -> torch.device:
+        return self._device
+
+    def to(self, device: torch.device | str) -> "TitanSlideEncoder":
+        self._device = torch.device(device)
+        self._model = self._model.to(self._device)
+        return self
+
+    def encode_slide(
+        self,
+        tile_features: torch.Tensor,
+        coordinates: torch.Tensor | None = None,
+        *,
+        tile_size_lv0: int | None = None,
+    ) -> torch.Tensor:
+        if coordinates is None or tile_size_lv0 is None:
+            raise ValueError("TITAN slide encoding requires coordinates and tile_size_lv0")
+        if tile_features.ndim == 2:
+            tile_features = tile_features.unsqueeze(0)
+        if coordinates.ndim == 2:
+            coordinates = coordinates.unsqueeze(0)
+        if not self._remote_code_patched:
+            # fallback for unrecognized remote code: fp32 features satisfy its fp32
+            # grid index_add_, at the cost of the reference memory behavior
+            tile_features = tile_features.float()
+        return self._model.encode_slide_from_patch_features(
+            tile_features,
+            coordinates.long(),
+            np.int64(tile_size_lv0),
+        ).squeeze(0)

{slide2vec-5.8.0 → slide2vec-5.8.2}/slide2vec.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: slide2vec
-Version: 5.8.0
+Version: 5.8.2
 Summary: Embedding of whole slide images with Foundation Models
 Author-email: Clément Grisi <clement.grisi@radboudumc.nl>
 License-Expression: Apache-2.0
@@ -15,7 +15,7 @@ Classifier: Programming Language :: Python :: 3.13
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: hs2p[asap,cucim,openslide,sam2,vips]>=4.4.1
+Requires-Dist: hs2p[asap,cucim,openslide,sam2,vips]>=4.4.2
 Requires-Dist: omegaconf
 Requires-Dist: matplotlib
 Requires-Dist: numpy<2
@@ -75,7 +75,7 @@ Requires-Dist: numpy<2; extra == "fm"
 Requires-Dist: pandas; extra == "fm"
 Requires-Dist: pillow; extra == "fm"
 Requires-Dist: rich; extra == "fm"
-Requires-Dist: hs2p[asap,cucim,openslide,sam2,vips]>=4.4.1; extra == "fm"
+Requires-Dist: hs2p[asap,cucim,openslide,sam2,vips]>=4.4.2; extra == "fm"
 Requires-Dist: wandb; extra == "fm"
 Requires-Dist: torch<2.8,>=2.3; extra == "fm"
 Requires-Dist: torchvision>=0.18.0; extra == "fm"

{slide2vec-5.8.0 → slide2vec-5.8.2}/slide2vec.egg-info/SOURCES.txt

@@ -153,4 +153,5 @@ tests/test_sharding.py
 tests/test_slide_coordinate_preparation.py
 tests/test_soma_migration.py
 tests/test_tile_store.py
-tests/test_tiling_pipeline.py
+tests/test_tiling_pipeline.py
+tests/test_titan.py

{slide2vec-5.8.0 → slide2vec-5.8.2}/slide2vec.egg-info/requires.txt

@@ -1,4 +1,4 @@
-hs2p[asap,cucim,openslide,sam2,vips]>=4.4.1
+hs2p[asap,cucim,openslide,sam2,vips]>=4.4.2
 omegaconf
 matplotlib
 numpy<2
@@ -27,7 +27,7 @@ numpy<2
 pandas
 pillow
 rich
-hs2p[asap,cucim,openslide,sam2,vips]>=4.4.1
+hs2p[asap,cucim,openslide,sam2,vips]>=4.4.2
 wandb
 torch<2.8,>=2.3
 torchvision>=0.18.0

{slide2vec-5.8.0 → slide2vec-5.8.2}/tests/test_hs2p_package_cutover.py

@@ -23,7 +23,7 @@ def test_package_root_exports_api():
     assert hasattr(package, "DenseRegionArtifact")
 
 
-def test_dependency_declarations_require_hs2p_4_4_1_with_consistent_extras():
+def test_dependency_declarations_require_hs2p_4_4_2_with_consistent_extras():
     project = tomllib.loads(
         (Path(__file__).resolve().parents[1] / "pyproject.toml").read_text(
             encoding="utf-8"
@@ -39,8 +39,8 @@ def test_dependency_declarations_require_hs2p_4_4_1_with_consistent_extras():
     ]
 
     assert hs2p_dependencies == [
-        "hs2p[asap,cucim,openslide,sam2,vips]>=4.4.1",
-        "hs2p[asap,cucim,openslide,sam2,vips]>=4.4.1",
+        "hs2p[asap,cucim,openslide,sam2,vips]>=4.4.2",
+        "hs2p[asap,cucim,openslide,sam2,vips]>=4.4.2",
     ]
 
 

slide2vec-5.8.2/tests/test_titan.py

@@ -0,0 +1,134 @@
+"""Regression tests for the TITAN slide encoder input contract."""
+
+from types import SimpleNamespace
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+
+class _FakeTitan:
+    """Mirrors TITAN's preprocess_features: index_add_ of the input into an fp32
+    grid, which raises on fp16 features."""
+
+    def eval(self):
+        return self
+
+    def encode_slide_from_patch_features(self, patch_features, patch_coords, patch_size_lv0):
+        grid = torch.zeros(4, patch_features.size(-1))
+        indices = torch.zeros(patch_features.size(1), dtype=torch.long)
+        grid.index_add_(0, indices, patch_features.squeeze(0))
+        return torch.zeros(1, 768)
+
+
+def test_titan_encode_slide_accepts_fp16_features(monkeypatch):
+    # _FakeTitan has no vision_encoder, so the remote-code patch declines and the
+    # encoder must fall back to fp32 features to satisfy the fp32-grid index_add_
+    import slide2vec.encoders.models.titan as titan_mod
+
+    monkeypatch.setattr(
+        titan_mod,
+        "AutoModel",
+        SimpleNamespace(from_pretrained=lambda *args, **kwargs: _FakeTitan()),
+    )
+    encoder = titan_mod.TitanSlideEncoder()
+    assert encoder._remote_code_patched is False
+
+    features = torch.randn(5, 768, dtype=torch.float16)
+    coordinates = torch.zeros(5, 2, dtype=torch.int64)
+    embedding = encoder.encode_slide(features, coordinates, tile_size_lv0=512)
+
+    assert embedding.shape[-1] == 768
+
+
+def _reference_get_alibi(num_heads, w, h, bg_mask=None):
+    """Verbatim port of TITAN's get_alibi (revision dac6773) + the caller's fp16 cast."""
+    import math
+
+    import numpy as np
+
+    x, y = np.meshgrid(np.arange(w), np.arange(h), indexing="ij")
+    if bg_mask is not None:
+        x = x[bg_mask.cpu().squeeze(0)]
+        y = y[bg_mask.cpu().squeeze(0)]
+    points = np.stack([x.ravel(), y.ravel()], axis=1)
+    diffs = points[:, None, :] - points[None, :, :]
+    dists = np.sqrt(np.sum(diffs**2, axis=-1))
+
+    def get_slopes(n):
+        if math.log2(n).is_integer():
+            p = 2 ** (-(2 ** -(math.log2(n) - 3)))
+            return [p * (p**i) for i in range(n)]
+        nearest = 2 ** math.floor(math.log2(n))
+        return get_slopes(nearest) + get_slopes(2 * nearest)[0::2][: n - nearest]
+
+    slopes = torch.tensor(get_slopes(num_heads), dtype=torch.float32).view(num_heads, 1, 1)
+    n_patches = dists.shape[-1]
+    dists_tensor = torch.tensor(dists, dtype=torch.float32).view(1, n_patches, n_patches)
+    bias_matrix = dists_tensor * slopes * -1
+    embed_len = n_patches + 1
+    all_bias = torch.zeros(1, num_heads, embed_len, embed_len)
+    all_bias[:, :, 1:, 1:] = bias_matrix
+    return all_bias.half()
+
+
+@pytest.mark.parametrize("w,h,use_mask", [(7, 5, False), (9, 11, True), (40, 45, True)])
+def test_lean_alibi_bias_matches_reference(w, h, use_mask):
+    from slide2vec.encoders.models.titan import _lean_alibi_bias
+
+    torch.manual_seed(0)
+    bg_mask = (torch.rand(1, w, h) > 0.3) if use_mask else None
+    reference = _reference_get_alibi(12, w, h, bg_mask)
+    lean = _lean_alibi_bias(
+        SimpleNamespace(num_heads=12), w, h, bg_mask, device="cpu", dtype=torch.float16
+    )
+
+    assert torch.equal(lean, reference)
+    # padded row stride is what makes SDPA's memory-efficient kernel accept the bias
+    assert lean.stride(-2) % 8 == 0
+
+
+def test_patch_declines_unrecognized_remote_code():
+    from slide2vec.encoders.models.titan import _patch_titan_remote_code
+
+    assert _patch_titan_remote_code(_FakeTitan()) is False
+
+
+def test_patched_preprocess_keeps_features_dtype():
+    import sys
+    import types
+
+    from slide2vec.encoders.models.titan import _patch_titan_remote_code
+
+    module = types.ModuleType("_fake_titan_remote")
+
+    def preprocess_features(features, coords, patch_size_lv0):
+        # fp32 zero-grid + index_add_, as in the real remote code: raises on fp16
+        grid = torch.zeros(4, features.size(-1))
+        grid.index_add_(0, torch.arange(features.size(0)) % 4, features)
+        return grid, torch.zeros(4, 2, dtype=torch.int64), torch.ones(4, dtype=torch.bool)
+
+    module.preprocess_features = preprocess_features
+
+    class _FakeViT:
+        pos_encode_type = "alibi"
+        num_heads = 12
+        patch_embed = _pos_embed = norm_pre = blocks = norm = object()
+        masked_im_modeling = False
+
+        def forward_features(self, x, coords=None, mask=None, bg_mask=None):
+            return x
+
+    _FakeViT.__module__ = module.__name__
+    sys.modules[module.__name__] = module
+    try:
+        model = SimpleNamespace(vision_encoder=_FakeViT())
+        assert _patch_titan_remote_code(model) is True
+        # repeat call is an idempotent no-op, not a double wrap
+        assert _patch_titan_remote_code(model) is True
+
+        fp16_features = torch.randn(6, 32, dtype=torch.float16)
+        grid, _, _ = module.preprocess_features(fp16_features, None, 512)
+        assert grid.dtype == torch.float16
+    finally:
+        del sys.modules[module.__name__]

slide2vec-5.8.0/slide2vec/encoders/models/titan.py

@@ -1,58 +0,0 @@
-"""TITAN slide encoder implementation."""
-
-import numpy as np
-import torch
-from transformers import AutoModel
-
-from slide2vec.encoders.base import SlideEncoder, preferred_default_device, resolve_requested_output_variant
-from slide2vec.encoders.registry import register_encoder
-
-
-@register_encoder(
-    "titan",
-    level="slide",
-    tile_encoder="conchv15",
-    tile_encoder_output_variant="default",
-    output_variants={"default": {"encode_dim": 768}},
-    default_output_variant="default",
-    supported_spacing_um=0.5,
-    precision="fp16",
-    source="MahmoodLab/TITAN",
-)
-class TitanSlideEncoder(SlideEncoder):
-    def __init__(self, *, output_variant: str | None = None):
-        self._model = AutoModel.from_pretrained("MahmoodLab/TITAN", trust_remote_code=True).eval()
-        self._device = preferred_default_device()
-        self._output_variant = resolve_requested_output_variant(output_variant)
-
-    @property
-    def encode_dim(self) -> int:
-        return 768
-
-    @property
-    def device(self) -> torch.device:
-        return self._device
-
-    def to(self, device: torch.device | str) -> "TitanSlideEncoder":
-        self._device = torch.device(device)
-        self._model = self._model.to(self._device)
-        return self
-
-    def encode_slide(
-        self,
-        tile_features: torch.Tensor,
-        coordinates: torch.Tensor | None = None,
-        *,
-        tile_size_lv0: int | None = None,
-    ) -> torch.Tensor:
-        if coordinates is None or tile_size_lv0 is None:
-            raise ValueError("TITAN slide encoding requires coordinates and tile_size_lv0")
-        if tile_features.ndim == 2:
-            tile_features = tile_features.unsqueeze(0)
-        if coordinates.ndim == 2:
-            coordinates = coordinates.unsqueeze(0)
-        return self._model.encode_slide_from_patch_features(
-            tile_features,
-            coordinates.long(),
-            np.int64(tile_size_lv0),
-        ).squeeze(0)