Commit 6db777b3 authored by pythongosssss's avatar pythongosssss
Browse files

Added ability to save images to temp dir

parent a50e1118
import os import os
import sys import sys
import shutil
import threading import threading
import asyncio import asyncio
...@@ -53,7 +54,14 @@ def hijack_progress(server): ...@@ -53,7 +54,14 @@ def hijack_progress(server):
return v return v
setattr(tqdm, "update", wrapped_func) 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__": if __name__ == "__main__":
cleanup_temp()
loop = asyncio.new_event_loop() loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop) asyncio.set_event_loop(loop)
server = server.PromptServer(loop) server = server.PromptServer(loop)
...@@ -93,3 +101,4 @@ if __name__ == "__main__": ...@@ -93,3 +101,4 @@ if __name__ == "__main__":
else: else:
loop.run_until_complete(run(server, address=address, port=port, verbose=not dont_print, call_on_start=call_on_start)) loop.run_until_complete(run(server, address=address, port=port, verbose=not dont_print, call_on_start=call_on_start))
cleanup_temp()
...@@ -775,12 +775,14 @@ class KSamplerAdvanced: ...@@ -775,12 +775,14 @@ class KSamplerAdvanced:
class SaveImage: class SaveImage:
def __init__(self): def __init__(self):
self.output_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "output") 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 @classmethod
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"}),
"use_temp_file": (["yes", "no"], {"default" : "no"} ),},
"hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"}, "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
} }
...@@ -791,7 +793,7 @@ class SaveImage: ...@@ -791,7 +793,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", use_temp_file="no", 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]
...@@ -818,8 +820,21 @@ class SaveImage: ...@@ -818,8 +820,21 @@ class SaveImage:
if extra_pnginfo is not None: if extra_pnginfo is not None:
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)
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) paths.append(file)
counter += 1 counter += 1
return { "ui": { "images": paths } } return { "ui": { "images": paths } }
......
...@@ -113,7 +113,7 @@ class PromptServer(): ...@@ -113,7 +113,7 @@ class PromptServer():
async def view_image(request): async def view_image(request):
if "file" in request.match_info: if "file" in request.match_info:
type = request.rel_url.query.get("type", "output") 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) 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)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment