| import inspect | |
| import os | |
| from collections import namedtuple | |
| from typing import Optional, Dict, 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 CFGDenoiserParams: | |
| def __init__(self, x, image_cond, sigma, sampling_step, total_sampling_steps, text_cond, text_uncond): | |
| 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""" | |
| 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 | |
| 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_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=[], | |
| ) | |
| 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 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 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_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) | |