Address reviews to fix stale masks, prompt routing and device selection for interactive segmentation

micro_sam/precompute_state.py

@@ -538,6 +538,7 @@ def _cache_ais_state_v2(
     tile_shape: Optional[Tuple[int, int]] = None,
     halo: Optional[Tuple[int, int]] = None,
     device: Optional[str] = None,
+    devices: Optional[Union[str, torch.device, Sequence[Union[str, torch.device]]]] = None,
     z_block: Optional[int] = None,
     z_halo: Optional[int] = None,
     verbose: bool = True,
@@ -566,7 +567,9 @@ def _cache_ais_state_v2(
         is_tiled: Whether to use the tiled segmenter. By default inferred from the embeddings.
         tile_shape: The tile shape for the tiled segmenter.
         halo: The tile overlap for the tiled segmenter.
-        device: The device to run inference on.
+        device: The device the decoder lives on.
+        devices: The device or devices to run the decoder inference on. By default all visible CUDA
+            devices are used; pass a single device to pin inference to it.
         z_block: Number of slices to decode per z block for volumes.
         z_halo: Number of overlapping slices between z blocks for volumes.
         verbose: Whether to run verbose.
@@ -581,7 +584,9 @@ def _cache_ais_state_v2(
     if is_tiled is None:
         is_tiled = image_embeddings is not None and image_embeddings.get("input_size") is None
 
-    segmenter = get_unisam2_segmentation_generator(decoder, is_tiled=is_tiled, device=device)
+    segmenter = get_unisam2_segmentation_generator(
+        decoder, is_tiled=is_tiled, device=device, inference_device=devices
+    )
 
     key_index = i if state_index is None else state_index
     key, signature = None, None
@@ -634,14 +639,14 @@ def _resolve_unisam2_decoder(model_type, checkpoint_path, device):
 
 
 def _cache_autoseg_state_for_file(
-    predictor, decoder, model_type, image_data, embeddings, save_path, ndim, verbose,
+    predictor, decoder, model_type, image_data, embeddings, save_path, ndim, verbose, devices=None,
 ):
     """Cache the SAM2 automatic-segmentation state for one file: AIS if a decoder is given, else AMG."""
     if decoder is not None:  # AIS segments the whole image / volume in one pass.
         device = next(decoder.parameters()).device
         cache_autoseg_state(
             "ais", decoder, image_data, embeddings, save_path, ndim=ndim, model_type=model_type,
-            device=device, verbose=verbose,
+            device=device, devices=devices, verbose=verbose,
         )
     elif ndim == 2:  # AMG on a single 2d image.
         model = getattr(predictor, "model", predictor)
@@ -695,8 +700,9 @@ def precompute_state(
             instead of grid-based mask generation (AMG).
         batch_size: The number of tiles / slices per model call. Pass None to select a throughput-efficient
             value per device. Ignored by the model families that do not support batching (VFM encoders).
-        devices: The device or devices to compute the embeddings on. By default all visible CUDA devices
-            are used. Only supported for SAM2 ('hvit_*') models.
+        devices: The device or devices to compute the embeddings and the decoder-based (AIS)
+            automatic-segmentation state on. By default all visible CUDA devices are used. Only
+            supported for SAM2 ('hvit_*') models.
     """
     # Imported lazily to avoid a circular import ('_state' imports from this module).
     from micro_sam.sam_annotator._state import _get_sam_model
@@ -762,4 +768,5 @@ def precompute_state(
         if precompute_autoseg_state:
             _cache_autoseg_state_for_file(
                 predictor, decoder, model_type, image_data, embeddings, save_path, file_ndim, verbose=single,
+                devices=devices,
             )

micro_sam/sam_annotator/_annotator.py

@@ -154,7 +154,8 @@ def _create_widgets(self):
         shape_prompt_layer = self._viewer.layers["prompts"]
         linked_layers = [shape_prompt_layer] if "label" in shape_prompt_layer.current_properties else None
         self._prompt_widget = widgets.create_prompt_menu(
-            self._point_prompt_layer, self._point_labels, linked_layers=linked_layers
+            self._point_prompt_layer, self._point_labels, linked_layers=linked_layers,
+            viewer=self._viewer,
         )
 
         # Create the dictionary for the widgets and get the widgets of the child plugin.
@@ -180,9 +181,11 @@ def _segment_prompts(event):
         def _segment_point_prompts(event):
             self._widgets["segment"](self._viewer)
 
+        # The layer the key reached is the active one, so it goes first: only it relabels a selected
+        # scribble. The viewer-level fallback below has no active prompt layer, so it relabels none.
         @prompt_layer.bind_key("t", overwrite=True)
         def _toggle_shape_prompt_label(event=None):
-            vutil.toggle_label(self._point_prompt_layer, self._shape_prompt_layer)
+            vutil.toggle_label(self._shape_prompt_layer, self._point_prompt_layer)
 
         @point_prompt_layer.bind_key("t", overwrite=True)
         def _toggle_point_prompt_label(event=None):
@@ -325,7 +328,8 @@ def _rebuild_for_ndim(self, ndim, force=False):
         shape_prompt_layer = self._viewer.layers["prompts"]
         linked_layers = [shape_prompt_layer] if "label" in shape_prompt_layer.current_properties else None
         self._prompt_widget = widgets.create_prompt_menu(
-            self._point_prompt_layer, self._point_labels, linked_layers=linked_layers
+            self._point_prompt_layer, self._point_labels, linked_layers=linked_layers,
+            viewer=self._viewer,
         )
 
         # Rebuild the dimension-specific widgets, keeping the shared embedding widget.

micro_sam/sam_annotator/_state.py

@@ -3,9 +3,10 @@
 https://itnext.io/deciding-the-best-singleton-approach-in-python-65c61e90cdc4
 """
 
+import inspect
 from functools import partial
 from dataclasses import dataclass, field
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 import numpy as np
 from segment_anything import SamPredictor
@@ -96,6 +97,9 @@ class AnnotatorState(metaclass=Singleton):
     embedding_tmpdir: Optional[str] = None
     data_signature: Optional[str] = None
     skip_recomputing_embeddings: Optional[bool] = None
+    # The un-resolved device request, forwarded to every batched backend so inference stays on the
+    # selected device. None means 'auto', i.e. fan out over all visible GPUs.
+    inference_devices: Optional[Union[str, List[str]]] = None
     # Whether the tool showed the one-time CPU info popup this session (not reset on recompute).
     cpu_info_shown: Optional[bool] = None
 
@@ -233,6 +237,11 @@ def initialize_predictor(
                     print(f"Could not load a UniSAM2 decoder from '{decoder_source}': {e}")
                     self.decoder = None
 
+        # The inference devices follow the un-resolved request, not the resolved model placement:
+        # None / 'auto' fans out over every visible GPU, an explicit device stays pinned to it.
+        inference_devices = None if device in (None, "auto") else device
+        self.inference_devices = inference_devices
+
         # Compute the image embeddings.
         if isinstance(save_path, dict) and "features" in save_path:  # i.e. embeddings are precomputed
             self.image_embeddings = save_path
@@ -241,6 +250,11 @@ def initialize_predictor(
         else:  # Otherwise, compute the image embeddings.
             _comp_embed_fn = util.get_embedding_function(model_type)
 
+            # The SAM1 embedding function has no 'devices' parameter.
+            device_kwargs = {}
+            if "devices" in inspect.signature(_comp_embed_fn).parameters:
+                device_kwargs["devices"] = inference_devices
+
             # When no save path is given for a SAM2 volume or a tiled image, cache the embeddings to
             # an ephemeral on-disk zarr instead of holding the whole volume in RAM. All slices at once
             # cost about 200 MB per slice and run out of memory on large volumes. The disk cache lets
@@ -270,6 +284,7 @@ def initialize_predictor(
                 lazy_loading=lazy_loading,
                 pbar_init=pbar_init,
                 pbar_update=pbar_update,
+                **device_kwargs,
             )
             self.embedding_path = save_path
 
@@ -282,7 +297,7 @@ def initialize_predictor(
                 from micro_sam.v2.prompt_based_segmentation import TiledPromptableSegmentation3D
                 self.interactive_segmenter = TiledPromptableSegmentation3D(
                     predictor=self.predictor, volume=image_data,
-                    volume_embeddings=self.image_embeddings, device=device,
+                    volume_embeddings=self.image_embeddings, devices=inference_devices,
                 )
             else:
                 from micro_sam.v2.prompt_based_segmentation import PromptableSegmentation3D
@@ -349,7 +364,8 @@ def relabel_pbar_init(total, _description):
             device = next(self.decoder.parameters()).device
             cache_autoseg_state(
                 "ais", self.decoder, image_data, self.image_embeddings, save_path, ndim=ndim,
-                model_type=resolved_model_type, device=device, pbar_init=init_cb, pbar_update=pbar_update,
+                model_type=resolved_model_type, device=device, devices=self.inference_devices,
+                pbar_init=init_cb, pbar_update=pbar_update,
             )
         elif ndim == 2:  # AMG on a single 2d image.
             if pbar_init is not None:
@@ -448,6 +464,7 @@ def reset_state(self):
         self.ndim = None
         self.image_name = None
         self.embedding_path = None
+        self.inference_devices = None
         self.automatic_segmenter = None
         self.autoseg_state = None
         self.decoder = None

micro_sam/sam_annotator/_tooltips.py

@@ -8,7 +8,7 @@
         ),
         "cache_state": "Cache the automatic segmentation state to disk for faster (re)runs.",
         "custom_weights": "Select custom model weights, for example from a model that you finetuned.",
-        "device": "Select the computational device to use for processing.",
+        "device": "Select the computational device to use for processing. 'auto' picks the best available device and uses all visible GPUs; selecting a device runs on that one only.",  # noqa
         "embeddings_save_path": "Select path to save or load the computed image embeddings.",
         "halo": "Enter overlap values to compute tiled embeddings. Enter only the x-value for a square size.\n Active only when you use tiling.",  # noqa
         "image": "Select the napari image layer.",