preview replace save params

1 files changed, 9 insertions, 1 deletions

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]