mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-02-01 01:00:20 +08:00
104 lines
2.5 KiB
Python
104 lines
2.5 KiB
Python
from typing_extensions import override
|
|
|
|
from comfy_api.latest import ComfyExtension, io
|
|
|
|
|
|
class IntAdd(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="IntAdd",
|
|
display_name="Int Add",
|
|
category="math",
|
|
inputs=[
|
|
io.Int.Input("a"),
|
|
io.Int.Input("b"),
|
|
],
|
|
outputs=[io.Int.Output()],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, a: int, b: int) -> io.NodeOutput:
|
|
return io.NodeOutput(a + b)
|
|
|
|
|
|
class IntSubtract(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="IntSubtract",
|
|
display_name="Int Subtract",
|
|
category="math",
|
|
inputs=[
|
|
io.Int.Input("a"),
|
|
io.Int.Input("b"),
|
|
],
|
|
outputs=[io.Int.Output()],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, a: int, b: int) -> io.NodeOutput:
|
|
return io.NodeOutput(a - b)
|
|
|
|
|
|
class IntMultiply(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="IntMultiply",
|
|
display_name="Int Multiply",
|
|
category="math",
|
|
inputs=[
|
|
io.Int.Input("a"),
|
|
io.Int.Input("b"),
|
|
],
|
|
outputs=[io.Int.Output()],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, a: int, b: int) -> io.NodeOutput:
|
|
return io.NodeOutput(a * b)
|
|
|
|
|
|
class IntDivide(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls):
|
|
return io.Schema(
|
|
node_id="IntDivide",
|
|
display_name="Int Divide",
|
|
category="math",
|
|
inputs=[
|
|
io.Int.Input("a"),
|
|
io.Int.Input("b"),
|
|
],
|
|
outputs=[io.Int.Output()],
|
|
)
|
|
|
|
@classmethod
|
|
def validate_inputs(cls, a: int, b: int) -> bool:
|
|
if b == 0:
|
|
return "Division by zero is not allowed"
|
|
return True
|
|
|
|
@classmethod
|
|
def execute(cls, a: int, b: int) -> io.NodeOutput:
|
|
if b == 0:
|
|
raise ValueError("Division by zero is not allowed")
|
|
# Integer division (floor division)
|
|
return io.NodeOutput(a // b)
|
|
|
|
|
|
class MathExtension(ComfyExtension):
|
|
@override
|
|
async def get_node_list(self) -> list[type[io.ComfyNode]]:
|
|
return [
|
|
IntAdd,
|
|
IntSubtract,
|
|
IntMultiply,
|
|
IntDivide,
|
|
]
|
|
|
|
|
|
async def comfy_entrypoint() -> MathExtension:
|
|
return MathExtension()
|