From 8e653e2173d9fa6ba89062b9f3ee0eebbcc81b1a Mon Sep 17 00:00:00 2001 From: Will Miao <13051207myq@gmail.com> Date: Fri, 21 Mar 2025 17:28:20 +0800 Subject: [PATCH] Refactor recipe saving process to utilize workflow JSON and enhance Lora handling - Updated the recipe saving logic to accept a workflow JSON input instead of individual fields like name, tags, and metadata. - Implemented parsing of the workflow to extract generation parameters and Lora stack, improving the recipe creation process. - Enhanced error handling for missing workflow data and invalid Lora formats. - Removed deprecated code related to individual field handling, streamlining the recipe saving functionality. - Updated the front-end widget to send the workflow JSON directly, simplifying the data preparation process. --- py/routes/recipe_routes.py | 120 +++++++++++++++++++----------------- py/utils/exif_utils.py | 1 - web/comfyui/loras_widget.js | 38 +----------- 3 files changed, 68 insertions(+), 91 deletions(-) diff --git a/py/routes/recipe_routes.py b/py/routes/recipe_routes.py index 2a6d9ec6..9365cb96 100644 --- a/py/routes/recipe_routes.py +++ b/py/routes/recipe_routes.py @@ -739,39 +739,22 @@ class RecipeRoutes: reader = await request.multipart() # Process form data - name = None - tags = [] - metadata = None + workflow_json = 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() + if field.name == 'workflow_json': + workflow_text = await field.text() try: - tags = json.loads(tags_text) + workflow_json = json.loads(workflow_text) except: - tags = [] - - elif field.name == 'metadata': - metadata_text = await field.text() - try: - metadata = json.loads(metadata_text) - except: - metadata = {} + return web.json_response({"error": "Invalid workflow JSON"}, status=400) - 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) + if not workflow_json: + return web.json_response({"error": "Missing required workflow_json field"}, status=400) # Find the latest image in the temp directory temp_dir = config.temp_directory @@ -789,8 +772,38 @@ class RecipeRoutes: 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) + # Parse the workflow to extract generation parameters and loras + from ..workflow_params.workflow_parser import parse_workflow + # load_extensions=False to avoid loading extensions for now + parsed_workflow = parse_workflow(workflow_json, load_extensions=False) + + logger.info(f"Parsed workflow: {parsed_workflow}") + + if not parsed_workflow or not parsed_workflow.get("gen_params"): + return web.json_response({"error": "Could not extract generation parameters from workflow"}, status=400) + + # Get the lora stack from the parsed workflow + lora_stack = parsed_workflow.get("loras", "") + + # Parse the lora stack format: " ..." + import re + lora_matches = re.findall(r']+)>', lora_stack) + + # Check if any loras were found + if not lora_matches: + return web.json_response({"error": "No LoRAs found in the workflow"}, status=400) + + # Generate recipe name from the first 3 loras (or less if fewer are available) + loras_for_name = lora_matches[:3] # Take at most 3 loras for the name + + recipe_name_parts = [] + for lora_name, lora_strength in loras_for_name: + # Get the basename without path or extension + basename = os.path.basename(lora_name) + basename = os.path.splitext(basename)[0] + recipe_name_parts.append(f"{basename}:{lora_strength}") + + recipe_name = " ".join(recipe_name_parts) # Read the image with open(latest_image_path, 'rb') as f: @@ -819,30 +832,29 @@ class RecipeRoutes: with open(image_path, 'wb') as f: f.write(optimized_image) - # Format loras data from metadata + # Format loras data from the lora stack 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 + for lora_name, lora_strength in lora_matches: + try: + # Get lora info from scanner + 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_strength), + "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) + except Exception as e: + logger.warning(f"Error processing LoRA {lora_name}: {e}") + + # Get base model from lora scanner for the available loras 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", "")) @@ -859,18 +871,15 @@ class RecipeRoutes: recipe_data = { "id": recipe_id, "file_path": image_path, - "title": name, + "title": recipe_name, # Use generated recipe 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 + "gen_params": parsed_workflow.get("gen_params", {}), # Use the parsed workflow parameters + "loras_stack": lora_stack # Include the original lora stack string } - # 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) @@ -892,7 +901,8 @@ class RecipeRoutes: 'success': True, 'recipe_id': recipe_id, 'image_path': image_path, - 'json_path': json_path + 'json_path': json_path, + 'recipe_name': recipe_name # Include the generated recipe name in the response }) except Exception as e: diff --git a/py/utils/exif_utils.py b/py/utils/exif_utils.py index f31ec575..f1b304a1 100644 --- a/py/utils/exif_utils.py +++ b/py/utils/exif_utils.py @@ -44,7 +44,6 @@ class ExifUtils: return None except Exception as e: - logger.error(f"Error extracting EXIF data from {image_path}: {e}") return None @staticmethod diff --git a/web/comfyui/loras_widget.js b/web/comfyui/loras_widget.js index 362db08e..ebe162c9 100644 --- a/web/comfyui/loras_widget.js +++ b/web/comfyui/loras_widget.js @@ -769,28 +769,8 @@ export function addLorasWidget(node, name, opts, callback) { // Function to directly save the recipe without dialog async function saveRecipeDirectly(widget) { try { + // Get the workflow data from the ComfyUI app const prompt = await app.graphToPrompt(); - console.log("prompt", prompt); - // 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) { @@ -802,21 +782,9 @@ async function saveRecipeDirectly(widget) { }); } - // Prepare the data + // Prepare the data - only send workflow JSON 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)); + formData.append('workflow_json', JSON.stringify(prompt.output)); // Send the request const response = await fetch('/api/recipes/save-from-widget', {