FilesExpand file tree

master

/

nodes_math.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

123 lines (99 loc) · 3.65 KB

master

/

nodes_math.py

Copy path

Top

File metadata and controls

• Code
• Blame

123 lines (99 loc) · 3.65 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

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

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

85

86

87

88

89

90

91

92

93

94

95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

"""Math expression node using simpleeval for safe evaluation.

Provides a ComfyMathExpression node that evaluates math expressions

against dynamically-grown numeric inputs.

"""

import math

import string

from simpleeval import simple_eval

from typing_extensions import override

from comfy_api.latest import ComfyExtension, io

MAX_EXPONENT = 4000

def _variadic_sum(*args):

"""Support both sum(values) and sum(a, b, c)."""

if len(args) == 1 and hasattr(args[0], "__iter__"):

return sum(args[0])

return sum(args)

def _safe_pow(base, exp):

"""Wrap pow() with an exponent cap to prevent DoS via huge exponents.

The ** operator is already guarded by simpleeval's safe_power, but

pow() as a callable bypasses that guard.

"""

if abs(exp) > MAX_EXPONENT:

raise ValueError(f"Exponent {exp} exceeds maximum allowed ({MAX_EXPONENT})")

return pow(base, exp)

MATH_FUNCTIONS = {

"sum": _variadic_sum,

"min": min,

"max": max,

"abs": abs,

"round": round,

"pow": _safe_pow,

"sqrt": math.sqrt,

"ceil": math.ceil,

"floor": math.floor,

"log": math.log,

"log2": math.log2,

"log10": math.log10,

"sin": math.sin,

"cos": math.cos,

"tan": math.tan,

"int": int,

"float": float,

}

class MathExpressionNode(io.ComfyNode):

"""Evaluates a math expression against dynamically-grown inputs."""

@classmethod

def define_schema(cls) -> io.Schema:

autogrow = io.Autogrow.TemplateNames(

input=io.MultiType.Input("value", [io.Float, io.Int, io.Boolean]),

names=list(string.ascii_lowercase),

min=1,

)

return io.Schema(

node_id="ComfyMathExpression",

display_name="Math Expression",

category="utilities",

search_aliases=["math", "expression", "formula", "calculate", "calculator", "eval"],

inputs=[

io.String.Input("expression", default="a + b", multiline=True),

io.Autogrow.Input("values", template=autogrow),

],

outputs=[

io.Float.Output(display_name="FLOAT"),

io.Int.Output(display_name="INT"),

io.Boolean.Output(display_name="BOOL"),

],

)

@classmethod

def execute(

cls, expression: str, values: io.Autogrow.Type

) -> io.NodeOutput:

if not expression.strip():

raise ValueError("Expression cannot be empty.")

context: dict = dict(values)

context["values"] = list(values.values())

result = simple_eval(expression, names=context, functions=MATH_FUNCTIONS)

# bool check must come first because bool is a subclass of int in Python

if not isinstance(result, (int, float)):

raise ValueError(

f"Math Expression '{expression}' must evaluate to a numeric result, "

f"got {type(result).__name__}: {result!r}"

)

try:

float_result = float(result)

except OverflowError:

raise ValueError(

f"Math Expression '{expression}' produced a result too large to "

f"represent as a float: {result}"

) from None

if not math.isfinite(float_result):

raise ValueError(

f"Math Expression '{expression}' produced a non-finite result: {result}"

)

return io.NodeOutput(float_result, int(result), bool(result))

class MathExtension(ComfyExtension):

@override

async def get_node_list(self) -> list[type[io.ComfyNode]]:

return [MathExpressionNode]

async def comfy_entrypoint() -> MathExtension:

return MathExtension()