mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-03-23 22:22:11 -03:00
feat(persistent-cache): implement SQLite-based persistent model cache with loading and saving functionality
This commit is contained in:
Will Miao
@@ -345,14 +345,19 @@ class DownloadManager:
|
||||
self._progress['processed_models'].add(model_hash)
|
||||
return False # Return False to indicate no remote download happened
|
||||
|
||||
full_model = await MetadataUpdater.get_updated_model(
|
||||
model_hash, scanner
|
||||
)
|
||||
civitai_payload = (full_model or {}).get('civitai', {}) if full_model else {}
|
||||
|
||||
# If no local images, try to download from remote
|
||||
elif model.get('civitai') and model.get('civitai', {}).get('images'):
|
||||
images = model.get('civitai', {}).get('images', [])
|
||||
|
||||
if civitai_payload.get('images'):
|
||||
images = civitai_payload.get('images', [])
|
||||
|
||||
success, is_stale = await ExampleImagesProcessor.download_model_images(
|
||||
model_hash, model_name, images, model_dir, optimize, downloader
|
||||
)
|
||||
|
||||
|
||||
# If metadata is stale, try to refresh it
|
||||
if is_stale and model_hash not in self._progress['refreshed_models']:
|
||||
await MetadataUpdater.refresh_model_metadata(
|
||||
@@ -363,16 +368,17 @@ class DownloadManager:
|
||||
updated_model = await MetadataUpdater.get_updated_model(
|
||||
model_hash, scanner
|
||||
)
|
||||
updated_civitai = (updated_model or {}).get('civitai', {}) if updated_model else {}
|
||||
|
||||
if updated_model and updated_model.get('civitai', {}).get('images'):
|
||||
if updated_civitai.get('images'):
|
||||
# Retry download with updated metadata
|
||||
updated_images = updated_model.get('civitai', {}).get('images', [])
|
||||
updated_images = updated_civitai.get('images', [])
|
||||
success, _ = await ExampleImagesProcessor.download_model_images(
|
||||
model_hash, model_name, updated_images, model_dir, optimize, downloader
|
||||
)
|
||||
|
||||
self._progress['refreshed_models'].add(model_hash)
|
||||
|
||||
|
||||
# Mark as processed if successful, or as failed if unsuccessful after refresh
|
||||
if success:
|
||||
self._progress['processed_models'].add(model_hash)
|
||||
@@ -381,13 +387,13 @@ class DownloadManager:
|
||||
if model_hash in self._progress['refreshed_models']:
|
||||
self._progress['failed_models'].add(model_hash)
|
||||
logger.info(f"Marking model {model_name} as failed after metadata refresh")
|
||||
|
||||
|
||||
return True # Return True to indicate a remote download happened
|
||||
else:
|
||||
# No civitai data or images available, mark as failed to avoid future attempts
|
||||
self._progress['failed_models'].add(model_hash)
|
||||
logger.debug(f"No civitai images available for model {model_name}, marking as failed")
|
||||
|
||||
|
||||
# Save progress periodically
|
||||
if self._progress['completed'] % 10 == 0 or self._progress['completed'] == self._progress['total'] - 1:
|
||||
self._save_progress(output_dir)
|
||||
@@ -627,51 +633,59 @@ class DownloadManager:
|
||||
self._progress['processed_models'].add(model_hash)
|
||||
return False # Return False to indicate no remote download happened
|
||||
|
||||
full_model = await MetadataUpdater.get_updated_model(
|
||||
model_hash, scanner
|
||||
)
|
||||
civitai_payload = (full_model or {}).get('civitai', {}) if full_model else {}
|
||||
|
||||
# If no local images, try to download from remote
|
||||
elif model.get('civitai') and model.get('civitai', {}).get('images'):
|
||||
images = model.get('civitai', {}).get('images', [])
|
||||
|
||||
if civitai_payload.get('images'):
|
||||
images = civitai_payload.get('images', [])
|
||||
|
||||
success, is_stale, failed_images = await ExampleImagesProcessor.download_model_images_with_tracking(
|
||||
model_hash, model_name, images, model_dir, optimize, downloader
|
||||
)
|
||||
|
||||
|
||||
# If metadata is stale, try to refresh it
|
||||
if is_stale and model_hash not in self._progress['refreshed_models']:
|
||||
await MetadataUpdater.refresh_model_metadata(
|
||||
model_hash, model_name, scanner_type, scanner, self._progress
|
||||
)
|
||||
|
||||
|
||||
# Get the updated model data
|
||||
updated_model = await MetadataUpdater.get_updated_model(
|
||||
model_hash, scanner
|
||||
)
|
||||
|
||||
if updated_model and updated_model.get('civitai', {}).get('images'):
|
||||
updated_civitai = (updated_model or {}).get('civitai', {}) if updated_model else {}
|
||||
|
||||
if updated_civitai.get('images'):
|
||||
# Retry download with updated metadata
|
||||
updated_images = updated_model.get('civitai', {}).get('images', [])
|
||||
updated_images = updated_civitai.get('images', [])
|
||||
success, _, additional_failed_images = await ExampleImagesProcessor.download_model_images_with_tracking(
|
||||
model_hash, model_name, updated_images, model_dir, optimize, downloader
|
||||
)
|
||||
|
||||
|
||||
# Combine failed images from both attempts
|
||||
failed_images.extend(additional_failed_images)
|
||||
|
||||
|
||||
self._progress['refreshed_models'].add(model_hash)
|
||||
|
||||
|
||||
# For forced downloads, remove failed images from metadata
|
||||
if failed_images:
|
||||
# Create a copy of images excluding failed ones
|
||||
await self._remove_failed_images_from_metadata(
|
||||
model_hash, model_name, failed_images, scanner
|
||||
)
|
||||
|
||||
|
||||
# Mark as processed
|
||||
if success or failed_images: # Mark as processed if we successfully downloaded some images or removed failed ones
|
||||
self._progress['processed_models'].add(model_hash)
|
||||
|
||||
|
||||
return True # Return True to indicate a remote download happened
|
||||
else:
|
||||
logger.debug(f"No civitai images available for model {model_name}")
|
||||
|
||||
|
||||
return False
|
||||
|
||||
except Exception as e:
|
||||
|
||||
@@ -95,21 +95,35 @@ class MetadataUpdater:
|
||||
|
||||
@staticmethod
|
||||
async def get_updated_model(model_hash, scanner):
|
||||
"""Get updated model data
|
||||
|
||||
Args:
|
||||
model_hash: SHA256 hash of the model
|
||||
scanner: Scanner instance
|
||||
|
||||
Returns:
|
||||
dict: Updated model data or None if not found
|
||||
"""
|
||||
"""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:
|
||||
return item
|
||||
return None
|
||||
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user