"vscode:/vscode.git/clone" did not exist on "01ec3db9324a6c2a3ac2bdc6b99ed729528a10e1"
Commit 6db777b3 authored by pythongosssss's avatar pythongosssss
Browse files

Added ability to save images to temp dir

parent a50e1118
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()
......@@ -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 } }
......
......@@ -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)
......
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