From 734986dde3231416813f827242c111da212b2ccb Mon Sep 17 00:00:00 2001
From: Trung Ngo <codem01@gmail.com>
Date: Mon, 24 Oct 2022 01:17:09 -0500
Subject: add callback after image is saved

---
 modules/images.py           |  3 ++-
 modules/script_callbacks.py | 12 +++++++++++-
 2 files changed, 13 insertions(+), 2 deletions(-)

(limited to 'modules')

diff --git a/modules/images.py b/modules/images.py
index b9589563..01c60f89 100644
--- a/modules/images.py
+++ b/modules/images.py
@@ -12,7 +12,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)
@@ -467,6 +467,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 5bcccd67..5836e4b9 100644
--- a/modules/script_callbacks.py
+++ b/modules/script_callbacks.py
@@ -2,11 +2,12 @@
 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):
@@ -28,6 +29,10 @@ def ui_settings_callback():
         callback()
 
 
+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"""
@@ -51,3 +56,8 @@ 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(...)) """
     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)
-- 
cgit v1.2.1