From 2ee057e19b0adb4307d0ae54ed40633b732913d1 Mon Sep 17 00:00:00 2001 From: Will Miao <13051207myq@gmail.com> Date: Wed, 25 Jun 2025 11:50:10 +0800 Subject: [PATCH] feat: update metadata saving to ensure backup creation and support nested civitai structure --- py/routes/api_routes.py | 8 ++++++-- py/services/download_manager.py | 2 +- py/services/model_scanner.py | 6 +++--- py/utils/metadata_manager.py | 11 +++++++---- py/utils/routes_common.py | 2 +- 5 files changed, 18 insertions(+), 11 deletions(-) diff --git a/py/routes/api_routes.py b/py/routes/api_routes.py index 78aee080..01c6e0cf 100644 --- a/py/routes/api_routes.py +++ b/py/routes/api_routes.py @@ -792,9 +792,13 @@ class ApiRoutes: metadata['modelDescription'] = description metadata['tags'] = tags - metadata['creator'] = creator + # Ensure the civitai dict exists + if 'civitai' not in metadata: + metadata['civitai'] = {} + # Store creator in the civitai nested structure + metadata['civitai']['creator'] = creator - await MetadataManager.save_metadata(file_path, metadata) + await MetadataManager.save_metadata(file_path, metadata, True) except Exception as e: logger.error(f"Error saving model metadata: {e}") diff --git a/py/services/download_manager.py b/py/services/download_manager.py index 3ce4f083..2a0ae048 100644 --- a/py/services/download_manager.py +++ b/py/services/download_manager.py @@ -255,7 +255,7 @@ class DownloadManager: metadata.update_file_info(save_path) # 5. Final metadata update - await MetadataManager.save_metadata(save_path, metadata) + await MetadataManager.save_metadata(save_path, metadata, True) # 6. Update cache based on model type if model_type == "checkpoint": diff --git a/py/services/model_scanner.py b/py/services/model_scanner.py index 38096cb6..f9008013 100644 --- a/py/services/model_scanner.py +++ b/py/services/model_scanner.py @@ -788,7 +788,7 @@ class ModelScanner: metadata = self.model_class.from_civitai_info(version_info, file_info, file_path) metadata.preview_url = find_preview_file(file_name, os.path.dirname(file_path)) - await MetadataManager.save_metadata(file_path, metadata) + await MetadataManager.save_metadata(file_path, metadata, True) logger.debug(f"Created metadata from .civitai.info for {file_path}") except Exception as e: logger.error(f"Error creating metadata from .civitai.info for {file_path}: {e}") @@ -815,7 +815,7 @@ class ModelScanner: metadata.modelDescription = version_info['model']['description'] # Save the updated metadata - await MetadataManager.save_metadata(file_path, metadata) + await MetadataManager.save_metadata(file_path, metadata, True) logger.debug(f"Updated metadata with civitai info for {file_path}") except Exception as e: logger.error(f"Error restoring civitai data from .civitai.info for {file_path}: {e}") @@ -884,7 +884,7 @@ class ModelScanner: model_data['civitai']['creator'] = model_metadata['creator'] - await MetadataManager.save_metadata(file_path, model_data) + await MetadataManager.save_metadata(file_path, model_data, True) except Exception as e: logger.error(f"Failed to update metadata from Civitai for {file_path}: {e}") diff --git a/py/utils/metadata_manager.py b/py/utils/metadata_manager.py index 3bdcd542..8d3a7a27 100644 --- a/py/utils/metadata_manager.py +++ b/py/utils/metadata_manager.py @@ -91,14 +91,14 @@ class MetadataManager: return None @staticmethod - async def save_metadata(path: str, metadata: Union[BaseModelMetadata, Dict], create_backup: bool = True) -> bool: + async def save_metadata(path: str, metadata: Union[BaseModelMetadata, Dict], create_backup: bool = False) -> bool: """ Save metadata with atomic write operations and backup creation. Args: path: Path to the model file or directly to the metadata file metadata: Metadata to save (either BaseModelMetadata object or dict) - create_backup: Whether to create a backup of existing file + create_backup: Whether to create a new backup of existing file if a backup doesn't already exist Returns: bool: Success or failure @@ -114,10 +114,13 @@ class MetadataManager: backup_path = f"{metadata_path}.bak" try: - # Create backup if requested and file exists - if create_backup and os.path.exists(metadata_path): + # Create backup if file exists and either: + # 1. create_backup is True, OR + # 2. backup file doesn't already exist + if os.path.exists(metadata_path) and (create_backup or not os.path.exists(backup_path)): try: shutil.copy2(metadata_path, backup_path) + logger.debug(f"Created metadata backup at: {backup_path}") except Exception as e: logger.warning(f"Failed to create metadata backup: {str(e)}") diff --git a/py/utils/routes_common.py b/py/utils/routes_common.py index 56135289..d51419a0 100644 --- a/py/utils/routes_common.py +++ b/py/utils/routes_common.py @@ -154,7 +154,7 @@ class ModelRouteUtils: local_metadata['preview_nsfw_level'] = first_preview.get('nsfwLevel', 0) # Save updated metadata - await MetadataManager.save_metadata(metadata_path, local_metadata) + await MetadataManager.save_metadata(metadata_path, local_metadata, True) @staticmethod async def fetch_and_update_model(