mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-09-21 03:01:27 -03:00
Will Miao
7d963b27b5
Example videos added through the "Add examples" flow were stored with a hardcoded 720x1280 entry. The dimension probe next to it only ran for images (PIL cannot open .mp4/.webm files), so every video entry stayed portrait regardless of the source. The showcase viewer then sizes its container straight from that value (--media-aspect in showcase.css), so landscape clips were letterboxed inside a 9:16 box. CivitAI-sourced examples were unaffected because their dimensions come from the API. PIL cannot read video containers, so add a dependency-free reader that parses the container headers instead: moov/trak/tkhd for ISO base media (with the sample description as a fallback), Segment/Tracks/Pixel* for WebM/Matroska, and RIFF/WebP for animated examples saved with a video extension. The sniffed signature decides which reader runs, so a .mp4 that is really WebM still reports the right size; the extension is only a fallback. Both readers seek past mdat rather than reading it, so a large file costs the same as a small one. Imported entries now record the file's real size and keep the previous placeholder only when the file cannot be parsed. Existing libraries keep their wrong entries, so backfill them once via the existing naming migration: bump CURRENT_NAMING_VERSION to 3 and repair each model's empty-url entries from the files on disk, then sync the scanner cache. Only entries with no remote url are touched -- those have no other source, which makes the rewrite lossless -- and entries already carrying the right size are left byte-identical, so the pass is idempotent and a no-op for libraries that never imported a video.
711 lines
28 KiB
Python
711 lines
28 KiB
Python
import inspect
|
||
import logging
|
||
import os
|
||
import re
|
||
from typing import TYPE_CHECKING, Any, Dict, Mapping, MutableMapping, Optional
|
||
|
||
from ..recipes.constants import GEN_PARAM_KEYS
|
||
from ..services.metadata_service import get_default_metadata_provider, get_metadata_provider
|
||
from ..services.metadata_sync_service import MetadataSyncService
|
||
from ..services.preview_asset_service import PreviewAssetService
|
||
from ..services.settings_manager import get_settings_manager
|
||
from ..services.downloader import get_downloader
|
||
from ..utils.constants import SUPPORTED_MEDIA_EXTENSIONS
|
||
from ..utils.exif_utils import ExifUtils
|
||
from ..utils.metadata_manager import MetadataManager
|
||
from ..utils.video_metadata import get_video_dimensions
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# Placeholder dimensions written when the real ones cannot be determined.
|
||
# Kept for backwards compatibility with pre-existing metadata entries.
|
||
_DEFAULT_MEDIA_WIDTH = 720
|
||
_DEFAULT_MEDIA_HEIGHT = 1280
|
||
|
||
# Example metadata entries carry a marker: ``customImages`` use their ``id``
|
||
# while ``images`` use the positional index. Either way the marker must be a
|
||
# plain filename-safe token, never a path fragment.
|
||
_ENTRY_MARKER_PATTERN = re.compile(r"^(?:custom_|image_)?([^./\\]+)$")
|
||
|
||
_preview_service = PreviewAssetService(
|
||
metadata_manager=MetadataManager,
|
||
downloader_factory=get_downloader,
|
||
exif_utils=ExifUtils,
|
||
)
|
||
|
||
_metadata_sync_service: MetadataSyncService | None = None
|
||
_metadata_sync_service_settings: Optional["SettingsManager"] = None
|
||
|
||
if TYPE_CHECKING: # pragma: no cover - import for type checkers only
|
||
from ..services.settings_manager import SettingsManager
|
||
|
||
|
||
async def update_cache_from_metadata(
|
||
scanner: Any, file_path: str, metadata: Dict[str, Any]
|
||
) -> bool:
|
||
"""Update the scanner cache from a metadata dict using the in-place sync path.
|
||
|
||
``sync_cache_from_metadata`` patches the existing cache entry incrementally
|
||
(tag/hash/version indexes, targeted single-row SQL update) and only resorts
|
||
when a sort-key field changed. This avoids the ``O(n)`` full-list resort and
|
||
full cache rewrite that ``update_single_model_cache`` performs on every call,
|
||
which is critical for libraries with 100k+ models.
|
||
|
||
Falls back to the legacy full update when the scanner does not expose an
|
||
async ``sync_cache_from_metadata`` method.
|
||
|
||
Returns:
|
||
``True`` if the cache entry was updated, ``False`` otherwise.
|
||
"""
|
||
|
||
sync_method = getattr(scanner, "sync_cache_from_metadata", None)
|
||
if inspect.iscoroutinefunction(sync_method):
|
||
return await sync_method(file_path, metadata)
|
||
|
||
return await scanner.update_single_model_cache(file_path, file_path, metadata)
|
||
|
||
|
||
def _build_metadata_sync_service(settings_manager: "SettingsManager") -> MetadataSyncService:
|
||
"""Construct a metadata sync service bound to the provided settings."""
|
||
|
||
return MetadataSyncService(
|
||
metadata_manager=MetadataManager,
|
||
preview_service=_preview_service,
|
||
settings=settings_manager,
|
||
default_metadata_provider_factory=get_default_metadata_provider,
|
||
metadata_provider_selector=get_metadata_provider,
|
||
)
|
||
|
||
|
||
def _read_media_dimensions(path: str, is_video: bool) -> tuple[int, int]:
|
||
"""Return ``(width, height)`` for an example image or video file.
|
||
|
||
Videos are read from their container headers (PIL cannot open them) so the
|
||
showcase viewer sizes the gallery to the real aspect ratio. Falls back to
|
||
the legacy ``720x1280`` placeholder when the dimensions cannot be
|
||
determined — e.g. an unreadable file or an exotic codec — which only
|
||
affects the displayed aspect ratio, never the file itself.
|
||
"""
|
||
|
||
dimensions = None
|
||
|
||
if is_video:
|
||
dimensions = get_video_dimensions(path)
|
||
else:
|
||
try:
|
||
from PIL import Image
|
||
|
||
if os.path.exists(path):
|
||
with Image.open(path) as img:
|
||
dimensions = img.size
|
||
except Exception:
|
||
dimensions = None
|
||
|
||
if dimensions:
|
||
width, height = dimensions
|
||
if width > 0 and height > 0:
|
||
return int(width), int(height)
|
||
|
||
return _DEFAULT_MEDIA_WIDTH, _DEFAULT_MEDIA_HEIGHT
|
||
|
||
|
||
def _is_video_entry(file_path: Optional[str], entry: Mapping[str, Any]) -> bool:
|
||
"""Return True when an example entry points at a video file.
|
||
|
||
The local file extension wins over the recorded ``type`` because files in
|
||
the wild are frequently mislabelled (animated WebP saved as ``.mp4``);
|
||
``_read_media_dimensions`` handles that correctly either way.
|
||
"""
|
||
|
||
if file_path:
|
||
ext = os.path.splitext(file_path)[1].lower()
|
||
if ext in SUPPORTED_MEDIA_EXTENSIONS["videos"]:
|
||
return True
|
||
if ext in SUPPORTED_MEDIA_EXTENSIONS["images"]:
|
||
return False
|
||
return str(entry.get("type", "")).lower() == "video"
|
||
|
||
|
||
def _resolve_local_file(
|
||
entry: Mapping[str, Any],
|
||
index: int,
|
||
local_files: Mapping[str, str],
|
||
) -> Optional[str]:
|
||
"""Map a metadata entry onto its example file inside the model folder.
|
||
|
||
Reads the entry's own marker (``id`` for ``customImages``, positional
|
||
``index`` for ``images``) with an anchored regex, so the identifier can
|
||
never bleed into a neighbouring filename the way a prefix comparison can.
|
||
"""
|
||
|
||
marker = entry.get("id")
|
||
if not isinstance(marker, str) or not marker:
|
||
marker = str(index)
|
||
|
||
match = _ENTRY_MARKER_PATTERN.fullmatch(marker)
|
||
if not match:
|
||
return None
|
||
|
||
return local_files.get(match.group(1))
|
||
|
||
|
||
def repair_local_video_dimensions(
|
||
metadata: MutableMapping[str, Any],
|
||
local_files: Mapping[str, str],
|
||
*,
|
||
dry_run: bool = False,
|
||
) -> int:
|
||
"""Backfill real video dimensions for an entry that has local files.
|
||
|
||
Only entries with an empty ``url`` are considered: those have no remote
|
||
source, so the local file is the single source of truth for their size and
|
||
rewriting them cannot discard API-supplied data. Entries whose dimensions
|
||
already match the file are left byte-identical.
|
||
|
||
Args:
|
||
metadata: Raw metadata payload (mutated in place unless ``dry_run``).
|
||
local_files: ``{identifier: path}`` for files present in the model's
|
||
example folder, where the identifier is the entry's ``id`` (for
|
||
``customImages``) or its positional index (for ``images``).
|
||
dry_run: Count the fixes without mutating ``metadata``.
|
||
|
||
Returns:
|
||
The number of entries that were (or would be) repaired.
|
||
"""
|
||
|
||
civitai = metadata.get("civitai")
|
||
if not isinstance(civitai, dict):
|
||
return 0
|
||
|
||
repaired = 0
|
||
|
||
for key in ("customImages", "images"):
|
||
entries = civitai.get(key)
|
||
if not isinstance(entries, list) or not entries:
|
||
continue
|
||
|
||
for index, entry in enumerate(entries):
|
||
if not isinstance(entry, dict):
|
||
continue
|
||
if entry.get("url", "") != "":
|
||
# Remote-backed entry: never rebuilt from local state.
|
||
continue
|
||
|
||
file_path = _resolve_local_file(entry, index, local_files)
|
||
if not file_path or not os.path.isfile(file_path):
|
||
continue
|
||
|
||
dimensions = _read_media_dimensions(
|
||
file_path, _is_video_entry(file_path, entry)
|
||
)
|
||
width, height = dimensions
|
||
if width <= 0 or height <= 0:
|
||
continue
|
||
if entry.get("width") == width and entry.get("height") == height:
|
||
continue
|
||
|
||
if not dry_run:
|
||
entry["width"] = width
|
||
entry["height"] = height
|
||
repaired += 1
|
||
|
||
return repaired
|
||
|
||
|
||
def _get_metadata_sync_service() -> MetadataSyncService:
|
||
"""Return the shared metadata sync service, initialising it lazily."""
|
||
|
||
global _metadata_sync_service, _metadata_sync_service_settings
|
||
|
||
settings_manager = get_settings_manager()
|
||
|
||
if isinstance(_metadata_sync_service, MetadataSyncService):
|
||
if _metadata_sync_service_settings is not settings_manager:
|
||
_metadata_sync_service = _build_metadata_sync_service(settings_manager)
|
||
_metadata_sync_service_settings = settings_manager
|
||
elif _metadata_sync_service is None:
|
||
_metadata_sync_service = _build_metadata_sync_service(settings_manager)
|
||
_metadata_sync_service_settings = settings_manager
|
||
else:
|
||
# Tests may inject stand-ins that do not match the sync service type. Preserve
|
||
# those injections while still updating our cached settings reference so the
|
||
# next real service instantiation uses the current configuration.
|
||
_metadata_sync_service_settings = settings_manager
|
||
|
||
return _metadata_sync_service
|
||
|
||
|
||
class MetadataUpdater:
|
||
"""Handles updating model metadata related to example images"""
|
||
|
||
@staticmethod
|
||
async def refresh_model_metadata(model_hash, model_name, scanner_type, scanner, progress: dict[str, Any] | None = None):
|
||
"""Refresh model metadata from CivitAI
|
||
|
||
Args:
|
||
model_hash: SHA256 hash of the model
|
||
model_name: Model name (for logging)
|
||
scanner_type: Scanner type ('lora' or 'checkpoint')
|
||
scanner: Scanner instance for this model type
|
||
|
||
Returns:
|
||
bool: True if metadata was successfully refreshed, False otherwise
|
||
"""
|
||
try:
|
||
# Find the model in the scanner cache
|
||
cache = await scanner.get_cached_data()
|
||
model_data = None
|
||
|
||
for item in cache.raw_data:
|
||
if item.get('sha256') == model_hash:
|
||
model_data = item
|
||
break
|
||
|
||
if not model_data:
|
||
logger.warning(f"Model {model_name} with hash {model_hash} not found in cache")
|
||
return False
|
||
|
||
file_path = model_data.get('file_path')
|
||
if not file_path:
|
||
logger.warning(f"Model {model_name} has no file path")
|
||
return False
|
||
|
||
# Track that we're refreshing this model
|
||
if progress is not None:
|
||
progress['refreshed_models'].add(model_hash)
|
||
|
||
async def update_cache_func(old_path, new_path, metadata):
|
||
return await update_cache_from_metadata(scanner, 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,
|
||
model_data=model_data,
|
||
update_cache_func=update_cache_func,
|
||
)
|
||
|
||
if success:
|
||
logger.info(f"Successfully refreshed metadata for {model_name}")
|
||
return True
|
||
else:
|
||
logger.warning(f"Failed to refresh metadata for {model_name}, {error}")
|
||
return False
|
||
|
||
except Exception as e:
|
||
error_msg = f"Error refreshing metadata for {model_name}: {str(e)}"
|
||
logger.error(error_msg, exc_info=True)
|
||
if progress is not None:
|
||
progress['errors'].append(error_msg)
|
||
progress['last_error'] = error_msg
|
||
return False
|
||
|
||
@staticmethod
|
||
async def get_updated_model(model_hash, scanner):
|
||
"""Load the most recent metadata for a model identified by hash."""
|
||
cache = await scanner.get_cached_data()
|
||
target = None
|
||
for item in cache.raw_data:
|
||
if item.get('sha256') == model_hash:
|
||
target = item
|
||
break
|
||
|
||
if not target:
|
||
return None
|
||
|
||
file_path = target.get('file_path')
|
||
if not file_path:
|
||
return target
|
||
|
||
model_cls = getattr(scanner, 'model_class', None)
|
||
if model_cls is None:
|
||
metadata, should_skip = await MetadataManager.load_metadata(file_path)
|
||
else:
|
||
metadata, should_skip = await MetadataManager.load_metadata(file_path, model_cls)
|
||
|
||
if should_skip or metadata is None:
|
||
return target
|
||
|
||
rich_metadata = metadata.to_dict()
|
||
rich_metadata.setdefault('folder', target.get('folder', ''))
|
||
return rich_metadata
|
||
|
||
|
||
@staticmethod
|
||
async def update_metadata_from_local_examples(model_hash, model, scanner_type, scanner, model_dir):
|
||
"""Update model metadata with local example image information
|
||
|
||
Args:
|
||
model_hash: SHA256 hash of the model
|
||
model: Model data dictionary
|
||
scanner_type: Scanner type ('lora' or 'checkpoint')
|
||
scanner: Scanner instance for this model type
|
||
model_dir: Model images directory
|
||
|
||
Returns:
|
||
bool: True if metadata was successfully updated, False otherwise
|
||
"""
|
||
try:
|
||
# Collect local image paths
|
||
local_images_paths = []
|
||
if os.path.exists(model_dir):
|
||
for file in os.listdir(model_dir):
|
||
file_path = os.path.join(model_dir, file)
|
||
if os.path.isfile(file_path):
|
||
file_ext = os.path.splitext(file)[1].lower()
|
||
is_supported = (file_ext in SUPPORTED_MEDIA_EXTENSIONS['images'] or
|
||
file_ext in SUPPORTED_MEDIA_EXTENSIONS['videos'])
|
||
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 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
|
||
# Create images array
|
||
images = []
|
||
|
||
# Generate metadata for each local image/video
|
||
for path in local_images_paths:
|
||
# Determine if video or image
|
||
file_ext = os.path.splitext(path)[1].lower()
|
||
is_video = file_ext in SUPPORTED_MEDIA_EXTENSIONS['videos']
|
||
|
||
width, height = _read_media_dimensions(path, is_video)
|
||
|
||
# Create image metadata entry
|
||
image_entry = {
|
||
"url": "", # Empty URL as required
|
||
"nsfwLevel": 0,
|
||
"width": width,
|
||
"height": height,
|
||
"type": "video" if is_video else "image",
|
||
"meta": None,
|
||
"hasMeta": False,
|
||
"hasPositivePrompt": False
|
||
}
|
||
|
||
images.append(image_entry)
|
||
|
||
# Update the model's civitai.images field
|
||
civitai_data['images'] = images
|
||
|
||
# Save metadata to .metadata.json file
|
||
file_path = model.get('file_path')
|
||
model_copy: Optional[Dict[str, Any]] = None
|
||
try:
|
||
model_copy = model.copy()
|
||
if model_copy is not None:
|
||
model_copy.pop('folder', None)
|
||
await MetadataManager.save_metadata(file_path, model_copy)
|
||
logger.info(f"Saved metadata for {model.get('model_name')}")
|
||
except Exception as e:
|
||
logger.error(f"Failed to save metadata for {model.get('model_name')}: {str(e)}")
|
||
|
||
# Save updated metadata to scanner cache. sync_cache_from_metadata
|
||
# returns False both for "already in sync" and for actual failures,
|
||
# so the cache sync result is deliberately not treated as an error;
|
||
# the return value reflects whether the metadata was persisted.
|
||
if file_path and model_copy is not None:
|
||
await update_cache_from_metadata(scanner, file_path, model_copy)
|
||
logger.info(f"Successfully updated metadata for {model.get('model_name')} with {len(images)} local examples")
|
||
return True
|
||
|
||
logger.warning(f"Failed to update metadata for {model.get('model_name')}")
|
||
return False
|
||
|
||
return False
|
||
except Exception as e:
|
||
logger.error(f"Error updating metadata from local examples: {str(e)}", exc_info=True)
|
||
return False
|
||
|
||
@staticmethod
|
||
async def update_metadata_after_import(model_hash, model_data, scanner, newly_imported_paths):
|
||
"""Update model metadata after importing example images
|
||
|
||
Args:
|
||
model_hash: SHA256 hash of the model
|
||
model_data: Model data dictionary
|
||
scanner: Scanner instance (lora or checkpoint)
|
||
newly_imported_paths: List of paths to newly imported files
|
||
|
||
Returns:
|
||
tuple: (regular_images, custom_images) - Both image arrays
|
||
"""
|
||
try:
|
||
await MetadataManager.hydrate_model_data(model_data)
|
||
civitai_data = model_data.get('civitai')
|
||
|
||
if not isinstance(civitai_data, dict):
|
||
civitai_data = {}
|
||
model_data['civitai'] = civitai_data
|
||
|
||
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:
|
||
path, short_id = path_tuple
|
||
|
||
# Determine if video or image
|
||
file_ext = os.path.splitext(path)[1].lower()
|
||
is_video = file_ext in SUPPORTED_MEDIA_EXTENSIONS['videos']
|
||
|
||
width, height = _read_media_dimensions(path, is_video)
|
||
|
||
# Create image metadata entry
|
||
image_entry = {
|
||
"url": "", # Empty URL as requested
|
||
"id": short_id,
|
||
"nsfwLevel": 0,
|
||
"width": width,
|
||
"height": height,
|
||
"type": "video" if is_video else "image",
|
||
"meta": None,
|
||
"hasMeta": False,
|
||
"hasPositivePrompt": False
|
||
}
|
||
|
||
# Extract and parse metadata if this is an image
|
||
if not is_video:
|
||
try:
|
||
# Extract metadata from image
|
||
extracted_metadata = ExifUtils.extract_image_metadata(path)
|
||
|
||
if extracted_metadata:
|
||
# Parse the extracted metadata to get generation parameters
|
||
parsed_meta = MetadataUpdater._parse_image_metadata(extracted_metadata)
|
||
|
||
if parsed_meta:
|
||
image_entry["meta"] = parsed_meta
|
||
image_entry["hasMeta"] = True
|
||
image_entry["hasPositivePrompt"] = bool(parsed_meta.get("prompt", ""))
|
||
logger.debug(f"Extracted metadata from {os.path.basename(path)}")
|
||
except Exception as e:
|
||
logger.warning(f"Failed to extract metadata from {os.path.basename(path)}: {e}")
|
||
|
||
# Append to existing customImages array
|
||
custom_images.append(image_entry)
|
||
|
||
# Save metadata to .metadata.json file
|
||
file_path = model_data.get('file_path')
|
||
model_copy: Optional[Dict[str, Any]] = None
|
||
if file_path:
|
||
try:
|
||
model_copy = model_data.copy()
|
||
if model_copy is not None:
|
||
model_copy.pop('folder', None)
|
||
await MetadataManager.save_metadata(file_path, model_copy)
|
||
logger.info(f"Saved metadata for {model_data.get('model_name')}")
|
||
except Exception as e:
|
||
logger.error(f"Failed to save metadata: {str(e)}")
|
||
|
||
# Save updated metadata to scanner cache
|
||
if file_path and model_copy is not None:
|
||
await update_cache_from_metadata(scanner, file_path, model_copy)
|
||
|
||
# Get regular images array (might be None)
|
||
regular_images = civitai_data.get('images', [])
|
||
|
||
# Return both image arrays
|
||
return regular_images, custom_images
|
||
|
||
except Exception as e:
|
||
logger.error(f"Failed to update metadata after import: {e}", exc_info=True)
|
||
return [], []
|
||
|
||
@staticmethod
|
||
def _parse_image_metadata(user_comment):
|
||
"""Parse metadata from image to extract generation parameters
|
||
|
||
Args:
|
||
user_comment: Metadata string extracted from image
|
||
|
||
Returns:
|
||
dict: Parsed metadata with generation parameters
|
||
"""
|
||
if not user_comment:
|
||
return None
|
||
|
||
try:
|
||
# Initialize metadata dictionary
|
||
metadata = {}
|
||
|
||
# Split on Negative prompt if it exists
|
||
if "Negative prompt:" in user_comment:
|
||
parts = user_comment.split('Negative prompt:', 1)
|
||
prompt = parts[0].strip()
|
||
negative_and_params = parts[1] if len(parts) > 1 else ""
|
||
else:
|
||
# No negative prompt section
|
||
param_start = re.search(r'Steps: \d+', user_comment)
|
||
if param_start:
|
||
prompt = user_comment[:param_start.start()].strip()
|
||
negative_and_params = user_comment[param_start.start():]
|
||
else:
|
||
prompt = user_comment.strip()
|
||
negative_and_params = ""
|
||
|
||
# Add prompt if it's in GEN_PARAM_KEYS
|
||
if 'prompt' in GEN_PARAM_KEYS:
|
||
metadata['prompt'] = prompt
|
||
|
||
# Extract negative prompt and parameters
|
||
if negative_and_params:
|
||
# If we split on "Negative prompt:", check for params section
|
||
if "Negative prompt:" in user_comment:
|
||
param_start = re.search(r'Steps: ', negative_and_params)
|
||
if param_start:
|
||
neg_prompt = negative_and_params[:param_start.start()].strip()
|
||
if 'negative_prompt' in GEN_PARAM_KEYS:
|
||
metadata['negative_prompt'] = neg_prompt
|
||
params_section = negative_and_params[param_start.start():]
|
||
else:
|
||
if 'negative_prompt' in GEN_PARAM_KEYS:
|
||
metadata['negative_prompt'] = negative_and_params.strip()
|
||
params_section = ""
|
||
else:
|
||
# No negative prompt, entire section is params
|
||
params_section = negative_and_params
|
||
|
||
# Extract generation parameters
|
||
if params_section:
|
||
# Extract basic parameters
|
||
param_pattern = r'([A-Za-z\s]+): ([^,]+)'
|
||
params = re.findall(param_pattern, params_section)
|
||
|
||
for key, value in params:
|
||
clean_key = key.strip().lower().replace(' ', '_')
|
||
|
||
# Skip if not in recognized gen param keys
|
||
if clean_key not in GEN_PARAM_KEYS:
|
||
continue
|
||
|
||
# Convert numeric values
|
||
if clean_key in ['steps', 'seed']:
|
||
try:
|
||
metadata[clean_key] = int(value.strip())
|
||
except ValueError:
|
||
metadata[clean_key] = value.strip()
|
||
elif clean_key in ['cfg_scale']:
|
||
try:
|
||
metadata[clean_key] = float(value.strip())
|
||
except ValueError:
|
||
metadata[clean_key] = value.strip()
|
||
else:
|
||
metadata[clean_key] = value.strip()
|
||
|
||
# Extract size if available and add if a recognized key
|
||
size_match = re.search(r'Size: (\d+)x(\d+)', params_section)
|
||
if size_match and 'size' in GEN_PARAM_KEYS:
|
||
width, height = size_match.groups()
|
||
metadata['size'] = f"{width}x{height}"
|
||
|
||
# Return metadata if we have any entries
|
||
return metadata if metadata else None
|
||
|
||
except Exception as e:
|
||
logger.error(f"Error parsing image metadata: {e}", exc_info=True)
|
||
return None
|
||
|
||
@staticmethod
|
||
async def prune_stale_example_images(metadata) -> bool:
|
||
"""Remove example-image metadata entries whose files no longer exist on disk.
|
||
|
||
Checks ``civitai.customImages`` (by ``id``) and ``civitai.images`` entries
|
||
that have an empty ``url`` (no remote fallback) against actual files in
|
||
the model's example-image folder. Stale entries are removed in-place so
|
||
the caller can persist the cleaned metadata afterwards.
|
||
|
||
Args:
|
||
metadata: A ``BaseModelMetadata`` instance (modified in place).
|
||
|
||
Returns:
|
||
True if at least one entry was removed.
|
||
"""
|
||
from ..utils.example_images_paths import get_model_folder
|
||
|
||
model_hash = getattr(metadata, "sha256", None)
|
||
if not model_hash:
|
||
return False
|
||
|
||
model_folder = get_model_folder(model_hash)
|
||
if not model_folder or not os.path.isdir(model_folder):
|
||
return False
|
||
|
||
civitai = getattr(metadata, "civitai", None)
|
||
if not isinstance(civitai, dict):
|
||
return False
|
||
|
||
# Read the directory listing once so every image entry reuses it.
|
||
try:
|
||
dir_entries = os.listdir(model_folder)
|
||
except OSError:
|
||
dir_entries = []
|
||
|
||
has_changes = False
|
||
|
||
custom_images = civitai.get("customImages")
|
||
if isinstance(custom_images, list) and custom_images:
|
||
stale: list[int] = []
|
||
|
||
for idx, img in enumerate(custom_images):
|
||
img_id = img.get("id", "")
|
||
if not img_id:
|
||
continue
|
||
|
||
prefix = f"custom_{img_id}"
|
||
found = any(
|
||
f.startswith(prefix) and os.path.isfile(
|
||
os.path.join(model_folder, f)
|
||
)
|
||
for f in dir_entries
|
||
)
|
||
if not found:
|
||
stale.append(idx)
|
||
|
||
if stale:
|
||
for idx in reversed(stale):
|
||
custom_images.pop(idx)
|
||
has_changes = True
|
||
logger.info(
|
||
"Pruned %d stale custom image(s) for %s",
|
||
len(stale),
|
||
getattr(metadata, "model_name", model_hash),
|
||
)
|
||
|
||
images = civitai.get("images")
|
||
if isinstance(images, list) and images:
|
||
stale_images: list[int] = []
|
||
|
||
for idx, img in enumerate(images):
|
||
if img.get("url", ""):
|
||
# Has a remote fallback – keep it even if the local copy
|
||
# is gone.
|
||
continue
|
||
|
||
prefix = f"image_{idx}."
|
||
if not any(f.startswith(prefix) for f in dir_entries):
|
||
stale_images.append(idx)
|
||
|
||
if stale_images:
|
||
for idx in reversed(stale_images):
|
||
images.pop(idx)
|
||
has_changes = True
|
||
logger.info(
|
||
"Pruned %d stale image entry(ies) for %s",
|
||
len(stale_images),
|
||
getattr(metadata, "model_name", model_hash),
|
||
)
|
||
|
||
return has_changes
|