121 lines
4.2 KiB
Python
121 lines
4.2 KiB
Python
import os
|
|
import random
|
|
import folder_paths
|
|
|
|
class CheckpointSelector:
|
|
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 folders in 'checkpoints' paths
|
|
checkpoint_paths = folder_paths.get_folder_paths("checkpoints")
|
|
checkpoint_folders = []
|
|
|
|
for path in checkpoint_paths:
|
|
if os.path.exists(path) and os.path.isdir(path):
|
|
# List subdirectories
|
|
try:
|
|
subdirs = [d for d in os.listdir(path) if os.path.isdir(os.path.join(path, d))]
|
|
checkpoint_folders.extend(subdirs)
|
|
except Exception:
|
|
pass
|
|
|
|
# Remove duplicates and sort
|
|
checkpoint_folders = sorted(list(set(checkpoint_folders)))
|
|
|
|
if not checkpoint_folders:
|
|
checkpoint_folders = ["None"]
|
|
|
|
return {
|
|
"required": {
|
|
"folder": (checkpoint_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 = ("checkpoint_name", "total_checkpoints")
|
|
FUNCTION = "select_checkpoint"
|
|
CATEGORY = "AODH Pack"
|
|
|
|
def select_checkpoint(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
|
|
checkpoint_paths = folder_paths.get_folder_paths("checkpoints")
|
|
target_path = None
|
|
|
|
for path in checkpoint_paths:
|
|
check_path = os.path.join(path, 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 checkpoint files (.safetensors, .ckpt)
|
|
extensions = ('.safetensors', '.ckpt')
|
|
checkpoint_files = [f for f in os.listdir(target_path) if f.lower().endswith(extensions)]
|
|
checkpoint_files.sort()
|
|
|
|
count = len(checkpoint_files)
|
|
if count == 0:
|
|
return ("", 0)
|
|
|
|
selected_checkpoint = ""
|
|
|
|
# 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:
|
|
rng = random.Random()
|
|
self.last_selection = rng.choice(checkpoint_files)
|
|
self.current_count = 0
|
|
|
|
selected_checkpoint = 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_checkpoint = checkpoint_files[(self.current_index + manual_index) % count]
|
|
self.current_count += 1
|
|
|
|
elif mode == "Manual":
|
|
selected_checkpoint = checkpoint_files[manual_index % count]
|
|
|
|
# Construct the relative path for ComfyUI Checkpoint loader
|
|
# Use forward slashes for consistency with ComfyUI paths
|
|
full_checkpoint_name = f"{folder}/{selected_checkpoint}"
|
|
|
|
return (full_checkpoint_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}"
|