88 lines
3.5 KiB
Python
88 lines
3.5 KiB
Python
"""Toolbox — ComfyUI extension with model download support."""
|
|
import os
|
|
import asyncio
|
|
import aiohttp
|
|
from aiohttp import web
|
|
|
|
WEB_DIRECTORY = "./web"
|
|
NODE_CLASS_MAPPINGS = {}
|
|
__all__ = ["NODE_CLASS_MAPPINGS", "WEB_DIRECTORY"]
|
|
|
|
_downloads = {}
|
|
|
|
def _models_dir():
|
|
return os.path.join(
|
|
os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))),
|
|
"models"
|
|
)
|
|
|
|
try:
|
|
from server import PromptServer
|
|
|
|
@PromptServer.instance.routes.post("/toolbox/download")
|
|
async def start_download(request):
|
|
data = await request.json()
|
|
url = data.get("url", "").strip()
|
|
folder = data.get("folder", "").strip()
|
|
filename = data.get("filename", "").strip()
|
|
if not url or not folder or not filename:
|
|
return web.json_response({"error": "url, folder, filename required"}, status=400)
|
|
|
|
target_dir = os.path.join(_models_dir(), folder)
|
|
os.makedirs(target_dir, exist_ok=True)
|
|
target_path = os.path.join(target_dir, filename)
|
|
|
|
if os.path.exists(target_path):
|
|
return web.json_response({"error": "file already exists", "path": target_path}, status=409)
|
|
|
|
dl_id = f"{folder}/{filename}"
|
|
if dl_id in _downloads and _downloads[dl_id].get("status") == "downloading":
|
|
return web.json_response({"error": "already downloading"}, status=409)
|
|
|
|
_downloads[dl_id] = {"status": "starting", "progress": 0, "total": 0, "error": None}
|
|
asyncio.ensure_future(_do_download(dl_id, url, target_path))
|
|
return web.json_response({"id": dl_id, "status": "started", "path": target_path})
|
|
|
|
async def _do_download(dl_id, url, target_path):
|
|
tmp = target_path + ".part"
|
|
try:
|
|
_downloads[dl_id]["status"] = "downloading"
|
|
async with aiohttp.ClientSession() as session:
|
|
async with session.get(url, allow_redirects=True) as resp:
|
|
if resp.status != 200:
|
|
_downloads[dl_id] = {"status": "error", "error": f"HTTP {resp.status}", "progress": 0, "total": 0}
|
|
return
|
|
total = int(resp.headers.get("content-length", 0))
|
|
_downloads[dl_id]["total"] = total
|
|
done = 0
|
|
with open(tmp, "wb") as f:
|
|
async for chunk in resp.content.iter_chunked(1024 * 1024):
|
|
f.write(chunk)
|
|
done += len(chunk)
|
|
_downloads[dl_id]["progress"] = done
|
|
os.rename(tmp, target_path)
|
|
_downloads[dl_id]["status"] = "done"
|
|
_downloads[dl_id]["progress"] = _downloads[dl_id]["total"]
|
|
except Exception as e:
|
|
_downloads[dl_id]["status"] = "error"
|
|
_downloads[dl_id]["error"] = str(e)
|
|
if os.path.exists(tmp):
|
|
try: os.remove(tmp)
|
|
except: pass
|
|
|
|
@PromptServer.instance.routes.get("/toolbox/download/status")
|
|
async def download_status(request):
|
|
return web.json_response(_downloads)
|
|
|
|
@PromptServer.instance.routes.post("/toolbox/download/cancel")
|
|
async def cancel_download(request):
|
|
# Mark as cancelled — the download loop checks this
|
|
data = await request.json()
|
|
dl_id = data.get("id", "")
|
|
if dl_id in _downloads:
|
|
_downloads[dl_id]["status"] = "cancelled"
|
|
return web.json_response({"ok": True})
|
|
|
|
except ImportError:
|
|
print("[toolbox] Could not import PromptServer — download features disabled")
|