From 1100363427cbf6e645551b9e62459ca15d8fc1fb Mon Sep 17 00:00:00 2001 From: Will Miao <13051207myq@gmail.com> Date: Sat, 25 Jan 2025 22:18:36 +0800 Subject: [PATCH] Update metadata structure --- utils/models.py | 36 ++++++++++++++++++++++++++++++++++++ 1 file changed, 36 insertions(+) create mode 100644 utils/models.py diff --git a/utils/models.py b/utils/models.py new file mode 100644 index 00000000..e5e43c7f --- /dev/null +++ b/utils/models.py @@ -0,0 +1,36 @@ +from dataclasses import dataclass, asdict +from typing import Dict, Optional +from datetime import datetime + +@dataclass +class LoraMetadata: + """Represents the metadata structure for a Lora model""" + file_name: str # The filename without extension of the lora + model_name: str # The lora's name defined by the creator, initially same as file_name + file_path: str # Full path to the safetensors file + size: int # File size in bytes + modified: float # Last modified timestamp + sha256: str # SHA256 hash of the file + base_model: str # Base model (SD1.5/SD2.1/SDXL/etc.) + preview_url: str # Preview image URL + civitai: Optional[Dict] = None # Civitai API data if available + + @classmethod + def from_dict(cls, data: Dict) -> 'LoraMetadata': + """Create LoraMetadata instance from dictionary""" + # Create a copy of the data to avoid modifying the input + data_copy = data.copy() + return cls(**data_copy) + + def to_dict(self) -> Dict: + """Convert to dictionary for JSON serialization""" + return asdict(self) + + @property + def modified_datetime(self) -> datetime: + """Convert modified timestamp to datetime object""" + return datetime.fromtimestamp(self.modified) + + def update_civitai_info(self, civitai_data: Dict) -> None: + """Update Civitai information""" + self.civitai = civitai_data \ No newline at end of file