mirror of
https://github.com/invoke-ai/InvokeAI
synced 2025-07-26 05:17:55 +00:00
Extract util and fix model image logic
This commit is contained in:
Kent Keirsey
committed by
psychedelicious
psychedelicious
parent
dd35ab026a
commit
b14d841d57
@ -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()
|
||||
|
||||
|
||||
145
invokeai/backend/model_manager/util/lora_metadata_extractor.py
Normal file
145
invokeai/backend/model_manager/util/lora_metadata_extractor.py
Normal file
@ -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
|
||||
Reference in New Issue
Block a user