diff --git a/src/more_math/VideoMathNode.py b/src/more_math/VideoMathNode.py new file mode 100644 index 0000000..baf2a97 --- /dev/null +++ b/src/more_math/VideoMathNode.py @@ -0,0 +1,139 @@ +from inspect import cleandoc + +from comfy_api.latest import ComfyExtension, io +from comfy_api.input_impl import VideoFromComponents, VideoFromFile +from comfy_api.util import VideoCodec, VideoComponents, VideoContainer + + +from antlr4 import CommonTokenStream, InputStream +import torch + +from .helper_functions import ThrowingErrorListener, getIndexTensorAlongDim + +from .Parser.MathExprParser import MathExprParser +from .Parser.MathExprLexer import MathExprLexer +from .Parser.TensorEvalVisitor import TensorEvalVisitor + +class VideoMathNode(io.ComfyNode): + """ + This node enables the use of math expressions on Latents. + inputs: + a, b, c, d: + Latent, bound to variables with the same name. Defaults to zero latent if not provided. + w, x, y, z: + Floats, bound to variables of the expression. Defaults to 0.0 if not provided. + Latent expression: + String, describing expression to aply to latents. + + outputs: + LATENT: + Returns a LATENT 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_VideoMathNode", + display_name="Video math", + category="More math", + inputs=[ + io.Video.Input(id="a"), + io.Video.Input(id="b", optional=True), + io.Video.Input(id="c", optional=True), + io.Video.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="Audio", default="a*(1-w)+b*w", tooltip="Expression to apply on audio part of video"), + io.String.Input(id="Images", default="a*(1-w)+b*w", tooltip="Expression to apply on image part of video"), + ], + outputs=[ + io.Video.Output(), + ], + ) + + #RETURN_NAMES = ("image_output_name",) + tooltip = cleandoc(__doc__) + + #OUTPUT_NODE = False + #OUTPUT_TOOLTIPS = ("",) # Tooltips for the output node + + + @classmethod + def execute(cls, Audio,Images, a, b=None, c=None, d=None, w=0.0, x=0.0, y=0.0, z=0.0) -> io.NodeOutput: + + ac = a.get_components() + bc = b.get_components() if b is not None else VideoComponents(images=torch.zeros_like(ac.images), audio={'waveform':torch.zeros_like(ac.audio['waveform']),'sample_rate':ac.audio['sample_rate']}, frame_rate=ac.frame_rate,metadata=None) + cc = c.get_components() if c is not None else VideoComponents(images=torch.zeros_like(ac.images), audio={'waveform':torch.zeros_like(ac.audio['waveform']),'sample_rate':ac.audio['sample_rate']}, frame_rate=ac.frame_rate,metadata=None) + dc = d.get_components() if d is not None else VideoComponents(images=torch.zeros_like(ac.images), audio={'waveform':torch.zeros_like(ac.audio['waveform']),'sample_rate':ac.audio['sample_rate']}, frame_rate=ac.frame_rate,metadata=None) + + + B = getIndexTensorAlongDim(ac.images, 0) + W = getIndexTensorAlongDim(ac.images, 2) + H = getIndexTensorAlongDim(ac.images, 1) + C = getIndexTensorAlongDim(ac.images, 3) + R = torch.full_like(H, float(ac.frame_rate), dtype=torch.float32) + + variables = {'a': ac.images, 'b': bc.images, 'c': cc.images, 'd': dc.images, 'w': w, 'x': x, 'y': y, 'z': z,'B':B,'X':W,'Y':H,'C':C,'R':R,'W':ac.images.shape[1],'H':ac.images.shape[2] } + input_stream = InputStream(Images) + lexer = MathExprLexer(input_stream) + stream = CommonTokenStream(lexer) + parser = MathExprParser(stream) + parser.addErrorListener(ThrowingErrorListener()) + tree = parser.expr() + visitor = TensorEvalVisitor(variables,ac.images.shape) + imgs = visitor.visit(tree) + + + B = getIndexTensorAlongDim(ac.audio['waveform'], 0) + C = getIndexTensorAlongDim(ac.audio['waveform'], 1) + S = getIndexTensorAlongDim(ac.audio['waveform'], 2) + R = torch.full_like(S, ac.audio['sample_rate'], dtype=torch.float32) + T = torch.full_like(S, ac.audio['waveform'].shape[2], dtype=torch.float32) + + variables = { + 'a': ac.audio['waveform'], 'b': bc.audio['waveform'], 'c': cc.audio['waveform'], 'd': dc.audio['waveform'], + 'w': w, 'x': x, 'y': y, 'z': z, + 'B': B, 'C': C, 'S': S,'R': R, 'T' : T + } + + input_stream = InputStream(Audio) + lexer = MathExprLexer(input_stream) + stream = CommonTokenStream(lexer) + parser = MathExprParser(stream) + tree = parser.expr() + + visitor = TensorEvalVisitor(variables, ac.audio['waveform'].shape) + result_tensor = visitor.visit(tree) + + # Create output dictionary with the same sample rate + audioo = { + 'waveform': result_tensor, + 'sample_rate': ac.audio['sample_rate'] + } + + + + + + + out = VideoFromComponents(VideoComponents(images=imgs, audio=audioo, frame_rate=ac.frame_rate,metadata=ac.metadata)) + return (out,) + + + """ + 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 "" diff --git a/src/more_math/nodes.py b/src/more_math/nodes.py index 134a8a3..9998681 100644 --- a/src/more_math/nodes.py +++ b/src/more_math/nodes.py @@ -6,6 +6,7 @@ from .ConditioningMathNode import ConditioningMathNode from .LatentMathNode import LatentMathNode from .ImageMathNode import ImageMathNode from .AudioMathNode import AudioMathNode +from .VideoMathNode import VideoMathNode from comfy_api.latest import ComfyExtension, io @@ -70,6 +71,7 @@ class MoreMathExtension(ComfyExtension): IntToFloatNode, FloatToIntNode, AudioMathNode, + VideoMathNode ] async def comfy_entrypoint() -> MoreMathExtension: return MoreMathExtension()