diff --git a/__init__.py b/__init__.py index cb07449..2078d31 100644 --- a/__init__.py +++ b/__init__.py @@ -14,10 +14,13 @@ from server import PromptServer active_downloads = {} # Download control (for pause/resume) download_control = {} +# Download queue management +download_queue = [] +current_download_task = None # Only one download at a time -# Configuration -CHUNK_SIZE = 10 * 1024 * 1024 # 10MB chunks for faster downloads -NUM_CONNECTIONS = 4 # Number of parallel connections per file +# Configuration optimized for datacenter connections (RunPod) +CHUNK_SIZE = 32 * 1024 * 1024 # 32MB chunks - balanced for 500MB to 30GB+ files +NUM_CONNECTIONS = 8 # 8 parallel connections - optimal for DC bandwidth @PromptServer.instance.routes.post("/server_download/start") @@ -56,7 +59,7 @@ async def start_download(request): # Create directory if it doesn't exist os.makedirs(os.path.dirname(output_path), exist_ok=True) - # Mark as downloading + # Mark as queued download_id = f"{save_path}/{filename}" active_downloads[download_id] = { "url": url, @@ -64,17 +67,24 @@ async def start_download(request): "save_path": save_path, "output_path": output_path, "progress": 0, - "status": "downloading" + "status": "queued", + "priority": None } - # Start download in background - import asyncio - asyncio.create_task(download_file(url, output_path, download_id)) + # Add to queue + download_queue.append({ + "download_id": download_id, + "url": url, + "output_path": output_path + }) + + # Process queue (will start download if slot available) + asyncio.create_task(process_download_queue()) return web.json_response({ "success": True, "download_id": download_id, - "message": "Download started" + "message": "Download queued" }) except Exception as e: @@ -85,6 +95,58 @@ async def start_download(request): ) +async def process_download_queue(): + """Process the download queue - one download at a time""" + global download_queue, current_download_task + + # Check if already downloading + if current_download_task is not None and not current_download_task.done(): + logging.info("[RunpodDirect] Download already in progress, waiting...") + return # Already downloading + + if len(download_queue) == 0: + logging.info("[RunpodDirect] Queue is empty") + return # Nothing to process + + # Get next download from queue + download_item = download_queue.pop(0) + download_id = download_item["download_id"] + url = download_item["url"] + output_path = download_item["output_path"] + + # Set status to downloading + active_downloads[download_id]["status"] = "downloading" + active_downloads[download_id]["progress"] = 0 + active_downloads[download_id]["downloaded"] = 0 + + logging.info(f"[RunpodDirect] Starting download {download_id} with {NUM_CONNECTIONS} connections (full speed)") + + # Notify frontend that download is starting + await PromptServer.instance.send("server_download_progress", { + "download_id": download_id, + "progress": 0, + "downloaded": 0, + "total": 0 + }) + + # Start download task + current_download_task = asyncio.create_task(download_file(url, output_path, download_id)) + + # Add completion callback to process next in queue + current_download_task.add_done_callback(lambda t: on_download_complete(download_id)) + + +def on_download_complete(download_id): + """Called when a download completes - processes next in queue""" + global current_download_task + + current_download_task = None + logging.info(f"[RunpodDirect] Download completed: {download_id}, processing next in queue...") + + # Process next in queue + asyncio.create_task(process_download_queue()) + + async def download_chunk(session, url, start, end, output_path, chunk_index, download_id): """Download a specific chunk of the file""" headers = {'Range': f'bytes={start}-{end}'} @@ -111,6 +173,8 @@ async def download_file(url, output_path, download_id): """Download file with multi-connection support and progress tracking""" import aiohttp + logging.info(f"[RunpodDirect] Download {download_id} using {NUM_CONNECTIONS} connections (full speed)") + try: # Initialize control for this download download_control[download_id] = { @@ -411,6 +475,8 @@ async def resume_download(request): @PromptServer.instance.routes.post("/server_download/cancel") async def cancel_download(request): """Cancel an active download""" + global download_queue, current_download_task + try: json_data = await request.json() download_id = json_data.get("download_id") @@ -421,14 +487,16 @@ async def cancel_download(request): status=400 ) - if download_id not in download_control: - return web.json_response( - {"error": "Download not found or already completed"}, - status=404 - ) + # Check if download is queued (not started yet) + download_queue[:] = [d for d in download_queue if d["download_id"] != download_id] - download_control[download_id]["cancelled"] = True - active_downloads[download_id]["status"] = "cancelled" + # Check if download is active + if download_id in download_control: + download_control[download_id]["cancelled"] = True + + # Update status + if download_id in active_downloads: + active_downloads[download_id]["status"] = "cancelled" await PromptServer.instance.send("server_download_cancelled", { "download_id": download_id @@ -459,7 +527,7 @@ async def serve_js_with_version(request): WEB_DIRECTORY = "./web" # Version for cache busting - increment this when you update the JS -__version__ = "1.0.2" +__version__ = "1.0.3" NODE_CLASS_MAPPINGS = {} NODE_DISPLAY_NAME_MAPPINGS = {} diff --git a/pyproject.toml b/pyproject.toml index 144d9d4..efaf265 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "comfyui-runpoddirect" description = "Server-side model download extension for ComfyUI on Runpod - Download missing models directly to your pod with multi-connection downloads, queue management, and real-time progress tracking" -version = "1.0.2" +version = "1.0.3" license = {file = "LICENSE"} authors = [ {name = "Madiator2011"}, @@ -32,5 +32,4 @@ Documentation = "https://github.com/MadiatorLabs/ComfyUI-RunpodDirect#readme" PublisherId = "madiator2011" DisplayName = "Madiator2011" Icon = "" -includes = [] - +includes = [] \ No newline at end of file diff --git a/web/serverDownload.js b/web/serverDownload.js index c29f6f0..8b3c57d 100644 --- a/web/serverDownload.js +++ b/web/serverDownload.js @@ -2,8 +2,8 @@ import { app } from "../../scripts/app.js"; import { api } from "../../scripts/api.js"; // ComfyUI RunpodDirect Extension -// Version: 1.0.1 -console.log('[RunpodDirect] v1.0.1'); +// Version: 1.0.3 +console.log('[RunpodDirect] v1.0.3'); // Track download states const downloadStates = new Map(); @@ -60,6 +60,13 @@ api.addEventListener("server_download_progress", ({ detail }) => { api.addEventListener("server_download_complete", ({ detail }) => { const { download_id, path, size } = detail; + + // Increment counter BEFORE updating UI + if (isDownloadingAll) { + completedDownloads++; + console.log(`[RunpodDirect] Progress: ${completedDownloads}/${totalDownloads} completed`); + } + downloadStates.set(download_id, { status: 'completed', progress: 100, @@ -73,15 +80,23 @@ api.addEventListener("server_download_complete", ({ detail }) => { console.log(`Download completed: ${download_id} -> ${path}`); - // Process next in queue if downloading all - if (isDownloadingAll) { - completedDownloads++; - processDownloadQueue(); + // Check if all downloads are done + if (isDownloadingAll && completedDownloads >= totalDownloads) { + console.log('[RunpodDirect] All downloads completed!'); + isDownloadingAll = false; + showRefreshPrompt(); } }); api.addEventListener("server_download_error", ({ detail }) => { const { download_id, error } = detail; + + // Increment counter BEFORE updating UI + if (isDownloadingAll) { + completedDownloads++; + console.log(`[RunpodDirect] Progress: ${completedDownloads}/${totalDownloads} completed (1 error)`); + } + downloadStates.set(download_id, { status: 'error', error @@ -93,16 +108,31 @@ api.addEventListener("server_download_error", ({ detail }) => { console.error(`Download error: ${download_id} - ${error}`); - // Continue with next download even if one fails - if (isDownloadingAll) { - completedDownloads++; - processDownloadQueue(); + // Check if all downloads are done (including failed ones) + if (isDownloadingAll && completedDownloads >= totalDownloads) { + console.log('[RunpodDirect] All downloads completed!'); + isDownloadingAll = false; + showRefreshPrompt(); } }); // Function to start a server download -async function startServerDownload(url, savePath, filename) { +async function startServerDownload(url, savePath, filename, markAsQueued = false) { try { + const download_id = `${savePath}/${filename}`; + + // Mark as queued immediately if requested (for Download All) + if (markAsQueued) { + downloadStates.set(download_id, { + status: 'queued', + progress: 0 + }); + + window.dispatchEvent(new CustomEvent('serverDownloadUpdate', { + detail: { download_id, ...downloadStates.get(download_id) } + })); + } + const response = await api.fetchApi("/server_download/start", { method: "POST", headers: { @@ -118,15 +148,17 @@ async function startServerDownload(url, savePath, filename) { const result = await response.json(); if (response.ok) { - const download_id = result.download_id; - downloadStates.set(download_id, { - status: 'downloading', - progress: 0 - }); + // If not already marked as queued, set as queued now + if (!markAsQueued) { + downloadStates.set(download_id, { + status: 'queued', + progress: 0 + }); - window.dispatchEvent(new CustomEvent('serverDownloadUpdate', { - detail: { download_id, ...downloadStates.get(download_id) } - })); + window.dispatchEvent(new CustomEvent('serverDownloadUpdate', { + detail: { download_id, ...downloadStates.get(download_id) } + })); + } return { success: true, download_id }; } else { @@ -194,24 +226,27 @@ async function cancelDownload(downloadId) { } } -// Process download queue +// Process download queue - Sends all downloads to backend which handles queue management async function processDownloadQueue() { if (downloadQueue.length === 0) { - // All downloads complete - isDownloadingAll = false; - console.log('[RunpodDirect] All downloads completed!'); - - // Show refresh prompt - if (completedDownloads > 0) { - showRefreshPrompt(); - } + console.log('[RunpodDirect] No downloads in queue'); return; } - const nextDownload = downloadQueue.shift(); - console.log(`[RunpodDirect] Starting download ${completedDownloads + 1}/${totalDownloads}: ${nextDownload.filename}`); + // Send all downloads to the backend (backend handles queue and priorities) + console.log(`[RunpodDirect] Starting ${downloadQueue.length} downloads (backend will queue and prioritize)`); - await startServerDownload(nextDownload.url, nextDownload.directory, nextDownload.filename); + const downloadsToStart = [...downloadQueue]; + downloadQueue = []; // Clear queue as we're sending all to backend + + // Start all downloads - backend will queue and manage priorities + // Pass markAsQueued=true so buttons show "Queued" status immediately + for (const download of downloadsToStart) { + console.log(`[RunpodDirect] Queuing download ${download.filename}`); + await startServerDownload(download.url, download.directory, download.filename, true); + } + + console.log(`[RunpodDirect] All ${downloadsToStart.length} downloads queued on backend`); } // Show refresh prompt @@ -223,8 +258,15 @@ function showRefreshPrompt() { const dialogContent = dialog.querySelector('.p-dialog-content'); if (!dialogContent) return; + // Check if prompt already exists + if (document.querySelector('.server-download-refresh-prompt')) { + console.log('[RunpodDirect] Refresh prompt already shown'); + return; + } + // Create refresh prompt const refreshPrompt = document.createElement('div'); + refreshPrompt.className = 'server-download-refresh-prompt'; refreshPrompt.style.cssText = ` margin-top: 20px; padding: 16px; @@ -249,7 +291,7 @@ function showRefreshPrompt() { dialogContent.appendChild(refreshPrompt); } -// Create global progress area +// Create global progress area with individual progress bars for each download function createProgressArea(listbox) { // Remove existing progress area if any const existing = document.querySelector('.server-download-progress-area'); @@ -270,135 +312,134 @@ function createProgressArea(listbox) {