Skip to content

Commit c01ac61

Browse files
committed
Added a function to load the directories in the optional file: extra_model_paths.yaml
1 parent 48248e0 commit c01ac61

File tree

3 files changed

+43
-5
lines changed

3 files changed

+43
-5
lines changed

.gitignore

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -145,3 +145,6 @@ cython_debug/
145145
*.sql
146146
*.sqlite
147147
*.xml
148+
/workflow_api.json
149+
/workflow_api.py
150+
/custom_nodes/

comfyui_to_python.py

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -10,9 +10,10 @@
1010

1111
import black
1212

13-
from utils import import_custom_nodes, add_comfyui_directory_to_sys_path, get_value_at_index
14-
1513
sys.path.append('../')
14+
sys.path.append('.')
15+
16+
from utils import import_custom_nodes, add_comfyui_directory_to_sys_path, get_value_at_index, add_extra_model_paths
1617

1718
from nodes import NODE_CLASS_MAPPINGS
1819

@@ -306,11 +307,11 @@ def assemble_python_code(self, import_statements: set, speical_functions_code: L
306307
"""
307308
# Get the source code of the utils functions as a string
308309
func_strings = []
309-
for func in [add_comfyui_directory_to_sys_path, get_value_at_index]:
310+
for func in [add_comfyui_directory_to_sys_path, get_value_at_index, add_extra_model_paths]:
310311
func_strings.append(f'\n{inspect.getsource(func)}')
311312
# Define static import statements required for the script
312-
static_imports = ['import os', 'import random', 'import sys', 'from typing import Sequence, Mapping, Any, Union',
313-
'import torch'] + func_strings + ['\n\nadd_comfyui_directory_to_sys_path()']
313+
static_imports = ['import os', 'import random', 'from pathlib import Path', 'import sys', 'from typing import Sequence, Mapping, Any, Union',
314+
'import torch'] + func_strings + ['\n\nadd_comfyui_directory_to_sys_path()\nadd_extra_model_paths()\n']
314315
# Check if custom nodes should be included
315316
if custom_nodes:
316317
static_imports.append(f'\n{inspect.getsource(import_custom_nodes)}\n')

utils.py

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import os
22
from typing import Sequence, Mapping, Any, Union
3+
from pathlib import Path
34
import sys
45

56
sys.path.append('../')
@@ -54,6 +55,39 @@ def search_directory(path: str) -> None:
5455
# Start the search from the current working directory
5556
search_directory(start_path)
5657

58+
def add_extra_model_paths() -> Path:
59+
"""
60+
Parse the optional extra_model_paths.yaml file and add the parsed paths to the sys.path.
61+
"""
62+
from pathlib import Path
63+
from main import load_extra_path_config
64+
65+
def find_config_file(path: str, name: str = "extra_model_paths.yaml") -> Path:
66+
# Check if the current directory contains the file
67+
if name in os.listdir(path):
68+
directory_path = os.path.join(path, name)
69+
print(f"{name} found: {directory_path}")
70+
return Path(directory_path)
71+
72+
# Get the parent directory
73+
parent_directory = os.path.dirname(path)
74+
75+
# If the parent directory is the same as the current directory, we've reached the root and stop the search
76+
if parent_directory == path:
77+
return
78+
79+
# Recursively call the function with the parent directory
80+
return find_config_file(parent_directory)
81+
82+
start_path = os.getcwd() # Get the current working directory
83+
file = find_config_file(start_path)
84+
85+
if os.path.isfile(file):
86+
load_extra_path_config(file)
87+
else:
88+
print("Could not find the extra_model_paths config file.")
89+
90+
5791

5892
def get_value_at_index(obj: Union[Sequence, Mapping], index: int) -> Any:
5993
"""Returns the value at the given index of a sequence or mapping.

0 commit comments

Comments
 (0)