refactor: centralize base model token collection
This commit is contained in:
+13
-11
@@ -6,7 +6,7 @@ Curated list of common models with known download URLs.
|
||||
|
||||
import os
|
||||
import re
|
||||
from typing import Any, Dict, List, Optional
|
||||
from typing import Any, Dict, List, Optional, Set
|
||||
|
||||
from ..log_system import create_module_logger
|
||||
|
||||
@@ -257,6 +257,16 @@ def generate_aliases(name: str) -> List[str]:
|
||||
return sorted(list({a for a in aliases if len(a) > 1}))
|
||||
|
||||
|
||||
def _collect_normalized_base_model_tokens(base_models: List[Dict[str, Any]]) -> Set[str]:
|
||||
"""Collect normalized names and aliases from the base-model catalog."""
|
||||
normalized_tokens = set()
|
||||
for model in base_models:
|
||||
normalized_tokens.add(normalize_alphanumeric_key(model.get("name", "")))
|
||||
for alias in model.get("aliases", []):
|
||||
normalized_tokens.add(normalize_alphanumeric_key(alias))
|
||||
return normalized_tokens
|
||||
|
||||
|
||||
def get_base_models_status(check_remote: bool = False) -> Dict[str, Any]:
|
||||
"""Return local base-models metadata and optionally compare with CivitAI."""
|
||||
local_data = _read_base_models_file()
|
||||
@@ -281,11 +291,7 @@ def get_base_models_status(check_remote: bool = False) -> Dict[str, Any]:
|
||||
remote_models = enums.get("BaseModel", [])
|
||||
|
||||
# Check if there are new models by comparing normalized names and aliases
|
||||
all_existing_normalized = set()
|
||||
for m in base_models:
|
||||
all_existing_normalized.add(normalize_alphanumeric_key(m.get("name", "")))
|
||||
for alias in m.get("aliases", []):
|
||||
all_existing_normalized.add(normalize_alphanumeric_key(alias))
|
||||
all_existing_normalized = _collect_normalized_base_model_tokens(base_models)
|
||||
|
||||
new_models_found = False
|
||||
for remote_name in remote_models:
|
||||
@@ -318,11 +324,7 @@ def update_base_models_from_remote() -> Dict[str, Any]:
|
||||
base_models = local_data.get("base_models", [])
|
||||
|
||||
# 1. Build a set of all normalized aliases currently in base-models.json
|
||||
all_known_normalized = set()
|
||||
for m in base_models:
|
||||
for alias in m.get("aliases", []):
|
||||
all_known_normalized.add(normalize_alphanumeric_key(alias))
|
||||
all_known_normalized.add(normalize_alphanumeric_key(m.get("name", "")))
|
||||
all_known_normalized = _collect_normalized_base_model_tokens(base_models)
|
||||
|
||||
updated_models = list(base_models)
|
||||
new_added_count = 0
|
||||
|
||||
@@ -151,6 +151,44 @@ class TestRefactoringUnification(unittest.TestCase):
|
||||
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.request_source_json")
|
||||
@patch("core.sources.popular._read_base_models_file")
|
||||
@patch("core.sources.popular._read_base_models_meta")
|
||||
@patch("core.sources.popular.base_models_mgr")
|
||||
@patch("core.sources.popular.reload_databases")
|
||||
def test_base_model_alias_matching_is_consistent_for_status_and_update(
|
||||
self,
|
||||
mock_reload_databases,
|
||||
mock_base_models_mgr,
|
||||
mock_read_meta,
|
||||
mock_read_file,
|
||||
mock_request_source_json,
|
||||
):
|
||||
from core.sources.popular import get_base_models_status, update_base_models_from_remote
|
||||
|
||||
local_data = {
|
||||
"base_models": [{"name": "Known Model", "aliases": ["Known Alias"]}]
|
||||
}
|
||||
mock_read_file.return_value = local_data
|
||||
mock_read_meta.return_value = {}
|
||||
mock_request_source_json.return_value = {
|
||||
"BaseModel": ["Known Alias", "New Model"]
|
||||
}
|
||||
|
||||
status = get_base_models_status(check_remote=True)
|
||||
self.assertTrue(status["update_available"])
|
||||
|
||||
result = update_base_models_from_remote()
|
||||
|
||||
self.assertEqual(result["new_models_added"], 1)
|
||||
self.assertEqual(result["new_models_added_list"], ["New Model"])
|
||||
saved_data = mock_base_models_mgr.sync_catalog.call_args[0][0]
|
||||
self.assertEqual(
|
||||
[model["name"] for model in saved_data["base_models"]],
|
||||
["Known Model", "New Model"],
|
||||
)
|
||||
mock_reload_databases.assert_called_once()
|
||||
|
||||
@patch("core.sources.popular.base_models_mgr")
|
||||
@patch("requests.get")
|
||||
def test_update_base_models_from_remote_success(self, mock_get, mock_base_models_mgr):
|
||||
|
||||
Reference in New Issue
Block a user