# Copyright (C) 2025 AIDC-AI # Licensed under the MIT License. import json import asyncio import time from typing import Optional, TypedDict, List, Union import threading from collections import defaultdict from sqlalchemy.orm import identity from ..utils.globals import set_language from ..utils.auth_utils import extract_and_store_api_key import server from aiohttp import web import base64 # Import the MCP client function import os import shutil from ..service.debug_agent import debug_workflow_errors from ..dao.workflow_table import save_workflow_data, get_workflow_data_by_id, update_workflow_ui_by_id from ..service.mcp_client import comfyui_agent_invoke from ..utils.request_context import set_request_context, get_session_id from ..utils.logger import log from ..utils.modelscope_gateway import ModelScopeGateway import folder_paths # 全局下载进度存储 download_progress = {} download_lock = threading.Lock() # 不再使用内存存储会话消息，改为从前端传递历史消息 # 在文件开头添加 STATIC_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)), "public") # Define types using TypedDict class Node(TypedDict): name: str # 节点名称 description: str # 节点描述 image: str # 节点图片url，可为空 github_url: str # 节点github地址 from_index: int # 节点在列表中的位置 to_index: int # 节点在列表中的位置 class NodeInfo(TypedDict): existing_nodes: List[Node] # 已安装的节点 missing_nodes: List[Node] # 未安装的节点 class Workflow(TypedDict, total=False): id: Optional[int] # 工作流id name: Optional[str] # 工作流名称 description: Optional[str] # 工作流描述 image: Optional[str] # 工作流图片 workflow: Optional[str] # 工作流 class ExtItem(TypedDict): type: str # 扩展类型 data: Union[dict, list] # 扩展数据 class ChatResponse(TypedDict): session_id: str # 会话id text: Optional[str] # 返回文本 finished: bool # 是否结束 type: str # 返回的类型 format: str # 返回的格式 ext: Optional[List[ExtItem]] # 扩展信息 # 下载进度回调类 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()) async def upload_to_oss(file_data: bytes, filename: str) -> str: # TODO: Implement your OSS upload logic here # For now, save locally and return a placeholder URL try: # Create uploads directory if it doesn't exist uploads_dir = os.path.join(os.path.dirname(__file__), '..', '..', 'uploads') os.makedirs(uploads_dir, exist_ok=True) # Generate unique filename to avoid conflicts import uuid unique_filename = f"{uuid.uuid4().hex}_{filename}" file_path = os.path.join(uploads_dir, unique_filename) # Save file locally with open(file_path, 'wb') as f: f.write(file_data) # Return a local URL or base64 data URL for now # In production, replace this with actual OSS URL base64_data = base64.b64encode(file_data).decode('utf-8') # Determine MIME type based on file extension if filename.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.webp')): mime_type = f"image/{filename.split('.')[-1].lower()}" if mime_type == "image/jpg": mime_type = "image/jpeg" else: mime_type = "image/jpeg" # Default return f"data:{mime_type};base64,{base64_data}" except Exception as e: log.error(f"Error uploading file {filename}: {str(e)}") # Return original base64 data if upload fails base64_data = base64.b64encode(file_data).decode('utf-8') return f"data:image/jpeg;base64,{base64_data}" # 关联用户消息和AI响应的checkpoint信息 def processMessagesWithCheckpoints(messages): # ... existing code ... pass @server.PromptServer.instance.routes.post("/api/chat/invoke") async def invoke_chat(request): log.info("Received invoke_chat request") # Extract and store API key from Authorization header extract_and_store_api_key(request) req_json = await request.json() log.info("Request JSON:", req_json) response = web.StreamResponse( status=200, reason='OK', headers={ 'Content-Type': 'application/json', 'X-Content-Type-Options': 'nosniff' } ) await response.prepare(request) session_id = req_json.get('session_id') prompt = req_json.get('prompt') images = req_json.get('images', []) intent = req_json.get('intent') ext = req_json.get('ext') historical_messages = req_json.get('messages', []) workflow_checkpoint_id = req_json.get('workflow_checkpoint_id') # 获取当前语言 language = request.headers.get('Accept-Language', 'en') set_language(language) # 构建配置信息 config = { "session_id": session_id, "workflow_checkpoint_id": workflow_checkpoint_id, "openai_api_key": request.headers.get('Openai-Api-Key'), "openai_base_url": request.headers.get('Openai-Base-Url'), # Workflow LLM settings (optional, used by tools/agents that need a different LLM) "workflow_llm_api_key": request.headers.get('Workflow-LLM-Api-Key'), "workflow_llm_base_url": request.headers.get('Workflow-LLM-Base-Url'), "workflow_llm_model": request.headers.get('Workflow-LLM-Model'), "model_select": next((x['data'][0] for x in ext if x['type'] == 'model_select' and x.get('data')), None) } # 设置请求上下文 - 这里建立context隔离 set_request_context(session_id, workflow_checkpoint_id, config) # 图片处理已移至前端，图片信息现在包含在历史消息的OpenAI格式中 # 保留空的处理逻辑以向后兼容，但实际不再使用 processed_images = [] if images and len(images) > 0: log.info(f"Note: Received {len(images)} images in legacy format, but using OpenAI message format instead") # 历史消息已经从前端传递过来，格式为OpenAI格式，直接使用 log.info(f"-- Received {len(historical_messages)} historical messages") # Log workflow checkpoint ID if provided (workflow is now pre-saved before invoke) if workflow_checkpoint_id: log.info(f"Using workflow checkpoint ID: {workflow_checkpoint_id} for session {session_id}") else: log.info(f"No workflow checkpoint ID provided for session {session_id}") # 历史消息已经从前端传递过来，包含了正确格式的OpenAI消息（包括图片） # 直接使用前端传递的历史消息，无需重新构建当前消息 openai_messages = historical_messages # 不再需要创建用户消息存储到后端，前端负责消息存储 try: # Call the MCP client to get streaming response with historical messages and image support # Pass OpenAI-formatted messages and processed images to comfyui_agent_invoke accumulated_text = "" ext_data = None finished = True # Default to True has_sent_response = False previous_text_length = 0 log.info(f"config: {config}") # Pass messages in OpenAI format (images are now included in messages) # Config is now available through request context async for result in comfyui_agent_invoke(openai_messages, None): # The MCP client now returns tuples (text, ext_with_finished) where ext_with_finished includes finished status if isinstance(result, tuple) and len(result) == 2: text, ext_with_finished = result if text: accumulated_text = text # text from MCP is already accumulated if ext_with_finished: # Extract ext data and finished status from the structured response ext_data = ext_with_finished.get("data") finished = ext_with_finished.get("finished", True) log.info(f"-- Received ext data: {ext_data}, finished: {finished}") else: # Handle single text chunk (backward compatibility) text_chunk = result if text_chunk: accumulated_text += text_chunk log.info(f"-- Received text chunk: '{text_chunk}', total length: {len(accumulated_text)}") # Send streaming response if we have new text content # Only send intermediate responses during streaming (not the final one) if accumulated_text and len(accumulated_text) > previous_text_length: chat_response = ChatResponse( session_id=session_id, text=accumulated_text, finished=False, # Always false during streaming type="message", format="markdown", ext=None # ext is only sent in final response ) await response.write(json.dumps(chat_response).encode() + b"\n") previous_text_length = len(accumulated_text) await asyncio.sleep(0.01) # Small delay for streaming effect # Send final response with proper finished logic from MCP client log.info(f"-- Sending final response: {len(accumulated_text)} chars, ext: {bool(ext_data)}, finished: {finished}") final_response = ChatResponse( session_id=session_id, text=accumulated_text, finished=finished, # Use finished status from MCP client type="message", format="markdown", ext=ext_data ) await response.write(json.dumps(final_response).encode() + b"\n") # AI响应不再存储到后端，前端负责消息存储 except Exception as e: log.error(f"Error in invoke_chat: {str(e)}") error_response = ChatResponse( session_id=session_id, text=f"I apologize, but an error occurred: {str(e)}", finished=True, # Always finish on error type="message", format="text", ext=None ) await response.write(json.dumps(error_response).encode() + b"\n") await response.write_eof() return response @server.PromptServer.instance.routes.post("/api/save-workflow-checkpoint") async def save_workflow_checkpoint(request): """ Save workflow checkpoint for restore functionality """ log.info("Received save-workflow-checkpoint request") req_json = await request.json() try: session_id = req_json.get('session_id') workflow_api = req_json.get('workflow_api') # API format workflow workflow_ui = req_json.get('workflow_ui') # UI format workflow checkpoint_type = req_json.get('checkpoint_type', 'debug_start') # debug_start, debug_complete, user_message_checkpoint message_id = req_json.get('message_id') # User message ID for linking (optional) if not session_id or not workflow_api: return web.json_response({ "success": False, "message": "Missing required parameters: session_id and workflow_api" }) # Save workflow with checkpoint type in attributes attributes = { "checkpoint_type": checkpoint_type, "timestamp": time.time() } # Set description and additional attributes based on checkpoint type if checkpoint_type == "user_message_checkpoint" and message_id: attributes.update({ "description": f"Workflow checkpoint before user message {message_id}", "message_id": message_id, "source": "user_message_pre_invoke" }) else: attributes["description"] = f"Workflow checkpoint: {checkpoint_type}" version_id = save_workflow_data( session_id=session_id, workflow_data=workflow_api, workflow_data_ui=workflow_ui, attributes=attributes ) log.info(f"Workflow checkpoint saved with version ID: {version_id}") # Return response format based on checkpoint type response_data = { "version_id": version_id, "checkpoint_type": checkpoint_type } if checkpoint_type == "user_message_checkpoint" and message_id: response_data.update({ "checkpoint_id": version_id, # Add checkpoint_id alias for user message checkpoints "message_id": message_id }) return web.json_response({ "success": True, "data": response_data, "message": f"Workflow checkpoint saved successfully" }) except Exception as e: log.error(f"Error saving workflow checkpoint: {str(e)}") return web.json_response({ "success": False, "message": f"Failed to save workflow checkpoint: {str(e)}" }) @server.PromptServer.instance.routes.get("/api/restore-workflow-checkpoint") async def restore_workflow_checkpoint(request): """ Restore workflow checkpoint by version ID """ log.info("Received restore-workflow-checkpoint request") try: version_id = request.query.get('version_id') if not version_id: return web.json_response({ "success": False, "message": "Missing required parameter: version_id" }) try: version_id = int(version_id) except ValueError: return web.json_response({ "success": False, "message": "Invalid version_id format" }) # Get workflow data by version ID workflow_version = get_workflow_data_by_id(version_id) if not workflow_version: return web.json_response({ "success": False, "message": f"Workflow version {version_id} not found" }) log.info(f"Restored workflow checkpoint version ID: {version_id}") return web.json_response({ "success": True, "data": { "version_id": version_id, "workflow_data": workflow_version.get('workflow_data'), "workflow_data_ui": workflow_version.get('workflow_data_ui'), "attributes": workflow_version.get('attributes'), "created_at": workflow_version.get('created_at') }, "message": f"Workflow checkpoint restored successfully" }) except Exception as e: log.error(f"Error restoring workflow checkpoint: {str(e)}") return web.json_response({ "success": False, "message": f"Failed to restore workflow checkpoint: {str(e)}" }) @server.PromptServer.instance.routes.post("/api/debug-agent") async def invoke_debug(request): """ Debug agent endpoint for analyzing ComfyUI workflow errors """ log.info("Received debug-agent request") # Extract and store API key from Authorization header extract_and_store_api_key(request) req_json = await request.json() response = web.StreamResponse( status=200, reason='OK', headers={ 'Content-Type': 'application/json', 'X-Content-Type-Options': 'nosniff' } ) await response.prepare(request) session_id = req_json.get('session_id') workflow_data = req_json.get('workflow_data') # Get configuration from headers (OpenAI settings) config = { "session_id": session_id, "model": "gemini-2.5-flash", # Default model for debug agents "openai_api_key": request.headers.get('Openai-Api-Key'), "openai_base_url": request.headers.get('Openai-Base-Url'), # Workflow LLM settings (optional) "workflow_llm_api_key": request.headers.get('Workflow-LLM-Api-Key'), "workflow_llm_base_url": request.headers.get('Workflow-LLM-Base-Url'), "workflow_llm_model": request.headers.get('Workflow-LLM-Model'), } # 获取当前语言 language = request.headers.get('Accept-Language', 'en') set_language(language) # 设置请求上下文 - 为debug请求建立context隔离 set_request_context(session_id, None, config) log.info(f"Debug agent config: {config}") log.info(f"Session ID: {session_id}") log.info(f"Workflow nodes: {list(workflow_data.keys()) if workflow_data else 'None'}") try: # Call the debug agent with streaming response accumulated_text = "" final_ext_data = None finished = False async for result in debug_workflow_errors(workflow_data): # Stream the response if isinstance(result, tuple) and len(result) == 2: text, ext = result if text: accumulated_text = text # text is already accumulated from debug agent # Handle new ext format from debug_agent matching mcp-client if ext: if isinstance(ext, dict) and "data" in ext and "finished" in ext: # New format: {"data": ext, "finished": finished} final_ext_data = ext["data"] finished = ext["finished"] # 检查是否包含需要实时发送的workflow_update或param_update has_realtime_ext = False if final_ext_data: for ext_item in final_ext_data: if ext_item.get("type") in ["workflow_update", "param_update"]: has_realtime_ext = True break # 如果包含实时ext数据或者已完成，则发送响应 if has_realtime_ext or finished: chat_response = ChatResponse( session_id=session_id, text=accumulated_text, finished=finished, type="message", format="markdown", ext=final_ext_data # 发送ext数据 ) await response.write(json.dumps(chat_response).encode() + b"\n") elif not finished: # 只有文本更新，不发送ext数据 chat_response = ChatResponse( session_id=session_id, text=accumulated_text, finished=False, type="message", format="markdown", ext=None ) await response.write(json.dumps(chat_response).encode() + b"\n") else: # Legacy format: direct ext data (for backward compatibility) final_ext_data = ext finished = False # Create streaming response chat_response = ChatResponse( session_id=session_id, text=accumulated_text, finished=False, type="message", format="markdown", ext=ext ) await response.write(json.dumps(chat_response).encode() + b"\n") else: # No ext data, just text streaming chat_response = ChatResponse( session_id=session_id, text=accumulated_text, finished=False, type="message", format="markdown", ext=None ) await response.write(json.dumps(chat_response).encode() + b"\n") await asyncio.sleep(0.01) # Small delay for streaming effect # Send final response final_response = ChatResponse( session_id=session_id, text=accumulated_text, finished=True, type="message", format="markdown", ext=final_ext_data if final_ext_data else [{"type": "debug_complete", "data": {"status": "completed"}}] ) # Save workflow checkpoint after debug completion if we have workflow_data if workflow_data and accumulated_text: try: current_session_id = get_session_id() checkpoint_id = save_workflow_data( session_id=current_session_id, workflow_data=workflow_data, workflow_data_ui=None, # UI format not available in debug agent attributes={ "checkpoint_type": "debug_complete", "description": "Workflow state after debug completion", "timestamp": time.time() } ) # Add checkpoint info to ext data if final_response["ext"]: final_response["ext"].append({ "type": "debug_checkpoint", "data": { "checkpoint_id": checkpoint_id, "checkpoint_type": "debug_complete" } }) else: final_response["ext"] = [{ "type": "debug_checkpoint", "data": { "checkpoint_id": checkpoint_id, "checkpoint_type": "debug_complete" } }] log.info(f"Debug completion checkpoint saved with ID: {checkpoint_id}") except Exception as checkpoint_error: log.error(f"Failed to save debug completion checkpoint: {checkpoint_error}") await response.write(json.dumps(final_response).encode() + b"\n") log.info("Debug agent processing complete") except Exception as e: log.error(f"Error in debug agent: {str(e)}") import traceback traceback.print_exc() error_response = ChatResponse( session_id=session_id, text=f"❌ Debug agent error: {str(e)}", finished=True, type="message", format="text", ext=[{"type": "error", "data": {"error": str(e)}}] ) await response.write(json.dumps(error_response).encode() + b"\n") await response.write_eof() return response @server.PromptServer.instance.routes.post("/api/update-workflow-ui") async def update_workflow_ui(request): """ Update workflow_data_ui field for a specific checkpoint without affecting other fields """ log.info("Received update-workflow-ui request") req_json = await request.json() try: checkpoint_id = req_json.get('checkpoint_id') workflow_data_ui = req_json.get('workflow_data_ui') if not checkpoint_id or not workflow_data_ui: return web.json_response({ "success": False, "message": "Missing required parameters: checkpoint_id and workflow_data_ui" }) try: checkpoint_id = int(checkpoint_id) except ValueError: return web.json_response({ "success": False, "message": "Invalid checkpoint_id format" }) # Update only the workflow_data_ui field success = update_workflow_ui_by_id(checkpoint_id, workflow_data_ui) if success: log.info(f"Successfully updated workflow_data_ui for checkpoint ID: {checkpoint_id}") return web.json_response({ "success": True, "message": f"Workflow UI data updated successfully for checkpoint {checkpoint_id}" }) else: return web.json_response({ "success": False, "message": f"Checkpoint {checkpoint_id} not found" }) except Exception as e: log.error(f"Error updating workflow UI data: {str(e)}") return web.json_response({ "success": False, "message": f"Failed to update workflow UI data: {str(e)}" }) @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)}" })