From aaacf4823241450d88315af9d465d6815119fe0d Mon Sep 17 00:00:00 2001
From: CodeHatchling <steve@codehatch.com>
Date: Mon, 4 Dec 2023 01:27:22 -0700
Subject: 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.

---
 modules/sd_samplers_cfg_denoiser.py | 35 ++++++++++++++++++++++++-----------
 1 file changed, 24 insertions(+), 11 deletions(-)

(limited to 'modules/sd_samplers_cfg_denoiser.py')

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)
 
-- 
cgit v1.2.3