From 12f88d83421aaee9cdf72f9ba4c2587c7f0f44e3 Mon Sep 17 00:00:00 2001
From: Robin Huang <robin.j.huang@gmail.com>
Date: Tue, 3 Sep 2024 17:29:47 -0700
Subject: [PATCH] Override user directory.

---
 app/user_manager.py |  1 +
 comfy/cli_args.py   |  1 +
 folder_paths.py     | 15 +++++++++++++++
 main.py             |  5 +++++
 4 files changed, 22 insertions(+)

diff --git a/app/user_manager.py b/app/user_manager.py
index 95ad157c3..9d99204b3 100644
--- a/app/user_manager.py
+++ b/app/user_manager.py
@@ -20,6 +20,7 @@ class UserManager():
 
         self.settings = AppSettings(self)
         if not os.path.exists(user_directory):
+            print('Creating user directory: ', user_directory)
             os.mkdir(user_directory)
             if not args.multi_user:
                 print("****** User settings have been changed to be stored on the server instead of browser storage. ******")
diff --git a/comfy/cli_args.py b/comfy/cli_args.py
index 65b879907..75b19c009 100644
--- a/comfy/cli_args.py
+++ b/comfy/cli_args.py
@@ -47,6 +47,7 @@ parser.add_argument("--extra-model-paths-config", type=str, default=None, metava
 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("--input-directory", type=str, default=None, help="Set the ComfyUI input directory.")
+parser.add_argument("--user-directory", type=str, default=None, help="Set the ComfyUI user 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 d7a7e0c3d..949b138c2 100644
--- a/folder_paths.py
+++ b/folder_paths.py
@@ -84,6 +84,14 @@ def get_input_directory() -> str:
     global input_directory
     return input_directory
 
+def get_user_directory() -> str:
+    global user_directory
+    return user_directory
+
+def set_user_directory(user_dir: str) -> None:
+    global user_directory
+    user_directory = user_dir
+
 
 #NOTE: used in http server so don't put folders that should not be accessed remotely
 def get_directory_by_type(type_name: str) -> str | None:
@@ -166,6 +174,13 @@ def add_model_folder_path(folder_name: str, full_folder_path: str) -> None:
     else:
         folder_names_and_paths[folder_name] = ([full_folder_path], set())
 
+def replace_folder_path(folder_name: str, full_folder_path: str) -> None:
+    global folder_names_and_paths
+    folder_name = map_legacy(folder_name)
+    if folder_name in folder_names_and_paths:
+        folder_names_and_paths[folder_name] = full_folder_path
+
+
 def get_folder_paths(folder_name: str) -> list[str]:
     folder_name = map_legacy(folder_name)
     return folder_names_and_paths[folder_name][0][:]
diff --git a/main.py b/main.py
index 748846b4a..00b49fc4c 100644
--- a/main.py
+++ b/main.py
@@ -232,6 +232,11 @@ if __name__ == "__main__":
         input_dir = os.path.abspath(args.input_directory)
         logging.info(f"Setting input directory to: {input_dir}")
         folder_paths.set_input_directory(input_dir)
+    
+    if args.user_directory:
+        user_dir = os.path.abspath(args.user_directory)
+        logging.info(f"Setting user directory to: {user_dir}")
+        folder_paths.set_user_directory(user_dir)
 
     if args.quick_test_for_ci:
         exit(0)