113 lines
5.3 KiB
Python
113 lines
5.3 KiB
Python
from inspect import cleandoc
|
|
from antlr4 import CommonTokenStream, InputStream
|
|
import torch
|
|
|
|
from .Parser.MathExprParser import MathExprParser
|
|
from .Parser.MathExprLexer import MathExprLexer
|
|
from .Parser.TensorEvalVisitor import TensorEvalVisitor
|
|
|
|
from comfy_api.latest import ComfyExtension, io
|
|
|
|
class ConditioningMathNode(io.ComfyNode):
|
|
"""
|
|
This node enables the use of math on conditionings. It is recommended to keep the Tensor expression the same as the pooled_output expression.
|
|
inputs:
|
|
a, b, c, d:
|
|
Conditioning, bound to variables with the same name. Defaults to zero conditioning if not provided.
|
|
w, x, y, z:
|
|
Floats, bound to variables of the expression. Defaults to 0.0 if not provided.
|
|
Tensor expression:
|
|
String, describing expression to mix tensor part of conditioning. Tensor probably describes the composition of the image as it has the say what is on the image. Valid functions are sin, cos, tan, abs, sqrt, min, max, norm. Valid operators are +, -, *, /, ^, %. Usable constants are e and pi.
|
|
Pooled output expression:
|
|
String, describing expression to mix pooled_output part of conditioning. pooled_output does not have that much say in the image but can change some details of the image or completly change the image in some cases, probably also used for positional and temporal conditioning. Expression uses same syntax as Tensor expression.
|
|
|
|
outputs:
|
|
CONDITIONING:
|
|
Returns a CONDITIONING object that contains the result of the math expression applied to the input conditionings.
|
|
"""
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id="mrmth_ConditioningMathNode",
|
|
category="More math",
|
|
inputs=[
|
|
io.Conditioning.Input(id="a"),
|
|
io.Conditioning.Input(id="b", optional=True),
|
|
io.Conditioning.Input(id="c", optional=True),
|
|
io.Conditioning.Input(id="d", optional=True),
|
|
io.Float.Input(id="w", default=0.0,optional=True, force_input=True),
|
|
io.Float.Input(id="x", default=0.0,optional=True, force_input=True),
|
|
io.Float.Input(id="y", default=0.0,optional=True, force_input=True),
|
|
io.Float.Input(id="z", default=0.0,optional=True, force_input=True),
|
|
io.String.Input(id="Tensor", default="a*(1-w)+b*w", tooltip="Describes composition of the image."),
|
|
io.String.Input(id="pooled_output", default="a*(1-w)+b*w", tooltip="Can change some details of the image. Idk."),
|
|
],
|
|
outputs=[
|
|
io.Conditioning.Output(),
|
|
],
|
|
)
|
|
|
|
tooltip = cleandoc(__doc__)
|
|
|
|
#OUTPUT_NODE = False
|
|
#OUTPUT_TOOLTIPS = ("",) # Tooltips for the output node
|
|
|
|
@classmethod
|
|
def execute(self,cls, Tensor,pooled_output, a, b=None, c=None, d=None,w=0.0,x=0.0,y=0.0,z=0.0):
|
|
if b is None:
|
|
b = [[torch.zeros_like(a[0][0]), {"pooled_output": torch.zeros_like(a[0][1]["pooled_output"])}]]
|
|
if c is None:
|
|
c = [[torch.zeros_like(a[0][0]), {"pooled_output": torch.zeros_like(a[0][1]["pooled_output"])}]]
|
|
if d is None:
|
|
d = [[torch.zeros_like(a[0][0]), {"pooled_output": torch.zeros_like(a[0][1]["pooled_output"])}]]
|
|
|
|
|
|
ta = a[0][0].clone()
|
|
tb = b[0][0].clone()
|
|
pa = a[0][1]["pooled_output"].clone()
|
|
pb = b[0][1]["pooled_output"].clone()
|
|
tc = c[0][0].clone()
|
|
td = d[0][0].clone()
|
|
pc = c[0][1]["pooled_output"].clone()
|
|
pd = d[0][1]["pooled_output"].clone()
|
|
|
|
variables = {'a': ta, 'b': tb, 'c': tc, 'd': td, 'w': w, 'x': x, 'y': y, 'z': z}
|
|
input_stream = InputStream(Tensor)
|
|
lexer = MathExprLexer(input_stream)
|
|
stream = CommonTokenStream(lexer)
|
|
parser = MathExprParser(stream)
|
|
tree = parser.expr()
|
|
print("Tensor\n"+tree.toStringTree(recog=parser))
|
|
visitor = TensorEvalVisitor(variables,ta.shape)
|
|
result1 = visitor.visit(tree)
|
|
|
|
|
|
variables = {'a': pa, 'b': pb, 'c': pc, 'd': pd, 'w': w, 'x': x, 'y': y, 'z': z}
|
|
input_stream = InputStream(Tensor)
|
|
lexer = MathExprLexer(input_stream)
|
|
stream = CommonTokenStream(lexer)
|
|
parser = MathExprParser(stream)
|
|
tree = parser.expr()
|
|
print("pooled_output\n"+tree.toStringTree(recog=parser))
|
|
visitor = TensorEvalVisitor(variables,pa.shape)
|
|
result2 = visitor.visit(tree)
|
|
|
|
|
|
result = [[result1, {"pooled_output": result2}]]
|
|
return (result,)
|
|
|
|
"""
|
|
The node will always be re executed if any of the inputs change but
|
|
this method can be used to force the node to execute again even when the inputs don't change.
|
|
You can make this node return a number or a string. This value will be compared to the one returned the last time the node was
|
|
executed, if it is different the node will be executed again.
|
|
This method is used in the core repo for the LoadImage node where they return the image hash as a string, if the image hash
|
|
changes between executions the LoadImage node is executed again.
|
|
"""
|
|
#@classmethod
|
|
#def IS_CHANGED(s, image, string_field, int_field, float_field, print_to_screen):
|
|
# return ""
|