From aea5b2510ed4bd9150cea67b6036c837f7df2750 Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Wed, 28 Sep 2022 17:05:23 +0300
Subject: save parameters for images when using the Save button.

---
 modules/ui.py | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

(limited to 'modules/ui.py')

diff --git a/modules/ui.py b/modules/ui.py
index 7db8edbd..484be762 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -12,7 +12,7 @@ import traceback
 
 import numpy as np
 import torch
-from PIL import Image
+from PIL import Image, PngImagePlugin
 
 import gradio as gr
 import gradio.utils
@@ -97,10 +97,11 @@ def save_files(js_data, images, index):
     filenames = []
 
     data = json.loads(js_data)
-    
-    if index > -1 and opts.save_selected_only and (index > 0 or not opts.return_grid): # ensures we are looking at a specific non-grid picture, and we have save_selected_only
+    if index > -1 and opts.save_selected_only and (index >= data["index_of_first_image"]):  # ensures we are looking at a specific non-grid picture, and we have save_selected_only
         images = [images[index]]
-        data["seed"] += (index - 1 if opts.return_grid else index)
+        infotexts = [data["infotexts"][index]]
+    else:
+        infotexts = data["infotexts"]
 
     with open(os.path.join(opts.outdir_save, "log.csv"), "a", encoding="utf8", newline='') as file:
         at_start = file.tell() == 0
@@ -116,8 +117,11 @@ def save_files(js_data, images, index):
             if filedata.startswith("data:image/png;base64,"):
                 filedata = filedata[len("data:image/png;base64,"):]
 
-            with open(filepath, "wb") as imgfile:
-                imgfile.write(base64.decodebytes(filedata.encode('utf-8')))
+            pnginfo = PngImagePlugin.PngInfo()
+            pnginfo.add_text('parameters', infotexts[i])
+
+            image = Image.open(io.BytesIO(base64.decodebytes(filedata.encode('utf-8'))))
+            image.save(filepath, quality=opts.jpeg_quality, pnginfo=pnginfo)
 
             filenames.append(filename)
 
-- 
cgit v1.2.3