From 4ee32f02c5ee896c0d8f7548bd8a2f98ecf63fb4 Mon Sep 17 00:00:00 2001 From: Will Miao <13051207myq@gmail.com> Date: Fri, 21 Mar 2025 11:11:09 +0800 Subject: [PATCH] Add functionality to save recipes from the LoRAs widget - Introduced a new API endpoint to save recipes directly from the LoRAs widget. - Implemented logic to handle recipe data, including image processing and metadata extraction. - Enhanced error handling for missing fields and image retrieval. - Updated the ExifUtils to extract generation parameters from images for recipe creation. - Added a direct save option in the widget, improving user experience. --- py/config.py | 1 + py/routes/recipe_routes.py | 168 ++++++++++++++++++++++++++++++++++++ py/services/lora_scanner.py | 16 ++++ py/utils/exif_utils.py | 150 +++++++++++++++++++++++++++++++- py/utils/recipe_parsers.py | 4 - web/comfyui/loras_widget.js | 106 +++++++++++++++++++++-- 6 files changed, 433 insertions(+), 12 deletions(-) diff --git a/py/config.py b/py/config.py index 516a3921..bae54b63 100644 --- a/py/config.py +++ b/py/config.py @@ -17,6 +17,7 @@ class Config: # 静态路由映射字典, target to route mapping self._route_mappings = {} self.loras_roots = self._init_lora_paths() + self.temp_directory = folder_paths.get_temp_directory() # 在初始化时扫描符号链接 self._scan_symbolic_links() diff --git a/py/routes/recipe_routes.py b/py/routes/recipe_routes.py index 6e2e49fe..2a6d9ec6 100644 --- a/py/routes/recipe_routes.py +++ b/py/routes/recipe_routes.py @@ -46,6 +46,8 @@ class RecipeRoutes: # Start cache initialization app.on_startup.append(routes._init_cache) + + app.router.add_post('/api/recipes/save-from-widget', routes.save_recipe_from_widget) async def _init_cache(self, app): """Initialize cache on startup""" @@ -730,3 +732,169 @@ class RecipeRoutes: del self._shared_recipes[rid] except Exception as e: logger.error(f"Error cleaning up shared recipe {rid}: {e}") + + async def save_recipe_from_widget(self, request: web.Request) -> web.Response: + """Save a recipe from the LoRAs widget""" + try: + reader = await request.multipart() + + # Process form data + name = None + tags = [] + metadata = None + + while True: + field = await reader.next() + if field is None: + break + + if field.name == 'name': + name = await field.text() + + elif field.name == 'tags': + tags_text = await field.text() + try: + tags = json.loads(tags_text) + except: + tags = [] + + elif field.name == 'metadata': + metadata_text = await field.text() + try: + metadata = json.loads(metadata_text) + except: + metadata = {} + + missing_fields = [] + if not name: + missing_fields.append("name") + if not metadata: + missing_fields.append("metadata") + if missing_fields: + return web.json_response({"error": f"Missing required fields: {', '.join(missing_fields)}"}, status=400) + + # Find the latest image in the temp directory + temp_dir = config.temp_directory + image_files = [] + + for file in os.listdir(temp_dir): + if file.lower().endswith(('.png', '.jpg', '.jpeg', '.webp')): + file_path = os.path.join(temp_dir, file) + image_files.append((file_path, os.path.getmtime(file_path))) + + if not image_files: + return web.json_response({"error": "No recent images found to use for recipe"}, status=400) + + # Sort by modification time (newest first) + image_files.sort(key=lambda x: x[1], reverse=True) + latest_image_path = image_files[0][0] + + # Extract ComfyUI generation parameters from the latest image + gen_params = ExifUtils.extract_comfyui_gen_params(latest_image_path) + + # Read the image + with open(latest_image_path, 'rb') as f: + image = f.read() + + # Create recipes directory if it doesn't exist + recipes_dir = self.recipe_scanner.recipes_dir + os.makedirs(recipes_dir, exist_ok=True) + + # Generate UUID for the recipe + import uuid + recipe_id = str(uuid.uuid4()) + + # Optimize the image (resize and convert to WebP) + optimized_image, extension = ExifUtils.optimize_image( + image_data=image, + target_width=480, + format='webp', + quality=85, + preserve_metadata=True + ) + + # Save the optimized image + image_filename = f"{recipe_id}{extension}" + image_path = os.path.join(recipes_dir, image_filename) + with open(image_path, 'wb') as f: + f.write(optimized_image) + + # Format loras data from metadata + loras_data = [] + for lora in metadata.get("loras", []): + # Skip inactive LoRAs + if not lora.get("active", True): + continue + + # Get lora info from scanner + lora_name = lora.get("name", "") + lora_info = await self.recipe_scanner._lora_scanner.get_lora_info_by_name(lora_name) + + # Create lora entry + lora_entry = { + "file_name": lora_name, + "hash": lora_info.get("sha256", "").lower() if lora_info else "", + "strength": float(lora.get("weight", 1.0)), + "modelVersionId": lora_info.get("civitai", {}).get("id", "") if lora_info else "", + "modelName": lora_info.get("civitai", {}).get("model", {}).get("name", "") if lora_info else lora_name, + "modelVersionName": lora_info.get("civitai", {}).get("name", "") if lora_info else "", + "isDeleted": False + } + loras_data.append(lora_entry) + + # Get base model from lora scanner + base_model_counts = {} + for lora in loras_data: + lora_info = await self.recipe_scanner._lora_scanner.get_lora_info_by_name(lora.get("file_name", "")) + if lora_info and "base_model" in lora_info: + base_model = lora_info["base_model"] + base_model_counts[base_model] = base_model_counts.get(base_model, 0) + 1 + + # Get most common base model + most_common_base_model = "" + if base_model_counts: + most_common_base_model = max(base_model_counts.items(), key=lambda x: x[1])[0] + + # Create the recipe data structure + recipe_data = { + "id": recipe_id, + "file_path": image_path, + "title": name, + "modified": time.time(), + "created_date": time.time(), + "base_model": most_common_base_model, + "loras": loras_data, + "gen_params": gen_params # Directly use the extracted params + } + + # Add tags if provided + if tags: + recipe_data["tags"] = tags + + # Save the recipe JSON + json_filename = f"{recipe_id}.recipe.json" + json_path = os.path.join(recipes_dir, json_filename) + with open(json_path, 'w', encoding='utf-8') as f: + json.dump(recipe_data, f, indent=4, ensure_ascii=False) + + # Add recipe metadata to the image + ExifUtils.append_recipe_metadata(image_path, recipe_data) + + # Update cache + if self.recipe_scanner._cache is not None: + # Add the recipe to the raw data if the cache exists + self.recipe_scanner._cache.raw_data.append(recipe_data) + # Schedule a background task to resort the cache + asyncio.create_task(self.recipe_scanner._cache.resort()) + logger.info(f"Added recipe {recipe_id} to cache") + + return web.json_response({ + 'success': True, + 'recipe_id': recipe_id, + 'image_path': image_path, + 'json_path': json_path + }) + + except Exception as e: + logger.error(f"Error saving recipe from widget: {e}", exc_info=True) + return web.json_response({"error": str(e)}, status=500) diff --git a/py/services/lora_scanner.py b/py/services/lora_scanner.py index 8bc7cd1b..bf58029c 100644 --- a/py/services/lora_scanner.py +++ b/py/services/lora_scanner.py @@ -721,3 +721,19 @@ class LoraScanner: test_hash_result = self._hash_index.get_hash(test_path) print(f"Test reverse lookup: {test_path} -> {test_hash_result[:8]}...\n\n", file=sys.stderr) + async def get_lora_info_by_name(self, name): + """Get LoRA information by name""" + try: + # Get cached data + cache = await self.get_cached_data() + + # Find the LoRA by name + for lora in cache.raw_data: + if lora.get("file_name") == name: + return lora + + return None + except Exception as e: + logger.error(f"Error getting LoRA info by name: {e}", exc_info=True) + return None + diff --git a/py/utils/exif_utils.py b/py/utils/exif_utils.py index 45831ec1..6c11333e 100644 --- a/py/utils/exif_utils.py +++ b/py/utils/exif_utils.py @@ -278,4 +278,152 @@ class ExifUtils: if isinstance(image_data, str) and os.path.exists(image_data): with open(image_data, 'rb') as f: return f.read(), os.path.splitext(image_data)[1] - return image_data, '.jpg' \ No newline at end of file + return image_data, '.jpg' + + @staticmethod + def _parse_comfyui_workflow(workflow_data: Any) -> Dict[str, Any]: + """ + Parse ComfyUI workflow data and extract relevant generation parameters + + Args: + workflow_data: Raw workflow data (string or dict) + + Returns: + Formatted generation parameters dictionary + """ + try: + # If workflow_data is a string, try to parse it as JSON + if isinstance(workflow_data, str): + try: + workflow_data = json.loads(workflow_data) + except json.JSONDecodeError: + logger.error("Failed to parse workflow data as JSON") + return {} + + # Now workflow_data should be a dictionary + if not isinstance(workflow_data, dict): + logger.error(f"Workflow data is not a dictionary: {type(workflow_data)}") + return {} + + # Initialize parameters dictionary with only the required fields + gen_params = { + "prompt": "", + "negative_prompt": "", + "steps": "", + "sampler": "", + "cfg_scale": "", + "seed": "", + "size": "", + "clip_skip": "" + } + + # Process each node in the workflow to extract parameters + for node_id, node_data in workflow_data.items(): + if not isinstance(node_data, dict): + continue + + # Extract node inputs if available + inputs = node_data.get("inputs", {}) + if not inputs: + continue + + # KSampler nodes contain most generation parameters + if "KSampler" in node_data.get("class_type", ""): + # Extract basic sampling parameters + gen_params["steps"] = inputs.get("steps", "") + gen_params["cfg_scale"] = inputs.get("cfg", "") + gen_params["sampler"] = inputs.get("sampler_name", "") + gen_params["seed"] = inputs.get("seed", "") + if isinstance(gen_params["seed"], list) and len(gen_params["seed"]) > 1: + gen_params["seed"] = gen_params["seed"][1] # Use the actual value if it's a list + + # CLIP Text Encode nodes contain prompts + elif "CLIPTextEncode" in node_data.get("class_type", ""): + # Check for negative prompt nodes + title = node_data.get("_meta", {}).get("title", "").lower() + prompt_text = inputs.get("text", "") + + if "negative" in title: + gen_params["negative_prompt"] = prompt_text + elif prompt_text and not "negative" in title and gen_params["prompt"] == "": + gen_params["prompt"] = prompt_text + + # CLIPSetLastLayer contains clip_skip information + elif "CLIPSetLastLayer" in node_data.get("class_type", ""): + gen_params["clip_skip"] = inputs.get("stop_at_clip_layer", "") + if isinstance(gen_params["clip_skip"], int) and gen_params["clip_skip"] < 0: + # Convert negative layer index to positive clip skip value + gen_params["clip_skip"] = abs(gen_params["clip_skip"]) + + # Look for resolution information + elif "LatentImage" in node_data.get("class_type", "") or "Empty" in node_data.get("class_type", ""): + width = inputs.get("width", 0) + height = inputs.get("height", 0) + if width and height: + gen_params["size"] = f"{width}x{height}" + + # Some nodes have resolution as a string like "832x1216 (0.68)" + resolution = inputs.get("resolution", "") + if isinstance(resolution, str) and "x" in resolution: + gen_params["size"] = resolution.split(" ")[0] # Extract just the dimensions + + return gen_params + + except Exception as e: + logger.error(f"Error parsing ComfyUI workflow: {e}", exc_info=True) + return {} + + @staticmethod + def extract_comfyui_gen_params(image_path: str) -> Dict[str, Any]: + """ + Extract ComfyUI workflow data from PNG images and format for recipe data + Only extracts the specific generation parameters needed for recipes. + + Args: + image_path: Path to the ComfyUI-generated PNG image + + Returns: + Dictionary containing formatted generation parameters + """ + try: + # Check if the file exists and is accessible + if not os.path.exists(image_path): + logger.error(f"Image file not found: {image_path}") + return {} + + # Open the image to extract embedded workflow data + with Image.open(image_path) as img: + workflow_data = None + + # For PNG images, look for the ComfyUI workflow data in PNG chunks + if img.format == 'PNG': + # Check standard metadata fields that might contain workflow + if 'parameters' in img.info: + workflow_data = img.info['parameters'] + elif 'prompt' in img.info: + workflow_data = img.info['prompt'] + else: + # Look for other potential field names that might contain workflow data + for key in img.info: + if isinstance(key, str) and ('workflow' in key.lower() or 'comfy' in key.lower()): + workflow_data = img.info[key] + break + + # If no workflow data found in PNG chunks, try EXIF as fallback + if not workflow_data: + user_comment = ExifUtils.extract_user_comment(image_path) + if user_comment and '{' in user_comment and '}' in user_comment: + # Try to extract JSON part + json_start = user_comment.find('{') + json_end = user_comment.rfind('}') + 1 + workflow_data = user_comment[json_start:json_end] + + # Parse workflow data if found + if workflow_data: + return ExifUtils._parse_comfyui_workflow(workflow_data) + + return {} + + except Exception as e: + logger.error(f"Error extracting ComfyUI gen params from {image_path}: {e}", exc_info=True) + return {} \ No newline at end of file diff --git a/py/utils/recipe_parsers.py b/py/utils/recipe_parsers.py index 0b017220..581eda74 100644 --- a/py/utils/recipe_parsers.py +++ b/py/utils/recipe_parsers.py @@ -517,14 +517,10 @@ class RecipeParserFactory: Appropriate RecipeMetadataParser implementation """ if RecipeFormatParser().is_metadata_matching(user_comment): - logger.info("RecipeFormatParser") return RecipeFormatParser() elif StandardMetadataParser().is_metadata_matching(user_comment): - logger.info("StandardMetadataParser") return StandardMetadataParser() elif A1111MetadataParser().is_metadata_matching(user_comment): - logger.info("A1111MetadataParser") return A1111MetadataParser() else: - logger.info("No parser found for this image") return None \ No newline at end of file diff --git a/web/comfyui/loras_widget.js b/web/comfyui/loras_widget.js index 7d544c0b..cbd89516 100644 --- a/web/comfyui/loras_widget.js +++ b/web/comfyui/loras_widget.js @@ -1,4 +1,5 @@ import { api } from "../../scripts/api.js"; +import { app } from "../../scripts/app.js"; export function addLorasWidget(node, name, opts, callback) { // Create container for loras @@ -376,16 +377,16 @@ export function addLorasWidget(node, name, opts, callback) { } ); - // Save recipe option with bookmark icon (WIP) + // Save recipe option with bookmark icon const saveOption = createMenuItem( - 'Save Recipe (WIP)', + 'Save Recipe', '', - null + () => { + menu.remove(); + document.removeEventListener('click', closeMenu); + saveRecipeDirectly(widget); + } ); - Object.assign(saveOption.style, { - opacity: '0.6', - cursor: 'default', - }); // Add separator const separator = document.createElement('div'); @@ -764,3 +765,94 @@ export function addLorasWidget(node, name, opts, callback) { return { minWidth: 400, minHeight: 200, widget }; } + +// Function to directly save the recipe without dialog +async function saveRecipeDirectly(widget) { + try { + // Filter active loras + const activeLoras = widget.value.filter(lora => lora.active); + + if (activeLoras.length === 0) { + // Show toast notification for no active LoRAs + if (app && app.extensionManager && app.extensionManager.toast) { + app.extensionManager.toast.add({ + severity: 'warn', + summary: 'No Active LoRAs', + detail: 'Please activate at least one LoRA to save a recipe', + life: 3000 + }); + } + return; + } + + // Generate a name based on active LoRAs + const recipeName = activeLoras.map(lora => + `${lora.name.split('/').pop().split('\\').pop()}:${lora.strength}` + ).join(' '); + + // Show loading toast + if (app && app.extensionManager && app.extensionManager.toast) { + app.extensionManager.toast.add({ + severity: 'info', + summary: 'Saving Recipe', + detail: 'Please wait...', + life: 2000 + }); + } + + // Prepare the data + const formData = new FormData(); + formData.append('name', recipeName); + formData.append('tags', JSON.stringify([])); + + // Prepare metadata with loras + const metadata = { + loras: activeLoras.map(lora => ({ + name: lora.name, + weight: parseFloat(lora.strength), + active: true + })) + }; + + formData.append('metadata', JSON.stringify(metadata)); + + // Send the request + const response = await fetch('/api/recipes/save-from-widget', { + method: 'POST', + body: formData + }); + + const result = await response.json(); + + // Show result toast + if (app && app.extensionManager && app.extensionManager.toast) { + if (result.success) { + app.extensionManager.toast.add({ + severity: 'success', + summary: 'Recipe Saved', + detail: 'Recipe has been saved successfully', + life: 3000 + }); + } else { + app.extensionManager.toast.add({ + severity: 'error', + summary: 'Error', + detail: result.error || 'Failed to save recipe', + life: 5000 + }); + } + } + } catch (error) { + console.error('Error saving recipe:', error); + + // Show error toast + if (app && app.extensionManager && app.extensionManager.toast) { + app.extensionManager.toast.add({ + severity: 'error', + summary: 'Error', + detail: 'Failed to save recipe: ' + (error.message || 'Unknown error'), + life: 5000 + }); + } + } +}