feat: add AODH Image Saver (Metadata), Lora Selector, Checkpoint Selector, and various node improvements

nodes/checkpoint_selector/checkpoint_selector.py

@@ -0,0 +1,120 @@
+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}"