mirror of
https://github.com/Comfy-Org/ComfyUI.git
synced 2026-08-24 18:10:27 +08:00
feat: add model download API gated behind --enable-download-api
Add a new server-side download API that allows frontends and desktop apps
to download models directly into ComfyUI's models directory, eliminating
the need for DOM scraping of the frontend UI.
New files:
- app/download_manager.py: Async download manager with streaming downloads,
pause/resume/cancel, manual redirect following with per-hop host validation,
sidecar metadata for safe resume, and concurrency limiting.
API endpoints (all under /download/, also mirrored at /api/download/):
- POST /download/model - Start a download (url, directory, filename)
- GET /download/status - List all downloads (filterable by client_id)
- GET /download/status/{id} - Get single download status
- POST /download/pause/{id} - Pause (cancels transfer, keeps temp)
- POST /download/resume/{id} - Resume (new request with Range header)
- POST /download/cancel/{id} - Cancel and clean up temp files
Security:
- Gated behind --enable-download-api CLI flag (403 if disabled)
- HTTPS-only with exact host allowlist (huggingface.co, civitai.com + CDNs)
- Manual redirect following with per-hop host validation (no SSRF)
- Path traversal protection via realpath + commonpath
- Extension allowlist (.safetensors, .sft)
- Filename sanitization (no separators, .., control chars)
- Destination re-checked before final rename
- Progress events scoped to initiating client_id
Closes Comfy-Org/ComfyUI-Desktop-2.0-Beta#293
Amp-Thread-ID: https://ampcode.com/threads/T-019d2344-139e-77a5-9f24-1cbb3b26a8ec
Co-authored-by: Amp <amp@ampcode.com>
This commit is contained in:
84
server.py
84
server.py
@@ -43,6 +43,7 @@ from app.model_manager import ModelFileManager
|
||||
from app.custom_node_manager import CustomNodeManager
|
||||
from app.subgraph_manager import SubgraphManager
|
||||
from app.node_replace_manager import NodeReplaceManager
|
||||
from app.download_manager import DownloadManager
|
||||
from typing import Optional, Union
|
||||
from api_server.routes.internal.internal_routes import InternalRoutes
|
||||
from protocol import BinaryEventTypes
|
||||
@@ -205,6 +206,7 @@ class PromptServer():
|
||||
self.subgraph_manager = SubgraphManager()
|
||||
self.node_replace_manager = NodeReplaceManager()
|
||||
self.internal_routes = InternalRoutes(self)
|
||||
self.download_manager = DownloadManager(self) if args.enable_download_api else None
|
||||
self.supports = ["custom_nodes_from_web"]
|
||||
self.prompt_queue = execution.PromptQueue(self)
|
||||
self.loop = loop
|
||||
@@ -1028,9 +1030,91 @@ class PromptServer():
|
||||
|
||||
return web.Response(status=200)
|
||||
|
||||
# -- Download API (gated behind --enable-download-api) --
|
||||
|
||||
def _require_download_api(handler):
|
||||
async def wrapper(request):
|
||||
if self.download_manager is None:
|
||||
return web.json_response(
|
||||
{"error": "Download API is not enabled. Start ComfyUI with --enable-download-api."},
|
||||
status=403,
|
||||
)
|
||||
return await handler(request)
|
||||
return wrapper
|
||||
|
||||
@routes.post("/download/model")
|
||||
@_require_download_api
|
||||
async def post_download_model(request):
|
||||
json_data = await request.json()
|
||||
url = json_data.get("url")
|
||||
directory = json_data.get("directory")
|
||||
filename = json_data.get("filename")
|
||||
client_id = json_data.get("client_id")
|
||||
|
||||
if not url or not directory or not filename:
|
||||
return web.json_response(
|
||||
{"error": "Missing required fields: url, directory, filename"},
|
||||
status=400,
|
||||
)
|
||||
|
||||
task, err = await self.download_manager.start_download(url, directory, filename, client_id=client_id)
|
||||
if err:
|
||||
status = 409 if "already" in err.lower() else 400
|
||||
return web.json_response({"error": err}, status=status)
|
||||
|
||||
return web.json_response(task.to_dict(), status=201)
|
||||
|
||||
@routes.get("/download/status")
|
||||
@_require_download_api
|
||||
async def get_download_status(request):
|
||||
client_id = request.rel_url.query.get("client_id")
|
||||
return web.json_response(self.download_manager.get_all_tasks(client_id=client_id))
|
||||
|
||||
@routes.get("/download/status/{task_id}")
|
||||
@_require_download_api
|
||||
async def get_download_task_status(request):
|
||||
task_id = request.match_info["task_id"]
|
||||
task_data = self.download_manager.get_task(task_id)
|
||||
if task_data is None:
|
||||
return web.json_response({"error": "Download not found"}, status=404)
|
||||
return web.json_response(task_data)
|
||||
|
||||
@routes.post("/download/pause/{task_id}")
|
||||
@_require_download_api
|
||||
async def post_download_pause(request):
|
||||
task_id = request.match_info["task_id"]
|
||||
err = self.download_manager.pause_download(task_id)
|
||||
if err:
|
||||
return web.json_response({"error": err}, status=400)
|
||||
return web.json_response({"status": "paused"})
|
||||
|
||||
@routes.post("/download/resume/{task_id}")
|
||||
@_require_download_api
|
||||
async def post_download_resume(request):
|
||||
task_id = request.match_info["task_id"]
|
||||
err = self.download_manager.resume_download(task_id)
|
||||
if err:
|
||||
return web.json_response({"error": err}, status=400)
|
||||
return web.json_response({"status": "resumed"})
|
||||
|
||||
@routes.post("/download/cancel/{task_id}")
|
||||
@_require_download_api
|
||||
async def post_download_cancel(request):
|
||||
task_id = request.match_info["task_id"]
|
||||
err = self.download_manager.cancel_download(task_id)
|
||||
if err:
|
||||
return web.json_response({"error": err}, status=400)
|
||||
return web.json_response({"status": "cancelled"})
|
||||
|
||||
async def setup(self):
|
||||
timeout = aiohttp.ClientTimeout(total=None) # no timeout
|
||||
self.client_session = aiohttp.ClientSession(timeout=timeout)
|
||||
if self.download_manager is not None:
|
||||
self.app.on_cleanup.append(self._cleanup_download_manager)
|
||||
|
||||
async def _cleanup_download_manager(self, app):
|
||||
if self.download_manager is not None:
|
||||
await self.download_manager.close()
|
||||
|
||||
def add_routes(self):
|
||||
self.user_manager.add_routes(self.routes)
|
||||
|
||||
Reference in New Issue
Block a user