Add stack passing

more_math/MaskMathNode.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 MaskMathNode(io.ComfyNode):
@@ -32,15 +33,17 @@ class MaskMathNode(io.ComfyNode):
                     options=["tile", "error", "pad"],
                     default="error",
                     tooltip="How to handle mismatched mask 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.Mask.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)
@@ -72,7 +75,7 @@ class MaskMathNode(io.ComfyNode):
         return needed1
 
     @classmethod
-    def execute(cls, V, F, Expression, length_mismatch="tile"):
+    def execute(cls, V, F, Expression, length_mismatch="tile",stack=[]):
         # Identify all present tensors and their keys
         tensor_keys = [k for k, v in V.items() if v is not None]
         if not tensor_keys:
@@ -139,7 +142,7 @@ class MaskMathNode(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)