Add stack passing

Author: mcDandy

more_math/AudioMathNode.py

@@ -14,6 +14,7 @@
 from .Parser.MathExprLexer import MathExprLexer
 from .Parser.MathExprParser import MathExprParser
 import re
+from .Stack import MrmthStack
 
 class AudioMathNode(io.ComfyNode):
     """
@@ -40,15 +41,17 @@ def define_schema(cls) -> io.Schema:
                     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.Audio.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)
@@ -80,7 +83,7 @@ def check_lazy_status(cls, Expression, V, F, length_mismatch="tile"):
         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 audio inputs and their keys
         tensor_keys = [k for k, v in V.items() if v is not None and isinstance(v, dict) and "waveform" in v]
         if not tensor_keys:
@@ -145,7 +148,7 @@ def execute(cls, V, F, Expression, length_mismatch="tile"):
             variables[k] = val if val is not None else 0.0
 
         tree = parse_expr(Expression);
-        visitor = UnifiedMathVisitor(variables, a_w.shape,a_w.device)
+        visitor = UnifiedMathVisitor(variables, a_w.shape,a_w.device,state_storage=stack)
         result = visitor.visit(tree)
         result = as_tensor(result, a_w.shape)
-        return ({"waveform":result,"sample_rate":sample_rate},)
+        return ({"waveform":result,"sample_rate":sample_rate},stack)

more_math/ClipMathNode.py

@@ -4,6 +4,7 @@
 from .Parser.MathExprLexer import MathExprLexer
 from .Parser.MathExprParser import MathExprParser
 import re
+from .Stack import MrmthStack
 
 
 class CLIPMathNode(io.ComfyNode):
@@ -26,17 +27,19 @@ def define_schema(cls) -> io.Schema:
                     options=["tile", "error", "pad"],
                     default="error",
                     tooltip="How to handle mismatched layer counts. For models, this usually defaults to broadcast (zero for missing layers)."
-                )
+                ),
+                MrmthStack.Input(id="stack", tooltip="Access stack between nodes",optional=True)
             ],
             outputs=[
                 io.Clip.Output(),
+                MrmthStack.Output(),
             ],
         )
 
     tooltip = cleandoc(__doc__)
 
     @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)
@@ -68,7 +71,7 @@ def check_lazy_status(cls, Expression, V, F, length_mismatch="tile"):
         return needed1
 
     @classmethod
-    def execute(cls, V, F, Expression, length_mismatch="tile") -> io.NodeOutput:
+    def execute(cls, V, F, Expression, length_mismatch="tile",stack=[]) -> io.NodeOutput:
         # Determine reference CLIP
         a = V.get("V0")
         if a is None:
@@ -93,9 +96,9 @@ def execute(cls, V, F, Expression, length_mismatch="tile") -> io.NodeOutput:
         # The prompt says aliases are supported in check_lazy_status. Variables map in helper handles logic.
         aliases = {"a": "V0", "b": "V1", "c": "V2", "d": "V3", "w": "F0", "x": "F1", "y": "F2", "z": "F3"}
 
-        patches = calculate_patches_autogrow(Expression, V=patchers_V, F=F, mapping=aliases)
+        patches = calculate_patches_autogrow(Expression, V=patchers_V, F=F, mapping=aliases,stack=stack)
 
         out_clip = a.clone()
         if patches:
             out_clip.add_patches(patches, 1.0, 1.0)
-        return (out_clip,)
+        return (out_clip,stack)

more_math/ConditioningMathNode.py

@@ -8,6 +8,7 @@
 from .Parser.MathExprParser import MathExprParser
 import re
 import copy
+from .Stack import MrmthStack
 
 class ConditioningMathNode(io.ComfyNode):
     """
@@ -36,15 +37,17 @@ def define_schema(cls) -> io.Schema:
                     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."
                 ),
-                io.Int.Input(id="batching")
+                io.Int.Input(id="batching"),
+                MrmthStack.Input(id="stack",optional=True)
             ],
             outputs=[
                 io.Conditioning.Output(is_output_list=True),
+                MrmthStack.Output()
             ],
         )
 
     @classmethod
-    def check_lazy_status(cls, Expression,Expression_pi, V, F,batching, length_mismatch="tile"):
+    def check_lazy_status(cls, Expression,Expression_pi, V, F,batching, length_mismatch="tile",stack=[]):
 
         input_stream = InputStream(Expression)
         lexer = MathExprLexer(input_stream)
