From 3cd0c1e00afd78140f5997d50969332ecf167f22 Mon Sep 17 00:00:00 2001
From: digital-garbage <133010426+digital-garbage@users.noreply.github.com>
Date: Sun, 27 Sep 2026 17:59:18 +0300
Subject: [PATCH] fix(v5): block repeat and shadow negative sample with the
 model compiler off

Both replace what a block allocates, which crashes aimdo's compiler (seen on a
rental with v4's shadow negative). Attached to the model, so any sampler.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 core/dit_hooks.py                            | 36 ++++++++++++++++++++
 core/tests/test_dit_hooks.py                 | 26 ++++++++++++++
 modules/sampling/block_repeat/__init__.py    |  2 ++
 modules/sampling/shadow_negative/__init__.py |  2 ++
 4 files changed, 66 insertions(+)

diff --git a/core/dit_hooks.py b/core/dit_hooks.py
index 2713be7d..5c1d10f7 100644
--- a/core/dit_hooks.py
+++ b/core/dit_hooks.py
@@ -183,3 +183,39 @@ def target_rows(mod_segments, seq_len, device, stream):
         if mask is not None:
             return mask
     return None
+
+
+def without_compiler(patcher, key: str) -> None:
+    """Run this model's sampling with ComfyUI's model compiler off.
+
+    The compiler (aimdo malloc-graph) records what each block allocates and
+    assumes it never changes. A hook that REPLACES a block's forward with one
+    allocating differently -- a repeated block, a second attention stream --
+    crashes it with "aimdo memory compile error" (v4, on rentals: block repeat,
+    REINS injection, shadow negative). Comfy-Org's own fix is
+    --disable-comfy-compiler; it is read live on every scope, so switching it
+    for the one sampling call is enough and the next run keeps the compiler.
+
+    Attached to the MODEL as an OUTER_SAMPLE wrapper, so it holds for any
+    sampler, not only FunPack's.
+    """
+    from comfy.patcher_extension import WrappersMP
+
+    def outer(executor, *args, **kwargs):
+        import comfy.cli_args
+        cli = comfy.cli_args.args
+        had = hasattr(cli, "disable_comfy_compiler")
+        prior = getattr(cli, "disable_comfy_compiler", None)
+        cli.disable_comfy_compiler = True
+        try:
+            return executor(*args, **kwargs)
+        finally:
+            if had:
+                cli.disable_comfy_compiler = prior
+            else:
+                try:
+                    del cli.disable_comfy_compiler
+                except AttributeError:
+                    pass
+
+    patcher.add_wrapper_with_key(WrappersMP.OUTER_SAMPLE, key, outer)
diff --git a/core/tests/test_dit_hooks.py b/core/tests/test_dit_hooks.py
index b0d29c29..aea4069c 100644
--- a/core/tests/test_dit_hooks.py
+++ b/core/tests/test_dit_hooks.py
@@ -100,3 +100,29 @@ def test_parse_blocks_reports_what_it_ignores(spec, expect, problems):
     from core import dit_hooks
     got, said = dit_hooks.parse_blocks(spec, 4)
     assert got == expect and len(said) == problems
+
+
+def test_without_compiler_switches_it_off_for_the_sample_and_back(tiny_h3):
+    import comfy.cli_args
+    from comfy.patcher_extension import WrappersMP
+    from core import dit_hooks
+    p = tiny_h3.patcher.clone()
+    dit_hooks.without_compiler(p, "funpack.x")
+    (outer,), = p.wrappers[WrappersMP.OUTER_SAMPLE].values()
+    had = hasattr(comfy.cli_args.args, "disable_comfy_compiler")
+    prior = getattr(comfy.cli_args.args, "disable_comfy_compiler", None)
+    seen = outer(lambda: comfy.cli_args.args.disable_comfy_compiler)
+    assert seen is True
+    assert hasattr(comfy.cli_args.args, "disable_comfy_compiler") == had
+    assert getattr(comfy.cli_args.args, "disable_comfy_compiler", None) == prior
+
+
+@pytest.mark.parametrize("module,settings", [
+    ("block_repeat", {"enabled": True, "blocks": "1"}),
+    ("shadow_negative", {"enabled": True}),
+])
+def test_block_replacing_modules_turn_the_compiler_off(tiny_h3, module, settings):
+    from comfy.patcher_extension import WrappersMP
+    from modules.sampling.modifiers.nodes import FunPackLoadModifiers
+    patched, _ = FunPackLoadModifiers.execute(tiny_h3.patcher, {module: settings}).result
+    assert patched.wrappers.get(WrappersMP.OUTER_SAMPLE), f"{module} left the compiler on"
diff --git a/modules/sampling/block_repeat/__init__.py b/modules/sampling/block_repeat/__init__.py
index 1788261a..91511ac1 100644
--- a/modules/sampling/block_repeat/__init__.py
+++ b/modules/sampling/block_repeat/__init__.py
@@ -171,6 +171,8 @@ def install(patcher, values, key):
             dit_hooks.add_block_hook(patcher, key, i, repeat)
         what = f"blocks {','.join(map(str, blocks))} run {times + 1}x each"
 
+    # Replaces what the hooked blocks allocate: see dit_hooks.without_compiler.
+    dit_hooks.without_compiler(patcher, key)
     if video_only:
         what += ", video only"
     if steps:
diff --git a/modules/sampling/shadow_negative/__init__.py b/modules/sampling/shadow_negative/__init__.py
index cf54c500..888081df 100644
--- a/modules/sampling/shadow_negative/__init__.py
+++ b/modules/sampling/shadow_negative/__init__.py
@@ -153,6 +153,8 @@ def install(patcher, values, key):
     for i, block in enumerate(dm.blocks):
         dit_hooks.add_block_hook(patcher, key, i, make(block))
 
+    # Replaces what the hooked blocks allocate: see dit_hooks.without_compiler.
+    dit_hooks.without_compiler(patcher, key)
     note = f"picture {video_scale:g}, sound {audio_scale:g}, steps {start:.2f}-{end:.2f}"
     if taken:
         note += (f"; leaves block(s) {','.join(map(str, taken))} to the feature already there"