From ac5b9b0fc8635fc45cda1e5d9eb69ecdec8443d4 Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Thu, 15 Oct 2026 15:43:22 +0300 Subject: [PATCH] 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 --- modules/sampling/block_influence/measure.py | 55 ++++++++++++------- .../tests/test_block_influence.py | 29 ++++++++++ 2 files changed, 64 insertions(+), 20 deletions(-) diff --git a/modules/sampling/block_influence/measure.py b/modules/sampling/block_influence/measure.py index f0a0c8f9..2420e026 100644 --- a/modules/sampling/block_influence/measure.py +++ b/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: diff --git a/modules/sampling/block_influence/tests/test_block_influence.py b/modules/sampling/block_influence/tests/test_block_influence.py index 1fac5953..7cca86a6 100644 --- a/modules/sampling/block_influence/tests/test_block_influence.py +++ b/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"