| from __future__ import annotations
|
|
|
| import os
|
| import base64
|
| import json
|
| import time
|
| import logging
|
| import folder_paths
|
| import glob
|
| import comfy.utils
|
| from aiohttp import web
|
| from PIL import Image
|
| from io import BytesIO
|
| from folder_paths import map_legacy, filter_files_extensions, filter_files_content_types
|
|
|
|
|
| class ModelFileManager:
|
| def __init__(self) -> None:
|
| self.cache: dict[str, tuple[list[dict], dict[str, float], float]] = {}
|
|
|
| def get_cache(self, key: str, default=None) -> tuple[list[dict], dict[str, float], float] | None:
|
| return self.cache.get(key, default)
|
|
|
| def set_cache(self, key: str, value: tuple[list[dict], dict[str, float], float]):
|
| self.cache[key] = value
|
|
|
| def clear_cache(self):
|
| self.cache.clear()
|
|
|
| def add_routes(self, routes):
|
|
|
| @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)
|
|
|
|
|
| @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)
|
|
|
| @routes.get("/experiment/models/preview/{folder}/{path_index}/{filename:.*}")
|
| async def get_model_preview(request):
|
| folder_name = request.match_info.get("folder", None)
|
| path_index = int(request.match_info.get("path_index", None))
|
| filename = request.match_info.get("filename", None)
|
|
|
| if not folder_name in folder_paths.folder_names_and_paths:
|
| return web.Response(status=404)
|
|
|
| folders = folder_paths.folder_names_and_paths[folder_name]
|
| folder = folders[0][path_index]
|
| full_filename = os.path.join(folder, filename)
|
|
|
| previews = self.get_model_previews(full_filename)
|
| default_preview = previews[0] if len(previews) > 0 else None
|
| if default_preview is None or (isinstance(default_preview, str) and not os.path.isfile(default_preview)):
|
| return web.Response(status=404)
|
|
|
| try:
|
| with Image.open(default_preview) as img:
|
| img_bytes = BytesIO()
|
| img.save(img_bytes, format="WEBP")
|
| img_bytes.seek(0)
|
| return web.Response(body=img_bytes.getvalue(), content_type="image/webp")
|
| except:
|
| return web.Response(status=404)
|
|
|
| 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]):
|
| if not os.path.isdir(folder):
|
| continue
|
| 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 None:
|
| return None
|
| if not os.path.isdir(folder):
|
| return None
|
| if os.path.getmtime(folder) != model_file_list_cache[1]:
|
| return None
|
| for x in model_file_list_cache[1]:
|
| time_modified = model_file_list_cache[1][x]
|
| folder = x
|
| if os.path.getmtime(folder) != time_modified:
|
| return None
|
|
|
| return model_file_list_cache
|
|
|
| def recursive_search_models_(self, directory: str, pathIndex: int) -> tuple[list[str], dict[str, float], float]:
|
| if not os.path.isdir(directory):
|
| return [], {}, time.perf_counter()
|
|
|
| excluded_dir_names = [".git"]
|
|
|
| include_hidden_files = False
|
|
|
| result: list[str] = []
|
| dirs: dict[str, float] = {}
|
|
|
| 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
|
|
|
| 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
|
|
|
| return [{"name": f, "pathIndex": pathIndex} for f in result], dirs, time.perf_counter()
|
|
|
| def get_model_previews(self, filepath: str) -> list[str | BytesIO]:
|
| dirname = os.path.dirname(filepath)
|
|
|
| if not os.path.exists(dirname):
|
| return []
|
|
|
| basename = os.path.splitext(filepath)[0]
|
| match_files = glob.glob(f"{basename}.*", recursive=False)
|
| image_files = filter_files_content_types(match_files, "image")
|
| safetensors_file = next(filter(lambda x: x.endswith(".safetensors"), match_files), None)
|
| safetensors_metadata = {}
|
|
|
| result: list[str | BytesIO] = []
|
|
|
| for filename in image_files:
|
| _basename = os.path.splitext(filename)[0]
|
| if _basename == basename:
|
| result.append(filename)
|
| if _basename == f"{basename}.preview":
|
| result.append(filename)
|
|
|
| if safetensors_file:
|
| safetensors_filepath = os.path.join(dirname, safetensors_file)
|
| header = comfy.utils.safetensors_header(safetensors_filepath, max_size=8*1024*1024)
|
| if header:
|
| safetensors_metadata = json.loads(header)
|
| safetensors_images = safetensors_metadata.get("__metadata__", {}).get("ssmd_cover_images", None)
|
| if safetensors_images:
|
| safetensors_images = json.loads(safetensors_images)
|
| for image in safetensors_images:
|
| result.append(BytesIO(base64.b64decode(image)))
|
|
|
| return result
|
|
|
| def __exit__(self, exc_type, exc_value, traceback):
|
| self.clear_cache()
|
|
|
