From e24621a0af4008aad78d159ce5fb6fbe19828fac Mon Sep 17 00:00:00 2001 From: Will Miao Date: Fri, 21 Nov 2025 11:27:09 +0800 Subject: [PATCH] feat(recipe-scanner): add version index fallback for LoRA enrichment Add _get_lora_from_version_index method to fetch cached LoRA entries by modelVersionId when hash is unavailable. This improves LoRA enrichment by using version index as fallback when hash is missing, ensuring proper library status, file paths, and preview URLs are set even without hash values. Update test suite to include version_index in stub cache and add test coverage for version-based lookup functionality. --- py/services/recipe_scanner.py | 49 ++++++++++++++++++++++++--- tests/services/test_recipe_scanner.py | 33 +++++++++++++++++- 2 files changed, 76 insertions(+), 6 deletions(-) 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"]