Merge 9c957977d0 into 22ad513c72

app/output_manager.py

@@ -0,0 +1,156 @@
+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
+from PIL import Image
+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):
+            pathname = request.match_info.get("pathname", None)
+            try:
+                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_folder_items(filepath)
+                    return web.json_response(files)
+
+                return web.Response(status=404)
+            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:
+                filepath = self.get_output_filepath(pathname)
+
+                if os.path.isfile(filepath):
+                    os.remove(filepath)
+                elif os.path.isdir(filepath):
+                    shutil.rmtree(filepath)
+                    self.rm_cache(filepath)
+                return web.Response(status=200)
+            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):
+        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]):
+            filepath = entry.path
+            is_dir = entry.is_dir()
+
+            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,
+                "createTime": round(stat.st_ctime_ns / 1000000),
+                "modifyTime": 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
+                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"]):
+        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

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
@@ -155,6 +156,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"]
@@ -715,6 +717,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())