From a18e54ecd756a4101e16e42fc313df259542e07b Mon Sep 17 00:00:00 2001
From: w-e-w <40751091+w-e-w@users.noreply.github.com>
Date: Sun, 18 Feb 2024 00:38:05 +0900
Subject: [PATCH] option "open image button" open the actual dir

---
 modules/shared_options.py |  2 ++
 modules/ui_common.py      | 54 +++++++++++++++++++++++++++------------
 modules/ui_tempdir.py     | 15 +++++++++++
 3 files changed, 54 insertions(+), 17 deletions(-)

diff --git a/modules/shared_options.py b/modules/shared_options.py
index 25b47aa19..7571a7d1d 100644
--- a/modules/shared_options.py
+++ b/modules/shared_options.py
@@ -284,6 +284,8 @@ options_templates.update(options_section(('ui_gallery', "Gallery", "ui"), {
     "sd_webui_modal_lightbox_icon_opacity": OptionInfo(1, "Full page image viewer: control icon unfocused opacity", gr.Slider, {"minimum": 0.0, "maximum": 1, "step": 0.01}, onchange=shared.reload_gradio_theme).info('for mouse only').needs_reload_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(),
+    "button_open_image_actual_dir": OptionInfo(True, '"Open images output directory" button opens the actual directory of the image rather than the output root folder'),
+    "button_open_image_actual_dir_temp": OptionInfo(False, '"Open images output directory" button opens the actual directory even for temp images'),
 }))
 
 options_templates.update(options_section(('ui_alternatives', "UI alternatives", "ui"), {
diff --git a/modules/ui_common.py b/modules/ui_common.py
index 29fe7d0e9..78481c6fb 100644
--- a/modules/ui_common.py
+++ b/modules/ui_common.py
@@ -9,7 +9,7 @@ import sys
 import gradio as gr
 import subprocess as sp
 
-from modules import call_queue, shared
+from modules import call_queue, shared, ui_tempdir
 from modules.infotext_utils import image_from_url_text
 import modules.images
 from modules.ui_components import ToolButton
@@ -164,29 +164,45 @@ class OutputPanel:
 def create_output_panel(tabname, outdir, toprow=None):
     res = OutputPanel()
 
-    def open_folder(f):
+    def open_folder(f, images=None, index=None):
+        if shared.cmd_opts.hide_ui_dir_config:
+            return
+
+        try:
+            if shared.opts.button_open_image_actual_dir and 0 <= index < len(images):
+                image = images[index]
+                image_path = image["name"].rsplit('?', 1)[0]
+                image_dir = os.path.split(image_path)[0]
+                if shared.opts.button_open_image_actual_dir_temp or not ui_tempdir.is_gradio_temp_path(image_dir):
+                    f = image_dir
+        except Exception:
+            pass
+
         if not os.path.exists(f):
-            print(f'Folder "{f}" does not exist. After you create an image, the folder will be created.')
+            msg = f'Folder "{f}" does not exist. After you create an image, the folder will be created.'
+            print(msg)
+            gr.Info(msg)
             return
         elif not os.path.isdir(f):
-            print(f"""
+            msg = f"""
 WARNING
 An open_folder request was made with an argument that is not a folder.
 This could be an error or a malicious attempt to run code on your computer.
 Requested path was: {f}
-""", file=sys.stderr)
+"""
+            print(msg, file=sys.stderr)
+            gr.Warning(msg)
             return
 
-        if not shared.cmd_opts.hide_ui_dir_config:
-            path = os.path.normpath(f)
-            if platform.system() == "Windows":
-                os.startfile(path)
-            elif platform.system() == "Darwin":
-                sp.Popen(["open", path])
-            elif "microsoft-standard-WSL2" in platform.uname().release:
-                sp.Popen(["wsl-open", path])
-            else:
-                sp.Popen(["xdg-open", path])
+        path = os.path.normpath(f)
+        if platform.system() == "Windows":
+            os.startfile(path)
+        elif platform.system() == "Darwin":
+            sp.Popen(["open", path])
+        elif "microsoft-standard-WSL2" in platform.uname().release:
+            sp.Popen(["wsl-open", path])
+        else:
+            sp.Popen(["xdg-open", path])
 
     with gr.Column(elem_id=f"{tabname}_results"):
         if toprow:
@@ -213,8 +229,12 @@ Requested path was: {f}
                     res.button_upscale = ToolButton('✨', elem_id=f'{tabname}_upscale', tooltip="Create an upscaled version of the current image using hires fix settings.")
 
             open_folder_button.click(
-                fn=lambda: open_folder(shared.opts.outdir_samples or outdir),
-                inputs=[],
+                fn=lambda images, index: open_folder(shared.opts.outdir_samples or outdir, images, index),
+                _js="(y, w) => [y, selected_gallery_index()]",
+                inputs=[
+                    res.gallery,
+                    open_folder_button,  # placeholder for index
+                ],
                 outputs=[],
             )
 
diff --git a/modules/ui_tempdir.py b/modules/ui_tempdir.py
index 91f40ea42..621ed1eca 100644
--- a/modules/ui_tempdir.py
+++ b/modules/ui_tempdir.py
@@ -81,3 +81,18 @@ def cleanup_tmpdr():
 
             filename = os.path.join(root, name)
             os.remove(filename)
+
+
+def is_gradio_temp_path(path):
+    """
+    Check if the path is a temp dir used by gradio
+    """
+    path = Path(path)
+    if shared.opts.temp_dir and path.is_relative_to(shared.opts.temp_dir):
+        return True
+    if gradio_temp_dir := os.environ.get("GRADIO_TEMP_DIR"):
+        if path.is_relative_to(gradio_temp_dir):
+            return True
+    if path.is_relative_to(Path(tempfile.gettempdir()) / "gradio"):
+        return True
+    return False