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

refactor: unify model download URL collection

Browse Source
This commit is contained in:
Dariusz L
2026-08-05 02:47:37 +02:00
parent 188809c31f
commit eb86f647c8
5 changed files with 190 additions and 126 deletions
Download Patch File Download Diff File Expand all files Collapse all files
core/services/civitai_search_service.py
+9 -48
View File
@@ -2,7 +2,7 @@
from .. import path_utils
from ..routes.context import RouteContext
from ..sources.common import is_remote_link_marked_dead
from ..sources.common import collect_download_urls
from .model_utils import CivitAISearchDependencies, ModelServiceDependencies
@@ -361,53 +361,14 @@ class CivitAISearchService(ModelServiceDependencies):
)
def collect_result_download_urls(result):
urls = []
if is_remote_link_marked_dead(result):
return urls
expected_filename = result.get("filename") or filename
dead_urls = set()
mirrors = result.get("mirrors") or []
if not isinstance(mirrors, list):
mirrors = [mirrors]
for mirror in mirrors:
if not isinstance(mirror, dict):
continue
if is_remote_link_marked_dead(mirror):
dead_url = str(mirror.get("url") or "").strip()
if dead_url.startswith(("http://", "https://")):
dead_urls.add(dead_url)
continue
url = str(mirror.get("url") or "").strip()
mirror_filename = (
mirror.get("filename")
or mirror.get("name")
or expected_filename
)
if (
looks_like_model_file(url, mirror_filename)
and url not in urls
):
urls.append(url)
raw_urls = result.get("download_urls") or []
if not isinstance(raw_urls, list):
raw_urls = [raw_urls]
for raw_url in raw_urls:
url = str(raw_url or "").strip()
if (
looks_like_model_file(url, expected_filename)
and url not in dead_urls
and url not in urls
):
urls.append(url)
for key in ("download_url", "downloadUrl"):
url = str(result.get(key) or "").strip()
if (
looks_like_model_file(url, expected_filename)
and url not in dead_urls
and url not in urls
):
urls.append(url)
return urls
return collect_download_urls(
result,
model_file_predicate=looks_like_model_file,
expected_filename=filename,
include_download_urls=True,
skip_dead_item=True,
download_url_keys=("download_url", "downloadUrl"),
)
def remote_download_url_is_alive(url):
headers = {
core/sources/civarchive.py
+19 -60
View File
@@ -51,6 +51,7 @@ from ..type_utils import (
from .common import (
build_custom_url_result,
build_unified_search_result,
collect_download_urls,
is_remote_link_marked_dead,
)
@@ -676,7 +677,7 @@ def _resolve_file_size_bytes(
if size:
return size
urls = download_urls if download_urls is not None else _collect_download_urls_unified(
urls = download_urls if download_urls is not None else _collect_archive_download_urls(
file_info,
**_NORMALIZED_DOWNLOAD_URL_OPTIONS,
)
@@ -874,64 +875,22 @@ def _select_primary_file(files: List[Dict[str, Any]]) -> Optional[Dict[str, Any]
return select_primary_model_file(files)
def _collect_download_urls_unified(
def _collect_archive_download_urls(
file_info: Dict[str, Any],
prioritize_civitai_last: bool = False,
skip_file_if_dead: bool = False,
check_download_urls_list: bool = False,
download_url_keys: Tuple[str, ...] = ("downloadUrl",),
) -> List[str]:
urls: List[str] = []
if skip_file_if_dead and is_remote_link_marked_dead(file_info):
return urls
expected_filename = file_info.get("filename") or file_info.get("name") or ""
dead_urls = set()
mirrors = file_info.get("mirrors") or []
if not isinstance(mirrors, list):
mirrors = [mirrors]
for mirror in mirrors:
if not isinstance(mirror, dict):
continue
if is_remote_link_marked_dead(mirror):
dead_url = _normalize_download_url(mirror.get("url"))
if dead_url:
dead_urls.add(dead_url)
continue
url = _normalize_download_url(mirror.get("url"))
mirror_filename = mirror.get("filename") or mirror.get("name") or expected_filename
if (
url
and _download_url_looks_like_model_file(url, mirror_filename)
and url not in urls
):
urls.append(url)
if check_download_urls_list:
raw_urls = file_info.get("download_urls") or []
if not isinstance(raw_urls, list):
raw_urls = [raw_urls]
for raw_url in raw_urls:
url = _normalize_download_url(raw_url)
if (
url
and url not in dead_urls
and _download_url_looks_like_model_file(url, expected_filename)
and url not in urls
):
urls.append(url)
if not is_remote_link_marked_dead(file_info):
for key in download_url_keys:
url = _normalize_download_url(file_info.get(key))
if (
url
and url not in dead_urls
and _download_url_looks_like_model_file(url, expected_filename)
and url not in urls
):
urls.append(url)
urls = collect_download_urls(
file_info,
model_file_predicate=_download_url_looks_like_model_file,
expected_filename=file_info.get("filename") or file_info.get("name") or "",
url_normalizer=_normalize_download_url,
include_download_urls=check_download_urls_list,
skip_dead_item=skip_file_if_dead,
download_url_keys=download_url_keys,
)
if prioritize_civitai_last:
non_civitai_urls = [
@@ -1010,7 +969,7 @@ def _normalize_archive_mirrors(file_info: Dict[str, Any]) -> List[Dict[str, Any]
def _normalize_archive_file(file_info: Dict[str, Any], model_id: Optional[int], version_id: Optional[int]) -> Dict[str, Any]:
transformed = _transform_file_entry(file_info)
download_urls = _collect_download_urls_unified(
download_urls = _collect_archive_download_urls(
transformed,
**_ARCHIVE_DOWNLOAD_URL_OPTIONS,
)
@@ -1211,7 +1170,7 @@ def _select_primary_model_file(files: List[Dict[str, Any]]) -> Optional[Dict[str
file_info
for file_info in files
if isinstance(file_info, dict)
and _collect_download_urls_unified(
and _collect_archive_download_urls(
file_info,
**_NORMALIZED_DOWNLOAD_URL_OPTIONS,
)
@@ -1225,7 +1184,7 @@ def _build_result_from_normalized_version(
file_info: Dict[str, Any],
match_type: str,
) -> Optional[Dict[str, Any]]:
download_urls = _collect_download_urls_unified(
download_urls = _collect_archive_download_urls(
file_info,
**_NORMALIZED_DOWNLOAD_URL_OPTIONS,
)
@@ -1464,7 +1423,7 @@ def _select_hash_page_file(
continue
if is_remote_link_marked_dead(file_info):
continue
if not _collect_download_urls_unified(
if not _collect_archive_download_urls(
file_info,
**_NORMALIZED_DOWNLOAD_URL_OPTIONS,
):
@@ -1500,7 +1459,7 @@ def _select_model_details_file(
for file_info in files:
if not isinstance(file_info, dict):
continue
if not _collect_download_urls_unified(
if not _collect_archive_download_urls(
file_info,
**_NORMALIZED_DOWNLOAD_URL_OPTIONS,
):
@@ -1561,7 +1520,7 @@ def _prefer_query_matching_mirror(
download_urls = []
if preferred_url:
download_urls.append(preferred_url)
for url in _collect_download_urls_unified(
for url in _collect_archive_download_urls(
file_info,
**_NORMALIZED_DOWNLOAD_URL_OPTIONS,
):
@@ -1657,7 +1616,7 @@ def _build_result_from_payload(
if exact_only and best_confidence < 100.0:
return None
download_urls = _collect_download_urls_unified(
download_urls = _collect_archive_download_urls(
selected_file,
**_ARCHIVE_DOWNLOAD_URL_OPTIONS,
)
core/sources/common.py
+72 -1
View File
@@ -2,7 +2,7 @@
Common utilities and result builder functions for external model metadata sources.
"""
from typing import Any, Dict, List, Optional
from typing import Any, Callable, Dict, List, Optional, Tuple
from ..type_utils import build_search_result, normalize_sha256
@@ -24,6 +24,77 @@ def is_remote_link_marked_dead(item: Any) -> bool:
)
def collect_download_urls(
item: Dict[str, Any],
*,
model_file_predicate: Callable[[str, str], bool],
expected_filename: str = "",
url_normalizer: Optional[Callable[[Any], Optional[str]]] = None,
include_download_urls: bool = False,
skip_dead_item: bool = False,
download_url_keys: Tuple[str, ...] = ("downloadUrl",),
) -> List[str]:
"""Collect valid, ordered, and deduplicated model download URLs."""
if not isinstance(item, dict):
return []
if skip_dead_item and is_remote_link_marked_dead(item):
return []
def normalize_url(value: Any) -> Optional[str]:
normalized = (
url_normalizer(value)
if url_normalizer is not None
else str(value or "").strip()
)
return normalized or None
expected = item.get("filename") or expected_filename
urls: List[str] = []
dead_urls = set()
mirrors = item.get("mirrors") or []
if not isinstance(mirrors, list):
mirrors = [mirrors]
for mirror in mirrors:
if not isinstance(mirror, dict):
continue
url = normalize_url(mirror.get("url"))
if is_remote_link_marked_dead(mirror):
if url and url.startswith(("http://", "https://")):
dead_urls.add(url)
continue
mirror_filename = mirror.get("filename") or mirror.get("name") or expected
if url and model_file_predicate(url, mirror_filename) and url not in urls:
urls.append(url)
if include_download_urls:
raw_urls = item.get("download_urls") or []
if not isinstance(raw_urls, list):
raw_urls = [raw_urls]
for raw_url in raw_urls:
url = normalize_url(raw_url)
if (
url
and url not in dead_urls
and model_file_predicate(url, expected)
and url not in urls
):
urls.append(url)
if not is_remote_link_marked_dead(item):
for key in download_url_keys:
url = normalize_url(item.get(key))
if (
url
and url not in dead_urls
and model_file_predicate(url, expected)
and url not in urls
):
urls.append(url)
return urls
def normalize_hashes_dict(hashes: Optional[Dict[str, Any]]) -> Dict[str, str]:
"""
Normalize a hashes dictionary (e.g. from CivitAI API or CivArchive) so that
tests/test_refactoring_targets.py
+73
View File
@@ -2,6 +2,7 @@ import importlib
import json
import os
import sys
import tempfile
import unittest
from typing import ClassVar
from unittest.mock import AsyncMock, MagicMock, patch
@@ -1290,6 +1291,78 @@ class TestRefactoringTargets(unittest.IsolatedAsyncioTestCase):
"resolved_path is outside configured model directories",
)
async def test_civitai_service_filters_dead_and_duplicate_download_urls(self):
expected_hash = "a" * 64
dead_url = "https://dead.example/model.safetensors"
mirror_url = "https://mirror.example/model.safetensors"
listed_url = "https://listed.example/model.safetensors"
direct_url = "https://direct.example/model.safetensors"
result = {
"source": "civarchive",
"from_metadata": True,
"filename": "model.safetensors",
"sha256": expected_hash,
"mirrors": [
{"url": dead_url, "isDead": True},
{"url": mirror_url},
],
"download_urls": [dead_url, listed_url, mirror_url],
"download_url": direct_url,
}
def to_bool(value, default=False):
return default if value is None else bool(value)
response = MagicMock(status_code=200)
request_public_url = MagicMock(return_value=(response, mirror_url, {}))
resolve_civarchive_by_hash = MagicMock(return_value=result)
service = self._build_model_service_for_unit_test(
service_module="civitai_search_service",
service_class="CivitAISearchService",
download_available=True,
extract_sha256_from_metadata=lambda value: value.get("sha256", ""),
find_local_file_path=lambda filename, category=None: None,
get_filename_from_path=lambda value: str(value or "").replace(
"\\", "/"
).rsplit("/", 1)[-1],
is_path_in_configured_model_roots=lambda path: True,
looks_like_model_file=lambda url, expected_filename="": str(url).endswith(
".safetensors"
),
normalize_category_to_model_type=lambda value: value,
normalize_sha256=lambda value: str(value or ""),
request_public_url=request_public_url,
resolve_civarchive_by_hash=resolve_civarchive_by_hash,
search_huggingface_for_file=MagicMock(return_value=None),
to_bool=to_bool,
)
with tempfile.NamedTemporaryFile(
suffix=".safetensors", delete=False
) as handle:
model_path = handle.name
request = AsyncMock()
request.json.return_value = {
"filename": "model.safetensors",
"resolved_path": model_path,
"sha256": expected_hash,
}
try:
with patch(
"core.sources.civitai.get_model_info_by_hash",
return_value=None,
):
api_response = await service.civitai_search(request)
finally:
os.unlink(model_path)
self.assertEqual(api_response.status, 200)
resolve_civarchive_by_hash.assert_called_once()
request_public_url.assert_called_once()
self.assertEqual(request_public_url.call_args.args[1], mirror_url)
payload = json.loads(api_response.text)
self.assertEqual(payload["download_url"], mirror_url)
async def test_custom_url_service_rejects_unsafe_url(self):
class TestUnsafeUrlError(Exception):
pass
tests/test_source_internals.py
+17 -17
View File
@@ -22,7 +22,7 @@ import requests
# civarchive internals
# ---------------------------------------------------------------------------
from core.sources.civarchive import (
_collect_download_urls_unified,
_collect_archive_download_urls,
_fetch_remote_file_size_bytes,
_normalize_archive_version,
_normalize_download_url,
@@ -560,9 +560,9 @@ class RemoteLinkIsDeadTests(unittest.TestCase):
# ===========================================================================
# civarchive - _collect_download_urls_unified (new consolidated helper)
# civarchive - _collect_archive_download_urls (common collector adapter)
# ===========================================================================
class CollectDownloadUrlsUnifiedTests(unittest.TestCase):
class CollectArchiveDownloadUrlsTests(unittest.TestCase):
def _make_file_info(self, mirrors=None, download_url=None,
download_urls=None, dead=False):
@@ -581,14 +581,14 @@ class CollectDownloadUrlsUnifiedTests(unittest.TestCase):
fi = self._make_file_info(mirrors=[
{"url": "https://civarchive.com/api/download/models/1.safetensors"},
])
result = _collect_download_urls_unified(fi)
result = _collect_archive_download_urls(fi)
self.assertIn("https://civarchive.com/api/download/models/1.safetensors", result)
def test_skips_dead_mirrors(self):
fi = self._make_file_info(mirrors=[
{"url": "https://civarchive.com/api/download/models/1.safetensors", "isDead": True},
])
result = _collect_download_urls_unified(fi)
result = _collect_archive_download_urls(fi)
self.assertEqual(result, [])
def test_skip_file_if_dead_returns_empty(self):
@@ -596,14 +596,14 @@ class CollectDownloadUrlsUnifiedTests(unittest.TestCase):
download_url="https://civarchive.com/api/download/models/1.safetensors",
dead=True,
)
result = _collect_download_urls_unified(fi, skip_file_if_dead=True)
result = _collect_archive_download_urls(fi, skip_file_if_dead=True)
self.assertEqual(result, [])
def test_check_download_urls_list_includes_extra_urls(self):
fi = self._make_file_info(download_urls=[
"https://civarchive.com/api/download/models/2.safetensors"
])
result = _collect_download_urls_unified(fi, check_download_urls_list=True,
result = _collect_archive_download_urls(fi, check_download_urls_list=True,
skip_file_if_dead=True)
self.assertIn("https://civarchive.com/api/download/models/2.safetensors", result)
@@ -614,7 +614,7 @@ class CollectDownloadUrlsUnifiedTests(unittest.TestCase):
{"url": civitai_url},
{"url": other_url},
])
result = _collect_download_urls_unified(fi, prioritize_civitai_last=True)
result = _collect_archive_download_urls(fi, prioritize_civitai_last=True)
self.assertEqual(result[-1], civitai_url)
self.assertEqual(result[0], other_url)
@@ -624,23 +624,23 @@ class CollectDownloadUrlsUnifiedTests(unittest.TestCase):
mirrors=[{"url": url}],
download_url=url,
)
result = _collect_download_urls_unified(fi)
result = _collect_archive_download_urls(fi)
self.assertEqual(result.count(url), 1)
def test_empty_file_info_returns_empty(self):
self.assertEqual(_collect_download_urls_unified({}), [])
self.assertEqual(_collect_archive_download_urls({}), [])
# ===========================================================================
# civarchive - archive download URL collection strategy
# ===========================================================================
class CollectDownloadUrlsStrategyTests(unittest.TestCase):
class CollectArchiveDownloadUrlsStrategyTests(unittest.TestCase):
def test_civitai_url_sorted_last(self):
civitai = "https://civitai.com/api/download/models/99.safetensors"
archive = "https://civarchive.com/api/download/models/77.safetensors"
fi = {"mirrors": [{"url": civitai}, {"url": archive}]}
result = _collect_download_urls_unified(
result = _collect_archive_download_urls(
fi,
prioritize_civitai_last=True,
download_url_keys=("downloadUrl",),
@@ -649,7 +649,7 @@ class CollectDownloadUrlsStrategyTests(unittest.TestCase):
self.assertEqual(result[0], archive)
def test_returns_list(self):
result = _collect_download_urls_unified(
result = _collect_archive_download_urls(
{},
prioritize_civitai_last=True,
download_url_keys=("downloadUrl",),
@@ -660,7 +660,7 @@ class CollectDownloadUrlsStrategyTests(unittest.TestCase):
# ===========================================================================
# civarchive - normalized download URL collection strategy
# ===========================================================================
class CollectNormalizedDownloadUrlsTests(unittest.TestCase):
class CollectNormalizedArchiveDownloadUrlsTests(unittest.TestCase):
def test_skips_dead_file(self):
fi = {
@@ -668,7 +668,7 @@ class CollectNormalizedDownloadUrlsTests(unittest.TestCase):
"downloadUrl": "https://civarchive.com/api/download/models/1.safetensors",
}
self.assertEqual(
_collect_download_urls_unified(
_collect_archive_download_urls(
fi,
skip_file_if_dead=True,
check_download_urls_list=True,
@@ -682,7 +682,7 @@ class CollectNormalizedDownloadUrlsTests(unittest.TestCase):
"mirrors": [{"url": "https://civarchive.com/api/download/models/1.safetensors"}],
"download_urls": ["https://civarchive.com/api/download/models/2.safetensors"],
}
result = _collect_download_urls_unified(
result = _collect_archive_download_urls(
fi,
skip_file_if_dead=True,
check_download_urls_list=True,
@@ -696,7 +696,7 @@ class CollectNormalizedDownloadUrlsTests(unittest.TestCase):
"mirrors": [{"url": dead_url, "isDead": True}],
"download_urls": [dead_url],
}
result = _collect_download_urls_unified(
result = _collect_archive_download_urls(
fi,
skip_file_if_dead=True,
check_download_urls_list=True,