From d32b5e54581a3dbca480e47cad83a63995d0a677 Mon Sep 17 00:00:00 2001
From: lllyasviel <lyuminzhang@outlook.com>
Date: Sun, 4 Feb 2024 03:34:20 -0800
Subject: [PATCH] fix ddim inpaint

---
 modules/sd_samplers_cfg_denoiser.py | 23 ++++++++---------------
 1 file changed, 8 insertions(+), 15 deletions(-)

diff --git a/modules/sd_samplers_cfg_denoiser.py b/modules/sd_samplers_cfg_denoiser.py
index b0103492..7475209b 100644
--- a/modules/sd_samplers_cfg_denoiser.py
+++ b/modules/sd_samplers_cfg_denoiser.py
@@ -158,9 +158,7 @@ class CFGDenoiser(torch.nn.Module):
         if classic_ddim_eps_estimation:
             acd = self.inner_model.inner_model.alphas_cumprod
             fake_sigmas = ((1 - acd) / acd) ** 0.5
-            tid = sigma.round().long().clip(0, int(fake_sigmas.shape[0]))
-            real_sigma = fake_sigmas[tid]
-            sqrt_alpha_cumprod = torch.sqrt(acd[tid])
+            real_sigma = fake_sigmas[sigma.round().long().clip(0, int(fake_sigmas.shape[0]))]
             real_sigma_data = 1.0
             x = x * (real_sigma ** 2.0 + real_sigma_data ** 2.0) ** 0.5
             sigma = real_sigma
@@ -174,24 +172,19 @@ class CFGDenoiser(torch.nn.Module):
 
         # If we use masks, blending between the denoised and original latent images occurs here.
         def apply_blend(current_latent):
-            init_latent = self.init_latent
-
-            if classic_ddim_eps_estimation:
-                init_latent = init_latent / sqrt_alpha_cumprod
-
-            blended_latent = current_latent * self.nmask + init_latent * self.mask
+            blended_latent = current_latent * self.nmask + self.init_latent * self.mask
 
             if self.p.scripts is not None:
                 from modules import scripts
-                mba = scripts.MaskBlendArgs(current_latent, self.nmask, init_latent, self.mask, blended_latent, denoiser=self, sigma=sigma)
+                mba = scripts.MaskBlendArgs(current_latent, self.nmask, self.init_latent, self.mask, blended_latent, denoiser=self, sigma=sigma)
                 self.p.scripts.on_mask_blend(self.p, mba)
                 blended_latent = mba.blended_latent
 
             return blended_latent
 
-        # Blend in the original latents (before)
-        if self.mask_before_denoising and self.mask is not None:
-            x = apply_blend(x)
+        # # Blend in the original latents (before, wrong method)
+        # if self.mask is not None:
+        #     x = apply_blend(x)
 
         denoiser_params = CFGDenoiserParams(x, image_cond, sigma, state.sampling_step, state.sampling_steps, cond, uncond, self)
         cfg_denoiser_callback(denoiser_params)
@@ -199,8 +192,8 @@ class CFGDenoiser(torch.nn.Module):
         denoised = forge_sampler.forge_sample(self, denoiser_params=denoiser_params,
                                               cond_scale=cond_scale, cond_composition=cond_composition)
 
-        # Blend in the original latents (after)
-        if not self.mask_before_denoising and self.mask is not None:
+        # Blend in the original latents (after, correct method)
+        if self.mask is not None:
             denoised = apply_blend(denoised)
 
         preview = self.sampler.last_latent = denoised