From d407e5af4093fabc8cf8d03ea86a3053e25ecc02 Mon Sep 17 00:00:00 2001
From: Daniel Martinek <danda.martinek@gmail.com>
Date: Sat, 21 Mar 2026 23:00:54 +0100
Subject: [PATCH] add stach batch toggle.

---
 more_math/FloatMathNode.py | 19 +++++++++++++++----
 1 file changed, 15 insertions(+), 4 deletions(-)

diff --git a/more_math/FloatMathNode.py b/more_math/FloatMathNode.py
index be6de25..7f857e4 100644
--- a/more_math/FloatMathNode.py
+++ b/more_math/FloatMathNode.py
@@ -37,6 +37,14 @@ class FloatMathNode(io.ComfyNode):
                     types=[io.String, MrmthParseTree],
                     tooltip="Expression to use on inputs",
                 ),
+                io.Bool.Input(
+                    id="remember_stack",
+                    default=False,
+                    display_name="Remember stack across batch",
+                    tooltip=(
+                        "If enabled, stack is copied at output leading for changes being remembered during batch operations (node runs multiple times in sucession). If disabled each batch gets it's own copy of the stack."
+                    ),
+                ),
                 MrmthStack.Input(
                     id="stack", tooltip="Access stack between nodes", optional=True
                 ),
@@ -51,11 +59,13 @@ class FloatMathNode(io.ComfyNode):
 
     @classmethod
     def check_lazy_status(cls, FloatFunc, V,stack={}):
-            return checkLazyNew(FloatFunc,V,V)
-
+    def check_lazy_status(cls, FloatFunc, V, remember_stack=False, stack={}):
+        # remember_stack ani stack nemění lazy logiku
 
     @classmethod
     def execute(cls, FloatFunc, V,stack={}):
+    def execute(cls, FloatFunc, V, remember_stack=False, stack={}):
+        work_stack = stack if remember_stack else (copy.deepcopy(stack) if stack is not None else {})
 
         variables = {}
         # Populate aliases
@@ -84,10 +94,11 @@ class FloatMathNode(io.ComfyNode):
         else:
             tree = FloatFunc
         # scalar execution
-        visitor = UnifiedMathVisitor(variables, [1],state_storage=stack)
+        visitor = UnifiedMathVisitor(variables, [1], state_storage=work_stack)
         result = visitor.visit(tree)
 
         # Result might be float or tensor(scalar)
         if torch.is_tensor(result):
-             result = result[0].item()
         return (float(result),stack)
+        returned_stack = work_stack if remember_stack else copy.deepcopy(work_stack)
+        return (float(result), returned_stack)
\ No newline at end of file