diff --git a/model_filemanager/__init__.py b/model_filemanager/__init__.py index b7ac16256..f5d0c7bbc 100644 --- a/model_filemanager/__init__.py +++ b/model_filemanager/__init__.py @@ -1,2 +1,3 @@ # model_manager/__init__.py from .download_models import download_model, DownloadModelStatus, DownloadStatusType, create_model_path, check_file_exists, track_download_progress, validate_filename +from .model_filemanager import ModelFileManager \ No newline at end of file diff --git a/model_filemanager/model_filemanager.py b/model_filemanager/model_filemanager.py new file mode 100644 index 000000000..179de8a93 --- /dev/null +++ b/model_filemanager/model_filemanager.py @@ -0,0 +1,22 @@ +import folder_paths +from aiohttp import web + + +class ModelFileManager: + def __init__(self) -> None: + pass + + def add_routes(self, routes): + @routes.get("/models") + def list_model_types(request): + model_types = list(folder_paths.folder_names_and_paths.keys()) + + return web.json_response(model_types) + + @routes.get("/models/{folder}") + async def get_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 = folder_paths.get_filename_list(folder) + return web.json_response(files) diff --git a/server.py b/server.py index e663095bc..a87766df9 100644 --- a/server.py +++ b/server.py @@ -29,7 +29,7 @@ import comfy.model_management import node_helpers from app.frontend_management import FrontendManager from app.user_manager import UserManager -from model_filemanager import download_model, DownloadModelStatus +from model_filemanager import download_model, DownloadModelStatus, ModelFileManager from typing import Optional from api_server.routes.internal.internal_routes import InternalRoutes @@ -152,6 +152,7 @@ class PromptServer(): mimetypes.types_map['.js'] = 'application/javascript; charset=utf-8' self.user_manager = UserManager() + self.model_file_manager = ModelFileManager() self.internal_routes = InternalRoutes(self) self.supports = ["custom_nodes_from_web"] self.prompt_queue = None @@ -221,20 +222,6 @@ class PromptServer(): def get_embeddings(self): embeddings = folder_paths.get_filename_list("embeddings") return web.json_response(list(map(lambda a: os.path.splitext(a)[0], embeddings))) - - @routes.get("/models") - def list_model_types(request): - model_types = list(folder_paths.folder_names_and_paths.keys()) - - return web.json_response(model_types) - - @routes.get("/models/{folder}") - async def get_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 = folder_paths.get_filename_list(folder) - return web.json_response(files) @routes.get("/extensions") async def get_extensions(request): @@ -713,6 +700,7 @@ class PromptServer(): def add_routes(self): self.user_manager.add_routes(self.routes) + self.model_file_manager.add_routes(self.routes) self.app.add_subapp('/internal', self.internal_routes.get_app()) # Prefix every route with /api for easier matching for delegation.