From 69b5adb6380c4fe7b0d92d9ee2d2b19a4e72b69e Mon Sep 17 00:00:00 2001 From: sn0w12 Date: Fri, 16 Feb 2024 14:02:13 +0100 Subject: [PATCH] Added load lora concept, added print_sn0w --- README.md | 9 ++++ __init__.py | 3 ++ src/load_lora_character.py | 13 ++--- src/load_lora_concept.py | 101 +++++++++++++++++++++++++++++++++++++ src/print_sn0w.py | 4 ++ 5 files changed, 122 insertions(+), 8 deletions(-) create mode 100644 src/load_lora_concept.py create mode 100644 src/print_sn0w.py diff --git a/README.md b/README.md index aa42d99..b93bee1 100644 --- a/README.md +++ b/README.md @@ -16,6 +16,14 @@ A collection of nodes and improvements created mainly for testing loras and gene ``` 3. Start up ComfyUI. +# Important Note +Many of the nodes require these paths to be in your extra_model_paths.yaml (with your real directories of course). + +- loras_xl: C:/path/XL +- loras_15: C:/path/1.5 + +![Example](./imgs/lora_paths_example.png) + # Nodes ## Find Resolution @@ -83,6 +91,7 @@ A collection of nodes and improvements created mainly for testing loras and gene > ℹ️ See More Information > > - Basically Lora Selector but manual. +> - Specifically for the lora testers. > > ![Lora Stacker](./imgs/lora_stacker.png) > diff --git a/__init__.py b/__init__.py index 3555e4f..2498b35 100644 --- a/__init__.py +++ b/__init__.py @@ -13,6 +13,7 @@ from .src.load_lora_character import LoadLoraCharacterNode from .src.get_font_size import GetFontSizeNode from .src.get_all_styles import GetStylesNode from .src.prompt_selector import PromptSelectNode +from .src.load_lora_concept import LoadLoraConceptNode NODE_CLASS_MAPPINGS = { "Find SDXL Resolution": FindResolutionNode, @@ -30,6 +31,7 @@ NODE_CLASS_MAPPINGS = { "Get Font Size": GetFontSizeNode, "Get All Styles": GetStylesNode, "Prompt Selector": PromptSelectNode, + "Load Lora Concept": LoadLoraConceptNode, } NODE_DISPLAY_NAME_MAPPINGS = { @@ -48,4 +50,5 @@ NODE_DISPLAY_NAME_MAPPINGS = { "Get Font Size": "Get Font Size", "Get All Styles": "Get All Styles", "Prompt Selector": "Prompt Selector", + "Load Lora Concept": "Load Lora Concept", } diff --git a/src/load_lora_character.py b/src/load_lora_character.py index 53734b4..7666da9 100644 --- a/src/load_lora_character.py +++ b/src/load_lora_character.py @@ -3,6 +3,7 @@ import json import re import folder_paths from nodes import LoraLoader +from .print_sn0w import print_sn0w class LoadLoraCharacterNode: @classmethod @@ -53,10 +54,7 @@ class LoadLoraCharacterNode: full_lora_path = folder_paths.get_filename_list("loras") # Select the appropriate lora list based on 'xl' - if xl: - lora_paths = folder_paths.get_filename_list("loras_xl") - else: - lora_paths = folder_paths.get_filename_list("loras_15") + lora_paths = folder_paths.get_filename_list("loras_xl" if xl else "loras_15") # Extract just the filenames for comparison lora_filenames = [path.split('\\')[-1] for path in lora_paths] @@ -75,10 +73,9 @@ class LoadLoraCharacterNode: if lora_path: break - print(lora_path) + print_sn0w(lora_path) if lora_path: - modified_model, modified_clip = lora_loader.load_lora(model, clip, lora_path, lora_strength, lora_strength) - return (modified_model, modified_clip, ) + model, clip = lora_loader.load_lora(model, clip, lora_path, lora_strength, lora_strength) - return (model, clip, ) # if character lora not found, just return the normal model. \ No newline at end of file + return (model, clip, ) \ No newline at end of file diff --git a/src/load_lora_concept.py b/src/load_lora_concept.py new file mode 100644 index 0000000..56e1376 --- /dev/null +++ b/src/load_lora_concept.py @@ -0,0 +1,101 @@ +import os +import json +import re +import folder_paths +from nodes import LoraLoader +from .print_sn0w import print_sn0w + +class LoadLoraConceptNode: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model": ("MODEL",), + "clip": ("CLIP", ), + "prompt": ("STRING", {"default": ""}), + "xl": ("BOOLEAN", {"default": False},), + "lora_strength": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}), + "separator": ("STRING", {"default": ", "}), + }, + } + + RETURN_TYPES = ("MODEL", "CLIP",) + RETURN_NAMES = ("MODEL", "CLIP",) + FUNCTION = "find_and_apply_lora" + CATEGORY = "sn0w" + + def clean_string(self, input_string): + cleaned_string = re.sub(r'[\\()]', '', input_string) + cleaned_string = re.sub(r':\d+(\.\d+)?', '', cleaned_string) + cleaned_string = re.sub(r',$', '', cleaned_string.strip()) + return cleaned_string.lower() + + def similarity_ratio(self, str1, str2): + """ + Calculates a basic similarity ratio between two strings based on the length + of the longest common subsequence. This is a simplified and less accurate + approach compared to algorithms like Levenshtein distance. + + Args: + - str1: First string. + - str2: Second string. + + Returns: + - A percentage representing the similarity between the two strings. + """ + def longest_common_subsequence(s1, s2): + m, n = len(s1), len(s2) + dp = [[0] * (n + 1) for _ in range(m + 1)] + + for i in range(1, m + 1): + for j in range(1, n + 1): + if s1[i - 1] == s2[j - 1]: + dp[i][j] = dp[i - 1][j - 1] + 1 + else: + dp[i][j] = max(dp[i - 1][j], dp[i][j - 1]) + + return dp[m][n] + + # Normalize and prepare strings (remove file extension, replace underscores) + normalized_str1 = str1.replace(".safetensors", "").replace("_", " ").lower() + normalized_str2 = str2.replace(".safetensors", "").replace("_", " ").lower() + + # Calculate the length of the longest common subsequence + lcs_len = longest_common_subsequence(normalized_str1, normalized_str2) + max_len = max(len(normalized_str1), len(normalized_str2)) + + # Return the similarity ratio + return (lcs_len / max_len) * 100 if max_len > 0 else 0 + + def find_and_apply_lora(self, model, clip, prompt, xl, lora_strength, separator): + if prompt is None: + prompt = '' + prompt_parts = [self.clean_string(part) for part in prompt.split(separator)] + + # Retrieve the full list of lora paths + full_lora_paths = folder_paths.get_filename_list("loras") + lora_paths = folder_paths.get_filename_list("loras_xl" if xl else "loras_15") + # lora_paths = [path for path in filtered_lora_paths if "concept" in path.lower()] + + found_full_lora_paths = [] + + for prompt_part in prompt_parts: + for lora_path in lora_paths: + lora_filename = os.path.split(lora_path)[-1].lower() + # Remove file extension and replace underscores with spaces for comparison + processed_lora_filename = lora_filename.replace(".safetensors", "").replace("_", " ") + # Calculate similarity ratio + similarity_ratio = self.similarity_ratio(prompt_part, processed_lora_filename) + if similarity_ratio >= 85: + # Find the full path from the full_lora_paths list + for full_path in full_lora_paths: + if lora_filename in full_path.lower(): + found_full_lora_paths.append(full_path) + break + + # Print and apply found lora paths + for path in found_full_lora_paths: + print_sn0w(path) + model, clip = LoraLoader().load_lora(model, clip, path, lora_strength, lora_strength) + + return (model, clip, ) \ No newline at end of file diff --git a/src/print_sn0w.py b/src/print_sn0w.py new file mode 100644 index 0000000..8c2a653 --- /dev/null +++ b/src/print_sn0w.py @@ -0,0 +1,4 @@ +def print_sn0w(message, prefix="[sn0w] "): + RED_TEXT = "\033[0;35m" + RESET_TEXT = "\033[0m" + print(f"{RED_TEXT}{prefix}{RESET_TEXT}{message}")