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