diff --git a/py/manager.py b/py/manager.py index 0d52a0b..7a3efcc 100644 --- a/py/manager.py +++ b/py/manager.py @@ -1,6 +1,7 @@ import os import folder_paths from aiohttp import web +from concurrent.futures import ThreadPoolExecutor, as_completed from . import utils @@ -115,43 +116,57 @@ class ModelManager: def scan_models(self, folder: str, request): result = [] + include_hidden_files = utils.get_setting_value(request, "scan.include_hidden_files", False) folders, extensions = folder_paths.folder_names_and_paths[folder] + + def get_file_info(entry: os.DirEntry[str], base_path: str, path_index: int): + fullname = entry.path.replace(base_path, "") + basename = os.path.splitext(fullname)[0] + extension = os.path.splitext(fullname)[1] + + if extension not in folder_paths.supported_pt_extensions: + return None + + model_preview = f"/model-manager/preview/{folder}/{path_index}/{basename}.webp" + + stat = entry.stat() + return { + "fullname": fullname, + "basename": basename, + "extension": extension, + "type": folder, + "pathIndex": path_index, + "sizeBytes": stat.st_size, + "preview": model_preview, + "createdAt": round(stat.st_ctime_ns / 1000000), + "updatedAt": round(stat.st_mtime_ns / 1000000), + } + + def get_all_files_entry(directory: str): + files = [] + with os.scandir(directory) as it: + for entry in it: + # Skip hidden files + if not include_hidden_files: + if entry.name.startswith("."): + continue + if entry.is_dir(): + files.extend(get_all_files_entry(entry.path)) + elif entry.is_file(): + files.append(entry) + return files + for path_index, base_path in enumerate(folders): - files = utils.recursive_search_files(base_path, request) - - models = folder_paths.filter_files_extensions(files, folder_paths.supported_pt_extensions) - - for fullname in models: - fullname = utils.normalize_path(fullname) - basename = os.path.splitext(fullname)[0] - extension = os.path.splitext(fullname)[1] - - abs_path = utils.join_path(base_path, fullname) - file_stats = os.stat(abs_path) - - # Resolve preview - image_name = utils.get_model_preview_name(abs_path) - image_name = utils.join_path(os.path.dirname(fullname), image_name) - abs_image_path = utils.join_path(base_path, image_name) - if os.path.isfile(abs_image_path): - image_state = os.stat(abs_image_path) - image_timestamp = round(image_state.st_mtime_ns / 1000000) - image_name = f"{image_name}?ts={image_timestamp}" - model_preview = f"/model-manager/preview/{folder}/{path_index}/{image_name}" - - model_info = { - "fullname": fullname, - "basename": basename, - "extension": extension, - "type": folder, - "pathIndex": path_index, - "sizeBytes": file_stats.st_size, - "preview": model_preview, - "createdAt": round(file_stats.st_ctime_ns / 1000000), - "updatedAt": round(file_stats.st_mtime_ns / 1000000), - } - - result.append(model_info) + if not os.path.exists(base_path): + continue + file_entries = get_all_files_entry(base_path) + with ThreadPoolExecutor() as executor: + futures = {executor.submit(get_file_info, entry, base_path, path_index): entry for entry in file_entries} + for future in as_completed(futures): + file_info = future.result() + if file_info is None: + continue + result.append(file_info) return result