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

refactor: unify provider JSON requests and move settings schema to dynamic loading

Browse Source
This commit is contained in:
Dariusz L
2026-07-18 15:10:34 +02:00
parent 8cc02010b0
commit 79e7aa3743
7 changed files with 83 additions and 78 deletions
Download Patch File Download Diff File Expand all files Collapse all files
core/sources/civarchive.py
+4 -24
View File
@@ -9,7 +9,7 @@ import json
import re
import time
from typing import Any, Callable, Dict, List, Optional, Tuple
from urllib.parse import urljoin, urlparse
from urllib.parse import urljoin
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 host_matches_domain, request_source_response
from ..network_utils import request_source_response
from ..path_utils import get_filename_from_path
from ..progress import get_progress_reporter
from ..type_utils import (
@@ -535,31 +535,11 @@ def _archive_link_is_dead(item: Dict[str, Any]) -> bool:
)
def _prepare_size_probe_url(url: Any) -> Optional[str]:
def _fetch_remote_file_size_bytes(url: Any, timeout: int = 15) -> Optional[int]:
normalized = _normalize_download_url(url)
if not normalized:
return None
normalized = prepare_remote_size_probe_url(normalized)
if not normalized:
return None
parsed = urlparse(normalized)
host = parsed.hostname
path = parsed.path or ""
if host_matches_domain(host, "civarchive.com") and not path.startswith("/api/download/"):
return None
if host_matches_domain(host, "civitai.com", "civitai.red") and not path.startswith("/api/download/"):
return None
return normalized
def _fetch_remote_file_size_bytes(url: Any, timeout: int = 15) -> Optional[int]:
probe_url = _prepare_size_probe_url(url)
probe_url = prepare_remote_size_probe_url(normalized, ["civarchive.com", "civitai.com", "civitai.red"])
if not probe_url:
return None
return fetch_remote_file_size_cached(probe_url, headers=REQUEST_HEADERS, timeout=timeout)
core/sources/civitai.py
+7 -7
View File
@@ -28,7 +28,7 @@ from ..matcher import (
from ..matcher import (
normalize_base_model as _normalize_base_model,
)
from ..network_utils import execute_provider_json_request, request_source_json, request_source_response
from ..network_utils import execute_provider_json_request, request_source_response
from ..path_utils import (
calculate_file_sha256,
extract_safetensors_header_metadata,
@@ -1056,12 +1056,12 @@ def search_civitai(
if civitai_type:
params["types"] = civitai_type
headers = {}
if api_key:
headers["Authorization"] = f"Bearer {api_key}"
data = request_source_json(
f"{CIVITAI_API_URL}/models", params=params, headers=headers, timeout=15, log_name="CivitAI search models"
data = execute_provider_json_request(
"CivitAI search models",
f"{CIVITAI_API_URL}/models",
params=params,
api_key=api_key,
timeout=15,
)
if data:
core/sources/huggingface.py
+17 -12
View File
@@ -17,7 +17,7 @@ from ..progress import get_progress_reporter
log = create_module_logger(__name__)
from ..matcher import build_filename_search_queries
from ..network_utils import host_matches_domain, request_source_json, request_source_response
from ..network_utils import execute_provider_json_request, host_matches_domain, request_source_response
from ..path_utils import METADATA_DIR, get_filename_from_path, read_json_safe, write_json_atomic
from ..type_utils import (
build_search_result,
@@ -696,14 +696,19 @@ def search_huggingface_for_file(
query_count=total_queries,
)
search_url = f"{HF_API_URL}/models?search={quote(search_query)}&limit=20"
response = request_source_response(search_url, headers=headers, timeout=10, log_name="HuggingFace models search")
if response.status_code != 200:
repos = execute_provider_json_request(
"HuggingFace models search",
search_url,
api_key=token,
timeout=10,
)
if repos is None:
log.warning(
f"HuggingFace API status={response.status_code} query={search_query}"
f"HuggingFace API query={search_query} failed"
)
continue
repos = response.json()
if not repos:
continue
log.info(f"HuggingFace API query={search_query} repos={len(repos)}")
for repo in repos:
@@ -970,14 +975,14 @@ def search_huggingface(
results = []
try:
headers = {}
if token:
headers["Authorization"] = f"Bearer {token}"
params = {"search": query, "limit": limit, "full": "true"}
data = request_source_json(
f"{HF_API_URL}/models", params=params, headers=headers, timeout=15, log_name="HuggingFace general search"
data = execute_provider_json_request(
"HuggingFace general search",
f"{HF_API_URL}/models",
params=params,
api_key=token,
timeout=15,
)
if data:
core/type_utils.py
+7 -1
View File
@@ -893,7 +893,7 @@ def extract_trained_words(*values: Any) -> List[str]:
_remote_size_cache: Dict[tuple[str, Optional[str]], Optional[int]] = {}
def prepare_remote_size_probe_url(url: Any) -> Optional[str]:
def prepare_remote_size_probe_url(url: Any, allowed_domains: Optional[List[str]] = None) -> Optional[str]:
"""
Normalize a remote URL to be suitable for file size probing.
Converts HuggingFace '/blob/' URLs to '/resolve/' URLs.
@@ -916,6 +916,12 @@ def prepare_remote_size_probe_url(url: Any) -> Optional[str]:
except Exception:
return None
if allowed_domains and host:
if not host_matches_domain(host, *allowed_domains):
return None
if host_matches_domain(host, "civarchive.com", "civitai.com", "civitai.red") and not path.startswith("/api/download/"):
return None
if host and host_matches_domain(host, "huggingface.co") and "/blob/" in path:
value = value.replace("/blob/", "/resolve/", 1)
tests/test_refactoring_unification.py
+38
View File
@@ -283,6 +283,44 @@ class TestRefactoringUnification(unittest.TestCase):
self.assertEqual(data, {"key": "value"})
mock_request_source_json.assert_called_once()
def test_prepare_remote_size_probe_url_with_allowed_domains(self):
from core.type_utils import prepare_remote_size_probe_url
# Test default behavior
self.assertEqual(
prepare_remote_size_probe_url("https://huggingface.co/user/repo/blob/main/model.safetensors"),
"https://huggingface.co/user/repo/resolve/main/model.safetensors"
)
# Test allowed domains filtering
self.assertEqual(
prepare_remote_size_probe_url("https://civitai.com/api/download/models/123", ["civitai.com", "civitai.red"]),
"https://civitai.com/api/download/models/123"
)
self.assertIsNone(
prepare_remote_size_probe_url("https://otherdomain.com/model.safetensors", ["civitai.com", "civitai.red"])
)
@patch("core.network_utils.request_source_json")
def test_execute_provider_json_request(self, mock_request_source_json):
from core.network_utils import execute_provider_json_request
mock_request_source_json.return_value = {"response": "ok"}
# Test Bearer Authorization and Cookie header injection
res = execute_provider_json_request(
"Test Provider",
"https://example.com/api",
api_key="my_secret_token",
session_token="my_session"
)
self.assertEqual(res, {"response": "ok"})
mock_request_source_json.assert_called_once()
call_args = mock_request_source_json.call_args
headers = call_args[1]["headers"]
self.assertEqual(headers["Authorization"], "Bearer my_secret_token")
self.assertIn("__Secure-civ-token=my_session", headers["Cookie"])
if __name__ == "__main__":
tests/test_source_internals.py
+9 -6
View File
@@ -28,8 +28,8 @@ from core.sources.civarchive import (
_collect_download_urls_unified,
_normalize_archive_version,
parse_civarchive_url,
_prepare_size_probe_url,
)
from core.type_utils import prepare_remote_size_probe_url
# ---------------------------------------------------------------------------
# civitai internals
@@ -462,24 +462,27 @@ class ParseCivarchiveUrlTests(unittest.TestCase):
# ===========================================================================
class PrepareSizeProbeUrlTests(unittest.TestCase):
def _prepare_size_probe_url(self, url):
return prepare_remote_size_probe_url(url, ["civarchive.com", "civitai.com", "civitai.red"])
def test_civarchive_api_download_url_passes(self):
url = "https://civarchive.com/api/download/models/77.safetensors"
self.assertEqual(_prepare_size_probe_url(url), url)
self.assertEqual(self._prepare_size_probe_url(url), url)
def test_civarchive_non_download_url_blocked(self):
url = "https://civarchive.com/models/77"
self.assertIsNone(_prepare_size_probe_url(url))
self.assertIsNone(self._prepare_size_probe_url(url))
def test_civitai_non_download_url_blocked(self):
url = "https://civitai.com/models/123"
self.assertIsNone(_prepare_size_probe_url(url))
self.assertIsNone(self._prepare_size_probe_url(url))
def test_civitai_api_download_url_passes(self):
url = "https://civitai.com/api/download/models/99.safetensors"
self.assertEqual(_prepare_size_probe_url(url), url)
self.assertEqual(self._prepare_size_probe_url(url), url)
def test_none_returns_none(self):
self.assertIsNone(_prepare_size_probe_url(None))
self.assertIsNone(self._prepare_size_probe_url(None))
web/resolver/views/options_methods.js
+1 -28
View File
@@ -5,34 +5,7 @@ import { getSvgIcon } from "../../utils/icon_utils.js";
import { LOG_LEVEL as DEFAULT_FRONTEND_LOG_LEVEL } from "../../log_system/config.js";
import { logger as frontendLogger } from "../../log_system/logger.js";
import { pollBackgroundTask, safeStorage } from "../utils/html_utils.js";
const SETTINGS_MAP = [
{ serverKey: 'civitai_key', localKey: 'ModelResolver.civitaiApiKey', type: 'string', default: '' },
{ serverKey: 'civitai_session_token', localKey: 'ModelResolver.civitaiSessionToken', type: 'string', default: '' },
{ serverKey: 'civitai_use_trpc_search', localKey: 'ModelResolver.civitaiUseTrpcSearch', type: 'boolean', default: true },
{ serverKey: 'civitai_use_api_search', localKey: 'ModelResolver.civitaiUseApiSearch', type: 'boolean', default: true },
{ serverKey: 'civitai_use_html_fallback', localKey: 'ModelResolver.civitaiUseHtmlFallback', type: 'boolean', default: true },
{ serverKey: 'hf_token', localKey: 'ModelResolver.huggingFaceToken', type: 'string', default: '' },
{ serverKey: 'brave_search_api_key', localKey: 'ModelResolver.braveSearchApiKey', type: 'string', default: '' },
{ serverKey: 'hf_use_api_search', localKey: 'ModelResolver.hfUseApiSearch', type: 'boolean', default: true },
{ serverKey: 'hf_use_comfy_org_fallback', localKey: 'ModelResolver.hfUseComfyOrgFallback', type: 'boolean', default: true },
{ serverKey: 'hf_use_brave_fallback', localKey: 'ModelResolver.hfUseBraveFallback', type: 'boolean', default: true },
{ serverKey: 'auto_fill_base_model', localKey: 'ModelResolver.autoFillBaseModel', type: 'boolean', default: true },
{ serverKey: 'auto_fill_subfolder', localKey: 'ModelResolver.autoFillSubfolder', type: 'boolean', default: true },
{ serverKey: 'auto_refresh_comfy_models_after_apply', localKey: 'ModelResolver.autoRefreshComfyModelsAfterApply', type: 'boolean', default: true },
{ serverKey: 'workflow_hash_metadata_enabled', localKey: 'ModelResolver.workflowHashMetadataEnabled', type: 'boolean', default: true },
{ serverKey: 'workflow_dependency_marker_enabled', localKey: 'ModelResolver.workflowDependencyMarkerEnabled', type: 'boolean', default: false },
{ serverKey: 'download_backend', localKey: 'ModelResolver.downloadBackend', type: 'backend', default: 'python' },
{ serverKey: 'aria2c_path', localKey: 'ModelResolver.aria2cPath', type: 'string', default: '' },
{ serverKey: 'aria2_auto_stop_daemon', localKey: 'ModelResolver.aria2AutoStopDaemon', type: 'boolean', default: true },
{ serverKey: 'download_path_mode', localKey: 'ModelResolver.downloadPathMode', type: 'pathMode', default: 'suggested' },
{ serverKey: 'download_path_templates', localKey: 'ModelResolver.downloadPathTemplates', type: 'json', default: {} },
{ serverKey: 'base_model_path_mappings', localKey: 'ModelResolver.baseModelPathMappings', type: 'json', default: {} },
{ serverKey: 'frontend_logs_enabled', localKey: 'ModelResolver.frontendLogsEnabled', type: 'frontendLogsEnabled', default: true },
{ serverKey: 'backend_logs_enabled', localKey: 'ModelResolver.backendLogsEnabled', type: 'backendLogsEnabled', default: true },
{ serverKey: 'frontend_log_level', localKey: 'ModelResolver.frontendLogLevel', type: 'string', default: 'ERROR' },
{ serverKey: 'backend_log_level', localKey: 'ModelResolver.backendLogLevel', type: 'string', default: 'ERROR' },
{ serverKey: 'civitai_candidate_limit', localKey: 'ModelResolver.civitaiCandidateLimit', type: 'candidateLimit', default: 5 },
];
const SETTINGS_MAP = [];
export const optionsMethods = {
displayOptions() {