Merge branch 'main' into dev

py/utils/constants.py

@@ -0,0 +1,8 @@
+NSFW_LEVELS = {
+    "PG": 1,
+    "PG13": 2,
+    "R": 4,
+    "X": 8,
+    "XXX": 16,
+    "Blocked": 32, # Probably not actually visible through the API without being logged in on model owner account?
+}

py/utils/file_utils.py

@@ -4,6 +4,8 @@ import hashlib
 import json
 from typing import Dict, Optional
 
+from .model_utils import determine_base_model
+
 from .lora_metadata import extract_lora_metadata
 from .models import LoraMetadata
 
@@ -69,6 +71,8 @@ async def get_file_info(file_path: str) -> Optional[LoraMetadata]:
             notes="",
             from_civitai=True,
             preview_url=normalize_path(preview_url),
+            tags=[],
+            modelDescription=""
         )
 
         # create metadata file
@@ -103,9 +107,18 @@ async def load_metadata(file_path: str) -> Optional[LoraMetadata]:
                 data = json.load(f)
                 
                 needs_update = False
+
+                # Check and normalize base model name
+                normalized_base_model = determine_base_model(data['base_model'])
+                if data['base_model'] != normalized_base_model:
+                    data['base_model'] = normalized_base_model
+                    needs_update = True
                 
-                if data['file_path'] != normalize_path(data['file_path']):
-                    data['file_path'] = normalize_path(data['file_path'])
+                # Compare paths without extensions
+                stored_path_base = os.path.splitext(data['file_path'])[0]
+                current_path_base = os.path.splitext(normalize_path(file_path))[0]
+                if stored_path_base != current_path_base:
+                    data['file_path'] = normalize_path(file_path)
                     needs_update = True
                 
                 preview_url = data.get('preview_url', '')
@@ -116,8 +129,21 @@ async def load_metadata(file_path: str) -> Optional[LoraMetadata]:
                     if new_preview_url != preview_url:
                         data['preview_url'] = new_preview_url
                         needs_update = True
-                elif preview_url != normalize_path(preview_url):
-                    data['preview_url'] = normalize_path(preview_url)
+                else:
+                    # Compare preview paths without extensions
+                    stored_preview_base = os.path.splitext(preview_url)[0]
+                    current_preview_base = os.path.splitext(normalize_path(preview_url))[0]
+                    if stored_preview_base != current_preview_base:
+                        data['preview_url'] = normalize_path(preview_url)
+                        needs_update = True
+
+                # Ensure all fields are present
+                if 'tags' not in data:
+                    data['tags'] = []
+                    needs_update = True
+                    
+                if 'modelDescription' not in data:
+                    data['modelDescription'] = ""
                     needs_update = True
                 
                 if needs_update:

py/utils/model_utils.py

@@ -8,7 +8,8 @@ BASE_MODEL_MAPPING = {
     "sd-v2": "SD 2.0",
     "flux1": "Flux.1 D",
     "flux.1 d": "Flux.1 D",
-    "illustrious": "IL",
+    "illustrious": "Illustrious",
+    "il": "Illustrious",
     "pony": "Pony",
     "Hunyuan Video": "Hunyuan Video"
 }

py/utils/models.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass, asdict
-from typing import Dict, Optional
+from typing import Dict, Optional, List
 from datetime import datetime
 import os
 from .model_utils import determine_base_model
@@ -15,10 +15,18 @@ class LoraMetadata:
     sha256: str                 # SHA256 hash of the file
     base_model: str             # Base model (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 the lora is from Civitai
     civitai: Optional[Dict] = None  # Civitai API data if available
+    tags: List[str] = None      # Model tags
+    modelDescription: str = ""  # Full model description
+
+    def __post_init__(self):
+        # Initialize empty lists to avoid mutable default parameter issue
+        if self.tags is None:
+            self.tags = []
 
     @classmethod
     def from_dict(cls, data: Dict) -> 'LoraMetadata':
@@ -42,6 +50,7 @@ class LoraMetadata:
             sha256=file_info['hashes'].get('SHA256', ''),
             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
         )