diff --git a/py/services/model_scanner.py b/py/services/model_scanner.py index a8329b6b..92fd186a 100644 --- a/py/services/model_scanner.py +++ b/py/services/model_scanner.py @@ -2140,8 +2140,97 @@ class ModelScanner: return sorted_models return sorted_models[:limit] - async def get_model_info_by_name(self, name): - """Get model information by name""" + @staticmethod + 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: cache = await self.get_cached_data() diff --git a/py/services/recipe_scanner.py b/py/services/recipe_scanner.py index 94d26b08..476b8586 100644 --- a/py/services/recipe_scanner.py +++ b/py/services/recipe_scanner.py @@ -2934,44 +2934,19 @@ class RecipeScanner: if not self._lora_scanner or not name: return None - normalized_name = str(name).replace("\\", "/").casefold() - for extension in (".safetensors", ".ckpt", ".pt", ".bin"): - if normalized_name.endswith(extension): - normalized_name = normalized_name[: -len(extension)] - break - has_path = "/" in normalized_name - basename = normalized_name.rsplit("/", 1)[-1] + return await self._lora_scanner.get_model_info_by_name( + name, require_unique=True, base_model=base_model + ) - cached_data = await self._lora_scanner.get_cached_data() - matches = [] - for model in cached_data.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 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) + async def find_local_loras_by_name( + self, name: str, base_model: Optional[str] = None + ) -> List[Dict[str, Any]]: + """Return every local LoRA matching ``name`` (used to explain lookup misses).""" - if len(matches) != 1: - return None + if not self._lora_scanner or not name: + return [] - match = matches[0] - 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 + return await self._lora_scanner.find_models_by_name(name, base_model=base_model) 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.""" diff --git a/py/services/recipes/persistence_service.py b/py/services/recipes/persistence_service.py index 75d81e49..bfb62429 100644 --- a/py/services/recipes/persistence_service.py +++ b/py/services/recipes/persistence_service.py @@ -426,8 +426,21 @@ class RecipePersistenceService: if not recipe_path or not os.path.exists(recipe_path): 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: + 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}") recipe_data, updated_lora = await recipe_scanner.update_lora_entry( diff --git a/tests/services/test_recipe_scanner.py b/tests/services/test_recipe_scanner.py index 015920be..3a066cf1 100644 --- a/tests/services/test_recipe_scanner.py +++ b/tests/services/test_recipe_scanner.py @@ -57,9 +57,33 @@ class StubLoraScanner: meta = self._hash_meta.get(hash_value.lower()) 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) + 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: self._models_by_name[name] = info 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"]} 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] + # 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("other/style.safetensors") is None assert await scanner.get_local_lora_by_hash("b" * 64) is models[1] diff --git a/tests/services/test_recipe_services.py b/tests/services/test_recipe_services.py index c05a0bd5..f63c56e4 100644 --- a/tests/services/test_recipe_services.py +++ b/tests/services/test_recipe_services.py @@ -17,6 +17,7 @@ from py.services.recipes.errors import ( RecipeValidationError, ) 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.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["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" + )