feat(recipe-parser): enhance LoRA metadata with local file matching

Add comprehensive local file matching for LoRA entries in recipe metadata:
- Add modelVersionId-based lookup via new _get_lora_from_version_index method
- Extend LoRA entry with additional fields: existsLocally, inLibrary, localPath, thumbnailUrl, size
- Improve local file detection by checking both SHA256 hash and modelVersionId
- Set default thumbnail URL and size values for missing LoRA files
- Add proper typing with Optional imports for better code clarity

This provides more accurate local file status and metadata for LoRA entries in recipes.
This commit is contained in:
Will Miao
2025-11-29 08:29:05 +08:00
parent 53c4165d82
commit 073fb3a94a
2 changed files with 158 additions and 35 deletions

View File

@@ -3,7 +3,7 @@
import re import re
import json import json
import logging import logging
from typing import Dict, Any from typing import Dict, Any, Optional
from ...config import config from ...config import config
from ..base import RecipeMetadataParser from ..base import RecipeMetadataParser
from ..constants import GEN_PARAM_KEYS from ..constants import GEN_PARAM_KEYS
@@ -16,6 +16,28 @@ class RecipeFormatParser(RecipeMetadataParser):
# Regular expression pattern for extracting recipe metadata # Regular expression pattern for extracting recipe metadata
METADATA_MARKER = r'Recipe metadata: (\{.*\})' METADATA_MARKER = r'Recipe metadata: (\{.*\})'
async def _get_lora_from_version_index(self, recipe_scanner, model_version_id: Any) -> Optional[Dict[str, Any]]:
"""Return a cached LoRA entry by modelVersionId if available."""
if not recipe_scanner or not getattr(recipe_scanner, "_lora_scanner", None):
return None
try:
normalized_id = int(model_version_id)
except (TypeError, ValueError):
return None
try:
cache = await recipe_scanner._lora_scanner.get_cached_data()
except Exception as exc: # pragma: no cover - defensive logging
logger.debug("Unable to load lora cache for version lookup: %s", exc)
return None
if not cache or not getattr(cache, "version_index", None):
return None
return cache.version_index.get(normalized_id)
def is_metadata_matching(self, user_comment: str) -> bool: def is_metadata_matching(self, user_comment: str) -> bool:
"""Check if the user comment matches the metadata format""" """Check if the user comment matches the metadata format"""
@@ -53,45 +75,69 @@ class RecipeFormatParser(RecipeMetadataParser):
'type': 'lora', 'type': 'lora',
'weight': lora.get('strength', 1.0), 'weight': lora.get('strength', 1.0),
'file_name': lora.get('file_name', ''), 'file_name': lora.get('file_name', ''),
'hash': lora.get('hash', '') 'hash': lora.get('hash', ''),
'existsLocally': False,
'inLibrary': False,
'localPath': None,
'thumbnailUrl': '/loras_static/images/no-preview.png',
'size': 0
} }
# Check if this LoRA exists locally by SHA256 hash # Check if this LoRA exists locally by SHA256 hash
if lora.get('hash') and recipe_scanner: if recipe_scanner:
lora_scanner = recipe_scanner._lora_scanner lora_scanner = recipe_scanner._lora_scanner
exists_locally = lora_scanner.has_hash(lora['hash'])
if exists_locally: if lora.get('hash'):
lora_cache = await lora_scanner.get_cached_data() exists_locally = lora_scanner.has_hash(lora['hash'])
lora_item = next((item for item in lora_cache.raw_data if item['sha256'].lower() == lora['hash'].lower()), None) if exists_locally:
if lora_item: lora_cache = await lora_scanner.get_cached_data()
lora_item = next((item for item in lora_cache.raw_data if item['sha256'].lower() == lora['hash'].lower()), None)
if lora_item:
lora_entry['existsLocally'] = True
lora_entry['inLibrary'] = True
lora_entry['localPath'] = lora_item['file_path']
lora_entry['file_name'] = lora_item['file_name']
lora_entry['size'] = lora_item['size']
lora_entry['thumbnailUrl'] = config.get_preview_static_url(lora_item['preview_url'])
else:
lora_entry['existsLocally'] = False
lora_entry['inLibrary'] = False
lora_entry['localPath'] = None
# If we still don't have a local match, try matching by modelVersionId
if not lora_entry['existsLocally'] and lora.get('modelVersionId') is not None:
cached_lora = await self._get_lora_from_version_index(recipe_scanner, lora.get('modelVersionId'))
if cached_lora:
lora_entry['existsLocally'] = True lora_entry['existsLocally'] = True
lora_entry['localPath'] = lora_item['file_path'] lora_entry['inLibrary'] = True
lora_entry['file_name'] = lora_item['file_name'] lora_entry['localPath'] = cached_lora.get('file_path')
lora_entry['size'] = lora_item['size'] lora_entry['file_name'] = cached_lora.get('file_name') or lora_entry['file_name']
lora_entry['thumbnailUrl'] = config.get_preview_static_url(lora_item['preview_url']) lora_entry['size'] = cached_lora.get('size', lora_entry['size'])
if cached_lora.get('sha256'):
else: lora_entry['hash'] = cached_lora['sha256']
lora_entry['existsLocally'] = False preview_url = cached_lora.get('preview_url')
lora_entry['localPath'] = None if preview_url:
lora_entry['thumbnailUrl'] = config.get_preview_static_url(preview_url)
# Try to get additional info from Civitai if we have a model version ID
if lora.get('modelVersionId') and metadata_provider: # Try to get additional info from Civitai if we have a model version ID and still missing locally
try: if not lora_entry['existsLocally'] and lora.get('modelVersionId') and metadata_provider:
civitai_info_tuple = await metadata_provider.get_model_version_info(lora['modelVersionId']) try:
# Populate lora entry with Civitai info civitai_info_tuple = await metadata_provider.get_model_version_info(lora['modelVersionId'])
populated_entry = await self.populate_lora_from_civitai( # Populate lora entry with Civitai info
lora_entry, populated_entry = await self.populate_lora_from_civitai(
civitai_info_tuple, lora_entry,
recipe_scanner, civitai_info_tuple,
None, # No need to track base model counts recipe_scanner,
lora['hash'] None, # No need to track base model counts
) lora_entry.get('hash', '')
if populated_entry is None: )
continue # Skip invalid LoRA types if populated_entry is None:
lora_entry = populated_entry continue # Skip invalid LoRA types
except Exception as e: lora_entry = populated_entry
logger.error(f"Error fetching Civitai info for LoRA: {e}") except Exception as e:
lora_entry['thumbnailUrl'] = '/loras_static/images/no-preview.png' logger.error(f"Error fetching Civitai info for LoRA: {e}")
lora_entry['thumbnailUrl'] = '/loras_static/images/no-preview.png'
loras.append(lora_entry) loras.append(lora_entry)

