From 22ba261dcfc97098c90d78edceb01b30cfb6e8ae Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lucas=20=20Sk=C3=B6ld?= Date: Tue, 27 Feb 2024 18:55:48 +0100 Subject: [PATCH] added custom_lora_loaders --- README.md | 11 +++++++--- __init__.py | 36 ++++++++++++++++++++++++++++++++ config.json.example | 6 +++++- src/dynamic_lora_loader.py | 42 ++++++++++++++++++++++++++++++++++++++ 4 files changed, 91 insertions(+), 4 deletions(-) create mode 100644 src/dynamic_lora_loader.py diff --git a/README.md b/README.md index cf4b976..3263fb4 100644 --- a/README.md +++ b/README.md @@ -33,9 +33,13 @@ To get started, create a `config.json` file in the root directory of the `sn0w-s ```json { - "logging_level": "GENERAL", - "animagine_formatting": "TRUE", - "excluded_lora_folders": "character, concept" + "logging_level": "GENERAL", + "animagine_formatting": "TRUE", + "excluded_lora_folders": "character, concept", + "custom_lora_loaders_xl": { + "Loader1": "character, style", + "name": "included folders" + } } ``` @@ -48,6 +52,7 @@ Below is a table detailing the configuration options available in `config.json`, | `logging_level` | Determines the verbosity of logs generated by the scripts. | `NONE`, `GENERAL`, `ALL` | `GENERAL` | | `animagine_formatting` | Enables or disables specific formatting rules for prompts, such as prioritizing certain tags. | `TRUE`, `FALSE` | `FALSE` | | `excluded_lora_folders` | Removes folders from custom lora loaders. Folders are seperated by comma. | `any`, `folder`, `style` | | +| `custom_lora_loaders(_xl)` | Creates custom lora loaders, filtering out all folders except ones in the config. | `any`, `folder`, `style` | | # Nodes diff --git a/__init__.py b/__init__.py index 8e27475..05468d2 100644 --- a/__init__.py +++ b/__init__.py @@ -1,3 +1,7 @@ +import re +from .src.dynamic_lora_loader import generate_lora_node_class +from .src.sn0w import ConfigReader, Logger + from .src.find_resolution import FindResolutionNode from .src.lora_selector import LoraSelectorNode from .src.lora_tester_xl import LoraTestXLNode @@ -52,3 +56,35 @@ NODE_DISPLAY_NAME_MAPPINGS = { "Prompt Selector": "Prompt Selector", "Load Lora Concept": "Load Lora Concept", } + +current_unique_id = 0 # Global variable to track the unique ID + +def parse_custom_lora_loaders(custom_lora_loaders): + required_folders_with_names = [] + for name, folders_str in custom_lora_loaders.items(): + # Treat every value as a comma-separated list + folders = [folder.strip() for folder in folders_str.split(',')] + required_folders_with_names.append((name, folders)) + return required_folders_with_names + + +def generate_and_register_lora_node(lora_type, setting): + global current_unique_id # Use the global variable to keep track of IDs + logger = Logger() + + custom_lora_loaders = ConfigReader.get_setting(setting, {}) + required_folders_with_names = parse_custom_lora_loaders(custom_lora_loaders) + + for name, folders in required_folders_with_names: + current_unique_id += 1 # Increment the unique ID for each new class + unique_id_with_name = str(f"{name}_{current_unique_id}") + + logger.log(f"Adding custom lora loader: {folders}, {unique_id_with_name}", "ALL") + DynamicLoraNode = generate_lora_node_class(lora_type, folders) + + # Update NODE_CLASS_MAPPINGS and NODE_DISPLAY_NAME_MAPPINGS for each generated class + NODE_CLASS_MAPPINGS[unique_id_with_name] = DynamicLoraNode + NODE_DISPLAY_NAME_MAPPINGS[unique_id_with_name] = f"{name}" + +generate_and_register_lora_node("loras_xl", "custom_lora_loaders_xl") +generate_and_register_lora_node("loras_15", "custom_lora_loaders") \ No newline at end of file diff --git a/config.json.example b/config.json.example index 04531c5..a628d8d 100644 --- a/config.json.example +++ b/config.json.example @@ -1,5 +1,9 @@ { "logging_level": "GENERAL", "animagine_formatting": "TRUE", - "excluded_lora_folders": "character, concept" + "excluded_lora_folders": "character, concept", + "custom_lora_loaders_xl": { + "Loader1": "character, style", + "Loader2": "concept" + } } \ No newline at end of file diff --git a/src/dynamic_lora_loader.py b/src/dynamic_lora_loader.py new file mode 100644 index 0000000..d03560c --- /dev/null +++ b/src/dynamic_lora_loader.py @@ -0,0 +1,42 @@ +import folder_paths +from nodes import LoraLoader +from pathlib import Path +from .sn0w import ConfigReader + +def generate_lora_node_class(lora_type, required_folders): + class DynamicLoraNode: + @classmethod + def INPUT_TYPES(cls): + # Sort the loras_xl list alphabetically + loras_xl = folder_paths.get_filename_list(lora_type) + sorted_loras_xl = sorted(loras_xl, key=lambda p: [part.lower() for part in Path(p).parts]) + + # Filter sorted_loras_xl to include items containing any of the required folder names + filtered_sorted_loras_xl = [ + lora for lora in sorted_loras_xl + if any(folder.lower() in (part.lower() for part in Path(lora).parts) for folder in required_folders) + ] + + return { + "required": { + "model": ("MODEL",), + "clip": ("CLIP",), + "lora": (['None'] + filtered_sorted_loras_xl,), + "lora_strength": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}), + }, + } + + RETURN_TYPES = ("MODEL", "CLIP",) + RETURN_NAMES = ("MODEL", "CLIP",) + FUNCTION = "find_lora" + CATEGORY = "sn0w" + OUTPUT_NODE = True + + def find_lora(self, model, clip, lora, lora_strength): + lora_loader = LoraLoader() + full_loras_list = folder_paths.get_filename_list("loras") + full_lora_path = next((full_path for full_path in full_loras_list if lora in full_path), None) + modified_model, modified_clip = lora_loader.load_lora(model, clip, full_lora_path, lora_strength, lora_strength) + return (modified_model, modified_clip, ) + + return DynamicLoraNode