diff --git a/src/load_lora_from_folder.py b/src/load_lora_from_folder.py index 674e7cb..443f548 100644 --- a/src/load_lora_from_folder.py +++ b/src/load_lora_from_folder.py @@ -2,7 +2,7 @@ import os import re import folder_paths from nodes import LoraLoader -from .sn0w import Logger, Utility +from .sn0w import Logger, Utility, ConfigReader class LoadLoraFolderNode: logger = Logger() @@ -57,51 +57,60 @@ class LoadLoraFolderNode: return master_folder, include_folders, exclude_folders def find_and_apply_lora(self, model, clip, prompt, folders, xl, lora_strength, separator): + # Set default prompt if None provided if prompt is None: prompt = '' + # Clean and split the prompt into parts prompt_parts = [self.clean_string(part) for part in prompt.split(separator)] - # Retrieve paths and apply folder constraints + # Retrieve and filter Lora file paths full_lora_paths = folder_paths.get_filename_list("loras") filtered_lora_paths = folder_paths.get_filename_list("loras_xl" if xl else "loras_15") master_folder, include_folders, exclude_folders = self.parse_folders(folders) - # Filter by master folder if specified + # Include only paths in the master folder, if specified if master_folder: filtered_lora_paths = [path for path in filtered_lora_paths if master_folder in self.normalize_folder_name(path)] - # Filter by included and excluded folders + # Further filter paths to include and exclude specific folders lora_paths = [path for path in filtered_lora_paths if any(folder in self.normalize_folder_name(path) for folder in include_folders)] lora_paths = [path for path in lora_paths if not any(exclude in self.normalize_folder_name(path) for exclude in exclude_folders)] lora_candidates = {} loaded_loras = set() + max_distance = ConfigReader.get_setting('sn0w.LoraFolderMinDistance', 5) + + # Match prompt parts with Lora filenames for prompt_part in prompt_parts: for lora_path in lora_paths: lora_filename = os.path.split(lora_path)[-1].lower() + # Clean up the filename for matching processed_lora_filename = lora_filename.replace(".safetensors", "").replace("_", " ") + # Check if prompt part matches the cleaned filename if prompt_part in processed_lora_filename: distance = Utility.levenshtein_distance(prompt_part, processed_lora_filename) + # Store matching Loras if within acceptable distance for full_path in full_lora_paths: - if lora_filename in full_path.lower(): - if distance <= 5: - if prompt_part not in lora_candidates: - lora_candidates[prompt_part] = [] - lora_candidates[prompt_part].append({'full_path': full_path, 'distance': distance}) - self.logger.log("Distance: " + str(distance) + " Lora: " + lora_filename + " Tag: " + prompt_part, "ALL") - break + if lora_filename in full_path.lower() and distance <= max_distance: + if prompt_part not in lora_candidates: + lora_candidates[prompt_part] = [] + lora_candidates[prompt_part].append({'full_path': full_path, 'distance': distance}) + self.logger.log("Distance: " + str(distance) + " Lora: " + lora_filename + " Tag: " + prompt_part, "ALL") + break + # Load the best candidate Lora for each prompt part for prompt_part, candidates in lora_candidates.items(): if candidates: selected_candidate = min(candidates, key=lambda x: x['distance']) folder_name = self.normalize_folder_name(os.path.dirname(selected_candidate['full_path'])) + # Load Lora if not already loaded and within the allowed number per folder if selected_candidate['full_path'] not in loaded_loras and (folder_name not in include_folders or len(loaded_loras) < include_folders[folder_name]): self.logger.log(f"Loading Lora: {os.path.split(selected_candidate['full_path'])[-1]}", "GENERAL") - model, clip = LoraLoader().load_lora(model, clip, selected_candidate['full_path'], lora_strength, lora_strength) + model, clip = LoraLoader().load_lora(model, clip, selected_candidate['full_path'], lora_strength) loaded_loras.add(selected_candidate['full_path']) else: self.logger.log(f"No matching Lora found for any tags.", "GENERAL") - return (model, clip,) + return (model, clip)