Update app.py

app.py

@@ -2,49 +2,79 @@ import gradio as gr
 from PIL import Image
 import json
 
-def extract_parameters(image):
+def extract_parameters(image, show_full_metadata=False):
     """
-    Extract metadata/parameters from a ComfyUI-generated image.
-    Returns a formatted string with the extracted info.
+    Extract and format parameters from a ComfyUI image.
+    If show_full_metadata is False, shows only the main prompt cleanly.
+    If True, shows the full metadata with better formatting.
     """
     try:
         # Open the uploaded image
         img = Image.open(image)
         # Get metadata from the image's info dictionary
         metadata = img.info
-        
-        # Check common ComfyUI metadata keys (adjust based on your images)
+
+        if not metadata:
+            return "No parameters found in image metadata."
+
+        # Check if "prompt" key exists (common in ComfyUI)
         if "prompt" in metadata:
-            prompt = metadata["prompt"]
-            # If it's a JSON string, try to parse it
+            raw_prompt = metadata["prompt"]
             try:
-                prompt_data = json.loads(prompt)
-                # Format JSON nicely
-                output = "Extracted Parameters:\n"
-                for key, value in prompt_data.items():
-                    output += f"{key}: {value}\n"
-                return output
+                # Parse as JSON if possible
+                prompt_data = json.loads(raw_prompt)
+                # Look for the main prompt text (often in a node like '51' from your output)
+                main_prompt = None
+                for node_id, node in prompt_data.items():
+                    if "inputs" in node and "text" in node["inputs"]:
+                        # Prioritize nodes with meaningful text input
+                        text = node["inputs"]["text"]
+                        if isinstance(text, str) and text.strip():
+                            main_prompt = text
+                            break
+                
+                if not main_prompt:
+                    main_prompt = "No clear prompt found in workflow."
+
+                if show_full_metadata:
+                    # Format full metadata with spacing
+                    full_output = "Full Workflow Metadata:\n\n"
+                    for node_id, node in prompt_data.items():
+                        full_output += f"Node {node_id}:\n"
+                        for key, value in node.items():
+                            full_output += f"  {key}: {value}\n"
+                        full_output += "\n"
+                    return f"Main Prompt:\n{main_prompt}\n\n{full_output}"
+                else:
+                    # Just the cleaned-up main prompt
+                    return f"Main Prompt:\n{main_prompt}"
             except json.JSONDecodeError:
-                # If not JSON, return raw prompt
-                return f"Prompt: {prompt}"
-        elif metadata:
-            # If no "prompt" key, dump all metadata
-            output = "Metadata Found:\n"
-            for key, value in metadata.items():
-                output += f"{key}: {value}\n"
-            return output
+                # If not JSON, treat as raw text
+                if show_full_metadata:
+                    return f"Raw Prompt:\n{raw_prompt}\n\nFull Metadata:\n{metadata}"
+                return f"Main Prompt:\n{raw_prompt}"
         else:
-            return "No parameters found in image metadata."
+            # No "prompt" key, show all metadata if toggled
+            if show_full_metadata:
+                output = "Full Metadata:\n\n"
+                for key, value in metadata.items():
+                    output += f"{key}:\n  {value}\n\n"
+                return output
+            return "No prompt found, toggle 'Show Full Metadata' for details."
+
     except Exception as e:
         return f"Error processing image: {str(e)}"
 
 # Define the Gradio interface
 interface = gr.Interface(
     fn=extract_parameters,
-    inputs=gr.Image(type="filepath", label="Upload a ComfyUI Image"),
+    inputs=[
+        gr.Image(type="filepath", label="Upload a ComfyUI Image"),
+        gr.Checkbox(label="Show Full Metadata", default=False)
+    ],
     outputs=gr.Textbox(label="Extracted Parameters", lines=10),
     title="ComfyUI Image to Parameters",
-    description="Upload an image generated by ComfyUI to extract its embedded parameters or prompt.",
+    description="Upload a ComfyUI image to extract its prompt. Toggle 'Show Full Metadata' for the complete workflow.",
     theme="default"
 )