can specify a subfolder in the SaveImage node

This commit is contained in:
m957ymj75urz 2023-03-12 19:51:39 +01:00
parent 6d6758e9e4
commit 195d7aec9f
3 changed files with 19 additions and 10 deletions

View File

@ -780,7 +780,8 @@ class SaveImage:
def INPUT_TYPES(s): def INPUT_TYPES(s):
return {"required": return {"required":
{"images": ("IMAGE", ), {"images": ("IMAGE", ),
"filename_prefix": ("STRING", {"default": "ComfyUI"})}, "filename_prefix": ("STRING", {"default": "ComfyUI"}),
"subfolder": ("STRING", {})},
"hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"}, "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
} }
@ -791,7 +792,7 @@ class SaveImage:
CATEGORY = "image" 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): def map_filename(filename):
prefix_len = len(filename_prefix) prefix_len = len(filename_prefix)
prefix = filename[:prefix_len + 1] prefix = filename[:prefix_len + 1]
@ -800,12 +801,15 @@ class SaveImage:
except: except:
digits = 0 digits = 0
return (digits, prefix) return (digits, prefix)
outputDirectory = os.path.join(os.path.dirname(os.path.realpath(__file__)), "output", subfolder);
try: 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: except ValueError:
counter = 1 counter = 1
except FileNotFoundError: except FileNotFoundError:
os.mkdir(self.output_dir) os.mkdir(outputDirectory)
counter = 1 counter = 1
paths = list() paths = list()
@ -819,8 +823,8 @@ class SaveImage:
for x in extra_pnginfo: for x in extra_pnginfo:
metadata.add_text(x, json.dumps(extra_pnginfo[x])) metadata.add_text(x, json.dumps(extra_pnginfo[x]))
file = f"{filename_prefix}_{counter:05}_.png" file = f"{filename_prefix}_{counter:05}_.png"
img.save(os.path.join(self.output_dir, file), pnginfo=metadata, optimize=True) img.save(os.path.join(outputDirectory, file), pnginfo=metadata, optimize=True)
paths.append(file) paths.append(os.path.join(subfolder, file))
counter += 1 counter += 1
return { "ui": { "images": paths } } return { "ui": { "images": paths } }

View File

@ -109,15 +109,18 @@ class PromptServer():
return web.Response(status=400) return web.Response(status=400)
@routes.get("/view/{file}") @routes.get("/view")
async def view_image(request): 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") type = request.rel_url.query.get("type", "output")
if type != "output" and type != "input": if type != "output" and type != "input":
return web.Response(status=400) return web.Response(status=400)
output_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), type) 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.basename(file)
file = os.path.join(output_dir, file) file = os.path.join(output_dir, file)
if os.path.isfile(file): if os.path.isfile(file):

View File

@ -110,7 +110,9 @@ class ComfyApp {
const img = new Image(); const img = new Image();
img.onload = () => r(img); img.onload = () => r(img);
img.onerror = () => r(null); 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) => { ).then((imgs) => {