Update pyproject.toml

Update Github Action for Publishing to Comfy Registry

.github/workflows/publish_action.yml

@@ -7,14 +7,18 @@ on:
     paths:
       - "pyproject.toml"
 
+permissions:
+  issues: write
+
 jobs:
   publish-node:
     name: Publish Custom Node to registry
     runs-on: ubuntu-latest
+    if: ${{ github.repository_owner == 'scraed' }}
     steps:
       - name: Check out code
         uses: actions/checkout@v4
       - name: Publish Custom Node
-        uses: Comfy-Org/publish-node-action@main
+        uses: Comfy-Org/publish-node-action@v1
         with:
           personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }} ## Add your own personal access token to your Github Repository secrets and reference it here.

README.md

@@ -1,37 +1,91 @@
 # 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. We encourage you to try it out and share your feedback through issues or discussions, as your input will help us enhance the algorithm's performance and stability.
+Unlock precise inpainting without additional training. LanPaint lets the model "think" through multiple iterations before denoising, aiming for seamless and accurate results.
+
+
+We encourage you to try it out and share your feedback through issues or discussions, as your input will help us enhance the algorithm's performance and stability.
 
 ## Features
 
-- 🎨 **Zero-Training Inpainting** - Works immediately with ANY SD model, even custom models you've trained yourself
+- 🎨 **Zero-Training Inpainting** - Works immediately with ANY SD model (with/without ControlNet), and Flux model! even custom models you've trained yourself
 - 🛠️ **Simple Integration** - Same workflow as standard ComfyUI KSampler
-- 🚀 **Quality Enhancements** - High quality and seamless inpainting
+- 🎯 **True Blank-Slate Generation** - No need to set default denoise at 0.7 (preserving 30% original pixels in masks) used in conventional methods: 100% **new content creation**, No "painting over" existing content.
+- 🌈 **Not only inpaint**: You can even use it as a simple way to generate consistent characters.
+
+## How It Works  
+
+LanPaint introduces **two-way alignment** between masked and unmasked areas. It continuously evaluates:  
+- *"Does the new content make sense with the existing elements?"*  
+- *"Do the existing elements support the new creation?"*  
+
+Based on this evaluation, LanPaint iteratively updates the noise in both the masked and unmasked regions.  
+  
+## Updates
+- 2025/04/16
+    - Added Primary HiDream support
+- 2025/03/22
+    - Added Primary Flux support
+    - Added Tease Mode
+- 2025/03/10
+    - LanPaint has received a major update! All examples now use the LanPaint K Sampler, offering a simplified interface with enhanced performance and stability.
 
 ## Example Results
