From 1cd5604e7bece0082467adf9f3f7243fe5847956 Mon Sep 17 00:00:00 2001 From: Great <44393369+yaosheng216@users.noreply.github.com> Date: Fri, 14 Feb 2025 16:45:00 +0800 Subject: [PATCH] Create calculator_node.py --- scepter/workflow/calculator_node.py | 64 +++++++++++++++++++++++++++++ 1 file changed, 64 insertions(+) create mode 100644 scepter/workflow/calculator_node.py diff --git a/scepter/workflow/calculator_node.py b/scepter/workflow/calculator_node.py new file mode 100644 index 0000000..2c9ecd1 --- /dev/null +++ b/scepter/workflow/calculator_node.py @@ -0,0 +1,64 @@ +# -*- coding: utf-8 -*- +# Copyright (c) Alibaba, Inc. and its affiliates. +import math +from .constant import WORKFLOW_CONFIG + + +class CalculatorNode: + def __init__(self): + self.cfg = WORKFLOW_CONFIG.workflow_config + + CATEGORY = '🪄 ComfyUI-Scepter' + + @classmethod + def INPUT_TYPES(s): + return { + 'required': { + 'parameter': ('INT',), + 'type': (list(s().cfg['CALCULATOR']['TYPE']),), + 'value': ('FLOAT',), + 'round_method': (list(s().cfg['CALCULATOR']['ROUND']),) + } + } + + OUTPUT_NODE = True + RETURN_TYPES = ('INT',) + RETURN_NAMES = ('INT',) + FUNCTION = 'execute' + + def execute(self, parameter, type, value, round_method): + _OPERATIONS = { + 'add': lambda a, b: a + b, + 'sub': lambda a, b: a - b, + 'mul': lambda a, b: a * b, + 'div': lambda a, b: a / b if b != 0 else _raise_zero_division() + } + + _ROUND_METHODS = { + 'ceil': math.ceil, + 'floor': math.floor, + 'round': round + } + + def _raise_zero_division(): + raise ValueError("Division by zero is not allowed.") + + if not isinstance(parameter, (int, float)) or not isinstance(value, (int, float)): + raise TypeError("Parameters must be int or float.") + + try: + operation = _OPERATIONS[type] + except KeyError: + raise ValueError(f"Invalid type: {type}") from None + + try: + res = operation(parameter, value) + except ZeroDivisionError: + raise ValueError("Division by zero is not allowed") from None + + try: + round_func = _ROUND_METHODS[round_method] + except KeyError: + raise ValueError(f"Invalid rounding method: {round_method}") from None + + return (round_func(res),)