nodes/sn0w12-ComfyUI-Sn0w-Scripts
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Added comments and made max distance customizable

Browse Source
This commit is contained in:
Lucas Sköld
2024-04-23 21:42:49 +02:00
parent dd2c595b98
commit f37deb085e
1 changed files with 22 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/load_lora_from_folder.py
+22 -13
View File
@@ -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)