-All examples use random seed 0 to ensure fair comparison.
-### Example 1: Basket to Basket Ball (LanPaint K Sampler, It is fast).
+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)
+![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 install [ComfyUI GGUF](https://github.com/city96/ComfyUI-GGUF) in order to load the models. Make sure you have the latest (nightly at 2025/04/16) comfyui installed. The following models are needed for Hidream:
+- [clip_g_hidream.safetensors](https://huggingface.co/Comfy-Org/HiDream-I1_ComfyUI/blob/main/split_files/text_encoders/clip_g_hidream.safetensors)
+- [clip_l_hidream.safetensors](https://huggingface.co/Comfy-Org/HiDream-I1_ComfyUI/blob/main/split_files/text_encoders/clip_l_hidream.safetensors)
+- [T5 GGUF](https://huggingface.co/city96/t5-v1_1-xxl-encoder-gguf/tree/main)
+- [Llama 3.1](https://huggingface.co/bartowski/Meta-Llama-3.1-8B-Instruct-GGUF/tree/main)
+- [Flux VAE](https://huggingface.co/StableDiffusionVN/Flux/blob/main/Vae/flux_vae.safetensors)
+
+
+### Example 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) 
 [Model Used in This Example](https://civitai.com/models/1188071?modelVersionId=1408658) 
-### Example 2: White Shirt to Blue Shirt (LanPaint K Sampler (Advanced), it is slower.)
+### Example 2: White Shirt to Blue Shirt (LanPaint K Sampler, 5 steps of thinking)
 ![Inpainting Result 2](https://github.com/scraed/LanPaint/blob/master/examples/InpaintChara_05.jpg)  
 [View Workflow & Masks](https://github.com/scraed/LanPaint/tree/master/examples/Example_2)
 [Model Used in This Example](https://civitai.com/models/1188071?modelVersionId=1408658)
-### Example 3: Smile to Sad (LanPaint K Sampler (Advanced))
+### Example 3: Smile to Sad (LanPaint K Sampler, 5 steps of thinking)
 ![Inpainting Result 3](https://github.com/scraed/LanPaint/blob/master/examples/InpaintChara_06.jpg)  
 [View Workflow & Masks](https://github.com/scraed/LanPaint/tree/master/examples/Example_3)
 [Model Used in This Example](https://civitai.com/models/133005/juggernaut-xl)
-### Example 4: Damage Restoration (LanPaint K Sampler (Advanced))
-![Inpainting Result 3](https://github.com/scraed/LanPaint/blob/master/examples/InpaintChara_07.jpg)  
+### Example 4: Damage Restoration (LanPaint K Sampler, 5 steps of thinking)
+![Inpainting Result 4](https://github.com/scraed/LanPaint/blob/master/examples/InpaintChara_07.jpg)  
 [View Workflow & Masks](https://github.com/scraed/LanPaint/tree/master/examples/Example_4)
 [Model Used in This Example](https://civitai.com/models/133005/juggernaut-xl)
-### Example 5: Huge Damage Restoration (LanPaint K Sampler (Advanced))
-![Inpainting Result 3](https://github.com/scraed/LanPaint/blob/master/examples/InpaintChara_08.jpg)  
+### Example 5: Huge Damage Restoration (LanPaint K Sampler, 20 steps of thinking)
+![Inpainting Result 5](https://github.com/scraed/LanPaint/blob/master/examples/InpaintChara_08.jpg)  
 [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) (The key is increase LanPaint_stepsize to 0.5)
+[Model Used in This Example](https://civitai.com/models/133005/juggernaut-xl)
+### Example 6: Character Consistency (Side View Generation) (LanPaint K Sampler, 5 steps of thinking)
+![Inpainting Result 6](https://github.com/scraed/LanPaint/blob/master/examples/InpaintChara_09.jpg)  
+[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) 
 
-**How to Use These Examples:**  
+(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!)
+
+### Example 7: Flux Model InPaint(LanPaint K Sampler, 5 steps of thinking)
+![Inpainting Result 7](https://github.com/scraed/LanPaint/blob/master/examples/InpaintChara_10.jpg)  
+[View Workflow & Masks](https://github.com/scraed/LanPaint/tree/master/examples/Example_7)
+[Model Used in This Example](https://huggingface.co/Comfy-Org/flux1-dev/blob/main/flux1-dev-fp8.safetensors) 
+(Note: Use CFG scale 1.0 for Flux as it don't use CFG. LanPaint_cfg_BIG is also disabled on Flux)
+
+
+
+
+## **How to Use These Examples:**  
 1. Navigate to the **example** folder (i.e example_1) by clicking **View Workflow & Masks**, download all pictures.  
 2. Drag **InPainted_Drag_Me_to_ComfyUI.png** into ComfyUI to load the workflow.  
 3. Download the required model from Civitai by clicking **Model Used in This Example**.  
@@ -49,7 +103,7 @@ Compare and explore the results from each method!
 
 ## Quickstart
 
-1. **Install ComfyUI**: Follow the official [ComfyUI installation guide](https://docs.comfy.org/get_started) to set up ComfyUI on your system.  
+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.  
@@ -71,26 +125,30 @@ Same as default ComfyUI KSampler - simply replace with LanPaint KSampler nodes.
 - LanPaint requires binary masks (values of 0 or 1) without opacity or smoothing. To ensure compatibility, set the mask's **opacity and hardness to maximum** in your mask editor. During inpainting, any mask with smoothing or gradients will automatically be converted to a binary mask.
 - LanPaint relies heavily on your text prompts to guide inpainting - explicitly describe the content you want generated in the masked area. If results show artifacts or mismatched elements, counteract them with targeted negative prompts.
 
-### Basic Sampler
+## Basic Sampler
 ![Samplers](https://github.com/scraed/LanPaint/blob/master/Nodes.JPG)  
-**LanPaint KSampler**  
+### LanPaint KSampler
 Simplified interface with recommended defaults:
 
 - Steps: 50+ recommended
-- LanPaint NumSteps: 1-10 (complexity of edits)
-- Built-in parameter presets
+- LanPaint NumSteps: The turns of thinking before denoising. Recommend 5 for most of tasks.
+- LanPaint EndSigma: The noise level below which thinking is disabled. Recommend 0.6 for realistic style (tested on Juggernaut-xl), 3.0 for anime style (tested on Animagine XL 4.0)
 
-**LanPaint KSampler (Advanced)**  
+The default settings are tested on Animagine XL 4.0 and Juggernaut-xl. Other model might need some paramter tuning. Please raise issue or share your own setting if it doesn't work on your model.
+
+
+### LanPaint KSampler (Advanced)
 Full parameter control:
-## Key Parameters
+**Key Parameters**
 
 | Parameter | Range | Description |
 |-----------|-------|-------------|
 | `Steps` | 0-100 | Total steps of diffusion sampling. Higher means better inpainting. Recommend 50. |
 | `LanPaint_NumSteps` | 0-20 | Reasoning iterations per denoising step ("thinking depth"). Easy task: 1-2. Hard task: 5-10 |
-| `LanPaint_Lambda` | 0.1-50 | Content alignment strength (higher = stricter). Recommend 6.0 |
-| `LanPaint_StepSize` | 0.1-1.0 | The StepSize of each thinking step. Recommend 0.1 for most cases, 0.5 for some difficult cases. |
-| `LanPaint_cfg_BIG` | 0-20 | CFG scale used when aligning masked and unmasked region (higher = better alignment). Recommend 8 for seamless inpaint (i.e limbs, faces), 0-1 for character consistency (i.e multiple view) |
+| `LanPaint_Lambda` | 0.1-50 | Content alignment strength (higher = stricter). Recommend 8.0 |
+| `LanPaint_StepSize` | 0.1-1.0 | The StepSize of each thinking step. Recommend 0.5. |
+| `LanPaint_EndSigma` | 0.0-20.0 | The noise level below which thinking is disabled. recommend 0.3 - 3. High value is faster, but may damage quality. Low value gives more thinking but might make the output blurry. |
+| `LanPaint_cfg_BIG` | -20-20 | CFG scale used when aligning masked and unmasked region (positive value tends to ignores promts, negative value enhances prompts.). Recommend 8 for seamless inpaint (i.e limbs, faces) when prompt is not important. -0.5 when prompt is important, like character consistency (i.e multiple view) |
 
 For detailed descriptions of each parameter, simply hover your mouse over the corresponding input field to view tooltips with additional information.
 
@@ -100,18 +158,25 @@ For detailed descriptions of each parameter, simply hover your mouse over the co
 For challenging inpainting tasks:  
 
 1️⃣ **Primary Adjustments**:
-- Increase **steps**, **LanPaint_StepSize**，**LanPaint_NumSteps** (thinking iterations), and **LanPaint_cfg_BIG** (guidance scale).
+- Decrease **LanPaint_endsigma** increase **LanPaint_NumSteps** (thinking iterations) if the inpainted area is not seamless.
+
   
 2️⃣ **Secondary Tweaks**:  
-- Boost **LanPaint_Lambda** (spatial constraint strength) or **LanPaint_StepSize** (denoising aggressiveness).
+- Boost **LanPaint_Lambda** (bidirectional guidance scale) will force the masked/unmasked region to align more closely.
+- If the output is blurry, increase **LanPaint_endsigma** to turn off thinking at the end of denoising. OR decrease **LanPaint_StepSize** to decrease thinking step size.
+- If prompt is not that important, try increase **LanPaint_cfg_BIG**(cfg scale used for unmasked region, default -0.5 ) to 8 for better inpainting.
     
 3️⃣ **Balance Speed vs Stability**:  
 - Reduce **LanPaint_Friction** to prioritize faster results with fewer "thinking" steps (*may risk instability*).  
-- Increase **LanPaint_Tamed** (noise normalization onto a sphere) or **LanPaint_Alpha** (constraint the friction of underdamped Langevin dynamics) to suppress artifacts.
+- Increase **LanPaint_Tamed** (noise normalization onto a sphere) or **LanPaint_Alpha** (constraint the friction of underdamped Langevin dynamics) to suppress artifacts like blurry/wired texture.
 
 ⚠️ **Notes**:  
 - Optimal parameters vary depending on the **model** and the **size of the inpainting area**.  
-- For effective tuning, **fix the seed** and adjust parameters incrementally while observing the results. This helps isolate the impact of each setting.  
+- For effective tuning, **fix the seed** and adjust parameters incrementally while observing the results. This helps isolate the impact of each setting.  Better to do it with a batche of images to avoid overfitting on a single image.
+
+## ToDo
+- SD 3.5 also have problems
+- Try Implement Detailer
 
 ## Contribute
 
@@ -119,6 +184,21 @@ For challenging inpainting tasks:
 
 Help us improve LanPaint! 🚀 **Report bugs**, share **example cases**, or contribute your **personal parameter settings** to benefit the community. 
 
+## Citation
+
+```
+@misc{zheng2025lanpainttrainingfreediffusioninpainting,
+      title={Lanpaint: Training-Free Diffusion Inpainting with Exact and Fast Conditional Inference}, 
+      author={Candi Zheng and Yuan Lan and Yang Wang},
+      year={2025},
+      eprint={2502.03491},
+      archivePrefix={arXiv},
+      primaryClass={eess.IV},
+      url={https://arxiv.org/abs/2502.03491}, 
+}
+```
+
+
 
 
 

pyproject.toml

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

src/LanPaint/nodes.py

@@ -8,6 +8,7 @@ import latent_preview
 from functools import partial
 from comfy.utils import repeat_to_batch_size
 from comfy.samplers import *
+from comfy.model_base import ModelType
 # Monkey patch comfy.samplers module by importing with absolute package path
 #exec(inspect.getsource(comfy.samplers).replace("from .", "from comfy."))
 
@@ -76,10 +77,23 @@ class KSamplerX0Inpaint:
         self.sigmas = sigmas
         self.model_sigmas = torch.cat( (torch.tensor([0.], device = sigmas.device) , torch.tensor( self.inner_model.model_patcher.get_model_object("model_sampling").sigmas, device = sigmas.device) ) )
         self.model_sigmas = torch.tensor( self.model_sigmas, dtype = self.sigmas.dtype )
-    def __call__(self, x, sigma, denoise_mask, model_options={}, seed=None):
+    def __call__(self, x, sigma, denoise_mask, model_options={}, seed=None,**kwargs):
+        ### For 1.5 and XL model
         # x is x_t in the notation of variance exploding diffusion model, x_t = x_0 + sigma * noise
         # sigma is the noise level
-        # print what is inside model_options
+        ### For flux model 
+        # x is rectified flow x_t = sigma * noise + (1.0 - sigma) * x_0
+
+        IS_FLUX = self.inner_model.inner_model.model_type == ModelType.FLUX
+        IS_FLOW = self.inner_model.inner_model.model_type == ModelType.FLOW
+
+        # unify the notations into variance exploding diffusion model
+        if IS_FLUX or IS_FLOW:
+            LanPaint_Sigma = sigma / ( torch.maximum( 1 - sigma , sigma*0 + 5e-2  ))
+            self.LanPaint_Sigmas = self.sigmas / ( torch.maximum( 1 - self.sigmas , self.sigmas*0 + 5e-2  ))
+        else:
+            LanPaint_Sigma = sigma 
+
         if denoise_mask is not None:
             if "denoise_mask_function" in model_options:
                 denoise_mask = model_options["denoise_mask_function"](sigma, denoise_mask, extra_options={"model": self.inner_model, "sigmas": self.sigmas})
@@ -88,44 +102,76 @@ class KSamplerX0Inpaint:
 
             latent_mask = 1 - denoise_mask
 
-            abt = 1/( 1+sigma**2 )
+            abt = 1/( 1+LanPaint_Sigma**2 )
+
+            print("sigma", LanPaint_Sigma, "abt", abt)
+
+            
 
             if self.step_time_schedule == "dual_shrink":
                 step_size = self.step_size * (1 - abt) ** 0.5 * abt ** 0.5
             elif self.step_time_schedule == "follow_sampler":
-                time_ind = torch.argmin(torch.abs(self.sigmas - sigma))
-                times = torch.log( 1+ self.sigmas**2)
+                time_ind = torch.argmin(torch.abs(self.LanPaint_Sigmas - LanPaint_Sigma))
+                times = torch.log( 1+ self.LanPaint_Sigmas**2)
                 time_intervals = times[1:] - times[:-1]
                 time_intervals = time_intervals / time_intervals[0]
                 step_size = time_intervals[time_ind] * self.step_size 
             else:
                 step_size = self.step_size * (1 - abt) ** 0.5
 
+
+            #step_size = self.step_size * (1 - abt) ** b * abt ** a / ( ((a/(a+b))**a*(b/(a+b))**b) )
+            abt_end = 1/( 1+self.end_sigma**2 ) 
+            step_size = self.step_size * (1 -  torch.minimum(abt/abt_end, abt**0) ) ** 0.5
+            step_size = step_size[:, None, None, None]
+
             
-            current_times = (sigma, abt)
+            current_times = (LanPaint_Sigma, abt)
 
             # self.inner_model.inner_model.scale_latent_inpaint returns variance exploding x_t values
             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
-            x_t = x #/ ( 1+sigma**2 )**0.5 # switch to variance perserving x_t values
+            
+            if IS_FLUX or IS_FLOW:
+                x_t = x * ( 1 + LanPaint_Sigma[:, None,None,None])
+            else:
+                x_t = x #/ ( 1+sigma**2 )**0.5 # switch to variance perserving x_t values
+
+
             # 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):
-
-                if sigma > self.start_sigma or sigma < self.end_sigma:
+                if torch.mean(LanPaint_Sigma) > self.start_sigma or torch.mean(LanPaint_Sigma) < self.end_sigma:
+                    
                     break
 
-                score_func = partial( self.score_model, y = self.latent_image, mask = latent_mask, abt = abt, sigma = sigma, model_options = model_options, seed = seed )
+                score_func = partial( self.score_model, y = self.latent_image, mask = latent_mask, abt = abt[:, None,None,None], sigma = LanPaint_Sigma[:, None,None,None], model_options = model_options, seed = seed )
                 if self.step_size_schedule == "linear":
                     step_size_i = step_size * (1 - i/(self.n_steps) )
                 else:
                     step_size_i = step_size 
-                x_t, args = self.langevin_dynamics(x_t, score_func , latent_mask, step_size_i , current_times, sigma_x = self.sigma_x(abt), sigma_y = self.sigma_y(abt), args = args)  
-            x = x_t #* ( 1+sigma**2 )**0.5
+                x_t, args = self.langevin_dynamics(x_t, score_func , latent_mask, step_size_i , current_times, sigma_x = self.sigma_x(abt)[:, None,None,None], sigma_y = self.sigma_y(abt)[:, None,None,None], args = args)  
+            if IS_FLUX or IS_FLOW:
+                x = x_t / ( 1 + LanPaint_Sigma[:, None,None,None] )
+            else:
+                x = x_t #/ ( 1+sigma**2 )**0.5 # switch to variance perserving x_t values
+
             # out is x_0
             out, _ = self.inner_model(x, sigma, model_options=model_options, seed=seed)
             out = out * denoise_mask + self.latent_image * latent_mask
         else:
             out, _ = self.inner_model(x, sigma, model_options=model_options, seed=seed)
+        
+        # Add TAESD preview support - directly use the latent_preview module
+        current_step = model_options.get("i", kwargs.get("i", 0))
+        total_steps = model_options.get("total_steps", 0)
+
+        # Only show preview every few steps to improve performance
+        if current_step % 2 == 0:
+            # Directly call the preview callback if it exists
+            callback = model_options.get("callback", None)
+            if callback is not None:
+                callback({"i": current_step, "denoised": out, "x": x})
+    
         return out
     def mid_times(self, current_times, step_size):
         sigma, abt = current_times
@@ -144,7 +190,13 @@ class KSamplerX0Inpaint:
         lamb = self.chara_lamb
         beta = self.chara_beta * (1-abt)**0.5
 
-        x_0, x_0_BIG = self.inner_model(x_t, sigma, model_options=model_options, seed=seed)
+        IS_FLUX = self.inner_model.inner_model.model_type == ModelType.FLUX
+        IS_FLOW = self.inner_model.inner_model.model_type == ModelType.FLOW
+        if IS_FLUX or IS_FLOW:
+            x_0, x_0_BIG = self.inner_model(x_t / ( 1 + sigma ), sigma[:, 0,0,0] / ( 1 + sigma[:, 0,0,0] ), model_options=model_options, seed=seed)
+        else:
+            x_0, x_0_BIG = self.inner_model(x_t, sigma[:, 0,0,0], model_options=model_options, seed=seed)
+
         e_t = x_t / ((1 - abt) ** 0.5 * (1 + sigma**2) ** 0.5 )- (abt ** 0.5  / (1 - abt) ** 0.5) * x_0
         e_t_BIG = x_t / ((1 - abt) ** 0.5 * (1 + sigma**2) ** 0.5 )- (abt ** 0.5  / (1 - abt) ** 0.5) * x_0_BIG
         
@@ -153,7 +205,7 @@ class KSamplerX0Inpaint:
         return score_x * (1 - mask) + score_y * mask
     def sigma_x(self, abt):
         # the time scale for the x_t update
-        return 1
+        return abt**0
     def sigma_y(self, abt):
         # the time scale for the y_t update
         if self.beta_scale == "shrink":
@@ -163,33 +215,43 @@ class KSamplerX0Inpaint:
         elif self.beta_scale == "back_shrink":
             beta = self.chara_beta * abt ** 0.5
         else:
-            beta = self.chara_beta
+            beta = self.chara_beta * abt ** 0
         return beta
     def langevin_dynamics(self, x_t, score, mask, step_size, current_times, sigma_x=1, sigma_y=0, args=None):
         # -------------------------------------------------------------------------
         # Unpack current times parameters (sigma and abt)
         sigma, abt = current_times
-
+        sigma = sigma[:, None,None,None]
+        abt = abt[:, None,None,None]
         # Compute time step (dtx, dty) for x and y branches.
         dtx = 2 * step_size * sigma_x
         dty = 2 * step_size * sigma_y
 
-        if self.step_time_schedule == "dual_shrink":
-            ref_dt = 0.1 * (1-abt)**0.5 * abt ** 0.5
-        else:
-            ref_dt = 0.1 * (1-abt)**0.5
 
+
+
+        #ref_dt = 0.1 * (1 - abt) ** b * abt ** a / ( ((a/(a+b))**a*(b/(a+b))**b) )
+        abt_end = 1/( 1+self.end_sigma**2 ) 
+        ref_dt = 0.1 * (1 -  torch.minimum(abt/abt_end, abt**0) ) ** 0.5
         # -------------------------------------------------------------------------
         # Define friction parameter Gamma_hat for each branch.
         # Using dtx**0 provides a tensor of the proper device/dtype.
+
         Gamma_hat_x = self.friction * dtx / (1e-4+ 2 * sigma_x * ref_dt)
         Gamma_hat_y = self.friction * dty / (1e-4+ 2 * sigma_y * ref_dt)
 
         # Get mid time parameters (sigma_mid and abt_mid) for each branch.
-        sigma_mid_x, abt_mid_x = self.mid_times(current_times, dtx)
-        sigma_mid_y, abt_mid_y = self.mid_times(current_times, dty)
+        sigma_mid_x, abt_mid_x = self.mid_times(current_times, torch.squeeze(dtx))
+        sigma_mid_y, abt_mid_y = self.mid_times(current_times, torch.squeeze(dty))
 
-        if sigma_mid_x >= sigma or sigma_mid_y >= sigma:
+        sigma_mid_x = sigma_mid_x[:, None,None,None]
+        sigma_mid_y = sigma_mid_y[:, None,None,None]
+        abt_mid_x = abt_mid_x[:, None,None,None]
+        abt_mid_y = abt_mid_y[:, None,None,None]
+
+
+
+        if torch.mean(sigma_mid_x) >= torch.mean(sigma) or torch.mean(sigma_mid_y) >= torch.mean(sigma):
             return x_t, args
 
         # -------------------------------------------------------------------------
@@ -226,9 +288,9 @@ class KSamplerX0Inpaint:
 
         # tamed
         eps_model_x = eps_denoise* (1 - mask)
-        eps_model_x = eps_model_x* (torch.sum(1 - mask, dim = (1,2,3))/torch.sum(eps_model_x**2, dim = (1,2,3))) **0.5 ** torch.minimum(self.tamed*(dtx),sigma**0)#/( 1 + self.tamed*(sigma - sigma_mid_x) * (torch.sum(eps_model_x**2)/torch.sum((1 - mask)))**0.5 )
+        eps_model_x = eps_model_x* (torch.sum(1 - mask, dim = (1,2,3), keepdim = True)/torch.sum(eps_model_x**2, dim = (1,2,3), keepdim = True)) **0.5 ** torch.minimum(self.tamed*(dtx),sigma**0)#/( 1 + self.tamed*(sigma - sigma_mid_x) * (torch.sum(eps_model_x**2)/torch.sum((1 - mask)))**0.5 )
         eps_model_y = eps_denoise* mask
-        eps_model_y = eps_model_y* (torch.sum(mask, dim = (1,2,3))/torch.sum(eps_model_y**2, dim = (1,2,3))) **0.5 ** torch.minimum(self.tamed*(dty),sigma**0)#/( 1 + self.tamed*(sigma - sigma_mid_y) * (torch.sum(eps_model_y**2)/torch.sum(mask))**0.5 )
+        eps_model_y = eps_model_y* (torch.sum(mask, dim = (1,2,3), keepdim = True)/torch.sum(eps_model_y**2, dim = (1,2,3), keepdim = True)) **0.5 ** torch.minimum(self.tamed*(dty),sigma**0)#/( 1 + self.tamed*(sigma - sigma_mid_y) * (torch.sum(eps_model_y**2)/torch.sum(mask))**0.5 )
         eps_denoise = eps_model_x * (1 - mask) + eps_model_y * mask
 
 
@@ -299,7 +361,15 @@ class KSAMPLER(comfy.samplers.KSAMPLER):
             model_k.noise = torch.randn(noise.shape, generator=generator, device="cpu").to(noise.dtype).to(noise.device)
         else:
             model_k.noise = noise
-        model_wrap.cfg_BIG = model_wrap.model_patcher.LanPaint_cfg_BIG
+
+        IS_FLUX = model_wrap.inner_model.model_type == ModelType.FLUX
+
+        # unify the notations into variance exploding diffusion model
+        if IS_FLUX:
+            model_wrap.cfg_BIG = 1.0
+        else:
+            model_wrap.cfg_BIG = model_wrap.model_patcher.LanPaint_cfg_BIG
+
         model_k.step_size = model_wrap.model_patcher.LanPaint_StepSize
         model_k.chara_lamb = model_wrap.model_patcher.LanPaint_Lambda
         model_k.chara_beta = model_wrap.model_patcher.LanPaint_Beta
@@ -312,6 +382,7 @@ class KSAMPLER(comfy.samplers.KSAMPLER):
         model_k.step_time_schedule = model_wrap.model_patcher.LanPaint_StepTimeSchedule
         model_k.start_sigma = model_wrap.model_patcher.LanPaint_StartSigma
         model_k.end_sigma = model_wrap.model_patcher.LanPaint_EndSigma
+
         noise = model_wrap.inner_model.model_sampling.noise_scaling(sigmas[0], noise, latent_image, self.max_denoise(model_wrap, sigmas))
         #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.
@@ -319,10 +390,13 @@ class KSAMPLER(comfy.samplers.KSAMPLER):
         total_steps = len(sigmas) - 1
         if callback is not None:
             k_callback = lambda x: callback(x["i"], x["denoised"], x["x"], total_steps)
-        print("LanPaint KSampler call sampler_function", self.sampler_function)
+        #print("LanPaint KSampler call sampler_function", self.sampler_function)
         # The main loop!
+        #print("##########")
+        #print("Sampling with ", self.sampler_function)
+        #print("##########")
         samples = self.sampler_function(model_k, noise, sigmas, extra_args=extra_args, callback=k_callback, disable=disable_pbar, **self.extra_options)
-        print("LanPaint KSampler end sampler_function")
+        #print("LanPaint KSampler end sampler_function")
         samples = model_wrap.inner_model.model_sampling.inverse_noise_scaling(sigmas[-1], samples)
         return samples
 
@@ -345,8 +419,27 @@ def override_sample_function():
         comfy.samplers.CFGGuider.outer_sample = original_outer_sample
 
 
+class LanPaint_UpSale_LatentNoiseMask:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": { "samples": ("LATENT",),
+                              "scale": ("INT", {"default": 2, "min": 2, "max": 8, "step": 1}),
+                              }}
+    RETURN_TYPES = ("LATENT",)
+    FUNCTION = "set_mask"
 
 
+    CATEGORY = "latent/inpaint"
+
+    def set_mask(self, samples, scale):
+        s = samples.copy()
+        samples = s['samples']
+        # generate a mask with every scaleth pixel set to 1
+        mask = torch.zeros(samples.shape[0], 1, samples.shape[2], samples.shape[3], device=samples.device) + 1
+        mask[:, :, ::scale, ::scale] = 0
+        s["noise_mask"] = mask
+        return (s,)
+
 KSAMPLER_NAMES = ["euler", "dpmpp_2m", "uni_pc"]
 
 class LanPaint_KSampler():
@@ -359,13 +452,14 @@ class LanPaint_KSampler():
                 "steps": ("INT", {"default": 50, "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, {"tooltip": "The scheduler controls how noise is gradually removed to form the image."}),
+                "scheduler": (comfy.samplers.KSampler.SCHEDULERS, {"default": "karras", "tooltip": "The scheduler controls how noise is gradually removed to form the image."}),
                 "positive": ("CONDITIONING", {"tooltip": "The conditioning describing the attributes you want to include in the image."}),
                 "negative": ("CONDITIONING", {"tooltip": "The conditioning describing the attributes you want to exclude from the image."}),
                 "latent_image": ("LATENT", {"tooltip": "The latent image to denoise."}),
                 "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "The amount of denoising applied, lower values will maintain the structure of the initial image allowing for image to image sampling."}),
-                "LanPaint_NumSteps": ("INT", {"default": 10, "min": 0, "max": 20, "tooltip": "The number of steps for the Langevin dynamics, representing the turns of thinking per step."}),  
-                "LanPaint_Info": ("STRING", {"default": "LanPaint KSampler. Recommend steps 50 ( increase steps boosts performance ), LanPaint NumSteps 1-10 depending on the difficulty of task. For more information, visit https://github.com/scraed/LanPaint", "multiline": True}),
+                "LanPaint_NumSteps": ("INT", {"default": 5, "min": 0, "max": 20, "tooltip": "The number of steps for the Langevin dynamics, representing the turns of thinking per step."}),  
+                "LanPaint_EndSigma": ("FLOAT", {"default": 3.0, "min": 0.0, "max": 20.0, "step": 0.01, "tooltip": "The noise level at which the thinking process stops. Higher value give less thinking but helps to deal with blurring when turns of thinking is too high."}),
+                "LanPaint_Info": ("STRING", {"default": "LanPaint KSampler. Recommend steps 50, LanPaint NumSteps 1-20 depending on the difficulty of task. LanPaint_EndSigma = 3.0 for anime style, 0.6 for realistic style. For more information, visit https://github.com/scraed/LanPaint", "multiline": True}),
                   }
         }
 
@@ -376,20 +470,20 @@ class LanPaint_KSampler():
     CATEGORY = "sampling"
     DESCRIPTION = "Uses the provided model, positive and negative conditioning to denoise the latent image."
 
-    def sample(self, model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=1.0, LanPaint_StepSize=0.05, LanPaint_NumSteps=5, LanPaint_Info=""):
-        model.LanPaint_StepSize = 0.3
-        model.LanPaint_Lambda = 6.0
-        model.LanPaint_Beta = 0.6
+    def sample(self, model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=1.0, LanPaint_StepSize=0.05, LanPaint_NumSteps=5, LanPaint_EndSigma = 3., LanPaint_Info=""):
+        model.LanPaint_StepSize = 0.5
+        model.LanPaint_Lambda = 8.0
+        model.LanPaint_Beta = 1.2
         model.LanPaint_NumSteps = LanPaint_NumSteps
-        model.LanPaint_Friction = 10.
-        model.LanPaint_Alpha = 0.5
-        model.LanPaint_Tamed = 0.1
+        model.LanPaint_Friction = 5.
+        model.LanPaint_Alpha = 0.9
+        model.LanPaint_Tamed = 1.
         model.LanPaint_BetaScale = "shrink"
         model.LanPaint_StepSizeSchedule = "linear"
         model.LanPaint_StepTimeSchedule = "shrink"
         model.LanPaint_StartSigma = 20.
-        model.LanPaint_EndSigma = 1.
-        model.LanPaint_cfg_BIG = cfg
+        model.LanPaint_EndSigma = LanPaint_EndSigma
+        model.LanPaint_cfg_BIG = -0.5
         with override_sample_function():
             return nodes.common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise)
 class LanPaint_KSamplerAdvanced:
@@ -409,19 +503,19 @@ class LanPaint_KSamplerAdvanced:
                     "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": 10, "min": 0, "max": 20, "tooltip": "The number of steps for the Langevin dynamics, representing the turns of thinking per step."}),
-                "LanPaint_Lambda": ("FLOAT", {"default": 6., "min": 0.1, "max": 50.0, "step": 0.1, "round": 0.1, "tooltip": "The lambda parameter for the bidirectional guidance. Higher values align with known regions more closely, but may result in instability."}),
-                "LanPaint_StepSize": ("FLOAT", {"default": 0.3, "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": 0.6, "min": 0.0001, "max": 5, "step": 0.1, "round": 0.1, "tooltip": "The beta parameter for the bidirectional guidance. Scale the step size for the known region independently for the Langevin dynamics. Higher values result in faster convergence but may be unstable."}),
-                "LanPaint_Friction": ("FLOAT", {"default": 10., "min": 1., "max": 50.0, "step": 0.1, "round": 0.1, "tooltip": "The friction parameter for the underdamped Langevin dynamics, higher values result in faster convergence but may be unstable."}),
-                "LanPaint_Alpha": ("FLOAT", {"default": 0.5, "min": 0.0001, "max": 1., "step": 0.1, "round": 0.1, "tooltip": "The (rescaled) alpha parameter for the HFHR langevin dynamics, mixes Langevin dynamics and underdamped Langevin dynamics with a friction term. 0 corresponds to Langevin dynamics, 1 corresponds to underdamped Langevin dynamics."}),
+                "LanPaint_NumSteps": ("INT", {"default": 5, "min": 0, "max": 20, "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 lambda parameter for the bidirectional guidance. Higher values align with known regions more closely, but may result in instability."}),
+                "LanPaint_StepSize": ("FLOAT", {"default": 0.5, "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.2, "min": 0.0001, "max": 5, "step": 0.1, "round": 0.1, "tooltip": "The beta parameter for the bidirectional guidance. Scale the step size for the known region independently for the Langevin dynamics. Higher values result in faster convergence but may be unstable."}),
+                "LanPaint_Friction": ("FLOAT", {"default": 5., "min": 1., "max": 50.0, "step": 0.1, "round": 0.1, "tooltip": "The friction parameter for the underdamped Langevin dynamics, higher values result in faster convergence but may be unstable."}),
+                "LanPaint_Alpha": ("FLOAT", {"default": 0.9, "min": 0.0001, "max": 1., "step": 0.1, "round": 0.1, "tooltip": "The (rescaled) alpha parameter for the HFHR langevin dynamics, mixes Langevin dynamics and underdamped Langevin dynamics with a friction term. 0 corresponds to Langevin dynamics, 1 corresponds to underdamped Langevin dynamics."}),
                 "LanPaint_Tamed": ("FLOAT", {"default": 1., "min": 0.000, "max": 20., "step": 0.1, "round": 0.1, "tooltip": "The tame strength for the noise, normalize and projects the noise onto unit sphere to enhance stability."}),
                 "LanPaint_BetaScale": (["shrink", "fixed", "dual_shrink", "back_shrink"], {"default": "shrink", "tooltip": "The beta scale, determines how the beta parameter changes over time. Shrink: beta = beta * (1 - alpha bar) ** 0.5; Fixed: beta = beta; Dual_shrink: beta = beta * (1 - alpha bar) ** 0.5 *  alpha bar ** 0.5; Back_shrink: beta = beta *  alpha bar ** 0.5; Alpha bar: the alpha cumprod."}),
                 "LanPaint_StepSizeSchedule": (["const", "linear"], {"default": "linear", "tooltip": "The step size schedule for the Langevin dynamics, const: constant step size, linear: linearly decreasing step size."}),
                 "LanPaint_StepTimeSchedule": (["shrink", "dual_shrink", "follow_sampler"], {"default": "shrink", "tooltip": "The step size schedule for the first step of Langevin dynamics during diffusion sampling, shrink: step size = step size * (1 - alpha bar) ** 0.5; Dual_shrink: step size = step size * (1 - alpha bar) ** 0.5 *  alpha bar ** 0.5; Follow_sampler: scale with the sampler step size."}),
                 "LanPaint_StartSigma": ("FLOAT", {"default": 20., "min": 0.0001, "max": 20.0, "step": 0.1, "round": 0.1, "tooltip": "Start 'thinking' with Langevin dynamics at this sigma value."}),
-                "LanPaint_EndSigma": ("FLOAT", {"default": 1., "min": 0.000, "max": 20.0, "step": 0.1, "round": 0.1, "tooltip": "Stop 'thinking' with Langevin dynamics at this sigma value."}),
-                "LanPaint_cfg_BIG": ("FLOAT", {"default": 8., "min": 0., "max": 20.0, "step": 0.1, "round": 0.1, "tooltip": "The CFG scale used in the bidirectional guidance (for the known region only). Higher value results in more closely matching the known region."}),
+                "LanPaint_EndSigma": ("FLOAT", {"default": 3., "min": 0.000, "max": 20.0, "step": 0.1, "round": 0.1, "tooltip": "Stop 'thinking' with Langevin dynamics at this sigma value."}),
+                "LanPaint_cfg_BIG": ("FLOAT", {"default": -0.5, "min": -20, "max": 20.0, "step": 0.1, "round": 0.1, "tooltip": "The CFG scale used in the bidirectional guidance (for the known region only). Higher value results in more closely matching the known region."}),
                 "LanPaint_Info": ("STRING", {"default": "LanPaint KSampler Advanced. For difficult tasks, first try increasing steps, LanPaint_NumSteps, and LanPaint_cfg_BIG. Then try increase LanPaint_Lambda or LanPaint_StepSize. Decrease LanPaint_Friction if you want to obtain good results with fewer turns of thinking (LanPaint_NumSteps) at the risk of irregular behavior. Increase LanPaint_Tamed or LanPaint_Alpha can suppress irregular behavior. For more information, visit https://github.com/scraed/LanPaint", "multiline": True}),
                      },
                 }
@@ -460,10 +554,12 @@ class LanPaint_KSamplerAdvanced:
 NODE_CLASS_MAPPINGS = {
     "LanPaint_KSampler": LanPaint_KSampler,
     "LanPaint_KSamplerAdvanced": LanPaint_KSamplerAdvanced,
+#    "LanPaint_UpSale_LatentNoiseMask": LanPaint_UpSale_LatentNoiseMask,
 }
 
 # A dictionary that contains the friendly/humanly readable titles for the nodes
 NODE_DISPLAY_NAME_MAPPINGS = {
     "LanPaint_KSampler": "LanPaint KSampler",
-    "LanPaint_KSamplerAdvanced": "LanPaint KSampler (Advanced)"
+    "LanPaint_KSamplerAdvanced": "LanPaint KSampler (Advanced)",
+#    "LanPaint_UpSale_LatentNoiseMask": "LanPaint UpSale Latent Noise Mask"
 }