Adding some suggests from the author

2 years ago · 493671574e
5 changed files with 188 additions and 3 deletions
--- a/custom_nodes/example.py
+++ b/custom_nodes/example.py
@ -0,0 +1,87 @@
+class Example:
+    """
+    A example node
+
+    Class methods
+    -------------
+    INPUT_TYPES (dict): 
+        Tell the main program input parameters of nodes.
+
+    Attributes
+    ----------
+    RETURN_TYPES (`tuple`): 
+        The type of each element in the output tulple.
+    FUNCTION (`str`):
+        The name of the entry-point method which will return a tuple. For example, if `FUNCTION = "execute"` then it will run Example().execute()
+    OUTPUT_NODE ([`bool`]):
+        WIP
+    CATEGORY (`str`):
+        WIP
+    execute(s) -> tuple || None:
+        The entry point method. The name of this method must be the same as the value of property `FUNCTION`.
+        For example, if `FUNCTION = "execute"` then this method's name must be `execute`, if `FUNCTION = "foo"` then it must be `foo`.
+    """
+    def __init__(self):
+        pass
+    
+    @classmethod
+    def INPUT_TYPES(s):
+        """
+            Return a dictionary which contains config for all input fields.
+            The type can be a string indicate a type or a list indicate selection.
+            Prebuilt types (string): "MODEL", "VAE", "CLIP", "CONDITIONING", "LATENT", "IMAGE", "INT", "STRING", "FLOAT".
+            Input in type "INT", "STRING" or "FLOAT" will be converted automatically from a string to the corresponse Python type before passing and have special config
+            Argument: s (`None`): Useless ig
+            Returns: `dict`:
+                - Key input_fields_group (`string`): Can be either required, hidden or optional. A node class must have property `required`
+                - Value input_fields (`dict`): Contains input fields config:
+                    * Key field_name (`string`): Name of a entry-point method's argument
+                    * Value field_config (`tuple`):
+                        + First value is a string indicate the type of field or a list for selection.
+                        + Secound value is a config for type "INT", "STRING" or "FLOAT".
+        """
+        return {
+            "required": {
+                "string_field": ("STRING", {
+                    "multiline": True, #Allow the input to be multilined
+                    "default": "Hello World!"
+                }),
+                "int_field": ("INT", {
+                    "default": 0, 
+                    "min": 0, #Minimum value
+                    "max": 4096, #Maximum value
+                    "step": 64 #Slider's step
+                }),
+                #Like INT
+                "float_field": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.01}),
+                "print_to_screen": (["Enable", "Disable"], {"default": "Enable"})
+            },
+            #"hidden": {
+            #    "prompt": "PROMPT", 
+            #    "extra_pnginfo": "EXTRA_PNGINFO"
+            #},
+        }
+
+    RETURN_TYPES = ("STRING", "INT", "FLOAT", "STRING")
+    FUNCTION = "test"
+
+    #OUTPUT_NODE = True
+
+    CATEGORY = "Example"
+
+    def test(self, string_field, int_field, float_field, print_to_screen):
+        if print_to_screen == "Enable":
+            print(f"""Your input contains:
+                string_field aka input text: {string_field}
+                int_field: {int_field}
+                float_field: {float_field}
+            """)
+        return (string_field, int_field, float_field, print_to_screen)
+
+
+NODE_CLASS_MAPPINGS = {
+    "Example": Example
+}
+"""
+NODE_CLASS_MAPPINGS (dict): A dictionary contains all nodes you want to export
+"""
--- a/custom_nodes/example_folder/init.py
+++ b/custom_nodes/example_folder/init.py
@ -0,0 +1,87 @@
+from utils import waste_cpu_resource
+class ExampleFolder:
+    """
+    A example node
+
+    Class methods
+    -------------
+    INPUT_TYPES (dict): 
+        Tell the main program input parameters of nodes.
+
+    Attributes
+    ----------
+    RETURN_TYPES (`tuple`): 
+        The type of each element in the output tulple.
+    FUNCTION (`str`):
+        The name of the entry-point method which will return a tuple. For example, if `FUNCTION = "execute"` then it will run Example().execute()
+    OUTPUT_NODE ([`bool`]):
+        WIP
+    CATEGORY (`str`):
+        WIP
+    execute(s) -> tuple || None:
+        The entry point method. The name of this method must be the same as the value of property `FUNCTION`.
+        For example, if `FUNCTION = "execute"` then this method's name must be `execute`, if `FUNCTION = "foo"` then it must be `foo`.
+    """
+    def __init__(self):
+        pass
+    
+    @classmethod
+    def INPUT_TYPES(s):
+        """
+            Return a dictionary which contains config for all input fields.
+            The type can be a string indicate a type or a list indicate selection.
+            Prebuilt types (string): "MODEL", "VAE", "CLIP", "CONDITIONING", "LATENT", "IMAGE", "INT", "STRING", "FLOAT".
+            Input in type "INT", "STRING" or "FLOAT" will be converted automatically from a string to the corresponse Python type before passing and have special config
+            Argument: s (`None`): Useless ig
+            Returns: `dict`:
+                - Key input_fields_group (`string`): Can be either required, hidden or optional. A node class must have property `required`
+                - Value input_fields (`dict`): Contains input fields config:
+                    * Key field_name (`string`): Name of a entry-point method's argument
+                    * Value field_config (`tuple`):
+                        + First value is a string indicate the type of field or a list for selection.
+                        + Secound value is a config for type "INT", "STRING" or "FLOAT".
+        """
+        return {
+            "required": {
+                "string_field": ("STRING", {
+                    "multiline": True, #Allow the input to be multilined
+                    "default": "Hello World!"
+                }),
+                "int_field": ("INT", {
+                    "default": 0, 
+                    "min": 0, #Minimum value
+                    "max": 4096, #Maximum value
+                    "step": 64 #Slider's step
+                }),
+                #Like INT
+                "print_to_screen": (["Enable", "Disable"], {"default": "Enable"})
+            },
+            #"hidden": {
+            #    "prompt": "PROMPT", 
+            #    "extra_pnginfo": "EXTRA_PNGINFO"
+            #},
+        }
+
+    RETURN_TYPES = ("STRING", "INT", "FLOAT", "STRING")
+    FUNCTION = "test"
+
+    #OUTPUT_NODE = True
+
+    CATEGORY = "Example"
+
+    def test(self, string_field, int_field, print_to_screen):
+        rand_float = waste_cpu_resource()
+        if print_to_screen == "Enable":
+            print(f"""Your input contains:
+                string_field aka input text: {string_field}
+                int_field: {int_field}
+                A random float number: {rand_float}
+            """)
+        return (string_field, int_field, rand_float, print_to_screen)
+
+NODE_CLASS_MAPPINGS = {
+    "ExampleFolder": ExampleFolder
+}
+"""
+NODE_CLASS_MAPPINGS (dict): A dictionary contains all nodes you want to export
+"""
--- a/custom_nodes/example_folder/utils.py
+++ b/custom_nodes/example_folder/utils.py
@ -0,0 +1,4 @@
+import torch
+def waste_cpu_resource():
+    x = torch.rand(1, 1e6, dtype=torch.float64).cpu()
+    return x.numpy()[0, 1]
--- a/custom_nodes/put_custom_nodes_here
+++ b/custom_nodes/put_custom_nodes_here
--- a/nodes.py
+++ b/nodes.py
@ -598,12 +598,19 @@ NODE_CLASS_MAPPINGS = {
    "CLIPLoader": CLIPLoader,
 }

+CUSTOM_NODE_PATH = os.path.join(os.path.dirname(os.path.realpath(__file__)), "custom_nodes")
 def load_custom_nodes():
-    possible_modules = os.listdir("custom_nodes")
-    possible_modules.remove("put_custom_nodes_here")
+    possible_modules = os.listdir(CUSTOM_NODE_PATH)
+    try:
+        possible_modules.remove("example.py")
+        possible_modules.remove("example_folder")
+    except ValueError: pass
+
    for possible_module in possible_modules:
+        module_path = os.path.join(CUSTOM_NODE_PATH, possible_module)
+        if os.path.isfile(module_path) and os.path.splitext(module_path)[1] != ".py": continue
        try:
-            custom_nodes = import_module(possible_module, "custom_nodes")
+            custom_nodes = import_module(possible_module, CUSTOM_NODE_PATH)
            if getattr(custom_nodes, "NODE_CLASS_MAPPINGS") is not None:
                NODE_CLASS_MAPPINGS.update(custom_nodes.NODE_CLASS_MAPPINGS)
            else: