math-eval-node/math_eval.py at main · JPPhoto/math-eval-node · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
# Copyright (c) 2023 Jonathan S. Pollack (https://github.com/JPPhoto)

import math

from invokeai.invocation_api import (
    BaseInvocation,
    BaseInvocationOutput,
    InputField,
    InvocationContext,
    OutputField,
    invocation,
    invocation_output,
)


@invocation_output("math_eval_output")
class MathEvalOutput(BaseInvocationOutput):
    """Base class for math output"""

    a: float = OutputField(description="Output a")
    b: float = OutputField(description="Output b")
    c: float = OutputField(description="Output c")
    d: float = OutputField(description="Output d")


@invocation("math_eval", title="MathEval", tags=["math_eval"], version="1.1.1")
class MathEvalInvocation(BaseInvocation):
    """Performs arbitrary user-defined calculations on x, y, z, and w variables"""

    x: float = InputField(default=0.0, description="Input x")
    y: float = InputField(default=0.0, description="Input y")
    z: float = InputField(default=0.0, description="Input z")
    w: float = InputField(default=0.0, description="Input w")
    equation_a: str = InputField(
        default="math.sin(x)", description="A basic mathematical equation for a involving x, y, z, and w"
    )
    equation_b: str = InputField(
        default="1+(x*y)", description="A basic mathematical equation for b involving x, y, z, and w"
    )
    equation_c: str = InputField(default="", description="A basic mathematical equation for c involving x, y, z, and w")
    equation_d: str = InputField(default="", description="A basic mathematical equation for d involving x, y, z, and w")

    def safe_eval(self, equation):
        result = 0
        if len(equation) > 0:
            my_dict = {}
            my_dict["math"] = math
            my_dict["min"] = min
            my_dict["max"] = max
            my_dict["x"] = self.x
            my_dict["y"] = self.y
            my_dict["z"] = self.z
            my_dict["w"] = self.w
            result = eval(equation, {"__builtins__": None}, my_dict)
        return result

    def invoke(self, context: InvocationContext) -> MathEvalOutput:
        return MathEvalOutput(
            a=self.safe_eval(self.equation_a),
            b=self.safe_eval(self.equation_b),
            c=self.safe_eval(self.equation_c),
            d=self.safe_eval(self.equation_d),
        )