update version

README.md

@@ -1,26 +1,55 @@
 # 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, aiming for seamless and accurate results. 
+![Inpainting Result 13](https://github.com/scraed/LanPaint/blob/master/examples/InpaintChara_13.jpg) 
 
-
-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.
+This is the official implementation of ["Lanpaint: Training-Free Diffusion Inpainting with Exact and Fast Conditional Inference"](https://arxiv.org/abs/2502.03491). 
 
 ## Features
 
-- 🎨 **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
-- 🎯 **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.
+- **Universal Compatibility** – Works instantly with almost any model (SD 1.5, XL, 3.5, Flux, HiDream, or custom LoRAs) and ControlNet.  
+- **No Training Needed** – Works out of the box with your existing model.  
+- **Easy to Use** – Same workflow as standard ComfyUI KSampler.  
+- **Flexible Masking** – Supports any mask shape, size, or position for inpainting/outpainting.  
+- **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. 
 
 ## How It Works  
+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 **two-way alignment** between masked and unmasked areas. It continuously evaluates:  
+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?"*  
 - *"Do the existing elements support the new creation?"*  
-
 Based on this evaluation, LanPaint iteratively updates the noise in both the masked and unmasked regions.  
   
+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/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.
+    - Fix performance issue on Flux and SD 3.5
 - 2025/04/16
     - Added Primary HiDream support
 - 2025/03/22
@@ -29,60 +58,61 @@ Based on this evaluation, LanPaint iteratively updates the noise in both the mas
 - 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
+## 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)
 ![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)
+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 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)
 
-### 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, 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, 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, 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, 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)
-### Example 6: Character Consistency (Side View Generation) (LanPaint K Sampler, 5 steps of thinking)
+You need to follow the ComfyUI version of [SD 3.5 workflow](https://comfyui-wiki.com/en/tutorial/advanced/stable-diffusion-3-5-comfyui-workflow) to download and install the model.
+
+### Example Flux: 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: Prompt First mode is disabled on Flux. As it does not use CFG guidance.)
+
+### Example SDXL 0: 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) 
 
 (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 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 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)
 
+### 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) 
+[Model Used in This Example](https://civitai.com/models/1188071?modelVersionId=1408658) 
+### Example SDXL 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 SDXL 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 SDXL 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 SDXL 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)
 
+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:**  
@@ -101,20 +131,6 @@ 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
 
@@ -127,15 +143,16 @@ Same as default ComfyUI KSampler - simply replace with LanPaint KSampler nodes.
 
 ## Basic Sampler
 ![Samplers](https://github.com/scraed/LanPaint/blob/master/Nodes.JPG)  
+
+- LanPaint KSampler: The most basic and easy to use sampler for inpainting.
+- LanPaint KSampler (Advanced): Full control of all parameters.
+
 ### LanPaint KSampler
 Simplified interface with recommended defaults:
 
-- Steps: 50+ recommended
-- 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)
-
-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.
-
+- 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, 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:
@@ -143,12 +160,13 @@ Full parameter control:
 
 | 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 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) |
+| `Steps` | 0-100 | Total steps of diffusion sampling. Higher means better inpainting. Recommend 20-50. |
+| `LanPaint_NumSteps` | 0-20 | Reasoning iterations per denoising step ("thinking depth"). Easy task: 2-5. Hard task: 5-10 |
+| `LanPaint_Lambda` | 0.1-50 | Content alignment strength (higher = stricter). Recommend 4.0 - 10.0 |
+| `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 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.
 
@@ -157,32 +175,31 @@ For detailed descriptions of each parameter, simply hover your mouse over the co
 ## LanPaint KSampler (Advanced) Tuning Guide
 For challenging inpainting tasks:  
 
-1️⃣ **Primary Adjustments**:
-- Decrease **LanPaint_endsigma** increase **LanPaint_NumSteps** (thinking iterations) if the inpainted area is not seamless.
-
+1️⃣ **Boost Quality**
+Increase total number of sampling steps, **LanPaint_NumSteps** (thinking iterations) or **LanPaint_Lambda** if the inpainted result does not meet your expectations.
   
-2️⃣ **Secondary Tweaks**:  
-- 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.
+2️⃣ **Boost Speed**
+If you want better results but still need fewer steps, consider:
+    - **Increasing LanPaint_StepSize** to speed up the thinking process.
+    - **Decreasing LanPaint_Friction** to make the Langevin dynamics converges more faster.
     
-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 like blurry/wired texture.
+3️⃣ **Fix Unstability**:  
+If you find the results have wired texture, try
+- Reduce **LanPaint_Friction** to make the Langevin dynamics more stable. 
+- Reduce **LanPaint_StepSize** to use smaller step size.
+- Reduce **LanPaint_Beta** if you are using a high lambda value.
 
 ⚠️ **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.  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
+- Provide inference code on without GUI.
 
 ## Contribute
 
 - 2025/03/06: Bug Fix for str not callable error and unpack error. Big thanks to [jamesWalker55](https://github.com/jamesWalker55) and [EricBCoding](https://github.com/EricBCoding).
 
-Help us improve LanPaint! 🚀 **Report bugs**, share **example cases**, or contribute your **personal parameter settings** to benefit the community. 
 
 ## Citation
 

pyproject.toml

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

src/LanPaint/lanpaint.py

@@ -0,0 +1,145 @@
+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
+        self.chara_lamb = Lambda
+        self.IS_FLUX = IS_FLUX
+        self.IS_FLOW = IS_FLOW
+        self.step_size = StepSize
+        self.inner_model = Model
+        self.friction = Friction
+        self.chara_beta = Beta
+        
+    def __call__(self, x, latent_image, noise, sigma, latent_mask, current_times, model_options, seed, n_steps=None):
+        self.latent_image = latent_image
+        self.noise = noise
+        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]
+        # 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 )
+        else:
+            x_t = x / ( 1+VE_Sigma[:, None,None,None]**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(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)  
+        if IS_FLUX or IS_FLOW:
+            x = x_t / ( abt[:, None,None,None]**0.5 + (1-abt[:, None,None,None])**0.5 )
+        else:
+            x = x_t * ( 1+VE_Sigma[:, None,None,None]**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)
+        out = out * (1-latent_mask) + self.latent_image * latent_mask
+        return out
+
+    def score_model(self, x_t, y, mask, abt, sigma, tflow, model_options, seed):
+        
+        lamb = self.chara_lamb
+
+        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)
+        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)
+
+        score_x = -(x_t - x_0)
+        score_y =  - (1 + lamb) * ( x_t - y )  + lamb * (x_t - x_0_BIG)  
+        return score_x * (1 - mask) + score_y * mask
+    def sigma_x(self, abt):
+        # the time scale for the x_t update
+        return abt**0
+    def sigma_y(self, abt):
+        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):
+        # prepare the step size and time parameters
+        with torch.autocast(device_type=x_t.device.type, dtype=torch.float32):
+            step_sizes = self.prepare_step_size(current_times, step_size, sigma_x, sigma_y)
+            sigma, abt, dtx, dty, Gamma_x, Gamma_y, A_x, A_y, D_x, D_y = step_sizes
+        # print('mask',mask.device)
+        if torch.mean(dtx) <= 0.:
+            return x_t, args
+        # -------------------------------------------------------------------------
+        # Compute the Langevin dynamics update in variance perserving notation
+        # -------------------------------------------------------------------------
+        x0 = self.x0_evalutation(x_t, score, sigma, args)
+        C = abt**0.5 * x0 / (1-abt)
+        A = A_x * (1-mask) + A_y * mask
+        D = D_x * (1-mask) + D_y * mask
+        dt = dtx * (1-mask) + dty * mask
+        Gamma = Gamma_x * (1-mask) + Gamma_y * mask
+
+
+
+        if args is None:
+            #v = torch.zeros_like(x_t)
+            v = None
+        else:
+            v, = 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 )
+  
+        return x_t, (v,)
+
+    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]
+        # Compute time step (dtx, dty) for x and y branches.
+        dtx = 2 * step_size * sigma_x
+        dty = 2 * step_size * sigma_y
+        
+        # -------------------------------------------------------------------------
+        # Define friction parameter Gamma_hat for each branch.
+        # Using dtx**0 provides a tensor of the proper device/dtype.
+
+        Gamma_hat_x = self.friction **2 * self.step_size * sigma_x / 0.1 * sigma**0
+        Gamma_hat_y = self.friction **2 * self.step_size * sigma_y / 0.1 * sigma**0
+        #print("Gamma_hat_x", torch.mean(Gamma_hat_x).item(), "Gamma_hat_y", torch.mean(Gamma_hat_y).item())
+        # 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_x = A_t_x / (dtx/2)
+        A_y = A_t_y / (dty/2)
+        Gamma_x = Gamma_hat_x / (dtx/2)
+        Gamma_y = Gamma_hat_y / (dty/2)
+
+        #D_x = (2 * (1 + sigma**2) )**0.5
+        #D_y = (2 * (1 + sigma**2) )**0.5
+        D_x = (2 * abt**0 )**0.5
+        D_y = (2 * abt**0 )**0.5
+        return sigma, abt, dtx/2, dty/2, Gamma_x, Gamma_y, A_x, A_y, D_x, D_y
+
+
+
+    def x0_evalutation(self, x_t, score, sigma, args):
+        x0 = x_t + score(x_t)
+        return x0

src/LanPaint/nodes.py

@@ -9,6 +9,8 @@ from functools import partial
 from comfy.utils import repeat_to_batch_size
 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."))
 
@@ -89,10 +91,16 @@ class KSamplerX0Inpaint:
 
         # 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  ))
+            Flow_t = sigma
+            abt = (1 - Flow_t)**2 / ((1 - Flow_t)**2 + Flow_t**2 )
+            VE_Sigma = Flow_t / (1 - Flow_t)
+            #print("t", torch.mean( sigma ).item(), "VE_Sigma", torch.mean( VE_Sigma ).item())
+            
+
         else:
-            LanPaint_Sigma = sigma 
+            VE_Sigma = sigma 
+            abt = 1/( 1+VE_Sigma**2 )
+            Flow_t = (1-abt)**0.5 / ( (1-abt)**0.5 + abt**0.5  )
 
         if denoise_mask is not None:
             if "denoise_mask_function" in model_options:
@@ -101,63 +109,15 @@ class KSamplerX0Inpaint:
             denoise_mask = (denoise_mask > 0.5).float()
 
             latent_mask = 1 - denoise_mask
+            current_times = (VE_Sigma, abt, Flow_t)
 
-            abt = 1/( 1+LanPaint_Sigma**2 )
+            current_step = torch.argmin( torch.abs( self.sigmas - torch.mean(sigma) ) )
+            total_steps = len(self.sigmas)-1
 
-            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.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 
+            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:
-                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 = (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
-            
-            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 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[:, 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)[:, 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
+                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)
         
@@ -173,181 +133,7 @@ class KSamplerX0Inpaint:
                 callback({"i": current_step, "denoised": out, "x": x})
     
         return out
-    def mid_times(self, current_times, step_size):
-        sigma, abt = current_times
-        tt = torch.log(1+sigma**2)
-        tt_mid = torch.max( tt - step_size, tt*0 )
-        sigma_mid = (torch.exp(tt_mid) - 1) ** 0.5
-        sigma_mid_prev = sigma_mid
-        # find the closest sigma to sigma_mid from self.sigmas
-        #sigma_mid = self.model_sigmas[torch.argmin(torch.abs(self.model_sigmas - sigma_mid))]
 
-        abt_mid = 1/(1+sigma_mid**2)
-        return sigma_mid, abt_mid
-    def score_model(self, x_t, y, mask, abt, sigma, model_options, seed):
-        
-        # the score function for the Langevin dynamics
-        lamb = self.chara_lamb
-        beta = self.chara_beta * (1-abt)**0.5
-
-        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
-        
-        score_x = -e_t
-        score_y = - (1 + lamb) * ( x_t/ ((1 + sigma**2) ** 0.5 *(1 - abt)**0.5) - abt**0.5 /(1 - abt)**0.5 * y )  + lamb *  e_t_BIG
-        return score_x * (1 - mask) + score_y * mask
-    def sigma_x(self, abt):
-        # the time scale for the x_t update
-        return abt**0
-    def sigma_y(self, abt):
-        # the time scale for the y_t update
-        if self.beta_scale == "shrink":
-            beta = self.chara_beta * (1-abt)**0.5
-        elif self.beta_scale == "dual_shrink":
-            beta = self.chara_beta * (1-abt)**0.5 * abt ** 0.5
-        elif self.beta_scale == "back_shrink":
-            beta = self.chara_beta * abt ** 0.5
-        else:
-            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
-
-
-
-
-        #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, torch.squeeze(dtx))
-        sigma_mid_y, abt_mid_y = self.mid_times(current_times, torch.squeeze(dty))
-
-        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
-
-        # -------------------------------------------------------------------------
-        # A: Update epsilon (score estimate and noise initialization)
-        # -------------------------------------------------------------------------
-        # Compute the score-based epsilon (scaled as sqrt(1-abt))
-        score_model = score(x_t)
-        eps_model = -score_model 
-
-
-
-        # Initialize epsilon and Z if not provided in args.
-        if args is None:
-            eps = eps_model
-            Z = torch.randn_like(x_t)
-        else:
-            eps, Z = args
-
-        # -------------------------------------------------------------------------
-        # B: Update epsilon mean dynamics and compute the mid-point in z-space.
-        # -------------------------------------------------------------------------
-        # Compute the weighted combination term for epsilon mean update:
-        # term = (2/Γ_hat)*(1-exp(-0.5*Γ_hat))
-        term_x = 2.0 / (Gamma_hat_x + 1e-4) * (1 - torch.exp(-0.5 * Gamma_hat_x))
-        term_y = 2.0 / (Gamma_hat_y + 1e-4) * (1 - torch.exp(-0.5 * Gamma_hat_y))
-        eps_bar_x = term_x * eps + (1 - term_x) * eps_model
-        eps_bar_y = term_y * eps + (1 - term_y) * eps_model
-        # Combine branches according to mask.
-        eps_bar = eps_bar_x * (1 - mask) + eps_bar_y * mask
-
-        # Form the denoised epsilon using self.alpha (assumed to be 1/Ψ)
-        eps_denoise = self.alpha * eps_bar + (1 - self.alpha) * eps_model
-
-
-        # tamed
-        eps_model_x = eps_denoise* (1 - mask)
-        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), 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
-
-
-        # Update the mean epsilon for the next step:
-        eps_x = eps * torch.exp(-0.5 * Gamma_hat_x) + eps_model * (1 - torch.exp(-0.5 * Gamma_hat_x))
-        eps_y = eps * torch.exp(-0.5 * Gamma_hat_y) + eps_model * (1 - torch.exp(-0.5 * Gamma_hat_y))
-        eps = eps_x * (1 - mask) + eps_y * mask
-
-        # Transform x to z using z = x * sqrt(1+sigma^2). Here we have already set x to z to avoid floating point stability issue.
-        z_t = x_t #* (1 + sigma**2) ** 0.5
-
-        # Compute the mid-point update in z-space for each branch:
-        z_mid_x = z_t + eps_denoise * (sigma_mid_x - sigma)
-        z_mid_y = z_t + eps_denoise * (sigma_mid_y - sigma)
-        z_mid = z_mid_x * (1 - mask) + z_mid_y * mask
-
-        # -------------------------------------------------------------------------
-        # C: Update noise terms and finalize the x update.
-        # -------------------------------------------------------------------------
-        # Generate auxiliary noise terms.
-        Z_q     = torch.randn_like(x_t)
-        Z_q_avg = torch.randn_like(x_t)
-        Z_z     = torch.randn_like(x_t)
-
-        # Update Z for each branch:
-        Z_x = torch.exp(-0.5 * Gamma_hat_x) * Z + (1 - torch.exp(-Gamma_hat_x)) ** 0.5 * Z_q
-        Z_y = torch.exp(-0.5 * Gamma_hat_y) * Z + (1 - torch.exp(-Gamma_hat_y)) ** 0.5 * Z_q
-        Z_next = Z_x * (1 - mask) + Z_y * mask
-
-        # Compute the combined noise update following the scheme:
-        Z_comb_x = (
-            (1 - torch.exp(-Gamma_hat_x / 2)) / torch.sqrt(Gamma_hat_x + 1e-4) *
-            (Z + torch.sqrt(torch.tanh(Gamma_hat_x / 4)) * Z_q)
-            + torch.sqrt(1 - (4 / (Gamma_hat_x + 1e-4)) * torch.tanh(Gamma_hat_x / 4)) * Z_q_avg
-        )
-        Z_comb_y = (
-            (1 - torch.exp(-Gamma_hat_y / 2)) / torch.sqrt(Gamma_hat_y + 1e-4) *
-            (Z + torch.sqrt(torch.tanh(Gamma_hat_y / 4)) * Z_q)
-            + torch.sqrt(1 - (4 / (Gamma_hat_y + 1e-4)) * torch.tanh(Gamma_hat_y / 4)) * Z_q_avg
-        )
-        Z_comb = Z_comb_x * (1 - mask) + Z_comb_y * mask
-
-        # Combine with an additional noise term using self.alpha.
-        Z_comb = self.alpha ** 0.5 * Z_comb + (1 - self.alpha) ** 0.5 * Z_z
-
-        # Compute the change in sigma (dsigma = sqrt(sigma^2 - sigma_mid^2)).
-        dsigma_x = sigma * torch.sqrt(1 - (sigma_mid_x / sigma) ** 2)
-        dsigma_y = sigma * torch.sqrt(1 - (sigma_mid_y / sigma) ** 2)
-        dsigma = dsigma_x * (1 - mask) + dsigma_y * mask
-
-        # Final z update.
-        z_final = z_mid + Z_comb * dsigma
-
-        # Transform back to x-space: x = z / sqrt(1+sigma^2)
-        x_t = z_final #/ (1 + sigma**2) ** 0.5
-
-        return x_t, (eps, Z_next)
 # Custom sampler class extending ComfyUI's KSAMPLER for LanPaint
 class KSAMPLER(comfy.samplers.KSAMPLER):
     def sample(self, model_wrap, sigmas, extra_args, callback, noise, latent_image=None, denoise_mask=None, disable_pbar=False):
@@ -363,27 +149,23 @@ class KSAMPLER(comfy.samplers.KSAMPLER):
             model_k.noise = noise
 
         IS_FLUX = model_wrap.inner_model.model_type == ModelType.FLUX
-
+        IS_FLOW = model_wrap.inner_model.model_type == ModelType.FLOW
         # 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
-        model_k.n_steps = model_wrap.model_patcher.LanPaint_NumSteps
-        model_k.friction = model_wrap.model_patcher.LanPaint_Friction
-        model_k.alpha = model_wrap.model_patcher.LanPaint_Alpha
-        model_k.tamed = model_wrap.model_patcher.LanPaint_Tamed
-        model_k.beta_scale = model_wrap.model_patcher.LanPaint_BetaScale
-        model_k.step_size_schedule = model_wrap.model_patcher.LanPaint_StepSizeSchedule
-        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))
+
+        model_k.PaintMethod = LanPaint(model_k.inner_model, 
+                                       model_wrap.model_patcher.LanPaint_NumSteps,
+                                       model_wrap.model_patcher.LanPaint_Friction,
+                                       model_wrap.model_patcher.LanPaint_Lambda,
+                                       model_wrap.model_patcher.LanPaint_Beta,
+                                       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
@@ -440,7 +222,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
@@ -449,7 +236,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."}),
@@ -457,9 +244,9 @@ class LanPaint_KSampler():
                 "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": 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}),
+                "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_PromptMode": (["Image First", "Prompt First"], {"tooltip": "Image First: emphasis image quality, Prompt First: emphasis prompt following"}),
+                "LanPaint_Info": ("STRING", {"default": "LanPaint KSampler. For more info, visit https://github.com/scraed/LanPaint. If you find it useful, please give a star ⭐️!", "multiline": True}),
                   }
         }
 
@@ -470,20 +257,18 @@ 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_EndSigma = 3., LanPaint_Info=""):
-        model.LanPaint_StepSize = 0.5
+    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.2
+        model.LanPaint_Beta = 1.0
         model.LanPaint_NumSteps = LanPaint_NumSteps
-        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 = LanPaint_EndSigma
-        model.LanPaint_cfg_BIG = -0.5
+        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():
             return nodes.common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=denoise)
 class LanPaint_KSamplerAdvanced:
@@ -493,7 +278,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, ),
@@ -503,20 +288,14 @@ 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": 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": 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}),
+                "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_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}),
                      },
                 }
 
@@ -525,7 +304,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_Alpha=1, LanPaint_Tamed=0., LanPaint_BetaScale="fixed", LanPaint_StepSizeSchedule = "const", LanPaint_StepTimeSchedule = "shrink",  LanPaint_StartSigma=20, LanPaint_EndSigma=0, LanPaint_cfg_BIG = 5., 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
@@ -537,14 +316,12 @@ class LanPaint_KSamplerAdvanced:
         model.LanPaint_Beta = LanPaint_Beta
         model.LanPaint_NumSteps = LanPaint_NumSteps
         model.LanPaint_Friction = LanPaint_Friction
-        model.LanPaint_Alpha = LanPaint_Alpha
-        model.LanPaint_Tamed = LanPaint_Tamed
-        model.LanPaint_BetaScale = LanPaint_BetaScale
-        model.LanPaint_StepSizeSchedule = LanPaint_StepSizeSchedule
-        model.LanPaint_StepTimeSchedule = LanPaint_StepTimeSchedule
-        model.LanPaint_StartSigma = LanPaint_StartSigma
-        model.LanPaint_EndSigma = LanPaint_EndSigma
-        model.LanPaint_cfg_BIG = LanPaint_cfg_BIG
+        model.LanPaint_EarlyStop = LanPaint_EarlyStop
+        if LanPaint_PromptMode == "Image First":
+            model.LanPaint_cfg_BIG = cfg
+        else:
+            model.LanPaint_cfg_BIG = 0*cfg - 0.5
+
         with override_sample_function():
             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)
 

src/LanPaint/utils.py

@@ -0,0 +1,288 @@
+import torch
+def epxm1_x(x):
+    # Compute the (exp(x) - 1) / x term with a small value to avoid division by zero.
+    result = torch.special.expm1(x) / x
+    # replace NaN or inf values with 0
+    result = torch.where(torch.isfinite(result), result, torch.zeros_like(result))
+    mask = torch.abs(x) < 1e-2
+    result = torch.where(mask, 1 + x/2. + x**2 / 6., result)
+    return result
+def epxm1mx_x2(x):
+    # Compute the (exp(x) - 1 - x) / x**2 term with a small value to avoid division by zero.
+    result = (torch.special.expm1(x) - x) / x**2
+    # replace NaN or inf values with 0
+    result = torch.where(torch.isfinite(result), result, torch.zeros_like(result))
+    mask = torch.abs(x**2) < 1e-2
+    result = torch.where(mask, 1/2. + x/6 + x**2 / 24 + x**3 / 120, result)
+    return result
+
+def expm1mxmhx2_x3(x):
+    # Compute the (exp(x) - 1 - x - x**2 / 2) / x**3 term with a small value to avoid division by zero.
+    result = (torch.special.expm1(x) - x - x**2 / 2) / x**3
+    # replace NaN or inf values with 0
+    result = torch.where(torch.isfinite(result), result, torch.zeros_like(result))
+    mask = torch.abs(x**3) < 1e-2
+    result = torch.where(mask, 1/6 + x/24 + x**2 / 120 + x**3 / 720 + x**4 / 5040, result)
+    return result
+
+def exp_1mcosh_GD(gamma_t, delta):
+    """
+    Compute e^(-Γt) * (1 - cosh(Γt√Δ))/ ( (Γt)**2 Δ )
+    
+    Parameters:
+    gamma_t: Γ*t term (could be a scalar or tensor)
+    delta: Δ term (could be a scalar or tensor)
+    
+    Returns:
+    Result of the computation with numerical stability handling
+    """
+    # Main computation
+    is_positive = delta > 0
+    sqrt_abs_delta = torch.sqrt(torch.abs(delta))
+    gamma_t_sqrt_delta = gamma_t * sqrt_abs_delta
+    numerator_pos =  torch.exp(-gamma_t) - (torch.exp(gamma_t * (sqrt_abs_delta - 1)) + torch.exp(gamma_t * (-sqrt_abs_delta - 1))) / 2
+    numerator_neg = torch.exp(-gamma_t) * ( 1 -  torch.cos(gamma_t * sqrt_abs_delta ) )
+    numerator = torch.where(is_positive, numerator_pos, numerator_neg)
+    result =  numerator / (delta * gamma_t**2 )
+    # Handle NaN/inf cases
+    result = torch.where(torch.isfinite(result), result, torch.zeros_like(result))
+    # Handle numerical instability for small delta
+    mask = torch.abs(gamma_t_sqrt_delta**2) < 5e-2
+    taylor = ( -0.5  - gamma_t**2 / 24 * delta - gamma_t**4 / 720 * delta**2 ) * torch.exp(-gamma_t)
+    result = torch.where(mask, taylor, result)
+    return result
+
+def exp_sinh_GsqrtD(gamma_t, delta):
+    """
+    Compute e^(-Γt) * sinh(Γt√Δ) / (Γt√Δ)
+
+    Parameters:
+    gamma_t: Γ*t term (could be a scalar or tensor)
+    delta: Δ term (could be a scalar or tensor)
+
+    Returns:
+    Result of the computation with numerical stability handling
+    """
+    # Main computation
+    is_positive = delta > 0
+    sqrt_abs_delta = torch.sqrt(torch.abs(delta))
+    gamma_t_sqrt_delta = gamma_t * sqrt_abs_delta
+    numerator_pos =  (torch.exp(gamma_t * (sqrt_abs_delta - 1)) - torch.exp(gamma_t * (-sqrt_abs_delta - 1))) / 2
+    denominator_pos = gamma_t_sqrt_delta
+    result_pos = numerator_pos / gamma_t_sqrt_delta
+    result_pos = torch.where(torch.isfinite(result_pos), result_pos, torch.zeros_like(result_pos))
+
+    # Taylor expansion for small gamma_t_sqrt_delta
+    mask = torch.abs(gamma_t_sqrt_delta) < 1e-2
+    taylor = ( 1  + gamma_t**2 / 6 * delta + gamma_t**4 / 120 * delta**2 ) * torch.exp(-gamma_t)
+    result_pos = torch.where(mask, taylor, result_pos)
+
+    # Handle negative delta
+    result_neg = torch.exp(-gamma_t) * torch.special.sinc(gamma_t_sqrt_delta/torch.pi)
+    result = torch.where(is_positive, result_pos, result_neg)
+    return result
+
+def exp_cosh(gamma_t, delta):
+    """
+    Compute e^(-Γt) * cosh(Γt√Δ)
+
+    Parameters:
+    gamma_t: Γ*t term (could be a scalar or tensor)
+    delta: Δ term (could be a scalar or tensor)
+
+    Returns:
+    Result of the computation with numerical stability handling
+    """
+    exp_1mcosh_GD_result = exp_1mcosh_GD(gamma_t, delta) # e^(-Γt) * (1 - cosh(Γt√Δ))/ ( (Γt)**2 Δ )
+    result = torch.exp(-gamma_t) - gamma_t**2 * delta * exp_1mcosh_GD_result
+    return result
+def exp_sinh_sqrtD(gamma_t, delta):
+    """
+    Compute e^(-Γt) * sinh(Γt√Δ) / √Δ
+    Parameters:
+    gamma_t: Γ*t term (could be a scalar or tensor)
+    delta: Δ term (could be a scalar or tensor)
+    Returns:
+    Result of the computation with numerical stability handling
+    """
+    exp_sinh_GsqrtD_result = exp_sinh_GsqrtD(gamma_t, delta) # e^(-Γt) * sinh(Γt√Δ) / (Γt√Δ)
+    result = gamma_t * exp_sinh_GsqrtD_result
+    return result
+
+
+
+def zeta1(gamma_t, delta):
+    # Compute hyperbolic terms and exponential
+    half_gamma_t = gamma_t / 2
+    exp_cosh_term = exp_cosh(half_gamma_t, delta)
+    exp_sinh_term = exp_sinh_sqrtD(half_gamma_t, delta)
+
+    
+    # Main computation
+    numerator = 1 - (exp_cosh_term + exp_sinh_term)
+    denominator = gamma_t * (1 - delta) / 4
+    result = 1 - numerator / denominator
+    
+    # Handle numerical instability
+    result = torch.where(torch.isfinite(result), result, torch.zeros_like(result))
+    
+    # Taylor expansion for small x (similar to your epxm1Dx approach)
+    mask = torch.abs(denominator) < 5e-3
+    term1 = epxm1_x(-gamma_t)
+    term2 = epxm1mx_x2(-gamma_t)
+    term3 = expm1mxmhx2_x3(-gamma_t)
+    taylor = term1 + (1/2.+ term1-3*term2)*denominator + (-1/6. + term1/2 - 4 * term2 + 10 * term3) * denominator**2
+    result = torch.where(mask, taylor, result)
+    
+    return result
+
+def exp_cosh_minus_terms(gamma_t, delta):
+    """
+    Compute E^(-tΓ) * (Cosh[tΓ] - 1 - (Cosh[tΓ√Δ] - 1)/Δ) / (tΓ(1 - Δ))
+    
+    Parameters:
+    gamma_t: Γ*t term (could be a scalar or tensor)
+    delta: Δ term (could be a scalar or tensor)
+    
+    Returns:
+    Result of the computation with numerical stability handling
+    """
+    exp_term = torch.exp(-gamma_t)
+    # Compute individual terms
+    exp_cosh_term = exp_cosh(gamma_t, gamma_t**0) - exp_term # E^(-tΓ) (Cosh[tΓ] - 1) term
+    exp_cosh_delta_term = - gamma_t**2 * exp_1mcosh_GD(gamma_t, delta)  # E^(-tΓ) (Cosh[tΓ√Δ] - 1)/Δ term
+    
+    #exp_1mcosh_GD e^(-Γt) * (1 - cosh(Γt√Δ))/ ( (Γt)**2 Δ )
+    # Main computation
+    numerator = exp_cosh_term - exp_cosh_delta_term
+    denominator = gamma_t * (1 - delta)
+    
+    result = numerator / denominator
+    
+    # Handle numerical instability
+    result = torch.where(torch.isfinite(result), result, torch.zeros_like(result))
+    
+    # Taylor expansion for small gamma_t and delta near 1
+    mask = (torch.abs(denominator) < 1e-1)
+    exp_1mcosh_GD_term = exp_1mcosh_GD(gamma_t, delta**0)
+    taylor = (
+       gamma_t*exp_1mcosh_GD_term + 0.5 * gamma_t * exp_sinh_GsqrtD(gamma_t, delta**0) 
+       - denominator / 4 * ( 0.5 * exp_cosh(gamma_t, delta**0) - 4 * exp_1mcosh_GD_term - 5 /2 * exp_sinh_GsqrtD(gamma_t, delta**0) )
+    )
+    result = torch.where(mask, taylor, result)
+    
+    return result
+
+
+def zeta2(gamma_t, delta):
+    half_gamma_t = gamma_t / 2
+    return exp_sinh_GsqrtD(half_gamma_t, delta)
+
+def sig11(gamma_t, delta):
+    return 1 - torch.exp(-gamma_t) + gamma_t**2 * exp_1mcosh_GD(gamma_t, delta) + exp_sinh_sqrtD(gamma_t, delta)
+
+
+def Zcoefs(gamma_t, delta):
+    Zeta1 = zeta1(gamma_t, delta)
+    Zeta2 = zeta2(gamma_t, delta)
+    
+    sq_total = 1 - Zeta1 + gamma_t * (delta - 1) * (Zeta1 - 1)**2 / 8
+    amplitude = torch.sqrt(sq_total)
+    Zcoef1 = ( gamma_t**0.5 * Zeta2 / 2 **0.5 ) / amplitude
+    Zcoef2 = Zcoef1 * gamma_t *( - 2 * exp_1mcosh_GD(gamma_t, delta)  / sig11(gamma_t, delta)  ) ** 0.5 
+    #cterm = exp_cosh_minus_terms(gamma_t, delta)
+    #sterm = exp_sinh_sqrtD(gamma_t, delta**0) + exp_sinh_sqrtD(gamma_t, delta)
+    #Zcoef3 = 2 * torch.sqrt(  cterm / ( gamma_t * (1 - delta) * cterm + sterm ) )
+    Zcoef3 = torch.sqrt( torch.maximum(1 - Zcoef1**2 - Zcoef2**2, sq_total.new_zeros(sq_total.shape)) )
+
+    return Zcoef1 * amplitude, Zcoef2 * amplitude, Zcoef3 * amplitude, amplitude
+
+def Zcoefs_asymp(gamma_t, delta):
+    A_t = (gamma_t * (1 - delta) )/4
+    return epxm1_x(- 2 * A_t)
+
+class StochasticHarmonicOscillator:
+    """
+    Simulates a stochastic harmonic oscillator governed by the equations:
+        dy(t) = q(t) dt
+        dq(t) = -Γ A y(t) dt + Γ C dt + Γ D dw(t) - Γ q(t) dt
+
+    Also define v(t) = q(t) / √Γ, which is numerically more stable.
+        
+    Where:
+        y(t) - Position variable
+        q(t) - Velocity variable
+        Γ - Damping coefficient
+        A - Harmonic potential strength
+        C - Constant force term
+        D - Noise amplitude
+        dw(t) - Wiener process (Brownian motion)
+    """
+    def __init__(self, Gamma, A, C, D):
+        self.Gamma = Gamma
+        self.A = A
+        self.C = C
+        self.D = D
+        self.Delta = 1 - 4 * A / Gamma
+    def sig11(self, gamma_t, delta):
+        return 1 - torch.exp(-gamma_t) + gamma_t**2 * exp_1mcosh_GD(gamma_t, delta) + exp_sinh_sqrtD(gamma_t, delta)
+    def sig22(self, gamma_t, delta):
+        return 1- zeta1(2*gamma_t, delta) + 2 * gamma_t * exp_1mcosh_GD(gamma_t, delta) 
+    def dynamics(self, y0, v0, t):
+        """
+        Calculates the position and velocity variables at time t.
+
+        Parameters:
+            y0 (float): Initial position
+            v0 (float): Initial velocity v(0) = q(0) / √Γ
+            t (float): Time at which to evaluate the dynamics
+        Returns:
+            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
+        A = self.A + dummyzero
+        C = self.C + dummyzero
+        D = self.D + dummyzero
+        Gamma = self.Gamma + dummyzero
+        zeta_1 = zeta1( Gamma_hat, Delta) 
+        zeta_2 = zeta2( Gamma_hat, Delta)
+        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
+        
+        # Calculate mean position and velocity
+        term1 = (1 - zeta_1) * (C * t - A * t * y0) + zeta_2 * (Gamma ** 0.5) * v0 * t
+        y_mean = term1 + y0
+        v_mean =  (1 - EE)*(C - A * y0) / (Gamma ** 0.5) + (EE - A * t * (1 - zeta_1)) * v0
+        
+        cov_yy = D**2 * t * self.sig22(Gamma_hat, Delta)
+        cov_vv = D**2 * self.sig11(Gamma_hat, Delta) / 2
+        cov_yv = (zeta2(Gamma_hat, Delta) * Gamma_hat * D ) **2 / 2 / (Gamma ** 0.5)
+
+        # sample new position and velocity with multivariate normal distribution
+
+        batch_shape = y0.shape
+        cov_matrix = torch.zeros(*batch_shape, 2, 2, device=y0.device, dtype=y0.dtype)
+        cov_matrix[..., 0, 0] = cov_yy
+        cov_matrix[..., 0, 1] = cov_yv
+        cov_matrix[..., 1, 0] = cov_yv  # symmetric
+        cov_matrix[..., 1, 1] = cov_vv
+        
+        # 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
+        ).sample()
+
+        return new_yv[...,0], new_yv[...,1]
+
+