TBNilles/ComfyUI-Model-Manager
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: download module error (#141)

Browse Source
This commit is contained in:
Hayden
2025-02-19 14:37:27 +08:00
committed by GitHub
parent ea26ec5098
commit 05fa31f2c5
1 changed files with 322 additions and 329 deletions
Download Patch File Download Diff File Expand all files Collapse all files

241
py/download.py
View File

@@ -3,8 +3,13 @@ import uuid
import time
import requests
import folder_paths
from typing import Callable, Awaitable, Any, Literal, Union, Optional
from dataclasses import dataclass
from aiohttp import web
from . import config
from . import utils
from . import thread
@@ -87,17 +92,95 @@ class TaskContent:
}
download_model_task_status: dict[str, TaskStatus] = {}
download_thread_pool = thread.DownloadThreadPool()
class ModelDownload:
def add_routes(self, routes):
@routes.get("/model-manager/download/task")
async def scan_download_tasks(request):
"""
Read download task list.
"""
try:
result = await self.scan_model_download_task_list()
return web.json_response({"success": True, "data": result})
except Exception as e:
error_msg = f"Read download task list failed: {e}"
utils.print_error(error_msg)
return web.json_response({"success": False, "error": error_msg})
def set_task_content(task_id: str, task_content: Union[TaskContent, dict]):
@routes.put("/model-manager/download/{task_id}")
async def resume_download_task(request):
"""
Toggle download task status.
"""
try:
task_id = request.match_info.get("task_id", None)
if task_id is None:
raise web.HTTPBadRequest(reason="Invalid task id")
json_data = await request.json()
status = json_data.get("status", None)
if status == "pause":
await self.pause_model_download_task(task_id)
elif status == "resume":
await self.download_model(task_id, request)
else:
raise web.HTTPBadRequest(reason="Invalid status")
return web.json_response({"success": True})
except Exception as e:
error_msg = f"Resume download task failed: {str(e)}"
utils.print_error(error_msg)
return web.json_response({"success": False, "error": error_msg})
@routes.delete("/model-manager/download/{task_id}")
async def delete_model_download_task(request):
"""
Delete download task.
"""
task_id = request.match_info.get("task_id", None)
try:
await self.delete_model_download_task(task_id)
return web.json_response({"success": True})
except Exception as e:
error_msg = f"Delete download task failed: {str(e)}"
utils.print_error(error_msg)
return web.json_response({"success": False, "error": error_msg})
@routes.post("/model-manager/model")
async def create_model(request):
"""
Create a new model.
request body: x-www-form-urlencoded
- type: model type.
- pathIndex: index of the model folders.
- fullname: filename that relative to the model folder.
- previewFile: preview file.
- description: description.
- downloadPlatform: download platform.
- downloadUrl: download url.
- hash: a JSON string containing the hash value of the downloaded model.
"""
task_data = await request.post()
task_data = dict(task_data)
try:
task_id = await self.create_model_download_task(task_data, request)
return web.json_response({"success": True, "data": {"taskId": task_id}})
except Exception as e:
error_msg = f"Create model download task failed: {str(e)}"
utils.print_error(error_msg)
return web.json_response({"success": False, "error": error_msg})
download_model_task_status: dict[str, TaskStatus] = {}
download_thread_pool = thread.DownloadThreadPool()
def set_task_content(self, task_id: str, task_content: Union[TaskContent, dict]):
download_path = utils.get_download_path()
task_file_path = utils.join_path(download_path, f"{task_id}.task")
utils.save_dict_pickle_file(task_file_path, task_content)
def get_task_content(task_id: str):
def get_task_content(self, task_id: str):
download_path = utils.get_download_path()
task_file = utils.join_path(download_path, f"{task_id}.task")
if not os.path.isfile(task_file):
@@ -107,13 +190,12 @@ def get_task_content(task_id: str):
return task_content
return TaskContent(**task_content)
def get_task_status(task_id: str):
task_status = download_model_task_status.get(task_id, None)
def get_task_status(self, task_id: str):
task_status = self.download_model_task_status.get(task_id, None)
if task_status is None:
download_path = utils.get_download_path()
task_content = get_task_content(task_id)
task_content = self.get_task_content(task_id)
download_file = utils.join_path(download_path, f"{task_id}.download")
download_size = 0
if os.path.exists(download_file):
@@ -131,16 +213,14 @@ def get_task_status(task_id: str):
progress=download_size / total_size * 100 if total_size > 0 else 0,
)
download_model_task_status[task_id] = task_status
self.download_model_task_status[task_id] = task_status
return task_status
def delete_task_status(self, task_id: str):
self.download_model_task_status.pop(task_id, None)
def delete_task_status(task_id: str):
download_model_task_status.pop(task_id, None)
async def scan_model_download_task_list():
async def scan_model_download_task_list(self):
"""
Scan the download directory and send the task list to the client.
"""
@@ -155,13 +235,12 @@ async def scan_model_download_task_list():
task_list: list[dict] = []
for task_file in task_files:
task_id = task_file.replace(".task", "")
task_status = get_task_status(task_id)
task_status = self.get_task_status(task_id)
task_list.append(task_status.to_dict())
return task_list
async def create_model_download_task(task_data: dict, request):
async def create_model_download_task(self, task_data: dict, request):
"""
Creates a download task for the given data.
"""
@@ -185,7 +264,7 @@ async def create_model_download_task(task_data: dict, request):
try:
preview_file = task_data.pop("previewFile", None)
utils.save_model_preview_image(task_path, preview_file, download_platform)
set_task_content(task_id, task_data)
self.set_task_content(task_id, task_data)
task_status = TaskStatus(
taskId=task_id,
type=model_type,
@@ -194,23 +273,21 @@ async def create_model_download_task(task_data: dict, request):
platform=download_platform,
totalSize=float(task_data.get("sizeBytes", 0)),
)
download_model_task_status[task_id] = task_status
self.download_model_task_status[task_id] = task_status
await utils.send_json("create_download_task", task_status.to_dict())
except Exception as e:
await delete_model_download_task(task_id)
await self.delete_model_download_task(task_id)
raise RuntimeError(str(e)) from e
await download_model(task_id, request)
await self.download_model(task_id, request)
return task_id
async def pause_model_download_task(task_id: str):
task_status = get_task_status(task_id=task_id)
async def pause_model_download_task(self, task_id: str):
task_status = self.get_task_status(task_id=task_id)
task_status.status = "pause"
async def delete_model_download_task(task_id: str):
task_status = get_task_status(task_id)
async def delete_model_download_task(self, task_id: str):
task_status = self.get_task_status(task_id)
is_running = task_status.status == "doing"
task_status.status = "waiting"
await utils.send_json("delete_download_task", task_id)
@@ -225,20 +302,19 @@ async def delete_model_download_task(task_id: str):
for task_file in task_file_list:
task_file_target = os.path.splitext(task_file)[0]
if task_file_target == task_id:
delete_task_status(task_id)
self.delete_task_status(task_id)
os.remove(utils.join_path(download_dir, task_file))
await utils.send_json("delete_download_task", task_id)
async def download_model(task_id: str, request):
async def download_model(self, task_id: str, request):
async def download_task(task_id: str):
async def report_progress(task_status: TaskStatus):
await utils.send_json("update_download_task", task_status.to_dict())
try:
# When starting a task from the queue, the task may not exist
task_status = get_task_status(task_id)
task_status = self.get_task_status(task_id)
except:
return
@@ -263,7 +339,7 @@ async def download_model(task_id: str, request):
headers["Authorization"] = f"Bearer {api_key}"
progress_interval = 1.0
await download_model_file(
await self.download_model_file(
task_id=task_id,
headers=headers,
progress_callback=report_progress,
@@ -277,9 +353,9 @@ async def download_model(task_id: str, request):
utils.print_error(str(e))
try:
status = download_thread_pool.submit(download_task, task_id)
status = self.download_thread_pool.submit(download_task, task_id)
if status == "Waiting":
task_status = get_task_status(task_id)
task_status = self.get_task_status(task_id)
task_status.status = "waiting"
await utils.send_json("update_download_task", task_status.to_dict())
except Exception as e:
@@ -289,13 +365,13 @@ async def download_model(task_id: str, request):
task_status.error = None
utils.print_error(str(e))
async def download_model_file(
async def download_model_file(
self,
task_id: str,
headers: dict,
progress_callback: Callable[[TaskStatus], Awaitable[Any]],
interval: float = 1.0,
):
):
async def download_complete():
"""
@@ -331,8 +407,8 @@ async def download_model_file(
last_update_time = time.time()
last_downloaded_size = downloaded_size
task_status = get_task_status(task_id)
task_content = get_task_content(task_id)
task_status = self.get_task_status(task_id)
task_content = self.get_task_content(task_id)
# Check download uri
model_url = task_content.downloadUrl
@@ -385,7 +461,7 @@ async def download_model_file(
total_size = float(response.headers.get("content-length", 0))
task_content.sizeBytes = total_size
task_status.totalSize = total_size
set_task_content(task_id, task_content)
self.set_task_content(task_id, task_content)
await utils.send_json("update_download_task", task_content.to_dict())
with open(download_tmp_file, "ab") as f:
@@ -406,86 +482,3 @@ async def download_model_file(
else:
task_status.status = "pause"
await utils.send_json("update_download_task", task_status.to_dict())
from aiohttp import web
class ModelDownload:
def add_routes(self, routes):
@routes.get("/model-manager/download/task")
async def scan_download_tasks(request):
"""
Read download task list.
"""
try:
result = await scan_model_download_task_list()
return web.json_response({"success": True, "data": result})
except Exception as e:
error_msg = f"Read download task list failed: {e}"
utils.print_error(error_msg)
return web.json_response({"success": False, "error": error_msg})
@routes.put("/model-manager/download/{task_id}")
async def resume_download_task(request):
"""
Toggle download task status.
"""
try:
task_id = request.match_info.get("task_id", None)
if task_id is None:
raise web.HTTPBadRequest(reason="Invalid task id")
json_data = await request.json()
status = json_data.get("status", None)
if status == "pause":
await pause_model_download_task(task_id)
elif status == "resume":
await download_model(task_id, request)
else:
raise web.HTTPBadRequest(reason="Invalid status")
return web.json_response({"success": True})
except Exception as e:
error_msg = f"Resume download task failed: {str(e)}"
utils.print_error(error_msg)
return web.json_response({"success": False, "error": error_msg})
@routes.delete("/model-manager/download/{task_id}")
async def delete_model_download_task(request):
"""
Delete download task.
"""
task_id = request.match_info.get("task_id", None)
try:
await delete_model_download_task(task_id)
return web.json_response({"success": True})
except Exception as e:
error_msg = f"Delete download task failed: {str(e)}"
utils.print_error(error_msg)
return web.json_response({"success": False, "error": error_msg})
@routes.post("/model-manager/model")
async def create_model(request):
"""
Create a new model.
request body: x-www-form-urlencoded
- type: model type.
- pathIndex: index of the model folders.
- fullname: filename that relative to the model folder.
- previewFile: preview file.
- description: description.
- downloadPlatform: download platform.
- downloadUrl: download url.
- hash: a JSON string containing the hash value of the downloaded model.
"""
task_data = await request.post()
task_data = dict(task_data)
try:
task_id = await create_model_download_task(task_data, request)
return web.json_response({"success": True, "data": {"taskId": task_id}})
except Exception as e:
error_msg = f"Create model download task failed: {str(e)}"
utils.print_error(error_msg)
return web.json_response({"success": False, "error": error_msg})