mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-08-06 22:10:14 -03:00
fix(example-images): use in-place cache sync and bulk pending-check index for large libraries
This commit is contained in:
@@ -14,11 +14,16 @@ from ..services.service_registry import ServiceRegistry
|
|||||||
from ..utils.example_images_paths import (
|
from ..utils.example_images_paths import (
|
||||||
ExampleImagePathResolver,
|
ExampleImagePathResolver,
|
||||||
ensure_library_root_exists,
|
ensure_library_root_exists,
|
||||||
|
get_example_images_root,
|
||||||
|
is_hash_folder,
|
||||||
uses_library_scoped_folders,
|
uses_library_scoped_folders,
|
||||||
)
|
)
|
||||||
from ..utils.metadata_manager import MetadataManager
|
from ..utils.metadata_manager import MetadataManager
|
||||||
from .example_images_processor import ExampleImagesProcessor
|
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.downloader import get_downloader
|
||||||
from ..services.settings_manager import get_settings_manager
|
from ..services.settings_manager import get_settings_manager
|
||||||
|
|
||||||
@@ -87,6 +92,13 @@ class _DownloadProgress(dict):
|
|||||||
return snapshot
|
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:
|
def _model_directory_has_files(path: str) -> bool:
|
||||||
"""Return True when the provided directory exists and contains entries."""
|
"""Return True when the provided directory exists and contains entries."""
|
||||||
|
|
||||||
@@ -103,6 +115,36 @@ def _model_directory_has_files(path: str) -> bool:
|
|||||||
return False
|
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:
|
class DownloadManager:
|
||||||
"""Manages downloading example images for models."""
|
"""Manages downloading example images for models."""
|
||||||
|
|
||||||
@@ -410,14 +452,49 @@ class DownloadManager:
|
|||||||
# Calculate pending count: check which models actually need processing.
|
# 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,
|
# 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.
|
# and its folder doesn't exist or is empty.
|
||||||
pending_hashes = set()
|
candidate_hashes = [
|
||||||
for model_hash, model_name in all_models_with_hash:
|
model_hash
|
||||||
if model_hash not in processed_models and model_hash not in failed_models:
|
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_dir = ExampleImagePathResolver.get_model_folder(
|
||||||
model_hash, active_library
|
model_hash, active_library
|
||||||
)
|
)
|
||||||
if not _model_directory_has_files(model_dir):
|
if not _model_directory_has_files(model_dir):
|
||||||
pending_hashes.add(model_hash)
|
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)
|
pending_count = len(pending_hashes)
|
||||||
|
|
||||||
@@ -1343,8 +1420,8 @@ class DownloadManager:
|
|||||||
await MetadataManager.save_metadata(file_path, model_copy)
|
await MetadataManager.save_metadata(file_path, model_copy)
|
||||||
|
|
||||||
try:
|
try:
|
||||||
await scanner.update_single_model_cache(
|
await update_cache_from_metadata(
|
||||||
file_path, file_path, model_data
|
scanner, file_path, model_copy
|
||||||
)
|
)
|
||||||
except AttributeError:
|
except AttributeError:
|
||||||
logger.debug(
|
logger.debug(
|
||||||
|
|||||||
@@ -1,3 +1,4 @@
|
|||||||
|
import inspect
|
||||||
import logging
|
import logging
|
||||||
import os
|
import os
|
||||||
import re
|
import re
|
||||||
@@ -28,6 +29,31 @@ if TYPE_CHECKING: # pragma: no cover - import for type checkers only
|
|||||||
from ..services.settings_manager import SettingsManager
|
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:
|
def _build_metadata_sync_service(settings_manager: "SettingsManager") -> MetadataSyncService:
|
||||||
"""Construct a metadata sync service bound to the provided settings."""
|
"""Construct a metadata sync service bound to the provided settings."""
|
||||||
|
|
||||||
@@ -103,7 +129,7 @@ class MetadataUpdater:
|
|||||||
progress['refreshed_models'].add(model_hash)
|
progress['refreshed_models'].add(model_hash)
|
||||||
|
|
||||||
async def update_cache_func(old_path, new_path, metadata):
|
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)
|
await MetadataManager.hydrate_model_data(model_data)
|
||||||
success, error = await _get_metadata_sync_service().fetch_and_update_model(
|
success, error = await _get_metadata_sync_service().fetch_and_update_model(
|
||||||
@@ -234,6 +260,7 @@ class MetadataUpdater:
|
|||||||
|
|
||||||
# Save metadata to .metadata.json file
|
# Save metadata to .metadata.json file
|
||||||
file_path = model.get('file_path')
|
file_path = model.get('file_path')
|
||||||
|
model_copy: Optional[Dict[str, Any]] = None
|
||||||
try:
|
try:
|
||||||
model_copy = model.copy()
|
model_copy = model.copy()
|
||||||
model_copy.pop('folder', None)
|
model_copy.pop('folder', None)
|
||||||
@@ -242,13 +269,17 @@ class MetadataUpdater:
|
|||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Failed to save metadata for {model.get('model_name')}: {str(e)}")
|
logger.error(f"Failed to save metadata for {model.get('model_name')}: {str(e)}")
|
||||||
|
|
||||||
# Save updated metadata to scanner cache
|
# Save updated metadata to scanner cache. sync_cache_from_metadata
|
||||||
success = await scanner.update_single_model_cache(file_path, file_path, model)
|
# returns False both for "already in sync" and for actual failures,
|
||||||
if success:
|
# 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")
|
logger.info(f"Successfully updated metadata for {model.get('model_name')} with {len(images)} local examples")
|
||||||
return True
|
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
|
return False
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
@@ -336,6 +367,7 @@ class MetadataUpdater:
|
|||||||
|
|
||||||
# Save metadata to .metadata.json file
|
# Save metadata to .metadata.json file
|
||||||
file_path = model_data.get('file_path')
|
file_path = model_data.get('file_path')
|
||||||
|
model_copy: Optional[Dict[str, Any]] = None
|
||||||
if file_path:
|
if file_path:
|
||||||
try:
|
try:
|
||||||
model_copy = model_data.copy()
|
model_copy = model_data.copy()
|
||||||
@@ -346,8 +378,8 @@ class MetadataUpdater:
|
|||||||
logger.error(f"Failed to save metadata: {str(e)}")
|
logger.error(f"Failed to save metadata: {str(e)}")
|
||||||
|
|
||||||
# Save updated metadata to scanner cache
|
# Save updated metadata to scanner cache
|
||||||
if file_path:
|
if file_path and model_copy is not None:
|
||||||
await scanner.update_single_model_cache(file_path, file_path, model_data)
|
await update_cache_from_metadata(scanner, file_path, model_copy)
|
||||||
|
|
||||||
# Get regular images array (might be None)
|
# Get regular images array (might be None)
|
||||||
regular_images = civitai_data.get('images', [])
|
regular_images = civitai_data.get('images', [])
|
||||||
|
|||||||
@@ -15,6 +15,7 @@ from ..utils.example_images_paths import (
|
|||||||
)
|
)
|
||||||
from ..utils.metadata_manager import MetadataManager
|
from ..utils.metadata_manager import MetadataManager
|
||||||
from ..utils.example_images_processor import ExampleImagesProcessor
|
from ..utils.example_images_processor import ExampleImagesProcessor
|
||||||
|
from ..utils.example_images_metadata import update_cache_from_metadata
|
||||||
from ..utils.constants import SUPPORTED_MEDIA_EXTENSIONS
|
from ..utils.constants import SUPPORTED_MEDIA_EXTENSIONS
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
@@ -421,7 +422,7 @@ class ExampleImagesMigration:
|
|||||||
await MetadataManager.save_metadata(file_path, model_copy)
|
await MetadataManager.save_metadata(file_path, model_copy)
|
||||||
|
|
||||||
# Update scanner cache
|
# 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
|
updated_models += 1
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
|
|||||||
@@ -9,7 +9,7 @@ from ..utils.constants import SUPPORTED_MEDIA_EXTENSIONS
|
|||||||
from ..services.service_registry import ServiceRegistry
|
from ..services.service_registry import ServiceRegistry
|
||||||
from ..services.settings_manager import get_settings_manager
|
from ..services.settings_manager import get_settings_manager
|
||||||
from ..utils.example_images_paths import get_model_folder, get_model_relative_path
|
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
|
from ..utils.metadata_manager import MetadataManager
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
@@ -644,7 +644,7 @@ class ExampleImagesProcessor:
|
|||||||
}, status=500)
|
}, status=500)
|
||||||
|
|
||||||
# Update cache
|
# 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)
|
# Get regular images array (might be None)
|
||||||
regular_images = civitai_data.get('images', [])
|
regular_images = civitai_data.get('images', [])
|
||||||
@@ -759,7 +759,7 @@ class ExampleImagesProcessor:
|
|||||||
model_copy = model_data.copy()
|
model_copy = model_data.copy()
|
||||||
model_copy.pop('folder', None)
|
model_copy.pop('folder', None)
|
||||||
await MetadataManager.save_metadata(file_path, model_copy)
|
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({
|
return web.json_response({
|
||||||
'success': True,
|
'success': True,
|
||||||
|
|||||||
@@ -363,6 +363,148 @@ async def test_check_pending_models_handles_corrupted_progress_file(
|
|||||||
assert result["pending_count"] == 1
|
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/<hash> 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
|
@pytest.fixture
|
||||||
def settings_manager():
|
def settings_manager():
|
||||||
return get_settings_manager()
|
return get_settings_manager()
|
||||||
|
|||||||
@@ -26,6 +26,7 @@ class StubScanner:
|
|||||||
|
|
||||||
def __init__(self, models: list[dict]) -> None:
|
def __init__(self, models: list[dict]) -> None:
|
||||||
self._cache = SimpleNamespace(raw_data=models)
|
self._cache = SimpleNamespace(raw_data=models)
|
||||||
|
self.sync_calls: list[tuple[str, dict]] = []
|
||||||
|
|
||||||
async def get_cached_data(self):
|
async def get_cached_data(self):
|
||||||
return self._cache
|
return self._cache
|
||||||
@@ -38,6 +39,14 @@ class StubScanner:
|
|||||||
break
|
break
|
||||||
return True
|
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:
|
def _patch_scanner(monkeypatch: pytest.MonkeyPatch, scanner: StubScanner) -> None:
|
||||||
async def _get_lora_scanner(cls):
|
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 missing_url in downloader.calls
|
||||||
assert manager._progress["failed_models"] == {model_hash}
|
assert manager._progress["failed_models"] == {model_hash}
|
||||||
assert model_hash in manager._progress["processed_models"]
|
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"]
|
remaining_images = model_metadata["civitai"]["images"]
|
||||||
assert remaining_images == [
|
assert remaining_images == [
|
||||||
|
|||||||
@@ -15,6 +15,7 @@ class StubScanner:
|
|||||||
def __init__(self, cache_items: List[Dict[str, Any]]) -> None:
|
def __init__(self, cache_items: List[Dict[str, Any]]) -> None:
|
||||||
self.cache = SimpleNamespace(raw_data=cache_items)
|
self.cache = SimpleNamespace(raw_data=cache_items)
|
||||||
self.updates: List[Tuple[str, str, Dict[str, Any]]] = []
|
self.updates: List[Tuple[str, str, Dict[str, Any]]] = []
|
||||||
|
self.sync_updates: List[Tuple[str, Dict[str, Any]]] = []
|
||||||
|
|
||||||
async def get_cached_data(self):
|
async def get_cached_data(self):
|
||||||
return self.cache
|
return self.cache
|
||||||
@@ -23,6 +24,10 @@ class StubScanner:
|
|||||||
self.updates.append((old_path, new_path, metadata))
|
self.updates.append((old_path, new_path, metadata))
|
||||||
return True
|
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)
|
@pytest.fixture(autouse=True)
|
||||||
def patch_metadata_manager(monkeypatch: pytest.MonkeyPatch):
|
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 custom[0]["type"] == "image"
|
||||||
|
|
||||||
assert Path(patch_metadata_manager[0][0]) == model_file
|
assert Path(patch_metadata_manager[0][0]) == model_file
|
||||||
assert scanner.updates
|
assert scanner.sync_updates
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
@@ -151,8 +156,8 @@ async def test_update_metadata_after_import_preserves_existing_metadata(
|
|||||||
assert saved_payload["civitai"]["trainedWords"] == ["foo"]
|
assert saved_payload["civitai"]["trainedWords"] == ["foo"]
|
||||||
assert {entry["id"] for entry in saved_payload["civitai"]["customImages"]} == {"existing-id", "new-id"}
|
assert {entry["id"] for entry in saved_payload["civitai"]["customImages"]} == {"existing-id", "new-id"}
|
||||||
|
|
||||||
assert scanner.updates
|
assert scanner.sync_updates
|
||||||
updated_metadata = scanner.updates[-1][2]
|
updated_metadata = scanner.sync_updates[-1][1]
|
||||||
assert updated_metadata["civitai"]["images"] == existing_payload["civitai"]["images"]
|
assert updated_metadata["civitai"]["images"] == existing_payload["civitai"]["images"]
|
||||||
assert {entry["id"] for entry in updated_metadata["civitai"]["customImages"]} == {"existing-id", "new-id"}
|
assert {entry["id"] for entry in updated_metadata["civitai"]["customImages"]} == {"existing-id", "new-id"}
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user