Upload app.py

app.py

@@ -3,53 +3,13 @@ from PIL import Image
 from PIL.ExifTags import TAGS
 import numpy as np
 import re
+import json
+import tempfile
+import os
 
-def main():
-    # Create Gradio User Interface
-    with gr.Blocks() as iface:
-        gr.Markdown("# Image Metadata Extractor")
-        gr.Markdown("Upload an image (PNG, JPEG, TIFF) to extract its metadata and parse it for prompts.")
-        with gr.Row():
-            with gr.Column():  # First column for image upload and preview
-                image_input = gr.File(label="Upload Image", type="filepath")
-                image_output = gr.Image(label="Image Preview")
-                original_metadata_output = gr.Textbox(label="Original Metadata", lines=20)
-
-            with gr.Column():  # Second column for metadata outputs
-                prompt_output = gr.Textbox(label="Prompt", lines=4, show_copy_button=True)
-                negative_prompt_output = gr.Textbox(label="Negative Prompt", lines=4, show_copy_button=True)
-                seed_output = gr.Textbox(label="Seed Number", lines=1, show_copy_button=True)
-                adetailer_prompt_1_output = gr.Textbox(label="ADetailer Prompt 1", lines=3, show_copy_button=True)
-                adetailer_prompt_2_output = gr.Textbox(label="ADetailer Prompt 2", lines=3, show_copy_button=True)
-            with gr.Column():  # Second column for metadata outputs
-                with gr.Row():
-                    comfy_prompt_output = gr.Textbox(label="Comfy Prompt", lines=20)
-                with gr.Row():
-                    comfy_workflow_output = gr.Textbox(label="Comfy Workflow", lines=20)
-
-        # Set up the input-output relationships
-        image_input.change(
-            fn=extract_metadata,
-            inputs=image_input,
-            outputs=[
-                image_output,
-                prompt_output,
-                negative_prompt_output,
-                seed_output,
-                adetailer_prompt_1_output,
-                adetailer_prompt_2_output,
-                original_metadata_output,
-                comfy_prompt_output, 
-                comfy_workflow_output,
-            ]
-        )
-    iface.launch()
-
-def extract_metadata(image_file):
-    img = Image.open(image_file)
+def get_image_metadata(img):
+    """Extract metadata based on image format."""
     metadata = {}
-
-    # Extract metadata
     if img.format == 'PNG':
         metadata = img.info
     elif img.format in ['JPEG', 'TIFF']:
@@ -58,66 +18,82 @@ def extract_metadata(image_file):
             for tag, value in exif_data.items():
                 tag_name = TAGS.get(tag, tag)
                 metadata[tag_name] = value
+    return metadata
 
-    # Convert image to NumPy array for display
-    img_display = np.array(img)
-
-    # Initialize outputs
+def parse_parameters(parameters):
+    """Parse the 'parameters' field to extract prompts and seed number."""
     prompt = 'N/A'
     negative_prompt = 'N/A'
     adetailer_prompt_1 = 'N/A'
     adetailer_prompt_2 = 'N/A'
     seed_number = -1
-    comfy_prompt = 'N/A'
-    comfy_workflow = 'N/A'
 
