mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-02-16 16:32:34 +08:00
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
parent
4c1d4196eb
commit
be88ae7a64
169
app/simple_downloader.py
Normal file
169
app/simple_downloader.py
Normal 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()
|
||||
55
server.py
55
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):
|
||||
|
||||
Loading…
Reference in New Issue
Block a user