From 64c8d6db647e2384101f3c761dcb62835703a8c5 Mon Sep 17 00:00:00 2001
From: City <125218114+city96@users.noreply.github.com>
Date: Fri, 5 Jan 2024 15:17:03 +0100
Subject: [PATCH] Fix output selection logic

---
 README.md        |  2 +-
 core/dispatch.py |  4 ++--
 core/fetch.py    | 10 +++++++++-
 3 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index e8cc3cc..34b9754 100644
--- a/README.md
+++ b/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)
 
diff --git a/core/dispatch.py b/core/dispatch.py
index 661b001..bb3dbb3 100644
--- a/core/dispatch.py
+++ b/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:
diff --git a/core/fetch.py b/core/fetch.py
index 96ed4ae..7c64f21 100644
--- a/core/fetch.py
+++ b/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