diff --git a/py/services/recipe_scanner.py b/py/services/recipe_scanner.py index b5f1ce2f..f1fdfd2b 100644 --- a/py/services/recipe_scanner.py +++ b/py/services/recipe_scanner.py @@ -564,6 +564,27 @@ class RecipeScanner: logger.error(f"Error getting hash from Civitai: {e}") return None, False + def _get_lora_from_version_index(self, model_version_id: Any) -> Optional[Dict[str, Any]]: + """Quickly fetch a cached LoRA entry by modelVersionId using the version index.""" + + if not self._lora_scanner: + return None + + cache = getattr(self._lora_scanner, "_cache", None) + if cache is None: + return None + + version_index = getattr(cache, "version_index", None) + if not version_index: + return None + + try: + normalized_id = int(model_version_id) + except (TypeError, ValueError): + return None + + return version_index.get(normalized_id) + async def _determine_base_model(self, loras: List[Dict]) -> Optional[str]: """Determine the most common base model among LoRAs""" base_models = {} @@ -609,13 +630,31 @@ class RecipeScanner: return lora hash_value = (lora.get('hash') or '').lower() - if not hash_value: - return lora + version_entry = None + if not hash_value and lora.get('modelVersionId') is not None: + version_entry = self._get_lora_from_version_index(lora.get('modelVersionId')) try: - lora['inLibrary'] = self._lora_scanner.has_hash(hash_value) - lora['preview_url'] = self._lora_scanner.get_preview_url_by_hash(hash_value) - lora['localPath'] = self._lora_scanner.get_path_by_hash(hash_value) + if hash_value: + lora['inLibrary'] = self._lora_scanner.has_hash(hash_value) + lora['preview_url'] = self._lora_scanner.get_preview_url_by_hash(hash_value) + lora['localPath'] = self._lora_scanner.get_path_by_hash(hash_value) + elif version_entry: + lora['inLibrary'] = True + cached_path = version_entry.get('file_path') or version_entry.get('path') + if cached_path: + lora.setdefault('localPath', cached_path) + if not lora.get('file_name'): + lora['file_name'] = os.path.splitext(os.path.basename(cached_path))[0] + + if version_entry.get('sha256') and not lora.get('hash'): + lora['hash'] = version_entry.get('sha256') + + preview_url = version_entry.get('preview_url') + if preview_url: + lora.setdefault('preview_url', preview_url) + else: + lora.setdefault('inLibrary', False) except Exception as exc: # pragma: no cover - defensive logging logger.debug("Error enriching lora entry %s: %s", hash_value, exc) diff --git a/tests/services/test_recipe_scanner.py b/tests/services/test_recipe_scanner.py index 20f25727..d5da17c7 100644 --- a/tests/services/test_recipe_scanner.py +++ b/tests/services/test_recipe_scanner.py @@ -24,7 +24,7 @@ class StubLoraScanner: self._hash_index = StubHashIndex() self._hash_meta: dict[str, dict[str, str]] = {} self._models_by_name: dict[str, dict] = {} - self._cache = SimpleNamespace(raw_data=[]) + self._cache = SimpleNamespace(raw_data=[], version_index={}) async def get_cached_data(self): return self._cache @@ -46,12 +46,20 @@ class StubLoraScanner: def register_model(self, name: str, info: dict) -> None: self._models_by_name[name] = info hash_value = (info.get("sha256") or "").lower() + version_id = info.get("civitai", {}).get("id") if hash_value: self._hash_meta[hash_value] = { "path": info.get("file_path", ""), "preview_url": info.get("preview_url", ""), } self._hash_index._hash_to_path[hash_value] = info.get("file_path", "") + if version_id is not None: + self._cache.version_index[int(version_id)] = { + "file_path": info.get("file_path", ""), + "sha256": hash_value, + "preview_url": info.get("preview_url", ""), + "civitai": info.get("civitai", {}), + } self._cache.raw_data.append({ "sha256": info.get("sha256", ""), "path": info.get("file_path", ""), @@ -216,3 +224,26 @@ async def test_load_recipe_rewrites_missing_image_path(tmp_path: Path, recipe_sc persisted = json.loads(recipe_path.read_text()) assert persisted["file_path"] == expected_path + + +def test_enrich_uses_version_index_when_hash_missing(recipe_scanner): + scanner, stub = recipe_scanner + version_id = 77 + file_path = str(Path(config.loras_roots[0]) / "loras" / "version-entry.safetensors") + registered = { + "sha256": "deadbeef", + "file_path": file_path, + "preview_url": "preview-from-cache.png", + "civitai": {"id": version_id}, + } + stub.register_model("version-entry", registered) + + lora = {"hash": "", "file_name": "", "modelVersionId": version_id, "strength": 0.5} + + enriched = scanner._enrich_lora_entry(dict(lora)) + + assert enriched["inLibrary"] is True + assert enriched["hash"] == registered["sha256"] + assert enriched["localPath"] == file_path + assert enriched["file_name"] == Path(file_path).stem + assert enriched["preview_url"] == registered["preview_url"]