nodes/Azornes-Comfyui-Model-Resolver
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add source-filtered search and URN CivitAI support

Browse Source 
Backend: add explicit URN handling to resolve CivitAI model/version into a direct download source; normalize and accept a `sources` parameter (local, huggingface, civitai, all) and drive per-source searching. Reorder and tighten search logic (local/popular, model-list, HuggingFace, CivitAI) and include `searched_sources` in results and improved logging. Frontend: add per-missing-item search result cache and state, UI buttons to pick search sources, merge incremental results instead of overwriting, show selected-source helper text, pass selected source to the server, clear cache on refresh, and improve search status messages and UX. Overall: enables targeted online searches, preserves previous results, and surfaces direct CivitAI downloads for URNs.
This commit is contained in:
Dariusz L
2026-04-24 19:27:53 +02:00
parent 4fef6c66d9
commit c8272a8a45
2 changed files with 224 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+86 -21
View File
@@ -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
web/linker.js
+138 -11
View File
@@ -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 += `<div class="ml-download-section">`;
const searchSourcesId = `search-sources-${missing.node_id}-${missing.widget_index}`;
const searchInfoId = `search-source-info-${missing.node_id}-${missing.widget_index}`;
html += `<div id="${searchSourcesId}" style="display: flex; flex-wrap: wrap; gap: 6px; margin-bottom: 8px;">`;
html += `<button type="button" class="ml-btn ml-btn-secondary ml-btn-sm ml-search-source-btn" data-search-source="local">Local DB</button>`;
html += `<button type="button" class="ml-btn ml-btn-secondary ml-btn-sm ml-search-source-btn" data-search-source="huggingface">HuggingFace</button>`;
html += `<button type="button" class="ml-btn ml-btn-secondary ml-btn-sm ml-search-source-btn" data-search-source="civitai">CivitAI</button>`;
html += `<button type="button" class="ml-btn ml-btn-secondary ml-btn-sm ml-search-source-btn" data-search-source="all">All</button>`;
html += `</div>`;
html += `<button id="search-${missing.node_id}-${missing.widget_index}" class="ml-btn ml-btn-link">`;
html += `<span class="ml-btn-icon">🔍</span> Search Online`;
html += `</button>`;
html += `<div class="ml-download-info">Search HuggingFace & CivitAI</div>`;
html += `<div id="${searchInfoId}" class="ml-download-info">Selected source: Everything</div>`;
html += `</div>`;
html += `<div id="search-results-${missing.node_id}-${missing.widget_index}" style="margin-top: 8px; display: none;"></div>`;
}
@@ -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 = '<span style="color: #2196F3;">Searching HuggingFace and CivitAI...</span>';
resultsDiv.innerHTML = `<span style="color: #2196F3;">Searching ${selectedSourceLabel}...</span>`;
}
// 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 = '<div style="margin-top: 8px; display: flex; flex-direction: column; gap: 8px;">';
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 += `<div class="ml-status ml-status-success" style="flex-direction: column; align-items: flex-start;">`;