diff --git a/comfy/cli_args.py b/comfy/cli_args.py index 4bef096fb..0d7df5e13 100644 --- a/comfy/cli_args.py +++ b/comfy/cli_args.py @@ -225,6 +225,7 @@ parser.add_argument( ) parser.add_argument("--user-directory", type=is_valid_directory, default=None, help="Set the ComfyUI user directory with an absolute path. Overrides --base-directory.") +parser.add_argument("--models-directory", type=is_valid_directory, default=None, help="Set the ComfyUI models directory. Overrides the models folder in --base-directory.") parser.add_argument("--enable-compress-response-body", action="store_true", help="Enable compressing response body.") diff --git a/folder_paths.py b/folder_paths.py index ee048b0f2..937428c18 100644 --- a/folder_paths.py +++ b/folder_paths.py @@ -17,7 +17,11 @@ if args.base_directory: else: base_path = os.path.dirname(os.path.realpath(__file__)) -models_dir = os.path.join(base_path, "models") +if args.models_directory: + models_dir = os.path.abspath(args.models_directory) +else: + models_dir = os.path.join(base_path, "models") + folder_names_and_paths["checkpoints"] = ([os.path.join(models_dir, "checkpoints")], supported_pt_extensions) folder_names_and_paths["configs"] = ([os.path.join(models_dir, "configs")], [".yaml"]) diff --git a/main.py b/main.py index 20ec83c9e..580074b19 100644 --- a/main.py +++ b/main.py @@ -131,6 +131,10 @@ def apply_custom_paths(): if args.base_directory: logging.info(f"Setting base directory to: {folder_paths.base_path}") + # --models-directory + if args.models_directory: + logging.info(f"Setting models directory to: {folder_paths.models_dir}") + # --output-directory, --input-directory, --user-directory if args.output_directory: output_dir = os.path.abspath(args.output_directory) diff --git a/tests-unit/comfy_test/folder_path_test.py b/tests-unit/comfy_test/folder_path_test.py index 3b398e60b..a0ef17a4c 100644 --- a/tests-unit/comfy_test/folder_path_test.py +++ b/tests-unit/comfy_test/folder_path_test.py @@ -163,3 +163,20 @@ def test_base_path_change_clears_old(set_base_dir): for name in ["controlnet", "diffusion_models", "text_encoders"]: assert len(folder_paths.get_folder_paths(name)) == 2 + + +def test_models_directory_cli_and_getters(temp_dir): + try: + with patch.object(sys, 'argv', ["main.py", "--models-directory", temp_dir]): + reload(comfy.cli_args) + reload(folder_paths) + + assert folder_paths.models_dir == os.path.abspath(temp_dir) + + with pytest.raises(Exception): + comfy.cli_args.is_valid_directory(os.path.join(temp_dir, "non_existent_folder_path")) + finally: + with patch.object(sys, 'argv', ["main.py"]): + reload(comfy.cli_args) + reload(folder_paths) +