diff --git a/py/services/lora_scanner.py b/py/services/lora_scanner.py index 18cd45df..f322c9f2 100644 --- a/py/services/lora_scanner.py +++ b/py/services/lora_scanner.py @@ -5,10 +5,11 @@ import asyncio import shutil import time from typing import List, Dict, Optional -from dataclasses import dataclass -from operator import itemgetter + +from ..utils.models import LoraMetadata from ..config import config -from ..utils.file_utils import load_metadata, get_file_info +from ..utils.file_utils import load_metadata, get_file_info, normalize_path, find_preview_file, save_metadata +from ..utils.lora_metadata import extract_lora_metadata from .lora_cache import LoraCache from .lora_hash_index import LoraHashIndex from .settings_manager import settings @@ -332,8 +333,30 @@ class LoraScanner: metadata = await load_metadata(file_path) if metadata is None: - # Create new metadata if none exists - metadata = await get_file_info(file_path) + # Try to find and use .civitai.info file first + civitai_info_path = f"{os.path.splitext(file_path)[0]}.civitai.info" + if os.path.exists(civitai_info_path): + try: + with open(civitai_info_path, 'r', encoding='utf-8') as f: + version_info = json.load(f) + + file_info = next((f for f in version_info.get('files', []) if f.get('primary')), None) + if file_info: + # Create a minimal file_info with the required fields + file_name = os.path.splitext(os.path.basename(file_path))[0] + file_info['name'] = file_name + + # Use from_civitai_info to create metadata + metadata = LoraMetadata.from_civitai_info(version_info, file_info, file_path) + metadata.preview_url = find_preview_file(file_name, os.path.dirname(file_path)) + await save_metadata(file_path, metadata) + logger.debug(f"Created metadata from .civitai.info for {file_path}") + except Exception as e: + logger.error(f"Error creating metadata from .civitai.info for {file_path}: {e}") + + # If still no metadata, create new metadata using get_file_info + if metadata is None: + metadata = await get_file_info(file_path) # Convert to dict and add folder info lora_data = metadata.to_dict() @@ -344,7 +367,7 @@ class LoraScanner: lora_data['folder'] = folder.replace(os.path.sep, '/') return lora_data - + async def _fetch_missing_metadata(self, file_path: str, lora_data: Dict) -> None: """Fetch missing description and tags from Civitai if needed diff --git a/py/utils/file_utils.py b/py/utils/file_utils.py index 8aec9002..0f282051 100644 --- a/py/utils/file_utils.py +++ b/py/utils/file_utils.py @@ -19,7 +19,7 @@ async def calculate_sha256(file_path: str) -> str: sha256_hash.update(byte_block) return sha256_hash.hexdigest() -def _find_preview_file(base_name: str, dir_path: str) -> str: +def find_preview_file(base_name: str, dir_path: str) -> str: """Find preview file for given base name in directory""" preview_patterns = [ f"{base_name}.preview.png", @@ -56,16 +56,33 @@ async def get_file_info(file_path: str) -> Optional[LoraMetadata]: base_name = os.path.splitext(os.path.basename(file_path))[0] dir_path = os.path.dirname(file_path) - preview_url = _find_preview_file(base_name, dir_path) + preview_url = find_preview_file(base_name, dir_path) + + # Check if a .json file exists with SHA256 hash to avoid recalculation + json_path = f"{os.path.splitext(file_path)[0]}.json" + sha256 = None + if os.path.exists(json_path): + try: + with open(json_path, 'r', encoding='utf-8') as f: + json_data = json.load(f) + if 'sha256' in json_data: + sha256 = json_data['sha256'].lower() + logger.debug(f"Using SHA256 from .json file for {file_path}") + except Exception as e: + logger.error(f"Error reading .json file for {file_path}: {e}") try: + # If we didn't get SHA256 from the .json file, calculate it + if not sha256: + sha256 = await calculate_sha256(real_path) + metadata = LoraMetadata( file_name=base_name, model_name=base_name, file_path=normalize_path(file_path), size=os.path.getsize(real_path), modified=os.path.getmtime(real_path), - sha256=await calculate_sha256(real_path), + sha256=sha256, base_model="Unknown", # Will be updated later usage_tips="", notes="", @@ -125,7 +142,7 @@ async def load_metadata(file_path: str) -> Optional[LoraMetadata]: if not preview_url or not os.path.exists(preview_url): base_name = os.path.splitext(os.path.basename(file_path))[0] dir_path = os.path.dirname(file_path) - new_preview_url = normalize_path(_find_preview_file(base_name, dir_path)) + new_preview_url = normalize_path(find_preview_file(base_name, dir_path)) if new_preview_url != preview_url: data['preview_url'] = new_preview_url needs_update = True