diff --git a/nodes.py b/nodes.py index 1a567e01..1b22e18d 100644 --- a/nodes.py +++ b/nodes.py @@ -157,20 +157,23 @@ class VAEEncodeForInpaint: def encode(self, vae, pixels, mask): x = (pixels.shape[1] // 64) * 64 y = (pixels.shape[2] // 64) * 64 + mask = torch.nn.functional.interpolate(mask[None,None,], size=(pixels.shape[1], pixels.shape[2]), mode="bilinear")[0][0] + if pixels.shape[1] != x or pixels.shape[2] != y: pixels = pixels[:,:x,:y,:] mask = mask[:x,:y] - #shave off a few pixels to keep things seamless + #grow mask by a few pixels to keep things seamless in latent space kernel_tensor = torch.ones((1, 1, 6, 6)) - mask_erosion = torch.clamp(torch.nn.functional.conv2d((1.0 - mask.round())[None], kernel_tensor, padding=3), 0, 1) + mask_erosion = torch.clamp(torch.nn.functional.conv2d((mask.round())[None], kernel_tensor, padding=3), 0, 1) + m = (1.0 - mask.round()) for i in range(3): pixels[:,:,:,i] -= 0.5 - pixels[:,:,:,i] *= mask_erosion[0][:x,:y].round() + pixels[:,:,:,i] *= m pixels[:,:,:,i] += 0.5 t = vae.encode(pixels) - return ({"samples":t, "noise_mask": mask}, ) + return ({"samples":t, "noise_mask": (mask_erosion[0][:x,:y].round())}, ) class CheckpointLoader: models_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "models")