mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-03-21 21:22:11 -03:00
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:
@@ -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)
|
||||
|
||||
|
||||
@@ -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"]
|
||||
|
||||
Reference in New Issue
Block a user