From 4264dd19a84d61525c8cb72e9b16ee89031b6241 Mon Sep 17 00:00:00 2001 From: Will Miao <13051207myq@gmail.com> Date: Tue, 18 Mar 2025 16:49:04 +0800 Subject: [PATCH] Enhance recipe metadata handling in RecipeRoutes and ExifUtils - Added functionality to extract and process existing recipe metadata from images, including LoRA details and Civitai information. - Updated ExifUtils to manage recipe metadata more effectively, including appending and removing metadata from user comments. - Improved the ImportManager to utilize recipe metadata for setting default recipe names and tags when importing shared recipes. --- py/routes/recipe_routes.py | 92 ++++++++++++++++++++++++++--- py/utils/exif_utils.py | 85 ++++++++++++++++++-------- static/js/managers/ImportManager.js | 24 +++++++- 3 files changed, 167 insertions(+), 34 deletions(-) diff --git a/py/routes/recipe_routes.py b/py/routes/recipe_routes.py index 42c8b53d..f67131f2 100644 --- a/py/routes/recipe_routes.py +++ b/py/routes/recipe_routes.py @@ -220,7 +220,85 @@ class RecipeRoutes: "loras": [] # Return empty loras array to prevent client-side errors }, status=200) # Return 200 instead of 400 to handle gracefully - # Parse the recipe metadata + # First, check if this image has recipe metadata from a previous share + recipe_metadata = ExifUtils.extract_recipe_metadata(user_comment) + if recipe_metadata: + logger.info("Found existing recipe metadata in image") + + # Process the recipe metadata + loras = [] + for lora in recipe_metadata.get('loras', []): + # Convert recipe lora format to frontend format + lora_entry = { + 'id': lora.get('modelVersionId', ''), + 'name': lora.get('modelName', ''), + 'version': lora.get('modelVersionName', ''), + 'type': 'lora', + 'weight': lora.get('strength', 1.0), + 'file_name': lora.get('file_name', ''), + 'hash': lora.get('hash', '') + } + + # Check if this LoRA exists locally by SHA256 hash + if lora.get('hash'): + exists_locally = self.recipe_scanner._lora_scanner.has_lora_hash(lora['hash']) + if exists_locally: + lora_entry['existsLocally'] = True + + lora_cache = await self.recipe_scanner._lora_scanner.get_cached_data() + lora_item = next((item for item in lora_cache.raw_data if item['sha256'] == lora['hash']), None) + if lora_item: + lora_entry['localPath'] = lora_item['file_path'] + lora_entry['file_name'] = lora_item['file_name'] + lora_entry['size'] = lora_item['size'] + lora_entry['thumbnailUrl'] = config.get_preview_static_url(lora_item['preview_url']) + + else: + lora_entry['existsLocally'] = False + lora_entry['localPath'] = None + + # Try to get additional info from Civitai if we have a model version ID + if lora.get('modelVersionId'): + try: + civitai_info = await self.civitai_client.get_model_version_info(lora['modelVersionId']) + if civitai_info and civitai_info.get("error") != "Model not found": + # Get thumbnail URL from first image + if 'images' in civitai_info and civitai_info['images']: + lora_entry['thumbnailUrl'] = civitai_info['images'][0].get('url', '') + + # Get base model + lora_entry['baseModel'] = civitai_info.get('baseModel', '') + + # Get download URL + lora_entry['downloadUrl'] = civitai_info.get('downloadUrl', '') + + # Get size from files if available + if 'files' in civitai_info: + model_file = next((file for file in civitai_info.get('files', []) + if file.get('type') == 'Model'), None) + if model_file: + lora_entry['size'] = model_file.get('sizeKB', 0) * 1024 + else: + lora_entry['isDeleted'] = True + lora_entry['thumbnailUrl'] = '/loras_static/images/no-preview.png' + except Exception as e: + logger.error(f"Error fetching Civitai info for LoRA: {e}") + lora_entry['thumbnailUrl'] = '/loras_static/images/no-preview.png' + + loras.append(lora_entry) + + logger.info(f"Found {len(loras)} loras in recipe metadata") + + return web.json_response({ + 'base_model': recipe_metadata.get('base_model', ''), + 'loras': loras, + 'gen_params': recipe_metadata.get('gen_params', {}), + 'tags': recipe_metadata.get('tags', []), + 'title': recipe_metadata.get('title', ''), + 'from_recipe_metadata': True + }) + + # If no recipe metadata, parse the standard metadata metadata = ExifUtils.parse_recipe_metadata(user_comment) # Look for Civitai resources in the metadata @@ -233,10 +311,6 @@ class RecipeRoutes: "loras": [] # Return empty loras array }, status=200) # Return 200 instead of 400 - # Process the resources to get LoRA information - loras = [] - base_model = None - # Process LoRAs and collect base models base_model_counts = {} loras = [] @@ -248,9 +322,6 @@ class RecipeRoutes: if not model_version_id: continue - # Get additional info from Civitai - civitai_info = await self.civitai_client.get_model_version_info(model_version_id) - # Initialize lora entry with default values lora_entry = { 'id': model_version_id, @@ -269,6 +340,9 @@ class RecipeRoutes: 'isDeleted': False # New flag to indicate if the LoRA is deleted from Civitai } + # Get additional info from Civitai + civitai_info = await self.civitai_client.get_model_version_info(model_version_id) + # Check if this LoRA exists locally by SHA256 hash if civitai_info and civitai_info.get("error") != "Model not found": # LoRA exists on Civitai, process its information @@ -314,6 +388,7 @@ class RecipeRoutes: loras.append(lora_entry) # Set base_model to the most common one from civitai_info + base_model = None if base_model_counts: base_model = max(base_model_counts.items(), key=lambda x: x[1])[0] @@ -641,7 +716,6 @@ class RecipeRoutes: # Create a URL for the processed image # Use a timestamp to prevent caching timestamp = int(time.time()) - filename = os.path.basename(processed_path) url_path = f"/api/recipe/{recipe_id}/share/download?t={timestamp}" # Store the temp path in a dictionary to serve later diff --git a/py/utils/exif_utils.py b/py/utils/exif_utils.py index 60451b00..04a676ac 100644 --- a/py/utils/exif_utils.py +++ b/py/utils/exif_utils.py @@ -31,7 +31,7 @@ class ExifUtils: return None @staticmethod - def update_user_comment(image_path: str, user_comment: str) -> bool: + def update_user_comment(image_path: str, user_comment: str) -> str: """Update UserComment field in image EXIF data""" try: # Load the image and its EXIF data @@ -54,10 +54,10 @@ class ExifUtils: # Save the image with updated EXIF data img.save(image_path, exif=exif_bytes) - return True + return image_path except Exception as e: logger.error(f"Error updating EXIF data in {image_path}: {e}") - return False + return image_path @staticmethod def parse_recipe_metadata(user_comment: str) -> Dict[str, Any]: @@ -131,34 +131,71 @@ class ExifUtils: return None @staticmethod - def append_recipe_metadata(image_path: str, recipe_data: Dict) -> str: - """Append recipe metadata to image EXIF data and return the path to the modified image""" + def append_recipe_metadata(image_path, recipe_data) -> str: + """Append recipe metadata to an image's EXIF data""" try: - # Extract existing user comment - existing_comment = ExifUtils.extract_user_comment(image_path) or "" + # First, extract existing user comment + user_comment = ExifUtils.extract_user_comment(image_path) - # Prepare recipe metadata to append + # Check if there's already recipe metadata in the user comment + if user_comment: + # Remove any existing recipe metadata + user_comment = ExifUtils.remove_recipe_metadata(user_comment) + + # Prepare simplified loras data + simplified_loras = [] + for lora in recipe_data.get("loras", []): + simplified_lora = { + "file_name": lora.get("file_name", ""), + "hash": lora.get("hash", "").lower() if lora.get("hash") else "", + "strength": float(lora.get("strength", 1.0)), + "modelVersionId": lora.get("modelVersionId", ""), + "modelName": lora.get("modelName", ""), + "modelVersionName": lora.get("modelVersionName", ""), + } + simplified_loras.append(simplified_lora) + + # Create recipe metadata JSON recipe_metadata = { - "title": recipe_data.get("title", ""), - "base_model": recipe_data.get("base_model", ""), - "loras": recipe_data.get("loras", []), - "gen_params": recipe_data.get("gen_params", {}), - "tags": recipe_data.get("tags", []) + 'title': recipe_data.get('title', ''), + 'base_model': recipe_data.get('base_model', ''), + 'loras': simplified_loras, + 'gen_params': recipe_data.get('gen_params', {}), + 'tags': recipe_data.get('tags', []) } # Convert to JSON string - recipe_json = json.dumps(recipe_metadata, ensure_ascii=False) + recipe_metadata_json = json.dumps(recipe_metadata) - # Append to existing comment - if existing_comment and not existing_comment.endswith("\n"): - existing_comment += "\n" + # Create the recipe metadata marker + recipe_metadata_marker = f"Recipe metadata: {recipe_metadata_json}" - new_comment = existing_comment + "Recipe metadata: " + recipe_json + # Append to existing user comment or create new one + new_user_comment = user_comment + "\n" + recipe_metadata_marker if user_comment else recipe_metadata_marker - # Update the image with new comment - ExifUtils.update_user_comment(image_path, new_comment) - - return image_path + # Write back to the image + return ExifUtils.update_user_comment(image_path, new_user_comment) except Exception as e: - logger.error(f"Error appending recipe metadata: {e}") - return image_path # Return original path on error \ No newline at end of file + logger.error(f"Error appending recipe metadata: {e}", exc_info=True) + return image_path + + @staticmethod + def remove_recipe_metadata(user_comment): + """Remove recipe metadata from user comment""" + if not user_comment: + return "" + + # Find the recipe metadata marker + recipe_marker_index = user_comment.find("Recipe metadata: ") + if recipe_marker_index == -1: + return user_comment + + # Remove the recipe metadata part + # First, find where the metadata ends (next line or end of string) + next_line_index = user_comment.find("\n", recipe_marker_index) + if next_line_index == -1: + # Metadata is at the end of the string + return user_comment[:recipe_marker_index].rstrip() + else: + # Metadata is in the middle of the string + return user_comment[:recipe_marker_index] + user_comment[next_line_index:] \ No newline at end of file diff --git a/static/js/managers/ImportManager.js b/static/js/managers/ImportManager.js index 90d85e4a..f55f81c6 100644 --- a/static/js/managers/ImportManager.js +++ b/static/js/managers/ImportManager.js @@ -211,7 +211,21 @@ export class ImportManager { // Set default recipe name from prompt or image filename const recipeName = document.getElementById('recipeName'); - if (this.recipeData && this.recipeData.gen_params && this.recipeData.gen_params.prompt) { + + // Check if we have recipe metadata from a shared recipe + if (this.recipeData && this.recipeData.from_recipe_metadata) { + // Use title from recipe metadata + if (this.recipeData.title) { + recipeName.value = this.recipeData.title; + this.recipeName = this.recipeData.title; + } + + // Use tags from recipe metadata + if (this.recipeData.tags && Array.isArray(this.recipeData.tags)) { + this.recipeTags = [...this.recipeData.tags]; + this.updateTagsDisplay(); + } + } else if (this.recipeData && this.recipeData.gen_params && this.recipeData.gen_params.prompt) { // Use the first 15 words from the prompt as the default recipe name const promptWords = this.recipeData.gen_params.prompt.split(' '); const truncatedPrompt = promptWords.slice(0, 15).join(' '); @@ -232,6 +246,14 @@ export class ImportManager { this.recipeName = fileName; } + // Always set up click handler for easy editing if not already set + if (!recipeName.hasSelectAllHandler) { + recipeName.addEventListener('click', function() { + this.select(); + }); + recipeName.hasSelectAllHandler = true; + } + // Display the uploaded image in the preview const imagePreview = document.getElementById('recipeImagePreview'); if (imagePreview && this.recipeImage) {