feat/startup-script: Feature to avoid package installation errors when installing custom nodes. (Comfy-Org#856)

ltdrdata · web-flow · commit 99abcbef41b5 · 2023-07-11T02:33:21.000-04:00
* support startup script for installation without locking on windows

* modified: Instead of executing scripts from the startup-scripts directory, I will change it to execute the prestartup_script.py for each custom node.
diff --git a/.gitignore b/.gitignore
@@ -13,4 +13,5 @@ extra_model_paths.yaml
 venv/
 web/extensions/*
 !web/extensions/logging.js.example
-!web/extensions/core/
+!web/extensions/core/
+startup-scripts/
diff --git a/main.py b/main.py
@@ -1,6 +1,38 @@
+import os
+import importlib.util
+import folder_paths
+
+
+def execute_prestartup_script():
+    def execute_script(script_path):
+        if os.path.exists(script_path):
+            module_name = os.path.splitext(script_path)[0]
+            try:
+                spec = importlib.util.spec_from_file_location(module_name, script_path)
+                module = importlib.util.module_from_spec(spec)
+                spec.loader.exec_module(module)
+            except Exception as e:
+                print(f"Failed to execute startup-script: {script_path} / {e}")
+
+    node_paths = folder_paths.get_folder_paths("custom_nodes")
+    for custom_node_path in node_paths:
+        possible_modules = os.listdir(custom_node_path)
+
+        for possible_module in possible_modules:
+            module_path = os.path.join(custom_node_path, possible_module)
+            if os.path.isfile(module_path) or module_path.endswith(".disabled") or module_path == "__pycache__":
+                continue
+
+            script_path = os.path.join(module_path, "prestartup_script.py")
+            execute_script(script_path)
+
+
+execute_prestartup_script()
+
+
+# Main code
 import asyncio
 import itertools
-import os
 import shutil
 import threading
 import gc
@@ -22,7 +54,6 @@
 import yaml
 
 import execution
-import folder_paths
 import server
 from server import BinaryEventTypes
 from nodes import init_custom_nodes
