diff --git a/py/utils/example_images_metadata.py b/py/utils/example_images_metadata.py index 10e9b113..92b59eb6 100644 --- a/py/utils/example_images_metadata.py +++ b/py/utils/example_images_metadata.py @@ -475,13 +475,19 @@ class MetadataUpdater: return False model_folder = get_model_folder(model_hash) - if not model_folder: + 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") @@ -493,24 +499,15 @@ class MetadataUpdater: if not img_id: continue - if not os.path.isdir(model_folder): + 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) - else: - found = False - try: - prefix = f"custom_{img_id}" - for fname in os.listdir(model_folder): - if fname.startswith(prefix) and os.path.isfile( - os.path.join(model_folder, fname) - ): - found = True - break - except OSError: - stale.append(idx) - continue - - if not found: - stale.append(idx) if stale: for idx in reversed(stale): @@ -532,22 +529,9 @@ class MetadataUpdater: # is gone. continue - if not os.path.isdir(model_folder): + prefix = f"image_{idx}." + if not any(f.startswith(prefix) for f in dir_entries): stale.append(idx) - else: - found = False - try: - prefix = f"image_{idx}." - for fname in os.listdir(model_folder): - if fname.startswith(prefix): - found = True - break - except OSError: - stale.append(idx) - continue - - if not found: - stale.append(idx) if stale: for idx in reversed(stale): diff --git a/py/utils/example_images_migration.py b/py/utils/example_images_migration.py index 9eab2822..77b2777c 100644 --- a/py/utils/example_images_migration.py +++ b/py/utils/example_images_migration.py @@ -3,9 +3,16 @@ import logging import os import re import json +import shutil from ..services.settings_manager import get_settings_manager from ..services.service_registry import ServiceRegistry -from ..utils.example_images_paths import iter_library_roots +from ..utils.example_images_paths import ( + get_example_images_root, + is_hash_folder, + iter_library_roots, + uses_library_scoped_folders, + _library_folder_has_only_hash_dirs, +) from ..utils.metadata_manager import MetadataManager from ..utils.example_images_processor import ExampleImagesProcessor from ..utils.constants import SUPPORTED_MEDIA_EXTENSIONS @@ -36,6 +43,90 @@ settings = _SettingsProxy() class ExampleImagesMigration: """Handles migrations for example images naming conventions""" + @staticmethod + def _consolidate_library_folders(): + """Move hash folders from library-named subdirectories back to root. + + When a user switches from multi-library mode back to single-library + mode, example images previously stored under e.g. + ``/default//`` need to be moved back to + ``//``. Running this once at startup removes the need + for ``get_model_folder()`` to perform directory scans on every + request. + """ + if uses_library_scoped_folders(): + return + + root = get_example_images_root() + if not root or not os.path.isdir(root): + return + + moved: list[str] = [] + cleaned: list[str] = [] + + try: + for entry in os.listdir(root): + # Fast regex checks first — no filesystem I/O. + if is_hash_folder(entry) or entry == "_deleted": + continue + + entry_path = os.path.join(root, entry) + if not os.path.isdir(entry_path): + continue + if not _library_folder_has_only_hash_dirs(entry_path): + continue + + try: + for hash_entry in os.listdir(entry_path): + hash_path = os.path.join(entry_path, hash_entry) + if not os.path.isdir(hash_path) or not is_hash_folder(hash_entry): + continue + target = os.path.join(root, hash_entry) + if not os.path.exists(target): + try: + shutil.move(hash_path, target) + moved.append(hash_entry) + except (OSError, shutil.Error) as exc: + logger.error( + "Failed to move '%s' → '%s': %s", + hash_path, target, exc, + ) + except OSError as exc: + logger.error( + "Failed to list library subdirectory '%s': %s", + entry_path, exc, + ) + + try: + remaining = os.listdir(entry_path) + except OSError: + remaining = [] + if not remaining: + try: + os.rmdir(entry_path) + cleaned.append(entry) + except OSError as exc: + logger.debug( + "Could not remove empty library dir '%s': %s", + entry_path, exc, + ) + except OSError as exc: + logger.error( + "Failed to list example images root during consolidation: %s", + exc, + ) + + if moved: + logger.info( + "Consolidated %d example image folder(s) to root", + len(moved), + ) + if cleaned: + logger.info( + "Removed %d empty library directories", + len(cleaned), + ) + @staticmethod async def check_and_run_migrations(): """Check if migrations are needed and run them in background""" @@ -44,6 +135,10 @@ class ExampleImagesMigration: logger.debug("No example images path configured or path doesn't exist, skipping migrations") return + # Run library-to-root consolidation once at startup so the hot + # path (get_model_folder) stays a pure-path computation. + ExampleImagesMigration._consolidate_library_folders() + for library_name, library_path in iter_library_roots(): if not library_path or not os.path.exists(library_path): continue diff --git a/py/utils/example_images_paths.py b/py/utils/example_images_paths.py index 44a59f6a..de628d7b 100644 --- a/py/utils/example_images_paths.py +++ b/py/utils/example_images_paths.py @@ -83,7 +83,12 @@ def ensure_library_root_exists(library_name: Optional[str] = None) -> str: def get_model_folder(model_hash: str, library_name: Optional[str] = None) -> str: - """Return the folder path for a model's example images.""" + """Return the folder path for a model's example images. + + Multi-library ↔ single-library consolidation is handled once at startup by + ``ExampleImagesMigration._consolidate_library_folders`` — this function is a + pure path computation on the hot path (no directory scans). + """ if not model_hash: return "" @@ -113,35 +118,6 @@ def get_model_folder(model_hash: str, library_name: Optional[str] = None) -> str exc, ) return legacy_folder - elif not os.path.exists(resolved_folder): - # Reverse migration: when consolidating from multi-library to - # single-library mode (e.g. after "default" was cleaned up), look - # for existing example images inside library-named subdirectories - # and bring them back to the root level. - root = get_example_images_root() - if root: - try: - for entry in os.listdir(root): - entry_path = os.path.join(root, entry) - if not os.path.isdir(entry_path): - continue - if is_hash_folder(entry) or entry == "_deleted": - continue - if not _library_folder_has_only_hash_dirs(entry_path): - continue - legacy = os.path.join(entry_path, normalized_hash) - if os.path.exists(legacy): - shutil.move(legacy, resolved_folder) - logger.info( - "Consolidated example images from '%s' to '%s'", - legacy, resolved_folder, - ) - break - except OSError as exc: - logger.error( - "Failed to consolidate example images during " - "library merge: %s", exc, - ) return resolved_folder