Organized the settings and UI of soft inpainting to allow for toggling the...

import modules.processing as processing

from modules.ui import plaintext_to_html

import modules.scripts

import modules.soft_inpainting as si

def process_batch(p, input_dir, output_dir, inpaint_mask_dir, args, to_scale=False, scale_by=1.0, use_png_info=False, png_info_props=None, png_info_dir=None):

            sampler_name: str,

            mask_blur: int,

            mask_alpha: float,

            mask_blend_enabled: bool,

            mask_blend_power: float,

            mask_blend_scale: float,

            inpaint_detail_preservation: float,

    assert 0. <= denoising_strength <= 1., 'can only work with strength in [0.0, 1.0]'

    soft_inpainting = si.SoftInpaintingSettings(mask_blend_power, mask_blend_scale, inpaint_detail_preservation) \

        if mask_blend_enabled else None

    p = StableDiffusionProcessingImg2Img(

        sd_model=shared.sd_model,

        outpath_samples=opts.outdir_samples or opts.outdir_img2img_samples,

        init_images=[image],

        mask=mask,

        mask_blur=mask_blur,

        mask_blend_power=mask_blend_power,

        mask_blend_scale=mask_blend_scale,

        inpaint_detail_preservation=inpaint_detail_preservation,

        soft_inpainting=soft_inpainting,

        inpainting_fill=inpainting_fill,

        resize_mode=resize_mode,

        denoising_strength=denoising_strength,

    if mask:

        p.extra_generation_params["Mask blur"] = mask_blur

        p.extra_generation_params["Mask blending bias"] = mask_blend_power

        p.extra_generation_params["Mask blending preservation"] = mask_blend_scale

        p.extra_generation_params["Mask blending contrast boost"] = inpaint_detail_preservation

        if soft_inpainting is not None:

            soft_inpainting.add_generation_params(p.extra_generation_params)

    with closing(p):

        if is_batch:

import modules.sd_vae as sd_vae

from ldm.data.util import AddMiDaS

from ldm.models.diffusion.ddpm import LatentDepth2ImageDiffusion

import modules.soft_inpainting as si

from einops import repeat, rearrange

from blendmodes.blend import blendLayers, BlendType

    mask_blur_x: int = 4

    mask_blur_y: int = 4

    mask_blur: int = None

    mask_blend_power: float = 1

    mask_blend_scale: float = 0.5

    inpaint_detail_preservation: float = 4

    soft_inpainting: si.SoftInpaintingParameters = si.default

    inpainting_fill: int = 0

    inpaint_full_res: bool = True

    inpaint_full_res_padding: int = 0

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

import modules.soft_inpainting as si

def catenate_conds(conds):

        self.model_wrap = None

        self.mask = None

        self.nmask = None

        self.mask_blend_power = 1

        self.mask_blend_scale = 0.5

        self.inpaint_detail_preservation = 4

        self.soft_inpainting: si.SoftInpaintingParameters = None

        self.init_latent = None

        self.steps = None

        """number of steps as specified by user in UI"""

        self.sampler.sampler_extra_args['uncond'] = uc

    def forward(self, x, sigma, uncond, cond, cond_scale, s_min_uncond, image_cond):

        def latent_blend(a, b, t):

        def latent_blend(a, b, t, one_minus_t=None):

            """

            Interpolates two latent image representations according to the parameter t,

            where the interpolated vectors' magnitudes are also interpolated separately.

            """

            # NOTE: We use inplace operations wherever possible.

            if one_minus_t is None:

                one_minus_t = 1 - t

            if self.soft_inpainting is None:

                return a * one_minus_t + b * t

            # Linearly interpolate the image vectors.

            a_scaled = a * one_minus_t

            b_scaled = b * t

            current_magnitude = torch.norm(image_interp, p=2, dim=1).to(torch.float64).add_(0.00001)

            # Interpolate the powered magnitudes, then un-power them (bring them back to a power of 1).

            a_magnitude = torch.norm(a, p=2, dim=1).to(torch.float64).pow_(self.inpaint_detail_preservation) * one_minus_t

            b_magnitude = torch.norm(b, p=2, dim=1).to(torch.float64).pow_(self.inpaint_detail_preservation) * t

            a_magnitude = torch.norm(a, p=2, dim=1).to(torch.float64).pow_(self.soft_inpainting.inpaint_detail_preservation) * one_minus_t

            b_magnitude = torch.norm(b, p=2, dim=1).to(torch.float64).pow_(self.soft_inpainting.inpaint_detail_preservation) * t

            desired_magnitude = a_magnitude

            desired_magnitude.add_(b_magnitude).pow_(1 / self.inpaint_detail_preservation)

            desired_magnitude.add_(b_magnitude).pow_(1 / self.soft_inpainting.inpaint_detail_preservation)

            del a_magnitude, b_magnitude, one_minus_t

            # Change the linearly interpolated image vectors' magnitudes to the value we want.

            NOTE: "mask" is not used

            """

            return torch.pow(nmask, (_sigma ** self.mask_blend_power) * self.mask_blend_scale)

            if self.soft_inpainting is None:

                return nmask

            return torch.pow(nmask, (_sigma ** self.soft_inpainting.mask_blend_power) * self.soft_inpainting.mask_blend_scale)

        if state.interrupted or state.skipped:

            raise sd_samplers_common.InterruptedException

        # Blend in the original latents (before)

        if self.mask_before_denoising and self.mask is not None:

            if self.soft_inpainting is None:

                x = latent_blend(self.init_latent, x, self.nmask, self.mask)

            else:

                x = latent_blend(self.init_latent, x, get_modified_nmask(self.nmask, sigma))

        batch_size = len(conds_list)

        # Blend in the original latents (after)

        if not self.mask_before_denoising and self.mask is not None:

            if self.soft_inpainting is None:

                denoised = latent_blend(self.init_latent, denoised, self.nmask, self.mask)

            else:

                denoised = latent_blend(self.init_latent, denoised, get_modified_nmask(self.nmask, sigma))

        self.sampler.last_latent = self.get_pred_x0(torch.cat([x_in[i:i + 1] for i in denoised_image_indexes]), torch.cat([x_out[i:i + 1] for i in denoised_image_indexes]), sigma)

        self.model_wrap_cfg.p = p

        self.model_wrap_cfg.mask = p.mask if hasattr(p, 'mask') else None

        self.model_wrap_cfg.nmask = p.nmask if hasattr(p, 'nmask') else None

        self.model_wrap_cfg.mask_blend_power = p.mask_blend_power if hasattr(p, 'mask_blend_power') else None

        self.model_wrap_cfg.mask_blend_scale = p.mask_blend_scale if hasattr(p, 'mask_blend_scale') else None

        self.model_wrap_cfg.inpaint_detail_preservation = p.inpaint_detail_preservation if hasattr(p, 'inpaint_detail_preservation') else None

        self.model_wrap_cfg.soft_inpainting = p.soft_inpainting if hasattr(p, 'soft_inpainting') else None

        self.model_wrap_cfg.step = 0

        self.model_wrap_cfg.image_cfg_scale = getattr(p, 'image_cfg_scale', None)

        self.eta = p.eta if p.eta is not None else getattr(opts, self.eta_option_field, 0.0)

class SoftInpaintingSettings:

    def __init__(self, mask_blend_power, mask_blend_scale, inpaint_detail_preservation):

        self.mask_blend_power = mask_blend_power

        self.mask_blend_scale = mask_blend_scale

        self.inpaint_detail_preservation = inpaint_detail_preservation

    def get_paste_fields(self):

        return [

            (self.mask_blend_power, gen_param_labels.mask_blend_power),

            (self.mask_blend_scale, gen_param_labels.mask_blend_scale),

            (self.inpaint_detail_preservation, gen_param_labels.inpaint_detail_preservation),

        ]

    def add_generation_params(self, dest):

        dest[enabled_gen_param_label] = True

        dest[gen_param_labels.mask_blend_power] = self.mask_blend_power

        dest[gen_param_labels.mask_blend_scale] = self.mask_blend_scale

        dest[gen_param_labels.inpaint_detail_preservation] = self.inpaint_detail_preservation

enabled_ui_label = "Soft inpainting"

enabled_gen_param_label = "Soft inpainting enabled"

enabled_el_id = "soft_inpainting_enabled"

default = SoftInpaintingSettings(1, 0.5, 4)

ui_labels = SoftInpaintingSettings("Schedule bias", "Preservation strength", "Transition contrast boost")

ui_info = SoftInpaintingSettings(

    mask_blend_power="Shifts when preservation of original content occurs during denoising.",

                     # "Below 1: Stronger preservation near the end (with low sigma)\n"

                     # "1: Balanced (proportional to sigma)\n"

                     # "Above 1: Stronger preservation in the beginning (with high sigma)",

    mask_blend_scale="How strongly partially masked content should be preserved.",

                     # "Low values: Favors generated content.\n"

                     # "High values: Favors original content.",

    inpaint_detail_preservation="Amplifies the contrast that may be lost in partially masked regions.")

gen_param_labels = SoftInpaintingSettings("Soft inpainting schedule bias", "Soft inpainting preservation strength", "Soft inpainting transition contrast boost")

el_ids = SoftInpaintingSettings("mask_blend_power", "mask_blend_scale", "inpaint_detail_preservation")

def gradio_ui():

    import gradio as gr

    from modules.ui_components import InputAccordion

    """

            with InputAccordion(False, label="Refiner", elem_id=self.elem_id("enable")) as enable_refiner:

            with gr.Row():

                refiner_checkpoint = gr.Dropdown(label='Checkpoint', elem_id=self.elem_id("checkpoint"), choices=sd_models.checkpoint_tiles(), value='', tooltip="switch to another model in the middle of generation")

                create_refresh_button(refiner_checkpoint, sd_models.list_models, lambda: {"choices": sd_models.checkpoint_tiles()}, self.elem_id("checkpoint_refresh"))

                refiner_switch_at = gr.Slider(value=0.8, label="Switch at", minimum=0.01, maximum=1.0, step=0.01, elem_id=self.elem_id("switch_at"), tooltip="fraction of sampling steps when the switch to refiner model should happen; 1=never, 0.5=switch in the middle of generation")

    """

    with InputAccordion(False, label=enabled_ui_label, elem_id=enabled_el_id) as soft_inpainting_enabled:

        with gr.Group():

            gr.Markdown(

                """

                Soft inpainting allows you to **seamlessly blend original content with inpainted content** according to the mask opacity.

                **High _Mask blur_** values are recommended!

                """)

            result = SoftInpaintingSettings(

                gr.Slider(label=ui_labels.mask_blend_power,

                          info=ui_info.mask_blend_power,

                          minimum=0,

                          maximum=8,

                          step=0.1,

                          value=default.mask_blend_power,

                          elem_id=el_ids.mask_blend_power),

                gr.Slider(label=ui_labels.mask_blend_scale,

                          info=ui_info.mask_blend_scale,

                          minimum=0,

                          maximum=8,

                          step=0.05,

                          value=default.mask_blend_scale,

                          elem_id=el_ids.mask_blend_scale),

                gr.Slider(label=ui_labels.inpaint_detail_preservation,

                          info=ui_info.inpaint_detail_preservation,

                          minimum=1,

                          maximum=32,

                          step=0.5,

                          value=default.inpaint_detail_preservation,

                          elem_id=el_ids.inpaint_detail_preservation))

            with gr.Accordion("Help", open=False):

                gr.Markdown(

                    f"""

                    ### {ui_labels.mask_blend_power}

                    The blending strength of original content is scaled proportionally with the decreasing noise level values at each step (sigmas).

                    This ensures that the influence of the denoiser and original content preservation is roughly balanced at each step.

                    This balance can be shifted using this parameter, controlling whether earlier or later steps have stronger preservation.

                    - **Below 1**: Stronger preservation near the end (with low sigma)

                    - **1**: Balanced (proportional to sigma)

                    - **Above 1**: Stronger preservation in the beginning (with high sigma)

                    """)

                gr.Markdown(

                    f"""

                    ### {ui_labels.mask_blend_scale}

                    Skews whether partially masked image regions should be more likely to preserve the original content or favor inpainted content.

                    This may need to be adjusted depending on the {ui_labels.mask_blend_power}, CFG Scale, prompt and Denoising strength.

                    - **Low values**: Favors generated content.

                    - **High values**: Favors original content.

                    """)

                gr.Markdown(

                    f"""

                    ### {ui_labels.inpaint_detail_preservation}

                    This parameter controls how the original latent vectors and denoised latent vectors are interpolated.

                    With higher values, the magnitude of the resulting blended vector will be closer to the maximum of the two interpolated vectors.

                    This can prevent the loss of contrast that occurs with linear interpolation.

                    - **Low values**: Softer blending, details may fade.

                    - **High values**: Stronger contrast, may over-saturate colors.

                    """)

    return (

        [

            soft_inpainting_enabled,

            result.mask_blend_power,

            result.mask_blend_scale,

            result.inpaint_detail_preservation

        ],

        [

            (soft_inpainting_enabled, enabled_gen_param_label),

            (result.mask_blend_power, gen_param_labels.mask_blend_power),

            (result.mask_blend_scale, gen_param_labels.mask_blend_scale),

            (result.inpaint_detail_preservation, gen_param_labels.inpaint_detail_preservation)

        ]

    )