mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-03-21 21:22:11 -03:00
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:
@@ -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()
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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 {}
|
||||
@@ -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
|
||||
@@ -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
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user