Enhance Civitai model handling: add get_model_version method for detailed metadata retrieval, update routes to utilize new method, and improve URL handling in context menu for model re-linking.

py/services/civitai_client.py

@@ -224,6 +224,69 @@ class CivitaiClient:
         except Exception as e:
             logger.error(f"Error fetching model versions: {e}")
             return None
+            
+    async def get_model_version(self, model_id: str, version_id: str = "") -> Optional[Dict]:
+        """Get specific model version with additional metadata
+        
+        Args:
+            model_id: The Civitai model ID
+            version_id: Optional specific version ID to retrieve
+            
+        Returns:
+            Optional[Dict]: The model version data with additional fields or None if not found
+        """
+        try:
+            session = await self._ensure_fresh_session()
+            async with session.get(f"{self.base_url}/models/{model_id}") as response:
+                if response.status != 200:
+                    return None
+                    
+                data = await response.json()
+                model_versions = data.get('modelVersions', [])
+                
+                # Find matching version
+                matched_version = None
+                
+                if version_id:
+                    # If version_id provided, find exact match
+                    for version in model_versions:
+                        if str(version.get('id')) == str(version_id):
+                            matched_version = version
+                            break
+                else:
+                    # If no version_id then use the first version
+                    matched_version = model_versions[0] if model_versions else None
+                
+                # If no match found, return None
+                if not matched_version:
+                    return None
+                    
+                # Build result with modified fields
+                result = matched_version.copy()  # Copy to avoid modifying original
+                
+                # Replace index with modelId
+                if 'index' in result:
+                    del result['index']
+                result['modelId'] = model_id
+                
+                # Add model field with metadata from top level
+                result['model'] = {
+                    "name": data.get("name"),
+                    "type": data.get("type"),
+                    "nsfw": data.get("nsfw", False),
+                    "poi": data.get("poi", False),
+                    "description": data.get("description"),
+                    "tags": data.get("tags", [])
+                }
+                
+                # Add creator field from top level
+                result['creator'] = data.get("creator")
+                
+                return result
+                
+        except Exception as e:
+            logger.error(f"Error fetching model version: {e}")
+            return None
 
     async def get_model_version_info(self, version_id: str) -> Tuple[Optional[Dict], Optional[str]]:
         """Fetch model version metadata from Civitai

py/utils/routes_common.py

@@ -47,13 +47,30 @@ class ModelRouteUtils:
             if civitai_metadata.get('model', {}).get('name'):
                 local_metadata['model_name'] = civitai_metadata['model']['name']
         
-            # Fetch additional model metadata (description and tags) if we have model ID
-            model_id = civitai_metadata['modelId']
-            if model_id:
-                model_metadata, _ = await client.get_model_metadata(str(model_id))
-                if (model_metadata):
-                    local_metadata['modelDescription'] = model_metadata.get('description', '')
-                    local_metadata['tags'] = model_metadata.get('tags', [])
+            # Extract model metadata directly from civitai_metadata if available
+            model_metadata = None
+            
+            if 'model' in civitai_metadata and civitai_metadata.get('model'):
+                # Data is already available in the response from get_model_version
+                model_metadata = {
+                    'description': civitai_metadata.get('model', {}).get('description', ''),
+                    'tags': civitai_metadata.get('model', {}).get('tags', []),
+                    'creator': civitai_metadata.get('creator', {})
+                }
+            
+            # If we have modelId and don't have enough metadata, fetch additional data
+            if not model_metadata or not model_metadata.get('description'):
+                model_id = civitai_metadata.get('modelId')
+                if model_id:
+                    fetched_metadata, _ = await client.get_model_metadata(str(model_id))
+                    if fetched_metadata:
+                        model_metadata = fetched_metadata
+            
+            # Update local metadata with the model information
+            if model_metadata:
+                local_metadata['modelDescription'] = model_metadata.get('description', '')
+                local_metadata['tags'] = model_metadata.get('tags', [])
+                if 'creator' in model_metadata and model_metadata['creator']:
                     local_metadata['civitai']['creator'] = model_metadata['creator']
         
         # Update base model
@@ -607,7 +624,7 @@ class ModelRouteUtils:
 
     @staticmethod
     async def handle_relink_civitai(request: web.Request, scanner) -> web.Response:
-        """Handle CivitAI metadata re-linking request by model version ID
+        """Handle CivitAI metadata re-linking request by model ID and/or version ID
         
         Args:
             request: The aiohttp request
@@ -619,10 +636,11 @@ class ModelRouteUtils:
         try:
             data = await request.json()
             file_path = data.get('file_path')
+            model_id = data.get('model_id')
             model_version_id = data.get('model_version_id')
             
-            if not file_path or not model_version_id:
-                return web.json_response({"success": False, "error": "Both file_path and model_version_id are required"}, status=400)
+            if not file_path or not model_id:
+                return web.json_response({"success": False, "error": "Both file_path and model_id are required"}, status=400)
             
             metadata_path = os.path.splitext(file_path)[0] + '.metadata.json'
             
@@ -632,24 +650,24 @@ class ModelRouteUtils:
             # Create a client for fetching from Civitai
             client = await CivitaiClient.get_instance()
             try:
-                # Fetch metadata by model version ID
-                civitai_metadata, error = await client.get_model_version_info(model_version_id)
+                # Fetch metadata using get_model_version which includes more comprehensive data
+                civitai_metadata = await client.get_model_version(model_id, model_version_id)
                 if not civitai_metadata:
-                    error_msg = error or "Model version not found on CivitAI"
+                    error_msg = f"Model version not found on CivitAI for ID: {model_id}"
+                    if model_version_id:
+                        error_msg += f" with version: {model_version_id}"
                     return web.json_response({"success": False, "error": error_msg}, status=404)
                 
-                # Find the primary model file to get the correct SHA256 hash
+                # Try to find the primary model file to get the SHA256 hash
                 primary_model_file = None
                 for file in civitai_metadata.get('files', []):
                     if file.get('primary', False) and file.get('type') == 'Model':
                         primary_model_file = file
                         break
                 
-                if not primary_model_file or not primary_model_file.get('hashes', {}).get('SHA256'):
-                    return web.json_response({"success": False, "error": "No SHA256 hash found in model metadata"}, status=404)
-                
-                # Update the SHA256 hash in local metadata (convert to lowercase)
-                local_metadata['sha256'] = primary_model_file['hashes']['SHA256'].lower()
+                # Update the SHA256 hash in local metadata if available
+                if primary_model_file and primary_model_file.get('hashes', {}).get('SHA256'):
+                    local_metadata['sha256'] = primary_model_file['hashes']['SHA256'].lower()
                 
                 # Update metadata with CivitAI information
                 await ModelRouteUtils.update_model_metadata(metadata_path, local_metadata, civitai_metadata, client)
@@ -659,8 +677,9 @@ class ModelRouteUtils:
                 
                 return web.json_response({
                     "success": True,
-                    "message": f"Model successfully re-linked to Civitai version {model_version_id}",
-                    "hash": local_metadata['sha256']
+                    "message": f"Model successfully re-linked to Civitai model {model_id}" + 
+                               (f" version {model_version_id}" if model_version_id else ""),
+                    "hash": local_metadata.get('sha256', '')
                 })
                 
             finally: