From 49ea5435c7b85252dd3a729ac919b4087a18b25a Mon Sep 17 00:00:00 2001
From: Loner <yinjiakang@gmail.com>
Date: Wed, 13 Mar 2024 14:55:45 +0800
Subject: [PATCH] Update ComfyUI for private

---
 nodes.py             | 21 ++++++++++++-------
 web/scripts/app.js   | 50 ++++++++++++++++++++++++++------------------
 web/scripts/utils.js | 25 ++++++++++++++++++++++
 3 files changed, 68 insertions(+), 28 deletions(-)

diff --git a/nodes.py b/nodes.py
index 453f6e606..b28b39eb8 100644
--- a/nodes.py
+++ b/nodes.py
@@ -1907,22 +1907,27 @@ def load_custom_node(module_path, ignore=set()):
         logging.warning(f"Cannot import {module_path} module for custom nodes: {e}")
         return False
 
-def load_custom_nodes():
+def load_custom_nodes(public_mark: str = '_public/'):
     base_node_names = set(NODE_CLASS_MAPPINGS.keys())
     node_paths = folder_paths.get_folder_paths("custom_nodes")
     node_import_times = []
-    for custom_node_path in node_paths:
+    to_install_possible_modules = {}
+
+    for custom_node_path in sorted(node_paths, key=lambda x: public_mark in x, reverse=True):
         possible_modules = os.listdir(os.path.realpath(custom_node_path))
         if "__pycache__" in possible_modules:
             possible_modules.remove("__pycache__")
 
         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
-            if module_path.endswith(".disabled"): continue
-            time_before = time.perf_counter()
-            success = load_custom_node(module_path, base_node_names)
-            node_import_times.append((time.perf_counter() - time_before, module_path, success))
+            to_install_possible_modules[possible_module] = (custom_node_path, possible_module)
+
+    for custom_node_path, possible_module in to_install_possible_modules.values():
+        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
+        if module_path.endswith(".disabled"): continue
+        time_before = time.perf_counter()
+        success = load_custom_node(module_path, base_node_names)
+        node_import_times.append((time.perf_counter() - time_before, module_path, success))
 
     if len(node_import_times) > 0:
         logging.info("\nImport times for custom nodes:")
diff --git a/web/scripts/app.js b/web/scripts/app.js
index 77f29b8e5..28f1c4112 100644
--- a/web/scripts/app.js
+++ b/web/scripts/app.js
@@ -6,6 +6,7 @@ import { defaultGraph } from "./defaultGraph.js";
 import { getPngMetadata, getWebpMetadata, importA1111, getLatentMetadata } from "./pnginfo.js";
 import { addDomClippingSetting } from "./domWidget.js";
 import { createImageHost, calculateImageGrid } from "./ui/imagePreview.js"
+import { getWorkflow } from "./utils.js";
 
 export const ANIM_PREVIEW_WIDGET = "$$comfy_animation_preview"
 
@@ -1496,28 +1497,37 @@ export class ComfyApp {
 		await this.registerNodes();
 		initWidgets(this);
 
-		// Load previous workflow
-		let restored = false;
-		try {
-			const loadWorkflow = async (json) => {
-				if (json) {
-					const workflow = JSON.parse(json);
-					await this.loadGraphData(workflow);
-					return true;
-				}
-			};
-			const clientId = api.initialClientId ?? api.clientId;
-			restored =
-				(clientId && (await loadWorkflow(sessionStorage.getItem(`workflow:${clientId}`)))) ||
-				(await loadWorkflow(localStorage.getItem("workflow")));
-		} catch (err) {
-			console.error("Error loading previous workflow", err);
+		// Load prebuilt workflow
+		const workflow = await getWorkflow();
+
+		if (workflow) {
+			await this.loadGraphData(workflow);
+		}
+		else {
+			// Load previous workflow
+			let restored = false;
+			try {
+				const loadWorkflow = async (json) => {
+					if (json) {
+						const workflow = JSON.parse(json);
+						await this.loadGraphData(workflow);
+						return true;
+					}
+				};
+				const clientId = api.initialClientId ?? api.clientId;
+				restored =
+					(clientId && (await loadWorkflow(sessionStorage.getItem(`workflow:${clientId}`)))) ||
+					(await loadWorkflow(localStorage.getItem("workflow")));
+			} catch (err) {
+				console.error("Error loading previous workflow", err);
+			}
+
+			// We failed to restore a workflow so load the default
+			if (!restored) {
+				await this.loadGraphData();
+			}
 		}
 
-		// We failed to restore a workflow so load the default
-		if (!restored) {
-			await this.loadGraphData();
-		}
 
 		// Save current workflow automatically
 		setInterval(() => {
diff --git a/web/scripts/utils.js b/web/scripts/utils.js
index 01b988462..e80d17371 100644
--- a/web/scripts/utils.js
+++ b/web/scripts/utils.js
@@ -1,5 +1,30 @@
 import { $el } from "./ui.js";
 
+export function needLoadPrebuiltWorkflow(workflowId) {
+	var loaded = localStorage.getItem('PrebuiltWorkflowId' + workflowId);
+	if (loaded) {
+		return false
+	} else {
+		localStorage.setItem('PrebuiltWorkflowId' + workflowId, true);
+		return true
+	}
+}
+
+export async function getWorkflow() {
+	let flow_json = null;
+	const queryString = window.location.search;
+	const urlParams = new URLSearchParams(queryString);
+	const workflowId = urlParams.get('workflow');
+	if (workflowId && needLoadPrebuiltWorkflow(workflowId)) {
+		await fetch('../workflows/' + workflowId + '/' + workflowId + '.json').then(
+			response => {
+				flow_json = response.json()
+			}
+		)
+	}
+	return flow_json;
+}
+
 // Simple date formatter
 const parts = {
 	d: (d) => d.getDate(),