Mirrored_Repos/InvokeAI
1
0
Fork 0
mirror of https://github.com/invoke-ai/InvokeAI synced 2024-08-30 20:32:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix masked generation with inpaint models

Browse Source
This commit is contained in:
Sergey Borisov 2023-08-16 20:28:33 +03:00
parent a7ba142ad9
commit 5aa7bfebd4
1 changed files with 18 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
invokeai/backend/stable_diffusion/diffusers_pipeline.py
View File

@ -360,14 +360,6 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline):
latents = self.scheduler.add_noise(latents, noise, batched_t) latents = self.scheduler.add_noise(latents, noise, batched_t)
if mask is not None: if mask is not None:
if is_inpainting_model(self.unet):
# You'd think the inpainting model wouldn't be paying attention to the area it is going to repaint
# (that's why there's a mask!) but it seems to really want that blanked out.
# masked_latents = latents * torch.where(mask < 0.5, 1, 0) TODO: inpaint/outpaint/infill
# TODO: we should probably pass this in so we don't have to try/finally around setting it.
self.invokeai_diffuser.model_forward_callback = AddsMaskLatents(self._unet_forward, mask, orig_latents)
else:
# if no noise provided, noisify unmasked area based on seed(or 0 as fallback) # if no noise provided, noisify unmasked area based on seed(or 0 as fallback)
if noise is None: if noise is None:
noise = torch.randn( noise = torch.randn(
@ -382,6 +374,16 @@ class StableDiffusionGeneratorPipeline(StableDiffusionPipeline):
orig_latents, latents.to(dtype=orig_latents.dtype), mask.to(dtype=orig_latents.dtype) orig_latents, latents.to(dtype=orig_latents.dtype), mask.to(dtype=orig_latents.dtype)
) )
if is_inpainting_model(self.unet):
# You'd think the inpainting model wouldn't be paying attention to the area it is going to repaint
# (that's why there's a mask!) but it seems to really want that blanked out.
masked_latents = orig_latents * torch.where(mask < 0.5, 1, 0)
# TODO: we should probably pass this in so we don't have to try/finally around setting it.
self.invokeai_diffuser.model_forward_callback = AddsMaskLatents(
self._unet_forward, mask, masked_latents
)
else:
additional_guidance.append(AddsMaskGuidance(mask, orig_latents, self.scheduler, noise)) additional_guidance.append(AddsMaskGuidance(mask, orig_latents, self.scheduler, noise))
try: try: