From 5ce4b9eca3c601965f93852274b2e428b9a49179 Mon Sep 17 00:00:00 2001
From: Hacker 17082006 <tai17082006.tk@gmail.com>
Date: Mon, 13 Feb 2023 18:17:40 +0700
Subject: [PATCH] Add custom node feature

---
 custom_nodes/put_custom_nodes_here |  0
 nodes.py                           | 15 +++++++++++++++
 2 files changed, 15 insertions(+)
 create mode 100644 custom_nodes/put_custom_nodes_here

diff --git a/custom_nodes/put_custom_nodes_here b/custom_nodes/put_custom_nodes_here
new file mode 100644
index 00000000..e69de29b
diff --git a/nodes.py b/nodes.py
index aaa4f87a..8fa49932 100644
--- a/nodes.py
+++ b/nodes.py
@@ -16,6 +16,7 @@ sys.path.insert(0, os.path.join(sys.path[0], "comfy"))
 import comfy.samplers
 import comfy.sd
 import model_management
+from importlib import import_module
 
 supported_ckpt_extensions = ['.ckpt']
 supported_pt_extensions = ['.ckpt', '.pt', '.bin']
@@ -597,4 +598,18 @@ NODE_CLASS_MAPPINGS = {
     "CLIPLoader": CLIPLoader,
 }
 
+def load_custom_nodes():
+    possible_modules = os.listdir("custom_nodes")
+    possible_modules.remove("put_custom_nodes_here")
+    for possible_module in possible_modules:
+        try:
+            custom_nodes = import_module(possible_module, "custom_nodes")
+            if getattr(custom_nodes, "NODE_CLASS_MAPPINGS") is not None:
+                NODE_CLASS_MAPPINGS.update(custom_nodes.NODE_CLASS_MAPPINGS)
+            else:
+                NODE_CLASS_MAPPINGS[possible_module] = custom_nodes
+        except ImportError as e:
+            print(f"Cannot import {possible_module} node.")
+            print(e)
 
+load_custom_nodes()
\ No newline at end of file