149 lines
5.7 KiB
Python
149 lines
5.7 KiB
Python
import importlib
|
|
import os
|
|
from .src.dynamic_lora_loader import generate_lora_node_class
|
|
from .src.dynamic_scheduler_loader import generate_scheduler_node_class
|
|
from .src.check_folder_paths import check_lora_folders
|
|
from .sn0w import ConfigReader, Logger
|
|
|
|
from .src.lora_selector import LoraSelectorNode
|
|
from .src.lora_tester import LoraTestNode
|
|
from .src.character_select import CharacterSelectNode
|
|
from .src.prompt_combine import CombineStringNode
|
|
from .src.lora_stacker import LoraStackerNode
|
|
from .src.get_font_size import GetFontSizeNode
|
|
from .src.prompt_selector import PromptSelectNode
|
|
from .src.upscale_with_model_by import UpscaleImageBy
|
|
from .src.load_lora_from_folder import LoadLoraFolderNode
|
|
from .src.textbox import TextboxNode
|
|
from .src.simple_ksampler import SimpleSamplerCustom
|
|
from .src.filter_tags import FilterTags
|
|
|
|
# Constants
|
|
WEB_DIRECTORY = "./web"
|
|
CURRENT_UNIQUE_ID = 0 # Global variable to track the unique ID
|
|
logger = Logger()
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"Lora Selector": LoraSelectorNode,
|
|
"Lora Tester": LoraTestNode,
|
|
"Character Selector": CharacterSelectNode,
|
|
"Prompt Combine": CombineStringNode,
|
|
"Sn0w Lora Stacker": LoraStackerNode,
|
|
"Load Lora Sn0w": generate_lora_node_class("loras"),
|
|
"Get Font Size": GetFontSizeNode,
|
|
"Prompt Selector": PromptSelectNode,
|
|
"Upscale Image With Model By": UpscaleImageBy,
|
|
"Load Lora Folder": LoadLoraFolderNode,
|
|
"Copy/Paste Textbox": TextboxNode,
|
|
"Sn0w KSampler": SimpleSamplerCustom,
|
|
"Filter Tags": FilterTags,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"Lora Selector": "Lora Selector",
|
|
"Lora Tester": "Lora Tester",
|
|
"Character Selector": "Character Selector",
|
|
"Prompt Combine": "Prompt Combine",
|
|
"Sn0w Lora Stacker": "Sn0w Lora Stacker",
|
|
"Load Lora Sn0w": "Load Lora Sn0w",
|
|
"Get Font Size": "Get Font Size",
|
|
"Prompt Selector": "Prompt Selector",
|
|
"Upscale Image With Model By": "Upscale Image With Model By",
|
|
"Load Lora Folder": "Load Lora Folder",
|
|
"Copy/Paste Textbox": "Textbox",
|
|
"Sn0w KSampler": "Sn0w KSampler",
|
|
"Filter Tags": "Filter Tags",
|
|
}
|
|
|
|
# Function to check required folder paths
|
|
check_lora_folders()
|
|
|
|
|
|
def parse_custom_lora_loaders(custom_lora_loaders):
|
|
"""
|
|
Parse a string of custom Lora loaders into a list of tuples containing
|
|
the name, required folders, and the number of combinations.
|
|
"""
|
|
required_folders_with_names = []
|
|
entries = custom_lora_loaders.split("\n")
|
|
for entry in entries:
|
|
if entry.strip():
|
|
name, _, remainder = entry.partition(":")
|
|
if ":" in remainder:
|
|
value, _, combos = remainder.partition(":")
|
|
else:
|
|
value = remainder
|
|
combos = 1
|
|
required_folders_with_names.append((name.strip(), value.strip().split(","), int(combos)))
|
|
return required_folders_with_names
|
|
|
|
|
|
def generate_and_register_lora_node(lora_type, setting):
|
|
"""Generate and register the users custom lora loaders"""
|
|
global CURRENT_UNIQUE_ID
|
|
|
|
custom_lora_loaders = ConfigReader.get_setting(setting, None)
|
|
if custom_lora_loaders is not None:
|
|
required_folders_with_names = parse_custom_lora_loaders(custom_lora_loaders)
|
|
for name, folders, combos in required_folders_with_names:
|
|
CURRENT_UNIQUE_ID += 1
|
|
unique_id_with_name = f"{name}_{CURRENT_UNIQUE_ID}"
|
|
|
|
logger.log(f"Adding custom lora loader. Path: {folders}, Name: {name}, Inputs: {combos}", "INFORMATIONAL")
|
|
DynamicLoraNode = generate_lora_node_class(lora_type, folders, combos)
|
|
|
|
if name in NODE_CLASS_MAPPINGS:
|
|
NODE_CLASS_MAPPINGS[unique_id_with_name] = DynamicLoraNode
|
|
NODE_DISPLAY_NAME_MAPPINGS[unique_id_with_name] = f"{name}"
|
|
else:
|
|
NODE_CLASS_MAPPINGS[name] = DynamicLoraNode
|
|
NODE_DISPLAY_NAME_MAPPINGS[name] = f"{name}"
|
|
|
|
|
|
def generate_and_register_all_lora_nodes():
|
|
"""Generate and register all custom lora nodes."""
|
|
lora_types = [
|
|
("loras_xl", "sn0w.CustomLoraLoadersXL"),
|
|
("loras_15", "sn0w.CustomLoraLoaders1.5"),
|
|
("loras_3", "sn0w.CustomLoraLoaders3"),
|
|
]
|
|
for lora_type, setting in lora_types:
|
|
generate_and_register_lora_node(lora_type, setting)
|
|
|
|
|
|
def import_and_register_scheduler_nodes():
|
|
"""
|
|
Import and register custom scheduler nodes from Python files in the
|
|
'src/custom_schedulers' directory.
|
|
"""
|
|
custom_schedulers_path = os.path.join(os.path.dirname(__file__), "src", "custom_schedulers")
|
|
for filename in os.listdir(custom_schedulers_path):
|
|
if filename.endswith(".py") and filename != "custom_schedulers.py":
|
|
module_name = filename[:-3]
|
|
module = import_module_from_path(module_name, os.path.join(custom_schedulers_path, filename))
|
|
register_scheduler_node(module)
|
|
|
|
|
|
def import_module_from_path(module_name, file_path):
|
|
"""Helper function to import a module from a file path."""
|
|
spec = importlib.util.spec_from_file_location(module_name, file_path)
|
|
module = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
def register_scheduler_node(module):
|
|
"""Register a scheduler node if it contains settings and a get_sigmas function."""
|
|
settings = getattr(module, "settings", None)
|
|
get_sigmas_function = getattr(module, "get_sigmas", None)
|
|
if settings and get_sigmas_function:
|
|
DynamicSchedulerNode = generate_scheduler_node_class(settings, get_sigmas_function)
|
|
class_name = settings.get("name", "default").capitalize() + "Scheduler"
|
|
NODE_CLASS_MAPPINGS[class_name] = DynamicSchedulerNode
|
|
NODE_DISPLAY_NAME_MAPPINGS[class_name] = class_name
|
|
|
|
|
|
# Register custom nodes
|
|
generate_and_register_all_lora_nodes()
|
|
import_and_register_scheduler_nodes()
|