Add working backend model downloader

- 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 <noreply@anthropic.com>
This commit is contained in:
fragmede 2025-09-27 03:24:10 -07:00
parent 4c1d4196eb
commit be88ae7a64
No known key found for this signature in database
2 changed files with 215 additions and 9 deletions

169
app/simple_downloader.py Normal file
View File

@ -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()

View File

@ -35,7 +35,7 @@ from comfy_api.internal import _ComfyNodeInternal
from app.user_manager import UserManager from app.user_manager import UserManager
from app.model_manager import ModelFileManager from app.model_manager import ModelFileManager
from app.custom_node_manager import CustomNodeManager 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 typing import Optional, Union
from api_server.routes.internal.internal_routes import InternalRoutes from api_server.routes.internal.internal_routes import InternalRoutes
from protocol import BinaryEventTypes from protocol import BinaryEventTypes
@ -792,18 +792,49 @@ class PromptServer():
@routes.post("/models/download") @routes.post("/models/download")
async def start_model_download(request): async def start_model_download(request):
"""Start a new model download - placeholder.""" """Start a new model download."""
return web.json_response({"error": "Model download functionality not available"}, status=501) 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}") @routes.get("/models/download/{task_id}")
async def get_download_status(request): async def get_download_status(request):
"""Get status of a specific download - placeholder.""" """Get status of a specific download."""
return web.json_response({"error": "Model download functionality not available"}, status=501) 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") @routes.get("/models/downloads")
async def get_all_downloads(request): async def get_all_downloads(request):
"""Get status of all downloads - placeholder.""" """Get status of all downloads."""
return web.json_response([]) downloads = simple_downloader.get_all_downloads()
return web.json_response(downloads)
@routes.post("/models/download/{task_id}/pause") @routes.post("/models/download/{task_id}/pause")
async def pause_download(request): async def pause_download(request):
@ -817,8 +848,14 @@ class PromptServer():
@routes.post("/models/download/{task_id}/cancel") @routes.post("/models/download/{task_id}/cancel")
async def cancel_download(request): async def cancel_download(request):
"""Cancel a download - placeholder.""" """Cancel a download."""
return web.json_response({"error": "Model download functionality not available"}, status=501) 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") @routes.get("/models/download/history")
async def get_download_history(request): async def get_download_history(request):