From 195d7aec9f9c170ea94643ea71333ef8d1a6b84c Mon Sep 17 00:00:00 2001
From: m957ymj75urz <m957ymj75urz@gmail.com>
Date: Sun, 12 Mar 2023 19:51:39 +0100
Subject: [PATCH] can specify a subfolder in the SaveImage node

---
 nodes.py           | 16 ++++++++++------
 server.py          |  9 ++++++---
 web/scripts/app.js |  4 +++-
 3 files changed, 19 insertions(+), 10 deletions(-)

diff --git a/nodes.py b/nodes.py
index 0a0a0a9c..d7708e34 100644
--- a/nodes.py
+++ b/nodes.py
@@ -780,7 +780,8 @@ class SaveImage:
     def INPUT_TYPES(s):
         return {"required": 
                     {"images": ("IMAGE", ),
-                     "filename_prefix": ("STRING", {"default": "ComfyUI"})},
+                     "filename_prefix": ("STRING", {"default": "ComfyUI"}),
+                     "subfolder": ("STRING", {})},
                 "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
                 }
 
@@ -791,7 +792,7 @@ class SaveImage:
 
     CATEGORY = "image"
 
-    def save_images(self, images, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None):
+    def save_images(self, images, filename_prefix="ComfyUI", subfolder=None, prompt=None, extra_pnginfo=None):
         def map_filename(filename):
             prefix_len = len(filename_prefix)
             prefix = filename[:prefix_len + 1]
@@ -800,12 +801,15 @@ class SaveImage:
             except:
                 digits = 0
             return (digits, prefix)
+        
+        outputDirectory = os.path.join(os.path.dirname(os.path.realpath(__file__)), "output", subfolder);
+        
         try:
-            counter = max(filter(lambda a: a[1][:-1] == filename_prefix and a[1][-1] == "_", map(map_filename, os.listdir(self.output_dir))))[0] + 1
+            counter = max(filter(lambda a: a[1][:-1] == filename_prefix and a[1][-1] == "_", map(map_filename, os.listdir(outputDirectory))))[0] + 1
         except ValueError:
             counter = 1
         except FileNotFoundError:
-            os.mkdir(self.output_dir)
+            os.mkdir(outputDirectory)
             counter = 1
 
         paths = list()
@@ -819,8 +823,8 @@ class SaveImage:
                 for x in extra_pnginfo:
                     metadata.add_text(x, json.dumps(extra_pnginfo[x]))
             file = f"{filename_prefix}_{counter:05}_.png"
-            img.save(os.path.join(self.output_dir, file), pnginfo=metadata, optimize=True)
-            paths.append(file)
+            img.save(os.path.join(outputDirectory, file), pnginfo=metadata, optimize=True)
+            paths.append(os.path.join(subfolder, file))
             counter += 1
         return { "ui": { "images": paths } }
 
diff --git a/server.py b/server.py
index 5aba5761..05d1af04 100644
--- a/server.py
+++ b/server.py
@@ -109,15 +109,18 @@ class PromptServer():
                 return web.Response(status=400)
 
 
-        @routes.get("/view/{file}")
+        @routes.get("/view")
         async def view_image(request):
-            if "file" in request.match_info:
+            if "file" in request.rel_url.query:
                 type = request.rel_url.query.get("type", "output")
                 if type != "output" and type != "input":
                     return web.Response(status=400)
 
                 output_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), type)
-                file = request.match_info["file"]
+                if "subfolder" in request.rel_url.query:
+                    output_dir = os.path.join(output_dir, request.rel_url.query["subfolder"])
+
+                file = request.rel_url.query["file"]
                 file = os.path.basename(file)
                 file = os.path.join(output_dir, file)
                 if os.path.isfile(file):
diff --git a/web/scripts/app.js b/web/scripts/app.js
index e70e1c15..931eb1f7 100644
--- a/web/scripts/app.js
+++ b/web/scripts/app.js
@@ -110,7 +110,9 @@ class ComfyApp {
 									const img = new Image();
 									img.onload = () => r(img);
 									img.onerror = () => r(null);
-									img.src = "/view/" + src;
+									var filename = src.replace(/^.*[\\\/]/, '');
+									var subfolder = src.replace(filename, '');
+									img.src = "/view?file=" + filename + "&subfolder=" + subfolder;
 								});
 							})
 						).then((imgs) => {