Update app.py

app.py

@@ -1,52 +1,46 @@
 import gradio as gr
-from PIL import Image
+from PIL import Image, PngImagePlugin
 import io
-import re
 
 def extract_metadata(image):
     """Extracts Stable Diffusion metadata from a PIL Image."""
     if image:
-        metadata = image.text.get("parameters", "")
+        # Access metadata using the 'info' dictionary
+        metadata = image.info.get("parameters", "")
     else:
         metadata = ""
     return image, metadata
 
-
 def update_metadata(image, new_metadata):
-    """Updates the Stable Diffusion metadata in a PIL Image and saves it."""
+    """Updates the Stable Diffusion metadata in a PIL Image."""
     if image:
-        # Convert PIL Image to bytes
-        img_byte_arr = io.BytesIO()
-        image.save(img_byte_arr, format='PNG', pnginfo=image.text)
-        img_byte_arr = img_byte_arr.getvalue()
-
-        # Reload image to apply changes.  PIL is a bit finicky about modifying .text
-        img = Image.open(io.BytesIO(img_byte_arr))
+        # Create a PngInfo object to store metadata
+        pnginfo = PngImagePlugin.PngInfo()
+        pnginfo.add_text("parameters", new_metadata)
 
-        img.text["parameters"] = new_metadata
-        
-        # Save to a BytesIO object
+        # Save the image to a BytesIO object with the updated metadata
         output_bytes = io.BytesIO()
-        img.save(output_bytes, format="PNG")  # Save back as PNG with updated metadata
-        output_bytes.seek(0) # Important: Rewind the buffer to the beginning
+        image.save(output_bytes, format="PNG", pnginfo=pnginfo)
+        output_bytes.seek(0)
+
+        # Re-open the image from the BytesIO object to ensure Gradio displays the updated image
+        updated_image = Image.open(output_bytes)
 
-        return output_bytes, new_metadata
+        return updated_image, new_metadata # Return updated PIL Image
     else:
-      return None, ""
+        return None, ""
 
 def clear_metadata(image, metadata):
-    """Clears the Stable Diffusion Metadata"""
+    """Clears the Stable Diffusion Metadata."""
     return update_metadata(image, "")
 
 
 def copy_metadata(metadata: str):
-    """Copies the metadata to the clipboard (simulated)."""
-    # gradio.Info("Metadata copied to clipboard!")  # Informational message
-    metadata = metadata.strip()  # Clean up leading/trailing whitespace
-    # Clears Used Embeddings
+    """Copies the metadata to the clipboard (simulated in Gradio)."""
+    metadata = metadata.strip()
     embed_index = metadata.find("\nUsed embeddings: ")
     if embed_index < 0:
-        embed_index = metadata.find("\nUsed embeddings: ")
+        embed_index = metadata.find("\nUsed embeddings: ")  # Corrected typo
     if embed_index > 0:
         metadata = metadata[:embed_index]
 
@@ -57,7 +51,6 @@ def process_image(image):
     return extract_metadata(image)
 
 
-
 with gr.Blocks() as demo:
     gr.Markdown("# SD EXIF Editor (Gradio)")
     gr.Markdown("Upload a PNG image generated by Stable Diffusion WebUI (AUTOMATIC1111) to view and edit its metadata.")
@@ -65,8 +58,7 @@ with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column():
             image_input = gr.Image(type="pil", label="Upload PNG Image")
-            image_output = gr.Image(type="pil", label="Processed Image")
-            # image_output = gr.Image(type="filepath", label="Processed Image") #For local testing.
+            image_output = gr.Image(type="pil", label="Processed Image") # Keep as PIL
 
         with gr.Column():
             metadata_textbox = gr.Textbox(label="Metadata", lines=5)
@@ -75,12 +67,10 @@ with gr.Blocks() as demo:
                 clear_button = gr.Button("Clear Metadata")
                 save_button = gr.Button("Save Metadata")
 
-
     # --- Event Handling ---
     image_input.change(process_image, inputs=[image_input], outputs=[image_output, metadata_textbox])
-    save_button.click(update_metadata, inputs=[image_output, metadata_textbox], outputs=[image_output, metadata_textbox])
-    clear_button.click(clear_metadata, inputs=[image_output, metadata_textbox], outputs=[image_output, metadata_textbox])
+    save_button.click(update_metadata, inputs=[image_input, metadata_textbox], outputs=[image_output, metadata_textbox])  # Use image_input
+    clear_button.click(clear_metadata, inputs=[image_input, metadata_textbox], outputs=[image_output, metadata_textbox])  # Use image_input
     copy_button.click(copy_metadata, inputs=[metadata_textbox], outputs=[metadata_textbox])
 
-
-demo.launch()
+demo.launch(share=True) # Enable public link