Enhance metadata processing in ModelScanner: prevent intermediate writes, restore missing civitai data, and ensure base_model consistency. #185

py/services/model_scanner.py

@@ -6,6 +6,8 @@ import time
 import shutil
 from typing import List, Dict, Optional, Type, Set
 
+from ..utils.model_utils import determine_base_model
+
 from ..utils.models import BaseModelMetadata
 from ..config import config
 from ..utils.file_utils import load_metadata, get_file_info, find_preview_file, save_metadata
@@ -537,6 +539,34 @@ class ModelScanner:
     # Common methods shared between scanners
     async def _process_model_file(self, file_path: str, root_path: str) -> Dict:
         """Process a single model file and return its metadata"""
+        needs_metadata_update = False
+        original_save_metadata = save_metadata
+        
+        # Temporarily override save_metadata to prevent intermediate writes
+        async def no_op_save(*args, **kwargs):
+            nonlocal needs_metadata_update
+            needs_metadata_update = True
+            return None
+            
+        # Use a context manager to temporarily replace save_metadata
+        from contextlib import contextmanager
+        
+        @contextmanager
+        def prevent_metadata_writes():
+            nonlocal needs_metadata_update
+            # Replace the function temporarily
+            import sys
+            from .. import utils
+            original = utils.file_utils.save_metadata
+            utils.file_utils.save_metadata = no_op_save
+            try:
+                yield
+            finally:
+                # Restore the original function
+                utils.file_utils.save_metadata = original
+        
+        # Process with write prevention
+        with prevent_metadata_writes():
             metadata = await load_metadata(file_path, self.model_class)
             
             if metadata is None:
@@ -553,7 +583,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 save_metadata(file_path, metadata)
+                            needs_metadata_update = 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}")
@@ -568,38 +598,54 @@ class ModelScanner:
                             
                             logger.debug(f"Restoring missing civitai data from .civitai.info for {file_path}")
                             metadata.civitai = version_info
+                            needs_metadata_update = True
                             
                             # Ensure tags are also updated if they're missing
                             if (not metadata.tags or len(metadata.tags) == 0) and 'model' in version_info:
                                 if 'tags' in version_info['model']:
                                     metadata.tags = version_info['model']['tags']
+                                    needs_metadata_update = True
                             
                             # Also restore description if missing
                             if (not metadata.modelDescription or metadata.modelDescription == "") and 'model' in version_info:
                                 if 'description' in version_info['model']:
                                     metadata.modelDescription = version_info['model']['description']
-                        
+                                    needs_metadata_update = True
-                        # Save the updated metadata
-                        await save_metadata(file_path, metadata)
-                        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}")
                 
+                # Check if base_model is consistent with civitai baseModel
+                if metadata.civitai and 'baseModel' in metadata.civitai:
+                    civitai_base_model = determine_base_model(metadata.civitai['baseModel'])
+                    if metadata.base_model != civitai_base_model:
+                        logger.debug(f"Updating base_model from {metadata.base_model} to {civitai_base_model} for {file_path}")
+                        metadata.base_model = civitai_base_model
+                        needs_metadata_update = True
+            
             if metadata is None:
                 metadata = await self._get_file_info(file_path)
+                needs_metadata_update = True
         
-        model_data = metadata.to_dict()
+        # Continue processing
+        model_data = metadata.to_dict() if metadata else None
         
         # Skip excluded models
-        if model_data.get('exclude', False):
+        if model_data and model_data.get('exclude', False):
             self._excluded_models.append(model_data['file_path'])
             return None
             
+        # Fetch missing metadata from Civitai if needed (with write prevention)
+        with prevent_metadata_writes():
             await self._fetch_missing_metadata(file_path, model_data)
+            
         rel_path = os.path.relpath(file_path, root_path)
         folder = os.path.dirname(rel_path)
         model_data['folder'] = folder.replace(os.path.sep, '/')
         
+        # Only save metadata if needed
+        if needs_metadata_update and metadata:
+            await original_save_metadata(file_path, metadata)
+        
         return model_data
 
     async def _fetch_missing_metadata(self, file_path: str, model_data: Dict) -> None:
@@ -651,9 +697,9 @@ class ModelScanner:
 
                     model_data['civitai']['creator'] = model_metadata['creator']
                     
-                    metadata_path = os.path.splitext(file_path)[0] + '.metadata.json'
+                    # Create a metadata object and save it using save_metadata
-                    with open(metadata_path, 'w', encoding='utf-8') as f:
+                    metadata_obj = self.model_class.from_dict(model_data)
-                        json.dump(model_data, f, indent=2, ensure_ascii=False)
+                    await save_metadata(file_path, metadata_obj)
         except Exception as e:
             logger.error(f"Failed to update metadata from Civitai for {file_path}: {e}")