Initial commit

utils/__init__.py

@@ -0,0 +1 @@
+# Empty file to mark directory as Python package 

utils/file_utils.py

@@ -0,0 +1,48 @@
+import os
+import hashlib
+import json
+from typing import Dict, Optional
+
+async def calculate_sha256(file_path: str) -> str:
+    """Calculate SHA256 hash of a file"""
+    sha256_hash = hashlib.sha256()
+    with open(file_path, "rb") as f:
+        for byte_block in iter(lambda: f.read(4096), b""):
+            sha256_hash.update(byte_block)
+    return sha256_hash.hexdigest()
+
+async def get_file_info(file_path: str) -> Dict:
+    """Get basic file information"""
+    return {
+        "name": os.path.splitext(os.path.basename(file_path))[0],
+        "file_path": file_path,
+        "size": os.path.getsize(file_path),
+        "modified": os.path.getmtime(file_path),
+        "sha256": await calculate_sha256(file_path)
+    }
+
+async def save_metadata(file_path: str, metadata: Dict) -> None:
+    """Save metadata to .metadata.json file"""
+    metadata_path = f"{os.path.splitext(file_path)[0]}.metadata.json"
+    try:
+        with open(metadata_path, 'w', encoding='utf-8') as f:
+            json.dump(metadata, f, indent=2, ensure_ascii=False)
+    except Exception as e:
+        print(f"Error saving metadata to {metadata_path}: {str(e)}")
+
+async def load_metadata(file_path: str) -> Dict:
+    """Load metadata from .metadata.json file"""
+    metadata_path = f"{os.path.splitext(file_path)[0]}.metadata.json"
+    try:
+        if os.path.exists(metadata_path):
+            with open(metadata_path, 'r', encoding='utf-8') as f:
+                return json.load(f)
+    except Exception as e:
+        print(f"Error loading metadata from {metadata_path}: {str(e)}")
+    return {}
+
+async def update_civitai_metadata(file_path: str, civitai_data: Dict) -> None:
+    """Update metadata file with Civitai data"""
+    metadata = await load_metadata(file_path)
+    metadata['civitai'] = civitai_data
+    await save_metadata(file_path, metadata) 

utils/lora_metadata.py

@@ -0,0 +1,16 @@
+from safetensors import safe_open
+from typing import Dict
+from .model_utils import determine_base_model
+
+async def extract_lora_metadata(file_path: str) -> Dict:
+    """Extract essential metadata from safetensors file"""
+    try:
+        with safe_open(file_path, framework="pt", device="cpu") as f:
+            metadata = f.metadata()
+            if metadata:
+                # Only extract base_model from ss_base_model_version
+                base_model = determine_base_model(metadata.get("ss_base_model_version"))
+                return {"base_model": base_model}
+    except Exception as e:
+        print(f"Error reading metadata from {file_path}: {str(e)}")
+    return {"base_model": "Unknown"} 

utils/model_utils.py

@@ -0,0 +1,22 @@
+from typing import Dict, Optional
+
+# Base model mapping based on version string
+BASE_MODEL_MAPPING = {
+    "sd-v1-5": "SD1.5",
+    "sd-v2-1": "SD2.1",
+    "sdxl": "SDXL",
+    "sd-v2": "SD2.0",
+    "flux1": "Flux1.D",
+}
+
+def determine_base_model(version_string: Optional[str]) -> str:
+    """Determine base model from version string in safetensors metadata"""
+    if not version_string:
+        return "Unknown"
+    
+    version_lower = version_string.lower()
+    for key, value in BASE_MODEL_MAPPING.items():
+        if key in version_lower:
+            return value
+    
+    return "Unknown"