feat: update metadata saving to ensure backup creation and support nested civitai structure

This commit is contained in:
Will Miao
2025-06-25 11:50:10 +08:00
parent afc810f21f
commit 2ee057e19b
5 changed files with 18 additions and 11 deletions

View File

@@ -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}")

View File

@@ -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":

View File

@@ -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}")

View File

@@ -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)}")

View File

@@ -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(