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] 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 b2b3f1c4..889e2cef 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 0a0a0a9c..b201c352 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 a29d8597..eb685701 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)