diff --git a/invokeai/app/services/model_install/model_install_default.py b/invokeai/app/services/model_install/model_install_default.py index c6ce29f9bd..14232db45a 100644 --- a/invokeai/app/services/model_install/model_install_default.py +++ b/invokeai/app/services/model_install/model_install_default.py @@ -52,6 +52,7 @@ from invokeai.backend.model_manager.metadata import ( from invokeai.backend.model_manager.metadata.metadata_base import HuggingFaceMetadata from invokeai.backend.model_manager.search import ModelSearch from invokeai.backend.model_manager.taxonomy import ModelRepoVariant, ModelSourceType +from invokeai.backend.model_manager.util.lora_metadata_extractor import apply_lora_metadata from invokeai.backend.util import InvokeAILogger from invokeai.backend.util.catch_sigint import catch_sigint from invokeai.backend.util.devices import TorchDevice @@ -661,82 +662,6 @@ class ModelInstallService(ModelInstallServiceBase): except InvalidModelConfigException: return ModelConfigBase.classify(model_path, hash_algo, **fields) - def _check_for_lora_metadata(self, model_path: Path, info: "AnyModelConfig") -> None: - """ - Check for image files (PNG, JPG, WebP) or JSON metadata files with the same name as the LoRA model. - If found, extract relevant metadata and update the model configuration. - """ - from invokeai.backend.model_manager.config import ModelType - - # Only process LoRA models - if info.type != ModelType.LoRA: - return - - # Get the base name without extension - model_stem = model_path.stem - model_dir = model_path.parent - - # Look for image files with same name (PNG, JPG, WebP) - image_extensions = ['.png', '.jpg', '.jpeg', '.webp'] - preview_image_path = None - - for ext in image_extensions: - image_path = model_dir / f"{model_stem}{ext}" - if image_path.exists(): - preview_image_path = image_path - self._logger.info(f"Found preview image {image_path.name} for LoRA model {model_path.name}") - break - - # Set the preview image if found - if preview_image_path: - # Store the relative path to the image file - if preview_image_path.is_relative_to(self.app_config.models_path): - relative_path = preview_image_path.relative_to(self.app_config.models_path) - info.cover_image = relative_path.as_posix() - else: - info.cover_image = preview_image_path.as_posix() - self._logger.info(f"Set cover_image to {info.cover_image} for LoRA model {model_path.name}") - - # Look for JSON file with same name - json_path = model_dir / f"{model_stem}.json" - - if json_path.exists(): - try: - with open(json_path, 'r', encoding='utf-8') as f: - metadata = json.load(f) - - # Check if the JSON has any of the expected keys - expected_keys = { - "description", "sd version", "activation text", - "preferred weight", "negative text", "notes" - } - - if any(key in metadata for key in expected_keys): - # Map description + notes to model description - description_parts = [] - if "description" in metadata and metadata["description"]: - description_parts.append(str(metadata["description"]).strip()) - if "notes" in metadata and metadata["notes"]: - description_parts.append(str(metadata["notes"]).strip()) - - if description_parts: - combined_description = " | ".join(description_parts) - info.description = combined_description - - # Map activation text to trigger phrases - if "activation text" in metadata and metadata["activation text"]: - activation_text = str(metadata["activation text"]).strip() - if activation_text: - # Split on commas and clean up each phrase - phrases = [phrase.strip() for phrase in activation_text.split(',')] - phrases = [phrase for phrase in phrases if phrase] # Remove empty strings - if phrases: - info.trigger_phrases = set(phrases) - - self._logger.info(f"Applied metadata from {json_path.name} to LoRA model {model_path.name}") - - except (json.JSONDecodeError, IOError, Exception) as e: - self._logger.warning(f"Failed to read metadata from {json_path}: {e}") def _register( self, model_path: Path, config: Optional[ModelRecordChanges] = None, info: Optional[AnyModelConfig] = None @@ -745,11 +670,9 @@ class ModelInstallService(ModelInstallServiceBase): info = info or self._probe(model_path, config) - # Store the original resolved path for metadata checking - original_path = model_path.resolve() - - # Check for LoRA metadata files before finalizing the model config - self._check_for_lora_metadata(original_path, info) + # Apply LoRA metadata if applicable + model_images_path = self.app_config.models_path / "model_images" + apply_lora_metadata(info, model_path.resolve(), model_images_path) model_path = model_path.resolve() diff --git a/invokeai/backend/model_manager/util/lora_metadata_extractor.py b/invokeai/backend/model_manager/util/lora_metadata_extractor.py new file mode 100644 index 0000000000..0f879a0e8b --- /dev/null +++ b/invokeai/backend/model_manager/util/lora_metadata_extractor.py @@ -0,0 +1,145 @@ +"""Utility functions for extracting metadata from LoRA model files.""" + +import json +from pathlib import Path +from typing import Dict, Any, Optional, Set, Tuple +import logging + +from PIL import Image + +from invokeai.backend.model_manager.config import AnyModelConfig, ModelType +from invokeai.app.util.thumbnails import make_thumbnail + +logger = logging.getLogger(__name__) + + +def extract_lora_metadata(model_path: Path, model_key: str, model_images_path: Path) -> Tuple[Optional[str], Optional[Set[str]]]: + """ + Extract metadata for a LoRA model from associated JSON and image files. + + Args: + model_path: Path to the LoRA model file + model_key: Unique key for the model + model_images_path: Path to the model images directory + + Returns: + Tuple of (description, trigger_phrases) + """ + model_stem = model_path.stem + model_dir = model_path.parent + + # Find and process preview image + _process_preview_image(model_stem, model_dir, model_key, model_images_path) + + # Extract metadata from JSON + description, trigger_phrases = _extract_json_metadata(model_stem, model_dir) + + return description, trigger_phrases + + +def _process_preview_image(model_stem: str, model_dir: Path, model_key: str, model_images_path: Path) -> bool: + """Find and process a preview image for the model, saving it to the model images store.""" + image_extensions = ['.png', '.jpg', '.jpeg', '.webp'] + + for ext in image_extensions: + image_path = model_dir / f"{model_stem}{ext}" + if image_path.exists(): + try: + # Open the image + with Image.open(image_path) as img: + # Create thumbnail and save to model images directory + thumbnail = make_thumbnail(img, 256) + thumbnail_path = model_images_path / f"{model_key}.webp" + thumbnail.save(thumbnail_path, format="webp") + + logger.info(f"Processed preview image {image_path.name} for model {model_key}") + return True + + except Exception as e: + logger.warning(f"Failed to process preview image {image_path.name}: {e}") + return False + + return False + + +def _extract_json_metadata(model_stem: str, model_dir: Path) -> Tuple[Optional[str], Optional[Set[str]]]: + """Extract metadata from a JSON file with the same name as the model.""" + json_path = model_dir / f"{model_stem}.json" + + if not json_path.exists(): + return None, None + + try: + with open(json_path, 'r', encoding='utf-8') as f: + metadata = json.load(f) + + # Extract description + description = _build_description(metadata) + + # Extract trigger phrases + trigger_phrases = _extract_trigger_phrases(metadata) + + if description or trigger_phrases: + logger.info(f"Applied metadata from {json_path.name}") + + return description, trigger_phrases + + except (json.JSONDecodeError, IOError, Exception) as e: + logger.warning(f"Failed to read metadata from {json_path}: {e}") + return None, None + + +def _build_description(metadata: Dict[str, Any]) -> Optional[str]: + """Build a description from metadata fields.""" + description_parts = [] + + if description := metadata.get("description"): + description_parts.append(str(description).strip()) + + if notes := metadata.get("notes"): + description_parts.append(str(notes).strip()) + + return " | ".join(description_parts) if description_parts else None + + +def _extract_trigger_phrases(metadata: Dict[str, Any]) -> Optional[Set[str]]: + """Extract trigger phrases from metadata.""" + if not (activation_text := metadata.get("activation text")): + return None + + activation_text = str(activation_text).strip() + if not activation_text: + return None + + # Split on commas and clean up each phrase + phrases = [phrase.strip() for phrase in activation_text.split(',') if phrase.strip()] + + return set(phrases) if phrases else None + + +def apply_lora_metadata(info: AnyModelConfig, model_path: Path, model_images_path: Path) -> None: + """ + Apply extracted metadata to a LoRA model configuration. + + Args: + info: The model configuration to update + model_path: Path to the LoRA model file + model_images_path: Path to the model images directory + """ + # Only process LoRA models + if info.type != ModelType.LoRA: + return + + # Extract and apply metadata + description, trigger_phrases = extract_lora_metadata( + model_path, info.key, model_images_path + ) + + # We don't set cover_image path in the config anymore since images are stored + # separately in the model images store by model key + + if description: + info.description = description + + if trigger_phrases: + info.trigger_phrases = trigger_phrases \ No newline at end of file