From 1fe6904ad77c2fd5fc4da0181e08d9e669d62bac Mon Sep 17 00:00:00 2001 From: softcookiepp Date: Mon, 2 Nov 2026 09:16:04 -0800 Subject: [PATCH] lets hope its this simple --- __init__.py | 8 +-- common.py | 2 + nodes.py | 152 ------------------------------------------ pyproject.toml | 12 ++-- requirements.txt | 6 +- torchvulkan_loader.py | 33 +++++++++ 6 files changed, 49 insertions(+), 164 deletions(-) create mode 100644 common.py delete mode 100644 nodes.py create mode 100644 torchvulkan_loader.py diff --git a/__init__.py b/__init__.py index a2a3159..889a6ef 100644 --- a/__init__.py +++ b/__init__.py @@ -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 } diff --git a/common.py b/common.py new file mode 100644 index 0000000..5686baf --- /dev/null +++ b/common.py @@ -0,0 +1,2 @@ +import torchvulkan +import torch diff --git a/nodes.py b/nodes.py deleted file mode 100644 index 9f507e2..0000000 --- a/nodes.py +++ /dev/null @@ -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, ) diff --git a/pyproject.toml b/pyproject.toml index 9364ed6..b2b0ce8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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 = "" diff --git a/requirements.txt b/requirements.txt index aa145a9..68efc92 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,7 @@ # Update your requirements. # numpy==1.18.1 - +numpy +torch +torchaudio +torchvision +torchvulkan diff --git a/torchvulkan_loader.py b/torchvulkan_loader.py new file mode 100644 index 0000000..54ff4ac --- /dev/null +++ b/torchvulkan_loader.py @@ -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 +