From 389b3325d11b1992322be8c3dd6a035b314acb0b Mon Sep 17 00:00:00 2001 From: yt-koike Date: Fri, 30 Jan 2026 11:41:07 +0900 Subject: [PATCH 1/2] add model downloader endpoint --- app/model_manager.py | 44 +++++++++++++++++++++++++++++++++++++++++++- server.py | 2 +- 2 files changed, 44 insertions(+), 2 deletions(-) diff --git a/app/model_manager.py b/app/model_manager.py index f124d1117..a1c121d13 100644 --- a/app/model_manager.py +++ b/app/model_manager.py @@ -5,8 +5,10 @@ import base64 import json import time import logging +import requests import folder_paths import glob +from tqdm.auto import tqdm import comfy.utils from aiohttp import web from PIL import Image @@ -15,8 +17,9 @@ from folder_paths import map_legacy, filter_files_extensions, filter_files_conte class ModelFileManager: - def __init__(self) -> None: + def __init__(self, is_download_model_enabled: lambda: bool= lambda: False) -> None: self.cache: dict[str, tuple[list[dict], dict[str, float], float]] = {} + self.is_download_model_enabled = is_download_model_enabled def get_cache(self, key: str, default=None) -> tuple[list[dict], dict[str, float], float] | None: return self.cache.get(key, default) @@ -76,6 +79,45 @@ class ModelFileManager: except: return web.Response(status=404) + @routes.post("/download_model") + async def post_download_model(request): + if not self.is_download_model_enabled(): + print("Download Model endpoint is disabled") + return web.Response(status=403) + json_data = await request.json() + url = json_data.get("url", None) + if url is None: + print("URL is not provided") + return web.Response(status=401) + save_dir = json_data.get("save_dir", None) + if save_dir not in folder_paths.folder_names_and_paths: + print("Save directory is not valid") + return web.Response(status=401) + filename = json_data.get("filename", url.split("/")[-1]) + token = json_data.get("token", None) + + save_path = os.path.join(folder_paths.folder_names_and_paths[save_dir][0][0], filename) + tmp_path = save_path + ".tmp" + headers = {"Authorization": f"Bearer {token}"} if token else {} + try: + with requests.get(url, headers=headers,stream=True,timeout=10) as r: + r.raise_for_status() + total_size = int(r.headers.get('content-length', 0)) + with open(tmp_path, "wb") as f: + with tqdm(total=total_size, unit='iB', unit_scale=True, desc=filename) as pbar: + for chunk in r.iter_content(chunk_size=1024*1024): + if not chunk: + break + size = f.write(chunk) + pbar.update(size) + os.rename(tmp_path, save_path) + return web.Response(status=200) + except Exception as e: + print(f"Failed to download model: {e}") + if os.path.exists(tmp_path): + os.remove(tmp_path) + return web.Response(status=500) + def get_model_file_list(self, folder_name: str): folder_name = map_legacy(folder_name) folders = folder_paths.folder_names_and_paths[folder_name] diff --git a/server.py b/server.py index 2aee5cc06..bab4c0dda 100644 --- a/server.py +++ b/server.py @@ -201,7 +201,7 @@ class PromptServer(): mimetypes.add_type('image/webp', '.webp') self.user_manager = UserManager() - self.model_file_manager = ModelFileManager() + self.model_file_manager = ModelFileManager(is_download_model_enabled=lambda: self.user_manager.settings.get_settings(None).get("Comfy.ModelDownloadEnabled", False)) self.custom_node_manager = CustomNodeManager() self.subgraph_manager = SubgraphManager() self.internal_routes = InternalRoutes(self) From b65c1b15807824f2d817ee2eb2dbd1823c56ee2d Mon Sep 17 00:00:00 2001 From: yt-koike Date: Fri, 30 Jan 2026 12:14:28 +0900 Subject: [PATCH 2/2] replace print() to logging.error() --- app/model_manager.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/app/model_manager.py b/app/model_manager.py index a1c121d13..09a80f58b 100644 --- a/app/model_manager.py +++ b/app/model_manager.py @@ -82,16 +82,16 @@ class ModelFileManager: @routes.post("/download_model") async def post_download_model(request): if not self.is_download_model_enabled(): - print("Download Model endpoint is disabled") + logging.error("Download Model endpoint is disabled") return web.Response(status=403) json_data = await request.json() url = json_data.get("url", None) if url is None: - print("URL is not provided") + logging.error("URL is not provided") return web.Response(status=401) save_dir = json_data.get("save_dir", None) if save_dir not in folder_paths.folder_names_and_paths: - print("Save directory is not valid") + logging.error("Save directory is not valid") return web.Response(status=401) filename = json_data.get("filename", url.split("/")[-1]) token = json_data.get("token", None) @@ -113,7 +113,7 @@ class ModelFileManager: os.rename(tmp_path, save_path) return web.Response(status=200) except Exception as e: - print(f"Failed to download model: {e}") + logging.error(f"Failed to download model: {e}") if os.path.exists(tmp_path): os.remove(tmp_path) return web.Response(status=500)