From 6db777b3489798740c0ffdc6f503fe4279f2c435 Mon Sep 17 00:00:00 2001 From: pythongosssss <125205205+pythongosssss@users.noreply.github.com> Date: Mon, 13 Mar 2023 19:34:05 +0000 Subject: [PATCH 1/5] Added ability to save images to temp dir --- main.py | 9 +++++++++ nodes.py | 21 ++++++++++++++++++--- server.py | 2 +- 3 files changed, 28 insertions(+), 4 deletions(-) diff --git a/main.py b/main.py index b2b3f1c40..889e2cef7 100644 --- a/main.py +++ b/main.py @@ -1,5 +1,6 @@ import os import sys +import shutil import threading import asyncio @@ -53,7 +54,14 @@ def hijack_progress(server): return v setattr(tqdm, "update", wrapped_func) +def cleanup_temp(): + temp_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "temp") + if os.path.exists(temp_dir): + shutil.rmtree(temp_dir) + if __name__ == "__main__": + cleanup_temp() + loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) server = server.PromptServer(loop) @@ -93,3 +101,4 @@ if __name__ == "__main__": else: loop.run_until_complete(run(server, address=address, port=port, verbose=not dont_print, call_on_start=call_on_start)) + cleanup_temp() diff --git a/nodes.py b/nodes.py index 0a0a0a9cd..b201c352c 100644 --- a/nodes.py +++ b/nodes.py @@ -775,12 +775,14 @@ class KSamplerAdvanced: class SaveImage: def __init__(self): self.output_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "output") + self.temp_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "temp") @classmethod def INPUT_TYPES(s): return {"required": {"images": ("IMAGE", ), - "filename_prefix": ("STRING", {"default": "ComfyUI"})}, + "filename_prefix": ("STRING", {"default": "ComfyUI"}), + "use_temp_file": (["yes", "no"], {"default" : "no"} ),}, "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"}, } @@ -791,7 +793,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", use_temp_file="no", prompt=None, extra_pnginfo=None): def map_filename(filename): prefix_len = len(filename_prefix) prefix = filename[:prefix_len + 1] @@ -818,8 +820,21 @@ 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(self.output_dir, file), pnginfo=metadata, optimize=True) + + if use_temp_file == "yes": + if not os.path.exists(self.temp_dir): + os.makedirs(self.temp_dir) + dir = self.temp_dir + else: + dir = self.output_dir + + img.save(os.path.join(dir, file), pnginfo=metadata, optimize=True) + + if use_temp_file == "yes": + file += "?type=temp" + paths.append(file) counter += 1 return { "ui": { "images": paths } } diff --git a/server.py b/server.py index a29d85974..eb6857010 100644 --- a/server.py +++ b/server.py @@ -113,7 +113,7 @@ class PromptServer(): async def view_image(request): if "file" in request.match_info: type = request.rel_url.query.get("type", "output") - if type != "output" and type != "input": + if type not in ["output", "input", "temp"]: return web.Response(status=400) output_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), type) From e5318d918cde493eb2100bd308cf76425d828da1 Mon Sep 17 00:00:00 2001 From: pythongosssss <125205205+pythongosssss@users.noreply.github.com> Date: Mon, 13 Mar 2023 19:34:29 +0000 Subject: [PATCH 2/5] Combo support detault value --- web/scripts/app.js | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/web/scripts/app.js b/web/scripts/app.js index e70e1c157..8b832eba1 100644 --- a/web/scripts/app.js +++ b/web/scripts/app.js @@ -497,7 +497,11 @@ class ComfyApp { if (Array.isArray(type)) { // Enums e.g. latent rotation - this.addWidget("combo", inputName, type[0], () => {}, { values: type }); + let defaultValue = type[0]; + if (inputData[1] && inputData[1].default) { + defaultValue = inputData[1].default; + } + this.addWidget("combo", inputName, defaultValue, () => {}, { values: type }); } else if (`${type}:${inputName}` in widgets) { // Support custom widgets by Type:Name Object.assign(config, widgets[`${type}:${inputName}`](this, inputName, inputData, app) || {}); From 6a6256a75c39c0e945db40670e2855808c2f4663 Mon Sep 17 00:00:00 2001 From: pythongosssss <125205205+pythongosssss@users.noreply.github.com> Date: Tue, 14 Mar 2023 19:28:07 +0000 Subject: [PATCH 3/5] Changed flag to new node --- nodes.py | 38 ++++++++++++++++++++------------------ 1 file changed, 20 insertions(+), 18 deletions(-) diff --git a/nodes.py b/nodes.py index b201c352c..b5aa3efe0 100644 --- a/nodes.py +++ b/nodes.py @@ -775,14 +775,13 @@ class KSamplerAdvanced: class SaveImage: def __init__(self): self.output_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "output") - self.temp_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "temp") + self.url_suffix = "" @classmethod def INPUT_TYPES(s): return {"required": {"images": ("IMAGE", ), - "filename_prefix": ("STRING", {"default": "ComfyUI"}), - "use_temp_file": (["yes", "no"], {"default" : "no"} ),}, + "filename_prefix": ("STRING", {"default": "ComfyUI"}),}, "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"}, } @@ -810,6 +809,9 @@ class SaveImage: os.mkdir(self.output_dir) counter = 1 + if not os.path.exists(self.output_dir): + os.makedirs(self.output_dir) + paths = list() for image in images: i = 255. * image.cpu().numpy() @@ -820,25 +822,24 @@ 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" - - if use_temp_file == "yes": - if not os.path.exists(self.temp_dir): - os.makedirs(self.temp_dir) - dir = self.temp_dir - else: - dir = self.output_dir - - img.save(os.path.join(dir, file), pnginfo=metadata, optimize=True) - - if use_temp_file == "yes": - file += "?type=temp" - - paths.append(file) + img.save(os.path.join(self.output_dir, file), pnginfo=metadata, optimize=True) + paths.append(file + self.url_suffix) counter += 1 return { "ui": { "images": paths } } +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" + + @classmethod + def INPUT_TYPES(s): + return {"required": + {"images": ("IMAGE", ), }, + "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"}, + } + class LoadImage: input_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "input") @classmethod @@ -959,6 +960,7 @@ NODE_CLASS_MAPPINGS = { "EmptyLatentImage": EmptyLatentImage, "LatentUpscale": LatentUpscale, "SaveImage": SaveImage, + "PreviewImage": PreviewImage, "LoadImage": LoadImage, "LoadImageMask": LoadImageMask, "ImageScale": ImageScale, From 94a279373bf5104b9ad9ac2bc6da3e8bbe1e5141 Mon Sep 17 00:00:00 2001 From: pythongosssss <125205205+pythongosssss@users.noreply.github.com> Date: Tue, 14 Mar 2023 19:39:49 +0000 Subject: [PATCH 4/5] Better auto pos of images --- web/scripts/app.js | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/web/scripts/app.js b/web/scripts/app.js index 8b832eba1..445bc5d4b 100644 --- a/web/scripts/app.js +++ b/web/scripts/app.js @@ -142,7 +142,14 @@ class ComfyApp { if (numImages === 1 && !imageIndex) { this.imageIndex = imageIndex = 0; } - let shiftY = this.type === "SaveImage" ? 55 : this.imageOffset || 0; + + let shiftY; + if (this.imageOffset != null) { + shiftY = this.imageOffset; + } else { + shiftY = this.computeSize()[1]; + } + let dw = this.size[0]; let dh = this.size[1]; dh -= shiftY; From 8537ab6f45cdfe060e5fc3fb4aff87a35a0b9072 Mon Sep 17 00:00:00 2001 From: pythongosssss <125205205+pythongosssss@users.noreply.github.com> Date: Tue, 14 Mar 2023 19:42:28 +0000 Subject: [PATCH 5/5] tidy --- nodes.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/nodes.py b/nodes.py index b5aa3efe0..650d7f65d 100644 --- a/nodes.py +++ b/nodes.py @@ -781,7 +781,7 @@ class SaveImage: def INPUT_TYPES(s): return {"required": {"images": ("IMAGE", ), - "filename_prefix": ("STRING", {"default": "ComfyUI"}),}, + "filename_prefix": ("STRING", {"default": "ComfyUI"})}, "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"}, } @@ -792,7 +792,7 @@ class SaveImage: CATEGORY = "image" - def save_images(self, images, filename_prefix="ComfyUI", use_temp_file="no", 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 = filename[:prefix_len + 1]