From f4a6a08ec2ae80305c45a3fb181e8b006828de09 Mon Sep 17 00:00:00 2001
From: DenOfEquity <166248528+DenOfEquity@users.noreply.github.com>
Date: Sat, 16 Nov 2024 12:27:20 +0000
Subject: [PATCH] Chasing save issues (#2331)

* already_saved_as attribute for gallery items after hiresfix quickbutton
* different method to select image in gallery after hires quickbutton with insert option to avoid other changes
---
 modules/images.py  | 4 ++--
 modules/txt2img.py | 3 ++-
 modules/ui.py      | 8 +++++++-
 3 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/modules/images.py b/modules/images.py
index 2441940f..d6488cec 100644
--- a/modules/images.py
+++ b/modules/images.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-import datetime, time
+import datetime
 import functools
 import pytz
 import io
@@ -740,7 +740,7 @@ def save_image(image, path, basename, seed=None, prompt=None, extension='png', i
 
     fullfn_without_extension = _atomically_save_image(image, fullfn_without_extension, extension)
     fullfn = fullfn_without_extension + extension
-    image.already_saved_as = f"{fullfn}\?{time.process_time_ns()}"
+    image.already_saved_as = fullfn
 
     oversize = image.width > opts.target_side_length or image.height > opts.target_side_length
     if opts.export_for_4chan and (oversize or os.stat(fullfn).st_size > opts.img_downscale_threshold * 1024 * 1024):
diff --git a/modules/txt2img.py b/modules/txt2img.py
index f071233e..a1971dce 100644
--- a/modules/txt2img.py
+++ b/modules/txt2img.py
@@ -106,6 +106,7 @@ def txt2img_upscale_function(id_task: str, request: gr.Request, gallery, gallery
     new_gallery = []
     for i, image in enumerate(gallery):
         if insert or i != gallery_index:
+            image[0].already_saved_as = image[0].filename.rsplit('?', 1)[0]
             new_gallery.append(image)
         if i == gallery_index:
             new_gallery.extend(processed.images)
@@ -117,7 +118,7 @@ def txt2img_upscale_function(id_task: str, request: gr.Request, gallery, gallery
     else:
         geninfo["infotexts"][gallery_index] = processed.info
 
-    return gr.Gallery(new_gallery, selected_index=new_index), json.dumps(geninfo), plaintext_to_html(processed.info), plaintext_to_html(processed.comments, classname="comments")
+    return new_gallery, json.dumps(geninfo), plaintext_to_html(processed.info), plaintext_to_html(processed.comments, classname="comments")
 
 
 def txt2img_function(id_task: str, request: gr.Request, *args):
diff --git a/modules/ui.py b/modules/ui.py
index b65009a0..7bc8dd94 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -470,6 +470,12 @@ def create_ui():
             toprow.prompt.submit(**txt2img_args)
             toprow.submit.click(**txt2img_args)
 
+            def select_gallery_image(index):
+                index = int(index)
+                if getattr(shared.opts, 'hires_button_gallery_insert', False):
+                    index += 1
+                return gr.update(selected_index=index)
+            
             txt2img_upscale_inputs = txt2img_inputs[0:1] + [output_panel.gallery, dummy_component_number, output_panel.generation_info] + txt2img_inputs[1:]
             output_panel.button_upscale.click(
                 fn=wrap_gradio_gpu_call(modules.txt2img.txt2img_upscale, extra_outputs=[None, '', '']),
@@ -477,7 +483,7 @@ def create_ui():
                 inputs=txt2img_upscale_inputs,
                 outputs=txt2img_outputs,
                 show_progress=False,
-            )
+            ).then(fn=select_gallery_image, js="selected_gallery_index", inputs=[dummy_component], outputs=[output_panel.gallery])
 
             res_switch_btn.click(lambda w, h: (h, w), inputs=[width, height], outputs=[width, height], show_progress=False)