-    # Convert the entire metadata dictionary to a readable string
+    # Extract the prompt
+    prompt_match = re.search(r'(.*?)Negative prompt:', parameters, re.DOTALL | re.IGNORECASE)
+    if prompt_match:
+        prompt = prompt_match.group(1).strip()
+    else:
+        prompt_match = re.search(r'(.*?)(Steps:|$)', parameters, re.DOTALL | re.IGNORECASE)
+        if prompt_match:
+            prompt = prompt_match.group(1).strip()
+
+    # Extract the negative prompt
+    negative_prompt_match = re.search(r'Negative prompt:(.*?)(Steps:|$)', parameters, re.DOTALL | re.IGNORECASE)
+    if negative_prompt_match:
+        negative_prompt = negative_prompt_match.group(1).strip()
+
+    # Extract ADetailer prompts
+    adetailer_prompt_1_match = re.search(r'ADetailer prompt:\s*"(.*?)"', parameters, re.IGNORECASE)
+    if adetailer_prompt_1_match:
+        adetailer_prompt_1 = adetailer_prompt_1_match.group(1).strip()
+
+    adetailer_prompt_2_match = re.search(r'ADetailer negative prompt 2nd:\s*"(.*?)"', parameters, re.IGNORECASE)
+    if adetailer_prompt_2_match:
+        adetailer_prompt_2 = adetailer_prompt_2_match.group(1).strip()
+
+    # Extract Seed Number
+    seed_match = re.search(r'Seed:\s*(\d+)', parameters, re.IGNORECASE)
+    if seed_match:
+        seed_number = seed_match.group(1).strip()
+
+    return prompt, negative_prompt, adetailer_prompt_1, adetailer_prompt_2, seed_number
+
+def extract_metadata(image_file):
+    """Extract and parse metadata from the uploaded image."""
+    try:
+        img = Image.open(image_file)
+    except Exception as e:
+        # Return placeholders in case of error
+        return (None, 'Error: Unable to open image.', *(['N/A'] * 7))
+
+    metadata = get_image_metadata(img)
+    img_display = np.array(img)
+
+    # Convert metadata to string
     metadata_str = "\n".join([f"{key}: {value}" for key, value in metadata.items()])
 
-    # Parse metadata to extract prompts
+    # Initialize default values
+    prompt = negative_prompt = adetailer_prompt_1 = adetailer_prompt_2 = 'N/A'
+    seed_number = -1
+    comfy_prompt = comfy_workflow = 'N/A'
+
+    # Parse parameters if available
     if 'parameters' in metadata:
-        parameters = metadata['parameters']
+        parsed = parse_parameters(metadata['parameters'])
+        prompt, negative_prompt, adetailer_prompt_1, adetailer_prompt_2, seed_number = parsed
 
-        # Use regular expressions to extract prompts
-        # Extract the prompt (before 'Negative prompt:')
-        prompt_match = re.search(r'(.*?)Negative prompt:', parameters, re.DOTALL)
-        if prompt_match:
-            prompt = prompt_match.group(1).strip()
-        else:
-            # If 'Negative prompt:' not found, try to extract prompt up to 'Steps:' or end of string
-            prompt_match = re.search(r'(.*?)(Steps:|$)', parameters, re.DOTALL)
-            if prompt_match:
-                prompt = prompt_match.group(1).strip()
-
-        # Extract the negative prompt
-        negative_prompt_match = re.search(r'Negative prompt:(.*?)(Steps:|$)', parameters, re.DOTALL)
-        if negative_prompt_match:
-            negative_prompt = negative_prompt_match.group(1).strip()
-
-        # Extract ADetailer prompt 1
-        adetailer_prompt_1_match = re.search(r'ADetailer prompt:\s*"(.*?)"', parameters)
-        if adetailer_prompt_1_match:
-            adetailer_prompt_1 = adetailer_prompt_1_match.group(1).strip()
-
-        # Extract ADetailer prompt 2
-        adetailer_prompt_2_match = re.search(r'ADetailer negative prompt 2nd:\s*"(.*?)"', parameters)
-        if adetailer_prompt_2_match:
-            adetailer_prompt_2 = adetailer_prompt_2_match.group(1).strip()
-        
-        seed_match = re.search(r'Seed:\s*(\d+)', parameters)
-        if seed_match:
-            seed_number = seed_match.group(1).strip()
-
-    # Extract additional metadata fields directly if they exist
-    if 'prompt' in metadata:
+    # Direct extraction with validation for Comfy Prompt
+    if 'prompt' in metadata and isinstance(metadata['prompt'], str):
         comfy_prompt = metadata['prompt'].strip()
+    else:
+        comfy_prompt = 'N/A'  # Fallback if prompt is missing or not a string
 
-
-    if 'workflow' in metadata:
+    # Direct extraction with validation for Comfy Workflow
+    if 'workflow' in metadata and isinstance(metadata['workflow'], str):
         comfy_workflow = metadata['workflow'].strip()
