diff --git a/py/api.py b/py/api.py index 204d414..7791e87 100644 --- a/py/api.py +++ b/py/api.py @@ -317,6 +317,10 @@ class PromptManagerAPI: @routes.post("/prompt_manager/images/link") async def link_image_route(request): return await self.link_image_to_prompt(request) + + @routes.get("/prompt_manager/images/prompt/{image_path:.*}") + async def get_image_prompt_route(request): + return await self.get_image_prompt(request) @routes.delete("/prompt_manager/images/{image_id}") async def delete_image_route(request): @@ -2368,6 +2372,85 @@ class PromptManagerAPI: 'error': str(e) }, status=500) + async def get_image_prompt(self, request): + """Get prompt information for a specific image path.""" + try: + import urllib.parse + import os + import json + from pathlib import Path + + # Get the image path from URL + raw_image_path = request.match_info.get('image_path', '') + image_path = urllib.parse.unquote(raw_image_path) + + if not image_path: + return web.json_response({ + 'success': False, + 'error': 'Image path is required' + }, status=400) + + # Convert relative path to absolute if needed + if not os.path.isabs(image_path): + # If it's a relative path from ComfyUI output, make it absolute + output_dir = self._find_comfyui_output_dir() + if output_dir: + image_path = str(Path(output_dir) / image_path) + + # Look up the image in generated_images table + try: + with self.db.model.get_connection() as conn: + cursor = conn.execute( + """SELECT gi.prompt_id, p.text, p.category, p.tags, p.rating, p.notes, + gi.workflow_data, gi.prompt_metadata, gi.generation_time + FROM generated_images gi + JOIN prompts p ON gi.prompt_id = p.id + WHERE gi.image_path = ? OR gi.image_path LIKE ?""", + (image_path, f'%{os.path.basename(image_path)}') + ) + result = cursor.fetchone() + + if result: + # Convert to dict + prompt_data = { + 'prompt_id': result[0], + 'text': result[1], + 'category': result[2], + 'tags': json.loads(result[3]) if result[3] else [], + 'rating': result[4], + 'notes': result[5], + 'workflow_data': json.loads(result[6]) if result[6] else None, + 'prompt_metadata': json.loads(result[7]) if result[7] else None, + 'generation_time': result[8], + 'image_path': image_path + } + + return web.json_response({ + 'success': True, + 'prompt': prompt_data + }) + else: + # No linked prompt found - this is normal for many images + return web.json_response({ + 'success': False, + 'error': 'No prompt found for this image', + 'image_path': image_path + }) + + except Exception as db_error: + self.logger.error(f"Database error in get_image_prompt: {db_error}") + return web.json_response({ + 'success': False, + 'error': 'Database error occurred' + }, status=500) + + except Exception as e: + self.logger.error(f"Get image prompt error: {e}") + return web.json_response({ + 'success': False, + 'error': str(e) + }, status=500) + async def delete_image(self, request): """Delete an image record.""" try: diff --git a/web/gallery.html b/web/gallery.html index 9b7ee72..5e73781 100644 --- a/web/gallery.html +++ b/web/gallery.html @@ -1405,7 +1405,72 @@ async loadImageMetadata(imageUrl, container) { try { - const metadata = await this.extractImageMetadata(imageUrl); + // First try to get linked prompt from database + let databasePrompt = null; + try { + // Extract relative path from the URL for the API call + const urlParts = imageUrl.split('/prompt_manager/images/serve/'); + if (urlParts.length > 1) { + const relativePath = urlParts[1]; + const response = await fetch(`/prompt_manager/images/prompt/${encodeURIComponent(relativePath)}`); + if (response.ok) { + const data = await response.json(); + if (data.success && data.prompt) { + databasePrompt = data.prompt; + } + } + } + } catch (dbError) { + // No database prompt found, will fall back to metadata extraction + } + + // If we have a database prompt, use it; otherwise extract from metadata + let metadata; + if (databasePrompt) { + // Create metadata object from database prompt + metadata = { + positivePrompt: databasePrompt.text, + negativePrompt: 'No negative prompt found', // Database doesn't separate positive/negative + checkpoint: 'Unknown', + steps: 'Unknown', + cfgScale: 'Unknown', + sampler: 'Unknown', + seed: 'Unknown', + workflow: databasePrompt.workflow_data, + imagePath: imageUrl, + source: 'database', // Mark this as coming from database + promptId: databasePrompt.prompt_id, + category: databasePrompt.category, + tags: databasePrompt.tags, + rating: databasePrompt.rating, + notes: databasePrompt.notes + }; + + // Still try to extract technical parameters from PNG metadata if available + try { + const pngMetadata = await this.extractImageMetadata(imageUrl); + if (pngMetadata) { + metadata.checkpoint = pngMetadata.checkpoint; + metadata.steps = pngMetadata.steps; + metadata.cfgScale = pngMetadata.cfgScale; + metadata.sampler = pngMetadata.sampler; + metadata.seed = pngMetadata.seed; + // Use negative prompt from PNG if available + if (pngMetadata.negativePrompt && pngMetadata.negativePrompt !== 'No negative prompt found') { + metadata.negativePrompt = pngMetadata.negativePrompt; + } + } + } catch (pngError) { + // Could not extract PNG metadata, using database-only data + } + } else { + // Fall back to PNG metadata extraction + metadata = await this.extractImageMetadata(imageUrl); + if (metadata) { + metadata.source = 'png'; // Mark this as coming from PNG + } + } + this.displayMetadata(metadata, container); } catch (error) { console.error('Error loading metadata:', error); @@ -1894,46 +1959,85 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`; if (comfyData.prompt) { const promptNodes = comfyData.prompt; + // Track all text nodes and try to identify positive/negative + const textNodes = []; + for (const nodeId in promptNodes) { const node = promptNodes[nodeId]; - // Checkpoint - if (node.class_type === 'CheckpointLoaderSimple' && node.inputs) { - checkpoint = node.inputs.ckpt_name || checkpoint; + // Checkpoint - check multiple types + if ((node.class_type === 'CheckpointLoaderSimple' || + node.class_type === 'UNETLoader' || + node.class_type === 'DualCLIPLoader') && node.inputs) { + checkpoint = node.inputs.ckpt_name || node.inputs.unet_name || checkpoint; } - // Prompts - need to identify which is positive vs negative - if (node.class_type === 'PromptManager' && node.inputs && node.inputs.text) { - // PromptManager typically contains the positive prompt - positivePrompt = node.inputs.text; + // Collect all text nodes for analysis + if (node.inputs && node.inputs.text) { + textNodes.push({ + nodeId: nodeId, + classType: node.class_type, + text: node.inputs.text, + inputs: node.inputs + }); } - if (node.class_type === 'CLIPTextEncode' && node.inputs && node.inputs.text) { - // Check if this looks like a negative prompt - const text = node.inputs.text.toLowerCase(); - if (text.includes('bad anatomy') || text.includes('unfinished') || - text.includes('censored') || text.includes('weird anatomy') || - text.includes('negative') || text.includes('embedding:')) { - negativePrompt = node.inputs.text; - } else if (positivePrompt === 'No prompt found') { - // If we haven't found a positive prompt yet, this might be it + // Special handling for PromptManager nodes + if (node.class_type === 'PromptManager' && node.inputs) { + // PromptManager always contains the positive prompt + if (node.inputs.text) { positivePrompt = node.inputs.text; } + // Also check for loaded prompt + if (node.inputs.selected_prompt) { + positivePrompt = node.inputs.selected_prompt; + } } - // Sampling parameters - if (node.class_type === 'KSampler' && node.inputs) { - steps = node.inputs.steps || steps; - cfgScale = node.inputs.cfg || cfgScale; - sampler = node.inputs.sampler_name || sampler; - seed = node.inputs.seed || seed; + // PromptManagerText node + if (node.class_type === 'PromptManagerText' && node.inputs) { + if (node.inputs.text) { + positivePrompt = node.inputs.text; + } + if (node.inputs.selected_prompt) { + positivePrompt = node.inputs.selected_prompt; + } } - if (node.class_type === 'KSamplerAdvanced' && node.inputs) { + // Sampling parameters - check multiple sampler types + if ((node.class_type === 'KSampler' || + node.class_type === 'KSamplerAdvanced' || + node.class_type === 'SamplerCustom' || + node.class_type === 'SamplerCustomAdvanced') && node.inputs) { steps = node.inputs.steps || steps; cfgScale = node.inputs.cfg || cfgScale; sampler = node.inputs.sampler_name || sampler; - seed = node.inputs.noise_seed || seed; + seed = node.inputs.seed || node.inputs.noise_seed || seed; + } + } + + // Process collected text nodes to identify positive/negative prompts + // if we haven't found them from PromptManager + for (const textNode of textNodes) { + const text = textNode.text; + const textLower = text.toLowerCase(); + + // Skip if this is the PromptManager text we already have + if (textNode.classType === 'PromptManager' || textNode.classType === 'PromptManagerText') { + continue; + } + + // Identify negative prompts by common patterns + if (textNode.classType === 'CLIPTextEncode') { + if (textLower.includes('bad anatomy') || textLower.includes('unfinished') || + textLower.includes('censored') || textLower.includes('weird anatomy') || + textLower.includes('negative') || textLower.includes('embedding:') || + textLower.includes('worst quality') || textLower.includes('low quality')) { + negativePrompt = text; + } else if (positivePrompt === 'No prompt found') { + // If we haven't found a positive prompt yet, this might be it + positivePrompt = text; + } } } } @@ -1941,23 +2045,44 @@ Seed: ${this.currentMetadata.seed || 'Unknown'}`; // Also try to parse from workflow data if we didn't find everything in prompt if (comfyData.workflow && comfyData.workflow.nodes) { for (const node of comfyData.workflow.nodes) { - if (node.type === 'CheckpointLoaderSimple' && node.widgets_values) { + // Checkpoint loaders + if ((node.type === 'CheckpointLoaderSimple' || + node.type === 'UNETLoader' || + node.type === 'DualCLIPLoader') && node.widgets_values) { checkpoint = node.widgets_values[0] || checkpoint; } + // PromptManager nodes in workflow + if (node.type === 'PromptManager' && node.widgets_values) { + // The first widget value is usually the text + if (node.widgets_values[0] && positivePrompt === 'No prompt found') { + positivePrompt = node.widgets_values[0]; + } + } + + if (node.type === 'PromptManagerText' && node.widgets_values) { + if (node.widgets_values[0] && positivePrompt === 'No prompt found') { + positivePrompt = node.widgets_values[0]; + } + } + + // CLIPTextEncode nodes if (node.type === 'CLIPTextEncode' && node.widgets_values && node.widgets_values[0]) { const text = node.widgets_values[0]; const textLower = text.toLowerCase(); if (textLower.includes('bad anatomy') || textLower.includes('unfinished') || - textLower.includes('censored') || textLower.includes('negative')) { + textLower.includes('censored') || textLower.includes('negative') || + textLower.includes('worst quality') || textLower.includes('low quality')) { negativePrompt = text; } else if (positivePrompt === 'No prompt found') { positivePrompt = text; } } - if ((node.type === 'KSampler' || node.type === 'KSamplerAdvanced') && node.widgets_values) { + // Samplers + if ((node.type === 'KSampler' || node.type === 'KSamplerAdvanced' || + node.type === 'SamplerCustom' || node.type === 'SamplerCustomAdvanced') && node.widgets_values) { if (node.widgets_values.length >= 4) { seed = node.widgets_values[0] || seed; steps = node.widgets_values[1] || steps;