feat: wire progress callbacks into backend search sources
This commit is contained in:
+69
-3
@@ -1437,6 +1437,51 @@ class ModelResolverExtension:
|
||||
marked["requested_base_model"] = base_model_context
|
||||
return marked
|
||||
|
||||
def make_source_progress_callback(
|
||||
source_key,
|
||||
percent_min=None,
|
||||
percent_max=None,
|
||||
):
|
||||
def source_progress_callback(payload):
|
||||
if not isinstance(payload, dict):
|
||||
return
|
||||
|
||||
progress_payload = dict(payload)
|
||||
stage = progress_payload.pop("stage", "running")
|
||||
message = progress_payload.pop(
|
||||
"message", "Searching..."
|
||||
)
|
||||
percent = progress_payload.pop("percent", None)
|
||||
status = progress_payload.pop("status", "running")
|
||||
progress_payload.pop("source", None)
|
||||
|
||||
if (
|
||||
percent is not None
|
||||
and percent_min is not None
|
||||
and percent_max is not None
|
||||
):
|
||||
try:
|
||||
normalized_percent = max(
|
||||
0.0, min(100.0, float(percent))
|
||||
)
|
||||
percent = percent_min + (
|
||||
normalized_percent / 100.0
|
||||
) * (percent_max - percent_min)
|
||||
except (TypeError, ValueError):
|
||||
percent = None
|
||||
|
||||
update_search_progress(
|
||||
progress_id,
|
||||
source_key,
|
||||
stage,
|
||||
message,
|
||||
percent,
|
||||
status=status,
|
||||
**progress_payload,
|
||||
)
|
||||
|
||||
return source_progress_callback
|
||||
|
||||
def search_local_sources():
|
||||
source_results = {"popular": None, "model_list": None}
|
||||
source_found = False
|
||||
@@ -1527,13 +1572,16 @@ class ModelResolverExtension:
|
||||
use_comfy_org_fallback=hf_use_comfy_org_fallback,
|
||||
use_brave_fallback=hf_use_brave_fallback,
|
||||
force_refresh=force_search,
|
||||
progress_callback=make_source_progress_callback(
|
||||
"huggingface"
|
||||
),
|
||||
)
|
||||
update_search_progress(
|
||||
progress_id,
|
||||
"huggingface",
|
||||
"parse",
|
||||
"Checking HuggingFace result",
|
||||
86,
|
||||
"done",
|
||||
"HuggingFace checked",
|
||||
92,
|
||||
)
|
||||
log_search_result("huggingface", hf_result)
|
||||
return {"huggingface": hf_result}, bool(hf_result)
|
||||
@@ -1689,6 +1737,9 @@ class ModelResolverExtension:
|
||||
candidate_limit=civitai_candidate_limit,
|
||||
use_trpc_search=civitai_use_trpc_search,
|
||||
use_html_fallback=civitai_use_html_fallback,
|
||||
progress_callback=make_source_progress_callback(
|
||||
"civitai"
|
||||
),
|
||||
)
|
||||
log_search_result("civitai", civitai_result)
|
||||
if not civitai_result and base_model_context:
|
||||
@@ -1715,6 +1766,9 @@ class ModelResolverExtension:
|
||||
candidate_limit=civitai_candidate_limit,
|
||||
use_trpc_search=civitai_use_trpc_search,
|
||||
use_html_fallback=civitai_use_html_fallback,
|
||||
progress_callback=make_source_progress_callback(
|
||||
"civitai", 72, 92
|
||||
),
|
||||
)
|
||||
log_search_result(
|
||||
"civitai/any_model",
|
||||
@@ -1806,6 +1860,9 @@ class ModelResolverExtension:
|
||||
model_type=category,
|
||||
base_model_context=base_model_context or None,
|
||||
limit=civarchive_candidate_limit,
|
||||
progress_callback=make_source_progress_callback(
|
||||
"civarchive"
|
||||
),
|
||||
)
|
||||
log_search_result("civarchive", civarchive_result)
|
||||
if not civarchive_result and base_model_context:
|
||||
@@ -1829,6 +1886,9 @@ class ModelResolverExtension:
|
||||
model_type=category,
|
||||
base_model_context=None,
|
||||
limit=civarchive_candidate_limit,
|
||||
progress_callback=make_source_progress_callback(
|
||||
"civarchive", 72, 92
|
||||
),
|
||||
)
|
||||
log_search_result(
|
||||
"civarchive/any_model",
|
||||
@@ -1883,6 +1943,9 @@ class ModelResolverExtension:
|
||||
filename,
|
||||
model_type=category,
|
||||
base_model_context=base_model_context or None,
|
||||
progress_callback=make_source_progress_callback(
|
||||
"lora_manager_archive"
|
||||
),
|
||||
)
|
||||
)
|
||||
log_search_result(
|
||||
@@ -1910,6 +1973,9 @@ class ModelResolverExtension:
|
||||
filename,
|
||||
model_type=category,
|
||||
base_model_context=None,
|
||||
progress_callback=make_source_progress_callback(
|
||||
"lora_manager_archive", 72, 92
|
||||
),
|
||||
)
|
||||
)
|
||||
log_search_result(
|
||||
|
||||
+103
-3
@@ -9,7 +9,7 @@ import json
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
from typing import Any, Callable, Dict, List, Optional, Tuple
|
||||
from urllib.parse import parse_qs, urljoin, urlparse
|
||||
|
||||
import requests
|
||||
@@ -85,6 +85,28 @@ def clear_search_cache():
|
||||
_search_cache.clear()
|
||||
|
||||
|
||||
def _report_progress(
|
||||
progress_callback: Optional[Callable[[Dict[str, Any]], None]],
|
||||
stage: str,
|
||||
message: str,
|
||||
percent: Optional[float] = None,
|
||||
**extra: Any,
|
||||
) -> None:
|
||||
if not progress_callback:
|
||||
return
|
||||
|
||||
payload = {"stage": stage, "message": message}
|
||||
if percent is not None:
|
||||
payload["percent"] = percent
|
||||
if extra:
|
||||
payload.update(extra)
|
||||
|
||||
try:
|
||||
progress_callback(payload)
|
||||
except Exception as e:
|
||||
log_debug(f"CivArchive progress callback failed: {e}")
|
||||
|
||||
|
||||
def is_civarchive_available() -> bool:
|
||||
"""Return True when the CivArchive provider can be offered."""
|
||||
return True
|
||||
@@ -1510,6 +1532,7 @@ def search_civarchive_for_file(
|
||||
base_model_context: Optional[str] = None,
|
||||
exact_only: bool = False,
|
||||
limit: int = DEFAULT_CIVARCHIVE_CANDIDATE_LIMIT,
|
||||
progress_callback: Optional[Callable[[Dict[str, Any]], None]] = None,
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
"""
|
||||
Search CivArchive for the best downloadable match for a model filename.
|
||||
@@ -1524,16 +1547,36 @@ def search_civarchive_for_file(
|
||||
f"file::{normalized_filename.lower()}::{model_type or ''}::{base_model_key}::{exact_only}::{limit}"
|
||||
)
|
||||
if cache_key in _search_cache:
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"cache",
|
||||
"Using CivArchive cache",
|
||||
86,
|
||||
)
|
||||
return _search_cache[cache_key]
|
||||
|
||||
sha256 = _extract_sha256(normalized_filename)
|
||||
if sha256:
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"hash",
|
||||
"Resolving CivArchive hash",
|
||||
40,
|
||||
sha256=sha256,
|
||||
)
|
||||
result = resolve_civarchive_by_hash(
|
||||
sha256,
|
||||
query=normalized_filename,
|
||||
exact_only=exact_only,
|
||||
)
|
||||
_search_cache[cache_key] = result
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"done",
|
||||
"CivArchive checked",
|
||||
92,
|
||||
found=bool(result),
|
||||
)
|
||||
return result
|
||||
|
||||
best_match = None
|
||||
@@ -1545,12 +1588,44 @@ def search_civarchive_for_file(
|
||||
)
|
||||
|
||||
try:
|
||||
for search_query in _build_search_queries(normalized_filename):
|
||||
search_queries = _build_search_queries(normalized_filename)
|
||||
query_count = max(1, len(search_queries))
|
||||
for query_index, search_query in enumerate(search_queries, start=1):
|
||||
query_percent = 30 + ((query_index - 1) / query_count) * 18
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"query",
|
||||
f"CivArchive query {query_index}/{query_count}",
|
||||
query_percent,
|
||||
query=search_query,
|
||||
query_index=query_index,
|
||||
query_count=query_count,
|
||||
)
|
||||
candidates = _search_page(search_query, model_type=model_type)
|
||||
for candidate in candidates:
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"candidates",
|
||||
f"CivArchive candidates: {len(candidates)}",
|
||||
min(55, query_percent + 10),
|
||||
candidate_count=len(candidates),
|
||||
checked_candidate_count=len(seen),
|
||||
)
|
||||
candidate_count = len(candidates)
|
||||
for candidate_index, candidate in enumerate(candidates, start=1):
|
||||
identity = _candidate_identity(candidate)
|
||||
if identity in seen:
|
||||
continue
|
||||
checked_index = min(len(seen) + 1, limit)
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"candidate",
|
||||
f"Resolving CivArchive candidate {checked_index}/{limit}",
|
||||
58 + (checked_index / max(1, limit)) * 30,
|
||||
candidate_index=candidate_index,
|
||||
candidate_count=candidate_count,
|
||||
checked_candidate_count=checked_index,
|
||||
candidate_limit=limit,
|
||||
)
|
||||
seen.add(identity)
|
||||
|
||||
resolved = _resolve_search_candidate(
|
||||
@@ -1598,6 +1673,14 @@ def search_civarchive_for_file(
|
||||
|
||||
if confidence == 100.0 and base_model_matches:
|
||||
_search_cache[cache_key] = best_match
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"found",
|
||||
"Found CivArchive match",
|
||||
92,
|
||||
confidence=confidence,
|
||||
checked_candidate_count=len(seen),
|
||||
)
|
||||
return best_match
|
||||
|
||||
if len(seen) >= limit:
|
||||
@@ -1612,6 +1695,13 @@ def search_civarchive_for_file(
|
||||
raise
|
||||
except Exception as e:
|
||||
log_exception(f"CivArchive search error for {normalized_filename}: {e}")
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"error",
|
||||
f"CivArchive search error: {e}",
|
||||
100,
|
||||
status="error",
|
||||
)
|
||||
return None
|
||||
|
||||
if best_match and best_confidence < 40:
|
||||
@@ -1622,4 +1712,14 @@ def search_civarchive_for_file(
|
||||
best_match = None
|
||||
|
||||
_search_cache[cache_key] = best_match
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"found" if best_match else "done",
|
||||
"Found CivArchive match" if best_match else "CivArchive checked",
|
||||
92,
|
||||
found=bool(best_match),
|
||||
confidence=best_confidence if best_match else None,
|
||||
checked_candidate_count=len(seen),
|
||||
candidate_limit=limit,
|
||||
)
|
||||
return best_match
|
||||
|
||||
+112
-2
@@ -9,7 +9,7 @@ import re
|
||||
import json
|
||||
import hashlib
|
||||
import requests
|
||||
from typing import Dict, Any, Optional, List
|
||||
from typing import Dict, Any, Optional, List, Callable
|
||||
from urllib.parse import urlparse, parse_qs, quote
|
||||
|
||||
from ..matcher import calculate_similarity_with_normalization, normalize_filename
|
||||
@@ -66,6 +66,28 @@ def clear_search_cache():
|
||||
_hash_cache.clear()
|
||||
|
||||
|
||||
def _report_progress(
|
||||
progress_callback: Optional[Callable[[Dict[str, Any]], None]],
|
||||
stage: str,
|
||||
message: str,
|
||||
percent: Optional[float] = None,
|
||||
**extra: Any,
|
||||
) -> None:
|
||||
if not progress_callback:
|
||||
return
|
||||
|
||||
payload = {"stage": stage, "message": message}
|
||||
if percent is not None:
|
||||
payload["percent"] = percent
|
||||
if extra:
|
||||
payload.update(extra)
|
||||
|
||||
try:
|
||||
progress_callback(payload)
|
||||
except Exception as e:
|
||||
log_debug(f"CivitAI progress callback failed: {e}")
|
||||
|
||||
|
||||
def check_civitai_session_token(session_token: Optional[str]) -> Dict[str, Any]:
|
||||
"""Check whether a CivitAI browser session token is accepted by civitai.com."""
|
||||
token = (session_token or "").strip()
|
||||
@@ -932,6 +954,7 @@ def search_civitai_for_file(
|
||||
candidate_limit: int = DEFAULT_CIVITAI_CANDIDATE_LIMIT,
|
||||
use_trpc_search: bool = True,
|
||||
use_html_fallback: bool = True,
|
||||
progress_callback: Optional[Callable[[Dict[str, Any]], None]] = None,
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
"""
|
||||
Search CivitAI for a specific model file.
|
||||
@@ -958,11 +981,24 @@ def search_civitai_for_file(
|
||||
)
|
||||
if cache_key in _search_cache:
|
||||
log_debug(f"CivitAI search cache hit for {filename} (exact_only={exact_only})")
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"cache",
|
||||
"Using CivitAI cache",
|
||||
86,
|
||||
)
|
||||
return _search_cache[cache_key]
|
||||
|
||||
try:
|
||||
# Prefer the CivitAI.red tRPC search because it is much closer to what
|
||||
# the browser search UI returns than the broad public /models API.
|
||||
if use_trpc_search:
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"trpc",
|
||||
"Searching CivitAI tRPC",
|
||||
32,
|
||||
)
|
||||
trpc_candidates = (
|
||||
_search_civitai_trpc_candidates(
|
||||
filename,
|
||||
@@ -973,6 +1009,13 @@ def search_civitai_for_file(
|
||||
if use_trpc_search
|
||||
else []
|
||||
)
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"trpc_candidates",
|
||||
f"CivitAI tRPC candidates: {len(trpc_candidates)}",
|
||||
42,
|
||||
candidate_count=len(trpc_candidates),
|
||||
)
|
||||
candidates_to_check: List[Dict[str, Optional[int]]] = []
|
||||
seen_candidates = set()
|
||||
|
||||
@@ -989,19 +1032,51 @@ def search_civitai_for_file(
|
||||
add_candidates(trpc_candidates)
|
||||
|
||||
if use_html_fallback and len(candidates_to_check) < candidate_limit:
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"html",
|
||||
"Searching CivitAI HTML fallback",
|
||||
50,
|
||||
)
|
||||
html_candidates = _search_civitai_red_candidates(
|
||||
filename, limit=candidate_limit
|
||||
)
|
||||
add_candidates(html_candidates)
|
||||
else:
|
||||
html_candidates = []
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"html_candidates",
|
||||
f"CivitAI HTML candidates: {len(html_candidates)}",
|
||||
58,
|
||||
candidate_count=len(html_candidates),
|
||||
)
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"candidates",
|
||||
f"CivitAI candidates to check: {len(candidates_to_check)}",
|
||||
60,
|
||||
candidate_count=len(candidates_to_check),
|
||||
candidate_limit=candidate_limit,
|
||||
)
|
||||
|
||||
best_result = None
|
||||
best_confidence = 0.0
|
||||
|
||||
for candidate in candidates_to_check:
|
||||
total_candidates = len(candidates_to_check)
|
||||
for candidate_index, candidate in enumerate(candidates_to_check, start=1):
|
||||
model_id = candidate["model_id"]
|
||||
version_id = candidate.get("version_id")
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"candidate",
|
||||
f"Checking CivitAI candidate {candidate_index}/{total_candidates}",
|
||||
62 + (candidate_index / max(1, total_candidates)) * 24,
|
||||
model_id=model_id,
|
||||
version_id=version_id,
|
||||
candidate_index=candidate_index,
|
||||
candidate_count=total_candidates,
|
||||
)
|
||||
log_info(
|
||||
f"CivitAI candidate check: model_id={model_id}, preferred_version_id={version_id}, filename={filename}"
|
||||
)
|
||||
@@ -1017,6 +1092,15 @@ def search_civitai_for_file(
|
||||
confidence = float(result.get("confidence") or 0.0)
|
||||
if confidence >= 100.0:
|
||||
_search_cache[cache_key] = result
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"found",
|
||||
"Found exact CivitAI match",
|
||||
92,
|
||||
model_id=result.get("model_id"),
|
||||
version_id=result.get("version_id"),
|
||||
confidence=confidence,
|
||||
)
|
||||
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}"
|
||||
)
|
||||
@@ -1027,6 +1111,15 @@ def search_civitai_for_file(
|
||||
|
||||
if best_result:
|
||||
_search_cache[cache_key] = best_result
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"found",
|
||||
"Found CivitAI match",
|
||||
92,
|
||||
model_id=best_result.get("model_id"),
|
||||
version_id=best_result.get("version_id"),
|
||||
confidence=best_confidence,
|
||||
)
|
||||
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}"
|
||||
)
|
||||
@@ -1034,6 +1127,16 @@ def search_civitai_for_file(
|
||||
|
||||
# Not found
|
||||
_search_cache[cache_key] = None
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"done",
|
||||
"CivitAI checked",
|
||||
92,
|
||||
trpc_candidate_count=len(trpc_candidates),
|
||||
html_candidate_count=len(html_candidates),
|
||||
checked_candidate_count=len(candidates_to_check),
|
||||
candidate_limit=candidate_limit,
|
||||
)
|
||||
log_info(
|
||||
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}"
|
||||
)
|
||||
@@ -1041,6 +1144,13 @@ def search_civitai_for_file(
|
||||
|
||||
except Exception as e:
|
||||
log_exception(f"CivitAI search error for {filename}: {e}")
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"error",
|
||||
f"CivitAI search error: {e}",
|
||||
100,
|
||||
status="error",
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
|
||||
+169
-6
@@ -10,7 +10,7 @@ import re
|
||||
import threading
|
||||
import time
|
||||
import requests
|
||||
from typing import Dict, Any, Optional, List
|
||||
from typing import Dict, Any, Optional, List, Callable
|
||||
from urllib.parse import urlparse, quote
|
||||
|
||||
from ..log_system.log_funcs import (
|
||||
@@ -47,6 +47,28 @@ def clear_search_cache():
|
||||
_author_index_cache.clear()
|
||||
|
||||
|
||||
def _report_progress(
|
||||
progress_callback: Optional[Callable[[Dict[str, Any]], None]],
|
||||
stage: str,
|
||||
message: str,
|
||||
percent: Optional[float] = None,
|
||||
**extra: Any,
|
||||
) -> None:
|
||||
if not progress_callback:
|
||||
return
|
||||
|
||||
payload = {"stage": stage, "message": message}
|
||||
if percent is not None:
|
||||
payload["percent"] = percent
|
||||
if extra:
|
||||
payload.update(extra)
|
||||
|
||||
try:
|
||||
progress_callback(payload)
|
||||
except Exception as e:
|
||||
log_debug(f"HuggingFace progress callback failed: {e}")
|
||||
|
||||
|
||||
def check_huggingface_token(token: Optional[str]) -> Dict[str, Any]:
|
||||
"""Check whether a HuggingFace access token is accepted."""
|
||||
value = (token or "").strip()
|
||||
@@ -669,6 +691,7 @@ def search_huggingface_for_file(
|
||||
use_comfy_org_fallback: bool = True,
|
||||
use_brave_fallback: bool = True,
|
||||
force_refresh: bool = False,
|
||||
progress_callback: Optional[Callable[[Dict[str, Any]], None]] = None,
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
"""
|
||||
Search HuggingFace for a specific model file.
|
||||
@@ -696,6 +719,12 @@ def search_huggingface_for_file(
|
||||
cache_key = f"hf_{filename}_exact{exact_only}_{token_key}_{brave_key}_{methods_key}"
|
||||
if cache_key in _search_cache:
|
||||
log_debug(f"HuggingFace cache hit file={filename} exact={exact_only}")
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"cache",
|
||||
"Using HuggingFace cache",
|
||||
86,
|
||||
)
|
||||
return _search_cache[cache_key]
|
||||
|
||||
try:
|
||||
@@ -710,7 +739,17 @@ def search_huggingface_for_file(
|
||||
search_queries = _build_huggingface_search_queries(filename)
|
||||
|
||||
if use_api_search:
|
||||
for search_query in search_queries:
|
||||
total_queries = max(1, len(search_queries))
|
||||
for query_index, search_query in enumerate(search_queries, start=1):
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"api_search",
|
||||
f"HuggingFace API query {query_index}/{total_queries}",
|
||||
32 + ((query_index - 1) / total_queries) * 10,
|
||||
query=search_query,
|
||||
query_index=query_index,
|
||||
query_count=total_queries,
|
||||
)
|
||||
search_url = f"{HF_API_URL}/models?search={quote(search_query)}&limit=20"
|
||||
response = requests.get(search_url, headers=headers, timeout=10)
|
||||
if response.status_code != 200:
|
||||
@@ -728,7 +767,32 @@ def search_huggingface_for_file(
|
||||
seen_repos.add(repo_id)
|
||||
repos_to_check.append(repo_id)
|
||||
|
||||
for repo_id in repos_to_check:
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"api_candidates",
|
||||
f"HuggingFace API candidates: {len(repos_to_check)} repos",
|
||||
44,
|
||||
candidate_count=len(repos_to_check),
|
||||
)
|
||||
else:
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"api_skip",
|
||||
"Skipping HuggingFace API search",
|
||||
44,
|
||||
)
|
||||
|
||||
repo_count = len(repos_to_check)
|
||||
for repo_index, repo_id in enumerate(repos_to_check, start=1):
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"repo_files",
|
||||
f"Checking HF repo {repo_index}/{repo_count}",
|
||||
48 + (repo_index / max(1, repo_count)) * 14,
|
||||
repo=repo_id,
|
||||
candidate_index=repo_index,
|
||||
candidate_count=repo_count,
|
||||
)
|
||||
log_debug(f"HuggingFace check repo={repo_id} file={filename}")
|
||||
files = _get_repo_tree(repo_id, headers=headers)
|
||||
if not files:
|
||||
@@ -740,6 +804,14 @@ def search_huggingface_for_file(
|
||||
)
|
||||
if result:
|
||||
_search_cache[cache_key] = result
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"found",
|
||||
"Found HuggingFace match",
|
||||
92,
|
||||
repo=repo_id,
|
||||
match_type=result.get("match_type"),
|
||||
)
|
||||
log_info(
|
||||
f"HuggingFace found match={result['match_type']} file={filename} repo={repo_id} path={result['path']}"
|
||||
)
|
||||
@@ -748,7 +820,17 @@ def search_huggingface_for_file(
|
||||
author_fallback_repos = []
|
||||
author_fallback_repo_count = 0
|
||||
if use_comfy_org_fallback:
|
||||
for author in HF_AUTHOR_FALLBACKS:
|
||||
total_authors = max(1, len(HF_AUTHOR_FALLBACKS))
|
||||
for author_index, author in enumerate(HF_AUTHOR_FALLBACKS, start=1):
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"author_index",
|
||||
f"Checking HF fallback index {author_index}/{total_authors}",
|
||||
64 + ((author_index - 1) / total_authors) * 6,
|
||||
author=author,
|
||||
author_index=author_index,
|
||||
author_count=total_authors,
|
||||
)
|
||||
index = _get_author_index(
|
||||
author,
|
||||
headers={},
|
||||
@@ -761,6 +843,14 @@ def search_huggingface_for_file(
|
||||
author_fallback_repo_count += int(index.get("repo_count") or 0)
|
||||
if result:
|
||||
_search_cache[cache_key] = result
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"found",
|
||||
"Found HuggingFace fallback match",
|
||||
92,
|
||||
repo=result.get("repo_id"),
|
||||
match_type=result.get("match_type"),
|
||||
)
|
||||
log_info(
|
||||
f"HuggingFace found source=author_index match={result['match_type']} file={filename} repo={result['repo_id']} path={result['path']}"
|
||||
)
|
||||
@@ -778,8 +868,25 @@ def search_huggingface_for_file(
|
||||
seen_repos.add(repo_id)
|
||||
author_fallback_repos.append(repo_id)
|
||||
author_fallback_repo_count += 1
|
||||
else:
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"author_skip",
|
||||
"Skipping HuggingFace fallback index",
|
||||
70,
|
||||
)
|
||||
|
||||
for repo_id in author_fallback_repos:
|
||||
fallback_repo_count = len(author_fallback_repos)
|
||||
for repo_index, repo_id in enumerate(author_fallback_repos, start=1):
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"author_repo",
|
||||
f"Checking HF fallback repo {repo_index}/{fallback_repo_count}",
|
||||
70 + (repo_index / max(1, fallback_repo_count)) * 8,
|
||||
repo=repo_id,
|
||||
candidate_index=repo_index,
|
||||
candidate_count=fallback_repo_count,
|
||||
)
|
||||
log_debug(
|
||||
f"HuggingFace author_fallback check repo={repo_id} file={filename}"
|
||||
)
|
||||
@@ -792,22 +899,54 @@ def search_huggingface_for_file(
|
||||
)
|
||||
if result:
|
||||
_search_cache[cache_key] = result
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"found",
|
||||
"Found HuggingFace fallback match",
|
||||
92,
|
||||
repo=repo_id,
|
||||
match_type=result.get("match_type"),
|
||||
)
|
||||
log_info(
|
||||
f"HuggingFace found source=author_fallback match={result['match_type']} file={filename} repo={repo_id} path={result['path']}"
|
||||
)
|
||||
return result
|
||||
|
||||
if use_brave_fallback:
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"brave_search",
|
||||
"Searching Brave fallback",
|
||||
80,
|
||||
)
|
||||
brave_candidates = (
|
||||
_search_brave_for_huggingface_candidates(filename, brave_api_key or "")
|
||||
if use_brave_fallback
|
||||
else []
|
||||
)
|
||||
log_info(f"HuggingFace brave_fallback file={filename} candidates={len(brave_candidates)}")
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"brave_candidates",
|
||||
f"Brave candidates: {len(brave_candidates)}",
|
||||
84,
|
||||
candidate_count=len(brave_candidates),
|
||||
)
|
||||
|
||||
for candidate in brave_candidates:
|
||||
brave_count = len(brave_candidates)
|
||||
for candidate_index, candidate in enumerate(brave_candidates, start=1):
|
||||
repo_id = candidate.get("repo", "")
|
||||
branch = candidate.get("branch", "main")
|
||||
expected_file = candidate.get("filename", filename)
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"brave_repo",
|
||||
f"Checking Brave candidate {candidate_index}/{brave_count}",
|
||||
84 + (candidate_index / max(1, brave_count)) * 6,
|
||||
repo=repo_id,
|
||||
candidate_index=candidate_index,
|
||||
candidate_count=brave_count,
|
||||
)
|
||||
files = _get_repo_tree(repo_id, headers=headers, branch=branch)
|
||||
if not files:
|
||||
continue
|
||||
@@ -817,6 +956,14 @@ def search_huggingface_for_file(
|
||||
)
|
||||
if result and result.get("filename", "").lower() == filename.lower():
|
||||
_search_cache[cache_key] = result
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"found",
|
||||
"Found HuggingFace Brave match",
|
||||
92,
|
||||
repo=repo_id,
|
||||
match_type=result.get("match_type"),
|
||||
)
|
||||
log_info(
|
||||
f"HuggingFace found source=brave_fallback match=exact file={filename} repo={repo_id} path={result['path']}"
|
||||
)
|
||||
@@ -824,6 +971,15 @@ def search_huggingface_for_file(
|
||||
|
||||
# Not found
|
||||
_search_cache[cache_key] = None
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"done",
|
||||
"HuggingFace checked",
|
||||
92,
|
||||
candidate_count=len(repos_to_check),
|
||||
author_repo_count=author_fallback_repo_count,
|
||||
brave_candidate_count=len(brave_candidates),
|
||||
)
|
||||
log_info(
|
||||
f"HuggingFace miss file={filename} repos={len(repos_to_check)} author_repos={author_fallback_repo_count} brave={len(brave_candidates)}"
|
||||
)
|
||||
@@ -831,6 +987,13 @@ def search_huggingface_for_file(
|
||||
|
||||
except Exception as e:
|
||||
log_exception(f"HuggingFace search error for {filename}: {e}")
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"error",
|
||||
f"HuggingFace search error: {e}",
|
||||
100,
|
||||
status="error",
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
|
||||
@@ -9,7 +9,7 @@ import os
|
||||
import sqlite3
|
||||
import re
|
||||
import time
|
||||
from typing import Any, Dict, List, Optional
|
||||
from typing import Any, Callable, Dict, List, Optional
|
||||
|
||||
from ..matcher import calculate_similarity_with_normalization, normalize_filename
|
||||
from ..log_system.log_funcs import (
|
||||
@@ -50,6 +50,28 @@ def clear_search_cache():
|
||||
_db_path_cache = None
|
||||
|
||||
|
||||
def _report_progress(
|
||||
progress_callback: Optional[Callable[[Dict[str, Any]], None]],
|
||||
stage: str,
|
||||
message: str,
|
||||
percent: Optional[float] = None,
|
||||
**extra: Any,
|
||||
) -> None:
|
||||
if not progress_callback:
|
||||
return
|
||||
|
||||
payload = {"stage": stage, "message": message}
|
||||
if percent is not None:
|
||||
payload["percent"] = percent
|
||||
if extra:
|
||||
payload.update(extra)
|
||||
|
||||
try:
|
||||
progress_callback(payload)
|
||||
except Exception as e:
|
||||
log_debug(f"LoRA Manager archive progress callback failed: {e}")
|
||||
|
||||
|
||||
def get_lora_manager_archive_db_path() -> Optional[str]:
|
||||
"""Resolve the local comfyui-lora-manager archive database path."""
|
||||
global _db_path_cache
|
||||
@@ -548,6 +570,7 @@ def search_lora_manager_archive(
|
||||
model_type: Optional[str] = None,
|
||||
base_model_context: Optional[str] = None,
|
||||
limit: int = 10,
|
||||
progress_callback: Optional[Callable[[Dict[str, Any]], None]] = None,
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Search the local comfyui-lora-manager CivitAI archive by model name.
|
||||
@@ -570,27 +593,76 @@ def search_lora_manager_archive(
|
||||
f"search::{normalized_query.lower()}::{normalized_type}::{base_model_key}::{limit}"
|
||||
)
|
||||
if cache_key in _search_cache:
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"cache",
|
||||
"Using LoRA archive cache",
|
||||
76,
|
||||
result_count=len(_search_cache[cache_key] or []),
|
||||
)
|
||||
return list(_search_cache[cache_key])
|
||||
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"db",
|
||||
"Opening LoRA archive database",
|
||||
28,
|
||||
)
|
||||
conn = _connect_readonly()
|
||||
if conn is None:
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"unavailable",
|
||||
"LoRA archive database unavailable",
|
||||
92,
|
||||
)
|
||||
return []
|
||||
|
||||
try:
|
||||
started_at = time.perf_counter()
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"query",
|
||||
"Querying LoRA archive database",
|
||||
35,
|
||||
)
|
||||
candidate_rows = _query_candidate_rows(
|
||||
conn, normalized_query, normalized_type, limit
|
||||
)
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"candidates",
|
||||
f"LoRA archive rows: {len(candidate_rows)}",
|
||||
58,
|
||||
candidate_count=len(candidate_rows),
|
||||
)
|
||||
if not candidate_rows:
|
||||
elapsed = time.perf_counter() - started_at
|
||||
log_info(
|
||||
f"LoRA Manager archive query={normalized_query} type={normalized_type or 'all'} results=0 candidates=0 elapsed={elapsed:.3f}s"
|
||||
)
|
||||
_search_cache[cache_key] = []
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"results",
|
||||
"LoRA archive matches: 0",
|
||||
88,
|
||||
candidate_count=0,
|
||||
result_count=0,
|
||||
)
|
||||
return []
|
||||
|
||||
scored_candidates = []
|
||||
for row in candidate_rows:
|
||||
row_count = len(candidate_rows)
|
||||
for row_index, row in enumerate(candidate_rows, start=1):
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"score",
|
||||
f"Scoring LoRA archive row {row_index}/{row_count}",
|
||||
60 + (row_index / max(1, row_count)) * 10,
|
||||
candidate_index=row_index,
|
||||
candidate_count=row_count,
|
||||
)
|
||||
confidence = _calculate_confidence(
|
||||
normalized_query,
|
||||
row["model_name"] or "",
|
||||
@@ -616,10 +688,28 @@ def search_lora_manager_archive(
|
||||
row["version_id"]
|
||||
for _, _, row in scored_candidates[: max(limit * 3, limit)]
|
||||
]
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"details",
|
||||
f"Loading LoRA archive details: {len(selected_version_ids)}",
|
||||
72,
|
||||
candidate_count=len(selected_version_ids),
|
||||
)
|
||||
full_rows = _load_full_rows_for_versions(conn, selected_version_ids)
|
||||
|
||||
results = []
|
||||
for _, _, candidate_row in scored_candidates:
|
||||
scored_count = len(scored_candidates)
|
||||
for candidate_index, (_, _, candidate_row) in enumerate(
|
||||
scored_candidates, start=1
|
||||
):
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"build",
|
||||
f"Building LoRA archive match {candidate_index}/{scored_count}",
|
||||
74 + (candidate_index / max(1, scored_count)) * 12,
|
||||
candidate_index=candidate_index,
|
||||
candidate_count=scored_count,
|
||||
)
|
||||
version_id = candidate_row["version_id"]
|
||||
full_row = full_rows.get(version_id)
|
||||
if not full_row:
|
||||
@@ -647,9 +737,24 @@ def search_lora_manager_archive(
|
||||
log_info(
|
||||
f"LoRA Manager archive query={normalized_query} type={normalized_type or 'all'} results={len(results)} candidates={len(candidate_rows)} elapsed={elapsed:.3f}s"
|
||||
)
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"results",
|
||||
f"LoRA archive matches: {len(results)}",
|
||||
88,
|
||||
candidate_count=len(candidate_rows),
|
||||
result_count=len(results),
|
||||
)
|
||||
return results
|
||||
except Exception as e:
|
||||
log_exception(f"LoRA Manager archive search error for {query}: {e}")
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"error",
|
||||
f"LoRA archive search error: {e}",
|
||||
100,
|
||||
status="error",
|
||||
)
|
||||
return []
|
||||
finally:
|
||||
conn.close()
|
||||
@@ -661,6 +766,7 @@ def search_lora_manager_archive_for_file(
|
||||
base_model_context: Optional[str] = None,
|
||||
exact_only: bool = False,
|
||||
limit: int = 10,
|
||||
progress_callback: Optional[Callable[[Dict[str, Any]], None]] = None,
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
"""
|
||||
Search the archive for the best model/version match for a filename.
|
||||
@@ -674,6 +780,12 @@ def search_lora_manager_archive_for_file(
|
||||
f"file::{filename.lower()}::{_normalize_model_type(model_type)}::{base_model_key}::{exact_only}::{limit}"
|
||||
)
|
||||
if cache_key in _search_cache:
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"cache",
|
||||
"Using LoRA archive cache",
|
||||
86,
|
||||
)
|
||||
return _search_cache[cache_key]
|
||||
|
||||
candidates = search_lora_manager_archive(
|
||||
@@ -681,12 +793,29 @@ def search_lora_manager_archive_for_file(
|
||||
model_type=model_type,
|
||||
base_model_context=base_model_context,
|
||||
limit=limit,
|
||||
progress_callback=progress_callback,
|
||||
)
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"candidates",
|
||||
f"LoRA archive matches to check: {len(candidates)}",
|
||||
88 if candidates else 92,
|
||||
candidate_count=len(candidates),
|
||||
)
|
||||
best_match = None
|
||||
best_confidence = 0.0
|
||||
best_rank = -9999.0
|
||||
|
||||
for candidate in candidates:
|
||||
total_candidates = len(candidates)
|
||||
for candidate_index, candidate in enumerate(candidates, start=1):
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"candidate",
|
||||
f"Checking LoRA archive match {candidate_index}/{total_candidates}",
|
||||
88 + (candidate_index / max(1, total_candidates)) * 4,
|
||||
candidate_index=candidate_index,
|
||||
candidate_count=total_candidates,
|
||||
)
|
||||
candidate_filename = candidate.get("filename", "")
|
||||
confidence = _calculate_confidence(filename, candidate.get("name", ""), candidate.get("version_name", ""), candidate_filename)
|
||||
if exact_only and confidence < 100.0:
|
||||
@@ -706,11 +835,33 @@ def search_lora_manager_archive_for_file(
|
||||
best_match["confidence"] = confidence
|
||||
best_match["match_type"] = "exact" if confidence == 100.0 else "similar"
|
||||
if confidence == 100.0 and base_model_matches:
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"found",
|
||||
"Found LoRA archive match",
|
||||
92,
|
||||
confidence=confidence,
|
||||
)
|
||||
break
|
||||
|
||||
if best_match is None and exact_only:
|
||||
_search_cache[cache_key] = None
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"done",
|
||||
"LoRA archive checked",
|
||||
92,
|
||||
found=False,
|
||||
)
|
||||
return None
|
||||
|
||||
_search_cache[cache_key] = best_match
|
||||
_report_progress(
|
||||
progress_callback,
|
||||
"found" if best_match else "done",
|
||||
"Found LoRA archive match" if best_match else "LoRA archive checked",
|
||||
92,
|
||||
found=bool(best_match),
|
||||
confidence=best_confidence if best_match else None,
|
||||
)
|
||||
return best_match
|
||||
|
||||
File diff suppressed because one or more lines are too long
Reference in New Issue
Block a user