+    else:
+        comfy_workflow = 'N/A'  # Fallback if workflow is missing or not a string
 
-
-    # Return the outputs, replacing 'original_metadata_output' with 'metadata_str'
     return (
         img_display, 
         prompt, 
@@ -128,7 +104,96 @@ def extract_metadata(image_file):
         metadata_str, 
         comfy_prompt, 
         comfy_workflow
+    )
+
+def export_workflow(workflow_text):
+    """
+    Converts the workflow text into JSON format and returns it as a downloadable file.
+    """
+    print("Export workflow function called.")
+    if workflow_text == "N/A" or not workflow_text.strip():
+        print("No workflow data to export.")
+        return None, "No workflow data to export."
+
+    # Structure the JSON data
+    workflow_data = {
+        "comfy_workflow": workflow_text
+    }
+
+    # Create a temporary JSON file
+    try:
+        with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.json') as tmp_file:
+            json.dump(workflow_data, tmp_file, indent=4)
+            tmp_file_path = tmp_file.name
+        print(f"Temporary JSON file created at: {tmp_file_path}")
+    except Exception as e:
+        print(f"Error creating temporary file: {e}")
+        return None, "Failed to create workflow JSON file."
+
+    # Check if the file was created successfully
+    if os.path.exists(tmp_file_path):
+        message = "Workflow exported successfully."
+        print(message)
+        return tmp_file_path, message
+    else:
+        message = "Failed to export workflow."
+        print(message)
+        return None, message
+
+def main():
+    # Create Gradio User Interface
+    with gr.Blocks() as iface:
+        gr.Markdown("<h1>Automatic 1111/ Comfy Metadata Reader</h1>")
+        gr.Markdown("<p>Upload an image (PNG, JPEG) to extract its metadata and parse it for prompts.</p>")
+        with gr.Row():
+            with gr.Column():  # First column for image upload and preview
+                image_input = gr.File(label="Upload Image", type="filepath")
+                image_output = gr.Image(label="Image Preview")
+                original_metadata_output = gr.Textbox(label="Original Metadata", lines=20, interactive=False)
+
+            with gr.Column():  # Second column for metadata outputs
+                prompt_output = gr.Textbox(label="Prompt", lines=4, show_copy_button=True, interactive=False)
+                negative_prompt_output = gr.Textbox(label="Negative Prompt", lines=4, show_copy_button=True, interactive=False)
+                seed_output = gr.Textbox(label="Seed Number", lines=1, show_copy_button=True, interactive=False)
+                adetailer_prompt_1_output = gr.Textbox(label="ADetailer Prompt 1", lines=3, show_copy_button=True, interactive=False)
+                adetailer_prompt_2_output = gr.Textbox(label="ADetailer Prompt 2", lines=3, show_copy_button=True, interactive=False)
+
+            with gr.Column():  # Third column for additional metadata outputs
+                with gr.Row():
+                    comfy_prompt_output = gr.Textbox(label="Comfy Prompt", lines=4, value="N/A", interactive=False)
+                with gr.Row():
+                    comfy_workflow_output = gr.Textbox(label="Comfy Workflow", lines=20, value="N/A", interactive=False)
+                    export_button = gr.Button("Export Workflow as JSON")
+                    workflow_file = gr.File(label="Download Workflow JSON", visible=False)
+                with gr.Row():
+                    export_message = gr.Textbox(label="Export Status", lines=1, interactive=False, visible=False)
+
+        # Set up the input-output relationships
+        image_input.change(
+            fn=extract_metadata,
+            inputs=image_input,
+            outputs=[
+                image_output,
+                prompt_output,
+                negative_prompt_output,
+                seed_output,
+                adetailer_prompt_1_output,
+                adetailer_prompt_2_output,
+                original_metadata_output,
+                comfy_prompt_output, 
+                comfy_workflow_output,
+            ]
+        )
+
+        # Connect the export button to the export function
+        export_button.click(
+            fn=export_workflow,
+            inputs=comfy_workflow_output,
+            outputs=[workflow_file, export_message],
+            queue=False
         )
 
+    iface.launch()
+
 if __name__ == "__main__":
     main()