revert to applying mask before denoising for k-diffusion, like it was before
This commit is contained in:
parent
cda2f0a162
commit
c1a31ec9f7
|
@ -56,6 +56,7 @@ class CFGDenoiser(torch.nn.Module):
|
||||||
self.sampler = sampler
|
self.sampler = sampler
|
||||||
self.model_wrap = None
|
self.model_wrap = None
|
||||||
self.p = None
|
self.p = None
|
||||||
|
self.mask_before_denoising = False
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def inner_model(self):
|
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)"
|
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
|
x = self.init_latent * self.mask + self.nmask * x
|
||||||
|
|
||||||
batch_size = len(conds_list)
|
batch_size = len(conds_list)
|
||||||
|
@ -206,6 +207,9 @@ class CFGDenoiser(torch.nn.Module):
|
||||||
else:
|
else:
|
||||||
denoised = self.combine_denoised(x_out, conds_list, uncond, cond_scale)
|
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)
|
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":
|
if opts.live_preview_content == "Prompt":
|
||||||
|
|
|
@ -49,6 +49,7 @@ class CFGDenoiserTimesteps(CFGDenoiser):
|
||||||
super().__init__(sampler)
|
super().__init__(sampler)
|
||||||
|
|
||||||
self.alphas = shared.sd_model.alphas_cumprod
|
self.alphas = shared.sd_model.alphas_cumprod
|
||||||
|
self.mask_before_denoising = True
|
||||||
|
|
||||||
def get_pred_x0(self, x_in, x_out, sigma):
|
def get_pred_x0(self, x_in, x_out, sigma):
|
||||||
ts = sigma.to(dtype=int)
|
ts = sigma.to(dtype=int)
|
||||||
|
|
Loading…
Reference in New Issue