From e568f81f2b14fe0e661df95b1174366dcaa3739f Mon Sep 17 00:00:00 2001 From: Dariusz L Date: Fri, 24 Jul 2026 21:42:11 +0200 Subject: [PATCH] feat(civarchive): prioritize known mirrors for remote size probes --- core/sources/civarchive.py | 39 +++++++++++++++-- tests/test_source_internals.py | 77 +++++++++++++++++++++++++++++----- 2 files changed, 101 insertions(+), 15 deletions(-) diff --git a/core/sources/civarchive.py b/core/sources/civarchive.py index e473126..b7d1f9f 100644 --- a/core/sources/civarchive.py +++ b/core/sources/civarchive.py @@ -9,7 +9,7 @@ import json import re import time from typing import Any, Callable, Dict, List, Optional, Tuple -from urllib.parse import urljoin +from urllib.parse import urljoin, urlparse from ..log_system import create_module_logger from ..matcher import ( @@ -27,7 +27,7 @@ from ..matcher import ( from ..matcher import ( normalize_base_model as _normalize_base_model, ) -from ..network_utils import request_source_response +from ..network_utils import host_matches_domain, request_source_response from ..path_utils import get_filename_from_path from ..progress import get_progress_reporter from ..type_utils import ( @@ -64,6 +64,22 @@ DEFAULT_CIVARCHIVE_CANDIDATE_LIMIT = 10 MAX_CIVARCHIVE_CANDIDATE_LIMIT = 30 SEARCH_RESULT_DETAIL_LIMIT = 5 HASH_PAGE_MODEL_LINK_LIMIT = 5 +CIVARCHIVE_SIZE_PROBE_DOMAINS = ( + "huggingface.co", + "civarchive.com", + "civitai.com", + "civitai.red", + "modelscope.ai", + "modelscope.cn", +) +CIVARCHIVE_SIZE_PROBE_PRIORITY = ( + "huggingface.co", + "civarchive.com", + "civitai.com", + "civitai.red", + "modelscope.ai", + "modelscope.cn", +) @@ -540,12 +556,27 @@ def _fetch_remote_file_size_bytes(url: Any, timeout: int = 15) -> Optional[int]: normalized = _normalize_download_url(url) if not normalized: return None - probe_url = prepare_remote_size_probe_url(normalized, ["civarchive.com", "civitai.com", "civitai.red"]) + probe_url = prepare_remote_size_probe_url( + normalized, + list(CIVARCHIVE_SIZE_PROBE_DOMAINS), + ) if not probe_url: return None return fetch_remote_file_size_cached(probe_url, headers=REQUEST_HEADERS, timeout=timeout) +def _remote_size_probe_priority(url: Any) -> int: + try: + host = urlparse(str(url or "")).hostname + except ValueError: + return len(CIVARCHIVE_SIZE_PROBE_PRIORITY) + + for priority, domain in enumerate(CIVARCHIVE_SIZE_PROBE_PRIORITY): + if host_matches_domain(host, domain): + return priority + return len(CIVARCHIVE_SIZE_PROBE_PRIORITY) + + def _resolve_file_size_bytes( file_info: Dict[str, Any], @@ -556,7 +587,7 @@ def _resolve_file_size_bytes( return size urls = download_urls if download_urls is not None else _collect_normalized_download_urls(file_info) - for url in urls: + for url in sorted(urls, key=_remote_size_probe_priority): size = _fetch_remote_file_size_bytes(url) if size: return size diff --git a/tests/test_source_internals.py b/tests/test_source_internals.py index a514a7b..7668c38 100644 --- a/tests/test_source_internals.py +++ b/tests/test_source_internals.py @@ -10,8 +10,8 @@ extract_trained_words), or where a refactored unified function changes its output contract. """ -import os import json +import os import tempfile import unittest from unittest.mock import MagicMock, patch @@ -22,38 +22,39 @@ import requests # civarchive internals # --------------------------------------------------------------------------- from core.sources.civarchive import ( - _normalize_download_url, _archive_link_is_dead, - _download_url_looks_like_model_file, _collect_download_urls, - _collect_normalized_download_urls, _collect_download_urls_unified, + _collect_normalized_download_urls, + _fetch_remote_file_size_bytes, _normalize_archive_version, + _normalize_download_url, + _resolve_file_size_bytes, parse_civarchive_url, ) -from core.type_utils import prepare_remote_size_probe_url # --------------------------------------------------------------------------- # civitai internals # --------------------------------------------------------------------------- from core.sources.civitai import ( - clear_search_cache, - build_civitai_session_cookie, - parse_civitai_url, - _normalize_civitai_file, - _extract_model_images, _build_civitai_result_from_version, _enrich_model_info_with_details, - _find_matching_file_in_versions, + _extract_model_images, _extract_public_api_model_candidates, _extract_trpc_model_candidates, + _find_matching_file_in_versions, + _normalize_civitai_file, _search_civitai_public_api_candidates, _search_civitai_red_candidates, _search_civitai_trpc_candidates, + build_civitai_session_cookie, + clear_search_cache, get_model_info_by_hash, get_model_info_for_file, + parse_civitai_url, search_civitai_for_file, ) +from core.type_utils import prepare_remote_size_probe_url # --------------------------------------------------------------------------- # workflow_updater @@ -535,6 +536,60 @@ class PrepareSizeProbeUrlTests(unittest.TestCase): self.assertIsNone(self._prepare_size_probe_url(None)) +class CivArchiveMirrorSizeProbeTests(unittest.TestCase): + + @patch( + "core.sources.civarchive.fetch_remote_file_size_cached", + return_value=364853140, + ) + def test_huggingface_mirror_is_allowed(self, mock_fetch): + url = ( + "https://huggingface.co/unsloth/LTX-2.3-GGUF/resolve/main/vae/" + "ltx-2.3-22b-distilled_audio_vae.safetensors" + ) + + self.assertEqual(_fetch_remote_file_size_bytes(url), 364853140) + mock_fetch.assert_called_once() + self.assertEqual(mock_fetch.call_args.args[0], url) + + @patch( + "core.sources.civarchive.fetch_remote_file_size_cached", + return_value=364853140, + ) + def test_modelscope_mirror_is_allowed(self, mock_fetch): + url = ( + "https://www.modelscope.ai/unsloth/LTX-2.3-GGUF/resolve/master/vae/" + "ltx-2.3-22b-distilled_audio_vae.safetensors" + ) + + self.assertEqual(_fetch_remote_file_size_bytes(url), 364853140) + mock_fetch.assert_called_once() + self.assertEqual(mock_fetch.call_args.args[0], url) + + @patch("core.sources.civarchive._fetch_remote_file_size_bytes") + def test_huggingface_is_probed_before_modelscope(self, mock_fetch): + modelscope_url = "https://www.modelscope.ai/org/repo/resolve/master/model.safetensors" + huggingface_url = "https://huggingface.co/org/repo/resolve/main/model.safetensors" + mock_fetch.side_effect = lambda url: 123456 if url == huggingface_url else None + + size = _resolve_file_size_bytes( + {}, + [modelscope_url, huggingface_url], + ) + + self.assertEqual(size, 123456) + mock_fetch.assert_called_once_with(huggingface_url) + + @patch("core.sources.civarchive.fetch_remote_file_size_cached") + def test_untrusted_mirror_is_not_probed(self, mock_fetch): + self.assertIsNone( + _fetch_remote_file_size_bytes( + "https://example.com/models/model.safetensors" + ) + ) + mock_fetch.assert_not_called() + + # =========================================================================== # civitai - parse_civitai_url