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

feat: add CivitAI title-based search fallback for filename mismatches

Browse Source
This commit is contained in:
Dariusz L
2026-08-25 02:35:15 +02:00
parent 7627c61450
commit ddde0ac3ce
2 changed files with 192 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
core/sources/civitai.py
+123
View File
@@ -283,6 +283,51 @@ def _find_model_title_match_in_model(
)
def _build_civitai_title_search_queries(
filename: str,
model_type: Optional[str] = None,
) -> List[str]:
"""Build conservative model-title queries from a workflow filename.
Workflow filenames and provider model titles often differ. Moving the
final filename token to the front helps names such as
``Text Refusal Reduction Krea2`` find a provider title beginning with
``Krea2``. For LoRAs, adding the provider type also removes unrelated
text-focused models from the candidate set.
"""
basename = get_filename_from_path(filename).strip()
if not basename:
return []
stem = os.path.splitext(basename)[0].strip()
tokens = re.findall(r"[A-Za-z0-9]+", stem)
if len(tokens) < 3:
return []
title_queries: List[str] = []
def add_query(parts: List[str]) -> None:
query = " ".join(part for part in parts if part).strip()
if query and query not in title_queries:
title_queries.append(query)
civitai_type = CIVITAI_API_TYPE_MAP.get(str(model_type or "").lower())
type_label = "LoRA" if civitai_type == "LORA" else ""
final_token = tokens[-1]
without_first = [final_token, *tokens[1:-1]]
rotated = [final_token, *tokens[:-1]]
if type_label:
add_query([*without_first, type_label])
add_query(without_first)
if type_label:
add_query([*rotated, type_label])
add_query(rotated)
add_query(tokens)
return title_queries[:5]
def _filename_base_partial_match(target_base: str, candidate_base: str) -> bool:
"""Return True for meaningful filename-base containment matches.
@@ -1214,6 +1259,84 @@ def search_civitai_for_file(
)
return best_result
if not exact_only:
title_queries = _build_civitai_title_search_queries(
filename,
model_type=model_type,
)
title_candidate_limit = min(
MAX_CIVITAI_CANDIDATE_LIMIT,
max(candidate_limit, 10),
)
for title_index, title_query in enumerate(title_queries, start=1):
_report_progress(
progress_callback,
"title_fallback",
f"Searching CivitAI model title {title_index}/{len(title_queries)}",
72 + (title_index / max(1, len(title_queries))) * 16,
title_query=title_query,
)
title_candidates: List[Dict[str, Optional[int]]] = []
if use_trpc_search:
title_candidates = _search_civitai_trpc_candidates(
title_query,
model_type=model_type,
session_token=session_token,
limit=title_candidate_limit,
)
if not title_candidates and use_api_search:
title_candidates = _search_civitai_public_api_candidates(
title_query,
model_type=model_type,
api_key=api_key,
session_token=session_token,
limit=title_candidate_limit,
)
if not title_candidates and use_html_fallback:
title_candidates = _search_civitai_red_candidates(
title_query,
model_type=model_type,
session_token=session_token,
limit=title_candidate_limit,
)
for candidate in title_candidates:
model_id = candidate.get("model_id")
if not isinstance(model_id, int):
continue
title_result = _find_civitai_file_in_model(
model_id=model_id,
filename=title_query,
api_key=api_key,
exact_only=False,
preferred_version_id=candidate.get("version_id"),
base_model_context=base_model_context,
)
if not title_result or title_result.match_type != "model_title":
continue
title_result = title_result.with_extra(
title_fallback=True,
title_query=title_query,
)
_search_cache[cache_key] = title_result
_report_progress(
progress_callback,
"found",
"Found CivitAI model-title match",
92,
model_id=title_result.model_id,
version_id=title_result.version_id,
confidence=title_result.confidence,
)
log.info(
f"Found CivitAI model-title match for {filename}: model_id={title_result.model_id}, version_id={title_result.version_id}, title_query={title_query}, confidence={title_result.confidence}"
)
return title_result
# Not found
_search_cache[cache_key] = None
_report_progress(
tests/test_source_internals.py
+69
View File
@@ -48,6 +48,7 @@ from core.sources.civarchive import (
# ---------------------------------------------------------------------------
from core.sources.civitai import (
_build_civitai_result_from_version,
_build_civitai_title_search_queries,
_enrich_model_info_with_details,
_extract_model_images,
_extract_public_api_model_candidates,
@@ -2041,6 +2042,74 @@ class CivitaiFilenameMatchTests(unittest.TestCase):
clear_search_cache()
class CivitaiTitleFallbackTests(unittest.TestCase):
def test_title_queries_reorder_filename_and_add_lora_type(self):
queries = _build_civitai_title_search_queries(
"Text Refusal Reduction Krea2.safetensors",
model_type="loras",
)
self.assertEqual(
[
"Krea2 Refusal Reduction LoRA",
"Krea2 Refusal Reduction",
"Krea2 Text Refusal Reduction LoRA",
"Krea2 Text Refusal Reduction",
"Text Refusal Reduction Krea2",
],
queries,
)
@patch("core.sources.civitai._search_civitai_trpc_candidates")
@patch("core.sources.civitai._find_civitai_file_in_model")
def test_title_fallback_finds_model_when_filename_differs(
self,
mock_find_model,
mock_search_candidates,
):
clear_search_cache()
expected = SearchResult(
source="civitai",
model_id=2775340,
version_id=3125118,
name="Krea2 TextFusion Refusal-Reduction LoRA",
filename="Krea2_TextFusion_Refusal_Reduction.safetensors",
match_type="model_title",
confidence=83.6,
)
def search_candidates(query, **kwargs):
if query == "Krea2 Refusal Reduction LoRA":
return [{"model_id": 2775340, "version_id": 3125118}]
return []
def find_model(model_id, filename, **kwargs):
if filename == "Krea2 Refusal Reduction LoRA":
return expected
return None
mock_search_candidates.side_effect = search_candidates
mock_find_model.side_effect = find_model
result = search_civitai_for_file(
"Text Refusal Reduction Krea2.safetensors",
model_type="loras",
use_api_search=False,
use_html_fallback=False,
)
self.assertIsNotNone(result)
self.assertEqual(2775340, result.model_id)
self.assertEqual("model_title", result.match_type)
self.assertTrue(result.extra_value("title_fallback"))
self.assertEqual(
"Krea2 Refusal Reduction LoRA",
result.extra_value("title_query"),
)
clear_search_cache()
# ===========================================================================
# civitai - hash lookup details enrichment
# ===========================================================================