Cleanup

__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']

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"]

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))

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:

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()}