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) {
Download Progress
-
- - - -
-
Waiting to start...
-
Overall: 0/0 models completed
-
-
+
+ Overall: 0/${totalDownloads} models completed
-
- Speed: -- - Size: -- +
+
`; listbox.parentElement.appendChild(progressArea); - // Initialize progress bar to 0% - const progressBarInit = document.getElementById('server-download-progress-bar'); - if (progressBarInit) { - progressBarInit.style.width = '0%'; - } - - // Setup button handlers - let currentDownloadId = null; - - const pauseBtn = document.getElementById('server-download-pause-btn'); - const resumeBtn = document.getElementById('server-download-resume-btn'); - const cancelBtn = document.getElementById('server-download-cancel-btn'); - - if (pauseBtn) { - pauseBtn.onclick = async () => { - if (currentDownloadId) { - await pauseDownload(currentDownloadId); - } - }; - } - - if (resumeBtn) { - resumeBtn.onclick = async () => { - if (currentDownloadId) { - await resumeDownload(currentDownloadId); - } - }; - } - - if (cancelBtn) { - cancelBtn.onclick = async () => { - if (currentDownloadId && confirm('Are you sure you want to cancel this download?')) { - await cancelDownload(currentDownloadId); - // Clear queue - downloadQueue = []; - isDownloadingAll = false; - } - }; - } - // Listen for updates const updateHandler = (event) => { const { download_id, status, progress, downloaded, total, speed } = event.detail; - if (!isDownloadingAll) return; - - currentDownloadId = download_id; - - const currentFile = document.getElementById('server-download-current-file'); - const overallProgress = document.getElementById('server-download-overall-progress'); - const progressBar = document.getElementById('server-download-progress-bar'); - const speedInfo = document.getElementById('server-download-speed'); - const sizeInfo = document.getElementById('server-download-size'); - - if (currentFile) { - currentFile.textContent = `Current: ${download_id}`; + if (!isDownloadingAll) { + return; } + // Update overall progress + const overallProgress = document.getElementById('server-download-overall-progress'); if (overallProgress) { overallProgress.textContent = `Overall: ${completedDownloads}/${totalDownloads} models completed`; } - if (progressBar) { - if (status === 'downloading' && progress !== undefined) { - progressBar.style.width = `${progress}%`; - } else if (status === 'completed') { - progressBar.style.width = '100%'; - } else if (status === 'error' || status === 'cancelled') { - progressBar.style.width = '0%'; - } - } - - if (speedInfo && speed) { - speedInfo.textContent = `Speed: ${speed}`; - } - - if (sizeInfo && downloaded && total) { - sizeInfo.textContent = `Size: ${formatBytes(downloaded)} / ${formatBytes(total)}`; - } - - // Update button visibility based on status - if (pauseBtn && resumeBtn && cancelBtn) { - if (status === 'downloading') { - pauseBtn.style.display = 'inline-flex'; - resumeBtn.style.display = 'none'; - cancelBtn.style.display = 'inline-flex'; - } else if (status === 'paused') { - pauseBtn.style.display = 'none'; - resumeBtn.style.display = 'inline-flex'; - cancelBtn.style.display = 'inline-flex'; - } else { - pauseBtn.style.display = 'none'; - resumeBtn.style.display = 'none'; - cancelBtn.style.display = 'none'; - } - } + // Update or create individual progress item + updateDownloadProgressItem(download_id, status, progress, downloaded, total, speed); }; window.addEventListener('serverDownloadUpdate', updateHandler); } +// Update or create a progress item for a specific download +function updateDownloadProgressItem(download_id, status, progress, downloaded, total, speed) { + // Declare variables at function scope so they're accessible across try blocks + let item = null; + let container = null; + const itemId = `download-item-${download_id.replace(/\//g, '-')}`; + + try { + container = document.getElementById('server-download-items-container'); + if (!container) return; + + item = document.getElementById(itemId); + + // Don't show queued items + if (status === 'queued') { + if (item) item.remove(); + return; + } + + // Remove completed/error items after a delay + if (status === 'completed' || status === 'error') { + if (item && !item.dataset.removing) { + item.dataset.removing = 'true'; + setTimeout(() => { + try { + if (item && item.parentNode) item.remove(); + } catch (e) { + console.error('[RunpodDirect] Error removing progress item:', e); + } + }, 2000); + } + } + + // Create new item if it doesn't exist + if (!item) { + item = document.createElement('div'); + item.id = itemId; + item.style.cssText = ` + padding: 12px; + background: rgba(255, 255, 255, 0.05); + border-radius: 6px; + border: 1px solid rgba(255, 255, 255, 0.1); + `; + container.appendChild(item); + } + } catch (e) { + console.error('[RunpodDirect] Error in updateDownloadProgressItem:', e); + return; + } + + try { + // Status icon and color (no priority badge needed - downloading one at a time) + let statusIcon = ''; + let statusColor = '#2196F3'; + if (status === 'downloading') { + statusIcon = ''; + statusColor = '#2196F3'; + } else if (status === 'completed') { + statusIcon = ''; + statusColor = '#4CAF50'; + } else if (status === 'error') { + statusIcon = ''; + statusColor = '#ef4444'; + } else if (status === 'paused') { + statusIcon = ''; + statusColor = '#FF9800'; + } + + const progressPercent = progress || 0; + const speedText = speed || '--'; + const sizeText = downloaded && total ? `${formatBytes(downloaded)} / ${formatBytes(total)}` : '--'; + + if (item) { + item.innerHTML = ` +
+
+ ${statusIcon}${download_id} +
+
+ ${progressPercent.toFixed(1)}% +
+
+
+
+
+
+ Speed: ${speedText} + ${sizeText} +
+ `; + } + } catch (e) { + console.error('[RunpodDirect] Error updating progress item HTML:', e); + } +} + // Export functions for use in other modules window.serverDownload = { start: startServerDownload, @@ -602,9 +643,10 @@ function injectServerDownloadButtons() { const result = await startServerDownload(url, directory, filename); if (result.success) { - btnLabel.textContent = 'Downloading'; - statusIcon.className = 'pi pi-spin pi-spinner'; + btnLabel.textContent = 'Queued'; + statusIcon.className = 'pi pi-clock'; statusIcon.style.display = 'inline'; + statusIcon.style.color = '#FF9800'; } else { btnLabel.textContent = 'Error'; statusIcon.className = 'pi pi-times-circle'; @@ -619,7 +661,14 @@ function injectServerDownloadButtons() { if (event.detail.download_id === download_id) { const { status, error } = event.detail; - if (status === 'downloading') { + if (status === 'queued') { + serverDownloadBtn.disabled = true; + btnLabel.textContent = 'Queued'; + statusIcon.className = 'pi pi-clock'; + statusIcon.style.display = 'inline'; + statusIcon.style.color = '#FF9800'; + } else if (status === 'downloading') { + serverDownloadBtn.disabled = true; btnLabel.textContent = 'Downloading'; statusIcon.className = 'pi pi-spin pi-spinner'; statusIcon.style.display = 'inline';