124 lines
4.1 KiB
Python
124 lines
4.1 KiB
Python
import json
|
|
import os
|
|
|
|
import folder_paths
|
|
import nodes
|
|
|
|
|
|
class Loader:
|
|
def __init__(self):
|
|
pass
|
|
|
|
def __log(self, text):
|
|
print("\033[92m[Allor]\033[0m: " + text)
|
|
|
|
def __get_config(self):
|
|
root = os.path.dirname(os.path.abspath(__file__))
|
|
|
|
with open(os.path.join(root, "config.json"), "r") as f:
|
|
return json.load(f)
|
|
|
|
config = __get_config
|
|
|
|
def __get_fonts_folder_path(self):
|
|
return os.path.join(folder_paths.base_path, *self.config()["fonts_folder_path"])
|
|
|
|
def setup_rembg(self):
|
|
os.environ["U2NET_HOME"] = folder_paths.models_dir + "/onnx"
|
|
|
|
def setup_paths(self):
|
|
fonts_folder_path = self.__get_fonts_folder_path()
|
|
|
|
folder_paths.folder_names_and_paths["onnx"] = ([os.path.join(folder_paths.models_dir, "onnx")], {".onnx"})
|
|
folder_paths.folder_names_and_paths["fonts"] = ([fonts_folder_path], {".otf", ".ttf"})
|
|
|
|
if not os.path.exists(fonts_folder_path):
|
|
os.mkdir(fonts_folder_path)
|
|
|
|
def setup_override(self):
|
|
override_nodes_len = 0
|
|
|
|
def override(function):
|
|
start_len = nodes.NODE_CLASS_MAPPINGS.__len__()
|
|
|
|
nodes.NODE_CLASS_MAPPINGS = dict(
|
|
filter(function, nodes.NODE_CLASS_MAPPINGS.items())
|
|
)
|
|
|
|
return start_len - nodes.NODE_CLASS_MAPPINGS.__len__()
|
|
|
|
if self.config()["override"]["postprocessing"]:
|
|
override_nodes_len += override(lambda item: not item[1].CATEGORY.startswith("image/postprocessing"))
|
|
|
|
if self.config()["override"]["transform"]:
|
|
override_nodes_len += override(lambda item: not item[0] == "ImageScale" and not item[0] == "ImageScaleBy" and not item[0] == "ImageInvert")
|
|
|
|
if self.config()["override"]["debug"]:
|
|
nodes.VAEDecodeTiled.CATEGORY = "latent"
|
|
nodes.VAEEncodeTiled.CATEGORY = "latent"
|
|
|
|
override_nodes_len += override(lambda item: not item[1].CATEGORY.startswith("_for_testing"))
|
|
|
|
self.__log(str(override_nodes_len) + " standard nodes was overridden.")
|
|
|
|
def get_modules(self):
|
|
modules = dict()
|
|
|
|
if self.config()["modules"]["AlphaChanel"]:
|
|
from .modules import AlphaChanel
|
|
modules.update(AlphaChanel.NODE_CLASS_MAPPINGS)
|
|
|
|
if self.config()["modules"]["Clamp"]:
|
|
from .modules import Clamp
|
|
modules.update(Clamp.NODE_CLASS_MAPPINGS)
|
|
|
|
if self.config()["modules"]["ImageBatch"]:
|
|
from .modules import ImageBatch
|
|
modules.update(ImageBatch.NODE_CLASS_MAPPINGS)
|
|
|
|
if self.config()["modules"]["ImageComposite"]:
|
|
from .modules import ImageComposite
|
|
modules.update(ImageComposite.NODE_CLASS_MAPPINGS)
|
|
|
|
if self.config()["modules"]["ImageContainer"]:
|
|
from .modules import ImageContainer
|
|
modules.update(ImageContainer.NODE_CLASS_MAPPINGS)
|
|
|
|
if self.config()["modules"]["ImageDraw"]:
|
|
from .modules import ImageDraw
|
|
modules.update(ImageDraw.NODE_CLASS_MAPPINGS)
|
|
|
|
if self.config()["modules"]["ImageEffects"]:
|
|
from .modules import ImageEffects
|
|
modules.update(ImageEffects.NODE_CLASS_MAPPINGS)
|
|
|
|
if self.config()["modules"]["ImageFilter"]:
|
|
from .modules import ImageFilter
|
|
modules.update(ImageFilter.NODE_CLASS_MAPPINGS)
|
|
|
|
if self.config()["modules"]["ImageSegmentation"]:
|
|
from .modules import ImageSegmentation
|
|
modules.update(ImageSegmentation.NODE_CLASS_MAPPINGS)
|
|
|
|
if self.config()["modules"]["ImageText"]:
|
|
from .modules import ImageText
|
|
modules.update(ImageText.NODE_CLASS_MAPPINGS)
|
|
|
|
if self.config()["modules"]["ImageTransform"]:
|
|
from .modules import ImageTransform
|
|
modules.update(ImageTransform.NODE_CLASS_MAPPINGS)
|
|
|
|
modules_len = dict(
|
|
filter(
|
|
lambda item: item[1],
|
|
self.config()["modules"].items()
|
|
)
|
|
).__len__()
|
|
|
|
nodes_len = modules.__len__()
|
|
|
|
self.__log(str(modules_len) + " modules enabled.")
|
|
self.__log(str(nodes_len) + " nodes was loaded.")
|
|
|
|
return modules
|