add new masked blend node

README.md

@@ -126,6 +126,7 @@ Check more for use cases like inpaint on [fine tuned models](https://github.com/
    - **[VAE Encode for Inpainting](https://comfyanonymous.github.io/ComfyUI_examples/inpaint/)** (middle),  
    - **[Set Latent Noise Mask](https://comfyui-wiki.com/en/tutorial/basic/how-to-inpaint-an-image-in-comfyui)** (second from right),  
    - **LanPaint** (far right).  
+8. You also get an output from "masked blend" node, which copy the original image and paste onto the unmasked part of output. It is useful if you want unmasked region to match original picture pixel perfectly.
 
 Compare and explore the results from each method!
 
@@ -166,11 +167,13 @@ Full parameter control:
 | `LanPaint_StepSize` | 0.1-1.0 | The StepSize of each thinking step. Recommend 0.1-0.5. |
 | `LanPaint_Beta` | 0.1-2.0 | The StepSize ratio between masked / unmasked region. Small value can compensate high lambda values. Recommend 1.0 |
 | `LanPaint_Friction` | 0.0-100.0 | The friction of Langevin dynamics. Higher means more slow but stable, lower means fast but unstable. Recommend 10.0 - 20.0|
+| `LanPaint_EarlyStop` | 0-10 | Stop LanPaint iteration before the final sampling step. Helps to remove artifacts in some cases. Recommend 1-5|
 | `LanPaint_PromptMode` | Image First / Prompt First | Image First mode focuses on the image context, maybe ignore prompt. Prompt First mode focuses more on the prompt. |
 
 For detailed descriptions of each parameter, simply hover your mouse over the corresponding input field to view tooltips with additional information.
 
-
+### LanPaint Mask Blend
+This node blends the original image with the inpainted image based on the mask. It is useful if you want the unmasked region to match the original image pixel perfectly.
 
 ## LanPaint KSampler (Advanced) Tuning Guide
 For challenging inpainting tasks:  

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "LanPaint"
-version = "1.0.2"
+version = "1.0.3"
 description = "Achieve seamless inpainting results without needing a specialized inpainting model."
 authors = [
   {name = "LanPaint", email = "czhengac@connect.ust.hk"}

src/LanPaint/nodes.py

@@ -114,7 +114,7 @@ class KSamplerX0Inpaint:
             current_step = torch.argmin( torch.abs( self.sigmas - torch.mean(sigma) ) )
             total_steps = len(self.sigmas)-1
 
-            if total_steps - current_step < self.LanPaint_early_stop:
+            if total_steps - current_step <= self.LanPaint_early_stop:
                 out = self.PaintMethod(x, self.latent_image, self.noise, sigma, latent_mask, current_times, model_options, seed, n_steps=0)
             else:
                 out = self.PaintMethod(x, self.latent_image, self.noise, sigma, latent_mask, current_times, model_options, seed)
@@ -326,11 +326,64 @@ class LanPaint_KSamplerAdvanced:
             return nodes.common_ksampler(model, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise, disable_noise=disable_noise, start_step=start_at_step, last_step=end_at_step, force_full_denoise=force_full_denoise)
 
 
+class MaskBlend:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image1": ("IMAGE",),
+                "image2": ("IMAGE",),
+                "mask": ("MASK",),
+                "blend_overlap": ("INT", {"default": 1, "min": 1, "max": 51, "step": 2, "tooltip": "The number of pixels to blend between the two images."})
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "blend_images"
+
+    CATEGORY = "image/postprocessing"
+
+    def blend_images(self, image1: torch.Tensor, image2: torch.Tensor, mask: torch.Tensor, blend_overlap: int):
+        # smooth the binary 01 mask, keep 1 still 1, but smooth the transition from 1 to 0
+        # for each mask pixel, find out the nearest 1 pixel, and set the mask value to the distance between the two pixels
+        mask = mask.float()
+        mask = torch.nn.functional.max_pool2d(mask, kernel_size=blend_overlap, stride=1, padding=blend_overlap//2)
+        # apply Gaussian blur with kernel size blend_overlap
+        kernel = self.gaussian_kernel(blend_overlap)
+        kernel = kernel.to(image1.device)
+        kernel = kernel[None, None, ...]
+        
+        mask = torch.nn.functional.conv2d(mask[:,None,:,:], kernel, padding=blend_overlap//2)[:,0,:,:]
+
+
+        blended_image = image1 * (1 - mask[...,None]) + image2 * mask[...,None]
+        return (blended_image,)
+    def gaussian_kernel(self,kernel_size):
+        """
+        Creates a 2D Gaussian kernel with the given size and standard deviation (sigma).
+        """
+        sigma = (kernel_size - 1)/4
+        # Create a grid of (x, y) coordinates
+        x = torch.arange(kernel_size).float() - kernel_size // 2
+        y = torch.arange(kernel_size).float() - kernel_size // 2
+        x_grid, y_grid = torch.meshgrid(x, y, indexing='ij')
+
+        # Compute the Gaussian function
+        kernel = torch.exp(-(x_grid ** 2 + y_grid ** 2) / (2 * sigma ** 2))
+        kernel = kernel / kernel.sum()  # Normalize the kernel
+
+        return kernel
+
+
 # A dictionary that contains all nodes you want to export with their names
 # NOTE: names should be globally unique
 NODE_CLASS_MAPPINGS = {
     "LanPaint_KSampler": LanPaint_KSampler,
     "LanPaint_KSamplerAdvanced": LanPaint_KSamplerAdvanced,
+    "LanPaint_MaskBlend": MaskBlend,
 #    "LanPaint_UpSale_LatentNoiseMask": LanPaint_UpSale_LatentNoiseMask,
 }
 
@@ -338,5 +391,6 @@ NODE_CLASS_MAPPINGS = {
 NODE_DISPLAY_NAME_MAPPINGS = {
     "LanPaint_KSampler": "LanPaint KSampler",
     "LanPaint_KSamplerAdvanced": "LanPaint KSampler (Advanced)",
+    "LanPaint_MaskBlend": "LanPaint Mask Blend",
 #    "LanPaint_UpSale_LatentNoiseMask": "LanPaint UpSale Latent Noise Mask"
 }