aboutsummaryrefslogtreecommitdiff
path: root/modules/sd_samplers_cfg_denoiser.py
diff options
context:
space:
mode:
Diffstat (limited to 'modules/sd_samplers_cfg_denoiser.py')
-rw-r--r--modules/sd_samplers_cfg_denoiser.py35
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)