diff --git a/model_filemanager/model_filemanager.py b/model_filemanager/model_filemanager.py index 179de8a93..4784cc267 100644 --- a/model_filemanager/model_filemanager.py +++ b/model_filemanager/model_filemanager.py @@ -1,10 +1,23 @@ +import os +import time +import logging import folder_paths from aiohttp import web +from folder_paths import map_legacy, filter_files_extensions class ModelFileManager: def __init__(self) -> None: - pass + self.cache: dict[str, tuple[list[dict], float, float]] = {} + + def get_cache(self, key: str, default=None): + return self.cache.get(key, default) + + def set_cache(self, key: str, value: tuple[list[dict], float, float]): + self.cache[key] = value + + def clear_cache(self): + self.cache.clear() def add_routes(self, routes): @routes.get("/models") @@ -20,3 +33,80 @@ class ModelFileManager: return web.Response(status=404) files = folder_paths.get_filename_list(folder) return web.json_response(files) + + # NOTE: This is an experiment to replace `/models` + @routes.get("/experiment/models") + async def get_model_folders(request): + model_types = list(folder_paths.folder_names_and_paths.keys()) + folder_black_list = ["configs", "custom_nodes"] + output_folders: list[dict] = [] + for folder in model_types: + if folder in folder_black_list: + continue + output_folders.append({"name": folder, "folders": folder_paths.get_folder_paths(folder)}) + return web.json_response(output_folders) + + # NOTE: This is an experiment to replace `/models/{folder}` + @routes.get("/experiment/models/{folder}") + async def get_all_models(request): + folder = request.match_info.get("folder", None) + if not folder in folder_paths.folder_names_and_paths: + return web.Response(status=404) + files = self.get_model_file_list(folder) + return web.json_response(files) + + def get_model_file_list(self, folder_name: str): + folder_name = map_legacy(folder_name) + folders = folder_paths.folder_names_and_paths[folder_name] + output_list: list[dict] = [] + + for index, folder in enumerate(folders[0]): + out = self.cache_model_file_list_(folder) + if out is None: + out = self.recursive_search_models_(folder, index) + self.set_cache(folder, out) + output_list.extend(out[0]) + + return output_list + + def cache_model_file_list_(self, folder: str): + model_file_list_cache = self.get_cache(folder) + + if model_file_list_cache is not None: + if not os.path.isdir(folder): + return None + if os.path.getmtime(folder) != model_file_list_cache[1]: + return None + return model_file_list_cache + return None + + def recursive_search_models_(self, directory: str, pathIndex: int) -> tuple[list[str], float, float]: + if not os.path.isdir(directory): + return [], 0, time.perf_counter() + + excluded_dir_names = [".git"] + # TODO use settings + include_hidden_files = False + + result: list[str] = [] + + for dirpath, subdirs, filenames in os.walk(directory, followlinks=True, topdown=True): + subdirs[:] = [d for d in subdirs if d not in excluded_dir_names] + if not include_hidden_files: + subdirs[:] = [d for d in subdirs if not d.startswith(".")] + filenames = [f for f in filenames if not f.startswith(".")] + + filenames = filter_files_extensions(filenames, folder_paths.supported_pt_extensions) + + 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 + + return [{"name": f, "pathIndex": pathIndex} for f in result], os.path.getmtime(directory), time.perf_counter() + + def __exit__(self, exc_type, exc_value, traceback): + self.clear_cache()