From 2c8fd9f58909f87e4b4482a88b43501dfae59564 Mon Sep 17 00:00:00 2001
From: cerspense <cerspense@gmail.com>
Date: Mon, 22 Jul 2024 21:00:35 -0700
Subject: [PATCH] added sorting modes to image dir iterator

---
 cspnodes.py | 32 ++++++++++++++++++++++----------
 1 file changed, 22 insertions(+), 10 deletions(-)

diff --git a/cspnodes.py b/cspnodes.py
index c665d4c..be88b8c 100644
--- a/cspnodes.py
+++ b/cspnodes.py
@@ -4,6 +4,7 @@ from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
 from diffusers.utils import export_to_video
 from PIL import Image, ImageOps
 import numpy as np
+import random
 
 
 class TextFileLineIterator:
@@ -40,23 +41,34 @@ class ImageDirIterator:
         return {
             "required": {
                 "directory_path": ("STRING", {}),
-                "image_index": ("INT", {"default": 0})
+                "image_index": ("INT", {"default": 0}),
+                "sort_by": (["date_modified", "name", "size", "random"],),
+                "sort_order": (["ascending", "descending"],),
             }
         }
 
     RETURN_TYPES = ("IMAGE", "STRING")
-
     FUNCTION = "get_image_by_index"
-
     CATEGORY = "cspnodes"
 
-    def get_image_by_index(self, directory_path, image_index):
-        # Get list of image files sorted by modification time (most recent first)
-        image_files = sorted(
-            [os.path.join(directory_path, f) for f in os.listdir(directory_path)
-             if f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif'))],
-            key=lambda x: os.path.getmtime(x), reverse=True
-        )
+    def get_image_by_index(self, directory_path, image_index, sort_by, sort_order):
+        # Get list of image files
+        image_files = [os.path.join(directory_path, f) for f in os.listdir(directory_path)
+                       if f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif'))]
+
+        # Define sorting key functions
+        sort_functions = {
+            "date_modified": lambda x: os.path.getmtime(x),
+            "name": lambda x: int(''.join(filter(str.isdigit, os.path.basename(x))) or 0),
+            "size": lambda x: os.path.getsize(x),
+            "random": lambda x: random.random(),
+        }
+
+        # Sort the image files
+        if sort_by == "random":
+            random.shuffle(image_files)
+        else:
+            image_files.sort(key=sort_functions[sort_by], reverse=(sort_order == "descending"))
 
         # Wrap the index around using modulo
         image_index = image_index % len(image_files)