diff --git a/__init__.py b/__init__.py index 7628721..1fdd8d3 100644 --- a/__init__.py +++ b/__init__.py @@ -97,12 +97,16 @@ class ModelLinkerExtension: search_model_list, search_model_list_multiple, ) - from .core.sources.huggingface import search_huggingface_for_file + from .core.sources.huggingface import ( + search_huggingface_for_file, + clear_search_cache as clear_huggingface_search_cache, + ) from .core.sources.civitai import ( search_civitai_for_file, search_civitai, get_civitai_download_url, resolve_urn, + clear_search_cache as clear_civitai_search_cache, ) download_available = True @@ -804,6 +808,7 @@ class ModelLinkerExtension: category = data.get("category", "") # Handle both boolean and string forms is_urn_raw = data.get("is_urn", False) + civitai_session_token = data.get("civitai_session_token", "") is_urn = ( is_urn_raw if isinstance(is_urn_raw, bool) @@ -972,7 +977,8 @@ class ModelLinkerExtension: f"URN: No model_id/version_id, falling back to CivitAI search" ) civitai_results = search_civitai( - filename, model_type=category + filename, + model_type=category, ) log_search_result( "civitai/fallback", @@ -999,7 +1005,9 @@ class ModelLinkerExtension: results["found"] = True else: civitai_result = search_civitai_for_file( - filename, model_type=category + filename, + model_type=category, + session_token=civitai_session_token or None, ) log_search_result("civitai", civitai_result) if civitai_result: @@ -1015,6 +1023,18 @@ class ModelLinkerExtension: log_exception(f"Model Linker search error: {e}") return web.json_response({"error": str(e)}, status=500) + @routes.post("/model_linker/clear-search-cache") + async def clear_search_cache_route(request): + """Clear backend search caches after token/settings changes.""" + try: + clear_huggingface_search_cache() + clear_civitai_search_cache() + log_info("Cleared backend search caches") + return web.json_response({"success": True}) + except Exception as e: + log_exception(f"Clear search cache error: {e}") + return web.json_response({"error": str(e)}, status=500) + @routes.post("/model_linker/download") async def download_model(request): """Start downloading a model.""" diff --git a/core/sources/civitai.py b/core/sources/civitai.py index 623ad2a..bc93e9d 100644 --- a/core/sources/civitai.py +++ b/core/sources/civitai.py @@ -42,6 +42,12 @@ CIVITAI_TYPE_MAP = { } +def clear_search_cache(): + """Clear cached CivitAI search results.""" + global _search_cache + _search_cache.clear() + + def _build_civitai_result_from_version( model_id: int, model_name: str, @@ -210,6 +216,7 @@ def _extract_trpc_model_candidates( def _search_civitai_trpc_candidates( filename: str, model_type: Optional[str] = None, + session_token: Optional[str] = None, timeout: int = 15, limit: int = 5, ) -> List[Dict[str, Optional[int]]]: @@ -266,9 +273,11 @@ def _search_civitai_trpc_candidates( "x-client": "web", "x-client-version": "5.0.1657", } + if session_token: + headers["Cookie"] = f"__Secure-civitai-token={session_token}" log_info( - f"CivitAI tRPC search start: filename={filename}, model_type={model_type}, url={url}" + f"CivitAI tRPC search start: filename={filename}, model_type={model_type}, session_token={'yes' if session_token else 'no'}, url={url}" ) try: @@ -539,6 +548,7 @@ def search_civitai_for_file( api_key: Optional[str] = None, exact_only: bool = False, model_type: Optional[str] = None, + session_token: Optional[str] = None, ) -> Optional[Dict[str, Any]]: """ Search CivitAI for a specific model file. @@ -555,7 +565,11 @@ def search_civitai_for_file( """ global _search_cache - cache_key = f"civit_{filename}_exact{exact_only}" + 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}" + ) if cache_key in _search_cache: log_debug(f"CivitAI search cache hit for {filename} (exact_only={exact_only})") return _search_cache[cache_key] @@ -564,7 +578,7 @@ 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 + filename, model_type=model_type, session_token=session_token ) for candidate in trpc_candidates: model_id = candidate["model_id"] diff --git a/core/sources/huggingface.py b/core/sources/huggingface.py index c993878..b4c0bcb 100644 --- a/core/sources/huggingface.py +++ b/core/sources/huggingface.py @@ -24,6 +24,12 @@ HF_API_URL = "https://huggingface.co/api" _search_cache: Dict[str, Any] = {} +def clear_search_cache(): + """Clear cached HuggingFace search results.""" + global _search_cache + _search_cache.clear() + + def parse_huggingface_url(url: str) -> Optional[Dict[str, str]]: """ Parse a HuggingFace URL to extract repo and filename. diff --git a/web/linker.js b/web/linker.js index cf68098..fcceb47 100644 --- a/web/linker.js +++ b/web/linker.js @@ -3567,10 +3567,25 @@ class LinkerManagerDialog extends ComfyDialog { getStoredTokens() { return { civitai_key: localStorage.getItem('modelLinker.civitaiApiKey') || '', + civitai_session_token: localStorage.getItem('modelLinker.civitaiSessionToken') || '', hf_token: localStorage.getItem('modelLinker.huggingFaceToken') || '' }; } + async clearSearchCaches() { + this.searchResultCache.clear(); + try { + const response = await api.fetchApi('/model_linker/clear-search-cache', { + method: 'POST' + }); + if (!response.ok) { + throw new Error('Failed to clear backend search cache'); + } + } catch (error) { + console.error('Model Linker: Clear search cache error:', error); + } + } + displayOptions() { if (!this.contentElement) return; @@ -3579,13 +3594,18 @@ class LinkerManagerDialog extends ComfyDialog {

