Fix NaN not allowed in newer requests versions

Refractor

README.md

@@ -47,7 +47,7 @@ It also allows using a workflow JSON as an input. To allow any workflow to run,
 
 I have nodes to save/load the workflows, but ideally there would be some nodes to also edit them - search and replace seed, etc. PRs welcome ;P
 
-Workflow JSON: [NetDistAdvanced.json](https://github.com/city96/ComfyUI_NetDist/files/13825337/NetDistAdvanced.json)
+Workflow JSON: [NetDistAdvancedV2.json](https://github.com/city96/ComfyUI_NetDist/files/13843005/NetDistAdvancedV2.json)
 
 ![NetDistAdvanced](https://github.com/city96/ComfyUI_NetDist/assets/125218114/851c1ee6-edcf-4489-bab1-92ab9c5ef15e)
 
@@ -55,7 +55,6 @@ Workflow JSON: [NetDistAdvanced.json](https://github.com/city96/ComfyUI_NetDist/
 
 ![NetDistSaved](https://github.com/city96/ComfyUI_NetDist/assets/125218114/a39b5117-af1b-4f2c-a94e-5a330acc8ea4)
 
-
 ### Remote images
 The `LoadImageUrl` ('Load Image (URL)') Node acts just like the normal 'Load Image' node.
 
@@ -65,9 +64,25 @@ The `SaveImageUrl` ('Save Image (URL)') Node sends a POST request to the target
 - The filenames are **not** guaranteed to be unique across batches since they aren't saved locally. You should handle this server-side.
 - No data is written to disk on the server.
 
+### Remote latents
+
+This node pack has a set of nodes which should (in theory) allow you to pass latents between the nodes seamlessly. A node to save the input latent as a `.npy` file is provided. This node also returns the filename of the saved latent, which can then be loaded by the other instance.
+
+To load a latent from the other instance, you can plug the filename into this URL:
+
+```
+# change the filename with a string replacement node.
+http://127.0.0.1:8188/view?filename=ComfyUI_00001_.latent&type=output`
+# To load them from the input folder instead, change type to 'input'
+http://127.0.0.1:8188/view?filename=TestLatent.npy&type=input
+```
+
+The `LoadLatentNumpy` node can also load the default safetensor latents, the npy ones (simple numpy file containing just the latent in the standard torch format) as well as the sd_scripts npz cache files.
+
+![LatentSave](https://github.com/city96/ComfyUI_NetDist/assets/125218114/cd68d8dc-bd96-4018-82c9-400337fc5f80)
 
 ### Things you probably shouldn't do:
-- Queue a workflow on the same client multiple times.
+- Queue a workflow on the same remote worker multiple times from the same client.
 - ~~Expect this to work smoothly.~~
 
 ## Roadmap

__init__.py

@@ -15,6 +15,9 @@ else:
 	from .nodes.images import NODE_CLASS_MAPPINGS as ImgNodes
 	NODE_CLASS_MAPPINGS.update(ImgNodes)
 
+	from .nodes.latents import NODE_CLASS_MAPPINGS as LatNodes
+	NODE_CLASS_MAPPINGS.update(LatNodes)
+
 	from .nodes.workflows import NODE_CLASS_MAPPINGS as WrkNodes
 	NODE_CLASS_MAPPINGS.update(WrkNodes)
 

core/dispatch.py

@@ -88,7 +88,7 @@ def dispatch_to_remote(remote_url, prompt, job_id=f"{get_client_id()}-unknown",
 			else:
 				prompt[i]["inputs"]["enabled"] = "false"
 
-	banned = [] if outputs == "any" else get_output_nodes(remote_url)
+	banned = [] if outputs == "any" else ["PreviewImage", "SaveImage"] # get_output_nodes(remote_url)
 	output = None
 	for i in prompt.keys():
 		# only leave current fetch but replace with PreviewImage
@@ -97,10 +97,10 @@ def dispatch_to_remote(remote_url, prompt, job_id=f"{get_client_id()}-unknown",
 				output = {
 					"inputs": {"images": prompt[i]["inputs"]["final_image"]},
 					"class_type": 'PreviewImage',
+					"final_output": True, # might allow multiple outputs with an ID?
 				}
 			recursive_node_deletion(i)
 		# do not save output on remote
-		# todo: other output types
 		if prompt[i]["class_type"] in banned:
 			recursive_node_deletion(i)
 	if output:
@@ -130,6 +130,11 @@ def dispatch_to_remote(remote_url, prompt, job_id=f"{get_client_id()}-unknown",
 			"job_id": job_id,
 		}
 	}
-	ar = requests.post(f"{remote_url}/prompt", json=data, timeout=4)
+	ar = requests.post(
+		f"{remote_url}/prompt",
+		data    = json.dumps(data),
+		headers = {"Content-Type": "application/json"},
+		timeout = 4,
+	)
 	ar.raise_for_status()
 	return

core/fetch.py

@@ -7,6 +7,14 @@ from PIL import Image
 
 POLLING = 0.5
 
+def get_job_output(inputs, outputs):
+	output_id = list(outputs.keys())[-1] # fallback to last
+	for i,d in inputs.items():
+		if d.get("final_output") and i in outputs.keys():
+			output_id = i
+			break
+	return outputs[output_id].get("images", [])
+
 def wait_for_job(remote_url, job_id):
 	fail = 0
 	while fail <= 3:
@@ -25,7 +33,7 @@ def wait_for_job(remote_url, job_id):
 			if d["prompt"][3].get("job_id") == job_id:
 				# this needs to be less jank
 				if len(d["outputs"].keys()) > 0:
-					return d["outputs"][list(d["outputs"].keys())[-1]].get("images")
+					return get_job_output(d["prompt"][2], d["outputs"])
 				else:
 					return []
 		# todo: check if it's actually in the queue to avoid waiting forever

nodes/latents.py

@@ -0,0 +1,170 @@
+import os
+import torch
+import requests
+import safetensors.torch
+import numpy as np
+from io import BytesIO
+
+import folder_paths
+
+class LoadLatentNumpy:
+	def __init__(self):
+		pass
+
+	@classmethod
+	def INPUT_TYPES(s):
+		exts = [".latent", ".safetensors", ".npy", ".npz"]
+		input_dir = folder_paths.get_input_directory()
+		files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))]
+		files = [f for f in files if any([f.endswith(x) for x in exts])]
+		return {
+			"required": {
+				"latent": [sorted(files), ]
+			},
+		}
+
+	RETURN_TYPES = ("LATENT",)
+	FUNCTION = "load"
+	CATEGORY = "remote/latent"
+	TITLE = "Load Latent (Numpy)"
+
+	def load_comfy(self, file):
+		# From default node - renamed safetensors file
+		if type(file) == str:
+			data = safetensors.torch.load_file(file)
+		else:
+			data = safetensors.torch.load(file)
+
+		latent = data["latent_tensor"].to(torch.float32)
+		if "latent_format_version_0" not in data:
+			latent *= 1.0 / 0.18215 # XL?
+		return latent
+
+	def load_numpy(self, file):
+		# plain npy file - saved as-is
+		return torch.from_numpy(np.load(file))
+
+	def load_koyha(self, file):
+		# generated by sd_scripts - npz
+		if "latents" in data.keys():
+			latent = data["latents"]
+		else:
+			latent = [x for x in data.items() if x.shape > 3][0]
+		return torch.from_numpy(latent)
+
+	def load(self, latent):
+		path = folder_paths.get_annotated_filepath(latent)
+		name, ext = os.path.splitext(latent)
+
+		if ext in [".latent", ".safetensors"]:
+			latent = self.load_comfy(path)
+		elif ext == ".npy":
+			latent = self.load_numpy(path)
+		elif ext == ".npz":
+			latent = self.load_koyha(path)
+		else:
+			try:
+				latent = self.load_numpy(path)
+			except:
+				raise ValueError(f"Unknown latent extension '{ext}'")
+
+		if len(latent.shape) == 3:
+			latent = latent.unsqueeze(0)
+		print("asdasd", latent.shape)
+
+		return ({"samples": latent.to(torch.float32)},)
+
+	@classmethod
+	def IS_CHANGED(s, latent):
+		image_path = folder_paths.get_annotated_filepath(latent)
+		m = hashlib.sha256()
+		with open(image_path, 'rb') as f:
+			m.update(f.read())
+		return m.digest().hex()
+
+	@classmethod
+	def VALIDATE_INPUTS(s, latent):
+		if not folder_paths.exists_annotated_filepath(latent):
+			return f"Invalid latent file '{latent}'"
+		return True
+
+class LoadLatentUrl(LoadLatentNumpy):
+	def __init__(self):
+		pass
+
+	@classmethod
+	def INPUT_TYPES(s):
+		return {
+			"required": {
+				"url": ("STRING", { "multiline": False, })
+			}
+		}
+
+	RETURN_TYPES = ("LATENT",)
+	TITLE = "Load Latent (URL)"
+
+	def load(self, url):
+		buffer = BytesIO()
+		with requests.get(url, stream=True, timeout=16) as r:
+			r.raise_for_status()
+			buffer.write(r.content)
+		buffer.seek(0)
+
+		if ".latent" in url or ".safetensors" in url:
+			latent = self.load_comfy(buffer)
+		elif ".npy" in url:
+			latent = self.load_numpy(buffer)
+		elif ".npz" in url:
+			latent = self.load_koyha(buffer)
+		else:
+			try:
+				latent = self.load_comfy(buffer)
+			except:
+				raise ValueError(f"Unknown latent extension '{url}'")
+
+		if len(latent.shape) == 3:
+			latent = latent.unsqueeze(0)
+
+		del buffer
+		return ({"samples": latent.to(torch.float32)},)
+
+	@classmethod
+	def IS_CHANGED(s, url):
+		return str(url)
+
+	@classmethod
+	def VALIDATE_INPUTS(s, url):
+		return True
+
+class SaveLatentNumpy:
+	def __init__(self):
+		self.output_dir = folder_paths.get_output_directory()
+
+	@classmethod
+	def INPUT_TYPES(s):
+		return {
+			"required": {
+				"samples": ("LATENT",),
+				"filename_prefix": ("STRING", {"default": "latents/ComfyUI"})
+			}
+		}
+
+	RETURN_TYPES = ("STRING",)
+	RETURN_NAMES = ("filename",)
+	OUTPUT_NODE = True
+	FUNCTION = "save"
+	CATEGORY = "remote/latent"
+	TITLE = "Save Latent (Numpy)"
+
+	def save(self, samples, filename_prefix="ComfyUI"):
+		full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir)
+		fname = f"{filename}_{counter:05}_.npy"
+		path = os.path.join(full_output_folder, fname)
+		np.save(path, samples["samples"].numpy())
+		return (fname,)
+
+NODE_CLASS_MAPPINGS = {
+	"LoadLatentNumpy" : LoadLatentNumpy,
+	"LoadLatentUrl"   : LoadLatentUrl,
+	"SaveLatentNumpy" : SaveLatentNumpy,
+}

requirements.txt

@@ -0,0 +1 @@
+requests>=2.28.2