Skip to content

Commit 3a55fbe

Browse files
mcDandymcDandy
committed
Made inputs lazy when possible
1 parent 7ec9b9c commit 3a55fbe

File tree

10 files changed

+93
-72
lines changed

10 files changed

+93
-72
lines changed

‎more_math/AudioMathNode.py‎

Lines changed: 11 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
from .Parser.MathExprParser import MathExprParser
44
from .Parser.MathExprLexer import MathExprLexer
55
from .Parser.TensorEvalVisitor import TensorEvalVisitor
6-
from .helper_functions import getIndexTensorAlongDim
6+
from .helper_functions import getIndexTensorAlongDim, comonLazy
77

88
from comfy_api.latest import io
99

@@ -33,20 +33,22 @@ def define_schema(cls) -> io.Schema:
3333
display_name="Audio math",
3434
inputs=[
3535
io.Audio.Input(id="a", tooltip="Input audio tensor"),
36-
io.Audio.Input(id="b", optional=True, tooltip="Second input audio tensor"),
37-
io.Audio.Input(id="c", optional=True, tooltip="Third input audio tensor"),
38-
io.Audio.Input(id="d", optional=True, tooltip="Fourth input audio tensor"),
39-
io.Float.Input(id="w", default=0.0, optional=True, force_input=True),
40-
io.Float.Input(id="x", default=0.0, optional=True, force_input=True),
41-
io.Float.Input(id="y", default=0.0, optional=True, force_input=True),
42-
io.Float.Input(id="z", default=0.0, optional=True, force_input=True),
36+
io.Audio.Input(id="b", optional=True,lazy=True, tooltip="Second input audio tensor"),
37+
io.Audio.Input(id="c", optional=True,lazy=True, tooltip="Third input audio tensor"),
38+
io.Audio.Input(id="d", optional=True,lazy=True, tooltip="Fourth input audio tensor"),
39+
io.Float.Input(id="w", default=0.0, optional=True,lazy=True, force_input=True),
40+
io.Float.Input(id="x", default=0.0, optional=True,lazy=True, force_input=True),
41+
io.Float.Input(id="y", default=0.0, optional=True,lazy=True, force_input=True),
42+
io.Float.Input(id="z", default=0.0, optional=True,lazy=True, force_input=True),
4343
io.String.Input(id="AudioExpr", default="a*(1-w)+b*w", tooltip="Expression to apply on input audio tensors"),
4444
],
4545
outputs=[
4646
io.Audio.Output(),
4747
],
4848
)
49-
49+
@classmethod
50+
def check_lazy_status(cls, AudioExpr, a, b=[], c=[], d=[],w=0,x=0,y=0,z=0):
51+
return comonLazy(AudioExpr, a, b, c, d)
5052
@classmethod
5153
def execute(cls, a, AudioExpr, b=None, c=None, d=None, w=0.0, x=0.0, y=0.0, z=0.0):
5254

‎more_math/ClipMathNode.py‎

Lines changed: 11 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
from comfy_api.latest import io
22
from .modelLikeCommon import calculate_patches
33
from inspect import cleandoc
4+
from .helper_functions import comonLazy
45

