From 148d57a772c40a46004846206aaf068fc15a5c57 Mon Sep 17 00:00:00 2001
From: doctorpangloss <@hiddenswitch.com>
Date: Fri, 1 Mar 2024 08:14:21 -0800
Subject: [PATCH] Add extra_model_paths_config to the valid configuration for
 the comfyui-worker entry point

---
 comfy/cmd/extra_model_paths.py | 25 +++++++++++++++++++++++++
 comfy/cmd/main.py              | 27 +--------------------------
 comfy/cmd/worker.py            |  6 ++++++
 3 files changed, 32 insertions(+), 26 deletions(-)
 create mode 100644 comfy/cmd/extra_model_paths.py

diff --git a/comfy/cmd/extra_model_paths.py b/comfy/cmd/extra_model_paths.py
new file mode 100644
index 000000000..3eb02809e
--- /dev/null
+++ b/comfy/cmd/extra_model_paths.py
@@ -0,0 +1,25 @@
+import os
+import yaml
+
+
+def load_extra_path_config(yaml_path):
+    from . import folder_paths
+
+    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")
+        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)
+                print("Adding extra search path", x, full_path)
+                folder_paths.add_model_folder_path(x, full_path)
diff --git a/comfy/cmd/main.py b/comfy/cmd/main.py
index c8fa363e7..1385c8ae7 100644
--- a/comfy/cmd/main.py
+++ b/comfy/cmd/main.py
@@ -1,6 +1,3 @@
-import signal
-import sys
-
 from .. import options
 
 options.enable_args_parsing()
@@ -10,6 +7,7 @@ import importlib.util
 
 from ..cmd import cuda_malloc
 from ..cmd import folder_paths
+from .extra_model_paths import load_extra_path_config
 from ..analytics.analytics import initialize_event_tracking
 import time
 
@@ -78,8 +76,6 @@ if args.deterministic:
         os.environ['CUBLAS_WORKSPACE_CONFIG'] = ":4096:8"
 
 from .. import utils
-import yaml
-from contextlib import AsyncExitStack
 
 from ..cmd import execution
 from ..cmd import server as server_module
@@ -172,27 +168,6 @@ def cleanup_temp():
         pass
 
 
-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")
-        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)
-                print("Adding extra search path", x, full_path)
-                folder_paths.add_model_folder_path(x, full_path)
-
-
 def cuda_malloc_warning():
     device = model_management.get_torch_device()
     device_name = model_management.get_torch_device_name(device)
diff --git a/comfy/cmd/worker.py b/comfy/cmd/worker.py
index bfd67f9c7..74a0633ec 100644
--- a/comfy/cmd/worker.py
+++ b/comfy/cmd/worker.py
@@ -1,6 +1,8 @@
 import asyncio
+import itertools
 import os
 
+from .extra_model_paths import load_extra_path_config
 from .. import options
 
 options.enable_args_parsing()
@@ -45,6 +47,10 @@ async def main():
 
         folder_paths.set_temp_directory(temp_dir)
 
+    if args.extra_model_paths_config:
+        for config_path in itertools.chain(*args.extra_model_paths_config):
+            load_extra_path_config(config_path)
+
     from ..distributed.distributed_prompt_worker import DistributedPromptWorker
     async with DistributedPromptWorker(connection_uri=args.distributed_queue_connection_uri,
                                        queue_name=args.distributed_queue_name):