lets hope its this simple
This commit is contained in:
softcookiepp
+3
-5
@@ -1,12 +1,10 @@
|
||||
from .nodes import *
|
||||
from .common import *
|
||||
from .torchvulkan_loader import *
|
||||
|
||||
|
||||
# Map all your custom nodes classes with the names that will be displayed in the UI.
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"(Down)Load My Model": MyModelLoader,
|
||||
"Calculate Plus": CalculatePlus,
|
||||
"Calculate Minus": CalculateMinus,
|
||||
"Example Output Node": ExampleOutputNode,
|
||||
"Torchvulkan Loader": TorchVulkanLoader
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -1,152 +0,0 @@
|
||||
# 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, )
|
||||
+6
-6
@@ -1,17 +1,17 @@
|
||||
# Edit the following fields to match your project
|
||||
|
||||
[project]
|
||||
name = "comfyui-customnodes-template"
|
||||
description = "Template repository for custom nodes from [ComfyUI-CustomNodes-Template](https://github.com/jhj0517/ComfyUI-CustomNodes-Template)."
|
||||
name = "torchvulkan-comfyui-node"
|
||||
description = "Node for using the out-of-tree torchvulkan backend with comfyui"
|
||||
version = "1.0.0"
|
||||
license = {file = "LICENSE"}
|
||||
dependencies = ["",] # e.g.) ["numpy==1.18.1", ]
|
||||
dependencies = ["torchvulkan"] # e.g.) ["numpy==1.18.1", ]
|
||||
|
||||
[project.urls]
|
||||
Repository = "https://github.com/jhj0517/ComfyUI-CustomNodes-Template"
|
||||
Repository = "https://github.com/softcookiepp/torchvulkan-comfyui-node"
|
||||
# Used by Comfy Registry https://comfyregistry.org
|
||||
|
||||
[tool.comfy]
|
||||
PublisherId = "jhj0517"
|
||||
DisplayName = "ComfyUI-CustomNodes-Template"
|
||||
PublisherId = "softcookiepp"
|
||||
DisplayName = "torchvulkan-comfyui-node"
|
||||
Icon = ""
|
||||
|
||||
+5
-1
@@ -1,3 +1,7 @@
|
||||
# Update your requirements.
|
||||
# numpy==1.18.1
|
||||
|
||||
numpy
|
||||
torch
|
||||
torchaudio
|
||||
torchvision
|
||||
torchvulkan
|
||||
|
||||
@@ -0,0 +1,33 @@
|
||||
from .common import *
|
||||
|
||||
class TorchVulkanLoader:
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(self):
|
||||
devices = [i for i in range(torchvulkan.device_count())]
|
||||
return {
|
||||
"required": {
|
||||
"device": (devices, {"default": "default ( local-task:// )"})
|
||||
},
|
||||
"optional": {
|
||||
"model": ("MODEL",),
|
||||
"clip": ("CLIP",),
|
||||
"vae": ("VAE",)
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("MODEL", "CLIP", "VAE")
|
||||
|
||||
FUNCTION = "load_to_torchvulkan"
|
||||
|
||||
CATEGORY = "Backends"
|
||||
|
||||
def load_to_torchvulkan(self, device, model = None, clip = None, vae = None):
|
||||
device = torch.device(f"vulkan:{device}")
|
||||
if model: model.to(device)
|
||||
if clip: clip.to(device)
|
||||
if vae: vae.to(device)
|
||||
return model, clip, vae
|
||||
|
||||
Reference in New Issue
Block a user