Merge remote-tracking branch 'auto1111/dev' into shared-hires-prompt-test

1 files changed, 10 insertions, 2 deletions

diff --git a/modules/sd_samplers_cfg_denoiser.py b/modules/sd_samplers_cfg_denoiser.py
index a532e013..bc9b97e4 100644
--- a/modules/sd_samplers_cfg_denoiser.py
+++ b/modules/sd_samplers_cfg_denoiser.py
@@ -45,18 +45,23 @@ class CFGDenoiser(torch.nn.Module):
         self.nmask = None

         self.init_latent = None

         self.steps = None

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

+

+        self.total_steps = None

+        """expected number of calls to denoiser calculated from self.steps and specifics of the selected sampler"""

+

         self.step = 0

         self.image_cfg_scale = None

         self.padded_cond_uncond = False

         self.sampler = sampler

         self.model_wrap = None

         self.p = None

+        self.mask_before_denoising = False

 

     @property

     def inner_model(self):

         raise NotImplementedError()

 

-

     def combine_denoised(self, x_out, conds_list, uncond, cond_scale):

         denoised_uncond = x_out[-uncond.shape[0]:]

         denoised = torch.clone(denoised_uncond)

@@ -100,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)

@@ -202,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":