diff --git a/nodes.py b/nodes.py index 1ac817a2..300cfcf1 100644 --- a/nodes.py +++ b/nodes.py @@ -382,6 +382,7 @@ class InpaintModelConditioning: "vae": ("VAE", ), "pixels": ("IMAGE", ), "mask": ("MASK", ), + "add_noise_mask": ("BOOLEAN", {"default": True, "tooltip": "Add a noise mask to the latent so sampling will only happen within the mask. Might improve results or completely break things depending on the model."}), }} RETURN_TYPES = ("CONDITIONING","CONDITIONING","LATENT") @@ -390,7 +391,7 @@ class InpaintModelConditioning: CATEGORY = "conditioning/inpaint" - def encode(self, positive, negative, pixels, vae, mask): + def encode(self, positive, negative, pixels, vae, mask, add_noise_mask): x = (pixels.shape[1] // 8) * 8 y = (pixels.shape[2] // 8) * 8 mask = torch.nn.functional.interpolate(mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1])), size=(pixels.shape[1], pixels.shape[2]), mode="bilinear") @@ -414,7 +415,8 @@ class InpaintModelConditioning: out_latent = {} out_latent["samples"] = orig_latent - out_latent["noise_mask"] = mask + if add_noise_mask: + out_latent["noise_mask"] = mask out = [] for conditioning in [positive, negative]: