From 5dd1fa7410b1e4d897d8718aee35214b2905c1e1 Mon Sep 17 00:00:00 2001
From: City <125218114+city96@users.noreply.github.com>
Date: Sat, 7 Oct 2023 16:56:12 +0200
Subject: [PATCH] Cleanup

---
 __init__.py                             | 56 ++++---------------------
 nodes/{remote_control.py => control.py} |  4 ++
 nodes/{remote_images.py => images.py}   |  2 +
 nodes/misc.py                           |  1 +
 nodes/nodes.py                          | 14 +++++++
 5 files changed, 29 insertions(+), 48 deletions(-)
 rename nodes/{remote_control.py => control.py} (97%)
 rename nodes/{remote_images.py => images.py} (94%)
 create mode 100644 nodes/nodes.py

diff --git a/__init__.py b/__init__.py
index ffdd15c..3f9919a 100644
--- a/__init__.py
+++ b/__init__.py
@@ -1,48 +1,8 @@
-NODE_CLASS_MAPPINGS = {}
-NODE_DISPLAY_NAME_MAPPINGS = {}
-
-def remote_control():
-	global NODE_CLASS_MAPPINGS
-	global NODE_DISPLAY_NAME_MAPPINGS
-	from .nodes.remote_control import QueueRemoteChainStart, QueueRemoteChainEnd, QueueRemote, FetchRemote
-	NODE_CLASS_MAPPINGS.update({
-		"QueueRemoteChainStart": QueueRemoteChainStart,
-		"QueueRemoteChainEnd": QueueRemoteChainEnd,
-		"QueueRemote": QueueRemote,
-		"FetchRemote": FetchRemote,
-	})
-	NODE_DISPLAY_NAME_MAPPINGS.update({
-		"QueueRemoteChainStart": "Queue on remote (start of chain)",
-		"QueueRemoteChainEnd": "Queue on remote (end of chain)",
-		"QueueRemote": "Queue on remote",
-		"FetchRemote": "Fetch from remote",
-	})
-
-def remote_images():
-	global NODE_CLASS_MAPPINGS
-	global NODE_DISPLAY_NAME_MAPPINGS
-	from .nodes.remote_images import LoadImageUrl, SaveImageUrl
-	NODE_CLASS_MAPPINGS.update({
-		"LoadImageUrl": LoadImageUrl,
-		"SaveImageUrl": SaveImageUrl,
-	})
-	NODE_DISPLAY_NAME_MAPPINGS.update({
-		"LoadImageUrl": "Load Image (URL)",
-		"SaveImageUrl": "Save Image (URL)",
-	})
-
-def remote_misc():
-	global NODE_CLASS_MAPPINGS
-	global NODE_DISPLAY_NAME_MAPPINGS
-	from .nodes.misc import CombineImageBatch
-	NODE_CLASS_MAPPINGS.update({
-		"CombineImageBatch": CombineImageBatch,
-	})
-	NODE_DISPLAY_NAME_MAPPINGS.update({
-		"CombineImageBatch": "Combine images",
-	})
-
-print("Loading network distribution node pack")
-remote_control()
-remote_images()
-remote_misc()
+# only import if running as a custom node
+try:
+	import comfy.utils
+except ImportError:
+	pass
+else:
+	from .nodes.nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+	__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
diff --git a/nodes/remote_control.py b/nodes/control.py
similarity index 97%
rename from nodes/remote_control.py
rename to nodes/control.py
index b72a380..65523c5 100644
--- a/nodes/remote_control.py
+++ b/nodes/control.py
@@ -24,6 +24,7 @@ class FetchRemote():
 	RETURN_TYPES = ("IMAGE",)
 	FUNCTION = "get_remote_job"
 	CATEGORY = "remote"
+	TITLE = "Fetch from remote"
 
 	def wait_for_job(self,remote_url,job_id):
 		url = remote_url + "history"
@@ -94,6 +95,7 @@ class QueueRemoteChainStart:
 	RETURN_NAMES = ("remote_chain_start",)
 	FUNCTION = "chain_start"
 	CATEGORY = "remote"
+	TITLE = "Queue on remote (start of chain)"
 
 	def chain_start(self, workflow, trigger, batch, seed, prompt):
 		remote_chain = {
@@ -128,6 +130,7 @@ class QueueRemoteChainEnd:
 	RETURN_NAMES = ("seed", "batch")
 	FUNCTION = "chain_end"
 	CATEGORY = "remote"
+	TITLE = "Queue on remote (end of chain)"
 
 	def chain_end(self, remote_chain_end):
 		seed = remote_chain_end["current_seed"]
@@ -163,6 +166,7 @@ class QueueRemote:
 	RETURN_NAMES = ("remote_chain", "remote_info")
 	FUNCTION = "queue_on_remote"
 	CATEGORY = "remote"
+	TITLE = "Queue on remote"
 
 	def queue_on_remote(self, remote_chain, remote_url, system, batch_override, enabled):
 		batch = batch_override if batch_override > 0 else remote_chain["batch"]
diff --git a/nodes/remote_images.py b/nodes/images.py
similarity index 94%
rename from nodes/remote_images.py
rename to nodes/images.py
index ca8d32a..52b5afc 100644
--- a/nodes/remote_images.py
+++ b/nodes/images.py
@@ -23,6 +23,7 @@ class LoadImageUrl:
 	RETURN_TYPES = ("IMAGE", "MASK")
 	FUNCTION = "load_image_url"
 	CATEGORY = "remote"
+	TITLE = "Load Image (URL)"
 
 	def load_image_url(self, url):
 		with requests.get(url, stream=True) as r:
@@ -58,6 +59,7 @@ class SaveImageUrl:
 	OUTPUT_NODE = True
 	FUNCTION = "save_images"
 	CATEGORY = "remote"
+	TITLE = "Save Image (URL)"
 	
 	def save_images(self, images, url, data_format, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None):
 		filename = os.path.basename(os.path.normpath(filename_prefix))
diff --git a/nodes/misc.py b/nodes/misc.py
index 6e74af7..e490bac 100644
--- a/nodes/misc.py
+++ b/nodes/misc.py
@@ -18,6 +18,7 @@ class CombineImageBatch:
 	RETURN_NAMES = ("images",)
 	FUNCTION = "combine_images"
 	CATEGORY = "remote"
+	TITLE = "Combine images"
 
 	def combine_images(self,images_a,images_b):
 		try:
diff --git a/nodes/nodes.py b/nodes/nodes.py
new file mode 100644
index 0000000..64e21b3
--- /dev/null
+++ b/nodes/nodes.py
@@ -0,0 +1,14 @@
+from .control import QueueRemoteChainStart, QueueRemoteChainEnd, QueueRemote, FetchRemote
+from .images import LoadImageUrl, SaveImageUrl
+from .misc import CombineImageBatch
+
+NODE_CLASS_MAPPINGS = {
+	"QueueRemoteChainStart": QueueRemoteChainStart,
+	"QueueRemoteChainEnd": QueueRemoteChainEnd,
+	"QueueRemote": QueueRemote,
+	"FetchRemote": FetchRemote,
+	"LoadImageUrl": LoadImageUrl,
+	"SaveImageUrl": SaveImageUrl,
+	"CombineImageBatch": CombineImageBatch,
+}
+NODE_DISPLAY_NAME_MAPPINGS = {k:v.TITLE for k,v in NODE_CLASS_MAPPINGS.items()}