Add Qwen image support and new inpainting examples

Updated the README to announce Qwen image support and provide usage instructions. Added new example images and workflow files for Qwen inpainting in the examples directory.

README.md

@@ -1,8 +1,10 @@
 # LanPaint (Thinking mode Inpaint)
 
-Unlock precise inpainting without additional training. LanPaint lets the model "think" through multiple iterations before denoising, aiming for seamless and accurate results. 
+Unlock precise inpainting without additional training. LanPaint lets the model "think" through multiple iterations before denoising, enabling you to invest more computation time for superior quality.  
 ![Inpainting Result 13](https://github.com/scraed/LanPaint/blob/master/examples/InpaintChara_13.jpg) 
 
+This is the official implementation of ["Lanpaint: Training-Free Diffusion Inpainting with Exact and Fast Conditional Inference"](https://arxiv.org/abs/2502.03491). 
+
 ## Features
 
 - **Universal Compatibility** – Works instantly with almost any model (SD 1.5, XL, 3.5, Flux, HiDream, or custom LoRAs) and ControlNet.  
@@ -12,8 +14,10 @@ Unlock precise inpainting without additional training. LanPaint lets the model "
 - **No Workarounds** – Generates 100% new content (no blending or smoothing) without relying on partial denoising.  
 - **Beyond Inpainting** – You can even use it as a simple way to generate consistent characters. 
 
+**Warning**: LanPaint has degraded performance on distillation models, such as Flux.dev, due to a similar [issue with LORA training](https://medium.com/@zhiwangshi28/why-flux-lora-so-hard-to-train-and-how-to-overcome-it-a0c70bc59eaf). Please use low flux guidance (1.0-2.0) to mitigate this [issue](https://github.com/scraed/LanPaint/issues/30).
+
 ## How It Works  
-This is the official implementation of "Lanpaint: Training-Free Diffusion Inpainting with Exact and Fast Conditional Inference". LanPaint uses Langevin Dynamics as "thinking" steps, which digs deeper into the diffusion process and allows the model to generate more consistent results.
+LanPaint uses Langevin Dynamics as "thinking" steps, which digs deeper into the diffusion process and allows the model to generate more consistent results.
 
 LanPaint introduces "BIG score" that creates a **two-way alignment** between masked and unmasked areas. It continuously evaluates:  
 - *"Does the new content make sense with the existing elements?"*  
@@ -22,7 +26,34 @@ Based on this evaluation, LanPaint iteratively updates the noise in both the mas
   
 LanPaint also implements an accurate, robust, and fast Langevin dynamics solver.
 
+
+## Quickstart
+
+1. **Install ComfyUI**: Follow the official [ComfyUI installation guide](https://docs.comfy.org/get_started) to set up ComfyUI on your system. Or ensure your ComfyUI version > 0.3.11.
+2. **Install ComfyUI-Manager**: Add the [ComfyUI-Manager](https://github.com/ltdrdata/ComfyUI-Manager) for easy extension management.  
+3. **Install LanPaint Nodes**:  
+   - **Via ComfyUI-Manager**: Search for "[LanPaint](https://registry.comfy.org/publishers/scraed/nodes/LanPaint)" in the manager and install it directly.  
+   - **Manually**: Click "Install via Git URL" in ComfyUI-Manager and input the GitHub repository link:  
+     ```
+     https://github.com/scraed/LanPaint.git
+     ```  
+     Alternatively, clone this repository into the `ComfyUI/custom_nodes` folder.  
+4. **Restart ComfyUI**: Restart ComfyUI to load the LanPaint nodes.  
+
+Once installed, you'll find the LanPaint nodes under the "sampling" category in ComfyUI. Use them just like the default KSampler for high-quality inpainting!
+
+
+
 ## Updates
+- 2025/08/08
+    - Add Qwen image support
+- 2025/06/21
+    - Update the algorithm with enhanced stability and outpaint performance.
+    - Add outpaint example
+    - Supports Sampler Custom (Thanks to [MINENEMA](https://github.com/MINENEMA))
+- 2025/06/04
+    - Add more sampler support.
+    - Add early stopping to advanced sampler.
 - 2025/05/28
     - Major update on the Langevin solver. It is now much faster and more stable.
     - Greatly simplified the parameters for advanced sampler.
@@ -36,14 +67,29 @@ LanPaint also implements an accurate, robust, and fast Langevin dynamics solver.
     - LanPaint has received a major update! All examples now use the LanPaint K Sampler, offering a simplified interface with enhanced performance and stability.
 
 ## Examples
-All examples use a random seed 0 to generate batch of 4 images for fair comparison. (Warning: Generating 4 images may exceed your GPU memory; adjust batch size as needed.)
 
-### Example HiDream: InPaint(LanPaint K Sampler, 5 steps of thinking)
+### Example Qwen Image: InPaint(LanPaint K Sampler, 5 steps of thinking)
+We are excited to announce that LanPaint now supports Qwen Image, providing powerful inpainting capabilities for image editing.
+
+![Inpainting Result 14](https://github.com/scraed/LanPaint/blob/master/examples/InpaintChara_14.jpg)  
+[View Workflow & Masks](https://github.com/scraed/LanPaint/tree/master/examples/Example_11)
+
+You need to follow the ComfyUI version of [Qwen Image workflow](https://docs.comfy.org/tutorials/image/qwen/qwen-image) to download and install the model.
+
+The following examples utilize a random seed of 0 to generate a batch of 4 images for variance demonstration and fair comparison. (Note: Generating 4 images may exceed your GPU memory; please adjust the batch size as necessary.)
+
+### Example HiDream: InPaint (LanPaint K Sampler, 5 steps of thinking)
 ![Inpainting Result 8](https://github.com/scraed/LanPaint/blob/master/examples/InpaintChara_11.jpg)  
 [View Workflow & Masks](https://github.com/scraed/LanPaint/tree/master/examples/Example_8)
 
 You need to follow the ComfyUI version of [HiDream workflow](https://docs.comfy.org/tutorials/image/hidream/hidream-i1) to download and install the model.
 
+### Example HiDream: OutPaint(LanPaint K Sampler, 5 steps of thinking)
+![Inpainting Result 8](https://github.com/scraed/LanPaint/blob/master/examples/InpaintChara_13(1).jpg)  
+[View Workflow & Masks](https://github.com/scraed/LanPaint/tree/master/examples/Example_10)
+
+You need to follow the ComfyUI version of [HiDream workflow](https://docs.comfy.org/tutorials/image/hidream/hidream-i1) to download and install the model. Thanks [Amazon90](https://github.com/Amazon90) for providing this example.
+
 ### Example SD 3.5: InPaint(LanPaint K Sampler, 5 steps of thinking)
 ![Inpainting Result 8](https://github.com/scraed/LanPaint/blob/master/examples/InpaintChara_12.jpg)  
 [View Workflow & Masks](https://github.com/scraed/LanPaint/tree/master/examples/Example_9)
@@ -61,6 +107,13 @@ You need to follow the ComfyUI version of [SD 3.5 workflow](https://comfyui-wiki
 [View Workflow & Masks](https://github.com/scraed/LanPaint/tree/master/examples/Example_6)
 [Model Used in This Example](https://civitai.com/models/1188071?modelVersionId=1408658) 
 
+(Tricks 1: You can emphasize the character by copy it's image multiple times with Photoshop. Here I have made one extra copy.)
+
+(Tricks 2: Use prompts like multiple views, multiple angles, clone, turnaround. Use LanPaint's Prompt first mode (does not support Flux))
+
+(Tricks 3: Remeber LanPaint can in-paint: Mask non-consistent regions and try again!)
+
+
 ### Example SDXL 1: Basket to Basket Ball (LanPaint K Sampler, 2 steps of thinking).
 ![Inpainting Result 1](https://github.com/scraed/LanPaint/blob/master/examples/InpaintChara_04.jpg)  
 [View Workflow & Masks](https://github.com/scraed/LanPaint/tree/master/examples/Example_1) 
@@ -82,13 +135,7 @@ You need to follow the ComfyUI version of [SD 3.5 workflow](https://comfyui-wiki
 [View Workflow & Masks](https://github.com/scraed/LanPaint/tree/master/examples/Example_5)
 [Model Used in This Example](https://civitai.com/models/133005/juggernaut-xl)
 
-
-(Tricks 1: You can emphasize the character by copy it's image multiple times with Photoshop. Here I have made one extra copy.)
-
-(Tricks 2: Use prompts like multiple views, multiple angles, clone, turnaround.)
-
-(Tricks 3: Remeber LanPaint can in-paint: Mask non-consistent regions and try again!)
-
+Check more for use cases like inpaint on [fine tuned models](https://github.com/scraed/LanPaint/issues/12#issuecomment-2938662021) and [face swapping](https://github.com/scraed/LanPaint/issues/12#issuecomment-2938723501), thanks to [Amazon90](https://github.com/Amazon90).
 
 
 ## **How to Use These Examples:**  
@@ -102,25 +149,12 @@ You need to follow the ComfyUI version of [SD 3.5 workflow](https://comfyui-wiki
    - **[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!
 
 ![WorkFlow](https://github.com/scraed/LanPaint/blob/master/Example.JPG)  
 
-## Quickstart
-
-1. **Install ComfyUI**: Follow the official [ComfyUI installation guide](https://docs.comfy.org/get_started) to set up ComfyUI on your system. Or ensure your ComfyUI version > 0.3.11.
-2. **Install ComfyUI-Manager**: Add the [ComfyUI-Manager](https://github.com/ltdrdata/ComfyUI-Manager) for easy extension management.  
-3. **Install LanPaint Nodes**:  
-   - **Via ComfyUI-Manager**: Search for "[LanPaint](https://registry.comfy.org/publishers/scraed/nodes/LanPaint)" in the manager and install it directly.  
-   - **Manually**: Click "Install via Git URL" in ComfyUI-Manager and input the GitHub repository link:  
-     ```
-     https://github.com/scraed/LanPaint.git
-     ```  
-     Alternatively, clone this repository into the `ComfyUI/custom_nodes` folder.  
-4. **Restart ComfyUI**: Restart ComfyUI to load the LanPaint nodes.  
-
-Once installed, you'll find the LanPaint nodes under the "sampling" category in ComfyUI. Use them just like the default KSampler for high-quality inpainting!
 
 ## Usage
 
@@ -142,7 +176,7 @@ Simplified interface with recommended defaults:
 
 - Steps: 20 - 50. More steps will give more "thinking" and better results.
 - LanPaint NumSteps: The turns of thinking before denoising. Recommend 5 for most of tasks ( which means 5 times slower than sampling without thinking). Use 10 for more challenging tasks. 
-- LanPaint Prompt mode: Image First mode and Prompt First mode. Image First mode focuses on the image, while Prompt First mode focuses more on the prompt. Use Prompt First mode for tasks like character consistency. (Technically, it Prompt First mode change CFG scale to negative value in the BIG score to emphasis prompt, which will costs image quality.)
+- LanPaint Prompt mode: Image First mode and Prompt First mode. Image First mode focuses on the image, inpaint based on image context (maybe ignore prompt), while Prompt First mode focuses more on the prompt. Use Prompt First mode for tasks like character consistency. (Technically, it Prompt First mode change CFG scale to negative value in the BIG score to emphasis prompt, which will costs image quality.)
 
 ### LanPaint KSampler (Advanced)
 Full parameter control:
@@ -156,17 +190,19 @@ 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_PromptMode` | Image First / Prompt First | Image First mode focuses on the image, while Prompt First mode focuses more on the prompt. |
+| `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:  
 
 1️⃣ **Boost Quality**
-Increase **LanPaint_NumSteps** (thinking iterations) or **LanPaint_Lambda** if the inpainted result does not meet your expectations.
+Increase **total number of sampling steps** (very important!), **LanPaint_NumSteps** (thinking iterations) or **LanPaint_Lambda** if the inpainted result does not meet your expectations.
   
 2️⃣ **Boost Speed**
 If you want better results but still need fewer steps, consider:

pyproject.toml

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

src/LanPaint/lanpaint.py

@@ -1,6 +1,7 @@
 import torch
 from .utils import *
 from functools import partial
+
 class LanPaint():
     def __init__(self, Model, NSteps, Friction, Lambda, Beta, StepSize, IS_FLUX = False, IS_FLOW = False):
         self.n_steps = NSteps
@@ -9,44 +10,49 @@ class LanPaint():
         self.IS_FLOW = IS_FLOW
         self.step_size = StepSize
         self.inner_model = Model
-        #self.alpha = Alpha
-        #self.beta_scale = BetaSchedule
-        #self.step_size_schedule = IteStepSchedule
-        #self.step_time_schedule = StepSizeSchedule
-        #self.end_sigma = EndSigma
         self.friction = Friction
-        #self.LanPaint_Cap_Sigma = CapSigma
         self.chara_beta = Beta
-        
-    def __call__(self, x, latent_image, noise, sigma, latent_mask, current_times, model_options, seed):
+        self.img_dim_size = None
+    def add_none_dims(self, array):
+        # Create a tuple with ':' for the first dimension and 'None' repeated num_nones times
+        index = (slice(None),) + (None,) * (self.img_dim_size-1)
+        return array[index]
+    def remove_none_dims(self, array):
+        # Create a tuple with ':' for the first dimension and 'None' repeated num_nones times
+        index = (slice(None),) + (0,) * (self.img_dim_size-1)
+        return array[index]
+    def __call__(self, x, latent_image, noise, sigma, latent_mask, current_times, model_options, seed, n_steps=None):
+        self.img_dim_size = len(x.shape)
         self.latent_image = latent_image
         self.noise = noise
-        return self.LanPaint(x, sigma, latent_mask, current_times, model_options, seed, self.IS_FLUX, self.IS_FLOW)
-    def LanPaint(self, x, sigma, latent_mask, current_times, model_options, seed, IS_FLUX, IS_FLOW):
+        if n_steps is None:
+            n_steps = self.n_steps
+        return self.LanPaint(x, sigma, latent_mask, current_times, n_steps, model_options, seed, self.IS_FLUX, self.IS_FLOW)
+    def LanPaint(self, x, sigma, latent_mask, current_times, n_steps, model_options, seed, IS_FLUX, IS_FLOW):
         VE_Sigma, abt, Flow_t = current_times
 
         
         step_size = self.step_size * (1 - abt)
-        step_size = step_size[:, None, None, None]
+        step_size = self.add_none_dims(step_size)
         # self.inner_model.inner_model.scale_latent_inpaint returns variance exploding x_t values
         # This is the replace step
         x = x * (1 - latent_mask) +  self.inner_model.inner_model.scale_latent_inpaint(x=x, sigma=sigma, noise=self.noise, latent_image=self.latent_image)* latent_mask
         
         if IS_FLUX or IS_FLOW:
-            x_t = x * ( abt[:, None,None,None]**0.5 + (1-abt[:, None,None,None])**0.5 )
+            x_t = x * ( self.add_none_dims(abt)**0.5 + (1-self.add_none_dims(abt))**0.5 )
         else:
-            x_t = x / ( 1+VE_Sigma[:, None,None,None]**2 )**0.5 # switch to variance perserving x_t values
+            x_t = x / ( 1+self.add_none_dims(VE_Sigma)**2 )**0.5 # switch to variance perserving x_t values
 
         ############ LanPaint Iterations Start ###############
         # after noise_scaling, noise = latent_image + noise * sigma, which is x_t in the variance exploding diffusion model notation for the known region.
         args = None
-        for i in range(self.n_steps):
-            score_func = partial( self.score_model, y = self.latent_image, mask = latent_mask, abt = abt[:, None,None,None], sigma = VE_Sigma[:, None,None,None], tflow = Flow_t[:, None,None,None], model_options = model_options, seed = seed )
-            x_t, args = self.langevin_dynamics(x_t, score_func , latent_mask, step_size , current_times, sigma_x = self.sigma_x(abt)[:, None,None,None], sigma_y = self.sigma_y(abt)[:, None,None,None], args = args)  
+        for i in range(n_steps):
+            score_func = partial( self.score_model, y = self.latent_image, mask = latent_mask, abt = self.add_none_dims(abt), sigma = self.add_none_dims(VE_Sigma), tflow = self.add_none_dims(Flow_t), model_options = model_options, seed = seed )
+            x_t, args = self.langevin_dynamics(x_t, score_func , latent_mask, step_size , current_times, sigma_x = self.add_none_dims(self.sigma_x(abt)), sigma_y = self.add_none_dims(self.sigma_y(abt)), args = args)  
         if IS_FLUX or IS_FLOW:
-            x = x_t / ( abt[:, None,None,None]**0.5 + (1-abt[:, None,None,None])**0.5 )
+            x = x_t / ( self.add_none_dims(abt)**0.5 + (1-self.add_none_dims(abt))**0.5 )
         else:
-            x = x_t * ( 1+VE_Sigma[:, None,None,None]**2 )**0.5 # switch to variance perserving x_t values
+            x = x_t * ( 1+self.add_none_dims(VE_Sigma)**2 )**0.5 # switch to variance perserving x_t values
         ############ LanPaint Iterations End ###############
         # out is x_0
         out, _ = self.inner_model(x, sigma, model_options=model_options, seed=seed)
@@ -60,10 +66,10 @@ class LanPaint():
         if self.IS_FLUX or self.IS_FLOW:
             # compute t for flow model, with a small epsilon compensating for numerical error.
             x = x_t / ( abt**0.5 + (1-abt)**0.5 ) # switch to Gaussian flow matching
-            x_0, x_0_BIG = self.inner_model(x, tflow[:, 0,0,0], model_options=model_options, seed=seed)
+            x_0, x_0_BIG = self.inner_model(x, self.remove_none_dims(tflow), model_options=model_options, seed=seed)
         else:
             x = x_t * ( 1+sigma**2 )**0.5 # switch to variance exploding
-            x_0, x_0_BIG = self.inner_model(x, sigma[:, 0,0,0], model_options=model_options, seed=seed)
+            x_0, x_0_BIG = self.inner_model(x, self.remove_none_dims(sigma), model_options=model_options, seed=seed)
 
         score_x = -(x_t - x_0)
         score_y =  - (1 + lamb) * ( x_t - y )  + lamb * (x_t - x_0_BIG)  
@@ -94,25 +100,44 @@ class LanPaint():
         Gamma = Gamma_x * (1-mask) + Gamma_y * mask
 
 
-
+        def Coef_C(x_t):
+            x0 = self.x0_evalutation(x_t, score, sigma, args)
+            C = (abt**0.5 * x0  - x_t )/ (1-abt) + A * x_t
+            return C
+        def advance_time(x_t, v, dt, Gamma, A, C, D):
+            dtype = x_t.dtype
+            with torch.autocast(device_type=x_t.device.type, dtype=torch.float32):
+                osc = StochasticHarmonicOscillator(Gamma, A, C, D )
+                x_t, v = osc.dynamics(x_t, v, dt )
+            x_t = x_t.to(dtype)
+            v = v.to(dtype)
+            return x_t, v
         if args is None:
             #v = torch.zeros_like(x_t)
             v = None
+            C = Coef_C(x_t)
+            #print(torch.squeeze(dtx), torch.squeeze(dty))
+            x_t, v = advance_time(x_t, v, dt, Gamma, A, C, D)
         else:
-            v, = args
+            v, C = args
 
-        with torch.autocast(device_type=x_t.device.type, dtype=torch.float32):
-            osc = StochasticHarmonicOscillator(Gamma, A, C, D )
-            x_t, v = osc.dynamics(x_t, v, dt )
+            x_t, v = advance_time(x_t, v, dt/2, Gamma, A, C, D)
+
+            C_new = Coef_C(x_t)
+            v = v + Gamma**0.5 * ( C_new - C) *dt
+
+            x_t, v = advance_time(x_t, v, dt/2, Gamma, A, C, D)
+
+            C = C_new
   
-        return x_t, (v,)
+        return x_t, (v, C)
 
     def prepare_step_size(self, current_times, step_size, sigma_x, sigma_y):
         # -------------------------------------------------------------------------
         # Unpack current times parameters (sigma and abt)
         sigma, abt, flow_t = current_times
-        sigma = sigma[:, None,None,None]
-        abt = abt[:, None,None,None]
+        sigma = self.add_none_dims(sigma)
+        abt = self.add_none_dims(abt)
         # Compute time step (dtx, dty) for x and y branches.
         dtx = 2 * step_size * sigma_x
         dty = 2 * step_size * sigma_y
@@ -127,9 +152,9 @@ class LanPaint():
         # adjust dt to match denoise-addnoise steps sizes
         Gamma_hat_x /= 2.
         Gamma_hat_y /= 2.
-
         A_t_x = (1) / ( 1 - abt ) * dtx / 2
-        A_t_y =  (1) / ( 1 - abt ) * dty / 2
+        A_t_y =  (1+self.chara_lamb) / ( 1 - abt ) * dty / 2
+
 
         A_x = A_t_x / (dtx/2)
         A_y = A_t_y / (dty/2)

src/LanPaint/nodes.py

@@ -11,15 +11,13 @@ from comfy.samplers import *
 from comfy.model_base import ModelType
 from .utils import *
 from .lanpaint import LanPaint
-# Monkey patch comfy.samplers module by importing with absolute package path
-#exec(inspect.getsource(comfy.samplers).replace("from .", "from comfy."))
+
 
 def reshape_mask(input_mask, output_shape):
     dims = len(output_shape) - 2
 
-
     scale_mode = "nearest-exact"
-    mask = torch.nn.functional.interpolate(input_mask, size=output_shape[2:], mode=scale_mode)
+    mask = torch.nn.functional.interpolate(input_mask, size=output_shape[-2:], mode=scale_mode)
     if mask.shape[1] < output_shape[1]:
         mask = mask.repeat((1, output_shape[1]) + (1,) * dims)[:,:output_shape[1]]
     mask = repeat_to_batch_size(mask, output_shape[0])
@@ -109,12 +107,15 @@ class KSamplerX0Inpaint:
             denoise_mask = (denoise_mask > 0.5).float()
 
             latent_mask = 1 - denoise_mask
-
-            
-
             current_times = (VE_Sigma, abt, Flow_t)
 
-            out = self.PaintMethod(x, self.latent_image, self.noise, sigma, latent_mask, current_times, model_options, seed)
+            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:
+                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)
         else:
             out, _ = self.inner_model(x, sigma, model_options=model_options, seed=seed)
         
@@ -162,6 +163,7 @@ class KSAMPLER(comfy.samplers.KSAMPLER):
                                        model_wrap.model_patcher.LanPaint_StepSize, 
                                        IS_FLUX = IS_FLUX, 
                                        IS_FLOW = IS_FLOW)
+        model_k.LanPaint_early_stop = model_wrap.model_patcher.LanPaint_EarlyStop
         #if not inpainting, after noise_scaling, noise = noise * sigma, which is the noise added to the clean latent image in the variance exploding diffusion model notation.
         #if inpainting, after noise_scaling, noise = latent_image + noise * sigma, which is x_t in the variance exploding diffusion model notation for the known region.
         k_callback = None
@@ -218,7 +220,12 @@ class LanPaint_UpSale_LatentNoiseMask:
         s["noise_mask"] = mask
         return (s,)
 
-KSAMPLER_NAMES = ["euler", "dpmpp_2m", "uni_pc"]
+#KSAMPLER_NAMES = ["euler", "dpmpp_2m", "uni_pc"]
+KSAMPLER_NAMES = ["euler","euler_ancestral", "heun", "heunpp2","dpm_2", "dpm_2_ancestral",
+                "dpm_fast",  "dpmpp_sde", "dpmpp_sde_gpu",
+                  "dpmpp_2m", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "dpmpp_3m_sde", "dpmpp_3m_sde_gpu", "ddpm", 
+                   "deis", "res_multistep", "res_multistep_ancestral", 
+                  "gradient_estimation",  "er_sde", "seeds_2", "seeds_3"]
 
 class LanPaint_KSampler():
     @classmethod
@@ -227,7 +234,7 @@ class LanPaint_KSampler():
             "required": {
                 "model": ("MODEL", {"tooltip": "The model used for denoising the input latent."}),
                 "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "tooltip": "The random seed used for creating the noise."}),
-                "steps": ("INT", {"default": 50, "min": 1, "max": 10000, "tooltip": "The number of steps used in the denoising process."}),
+                "steps": ("INT", {"default": 30, "min": 1, "max": 10000, "tooltip": "The number of steps used in the denoising process."}),
                 "cfg": ("FLOAT", {"default": 5.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01, "tooltip": "The Classifier-Free Guidance scale balances creativity and adherence to the prompt. Higher values result in images more closely matching the prompt however too high values will negatively impact quality."}),
                 "sampler_name": (KSAMPLER_NAMES, {"tooltip": "Recommended: euler."}),
                 "scheduler": (comfy.samplers.KSampler.SCHEDULERS, {"default": "karras", "tooltip": "The scheduler controls how noise is gradually removed to form the image."}),
@@ -251,10 +258,11 @@ class LanPaint_KSampler():
     def sample(self, model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=1.0, LanPaint_NumSteps=5, LanPaint_PromptMode = "Image First", LanPaint_Info=""):
 
         model.LanPaint_StepSize = 0.15
-        model.LanPaint_Lambda = 8.0
-        model.LanPaint_Beta = 1.0
+        model.LanPaint_Lambda = 16.0
+        model.LanPaint_Beta = 1.
         model.LanPaint_NumSteps = LanPaint_NumSteps
         model.LanPaint_Friction = 15.
+        model.LanPaint_EarlyStop = 1
         if LanPaint_PromptMode == "Image First":
             model.LanPaint_cfg_BIG = cfg
         else:
@@ -268,7 +276,7 @@ class LanPaint_KSamplerAdvanced:
                     {"model": ("MODEL",),
                     "add_noise": (["enable", "disable"], ),
                     "noise_seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
-                    "steps": ("INT", {"default": 50, "min": 1, "max": 10000}),
+                    "steps": ("INT", {"default": 30, "min": 1, "max": 10000}),
                     "cfg": ("FLOAT", {"default": 5.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}),
                     "sampler_name": (KSAMPLER_NAMES, ),
                     "scheduler": (comfy.samplers.KSampler.SCHEDULERS, ),
@@ -279,11 +287,12 @@ class LanPaint_KSamplerAdvanced:
                     "end_at_step": ("INT", {"default": 10000, "min": 0, "max": 10000}),
                     "return_with_leftover_noise": (["disable", "enable"], ),
                 "LanPaint_NumSteps": ("INT", {"default": 5, "min": 0, "max": 100, "tooltip": "The number of steps for the Langevin dynamics, representing the turns of thinking per step."}),
-                "LanPaint_Lambda": ("FLOAT", {"default": 8., "min": 0.1, "max": 50.0, "step": 0.1, "round": 0.1, "tooltip": "The bidirectional guidance scale. Higher values align with known regions more closely, but may result in instability."}),
+                "LanPaint_Lambda": ("FLOAT", {"default": 16., "min": 0.1, "max": 50.0, "step": 0.1, "round": 0.1, "tooltip": "The bidirectional guidance scale. Higher values align with known regions more closely, but may result in instability."}),
                 "LanPaint_StepSize": ("FLOAT", {"default": 0.15, "min": 0.0001, "max": 1., "step": 0.01, "round": 0.001, "tooltip": "The step size for the Langevin dynamics. Higher values result in faster convergence but may be unstable."}),
                 "LanPaint_Beta": ("FLOAT", {"default": 1., "min": 0.0001, "max": 5, "step": 0.1, "round": 0.1, "tooltip": "The step size ratio between masked / unmasked regions. Lower value can compensate high values of LanPaint_Lambda."}),
                 "LanPaint_Friction": ("FLOAT", {"default": 15, "min": 0., "max": 50.0, "step": 0.1, "round": 0.1, "tooltip": "The friction parameter for fast langevin, lower values result in faster convergence but may be unstable."}),
                 "LanPaint_PromptMode": (["Image First", "Prompt First"], {"tooltip": "Image First: emphasis image quality, Prompt First: emphasis prompt following"}),
+                "LanPaint_EarlyStop": ("INT", {"default": 1, "min": 0, "max": 10000, "tooltip": "The number of steps to stop the LanPaint early, useful for preventing the image from irregular patterns."}),
                 "LanPaint_Info": ("STRING", {"default": "LanPaint KSampler Adv. For more info, visit https://github.com/scraed/LanPaint. If you find it useful, please give a star ⭐️!", "multiline": True}),
                      },
                 }
@@ -293,7 +302,7 @@ class LanPaint_KSamplerAdvanced:
 
     CATEGORY = "sampling"
 
-    def sample(self, model, add_noise, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, start_at_step, end_at_step, return_with_leftover_noise, denoise=1.0, LanPaint_StepSize=0.05, LanPaint_Lambda=5, LanPaint_Beta=1, LanPaint_NumSteps=5, LanPaint_Friction=5, LanPaint_PromptMode = "Image First", LanPaint_Info=""):
+    def sample(self, model, add_noise, noise_seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, start_at_step, end_at_step, return_with_leftover_noise, denoise=1.0, LanPaint_StepSize=0.05, LanPaint_Lambda=5, LanPaint_Beta=1, LanPaint_NumSteps=5, LanPaint_Friction=5, LanPaint_PromptMode = "Image First", LanPaint_EarlyStop = 1, LanPaint_Info=""):
         force_full_denoise = True
         if return_with_leftover_noise == "enable":
             force_full_denoise = False
@@ -305,6 +314,7 @@ class LanPaint_KSamplerAdvanced:
         model.LanPaint_Beta = LanPaint_Beta
         model.LanPaint_NumSteps = LanPaint_NumSteps
         model.LanPaint_Friction = LanPaint_Friction
+        model.LanPaint_EarlyStop = LanPaint_EarlyStop
         if LanPaint_PromptMode == "Image First":
             model.LanPaint_cfg_BIG = cfg
         else:
@@ -314,11 +324,233 @@ 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", {"tooltip": "Image before inpaint"}),
+                "image2": ("IMAGE", {"tooltip": "Image after inpaint"}),
+                "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
+        # check the size of mask and image1, image2, if not the same, assert error
+        if image1.shape[1] != image2.shape[1] or image1.shape[2] != image2.shape[2]:
+            raise ValueError("Make sure your image size is a multiple of 8. Otherwise the mask will not be aligned with the output image.")
+        
+        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
+
+class Noise_EmptyNoise:
+    def generate_noise(self, latent):
+        return torch.zeros_like(latent["samples"])
+
+class Noise_RandomNoise:
+    def __init__(self, seed):
+        self.seed = seed
+    def generate_noise(self, latent):
+        torch.manual_seed(self.seed)
+        return torch.randn_like(latent["samples"])
+
+# Custom sampler implementation mimmicking base comfy nodes_custom_sampler.py
+class LanPaint_SamplerCustom:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required":
+                    {"model": ("MODEL",),
+                     "add_noise": ("BOOLEAN", {"default": True}),
+                     "noise_seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, "control_after_generate": True}),
+                     "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step": 0.1, "round": 0.01}),
+                     "positive": ("CONDITIONING",),
+                     "negative": ("CONDITIONING",),
+                     "sampler": ("SAMPLER",),
+                     "sigmas": ("SIGMAS",),
+                     "latent_image": ("LATENT",),
+                     "LanPaint_NumSteps": ("INT", {"default": 5, "min": 0, "max": 100, "tooltip": "Number of steps for Langevin dynamics, representing turns of thinking per step."}),
+                     "LanPaint_PromptMode": (["Image First", "Prompt First"], {"tooltip": "Image First: prioritizes image quality; Prompt First: prioritizes prompt adherence."}),
+                     "LanPaint_Info": ("STRING", {"default": "LanPaint Custom Sampler. For more info, visit https://github.com/scraed/LanPaint. If you find it useful, please give a star ⭐️!", "multiline": True}),
+                      }
+               }
+
+    RETURN_TYPES = ("LATENT", "LATENT")
+    RETURN_NAMES = ("output", "denoised_output")
+    FUNCTION = "sample"
+    CATEGORY = "sampling/custom_sampling"
+
+    def sample(self, model, sampler, sigmas, add_noise, noise_seed, cfg, positive, negative, latent_image, LanPaint_NumSteps, LanPaint_PromptMode, LanPaint_Info=""):
+        model.LanPaint_StepSize = 0.15
+        model.LanPaint_Lambda = 16.0
+        model.LanPaint_Beta = 1.
+        model.LanPaint_NumSteps = LanPaint_NumSteps
+        model.LanPaint_Friction = 15.
+        model.LanPaint_EarlyStop = 1
+        if LanPaint_PromptMode == "Image First":
+            model.LanPaint_cfg_BIG = cfg
+        else:
+            model.LanPaint_cfg_BIG = 0 * cfg - 0.5
+        with override_sample_function():
+            latent = latent_image.copy()
+            latent_image = latent["samples"]
+            latent_image = comfy.sample.fix_empty_latent_channels(model, latent_image)
+            latent["samples"] = latent_image
+
+            if not add_noise:
+                noise = Noise_EmptyNoise().generate_noise(latent)
+            else:
+                noise = Noise_RandomNoise(noise_seed).generate_noise(latent)
+
+            noise_mask = None
+            if "noise_mask" in latent:
+                noise_mask = latent["noise_mask"]
+
+            x0_output = {}
+            callback = latent_preview.prepare_callback(model, sigmas.shape[-1] - 1, x0_output)
+            disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED
+
+            samples = comfy.sample.sample_custom(model, noise, cfg, sampler, sigmas, positive, negative, latent_image,noise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=noise_seed)
+
+            out = latent.copy()
+            out["samples"] = samples
+            if "x0" in x0_output:
+                out_denoised = latent.copy()
+                out_denoised["samples"] = model.model.process_latent_out(x0_output["x0"].cpu())
+            else:
+                out_denoised = out
+            return (out, out_denoised)
+        
+class LanPaint_SamplerCustomAdvanced:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required":
+                    {"noise": ("NOISE",),
+                     "guider": ("GUIDER",),
+                     "sampler": ("SAMPLER",),
+                     "sigmas": ("SIGMAS",),
+                     "latent_image": ("LATENT",),
+                     "start_at_step": ("INT", {"default": 0, "min": 0, "max": 10000}),
+                     "end_at_step": ("INT", {"default": 10000, "min": 0, "max": 10000}),
+                     "return_with_leftover_noise": (["disable", "enable"], ),
+                     "LanPaint_NumSteps": ("INT", {"default": 5, "min": 0, "max": 100, "tooltip": "Number of steps for Langevin dynamics, representing turns of thinking per step."}),
+                     "LanPaint_Lambda": ("FLOAT", {"default": 16.0, "min": 0.1, "max": 50.0, "step": 0.1, "tooltip": "Bidirectional guidance scale. Higher values align with known regions but may cause instability."}),
+                     "LanPaint_StepSize": ("FLOAT", {"default": 0.15, "min": 0.0001, "max": 1.0, "step": 0.01, "tooltip": "Step size for Langevin dynamics. Higher values speed convergence but may be unstable."}),
+                     "LanPaint_Beta": ("FLOAT", {"default": 1.0, "min": 0.0001, "max": 5.0, "step": 0.1, "tooltip": "Step size ratio between masked/unmasked regions. Lower values balance high Lambda."}),
+                     "LanPaint_Friction": ("FLOAT", {"default": 15.0, "min": 0.0, "max": 50.0, "step": 0.1, "tooltip": "Friction parameter for fast Langevin. Lower values speed convergence but may be unstable."}),
+                     "LanPaint_PromptMode": (["Image First", "Prompt First"], {"tooltip": "Image First: prioritizes image quality; Prompt First: prioritizes prompt adherence."}),
+                     "LanPaint_EarlyStop": ("INT", {"default": 1, "min": 0, "max": 10000, "tooltip": "Steps to stop LanPaint early, preventing irregular patterns."}),
+                     "LanPaint_Info": ("STRING", {"default": "LanPaint Custom Sampler Adv. For more info, visit https://github.com/scraed/LanPaint. If you find it useful, please give a star ⭐️!", "multiline": True}),
+                    }
+               }
+
+    RETURN_TYPES = ("LATENT", "LATENT")
+    RETURN_NAMES = ("output", "denoised_output")
+    FUNCTION = "sample"
+    CATEGORY = "sampling/custom_sampling"
+
+    def sample(self, noise, guider, sampler, sigmas, latent_image, start_at_step, end_at_step, return_with_leftover_noise, LanPaint_NumSteps, LanPaint_Lambda, LanPaint_StepSize, LanPaint_Beta, LanPaint_Friction, LanPaint_PromptMode, LanPaint_EarlyStop, LanPaint_Info=""):
+        force_full_denoise = True
+        if end_at_step <= start_at_step:
+            raise ValueError('end_at_step must be larger than start_at_step')
+        if return_with_leftover_noise == "enable":
+            force_full_denoise = False
+        model = guider.model_patcher
+        model.LanPaint_StepSize = LanPaint_StepSize
+        model.LanPaint_Lambda = LanPaint_Lambda
+        model.LanPaint_Beta = LanPaint_Beta
+        model.LanPaint_NumSteps = LanPaint_NumSteps
+        model.LanPaint_Friction = LanPaint_Friction
+        model.LanPaint_EarlyStop = LanPaint_EarlyStop
+        if LanPaint_PromptMode == "Image First":
+            model.LanPaint_cfg_BIG = guider.cfg
+        else:
+            model.LanPaint_cfg_BIG = 0 * guider.cfg - 0.5
+        with override_sample_function():
+            latent = latent_image.copy()
+            latent_image_samples = latent["samples"]
+            latent_image_samples = comfy.sample.fix_empty_latent_channels(model, latent_image_samples)
+            latent["samples"] = latent_image_samples
+
+            noise_mask = None
+            if "noise_mask" in latent:
+                noise_mask = latent["noise_mask"]
+
+            # From base comfy samplers.py
+            if end_at_step is not None and end_at_step < (len(sigmas) - 1):
+                sigmas = sigmas[:end_at_step + 1]
+                if force_full_denoise:
+                    sigmas[-1] = 0
+
+            if start_at_step is not None:
+                if start_at_step < (len(sigmas) - 1):
+                    sigmas = sigmas[start_at_step:]
+                else:
+                    if latent_image is not None:
+                        return latent_image
+                    else:
+                        return torch.zeros_like(noise)
+
+            x0_output = {}
+            callback = latent_preview.prepare_callback(model, sigmas.shape[-1] - 1, x0_output)
+            disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED
+
+            samples = guider.sample( noise.generate_noise(latent), latent_image_samples, sampler, sigmas, denoise_mask=noise_mask, callback=callback,disable_pbar=disable_pbar, seed=noise.seed
+            )
+
+            samples = samples.to(comfy.model_management.intermediate_device())
+            out = latent.copy()
+            out["samples"] = samples
+            if "x0" in x0_output:
+                out_denoised = latent.copy()
+                out_denoised["samples"] = model.model.process_latent_out(x0_output["x0"].cpu())
+            else:
+                out_denoised = out
+            return (out, out_denoised)
+
 # 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_SamplerCustom" : LanPaint_SamplerCustom,
+    "LanPaint_SamplerCustomAdvanced" : LanPaint_SamplerCustomAdvanced,
+    "LanPaint_MaskBlend": MaskBlend,
 #    "LanPaint_UpSale_LatentNoiseMask": LanPaint_UpSale_LatentNoiseMask,
 }
 
@@ -326,5 +558,8 @@ NODE_CLASS_MAPPINGS = {
 NODE_DISPLAY_NAME_MAPPINGS = {
     "LanPaint_KSampler": "LanPaint KSampler",
     "LanPaint_KSamplerAdvanced": "LanPaint KSampler (Advanced)",
+    "LanPaint_SamplerCustom" : "LanPaint Sampler Custom",
+    "LanPaint_SamplerCustomAdvanced" : "LanPaint Sampler Custom (Advanced)",
+    "LanPaint_MaskBlend": "LanPaint Mask Blend",
 #    "LanPaint_UpSale_LatentNoiseMask": "LanPaint UpSale Latent Noise Mask"
 }

src/LanPaint/utils.py

@@ -240,7 +240,6 @@ class StochasticHarmonicOscillator:
             tuple: (y(t), v(t))
         """
         
-
         dummyzero = y0.new_zeros(1) # convert scalar to tensor with same device and dtype as y0
         Delta = self.Delta + dummyzero
         Gamma_hat = self.Gamma * t + dummyzero
@@ -253,8 +252,8 @@ class StochasticHarmonicOscillator:
         EE = 1 - Gamma_hat * zeta_2
 
         if v0 is None:
-            #v0 = torch.randn_like(y0) * D / 2 ** 0.5
-            v0 = (C - A * y0)/Gamma**0.5
+            v0 = torch.randn_like(y0) * D / 2 ** 0.5
+            #v0 = (C - A * y0)/Gamma**0.5
         
         # Calculate mean position and velocity
         term1 = (1 - zeta_1) * (C * t - A * t * y0) + zeta_2 * (Gamma ** 0.5) * v0 * t
@@ -273,16 +272,30 @@ class StochasticHarmonicOscillator:
         cov_matrix[..., 0, 1] = cov_yv
         cov_matrix[..., 1, 0] = cov_yv  # symmetric
         cov_matrix[..., 1, 1] = cov_vv
+
+
         
+        # Compute the Cholesky decomposition to get scale_tril
+        #scale_tril = torch.linalg.cholesky(cov_matrix)
+        scale_tril = torch.zeros(*batch_shape, 2, 2, device=y0.device, dtype=y0.dtype)
+        tol = 1e-8
+        cov_yy = torch.clamp( cov_yy, min = tol )
+        sd_yy = torch.sqrt( cov_yy )
+        inv_sd_yy = 1/(sd_yy)
+
+        scale_tril[..., 0, 0] = sd_yy
+        scale_tril[..., 0, 1] = 0.
+        scale_tril[..., 1, 0] = cov_yv * inv_sd_yy
+        scale_tril[..., 1, 1] = torch.clamp( cov_vv - cov_yv**2 / cov_yy, min = tol ) ** 0.5
+        # check if it matches torch.linalg.
+        #assert torch.allclose(torch.linalg.cholesky(cov_matrix), scale_tril, atol = 1e-4, rtol = 1e-4 )
         # Sample correlated noise from multivariate normal
         mean = torch.zeros(*batch_shape, 2, device=y0.device, dtype=y0.dtype)
         mean[..., 0] = y_mean
         mean[..., 1] = v_mean
         new_yv = torch.distributions.MultivariateNormal(
             loc=mean,
-            covariance_matrix=cov_matrix
+            scale_tril=scale_tril
         ).sample()
 
-        return new_yv[...,0], new_yv[...,1]
-
-
+        return new_yv[...,0], new_yv[...,1]