Add stack passing

more_math/ImageMathNode.py

@@ -5,6 +5,7 @@ from antlr4 import InputStream, CommonTokenStream
 from .Parser.MathExprLexer import MathExprLexer
 from .Parser.MathExprParser import MathExprParser
 import re
+from .Stack import MrmthStack
 
 class ImageMathNode(io.ComfyNode):
     """
@@ -31,15 +32,17 @@ class ImageMathNode(io.ComfyNode):
                     options=["tile", "error", "pad"],
                     default="error",
                     tooltip="How to handle mismatched image batch sizes. tile: repeat shorter inputs; error: raise error on mismatch; pad: treat missing frames as zero."
-                )
+                ),
+                MrmthStack.Input(id="stack", tooltip="Access stack between nodes",optional=True)
             ],
             outputs=[
                 io.Image.Output(),
+                MrmthStack.Output(),
             ],
         )
 
     @classmethod
-    def check_lazy_status(cls, Expression, V, F, length_mismatch="tile"):
+    def check_lazy_status(cls, Expression, V, F, length_mismatch="tile",stack=[]):
 
         input_stream = InputStream(Expression)
         lexer = MathExprLexer(input_stream)
@@ -71,7 +74,7 @@ class ImageMathNode(io.ComfyNode):
         return needed1
 
     @classmethod
-    def execute(cls, V, F, Expression, length_mismatch="error"):
+    def execute(cls, V, F, Expression, length_mismatch="error",stack=[]):
         # I and F are Autogrow.Type which is dict[str, Any]
 
         # Identify all present tensors and their keys
@@ -144,7 +147,7 @@ class ImageMathNode(io.ComfyNode):
             variables[k] = val if val is not None else 0.0
 
         tree = parse_expr(Expression);
-        visitor = UnifiedMathVisitor(variables, ae.shape,ae.device)
+        visitor = UnifiedMathVisitor(variables, ae.shape,ae.device,state_storage=stack)
         result = visitor.visit(tree)
         result = as_tensor(result, ae.shape)
-        return (result,)
+        return (result,stack)