dir_loop.py

# dir_loop.py is a GirDir node that provides a function to iterate through a directory and return the files in it.
from aiohttp import web
import os
import re
import server
import random

# Utility function for filtering files


def filter_files(directory, filter_type, filter_value, sort_by="name", sort_order="asc"):
    matched_files = []
    files = os.listdir(directory)
    if sort_by == "name":
        files.sort()
    elif sort_by == "date_modified":
        files.sort(key=lambda x: os.path.getmtime(os.path.join(directory, x)))
    elif sort_by == "date_created":
        files.sort(key=lambda x: os.path.getctime(os.path.join(directory, x)))
    if sort_order == "desc":
        files.reverse()
    if sort_order == "random":
        random.shuffle(files)

    for file in files:
        if filter_type == "regex" and re.match(filter_value, file):
            matched_files.append(file)
        elif filter_type == "extension" and file.endswith(filter_value):
            matched_files.append(file)
    return matched_files


loop_indexes = {}


class LoopyDir:
    # Increment this index each time a file is served, reset when reaching the end or on specific conditions
    file_index = 0
    matched_files = []

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "directory": ("STRING", {"forceInput": True, "default": "", "dynamicPrompts": False}),
                # Dropdown for regex or extension
                "filter_type": (["regex", "extension"], {"default": "extension"}),
                # Input for regex pattern or file extension
                "filter_value": ("STRING", {"default": "", "dynamicPrompts": False}),
                # Dropdown for sorting by name or date (modified/created)
                "sort_by": (["name", "date_modified", "date_created"], {"default": "name"}),
                # Dropdown for ascending or descending
                "sort_order": (["asc", "desc", "random"], {"default": "asc"}),
                # External loop index
                "loop_index": ("INT", {"default": 0}),
            },
            "hidden": {
                "prompt": "PROMPT",
                "id": "UNIQUE_ID",
            }
        }

    # Outputs the index and filename
    RETURN_TYPES = ("INT", "INT", "STRING", "STRING", "COMBO")
    RETURN_NAMES = ("file_count", "current_index",
                    "current_file", "current_file_path", "all_files")

    OUTPUT_NODE = True

    FUNCTION = "iterate_directory"
    CATEGORY = "Dir Gir"

    def iterate_directory(cls, directory, filter_type, filter_value, sort_by, sort_order, loop_index, prompt, id):
        # Load or refresh the list of matched files
        cls.matched_files = filter_files(directory, filter_type, filter_value, sort_by, sort_order)

        if len(cls.matched_files) == 0:
            # No files found, reset index
            loop_indexes[id] = 0
            print("[ComfyUI-DirGir] No files found in directory" + directory)
            return (0, 0, "", "", [])

        # Ensure loop_index is within bounds
        loop_index = loop_indexes.get(id, 0)

        if loop_index >= len(cls.matched_files):
            # If the external loop index is beyond the available files, reset to 0
            loop_indexes[id] = 0

        # Serve the file at the current loop index
        current_file = cls.matched_files[loop_index]

        # Prepare outputs
        output = (len(cls.matched_files), loop_index, current_file,
                  os.path.join(directory, current_file), cls.matched_files)

        # Increment the external loop index or reset if at the end
        loop_indexes[id] = (loop_index + 1) % len(cls.matched_files)

        return output


@server.PromptServer.instance.routes.get("/gir-dir/loop-index")
async def get_last_index(request):
    return web.json_response({'loop_index': loop_indexes.get(request.rel_url.query.get('id', '')) or 0})

@server.PromptServer.instance.routes.get("/gir-dir/set-loop-index")
async def set_last_index(request):
    loop_indexes[request.rel_url.query.get('id', '')] = int(request.rel_url.query.get('index', 0))
    return web.json_response({'success': True})