checkpoint

py/utils/models.py

@@ -5,20 +5,19 @@ import os
 from .model_utils import determine_base_model
 
 @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
+class BaseModelMetadata:
+    """Base class for all model metadata structures"""
+    file_name: str              # The filename without extension
+    model_name: str             # The model's name defined by the creator
+    file_path: str              # Full path to the model 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.)
+    base_model: str             # Base model type (SD1.5/SD2.1/SDXL/etc.)
     preview_url: str            # Preview image URL
     preview_nsfw_level: int = 0 # NSFW level of the preview image
-    usage_tips: str = "{}"      # Usage tips for the model, json string
     notes: str = ""             # Additional notes
-    from_civitai: bool = True   # Whether the lora is from Civitai
+    from_civitai: bool = True   # Whether from Civitai
     civitai: Optional[Dict] = None  # Civitai API data if available
     tags: List[str] = None      # Model tags
     modelDescription: str = ""  # Full model description
@@ -29,32 +28,11 @@ class LoraMetadata:
             self.tags = []
 
     @classmethod
-    def from_dict(cls, data: Dict) -> 'LoraMetadata':
-        """Create LoraMetadata instance from dictionary"""
-        # Create a copy of the data to avoid modifying the input
+    def from_dict(cls, data: Dict) -> 'BaseModelMetadata':
+        """Create instance from dictionary"""
         data_copy = data.copy()
         return cls(**data_copy)
 
-    @classmethod
-    def from_civitai_info(cls, version_info: Dict, file_info: Dict, save_path: str) -> 'LoraMetadata':
-        """Create LoraMetadata instance from Civitai version info"""
-        file_name = file_info['name']
-        base_model = determine_base_model(version_info.get('baseModel', ''))
-        
-        return cls(
-            file_name=os.path.splitext(file_name)[0],
-            model_name=version_info.get('model').get('name', os.path.splitext(file_name)[0]),
-            file_path=save_path.replace(os.sep, '/'),
-            size=file_info.get('sizeKB', 0) * 1024,
-            modified=datetime.now().timestamp(),
-            sha256=file_info['hashes'].get('SHA256', '').lower(),
-            base_model=base_model,
-            preview_url=None,  # Will be updated after preview download
-            preview_nsfw_level=0, # Will be updated after preview download, it is decided by the nsfw level of the preview image
-            from_civitai=True,
-            civitai=version_info
-        )
-
     def to_dict(self) -> Dict:
         """Convert to dictionary for JSON serialization"""
         return asdict(self)
@@ -76,30 +54,54 @@ class LoraMetadata:
             self.file_path = file_path.replace(os.sep, '/')
 
 @dataclass
-class CheckpointMetadata:
-    """Represents the metadata structure for a Checkpoint model"""
-    file_name: str              # The filename without extension
-    model_name: str             # The checkpoint's name defined by the creator
-    file_path: str              # Full path to the model file
-    size: int                   # File size in bytes
-    modified: float             # Last modified timestamp
-    sha256: str                 # SHA256 hash of the file
-    base_model: str             # Base model type (SD1.5/SD2.1/SDXL/etc.)
-    preview_url: str            # Preview image URL
-    preview_nsfw_level: int = 0 # NSFW level of the preview image
-    model_type: str = "checkpoint"  # Model type (checkpoint, inpainting, etc.)
-    notes: str = ""             # Additional notes
-    from_civitai: bool = True   # Whether from Civitai
-    civitai: Optional[Dict] = None  # Civitai API data if available
-    tags: List[str] = None      # Model tags
-    modelDescription: str = ""  # Full model description
-    
-    # Additional checkpoint-specific fields
-    resolution: Optional[str] = None  # Native resolution (e.g., 512x512, 1024x1024)
-    vae_included: bool = False       # Whether VAE is included in the checkpoint
-    architecture: str = ""           # Model architecture (if known)
-    
-    def __post_init__(self):
-        if self.tags is None:
-            self.tags = []
+class LoraMetadata(BaseModelMetadata):
+    """Represents the metadata structure for a Lora model"""
+    usage_tips: str = "{}"      # Usage tips for the model, json string
+
+    @classmethod
+    def from_civitai_info(cls, version_info: Dict, file_info: Dict, save_path: str) -> 'LoraMetadata':
+        """Create LoraMetadata instance from Civitai version info"""
+        file_name = file_info['name']
+        base_model = determine_base_model(version_info.get('baseModel', ''))
+        
+        return cls(
+            file_name=os.path.splitext(file_name)[0],
+            model_name=version_info.get('model').get('name', os.path.splitext(file_name)[0]),
+            file_path=save_path.replace(os.sep, '/'),
+            size=file_info.get('sizeKB', 0) * 1024,
+            modified=datetime.now().timestamp(),
+            sha256=file_info['hashes'].get('SHA256', '').lower(),
+            base_model=base_model,
+            preview_url=None,  # Will be updated after preview download
+            preview_nsfw_level=0, # Will be updated after preview download
+            from_civitai=True,
+            civitai=version_info
+        )
+
+@dataclass
+class CheckpointMetadata(BaseModelMetadata):
+    """Represents the metadata structure for a Checkpoint model"""
+    model_type: str = "checkpoint"  # Model type (checkpoint, inpainting, etc.)
+
+    @classmethod
+    def from_civitai_info(cls, version_info: Dict, file_info: Dict, save_path: str) -> 'CheckpointMetadata':
+        """Create CheckpointMetadata instance from Civitai version info"""
+        file_name = file_info['name']
+        base_model = determine_base_model(version_info.get('baseModel', ''))
+        model_type = version_info.get('type', 'checkpoint')
+        
+        return cls(
+            file_name=os.path.splitext(file_name)[0],
+            model_name=version_info.get('model').get('name', os.path.splitext(file_name)[0]),
+            file_path=save_path.replace(os.sep, '/'),
+            size=file_info.get('sizeKB', 0) * 1024,
+            modified=datetime.now().timestamp(),
+            sha256=file_info['hashes'].get('SHA256', '').lower(),
+            base_model=base_model,
+            preview_url=None,  # Will be updated after preview download
+            preview_nsfw_level=0,
+            from_civitai=True,
+            civitai=version_info,
+            model_type=model_type
+        )