Organized the settings and UI of soft inpainting to allow for toggling the feature, and centralizes default values to reduce the amount of copy-pasta.

1 files changed, 24 insertions, 11 deletions

diff --git a/modules/sd_samplers_cfg_denoiser.py b/modules/sd_samplers_cfg_denoiser.py
index efbe7a40..0ee0b7dd 100644
--- a/modules/sd_samplers_cfg_denoiser.py
+++ b/modules/sd_samplers_cfg_denoiser.py
@@ -6,6 +6,7 @@ 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

+import modules.soft_inpainting as si

 

 

 def catenate_conds(conds):

@@ -43,9 +44,7 @@ class CFGDenoiser(torch.nn.Module):
         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"""

@@ -95,7 +94,8 @@ class CFGDenoiser(torch.nn.Module):
         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.

@@ -104,7 +104,11 @@ class CFGDenoiser(torch.nn.Module):
             """

             # NOTE: We use inplace operations wherever possible.

 

-            one_minus_t = 1 - t

+            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

@@ -119,10 +123,10 @@ class CFGDenoiser(torch.nn.Module):
             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.

@@ -156,7 +160,10 @@ class CFGDenoiser(torch.nn.Module):
 

             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

@@ -176,7 +183,10 @@ class CFGDenoiser(torch.nn.Module):
 

         # Blend in the original latents (before)

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

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

+            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)

         repeats = [len(conds_list[i]) for i in range(batch_size)]

@@ -279,7 +289,10 @@ class CFGDenoiser(torch.nn.Module):
 

         # Blend in the original latents (after)

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

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

+            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)