API Tokens

-

Stored locally in your browser. Used only for downloads that require authentication.

+

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.

Used for direct CivitAI downloads that otherwise return HTTP 401 or 403.
+
+ + +
Used only for CivitAI web search (civitai.red) to include results available to your logged-in session, including NSFW. Keep it private.
+
@@ -3602,28 +3622,34 @@ class LinkerManagerDialog extends ComfyDialog { `; const civitaiInput = this.contentElement.querySelector('#ml-options-civitai'); + const civitaiSessionInput = this.contentElement.querySelector('#ml-options-civitai-session'); const hfInput = this.contentElement.querySelector('#ml-options-hf'); const status = this.contentElement.querySelector('#ml-options-status'); const saveBtn = this.contentElement.querySelector('#ml-options-save'); const clearBtn = this.contentElement.querySelector('#ml-options-clear'); if (saveBtn) { - saveBtn.addEventListener('click', () => { + saveBtn.addEventListener('click', async () => { localStorage.setItem('modelLinker.civitaiApiKey', civitaiInput?.value || ''); + localStorage.setItem('modelLinker.civitaiSessionToken', civitaiSessionInput?.value || ''); localStorage.setItem('modelLinker.huggingFaceToken', hfInput?.value || ''); + await this.clearSearchCaches(); if (status) status.textContent = 'Tokens saved locally.'; - this.showNotification('API tokens saved', 'success'); + this.showNotification('API tokens saved and search cache cleared', 'success'); }); } if (clearBtn) { - clearBtn.addEventListener('click', () => { + clearBtn.addEventListener('click', async () => { localStorage.removeItem('modelLinker.civitaiApiKey'); + localStorage.removeItem('modelLinker.civitaiSessionToken'); localStorage.removeItem('modelLinker.huggingFaceToken'); if (civitaiInput) civitaiInput.value = ''; + if (civitaiSessionInput) civitaiSessionInput.value = ''; if (hfInput) hfInput.value = ''; + await this.clearSearchCaches(); if (status) status.textContent = 'Tokens cleared.'; - this.showNotification('API tokens cleared', 'info'); + this.showNotification('API tokens cleared and search cache reset', 'info'); }); } } @@ -5878,7 +5904,14 @@ class LinkerManagerDialog extends ComfyDialog { } // For URNs, include model_id and version_id for direct download - const searchData = { filename, category, is_urn: isUrn, sources: [selectedSource] }; + const tokens = this.getStoredTokens(); + const searchData = { + filename, + category, + is_urn: isUrn, + sources: [selectedSource], + civitai_session_token: tokens.civitai_session_token + }; if (isUrn && missing.urn) { searchData.model_id = missing.urn.model_id; searchData.version_id = missing.urn.version_id;