From 4d94beabe9707e1e089b52e9cef28b66db3d87fc Mon Sep 17 00:00:00 2001
From: Vladimir Mandic <mandic00@live.com>
Date: Tue, 12 Sep 2023 14:56:09 -0400
Subject: [PATCH] add callback

---
 modules/script_callbacks.py | 19 +++++++++++++++++++
 modules/ui_common.py        |  4 ++++
 2 files changed, 23 insertions(+)

diff --git a/modules/script_callbacks.py b/modules/script_callbacks.py
index 29cd9470a..28b906aaf 100644
--- a/modules/script_callbacks.py
+++ b/modules/script_callbacks.py
@@ -98,6 +98,7 @@ callback_map = dict(
     callbacks_ui_settings=[],
     callbacks_before_image_saved=[],
     callbacks_image_saved=[],
+    callbacks_image_save_btn=[],
     callbacks_cfg_denoiser=[],
     callbacks_cfg_denoised=[],
     callbacks_cfg_after_cfg=[],
@@ -205,6 +206,16 @@ def image_saved_callback(params: ImageSaveParams):
             report_exception(e, c, 'image_saved_callback')
 
 
+def image_save_btn_callback(filename: str):
+    for c in callback_map['callbacks_image_save_btn']:
+        try:
+            t0 = time.time()
+            c.callback(filename)
+            timer(t0, c.script, 'image_save_btn')
+        except Exception as e:
+            report_exception(e, c, 'image_save_btn_callback')
+
+
 def cfg_denoiser_callback(params: CFGDenoiserParams):
     for c in callback_map['callbacks_cfg_denoiser']:
         try:
@@ -376,6 +387,14 @@ def on_image_saved(callback):
     add_callback(callback_map['callbacks_image_saved'], callback)
 
 
+def on_image_save_btn(callback):
+    """register a function to be called after an image save button is pressed.
+    The callback is called with one argument:
+        - params: ImageSaveParams - parameters the image was saved with. Changing fields in this object does nothing.
+    """
+    add_callback(callback_map['callbacks_image_save_btn'], callback)
+
+
 def on_cfg_denoiser(callback):
     """register a function to be called in the kdiffussion cfg_denoiser method after building the inner model inputs.
     The callback is called with one argument:
diff --git a/modules/ui_common.py b/modules/ui_common.py
index 58ea9f8f8..dd4d2225e 100644
--- a/modules/ui_common.py
+++ b/modules/ui_common.py
@@ -9,6 +9,7 @@ from modules import call_queue, shared
 from modules.generation_parameters_copypaste import image_from_url_text
 import modules.ui_symbols as symbols
 import modules.images
+import modules.script_callbacks
 
 
 def update_generation_info(generation_info, html_info, img_index):
@@ -115,6 +116,8 @@ def save_files(js_data, images, html_info, index):
                 os.makedirs(destination, exist_ok = True)
             shutil.copy(fullfn, destination)
             shared.log.info(f"Copying image: {fullfn} -> {destination}")
+            tgt_filename = os.path.join(destination, os.path.basename(fullfn))
+            modules.script_callbacks.image_save_btn_callback(tgt_filename)
         else:
             image = image_from_url_text(filedata)
             info = p.infotexts[i + 1] if len(p.infotexts) > len(p.all_seeds) else p.infotexts[i] # infotexts may be offset by 1 because the first image is the grid
@@ -127,6 +130,7 @@ def save_files(js_data, images, html_info, index):
             if txt_fullfn:
                 filenames.append(os.path.basename(txt_fullfn))
                 fullfns.append(txt_fullfn)
+            modules.script_callbacks.image_save_btn_callback(filename)
     if shared.opts.samples_save_zip and len(fullfns) > 1:
         zip_filepath = os.path.join(shared.opts.outdir_save, "images.zip")
         from zipfile import ZipFile