From 7bba21af47bcf2510be7811b38378dd1bb9c8804 Mon Sep 17 00:00:00 2001 From: hayden <48267247+hayden-fr@users.noreply.github.com> Date: Tue, 21 Jan 2025 15:33:11 +0800 Subject: [PATCH 1/6] Add output manager --- app/output_manager.py | 126 ++++++++++++++++++++++++++++++++++++++++++ server.py | 3 + 2 files changed, 129 insertions(+) create mode 100644 app/output_manager.py diff --git a/app/output_manager.py b/app/output_manager.py new file mode 100644 index 000000000..14b5a0034 --- /dev/null +++ b/app/output_manager.py @@ -0,0 +1,126 @@ +from __future__ import annotations + +import os +import folder_paths +import mimetypes +import shutil +from aiohttp import web +from io import BytesIO +from PIL import Image +from typing import Literal + + +class OutputManager: + def __init__(self) -> None: + self.output_uri = folder_paths.get_output_directory() + + def add_routes(self, routes) -> None: + @routes.get("/output{pathname:.*}") + async def get_output_file_or_files(request): + try: + pathname = request.match_info.get("pathname", None) + filepath = self.get_output_filepath(pathname) + + if os.path.isfile(filepath): + + preview_type = request.query.get("preview_type", None) + if not preview_type: + return web.FileResponse(filepath) + + # get image preview + if self.assert_file_type(filepath, ["image"]): + image_data = self.get_image_preview_data(filepath) + return web.Response(body=image_data.getvalue(), content_type="image/webp") + + # TODO get video cover preview + + elif os.path.isdir(filepath): + files = self.get_items_by_directory(filepath) + return web.json_response(files) + + return web.Response(status=404) + except Exception as e: + return web.Response(status=500) + + @routes.delete("/output{pathname:.*}") + async def delete_output_file_or_files(request): + try: + pathname = request.match_info.get("pathname", None) + filepath = self.get_output_filepath(pathname) + + if os.path.isfile(filepath): + os.remove(filepath) + elif os.path.isdir(filepath): + shutil.rmtree(filepath) + return web.Response(status=200) + except Exception as e: + return web.Response(status=500) + + def get_output_filepath(self, pathname: str): + return f"{self.output_uri}/{pathname}" + + def get_items_by_directory(self, pathname: str): + result = [] + items = os.listdir(pathname) + + for name in items: + filepath = os.path.join(pathname, name) + + if os.path.isfile(filepath) and not self.assert_file_type(filepath, ["image", "video", "audio"]): + continue + + state = os.stat(filepath) + is_dir = os.path.isdir(filepath) + + result.append( + { + "name": name, + "type": "folder" if is_dir else self.get_file_content_type(filepath), + "size": 0 if is_dir else state.st_size, + "createdAt": round(state.st_ctime_ns / 1000000), + "updatedAt": round(state.st_mtime_ns / 1000000), + } + ) + + return result + + def assert_file_type(self, filename: str, content_types: Literal["image", "video", "audio"]): + content_type = self.get_file_content_type(filename) + if not content_type: + return False + return content_type in content_types + + def get_file_content_type(self, filename: str): + extension_mimetypes_cache = folder_paths.extension_mimetypes_cache + + extension = filename.split(".")[-1] + content_type = None + if extension not in extension_mimetypes_cache: + mime_type, _ = mimetypes.guess_type(filename, strict=False) + if mime_type: + content_type = mime_type.split("/")[0] + extension_mimetypes_cache[extension] = content_type + else: + content_type = extension_mimetypes_cache[extension] + + return content_type + + def get_image_preview_data(self, filename: str): + with Image.open(filename) as img: + max_size = 128 + + old_width, old_height = img.size + scale = min(max_size / old_width, max_size / old_height) + + if scale >= 1: + new_width, new_height = old_width, old_height + else: + new_width = int(old_width * scale) + new_height = int(old_height * scale) + + img = img.resize((new_width, new_height), Image.Resampling.LANCZOS) + + img_byte_arr = BytesIO() + img.save(img_byte_arr, format="WEBP") + img_byte_arr.seek(0) + return img_byte_arr diff --git a/server.py b/server.py index 88c163fc7..0b6d5010b 100644 --- a/server.py +++ b/server.py @@ -31,6 +31,7 @@ from comfyui_version import __version__ from app.frontend_management import FrontendManager from app.user_manager import UserManager from app.model_manager import ModelFileManager +from app.output_manager import OutputManager from app.custom_node_manager import CustomNodeManager from typing import Optional from api_server.routes.internal.internal_routes import InternalRoutes @@ -140,6 +141,7 @@ class PromptServer(): self.user_manager = UserManager() self.model_file_manager = ModelFileManager() + self.output_manager = OutputManager() self.custom_node_manager = CustomNodeManager() self.internal_routes = InternalRoutes(self) self.supports = ["custom_nodes_from_web"] @@ -691,6 +693,7 @@ class PromptServer(): def add_routes(self): self.user_manager.add_routes(self.routes) self.model_file_manager.add_routes(self.routes) + self.output_manager.add_routes(self.routes) self.custom_node_manager.add_routes(self.routes, self.app, nodes.LOADED_MODULE_DIRS.items()) self.app.add_subapp('/internal', self.internal_routes.get_app()) From 0e8640519895fd8458c2708009a8157b256962cf Mon Sep 17 00:00:00 2001 From: hayden <48267247+hayden-fr@users.noreply.github.com> Date: Wed, 22 Jan 2025 15:58:41 +0800 Subject: [PATCH 2/6] Optimize file scanning using scandir --- app/output_manager.py | 37 ++++++++++++++++++------------------- 1 file changed, 18 insertions(+), 19 deletions(-) diff --git a/app/output_manager.py b/app/output_manager.py index 14b5a0034..3938f9565 100644 --- a/app/output_manager.py +++ b/app/output_manager.py @@ -35,7 +35,7 @@ class OutputManager: # TODO get video cover preview elif os.path.isdir(filepath): - files = self.get_items_by_directory(filepath) + files = self.get_folder_items(filepath) return web.json_response(files) return web.Response(status=404) @@ -59,28 +59,27 @@ class OutputManager: def get_output_filepath(self, pathname: str): return f"{self.output_uri}/{pathname}" - def get_items_by_directory(self, pathname: str): + def get_folder_items(self, folder: str): result = [] - items = os.listdir(pathname) - for name in items: - filepath = os.path.join(pathname, name) + with os.scandir(folder) as it: + for entry in it: + filepath = entry.path + is_dir = entry.is_dir() - if os.path.isfile(filepath) and not self.assert_file_type(filepath, ["image", "video", "audio"]): - continue + if not is_dir and not self.assert_file_type(filepath, ["image", "video", "audio"]): + continue - state = os.stat(filepath) - is_dir = os.path.isdir(filepath) - - result.append( - { - "name": name, - "type": "folder" if is_dir else self.get_file_content_type(filepath), - "size": 0 if is_dir else state.st_size, - "createdAt": round(state.st_ctime_ns / 1000000), - "updatedAt": round(state.st_mtime_ns / 1000000), - } - ) + state = entry.stat() + result.append( + { + "name": entry.name, + "type": "folder" if entry.is_dir() else self.get_file_content_type(filepath), + "size": 0 if is_dir else state.st_size, + "createdAt": round(state.st_ctime_ns / 1000000), + "updatedAt": round(state.st_mtime_ns / 1000000), + } + ) return result From 82c3afe07792e89900e2c74bad17c80f98dacb0f Mon Sep 17 00:00:00 2001 From: hayden <48267247+hayden-fr@users.noreply.github.com> Date: Thu, 23 Jan 2025 12:55:08 +0800 Subject: [PATCH 3/6] Optimize file scanning performance using multi-threading --- app/output_manager.py | 38 ++++++++++++++++++++++---------------- 1 file changed, 22 insertions(+), 16 deletions(-) diff --git a/app/output_manager.py b/app/output_manager.py index 3938f9565..63fe99aae 100644 --- a/app/output_manager.py +++ b/app/output_manager.py @@ -5,6 +5,7 @@ import folder_paths import mimetypes import shutil from aiohttp import web +from concurrent.futures import ThreadPoolExecutor, as_completed from io import BytesIO from PIL import Image from typing import Literal @@ -62,24 +63,29 @@ class OutputManager: def get_folder_items(self, folder: str): result = [] - with os.scandir(folder) as it: - for entry in it: - filepath = entry.path - is_dir = entry.is_dir() + def get_file_info(entry: os.DirEntry[str]): + filepath = entry.path + is_dir = entry.is_dir() - if not is_dir and not self.assert_file_type(filepath, ["image", "video", "audio"]): + if not is_dir and not self.assert_file_type(filepath, ["image", "video", "audio"]): + return None + + stat = entry.stat() + return { + "name": entry.name, + "type": "folder" if entry.is_dir() else self.get_file_content_type(filepath), + "size": 0 if is_dir else stat.st_size, + "createdAt": round(stat.st_ctime_ns / 1000000), + "updatedAt": round(stat.st_mtime_ns / 1000000), + } + + with os.scandir(folder) as it, ThreadPoolExecutor() as executor: + future_to_entry = {executor.submit(get_file_info, entry): entry for entry in it} + for future in as_completed(future_to_entry): + file_info = future.result() + if file_info is None: continue - - state = entry.stat() - result.append( - { - "name": entry.name, - "type": "folder" if entry.is_dir() else self.get_file_content_type(filepath), - "size": 0 if is_dir else state.st_size, - "createdAt": round(state.st_ctime_ns / 1000000), - "updatedAt": round(state.st_mtime_ns / 1000000), - } - ) + result.append(file_info) return result From d8eae1b24144e100a56787ec37455a40a291777c Mon Sep 17 00:00:00 2001 From: hayden <48267247+hayden-fr@users.noreply.github.com> Date: Thu, 23 Jan 2025 13:53:41 +0800 Subject: [PATCH 4/6] Add folder content caching --- app/output_manager.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/app/output_manager.py b/app/output_manager.py index 63fe99aae..749a82df6 100644 --- a/app/output_manager.py +++ b/app/output_manager.py @@ -13,8 +13,19 @@ from typing import Literal class OutputManager: def __init__(self) -> None: + self.cache: dict[str, tuple[list, float]] = {} self.output_uri = folder_paths.get_output_directory() + def get_cache(self, key: str): + return self.cache.get(key, ([], 0)) + + def set_cache(self, key: str, value: tuple[list, float]): + self.cache[key] = value + + def rm_cache(self, key: str): + if key in self.cache: + del self.cache[key] + def add_routes(self, routes) -> None: @routes.get("/output{pathname:.*}") async def get_output_file_or_files(request): @@ -53,6 +64,7 @@ class OutputManager: os.remove(filepath) elif os.path.isdir(filepath): shutil.rmtree(filepath) + self.rm_cache(filepath) return web.Response(status=200) except Exception as e: return web.Response(status=500) @@ -61,6 +73,12 @@ class OutputManager: return f"{self.output_uri}/{pathname}" def get_folder_items(self, folder: str): + result, m_time = self.get_cache(folder) + folder_m_time = os.path.getmtime(folder) + + if folder_m_time == m_time: + return result + result = [] def get_file_info(entry: os.DirEntry[str]): @@ -87,6 +105,7 @@ class OutputManager: continue result.append(file_info) + self.set_cache(folder, (result, os.path.getmtime(folder))) return result def assert_file_type(self, filename: str, content_types: Literal["image", "video", "audio"]): From 2cf95ed231d661f0e50d67c2947c7fec9373e38f Mon Sep 17 00:00:00 2001 From: hayden <48267247+hayden-fr@users.noreply.github.com> Date: Fri, 24 Jan 2025 14:37:34 +0800 Subject: [PATCH 5/6] Change property name --- app/output_manager.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/app/output_manager.py b/app/output_manager.py index 749a82df6..999adf9fb 100644 --- a/app/output_manager.py +++ b/app/output_manager.py @@ -93,8 +93,8 @@ class OutputManager: "name": entry.name, "type": "folder" if entry.is_dir() else self.get_file_content_type(filepath), "size": 0 if is_dir else stat.st_size, - "createdAt": round(stat.st_ctime_ns / 1000000), - "updatedAt": round(stat.st_mtime_ns / 1000000), + "createTime": round(stat.st_ctime_ns / 1000000), + "modifyTime": round(stat.st_mtime_ns / 1000000), } with os.scandir(folder) as it, ThreadPoolExecutor() as executor: From 9c957977d09abf730d9aaecc5420b15941496903 Mon Sep 17 00:00:00 2001 From: hayden <48267247+hayden-fr@users.noreply.github.com> Date: Fri, 24 Jan 2025 16:58:40 +0800 Subject: [PATCH 6/6] Repair the file operation abnormality and add a log record --- app/output_manager.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/app/output_manager.py b/app/output_manager.py index 999adf9fb..f60229ed3 100644 --- a/app/output_manager.py +++ b/app/output_manager.py @@ -1,9 +1,11 @@ from __future__ import annotations import os +import logging import folder_paths import mimetypes import shutil +import traceback from aiohttp import web from concurrent.futures import ThreadPoolExecutor, as_completed from io import BytesIO @@ -29,8 +31,8 @@ class OutputManager: def add_routes(self, routes) -> None: @routes.get("/output{pathname:.*}") async def get_output_file_or_files(request): + pathname = request.match_info.get("pathname", None) try: - pathname = request.match_info.get("pathname", None) filepath = self.get_output_filepath(pathname) if os.path.isfile(filepath): @@ -51,13 +53,15 @@ class OutputManager: return web.json_response(files) return web.Response(status=404) - except Exception as e: + except Exception: + logging.error(f"File '{pathname}' retrieval failed") + logging.error(traceback.format_exc()) return web.Response(status=500) @routes.delete("/output{pathname:.*}") async def delete_output_file_or_files(request): + pathname = request.match_info.get("pathname", None) try: - pathname = request.match_info.get("pathname", None) filepath = self.get_output_filepath(pathname) if os.path.isfile(filepath): @@ -66,7 +70,9 @@ class OutputManager: shutil.rmtree(filepath) self.rm_cache(filepath) return web.Response(status=200) - except Exception as e: + except Exception: + logging.error(f"File '{pathname}' deletion failed") + logging.error(traceback.format_exc()) return web.Response(status=500) def get_output_filepath(self, pathname: str):