From c1a31ec9f75c8dfe4ddcb0061f06e2704db98359 Mon Sep 17 00:00:00 2001
From: AUTOMATIC1111 <16777216c@gmail.com>
Date: Mon, 14 Aug 2023 08:59:15 +0300
Subject: revert to applying mask before denoising for k-diffusion, like it was
 before

---
 modules/sd_samplers_cfg_denoiser.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

(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 113425b2..bc9b97e4 100644
--- a/modules/sd_samplers_cfg_denoiser.py
+++ b/modules/sd_samplers_cfg_denoiser.py
@@ -56,6 +56,7 @@ class CFGDenoiser(torch.nn.Module):
         self.sampler = sampler
         self.model_wrap = None
         self.p = None
+        self.mask_before_denoising = False
 
     @property
     def inner_model(self):
@@ -104,7 +105,7 @@ class CFGDenoiser(torch.nn.Module):
 
         assert not is_edit_model or all(len(conds) == 1 for conds in conds_list), "AND is not supported for InstructPix2Pix checkpoint (unless using Image CFG scale = 1.0)"
 
-        if self.mask is not None:
+        if self.mask_before_denoising and self.mask is not None:
             x = self.init_latent * self.mask + self.nmask * x
 
         batch_size = len(conds_list)
@@ -206,6 +207,9 @@ class CFGDenoiser(torch.nn.Module):
         else:
             denoised = self.combine_denoised(x_out, conds_list, uncond, cond_scale)
 
+        if not self.mask_before_denoising and self.mask is not None:
+            denoised = self.init_latent * self.mask + self.nmask * denoised
+
         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)
 
         if opts.live_preview_content == "Prompt":
-- 
cgit v1.2.3