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

refactor: unify primary model file selection

Browse Source
This commit is contained in:
Dariusz L
2026-08-05 14:29:02 +02:00
parent 24d0dae0e0
commit 7a219bd2d2
7 changed files with 164 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files
core/services/search_providers.py
+6 -9
View File
@@ -4,6 +4,7 @@ from dataclasses import dataclass
from typing import Any
from ..sources.civitai import build_civitai_result_payload
from ..type_utils import select_primary_model_file
class SearchCancelled(BaseException):
@@ -358,15 +359,11 @@ class SearchProviderRunner:
"Selecting CivitAI file",
76,
)
primary_file = None
for file_info in model_info.get("files", []):
if file_info.get("name") == model_info.get(
"expected_filename"
):
primary_file = file_info
break
if primary_file is None:
primary_file = (model_info.get("files") or [{}])[0]
primary_file = select_primary_model_file(
model_info.get("files") or [],
expected_filename=model_info.get("expected_filename"),
fallback_to_first=True,
) or {}
download_url = self.owner.get_civitai_download_url(version_id)
result_payload = build_civitai_result_payload(
core/sources/civarchive.py
+1 -5
View File
@@ -871,10 +871,6 @@ def _extract_version_files(
return _merge_top_file_mirrors(files, top_files or [])
def _select_primary_file(files: List[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
return select_primary_model_file(files)
def _collect_archive_download_urls(
file_info: Dict[str, Any],
prioritize_civitai_last: bool = False,
@@ -1603,7 +1599,7 @@ def _build_result_from_payload(
selected_file = file_info
if selected_file is None:
selected_file = _select_primary_file(files)
selected_file = select_primary_model_file(files)
best_confidence = calculate_archived_model_confidence(
query_value,
context.get("name", ""),
core/sources/civitai.py
+13 -25
View File
@@ -738,13 +738,11 @@ def _find_civitai_file_in_model(
resolved: Dict[str, Any], version_id: int
) -> Dict[str, Any]:
expected_filename = resolved.get("expected_filename", "")
primary_file = None
for file_info in resolved.get("files", []):
if file_info.get("name") == expected_filename:
primary_file = file_info
break
if primary_file is None:
primary_file = (resolved.get("files") or [{}])[0]
primary_file = select_primary_model_file(
resolved.get("files") or [],
expected_filename=expected_filename,
fallback_to_first=True,
) or {}
return {
"source": "civitai",
@@ -1248,14 +1246,10 @@ def search_civitai(
version_id = latest.get("id")
files = latest.get("files", [])
primary_file = None
for f in files:
if f.get("primary", False) or f.get("type") == "Model":
primary_file = f
break
if not primary_file and files:
primary_file = files[0]
primary_file = select_primary_model_file(
files,
prefer_first_marked=True,
)
result = {
"source": "civitai",
@@ -1533,16 +1527,10 @@ def resolve_urn(
return None
files = target_version.get("files", [])
primary_file = None
# Prefer primary file or type=='Model'
for f in files:
if f.get("primary") or f.get("type") == "Model":
primary_file = f
break
if not primary_file and files:
primary_file = files[0] # Fallback to first
primary_file = select_primary_model_file(
files,
prefer_first_marked=True,
)
if not primary_file:
log.warning(f"No files found for version {version_id}")
core/type_utils.py
+14
View File
@@ -409,10 +409,16 @@ def select_primary_model_file(
files: List[Dict[str, Any]],
expected_filename: Optional[str] = None,
require_download: bool = False,
fallback_to_first: bool = False,
prefer_first_marked: bool = False,
) -> Optional[Dict[str, Any]]:
"""
Selects the primary file from a list of model version files,
optionally filtering by expected filename or requiring a download URL.
``fallback_to_first`` preserves callers that use an exact filename match
followed by the first available file. ``prefer_first_marked`` preserves
callers that select the first file marked primary or of type ``Model``.
"""
if not isinstance(files, list):
files = [files] if files else []
@@ -440,6 +446,14 @@ def select_primary_model_file(
for file_info in candidates:
if get_filename_from_path(file_name(file_info)).lower() == expected:
return file_info
if fallback_to_first:
return candidates[0]
if prefer_first_marked:
for file_info in candidates:
if file_info.get("primary") or str(file_info.get("type") or "").lower() == "model":
return file_info
return candidates[0]
# 1. Look for file marked primary and of type Model
for file_info in candidates:
tests/test_refactoring_targets.py
+20
View File
@@ -1082,6 +1082,26 @@ class TestRefactoringTargets(unittest.IsolatedAsyncioTestCase):
# Require download excludes main.safetensors because it has no download url
self.assertEqual(select_primary_model_file(files, require_download=True)["id"], 1)
marked_files = [
{"id": 4, "name": "first.safetensors", "type": "Model"},
{"id": 5, "name": "primary.safetensors", "primary": True, "type": "Model"},
]
self.assertEqual(
select_primary_model_file(
marked_files,
prefer_first_marked=True,
)["id"],
4,
)
self.assertEqual(
select_primary_model_file(
marked_files,
expected_filename="missing.safetensors",
fallback_to_first=True,
)["id"],
4,
)
def _build_model_service_for_unit_test(
self,
service_module="model_service",
tests/test_search_service_components.py
+42
View File
@@ -389,6 +389,48 @@ def test_search_provider_runner_resolves_civitai_urn():
assert result["civitai"]["hashes"] == {}
def test_search_provider_runner_keeps_first_file_when_expected_name_is_missing():
owner = SimpleNamespace(
CivArchiveSearchError=Exception,
logger=MagicMock(),
search_tracker=MagicMock(),
format_log_fields=MagicMock(return_value="ids=10@20"),
resolve_urn=MagicMock(
return_value={
"model_name": "URN Model",
"version_name": "v1",
"expected_filename": "missing.safetensors",
"base_model": "SDXL",
"files": [
{"name": "fallback.safetensors", "sha256": "a" * 64},
{
"name": "primary.safetensors",
"primary": True,
"type": "Model",
"sha256": "b" * 64,
},
],
}
),
get_civitai_download_url=MagicMock(return_value="https://example.test/download"),
build_search_result=MagicMock(side_effect=lambda source, **fields: {"source": source, **fields}),
)
owner.search_tracker.is_cancelled.return_value = False
owner.log_search_result = MagicMock()
runner = SearchProviderRunner(owner)
result, found = runner.search_civitai_source_task(
_request(
data={"model_id": 10, "version_id": 20},
is_urn=True,
base_model_context="",
)
)
assert found is True
assert result["civitai"]["sha256"] == "a" * 64
def test_search_provider_runner_falls_back_for_civitai_urn_without_ids():
owner = SimpleNamespace(
logger=MagicMock(),
tests/test_source_internals.py
+68
View File
@@ -53,6 +53,8 @@ from core.sources.civitai import (
get_model_info_by_hash,
get_model_info_for_file,
parse_civitai_url,
resolve_urn,
search_civitai,
search_civitai_for_file,
)
from core.sources.common import is_remote_link_marked_dead
@@ -125,6 +127,72 @@ class CivitaiResultBuilderTests(unittest.TestCase):
self.assertEqual(10 * 1024, result["size"])
class CivitaiPrimaryFileSelectionTests(unittest.TestCase):
@patch("core.sources.civitai.execute_provider_json_request")
def test_resolve_urn_preserves_first_marked_file_order(self, mock_request):
clear_search_cache()
mock_request.return_value = {
"name": "Example model",
"tags": [],
"modelVersions": [
{
"id": 456,
"name": "v1",
"baseModel": "SDXL",
"files": [
{"name": "first.safetensors", "type": "Model", "sizeKB": 1},
{
"name": "primary.safetensors",
"primary": True,
"type": "Model",
"sizeKB": 2,
},
],
}
],
}
result = resolve_urn(123, 456)
self.assertEqual("first.safetensors", result["expected_filename"])
clear_search_cache()
@patch("core.sources.civitai.execute_provider_json_request")
def test_general_search_preserves_first_marked_file_order(self, mock_request):
clear_search_cache()
mock_request.return_value = {
"items": [
{
"id": 123,
"name": "Example model",
"type": "Checkpoint",
"stats": {"downloadCount": 1},
"modelVersions": [
{
"id": 456,
"baseModel": "SDXL",
"files": [
{"name": "first.safetensors", "type": "Model", "sizeKB": 1},
{
"name": "primary.safetensors",
"primary": True,
"type": "Model",
"sizeKB": 2,
},
],
}
],
}
]
}
results = search_civitai("example", limit=1)
self.assertEqual("first.safetensors", results[0]["filename"])
clear_search_cache()
class CustomUrlResultBuilderTests(unittest.TestCase):
def test_civitai_custom_result_preserves_all_result_fields(self):
result = build_civitai_custom_result({