diff --git a/more_math/MaskMathNode.py b/more_math/MaskMathNode.py new file mode 100644 index 0000000..808c61c --- /dev/null +++ b/more_math/MaskMathNode.py @@ -0,0 +1,71 @@ +from .helper_functions import generate_dim_variables, getIndexTensorAlongDim, eval_tensor_expr, as_tensor + +from comfy_api.latest import io + + +from .MathNodeBase import MathNodeBase + + +class MaskMathNode(MathNodeBase): + """ + Enables math expressions on Images. + + Inputs: + a, b, c, d: Mask inputs (b, c, d default to zero if not provided) + w, x, y, z: Float variables for expressions + mask: Expression to apply on input mask + + Outputs: + MASK: Result of applying expression to input masks + """ + + @classmethod + def define_schema(cls) -> io.Schema: + return io.Schema( + node_id="mrmth_MaskMathNode", + category="More math", + display_name="Mask math", + inputs=[ + io.Mask.Input(id="a"), + io.Mask.Input(id="b", optional=True, lazy=True), + io.Mask.Input(id="c", optional=True, lazy=True), + io.Mask.Input(id="d", optional=True, lazy=True), + io.Float.Input(id="w", default=0.0, optional=True, lazy=True, force_input=True), + io.Float.Input(id="x", default=0.0, optional=True, lazy=True, force_input=True), + io.Float.Input(id="y", default=0.0, optional=True, lazy=True, force_input=True), + io.Float.Input(id="z", default=0.0, optional=True, lazy=True, force_input=True), + io.String.Input(id="Mask", default="a*(1-w)+b*w", tooltip="Expression to apply on input images"), + ], + outputs=[ + io.Mask.Output(), + ], + ) + + @classmethod + def execute(cls, Mask, a, b=None, c=None, d=None, w=0.0, x=0.0, y=0.0, z=0.0): + a, b, c, d = cls.prepare_inputs(a, b, c, d) + + variables = { + "a": a, + "b": b, + "c": c, + "d": d, + "w": w, + "x": x, + "y": y, + "z": z, + "X": getIndexTensorAlongDim(a, 2), + "Y": getIndexTensorAlongDim(a, 1), + "B": getIndexTensorAlongDim(a, 0), + "batch": getIndexTensorAlongDim(a, 0), + "W": a.shape[2], + "width": a.shape[2], + "H": a.shape[1], + "height": a.shape[1], + "T": a.shape[0], + "batch_count": a.shape[0], + } | generate_dim_variables(a) + + result = eval_tensor_expr(Mask, variables, a.shape) + + return (as_tensor(result, a.shape),) diff --git a/more_math/nodes.py b/more_math/nodes.py index b6a6676..b4e884b 100644 --- a/more_math/nodes.py +++ b/more_math/nodes.py @@ -3,6 +3,7 @@ from .FloatMathNode import FloatMathNode from .ConditioningMathNode import ConditioningMathNode from .LatentMathNode import LatentMathNode from .ImageMathNode import ImageMathNode +from .MaskMathNode import MaskMathNode from .AudioMathNode import AudioMathNode from .VideoMathNode import VideoMathNode from .ModelMathNode import ModelMathNode @@ -77,6 +78,7 @@ class MoreMathExtension(ComfyExtension): VAEMathNode, LatentMathNode, ImageMathNode, + MaskMathNode, FloatMathNode, NoiseMathNode, IntToFloatNode,