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.
This commit is contained in:
Will Miao
2025-03-21 11:11:09 +08:00
parent 2cf4440a1e
commit 4ee32f02c5
6 changed files with 433 additions and 12 deletions

View File

@@ -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()

View File

@@ -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)

View File

@@ -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

View File

@@ -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'
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 {}

View File

@@ -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

View File

@@ -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',
'<svg width="14" height="14" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2"><path d="M19 21l-7-5-7 5V5a2 2 0 0 1 2-2h10a2 2 0 0 1 2 2z"></path></svg>',
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
});
}
}
}