refactor(services): unify local model name matching with uniqueness and base-model guards (#1065)

Consolidate the duplicate name-matching logic into ModelScanner:
find_matching_models is now the single core, using each scanner's own
file_extensions for suffix stripping. get_model_info_by_name gains
require_unique/base_model kwargs while legacy route behavior is kept
byte-identical. reconnect_lora passes the recipe base model as a guard
and distinguishes ambiguous, base-model-mismatched, and missing LoRAs
in its error messages.
This commit is contained in:
Will Miao
2026-08-19 21:02:04 +08:00
parent 7fc3b7e5be
commit 8a16034135
5 changed files with 214 additions and 39 deletions
+91 -2
View File
@@ -2140,8 +2140,97 @@ class ModelScanner:
return sorted_models return sorted_models
return sorted_models[:limit] return sorted_models[:limit]
async def get_model_info_by_name(self, name): @staticmethod
"""Get model information by name""" def find_matching_models(
raw_data: List[Dict[str, Any]],
name: str,
*,
base_model: Optional[str] = None,
extensions: Optional[Set[str]] = None,
) -> List[Dict[str, Any]]:
"""Return all cached models matching ``name`` (case-insensitive).
A name containing a path separator must equal the model's
folder-relative path; a bare name matches on basename. When
``base_model`` is given, confident mismatches are rejected while
unknowns on either side stay eligible (lenient guard).
``extensions`` should be the scanner's own ``file_extensions`` so
suffix stripping only covers formats the scanner actually indexes.
"""
# Longest first so overlapping suffixes strip correctly.
exts = sorted(extensions or (".safetensors", ".ckpt", ".pt", ".bin"), key=len, reverse=True)
normalized_name = str(name).replace("\\", "/").casefold()
for ext in exts:
if normalized_name.endswith(ext):
normalized_name = normalized_name[: -len(ext)]
break
has_path = "/" in normalized_name
basename = normalized_name.rsplit("/", 1)[-1]
matches = []
for model in raw_data:
file_name = str(model.get("file_name") or "").replace("\\", "/")
folder = str(model.get("folder") or "").replace("\\", "/").strip("/")
model_path = f"{folder}/{file_name}" if folder else file_name
for ext in exts:
if model_path.casefold().endswith(ext):
model_path = model_path[: -len(ext)]
break
if (has_path and model_path.casefold() == normalized_name) or (
not has_path and model_path.rsplit("/", 1)[-1].casefold() == basename
):
matches.append(model)
expected_base = str(base_model or "").strip().casefold()
if expected_base and expected_base != "unknown":
matches = [
model
for model in matches
if str(model.get("base_model") or "").strip().casefold()
in ("", "unknown", expected_base)
]
return matches
async def find_models_by_name(
self, name: str, *, base_model: Optional[str] = None
) -> List[Dict[str, Any]]:
"""Return every cached model matching ``name`` (see ``find_matching_models``)."""
try:
cache = await self.get_cached_data()
return self.find_matching_models(
cache.raw_data,
name,
base_model=base_model,
extensions=self.file_extensions,
)
except Exception as e:
logger.error(f"Error finding models by name: {e}", exc_info=True)
return []
async def get_model_info_by_name(
self,
name: str,
*,
require_unique: bool = False,
base_model: Optional[str] = None,
):
"""Get model information by name.
Default mode keeps the legacy first-match/fallback semantics. With
``require_unique`` an ambiguous name is a miss, and ``base_model``
rejects confident base-model mismatches (unknowns stay eligible).
"""
if require_unique or base_model:
try:
matches = await self.find_models_by_name(name, base_model=base_model)
if require_unique and len(matches) != 1:
return None
return matches[0] if matches else None
except Exception as e:
logger.error(f"Error getting model info by name: {e}", exc_info=True)
return None
try: try:
cache = await self.get_cached_data() cache = await self.get_cached_data()
+10 -35
View File
@@ -2934,44 +2934,19 @@ class RecipeScanner:
if not self._lora_scanner or not name: if not self._lora_scanner or not name:
return None return None
normalized_name = str(name).replace("\\", "/").casefold() return await self._lora_scanner.get_model_info_by_name(
for extension in (".safetensors", ".ckpt", ".pt", ".bin"): name, require_unique=True, base_model=base_model
if normalized_name.endswith(extension): )
normalized_name = normalized_name[: -len(extension)]
break
has_path = "/" in normalized_name
basename = normalized_name.rsplit("/", 1)[-1]
cached_data = await self._lora_scanner.get_cached_data() async def find_local_loras_by_name(
matches = [] self, name: str, base_model: Optional[str] = None
for model in cached_data.raw_data: ) -> List[Dict[str, Any]]:
file_name = str(model.get("file_name") or "").replace("\\", "/") """Return every local LoRA matching ``name`` (used to explain lookup misses)."""
folder = str(model.get("folder") or "").replace("\\", "/").strip("/")
model_path = f"{folder}/{file_name}" if folder else file_name
for extension in (".safetensors", ".ckpt", ".pt", ".bin"):
if model_path.casefold().endswith(extension):
model_path = model_path[: -len(extension)]
break
if (has_path and model_path.casefold() == normalized_name) or (
not has_path and model_path.rsplit("/", 1)[-1].casefold() == basename
):
matches.append(model)
if len(matches) != 1: if not self._lora_scanner or not name:
return None return []
match = matches[0] return await self._lora_scanner.find_models_by_name(name, base_model=base_model)
expected_base = str(base_model or "").strip().casefold()
actual_base = str(match.get("base_model") or "").strip().casefold()
if (
expected_base
and expected_base != "unknown"
and actual_base
and actual_base != "unknown"
and expected_base != actual_base
):
return None
return match
async def get_local_lora_by_hash(self, hash_value: str) -> Optional[Dict[str, Any]]: async def get_local_lora_by_hash(self, hash_value: str) -> Optional[Dict[str, Any]]:
"""Lookup a local LoRA through the scanner's hash index.""" """Lookup a local LoRA through the scanner's hash index."""
+14 -1
View File
@@ -426,8 +426,21 @@ class RecipePersistenceService:
if not recipe_path or not os.path.exists(recipe_path): if not recipe_path or not os.path.exists(recipe_path):
raise RecipeNotFoundError("Recipe not found") raise RecipeNotFoundError("Recipe not found")
target_lora = await recipe_scanner.get_local_lora(target_name) with open(recipe_path, "r", encoding="utf-8") as file_obj:
recipe_base_model = json.load(file_obj).get("base_model", "")
target_lora = await recipe_scanner.get_local_lora(target_name, recipe_base_model)
if not target_lora: if not target_lora:
matches = await recipe_scanner.find_local_loras_by_name(target_name)
if len(matches) > 1:
raise RecipeValidationError(
f"Multiple local LoRAs match '{target_name}'; "
"include the folder path to disambiguate"
)
if len(matches) == 1:
raise RecipeValidationError(
f"Local LoRA '{target_name}' has a different base model than the recipe"
)
raise RecipeNotFoundError(f"Local LoRA not found with name: {target_name}") raise RecipeNotFoundError(f"Local LoRA not found with name: {target_name}")
recipe_data, updated_lora = await recipe_scanner.update_lora_entry( recipe_data, updated_lora = await recipe_scanner.update_lora_entry(
+29 -1
View File
@@ -57,9 +57,33 @@ class StubLoraScanner:
meta = self._hash_meta.get(hash_value.lower()) meta = self._hash_meta.get(hash_value.lower())
return meta.get("path") if meta else None return meta.get("path") if meta else None
async def get_model_info_by_name(self, name: str): async def get_model_info_by_name(
self,
name: str,
*,
require_unique: bool = False,
base_model: str | None = None,
):
if require_unique or base_model:
matches = ModelScanner.find_matching_models(
self._cache.raw_data,
name,
base_model=base_model,
extensions={".safetensors"},
)
if require_unique and len(matches) != 1:
return None
return matches[0] if matches else None
return self._models_by_name.get(name) return self._models_by_name.get(name)
async def find_models_by_name(self, name: str, *, base_model: str | None = None):
return ModelScanner.find_matching_models(
self._cache.raw_data,
name,
base_model=base_model,
extensions={".safetensors"},
)
def register_model(self, name: str, info: Dict[str, Any]) -> None: def register_model(self, name: str, info: Dict[str, Any]) -> None:
self._models_by_name[name] = info self._models_by_name[name] = info
hash_value = (info.get("sha256") or "").lower() hash_value = (info.get("sha256") or "").lower()
@@ -130,7 +154,11 @@ async def test_local_lora_lookup_requires_unambiguous_name_and_matching_base_mod
stub._hash_meta["b" * 64] = {"path": models[1]["file_path"]} stub._hash_meta["b" * 64] = {"path": models[1]["file_path"]}
assert await scanner.get_local_lora("style") is None assert await scanner.get_local_lora("style") is None
assert await scanner.get_local_lora("style", "SDXL 1.0") is models[1]
assert await scanner.get_local_lora("sdxl/style.safetensors", "SDXL 1.0") is models[1] assert await scanner.get_local_lora("sdxl/style.safetensors", "SDXL 1.0") is models[1]
# The lora scanner only indexes .safetensors, so a .pt name must not be
# stripped into a cross-extension match.
assert await scanner.get_local_lora("sdxl/style.pt", "SDXL 1.0") is None
assert await scanner.get_local_lora("sdxl/style.safetensors", "SD 1.5") is None assert await scanner.get_local_lora("sdxl/style.safetensors", "SD 1.5") is None
assert await scanner.get_local_lora("other/style.safetensors") is None assert await scanner.get_local_lora("other/style.safetensors") is None
assert await scanner.get_local_lora_by_hash("b" * 64) is models[1] assert await scanner.get_local_lora_by_hash("b" * 64) is models[1]
+70
View File
@@ -17,6 +17,7 @@ from py.services.recipes.errors import (
RecipeValidationError, RecipeValidationError,
) )
from py.services.recipes.persistence_service import RecipePersistenceService from py.services.recipes.persistence_service import RecipePersistenceService
from py.services.model_scanner import ModelScanner
from py.recipes.parsers.civitai_image import CivitaiApiMetadataParser from py.recipes.parsers.civitai_image import CivitaiApiMetadataParser
from py.utils.exif_utils import ExifUtils from py.utils.exif_utils import ExifUtils
@@ -1156,3 +1157,72 @@ async def test_analyze_local_image_fingerprint_uses_sha256_normalized_hash(tmp_p
assert result.payload["loras"][0]["hash"] == sha256 assert result.payload["loras"][0]["hash"] == sha256
assert result.payload["fingerprint"] == f"{sha256}:1.0" assert result.payload["fingerprint"] == f"{sha256}:1.0"
@pytest.mark.asyncio
async def test_reconnect_lora_distinguishes_ambiguous_mismatched_and_missing(tmp_path):
service = RecipePersistenceService(
exif_utils=DummyExifUtils(),
card_preview_width=512,
logger=logging.getLogger("test"),
)
models = [
{
"file_name": "style.safetensors",
"folder": "sd15",
"file_path": "/models/loras/sd15/style.safetensors",
"base_model": "SD 1.5",
},
{
"file_name": "style.safetensors",
"folder": "sdxl",
"file_path": "/models/loras/sdxl/style.safetensors",
"base_model": "SDXL 1.0",
},
]
class DummyScanner:
def __init__(self, recipe_path):
self._recipe_path = recipe_path
async def get_recipe_json_path(self, recipe_id):
return str(self._recipe_path)
async def get_local_lora(self, name, base_model=None):
matches = ModelScanner.find_matching_models(models, name, base_model=base_model)
return matches[0] if len(matches) == 1 else None
async def find_local_loras_by_name(self, name, base_model=None):
return ModelScanner.find_matching_models(models, name, base_model=base_model)
def write_recipe(base_model):
recipe_path = tmp_path / "recipe.json"
recipe_path.write_text(
json.dumps({"id": "r1", "base_model": base_model, "loras": []})
)
return DummyScanner(recipe_path)
# Ambiguous bare name: two candidates survive (recipe base model unknown)
scanner = write_recipe("")
with pytest.raises(RecipeValidationError, match="include the folder path"):
await service.reconnect_lora(
recipe_scanner=scanner, recipe_id="r1", lora_index=0, target_name="style"
)
# Confident base-model mismatch: the only candidate belongs to another family
scanner = write_recipe("SD 1.5")
with pytest.raises(RecipeValidationError, match="different base model"):
await service.reconnect_lora(
recipe_scanner=scanner,
recipe_id="r1",
lora_index=0,
target_name="sdxl/style",
)
# No candidate at all
scanner = write_recipe("SDXL 1.0")
with pytest.raises(RecipeNotFoundError, match="not found"):
await service.reconnect_lora(
recipe_scanner=scanner, recipe_id="r1", lora_index=0, target_name="missing"
)