@@ -81,7 +84,7 @@ def check_lazy_status(cls, Expression,Expression_pi, V, F,batching, length_misma
         return needed1
 
     @classmethod
-    def execute(cls, V, F, Expression, Expression_pi,batching, length_mismatch="tile"):
+    def execute(cls, V, F, Expression, Expression_pi,batching, length_mismatch="tile",stack=[]):
         # Identify all present conditioning inputs
         tensor_keys = [k for k, v in V.items() if v is not None and isinstance(v, list) and len(v) > 0]
         if not tensor_keys:
@@ -90,7 +93,6 @@ def execute(cls, V, F, Expression, Expression_pi,batching, length_mismatch="tile
         # Extract tensors and pooled outputs
         tensors = {}
         pooled_outputs = {}
-        ss = dict()
         for key in tensor_keys:
              conditioning = V[key]
              tensors[key] = conditioning[0][0]
@@ -152,7 +154,7 @@ def execute(cls, V, F, Expression, Expression_pi,batching, length_mismatch="tile
 
         # Execute Expression (Main Tensor)
         tree = parse_expr(Expression)
-        visitor = UnifiedMathVisitor(variables, a.shape,a.device, state_storage=ss)
+        visitor = UnifiedMathVisitor(variables, a.shape,a.device, state_storage=stack)
         rtensor = visitor.visit(tree)
         rtensor = as_tensor(rtensor, a.shape)
 
@@ -193,7 +195,7 @@ def execute(cls, V, F, Expression, Expression_pi,batching, length_mismatch="tile
 
         # Execute Expression_pi (Pooled Output)
         tree_pi = parse_expr(Expression_pi)
-        visitor_pi = UnifiedMathVisitor(variables_pi, a_p.shape,a_p.device, state_storage=ss)
+        visitor_pi = UnifiedMathVisitor(variables_pi, a_p.shape,a_p.device, state_storage=stack)
         rpooled_raw = visitor_pi.visit(tree_pi)
         rpooled = as_tensor(rpooled_raw, a_p.shape)
 

more_math/FloatMathNode.py

@@ -9,6 +9,7 @@
 from .Parser.MathExprLexer import MathExprLexer
 from .Parser.MathExprParser import MathExprParser
 import re
+from .Stack import MrmthStack
 
 
 class FloatMathNode(io.ComfyNode):
@@ -29,16 +30,18 @@ def define_schema(cls) -> io.Schema:
             inputs=[
                 io.Autogrow.Input(id="V",template=io.Autogrow.TemplatePrefix(io.Float.Input("values"), prefix="V", min=1, max=50)),
                 io.String.Input(id="FloatFunc", default="a*(1-w)+b*w", tooltip="Expression to use on inputs"),
+                MrmthStack.Input(id="stack", tooltip="Access stack between nodes",optional=True)
             ],
             outputs=[
                 io.Float.Output(),
+                MrmthStack.Output(),
             ],
         )
 
     tooltip = cleandoc(__doc__)
 
     @classmethod
-    def check_lazy_status(cls, FloatFunc, V):
+    def check_lazy_status(cls, FloatFunc, V,stack=[]):
         input_stream = InputStream(FloatFunc)
         lexer = MathExprLexer(input_stream)
         stream = CommonTokenStream(lexer)
@@ -69,7 +72,7 @@ def check_lazy_status(cls, FloatFunc, V):
         return needed1
 
     @classmethod
-    def execute(cls, FloatFunc, V):
+    def execute(cls, FloatFunc, V,stack=[]):
 
         variables = {}
         # Populate aliases
@@ -95,9 +98,9 @@ def execute(cls, FloatFunc, V):
         tree = parse_expr(FloatFunc);
         # scalar execution
         # UnifiedMathVisitor expects variables and a shape. Shape [1] for scalar?
-        visitor = UnifiedMathVisitor(variables, [1])
+        visitor = UnifiedMathVisitor(variables, [1],state_storage=stack)
         result = visitor.visit(tree)
         # Result might be float or tensor(scalar)
         if torch.is_tensor(result):
              result = result[0].item()
-        return (float(result),)
+        return (float(result),stack)

more_math/GuiderMathNode.py

@@ -1,3 +1,4 @@
+from numpy import stack
 import torch
 import re
 from antlr4 import InputStream, CommonTokenStream
@@ -19,6 +20,7 @@
 import comfy.utils
 import comfy.hooks
 import comfy.samplers
+from .Stack import MrmthStack
 
 
 class GuiderMathNode(io.ComfyNode):
@@ -37,14 +39,16 @@ def define_schema(cls) -> io.Schema:
                 io.Autogrow.Input(id="F", template=io.Autogrow.TemplatePrefix(io.Float.Input("float", default=0.0, optional=True, lazy=True, force_input=True), prefix="F", min=1, max=50)),
                 io.String.Input(id="Expression", default="G0*(1-F0)+G1*F0", tooltip="Expression to apply on input guiders. Aliases: a=G0, b=G1, c=G2, d=G3, w=F0, x=F1, y=F2, z=F3. Context: steps, current_step"),
                 io.String.Input(id="Expression1", default="G0*(1-F0)+G1*F0", tooltip="Expression to apply after generation finishes."),
+                MrmthStack.Input(id="stack", tooltip="Access stack between nodes",optional=True)
             ],
             outputs=[
                 io.Guider.Output(),
+                MrmthStack.Output()
             ],
         )
 
     @classmethod
-    def check_lazy_status(cls, Expression,Expression1, V, F):
+    def check_lazy_status(cls, Expression,Expression1, V, F,stack=[]):
         input_stream = InputStream(Expression)
         input_stream1 = InputStream(Expression1)
         lexer = MathExprLexer(input_stream)
@@ -79,12 +83,12 @@ def check_lazy_status(cls, Expression,Expression1, V, F):
         return needed1
 
     @classmethod
-    def execute(cls, V, F, Expression,Expression1):
-        return (MathGuider(V, F, Expression,Expression1),)
+    def execute(cls, V, F, Expression,Expression1,stack=[]):
+        return (MathGuider(V, F, Expression,Expression1),stack)
 
 
 class MathGuider:
-    def __init__(self, V, F, expression,expression1):
+    def __init__(self, V, F, expression,expression1,stack=[]):
         self.V = V
         self.F = F
         self.expression = expression
@@ -94,7 +98,7 @@ def __init__(self, V, F, expression,expression1):
         self.sigmas = None
         self.current_step = 0
         self.steps = 0
-        self.stck = {}
+        self.stck = stack
 
     @property
     def model_patcher(self):
@@ -167,7 +171,7 @@ def setVars(self, x, sigma, seed, g_results):
                 "c": g_results.get("V2", make_zero_like(eval_samples)),
                 "d": g_results.get("V3", make_zero_like(eval_samples)),
             })
