From 592b377ac46ea931bfb6311cced636b1c68f103f Mon Sep 17 00:00:00 2001
From: pythongosssss <125205205+pythongosssss@users.noreply.github.com>
Date: Fri, 3 Mar 2023 19:05:39 +0000
Subject: [PATCH] Added dynamic loading of extensions

---
 server.py                         |  6 ++++++
 web/extensions/logging.js.example |  1 +
 web/index.html                    |  3 ---
 web/scripts/api.js                |  9 +++++++++
 web/scripts/app.js                | 28 ++++++++++++++++++----------
 5 files changed, 34 insertions(+), 13 deletions(-)

diff --git a/server.py b/server.py
index cfc103cb..64b2f70a 100644
--- a/server.py
+++ b/server.py
@@ -5,6 +5,7 @@ import nodes
 import execution
 import uuid
 import json
+import glob
 
 try:
     import aiohttp
@@ -52,6 +53,11 @@ class PromptServer():
         async def get_root(request):
             return web.FileResponse(os.path.join(self.web_root, "index.html"))
 
+        @routes.get("/extensions")
+        async def get_extensions(request):
+            files = glob.glob(os.path.join(self.web_root, 'extensions/**/*.js'), recursive=True)
+            return web.json_response(list(map(lambda f: "/" + os.path.relpath(f, self.web_root).replace("\\", "/"), files)))
+
         @routes.get("/view/{file}")
         async def view_image(request):
             if "file" in request.match_info:
diff --git a/web/extensions/logging.js.example b/web/extensions/logging.js.example
index ce0b346e..d015096a 100644
--- a/web/extensions/logging.js.example
+++ b/web/extensions/logging.js.example
@@ -1,6 +1,7 @@
 import { app } from "../scripts/app.js";
 
 const ext = {
+	// Unique name for the extension
 	name: "Example.LoggingExtension",
 	async init(app) {
 		// Any initial setup to run as soon as the page loads
diff --git a/web/index.html b/web/index.html
index 1236d754..cd95594f 100644
--- a/web/index.html
+++ b/web/index.html
@@ -6,9 +6,6 @@
 		<script type="text/javascript" src="lib/litegraph.core.js"></script>
 
 		<script type="module">
-			import "/extensions/core/dynamicPrompts.js";
-			import "/extensions/core/rerouteNode.js";
-
 			import { app } from "/scripts/app.js";
 			await app.setup();
 			window.app = app;
diff --git a/web/scripts/api.js b/web/scripts/api.js
index 1b3bbecd..6ffff4a0 100644
--- a/web/scripts/api.js
+++ b/web/scripts/api.js
@@ -88,6 +88,15 @@ class ComfyApi extends EventTarget {
 		this.#createSocket();
 	}
 
+	/**
+	 * Gets a list of extension urls
+	 * @returns An array of script urls to import
+	 */
+	async getExtensions() {
+		const resp = await fetch("/extensions", { cache: "no-store" });
+		return await resp.json();
+	}
+
 	/**
 	 * Loads node object definitions for the graph
 	 * @returns The node definitions
diff --git a/web/scripts/app.js b/web/scripts/app.js
index 071742c9..3045d436 100644
--- a/web/scripts/app.js
+++ b/web/scripts/app.js
@@ -11,14 +11,6 @@ class ComfyApp {
 		this.nodeOutputs = {};
 	}
 
-	#log(message, ...other) {
-		console.log("[comfy]", message, ...other);
-	}
-
-	#error(message, ...other) {
-		console.error("[comfy]", message, ...other);
-	}
-
 	/**
 	 * Invoke an extension callback
 	 * @param {string} method The extension callback to execute
@@ -32,7 +24,7 @@ class ComfyApp {
 				try {
 					results.push(ext[method](...args, this));
 				} catch (error) {
-					this.#error(
+					console.error(
 						`Error calling extension '${ext.name}' method '${method}'`,
 						{ error },
 						{ extension: ext },
@@ -58,7 +50,7 @@ class ComfyApp {
 					try {
 						return await ext[method](...args, this);
 					} catch (error) {
-						this.#error(
+						console.error(
 							`Error calling extension '${ext.name}' method '${method}'`,
 							{ error },
 							{ extension: ext },
@@ -413,10 +405,26 @@ class ComfyApp {
 		api.init();
 	}
 
+	/**
+	 * Loads all extensions from the API into the window
+	 */
+	async #loadExtensions() {
+		const extensions = await api.getExtensions();
+		for (const ext of extensions) {
+			try {
+				await import(ext);
+			} catch (error) {
+				console.error("Error loading extension", ext, error);
+			}
+		}
+	}
+
 	/**
 	 * Set up the app on the page
 	 */
 	async setup() {
+		await this.#loadExtensions();
+
 		// Create and mount the LiteGraph in the DOM
 		const canvasEl = Object.assign(document.createElement("canvas"), { id: "graph-canvas" });
 		document.body.prepend(canvasEl);