Fix output selection logic

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)
 

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:

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