From 3518e78b664359131a936aa1c4155735a304b0a9 Mon Sep 17 00:00:00 2001
From: ohmygoobness <82442572+ohmygoobness@users.noreply.github.com>
Date: Tue, 8 Aug 2023 06:04:45 +0530
Subject: [PATCH] Added PNGInfo tab to app.py

will need to install pillow
---
 app.py | 49 ++++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 48 insertions(+), 1 deletion(-)

diff --git a/app.py b/app.py
index c07df3ba..63e3d5de 100644
--- a/app.py
+++ b/app.py
@@ -27,6 +27,8 @@
 import tqdm
 import re
 import shutil
+from PIL import Image
+from urllib.request import Request, urlopen
 
 sample_idx     = 0
 max_LoRAs      = 5
@@ -90,6 +92,29 @@ def download_checkpoint(url, progress=gr.Progress(track_tqdm=True)):
 def download_loras(url, progress=gr.Progress(track_tqdm=True)):
     return download_url(url=url, destination_path=controller.loras_dir)
 
+def pnginfo_from_url(url, input_image):
+    if url == '' and input_image is None:
+        return None, "", ""
+
+    image = None
+    if url != '':
+        req = Request(
+            url=url, 
+            headers={'User-Agent': 'Mozilla/5.0'}
+        )
+        res = urlopen(req)
+        image = Image.open(res)
+        image.load()
+
+
+    if input_image is not None:
+        image = input_image
+
+    parameters = "Parameters have been erased from this image or unsupported format"
+    if 'parameters' in image.info:
+        parameters = image.info['parameters']
+
+    return image, parameters, image.info
 
 class ProjectConfigs:
     def __init__(self):
@@ -615,6 +640,25 @@ def configuration_tab_ui():
         """
     )
 
+def pnginfo_tab_ui():
+    blocks = gr.Blocks(css="#out_image {height: 400px}")
+    with blocks as png_info:
+        with gr.Row().style(equal_height=False):
+            with gr.Column():
+                in_url = gr.Textbox(label="Source URL")
+                in_image = gr.Image(label="Source Image", type='pil')
+                with gr.Row():
+                    btn_submit = gr.Button("Submit", variant="primary")
+
+            with gr.Column():
+                out_info = gr.Textbox(label="Generation Parameters")
+                out_meta = gr.Textbox(label="Metadata")
+                out_image = gr.Image(type='pil', elem_id="out_image")
+        
+        btn_submit.click(fn=pnginfo_from_url,
+        inputs=[in_url, in_image],
+        outputs=[out_image, out_info, out_meta])
+
 def credits_tab_ui():
     gr.Markdown(
         """
@@ -641,6 +685,9 @@ def ui():
         with gr.Tab(label="Configurations"):
             configuration_tab_ui()
 
+        with gr.Tab(label="PNGInfo"):
+            pnginfo_tab_ui()
+
         with gr.Tab(label="Credits"):
             credits_tab_ui()
 
@@ -650,4 +697,4 @@ def ui():
 if __name__ == "__main__":
     demo = ui()
     demo.queue(concurrency_count=3)
-    demo.launch(share=True, debug=True)
\ No newline at end of file
+    demo.launch(share=True, debug=True)