View File

@@ -2,6 +2,7 @@ import json
import pytest import pytest
from py.recipes.parsers.recipe_format import RecipeFormatParser from py.recipes.parsers.recipe_format import RecipeFormatParser
from py.config import config
@pytest.mark.asyncio @pytest.mark.asyncio
@@ -65,3 +66,79 @@ async def test_recipe_format_parser_populates_checkpoint(monkeypatch):
assert checkpoint["file_name"] == "Z_Image_Turbo" assert checkpoint["file_name"] == "Z_Image_Turbo"
assert result["base_model"] == "sdxl" assert result["base_model"] == "sdxl"
assert result["model"] == checkpoint assert result["model"] == checkpoint
@pytest.mark.asyncio
async def test_recipe_format_parser_marks_lora_in_library_by_version(monkeypatch):
async def fake_metadata_provider():
class Provider:
async def get_model_version_info(self, version_id):
assert version_id == 1244133
return None, None
return Provider()
monkeypatch.setattr(
"py.recipes.parsers.recipe_format.get_default_metadata_provider",
fake_metadata_provider,
)
cached_entry = {
"file_path": "/loras/moriimee.safetensors",
"file_name": "MoriiMee Gothic Niji | LoRA Style",
"size": 4096,
"sha256": "abc123",
"preview_url": "/previews/moriimee.png",
}
class FakeCache:
def __init__(self, entry):
self.raw_data = [entry]
self.version_index = {1244133: entry}
class FakeLoraScanner:
def __init__(self, entry):
self._cache = FakeCache(entry)
def has_hash(self, sha256):
return False
async def get_cached_data(self):
return self._cache
class FakeRecipeScanner:
def __init__(self, entry):
self._lora_scanner = FakeLoraScanner(entry)
parser = RecipeFormatParser()
recipe_metadata = {
"title": "Semi-realism",
"base_model": "Illustrious",
"loras": [
{
"modelVersionId": 1244133,
"modelName": "MoriiMee Gothic Niji | LoRA Style",
"modelVersionName": "V1 Ilustrious",
"strength": 0.5,
"hash": "",
}
],
"gen_params": {"steps": 29},
"tags": ["woman"],
}
metadata_text = f"Recipe metadata: {json.dumps(recipe_metadata)}"
result = await parser.parse_metadata(
metadata_text, recipe_scanner=FakeRecipeScanner(cached_entry)
)
lora_entry = result["loras"][0]
assert lora_entry["existsLocally"] is True
assert lora_entry["inLibrary"] is True
assert lora_entry["localPath"] == cached_entry["file_path"]
assert lora_entry["file_name"] == cached_entry["file_name"]
assert lora_entry["hash"] == cached_entry["sha256"]
assert lora_entry["size"] == cached_entry["size"]
assert lora_entry["thumbnailUrl"] == config.get_preview_static_url(
cached_entry["preview_url"]
)