1 files changed, 12 insertions, 4 deletions
diff --git a/modules/ui_extra_networks.py b/modules/ui_extra_networks.py
index daea03d6..aa2f5d1b 100644
--- a/modules/ui_extra_networks.py
+++ b/modules/ui_extra_networks.py
@@ -2,8 +2,10 @@ import glob
 import os.path
 import urllib.parse
 from pathlib import Path
+from PIL import PngImagePlugin
 
 from modules import shared
+from modules.images import read_info_from_image
 import gradio as gr
 import json
 import html
@@ -239,7 +241,7 @@ def create_ui(container, button, tabname):
 
     with gr.Tabs(elem_id=tabname+"_extra_tabs") as tabs:
         for page in ui.stored_extra_pages:
-            with gr.Tab(page.title):
+            with gr.Tab(page.title, id=page.title.lower().replace(" ", "_")):
 
                 page_elem = gr.HTML(page.create_html(ui.tabname))
                 ui.pages.append(page_elem)
@@ -252,10 +254,10 @@ def create_ui(container, button, tabname):
 
     def toggle_visibility(is_visible):
         is_visible = not is_visible
-        return is_visible, gr.update(visible=is_visible)
+        return is_visible, gr.update(visible=is_visible), gr.update(variant=("secondary-down" if is_visible else "secondary"))
 
     state_visible = gr.State(value=False)
-    button.click(fn=toggle_visibility, inputs=[state_visible], outputs=[state_visible, container])
+    button.click(fn=toggle_visibility, inputs=[state_visible], outputs=[state_visible, container, button])
 
     def refresh():
         res = []
@@ -290,6 +292,7 @@ def setup_ui(ui, gallery):
 
         img_info = images[index if index >= 0 else 0]
         image = image_from_url_text(img_info)
+        geninfo, items = read_info_from_image(image)
 
         is_allowed = False
         for extra_page in ui.stored_extra_pages:
@@ -299,7 +302,12 @@ def setup_ui(ui, gallery):
 
         assert is_allowed, f'writing to {filename} is not allowed'
 
-        image.save(filename)
+        if geninfo:
+            pnginfo_data = PngImagePlugin.PngInfo()
+            pnginfo_data.add_text('parameters', geninfo)
+            image.save(filename, pnginfo=pnginfo_data)
+        else:
+            image.save(filename)
 
         return [page.create_html(ui.tabname) for page in ui.stored_extra_pages]