import threading import asyncio import time import server from aiohttp import web import os import shutil from ..utils.logger import log from ..utils.modelscope_gateway import ModelScopeGateway import folder_paths # 全局下载进度存储 download_progress = {} download_lock = threading.Lock() # 下载进度回调类 class DownloadProgressCallback: def __init__(self, id: str, filename: str, file_size: int, download_id: str): self.id = id self.filename = filename self.file_size = file_size self.download_id = download_id self.progress = 0 self.status = "downloading" # downloading, completed, failed self.error_message = None self.start_time = time.time() # 初始化进度记录 with download_lock: download_progress[download_id] = { "id": id, "filename": filename, "file_size": file_size, "progress": 0, "percentage": 0.0, "status": "downloading", "start_time": self.start_time, "estimated_time": None, "speed": 0.0, "error_message": None } def update(self, size: int): """更新下载进度""" self.progress += size current_time = time.time() elapsed_time = current_time - self.start_time # 计算下载速度和预估时间 if elapsed_time > 0: speed = self.progress / elapsed_time # bytes per second if speed > 0 and self.progress < self.file_size: remaining_bytes = self.file_size - self.progress estimated_time = remaining_bytes / speed else: estimated_time = None else: speed = 0.0 estimated_time = None percentage = (self.progress / self.file_size) * 100 if self.file_size > 0 else 0 # 更新全局进度 with download_lock: if self.download_id in download_progress: # 直接更新字典的值，而不是调用update方法 progress_dict = download_progress[self.download_id] progress_dict["progress"] = self.progress progress_dict["percentage"] = round(percentage, 2) progress_dict["speed"] = round(speed, 2) progress_dict["estimated_time"] = round(estimated_time, 2) if estimated_time else None def end(self, success: bool = True, error_message: str = None): """下载结束回调""" current_time = time.time() total_time = current_time - self.start_time if success: self.status = "completed" # 验证下载完整性 assert self.progress == self.file_size, f"Download incomplete: {self.progress}/{self.file_size}" else: self.status = "failed" self.error_message = error_message # 更新最终状态 with download_lock: if self.download_id in download_progress: # 直接更新字典的值，而不是调用update方法 progress_dict = download_progress[self.download_id] progress_dict["status"] = self.status progress_dict["progress"] = self.file_size if success else self.progress if self.file_size > 0: progress_dict["percentage"] = 100.0 if success else (self.progress / self.file_size) * 100 else: progress_dict["percentage"] = 0.0 progress_dict["total_time"] = round(total_time, 2) progress_dict["error_message"] = self.error_message def fail(self, error_message: str): """下载失败回调""" self.end(success=False, error_message=error_message) # 生成唯一下载ID def generate_download_id() -> str: import uuid return str(uuid.uuid4()) @server.PromptServer.instance.routes.post("/api/download-model") async def download_model(request): """ Download model from ModelScope using SDK """ log.info("Received download-model request") req_json = await request.json() try: id = req_json.get('id') model_id = req_json.get('model_id') model_type = req_json.get('model_type') dest_dir = req_json.get('dest_dir') # 验证必需参数 if not id: return web.json_response({ "success": False, "message": "Missing required parameter: id" }) if not model_id: return web.json_response({ "success": False, "message": "Missing required parameter: model_id" }) if not model_type: return web.json_response({ "success": False, "message": "Missing required parameter: model_type" }) log.info(f"Downloading model: {model_id} (type: {model_type})") # 生成下载ID download_id = generate_download_id() # 计算目标目录：优先使用传入的dest_dir，否则使用ComfyUI的models目录下对应类型 resolved_dest_dir = None if dest_dir: resolved_dest_dir = os.path.abspath(os.path.expanduser(f"models/{dest_dir}")) else: try: model_type_paths = folder_paths.get_folder_paths(model_type) resolved_dest_dir = model_type_paths[0] if model_type_paths else os.path.join(folder_paths.models_dir, model_type) except Exception: resolved_dest_dir = os.path.join(folder_paths.models_dir, model_type) # 创建进度回调 progress_callback = DownloadProgressCallback( id=id, filename=f"{model_id}.{model_type}", file_size=0, # 实际大小会在下载过程中获取 download_id=download_id ) # 启动下载任务（异步执行） async def download_task(): try: # 调用下载方法 - 使用snapshot_download from modelscope import snapshot_download # 创建进度回调包装器 - 实现ModelScope期望的接口 class ProgressWrapper: """Factory that returns a per-file progress object with update/end.""" def __init__(self, download_id: str): self.download_id = download_id def __call__(self, file_name: str, file_size: int): # Create a per-file progress tracker expected by ModelScope download_id = self.download_id class _PerFileProgress: def __init__(self, fname: str, fsize: int): self.file_name = fname self.file_size = max(int(fsize or 0), 0) self.progress = 0 self.last_update_time = time.time() self.last_downloaded = 0 with download_lock: if download_id in download_progress: # If unknown size, keep 0 to avoid div-by-zero download_progress[download_id]["file_size"] = self.file_size def update(self, size: int): try: self.progress += int(size or 0) now = time.time() # Update global progress with download_lock: if download_id in download_progress: dp = download_progress[download_id] dp["progress"] = self.progress if self.file_size > 0: dp["percentage"] = round(self.progress * 100.0 / self.file_size, 2) # speed elapsed = max(now - self.last_update_time, 1e-6) speed = (self.progress - self.last_downloaded) / elapsed dp["speed"] = round(speed, 2) self.last_update_time = now self.last_downloaded = self.progress except Exception as e: log.error(f"Error in progress update: {e}") def end(self): # Called by modelscope when a file finishes with download_lock: if download_id in download_progress: dp = download_progress[download_id] if self.file_size > 0: dp["progress"] = self.file_size dp["percentage"] = 100.0 return _PerFileProgress(file_name, file_size) progress_wrapper = ProgressWrapper(download_id) # 添加调试日志 log.info(f"Starting download with progress wrapper: {download_id}") # 在线程中执行阻塞的下载，避免阻塞事件循环 from functools import partial local_dir = await asyncio.to_thread( partial( snapshot_download, model_id=model_id, cache_dir=resolved_dest_dir, progress_callbacks=[progress_wrapper] ) ) # 下载完成 progress_callback.end(success=True) log.info(f"Model downloaded successfully to: {local_dir}") # 下载后遍历目录，将所有重要权重/资源文件移动到最外层（与目录同级，即 resolved_dest_dir） try: moved_count = 0 allowed_exts = { ".safetensors", ".ckpt", ".pt", ".pth", ".bin" # ".msgpack", ".json", ".yaml", ".yml", ".toml", ".png", ".onnx" } for root, dirs, files in os.walk(local_dir): for name in files: ext = os.path.splitext(name)[1].lower() if ext in allowed_exts: src_path = os.path.join(root, name) target_dir = resolved_dest_dir os.makedirs(target_dir, exist_ok=True) target_path = os.path.join(target_dir, name) # 如果已经在目标目录则跳过 if os.path.abspath(os.path.dirname(src_path)) == os.path.abspath(target_dir): continue # 处理重名情况：自动追加 _1, _2 ... if os.path.exists(target_path): base, ext_real = os.path.splitext(name) idx = 1 while True: candidate = f"{base}_{idx}{ext_real}" candidate_path = os.path.join(target_dir, candidate) if not os.path.exists(candidate_path): target_path = candidate_path break idx += 1 shutil.move(src_path, target_path) moved_count += 1 log.info( f"Moved {moved_count} files with extensions {sorted(list(allowed_exts))} to: {resolved_dest_dir}") except Exception as move_err: log.error(f"Post-download move failed: {move_err}") except Exception as e: progress_callback.fail(str(e)) log.error(f"Download failed: {str(e)}") # 启动异步下载任务 asyncio.create_task(download_task()) return web.json_response({ "success": True, "data": { "download_id": download_id, "id": id, "model_id": model_id, "model_type": model_type, "dest_dir": resolved_dest_dir, "status": "started" }, "message": f"Download started for model '{model_id}'" }) except ImportError as e: log.error(f"ModelScope SDK not installed: {str(e)}") return web.json_response({ "success": False, "message": "ModelScope SDK not installed. Please install with: pip install modelscope" }) except Exception as e: log.error(f"Error starting download: {str(e)}") import traceback traceback.print_exc() return web.json_response({ "success": False, "message": f"Failed to start download: {str(e)}" }) @server.PromptServer.instance.routes.get("/api/download-progress/{download_id}") async def get_download_progress(request): """ Get download progress by download ID """ download_id = request.match_info.get('download_id') if not download_id: return web.json_response({ "success": False, "message": "Missing download_id parameter" }) with download_lock: progress_info = download_progress.get(download_id) if not progress_info: return web.json_response({ "success": False, "message": f"Download ID {download_id} not found" }) return web.json_response({ "success": True, "data": progress_info, "message": "Download progress retrieved successfully" }) @server.PromptServer.instance.routes.get("/api/download-progress") async def list_downloads(request): """ List all active downloads """ with download_lock: downloads = list(download_progress.keys()) return web.json_response({ "success": True, "data": { "downloads": downloads, "count": len(downloads) }, "message": "Download list retrieved successfully" }) @server.PromptServer.instance.routes.delete("/api/download-progress/{download_id}") async def clear_download_progress(request): """ Clear download progress record (for cleanup) """ download_id = request.match_info.get('download_id') if not download_id: return web.json_response({ "success": False, "message": "Missing download_id parameter" }) with download_lock: if download_id in download_progress: del download_progress[download_id] return web.json_response({ "success": True, "message": f"Download progress {download_id} cleared successfully" }) else: return web.json_response({ "success": False, "message": f"Download ID {download_id} not found" }) @server.PromptServer.instance.routes.get("/api/model-searchs") async def model_suggests(request): """ Get model search list by keyword """ log.info("Received model-search request") try: keyword = request.query.get('keyword') if not keyword: return web.json_response({ "success": False, "message": "Missing required parameter: keyword" }) # 创建ModelScope网关实例 gateway = ModelScopeGateway() suggests = gateway.search(name=keyword) list = suggests["data"] if suggests.get("data") else [] return web.json_response({ "success": True, "data": { "searchs": list, "total": len(list) }, "message": f"Get searchs successfully" }) except ImportError as e: log.error(f"ModelScope SDK not installed: {str(e)}") return web.json_response({ "success": False, "message": "ModelScope SDK not installed. Please install with: pip install modelscope" }) except Exception as e: log.error(f"Error get model searchs: {str(e)}") import traceback traceback.print_exc() return web.json_response({ "success": False, "message": f"Get model searchs failed: {str(e)}" }) @server.PromptServer.instance.routes.get("/api/model-paths") async def model_paths(request): """ Get model paths by type """ log.info("Received model-paths request") try: model_paths = list(folder_paths.folder_names_and_paths.keys()) return web.json_response({ "success": True, "data": { "paths": model_paths, }, "message": f"Get paths successfully" }) except Exception as e: log.error(f"Error get model path: {str(e)}") import traceback traceback.print_exc() return web.json_response({ "success": False, "message": f"Get model failed: {str(e)}" })