-            
+
             v_stacked, v_cnt = get_v_variable(g_results)
             if v_stacked is not None:
                  variables["V"] = v_stacked

more_math/ImageMathNode.py

@@ -5,6 +5,7 @@
 from .Parser.MathExprLexer import MathExprLexer
 from .Parser.MathExprParser import MathExprParser
 import re
+from .Stack import MrmthStack
 
 class ImageMathNode(io.ComfyNode):
     """
@@ -31,15 +32,17 @@ def define_schema(cls) -> io.Schema:
                     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 @@ def check_lazy_status(cls, Expression, V, F, length_mismatch="tile"):
         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 @@ def execute(cls, V, F, Expression, length_mismatch="error"):
             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)

more_math/LatentMathNode.py

@@ -17,6 +17,7 @@
 from .Parser.MathExprParser import MathExprParser
 import re
 from comfy.nested_tensor import NestedTensor
+from .Stack import MrmthStack
 
 class LatentMathNode(io.ComfyNode):
     """
@@ -43,17 +44,20 @@ def define_schema(cls) -> io.Schema:
                     default="error",
                     tooltip="How to handle mismatched latent batch sizes. tile: repeat shorter inputs; error: raise error on mismatch; pad: treat missing frames as zero."
                 ),
-                io.Int.Input(id="batching")
+                io.Int.Input(id="batching"),
+                MrmthStack.Input(id="stack", tooltip="Access stack between nodes",optional=True)
             ],
             outputs=[
                 io.Latent.Output(is_output_list=True),
+                MrmthStack.Output(),
+
             ],
         )
 
     tooltip = cleandoc(__doc__)
 
     @classmethod
-    def check_lazy_status(cls, Expression, V, F,batching, length_mismatch="tile"):
+    def check_lazy_status(cls, Expression, V, F,batching, length_mismatch="tile",stack=[]):
 
         input_stream = InputStream(Expression)
         lexer = MathExprLexer(input_stream)
@@ -85,7 +89,7 @@ def check_lazy_status(cls, Expression, V, F,batching, length_mismatch="tile"):
         return needed1
 
     @classmethod
-    def execute(cls, V, F, Expression,batching, length_mismatch="tile") -> io.NodeOutput:
+    def execute(cls, V, F, Expression,batching, length_mismatch="tile",stack=[]) -> io.NodeOutput:
         # Determine reference latent
         ref_latent = None
         for lat in V.values():
@@ -198,7 +202,7 @@ def execute(cls, V, F, Expression,batching, length_mismatch="tile") -> io.NodeOu
         for k, v in F.items():
             variables[k] = v if v is not None else 0.0
 
-        visitor = UnifiedMathVisitor(variables, ae.shape,ae.device)
+        visitor = UnifiedMathVisitor(variables, ae.shape,ae.device,state_storage=stack)
         result_t = as_tensor(visitor.visit(tree), ae.shape)
 
         result_latent = ref_latent.copy()
@@ -221,7 +225,7 @@ def execute(cls, V, F, Expression,batching, length_mismatch="tile") -> io.NodeOu
                 else:
                     rl["samples"] = result_t
                 results1.append(rl)
-            return (results1,)
+            return (results1,stack)
         rl = result_latent.copy()
         rl["samples"] = result_t
-        return ([rl],)
+        return ([rl],stack)