From 8e8b80ddcf7b49c1cb265ca083052c95b6629da9 Mon Sep 17 00:00:00 2001 From: Will Miao <13051207myq@gmail.com> Date: Sun, 2 Feb 2025 23:30:06 +0800 Subject: [PATCH] Add preview URL update mechanism for LoRA scanner cache --- routes/api_routes.py | 3 +++ services/lora_scanner.py | 46 ++++++++++++++++++++++++++++++++++++++++ utils/file_utils.py | 39 +++++++++++++++++++++++++++------- 3 files changed, 80 insertions(+), 8 deletions(-) diff --git a/routes/api_routes.py b/routes/api_routes.py index df5dde84..372a27bc 100644 --- a/routes/api_routes.py +++ b/routes/api_routes.py @@ -84,6 +84,9 @@ class ApiRoutes: preview_path = await self._save_preview_file(model_path, preview_data, content_type) await self._update_preview_metadata(model_path, preview_path) + # Update preview URL in scanner cache + await self.scanner.update_preview_in_cache(model_path, preview_path) + return web.json_response({ "success": True, "preview_url": config.get_preview_static_url(preview_path) diff --git a/services/lora_scanner.py b/services/lora_scanner.py index 33fb84a9..55aa8177 100644 --- a/services/lora_scanner.py +++ b/services/lora_scanner.py @@ -19,6 +19,37 @@ class LoraCache: folders: List[str] timestamp: float + def update_preview_url(self, file_path: str, preview_url: str) -> bool: + """Update preview_url for a specific lora in all cached data + + Args: + file_path: The file path of the lora to update + preview_url: The new preview URL + + Returns: + bool: True if the update was successful, False if the lora wasn't found + """ + # Update in raw_data + for item in self.raw_data: + if item['file_path'] == file_path: + item['preview_url'] = preview_url + break + else: + return False # Lora not found + + # Update in sorted lists (references to the same dict objects) + for item in self.sorted_by_name: + if item['file_path'] == file_path: + item['preview_url'] = preview_url + break + + for item in self.sorted_by_date: + if item['file_path'] == file_path: + item['preview_url'] = preview_url + break + + return True + class LoraScanner: """Service for scanning and managing LoRA files""" @@ -134,3 +165,18 @@ class LoraScanner: lora_data['folder'] = folder.replace(os.path.sep, '/') return lora_data + + async def update_preview_in_cache(self, file_path: str, preview_url: str) -> bool: + """Update preview URL in cache for a specific lora + + Args: + file_path: The file path of the lora to update + preview_url: The new preview URL + + Returns: + bool: True if the update was successful, False if cache doesn't exist or lora wasn't found + """ + if self._cache is None: + return False + + return self._cache.update_preview_url(file_path, preview_url) diff --git a/utils/file_utils.py b/utils/file_utils.py index e2cc9380..b3c9af76 100644 --- a/utils/file_utils.py +++ b/utils/file_utils.py @@ -31,6 +31,10 @@ def _find_preview_file(base_name: str, dir_path: str) -> str: return full_pattern.replace(os.sep, "/") return "" +def normalize_path(path: str) -> str: + """Normalize file path to use forward slashes""" + return path.replace(os.sep, "/") if path else path + async def get_file_info(file_path: str) -> LoraMetadata: """Get basic file information as LoraMetadata object""" base_name = os.path.splitext(os.path.basename(file_path))[0] @@ -41,21 +45,25 @@ async def get_file_info(file_path: str) -> LoraMetadata: return LoraMetadata( file_name=base_name, model_name=base_name, - file_path=file_path, + file_path=normalize_path(file_path), size=os.path.getsize(file_path), modified=os.path.getmtime(file_path), sha256=await calculate_sha256(file_path), base_model="Unknown", # Will be updated later from_civitai=True, - preview_url=preview_url, + preview_url=normalize_path(preview_url), ) async def save_metadata(file_path: str, metadata: LoraMetadata) -> None: """Save metadata to .metadata.json file""" metadata_path = f"{os.path.splitext(file_path)[0]}.metadata.json" try: + metadata_dict = metadata.to_dict() + metadata_dict['file_path'] = normalize_path(metadata_dict['file_path']) + metadata_dict['preview_url'] = normalize_path(metadata_dict['preview_url']) + with open(metadata_path, 'w', encoding='utf-8') as f: - json.dump(metadata.to_dict(), f, indent=2, ensure_ascii=False) + json.dump(metadata_dict, f, indent=2, ensure_ascii=False) except Exception as e: print(f"Error saving metadata to {metadata_path}: {str(e)}") @@ -66,16 +74,31 @@ async def load_metadata(file_path: str) -> Optional[LoraMetadata]: if os.path.exists(metadata_path): with open(metadata_path, 'r', encoding='utf-8') as f: data = json.load(f) + + needs_update = False + + if data['file_path'] != normalize_path(data['file_path']): + data['file_path'] = normalize_path(data['file_path']) + needs_update = True + preview_url = data.get('preview_url', '') - # preview_path = os.path.join(loras_root, preview_url) if preview_url else '' 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) - data['preview_url'] = _find_preview_file(base_name, dir_path) - if data['preview_url']: - with open(metadata_path, 'w', encoding='utf-8') as f: - json.dump(data, f, indent=2, ensure_ascii=False) + 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 + elif preview_url != normalize_path(preview_url): + data['preview_url'] = normalize_path(preview_url) + needs_update = True + + if needs_update: + with open(metadata_path, 'w', encoding='utf-8') as f: + json.dump(data, f, indent=2, ensure_ascii=False) + return LoraMetadata.from_dict(data) + except Exception as e: print(f"Error loading metadata from {metadata_path}: {str(e)}") return None