mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-03-25 07:05:43 -03:00
feat: improve civitai data handling and type safety, fixes #565
- Replace setdefault with get and explicit dict initialization in MetadataUpdater - Change civitai field type from Optional[Dict] to Dict[str, Any] with default_factory - Add None check and dict initialization in BaseModelMetadata.__post_init__ - Ensures civitai data is always a dictionary, preventing type errors and improving code reliability
This commit is contained in:
@@ -270,7 +270,12 @@ class MetadataUpdater:
|
|||||||
"""
|
"""
|
||||||
try:
|
try:
|
||||||
await MetadataManager.hydrate_model_data(model_data)
|
await MetadataManager.hydrate_model_data(model_data)
|
||||||
civitai_data = model_data.setdefault('civitai', {})
|
civitai_data = model_data.get('civitai')
|
||||||
|
|
||||||
|
if not isinstance(civitai_data, dict):
|
||||||
|
civitai_data = {}
|
||||||
|
model_data['civitai'] = civitai_data
|
||||||
|
|
||||||
custom_images = civitai_data.get('customImages')
|
custom_images = civitai_data.get('customImages')
|
||||||
|
|
||||||
if not isinstance(custom_images, list):
|
if not isinstance(custom_images, list):
|
||||||
@@ -446,4 +451,4 @@ class MetadataUpdater:
|
|||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"Error parsing image metadata: {e}", exc_info=True)
|
logger.error(f"Error parsing image metadata: {e}", exc_info=True)
|
||||||
return None
|
return None
|
||||||
|
|||||||
@@ -18,7 +18,7 @@ class BaseModelMetadata:
|
|||||||
preview_nsfw_level: int = 0 # NSFW level of the preview image
|
preview_nsfw_level: int = 0 # NSFW level of the preview image
|
||||||
notes: str = "" # Additional notes
|
notes: str = "" # Additional notes
|
||||||
from_civitai: bool = True # Whether from Civitai
|
from_civitai: bool = True # Whether from Civitai
|
||||||
civitai: Optional[Dict] = None # Civitai API data if available
|
civitai: Dict[str, Any] = field(default_factory=dict) # Civitai API data if available
|
||||||
tags: List[str] = None # Model tags
|
tags: List[str] = None # Model tags
|
||||||
modelDescription: str = "" # Full model description
|
modelDescription: str = "" # Full model description
|
||||||
civitai_deleted: bool = False # Whether deleted from Civitai
|
civitai_deleted: bool = False # Whether deleted from Civitai
|
||||||
@@ -31,6 +31,9 @@ class BaseModelMetadata:
|
|||||||
|
|
||||||
def __post_init__(self):
|
def __post_init__(self):
|
||||||
# Initialize empty lists to avoid mutable default parameter issue
|
# Initialize empty lists to avoid mutable default parameter issue
|
||||||
|
if self.civitai is None:
|
||||||
|
self.civitai = {}
|
||||||
|
|
||||||
if self.tags is None:
|
if self.tags is None:
|
||||||
self.tags = []
|
self.tags = []
|
||||||
|
|
||||||
|
|||||||
37
tests/utils/test_metadata_defaults.py
Normal file
37
tests/utils/test_metadata_defaults.py
Normal file
@@ -0,0 +1,37 @@
|
|||||||
|
import json
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from py.utils.metadata_manager import MetadataManager
|
||||||
|
from py.utils.models import BaseModelMetadata
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_base_model_metadata_sets_empty_civitai_dict():
|
||||||
|
metadata = BaseModelMetadata(
|
||||||
|
file_name="model",
|
||||||
|
model_name="Model",
|
||||||
|
file_path="/tmp/model.safetensors",
|
||||||
|
size=0,
|
||||||
|
modified=0.0,
|
||||||
|
sha256="deadbeef",
|
||||||
|
base_model="Unknown",
|
||||||
|
preview_url="",
|
||||||
|
)
|
||||||
|
|
||||||
|
assert metadata.civitai == {}
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_create_default_metadata_uses_empty_civitai(tmp_path):
|
||||||
|
model_path = tmp_path / "example.safetensors"
|
||||||
|
model_path.write_bytes(b"stub")
|
||||||
|
|
||||||
|
metadata = await MetadataManager.create_default_metadata(str(model_path))
|
||||||
|
|
||||||
|
assert metadata is not None
|
||||||
|
assert metadata.civitai == {}
|
||||||
|
|
||||||
|
metadata_path = model_path.with_suffix(".metadata.json")
|
||||||
|
payload = json.loads(metadata_path.read_text(encoding="utf-8"))
|
||||||
|
|
||||||
|
assert payload.get("civitai") == {}
|
||||||
Reference in New Issue
Block a user