diff --git a/CHANGELOG.md b/CHANGELOG.md index f0bb851..818e46b 100644 --- a/CHANGELOG.md +++ b/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. diff --git a/MANIFEST.json b/MANIFEST.json index ca25377..ca82afe 100644 --- a/MANIFEST.json +++ b/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", diff --git a/_core/runtime_bootstrap.py b/_core/runtime_bootstrap.py index 911f5e1..1564944 100644 --- a/_core/runtime_bootstrap.py +++ b/_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) diff --git a/pyproject.toml b/pyproject.toml index 201642e..6173d50 100644 --- a/pyproject.toml +++ b/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" diff --git a/tests/test_sam_image_cache.py b/tests/test_sam_image_cache.py new file mode 100644 index 0000000..01c6635 --- /dev/null +++ b/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() +