mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-03-23 22:22:11 -03:00
feat: enhance Civitai metadata handling and image URL processing
- Import rewrite_preview_url utility for optimized image URL handling - Update thumbnail URL processing for both LoRA and checkpoint entries to use rewritten URLs - Expand checkpoint metadata with modelId, file size, SHA256 hash, and file name - Improve error handling and data validation for Civitai API responses - Maintain backward compatibility with existing data structures
This commit is contained in:
Will Miao
@@ -8,6 +8,7 @@ from typing import Dict, List, Any, Optional, Tuple
|
||||
from abc import ABC, abstractmethod
|
||||
from ..config import config
|
||||
from ..utils.constants import VALID_LORA_TYPES
|
||||
from ..utils.civitai_utils import rewrite_preview_url
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -78,7 +79,7 @@ class RecipeMetadataParser(ABC):
|
||||
# Update model name if available
|
||||
if 'model' in civitai_info and 'name' in civitai_info['model']:
|
||||
lora_entry['name'] = civitai_info['model']['name']
|
||||
|
||||
|
||||
lora_entry['id'] = civitai_info.get('id')
|
||||
lora_entry['modelId'] = civitai_info.get('modelId')
|
||||
|
||||
@@ -88,7 +89,10 @@ class RecipeMetadataParser(ABC):
|
||||
|
||||
# Get thumbnail URL from first image
|
||||
if 'images' in civitai_info and civitai_info['images']:
|
||||
lora_entry['thumbnailUrl'] = civitai_info['images'][0].get('url', '')
|
||||
image_url = civitai_info['images'][0].get('url')
|
||||
if image_url:
|
||||
rewritten_image_url, _ = rewrite_preview_url(image_url, media_type='image')
|
||||
lora_entry['thumbnailUrl'] = rewritten_image_url or image_url
|
||||
|
||||
# Get base model
|
||||
current_base_model = civitai_info.get('baseModel', '')
|
||||
@@ -151,33 +155,59 @@ class RecipeMetadataParser(ABC):
|
||||
|
||||
Args:
|
||||
checkpoint: The checkpoint entry to populate
|
||||
civitai_info: The response from Civitai API
|
||||
civitai_info: The response from Civitai API or a (data, error_msg) tuple
|
||||
|
||||
Returns:
|
||||
The populated checkpoint dict
|
||||
"""
|
||||
try:
|
||||
if civitai_info and civitai_info.get("error") != "Model not found":
|
||||
# Update model name if available
|
||||
if 'model' in civitai_info and 'name' in civitai_info['model']:
|
||||
checkpoint['name'] = civitai_info['model']['name']
|
||||
|
||||
# Update version if available
|
||||
if 'name' in civitai_info:
|
||||
checkpoint['version'] = civitai_info.get('name', '')
|
||||
|
||||
# Get thumbnail URL from first image
|
||||
if 'images' in civitai_info and civitai_info['images']:
|
||||
checkpoint['thumbnailUrl'] = civitai_info['images'][0].get('url', '')
|
||||
|
||||
# Get base model
|
||||
checkpoint['baseModel'] = civitai_info.get('baseModel', '')
|
||||
|
||||
# Get download URL
|
||||
checkpoint['downloadUrl'] = civitai_info.get('downloadUrl', '')
|
||||
else:
|
||||
# Model not found or deleted
|
||||
civitai_data, error_msg = (
|
||||
(civitai_info, None)
|
||||
if not isinstance(civitai_info, tuple)
|
||||
else civitai_info
|
||||
)
|
||||
|
||||
if not civitai_data or error_msg == "Model not found":
|
||||
checkpoint['isDeleted'] = True
|
||||
return checkpoint
|
||||
|
||||
if 'model' in civitai_data and 'name' in civitai_data['model']:
|
||||
checkpoint['name'] = civitai_data['model']['name']
|
||||
|
||||
if 'name' in civitai_data:
|
||||
checkpoint['version'] = civitai_data.get('name', '')
|
||||
|
||||
if 'images' in civitai_data and civitai_data['images']:
|
||||
image_url = civitai_data['images'][0].get('url')
|
||||
if image_url:
|
||||
rewritten_image_url, _ = rewrite_preview_url(image_url, media_type='image')
|
||||
checkpoint['thumbnailUrl'] = rewritten_image_url or image_url
|
||||
|
||||
checkpoint['baseModel'] = civitai_data.get('baseModel', '')
|
||||
checkpoint['downloadUrl'] = civitai_data.get('downloadUrl', '')
|
||||
|
||||
checkpoint['modelId'] = civitai_data.get('modelId', checkpoint.get('modelId', 0))
|
||||
|
||||
if 'files' in civitai_data:
|
||||
model_file = next(
|
||||
(
|
||||
file
|
||||
for file in civitai_data.get('files', [])
|
||||
if file.get('type') == 'Model'
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
if model_file:
|
||||
checkpoint['size'] = model_file.get('sizeKB', 0) * 1024
|
||||
|
||||
sha256 = model_file.get('hashes', {}).get('SHA256')
|
||||
if sha256:
|
||||
checkpoint['hash'] = sha256.lower()
|
||||
|
||||
file_name = model_file.get('name', '')
|
||||
if file_name:
|
||||
checkpoint['file_name'] = os.path.splitext(file_name)[0]
|
||||
except Exception as e:
|
||||
logger.error(f"Error populating checkpoint from Civitai info: {e}")
|
||||
|
||||
|
||||
@@ -50,6 +50,7 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
|
||||
result = {
|
||||
'base_model': None,
|
||||
'loras': [],
|
||||
'model': None,
|
||||
'gen_params': {},
|
||||
'from_civitai_image': True
|
||||
}
|
||||
@@ -174,13 +175,48 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
|
||||
# Process civitaiResources array
|
||||
if "civitaiResources" in metadata and isinstance(metadata["civitaiResources"], list):
|
||||
for resource in metadata["civitaiResources"]:
|
||||
# Get unique identifier for deduplication
|
||||
# Get resource type and identifier
|
||||
resource_type = str(resource.get("type") or "").lower()
|
||||
version_id = str(resource.get("modelVersionId", ""))
|
||||
|
||||
|
||||
if resource_type == "checkpoint":
|
||||
checkpoint_entry = {
|
||||
'id': resource.get("modelVersionId", 0),
|
||||
'modelId': resource.get("modelId", 0),
|
||||
'name': resource.get("modelName", "Unknown Checkpoint"),
|
||||
'version': resource.get("modelVersionName", ""),
|
||||
'type': resource.get("type", "checkpoint"),
|
||||
'existsLocally': False,
|
||||
'localPath': None,
|
||||
'file_name': resource.get("modelName", ""),
|
||||
'hash': resource.get("hash", "") or "",
|
||||
'thumbnailUrl': '/loras_static/images/no-preview.png',
|
||||
'baseModel': '',
|
||||
'size': 0,
|
||||
'downloadUrl': '',
|
||||
'isDeleted': False
|
||||
}
|
||||
|
||||
if version_id and metadata_provider:
|
||||
try:
|
||||
civitai_info = await metadata_provider.get_model_version_info(version_id)
|
||||
|
||||
checkpoint_entry = await self.populate_checkpoint_from_civitai(
|
||||
checkpoint_entry,
|
||||
civitai_info
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(f"Error fetching Civitai info for checkpoint version {version_id}: {e}")
|
||||
|
||||
if result["model"] is None:
|
||||
result["model"] = checkpoint_entry
|
||||
|
||||
continue
|
||||
|
||||
# Skip if we've already added this LoRA
|
||||
if version_id and version_id in added_loras:
|
||||
continue
|
||||
|
||||
|
||||
# Initialize lora entry
|
||||
lora_entry = {
|
||||
'id': resource.get("modelVersionId", 0),
|
||||
@@ -196,31 +232,31 @@ class CivitaiApiMetadataParser(RecipeMetadataParser):
|
||||
'downloadUrl': '',
|
||||
'isDeleted': False
|
||||
}
|
||||
|
||||
|
||||
# Try to get info from Civitai if modelVersionId is available
|
||||
if version_id and metadata_provider:
|
||||
try:
|
||||
# Use get_model_version_info instead of get_model_version
|
||||
civitai_info = await metadata_provider.get_model_version_info(version_id)
|
||||
|
||||
|
||||
populated_entry = await self.populate_lora_from_civitai(
|
||||
lora_entry,
|
||||
civitai_info,
|
||||
recipe_scanner,
|
||||
base_model_counts
|
||||
)
|
||||
|
||||
|
||||
if populated_entry is None:
|
||||
continue # Skip invalid LoRA types
|
||||
|
||||
|
||||
lora_entry = populated_entry
|
||||
except Exception as e:
|
||||
logger.error(f"Error fetching Civitai info for model version {version_id}: {e}")
|
||||
|
||||
|
||||
# Track this LoRA in our deduplication dict
|
||||
if version_id:
|
||||
added_loras[version_id] = len(result["loras"])
|
||||
|
||||
|
||||
result["loras"].append(lora_entry)
|
||||
|
||||
# Process additionalResources array
|
||||
|
||||
Reference in New Issue
Block a user