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

Add CivitAI session token & clear search cache

Browse Source 
Add support for an optional CivitAI session token and cache-management for search results. UI: new input for __Secure-civitai-token saved to localStorage (modelLinker.civitaiSessionToken) and cleared via the options dialog; saving/clearing tokens triggers a backend cache clear and updates notifications. Backend: add clear_search_cache() helpers in core.sources.huggingface and core.sources.civitai, register /model_linker/clear-search-cache route to flush both caches, and wire the civitai session token into search flows. CivitAI changes include sending the session token as a Cookie header for tRPC searches, incorporating session and model_type into the search cache key, and passing session_token through search_civitai_for_file and related calls. Purpose: ensure authenticated/session-aware search results (including NSFW/visibility differences) and avoid stale cached search results after token/settings changes.
This commit is contained in:
Dariusz L
2026-04-26 02:34:26 +02:00
parent cf0cd8491a
commit 233ed15222
4 changed files with 85 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+23 -3
View File
@@ -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."""
core/sources/civitai.py
+17 -3
View File
@@ -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"]
core/sources/huggingface.py
+6
View File
@@ -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.
web/linker.js
+39 -6
View File
@@ -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 {
<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. Used only for downloads that require authentication.</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.</p>
<div class="ml-options-grid">
<div class="ml-options-field">
<label for="ml-options-civitai" class="ml-options-label">CivitAI API Key</label>
<input id="ml-options-civitai" class="ml-options-input" type="password" placeholder="Paste CivitAI API key" value="${tokens.civitai_key}">
<div class="ml-options-help">Used for direct CivitAI downloads that otherwise return HTTP 401 or 403.</div>
</div>
<div class="ml-options-field">
<label for="ml-options-civitai-session" class="ml-options-label">CivitAI Session Token</label>
<input id="ml-options-civitai-session" class="ml-options-input" type="password" placeholder="Paste __Secure-civitai-token" value="${tokens.civitai_session_token}">
<div class="ml-options-help">Used only for CivitAI web search (civitai.red) to include results available to your logged-in session, including NSFW. Keep it private.</div>
</div>
<div class="ml-options-field">
<label for="ml-options-hf" class="ml-options-label">HuggingFace Token</label>
<input id="ml-options-hf" class="ml-options-input" type="password" placeholder="Paste HuggingFace token" value="${tokens.hf_token}">
@@ -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;