Skip to content

Commit

Permalink
Add cli arg to override user directory (comfyanonymous#4856)
Browse files Browse the repository at this point in the history
* 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.
  • Loading branch information
robinjhuang authored Sep 12, 2024
1 parent d0b7ab8 commit b962db9
Show file tree
Hide file tree
Showing 4 changed files with 23 additions and 8 deletions.
17 changes: 9 additions & 8 deletions app/user_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,16 +7,15 @@
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):
Expand All @@ -26,14 +25,17 @@ def __init__(self):
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:
Expand All @@ -45,7 +47,7 @@ def get_request_user_id(self, request):
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
Expand Down Expand Up @@ -85,8 +87,7 @@ def add_user(self, name):

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
Expand Down
2 changes: 2 additions & 0 deletions comfy/cli_args.py
Original file line number Diff line number Diff line change
Expand Up @@ -171,6 +171,8 @@ def is_valid_directory(path: Optional[str]) -> Optional[str]:
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:
Expand Down
7 changes: 7 additions & 0 deletions folder_paths.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
5 changes: 5 additions & 0 deletions main.py
Original file line number Diff line number Diff line change
Expand Up @@ -232,6 +232,11 @@ def cleanup_temp():
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)
Expand Down

0 comments on commit b962db9

Please sign in to comment.