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

Add configurable CivitAI candidate limit

Browse Source 
Allow the number of CivitAI search candidates to be configured (default 5, bounds 1–20) and thread this limit through the frontend and backend. Backend: introduce DEFAULT/MAX constants, validate candidate_limit, include it in cache keys and logs, collect and deduplicate candidates from tRPC and HTML searches up to the limit, prefer exact 100% matches and otherwise pick the best-confidence result. Frontend: persist civitaiCandidateLimit in localStorage, add an input to options UI, clamp/save the value, include civitai_candidate_limit in search payloads, and apply suggested subfolder info from returned civitai_search_result. Misc: UI copy updated and minor log/message tweaks.
This commit is contained in:
Dariusz L
2026-04-26 03:09:03 +02:00
parent ef567a449e
commit f4f28fc450
3 changed files with 118 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+14 -1
View File
@@ -806,6 +806,18 @@ class ModelLinkerExtension:
data = await request.json()
filename = data.get("filename", "")
category = data.get("category", "")
civitai_candidate_limit_raw = data.get(
"civitai_candidate_limit", 5
)
try:
civitai_candidate_limit = int(
civitai_candidate_limit_raw
)
except (TypeError, ValueError):
civitai_candidate_limit = 5
civitai_candidate_limit = max(
1, min(civitai_candidate_limit, 20)
)
# Handle both boolean and string forms
is_urn_raw = data.get("is_urn", False)
civitai_session_token = data.get("civitai_session_token", "")
@@ -848,7 +860,7 @@ class ModelLinkerExtension:
search_civitai_source = "civitai" in normalized_sources
log_info(
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)}"
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)}, civitai_candidate_limit={civitai_candidate_limit}"
)
results = {
@@ -1008,6 +1020,7 @@ class ModelLinkerExtension:
filename,
model_type=category,
session_token=civitai_session_token or None,
candidate_limit=civitai_candidate_limit,
)
log_search_result("civitai", civitai_result)
if civitai_result:
core/sources/civitai.py
+54 -31
View File
@@ -27,6 +27,8 @@ CIVITAI_API_URL = "https://civitai.com/api/v1"
_search_cache: Dict[str, Any] = {}
_urn_cache: Dict[tuple[int, int], Dict[str, Any]] = {}
_hash_cache: Dict[str, Dict[str, Any]] = {}
DEFAULT_CIVITAI_CANDIDATE_LIMIT = 5
MAX_CIVITAI_CANDIDATE_LIMIT = 20
CIVITAI_TYPE_MAP = {
"checkpoint": "Checkpoint",
@@ -549,6 +551,7 @@ def search_civitai_for_file(
exact_only: bool = False,
model_type: Optional[str] = None,
session_token: Optional[str] = None,
candidate_limit: int = DEFAULT_CIVITAI_CANDIDATE_LIMIT,
) -> Optional[Dict[str, Any]]:
"""
Search CivitAI for a specific model file.
@@ -565,10 +568,11 @@ def search_civitai_for_file(
"""
global _search_cache
candidate_limit = max(1, min(int(candidate_limit), MAX_CIVITAI_CANDIDATE_LIMIT))
session_key = "session" if session_token else "anon"
model_type_key = str(model_type or "").lower()
cache_key = (
f"civit_{filename}_exact{exact_only}_type{model_type_key}_session{session_key}"
f"civit_{filename}_exact{exact_only}_type{model_type_key}_session{session_key}_limit{candidate_limit}"
)
if cache_key in _search_cache:
log_debug(f"CivitAI search cache hit for {filename} (exact_only={exact_only})")
@@ -578,35 +582,42 @@ def search_civitai_for_file(
# Prefer the CivitAI.red tRPC search because it is much closer to what
# the browser search UI returns than the broad public /models API.
trpc_candidates = _search_civitai_trpc_candidates(
filename, model_type=model_type, session_token=session_token
filename,
model_type=model_type,
session_token=session_token,
limit=candidate_limit,
)
for candidate in trpc_candidates:
model_id = candidate["model_id"]
version_id = candidate.get("version_id")
log_info(
f"CivitAI tRPC candidate check: model_id={model_id}, preferred_version_id={version_id}, filename={filename}"
)
result = _find_civitai_file_in_model(
model_id=model_id,
filename=filename,
api_key=api_key,
exact_only=exact_only,
preferred_version_id=version_id,
)
if result:
_search_cache[cache_key] = result
log_info(
f"Found {filename} via CivitAI tRPC fallback: model_id={result.get('model_id')}, version_id={result.get('version_id')}"
)
return result
candidates_to_check: List[Dict[str, Optional[int]]] = []
seen_candidates = set()
# Fallback: civitai.red HTML search, then inspect all submodels/versions
candidates = _search_civitai_red_candidates(filename)
for candidate in candidates:
def add_candidates(candidates: List[Dict[str, Optional[int]]]) -> None:
for candidate in candidates:
key = (candidate.get("model_id"), candidate.get("version_id"))
if key in seen_candidates:
continue
seen_candidates.add(key)
candidates_to_check.append(candidate)
if len(candidates_to_check) >= candidate_limit:
break
add_candidates(trpc_candidates)
if len(candidates_to_check) < candidate_limit:
html_candidates = _search_civitai_red_candidates(
filename, limit=candidate_limit
)
add_candidates(html_candidates)
else:
html_candidates = []
best_result = None
best_confidence = 0.0
for candidate in candidates_to_check:
model_id = candidate["model_id"]
version_id = candidate.get("version_id")
log_info(
f"CivitAI.red candidate check: model_id={model_id}, preferred_version_id={version_id}, filename={filename}"
f"CivitAI candidate check: model_id={model_id}, preferred_version_id={version_id}, filename={filename}"
)
result = _find_civitai_file_in_model(
model_id=model_id,
@@ -616,16 +627,28 @@ def search_civitai_for_file(
preferred_version_id=version_id,
)
if result:
_search_cache[cache_key] = result
log_info(
f"Found {filename} via CivitAI.red fallback: model_id={result.get('model_id')}, version_id={result.get('version_id')}"
)
return result
confidence = float(result.get("confidence") or 0.0)
if confidence >= 100.0:
_search_cache[cache_key] = result
log_info(
f"Found exact CivitAI match for {filename}: model_id={result.get('model_id')}, version_id={result.get('version_id')}, candidate_limit={candidate_limit}"
)
return result
if confidence > best_confidence:
best_result = result
best_confidence = confidence
if best_result:
_search_cache[cache_key] = best_result
log_info(
f"Found best CivitAI match for {filename}: model_id={best_result.get('model_id')}, version_id={best_result.get('version_id')}, confidence={best_confidence}, candidate_limit={candidate_limit}"
)
return best_result
# Not found
_search_cache[cache_key] = None
log_info(
f"CivitAI search no result: filename={filename}, trpc_candidates={len(trpc_candidates)}, html_candidates={len(candidates) if 'candidates' in locals() else 0}"
f"CivitAI search no result: filename={filename}, trpc_candidates={len(trpc_candidates)}, html_candidates={len(html_candidates)}, checked_candidates={len(candidates_to_check)}, candidate_limit={candidate_limit}"
)
return None
web/linker.js
+50 -6
View File
@@ -3412,6 +3412,7 @@ class LinkerManagerDialog extends ComfyDialog {
const civitaiData = {
...(missing?.civitai_info || {}),
...(missing?.civitai_search_result || {}),
...(missing?.download_source || {})
};
const baseModel = civitaiData.base_model || '';
@@ -3473,6 +3474,13 @@ class LinkerManagerDialog extends ComfyDialog {
}
}
applySearchResultSuggestion(missing) {
const categoryEl = this.contentElement?.querySelector(`#download-category-${missing.node_id}-${missing.widget_index}`);
const subfolderEl = this.contentElement?.querySelector(`#download-subfolder-${missing.node_id}-${missing.widget_index}`);
if (!categoryEl || !subfolderEl) return;
this.applySuggestedCivitaiSubfolder(missing, categoryEl, subfolderEl);
}
async ensureDownloadSubfoldersLoaded(category = '') {
const key = (category || '').trim().toLowerCase();
if (!key) return [];
@@ -3598,10 +3606,16 @@ class LinkerManagerDialog extends ComfyDialog {
}
getStoredTokens() {
const civitaiCandidateLimitRaw = parseInt(localStorage.getItem('modelLinker.civitaiCandidateLimit') || '5', 10);
const civitai_candidate_limit = Number.isFinite(civitaiCandidateLimitRaw)
? Math.min(20, Math.max(1, civitaiCandidateLimitRaw))
: 5;
return {
civitai_key: localStorage.getItem('modelLinker.civitaiApiKey') || '',
civitai_session_token: localStorage.getItem('modelLinker.civitaiSessionToken') || '',
hf_token: localStorage.getItem('modelLinker.huggingFaceToken') || ''
hf_token: localStorage.getItem('modelLinker.huggingFaceToken') || '',
civitai_candidate_limit
};
}
@@ -3627,7 +3641,7 @@ class LinkerManagerDialog extends ComfyDialog {
<div class="ml-options-wrap">
<div class="ml-options-card">
<h3 class="ml-options-title">API Tokens</h3>
<p class="ml-options-subtitle">Stored locally in your browser. API keys are used for downloads. Session token can improve CivitAI web search results, including NSFW items visible to your logged-in account.</p>
<p class="ml-options-subtitle">Stored locally in your browser. API keys are used for downloads. Session token can improve CivitAI web search results, including NSFW items visible to your logged-in account. You can also choose how many CivitAI search results should be inspected before the best match is picked.</p>
<div class="ml-options-grid">
<div class="ml-options-field">
<label for="ml-options-civitai" class="ml-options-label">CivitAI API Key</label>
@@ -3684,6 +3698,13 @@ class LinkerManagerDialog extends ComfyDialog {
</div>
<div class="ml-options-help">Used for gated HuggingFace repos that need authorization during download.</div>
</div>
<div class="ml-options-field">
<label for="ml-options-civitai-limit" class="ml-options-label">CivitAI Models To Inspect</label>
<div class="ml-options-input-row">
<input id="ml-options-civitai-limit" class="ml-options-input" type="number" min="1" max="20" step="1" value="${tokens.civitai_candidate_limit}">
</div>
<div class="ml-options-help">Checks the first N CivitAI search results in order. Stops early when an exact 100% filename match is found.</div>
</div>
</div>
<div class="ml-options-actions">
<button id="ml-options-save" class="ml-btn ml-btn-primary">Save Tokens</button>
@@ -3699,6 +3720,7 @@ class LinkerManagerDialog extends ComfyDialog {
const civitaiToggle = this.contentElement.querySelector('#ml-options-civitai-toggle');
const civitaiSessionToggle = this.contentElement.querySelector('#ml-options-civitai-session-toggle');
const hfToggle = this.contentElement.querySelector('#ml-options-hf-toggle');
const civitaiLimitInput = this.contentElement.querySelector('#ml-options-civitai-limit');
const status = this.contentElement.querySelector('#ml-options-status');
const saveBtn = this.contentElement.querySelector('#ml-options-save');
const civitaiHelpBtn = this.contentElement.querySelector('#ml-options-civitai-help');
@@ -3745,12 +3767,20 @@ class LinkerManagerDialog extends ComfyDialog {
if (saveBtn) {
saveBtn.addEventListener('click', async () => {
const civitaiCandidateLimitRaw = parseInt(civitaiLimitInput?.value || `${tokens.civitai_candidate_limit}`, 10);
const civitaiCandidateLimit = Number.isFinite(civitaiCandidateLimitRaw)
? Math.min(20, Math.max(1, civitaiCandidateLimitRaw))
: 5;
localStorage.setItem('modelLinker.civitaiApiKey', civitaiInput?.value || '');
localStorage.setItem('modelLinker.civitaiSessionToken', civitaiSessionInput?.value || '');
localStorage.setItem('modelLinker.huggingFaceToken', hfInput?.value || '');
localStorage.setItem('modelLinker.civitaiCandidateLimit', `${civitaiCandidateLimit}`);
if (civitaiLimitInput) {
civitaiLimitInput.value = `${civitaiCandidateLimit}`;
}
await this.clearSearchCaches();
if (status) status.textContent = 'Tokens saved locally.';
this.showNotification('API tokens saved and search cache cleared', 'success');
if (status) status.textContent = 'Options saved locally.';
this.showNotification('Options saved and search cache cleared', 'success');
});
}
@@ -6027,7 +6057,8 @@ class LinkerManagerDialog extends ComfyDialog {
category,
is_urn: isUrn,
sources: [selectedSource],
civitai_session_token: tokens.civitai_session_token
civitai_session_token: tokens.civitai_session_token,
civitai_candidate_limit: tokens.civitai_candidate_limit
};
if (isUrn && missing.urn) {
searchData.model_id = missing.urn.model_id;
@@ -6051,7 +6082,18 @@ class LinkerManagerDialog extends ComfyDialog {
state.results = this.mergeSearchResults(state.results, data);
state.lastAttemptSources = Array.isArray(data.searched_sources) ? data.searched_sources : [selectedSource];
state.lastAttemptFound = this.hasSearchResults(data);
if (data.civitai) {
missing.civitai_search_result = {
base_model: data.civitai.base_model,
tags: data.civitai.tags || [],
trained_words: data.civitai.trained_words || [],
filename: data.civitai.filename,
name: data.civitai.name,
type: data.civitai.type
};
}
this.displaySearchResults(missing, state, resultsDiv);
this.applySearchResultSuggestion(missing);
} catch (error) {
console.error('Model Linker: Search error:', error);
@@ -6077,13 +6119,15 @@ class LinkerManagerDialog extends ComfyDialog {
}
try {
const tokens = this.getStoredTokens();
const payload = {
filename: modelId + '_' + versionId,
category: '',
is_urn: true,
sources: ['civitai'],
model_id: modelId,
version_id: versionId
version_id: versionId,
civitai_candidate_limit: tokens.civitai_candidate_limit
};
console.log('resolveUrnAsync payload:', JSON.stringify(payload));