import dataclasses import inspect import os from collections import namedtuple from typing import Optional, Any from fastapi import FastAPI from gradio import Blocks from modules import errors, timer def report_exception(c, job): errors.report(f"Error executing callback {job} for {c.script}", exc_info=True) 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'""" class ExtraNoiseParams: def __init__(self, noise, x, xi): self.noise = noise """Random noise generated by the seed""" self.x = x """Latent representation of the image""" self.xi = xi """Noisy latent representation of the image""" class CFGDenoiserParams: def __init__(self, x, image_cond, sigma, sampling_step, total_sampling_steps, text_cond, text_uncond, denoiser=None): self.x = x """Latent image representation in the process of being denoised""" self.image_cond = image_cond """Conditioning image""" self.sigma = sigma """Current sigma noise step value""" self.sampling_step = sampling_step """Current Sampling step number""" self.total_sampling_steps = total_sampling_steps """Total number of sampling steps planned""" self.text_cond = text_cond """ Encoder hidden states of text conditioning from prompt""" self.text_uncond = text_uncond """ Encoder hidden states of text conditioning from negative prompt""" self.denoiser = denoiser """Current CFGDenoiser object with processing parameters""" class CFGDenoisedParams: def __init__(self, x, sampling_step, total_sampling_steps, inner_model): self.x = x """Latent image representation in the process of being denoised""" self.sampling_step = sampling_step """Current Sampling step number""" self.total_sampling_steps = total_sampling_steps """Total number of sampling steps planned""" self.inner_model = inner_model """Inner model reference used for denoising""" class AfterCFGCallbackParams: def __init__(self, x, sampling_step, total_sampling_steps): self.x = x """Latent image representation in the process of being denoised""" self.sampling_step = sampling_step """Current Sampling step number""" self.total_sampling_steps = total_sampling_steps """Total number of sampling steps planned""" class UiTrainTabParams: def __init__(self, txt2img_preview_params): self.txt2img_preview_params = txt2img_preview_params class ImageGridLoopParams: def __init__(self, imgs, cols, rows): self.imgs = imgs self.cols = cols self.rows = rows @dataclasses.dataclass class BeforeTokenCounterParams: prompt: str steps: int styles: list is_positive: bool = True ScriptCallback = namedtuple("ScriptCallback", ["script", "callback"]) callback_map = dict( callbacks_app_started=[], callbacks_model_loaded=[], callbacks_ui_tabs=[], callbacks_ui_train_tabs=[], callbacks_ui_settings=[], callbacks_before_image_saved=[], callbacks_image_saved=[], callbacks_extra_noise=[], callbacks_cfg_denoiser=[], callbacks_cfg_denoised=[], callbacks_cfg_after_cfg=[], callbacks_before_component=[], callbacks_after_component=[], callbacks_image_grid=[], callbacks_infotext_pasted=[], callbacks_script_unloaded=[], callbacks_before_ui=[], callbacks_on_reload=[], callbacks_list_optimizers=[], callbacks_list_unets=[], callbacks_before_token_counter=[], ) def clear_callbacks(): for callback_list in callback_map.values(): callback_list.clear() def app_started_callback(demo: Optional[Blocks], app: FastAPI): for c in callback_map['callbacks_app_started']: try: c.callback(demo, app) timer.startup_timer.record(os.path.basename(c.script)) except Exception: report_exception(c, 'app_started_callback') def app_reload_callback(): for c in callback_map['callbacks_on_reload']: try: c.callback() except Exception: report_exception(c, 'callbacks_on_reload') def model_loaded_callback(sd_model): for c in callback_map['callbacks_model_loaded']: try: c.callback(sd_model) except Exception: report_exception(c, 'model_loaded_callback') def ui_tabs_callback(): res = [] for c in callback_map['callbacks_ui_tabs']: try: res += c.callback() or [] except Exception: report_exception(c, 'ui_tabs_callback') return res def ui_train_tabs_callback(params: UiTrainTabParams): for c in callback_map['callbacks_ui_train_tabs']: try: c.callback(params) except Exception: report_exception(c, 'callbacks_ui_train_tabs') def ui_settings_callback(): for c in callback_map['callbacks_ui_settings']: try: c.callback() except Exception: report_exception(c, 'ui_settings_callback') def before_image_saved_callback(params: ImageSaveParams): for c in callback_map['callbacks_before_image_saved']: try: c.callback(params) except Exception: report_exception(c, 'before_image_saved_callback') def image_saved_callback(params: ImageSaveParams): for c in callback_map['callbacks_image_saved']: try: c.callback(params) except Exception: report_exception(c, 'image_saved_callback') def extra_noise_callback(params: ExtraNoiseParams): for c in callback_map['callbacks_extra_noise']: try: c.callback(params) except Exception: report_exception(c, 'callbacks_extra_noise') def cfg_denoiser_callback(params: CFGDenoiserParams): for c in callback_map['callbacks_cfg_denoiser']: try: c.callback(params) except Exception: report_exception(c, 'cfg_denoiser_callback') def cfg_denoised_callback(params: CFGDenoisedParams): for c in callback_map['callbacks_cfg_denoised']: try: c.callback(params) except Exception: report_exception(c, 'cfg_denoised_callback') def cfg_after_cfg_callback(params: AfterCFGCallbackParams): for c in callback_map['callbacks_cfg_after_cfg']: try: c.callback(params) except Exception: report_exception(c, 'cfg_after_cfg_callback') def before_component_callback(component, **kwargs): for c in callback_map['callbacks_before_component']: try: c.callback(component, **kwargs) except Exception: report_exception(c, 'before_component_callback') def after_component_callback(component, **kwargs): for c in callback_map['callbacks_after_component']: try: c.callback(component, **kwargs) except Exception: report_exception(c, 'after_component_callback') def image_grid_callback(params: ImageGridLoopParams): for c in callback_map['callbacks_image_grid']: try: c.callback(params) except Exception: report_exception(c, 'image_grid') def infotext_pasted_callback(infotext: str, params: dict[str, Any]): for c in callback_map['callbacks_infotext_pasted']: try: c.callback(infotext, params) except Exception: report_exception(c, 'infotext_pasted') def script_unloaded_callback(): for c in reversed(callback_map['callbacks_script_unloaded']): try: c.callback() except Exception: report_exception(c, 'script_unloaded') def before_ui_callback(): for c in reversed(callback_map['callbacks_before_ui']): try: c.callback() except Exception: report_exception(c, 'before_ui') def list_optimizers_callback(): res = [] for c in callback_map['callbacks_list_optimizers']: try: c.callback(res) except Exception: report_exception(c, 'list_optimizers') return res def list_unets_callback(): res = [] for c in callback_map['callbacks_list_unets']: try: c.callback(res) except Exception: report_exception(c, 'list_unets') return res def before_token_counter_callback(params: BeforeTokenCounterParams): for c in callback_map['callbacks_before_token_counter']: try: c.callback(params) except Exception: report_exception(c, 'before_token_counter') def add_callback(callbacks, fun): stack = [x for x in inspect.stack() if x.filename != __file__] filename = stack[0].filename if stack else 'unknown file' callbacks.append(ScriptCallback(filename, fun)) def remove_current_script_callbacks(): stack = [x for x in inspect.stack() if x.filename != __file__] filename = stack[0].filename if stack else 'unknown file' if filename == 'unknown file': return for callback_list in callback_map.values(): for callback_to_remove in [cb for cb in callback_list if cb.script == filename]: callback_list.remove(callback_to_remove) def remove_callbacks_for_function(callback_func): for callback_list in callback_map.values(): for callback_to_remove in [cb for cb in callback_list if cb.callback == callback_func]: callback_list.remove(callback_to_remove) def on_app_started(callback): """register a function to be called when the webui started, the gradio `Block` component and fastapi `FastAPI` object are passed as the arguments""" add_callback(callback_map['callbacks_app_started'], callback) def on_before_reload(callback): """register a function to be called just before the server reloads.""" add_callback(callback_map['callbacks_on_reload'], callback) 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; this function is also called when the script is reloaded. """ add_callback(callback_map['callbacks_model_loaded'], callback) def on_ui_tabs(callback): """register a function to be called when the UI is creating new tabs. The function must either return a None, which means no new tabs to be added, or a list, where each element is a tuple: (gradio_component, title, elem_id) gradio_component is a gradio component to be used for contents of the tab (usually gr.Blocks) title is tab text displayed to user in the UI elem_id is HTML id for the tab """ add_callback(callback_map['callbacks_ui_tabs'], callback) def on_ui_train_tabs(callback): """register a function to be called when the UI is creating new tabs for the train tab. Create your new tabs with gr.Tab. """ add_callback(callback_map['callbacks_ui_train_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(...)) """ add_callback(callback_map['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(callback_map['callbacks_before_image_saved'], callback) def on_image_saved(callback): """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(callback_map['callbacks_image_saved'], callback) def on_extra_noise(callback): """register a function to be called before adding extra noise in img2img or hires fix; The callback is called with one argument: - params: ExtraNoiseParams - contains noise determined by seed and latent representation of image """ add_callback(callback_map['callbacks_extra_noise'], callback) def on_cfg_denoiser(callback): """register a function to be called in the kdiffussion cfg_denoiser method after building the inner model inputs. The callback is called with one argument: - params: CFGDenoiserParams - parameters to be passed to the inner model and sampling state details. """ add_callback(callback_map['callbacks_cfg_denoiser'], callback) def on_cfg_denoised(callback): """register a function to be called in the kdiffussion cfg_denoiser method after building the inner model inputs. The callback is called with one argument: - params: CFGDenoisedParams - parameters to be passed to the inner model and sampling state details. """ add_callback(callback_map['callbacks_cfg_denoised'], callback) def on_cfg_after_cfg(callback): """register a function to be called in the kdiffussion cfg_denoiser method after cfg calculations are completed. The callback is called with one argument: - params: AfterCFGCallbackParams - parameters to be passed to the script for post-processing after cfg calculation. """ add_callback(callback_map['callbacks_cfg_after_cfg'], callback) def on_before_component(callback): """register a function to be called before a component is created. The callback is called with arguments: - component - gradio component that is about to be created. - **kwargs - args to gradio.components.IOComponent.__init__ function Use elem_id/label fields of kwargs to figure out which component it is. This can be useful to inject your own components somewhere in the middle of vanilla UI. """ add_callback(callback_map['callbacks_before_component'], callback) def on_after_component(callback): """register a function to be called after a component is created. See on_before_component for more.""" add_callback(callback_map['callbacks_after_component'], callback) def on_image_grid(callback): """register a function to be called before making an image grid. The callback is called with one argument: - params: ImageGridLoopParams - parameters to be used for grid creation. Can be modified. """ add_callback(callback_map['callbacks_image_grid'], callback) def on_infotext_pasted(callback): """register a function to be called before applying an infotext. The callback is called with two arguments: - infotext: str - raw infotext. - result: dict[str, any] - parsed infotext parameters. """ add_callback(callback_map['callbacks_infotext_pasted'], callback) def on_script_unloaded(callback): """register a function to be called before the script is unloaded. Any hooks/hijacks/monkeying about that the script did should be reverted here""" add_callback(callback_map['callbacks_script_unloaded'], callback) def on_before_ui(callback): """register a function to be called before the UI is created.""" add_callback(callback_map['callbacks_before_ui'], callback) def on_list_optimizers(callback): """register a function to be called when UI is making a list of cross attention optimization options. The function will be called with one argument, a list, and shall add objects of type modules.sd_hijack_optimizations.SdOptimization to it.""" add_callback(callback_map['callbacks_list_optimizers'], callback) def on_list_unets(callback): """register a function to be called when UI is making a list of alternative options for unet. The function will be called with one argument, a list, and shall add objects of type modules.sd_unet.SdUnetOption to it.""" add_callback(callback_map['callbacks_list_unets'], callback) def on_before_token_counter(callback): """register a function to be called when UI is counting tokens for a prompt. The function will be called with one argument of type BeforeTokenCounterParams, and should modify its fields if necessary.""" add_callback(callback_map['callbacks_before_token_counter'], callback)