From 38192f34907058c3db97cf328c9318176249d6b4 Mon Sep 17 00:00:00 2001
From: Robin Huang <robin.j.huang@gmail.com>
Date: Mon, 9 Sep 2024 17:54:56 +0900
Subject: [PATCH] Refactor.

---
 main.py                                       | 28 ++-----------------
 .../extra_config_test.py}                     |  2 +-
 utils/__init__.py                             |  0
 utils/extra_config.py                         | 25 +++++++++++++++++
 4 files changed, 29 insertions(+), 26 deletions(-)
 rename tests-unit/{main_test.py => utils/extra_config_test.py} (97%)
 create mode 100644 utils/__init__.py
 create mode 100644 utils/extra_config.py

diff --git a/main.py b/main.py
index 1bd6a6853..a1db97cd0 100644
--- a/main.py
+++ b/main.py
@@ -63,6 +63,7 @@ import threading
 import gc
 
 import logging
+from utils import extra_config
 
 if os.name == "nt":
     logging.getLogger("xformers").addFilter(lambda record: 'A matching Triton is not available' not in record.getMessage())
@@ -85,7 +86,6 @@ if args.windows_standalone_build:
         pass
 
 import comfy.utils
-import yaml
 
 import execution
 import server
@@ -180,28 +180,6 @@ def cleanup_temp():
         shutil.rmtree(temp_dir, ignore_errors=True)
 
 
-def load_extra_path_config(yaml_path):
-    with open(yaml_path, 'r') as stream:
-        config = yaml.safe_load(stream)
-    for c in config:
-        conf = config[c]
-        if conf is None:
-            continue
-        base_path = None
-        if "base_path" in conf:
-            base_path = conf.pop("base_path")
-            base_path = os.path.expanduser(base_path)
-        for x in conf:
-            for y in conf[x].split("\n"):
-                if len(y) == 0:
-                    continue
-                full_path = y
-                if base_path is not None:
-                    full_path = os.path.join(base_path, full_path)
-                logging.info("Adding extra search path {} {}".format(x, full_path))
-                folder_paths.add_model_folder_path(x, full_path)
-
-
 if __name__ == "__main__":
     if args.temp_directory:
         temp_dir = os.path.join(os.path.abspath(args.temp_directory), "temp")
@@ -223,11 +201,11 @@ if __name__ == "__main__":
 
     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):
-        load_extra_path_config(extra_model_paths_config_path)
+        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):
-            load_extra_path_config(config_path)
+            extra_config.load_extra_path_config(config_path)
 
     nodes.init_extra_nodes(init_custom_nodes=not args.disable_all_custom_nodes)
 
diff --git a/tests-unit/main_test.py b/tests-unit/utils/extra_config_test.py
similarity index 97%
rename from tests-unit/main_test.py
rename to tests-unit/utils/extra_config_test.py
index 836c94995..f56dd3e2e 100644
--- a/tests-unit/main_test.py
+++ b/tests-unit/utils/extra_config_test.py
@@ -3,7 +3,7 @@ import yaml
 import os
 from unittest.mock import Mock, patch, mock_open
 
-from main import load_extra_path_config
+from utils.extra_config import load_extra_path_config
 import folder_paths
 
 @pytest.fixture
diff --git a/utils/__init__.py b/utils/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/utils/extra_config.py b/utils/extra_config.py
new file mode 100644
index 000000000..23c2d791c
--- /dev/null
+++ b/utils/extra_config.py
@@ -0,0 +1,25 @@
+import os
+import yaml
+import folder_paths
+import logging
+
+def load_extra_path_config(yaml_path):
+    with open(yaml_path, 'r') as stream:
+        config = yaml.safe_load(stream)
+    for c in config:
+        conf = config[c]
+        if conf is None:
+            continue
+        base_path = None
+        if "base_path" in conf:
+            base_path = conf.pop("base_path")
+            base_path = os.path.expanduser(base_path)
+        for x in conf:
+            for y in conf[x].split("\n"):
+                if len(y) == 0:
+                    continue
+                full_path = y
+                if base_path is not None:
+                    full_path = os.path.join(base_path, full_path)
+                logging.info("Adding extra search path {} {}".format(x, full_path))
+                folder_paths.add_model_folder_path(x, full_path)