fix: Block influence no longer breaks ComfyUI's memory recorder when a Taste key is set (every run died with 'aimdo memory compile error' unless another module switched it off); found on a real GPU

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>

modules/sampling/block_influence/measure.py

@@ -1,5 +1,7 @@
 """The probe's arithmetic and its stored profile: no model, no routes."""
 
+import contextlib
+
 import torch
 
 from ..._core import config, dit_hooks, registry
@@ -27,6 +29,17 @@ def set_enabled(on: bool) -> bool:
     return bool(on)
 
 
+def _aside():
+    """The measurement's own tensors stay out of ComfyUI's memory recorder (aimdo), which assumes every step
+    allocates the same: the running sums grow differently from step to step and broke it ("aimdo memory compile
+    error") whenever no other module had switched the recorder off. The block itself still runs inside it."""
+    try:
+        from comfy.model_prefetch import pause_malloc_graph
+    except ImportError:                          # a ComfyUI without the recorder has nothing to pause
+        return contextlib.nullcontext()
+    return pause_malloc_graph()
+
+
 class Tally:
     """One run's per-block sums, kept on the device until the run ends.
 
@@ -58,27 +71,29 @@ class Tally:
         """Run the block and record how far it moved the picture rows."""
         src = args.get("img")
         sel = before = None
-        if torch.is_tensor(src) and src.dim() == 2:
-            span = self._span(args, src)
-            if span is not None:
-                sel = slice(span[0], span[1], max(1, (span[1] - span[0]) // self.max_rows))
-                # A COPY, read before the block runs: a block that writes its residual in
-                # place leaves `src` already holding the result, and every delta is zero.
-                before = src[sel].detach().to(torch.float32, copy=True)
+        with _aside():
+            if torch.is_tensor(src) and src.dim() == 2:
+                span = self._span(args, src)
+                if span is not None:
+                    sel = slice(span[0], span[1], max(1, (span[1] - span[0]) // self.max_rows))
+                    # A COPY, read before the block runs: a block that writes its residual in
+                    # place leaves `src` already holding the result, and every delta is zero.
+                    before = src[sel].detach().to(torch.float32, copy=True)
         out = original(args)
-        if before is not None and torch.is_tensor(out.get("img")) \
-                and out["img"].shape[0] == src.shape[0]:
-            delta = out["img"][sel].detach().float() - before
-            push = delta.norm()
-            self._add("raw", block, push)
-            self._add("ratio", block, push / before.norm().clamp(min=1e-8))
-            last, last_block = self._prev
-            # A block index that does not exceed the last one means a new step began:
-            # pairing across it would compare the end of one step to the start of the next.
-            if last is not None and block > last_block and last.shape == delta.shape:
-                self._add("novelty", block, torch.nn.functional.cosine_similarity(
-                    delta.flatten(), last.flatten(), dim=0))
-            self._prev = (delta, block)
+        with _aside():
+            if before is not None and torch.is_tensor(out.get("img")) \
+                    and out["img"].shape[0] == src.shape[0]:
+                delta = out["img"][sel].detach().float() - before
+                push = delta.norm()
+                self._add("raw", block, push)
+                self._add("ratio", block, push / before.norm().clamp(min=1e-8))
+                last, last_block = self._prev
+                # A block index that does not exceed the last one means a new step began:
+                # pairing across it would compare the end of one step to the start of the next.
+                if last is not None and block > last_block and last.shape == delta.shape:
+                    self._add("novelty", block, torch.nn.functional.cosine_similarity(
+                        delta.flatten(), last.flatten(), dim=0))
+                self._prev = (delta, block)
         return out
 
     def blocks_seen(self) -> int:

modules/sampling/block_influence/tests/test_block_influence.py

@@ -358,3 +358,32 @@ def test_flatness_reads_raw_push_so_a_growing_stream_with_equal_pushes_is_flat()
     rows = [{"reward": w, "rows": {"ratio": torch.tensor(ratio), "raw": torch.ones(10),
                                    "novelty": torch.full((10,), 0.5)}} for w in (1.0, 1.0, -1.0, -1.0)]
     assert profile(rows)["flatness"] == pytest.approx(0.0, abs=1e-6)
+
+
+def test_the_measurement_stays_out_of_comfys_memory_recorder_but_the_block_does_not(monkeypatch):
+    """With a Taste key set the recording is on in every run. Its running sums allocate differently from step to
+    step, which broke ComfyUI's memory recorder ("aimdo memory compile error", seen on a real GPU) unless another
+    module had switched the recorder off. The measuring is paused; the block itself must still run recorded."""
+    import contextlib
+    import comfy.model_prefetch as mp
+    from modules.sampling.block_influence import measure
+    paused = {"now": False, "block_ran_paused": None, "entered": 0}
+
+    @contextlib.contextmanager
+    def pause(sync=False):
+        paused["now"], paused["entered"] = True, paused["entered"] + 1
+        try:
+            yield
+        finally:
+            paused["now"] = False
+
+    monkeypatch.setattr(mp, "pause_malloc_graph", pause, raising=False)     # older ComfyUIs lack it
+    tally = measure.Tally(2)
+
+    def block(args):
+        paused["block_ran_paused"] = paused["now"]
+        return {"img": args["img"] + 1.0}
+
+    tally.measure(0, {"img": torch.zeros(8, 4)}, block)
+    assert paused["block_ran_paused"] is False, "the block must run inside the recorder, as it always did"
+    assert paused["entered"] >= 1, "the measurement's own allocations must be paused"