Commit 734986dd authored by Trung Ngo's avatar Trung Ngo
Browse files

add callback after image is saved

parent e9a410b5
Loading
Loading
Loading
Loading
+2 −1
Original line number Diff line number Diff line
@@ -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


+11 −1
Original line number Diff line number Diff line
@@ -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)