56
class CLIPMathNode(io.ComfyNode):
67
"""
@@ -14,13 +15,13 @@ def define_schema(cls) -> io.Schema:
1415
category="More math",
1516
inputs=[
1617
io.Clip.Input(id="a", tooltip="Main CLIP (base)"),
17-
io.Clip.Input(id="b", optional=True, tooltip="Optional 2nd CLIP"),
18-
io.Clip.Input(id="c", optional=True, tooltip="Optional 3rd CLIP"),
19-
io.Clip.Input(id="d", optional=True, tooltip="Optional 4th CLIP"),
20-
io.Float.Input(id="w", default=0.0, optional=True, force_input=True),
21-
io.Float.Input(id="x", default=0.0, optional=True, force_input=True),
22-
io.Float.Input(id="y", default=0.0, optional=True, force_input=True),
23-
io.Float.Input(id="z", default=0.0, optional=True, force_input=True),
18+
io.Clip.Input(id="b", optional=True,lazy=True, tooltip="Optional 2nd CLIP"),
19+
io.Clip.Input(id="c", optional=True,lazy=True, tooltip="Optional 3rd CLIP"),
20+
io.Clip.Input(id="d", optional=True,lazy=True, tooltip="Optional 4th CLIP"),
21+
io.Float.Input(id="w", default=0.0, optional=True,lazy=True, force_input=True),
22+
io.Float.Input(id="x", default=0.0, optional=True,lazy=True, force_input=True),
23+
io.Float.Input(id="y", default=0.0, optional=True,lazy=True, force_input=True),
24+
io.Float.Input(id="z", default=0.0, optional=True,lazy=True, force_input=True),
2425
io.String.Input(id="Model", default="a*(1-w)+b*w", tooltip="Expression to apply on weights"),
2526
],
2627
outputs=[
@@ -29,6 +30,9 @@ def define_schema(cls) -> io.Schema:
2930
)
3031

3132
tooltip = cleandoc(__doc__)
33+
@classmethod
34+
def check_lazy_status(cls, Model, a, b=[], c=[], d=[],w=0,x=0,y=0,z=0):
35+
return comonLazy(Model, a, b, c, d)
3236

3337
@classmethod
3438
def execute(cls, Model, a, b=None, c=None, d=None, w=0.0, x=0.0, y=0.0, z=0.0) -> io.NodeOutput:

‎more_math/ConditioningMathNode.py‎

Lines changed: 12 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
from .Parser.MathExprParser import MathExprParser
66
from .Parser.MathExprLexer import MathExprLexer
77
from .Parser.TensorEvalVisitor import TensorEvalVisitor
8+
from .helper_functions import ThrowingErrorListener, comonLazy
89

910
from comfy_api.latest import io
1011

@@ -36,13 +37,13 @@ def define_schema(cls) -> io.Schema:
3637
category="More math",
3738
inputs=[
3839
io.Conditioning.Input(id="a"),
39-
io.Conditioning.Input(id="b", optional=True),
40-
io.Conditioning.Input(id="c", optional=True),
41-
io.Conditioning.Input(id="d", optional=True),
42-
io.Float.Input(id="w", default=0.0,optional=True, force_input=True),
43-
io.Float.Input(id="x", default=0.0,optional=True, force_input=True),
44-
io.Float.Input(id="y", default=0.0,optional=True, force_input=True),
45-
io.Float.Input(id="z", default=0.0,optional=True, force_input=True),
40+
io.Conditioning.Input(id="b", optional=True,lazy=True),
41+
io.Conditioning.Input(id="c", optional=True,lazy=True),
42+
io.Conditioning.Input(id="d", optional=True,lazy=True),
43+
io.Float.Input(id="w", default=0.0,optional=True,lazy=True, force_input=True),
44+
io.Float.Input(id="x", default=0.0,optional=True,lazy=True, force_input=True),
45+
io.Float.Input(id="y", default=0.0,optional=True,lazy=True, force_input=True),
46+
io.Float.Input(id="z", default=0.0,optional=True,lazy=True, force_input=True),
4647
io.String.Input(id="Tensor", default="a*(1-w)+b*w", tooltip="Describes composition of the image."),
4748
io.String.Input(id="pooled_output", default="a*(1-w)+b*w", tooltip="Composition of the image condensed into one vector"),
4849
],
@@ -55,7 +56,9 @@ def define_schema(cls) -> io.Schema:
5556

5657
#OUTPUT_NODE = False
5758
#OUTPUT_TOOLTIPS = ("",) # Tooltips for the output node
58-
59+
@classmethod
60+
def check_lazy_status(cls, Tensor,pooled_output, a, b=[], c=[], d=[],w=0,x=0,y=0,z=0):
61+
return list(set(comonLazy(Tensor, a, b, c, d)).union(comonLazy(pooled_output, a, b, c, d)))
5962
@classmethod
6063
def execute(cls, Tensor,pooled_output, a, b=None, c=None, d=None,w=0.0,x=0.0,y=0.0,z=0.0):
6164
if b is None:
@@ -87,6 +90,7 @@ def execute(cls, Tensor,pooled_output, a, b=None, c=None, d=None,w=0.0,x=0.0,y=0
8790
lexer = MathExprLexer(input_stream)
8891
stream = CommonTokenStream(lexer)
8992
parser = MathExprParser(stream)
93+
parser.addErrorListener(ThrowingErrorListener())
9094
tree = parser.expr()
9195
visitor = TensorEvalVisitor(variables,ta.shape)
9296
result1 = visitor.visit(tree)

‎more_math/FloatMathNode.py‎

Lines changed: 14 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,11 @@
22

33
from antlr4 import CommonTokenStream, InputStream
44

5-
from .helper_functions import ThrowingErrorListener
5+
from .helper_functions import ThrowingErrorListener, comonLazy
66

77
from .Parser.MathExprParser import MathExprParser
88
from .Parser.MathExprLexer import MathExprLexer
99
from .Parser.FloatEvalVisitor import FloatEvalVisitor
10-
1110
from comfy_api.latest import io
1211

1312
class FloatMathNode(io.ComfyNode):
@@ -27,7 +26,9 @@ class FloatMathNode(io.ComfyNode):
2726
"""
2827
def __init__(self):
2928
pass
30-
29+
@classmethod
30+
def check_lazy_status(cls, Model, a, b=[], c=[], d=[],w=0,x=0,y=0,z=0):
31+
return comonLazy(Model, a, b, c, d)
3132
@classmethod
3233
def define_schema(cls) -> io.Schema:
3334
"""
@@ -38,13 +39,13 @@ def define_schema(cls) -> io.Schema:
3839
display_name="Float math",
3940
inputs=[
4041
io.Float.Input(id="a", force_input=True),
41-
io.Float.Input(id="b", default=0.0,optional=True, force_input=True),
42-
io.Float.Input(id="c", default=0.0,optional=True, force_input=True),
43-
io.Float.Input(id="d", default=0.0,optional=True, force_input=True),
44-
io.Float.Input(id="w", default=0.0,optional=True, force_input=True),
45-
io.Float.Input(id="x", default=0.0,optional=True, force_input=True),
46-
io.Float.Input(id="y", default=0.0,optional=True, force_input=True),
47-
io.Float.Input(id="z", default=0.0,optional=True, force_input=True),
42+
io.Float.Input(id="b", default=0.0,optional=True,lazy=True, force_input=True),
43+
io.Float.Input(id="c", default=0.0,optional=True,lazy=True, force_input=True),
44+
io.Float.Input(id="d", default=0.0,optional=True,lazy=True, force_input=True),
45+
io.Float.Input(id="w", default=0.0,optional=True,lazy=True, force_input=True),
46+
io.Float.Input(id="x", default=0.0,optional=True,lazy=True, force_input=True),
47+
io.Float.Input(id="y", default=0.0,optional=True,lazy=True, force_input=True),
48+
io.Float.Input(id="z", default=0.0,optional=True,lazy=True, force_input=True),
4849
io.String.Input(id="FloatFunc", default="a*(1-w)+b*w", tooltip="Expression to use on inputs"),
4950
],
5051
outputs=[
@@ -58,10 +59,11 @@ def define_schema(cls) -> io.Schema:
5859
#OUTPUT_NODE = False
5960
#OUTPUT_TOOLTIPS = ("",) # Tooltips for the output node
6061
@classmethod
62+
def check_lazy_status(cls, FloatFunc, a, b=[], c=[], d=[],w=0,x=0,y=0,z=0):
63+
return comonLazy(FloatFunc, a, b, c, d)
64+
@classmethod
6165
def execute(cls, FloatFunc, a, b=0.0, c=0.0, d=0.0, w=0.0, x=0.0, y=0.0, z=0.0):
6266

63-
64-
6567
variables = {'a': a, 'b': b, 'c': c, 'd': d, 'w': w, 'x': x, 'y': y, 'z': z}
6668
input_stream = InputStream(FloatFunc)
6769
lexer = MathExprLexer(input_stream)

‎more_math/ImageMathNode.py‎

Lines changed: 11 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
from antlr4 import CommonTokenStream
33
from antlr4.atn.LexerActionExecutor import InputStream
44
import torch
5-
from .helper_functions import ThrowingErrorListener, getIndexTensorAlongDim
5+
from .helper_functions import ThrowingErrorListener, getIndexTensorAlongDim, comonLazy
66

77
from .Parser.MathExprParser import MathExprParser
88
from .Parser.MathExprLexer import MathExprLexer
@@ -38,20 +38,22 @@ def define_schema(cls) -> io.Schema:
3838
display_name="Image math",
3939
inputs=[
4040
io.Image.Input(id="a"),
41-
io.Image.Input(id="b", optional=True),
42-
io.Image.Input(id="c", optional=True),
43-
io.Image.Input(id="d", optional=True),
44-
io.Float.Input(id="w", default=0.0,optional=True, force_input=True),
45-
io.Float.Input(id="x", default=0.0,optional=True, force_input=True),
46-
io.Float.Input(id="y", default=0.0,optional=True, force_input=True),
47-
io.Float.Input(id="z", default=0.0,optional=True, force_input=True),
41+
io.Image.Input(id="b", optional=True,lazy=True),
42+
io.Image.Input(id="c", optional=True,lazy=True),
43+
io.Image.Input(id="d", optional=True,lazy=True),
44+
io.Float.Input(id="w", default=0.0,optional=True,lazy=True, force_input=True),
45+
io.Float.Input(id="x", default=0.0,optional=True,lazy=True, force_input=True),
46+
io.Float.Input(id="y", default=0.0,optional=True,lazy=True, force_input=True),
47+
io.Float.Input(id="z", default=0.0,optional=True,lazy=True, force_input=True),
4848
io.String.Input(id="Image", default="a*(1-w)+b*w", tooltip="Expression to apply on input images"),
4949
],
5050
outputs=[
5151
io.Image.Output(),
5252
],
5353
)
54-
54+
@classmethod
55+
def check_lazy_status(cls, Image, a, b=[], c=[], d=[],w=0,x=0,y=0,z=0):
56+
return comonLazy(Image, a, b, c, d)
5557
@classmethod
5658
def execute(scls, Image, a, b=None, c=None, d=None, w=0.0, x=0.0, y=0.0, z=0.0):
5759
b = torch.zeros_like(a) if b is None else b

‎more_math/LatentMathNode.py‎

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
from antlr4 import CommonTokenStream, InputStream
66
import torch
77

8-
from .helper_functions import ThrowingErrorListener, getIndexTensorAlongDim
8+
from .helper_functions import ThrowingErrorListener, getIndexTensorAlongDim,comonLazy
99

1010
from .Parser.MathExprParser import MathExprParser
1111
from .Parser.MathExprLexer import MathExprLexer
@@ -50,10 +50,10 @@ def define_schema(cls) -> io.Schema:
5050
io.Latent.Input(id="b", optional=True, lazy=True),
5151
io.Latent.Input(id="c", optional=True, lazy=True),
5252
io.Latent.Input(id="d", optional=True, lazy=True),
53-
io.Float.Input(id="w", default=0.0,optional=True, force_input=True),
54-
io.Float.Input(id="x", default=0.0,optional=True, force_input=True),
55-
io.Float.Input(id="y", default=0.0,optional=True, force_input=True),
56-
io.Float.Input(id="z", default=0.0,optional=True, force_input=True),
53+
io.Float.Input(id="w", default=0.0,optional=True,lazy=True, force_input=True),
54+
io.Float.Input(id="x", default=0.0,optional=True,lazy=True, force_input=True),
55+
io.Float.Input(id="y", default=0.0,optional=True,lazy=True, force_input=True),
56+
io.Float.Input(id="z", default=0.0,optional=True,lazy=True, force_input=True),
5757
io.String.Input(id="Latent", default="a*(1-w)+b*w", tooltip="Expression to apply on input latents"),
5858
],
5959
outputs=[
@@ -66,7 +66,9 @@ def define_schema(cls) -> io.Schema:
6666

6767
#OUTPUT_NODE = False
6868
#OUTPUT_TOOLTIPS = ("",) # Tooltips for the output node
69-
69+
@classmethod
70+
def check_lazy_status(cls, Latent, a, b=[], c=[], d=[],w=0,x=0,y=0,z=0):
71+
return comonLazy(Latent, a, b, c, d)
7072
@classmethod
7173
def execute(cls, Latent, a, b=None, c=None, d=None, w=0.0, x=0.0, y=0.0, z=0.0) -> io.NodeOutput:
7274
# Extract raw sample tensors (may be Tensor or NestedTensor)

‎more_math/ModelMathNode.py‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,8 @@
11
from inspect import cleandoc
22
from comfy_api.latest import io
3-
from antlr4 import CommonTokenStream, InputStream
4-
from .Parser.MathExprLexer import MathExprLexer
53
from .helper_functions import comonLazy
64
from .modelLikeCommon import calculate_patches
75

8-
import copy
96

107

118
class ModelMathNode(io.ComfyNode):

‎more_math/NoiseMathNode.py‎

Lines changed: 11 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
from antlr4 import CommonTokenStream, InputStream
44
import torch
55

6-
from .helper_functions import ThrowingErrorListener, getIndexTensorAlongDim
6+
from .helper_functions import ThrowingErrorListener, getIndexTensorAlongDim,comonLazy
77

88
from .Parser.MathExprParser import MathExprParser
99
from .Parser.MathExprLexer import MathExprLexer
@@ -43,13 +43,13 @@ def define_schema(cls) -> io.Schema:
4343
display_name="Noise math",
4444
inputs=[
4545
io.Noise.Input(id="a"),
46-
io.Noise.Input(id="b", optional=True),
47-
io.Noise.Input(id="c", optional=True),
48-
io.Noise.Input(id="d", optional=True),
49-
io.Float.Input(id="w", default=0.0,optional=True, force_input=True),
50-
io.Float.Input(id="x", default=0.0,optional=True, force_input=True),
51-
io.Float.Input(id="y", default=0.0,optional=True, force_input=True),
52-
io.Float.Input(id="z", default=0.0,optional=True, force_input=True),
46+
io.Noise.Input(id="b", optional=True,lazy=True),
47+
io.Noise.Input(id="c", optional=True,lazy=True),
48+
io.Noise.Input(id="d", optional=True,lazy=True),
49+
io.Float.Input(id="w", default=0.0,optional=True,lazy=True, force_input=True),
50+
io.Float.Input(id="x", default=0.0,optional=True,lazy=True, force_input=True),
51+
io.Float.Input(id="y", default=0.0,optional=True,lazy=True, force_input=True),
52+
io.Float.Input(id="z", default=0.0,optional=True,lazy=True, force_input=True),
5353
io.String.Input(id="Noise", default="a*(1-w)+b*w", tooltip="Expression to apply on input noise generators"),
5454
],
5555
outputs=[
@@ -62,7 +62,9 @@ def define_schema(cls) -> io.Schema:
6262
#OUTPUT_TOOLTIPS = ("",) # Tooltips for the output node
6363

6464
CATEGORY = "More math"
65-
65+
@classmethod
66+
def check_lazy_status(cls, Noise, a, b=[], c=[], d=[],w=0,x=0,y=0,z=0):
67+
return comonLazy(Noise, a, b, c, d)
6668
@classmethod
6769
def execute(cls, Noise, a, b=None, c=None, d=None, w=0.0, x=0.0, y=0.0, z=0.0):
6870
return (NoiseExecutor(a, b, c, d, w, x, y, z, Noise),)

‎more_math/VaeMathNode.py‎

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
from comfy_api.latest import io
33
import copy
44
from .modelLikeCommon import calculate_patches
5+
from .helper_functions import comonLazy
56

67
class VAEMathNode(io.ComfyNode):
78
"""
@@ -15,13 +16,13 @@ def define_schema(cls) -> io.Schema:
1516
category="More math",
1617
inputs=[
1718
io.Vae.Input(id="a", tooltip="Main VAE (base)"),
18-
io.Vae.Input(id="b", optional=True, tooltip="Optional 2nd VAE"),
19-
io.Vae.Input(id="c", optional=True, tooltip="Optional 3rd VAE"),
20-
io.Vae.Input(id="d", optional=True, tooltip="Optional 4th VAE"),
21-
io.Float.Input(id="w", default=0.0, optional=True, force_input=True),
22-
io.Float.Input(id="x", default=0.0, optional=True, force_input=True),
23-
io.Float.Input(id="y", default=0.0, optional=True, force_input=True),
24-
io.Float.Input(id="z", default=0.0, optional=True, force_input=True),
19+
io.Vae.Input(id="b", optional=True,lazy=True, tooltip="Optional 2nd VAE"),
20+
io.Vae.Input(id="c", optional=True,lazy=True, tooltip="Optional 3rd VAE"),
21+
io.Vae.Input(id="d", optional=True,lazy=True, tooltip="Optional 4th VAE"),
22+
io.Float.Input(id="w", default=0.0,lazy=True, optional=True, force_input=True),
23+
io.Float.Input(id="x", default=0.0,lazy=True, optional=True, force_input=True),
24+
io.Float.Input(id="y", default=0.0,lazy=True, optional=True, force_input=True),
25+
io.Float.Input(id="z", default=0.0,lazy=True, optional=True, force_input=True),
2526
io.String.Input(id="Model", default="a*(1-w)+b*w", tooltip="Expression to apply on weights"),
2627
],
2728
outputs=[
@@ -30,7 +31,9 @@ def define_schema(cls) -> io.Schema:
3031
)
3132

