fix(example-images): use in-place cache sync and bulk pending-check index for large libraries

This commit is contained in:
Will Miao
2026-08-03 12:04:56 +08:00
parent 75e63c758b
commit 28e93d12ff
7 changed files with 294 additions and 25 deletions

View File

@@ -14,11 +14,16 @@ from ..services.service_registry import ServiceRegistry
from ..utils.example_images_paths import (
ExampleImagePathResolver,
ensure_library_root_exists,
get_example_images_root,
is_hash_folder,
uses_library_scoped_folders,
)
from ..utils.metadata_manager import MetadataManager
from .example_images_processor import ExampleImagesProcessor
from .example_images_metadata import MetadataUpdater
from .example_images_metadata import (
MetadataUpdater,
update_cache_from_metadata,
)
from ..services.downloader import get_downloader
from ..services.settings_manager import get_settings_manager
@@ -87,6 +92,13 @@ class _DownloadProgress(dict):
return snapshot
# When fewer candidates than this remain in check_pending_models, probe each
# model folder directly (preserving legacy-folder migration semantics). Above
# it, build a folder index with a single directory scan so libraries with
# 100k+ models do not pay one syscall per candidate.
_BULK_LOOKUP_THRESHOLD = 1000
def _model_directory_has_files(path: str) -> bool:
"""Return True when the provided directory exists and contains entries."""
@@ -103,6 +115,36 @@ def _model_directory_has_files(path: str) -> bool:
return False
def _build_example_folder_index(output_dir: str) -> dict[str, bool]:
"""Build a ``{hash: has_files}`` index for a library's example-image folders.
A single directory scan over the library root replaces ``O(candidates)``
per-folder ``os.scandir`` calls, which is required for libraries with
100k+ models. Each hash folder is classified by whether it contains any
entries, matching the semantics of ``_model_directory_has_files``.
"""
index: dict[str, bool] = {}
if not output_dir or not os.path.isdir(output_dir):
return index
try:
with os.scandir(output_dir) as entries:
for entry in entries:
name = entry.name
if not entry.is_dir() or not is_hash_folder(name):
continue
try:
with os.scandir(entry.path) as subentries:
index[name.lower()] = any(subentries)
except OSError:
index[name.lower()] = False
except OSError:
pass
return index
class DownloadManager:
"""Manages downloading example images for models."""
@@ -410,14 +452,49 @@ class DownloadManager:
# Calculate pending count: check which models actually need processing.
# A model is pending if it has a hash, is not already processed or known-failed,
# and its folder doesn't exist or is empty.
pending_hashes = set()
for model_hash, model_name in all_models_with_hash:
if model_hash not in processed_models and model_hash not in failed_models:
candidate_hashes = [
model_hash
for model_hash, _ in all_models_with_hash
if model_hash not in processed_models
and model_hash not in failed_models
]
pending_hashes: set[str] = set()
# For small candidate counts the existing per-folder check is fine
# and handles legacy folder migration.
# For large libraries, scan the library root once and do set lookups.
if len(candidate_hashes) <= _BULK_LOOKUP_THRESHOLD or not output_dir:
for model_hash in candidate_hashes:
model_dir = ExampleImagePathResolver.get_model_folder(
model_hash, active_library
)
if not _model_directory_has_files(model_dir):
pending_hashes.add(model_hash)
else:
folder_index = await asyncio.get_event_loop().run_in_executor(
None, _build_example_folder_index, output_dir
)
# In multi-library mode, folders that have not been consolidated
# into the library root yet (startup migration skipped, failed
# move, or created at the legacy path afterwards) still live at
# the legacy root/<hash> location. Only scan that root when at
# least one candidate is missing from the library-root index, so
# the fully-consolidated case does not pay an extra directory
# pass on every call.
if uses_library_scoped_folders() and any(
not folder_index.get(model_hash, False)
for model_hash in candidate_hashes
):
legacy_root = get_example_images_root()
if legacy_root and legacy_root != output_dir:
legacy_index = await asyncio.get_event_loop().run_in_executor(
None, _build_example_folder_index, legacy_root
)
for hash_key, has_files in legacy_index.items():
folder_index.setdefault(hash_key, has_files)
for model_hash in candidate_hashes:
if not folder_index.get(model_hash, False):
pending_hashes.add(model_hash)
pending_count = len(pending_hashes)
@@ -1343,8 +1420,8 @@ class DownloadManager:
await MetadataManager.save_metadata(file_path, model_copy)
try:
await scanner.update_single_model_cache(
file_path, file_path, model_data
await update_cache_from_metadata(
scanner, file_path, model_copy
)
except AttributeError:
logger.debug(

View File

@@ -1,3 +1,4 @@
import inspect
import logging
import os
import re
@@ -28,6 +29,31 @@ 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."""
@@ -103,8 +129,8 @@ class MetadataUpdater:
progress['refreshed_models'].add(model_hash)
async def update_cache_func(old_path, new_path, metadata):
return await scanner.update_single_model_cache(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,
@@ -234,6 +260,7 @@ class MetadataUpdater:
# Save metadata to .metadata.json file
file_path = model.get('file_path')
model_copy: Optional[Dict[str, Any]] = None
try:
model_copy = model.copy()
model_copy.pop('folder', None)
@@ -241,14 +268,18 @@ class MetadataUpdater:
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
success = await scanner.update_single_model_cache(file_path, file_path, model)
if success:
# 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
else:
logger.warning(f"Failed to update metadata for {model.get('model_name')}")
logger.warning(f"Failed to update metadata for {model.get('model_name')}")
return False
return False
except Exception as e:
@@ -336,6 +367,7 @@ class MetadataUpdater:
# 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()
@@ -344,11 +376,11 @@ class MetadataUpdater:
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:
await scanner.update_single_model_cache(file_path, file_path, model_data)
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', [])

View File

@@ -15,6 +15,7 @@ from ..utils.example_images_paths import (
)
from ..utils.metadata_manager import MetadataManager
from ..utils.example_images_processor import ExampleImagesProcessor
from ..utils.example_images_metadata import update_cache_from_metadata
from ..utils.constants import SUPPORTED_MEDIA_EXTENSIONS
logger = logging.getLogger(__name__)
@@ -421,7 +422,7 @@ class ExampleImagesMigration:
await MetadataManager.save_metadata(file_path, model_copy)
# Update scanner cache
await scanner.update_single_model_cache(file_path, file_path, model_metadata)
await update_cache_from_metadata(scanner, file_path, model_copy)
updated_models += 1
except Exception as e:

View File

@@ -9,7 +9,7 @@ from ..utils.constants import SUPPORTED_MEDIA_EXTENSIONS
from ..services.service_registry import ServiceRegistry
from ..services.settings_manager import get_settings_manager
from ..utils.example_images_paths import get_model_folder, get_model_relative_path
from .example_images_metadata import MetadataUpdater
from .example_images_metadata import MetadataUpdater, update_cache_from_metadata
from ..utils.metadata_manager import MetadataManager
logger = logging.getLogger(__name__)
@@ -644,7 +644,7 @@ class ExampleImagesProcessor:
}, status=500)
# Update cache
await scanner.update_single_model_cache(file_path, file_path, model_data)
await update_cache_from_metadata(scanner, file_path, model_data)
# Get regular images array (might be None)
regular_images = civitai_data.get('images', [])
@@ -759,7 +759,7 @@ class ExampleImagesProcessor:
model_copy = model_data.copy()
model_copy.pop('folder', None)
await MetadataManager.save_metadata(file_path, model_copy)
await scanner.update_single_model_cache(file_path, file_path, model_data)
await update_cache_from_metadata(scanner, file_path, model_copy)
return web.json_response({
'success': True,