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.
This commit is contained in:
Will Miao
2025-11-21 11:27:09 +08:00
parent 7173a2b9d6
commit e24621a0af
2 changed files with 76 additions and 6 deletions

View File

@@ -564,6 +564,27 @@ class RecipeScanner:
logger.error(f"Error getting hash from Civitai: {e}") logger.error(f"Error getting hash from Civitai: {e}")
return None, False 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]: async def _determine_base_model(self, loras: List[Dict]) -> Optional[str]:
"""Determine the most common base model among LoRAs""" """Determine the most common base model among LoRAs"""
base_models = {} base_models = {}
@@ -609,13 +630,31 @@ class RecipeScanner:
return lora return lora
hash_value = (lora.get('hash') or '').lower() hash_value = (lora.get('hash') or '').lower()
if not hash_value: version_entry = None
return lora if not hash_value and lora.get('modelVersionId') is not None:
version_entry = self._get_lora_from_version_index(lora.get('modelVersionId'))
try: try:
lora['inLibrary'] = self._lora_scanner.has_hash(hash_value) if hash_value:
lora['preview_url'] = self._lora_scanner.get_preview_url_by_hash(hash_value) lora['inLibrary'] = self._lora_scanner.has_hash(hash_value)
lora['localPath'] = self._lora_scanner.get_path_by_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 except Exception as exc: # pragma: no cover - defensive logging
logger.debug("Error enriching lora entry %s: %s", hash_value, exc) logger.debug("Error enriching lora entry %s: %s", hash_value, exc)

View File

@@ -24,7 +24,7 @@ class StubLoraScanner:
self._hash_index = StubHashIndex() self._hash_index = StubHashIndex()
self._hash_meta: dict[str, dict[str, str]] = {} self._hash_meta: dict[str, dict[str, str]] = {}
self._models_by_name: dict[str, dict] = {} 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): async def get_cached_data(self):
return self._cache return self._cache
@@ -46,12 +46,20 @@ class StubLoraScanner:
def register_model(self, name: str, info: dict) -> None: def register_model(self, name: str, info: dict) -> 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()
version_id = info.get("civitai", {}).get("id")
if hash_value: if hash_value:
self._hash_meta[hash_value] = { self._hash_meta[hash_value] = {
"path": info.get("file_path", ""), "path": info.get("file_path", ""),
"preview_url": info.get("preview_url", ""), "preview_url": info.get("preview_url", ""),
} }
self._hash_index._hash_to_path[hash_value] = info.get("file_path", "") 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({ self._cache.raw_data.append({
"sha256": info.get("sha256", ""), "sha256": info.get("sha256", ""),
"path": info.get("file_path", ""), "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()) persisted = json.loads(recipe_path.read_text())
assert persisted["file_path"] == expected_path 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"]