From b962db9952d45b0ed8185ce7d18cf7d960f99de4 Mon Sep 17 00:00:00 2001 From: Robin Huang Date: Thu, 12 Sep 2024 21:10:27 +0900 Subject: [PATCH] Add cli arg to override user directory (#4856) * Override user directory. * Use overridden user directory. * Remove prints. * Remove references to global user_files. * Remove unused replace_folder function. * Remove newline. * Remove global during get_user_directory. * Add validation. --- app/user_manager.py | 17 +++++++++-------- comfy/cli_args.py | 2 ++ folder_paths.py | 7 +++++++ main.py | 5 +++++ 4 files changed, 23 insertions(+), 8 deletions(-) diff --git a/app/user_manager.py b/app/user_manager.py index 95ad157c..62c22cde 100644 --- a/app/user_manager.py +++ b/app/user_manager.py @@ -7,16 +7,15 @@ import shutil from aiohttp import web from urllib import parse from comfy.cli_args import args -from folder_paths import user_directory +import folder_paths from .app_settings import AppSettings default_user = "default" -users_file = os.path.join(user_directory, "users.json") class UserManager(): def __init__(self): - global user_directory + user_directory = folder_paths.get_user_directory() self.settings = AppSettings(self) if not os.path.exists(user_directory): @@ -26,14 +25,17 @@ class UserManager(): print("****** For multi-user setups add the --multi-user CLI argument to enable multiple user profiles. ******") if args.multi_user: - if os.path.isfile(users_file): - with open(users_file) as f: + if os.path.isfile(self.get_users_file()): + with open(self.get_users_file()) as f: self.users = json.load(f) else: self.users = {} else: self.users = {"default": "default"} + def get_users_file(self): + return os.path.join(folder_paths.get_user_directory(), "users.json") + def get_request_user_id(self, request): user = "default" if args.multi_user and "comfy-user" in request.headers: @@ -45,7 +47,7 @@ class UserManager(): return user def get_request_user_filepath(self, request, file, type="userdata", create_dir=True): - global user_directory + user_directory = folder_paths.get_user_directory() if type == "userdata": root_dir = user_directory @@ -85,8 +87,7 @@ class UserManager(): self.users[user_id] = name - global users_file - with open(users_file, "w") as f: + with open(self.get_users_file(), "w") as f: json.dump(self.users, f) return user_id diff --git a/comfy/cli_args.py b/comfy/cli_args.py index 65b87990..d0f70643 100644 --- a/comfy/cli_args.py +++ b/comfy/cli_args.py @@ -171,6 +171,8 @@ parser.add_argument( help="The local filesystem path to the directory where the frontend is located. Overrides --front-end-version.", ) +parser.add_argument("--user-directory", type=is_valid_directory, default=None, help="Set the ComfyUI user directory with an absolute path.") + if comfy.options.args_parsing: args = parser.parse_args() else: diff --git a/folder_paths.py b/folder_paths.py index d7a7e0c3..263704a4 100644 --- a/folder_paths.py +++ b/folder_paths.py @@ -84,6 +84,13 @@ def get_input_directory() -> str: global input_directory return input_directory +def get_user_directory() -> str: + 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: diff --git a/main.py b/main.py index 748846b4..00b49fc4 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)