From 95021aebec79ac12d222e8e61622c5951add1f5d Mon Sep 17 00:00:00 2001
From: DenOfEquity <166248528+DenOfEquity@users.noreply.github.com>
Date: Fri, 15 Nov 2024 12:16:29 +0000
Subject: [PATCH] Hiresfix quickbutton insert option (#2325)

add gallery insert option for hiresfix quickbutton, and select the new image in the new gallery
---
 modules/shared_options.py |  1 +
 modules/txt2img.py        | 16 +++++++++++-----
 2 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/modules/shared_options.py b/modules/shared_options.py
index 4049758b..c205da1a 100644
--- a/modules/shared_options.py
+++ b/modules/shared_options.py
@@ -319,6 +319,7 @@ options_templates.update(options_section(('ui_gallery', "Gallery", "ui"), {
     "sd_webui_modal_lightbox_toolbar_opacity": OptionInfo(0.9, "Full page image viewer: tool bar opacity", gr.Slider, {"minimum": 0.0, "maximum": 1, "step": 0.01}, onchange=shared.reload_gradio_theme).info('for mouse only').needs_reload_ui(),
     "gallery_height": OptionInfo("", "Gallery height", gr.Textbox).info("can be any valid CSS value, for example 768px or 20em").needs_reload_ui(),
     "open_dir_button_choice": OptionInfo("Subdirectory", "What directory the [📂] button opens", gr.Radio, {"choices": ["Output Root", "Subdirectory", "Subdirectory (even temp dir)"]}),
+    "hires_button_gallery_insert": OptionInfo(False, "Insert [✨] hires button results into gallery").info("Default: original image will be replaced"),
 }))
 
 options_templates.update(options_section(('ui_alternatives', "UI alternatives", "ui"), {
diff --git a/modules/txt2img.py b/modules/txt2img.py
index aceca949..f071233e 100644
--- a/modules/txt2img.py
+++ b/modules/txt2img.py
@@ -102,16 +102,22 @@ def txt2img_upscale_function(id_task: str, request: gr.Request, gallery, gallery
 
     shared.total_tqdm.clear()
 
+    insert = getattr(shared.opts, 'hires_button_gallery_insert', False)
     new_gallery = []
     for i, image in enumerate(gallery):
+        if insert or i != gallery_index:
+            new_gallery.append(image)
         if i == gallery_index:
             new_gallery.extend(processed.images)
-        else:
-            new_gallery.append(image)
+        
+    new_index = gallery_index
+    if insert:
+        new_index += 1
+        geninfo["infotexts"].insert(new_index, processed.info)
+    else:
+        geninfo["infotexts"][gallery_index] = processed.info
 
-    geninfo["infotexts"][gallery_index] = processed.info
-
-    return new_gallery, json.dumps(geninfo), plaintext_to_html(processed.info), plaintext_to_html(processed.comments, classname="comments")
+    return gr.Gallery(new_gallery, selected_index=new_index), 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):