From 313f1f83a6f41ccff589663850c22c4e71e2819f Mon Sep 17 00:00:00 2001
From: pythongosssss <125205205+pythongosssss@users.noreply.github.com>
Date: Sat, 1 Apr 2023 12:44:29 +0100
Subject: [PATCH] Tweak server/custom node load order  - Load custom nodes
 after creating server  - Add routes after loading custom nodes Custom nodes
 can now add routes via PromptServer.instance

---
 main.py   | 4 +++-
 nodes.py  | 6 +++---
 server.py | 6 ++++--
 3 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/main.py b/main.py
index c9809137..824530fb 100644
--- a/main.py
+++ b/main.py
@@ -40,6 +40,7 @@ if __name__ == "__main__":
     except:
         pass
 
+from nodes import init_custom_nodes
 import execution
 import server
 import folder_paths
@@ -98,6 +99,8 @@ if __name__ == "__main__":
     server = server.PromptServer(loop)
     q = execution.PromptQueue(server)
 
+    init_custom_nodes()
+    server.add_routes()
     hijack_progress(server)
 
     threading.Thread(target=prompt_worker, daemon=True, args=(q,server,)).start()
@@ -113,7 +116,6 @@ if __name__ == "__main__":
     except:
         address = '127.0.0.1'
 
-
     dont_print = False
     if '--dont-print-server' in sys.argv:
         dont_print = True
diff --git a/nodes.py b/nodes.py
index 6fb7f017..b422f2cb 100644
--- a/nodes.py
+++ b/nodes.py
@@ -1050,6 +1050,6 @@ def load_custom_nodes():
         if os.path.isfile(module_path) and os.path.splitext(module_path)[1] != ".py": continue
         load_custom_node(module_path)
 
-load_custom_nodes()
-
-load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_upscale_model.py"))
\ No newline at end of file
+def init_custom_nodes():
+    load_custom_nodes()
+    load_custom_node(os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy_extras"), "nodes_upscale_model.py"))
\ No newline at end of file
diff --git a/server.py b/server.py
index 80fb2dc7..963daeff 100644
--- a/server.py
+++ b/server.py
@@ -42,6 +42,7 @@ class PromptServer():
         self.web_root = os.path.join(os.path.dirname(
             os.path.realpath(__file__)), "web")
         routes = web.RouteTableDef()
+        self.routes = routes
         self.last_node_id = None
         self.client_id = None
 
@@ -239,8 +240,9 @@ class PromptServer():
                     self.prompt_queue.delete_history_item(id_to_delete)
 
             return web.Response(status=200)
-
-        self.app.add_routes(routes)
+        
+    def add_routes(self):
+        self.app.add_routes(self.routes)
         self.app.add_routes([
             web.static('/', self.web_root),
         ])