From ddde0ac3ce901bb6daf5283e666b30625d4b1c11 Mon Sep 17 00:00:00 2001 From: Dariusz L Date: Tue, 25 Aug 2026 02:35:15 +0200 Subject: [PATCH] feat: add CivitAI title-based search fallback for filename mismatches --- core/sources/civitai.py | 123 +++++++++++++++++++++++++++++++++ tests/test_source_internals.py | 69 ++++++++++++++++++ 2 files changed, 192 insertions(+) diff --git a/core/sources/civitai.py b/core/sources/civitai.py index ccf73ac..39f0788 100644 --- a/core/sources/civitai.py +++ b/core/sources/civitai.py @@ -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( diff --git a/tests/test_source_internals.py b/tests/test_source_internals.py index cb99234..6617915 100644 --- a/tests/test_source_internals.py +++ b/tests/test_source_internals.py @@ -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 # ===========================================================================