diff --git a/__init__.py b/__init__.py index 1fdd8d3..c95637d 100644 --- a/__init__.py +++ b/__init__.py @@ -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: diff --git a/core/sources/civitai.py b/core/sources/civitai.py index bc93e9d..a8d58f1 100644 --- a/core/sources/civitai.py +++ b/core/sources/civitai.py @@ -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 diff --git a/web/linker.js b/web/linker.js index 2be801b..796194c 100644 --- a/web/linker.js +++ b/web/linker.js @@ -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 {

API Tokens

-

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.

+

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.

@@ -3684,6 +3698,13 @@ class LinkerManagerDialog extends ComfyDialog {
Used for gated HuggingFace repos that need authorization during download.
+
+ +
+ +
+
Checks the first N CivitAI search results in order. Stops early when an exact 100% filename match is found.
+
@@ -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));