Merge pull request #3549 from tsngo/on-image-saved-callback

add callback after image is saved

2 files changed, 14 insertions, 3 deletions

diff --git a/modules/images.py b/modules/images.py
index 9a8fe3ed..0e044af2 100644
--- a/modules/images.py
+++ b/modules/images.py
@@ -16,7 +16,7 @@ from PIL import Image, ImageFont, ImageDraw, PngImagePlugin
 from fonts.ttf import Roboto

 import string

 

-from modules import sd_samplers, shared

+from modules import sd_samplers, shared, script_callbacks

 from modules.shared import opts, cmd_opts

 

 LANCZOS = (Image.Resampling.LANCZOS if hasattr(Image, 'Resampling') else Image.LANCZOS)

@@ -539,6 +539,7 @@ def save_image(image, path, basename, seed=None, prompt=None, extension='png', i
     else:

         txt_fullfn = None

 

+    script_callbacks.image_saved_callback(image, p, fullfn, txt_fullfn)

     return fullfn, txt_fullfn

 

 

diff --git a/modules/script_callbacks.py b/modules/script_callbacks.py
index 5e9a6d4b..9933fa38 100644
--- a/modules/script_callbacks.py
+++ b/modules/script_callbacks.py
@@ -13,11 +13,12 @@ ScriptCallback = namedtuple("ScriptCallback", ["script", "callback"])
 callbacks_model_loaded = []

 callbacks_ui_tabs = []

 callbacks_ui_settings = []

-

+callbacks_image_saved = []

 

 def clear_callbacks():

     callbacks_model_loaded.clear()

     callbacks_ui_tabs.clear()

+    callbacks_image_saved.clear()

 

 

 def model_loaded_callback(sd_model):

@@ -55,6 +56,10 @@ def add_callback(callbacks, fun):
     callbacks.append(ScriptCallback(filename, fun))

 

 

+def image_saved_callback(image, p, fullfn, txt_fullfn):

+    for callback in callbacks_image_saved:

+        callback(image, p, fullfn, txt_fullfn)

+

 def on_model_loaded(callback):

     """register a function to be called when the stable diffusion model is created; the model is

     passed as an argument"""

@@ -77,4 +82,9 @@ def on_ui_tabs(callback):
 def on_ui_settings(callback):

     """register a function to be called before UI settings are populated; add your settings

     by using shared.opts.add_option(shared.OptionInfo(...)) """

-    add_callback(callbacks_ui_settings, callback)

+    callbacks_ui_settings.append(callback)

+

+

+def on_save_imaged(callback):

+    """register a function to call after modules.images.save_image is called returning same values, original image and p """

+    callbacks_image_saved.append(callback)