3233
tooltip = cleandoc(__doc__)
33-
34+
@classmethod
35+
def check_lazy_status(cls, Model, a, b=[], c=[], d=[],w=0,x=0,y=0,z=0):
36+
return comonLazy(Model, a, b, c, d)
3437
@classmethod
3538
def execute(cls, Model, a, b=None, c=None, d=None, w=0.0, x=0.0, y=0.0, z=0.0) -> io.NodeOutput:
3639
patcher_a = a.patcher

‎more_math/helper_functions.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,8 @@
11
from antlr4.error.ErrorListener import ErrorListener
22
from antlr4 import InputStream
3+
from .Parser.MathExprLexer import MathExprLexer
4+
from .Parser.MathExprParser import MathExprParser
5+
from antlr4 import CommonTokenStream
36

47
import torch
58

@@ -31,7 +34,7 @@ def freq_to_time(element: torch.Tensor) -> torch.Tensor:
3134

3235
class ThrowingErrorListener(ErrorListener):
3336
def syntaxError(self, recognizer, offendingSymbol, line, column, msg, e):
34-
raise ValueError(f"Syntax error in AudioExpr at line {line}, col {column}: {msg}")
37+
raise ValueError(f"Syntax error in expression at line {line}, col {column}: {msg}")
3538

3639

3740
def comonLazy(expr, a, b=None, c=None, d=None, w=0.0, x=0.0, y=0.0, z=0.0):

0 commit comments

Comments
 (0)