153 lines
4.7 KiB
Python
153 lines
4.7 KiB
Python
# Package Modules
|
|
import os
|
|
from typing import Union, BinaryIO, Dict, List, Tuple, Optional
|
|
import time
|
|
|
|
# ComfyUI Modules
|
|
import folder_paths
|
|
from comfy.utils import ProgressBar
|
|
|
|
# Your Modules
|
|
from .modules.calculator import CalculatorModel
|
|
|
|
|
|
# Basic practice to get paths from ComfyUI
|
|
custom_nodes_script_dir = os.path.dirname(os.path.abspath(__file__))
|
|
custom_nodes_model_dir = os.path.join(folder_paths.models_dir, "my-custom-nodes")
|
|
custom_nodes_output_dir = os.path.join(folder_paths.get_output_directory(), "my-custom-nodes")
|
|
|
|
|
|
# These are example nodes that only contains basic functionalities with some comments.
|
|
# If you need detailed explanation, please refer to : https://docs.comfy.org/essentials/custom_node_walkthrough
|
|
# First Node:
|
|
class MyModelLoader:
|
|
# Define the input parameters of the node here.
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
my_models = ["Model A", "Model B", "Model C"]
|
|
|
|
return {
|
|
# If the key is "required", the value must be filled.
|
|
"required": {
|
|
# `my_models` is the list, so it will be shown as a dropdown menu in the node. ( So that user can select one of them. )
|
|
# You must provide the value in the tuple format. e.g. ("value",) or (3,) or ([1, 2],) etc.
|
|
"model": (my_models,),
|
|
"device": (['cuda', 'cpu', 'auto'],),
|
|
},
|
|
# If the key is "optional", the value is optional.
|
|
"optional": {
|
|
"compute_type": (['float32', 'float16'],),
|
|
}
|
|
}
|
|
|
|
# Define these constants inside the node.
|
|
# `RETURN_TYPES` is important, as it limits the parameter types that can be passed to the next node, in `INPUT_TYPES()` above.
|
|
RETURN_TYPES = ("MY_MODEL",)
|
|
RETURN_NAMES = ("my_model",)
|
|
# `FUNCTION` is the function name that will be called in the node.
|
|
FUNCTION = "load_model"
|
|
# `CATEGORY` is the category name that will be used when user searches the node.
|
|
CATEGORY = "CustomNodesTemplate"
|
|
|
|
# In the function, use same parameter names as you specified in `INPUT_TYPES()`
|
|
def load_model(self,
|
|
model: str,
|
|
device: str,
|
|
compute_type: Optional[str] = None,
|
|
) -> Tuple[CalculatorModel]:
|
|
calculator_model = CalculatorModel()
|
|
calculator_model.load_model(model, device, compute_type)
|
|
|
|
# You can use `comfy.utils.ProgressBar` to show the progress of the process.
|
|
# First, initialize the total amount of the process.
|
|
total_steps = 5
|
|
comfy_pbar = ProgressBar(total_steps)
|
|
# Then, update the progress.
|
|
for i in range(1, total_steps):
|
|
time.sleep(1)
|
|
comfy_pbar.update(i) # Alternatively, you can use `comfy_pbar.update_absolute(value)` to update the progress with absolute value.
|
|
|
|
# Return the model as a tuple.
|
|
return (calculator_model, )
|
|
|
|
|
|
# Second Node
|
|
class CalculatePlus:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ("MY_MODEL", ),
|
|
},
|
|
# Specify the parameters with type and default value.
|
|
"optional": {
|
|
"a": ("INT", {"default": 5}),
|
|
"b": ("INT", {"default": 10}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("INT",)
|
|
RETURN_NAMES = ("plus_value",)
|
|
FUNCTION = "plus"
|
|
CATEGORY = "CustomNodesTemplate"
|
|
|
|
def plus(self,
|
|
model: CalculatorModel,
|
|
a: Optional[int],
|
|
b: Optional[int],
|
|
) -> Tuple[int]:
|
|
result = model.plus(a, b)
|
|
return (result, )
|
|
|
|
|
|
|
|
# Third Node
|
|
class CalculateMinus:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ("MY_MODEL", ),
|
|
"a": ("INT", ),
|
|
},
|
|
"optional": {
|
|
"b": ("INT", {"default": 10}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("INT",)
|
|
RETURN_NAMES = ("minus_value",)
|
|
FUNCTION = "minus"
|
|
CATEGORY = "CustomNodesTemplate"
|
|
|
|
def minus(self,
|
|
model: CalculatorModel,
|
|
a: Optional[int],
|
|
b: Optional[int],
|
|
) -> Tuple[int]:
|
|
result = model.minus(a, b)
|
|
return (result, )
|
|
|
|
|
|
|
|
# Output Node
|
|
class ExampleOutputNode:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"value": ("INT", ),
|
|
},
|
|
}
|
|
|
|
# If the node is output node, set this to True.
|
|
OUTPUT_NODE = True
|
|
RETURN_TYPES = ("INT",)
|
|
RETURN_NAMES = ("int",)
|
|
FUNCTION = "result"
|
|
CATEGORY = "CustomNodesTemplate"
|
|
|
|
def result(self,
|
|
value: int,) -> Tuple[int]:
|
|
return (value, )
|