From 558b7d8b22da91d575ba93eb52d69010257b127f Mon Sep 17 00:00:00 2001 From: "Dr.Lt.Data" <128333288+ltdrdata@users.noreply.github.com> Date: Sat, 14 Dec 2024 08:21:32 +0900 Subject: [PATCH] fix: prestartup script is not applied due to `extra_model_paths.yaml` and ensure custom paths are used during startup (#5872) * fix: The custom nodes installed in the paths specified in `extra_model_paths.yaml` encounter a bug where the prestartup script is not imported. * Ensure custom paths are used during startup https://github.com/comfyanonymous/ComfyUI/pull/5794 --- main.py | 71 ++++++++++++++++++++++++++++++--------------------------- 1 file changed, 38 insertions(+), 33 deletions(-) diff --git a/main.py b/main.py index dd6526e1..9d163263 100644 --- a/main.py +++ b/main.py @@ -7,6 +7,9 @@ import folder_paths import time from comfy.cli_args import args from app.logger import setup_logger +import itertools +import utils.extra_config +import logging if __name__ == "__main__": #NOTE: These do not do anything on core ComfyUI which should already have no communication with the internet, they are for custom nodes. @@ -16,6 +19,40 @@ if __name__ == "__main__": setup_logger(log_level=args.verbose) +def apply_custom_paths(): + # extra model paths + extra_model_paths_config_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "extra_model_paths.yaml") + if os.path.isfile(extra_model_paths_config_path): + utils.extra_config.load_extra_path_config(extra_model_paths_config_path) + + if args.extra_model_paths_config: + for config_path in itertools.chain(*args.extra_model_paths_config): + utils.extra_config.load_extra_path_config(config_path) + + # --output-directory, --input-directory, --user-directory + if args.output_directory: + output_dir = os.path.abspath(args.output_directory) + logging.info(f"Setting output directory to: {output_dir}") + folder_paths.set_output_directory(output_dir) + + # These are the default folders that checkpoints, clip and vae models will be saved to when using CheckpointSave, etc.. nodes + folder_paths.add_model_folder_path("checkpoints", os.path.join(folder_paths.get_output_directory(), "checkpoints")) + folder_paths.add_model_folder_path("clip", os.path.join(folder_paths.get_output_directory(), "clip")) + folder_paths.add_model_folder_path("vae", os.path.join(folder_paths.get_output_directory(), "vae")) + folder_paths.add_model_folder_path("diffusion_models", + os.path.join(folder_paths.get_output_directory(), "diffusion_models")) + folder_paths.add_model_folder_path("loras", os.path.join(folder_paths.get_output_directory(), "loras")) + + if args.input_directory: + 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) + def execute_prestartup_script(): def execute_script(script_path): @@ -57,18 +94,16 @@ def execute_prestartup_script(): print("{:6.1f} seconds{}:".format(n[0], import_message), n[1]) print() +apply_custom_paths() execute_prestartup_script() # Main code import asyncio -import itertools import shutil import threading import gc -import logging -import utils.extra_config if os.name == "nt": logging.getLogger("xformers").addFilter(lambda record: 'A matching Triton is not available' not in record.getMessage()) @@ -209,14 +244,6 @@ if __name__ == "__main__": server = server.PromptServer(loop) q = execution.PromptQueue(server) - extra_model_paths_config_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "extra_model_paths.yaml") - if os.path.isfile(extra_model_paths_config_path): - utils.extra_config.load_extra_path_config(extra_model_paths_config_path) - - if args.extra_model_paths_config: - for config_path in itertools.chain(*args.extra_model_paths_config): - utils.extra_config.load_extra_path_config(config_path) - nodes.init_extra_nodes(init_custom_nodes=not args.disable_all_custom_nodes) cuda_malloc_warning() @@ -226,28 +253,6 @@ if __name__ == "__main__": threading.Thread(target=prompt_worker, daemon=True, args=(q, server,)).start() - if args.output_directory: - output_dir = os.path.abspath(args.output_directory) - logging.info(f"Setting output directory to: {output_dir}") - folder_paths.set_output_directory(output_dir) - - #These are the default folders that checkpoints, clip and vae models will be saved to when using CheckpointSave, etc.. nodes - folder_paths.add_model_folder_path("checkpoints", os.path.join(folder_paths.get_output_directory(), "checkpoints")) - folder_paths.add_model_folder_path("clip", os.path.join(folder_paths.get_output_directory(), "clip")) - folder_paths.add_model_folder_path("vae", os.path.join(folder_paths.get_output_directory(), "vae")) - folder_paths.add_model_folder_path("diffusion_models", os.path.join(folder_paths.get_output_directory(), "diffusion_models")) - folder_paths.add_model_folder_path("loras", os.path.join(folder_paths.get_output_directory(), "loras")) - - if args.input_directory: - 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)