Clean up load extra models code.

This commit is contained in:
Peyton
2023-08-24 21:38:50 -05:00
parent ffe21d24c3
commit 27f7a84eec
2 changed files with 10 additions and 17 deletions
+5 -13
View File
@@ -2,7 +2,6 @@ import copy
import glob
import inspect
import json
import logging
import os
import random
import sys
@@ -11,20 +10,13 @@ from typing import Dict, List, Any, Callable, Tuple
import black
from utils import import_custom_nodes, find_path, add_comfyui_directory_to_sys_path, add_extra_model_paths, get_value_at_index
sys.path.append('../')
try:
from utils import import_custom_nodes, find_path, add_comfyui_directory_to_sys_path, add_extra_model_paths, get_value_at_index
except ImportError:
sys.path.append('.')
from utils import import_custom_nodes, find_path, add_comfyui_directory_to_sys_path, add_extra_model_paths, get_value_at_index
from nodes import NODE_CLASS_MAPPINGS
logging.basicConfig(level=logging.INFO)
class FileHandler:
"""Handles reading and writing files.
@@ -311,10 +303,10 @@ class CodeGenerator:
"""
# Get the source code of the utils functions as a string
func_strings = []
for func in [find_path, add_comfyui_directory_to_sys_path, add_extra_model_paths, get_value_at_index]:
for func in [get_value_at_index, find_path, add_comfyui_directory_to_sys_path, add_extra_model_paths]:
func_strings.append(f'\n{inspect.getsource(func)}')
# Define static import statements required for the script
static_imports = ['import os', 'import random', 'from pathlib import Path', 'import sys', 'from typing import Sequence, Mapping, Any, Union',
static_imports = ['import os', 'import random', 'import sys', 'from typing import Sequence, Mapping, Any, Union',
'import torch'] + func_strings + ['\n\nadd_comfyui_directory_to_sys_path()\nadd_extra_model_paths()\n']
# Check if custom nodes should be included
if custom_nodes:
+5 -4
View File
@@ -1,6 +1,5 @@
import os
from typing import Sequence, Mapping, Any, Union
from pathlib import Path
import sys
sys.path.append('../')
@@ -28,7 +27,8 @@ def import_custom_nodes() -> None:
# Initializing custom nodes
init_custom_nodes()
def find_path(name: str, path: str = None) -> Union[Path, None]:
def find_path(name: str, path: str = None) -> str:
"""
Recursively looks at parent folders starting from the given path until it finds the given name.
Returns the path as a Path object if found, or None otherwise.
@@ -41,7 +41,7 @@ def find_path(name: str, path: str = None) -> Union[Path, None]:
if name in os.listdir(path):
path_name = os.path.join(path, name)
print(f"{name} found: {path_name}")
return Path(path_name)
return path_name
# Get the parent directory
parent_directory = os.path.dirname(path)
@@ -63,7 +63,8 @@ def add_comfyui_directory_to_sys_path() -> None:
sys.path.append(comfyui_path)
print(f"'{comfyui_path}' added to sys.path")
def add_extra_model_paths() -> Path:
def add_extra_model_paths() -> None:
"""
Parse the optional extra_model_paths.yaml file and add the parsed paths to the sys.path.
"""