Merge 0c509600a4 into 7cf4e78335

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 --base-directory.")
 
 parser.add_argument("--enable-compress-response-body", action="store_true", help="Enable compressing response body.")
 

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"])
 
@@ -144,6 +148,14 @@ def set_user_directory(user_dir: str) -> None:
     global user_directory
     user_directory = user_dir
 
+def get_models_directory() -> str:
+    global models_dir
+    return models_dir
+
+def set_models_directory(models_directory: str) -> None:
+    global models_dir
+    models_dir = models_directory
+
 
 # System User Protection - Protects system directories from HTTP endpoint access
 # System Users are internal-only users that cannot be accessed via HTTP endpoints.

tests-unit/comfy_test/folder_path_test.py

@@ -163,3 +163,30 @@ 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.get_models_directory() == os.path.abspath(temp_dir)
+        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)
+
+
+def test_models_dir_getter_setter():
+    orig_models_dir = folder_paths.get_models_directory()
+    try:
+        folder_paths.set_models_directory("/custom/models/path")
+        assert folder_paths.get_models_directory() == "/custom/models/path"
+    finally:
+        folder_paths.set_models_directory(orig_models_dir)
+