diff --git a/__init__.py b/__init__.py index b7c31e1..5492320 100644 --- a/__init__.py +++ b/__init__.py @@ -258,15 +258,17 @@ class ModelResolverExtension: ) -> None: if not loaded_id: return - self.loaded_progress.update( + self.loaded_progress.update_from_payload( loaded_id, - status=status, - stage=stage, - message=message, - percent=percent, - current=current, - total=total, - **payload, + { + "stage": stage, + "message": message, + "percent": percent, + "status": status, + "current": current, + "total": total, + **payload, + }, ) def _update_workflow_analysis_progress( @@ -657,13 +659,15 @@ class ModelResolverExtension: ) if analysis_id: - self.analysis_progress.update( + self._update_analysis_progress( analysis_id, - status="starting", - stage="starting", - message="Starting analysis...", - current=0, - total=0, + { + "status": "starting", + "stage": "starting", + "message": "Starting analysis...", + "current": 0, + "total": 0, + }, ) def update_analysis_progress(payload): diff --git a/core/catalog_manager.py b/core/catalog_manager.py index b27e5fb..21a767b 100644 --- a/core/catalog_manager.py +++ b/core/catalog_manager.py @@ -21,3 +21,27 @@ class CatalogManager: def save(self, data: Any, meta: Any, indent: int = 2) -> None: save_catalog_with_backup(self.data_file, data, self.meta_file, meta, indent=indent) + + def get_items(self) -> list: + data = self.read_data() + if isinstance(data, dict): + items = data.get(self.root_key, []) + return items if isinstance(items, list) else [] + return [] + + def get_local_count(self) -> int: + return len(self.get_items()) + + def sync_catalog( + self, + new_data: Any, + meta_updates: Dict[str, Any], + indent: int = 2, + ) -> Dict[str, Any]: + """Save updated catalog payload and sidecar metadata.""" + data = {self.root_key: new_data} if not isinstance(new_data, dict) else new_data + self.save(data, meta_updates, indent=indent) + return { + "updated": True, + **meta_updates, + } diff --git a/core/matcher.py b/core/matcher.py index 42c17ef..f6c3e0b 100644 --- a/core/matcher.py +++ b/core/matcher.py @@ -239,8 +239,8 @@ def normalize_base_model(value: str) -> str: Returns: Normalized base model name """ - from .type_utils import normalize_alphanumeric_lower - return normalize_alphanumeric_lower(value) + from .type_utils import normalize_alphanumeric_key + return normalize_alphanumeric_key(value) def base_model_matches(candidate: str, preferred: Optional[str]) -> bool: diff --git a/core/path_templates.py b/core/path_templates.py index db0c4d8..c211be5 100644 --- a/core/path_templates.py +++ b/core/path_templates.py @@ -46,8 +46,8 @@ TAG_FOLDER_TOKENS = { def _normalize_token(value: Any) -> str: - from .type_utils import normalize_alphanumeric_lower - return normalize_alphanumeric_lower(value) + from .type_utils import normalize_alphanumeric_key + return normalize_alphanumeric_key(value) def _base_model_token_variants(value: Any) -> List[str]: diff --git a/core/path_utils.py b/core/path_utils.py index 5274bb6..d5c11b9 100644 --- a/core/path_utils.py +++ b/core/path_utils.py @@ -406,8 +406,8 @@ def extract_safetensors_header_sha256( def _normalize_base_model_token(value: Any) -> str: - from .type_utils import normalize_alphanumeric_lower - return normalize_alphanumeric_lower(value) + from .type_utils import normalize_alphanumeric_key + return normalize_alphanumeric_key(value) def _load_base_model_aliases() -> List[Tuple[str, str]]: diff --git a/core/settings.py b/core/settings.py index 75898f6..2319b9a 100644 --- a/core/settings.py +++ b/core/settings.py @@ -238,8 +238,8 @@ def _listify_tags(value: Any) -> List[str]: def _normalize_tag(value: str) -> str: - from .type_utils import normalize_alphanumeric_lower - return normalize_alphanumeric_lower(value) + from .type_utils import normalize_alphanumeric_key + return normalize_alphanumeric_key(value) @@ -274,7 +274,8 @@ def _resolve_base_model_mapping(mappings: Mapping[str, str], base_model: Any) -> def _first_tag(tags: Any) -> str: - tag_list = list(_listify_tags(tags)) + from .type_utils import as_list + tag_list = list(as_list(tags)) if not tag_list: return "no tags" normalized_tags = {tag: _normalize_tag(tag) for tag in tag_list} diff --git a/core/sources/civarchive.py b/core/sources/civarchive.py index 320a0a5..e473126 100644 --- a/core/sources/civarchive.py +++ b/core/sources/civarchive.py @@ -37,6 +37,7 @@ from ..type_utils import ( extract_trained_words, fetch_remote_file_size_cached, looks_like_model_file, + normalize_alphanumeric_key, normalize_model_file_info, normalize_model_image, normalize_sha256, @@ -126,14 +127,14 @@ def _request_json( params: Optional[Dict[str, Any]] = None, timeout: int = 20, ) -> Optional[Dict[str, Any]]: - from ..network_utils import request_source_json + from ..network_utils import execute_provider_json_request url = f"{CIVARCHIVE_API_URL}{path}" - return request_source_json( + return execute_provider_json_request( + "CivArchive API", url, params=params, headers=REQUEST_HEADERS, timeout=timeout, - log_name="CivArchive API" ) @@ -178,7 +179,7 @@ def _normalize_civarchive_type(value: Any) -> str: return "" text = str(value).strip() mapped = resolve_model_category(text, target_format="civarchive") - return re.sub(r"[^a-z0-9]+", "", str(mapped).lower()) + return normalize_alphanumeric_key(mapped) def _extract_hash_page_model_cards(next_data: Dict[str, Any]) -> List[Dict[str, Any]]: diff --git a/core/sources/civitai.py b/core/sources/civitai.py index 62d8d3e..fbe3ea6 100644 --- a/core/sources/civitai.py +++ b/core/sources/civitai.py @@ -49,6 +49,7 @@ from ..type_utils import ( normalize_model_file_info, normalize_model_image, parse_size_to_bytes, + prepare_remote_size_probe_url, resolve_model_category, select_primary_model_file, ) @@ -58,6 +59,15 @@ log = create_module_logger(__name__) CIVITAI_API_URL = "https://civitai.com/api/v1" + +def _fetch_civitai_remote_file_size(url: str, timeout: int = 15) -> Optional[int]: + """Probe file size for Civitai remote URLs.""" + from ..type_utils import fetch_remote_file_size_cached + probe_url = prepare_remote_size_probe_url(url, ["civitai.com", "civitai.red"]) + if not probe_url: + return None + return fetch_remote_file_size_cached(probe_url, timeout=timeout) + # Cache for search results and URN resolutions _search_cache: Dict[str, Any] = {} _urn_cache: Dict[tuple[int, int], Dict[str, Any]] = {} @@ -1135,21 +1145,14 @@ def get_civitai_model_details( if not model_id: return None - headers = {} - if api_key: - headers["Authorization"] = f"Bearer {api_key}" - - response = request_source_response( + data = execute_provider_json_request( + "CivitAI model details", f"{CIVITAI_API_URL}/models/{model_id}", - headers=headers, + api_key=api_key, timeout=20, - log_name="CivitAI model details" ) - if response.status_code != 200: - log.warning(f"CivitAI model details returned {response.status_code}: model_id={model_id}") + if not data: return None - - data = response.json() versions = data.get("modelVersions") or [] normalized_versions = [] selected_version = None @@ -1273,21 +1276,16 @@ def resolve_urn( return _urn_cache[cache_key] try: - headers = {} - if api_key: - headers["Authorization"] = f"Bearer {api_key}" - - url = f"{CIVITAI_API_URL}/models/{model_id}" - params = {"modelVersionId": version_id} - - response = request_source_response(url, headers=headers, params=params, timeout=15, log_name="CivitAI URN resolve") - - if response.status_code != 200: - log.warning(f"CivitAI URN resolve failed: {response.status_code}") + data = execute_provider_json_request( + "CivitAI URN resolve", + f"{CIVITAI_API_URL}/models/{model_id}", + params={"modelVersionId": version_id}, + api_key=api_key, + timeout=15, + ) + if not data: _urn_cache[cache_key] = None return None - - data = response.json() versions = data.get("modelVersions", []) if not versions: @@ -1387,15 +1385,14 @@ def get_model_info_by_hash( api_url = f"{CIVITAI_API_URL}/model-versions/by-hash/{file_hash}" try: - headers = {} - if api_key: - headers["Authorization"] = f"Bearer {api_key}" - - response = request_source_response(api_url, headers=headers, timeout=15, log_name="CivitAI by-hash") - - if response.status_code == 200: - data = response.json() + data = execute_provider_json_request( + "CivitAI by-hash", + api_url, + api_key=api_key, + timeout=15, + ) + if data: # Extract useful info from the response model_info = data.get("model", {}) version_info = data @@ -1438,14 +1435,9 @@ def get_model_info_by_hash( log.info(f"Found model by hash {file_hash}: {result.get('model_name')}") return result - elif response.status_code == 404: - log.info(f"Model not found on CivitAI for hash {file_hash}") - _hash_cache[cache_key] = None - return None else: - log.warning( - f"CivitAI hash lookup returned {response.status_code} for {file_hash}" - ) + log.info(f"Model not found or lookup failed on CivitAI for hash {file_hash}") + _hash_cache[cache_key] = None return None except Exception as e: diff --git a/core/sources/huggingface.py b/core/sources/huggingface.py index c50efc6..7cd69c0 100644 --- a/core/sources/huggingface.py +++ b/core/sources/huggingface.py @@ -196,23 +196,18 @@ def _build_author_index_from_models( def _fetch_author_index(author: str, headers: Dict[str, str], limit: int = 200): try: - response = request_source_response( + repos = execute_provider_json_request( + "HuggingFace author index", f"{HF_API_URL}/models", params={"author": author, "limit": limit, "full": "true"}, headers=headers, timeout=15, - log_name="HuggingFace author index" ) - if response.status_code != 200: - log.debug( - f"HuggingFace author index returned {response.status_code} for author={author}" - ) - return None - - repos = response.json() if not isinstance(repos, list): return None + return repos + index = _build_author_index_from_models(author, repos) log.info( f"HuggingFace author index refreshed for author={author}: repos={index['repo_count']}, files={index['file_count']}" @@ -393,19 +388,16 @@ def _get_repo_tree( files_url = f"{HF_API_URL}/models/{repo_id}/tree/{quote(branch, safe='')}" try: - response = request_source_response( + data = execute_provider_json_request( + "HuggingFace tree", files_url, params={"recursive": "1"}, headers=headers, timeout=15, - log_name="HuggingFace tree" ) - if response.status_code != 200: - log.debug( - f"HuggingFace tree request returned {response.status_code} for repo={repo_id}, branch={branch}" - ) - return None - return response.json() + if isinstance(data, list): + return data + return None except Exception as e: log.debug(f"Error getting HuggingFace tree for {repo_id}@{branch}: {e}") return None diff --git a/core/sources/model_list.py b/core/sources/model_list.py index c33ff58..7cef5fb 100644 --- a/core/sources/model_list.py +++ b/core/sources/model_list.py @@ -163,7 +163,7 @@ def update_model_list_from_remote() -> Dict[str, Any]: "model_count": len(models), "updated_at": _utc_now_iso(), } - catalog_mgr.save(data, meta, indent=2) + catalog_mgr.sync_catalog(data, meta, indent=2) reload_model_list() return { diff --git a/core/sources/popular.py b/core/sources/popular.py index d419172..0b5ac27 100644 --- a/core/sources/popular.py +++ b/core/sources/popular.py @@ -369,9 +369,9 @@ def update_base_models_from_remote() -> Dict[str, Any]: "last_checked_at": now_str, "local_count": len(updated_models), "new_models_added": new_added_count, - "new_models_added_list": new_added_names + "new_models_added_list": new_added_names, } - base_models_mgr.save({"base_models": updated_models}, meta, indent=2) + base_models_mgr.sync_catalog({"base_models": updated_models}, meta, indent=2) reload_databases() return { diff --git a/core/type_utils.py b/core/type_utils.py index f12c63f..773641f 100644 --- a/core/type_utils.py +++ b/core/type_utils.py @@ -21,17 +21,18 @@ def as_dict(value: Any) -> Dict[str, Any]: return value if isinstance(value, dict) else {} -def as_list(value: Any) -> List[Any]: +def as_list(value: Any, pattern: str = r"[,;]+") -> List[Any]: """ Safely cast value to a list. Supports: - lists (returns list filtered of None/empty strings) - tuples and sets (casts to list and filters) - - comma-separated or semicolon-separated strings (splits, strips, and filters) + - delimited strings (splits by pattern, strips, and filters) Args: value: Value to cast + pattern: Regex pattern for string splitting (default: r"[,;]+") Returns: A list of elements. @@ -41,7 +42,7 @@ def as_list(value: Any) -> List[Any]: if isinstance(value, (tuple, set)): return [item for item in value if item not in (None, "")] if isinstance(value, str): - return [item.strip() for item in re.split(r"[,;]+", value) if item.strip()] + return [item.strip() for item in re.split(pattern, value) if item.strip()] return [] @@ -1290,11 +1291,15 @@ def parse_provider_model_url(url: str, allowed_domains: list[str]) -> Optional[D return parse_civitai_model_path(parsed.path, parsed.query) -def normalize_alphanumeric_lower(value: Any) -> str: +def normalize_alphanumeric_key(value: Any) -> str: """Normalize a string or token to standard lowercase alphanumeric format.""" return re.sub(r"[^a-z0-9]+", "", str(value or "").lower()) +normalize_alphanumeric_lower = normalize_alphanumeric_key + + + def utc_now_iso() -> str: """Return the current UTC timestamp formatted as ISO 8601 string without microseconds.""" from datetime import datetime, timezone diff --git a/tests/test_refactoring_unification.py b/tests/test_refactoring_unification.py index 5bd834d..5f76e1f 100644 --- a/tests/test_refactoring_unification.py +++ b/tests/test_refactoring_unification.py @@ -127,8 +127,8 @@ class TestRefactoringUnification(unittest.TestCase): self.assertIn("NewModel", res.get("new_models_added_list", [])) # Verify it saved the updated catalog - mock_base_models_mgr.save.assert_called_once() - saved_data = mock_base_models_mgr.save.call_args[0][0] + mock_base_models_mgr.sync_catalog.assert_called_once() + saved_data = mock_base_models_mgr.sync_catalog.call_args[0][0] self.assertEqual(len(saved_data["base_models"]), 3) @@ -171,7 +171,7 @@ class TestRefactoringUnification(unittest.TestCase): mock_fetch_json.return_value = {"models": [{"name": "ModelA"}, {"name": "ModelB"}]} res = update_model_list_from_remote() self.assertTrue(res.get("updated")) - self.assertTrue(mock_catalog_mgr.save.called) + self.assertTrue(mock_catalog_mgr.sync_catalog.called) def test_tracker_progress_updates(self): diff --git a/web/resolver/utils/html_utils.js b/web/resolver/utils/html_utils.js index 724c737..0ab08b6 100644 --- a/web/resolver/utils/html_utils.js +++ b/web/resolver/utils/html_utils.js @@ -16,6 +16,23 @@ export function escapeHtml(value) { .replace(/'/g, '''); } +/** + * Tagged template literal for auto-escaping interpolation values in HTML strings. + * @param {TemplateStringsArray} strings + * @param {...any} values + * @returns {string} + */ +export function html(strings, ...values) { + return strings.reduce((result, string, i) => { + const value = values[i - 1]; + let valStr = ''; + if (value !== undefined && value !== null) { + valStr = Array.isArray(value) ? value.join('') : escapeHtml(value); + } + return result + valStr + string; + }); +} + /** * Escapes a value to be safely embedded inside inline JavaScript. * @param {any} value diff --git a/web/resolver/views/options_methods.js b/web/resolver/views/options_methods.js index 1b0c40f..d36ebdf 100644 --- a/web/resolver/views/options_methods.js +++ b/web/resolver/views/options_methods.js @@ -4,10 +4,13 @@ import { $el } from "../../../../../scripts/ui.js"; 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"; +import { escapeHtml, pollBackgroundTask, safeStorage } from "../utils/html_utils.js"; const SETTINGS_MAP = []; export const optionsMethods = { + escapeHtml(value) { + return escapeHtml(value); + }, displayOptions() { if (!this.contentElement) return; this.contentElement.style.overflowY = 'hidden'; diff --git a/web/resolver/views/tabs_loaded_methods.js b/web/resolver/views/tabs_loaded_methods.js index f6f7678..9aeea39 100644 --- a/web/resolver/views/tabs_loaded_methods.js +++ b/web/resolver/views/tabs_loaded_methods.js @@ -2,7 +2,11 @@ import { app } from "../../../../../scripts/app.js"; import { api } from "../../../../../scripts/api.js"; import { $el } from "../../../../../scripts/ui.js"; import { getSvgIcon } from "../../utils/icon_utils.js"; +import { escapeHtml } from "../utils/html_utils.js"; export const tabsLoadedMethods = { + escapeHtml(value) { + return escapeHtml(value); + }, getTabButton(tab) { return { missing: this.missingTab,