From 734986dde3231416813f827242c111da212b2ccb Mon Sep 17 00:00:00 2001 From: Trung Ngo Date: Mon, 24 Oct 2022 01:17:09 -0500 Subject: add callback after image is saved --- modules/script_callbacks.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) (limited to 'modules/script_callbacks.py') 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.3 From 77a320f406a76425176b8ca4c034c362b6734713 Mon Sep 17 00:00:00 2001 From: AUTOMATIC <16777216c@gmail.com> Date: Mon, 24 Oct 2022 17:23:51 +0300 Subject: do not stop execution when script's callback misbehaves and report which script it was --- modules/script_callbacks.py | 45 ++++++++++++++++++++++++++++++++++++--------- 1 file changed, 36 insertions(+), 9 deletions(-) (limited to 'modules/script_callbacks.py') diff --git a/modules/script_callbacks.py b/modules/script_callbacks.py index f46d3d9a..5e9a6d4b 100644 --- a/modules/script_callbacks.py +++ b/modules/script_callbacks.py @@ -1,4 +1,15 @@ +import sys +import traceback +from collections import namedtuple +import inspect + +def report_exception(c, job): + print(f"Error executing callback {job} for {c.script}", file=sys.stderr) + print(traceback.format_exc(), file=sys.stderr) + + +ScriptCallback = namedtuple("ScriptCallback", ["script", "callback"]) callbacks_model_loaded = [] callbacks_ui_tabs = [] callbacks_ui_settings = [] @@ -10,28 +21,44 @@ def clear_callbacks(): def model_loaded_callback(sd_model): - for callback in callbacks_model_loaded: - callback(sd_model) + for c in callbacks_model_loaded: + try: + c.callback(sd_model) + except Exception: + report_exception(c, 'model_loaded_callback') def ui_tabs_callback(): res = [] - for callback in callbacks_ui_tabs: - res += callback() or [] + for c in callbacks_ui_tabs: + try: + res += c.callback() or [] + except Exception: + report_exception(c, 'ui_tabs_callback') return res def ui_settings_callback(): - for callback in callbacks_ui_settings: - callback() + for c in callbacks_ui_settings: + try: + c.callback() + except Exception: + report_exception(c, 'ui_settings_callback') + + +def add_callback(callbacks, fun): + stack = [x for x in inspect.stack() if x.filename != __file__] + filename = stack[0].filename if len(stack) > 0 else 'unknown file' + + callbacks.append(ScriptCallback(filename, fun)) 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""" - callbacks_model_loaded.append(callback) + add_callback(callbacks_model_loaded, callback) def on_ui_tabs(callback): @@ -44,10 +71,10 @@ def on_ui_tabs(callback): title is tab text displayed to user in the UI elem_id is HTML id for the tab """ - callbacks_ui_tabs.append(callback) + add_callback(callbacks_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(...)) """ - callbacks_ui_settings.append(callback) + add_callback(callbacks_ui_settings, callback) -- cgit v1.2.3 From 3e15f8e0f5cc87507f77546d92435670644dbd18 Mon Sep 17 00:00:00 2001 From: AUTOMATIC <16777216c@gmail.com> Date: Tue, 25 Oct 2022 12:16:17 +0300 Subject: update callbacks code for #3549 --- modules/script_callbacks.py | 22 ++++++++++++++++------ 1 file changed, 16 insertions(+), 6 deletions(-) (limited to 'modules/script_callbacks.py') diff --git a/modules/script_callbacks.py b/modules/script_callbacks.py index 9933fa38..dc520abc 100644 --- a/modules/script_callbacks.py +++ b/modules/script_callbacks.py @@ -49,6 +49,14 @@ def ui_settings_callback(): report_exception(c, 'ui_settings_callback') +def image_saved_callback(image, p, fullfn, txt_fullfn): + for c in callbacks_image_saved: + try: + c.callback(image, p, fullfn, txt_fullfn) + except Exception: + report_exception(c, 'image_saved_callback') + + def add_callback(callbacks, fun): stack = [x for x in inspect.stack() if x.filename != __file__] filename = stack[0].filename if len(stack) > 0 else 'unknown file' @@ -56,9 +64,6 @@ 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 @@ -82,9 +87,14 @@ 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(...)) """ - callbacks_ui_settings.append(callback) + add_callback(callbacks_ui_settings, 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) + """register a function to be called after modules.images.save_image is called. + The callback is called with three arguments: + - p - procesing object (or a dummy object with same fields if the image is saved using save button) + - fullfn - image filename + - txt_fullfn - text file with parameters; may be None + """ + add_callback(callbacks_image_saved, callback) -- cgit v1.2.3 From 9d82c351ac36d1511f5f65b24443c60250ee3e9e Mon Sep 17 00:00:00 2001 From: AUTOMATIC <16777216c@gmail.com> Date: Wed, 26 Oct 2022 09:56:25 +0300 Subject: fix typo in on_save_imaged/on_image_saved; hope no extension is using it yet --- modules/script_callbacks.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) (limited to 'modules/script_callbacks.py') diff --git a/modules/script_callbacks.py b/modules/script_callbacks.py index dc520abc..6803d57b 100644 --- a/modules/script_callbacks.py +++ b/modules/script_callbacks.py @@ -90,7 +90,7 @@ def on_ui_settings(callback): add_callback(callbacks_ui_settings, callback) -def on_save_imaged(callback): +def on_image_saved(callback): """register a function to be called after modules.images.save_image is called. The callback is called with three arguments: - p - procesing object (or a dummy object with same fields if the image is saved using save button) -- cgit v1.2.3 From 0cd74602531a40f72d1a75b471a8a9166135d333 Mon Sep 17 00:00:00 2001 From: AUTOMATIC <16777216c@gmail.com> Date: Wed, 26 Oct 2022 13:12:44 +0300 Subject: add script callback for before image save and change callback for after image save to use a class with parameters --- modules/script_callbacks.py | 48 +++++++++++++++++++++++++++++++++++++-------- 1 file changed, 40 insertions(+), 8 deletions(-) (limited to 'modules/script_callbacks.py') diff --git a/modules/script_callbacks.py b/modules/script_callbacks.py index 6803d57b..6ea58d61 100644 --- a/modules/script_callbacks.py +++ b/modules/script_callbacks.py @@ -9,15 +9,34 @@ def report_exception(c, job): print(traceback.format_exc(), file=sys.stderr) +class ImageSaveParams: + def __init__(self, image, p, filename, pnginfo): + self.image = image + """the PIL image itself""" + + self.p = p + """p object with processing parameters; either StableDiffusionProcessing or an object with same fields""" + + self.filename = filename + """name of file that the image would be saved to""" + + self.pnginfo = pnginfo + """dictionary with parameters for image's PNG info data; infotext will have the key 'parameters'""" + + ScriptCallback = namedtuple("ScriptCallback", ["script", "callback"]) callbacks_model_loaded = [] callbacks_ui_tabs = [] callbacks_ui_settings = [] +callbacks_before_image_saved = [] callbacks_image_saved = [] + def clear_callbacks(): callbacks_model_loaded.clear() callbacks_ui_tabs.clear() + callbacks_ui_settings.clear() + callbacks_before_image_saved.clear() callbacks_image_saved.clear() @@ -49,10 +68,18 @@ def ui_settings_callback(): report_exception(c, 'ui_settings_callback') -def image_saved_callback(image, p, fullfn, txt_fullfn): +def before_image_saved_callback(params: ImageSaveParams): for c in callbacks_image_saved: try: - c.callback(image, p, fullfn, txt_fullfn) + c.callback(params) + except Exception: + report_exception(c, 'before_image_saved_callback') + + +def image_saved_callback(params: ImageSaveParams): + for c in callbacks_image_saved: + try: + c.callback(params) except Exception: report_exception(c, 'image_saved_callback') @@ -64,7 +91,6 @@ def add_callback(callbacks, fun): callbacks.append(ScriptCallback(filename, fun)) - 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""" @@ -90,11 +116,17 @@ def on_ui_settings(callback): add_callback(callbacks_ui_settings, callback) +def on_before_image_saved(callback): + """register a function to be called before an image is saved to a file. + The callback is called with one argument: + - params: ImageSaveParams - parameters the image is to be saved with. You can change fields in this object. + """ + add_callback(callbacks_before_image_saved, callback) + + def on_image_saved(callback): - """register a function to be called after modules.images.save_image is called. - The callback is called with three arguments: - - p - procesing object (or a dummy object with same fields if the image is saved using save button) - - fullfn - image filename - - txt_fullfn - text file with parameters; may be None + """register a function to be called after an image is saved to a file. + The callback is called with one argument: + - params: ImageSaveParams - parameters the image was saved with. Changing fields in this object does nothing. """ add_callback(callbacks_image_saved, callback) -- cgit v1.2.3