From abc4d3a6934f2024b7e8c217ae8d5c90bf7c20ed Mon Sep 17 00:00:00 2001
From: SirFrags <85357620+SirFrags@users.noreply.github.com>
Date: Sat, 25 Mar 2023 15:59:50 -0400
Subject: [PATCH] preview replace save params

---
 modules/ui_extra_networks.py | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/modules/ui_extra_networks.py b/modules/ui_extra_networks.py
index daea03d6..aaa93969 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
@@ -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]