From b64597fd3f7faeadad47cfaedc57b341b692d426 Mon Sep 17 00:00:00 2001 From: Dariusz L Date: Fri, 24 Apr 2026 20:38:11 +0200 Subject: [PATCH] Add CivitAI URN download UI and metadata Move direct URN->CivitAI download handling into the frontend and enrich backend metadata. Removed server-side immediate URN download exposure and adjusted comment numbering. Backend now selects a primary file from resolved CivitAI model info, includes model_id, version_id, size, and a modelVersionId query on model URLs. Frontend: add renderDownloadSourceSection to render download buttons and CivitAI model-card links, display file sizes, and format names. Store missingModels for later updates; render a placeholder URN download section while resolving; when URN resolution returns, inject civitai info into the missing entry, build a download_source object, replace the placeholder with the rendered download UI, and attach the download button handler. Improved error handling and model URL construction to include modelVersionId when available. --- __init__.py | 69 +++++++------------------ web/linker.js | 139 ++++++++++++++++++++++++++++++++++++-------------- 2 files changed, 120 insertions(+), 88 deletions(-) diff --git a/__init__.py b/__init__.py index bb90e10..994fa25 100644 --- a/__init__.py +++ b/__init__.py @@ -210,54 +210,7 @@ class ModelLinkerExtension: } continue - # 1. For URNs, we already have an exact CivitAI identity. - # Expose it immediately as a direct download source. - if missing.get("is_urn"): - urn_model_id = missing.get("urn_model_id") or ( - missing.get("urn", {}) or {} - ).get("model_id") - urn_version_id = missing.get( - "urn_version_id" - ) or (missing.get("urn", {}) or {}).get( - "version_id" - ) - - if urn_model_id and urn_version_id: - civitai_info = resolve_urn( - urn_model_id, urn_version_id - ) - if civitai_info: - missing["civitai_info"] = civitai_info - expected_filename = ( - civitai_info.get("expected_filename") - or filename - ) - primary_file = ( - civitai_info.get("files") or [{}] - )[0] - - missing["download_source"] = { - "source": "civitai", - "url": get_civitai_download_url( - urn_version_id - ), - "filename": expected_filename, - "name": civitai_info.get( - "model_name" - ), - "type": missing.get("category"), - "directory": missing.get( - "category", "checkpoints" - ), - "match_type": "exact", - "size": primary_file.get("size"), - "model_id": urn_model_id, - "version_id": urn_version_id, - "model_url": f"https://civitai.com/models/{urn_model_id}?modelVersionId={urn_version_id}", - } - continue - - # 2. Check popular models (always exact match) + # 1. Check popular models (always exact match) popular_info = get_popular_model_url(filename) if popular_info: missing["download_source"] = { @@ -270,7 +223,7 @@ class ModelLinkerExtension: } continue - # 3. Check model list (ComfyUI Manager database) + # 2. Check model list (ComfyUI Manager database) # Use exact_only=True to avoid confusing fuzzy matches for downloads model_list_result = search_model_list( filename, exact_only=True @@ -772,6 +725,19 @@ class ModelLinkerExtension: # Use resolve_urn to get model info (cached) model_info = resolve_urn(model_id, version_id) if model_info: + primary_file = None + for file_info in model_info.get("files", []): + if ( + file_info.get("name") + == model_info.get("expected_filename") + ): + primary_file = file_info + break + if primary_file is None: + primary_file = ( + model_info.get("files") or [{}] + )[0] + download_url = get_civitai_download_url( version_id ) @@ -783,7 +749,10 @@ class ModelLinkerExtension: ), "type": category, "download_url": download_url, - "url": f"https://civitai.com/models/{model_id}", + "url": f"https://civitai.com/models/{model_id}?modelVersionId={version_id}", + "model_id": model_id, + "version_id": version_id, + "size": primary_file.get("size"), } results["found"] = True elif category: diff --git a/web/linker.js b/web/linker.js index c871f66..c7fb2c5 100644 --- a/web/linker.js +++ b/web/linker.js @@ -204,6 +204,52 @@ class LinkerManagerDialog extends ComfyDialog { this.syncSearchSourceUi(missing, container); } + renderDownloadSourceSection(missing, downloadSource) { + const filename = missing.original_path?.split('/').pop()?.split('\\').pop() || ''; + const isFromWorkflow = downloadSource.url_source === 'workflow'; + const isCivitaiSource = downloadSource.source === 'civitai'; + const sourceLabels = { + popular: 'Popular Models', + model_list: 'Model Database', + huggingface: 'HuggingFace', + civitai: 'CivitAI', + workflow: 'Workflow' + }; + const sourceLabel = isFromWorkflow ? 'Workflow' : (sourceLabels[downloadSource.source] || 'Online'); + const downloadFilename = downloadSource.filename || filename; + const formattedDownloadName = this.formatFilename(downloadFilename, 45); + + let sizeDisplay = ''; + if (downloadSource.size) { + sizeDisplay = typeof downloadSource.size === 'number' + ? this.formatBytes(downloadSource.size) + : downloadSource.size; + } + + const modelCardUrl = downloadSource.model_url || this.getModelCardUrl(downloadSource.url); + + let html = `
`; + html += ``; + if (isCivitaiSource && modelCardUrl) { + html += ` Open on CivitAI`; + } + html += `
`; + html += `${isFromWorkflow ? 'URL from workflow' : sourceLabel}`; + if (modelCardUrl) { + html += `
${formattedDownloadName.display}`; + } else { + html += `
${formattedDownloadName.display}`; + } + if (isCivitaiSource) { + html += `
Direct CivitAI download can require an API key for some models.`; + } + html += `
`; + html += `
`; + return html; + } + /** * Handle click outside context menu to hide it */ @@ -3073,6 +3119,7 @@ class LinkerManagerDialog extends ComfyDialog { displayMissingModels(container, data) { const missingModels = data.missing_models || []; const totalMissing = data.total_missing || 0; + this.missingModels = missingModels; // Check if there are active downloads const activeCount = Object.keys(this.activeDownloads).length; @@ -3504,48 +3551,16 @@ class LinkerManagerDialog extends ComfyDialog { const filename = missing.original_path?.split('/').pop()?.split('\\').pop() || ''; const downloadSource = missing.download_source; + const urnDownloadId = `urn-download-${missing.node_id}-${missing.widget_index}`; if (perfectMatches.length > 0) { // Has perfect local match - download not needed html += `
Not needed - exact local match available
`; } else if (downloadSource && downloadSource.url) { - // We have a known download URL - const isExact = downloadSource.match_type === 'exact' || downloadSource.source === 'popular' || downloadSource.source === 'huggingface' || downloadSource.source === 'civitai'; - const isFromWorkflow = downloadSource.url_source === 'workflow'; - const sourceLabels = { - 'popular': 'Popular Models', - 'model_list': 'Model Database', - 'huggingface': 'HuggingFace', - 'civitai': 'CivitAI', - 'workflow': 'Workflow' - }; - const sourceLabel = isFromWorkflow ? 'Workflow' : (sourceLabels[downloadSource.source] || 'Online'); - const downloadFilename = downloadSource.filename || filename; - const formattedDownloadName = this.formatFilename(downloadFilename, 45); - - // Format file size - let sizeDisplay = ''; - if (downloadSource.size) { - if (typeof downloadSource.size === 'number') { - sizeDisplay = this.formatBytes(downloadSource.size); - } else { - sizeDisplay = downloadSource.size; - } - } - - html += `
`; - html += ``; - html += `
`; - html += `${isFromWorkflow ? 'URL from workflow' : sourceLabel}`; - const modelCardUrl = this.getModelCardUrl(downloadSource.url); - if (modelCardUrl) { - html += `
${formattedDownloadName.display}`; - } else { - html += `
${formattedDownloadName.display}`; - } - html += `
`; + html += this.renderDownloadSourceSection(missing, downloadSource); + } else if (missing.is_urn) { + html += `
`; + html += `
Resolving CivitAI download for this URN...
`; html += `
`; } else { // No known download - offer search @@ -4406,12 +4421,55 @@ class LinkerManagerDialog extends ComfyDialog { loadingEl.textContent = 'Not found'; loadingEl.style.color = 'var(--ml-text-muted)'; } + + const downloadContainerId = loadingElementId.replace('urn-loading-', 'urn-download-'); + const downloadEl = document.getElementById(downloadContainerId); + if (downloadEl && data.civitai) { + const missing = this.missingModels.find(m => + `urn-download-${m.node_id}-${m.widget_index}` === downloadContainerId + ); + if (missing) { + missing.civitai_info = { + model_name: data.civitai.name, + version_name: data.civitai.version_name, + expected_filename: data.civitai.filename + }; + missing.download_source = { + source: 'civitai', + url: data.civitai.download_url, + filename: data.civitai.filename, + name: data.civitai.name, + type: data.civitai.type, + directory: missing.category || 'checkpoints', + match_type: 'exact', + size: data.civitai.size, + model_id: data.civitai.model_id || modelId, + version_id: data.civitai.version_id || versionId, + model_url: data.civitai.url || `https://civitai.com/models/${modelId}?modelVersionId=${versionId}` + }; + downloadEl.outerHTML = this.renderDownloadSourceSection(missing, missing.download_source); + + const refreshedBtn = this.contentElement?.querySelector(`#download-${missing.node_id}-${missing.widget_index}`); + if (refreshedBtn) { + refreshedBtn.addEventListener('click', () => { + this.downloadModel(missing); + }); + } + } + } else if (downloadEl) { + downloadEl.innerHTML = `
Unable to resolve direct download for this URN.
`; + } } else { const loadingEl = document.getElementById(loadingElementId); if (loadingEl) { loadingEl.textContent = 'Error'; loadingEl.style.color = '#f44336'; } + const downloadContainerId = loadingElementId.replace('urn-loading-', 'urn-download-'); + const downloadEl = document.getElementById(downloadContainerId); + if (downloadEl) { + downloadEl.innerHTML = `
Failed to resolve URN download.
`; + } } } catch (error) { console.error('Model Linker: URN resolve error:', error); @@ -4420,6 +4478,11 @@ class LinkerManagerDialog extends ComfyDialog { loadingEl.textContent = 'Error'; loadingEl.style.color = '#f44336'; } + const downloadContainerId = loadingElementId.replace('urn-loading-', 'urn-download-'); + const downloadEl = document.getElementById(downloadContainerId); + if (downloadEl) { + downloadEl.innerHTML = `
Failed to resolve URN download.
`; + } } } @@ -4505,7 +4568,7 @@ class LinkerManagerDialog extends ComfyDialog { // CivitAI result if (civitaiResult && civitaiResult.download_url) { - const modelUrl = civitaiResult.url || `https://civitai.com/models/${civitaiResult.model_id || ''}`; + const modelUrl = civitaiResult.url || (civitaiResult.model_id ? `https://civitai.com/models/${civitaiResult.model_id}${civitaiResult.version_id ? `?modelVersionId=${civitaiResult.version_id}` : ''}` : ''); // Use expected_filename from URN resolution, or fallback to civitaiResult filename const downloadFilename = missing.civitai_info?.expected_filename || civitaiResult.filename || civitaiResult.name; // Build display name with version if available