Merge pull request #10357 from catboxanon/sag

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.total_sampling_steps = total_sampling_steps

        """Total number of sampling steps planned"""

        self.output_altered = False

        """A flag for CFGDenoiser indicating whether the output has been altered by the callback"""

class UiTrainTabParams:

    def __init__(self, txt2img_preview_params):

    callbacks_image_saved=[],

    callbacks_cfg_denoiser=[],

    callbacks_cfg_denoised=[],

    callbacks_cfg_after_cfg=[],

    callbacks_before_component=[],

    callbacks_after_component=[],

    callbacks_image_grid=[],

            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:

    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:

import modules.shared as shared

from modules.script_callbacks import CFGDenoiserParams, cfg_denoiser_callback

from modules.script_callbacks import CFGDenoisedParams, cfg_denoised_callback

from modules.script_callbacks import AfterCFGCallbackParams, cfg_after_cfg_callback

samplers_k_diffusion = [

    ('Euler a', 'sample_euler_ancestral', ['k_euler_a', 'k_euler_ancestral'], {}),

            fake_uncond = torch.cat([x_out[i:i+1] for i in denoised_image_indexes])

            x_out = torch.cat([x_out, fake_uncond])  # we skipped uncond denoising, so we put cond-denoised image to where the uncond-denoised image should be

        denoised_params = CFGDenoisedParams(x_out, state.sampling_step, state.sampling_steps)

        denoised_params = CFGDenoisedParams(x_out, state.sampling_step, state.sampling_steps, self.inner_model)

        cfg_denoised_callback(denoised_params)

        devices.test_for_nans(x_out, "unet")

        if self.mask is not None:

            denoised = self.init_latent * self.mask + self.nmask * denoised

        after_cfg_callback_params = AfterCFGCallbackParams(denoised, state.sampling_step, state.sampling_steps)

        cfg_after_cfg_callback(after_cfg_callback_params)

        if after_cfg_callback_params.output_altered:

            denoised = after_cfg_callback_params.x

        self.step += 1

        return denoised