From 2bc12d3d22efb5c63ae3a7fc342bb2dd16b31735 Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Fri, 11 Aug 2023 05:00:25 -0400
Subject: [PATCH] Add --temp-directory argument to set temp directory.

---
 comfy/cli_args.py | 1 +
 folder_paths.py   | 4 ++++
 main.py           | 6 +++++-
 3 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/comfy/cli_args.py b/comfy/cli_args.py
index 81bbc479..ec7d34a5 100644
--- a/comfy/cli_args.py
+++ b/comfy/cli_args.py
@@ -38,6 +38,7 @@ parser.add_argument("--port", type=int, default=8188, help="Set the listen port.
 parser.add_argument("--enable-cors-header", type=str, default=None, metavar="ORIGIN", nargs="?", const="*", help="Enable CORS (Cross-Origin Resource Sharing) with optional origin or allow all with default '*'.")
 parser.add_argument("--extra-model-paths-config", type=str, default=None, metavar="PATH", nargs='+', action='append', help="Load one or more extra_model_paths.yaml files.")
 parser.add_argument("--output-directory", type=str, default=None, help="Set the ComfyUI output directory.")
+parser.add_argument("--temp-directory", type=str, default=None, help="Set the ComfyUI temp directory (default is in the ComfyUI directory).")
 parser.add_argument("--auto-launch", action="store_true", help="Automatically launch ComfyUI in the default browser.")
 parser.add_argument("--disable-auto-launch", action="store_true", help="Disable auto launching the browser.")
 parser.add_argument("--cuda-device", type=int, default=None, metavar="DEVICE_ID", help="Set the id of the cuda device this instance will use.")
diff --git a/folder_paths.py b/folder_paths.py
index d3b76d0d..e321690d 100644
--- a/folder_paths.py
+++ b/folder_paths.py
@@ -43,6 +43,10 @@ def set_output_directory(output_dir):
     global output_directory
     output_directory = output_dir
 
+def set_temp_directory(temp_dir):
+    global temp_directory
+    temp_directory = temp_dir
+
 def get_output_directory():
     global output_directory
     return output_directory
diff --git a/main.py b/main.py
index 07ebbd70..1571376b 100644
--- a/main.py
+++ b/main.py
@@ -100,7 +100,7 @@ def hijack_progress(server):
 
 
 def cleanup_temp():
-    temp_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "temp")
+    temp_dir = folder_paths.get_temp_directory()
     if os.path.exists(temp_dir):
         shutil.rmtree(temp_dir, ignore_errors=True)
 
@@ -127,6 +127,10 @@ def load_extra_path_config(yaml_path):
 
 
 if __name__ == "__main__":
+    if args.temp_directory:
+        temp_dir = os.path.join(os.path.abspath(args.temp_directory), "temp")
+        print(f"Setting temp directory to: {temp_dir}")
+        folder_paths.set_temp_directory(temp_dir)
     cleanup_temp()
 
     loop = asyncio.new_event_loop()