From 76d0a3aa7cea8939754577df100fad99988caaaa Mon Sep 17 00:00:00 2001 From: Dariusz L Date: Thu, 6 Aug 2026 22:25:42 +0200 Subject: [PATCH] refactor shared remote file size resolution --- core/sources/civarchive.py | 17 +++++----- core/sources/common.py | 23 ++++++++++++-- core/sources/huggingface.py | 21 ++++++++----- tests/test_huggingface.py | 18 +++++++++++ tests/test_source_internals.py | 57 +++++++++++++++++++++++++++++++++- 5 files changed, 115 insertions(+), 21 deletions(-) diff --git a/core/sources/civarchive.py b/core/sources/civarchive.py index eec5ec8..340a60e 100644 --- a/core/sources/civarchive.py +++ b/core/sources/civarchive.py @@ -51,6 +51,7 @@ from .common import ( build_custom_result_fields, collect_download_urls, is_remote_link_marked_dead, + resolve_file_size, ) log = create_module_logger(__name__) @@ -671,20 +672,16 @@ def _resolve_file_size_bytes( file_info: Dict[str, Any], download_urls: Optional[List[str]] = None, ) -> Optional[int]: - size = extract_file_size(file_info) - if size: - return size - urls = download_urls if download_urls is not None else _collect_archive_download_urls( file_info, **_NORMALIZED_DOWNLOAD_URL_OPTIONS, ) - for url in sorted(urls, key=_remote_size_probe_priority): - size = _fetch_remote_file_size_bytes(url) - if size: - return size - - return None + ordered_urls = sorted(urls, key=_remote_size_probe_priority) + return resolve_file_size( + file_info, + ordered_urls, + probe=_fetch_remote_file_size_bytes, + ) def _mirror_from_top_file(file_data: Dict[str, Any]) -> Optional[Dict[str, Any]]: diff --git a/core/sources/common.py b/core/sources/common.py index 8103c99..edc030b 100644 --- a/core/sources/common.py +++ b/core/sources/common.py @@ -2,9 +2,28 @@ Common utilities for external model metadata sources. """ -from typing import Any, Callable, Dict, List, Optional, Tuple +from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple -from ..type_utils import extract_file_sha256 +from ..type_utils import extract_file_sha256, extract_file_size + + +def resolve_file_size( + file_info: Dict[str, Any], + candidate_urls: Iterable[str], + *, + probe: Callable[[str], Optional[int]], +) -> Optional[int]: + """Resolve a model file size from metadata and ordered remote candidates.""" + size = extract_file_size(file_info) + if size: + return size + + for url in candidate_urls: + size = probe(url) + if size: + return size + + return None def build_custom_result_fields( diff --git a/core/sources/huggingface.py b/core/sources/huggingface.py index 2f9ab80..5575890 100644 --- a/core/sources/huggingface.py +++ b/core/sources/huggingface.py @@ -31,6 +31,7 @@ from ..type_utils import ( normalize_sha256, prepare_remote_size_probe_url, ) +from .common import resolve_file_size HF_API_URL = "https://huggingface.co/api" HF_AUTHOR_FALLBACKS = ["Comfy-Org", "Kijai"] @@ -721,9 +722,11 @@ def _build_huggingface_result( ) -> Dict[str, Any]: sha256 = _extract_huggingface_file_sha256(file_info) download_url = get_huggingface_download_url(repo_id, file_path) - size = extract_file_size(file_info) - if not size: - size = _fetch_remote_file_size_bytes(download_url, headers=headers) + size = resolve_file_size( + file_info, + [download_url], + probe=lambda url: _fetch_remote_file_size_bytes(url, headers=headers), + ) return build_model_result( "huggingface", model_id=repo_id, @@ -1552,13 +1555,15 @@ def build_huggingface_custom_result( file_info = candidate break - size = extract_file_size(file_info) - if not size: - size = fetch_remote_file_size_cached( - download_url, + size = resolve_file_size( + file_info, + [download_url], + probe=lambda url: fetch_remote_file_size_cached( + url, headers=headers, timeout=10, - ) + ), + ) sha256 = _extract_huggingface_file_sha256(file_info) def quote_url_path(val): diff --git a/tests/test_huggingface.py b/tests/test_huggingface.py index c01f777..88f0cf1 100644 --- a/tests/test_huggingface.py +++ b/tests/test_huggingface.py @@ -11,6 +11,7 @@ from core.sources.huggingface import ( HF_AUTHOR_FALLBACKS, _build_author_index_from_models, _fetch_author_index, + _fetch_remote_file_size_bytes, _find_matching_file_in_author_index, _get_author_index, _is_author_index_fresh, @@ -589,6 +590,23 @@ class HuggingFaceSourceTests(unittest.TestCase): normalized = _normalize_huggingface_size_probe_url(url) self.assertEqual(normalized, "https://huggingface.co/user/repo/resolve/main/model.safetensors") + @patch( + "core.sources.huggingface.fetch_remote_file_size_cached", + return_value=789, + ) + def test_remote_size_probe_preserves_headers_and_timeout(self, mock_fetch): + url = "https://huggingface.co/user/repo/blob/main/model.safetensors" + headers = {"Authorization": "Bearer test"} + + size = _fetch_remote_file_size_bytes(url, headers=headers, timeout=7) + + self.assertEqual(size, 789) + mock_fetch.assert_called_once_with( + "https://huggingface.co/user/repo/resolve/main/model.safetensors", + headers=headers, + timeout=7, + ) + def test_custom_result_uses_sha256_from_exact_huggingface_file(self): sha256 = "a" * 64 url = ( diff --git a/tests/test_source_internals.py b/tests/test_source_internals.py index ffd444d..020b89a 100644 --- a/tests/test_source_internals.py +++ b/tests/test_source_internals.py @@ -62,7 +62,11 @@ from core.sources.civitai import ( search_civitai, search_civitai_for_file, ) -from core.sources.common import build_custom_result_fields, is_remote_link_marked_dead +from core.sources.common import ( + build_custom_result_fields, + is_remote_link_marked_dead, + resolve_file_size, +) from core.type_utils import prepare_remote_size_probe_url # --------------------------------------------------------------------------- @@ -1152,6 +1156,57 @@ class PrepareSizeProbeUrlTests(unittest.TestCase): class CivArchiveMirrorSizeProbeTests(unittest.TestCase): + def test_shared_file_size_resolver_prefers_metadata(self): + probe = MagicMock(return_value=654321) + + size = resolve_file_size( + {"size": 123456}, + ["https://example.com/model.safetensors"], + probe=probe, + ) + + self.assertEqual(size, 123456) + probe.assert_not_called() + + def test_shared_file_size_resolver_uses_ordered_probe_until_success(self): + probe = MagicMock(side_effect=[None, 654321]) + urls = [ + "https://example.com/first.safetensors", + "https://example.com/second.safetensors", + "https://example.com/third.safetensors", + ] + + size = resolve_file_size({}, urls, probe=probe) + + self.assertEqual(size, 654321) + probe.assert_has_calls([ + ((urls[0],), {}), + ((urls[1],), {}), + ]) + self.assertEqual(probe.call_count, 2) + + def test_metadata_size_is_preferred_over_remote_probe(self): + with patch("core.sources.civarchive._fetch_remote_file_size_bytes") as mock_fetch: + size = _resolve_file_size_bytes( + {"size": 123456}, + ["https://huggingface.co/org/repo/resolve/main/model.safetensors"], + ) + + self.assertEqual(size, 123456) + mock_fetch.assert_not_called() + + @patch( + "core.sources.civarchive._fetch_remote_file_size_bytes", + return_value=654321, + ) + def test_zero_metadata_size_falls_back_to_remote_probe(self, mock_fetch): + url = "https://huggingface.co/org/repo/resolve/main/model.safetensors" + + size = _resolve_file_size_bytes({"size": 0}, [url]) + + self.assertEqual(size, 654321) + mock_fetch.assert_called_once_with(url) + @patch( "core.sources.civarchive.fetch_remote_file_size_cached", return_value=364853140,