From 30a0829f33c38444045e75c65970c6cbdb9aa63b Mon Sep 17 00:00:00 2001 From: XchanBik Date: Sat, 10 May 2025 18:27:05 +0200 Subject: [PATCH] Update Utils.py --- Utils.py | 104 +++++++++++++++++++------------------------------------ 1 file changed, 36 insertions(+), 68 deletions(-) diff --git a/Utils.py b/Utils.py index 527f921..00ccf8f 100644 --- a/Utils.py +++ b/Utils.py @@ -1,76 +1,44 @@ import os -import time -import mimetypes import logging -from typing import Literal, List -from comfy.cli_args import args -from collections.abc import Collection +from comfy.cli_args import args # For your custom path logic -# --base-directory - Resets all default paths configured in folder_paths with a new base path if args.base_directory: - base_path = os.path.abspath(args.base_directory) + user_base_path = os.path.abspath(args.base_directory) else: - base_path = os.path.dirname(os.path.realpath(__file__)) -models_dir = os.path.join(base_path, "models") + user_base_path = os.path.dirname(os.path.realpath(__file__)) +user_models_dir = os.path.join(user_base_path, "models") -def recursive_search(directory: str) -> tuple[list[str], dict[str, float]]: - if not os.path.isdir(directory): - return [], {} - result = [] - dirs = {} - # Attempt to add the initial directory to dirs with error handling - try: - dirs[directory] = os.path.getmtime(directory) - except FileNotFoundError: - logging.warning(f"Warning: Unable to access {directory}. Skipping this path.") +# Log the determined models_dir for easier debugging by the +logging.info(f"[SimpleTextFileSelector] Initialized. Custom 'user_models_dir' is set to: {user_models_dir}") - logging.debug("recursive file list on directory {}".format(directory)) - dirpath: str - subdirs: list[str] - filenames: list[str] - - for dirpath, subdirs, filenames in os.walk(directory, followlinks=True, topdown=True): - subdirs[:] = [d for d in subdirs] - for file_name in filenames: - try: - relative_path = os.path.relpath(os.path.join(dirpath, file_name), directory) - result.append(relative_path) - except: - logging.warning(f"Warning: Unable to access {file_name}. Skipping this file.") - continue - - for d in subdirs: - path: str = os.path.join(dirpath, d) - try: - dirs[path] = os.path.getmtime(path) - except FileNotFoundError: - logging.warning(f"Warning: Unable to access {path}. Skipping this path.") - continue - logging.debug("found {} files".format(len(result))) - return result, dirs - -def filter_files_by_extension(files: Collection[str], extensions: Collection[str]) -> list[str]: - return sorted(list(filter(lambda a: os.path.splitext(a)[-1].lower() in extensions or len(extensions) == 0, files))) - -def get_filename_list(folder_name: str, extension: str) -> tuple[list[str], dict[str, float], float]: - full_path = os.path.join(models_dir, folder_name) - output_list = set - folders = full_path - output_folders = {} +# --- Utility Function to List .txt Files --- +def get_files_with_ext(folder_name: str, extension: str) -> list[str]: + """ + Scans a specific folder within the 'user_models_dir' for .{extension} files. + Returns a list of filenames (relative to that subfolder) for the dropdown. + """ + # user_models_dir is the global one defined above + full_scan_path = os.path.join(user_models_dir, folder_name) - # Let's manually handle the extension filtering inside this loop - files, folders_all = recursive_search(full_path) - - # Debug print to verify filenames - print(f"[get_filename_list] Files found: {files}") - - # Instead of using filter_files_extensions, we now manually filter files by extensions - extension_filtered_files = filter_files_by_extension(files, extension) # Assuming folders[1] contains the extensions - - # Update the output list with the filtered files - output_list.update(extension_filtered_files) - - # Merge folder info into output_folders - output_folders.update(folders_all) - - return sorted(list(output_list)), output_folders, time.perf_counter() + if not os.path.isdir(full_scan_path): + logging.warning(f"[SimpleTextFileSelector] Target subfolder for .txt files not found: {full_scan_path}. " + f"(This path is 'user_models_dir' / '{target_subfolder_name}')") + return ["None"] # ComfyUI dropdowns expect a list; "None" is a safe default. + + txt_files_found = [] + try: + for dirpath, _, filenames in os.walk(full_scan_path, followlinks=True): + for filename in filenames: + if filename.lower().endswith(extension): + relative_file_path = os.path.relpath(os.path.join(dirpath, filename), full_scan_path) + txt_files_found.append(relative_file_path) + except Exception as e: + logging.error(f"[SimpleTextFileSelector] Error occurred while scanning directory {full_scan_path}: {e}") + return ["None"] + + if not txt_files_found: + logging.info(f"[SimpleTextFileSelector] No .txt files were found in {full_scan_path}.") + return ["None"] + + logging.debug(f"[SimpleTextFileSelector] Found these .txt files in {full_scan_path}: {sorted(txt_files_found)}") + return sorted(txt_files_found)