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)