alphard/ComfyUI-Lora-Manager
1
0
Fork 0
mirror of https://github.com/willmiao/ComfyUI-Lora-Manager.git synced 2026-03-24 14:42:11 -03:00
Code Issues Packages Projects Releases Wiki Activity

feat(metadata): implement model data hydration and enhance metadata handling across services, fixes #547

Browse Source
This commit is contained in:
Will Miao
2025-10-09 22:15:07 +08:00
parent f542ade628
commit c9a65c7347
9 changed files with 443 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files

38
py/utils/example_images_metadata.py
View File

@@ -1,7 +1,7 @@
import logging
import os
import re
from typing import TYPE_CHECKING, Optional
from typing import TYPE_CHECKING, Any, Dict, Optional
from ..recipes.constants import GEN_PARAM_KEYS
from ..services.metadata_service import get_default_metadata_provider, get_metadata_provider
@@ -105,6 +105,7 @@ class MetadataUpdater:
async def update_cache_func(old_path, new_path, metadata):
return await scanner.update_single_model_cache(old_path, new_path, metadata)
await MetadataManager.hydrate_model_data(model_data)
success, error = await _get_metadata_sync_service().fetch_and_update_model(
sha256=model_hash,
file_path=file_path,
@@ -185,16 +186,16 @@ class MetadataUpdater:
if is_supported:
local_images_paths.append(file_path)
await MetadataManager.hydrate_model_data(model)
civitai_data = model.setdefault('civitai', {})
# Check if metadata update is needed (no civitai field or empty images)
needs_update = not model.get('civitai') or not model.get('civitai', {}).get('images')
needs_update = not civitai_data or not civitai_data.get('images')
if needs_update and local_images_paths:
logger.debug(f"Found {len(local_images_paths)} local example images for {model.get('model_name')}, updating metadata")
# Create or get civitai field
if not model.get('civitai'):
model['civitai'] = {}
# Create images array
images = []
@@ -229,16 +230,13 @@ class MetadataUpdater:
images.append(image_entry)
# Update the model's civitai.images field
model['civitai']['images'] = images
civitai_data['images'] = images
# Save metadata to .metadata.json file
file_path = model.get('file_path')
try:
# Create a copy of model data without 'folder' field
model_copy = model.copy()
model_copy.pop('folder', None)
# Write metadata to file
await MetadataManager.save_metadata(file_path, model_copy)
logger.info(f"Saved metadata for {model.get('model_name')}")
except Exception as e:
@@ -271,16 +269,13 @@ class MetadataUpdater:
tuple: (regular_images, custom_images) - Both image arrays
"""
try:
# Ensure civitai field exists in model_data
if not model_data.get('civitai'):
model_data['civitai'] = {}
# Ensure customImages array exists
if not model_data['civitai'].get('customImages'):
model_data['civitai']['customImages'] = []
# Get current customImages array
custom_images = model_data['civitai']['customImages']
await MetadataManager.hydrate_model_data(model_data)
civitai_data = model_data.setdefault('civitai', {})
custom_images = civitai_data.get('customImages')
if not isinstance(custom_images, list):
custom_images = []
civitai_data['customImages'] = custom_images
# Add new image entry for each imported file
for path_tuple in newly_imported_paths:
@@ -338,11 +333,8 @@ class MetadataUpdater:
file_path = model_data.get('file_path')
if file_path:
try:
# Create a copy of model data without 'folder' field
model_copy = model_data.copy()
model_copy.pop('folder', None)
# Write metadata to file
await MetadataManager.save_metadata(file_path, model_copy)
logger.info(f"Saved metadata for {model_data.get('model_name')}")
except Exception as e:
@@ -353,7 +345,7 @@ class MetadataUpdater:
await scanner.update_single_model_cache(file_path, file_path, model_data)
# Get regular images array (might be None)
regular_images = model_data['civitai'].get('images', [])
regular_images = civitai_data.get('images', [])
# Return both image arrays
return regular_images, custom_images

18
py/utils/example_images_processor.py
View File

@@ -475,15 +475,17 @@ class ExampleImagesProcessor:
'error': f"Model with hash {model_hash} not found in cache"
}, status=404)
# Check if model has custom images
if not model_data.get('civitai', {}).get('customImages'):
await MetadataManager.hydrate_model_data(model_data)
civitai_data = model_data.setdefault('civitai', {})
custom_images = civitai_data.get('customImages')
if not isinstance(custom_images, list) or not custom_images:
return web.json_response({
'success': False,
'error': f"Model has no custom images"
}, status=404)
# Find the custom image with matching short_id
custom_images = model_data['civitai']['customImages']
matching_image = None
new_custom_images = []
@@ -527,17 +529,15 @@ class ExampleImagesProcessor:
logger.warning(f"File for custom example with id {short_id} not found, but metadata will still be updated")
# Update metadata
model_data['civitai']['customImages'] = new_custom_images
civitai_data['customImages'] = new_custom_images
model_data.setdefault('civitai', {})['customImages'] = new_custom_images
# Save updated metadata to file
file_path = model_data.get('file_path')
if file_path:
try:
# Create a copy of model data without 'folder' field
model_copy = model_data.copy()
model_copy.pop('folder', None)
# Write metadata to file
await MetadataManager.save_metadata(file_path, model_copy)
logger.debug(f"Saved updated metadata for {model_data.get('model_name')}")
except Exception as e:
@@ -551,7 +551,7 @@ class ExampleImagesProcessor:
await scanner.update_single_model_cache(file_path, file_path, model_data)
# Get regular images array (might be None)
regular_images = model_data['civitai'].get('images', [])
regular_images = civitai_data.get('images', [])
return web.json_response({
'success': True,
@@ -568,4 +568,4 @@ class ExampleImagesProcessor:
}, status=500)

66
py/utils/metadata_manager.py
View File

@@ -2,7 +2,7 @@ from datetime import datetime
import os
import json
import logging
from typing import Dict, Optional, Type, Union
from typing import Any, Dict, Optional, Type, Union
from .models import BaseModelMetadata, LoraMetadata
from .file_utils import normalize_path, find_preview_file, calculate_sha256
@@ -53,6 +53,70 @@ class MetadataManager:
error_type = "Invalid JSON" if isinstance(e, json.JSONDecodeError) else "Parse error"
logger.error(f"{error_type} in metadata file: {metadata_path}. Error: {str(e)}. Skipping model to preserve existing data.")
return None, True # should_skip = True
@staticmethod
async def load_metadata_payload(file_path: str) -> Dict:
"""
Load metadata and return it as a dictionary, including any unknown fields.
Falls back to reading the raw JSON file if parsing into a model class fails.
"""
payload: Dict = {}
metadata_obj, should_skip = await MetadataManager.load_metadata(file_path)
if metadata_obj:
payload = metadata_obj.to_dict()
unknown_fields = getattr(metadata_obj, "_unknown_fields", None)
if isinstance(unknown_fields, dict):
payload.update(unknown_fields)
else:
if not should_skip:
metadata_path = (
file_path
if file_path.endswith(".metadata.json")
else f"{os.path.splitext(file_path)[0]}.metadata.json"
)
if os.path.exists(metadata_path):
try:
with open(metadata_path, "r", encoding="utf-8") as handle:
raw = json.load(handle)
if isinstance(raw, dict):
payload = raw
except json.JSONDecodeError:
logger.warning(
"Failed to parse metadata file %s while loading payload",
metadata_path,
)
except Exception as exc: # pragma: no cover - defensive logging
logger.warning("Failed to read metadata file %s: %s", metadata_path, exc)
if not isinstance(payload, dict):
payload = {}
if file_path:
payload.setdefault("file_path", normalize_path(file_path))
return payload
@staticmethod
async def hydrate_model_data(model_data: Dict[str, Any]) -> Dict[str, Any]:
"""
Replace the provided model data with the authoritative payload from disk.
Preserves the cached folder entry if present.
"""
file_path = model_data.get("file_path")
if not file_path:
return model_data
folder = model_data.get("folder")
payload = await MetadataManager.load_metadata_payload(file_path)
if folder is not None:
payload["folder"] = folder
model_data.clear()
model_data.update(payload)
return model_data
@staticmethod
async def save_metadata(path: str, metadata: Union[BaseModelMetadata, Dict]) -> bool: