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

refactor: complete migration to centralized HTTP helpers and shared normalization

Browse Source 
- Replace remaining direct requests.get calls with request_public_url,
  request_source_response, and request_source_json in downloader,
  civarchive, civitai, and popular sources
- Use normalize_model_file_info in CivitAI source instead of inline logic
- Update huggingface-author-index.json metadata
- Add tests for civarchive page/search requests and popular base model updates
- Adjust downloader redirect test for new request_public_url behavior
This commit is contained in:
Dariusz L
2026-07-15 13:04:00 +02:00
parent a5e08653af
commit 7e724d373f
7 changed files with 1061 additions and 1014 deletions
Download Patch File Download Diff File Expand all files Collapse all files
core/downloader.py
+5 -18
View File
@@ -228,34 +228,21 @@ def _resolve_civitai_download_url_for_aria2(
if not _is_civitai_api_download_url(validated_url):
return validated_url
response = None
try:
response = requests.get(
response, resolved_url, _ = request_public_url(
"GET",
validated_url,
headers=build_download_headers(validated_url, headers),
allow_redirects=False,
stream=True,
timeout=20,
stream=True,
max_redirects=1,
)
if response.status_code not in {301, 302, 303, 307, 308}:
return url
location = response.headers.get("Location") or response.headers.get("location")
if not location:
return url
resolved_url = validate_public_http_url(urljoin(validated_url, location.strip()))
response.close()
log.debug("Resolved CivitAI download redirect for aria2")
return resolved_url
except Exception as exc:
log.warning(f"Could not pre-resolve CivitAI download URL for aria2: {exc}")
return validated_url
finally:
if response is not None:
try:
response.close()
except Exception:
pass
def _resolve_download_url_for_aria2(
core/sources/civarchive.py
+20 -21
View File
@@ -14,7 +14,7 @@ from urllib.parse import parse_qs, unquote, urljoin, urlparse
import requests
from ..network_utils import host_matches_domain
from ..network_utils import host_matches_domain, request_source_response
from ..matcher import (
calculate_archived_model_confidence,
calculate_model_title_confidence,
@@ -155,16 +155,15 @@ def _request_page_text(
timeout: int = 20,
) -> Optional[str]:
url = urljoin(CIVARCHIVE_BASE_URL, path_or_url)
request_params = {k: v for k, v in (params or {}).items() if v is not None}
try:
response = requests.get(
url,
params=request_params,
headers=REQUEST_HEADERS,
timeout=timeout,
)
except Exception as e:
log.warning(f"CivArchive page request failed: url={url}, error={e}")
response = request_source_response(
url,
method="GET",
headers=REQUEST_HEADERS,
params=params,
timeout=timeout,
log_name="CivArchive page",
)
if response is None:
return None
if response.status_code != 200:
@@ -417,16 +416,16 @@ def _search_page(
# (especially VAEs) are frequently misclassified (e.g. as Checkpoint or null type) in their database.
# We rely on name confidence sorting instead.
try:
response = requests.get(
f"{CIVARCHIVE_BASE_URL}/search",
params=params,
headers=REQUEST_HEADERS,
timeout=timeout,
)
except Exception as e:
log.warning(f"CivArchive search request failed: query={query}, error={e}")
raise CivArchiveSearchError(str(e)) from e
response = request_source_response(
f"{CIVARCHIVE_BASE_URL}/search",
method="GET",
headers=REQUEST_HEADERS,
params=params,
timeout=timeout,
log_name="CivArchive search",
)
if response is None:
raise CivArchiveSearchError("CivArchive search request failed (network error or timeout)")
if response.status_code != 200:
log.warning(
core/sources/civitai.py
+8 -16
View File
@@ -35,6 +35,7 @@ from ..type_utils import (
as_list,
extract_trained_words,
build_search_result,
normalize_model_file_info,
)
from ..progress import report_progress, get_progress_reporter
from ..path_utils import (
@@ -1129,22 +1130,13 @@ def _normalize_civitai_file(
version_id: Optional[int],
api_key: Optional[str] = None,
) -> Dict[str, Any]:
hashes = file_info.get("hashes") if isinstance(file_info.get("hashes"), dict) else {}
return {
"id": file_info.get("id"),
"name": file_info.get("name") or file_info.get("filename"),
"type": file_info.get("type"),
"size": file_info.get("sizeKB", 0) * 1024 if file_info.get("sizeKB") else file_info.get("size"),
"download_url": file_info.get("downloadUrl") or (
get_civitai_download_url(version_id, api_key) if version_id else None
),
"primary": bool(file_info.get("primary", False)),
"sha256": file_info.get("sha256") or hashes.get("SHA256") or hashes.get("sha256"),
"hashes": hashes,
"metadata": file_info.get("metadata") or {},
"model_id": model_id,
"version_id": version_id,
}
fallback_url = get_civitai_download_url(version_id, api_key) if version_id else None
return normalize_model_file_info(
file_info,
model_id=model_id,
version_id=version_id,
fallback_download_url=fallback_url
)
def get_civitai_model_details(
core/sources/popular.py
+7 -11
View File
@@ -17,6 +17,7 @@ log = create_module_logger(__name__)
from ..path_utils import METADATA_DIR, read_json_safe, write_json_atomic, save_catalog_with_backup
from ..network_utils import request_source_json
POPULAR_MODELS_FILE = os.path.join(METADATA_DIR, "popular-models.json")
MODEL_ALIASES_FILE = os.path.join(METADATA_DIR, "model-aliases.json")
BASE_MODELS_FILE = os.path.join(METADATA_DIR, "base-models.json")
@@ -259,11 +260,9 @@ def get_base_models_status(check_remote: bool = False) -> Dict[str, Any]:
if check_remote:
try:
import requests
# Fetch from CivitAI API
response = requests.get("https://civitai.com/api/v1/enums", timeout=15)
if response.status_code == 200:
enums = response.json()
# Fetch from CivitAI API using unified helper
enums = request_source_json("https://civitai.com/api/v1/enums", timeout=15, log_name="CivitAI Enums")
if enums:
remote_models = enums.get("BaseModel", [])
# Check if there are new models by comparing normalized names and aliases
@@ -292,12 +291,9 @@ def get_base_models_status(check_remote: bool = False) -> Dict[str, Any]:
def update_base_models_from_remote() -> Dict[str, Any]:
"""Fetch live BaseModel enums from CivitAI and merge into base-models.json."""
import requests
response = requests.get("https://civitai.com/api/v1/enums", timeout=15)
if response.status_code != 200:
raise ValueError(f"CivitAI enums endpoint returned HTTP {response.status_code}")
enums = response.json()
enums = request_source_json("https://civitai.com/api/v1/enums", timeout=15, log_name="CivitAI Enums", raise_on_error=True)
if not enums:
raise ValueError("CivitAI enums did not return a valid response")
remote_names = enums.get("BaseModel", [])
if not remote_names or not isinstance(remote_names, list):
raise ValueError("CivitAI enums did not return a valid BaseModel list")
metadata/huggingface-author-index.json
+943 -943
View File
File diff suppressed because it is too large Load Diff
tests/test_refactoring_unification.py
+67
View File
@@ -91,6 +91,73 @@ class TestRefactoringUnification(unittest.TestCase):
res = search_civitai_by_hash("fake_hash")
self.assertIsNone(res)
@patch("requests.get")
def test_request_page_text_success(self, mock_get):
from core.sources.civarchive import _request_page_text
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.text = "<html>content</html>"
mock_get.return_value = mock_response
res = _request_page_text("test_path")
self.assertEqual(res, "<html>content</html>")
mock_get.assert_called_once()
@patch("requests.get")
def test_request_page_text_fail(self, mock_get):
from core.sources.civarchive import _request_page_text
mock_response = MagicMock()
mock_response.status_code = 404
mock_get.return_value = mock_response
res = _request_page_text("test_path")
self.assertIsNone(res)
mock_get.side_effect = Exception("network error")
res2 = _request_page_text("test_path")
self.assertIsNone(res2)
@patch("requests.get")
def test_get_base_models_status_remote(self, mock_get):
from core.sources.popular import get_base_models_status
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {"BaseModel": ["SD 1.5", "SDXL 1.0"]}
mock_get.return_value = mock_response
res = get_base_models_status(check_remote=True)
self.assertIn("local_count", res)
self.assertIn("update_available", res)
mock_get.assert_called_once_with("https://civitai.com/api/v1/enums", params={}, headers=None, timeout=15)
@patch("core.sources.popular.save_catalog_with_backup")
@patch("core.sources.popular.read_json_safe")
@patch("requests.get")
def test_update_base_models_from_remote_success(self, mock_get, mock_read_json, mock_save_catalog):
from core.sources.popular import update_base_models_from_remote
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.json.return_value = {"BaseModel": ["SD 1.5", "SDXL 1.0", "NewModel"]}
mock_get.return_value = mock_response
# Mock reading base-models.json
mock_read_json.return_value = {
"base_models": [
{"name": "SD 1.5", "aliases": ["sd1.5"]},
{"name": "SDXL 1.0", "aliases": ["sdxl"]},
]
}
res = update_base_models_from_remote()
self.assertTrue(res.get("updated"))
self.assertEqual(res.get("new_models_added"), 1)
self.assertIn("NewModel", res.get("new_models_added_list", []))
# Verify it saved the updated catalog
mock_save_catalog.assert_called_once()
saved_data = mock_save_catalog.call_args[0][1]
self.assertEqual(len(saved_data["base_models"]), 3)
if __name__ == "__main__":
unittest.main()
tests/test_resolver_robustness.py
+11 -5
View File
@@ -314,17 +314,23 @@ class ModelResolverRobustnessTests(unittest.TestCase):
api_url = "https://civitai.com/api/download/models/123"
redirect_url = "https://civitai-delivery.net/models/123/model.safetensors?token=abc"
mock_response = MagicMock()
mock_response.status_code = 302
mock_response.headers = {"Location": redirect_url}
mock_response1 = MagicMock()
mock_response1.status_code = 302
mock_response1.headers = {"Location": redirect_url}
with patch("requests.get", return_value=mock_response) as mock_get, patch(
mock_response2 = MagicMock()
mock_response2.status_code = 200
with patch("requests.get", side_effect=[mock_response1, mock_response2]) as mock_get, patch(
"core.downloader.validate_public_http_url",
side_effect=lambda value: value,
), patch(
"core.network_utils.validate_public_http_url",
side_effect=lambda value: value,
):
resolved = _resolve_civitai_download_url_for_aria2(api_url)
self.assertEqual(resolved, redirect_url)
mock_get.assert_called_once()
self.assertEqual(mock_get.call_count, 2)
# Scenario 2: Non-CivitAI URL (e.g. HuggingFace) should bypass redirects immediately
hf_url = "https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/model.safetensors"