Files
sn0w12-ComfyUI-Sn0w-Scripts/src/lora_stacker.py
T
2024-06-29 18:03:19 +02:00

61 lines
2.4 KiB
Python

import os
import folder_paths
class LoraStackerNode:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"lora_a": (["None"] + folder_paths.get_filename_list("loras"),),
"lora_strength_a": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
"lora_b": (["None"] + folder_paths.get_filename_list("loras"),),
"lora_strength_b": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
"lora_c": (["None"] + folder_paths.get_filename_list("loras"),),
"lora_strength_c": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
"lora_d": (["None"] + folder_paths.get_filename_list("loras"),),
"lora_strength_d": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
"add_default_generation": ("BOOLEAN", {"default": False}),
},
}
RETURN_TYPES = ("STRING", "INT")
RETURN_NAMES = ("LORA_INFO", "TOTAL_LORAS")
FUNCTION = "process_loras"
CATEGORY = "sn0w/lora"
OUTPUT_NODE = True
def process_loras(
self, lora_a, lora_strength_a, lora_b, lora_strength_b, lora_c, lora_strength_c, lora_d, lora_strength_d, add_default_generation
):
# Function to extract lora string and format with strength
def format_lora(lora, strength):
if lora.lower() == "none": # Skip "None" loras
return None
lora_string = os.path.splitext(os.path.basename(lora))[0]
return f"<lora:{lora_string}:{strength:.1f}>"
# Process each lora and its strength, filtering out "None"
loras = [
format_lora(lora_a, lora_strength_a),
format_lora(lora_b, lora_strength_b),
format_lora(lora_c, lora_strength_c),
format_lora(lora_d, lora_strength_d),
]
# Filter out None values from the loras list
loras = [lora for lora in loras if lora is not None]
# Initialize total_loras count based on filtered loras
total_loras = len(loras)
# Handle additional generation if required
if add_default_generation:
total_loras += 1
loras.append("Nothing")
# Join the lora strings into a single string separated by semicolons
lora_output = ";".join(loras)
return (lora_output, total_loras)