88 lines
3.0 KiB
Python
88 lines
3.0 KiB
Python
# 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
|
|
|
|
# Utility function for filtering files
|
|
|
|
|
|
def filter_files(directory, filter_type, filter_value):
|
|
matched_files = []
|
|
for file in os.listdir(directory):
|
|
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}),
|
|
"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, loop_index, prompt, id):
|
|
# Load or refresh the list of matched files
|
|
cls.matched_files = filter_files(directory, filter_type, filter_value)
|
|
|
|
if len(cls.matched_files) == 0:
|
|
# No files found, reset index
|
|
loop_indexes[id] = 0
|
|
print("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})
|