From 28e93d12ff6051185c7e9c51a46510aee0da72ef Mon Sep 17 00:00:00 2001 From: Will Miao Date: Mon, 3 Aug 2026 12:04:56 +0800 Subject: [PATCH] fix(example-images): use in-place cache sync and bulk pending-check index for large libraries --- py/utils/example_images_download_manager.py | 89 ++++++++++- py/utils/example_images_metadata.py | 56 +++++-- py/utils/example_images_migration.py | 3 +- py/utils/example_images_processor.py | 6 +- tests/services/test_check_pending_models.py | 142 ++++++++++++++++++ ...t_example_images_download_manager_async.py | 12 ++ tests/utils/test_example_images_metadata.py | 11 +- 7 files changed, 294 insertions(+), 25 deletions(-) diff --git a/py/utils/example_images_download_manager.py b/py/utils/example_images_download_manager.py index 994ca62d..fd46a921 100644 --- a/py/utils/example_images_download_manager.py +++ b/py/utils/example_images_download_manager.py @@ -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/ 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( diff --git a/py/utils/example_images_metadata.py b/py/utils/example_images_metadata.py index 92b59eb6..42c1529f 100644 --- a/py/utils/example_images_metadata.py +++ b/py/utils/example_images_metadata.py @@ -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', []) diff --git a/py/utils/example_images_migration.py b/py/utils/example_images_migration.py index 77b2777c..99786f1e 100644 --- a/py/utils/example_images_migration.py +++ b/py/utils/example_images_migration.py @@ -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: diff --git a/py/utils/example_images_processor.py b/py/utils/example_images_processor.py index da617706..b25d9c97 100644 --- a/py/utils/example_images_processor.py +++ b/py/utils/example_images_processor.py @@ -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, diff --git a/tests/services/test_check_pending_models.py b/tests/services/test_check_pending_models.py index 518e665f..a5c4e790 100644 --- a/tests/services/test_check_pending_models.py +++ b/tests/services/test_check_pending_models.py @@ -363,6 +363,148 @@ async def test_check_pending_models_handles_corrupted_progress_file( assert result["pending_count"] == 1 +@pytest.mark.asyncio +@pytest.mark.usefixtures("tmp_path") +async def test_check_pending_models_uses_bulk_folder_index_for_large_libraries( + monkeypatch: pytest.MonkeyPatch, + tmp_path, + settings_manager, +): + """For >1000 candidates the pre-check scans the library root once instead of + probing every folder individually.""" + + ws_manager = RecordingWebSocketManager() + manager = download_module.DownloadManager(ws_manager=ws_manager) + + monkeypatch.setitem(settings_manager.settings, "example_images_path", str(tmp_path)) + + # 1500 unprocessed models triggers the bulk lookup path + models = [ + {"sha256": f"{i:064x}", "model_name": f"Model {i}"} + for i in range(1500) + ] + + # Create folders with files for the first 500 models + for i in range(500): + model_dir = tmp_path / f"{i:064x}" + model_dir.mkdir() + (model_dir / "image_0.png").write_text("data") + + _patch_scanners(monkeypatch, lora_scanner=StubScanner(models)) + + per_model_checks = 0 + + def counting_model_directory_has_files(path: str) -> bool: + nonlocal per_model_checks + per_model_checks += 1 + return False + + monkeypatch.setattr( + download_module, + "_model_directory_has_files", + counting_model_directory_has_files, + ) + + result = await manager.check_pending_models(["lora"]) + + assert result["success"] is True + assert result["total_models"] == 1500 + assert result["pending_count"] == 1000 + assert result["needs_download"] is True + # The per-folder check should not be used once we cross the threshold. + assert per_model_checks == 0 + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("tmp_path") +async def test_check_pending_models_uses_per_folder_check_for_small_candidate_sets( + monkeypatch: pytest.MonkeyPatch, + tmp_path, + settings_manager, +): + """For <=1000 candidates the pre-check keeps the accurate per-folder path.""" + + ws_manager = RecordingWebSocketManager() + manager = download_module.DownloadManager(ws_manager=ws_manager) + + monkeypatch.setitem(settings_manager.settings, "example_images_path", str(tmp_path)) + + models = [ + {"sha256": f"{i:064x}", "model_name": f"Model {i}"} + for i in range(500) + ] + + # Create folders with files for the first 200 models + for i in range(200): + model_dir = tmp_path / f"{i:064x}" + model_dir.mkdir() + (model_dir / "image_0.png").write_text("data") + + _patch_scanners(monkeypatch, lora_scanner=StubScanner(models)) + + per_model_checks = 0 + original_has_files = download_module._model_directory_has_files + + def counting_model_directory_has_files(path: str) -> bool: + nonlocal per_model_checks + per_model_checks += 1 + return original_has_files(path) + + monkeypatch.setattr( + download_module, + "_model_directory_has_files", + counting_model_directory_has_files, + ) + + result = await manager.check_pending_models(["lora"]) + + assert result["success"] is True + assert result["total_models"] == 500 + assert result["pending_count"] == 300 + assert result["needs_download"] is True + # Per-folder path should run once per candidate. + assert per_model_checks == 500 + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("tmp_path") +async def test_check_pending_models_bulk_index_includes_legacy_folders( + monkeypatch: pytest.MonkeyPatch, + tmp_path, + settings_manager, +): + """In multi-library mode the bulk index also scans the legacy root so models + whose folders have not been consolidated yet are not reported pending.""" + + ws_manager = RecordingWebSocketManager() + manager = download_module.DownloadManager(ws_manager=ws_manager) + + monkeypatch.setitem(settings_manager.settings, "example_images_path", str(tmp_path)) + monkeypatch.setitem(settings_manager.settings, "libraries", {"default": {}, "extra": {}}) + monkeypatch.setitem(settings_manager.settings, "active_library", "extra") + + # 1500 unprocessed models triggers the bulk lookup path + models = [ + {"sha256": f"{i:064x}", "model_name": f"Model {i}"} + for i in range(1500) + ] + + # Folders live at the LEGACY root/ path (not yet consolidated) + for i in range(500): + model_dir = tmp_path / f"{i:064x}" + model_dir.mkdir() + (model_dir / "image_0.png").write_text("data") + + _patch_scanners(monkeypatch, lora_scanner=StubScanner(models)) + + result = await manager.check_pending_models(["lora"]) + + assert result["success"] is True + assert result["total_models"] == 1500 + assert result["pending_count"] == 1000 + assert result["needs_download"] is True + + @pytest.fixture def settings_manager(): return get_settings_manager() diff --git a/tests/services/test_example_images_download_manager_async.py b/tests/services/test_example_images_download_manager_async.py index 40de65b6..a7fb553a 100644 --- a/tests/services/test_example_images_download_manager_async.py +++ b/tests/services/test_example_images_download_manager_async.py @@ -26,6 +26,7 @@ class StubScanner: def __init__(self, models: list[dict]) -> None: self._cache = SimpleNamespace(raw_data=models) + self.sync_calls: list[tuple[str, dict]] = [] async def get_cached_data(self): return self._cache @@ -38,6 +39,14 @@ class StubScanner: break return True + async def sync_cache_from_metadata(self, file_path: str, metadata: dict) -> bool: + self.sync_calls.append((file_path, metadata)) + for index, model in enumerate(self._cache.raw_data): + if model.get("file_path") == metadata.get("file_path"): + self._cache.raw_data[index] = metadata + break + return True + def _patch_scanner(monkeypatch: pytest.MonkeyPatch, scanner: StubScanner) -> None: async def _get_lora_scanner(cls): @@ -588,6 +597,9 @@ async def test_not_found_example_images_are_cleaned( assert missing_url in downloader.calls assert manager._progress["failed_models"] == {model_hash} assert model_hash in manager._progress["processed_models"] + assert scanner.sync_calls + assert len(scanner.sync_calls) == 1 + assert scanner.sync_calls[0][0] == str(model_path) remaining_images = model_metadata["civitai"]["images"] assert remaining_images == [ diff --git a/tests/utils/test_example_images_metadata.py b/tests/utils/test_example_images_metadata.py index 4e4bd513..7498d679 100644 --- a/tests/utils/test_example_images_metadata.py +++ b/tests/utils/test_example_images_metadata.py @@ -15,6 +15,7 @@ class StubScanner: def __init__(self, cache_items: List[Dict[str, Any]]) -> None: self.cache = SimpleNamespace(raw_data=cache_items) self.updates: List[Tuple[str, str, Dict[str, Any]]] = [] + self.sync_updates: List[Tuple[str, Dict[str, Any]]] = [] async def get_cached_data(self): return self.cache @@ -23,6 +24,10 @@ class StubScanner: self.updates.append((old_path, new_path, metadata)) return True + async def sync_cache_from_metadata(self, file_path: str, metadata: Dict[str, Any]) -> bool: + self.sync_updates.append((file_path, metadata)) + return True + @pytest.fixture(autouse=True) def patch_metadata_manager(monkeypatch: pytest.MonkeyPatch): @@ -83,7 +88,7 @@ async def test_update_metadata_after_import_enriches_entries(monkeypatch: pytest assert custom[0]["type"] == "image" assert Path(patch_metadata_manager[0][0]) == model_file - assert scanner.updates + assert scanner.sync_updates @pytest.mark.asyncio @@ -151,8 +156,8 @@ async def test_update_metadata_after_import_preserves_existing_metadata( assert saved_payload["civitai"]["trainedWords"] == ["foo"] assert {entry["id"] for entry in saved_payload["civitai"]["customImages"]} == {"existing-id", "new-id"} - assert scanner.updates - updated_metadata = scanner.updates[-1][2] + assert scanner.sync_updates + updated_metadata = scanner.sync_updates[-1][1] assert updated_metadata["civitai"]["images"] == existing_payload["civitai"]["images"] assert {entry["id"] for entry in updated_metadata["civitai"]["customImages"]} == {"existing-id", "new-id"}