diff --git a/modules/script_callbacks.py b/modules/script_callbacks.py index ce2646906..da88635b1 100644 --- a/modules/script_callbacks.py +++ b/modules/script_callbacks.py @@ -26,6 +26,24 @@ class ImageSaveParams: """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): + 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""" + + ScriptCallback = namedtuple("ScriptCallback", ["script", "callback"]) callbacks_app_started = [] callbacks_model_loaded = [] @@ -33,6 +51,7 @@ callbacks_ui_tabs = [] callbacks_ui_settings = [] callbacks_before_image_saved = [] callbacks_image_saved = [] +callbacks_cfg_denoiser = [] def clear_callbacks(): @@ -41,7 +60,7 @@ def clear_callbacks(): callbacks_ui_settings.clear() callbacks_before_image_saved.clear() callbacks_image_saved.clear() - + callbacks_cfg_denoiser.clear() def app_started_callback(demo: Blocks, app: FastAPI): for c in callbacks_app_started: @@ -95,6 +114,14 @@ def image_saved_callback(params: ImageSaveParams): report_exception(c, 'image_saved_callback') +def cfg_denoiser_callback(params: CFGDenoiserParams): + for c in callbacks_cfg_denoiser: + try: + c.callback(params) + except Exception: + report_exception(c, 'cfg_denoiser_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' @@ -147,3 +174,12 @@ def on_image_saved(callback): - params: ImageSaveParams - parameters the image was saved with. Changing fields in this object does nothing. """ add_callback(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(callbacks_cfg_denoiser, callback) + diff --git a/modules/sd_samplers.py b/modules/sd_samplers.py index 8772db564..44d4c189b 100644 --- a/modules/sd_samplers.py +++ b/modules/sd_samplers.py @@ -12,6 +12,7 @@ from modules import prompt_parser, devices, processing, images from modules.shared import opts, cmd_opts, state import modules.shared as shared +from modules.script_callbacks import CFGDenoiserParams, cfg_denoiser_callback SamplerData = namedtuple('SamplerData', ['name', 'constructor', 'aliases', 'options']) @@ -280,6 +281,12 @@ class CFGDenoiser(torch.nn.Module): image_cond_in = torch.cat([torch.stack([image_cond[i] for _ in range(n)]) for i, n in enumerate(repeats)] + [image_cond]) sigma_in = torch.cat([torch.stack([sigma[i] for _ in range(n)]) for i, n in enumerate(repeats)] + [sigma]) + denoiser_params = CFGDenoiserParams(x_in, image_cond_in, sigma_in, state.sampling_step, state.sampling_steps) + cfg_denoiser_callback(denoiser_params) + x_in = denoiser_params.x + image_cond_in = denoiser_params.image_cond + sigma_in = denoiser_params.sigma + if tensor.shape[1] == uncond.shape[1]: cond_in = torch.cat([tensor, uncond])