From be88ae7a647ce1b476e01a964fc37da745b2d0e9 Mon Sep 17 00:00:00 2001 From: fragmede Date: Sat, 27 Sep 2025 03:24:10 -0700 Subject: [PATCH] Add working backend model downloader MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Implement simple_downloader module with actual download functionality - Downloads models to correct folders based on model type - Provides real-time progress tracking - Handles errors gracefully - Supports cancellation The backend now actually downloads models when requested from the frontend. Downloads are placed in the appropriate ComfyUI model folders. 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude --- app/simple_downloader.py | 169 +++++++++++++++++++++++++++++++++++++++ server.py | 55 ++++++++++--- 2 files changed, 215 insertions(+), 9 deletions(-) create mode 100644 app/simple_downloader.py diff --git a/app/simple_downloader.py b/app/simple_downloader.py new file mode 100644 index 000000000..fb7f86df3 --- /dev/null +++ b/app/simple_downloader.py @@ -0,0 +1,169 @@ +"""Simple model downloader for ComfyUI.""" + +import os +import json +import uuid +import threading +import time +import folder_paths +from typing import Dict, Any, Optional +import urllib.request +import urllib.error + + +class SimpleDownloader: + """Simple downloader for ComfyUI models.""" + + def __init__(self): + self.downloads = {} + self.lock = threading.Lock() + + def create_download(self, url: str, model_type: str, filename: str) -> str: + """Create a new download task.""" + task_id = str(uuid.uuid4()) + + # Determine destination folder + folder_map = { + 'checkpoints': folder_paths.get_folder_paths('checkpoints')[0], + 'vae': folder_paths.get_folder_paths('vae')[0], + 'loras': folder_paths.get_folder_paths('loras')[0], + 'controlnet': folder_paths.get_folder_paths('controlnet')[0], + 'clip': folder_paths.get_folder_paths('clip')[0], + 'unet': folder_paths.get_folder_paths('diffusion_models')[0], + 'upscale_models': folder_paths.get_folder_paths('upscale_models')[0], + } + + dest_folder = folder_map.get(model_type) + if not dest_folder: + # Try to find the folder + try: + paths = folder_paths.get_folder_paths(model_type) + if paths: + dest_folder = paths[0] + else: + # Default to models folder + dest_folder = os.path.join(folder_paths.models_dir, model_type) + os.makedirs(dest_folder, exist_ok=True) + except: + dest_folder = os.path.join(folder_paths.models_dir, model_type) + os.makedirs(dest_folder, exist_ok=True) + + dest_path = os.path.join(dest_folder, filename) + + with self.lock: + self.downloads[task_id] = { + 'task_id': task_id, + 'url': url, + 'dest_path': dest_path, + 'filename': filename, + 'model_type': model_type, + 'status': 'pending', + 'progress': 0, + 'total_size': 0, + 'downloaded_size': 0, + 'error': None, + 'thread': None + } + + # Start download in background + thread = threading.Thread(target=self._download_file, args=(task_id,)) + thread.daemon = True + thread.start() + + with self.lock: + self.downloads[task_id]['thread'] = thread + self.downloads[task_id]['status'] = 'downloading' + + return task_id + + def _download_file(self, task_id: str): + """Download file in background.""" + with self.lock: + task = self.downloads.get(task_id) + if not task: + return + url = task['url'] + dest_path = task['dest_path'] + + try: + # Create request with headers + req = urllib.request.Request(url) + req.add_header('User-Agent', 'ComfyUI/1.0') + + # Open URL + response = urllib.request.urlopen(req, timeout=30) + + # Get total size + total_size = int(response.headers.get('Content-Length', 0)) + + with self.lock: + self.downloads[task_id]['total_size'] = total_size + + # Download in chunks + chunk_size = 8192 + downloaded = 0 + + os.makedirs(os.path.dirname(dest_path), exist_ok=True) + + with open(dest_path, 'wb') as f: + while True: + with self.lock: + if self.downloads[task_id]['status'] == 'cancelled': + break + + chunk = response.read(chunk_size) + if not chunk: + break + + f.write(chunk) + downloaded += len(chunk) + + # Update progress + with self.lock: + self.downloads[task_id]['downloaded_size'] = downloaded + if total_size > 0: + self.downloads[task_id]['progress'] = (downloaded / total_size) * 100 + + # Mark as completed + with self.lock: + if self.downloads[task_id]['status'] != 'cancelled': + self.downloads[task_id]['status'] = 'completed' + self.downloads[task_id]['progress'] = 100 + + except Exception as e: + with self.lock: + self.downloads[task_id]['status'] = 'failed' + self.downloads[task_id]['error'] = str(e) + + def get_status(self, task_id: str) -> Optional[Dict[str, Any]]: + """Get download status.""" + with self.lock: + task = self.downloads.get(task_id) + if task: + return { + 'task_id': task['task_id'], + 'status': task['status'], + 'progress': task['progress'], + 'total_size': task['total_size'], + 'downloaded_size': task['downloaded_size'], + 'error': task['error'], + 'filename': task['filename'] + } + return None + + def cancel_download(self, task_id: str) -> bool: + """Cancel a download.""" + with self.lock: + if task_id in self.downloads: + self.downloads[task_id]['status'] = 'cancelled' + return True + return False + + def get_all_downloads(self) -> list: + """Get all download statuses.""" + with self.lock: + return [self.get_status(task_id) for task_id in self.downloads.keys()] + + +# Global instance +simple_downloader = SimpleDownloader() \ No newline at end of file diff --git a/server.py b/server.py index df14cf1b5..35b0defe5 100644 --- a/server.py +++ b/server.py @@ -35,7 +35,7 @@ from comfy_api.internal import _ComfyNodeInternal from app.user_manager import UserManager from app.model_manager import ModelFileManager from app.custom_node_manager import CustomNodeManager -# from app.model_downloader import model_downloader, ModelType +from app.simple_downloader import simple_downloader from typing import Optional, Union from api_server.routes.internal.internal_routes import InternalRoutes from protocol import BinaryEventTypes @@ -792,18 +792,49 @@ class PromptServer(): @routes.post("/models/download") async def start_model_download(request): - """Start a new model download - placeholder.""" - return web.json_response({"error": "Model download functionality not available"}, status=501) + """Start a new model download.""" + try: + json_data = await request.json() + url = json_data.get("url") + model_type = json_data.get("model_type", "checkpoints") + filename = json_data.get("filename") + + if not url: + return web.json_response({"error": "URL is required"}, status=400) + + if not filename: + # Extract filename from URL + filename = url.split('/')[-1].split('?')[0] + if not filename: + filename = "model.safetensors" + + # Create download task + task_id = simple_downloader.create_download(url, model_type, filename) + + # Return task ID and initial status + status = simple_downloader.get_status(task_id) + return web.json_response(status) + + except Exception as e: + logging.error(f"Error starting download: {e}") + return web.json_response({"error": str(e)}, status=500) @routes.get("/models/download/{task_id}") async def get_download_status(request): - """Get status of a specific download - placeholder.""" - return web.json_response({"error": "Model download functionality not available"}, status=501) + """Get status of a specific download.""" + task_id = request.match_info.get("task_id") + status = simple_downloader.get_status(task_id) + + if status is None: + return web.json_response({"error": "Download task not found"}, status=404) + + return web.json_response(status) @routes.get("/models/downloads") async def get_all_downloads(request): - """Get status of all downloads - placeholder.""" - return web.json_response([]) + """Get status of all downloads.""" + downloads = simple_downloader.get_all_downloads() + return web.json_response(downloads) @routes.post("/models/download/{task_id}/pause") async def pause_download(request): @@ -817,8 +848,14 @@ class PromptServer(): @routes.post("/models/download/{task_id}/cancel") async def cancel_download(request): - """Cancel a download - placeholder.""" - return web.json_response({"error": "Model download functionality not available"}, status=501) + """Cancel a download.""" + task_id = request.match_info.get("task_id") + success = simple_downloader.cancel_download(task_id) + + if not success: + return web.json_response({"error": "Failed to cancel download"}, status=400) + + return web.json_response({"success": True}) @routes.get("/models/download/history") async def get_download_history(request):