diff --git a/__init__.py b/__init__.py index 7e0b654..bb90e10 100644 --- a/__init__.py +++ b/__init__.py @@ -210,7 +210,54 @@ class ModelLinkerExtension: } continue - # 1. Check popular models (always exact match) + # 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) popular_info = get_popular_model_url(filename) if popular_info: missing["download_source"] = { @@ -223,7 +270,7 @@ class ModelLinkerExtension: } continue - # 2. Check model list (ComfyUI Manager database) + # 3. 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 @@ -650,9 +697,30 @@ class ModelLinkerExtension: status=400, ) + raw_sources = data.get("sources", ["all"]) + if isinstance(raw_sources, str): + raw_sources = [raw_sources] + elif not isinstance(raw_sources, list): + raw_sources = ["all"] + + normalized_sources = { + str(source).strip().lower() + for source in raw_sources + if str(source).strip() + } + if not normalized_sources: + normalized_sources = {"all"} + + if "all" in normalized_sources: + normalized_sources = {"local", "huggingface", "civitai"} + + search_local = "local" in normalized_sources + search_huggingface_source = "huggingface" in normalized_sources + search_civitai_source = "civitai" in normalized_sources + # Debug logging self.logger.info( - f"Search request: filename={filename}, category={category}, is_urn={is_urn}, model_id={data.get('model_id')}, version_id={data.get('version_id')}" + f"Search request: filename={filename}, category={category}, is_urn={is_urn}, model_id={data.get('model_id')}, version_id={data.get('version_id')}, sources={sorted(normalized_sources)}" ) results = { @@ -661,26 +729,23 @@ class ModelLinkerExtension: "huggingface": None, "civitai": None, "found": False, + "searched_sources": sorted(normalized_sources), } - # 1. Check popular models first (curated database) - popular_info = get_popular_model_url(filename) - if popular_info: - results["popular"] = { - "source": "popular", - "filename": filename, - **popular_info, - } - results["found"] = True + # 1. Search local databases (curated + model-list) + if search_local: + popular_info = get_popular_model_url(filename) + if popular_info: + results["popular"] = { + "source": "popular", + "filename": filename, + **popular_info, + } + results["found"] = True - # 2. Search model-list.json (ComfyUI Manager database with fuzzy matching) - # Only accept if confidence >= 70% (or if not URN, allow any) - if not results["found"]: model_list_result = search_model_list(filename) if model_list_result: confidence = model_list_result.get("confidence", 0) - # For URNs, only accept high-confidence matches (>70%) - # Otherwise continue to CivitAI search for exact model if is_urn and confidence >= 70: results["model_list"] = model_list_result results["found"] = True @@ -688,15 +753,15 @@ class ModelLinkerExtension: results["model_list"] = model_list_result results["found"] = True - # 3. Search HuggingFace for exact file match - if not results["found"]: + # 2. Search HuggingFace for exact file match + if search_huggingface_source: hf_result = search_huggingface_for_file(filename) if hf_result: results["huggingface"] = hf_result results["found"] = True - # 4. Search CivitAI - use direct download for URNs - if not results["found"]: + # 3. Search CivitAI - use direct download for URNs + if search_civitai_source: # For URNs, use direct model_id/version_id to get download URL if is_urn: # Get model_id and version_id from request data diff --git a/web/linker.js b/web/linker.js index 3a11edd..4646926 100644 --- a/web/linker.js +++ b/web/linker.js @@ -19,6 +19,7 @@ class LinkerManagerDialog extends ComfyDialog { this.pendingResolutions = []; this.pendingIndex = new Map(); // key -> index in pendingResolutions this.activeDownloads = {}; // Track active downloads + this.searchResultCache = new Map(); this.boundHandleOutsideClick = this.handleOutsideClick.bind(this); this.activeTab = 'missing'; // Default tab this.fullscreen = false; @@ -113,6 +114,96 @@ class LinkerManagerDialog extends ComfyDialog { document.addEventListener('click', this.boundHandleContextMenuClick); } + /** + * Build stable cache key for a missing model entry + */ + getMissingSearchKey(missing) { + return `${missing.node_id}:${missing.widget_index}`; + } + + /** + * Get or initialize search state for a missing model entry + */ + getSearchState(missing) { + const key = this.getMissingSearchKey(missing); + if (!this.searchResultCache.has(key)) { + this.searchResultCache.set(key, { + selectedSource: 'all', + results: { + popular: null, + model_list: null, + huggingface: null, + civitai: null + }, + lastAttemptSources: [], + lastAttemptFound: null + }); + } + return this.searchResultCache.get(key); + } + + /** + * Merge new search results into cached per-source results. + * Empty responses do not delete previously found results. + */ + mergeSearchResults(existingResults = {}, newResults = {}) { + return { + popular: newResults.popular || existingResults.popular || null, + model_list: newResults.model_list || existingResults.model_list || null, + huggingface: newResults.huggingface || existingResults.huggingface || null, + civitai: newResults.civitai || existingResults.civitai || null + }; + } + + /** + * Return true when at least one downloadable source was found + */ + hasSearchResults(data = {}) { + return !!(data.popular || data.model_list || data.huggingface || data.civitai); + } + + /** + * Convert source ids to readable labels + */ + getSearchSourceLabel(source) { + const labels = { + all: 'Everything', + local: 'Local Database', + huggingface: 'HuggingFace', + civitai: 'CivitAI' + }; + return labels[source] || source; + } + + /** + * Update source selector buttons and helper text for one card + */ + syncSearchSourceUi(missing, container) { + if (!container) return; + + const state = this.getSearchState(missing); + const buttons = container.querySelectorAll('.ml-search-source-btn'); + buttons.forEach(btn => { + const isActive = btn.dataset.searchSource === state.selectedSource; + btn.classList.toggle('ml-btn-primary', isActive); + btn.classList.toggle('ml-btn-secondary', !isActive); + }); + + const infoEl = container.querySelector(`#search-source-info-${missing.node_id}-${missing.widget_index}`); + if (infoEl) { + infoEl.textContent = `Selected source: ${this.getSearchSourceLabel(state.selectedSource)}`; + } + } + + /** + * Set current search source for one card + */ + setSearchSource(missing, source, container) { + const state = this.getSearchState(missing); + state.selectedSource = source || 'all'; + this.syncSearchSourceUi(missing, container); + } + /** * Handle click outside context menu to hide it */ @@ -2644,6 +2735,7 @@ class LinkerManagerDialog extends ComfyDialog { } const data = await response.json(); + this.searchResultCache.clear(); this.displayMissingModels(this.contentElement, data); // Reconnect any active downloads to their new progress divs @@ -2942,6 +3034,16 @@ class LinkerManagerDialog extends ComfyDialog { this.searchOnline(missing); }); } + + const sourceButtons = container.querySelectorAll('.ml-search-source-btn'); + sourceButtons.forEach(btn => { + btn.addEventListener('click', () => { + this.setSearchSource(missing, btn.dataset.searchSource, container); + }); + }); + if (sourceButtons.length > 0) { + this.syncSearchSourceUi(missing, container); + } // Wire Locate button (only for top-level nodes) const locateId = `locate-${missing.node_id}-${missing.widget_index}`; @@ -3248,10 +3350,18 @@ class LinkerManagerDialog extends ComfyDialog { } else { // No known download - offer search html += `
`; + const searchSourcesId = `search-sources-${missing.node_id}-${missing.widget_index}`; + const searchInfoId = `search-source-info-${missing.node_id}-${missing.widget_index}`; + html += `
`; + html += ``; + html += ``; + html += ``; + html += ``; + html += `
`; html += ``; - html += `
Search HuggingFace & CivitAI
`; + html += `
Selected source: Everything
`; html += `
`; html += ``; } @@ -3968,6 +4078,9 @@ class LinkerManagerDialog extends ComfyDialog { async searchOnline(missing) { let filename = missing.original_path?.split('/').pop()?.split('\\').pop() || ''; let category = missing.category || ''; + const state = this.getSearchState(missing); + const selectedSource = state.selectedSource || 'all'; + const selectedSourceLabel = this.getSearchSourceLabel(selectedSource); // For URNs, use the CivitAI model name for searching instead of the URN itself // and pass the URN type as category (CivitAI expects specific type names) @@ -4003,15 +4116,15 @@ class LinkerManagerDialog extends ComfyDialog { try { if (searchBtn) { searchBtn.disabled = true; - searchBtn.textContent = '🔍 Searching...'; + searchBtn.textContent = `🔍 Searching ${selectedSourceLabel}...`; } if (resultsDiv) { resultsDiv.style.display = 'block'; - resultsDiv.innerHTML = 'Searching HuggingFace and CivitAI...'; + resultsDiv.innerHTML = `Searching ${selectedSourceLabel}...`; } // For URNs, include model_id and version_id for direct download - const searchData = { filename, category, is_urn: isUrn }; + const searchData = { filename, category, is_urn: isUrn, sources: [selectedSource] }; if (isUrn && missing.urn) { searchData.model_id = missing.urn.model_id; searchData.version_id = missing.urn.version_id; @@ -4031,7 +4144,10 @@ class LinkerManagerDialog extends ComfyDialog { const data = await response.json(); console.log('Model Linker: Search response:', JSON.stringify(data)); - this.displaySearchResults(missing, data, resultsDiv); + state.results = this.mergeSearchResults(state.results, data); + state.lastAttemptSources = Array.isArray(data.searched_sources) ? data.searched_sources : [selectedSource]; + state.lastAttemptFound = this.hasSearchResults(data); + this.displaySearchResults(missing, state, resultsDiv); } catch (error) { console.error('Model Linker: Search error:', error); @@ -4061,6 +4177,7 @@ class LinkerManagerDialog extends ComfyDialog { filename: modelId + '_' + versionId, category: '', is_urn: true, + sources: ['civitai'], model_id: modelId, version_id: versionId }; @@ -4109,22 +4226,32 @@ class LinkerManagerDialog extends ComfyDialog { /** * Display search results */ - displaySearchResults(missing, data, container) { + displaySearchResults(missing, state, container) { if (!container) return; - const popular = data.popular; - const modelListResult = data.model_list; - const hfResult = data.huggingface ? (Array.isArray(data.huggingface) ? data.huggingface[0] : data.huggingface) : null; - const civitaiResult = data.civitai ? (Array.isArray(data.civitai) ? data.civitai[0] : data.civitai) : null; + const results = state?.results || {}; + const popular = results.popular; + const modelListResult = results.model_list; + const hfResult = results.huggingface ? (Array.isArray(results.huggingface) ? results.huggingface[0] : results.huggingface) : null; + const civitaiResult = results.civitai ? (Array.isArray(results.civitai) ? results.civitai[0] : results.civitai) : null; const hasResults = popular || modelListResult || hfResult || civitaiResult; if (!hasResults) { - container.innerHTML = this.renderStatusMessage('No matches found online for this model.', 'warning'); + const searchedLabel = (state?.lastAttemptSources || []).map(source => this.getSearchSourceLabel(source)).join(', '); + container.innerHTML = this.renderStatusMessage( + searchedLabel ? `No matches found in ${searchedLabel}.` : 'No matches found online for this model.', + 'warning' + ); return; } let html = '
'; + if (state?.lastAttemptFound === false) { + const searchedLabel = (state.lastAttemptSources || []).map(source => this.getSearchSourceLabel(source)).join(', '); + html += this.renderStatusMessage(`No new matches found in ${searchedLabel}. Existing results are kept below.`, 'warning'); + } + // Popular models result (highest priority) if (popular) { html += `
`;