From 7713cecba6336dd930b04c4c3ef174e8f1a43836 Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Fri, 3 Apr 2026 08:24:45 -0700 Subject: [PATCH] fix: handle null civitai field, link all preview images (#52) - Guard against civitai: null in metadata (was crashing the scan) - Link all local preview images per LoRA, not just the first - Add get_civitai_image_urls() for future remote image support --- py/api/lora_integration.py | 25 ++++++++------- py/lora_utils.py | 63 ++++++++++++++++++++++++++++---------- 2 files changed, 58 insertions(+), 30 deletions(-) diff --git a/py/api/lora_integration.py b/py/api/lora_integration.py index 942012c..f5f910f 100644 --- a/py/api/lora_integration.py +++ b/py/api/lora_integration.py @@ -155,7 +155,7 @@ class LoraIntegrationMixin: from ..lora_utils import ( find_lora_directories, get_example_prompt_from_metadata, - get_preview_image_from_metadata, + get_preview_images_from_metadata, get_trigger_words_from_metadata, get_model_name_from_metadata, read_lora_metadata, @@ -225,8 +225,8 @@ class LoraIntegrationMixin: model_name = get_model_name_from_metadata(metadata) trigger_words = get_trigger_words_from_metadata(metadata) - preview_path = await self._run_in_executor( - get_preview_image_from_metadata, metadata, meta_file + preview_paths = await self._run_in_executor( + get_preview_images_from_metadata, metadata, meta_file ) # Build prompt text: prefer example prompt, then model name @@ -250,12 +250,12 @@ class LoraIntegrationMixin: ) if existing: - # Link preview image if we have one and it's not already linked - if preview_path: + # Link all preview images + for pp in preview_paths: await self._run_in_executor( self.db.link_image_to_prompt, existing["id"], - preview_path, + pp, ) skipped += 1 else: @@ -269,14 +269,13 @@ class LoraIntegrationMixin: prompt_hash, ) - if prompt_id and preview_path: - await self._run_in_executor( - self.db.link_image_to_prompt, - prompt_id, - preview_path, - ) - if prompt_id: + for pp in preview_paths: + await self._run_in_executor( + self.db.link_image_to_prompt, + prompt_id, + pp, + ) imported += 1 else: skipped += 1 diff --git a/py/lora_utils.py b/py/lora_utils.py index bdc2b4a..6bd4e80 100644 --- a/py/lora_utils.py +++ b/py/lora_utils.py @@ -197,11 +197,14 @@ def read_lora_metadata(metadata_path: Path) -> Optional[Dict]: return None +def _get_civitai(metadata: Dict) -> Dict: + """Safely get the civitai dict, handling None values.""" + return metadata.get("civitai") or {} + + def get_trigger_words_from_metadata(metadata: Dict) -> List[str]: """Extract trigger words from a parsed LoraManager metadata dict.""" - civitai = metadata.get("civitai", {}) - if not civitai: - return [] + civitai = _get_civitai(metadata) words = civitai.get("trainedWords", []) if isinstance(words, list): return [w.strip() for w in words if isinstance(w, str) and w.strip()] @@ -213,8 +216,8 @@ def get_example_prompt_from_metadata(metadata: Dict) -> Optional[str]: Looks at civitai.images[].meta.prompt for the first available example. """ - civitai = metadata.get("civitai", {}) - images = civitai.get("images", []) + civitai = _get_civitai(metadata) + images = civitai.get("images", []) or [] for img in images: if not isinstance(img, dict): continue @@ -226,43 +229,69 @@ def get_example_prompt_from_metadata(metadata: Dict) -> Optional[str]: return None +def get_civitai_image_urls(metadata: Dict) -> List[str]: + """Extract all civitai example image URLs from metadata.""" + civitai = _get_civitai(metadata) + urls = [] + for img in civitai.get("images", []) or []: + if not isinstance(img, dict): + continue + url = img.get("url", "") + if isinstance(url, str) and url.strip(): + urls.append(url.strip()) + return urls + + def get_model_name_from_metadata(metadata: Dict) -> str: """Extract the model display name from metadata.""" - # Try civitai model name first, then file_name, then fallback name = metadata.get("model_name", "") if not name: - civitai = metadata.get("civitai", {}) - model = civitai.get("model", {}) + civitai = _get_civitai(metadata) + model = civitai.get("model") or {} name = model.get("name", "") if not name: name = metadata.get("file_name", "unknown") return name -def get_preview_image_from_metadata( - metadata: Dict, metadata_path: Path -) -> Optional[str]: - """Find the preview image path for a LoRA from its metadata. +def get_preview_images_from_metadata(metadata: Dict, metadata_path: Path) -> List[str]: + """Find all local preview/example image paths for a LoRA. Returns: - Absolute path string to the preview image, or None. + List of absolute path strings to image files. """ + results = [] lora_dir = metadata_path.parent file_name = metadata.get("file_name", "") if not file_name: - # Derive from metadata filename: "foo.metadata.json" -> "foo" stem = metadata_path.name.replace(".metadata.json", "") file_name = stem base_name = Path(file_name).stem # Check standard preview naming conventions - for ext in (".png", ".jpg", ".jpeg", ".preview.png", ".preview.jpg"): + for ext in ( + ".png", + ".jpg", + ".jpeg", + ".webp", + ".preview.png", + ".preview.jpg", + ".preview.jpeg", + ): candidate = lora_dir / f"{base_name}{ext}" if candidate.exists(): - return str(candidate.resolve()) + results.append(str(candidate.resolve())) - return None + return results + + +def get_preview_image_from_metadata( + metadata: Dict, metadata_path: Path +) -> Optional[str]: + """Find the first preview image path for a LoRA (backward compat).""" + images = get_preview_images_from_metadata(metadata, metadata_path) + return images[0] if images else None def get_example_images_dir(lora_manager_path: str) -> Optional[str]: