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.
This commit is contained in:
Robin Huang 2024-09-12 21:10:27 +09:00 committed by GitHub
parent d0b7ab88ba
commit b962db9952
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 23 additions and 8 deletions

View File

@ -7,16 +7,15 @@ import shutil
from aiohttp import web from aiohttp import web
from urllib import parse from urllib import parse
from comfy.cli_args import args from comfy.cli_args import args
from folder_paths import user_directory import folder_paths
from .app_settings import AppSettings from .app_settings import AppSettings
default_user = "default" default_user = "default"
users_file = os.path.join(user_directory, "users.json")
class UserManager(): class UserManager():
def __init__(self): def __init__(self):
global user_directory user_directory = folder_paths.get_user_directory()
self.settings = AppSettings(self) self.settings = AppSettings(self)
if not os.path.exists(user_directory): 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. ******") print("****** For multi-user setups add the --multi-user CLI argument to enable multiple user profiles. ******")
if args.multi_user: if args.multi_user:
if os.path.isfile(users_file): if os.path.isfile(self.get_users_file()):
with open(users_file) as f: with open(self.get_users_file()) as f:
self.users = json.load(f) self.users = json.load(f)
else: else:
self.users = {} self.users = {}
else: else:
self.users = {"default": "default"} 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): def get_request_user_id(self, request):
user = "default" user = "default"
if args.multi_user and "comfy-user" in request.headers: if args.multi_user and "comfy-user" in request.headers:
@ -45,7 +47,7 @@ class UserManager():
return user return user
def get_request_user_filepath(self, request, file, type="userdata", create_dir=True): 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": if type == "userdata":
root_dir = user_directory root_dir = user_directory
@ -85,8 +87,7 @@ class UserManager():
self.users[user_id] = name self.users[user_id] = name
global users_file with open(self.get_users_file(), "w") as f:
with open(users_file, "w") as f:
json.dump(self.users, f) json.dump(self.users, f)
return user_id return user_id

View File

@ -171,6 +171,8 @@ parser.add_argument(
help="The local filesystem path to the directory where the frontend is located. Overrides --front-end-version.", 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: if comfy.options.args_parsing:
args = parser.parse_args() args = parser.parse_args()
else: else:

View File

@ -84,6 +84,13 @@ def get_input_directory() -> str:
global input_directory global input_directory
return 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 #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: def get_directory_by_type(type_name: str) -> str | None:

View File

@ -232,6 +232,11 @@ if __name__ == "__main__":
input_dir = os.path.abspath(args.input_directory) input_dir = os.path.abspath(args.input_directory)
logging.info(f"Setting input directory to: {input_dir}") logging.info(f"Setting input directory to: {input_dir}")
folder_paths.set_input_directory(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: if args.quick_test_for_ci:
exit(0) exit(0)