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

Optimize CivitAI lookups by caching model data and skipping unnecessary version requests

Browse Source 
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.
This commit is contained in:
Dariusz L
2026-09-10 02:24:24 +02:00
parent 63050819fe
commit 9f537eaef0
2 changed files with 130 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files
core/sources/civitai.py
+74 -50
View File
@@ -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
tests/test_source_internals.py
+56
View File
@@ -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)