From 195d7aec9f9c170ea94643ea71333ef8d1a6b84c Mon Sep 17 00:00:00 2001 From: m957ymj75urz Date: Sun, 12 Mar 2023 19:51:39 +0100 Subject: [PATCH 1/9] 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) => { From 70f9bd1408ef58733f6072e549cbf9a9e97bed75 Mon Sep 17 00:00:00 2001 From: m957ymj75urz Date: Sun, 12 Mar 2023 20:23:46 +0100 Subject: [PATCH 2/9] fix makedirs on save --- nodes.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nodes.py b/nodes.py index d7708e34..7a37342f 100644 --- a/nodes.py +++ b/nodes.py @@ -809,7 +809,7 @@ class SaveImage: except ValueError: counter = 1 except FileNotFoundError: - os.mkdir(outputDirectory) + os.makedirs(outputDirectory, exist_ok=True) counter = 1 paths = list() From 5b425aaa4012e9dd0a7dddff4f43fecacf2c06f8 Mon Sep 17 00:00:00 2001 From: m957ymj75urz Date: Tue, 14 Mar 2023 09:08:54 +0100 Subject: [PATCH 3/9] remove subfolder widget and use filename_prefix --- nodes.py | 24 +++++++++++++++--------- 1 file changed, 15 insertions(+), 9 deletions(-) diff --git a/nodes.py b/nodes.py index 7a37342f..dbbe9eba 100644 --- a/nodes.py +++ b/nodes.py @@ -780,8 +780,7 @@ class SaveImage: def INPUT_TYPES(s): return {"required": {"images": ("IMAGE", ), - "filename_prefix": ("STRING", {"default": "ComfyUI"}), - "subfolder": ("STRING", {})}, + "filename_prefix": ("STRING", {"default": "ComfyUI"})}, "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"}, } @@ -792,9 +791,9 @@ class SaveImage: CATEGORY = "image" - def save_images(self, images, filename_prefix="ComfyUI", subfolder=None, prompt=None, extra_pnginfo=None): + def save_images(self, images, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None): def map_filename(filename): - prefix_len = len(filename_prefix) + prefix_len = len(os.path.basename(filename_prefix)) prefix = filename[:prefix_len + 1] try: digits = int(filename[prefix_len + 1:].split('_')[0]) @@ -802,14 +801,21 @@ class SaveImage: digits = 0 return (digits, prefix) - outputDirectory = os.path.join(os.path.dirname(os.path.realpath(__file__)), "output", subfolder); + subfolder = os.path.dirname(filename_prefix) + filename = os.path.basename(filename_prefix) + comfy_output_folder = os.path.join(os.path.dirname(os.path.realpath(__file__)), "output") + full_output_folder = os.path.join(comfy_output_folder, subfolder) + + if os.path.commonpath((comfy_output_folder, os.path.abspath(full_output_folder))) != comfy_output_folder: + print("Saving image outside the output folder is not allowed.") + return try: - counter = max(filter(lambda a: a[1][:-1] == filename_prefix and a[1][-1] == "_", map(map_filename, os.listdir(outputDirectory))))[0] + 1 + counter = max(filter(lambda a: a[1][:-1] == filename and a[1][-1] == "_", map(map_filename, os.listdir(full_output_folder))))[0] + 1 except ValueError: counter = 1 except FileNotFoundError: - os.makedirs(outputDirectory, exist_ok=True) + os.makedirs(full_output_folder, exist_ok=True) counter = 1 paths = list() @@ -822,8 +828,8 @@ class SaveImage: if extra_pnginfo is not None: 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(outputDirectory, file), pnginfo=metadata, optimize=True) + 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)) counter += 1 return { "ui": { "images": paths } } From b1294fa49f3a53981e8a194d3d2797be7a7714a0 Mon Sep 17 00:00:00 2001 From: m957ymj75urz Date: Tue, 14 Mar 2023 09:27:17 +0100 Subject: [PATCH 4/9] fix path traversal for /view --- server.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/server.py b/server.py index 05d1af04..be2fc056 100644 --- a/server.py +++ b/server.py @@ -118,11 +118,15 @@ class PromptServer(): output_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), type) if "subfolder" in request.rel_url.query: - output_dir = os.path.join(output_dir, request.rel_url.query["subfolder"]) + full_output_dir = os.path.join(output_dir, request.rel_url.query["subfolder"]) + if os.path.commonpath((os.path.realpath(full_output_dir), output_dir)) != output_dir: + return web.Response(status=403) + output_dir = full_output_dir file = request.rel_url.query["file"] file = os.path.basename(file) file = os.path.join(output_dir, file) + if os.path.isfile(file): return web.FileResponse(file) From 6daf9bb22c6b4cc268bf818a11358ee6629098c6 Mon Sep 17 00:00:00 2001 From: m957ymj75urz Date: Tue, 14 Mar 2023 09:27:53 +0100 Subject: [PATCH 5/9] switch to realpath to check path traversal --- nodes.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nodes.py b/nodes.py index dbbe9eba..e7c212c8 100644 --- a/nodes.py +++ b/nodes.py @@ -806,7 +806,7 @@ class SaveImage: comfy_output_folder = os.path.join(os.path.dirname(os.path.realpath(__file__)), "output") full_output_folder = os.path.join(comfy_output_folder, subfolder) - if os.path.commonpath((comfy_output_folder, os.path.abspath(full_output_folder))) != comfy_output_folder: + if os.path.commonpath((comfy_output_folder, os.path.realpath(full_output_folder))) != comfy_output_folder: print("Saving image outside the output folder is not allowed.") return From 629272c0cae419e7610411f2a162fa111e6eb040 Mon Sep 17 00:00:00 2001 From: m957ymj75urz Date: Wed, 15 Mar 2023 12:10:45 +0100 Subject: [PATCH 6/9] resolve conflict with the new PreviewImage node --- nodes.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/nodes.py b/nodes.py index 8c360311..7b27ef54 100644 --- a/nodes.py +++ b/nodes.py @@ -804,10 +804,9 @@ class SaveImage: subfolder = os.path.dirname(filename_prefix) filename = os.path.basename(filename_prefix) - comfy_output_folder = os.path.join(os.path.dirname(os.path.realpath(__file__)), "output") - full_output_folder = os.path.join(comfy_output_folder, subfolder) + full_output_folder = os.path.join(self.output_dir, subfolder) - if os.path.commonpath((comfy_output_folder, os.path.realpath(full_output_folder))) != comfy_output_folder: + if os.path.commonpath((self.output_dir, os.path.realpath(full_output_folder))) != self.output_dir: print("Saving image outside the output folder is not allowed.") return @@ -842,7 +841,7 @@ class SaveImage: 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.url_suffix = "&type=temp" @classmethod def INPUT_TYPES(s): From 16001e94298d02756e2029781927694097ec4284 Mon Sep 17 00:00:00 2001 From: m957ymj75urz Date: Thu, 16 Mar 2023 19:48:59 +0100 Subject: [PATCH 7/9] fixes for specific paths --- nodes.py | 5 +++-- web/scripts/app.js | 4 +++- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/nodes.py b/nodes.py index 7b27ef54..b4901a85 100644 --- a/nodes.py +++ b/nodes.py @@ -802,8 +802,9 @@ class SaveImage: digits = 0 return (digits, prefix) - subfolder = os.path.dirname(filename_prefix) - filename = os.path.basename(filename_prefix) + subfolder = os.path.dirname(os.path.normpath(filename_prefix)) + filename = os.path.basename(os.path.normpath(filename_prefix)) + full_output_folder = os.path.join(self.output_dir, subfolder) if os.path.commonpath((self.output_dir, os.path.realpath(full_output_folder))) != self.output_dir: diff --git a/web/scripts/app.js b/web/scripts/app.js index 31c575d7..1d777130 100644 --- a/web/scripts/app.js +++ b/web/scripts/app.js @@ -110,9 +110,11 @@ class ComfyApp { 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; + var params = new URLSearchParams({file: filename, subfolder}); + img.src = "/view?" + params.toString(); }); }) ).then((imgs) => { From 5fa514fa512f321574e7385097345c05198972e9 Mon Sep 17 00:00:00 2001 From: m957ymj75urz Date: Sat, 18 Mar 2023 20:13:25 +0100 Subject: [PATCH 8/9] rollback urlparams --- web/scripts/app.js | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/web/scripts/app.js b/web/scripts/app.js index 1d777130..51e67b92 100644 --- a/web/scripts/app.js +++ b/web/scripts/app.js @@ -113,8 +113,7 @@ class ComfyApp { var filename = src.replace(/^.*[\\\/]/, ''); var subfolder = src.replace(filename, ''); - var params = new URLSearchParams({file: filename, subfolder}); - img.src = "/view?" + params.toString(); + img.src = "/view?file=" + filename + "&subfolder=" + subfolder; }); }) ).then((imgs) => { From d1138e8ba0a8fd007c4277b2a4bb89e1c6d82349 Mon Sep 17 00:00:00 2001 From: m957ymj75urz Date: Sun, 19 Mar 2023 12:54:29 +0100 Subject: [PATCH 9/9] 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