From d3bc62cbb903ce6ab66f5b03a531bd22558f9fd2 Mon Sep 17 00:00:00 2001 From: mcDandy Date: Sat, 21 Feb 2026 17:27:56 +0100 Subject: [PATCH] Add ability to use custom widget everywhere --- more_math/ClipMathNode.py | 41 ++++------------ more_math/ConditioningMathNode.py | 68 ++++++++++---------------- more_math/FloatMathNode.py | 47 +++++------------- more_math/GuiderMathNode.py | 62 +++++++++--------------- more_math/ImageMathNode.py | 49 +++++-------------- more_math/LatentMathNode.py | 50 ++++++------------- more_math/MaskMathNode.py | 49 +++++-------------- more_math/ModelMathNode.py | 42 ++++------------ more_math/NoiseMathNode.py | 46 +++++------------- more_math/SigmasMathNode.py | 49 +++++-------------- more_math/VaeMathNode.py | 43 ++++------------- more_math/VideoMathNode.py | 80 ++++++++++++------------------- more_math/modelLikeCommon.py | 8 +++- 13 files changed, 188 insertions(+), 446 deletions(-) diff --git a/more_math/ClipMathNode.py b/more_math/ClipMathNode.py index 3181d4d..d9a20a2 100644 --- a/more_math/ClipMathNode.py +++ b/more_math/ClipMathNode.py @@ -1,10 +1,8 @@ from comfy_api.latest import io from inspect import cleandoc -from antlr4 import InputStream, CommonTokenStream -from .Parser.MathExprLexer import MathExprLexer -from .Parser.MathExprParser import MathExprParser -import re +from .helper_functions import checkLazyNew from .Stack import MrmthStack +from .ParseTree import MrmthParseTree import copy import comfy.utils @@ -24,7 +22,11 @@ class CLIPMathNode(io.ComfyNode): inputs=[ io.Autogrow.Input(id="V",template=io.Autogrow.TemplatePrefix(io.Clip.Input("values"), prefix="V", min=1, max=50)), 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="I0*(1-F0)+I1*F0", tooltip="Expression to apply on weights"), + io.MultiType.Input( + io.String.Input("Expression", default="I0*(1-F0)+I1*F0", multiline=False), + types=[io.String,MrmthParseTree], + tooltip="Expression to apply on weights", + ), io.Combo.Input( id="length_mismatch", options=["do nothing","error","tile", "pad"], @@ -44,35 +46,8 @@ class CLIPMathNode(io.ComfyNode): @classmethod def check_lazy_status(cls, Expression, V, F, length_mismatch="tile",stack={}): + return checkLazyNew(Expression,V,F) - input_stream = InputStream(Expression) - lexer = MathExprLexer(input_stream) - stream = CommonTokenStream(lexer) - stream.fill() - - # Support aliases - aliases_img = {"a": "V0", "b": "V1", "c": "V2", "d": "V3"} - aliases_flt = {"w": "F0", "x": "F1", "y": "F2", "z": "F3"} - - needed = [] - needed1 = [] - for token in filter(lambda t: t.type == MathExprParser.VARIABLE, stream.tokens): - var_name = token.text - - if re.match(r"[VF][0-9]+", var_name): - needed.append(var_name) - elif var_name in aliases_img: - needed.append(aliases_img[var_name]) - elif var_name in aliases_flt: - needed.append(aliases_flt[var_name]) - for v in needed: - if v.startswith("V"): - if v not in V or V[v] is None: - needed1.append(v) - elif v.startswith("F"): - if v not in F or F[v] is None: - needed1.append(v) - return needed1 @classmethod def execute(cls, V, F, Expression, length_mismatch="tile",stack={}) -> io.NodeOutput: diff --git a/more_math/ConditioningMathNode.py b/more_math/ConditioningMathNode.py index 1271abe..5f84249 100644 --- a/more_math/ConditioningMathNode.py +++ b/more_math/ConditioningMathNode.py @@ -1,13 +1,11 @@ +from tkinter import E import torch -from .helper_functions import generate_dim_variables, parse_expr, getIndexTensorAlongDim, as_tensor, normalize_to_common_shape, make_zero_like, get_v_variable, get_f_variable +from .helper_functions import checkLazyNew, generate_dim_variables, parse_expr, getIndexTensorAlongDim, as_tensor, normalize_to_common_shape, make_zero_like, get_v_variable, get_f_variable from .Parser.UnifiedMathVisitor import UnifiedMathVisitor from comfy_api.latest import io -from antlr4 import InputStream, CommonTokenStream -from .Parser.MathExprLexer import MathExprLexer -from .Parser.MathExprParser import MathExprParser -import re import copy from .Stack import MrmthStack +from .ParseTree import MrmthParseTree class ConditioningMathNode(io.ComfyNode): """ @@ -28,8 +26,16 @@ class ConditioningMathNode(io.ComfyNode): inputs=[ io.Autogrow.Input(id="V",template=io.Autogrow.TemplatePrefix(io.Conditioning.Input("values"), prefix="V", min=1, max=50)), 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",display_name="Tensor expr.", default="I0*(1-F0)+I1*F0", tooltip="Expression to apply on tensor part of conditioning"), - io.String.Input(id="Expression_pi",display_name="pooled output expr.", default="I0*(1-F0)+I1*F0", tooltip="Expression to apply on pooled_input part of conditioning"), + io.MultiType.Input( + io.String.Input("Expression", display_name="Tensor expr.", default="I0*(1-F0)+I1*F0", multiline=False), + types=[io.String,MrmthParseTree], + tooltip="Expression to apply on tensor part of conditioning", + ), + io.MultiType.Input( + io.String.Input("Expression_pi", display_name="pooled output expr.", default="I0*(1-F0)+I1*F0", multiline=False), + types=[io.String,MrmthParseTree], + tooltip="Expression to apply on pooled_input part of conditioning", + ), io.Combo.Input( id="length_mismatch", options=["do nothing","error","tile", "pad"], @@ -48,40 +54,10 @@ class ConditioningMathNode(io.ComfyNode): @classmethod def check_lazy_status(cls, Expression,Expression_pi, V, F,batching, length_mismatch="tile",stack={}): + d = checkLazyNew(Expression,V,F) + b = checkLazyNew(Expression_pi,V,F) + return d|b - input_stream = InputStream(Expression) - lexer = MathExprLexer(input_stream) - stream = CommonTokenStream(lexer) - stream.fill() - - input_stream = InputStream(Expression_pi) - lexer = MathExprLexer(input_stream) - stream1 = CommonTokenStream(lexer) - stream1.fill() - - # Support aliases - aliases_img = {"a": "V0", "b": "V1", "c": "V2", "d": "V3"} - aliases_flt = {"w": "F0", "x": "F1", "y": "F2", "z": "F3"} - - needed = set() - needed1 = set() - for token in filter(lambda t: t.type == MathExprParser.VARIABLE, stream.tokens + stream1.tokens): - var_name = token.text - - if re.match(r"[VF][0-9]+", var_name): - needed.add(var_name) - elif var_name in aliases_img: - needed.add(aliases_img[var_name]) - elif var_name in aliases_flt: - needed.add(aliases_flt[var_name]) - for v in needed: - if v.startswith("V"): - if v not in V or V[v] is None: - needed1.add(v) - elif v.startswith("F"): - if v not in F or F[v] is None: - needed1.add(v) - return needed1 @classmethod def execute(cls, V, F, Expression, Expression_pi,batching, length_mismatch="tile",stack={}): @@ -153,7 +129,11 @@ class ConditioningMathNode(io.ComfyNode): variables[k] = val if val is not None else 0.0 # Execute Expression (Main Tensor) - tree = parse_expr(Expression) + tree = None + if isinstance(Expression,str): + tree = parse_expr(Expression) + else: + tree = Expression visitor = UnifiedMathVisitor(variables, a.shape,a.device, state_storage=stack) rtensor = visitor.visit(tree) rtensor = as_tensor(rtensor, a.shape) @@ -194,7 +174,11 @@ class ConditioningMathNode(io.ComfyNode): variables_pi[k] = val if val is not None else 0.0 # Execute Expression_pi (Pooled Output) - tree_pi = parse_expr(Expression_pi) + tree_pi = None + if isinstance(Expression_pi,str): + tree_pi = parse_expr(Expression_pi) + else: + tree_pi = Expression_pi 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) diff --git a/more_math/FloatMathNode.py b/more_math/FloatMathNode.py index 579c7fd..a55634a 100644 --- a/more_math/FloatMathNode.py +++ b/more_math/FloatMathNode.py @@ -1,15 +1,12 @@ from inspect import cleandoc import torch -from .helper_functions import parse_expr, get_v_variable +from .helper_functions import parse_expr, get_v_variable, checkLazyNew from .Parser.UnifiedMathVisitor import UnifiedMathVisitor from comfy_api.latest import io -from antlr4 import InputStream, CommonTokenStream -from .Parser.MathExprLexer import MathExprLexer -from .Parser.MathExprParser import MathExprParser -import re from .Stack import MrmthStack +from .ParseTree import MrmthParseTree import copy @@ -31,7 +28,11 @@ class FloatMathNode(io.ComfyNode): display_name="Float math", 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"), + io.MultiType.Input( + io.String.Input("FloatFunc", default="a*(1-w)+b*w", multiline=False), + types=[io.String,MrmthParseTree], + tooltip="Expression to use on inputs", + ), MrmthStack.Input(id="stack", tooltip="Access stack between nodes",optional=True) ], outputs=[ @@ -44,34 +45,8 @@ class FloatMathNode(io.ComfyNode): @classmethod def check_lazy_status(cls, FloatFunc, V,stack={}): - input_stream = InputStream(FloatFunc) - lexer = MathExprLexer(input_stream) - stream = CommonTokenStream(lexer) - stream.fill() + return checkLazyNew(FloatFunc,V,V) - # Support aliases - # Legacy FloatMathNode mapped a,b,c,d,w,x,y,z to V0-V7 roughly? - # Actually Step 37 showed explicit mapping: - # a->V0, b->V1, c->V2, d->V3, w->V4, x->V5, y->V6, z->V7 - aliases = { - "a": "V0", "b": "V1", "c": "V2", "d": "V3", - "w": "V4", "x": "V5", "y": "V6", "z": "V7" - } - - needed = [] - needed1 = [] - for token in filter(lambda t: t.type == MathExprParser.VARIABLE, stream.tokens): - var_name = token.text - - if re.match(r"V[0-9]+", var_name): - needed.append(var_name) - elif var_name in aliases: - needed.append(aliases[var_name]) - - for v in needed: - if v not in V or V[v] is None: - needed1.append(v) - return needed1 @classmethod def execute(cls, FloatFunc, V,stack={}): @@ -98,7 +73,11 @@ class FloatMathNode(io.ComfyNode): variables["Vcnt"] = float(v_cnt) variables["V_count"] = float(v_cnt) - tree = parse_expr(FloatFunc); + tree = None + if isinstance(FloatFunc,str): + tree = parse_expr(FloatFunc) + else: + tree = FloatFunc # scalar execution visitor = UnifiedMathVisitor(variables, [1],state_storage=stack) result = visitor.visit(tree) diff --git a/more_math/GuiderMathNode.py b/more_math/GuiderMathNode.py index 08941ea..7ad7371 100644 --- a/more_math/GuiderMathNode.py +++ b/more_math/GuiderMathNode.py @@ -1,10 +1,7 @@ import torch -import re -from antlr4 import InputStream, CommonTokenStream -from .Parser.MathExprLexer import MathExprLexer -from .Parser.MathExprParser import MathExprParser from .Parser.UnifiedMathVisitor import UnifiedMathVisitor from .helper_functions import ( + checkLazyNew, generate_dim_variables, getIndexTensorAlongDim, parse_expr, @@ -20,6 +17,7 @@ import comfy.utils import comfy.hooks import comfy.samplers from .Stack import MrmthStack +from .ParseTree import MrmthParseTree import copy @@ -37,8 +35,16 @@ class GuiderMathNode(io.ComfyNode): inputs=[ io.Autogrow.Input(id="V", template=io.Autogrow.TemplatePrefix(io.Guider.Input("values"), prefix="V", min=1, max=50)), 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."), + io.MultiType.Input( + io.String.Input("Expression", default="G0*(1-F0)+G1*F0", multiline=False), + types=[io.String,MrmthParseTree], + 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.MultiType.Input( + io.String.Input("Expression1", default="G0*(1-F0)+G1*F0", multiline=False), + types=[io.String,MrmthParseTree], + tooltip="Expression to apply after generation finishes.", + ), MrmthStack.Input(id="stack", tooltip="Access stack between nodes",optional=True) ], outputs=[ @@ -49,38 +55,10 @@ class GuiderMathNode(io.ComfyNode): @classmethod def check_lazy_status(cls, Expression,Expression1, V, F,stack={}): - input_stream = InputStream(Expression) - input_stream1 = InputStream(Expression1) - lexer = MathExprLexer(input_stream) - lexer1 = MathExprLexer(input_stream1) - stream = CommonTokenStream(lexer) - stream1 = CommonTokenStream(lexer1) - stream.fill() - stream1.fill() + d = checkLazyNew(Expression,V,F) + b = checkLazyNew(Expression1,V,F) + return d|b - # Support aliases - aliases_smp = {"a": "V0", "b": "V1", "c": "V2", "d": "V3"} - aliases_flt = {"w": "F0", "x": "F1", "y": "F2", "z": "F3"} - - needed = [] - needed1 = [] - for token in filter(lambda t: t.type == MathExprParser.VARIABLE, stream.tokens+stream1.tokens): - var_name = token.text - if re.match(r"[VF][0-9]+", var_name): - needed.append(var_name) - elif var_name in aliases_smp: - needed.append(aliases_smp[var_name]) - elif var_name in aliases_flt: - needed.append(aliases_flt[var_name]) - - for v in needed: - if v.startswith("V"): - if v not in V or V[v] is None: - needed1.append(v) - elif v.startswith("F"): - if v not in F or F[v] is None: - needed1.append(v) - return needed1 @classmethod def execute(cls, V, F, Expression,Expression1,stack={}): @@ -93,8 +71,14 @@ class MathGuider: self.V = V self.F = F self.expression = expression - self.tree = parse_expr(expression) - self.tree1 = parse_expr(expression1) + if isinstance(expression,str): + self.tree = parse_expr(expression) + else: + self.tree = expression + if isinstance(expression1,str): + self.tree1 = parse_expr(expression1) + else: + self.tree1 = expression1 self.inner_model = None # Will be set during sample self.sigmas = None self.current_step = 0 diff --git a/more_math/ImageMathNode.py b/more_math/ImageMathNode.py index ce7d7dc..5361e4f 100644 --- a/more_math/ImageMathNode.py +++ b/more_math/ImageMathNode.py @@ -1,12 +1,9 @@ -from .helper_functions import generate_dim_variables, parse_expr, getIndexTensorAlongDim, as_tensor, normalize_to_common_shape, make_zero_like, get_v_variable, get_f_variable +from .helper_functions import generate_dim_variables, parse_expr, getIndexTensorAlongDim, as_tensor, normalize_to_common_shape, make_zero_like, get_v_variable, get_f_variable, checkLazyNew from .Parser.UnifiedMathVisitor import UnifiedMathVisitor from comfy_api.latest import io -from antlr4 import InputStream, CommonTokenStream -from .Parser.MathExprLexer import MathExprLexer -from .Parser.MathExprParser import MathExprParser -import re import torch from .Stack import MrmthStack +from .ParseTree import MrmthParseTree import copy @@ -29,7 +26,11 @@ class ImageMathNode(io.ComfyNode): inputs=[ io.Autogrow.Input(id="V",template=io.Autogrow.TemplatePrefix(io.Image.Input("values"), prefix="V", min=1, max=50)), 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="I0*(1-F0)+I1*F0", tooltip="Expression to apply on input images"), + io.MultiType.Input( + io.String.Input("Expression", default="I0*(1-F0)+I1*F0", multiline=False), + types=[io.String,MrmthParseTree], + tooltip="Expression to apply on input images", + ), io.Combo.Input( id="length_mismatch", options=["do nothing","error","tile", "pad"], @@ -48,35 +49,7 @@ class ImageMathNode(io.ComfyNode): @classmethod def check_lazy_status(cls, Expression, V, F, length_mismatch="tile",batching=0,stack={}): - - input_stream = InputStream(Expression) - lexer = MathExprLexer(input_stream) - stream = CommonTokenStream(lexer) - stream.fill() - - # Support aliases - aliases_img = {"a": "V0", "b": "V1", "c": "V2", "d": "V3"} - aliases_flt = {"w": "F0", "x": "F1", "y": "F2", "z": "F3"} - - needed = [] - needed1 = [] - for token in filter(lambda t: t.type == MathExprParser.VARIABLE, stream.tokens): - var_name = token.text - - if re.match(r"[VF][0-9]+", var_name): - needed.append(var_name) - elif var_name in aliases_img: - needed.append(aliases_img[var_name]) - elif var_name in aliases_flt: - needed.append(aliases_flt[var_name]) - for v in needed: - if v.startswith("V"): - if v not in V or V[v] is None: - needed1.append(v) - elif v.startswith("F"): - if v not in F or F[v] is None: - needed1.append(v) - return needed1 + return checkLazyNew(Expression,V,F) @classmethod def execute(cls, V, F, Expression, length_mismatch="error",batching=0,stack={}): @@ -152,7 +125,11 @@ class ImageMathNode(io.ComfyNode): for k, val in F.items(): variables[k] = val if val is not None else 0.0 - tree = parse_expr(Expression); + tree = None + if isinstance(Expression,str): + tree = parse_expr(Expression) + else: + tree = Expression visitor = UnifiedMathVisitor(variables, ae.shape,ae.device,state_storage=stack) result = visitor.visit(tree) result = as_tensor(result, ae.shape) diff --git a/more_math/LatentMathNode.py b/more_math/LatentMathNode.py index 1d2f889..3bcba23 100644 --- a/more_math/LatentMathNode.py +++ b/more_math/LatentMathNode.py @@ -8,16 +8,14 @@ from .helper_functions import ( normalize_to_common_shape, make_zero_like, get_v_variable, - get_f_variable + get_f_variable, + checkLazyNew ) from .Parser.UnifiedMathVisitor import UnifiedMathVisitor import torch -from antlr4 import InputStream, CommonTokenStream -from .Parser.MathExprLexer import MathExprLexer -from .Parser.MathExprParser import MathExprParser -import re from comfy.nested_tensor import NestedTensor from .Stack import MrmthStack +from .ParseTree import MrmthParseTree import copy class LatentMathNode(io.ComfyNode): @@ -38,7 +36,11 @@ class LatentMathNode(io.ComfyNode): inputs=[ io.Autogrow.Input(id="V",template=io.Autogrow.TemplatePrefix(io.Latent.Input("values"), prefix="V", min=1, max=50)), 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="I0*(1-F0)+I1*F0", tooltip="Expression to apply on input latents"), + io.MultiType.Input( + io.String.Input("Expression", default="I0*(1-F0)+I1*F0", multiline=False), + types=[io.String,MrmthParseTree], + tooltip="Expression to apply on input latents", + ), io.Combo.Input( id="length_mismatch", options=["do nothing","error","tile", "pad"], @@ -60,35 +62,7 @@ class LatentMathNode(io.ComfyNode): @classmethod def check_lazy_status(cls, Expression, V, F,batching, length_mismatch="tile",stack={}): - - input_stream = InputStream(Expression) - lexer = MathExprLexer(input_stream) - stream = CommonTokenStream(lexer) - stream.fill() - - # Support aliases - aliases_img = {"a": "V0", "b": "V1", "c": "V2", "d": "V3"} - aliases_flt = {"w": "F0", "x": "F1", "y": "F2", "z": "F3"} - - needed = [] - needed1 = [] - for token in filter(lambda t: t.type == MathExprParser.VARIABLE, stream.tokens): - var_name = token.text - - if re.match(r"[VF][0-9]+", var_name): - needed.append(var_name) - elif var_name in aliases_img: - needed.append(aliases_img[var_name]) - elif var_name in aliases_flt: - needed.append(aliases_flt[var_name]) - for v in needed: - if v.startswith("V"): - if v not in V or V[v] is None: - needed1.append(v) - elif v.startswith("F"): - if v not in F or F[v] is None: - needed1.append(v) - return needed1 + return checkLazyNew(Expression,V,F) @classmethod def execute(cls, V, F, Expression,batching, length_mismatch="tile",stack={}) -> io.NodeOutput: @@ -148,7 +122,11 @@ class LatentMathNode(io.ComfyNode): raise ValueError(f"Input '{name}' has shape {V[name]['samples'].shape[0]}, expected {ae.shape[0]} to match input.") # parse expression once - tree = parse_expr(Expression) + tree = None + if isinstance(Expression,str): + tree = parse_expr(Expression) + else: + tree = Expression ndim = ae.ndim batch_dim = 0 diff --git a/more_math/MaskMathNode.py b/more_math/MaskMathNode.py index cfc992b..7584e79 100644 --- a/more_math/MaskMathNode.py +++ b/more_math/MaskMathNode.py @@ -1,12 +1,9 @@ -from .helper_functions import generate_dim_variables,parse_expr, getIndexTensorAlongDim, as_tensor, normalize_to_common_shape,make_zero_like, get_v_variable, get_f_variable +from .helper_functions import generate_dim_variables,parse_expr, getIndexTensorAlongDim, as_tensor, normalize_to_common_shape,make_zero_like, get_v_variable, get_f_variable, checkLazyNew from .Parser.UnifiedMathVisitor import UnifiedMathVisitor from comfy_api.latest import io -from antlr4 import InputStream, CommonTokenStream -from .Parser.MathExprLexer import MathExprLexer -from .Parser.MathExprParser import MathExprParser -import re import torch from .Stack import MrmthStack +from .ParseTree import MrmthParseTree import copy @@ -29,7 +26,11 @@ class MaskMathNode(io.ComfyNode): inputs=[ io.Autogrow.Input(id="V",template=io.Autogrow.TemplatePrefix(io.Mask.Input("values"), prefix="V", min=1, max=50)), 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="I0*(1-F0)+I1*F0", tooltip="Expression to apply on input masks"), + io.MultiType.Input( + io.String.Input("Expression", default="I0*(1-F0)+I1*F0", multiline=False), + types=[io.String,MrmthParseTree], + tooltip="Expression to apply on input masks", + ), io.Combo.Input( id="length_mismatch", options=["do nothing","error","tile", "pad"], @@ -48,35 +49,7 @@ class MaskMathNode(io.ComfyNode): @classmethod def check_lazy_status(cls, Expression, V, F, length_mismatch="tile",batching=0,stack={}): - - input_stream = InputStream(Expression) - lexer = MathExprLexer(input_stream) - stream = CommonTokenStream(lexer) - stream.fill() - - # Support aliases - aliases_img = {"a": "V0", "b": "V1", "c": "V2", "d": "V3"} - aliases_flt = {"w": "F0", "x": "F1", "y": "F2", "z": "F3"} - - needed = [] - needed1 = [] - for token in filter(lambda t: t.type == MathExprParser.VARIABLE, stream.tokens): - var_name = token.text - - if re.match(r"[VF][0-9]+", var_name): - needed.append(var_name) - elif var_name in aliases_img: - needed.append(aliases_img[var_name]) - elif var_name in aliases_flt: - needed.append(aliases_flt[var_name]) - for v in needed: - if v.startswith("V"): - if v not in V or V[v] is None: - needed1.append(v) - elif v.startswith("F"): - if v not in F or F[v] is None: - needed1.append(v) - return needed1 + return checkLazyNew(Expression,V,F) @classmethod def execute(cls, V, F, Expression, length_mismatch="tile",batching=0,stack={}): @@ -145,7 +118,11 @@ class MaskMathNode(io.ComfyNode): for k, val in F.items(): variables[k] = val if val is not None else 0.0 - tree = parse_expr(Expression); + tree = None + if isinstance(Expression,str): + tree = parse_expr(Expression) + else: + tree = Expression visitor = UnifiedMathVisitor(variables, ae.shape,ae.device,state_storage=stack) result = visitor.visit(tree) result = as_tensor(result, ae.shape) diff --git a/more_math/ModelMathNode.py b/more_math/ModelMathNode.py index 3e42a70..09af282 100644 --- a/more_math/ModelMathNode.py +++ b/more_math/ModelMathNode.py @@ -1,10 +1,8 @@ from inspect import cleandoc from comfy_api.latest import io -from antlr4 import InputStream, CommonTokenStream -from .Parser.MathExprLexer import MathExprLexer -from .Parser.MathExprParser import MathExprParser -import re +from .helper_functions import checkLazyNew from .Stack import MrmthStack +from .ParseTree import MrmthParseTree import copy import comfy.utils from .modelLikeCommon import calculate_patches_autogrow @@ -25,7 +23,11 @@ class ModelMathNode(io.ComfyNode): inputs=[ io.Autogrow.Input(id="V",template=io.Autogrow.TemplatePrefix(io.Model.Input("values"), prefix="V", min=1, max=50)), 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="I0*(1-F0)+I1*F0", tooltip="Expression to apply on weights"), + io.MultiType.Input( + io.String.Input("Expression", default="I0*(1-F0)+I1*F0", multiline=False), + types=[io.String,MrmthParseTree], + tooltip="Expression to apply on weights", + ), io.Combo.Input( id="length_mismatch", options=["do nothing","error","tile", "pad"], @@ -46,35 +48,7 @@ class ModelMathNode(io.ComfyNode): @classmethod def check_lazy_status(cls, Expression, V, F, length_mismatch="tile",stack={}): - - input_stream = InputStream(Expression) - lexer = MathExprLexer(input_stream) - stream = CommonTokenStream(lexer) - stream.fill() - - # Support aliases - aliases_img = {"a": "V0", "b": "V1", "c": "V2", "d": "V3"} - aliases_flt = {"w": "F0", "x": "F1", "y": "F2", "z": "F3"} - - needed = [] - needed1 = [] - for token in filter(lambda t: t.type == MathExprParser.VARIABLE, stream.tokens): - var_name = token.text - - if re.match(r"[VF][0-9]+", var_name): - needed.append(var_name) - elif var_name in aliases_img: - needed.append(aliases_img[var_name]) - elif var_name in aliases_flt: - needed.append(aliases_flt[var_name]) - for v in needed: - if v.startswith("V"): - if v not in V or V[v] is None: - needed1.append(v) - elif v.startswith("F"): - if v not in F or F[v] is None: - needed1.append(v) - return needed1 + return checkLazyNew(Expression,V,F) @classmethod def execute(cls, V, F, Expression, length_mismatch="tile",stack={}) -> io.NodeOutput: diff --git a/more_math/NoiseMathNode.py b/more_math/NoiseMathNode.py index f61ce37..f724236 100644 --- a/more_math/NoiseMathNode.py +++ b/more_math/NoiseMathNode.py @@ -1,11 +1,9 @@ -from .helper_functions import generate_dim_variables, as_tensor, parse_expr, getIndexTensorAlongDim, make_zero_like, get_v_variable, get_f_variable +from .helper_functions import generate_dim_variables, as_tensor, parse_expr, getIndexTensorAlongDim, make_zero_like, get_v_variable, get_f_variable, checkLazyNew from comfy_api.latest import io import torch -from .Parser.MathExprParser import MathExprParser,InputStream,CommonTokenStream -from .Parser.MathExprLexer import MathExprLexer -import re from .Parser.UnifiedMathVisitor import UnifiedMathVisitor from .Stack import MrmthStack +from .ParseTree import MrmthParseTree import copy class NoiseMathNode(io.ComfyNode): @@ -33,7 +31,11 @@ class NoiseMathNode(io.ComfyNode): inputs=[ io.Autogrow.Input(id="V",template=io.Autogrow.TemplatePrefix(io.Noise.Input("values"), prefix="V", min=1, max=50)), 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="Noise", default="a*(1-w)+b*w"), + io.MultiType.Input( + io.String.Input("Noise", default="a*(1-w)+b*w", multiline=False), + types=[io.String,MrmthParseTree], + tooltip="Expression for noise", + ), MrmthStack.Input(id="stack", tooltip="Access stack between nodes",optional=True) ], outputs=[ @@ -44,34 +46,7 @@ class NoiseMathNode(io.ComfyNode): @classmethod def check_lazy_status(cls, Noise, V, F,stack={}): - input_stream = InputStream(Noise) - lexer = MathExprLexer(input_stream) - stream = CommonTokenStream(lexer) - stream.fill() - - # Support aliases - aliases_smp = {"a": "V0", "b": "V1", "c": "V2", "d": "V3"} - aliases_flt = {"w": "F0", "x": "F1", "y": "F2", "z": "F3"} - - needed = [] - needed1 = [] - for token in filter(lambda t: t.type == MathExprParser.VARIABLE, stream.tokens): - var_name = token.text - if re.match(r"[VF][0-9]+", var_name): - needed.append(var_name) - elif var_name in aliases_smp: - needed.append(aliases_smp[var_name]) - elif var_name in aliases_flt: - needed.append(aliases_flt[var_name]) - - for v in needed: - if v.startswith("V"): - if v not in V or V[v] is None: - needed1.append(v) - elif v.startswith("F"): - if v not in F or F[v] is None: - needed1.append(v) - return needed1 + return checkLazyNew(Noise,V,F) @classmethod def execute(cls, Noise, V,F,stack={}): @@ -83,7 +58,10 @@ class NoiseExecutor: def __init__(self, V,F, expr,stack): self.V = V self.F = F - self.tree = parse_expr(expr) + if isinstance(expr,str): + self.tree = parse_expr(expr) + else: + self.tree = expr self.stack = stack seed = -1 diff --git a/more_math/SigmasMathNode.py b/more_math/SigmasMathNode.py index cd72a94..1cf983a 100644 --- a/more_math/SigmasMathNode.py +++ b/more_math/SigmasMathNode.py @@ -1,11 +1,8 @@ -from .helper_functions import generate_dim_variables, parse_expr, getIndexTensorAlongDim, as_tensor, prepare_inputs, make_zero_like, normalize_to_common_shape +from .helper_functions import generate_dim_variables, parse_expr, getIndexTensorAlongDim, as_tensor, prepare_inputs, make_zero_like, normalize_to_common_shape, checkLazyNew from .Parser.UnifiedMathVisitor import UnifiedMathVisitor from comfy_api.latest import io -from antlr4 import InputStream, CommonTokenStream -from .Parser.MathExprLexer import MathExprLexer -from .Parser.MathExprParser import MathExprParser -import re from .Stack import MrmthStack +from .ParseTree import MrmthParseTree import copy class SigmasMathNode(io.ComfyNode): @@ -27,7 +24,11 @@ class SigmasMathNode(io.ComfyNode): inputs=[ io.Autogrow.Input(id="V",template=io.Autogrow.TemplatePrefix(io.Sigmas.Input("values"), prefix="V", min=1, max=50)), 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="I0*(1-F0)+I1*F0", tooltip="Expression to apply on input images"), + io.MultiType.Input( + io.String.Input("Expression", default="I0*(1-F0)+I1*F0", multiline=False), + types=[io.String,MrmthParseTree], + tooltip="Expression to apply on input images", + ), io.Combo.Input( id="length_mismatch", options=["do nothing","error","tile", "pad"], @@ -45,35 +46,7 @@ class SigmasMathNode(io.ComfyNode): @classmethod def check_lazy_status(cls, Expression, V, F, length_mismatch="tile",stack={}): - - input_stream = InputStream(Expression) - lexer = MathExprLexer(input_stream) - stream = CommonTokenStream(lexer) - stream.fill() - - # Support aliases - aliases_img = {"a": "V0", "b": "V1", "c": "V2", "d": "V3"} - aliases_flt = {"w": "F0", "x": "F1", "y": "F2", "z": "F3"} - - needed = [] - needed1 = [] - for token in filter(lambda t: t.type == MathExprParser.VARIABLE, stream.tokens): - var_name = token.text - - if re.match(r"[VF][0-9]+", var_name): - needed.append(var_name) - elif var_name in aliases_img: - needed.append(aliases_img[var_name]) - elif var_name in aliases_flt: - needed.append(aliases_flt[var_name]) - for v in needed: - if v.startswith("V"): - if v not in V or V[v] is None: - needed1.append(v) - elif v.startswith("F"): - if v not in F or F[v] is None: - needed1.append(v) - return needed1 + return checkLazyNew(Expression,V,F) @classmethod def execute(cls, V, F, Expression, length_mismatch="tile",stack={}): @@ -129,7 +102,11 @@ class SigmasMathNode(io.ComfyNode): for k, v in F.items(): variables[k] = v if v is not None else 0.0 - tree = parse_expr(Expression); + tree = None + if isinstance(Expression,str): + tree = parse_expr(Expression) + else: + tree = Expression visitor = UnifiedMathVisitor(variables, ae.shape,ae.device,state_storage=stack) result = visitor.visit(tree) result = as_tensor(result, ae.shape) diff --git a/more_math/VaeMathNode.py b/more_math/VaeMathNode.py index 41e3ee6..e17fa5f 100644 --- a/more_math/VaeMathNode.py +++ b/more_math/VaeMathNode.py @@ -1,12 +1,9 @@ from inspect import cleandoc from comfy_api.latest import io import copy -from antlr4 import InputStream, CommonTokenStream - -from .Parser.MathExprLexer import MathExprLexer -from .Parser.MathExprParser import MathExprParser -import re +from .helper_functions import checkLazyNew from .Stack import MrmthStack +from .ParseTree import MrmthParseTree import comfy.utils class VAEMathNode(io.ComfyNode): @@ -23,7 +20,11 @@ class VAEMathNode(io.ComfyNode): inputs=[ io.Autogrow.Input(id="V",template=io.Autogrow.TemplatePrefix(io.Vae.Input("values"), prefix="V", min=1, max=50)), 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="I0*(1-F0)+I1*F0", tooltip="Expression to apply on weights"), + io.MultiType.Input( + io.String.Input("Expression", default="I0*(1-F0)+I1*F0", multiline=False), + types=[io.String,MrmthParseTree], + tooltip="Expression to apply on weights", + ), io.Combo.Input( id="length_mismatch", options=["do nothing","error","tile", "pad"], @@ -43,35 +44,7 @@ class VAEMathNode(io.ComfyNode): @classmethod def check_lazy_status(cls, Expression, V, F, length_mismatch="tile",stack={}): - - input_stream = InputStream(Expression) - lexer = MathExprLexer(input_stream) - stream = CommonTokenStream(lexer) - stream.fill() - - # Support aliases - aliases_img = {"a": "V0", "b": "V1", "c": "V2", "d": "V3"} - aliases_flt = {"w": "F0", "x": "F1", "y": "F2", "z": "F3"} - - needed = [] - needed1 = [] - for token in filter(lambda t: t.type == MathExprParser.VARIABLE, stream.tokens): - var_name = token.text - - if re.match(r"[VF][0-9]+", var_name): - needed.append(var_name) - elif var_name in aliases_img: - needed.append(aliases_img[var_name]) - elif var_name in aliases_flt: - needed.append(aliases_flt[var_name]) - for v in needed: - if v.startswith("V"): - if v not in V or V[v] is None: - needed1.append(v) - elif v.startswith("F"): - if v not in F or F[v] is None: - needed1.append(v) - return needed1 + return checkLazyNew(Expression,V,F) @classmethod def execute(cls, V, F, Expression, length_mismatch="tile",stack={}) -> io.NodeOutput: diff --git a/more_math/VideoMathNode.py b/more_math/VideoMathNode.py index 6a8b987..63cd208 100644 --- a/more_math/VideoMathNode.py +++ b/more_math/VideoMathNode.py @@ -1,11 +1,8 @@ -from .helper_functions import generate_dim_variables, parse_expr, getIndexTensorAlongDim, as_tensor, normalize_to_common_shape, make_zero_like, get_v_variable, get_f_variable +from .helper_functions import generate_dim_variables, parse_expr, getIndexTensorAlongDim, as_tensor, normalize_to_common_shape, make_zero_like, get_v_variable, get_f_variable, checkLazyNew from .Parser.UnifiedMathVisitor import UnifiedMathVisitor from comfy_api.latest import io -from antlr4 import InputStream, CommonTokenStream -from .Parser.MathExprLexer import MathExprLexer -from .Parser.MathExprParser import MathExprParser -import re from .Stack import MrmthStack +from .ParseTree import MrmthParseTree import copy class VideoMathNode(io.ComfyNode): @@ -27,8 +24,16 @@ class VideoMathNode(io.ComfyNode): inputs=[ io.Autogrow.Input(id="V",template=io.Autogrow.TemplatePrefix(io.Video.Input("values"), prefix="V", min=1, max=50)), 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="I0*(1-F0)+I1*F0", tooltip="Expression to apply on tensor part of conditioning"), - io.String.Input(id="Expression_pi", default="I0*(1-F0)+I1*F0", tooltip="Expression to apply on pooled_input part of conditioning"), + io.MultiType.Input( + io.String.Input("Expression", default="I0*(1-F0)+I1*F0", multiline=False), + types=[io.String,MrmthParseTree], + tooltip="Expression to apply on tensor part of conditioning", + ), + io.MultiType.Input( + io.String.Input("Expression_pi", default="I0*(1-F0)+I1*F0", multiline=False), + types=[io.String,MrmthParseTree], + tooltip="Expression to apply on pooled_input part of conditioning", + ), io.Combo.Input( id="length_mismatch", options=["do nothing","error","tile", "pad"], @@ -46,40 +51,9 @@ class VideoMathNode(io.ComfyNode): @classmethod def check_lazy_status(cls, Expression,Expression_pi, V, F, length_mismatch="tile",stack={}): - - input_stream = InputStream(Expression) - lexer = MathExprLexer(input_stream) - stream = CommonTokenStream(lexer) - stream.fill() - - input_stream = InputStream(Expression_pi) - lexer = MathExprLexer(input_stream) - stream1 = CommonTokenStream(lexer) - stream1.fill() - - # Support aliases - aliases_img = {"a": "V0", "b": "V1", "c": "V2", "d": "V3"} - aliases_flt = {"w": "F0", "x": "F1", "y": "F2", "z": "F3"} - - needed = set() - needed1 = set() - for token in filter(lambda t: t.type == MathExprParser.VARIABLE, stream.tokens + stream1.tokens): - var_name = token.text - - if re.match(r"[VF][0-9]+", var_name): - needed.add(var_name) - elif var_name in aliases_img: - needed.add(aliases_img[var_name]) - elif var_name in aliases_flt: - needed.add(aliases_flt[var_name]) - for v in needed: - if v.startswith("V"): - if v not in V or V[v] is None: - needed1.add(v) - elif v.startswith("F"): - if v not in F or F[v] is None: - needed1.add(v) - return needed1 + d = checkLazyNew(Expression,V,F) + b = checkLazyNew(Expression_pi,V,F) + return d|b @classmethod def execute(cls, V, F, Expression, Expression_pi, length_mismatch="tile",stack={}): @@ -152,10 +126,14 @@ class VideoMathNode(io.ComfyNode): for k, val in F.items(): variables[k] = val if val is not None else 0.0 - tree = parse_expr(Expression); + tree = None + if isinstance(Expression,str): + tree = parse_expr(Expression) + else: + tree = Expression visitor = UnifiedMathVisitor(variables, ae.shape,ae.device,state_storage=stack) - result = visitor.visit(tree) - result = as_tensor(result, ae.shape) + result_video = visitor.visit(tree) + result_video = as_tensor(result_video, ae.shape) @@ -219,9 +197,13 @@ class VideoMathNode(io.ComfyNode): for k, val in F.items(): variables[k] = val if val is not None else 0.0 - tree = parse_expr(Expression); - visitor = UnifiedMathVisitor(variables, a_w.shape,state_storage=stack) - result1 = visitor.visit(tree) - result1 = as_tensor(result, a_w.shape) + tree_pi = None + if isinstance(Expression_pi,str): + tree_pi = parse_expr(Expression_pi) + else: + tree_pi = Expression_pi + visitor = UnifiedMathVisitor(variables, a_w.shape,a_w.device,state_storage=stack) + result_waveform = visitor.visit(tree_pi) + result_waveform = as_tensor(result_waveform, a_w.shape) - return ([result,{"waveform":result1,"sample_rate":sample_rate}],stack) + return ([result_video,{"waveform":result_waveform,"sample_rate":sample_rate}],stack) diff --git a/more_math/modelLikeCommon.py b/more_math/modelLikeCommon.py index 8db4356..363a43f 100644 --- a/more_math/modelLikeCommon.py +++ b/more_math/modelLikeCommon.py @@ -13,7 +13,7 @@ def calculate_patches_autogrow(Expr, V, F,pbar, mapping=None,stack = []): Iterates over the UNION of keys from all input models to support merging disjoint architectures/patches. Args: - Expr: The math expression string. + Expr: The math expression string or parse tree. V: Dictionary of input models (Autogrow Input). F: Dictionary of input floats (Autogrow Float). mapping: Optional dict mapping legacy alias names to V keys (e.g. {"a": "V0"}). @@ -53,7 +53,11 @@ def calculate_patches_autogrow(Expr, V, F,pbar, mapping=None,stack = []): return sd.get(key, None) return None - tree = parse_expr(Expr) + tree = None + if isinstance(Expr,str): + tree = parse_expr(Expr) + else: + tree = Expr patches = {} # Progress bar if possible (comfy.utils.ProgressBar might assume unthreaded?)