feat: wire progress callbacks into backend search sources

This commit is contained in:
Dariusz L
2026-06-23 03:33:28 +02:00
parent d8124d03da
commit 48bc3bf5f1
6 changed files with 609 additions and 19 deletions
+69 -3
View File
@@ -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
View File
@@ -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
View File
@@ -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
View File
@@ -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
+155 -4
View File
@@ -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