refactor model list result construction
This commit is contained in:
+16
-28
@@ -177,6 +177,18 @@ def update_model_list_from_remote() -> Dict[str, Any]:
|
||||
|
||||
|
||||
|
||||
def _build_model_list_result(model: Dict[str, Any]) -> Dict[str, Any]:
|
||||
return {
|
||||
"source": "model_list",
|
||||
"filename": model.get("filename", ""),
|
||||
"url": model.get("url", ""),
|
||||
"name": model.get("name", ""),
|
||||
"type": model.get("type", ""),
|
||||
"directory": model.get("save_path", "checkpoints"),
|
||||
"size": model.get("size", ""),
|
||||
}
|
||||
|
||||
|
||||
def search_model_list(
|
||||
filename: str, exact_only: bool = False
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
@@ -207,13 +219,7 @@ def search_model_list(
|
||||
url = model.get("url", "")
|
||||
if url:
|
||||
return {
|
||||
"source": "model_list",
|
||||
"filename": model_filename,
|
||||
"url": url,
|
||||
"name": model.get("name", ""),
|
||||
"type": model.get("type", ""),
|
||||
"directory": model.get("save_path", "checkpoints"),
|
||||
"size": model.get("size", ""),
|
||||
**_build_model_list_result(model),
|
||||
"match_type": "exact",
|
||||
}
|
||||
|
||||
@@ -237,13 +243,7 @@ def search_model_list(
|
||||
if url:
|
||||
score = calculate_similarity(filename_norm, normalize_filename(model_filename))
|
||||
return {
|
||||
"source": "model_list",
|
||||
"filename": model_filename,
|
||||
"url": url,
|
||||
"name": model.get("name", ""),
|
||||
"type": model.get("type", ""),
|
||||
"directory": model.get("save_path", "checkpoints"),
|
||||
"size": model.get("size", ""),
|
||||
**_build_model_list_result(model),
|
||||
"match_type": "fuzzy",
|
||||
"confidence": round(score * 100, 1),
|
||||
}
|
||||
@@ -265,13 +265,7 @@ def search_model_list(
|
||||
if url:
|
||||
best_score = score
|
||||
best_match = {
|
||||
"source": "model_list",
|
||||
"filename": model_filename,
|
||||
"url": url,
|
||||
"name": model.get("name", ""),
|
||||
"type": model.get("type", ""),
|
||||
"directory": model.get("save_path", "checkpoints"),
|
||||
"size": model.get("size", ""),
|
||||
**_build_model_list_result(model),
|
||||
"match_type": "similar",
|
||||
"confidence": round(score * 100, 1),
|
||||
}
|
||||
@@ -310,13 +304,7 @@ def search_model_list_multiple(filename: str, limit: int = 5) -> List[Dict[str,
|
||||
if url:
|
||||
results.append(
|
||||
{
|
||||
"source": "model_list",
|
||||
"filename": model_filename,
|
||||
"url": url,
|
||||
"name": model.get("name", ""),
|
||||
"type": model.get("type", ""),
|
||||
"directory": model.get("save_path", "checkpoints"),
|
||||
"size": model.get("size", ""),
|
||||
**_build_model_list_result(model),
|
||||
"confidence": round(score * 100, 1),
|
||||
}
|
||||
)
|
||||
|
||||
+50
-1
@@ -1,4 +1,6 @@
|
||||
import unittest
|
||||
from unittest.mock import patch
|
||||
|
||||
from core.matcher import (
|
||||
build_filename_search_queries,
|
||||
calculate_archived_model_confidence,
|
||||
@@ -8,8 +10,8 @@ from core.matcher import (
|
||||
normalize_filename,
|
||||
normalize_model_family_filename,
|
||||
)
|
||||
from core.sources.model_list import search_model_list, search_model_list_multiple
|
||||
from core.progress import report_progress
|
||||
from core.sources.model_list import search_model_list, search_model_list_multiple
|
||||
from core.type_utils import (
|
||||
MODEL_CONTAINER_SUFFIXES,
|
||||
MODEL_VARIANT_SUFFIXES,
|
||||
@@ -21,6 +23,8 @@ from core.type_utils import (
|
||||
as_list,
|
||||
format_size_bytes,
|
||||
)
|
||||
|
||||
|
||||
class MatcherTests(unittest.TestCase):
|
||||
def test_normalize_filename_removes_extension_and_converts_lowercase(self):
|
||||
self.assertEqual("test model", normalize_filename("Test_Model.safetensors"))
|
||||
@@ -171,6 +175,51 @@ class MatcherTests(unittest.TestCase):
|
||||
except Exception as e:
|
||||
self.fail(f"search_model_list failed with exception: {e}")
|
||||
|
||||
def test_model_list_result_fields_preserve_match_mode_contracts(self):
|
||||
models = [
|
||||
{
|
||||
"filename": "exact.safetensors",
|
||||
"url": "https://example.test/exact",
|
||||
"name": "Exact model",
|
||||
"type": "checkpoints",
|
||||
"save_path": "checkpoints",
|
||||
"size": "12 MB",
|
||||
},
|
||||
{
|
||||
"filename": "target-model.safetensors",
|
||||
"url": "https://example.test/target",
|
||||
"name": "Target model",
|
||||
"type": "loras",
|
||||
"save_path": "loras",
|
||||
"size": 24,
|
||||
},
|
||||
]
|
||||
|
||||
with patch("core.sources.model_list._load_model_list", return_value=models):
|
||||
exact = search_model_list("exact.safetensors", exact_only=True)
|
||||
fuzzy = search_model_list("target-model-extra.safetensors")
|
||||
multiple = search_model_list_multiple("exact", limit=1)
|
||||
|
||||
self.assertEqual(
|
||||
exact,
|
||||
{
|
||||
"source": "model_list",
|
||||
"filename": "exact.safetensors",
|
||||
"url": "https://example.test/exact",
|
||||
"name": "Exact model",
|
||||
"type": "checkpoints",
|
||||
"directory": "checkpoints",
|
||||
"size": "12 MB",
|
||||
"match_type": "exact",
|
||||
},
|
||||
)
|
||||
self.assertEqual(fuzzy["filename"], "target-model.safetensors")
|
||||
self.assertEqual(fuzzy["match_type"], "fuzzy")
|
||||
self.assertIn("confidence", fuzzy)
|
||||
self.assertEqual(multiple[0]["filename"], "exact.safetensors")
|
||||
self.assertIn("confidence", multiple[0])
|
||||
self.assertNotIn("match_type", multiple[0])
|
||||
|
||||
def test_build_filename_search_queries(self):
|
||||
# Verify that common precisions are extracted correctly
|
||||
queries = build_filename_search_queries("sdxl_ep10_fp16.safetensors")
|
||||
|
||||
Reference in New Issue
Block a user