From 1b225753cd3f96937232bdd1071a0c1964d15e09 Mon Sep 17 00:00:00 2001 From: mcDandy Date: Sun, 18 Oct 2026 12:55:08 +0200 Subject: [PATCH] Remove redundant stack copies after execution The input stack is already copied when remember_stack is off, and the removed line is a no-op when it is on. Co-Authored-By: Claude Opus 5.5 --- more_math/AudioMathNode.py | 1 - more_math/ConditioningMathNode.py | 1 - more_math/ImageMathNode.py | 2 -- more_math/LatentMathNode.py | 4 ---- more_math/MaskMathNode.py | 2 -- more_math/StringMathNode.py | 2 -- tests/test_unified_visitor_regressions.py | 25 +++++++++++++++++++++++ 7 files changed, 25 insertions(+), 12 deletions(-) diff --git a/more_math/AudioMathNode.py b/more_math/AudioMathNode.py index beea704..0312840 100644 --- a/more_math/AudioMathNode.py +++ b/more_math/AudioMathNode.py @@ -167,6 +167,5 @@ class AudioMathNode(io.ComfyNode): res_list.append({"waveform": result_chunk, "sample_rate": sample_rate}) return (res_list, stack) else: - stack = stack if remember_stack else copy.deepcopy(stack) return ([{"waveform": result, "sample_rate": sample_rate}], stack) diff --git a/more_math/ConditioningMathNode.py b/more_math/ConditioningMathNode.py index 7cd2480..8f4ba3b 100644 --- a/more_math/ConditioningMathNode.py +++ b/more_math/ConditioningMathNode.py @@ -251,5 +251,4 @@ class ConditioningMathNode(io.ComfyNode): base[0] = (rtensor, new_dict) res_list = [base] - stack = stack if remember_stack else copy.deepcopy(stack) return (move_to_device(res_list, original_device),stack) diff --git a/more_math/ImageMathNode.py b/more_math/ImageMathNode.py index 3820b0f..1bd9875 100644 --- a/more_math/ImageMathNode.py +++ b/more_math/ImageMathNode.py @@ -159,8 +159,6 @@ class ImageMathNode(io.ComfyNode): res_list = [] for result_chunk in res: res_list.append(result_chunk) - stack = stack if remember_stack else copy.deepcopy(stack) return (res_list, stack) else: - stack = stack if remember_stack else copy.deepcopy(stack) return ([result], stack) diff --git a/more_math/LatentMathNode.py b/more_math/LatentMathNode.py index 8e42e70..fccc66f 100644 --- a/more_math/LatentMathNode.py +++ b/more_math/LatentMathNode.py @@ -242,7 +242,6 @@ class LatentMathNode(io.ComfyNode): if getattr(result_t, "is_nested", False): rl = result_latent.copy() rl["samples"] = result_t - stack = stack if remember_stack else copy.deepcopy(stack) return (move_to_device([rl], original_device), stack) # If the visitor produced a list/tuple, emit each element as a separate latent. @@ -252,7 +251,6 @@ class LatentMathNode(io.ComfyNode): rl = result_latent.copy() rl["samples"] = comp results.append(rl) - stack = stack if remember_stack else copy.deepcopy(stack) return (move_to_device(results, original_device), stack) if batching > 0: @@ -262,9 +260,7 @@ class LatentMathNode(io.ComfyNode): rl = result_latent.copy() rl["samples"] = result_tensor results.append(rl) - stack = stack if remember_stack else copy.deepcopy(stack) return (move_to_device(results, original_device), stack) rl = result_latent.copy() rl["samples"] = result_t - stack = stack if remember_stack else copy.deepcopy(stack) return (move_to_device([rl], original_device), stack) diff --git a/more_math/MaskMathNode.py b/more_math/MaskMathNode.py index f6414db..8f2fc38 100644 --- a/more_math/MaskMathNode.py +++ b/more_math/MaskMathNode.py @@ -140,8 +140,6 @@ class MaskMathNode(io.ComfyNode): res_list = [] for result_chunk in res: res_list.append(result_chunk) - stack = stack if remember_stack else copy.deepcopy(stack) return (res_list, stack) else: - stack = stack if remember_stack else copy.deepcopy(stack) return ([result], stack) diff --git a/more_math/StringMathNode.py b/more_math/StringMathNode.py index 9fed137..358f5d6 100644 --- a/more_math/StringMathNode.py +++ b/more_math/StringMathNode.py @@ -105,9 +105,7 @@ class StringMathNode(io.ComfyNode): for start in range(0, len(s), n): yield s[start:start+n] - stack = stack if remember_stack else copy.deepcopy(stack) return (chunks(result,batching), stack) else: - stack = stack if remember_stack else copy.deepcopy(stack) return ([result], stack) diff --git a/tests/test_unified_visitor_regressions.py b/tests/test_unified_visitor_regressions.py index 0e9bc97..97bfafd 100644 --- a/tests/test_unified_visitor_regressions.py +++ b/tests/test_unified_visitor_regressions.py @@ -9,6 +9,9 @@ from PIL import ImageFont from more_math.Parser.UnifiedMathVisitor import ChangeTrackingDict, MathDict, UnifiedMathVisitor from more_math.helper_functions import checkLazyNew, parse_expr from more_math.FloatMathNode import FloatMathNode +from more_math.ImageMathNode import ImageMathNode +from more_math.LatentMathNode import LatentMathNode +from more_math.MaskMathNode import MaskMathNode def _visit(expression, variables): @@ -284,6 +287,28 @@ def test_unconnected_stack_is_not_shared_between_runs(): assert stack == {0: [1.0]} +@pytest.mark.parametrize("node, value, extra", [ + (ImageMathNode, torch.zeros(2, 4, 4, 3), {"use_compute_device": False}), + (MaskMathNode, torch.zeros(2, 4, 4), {}), + (LatentMathNode, {"samples": torch.zeros(2, 4, 4, 4)}, {"use_compute_device": False}), +]) +@pytest.mark.parametrize("batching", [0, 1]) +def test_stack_across_list_runs(node, value, extra, batching): + # ComfyUI runs the node once per list item and passes the same stack object every time + def run(stack, remember): + return node.execute(V={"V0": value}, F={}, Expression="stack_push(0, 1); V0", batching=batching, remember_stack=remember, stack=stack, **extra)[-1] + + shared = {} + run(shared, True) + assert run(shared, True) == {0: [1.0, 1.0]} + + fresh = {} + first, second = run(fresh, False), run(fresh, False) + assert fresh == {} + assert first == second == {0: [1.0]} + assert first is not second + + def test_top_level_return_stops_script(): assert _visit("x = 1; return 5; x = 2; x", {}) == 5