Fix SAM image cache invalidation (0.3.2)

CHANGELOG.md

@@ -2,6 +2,11 @@
 
 User-visible changes are recorded here. This file describes package content; publication tags and release channels are separate facts.
 
+## 0.3.2 - development candidate - 2026-09-09
+
+- Refresh Eye Mask SAM embeddings when the input image changes or its pixels are modified in place. Preserve reuse for repeated refinement of the same unchanged image.
+- Prevent Python object-identity reuse from selecting a previous image's embedding.
+
 ## 0.3.1 - development candidate - 2026-09-09
 
 - Metadata Killer displays a small, metadata-verified preview while keeping the verified full-resolution original for opening and downloading.

MANIFEST.json

@@ -132,7 +132,7 @@
     "MATRIX_SpectralSampler": "MATRIX LAB/Sampling & Detail",
     "MATRIX_AIInfluencerResolution2K4K": "MATRIX LAB/Resolution & Layout"
   },
-  "compiler_sha256": "fcc055cedd826ac4f38f1da2ae7097472f029b7fcc52d7a075b53bbbf7d8ed53",
+  "compiler_sha256": "ecb4e0a37e55c7eb60f4263e1dfd0f6b0e7888c32c812922570b60e8ad5dd882",
   "declaration_hashes": {
     "MATRIX_AIInfluencerResolution": "9dd100da79f07f132edda489d044fc2389017ac74ec7401f33d63f0f2ece3ce7",
     "MATRIX_AutoPrompter": "1f9b62b8c5b9d5eba5596aeb3493b1d62b7978b5a989a904f09f5bf1af9006b9",
@@ -178,11 +178,11 @@
   "pack_id": "matrix-lab-nodes",
   "runtime_asset_registry_sha256": "ab1955cfe5e8d05b832a4f63cae121d9afd6f74d3a9048b0107cf634659cada8",
   "status": "development",
-  "version": "0.3.1",
+  "version": "0.3.2",
   "candidate_provenance": {
     "strategy": "current-canonical-base-plus-additive-factory-node",
     "base_builder_sha256": "831ce90d7ec059f0538c8b690e2078339fddd81ecf8197fb8ec4767f920606b4",
-    "compiler_sha256": "fcc055cedd826ac4f38f1da2ae7097472f029b7fcc52d7a075b53bbbf7d8ed53",
+    "compiler_sha256": "ecb4e0a37e55c7eb60f4263e1dfd0f6b0e7888c32c812922570b60e8ad5dd882",
     "builder_sha256": "f60506d9d5a1597c860f08d6b8a8119a265b4e50ef42e4930defc60e9bc55ab3",
     "existing_node_implementation_policy": "fresh-canonical-base-build",
     "runtime_registry_policy": "owned-by-canonical-base-builder",

_core/runtime_bootstrap.py

@@ -610,7 +610,10 @@ def segment_person(prepared):
 
 def _eye_mask_adapters(folder_paths, model_management):
     """Own ultralytics call and SAM refinement for mask.eye-region (no Impact code path)."""
-    state = {}
+    import threading
+    import weakref
+
+    state = {"predictor_lock": threading.RLock()}
 
     def _models():
         if "yolo" not in state:
@@ -634,7 +637,8 @@ def _models():
             state["yolo"] = YOLO(str(bbox_path))
             state["predictor"] = SamPredictor(sam)
             state["device"] = device
-            state["image_key"] = None
+            state["image_ref"] = None
+            state["image_snapshot"] = None
         return state
 
     def detect(detector_id, image_rgb_uint8, *, conf, imgsz):
@@ -664,16 +668,27 @@ def refine(image_rgb_uint8, box_xyxy, point_xy):
         import numpy as np
         models = _models()
         pixels = np.ascontiguousarray(image_rgb_uint8)
-        key = (id(image_rgb_uint8), pixels.shape)
-        if models["image_key"] != key:
-            models["predictor"].set_image(pixels)
-            models["image_key"] = key
-        masks, scores, _ = models["predictor"].predict(
-            point_coords=np.asarray([point_xy], dtype=np.float32),
-            point_labels=np.asarray([1], dtype=np.int64),
-            box=np.asarray(box_xyxy, dtype=np.float32),
-            multimask_output=True,
-        )
+        with models["predictor_lock"]:
+            image_ref = models["image_ref"]
+            cached_image = image_ref() if image_ref is not None else None
+            snapshot = models["image_snapshot"]
+            unchanged = (
+                cached_image is image_rgb_uint8
+                and snapshot is not None
+                and snapshot.shape == pixels.shape
+                and snapshot.dtype == pixels.dtype
+                and np.array_equal(snapshot, pixels)
+            )
+            if not unchanged:
+                models["predictor"].set_image(pixels)
+                models["image_ref"] = weakref.ref(image_rgb_uint8)
+                models["image_snapshot"] = pixels.copy()
+            masks, scores, _ = models["predictor"].predict(
+                point_coords=np.asarray([point_xy], dtype=np.float32),
+                point_labels=np.asarray([1], dtype=np.int64),
+                box=np.asarray(box_xyxy, dtype=np.float32),
+                multimask_output=True,
+            )
         if len(masks) == 0:
             return np.zeros(pixels.shape[:2], dtype=np.float32)
         return masks[int(np.argmax(scores))].astype(np.float32)

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "matrix-lab-nodes"
-version = "0.3.1"
+version = "0.3.2"
 description = "MATRIX LAB custom nodes for ComfyUI."
 readme = "README.md"
 requires-python = ">=3.10"

tests/test_sam_image_cache.py

@@ -0,0 +1,155 @@
+"""Portable regression for the shipped SAM image-embedding cache.
+
+The predictor is a CPU fake; tests execute the actual distributed adapter function.
+They do not assert detector quality, GPU behavior, or unrelated runtime readiness.
+"""
+from __future__ import annotations
+
+import ast
+import gc
+from pathlib import Path
+import types
+import unittest
+from unittest import mock
+import weakref
+import threading
+from concurrent.futures import ThreadPoolExecutor, TimeoutError as FutureTimeoutError
+
+import numpy as np
+
+
+class SamImageCacheTests(unittest.TestCase):
+    def setUp(self):
+        runtime = Path(__file__).resolve().parents[1] / "_core" / "runtime_bootstrap.py"
+        tree = ast.parse(runtime.read_text(encoding="utf-8"))
+        function = next(node for node in tree.body
+                        if isinstance(node, ast.FunctionDef) and node.name == "_eye_mask_adapters")
+        namespace = {
+            "_load_registry": lambda: {},
+            "_asset_by_role": lambda registry, role: {"logical_id": role},
+            "_verified_asset": lambda paths, entry: (None, Path("fixture-model")),
+            "id": lambda value: 7,  # deterministically collide the former id+shape key
+        }
+        exec(compile(ast.Module(body=[function], type_ignores=[]), str(runtime), "exec"), namespace)
+        instances = []
+
+        class Model:
+            def to(self, **kwargs):
+                return self
+            def eval(self):
+                return self
+
+        class Predictor:
+            def __init__(self, model):
+                self.images = []
+                self.fail_next = False
+                instances.append(self)
+
+            def set_image(self, pixels):
+                if self.fail_next:
+                    self.fail_next = False
+                    raise RuntimeError("injected encoder failure")
+                self.images.append(pixels.copy())
+
+            def predict(self, **kwargs):
+                pixels = self.images[-1]
+                mask = np.full(pixels.shape[:2], int(pixels[0, 0, 0]), dtype=np.float32)
+                return np.stack([mask]), np.array([1.0]), None
+
+        ultra = types.ModuleType("ultralytics")
+        ultra.YOLO = lambda path: object()
+        segment = types.ModuleType("segment_anything")
+        segment.SamPredictor = Predictor
+        segment.sam_model_registry = {"vit_b": lambda **kwargs: Model()}
+        self.imports = mock.patch.dict("sys.modules", {"ultralytics": ultra, "segment_anything": segment})
+        self.imports.start()
+        self.addCleanup(self.imports.stop)
+        _, self.refine = namespace["_eye_mask_adapters"](
+            None, types.SimpleNamespace(get_torch_device=lambda: "cpu"))
+        self.instances = instances
+
+    def run_refine(self, image):
+        return self.refine(image, (0, 0, 4, 4), (2, 2))
+
+    def test_unchanged_reuses_but_mutation_and_new_image_refresh(self):
+        first = np.zeros((8, 8, 3), dtype=np.uint8)
+        self.run_refine(first)
+        self.run_refine(first)
+        first[0, 0, 0] = 1
+        self.assertEqual(float(self.run_refine(first)[0, 0]), 1)
+        second = np.full(first.shape, 2, dtype=np.uint8)
+        self.assertEqual(float(self.run_refine(second)[0, 0]), 2)
+        self.run_refine(second)
+        self.assertEqual([int(x[0, 0, 0]) for x in self.instances[0].images], [0, 1, 2])
+
+    def test_noncontiguous_same_view_reuses_and_mutation_refreshes(self):
+        base = np.zeros((8, 16, 3), dtype=np.uint8)
+        image = base[:, ::2]
+        self.assertFalse(image.flags.c_contiguous)
+        self.run_refine(image)
+        self.run_refine(image)
+        base[0, 0, 0] = 9
+        self.assertEqual(float(self.run_refine(image)[0, 0]), 9)
+        self.assertEqual(len(self.instances[0].images), 2)
+
+    def test_cache_does_not_keep_input_alive_and_new_object_refreshes(self):
+        first = np.zeros((8, 8, 3), dtype=np.uint8)
+        pointer = weakref.ref(first)
+        self.run_refine(first)
+        del first
+        gc.collect()
+        self.assertIsNone(pointer())
+        self.run_refine(np.full((8, 8, 3), 3, dtype=np.uint8))
+        self.assertEqual(len(self.instances[0].images), 2)
+
+    def test_concurrent_refinements_cannot_swap_embeddings(self):
+        self.run_refine(np.zeros((8, 8, 3), dtype=np.uint8))
+        predictor = self.instances[0]
+        predict_original = predictor.predict
+        first_inside = threading.Event()
+        second_started = threading.Event()
+        allow_first = threading.Event()
+
+        def predict(**kwargs):
+            if threading.current_thread().name.startswith("sam-first"):
+                first_inside.set()
+                if not allow_first.wait(3):
+                    raise RuntimeError("test coordination timeout")
+            return predict_original(**kwargs)
+
+        predictor.predict = predict
+        first = np.full((8, 8, 3), 1, dtype=np.uint8)
+        second = np.full(first.shape, 2, dtype=np.uint8)
+
+        def second_call():
+            second_started.set()
+            return self.run_refine(second)
+
+        with ThreadPoolExecutor(max_workers=1, thread_name_prefix="sam-first") as a, \
+             ThreadPoolExecutor(max_workers=1, thread_name_prefix="sam-second") as b:
+            result_a = a.submit(self.run_refine, first)
+            self.assertTrue(first_inside.wait(3))
+            result_b = b.submit(second_call)
+            self.assertTrue(second_started.wait(3))
+            try:
+                with self.assertRaises(FutureTimeoutError):
+                    result_b.result(timeout=0.05)
+            finally:
+                allow_first.set()
+            self.assertEqual(float(result_a.result(timeout=3)[0, 0]), 1)
+            self.assertEqual(float(result_b.result(timeout=3)[0, 0]), 2)
+
+    def test_failed_refresh_retries_before_prediction(self):
+        first = np.zeros((8, 8, 3), dtype=np.uint8)
+        self.run_refine(first)
+        second = np.full(first.shape, 5, dtype=np.uint8)
+        self.instances[0].fail_next = True
+        with self.assertRaisesRegex(RuntimeError, "injected encoder failure"):
+            self.run_refine(second)
+        self.assertEqual(float(self.run_refine(second)[0, 0]), 5)
+        self.assertEqual(len(self.instances[0].images), 2)
+
+
+if __name__ == "__main__":
+    unittest.main()
+