122 lines
4.2 KiB
Python
122 lines
4.2 KiB
Python
import os
|
|
import random
|
|
import folder_paths
|
|
|
|
class LoraSelector:
|
|
def __init__(self):
|
|
self.current_index = 0
|
|
self.current_count = 0
|
|
self.last_folder = None
|
|
self.last_selection = None
|
|
self.last_manual_index = 0
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
# Find 'Illustrious' folders in all lora paths
|
|
lora_paths = folder_paths.get_folder_paths("loras")
|
|
illustrious_folders = []
|
|
|
|
for path in lora_paths:
|
|
illustrious_path = os.path.join(path, "Illustrious")
|
|
if os.path.exists(illustrious_path) and os.path.isdir(illustrious_path):
|
|
# List subdirectories
|
|
try:
|
|
subdirs = [d for d in os.listdir(illustrious_path) if os.path.isdir(os.path.join(illustrious_path, d))]
|
|
illustrious_folders.extend(subdirs)
|
|
except Exception:
|
|
pass
|
|
|
|
# Remove duplicates and sort
|
|
illustrious_folders = sorted(list(set(illustrious_folders)))
|
|
|
|
if not illustrious_folders:
|
|
illustrious_folders = ["None"]
|
|
|
|
return {
|
|
"required": {
|
|
"folder": (illustrious_folders, ),
|
|
"mode": (["Random", "Sequential", "Manual"], {"default": "Random"}),
|
|
"repeat_count": ("INT", {"default": 1, "min": 1, "max": 100, "step": 1}),
|
|
"manual_index": ("INT", {"default": 0, "min": 0, "max": 10000, "step": 1}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "INT")
|
|
RETURN_NAMES = ("lora_name", "total_loras")
|
|
FUNCTION = "select_lora"
|
|
CATEGORY = "AODH Pack"
|
|
|
|
def select_lora(self, folder, mode, repeat_count, manual_index):
|
|
repeat_count = max(1, repeat_count)
|
|
if folder == "None":
|
|
return ("", 0)
|
|
|
|
# Find the full path for the selected folder
|
|
lora_paths = folder_paths.get_folder_paths("loras")
|
|
target_path = None
|
|
|
|
for path in lora_paths:
|
|
check_path = os.path.join(path, "Illustrious", folder)
|
|
if os.path.exists(check_path) and os.path.isdir(check_path):
|
|
target_path = check_path
|
|
break
|
|
|
|
if not target_path:
|
|
return ("", 0)
|
|
|
|
# Find .safetensors files
|
|
lora_files = [f for f in os.listdir(target_path) if f.endswith('.safetensors')]
|
|
lora_files.sort() # Ensure consistent order
|
|
|
|
count = len(lora_files)
|
|
if count == 0:
|
|
return ("", 0)
|
|
|
|
selected_lora = ""
|
|
|
|
# Reset state if folder changes
|
|
if self.last_folder != folder:
|
|
self.current_index = 0
|
|
self.current_count = 0
|
|
self.last_selection = None
|
|
self.last_folder = folder
|
|
|
|
# Reset state if manual_index changes (for Sequential mode)
|
|
if self.last_manual_index != manual_index:
|
|
self.current_index = 0
|
|
self.current_count = 0
|
|
self.last_manual_index = manual_index
|
|
|
|
if mode == "Random":
|
|
if self.current_count >= repeat_count or self.last_selection is None:
|
|
# Use a local Random instance to ensure randomness regardless of global seed
|
|
rng = random.Random()
|
|
self.last_selection = rng.choice(lora_files)
|
|
self.current_count = 0
|
|
|
|
selected_lora = self.last_selection
|
|
self.current_count += 1
|
|
|
|
elif mode == "Sequential":
|
|
if self.current_count >= repeat_count:
|
|
self.current_index += 1
|
|
self.current_count = 0
|
|
|
|
selected_lora = lora_files[(self.current_index + manual_index) % count]
|
|
self.current_count += 1
|
|
|
|
elif mode == "Manual":
|
|
selected_lora = lora_files[manual_index % count]
|
|
|
|
# Construct the relative path for ComfyUI Lora loader
|
|
# Use forward slashes for consistency
|
|
full_lora_name = f"Illustrious/{folder}/{selected_lora}"
|
|
|
|
return (full_lora_name, count)
|
|
|
|
@classmethod
|
|
def IS_CHANGED(s, folder, mode, repeat_count, manual_index):
|
|
if mode == "Random" or mode == "Sequential":
|
|
return float("nan")
|
|
return f"{folder}_{manual_index}"
|