From 9f537eaef06f97280ee848c020a72a53f052211f Mon Sep 17 00:00:00 2001 From: Dariusz L Date: Thu, 10 Sep 2026 02:24:24 +0200 Subject: [PATCH] Optimize CivitAI lookups by caching model data and skipping unnecessary version requests Reuse model metadata across _find_civitai_file_in_model calls within a single search, avoid resolve_urn when modelVersions already contains usable file lists, and add tests for request-count behavior. --- core/sources/civitai.py | 124 ++++++++++++++++++++------------- tests/test_source_internals.py | 56 +++++++++++++++ 2 files changed, 130 insertions(+), 50 deletions(-) diff --git a/core/sources/civitai.py b/core/sources/civitai.py index 39f0788..3ffc3b5 100644 --- a/core/sources/civitai.py +++ b/core/sources/civitai.py @@ -792,6 +792,8 @@ def _find_civitai_file_in_model( exact_only: bool = False, preferred_version_id: Optional[int] = None, base_model_context: Optional[str] = None, + model_data_cache: Optional[Dict[int, Any]] = None, + title_only: bool = False, ) -> Optional[SearchResult]: """Load one CivitAI model and search all its versions for the requested file.""" filename_lower = filename.lower() @@ -843,35 +845,40 @@ def _find_civitai_file_in_model( return "partial" return None - if preferred_version_id is not None: - resolved = resolve_urn(model_id, preferred_version_id, api_key) - resolved_match_type = ( - resolved_version_match_type(resolved) if resolved else None - ) - if ( - resolved - and resolved_match_type - and _base_model_matches(resolved.get("base_model"), base_model_context) - ): - return build_result_from_resolved_version( - resolved, - preferred_version_id, - match_type=resolved_match_type, - ) - best_resolved_result = None best_resolved_confidence = 0.0 - data = execute_provider_json_request( - "CivitAI model lookup", - f"{CIVITAI_API_URL}/models/{model_id}", - api_key=api_key, - timeout=15, - ) + if model_data_cache is None: + model_data_cache = {} + if model_id not in model_data_cache: + model_data_cache[model_id] = execute_provider_json_request( + "CivitAI model lookup", + f"{CIVITAI_API_URL}/models/{model_id}", + api_key=api_key, + timeout=15, + ) + data = model_data_cache[model_id] if data is None: + if title_only: + return None + if preferred_version_id is not None: + resolved = resolve_urn(model_id, preferred_version_id, api_key) + resolved_match_type = ( + resolved_version_match_type(resolved) if resolved else None + ) + if ( + resolved + and resolved_match_type + and _base_model_matches(resolved.get("base_model"), base_model_context) + ): + return build_result_from_resolved_version( + resolved, + preferred_version_id, + match_type=resolved_match_type, + ) return None versions = data.get("modelVersions", []) - if allow_model_title_match: + if allow_model_title_match or title_only: title_match = _find_model_title_match_in_model( model_id=model_id, model_data=data, @@ -881,6 +888,8 @@ def _find_civitai_file_in_model( ) if title_match: return title_match + if title_only: + return None if preferred_version_id is not None: preferred = [v for v in versions if v.get("id") == preferred_version_id] @@ -891,11 +900,49 @@ def _find_civitai_file_in_model( f"CivitAI model lookup model_id={model_id} returned {len(versions)} versions" ) + matching_versions = [ + version for version in versions + if _base_model_matches(version.get("baseModel"), base_model_context) + ] + match = _find_matching_file_in_versions(matching_versions, filename, exact_only=exact_only) + if match: + version = match["version"] + file_info = match["file_info"] + result = _build_civitai_result_from_version( + model_id=model_id, + model_name=data.get("name", ""), + model_type=data.get("type", ""), + version=version, + file_info=file_info, + tags=data.get("tags", []), + match_type=match["match_type"], + ) + result = result.with_updates( + confidence=match.get( + "confidence", + calculate_filename_confidence(filename, result.filename), + ) + ) + if match["match_type"] == "similar": + log.info( + f"CivitAI version-list probable match: model_id={model_id}, version_id={result.version_id}, filename={result.filename}, confidence={result.confidence}" + ) + if _base_model_matches(result.base_model, base_model_context): + return result + if not base_model_context: + return result + + for version in versions: version_id = version.get("id") if not version_id: continue + # Complete file lists were already checked without extra requests. + # Fetch details only when this endpoint omitted usable file names. + files = version.get("files") or [] + if files and all(isinstance(item, dict) and item.get("name") for item in files): + continue resolved = resolve_urn(model_id, version_id, api_key) if resolved: resolved_match_type = resolved_version_match_type(resolved) @@ -940,33 +987,6 @@ def _find_civitai_file_in_model( ) return best_resolved_result - match = _find_matching_file_in_versions(versions, filename, exact_only=exact_only) - if match: - version = match["version"] - file_info = match["file_info"] - result = _build_civitai_result_from_version( - model_id=model_id, - model_name=data.get("name", ""), - model_type=data.get("type", ""), - version=version, - file_info=file_info, - tags=data.get("tags", []), - match_type=match["match_type"], - ) - result = result.with_updates( - confidence=match.get( - "confidence", - calculate_filename_confidence(filename, result.filename), - ) - ) - if match["match_type"] == "similar": - log.info( - f"CivitAI version-list probable match: model_id={model_id}, version_id={result.version_id}, filename={result.filename}, confidence={result.confidence}" - ) - if _base_model_matches(result.base_model, base_model_context): - return result - if not base_model_context: - return result return None @@ -1040,6 +1060,7 @@ def search_civitai_for_file( ) return _search_cache[cache_key] + model_data_cache: Dict[int, Any] = {} try: if requested_sha256: _report_progress( @@ -1221,6 +1242,7 @@ def search_civitai_for_file( exact_only=exact_only, preferred_version_id=version_id, base_model_context=base_model_context, + model_data_cache=model_data_cache, ) if result: confidence = result.confidence @@ -1314,6 +1336,8 @@ def search_civitai_for_file( exact_only=False, preferred_version_id=candidate.get("version_id"), base_model_context=base_model_context, + model_data_cache=model_data_cache, + title_only=True, ) if not title_result or title_result.match_type != "model_title": continue diff --git a/tests/test_source_internals.py b/tests/test_source_internals.py index 6617915..84f249d 100644 --- a/tests/test_source_internals.py +++ b/tests/test_source_internals.py @@ -1979,8 +1979,10 @@ class CivitaiFilenameMatchTests(unittest.TestCase): self.assertEqual(result["match_type"], "exact") @patch("core.sources.civitai.resolve_urn") + @patch("core.sources.civitai.execute_provider_json_request", return_value=None) def test_resolved_meaningful_partial_filename_keeps_partial_type( self, + mock_request, mock_resolve_urn, ): mock_resolve_urn.return_value = { @@ -2196,3 +2198,57 @@ class GetBaseDirectoryForModelTests(unittest.TestCase): if __name__ == "__main__": unittest.main() + + +class CivitaiRequestOptimizationTests(unittest.TestCase): + @patch("core.sources.civitai.resolve_urn") + @patch("core.sources.civitai.execute_provider_json_request") + def test_reuses_model_data_across_queries_and_skips_unrelated_versions(self, request, resolve): + request.return_value = { + "name": "Unrelated", + "modelVersions": [{"id": n, "files": [{"name": "zzzzzz.safetensors"}]} for n in range(20)], + } + cache = {} + for query in ("target_model.safetensors", "target model", "target model LoRA"): + _find_civitai_file_in_model(1, query, model_data_cache=cache, title_only=True) + request.assert_called_once() + resolve.assert_not_called() + + @patch("core.sources.civitai.resolve_urn") + @patch("core.sources.civitai.execute_provider_json_request") + def test_finds_non_primary_file_without_version_requests(self, request, resolve): + request.return_value = { + "name": "Example", "type": "LORA", + "modelVersions": [{"id": 2, "baseModel": "Krea 2", "files": [ + {"name": "other.safetensors", "primary": True}, + {"name": "target.safetensors", "sizeKB": 12, "hashes": {"SHA256": "a" * 64}}, + ]}], + } + result = _find_civitai_file_in_model(1, "target.safetensors", preferred_version_id=2, exact_only=True) + self.assertEqual(result.filename, "target.safetensors") + self.assertEqual(result.match_type, "exact") + request.assert_called_once() + resolve.assert_not_called() + + @patch("core.sources.civitai.resolve_urn") + @patch("core.sources.civitai.execute_provider_json_request") + def test_unrelated_complete_file_lists_need_no_version_requests(self, request, resolve): + request.return_value = { + "name": "Unrelated", + "modelVersions": [{"id": n, "files": [{"name": "zzzzzz.safetensors"}]} for n in range(20)], + } + result = _find_civitai_file_in_model(1, "target.safetensors", exact_only=True) + self.assertIsNone(result) + resolve.assert_not_called() + + @patch("core.sources.civitai.resolve_urn") + @patch("core.sources.civitai.execute_provider_json_request") + def test_missing_file_list_falls_back_to_version_details(self, request, resolve): + request.return_value = {"name": "Example", "modelVersions": [{"id": 2}]} + resolve.return_value = { + "expected_filename": "target.safetensors", + "files": [{"name": "target.safetensors", "size": 12}], + } + result = _find_civitai_file_in_model(1, "target.safetensors", exact_only=True) + self.assertEqual(result.filename, "target.safetensors") + resolve.assert_called_once_with(1, 2, None)