From d1138e8ba0a8fd007c4277b2a4bb89e1c6d82349 Mon Sep 17 00:00:00 2001
From: m957ymj75urz <m957ymj75urz@gmail.com>
Date: Sun, 19 Mar 2023 12:54:29 +0100
Subject: [PATCH] rework payload from server

---
 nodes.py               | 15 ++++++++++-----
 server.py              |  4 ++--
 web/scripts/app.js     |  7 ++-----
 web/scripts/widgets.js |  2 +-
 4 files changed, 15 insertions(+), 13 deletions(-)

diff --git a/nodes.py b/nodes.py
index b4901a85..f758c8d0 100644
--- a/nodes.py
+++ b/nodes.py
@@ -775,7 +775,7 @@ class KSamplerAdvanced:
 class SaveImage:
     def __init__(self):
         self.output_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "output")
-        self.url_suffix = ""
+        self.type = "output"
 
     @classmethod
     def INPUT_TYPES(s):
@@ -822,7 +822,7 @@ class SaveImage:
         if not os.path.exists(self.output_dir):
             os.makedirs(self.output_dir)
 
-        paths = list()
+        results = list()
         for image in images:
             i = 255. * image.cpu().numpy()
             img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
@@ -835,14 +835,19 @@ class SaveImage:
 
             file = f"{filename}_{counter:05}_.png"
             img.save(os.path.join(full_output_folder, file), pnginfo=metadata, optimize=True)
-            paths.append(os.path.join(subfolder, file + self.url_suffix))
+            results.append({
+                "filename": file,
+                "subfolder": subfolder,
+                "type": self.type
+            });
             counter += 1
-        return { "ui": { "images": paths } }
+        
+        return { "ui": { "images": results } }
 
 class PreviewImage(SaveImage):
     def __init__(self):
         self.output_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "temp")
-        self.url_suffix = "&type=temp"
+        self.type = "temp"
 
     @classmethod
     def INPUT_TYPES(s):
diff --git a/server.py b/server.py
index a58aeec7..94928ae2 100644
--- a/server.py
+++ b/server.py
@@ -111,7 +111,7 @@ class PromptServer():
 
         @routes.get("/view")
         async def view_image(request):
-            if "file" in request.rel_url.query:
+            if "filename" in request.rel_url.query:
                 type = request.rel_url.query.get("type", "output")
                 if type not in ["output", "input", "temp"]:
                     return web.Response(status=400)
@@ -123,7 +123,7 @@ class PromptServer():
                         return web.Response(status=403)
                     output_dir = full_output_dir
 
-                file = request.rel_url.query["file"]
+                file = request.rel_url.query["filename"]
                 file = os.path.basename(file)
                 file = os.path.join(output_dir, file)
 
diff --git a/web/scripts/app.js b/web/scripts/app.js
index 51e67b92..8d773b93 100644
--- a/web/scripts/app.js
+++ b/web/scripts/app.js
@@ -109,11 +109,8 @@ class ComfyApp {
 								return new Promise((r) => {
 									const img = new Image();
 									img.onload = () => r(img);
-									img.onerror = () => r(null);
-									
-									var filename = src.replace(/^.*[\\\/]/, '');
-									var subfolder = src.replace(filename, '');
-									img.src = "/view?file=" + filename + "&subfolder=" + subfolder;
+									img.onerror = () => r(null);									
+									img.src = "/view?" + new URLSearchParams(src).toString();
 								});
 							})
 						).then((imgs) => {
diff --git a/web/scripts/widgets.js b/web/scripts/widgets.js
index 55bdd8f1..86619552 100644
--- a/web/scripts/widgets.js
+++ b/web/scripts/widgets.js
@@ -141,7 +141,7 @@ export const ComfyWidgets = {
 				node.imgs = [img];
 				app.graph.setDirtyCanvas(true);
 			};
-			img.src = `/view/${name}?type=input`;
+			img.src = `/view?filename=${name}&type=input`;
 		}
 
 		// Add our own callback to the combo widget to render an image when it changes