add callback after image is saved

This commit is contained in:
Trung Ngo 2022-10-24 01:17:09 -05:00
parent e9a410b535
commit 734986dde3
2 changed files with 13 additions and 2 deletions

View file

@ -12,7 +12,7 @@ from PIL import Image, ImageFont, ImageDraw, PngImagePlugin
from fonts.ttf import Roboto from fonts.ttf import Roboto
import string import string
from modules import sd_samplers, shared from modules import sd_samplers, shared, script_callbacks
from modules.shared import opts, cmd_opts from modules.shared import opts, cmd_opts
LANCZOS = (Image.Resampling.LANCZOS if hasattr(Image, 'Resampling') else Image.LANCZOS) 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: else:
txt_fullfn = None txt_fullfn = None
script_callbacks.image_saved_callback(image, p, fullfn, txt_fullfn)
return fullfn, txt_fullfn return fullfn, txt_fullfn

View file

@ -2,11 +2,12 @@
callbacks_model_loaded = [] callbacks_model_loaded = []
callbacks_ui_tabs = [] callbacks_ui_tabs = []
callbacks_ui_settings = [] callbacks_ui_settings = []
callbacks_image_saved = []
def clear_callbacks(): def clear_callbacks():
callbacks_model_loaded.clear() callbacks_model_loaded.clear()
callbacks_ui_tabs.clear() callbacks_ui_tabs.clear()
callbacks_image_saved.clear()
def model_loaded_callback(sd_model): def model_loaded_callback(sd_model):
@ -28,6 +29,10 @@ def ui_settings_callback():
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): def on_model_loaded(callback):
"""register a function to be called when the stable diffusion model is created; the model is """register a function to be called when the stable diffusion model is created; the model is
passed as an argument""" 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 """register a function to be called before UI settings are populated; add your settings
by using shared.opts.add_option(shared.OptionInfo(...)) """ by using shared.opts.add_option(shared.OptionInfo(...)) """
callbacks_ui_settings.append(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)