diff --git a/app/model_manager.py b/app/model_manager.py index f124d1117..ff145c1a6 100644 --- a/app/model_manager.py +++ b/app/model_manager.py @@ -8,15 +8,33 @@ import logging import folder_paths import glob import comfy.utils +import uuid +from urllib.parse import urlparse +import aiohttp 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 +ALLOWED_MODEL_SOURCES = ( + "https://civitai.com/", + "https://huggingface.co/", + "http://localhost:", +) +ALLOWED_MODEL_SUFFIXES = (".safetensors", ".sft") +WHITELISTED_MODEL_URLS = { + "https://huggingface.co/stabilityai/stable-zero123/resolve/main/stable_zero123.ckpt", + "https://huggingface.co/TencentARC/T2I-Adapter/resolve/main/models/t2iadapter_depth_sd14v1.pth?download=true", + "https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth", +} +DOWNLOAD_CHUNK_SIZE = 1024 * 1024 +MAX_BULK_MODEL_DOWNLOADS = 200 + class ModelFileManager: - def __init__(self) -> None: + def __init__(self, prompt_server) -> None: self.cache: dict[str, tuple[list[dict], dict[str, float], float]] = {} + self.prompt_server = prompt_server def get_cache(self, key: str, default=None) -> tuple[list[dict], dict[str, float], float] | None: return self.cache.get(key, default) @@ -75,6 +93,144 @@ class ModelFileManager: return web.Response(body=img_bytes.getvalue(), content_type="image/webp") except: return web.Response(status=404) + + @routes.post("/experiment/models/download_missing") + async def download_missing_models(request: web.Request) -> web.Response: + print("download_missing_models") + try: + payload = await request.json() + except Exception: + return web.json_response({"error": "Invalid JSON body"}, status=400) + print("download_missing_models") + print(payload) + + models = payload.get("models") + if not isinstance(models, list): + return web.json_response({"error": "Field 'models' must be a list"}, status=400) + if len(models) > MAX_BULK_MODEL_DOWNLOADS: + return web.json_response( + {"error": f"Maximum of {MAX_BULK_MODEL_DOWNLOADS} models allowed per request"}, + status=400, + ) + + results = [] + downloaded = 0 + skipped = 0 + failed = 0 + + session = self.prompt_server.client_session + owns_session = False + if session is None or session.closed: + timeout = aiohttp.ClientTimeout(total=None) + session = aiohttp.ClientSession(timeout=timeout) + owns_session = True + + try: + for model_entry in models: + model_name, model_directory, model_url = _normalize_model_entry(model_entry) + + if not model_name or not model_directory or not model_url: + failed += 1 + results.append( + { + "name": model_name or "", + "directory": model_directory or "", + "url": model_url or "", + "status": "failed", + "error": "Each model must include non-empty name, directory, and url", + } + ) + continue + + if not _is_http_url(model_url): + failed += 1 + results.append( + { + "name": model_name, + "directory": model_directory, + "url": model_url, + "status": "failed", + "error": "URL must be an absolute HTTP/HTTPS URL", + } + ) + continue + + allowed, reason = _is_model_download_allowed(model_name, model_url) + if not allowed: + failed += 1 + results.append( + { + "name": model_name, + "directory": model_directory, + "url": model_url, + "status": "blocked", + "error": reason, + } + ) + continue + + try: + destination = _resolve_download_destination(model_directory, model_name) + except Exception as exc: + failed += 1 + results.append( + { + "name": model_name, + "directory": model_directory, + "url": model_url, + "status": "failed", + "error": str(exc), + } + ) + continue + + if os.path.exists(destination): + skipped += 1 + results.append( + { + "name": model_name, + "directory": model_directory, + "url": model_url, + "status": "skipped_existing", + } + ) + continue + + try: + await _download_file(session, model_url, destination) + downloaded += 1 + results.append( + { + "name": model_name, + "directory": model_directory, + "url": model_url, + "status": "downloaded", + } + ) + except Exception as exc: + failed += 1 + results.append( + { + "name": model_name, + "directory": model_directory, + "url": model_url, + "status": "failed", + "error": str(exc), + } + ) + finally: + if owns_session: + await session.close() + + return web.json_response( + { + "downloaded": downloaded, + "skipped": skipped, + "failed": failed, + "results": results, + }, + status=200, + ) def get_model_file_list(self, folder_name: str): folder_name = map_legacy(folder_name) @@ -193,3 +349,75 @@ class ModelFileManager: def __exit__(self, exc_type, exc_value, traceback): self.clear_cache() + +def _is_model_download_allowed(model_name: str, model_url: str) -> tuple[bool, str | None]: + if model_url in WHITELISTED_MODEL_URLS: + return True, None + + if not any(model_url.startswith(source) for source in ALLOWED_MODEL_SOURCES): + return ( + False, + f"Download not allowed from source '{model_url}'.", + ) + + if not any(model_name.endswith(suffix) for suffix in ALLOWED_MODEL_SUFFIXES): + return ( + False, + f"Only allowed suffixes are: {', '.join(ALLOWED_MODEL_SUFFIXES)}", + ) + + return True, None + + +def _resolve_download_destination(directory: str, model_name: str) -> str: + if directory not in folder_paths.folder_names_and_paths: + raise ValueError(f"Unknown model directory '{directory}'") + + model_paths = folder_paths.folder_names_and_paths[directory][0] + if not model_paths: + raise ValueError(f"No filesystem paths configured for '{directory}'") + + base_path = os.path.abspath(model_paths[0]) + normalized_name = os.path.normpath(model_name).lstrip("/\\") + if not normalized_name or normalized_name == ".": + raise ValueError("Model name cannot be empty") + + destination = os.path.abspath(os.path.join(base_path, normalized_name)) + if os.path.commonpath((base_path, destination)) != base_path: + raise ValueError("Model path escapes configured model directory") + + destination_parent = os.path.dirname(destination) + if destination_parent: + os.makedirs(destination_parent, exist_ok=True) + + return destination + + +async def _download_file(session: aiohttp.ClientSession, url: str, destination: str) -> None: + temp_file = f"{destination}.part-{uuid.uuid4().hex}" + try: + async with session.get(url, allow_redirects=True) as response: + response.raise_for_status() + with open(temp_file, "wb") as file_handle: + async for chunk in response.content.iter_chunked(DOWNLOAD_CHUNK_SIZE): + if chunk: + file_handle.write(chunk) + os.replace(temp_file, destination) + finally: + if os.path.exists(temp_file): + os.remove(temp_file) + + +def _normalize_model_entry(model_entry: object) -> tuple[str | None, str | None, str | None]: + if not isinstance(model_entry, dict): + return None, None, None + + model_name = str(model_entry.get("name", "")).strip() + model_directory = str(model_entry.get("directory", "")).strip() + model_url = str(model_entry.get("url", "")).strip() + return model_name, model_directory, model_url + + +def _is_http_url(url: str) -> bool: + parsed = urlparse(url) + return parsed.scheme in ("http", "https") and bool(parsed.netloc) \ No newline at end of file diff --git a/server.py b/server.py index 8882e43c4..c8f579548 100644 --- a/server.py +++ b/server.py @@ -202,7 +202,7 @@ class PromptServer(): mimetypes.add_type('image/webp', '.webp') self.user_manager = UserManager() - self.model_file_manager = ModelFileManager() + self.model_file_manager = ModelFileManager(self) self.custom_node_manager = CustomNodeManager() self.subgraph_manager = SubgraphManager() self.node_replace_manager = NodeReplaceManager() diff --git a/tests-unit/app_test/model_manager_test.py b/tests-unit/app_test/model_manager_test.py index ae59206f6..12a559fa7 100644 --- a/tests-unit/app_test/model_manager_test.py +++ b/tests-unit/app_test/model_manager_test.py @@ -12,9 +12,14 @@ pytestmark = ( pytest.mark.asyncio ) # This applies the asyncio mark to all test functions in the module +class DummyPromptServer: + def __init__(self): + self.client_session = None + @pytest.fixture def model_manager(): - return ModelFileManager() + prompt_server = DummyPromptServer() + return ModelFileManager(prompt_server) @pytest.fixture def app(model_manager):