nnInteractive 2.1.0__tar.gz → 2.3.0__tar.gz

{nninteractive-2.1.0 → nninteractive-2.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nnInteractive
-Version: 2.1.0
+Version: 2.3.0
 Summary: Inference code for nnInteractive
 Author: Helmholtz Imaging Applied Computer Vision Lab
 Author-email: Fabian Isensee <f.isensee@dkfz-heidelberg.de>
@@ -226,10 +226,14 @@ Requires-Dist: batchgenerators>=0.25.1
 Requires-Dist: fastapi>=0.110
 Requires-Dist: uvicorn[standard]>=0.27
 Requires-Dist: httpx>=0.27
+Requires-Dist: blosc2
 Provides-Extra: dev
 Requires-Dist: black; extra == "dev"
 Requires-Dist: ruff; extra == "dev"
 Requires-Dist: pre-commit; extra == "dev"
+Provides-Extra: client
+Requires-Dist: httpx>=0.27; extra == "client"
+Requires-Dist: blosc2; extra == "client"
 Dynamic: license-file
 
 <img src="imgs/nnInteractive_header_white.png">
@@ -545,6 +549,8 @@ Link: [![arXiv](https://img.shields.io/badge/arXiv-2503.08373-b31b1b.svg)](https
 # License
 Note that while this repository is available under Apache-2.0 license (see [LICENSE](./LICENSE)), the [model checkpoint](https://huggingface.co/nnInteractive/nnInteractive) is `Creative Commons Attribution Non Commercial Share Alike 4.0`! 
 
+Release model folders ship their own `LICENSE` file whose **first line is the license identifier** (e.g. `CC BY-NC-SA 4.0`); any following lines (such as a link to the full license) are ignored by the tool. At load time this first line is read and exposed as `session.license` so applications can display the model's license prominently. If a checkpoint folder has no `LICENSE` file, the official v1 checkpoint is assumed to be `CC BY-NC-SA 4.0` and any other checkpoint reports `!!MISSING!!`.
+
 # Changelog
 
 ### 1.1.2 - 2025-08-02

{nninteractive-2.1.0 → nninteractive-2.3.0}/nnInteractive/inference/inference_session.py

@@ -50,14 +50,16 @@ class nnInteractiveInferenceSession:
     ):
         """
         Only intended to work with nnInteractiveTrainerV2 and its derivatives
+
+        ``use_torch_compile``: compile the network with ``torch.compile``. The
+        first prediction after enabling this is slow (compilation happens lazily
+        on the first forward pass), but every subsequent prediction is faster.
+        This is recommended for the persistent inference server, where the
+        process is long-lived so the one-time compile cost is paid only once and
+        amortized across the whole session lifetime.
         """
         print("session initialized")
 
-        # set as part of initialization
-        assert use_torch_compile is False, (
-            "torch.compile is not supported. The blosc2-backed interaction tensor "
-            "requires numpy↔torch round-trips that break compile tracing."
-        )
         self.network = None
         self.label_manager = None
         self.dataset_json = None
@@ -77,6 +79,11 @@ class nnInteractiveInferenceSession:
         self.channel_mapping: dict = {}
         self.supports_initial_label: bool = True
         self.supports_zero_shot_label_refinement: bool = True
+        # License of the loaded model checkpoint. Set when the model is loaded
+        # (read from the LICENSE file in the checkpoint folder, or derived for
+        # legacy checkpoints without one). Exposed so GUIs can display it once
+        # the session is initialized. "!!MISSING!!" means the license is unknown.
+        self.license: Optional[str] = None
 
         # image specific
         self.interactions = None  # blosc2.NDArray once initialized
@@ -116,6 +123,31 @@ class nnInteractiveInferenceSession:
             and checkpoint.get("init_args", {}).get("configuration") == "3d_fullres_ps192_bs24"
         )
 
+    @classmethod
+    def _load_license(cls, model_training_output_dir: str, plans: dict, checkpoint: dict) -> str:
+        """Determine the license of the model being loaded.
+
+        Reads the ``LICENSE`` file from the checkpoint folder if present.
+        Expected format: the FIRST line is a short license identifier (e.g.
+        ``CC BY-NC-SA 4.0``); any following lines (URL, full text, …) are for
+        human readers and are ignored. Only the first non-empty line is
+        returned, so ``self.license`` stays a short, displayable string.
+
+        If the folder has no ``LICENSE`` file it is most likely a legacy model:
+        the official v1 checkpoint is CC BY-NC-SA 4.0, anything else is reported
+        as ``"!!MISSING!!"`` so callers (e.g. GUIs) can flag the unknown license.
+        """
+        license_file = join(model_training_output_dir, "LICENSE")
+        if isfile(license_file):
+            with open(license_file, "r", encoding="utf-8") as f:
+                for line in f:
+                    line = line.strip()
+                    if line:
+                        return line
+        if cls._is_official_checkpoint(plans, checkpoint):
+            return "CC BY-NC-SA 4.0"
+        return "!!MISSING!!"
+
     def _legacy_default_capability(self) -> dict:
         return {
             "supported_interactions": {
@@ -533,7 +565,13 @@ class nnInteractiveInferenceSession:
             dtype=np.float16,
             chunks=(1, *[min(64, s) for s in shape[1:]]),
             blocks=(1, *[min(32, s) for s in shape[1:]]),
-            cparams={"codec": blosc2.Codec.LZ4, "clevel": 5, "nthreads": min(self.torch_n_threads, os.cpu_count())},
+            # Interactions compress better with NOFILTER, which is also faster than SHUFFLE.
+            cparams={
+                "codec": blosc2.Codec.LZ4,
+                "clevel": 5,
+                "filters": [blosc2.Filter.NOFILTER],
+                "nthreads": min(self.torch_n_threads, os.cpu_count()),
+            },
             dparams={"nthreads": 4},
         )
         self._interactions_shape = shape
@@ -602,7 +640,13 @@ class nnInteractiveInferenceSession:
                 dtype=np.float16,
                 chunks=(1, *[min(64, s) for s in self._interactions_shape[1:]]),
                 blocks=(1, *[min(32, s) for s in self._interactions_shape[1:]]),
-                cparams={"codec": blosc2.Codec.LZ4, "clevel": 5, "nthreads": os.cpu_count()},
+                # Interactions compress better with NOFILTER, which is also faster than SHUFFLE.
+                cparams={
+                    "codec": blosc2.Codec.LZ4,
+                    "clevel": 5,
+                    "filters": [blosc2.Filter.NOFILTER],
+                    "nthreads": os.cpu_count(),
+                },
                 dparams={"nthreads": 4},
             )
         self.current_interaction_intensity = 1.0
@@ -883,6 +927,42 @@ class nnInteractiveInferenceSession:
         else:
             del initial_seg
 
+    @torch.inference_mode()
+    def warmup(self) -> bool:
+        """Run a single dummy forward pass to trigger lazy ``torch.compile`` compilation up front.
+
+        With ``torch.compile`` enabled the network is compiled lazily on its first
+        forward pass, which would otherwise make the user's *first* real prediction
+        slow. Every prediction path — the initial coarse pass, the zoom-out
+        iterations, and the refinement patches — feeds the network an input of
+        identical shape ``[1, num_input_channels + num_interaction_channels,
+        *patch_size]`` (``_build_network_input`` always resizes the crop to
+        ``patch_size``, and refinement crops at exactly ``patch_size``). So a single
+        dummy pass at that shape populates the compile cache and every subsequent
+        real prediction is fast.
+
+        Returns ``True`` if a warmup pass was run, ``False`` if it was a no-op
+        (network not compiled — there is nothing to pre-compile, so a dummy pass
+        would not save the user any time). Mirrors ``_predict``'s autocast/
+        inference-mode context and the float32 input dtype that ``torch.cat``
+        produces when concatenating the float32 image with the fp16 interactions.
+        """
+        if self.network is None or self.configuration_manager is None:
+            raise RuntimeError("warmup() requires an initialized network; call initialize_* first")
+        if not isinstance(self.network, OptimizedModule):
+            return False
+        num_input_channels = (
+            determine_num_input_channels(self.plans_manager, self.configuration_manager, self.dataset_json)
+            + self.num_interaction_channels
+        )
+        patch_size = self.configuration_manager.patch_size
+        dummy = torch.zeros((1, num_input_channels, *patch_size), dtype=torch.float32, device=self.device)
+        with torch.autocast(self.device.type, enabled=True) if self.device.type == "cuda" else dummy_context():
+            self.network(dummy)
+        del dummy
+        empty_cache(self.device)
+        return True
+
     @torch.inference_mode()
     def _predict(self, force_full_refine: bool = False):
         """
@@ -1296,6 +1376,16 @@ class nnInteractiveInferenceSession:
         """
         artifacts = self._load_model_artifacts_from_disk(model_training_output_dir, use_fold, checkpoint_name)
         self.initialize_from_loaded_artifacts(artifacts)
+        # With torch.compile the network is compiled lazily on the first forward pass. For a
+        # locally hosted model that lag would otherwise surface on the user's first real
+        # prediction, where it is far more noticeable than during initialization. Trigger the
+        # compilation now with a dummy forward pass so the cost is paid here instead. warmup()
+        # is a no-op when the network is not compiled. The server takes care of its own warmup
+        # explicitly (it shares one compiled network across sessions via
+        # initialize_from_loaded_artifacts), so we only do this on the direct, local entry point.
+        if self.use_torch_compile:
+            print("torch.compile enabled; warming up (compiling) the network now (this is slow once)...")
+            self.warmup()
 
     def _load_model_artifacts_from_disk(
         self,
@@ -1307,9 +1397,16 @@ class nnInteractiveInferenceSession:
 
         Returns an artifact dict that can be applied to this or any other freshly
         constructed session via :meth:`initialize_from_loaded_artifacts`. The
-        returned references (network, plans_manager, configuration_manager, ...)
-        are safe to share by reference across sessions — they are treated as
-        read-only after construction.
+        returned values are the actual objects (the ``nn.Module`` with its
+        weights and buffers, the plans/configuration managers, the dataset
+        json, the label manager) — not copies. Multiple sessions calling
+        :meth:`initialize_from_loaded_artifacts` with the same dict will all
+        end up with ``self.network`` pointing at the same module instance and
+        the same weight tensors on the GPU. This is safe as long as callers
+        treat these objects as read-only after construction; in the multi-
+        session server that is enforced by running inference under
+        ``@torch.inference_mode()`` and serializing predict calls with a
+        global GPU lock.
 
         Note: this also mutates ``self`` (applies capability, sets pad/decay/
         thickness) because ``num_interaction_channels`` is required to build the
@@ -1344,12 +1441,11 @@ class nnInteractiveInferenceSession:
         checkpoint = torch.load(
             join(model_training_output_dir, fold_folder, checkpoint_name), map_location=self.device, weights_only=False
         )
-        if self._is_official_checkpoint(plans, checkpoint):
-            print(
-                "License reminder: The official nnInteractive checkpoint is licensed under "
-                "Creative Commons Attribution Non Commercial Share Alike 4.0 (CC BY-NC-SA 4.0). "
-                "See the license note in readme.md (# License)."
-            )
+        self.license = self._load_license(model_training_output_dir, plans, checkpoint)
+        print("=" * 80)
+        print("Model license:")
+        print(self.license)
+        print("=" * 80)
         trainer_name = checkpoint["trainer_name"]
         configuration_name = checkpoint["init_args"]["configuration"]
 
@@ -1395,6 +1491,7 @@ class nnInteractiveInferenceSession:
             "dataset_json": dataset_json,
             "trainer_name": trainer_name,
             "label_manager": plans_manager.get_label_manager(dataset_json),
+            "license": self.license,
         }
 
     def initialize_from_loaded_artifacts(self, artifacts: dict):
@@ -1402,7 +1499,10 @@ class nnInteractiveInferenceSession:
 
         ``artifacts`` is the dict returned by :meth:`_load_model_artifacts_from_disk`.
         Useful for spawning multiple sessions that share one loaded model (e.g.
-        the multi-session inference server).
+        the multi-session inference server). All artifact entries — including
+        ``self.network`` — are stored by reference; passing the same dict to
+        multiple sessions does not duplicate the network or its weights in
+        memory.
         """
         self.preferred_scribble_thickness = artifacts["preferred_scribble_thickness"]
         self.interaction_decay = artifacts["interaction_decay"]
@@ -1415,6 +1515,7 @@ class nnInteractiveInferenceSession:
         self.dataset_json = artifacts["dataset_json"]
         self.trainer_name = artifacts["trainer_name"]
         self.label_manager = artifacts["label_manager"]
+        self.license = artifacts["license"]
         if self.use_torch_compile and not isinstance(self.network, OptimizedModule):
             print("Using torch.compile")
             self.network = torch.compile(self.network)

{nninteractive-2.1.0 → nninteractive-2.3.0}/nnInteractive/inference/remote/remote_session.py

@@ -10,9 +10,11 @@ from __future__ import annotations
 
 import json
 import os
+import threading
 import warnings
 from typing import List, Optional, Tuple, Union
 
+import blosc2
 import httpx
 import numpy as np
 import torch
@@ -165,8 +167,16 @@ class nnInteractiveRemoteInferenceSession:
         claim_info = claim_resp.json()
         self._lease_token = claim_info["lease_token"]
         self.idle_timeout_seconds: float = float(claim_info.get("idle_timeout_seconds", 0.0))
+        self.liveness_timeout_seconds: float = float(claim_info.get("liveness_timeout_seconds", 0.0))
         self._http.headers[LEASE_HEADER] = self._lease_token
 
+        # Background liveness heartbeat bookkeeping. Defined before any code that
+        # might raise so close()/__del__ can always reference them safely. The
+        # thread itself is started at the end of __init__, once construction has
+        # fully succeeded.
+        self._stop_heartbeat = threading.Event()
+        self._heartbeat_thread: Optional[threading.Thread] = None
+
         caps = self._get_json(PATH_CAPABILITIES)
 
         # Attributes that mirror the local session so the GUI can introspect them
@@ -182,11 +192,27 @@ class nnInteractiveRemoteInferenceSession:
         self.preferred_scribble_thickness = caps["preferred_scribble_thickness"]
         self.interaction_decay = caps["interaction_decay"]
         self.INFERENCE_SESSION_VERSION = caps["inference_session_version"]
+        # License of the model loaded on the server. Mirrors
+        # nnInteractiveInferenceSession.license so a GUI can display it
+        # regardless of whether it holds a local or remote session.
+        # "!!MISSING!!" means the server could not determine the license.
+        self.license: Optional[str] = caps.get("license")
 
         self.original_image_shape: Optional[Tuple[int, ...]] = None
         self.target_buffer: Union[np.ndarray, torch.Tensor, None] = None
         self.do_autozoom: bool = bool(caps.get("do_autozoom", True))
 
+        # Construction succeeded — start auto-heartbeating to keep the server
+        # from reaping us as a dead client. Beat at half the liveness timeout so
+        # one dropped request still leaves margin. Daemon thread: it never blocks
+        # interpreter exit, and close() joins it cleanly.
+        if self.liveness_timeout_seconds > 0:
+            self._heartbeat_interval = max(5.0, self.liveness_timeout_seconds / 2.0)
+            self._heartbeat_thread = threading.Thread(
+                target=self._heartbeat_loop, name="nnInteractive-heartbeat", daemon=True
+            )
+            self._heartbeat_thread.start()
+
     # ------------------------------------------------------------------ #
     #                       HTTP helpers (private)                       #
     # ------------------------------------------------------------------ #
@@ -272,7 +298,10 @@ class nnInteractiveRemoteInferenceSession:
         self.target_buffer = target_buffer
         self._post_json(
             PATH_SET_TARGET_BUFFER,
-            {"shape": list(target_buffer.shape), "dtype": _buffer_dtype_str(target_buffer)},
+            {
+                "shape": list(target_buffer.shape),
+                "dtype": _buffer_dtype_str(target_buffer),
+            },
         )
 
     def set_do_autozoom(self, do_autozoom: bool) -> None:
@@ -372,7 +401,8 @@ class nnInteractiveRemoteInferenceSession:
             "override_capability_checks": bool(override_capability_checks),
             "interaction_bbox": ([list(b) for b in interaction_bbox] if interaction_bbox is not None else None),
         }
-        resp = self._post_binary(path, meta, pack_array(mask_image))
+        # Interactions (scribble/lasso masks) compress best with NOFILTER; skip auto-selection.
+        resp = self._post_binary(path, meta, pack_array(mask_image, filters=[blosc2.Filter.NOFILTER]))
         self._apply_prediction_response(resp)
 
     def add_initial_seg_interaction(
@@ -397,7 +427,8 @@ class nnInteractiveRemoteInferenceSession:
             "run_prediction": bool(run_prediction),
             "override_capability_checks": bool(override_capability_checks),
         }
-        resp = self._post_binary(PATH_ADD_INITIAL_SEG, meta, pack_array(initial_seg))
+        # Segmentations compress best with NOFILTER; skip auto-selection.
+        resp = self._post_binary(PATH_ADD_INITIAL_SEG, meta, pack_array(initial_seg, filters=[blosc2.Filter.NOFILTER]))
         self._apply_prediction_response(resp)
 
     # ------------------------------------------------------------------ #
@@ -420,17 +451,44 @@ class nnInteractiveRemoteInferenceSession:
             return False
 
     def heartbeat(self) -> float:
-        """Extend this session's idle timeout. Returns remaining seconds.
+        """Tell the server this client is still alive. Returns remaining seconds
+        until the *idle* timeout.
+
+        This proves liveness only: it stops the server from reaping the session
+        as a crashed/dead client, but it does NOT postpone the idle timeout —
+        that is refreshed solely by real user actions (``set_image``,
+        ``add_*_interaction``, …). A session left untouched will therefore still
+        be reaped at the idle timeout even while heartbeats keep arriving.
 
-        For GUIs that keep the app open across long idle stretches: drive
-        this from a timer (e.g. every 60 s) to avoid SessionExpiredError
-        after the configured idle timeout.
+        You normally never call this yourself: the session auto-heartbeats from
+        a background thread for the lifetime of the object.
         """
         resp = self._http.post(PATH_HEARTBEAT)
         _raise_for_lease_errors(resp)
         resp.raise_for_status()
         return float(resp.json().get("remaining_seconds", 0.0))
 
+    def _heartbeat_loop(self) -> None:
+        """Background daemon: prove liveness every ``_heartbeat_interval`` seconds.
+
+        Stops when the session is closed or once the lease is gone. Transient
+        network errors are swallowed so a brief blip doesn't kill the heartbeat;
+        the server's liveness timeout tolerates a few missed beats. Lease expiry
+        (idle reap or server restart) is surfaced to the user on their next real
+        call, not from this thread.
+        """
+        while not self._stop_heartbeat.wait(self._heartbeat_interval):
+            try:
+                self.heartbeat()
+            except SessionExpiredError:
+                break
+            except httpx.HTTPError:
+                continue
+            except Exception:
+                # Never let the daemon thread die noisily (e.g. client closing
+                # concurrently). Bail out quietly.
+                break
+
     def lease_status(self) -> dict:
         """Read-only probe: how much time is left before this session is reaped?
 
@@ -444,6 +502,14 @@ class nnInteractiveRemoteInferenceSession:
         return resp.json()
 
     def close(self) -> None:
+        # Stop the heartbeat thread first so it can't use self._http after we
+        # close it. join() with a short timeout: the thread spends almost all
+        # its time in Event.wait(), which the set() interrupts immediately.
+        self._stop_heartbeat.set()
+        if self._heartbeat_thread is not None:
+            self._heartbeat_thread.join(timeout=5.0)
+            self._heartbeat_thread = None
+
         # Best-effort release so the server can free our slot for other users
         # without waiting for the idle reaper. Swallow errors: the server may
         # already be gone, our lease may already be expired, etc. close()

{nninteractive-2.1.0 → nninteractive-2.3.0}/nnInteractive/inference/remote/serialization.py

@@ -47,8 +47,64 @@ _CODEC_ID = {
 _ID_CODEC = {v: k for k, v in _CODEC_ID.items()}
 
 
-def pack_array(arr: np.ndarray, codec: blosc2.Codec = blosc2.Codec.ZSTD, clevel: int = 3) -> bytes:
-    """Serialize a numpy array to a self-describing compressed byte string."""
+# Fraction of each axis used for the center crop that the filter heuristic compresses.
+_SELECT_FILTER_CROP_FRACTION = 0.25
+
+
+def _compress_all(raw: memoryview, total: int, codec: blosc2.Codec, clevel: int, filters: list) -> int:
+    """Compressed byte length of ``raw`` under ``filters``, chunked exactly as pack_array does."""
+    size = 0
+    nchunks = (total + _CHUNK_SIZE - 1) // _CHUNK_SIZE
+    for i in range(nchunks):
+        start = i * _CHUNK_SIZE
+        end = min(start + _CHUNK_SIZE, total)
+        size += len(blosc2.compress2(raw[start:end], codec=codec, clevel=clevel, filters=filters))
+    return size
+
+
+def _select_filter(arr: np.ndarray, codec: blosc2.Codec, clevel: int) -> "blosc2.Filter":
+    """Pick NOFILTER vs SHUFFLE for ``arr`` by trial-compressing a small centered crop.
+
+    Uses ``compress2`` on the raw bytes — exactly the path pack_array takes — so the decision
+    is consistent with how the whole array is actually compressed. The crop is
+    ``_SELECT_FILTER_CROP_FRACTION`` of each axis (centered), keeping the trial cheap and
+    representative (lands on foreground). Ties go to NOFILTER; any failure falls back to it.
+    """
+    try:
+        crop_shape = [max(1, int(s * _SELECT_FILTER_CROP_FRACTION)) for s in arr.shape]
+        slices = tuple(slice((s - cs) // 2, (s - cs) // 2 + cs) for s, cs in zip(arr.shape, crop_shape))
+        crop = np.ascontiguousarray(arr[slices])
+        raw = memoryview(crop).cast("B")
+        total = raw.nbytes
+
+        best_filter, best_bytes = blosc2.Filter.NOFILTER, None
+        for f in (blosc2.Filter.NOFILTER, blosc2.Filter.SHUFFLE):
+            cb = _compress_all(raw, total, codec, clevel, [f])
+            if best_bytes is None or cb < best_bytes:
+                best_bytes, best_filter = cb, f
+        return best_filter
+    except Exception as e:
+        from warnings import warn
+
+        warn(f"_select_filter failed ({e!r}); falling back to NOFILTER.")
+        return blosc2.Filter.NOFILTER
+
+
+def pack_array(
+    arr: np.ndarray,
+    codec: blosc2.Codec = blosc2.Codec.ZSTD,
+    clevel: int = 3,
+    filters: Optional[list] = None,
+) -> bytes:
+    """Serialize a numpy array to a self-describing compressed byte string.
+
+    ``filters`` is the blosc2 filter pipeline to apply. If ``None`` (the default), the
+    better of NOFILTER/SHUFFLE is auto-selected by trial-compressing a cheap, representative
+    slab — appropriate for images, whose optimum depends on the data. Callers that already
+    know the optimum (interactions and segmentations compress best with NOFILTER) should pass
+    ``[blosc2.Filter.NOFILTER]`` to skip the selection. The chosen filter is self-describing
+    inside the blosc2 frame, so unpack_array (decompress2) needs no changes.
+    """
     arr = np.ascontiguousarray(arr)
     dtype_str = arr.dtype.str.lstrip("<>|=").encode("ascii")
     if arr.dtype.byteorder not in ("=", "|", "<"):
@@ -77,11 +133,17 @@ def pack_array(arr: np.ndarray, codec: blosc2.Codec = blosc2.Codec.ZSTD, clevel:
     raw = memoryview(arr).cast("B")
     total = raw.nbytes
     nchunks = (total + _CHUNK_SIZE - 1) // _CHUNK_SIZE
+
+    if filters is None:
+        # Auto-select the better filter from a small centered crop, using the same
+        # compress2 path as below for consistency.
+        filters = [_select_filter(arr, codec, clevel)]
+
     parts = [header, shape_bytes, struct.pack("<I", nchunks)]
     for i in range(nchunks):
         start = i * _CHUNK_SIZE
         end = min(start + _CHUNK_SIZE, total)
-        chunk = blosc2.compress2(raw[start:end], codec=codec, clevel=clevel)
+        chunk = blosc2.compress2(raw[start:end], codec=codec, clevel=clevel, filters=filters)
         parts.append(struct.pack("<QQ", end - start, len(chunk)))
         parts.append(chunk)
     return b"".join(parts)