diff --git a/__init__.py b/__init__.py index 484a287..7cb9a21 100644 --- a/__init__.py +++ b/__init__.py @@ -31,6 +31,7 @@ async def start_download(request): url = json_data.get("url") save_path = json_data.get("save_path") # e.g., "checkpoints" filename = json_data.get("filename") # e.g., "model.safetensors" + token = json_data.get("token") # Optional HF token for gated models if not url or not save_path or not filename: return web.json_response( @@ -103,11 +104,15 @@ async def start_download(request): "priority": None } + # Resolve token: prefer explicit token, fall back to HF_TOKEN env var + resolved_token = token or os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") + # Add to queue download_queue.append({ "download_id": download_id, "url": url, - "output_path": output_path + "output_path": output_path, + "token": resolved_token }) # Process queue (will start download if slot available) @@ -145,6 +150,7 @@ async def process_download_queue(): download_id = download_item["download_id"] url = download_item["url"] output_path = download_item["output_path"] + token = download_item.get("token") # Set status to downloading active_downloads[download_id]["status"] = "downloading" @@ -162,7 +168,7 @@ async def process_download_queue(): }) # Start download task - current_download_task = asyncio.create_task(download_file(url, output_path, download_id)) + current_download_task = asyncio.create_task(download_file(url, output_path, download_id, token=token)) # Add completion callback to process next in queue current_download_task.add_done_callback(lambda t: on_download_complete(download_id)) @@ -201,7 +207,7 @@ async def download_chunk(session, url, start, end, output_path, chunk_index, dow return None -async def download_file(url, output_path, download_id): +async def download_file(url, output_path, download_id, token=None): """Download file with multi-connection support and progress tracking""" import aiohttp @@ -216,19 +222,44 @@ async def download_file(url, output_path, download_id): "lock": asyncio.Lock() # Lock for thread-safe updates } + # Build auth headers for gated HF models + auth_headers = {} + if token and 'huggingface.co' in url: + auth_headers['Authorization'] = f'Bearer {token}' + logging.info(f"[RunpodDirect] Using HF token for {download_id}") + timeout = aiohttp.ClientTimeout(total=None) - async with aiohttp.ClientSession(timeout=timeout) as session: + async with aiohttp.ClientSession(timeout=timeout, headers=auth_headers) as session: # Get file size - try HEAD first, then fall back to GET with Range total_size = 0 supports_range = False + # Helper to detect gated/auth errors and produce a useful message + def _check_gated_error(status_code, request_url): + if status_code in (401, 403, 451) and 'huggingface.co' in request_url: + # Convert download URL to repo URL for the user + # e.g. https://huggingface.co/org/repo/resolve/main/file.ext -> https://huggingface.co/org/repo + repo_url = request_url + resolve_idx = request_url.find('/resolve/') + if resolve_idx != -1: + repo_url = request_url[:resolve_idx] + if status_code == 401: + raise Exception(f"Authentication required. Provide a valid HF token. Repo: {repo_url}") + elif status_code == 403: + raise Exception(f"Access denied — you may need to accept the model's terms at {repo_url}") + elif status_code == 451: + raise Exception(f"Model is restricted. Accept the license agreement at {repo_url}") + try: # Try HEAD request first async with session.head(url, allow_redirects=True) as response: + _check_gated_error(response.status, url) if response.status == 200: total_size = int(response.headers.get('content-length', 0)) supports_range = response.headers.get('accept-ranges') == 'bytes' except Exception as e: + if 'Accept the license' in str(e) or 'Access denied' in str(e) or 'Authentication required' in str(e): + raise logging.warning(f"HEAD request failed for {download_id}: {e}") # If HEAD didn't give us the size, try GET with Range header @@ -237,6 +268,7 @@ async def download_file(url, output_path, download_id): try: headers = {'Range': 'bytes=0-0'} async with session.get(url, headers=headers, allow_redirects=True) as response: + _check_gated_error(response.status, url) if response.status in [200, 206]: # Try to get size from Content-Range header first content_range = response.headers.get('content-range', '') @@ -251,6 +283,8 @@ async def download_file(url, output_path, download_id): if total_size == 0: total_size = int(response.headers.get('content-length', 0)) except Exception as e: + if 'Accept the license' in str(e) or 'Access denied' in str(e) or 'Authentication required' in str(e): + raise logging.warning(f"GET with Range failed for {download_id}: {e}") if total_size == 0: @@ -340,6 +374,9 @@ async def download_chunk_with_progress(session, url, start, end, output_path, ch try: async with session.get(url, headers=headers) as response: + if response.status in (401, 403, 451) and 'huggingface.co' in url: + repo_url = url[:url.find('/resolve/')] if '/resolve/' in url else url + raise Exception(f"Access denied — accept the model's terms at {repo_url}") if response.status not in [200, 206]: raise Exception(f"HTTP {response.status} for chunk {chunk_index}") @@ -391,6 +428,9 @@ async def download_single_connection(session, url, output_path, download_id, tot downloaded_size = 0 async with session.get(url) as response: + if response.status in (401, 403, 451) and 'huggingface.co' in url: + repo_url = url[:url.find('/resolve/')] if '/resolve/' in url else url + raise Exception(f"Access denied — accept the model's terms at {repo_url}") if response.status != 200: raise Exception(f"HTTP {response.status}") @@ -540,6 +580,77 @@ async def cancel_download(request): return web.json_response({"error": str(e)}, status=500) +@PromptServer.instance.routes.get("/server_download/hf_token_status") +async def hf_token_status(request): + """Check if HF_TOKEN environment variable is set (without exposing the value)""" + has_token = bool(os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")) + return web.json_response({"has_token": has_token}) + + +@PromptServer.instance.routes.post("/server_download/validate_hf_token") +async def validate_hf_token(request): + """Validate a Hugging Face token and optionally check access to specific model URLs""" + import aiohttp + + try: + json_data = await request.json() + token = json_data.get("token", "") + urls = json_data.get("urls", []) # Optional: list of model URLs to check access + + # If token is the sentinel '__env__', use the environment variable + if token == "__env__": + token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") or "" + + if not token: + return web.json_response({"valid": False, "error": "No token provided"}, status=400) + + # Validate token against HF API + timeout = aiohttp.ClientTimeout(total=10) + async with aiohttp.ClientSession(timeout=timeout) as session: + # Check token validity via whoami endpoint + async with session.get( + "https://huggingface.co/api/whoami-v2", + headers={"Authorization": f"Bearer {token}"} + ) as response: + if response.status != 200: + return web.json_response({"valid": False, "error": "Invalid token"}) + user_data = await response.json() + username = user_data.get("name", "unknown") + + # If URLs provided, check access to each + url_access = {} + for url in urls[:10]: # Limit to 10 URLs + try: + async with session.head( + url, + headers={"Authorization": f"Bearer {token}"}, + allow_redirects=True + ) as resp: + if resp.status == 200: + url_access[url] = {"accessible": True} + elif resp.status in (401, 403, 451): + repo_url = url[:url.find('/resolve/')] if '/resolve/' in url else url + url_access[url] = { + "accessible": False, + "reason": "terms_not_accepted", + "repo_url": repo_url + } + else: + url_access[url] = {"accessible": False, "reason": f"HTTP {resp.status}"} + except Exception: + url_access[url] = {"accessible": False, "reason": "request_failed"} + + return web.json_response({ + "valid": True, + "username": username, + "url_access": url_access + }) + + except Exception as e: + logging.error(f"Error validating HF token: {e}") + return web.json_response({"valid": False, "error": str(e)}, status=500) + + @PromptServer.instance.routes.get("/extensions/ComfyUI-RunpodDirect/serverDownload.js") async def serve_js_with_version(request): """Serve JS file with cache-busting headers""" @@ -559,7 +670,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.6" +__version__ = "1.0.7" NODE_CLASS_MAPPINGS = {} NODE_DISPLAY_NAME_MAPPINGS = {} diff --git a/pyproject.toml b/pyproject.toml index 76bef85..26ba56e 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.6" +version = "1.0.7" license = {file = "LICENSE"} authors = [ {name = "Madiator2011"}, diff --git a/web/serverDownload.js b/web/serverDownload.js index ba22573..d96981a 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.6 -console.log('[RunpodDirect] v1.0.6'); +// Version: 1.0.7 - Updated for ComfyUI v0.17.0+ (frontend v1.42+) +console.log('[RunpodDirect] v1.0.7'); // Track download states const downloadStates = new Map(); @@ -13,7 +13,48 @@ let completedDownloads = 0; let totalDownloads = 0; let downloadStartTimes = new Map(); -// Format bytes to human readable +// Session-only HF token (never persisted to disk) +let sessionHfToken = null; +let envHasHfToken = false; + +// Badge label -> folder_paths directory name mapping +const BADGE_TO_DIRECTORY = { + 'VAE': 'vae', + 'DIFFUSION': 'diffusion_models', + 'TEXT ENCODER': 'text_encoders', + 'LORA': 'loras', + 'CHECKPOINT': 'checkpoints', + 'CLIP': 'clip', + 'CLIP_VISION': 'clip_vision', + 'CONTROLNET': 'controlnet', + 'UPSCALE_MODELS': 'upscale_models', + 'LATENT_UPSCALE_MODELS': 'latent_upscale_models', + 'EMBEDDINGS': 'embeddings', + 'HYPERNETWORKS': 'hypernetworks', + 'STYLE_MODELS': 'style_models', + 'GLIGEN': 'gligen', + 'UNET': 'unet', +}; + +const THEME = { + // Status colors + primary: 'var(--primary-background)', + primaryHover:'var(--primary-background-hover)', + success: 'var(--success-background)', + error: 'var(--destructive-background)', + warning: 'var(--warning-background)', + // Text + foreground: 'var(--base-foreground)', + muted: 'var(--muted-foreground)', + // Backgrounds + baseBg: 'var(--base-background)', + secondaryBg: 'var(--secondary-background)', + secondaryBgHover: 'var(--secondary-background-hover)', + // Borders + border: 'var(--border-default)', + borderSubtle:'var(--border-subtle)', +}; + function formatBytes(bytes) { if (bytes === 0) return '0 B'; const k = 1024; @@ -22,37 +63,32 @@ function formatBytes(bytes) { return Math.round((bytes / Math.pow(k, i)) * 100) / 100 + ' ' + sizes[i]; } -// Calculate download speed function calculateSpeed(downloadId, downloaded) { const startTime = downloadStartTimes.get(downloadId); if (!startTime) return '0 MB/s'; - const elapsedSeconds = (Date.now() - startTime) / 1000; if (elapsedSeconds < 1) return 'Calculating...'; - const bytesPerSecond = downloaded / elapsedSeconds; return formatBytes(bytesPerSecond) + '/s'; } -// Listen for server download events +function statusColor(status) { + if (status === 'downloading') return THEME.primary; + if (status === 'completed') return THEME.success; + if (status === 'error') return THEME.error; + if (status === 'paused' || status === 'queued') return THEME.warning; + return THEME.primary; +} + +// --- WebSocket event listeners --- + api.addEventListener("server_download_progress", ({ detail }) => { const { download_id, progress, downloaded, total } = detail; - if (!downloadStartTimes.has(download_id)) { downloadStartTimes.set(download_id, Date.now()); } - const speed = calculateSpeed(download_id, downloaded); - - downloadStates.set(download_id, { - status: 'downloading', - progress, - downloaded, - total, - speed - }); - - // Trigger update event for UI components + downloadStates.set(download_id, { status: 'downloading', progress, downloaded, total, speed }); window.dispatchEvent(new CustomEvent('serverDownloadUpdate', { detail: { download_id, ...downloadStates.get(download_id) } })); @@ -60,106 +96,69 @@ 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, - path, - size - }); - + downloadStates.set(download_id, { status: 'completed', progress: 100, path, size }); window.dispatchEvent(new CustomEvent('serverDownloadUpdate', { detail: { download_id, ...downloadStates.get(download_id) } })); - console.log(`Download completed: ${download_id} -> ${path}`); - - // Check if all downloads are done if (isDownloadingAll && completedDownloads >= totalDownloads) { console.log('[RunpodDirect] All downloads completed!'); isDownloadingAll = false; - showRefreshPrompt(); + window.dispatchEvent(new CustomEvent('serverDownloadAllDone')); } }); 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 - }); - + downloadStates.set(download_id, { status: 'error', error }); window.dispatchEvent(new CustomEvent('serverDownloadUpdate', { detail: { download_id, ...downloadStates.get(download_id) } })); - console.error(`Download error: ${download_id} - ${error}`); - - // Check if all downloads are done (including failed ones) if (isDownloadingAll && completedDownloads >= totalDownloads) { console.log('[RunpodDirect] All downloads completed!'); isDownloadingAll = false; - showRefreshPrompt(); + window.dispatchEvent(new CustomEvent('serverDownloadAllDone')); } }); -// Function to start a server download +// --- API functions --- + 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 - }); - + downloadStates.set(download_id, { status: 'queued', progress: 0 }); window.dispatchEvent(new CustomEvent('serverDownloadUpdate', { detail: { download_id, ...downloadStates.get(download_id) } })); } - + const body = { url, save_path: savePath, filename }; + // Pass token for HF downloads (skip if using env var — backend reads it directly) + if (sessionHfToken && sessionHfToken !== '__env__' && url.includes('huggingface.co')) { + body.token = sessionHfToken; + } const response = await api.fetchApi("/server_download/start", { method: "POST", - headers: { - "Content-Type": "application/json", - }, - body: JSON.stringify({ - url, - save_path: savePath, - filename - }) + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(body) }); - const result = await response.json(); - if (response.ok) { - // If not already marked as queued, set as queued now if (!markAsQueued) { - downloadStates.set(download_id, { - status: 'queued', - progress: 0 - }); - + downloadStates.set(download_id, { status: 'queued', progress: 0 }); window.dispatchEvent(new CustomEvent('serverDownloadUpdate', { detail: { download_id, ...downloadStates.get(download_id) } })); } - return { success: true, download_id }; } else { return { success: false, error: result.error }; @@ -170,12 +169,10 @@ async function startServerDownload(url, savePath, filename, markAsQueued = false } } -// Get download status function getDownloadStatus(downloadId) { return downloadStates.get(downloadId) || null; } -// Pause download async function pauseDownload(downloadId) { try { const response = await api.fetchApi("/server_download/pause", { @@ -183,7 +180,6 @@ async function pauseDownload(downloadId) { headers: { "Content-Type": "application/json" }, body: JSON.stringify({ download_id: downloadId }) }); - const result = await response.json(); return { success: response.ok, ...result }; } catch (error) { @@ -192,7 +188,6 @@ async function pauseDownload(downloadId) { } } -// Resume download async function resumeDownload(downloadId) { try { const response = await api.fetchApi("/server_download/resume", { @@ -200,7 +195,6 @@ async function resumeDownload(downloadId) { headers: { "Content-Type": "application/json" }, body: JSON.stringify({ download_id: downloadId }) }); - const result = await response.json(); return { success: response.ok, ...result }; } catch (error) { @@ -209,7 +203,6 @@ async function resumeDownload(downloadId) { } } -// Cancel download async function cancelDownload(downloadId) { try { const response = await api.fetchApi("/server_download/cancel", { @@ -217,7 +210,6 @@ async function cancelDownload(downloadId) { headers: { "Content-Type": "application/json" }, body: JSON.stringify({ download_id: downloadId }) }); - const result = await response.json(); return { success: response.ok, ...result }; } catch (error) { @@ -226,499 +218,892 @@ async function cancelDownload(downloadId) { } } -// Process download queue - Sends all downloads to backend which handles queue management async function processDownloadQueue() { if (downloadQueue.length === 0) { console.log('[RunpodDirect] No downloads in queue'); return; } - - // Send all downloads to the backend (backend handles queue and priorities) - console.log(`[RunpodDirect] Starting ${downloadQueue.length} downloads (backend will queue and prioritize)`); - + console.log(`[RunpodDirect] Starting ${downloadQueue.length} downloads`); 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 + downloadQueue = []; 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 -function showRefreshPrompt() { - const dialog = document.querySelector('[role="dialog"]'); - if (!dialog) return; +// --- DOM helpers (safe, no innerHTML) --- - // Find the dialog content area - const dialogContent = dialog.querySelector('.p-dialog-content'); - if (!dialogContent) return; +function createEl(tag, styles, textContent) { + const el = document.createElement(tag); + if (styles) Object.assign(el.style, styles); + if (textContent) el.textContent = textContent; + return el; +} - // Check if prompt already exists - if (document.querySelector('.server-download-refresh-prompt')) { - console.log('[RunpodDirect] Refresh prompt already shown'); - return; +function extractModelsFromPinia() { + // Find the Vue app root element (mounted on #vue-app) + const rootEl = document.getElementById('vue-app') + || document.querySelector('[data-v-app]') + || document.getElementById('app'); + if (!rootEl) { + console.log('[RunpodDirect] No Vue root element found'); + return null; } - // Create refresh prompt - const refreshPrompt = document.createElement('div'); - refreshPrompt.className = 'server-download-refresh-prompt'; - refreshPrompt.style.cssText = ` - margin-top: 20px; - padding: 16px; - background: #4caf50; - color: white; - border-radius: 8px; - text-align: center; - font-weight: 500; - `; + const vueApp = rootEl.__vue_app__; + if (!vueApp) { + console.log('[RunpodDirect] No __vue_app__ on root element'); + return null; + } - refreshPrompt.innerHTML = ` -
- ✅ All models downloaded successfully! -
- - `; + // Access Pinia through Vue app's global properties + const pinia = vueApp.config.globalProperties.$pinia; + if (!pinia || !pinia._s) { + console.log('[RunpodDirect] No Pinia store registry found'); + return null; + } - dialogContent.appendChild(refreshPrompt); + // Get the dialog store (registered as 'dialog') + const dialogStore = pinia._s.get('dialog'); + if (!dialogStore) { + console.log('[RunpodDirect] Dialog store not found in Pinia'); + return null; + } + + // Find the missing models dialog in the stack + const stack = dialogStore.dialogStack; + if (!stack || !Array.isArray(stack)) { + console.log('[RunpodDirect] dialogStack not found or not an array'); + return null; + } + + const dialog = stack.find(d => d.key === 'global-missing-models-warning'); + if (!dialog) { + console.log('[RunpodDirect] Missing models dialog not found in stack'); + return null; + } + + const missingModels = dialog.contentProps?.missingModels; + if (!missingModels || !Array.isArray(missingModels)) { + console.log('[RunpodDirect] No missingModels in dialog contentProps'); + return null; + } + + console.log(`[RunpodDirect] Found ${missingModels.length} models from Pinia dialog store`); + return missingModels.map(m => ({ + filename: m.name, + url: m.url, + directory: m.directory, + })); } -// 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'); - if (existing) existing.remove(); +// Fallback: extract model info from DOM when Vue internals aren't accessible +function extractModelsFromDOM(container) { + const rows = getModelRows(container); + const models = []; + rows.forEach((row) => { + const leftSide = row.querySelector('[class*="overflow-hidden"]'); + if (!leftSide) return; - const progressArea = document.createElement('div'); - progressArea.className = 'server-download-progress-area'; - progressArea.style.cssText = ` - margin-top: 20px; - padding: 16px; - background: var(--p-content-background, #1e1e1e); - border: 1px solid var(--p-content-border-color, #333); - border-radius: 8px; - `; + const nameSpan = leftSide.querySelector('span[title]'); + if (!nameSpan) return; + const filename = nameSpan.getAttribute('title') || nameSpan.textContent.trim(); - progressArea.innerHTML = ` -
-
- Download Progress -
-
-
- Overall: 0/${totalDownloads} models completed -
-
- -
- `; - - listbox.parentElement.appendChild(progressArea); - - // Listen for updates - const updateHandler = (event) => { - const { download_id, status, progress, downloaded, total, speed } = event.detail; - - if (!isDownloadingAll) { - return; + const badgeSpan = leftSide.querySelector('span[class*="rounded-full"]'); + let directory = null; + if (badgeSpan) { + const badgeText = badgeSpan.textContent.trim().toUpperCase(); + directory = BADGE_TO_DIRECTORY[badgeText] || badgeText.toLowerCase(); } - // Update overall progress - const overallProgress = document.getElementById('server-download-overall-progress'); - if (overallProgress) { - overallProgress.textContent = `Overall: ${completedDownloads}/${totalDownloads} models completed`; - } - - // 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); + // Try button title or anchor href for URL + const rightSide = row.querySelector('[class*="shrink-0"]'); + let url = null; + if (rightSide) { + const urlButton = rightSide.querySelector('button[title]'); + if (urlButton) url = urlButton.getAttribute('title'); + if (!url) { + const urlAnchor = rightSide.querySelector('a[href]'); + if (urlAnchor) url = urlAnchor.getAttribute('href'); } } - // 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); + if (filename && directory && url) { + models.push({ filename, directory, url }); } - } catch (e) { - console.error('[RunpodDirect] Error in updateDownloadProgressItem:', e); - return; - } + }); + return models.length > 0 ? models : null; +} + +function extractModelsFromFooter() { + return null; +} + +// Check if HF_TOKEN env var is set on the backend +async function checkEnvHfToken() { 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} -
- `; + const response = await api.fetchApi("/server_download/hf_token_status"); + if (response.ok) { + const data = await response.json(); + envHasHfToken = data.has_token; } } catch (e) { - console.error('[RunpodDirect] Error updating progress item HTML:', e); + // Ignore - endpoint may not exist on older backend } } -// Export functions for use in other modules +// Detect gated models from the DOM (rows with "Accept terms" links) +function detectGatedModels(container) { + const gated = []; + const rows = getModelRows(container); + for (const row of rows) { + const link = row.querySelector('a[target="_blank"]'); + if (link && link.textContent.toLowerCase().includes('accept')) { + const leftSide = row.querySelector('[class*="overflow-hidden"]'); + if (!leftSide) continue; + const nameSpan = leftSide.querySelector('span[title]'); + if (!nameSpan) continue; + const filename = nameSpan.getAttribute('title') || nameSpan.textContent.trim(); + const badgeSpan = leftSide.querySelector('span[class*="rounded-full"]'); + let directory = null; + if (badgeSpan) { + const badgeText = badgeSpan.textContent.trim().toUpperCase(); + directory = BADGE_TO_DIRECTORY[badgeText] || badgeText.toLowerCase(); + } + gated.push({ filename, directory, repoUrl: link.href }); + // Remove the native "Accept terms" link — our token section handles it + link.remove(); + } + } + return gated; +} + +// Validate HF token against backend and check access to specific URLs +async function validateHfToken(token, urls) { + try { + const response = await api.fetchApi("/server_download/validate_hf_token", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ token, urls }) + }); + if (response.ok) return await response.json(); + return { valid: false, error: 'Validation request failed' }; + } catch (e) { + return { valid: false, error: e.message }; + } +} + +// Token section for gated models. +// Sits below the model list. Controls whether the download button gets enabled. +// +// States: +// 1. Initial: show input + Verify btn. Download btn disabled. +// 2. Token invalid: show error. Download btn disabled. +// 3. Token valid, ALL gated accessible: green success. Download btn enabled for all. +// 4. Token valid, SOME need terms: show which models need terms (with links). Download btn disabled. +// User can click Accept terms → go to HF → come back → click Verify again. +function createTokenSection(dialog, gatedModels, gatedWithUrls, callbacks) { + if (document.querySelector('.server-download-token-section')) return; + + const modelList = dialog.querySelector('[class*="scrollbar-custom"][class*="overflow-y-auto"][class*="rounded-lg"]'); + if (!modelList) return; + + const section = createEl('div', { + borderRadius: '0.5rem', + backgroundColor: THEME.secondaryBg, + overflow: 'hidden', + }); + section.className = 'server-download-token-section'; + + // Header + const header = createEl('div', { + padding: '8px 12px', + display: 'flex', + justifyContent: 'space-between', + alignItems: 'center', + }); + const headerTitle = createEl('span', { + fontSize: '0.75rem', + fontWeight: '600', + color: THEME.foreground, + }, 'HF Token Required'); + const headerStatus = createEl('span', { + fontSize: '0.6875rem', + color: THEME.muted, + }, `${gatedModels.length} gated model${gatedModels.length > 1 ? 's' : ''}`); + header.appendChild(headerTitle); + header.appendChild(headerStatus); + section.appendChild(header); + + // Body area — holds input, status messages, terms list + const body = createEl('div', { + padding: '0 12px 10px 12px', + display: 'flex', + flexDirection: 'column', + gap: '8px', + }); + section.appendChild(body); + + // Status/message area + const statusEl = createEl('div', { + fontSize: '0.6875rem', + color: THEME.muted, + lineHeight: '1.4', + }); + + // Terms list container (shown when some models need terms) + const termsListEl = createEl('div', { + display: 'none', + flexDirection: 'column', + gap: '4px', + }); + + function showTermsList(deniedModels) { + while (termsListEl.firstChild) termsListEl.removeChild(termsListEl.firstChild); + termsListEl.style.display = 'flex'; + for (const m of deniedModels) { + const row = createEl('div', { + display: 'flex', + justifyContent: 'space-between', + alignItems: 'center', + padding: '4px 0', + }); + const name = createEl('span', { + fontSize: '0.75rem', + color: THEME.foreground, + overflow: 'hidden', + textOverflow: 'ellipsis', + whiteSpace: 'nowrap', + minWidth: '0', + }, m.filename); + const btn = createEl('button', { + backgroundColor: THEME.secondaryBgHover, + color: THEME.foreground, + border: 'none', + height: '22px', + padding: '0 8px', + fontSize: '0.6875rem', + fontWeight: '500', + borderRadius: '0.375rem', + cursor: 'pointer', + whiteSpace: 'nowrap', + flexShrink: '0', + marginLeft: '8px', + }, 'Accept terms'); + btn.type = 'button'; + btn.onclick = () => window.open(m.repoUrl, '_blank', 'noopener,noreferrer'); + row.appendChild(name); + row.appendChild(btn); + termsListEl.appendChild(row); + } + } + + function hideTermsList() { + termsListEl.style.display = 'none'; + } + + // Input row builder + function buildInputRow() { + const row = createEl('div', { + display: 'flex', + gap: '6px', + alignItems: 'center', + }); + + const input = document.createElement('input'); + input.type = 'password'; + input.placeholder = 'hf_...'; + input.autocomplete = 'off'; + input.spellcheck = false; + Object.assign(input.style, { + flex: '1', + height: '28px', + padding: '0 8px', + fontSize: '0.75rem', + borderRadius: '0.375rem', + border: `1px solid ${THEME.border}`, + backgroundColor: THEME.baseBg, + color: THEME.foreground, + outline: 'none', + fontFamily: 'monospace', + minWidth: '0', + }); + input.onfocus = () => { input.style.borderColor = THEME.primary; }; + input.onblur = () => { input.style.borderColor = THEME.border; }; + + const verifyBtn = createEl('button', { + backgroundColor: THEME.primary, + color: THEME.foreground, + border: 'none', + height: '28px', + padding: '0 10px', + fontSize: '0.75rem', + fontWeight: '500', + borderRadius: '0.375rem', + cursor: 'pointer', + whiteSpace: 'nowrap', + flexShrink: '0', + }, 'Verify'); + verifyBtn.type = 'button'; + + verifyBtn.onclick = () => doVerify(input, verifyBtn); + input.onkeydown = (e) => { if (e.key === 'Enter') doVerify(input, verifyBtn); }; + + row.appendChild(input); + row.appendChild(verifyBtn); + return row; + } + + async function doVerify(input, verifyBtn) { + const val = input.value.trim(); + if (!val || !val.startsWith('hf_')) { + statusEl.textContent = 'Token must start with hf_'; + statusEl.style.color = THEME.error; + return; + } + + verifyBtn.disabled = true; + verifyBtn.textContent = 'Verifying...'; + verifyBtn.style.opacity = '0.5'; + statusEl.textContent = ''; + hideTermsList(); + + const result = await validateHfToken(val, gatedWithUrls.map(m => m.url)); + + verifyBtn.disabled = false; + verifyBtn.textContent = 'Verify'; + verifyBtn.style.opacity = '1'; + + if (!result.valid) { + statusEl.textContent = result.error || 'Invalid token'; + statusEl.style.color = THEME.error; + callbacks.onFail(); + return; + } + + // Token valid — check per-model access + sessionHfToken = val; + const accessible = []; + const denied = []; + for (const m of gatedWithUrls) { + const access = result.url_access?.[m.url]; + if (access?.accessible) { + accessible.push(m); + } else { + denied.push({ ...m, repoUrl: access?.repo_url || m.url }); + } + } + + if (denied.length === 0) { + // All good — green + headerStatus.textContent = 'all accessible'; + headerStatus.style.color = THEME.success; + statusEl.textContent = `Verified as ${result.username}`; + statusEl.style.color = THEME.success; + input.disabled = true; + input.style.opacity = '0.5'; + verifyBtn.textContent = 'OK'; + verifyBtn.disabled = true; + verifyBtn.style.backgroundColor = THEME.success; + verifyBtn.style.opacity = '0.7'; + callbacks.onAllAccessible(gatedWithUrls); + } else { + // Some need terms + statusEl.textContent = `Verified as ${result.username} — ${denied.length} model${denied.length > 1 ? 's' : ''} need terms accepted:`; + statusEl.style.color = THEME.warning; + showTermsList(denied); + // Keep input active so user can re-verify after accepting terms + callbacks.onPartialAccess(accessible, denied); + } + } + + // Auto-validate env token + async function autoValidateEnv() { + statusEl.textContent = 'HF_TOKEN found in environment. Validating...'; + statusEl.style.color = THEME.muted; + sessionHfToken = '__env__'; + + const result = await validateHfToken('__env__', gatedWithUrls.map(m => m.url)); + + if (!result.valid) { + statusEl.textContent = 'Environment HF_TOKEN is invalid.'; + statusEl.style.color = THEME.error; + sessionHfToken = null; + // Show manual input as fallback + body.insertBefore(buildInputRow(), statusEl); + callbacks.onFail(); + return; + } + + const accessible = []; + const denied = []; + for (const m of gatedWithUrls) { + const access = result.url_access?.[m.url]; + if (access?.accessible) { + accessible.push(m); + } else { + denied.push({ ...m, repoUrl: access?.repo_url || m.url }); + } + } + + if (denied.length === 0) { + headerStatus.textContent = 'all accessible'; + headerStatus.style.color = THEME.success; + statusEl.textContent = `Verified as ${result.username}`; + statusEl.style.color = THEME.success; + callbacks.onAllAccessible(gatedWithUrls); + } else { + statusEl.textContent = `Verified as ${result.username} — ${denied.length} model${denied.length > 1 ? 's' : ''} need terms accepted:`; + statusEl.style.color = THEME.warning; + showTermsList(denied); + callbacks.onPartialAccess(accessible, denied); + } + } + + // Build the body + if (envHasHfToken) { + body.appendChild(statusEl); + body.appendChild(termsListEl); + } else { + body.appendChild(buildInputRow()); + body.appendChild(statusEl); + body.appendChild(termsListEl); + } + + // Insert after model list + modelList.parentElement.insertBefore(section, modelList.nextSibling); + + // Auto-validate if env token exists + if (envHasHfToken) { + autoValidateEnv(); + } + + return section; +} + +function findMissingModelsContainer() { + const containers = document.querySelectorAll('[class*="scrollbar-custom"][class*="overflow-y-auto"][class*="rounded-lg"]'); + for (const container of containers) { + const rows = container.querySelectorAll(':scope > div'); + if (rows.length > 0) { + for (const row of rows) { + if (row.querySelector('span[class*="rounded-full"]')) { + return container; + } + } + } + } + + const dialogs = document.querySelectorAll('[role="dialog"]'); + for (const dialog of dialogs) { + const text = dialog.textContent || ''; + if (text.includes('missing models') || text.includes('Missing models')) { + const scrollable = dialog.querySelector('[class*="overflow-y-auto"]'); + if (scrollable) return scrollable; + } + } + return null; +} + +function getModelRows(container) { + const allRows = container.querySelectorAll(':scope > div'); + return Array.from(allRows).filter(row => !row.classList.contains('sticky')); +} + +// --- UI injection (theme-aware, matches ComfyUI design system) --- + +function createProgressArea(container) { + const existing = document.querySelector('.server-download-progress-area'); + if (existing) existing.remove(); + + // Match the model list container style: rounded-lg bg-secondary-background + const area = createEl('div', { + borderRadius: '0.5rem', + backgroundColor: THEME.secondaryBg, + overflow: 'hidden', + }); + area.className = 'server-download-progress-area'; + + // Header row matching the sticky bottom row style from the model list + const header = createEl('div', { + padding: '8px 12px', + display: 'flex', + justifyContent: 'space-between', + alignItems: 'center', + borderBottom: `1px solid ${THEME.border}`, + }); + const headerTitle = createEl('span', { + fontSize: '0.75rem', + fontWeight: '600', + color: THEME.foreground, + }, 'Download Progress'); + const headerStatus = createEl('span', { + fontSize: '0.75rem', + color: THEME.muted, + }, `0/${totalDownloads} completed`); + headerStatus.id = 'server-download-overall-progress'; + header.appendChild(headerTitle); + header.appendChild(headerStatus); + area.appendChild(header); + + const itemsContainer = createEl('div', { + display: 'flex', + flexDirection: 'column', + }); + itemsContainer.id = 'server-download-items-container'; + area.appendChild(itemsContainer); + + // Insert after the model list container (inside the same parent flex column) + container.parentElement.insertBefore(area, container.nextSibling); + + window.addEventListener('serverDownloadUpdate', (event) => { + const { download_id, status, progress, downloaded, total, speed } = event.detail; + if (!isDownloadingAll) return; + + const overallEl = document.getElementById('server-download-overall-progress'); + if (overallEl) { + overallEl.textContent = `${completedDownloads}/${totalDownloads} completed`; + } + updateDownloadProgressItem(download_id, status, progress, downloaded, total, speed); + }); +} + +function updateDownloadProgressItem(download_id, status, progress, downloaded, total, speed) { + const itemId = `download-item-${download_id.replace(/\//g, '-')}`; + const container = document.getElementById('server-download-items-container'); + if (!container) return; + + let item = document.getElementById(itemId); + + if (status === 'queued') { + if (item) item.remove(); + return; + } + + if ((status === 'completed' || status === 'error') && item && !item.dataset.removing) { + item.dataset.removing = 'true'; + setTimeout(() => { try { if (item && item.parentNode) item.remove(); } catch (e) { /* */ } }, 2000); + } + + if (!item) { + // Match model list row style: px-3 py-2 with no extra bg/border + item = createEl('div', { padding: '8px 12px' }); + item.id = itemId; + container.appendChild(item); + } + + const progressPercent = progress || 0; + const speedText = speed || '--'; + const sizeText = downloaded && total ? `${formatBytes(downloaded)} / ${formatBytes(total)}` : '--'; + + while (item.firstChild) item.removeChild(item.firstChild); + + // Row 1: filename + percentage (matches model row: name left, info right) + const nameRow = createEl('div', { + display: 'flex', justifyContent: 'space-between', alignItems: 'center', marginBottom: '4px', + }); + const filename = download_id.split('/').pop() || download_id; + const nameEl = createEl('span', { + fontSize: '0.875rem', color: THEME.foreground, fontWeight: '400', + overflow: 'hidden', textOverflow: 'ellipsis', whiteSpace: 'nowrap', minWidth: '0', + }, filename); + const pctEl = createEl('span', { + fontSize: '0.75rem', color: THEME.muted, flexShrink: '0', paddingLeft: '8px', + }, progressPercent.toFixed(1) + '%'); + nameRow.appendChild(nameEl); + nameRow.appendChild(pctEl); + item.appendChild(nameRow); + + // Row 2: progress bar + const barOuter = createEl('div', { + width: '100%', height: '4px', + backgroundColor: THEME.secondaryBgHover, + borderRadius: '9999px', overflow: 'hidden', marginBottom: '4px', + }); + const color = statusColor(status); + const barInner = createEl('div', { + height: '100%', backgroundColor: color, + borderRadius: '9999px', + width: progressPercent + '%', transition: 'width 0.3s', + }); + barOuter.appendChild(barInner); + item.appendChild(barOuter); + + // Row 3: speed + size + const infoRow = createEl('div', { + display: 'flex', justifyContent: 'space-between', + fontSize: '0.6875rem', color: THEME.muted, + }); + infoRow.appendChild(createEl('span', {}, speedText)); + infoRow.appendChild(createEl('span', {}, sizeText)); + item.appendChild(infoRow); +} + +// Export functions window.serverDownload = { start: startServerDownload, getStatus: getDownloadStatus, states: downloadStates }; -// Helper to inject buttons using MutationObserver +// --- Main injection logic --- + +function injectServerDownloadButtons() { + console.log('[RunpodDirect] injectServerDownloadButtons called'); + + const container = findMissingModelsContainer(); + if (!container) { + console.log('[RunpodDirect] Missing models container not found'); + return; + } + + console.log('[RunpodDirect] Found missing models container'); + + if (document.querySelector('.server-download-all-btn')) { + console.log('[RunpodDirect] Buttons already injected'); + return; + } + + // Extract model data: try Pinia store first (most reliable - URLs aren't in DOM), + // then fall back to DOM parsing + let models = null; + + // Strategy 1: Pinia dialog store (most reliable - works in production builds) + const piniaModels = extractModelsFromPinia(); + if (piniaModels && piniaModels.length > 0) { + models = piniaModels; + } + + // Strategy 2: DOM parsing (fallback - works when buttons with URLs are visible) + if (!models) { + const domModels = extractModelsFromDOM(container); + if (domModels && domModels.length > 0) { + models = domModels; + console.log(`[RunpodDirect] Got ${models.length} models from DOM parsing`); + } + } + + // Strategy 3: Footer component fallback + if (!models) { + const footerData = extractModelsFromFooter(); + if (footerData && footerData.models?.length > 0) { + models = footerData.models; + console.log(`[RunpodDirect] Got ${models.length} models from footer component`); + } + } + + if (!models || models.length === 0) { + console.log('[RunpodDirect] No models could be extracted, will retry...'); + // Retry after a short delay - Pinia store may not be populated yet + if (!container.dataset.retryCount || parseInt(container.dataset.retryCount) < 5) { + container.dataset.retryCount = (parseInt(container.dataset.retryCount || '0') + 1).toString(); + setTimeout(() => injectServerDownloadButtons(), 500); + } else { + console.log('[RunpodDirect] Max retries reached, giving up'); + } + return; + } + + // Separate downloadable models from gated ones + const gatedModels = detectGatedModels(container); + const gatedFilenames = new Set(gatedModels.map(g => g.filename)); + + // Models from Pinia include ALL models (including gated ones). + const downloadableModels = models.filter(m => !gatedFilenames.has(m.filename)); + const gatedWithUrls = models.filter(m => gatedFilenames.has(m.filename)); + + downloadableModels.forEach((m, i) => { + console.log(`[RunpodDirect] Model ${i + 1}: ${m.directory}/${m.filename} -> ${m.url}`); + }); + if (gatedWithUrls.length > 0) { + console.log(`[RunpodDirect] ${gatedWithUrls.length} gated model(s) detected`); + } + + // Mutable list — starts with only non-gated, grows when token is validated + let allModelsToDownload = [...downloadableModels]; + // Track whether gated models block the download button + const hasGated = gatedModels.length > 0; + let gatedVerified = false; + + const dialog = container.closest('[role="dialog"]'); + const footerBtnRow = dialog?.querySelector('div[class*="justify-end"][class*="gap"]'); + + // Show gated models section if needed + if (hasGated && dialog) { + createTokenSection(dialog, gatedModels, gatedWithUrls, { + onAllAccessible(accessibleModels) { + // All gated models verified — enable download for everything + gatedVerified = true; + allModelsToDownload = [...downloadableModels, ...accessibleModels]; + updateBtnCount(); + console.log(`[RunpodDirect] All gated models accessible, ${allModelsToDownload.length} total`); + }, + onPartialAccess(_accessible, denied) { + // Some models need terms — keep button disabled + gatedVerified = false; + allModelsToDownload = [...downloadableModels]; + updateBtnCount(); + console.log(`[RunpodDirect] ${denied.length} model(s) need terms accepted`); + }, + onFail() { + // Invalid token — keep button disabled + gatedVerified = false; + allModelsToDownload = [...downloadableModels]; + updateBtnCount(); + console.log('[RunpodDirect] Token validation failed'); + }, + }); + } + + // Match ComfyUI's primary button: bg-primary-background text-base-foreground h-8 rounded-lg text-xs + const downloadAllBtn = createEl('button', { + backgroundColor: THEME.primary, + color: THEME.foreground, + border: 'none', + height: '32px', + padding: '0 8px', + fontSize: '0.75rem', + fontWeight: '500', + borderRadius: '0.5rem', + cursor: 'pointer', + display: 'inline-flex', + alignItems: 'center', + justifyContent: 'center', + whiteSpace: 'nowrap', + transition: 'background-color 0.15s', + }); + downloadAllBtn.className = 'server-download-all-btn'; + downloadAllBtn.type = 'button'; + downloadAllBtn.title = 'Download all models directly to this RunPod instance (server-side)'; + + function updateBtnCount() { + const count = allModelsToDownload.length; + downloadAllBtn.textContent = `Download to Pod (${count})`; + // Disable when: no models OR gated models exist but token not verified + const shouldDisable = count === 0 || (hasGated && !gatedVerified); + if (shouldDisable) { + downloadAllBtn.disabled = true; + downloadAllBtn.style.opacity = '0.5'; + downloadAllBtn.style.cursor = 'default'; + } else { + downloadAllBtn.disabled = false; + downloadAllBtn.style.opacity = '1'; + downloadAllBtn.style.cursor = 'pointer'; + } + } + updateBtnCount(); + + downloadAllBtn.onmouseenter = () => { if (!downloadAllBtn.disabled) downloadAllBtn.style.backgroundColor = THEME.primaryHover; }; + downloadAllBtn.onmouseleave = () => { if (!downloadAllBtn.disabled) downloadAllBtn.style.backgroundColor = THEME.primary; }; + + function setButtonRefresh() { + downloadAllBtn.disabled = false; + downloadAllBtn.style.opacity = '1'; + downloadAllBtn.style.pointerEvents = 'auto'; + downloadAllBtn.style.backgroundColor = THEME.success; + downloadAllBtn.style.color = THEME.foreground; + downloadAllBtn.textContent = 'Refresh Page'; + downloadAllBtn.onmouseenter = null; + downloadAllBtn.onmouseleave = null; + downloadAllBtn.onclick = () => location.reload(); + } + + downloadAllBtn.onclick = async (e) => { + e.stopPropagation(); + downloadAllBtn.disabled = true; + downloadAllBtn.style.opacity = '0.5'; + downloadAllBtn.style.cursor = 'default'; + downloadAllBtn.textContent = 'Downloading...'; + + downloadQueue = allModelsToDownload.map(m => ({ url: m.url, directory: m.directory, filename: m.filename })); + totalDownloads = allModelsToDownload.length; + completedDownloads = 0; + isDownloadingAll = true; + + createProgressArea(container); + + if (downloadQueue.length > 0) { + processDownloadQueue(); + } + }; + + // Listen for all downloads completing to show refresh button and remove progress area + window.addEventListener('serverDownloadAllDone', () => { + setButtonRefresh(); + const progressArea = document.querySelector('.server-download-progress-area'); + if (progressArea) progressArea.remove(); + }); + + if (footerBtnRow) { + // Insert before the native "Download all" button + footerBtnRow.insertBefore(downloadAllBtn, footerBtnRow.firstChild); + } else { + // Fallback: place above the model list + const fallbackContainer = createEl('div', { + padding: '0 16px 8px 16px', + display: 'flex', + justifyContent: 'center', + }); + fallbackContainer.appendChild(downloadAllBtn); + container.parentElement.insertBefore(fallbackContainer, container); + } + + console.log('[RunpodDirect] Button injection complete'); +} + +// MutationObserver for detecting the missing models dialog function setupDialogObserver() { console.log('[RunpodDirect] Setting up dialog observer'); - // Watch for the missing models dialog to appear const observer = new MutationObserver((mutations) => { for (const mutation of mutations) { if (mutation.addedNodes.length) { mutation.addedNodes.forEach(node => { - if (node.nodeType === Node.ELEMENT_NODE) { - // Check if this node or its descendants contain the missing models dialog - const hasDialog = node.querySelector && ( - node.querySelector('.comfy-missing-models') || - node.classList?.contains('comfy-missing-models') - ); + if (node.nodeType !== Node.ELEMENT_NODE) return; - if (hasDialog) { - console.log('[RunpodDirect] Detected missing models dialog, injecting buttons...'); - setTimeout(() => { - injectServerDownloadButtons(); - }, 500); - } + const isDialog = node.getAttribute?.('role') === 'dialog' || + node.querySelector?.('[role="dialog"]'); + + if (isDialog) { + setTimeout(() => { + const dialog = node.getAttribute?.('role') === 'dialog' + ? node + : node.querySelector('[role="dialog"]'); + if (dialog) { + const text = dialog.textContent || ''; + if (text.includes('missing models') || text.includes('Missing models')) { + console.log('[RunpodDirect] Detected missing models dialog'); + setTimeout(() => injectServerDownloadButtons(), 500); + } + } + }, 300); } }); } } }); - // Start observing the document body for dialog additions - observer.observe(document.body, { - childList: true, - subtree: true - }); - + observer.observe(document.body, { childList: true, subtree: true }); console.log('[RunpodDirect] Observer active'); } -function injectServerDownloadButtons() { - console.log('[RunpodDirect] injectServerDownloadButtons called'); - - // Find the missing models listbox - const listbox = document.querySelector('.comfy-missing-models'); - if (!listbox) { - console.log('[RunpodDirect] Missing models listbox not found'); - return; - } - - console.log('[RunpodDirect] Found .comfy-missing-models listbox'); - - // Check if we already added our UI - if (document.querySelector('.server-download-all-btn')) { - console.log('[RunpodDirect] Buttons already injected'); - return; - } - - const listItems = listbox.querySelectorAll('.p-listbox-option'); - console.log(`[RunpodDirect] Found ${listItems.length} list items`); - - if (listItems.length === 0) { - console.log('[RunpodDirect] No list items found'); - return; - } - - // Add "Download All" button before the listbox - const downloadAllContainer = document.createElement('div'); - downloadAllContainer.style.cssText = 'margin-bottom: 16px; display: flex; justify-content: center;'; - - const downloadAllBtn = document.createElement('button'); - downloadAllBtn.className = 'server-download-all-btn p-button p-component p-button-sm'; - downloadAllBtn.type = 'button'; - downloadAllBtn.style.cssText = 'background: #2196F3; color: white; border: none; padding: 10px 20px; font-weight: 600;'; - - const downloadAllLabel = document.createElement('span'); - downloadAllLabel.className = 'p-button-label'; - downloadAllLabel.textContent = `Download All Models to Pod (${listItems.length})`; - downloadAllBtn.appendChild(downloadAllLabel); - - downloadAllBtn.onclick = async (e) => { - e.stopPropagation(); - downloadAllBtn.disabled = true; - downloadAllLabel.textContent = 'Starting downloads...'; - - // Collect all models - downloadQueue = []; - const models = []; - - listItems.forEach((item) => { - const labelElement = item.querySelector('[title]'); - if (!labelElement) return; - - const label = labelElement.textContent.trim(); - const url = labelElement.getAttribute('title'); - const parts = label.split('/').map(p => p.trim()); - if (parts.length !== 2) return; - - const directory = parts[0]; - const filename = parts[1]; - - models.push({ url, directory, filename }); - }); - - downloadQueue = [...models]; - totalDownloads = models.length; - completedDownloads = 0; - isDownloadingAll = true; - - // Create progress area - createProgressArea(listbox); - - // Start first download - if (downloadQueue.length > 0) { - processDownloadQueue(); - } - }; - - downloadAllContainer.appendChild(downloadAllBtn); - listbox.parentElement.insertBefore(downloadAllContainer, listbox); - - listItems.forEach((item, index) => { - console.log(`[RunpodDirect] Processing item ${index + 1}`); - - // Check if we already added the button - if (item.querySelector('.server-download-btn')) { - console.log(`[RunpodDirect] Item ${index + 1} already has server download button, skipping`); - return; - } - - // The HTML structure is: - //
  • - //
    <- main container - //
    model info
    - //
    - //
    - //
    - //
    • - - // Find the main flex container - const mainContainer = item.querySelector('.flex.flex-row.items-center.gap-2'); - if (!mainContainer) { - console.log('[RunpodDirect] Main flex container not found'); - return; - } - - console.log('[RunpodDirect] Found main container'); - - // We'll create a new div for our button (following the same pattern) - const buttonWrapper = document.createElement('div'); - - // Get model info from the item - const labelElement = item.querySelector('[title]'); - if (!labelElement) { - console.log('[RunpodDirect] No title element found'); - return; - } - - const label = labelElement.textContent.trim(); - const url = labelElement.getAttribute('title'); - console.log(`[RunpodDirect] Model: ${label}, URL: ${url}`); - - // Parse "checkpoints / model.safetensors" format - const parts = label.split('/').map(p => p.trim()); - if (parts.length !== 2) { - console.log(`[RunpodDirect] Could not parse label format: ${label}`); - return; - } - - const directory = parts[0]; - const filename = parts[1]; - const download_id = `${directory}/${filename}`; - console.log(`[RunpodDirect] Creating button for ${download_id}`); - - // Create server download button - const serverDownloadBtn = document.createElement('button'); - serverDownloadBtn.className = 'server-download-btn p-button p-component p-button-outlined p-button-sm'; - serverDownloadBtn.type = 'button'; - - // Create button content - const btnLabel = document.createElement('span'); - btnLabel.className = 'p-button-label'; - btnLabel.textContent = 'Download to Pod'; - serverDownloadBtn.appendChild(btnLabel); - - // Status indicator (icon) - const statusIcon = document.createElement('i'); - statusIcon.style.cssText = 'margin-left: 6px; font-size: 14px; display: none;'; - serverDownloadBtn.appendChild(statusIcon); - - // Button click handler - serverDownloadBtn.onclick = async (e) => { - e.stopPropagation(); - serverDownloadBtn.disabled = true; - btnLabel.textContent = 'Starting...'; - - const result = await startServerDownload(url, directory, filename); - - if (result.success) { - 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'; - statusIcon.style.display = 'inline'; - statusIcon.style.color = '#ef4444'; - console.error('Download start failed:', result.error); - } - }; - - // Listen for download updates - const updateHandler = (event) => { - if (event.detail.download_id === download_id) { - const { status, error } = event.detail; - - 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'; - statusIcon.style.color = '#2196F3'; - } else if (status === 'completed') { - btnLabel.textContent = 'Completed'; - statusIcon.className = 'pi pi-check-circle'; - statusIcon.style.display = 'inline'; - statusIcon.style.color = '#4caf50'; - serverDownloadBtn.style.borderColor = '#4caf50'; - } else if (status === 'error') { - btnLabel.textContent = 'Failed'; - statusIcon.className = 'pi pi-times-circle'; - statusIcon.style.display = 'inline'; - statusIcon.style.color = '#ef4444'; - serverDownloadBtn.title = error; - } - } - }; - - window.addEventListener('serverDownloadUpdate', updateHandler); - - // Add button to wrapper div - buttonWrapper.appendChild(serverDownloadBtn); - - // Add wrapper to main container (alongside Download and Copy URL) - mainContainer.appendChild(buttonWrapper); - console.log(`[RunpodDirect] Button added to main container for ${download_id}`); - }); - - console.log('[RunpodDirect] Button injection complete'); -} - -// Register the extension +// Register extension app.registerExtension({ name: "ComfyUI.RunpodDirect", async setup() { console.log("[RunpodDirect] Extension setup starting"); - - // Set up observer to watch for missing models dialog + checkEnvHfToken(); setupDialogObserver(); - // Also try to inject immediately if dialog already exists setTimeout(() => { console.log('[RunpodDirect] Checking for existing dialog...'); injectServerDownloadButtons(); }, 1000); - // Try again after a longer delay setTimeout(() => { console.log('[RunpodDirect] Second check for dialog...'); injectServerDownloadButtons();